"""
Parts of this code were inspired by JuNkIE (https://bitbucket.org/rfg_lab/junkie/src/master/junkie.py)
Parts of this code were generated by PyCharm's AI assist autofill
"""

import matplotlib.pyplot as plt
import iohub
from iohub import open_ome_zarr
from viscy.light.engine import VSUNet
from ipyfilechooser import FileChooser
import IPython
import os
import numpy as np
from typing import NamedTuple

import torch
from skimage.exposure import rescale_intensity
from sklearn.decomposition import PCA

IPython.get_ipython().run_line_magic('matplotlib', 'widget')
import IPython.display as IPyd
import ipywidgets as ipyw


class Color(NamedTuple):
    r: float
    g: float
    b: float

BOP_ORANGE = Color(0.972549, 0.6784314, 0.1254902)
BOP_BLUE = Color(BOP_ORANGE.b, BOP_ORANGE.g, BOP_ORANGE.r)
GREEN = Color(0.0, 1.0, 0.0)
MAGENTA = Color(1.0, 0.0, 1.0)

class VirtualStainViewer:
    model: VSUNet = None
    config = dict(
        in_channels=1,
        out_channels=2,
        encoder_blocks=[3, 3, 9, 3],
        dims=[96, 192, 384, 768],
        decoder_conv_blocks=2,
        stem_kernel_size=(1, 2, 2),
        in_stack_depth=1,
        pretraining=False,
    )
    zarr: iohub.ngff.Plate = None
    crop: int = 384
    all_positions: list = ['0']
    input: np.ndarray = np.zeros([1, 1, 1, 64, 64])
    output: np.ndarray = np.zeros([64, 64])
    target: np.ndarray = np.zeros([64, 64])
    encoder_feature_maps: list = [np.zeros([64, 64]) for _ in range(4)]
    decoder_feature_maps: list = [np.zeros([64, 64]) for _ in range(3)]

    def __init__(self):
        self.fig, self.ax = plt.subplots(2, 5, figsize=(10, 4))
        self._zarr_upload = FileChooser(os.getcwd())
        self._zarr_upload.show_only_dirs = True
        self._zarr_upload.title = "Upload .zarr"
        self._zarr_upload.register_callback(self._on_zarr_upload)

        self._model_upload = FileChooser(os.getcwd())
        self._model_upload.filter_pattern = "*.ckpt"
        self._model_upload.title = "Upload model"
        self._model_upload.register_callback(self._on_model_upload)

        self._select_position = ipyw.Select(
            options=self.all_positions,
            value=self.all_positions[0],
            description="Position: ",
            disabled=True
        )
        self._update_position = ipyw.Button(description="Update", tooltip="Update position", disabled=True)
        self._update_position.on_click(self._on_position_change)

        self._apply_button = ipyw.Button(description='Apply', disabled=True, tooltip='Apply model to selected store')
        self._apply_button.on_click(self._on_apply)

        self.toolbar = ipyw.VBox([self._zarr_upload,
                                  self._model_upload,
                                  ipyw.HBox([self._select_position,self._update_position]),
                                  self._apply_button])
        IPyd.display(self.toolbar)

        self.update_plots()

    def _on_zarr_upload(self):
        path = self._zarr_upload.selected_path
        if not path.endswith('.zarr'):
            print("Invalid .zarr path")
            return
        self.zarr = open_ome_zarr(path)
        self.all_positions = [p[0].split('/')[-1] for p in self.zarr.positions()]
        self._select_position.set_trait("options", self.all_positions)
        self._select_position.set_trait("index", 0)
        self._select_position.set_trait("disabled", False)
        self._update_position.set_trait("disabled", False)

        if self.zarr is not None and self.model is not None:
            self._apply_button.set_trait('disabled', False)
        self.grab_images()
        self.reset_outputs()
        self.update_plots()

    def _on_model_upload(self):
        path = self._model_upload.selected
        self.model = VSUNet.load_from_checkpoint(
                                                path,
                                                architecture="UNeXt2_2D",
                                                model_config=self.config.copy(),
                                                accelerator="gpu"
                                                )
        if self.zarr is not None and self.model is not None:
            self._apply_button.set_trait('disabled', False)

    def _on_apply(self, button):
        n_components = 3
        self._apply_button.set_trait("description", "Applying ...")
        self._apply_button.set_trait("disabled", True)
        with torch.inference_mode():
            # encoder
            encoder_features = self.model.model.encoder(torch.from_numpy(self.input.astype(np.float32)).to(self.model.device))[0]
            encoder_features_np = [f.detach().cpu().numpy() for f in encoder_features]
            self.encoder_feature_maps = [self.pcs_to_rgb(feat, n_components=n_components) for feat in encoder_features_np]

            # decoder
            features = encoder_features.copy()
            features.reverse()
            feat = features[0]
            features.append(None)
            decoder_features_np = []
            for skip, stage in zip(features[1:], self.model.model.decoder.decoder_stages):
                feat = stage(feat, skip)
                decoder_features_np.append(feat.detach().cpu().numpy())
            self.decoder_feature_maps = [self.pcs_to_rgb(feat, n_components=n_components) for feat in decoder_features_np]
            self.output = self.model.model.head(feat).detach().cpu().numpy()
            self.output = self.clip_p(self.composite_nuc_mem(self.output[0, :, 0], BOP_BLUE, BOP_ORANGE))
        self.update_plots()
        self._apply_button.set_trait("description", "Apply")
        self._apply_button.set_trait("disabled", False)

    def _on_position_change(self, _):
        self._select_position.set_trait("disabled", True)
        self._update_position.set_trait("disabled", True)
        self.grab_images()
        self.reset_outputs()
        self.update_plots()
        self._select_position.set_trait("disabled", False)
        self._update_position.set_trait("disabled", False)

    def grab_images(self):
        norm_meta = self.zarr.zattrs["normalization"]["Phase3D"]["dataset_statistics"]
        Y, X = self.zarr[f"0/0/{self._select_position.value}"].data.shape[-2:]
        phase_idx = self.zarr.channel_names.index('Phase3D')
        self.input = self.zarr[f"0/0/{self._select_position.value}/0"][:, phase_idx:phase_idx + 1, 0:1,
                     Y // 2 - self.crop // 2: Y // 2 + self.crop // 2, X // 2 - self.crop // 2: X // 2 + self.crop // 2]
        self.input = (self.input - norm_meta["median"]) / norm_meta["iqr"]
        self.target = self.zarr[f"0/0/{self._select_position.value}/0"][0, 1:3, 0,
                      Y // 2 - self.crop // 2: Y // 2 + self.crop // 2, X // 2 - self.crop // 2: X // 2 + self.crop // 2]
        self.target = self.clip_p(self.composite_nuc_mem(self.target, GREEN, MAGENTA))

    def reset_outputs(self):
        self.output: np.ndarray = np.zeros([64, 64])
        self.encoder_feature_maps: list = [np.zeros([64, 64]) for _ in range(4)]
        self.decoder_feature_maps: list = [np.zeros([64, 64]) for _ in range(3)]

    def update_plots(self):
        plt.figure(self.fig)
        self.ax[0][0].imshow(self.input[0,0,0], cmap='gray')
        self.ax[0][0].set_title("phase", fontsize=8)

        self.ax[1][0].imshow(self.target)
        self.ax[1][0].set_title("fluorescence", fontsize=8)

        for level, feat in enumerate(self.encoder_feature_maps):
            self.ax[0][level + 1].imshow(feat)
            self.ax[0][level + 1].set_title(f"encoder stage {level + 1}", fontsize=8)

        for level, feat in enumerate(self.decoder_feature_maps):
            self.ax[1][level + 1].imshow(feat)
            self.ax[1][level + 1].set_title(f"decoder stage {level + 1}", fontsize=8)


        self.ax[1][-1].imshow(self.output)
        self.ax[1][-1].set_title(f"prediction", fontsize=8)

        for a in self.ax.ravel():
            a.axis("off")
        plt.tight_layout()

    @staticmethod
    def feature_map_pca(feature_map: np.array, n_components: int = 8) -> PCA:
        """
        Compute PCA on a feature map.
        :param np.array feature_map: (C, H, W) feature map
        :param int n_components: number of components to keep
        :return: PCA: fit sklearn PCA object
        """
        # (C, H, W) -> (C, H*W)
        feat = feature_map.reshape(feature_map.shape[0], -1)
        pca = PCA(n_components=n_components)
        pca.fit(feat)
        return pca

    def pcs_to_rgb(self, feat: np.ndarray, n_components: int = 8) -> np.ndarray:
        pca = self.feature_map_pca(feat[0], n_components=n_components)
        pc_first_3 = pca.components_[:3].reshape(3, *feat.shape[-2:])
        return np.stack(
            [rescale_intensity(pc, out_range=(0, 1)) for pc in pc_first_3], axis=-1
        )

    @staticmethod
    def rescale_clip(image: torch.Tensor) -> np.ndarray:
        return rescale_intensity(image, out_range=(0, 1))[..., None].repeat(3, axis=-1)

    def composite_nuc_mem(self, image: torch.Tensor, nuc_color: Color, mem_color: Color
    ) -> np.ndarray:
        c_nuc = self.rescale_clip(image[0]) * nuc_color
        c_mem = self.rescale_clip(image[1]) * mem_color
        return rescale_intensity(c_nuc + c_mem, out_range=(0, 1))

    @staticmethod
    def clip_p(image: np.ndarray) -> np.ndarray:
        return rescale_intensity(image.clip(*np.percentile(image, [1, 99])))

    @staticmethod
    def clip_highlight(image: np.ndarray) -> np.ndarray:
        return rescale_intensity(image.clip(0, np.percentile(image, 99.5)))
