"""
Wavelet band decomposition (CWT + synchrosqueezed variant), ported near-verbatim from
c:\\DataS\\InstrumentSimulator\\backend\\app\\wavelet.py (project "Macro Causal Lab").

No project-specific dependencies — pure numpy/ssqueezepy, safe to call from any service.
"""
import warnings

import numpy as np
from ssqueezepy import cwt, icwt
from ssqueezepy.utils.cwt_utils import center_frequency

MIN_SCALES_PER_BAND = 4

# Fixed period (days) lower-bounds for bands 0..5, independent of the analysis
# window's length. Without this, dividing [scales.min(), scales.max()] into
# num_levels *equal* log-spaced buckets means "band i" covers a totally
# different real-world period range depending on how much history is fed in
# (a longer window resolves longer scales, which shifts every intermediate
# boundary upward) - so the same band index is not comparable across window
# lengths. Anchoring to absolute day-periods keeps "band i" meaning the same
# oscillation whether the window is 3 months or 3 years; only the last band's
# upper bound stays open, since a longer window can genuinely resolve slower
# cycles a short one cannot.
CANONICAL_PERIOD_EDGES_DAYS = [0.3, 1.5, 4.0, 12.0, 35.0, 90.0]


def _period_to_log_scale_fn(wavelet, scales, n, num_calib_points=16):
    """Build an interpolation from log(period_days) to log(scale) by sampling
    center_frequency at a handful of scales (not every scale - it's not
    vectorized and can be slow over hundreds/thousands of scales).
    """
    log_scales_full = np.log(scales)
    calib_log_scales = np.linspace(log_scales_full.min(), log_scales_full.max(), min(num_calib_points, len(scales)))
    calib_periods = np.array([1 / center_frequency(wavelet, scale=float(np.exp(s)), N=n) for s in calib_log_scales])
    calib_log_periods = np.log(calib_periods)

    def period_to_log_scale(period_days):
        return float(np.interp(np.log(period_days), calib_log_periods, calib_log_scales))

    return period_to_log_scale


def band_decompose(
    values: list[float],
    dates: list[str],
    num_levels: int = 4,
    wavelet: str = "gmw",
) -> dict:
    x = np.asarray(values, dtype=float)
    n = len(x)
    if n < 32:
        raise ValueError("Serie trop courte pour une analyse ondelette (32 points minimum).")

    x_mean = float(x.mean())
    xc = x - x_mean

    Wx, scales = cwt(xc, wavelet=wavelet)
    log_scales = np.log(scales)

    period_to_log_scale = _period_to_log_scale_fn(wavelet, scales, n)
    period_edges = list(CANONICAL_PERIOD_EDGES_DAYS[:num_levels])
    while len(period_edges) < num_levels:
        period_edges.append(period_edges[-1] * 2.5)

    # scales[0] = smallest scale = highest frequency = shortest period (impulses/noise);
    # scales[-1] = largest scale = lowest frequency = longest period (slow trend).
    bands = []
    for i in range(num_levels):
        lo = period_to_log_scale(period_edges[i])
        hi = period_to_log_scale(period_edges[i + 1]) if i + 1 < len(period_edges) else log_scales.max()
        mask = (log_scales >= lo) & (log_scales <= hi if i == num_levels - 1 else log_scales < hi)
        scales_band = scales[mask]

        if len(scales_band) >= MIN_SCALES_PER_BAND:
            try:
                recon = icwt(Wx[mask, :], wavelet=wavelet, scales=scales_band, x_len=n)
            except Exception:
                # ssqueezepy's internal scale-type inference (log vs linear spacing
                # detection) can raise on some narrow slices depending on how many
                # scales land in this bucket for this specific series length.
                # Treat as a negligible band instead of failing the whole request.
                recon = np.zeros(n)
            scale_lo, scale_hi = float(scales_band.min()), float(scales_band.max())
        else:
            # Too few scales in this bucket for a stable reconstruction; report a
            # flat zero band (its content ends up folded into the residual) so the
            # band list always stays num_levels long.
            recon = np.zeros(n)
            scale_lo, scale_hi = float(np.exp(lo)), float(np.exp(hi))

        freq_at_scale_hi = center_frequency(wavelet, scale=scale_hi, N=n)
        freq_at_scale_lo = center_frequency(wavelet, scale=scale_lo, N=n)
        period_low_days = round(1 / freq_at_scale_lo, 1) if freq_at_scale_lo else None
        period_high_days = round(1 / freq_at_scale_hi, 1) if freq_at_scale_hi else None
        label = (
            f"{period_low_days}-{period_high_days}j"
            if period_low_days is not None and period_high_days is not None
            else f"bande {i + 1}"
        )
        bands.append(
            {
                "index": i,
                "label": label,
                "period_low_days": period_low_days,
                "period_high_days": period_high_days,
                "series": [round(float(v), 6) for v in recon],
            }
        )

    reconstructed_total = x_mean + sum(np.array(band["series"]) for band in bands)
    residual = x - reconstructed_total

    return {
        "dates": dates,
        "original": [round(float(v), 6) for v in x],
        "mean": round(x_mean, 6),
        "bands": bands,
        "residual": [round(float(v), 6) for v in residual],
        "wavelet": wavelet,
    }


def windowed_band_decompose(
    values: list[float],
    dates: list[str],
    window_size: int,
    num_levels: int = 4,
    wavelet: str = "gmw",
) -> dict:
    """Decompose a long series by running band_decompose independently on
    consecutive slices of `window_size` points, then concatenating the
    results. Each slice is analyzed on its own terms (own local mean, own
    CWT scale range), so precision near the reconstructed curve stays high
    even over a long history; the trade-off is a visible discontinuity at
    each slice boundary, which is acceptable for this diagnostic view.
    """
    n = len(values)
    window_size = max(32, min(window_size, n)) if window_size else n
    if n <= window_size:
        result = band_decompose(values, dates, num_levels, wavelet)
        result["window_size"] = window_size
        result["chunks"] = 1
        return result

    all_dates: list[str] = []
    all_original: list[float] = []
    all_residual: list[float] = []
    band_series: list[list[float]] = [[] for _ in range(num_levels)]
    band_meta: list[dict | None] = [None] * num_levels

    chunk_count = 0
    start = 0
    while start < n:
        end = min(start + window_size, n)
        if 0 < n - end < 32:
            # Don't leave a too-short trailing slice: fold the remainder in.
            end = n
        chunk = band_decompose(values[start:end], dates[start:end], num_levels, wavelet)
        chunk_count += 1
        all_dates.extend(chunk["dates"])
        all_original.extend(chunk["original"])
        all_residual.extend(chunk["residual"])
        for band in chunk["bands"]:
            band_series[band["index"]].extend(band["series"])
            band_meta[band["index"]] = band
        start = end

    bands = []
    for i in range(num_levels):
        meta = band_meta[i] or {"label": f"bande {i + 1}", "period_low_days": None, "period_high_days": None}
        bands.append(
            {
                "index": i,
                "label": meta["label"],
                "period_low_days": meta["period_low_days"],
                "period_high_days": meta["period_high_days"],
                "series": band_series[i],
            }
        )

    return {
        "dates": all_dates,
        "original": all_original,
        "mean": round(float(np.mean(all_original)), 6) if all_original else 0.0,
        "bands": bands,
        "residual": all_residual,
        "wavelet": wavelet,
        "window_size": window_size,
        "chunks": chunk_count,
    }


def rolling_causal_bands(
    values: list[float],
    dates: list[str],
    start_idx: int,
    lookback: int,
    num_levels: int = 4,
    wavelet: str = "gmw",
    step: int = 1,
) -> dict:
    """Walk-forward band decomposition with no look-ahead: the band value
    reported for day `t` is computed from a CWT run only on the trailing
    `lookback` points ending at `t` (never anything after `t`). This is
    deliberately much slower than `band_decompose` (one CWT per day instead
    of one for the whole range) - that cost is the point: a single
    whole-range CWT lets every point "see" the entire future through the
    transform's global/symmetric support, which makes any backtest built on
    it meaningless (the peaks/troughs it finds were computed with hindsight).
    `step` > 1 re-runs the CWT every `step` days and holds the last computed
    value in between, trading fidelity for speed on long ranges.
    """
    n = len(values)
    out_dates: list[str] = []
    out_original: list[float] = []
    band_series: list[list[float]] = [[] for _ in range(num_levels)]
    band_labels = [f"bande {i + 1}" for i in range(num_levels)]
    last_tip: list[float] | None = None
    recomputations = 0

    for t in range(start_idx, n):
        window_start = t - lookback + 1
        if window_start < 0:
            continue  # not enough trailing history yet for a full lookback window
        if last_tip is None or (t - start_idx) % step == 0:
            window_values = values[window_start:t + 1]
            window_dates = dates[window_start:t + 1]
            result = band_decompose(window_values, window_dates, num_levels, wavelet)
            last_tip = [band["series"][-1] for band in result["bands"]]
            band_labels = [band["label"] for band in result["bands"]]
            recomputations += 1
        out_dates.append(dates[t])
        out_original.append(values[t])
        for i in range(num_levels):
            band_series[i].append(last_tip[i])

    bands = [
        {
            "index": i,
            "label": band_labels[i],
            "period_low_days": None,
            "period_high_days": None,
            "series": band_series[i],
        }
        for i in range(num_levels)
    ]

    return {
        "dates": out_dates,
        "original": out_original,
        "bands": bands,
        "wavelet": wavelet,
        "lookback": lookback,
        "step": step,
        "recomputations": recomputations,
    }


# ---------------------------------------------------------------------------
# Synchrosqueezed variant (opt-in, method="ssq"). Everything above this line
# is completely untouched by what follows: band_decompose, windowed_band_decompose
# and rolling_causal_bands are the exact same functions used when method="cwt"
# (the default), so any existing simulation re-run with its original config
# hits the same code path and produces byte-identical results as before.
#
# Rationale: plain cwt/icwt smears a single true frequency's energy across
# several neighboring scales (that's *why* band_decompose needs canonical period
# edges and a minimum-scales guard - the boundaries are fuzzy). Synchrosqueezing
# reassigns that smeared energy onto its estimated true instantaneous frequency
# before splitting into bands, which should reduce cross-band leakage. It also
# exposes two things plain cwt does not: per-band energy |Tx|^2 (how dominant
# that cycle is *right now*, independent of direction) and a dominant-frequency
# ridge (which cycle length is currently winning, tracked over time).
#
# Separate period edges from the cwt path's CANONICAL_PERIOD_EDGES_DAYS: with
# fs=1.0 (one sample = one day, matching our daily data), ssq_cwt has a hard
# Nyquist floor at exactly 2.0 days (max resolvable frequency = fs/2). The cwt
# path's edges start at 0.3 days - below that floor, so reusing them here would
# leave the fastest band permanently empty. This does mean "band 0" covers a
# different absolute range under ssq than under cwt; that's an inherent
# consequence of ssq's stricter frequency floor, not a rounding choice.
# ---------------------------------------------------------------------------

CANONICAL_PERIOD_EDGES_DAYS_SSQ = [2.0, 4.0, 10.0, 25.0, 60.0, 120.0]


def band_decompose_ssq(
    values: list[float],
    dates: list[str],
    num_levels: int = 4,
    wavelet: str = "gmw",
) -> dict:
    from ssqueezepy import issq_cwt, ssq_cwt
    from ssqueezepy.ridge_extraction import extract_ridges

    x = np.asarray(values, dtype=float)
    n = len(x)
    if n < 32:
        raise ValueError("Serie trop courte pour une analyse ondelette (32 points minimum).")

    x_mean = float(x.mean())
    xc = x - x_mean

    with warnings.catch_warnings():
        warnings.simplefilter("ignore", RuntimeWarning)
        # fs=1.0: one sample = one day, so ssq_freqs comes out directly in
        # cycles/day and 1/ssq_freqs is a period in days with no extra
        # calibration step (unlike band_decompose's scale->period conversion,
        # which needs _period_to_log_scale_fn precisely because plain cwt's
        # scale axis has no absolute unit by itself).
        Tx, _Wx, ssq_freqs, _scales = ssq_cwt(xc, wavelet=wavelet, fs=1.0)

    ssq_freqs = np.asarray(ssq_freqs)
    periods = 1.0 / ssq_freqs

    period_edges = list(CANONICAL_PERIOD_EDGES_DAYS_SSQ[:num_levels])
    while len(period_edges) < num_levels:
        period_edges.append(period_edges[-1] * 2.5)

    bands = []
    for i in range(num_levels):
        lo_period = period_edges[i]
        hi_period = period_edges[i + 1] if i + 1 < len(period_edges) else float(periods.max())
        mask = (periods >= lo_period) & (periods <= hi_period if i == num_levels - 1 else periods < hi_period)

        if mask.any():
            try:
                Tx_band = np.zeros_like(Tx)
                Tx_band[mask, :] = Tx[mask, :]
                with warnings.catch_warnings():
                    warnings.simplefilter("ignore", RuntimeWarning)
                    recon = np.real(issq_cwt(Tx_band, wavelet=wavelet))
            except Exception:
                # Defensive, mirroring band_decompose: never let one band's edge
                # case fail the whole request.
                recon = np.zeros(n)
            energy = np.sum(np.abs(Tx[mask, :]) ** 2, axis=0)
        else:
            recon = np.zeros(n)
            energy = np.zeros(n)

        period_low_days = round(float(lo_period), 1)
        period_high_days = round(float(hi_period), 1)
        bands.append(
            {
                "index": i,
                "label": f"{period_low_days}-{period_high_days}j",
                "period_low_days": period_low_days,
                "period_high_days": period_high_days,
                "series": [round(float(v), 6) for v in recon],
                "energy": [round(float(v), 8) for v in energy],
            }
        )

    reconstructed_total = x_mean + sum(np.array(band["series"]) for band in bands)
    residual = x - reconstructed_total

    try:
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", RuntimeWarning)
            ridge_idxs = extract_ridges(Tx, ssq_freqs, n_ridges=1, bw=4)
        ridge_periods = 1.0 / ssq_freqs[ridge_idxs.ravel()]
        ridge_period_days = [round(float(v), 3) for v in ridge_periods]
    except Exception:
        ridge_period_days = [None] * n

    return {
        "dates": dates,
        "original": [round(float(v), 6) for v in x],
        "mean": round(x_mean, 6),
        "bands": bands,
        "residual": [round(float(v), 6) for v in residual],
        "wavelet": wavelet,
        "ridge_period_days": ridge_period_days,
    }


def rolling_causal_bands_ssq(
    values: list[float],
    dates: list[str],
    start_idx: int,
    lookback: int,
    num_levels: int = 4,
    wavelet: str = "gmw",
    step: int = 1,
) -> dict:
    """Synchrosqueezed counterpart to rolling_causal_bands - identical
    walk-forward, no-look-ahead structure (same day-by-day trailing window,
    same step logic), just calling band_decompose_ssq instead of
    band_decompose. rolling_causal_bands itself is untouched.
    """
    n = len(values)
    out_dates: list[str] = []
    out_original: list[float] = []
    band_series: list[list[float]] = [[] for _ in range(num_levels)]
    energy_series: list[list[float]] = [[] for _ in range(num_levels)]
    ridge_series: list[float | None] = []
    band_labels = [f"bande {i + 1}" for i in range(num_levels)]
    last_tip: list[float] | None = None
    last_energy_tip: list[float] | None = None
    last_ridge_tip: float | None = None
    recomputations = 0

    for t in range(start_idx, n):
        window_start = t - lookback + 1
        if window_start < 0:
            continue
        if last_tip is None or (t - start_idx) % step == 0:
            window_values = values[window_start:t + 1]
            window_dates = dates[window_start:t + 1]
            result = band_decompose_ssq(window_values, window_dates, num_levels, wavelet)
            last_tip = [band["series"][-1] for band in result["bands"]]
            last_energy_tip = [band["energy"][-1] for band in result["bands"]]
            last_ridge_tip = result["ridge_period_days"][-1]
            band_labels = [band["label"] for band in result["bands"]]
            recomputations += 1
        out_dates.append(dates[t])
        out_original.append(values[t])
        for i in range(num_levels):
            band_series[i].append(last_tip[i])
            energy_series[i].append(last_energy_tip[i])
        ridge_series.append(last_ridge_tip)

    bands = [
        {
            "index": i,
            "label": band_labels[i],
            "period_low_days": None,
            "period_high_days": None,
            "series": band_series[i],
            "energy": energy_series[i],
        }
        for i in range(num_levels)
    ]

    return {
        "dates": out_dates,
        "original": out_original,
        "bands": bands,
        "wavelet": wavelet,
        "lookback": lookback,
        "step": step,
        "recomputations": recomputations,
        "method": "ssq",
        "ridge_period_days": ridge_series,
    }