from typing import Mapping, Optional

import pytest

import networkx as nx
import numpy as np

from moscot.datasets import simulate_data


class TestSimulateData:
    @pytest.mark.fast
    @pytest.mark.parametrize("n_distributions", [2, 4])
    @pytest.mark.parametrize("key", ["batch", "day"])
    def test_n_distributions(self, n_distributions: int, key: str):
        adata = simulate_data(n_distributions=n_distributions, key=key)
        assert key in adata.obs.columns
        assert adata.obs[key].nunique() == n_distributions

    @pytest.mark.fast
    @pytest.mark.parametrize("obs_to_add", [{"celltype": 2}, {"celltype": 5, "cluster": 4}])
    def test_obs_to_add(self, obs_to_add: Mapping[str, int]):
        adata = simulate_data(obs_to_add=obs_to_add)

        for colname, k in obs_to_add.items():
            assert colname in adata.obs.columns
            assert adata.obs[colname].nunique() == k

    @pytest.mark.fast
    @pytest.mark.parametrize("spatial_dim", [None, 2, 3])
    def test_quad_term_spatial(self, spatial_dim: Optional[int]):
        kwargs = {}
        if spatial_dim is not None:
            kwargs["spatial_dim"] = spatial_dim
        adata = simulate_data(quad_term="spatial", **kwargs)

        assert isinstance(adata.obsm["spatial"], np.ndarray)
        if spatial_dim is None:
            assert adata.obsm["spatial"].shape[1] == 2
        else:
            assert adata.obsm["spatial"].shape[1] == spatial_dim

    @pytest.mark.fast
    @pytest.mark.parametrize("n_intBCs", [None, 4, 7])
    @pytest.mark.parametrize("barcode_dim", [None, 5, 8])
    def test_quad_term_barcode(self, n_intBCs: Optional[int], barcode_dim: Optional[int]):
        kwargs = {}
        if n_intBCs is not None:
            kwargs["n_intBCs"] = n_intBCs
        if barcode_dim is not None:
            kwargs["barcode_dim"] = barcode_dim

        adata = simulate_data(quad_term="barcode", **kwargs)

        assert isinstance(adata.obsm["barcode"], np.ndarray)
        if barcode_dim is None:
            assert adata.obsm["barcode"].shape[1] == 10
        else:
            assert adata.obsm["barcode"].shape[1] == barcode_dim

        if n_intBCs is None:
            assert len(np.unique(adata.obsm["barcode"])) <= 20
        else:
            assert len(np.unique(adata.obsm["barcode"])) <= n_intBCs

    @pytest.mark.fast
    @pytest.mark.parametrize("n_initial_nodes", [None, 4, 7])
    @pytest.mark.parametrize("n_distributions", [3, 6])
    def test_quad_term_tree(self, n_initial_nodes: Optional[int], n_distributions: int):
        kwargs = {}
        if n_initial_nodes is not None:
            kwargs["n_initial_nodes"] = n_initial_nodes
        adata = simulate_data(quad_term="tree", key="day", n_distributions=n_distributions, **kwargs)

        assert isinstance(adata.uns["trees"], dict)
        assert len(adata.uns["trees"]) == n_distributions

        for i in range(len(adata.uns["trees"])):
            assert isinstance(adata.uns["trees"][i], nx.DiGraph)