from dataclasses import dataclass, field
from typing import Callable

import click
import labmodule as lab
import numpy as np
from customfunc.app import PlotApp
from scipy.optimize import curve_fit
from scipy.interpolate import interp1d

LN2 = np.log(2)


@click.group()
def main():
    pass


@dataclass
class FitResult:
    curve: np.ndarray = field(repr=False)
    covariance: np.ndarray
    amp: float
    width: float
    x0: float = 0
    params: tuple[float, ...] = field(default_factory=tuple)


@dataclass
class VoigtFitResult:
    name: str
    curve: np.ndarray = field(repr=False)
    covariance: np.ndarray
    amp: float
    gaussian_width: float
    lorentzian_width: float

    @property
    def descr(self) -> str:
        return (
            f"{self.name} - Gaussian width : {self.gaussian_width:.2f}, "
            f"Lorentzian width : {self.lorentzian_width:.2f}"
        )


def setup_axes(n: int, lims=(-10, 10)) -> tuple[np.ndarray, float, np.ndarray, float]:
    """Creates a time-like and a frequency-like array and returns their respective spacing"""
    x, dx = np.linspace(*lims, n, retstep=True)
    f = np.fft.fftfreq(n, dx)
    return x, dx, f, f[1] - f[0]


def freq_axis(x: np.ndarray) -> tuple[np.ndarray, np.ndarray, float]:
    """
    given an array with equally spaced values, returns a corresponding frequency array,
    an array of indices to automatically sort ffts and the x spacing
    """
    dx = np.diff(x).mean()
    f = np.fft.fftfreq(len(x), dx)
    ind = np.argsort(f)
    return f[ind], ind, dx


def lorentzian(x: np.ndarray, amp: float, width: float, x0: float) -> np.ndarray:
    """Lorentzian function of max `amp`, FWHM `width` and position `x0`"""
    gamma = (width * 0.5) ** 2
    return amp * gamma / ((x - x0) ** 2 + gamma)


def unit_lorentzian(x: np.ndarray, width: float, x0: float) -> np.ndarray:
    """Normalized Lorentzian function of FWHM `width` and position `x0`"""
    return lorentzian(x, 1 / (np.pi * width * 0.5), width, x0)


def peak_func(x: np.ndarray, amp: float, width: float, x0: float) -> np.ndarray:
    """2-sided decaying exp of max `amp`, FWHM `width` and position `x0`"""
    a = 2 * LN2 / width
    return amp * np.exp(-np.abs(x - x0) * a)


def fft_lorentzian(f: np.ndarray, dx: float, amp: float, width: float) -> np.ndarray:
    """
    Same as `peak_func` but parametrized as the fourier transform
    of a Lorentzian function of max `amp` and FWHM `width`
    """
    gamma_sqrt = width / 2
    return amp * np.pi * gamma_sqrt * np.exp(-2 * np.pi * gamma_sqrt * np.abs(f)) / dx


def gaussian(x: np.ndarray, amp: float, width: float, x0: float) -> np.ndarray:
    """Gaussian function of max `amp`, FWHM `width` and position `x0`"""
    sig2 = width ** 2 / (8 * LN2)
    return amp * np.exp(-((x - x0) ** 2) / (2 * sig2))


def elevated_gaussian(x: np.ndarray, amp: float, width: float, x0: float, y0: float) -> np.ndarray:
    """Gaussian function of max `amp`, FWHM `width`, position `x0` and baseline `y0`"""
    sig2 = width ** 2 / (8 * LN2)
    return amp * np.exp(-((x - x0) ** 2) / (2 * sig2)) + y0


def unit_gaussian(x: np.ndarray, width: float, x0: float) -> np.ndarray:
    """Normalized Gaussian function of FWHM `width` and position `x0`"""
    return gaussian(x, np.sqrt(4 * LN2 / np.pi) / width, width, x0)


def fft_gaussian(f: np.ndarray, dx: float, amp: float, width: float) -> np.ndarray:
    """
    Same as `gaussian` but parametrized as the fourier transform
    of another Gaussian function of max `amp` and FWHM `width`
    """  # exp_param =
    freq_amp = amp * width / 2 * np.sqrt(np.pi / LN2) / dx
    freq_width = 4 * LN2 / (width * np.pi)
    return gaussian(f, freq_amp, freq_width, 0)


def voigt(x: np.ndarray, amp: float, gaussian_width: float, lorentzian_width: float) -> np.ndarray:
    """Voigt profile of max 1 with specified Gaussian and Lorentzian FWHM"""
    conv = convolve(unit_gaussian(x, gaussian_width, 0), unit_lorentzian(x, lorentzian_width, 0))
    if conv.max() > 0:
        return conv * amp / conv.max()
    return conv


def log_voigt(
    x: np.ndarray, amp: float, gaussian_width: float, lorentzian_width: float
) -> np.ndarray:
    """log of `voigt`"""
    return np.log(voigt(x, amp, gaussian_width, lorentzian_width))


def fft_voigt(
    f: np.ndarray, dx: float, amp: float, gaussian_width: float, lorentzian_width: float
) -> np.ndarray:
    """
    returns the product of a `peak_func` and a `gaussian`
    parameterized as the Fourier transform of a Voigt profile
    given by `gaussian_width` and `lorentzian_width`
    """
    profile = fft_gaussian(f, dx, 1, gaussian_width) * fft_lorentzian(f, dx, 1, lorentzian_width)
    fac = amp / np.abs(np.fft.ifft(np.fft.fftshift(profile))).max()
    return profile * fac


def single_func_fit(
    x: np.ndarray,
    y: np.ndarray,
    fct: Callable[[np.ndarray, float, float, float], np.ndarray],
    force_center=False,
    use_weights=False,
    extra_params: int = 0,
) -> FitResult:
    """fits a single bell-like function (Gaussian or Lorentian, not Voigt)

    Parameters
    ----------
    x : np.ndarray
        x input
    y : np.ndarray
        corresponding y data
    fct : Callable[[np.ndarray, float, float, float], np.ndarray]
        function to fit
    force_center : bool, optional
        assume the bell is centered at x=0, by default False
    use_weights : bool, optional
        give more importance to higher y values, by default False

    Returns
    -------
    FitResult
        resuts of the fit
    """
    if force_center:

        def _fct(_x: np.ndarray, amp: float, width: float):
            return fct(_x, amp, width, 0)

        fct = _fct
    am = y.argmax()
    params, cov = curve_fit(
        fct, x, y, sigma=1 / y if use_weights else None, p0=[y[am], 1, x[am]] + [0] * extra_params
    )
    curve = fct(x, *params)
    return FitResult(curve, cov, params[0], params[1], params[2], tuple(params))


def linear_direct_voigt_fit(x: np.ndarray, y: np.ndarray, use_weights=False) -> VoigtFitResult:
    params, cov = curve_fit(voigt, x, y, sigma=1 / y if use_weights else None)
    curve = voigt(x, *params)
    return VoigtFitResult("Linear direct fit", curve, cov, *params)


def log_direct_voigt_fit(x: np.ndarray, y: np.ndarray) -> VoigtFitResult:
    ind = y > 0
    y_fit = np.log(y[ind])
    x_fit = x[ind]
    params, cov = curve_fit(log_voigt, x_fit, y_fit, bounds=[1e-12, np.inf])
    curve = voigt(x, *params)
    return VoigtFitResult("Log direct fit", curve, cov, *params)


def linear_fft_voigt_fit(f: np.ndarray, A: np.ndarray, x: np.ndarray, dx: float) -> VoigtFitResult:
    params, cov = curve_fit(
        lambda _f, _a, _g, _l: fft_voigt(_f, dx, _a, _g, _l), f, A, bounds=[1e-12, np.inf]
    )
    curve = voigt(x, *params)
    return VoigtFitResult("Linear FFT fit", curve, cov, *params)


def convolve(gau: np.ndarray, lor: np.ndarray) -> np.ndarray:
    return np.convolve(gau, lor, mode="same")


def add_noise(arr: np.ndarray, abs_fac: float, rel_fac: float):
    return (
        arr + abs_fac * np.random.rand(len(arr)) + rel_fac * arr * (np.random.rand(len(arr)) - 0.5)
    )


@main.command(help="fit a lorentzian with a gaussian")
def demo1():
    x = np.linspace(-10, 10, 501)
    with PlotApp(
        FWHM=np.geomspace(0.1, 10),
        relative_noise=np.linspace(0, 2),
        absolute_noise=[0, *np.geomspace(0.001, 1)],
    ) as app:
        app.set_antialiasing(True)
        app.params["FWHM"].value = 3.7

        @app.update
        def draw(FWHM, relative_noise, absolute_noise):
            lo = lorentzian(x, 1, FWHM, 0)
            lo = add_noise(lo, absolute_noise, relative_noise * 0.04)
            fit = single_func_fit(x, lo, gaussian)
            app[0].set_line_data("Lorentzian", x, lo, width=2)
            app[0].set_line_data(
                "gaus fit", x, fit.curve, width=2, label=f"Gaussian fit : FWHM = {fit.width}"
            )


@main.command(name="default", help="illustrate the influence of noises and fit a voigt profile")
def demo2():
    with PlotApp(
        n=np.arange(256, 2046),
        lorentz_fwhm=np.geomspace(0.1, 1),
        gauss_fwhm=np.linspace(1, 10),
        relative_noise=np.linspace(0, 2),
        absolute_noise=[0, *np.geomspace(1e-6, 1, 201)],
    ) as app:
        app.set_antialiasing(True)
        ax = app["Main Plot"]

        app.params["gauss_fwhm"].value = 5
        app.params["n"].value = 1024
        app.params_layout.setSpacing(0)

        @app.update
        def draw(n, lorentz_fwhm, gauss_fwhm, absolute_noise, relative_noise):
            x, dx, f, df = setup_axes(n, (-20, 20))
            gaus = gaussian(x, 1, gauss_fwhm, 0)
            lore = lorentzian(x, 1, lorentz_fwhm, 0)
            ax.set_line_data("Gaussian", x, gaus / gaus.max(), width=2)
            ax.set_line_data("Lorentz", x, lore / lore.max(), width=2)

            conv = convolve(gaus, lore)
            conv /= conv.max()
            conv_noise = add_noise(conv, absolute_noise, relative_noise * 0.1)
            gaus_noise = add_noise(gaus, absolute_noise, relative_noise * 0.1)
            ax.set_line_data("Voigt", x, conv, width=2)
            ax.set_line_data("Noisy Voigt", x, conv_noise, width=2)
            ax.set_line_data("Noisy Gaussian", x, gaus_noise, width=2)

            lin_fit = linear_direct_voigt_fit(x, conv_noise)
            log_fit = log_direct_voigt_fit(x, conv_noise)
            transfo = np.fft.fft(np.fft.fftshift(conv_noise)).real
            fft_fit = linear_fft_voigt_fit(f, transfo / transfo.max(), x, dx)
            for fit in lin_fit, log_fit, fft_fit:
                display_fit(x, fit)

        def display_fit(x: np.ndarray, fit: VoigtFitResult):
            ax.set_line_data(fit.name, x, fit.curve)
            ax.set_line_data(
                fit.name + "gauss",
                x,
                gaussian(x, 1, fit.gaussian_width, 0),
                label=f"{fit.name} - Gaussian {fit.gaussian_width:.2f}",
                width=1.5,
            )
            ax.set_line_data(
                fit.name + "loren",
                x,
                lorentzian(x, 1, fit.lorentzian_width, 0),
                label=f"{fit.name} - Lorentzian {fit.lorentzian_width:.2f}",
                width=1.5,
            )


@main.command(help="Explore fft fitting")
def demo3():
    x, dx, f, df = setup_axes(1001, (-20, 20))
    with PlotApp(
        lorentz_fwhm=np.geomspace(0.1, 1),
        gauss_fwhm=np.linspace(1, 10),
        relative_noise=np.linspace(0, 2),
        absolute_noise=[0, *np.geomspace(1e-6, 1, 201)],
    ) as app:
        app.set_antialiasing(True)

        ax = app[0]
        ind = np.argsort(f)

        @app.update
        def draw(lorentz_fwhm, gauss_fwhm, absolute_noise, relative_noise):
            gaus_clean = gaussian(x, 1, gauss_fwhm, 0)
            lore_clean = lorentzian(x, 1, lorentz_fwhm, 0)
            voig_clean = convolve(gaus_clean, lore_clean)

            gaus = add_noise(gaus_clean, absolute_noise, relative_noise * 0.1)
            lore = add_noise(lore_clean, absolute_noise, relative_noise * 0.1)
            voig = add_noise(voig_clean, absolute_noise, relative_noise * 0.1)

            gaus_f = np.fft.fft(np.fft.fftshift(gaus)).real
            lore_f = np.fft.fft(np.fft.fftshift(lore)).real
            voig_f = np.fft.fft(np.fft.fftshift(voig)).real

            ax.set_line_data("Transformed Gaussian", f[ind], gaus_f[ind])
            ax.set_line_data("Transformed Lorentz", f[ind], lore_f[ind])
            ax.set_line_data("Transformed Voigt", f[ind], voig_f[ind])

            gaus_ff = fft_gaussian(f, dx, 1, gauss_fwhm)
            lore_ff = fft_lorentzian(f, dx, 1, lorentz_fwhm)
            voig_ff = fft_voigt(f, dx, gauss_fwhm, lorentz_fwhm) * voig_clean.sum()

            ax.set_line_data("Expected Gaussian", f[ind], gaus_ff[ind], width=1)
            ax.set_line_data("Expected Lorentz", f[ind], lore_ff[ind], width=1)
            ax.set_line_data("Expected Voigt", f[ind], voig_ff[ind], width=1)


@click.argument("file", type=click.Path(True, dir_okay=False))
@click.option("--all", "show_all", default=False, is_flag=True)
@main.command(name="fit")
def fit_measurement(file, show_all):
    raw_wl, intens = lab.osa.load_osa_spectrum(file)
    raw_wl *= 1e9
    gauss_fit = single_func_fit(raw_wl, intens, elevated_gaussian, extra_params=1)
    wl0 = gauss_fit.x0
    baseline = gauss_fit.params[-1]
    intens -= baseline
    gauss_fit.curve -= baseline
    intens[intens < 0] = 0

    lim = np.max(np.abs(raw_wl - wl0))
    new_wl = np.linspace(-lim, lim, 2048)
    intens = interp1d(raw_wl - wl0, intens, bounds_error=False, fill_value=0)(new_wl)
    wl = new_wl

    with PlotApp(n=np.arange(len(wl) // 2)) as app:
        app.set_antialiasing(True)

        ax1 = app["Measurement"]
        ax2 = app["Fourier domain"]
        ax1.set_line_data("Measured spectrum", wl, intens, width=2, color="red")
        ax1.set_line_data(
            "Initial Guassian fit", raw_wl - wl0, gauss_fit.curve, width=2, color="blue"
        )
        # ax1.plot_widget.plotItem.setLogMode(y=True)

        @app.update
        def draw(n):
            s = slice(n, -n) if n else slice(None)

            wl_fit = wl[s]
            intens_fit = intens[s] - intens[s].min()
            f, ind, dwl = freq_axis(wl_fit)

            ax1.set_line_data("Fitted spectrum", wl_fit, intens_fit)

            trans = np.abs(np.fft.fft(intens_fit))[ind]
            ax2.set_line_data("Transformed", f, trans)
            ax2.lines["Transformed"].setZValue(10)

            fft_fit = linear_fft_voigt_fit(f, trans, wl_fit, dwl)
            lin_fit = linear_direct_voigt_fit(wl_fit, intens_fit)
            log_fit = log_direct_voigt_fit(wl_fit, intens_fit)
            for fit in fft_fit, lin_fit, log_fit:
                display_fit(wl_fit, dwl, f, fit)

        def display_fit(x: np.ndarray, dx: float, f: np.ndarray, fit: VoigtFitResult):
            ax2.set_line_data(
                fit.name,
                f,
                fft_voigt(f, dx, fit.amp, fit.gaussian_width, fit.lorentzian_width),
            )
            ax1.set_line_data(fit.name, x, fit.curve)
            if show_all:
                ax1.set_line_data(
                    fit.name + "gauss",
                    x,
                    gaussian(x, 1, fit.gaussian_width, 0),
                    label=f"{fit.name} - Gaussian {fit.gaussian_width:.3f}",
                    width=1.5,
                )
                ax1.set_line_data(
                    fit.name + "lorentz",
                    x,
                    lorentzian(x, 1, fit.lorentzian_width, 0),
                    label=f"{fit.name} - Lorentzian {fit.lorentzian_width:.3f}",
                    width=1.5,
                )


@main.command(name="test", help="test scaling of voigt profile")
def test_amp():
    x, dx, f, df = setup_axes(1001, (-20, 20))
    with PlotApp(
        lorentz_fwhm=np.geomspace(0.1, 1), gauss_fwhm=np.linspace(1, 10), amp=np.linspace(1, 7)
    ) as app:
        app.set_antialiasing(True)

        ax1 = app["Main"]
        ax2 = app["Fourrier"]
        ind = np.argsort(f)

        @app.update
        def draw(lorentz_fwhm, gauss_fwhm, amp):
            vo = voigt(x, amp, gauss_fwhm, lorentz_fwhm)
            trans = np.fft.fft(np.fft.fftshift(vo))
            expected = fft_voigt(f, dx, amp, gauss_fwhm, lorentz_fwhm)
            back_transformed = np.fft.fftshift(np.fft.ifft(expected).real)

            ax1.set_line_data("original", x, vo)
            ax1.set_line_data("back tranformed", x, back_transformed)
            ax2.set_line_data("expected", f[ind], expected[ind])
            ax2.set_line_data("original transformed", f[ind], trans[ind].real)


if __name__ == "__main__":
    main()