#!/usr/bin/env python3
"""
Phase 2: re-simulate the existing trade pool under different SL/TP geometries
without re-running the engine.

The engine already gave us, per trade: entry_time, entry_price, direction,
and the original SL distance (= 1.5 × ATR by default). That distance lets
us recover ATR and re-derive new SL/TP for any (sl_mult, rr) combo. We
then walk forward through the saved OHLCV bars, applying first-touch
simulation with worst-case (SL first if both touch same bar).

This is faster than re-running the engine because the signal output (when
to trade, in which direction, at what price) is independent of SL/TP
geometry — only the exits change.

Discipline: we report TRAIN and TEST stats for every cell, time-split at
70/30 by entry_time. A cell wins only if it's positive on the held-out
TEST slice.

Usage:
    python geometry_grid.py
    python geometry_grid.py --apply-phase1-filters
    python geometry_grid.py --pairs USDJPY  # single-pair grid
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path
from typing import Dict, Tuple

import numpy as np
import pandas as pd

# Reuse history loader so we don't re-implement
from backtest import load_history
from config import Timeframe

# trades.csv stores TF as Twelve-Data interval string ("15min", "1h", "1d", …)
_TFSTR_TO_ENUM = {
    "1min":  Timeframe.M1,  "5min":  Timeframe.M5,
    "15min": Timeframe.M15, "30min": Timeframe.M30,
    "1h":    Timeframe.H1,  "4h":    Timeframe.H4,
    "1d":    Timeframe.D1,  "1wk":   Timeframe.W1,
}

# Filters carried over from Phase 1
PHASE1_BAD_HOURS = {10, 17, 19, 21}
PHASE1_BAD_PAIRS = {"NZD/USD", "AUD/USD", "GBP/USD"}
PHASE1_BAD_DOWS  = {5, 6}  # Sat / Sun


def simulate_one(
    direction: str,
    entry: float,
    new_sl: float,
    new_tp: float,
    forward_bars: pd.DataFrame,
    max_bars: int,
    spread_pips: float,
    pip_value: float,
) -> Tuple[float, str]:
    """First-touch simulation. Returns (r_multiple, outcome)."""
    sl_distance = abs(entry - new_sl)
    if sl_distance == 0:
        return 0.0, "INVALID"
    risk_pips = sl_distance / pip_value
    spread_r = spread_pips / risk_pips if risk_pips > 0 else 0.0
    rr_target = abs(new_tp - entry) / sl_distance

    bars_examined = 0
    for ts, bar in forward_bars.iterrows():
        bars_examined += 1
        if bars_examined > max_bars:
            break
        bar_high = bar["High"]
        bar_low  = bar["Low"]
        if direction == "BUY":
            sl_hit = bar_low <= new_sl
            tp_hit = bar_high >= new_tp
        else:
            sl_hit = bar_high >= new_sl
            tp_hit = bar_low <= new_tp
        if sl_hit:
            return -1.0 - spread_r, "LOSS_SL"
        if tp_hit:
            return rr_target - spread_r, "WIN_TP1"

    if bars_examined == 0:
        return 0.0, "INVALID"
    last = forward_bars.iloc[bars_examined - 1]
    pnl = (last["Close"] - entry) if direction == "BUY" else (entry - last["Close"])
    return (pnl / sl_distance) - spread_r, "TIMEOUT"


def pip_value_for(symbol: str) -> float:
    """Return pip value: 0.01 for JPY pairs, 0.0001 elsewhere."""
    return 0.01 if "JPY" in symbol else 0.0001


def grid_run(
    df: pd.DataFrame,
    history: Dict[Tuple[str, str], pd.DataFrame],
    sl_mults: list,
    rr_ratios: list,
    max_hold: int = 100,
    spread_pips: float = 1.0,
) -> pd.DataFrame:
    """For each (sl_mult, rr) cell, re-simulate every trade and report stats."""
    # Pre-index history by (symbol, tf_str) for fast lookup
    rows = []
    n_total = len(df) * len(sl_mults) * len(rr_ratios)
    print(f"   evaluating {n_total:,} simulations "
          f"({len(df):,} trades × {len(sl_mults)} SL × {len(rr_ratios)} RR)…")

    for sl_mult in sl_mults:
        for rr in rr_ratios:
            r_multiples = []
            outcomes = []
            entry_times = []
            pairs = []
            tfs = []
            for trade in df.itertuples(index=False):
                # Recover ATR from original SL distance (engine used 1.5×ATR)
                atr = abs(trade.entry - trade.sl) / 1.5
                new_sl_dist = sl_mult * atr
                if trade.direction == "BUY":
                    new_sl = trade.entry - new_sl_dist
                    new_tp = trade.entry + new_sl_dist * rr
                else:
                    new_sl = trade.entry + new_sl_dist
                    new_tp = trade.entry - new_sl_dist * rr

                hist_df = history.get((trade.symbol, trade.tf))
                if hist_df is None:
                    continue
                ts = pd.Timestamp(trade.entry_time)
                # Forward bars (strictly after entry)
                forward = hist_df[hist_df.index > ts].head(max_hold)
                if forward.empty:
                    continue

                pip = pip_value_for(trade.symbol)
                r, outcome = simulate_one(
                    trade.direction, trade.entry,
                    new_sl, new_tp, forward,
                    max_bars=max_hold, spread_pips=spread_pips,
                    pip_value=pip,
                )
                r_multiples.append(r)
                outcomes.append(outcome)
                entry_times.append(ts)
                pairs.append(trade.pair)
                tfs.append(trade.tf)

            arr = np.array(r_multiples) if r_multiples else np.array([0.0])
            wins = (arr > 0).sum()
            losses = (arr <= 0).sum()
            gross_win = arr[arr > 0].sum()
            gross_loss = -arr[arr <= 0].sum()
            pf = gross_win / gross_loss if gross_loss > 0 else float("inf")
            cum = np.cumsum(arr)
            running_peak = np.maximum.accumulate(cum)
            max_dd = (cum - running_peak).min()

            rows.append({
                "sl_mult":  sl_mult,
                "rr":       rr,
                "n":        len(arr),
                "wins":     int(wins),
                "wr":       wins / len(arr) * 100 if len(arr) else 0,
                "avg_r":    arr.mean() if len(arr) else 0,
                "total_r":  arr.sum(),
                "pf":       pf,
                "max_dd":   max_dd,
                # Carry the per-trade results so we can split train/test downstream
                "_arr":         arr,
                "_entry_times": entry_times,
                "_pairs":       pairs,
                "_outcomes":    outcomes,
            })
    return pd.DataFrame(rows)


def split_summary(cell: pd.Series, train_frac: float = 0.7) -> dict:
    """Split a cell's per-trade results into train/test by entry_time."""
    arr = cell["_arr"]
    times = cell["_entry_times"]
    if len(arr) == 0:
        return {"train_total": 0, "train_wr": 0, "test_total": 0, "test_wr": 0}
    order = np.argsort(times)
    cutoff = int(len(arr) * train_frac)
    train_idx = order[:cutoff]
    test_idx  = order[cutoff:]
    train_arr = arr[train_idx]
    test_arr  = arr[test_idx]
    return {
        "train_n":     len(train_arr),
        "train_wr":    (train_arr > 0).mean() * 100 if len(train_arr) else 0,
        "train_avg":   train_arr.mean() if len(train_arr) else 0,
        "train_total": train_arr.sum(),
        "test_n":      len(test_arr),
        "test_wr":     (test_arr > 0).mean() * 100 if len(test_arr) else 0,
        "test_avg":    test_arr.mean() if len(test_arr) else 0,
        "test_total":  test_arr.sum(),
    }


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--trades", default="backtest_results_all/trades.csv")
    ap.add_argument("--out", default="backtest_results_all/geometry_grid.txt")
    ap.add_argument("--max-hold", type=int, default=100)
    ap.add_argument("--spread-pips", type=float, default=1.0)
    ap.add_argument("--apply-phase1-filters", action="store_true",
                    help="Drop weekend trades, worst 4 hours, worst 3 pairs before grid")
    ap.add_argument("--sl-mults", nargs="+", type=float,
                    default=[1.0, 1.5, 2.0, 2.5, 3.0])
    ap.add_argument("--rr-ratios", nargs="+", type=float,
                    default=[1.0, 1.5, 2.0, 2.5])
    args = ap.parse_args()

    p = Path(args.trades)
    if not p.exists():
        print(f"❌ {p} not found")
        return 1

    df = pd.read_csv(p)
    df["entry_time"] = pd.to_datetime(df["entry_time"])
    df["hour"] = df["entry_time"].dt.hour
    df["dow"] = df["entry_time"].dt.dayofweek

    print(f"📊 Loaded {len(df):,} trades from {p}")

    if args.apply_phase1_filters:
        before = len(df)
        df = df[~df["hour"].isin(PHASE1_BAD_HOURS)]
        df = df[~df["pair"].isin(PHASE1_BAD_PAIRS)]
        df = df[~df["dow"].isin(PHASE1_BAD_DOWS)]
        df = df.reset_index(drop=True)
        print(f"   Phase 1 filters applied: {before:,} → {len(df):,} trades "
              f"({len(df)/before*100:.0f}% retained)")

    # Load history for forward simulation
    print(f"📁 Loading history…")
    pairs = sorted(df["symbol"].unique().tolist())
    tfs   = sorted({_TFSTR_TO_ENUM[t] for t in df["tf"].unique() if t in _TFSTR_TO_ENUM},
                   key=lambda x: x.value)
    history_raw = load_history(pairs, tfs)
    # Rekey by (symbol, tf_string) to match trades.csv
    history = {}
    for (sym, tf_enum), hdf in history_raw.items():
        history[(sym, tf_enum.value)] = hdf

    print(f"\n🧮 Grid search:")
    print(f"   SL multipliers (× ATR): {args.sl_mults}")
    print(f"   RR ratios:               {args.rr_ratios}\n")

    grid = grid_run(
        df, history, args.sl_mults, args.rr_ratios,
        max_hold=args.max_hold, spread_pips=args.spread_pips,
    )

    # ── Heatmap (full-data total R) ─────────────────────────────────
    pivot_total = grid.pivot(index="sl_mult", columns="rr", values="total_r").round(1)
    pivot_wr    = grid.pivot(index="sl_mult", columns="rr", values="wr").round(1)
    pivot_avg   = grid.pivot(index="sl_mult", columns="rr", values="avg_r").round(3)
    pivot_pf    = grid.pivot(index="sl_mult", columns="rr", values="pf").round(2)

    out_lines = []
    def emit(line: str = ""):
        print(line)
        out_lines.append(line)

    emit("=" * 80)
    emit(f"GEOMETRY GRID — {len(df):,} trades")
    emit(f"  filters applied: {args.apply_phase1_filters}")
    emit(f"  spread cost:     {args.spread_pips} pip")
    emit(f"  max hold:        {args.max_hold} bars")
    emit("=" * 80)
    emit()
    emit("─── Total R (rows = SL multiplier × ATR; cols = RR ratio) ───")
    emit(pivot_total.to_string())
    emit()
    emit("─── Win rate % ───")
    emit(pivot_wr.to_string())
    emit()
    emit("─── Avg R per trade ───")
    emit(pivot_avg.to_string())
    emit()
    emit("─── Profit factor ───")
    emit(pivot_pf.to_string())
    emit()

    # ── Train/Test discipline on top cells ──────────────────────────
    grid_sorted = grid.sort_values("total_r", ascending=False)
    emit("=" * 80)
    emit("TOP 5 CELLS (by full-period total R) — with TRAIN / TEST split discipline")
    emit("=" * 80)
    for _, cell in grid_sorted.head(5).iterrows():
        s = split_summary(cell)
        emit(
            f"  SL={cell['sl_mult']:.1f}×ATR  RR=1:{cell['rr']:.1f}  "
            f"|  full: {cell['n']} trades  WR {cell['wr']:.1f}%  "
            f"avgR {cell['avg_r']:+.3f}  totalR {cell['total_r']:+.1f}  "
            f"PF {cell['pf']:.2f}"
        )
        emit(
            f"      train ({s['train_n']}):  "
            f"WR {s['train_wr']:.1f}%  avgR {s['train_avg']:+.3f}  totalR {s['train_total']:+.1f}"
        )
        emit(
            f"      test  ({s['test_n']}):   "
            f"WR {s['test_wr']:.1f}%  avgR {s['test_avg']:+.3f}  totalR {s['test_total']:+.1f}"
        )
        emit()

    # ── Bottom 3 for contrast ───────────────────────────────────────
    emit("─── BOTTOM 3 cells (worst total R) ───")
    for _, cell in grid_sorted.tail(3).iterrows():
        emit(
            f"  SL={cell['sl_mult']:.1f}×ATR  RR=1:{cell['rr']:.1f}  "
            f"totalR {cell['total_r']:+.1f}  WR {cell['wr']:.1f}%"
        )
    emit()

    # ── Verdict ─────────────────────────────────────────────────────
    emit("=" * 80)
    emit("VERDICT")
    emit("=" * 80)
    best = grid_sorted.iloc[0]
    s = split_summary(best)
    emit(
        f"  Best cell: SL={best['sl_mult']:.1f}×ATR, RR=1:{best['rr']:.1f}  "
        f"(full totalR {best['total_r']:+.1f})"
    )
    emit(
        f"  Test slice: {s['test_n']} trades, WR {s['test_wr']:.1f}%, "
        f"totalR {s['test_total']:+.1f}, avgR {s['test_avg']:+.3f}"
    )
    if s["test_total"] > 0:
        emit("  ✅ Best geometry is profitable on held-out test set.")
        emit("     → Phase 3A: hardcode in _generate_trade_setup, re-run end-to-end.")
    elif s["test_avg"] > -0.05:
        emit("  ⚠️  Best geometry is near-break-even on test set.")
        emit("     → Phase 3A worth attempting; combine with Phase 1 filters for margin.")
    else:
        emit("  ❌ No geometry is profitable on held-out test set.")
        emit("     → Phase 3B: M15 confluence is structurally not edge-bearing.")
        emit("        Consider H4 primary, mean-reversion-only, or indicator-level meta-model.")

    Path(args.out).parent.mkdir(parents=True, exist_ok=True)
    Path(args.out).write_text("\n".join(out_lines))
    emit()
    emit(f"📁 Saved: {args.out}")
    return 0


if __name__ == "__main__":
    sys.exit(main())