#!/usr/bin/env python3
from __future__ import annotations

import argparse
import math
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import pandas as pd

CONTRACT_MULT = 100


@dataclass
class LegSpec:
    right: str  # "C" or "P"
    target_delta: float
    dte: int
    qty: int  # + for long, - for short


@dataclass
class Trade:
    ticker: str
    strategy: str
    entry_date: pd.Timestamp
    exit_date: pd.Timestamp
    pnl_usd: float
    spot_entry: float = 0.0  # approximate underlying price at entry (ATM call strike proxy)


def _session_dates_series(quote_datetime: pd.Series) -> pd.Series:
    qd = pd.to_datetime(quote_datetime, utc=False)
    if qd.dt.tz is None:
        qd = qd.dt.tz_localize("America/New_York", ambiguous="NaT", nonexistent="shift_forward")
    else:
        qd = qd.dt.tz_convert("America/New_York")
    return pd.to_datetime(qd.dt.date)


def load_option_rows(theta_dir: Path, ticker: str, start: pd.Timestamp, end: pd.Timestamp) -> pd.DataFrame:
    files = sorted(theta_dir.glob(f"{ticker.lower()}_1545_*.parquet"))
    if not files:
        return pd.DataFrame()
    cols = ["quote_datetime", "expiration", "strike", "right", "bid", "ask", "delta"]
    parts: list[pd.DataFrame] = []
    for p in files:
        df = pd.read_parquet(p, columns=cols)
        if df.empty:
            continue
        sd = _session_dates_series(df["quote_datetime"])
        m = (sd >= start) & (sd <= end)
        if not bool(m.any()):
            continue
        sub = df.loc[m].copy()
        sub["session_date"] = sd.loc[m].values
        parts.append(sub)
    if not parts:
        return pd.DataFrame()
    out = pd.concat(parts, ignore_index=True)
    out["session_date"] = pd.to_datetime(out["session_date"]).dt.normalize()
    out["expiration"] = pd.to_datetime(out["expiration"], errors="coerce").dt.normalize()
    out["right"] = out["right"].astype(str).str.upper().str.strip().str[0]
    for c in ("bid", "ask", "strike", "delta"):
        out[c] = pd.to_numeric(out[c], errors="coerce")
    if bool(out["strike"].max(skipna=True) < 150):
        out["strike"] = out["strike"] * 10.0
    out = out.dropna(subset=["session_date", "expiration", "strike", "bid", "ask"])
    out = out.loc[(out["bid"] > 0) & (out["ask"] > 0) & out["right"].isin(["C", "P"])]
    return out


def first_session_each_month(opt: pd.DataFrame) -> list[pd.Timestamp]:
    if opt.empty:
        return []
    d = opt["session_date"].dropna().sort_values().drop_duplicates()
    tmp = pd.DataFrame({"d": d})
    tmp["ym"] = tmp["d"].dt.to_period("M")
    out = tmp.groupby("ym", as_index=False)["d"].min()["d"]
    return sorted(pd.to_datetime(out).dt.normalize().unique())


def nearest_expiration(chain: pd.DataFrame, as_of: pd.Timestamp, target_dte: int) -> pd.Timestamp | None:
    c = chain.copy()
    c["dte"] = (c["expiration"] - as_of).dt.days
    c = c.loc[(c["dte"] >= max(7, target_dte - 12)) & (c["dte"] <= target_dte + 12)]
    if c.empty:
        return None
    row = c.iloc[(c["dte"] - target_dte).abs().argmin()]
    return pd.Timestamp(row["expiration"]).normalize()


def pick_contract(chain: pd.DataFrame, as_of: pd.Timestamp, spec: LegSpec) -> tuple[pd.Timestamp, float, str] | None:
    exp = nearest_expiration(chain, as_of, spec.dte)
    if exp is None:
        return None
    c = chain.loc[(chain["expiration"] == exp) & (chain["right"] == spec.right)].copy()
    if c.empty:
        return None
    # Prefer delta-based strike targeting when available.
    if bool(c["delta"].notna().any()):
        c["dd"] = (c["delta"] - spec.target_delta).abs()
        row = c.sort_values(["dd", "strike"], ascending=[True, True]).iloc[0]
    else:
        # Fallback: use median strike if no delta.
        med = float(c["strike"].median())
        c["sd"] = (c["strike"] - med).abs()
        row = c.sort_values(["sd", "strike"], ascending=[True, True]).iloc[0]
    return pd.Timestamp(exp).normalize(), float(row["strike"]), str(spec.right)


def leg_price(chain: pd.DataFrame, exp: pd.Timestamp, strike: float, right: str, side: str) -> float | None:
    leg = chain.loc[(chain["expiration"] == exp) & (chain["strike"] == strike) & (chain["right"] == right)]
    if leg.empty:
        return None
    row = leg.iloc[0]
    bid = float(row["bid"])
    ask = float(row["ask"])
    mid = 0.5 * (bid + ask)
    if side == "buy":
        return ask if math.isfinite(ask) and ask > 0 else mid
    return bid if math.isfinite(bid) and bid > 0 else mid


def approx_leg_price(chain: pd.DataFrame, exp: pd.Timestamp, strike: float, right: str, side: str) -> float | None:
    """
    Approximate leg lookup when exact contract is missing:
    1) same expiry/right nearest strike
    2) nearest expiry/right nearest strike
    """
    c = chain.loc[chain["right"] == right].copy()
    if c.empty:
        return None
    c_same = c.loc[c["expiration"] == exp].copy()
    if c_same.empty:
        c["exp_dist"] = (pd.to_datetime(c["expiration"]).dt.normalize() - pd.Timestamp(exp).normalize()).dt.days.abs()
        nearest_exp = pd.Timestamp(c.sort_values(["exp_dist", "expiration"]).iloc[0]["expiration"]).normalize()
        c_same = c.loc[c["expiration"] == nearest_exp].copy()
        if c_same.empty:
            return None
    c_same["k_dist"] = (c_same["strike"] - float(strike)).abs()
    row = c_same.sort_values(["k_dist", "strike"], ascending=[True, True]).iloc[0]
    bid = float(row["bid"])
    ask = float(row["ask"])
    mid = 0.5 * (bid + ask)
    if side == "buy":
        px = ask if math.isfinite(ask) and ask > 0 else mid
    else:
        px = bid if math.isfinite(bid) and bid > 0 else mid
    return float(px) if math.isfinite(px) and px >= 0 else None


def strategy_specs(name: str) -> list[LegSpec]:
    if name == "put_diagonal":
        return [LegSpec("P", -0.30, 21, -1), LegSpec("P", -0.15, 45, +1)]
    if name == "put_credit_spread":
        return [LegSpec("P", -0.25, 30, -1), LegSpec("P", -0.10, 30, +1)]
    if name == "iron_condor":
        return [LegSpec("P", -0.16, 30, -1), LegSpec("P", -0.08, 30, +1), LegSpec("C", 0.16, 30, -1), LegSpec("C", 0.08, 30, +1)]
    if name == "buy_write_pmcc":
        return [LegSpec("C", 0.80, 45, +1), LegSpec("C", 0.30, 30, -1)]
    if name == "long_strangle":
        return [LegSpec("P", -0.20, 30, +1), LegSpec("C", 0.20, 30, +1)]
    if name == "jade_lizard":
        return [LegSpec("P", -0.20, 30, -1), LegSpec("C", 0.20, 30, -1), LegSpec("C", 0.10, 30, +1)]
    if name == "butterfly_spread":
        return [LegSpec("C", 0.35, 30, +1), LegSpec("C", 0.20, 30, -2), LegSpec("C", 0.10, 30, +1)]
    if name == "bull_call_spread":
        return [LegSpec("C", 0.40, 30, +1), LegSpec("C", 0.20, 30, -1)]
    raise ValueError(name)


def run_strategy(
    opt: pd.DataFrame,
    ticker: str,
    strat: str,
    start: pd.Timestamp,
    end: pd.Timestamp,
    *,
    approximate_missing: bool,
) -> list[Trade]:
    if opt.empty:
        return []
    by_date = {pd.Timestamp(d).normalize(): g for d, g in opt.groupby("session_date", sort=False)}
    roll_dates = [d for d in first_session_each_month(opt) if d >= start and d <= end]
    specs = strategy_specs(strat)
    out: list[Trade] = []
    for d in roll_dates:
        chain = by_date.get(d)
        if chain is None or chain.empty:
            continue
        picks = []
        for sp in specs:
            p = pick_contract(chain, d, sp)
            if p is None:
                picks = []
                break
            picks.append((sp, *p))
        if not picks:
            continue

        # exit at the earliest expiration among legs (short cycle roll principle for multi-expiry structures).
        exit_date = min(exp for _, exp, _, _ in picks)
        if exit_date > end:
            continue
        exit_chain = by_date.get(pd.Timestamp(exit_date).normalize())
        if exit_chain is None or exit_chain.empty:
            continue

        pnl = 0.0
        ok = True
        for sp, exp, strike, right in picks:
            entry_side = "buy" if sp.qty > 0 else "sell"
            ep = leg_price(chain, exp, strike, right, entry_side)
            if ep is None and approximate_missing:
                ep = approx_leg_price(chain, exp, strike, right, entry_side)
            if ep is None or not math.isfinite(ep):
                ok = False
                break
            close_side = "sell" if sp.qty > 0 else "buy"
            xp = leg_price(exit_chain, exp, strike, right, close_side)
            if xp is None and approximate_missing:
                xp = approx_leg_price(exit_chain, exp, strike, right, close_side)
            if xp is None or not math.isfinite(xp):
                ok = False
                break
            # qty sign handles long/short and ratio.
            pnl += float(sp.qty) * (xp - ep) * CONTRACT_MULT * (-1.0)
        if not ok:
            continue
        # Approximate spot from the ATM call strike on the entry date.
        atm_spot = 0.0
        calls = chain.loc[(chain["right"] == "C") & chain["delta"].notna()]
        if not calls.empty:
            idx_atm = (calls["delta"] - 0.5).abs().idxmin()
            atm_spot = float(calls.loc[idx_atm, "strike"])
        out.append(Trade(ticker=ticker, strategy=strat, entry_date=d, exit_date=exit_date, pnl_usd=float(pnl), spot_entry=atm_spot))
    return out


def summarize(trades: list[Trade], start_cap: float, start: pd.Timestamp, end: pd.Timestamp) -> dict:
    if not trades:
        return {"trades": 0, "ending_capital": start_cap, "total_return_pct": 0.0, "cagr_pct": 0.0, "sharpe": float("nan"), "max_dd_pct": 0.0}
    idx = pd.bdate_range(start, end)
    df = pd.DataFrame([t.__dict__ for t in trades])
    df["exit_date"] = pd.to_datetime(df["exit_date"]).dt.normalize()
    pnl_by = df.groupby("exit_date")["pnl_usd"].sum()
    run = float(start_cap)
    vals = []
    for d in idx:
        run += float(pnl_by.get(pd.Timestamp(d).normalize(), 0.0))
        vals.append(run)
    eq = pd.Series(vals, index=idx)
    ret = eq.pct_change().replace([np.inf, -np.inf], np.nan).dropna()
    years = max((idx[-1] - idx[0]).days / 365.25, 1e-9)
    total = float(eq.iloc[-1] / start_cap - 1.0)
    cagr = float((eq.iloc[-1] / start_cap) ** (1.0 / years) - 1.0) if eq.iloc[-1] > 0 else -1.0
    sharpe = float((ret.mean() / ret.std(ddof=1)) * np.sqrt(252.0)) if len(ret) > 2 and ret.std(ddof=1) > 0 else float("nan")
    max_dd = float((eq / eq.cummax() - 1.0).min())
    return {
        "trades": int(len(df)),
        "ending_capital": float(eq.iloc[-1]),
        "total_return_pct": total * 100.0,
        "cagr_pct": cagr * 100.0,
        "sharpe": sharpe,
        "max_dd_pct": max_dd * 100.0,
    }


def main() -> None:
    ap = argparse.ArgumentParser(description="Cross-ticker strategy summary for multiple option structures.")
    ap.add_argument("--theta-dir", type=Path, default=Path("RenTech/data/theta_chunks"))
    ap.add_argument("--tickers", type=str, default="SPY,TLT,GLD,IWM,QQQ,USO")
    ap.add_argument("--start-date", type=str, default="2016-04-01")
    ap.add_argument("--end-date", type=str, default="2026-04-30")
    ap.add_argument("--starting-capital", type=float, default=100000.0)
    ap.add_argument(
        "--strict-legs",
        action="store_true",
        help="Disable nearest-contract approximation fallback; require exact leg matches.",
    )
    ap.add_argument("--output-csv", type=Path, default=Path("RenTech/data/logs/strategy_by_ticker_summary.csv"))
    args = ap.parse_args()

    tickers = [x.strip().upper() for x in args.tickers.split(",") if x.strip()]
    start = pd.Timestamp(args.start_date).normalize()
    end = pd.Timestamp(args.end_date).normalize()
    strats = [
        "put_diagonal",
        "put_credit_spread",
        "iron_condor",
        "buy_write_pmcc",
        "long_strangle",
        "jade_lizard",
        "butterfly_spread",
        "bull_call_spread",
    ]

    rows: list[dict] = []
    for t in tickers:
        opt = load_option_rows(args.theta_dir, t, start, end)
        for s in strats:
            trades = run_strategy(
                opt,
                t,
                s,
                start,
                end,
                approximate_missing=not bool(args.strict_legs),
            )
            m = summarize(trades, float(args.starting_capital), start, end)
            rows.append({"ticker": t, "strategy": s, **m})

    out = pd.DataFrame(rows)
    args.output_csv.parent.mkdir(parents=True, exist_ok=True)
    out.to_csv(args.output_csv, index=False)
    print(out.to_string(index=False, float_format=lambda x: f"{x:,.4f}"))
    print(f"\nWrote summary: {args.output_csv}")


if __name__ == "__main__":
    main()