"""
Backtest: static grid vs adaptive grid (with Kronos advisor).

Использование:
    python -m kronos.backtest --csv ./data/BTCUSDT-1h.csv --tf 60 --step 0.005 --levels 10

Что делает:
  1) Загружает историю.
  2) На каждом шаге t (от lookback до конца истории):
       - запускает Kronos на окне [t-lookback .. t] с прогнозом на pred_len свечей
       - вычисляет adaptive step
       - симулирует сделку на следующих M свечей:
            * static  — фиксированный шаг, центр = close[t]
            * adaptive — шаг от Kronos, центр = close[t] + center_offset
  3) Сравнивает метрики: PnL, max drawdown, число сделок, % прибыльных.
"""

from __future__ import annotations

import argparse
import logging
import math
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

import numpy as np
import pandas as pd

logger = logging.getLogger(__name__)


@dataclass
class BTResult:
    name: str
    total_pnl_pct: float
    n_trades: int
    winrate: float
    max_dd_pct: float
    sharpe: float


def simulate_grid(
    df: pd.DataFrame,
    start: int,
    horizon: int,
    step_pct: float,
    levels: int,
    qty: float = 0.001,
    center_offset_pct: float = 0.0,
    fee_pct: float = 0.001,   # 0.1% taker fee per order (типично для Tradernet)
) -> tuple[float, int, int]:
    """
    Простая сеточная симуляция с реалистичной логикой матчинга.

    Правила:
      - Одна сделка на свечу (BUY или SELL, не обе).
      - BUY разрешён, если pos_open is None.
      - SELL разрешён, если pos_open is not None.
      - При исполнении снимаем комиссию (fee_pct × qty × price).
      - Возвращаем (pnl_pct, n_roundtrips, n_winning_roundtrips).
    """
    if start >= len(df) - 1:
        return 0.0, 0, 0

    center = float(df["close"].iloc[start]) * (1.0 + center_offset_pct)
    step = center * step_pct
    if step <= 0:
        return 0.0, 0, 0

    buy_levels = sorted([center - n * step for n in range(1, levels + 1) if center - n * step > 0])
    sell_levels = sorted([center + n * step for n in range(1, levels + 1)])

    end = min(start + horizon, len(df))
    pnl = 0.0
    n_roundtrips = 0
    n_wins = 0
    pos_open: Optional[float] = None

    for i in range(start + 1, end):
        hi = float(df["high"].iloc[i])
        lo = float(df["low"].iloc[i])

        if pos_open is None:
            # Ищем ближайший BUY-уровень, который достигнут
            for bl in buy_levels:
                if lo <= bl:
                    # fee на вход
                    pnl -= bl * qty * fee_pct
                    pos_open = bl
                    break
        else:
            # Ищем ближайший SELL-уровень, который достигнут
            for sl in sell_levels:
                if hi >= sl:
                    gross = (sl - pos_open) * qty
                    fee = sl * qty * fee_pct + pos_open * qty * fee_pct
                    net = gross - fee
                    pnl += net
                    if net > 0:
                        n_wins += 1
                    n_roundtrips += 1
                    pos_open = None
                    break

    # PnL в долях от notional (center * qty)
    pnl_pct = pnl / (center * qty) if center > 0 else 0.0
    return pnl_pct, n_roundtrips, n_wins


def run_backtest(
    df: pd.DataFrame,
    tf_min: int,
    lookback: int,
    pred_len: int,
    static_step_pct: float,
    levels: int,
    advisor=None,
    stride: int = 24,
    regime_window_hours: int = 24 * 30,  # окно для классификации режима (тренд/боковик)
) -> list[BTResult]:
    """
    Прогоняет статическую и (если задан advisor) адаптивную сетки.
    stride — через сколько свечей делаем новый «тик» сетки.

    Дополнительно возвращает _raw_arrays в results[0].__dict__:
      - regimes_at_t: список ('trending'|'sideways') для каждого тика
      - static_pnls_per_t, adapt_pnls_per_t
    """
    n = len(df)
    if n < lookback + pred_len + 1:
        raise ValueError(f"need at least {lookback + pred_len + 1} candles, got {n}")

    static_pnls, static_n, static_w = [], [], []
    adapt_pnls, adapt_n, adapt_w = [], [], []
    regimes = []
    pause_hits = 0

    peak_equity_s, peak_equity_a = 1.0, 1.0
    eq_s, eq_a = 1.0, 1.0
    max_dd_s, max_dd_a = 0.0, 0.0

    for t in range(lookback, n - pred_len, stride):
        horizon = min(pred_len, stride * 2)
        center = float(df["close"].iloc[t])

        # Режим рынка: в окне [t-regime_window, t] считаем |return|
        win_start = max(0, t - regime_window_hours)
        win = df["close"].iloc[win_start:t + 1]
        ret = (win.iloc[-1] - win.iloc[0]) / win.iloc[0] if len(win) > 1 and win.iloc[0] > 0 else 0.0
        regime = "trending" if abs(ret) > 0.05 else "sideways"   # ±5% за месяц → тренд
        regimes.append(regime)

        # STATIC
        s_pnl, s_nt, s_w = simulate_grid(
            df, t, horizon,
            step_pct=static_step_pct,
            levels=levels,
            center_offset_pct=0.0,
        )
        static_pnls.append(s_pnl)
        static_n.append(s_nt)
        static_w.append(s_w)
        eq_s *= (1.0 + s_pnl)
        peak_equity_s = max(peak_equity_s, eq_s)
        dd = (eq_s - peak_equity_s) / peak_equity_s
        max_dd_s = min(max_dd_s, dd)

        # ADAPTIVE
        if advisor is not None:
            try:
                advice = advisor.advise(df.iloc[:t + 1], lookback=lookback, pred_len=pred_len, tf_min=tf_min)
                if advice.pause_grid and advice.confidence > 0.5:
                    pause_hits += 1
                    adapt_pnls.append(0.0)
                    adapt_n.append(0)
                    adapt_w.append(0)
                else:
                    a_pnl, a_nt, a_w = simulate_grid(
                        df, t, horizon,
                        step_pct=max(static_step_pct * 0.5, advice.step_percent),
                        levels=levels,
                        center_offset_pct=advice.center_offset_pct,
                    )
                    adapt_pnls.append(a_pnl)
                    adapt_n.append(a_nt)
                    adapt_w.append(a_w)
                    eq_a *= (1.0 + a_pnl)
                    peak_equity_a = max(peak_equity_a, eq_a)
                    dd = (eq_a - peak_equity_a) / peak_equity_a
                    max_dd_a = min(max_dd_a, dd)
            except Exception as e:
                logger.warning("Advisor failed at t=%d: %s", t, e)
                continue

    def _agg(name, pnls, ns, ws, dd):
        if not pnls:
            return BTResult(name, 0.0, 0, 0.0, 0.0, 0.0)
        total = float(np.prod(1.0 + np.array(pnls)) - 1.0)
        n_total = sum(ns)
        n_w = sum(ws)
        wr = n_w / n_total if n_total else 0.0
        if len(pnls) > 1:
            sharpe = float(np.mean(pnls) / (np.std(pnls) + 1e-9) * math.sqrt(252 / max(1, len(pnls) // 24)))
        else:
            sharpe = 0.0
        return BTResult(name, total, n_total, wr, dd, sharpe)

    results = [_agg("static", static_pnls, static_n, static_w, max_dd_s)]
    if advisor is not None:
        results.append(_agg("adaptive", adapt_pnls, adapt_n, adapt_w, max_dd_a))
        # Метрики по режимам
        for regime in ("trending", "sideways"):
            mask = [i for i, r in enumerate(regimes[:len(static_pnls)]) if r == regime]
            if not mask:
                continue
            s_pnls_r = [static_pnls[i] for i in mask]
            a_pnls_r = [adapt_pnls[i] for i in mask]
            s_total = float(np.prod(1.0 + np.array(s_pnls_r)) - 1.0) if s_pnls_r else 0.0
            a_total = float(np.prod(1.0 + np.array(a_pnls_r)) - 1.0) if a_pnls_r else 0.0
            logger.info(
                "regime=%-9s n_ticks=%-4d static_pnl=%+7.2f%%  adaptive_pnl=%+7.2f%%  delta=%+.2f%%  pause_hits_total=%d",
                regime, len(mask), s_total * 100, a_total * 100,
                (a_total - s_total) * 100, pause_hits,
            )
    return results


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--csv", required=True, help="CSV с колонками open,high,low,close,volume")
    ap.add_argument("--tf", type=int, default=60)
    ap.add_argument("--lookback", type=int, default=400)
    ap.add_argument("--pred-len", type=int, default=24)
    ap.add_argument("--step", type=float, default=0.005)
    ap.add_argument("--levels", type=int, default=10)
    ap.add_argument("--stride", type=int, default=24)
    ap.add_argument("--with-advisor", action="store_true")
    ap.add_argument("--model", default="kronos-mini")
    ap.add_argument("--device", default="cpu")
    ap.add_argument("--out", default=None)
    args = ap.parse_args()

    logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
    log = logging.getLogger("backtest")

    df = pd.read_csv(args.csv)
    # нормализуем имена колонок
    df.columns = [c.strip().lower() for c in df.columns]
    if "timestamps" in df.columns:
        df["timestamps"] = pd.to_datetime(df["timestamps"])
        df = df.set_index("timestamps")
    elif "timestamp" in df.columns:
        df["timestamps"] = pd.to_datetime(df["timestamp"])
        df = df.set_index("timestamps")
    elif "open_time" in df.columns:
        df.index = pd.to_datetime(df["open_time"], unit="ms", utc=True).tz_convert(None)
    elif "date" in df.columns:
        df["timestamps"] = pd.to_datetime(df["date"])
        df = df.set_index("timestamps")
    df = df[["open", "high", "low", "close", "volume"]].astype(float).sort_index()
    log.info("Loaded %d candles, last close=%.2f", len(df), float(df["close"].iloc[-1]))

    advisor = None
    if args.with_advisor:
        from kronos import KronosAdvisor
        advisor = KronosAdvisor(model_name=args.model, device=args.device)
        log.info("Loaded advisor %s", args.model)

    results = run_backtest(
        df,
        tf_min=args.tf,
        lookback=args.lookback,
        pred_len=args.pred_len,
        static_step_pct=args.step,
        levels=args.levels,
        advisor=advisor,
        stride=args.stride,
    )

    print()
    print("=" * 72)
    print(f"{'MODE':<12} {'PnL %':>10} {'#Trades':>10} {'WinRate':>10} {'MaxDD %':>10} {'Sharpe':>10}")
    print("-" * 72)
    for r in results:
        print(f"{r.name:<12} {r.total_pnl_pct*100:>9.2f}% {r.n_trades:>10d} "
              f"{r.winrate*100:>9.1f}% {r.max_dd_pct*100:>9.2f}% {r.sharpe:>9.2f}")
    print("=" * 72)

    if args.out:
        import json
        Path(args.out).write_text(json.dumps([r.__dict__ for r in results], indent=2))
        log.info("Wrote %s", args.out)


if __name__ == "__main__":
    main()