import pytest
import numpy as np
import pandas as pd
from ..get_scores import *
from pkg_resources import resource_filename

diamond_output = resource_filename(__name__,
                                   'test_data/tiny_test.diamond.out')
diamond_df = read_diamond_output(diamond_output)
ref_base_data_path = resource_filename(__name__,
                                       'test_data/tiny_test.base_data')
ref_base_data = pd.read_csv(ref_base_data_path)


def test_read_diamond_output():
    empty_file = resource_filename(__name__, '__init__.py')
    assert len(read_diamond_output(empty_file)) == 0
    diamond_df = read_diamond_output(diamond_output)
    assert len(diamond_df) == 17
    assert len(diamond_df.columns) == 4


def test_get_n_effective_surplus_clades():
    assert get_n_effective_surplus_clades([1]) == 0
    assert get_n_effective_surplus_clades([2, 8, 1, 1, 3]) == pytest.approx(1.84,
                                                                            rel=1e-2)


def test_calc_expected_conditional_entropy():
    s = pd.Series
    assert calc_conditional_entropy(s([0, 0, 1, 1]), s([1, 1, 0, 0])) == 0
    assert calc_conditional_entropy(s([0, 0, 1, 1]), s([0, 1, 0, 1])) > 0.1


def test_create_base_data():
    new_base_data = create_base_data(diamond_df)
    pd.testing.assert_frame_equal(ref_base_data, new_base_data)


def test_get_stats():
    assert get_stats(diamond_df) == (17, 15)


def test_get_abundant_lineages_cutoff():
    assert get_abundant_lineages_cutoff(True, 3) == 10
    assert get_abundant_lineages_cutoff(False, 10) == 0.2


def test_calc_contamination_portion():
    counts = pd.Series({'a': 1, 'b': 2, 'c': 3, 'd': 4})
    assert calc_contamination_portion(counts) == 0.6


def test_mean():
    assert mean([1, 2, 3]) == 2
    assert mean([0]) == 0
    assert mean([1, 1, 1]) == 1


def test_calc_mean_hit_identity():
    assert calc_mean_hit_identity([1, 2, 3]) == 0.02


def test_is_chimeric():
    assert is_chimeric(0) is False
    assert is_chimeric(1) is True
    assert is_chimeric(0.44) is False
    assert is_chimeric(0.45) is False
    assert is_chimeric(0.46) is True


def test_get_scores_for_taxlevel():
    # Test kingdom level
    data = get_scores_for_taxlevel(ref_base_data,
                                   'kingdom',
                                   0.34,
                                   'test',
                                   35,
                                   17,
                                   15,
                                   11)
    assert data['genome'] == 'test'
    assert data['n_contigs'] == 15
    assert data['n_genes_called'] == 35
    assert data['n_genes_mapped'] == 17
    assert data['taxonomic_level'] == 'kingdom'
    assert data['clade_separation_score'] == 0
    assert data['contamination_portion'] == 0
    assert data['n_effective_surplus_clades'] == 0
    assert data['proportion_genes_retained_in_major_clades'] == 1
    assert round(data['mean_hit_identity'], 2) == 0.92
    assert round(data['genes_retained_index'], 2) == 0.49
    assert round(data['reference_representation_score'], 2) == 0.45
    assert data['pass.GUNC']

    # Test species level
    data = get_scores_for_taxlevel(ref_base_data,
                                   'species',
                                   0.34,
                                   'test',
                                   35,
                                   17,
                                   15,
                                   11)
    assert data['genome'] == 'test'
    assert data['n_contigs'] == 15
    assert data['n_genes_called'] == 35
    assert data['n_genes_mapped'] == 17
    assert data['taxonomic_level'] == 'species'
    assert round(data['clade_separation_score'], 2) == 1
    assert round(data['contamination_portion'], 2) == 0.35
    assert round(data['n_effective_surplus_clades'], 2) == 1.28
    assert data['proportion_genes_retained_in_major_clades'] == 1
    assert round(data['mean_hit_identity'], 2) == 0.92
    assert round(data['genes_retained_index'], 2) == 0.49
    assert round(data['reference_representation_score'], 2) == 0.45
    assert data['pass.GUNC'] is False

    # Test case where nothing is left after minor clade filtering
    data = get_scores_for_taxlevel(ref_base_data,
                                   'class',
                                   1000,
                                   'test',
                                   35,
                                   17,
                                   15,
                                   11)
    assert data['genome'] == 'test'
    assert data['n_contigs'] == 15
    assert data['n_genes_called'] == 35
    assert data['n_genes_mapped'] == 17
    assert data['taxonomic_level'] == 'class'
    np.testing.assert_equal(data['clade_separation_score'], np.nan)
    np.testing.assert_equal(data['contamination_portion'], np.nan)
    np.testing.assert_equal(data['n_effective_surplus_clades'], np.nan)
    np.testing.assert_equal(data['proportion_genes_retained_in_major_clades'], np.nan)
    np.testing.assert_equal(data['mean_hit_identity'], np.nan)
    np.testing.assert_equal(data['genes_retained_index'], np.nan)
    np.testing.assert_equal(data['reference_representation_score'], np.nan)
    np.testing.assert_equal(data['pass.GUNC'], np.nan)


def test_chim_score():
    diamond_file_path = resource_filename(__name__,
                                          'test_data/test_genome.fa.diamond.out')
    data, _ = chim_score(diamond_file_path, genes_called=1832)

    expected_data_path = resource_filename(__name__,
                                           'test_data/test_genome.fa.diamond.out.chimerism_scores')
    expected_data = pd.read_csv(expected_data_path, sep='\t')
    pd.testing.assert_frame_equal(data, expected_data)
