#!/usr/bin/env python3
"""
Neural Gating Network
======================
MLP-based gating network that learns to weight different strategy predictions
based on market context. Outputs softmax weights that sum to 1.

Architecture:
    Input: [strategy_outputs (n_strategies * d_strategy), market_context (d_context)]
    -> Linear(input_dim, 256) + BatchNorm + ReLU + Dropout(0.3)
    -> Linear(256, 128) + BatchNorm + ReLU + Dropout(0.3)
    -> Linear(128, n_strategies)
    -> Softmax
    Output: weights (n_strategies,) that sum to 1

The gating network learns which strategy performs best in different market conditions.

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

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass
import numpy as np
from loguru import logger


@dataclass
class GatingConfig:
    """Configuration for the Gating Network."""

    n_strategies: int = 5  # PVA, MRD, VBP, MSA, MTS
    strategy_output_dim: int = 8  # Features per strategy output
    context_dim: int = 32  # Market context features
    hidden_dims: Tuple[int, ...] = (256, 128)
    dropout: float = 0.3
    temperature: float = 1.0  # Softmax temperature
    entropy_weight: float = 0.1  # Regularization weight for entropy
    min_weight: float = 0.02  # Minimum weight per strategy (anti-collapse)


class GatingNetwork(nn.Module):
    """
    Neural Gating Network for strategy weighting.

    Learns to dynamically weight different trading strategies based on
    current market context. Uses entropy regularization to prevent
    collapse to a single strategy.

    Architecture:
        - 3-layer MLP with BatchNorm and Dropout
        - Softmax output with temperature scaling
        - Optional entropy bonus for regularization

    Usage:
        config = GatingConfig(n_strategies=5)
        gating = GatingNetwork(config)

        # Strategy outputs: dict mapping strategy name to encoded features
        strategy_outputs = {
            'pva': torch.randn(batch, 8),
            'mrd': torch.randn(batch, 8),
            ...
        }
        context = torch.randn(batch, 32)  # Market context features

        weights = gating(strategy_outputs, context)
        # weights: (batch, n_strategies) summing to 1
    """

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

    def __init__(self, config: Optional[GatingConfig] = None):
        """
        Initialize the Gating Network.

        Args:
            config: Gating configuration
        """
        super().__init__()

        self.config = config or GatingConfig()

        # Calculate input dimension
        strategy_input_dim = self.config.n_strategies * self.config.strategy_output_dim
        total_input_dim = strategy_input_dim + self.config.context_dim

        # Build MLP layers
        layers = []
        in_dim = total_input_dim

        for hidden_dim in self.config.hidden_dims:
            layers.extend([
                nn.Linear(in_dim, hidden_dim),
                nn.BatchNorm1d(hidden_dim),
                nn.ReLU(inplace=True),
                nn.Dropout(self.config.dropout)
            ])
            in_dim = hidden_dim

        # Output layer
        layers.append(nn.Linear(in_dim, self.config.n_strategies))

        self.mlp = nn.Sequential(*layers)

        # Strategy feature encoders (project each strategy output to common dim)
        self.strategy_encoders = nn.ModuleDict({
            name: nn.Linear(self.config.strategy_output_dim, self.config.strategy_output_dim)
            for name in self.STRATEGY_NAMES
        })

        # Context encoder
        self.context_encoder = nn.Sequential(
            nn.Linear(self.config.context_dim, self.config.context_dim),
            nn.LayerNorm(self.config.context_dim),
            nn.ReLU()
        )

        # Initialize weights
        self._init_weights()

        logger.info(
            f"GatingNetwork initialized: "
            f"input_dim={total_input_dim}, "
            f"n_strategies={self.config.n_strategies}, "
            f"hidden_dims={self.config.hidden_dims}"
        )

    def _init_weights(self):
        """Initialize network weights."""
        for module in self.modules():
            if isinstance(module, nn.Linear):
                nn.init.kaiming_normal_(module.weight, mode='fan_out', nonlinearity='relu')
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
            elif isinstance(module, nn.BatchNorm1d):
                nn.init.ones_(module.weight)
                nn.init.zeros_(module.bias)

        # Initialize output layer to produce uniform weights initially
        with torch.no_grad():
            for layer in self.mlp:
                if isinstance(layer, nn.Linear) and layer.out_features == self.config.n_strategies:
                    nn.init.zeros_(layer.weight)
                    nn.init.zeros_(layer.bias)

    def forward(
        self,
        strategy_outputs: Dict[str, torch.Tensor],
        market_context: torch.Tensor,
        return_logits: bool = False
    ) -> torch.Tensor:
        """
        Forward pass to compute strategy weights.

        Args:
            strategy_outputs: Dictionary mapping strategy name to output tensor
                             Each tensor should be (batch, strategy_output_dim)
            market_context: Market context features (batch, context_dim)
            return_logits: If True, return raw logits instead of softmax weights

        Returns:
            weights: Strategy weights (batch, n_strategies) summing to 1
                    or logits if return_logits=True
        """
        batch_size = market_context.size(0)
        device = market_context.device

        # Encode and concatenate strategy outputs
        encoded_strategies = []
        for name in self.STRATEGY_NAMES:
            if name in strategy_outputs:
                encoded = self.strategy_encoders[name](strategy_outputs[name])
            else:
                # Use zeros for missing strategies
                encoded = torch.zeros(
                    batch_size, self.config.strategy_output_dim, device=device
                )
            encoded_strategies.append(encoded)

        # Concatenate all strategy outputs: (batch, n_strategies * strategy_output_dim)
        strategy_features = torch.cat(encoded_strategies, dim=1)

        # Encode context
        context_features = self.context_encoder(market_context)

        # Concatenate strategy and context features
        combined = torch.cat([strategy_features, context_features], dim=1)

        # Forward through MLP
        logits = self.mlp(combined)

        if return_logits:
            return logits

        # Apply temperature-scaled softmax
        weights = F.softmax(logits / self.config.temperature, dim=1)

        # Apply minimum weight constraint (anti-collapse)
        weights = self._apply_min_weight(weights)

        return weights

    def _apply_min_weight(self, weights: torch.Tensor) -> torch.Tensor:
        """
        Apply minimum weight constraint to prevent strategy collapse.

        Ensures each strategy has at least min_weight contribution.

        Args:
            weights: Raw softmax weights (batch, n_strategies)

        Returns:
            Adjusted weights still summing to 1
        """
        min_w = self.config.min_weight
        n = self.config.n_strategies

        # Calculate available weight for redistribution
        floor = min_w * n  # Total minimum allocation
        if floor >= 1.0:
            # If minimum total exceeds 1, return uniform
            return torch.ones_like(weights) / n

        # Clamp to minimum and rescale remaining weight
        clamped = torch.clamp(weights, min=min_w)
        excess = (clamped.sum(dim=1, keepdim=True) - 1.0)

        # Redistribute excess proportionally from weights above minimum
        above_min = clamped - min_w
        above_min_sum = above_min.sum(dim=1, keepdim=True) + 1e-8
        adjusted = clamped - excess * (above_min / above_min_sum)

        # Ensure non-negative and renormalize
        adjusted = F.relu(adjusted) + 1e-8
        adjusted = adjusted / adjusted.sum(dim=1, keepdim=True)

        return adjusted

    def compute_entropy_loss(self, weights: torch.Tensor) -> torch.Tensor:
        """
        Compute entropy bonus/penalty for regularization.

        Higher entropy = more uniform distribution = less specialization
        Lower entropy = more concentrated = stronger specialization

        We want moderate entropy to prevent collapse while allowing specialization.

        Args:
            weights: Strategy weights (batch, n_strategies)

        Returns:
            Negative entropy (to minimize) scaled by entropy_weight
        """
        # Entropy: -sum(p * log(p))
        entropy = -torch.sum(weights * torch.log(weights + 1e-8), dim=1)

        # Target entropy: ~70% of max entropy (balanced specialization)
        max_entropy = np.log(self.config.n_strategies)
        target_entropy = 0.7 * max_entropy

        # Loss: penalize deviation from target entropy
        entropy_loss = (entropy - target_entropy).pow(2).mean()

        return self.config.entropy_weight * entropy_loss

    def get_strategy_weights(
        self,
        strategy_outputs: Dict[str, torch.Tensor],
        market_context: torch.Tensor
    ) -> Dict[str, float]:
        """
        Get strategy weights as a dictionary.

        Args:
            strategy_outputs: Strategy output tensors
            market_context: Market context tensor

        Returns:
            Dictionary mapping strategy name to weight
        """
        self.eval()
        with torch.no_grad():
            weights = self.forward(strategy_outputs, market_context)
            weights = weights.mean(dim=0).cpu().numpy()

        return {
            name: float(weights[i])
            for i, name in enumerate(self.STRATEGY_NAMES)
        }


class ContextEncoder(nn.Module):
    """
    Encoder for market context features.

    Takes raw market data and produces a context embedding that captures
    current market conditions relevant for strategy selection.

    Features extracted:
    - Volatility metrics (ATR, realized vol, implied vol proxy)
    - Trend strength (ADX, directional movement)
    - Market structure (range vs trending)
    - Volume characteristics
    - Recent performance metrics
    """

    def __init__(
        self,
        input_dim: int = 50,
        output_dim: int = 32,
        hidden_dim: int = 64
    ):
        """
        Initialize context encoder.

        Args:
            input_dim: Number of raw context features
            output_dim: Output embedding dimension
            hidden_dim: Hidden layer dimension
        """
        super().__init__()

        self.encoder = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(hidden_dim, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(hidden_dim, output_dim),
            nn.LayerNorm(output_dim)
        )

        self.output_dim = output_dim

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Encode market context.

        Args:
            x: Raw context features (batch, input_dim)

        Returns:
            Context embedding (batch, output_dim)
        """
        return self.encoder(x)


class StrategyOutputEncoder(nn.Module):
    """
    Encodes strategy prediction outputs to a common representation.

    Each strategy produces different outputs:
    - PVA: direction, magnitude, confidence
    - MRD: regime, duration, continuation_prob
    - VBP: breakout_prob, direction, magnitude
    - MSA: next_bos_direction, poi_reaction_prob, structure_continuation
    - MTS: unified_direction, confidence_by_alignment, optimal_entry_tf

    This encoder normalizes them to a common embedding space.
    """

    STRATEGY_INPUT_DIMS = {
        'pva': 4,   # direction, magnitude, confidence, raw_prediction
        'mrd': 5,   # regime, duration, continuation_prob, reversal_prob, regime_confidence
        'vbp': 5,   # breakout_prob, direction, magnitude, confidence, is_breakout
        'msa': 5,   # bos_conf, poi_reaction_prob, structure_continuation_prob, is_bullish, is_bearish
        'mts': 5,   # unified_direction, confidence, confidence_by_alignment, signal_strength, is_buy
    }

    def __init__(self, output_dim: int = 8):
        """
        Initialize strategy output encoder.

        Args:
            output_dim: Common output dimension for all strategies
        """
        super().__init__()

        self.output_dim = output_dim

        # Create encoder for each strategy
        self.encoders = nn.ModuleDict({
            name: nn.Sequential(
                nn.Linear(in_dim, 16),
                nn.LayerNorm(16),
                nn.ReLU(),
                nn.Linear(16, output_dim),
                nn.LayerNorm(output_dim)
            )
            for name, in_dim in self.STRATEGY_INPUT_DIMS.items()
        })

    def forward(
        self,
        strategy_name: str,
        features: torch.Tensor
    ) -> torch.Tensor:
        """
        Encode strategy output to common space.

        Args:
            strategy_name: Name of the strategy
            features: Strategy output features

        Returns:
            Encoded features (batch, output_dim)
        """
        if strategy_name not in self.encoders:
            raise ValueError(f"Unknown strategy: {strategy_name}")

        return self.encoders[strategy_name](features)

    def encode_all(
        self,
        strategy_outputs: Dict[str, torch.Tensor]
    ) -> Dict[str, torch.Tensor]:
        """
        Encode all strategy outputs.

        Args:
            strategy_outputs: Dictionary of strategy outputs

        Returns:
            Dictionary of encoded outputs
        """
        return {
            name: self.forward(name, features)
            for name, features in strategy_outputs.items()
            if name in self.encoders
        }


if __name__ == "__main__":
    # Test the gating network
    torch.manual_seed(42)

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

    config = GatingConfig(
        n_strategies=5,
        strategy_output_dim=8,
        context_dim=32
    )

    gating = GatingNetwork(config)

    # Create dummy inputs
    batch_size = 16

    strategy_outputs = {
        'pva': torch.randn(batch_size, 8),
        'mrd': torch.randn(batch_size, 8),
        'vbp': torch.randn(batch_size, 8),
        'msa': torch.randn(batch_size, 8),
        'mts': torch.randn(batch_size, 8),
    }

    context = torch.randn(batch_size, 32)

    # Forward pass
    weights = gating(strategy_outputs, context)

    print(f"Input batch size: {batch_size}")
    print(f"Output weights shape: {weights.shape}")
    print(f"Weights sum: {weights.sum(dim=1)}")
    print(f"Sample weights: {weights[0].detach().numpy()}")

    # Check entropy loss
    entropy_loss = gating.compute_entropy_loss(weights)
    print(f"Entropy loss: {entropy_loss.item():.4f}")

    # Test with missing strategy
    print("\n--- Testing with missing strategy ---")
    partial_outputs = {k: v for k, v in strategy_outputs.items() if k != 'mts'}
    weights_partial = gating(partial_outputs, context)
    print(f"Partial weights: {weights_partial[0].detach().numpy()}")

    # Test minimum weight constraint
    print("\n--- Testing minimum weight constraint ---")
    print(f"Min weight per strategy: {config.min_weight}")
    print(f"All weights >= min: {(weights >= config.min_weight - 0.01).all()}")

    print("\nGatingNetwork test complete!")
