"""Synthetic OHLCV price path generation for scenario backtesting.

Implements Regime-Switching GBM with optional jump-diffusion (Merton model).
Individual stock paths are generated as market factor + beta + idiosyncratic noise
to produce realistic cross-sectional correlations.
"""
from __future__ import annotations

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

import numpy as np


@dataclass
class PriceRegime:
    """Parameters for a single market regime segment."""

    annualized_drift: float
    """Expected annual return, e.g. 0.15 for bull, -0.20 for bear."""
    annualized_vol: float
    """Annual volatility, e.g. 0.15 for calm, 0.35 for stressed."""
    duration_days: int
    """Number of trading days this regime lasts."""
    jump_prob: float = 0.0
    """Per-day probability of a Merton jump event."""
    jump_mean: float = 0.0
    """Mean log-jump size. Negative for crash-type regimes."""
    jump_std: float = 0.02
    """Standard deviation of log-jump size."""


def _generate_regime_log_returns(
    regimes: list[PriceRegime],
    n_days: int,
    rng: np.random.Generator,
) -> list[float]:
    """Generate n_days market-level daily log returns following the regime sequence."""
    log_rets: list[float] = []
    dt_step = 1.0 / 252

    for regime in regimes:
        drift = (regime.annualized_drift - 0.5 * regime.annualized_vol ** 2) * dt_step
        diffusion = regime.annualized_vol * math.sqrt(dt_step)
        days_this_regime = min(regime.duration_days, n_days - len(log_rets))

        for _ in range(days_this_regime):
            lr = drift + diffusion * rng.standard_normal()
            if regime.jump_prob > 0 and rng.random() < regime.jump_prob:
                lr += rng.normal(regime.jump_mean, max(regime.jump_std, 1e-6))
            log_rets.append(lr)

    # Extend with last regime if regimes run short
    last = regimes[-1]
    drift = (last.annualized_drift - 0.5 * last.annualized_vol ** 2) * dt_step
    diffusion = last.annualized_vol * math.sqrt(dt_step)
    while len(log_rets) < n_days:
        log_rets.append(drift + diffusion * rng.standard_normal())

    return log_rets[:n_days]


def _bars_from_log_returns(
    log_rets: list[float],
    trading_dates: list[dt.date],
    initial_price: float,
    intraday_vol_scale: float,
    base_volume: int,
    rng: np.random.Generator,
) -> dict[dt.date, dict[str, Any]]:
    """Build OHLCV bars from a sequence of daily log-returns.

    Ensures OHLCV consistency: low <= min(open, close), high >= max(open, close).
    Volume is log-normally distributed and positively correlated with |return|.
    """
    bars: dict[dt.date, dict[str, Any]] = {}
    prev_close = initial_price

    for date, lr in zip(trading_dates, log_rets):
        close = max(prev_close * math.exp(lr), 0.01)

        # Open: prev_close * small gap (mean-zero noise)
        gap = rng.normal(0, intraday_vol_scale * 0.5)
        open_ = max(prev_close * math.exp(gap), 0.01)

        # Intraday range around open/close extremes
        intraday_noise = abs(rng.normal(0, intraday_vol_scale * 0.7))
        hi_raw = max(open_, close) * (1.0 + intraday_noise)
        lo_raw = min(open_, close) * max(1.0 - intraday_noise, 0.001)

        high = max(hi_raw, open_, close)
        low = min(lo_raw, open_, close)
        low = max(low, 0.01)

        # Volume: log-normal, amplified by absolute return
        vol_factor = 1.0 + 3.0 * abs(math.exp(lr) - 1)
        volume = max(1000, int(rng.lognormal(math.log(base_volume), 0.4) * vol_factor))

        bars[date] = {
            "date": date,
            "open": round(open_, 4),
            "high": round(high, 4),
            "low": round(low, 4),
            "close": round(close, 4),
            "volume": volume,
        }
        prev_close = close

    return bars


def generate_price_paths(
    n_symbols: int,
    initial_prices: list[float] | None,
    regimes: list[PriceRegime],
    trading_dates: list[dt.date],
    market_beta_range: tuple[float, float] = (0.6, 1.2),
    rng: np.random.Generator | None = None,
    tickers: list[str] | None = None,
    return_market_log_rets: bool = False,
) -> (
    tuple[dict[str, dict[dt.date, dict[str, Any]]], list[float]]
    | dict[str, dict[dt.date, dict[str, Any]]]
):
    """Generate correlated OHLCV bars for n_symbols stocks.

    Each stock has a random beta to a shared market factor plus idiosyncratic noise.
    Returns bars_by_symbol_date dict compatible with SnapshotStore.

    Args:
        n_symbols: Number of synthetic stocks to generate.
        initial_prices: Optional list of starting prices (defaults to random $20-$200).
        regimes: Sequence of PriceRegime objects defining the market environment.
        trading_dates: Ordered list of NYSE trading dates (from calendar.get_trading_days).
        market_beta_range: (min, max) range for individual stock betas.
        rng: NumPy random generator (seeded for reproducibility).
        tickers: Optional list of ticker symbols (auto-generated if None).
        return_market_log_rets: If True, also return the market log-return series.

    Returns:
        bars_by_symbol_date dict, or (dict, market_log_rets) if return_market_log_rets=True.
    """
    if rng is None:
        rng = np.random.default_rng()

    n = len(trading_dates)
    if tickers is None:
        tickers = [f"SYM{i:03d}" for i in range(n_symbols)]

    # Typical vol across all regimes (for intraday range scaling)
    avg_vol = sum(r.annualized_vol for r in regimes) / max(len(regimes), 1)
    dt_step = 1.0 / 252

    # Market-level log returns (shared factor)
    market_log_rets = _generate_regime_log_returns(regimes, n, rng)

    bars_by_symbol: dict[str, dict[dt.date, dict[str, Any]]] = {}

    for i, ticker in enumerate(tickers[:n_symbols]):
        if initial_prices and i < len(initial_prices):
            init_price = initial_prices[i]
        else:
            init_price = float(rng.uniform(20.0, 200.0))

        beta = float(rng.uniform(*market_beta_range))
        idio_vol = avg_vol * float(rng.uniform(0.3, 0.8))

        # Build stock log returns: beta * market + idiosyncratic
        stock_log_rets: list[float] = []
        for mkt_lr in market_log_rets:
            idio = rng.normal(0, idio_vol * math.sqrt(dt_step))
            stock_log_rets.append(beta * mkt_lr + idio)

        intraday_scale = (idio_vol + avg_vol * beta) * math.sqrt(dt_step) * 0.5
        base_vol = int(rng.uniform(500_000, 10_000_000))

        bars = _bars_from_log_returns(
            stock_log_rets,
            trading_dates,
            initial_price=init_price,
            intraday_vol_scale=intraday_scale,
            base_volume=base_vol,
            rng=rng,
        )
        bars_by_symbol[ticker] = bars

    if return_market_log_rets:
        return bars_by_symbol, market_log_rets
    return bars_by_symbol


def generate_market_etf_paths(
    regimes: list[PriceRegime],
    trading_dates: list[dt.date],
    rng: np.random.Generator,
    spy_initial: float = 480.0,
    qqq_initial: float = 415.0,
) -> tuple[dict[dt.date, dict[str, Any]], dict[dt.date, dict[str, Any]], list[float]]:
    """Generate SPY and QQQ synthetic price bars plus market log-returns.

    QQQ has slightly higher vol and beta than SPY to reflect tech concentration.

    Returns:
        (spy_bars, qqq_bars, market_log_rets)
    """
    n = len(trading_dates)
    dt_step = 1.0 / 252
    avg_vol = sum(r.annualized_vol for r in regimes) / max(len(regimes), 1)

    market_log_rets = _generate_regime_log_returns(regimes, n, rng)

    # SPY ≈ market (beta ~1.0, low idio noise)
    spy_log_rets: list[float] = []
    for lr in market_log_rets:
        spy_log_rets.append(lr + rng.normal(0, avg_vol * 0.05 * math.sqrt(dt_step)))

    spy_bars = _bars_from_log_returns(
        spy_log_rets, trading_dates,
        initial_price=spy_initial,
        intraday_vol_scale=avg_vol * math.sqrt(dt_step) * 0.4,
        base_volume=80_000_000,
        rng=rng,
    )

    # QQQ ≈ market * 1.1 beta + higher idio noise
    qqq_log_rets: list[float] = []
    for lr in market_log_rets:
        qqq_log_rets.append(
            1.1 * lr + rng.normal(0, avg_vol * 0.08 * math.sqrt(dt_step))
        )

    qqq_bars = _bars_from_log_returns(
        qqq_log_rets, trading_dates,
        initial_price=qqq_initial,
        intraday_vol_scale=avg_vol * math.sqrt(dt_step) * 0.45,
        base_volume=60_000_000,
        rng=rng,
    )

    return spy_bars, qqq_bars, market_log_rets