from math import cos, sin

from typing import Literal, Optional, Tuple, Union

import pytest

import jax.numpy as jnp

import numpy as np

import pandas as pd

from jax import config

from scipy.sparse import csr_matrix

import matplotlib.pyplot as plt

import anndata as ad

import scanpy as sc

from anndata import AnnData

from tests._utils import Geom_t, _make_adata, _make_grid

ANGLES = (0, 30, 60)

# TODO(michalk8): consider passing this via env

config.update("jax_enable_x64", True)

_gt_temporal_adata = sc.read("tests/data/moscot_temporal_tests.h5ad")

def pytest_sessionstart() -> None:

    sc.pl.set_rcParams_defaults()

    sc.set_figure_params(dpi=40, color_map="viridis")

@pytest.fixture(autouse=True)

def _close_figure():

    # prevent `RuntimeWarning: More than 20 figures have been opened.`

    yield

    plt.close()

@pytest.fixture

def x() -> Geom_t:

    rng = np.random.RandomState(0)

    n = 20  # number of points in the first distribution

    sig = 1  # std of first distribution

    phi = np.arange(n)[:, None]

    xs = phi + sig * rng.randn(n, 1)

    return jnp.asarray(xs)

@pytest.fixture

def y() -> Geom_t:

    rng = np.random.RandomState(1)

    n2 = 30  # number of points in the second distribution

    sig = 1  # std of first distribution

    phi2 = np.arange(n2)[:, None]

    xt = phi2 + sig * rng.randn(n2, 1)

    return jnp.asarray(xt)

@pytest.fixture

def xy() -> Tuple[Geom_t, Geom_t]:

    rng = np.random.RandomState(2)

    n = 20  # number of points in the first distribution

    n2 = 30  # number of points in the second distribution

    sig = 1  # std of first distribution

    sig2 = 0.1  # std of second distribution

    phi = np.arange(n)[:, None]

    phi + sig * rng.randn(n, 1)

    ys = np.vstack((np.ones((n // 2, 1)), 0 * np.ones((n // 2, 1)))) + sig2 * rng.randn(n, 1)

    phi2 = np.arange(n2)[:, None]

    phi2 + sig * rng.randn(n2, 1)

    yt = np.vstack((np.ones((n2 // 2, 1)), 0 * np.ones((n2 // 2, 1)))) + sig2 * rng.randn(n2, 1)

    yt = yt[::-1, :]

    return jnp.asarray(ys), jnp.asarray(yt)

@pytest.fixture

def ab() -> Tuple[np.ndarray, np.ndarray]:

    rng = np.random.RandomState(42)

    return rng.normal(size=(20, 2)), rng.normal(size=(30, 4))

@pytest.fixture

def x_cost(x: Geom_t) -> jnp.ndarray:

    return ((x[:, None, :] - x[None, ...]) ** 2).sum(-1)

@pytest.fixture

def y_cost(y: Geom_t) -> jnp.ndarray:

    return ((y[:, None, :] - y[None, ...]) ** 2).sum(-1)

@pytest.fixture

def xy_cost(xy: Geom_t) -> jnp.ndarray:

    x, y = xy

    return ((x[:, None, :] - y[None, ...]) ** 2).sum(-1)

@pytest.fixture

def adata_x(x: Geom_t) -> AnnData:

    rng = np.random.RandomState(43)

    pc = rng.normal(size=(len(x), 4))

    return AnnData(X=np.asarray(x, dtype=float), obsm={"X_pca": pc})

@pytest.fixture

def adata_y(y: Geom_t) -> AnnData:

    rng = np.random.RandomState(44)

    pc = rng.normal(size=(len(y), 4))

    return AnnData(X=np.asarray(y, dtype=float), obsm={"X_pca": pc})

def creat_prob(n: int, *, uniform: bool = False, seed: Optional[int] = None) -> Geom_t:

    rng = np.random.RandomState(seed)

    a = np.ones((n,)) if uniform else np.abs(rng.normal(size=(n,)))

    a /= np.sum(a)

    return jnp.asarray(a)

@pytest.fixture

def adata_time() -> AnnData:

    rng = np.random.RandomState(42)

    adatas = [

        AnnData(

            X=csr_matrix(rng.normal(size=(96, 60))),

            obs={

                "left_marginals_balanced": creat_prob(96, seed=42),

                "right_marginals_balanced": creat_prob(96, seed=42),

            },

        )

        for _ in range(3)

    ]

    adata = ad.concat(adatas, label="time", index_unique="-")

    adata.obs["time"] = pd.to_numeric(adata.obs["time"]).astype("category")

    adata.obs["batch"] = rng.choice((0, 1, 2), len(adata))

    adata.obs["left_marginals_unbalanced"] = np.ones(len(adata))

    adata.obs["right_marginals_unbalanced"] = np.ones(len(adata))

    adata.obs["celltype"] = rng.choice(["A", "B", "C"], size=len(adata))

    # genes from mouse/human proliferation/apoptosis

    genes = ["ANLN", "ANP32E", "ATAD2", "Mcm4", "Smc4", "Gtse1", "ADD1", "AIFM3", "ANKH", "Ercc5", "Serpinb5", "Inhbb"]

    # genes which are transcription factors, 3 from drosophila, 2 from human, 1 from mouse

    genes += ["Cf2", "Dlip3", "Dref", "KLF12", "ZNF143", "Zic5"]

    adata.var.index = ["gene_" + el if i > len(genes) - 1 else genes[i] for i, el in enumerate(adata.var.index)]

    adata.obsm["X_umap"] = rng.randn(len(adata), 2)

    sc.pp.pca(adata)

    return adata

@pytest.fixture

def gt_temporal_adata() -> AnnData:

    adata = _gt_temporal_adata.copy()

    # TODO(michalk8): remove both lines once data has been regenerated

    adata.obs["day"] = pd.to_numeric(adata.obs["day"]).astype("category")

    adata.obs_names_make_unique()

    return adata

@pytest.fixture

def adata_space_rotate() -> AnnData:

    rng = np.random.RandomState(31)

    grid = _make_grid(10)

    adatas = _make_adata(grid, n=len(ANGLES), seed=32)

    for adata, angle in zip(adatas, ANGLES):

        theta = np.deg2rad(angle)

        rot = np.array([[cos(theta), -sin(theta)], [sin(theta), cos(theta)]])

        adata.obsm["spatial"] = np.dot(adata.obsm["spatial"], rot)

    adata = ad.concat(adatas, label="batch", index_unique="-")

    adata.obs["celltype"] = rng.choice(["A", "B", "C"], size=len(adata))

    adata.uns["spatial"] = {}

    sc.pp.pca(adata)

    return adata

@pytest.fixture

def adata_mapping() -> AnnData:

    grid = _make_grid(10)

    adataref, adata1, adata2 = _make_adata(grid, n=3, seed=17, cat_key="covariate", num_categories=3)

    sc.pp.pca(adataref, n_comps=30)

    return ad.concat([adataref, adata1, adata2], label="batch", join="outer", index_unique="-")

@pytest.fixture

def adata_translation() -> AnnData:

    rng = np.random.RandomState(31)

    adatas = [AnnData(X=csr_matrix(rng.normal(size=(100, 60)))) for _ in range(3)]

    adata = ad.concat(adatas, label="batch", index_unique="-")

    adata.obs["celltype"] = rng.choice(["A", "B", "C"], size=len(adata))

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

    adata.layers["counts"] = adata.X.toarray()

    sc.pp.pca(adata)

    return adata

@pytest.fixture

def adata_translation_split(adata_translation) -> Tuple[AnnData, AnnData]:

    rng = np.random.RandomState(15)

    adata_src = adata_translation[adata_translation.obs.batch != "0"].copy()

    adata_tgt = adata_translation[adata_translation.obs.batch == "0"].copy()

    adata_src.obsm["emb_src"] = rng.normal(size=(adata_src.shape[0], 5))

    adata_tgt.obsm["emb_tgt"] = rng.normal(size=(adata_tgt.shape[0], 15))

    return adata_src, adata_tgt

@pytest.fixture

def adata_anno(

    problem_kind: Literal["temporal", "cross_modality", "alignment", "mapping"],

) -> Union[AnnData, Tuple[AnnData, AnnData]]:

    rng = np.random.RandomState(31)

    adata_src = AnnData(X=csr_matrix(rng.normal(size=(10, 60))))

    rng_src = rng.choice(["A", "B", "C"], size=5).tolist()

    adata_src.obs["celltype1"] = ["C", "C", "A", "B", "B"] + rng_src

    adata_src.obs["celltype1"] = adata_src.obs["celltype1"].astype("category")

    adata_src.uns["expected_max1"] = ["C", "C", "A", "B", "B"] + rng_src + rng_src

    adata_src.uns["expected_sum1"] = ["C", "C", "B", "B", "B"] + rng_src + rng_src

    adata_tgt = AnnData(X=csr_matrix(rng.normal(size=(15, 60))))

    rng_tgt = rng.choice(["A", "B", "C"], size=5).tolist()

    adata_tgt.obs["celltype2"] = ["C", "C", "A", "B", "B"] + rng_tgt + rng_tgt

    adata_tgt.obs["celltype2"] = adata_tgt.obs["celltype2"].astype("category")

    adata_tgt.uns["expected_max2"] = ["C", "C", "A", "B", "B"] + rng_tgt

    adata_tgt.uns["expected_sum2"] = ["C", "C", "B", "B", "B"] + rng_tgt

    if problem_kind == "cross_modality":

        adata_src.obs["batch"] = "0"

        adata_tgt.obs["batch"] = "1"

        adata_src.obsm["emb_src"] = rng.normal(size=(adata_src.shape[0], 5))

        adata_tgt.obsm["emb_tgt"] = rng.normal(size=(adata_tgt.shape[0], 15))

        sc.pp.pca(adata_src)

        sc.pp.pca(adata_tgt)

        return adata_src, adata_tgt

    if problem_kind == "mapping":

        adata_src.obs["batch"] = "0"

        adata_tgt.obs["batch"] = "1"

        sc.pp.pca(adata_src)

        sc.pp.pca(adata_tgt)

        adata_tgt.obsm["spatial"] = rng.normal(size=(adata_tgt.n_obs, 2))

        return adata_src, adata_tgt

    if problem_kind == "alignment":

        adata_src.obsm["spatial"] = rng.normal(size=(adata_src.n_obs, 2))

        adata_tgt.obsm["spatial"] = rng.normal(size=(adata_tgt.n_obs, 2))

    key = "day" if problem_kind == "temporal" else "batch"

    adatas = [adata_src, adata_tgt]

    adata = ad.concat(adatas, join="outer", label=key, index_unique="-", uns_merge="unique")

    adata.obs[key] = (pd.to_numeric(adata.obs[key]) if key == "day" else adata.obs[key]).astype("category")

    adata.layers["counts"] = adata.X.toarray()

    sc.pp.pca(adata)

    return adata

@pytest.fixture

def gt_tm_annotation() -> np.ndarray:

    tm = np.zeros((10, 15))

    for i in range(10):

        tm[i][i] = 1

    for i in range(10, 15):

        tm[i - 5][i] = 1

    for j in range(2, 5):

        for i in range(2, 5):

            tm[i][j] = 0.3 if i != j else 0.4

    return tm