import numpy as np
from scipy import linalg, signal

from meep import CustomSource


class FilteredSource(CustomSource):
    def __init__(
        self,
        center_frequency,
        frequencies,
        frequency_response,
        dt,
        time_src=None,
    ):
        dt = dt / 2  # divide by two to compensate for staggered E,H time interval
        self.dt = dt
        self.frequencies = frequencies
        self.center_frequencies = frequencies

        # For now, the basis functions cannot overlap much in the frequency domain. Otherwise, the
        # resulting nodes are wildly large and induce numerical precision errors. We can always
        # produce a safe simulation by forcing the length of each basis function to meet the minimum
        # frequency requirements. This method still minimizes storage requirements.
        self.T = np.max(np.abs(1 / np.diff(frequencies)))
        self.N = np.rint(self.T / self.dt)
        self.t = np.arange(0, dt * (self.N), dt)
        self.n = np.arange(self.N)
        f = self.func()

        # frequency bandwidth of the Nuttall window function
        fwidth = self.nuttall_bandwidth()

        self.bf = [
            lambda t, i=i: 0
            if t > self.T
            else (
                self.nuttall(t, self.center_frequencies)
                / (self.dt / np.sqrt(2 * np.pi))
            )[i]
            for i in range(len(self.center_frequencies))
        ]
        self.time_src_bf = [
            CustomSource(
                src_func=bfi,
                center_frequency=center_frequency,
                is_integrated=False,
                end_time=self.T,
                fwidth=fwidth,
            )
            for bfi in self.bf
        ]

        if time_src:
            # get the cutoff of the input signal
            signal_t = np.array(
                [time_src.swigobj.current(ti, dt) for ti in self.t]
            )  # time domain signal
            signal_dtft = self.dtft(signal_t, self.frequencies)
        else:
            signal_dtft = 1

        # multiply sampled dft of input signal with filter transfer function
        H = signal_dtft * frequency_response

        # estimate the impulse response using a sinc function RBN
        self.nodes, self.err = self.estimate_impulse_response(H)

        # initialize super
        super().__init__(
            src_func=f,
            center_frequency=center_frequency,
            is_integrated=False,
            end_time=self.T,
            fwidth=fwidth,
        )

    def cos_window_td(self, a, t, f0):
        cos_sum = np.sum(
            [
                (-1) ** k * a[k] * np.cos(2 * np.pi * t * k / self.T)
                for k in range(len(a))
            ],
            axis=0,
        )
        return np.exp(-1j * 2 * np.pi * f0 * t) * (cos_sum)

    def cos_window_fd(self, a, f, f0):
        df = 1 / (self.N * self.dt)
        cos_sum = a[0] * self.sinc(f, f0)
        for k in range(1, len(a)):
            cos_sum += (-1) ** k * a[k] / 2 * self.sinc(f, f0 - k * df) + (-1) ** k * a[
                k
            ] / 2 * self.sinc(f, f0 + k * df)
        return cos_sum

    def sinc(self, f, f0):
        num = np.where(
            f == f0,
            self.N + 1,
            (1 - np.exp(1j * (self.N + 1) * (2 * np.pi) * (f - f0) * self.dt)),
        )
        den = np.where(f == f0, 1, (1 - np.exp(1j * (2 * np.pi) * (f - f0) * self.dt)))
        return num / den

    def rect(self, t, f0):
        n = np.rint((t) / self.dt)
        return np.where(
            n.any() < 0.0 or n.any() > self.N, 0, np.exp(-1j * 2 * np.pi * f0 * t)
        )

    def hann(self, t, f0):
        a = [0.5, 0.5]
        return self.cos_window_td(a, t, f0)

    def hann_dtft(self, f, f0):
        a = [0.5, 0.5]
        return self.cos_window_fd(a, f, f0)

    def nuttall(self, t, f0):
        a = [0.355768, 0.4873960, 0.144232, 0.012604]
        return self.cos_window_td(a, t, f0)

    def nuttall_dtft(self, f, f0):
        a = [0.355768, 0.4873960, 0.144232, 0.012604]
        return self.cos_window_fd(a, f, f0)

    ## compute the bandwidth of the DTFT of the Nuttall window function
    ## (magnitude) assuming it has decayed from its peak value by some
    ## tolerance by fitting it to an asymptotic power law of the form
    ## C / f^3 where C is a constant and f is the frequency
    def nuttall_bandwidth(self):
        tol = 1e-7
        fwidth = 1 / (self.N * self.dt)
        frq_inf = 10000 * fwidth
        na_dtft = self.nuttall_dtft(frq_inf, 0)
        coeff = frq_inf**3 * np.abs(na_dtft)
        na_dtft_max = self.nuttall_dtft(0, 0)
        bw = 2 * np.power(coeff / (tol * na_dtft_max), 1 / 3)
        return bw.real

    def dtft(self, y, f):
        return (
            np.matmul(
                np.exp(1j * 2 * np.pi * f[:, np.newaxis] * np.arange(y.size) * self.dt),
                y,
            )
            * self.dt
            / np.sqrt(2 * np.pi)
        )

    def __call__(self, t):
        if t > self.T:
            return 0
        vec = self.nuttall(t, self.center_frequencies) / (
            self.dt / np.sqrt(2 * np.pi)
        )  # compensate for meep dtft
        return np.inner(vec, self.nodes)

    def func(self):
        def _f(t):
            return self(t)

        return _f

    def estimate_impulse_response(self, H):
        # Use vandermonde matrix to calculate weights of each gaussian. Each window is centered at each frequency point.
        # TODO: come up with a more sophisticated way to choose temporal window size and basis locations
        # that will minimize l2 estimation error and the node weights (since matrix is ill-conditioned)
        vandermonde = self.nuttall_dtft(
            self.frequencies[:, np.newaxis], self.center_frequencies[np.newaxis, :]
        )
        nodes = np.matmul(linalg.pinv(vandermonde), H.T)
        H_hat = np.matmul(vandermonde, nodes)
        l2_err = np.sum(np.abs(H - H_hat.T) ** 2 / np.abs(H) ** 2)
        return nodes, l2_err
