import os
import sys
from pathlib import Path
from datetime import datetime, timedelta
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from dotenv import load_dotenv

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

from src.utils.logger import log
from src.data.storage import StorageManager
from src.core.engine import Engine
from src.strategies import MovingAverageCrossover
from src.risk.sizing.percent_risk import PercentRiskSizer
from src.risk.stops.fixed_stop import FixedStop
from src.risk.stops.trailing_stop import TrailingStop
from src.risk.stops.atr_stop import ATRStop

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

TRAIN_DAYS = 120
TEST_DAYS = 30
STEP_DAYS = 30

INITIAL_CAPITAL = 10_000
COMMISSION = 0.001
SLIPPAGE = 0.0005

RISK_FRACTION = 0.01

OUT_DIR = Path("scripts/research/output/wf_stops") / f"{SYMBOL.replace('/', '_')}_{TIMEFRAME}"
OUT_DIR.mkdir(parents=True, exist_ok=True)

STOPS = {
    "Fixed 2%": FixedStop(0.02),
    "Trailing 2%": TrailingStop(0.02),
    "ATR 14 x 2.0": ATRStop(14, 2.0),
}

# --------------------------------------------------
# Helpers
# --------------------------------------------------
def make_strategy():
    return MovingAverageCrossover(
        fast_period=10,
        slow_period=30,
        ma_type="ema",
        use_adx=False,
        adx_threshold=20.0,
    )


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=SYMBOL,
        timeframe=TIMEFRAME,
        use_cache=True,
    )

    storage.close()

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

    return data


# --------------------------------------------------
# Walk Forward
# --------------------------------------------------
def run():
    log.info("=" * 80)
    log.info("🔁 WALK-FORWARD – STOP COMPARISON (120/30/30, 1h)")
    log.info("=" * 80)

    data = load_data()

    wf_results = []
    equity_curves = {name: [] for name in STOPS.keys()}

    start_time = data.index[0]
    end_time = data.index[-1]

    window_id = 0
    current_train_start = start_time

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

        if test_end > end_time:
            break

        train_df = data.loc[current_train_start:train_end]
        test_df = data.loc[test_start:test_end]

        if len(test_df) < 50:
            break

        window_id += 1

        print()
        print(
            f"WF Window {window_id:02d} | "
            f"TRAIN {train_df.index[0].date()} → {train_df.index[-1].date()} | "
            f"TEST {test_df.index[0].date()} → {test_df.index[-1].date()} | "
            f"bars_test={len(test_df)}"
        )

        for stop_name, stop in STOPS.items():
            engine = Engine(
                strategy=make_strategy(),
                initial_capital=INITIAL_CAPITAL,
                commission=COMMISSION,
                slippage=SLIPPAGE,
                position_sizer=PercentRiskSizer(RISK_FRACTION),
                stop_loss=stop,
            )

            res = engine.run(test_df)

            wf_results.append({
                "window": window_id,
                "stop": stop_name,
                "train_start": train_df.index[0],
                "train_end": train_df.index[-1],
                "test_start": test_df.index[0],
                "test_end": test_df.index[-1],
                "trades": res["total_trades"],
                "max_dd_pct": res["max_drawdown_pct"],
                "return_pct": res["total_return_pct"],
                "final_equity": res["final_equity"],
            })

            equity_curves[stop_name].append(
                pd.Series(res["equity_curve"], index=res["timestamps"])
            )

            print(
                f"  {stop_name:<13} | "
                f"Trades: {res['total_trades']:>3} | "
                f"MaxDD: {res['max_drawdown_pct']:>7.2f}% | "
                f"Return: {res['total_return_pct']:>7.2f}%"
            )

        current_train_start += timedelta(days=STEP_DAYS)

    # --------------------------------------------------
    # Save results
    # --------------------------------------------------
    df = pd.DataFrame(wf_results)
    df.to_csv(OUT_DIR / "wf_results.csv", index=False)

    print()
    print("=" * 80)
    print("📊 WF SUMMARY (aggregated)")
    print("=" * 80)

    summary = (
        df.groupby("stop")
        .agg(
            windows=("window", "nunique"),
            trades_avg=("trades", "mean"),
            max_dd_worst=("max_dd_pct", "min"),
            return_mean=("return_pct", "mean"),
            return_median=("return_pct", "median"),
        )
        .round(2)
    )

    print(summary)
    print("=" * 80)

    # --------------------------------------------------
    # Plot equity curves (visual comparison)
    # --------------------------------------------------
    plt.figure(figsize=(14, 7))

    for stop_name, curves in equity_curves.items():
        if not curves:
            continue
        concat_curve = pd.concat(curves)
        plt.plot(concat_curve.index, concat_curve.values, label=stop_name)

    plt.title(f"WF Equity Comparison – {SYMBOL} {TIMEFRAME}")
    plt.xlabel("Time")
    plt.ylabel("Equity (per-window)")
    plt.legend()
    plt.grid(alpha=0.3)
    plt.tight_layout()

    plt.savefig(OUT_DIR / "wf_equity_comparison.png")
    plt.close()


if __name__ == "__main__":
    run()