"""
Wavelet transform
"""

from __future__ import annotations
from ..backend import wave_signif, wavelet
import xarray as xr
import pandas as pd
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import statsmodels.api as sm
import scipy

__all__ = [
    "calc_timeseries_wavelet_transform",
    "draw_global_wavelet_spectrum",
    "draw_wavelet_transform",
]


def calc_timeseries_wavelet_transform(
    timeseries_data: xr.DataArray,
    dt: float,
    time_dim: str = "time",
    pad: float = 1,
    dj: float = 0.25,
    j1_n_div_dj: int = 7,
    s0=None,
    lag1: float = None,
    mother: str = "morlet",
    mother_param=None,
    sigtest_wavelet: str = "regular chi-square test",
    sigtest_global: str = "time-average test",
    significance_level: float = 0.95,
) -> xr.Dataset:
    """
    Wavelet transform parameters calculation.

    Parameters
    ----------
    timeseries_data: :py:class:`xarray.DataArray<xarray.DataArray>`.
        Timeseries data.
    dt: :py:class:`float<float>`.
        Amount of time between each timeseries data value, i.e. the sampling time.
    time_dim: :py:class:`str <str>`.
        The time coordinate dimension name.
    pad: :py:class:`float<float>`, default: `1`.
        if set to 1 (default is 0), pad time series with zeroes to get N up to the next higher power of 2.
        This prevents wraparound from the end of the time series to the beginning, and also
        speeds up the FFT's used to do the wavelet transform. This will not eliminate all edge effects (see COI below).
    dj: :py:class:`float<float>`, default: `0.25`.
        The spacing between discrete scales. A smaller `dj` will give better scale resolution, but be slower to plot.
    j1_n_div_dj: :py:class:`int<int>`, default: `7`.
        Do `j1_n_div_dj` powers-of-two with `dj` sub-octaves each.

        .. math::
            j1 = \\mathrm{j1\\underline{ }n\\underline{ }div\\underline{ }dj} / dj

        This can adjust the size of the period.

    s0: :py:class:`float<float>`, default: :math:`2 \\times \\mathrm{d}t`.
        The smallest scale of the wavelet.
    lag1: :py:class:`float<float>`, default: `None`.
        Lag-1 autocorrelation for red noise background. The value is generated by :py:func:`statsmodels.api.tsa.acf <statsmodels.api.tsa.acf>`.
    mother: :py:class:`str<str>`, {'morlet', 'paul', 'dog'}, default: 'morlet'.
        The mother wavelet function.

        .. tab:: Morlet

            - Name: Morlet (:math:`\\omega_0` = frequency)
            - :math:`\\psi_0(\\eta)`: :math:`\\pi^{-1/4} e^{i \\omega_{0} \\eta} e^{-\\eta^2/2}`.

        .. tab:: Paul

            - Name: Paul (:math:`m` = order)
            - :math:`\\psi_0(\\eta)`: :math:`\\frac{2^m i^m m!}{\\sqrt{\\pi (2m)!}} (1-i \\eta)^{-(m+1)}`.

        .. tab:: DOG

            - Name: DOG (:math:`m` = derivative)
            - :math:`\\psi_0(\\eta)`: :math:`\\frac{(-1)^{m+1}}{\\sqrt{\\Gamma (m+\\frac{1}{2})}} \\frac{d^m}{d \\eta^m} (e^{-\\eta^2 /2})`.


    mother_param: :py:class:`float<float>`.
        The mother wavelet parameter.
            - For 'morlet' this is :math:`k_0` (wavenumber), default is 6.
            - For 'paul' this is :math:`m` (order), default is 4.
            - For 'dog' this is :math:`m` (m-th derivative), default is 2.

    sigtest_wavelet: {'regular chi-square test', 'time-average test', 'scale-average test'}, default: `'regular chi-square test'`.
        The type of significance test.

        1. Regular chi-square test
        i.e. Eqn (18) from Torrence & Compo.

        .. math::

            \\frac{\\left|W_n(s)\\right|^2}{\\sigma^2}\\Longrightarrow\\frac{1}{2} P_k\\chi_2^2

        2. The "time-average" test, i.e. Eqn (23).

        .. math::

            \\nu=2\\sqrt{1+\\left(\\frac{n_a\\delta t}{\\gamma s}\\right)^2}

        In this case, DOF should be set to NA, the number
        of local wavelet spectra that were averaged together.
        For the Global Wavelet Spectrum, this would be NA=N,
        where N is the number of points in your time series.

        3. The "scale-average" test, i.e. Eqns (25)-(28).

        .. math::

            \\overline{P}=S_{\\mathrm{avg}}\\sum_{j=j_1}^{j_2}\\frac{P_j}{S_j}, \\ \\mathrm{where} \\ S_{\\mathrm{avg}}=\\left(\\sum_{j=j_1}^{j_2}\\frac1{s_j}\\right)^{-1}, \\frac{C_\\delta S_\\mathrm{avg}}{\\delta j\\delta t\\sigma^2}\\overline{W}_n^2\\Rightarrow\\overline{P}\\frac{\\chi_\\nu^2}\\nu, \\nu=\\frac{2n_aS_{\\mathrm{avg}}}{S_{\\mathrm{mid}}}\\sqrt{1+\\left(\\frac{n_a\\delta j}{\\delta j_0}\\right)^2}.

        In this case, DOF should be set to a
        two-element vector [S1,S2], which gives the scale
        range that was averaged together.
        e.g. if one scale-averaged scales between 2 and 8,
        then DOF=[2,8].

    sigtest_global: {'regular chi-square test', 'time-average test', 'scale-average test'}, default: `'time-average test'`.
        See also the description of `sigtest_wavelet`.

    significance_level: :py:class:`float<float>`, default: `0.95`.
        Significance level to use.

    Returns
    -------
    Timeseries wavelet transform result (:py:class:`xarray.Dataset<xarray.Dataset>`).

    .. seealso::
        - https://github.com/ct6502/wavelets
        - https://github.com/regeirk/pycwt, https://pycwt.readthedocs.io/en/latest/index.html
        - http://nicolasfauchereau.github.io/climatecode/posts/wavelet-analysis-in-python/
        - https://blog.csdn.net/weixin_43304836/article/details/119752767

    Reference
    --------------
    - Torrence, C., & Compo, G. P. (1998). A Practical Guide to Wavelet Analysis. Bulletin of the American Meteorological Society, 79(1), 61-78. https://doi.org/10.1175/1520-0477(1998)079<0061:APGTWA>2.0.CO;2
    - Torrence, C., & Webster, P. J. (1999). Interdecadal Changes in the ENSO–Monsoon System. Journal of Climate, 12(8), 2679-2690. https://doi.org/10.1175/1520-0442(1999)012<2679:ICITEM>2.0.CO;2
    - Grinsted, A., Moore, J. C., and Jevrejeva, S.: Application of the cross wavelet transform and wavelet coherence to geophysical time series, Nonlin. Processes Geophys., 11, 561–566, https://doi.org/10.5194/npg-11-561-2004, 2004.

    .. minigallery::
        :add-heading: Example(s) related to the function

        ./dynamic_docs/plot_wavelet.py
    """
    # variance
    timeseries_data_numpy = timeseries_data.data
    variance = float(timeseries_data.var().data)
    n = timeseries_data[time_dim].shape[0]

    if s0 is None:
        # this says start at a scale of 6 months
        s0 = 2 * dt

    if mother_param is None:
        mother_param_value = -1
    else:
        mother_param_value = mother_param

    if lag1 is None:
        lag1 = sm.tsa.acf(timeseries_data_numpy)[1]
        print(
            "Allen and Smith AR(1) model estimate of the lag-one autocorrelation: lag1 = ",
            lag1,
        )

    # this says do `j1_n_div_dj` powers-of-two with dj sub-octaves each
    j1 = j1_n_div_dj / dj

    match mother:
        case "morlet":
            mother_value = "MORLET"
        case "paul":
            mother_value = "PAUL"
        case "dog":
            mother_value = "DOG"

    match sigtest_wavelet:
        case "regular chi-square test":
            sigtest_wavelet_value = 0
        case "time-average test":
            sigtest_wavelet_value = 1
        case "scale-average test":
            sigtest_wavelet_value = 2

    match sigtest_global:
        case "regular chi-square test":
            sigtest_global_value = 0
        case "time-average test":
            sigtest_global_value = 1
        case "scale-average test":
            sigtest_global_value = 2

    # ------------------------------------------------
    # Wavelet transform
    wave, period, scale, coi = wavelet(
        timeseries_data_numpy, dt, pad, dj, s0, j1, mother_value, mother_param_value
    )
    power = (np.abs(wave)) ** 2  # Compute wavelet power spectrum
    global_ws = np.sum(power, axis=1) / n  # time-average over all times

    # ------------------------------------------------
    # Significance levels
    signif = wave_signif(
        ([variance]),
        dt=dt,
        sigtest=sigtest_wavelet_value,
        scale=scale,
        lag1=lag1,
        mother=mother_value,
        siglvl=significance_level,
        param=mother_param,
    )

    # expand signif --> (J+1)x(N) array
    sig95 = signif[:, np.newaxis].dot(np.ones(n)[np.newaxis, :])
    sig95 = power / sig95  # Where ratio > 1, power is significant

    # ------------------------------------------------
    # Global wavelet spectrum & significance levels
    dof = n - scale  # the -scale corrects for padding at edges
    global_signif = wave_signif(
        variance,
        dt=dt,
        scale=scale,
        sigtest=sigtest_global_value,
        lag1=lag1,
        dof=dof,
        mother=mother_value,
        siglvl=significance_level,
        param=mother_param,
    )

    # ------------------------------------------------
    # Create xarray DataArray
    power_xarray = xr.DataArray(
        power,
        dims=["period", time_dim],
        coords={"period": period, time_dim: timeseries_data[time_dim]},
        name="power",
    )

    sig95_xarray = xr.DataArray(
        sig95,
        dims=["period", time_dim],
        coords={"period": period, time_dim: timeseries_data[time_dim]},
        name="sig",
    )

    global_ws_xarray = xr.DataArray(
        global_ws, dims=["period"], coords={"period": period}, name="global_ws"
    )

    global_signif_xarray = xr.DataArray(
        global_signif, dims=["period"], coords={"period": period}, name="global_signif"
    )

    coi_xarray = xr.DataArray(
        coi, dims=[time_dim], coords={time_dim: timeseries_data[time_dim]}, name="coi"
    )

    coi_bottom_xarray = xr.DataArray(
        coi * 0 + period[-1],
        dims=[time_dim],
        coords={time_dim: timeseries_data[time_dim]},
        name="coi",
    )

    result = xr.Dataset(
        data_vars={
            "power": power_xarray,
            "sig": sig95_xarray,
            "global_ws": global_ws_xarray,
            "global_signif": global_signif_xarray,
            "coi": coi_xarray,
            "coi_bottom": coi_bottom_xarray,
        }
    )
    result.attrs["dt"] = dt
    result.attrs["scale"] = scale
    result.attrs["sigtest_global"] = sigtest_global
    result.attrs["sigtest_wavelet"] = sigtest_wavelet
    result.attrs["lag1"] = lag1
    result.attrs["dof"] = dof
    result.attrs["mother"] = mother
    result.attrs["significance_level"] = significance_level
    return result


def draw_global_wavelet_spectrum(
    timeseries_wavelet_transform_result: xr.Dataset,
    ax: matplotlib.axes.Axes = None,
    global_ws_kwargs: dict = {},
    global_signif_kwargs: dict = {"ls": "--"},
):
    """
    Draw global wavelet spectrum

    Parameters
    ----------
    timeseries_wavelet_transform_result: :py:class:`xarray.Dataset<xarray.Dataset>`.
        Timeseries wavelet transform result.
    ax: :py:class:`matplotlib.axes.Axes`
        The axes to which the boundary will be applied.
    **global_ws_kwargs, :py:class:`dict <dict>`, optional:
        Additional keyword arguments to :py:func:`xarray.DataArray.plot.line<xarray.DataArray.plot.line>` for ploting `global_ws`.
    **global_signif_kwargs, :py:class:`dict <dict>`, optional, default {'ls': '--'}:
        Additional keyword arguments to :py:func:`xarray.DataArray.plot.line<xarray.DataArray.plot.line>` for ploting `global_signif`.

    .. minigallery::
        :add-heading: Example(s) related to the function

        ./dynamic_docs/plot_wavelet.py
    """
    if ax == None:
        ax = plt.gca()

    period = timeseries_wavelet_transform_result["period"]

    timeseries_wavelet_transform_result.global_ws.plot(ax=ax, **global_ws_kwargs)
    timeseries_wavelet_transform_result.global_signif.plot(
        ax=ax, **global_signif_kwargs
    )

    ax.set_xscale("log", base=2, subs=None)
    ax.xaxis.set_major_formatter(ticker.ScalarFormatter())
    ax.set_xlim([np.min(period), np.max(period)])

    ax.set_ylim([0, 1.25 * timeseries_wavelet_transform_result.global_ws.max()])

    ax.set_xlabel("Period")
    ax.set_ylabel("Power")


def draw_wavelet_transform(
    timeseries_wavelet_transform_result: xr.Dataset,
    ax: matplotlib.axes.Axes = None,
    power_kwargs: dict = {
        "levels": [0, 0.5, 1, 2, 4, 999],
        "colors": ["white", "bisque", "orange", "orangered", "darkred"],
    },
    sig_kwargs: dict = {"levels": [-99, 1], "colors": "k"},
    coi_kwargs: dict = {"color": "k"},
    fill_between_kwargs: dict = {
        "facecolor": "none",
        "edgecolor": "#00000040",
        "hatch": "x",
    },
):
    """
    Draw wavelet transform

    Parameters
    ----------
    timeseries_wavelet_transform_result: :py:class:`xarray.Dataset<xarray.Dataset>`.
        Timeseries wavelet transform result.
    ax : :py:class:`matplotlib.axes.Axes`
        The axes to which the boundary will be applied.
    **power_kwargs, optional, :py:class:`dict <dict>`, default {'levels': [0, 0.5, 1, 2, 4, 999], 'colors': ['white', 'bisque', 'orange', 'orangered', 'darkred']}:
        Additional keyword arguments to :py:func:`xarray.DataArray.plot.contourf<xarray.DataArray.plot.contourf>` for ploting `power`.
    **sig_kwargs, optional, :py:class:`dict <dict>`, default {'levels': [-99, 1], 'colors': 'k'}:
        Additional keyword arguments to :py:func:`xarray.DataArray.plot.contourf<xarray.DataArray.plot.contourf>` for ploting `sig`.
    **coi_kwargs, optional, :py:class:`dict <dict>`, default {'color': 'k'}:
        Additional keyword arguments to :py:func:`xarray.DataArray.plot.contourf<xarray.DataArray.plot.contourf>` for ploting `coi`.
    **fill_between_kwargs, :py:class:`dict <dict>`, optional, default {'facecolor': 'none', 'edgecolor': '#00000040', 'hatch': 'x'}:
        Additional keyword arguments to :py:func:`matplotlib.pyplot.fill_between<matplotlib.pyplot.fill_between>`.

    .. minigallery::
        :add-heading: Example(s) related to the function

        ./dynamic_docs/plot_wavelet.py
    """
    if ax == None:
        ax = plt.gca()

    period = timeseries_wavelet_transform_result["period"]
    power_xarray = timeseries_wavelet_transform_result["power"]
    sig95_xarray = timeseries_wavelet_transform_result["sig"]
    coi_xarray = timeseries_wavelet_transform_result["coi"]
    coi_bottom_xarray = timeseries_wavelet_transform_result["coi_bottom"]

    power_xarray.plot.contourf(**power_kwargs)
    sig95_xarray.plot.contour(**sig_kwargs)
    coi_xarray.plot(**coi_kwargs)

    plt.fill_between(
        power_xarray.time, coi_bottom_xarray, coi_xarray, **fill_between_kwargs
    )

    ax.set_yscale("log", base=2, subs=None)
    ax.yaxis.set_major_formatter(ticker.ScalarFormatter())
    ax.set_ylim([np.min(period), np.max(period)])
    ax.ticklabel_format(axis="y", style="plain")

    ax.set_xlabel("Time")
    ax.set_ylabel("Period")
