from typing import List, Optional

import pytest

import numpy as np

import matplotlib as mpl

from anndata import AnnData

import moscot as msc
from moscot.plotting._utils import _input_to_adatas
from tests._utils import Problem


class TestMoscotPl:
    def test_input_to_adatas_problem(self, adata_time: AnnData):
        p = Problem(adata_time)
        adata1, adata2 = _input_to_adatas(p)
        assert isinstance(adata1, AnnData)
        assert isinstance(adata2, AnnData)
        np.testing.assert_array_equal(adata1.X.toarray(), adata_time.X.toarray())
        np.testing.assert_array_equal(adata2.X.toarray(), adata_time.X.toarray())

    def test_input_to_adatas_adata(self, adata_time: AnnData):
        adata1, adata2 = _input_to_adatas(adata_time)
        assert isinstance(adata1, AnnData)
        assert isinstance(adata2, AnnData)
        np.testing.assert_array_equal(adata1.X.toarray(), adata_time.X.toarray())
        np.testing.assert_array_equal(adata2.X.toarray(), adata_time.X.toarray())

    @pytest.mark.parametrize("return_fig", [True, False])
    def test_cell_transition(self, adata_pl_cell_transition: AnnData, return_fig: bool):
        fig = msc.plotting.cell_transition(adata_pl_cell_transition, return_fig=return_fig)
        if return_fig:
            assert fig is not None
            assert isinstance(fig, mpl.figure.Figure)
        else:
            assert fig is None

    @pytest.mark.parametrize("time_points", [None, [0]])
    @pytest.mark.parametrize("return_fig", [True, False])
    def test_push(self, adata_pl_push: AnnData, time_points: Optional[List[int]], return_fig: bool):
        fig = msc.plotting.push(adata_pl_push, time_points=time_points, return_fig=return_fig)

        if return_fig:
            assert fig is not None
            assert isinstance(fig, mpl.figure.Figure)
        else:
            assert fig is None

    @pytest.mark.parametrize("time_points", [None, [0]])
    @pytest.mark.parametrize("return_fig", [True, False])
    def test_pull(
        self,
        adata_pl_pull: AnnData,
        time_points: Optional[List[int]],
        return_fig: bool,
    ):
        fig = msc.plotting.pull(adata_pl_pull, time_points=time_points, return_fig=return_fig)
        if return_fig:
            assert fig is not None
            assert isinstance(fig, mpl.figure.Figure)
        else:
            assert fig is None

    @pytest.mark.parametrize("return_fig", [True, False])
    @pytest.mark.parametrize("interpolate_color", [True, False])
    def test_sankey(self, adata_pl_sankey: AnnData, return_fig: bool, interpolate_color: bool):
        fig = msc.plotting.sankey(
            adata_pl_sankey,
            return_fig=return_fig,
            interpolate_color=interpolate_color,
        )
        if return_fig:
            assert fig is not None
            assert isinstance(fig, mpl.figure.Figure)
        else:
            assert fig is None