# %%
from typing import Optional, Sequence

import matplotlib.pyplot as plt
import torch
from matplotlib.figure import Figure
from scaffolding.config import Config
from scaffolding.model.classification import ClassificationModule
from torch.utils.data import Dataset

from mlbnb.paths import ExperimentPath


class Plotter:
    def __init__(
        self,
        cfg: Config,
        test_data: Dataset,
        save_to: Optional[ExperimentPath] = None,
        sample_indices: Sequence[int] = [0],
    ):
        self._cfg = cfg
        self._sample_tasks: list = [test_data[i] for i in sample_indices]
        self._num_samples = len(sample_indices)
        self._test_data = test_data
        self._dir = save_to

    def plot_prediction(
        self, model: ClassificationModule, samples_seen: int = 0
    ) -> Optional[Figure]:
        fig = self._plot(model)

        if fig:
            self._save_or_show(fig, f"{samples_seen}_prediction.png")

    @torch.no_grad()
    def _plot(self, model: ClassificationModule) -> Optional[Figure]:
        fig, axs = plt.subplots(
            1, self._num_samples, figsize=(10, 5 * self._num_samples), squeeze=False
        )
        for i, task in enumerate(self._sample_tasks):
            img, target = task
            img = img.to(self._cfg.runtime.device)
            # Squeeze and/or unsqueeze to ensure image is C_W_H:
            if img.dim() == 2:
                img = img.unsqueeze(0).unsqueeze(0)
            elif img.dim() == 3:
                img = img
            pred = model.predict(img.unsqueeze(0)).argmax(dim=1).item()
            # Normalize image to [0, 1]:
            img = (img - img.min()) / (img.max() - img.min())
            img = img.permute(1, 2, 0)
            axs[0, i].imshow(img.squeeze().cpu().numpy(), cmap="grey")
            axs[0, i].set_title(f"True: {target}, Pred: {pred}")
            if target != pred:
                axs[0, i].title.set_color("red")
        return fig

    def _save_or_show(self, fig: Figure, fname: str) -> None:
        if self._dir:
            fig.savefig(self._dir / fname, bbox_inches="tight", dpi=300)
        else:
            plt.show()


if __name__ == "__main__":
    from scaffolding.util.instantiate import Experiment, load_config

    cfg = Config()
    data = "mnist"

    cfg = load_config(data=data)
    exp = Experiment.from_config(cfg)

    plotter = Plotter(cfg, exp.val_loader.dataset, sample_indices=[0, 1, 2])

    assert isinstance(exp.model, ClassificationModule)

    plotter.plot_prediction(exp.model, 0)
