from context import models, pl, tl, score

import mudata as md

import anndata as ad

import torch

import numpy as np

# Define some gene names (useful for enrichment analysis).

gene_names = [

    "ENSG00000125877",

    "ENSG00000184840",

    "ENSG00000164440",

    "ENSG00000177144",

    "ENSG00000186815",

    "ENSG00000079974",

    "ENSG00000136159",

    "ENSG00000177243",

    "ENSG00000163932",

    "ENSG00000112799",

    "ENSG00000075618",

    "ENSG00000092531",

    "ENSG00000171408",

    "ENSG00000150527",

    "ENSG00000202429",

    "ENSG00000140807",

    "ENSG00000154589",

    "ENSG00000166263",

    "ENSG00000205268",

    "ENSG00000115008",

]

n_cells, n_genes, n_peaks = 20, len(gene_names), 5

latent_dim = 5

# Create a random anndata object for RNA.

rna = ad.AnnData(np.random.rand(n_cells, n_genes))

rna.var["highly_variable"] = True

# Create a random anndata object for ATAC.

atac = ad.AnnData(np.random.rand(n_cells, n_peaks))

atac.var["highly_variable"] = True

# Create a MuData object combining RNA and ATAC.

mdata = md.MuData({"rna": rna, "atac": atac})

mdata.obs["rna:mod_weight"] = 0.5

mdata.obs["atac:mod_weight"] = 0.5

mdata.obs["label"] = np.random.choice(["A", "B", "C"], size=n_cells)

def test_default_params():

    # Initialize the Mowgli model.

    model = models.MowgliModel(

        latent_dim=latent_dim,

        cost_path={

            "rna": "cost_rna.npy",

            "atac": "cost_atac.npy",

        },

    )

    # Train the model.

    model.train(mdata)

    # Check the size of the embedding.

    assert mdata.obsm["W_OT"].shape == (n_cells, latent_dim)

    # Check the size of the dictionaries.

    assert mdata["rna"].uns["H_OT"].shape == (n_genes, latent_dim)

    assert mdata["atac"].uns["H_OT"].shape == (n_peaks, latent_dim)

def test_custom_params():

    # Initialize the Mowgli model.

    model = models.MowgliModel(

        latent_dim=latent_dim,

        h_regularization={"rna": 0.1, "atac": 0.1},

        use_mod_weight=True,

        pca_cost=True,

        cost_path={

            "rna": "cost_rna.npy",

            "atac": "cost_atac.npy",

        },

    )

    model.init_parameters(

        mdata,

        force_recompute=True,

        normalize_rows=True,

        dtype=torch.float,

        device="cpu",

    )

    # Train the model.

    model.train(mdata, optim_name="adam")

    # Check the size of the embedding.

    assert mdata.obsm["W_OT"].shape == (n_cells, latent_dim)

    # Check the size of the dictionaries.

    assert mdata["rna"].uns["H_OT"].shape == (n_genes, latent_dim)

    assert mdata["atac"].uns["H_OT"].shape == (n_peaks, latent_dim)

def test_plotting():

    # Make a clustermap.

    pl.clustermap(mdata, show=False)

    # Make a violin plot.

    pl.factor_violin(mdata, groupby="label", dim=0, show=False)

    # Make a heatmap.

    pl.heatmap(mdata, groupby="label", show=False)

def test_tools():

    # Compute top genes.

    tl.top_features(mdata, mod="rna", dim=0, threshold=0.2)

    # Compute top peaks.

    tl.top_features(mdata, mod="atac", dim=0, threshold=0.2)

    # Compute enrichment.

    tl.enrich(mdata, n_genes=10, ordered=False)

def test_score():

    # Compute a silhouette score.

    score.embedding_silhouette_score(

        embedding=mdata.obsm["W_OT"],

        labels=mdata.obs["label"],

        metric="euclidean",

    )

    # Compute leiden clustering across resolutions.

    score.embedding_leiden_across_resolutions(

        embedding=mdata.obsm["W_OT"],

        labels=mdata.obs["label"],

        n_neighbors=10,

        resolutions=[0.1, 0.5, 1.0],

    )

    # Compute a knn from the embedding.

    knn = score.embedding_to_knn(embedding=mdata.obsm["W_OT"], k=15, metric="euclidean")

    # Compute the knn purity score.

    score.knn_purity_score(knn=knn, labels=mdata.obs["label"])