#!/usr/bin/env python3
"""
Metamodel Trainer
==================
Training pipeline for the Neural Gating Metamodel including:
1. Strategy prediction generation
2. Gating network training with walk-forward validation
3. Confidence calibration
4. Entropy regularization to prevent strategy collapse

Supports:
- Walk-forward training (time-series cross-validation)
- Online learning (incremental updates)
- Multi-symbol training

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

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from typing import Dict, List, Optional, Tuple, Any, Generator
from dataclasses import dataclass, field
from pathlib import Path
from loguru import logger
import time
from collections import defaultdict
import json

from .model import NeuralGatingMetamodel, MetamodelConfig, MetamodelPrediction
from .gating_network import GatingNetwork, GatingConfig
from .calibration import ConfidenceCalibrator
from .ensemble_pipeline import EnsemblePipeline, StrategyPredictions


@dataclass
class TrainerConfig:
    """Configuration for the Metamodel Trainer."""

    # Training parameters
    epochs: int = 50
    batch_size: int = 32
    learning_rate: float = 1e-3
    weight_decay: float = 1e-4
    gradient_clip: float = 1.0

    # Walk-forward parameters
    n_folds: int = 5
    min_train_size: int = 500
    validation_ratio: float = 0.2

    # Regularization
    entropy_weight: float = 0.1
    entropy_target: float = 0.7  # Target entropy as fraction of max
    diversity_weight: float = 0.05  # Penalize correlation between strategies

    # Early stopping
    early_stopping_patience: int = 10
    early_stopping_min_delta: float = 1e-4

    # Data parameters
    sequence_length: int = 100
    prediction_horizon: int = 5  # Candles ahead to predict

    # Device
    device: Optional[str] = None

    def to_dict(self) -> Dict[str, Any]:
        return {
            'epochs': self.epochs,
            'batch_size': self.batch_size,
            'learning_rate': self.learning_rate,
            'weight_decay': self.weight_decay,
            'gradient_clip': self.gradient_clip,
            'n_folds': self.n_folds,
            'min_train_size': self.min_train_size,
            'validation_ratio': self.validation_ratio,
            'entropy_weight': self.entropy_weight,
            'entropy_target': self.entropy_target,
            'diversity_weight': self.diversity_weight,
            'early_stopping_patience': self.early_stopping_patience,
            'early_stopping_min_delta': self.early_stopping_min_delta,
            'sequence_length': self.sequence_length,
            'prediction_horizon': self.prediction_horizon,
            'device': self.device
        }


@dataclass
class TrainingMetrics:
    """Metrics from training run."""

    # Loss metrics
    final_loss: float = 0.0
    best_loss: float = float('inf')
    direction_loss: float = 0.0
    entropy_loss: float = 0.0
    diversity_loss: float = 0.0

    # Accuracy metrics
    direction_accuracy: float = 0.0
    calibration_ece: float = 0.0

    # Walk-forward metrics
    fold_losses: List[float] = field(default_factory=list)
    fold_accuracies: List[float] = field(default_factory=list)

    # Strategy weight statistics
    weight_means: Dict[str, float] = field(default_factory=dict)
    weight_stds: Dict[str, float] = field(default_factory=dict)

    # Training info
    epochs_trained: int = 0
    training_time_s: float = 0.0

    def to_dict(self) -> Dict[str, Any]:
        return {
            'final_loss': self.final_loss,
            'best_loss': self.best_loss,
            'direction_loss': self.direction_loss,
            'entropy_loss': self.entropy_loss,
            'diversity_loss': self.diversity_loss,
            'direction_accuracy': self.direction_accuracy,
            'calibration_ece': self.calibration_ece,
            'fold_losses': self.fold_losses,
            'fold_accuracies': self.fold_accuracies,
            'weight_means': self.weight_means,
            'weight_stds': self.weight_stds,
            'epochs_trained': self.epochs_trained,
            'training_time_s': self.training_time_s
        }


class StrategyPredictionDataset(Dataset):
    """PyTorch Dataset for strategy predictions."""

    def __init__(
        self,
        strategy_preds: List[StrategyPredictions],
        contexts: List[torch.Tensor],
        targets: List[float],
        device: str = 'cpu'
    ):
        """
        Initialize dataset.

        Args:
            strategy_preds: List of strategy predictions
            contexts: List of context tensors
            targets: Target direction labels
            device: Device for tensors
        """
        self.strategy_preds = strategy_preds
        self.contexts = contexts
        self.targets = targets
        self.device = device

    def __len__(self) -> int:
        return len(self.targets)

    def __getitem__(self, idx: int) -> Tuple[StrategyPredictions, torch.Tensor, float]:
        return self.strategy_preds[idx], self.contexts[idx], self.targets[idx]


class MetamodelTrainer:
    """
    Trainer for the Neural Gating Metamodel.

    Implements:
    - Strategy prediction generation
    - Walk-forward training
    - Entropy regularization
    - Confidence calibration

    Usage:
        trainer = MetamodelTrainer(model)
        metrics = trainer.train(df_train, symbol='EURUSD')

        # Walk-forward evaluation
        wf_metrics = trainer.walk_forward_train(df, symbol, n_folds=5)
    """

    STRATEGY_NAMES = ['pva', 'mrd', 'vbp', 'msa', 'mts']

    def __init__(
        self,
        model: NeuralGatingMetamodel,
        config: Optional[TrainerConfig] = None
    ):
        """
        Initialize the trainer.

        Args:
            model: NeuralGatingMetamodel instance
            config: Trainer configuration
        """
        self.model = model
        self.config = config or TrainerConfig()
        self.device = self.config.device or model.device

        # Training history
        self._history: Dict[str, List[float]] = defaultdict(list)

        logger.info(f"MetamodelTrainer initialized: device={self.device}")

    def generate_strategy_predictions(
        self,
        df: pd.DataFrame,
        symbols: List[str],
        verbose: bool = True
    ) -> Tuple[List[StrategyPredictions], List[torch.Tensor], List[float]]:
        """
        Generate predictions from all strategies for training data.

        Args:
            df: OHLCV DataFrame
            symbols: Symbols to process
            verbose: Print progress

        Returns:
            Tuple of (strategy_predictions, contexts, targets)
        """
        logger.info(f"Generating strategy predictions for {len(df)} samples...")

        # Ensure strategies are loaded
        if not self.model.ensemble_pipeline.strategies:
            self.model.load_strategies(symbols)

        all_strategy_preds = []
        all_contexts = []
        all_targets = []

        seq_len = self.config.sequence_length
        horizon = self.config.prediction_horizon

        n_samples = len(df) - seq_len - horizon
        if n_samples <= 0:
            raise ValueError(f"Not enough data: need at least {seq_len + horizon} rows")

        for i in range(n_samples):
            if verbose and (i + 1) % 100 == 0:
                logger.info(f"Processing sample {i + 1}/{n_samples}")

            # Extract window
            start_idx = i
            end_idx = start_idx + seq_len
            sample_df = df.iloc[start_idx:end_idx].copy()

            # Get strategy predictions
            strategy_preds = self.model.ensemble_pipeline.predict_all_strategies(sample_df)
            all_strategy_preds.append(strategy_preds)

            # Get context
            context = self.model.ensemble_pipeline.extract_market_context(sample_df)
            all_contexts.append(context)

            # Compute target (future return direction)
            current_close = df['close'].iloc[end_idx - 1]
            future_close = df['close'].iloc[end_idx + horizon - 1]
            target = 1.0 if future_close > current_close else -1.0 if future_close < current_close else 0.0
            all_targets.append(target)

        logger.info(f"Generated {len(all_strategy_preds)} prediction samples")

        return all_strategy_preds, all_contexts, all_targets

    def train_gating(
        self,
        strategy_preds: List[StrategyPredictions],
        contexts: List[torch.Tensor],
        targets: List[float],
        val_strategy_preds: Optional[List[StrategyPredictions]] = None,
        val_contexts: Optional[List[torch.Tensor]] = None,
        val_targets: Optional[List[float]] = None,
        verbose: bool = True
    ) -> TrainingMetrics:
        """
        Train the gating network.

        Args:
            strategy_preds: Training strategy predictions
            contexts: Training context tensors
            targets: Training targets
            val_strategy_preds: Optional validation predictions
            val_contexts: Optional validation contexts
            val_targets: Optional validation targets
            verbose: Print progress

        Returns:
            TrainingMetrics object
        """
        logger.info(f"Training gating network with {len(strategy_preds)} samples...")
        start_time = time.time()

        # Set up optimizer
        all_params = (
            list(self.model.gating_network.parameters()) +
            list(self.model.context_encoder.parameters()) +
            list(self.model.strategy_encoder.parameters())
        )

        optimizer = optim.AdamW(
            all_params,
            lr=self.config.learning_rate,
            weight_decay=self.config.weight_decay
        )

        scheduler = optim.lr_scheduler.CosineAnnealingLR(
            optimizer,
            T_max=self.config.epochs
        )

        # Training loop
        best_loss = float('inf')
        patience_counter = 0
        metrics = TrainingMetrics()

        for epoch in range(self.config.epochs):
            self.model.gating_network.train()
            self.model.context_encoder.train()
            self.model.strategy_encoder.train()

            epoch_losses = {
                'total': 0.0,
                'direction': 0.0,
                'entropy': 0.0,
                'diversity': 0.0
            }

            # Shuffle indices
            n_samples = len(strategy_preds)
            indices = np.random.permutation(n_samples)

            n_batches = (n_samples + self.config.batch_size - 1) // self.config.batch_size

            for batch_idx in range(n_batches):
                batch_start = batch_idx * self.config.batch_size
                batch_end = min(batch_start + self.config.batch_size, n_samples)
                batch_indices = indices[batch_start:batch_end]

                optimizer.zero_grad()

                batch_loss, batch_dir_loss, batch_ent_loss, batch_div_loss = \
                    self._compute_batch_loss(
                        strategy_preds,
                        contexts,
                        targets,
                        batch_indices
                    )

                batch_loss.backward()

                # Gradient clipping
                torch.nn.utils.clip_grad_norm_(all_params, self.config.gradient_clip)

                optimizer.step()

                epoch_losses['total'] += batch_loss.item()
                epoch_losses['direction'] += batch_dir_loss
                epoch_losses['entropy'] += batch_ent_loss
                epoch_losses['diversity'] += batch_div_loss

            # Average losses
            for key in epoch_losses:
                epoch_losses[key] /= n_batches
                self._history[f'train_{key}'].append(epoch_losses[key])

            scheduler.step()

            # Validation
            val_loss = None
            if val_strategy_preds is not None:
                val_loss = self._compute_validation_loss(
                    val_strategy_preds,
                    val_contexts,
                    val_targets
                )
                self._history['val_loss'].append(val_loss)

            # Early stopping check
            current_loss = val_loss if val_loss is not None else epoch_losses['total']
            if current_loss < best_loss - self.config.early_stopping_min_delta:
                best_loss = current_loss
                patience_counter = 0
            else:
                patience_counter += 1

            if patience_counter >= self.config.early_stopping_patience:
                logger.info(f"Early stopping at epoch {epoch + 1}")
                break

            if verbose and (epoch + 1) % 5 == 0:
                val_str = f", Val: {val_loss:.4f}" if val_loss else ""
                logger.info(
                    f"Epoch {epoch + 1}/{self.config.epochs} - "
                    f"Loss: {epoch_losses['total']:.4f} "
                    f"(Dir: {epoch_losses['direction']:.4f}, "
                    f"Ent: {epoch_losses['entropy']:.4f}, "
                    f"Div: {epoch_losses['diversity']:.4f})"
                    f"{val_str}"
                )

            metrics.epochs_trained = epoch + 1

        # Finalize metrics
        metrics.final_loss = self._history['train_total'][-1]
        metrics.best_loss = best_loss
        metrics.direction_loss = self._history['train_direction'][-1]
        metrics.entropy_loss = self._history['train_entropy'][-1]
        metrics.diversity_loss = self._history['train_diversity'][-1]
        metrics.training_time_s = time.time() - start_time

        # Compute accuracy
        accuracy = self._compute_direction_accuracy(strategy_preds, contexts, targets)
        metrics.direction_accuracy = accuracy

        # Compute weight statistics
        weight_stats = self._compute_weight_statistics(strategy_preds, contexts)
        metrics.weight_means = weight_stats['means']
        metrics.weight_stds = weight_stats['stds']

        self.model.gating_network.eval()
        self.model.context_encoder.eval()
        self.model.strategy_encoder.eval()

        logger.info(f"Gating training complete in {metrics.training_time_s:.1f}s")

        return metrics

    def _compute_batch_loss(
        self,
        strategy_preds: List[StrategyPredictions],
        contexts: List[torch.Tensor],
        targets: List[float],
        indices: np.ndarray
    ) -> Tuple[torch.Tensor, float, float, float]:
        """Compute loss for a batch of samples."""
        total_loss = torch.tensor(0.0, device=self.device, requires_grad=True)
        direction_loss_sum = 0.0
        entropy_loss_sum = 0.0
        diversity_loss_sum = 0.0

        all_weights = []

        for idx in indices:
            strategy_pred = strategy_preds[idx]
            context = contexts[idx]
            target = targets[idx]

            # Encode strategy outputs
            strategy_tensors = self._encode_strategy_predictions(strategy_pred)

            # Get gating weights
            weights = self.model.gating_network(strategy_tensors, context)
            all_weights.append(weights)

            # Direction prediction
            pred_direction = self._compute_weighted_direction(strategy_pred, weights)

            # Direction loss
            target_tensor = torch.tensor(
                [float(target)],
                device=self.device,
                dtype=torch.float32
            )
            dir_loss = nn.functional.mse_loss(pred_direction, target_tensor)

            # Entropy loss
            ent_loss = self.model.gating_network.compute_entropy_loss(weights)

            # Accumulate
            sample_loss = dir_loss + ent_loss
            total_loss = total_loss + sample_loss

            direction_loss_sum += dir_loss.item()
            entropy_loss_sum += ent_loss.item()

        # Diversity loss (across batch)
        if len(all_weights) > 1:
            weights_stacked = torch.cat(all_weights, dim=0)  # (batch, n_strategies)
            diversity_loss = self._compute_diversity_loss(weights_stacked)
            total_loss = total_loss + self.config.diversity_weight * diversity_loss
            diversity_loss_sum = diversity_loss.item()

        # Average
        batch_size = len(indices)
        total_loss = total_loss / batch_size

        return (
            total_loss,
            direction_loss_sum / batch_size,
            entropy_loss_sum / batch_size,
            diversity_loss_sum
        )

    def _compute_diversity_loss(self, weights: torch.Tensor) -> torch.Tensor:
        """
        Compute diversity loss to encourage different weights across samples.

        Penalizes high correlation between strategy weights across batch.
        """
        if weights.size(0) < 2:
            return torch.tensor(0.0, device=self.device)

        # Normalize weights
        weights_centered = weights - weights.mean(dim=0, keepdim=True)

        # Compute correlation matrix
        weights_norm = weights_centered / (weights_centered.std(dim=0, keepdim=True) + 1e-8)
        corr_matrix = torch.mm(weights_norm.T, weights_norm) / weights.size(0)

        # Penalize off-diagonal correlations
        n_strategies = weights.size(1)
        eye = torch.eye(n_strategies, device=self.device)
        off_diag = corr_matrix * (1 - eye)

        diversity_loss = off_diag.abs().mean()

        return diversity_loss

    def _encode_strategy_predictions(
        self,
        strategy_preds: StrategyPredictions
    ) -> Dict[str, torch.Tensor]:
        """Convert strategy predictions to tensor format."""
        strategy_tensors = {}
        available = strategy_preds.get_available()

        for name, pred in available.items():
            values = list(pred.values())[:5]
            while len(values) < 5:
                values.append(0.0)

            tensor = torch.tensor(
                values,
                dtype=torch.float32,
                device=self.device
            ).unsqueeze(0)

            encoded = self.model.strategy_encoder(name, tensor)
            strategy_tensors[name] = encoded

        return strategy_tensors

    def _compute_weighted_direction(
        self,
        strategy_preds: StrategyPredictions,
        weights: torch.Tensor
    ) -> torch.Tensor:
        """Compute weighted direction prediction."""
        available = strategy_preds.get_available()
        direction = torch.zeros(1, device=self.device)

        for i, name in enumerate(self.STRATEGY_NAMES):
            if name not in available:
                continue

            pred = available[name]
            weight = weights[0, i]

            # Extract direction
            if name == 'pva':
                strat_dir = pred.get('direction', 0.0)
            elif name == 'mrd':
                regime = pred.get('regime', 1.0)
                strat_dir = {0: -1.0, 1: 0.0, 2: 1.0}.get(int(regime), 0.0)
                strat_dir *= pred.get('continuation_prob', 0.5)
            elif name == 'vbp':
                strat_dir = pred.get('direction', 0.0) * pred.get('breakout_prob', 0.0)
            elif name == 'msa':
                strat_dir = pred.get('is_bullish', 0.0) - pred.get('is_bearish', 0.0)
                strat_dir *= pred.get('bos_conf', 0.0)
            elif name == 'mts':
                strat_dir = pred.get('unified_direction', 0.0)
            else:
                strat_dir = 0.0

            direction = direction + weight * strat_dir

        return direction

    def _compute_validation_loss(
        self,
        strategy_preds: List[StrategyPredictions],
        contexts: List[torch.Tensor],
        targets: List[float]
    ) -> float:
        """Compute validation loss."""
        self.model.gating_network.eval()

        total_loss = 0.0
        n_samples = len(strategy_preds)

        with torch.no_grad():
            for i in range(n_samples):
                strategy_pred = strategy_preds[i]
                context = contexts[i]
                target = targets[i]

                strategy_tensors = self._encode_strategy_predictions(strategy_pred)
                weights = self.model.gating_network(strategy_tensors, context)
                pred_direction = self._compute_weighted_direction(strategy_pred, weights)

                target_tensor = torch.tensor(
                    [float(target)],
                    device=self.device,
                    dtype=torch.float32
                )
                loss = nn.functional.mse_loss(pred_direction, target_tensor)
                total_loss += loss.item()

        return total_loss / n_samples

    def _compute_direction_accuracy(
        self,
        strategy_preds: List[StrategyPredictions],
        contexts: List[torch.Tensor],
        targets: List[float]
    ) -> float:
        """Compute direction prediction accuracy."""
        self.model.gating_network.eval()

        correct = 0
        total = 0

        with torch.no_grad():
            for i in range(len(strategy_preds)):
                strategy_pred = strategy_preds[i]
                context = contexts[i]
                target = targets[i]

                strategy_tensors = self._encode_strategy_predictions(strategy_pred)
                weights = self.model.gating_network(strategy_tensors, context)
                pred_direction = self._compute_weighted_direction(strategy_pred, weights)

                pred_sign = 1 if pred_direction.item() > 0 else -1 if pred_direction.item() < 0 else 0
                target_sign = 1 if target > 0 else -1 if target < 0 else 0

                if pred_sign == target_sign:
                    correct += 1
                total += 1

        return correct / total if total > 0 else 0.0

    def _compute_weight_statistics(
        self,
        strategy_preds: List[StrategyPredictions],
        contexts: List[torch.Tensor]
    ) -> Dict[str, Dict[str, float]]:
        """Compute statistics of strategy weights."""
        self.model.gating_network.eval()

        all_weights = {name: [] for name in self.STRATEGY_NAMES}

        with torch.no_grad():
            for i in range(len(strategy_preds)):
                strategy_pred = strategy_preds[i]
                context = contexts[i]

                strategy_tensors = self._encode_strategy_predictions(strategy_pred)
                weights = self.model.gating_network(strategy_tensors, context)

                for j, name in enumerate(self.STRATEGY_NAMES):
                    all_weights[name].append(weights[0, j].item())

        means = {name: np.mean(weights) for name, weights in all_weights.items()}
        stds = {name: np.std(weights) for name, weights in all_weights.items()}

        return {'means': means, 'stds': stds}

    def train_calibrator(
        self,
        strategy_preds: List[StrategyPredictions],
        contexts: List[torch.Tensor],
        targets: List[float]
    ) -> Dict[str, float]:
        """
        Train the confidence calibrator.

        Args:
            strategy_preds: Strategy predictions
            contexts: Context tensors
            targets: Target directions

        Returns:
            Calibration metrics
        """
        logger.info("Training confidence calibrator...")

        raw_confidences = []
        binary_targets = []

        self.model.gating_network.eval()

        with torch.no_grad():
            for i in range(len(strategy_preds)):
                strategy_pred = strategy_preds[i]
                context = contexts[i]
                target = targets[i]

                strategy_tensors = self._encode_strategy_predictions(strategy_pred)
                weights = self.model.gating_network(strategy_tensors, context)

                # Compute raw confidence
                raw_conf = self._compute_raw_confidence(strategy_pred, weights)
                raw_confidences.append(raw_conf)

                # Binary target
                pred_dir = self._compute_weighted_direction(strategy_pred, weights)
                correct = (pred_dir.item() > 0) == (target > 0)
                binary_targets.append(int(correct))

        # Fit calibrator
        raw_confidences = np.array(raw_confidences)
        binary_targets = np.array(binary_targets)

        self.model.calibrator.fit(binary_targets, raw_confidences)

        # Evaluate
        metrics = self.model.calibrator.evaluate(binary_targets, raw_confidences, calibrated=True)

        return {
            'calibration_ece': metrics.ece,
            'calibration_mce': metrics.mce,
            'calibration_brier': metrics.brier_score
        }

    def _compute_raw_confidence(
        self,
        strategy_preds: StrategyPredictions,
        weights: torch.Tensor
    ) -> float:
        """Compute raw confidence from strategy predictions."""
        available = strategy_preds.get_available()
        confidence = 0.0

        for i, name in enumerate(self.STRATEGY_NAMES):
            if name not in available:
                continue

            pred = available[name]
            weight = weights[0, i].item()

            if name == 'pva':
                conf = pred.get('confidence', 0.0)
            elif name == 'mrd':
                conf = pred.get('regime_confidence', 0.0)
            elif name == 'vbp':
                conf = pred.get('confidence', 0.0) * pred.get('breakout_prob', 0.0)
            elif name == 'msa':
                conf = pred.get('bos_conf', 0.0)
            elif name == 'mts':
                conf = pred.get('confidence', 0.0)
            else:
                conf = 0.0

            confidence += weight * conf

        return float(np.clip(confidence, 0.0, 1.0))

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

        This is the recommended training method for time-series data.
        Each fold trains on past data and validates on future data.

        Args:
            df: Full OHLCV DataFrame
            symbol: Trading symbol
            n_folds: Number of folds (default from config)
            verbose: Print progress

        Returns:
            Aggregated training metrics
        """
        n_folds = n_folds or self.config.n_folds

        logger.info(f"Walk-forward training with {n_folds} folds...")

        total_samples = len(df)
        fold_size = (total_samples - self.config.min_train_size) // n_folds

        metrics = TrainingMetrics()
        all_fold_losses = []
        all_fold_accuracies = []

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

            # Calculate indices
            train_end = self.config.min_train_size + fold * fold_size
            val_end = train_end + fold_size

            train_df = df.iloc[:train_end]
            val_df = df.iloc[train_end:val_end]

            if verbose:
                logger.info(f"Train: {len(train_df)}, Val: {len(val_df)}")

            # Generate predictions
            train_preds, train_contexts, train_targets = self.generate_strategy_predictions(
                train_df, [symbol], verbose=False
            )
            val_preds, val_contexts, val_targets = self.generate_strategy_predictions(
                val_df, [symbol], verbose=False
            )

            # Train gating
            fold_metrics = self.train_gating(
                train_preds, train_contexts, train_targets,
                val_preds, val_contexts, val_targets,
                verbose=verbose
            )

            all_fold_losses.append(fold_metrics.best_loss)
            all_fold_accuracies.append(fold_metrics.direction_accuracy)

            if verbose:
                logger.info(
                    f"Fold {fold + 1} - Loss: {fold_metrics.best_loss:.4f}, "
                    f"Accuracy: {fold_metrics.direction_accuracy:.4f}"
                )

        # Train calibrator on final fold
        cal_metrics = self.train_calibrator(train_preds, train_contexts, train_targets)

        # Aggregate metrics
        metrics.fold_losses = all_fold_losses
        metrics.fold_accuracies = all_fold_accuracies
        metrics.final_loss = np.mean(all_fold_losses)
        metrics.best_loss = np.min(all_fold_losses)
        metrics.direction_accuracy = np.mean(all_fold_accuracies)
        metrics.calibration_ece = cal_metrics['calibration_ece']

        # Weight statistics from last fold
        weight_stats = self._compute_weight_statistics(train_preds, train_contexts)
        metrics.weight_means = weight_stats['means']
        metrics.weight_stds = weight_stats['stds']

        self.model._is_trained = True

        logger.info(f"\nWalk-forward complete: "
                    f"Avg Loss: {metrics.final_loss:.4f}, "
                    f"Avg Accuracy: {metrics.direction_accuracy:.4f}")

        return metrics


if __name__ == "__main__":
    # Test the MetamodelTrainer
    torch.manual_seed(42)
    np.random.seed(42)

    print("Testing MetamodelTrainer")
    print("=" * 60)

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

    price = 100 * np.cumprod(1 + np.random.randn(n_samples) * 0.001)
    df = pd.DataFrame({
        'open': price * (1 + np.random.randn(n_samples) * 0.001),
        'high': price * (1 + np.abs(np.random.randn(n_samples) * 0.002)),
        'low': price * (1 - np.abs(np.random.randn(n_samples) * 0.002)),
        'close': price,
        'volume': np.random.randint(1000, 10000, n_samples)
    }, index=dates)

    # Initialize model and trainer
    model_config = MetamodelConfig()
    model = NeuralGatingMetamodel(model_config)
    model.load_strategies(['EURUSD'])

    trainer_config = TrainerConfig(
        epochs=10,
        n_folds=2,
        min_train_size=100,
        sequence_length=50
    )
    trainer = MetamodelTrainer(model, trainer_config)

    # Generate predictions
    print("\n1. Generating strategy predictions...")
    preds, contexts, targets = trainer.generate_strategy_predictions(
        df, ['EURUSD'], verbose=True
    )
    print(f"   Generated {len(preds)} samples")

    # Train gating
    print("\n2. Training gating network...")
    split_idx = int(len(preds) * 0.8)
    metrics = trainer.train_gating(
        preds[:split_idx], contexts[:split_idx], targets[:split_idx],
        preds[split_idx:], contexts[split_idx:], targets[split_idx:],
        verbose=True
    )

    print(f"\n   Final loss: {metrics.final_loss:.4f}")
    print(f"   Direction accuracy: {metrics.direction_accuracy:.4f}")
    print(f"   Weight means: {metrics.weight_means}")

    # Train calibrator
    print("\n3. Training calibrator...")
    cal_metrics = trainer.train_calibrator(preds, contexts, targets)
    print(f"   ECE: {cal_metrics['calibration_ece']:.4f}")

    # Walk-forward training
    print("\n4. Testing walk-forward training...")
    wf_metrics = trainer.walk_forward_train(df, 'EURUSD', n_folds=2, verbose=True)
    print(f"\n   Walk-forward metrics:")
    print(f"   Avg Loss: {wf_metrics.final_loss:.4f}")
    print(f"   Avg Accuracy: {wf_metrics.direction_accuracy:.4f}")
    print(f"   Fold losses: {wf_metrics.fold_losses}")

    print("\nMetamodelTrainer test complete!")
