import pytest

import numpy as np

import pandas as pd

from sklearn.metrics import pairwise_distances

import anndata as ad

from anndata import AnnData

from tests._utils import Geom_t

@pytest.fixture

def adata_with_cost_matrix(adata_x: Geom_t, adata_y: Geom_t) -> AnnData:

    adata = ad.concat([adata_x, adata_y], label="batch", index_unique="-")

    C = pairwise_distances(adata_x.obsm["X_pca"], adata_y.obsm["X_pca"]) ** 2

    adata.obs["batch"] = pd.to_numeric(adata.obs["batch"])

    adata.uns[0] = C / C.mean()  # TODO(@MUCDK) make a callback function and replace this part

    return adata

@pytest.fixture

def adata_time_with_tmap(adata_time: AnnData) -> AnnData:

    adata = adata_time[adata_time.obs["time"].isin([0, 1])].copy()

    rng = np.random.RandomState(42)

    cell_types = ["cell_A", "cell_B", "cell_C", "cell_D"]

    cell_d1 = rng.multinomial(len(adata[adata.obs["time"] == 0]), [1 / len(cell_types)] * len(cell_types))

    cell_d2 = rng.multinomial(len(adata[adata.obs["time"] == 0]), [1 / len(cell_types)] * len(cell_types))

    a1 = np.concatenate(

        [["cell_A"] * cell_d1[0], ["cell_B"] * cell_d1[1], ["cell_C"] * cell_d1[2], ["cell_D"] * cell_d1[3]]

    ).flatten()

    a2 = np.concatenate(

        [["cell_A"] * cell_d2[0], ["cell_B"] * cell_d2[1], ["cell_C"] * cell_d2[2], ["cell_D"] * cell_d2[3]]

    ).flatten()

    adata.obs["cell_type"] = np.concatenate([a1, a2])

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

    cell_numbers_source = dict(adata[adata.obs["time"] == 0].obs["cell_type"].value_counts())

    cell_numbers_target = dict(adata[adata.obs["time"] == 1].obs["cell_type"].value_counts())

    trans_matrix = np.abs(rng.randn(len(cell_types), len(cell_types)))

    trans_matrix = trans_matrix / trans_matrix.sum(axis=1, keepdims=1)

    cell_transition_gt = pd.DataFrame(data=trans_matrix, index=cell_types, columns=cell_types)

    blocks = []

    for cell_row in cell_types:

        block_row = []

        for cell_col in cell_types:

            sub_trans_matrix = np.abs(rng.randn(cell_numbers_source[cell_row], cell_numbers_target[cell_col]))

            sub_trans_matrix /= sub_trans_matrix.sum() * (1 / cell_transition_gt.loc[cell_row, cell_col])

            block_row.append(sub_trans_matrix)

        blocks.append(block_row)

    transport_matrix = np.block(blocks)

    adata.uns["transport_matrix"] = transport_matrix

    adata.uns["cell_transition_gt"] = cell_transition_gt

    return adata

# keys for marginals

@pytest.fixture(

    params=[

        (None, None),

        ("left_marginals_balanced", "right_marginals_balanced"),

    ],

    ids=["default", "balanced"],

)

def marginal_keys(request):

    return request.param

sinkhorn_args_1 = {

    "epsilon": 0.7,

    "tau_a": 1.0,

    "tau_b": 1.0,

    "rank": 7,

    "initializer": "rank2",

    "initializer_kwargs": {},

    "jit": False,

    "threshold": 2e-3,

    "lse_mode": True,

    "norm_error": 2,

    "inner_iterations": 3,

    "min_iterations": 4,

    "max_iterations": 9,

    "gamma": 9.4,

    "gamma_rescale": False,

    "batch_size": None,  # in to_LRC() `batch_size` cannot be passed so we expect None.

    "scale_cost": "max_cost",

}

sinkhorn_args_2 = {  # no gamma/gamma_rescale as these are LR-specific

    "epsilon": 0.8,

    "tau_a": 0.9,

    "tau_b": 0.8,

    "rank": -1,

    "batch_size": 125,

    "initializer": "gaussian",

    "initializer_kwargs": {},

    "jit": True,

    "threshold": 3e-3,

    "lse_mode": False,

    "norm_error": 3,

    "inner_iterations": 4,

    "min_iterations": 1,

    "max_iterations": 2,

    "scale_cost": "mean",

}

linear_solver_kwargs1 = {

    "inner_iterations": 1,

    "min_iterations": 5,

    "max_iterations": 7,

    "lse_mode": False,

    "threshold": 5e-2,

    "norm_error": 4,

}

gw_args_1 = {  # no gamma/gamma_rescale/tolerances/ranks as these are LR-specific

    "epsilon": 0.5,

    "tau_a": 0.7,

    "tau_b": 0.8,

    "scale_cost": "max_cost",

    "rank": -1,

    "batch_size": 122,

    "initializer": None,

    "initializer_kwargs": {},

    "jit": True,

    "threshold": 3e-2,

    "min_iterations": 3,

    "max_iterations": 4,

    "gw_unbalanced_correction": True,

    "ranks": 4,

    "tolerances": 2e-2,

    "warm_start": False,

    "linear_solver_kwargs": linear_solver_kwargs1,

}

linear_solver_kwargs2 = {

    "inner_iterations": 3,

    "min_iterations": 7,

    "max_iterations": 8,

    "lse_mode": True,

    "threshold": 4e-2,

    "norm_error": 3,

    "gamma": 9.4,

    "gamma_rescale": False,

}

gw_args_2 = {

    "alpha": 0.4,

    "epsilon": 0.7,

    "tau_a": 1.0,

    "tau_b": 1.0,

    "scale_cost": "max_cost",

    "rank": 7,

    "batch_size": 123,

    "initializer": "rank2",

    "initializer_kwargs": {},

    "jit": False,

    "threshold": 2e-3,

    "min_iterations": 2,

    "max_iterations": 3,

    "gw_unbalanced_correction": False,

    "ranks": 3,

    "tolerances": 3e-2,

    # "linear_solver_kwargs": linear_solver_kwargs2,

}

gw_args_2 = {**gw_args_2, **linear_solver_kwargs2}

fgw_args_1 = gw_args_1.copy()

fgw_args_1["alpha"] = 0.6

fgw_args_2 = gw_args_2.copy()

fgw_args_2["alpha"] = 0.4

gw_solver_args = {

    "epsilon": "epsilon",

    "rank": "rank",

    "threshold": "threshold",

    "min_iterations": "min_iterations",

    "max_iterations": "max_iterations",

    "warm_start": "warm_start",

    "initializer": "initializer",

}

gw_lr_solver_args = {

    "epsilon": "epsilon",

    "rank": "rank",

    "threshold": "threshold",

    "min_iterations": "min_iterations",

    "max_iterations": "max_iterations",

    "initializer": "initializer",

}

gw_linear_solver_args = {

    "lse_mode": "lse_mode",

    "inner_iterations": "inner_iterations",

    "threshold": "threshold",

    "norm_error": "norm_error",

    "max_iterations": "max_iterations",

    "min_iterations": "min_iterations",

}

gw_lr_linear_solver_args = {

    "lse_mode": "lse_mode",

    "inner_iterations": "inner_iterations",

    "threshold": "threshold",

    "norm_error": "norm_error",

    "max_iterations": "max_iterations",

    "min_iterations": "min_iterations",

    "gamma": "gamma",

    "gamma_rescale": "gamma_rescale",

}

quad_prob_args = {

    "tau_a": "tau_a",

    "tau_b": "tau_b",

    "gw_unbalanced_correction": "gw_unbalanced_correction",

    "ranks": "ranks",

    "tolerances": "tolerances",

}

geometry_args = {"epsilon": "_epsilon_init", "scale_cost": "_scale_cost"}

pointcloud_args = {

    "batch_size": "_batch_size",

    "scale_cost": "_scale_cost",

}

lr_pointcloud_args = {

    "batch_size": "batch_size",

    "scale_cost": "_scale_cost",

}

sinkhorn_solver_args = {  # dictionary with key = moscot arg name, value = ott-jax attribute

    "lse_mode": "lse_mode",

    "threshold": "threshold",

    "norm_error": "norm_error",

    "inner_iterations": "inner_iterations",

    "min_iterations": "min_iterations",

    "max_iterations": "max_iterations",

    "initializer": "initializer",

    "initializer_kwargs": "initializer_kwargs",

}

lr_sinkhorn_solver_args = sinkhorn_solver_args.copy()

lr_sinkhorn_solver_args["gamma"] = "gamma"

lr_sinkhorn_solver_args["gamma_rescale"] = "gamma_rescale"

lin_prob_args = {

    "tau_a": "tau_a",

    "tau_b": "tau_b",

}

neurallin_cond_args_1 = {

    "batch_size": 8,

    "seed": 0,

    "iterations": 2,

    "valid_freq": 4,

}