src/models/train_specialist.py

#!/usr/bin/env python3
"""
Regime Specialist Trainer

Trains a specialized PPO agent on a specific market regime:
BULL_TREND, BEAR_TREND, RANGE_CHOP, or HIGH_VOL_BREAKOUT.

The HMM Regime Classifier is used to filter the historical data
so the agent only sees state-transitions relevant to its specialty.

Usage:
    python -m src.models.train_specialist --asset BTCUSDT --regime BULL_TREND --timesteps 500000
"""

import os
import sys
import logging
import argparse
from pathlib import Path

import numpy as np
import pandas as pd
from stable_baselines3 import PPO
from stable_baselines3.common.callbacks import EvalCallback
from stable_baselines3.common.monitor import Monitor
from stable_baselines3.common.vec_env import DummyVecEnv

# Add project root
PROJECT_ROOT = Path(__file__).parent.parent.parent
sys.path.insert(0, str(PROJECT_ROOT))

from src.env.mtf_env import create_mtf_env
from src.models.regime_classifier import RegimeClassifier, REGIME_NAMES
from src.backtest.data_loader import download_binance_data

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


def create_regime_dataset(
    df: pd.DataFrame,
    regime: str,
    classifier: RegimeClassifier,
    context_window: int = 48,
) -> pd.DataFrame:
    """
    Filter dataframe to only include periods where the requested regime is active.
    
    CRITICAL: We can't just drop rows, because the environment needs contiguous
    time-series data to compute features (like RSI, MACD, etc.) properly.
    
    Instead, we find continuous 'chunks' of the regime. If a chunk is shorter
    than the context window, we discard it.
    """
    logger.info(f"Filtering dataset for {regime} regime...")
    
    indices = classifier.get_regime_filtered_indices(df, regime)
    
    if len(indices) == 0:
        logger.warning(f"No data found for regime {regime}")
        return pd.DataFrame()
        
    # Find contiguous chunks
    chunks = []
    current_chunk = [indices[0]]
    
    min_chunk_length = 6  # Require at least 6 consecutive hours in the regime
    
    for i in range(1, len(indices)):
        if indices[i] == indices[i-1] + 1:
            current_chunk.append(indices[i])
        else:
            if len(current_chunk) >= min_chunk_length:
                chunks.append(current_chunk)
            current_chunk = [indices[i]]
            
    if len(current_chunk) >= min_chunk_length:
        chunks.append(current_chunk)
        
    # Reassemble safe dataframe chunks
    # We pad the beginning of each chunk with context_window bars from BEFORE the regime started
    # so the environment can compute features right when the regime begins.
    safe_indices = []
    for chunk in chunks:
        start_idx = max(0, chunk[0] - context_window)
        end_idx = chunk[-1]
        safe_indices.extend(list(range(start_idx, end_idx + 1)))
        
    # Remove duplicates but keep order
    safe_indices = sorted(list(set(safe_indices)))
    
    regime_df = df.iloc[safe_indices].copy()
    
    logger.info(
        f"✅ Filtered to {len(regime_df)} bars "
        f"({len(regime_df)/len(df)*100:.1f}%), {len(chunks)} chunks"
    )
    
    return regime_df


def train_specialist(
    symbol: str = 'BTCUSDT',
    regime: str = 'BULL_TREND',
    timesteps: int = 500000,
    days: int = 730,
):
    """Train a regime specialist agent."""
    
    # ── Verify Regime ──
    valid_regimes = list(REGIME_NAMES.values())
    if regime not in valid_regimes:
        raise ValueError(f"Invalid regime {regime}. Must be one of {valid_regimes}")
        
    # ── Fetch Data ──
    clean_symbol = symbol.replace('USDT', '/USDT')
    logger.info(f"📥 Fetching {days} days of data for {clean_symbol}...")
    df = download_binance_data(symbol=clean_symbol, timeframe='1h', days=days)
    
    # ── Load Regime Classifier ──
    classifier = RegimeClassifier()
    if not classifier.load(symbol):
        logger.info(f"No existing classifier found. Training new HMM for {symbol}...")
        classifier.fit(df)
        classifier.save(symbol)
        
    # ── Filter Data ──
    regime_df = create_regime_dataset(df, regime, classifier, context_window=48)
    if len(regime_df) < 1000:
        logger.error(f"Not enough data to train {regime} specialist (only {len(regime_df)} bars)")
        return
        
    # ── Train/Val Split ──
    split_idx = int(len(regime_df) * 0.8)
    train_df = regime_df.iloc[:split_idx]
    val_df = regime_df.iloc[split_idx:]
    
    # ── Environments ──
    # Using MTF Env from Phase 11.2
    logger.info("Setting up Multi-Timeframe Environments...")
    train_env = create_mtf_env(train_df)
    train_env = Monitor(train_env)
    train_env = DummyVecEnv([lambda: train_env])
    
    val_env = create_mtf_env(val_df)
    val_env = Monitor(val_env)
    val_env = DummyVecEnv([lambda: val_env])
    
    # ── PPO Setup ──
    # We use a slightly smaller network since it's a specialist
    policy_kwargs = dict(net_arch=[dict(pi=[256, 128], vf=[256, 128])])
    
    model = PPO(
        "MlpPolicy",
        train_env,
        learning_rate=3e-4,
        n_steps=2048,
        batch_size=64,
        n_epochs=10,
        gamma=0.99,
        gae_lambda=0.95,
        clip_range=0.2,
        ent_coef=0.01,
        policy_kwargs=policy_kwargs,
        verbose=1,
        device="auto"
    )
    
    # ── Training ──
    logger.info(f"\n{'='*60}")
    logger.info(f"🚀 Training {regime} Specialist for {symbol}")
    logger.info(f"{'='*60}\n")
    
    model_dir = Path('./data/models/specialists')
    model_dir.mkdir(parents=True, exist_ok=True)
    
    model_path = model_dir / f"ppo_{symbol.lower()}_{regime.lower()}"
    best_model_dir = model_dir / f"best_{symbol.lower()}_{regime.lower()}"
    
    eval_callback = EvalCallback(
        val_env,
        best_model_save_path=str(best_model_dir),
        log_path=str(model_dir / "logs"),
        eval_freq=10000,
        deterministic=True,
        render=False,
    )
    
    model.learn(total_timesteps=timesteps, callback=eval_callback)
    
    # Save final model
    model.save(model_path)
    logger.info(f"✅ Final model saved to {model_path}.zip")
    logger.info(f"✅ Best model saved to {best_model_dir}/best_model.zip")


# ─── CLI ──────────────────────────────────────────────────────────────────────

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Train Regime Specialist')
    parser.add_argument('--asset', type=str, default='BTCUSDT')
    parser.add_argument('--regime', type=str, required=True, 
                        help='BULL_TREND, BEAR_TREND, RANGE_CHOP, or HIGH_VOL_BREAKOUT')
    parser.add_argument('--timesteps', type=int, default=500000)
    parser.add_argument('--days', type=int, default=730)
    args = parser.parse_args()
    
    train_specialist(
        symbol=args.asset,
        regime=args.regime.upper(),
        timesteps=args.timesteps,
        days=args.days,
    )