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


class direction_regressor(pl.LightningModule):
    """
    A PyTorch Lightning Module for direction 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 EdgeConv layers to be used in the model.

    Attributes
    ----------
    mlp : torch.nn.Sequential
        The initial MLP used in the first EdgeConv layer.
    convs : torch.nn.ModuleList
        A list of EdgeConv layers.
    bns : torch.nn.ModuleList
        A list of batch normalization layers corresponding to each EdgeConv layer.
    fc : torch.nn.Linear
        The fully connected layer for output.
    dropout_rate : float
        The dropout rate to be applied after global pooling.
    learning_rate : float
        The learning rate for the optimizer.
    r2_score : R2Score
        The R-squared score metric.


    """

    def __init__(self, input_dim, hidden_dim, learning_rate, dropout_rate, num_layers):
        super(direction_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))

        # Additional convolutional layers
        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, 2)  # Output dimension for alt and az

        # 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

    def training_step(self, batch, batch_idx):
        x = self.forward(batch.x, batch.edge_index, batch.batch)
        loss = F.mse_loss(x, batch.direction)
        self.log("train_loss", loss)
        return loss

    def validation_step(self, batch, batch_idx):
        x = self.forward(batch.x, batch.edge_index, batch.batch)
        print(f"x shape: {x.shape}")
        print(f"batch.direction shape: {batch.direction.shape}")

        loss = F.mse_loss(x, batch.direction)

        self.log("val_loss", loss)
        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
