import pytorch_lightning as pl
import torch
import torch.nn.functional as F
from torch_geometric.nn import EdgeConv, global_mean_pool as gap, global_max_pool as gmp
from torchmetrics import R2Score
from stereograph.utils import compute_r68

class energy_regressor(pl.LightningModule):
    """
    A PyTorch Lightning Module for energy regression using EdgeConv layers.
    Parameters
    ----------
    input_dim : int
        The dimension of the input features.
    hidden_dim : int
        The dimension of the hidden layers.
    learning_rate : float
        The learning rate for the optimizer.
    dropout_rate : float
        The dropout rate to be applied after global pooling.
    num_layers : int
        The number of convolutionnal layers.
    Attributes
    ----------
    mlp : torch.nn.Sequential
        Initial MLP used in the first EdgeConv layer.
    convs : torch.nn.ModuleList
        List of EdgeConv layers.
    bns : torch.nn.ModuleList
        List of batch normalization layers.
    fc : torch.nn.Linear
        Fully connected layer for final regression output.
    dropout_rate : float
        Dropout rate applied after global pooling.
    learning_rate : float
        Learning rate for the optimizer.
    r2_score : R2Score
        R-squared score metric.

    """
    
    def __init__(self, input_dim, hidden_dim, learning_rate, dropout_rate, num_layers):
        super(energy_regressor, self).__init__()

        # Initial MLP used in the first EdgeConv
        self.mlp = torch.nn.Sequential(
            torch.nn.Linear(2 * input_dim, hidden_dim),
            torch.nn.GELU(),
            torch.nn.Linear(hidden_dim, hidden_dim)
        )

        # Convolutional layers and  batch normalization
        self.convs = torch.nn.ModuleList()
        self.bns = torch.nn.ModuleList()

        # first convolutional layer
        self.convs.append(EdgeConv(nn=self.mlp))
        self.bns.append(torch.nn.BatchNorm1d(hidden_dim))

        for _ in range(num_layers - 1):
            self.convs.append(EdgeConv(nn=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.fc = torch.nn.Linear(hidden_dim * 2, 1)
        
        # Dropout Rate
        self.dropout_rate = dropout_rate
        
        # Learning Rate
        self.learning_rate = learning_rate

        # R-squared Score Metric
        self.r2_score = R2Score()

    def forward(self, x, edge_index, batch):
        # Apply convolutional layers with batch normalization and GELU activation
        for conv, bn in zip(self.convs, self.bns):
            x = F.gelu(bn(conv(x, edge_index)))
        
        # Global Pooling
        x = torch.cat([gmp(x, batch), gap(x, batch)], dim=1)
        
        # Dropout
        x = F.dropout(x, p=self.dropout_rate, training=self.training)
        
        # Fully Connected Layer
        x = self.fc(x)
        
        return x.view(-1)

    def training_step(self, batch, batch_idx):
        x = self.forward(batch.x, batch.edge_index, batch.batch)
        loss = F.mse_loss(x, batch.true_energy)
        r2 = self.r2_score(x, batch.true_energy)
        self.log('train_r2', r2)
        self.log('train_loss', loss)
        return loss

    def validation_step(self, batch, batch_idx):
        x = self.forward(batch.x, batch.edge_index, batch.batch)
        loss = F.mse_loss(x, batch.true_energy)
        r2 = self.r2_score(x, batch.true_energy)
        r68 = compute_r68(10**x, 10**batch.true_energy)  # Compute r68
        self.log('val_r2', r2)
        self.log('val_loss', loss)
        self.log('r68', r68)
        return loss

    def predict_step(self, batch, batch_idx):
        x = self.forward(batch.x, batch.edge_index, batch.batch)
        return x

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