"""Synthetic macro data generation for scenario backtesting.

Generates the full macro_by_date dict consumed by SnapshotStore, including:
  - SPY/QQQ rolling indicators (SMA, momentum, vol, entropy, Hurst, etc.)
  - VIX via Ornstein-Uhlenbeck process correlated with market returns
  - HY credit spread via OU correlated with VIX
  - Proxy ETF close prices (TQQQ, QQQM, SPYM, SGOV)

The rolling indicators exactly mirror SnapshotStore._fetch_spy_macro._merge_series()
so that allocator Gates (macro regime, VIX scaler, parking momentum) behave correctly.
"""
from __future__ import annotations

import datetime as dt
import math
from dataclasses import dataclass
from typing import Any

import numpy as np


@dataclass
class VIXConfig:
    """Parameters for the VIX Ornstein-Uhlenbeck simulation."""

    base_level: float = 18.0
    """Long-run mean (θ). Bull ~14, normal ~18, bear ~25, crash ~35+."""
    mean_reversion: float = 5.0
    """Speed of reversion (κ, annualized). Higher = faster mean-reversion."""
    volatility: float = 5.0
    """VIX process volatility (σ, per year)."""
    market_corr: float = -0.7
    """Correlation between VIX changes and market log-returns (negative)."""


@dataclass
class HYSpreadConfig:
    """Parameters for the HY credit spread OU simulation."""

    base_level: float = 4.0
    """Long-run mean spread in percentage points."""
    mean_reversion: float = 3.0
    """Speed of reversion (annualized)."""
    volatility: float = 1.5
    """Spread process volatility (per year)."""
    vix_corr: float = 0.6
    """Correlation with VIX level changes."""


# ---------------------------------------------------------------------------
# Internal: rolling indicator computation (mirrors _merge_series in snapshot_store.py)
# ---------------------------------------------------------------------------

_SMA_PERIODS = (10, 20, 30, 40, 50)
_ROLLING_HIGH_PERIODS = (20, 50, 100)
_MOMENTUM_PERIODS = (5, 10, 20, 50)
_VOL_PERIODS = (15, 20, 30, 50)
_EFFICIENCY_PERIODS = (10, 20)
_DOWNSIDE_VOL_PERIODS = (10, 20)
_ULCER_PERIODS = (10, 20)
_ENTROPY_PERIODS = (10, 20)
_DRAWDOWN_ACCEL_DAYS = 5


def _compute_rolling_indicators(
    prefix: str,
    bars: dict[dt.date, dict[str, Any]],
    result: dict[dt.date, dict[str, Any]],
) -> None:
    """Compute all rolling macro indicators for a price series prefix (spy, qqq, etc.).

    Fills result[date][f"{prefix}_*"] in-place, mirroring exactly the keys
    produced by SnapshotStore._fetch_spy_macro._merge_series().
    """
    sorted_dates = sorted(bars.keys())
    closes: list[tuple[dt.date, float]] = [
        (d, float(bars[d]["close"])) for d in sorted_dates
    ]

    for i, (date, close) in enumerate(closes):
        result.setdefault(date, {})
        bar = bars[date]
        result[date][f"{prefix}_close"] = close
        result[date][f"{prefix}_open"] = float(bar.get("open", close))
        result[date][f"{prefix}_high"] = float(bar.get("high", close))
        result[date][f"{prefix}_low"] = float(bar.get("low", close))
        result[date][f"{prefix}_volume"] = float(bar.get("volume", 0))

        # ---- SMAs ----
        for period in _SMA_PERIODS:
            sma = None
            if i >= period - 1:
                window = [c for _, c in closes[i - period + 1 : i + 1]]
                sma = sum(window) / len(window)
            result[date][f"{prefix}_sma_{period}"] = sma

        # ---- Rolling highs (for drawdown gates) ----
        for rh_p in _ROLLING_HIGH_PERIODS:
            rh = None
            if i >= rh_p - 1:
                rh = max(c for _, c in closes[i - rh_p + 1 : i + 1])
            result[date][f"{prefix}_high_{rh_p}"] = rh

        # ---- Momentum / N-day return ----
        for mom_p in _MOMENTUM_PERIODS:
            mom = None
            if i >= mom_p:
                prev = closes[i - mom_p][1]
                if prev > 0:
                    mom = (close - prev) / prev
            result[date][f"{prefix}_mom_{mom_p}"] = mom

        # ---- Realized volatility (annualised std of log returns) ----
        for vol_p in _VOL_PERIODS:
            vol = None
            if i >= vol_p:
                log_rets = [
                    math.log(closes[j][1] / closes[j - 1][1])
                    for j in range(i - vol_p + 1, i + 1)
                    if closes[j - 1][1] > 0
                ]
                if len(log_rets) >= vol_p - 1:
                    mean_r = sum(log_rets) / len(log_rets)
                    var_r = sum((r - mean_r) ** 2 for r in log_rets) / len(log_rets)
                    vol = math.sqrt(var_r * 252)
            result[date][f"{prefix}_vol_{vol_p}"] = vol

        # ---- Efficiency ratio (Kaufman) ----
        for eff_p in _EFFICIENCY_PERIODS:
            efficiency = None
            if i >= eff_p:
                net = abs(close - closes[i - eff_p][1])
                gross = sum(
                    abs(closes[j][1] - closes[j - 1][1])
                    for j in range(i - eff_p + 1, i + 1)
                )
                efficiency = net / gross if gross > 0 else 0.0
            result[date][f"{prefix}_efficiency_{eff_p}"] = efficiency

        # ---- Downside semi-volatility ----
        for dv_p in _DOWNSIDE_VOL_PERIODS:
            downside_vol = None
            if i >= dv_p:
                neg_sq = [
                    min(closes[j][1] / closes[j - 1][1] - 1, 0.0) ** 2
                    for j in range(i - dv_p + 1, i + 1)
                    if closes[j - 1][1] > 0
                ]
                if len(neg_sq) >= dv_p - 1:
                    downside_vol = math.sqrt(sum(neg_sq) / len(neg_sq) * 252)
            result[date][f"{prefix}_downside_vol_{dv_p}"] = downside_vol

        # ---- Shannon entropy (market predictability) ----
        for ent_p in _ENTROPY_PERIODS:
            entropy = None
            if i >= ent_p:
                daily_rets = [
                    closes[j][1] / closes[j - 1][1] - 1
                    for j in range(i - ent_p + 1, i + 1)
                    if closes[j - 1][1] > 0
                ]
                if len(daily_rets) >= ent_p - 1:
                    n_pos = sum(1 for r in daily_rets if r > 0.001)
                    n_neg = sum(1 for r in daily_rets if r < -0.001)
                    n_flat = len(daily_rets) - n_pos - n_neg
                    entropy = 0.0
                    for cnt in (n_pos, n_neg, n_flat):
                        if cnt > 0:
                            p = cnt / len(daily_rets)
                            entropy -= p * math.log2(p)
            result[date][f"{prefix}_entropy_{ent_p}"] = entropy

        # ---- Ulcer index + current drawdown ----
        for ulcer_p in _ULCER_PERIODS:
            ulcer = None
            current_dd = None
            if i >= ulcer_p - 1:
                window_closes = [c for _, c in closes[i - ulcer_p + 1 : i + 1]]
                peak = 0.0
                drawdowns: list[float] = []
                for wc in window_closes:
                    peak = max(peak, wc)
                    if peak > 0:
                        drawdowns.append(wc / peak - 1.0)
                if drawdowns:
                    ulcer = math.sqrt(sum(dd * dd for dd in drawdowns) / len(drawdowns))
                    current_dd = abs(drawdowns[-1])
            result[date][f"{prefix}_ulcer_{ulcer_p}"] = ulcer
            result[date][f"{prefix}_drawdown_{ulcer_p}"] = current_dd

        # ---- Drawdown acceleration ----
        dd_lb = 20
        dd_accel = None
        if i >= dd_lb - 1 + _DRAWDOWN_ACCEL_DAYS:
            cur_window = [c for _, c in closes[i - dd_lb + 1 : i + 1]]
            prev_i = i - _DRAWDOWN_ACCEL_DAYS
            prev_window = [c for _, c in closes[prev_i - dd_lb + 1 : prev_i + 1]]
            cur_peak = max(cur_window) if cur_window else 0.0
            prev_peak = max(prev_window) if prev_window else 0.0
            if cur_peak > 0 and prev_peak > 0:
                cur_dd = (cur_peak - close) / cur_peak
                prev_dd = (prev_peak - closes[prev_i][1]) / prev_peak
                dd_accel = cur_dd - prev_dd
        result[date][f"{prefix}_drawdown_accel_{_DRAWDOWN_ACCEL_DAYS}"] = dd_accel

        # ---- Hurst exponent (R/S analysis) ----
        hurst_lookback = 60
        hurst = None
        if i >= hurst_lookback + 1:
            h_rets = [
                (closes[j + 1][1] - closes[j][1]) / closes[j][1]
                for j in range(i - hurst_lookback, i)
                if closes[j][1] > 0
            ]
            if len(h_rets) >= 30:
                def _rs(series: list[float]) -> float:
                    n_ = len(series)
                    m_ = sum(series) / n_
                    devs = [x - m_ for x in series]
                    cum, s_ = [], 0.0
                    for d_ in devs:
                        s_ += d_
                        cum.append(s_)
                    r_ = max(cum) - min(cum)
                    std_ = (sum(d_ ** 2 for d_ in devs) / n_) ** 0.5
                    return r_ / std_ if std_ > 0 else 0.0

                win_sizes = [s for s in [8, 12, 16, 24, 32] if s <= len(h_rets) // 2]
                if len(win_sizes) >= 2:
                    log_n, log_rs = [], []
                    for w in win_sizes:
                        chunks = [h_rets[st: st + w] for st in range(0, len(h_rets) - w + 1, w) if len(h_rets[st: st + w]) == w]
                        rs_vals = [_rs(c) for c in chunks]
                        if rs_vals:
                            avg_rs = sum(rs_vals) / len(rs_vals)
                            if avg_rs > 0:
                                log_n.append(math.log(w))
                                log_rs.append(math.log(avg_rs))
                    if len(log_n) >= 2:
                        n_h = len(log_n)
                        xm = sum(log_n) / n_h
                        ym = sum(log_rs) / n_h
                        num = sum((log_n[k] - xm) * (log_rs[k] - ym) for k in range(n_h))
                        den = sum((log_n[k] - xm) ** 2 for k in range(n_h))
                        hurst = num / den if den > 0 else 0.5
        result[date][f"{prefix}_hurst_60"] = hurst

        # ---- Lag-1 autocorrelation (Lo, 2004) ----
        ac_lb = 20
        autocorr = None
        if i >= ac_lb + 1:
            ac_rets = [
                closes[j][1] / closes[j - 1][1] - 1
                for j in range(i - ac_lb, i + 1)
                if closes[j - 1][1] > 0
            ]
            if len(ac_rets) >= ac_lb:
                x_ac, y_ac = ac_rets[:-1], ac_rets[1:]
                n_ac = len(x_ac)
                mx, my = sum(x_ac) / n_ac, sum(y_ac) / n_ac
                cov_xy = sum((x_ac[k] - mx) * (y_ac[k] - my) for k in range(n_ac)) / n_ac
                sx = (sum((x_ac[k] - mx) ** 2 for k in range(n_ac)) / n_ac) ** 0.5
                sy = (sum((y_ac[k] - my) ** 2 for k in range(n_ac)) / n_ac) ** 0.5
                if sx > 1e-12 and sy > 1e-12:
                    autocorr = cov_xy / (sx * sy)
        result[date][f"{prefix}_autocorr_20"] = autocorr


def _compute_pair_correlation(
    left_prefix: str,
    right_prefix: str,
    output_key: str,
    result: dict[dt.date, dict[str, Any]],
) -> None:
    """Compute rolling 20-day cross-asset correlation (mirrors snapshot_store logic)."""
    corr_lb = 20
    sorted_dates = sorted(result.keys())
    for idx_c, d_c in enumerate(sorted_dates):
        corr_val = None
        if idx_c >= corr_lb:
            left_r, right_r = [], []
            for jj in range(idx_c - corr_lb + 1, idx_c + 1):
                d_j = sorted_dates[jj]
                d_prev = sorted_dates[jj - 1]
                lc = result.get(d_j, {}).get(f"{left_prefix}_close")
                lp = result.get(d_prev, {}).get(f"{left_prefix}_close")
                rc = result.get(d_j, {}).get(f"{right_prefix}_close")
                rp = result.get(d_prev, {}).get(f"{right_prefix}_close")
                if all(v and v > 0 for v in [lc, lp, rc, rp]):
                    left_r.append(lc / lp - 1)
                    right_r.append(rc / rp - 1)
            if len(left_r) >= corr_lb - 2:
                n_cr = len(left_r)
                ml = sum(left_r) / n_cr
                mr = sum(right_r) / n_cr
                cov = sum((left_r[k] - ml) * (right_r[k] - mr) for k in range(n_cr)) / n_cr
                sl = (sum((left_r[k] - ml) ** 2 for k in range(n_cr)) / n_cr) ** 0.5
                sr = (sum((right_r[k] - mr) ** 2 for k in range(n_cr)) / n_cr) ** 0.5
                if sl > 1e-12 and sr > 1e-12:
                    corr_val = cov / (sl * sr)
        result[d_c][output_key] = corr_val


# ---------------------------------------------------------------------------
# VIX and HY spread simulation
# ---------------------------------------------------------------------------

def _generate_vix_series(
    market_log_rets: list[float],
    trading_dates: list[dt.date],
    config: VIXConfig,
    rng: np.random.Generator,
) -> dict[dt.date, float]:
    """Generate VIX time series via OU process correlated with market returns.

    dVIX = κ(θ-VIX)dt + σ·dW  where dW is partially driven by market return sign.
    """
    dt_step = 1.0 / 252
    vix = config.base_level
    vix_series: dict[dt.date, float] = {}

    for date, mkt_lr in zip(trading_dates, market_log_rets):
        # Approximate market z-score from log return
        market_daily_std = config.base_level / 100.0 * math.sqrt(dt_step) + 1e-8
        mkt_z = mkt_lr / market_daily_std
        z_idio = rng.standard_normal()
        rho = config.market_corr
        # Negative market → positive VIX shock
        z_combined = rho * (-mkt_z) + math.sqrt(max(0.0, 1 - rho ** 2)) * z_idio

        kappa = config.mean_reversion
        theta = config.base_level
        sigma = config.volatility

        dVIX = kappa * (theta - vix) * dt_step + sigma * math.sqrt(dt_step) * z_combined
        vix = max(5.0, min(90.0, vix + dVIX))
        vix_series[date] = round(vix, 2)

    return vix_series


def _generate_hy_series(
    vix_series: dict[dt.date, float],
    trading_dates: list[dt.date],
    config: HYSpreadConfig,
    rng: np.random.Generator,
) -> dict[dt.date, float]:
    """Generate HY credit spread via OU process correlated with VIX changes."""
    dt_step = 1.0 / 252
    spread = config.base_level
    hy_series: dict[dt.date, float] = {}
    prev_vix = config.base_level

    for date in trading_dates:
        cur_vix = vix_series.get(date, config.base_level)
        vix_chg = (cur_vix - prev_vix) / (prev_vix + 1e-8)

        z_idio = rng.standard_normal()
        rho = config.vix_corr
        z_combined = rho * vix_chg * 5.0 + math.sqrt(max(0.0, 1 - rho ** 2)) * z_idio

        kappa = config.mean_reversion
        theta = config.base_level
        sigma = config.volatility

        dSpread = kappa * (theta - spread) * dt_step + sigma * math.sqrt(dt_step) * z_combined
        spread = max(1.5, min(30.0, spread + dSpread))
        hy_series[date] = round(spread, 3)
        prev_vix = cur_vix

    return hy_series


# ---------------------------------------------------------------------------
# Main entry point
# ---------------------------------------------------------------------------

def generate_macro_data(
    spy_bars: dict[dt.date, dict[str, Any]],
    qqq_bars: dict[dt.date, dict[str, Any]],
    vix_config: VIXConfig,
    hy_config: HYSpreadConfig,
    trading_dates: list[dt.date],
    rng: np.random.Generator,
    market_log_rets: list[float] | None = None,
) -> dict[dt.date, dict[str, Any]]:
    """Generate complete macro_by_date dict for SnapshotStore.

    Includes:
      - SPY/QQQ rolling indicators (matches _merge_series keys exactly)
      - spy_qqq_corr_20 cross-correlation
      - VIXCLS, macro_vix, macro_hy_spread
      - Parking ETF proxies: tqqq_close, qqqm_close, spym_close, sgov_close

    Args:
        spy_bars: SPY OHLCV bars {date: {open, high, low, close, volume}}.
        qqq_bars: QQQ OHLCV bars.
        vix_config: VIX OU simulation parameters.
        hy_config: HY spread OU simulation parameters.
        trading_dates: Ordered list of NYSE trading dates.
        rng: NumPy random generator.
        market_log_rets: Optional market log-return series for VIX correlation.

    Returns:
        macro_by_date dict compatible with SnapshotStore.
    """
    result: dict[dt.date, dict[str, Any]] = {}

    # Rolling indicators for SPY and QQQ
    _compute_rolling_indicators("spy", spy_bars, result)
    _compute_rolling_indicators("qqq", qqq_bars, result)

    # SPY-QQQ cross correlation (regime shift detection)
    _compute_pair_correlation("spy", "qqq", "spy_qqq_corr_20", result)

    # Generate VIX and HY spread
    if market_log_rets is None:
        market_log_rets = [0.0] * len(trading_dates)

    vix_series = _generate_vix_series(market_log_rets, trading_dates, vix_config, rng)
    hy_series = _generate_hy_series(vix_series, trading_dates, hy_config, rng)

    # Inject VIX, HY, and parking ETF proxies
    sgov_price = 100.0
    sgov_daily_yield = 0.0525 / 252  # ~5.25% money market yield

    for i, date in enumerate(trading_dates):
        result.setdefault(date, {})

        vix = vix_series.get(date)
        if vix is not None:
            result[date]["VIXCLS"] = vix
            result[date]["macro_vix"] = vix

        hy = hy_series.get(date)
        if hy is not None:
            result[date]["macro_hy_spread"] = hy

        qqq_close = result[date].get("qqq_close")
        spy_close = result[date].get("spy_close")

        if qqq_close is not None and qqq_close > 0:
            # TQQQ ≈ 3x leveraged QQQ (simplified proxy)
            result[date]["tqqq_close"] = round(qqq_close * 3.0 / 415.0 * 55.0, 4)
            # QQQM ≈ QQQ (slightly lower price, same ETF)
            result[date]["qqqm_close"] = round(qqq_close * 0.99, 4)

        if spy_close is not None and spy_close > 0:
            # SPYM ≈ SPY (2x leveraged; simplified as 1.9× SPY level / 480 × 95)
            result[date]["spym_close"] = round(spy_close * 0.98, 4)

        # SGOV ≈ T-bill ETF with daily accrual
        sgov_price *= (1 + sgov_daily_yield)
        result[date]["sgov_close"] = round(sgov_price, 4)

    return result
