import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.data import Data
from torch_geometric.nn import GATv2Conv, LayerNorm
from safetensors.torch import load_file

class BipartiteData(Data):
    """bipartite graph data structure for pyg batching compatibility"""
    def __init__(self, x_var=None, x_con=None, edge_index_c2v=None, edge_attr=None, **kwargs):
        super().__init__(**kwargs)
        self.x_var = x_var
        self.x_con = x_con
        self.edge_index_c2v = edge_index_c2v
        self.edge_attr = edge_attr
        if x_var is not None and x_con is not None:
            self.num_nodes = x_var.size(0) + x_con.size(0)

    def __inc__(self, key, value, *args, **kwargs):
        if key == 'edge_index_c2v':
            return torch.tensor([[self.x_con.size(0)], [self.x_var.size(0)]])
        return super().__inc__(key, value, *args, **kwargs)

class BipartiteTransformerLayer(nn.Module):
    """alternating bipartite graph attention layer"""
    def __init__(self, hidden_dim, heads=8, dropout=0.1):
        super().__init__()
        self.attn_c2v = GATv2Conv((hidden_dim, hidden_dim), hidden_dim // heads, heads=heads, edge_dim=1, dropout=dropout, add_self_loops=False)
        self.norm_v = LayerNorm(hidden_dim)
        self.ffn_v = nn.Sequential(nn.Linear(hidden_dim, hidden_dim * 2), nn.GELU(), nn.Linear(hidden_dim * 2, hidden_dim))
        self.norm_ffn_v = LayerNorm(hidden_dim)
        
        self.attn_v2c = GATv2Conv((hidden_dim, hidden_dim), hidden_dim // heads, heads=heads, edge_dim=1, dropout=dropout, add_self_loops=False)
        self.norm_c = LayerNorm(hidden_dim)
        self.ffn_c = nn.Sequential(nn.Linear(hidden_dim, hidden_dim * 2), nn.GELU(), nn.Linear(hidden_dim * 2, hidden_dim))
        self.norm_ffn_c = LayerNorm(hidden_dim)

    def forward(self, x_var, x_con, edge_index_c2v, edge_attr):
        h_var = self.attn_c2v((x_con, x_var), edge_index_c2v, edge_attr=edge_attr)
        x_var = self.norm_v(x_var + F.dropout(h_var, p=0.1, training=self.training))
        x_var = self.norm_ffn_v(x_var + self.ffn_v(x_var))

        edge_index_v2c = edge_index_c2v.flip(0) 
        h_con = self.attn_v2c((x_var, x_con), edge_index_v2c, edge_attr=edge_attr)
        x_con = self.norm_c(x_con + F.dropout(h_con, p=0.1, training=self.training))
        x_con = self.norm_ffn_c(x_con + self.ffn_c(x_con))

        return x_var, x_con

class LippersheyBase(nn.Module):
    """lippershey bipartite graph transformer base model"""
    def __init__(self, node_in_dim=5, hidden_dim=256, num_layers=6, heads=8):
        super().__init__()
        self.var_proj = nn.Linear(node_in_dim, hidden_dim)
        self.con_proj = nn.Linear(node_in_dim, hidden_dim)
        self.layers = nn.ModuleList([BipartiteTransformerLayer(hidden_dim, heads) for _ in range(num_layers)])
        self.mlm_head = nn.Sequential(nn.Linear(hidden_dim, hidden_dim), nn.GELU(), nn.Linear(hidden_dim, 1))

    def forward(self, x_var, x_con, edge_index_c2v, edge_attr):
        v = self.var_proj(x_var)
        c = self.con_proj(x_con)
        for layer in self.layers:
            v, c = layer(v, c, edge_index_c2v, edge_attr)
        return self.mlm_head(v)

    def extract_features(self, x_var, x_con, edge_index_c2v, edge_attr):
        """extract latent representations of variables for downstream ppo policy"""
        self.eval()
        with torch.no_grad():
            v = self.var_proj(x_var)
            c = self.con_proj(x_con)
            for layer in self.layers:
                v, c = layer(v, c, edge_index_c2v, edge_attr)
        return v

    @classmethod
    def from_pretrained_safetensors(cls, safetensors_path, node_in_dim=5, hidden_dim=256, num_layers=6, heads=8):
        """load model directly from safetensors format"""
        model = cls(node_in_dim=node_in_dim, hidden_dim=hidden_dim, num_layers=num_layers, heads=heads)
        state_dict = load_file(safetensors_path)
        model.load_state_dict(state_dict)
        return model