from typing import Any, List, Optional, Tuple, Type, Union

import numpy as np

import pandas as pd

from scipy.sparse import csr_matrix

from anndata import AnnData

from moscot._types import ArrayLike

from moscot.base.output import MatrixSolverOutput

from moscot.base.problems import AnalysisMixin, CompoundProblem, OTProblem

from moscot.base.problems.compound_problem import B

Geom_t = Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]

RTOL = 1e-6

ATOL = 1e-6

class CompoundProblemWithMixin(CompoundProblem, AnalysisMixin):

    @property

    def _base_problem_type(self) -> Type[B]:

        return OTProblem

    @property

    def _valid_policies(self) -> Tuple[str, ...]:

        return ()

class MockSolverOutput(MatrixSolverOutput):

    @property

    def cost(self) -> float:

        return 0.5

    @property

    def converged(self) -> bool:

        return True

    @property

    def is_linear(self) -> bool:

        return True

    @property

    def potentials(self) -> Tuple[Optional[ArrayLike], Optional[ArrayLike]]:

        return None, None

    def _ones(self, n: int) -> ArrayLike:

        return np.ones(n)

def _make_adata(grid: ArrayLike, n: int, seed, cat_key: str = "covariate", num_categories: int = 3) -> List[AnnData]:

    rng = np.random.RandomState(seed)

    n_cells = 100

    X = rng.normal(size=(n_cells, 60))

    # generate a categorical variable

    categories = [f"cat_{i+1}" for i in range(num_categories)]

    categorical_data = rng.choice(categories, size=n_cells)

    adatas = []

    for _ in range(n):

        obs_df = pd.DataFrame({cat_key: pd.Categorical(categorical_data)})

        adatas.append(AnnData(X=csr_matrix(X), obs=obs_df, obsm={"spatial": grid.copy()}))

    return adatas

def _adata_spatial_split(adata: AnnData) -> Tuple[AnnData, AnnData]:

    adata_ref = adata[adata.obs.batch == "0"].copy()

    adata_ref.obsm.pop("spatial")

    adata_sp = adata[adata.obs.batch != "0"].copy()

    return adata_ref, adata_sp

def _make_grid(grid_size: int) -> ArrayLike:

    xlimits = ylimits = [0, 10]

    x1s = np.linspace(*xlimits, num=grid_size)

    x2s = np.linspace(*ylimits, num=grid_size)

    X1, X2 = np.meshgrid(x1s, x2s)

    return np.vstack([X1.ravel(), X2.ravel()]).T

def _assert_marginals_set(adata_time, problem, key, marginal_keys):

    """Helper function to check if marginals are set correctly"""

    adata_time0 = adata_time[key[0] == adata_time.obs["time"]]

    adata_time1 = adata_time[key[1] == adata_time.obs["time"]]

    if marginal_keys[0] is not None:  # check if marginal keys are set

        a = adata_time0.obs[marginal_keys[0]].values

        b = adata_time1.obs[marginal_keys[1]].values

        assert np.allclose(problem[key].a, a)

        assert np.allclose(problem[key].b, b)

    else:  # otherwise check if marginals are uniform

        assert np.allclose(problem[key].a, 1.0 / adata_time0.shape[0])

        assert np.allclose(problem[key].b, 1.0 / adata_time1.shape[0])

class Problem(CompoundProblem[Any, OTProblem]):

    @property

    def _base_problem_type(self) -> Type[B]:

        return OTProblem

    @property

    def _valid_policies(self) -> Tuple[str, ...]:

        return ()