"""
audit_reproducibility.py
========================
Run this BEFORE trusting any Model 5 (ML Stacking) output.

Usage:
    python audit_reproducibility.py --tickers AAPL TLT JPM SPY --runs 3

What it does:
    Runs the full forecast+optimization pipeline N times with identical inputs
    and measures how much the outputs drift. If max weight deviation > 0.5pp,
    the run is flagged as non-deterministic and should not be trusted.

After applying the n_jobs=1 fix in models.py, this should show 0.000 drift
on every metric. If it still shows drift, there is another source of
non-determinism that needs to be found.
"""

import sys
import os
import argparse
import copy
import numpy as np
import pandas as pd
import sqlite3

_this_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, _this_dir)

from config import load_config, Color
from core_types import PortfolioState
from data import fetch_risk_free_rate


def run_single_forecast(returns_df, bench_rets, cfg, model=5, run_id=0):
    """Run one complete forecast+optimization and return the weights."""
    from solver import build_and_optimize

    tickers = list(returns_df.columns)
    state = PortfolioState.empty(tickers)

    opt_res = build_and_optimize(
        returns_df=returns_df,
        benchmark_rets=bench_rets,
        risk_input=5,
        risk_factor=3.0,
        state=state,
        cfg=cfg,
        model=model,
        allocation_engine=1,
        silent=True
    )

    weights = opt_res.weights
    exp_rets = opt_res.expected_returns
    cov_mat = opt_res.covariance_matrix

    w_risky = weights.drop(labels=["CASH"], errors="ignore")
    opt_vol = float(np.sqrt(
        w_risky.reindex(cov_mat.columns).fillna(0).values
        @ cov_mat.values
        @ w_risky.reindex(cov_mat.columns).fillna(0).values
    ))
    opt_ret = float(
        w_risky @ exp_rets.reindex(w_risky.index).fillna(0)
    ) + (float(weights.get("CASH", 0)) * cfg["risk_free_rate"])

    return {
        "run": run_id,
        "weights": weights,
        "exp_ret": opt_ret,
        "opt_vol": opt_vol,
        "exp_rets": exp_rets,
    }


def audit(tickers, runs, model):
    cfg = load_config()
    cfg.update({
        "garch_enabled": False,   # Isolate ML non-determinism only
        "cvar_enabled": False,
        "tax_enabled": False,
        "dynamic_risk": False,
        "hmm_regime": False,
        "sector_map": {t: "Other" for t in tickers},
        "sector_limit": 1.0,
        "single_asset_min": 0.0,
        "single_asset_max": 0.40,
    })

    # Load returns from local SQLite cache
    import os
    from config import OUTPUT_DIR
    conn = sqlite3.connect(os.path.join(OUTPUT_DIR, "finance_data.db"))
    all_rets = {}
    for t in tickers:
        df = pd.read_sql(
            "SELECT date, close_price FROM daily_prices WHERE ticker=? ORDER BY date",
            conn, params=(t,)
        )
        if not df.empty:
            df["date"] = pd.to_datetime(df["date"])
            s = df.set_index("date")["close_price"].pct_change().dropna()
            if len(s) > 504:
                all_rets[t] = s
    conn.close()

    if len(all_rets) < 2:
        print(f"{Color.RED}Not enough cached data. Run the main engine first to populate the database.{Color.RESET}")
        return

    returns_df = pd.DataFrame(all_rets).dropna()
    bench_rets = returns_df.mean(axis=1)  # equal-weight proxy if SPY missing

    cfg["risk_free_rate"] = fetch_risk_free_rate(
        cfg.get("benchmarks", {}).get("risk_free", "^TNX"), 0.04
    )

    print(f"\n{Color.CYAN}Running {runs} identical forecasts to measure reproducibility...{Color.RESET}")
    print(f" Tickers: {tickers} | Model: {model} | GARCH: OFF | HMM: OFF\n")

    results = []
    for i in range(runs):
        try:
            r = run_single_forecast(returns_df, bench_rets, copy.deepcopy(cfg), model=model, run_id=i + 1)
            results.append(r)
            w_str = "  ".join(f"{t}={float(r['weights'].get(t, 0))*100:.2f}%" for t in tickers)
            print(f" Run {i+1}: {w_str} | Ret={r['exp_ret']:+.4f} | Vol={r['opt_vol']:.4f}")
        except Exception as e:
            print(f" {Color.RED}Run {i+1} FAILED: {e}{Color.RESET}")

    if len(results) < 2:
        print(f"\n{Color.RED}Not enough successful runs to compare.{Color.RESET}")
        return

    # Compute drift across all runs
    all_weights = pd.DataFrame([r["weights"] for r in results]).fillna(0)
    max_drift = float(all_weights.std().max()) * 100  # in percentage points

    all_rets_vals = [r["exp_ret"] for r in results]
    ret_drift = (max(all_rets_vals) - min(all_rets_vals)) * 10000  # in bps

    print(f"\n{'='*55}")
    print(" REPRODUCIBILITY AUDIT RESULTS")
    print(f"{'='*55}")
    print(f" Max weight std dev across runs : {max_drift:.4f} pp")
    print(f" Expected return range          : {ret_drift:.2f} bps")

    if max_drift < 0.05:
        print(f"\n {Color.GREEN}✓ PASS — Results are deterministic. Output is trustworthy.{Color.RESET}")
    elif max_drift < 0.50:
        print(f"\n {Color.YELLOW}⚠ WARNING — Small drift ({max_drift:.2f}pp). Likely solver tolerance, not XGBoost.{Color.RESET}")
        print("   This is acceptable for practical use but investigate the source.")
    else:
        print(f"\n {Color.RED}✗ FAIL — Large drift ({max_drift:.2f}pp). Results are NOT trustworthy.{Color.RESET}")
        print("   Check that n_jobs=1 in XGBRegressor in models.py.")
        print("   If the fix is applied and drift persists, another RNG source exists.")

    print("\n Individual weight ranges:")
    for col in all_weights.columns:
        mn = all_weights[col].min() * 100
        mx = all_weights[col].max() * 100
        rng = mx - mn
        flag = f" {Color.RED}← DRIFTING{Color.RESET}" if rng > 0.5 else ""
        print(f"   {col:<8} {mn:.2f}% – {mx:.2f}%  (range: {rng:.3f}pp){flag}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Reproducibility audit for the portfolio engine.")
    parser.add_argument("--tickers", nargs="+", default=["SPY", "TLT", "AAPL", "JPM"])
    parser.add_argument("--runs", type=int, default=3)
    parser.add_argument("--model", type=int, choices=[1, 2, 3, 4, 5], default=5)
    args = parser.parse_args()
    audit(args.tickers, args.runs, args.model)