"""
Graph Neural Network models for retention time prediction.
Includes GCN, GIN, GAT, and ensemble models.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import (
    GCNConv, GINConv, GATConv, GINEConv, GlobalAttention,
    global_mean_pool, global_max_pool, global_add_pool
)
from torch_geometric.nn.norm import BatchNorm, GraphNorm, LayerNorm
from typing import Optional, List, Dict, Sequence, Any
import numpy as np


class GraphConvModel(nn.Module):
    """Base Graph Convolutional Network model."""
    
    def __init__(self, input_dim: int, hidden_dim: int = 128, output_dim: int = 1,
                 num_layers: int = 3, dropout: float = 0.2, 
                 num_labs: int = 23, lab_embed_dim: int = 16):
        super().__init__()
        
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        self.dropout = dropout
        
        # Lab embedding
        self.lab_embedding = nn.Embedding(num_labs, lab_embed_dim)
        
        # Graph convolution layers
        self.convs = nn.ModuleList()
        self.batch_norms = nn.ModuleList()
        
        # First layer
        self.convs.append(GCNConv(input_dim, hidden_dim))
        self.batch_norms.append(GraphNorm(hidden_dim))
        
        # Hidden layers
        for _ in range(num_layers - 1):
            self.convs.append(GCNConv(hidden_dim, hidden_dim))
            self.batch_norms.append(GraphNorm(hidden_dim))
        
        # Global pooling
        self.global_pool = global_mean_pool
        
        # Final prediction layers
        final_input_dim = hidden_dim + lab_embed_dim
        self.predictor = nn.Sequential(
            nn.Linear(final_input_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim // 2, hidden_dim // 4),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim // 4, output_dim)
        )
        
    def forward(self, x, edge_index, batch, lab_feature, edge_attr=None):
        # Graph convolutions
        for conv, norm in zip(self.convs, self.batch_norms):
            # GCN expects edge_weight, not edge_attr for individual edge features
            edge_weight = None
            if edge_attr is not None and edge_attr.size(1) == 1:
                edge_weight = edge_attr.squeeze(-1)
            
            x = conv(x, edge_index, edge_weight=edge_weight)
            x = norm(x, batch)
            x = F.relu(x)
            x = F.dropout(x, p=self.dropout, training=self.training)
        
        # Global pooling
        graph_repr = self.global_pool(x, batch)
        
        # Lab embedding - handle different input shapes
        if lab_feature.dim() == 1:
            lab_embed = self.lab_embedding(lab_feature)
        else:
            lab_embed = self.lab_embedding(lab_feature.squeeze(-1))
        
        # Concatenate graph and lab features
        combined = torch.cat([graph_repr, lab_embed], dim=1)
        
        # Final prediction
        output = self.predictor(combined)
        return output.squeeze()


class GATModel(nn.Module):
    """Graph Attention Network (GAT) model."""
    
    def __init__(self, input_dim: int, hidden_dim: int = 128, output_dim: int = 1,
                 num_layers: int = 3, dropout: float = 0.2, num_heads: int = 4,
                 num_labs: int = 23, lab_embed_dim: int = 16):
        super().__init__()
        
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        self.dropout = dropout
        self.num_heads = num_heads
        
        # Lab embedding
        self.lab_embedding = nn.Embedding(num_labs, lab_embed_dim)
        
        # GAT layers
        self.convs = nn.ModuleList()
        self.batch_norms = nn.ModuleList()
        
        # Calculate dimensions properly for multi-head attention
        head_dim = hidden_dim // num_heads
        
        # First layer: input -> hidden_dim (via multi-head)
        self.convs.append(GATConv(input_dim, head_dim, heads=num_heads, dropout=dropout, concat=True))
        self.batch_norms.append(GraphNorm(hidden_dim))  # head_dim * num_heads = hidden_dim
        
        # Hidden layers: hidden_dim -> hidden_dim (via multi-head)
        for _ in range(num_layers - 2):
            self.convs.append(GATConv(hidden_dim, head_dim, heads=num_heads, dropout=dropout, concat=True))
            self.batch_norms.append(GraphNorm(hidden_dim))
        
        # Last layer: hidden_dim -> hidden_dim (single head)
        if num_layers > 1:
            self.convs.append(GATConv(hidden_dim, hidden_dim, heads=1, dropout=dropout, concat=False))
            self.batch_norms.append(GraphNorm(hidden_dim))
        
        # Global pooling
        self.global_pool = global_mean_pool
        
        # Final prediction layers
        final_input_dim = hidden_dim + lab_embed_dim
        self.predictor = nn.Sequential(
            nn.Linear(final_input_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim // 2, hidden_dim // 4),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim // 4, output_dim)
        )
        
    def forward(self, x, edge_index, batch, lab_feature, edge_attr=None):
        # GAT convolutions
        for i, (conv, norm) in enumerate(zip(self.convs, self.batch_norms)):
            x = conv(x, edge_index)
            x = norm(x, batch)
            if i < len(self.convs) - 1:  # No ReLU on last layer
                x = F.relu(x)
            x = F.dropout(x, p=self.dropout, training=self.training)
        
        # Global pooling
        graph_repr = self.global_pool(x, batch)
        
        # Lab embedding - handle different input shapes
        if lab_feature.dim() == 1:
            lab_embed = self.lab_embedding(lab_feature)
        else:
            lab_embed = self.lab_embedding(lab_feature.squeeze(-1))
        
        # Concatenate graph and lab features
        combined = torch.cat([graph_repr, lab_embed], dim=1)
        
        # Final prediction
        output = self.predictor(combined)
        return output.squeeze()


class MoleculeMPNNModel(nn.Module):
    """Edge-aware message passing network tailored for molecular graphs."""

    def __init__(
        self,
        input_dim: int,
        hidden_dim: int = 256,
        output_dim: int = 1,
        num_layers: int = 4,
        dropout: float = 0.2,
        edge_dim: int = 4,
        num_labs: int = 23,
        lab_embed_dim: int = 16,
        use_batch_norm: bool = True,
    ):
        super().__init__()

        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        self.dropout = dropout
        self.use_batch_norm = use_batch_norm

        self.input_proj = nn.Linear(input_dim, hidden_dim)
        self.edge_encoder = nn.Linear(edge_dim, hidden_dim) if edge_dim > 0 else None
        self.lab_embedding = nn.Embedding(num_labs, lab_embed_dim)

        self.convs = nn.ModuleList()
        self.norms = nn.ModuleList()

        for _ in range(num_layers):
            mlp = nn.Sequential(
                nn.Linear(hidden_dim, hidden_dim),
                nn.ReLU(),
                nn.Linear(hidden_dim, hidden_dim),
            )
            self.convs.append(GINEConv(mlp, train_eps=True))
            if use_batch_norm:
                self.norms.append(BatchNorm(hidden_dim))
            else:
                self.norms.append(GraphNorm(hidden_dim))

        pooled_dim = hidden_dim * 2  # mean + max pooling
        final_input_dim = pooled_dim + lab_embed_dim

        self.predictor = nn.Sequential(
            nn.Linear(final_input_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim // 2, output_dim),
        )

    def forward(self, x, edge_index, batch, lab_feature, edge_attr=None):
        x = self.input_proj(x)
        encoded_edge_attr = edge_attr
        if edge_attr is not None and self.edge_encoder is not None:
            encoded_edge_attr = self.edge_encoder(edge_attr)

        for conv, norm in zip(self.convs, self.norms):
            x = conv(x, edge_index, encoded_edge_attr)
            if isinstance(norm, BatchNorm):
                x = norm(x)
            else:
                x = norm(x, batch)
            x = F.relu(x)
            x = F.dropout(x, p=self.dropout, training=self.training)

        mean_pool = global_mean_pool(x, batch)
        max_pool = global_max_pool(x, batch)
        graph_repr = torch.cat([mean_pool, max_pool], dim=1)

        if lab_feature.dim() == 1:
            lab_embed = self.lab_embedding(lab_feature)
        else:
            lab_embed = self.lab_embedding(lab_feature.squeeze(-1))

        combined = torch.cat([graph_repr, lab_embed], dim=1)
        output = self.predictor(combined)
        return output.squeeze()


# Backwards-compatible alias for previous naming
GINEModel = MoleculeMPNNModel


class HybridModel(nn.Module):
    """Hybrid model combining graph features with molecular descriptors."""

    def __init__(
        self,
        *,
    graph_model_class,
    descriptor_dim: int,
    graph_model_kwargs: Optional[Dict[str, Any]] = None,
    graph_feature_dim: Optional[int] = None,
        descriptor_hidden_dims: Optional[Sequence[int]] = None,
        final_hidden_dims: Optional[Sequence[int]] = None,
        dropout: float = 0.2,
        use_batch_norm: bool = True,
        output_dim: int = 1,
    ) -> None:
        super().__init__()

        graph_model_kwargs = dict(graph_model_kwargs or {})

        if graph_feature_dim is None:
            graph_feature_dim = graph_model_kwargs.get("hidden_dim")
        if graph_feature_dim is None:
            graph_feature_dim = 128

        graph_model_kwargs.setdefault("hidden_dim", graph_feature_dim)
        graph_model_kwargs.setdefault("output_dim", graph_feature_dim)

        self.graph_model = graph_model_class(**graph_model_kwargs)
        self.graph_feature_dim = graph_feature_dim
        self.dropout = dropout
        self.use_batch_norm = use_batch_norm

        if descriptor_hidden_dims is None or len(descriptor_hidden_dims) == 0:
            self.descriptor_net = nn.Identity()
            self.descriptor_output_dim = descriptor_dim
        else:
            descriptor_layers: List[nn.Module] = []
            in_dim = descriptor_dim
            for hidden_dim in descriptor_hidden_dims:
                descriptor_layers.append(nn.Linear(in_dim, hidden_dim))
                if use_batch_norm:
                    descriptor_layers.append(nn.BatchNorm1d(hidden_dim))
                descriptor_layers.append(nn.ReLU())
                if dropout > 0:
                    descriptor_layers.append(nn.Dropout(dropout))
                in_dim = hidden_dim
            self.descriptor_net = nn.Sequential(*descriptor_layers)
            self.descriptor_output_dim = in_dim

        if final_hidden_dims is None or len(final_hidden_dims) == 0:
            final_hidden_dims = [max(graph_feature_dim // 2, 1)]

        final_layers: List[nn.Module] = []
        in_dim = self.graph_feature_dim + self.descriptor_output_dim
        for hidden_dim in final_hidden_dims:
            final_layers.append(nn.Linear(in_dim, hidden_dim))
            if use_batch_norm:
                final_layers.append(nn.BatchNorm1d(hidden_dim))
            final_layers.append(nn.ReLU())
            if dropout > 0:
                final_layers.append(nn.Dropout(dropout))
            in_dim = hidden_dim
        final_layers.append(nn.Linear(in_dim, output_dim))

        self.final_predictor = nn.Sequential(*final_layers)

    def forward(self, x, edge_index, batch, lab_feature, descriptors, edge_attr=None):
        graph_features = self.graph_model(
            x,
            edge_index,
            batch,
            lab_feature,
            edge_attr,
        )
        if graph_features.dim() == 1:
            graph_features = graph_features.unsqueeze(0)

        if descriptors.dim() == 1:
            descriptors = descriptors.unsqueeze(0)

        desc_features = self.descriptor_net(descriptors)
        if isinstance(self.descriptor_net, nn.Identity):
            desc_features = descriptors

        combined = torch.cat([graph_features, desc_features], dim=1)
        output = self.final_predictor(combined)

        return output.squeeze(-1)


def create_model(model_type: str, input_dim: int, num_labs: int = 23, **kwargs):
    """Factory function to create models."""
    
    models = {
    'gcn': GraphConvModel,
    'gin': GINEModel,
    'gat': GATModel,
    'mpnn': MoleculeMPNNModel,
    }
    
    if model_type not in models:
        raise ValueError(f"Unknown model type: {model_type}")
    
    return models[model_type](input_dim=input_dim, num_labs=num_labs, **kwargs)


class EarlyStopping:
    """Early stopping utility."""
    
    def __init__(self, patience: int = 10, min_delta: float = 0.0, restore_best_weights: bool = True):
        self.patience = patience
        self.min_delta = min_delta
        self.restore_best_weights = restore_best_weights
        self.best_loss = None
        self.counter = 0
        self.best_weights = None
        
    def __call__(self, val_loss: float, model: nn.Module) -> bool:
        if self.best_loss is None:
            self.best_loss = val_loss
            self.save_checkpoint(model)
        elif val_loss < self.best_loss - self.min_delta:
            self.best_loss = val_loss
            self.counter = 0
            self.save_checkpoint(model)
        else:
            self.counter += 1
            
        if self.counter >= self.patience:
            if self.restore_best_weights and self.best_weights is not None:
                model.load_state_dict(self.best_weights)
            return True
        return False
    
    def save_checkpoint(self, model: nn.Module):
        """Save model weights."""
        self.best_weights = model.state_dict().copy()


if __name__ == "__main__":
    # Test model creation
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    # Create sample data
    batch_size = 32
    input_dim = 7  # Number of atom features
    num_nodes = 20
    num_edges = 40
    
    x = torch.randn(num_nodes, input_dim)
    edge_index = torch.randint(0, num_nodes, (2, num_edges))
    batch = torch.zeros(num_nodes, dtype=torch.long)
    lab_feature = torch.randint(0, 23, (1,))
    
    # Test different models
    models = ['gcn', 'gin', 'gat']
    
    for model_type in models:
        print(f"\nTesting {model_type.upper()} model:")
        model = create_model(model_type, input_dim=input_dim, hidden_dim=64)
        model.eval()
        
        with torch.no_grad():
            output = model(x, edge_index, batch, lab_feature)
            print(f"Output shape: {output.shape}")
            print(f"Output value: {output.item():.4f}")