from dataclasses import dataclass, field
from functools import partial
from typing import Any

import jax
import jax.numpy as jnp
import numpy as np

from .core.get_filter import get_filter
from .core.time_delay import (
    get_pixel_amplitude as _get_pixel_amplitude_td,
    get_td_amp as _get_td_amp,
    get_td_vec as _get_td_vec,
    set_td_filter as _set_td_filter,
    split_tf_map as _split_tf_map,
)
from .core.t2w import HAS_JAX as HAS_T2W_JAX, HAS_NUMBA as HAS_T2W_NUMBA, t2w_jax, t2w_numba
from .core.w2t import (
    HAS_JAX as HAS_W2T_JAX,
    HAS_NUMBA as HAS_W2T_NUMBA,
    w2t_jax,
    w2t_numba,
    w2tQ_jax,
    w2tQ_numba,
)
from .fourier_coeff import load_fourier_coeff
from gwpy.timeseries import TimeSeries
from wdm_wavelet.types import TimeFrequencyMap


@dataclass(eq=False)
class WDM:
    """
    Wilson-Daubechies-Meyer (WDM) wavelet transform.

    Implements the forward (time → TF map) and inverse (TF map → time) WDM
    transforms, plus time-delay amplitude queries useful for coincidence-based
    gravitational-wave searches.

    Parameters
    ----------
    M : int
        Number of frequency bands (Nyquist layer index).  The TF map has
        ``M + 1`` frequency layers: DC (0), interior (1…M-1), and Nyquist (M).
    K : int
        Meyer window transition-band sharpness.  Larger values give steeper
        roll-off and longer filters.
    beta_order : int
        Beta polynomial order used in the Meyer window construction.
    precision : int
        Number of decimal digits of reconstruction accuracy required when
        auto-computing the filter length.  Higher values produce longer filters.
    backend : str, optional
        Compute backend to use for this instance. Supported values are
        ``'jax'`` and ``'numba'``. The backend is locked at construction time.
    filter : array-like or None, optional
        Custom filter taps.  When ``None`` (default) the filter is computed
        automatically from *M*, *K*, *beta_order*, and *precision*.
    m_H : int or None, optional
        Number of filter taps to use.  Inferred from *filter* when ``None``.

    Notes
    -----
    For embedding transforms inside ``@jax.jit`` or ``@numba.njit`` contexts,
    use the pure-array functions in :mod:`wdm_wavelet.core.t2w` and
    :mod:`wdm_wavelet.core.w2t` directly rather than the class methods.

    Examples
    --------
    >>> wdm = WDM(M=32, K=64, beta_order=6, precision=10)
    >>> tf_map = wdm.t2w(signal, sample_rate=4096.0)
    >>> reconstructed = wdm.w2t(tf_map)
    """

    M: int
    K: int
    beta_order: int
    precision: int
    backend: str = "jax"
    filter: Any = field(default=None, repr=False)  # numpy / jax array of taps
    m_H: int | None = None
    # ------------------------------------------------------------------ #
    # Fields below are set in __post_init__ and excluded from repr/compare
    # ------------------------------------------------------------------ #
    td_filters: Any = field(default=None, init=False, repr=False, compare=False)
    _cos2: Any = field(init=False, repr=False, compare=False)
    _sincos: Any = field(init=False, repr=False, compare=False)
    _cos2_size: Any = field(init=False, repr=False, compare=False)
    _sincos_size: Any = field(init=False, repr=False, compare=False)

    def __post_init__(self):
        self.backend = self._validate_backend(self.backend)
        if self.backend == "jax" and not (HAS_T2W_JAX and HAS_W2T_JAX):
            raise ImportError("JAX backend requires jax to be installed.")
        if self.backend == "numba" and not (HAS_T2W_NUMBA and HAS_W2T_NUMBA):
            raise ImportError("Numba backend requires numba to be installed.")

        Cos, Cos2, SinCos, CosSize, Cos2Size, SinCosSize = load_fourier_coeff()
        self._cos2 = Cos2
        self._sincos = SinCos
        self._cos2_size = Cos2Size
        self._sincos_size = SinCosSize

        if self.filter is None:
            # Compute filter taps then truncate to the minimum length N such that:
            #   (a) reconstruction residual < 10^{-precision}
            #   (b) (N-1) is divisible by 2*M  (block-alignment requirement)
            #   (c) N // (2*M) >= 3            (at least 3 complete blocks)
            n_max = int(3e5)
            filt = get_filter(self.M, self.K, self.beta_order, n_max, Cos, CosSize)
            self.filter, self.m_H = self._truncate_filter(filt, self.M, self.precision)
        else:
            filt_arr = jnp.asarray(self.filter, dtype=jnp.float64).reshape(-1)
            m_h = int(filt_arr.shape[0]) if self.m_H is None else int(self.m_H)
            if m_h <= 0:
                raise ValueError("m_H must be positive when providing a custom filter.")
            if m_h > int(filt_arr.shape[0]):
                raise ValueError("m_H cannot be greater than the provided filter length.")
            self.filter = filt_arr[:m_h]
            self.m_H = m_h

    @staticmethod
    def _truncate_filter(filt, M, precision):
        """Return (filter_array, N) with N the minimum sufficient filter length."""
        residual = 1 - filt[0] ** 2
        prec = 10 ** (-precision)
        N = 1
        M2 = M * 2
        while residual > prec or (N - 1) % M2 or N // M2 < 3:
            residual -= 2 * filt[N] ** 2
            N += 1
        return jnp.asarray(filt[:N], dtype=jnp.float64), N

    @property
    def max_level(self):
        return self.M

    @staticmethod
    def _validate_backend(backend: str) -> str:
        """Validate and normalise the backend name at instance creation time."""
        name = str(backend).strip().lower()
        supported = {"jax", "numba"}
        if name not in supported:
            raise ValueError(
                f"Unsupported backend '{name}'. "
                f"Available backends: {sorted(supported)}."
            )
        return name

    @staticmethod
    def _coerce_timeseries_input(strain, sample_rate, t0):
        """Normalise heterogeneous strain inputs to (numpy.array, sample_rate, t0)."""
        # gwpy TimeSeries — carries sr and t0 as Quantities
        if isinstance(strain, TimeSeries):
            return (
                np.asarray(strain.value, dtype=np.float64),
                float(strain.sample_rate.value),
                float(strain.t0.value),
            )

        # PyCBC TimeSeries — detected by module + class name to avoid a hard import
        _t = type(strain)
        if _t.__module__ == "pycbc.types.timeseries" and _t.__name__ == "TimeSeries":
            try:
                return (
                    np.asarray(strain.data, dtype=np.float64),
                    float(strain.sample_rate),
                    float(strain.start_time),
                )
            except AttributeError as exc:
                raise ValueError(
                    "strain appears to be a pycbc TimeSeries but sample_rate or "
                    "start_time could not be extracted."
                ) from exc

        # Plain array — caller must supply sample_rate
        if sample_rate is None:
            raise ValueError(
                "sample_rate must be provided when strain is not a gwpy or pycbc TimeSeries."
            )
        return np.asarray(strain, dtype=np.float64), float(sample_rate), float(t0)

    @partial(jax.jit, static_argnames=("MM",))
    def t2w_jax(self, strain, MM=-1):
        """
        JAX-native forward transform.

        Returns
        -------
        m_L : int
        aligned_length : int
        tf_map : jax.Array, shape (2, n_time, M+1)
        """
        return t2w_jax(self.M, self.m_H, strain, self.filter, MM)

    @partial(jax.jit, static_argnames=("m_layer", "output_length"))
    def w2t_jax(self, tf_map_flattened, m_layer, output_length=None):
        """
        JAX-native inverse transform from flattened coefficient map.
        """
        return w2t_jax(tf_map_flattened, m_layer, self.filter, output_length=output_length)

    @partial(jax.jit, static_argnames=("m_layer", "output_length"))
    def w2tQ_jax(self, tf_map_flattened, m_layer, output_length=None):
        """
        JAX-native inverse transform from flattened quadrature coefficient map.
        """
        return w2tQ_jax(tf_map_flattened, m_layer, self.filter, output_length=output_length)

    def t2w(self, strain, sample_rate=None, t0=0, MM=-1):
        """
        Transform time series to WDM map.

        Uses the backend selected at instance creation (``self.backend``).
        For direct use inside a JAX or Numba JIT context, call the pure
        core functions in ``wdm_wavelet.core.t2w`` instead.
        """
        signal, sample_rate, t0 = self._coerce_timeseries_input(strain, sample_rate, t0)

        if self.backend == "jax":
            _, _, tf_map = t2w_jax(self.M, self.m_H, signal, self.filter, MM)
        else:
            _, _, tf_map = t2w_numba(self.M, self.m_H, signal, self.filter, MM)

        max_f_layer = int(tf_map.shape[2] - 1)
        max_t_layer = int(tf_map.shape[1])
        df = sample_rate / max_f_layer / 2.0

        t_len = int(signal.shape[0]) / sample_rate
        dt = t_len / max_t_layer

        tf_complex = np.asarray(tf_map[0], dtype=np.float64) + 1j * np.asarray(tf_map[1], dtype=np.float64)

        params = dict(self.params)
        params["has_quadrature"] = bool(MM < 0)
        params["t2w_mm"] = int(MM)
        params["t2w_backend"] = self.backend

        return TimeFrequencyMap(np.asarray(tf_complex.T), df, dt, t0, int(signal.shape[0]), params)

    def w2t(self, time_frequency_map):
        """
        Inverse WDM transform.  Uses the backend selected at instance creation.
        """
        data = np.stack(
            (
                np.asarray(time_frequency_map.data.real.T, dtype=np.float64),
                np.asarray(time_frequency_map.data.imag.T, dtype=np.float64),
            ),
            axis=0,
        ).reshape(-1)

        target_len = getattr(time_frequency_map, "len_timeseries", None)
        if self.backend == "jax":
            t = w2t_jax(data, time_frequency_map.n_freq, self.filter, output_length=target_len)
        else:
            t = w2t_numba(data, time_frequency_map.n_freq, self.filter, output_length=target_len)

        sample_rate = 2 * time_frequency_map.df * (time_frequency_map.n_freq - 1)
        return TimeSeries(np.asarray(t, dtype=np.float64), sample_rate=sample_rate, t0=time_frequency_map.t0)

    def w2tQ(self, time_frequency_map):
        """
        Inverse WDM transform using the 90-degree quadrature.  Uses the backend
        selected at instance creation.
        """
        params = getattr(time_frequency_map, "wdm_params", None)
        if isinstance(params, dict) and params.get("has_quadrature") is False:
            raise ValueError(
                "w2tQ requires 00/90 quadratures. "
                "This TF map was created in power-only mode (MM >= 0)."
            )

        data = np.stack(
            (
                np.asarray(time_frequency_map.data.real.T, dtype=np.float64),
                np.asarray(time_frequency_map.data.imag.T, dtype=np.float64),
            ),
            axis=0,
        ).reshape(-1)

        target_len = getattr(time_frequency_map, "len_timeseries", None)
        if self.backend == "jax":
            t = w2tQ_jax(data, time_frequency_map.n_freq, self.filter, output_length=target_len)
        else:
            t = w2tQ_numba(data, time_frequency_map.n_freq, self.filter, output_length=target_len)

        sample_rate = 2 * time_frequency_map.df * (time_frequency_map.n_freq - 1)
        return TimeSeries(np.asarray(t, dtype=np.float64), sample_rate=sample_rate, t0=time_frequency_map.t0)

    def set_td_filter(self, n_coeffs, L=1):
        """
        Build time-delay filters (paper Sec. 4, Eq. (27)-(32)).

        Parameters
        ----------
        n_coeffs : int
            Half-width of symmetric TD filters.
        L : int, optional
            Delay upsampling factor. Fundamental delay step is tau/L.
        """
        self.td_filters = _set_td_filter(
            self.M,
            self.K,
            self.beta_order,
            n_coeffs,
            L,
            self._cos2,
            self._sincos,
            self._cos2_size,
            self._sincos_size,
        )
        return self.td_filters

    def get_pixel_amplitude(self, time_frequency_map, m, n, dT, quad=False):
        """
        Return delayed amplitude for one TF pixel (m,n).
        """
        if self.td_filters is None:
            raise ValueError("TD filters are not initialized. Call set_td_filter(...) first.")
        self._ensure_td_map_compatible(time_frequency_map)
        tf00, tf90 = _split_tf_map(time_frequency_map, self.M)
        return _get_pixel_amplitude_td(tf00, tf90, self.td_filters, m, n, dT, quad=quad)

    def get_td_amp(self, time_frequency_map, pixel_index, delay_index, mode="p"):
        """
        Return one delayed amplitude/power sample for a TF pixel index.
        """
        if self.td_filters is None:
            raise ValueError("TD filters are not initialized. Call set_td_filter(...) first.")
        self._ensure_td_map_compatible(time_frequency_map)
        tf00, tf90 = _split_tf_map(time_frequency_map, self.M)
        return _get_td_amp(tf00, tf90, self.td_filters, pixel_index, delay_index, mode=mode)

    def get_td_vec(self, time_frequency_map, pixel_index, K, mode="p"):
        """
        Return delayed amplitudes/power over delay range [-K, ..., K].
        """
        if self.td_filters is None:
            raise ValueError("TD filters are not initialized. Call set_td_filter(...) first.")
        self._ensure_td_map_compatible(time_frequency_map)
        tf00, tf90 = _split_tf_map(time_frequency_map, self.M)
        return _get_td_vec(tf00, tf90, self.td_filters, pixel_index, K, mode=mode)

    @staticmethod
    def _ensure_td_map_compatible(time_frequency_map):
        """
        Validate that TF map carries both quadratures required by TD filters.
        """
        params = getattr(time_frequency_map, "wdm_params", None)
        if isinstance(params, dict) and params.get("has_quadrature") is False:
            raise ValueError(
                "Time-delay amplitudes require 00/90 quadratures. "
                "This TF map was created in power-only mode (MM >= 0)."
            )

    @property
    def params(self):
        return {
            "M": self.M,
            "K": self.K,
            "beta_order": self.beta_order,
            "precision": self.precision,
        }


jax.tree_util.register_dataclass(
    WDM,
    data_fields=["filter"],
    meta_fields=["M", "K", "beta_order", "precision", "backend", "m_H"],
    drop_fields=["td_filters", "_cos2", "_sincos", "_cos2_size", "_sincos_size"],
)
