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 sys

import multivelo as mv

import scanpy as sc

import scvelo as scv

scv.settings.verbosity = 3

scv.settings.presenter_view = True

scv.set_figure_params('scvelo')

np.set_printoptions(suppress=True)

@pytest.fixture(scope="session")

def result_data_2():

    # read in the data

    adata_atac = sc.read("test_files/fig2_for_test.h5ad")

    # aggregate peaks

    adata_atac = mv.aggregate_peaks_10x(adata_atac,

                                        'test_files/peak_annotation.tsv',

                                        'test_files/feature_linkage.bedpe')

    return adata_atac

# test the aggregate_peaks_10x function

def test_agg_peaks(result_data_2):

    # the data indices we want to check

    indices = [0, 100000, 200000, 300000, 400000]

    # the results we should get

    data = [8.0, 4.0, 2.0, 2.0, 2.0]

    rows = [0, 1157, 2333, 3531, 4724]

    cols = [9, 276, 291, 78, 201]

    # convert the atac data into coo form

    atac_coo = result_data_2.X.tocoo()

    # check that there are the expected number of datapoints

    assert len(atac_coo.data) == 412887

    # make sure that the data, rows, and columns all match

    for n, i in enumerate(indices):

        assert atac_coo.data[i] == pytest.approx(data[n])

        assert atac_coo.row[i] == rows[n]

        assert atac_coo.col[i] == cols[n]

def test_tfidf(result_data_2):

    tfidf_result = result_data_2.copy()

    # run tfidf

    mv.tfidf_norm(tfidf_result)

    # the data indices we want to check

    indices = [0, 100000, 200000, 300000, 400000]

    # the results we should get

    data = [66.66449, 29.424345, 85.36392, 42.239613, 26.855387]

    rows = np.array([0, 1157, 2333, 3531, 4724])

    cols = np.array([9, 276, 291, 78, 201])

    # convert the atac data into coo form

    atac_coo = tfidf_result.X.tocoo()

    # make sure that the length of the data array is correct

    assert len(atac_coo.data) == 412887

    # make sure that the data, rows, and columns all match

    for n, i in enumerate(indices):

        assert atac_coo.data[i] == pytest.approx(data[n])

        assert atac_coo.row[i] == rows[n]

        assert atac_coo.col[i] == cols[n]

def test_smooth(result_data_2):

    new_result = result_data_2.copy()

    # load in the smoothing matrices

    nn_idx = np.loadtxt("test_files/nn_idx.txt", delimiter=',')

    nn_dist = np.loadtxt("test_files/nn_dist.txt", delimiter=',')

    # subset the ATAC data to make sure we can use the matrices

    atac_smooth = new_result[:nn_idx.shape[0], :]

    # run knn_smooth_chrom

    mv.knn_smooth_chrom(atac_smooth, nn_idx, nn_dist)

    # the data indices we want to check

    indices = [0, 70000, 140000, 210000, 280000]

    # the results we should get

    data = [8.0, 4.0, 2.0, 2.0, 2.0]

    rows = [0, 809, 1615, 2453, 3295]

    cols = [9, 327, 56, 25, 137]

    # convert the atac data into coo form

    atac_coo = atac_smooth.X.tocoo()

    # make sure that the length of the data array is correct

    assert len(atac_coo.data) == 285810

    # make sure that the data, rows, and columns all match

    for n, i in enumerate(indices):

        assert atac_coo.data[i] == pytest.approx(data[n])

        assert atac_coo.row[i] == rows[n]

        assert atac_coo.col[i] == cols[n]