import os

import sys

current_path = os.path.dirname(__file__)

src_path = os.path.join(current_path, "..")

sys.path.append(src_path)

import pytest

import numpy as np

import multivelo as mv

import scanpy as sc

import scvelo as scv

import sys

sys.path.append("/..")

scv.settings.verbosity = 3

scv.settings.presenter_view = True

scv.set_figure_params('scvelo')

np.set_printoptions(suppress=True)

rna_path = "test_files/adata_postpro.h5ad"

atac_path = "test_files/adata_atac_postpro.h5ad"

@pytest.fixture(scope="session")

def result_data_4():

    # read in the original AnnData objects

    adata_rna = sc.read(rna_path)

    adata_atac = sc.read(atac_path)

    # subset genes to run faster

    gene_list = ["Shh", "Heg1", "Cux1", "Lef1"]

    # run our first function to test (recover_dynamics_chrom)

    adata_result = mv.recover_dynamics_chrom(adata_rna,

                                             adata_atac,

                                             gene_list=gene_list,

                                             max_iter=5,

                                             init_mode="invert",

                                             parallel=True,

                                             n_jobs=15,

                                             save_plot=False,

                                             rna_only=False,

                                             fit=True,

                                             n_anchors=500,

                                             extra_color_key='celltype')

    return adata_result

# the next three tests check to see if recover_dynamics_chrom calculated

# the correct parameters for each of our four genes

def test_alpha(result_data_4):

    alpha = result_data_4.var["fit_alpha"]

    assert alpha[0] == pytest.approx(0.45878197934025416)

    assert alpha[1] == pytest.approx(0.08032904996744818)

    assert alpha[2] == pytest.approx(1.5346878202804608)

    assert alpha[3] == pytest.approx(0.9652887906148591)

def test_beta(result_data_4):

    beta = result_data_4.var["fit_beta"]

    assert beta[0] == pytest.approx(0.28770367567423)

    assert beta[1] == pytest.approx(0.14497469719573167)

    assert beta[2] == pytest.approx(0.564865749852349)

    assert beta[3] == pytest.approx(0.2522643118709811)

def test_gamma(result_data_4):

    gamma = result_data_4.var["fit_gamma"]

    assert gamma[0] == pytest.approx(0.19648836445315102)

    assert gamma[1] == pytest.approx(0.07703610603664116)

    assert gamma[2] == pytest.approx(1.0079569101225154)

    assert gamma[3] == pytest.approx(0.7485734061079243)

def test_embedding_stream(result_data_4):

    mv.velocity_graph(result_data_4)

    ax = mv.velocity_embedding_stream(result_data_4, basis='umap',

                                      color='celltype', show=False)

    assert ax is not None

    assert ax.axis()[0] == pytest.approx(-2.0698418340618714)

    assert ax.axis()[1] == pytest.approx(8.961822542538197)

    assert ax.axis()[2] == pytest.approx(-14.418079041548095)

    assert ax.axis()[3] == pytest.approx(-7.789863798927619)

    assert ax.get_xlim()[0] == pytest.approx(-2.0698418340618714)

    assert ax.get_xlim()[1] == pytest.approx(8.961822542538197)

    assert ax.get_ylim()[0] == pytest.approx(-14.418079041548095)

    assert ax.get_ylim()[1] == pytest.approx(-7.789863798927619)

# tests the latent_time function

def test_latent_time(result_data_4):

    mv.velocity_graph(result_data_4)

    mv.latent_time(result_data_4)

    latent_time = result_data_4.obs["latent_time"]

    assert latent_time.shape[0] == 6436

# test the velocity_graph function

def test_velo_graph(result_data_4):

    mv.velocity_graph(result_data_4)

    digits = 8

    v_graph_mat = result_data_4.uns["velo_s_norm_graph"].tocoo()

    v_graph = v_graph_mat.data

    v_graph = v_graph.astype(float)

    v_graph = v_graph.round(decimals=digits)

    v_graph_rows = v_graph_mat.row

    v_graph_cols = v_graph_mat.col

    assert len(v_graph) == 1883599

    assert v_graph[0] == pytest.approx(1.0)

    assert v_graph[500000] == pytest.approx(1.0)

    assert v_graph[1005000] == pytest.approx(0.99999994)

    assert v_graph[1500000] == pytest.approx(1.0)

    assert v_graph_rows[0] == 0

    assert v_graph_rows[500000] == 1411

    assert v_graph_rows[1005000] == 2834

    assert v_graph_rows[1500000] == 4985

    assert v_graph_cols[0] == 7

    assert v_graph_cols[500000] == 2406

    assert v_graph_cols[1005000] == 2892

    assert v_graph_cols[1500000] == 2480

@pytest.fixture(scope="session")

def lrt_compute():

    # read in the original AnnData objects

    adata_rna = sc.read(rna_path)

    adata_atac = sc.read(atac_path)

    # subset genes to run faster

    gene_list = ["Shh", "Heg1", "Cux1", "Lef1"]

    # run our first function to test (LRT_decoupling)

    w_de, wo_de, res = mv.LRT_decoupling(adata_rna,

                                         adata_atac,

                                         gene_list=gene_list,

                                         max_iter=5,

                                         init_mode="invert",

                                         parallel=True,

                                         n_jobs=15,

                                         save_plot=False,

                                         rna_only=False,

                                         fit=True,

                                         n_anchors=500,

                                         extra_color_key='celltype')

    # w_de = with decoupling

    # wo_de = without decoupling

    # res = LRT stats

    return (w_de, wo_de, res)

def decouple_test(lrt_compute):

    w_decouple = lrt_compute[0]

    alpha_c = w_decouple.var["fit_alpha_c"]

    assert alpha_c[0] == pytest.approx(0.057961)

    assert alpha_c[1] == pytest.approx(0.039439)

    assert alpha_c[2] == pytest.approx(0.076731)

    assert alpha_c[3] == pytest.approx(0.063575)

    beta = w_decouple.var["fit_beta"]

    assert beta[0] == pytest.approx(0.287704)

    assert beta[1] == pytest.approx(0.144975)

    assert beta[2] == pytest.approx(0.564866)

    assert beta[3] == pytest.approx(0.252264)

    gamma = w_decouple.var["fit_gamma"]

    assert gamma[0] == pytest.approx(0.196488)

    assert gamma[1] == pytest.approx(0.077036)

    assert gamma[2] == pytest.approx(1.007957)

    assert gamma[3] == pytest.approx(0.748573)

def no_decouple_test(lrt_compute):

    print("No decouple test")

    wo_decouple = lrt_compute[1]

    alpha_c = wo_decouple.var["fit_alpha_c"]

    assert alpha_c[0] == pytest.approx(0.093752)

    assert alpha_c[1] == pytest.approx(0.041792)

    assert alpha_c[2] == pytest.approx(0.051228)

    assert alpha_c[3] == pytest.approx(0.050951)

    beta = wo_decouple.var["fit_beta"]

    assert beta[0] == pytest.approx(0.840938)

    assert beta[1] == pytest.approx(0.182773)

    assert beta[2] == pytest.approx(0.326623)

    assert beta[3] == pytest.approx(0.232073)

    gamma = wo_decouple.var["fit_gamma"]

    assert gamma[0] == pytest.approx(0.561730)

    assert gamma[1] == pytest.approx(0.106799)

    assert gamma[2] == pytest.approx(0.783257)

    assert gamma[3] == pytest.approx(0.705256)

def lrt_res_test(lrt_compute):

    res = lrt_compute[2]

    likelihood_c_w_decoupled = res["likelihood_c_w_decoupled"]

    assert likelihood_c_w_decoupled[0] == pytest.approx(0.279303)

    assert likelihood_c_w_decoupled[1] == pytest.approx(0.186213)

    assert likelihood_c_w_decoupled[2] == pytest.approx(0.295591)

    assert likelihood_c_w_decoupled[3] == pytest.approx(0.144158)

    likelihood_c_wo_decoupled = res["likelihood_c_wo_decoupled"]

    assert likelihood_c_wo_decoupled[0] == pytest.approx(0.270491)

    assert likelihood_c_wo_decoupled[1] == pytest.approx(0.180695)

    assert likelihood_c_wo_decoupled[2] == pytest.approx(0.294631)

    assert likelihood_c_wo_decoupled[3] == pytest.approx(0.175622)

    LRT_c = res["LRT_c"]

    assert LRT_c[0] == pytest.approx(412.637730)

    assert LRT_c[1] == pytest.approx(387.177688)

    assert LRT_c[2] == pytest.approx(41.850304)

    assert LRT_c[3] == pytest.approx(-2541.289231)

    pval_c = res["pval_c"]

    assert pval_c[0] == pytest.approx(9.771580e-92)

    assert pval_c[1] == pytest.approx(3.406463e-86)

    assert pval_c[2] == pytest.approx(9.853544e-11)

    assert pval_c[3] == pytest.approx(1.000000e+00)

    likelihood_w_decoupled = res["likelihood_w_decoupled"]

    assert likelihood_w_decoupled[0] == pytest.approx(0.177979)

    assert likelihood_w_decoupled[1] == pytest.approx(0.008453)

    assert likelihood_w_decoupled[2] == pytest.approx(0.140156)

    assert likelihood_w_decoupled[3] == pytest.approx(0.005029)

    likelihood_wo_decoupled = res["likelihood_wo_decoupled"]

    assert likelihood_wo_decoupled[0] == pytest.approx(0.181317)

    assert likelihood_wo_decoupled[1] == pytest.approx(0.009486)

    assert likelihood_wo_decoupled[2] == pytest.approx(0.141367)

    assert likelihood_wo_decoupled[3] == pytest.approx(0.008299)

    LRT = res["LRT"]

    assert LRT[0] == pytest.approx(-239.217562)

    assert LRT[1] == pytest.approx(-1485.199859)

    assert LRT[2] == pytest.approx(-110.788912)

    assert LRT[3] == pytest.approx(-6447.599212)

    pval = res["pval"]

    assert pval[0] == pytest.approx(1.0)

    assert pval[1] == pytest.approx(1.0)

    assert pval[2] == pytest.approx(1.0)

    assert pval[3] == pytest.approx(1.0)

    c_likelihood = res["likelihood_c_w_decoupled"]

    assert c_likelihood[0] == pytest.approx(0.279303)

    assert c_likelihood[1] == pytest.approx(0.186213)

    assert c_likelihood[2] == pytest.approx(0.295591)

    assert c_likelihood[3] == pytest.approx(0.144158)

    def test_qc_metrics():

        adata_rna = sc.read(rna_path)

        mv.calculate_qc_metrics(adata_rna)

        total_unspliced = adata_rna.obs["total_unspliced"]

        assert total_unspliced.shape == (6436,)

        assert total_unspliced[0] == pytest.approx(91.709404)

        assert total_unspliced[1500] == pytest.approx(115.21283)

        assert total_unspliced[3000] == pytest.approx(61.402004)

        assert total_unspliced[4500] == pytest.approx(84.03409)

        assert total_unspliced[6000] == pytest.approx(61.26761)

        total_spliced = adata_rna.obs["total_spliced"]

        assert total_spliced.shape == (6436,)

        assert total_spliced[0] == pytest.approx(91.514175)

        assert total_spliced[1500] == pytest.approx(66.045616)

        assert total_spliced[3000] == pytest.approx(87.05275)

        assert total_spliced[4500] == pytest.approx(83.82857)

        assert total_spliced[6000] == pytest.approx(62.019516)

        unspliced_ratio = adata_rna.obs["unspliced_ratio"]

        assert unspliced_ratio.shape == (6436,)

        assert unspliced_ratio[0] == pytest.approx(0.5005328)

        assert unspliced_ratio[1500] == pytest.approx(0.6356273)

        assert unspliced_ratio[3000] == pytest.approx(0.4136075)

        assert unspliced_ratio[4500] == pytest.approx(0.50061214)

        assert unspliced_ratio[6000] == pytest.approx(0.4969506)

        cell_cycle_score = adata_rna.obs["cell_cycle_score"]

        assert cell_cycle_score.shape == (6436,)

        assert cell_cycle_score[0] == pytest.approx(-0.24967776384597046)

        assert cell_cycle_score[1500] == pytest.approx(0.5859756395543293)

        assert cell_cycle_score[3000] == pytest.approx(0.06501555292615813)

        assert cell_cycle_score[4500] == pytest.approx(0.1406775909466575)

        assert cell_cycle_score[6000] == pytest.approx(-0.33825528386759895)