"""Research probe for non-parking idle-capital ideas.

This tool does not change the core backtest engine.

Supported probes:
  - overnight: invest idle cash near the close and unwind at next open
  - pairs: simple market-neutral pair-trading sanity checks
  - weekday-collateral: intraday weekday sleeve funded by buying power or reserved cash
"""
from __future__ import annotations

import argparse
import datetime as dt
import json
import math
import statistics
from dataclasses import dataclass
from itertools import product
from typing import Any

from apps.backtester.run import BacktestRunner, _build_merged_snapshot_store, load_manifest, resolve_config
from apps.tools.put_spread_overlay_probe import (
    _build_runner,
    _compute_probe_metrics,
    _parse_date,
)
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 PairParams:
    lookback: int = 60
    enter_z: float = 2.0
    exit_z: float = 0.5
    max_hold: int = 20


def _simulate_overnight_overlay(
    curve: list[DailyPortfolioState],
    store: Any,
    *,
    symbol_prefix: str,
    min_cash_ratio: float,
    require_no_open_positions: bool,
) -> tuple[list[DailyPortfolioState], int, float]:
    dates = [state.date for state in curve]
    realized_pnl = 0.0
    adjusted_curve: list[DailyPortfolioState] = []
    traded_nights = 0

    for idx, state in enumerate(curve):
        adjusted_curve.append(
            state.model_copy(update={"equity": state.equity + (realized_pnl if idx > 0 else 0.0)})
        )
        if idx >= len(dates) - 1 or state.equity <= 0:
            continue
        if state.cash_available / state.equity < min_cash_ratio:
            continue
        if require_no_open_positions and state.open_positions:
            continue

        next_date = dates[idx + 1]
        curr_macro = store.get_macro_for_date(state.date)
        next_macro = store.get_macro_for_date(next_date)
        close_px = curr_macro.get(f"{symbol_prefix}_close")
        open_px = next_macro.get(f"{symbol_prefix}_open")
        if not close_px or not open_px or close_px <= 0 or open_px <= 0:
            continue

        overnight_ret = float(open_px) / float(close_px) - 1.0
        realized_pnl += state.cash_available * overnight_ret
        traded_nights += 1

    return adjusted_curve, traded_nights, round(realized_pnl, 2)


def _run_overnight_probe(args: argparse.Namespace) -> list[dict[str, Any]]:
    start_date = _parse_date(args.start)
    end_date = _parse_date(args.end, is_end=True)
    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,
        )
        base_curve = base_runner._equity_curve
        base_metrics = _compute_probe_metrics(base_curve)
        ref_curve = None
        ref_metrics = None
        if args.parking_preset:
            ref_runner, _ = _build_runner(
                config_path,
                start_date,
                end_date,
                args.capital,
                parking_preset=args.parking_preset,
            )
            ref_curve = ref_runner._equity_curve
            ref_metrics = _compute_probe_metrics(ref_curve)

        adjusted_curve, traded_nights, overlay_pnl = _simulate_overnight_overlay(
            base_curve,
            store,
            symbol_prefix=args.symbol.lower(),
            min_cash_ratio=args.min_cash_pct,
            require_no_open_positions=args.require_no_open_positions,
        )
        overlay_metrics = _compute_probe_metrics(adjusted_curve)
        row: dict[str, Any] = {
            "config": config_path,
            "symbol": args.symbol.upper(),
            "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_traded_nights": traded_nights,
            "overlay_realized_pnl": overlay_pnl,
        }
        if ref_metrics is not None:
            row["parking_reference"] = ref_metrics.__dict__
            row["overlay_vs_parking_return_pct"] = round(
                overlay_metrics.total_return_pct - ref_metrics.total_return_pct,
                2,
            )
            row["overlay_vs_parking_dd_pct"] = round(
                overlay_metrics.max_drawdown_pct - ref_metrics.max_drawdown_pct,
                2,
            )
        rows.append(row)
    return rows


def _parse_pair(value: str) -> tuple[str, str]:
    left, right = value.split("/", 1)
    return left.strip().upper(), right.strip().upper()


def _pair_curve(
    *,
    store: Any,
    left_symbol: str,
    right_symbol: str,
    start_date: dt.date,
    end_date: dt.date,
    params: PairParams,
    capital: float,
) -> dict[str, Any] | None:
    dates: list[dt.date] = []
    left_closes: list[float] = []
    right_closes: list[float] = []
    for date in store.all_trading_days():
        if date < start_date or date > end_date:
            continue
        left_bar = store.get_bar(left_symbol, date)
        right_bar = store.get_bar(right_symbol, date)
        if not left_bar or not right_bar:
            continue
        dates.append(date)
        left_closes.append(float(left_bar["close"]))
        right_closes.append(float(right_bar["close"]))

    if len(dates) < params.lookback + 5:
        return None

    equity = capital
    curve: list[DailyPortfolioState] = []
    open_position: dict[str, float | int] | None = None
    prev_left = None
    prev_right = None
    trades = 0

    for idx, date in enumerate(dates):
        left_close = left_closes[idx]
        right_close = right_closes[idx]
        if open_position is not None and prev_left and prev_right:
            left_ret = left_close / prev_left - 1.0
            right_ret = right_close / prev_right - 1.0
            pnl_fraction = (
                float(open_position["sign_left"]) * 0.5 * left_ret
                + float(open_position["sign_right"]) * 0.5 * right_ret
            )
            equity *= 1.0 + pnl_fraction
            open_position["days"] = int(open_position["days"]) + 1

        curve.append(
            DailyPortfolioState(
                date=date,
                equity=equity,
                sizing_equity=equity,
                cash_available=equity,
                gross_exposure=100.0 if open_position else 0.0,
                net_exposure=0.0,
                reserved_risk_budget=0.0,
                unrealized_pnl=0.0,
                realized_pnl=0.0,
                open_positions=["pair"] if open_position else [],
                daily_new_risk_used=0.0,
                peak_equity=max(equity, curve[-1].peak_equity if curve else equity),
                current_drawdown_pct=0.0,
            )
        )

        if idx >= params.lookback:
            window = [
                math.log(left_closes[j] / right_closes[j])
                for j in range(idx - params.lookback, idx)
            ]
            mean_spread = statistics.mean(window)
            spread_std = statistics.pstdev(window)
            if spread_std > 1e-9:
                zscore = (math.log(left_close / right_close) - mean_spread) / spread_std
                if open_position is None:
                    if zscore > params.enter_z:
                        open_position = {"sign_left": -1.0, "sign_right": 1.0, "days": 0}
                        trades += 1
                    elif zscore < -params.enter_z:
                        open_position = {"sign_left": 1.0, "sign_right": -1.0, "days": 0}
                        trades += 1
                elif abs(zscore) < params.exit_z or int(open_position["days"]) >= params.max_hold:
                    open_position = None

        prev_left = left_close
        prev_right = right_close

    return {
        "pair": f"{left_symbol}/{right_symbol}",
        "params": params.__dict__,
        "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),
        "trades": trades,
    }


def _run_pairs_probe(args: argparse.Namespace) -> list[dict[str, Any]]:
    start_date = _parse_date(args.start)
    end_date = _parse_date(args.end, is_end=True)
    base_runner, store = _build_runner(
        args.config,
        start_date,
        end_date,
        args.capital,
        parking_preset=None,
    )
    del base_runner

    pairs = [_parse_pair(value) for value in args.pair]
    rows: list[dict[str, Any]] = []
    if args.grid:
        grid = list(product([40, 60], [1.5, 2.0], [0.25, 0.5], [10, 20]))
        for left_symbol, right_symbol in pairs:
            best: dict[str, Any] | None = None
            for lookback, enter_z, exit_z, max_hold in grid:
                row = _pair_curve(
                    store=store,
                    left_symbol=left_symbol,
                    right_symbol=right_symbol,
                    start_date=start_date,
                    end_date=end_date,
                    params=PairParams(
                        lookback=lookback,
                        enter_z=enter_z,
                        exit_z=exit_z,
                        max_hold=max_hold,
                    ),
                    capital=args.capital,
                )
                if row is None:
                    continue
                if best is None or (
                    row["total_return_pct"],
                    -row["max_drawdown_pct"],
                    row["sharpe_ratio"],
                ) > (
                    best["total_return_pct"],
                    -best["max_drawdown_pct"],
                    best["sharpe_ratio"],
                ):
                    best = row
            if best is not None:
                rows.append(best)
        return rows

    params = PairParams(
        lookback=args.lookback,
        enter_z=args.enter_z,
        exit_z=args.exit_z,
        max_hold=args.max_hold,
    )
    for left_symbol, right_symbol in pairs:
        row = _pair_curve(
            store=store,
            left_symbol=left_symbol,
            right_symbol=right_symbol,
            start_date=start_date,
            end_date=end_date,
            params=params,
            capital=args.capital,
        )
        if row is not None:
            rows.append(row)
    return rows


def _run_micro_event_probe(args: argparse.Namespace) -> list[dict[str, Any]]:
    start_date = _parse_date(args.start)
    end_date = _parse_date(args.end, is_end=True)
    manifest = load_manifest(args.config)
    base_config = resolve_config(manifest)
    store = _build_merged_snapshot_store(
        manifest,
        base_config,
        snapshot_dir_override=None,
    ).slice_by_date_range(start_date, end_date)

    def _run_config(config: Any) -> tuple[ProbeMetrics, Any]:
        runner = BacktestRunner(
            manifest=manifest,
            config=config,
            store=store,
            initial_equity=args.capital,
            split_name="idle_micro_event_probe",
        )
        runner.run(output_root=None)
        return _compute_probe_metrics(runner._equity_curve), runner

    base_metrics, base_runner = _run_config(base_config)
    row: dict[str, Any] = {
        "config": args.config,
        "mode": args.mode,
        "engine_ids": list(args.engine_id),
        "base": base_metrics.__dict__,
    }

    variant_config = resolve_config(manifest)
    selected_ids = set(args.engine_id)
    if args.mode == "clone":
        lookup = {engine.engine_id: engine for engine in variant_config.strategy_engines}
        new_engines = list(variant_config.strategy_engines)
        for engine_id in args.engine_id:
            source = lookup[engine_id]
            new_engines.append(
                source.model_copy(
                    update={
                        "engine_id": f"{engine_id}_micro",
                        "selection_priority": args.selection_priority,
                        "capital_bucket_id": args.bucket_id,
                        "capital_bucket_allocation_pct": args.bucket_allocation_pct,
                        "engine_risk_budget_pct": min(float(source.engine_risk_budget_pct), 0.20),
                        "per_trade_risk_pct_override": (
                            args.per_trade_risk_pct
                            if args.per_trade_risk_pct is not None
                            else source.per_trade_risk_pct_override
                        ),
                        "residual_reserve_selected": False,
                        "recycle_on_cash_block": False,
                    }
                )
            )
        variant_config.strategy_engines = new_engines
        tracked_engine_ids = {f"{engine_id}_micro" for engine_id in args.engine_id}
    else:
        new_engines = []
        for engine in variant_config.strategy_engines:
            if engine.engine_id in selected_ids:
                new_engines.append(
                    engine.model_copy(
                        update={
                            "capital_bucket_id": args.bucket_id,
                            "capital_bucket_allocation_pct": args.bucket_allocation_pct,
                            "per_trade_risk_pct_override": (
                                args.per_trade_risk_pct
                                if args.per_trade_risk_pct is not None
                                else engine.per_trade_risk_pct_override
                            ),
                        }
                    )
                )
            else:
                new_engines.append(engine)
        variant_config.strategy_engines = new_engines
        tracked_engine_ids = selected_ids

    variant_metrics, variant_runner = _run_config(variant_config)
    row["variant"] = variant_metrics.__dict__
    row["variant_delta_return_pct"] = round(
        variant_metrics.total_return_pct - base_metrics.total_return_pct,
        2,
    )
    row["variant_delta_dd_pct"] = round(
        variant_metrics.max_drawdown_pct - base_metrics.max_drawdown_pct,
        2,
    )
    row["variant_trade_count"] = len(variant_runner._closed_trades)
    row["variant_tracked_trade_count"] = sum(
        1 for trade in variant_runner._closed_trades if trade.engine_id in tracked_engine_ids
    )
    row["variant_tracked_pnl"] = round(
        sum(float(trade.net_pnl) for trade in variant_runner._closed_trades if trade.engine_id in tracked_engine_ids),
        2,
    )

    if args.parking_preset:
        ref_runner, _ = _build_runner(
            args.config,
            start_date,
            end_date,
            args.capital,
            parking_preset=args.parking_preset,
        )
        ref_metrics = _compute_probe_metrics(ref_runner._equity_curve)
        row["parking_reference"] = ref_metrics.__dict__
        row["variant_vs_parking_return_pct"] = round(
            variant_metrics.total_return_pct - ref_metrics.total_return_pct,
            2,
        )
        row["variant_vs_parking_dd_pct"] = round(
            variant_metrics.max_drawdown_pct - ref_metrics.max_drawdown_pct,
            2,
        )

    return [row]


def _weekday_filter_allows(store: Any, date: dt.date, filter_name: str) -> bool:
    macro = store.get_macro_for_date(date)
    if filter_name == "none":
        return True
    if filter_name == "trend50":
        close = macro.get("spy_close")
        sma = macro.get("spy_sma_50")
        return close is not None and sma is not None and float(close) >= float(sma)
    if filter_name == "mom20":
        mom = macro.get("spy_mom_20")
        return mom is not None and float(mom) > 0.0
    if filter_name == "vol15":
        vol = macro.get("spy_vol_20")
        return vol is not None and float(vol) < 0.15
    if filter_name == "trend50_vol15":
        close = macro.get("spy_close")
        sma = macro.get("spy_sma_50")
        vol = macro.get("spy_vol_20")
        return (
            close is not None
            and sma is not None
            and vol is not None
            and float(close) >= float(sma)
            and float(vol) < 0.15
        )
    raise ValueError(f"Unsupported weekday filter: {filter_name}")


def _simulate_weekday_collateral_overlay(
    curve: list[DailyPortfolioState],
    store: Any,
    *,
    symbol: str,
    weekday: int,
    filter_name: str,
    capital_source: str,
    reserve_pct: float,
) -> tuple[list[DailyPortfolioState], int, float, dict[int, float]]:
    realized_pnl = 0.0
    adjusted_curve: list[DailyPortfolioState] = []
    traded_days = 0
    pnl_by_year: dict[int, float] = {}

    for idx, state in enumerate(curve):
        adjusted_curve.append(
            state.model_copy(update={"equity": state.equity + (realized_pnl if idx > 0 else 0.0)})
        )
        if state.date.weekday() != weekday:
            continue
        if not _weekday_filter_allows(store, state.date, filter_name):
            continue
        bar = store.get_bar(symbol, state.date)
        if not bar:
            continue
        open_px = bar.get("open")
        close_px = bar.get("close")
        if not open_px or not close_px or open_px <= 0 or close_px <= 0:
            continue
        if capital_source == "cash_available":
            trade_capital = float(state.cash_available)
        else:
            trade_capital = float(state.equity) * reserve_pct
        if trade_capital <= 0:
            continue
        intraday_ret = float(close_px) / float(open_px) - 1.0
        day_pnl = trade_capital * intraday_ret
        realized_pnl += day_pnl
        traded_days += 1
        pnl_by_year[state.date.year] = pnl_by_year.get(state.date.year, 0.0) + day_pnl

    return adjusted_curve, traded_days, round(realized_pnl, 2), pnl_by_year


def _run_weekday_collateral_probe(args: argparse.Namespace) -> list[dict[str, Any]]:
    start_date = _parse_date(args.start)
    end_date = _parse_date(args.end, is_end=True)
    rows: list[dict[str, Any]] = []
    for config_path in args.config:
        if args.capital_source == "reserve_pct":
            manifest = load_manifest(config_path)
            config = resolve_config(manifest)
            if args.parking_preset:
                config.risk.cash_parking_enabled = True
                config.risk.cash_parking_preset = args.parking_preset
                config.risk.apply_parking_preset()
            config.risk.cash_parking_reserve_pct = args.reserve_pct
            store = _build_merged_snapshot_store(
                manifest,
                config,
                snapshot_dir_override=None,
            ).slice_by_date_range(start_date, end_date)
            base_runner = BacktestRunner(
                manifest=manifest,
                config=config,
                store=store,
                initial_equity=args.capital,
                split_name="weekday_collateral_probe",
            )
            base_runner.run(output_root=None)
        else:
            base_runner, store = _build_runner(
                config_path,
                start_date,
                end_date,
                args.capital,
                parking_preset=args.parking_preset,
            )
        base_curve = base_runner._equity_curve
        base_metrics = _compute_probe_metrics(base_curve)
        (
            adjusted_curve,
            traded_days,
            overlay_pnl,
            pnl_by_year,
        ) = _simulate_weekday_collateral_overlay(
            base_curve,
            store,
            symbol=args.symbol.upper(),
            weekday=args.weekday,
            filter_name=args.filter,
            capital_source=args.capital_source,
            reserve_pct=args.reserve_pct,
        )
        overlay_metrics = _compute_probe_metrics(adjusted_curve)
        row: dict[str, Any] = {
            "config": config_path,
            "parking_preset": args.parking_preset,
            "symbol": args.symbol.upper(),
            "weekday": args.weekday,
            "filter": args.filter,
            "capital_source": args.capital_source,
            "reserve_pct": args.reserve_pct,
            "base": 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_traded_days": traded_days,
            "overlay_realized_pnl": overlay_pnl,
            "overlay_pnl_by_year": {str(year): round(value, 2) for year, value in sorted(pnl_by_year.items())},
        }
        rows.append(row)
    return rows


def main() -> None:
    parser = argparse.ArgumentParser(description="Probe non-parking idle-capital ideas")
    subparsers = parser.add_subparsers(dest="command", required=True)

    overnight_parser = subparsers.add_parser("overnight", help="Idle-cash overnight drift overlay")
    overnight_parser.add_argument("--config", action="append", required=True)
    overnight_parser.add_argument("--start", required=True)
    overnight_parser.add_argument("--end", required=True)
    overnight_parser.add_argument("--capital", type=float, default=10_000.0)
    overnight_parser.add_argument("--symbol", choices=("SPY", "QQQ"), default="SPY")
    overnight_parser.add_argument("--parking-preset", default=None)
    overnight_parser.add_argument("--min-cash-pct", type=float, default=0.0)
    overnight_parser.add_argument("--require-no-open-positions", action="store_true")
    overnight_parser.add_argument("--json", action="store_true")

    pairs_parser = subparsers.add_parser("pairs", help="Simple mean-reversion pair sanity check")
    pairs_parser.add_argument("--config", required=True)
    pairs_parser.add_argument("--start", required=True)
    pairs_parser.add_argument("--end", required=True)
    pairs_parser.add_argument("--capital", type=float, default=10_000.0)
    pairs_parser.add_argument(
        "--pair",
        action="append",
        required=True,
        help="Pair in SYMBOL_A/SYMBOL_B form. Repeatable.",
    )
    pairs_parser.add_argument("--lookback", type=int, default=60)
    pairs_parser.add_argument("--enter-z", type=float, default=2.0)
    pairs_parser.add_argument("--exit-z", type=float, default=0.5)
    pairs_parser.add_argument("--max-hold", type=int, default=20)
    pairs_parser.add_argument("--grid", action="store_true")
    pairs_parser.add_argument("--json", action="store_true")

    micro_parser = subparsers.add_parser(
        "micro-event",
        help="Test small capital-bucket event sleeves using existing strategy engines",
    )
    micro_parser.add_argument("--config", required=True)
    micro_parser.add_argument("--start", required=True)
    micro_parser.add_argument("--end", required=True)
    micro_parser.add_argument("--capital", type=float, default=10_000.0)
    micro_parser.add_argument(
        "--engine-id",
        action="append",
        required=True,
        help="Existing engine_id to clone or bucket. Repeatable.",
    )
    micro_parser.add_argument("--mode", choices=("clone", "bucket"), default="clone")
    micro_parser.add_argument("--bucket-id", default="micro_event")
    micro_parser.add_argument("--bucket-allocation-pct", type=float, required=True)
    micro_parser.add_argument("--per-trade-risk-pct", type=float, default=None)
    micro_parser.add_argument("--selection-priority", type=int, default=-1)
    micro_parser.add_argument("--parking-preset", default=None)
    micro_parser.add_argument("--json", action="store_true")

    weekday_parser = subparsers.add_parser(
        "weekday-collateral",
        help="Intraday weekday sleeve funded by buying power or reserved cash",
    )
    weekday_parser.add_argument("--config", action="append", required=True)
    weekday_parser.add_argument("--start", required=True)
    weekday_parser.add_argument("--end", required=True)
    weekday_parser.add_argument("--capital", type=float, default=10_000.0)
    weekday_parser.add_argument("--parking-preset", default="qqqm_low_dd")
    weekday_parser.add_argument("--symbol", choices=("SPY", "QQQ"), default="SPY")
    weekday_parser.add_argument("--weekday", type=int, choices=range(5), default=0)
    weekday_parser.add_argument(
        "--filter",
        choices=("none", "trend50", "mom20", "vol15", "trend50_vol15"),
        default="none",
    )
    weekday_parser.add_argument(
        "--capital-source",
        choices=("cash_available", "reserve_pct"),
        default="cash_available",
    )
    weekday_parser.add_argument("--reserve-pct", type=float, default=0.0)
    weekday_parser.add_argument("--json", action="store_true")

    args = parser.parse_args()
    configure_logging("WARNING")

    if args.command == "overnight":
        rows = _run_overnight_probe(args)
    elif args.command == "pairs":
        rows = _run_pairs_probe(args)
    elif args.command == "weekday-collateral":
        rows = _run_weekday_collateral_probe(args)
    else:
        rows = _run_micro_event_probe(args)

    if args.json:
        print(json.dumps(rows, indent=2))
    else:
        for row in rows:
            print(row)


if __name__ == "__main__":
    main()
