import math

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import GATConv

from config import config
from src.amd.optimizations import AMDGATConv, AMDSparseAttention
from src.models.gnn_model import TemporalAttention


class IntradayGNN(nn.Module):
    """
    Intraday GNN model optimized for AMD Radeon R9700 AI Pro
    """

    def __init__(
        self, num_features: int, sequence_length: int = config.SEQUENCE_LENGTH
    ):
        super(IntradayGNN, self).__init__()

        self.sequence_length = sequence_length
        self.num_features = num_features

        # Temporal attention for sequence processing
        self.temporal_attention = TemporalAttention(num_features, num_heads=8)

        # Feature processing modules with AMD optimizations
        self.feature_processor = nn.Sequential(
            nn.Linear(num_features, config.HIDDEN_CHANNELS),
            nn.SiLU(),  # Swish activation often works better on AMD GPUs
            nn.Linear(config.HIDDEN_CHANNELS, config.HIDDEN_CHANNELS),
            nn.LayerNorm(config.HIDDEN_CHANNELS),
        )

        # Graph attention layers with AMD optimizations
        try:
            self.conv1 = AMDGATConv(
                config.HIDDEN_CHANNELS,
                config.HIDDEN_CHANNELS,
                heads=config.NUM_HEADS,
                concat=True,
                dropout=config.DROPOUT,
                add_self_loops=True,
            )
        except Exception:
            self.conv1 = GATConv(
                config.HIDDEN_CHANNELS,
                config.HIDDEN_CHANNELS,
                heads=config.NUM_HEADS,
                concat=True,
                dropout=config.DROPOUT,
                add_self_loops=True,
            )

        try:
            self.conv2 = AMDGATConv(
                config.HIDDEN_CHANNELS * config.NUM_HEADS,
                config.HIDDEN_CHANNELS,
                heads=1,
                concat=False,
                dropout=config.DROPOUT,
                add_self_loops=True,
            )
        except Exception:
            self.conv2 = GATConv(
                config.HIDDEN_CHANNELS * config.NUM_HEADS,
                config.HIDDEN_CHANNELS,
                heads=1,
                concat=False,
                dropout=config.DROPOUT,
                add_self_loops=True,
            )

        # LSTM for temporal dependencies with AMD optimizations
        self.lstm = nn.LSTM(
            input_size=config.HIDDEN_CHANNELS,
            hidden_size=config.HIDDEN_CHANNELS,
            num_layers=2,
            batch_first=True,
            dropout=config.DROPOUT if config.NUM_HEADS > 1 else 0,
            bidirectional=False,
        )

        # Initialize LSTM weights
        self._init_lstm_weights()

        # Attention mechanism for final prediction
        self.attention = nn.Sequential(
            nn.Linear(config.HIDDEN_CHANNELS, config.HIDDEN_CHANNELS),
            nn.SiLU(),
            nn.Linear(config.HIDDEN_CHANNELS, 1),
            nn.Softmax(dim=1),
        )

        # Final prediction layer
        self.linear = nn.Linear(config.HIDDEN_CHANNELS, 1)

        # State for online learning
        self.hidden_state = None

        # Initialize weights
        self._init_weights()

    def _init_weights(self):
        """Initialize weights with AMD-friendly initialization"""
        # Initialize linear layer weights
        nn.init.xavier_uniform_(self.linear.weight)
        nn.init.zeros_(self.linear.bias)

        # Initialize attention layers
        for layer in self.attention:
            if isinstance(layer, nn.Linear):
                nn.init.xavier_uniform_(layer.weight)
                nn.init.zeros_(layer.bias)

    def _init_lstm_weights(self):
        """Initialize LSTM weights with orthogonal initialization for better convergence"""
        for name, param in self.lstm.named_parameters():
            if "weight_ih" in name:
                nn.init.orthogonal_(param)
            elif "weight_hh" in name:
                nn.init.orthogonal_(param)
            elif "bias" in name:
                nn.init.zeros_(param)
                # Set forget gate bias to 1 for better gradient flow
                n = param.size(0)
                param.data[n // 4 : n // 2].fill_(1)

    def forward(self, data):
        x, edge_index, edge_attr = data.x, data.edge_index, data.edge_attr

        # x shape: (num_stocks, sequence_length, num_features)
        batch_size, seq_len, num_features = x.size()

        # Apply temporal attention to each stock's sequence
        temporal_features = []
        for i in range(batch_size):
            stock_sequence = x[i].unsqueeze(0)  # (1, sequence_length, num_features)
            temporal_feature = self.temporal_attention(stock_sequence)
            temporal_features.append(temporal_feature)

        # Stack temporal features
        temporal_features = torch.cat(
            temporal_features, dim=0
        )  # (num_stocks, feature_dim)

        # Process features
        processed_features = self.feature_processor(temporal_features)

        # Process through GNN with gradient checkpointing for memory efficiency
        x = self._gnn_forward(processed_features, edge_index, edge_attr)

        # Process through LSTM for temporal dependencies
        # Reshape for LSTM: (num_stocks, 1, hidden_size)
        lstm_input = x.unsqueeze(1)

        # If we have a hidden state from previous prediction, use it
        if self.hidden_state is not None and config.STATEFUL_PREDICTION:
            lstm_out, self.hidden_state = self.lstm(lstm_input, self.hidden_state)
        else:
            lstm_out, self.hidden_state = self.lstm(lstm_input)

        # Remove sequence dimension
        lstm_out = lstm_out.squeeze(1)

        # Apply attention to LSTM outputs
        attention_weights = self.attention(lstm_out)
        attended = (lstm_out * attention_weights).sum(dim=1, keepdim=True)

        # Final prediction
        return self.linear(attended)

    def _gnn_forward(self, x, edge_index, edge_attr):
        """Forward pass through GNN with gradient checkpointing for memory efficiency"""
        # Gradient checkpointing for memory efficiency on AMD GPUs
        x = torch.utils.checkpoint.checkpoint(
            self.conv1, x, edge_index, edge_attr, preserve_rng_state=False
        )
        x = F.silu(x)  # Swish activation often works better than ReLU on AMD GPUs
        x = F.dropout(x, p=config.DROPOUT, training=self.training)

        x = torch.utils.checkpoint.checkpoint(
            self.conv2, x, edge_index, edge_attr, preserve_rng_state=False
        )

        return x

    def reset_state(self):
        """Reset the hidden state of the LSTM"""
        self.hidden_state = None

    def get_attention_weights(self, data):
        """Get attention weights for interpretability"""
        x, edge_index, edge_attr = data.x, data.edge_index, data.edge_attr

        # Process through the model up to attention
        batch_size, seq_len, num_features = x.size()
        temporal_features = []

        for i in range(batch_size):
            stock_sequence = x[i].unsqueeze(0)
            temporal_feature = self.temporal_attention(stock_sequence)
            temporal_features.append(temporal_feature)

        temporal_features = torch.cat(temporal_features, dim=0)
        processed_features = self.feature_processor(temporal_features)

        x = self.conv1(processed_features, edge_index, edge_attr)
        x = F.silu(x)
        x = self.conv2(x, edge_index, edge_attr)

        lstm_input = x.unsqueeze(1)
        lstm_out, _ = self.lstm(lstm_input)
        lstm_out = lstm_out.squeeze(1)

        attention_weights = self.attention(lstm_out)

        return attention_weights