#!/usr/bin/env python3
"""
Neural Gating Metamodel (Nivel 2 - Versión 2)
==============================================
Neural network-based metamodel that learns dynamic weights for combining
5m and 15m predictions using a gating mechanism.

Architecture:
    alpha = sigmoid(MLP([attention_5m, attention_15m, attention_class, context]))
    pred_final = alpha * pred_5m + (1-alpha) * pred_15m + residual_correction

Key Features:
1. Gating Network: Learns when to trust 5m vs 15m predictions
2. Residual Correction: Fine-tunes the weighted average
3. Confidence Head: Binary classifier for tradeable signals
4. Per-Asset: One model per trading asset

This is an alternative to XGBoost Stacking for comparison.

Author: ML Pipeline
Version: 1.0.0
Created: 2026-01-07
"""

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

try:
    import torch
    import torch.nn as nn
    import torch.optim as optim
    from torch.utils.data import DataLoader, TensorDataset
    HAS_TORCH = True
except ImportError:
    HAS_TORCH = False
    logger.warning("PyTorch not available - Neural Gating disabled")

from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
from sklearn.metrics import accuracy_score, f1_score
from sklearn.preprocessing import StandardScaler


@dataclass
class NeuralGatingConfig:
    """Configuration for Neural Gating Metamodel"""

    # Feature configuration (same as XGBoost version)
    prediction_features: List[str] = field(default_factory=lambda: [
        'pred_high_5m', 'pred_low_5m',
        'pred_high_15m', 'pred_low_15m'
    ])

    attention_features: List[str] = field(default_factory=lambda: [
        'attention_5m', 'attention_15m',
        'attention_class_5m', 'attention_class_15m'
    ])

    context_features: List[str] = field(default_factory=lambda: [
        'ATR_ratio', 'volume_z'
    ])

    # Neural network architecture
    gating_hidden_dims: List[int] = field(default_factory=lambda: [32, 16])
    residual_hidden_dims: List[int] = field(default_factory=lambda: [64, 32])
    confidence_hidden_dims: List[int] = field(default_factory=lambda: [32, 16])

    # Training parameters
    learning_rate: float = 0.001
    batch_size: int = 256
    epochs: int = 100
    early_stopping_patience: int = 10
    val_split: float = 0.15
    min_train_samples: int = 2000
    dropout: float = 0.2

    # Device
    device: str = 'auto'  # 'auto', 'cuda', or 'cpu'


@dataclass
class NeuralGatingPrediction:
    """Result of neural gating prediction"""
    delta_high_final: np.ndarray
    delta_low_final: np.ndarray
    confidence: np.ndarray
    confidence_proba: np.ndarray
    alpha_high: np.ndarray  # Gating weight for high (5m weight)
    alpha_low: np.ndarray   # Gating weight for low (5m weight)

    def to_dict(self) -> Dict[str, Any]:
        def to_list(arr):
            return arr.tolist() if isinstance(arr, np.ndarray) else [arr]
        return {
            'delta_high_final': to_list(self.delta_high_final),
            'delta_low_final': to_list(self.delta_low_final),
            'confidence': to_list(self.confidence),
            'confidence_proba': to_list(self.confidence_proba),
            'alpha_high': to_list(self.alpha_high),
            'alpha_low': to_list(self.alpha_low)
        }


class GatingNetwork(nn.Module):
    """
    Gating network that learns when to weight 5m vs 15m predictions.

    Input: attention_5m, attention_15m, attention_class_5m, attention_class_15m,
           ATR_ratio, volume_z (6 features)
    Output: alpha_high, alpha_low (weights for 5m predictions, 0-1)
    """

    def __init__(self, input_dim: int, hidden_dims: List[int], dropout: float = 0.2):
        super().__init__()

        layers = []
        prev_dim = input_dim

        for hidden_dim in hidden_dims:
            layers.extend([
                nn.Linear(prev_dim, hidden_dim),
                nn.BatchNorm1d(hidden_dim),
                nn.LeakyReLU(0.1),
                nn.Dropout(dropout)
            ])
            prev_dim = hidden_dim

        # Output: alpha_high, alpha_low (sigmoid for 0-1 range)
        layers.append(nn.Linear(prev_dim, 2))

        self.network = nn.Sequential(*layers)

    def forward(self, x):
        return torch.sigmoid(self.network(x))


class ResidualNetwork(nn.Module):
    """
    Residual correction network that fine-tunes the weighted average.

    Input: All 10 meta features
    Output: residual_high, residual_low (additive corrections)
    """

    def __init__(self, input_dim: int, hidden_dims: List[int], dropout: float = 0.2):
        super().__init__()

        layers = []
        prev_dim = input_dim

        for hidden_dim in hidden_dims:
            layers.extend([
                nn.Linear(prev_dim, hidden_dim),
                nn.BatchNorm1d(hidden_dim),
                nn.LeakyReLU(0.1),
                nn.Dropout(dropout)
            ])
            prev_dim = hidden_dim

        # Output: residual_high, residual_low (unbounded)
        layers.append(nn.Linear(prev_dim, 2))

        self.network = nn.Sequential(*layers)

    def forward(self, x):
        return self.network(x)


class ConfidenceNetwork(nn.Module):
    """
    Binary classifier for tradeable signals.

    Input: All meta features + gating weights + residual corrections
    Output: confidence probability (0-1)
    """

    def __init__(self, input_dim: int, hidden_dims: List[int], dropout: float = 0.2):
        super().__init__()

        layers = []
        prev_dim = input_dim

        for hidden_dim in hidden_dims:
            layers.extend([
                nn.Linear(prev_dim, hidden_dim),
                nn.BatchNorm1d(hidden_dim),
                nn.LeakyReLU(0.1),
                nn.Dropout(dropout)
            ])
            prev_dim = hidden_dim

        layers.append(nn.Linear(prev_dim, 1))

        self.network = nn.Sequential(*layers)

    def forward(self, x):
        return torch.sigmoid(self.network(x))


class NeuralGatingMetamodel(nn.Module):
    """
    Complete Neural Gating Metamodel.

    Architecture:
        1. Gating Network: attention + context -> alpha (5m weight)
        2. Weighted Average: alpha * pred_5m + (1-alpha) * pred_15m
        3. Residual Network: all_features -> residual correction
        4. Final: weighted_avg + residual
        5. Confidence Network: features + final_pred -> confidence
    """

    def __init__(self, config: NeuralGatingConfig):
        super().__init__()
        self.config = config

        # Input dimensions
        gating_input_dim = len(config.attention_features) + len(config.context_features)  # 6
        all_features_dim = len(config.prediction_features) + gating_input_dim  # 10

        # Networks
        self.gating = GatingNetwork(
            gating_input_dim,
            config.gating_hidden_dims,
            config.dropout
        )

        self.residual = ResidualNetwork(
            all_features_dim,
            config.residual_hidden_dims,
            config.dropout
        )

        # Confidence network gets all features + gating weights + final predictions
        confidence_input_dim = all_features_dim + 2 + 2  # +2 alphas +2 finals
        self.confidence = ConfidenceNetwork(
            confidence_input_dim,
            config.confidence_hidden_dims,
            config.dropout
        )

    def forward(self, x):
        """
        Forward pass.

        Args:
            x: Tensor of shape (batch, 10) with all meta features
               Order: pred_high_5m, pred_low_5m, pred_high_15m, pred_low_15m,
                      attention_5m, attention_15m, attention_class_5m, attention_class_15m,
                      ATR_ratio, volume_z

        Returns:
            delta_high_final, delta_low_final, confidence_proba, alpha_high, alpha_low
        """
        # Extract components
        pred_high_5m = x[:, 0:1]
        pred_low_5m = x[:, 1:2]
        pred_high_15m = x[:, 2:3]
        pred_low_15m = x[:, 3:4]

        # Gating input: attention + context features (indices 4-9)
        gating_input = x[:, 4:]  # attention_5m, attention_15m, classes, ATR_ratio, volume_z

        # Get gating weights
        alphas = self.gating(gating_input)  # (batch, 2) -> alpha_high, alpha_low
        alpha_high = alphas[:, 0:1]
        alpha_low = alphas[:, 1:2]

        # Weighted average
        weighted_high = alpha_high * pred_high_5m + (1 - alpha_high) * pred_high_15m
        weighted_low = alpha_low * pred_low_5m + (1 - alpha_low) * pred_low_15m

        # Residual correction
        residuals = self.residual(x)  # (batch, 2)
        residual_high = residuals[:, 0:1]
        residual_low = residuals[:, 1:2]

        # Final predictions
        delta_high_final = weighted_high + residual_high
        delta_low_final = weighted_low + residual_low

        # Ensure non-negative (using softplus for smooth gradient)
        delta_high_final = nn.functional.softplus(delta_high_final)
        delta_low_final = nn.functional.softplus(delta_low_final)

        # Confidence
        conf_input = torch.cat([
            x, alphas, delta_high_final, delta_low_final
        ], dim=1)
        confidence_proba = self.confidence(conf_input)

        return delta_high_final, delta_low_final, confidence_proba, alpha_high, alpha_low


class NeuralGatingMetamodelWrapper:
    """
    Wrapper class for training and inference with NeuralGatingMetamodel.

    This provides a scikit-learn-like interface for consistency with
    the XGBoost version.
    """

    def __init__(self, symbol: str, config: NeuralGatingConfig = None):
        if not HAS_TORCH:
            raise ImportError("PyTorch required for NeuralGatingMetamodel")

        self.symbol = symbol
        self.config = config or NeuralGatingConfig()

        # Feature names
        self.feature_names = (
            self.config.prediction_features +
            self.config.attention_features +
            self.config.context_features
        )

        # Model
        self.model: Optional[NeuralGatingMetamodel] = None
        self.scaler: Optional[StandardScaler] = None

        # Device
        if self.config.device == 'auto':
            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        else:
            self.device = torch.device(self.config.device)

        # Training state
        self.is_fitted = False
        self.training_history: List[Dict] = []
        self.best_val_loss = float('inf')

        logger.info(f"Initialized NeuralGatingMetamodel for {symbol}")
        logger.info(f"  Device: {self.device}")
        logger.info(f"  Features: {len(self.feature_names)}")

    def _validate_features(self, X: pd.DataFrame) -> pd.DataFrame:
        """Validate and select features."""
        missing = set(self.feature_names) - set(X.columns)
        if missing:
            logger.warning(f"Missing features: {missing}")
            for feat in missing:
                X[feat] = 0.0
        return X[self.feature_names]

    def _compute_confidence_target(
        self,
        target_high: np.ndarray,
        target_low: np.ndarray,
        pred_high: np.ndarray,
        pred_low: np.ndarray
    ) -> np.ndarray:
        """Compute binary confidence target."""
        long_favorable = target_high >= pred_high * 0.8
        short_favorable = target_low >= pred_low * 0.8
        return (long_favorable | short_favorable).astype(np.float32)

    def fit(
        self,
        meta_features: pd.DataFrame,
        target_high: np.ndarray,
        target_low: np.ndarray,
        sample_weights: np.ndarray = None
    ) -> 'NeuralGatingMetamodelWrapper':
        """
        Train the neural gating metamodel.
        """
        logger.info(f"\n{'='*60}")
        logger.info(f"Training NeuralGatingMetamodel for {self.symbol}")
        logger.info(f"{'='*60}")

        # Validate features
        X = self._validate_features(meta_features)
        X_values = X.values.astype(np.float32)

        # Remove invalid samples
        valid_mask = (
            ~np.isnan(X_values).any(axis=1) &
            ~np.isnan(target_high) &
            ~np.isnan(target_low)
        )

        X_valid = X_values[valid_mask]
        y_high_valid = target_high[valid_mask].astype(np.float32)
        y_low_valid = target_low[valid_mask].astype(np.float32)

        n_valid = len(X_valid)
        logger.info(f"Valid samples: {n_valid}")

        if n_valid < self.config.min_train_samples:
            raise ValueError(f"Insufficient data: {n_valid} < {self.config.min_train_samples}")

        # Scale features
        self.scaler = StandardScaler()
        X_scaled = self.scaler.fit_transform(X_valid)

        # Train/val split
        split_idx = int(n_valid * (1 - self.config.val_split))
        X_train, X_val = X_scaled[:split_idx], X_scaled[split_idx:]
        y_high_train, y_high_val = y_high_valid[:split_idx], y_high_valid[split_idx:]
        y_low_train, y_low_val = y_low_valid[:split_idx], y_low_valid[split_idx:]

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

        # Create model
        self.model = NeuralGatingMetamodel(self.config).to(self.device)

        # Get initial predictions for confidence target
        # Use simple average as initial prediction
        pred_high_init = (X_valid[:, 0] + X_valid[:, 2]) / 2  # avg of 5m and 15m
        pred_low_init = (X_valid[:, 1] + X_valid[:, 3]) / 2

        y_conf_valid = self._compute_confidence_target(
            y_high_valid, y_low_valid,
            pred_high_init, pred_low_init
        )
        y_conf_train, y_conf_val = y_conf_valid[:split_idx], y_conf_valid[split_idx:]

        # Create DataLoaders
        train_dataset = TensorDataset(
            torch.FloatTensor(X_train),
            torch.FloatTensor(y_high_train.reshape(-1, 1)),
            torch.FloatTensor(y_low_train.reshape(-1, 1)),
            torch.FloatTensor(y_conf_train.reshape(-1, 1))
        )

        val_dataset = TensorDataset(
            torch.FloatTensor(X_val),
            torch.FloatTensor(y_high_val.reshape(-1, 1)),
            torch.FloatTensor(y_low_val.reshape(-1, 1)),
            torch.FloatTensor(y_conf_val.reshape(-1, 1))
        )

        train_loader = DataLoader(
            train_dataset,
            batch_size=self.config.batch_size,
            shuffle=True
        )
        val_loader = DataLoader(
            val_dataset,
            batch_size=self.config.batch_size,
            shuffle=False
        )

        # Optimizer and loss
        optimizer = optim.Adam(self.model.parameters(), lr=self.config.learning_rate)
        scheduler = optim.lr_scheduler.ReduceLROnPlateau(
            optimizer, mode='min', factor=0.5, patience=5
        )

        mse_loss = nn.MSELoss()
        bce_loss = nn.BCELoss()

        # Training loop
        best_val_loss = float('inf')
        patience_counter = 0
        best_state = None

        for epoch in range(self.config.epochs):
            # Training
            self.model.train()
            train_loss = 0.0

            for batch_x, batch_y_high, batch_y_low, batch_y_conf in train_loader:
                batch_x = batch_x.to(self.device)
                batch_y_high = batch_y_high.to(self.device)
                batch_y_low = batch_y_low.to(self.device)
                batch_y_conf = batch_y_conf.to(self.device)

                optimizer.zero_grad()

                pred_high, pred_low, pred_conf, _, _ = self.model(batch_x)

                loss_high = mse_loss(pred_high, batch_y_high)
                loss_low = mse_loss(pred_low, batch_y_low)
                loss_conf = bce_loss(pred_conf, batch_y_conf)

                # Combined loss (weighted)
                loss = loss_high + loss_low + 0.5 * loss_conf

                loss.backward()
                optimizer.step()

                train_loss += loss.item()

            train_loss /= len(train_loader)

            # Validation
            self.model.eval()
            val_loss = 0.0

            with torch.no_grad():
                for batch_x, batch_y_high, batch_y_low, batch_y_conf in val_loader:
                    batch_x = batch_x.to(self.device)
                    batch_y_high = batch_y_high.to(self.device)
                    batch_y_low = batch_y_low.to(self.device)
                    batch_y_conf = batch_y_conf.to(self.device)

                    pred_high, pred_low, pred_conf, _, _ = self.model(batch_x)

                    loss_high = mse_loss(pred_high, batch_y_high)
                    loss_low = mse_loss(pred_low, batch_y_low)
                    loss_conf = bce_loss(pred_conf, batch_y_conf)

                    loss = loss_high + loss_low + 0.5 * loss_conf
                    val_loss += loss.item()

            val_loss /= len(val_loader)

            scheduler.step(val_loss)

            self.training_history.append({
                'epoch': epoch + 1,
                'train_loss': train_loss,
                'val_loss': val_loss,
                'lr': optimizer.param_groups[0]['lr']
            })

            if (epoch + 1) % 10 == 0:
                logger.info(f"  Epoch {epoch + 1}: train_loss={train_loss:.4f}, val_loss={val_loss:.4f}")

            # Early stopping
            if val_loss < best_val_loss:
                best_val_loss = val_loss
                patience_counter = 0
                best_state = {k: v.cpu().clone() for k, v in self.model.state_dict().items()}
            else:
                patience_counter += 1
                if patience_counter >= self.config.early_stopping_patience:
                    logger.info(f"  Early stopping at epoch {epoch + 1}")
                    break

        # Restore best model
        if best_state:
            self.model.load_state_dict(best_state)

        self.best_val_loss = best_val_loss

        # Evaluate final metrics
        self._evaluate_final(X_val, y_high_val, y_low_val, y_conf_val)

        self.is_fitted = True
        logger.info(f"NeuralGatingMetamodel training complete")

        return self

    def _evaluate_final(
        self,
        X_val: np.ndarray,
        y_high_val: np.ndarray,
        y_low_val: np.ndarray,
        y_conf_val: np.ndarray
    ):
        """Evaluate final model performance."""
        self.model.eval()

        with torch.no_grad():
            X_tensor = torch.FloatTensor(X_val).to(self.device)
            pred_high, pred_low, pred_conf, alpha_high, alpha_low = self.model(X_tensor)

            pred_high = pred_high.cpu().numpy().flatten()
            pred_low = pred_low.cpu().numpy().flatten()
            pred_conf = pred_conf.cpu().numpy().flatten()
            alpha_high = alpha_high.cpu().numpy().flatten()
            alpha_low = alpha_low.cpu().numpy().flatten()

        # Regression metrics
        mae_high = mean_absolute_error(y_high_val, pred_high)
        mae_low = mean_absolute_error(y_low_val, pred_low)
        rmse_high = np.sqrt(mean_squared_error(y_high_val, pred_high))
        rmse_low = np.sqrt(mean_squared_error(y_low_val, pred_low))
        r2_high = r2_score(y_high_val, pred_high)
        r2_low = r2_score(y_low_val, pred_low)

        logger.info(f"\nFinal Validation Metrics:")
        logger.info(f"  HIGH MAE: {mae_high:.4f}, RMSE: {rmse_high:.4f}, R2: {r2_high:.4f}")
        logger.info(f"  LOW MAE: {mae_low:.4f}, RMSE: {rmse_low:.4f}, R2: {r2_low:.4f}")

        # Gating analysis
        logger.info(f"\nGating Analysis:")
        logger.info(f"  Alpha HIGH (5m weight): mean={alpha_high.mean():.3f}, std={alpha_high.std():.3f}")
        logger.info(f"  Alpha LOW (5m weight): mean={alpha_low.mean():.3f}, std={alpha_low.std():.3f}")

        # Confidence metrics
        conf_pred = (pred_conf > 0.5).astype(int)
        conf_acc = accuracy_score(y_conf_val, conf_pred)
        conf_f1 = f1_score(y_conf_val, conf_pred)

        logger.info(f"\nConfidence Metrics:")
        logger.info(f"  Accuracy: {conf_acc:.4f}")
        logger.info(f"  F1 Score: {conf_f1:.4f}")

        # Compare to simple average
        # Unscale predictions to get original 5m/15m values
        X_unscaled = self.scaler.inverse_transform(X_val)
        avg_high = (X_unscaled[:, 0] + X_unscaled[:, 2]) / 2
        avg_low = (X_unscaled[:, 1] + X_unscaled[:, 3]) / 2

        mae_avg_high = mean_absolute_error(y_high_val, avg_high)
        mae_avg_low = mean_absolute_error(y_low_val, avg_low)

        improvement_high = (mae_avg_high - mae_high) / mae_avg_high * 100
        improvement_low = (mae_avg_low - mae_low) / mae_avg_low * 100

        logger.info(f"\nImprovement over simple average:")
        logger.info(f"  HIGH: {improvement_high:.1f}%")
        logger.info(f"  LOW: {improvement_low:.1f}%")

        self.training_metrics = {
            'mae_high': mae_high,
            'mae_low': mae_low,
            'rmse_high': rmse_high,
            'rmse_low': rmse_low,
            'r2_high': r2_high,
            'r2_low': r2_low,
            'conf_accuracy': conf_acc,
            'conf_f1': conf_f1,
            'improvement_high': improvement_high,
            'improvement_low': improvement_low,
            'alpha_high_mean': float(alpha_high.mean()),
            'alpha_low_mean': float(alpha_low.mean())
        }

    def predict(
        self,
        meta_features: Union[pd.DataFrame, np.ndarray]
    ) -> NeuralGatingPrediction:
        """Generate predictions."""
        if not self.is_fitted:
            raise ValueError("Model not fitted")

        # Handle input
        if isinstance(meta_features, pd.DataFrame):
            X = self._validate_features(meta_features).values.astype(np.float32)
        else:
            X = meta_features.astype(np.float32)
            if X.ndim == 1:
                X = X.reshape(1, -1)

        # Handle NaN
        X = np.nan_to_num(X, nan=0.0)

        # Scale
        X_scaled = self.scaler.transform(X)

        # Predict
        self.model.eval()
        with torch.no_grad():
            X_tensor = torch.FloatTensor(X_scaled).to(self.device)
            pred_high, pred_low, pred_conf, alpha_high, alpha_low = self.model(X_tensor)

            pred_high = pred_high.cpu().numpy().flatten()
            pred_low = pred_low.cpu().numpy().flatten()
            pred_conf = pred_conf.cpu().numpy().flatten()
            alpha_high = alpha_high.cpu().numpy().flatten()
            alpha_low = alpha_low.cpu().numpy().flatten()

        # Ensure non-negative
        pred_high = np.maximum(pred_high, 0)
        pred_low = np.maximum(pred_low, 0)

        return NeuralGatingPrediction(
            delta_high_final=pred_high,
            delta_low_final=pred_low,
            confidence=(pred_conf > 0.5).astype(int),
            confidence_proba=pred_conf,
            alpha_high=alpha_high,
            alpha_low=alpha_low
        )

    def predict_single(
        self,
        meta_features: Union[pd.Series, Dict, np.ndarray]
    ) -> Tuple[float, float, int, float, float, float]:
        """Predict for a single sample."""
        if isinstance(meta_features, dict):
            meta_features = pd.DataFrame([meta_features])
        elif isinstance(meta_features, pd.Series):
            meta_features = pd.DataFrame([meta_features.to_dict()])
        elif isinstance(meta_features, np.ndarray):
            if meta_features.ndim == 1:
                meta_features = meta_features.reshape(1, -1)

        pred = self.predict(meta_features)
        return (
            float(pred.delta_high_final[0]),
            float(pred.delta_low_final[0]),
            int(pred.confidence[0]),
            float(pred.confidence_proba[0]),
            float(pred.alpha_high[0]),
            float(pred.alpha_low[0])
        )

    def get_training_summary(self) -> Dict[str, Any]:
        """Get training summary."""
        if not self.is_fitted:
            return {'is_fitted': False}

        return {
            'is_fitted': True,
            'symbol': self.symbol,
            'architecture': 'neural_gating',
            'device': str(self.device),
            'metrics': self.training_metrics,
            'training_history': self.training_history[-10:],  # Last 10 epochs
            'best_val_loss': self.best_val_loss,
            'config': {
                'gating_hidden_dims': self.config.gating_hidden_dims,
                'residual_hidden_dims': self.config.residual_hidden_dims,
                'epochs': len(self.training_history),
                'learning_rate': self.config.learning_rate
            }
        }

    def save(self, path: str):
        """Save model to disk."""
        path = Path(path)
        path.mkdir(parents=True, exist_ok=True)

        # Save PyTorch model
        torch.save(self.model.state_dict(), path / 'model.pt')

        # Save scaler
        joblib.dump(self.scaler, path / 'scaler.joblib')

        # Save metadata
        metadata = {
            'symbol': self.symbol,
            'config': {
                'prediction_features': self.config.prediction_features,
                'attention_features': self.config.attention_features,
                'context_features': self.config.context_features,
                'gating_hidden_dims': self.config.gating_hidden_dims,
                'residual_hidden_dims': self.config.residual_hidden_dims,
                'confidence_hidden_dims': self.config.confidence_hidden_dims,
                'dropout': self.config.dropout
            },
            'feature_names': self.feature_names,
            'training_metrics': self.training_metrics if hasattr(self, 'training_metrics') else None,
            'training_history': self.training_history,
            'best_val_loss': self.best_val_loss,
            'is_fitted': self.is_fitted
        }
        joblib.dump(metadata, path / 'metadata.joblib')

        logger.info(f"NeuralGatingMetamodel saved to {path}")

    @classmethod
    def load(cls, path: str) -> 'NeuralGatingMetamodelWrapper':
        """Load model from disk."""
        path = Path(path)

        # Load metadata
        metadata = joblib.load(path / 'metadata.joblib')

        # Reconstruct config
        config = NeuralGatingConfig(**metadata['config'])

        # Create instance
        wrapper = cls(metadata['symbol'], config)
        wrapper.is_fitted = metadata['is_fitted']
        wrapper.training_history = metadata['training_history']
        wrapper.best_val_loss = metadata['best_val_loss']

        if metadata['training_metrics']:
            wrapper.training_metrics = metadata['training_metrics']

        # Load scaler
        wrapper.scaler = joblib.load(path / 'scaler.joblib')

        # Create and load model
        wrapper.model = NeuralGatingMetamodel(config).to(wrapper.device)
        wrapper.model.load_state_dict(torch.load(path / 'model.pt', map_location=wrapper.device))
        wrapper.model.eval()

        logger.info(f"NeuralGatingMetamodel loaded from {path}")
        return wrapper


if __name__ == "__main__":
    print("Testing NeuralGatingMetamodel...")

    if not HAS_TORCH:
        print("PyTorch not available - skipping test")
    else:
        np.random.seed(42)
        torch.manual_seed(42)

        n = 5000

        # Simulate meta features
        pred_high_5m = np.random.randn(n) * 5 + 10
        pred_low_5m = np.random.randn(n) * 4 + 8
        pred_high_15m = np.random.randn(n) * 6 + 12
        pred_low_15m = np.random.randn(n) * 5 + 9
        attention_5m = np.random.rand(n) * 2 + 0.5
        attention_15m = np.random.rand(n) * 2 + 0.5
        attention_class_5m = np.random.choice([0, 1, 2], n).astype(float)
        attention_class_15m = np.random.choice([0, 1, 2], n).astype(float)
        atr_ratio = np.random.rand(n) + 0.5
        volume_z = np.random.randn(n)

        meta_features = pd.DataFrame({
            'pred_high_5m': pred_high_5m,
            'pred_low_5m': pred_low_5m,
            'pred_high_15m': pred_high_15m,
            'pred_low_15m': pred_low_15m,
            'attention_5m': attention_5m,
            'attention_15m': attention_15m,
            'attention_class_5m': attention_class_5m,
            'attention_class_15m': attention_class_15m,
            'ATR_ratio': atr_ratio,
            'volume_z': volume_z
        })

        # Simulate targets
        target_high = (pred_high_5m + pred_high_15m) / 2 + np.random.randn(n) * 3
        target_low = (pred_low_5m + pred_low_15m) / 2 + np.random.randn(n) * 3
        target_high = np.maximum(target_high, 0)
        target_low = np.maximum(target_low, 0)

        # Test training
        config = NeuralGatingConfig(
            min_train_samples=1000,
            epochs=50,
            early_stopping_patience=5
        )

        model = NeuralGatingMetamodelWrapper('XAUUSD', config)
        model.fit(meta_features, target_high, target_low)

        # Test prediction
        print("\nTesting predictions...")
        test_features = meta_features.iloc[-100:]
        prediction = model.predict(test_features)

        print(f"Delta high: mean={prediction.delta_high_final.mean():.2f}")
        print(f"Delta low: mean={prediction.delta_low_final.mean():.2f}")
        print(f"Alpha high (5m weight): mean={prediction.alpha_high.mean():.3f}")
        print(f"Alpha low (5m weight): mean={prediction.alpha_low.mean():.3f}")
        print(f"Confidence: {prediction.confidence.mean():.1%}")

        # Test save/load
        print("\nTesting save/load...")
        model.save('/tmp/test_neural_gating')
        loaded = NeuralGatingMetamodelWrapper.load('/tmp/test_neural_gating')

        pred2 = loaded.predict(test_features.iloc[:5])
        print(f"Loaded model alpha: {pred2.alpha_high}")

        print("\nTest complete!")
