# scripts/research/wf_compare_strategies.py

import os
import sys
from pathlib import Path
from dotenv import load_dotenv
from datetime import timedelta

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

sys.path.insert(0, str(Path(__file__).parent.parent.parent))

from src.core.engine import Engine
from src.data.storage import StorageManager
from src.risk.sizing.percent_risk import PercentRiskSizer
from src.risk.stops.trailing_stop import TrailingStop

from src.strategies.moving_average import MovingAverageCrossover
from src.strategies.trend_filtered import TrendFilteredMACrossover
from src.strategies.breakout import DonchianBreakout
from src.strategies.mean_reversion import RSIMeanReversion

from src.utils.logger import log


# --------------------------------------------------
# CONFIG
# --------------------------------------------------
SYMBOL = "BTC/USDT"
TIMEFRAME = "1h"

TRAIN_DAYS = 120
TEST_DAYS = 30
STEP_DAYS = 30

INITIAL_CAPITAL = 10_000

RISK = PercentRiskSizer(0.01)
STOP = TrailingStop(0.02)

OUTPUT_DIR = Path(__file__).parent / "output/wf_strategies"
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)


# --------------------------------------------------
# STRATEGIES
# --------------------------------------------------
STRATEGIES = {
    "MA_Crossover": lambda: MovingAverageCrossover(20, 50, "ema", False),
    "TrendFiltered_MA": lambda: TrendFilteredMACrossover(20, 50, "ema", 20),
    "Donchian": lambda: DonchianBreakout(20),
    "RSI_Reversion": lambda: RSIMeanReversion(14, 30, 70),
}


# --------------------------------------------------
def setup_env():
    env_path = Path(__file__).parent.parent.parent / "config" / "secrets.env"
    load_dotenv(env_path)


def load_data():
    setup_env()

    storage = StorageManager(
        db_host=os.getenv("DB_HOST"),
        db_port=int(os.getenv("DB_PORT", 5432)),
        db_name=os.getenv("DB_NAME"),
        db_user=os.getenv("DB_USER"),
        db_password=os.getenv("DB_PASSWORD"),
    )

    data = storage.load_ohlcv(SYMBOL, TIMEFRAME, use_cache=True)
    storage.close()

    if data.empty:
        raise RuntimeError("No data loaded")

    return data


# --------------------------------------------------
def compute_percentiles(curves, n_points=100):
    x_common = np.linspace(0, 1, n_points)
    interpolated = []

    for curve in curves:
        x = np.linspace(0, 1, len(curve))
        interpolated.append(np.interp(x_common, x, curve.values))

    arr = np.vstack(interpolated)
    return (
        x_common,
        np.percentile(arr, 10, axis=0),
        np.percentile(arr, 50, axis=0),
        np.percentile(arr, 90, axis=0),
    )


# --------------------------------------------------
def run():
    log.info("🧪 WF MULTI-STRATEGY TEST (FULL VISUAL + STATS)")

    data = load_data()

    results_summary = []
    all_accumulated = {}
    all_dispersion_percentiles = {}

    for name, strat_factory in STRATEGIES.items():
        log.info(f"▶ Strategy: {name}")

        capital = INITIAL_CAPITAL
        equity_accumulated = []
        dispersion_curves = []
        window_returns = []

        worst_dd = 0.0
        start = data.index[0]

        while True:
            train_end = start + timedelta(days=TRAIN_DAYS)
            test_end = train_end + timedelta(days=TEST_DAYS)

            if test_end > data.index[-1]:
                break

            test_data = data.loc[train_end:test_end]

            engine = Engine(
                strategy=strat_factory(),
                initial_capital=capital,
                position_sizer=RISK,
                stop_loss=STOP,
                commission=0.001,
                slippage=0.0005,
            )

            res = engine.run(test_data)

            # ---- accumulated
            capital *= (1 + res["total_return_pct"] / 100)
            equity_accumulated.append((res["timestamps"][-1], capital))

            # ---- dispersion (normalized reset)
            eq = pd.Series(res["equity_curve"], index=res["timestamps"])
            dispersion_curves.append(eq / eq.iloc[0])

            window_returns.append(res["total_return_pct"])
            worst_dd = min(worst_dd, res["max_drawdown_pct"])

            start += timedelta(days=STEP_DAYS)

        # ===============================
        # STORE ACCUMULATED
        # ===============================
        acc_df = pd.DataFrame(equity_accumulated, columns=["time", "capital"]).set_index("time")
        all_accumulated[name] = acc_df

        # ===============================
        # 1️⃣ DISPERSION – RAW (OLD)
        # ===============================
        plt.figure(figsize=(12, 5))
        for eq in dispersion_curves:
            plt.plot(eq.index, eq.values, alpha=0.3)

        plt.title(f"WF Equity Dispersion (Raw) – {name}")
        plt.ylabel("Normalized Equity")
        plt.grid(alpha=0.3)

        path = OUTPUT_DIR / f"wf_equity_dispersion_raw_{name}.png"
        plt.tight_layout()
        plt.savefig(path)
        plt.close()

        # ===============================
        # 2️⃣ DISPERSION – PERCENTILES (NEW)
        # ===============================
        x, p10, p50, p90 = compute_percentiles(dispersion_curves)
        all_dispersion_percentiles[name] = (x, p10, p50, p90)

        plt.figure(figsize=(12, 5))
        plt.plot(x, p50, label="P50", linewidth=2)
        plt.fill_between(x, p10, p90, alpha=0.3, label="P10–P90")

        plt.title(f"WF Equity Dispersion (Percentiles) – {name}")
        plt.ylabel("Normalized Equity")
        plt.xlabel("Window Progress")
        plt.legend()
        plt.grid(alpha=0.3)

        path = OUTPUT_DIR / f"wf_equity_dispersion_percentiles_{name}.png"
        plt.tight_layout()
        plt.savefig(path)
        plt.close()

        # ===============================
        # 3️⃣ RETURN DISTRIBUTION (OLD)
        # ===============================
        plt.figure(figsize=(10, 5))
        plt.hist(window_returns, bins=20, density=True, alpha=0.7)

        plt.title(f"WF Return Distribution – {name}")
        plt.xlabel("Return per window (%)")
        plt.ylabel("Density")
        plt.grid(alpha=0.3)

        path = OUTPUT_DIR / f"wf_return_distribution_{name}.png"
        plt.tight_layout()
        plt.savefig(path)
        plt.close()

        # ===============================
        # SUMMARY
        # ===============================
        results_summary.append({
            "strategy": name,
            "windows": len(window_returns),
            "return_mean": round(np.mean(window_returns), 2),
            "return_median": round(np.median(window_returns), 2),
            "max_dd_worst": round(worst_dd, 2),
            "final_capital": round(capital, 2),
        })

    # ===============================
    # 4️⃣ ACCUMULATED EQUITY – ALL
    # ===============================
    plt.figure(figsize=(13, 6))
    for name, acc_df in all_accumulated.items():
        plt.plot(acc_df.index, acc_df["capital"], linewidth=2, label=name)

    plt.title("WF Accumulated Equity – Strategy Comparison")
    plt.ylabel("Capital")
    plt.legend()
    plt.grid(alpha=0.3)

    path = OUTPUT_DIR / "wf_equity_accumulated_ALL.png"
    plt.tight_layout()
    plt.savefig(path)
    plt.close()

    # ===============================
    # 5️⃣ DISPERSION COMPARISON – P50
    # ===============================
    plt.figure(figsize=(12, 6))
    for name, (x, _, p50, _) in all_dispersion_percentiles.items():
        plt.plot(x, p50, linewidth=2, label=name)

    plt.title("WF Dispersion Comparison (Median – P50)")
    plt.ylabel("Normalized Equity")
    plt.xlabel("Window Progress")
    plt.legend()
    plt.grid(alpha=0.3)

    path = OUTPUT_DIR / "wf_dispersion_comparison_P50.png"
    plt.tight_layout()
    plt.savefig(path)
    plt.close()

    # ===============================
    # FINAL TABLE
    # ===============================
    summary_df = pd.DataFrame(results_summary).set_index("strategy")

    print("\n" + "=" * 80)
    print("📊 WF SUMMARY (ACCUMULATED)")
    print("=" * 80)
    print(summary_df)
    print("=" * 80)


if __name__ == "__main__":
    run()
