"""
ForexBot2 — Walk-Forward Backtester + Parameter Optimizer
==========================================================
Replays Twelve Data USD/JPY history bar-by-bar through the same
SignalAggregator that drives live mode. The same code path that
trades USD/JPY in production is what produces these metrics — so
the backtest validates the engine, not a parallel implementation.
"""

from __future__ import annotations

import copy
import itertools
import logging
import multiprocessing as mp
import os
from dataclasses import replace
from typing import Dict, List, Optional, Tuple

import pandas as pd

from config import PlatformConfig, Timeframe, USDJPY, PIP_VALUE
from data_provider import ForexDataProvider
from metrics import BacktestResult, Trade, compute_metrics, trades_to_dataframe
from signal_aggregator import SignalAggregator, SignalReport, TradeSetup

logger = logging.getLogger(__name__)


# Period in minutes for each timeframe (used for close-time computation)
TF_MINUTES: Dict[Timeframe, int] = {
    Timeframe.M1:  1,
    Timeframe.M5:  5,
    Timeframe.M15: 15,
    Timeframe.M30: 30,
    Timeframe.H1:  60,
    Timeframe.H4:  240,
    Timeframe.D1:  1440,
    Timeframe.W1:  10080,
}

GRADE_ORDER = ["A+", "A", "B", "C", "D", "F"]


def _grade_passes(grade: str, min_grade: str) -> bool:
    """True if `grade` is at least as good as `min_grade`."""
    if grade not in GRADE_ORDER or min_grade not in GRADE_ORDER:
        return False
    return GRADE_ORDER.index(grade) <= GRADE_ORDER.index(min_grade)


# ══════════════════════════════════════════════════════════════════
#  WALK-FORWARD BACKTESTER
# ══════════════════════════════════════════════════════════════════

class WalkForwardBacktester:
    """Single-pair, single-position event-driven backtester."""

    def __init__(self, config: PlatformConfig):
        self.config = config
        self.aggregator = SignalAggregator(config)
        self.data_provider = self.aggregator.data_provider
        self._regime_adx: Optional[pd.Series] = None
        self._regime_tf: Optional[Timeframe] = None

    # ── Public ─────────────────────────────────────────────────────

    def load_data(self) -> Dict[Timeframe, pd.DataFrame]:
        """Pull (or read from disk cache) all needed timeframes."""
        timeframes = [self.config.primary_timeframe] + list(
            self.config.confirmation_timeframes
        )
        bt = self.config.backtest
        if bt.regime_filter and bt.regime_filter_timeframe not in timeframes:
            timeframes.append(bt.regime_filter_timeframe)

        dfs: Dict[Timeframe, pd.DataFrame] = {}
        for tf in timeframes:
            df = self.data_provider.fetch_full_history(
                symbol=self.config.symbol,
                timeframe=tf,
                max_bars=bt.max_bars,
                start_date=bt.start_date,
                end_date=bt.end_date,
            )
            df = self._apply_window(df)
            dfs[tf] = df
        return dfs

    def run(
        self,
        dfs: Optional[Dict[Timeframe, pd.DataFrame]] = None,
        progress: bool = True,
    ) -> BacktestResult:
        if dfs is None:
            dfs = self.load_data()

        self._precompute_regime_filter(dfs)

        primary_tf = self.config.primary_timeframe
        primary_df = dfs.get(primary_tf)
        if primary_df is None or primary_df.empty:
            return BacktestResult(initial_balance=self.config.backtest.initial_balance)

        # Warmup: enough bars for ema_trend & ichimoku_senkou_b on the primary TF
        warmup = max(
            self.config.technical.ema_trend,
            self.config.technical.ichimoku_senkou_b,
            200,
        ) + 30

        if len(primary_df) <= warmup + 1:
            logger.warning(
                f"Insufficient data: {len(primary_df)} primary bars, "
                f"need >{warmup}."
            )
            return BacktestResult(initial_balance=self.config.backtest.initial_balance)

        bt = self.config.backtest
        balance = bt.initial_balance
        equity_points: List[Tuple[pd.Timestamp, float]] = [
            (primary_df.index[warmup], balance)
        ]
        trades: List[Trade] = []

        open_position: Optional[Dict] = None

        # Optional progress bar
        iterator = range(warmup, len(primary_df) - 1)
        if progress:
            try:
                from tqdm import tqdm  # type: ignore
                iterator = tqdm(iterator, desc="Backtesting", unit="bar", leave=False)
            except ImportError:
                pass

        for t in iterator:
            decision_time = self._bar_close_time(primary_df.index[t], primary_tf)

            # ── Step 1: manage any open position ───────────────────
            if open_position is not None:
                next_bar = primary_df.iloc[t]
                exit_info = self._check_exit(open_position, next_bar)
                if exit_info is not None:
                    closed = self._close_trade(
                        open_position, next_bar.name, exit_info, balance
                    )
                    trades.append(closed)
                    balance += closed.pnl_currency
                    equity_points.append((next_bar.name, balance))
                    open_position = None

            # ── Step 2: maybe open a new position ──────────────────
            if open_position is None or not bt.one_position_only:
                # Slice each TF's df up to (and including) bars closed by decision_time
                sliced = self._slice_dfs(dfs, decision_time)
                report = self.aggregator.analyze_from_dataframes(
                    symbol=self.config.symbol,
                    dfs=sliced,
                    use_news=False,
                )

                if (
                    _grade_passes(report.signal_quality, bt.min_trade_grade)
                    and abs(float(report.final_score)) >= bt.min_abs_score
                    and report.trade_setup is not None
                    and t + 1 < len(primary_df)
                    and self._regime_passes(decision_time)
                ):
                    next_bar = primary_df.iloc[t + 1]
                    open_position = self._open_trade(report, next_bar, balance)

        # ── Force-close any still-open position at end of data ────
        if open_position is not None:
            last_bar = primary_df.iloc[-1]
            closed = self._close_trade(
                open_position,
                last_bar.name,
                {"price": float(last_bar["Close"]), "reason": "end_of_data"},
                balance,
            )
            trades.append(closed)
            balance += closed.pnl_currency
            equity_points.append((last_bar.name, balance))

        # Build equity curve
        eq_index = [p[0] for p in equity_points]
        eq_values = [p[1] for p in equity_points]
        equity_curve = pd.Series(eq_values, index=pd.DatetimeIndex(eq_index))

        elapsed_days: Optional[float] = None
        if len(equity_curve) >= 2:
            elapsed_days = (equity_curve.index[-1] - equity_curve.index[0]).total_seconds() / 86400

        metrics = compute_metrics(
            trades, equity_curve, bt.initial_balance, elapsed_days=elapsed_days
        )

        return BacktestResult(
            trades=trades,
            equity_curve=equity_curve,
            initial_balance=bt.initial_balance,
            final_balance=balance,
            metrics=metrics,
        )

    # ── Private: regime gate ───────────────────────────────────────

    def _precompute_regime_filter(self, dfs: Dict[Timeframe, pd.DataFrame]) -> None:
        """Pre-compute ADX on the regime-filter TF; cached for the whole run."""
        bt = self.config.backtest
        if not bt.regime_filter:
            self._regime_adx = None
            self._regime_tf = None
            return

        tf = bt.regime_filter_timeframe
        df = dfs.get(tf)
        if df is None or df.empty:
            logger.warning(
                f"Regime filter requested ({tf.value}) but no data available; "
                f"all signals will be skipped."
            )
            self._regime_adx = pd.Series(dtype=float)
            self._regime_tf = tf
            return

        self._regime_tf = tf
        self._regime_adx = self._compute_adx_series(df, self.config.technical.adx_period)

    def _regime_passes(self, decision_time: pd.Timestamp) -> bool:
        bt = self.config.backtest
        if not bt.regime_filter:
            return True
        if self._regime_adx is None or self._regime_adx.empty:
            return False
        period_min = TF_MINUTES[self._regime_tf]
        close_times = self._regime_adx.index + pd.Timedelta(minutes=period_min)
        valid = self._regime_adx[close_times <= decision_time]
        if valid.empty:
            return False
        last = valid.iloc[-1]
        if pd.isna(last):
            return False
        return float(last) >= bt.regime_filter_min_adx

    @staticmethod
    def _compute_adx_series(df: pd.DataFrame, period: int = 14) -> pd.Series:
        """Wilder's ADX series — same algorithm as engine._compute_adx."""
        import numpy as np
        high, low, close = df["High"], df["Low"], df["Close"]
        tr = pd.concat([
            high - low,
            (high - close.shift()).abs(),
            (low - close.shift()).abs(),
        ], axis=1).max(axis=1)
        plus_dm = high.diff()
        minus_dm = -low.diff()
        plus_dm = plus_dm.where((plus_dm > minus_dm) & (plus_dm > 0), 0.0)
        minus_dm = minus_dm.where((minus_dm > plus_dm) & (minus_dm > 0), 0.0)
        atr_s = tr.ewm(alpha=1/period, min_periods=period).mean()
        plus_di = 100 * (plus_dm.ewm(alpha=1/period, min_periods=period).mean() / atr_s)
        minus_di = 100 * (minus_dm.ewm(alpha=1/period, min_periods=period).mean() / atr_s)
        dx = 100 * (plus_di - minus_di).abs() / (plus_di + minus_di).replace(0, np.nan)
        return dx.ewm(alpha=1/period, min_periods=period).mean()

    # ── Private: data slicing ──────────────────────────────────────

    def _apply_window(self, df: pd.DataFrame) -> pd.DataFrame:
        bt = self.config.backtest
        if bt.start_date:
            start = pd.Timestamp(bt.start_date, tz="UTC")
            df = df[df.index >= start]
        if bt.end_date:
            end = pd.Timestamp(bt.end_date, tz="UTC")
            df = df[df.index <= end]
        return df

    @staticmethod
    def _bar_close_time(open_time: pd.Timestamp, tf: Timeframe) -> pd.Timestamp:
        return open_time + pd.Timedelta(minutes=TF_MINUTES[tf])

    def _slice_dfs(
        self, dfs: Dict[Timeframe, pd.DataFrame], decision_time: pd.Timestamp
    ) -> Dict[Timeframe, pd.DataFrame]:
        """Return only bars whose close time ≤ decision_time (no lookahead)."""
        sliced: Dict[Timeframe, pd.DataFrame] = {}
        for tf, df in dfs.items():
            close_times = df.index + pd.Timedelta(minutes=TF_MINUTES[tf])
            sliced[tf] = df[close_times <= decision_time]
        return sliced

    # ── Private: trade execution ───────────────────────────────────

    def _open_trade(
        self, report: SignalReport, next_bar: pd.Series, balance: float
    ) -> Dict:
        """Open at next bar's open, applying spread + slippage."""
        ts: TradeSetup = report.trade_setup  # type: ignore[assignment]
        bt = self.config.backtest

        raw_entry = float(next_bar["Open"])

        # Sanity-test inversion: flip direction and mirror SL/TP around entry
        direction = ts.direction
        sl = ts.stop_loss
        tp = ts.take_profit_1
        if bt.invert_signal:
            direction = "SELL" if direction == "BUY" else "BUY"
            sl, tp = (2 * raw_entry - sl), (2 * raw_entry - tp)

        # Cost in price units = (half-spread + slippage) × pip_value
        half_spread_pips = 0.5 if bt.use_spread_cost else 0.0
        cost_pips = half_spread_pips + bt.slippage_pips
        cost_price = cost_pips * PIP_VALUE

        if direction == "BUY":
            entry_price = raw_entry + cost_price  # buy higher (worse)
        else:
            entry_price = raw_entry - cost_price  # sell lower (worse)

        # Position sizing — risk a fixed % of current balance
        risk_pct = bt.risk_per_trade_pct / 100.0
        risk_amount = balance * risk_pct
        risk_pips = max(ts.risk_pips, 1e-9)
        pip_currency_value = risk_amount / risk_pips  # $ per pip

        return {
            "direction": direction,
            "entry_time": next_bar.name,
            "entry_price": entry_price,
            "stop_loss": sl,
            "take_profit": tp,
            "risk_pips": risk_pips,
            "pip_currency_value": pip_currency_value,
            "balance_at_open": balance,
            "grade": report.signal_quality,
            "score_at_entry": float(report.final_score),
            "confidence_at_entry": float(report.final_confidence),
            "bar_index_open": next_bar.name,
        }

    def _check_exit(
        self, position: Dict, bar: pd.Series
    ) -> Optional[Dict]:
        """Return {price, reason} if SL or TP hit, else None.

        On the entry bar (bar with same timestamp as position's entry_time),
        we already filled at the open — only check subsequent bars' OHLC.
        """
        if bar.name == position["bar_index_open"]:
            return None  # entry bar — no exit check

        bt = self.config.backtest
        high = float(bar["High"])
        low = float(bar["Low"])

        sl = position["stop_loss"]
        tp = position["take_profit"]
        slip = bt.slippage_pips * PIP_VALUE
        half_spread = (0.5 if bt.use_spread_cost else 0.0) * PIP_VALUE

        if position["direction"] == "BUY":
            sl_hit = low <= sl
            tp_hit = high >= tp
            if sl_hit and tp_hit:
                # Conservative: SL first
                return {"price": sl - slip - half_spread, "reason": "sl"} if bt.sl_first_on_tie else \
                       {"price": tp - slip - half_spread, "reason": "tp"}
            if sl_hit:
                return {"price": sl - slip - half_spread, "reason": "sl"}
            if tp_hit:
                return {"price": tp - slip - half_spread, "reason": "tp"}
        else:  # SELL
            sl_hit = high >= sl
            tp_hit = low <= tp
            if sl_hit and tp_hit:
                return {"price": sl + slip + half_spread, "reason": "sl"} if bt.sl_first_on_tie else \
                       {"price": tp + slip + half_spread, "reason": "tp"}
            if sl_hit:
                return {"price": sl + slip + half_spread, "reason": "sl"}
            if tp_hit:
                return {"price": tp + slip + half_spread, "reason": "tp"}
        return None

    def _close_trade(
        self, position: Dict, exit_time: pd.Timestamp, exit_info: Dict, balance: float
    ) -> Trade:
        exit_price = exit_info["price"]
        reason = exit_info["reason"]

        if position["direction"] == "BUY":
            pnl_pips = (exit_price - position["entry_price"]) / PIP_VALUE
        else:
            pnl_pips = (position["entry_price"] - exit_price) / PIP_VALUE

        pnl_currency = pnl_pips * position["pip_currency_value"]
        pnl_pct = pnl_currency / position["balance_at_open"] * 100

        # Duration in primary-TF bars (timestamp diff / period)
        period_min = TF_MINUTES[self.config.primary_timeframe]
        duration_minutes = (exit_time - position["entry_time"]).total_seconds() / 60
        duration_bars = int(round(duration_minutes / period_min))

        return Trade(
            entry_time=position["entry_time"],
            exit_time=exit_time,
            direction=position["direction"],
            entry_price=position["entry_price"],
            exit_price=exit_price,
            stop_loss=position["stop_loss"],
            take_profit=position["take_profit"],
            exit_reason=reason,
            risk_pips=position["risk_pips"],
            pnl_pips=pnl_pips,
            pnl_pct=pnl_pct,
            pnl_currency=pnl_currency,
            duration_bars=duration_bars,
            grade=position["grade"],
            score_at_entry=position["score_at_entry"],
            confidence_at_entry=position["confidence_at_entry"],
        )


# ══════════════════════════════════════════════════════════════════
#  PARAMETER OPTIMIZER (PARALLEL)
# ══════════════════════════════════════════════════════════════════

def _apply_overrides(config: PlatformConfig, overrides: Dict[str, float]) -> PlatformConfig:
    """Return a deep-copied config with parameter overrides applied."""
    new_config = copy.deepcopy(config)

    if "atr_sl_multiplier" in overrides:
        new_config.signals.atr_sl_multiplier = float(overrides["atr_sl_multiplier"])
    if "default_rr_ratio" in overrides:
        new_config.signals.default_rr_ratio = float(overrides["default_rr_ratio"])
    if "min_trade_grade" in overrides:
        new_config.backtest.min_trade_grade = str(overrides["min_trade_grade"])
    if "min_abs_score" in overrides:
        new_config.backtest.min_abs_score = float(overrides["min_abs_score"])
    if "invert_signal" in overrides:
        new_config.backtest.invert_signal = bool(overrides["invert_signal"])
    if "regime_filter" in overrides:
        new_config.backtest.regime_filter = bool(overrides["regime_filter"])
    if "regime_filter_min_adx" in overrides:
        new_config.backtest.regime_filter_min_adx = float(overrides["regime_filter_min_adx"])
    if "risk_per_trade_pct" in overrides:
        new_config.backtest.risk_per_trade_pct = float(overrides["risk_per_trade_pct"])

    return new_config


def _run_single_combo(args: Tuple[PlatformConfig, Dict[str, float]]) -> BacktestResult:
    """Worker function — must be top-level for multiprocessing.Pool."""
    base_config, overrides = args
    config = _apply_overrides(base_config, overrides)
    backtester = WalkForwardBacktester(config)
    # Workers reuse the disk parquet cache, so no Twelve Data fetch needed
    result = backtester.run(progress=False)
    result.params = overrides
    return result


class Optimizer:
    """Embarrassingly parallel grid search over BacktestConfig.param_grid."""

    def __init__(self, config: PlatformConfig):
        self.config = config

    def run_grid(
        self,
        param_grid: Optional[Dict[str, List[float]]] = None,
        prewarm_cache: bool = True,
    ) -> List[BacktestResult]:
        if param_grid is None:
            param_grid = self.config.backtest.param_grid

        # Prewarm the disk cache once before fanning out — avoids each worker
        # re-fetching from Twelve Data.
        if prewarm_cache:
            logger.info("Prewarming disk cache for backtest data...")
            warmer = WalkForwardBacktester(self.config)
            warmer.load_data()

        keys = list(param_grid.keys())
        value_lists = [param_grid[k] for k in keys]
        combos = [
            dict(zip(keys, vals))
            for vals in itertools.product(*value_lists)
        ]

        n_workers = min(
            self.config.backtest.parallel_workers,
            len(combos),
            os.cpu_count() or 1,
        )
        logger.info(
            f"Optimizer: {len(combos)} combos × {n_workers} workers"
        )

        tasks = [(self.config, c) for c in combos]

        if n_workers <= 1:
            results = [_run_single_combo(t) for t in tasks]
        else:
            with mp.Pool(processes=n_workers) as pool:
                results = pool.map(_run_single_combo, tasks)

        # Sort by total return descending
        results.sort(
            key=lambda r: r.metrics.get("total_return_pct", -float("inf")),
            reverse=True,
        )
        return results