"""V7+VT36 full trade log: regime changes, instrument switches, PT events,
vol-target scale, contribution analysis.

Run from /home/gahow/projects/quant:
    uv run python research/v7_trade_audit.py
"""
from __future__ import annotations

import sys
sys.path.insert(0, ".")

import numpy as np
import pandas as pd

import data_manager
import metrics
from main import backtest
from strategies.trend_rider_v7 import TrendRiderV7

YEARS = 10
CAPITAL = 100_000
TX_COST = 0.001
FIXED_FEE = 2.0


def load_data():
    tickers = sorted(set(["SPY", "TQQQ", "UPRO", "GLD", "DBC", "SHY"]))
    data = data_manager.update("etfs", tickers, with_open=False)
    if isinstance(data, tuple):
        data = data[0]
    cutoff = data.index[-1] - pd.DateOffset(years=YEARS)
    data = data[data.index >= cutoff]
    cols = [c for c in tickers if c in data.columns]
    return data[cols]


def trace_v7(data: pd.DataFrame):
    """Replicate V7+VT36 signal generation step-by-step, logging every event."""
    v7 = TrendRiderV7(target_vol=0.36, min_lev=0.75)

    # --- Layer 1: V3 regime weights ---
    w_v3 = v7.v3.generate_signals(data)

    # --- Layer 2: Vol-target overlay ---
    daily_ret = data.pct_change(fill_method=None).fillna(0.0)
    port_rets_v3 = (w_v3 * daily_ret).sum(axis=1)
    realized_vol = (
        port_rets_v3.rolling(v7.vol_window, min_periods=21).std() * np.sqrt(252)
    )
    scale_raw = (v7.target_vol / realized_vol).clip(lower=v7.min_lev, upper=v7.max_lev)
    scale = scale_raw.shift(1).fillna(1.0)
    w_vt = w_v3.mul(scale, axis=0)

    # --- Layer 3: Profit-take (replicate the loop, logging events) ---
    w = w_vt.copy()
    held = w.idxmax(axis=1)
    max_w = w.max(axis=1)
    held[max_w < 1e-8] = ""

    park_col = v7.pt_park if v7.pt_park in w.columns else ""
    entry_price = None
    current_sym = None
    is_stopped = False
    restore_level = v7.pt_threshold - v7.pt_band

    pt_events = []  # (date, type, sym, entry_price, exit_price, gain%)
    for i in range(len(w)):
        sym = held.iloc[i]
        if not sym or max_w.iloc[i] < 1e-8:
            current_sym = None
            entry_price = None
            is_stopped = False
            continue

        if sym != current_sym:
            current_sym = sym
            entry_price = float(data[sym].iloc[i - 1]) if i > 0 and sym in data.columns else None
            is_stopped = False
            continue

        if entry_price is None or entry_price <= 0 or sym not in data.columns:
            continue

        yesterday = float(data[sym].iloc[i - 1]) if i > 0 else float(data[sym].iloc[i])
        gain = yesterday / entry_price - 1.0

        if is_stopped:
            if gain < restore_level:
                is_stopped = False
                pt_events.append((data.index[i], "RESTORE", sym, entry_price, yesterday, gain))
            else:
                w.iloc[i] = 0.0
                if park_col:
                    w.at[w.index[i], park_col] = scale.iloc[i]
        else:
            if gain >= v7.pt_threshold:
                is_stopped = True
                pt_events.append((data.index[i], "CLEAR", sym, entry_price, yesterday, gain))
                w.iloc[i] = 0.0
                if park_col:
                    w.at[w.index[i], park_col] = scale.iloc[i]

    return w_v3, w_vt, w, scale, pt_events


def analyze_regime(w_v3: pd.DataFrame, data: pd.DataFrame):
    """Trace regime changes from V3 weights."""
    held_v3 = w_v3.idxmax(axis=1)
    max_w = w_v3.max(axis=1)
    held_v3[max_w < 1e-8] = ""

    risk_on_syms = {"TQQQ", "UPRO"}
    risk_off_syms = {"GLD", "DBC"}

    events = []
    prev_regime = None
    prev_sym = None
    regime_start = None

    for i in range(len(held_v3)):
        sym = held_v3.iloc[i]
        if not sym:
            continue

        regime = "risk_on" if sym in risk_on_syms else ("risk_off" if sym in risk_off_syms else "other")

        if regime != prev_regime:
            if prev_regime is not None:
                events.append({
                    "date": data.index[i],
                    "type": "REGIME_CHANGE",
                    "from": prev_regime,
                    "to": regime,
                    "from_sym": prev_sym,
                    "to_sym": sym,
                })
            prev_regime = regime
            regime_start = data.index[i]
            prev_sym = sym
        elif sym != prev_sym:
            events.append({
                "date": data.index[i],
                "type": "INSTRUMENT_SWITCH",
                "regime": regime,
                "from_sym": prev_sym,
                "to_sym": sym,
            })
            prev_sym = sym

    return events, held_v3


def compute_contributions(data, w_v3, w_vt, w_final, scale, pt_events):
    """Return attribution of returns to risk-on, risk-off, PT-park, and vol-target."""
    daily_ret = data.pct_change(fill_method=None).fillna(0.0)

    # Identify which days are in PT-park (cleared to cash by PT layer)
    # NOTE: SHY is not in V3's output columns, so PT clears to 0% (cash), not SHY.
    risk_on_syms = {"TQQQ", "UPRO"}
    risk_off_syms = {"GLD", "DBC"}

    vt_sum = w_vt.abs().sum(axis=1)
    final_sum = w_final.abs().sum(axis=1)

    held_final = w_final.idxmax(axis=1)
    max_w_final = w_final.max(axis=1)
    held_final[max_w_final < 1e-8] = ""

    held_vt = w_vt.idxmax(axis=1)
    max_w_vt = w_vt.max(axis=1)
    held_vt[max_w_vt < 1e-8] = ""

    # Classify each day
    day_class = pd.Series("none", index=data.index)
    for i in range(len(data)):
        sym_vt = held_vt.iloc[i]
        sym_final = held_final.iloc[i]
        # PT-park: VT layer has allocation but final layer cleared to 0
        if vt_sum.iloc[i] > 0.01 and final_sum.iloc[i] < 0.01:
            day_class.iloc[i] = "pt_park"
        elif not sym_final:
            continue
        elif sym_final in risk_on_syms:
            day_class.iloc[i] = "risk_on"
        elif sym_final in risk_off_syms:
            day_class.iloc[i] = "risk_off"
        elif sym_final == "SHY":
            day_class.iloc[i] = "risk_off_shy"
        else:
            day_class.iloc[i] = "other"

    # Daily portfolio return for each layer
    port_ret_v3 = (w_v3 * daily_ret).sum(axis=1)
    port_ret_vt = (w_vt * daily_ret).sum(axis=1)
    port_ret_final = (w_final * daily_ret).sum(axis=1)

    # Contribution by regime state
    contrib = {}
    for cls in ["risk_on", "risk_off", "risk_off_shy", "pt_park", "none"]:
        mask = day_class == cls
        contrib[cls] = {
            "days": int(mask.sum()),
            "cum_return": float((1 + port_ret_final[mask]).prod() - 1) if mask.any() else 0.0,
            "avg_daily_ret": float(port_ret_final[mask].mean()) if mask.any() else 0.0,
        }

    # Vol-target impact: compare V3-only vs V3+VT
    eq_v3 = (1 + port_ret_v3).cumprod() * CAPITAL
    eq_vt = (1 + port_ret_vt).cumprod() * CAPITAL
    eq_final = (1 + port_ret_final).cumprod() * CAPITAL

    return day_class, contrib, eq_v3, eq_vt, eq_final


def main():
    data = load_data()
    print(f"Data: {data.index[0].date()} to {data.index[-1].date()}, {len(data)} days")
    print(f"Columns: {list(data.columns)}")

    w_v3, w_vt, w_final, scale, pt_events = trace_v7(data)

    # =========================================================================
    # 1. Regime changes and instrument switches
    # =========================================================================
    regime_events, held_v3 = analyze_regime(w_v3, data)
    print(f"\n{'='*90}")
    print("  REGIME CHANGES AND INSTRUMENT SWITCHES")
    print(f"{'='*90}")
    regime_changes = [e for e in regime_events if e["type"] == "REGIME_CHANGE"]
    instrument_switches = [e for e in regime_events if e["type"] == "INSTRUMENT_SWITCH"]

    print(f"\nTotal regime changes: {len(regime_changes)}")
    print(f"Total instrument switches (within regime): {len(instrument_switches)}")

    print(f"\n--- Regime Changes ---")
    for e in regime_changes:
        print(f"  {e['date'].date()}  {e['from']:>9} -> {e['to']:<9}  ({e['from_sym']} -> {e['to_sym']})")

    print(f"\n--- Instrument Switches ---")
    for e in instrument_switches:
        print(f"  {e['date'].date()}  [{e['regime']:>9}]  {e['from_sym']} -> {e['to_sym']}")

    # Regime holding periods
    risk_on_syms = {"TQQQ", "UPRO"}
    risk_off_syms = {"GLD", "DBC"}
    regime_series = pd.Series("none", index=data.index)
    for i in range(len(held_v3)):
        sym = held_v3.iloc[i]
        if sym in risk_on_syms:
            regime_series.iloc[i] = "risk_on"
        elif sym in risk_off_syms:
            regime_series.iloc[i] = "risk_off"

    active = regime_series[regime_series != "none"]
    if len(active) > 0:
        pct_risk_on = (active == "risk_on").mean()
        pct_risk_off = (active == "risk_off").mean()
        print(f"\n  Time in risk_on:  {pct_risk_on:.1%}")
        print(f"  Time in risk_off: {pct_risk_off:.1%}")

        # Average holding period per regime stint
        regime_shifts = active != active.shift(1)
        stint_id = regime_shifts.cumsum()
        stint_lengths = stint_id.groupby(stint_id).count()
        avg_stint = stint_lengths.mean()
        print(f"  Avg regime stint: {avg_stint:.0f} trading days")

    # =========================================================================
    # 2. Profit-take events
    # =========================================================================
    print(f"\n{'='*90}")
    print("  PROFIT-TAKE EVENTS")
    print(f"{'='*90}")
    print(f"\nTotal PT events: {len(pt_events)}")
    clears = [e for e in pt_events if e[1] == "CLEAR"]
    restores = [e for e in pt_events if e[1] == "RESTORE"]
    print(f"  CLEAR events:   {len(clears)}")
    print(f"  RESTORE events: {len(restores)}")

    if clears:
        print(f"\n--- CLEAR Events ---")
        gains = []
        for date, typ, sym, entry, exit_p, gain in clears:
            print(f"  {date.date()}  {sym:>5}  entry=${entry:.2f}  exit=${exit_p:.2f}  gain={gain:+.1%}")
            gains.append(gain)
        print(f"\n  Avg gain at CLEAR: {np.mean(gains):.1%}")
        print(f"  Min gain at CLEAR: {np.min(gains):.1%}")
        print(f"  Max gain at CLEAR: {np.max(gains):.1%}")

    if restores:
        print(f"\n--- RESTORE Events ---")
        for date, typ, sym, entry, exit_p, gain in restores:
            print(f"  {date.date()}  {sym:>5}  entry=${entry:.2f}  price=${exit_p:.2f}  gain={gain:+.1%}")

    # PT-park days
    day_class, contrib, eq_v3, eq_vt, eq_final = compute_contributions(
        data, w_v3, w_vt, w_final, scale, pt_events)

    pt_park_days = (day_class == "pt_park").sum()
    total_active_days = (day_class != "none").sum()
    print(f"\n  Days in PT-park (SHY): {pt_park_days} ({pt_park_days/total_active_days:.1%} of active days)")

    # =========================================================================
    # 3. Vol-target scale analysis
    # =========================================================================
    print(f"\n{'='*90}")
    print("  VOL-TARGET SCALE ANALYSIS")
    print(f"{'='*90}")
    active_scale = scale[scale < 1.0 - 1e-6]
    print(f"\n  Scale stats (when < 1.0):")
    print(f"    Days at full scale (1.0):    {(scale >= 1.0 - 1e-6).sum()}")
    print(f"    Days below 1.0:              {len(active_scale)}")
    print(f"    Days below 0.90:             {(scale < 0.90).sum()}")
    print(f"    Days below 0.80:             {(scale < 0.80).sum()}")
    print(f"    Days at floor (0.75):        {(scale <= 0.75 + 1e-6).sum()}")
    print(f"    Min scale:                   {scale.min():.3f}")
    print(f"    Mean scale:                  {scale.mean():.3f}")

    # When did scale hit the floor?
    at_floor = scale[scale <= 0.75 + 1e-6]
    if len(at_floor) > 0:
        print(f"\n  Periods at floor (scale=0.75):")
        floor_mask = scale <= 0.75 + 1e-6
        shifts = floor_mask != floor_mask.shift(1)
        stint_ids = shifts.cumsum()
        floor_stints = stint_ids[floor_mask]
        for stint_id_val in floor_stints.unique():
            stint_dates = floor_stints[floor_stints == stint_id_val].index
            if len(stint_dates) > 0:
                print(f"    {stint_dates[0].date()} to {stint_dates[-1].date()} ({len(stint_dates)} days)")

    # =========================================================================
    # 4. Contribution analysis
    # =========================================================================
    print(f"\n{'='*90}")
    print("  RETURN CONTRIBUTION BY STATE")
    print(f"{'='*90}")

    for cls in ["risk_on", "risk_off", "risk_off_shy", "pt_park"]:
        c = contrib[cls]
        print(f"\n  {cls:>14}:  {c['days']:>5} days  "
              f"cum_return={c['cum_return']:>+8.1%}  "
              f"avg_daily={c['avg_daily_ret']*10000:>+6.1f}bps")

    # =========================================================================
    # 5. With vs Without vol-target
    # =========================================================================
    print(f"\n{'='*90}")
    print("  VOL-TARGET IMPACT (V3 alone vs V3+VT)")
    print(f"{'='*90}")

    daily_ret = data.pct_change(fill_method=None).fillna(0.0)

    # V3-only with turnover cost
    eq_v3_bt = backtest(
        type("V3Only", (), {"generate_signals": lambda self, d: TrendRiderV7(target_vol=0.36, min_lev=0.75).v3.generate_signals(d)})(),
        data, initial_capital=CAPITAL, transaction_cost=TX_COST, fixed_fee=FIXED_FEE)
    m_v3 = metrics.raw_summary(eq_v3_bt)

    # V7+VT36 full
    eq_full = backtest(
        TrendRiderV7(target_vol=0.36, min_lev=0.75),
        data, initial_capital=CAPITAL, transaction_cost=TX_COST, fixed_fee=FIXED_FEE)
    m_full = metrics.raw_summary(eq_full)

    # V7 with VT but no PT
    eq_no_pt = backtest(
        TrendRiderV7(target_vol=0.36, min_lev=0.75, pt_threshold=0.0),
        data, initial_capital=CAPITAL, transaction_cost=TX_COST, fixed_fee=FIXED_FEE)
    m_no_pt = metrics.raw_summary(eq_no_pt)

    print(f"\n  {'Variant':<30} {'Ann':>7} {'Sharpe':>7} {'MaxDD':>7} {'Sortino':>8} {'Calmar':>7}")
    print(f"  {'-'*67}")
    for label, m in [("V3 only (no VT, no PT)", m_v3),
                     ("V3 + VT (no PT)", m_no_pt),
                     ("V7+VT36 full (VT+PT)", m_full)]:
        print(f"  {label:<30} {m['annualizedReturn']*100:>6.1f}% {m['sharpeRatio']:>7.2f} "
              f"{m['maxDrawdown']*100:>6.1f}% {m['sortinoRatio']:>8.2f} {m['calmarRatio']:>7.2f}")

    # =========================================================================
    # 6. Drawdown analysis
    # =========================================================================
    print(f"\n{'='*90}")
    print("  WORST DRAWDOWN PERIODS")
    print(f"{'='*90}")

    rolling_peak = eq_full.cummax()
    dd = (eq_full - rolling_peak) / rolling_peak

    # Find top 5 drawdown troughs
    dd_sorted = dd.sort_values()
    seen_windows = []
    top_dds = []
    for date, dd_val in dd_sorted.items():
        too_close = False
        for prev_date in seen_windows:
            if abs((date - prev_date).days) < 60:
                too_close = True
                break
        if not too_close:
            top_dds.append((date, dd_val))
            seen_windows.append(date)
        if len(top_dds) >= 5:
            break

    for trough_date, trough_dd in top_dds:
        # Find start of this drawdown (last peak before trough)
        pre_trough = rolling_peak.loc[:trough_date]
        peak_val = pre_trough.iloc[-1]
        peak_dates = eq_full[eq_full >= peak_val * 0.999].loc[:trough_date]
        if len(peak_dates) > 0:
            peak_date = peak_dates.index[0]
        else:
            peak_date = trough_date

        # What was the strategy holding?
        held_during = day_class.loc[peak_date:trough_date]
        regime_during = held_during.value_counts()

        # Recovery date
        post_trough = eq_full.loc[trough_date:]
        recovered = post_trough[post_trough >= peak_val]
        recovery_date = recovered.index[0] if len(recovered) > 0 else None

        print(f"\n  DD {trough_dd:.1%} | {peak_date.date()} -> {trough_date.date()}"
              f" | Recovery: {recovery_date.date() if recovery_date else 'N/A'}")
        print(f"    State during DD: {dict(regime_during)}")

        # What instrument was held at the trough?
        held_final = w_final.idxmax(axis=1)
        max_w_final = w_final.max(axis=1)
        held_final[max_w_final < 1e-8] = ""
        trough_sym = held_final.loc[trough_date] if trough_date in held_final.index else "?"
        print(f"    Held at trough: {trough_sym}")

    # =========================================================================
    # 7. Disproportionate trades
    # =========================================================================
    print(f"\n{'='*90}")
    print("  TOP 10 BEST AND WORST SINGLE-DAY RETURNS")
    print(f"{'='*90}")

    port_ret_final = (w_final * daily_ret).sum(axis=1)
    # Apply turnover cost
    turnover = w_final.diff().abs().sum(axis=1).fillna(0.0)
    port_ret_final_net = port_ret_final - turnover * TX_COST

    held_final = w_final.idxmax(axis=1)
    max_w_final = w_final.max(axis=1)
    held_final[max_w_final < 1e-8] = ""

    best = port_ret_final_net.nlargest(10)
    worst = port_ret_final_net.nsmallest(10)

    print("\n  BEST DAYS:")
    for date, ret in best.items():
        sym = held_final.loc[date] if date in held_final.index else "?"
        sc = scale.loc[date] if date in scale.index else 1.0
        print(f"    {date.date()}  {ret:>+7.2%}  holding={sym}  scale={sc:.2f}")

    print("\n  WORST DAYS:")
    for date, ret in worst.items():
        sym = held_final.loc[date] if date in held_final.index else "?"
        sc = scale.loc[date] if date in scale.index else 1.0
        print(f"    {date.date()}  {ret:>+7.2%}  holding={sym}  scale={sc:.2f}")

    # Compound impact of top/bottom 20 days
    sorted_rets = port_ret_final_net.sort_values()
    total_days = len(sorted_rets)
    eq_without_top20 = (1 + sorted_rets.iloc[:-20]).prod()
    eq_without_bottom20 = (1 + sorted_rets.iloc[20:]).prod()
    eq_all = (1 + sorted_rets).prod()

    print(f"\n  Total compound growth factor: {eq_all:.2f}x")
    print(f"  Without best 20 days:         {eq_without_top20:.2f}x")
    print(f"  Without worst 20 days:        {eq_without_bottom20:.2f}x")

    # =========================================================================
    # 8. Annual returns breakdown
    # =========================================================================
    print(f"\n{'='*90}")
    print("  ANNUAL RETURNS BREAKDOWN")
    print(f"{'='*90}")

    yearly = port_ret_final_net.groupby(port_ret_final_net.index.year)
    print(f"\n  {'Year':>6} {'Return':>8} {'Vol':>7} {'MaxDD':>7} {'Days_ON':>8} {'Days_OFF':>9} {'Days_PT':>8}")
    print(f"  {'-'*55}")
    for year, rets in yearly:
        ann_r = (1 + rets).prod() ** (252 / len(rets)) - 1 if len(rets) > 0 else 0
        vol = rets.std() * np.sqrt(252) if len(rets) > 1 else 0
        eq_yr = (1 + rets).cumprod()
        mdd = ((eq_yr / eq_yr.cummax()) - 1).min()
        yr_class = day_class.loc[rets.index]
        n_on = (yr_class == "risk_on").sum()
        n_off = ((yr_class == "risk_off") | (yr_class == "risk_off_shy")).sum()
        n_pt = (yr_class == "pt_park").sum()
        print(f"  {year:>6} {ann_r:>+7.1%} {vol:>6.1%} {mdd:>+6.1%} {n_on:>8} {n_off:>9} {n_pt:>8}")


if __name__ == "__main__":
    main()