import numpy as np
import matplotlib.pyplot as plt
import os
from tqdm import tqdm
import cv2 as cv
import matplotlib.patches as patches
import pandas as pd
from typing import List, Dict, Tuple, Optional, Union


class Plot_Visium:
    def __init__(
        self,
        segmentation,
        boundary_dict: Dict,
        type_list: List[str],
        colors: Optional[np.ndarray] = None,
    ):
        """
        Args:
            segmentation: Object of spotiphy.segmentation.Segmentation.
            boundary_dict: Output of function spotiphy.segmentation.cell_boundary.
            type_list: List of the cell types.
            colors: Colors of each cell type.
        """
        self.img = segmentation.img  # Original image.
        if np.max(self.img) <= 1:
            self.img = (self.img * 255).astype(np.int32)
        self.img_seg = segmentation.label  # Segmented image.
        self.n_cell_df = (
            segmentation.n_cell_df
        )  # Dataframe shown the cells in each spot.
        self.nucleus_boundary = segmentation.nucleus_boundary.astype(
            np.int32
        )  # Segmented nucleus boundaries.
        self.nucleus_df = segmentation.nucleus_df
        self.spot_radius = segmentation.spot_radius
        self.spot_centers = segmentation.spot_center

        self.nucleus_centers = self.nucleus_df[["x", "y"]].values.astype(
            np.int32
        )  # Segmented nucleus centers.
        self.img_cell = boundary_dict["img_cell"].astype(
            np.int32
        )  # Inferred cell location.
        self.type_mask = None
        self.img_boundary = boundary_dict["cell_boundary"]
        self.individual_boundary = boundary_dict["individual_boundary"]
        self.img_size = self.img_seg.shape[:2]
        self.type_list = type_list

        if colors is None:
            self.colors = np.array(
                list(plt.get_cmap("tab20b").colors)
                + list(plt.get_cmap("tab20c").colors)
            )
            self.colors = (self.colors * 255).astype(np.int32)
        else:
            self.colors = colors
        self.color_dict = {
            type_: list(self.colors[i]) for i, type_ in enumerate(type_list)
        }

    def plot(
        self,
        background: bool = False,
        cell: str = "both",
        shape: str = "cell",
        circle_size: int = 10,
        boundary: Optional[str] = None,
        save: Optional[str] = "Visium_plot.png",
        background_alpha: float = 0.5,
        spot: bool = True,
        spot_width: int = 2,
        spot_color: Tuple[int, int, int] = (0, 0, 255),
        cell_boundary_color: Tuple[int, int, int] = (100, 100, 100),
        dpi: int = 300,
    ) -> None:
        """
        Args:
            background: If show the background.
            cell: Which group of cell shapes to plot? [both, in, out]
            shape: Which shape to plot? [cell, nucleus, circle]
            circle_size: Size of the nuclei.
            boundary: Which group of cell boundary to plot? [both, in, out]
            save: If not none, save the figure to the path.
            background_alpha: Opacity of the background figure.
            spot: If plot the spot.
            spot_width: Width of the spot.
            spot_color: Color of the spot.
            cell_boundary_color: Color of the cell boundaries.
            dpi: DPI of the image plotted.
        """
        img = (
            self.img.copy() * background_alpha
            if background
            else np.zeros(self.img.shape)
        )
        img = img.astype(np.int32)

        group_dict = {
            "both": [True, False],
            "in": [True],
            "out": [False],
            None: [],
            "None": [],
        }
        type_annotation = list(self.nucleus_df["cell_type"])
        in_spot = list(self.nucleus_df["in_spot"])

        print("Adding cells.")
        if shape == "cell":
            d = {t: i for i, t in enumerate(self.type_list)}
            type_mask = np.zeros(
                (img.shape[0], img.shape[1], len(self.type_list)), dtype=bool
            )
            for i in tqdm(range(len(self.img_cell))):
                for j in range(len(self.img_cell[0])):
                    k = self.img_cell[i, j]
                    if (
                        0 < k <= len(type_annotation)
                        and type_annotation[k - 1] in d.keys()
                        and in_spot[k - 1] in group_dict[cell]
                    ):
                        type_mask[i, j, d[type_annotation[k - 1]]] = True
            for i, type_ in tqdm(enumerate(self.type_list)):
                color = np.array(self.color_dict[type_]).astype(np.int32)
                img[type_mask[:, :, i]] = color
            # img.astype(np.int32)

        for i, type_ in tqdm(enumerate(type_annotation)):
            if in_spot[i] in group_dict[cell] and type_ in self.color_dict.keys():
                color = [int(i) for i in self.color_dict[type_]]
                if shape == "circle":
                    img = cv.circle(
                        img, self.nucleus_centers[i], circle_size, color, -1
                    )
                elif shape == "nucleus":
                    boundary_temp = self.nucleus_boundary[i].reshape((-1, 1, 2))
                    img = cv.fillPoly(img, [boundary_temp], color=color)

        boundary_mask = np.zeros((img.shape[0], img.shape[1]), dtype=bool)
        for i in tqdm(range(len(type_annotation))):
            if (
                in_spot[i] in group_dict[boundary]
                and len(self.individual_boundary[i + 1]) > 0
                and type_annotation[i]
            ) in self.color_dict.keys():
                index_temp = np.array(self.individual_boundary[i + 1])
                boundary_mask[index_temp[:, 0], index_temp[:, 1]] = True
        img[boundary_mask] = cell_boundary_color

        if spot:
            print("Adding spots.")
            for i in range(len(self.spot_centers)):
                img = cv.circle(
                    img,
                    self.spot_centers[i],
                    int(self.spot_radius),
                    spot_color,
                    spot_width,
                )
        if save is not None:
            print("Saving the image.")
            img1 = img[:, :, [2, 1, 0]]
            cv.imwrite(save, img1)
        plt.figure(dpi=dpi)
        plt.imshow(img)

    def plot_legend(
        self,
        save: Optional[str] = None,
        dpi: int = 300,
        background: Tuple[int, int, int] = (0, 0, 0),
        ncol: int = 1,
        word_color: Tuple[int, int, int] = (255, 255, 255),
        fontsize: int = 6,
        horizontal_distance: int = 380,
        show: bool = False,
    ) -> None:
        nrow = len(self.type_list) // ncol
        if nrow * ncol < len(self.type_list):
            nrow += 1
        img = np.ones((nrow * 60 + 40, ncol * horizontal_distance - 30, 3)) * np.array(
            background
        )
        img = img.astype(np.int32)

        fig, ax = plt.subplots(dpi=dpi)
        plt.imshow(img)
        i = 0  # col index
        j = 0  # row index
        word_color = (word_color[0] / 255, word_color[1] / 255, word_color[2] / 255)
        for k, v in self.color_dict.items():
            ax.add_patch(
                patches.Rectangle(
                    (i * horizontal_distance + 20, j * 60 + 20),
                    100,
                    45,
                    facecolor=np.array(v) / 255,
                    edgecolor="none",
                )
            )
            ax.text(
                i * horizontal_distance + 145,
                j * 60 + 43,
                k,
                va="center",
                ha="left",
                fontsize=fontsize,
                color=word_color,
            )
            j += 1
            if j == nrow:
                j = 0
                i += 1
        ax.axis("off")
        if save is not None:
            plt.savefig(save, bbox_inches="tight", pad_inches=0)
        if show:
            plt.show()


class Plot_Xenium:
    def __init__(
        self,
        Xenium_img: np.ndarray,
        cell_boundaries: pd.DataFrame,
        nucleus_boundaries: pd.DataFrame,
        type_list: List[str],
        cell_type: List[str],
        nucleus_centers: np.ndarray,
    ):
        """
        Args:
            Xenium_img: Image of Xenium
            cell_boundaries: Coordinates of the cell boundaries on the image.
            nucleus_boundaries: Coordinates of the nucleus boundaries on the image.
            type_list: List of the cell types.
            cell_type: Annotation.
            nucleus_centers: Centers of the nuclei.
        """
        self.img = Xenium_img
        self.type_list = type_list
        self.cell_type = cell_type
        self.n_cell = len(cell_type)
        self.nucleus_centers = nucleus_centers.astype(np.int32)

        cell_boundaries_temp = cell_boundaries.iloc[:, 1:3].values.astype(np.int32)
        idx = [[] for _ in range(self.n_cell)]
        for i in tqdm(range(len(cell_boundaries))):
            cell_idx = int(cell_boundaries.iloc[i, 0] - 1)
            idx[cell_idx] += [i]
        self.cell_boundaries = [
            (
                cell_boundaries_temp[np.array(idx[i])].reshape((-1, 1, 2))
                if len(idx[i]) > 0
                else []
            )
            for i in tqdm(range(self.n_cell))
        ]

        nucleus_boundaries_temp = nucleus_boundaries.iloc[:, 1:3].values.astype(
            np.int32
        )
        idx = [[] for _ in range(self.n_cell)]
        for i in tqdm(range(len(nucleus_boundaries))):
            cell_idx = int(nucleus_boundaries.iloc[i, 0] - 1)
            idx[cell_idx] += [i]
        self.nucleus_boundaries = [
            (
                nucleus_boundaries_temp[np.array(idx[i])].reshape((-1, 1, 2))
                if len(idx[i]) > 0
                else []
            )
            for i in tqdm(range(self.n_cell))
        ]

        self.colors = np.array(
            list(plt.get_cmap("tab20b").colors) + list(plt.get_cmap("tab20c").colors)
        )
        self.colors = (self.colors * 255).astype(np.int32)
        self.color_dict = {
            type_: list(self.colors[i]) for i, type_ in enumerate(type_list)
        }

    def plot(
        self,
        background: bool = False,
        shape: str = "cell",
        save: Optional[str] = "Xenium_plot.png",
        cell_boundaries: bool = False,
        background_alpha: float = 0.8,
        circle_size: int = 10,
        cell_boundary_color: Tuple[int, int, int] = (100, 100, 100),
        cell_boundary_thickness: int = 2,
    ) -> None:
        """
        Args:
            background: If show the background.
            shape: Which shape to plot? [cell, nucleus, circle]
            save: If not none, save the figure to the path.
            cell_boundaries: If plot the cell boundaries.
            background_alpha: Opacity of the background figure.
            circle_size: Size of the circle.
            cell_boundary_color: Color of the cell boundaries.
            cell_boundary_thickness: Thickness of the cell boundaries.
        """
        img = (
            self.img.copy() * background_alpha
            if background
            else np.zeros(self.img.shape)
        )
        img = img.astype(np.int32)
        print("Adding cellls.")
        for i, type_ in tqdm(enumerate(self.cell_type)):
            if type_ in self.color_dict.keys():
                color = [int(i) for i in self.color_dict[type_]]
                if shape == "cell":
                    if len(self.cell_boundaries[i]) > 0:
                        img = cv.fillPoly(img, [self.cell_boundaries[i]], color=color)
                elif shape == "nucleus":
                    if len(self.nucleus_boundaries[i]) > 0:
                        img = cv.fillPoly(
                            img, [self.nucleus_boundaries[i]], color=color
                        )
                elif shape == "circle":
                    img = cv.circle(
                        img, self.nucleus_centers[i], circle_size, color, -1
                    )
        if cell_boundaries:
            for i, type_ in tqdm(enumerate(self.cell_type)):
                if len(self.cell_boundaries[i]) > 0:
                    img = cv.polylines(
                        img,
                        [self.cell_boundaries[i]],
                        color=cell_boundary_color,
                        thickness=cell_boundary_thickness,
                        isClosed=True,
                    )

        if save is not None:
            print("Saving the image.")
            img1 = img[:, :, [2, 1, 0]]
            cv.imwrite(save, img1)
        plt.imshow(img)
