#!/usr/bin/env python3
"""
Neural Gating Metamodel - Complete Model
=========================================
Full metamodel that integrates:
1. GatingNetwork - Learns strategy weights from market context
2. EnsemblePipeline - Orchestrates 5 strategy predictions
3. ConfidenceCalibrator - Calibrates final confidence scores

Produces interpretable predictions with:
- direction: -1 to 1 (bearish to bullish)
- magnitude: Expected return percentage
- confidence: 0 to 1 (calibrated)
- strategy_weights: Contribution of each strategy
- dominant_strategy: Most influential strategy

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
from typing import Dict, List, Optional, Tuple, Any, Union
from dataclasses import dataclass, field
from pathlib import Path
import joblib
from loguru import logger
import json
import time

from .gating_network import GatingNetwork, GatingConfig, ContextEncoder, StrategyOutputEncoder
from .ensemble_pipeline import EnsemblePipeline, EnsemblePrediction, StrategyPredictions
from .calibration import ConfidenceCalibrator, CalibrationMetrics


@dataclass
class MetamodelConfig:
    """Configuration for the Neural Gating Metamodel."""

    # Gating network configuration
    n_strategies: int = 5
    strategy_output_dim: int = 8
    context_dim: int = 32
    gating_hidden_dims: Tuple[int, ...] = (256, 128)
    gating_dropout: float = 0.3
    gating_temperature: float = 1.0
    entropy_weight: float = 0.1
    min_weight: float = 0.02

    # Calibration configuration
    calibration_method: str = 'isotonic'
    calibration_bins: int = 15

    # Model paths
    models_dir: Optional[str] = None

    # Inference settings
    parallel_inference: bool = True
    device: Optional[str] = None

    def to_dict(self) -> Dict[str, Any]:
        return {
            'n_strategies': self.n_strategies,
            'strategy_output_dim': self.strategy_output_dim,
            'context_dim': self.context_dim,
            'gating_hidden_dims': self.gating_hidden_dims,
            'gating_dropout': self.gating_dropout,
            'gating_temperature': self.gating_temperature,
            'entropy_weight': self.entropy_weight,
            'min_weight': self.min_weight,
            'calibration_method': self.calibration_method,
            'calibration_bins': self.calibration_bins,
            'models_dir': self.models_dir,
            'parallel_inference': self.parallel_inference,
            'device': self.device
        }

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> 'MetamodelConfig':
        return cls(**{k: v for k, v in d.items() if k in cls.__dataclass_fields__})


@dataclass
class MetamodelPrediction:
    """Complete prediction from the Neural Gating Metamodel."""

    # Core predictions
    direction: float  # -1 (bearish) to 1 (bullish)
    magnitude: float  # Expected return percentage
    confidence: float  # 0 to 1, calibrated
    raw_confidence: float  # Before calibration

    # Strategy analysis
    strategy_weights: Dict[str, float]
    dominant_strategy: str
    strategy_contributions: Dict[str, float]  # Weighted contribution to final

    # Individual predictions
    strategy_predictions: Dict[str, Dict[str, float]]

    # Metadata
    timestamp: Optional[pd.Timestamp] = None
    symbol: Optional[str] = None
    processing_time_ms: float = 0.0
    calibration_applied: bool = True

    def to_dict(self) -> Dict[str, Any]:
        return {
            'direction': float(self.direction),
            'magnitude': float(self.magnitude),
            'confidence': float(self.confidence),
            'raw_confidence': float(self.raw_confidence),
            'strategy_weights': self.strategy_weights,
            'dominant_strategy': self.dominant_strategy,
            'strategy_contributions': self.strategy_contributions,
            'strategy_predictions': self.strategy_predictions,
            'timestamp': str(self.timestamp) if self.timestamp else None,
            'symbol': self.symbol,
            'processing_time_ms': self.processing_time_ms,
            'calibration_applied': self.calibration_applied
        }

    @property
    def signal(self) -> str:
        """Get trading signal."""
        if self.confidence < 0.5:
            return 'HOLD'
        if self.direction > 0.3:
            return 'BUY'
        elif self.direction < -0.3:
            return 'SELL'
        return 'HOLD'

    @property
    def signal_strength(self) -> float:
        """Get signal strength (0 to 1)."""
        return abs(self.direction) * self.confidence

    @property
    def risk_adjusted_signal(self) -> float:
        """Get risk-adjusted signal considering confidence."""
        return self.direction * self.confidence

    def get_trading_recommendation(self) -> Dict[str, Any]:
        """Get structured trading recommendation."""
        return {
            'action': self.signal,
            'direction': 'LONG' if self.direction > 0 else 'SHORT' if self.direction < 0 else 'NEUTRAL',
            'strength': self.signal_strength,
            'confidence': self.confidence,
            'dominant_strategy': self.dominant_strategy,
            'risk_level': 'LOW' if self.confidence > 0.7 else 'MEDIUM' if self.confidence > 0.5 else 'HIGH'
        }


class NeuralGatingMetamodel:
    """
    Neural Gating Metamodel - Combines 5 trading strategies with learned weighting.

    This model learns which strategy performs best in different market conditions
    and dynamically adjusts weights. It uses entropy regularization to prevent
    collapse to a single strategy.

    Components:
    1. EnsemblePipeline: Orchestrates 5 strategies (PVA, MRD, VBP, MSA, MTS)
    2. GatingNetwork: Learns context-dependent strategy weights
    3. ConfidenceCalibrator: Ensures confidence scores are well-calibrated

    Key Features:
    - Dynamic strategy weighting based on market context
    - Entropy regularization prevents strategy collapse
    - Calibrated confidence scores
    - Interpretable strategy contributions
    - Support for walk-forward training

    Usage:
        model = NeuralGatingMetamodel()
        model.fit(training_data, training_targets)

        prediction = model.predict(df_ohlcv)
        print(f"Direction: {prediction.direction:.3f}")
        print(f"Confidence: {prediction.confidence:.3f}")
        print(f"Dominant: {prediction.dominant_strategy}")
    """

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

    def __init__(self, config: Optional[MetamodelConfig] = None):
        """
        Initialize the Neural Gating Metamodel.

        Args:
            config: Model configuration
        """
        self.config = config or MetamodelConfig()
        self.device = self.config.device or ('cuda' if torch.cuda.is_available() else 'cpu')

        # Initialize gating configuration
        gating_config = GatingConfig(
            n_strategies=self.config.n_strategies,
            strategy_output_dim=self.config.strategy_output_dim,
            context_dim=self.config.context_dim,
            hidden_dims=self.config.gating_hidden_dims,
            dropout=self.config.gating_dropout,
            temperature=self.config.gating_temperature,
            entropy_weight=self.config.entropy_weight,
            min_weight=self.config.min_weight
        )

        # Initialize components
        self.ensemble_pipeline = EnsemblePipeline(
            gating_config=gating_config,
            models_dir=self.config.models_dir,
            device=self.device,
            parallel_inference=self.config.parallel_inference
        )

        self.gating_network = GatingNetwork(gating_config).to(self.device)
        self.context_encoder = ContextEncoder(output_dim=self.config.context_dim).to(self.device)
        self.strategy_encoder = StrategyOutputEncoder(
            output_dim=self.config.strategy_output_dim
        ).to(self.device)

        self.calibrator = ConfidenceCalibrator(
            method=self.config.calibration_method,
            n_bins=self.config.calibration_bins
        )

        # Training state
        self._is_trained = False
        self._training_metrics: Dict[str, float] = {}
        self._loaded_symbols: List[str] = []

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

    def load_strategies(
        self,
        symbols: List[str],
        strategies_to_load: Optional[List[str]] = None
    ) -> Dict[str, bool]:
        """
        Load strategy models for given symbols.

        Args:
            symbols: Trading symbols to load
            strategies_to_load: Specific strategies (default: all)

        Returns:
            Load status per strategy
        """
        status = self.ensemble_pipeline.load_strategies(symbols, strategies_to_load)
        self._loaded_symbols = symbols
        return status

    def fit(
        self,
        training_data: List[Tuple[pd.DataFrame, np.ndarray]],
        val_data: Optional[List[Tuple[pd.DataFrame, np.ndarray]]] = None,
        epochs: int = 50,
        batch_size: int = 32,
        learning_rate: float = 1e-3,
        verbose: bool = True
    ) -> Dict[str, float]:
        """
        Train the complete metamodel.

        This method:
        1. Generates predictions from all 5 strategies
        2. Trains the gating network to weight strategies
        3. Trains the confidence calibrator

        Args:
            training_data: List of (df_ohlcv, y_direction) tuples
            val_data: Optional validation data
            epochs: Training epochs for gating network
            batch_size: Batch size
            learning_rate: Learning rate
            verbose: Print progress

        Returns:
            Training metrics
        """
        logger.info(f"Training metamodel with {len(training_data)} samples...")
        start_time = time.time()

        # Step 1: Generate strategy predictions for all training samples
        all_strategy_preds = []
        all_contexts = []
        all_targets = []

        for i, (df, y_target) in enumerate(training_data):
            if verbose and (i + 1) % 100 == 0:
                logger.info(f"Processing sample {i + 1}/{len(training_data)}")

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

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

            # Store target
            all_targets.append(y_target)

        # Step 2: Train gating network
        gating_metrics = self._train_gating(
            all_strategy_preds,
            all_contexts,
            all_targets,
            epochs,
            batch_size,
            learning_rate,
            verbose
        )

        # Step 3: Train calibrator
        cal_metrics = self._train_calibrator(
            all_strategy_preds,
            all_contexts,
            all_targets
        )

        # Combine metrics
        self._training_metrics = {
            **gating_metrics,
            **cal_metrics,
            'training_time_s': time.time() - start_time
        }

        self._is_trained = True

        if verbose:
            logger.info(f"Training complete in {self._training_metrics['training_time_s']:.1f}s")
            logger.info(f"Gating loss: {gating_metrics.get('final_loss', 'N/A'):.4f}")
            logger.info(f"Calibration ECE: {cal_metrics.get('calibration_ece', 'N/A'):.4f}")

        return self._training_metrics

    def _train_gating(
        self,
        strategy_preds: List[StrategyPredictions],
        contexts: List[torch.Tensor],
        targets: List[np.ndarray],
        epochs: int,
        batch_size: int,
        learning_rate: float,
        verbose: bool
    ) -> Dict[str, float]:
        """Train the gating network."""
        logger.info("Training gating network...")

        self.gating_network.train()
        self.context_encoder.train()
        self.strategy_encoder.train()

        # Combine all parameters
        all_params = (
            list(self.gating_network.parameters()) +
            list(self.context_encoder.parameters()) +
            list(self.strategy_encoder.parameters())
        )

        optimizer = torch.optim.AdamW(all_params, lr=learning_rate, weight_decay=1e-4)
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, epochs)

        # Prepare data
        n_samples = len(strategy_preds)
        best_loss = float('inf')
        history = []

        for epoch in range(epochs):
            epoch_loss = 0.0
            epoch_direction_loss = 0.0
            epoch_entropy_loss = 0.0

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

            for batch_start in range(0, n_samples, batch_size):
                batch_indices = indices[batch_start:batch_start + batch_size]

                optimizer.zero_grad()

                batch_loss = 0.0
                batch_direction_loss = 0.0
                batch_entropy_loss = 0.0

                for idx in batch_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.gating_network(strategy_tensors, context)

                    # Compute weighted direction prediction
                    pred_direction = self._compute_weighted_direction(
                        strategy_pred,
                        weights
                    )

                    # Direction loss (MSE to target direction)
                    target_direction = torch.tensor(
                        [float(target)],
                        device=self.device,
                        dtype=torch.float32
                    )
                    direction_loss = nn.functional.mse_loss(pred_direction, target_direction)

                    # Entropy loss (regularization)
                    entropy_loss = self.gating_network.compute_entropy_loss(weights)

                    # Total loss
                    loss = direction_loss + entropy_loss

                    batch_loss += loss
                    batch_direction_loss += direction_loss.item()
                    batch_entropy_loss += entropy_loss.item()

                # Average over batch
                batch_loss = batch_loss / len(batch_indices)
                batch_loss.backward()

                # Gradient clipping
                torch.nn.utils.clip_grad_norm_(all_params, max_norm=1.0)

                optimizer.step()

                epoch_loss += batch_loss.item()
                epoch_direction_loss += batch_direction_loss / len(batch_indices)
                epoch_entropy_loss += batch_entropy_loss / len(batch_indices)

            # Average over epoch
            n_batches = (n_samples + batch_size - 1) // batch_size
            epoch_loss /= n_batches
            epoch_direction_loss /= n_batches
            epoch_entropy_loss /= n_batches

            scheduler.step()
            history.append(epoch_loss)

            if epoch_loss < best_loss:
                best_loss = epoch_loss

            if verbose and (epoch + 1) % 10 == 0:
                logger.info(
                    f"Epoch {epoch + 1}/{epochs} - "
                    f"Loss: {epoch_loss:.4f}, "
                    f"Dir: {epoch_direction_loss:.4f}, "
                    f"Entropy: {epoch_entropy_loss:.4f}"
                )

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

        return {
            'final_loss': history[-1] if history else 0.0,
            'best_loss': best_loss,
            'gating_epochs': epochs
        }

    def _encode_strategy_predictions(
        self,
        strategy_preds: StrategyPredictions
    ) -> Dict[str, torch.Tensor]:
        """Convert strategy predictions to tensor format for gating."""
        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.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 from each strategy
            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 += weight * strat_dir

        return direction

    def _train_calibrator(
        self,
        strategy_preds: List[StrategyPredictions],
        contexts: List[torch.Tensor],
        targets: List[np.ndarray]
    ) -> Dict[str, float]:
        """Train the confidence calibrator."""
        logger.info("Training confidence calibrator...")

        # Generate raw confidence scores
        raw_confidences = []
        binary_targets = []

        self.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]

                # Get gating weights
                strategy_tensors = self._encode_strategy_predictions(strategy_pred)
                weights = self.gating_network(strategy_tensors, context)

                # Compute raw confidence (weighted average of individual confidences)
                raw_conf = self._compute_raw_confidence(strategy_pred, weights)
                raw_confidences.append(raw_conf)

                # Binary target: 1 if prediction direction matches 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.calibrator.fit(binary_targets, raw_confidences)

        # Evaluate
        metrics = self.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()

            # Extract confidence from each strategy
            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 predict(
        self,
        df: pd.DataFrame,
        symbol: Optional[str] = None
    ) -> MetamodelPrediction:
        """
        Generate prediction from the metamodel.

        Args:
            df: OHLCV DataFrame
            symbol: Trading symbol (optional)

        Returns:
            MetamodelPrediction with full analysis
        """
        start_time = time.time()

        # Get strategy predictions
        strategy_preds = self.ensemble_pipeline.predict_all_strategies(df, symbol)

        # Get context
        context = self.ensemble_pipeline.extract_market_context(df)

        # Get gating weights
        self.gating_network.eval()
        with torch.no_grad():
            strategy_tensors = self._encode_strategy_predictions(strategy_preds)
            weights = self.gating_network(strategy_tensors, context)

            # Compute weighted direction
            direction = self._compute_weighted_direction(strategy_preds, weights)
            direction = float(direction.item())

            # Compute raw confidence
            raw_confidence = self._compute_raw_confidence(strategy_preds, weights)

        # Calibrate confidence
        if self._is_trained and self.calibrator._is_fitted:
            confidence = float(self.calibrator.calibrate(np.array([raw_confidence]))[0])
            calibration_applied = True
        else:
            confidence = raw_confidence
            calibration_applied = False

        # Extract weights as dict
        weights_dict = {
            name: float(weights[0, i].item())
            for i, name in enumerate(self.STRATEGY_NAMES)
        }

        # Compute strategy contributions
        contributions = self._compute_strategy_contributions(
            strategy_preds,
            weights_dict
        )

        # Compute magnitude (weighted average)
        magnitude = self._compute_weighted_magnitude(strategy_preds, weights_dict)

        # Find dominant strategy
        dominant_strategy = max(weights_dict.items(), key=lambda x: x[1])[0]

        # Processing time
        processing_time = (time.time() - start_time) * 1000

        return MetamodelPrediction(
            direction=float(np.clip(direction, -1.0, 1.0)),
            magnitude=magnitude,
            confidence=float(np.clip(confidence, 0.0, 1.0)),
            raw_confidence=raw_confidence,
            strategy_weights=weights_dict,
            dominant_strategy=dominant_strategy,
            strategy_contributions=contributions,
            strategy_predictions={k: v for k, v in strategy_preds.get_available().items()},
            timestamp=df.index[-1] if hasattr(df.index, '__getitem__') else None,
            symbol=symbol,
            processing_time_ms=processing_time,
            calibration_applied=calibration_applied
        )

    def _compute_strategy_contributions(
        self,
        strategy_preds: StrategyPredictions,
        weights_dict: Dict[str, float]
    ) -> Dict[str, float]:
        """Compute each strategy's contribution to the final prediction."""
        available = strategy_preds.get_available()
        contributions = {}

        for name in self.STRATEGY_NAMES:
            if name not in available:
                contributions[name] = 0.0
                continue

            pred = available[name]
            weight = weights_dict.get(name, 0.0)

            # Direction contribution
            if name == 'pva':
                dir_contrib = abs(pred.get('direction', 0.0))
            elif name == 'mrd':
                regime = pred.get('regime', 1.0)
                dir_contrib = abs({0: -1.0, 1: 0.0, 2: 1.0}.get(int(regime), 0.0))
            elif name == 'vbp':
                dir_contrib = abs(pred.get('direction', 0.0))
            elif name == 'msa':
                dir_contrib = abs(pred.get('is_bullish', 0.0) - pred.get('is_bearish', 0.0))
            elif name == 'mts':
                dir_contrib = abs(pred.get('unified_direction', 0.0))
            else:
                dir_contrib = 0.0

            contributions[name] = weight * dir_contrib

        # Normalize contributions to sum to 1
        total = sum(contributions.values())
        if total > 0:
            contributions = {k: v / total for k, v in contributions.items()}

        return contributions

    def _compute_weighted_magnitude(
        self,
        strategy_preds: StrategyPredictions,
        weights_dict: Dict[str, float]
    ) -> float:
        """Compute weighted magnitude prediction."""
        available = strategy_preds.get_available()
        magnitude = 0.0
        total_weight = 0.0

        for name in self.STRATEGY_NAMES:
            if name not in available:
                continue

            pred = available[name]
            weight = weights_dict.get(name, 0.0)

            # Extract magnitude
            if name == 'pva':
                mag = pred.get('magnitude', 0.0)
            elif name == 'vbp':
                mag = pred.get('magnitude', 0.0)
            else:
                mag = 0.01  # Default 1% for strategies without magnitude

            magnitude += weight * mag
            total_weight += weight

        if total_weight > 0:
            magnitude /= total_weight

        return float(magnitude)

    def get_strategy_contributions(self) -> Dict[str, float]:
        """Get current strategy weight distribution."""
        if not hasattr(self, '_last_weights'):
            return {name: 1.0 / len(self.STRATEGY_NAMES) for name in self.STRATEGY_NAMES}
        return self._last_weights

    def save(self, path: str) -> None:
        """
        Save the metamodel to disk.

        Args:
            path: Directory path to save model
        """
        path = Path(path)
        path.mkdir(parents=True, exist_ok=True)

        # Save configuration
        config_path = path / 'config.json'
        with open(config_path, 'w') as f:
            json.dump(self.config.to_dict(), f, indent=2)

        # Save gating network
        gating_path = path / 'gating_network.pt'
        torch.save(self.gating_network.state_dict(), gating_path)

        # Save context encoder
        context_path = path / 'context_encoder.pt'
        torch.save(self.context_encoder.state_dict(), context_path)

        # Save strategy encoder
        strategy_path = path / 'strategy_encoder.pt'
        torch.save(self.strategy_encoder.state_dict(), strategy_path)

        # Save calibrator
        calibrator_path = path / 'calibrator.joblib'
        joblib.dump(self.calibrator, calibrator_path)

        # Save training state
        state_path = path / 'state.json'
        with open(state_path, 'w') as f:
            json.dump({
                'is_trained': self._is_trained,
                'training_metrics': self._training_metrics,
                'loaded_symbols': self._loaded_symbols
            }, f, indent=2)

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

    @classmethod
    def load(cls, path: str, device: Optional[str] = None) -> 'NeuralGatingMetamodel':
        """
        Load metamodel from disk.

        Args:
            path: Directory containing saved model
            device: Device to load to (overrides saved config)

        Returns:
            Loaded NeuralGatingMetamodel
        """
        path = Path(path)

        # Load configuration
        config_path = path / 'config.json'
        with open(config_path, 'r') as f:
            config_dict = json.load(f)

        if device:
            config_dict['device'] = device

        config = MetamodelConfig.from_dict(config_dict)

        # Create model
        model = cls(config)

        # Load gating network
        gating_path = path / 'gating_network.pt'
        model.gating_network.load_state_dict(
            torch.load(gating_path, map_location=model.device)
        )

        # Load context encoder
        context_path = path / 'context_encoder.pt'
        model.context_encoder.load_state_dict(
            torch.load(context_path, map_location=model.device)
        )

        # Load strategy encoder
        strategy_path = path / 'strategy_encoder.pt'
        model.strategy_encoder.load_state_dict(
            torch.load(strategy_path, map_location=model.device)
        )

        # Load calibrator
        calibrator_path = path / 'calibrator.joblib'
        if calibrator_path.exists():
            model.calibrator = joblib.load(calibrator_path)

        # Load training state
        state_path = path / 'state.json'
        if state_path.exists():
            with open(state_path, 'r') as f:
                state = json.load(f)
                model._is_trained = state.get('is_trained', False)
                model._training_metrics = state.get('training_metrics', {})
                model._loaded_symbols = state.get('loaded_symbols', [])

        logger.info(f"Metamodel loaded from {path}")
        return model


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

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

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

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

    # Initialize model
    config = MetamodelConfig()
    model = NeuralGatingMetamodel(config)

    # Load strategies (mock)
    load_status = model.load_strategies(['EURUSD'])
    print(f"Load status: {load_status}")

    # Generate synthetic training data
    training_data = []
    for i in range(50):
        start_idx = i * 3
        end_idx = start_idx + 100
        if end_idx > len(df):
            break
        sample_df = df.iloc[start_idx:end_idx].copy()
        target = np.random.choice([-1.0, 0.0, 1.0])  # Random direction
        training_data.append((sample_df, target))

    print(f"\nTraining with {len(training_data)} samples...")

    # Train model
    metrics = model.fit(training_data, epochs=10, verbose=True)

    print(f"\n--- Training Metrics ---")
    for k, v in metrics.items():
        print(f"  {k}: {v}")

    # Make prediction
    print(f"\n--- Making Prediction ---")
    prediction = model.predict(df.tail(100), symbol='EURUSD')

    print(f"Direction: {prediction.direction:.4f}")
    print(f"Magnitude: {prediction.magnitude:.4f}")
    print(f"Confidence: {prediction.confidence:.4f} (raw: {prediction.raw_confidence:.4f})")
    print(f"Signal: {prediction.signal}")
    print(f"Dominant Strategy: {prediction.dominant_strategy}")
    print(f"Processing Time: {prediction.processing_time_ms:.2f} ms")

    print(f"\n--- Strategy Weights ---")
    for name, weight in prediction.strategy_weights.items():
        print(f"  {name.upper()}: {weight:.4f}")

    print(f"\n--- Strategy Contributions ---")
    for name, contrib in prediction.strategy_contributions.items():
        print(f"  {name.upper()}: {contrib:.4f}")

    # Test save/load
    print(f"\n--- Testing Save/Load ---")
    import tempfile
    with tempfile.TemporaryDirectory() as tmpdir:
        save_path = Path(tmpdir) / 'metamodel'
        model.save(str(save_path))

        loaded_model = NeuralGatingMetamodel.load(str(save_path))
        loaded_pred = loaded_model.predict(df.tail(100), symbol='EURUSD')

        print(f"Original direction: {prediction.direction:.4f}")
        print(f"Loaded direction:   {loaded_pred.direction:.4f}")

    print("\nNeuralGatingMetamodel test complete!")
