"""
models.py

This module defines the GraphTransformer, GraphSAGE, GAT, GNN, HAN, and GCF classes, neural network models using various architectures for processing graph data.

Classes:
    GraphTransformer: A neural network model for graph data using Transformer architecture.
    GraphSAGE: A neural network model for graph data using GraphSAGE architecture.
    GAT: A neural network model for graph data using Graph Attention Network architecture.
    GNN: A general neural network model for graph data.
    HAN: A neural network model for heterogeneous graph data using Heterogeneous Graph Attention Network architecture.
    GCF: A neural network model for collaborative filtering using Graph Convolutional Factorization.

Usage:
    from engine.models import GraphTransformer, GraphSAGE, GAT, GNN, HAN, GCF

    # Example usage for GraphTransformer
    model = GraphTransformer(num_layers=2, d_model=128, num_heads=4, d_feedforward=512, input_dim=10)
    output = model(torch.randn(1, 10))

    # Example usage for GraphSAGE
    model = GraphSAGE(in_channels=10, out_channels=128, num_layers=2, aggr='mean')
    output = model(data)

    # Example usage for GAT
    model = GAT(in_channels=10, out_channels=128, num_layers=2, heads=4)
    output = model(data)

    # Example usage for GNN
    model = GNN(in_channels=10, hidden_channels=64, out_channels=128, num_layers=2)
    output = model(data)

    # Example usage for HAN
    model = HAN(in_channels=10, out_channels=128, num_heads=8, num_metapaths=3)
    output = model(data)

    # Example usage for GCF
    model = GCF(num_users=1000, num_items=500, embedding_dim=64, num_layers=2)
    output = model(user_indices, item_indices)
"""
from torch.nn import Module, Linear, TransformerEncoderLayer, TransformerEncoder, ReLU, ModuleList, Embedding
from torch_geometric.utils import degree
# from cr_pkg import sage_conv, gat_conv, gcn_conv, han_conv
from torch_geometric.nn import SAGEConv, GATConv, GCNConv, HANConv
import torch

class GraphTransformer(Module):
    def __init__(self, num_layers, d_model, num_heads, d_feedforward, input_dim, num_weights=10, use_weights=True):
        super(GraphTransformer, self).__init__()
        self.num_weights = num_weights
        self.use_weights = use_weights
        self.input_linear = Linear(input_dim, d_model)
        self.encoder_layer = TransformerEncoderLayer(d_model=d_model, nhead=num_heads, dim_feedforward=d_feedforward, batch_first=True)
        self.transformer_encoder = TransformerEncoder(self.encoder_layer, num_layers=num_layers)
        self.output_linear = Linear(d_model, input_dim)
        if self.use_weights:
            self.weight_linears = ModuleList([Linear(input_dim, d_model) for _ in range(num_weights)])

    def forward(self, x, weights=None):
        x = x.float()
        if self.use_weights:
            if weights is not None:
                weighted_x = torch.zeros_like(x)
                for i, weight in enumerate(weights):
                    weighted_x += self.weight_linears[i](x) * weight
                x = weighted_x
            else:
                x = self.input_linear(x)
        else:
            x = self.input_linear(x)
        x = self.transformer_encoder(x)
        x = self.output_linear(x)
        return x

class GraphSAGE(Module):
    def __init__(self, in_channels, out_channels, num_layers, aggr='mean'):
        super(GraphSAGE, self).__init__()
        self.convs = ModuleList()
        self.convs.append(SAGEConv(in_channels, out_channels, aggr=aggr))
        for _ in range(num_layers - 1):
            self.convs.append(SAGEConv(out_channels, out_channels, aggr=aggr))
        self.relu = ReLU()

    def forward(self, data):
        x, edge_index = data.x, data.edge_index
        for conv in self.convs:
            x = conv(x, edge_index)
            x = self.relu(x)
        return x

class GAT(Module):
    def __init__(self, in_channels, out_channels, num_layers, heads=1):
        super(GAT, self).__init__()
        self.convs = ModuleList()
        self.convs.append(GATConv(in_channels, out_channels // heads, heads=heads))
        for _ in range(num_layers - 1):
            self.convs.append(GATConv(out_channels, out_channels // heads, heads=heads))
        self.relu = ReLU()

    def forward(self, data):
        x, edge_index = data.x, data.edge_index
        for conv in self.convs:
            x = conv(x, edge_index)
            x = self.relu(x)
        return x
    
class GNN(Module):
    def __init__(self, in_channels, hidden_channels, out_channels, num_layers):
        super(GNN, self).__init__()
        self.convs = ModuleList()
        self.convs.append(GCNConv(in_channels, hidden_channels))
        for _ in range(num_layers - 2):
            self.convs.append(GCNConv(hidden_channels, hidden_channels))
        self.convs.append(GCNConv(hidden_channels, out_channels))
        self.relu = ReLU()

    def forward(self, data):
        x, edge_index = data.x, data.edge_index
        for conv in self.convs:
            x = conv(x, edge_index)
            x = self.relu(x)
        return x

class HAN(Module):
    def __init__(self, in_channels, out_channels, num_heads, num_metapaths):
        super(HAN, self).__init__()
        self.convs = ModuleList()
        for _ in range(num_metapaths):
            self.convs.append(HANConv(in_channels, out_channels, heads=num_heads))
        self.lin = Linear(out_channels * num_heads, out_channels)
        self.relu = ReLU()

    def forward(self, data):
        x, edge_index_dict = data.x, data.edge_index_dict
        out = []
        for conv, edge_index in zip(self.convs, edge_index_dict.values()):
            out.append(conv(x, edge_index))
        x = torch.cat(out, dim=1)
        x = self.lin(x)
        x = self.relu(x)
        return x

class GCF(Module):
    def __init__(self, num_users, num_items, embedding_dim, num_layers):
        super(GCF, self).__init__()
        self.user_embedding = Embedding(num_users, embedding_dim)
        self.item_embedding = Embedding(num_items, embedding_dim)
        self.convs = ModuleList()
        for _ in range(num_layers):
            self.convs.append(GCNConv(embedding_dim, embedding_dim))
        self.relu = ReLU()

    def forward(self, user_indices, item_indices):
        user_emb = self.user_embedding(user_indices)
        item_emb = self.item_embedding(item_indices)
        x = torch.cat([user_emb, item_emb], dim=0)
        edge_index = self.build_edge_index(user_indices, item_indices, x.size(0))
        for conv in self.convs:
            x = conv(x, edge_index)
            x = self.relu(x)
        user_emb, item_emb = torch.split(x, [user_emb.size(0), item_emb.size(0)], dim=0)
        return user_emb, item_emb

    def build_edge_index(self, user_indices, item_indices, num_nodes):
        user_indices = user_indices + num_nodes // 2
        edge_index = torch.stack([torch.cat([user_indices, item_indices]), torch.cat([item_indices, user_indices])], dim=0)
        return edge_index
    










######################-Work-in-Progress-##################### 
class GraphTransformerV2(Module):
    def __init__(self, num_layers, d_model, num_heads, d_feedforward, input_dim):
        super(GraphTransformerV2, self).__init__()
        self.input_dim = input_dim
        self.d_model = d_model
        # self.output_dim = output_dim

        # Define layers
        self.input_linear = Linear(input_dim, d_model)
        self.encoder_layer = TransformerEncoderLayer(d_model=d_model, nhead=num_heads, dim_feedforward=d_feedforward, batch_first=True)
        self.transformer_encoder = TransformerEncoder(self.encoder_layer, num_layers=num_layers)
        # self.output_linear = Linear(d_model + 1, output_dim)  # +1 for centrality score
        self.relu = ReLU()

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

        # Direct Connections: Aggregate features from direct neighbors
        row, col = edge_index
        agg_features = torch.zeros_like(x)
        agg_features.index_add_(0, row, x[col])

        # Apply input linear layer and activation
        x = self.input_linear(agg_features)
        x = self.relu(x)

        # Neighborhood Similarity: Aggregate features from neighbors of neighbors
        agg_features = torch.zeros_like(x)
        agg_features.index_add_(0, row, x[col])

        # Apply transformer encoder
        x = self.transformer_encoder(agg_features)

        # Graph Structure: Centrality score
        centrality = degree(edge_index[0], x.size(0), dtype=x.dtype).view(-1, 1)

        # Concatenate centrality score with node features
        x = torch.cat([x, centrality], dim=1)

        # Final linear layer
        x = self.output_linear(x)

        return x

