from pathlib import Path

import cupy as cp
import cupyx
import cupyx.scipy.fft
import cupyx.scipy.ndimage
import cupyx.scipy.signal
import cv2
import matplotlib
import numpy as np
import scipy
import scipy.signal
from cupyx.scipy.fft import dctn, idctn
from scipy.stats import norm as norms


def skip(f):
    def wrapper(self, **kwargs):
        if 'cond' not in kwargs:
            return f(self, **kwargs)
        if kwargs['cond']:
            del kwargs['cond']
            return f(self, **kwargs)
        else:
            return self

    return wrapper


class ScanData:
    @staticmethod
    def from_buffer(buf: bytes | memoryview, shape, p=np):
        return ScanData(p.frombuffer(buf, dtype=p.int16, count=np.prod(shape)).reshape(shape))

    @staticmethod
    def from_file(filename: Path | str, shape, p=np):
        filename = Path(filename)
        return ScanData.from_buffer(filename.read_bytes(), shape, p)

    @staticmethod
    def from_npy(filename: Path, p=np):
        return ScanData(p.load(filename))

    def __init__(self, mat: cp.ndarray):
        if isinstance(mat, np.ndarray):
            self.p = np
        elif isinstance(mat, cp.ndarray):
            self.p = cp
        else:
            raise Exception('wrong type mat!!')
        self.m = mat
        self.d = mat.shape.__len__()

    def save(self, path: Path):
        self.p.save(path.__str__(), self.m)

    # @property
    # def p(self) -> np:
    #     if self._p == 'np':
    #         return np
    #     if self._p == 'cp':
    #         return cp

    def get(self):
        return self.m

    def apply(self, f):
        f(self.m)
        return self

    def call(self, f, *args, **kwargs):
        return ScanData(f(self.m, *args, **kwargs))

    def copy(self):
        return ScanData(self.m.copy())

    def cpu(self):
        if self.p == cp:
            return ScanData(self.m.get())
        if self.p == np:
            return self.copy()

    def gpu(self):
        if self.p == np:
            return ScanData(cp.asarray(self.m))
        if self.p == cp:
            return self.copy()

    def norm(self):
        m = self.m.astype(self.p.float32)
        m -= m.min()
        mmax = m.max()
        if mmax == 0:
            return ScanData(self.p.zeros_like(m))
        m /= mmax
        return ScanData(m)

    @skip
    def sum(self, axis=0):
        return ScanData(self.m.sum(axis=axis))

    def rotate90(self):
        return ScanData(cv2.rotate(self.m, cv2.ROTATE_90_CLOCKWISE))

    def filter_min(self, t, v=None, bid=False):
        m = self.m
        if v is None:
            v = t
        if bid:
            m[(0 <= m) & (m <= t)] = v
            m[(-t <= m) & (m <= 0)] = -v
        else:
            m[(m <= t)] = v
        return self

    def filter_max(self, t, v=None, bid=False):
        m = self.m.copy()
        if v is None:
            v = t
        if bid:
            m[m >= t] = v
            m[-m >= t] = -v
        else:
            m[m >= t] = v
        return ScanData(m)

    def filter_max_persent(self, t, v=None, bid=False, mmax=None):
        if mmax is None:
            mmax = self.m.max()
        else:
            self.m[0, 0] = mmax
        return self.filter_max(mmax * t, v, bid)

    def astype(self, t):
        return ScanData(self.m.astype(t))

    def dctn_cp(self):
        return ScanData(dctn(self.m))

    def dctn_cp_(self):
        self.m = dctn(self.m)
        return self

    def idctn_cp_(self):
        self.m = idctn(self.m)
        return self

    def dct(self, mmin, mmax):
        dct_ = scipy.fft.dct
        idct = scipy.fft.idct
        if self.p == cp:
            dct_ = cupyx.scipy.fft.dct
            idct = cupyx.scipy.fft.idct
        m_dct = dct_(self.m)
        if self.d == 2:
            m_dct[:, mmax:] = 0
            m_dct[:, :mmin] = 0
        elif self.d == 3:
            m_dct[:, :, mmax:] = 0
            m_dct[:, :, :mmin] = 0
        else:
            raise NotImplementedError()
        return ScanData(idct(m_dct))

    def argrelextrema(self, axis=1):
        arg = scipy.signal.argrelextrema
        m = self.m
        p = self.p
        if p == cp:
            arg = cupyx.scipy.signal.argrelextrema
        rm = p.zeros_like(m)
        indies1 = arg(m, p.greater, axis=axis)
        bl = p.zeros_like(m, dtype=p.bool_)
        indies2 = arg(m, p.less, axis=axis)
        bl2 = p.zeros_like(m, dtype=p.bool_)
        bl[indies1] = True
        bl2[indies2] = True
        i1 = bl & (m > 0)
        i2 = bl2 & (m < 0)
        idx = i1 | i2
        rm[idx] = m[idx].__abs__()
        return ScanData(rm)

    def conv_guass(self, b=0.01, axis=1):
        cv = scipy.ndimage.convolve1d
        m = self.m
        p = self.p
        if p == cp:
            cv = cupyx.scipy.ndimage.convolve1d
        rv = norms(loc=0, scale=b)
        x2 = np.arange(-1, 1.1, 0.1)
        w = rv.pdf(x2)
        if p == cp:
            w = cp.asarray(w)
        rm = cv(m, w, axis=axis)
        return ScanData(rm)

    def clip(self, start, end):
        return ScanData(self.m[:, start:end])

    def print_contiguous(self):
        print(self.m.flags['C_CONTIGUOUS'], self.m.flags['F_CONTIGUOUS'])
        return self

    def rgb(self):
        return ScanData((self.norm().m * 255).astype(np.uint8)[:, :, None].repeat(3, axis=2))

    def pseudo_color(self):
        m = self.m * 0.7
        p = self.p
        h = m
        s = p.zeros_like(h) + 1
        v = p.zeros_like(h) + 1
        hsv = p.stack((h, s, v), axis=2)
        if p == cp:
            rgb = matplotlib.colors.hsv_to_rgb(hsv.get())
        else:
            rgb = matplotlib.colors.hsv_to_rgb(hsv)
        return ScanData((rgb * 255).astype(np.uint8))

    def pseudo_color_rgb(self):
        m = (1 - self.m) * 0.7
        p = self.p
        h = m
        s = p.zeros_like(h) + 1
        v = p.zeros_like(h) + 1
        hsv = p.stack((h, s, v), axis=2)
        if p == cp:
            rgb = matplotlib.colors.hsv_to_rgb(hsv.get())
        else:
            rgb = matplotlib.colors.hsv_to_rgb(hsv)
        return ScanData((rgb * 255).astype(np.uint8))

    def grey_color(self):
        m = self.norm().m
        return ScanData((m * 255).astype(np.uint8))

    def time_gain_compensation_linear(self, scale: float, start: int = 0):
        h = self.m.shape[-1]
        addend = self.p.zeros((1, h), dtype=np.int64)
        addend[:, start:] = self.p.arange(h - start) * scale
        self.m += addend
        return self

    def time_gain_compensation_linear_float(self, scale: float, start: int = 0):
        h = self.m.shape[-1]
        addend = self.p.zeros((1, h), dtype=self.p.float32)
        addend[:, start:] = (self.p.arange(h - start) * scale) + 1
        print(addend)
        return ScanData(self.m * addend)

    def time_gain_compensation_linear_max(self, scale: float, mmax: int | None = None, start: int = 0):
        if scale == 0:
            return self
        if mmax is None:
            mmax = self.m.max()
        h = self.m.shape[-1]
        self.m = self.m.astype(np.float64)
        mmax_arr = self.p.zeros(h) + mmax
        mmax_arr[start:] -= self.p.arange(h - start) * scale
        for i in range(h):
            # a[1, a[1, :] > 99] = 99
            self.m[self.m[:, i] > mmax_arr[i], i] = mmax_arr[i]
            self.m[:, i] *= (mmax / mmax_arr[i])
        self.m[self.m > mmax] = mmax
        self.m = self.m.astype(np.int64)
        return self

    def hib(self):
        if self.p == cp:
            hilbert = cupyx.scipy.signal.hilbert
        else:
            hilbert = scipy.signal.hilbert
        return ScanData(hilbert(self.m))

    def abs(self):
        return ScanData(self.p.abs(self.m))

    def log(self):
        return ScanData(self.p.log(self.m))

    def show(self, cmap='grey', w=40, h=20, no_axis=False, hsv=False):
        from matplotlib import pyplot as plt
        # fig = plt.figure(figsize=(w, h))
        fig, ax = plt.subplots(figsize=(w, h))
        if no_axis:
            ax.get_xaxis().set_ticks([])
            ax.get_yaxis().set_ticks([])
        image = self.m if self.p == np else self.m.get()
        if hsv:
            image = self.norm().pseudo_color_rgb().m
        plt.imshow(image, cmap=cmap)
        return self

    def plot_center_line(self):
        from matplotlib import pyplot as plt
        plt.figure(figsize=(40, 5))
        m = self.m if self.p == np else self.m.get()
        plt.plot(m[128, :])
        return self

    def windows(self, fn, center_f, width_f, sps_f):
        fft = scipy.fft.fft if self.p == np else cupyx.scipy.fft.fft
        ifft = scipy.fft.ifft if self.p == np else cupyx.scipy.fft.ifft

        seq_len = self.m.shape[-1]
        per_f = sps_f / seq_len

        center = int(center_f / per_f)
        width = int(width_f / 2 / per_f)

        start = center - width
        end = center + width

        fft_seq = fft(self.m)
        window = fn(end - start)
        filter = np.zeros(seq_len) + window[0]
        filter[start:end] = window
        filter[seq_len - end:seq_len - start] = window
        fft_seq *= filter
        return ScanData(self.p.real(ifft(fft_seq)))

    def windows_plot(self):
        pass

    def dctn(self, xmin, xmax, ymin, ymax):
        dctn = scipy.fft.dctn if self.p == np else cupyx.scipy.fft.dctn
        idctn = scipy.fft.idctn if self.p == np else cupyx.scipy.fft.idctn

        dctm = dctn(self.m)
        dctm[xmin:xmax, ymin:ymax] = 0
        return ScanData(idctn(dctm))

    def resize_XY(self, xdyd):
        print(self.m.shape[0])
        if self.p != np:
            raise NotImplementedError()
        return ScanData(cv2.resize(self.m, (int(self.m.shape[1] * xdyd), self.m.shape[0])))


def test1():
    s = ScanData(np.array([1, 1, 2])).norm().call(np.add, 1)
    print(s.m)


if __name__ == '__main__':
    test1()
