#!/usr/bin/env python3
"""
Quant Research: Run all example strategies + custom experimental strategies
against BTC and ETH data. Produce comparison metrics and identify what works.

Usage:
    cd challenge && python research/run_strategies.py
"""
import sys
import os
import time
import json
from datetime import datetime

# Ensure challenge dir is on path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from vanta.backtest import BacktestRunner, BacktestResult, _compute_metrics
from vanta.tracker import TrackerBase
from vanta.indicators import (
    SMA, EMA, RSI, MACD, BollingerBands, ATR, VWAP, Returns,
    OrderBookImbalance, SpreadTracker, FundingSignal, BasisTracker,
)
from vanta.examples import (
    MeanReversionTracker,
    TrendFollowingTracker,
    RSITracker,
    BollingerBreakoutTracker,
    VolatilityRegimeTracker,
    DualTimeframeTracker,
    VolumeProfileTracker,
)


# ── Experimental Strategies ───────────────────────────────────────────────────


class MultiTimeframeMomentum(TrackerBase):
    """Use 5m and 15m EMA crossovers for stronger trend confirmation."""

    def __init__(self):
        super().__init__()
        # 5m timeframe
        self.ema_5m_fast = EMA(period=5)
        self.ema_5m_slow = EMA(period=20)
        # 15m timeframe
        self.ema_15m_fast = EMA(period=5)
        self.ema_15m_slow = EMA(period=15)
        # 1m for entry timing
        self.rsi_1m = RSI(period=14)

    def tick(self, data: dict):
        super().tick(data)
        for c in data.get("candles_1m", []):
            self.rsi_1m.update(float(c["close"]))
        for c in data.get("candles_5m", []):
            self.ema_5m_fast.update(float(c["close"]))
            self.ema_5m_slow.update(float(c["close"]))
        for c in data.get("candles_15m", []):
            self.ema_15m_fast.update(float(c["close"]))
            self.ema_15m_slow.update(float(c["close"]))

    def predict(self, symbol, horizon_seconds, step_seconds):
        if not (self.ema_5m_fast.ready and self.ema_15m_fast.ready and self.rsi_1m.ready):
            return {"expected_return": 0.0}

        trend_5m = (self.ema_5m_fast.value - self.ema_5m_slow.value) / self.ema_5m_slow.value
        trend_15m = (self.ema_15m_fast.value - self.ema_15m_slow.value) / self.ema_15m_slow.value

        # Both timeframes agree on direction
        if trend_5m > 0 and trend_15m > 0 and self.rsi_1m.value < 65:
            signal = min(trend_5m, trend_15m) * 0.5
            return {"expected_return": max(0.0, min(signal, 0.005))}
        elif trend_5m < 0 and trend_15m < 0 and self.rsi_1m.value > 35:
            signal = max(trend_5m, trend_15m) * 0.5
            return {"expected_return": max(-0.005, min(0.0, signal))}

        return {"expected_return": 0.0}


class MACDDivergence(TrackerBase):
    """MACD histogram momentum with signal-line crossover."""

    def __init__(self):
        super().__init__()
        self.macd = MACD(fast=8, slow=21, signal=5)
        self.prev_hist = None

    def tick(self, data: dict):
        super().tick(data)
        for c in data.get("candles_1m", []):
            self.macd.update(float(c["close"]))

    def predict(self, symbol, horizon_seconds, step_seconds):
        if not self.macd.ready:
            return {"expected_return": 0.0}

        hist = self.macd.value.histogram

        # Histogram flipping direction (momentum change)
        if self.prev_hist is not None:
            if self.prev_hist < 0 and hist > 0:
                self.prev_hist = hist
                return {"expected_return": 0.003}
            elif self.prev_hist > 0 and hist < 0:
                self.prev_hist = hist
                return {"expected_return": -0.003}

        self.prev_hist = hist

        # Trend following on histogram magnitude
        signal = hist * 0.01  # scale down
        return {"expected_return": max(-0.005, min(0.005, signal))}


class MeanReversionBollinger(TrackerBase):
    """Aggressive mean reversion: fade Bollinger Band extremes with ATR sizing."""

    def __init__(self, bb_period=20, bb_std=2.0, atr_period=14):
        super().__init__()
        self.bb = BollingerBands(period=bb_period, num_std=bb_std)
        self.atr = ATR(period=atr_period)
        self.sma = SMA(period=bb_period)

    def tick(self, data: dict):
        super().tick(data)
        for c in data.get("candles_1m", []):
            price = float(c["close"])
            self.bb.update(price)
            self.sma.update(price)
            self.atr.update(float(c["high"]), float(c["low"]), price)

    def predict(self, symbol, horizon_seconds, step_seconds):
        if not (self.bb.ready and self.atr.ready and self.sma.ready):
            return {"expected_return": 0.0}

        pct_b = self.bb.value.percent_b
        width = self.bb.value.width

        # Only trade when bands are reasonably wide (not in squeeze)
        if width < 0.005:
            return {"expected_return": 0.0}

        # Fade extremes — scale by how extreme
        if pct_b < 0.1:
            signal = (0.1 - pct_b) * 0.05
            return {"expected_return": min(signal, 0.005)}
        elif pct_b > 0.9:
            signal = -(pct_b - 0.9) * 0.05
            return {"expected_return": max(signal, -0.005)}

        return {"expected_return": 0.0}


class AdaptiveRSI(TrackerBase):
    """RSI with adaptive thresholds based on recent volatility."""

    def __init__(self, rsi_period=14, vol_period=20):
        super().__init__()
        self.rsi = RSI(period=rsi_period)
        self.atr = ATR(period=vol_period)
        self.atr_sma = SMA(period=vol_period)

    def tick(self, data: dict):
        super().tick(data)
        for c in data.get("candles_1m", []):
            price = float(c["close"])
            self.rsi.update(price)
            self.atr.update(float(c["high"]), float(c["low"]), price)
            if self.atr.ready:
                self.atr_sma.update(self.atr.value)

    def predict(self, symbol, horizon_seconds, step_seconds):
        if not (self.rsi.ready and self.atr_sma.ready and self.atr_sma.value > 0):
            return {"expected_return": 0.0}

        # In high vol, widen RSI thresholds; in low vol, tighten them
        vol_ratio = self.atr.value / self.atr_sma.value
        oversold = 30 - (vol_ratio - 1.0) * 10  # e.g. 25 in high vol, 35 in low vol
        overbought = 70 + (vol_ratio - 1.0) * 10

        oversold = max(15, min(40, oversold))
        overbought = max(60, min(85, overbought))

        if self.rsi.value < oversold:
            strength = (oversold - self.rsi.value) / oversold
            return {"expected_return": strength * 0.004}
        elif self.rsi.value > overbought:
            strength = (self.rsi.value - overbought) / (100 - overbought)
            return {"expected_return": -strength * 0.004}

        return {"expected_return": 0.0}


class VWAPReversion(TrackerBase):
    """Mean-reversion to VWAP with volume confirmation."""

    def __init__(self):
        super().__init__()
        self.vwap = VWAP()
        self.vol_sma = SMA(period=20)
        self.price_sma = SMA(period=5)  # short-term price

    def tick(self, data: dict):
        super().tick(data)
        for c in data.get("candles_1m", []):
            price = float(c["close"])
            volume = float(c["volume"])
            self.vwap.update(price, volume)
            self.vol_sma.update(volume)
            self.price_sma.update(price)

    def predict(self, symbol, horizon_seconds, step_seconds):
        if not (self.vwap.ready and self.vol_sma.ready and self.price_sma.ready):
            return {"expected_return": 0.0}

        price = self.price_sma.value
        deviation = (price - self.vwap.value) / self.vwap.value

        # Only trade when deviation is meaningful
        if abs(deviation) < 0.001:
            return {"expected_return": 0.0}

        # Fade the deviation (revert to VWAP)
        signal = -deviation * 0.5
        return {"expected_return": max(-0.005, min(0.005, signal))}


class EnsembleTracker(TrackerBase):
    """Equal-weight ensemble of multiple signals."""

    def __init__(self):
        super().__init__()
        self.rsi = RSI(period=14)
        self.bb = BollingerBands(period=20)
        self.macd = MACD(fast=8, slow=21, signal=5)
        self.ema_fast = EMA(period=5)
        self.ema_slow = EMA(period=20)
        self.atr = ATR(period=14)
        self.atr_sma = SMA(period=14)

    def tick(self, data: dict):
        super().tick(data)
        for c in data.get("candles_1m", []):
            price = float(c["close"])
            self.rsi.update(price)
            self.bb.update(price)
            self.macd.update(price)
            self.ema_fast.update(price)
            self.ema_slow.update(price)
            self.atr.update(float(c["high"]), float(c["low"]), price)
            if self.atr.ready:
                self.atr_sma.update(self.atr.value)

    def predict(self, symbol, horizon_seconds, step_seconds):
        signals = []

        # RSI signal
        if self.rsi.ready:
            if self.rsi.value < 30:
                signals.append(0.003)
            elif self.rsi.value > 70:
                signals.append(-0.003)
            else:
                signals.append(0.0)

        # Bollinger signal
        if self.bb.ready:
            pct_b = self.bb.value.percent_b
            if pct_b < 0.15:
                signals.append(0.003)
            elif pct_b > 0.85:
                signals.append(-0.003)
            else:
                signals.append(0.0)

        # MACD signal
        if self.macd.ready:
            signals.append(max(-0.003, min(0.003, self.macd.value.histogram * 0.01)))

        # Trend signal
        if self.ema_fast.ready and self.ema_slow.ready:
            spread = (self.ema_fast.value - self.ema_slow.value) / self.ema_slow.value
            signals.append(max(-0.003, min(0.003, spread * 0.3)))

        if not signals:
            return {"expected_return": 0.0}

        avg_signal = sum(signals) / len(signals)

        # Vol-scale: reduce position in high vol
        if self.atr.ready and self.atr_sma.ready and self.atr_sma.value > 0:
            vol_ratio = self.atr.value / self.atr_sma.value
            if vol_ratio > 1.5:
                avg_signal *= 0.5

        return {"expected_return": max(-0.005, min(0.005, avg_signal))}


# ── Run ───────────────────────────────────────────────────────────────────────


def run_all():
    strategies = [
        # Built-in examples
        ("MeanReversion", MeanReversionTracker()),
        ("TrendFollowing", TrendFollowingTracker()),
        ("RSI", RSITracker()),
        ("BollingerBreakout", BollingerBreakoutTracker()),
        ("VolatilityRegime", VolatilityRegimeTracker()),
        ("DualTimeframe", DualTimeframeTracker()),
        ("VolumeProfile", VolumeProfileTracker()),
        # Experimental
        ("MultiTF_Momentum", MultiTimeframeMomentum()),
        ("MACD_Divergence", MACDDivergence()),
        ("MeanRev_Bollinger", MeanReversionBollinger()),
        ("AdaptiveRSI", AdaptiveRSI()),
        ("VWAP_Reversion", VWAPReversion()),
        ("Ensemble", EnsembleTracker()),
    ]

    # Test period: 6 months
    test_start = "2025-07-01"
    test_end = "2026-01-01"
    subject = "BTCUSDT"

    print(f"\n{'='*90}")
    print(f"  QUANT RESEARCH: {subject} | {test_start} → {test_end}")
    print(f"  Running {len(strategies)} strategies...")
    print(f"{'='*90}\n")

    results = {}
    for name, model in strategies:
        t0 = time.time()
        try:
            runner = BacktestRunner(model=model)
            result = runner.run(
                subject=subject,
                start=test_start,
                end=test_end,
                window_size=120,
                prediction_interval_seconds=60,  # 1 prediction per minute to speed up
                horizon_seconds=300,
            )
            elapsed = time.time() - t0
            results[name] = result
            m = result.metrics
            print(f"  ✓ {name:25s}  mean={m.get('mean_return',0):+.8f}  "
                  f"hit={m.get('hit_rate',0):.3f}  "
                  f"sortino={m.get('sortino_ratio',0):+.4f}  "
                  f"dd={m.get('max_drawdown',0):+.4f}  "
                  f"({elapsed:.1f}s)")
        except Exception as e:
            elapsed = time.time() - t0
            print(f"  ✗ {name:25s}  ERROR: {e} ({elapsed:.1f}s)")

    # Summary table sorted by mean_return
    print(f"\n{'='*90}")
    print(f"  LEADERBOARD (sorted by mean_return)")
    print(f"{'='*90}")

    metric_keys = ["mean_return", "hit_rate", "sortino_ratio", "max_drawdown",
                    "score_recent", "score_steady", "score_anchor", "ic"]

    header = f"{'Rank':>4}  {'Strategy':25s}"
    for k in metric_keys:
        header += f"  {k:>14s}"
    print(header)
    print("-" * len(header))

    sorted_names = sorted(results, key=lambda n: results[n].metrics.get("mean_return", 0), reverse=True)
    for rank, name in enumerate(sorted_names, 1):
        m = results[name].metrics
        row = f"{rank:>4}  {name:25s}"
        for k in metric_keys:
            row += f"  {m.get(k, 0):>14.6f}"
        print(row)

    print("-" * len(header))

    # Save results to JSON
    output = {
        "test_period": {"start": test_start, "end": test_end, "subject": subject},
        "results": {}
    }
    for name in sorted_names:
        m = results[name].metrics
        n_preds = len(results[name]._predictions)
        scored = sum(1 for p in results[name]._predictions if p.get("score") is not None)
        output["results"][name] = {
            "rank": sorted_names.index(name) + 1,
            "metrics": {k: round(v, 8) for k, v in m.items()},
            "predictions": n_preds,
            "scored": scored,
        }

    os.makedirs("research/output", exist_ok=True)
    with open("research/output/strategy_comparison.json", "w") as f:
        json.dump(output, f, indent=2, default=str)

    print(f"\nResults saved to research/output/strategy_comparison.json")
    return results


if __name__ == "__main__":
    run_all()
