"""
Training pipeline for the strategy filter model.

Workflow:
  1. Load dataset.csv
  2. Temporal split 70/15/15
  3. Feature engineering + scaling
  4. Train Hybrid Gated Fusion model
  5. Evaluate: classification metrics + trading simulation
"""

import os
import warnings
import json
import logging
import numpy as np
import pandas as pd
from sklearn.preprocessing import StandardScaler, RobustScaler
import tensorflow as tf

warnings.filterwarnings('ignore')
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
logging.basicConfig(level=logging.INFO, format='%(asctime)s  %(levelname)-7s %(message)s')
logger = logging.getLogger('train')


# ── Config ──
DATA_PATH = 'data/dataset.csv'
MODEL_DIR = 'models'
EXPIRIES = [1, 2, 3, 4, 5, 10]
BATCH_SIZE = 128
EPOCHS = 200
LR = 1e-3
WD = 1e-4

# ── Column groups ──
STRATEGY_NAMES = ['ema_rsi_trend', 'bb_pa', 'rsi_divergence', 'macd_stoch', 'breakout_retest']
SIGNAL_COLS = [f'{n}_signal' for n in STRATEGY_NAMES]
CONF_COLS = [f'{n}_confidence' for n in STRATEGY_NAMES]
CONTEXT_COLS = ['RSI14', 'ATR14', 'MACD_hist', 'body_pct', 'hour', 'day_of_week']
TARGET_COLS = [f'target_{exp}h' for exp in EXPIRIES]


def load_and_preprocess():
    """Load dataset, temporal split, scale features."""
    logger.info("Loading %s", DATA_PATH)
    df = pd.read_csv(DATA_PATH, parse_dates=['timestamp'])
    df = df.sort_values('timestamp').reset_index(drop=True)
    logger.info("Loaded %d rows, %s → %s", len(df), df['timestamp'].iloc[0], df['timestamp'].iloc[-1])

    # Instrument one-hot
    df['is_BITCOIN'] = (df['instrument'] == 'BITCOIN').astype(float)
    df['is_EURUSD'] = (df['instrument'] == 'EURUSD').astype(float)

    # Cyclical time encoding
    df['hour_sin'] = np.sin(2 * np.pi * df['hour'] / 24)
    df['hour_cos'] = np.cos(2 * np.pi * df['hour'] / 24)
    df['dow_sin'] = np.sin(2 * np.pi * df['day_of_week'] / 7)
    df['dow_cos'] = np.cos(2 * np.pi * df['day_of_week'] / 7)

    # Normalize confidence to [0, 1]
    for col in CONF_COLS:
        if col in df.columns:
            df[col] = df[col] / 100.0

    # Sample weight: rows with any strategy signal get weight 1.0, silent rows get 0.1
    df['has_signal'] = (df[SIGNAL_COLS] != 0).any(axis=1).astype(float)
    df['sample_weight'] = np.where(df['has_signal'] == 1, 1.0, 0.1)

    # Temporal split
    n = len(df)
    train_end = int(n * 0.70)
    val_end = int(n * 0.85)

    df_train = df.iloc[:train_end].copy()
    df_val = df.iloc[train_end:val_end].copy()
    df_test = df.iloc[val_end:].copy()

    logger.info("Split: train=%d  val=%d  test=%d", len(df_train), len(df_val), len(df_test))
    logger.info("Train: %s → %s", df_train['timestamp'].iloc[0], df_train['timestamp'].iloc[-1])
    logger.info("Val:   %s → %s", df_val['timestamp'].iloc[0], df_val['timestamp'].iloc[-1])
    logger.info("Test:  %s → %s", df_test['timestamp'].iloc[0], df_test['timestamp'].iloc[-1])

    # Scale: RobustScaler for extreme values, StandardScaler for bounded
    context_features = ['RSI14', 'ATR14', 'MACD_hist', 'body_pct', 'hour_sin', 'hour_cos', 'dow_sin', 'dow_cos']

    scalers = {}
    for col_name in context_features:
        if col_name in {'ATR14', 'MACD_hist'}:
            scaler = RobustScaler()
        elif col_name == 'body_pct':
            continue  # already [0, 1]
        else:
            scaler = StandardScaler()

        vals_train = df_train[col_name].values.reshape(-1, 1)
        scaler.fit(vals_train)
        scalers[col_name] = scaler

        df_train[col_name] = scaler.transform(vals_train).ravel()
        df_val[col_name] = scaler.transform(df_val[col_name].values.reshape(-1, 1)).ravel()
        df_test[col_name] = scaler.transform(df_test[col_name].values.reshape(-1, 1)).ravel()

    return df_train, df_val, df_test, scalers


def build_inputs_targets(df):
    """Extract model inputs and targets from DataFrame."""
    strategy_feat = df[SIGNAL_COLS + CONF_COLS].values.astype(np.float32)
    context_feat = df[['RSI14', 'ATR14', 'MACD_hist', 'body_pct',
                       'hour_sin', 'hour_cos', 'dow_sin', 'dow_cos']].values.astype(np.float32)
    instr_feat = df[['is_BITCOIN', 'is_EURUSD']].values.astype(np.float32)

    targets = {}
    for exp in EXPIRIES:
        targets[f'expiry_{exp}h'] = df[f'target_{exp}h'].values.astype(np.float32)

    return [strategy_feat, context_feat, instr_feat], targets


def train_model(df_train, df_val):
    """Build and train the model."""
    os.makedirs(MODEL_DIR, exist_ok=True)

    X_train, y_train = build_inputs_targets(df_train)
    X_val, y_val = build_inputs_targets(df_val)
    sw_train = df_train['sample_weight'].values.astype(np.float32)
    sw_val = df_val['sample_weight'].values.astype(np.float32)

    logger.info("X_train shapes: strategy=%s  context=%s  instr=%s",
                X_train[0].shape, X_train[1].shape, X_train[2].shape)
    logger.info("Target keys: %s", list(y_train.keys()))

    from model import build_model, compile_model

    model = build_model()
    model = compile_model(model, lr=LR, wd=WD)
    model.summary()

    callbacks = [
        tf.keras.callbacks.EarlyStopping(
            monitor='val_loss', patience=30, min_delta=1e-4,
            restore_best_weights=True, verbose=1
        ),
        tf.keras.callbacks.ReduceLROnPlateau(
            monitor='val_loss', factor=0.5, patience=15,
            min_lr=1e-6, verbose=1
        ),
        tf.keras.callbacks.ModelCheckpoint(
            filepath=os.path.join(MODEL_DIR, 'best_strategy_filter.keras'),
            monitor='val_loss', save_best_only=True, verbose=1
        ),
        tf.keras.callbacks.CSVLogger(os.path.join(MODEL_DIR, 'training_log.csv')),
        tf.keras.callbacks.TerminateOnNaN(),
    ]

    history = model.fit(
        X_train, y_train,
        validation_data=(X_val, y_val),
        sample_weight=sw_train,
        batch_size=BATCH_SIZE,
        epochs=EPOCHS,
        callbacks=callbacks,
        verbose=1,
    )

    # Load best checkpoint (without compile to avoid custom loss serialization issues)
    from model import FocalBinaryCrossentropy
    model = tf.keras.models.load_model(
        os.path.join(MODEL_DIR, 'best_strategy_filter.keras'),
        custom_objects={'FocalBinaryCrossentropy': FocalBinaryCrossentropy},
        compile=False,
    )
    from model import compile_model as cm
    model = cm(model, lr=LR, wd=WD)

    return model, history


def evaluate(model, df_test):
    """Evaluate on test set and run trading simulation."""
    X_test, y_test = build_inputs_targets(df_test)

    logger.info("\n=== Classification Metrics (Test) ===")
    results = model.evaluate(X_test, y_test, verbose=1, return_dict=True)
    for k, v in results.items():
        if 'auc' in k or 'acc' in k:
            logger.info("  %s: %.4f", k, v)

    # Trading simulation
    logger.info("\n=== Trading Simulation (Test) ===")
    preds = model.predict(X_test, verbose=0)
    simulate_trading(df_test, preds)

    return results


def simulate_trading(df, predictions):
    """Simulate binary options trading with model predictions."""
    has_close = 'Close' in df.columns

    for threshold in [0.50, 0.55, 0.60, 0.65, 0.70]:
        balance = 1000.0
        trades = 0
        wins = 0

        for i in range(len(df)):
            row = df.iloc[i]
            consensus = (row[SIGNAL_COLS] != 0).any()

            if not consensus:
                continue

            # Find best expiry prediction
            best_prob = 0
            best_exp = None
            for exp in EXPIRIES:
                prob = predictions[f'expiry_{exp}h'][i][0]
                if prob > best_prob:
                    best_prob = prob
                    best_exp = exp

            if best_prob < threshold:
                continue

            # Determine direction from weighted vote
            call_score = 0
            put_score = 0
            for sn in STRATEGY_NAMES:
                sig = row[f'{sn}_signal']
                conf = row[f'{sn}_confidence']
                if sig == 1:
                    call_score += conf
                elif sig == -1:
                    put_score += conf

            if call_score == 0 and put_score == 0:
                continue

            direction = 'CALL' if call_score > put_score else 'PUT'

            # Check outcome if price data available
            if has_close:
                try:
                    future_idx = df.index.get_loc(row.name) + best_exp
                    if future_idx < len(df):
                        future_close = df.iloc[future_idx]['Close']
                        current_close = row['Close']
                        win = (direction == 'CALL' and future_close > current_close) or \
                              (direction == 'PUT' and future_close < current_close)
                    else:
                        continue
                except Exception:
                    continue
            else:
                win = None  # no price data for outcome check

            pnl = 80 if win else -100
            balance += pnl
            trades += 1
            if win:
                wins += 1

        if trades > 0:
            wr = wins / trades * 100
            pf = (wins * 80) / ((trades - wins) * 100) if (trades - wins) > 0 else float('inf')
            logger.info("  threshold=%.2f  trades=%d  wr=%.1f%%  pf=%.2f  pnl=%+.0f  (of %d consensus rows)",
                        threshold, trades, wr, pf, balance - 1000, df['has_signal'].sum() if 'has_signal' in df.columns else 0)
        else:
            logger.info("  threshold=%.2f  trades=0", threshold)


def main():
    df_train, df_val, df_test, scalers = load_and_preprocess()
    model, history = train_model(df_train, df_val)
    evaluate(model, df_test)

    logger.info("\n=== Done ===")
    logger.info("Model: %s", os.path.join(MODEL_DIR, 'best_strategy_filter.keras'))


if __name__ == '__main__':
    main()