import snapatac2 as snap

import numpy as np
import anndata as ad
from pathlib import Path
from collections import defaultdict
import pytest
from hypothesis import given, settings, HealthCheck, strategies as st
from hypothesis.extra.numpy import *
from scipy.sparse import csr_matrix
import gzip

from pytest import fixture
import os
import shutil

@fixture
def datadir(tmpdir, request):
    """
    Fixture responsible for searching a folder with the same name of test
    module and, if available, moving all contents to a temporary directory so
    tests can use them freely.
    """
    filename = request.module.__file__
    test_dir, _ = os.path.splitext(filename)

    if os.path.isdir(test_dir):
        for item in os.listdir(test_dir):
            src = os.path.join(test_dir, item)
            dst = os.path.join(tmpdir, item)
            if os.path.isdir(src):
                shutil.copytree(src, dst)
            else:
                shutil.copy2(src, dst)
    return tmpdir


def h5ad(dir=Path("./")):
    import uuid
    dir.mkdir(exist_ok=True)
    return str(dir / Path(str(uuid.uuid4()) + ".h5ad"))

def test_aggregation1():
    x = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=np.float64)
    x = np.concat([x] * 3333, axis=0)

    groups = np.random.choice(["A", "B", "C", "D", "E"], size=x.shape[0])
    var_names_len = np.random.randint(1000, 5000, size=x.shape[1])
    var_names = [f"chr1:0-{i}" for i in var_names_len]

    adata = ad.AnnData(
        X=x,
        obs=dict(groups=groups),
    )
    adata.var_names = var_names

    expected = defaultdict(list)
    for g, v in zip(groups, list(x)):
        expected[g].append(v)
    for k in expected.keys():
        expected[k] = np.array(expected[k], dtype="float64").sum(axis=0)

    actual = snap.tl.aggregate_X(adata, groupby=groups)
    actual = {actual.obs_names[i]: np.ravel(actual.X[i, :]) for i in range(actual.n_obs)}
    for k, v in expected.items():
        np.testing.assert_array_almost_equal_nulp(v, actual[k])

    actual = snap.tl.aggregate_X(adata, groupby=groups, normalize='RPM')
    actual = {actual.obs_names[i]: np.ravel(actual.X[i, :]) for i in range(actual.n_obs)}
    for k, v in expected.items():
        v /= v.sum() / 1e6
        np.testing.assert_array_almost_equal_nulp(v, actual[k])

    actual = snap.tl.aggregate_X(adata, groupby=groups, normalize='RPKM')
    actual = {actual.obs_names[i]: np.ravel(actual.X[i, :]) for i in range(actual.n_obs)}
    for k, v in expected.items():
        v /= v.sum() / 1e6
        v /= (var_names_len / 1000.0)
        np.testing.assert_array_almost_equal_nulp(v, actual[k], nulp=3)

def test_aggregation2():
    x = np.random.poisson(1.0, (10_000, 50)).astype(np.float64)
    groups = np.random.choice(["A", "B", "C", "D", "E"], size=x.shape[0])
    obs_names = [str(i) for i in range(len(groups))]
    var_names = [str(i) for i in range(x.shape[1])]

    adata = ad.AnnData(
        X=x,
        obs=dict(ident=obs_names, groups=groups),
        var=dict(ident=var_names),
    )

    expected = defaultdict(list)
    for g, v in zip(groups, list(x)):
        expected[g].append(v)
    for k in expected.keys():
        expected[k] = np.array(expected[k], dtype="float64").sum(axis=0)
    expected = dict(expected.items())

    np.testing.assert_array_almost_equal_nulp(
        x.sum(axis=0),
        np.ravel(snap.tl.aggregate_X(adata).X),
    )

    np.testing.assert_array_almost_equal_nulp(
        np.array(list(expected.values())),
        snap.tl.aggregate_X(adata, groupby=groups).X,
    )

    adata.X = csr_matrix(adata.X[:])
    np.testing.assert_array_almost_equal_nulp(
        np.array(list(expected.values())),
        snap.tl.aggregate_X(adata, groupby=groups).X,
    )

def test_make_fragment(datadir, tmp_path):
    bam = str(datadir.join('test.bam'))
    bed = str(datadir.join('test.bed.gz'))
    output = str(tmp_path) + "/out.bed.gz"
    snap.pp.make_fragment_file(bam, output, True, barcode_regex="(^[ATCG]+):", chunk_size=5000)

    with gzip.open(bed, 'rt') as fl:
        expected = sorted(fl.readlines())

    with gzip.open(output, 'rt') as fl:
        actual = sorted(fl.readlines())
    
    assert expected == actual

@given(
    mat = arrays(
        np.float64, (50, 100),
        elements = {"allow_subnormal": False, "allow_nan": False, "allow_infinity": False, "min_value": 1, "max_value": 100},
    ),
)
@settings(deadline = None, max_examples=5, suppress_health_check = [HealthCheck.function_scoped_fixture])
def test_reproducibility(mat):
    adata = ad.AnnData(X=csr_matrix(mat))
    embeddings = []
    for _ in range(3):
        embeddings.append(snap.tl.spectral(adata, features=None, random_state=0, inplace=False)[1])
    for x in embeddings:
        np.testing.assert_array_equal(x, embeddings[0])

    snap.tl.spectral(adata, features=None, random_state=0)
    knn = []
    for _ in range(3):
        knn.append(snap.pp.knn(adata, random_state=0, n_neighbors=25, inplace=False).todense())
    for x in knn:
        np.testing.assert_array_equal(x, knn[0])

    snap.pp.knn(adata, random_state=0, n_neighbors=25)
    leiden = []
    for _ in range(3):
        leiden.append(snap.tl.leiden(adata, random_state=0, resolution=1, n_iterations=10, inplace=False))
    for x in leiden:
        np.testing.assert_array_equal(x, leiden[0])

def read_bed(bed_file):
    content = []
    with gzip.open(bed_file, 'rt') as f:
        for line in f:
            if line.startswith('chr'):
                items = line.strip().split('\t')
                items.pop(4)
                content.append(items)
        return sorted(content)

def test_import(datadir):
    test_files = [
        (snap.datasets.pbmc500(downsample=True), True),
        (str(datadir.join('test_clean.tsv.gz')), True),
        (str(datadir.join('test_single.tsv.gz')), False),
    ]

    for fl, paired in test_files:
        data = snap.pp.import_fragments(
            fl,
            chrom_sizes=snap.genome.hg38,
            is_paired=paired,
            min_num_fragments=0,
            sorted_by_barcode=False,
        )

        data.obs['group'] = 'test_import'
        outputs = snap.ex.export_fragments(data, groupby="group", out_dir=str(datadir), suffix='.bed.gz')

        gold = read_bed(fl)
        test = read_bed(list(outputs.values())[0])
        for i, (g, t) in enumerate(zip(gold, test)):
            assert g == t, f"Line {i} mismatch: {g} != {t}"

def test_tile_matrix(datadir):
    def total_count(adata, bin_size):
        return snap.pp.add_tile_matrix(
            adata,
            bin_size=bin_size,
            inplace=False,
            counting_strategy='insertion',
            ).X.sum()

    data = snap.pp.import_fragments(
        snap.datasets.pbmc500(downsample=True),
        chrom_sizes=snap.genome.hg38,
        min_num_fragments=0,
        sorted_by_barcode=False,
    )

    counts = [total_count(data, i) for i in [500, 1000, 5000, 10000]]
    for i in range(1, len(counts)):
        assert counts[i] == counts[i - 1], f"Bin size {i} failed"

def test_spectral():
    fragment_file = snap.datasets.pbmc500(downsample=True)
    data = snap.pp.import_fragments(
        fragment_file,
        chrom_sizes=snap.genome.hg38,
        sorted_by_barcode=False,
    )
    snap.pp.add_tile_matrix(data)
    snap.metrics.tsse(data, snap.genome.hg38)
    snap.pp.filter_cells(data)
    snap.pp.select_features(data, n_features=4000)
    sp1 = snap.tl.spectral(data, random_state=0, inplace=False)[0]
    sp2 = snap.tl.spectral(data, random_state=0, inplace=False)[0]
    np.testing.assert_array_equal(sp1, sp2)
 