#!/usr/bin/env python3
"""
Эксперимент: SL=3×ATR, TP=6×ATR (RR 1:2) vs текущий SL=1.5×ATR, TP=1.5×ATR (RR 1:1).

Запуск:
  python3 experiment_rr.py                           # обучить + бэктест
  python3 experiment_rr.py --backtest-only           # только бэктест (по сохранённой модели)
  python3 experiment_rr.py --ticker SBER --compare   # сравнить обе модели

Сравнение:
  - val_acc, WinRate, Sharpe, MaxDD, Total Return
"""
import sys, os, time, json, warnings, importlib
import numpy as np
import pandas as pd
from datetime import datetime

warnings.filterwarnings('ignore')
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import config
config.set_device('train')  # CUDA for training

from models.moe import MultiTimeframeMoE, _prepare_df, _get_feature_cols, _build_dataset, _flat_mask
from models.experts import ExpertEnsemble, HORIZONS
import xgboost as xgb

SAVE_DIR = config.SAVE_DIR
TICKER = 'SBER'

# ── Конфиги для сравнения ──
CONFIGS = {
    'rr1x1': {  # текущий
        'atr_mult_sl': 1.5,
        'atr_mult_tp': 1.5,
        'max_bars': 100,
        'label': 'RR 1:1 (SL=1.5×ATR, TP=1.5×ATR)',
        'suffix': '',
    },
    'rr1x2': {  # экспериментальный
        'atr_mult_sl': 3.0,
        'atr_mult_tp': 6.0,
        'max_bars': 200,
        'label': 'RR 1:2 (SL=3×ATR, TP=6×ATR)',
        'suffix': '_rr1x2',
    },
}

TARGET_KEYS = ['atr_mult_sl', 'atr_mult_tp', 'max_bars']  # только эти ключи уходят в TARGET_CONFIG


def train_with_config(ticker: str, target_cfg: dict, suffix: str) -> dict | None:
    """Тренирует MoE с заданным TARGET_CONFIG."""
    import config as cfg_module

    # Патчим конфиг IN-PLACE (только ключи для compute_dual_outcomes)
    global TARGET_KEYS
    old_cfg = cfg_module.TARGET_CONFIG.copy()
    cfg_module.TARGET_CONFIG.clear()
    cfg_module.TARGET_CONFIG.update({k: target_cfg[k] for k in TARGET_KEYS})

    print(f'\n{"="*60}')
    print(f'  ОБУЧЕНИЕ: {ticker} {target_cfg["label"]}')
    print(f'  SL={target_cfg["atr_mult_sl"]}×ATR, TP={target_cfg["atr_mult_tp"]}×ATR, max_bars={target_cfg["max_bars"]}')
    print(f'{"="*60}\n')

    t0 = time.time()

    # ── OOS фаза ──
    moe = MultiTimeframeMoE(ticker)
    ok = moe.train(limit=None, verbose=True, retrain_on_full=False)
    if not ok:
        print(f'  [ERROR] train() вернул False')
        cfg_module.TARGET_CONFIG = old_cfg
        moe_module.TARGET_CONFIG = old_moe_target
        return None

    oos_val_acc = getattr(moe, 'oos_val_acc', 0)
    oos_long_th = moe.optimal_thresholds.get('long', 0.55)
    oos_short_th = moe.optimal_thresholds.get('short', 0.55)
    oos_selected = moe.selected_experts

    print(f'\n  OOS: val_acc={oos_val_acc:.1%}, L≥{oos_long_th:.2f}, S≥{oos_short_th:.2f}')

    # ── Retrain на 100% ──
    print(f'\n  Retrain на 100% данных с {target_cfg["label"]}...')
    df_full = _prepare_df(ticker, limit=None)
    if df_full is None or len(df_full) < 200:
        print(f'  [ERROR] недостаточно данных: {len(df_full) if df_full is not None else 0}')
        cfg_module.TARGET_CONFIG = old_cfg
        moe_module.TARGET_CONFIG = old_moe_target
        return None

    all_features = _get_feature_cols()
    available = [c for c in all_features if c in df_full.columns]
    norm_stats = moe.ensemble.get_norm_stats()

    ensemble_full = ExpertEnsemble(n_features=len(available))
    ensemble_full.train_all(df_full, available, verbose=True, epochs=None)
    ensemble_full.set_norm_stats(*norm_stats)

    signals_full = ensemble_full.predict_all(df_full, available)
    X_full = _build_dataset(df_full, signals_full, ensemble_full, available,
                            keep_experts=oos_selected)
    y_full = np.column_stack([df_full['outcome_long'].values, df_full['outcome_short'].values])
    min_len = min(len(X_full), len(y_full))
    X_full, y_full = X_full[-min_len:], y_full[-min_len:]

    # Флет-фильтр
    flat_mask = _flat_mask(df_full).values[-len(X_full):]
    n_flat = int(flat_mask.sum())
    y_full[flat_mask] = 0.0
    print(f'  Флет-фильтр: {n_flat}/{len(flat_mask)} ({n_flat/len(flat_mask):.1%})')

    # Исходные success rate
    long_sr = df_full['outcome_long'].mean()
    short_sr = df_full['outcome_short'].mean()
    print(f'  Success rate: long={long_sr:.1%}, short={short_sr:.1%}')

    params_xgb = {
        'max_depth': 4, 'learning_rate': 0.03, 'n_estimators': 300,
        'subsample': 0.7, 'colsample_bytree': 0.7,
        'reg_alpha': 0.1, 'reg_lambda': 1.0, 'min_child_weight': 5,
        'random_state': 42, 'verbosity': 0, 'n_jobs': -1,
    }

    pl = params_xgb.copy()
    pos_l = y_full[:, 0].sum()
    neg_l = len(y_full) - pos_l
    pl['scale_pos_weight'] = min(max(neg_l / max(pos_l, 1), 1.0), 5.0)
    xgb_long = xgb.XGBClassifier(**pl)
    xgb_long.fit(X_full, y_full[:, 0])

    ps = params_xgb.copy()
    pos_s = y_full[:, 1].sum()
    neg_s = len(y_full) - pos_s
    scale_s = min(max(neg_s / max(pos_s, 1), 1.0), 5.0)
    scale_s *= config.TICKER_SHORT_WEIGHTS.get(ticker, 3.0)
    ps['scale_pos_weight'] = scale_s
    xgb_short = xgb.XGBClassifier(**ps)
    xgb_short.fit(X_full, y_full[:, 1])

    # Сборка финальной модели
    moe.ensemble = ensemble_full
    moe.rf_long = xgb_long
    moe.rf_short = xgb_short
    moe.rf_model = None
    moe.optimal_thresholds = {'long': oos_long_th, 'short': oos_short_th}
    moe.selected_experts = oos_selected
    moe.oos_val_acc = oos_val_acc
    moe.val_acc = oos_val_acc
    moe.feature_cols = available
    moe.n_base_features = len(available)

    # Сохраняем
    save_path = os.path.join(SAVE_DIR, f'{ticker.lower()}_moe_v12{suffix}.joblib')
    moe.save(save_path)

    elapsed = time.time() - t0
    print(f'\n  ✓ Сохранено: {save_path} ({elapsed/60:.1f} мин)')

    # Восстанавливаем конфиг (сохраняем только нужные ключи)
    cfg_module.TARGET_CONFIG.clear()
    cfg_module.TARGET_CONFIG.update({k: old_cfg[k] for k in TARGET_KEYS})

    return {
        'status': 'ok',
        'ticker': ticker,
        'config': target_cfg['label'],
        'suffix': suffix,
        'rows': len(df_full),
        'val_acc': oos_val_acc,
        'long_th': oos_long_th,
        'short_th': oos_short_th,
        'long_sr': float(long_sr),
        'short_sr': float(short_sr),
        'flat_pct': n_flat / len(flat_mask) * 100,
        'elapsed_min': round(elapsed / 60, 1),
        'path': save_path,
    }


def backtest_model(ticker: str, suffix: str, target_cfg: dict) -> dict:
    """Загружает обученную модель и проводит бэктест по всей истории."""
    print(f'\n  ── БЭКТЕСТ: {target_cfg["label"]} ──')

    model_path = os.path.join(SAVE_DIR, f'{ticker.lower()}_moe_v12{suffix}.joblib')
    if not os.path.exists(model_path):
        print(f'  [WARN] Модель не найдена: {model_path}')
        return {'status': 'error', 'error': 'model_not_found'}

    # Загружаем модель
    moe = MultiTimeframeMoE(ticker)
    ok = moe.load(model_path)
    if not ok:
        return {'status': 'error', 'error': 'load_failed'}

    # Предсказания на всей истории
    df = _prepare_df(ticker)
    if df is None:
        return {'status': 'error', 'error': 'no_data'}

    result = moe.predict_proba_aligned(df)
    if result is None or len(result) < 100:
        return {'status': 'error', 'error': 'no_predictions'}

    # Загружаем цены для ATR и close (убеждаемся что длина совпадает)
    n_result = len(result)
    close_vals = df['Close'].values[-n_result:]
    atr_vals = df['atr'].values[-n_result:]
    high_vals = df['High'].values[-n_result:]
    low_vals = df['Low'].values[-n_result:]

    # Симулируем сделки по сигналам
    sl_mult = target_cfg['atr_mult_sl']
    tp_mult = target_cfg['atr_mult_tp']
    max_bars = target_cfg['max_bars']

    trades = []
    in_trade = False
    entry_price = 0.0
    entry_atr = 0.0
    entry_idx = 0
    direction = ''
    sl_price = 0.0
    tp_price = 0.0

    for i in range(n_result):
        signal = result.iloc[i]['signal']
        close = close_vals[i]
        atr = atr_vals[i]

        if in_trade:
            bars_held = i - entry_idx
            high = high_vals[i]
            low = low_vals[i]

            # Проверка TP/SL
            hit = None
            exit_price = close

            if direction == 'LONG':
                if high >= tp_price:
                    hit = 'TP'
                    exit_price = tp_price
                elif low <= sl_price:
                    hit = 'SL'
                    exit_price = sl_price
            else:  # SHORT
                if low <= tp_price:
                    hit = 'TP'
                    exit_price = tp_price
                elif high >= sl_price:
                    hit = 'SL'
                    exit_price = sl_price

            # Time-stop
            if hit is None and bars_held >= max_bars:
                hit = 'TIME_STOP'
                exit_price = close

            if hit:
                if direction == 'LONG':
                    pnl = (exit_price - entry_price) / entry_price
                else:
                    pnl = (entry_price - exit_price) / entry_price
                trades.append({
                    'direction': direction,
                    'entry': entry_price,
                    'exit': exit_price,
                    'pnl_pct': pnl * 100,
                    'reason': hit,
                    'bars_held': bars_held,
                })
                in_trade = False

        if not in_trade and signal in ('BUY', 'SELL'):
            # Открываем сделку
            direction = 'LONG' if signal == 'BUY' else 'SHORT'
            entry_price = close
            entry_atr = atr
            entry_idx = i
            in_trade = True

            sl_dist = atr * sl_mult
            tp_dist = atr * tp_mult
            if direction == 'LONG':
                sl_price = close - sl_dist
                tp_price = close + tp_dist
            else:
                sl_price = close + sl_dist
                tp_price = close - tp_dist

    # Статистика
    total = len(trades)
    wins = sum(1 for t in trades if t['pnl_pct'] > 0)
    losses = sum(1 for t in trades if t['pnl_pct'] < 0)
    winrate = wins / max(total, 1) * 100
    total_return = sum(t['pnl_pct'] for t in trades)
    avg_win = np.mean([t['pnl_pct'] for t in trades if t['pnl_pct'] > 0]) if wins else 0
    avg_loss = np.mean([t['pnl_pct'] for t in trades if t['pnl_pct'] < 0]) if losses else 0
    profit_factor = abs(sum(t['pnl_pct'] for t in trades if t['pnl_pct'] > 0) / 
                        max(abs(sum(t['pnl_pct'] for t in trades if t['pnl_pct'] < 0)), 0.01))

    # Sharpe Ratio (approximation)
    returns = np.array([t['pnl_pct'] for t in trades])
    sharpe = returns.mean() / max(returns.std(), 0.01) * np.sqrt(365 * 24 / np.mean([t['bars_held'] for t in trades])) if total > 1 else 0

    tp_trades = sum(1 for t in trades if t['reason'] == 'TP')
    sl_trades = sum(1 for t in trades if t['reason'] == 'SL')
    ts_trades = sum(1 for t in trades if t['reason'] == 'TIME_STOP')

    print(f'  Всего сделок: {total}')
    print(f'  WinRate: {winrate:.1f}% ({wins}/{total})')
    print(f'  TP: {tp_trades}, SL: {sl_trades}, TimeStop: {ts_trades}')
    print(f'  Суммарная доходность: {total_return:+.2f}%')
    print(f'  Средняя прибыль: {avg_win:.2f}% / Средний убыток: {avg_loss:.2f}%')
    print(f'  Profit Factor: {profit_factor:.2f}')
    print(f'  Sharpe (год): {sharpe:.2f}')

    return {
        'status': 'ok',
        'ticker': ticker,
        'config': target_cfg['label'],
        'total_trades': total,
        'winrate': winrate,
        'wins': wins,
        'losses': losses,
        'tp': tp_trades,
        'sl': sl_trades,
        'time_stop': ts_trades,
        'total_return': total_return,
        'avg_win': avg_win,
        'avg_loss': avg_loss,
        'profit_factor': profit_factor,
        'sharpe': sharpe,
    }


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description='RR Experiment: 1:1 vs 1:2')
    parser.add_argument('--ticker', default='SBER', help='Тикер для эксперимента')
    parser.add_argument('--compare', action='store_true', help='Сравнить обе модели (если обучены)')
    parser.add_argument('--train-rr1x2', action='store_true', help='Обучить только RR 1:2')
    parser.add_argument('--train-rr1x1', action='store_true', help='Обучить только RR 1:1')
    args = parser.parse_args()

    TICKER = args.ticker.upper()

    results_train = {}
    results_bt = {}

    # Определяем, что обучать
    train_rr1x1 = args.train_rr1x1 or (not args.train_rr1x2 and not args.compare)
    train_rr1x2 = args.train_rr1x2 or (not args.train_rr1x1 and not args.compare)

    if not args.compare or train_rr1x1 or train_rr1x2:
        print(f'\n{"█"*70}')
        print(f'  ЭКСПЕРИМЕНТ RR: {TICKER}')
        print(f'  Сравнение: RR 1:1 (SL=1.5×ATR, TP=1.5×ATR) vs RR 1:2 (SL=3×ATR, TP=6×ATR)')
        print(f'  Дата: {datetime.now().strftime("%d.%m.%Y %H:%M")}')
        print(f'{"█"*70}\n')

    # ── ОБУЧЕНИЕ ──
    for key, cfg in CONFIGS.items():
        if (key == 'rr1x1' and not train_rr1x1) or (key == 'rr1x2' and not train_rr1x2):
            continue

        r = train_with_config(TICKER, cfg, cfg['suffix'])
        results_train[key] = r
        if r:
            results_bt[key] = backtest_model(TICKER, cfg['suffix'], cfg)
        else:
            results_bt[key] = {'status': 'error'}

    # ── СРАВНЕНИЕ ──
    if args.compare or (train_rr1x1 and train_rr1x2):
        print(f'\n{"="*70}')
        print(f'  СРАВНЕНИЕ РЕЗУЛЬТАТОВ: {TICKER}')
        print(f'{"="*70}')
        print(f'\n  {"Параметр":<30} {"RR 1:1 (1.5×ATR)":<25} {"RR 1:2 (3×/6×ATR)":<25}')
        print(f'  {"-"*30} {"-"*25} {"-"*25}')

        for key in ['rr1x1', 'rr1x2']:
            if key not in results_bt:
                continue

        row = results_bt.get('rr1x1', {})
        row2 = results_bt.get('rr1x2', {})

        for param, fmt in [
            ('val_acc', '.1%'), ('Success Rate Long', '.1%'), ('Success Rate Short', '.1%'),
            ('Всего сделок', '.0f'), ('WinRate', '.1f'), ('TP сделок', '.0f'), ('SL сделок', '.0f'),
            ('TimeStop', '.0f'), ('Сумм доходность', '.2f'), ('Profit Factor', '.2f'),
            ('Sharpe (год)', '.2f'),
        ]:
            v1 = row.get(param, '—') if isinstance(row, dict) else '—'
            v2 = row2.get(param, '—') if isinstance(row2, dict) else '—'
            if isinstance(v1, float) and isinstance(v2, float) and param in ('val_acc',):
                # особый случай
                pass
            print(f'  {param:<30} {v1 if isinstance(v1, str) else f"{v1:{fmt}}":<25} '
                  f'{v2 if isinstance(v2, str) else f"{v2:{fmt}}":<25}')

        # Показываем упрощённую таблицу
        print()
        for bt_key, label in [('rr1x1', 'RR 1:1'), ('rr1x2', 'RR 1:2')]:
            bt = results_bt.get(bt_key, {})
            tr = results_train.get(bt_key, {})
            if bt.get('status') == 'ok':
                print(f'  {label:>8}: val_acc={tr.get("val_acc","N/A")} '
                      f'| сделок={bt["total_trades"]} '
                      f'| WR={bt["winrate"]:.1f}% '
                      f'| return={bt["total_return"]:+.2f}% '
                      f'| PF={bt["profit_factor"]:.2f} '
                      f'| Sharpe={bt["sharpe"]:.2f}')

    print()