from typing import Literal

import matplotlib.cm as cm
import torch
from einops import rearrange
from lightning.pytorch import LightningModule
from lightning.pytorch.loggers import TensorBoardLogger
from lightning.pytorch.utilities.types import STEP_OUTPUT
from torch import Tensor, nn
from torch.distributions import (
    Categorical,
    Distribution,
    Independent,
    MixtureSameFamily,
)
from torch.optim import Optimizer
from torchmetrics import MeanSquaredError, MetricCollection
from torchvision.transforms.v2 import functional as F
from torchvision.utils import make_grid

from torch_uncertainty.metrics import (
    DistributionNLL,
    Log10,
    MeanAbsoluteErrorInverse,
    MeanGTRelativeAbsoluteError,
    MeanGTRelativeSquaredError,
    MeanSquaredErrorInverse,
    MeanSquaredLogError,
    SILog,
    ThresholdAccuracy,
)
from torch_uncertainty.models import (
    EPOCH_UPDATE_MODEL,
    STEP_UPDATE_MODEL,
)
from torch_uncertainty.utils.distributions import (
    dist_rearrange,
    dist_size,
    dist_squeeze,
)


class PixelRegressionRoutine(LightningModule):
    inv_norm_params = {
        "mean": [-0.485 / 0.229, -0.456 / 0.224, -0.406 / 0.255],
        "std": [1 / 0.229, 1 / 0.224, 1 / 0.255],
    }

    def __init__(
        self,
        model: nn.Module,
        output_dim: int,
        probabilistic: bool,
        loss: nn.Module,
        is_ensemble: bool = False,
        format_batch_fn: nn.Module | None = None,
        optim_recipe: dict | Optimizer | None = None,
        num_image_plot: int = 4,
        log_plots: bool = False,
    ) -> None:
        """Routine for training & testing on **pixel regression** tasks.

        Args:
            model (nn.Module): Model to train.
            output_dim (int): Number of outputs of the model.
            probabilistic (bool): Whether the model is probabilistic, i.e.,
                outputs a PyTorch distribution.
            loss (nn.Module): Loss function to optimize the :attr:`model`.
            is_ensemble (bool, optional): Whether the model is an ensemble.
                Defaults to ``False``.
            optim_recipe (dict or Optimizer, optional): The optimizer and
                optionally the scheduler to use. Defaults to ``None``.
            format_batch_fn (nn.Module, optional): The function to format the
                batch. Defaults to ``None``.
            num_image_plot (int, optional): Number of images to plot. Defaults to ``4``.
            log_plots (bool, optional): Indicates whether to log plots from
                metrics. Defaults to ``False``.
        """
        super().__init__()
        _depth_routine_checks(output_dim, num_image_plot, log_plots)

        self.model = model
        self.output_dim = output_dim
        self.one_dim_depth = output_dim == 1
        self.probabilistic = probabilistic
        self.loss = loss
        self.num_image_plot = num_image_plot
        self.is_ensemble = is_ensemble
        self.log_plots = log_plots

        self.needs_epoch_update = isinstance(model, EPOCH_UPDATE_MODEL)
        self.needs_step_update = isinstance(model, STEP_UPDATE_MODEL)

        if format_batch_fn is None:
            format_batch_fn = nn.Identity()

        self.optim_recipe = optim_recipe
        self.format_batch_fn = format_batch_fn

        depth_metrics = MetricCollection(
            {
                "reg/SILog": SILog(),
                "reg/log10": Log10(),
                "reg/ARE": MeanGTRelativeAbsoluteError(),
                "reg/RSRE": MeanGTRelativeSquaredError(squared=False),
                "reg/RMSE": MeanSquaredError(squared=False),
                "reg/RMSELog": MeanSquaredLogError(squared=False),
                "reg/iMAE": MeanAbsoluteErrorInverse(),
                "reg/iRMSE": MeanSquaredErrorInverse(squared=False),
                "reg/d1": ThresholdAccuracy(power=1),
                "reg/d2": ThresholdAccuracy(power=2),
                "reg/d3": ThresholdAccuracy(power=3),
            },
            compute_groups=False,
        )

        self.val_metrics = depth_metrics.clone(prefix="val/")
        self.test_metrics = depth_metrics.clone(prefix="test/")

        if self.probabilistic:
            depth_prob_metrics = MetricCollection(
                {"reg/NLL": DistributionNLL(reduction="mean")}
            )
            self.val_prob_metrics = depth_prob_metrics.clone(prefix="val/")
            self.test_prob_metrics = depth_prob_metrics.clone(prefix="test/")

    def configure_optimizers(self) -> Optimizer | dict:
        return self.optim_recipe

    def on_train_start(self) -> None:
        if self.logger is not None:  # coverage: ignore
            self.logger.log_hyperparams(
                self.hparams,
            )

    def on_validation_start(self) -> None:
        if self.needs_epoch_update and not self.trainer.sanity_checking:
            self.model.update_wrapper(self.current_epoch)
            if hasattr(self.model, "need_bn_update"):
                self.model.bn_update(
                    self.trainer.train_dataloader, device=self.device
                )

    def on_test_start(self) -> None:
        if hasattr(self.model, "need_bn_update"):
            self.model.bn_update(
                self.trainer.train_dataloader, device=self.device
            )

    def forward(self, inputs: Tensor) -> Tensor | Distribution:
        """Forward pass of the routine.

        The forward pass automatically squeezes the output if the regression
        is one-dimensional and if the routine contains a single model.

        Args:
            inputs (Tensor): The input tensor.

        Returns:
            Tensor: The output tensor.
        """
        pred = self.model(inputs)
        if self.probabilistic:
            if not self.is_ensemble:
                pred = dist_squeeze(pred, -1)
        else:
            if not self.is_ensemble:
                pred = pred.squeeze(-1)
        return pred

    def training_step(
        self, batch: tuple[Tensor, Tensor], batch_idx: int
    ) -> STEP_OUTPUT:
        inputs, target = self.format_batch_fn(batch)
        if self.one_dim_depth:
            target = target.unsqueeze(1)

        dists = self.model(inputs)
        if self.probabilistic:
            out_shape = dist_size(dists)[-2:]
        else:
            out_shape = dists.shape[-2:]
        target = F.resize(
            target, out_shape, interpolation=F.InterpolationMode.NEAREST
        )
        padding_mask = torch.isnan(target)
        if self.probabilistic:
            loss = self.loss(dists, target, padding_mask)
        else:
            loss = self.loss(dists[padding_mask], target[padding_mask])

        if self.needs_step_update:
            self.model.update_wrapper(self.current_epoch)
        self.log("train_loss", loss, prog_bar=True, logger=True)
        return loss

    def validation_step(
        self, batch: tuple[Tensor, Tensor], batch_idx: int
    ) -> None:
        inputs, targets = batch
        if self.one_dim_depth:
            targets = targets.unsqueeze(1)
        batch_size = targets.size(0)
        targets = rearrange(targets, "b c h w -> (b c h w)")
        preds = self.model(inputs)

        if self.probabilistic:
            ens_dist = Independent(
                dist_rearrange(
                    preds, "(m b) c h w -> (b c h w) m", b=batch_size
                ),
                0,
            )
            mix = Categorical(
                torch.ones(
                    (dist_size(preds)[0] // batch_size), device=self.device
                )
            )
            mixture = MixtureSameFamily(mix, ens_dist)
            preds = mixture.mean
        else:
            preds = rearrange(preds, "(m b) c h w -> (b c h w) m", b=batch_size)
            preds = preds.mean(dim=1)

        if batch_idx == 0 and self.log_plots:
            self._plot_depth(
                inputs[: self.num_image_plot, ...],
                preds[: self.num_image_plot, ...],
                targets[: self.num_image_plot, ...],
                stage="val",
            )

        padding_mask = torch.isnan(targets)
        self.val_metrics.update(preds[padding_mask], targets[padding_mask])
        if self.probabilistic:
            self.val_prob_metrics.update(mixture, targets, padding_mask)

    def test_step(
        self,
        batch: tuple[Tensor, Tensor],
        batch_idx: int,
        dataloader_idx: int = 0,
    ) -> None:
        if dataloader_idx != 0:
            raise NotImplementedError(
                "Depth OOD detection not implemented yet. Raise an issue "
                "if needed."
            )
        inputs, targets = batch
        if self.one_dim_depth:
            targets = targets.unsqueeze(1)
        batch_size = targets.size(0)
        targets = rearrange(targets, "b c h w -> (b c h w)")
        preds = self.model(inputs)

        if self.probabilistic:
            ens_dist = dist_rearrange(
                preds, "(m b) c h w -> (b c h w) m", b=batch_size
            )
            mix = Categorical(
                torch.ones(
                    (dist_size(preds)[0] // batch_size), device=self.device
                )
            )
            mixture = MixtureSameFamily(mix, ens_dist)
            preds = mixture.mean
        else:
            preds = rearrange(preds, "(m b) c h w -> (b c h w) m", b=batch_size)
            preds = preds.mean(dim=1)

        if batch_idx == 0 and self.log_plots:
            num_images = (
                self.num_image_plot
                if self.num_image_plot < inputs.size(0)
                else inputs.size(0)
            )
            self._plot_depth(
                inputs[:num_images, ...],
                preds[:num_images, ...],
                targets[:num_images, ...],
                stage="test",
            )

        padding_mask = torch.isnan(targets)
        self.test_metrics.update(preds[padding_mask], targets[padding_mask])
        if self.probabilistic:
            self.test_prob_metrics.update(mixture, targets, padding_mask)

    def on_validation_epoch_end(self) -> None:
        self.log_dict(self.val_metrics.compute(), sync_dist=True)
        self.val_metrics.reset()
        if self.probabilistic:
            self.log_dict(
                self.val_prob_metrics.compute(),
                sync_dist=True,
            )
            self.val_prob_metrics.reset()

    def on_test_epoch_end(self) -> None:
        self.log_dict(
            self.test_metrics.compute(),
            sync_dist=True,
        )
        self.test_metrics.reset()
        if self.probabilistic:
            self.log_dict(
                self.test_prob_metrics.compute(),
                sync_dist=True,
            )
            self.test_prob_metrics.reset()

    def _plot_depth(
        self,
        inputs: Tensor,
        preds: Tensor,
        target: Tensor,
        stage: Literal["val", "test"],
    ) -> None:
        if (
            self.logger is not None
            and isinstance(self.logger, TensorBoardLogger)
            and self.one_dim_depth
        ):
            all_imgs = []
            for i in range(inputs.size(0)):
                img = F.normalize(inputs[i, ...].cpu(), **self.inv_norm_params)
                pred = colorize(
                    preds[i, 0, ...].cpu(), vmin=0, vmax=self.model.max_depth
                )
                tgt = colorize(
                    target[i, 0, ...].cpu(), vmin=0, vmax=self.model.max_depth
                )
                all_imgs.extend([img, pred, tgt])

            self.logger.experiment.add_image(
                f"{stage}/samples",
                make_grid(torch.stack(all_imgs, dim=0), nrow=3),
                self.current_epoch,
            )


def colorize(
    value: Tensor,
    vmin: float | None = None,
    vmax: float | None = None,
    cmap: str = "magma",
):
    """Colorize a tensor of depth values.

    Args:
        value (Tensor): The tensor of depth values.
        vmin (float, optional): The minimum depth value. Defaults to None.
        vmax (float, optional): The maximum depth value. Defaults to None.
        cmap (str, optional): The colormap to use. Defaults to 'magma'.
    """
    vmin = value.min().item() if vmin is None else vmin
    vmax = value.max().item() if vmax is None else vmax
    if vmin == vmax:
        return torch.zeros_like(value)
    value = (value - vmin) / (vmax - vmin)
    cmapper = cm.get_cmap(cmap)
    value = cmapper(value.numpy(), bytes=True)
    img = value[:, :, :3]
    return torch.as_tensor(img).permute(2, 0, 1).float() / 255.0


def _depth_routine_checks(
    output_dim: int, num_image_plot: int, log_plots: bool
) -> None:
    if output_dim < 1:
        raise ValueError(f"output_dim must be positive, got {output_dim}.")
    if num_image_plot < 1 and log_plots:
        raise ValueError(
            f"num_image_plot must be positive, got {num_image_plot}."
        )
