from abc import ABC, ABCMeta

from functools import wraps

from pathlib import Path

from typing import Callable, Optional

import pytest

import numpy as np

import pandas as pd

import matplotlib.pyplot as plt

from matplotlib.testing.compare import compare_images

from anndata import AnnData

from moscot import _constants

from moscot.plotting._utils import set_plotting_vars

HERE: Path = Path(__file__).parent

EXPECTED = HERE / "expected_figures"

ACTUAL = HERE / "actual_figures"

TOL = 60

DPI = 40

@pytest.fixture

def adata_pl_cell_transition(gt_temporal_adata: AnnData) -> AnnData:

    plot_vars = {

        "transition_matrix": gt_temporal_adata.uns["cell_transition_10_105_forward"],

        "source_groups": "cell_type",

        "target_groups": "cell_type",

        "source": 0,

        "target": 1,

    }

    set_plotting_vars(gt_temporal_adata, _constants.CELL_TRANSITION, key=_constants.CELL_TRANSITION, value=plot_vars)

    return gt_temporal_adata

@pytest.fixture

def adata_pl_push(adata_time: AnnData) -> AnnData:

    rng = np.random.RandomState(0)

    plot_vars = {"key": "time", "data": "celltype", "subset": "A", "source": 0, "target": 1}

    adata_time.uns["celltype_colors"] = ["#cc1b1b", "#2ccc1b", "#cc1bcc"]

    adata_time.obs["celltype"] = adata_time.obs["celltype"].astype("category")

    set_plotting_vars(adata_time, _constants.PUSH, key=_constants.PUSH, value=plot_vars)

    push_initial_dist = np.zeros(

        shape=(len(adata_time[adata_time.obs["time"] == 0]),)

    )  # we need this for a cat. distr. in plots

    push_initial_dist[0:10] = 0.1

    nan2 = np.empty(len(adata_time[adata_time.obs["time"] == 2]))

    nan2[:] = np.nan

    adata_time.obs[_constants.PUSH] = np.hstack(

        (push_initial_dist, np.abs(rng.randn(len(adata_time[adata_time.obs["time"] == 1]))), nan2)

    )

    return adata_time

@pytest.fixture

def adata_pl_pull(adata_time: AnnData) -> AnnData:

    rng = np.random.RandomState(0)

    plot_vars = {"key": "time", "data": "celltype", "subset": "A", "source": 0, "target": 1}

    adata_time.uns["celltype_colors"] = ["#cc1b1b", "#2ccc1b", "#cc1bcc"]

    adata_time.obs["celltype"] = adata_time.obs["celltype"].astype("category")

    set_plotting_vars(adata_time, _constants.PULL, key=_constants.PULL, value=plot_vars)

    pull_initial_dist = np.zeros(

        shape=(len(adata_time[adata_time.obs["time"] == 1]),)

    )  # we need this for a cat. distr. in plots

    pull_initial_dist[0:10] = 0.1

    rand0 = np.abs(rng.randn(len(adata_time[adata_time.obs["time"] == 0])))

    nan2 = np.empty(len(adata_time[adata_time.obs["time"] == 2]))

    nan2[:] = np.nan

    adata_time.obs[_constants.PULL] = np.hstack((rand0, pull_initial_dist, nan2))

    return adata_time

@pytest.fixture

def adata_pl_sankey(adata_time: AnnData) -> AnnData:

    rng = np.random.RandomState(0)

    celltypes = ["A", "B", "C", "D", "E"]

    adata_time.obs["celltype"] = rng.choice(celltypes, size=len(adata_time))

    adata_time.obs["celltype"] = adata_time.obs["celltype"].astype("category")

    data1 = np.abs(rng.randn(5, 5))

    data2 = np.abs(rng.randn(5, 5))

    tm1 = pd.DataFrame(data=data1, index=celltypes, columns=celltypes)

    tm2 = pd.DataFrame(data=data2, index=celltypes, columns=celltypes)

    plot_vars = {"transition_matrices": [tm1, tm2], "captions": ["0", "1"], "key": "celltype"}

    set_plotting_vars(adata_time, _constants.SANKEY, key=_constants.SANKEY, value=plot_vars)

    return adata_time

def _decorate(fn: Callable, clsname: str, name: Optional[str] = None) -> Callable:

    @wraps(fn)

    def save_and_compare(self, *args, **kwargs):

        fn(self, *args, **kwargs)

        self.compare(fig_name)

    if not callable(fn):

        raise TypeError(f"Expected a `callable` for class `{clsname}`, found `{type(fn).__name__}`.")

    name = fn.__name__ if name is None else name

    if not name.startswith("test_plot_") or not clsname.startswith("Test"):

        return fn

    fig_name = f"{clsname[4:]}_{name[10:]}"

    return save_and_compare

class PlotTesterMeta(ABCMeta):

    def __new__(cls, clsname, superclasses, attributedict):

        for key, value in attributedict.items():

            if callable(value):

                attributedict[key] = _decorate(value, clsname, name=key)

        return super().__new__(cls, clsname, superclasses, attributedict)

# ideally, we would you metaclass=PlotTesterMeta and all plotting tests just subclass this

# but for some reason, pytest erases the metaclass info

class PlotTester(ABC):  # noqa: B024

    @classmethod

    def compare(cls, basename: str, tolerance: Optional[float] = None):

        ACTUAL.mkdir(parents=True, exist_ok=True)

        out_path = ACTUAL / f"{basename}.png"

        plt.savefig(out_path, dpi=DPI)

        plt.close()

        tolerance = TOL if tolerance is None else tolerance

        res = compare_images(str(EXPECTED / f"{basename}.png"), str(out_path), tolerance)

        assert res is None, res