"""
MSA (Market Structure Analysis) - Trainer
==========================================
Training utilities for MSA model.

Features:
- Data preparation with structure labels
- Walk-forward training
- Cross-validation
- Model evaluation and metrics

Author: ML-Specialist (NEXUS v4.0)
Created: 2026-01-25
Version: 1.0.0
"""

import numpy as np
import pandas as pd
from typing import Dict, List, Optional, Tuple, Any, Union
from dataclasses import dataclass, field
from pathlib import Path
from datetime import datetime
from loguru import logger
import json

from .structure_detector import StructureDetector, StructureType
from .feature_engineering import MSAFeatureEngineer, PointOfInterest, POIType
from .model import MSAModel, MSAPrediction, MSAMetrics, BOSDirection


@dataclass
class TrainingConfig:
    """Configuration for MSA training"""
    # Data preparation
    swing_order: int = 5
    lookback_periods: int = 50
    forward_periods: int = 20

    # Label generation
    bos_lookforward: int = 30
    poi_reaction_threshold: float = 0.005
    continuation_lookforward: int = 20

    # Training
    train_ratio: float = 0.8
    use_sample_weighting: bool = True
    use_gpu: bool = True
    use_gnn: bool = False

    # Walk-forward
    n_folds: int = 5
    min_train_size: int = 5000
    expanding_window: bool = False

    # Model config
    model_config: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        return {
            'swing_order': self.swing_order,
            'lookback_periods': self.lookback_periods,
            'forward_periods': self.forward_periods,
            'bos_lookforward': self.bos_lookforward,
            'poi_reaction_threshold': self.poi_reaction_threshold,
            'continuation_lookforward': self.continuation_lookforward,
            'train_ratio': self.train_ratio,
            'use_sample_weighting': self.use_sample_weighting,
            'use_gpu': self.use_gpu,
            'use_gnn': self.use_gnn,
            'n_folds': self.n_folds,
            'min_train_size': self.min_train_size,
            'expanding_window': self.expanding_window,
            'model_config': self.model_config
        }


@dataclass
class WalkForwardResult:
    """Result from a single walk-forward fold"""
    fold_id: int
    train_start: int
    train_end: int
    val_start: int
    val_end: int
    train_metrics: MSAMetrics
    val_metrics: MSAMetrics
    model_path: Optional[str] = None

    def to_dict(self) -> Dict[str, Any]:
        return {
            'fold_id': self.fold_id,
            'train_range': [self.train_start, self.train_end],
            'val_range': [self.val_start, self.val_end],
            'train_metrics': {
                'bos_accuracy': self.train_metrics.bos_accuracy,
                'bos_macro_f1': self.train_metrics.bos_macro_f1,
                'poi_f1': self.train_metrics.poi_f1,
                'continuation_f1': self.train_metrics.continuation_f1
            },
            'val_metrics': {
                'bos_accuracy': self.val_metrics.bos_accuracy,
                'bos_macro_f1': self.val_metrics.bos_macro_f1,
                'poi_f1': self.val_metrics.poi_f1,
                'continuation_f1': self.val_metrics.continuation_f1
            },
            'model_path': self.model_path
        }


class MSATrainer:
    """
    Trainer for MSA (Market Structure Analysis) model.

    Handles:
    - Data preparation and feature extraction
    - Label generation from structure analysis
    - Model training with walk-forward validation
    - Evaluation and metrics
    """

    def __init__(
        self,
        config: Optional[TrainingConfig] = None
    ):
        """
        Initialize MSA Trainer.

        Args:
            config: Training configuration
        """
        self.config = config or TrainingConfig()

        # Initialize components
        self.detector = StructureDetector(swing_order=self.config.swing_order)
        self.feature_engineer = MSAFeatureEngineer(
            swing_order=self.config.swing_order,
            lookback_periods=self.config.lookback_periods
        )

        # Training state
        self.model: Optional[MSAModel] = None
        self.feature_columns: List[str] = []
        self.training_history: List[Dict] = []

        logger.info(f"MSATrainer initialized with config: {self.config.to_dict()}")

    def prepare_structure_data(
        self,
        df: pd.DataFrame
    ) -> Tuple[pd.DataFrame, pd.DataFrame]:
        """
        Prepare features and labels for training.

        Args:
            df: OHLCV DataFrame

        Returns:
            Tuple of (features DataFrame, labels DataFrame)
        """
        logger.info(f"Preparing structure data from {len(df)} samples")

        # Extract features
        features = self.feature_engineer.compute_structure_features(df)
        self.feature_columns = features.columns.tolist()

        # Generate labels
        labels = self._generate_labels(df)

        # Remove rows with NaN
        valid_mask = features.notna().all(axis=1) & labels.notna().all(axis=1)
        features = features[valid_mask]
        labels = labels[valid_mask]

        # Align indices
        common_idx = features.index.intersection(labels.index)
        features = features.loc[common_idx]
        labels = labels.loc[common_idx]

        logger.info(f"Prepared {len(features)} valid samples")
        logger.info(f"Label distributions:")
        for col in labels.columns:
            dist = labels[col].value_counts().to_dict()
            logger.info(f"  {col}: {dist}")

        return features, labels

    def _generate_labels(
        self,
        df: pd.DataFrame
    ) -> pd.DataFrame:
        """
        Generate training labels from structure analysis.

        Labels:
        - y_bos: Next BOS direction (0=neutral, 1=bullish, 2=bearish)
        - y_poi: POI reaction (0=no, 1=yes)
        - y_continuation: Structure continuation (0=reversal, 1=continuation)
        """
        labels = pd.DataFrame(index=df.index)

        # Get structure analysis
        swing_highs, swing_lows = self.detector.detect_swing_points(df)
        classified_highs, classified_lows = self.detector.detect_higher_highs_lower_lows(
            swing_highs, swing_lows
        )
        bos_events = self.detector.detect_break_of_structure(df, classified_highs, classified_lows)
        choch_events = self.detector.detect_change_of_character(df, classified_highs, classified_lows)
        fvgs = self.detector.detect_fair_value_gaps(df)
        order_blocks = self.detector.detect_order_blocks(df, bos_events)

        close = df['close'].values
        high = df['high'].values
        low = df['low'].values

        # Initialize label arrays
        n = len(df)
        y_bos = np.zeros(n, dtype=int)  # 0=neutral
        y_poi = np.zeros(n, dtype=int)  # 0=no reaction
        y_continuation = np.zeros(n, dtype=int)  # 0=reversal

        # Generate BOS direction labels
        # Look forward for next BOS event
        bos_indices = {bos.index: bos for bos in bos_events}

        for i in range(n - self.config.bos_lookforward):
            # Find next BOS in lookforward window
            for j in range(i + 1, min(i + self.config.bos_lookforward + 1, n)):
                if j in bos_indices:
                    bos = bos_indices[j]
                    if bos.event_type == StructureType.BOS_BULLISH:
                        y_bos[i] = BOSDirection.BULLISH
                    elif bos.event_type == StructureType.BOS_BEARISH:
                        y_bos[i] = BOSDirection.BEARISH
                    break

        # Generate POI reaction labels
        # Check if price reacts at POIs (OBs and FVGs)
        pois = self._collect_pois_for_labels(fvgs, order_blocks, classified_highs, classified_lows)

        for i in range(n - self.config.forward_periods):
            current_price = close[i]

            for poi in pois:
                if poi.index < i and poi.active:
                    # Check if near POI
                    dist = abs(current_price - poi.midpoint) / current_price
                    if dist < 0.005:  # Within 0.5% of POI
                        # Check for reaction in forward window
                        future_high = high[i + 1:i + self.config.forward_periods + 1].max()
                        future_low = low[i + 1:i + self.config.forward_periods + 1].min()

                        # Bullish reaction (at support)
                        if current_price <= poi.midpoint:
                            reaction = (future_high - current_price) / current_price
                            if reaction >= self.config.poi_reaction_threshold:
                                y_poi[i] = 1
                                break

                        # Bearish reaction (at resistance)
                        if current_price >= poi.midpoint:
                            reaction = (current_price - future_low) / current_price
                            if reaction >= self.config.poi_reaction_threshold:
                                y_poi[i] = 1
                                break

        # Generate structure continuation labels
        # Check if current structure continues or reverses
        choch_indices = {choch.index: choch for choch in choch_events}

        for i in range(n - self.config.continuation_lookforward):
            has_choch = False
            for j in range(i + 1, min(i + self.config.continuation_lookforward + 1, n)):
                if j in choch_indices:
                    has_choch = True
                    break

            # 1 = continuation (no CHoCH), 0 = reversal (CHoCH occurs)
            y_continuation[i] = 0 if has_choch else 1

        labels['y_bos'] = y_bos
        labels['y_poi'] = y_poi
        labels['y_continuation'] = y_continuation

        return labels

    def _collect_pois_for_labels(
        self,
        fvgs,
        order_blocks,
        swing_highs,
        swing_lows
    ) -> List[PointOfInterest]:
        """Collect POIs for label generation"""
        pois = []

        for fvg in fvgs:
            if not fvg.filled:
                pois.append(PointOfInterest(
                    poi_type=POIType.FAIR_VALUE_GAP,
                    price_low=fvg.low,
                    price_high=fvg.high,
                    index=fvg.index,
                    strength=1.0,
                    active=True
                ))

        for ob in order_blocks:
            if ob.valid:
                pois.append(PointOfInterest(
                    poi_type=POIType.ORDER_BLOCK,
                    price_low=ob.low,
                    price_high=ob.high,
                    index=ob.index,
                    strength=ob.strength,
                    active=True
                ))

        return pois

    def train(
        self,
        df: pd.DataFrame,
        symbol: str = "UNKNOWN",
        save_path: Optional[str] = None,
        verbose: bool = True
    ) -> MSAMetrics:
        """
        Train MSA model on provided data.

        Args:
            df: OHLCV DataFrame
            symbol: Trading symbol (for logging)
            save_path: Path to save trained model
            verbose: Print training progress

        Returns:
            Training metrics
        """
        logger.info(f"Training MSAModel for {symbol} on {len(df)} samples")

        # Prepare data
        features, labels = self.prepare_structure_data(df)

        # Split data
        train_size = int(len(features) * self.config.train_ratio)
        X_train = features.iloc[:train_size]
        X_val = features.iloc[train_size:]
        y_train = labels.iloc[:train_size]
        y_val = labels.iloc[train_size:]

        logger.info(f"Train size: {len(X_train)}, Validation size: {len(X_val)}")

        # Compute sample weights if enabled
        sample_weight = None
        if self.config.use_sample_weighting:
            sample_weight = self._compute_sample_weights(X_train, y_train)

        # Initialize model
        self.model = MSAModel(
            config=self.config.model_config or None,
            use_gpu=self.config.use_gpu,
            use_gnn=self.config.use_gnn
        )

        # Train
        metrics = self.model.fit(
            X_train,
            y_train['y_bos'].values,
            y_train['y_poi'].values,
            y_train['y_continuation'].values,
            X_val,
            y_val['y_bos'].values,
            y_val['y_poi'].values,
            y_val['y_continuation'].values,
            sample_weight=sample_weight,
            verbose=verbose
        )

        # Save model if path provided
        if save_path:
            self.model.save(save_path)

        # Record training history
        self.training_history.append({
            'timestamp': datetime.now().isoformat(),
            'symbol': symbol,
            'n_samples': len(df),
            'train_size': len(X_train),
            'val_size': len(X_val),
            'metrics': {
                'bos_accuracy': metrics.bos_accuracy,
                'bos_macro_f1': metrics.bos_macro_f1,
                'poi_f1': metrics.poi_f1,
                'continuation_f1': metrics.continuation_f1
            }
        })

        return metrics

    def walk_forward_train(
        self,
        df: pd.DataFrame,
        symbol: str = "UNKNOWN",
        n_folds: Optional[int] = None,
        save_dir: Optional[str] = None,
        verbose: bool = True
    ) -> List[WalkForwardResult]:
        """
        Train model using walk-forward validation.

        Walk-forward validation trains on expanding/sliding windows
        and validates on out-of-sample data, simulating real trading.

        Args:
            df: OHLCV DataFrame
            symbol: Trading symbol
            n_folds: Number of folds (uses config if None)
            save_dir: Directory to save fold models
            verbose: Print progress

        Returns:
            List of WalkForwardResult for each fold
        """
        n_folds = n_folds or self.config.n_folds
        logger.info(f"Walk-forward training for {symbol} with {n_folds} folds")

        # Prepare all data once
        features, labels = self.prepare_structure_data(df)
        n_samples = len(features)

        # Calculate fold sizes
        step_size = n_samples // (n_folds + 1)
        if step_size < self.config.min_train_size:
            logger.warning(f"Step size {step_size} < min_train_size {self.config.min_train_size}")
            n_folds = max(1, n_samples // self.config.min_train_size - 1)
            step_size = n_samples // (n_folds + 1)
            logger.info(f"Reduced to {n_folds} folds")

        results = []

        for fold in range(n_folds):
            logger.info(f"\n=== Fold {fold + 1}/{n_folds} ===")

            # Calculate indices
            if self.config.expanding_window:
                train_start = 0
            else:
                train_start = fold * step_size

            train_end = (fold + 1) * step_size
            val_start = train_end
            val_end = min(val_start + int(step_size * 0.2), n_samples)

            if val_end > n_samples or (train_end - train_start) < self.config.min_train_size:
                logger.warning(f"Skipping fold {fold + 1}: insufficient data")
                continue

            # Split data
            X_train = features.iloc[train_start:train_end]
            X_val = features.iloc[val_start:val_end]
            y_train = labels.iloc[train_start:train_end]
            y_val = labels.iloc[val_start:val_end]

            logger.info(f"Train: [{train_start}:{train_end}] n={len(X_train)}")
            logger.info(f"Val: [{val_start}:{val_end}] n={len(X_val)}")

            # Compute sample weights
            sample_weight = None
            if self.config.use_sample_weighting:
                sample_weight = self._compute_sample_weights(X_train, y_train)

            # Initialize model for this fold
            model = MSAModel(
                config=self.config.model_config or None,
                use_gpu=self.config.use_gpu,
                use_gnn=self.config.use_gnn
            )

            # Train
            train_metrics = model.fit(
                X_train,
                y_train['y_bos'].values,
                y_train['y_poi'].values,
                y_train['y_continuation'].values,
                X_val,
                y_val['y_bos'].values,
                y_val['y_poi'].values,
                y_val['y_continuation'].values,
                sample_weight=sample_weight,
                verbose=verbose
            )

            # Calculate validation metrics
            val_metrics = model._calculate_metrics(
                X_val.values,
                y_val['y_bos'].values,
                y_val['y_poi'].values,
                y_val['y_continuation'].values
            )

            # Save model if directory provided
            model_path = None
            if save_dir:
                model_path = f"{save_dir}/fold_{fold + 1}"
                model.save(model_path)

            result = WalkForwardResult(
                fold_id=fold + 1,
                train_start=train_start,
                train_end=train_end,
                val_start=val_start,
                val_end=val_end,
                train_metrics=train_metrics,
                val_metrics=val_metrics,
                model_path=model_path
            )
            results.append(result)

            if verbose:
                logger.info(f"Fold {fold + 1} Val Metrics:")
                logger.info(f"  BOS Accuracy: {val_metrics.bos_accuracy:.2%}")
                logger.info(f"  BOS Macro F1: {val_metrics.bos_macro_f1:.4f}")
                logger.info(f"  POI F1: {val_metrics.poi_f1:.4f}")
                logger.info(f"  Continuation F1: {val_metrics.continuation_f1:.4f}")

        # Print summary
        if results:
            self._print_walk_forward_summary(results)

        return results

    def _compute_sample_weights(
        self,
        X: pd.DataFrame,
        y: pd.DataFrame
    ) -> np.ndarray:
        """Compute sample weights based on structure clarity"""
        n = len(X)
        weights = np.ones(n)

        # Weight samples with clear structure more heavily
        if 'bos_net_bias' in X.columns:
            # Higher weight for clear directional bias
            bos_bias = np.abs(X['bos_net_bias'].values)
            weights *= (1 + bos_bias * 0.5)

        if 'structure_trend_bias' in X.columns:
            # Higher weight for clear trend
            trend_bias = np.abs(X['structure_trend_bias'].values)
            weights *= (1 + trend_bias * 0.3)

        # Normalize
        weights = weights / weights.mean()

        return weights

    def _print_walk_forward_summary(
        self,
        results: List[WalkForwardResult]
    ):
        """Print walk-forward validation summary"""
        bos_accs = [r.val_metrics.bos_accuracy for r in results]
        bos_f1s = [r.val_metrics.bos_macro_f1 for r in results]
        poi_f1s = [r.val_metrics.poi_f1 for r in results]
        cont_f1s = [r.val_metrics.continuation_f1 for r in results]

        logger.info("\n=== Walk-Forward Summary ===")
        logger.info(f"BOS Accuracy: {np.mean(bos_accs):.2%} (+/- {np.std(bos_accs):.2%})")
        logger.info(f"BOS Macro F1: {np.mean(bos_f1s):.4f} (+/- {np.std(bos_f1s):.4f})")
        logger.info(f"POI F1: {np.mean(poi_f1s):.4f} (+/- {np.std(poi_f1s):.4f})")
        logger.info(f"Continuation F1: {np.mean(cont_f1s):.4f} (+/- {np.std(cont_f1s):.4f})")

    def evaluate_structure_prediction(
        self,
        y_true: np.ndarray,
        y_pred: np.ndarray,
        prediction_type: str = 'bos'
    ) -> Dict[str, float]:
        """
        Evaluate structure predictions.

        Args:
            y_true: Ground truth labels
            y_pred: Predicted labels
            prediction_type: 'bos', 'poi', or 'continuation'

        Returns:
            Dictionary of metrics
        """
        from sklearn.metrics import (
            accuracy_score, precision_score, recall_score,
            f1_score, classification_report
        )

        metrics = {
            'accuracy': accuracy_score(y_true, y_pred)
        }

        if prediction_type == 'bos':
            metrics['macro_f1'] = f1_score(y_true, y_pred, average='macro', zero_division=0)
            metrics['weighted_f1'] = f1_score(y_true, y_pred, average='weighted', zero_division=0)

            # Per-class metrics
            for i, name in enumerate(['neutral', 'bullish', 'bearish']):
                mask = y_true == i
                if mask.any():
                    pred_mask = y_pred == i
                    metrics[f'{name}_precision'] = precision_score(
                        y_true == i, y_pred == i, zero_division=0
                    )
                    metrics[f'{name}_recall'] = recall_score(
                        y_true == i, y_pred == i, zero_division=0
                    )
                    metrics[f'{name}_f1'] = f1_score(
                        y_true == i, y_pred == i, zero_division=0
                    )
        else:
            # Binary classification
            metrics['precision'] = precision_score(y_true, y_pred, zero_division=0)
            metrics['recall'] = recall_score(y_true, y_pred, zero_division=0)
            metrics['f1'] = f1_score(y_true, y_pred, zero_division=0)

        return metrics

    def save_training_history(self, path: str):
        """Save training history to JSON"""
        with open(path, 'w') as f:
            json.dump(self.training_history, f, indent=2)
        logger.info(f"Saved training history to {path}")

    def load_model(self, path: str):
        """Load a trained model"""
        self.model = MSAModel(
            config=self.config.model_config or None,
            use_gpu=self.config.use_gpu,
            use_gnn=self.config.use_gnn
        )
        self.model.load(path)
        logger.info(f"Loaded model from {path}")


if __name__ == "__main__":
    # Test MSATrainer
    np.random.seed(42)

    # Create sample OHLCV data
    n_samples = 2000
    dates = pd.date_range('2024-01-01', periods=n_samples, freq='5min')

    price = 2000.0
    prices = [price]
    for i in range(1, n_samples):
        trend = 0.5 * np.sin(i / 100) + np.random.randn() * 0.5
        price = price + trend
        prices.append(price)

    prices = np.array(prices)

    df = pd.DataFrame({
        'open': prices + np.random.randn(n_samples) * 0.5,
        'high': prices + abs(np.random.randn(n_samples)) * 2.0,
        'low': prices - abs(np.random.randn(n_samples)) * 2.0,
        'close': prices + np.random.randn(n_samples) * 0.5,
        'volume': np.random.randint(100, 10000, n_samples)
    }, index=dates)

    df['high'] = df[['open', 'high', 'close']].max(axis=1)
    df['low'] = df[['open', 'low', 'close']].min(axis=1)

    # Create trainer
    config = TrainingConfig(
        swing_order=5,
        n_folds=3,
        min_train_size=200,
        use_gpu=False
    )
    trainer = MSATrainer(config)

    # Simple train
    print("=== Simple Training ===")
    metrics = trainer.train(df, symbol="TEST", verbose=True)

    # Predictions
    if trainer.model:
        predictions = trainer.model.predict(
            trainer.feature_engineer.compute_structure_features(df.tail(100))
        )
        print(f"\n=== Sample Predictions (last 5) ===")
        for pred in predictions[-5:]:
            print(f"BOS: {pred.next_bos_direction} ({pred.bos_confidence:.2f}), "
                  f"POI: {pred.poi_reaction_prob:.2f}, "
                  f"Cont: {pred.structure_continuation_prob:.2f}")

    # Walk-forward (with smaller dataset for speed)
    print("\n=== Walk-Forward Training ===")
    wf_results = trainer.walk_forward_train(
        df,
        symbol="TEST",
        n_folds=3,
        verbose=True
    )
