from typing import Any, Union

import hydra
import lightning as pl
import torch
from omegaconf import DictConfig

from relik.retriever.common.model_inputs import ModelInputs


class GoldenRetrieverPLModule(pl.LightningModule):
    def __init__(
        self,
        model: Union[torch.nn.Module, DictConfig],
        optimizer: Union[torch.optim.Optimizer, DictConfig],
        lr_scheduler: Union[torch.optim.lr_scheduler.LRScheduler, DictConfig] = None,
        *args,
        **kwargs,
    ) -> None:
        super().__init__()
        self.save_hyperparameters(ignore=["model"])
        if isinstance(model, DictConfig):
            self.model = hydra.utils.instantiate(model)
        else:
            self.model = model

        self.optimizer_config = optimizer
        self.lr_scheduler_config = lr_scheduler

    def forward(self, **kwargs) -> dict:
        """
        Method for the forward pass.
        'training_step', 'validation_step' and 'test_step' should call
        this method in order to compute the output predictions and the loss.

        Returns:
            output_dict: forward output containing the predictions (output logits ecc...) and the loss if any.

        """
        return self.model(**kwargs)

    def training_step(self, batch: ModelInputs, batch_idx: int) -> torch.Tensor:
        forward_output = self.forward(**batch, return_loss=True)
        self.log(
            "loss",
            forward_output["loss"],
            batch_size=batch["questions"]["input_ids"].size(0),
            prog_bar=True,
        )
        return forward_output["loss"]

    def validation_step(self, batch: ModelInputs, batch_idx: int) -> None:
        forward_output = self.forward(**batch, return_loss=True)
        self.log(
            "val_loss",
            forward_output["loss"],
            batch_size=batch["questions"]["input_ids"].size(0),
        )

    def test_step(self, batch: ModelInputs, batch_idx: int) -> Any:
        forward_output = self.forward(**batch, return_loss=True)
        self.log(
            "test_loss",
            forward_output["loss"],
            batch_size=batch["questions"]["input_ids"].size(0),
        )

    def configure_optimizers(self):
        if isinstance(self.optimizer_config, DictConfig):
            param_optimizer = list(self.named_parameters())
            no_decay = ["bias", "LayerNorm.bias", "LayerNorm.weight"]
            optimizer_grouped_parameters = [
                {
                    "params": [
                        p for n, p in param_optimizer if "layer_norm_layer" in n
                    ],
                    "weight_decay": self.hparams.optimizer.weight_decay,
                    "lr": 1e-4,
                },
                {
                    "params": [
                        p
                        for n, p in param_optimizer
                        if all(nd not in n for nd in no_decay)
                        and "layer_norm_layer" not in n
                    ],
                    "weight_decay": self.hparams.optimizer.weight_decay,
                },
                {
                    "params": [
                        p
                        for n, p in param_optimizer
                        if "layer_norm_layer" not in n
                        and any(nd in n for nd in no_decay)
                    ],
                    "weight_decay": 0.0,
                },
            ]
            optimizer = hydra.utils.instantiate(
                self.optimizer_config,
                # params=self.parameters(),
                params=optimizer_grouped_parameters,
                _convert_="partial",
            )
        else:
            optimizer = self.optimizer_config

        if self.lr_scheduler_config is None:
            return optimizer

        if isinstance(self.lr_scheduler_config, DictConfig):
            lr_scheduler = hydra.utils.instantiate(
                self.lr_scheduler_config, optimizer=optimizer
            )
        else:
            lr_scheduler = self.lr_scheduler_config

        lr_scheduler_config = {
            "scheduler": lr_scheduler,
            "interval": "step",
            "frequency": 1,
        }
        return [optimizer], [lr_scheduler_config]