import pytorch_lightning as pl
import torch
import torch.nn.functional as F
from torch_geometric.nn import EdgeConv, global_mean_pool
from torchmetrics import Accuracy

class edge_conv_classifier(pl.LightningModule):
    """
    A PyTorch Lightning Module for edge convolution classification.
    Parameters
    ----------
    hidden_dim : int
        The dimension of the hidden layers.
    input_dim : int
        The dimension of the input features.
    dropout_rate : float
        The dropout rate to apply during training.
    learning_rate : float
        The learning rate for the optimizer.
    num_layers : int, optional
        The number of convolution layers
    Attributes
    ----------
    mlp : torch.nn.Sequential
        The multi-layer perceptron for EdgeConv.
    num_layers : int
        The number of EdgeConv layers.
    convs : torch.nn.ModuleList
        The list of EdgeConv layers.
    bns : torch.nn.ModuleList
        The list of batch normalization layers.
    lin : torch.nn.Linear
        The fully connected layer.
    dropout_rate : float
        The dropout rate to apply during training.
    learning_rate : float
        The learning rate for the optimizer.
    train_acc : Accuracy
        The accuracy metric for training.
    val_acc : Accuracy
        The accuracy metric for validation.
    test_acc : Accuracy
        The accuracy metric for testing.
    criterion : torch.nn.BCEWithLogitsLoss
        The binary cross-entropy loss with logits.

    """
    
    def __init__(self, hidden_dim, input_dim, dropout_rate, learning_rate, num_layers=3):
        super(edge_conv_classifier, self).__init__()
        self.save_hyperparameters()

        # Define MLP for EdgeConv
        self.mlp = torch.nn.Sequential(
            torch.nn.Linear(2 * input_dim, hidden_dim),
            torch.nn.GELU(),
            torch.nn.Linear(hidden_dim, hidden_dim)
        )

        # Initialize layers
        self.num_layers = num_layers
        self.convs = torch.nn.ModuleList()
        self.bns = torch.nn.ModuleList()

        # First EdgeConv layer
        self.convs.append(EdgeConv(self.mlp))
        self.bns.append(torch.nn.BatchNorm1d(hidden_dim))

        for _ in range(num_layers - 1):
            self.convs.append(EdgeConv(torch.nn.Sequential(
                torch.nn.Linear(2 * hidden_dim, hidden_dim),
                torch.nn.GELU(),
                torch.nn.Linear(hidden_dim, hidden_dim)
            )))
            self.bns.append(torch.nn.BatchNorm1d(hidden_dim))

        # Fully Connected Layer
        self.lin = torch.nn.Linear(hidden_dim, 1)
        
        # Dropout Rate
        self.dropout_rate = dropout_rate
        
        # Learning Rate
        self.learning_rate = learning_rate

        # Accuracy metric
        self.train_acc = Accuracy(task='binary')
        self.val_acc = Accuracy(task='binary')
        self.test_acc = Accuracy(task='binary')  
        
        # Binary Cross Entropy Loss with logits
        self.criterion = torch.nn.BCEWithLogitsLoss()

    def forward(self, x, edge_index, batch):
        print(f"edge_index shape: {edge_index.shape}")
        for conv, bn in zip(self.convs, self.bns):
            x = conv(x, edge_index)
            x = F.gelu(bn(x))
        
        x = global_mean_pool(x, batch)  # Global mean pooling
        x = F.dropout(x, p=self.dropout_rate, training=self.training)
        x = self.lin(x).squeeze(1)  # logits
        return x

    def training_step(self, batch, batch_idx):
        out = self(batch.x, batch.edge_index, batch.batch)
        loss = self.criterion(out, batch.particle_type.float())
        preds = torch.sigmoid(out)  # Probabilités
        binary_preds = preds.round()  

        self.train_acc(binary_preds, batch.particle_type.int())
        self.log('train_loss', loss)
        self.log('train_acc', self.train_acc, on_step=False, on_epoch=True)
        return loss

    def validation_step(self, batch, batch_idx):
        out = self(batch.x, batch.edge_index, batch.batch)
        loss = self.criterion(out, batch.particle_type.float())
        preds = torch.sigmoid(out)
        binary_preds = preds.round()

        self.val_acc(binary_preds, batch.particle_type.int())
        self.log('val_loss', loss, prog_bar=True)
        self.log('val_acc', self.val_acc, on_step=False, on_epoch=True, prog_bar=True)
        return loss

    def test_step(self, batch, batch_idx):
        out = self(batch.x, batch.edge_index, batch.batch)
        loss = self.criterion(out, batch.particle_type.float())
        preds = torch.sigmoid(out)
        binary_preds = preds.round()  

        self.test_acc(binary_preds, batch.particle_type.int())
        self.log('test_loss', loss)
        self.log('test_acc', self.test_acc, on_step=False, on_epoch=True)
        return loss

    def predict_step(self, batch, batch_idx, dataloader_idx=None):
        x, edge_index, batch = batch.x, batch.edge_index, batch.batch
        out = self(x, edge_index, batch)
        preds = torch.sigmoid(out)
        return preds

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=self.learning_rate)
