import numpy as np
from numpy import fft


class MonogenicFilter:
    def __init__(self, ysize, xsize, wl, f_ty="lg", param=None):
        ygrid, xgrid = self.freq_grid2(ysize, xsize)

        w = np.sqrt(ygrid ** 2 + xgrid ** 2)
        w[0, 0] = 1  # Trick to avoid division by zero at DC component
        num_filt = len(wl)

        if param is None:
            sigma_onf = 0.5
            ratio = 0.98
        else:
            assert 0.0 <= param <= 1.0, "Parameter must be in [0, 1]"
            sigma_onf = param
            ratio = param
        bp_filt = np.zeros((ysize, xsize, num_filt))
        for flt in range(num_filt):
            fo = 1.0 / wl[flt]
            w0 = fo
            if f_ty == "lg":
                a = -(np.log(w / w0) ** 2)
                b = 2 * np.log(sigma_onf) ** 2
                bp_filt[:, :, flt] = np.exp(a / b)
            elif f_ty == "gabor":
                raise NotImplementedError()
            elif f_ty == "gd":
                raise NotImplementedError()
            elif f_ty == "cau":
                raise NotImplementedError()
            elif f_ty == "dop":
                raise NotImplementedError()
            else:
                raise Exception(
                    'Filter must be one of ["lg", "gabor", "gd", "cau", "dop"]'
                )

            bp_filt[0, 0, flt] = 0

            if ysize % 2 == 0:
                bp_filt[(ysize // 2), :, flt] = 0

            if xsize % 2 == 0:
                bp_filt[:, (xsize // 2), flt] = 0

        sfilt = np.sum(bp_filt, 2)
        self.bp_filt = bp_filt / np.max(sfilt)
        self.num_filt = num_filt
        self.Reisz_filt = (1j * ygrid - xgrid) / w
        self.diff_filt = 2j * np.pi * ygrid - 2 * np.pi * xgrid
        self.wl = wl
        self.ty = f_ty
        self.sigma_onf = sigma_onf
        self.ratio = ratio

    @staticmethod
    def freq_grid2(ysize, xsize=None):
        xsize = xsize if xsize is not None else ysize

        ymid = ysize // 2
        xmid = xsize // 2

        ymax = ymid - (1 - (ysize % 2))
        xmax = xmid - (1 - (xsize % 2))

        ygrid, xgrid = np.meshgrid(
            np.arange(-ymid, ymax + 1), np.arange(-xmid, xmax + 1), indexing="ij"
        )  # equiv ndgrid
        return (fft.ifftshift(ygrid) / ysize, fft.ifftshift(xgrid) / xsize)

    def monogenic_signal(self, img):
        """img: 2d numpy.ndarray"""
        # might not work
        f = fft.fft2(img)
        f = f.reshape((*f.shape, 1))

        ffilt = f * self.bp_filt
        fm1 = np.real(fft.ifft2(ffilt, axes=(0, 1)))

        rf = self.Reisz_filt.reshape(f.shape)
        fmodd = fft.ifft2(ffilt * rf, axes=(0, 1))
        fm2 = np.real(fmodd)
        fm3 = np.imag(fmodd)
        return fm1, fm2, fm3
