"""Assumption-driven QQQ put-spread overlay probe for idle cash.

This is NOT a true options backtest.

It re-runs a strategy with cash parking disabled, then layers a conservative
QQQ short put spread overlay on top of the resulting daily equity curve.
Option prices are approximated with Black-Scholes using QQQ realized vol as an
IV proxy, plus conservative fill haircuts. Use this only for relative research.

Example:
    python -m apps.tools.put_spread_overlay_probe \
      --config configs/experiments/return_max_long_v11.49.json \
      --start 2022 --end 2026 \
      --parking-preset qqqm_low_dd \
      --period-start y2026=2026-01-01
"""
from __future__ import annotations

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

from apps.backtester.run import (
    BacktestRunner,
    _build_merged_snapshot_store,
    load_manifest,
    resolve_config,
)
from libs.backtest.domain import DailyPortfolioState
from libs.backtest.metrics import (
    compute_max_drawdown_pct,
    compute_sharpe_ratio,
    compute_total_return_pct,
)
from libs.common.logging import configure_logging


@dataclass(frozen=True)
class OverlayParams:
    short_otm_pct: float = 0.06
    spread_width_pct: float = 0.025
    dte_trading_days: int = 15
    min_cash_available_pct: float = 0.70
    max_margin_pct_of_equity: float = 0.15
    iv_proxy_floor: float = 0.18
    iv_proxy_multiplier: float = 1.20
    risk_free_rate: float = 0.04
    entry_credit_haircut_pct: float = 0.15
    exit_debit_markup_pct: float = 0.05
    commission_per_contract_leg: float = 1.00
    strike_rounding: float = 1.0
    require_no_open_positions: bool = True
    close_on_new_risk: bool = True
    close_on_open_positions: bool = True
    close_on_cash_recall_pct: float = 0.50
    idle_cash_annual_yield: float = 0.0


@dataclass
class OverlayTrade:
    entry_date: dt.date
    exit_date: dt.date
    exit_reason: str
    contracts: int
    short_strike: float
    long_strike: float
    entry_credit_per_spread: float
    exit_debit_per_spread: float
    pnl: float


@dataclass
class OverlaySummary:
    adjusted_curve: list[DailyPortfolioState]
    trades: list[OverlayTrade]
    total_realized_pnl: float
    open_unrealized_pnl: float
    idle_cash_carry_pnl: float
    max_margin_used: float
    wins: int
    losses: int
    forced_closes: int


@dataclass
class ProbeMetrics:
    total_return_pct: float
    max_drawdown_pct: float
    sharpe_ratio: float
    final_equity: float


def _parse_date(value: str, *, is_end: bool = False) -> dt.date:
    parts = value.split("-")
    if len(parts) == 1 and len(value) == 4 and value.isdigit():
        year = int(value)
        return dt.date(year, 12, 31) if is_end else dt.date(year, 1, 1)
    if len(parts) == 2 and all(part.isdigit() for part in parts):
        year = int(parts[0])
        month = int(parts[1])
        if is_end:
            next_month = dt.date(year + (month // 12), (month % 12) + 1, 1)
            return next_month - dt.timedelta(days=1)
        return dt.date(year, month, 1)
    return dt.date.fromisoformat(value)


def _parse_period_start(value: str) -> tuple[str, dt.date]:
    if "=" in value:
        label, raw_date = value.split("=", 1)
    else:
        label, raw_date = value, value
    return label, _parse_date(raw_date)


def _norm_cdf(x: float) -> float:
    return 0.5 * (1.0 + math.erf(x / math.sqrt(2.0)))


def _black_scholes_put_price(
    spot: float,
    strike: float,
    years_to_expiry: float,
    sigma: float,
    risk_free_rate: float,
) -> float:
    if spot <= 0 or strike <= 0:
        return 0.0
    if years_to_expiry <= 0 or sigma <= 0:
        return max(strike - spot, 0.0)
    sqrt_t = math.sqrt(years_to_expiry)
    d1 = (
        math.log(spot / strike) + (risk_free_rate + 0.5 * sigma * sigma) * years_to_expiry
    ) / (sigma * sqrt_t)
    d2 = d1 - sigma * sqrt_t
    discounted_strike = strike * math.exp(-risk_free_rate * years_to_expiry)
    return discounted_strike * _norm_cdf(-d2) - spot * _norm_cdf(-d1)


def _round_to_increment(value: float, increment: float) -> float:
    increment = max(increment, 1e-9)
    return round(value / increment) * increment


def _compute_probe_metrics(curve: list[DailyPortfolioState]) -> ProbeMetrics:
    return ProbeMetrics(
        total_return_pct=round(compute_total_return_pct(curve) or 0.0, 2),
        max_drawdown_pct=round(compute_max_drawdown_pct(curve) or 0.0, 2),
        sharpe_ratio=round(compute_sharpe_ratio(curve) or 0.0, 3),
        final_equity=round(curve[-1].equity if curve else 0.0, 2),
    )


def _period_stats(curve: list[DailyPortfolioState], start_date: dt.date) -> dict[str, Any] | None:
    points = [state for state in curve if state.date >= start_date]
    if not points:
        return None
    return {
        "start_date": points[0].date.isoformat(),
        "end_date": points[-1].date.isoformat(),
        "return_pct": round(compute_total_return_pct(points) or 0.0, 2),
        "max_dd_pct": round(compute_max_drawdown_pct(points) or 0.0, 2),
        "sharpe_ratio": round(compute_sharpe_ratio(points) or 0.0, 3),
    }


def _build_runner(
    config_path: str,
    start_date: dt.date,
    end_date: dt.date,
    capital: float,
    *,
    parking_preset: str | None,
) -> tuple[BacktestRunner, Any]:
    manifest = load_manifest(config_path)
    config = resolve_config(manifest)
    config.risk.cash_parking_enabled = parking_preset is not None
    config.risk.cash_parking_preset = parking_preset
    if parking_preset is not None:
        config.risk.apply_parking_preset()
    store = _build_merged_snapshot_store(
        manifest,
        config,
        snapshot_dir_override=None,
    ).slice_by_date_range(start_date, end_date)
    runner = BacktestRunner(
        manifest=manifest,
        config=config,
        store=store,
        initial_equity=capital,
        split_name="put_spread_overlay_probe",
    )
    runner.run(output_root=None)
    return runner, store


def _should_open_overlay(
    state: DailyPortfolioState,
    params: OverlayParams,
) -> bool:
    if state.equity <= 0:
        return False
    cash_ratio = state.cash_available / state.equity
    if cash_ratio < params.min_cash_available_pct:
        return False
    if params.require_no_open_positions and state.open_positions:
        return False
    if state.daily_new_risk_used > 0:
        return False
    return True


def _should_force_close(
    state: DailyPortfolioState,
    params: OverlayParams,
) -> str | None:
    if params.close_on_new_risk and state.daily_new_risk_used > 0:
        return "new_risk"
    if params.close_on_open_positions and state.open_positions:
        return "position_opened"
    if state.equity > 0:
        cash_ratio = state.cash_available / state.equity
        if cash_ratio < params.close_on_cash_recall_pct:
            return "cash_recall"
    return None


def _spread_liability_per_spread(
    spot: float,
    short_strike: float,
    long_strike: float,
    years_to_expiry: float,
    sigma: float,
    risk_free_rate: float,
    exit_debit_markup_pct: float,
) -> float:
    short_put = _black_scholes_put_price(spot, short_strike, years_to_expiry, sigma, risk_free_rate)
    long_put = _black_scholes_put_price(spot, long_strike, years_to_expiry, sigma, risk_free_rate)
    mid_debit = max(short_put - long_put, 0.0) * 100.0
    return mid_debit * (1.0 + exit_debit_markup_pct)


def simulate_put_spread_overlay(
    curve: list[DailyPortfolioState],
    store: Any,
    params: OverlayParams,
) -> OverlaySummary:
    adjusted_curve: list[DailyPortfolioState] = []
    trades: list[OverlayTrade] = []
    total_realized_pnl = 0.0
    open_unrealized_pnl = 0.0
    idle_cash_carry_pnl = 0.0
    max_margin_used = 0.0
    wins = 0
    losses = 0
    forced_closes = 0
    prev_free_cash_for_carry = 0.0
    trading_days = [state.date for state in curve]
    date_to_index = {date: idx for idx, date in enumerate(trading_days)}
    open_position: dict[str, Any] | None = None
    daily_carry_rate = (
        (1.0 + params.idle_cash_annual_yield) ** (1.0 / 252.0) - 1.0
        if params.idle_cash_annual_yield > 0
        else 0.0
    )

    for state in curve:
        date = state.date
        macro = store.get_macro_for_date(date)
        spot = macro.get("qqq_close")
        sigma = max(
            float(macro.get("qqq_vol_20") or 0.0) * params.iv_proxy_multiplier,
            params.iv_proxy_floor,
        )

        if open_position is not None:
            current_idx = date_to_index[date]
            expiry_idx = date_to_index[open_position["expiry_date"]]
            remaining_days = max(expiry_idx - current_idx, 0)
            years_to_expiry = remaining_days / 252.0
            if spot is None or spot <= 0:
                liability = open_position["last_liability_per_spread"]
            elif remaining_days == 0:
                intrinsic = max(open_position["short_strike"] - spot, 0.0) - max(
                    open_position["long_strike"] - spot, 0.0
                )
                liability = max(intrinsic, 0.0) * 100.0
            else:
                liability = _spread_liability_per_spread(
                    spot=float(spot),
                    short_strike=open_position["short_strike"],
                    long_strike=open_position["long_strike"],
                    years_to_expiry=years_to_expiry,
                    sigma=sigma,
                    risk_free_rate=params.risk_free_rate,
                    exit_debit_markup_pct=params.exit_debit_markup_pct,
                )
            open_position["last_liability_per_spread"] = liability
            open_unrealized_pnl = (
                open_position["contracts"]
                * (open_position["net_credit_per_spread"] - liability)
            )

            exit_reason: str | None = None
            if remaining_days == 0:
                exit_reason = "expiry"
            else:
                exit_reason = _should_force_close(state, params)

            if exit_reason is not None:
                exit_debit_per_spread = liability
                if exit_reason == "expiry":
                    close_fee = 0.0
                else:
                    close_fee = params.commission_per_contract_leg * 2.0
                    forced_closes += 1
                realized_trade_pnl = open_position["contracts"] * (
                    open_position["net_credit_per_spread"] - exit_debit_per_spread - close_fee
                )
                total_realized_pnl += realized_trade_pnl
                trade = OverlayTrade(
                    entry_date=open_position["entry_date"],
                    exit_date=date,
                    exit_reason=exit_reason,
                    contracts=open_position["contracts"],
                    short_strike=open_position["short_strike"],
                    long_strike=open_position["long_strike"],
                    entry_credit_per_spread=round(open_position["net_credit_per_spread"], 2),
                    exit_debit_per_spread=round(exit_debit_per_spread + close_fee, 2),
                    pnl=round(realized_trade_pnl, 2),
                )
                trades.append(trade)
                if realized_trade_pnl > 0:
                    wins += 1
                else:
                    losses += 1
                open_position = None
                open_unrealized_pnl = 0.0

        if open_position is None and _should_open_overlay(state, params):
            current_idx = date_to_index[date]
            expiry_idx = current_idx + params.dte_trading_days
            if expiry_idx < len(trading_days) and spot is not None and spot > 0:
                strike_rounding = params.strike_rounding
                short_strike = _round_to_increment(float(spot) * (1.0 - params.short_otm_pct), strike_rounding)
                width_points = max(
                    strike_rounding,
                    _round_to_increment(float(spot) * params.spread_width_pct, strike_rounding),
                )
                long_strike = max(strike_rounding, short_strike - width_points)
                years_to_expiry = params.dte_trading_days / 252.0
                entry_liability = _spread_liability_per_spread(
                    spot=float(spot),
                    short_strike=short_strike,
                    long_strike=long_strike,
                    years_to_expiry=years_to_expiry,
                    sigma=sigma,
                    risk_free_rate=params.risk_free_rate,
                    exit_debit_markup_pct=0.0,
                )
                gross_credit = max(entry_liability, 0.0) * (1.0 - params.entry_credit_haircut_pct)
                entry_fee = params.commission_per_contract_leg * 2.0
                net_credit_per_spread = max(gross_credit - entry_fee, 0.0)
                max_loss_per_spread = max((short_strike - long_strike) * 100.0 - net_credit_per_spread, 1.0)
                margin_budget = min(
                    state.cash_available,
                    state.equity * params.max_margin_pct_of_equity,
                )
                contracts = int(margin_budget // max_loss_per_spread)
                if contracts > 0:
                    max_margin_used = max(max_margin_used, contracts * max_loss_per_spread)
                    open_position = {
                        "entry_date": date,
                        "expiry_date": trading_days[expiry_idx],
                        "contracts": contracts,
                        "short_strike": short_strike,
                        "long_strike": long_strike,
                        "net_credit_per_spread": net_credit_per_spread,
                        "max_loss_per_spread": max_loss_per_spread,
                        "last_liability_per_spread": entry_liability,
                    }
                    open_unrealized_pnl = contracts * (net_credit_per_spread - entry_liability)

        current_margin_reserved = 0.0
        if open_position is not None:
            current_margin_reserved = (
                open_position["contracts"] * open_position["max_loss_per_spread"]
            )
        if daily_carry_rate > 0.0 and prev_free_cash_for_carry > 0.0:
            idle_cash_carry_pnl += prev_free_cash_for_carry * daily_carry_rate
        adjusted_equity = (
            state.equity
            + total_realized_pnl
            + open_unrealized_pnl
            + idle_cash_carry_pnl
        )
        adjusted_curve.append(state.model_copy(update={"equity": adjusted_equity}))
        prev_free_cash_for_carry = max(0.0, state.cash_available - current_margin_reserved)

    return OverlaySummary(
        adjusted_curve=adjusted_curve,
        trades=trades,
        total_realized_pnl=round(total_realized_pnl, 2),
        open_unrealized_pnl=round(open_unrealized_pnl, 2),
        idle_cash_carry_pnl=round(idle_cash_carry_pnl, 2),
        max_margin_used=round(max_margin_used, 2),
        wins=wins,
        losses=losses,
        forced_closes=forced_closes,
    )


def _build_row(
    config_path: str,
    base_curve: list[DailyPortfolioState],
    overlay: OverlaySummary,
    reference_curve: list[DailyPortfolioState] | None,
    period_starts: list[tuple[str, dt.date]],
) -> dict[str, Any]:
    base_metrics = _compute_probe_metrics(base_curve)
    overlay_metrics = _compute_probe_metrics(overlay.adjusted_curve)
    reference_metrics = _compute_probe_metrics(reference_curve) if reference_curve else None
    row: dict[str, Any] = {
        "config": config_path,
        "base_no_parking": base_metrics.__dict__,
        "overlay": overlay_metrics.__dict__,
        "overlay_delta_return_pct": round(
            overlay_metrics.total_return_pct - base_metrics.total_return_pct,
            2,
        ),
        "overlay_delta_dd_pct": round(
            overlay_metrics.max_drawdown_pct - base_metrics.max_drawdown_pct,
            2,
        ),
        "overlay_trades": len(overlay.trades),
        "overlay_wins": overlay.wins,
        "overlay_losses": overlay.losses,
        "overlay_forced_closes": overlay.forced_closes,
        "overlay_realized_pnl": overlay.total_realized_pnl,
        "overlay_open_unrealized_pnl": overlay.open_unrealized_pnl,
        "overlay_idle_cash_carry_pnl": overlay.idle_cash_carry_pnl,
        "overlay_max_margin_used": overlay.max_margin_used,
    }
    if reference_metrics is not None:
        row["parking_reference"] = reference_metrics.__dict__
        row["overlay_vs_parking_return_pct"] = round(
            overlay_metrics.total_return_pct - reference_metrics.total_return_pct,
            2,
        )
        row["overlay_vs_parking_dd_pct"] = round(
            overlay_metrics.max_drawdown_pct - reference_metrics.max_drawdown_pct,
            2,
        )
    for label, period_start in period_starts:
        row[f"{label}_base"] = _period_stats(base_curve, period_start)
        row[f"{label}_overlay"] = _period_stats(overlay.adjusted_curve, period_start)
        if reference_curve:
            row[f"{label}_parking"] = _period_stats(reference_curve, period_start)
    return row


def main() -> None:
    parser = argparse.ArgumentParser(description="Probe assumption-driven QQQ put spread overlay")
    parser.add_argument(
        "--config",
        action="append",
        required=True,
        help="Experiment manifest JSON path (repeatable)",
    )
    parser.add_argument("--start", required=True, help="Start date (YYYY, YYYY-MM, YYYY-MM-DD)")
    parser.add_argument("--end", required=True, help="End date (YYYY, YYYY-MM, YYYY-MM-DD)")
    parser.add_argument(
        "--period-start",
        action="append",
        default=[],
        help="Sub-period start. Format: label=YYYY-MM-DD (repeatable).",
    )
    parser.add_argument(
        "--parking-preset",
        default=None,
        help="Optional parking preset for reference comparison (e.g. qqqm_low_dd)",
    )
    parser.add_argument("--capital", type=float, default=10_000.0)
    parser.add_argument("--short-otm-pct", type=float, default=0.06)
    parser.add_argument("--spread-width-pct", type=float, default=0.025)
    parser.add_argument("--dte", type=int, default=15, help="Expiry in trading days")
    parser.add_argument("--min-cash-pct", type=float, default=0.70)
    parser.add_argument("--max-margin-pct", type=float, default=0.15)
    parser.add_argument("--iv-floor", type=float, default=0.18)
    parser.add_argument("--iv-mult", type=float, default=1.20)
    parser.add_argument("--entry-haircut-pct", type=float, default=0.15)
    parser.add_argument("--exit-markup-pct", type=float, default=0.05)
    parser.add_argument("--commission", type=float, default=1.0, help="Per contract leg commission")
    parser.add_argument("--strike-rounding", type=float, default=1.0)
    parser.add_argument("--allow-with-open-positions", action="store_true")
    parser.add_argument("--keep-through-new-risk", action="store_true")
    parser.add_argument("--keep-through-position-open", action="store_true")
    parser.add_argument("--cash-recall-pct", type=float, default=0.50)
    parser.add_argument(
        "--idle-cash-yield",
        type=float,
        default=0.0,
        help="Approx annual yield on unallocated idle cash (e.g. 0.05 for SGOV-like carry)",
    )
    parser.add_argument("--json", action="store_true")
    args = parser.parse_args()

    configure_logging("WARNING")

    start_date = _parse_date(args.start)
    end_date = _parse_date(args.end, is_end=True)
    period_starts = [_parse_period_start(value) for value in args.period_start]
    params = OverlayParams(
        short_otm_pct=args.short_otm_pct,
        spread_width_pct=args.spread_width_pct,
        dte_trading_days=args.dte,
        min_cash_available_pct=args.min_cash_pct,
        max_margin_pct_of_equity=args.max_margin_pct,
        iv_proxy_floor=args.iv_floor,
        iv_proxy_multiplier=args.iv_mult,
        entry_credit_haircut_pct=args.entry_haircut_pct,
        exit_debit_markup_pct=args.exit_markup_pct,
        commission_per_contract_leg=args.commission,
        strike_rounding=args.strike_rounding,
        require_no_open_positions=not args.allow_with_open_positions,
        close_on_new_risk=not args.keep_through_new_risk,
        close_on_open_positions=not args.keep_through_position_open,
        close_on_cash_recall_pct=args.cash_recall_pct,
        idle_cash_annual_yield=args.idle_cash_yield,
    )

    rows: list[dict[str, Any]] = []
    for config_path in args.config:
        base_runner, store = _build_runner(
            config_path,
            start_date,
            end_date,
            args.capital,
            parking_preset=None,
        )
        overlay = simulate_put_spread_overlay(base_runner._equity_curve, store, params)
        reference_curve: list[DailyPortfolioState] | None = None
        if args.parking_preset:
            reference_runner, _ = _build_runner(
                config_path,
                start_date,
                end_date,
                args.capital,
                parking_preset=args.parking_preset,
            )
            reference_curve = reference_runner._equity_curve
        rows.append(
            _build_row(
                config_path=config_path,
                base_curve=base_runner._equity_curve,
                overlay=overlay,
                reference_curve=reference_curve,
                period_starts=period_starts,
            )
        )

    if args.json:
        print(json.dumps(rows, ensure_ascii=False, indent=2))
        return

    print("NOTE: assumption-driven overlay only; not a real options backtest.")
    for row in rows:
        base = row["base_no_parking"]
        overlay = row["overlay"]
        print()
        print(row["config"])
        print(
            f"  base(no parking): return {base['total_return_pct']:.2f}% | "
            f"dd {base['max_drawdown_pct']:.2f}% | sharpe {base['sharpe_ratio']:.3f}"
        )
        print(
            f"  overlay:          return {overlay['total_return_pct']:.2f}% | "
            f"dd {overlay['max_drawdown_pct']:.2f}% | sharpe {overlay['sharpe_ratio']:.3f}"
        )
        print(
            f"  delta:            return {row['overlay_delta_return_pct']:+.2f}%p | "
            f"dd {row['overlay_delta_dd_pct']:+.2f}%p"
        )
        print(
            f"  trades {row['overlay_trades']} | wins {row['overlay_wins']} | "
            f"losses {row['overlay_losses']} | forced {row['overlay_forced_closes']} | "
            f"realized ${row['overlay_realized_pnl']:.2f} | "
            f"carry ${row['overlay_idle_cash_carry_pnl']:.2f} | "
            f"max_margin ${row['overlay_max_margin_used']:.2f}"
        )
        if "parking_reference" in row:
            parking = row["parking_reference"]
            print(
                f"  parking({args.parking_preset}): return {parking['total_return_pct']:.2f}% | "
                f"dd {parking['max_drawdown_pct']:.2f}% | sharpe {parking['sharpe_ratio']:.3f}"
            )
            print(
                f"  overlay vs parking: return {row['overlay_vs_parking_return_pct']:+.2f}%p | "
                f"dd {row['overlay_vs_parking_dd_pct']:+.2f}%p"
            )
        for label, _ in period_starts:
            base_period = row.get(f"{label}_base")
            overlay_period = row.get(f"{label}_overlay")
            if base_period and overlay_period:
                print(
                    f"  {label}: base {base_period['return_pct']:.2f}% / {base_period['max_dd_pct']:.2f}%dd | "
                    f"overlay {overlay_period['return_pct']:.2f}% / {overlay_period['max_dd_pct']:.2f}%dd"
                )
            parking_period = row.get(f"{label}_parking")
            if parking_period:
                print(
                    f"  {label}: parking {parking_period['return_pct']:.2f}% / "
                    f"{parking_period['max_dd_pct']:.2f}%dd"
                )


if __name__ == "__main__":
    main()
