import time

import torch
import torch.nn as nn
import pytorch_lightning as pl

from torch import optim
from torchmetrics import MetricCollection
from torchmetrics.classification import Accuracy


class ClassificationLightning(pl.LightningModule):
    def __init__(self,
                 model: nn.Module,
                 num_classes: int,
                 optimizer_class,
                 optimizer_kwargs: dict,
                 scheduler_class=None,
                 scheduler_kwargs: dict | None = None,
                 scheduler_monitor: str = "val/loss"):
        super().__init__()
        
        
        scheduler_kwargs = scheduler_kwargs or {}
        
        self.model = model
        self.optimizer_class = optimizer_class
        self.scheduler_class = scheduler_class
        
        num_parameters = sum(
            p.numel()
            for p in model.parameters()
        )
        
        num_trainable_parameters = sum(
            p.numel()
            for p in model.parameters()
            if p.requires_grad
        )
        
        self.save_hyperparameters({
            "model": model.__class__.__name__,
            "num_classes": num_classes,
        
            "num_parameters": num_parameters,
            "num_trainable_parameters": num_trainable_parameters,
        
            "optimizer": optimizer_class.__name__,
            "optimizer_kwargs": optimizer_kwargs,
        
            "scheduler": (
                scheduler_class.__name__
                if scheduler_class is not None
                else None
            ),
            "scheduler_kwargs": scheduler_kwargs,
            "scheduler_monitor": scheduler_monitor
        })
        
        
        self.criterion = nn.CrossEntropyLoss()
        self.train_metrics = MetricCollection({
            "acc_top1": Accuracy(task="multiclass", 
                                 num_classes=self.hparams.num_classes,
                                 top_k=1),
            "acc_top5": Accuracy(task="multiclass", 
                                 num_classes=self.hparams.num_classes,
                                 top_k=5),  
        })
        
        self.val_metrics = self.train_metrics.clone()
        self.test_metrics = self.train_metrics.clone()
        
    def forward(self, image):
        return self.model(image)
    
    def _shared_step(self,
                    sample,
                    stage,
                    metrics):
        x, y = sample
        
        logits = self(x)
        loss = self.criterion(logits, y)
        
        metrics.update(logits, y)
        
        self.log(f"{stage}/loss", 
                 loss, 
                 prog_bar=True, 
                 on_epoch=True, 
                 on_step=False
         )
        
        return loss
    
    def _log_metrics(self, stage, metrics):
        computed = metrics.compute()
        
        for name, value in computed.items():
            self.log(f"{stage}/{name}", value, prog_bar=True)
        
        metrics.reset()
    
    def on_train_epoch_start(self):
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        
        self._epoch_start_time = time.perf_counter()
    
    def training_step(self, batch, batch_idx):
        return self._shared_step(batch, "train", self.train_metrics)
    
    def on_train_epoch_end(self):
        self._log_metrics("train", self.train_metrics)
        
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        
        epoch_time = time.perf_counter() - self._epoch_start_time
        lr = self.optimizers().param_groups[0]["lr"]
        
        self.log("train/lr", lr, prog_bar=True)
        self.log("train/epoch_time", epoch_time)
    
    def validation_step(self, batch, batch_idx):
        return self._shared_step(batch, "val", self.val_metrics)
    
    def on_validation_epoch_end(self):
        self._log_metrics("val", self.val_metrics)
    
    def test_step(self, batch, batch_idx):
        return self._shared_step(batch, "test", self.test_metrics)
    
    def on_test_epoch_end(self):
       self._log_metrics("test", self.test_metrics)
    
    def configure_optimizers(self):
        optimizer = self.optimizer_class(
            self.parameters(),
            **self.hparams.optimizer_kwargs
        )
        
        if self.scheduler_class is None:
            return optimizer
        
        scheduler = self.scheduler_class(
            optimizer,
            **self.hparams.scheduler_kwargs
        )
        
        scheduler_config = {
            "scheduler": scheduler,
            "interval": "epoch",
            "frequency": 1
        }
        
        if isinstance(scheduler, optim.lr_scheduler.ReduceLROnPlateau):
            scheduler_config["monitor"] = self.hparams.scheduler_monitor
        
        return {
            "optimizer": optimizer,
            "lr_scheduler": scheduler_config
        }