#!/usr/bin/env python3
"""
Walk-forward validation (7 folds) for MOEX instruments.

Usage:
    python run_walkforward.py                    # all tickers × all TFs
    python run_walkforward.py --ticker SBER      # single ticker
    python run_walkforward.py --ticker SBER --tf D1 --folds 5
"""
import sys, os, warnings, time, json
import numpy as np
import pandas as pd
warnings.filterwarnings('ignore')

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import config
config.set_device('auto')  # auto-detect device

from config import MOEX_TICKERS, TARGET_CONFIG
from features.technical import FEATURE_COLS
from features.directional import DIRECTIONAL_COLS
from features.context import CONTEXT_COLS
from features.technical import INTERACTION_COLS, TEMPORAL_COLS
from sklearn.ensemble import RandomForestClassifier
from sklearn.multioutput import MultiOutputClassifier
# Purged walk-forward: manual splits with embargo gap

N_FOLDS = 7
TIMEFRAMES = ['H1', 'D1', 'W1']

def get_directional_feature_cols():
    cols = list(dict.fromkeys(DIRECTIONAL_COLS + FEATURE_COLS + CONTEXT_COLS + INTERACTION_COLS + TEMPORAL_COLS))
    return cols

def _prepare_df(ticker, tf, limit=None):
    """Подготавливает DataFrame для walk-forward валидации.
    Делегирует в features.pipeline.prepare_df()."""
    from features.pipeline import prepare_df as unified_prepare_df
    return unified_prepare_df(
        ticker, tf,
        with_mtf=False,
        with_short_specific=False,  # legacy — не использует short_specific
        with_targets=True,
        with_crypto_features=False,
        limit=limit,
        min_rows=50,
        shift_mtf=None,
        winsor_bounds=None,
        log_nan_conversion=False,
        include_str_dtype=False,
        reset_index=True,  # run_walkforward требует reset_index
    )

def run_walkforward(ticker, tf, folds=N_FOLDS):
    """Walk-forward: обучаем на растущем окне, валидируем на следующих 1/fold данных.
    
    Используется purged gap (MAX_BARS) между train и val для предотвращения
    boundary label leakage (Lopez de Prado).
    """
    df = _prepare_df(ticker, tf)
    if df is None or len(df) < 200:
        print(f'  ⏭ Too few rows: {len(df) if df is not None else 0}')
        return None

    MAX_BARS = TARGET_CONFIG.get('max_bars', 100)
    feature_cols = get_directional_feature_cols()
    available = [c for c in feature_cols if c in df.columns]
    X = df[available].values.astype(np.float32)
    X = np.nan_to_num(X, nan=0.0)
    y = np.column_stack([df['outcome_long'].values, df['outcome_short'].values])

    n = len(X)
    fold_size = max(int(n / (folds + 1)), MAX_BARS + 50)
    fold_results = []

    for fold in range(folds):
        val_start = (fold + 1) * fold_size
        val_end = min(val_start + fold_size, n)
        train_end = val_start - MAX_BARS  # purged gap
        train_start = 0
        train_idx = list(range(train_start, train_end))
        val_idx = list(range(val_start, val_end))
        if len(train_idx) < 50 or len(val_idx) < 20:
            continue

        X_train, X_val = X[train_idx], X[val_idx]
        y_train, y_val = y[train_idx], y[val_idx]

        model = MultiOutputClassifier(
            RandomForestClassifier(n_estimators=200, max_depth=10,
                                   n_jobs=-1, random_state=42,
                                   class_weight='balanced')
        )
        model.fit(X_train, y_train)

        # Evaluate
        y_pred = model.predict(X_val)
        long_acc = float(np.mean(y_pred[:, 0] == y_val[:, 0]))
        short_acc = float(np.mean(y_pred[:, 1] == y_val[:, 1]))

        # HiConf
        y_proba = model.predict_proba(X_val)
        if isinstance(y_proba, list):
            def _sp(p):
                return p[:, 1] if p.shape[1] > 1 else p[:, 0]
            p_long = _sp(y_proba[0])
            p_short = _sp(y_proba[1])
        else:
            p_long = y_proba[:, 1] if y_proba.shape[1] > 1 else y_proba[:, 0]
            p_short = 1.0 - p_long

        hi_conf = (np.maximum(p_long, p_short) >= 0.6)
        hi_conf_wr = 0.0
        hi_conf_n = 0
        if hi_conf.sum() > 0:
            hi_long = p_long[hi_conf] > p_short[hi_conf]
            wins = np.where(hi_long, y_val[hi_conf, 0], y_val[hi_conf, 1])
            hi_conf_wr = float(wins.mean())
            hi_conf_n = int(hi_conf.sum())

        fold_results.append({
            'fold': fold + 1,
            'train_size': len(X_train),
            'val_size': len(X_val),
            'long_acc': long_acc,
            'short_acc': short_acc,
            'hi_conf_wr': hi_conf_wr,
            'hi_conf_n': hi_conf_n,
        })

    if not fold_results:
        return None

    summary = {
        'ticker': ticker,
        'timeframe': tf,
        'folds': len(fold_results),
        'total_samples': len(df),
        'mean_long_acc': float(np.mean([r['long_acc'] for r in fold_results])),
        'mean_short_acc': float(np.mean([r['short_acc'] for r in fold_results])),
        'mean_hi_conf_wr': float(np.mean([r['hi_conf_wr'] for r in fold_results if r['hi_conf_n'] > 0])),
        'total_hi_conf': int(sum(r['hi_conf_n'] for r in fold_results)),
        'fold_results': fold_results,
    }
    return summary


# ================================================================
if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description='Walk-forward validation')
    parser.add_argument('--ticker', type=str, default=None, help='Single ticker')
    parser.add_argument('--tf', type=str, default=None, help='Single timeframe (H1/D1/W1)')
    parser.add_argument('--folds', type=int, default=N_FOLDS, help='Number of folds')
    args = parser.parse_args()

    tickers = [args.ticker] if args.ticker else MOEX_TICKERS
    tfs = [args.tf] if args.tf else TIMEFRAMES

    print(f'\n{"="*70}')
    print(f'  WALK-FORWARD VALIDATION — {args.folds} folds')
    print(f'{"="*70}')
    print()

    all_results = []
    t_start = time.time()

    for ticker in tickers:
        for tf in tfs:
            print(f'{ticker:6s} {tf:2s}: ', end='', flush=True)
            t0 = time.time()
            summary = run_walkforward(ticker, tf, folds=args.folds)
            elapsed = time.time() - t0

            if summary is None:
                print('⏭ skipped')
                continue

            print(f'{summary["folds"]} folds | '
                  f'L={summary["mean_long_acc"]:.1%} S={summary["mean_short_acc"]:.1%} '
                  f'HiConf({summary["total_hi_conf"]})={summary["mean_hi_conf_wr"]:.1%} '
                  f'({elapsed:.1f}s)')
            all_results.append(summary)

    # Final table
    print(f'\n{"="*70}')
    print('  СВОДНАЯ ТАБЛИЦА')
    print(f'{"="*70}')
    print(f'  {"Ticker":6s} {"TF":2s} {"Folds":5s} {"Long":8s} {"Short":8s} {"HiConf WR":9s} {"N HiConf":8s} {"Samples":7s}')
    print('  ' + '-'*60)

    for r in sorted(all_results, key=lambda x: (x['ticker'], x['timeframe'])):
        print(f'  {r["ticker"]:6s} {r["timeframe"]:2s} {r["folds"]:3d}   '
              f'{r["mean_long_acc"]:7.1%} {r["mean_short_acc"]:7.1%}  '
              f'{r["mean_hi_conf_wr"]:7.1%}  {r["total_hi_conf"]:5d}    {r["total_samples"]:5d}')

    total_time = time.time() - t_start
    print(f'\n  Время: {total_time:.0f}s ({total_time/60:.1f} мин)')

    # Save
    out_path = 'walkforward_results.json'
    with open(out_path, 'w') as f:
        json.dump(all_results, f, indent=2, default=str)
    print(f'  Результаты сохранены: {out_path}')
    print()