from itertools import product
from unittest.mock import patch
import pytest

import numpy as np
import pandas as pd
from pandas.util.testing import assert_frame_equal

from sm.engine.annotation.fdr import FDR, run_fdr_ranking
from sm.engine.formula_parser import format_modifiers

FDR_CONFIG = {'decoy_sample_size': 2}


@patch('sm.engine.annotation.fdr.DECOY_ADDUCTS', ['+He', '+Li'])
def test_fdr_decoy_adduct_selection_saves_corr():
    fdr = FDR(
        fdr_config=FDR_CONFIG,
        chem_mods=[],
        neutral_losses=[],
        target_adducts=['+H', '+K', '[M]+'],
        analysis_version=1,
    )

    exp_target_decoy_df = pd.DataFrame(
        [
            ('H2O', '+H', '+He'),
            ('H2O', '+H', '+Li'),
            ('H2O', '+K', '+He'),
            ('H2O', '+K', '+Li'),
            ('H2O', '', '+He'),
            ('H2O', '', '+Li'),
        ],
        columns=['formula', 'tm', 'dm'],
    )

    fdr.decoy_adducts_selection(target_formulas=['H2O'])

    assert_frame_equal(
        fdr.td_df.sort_values(by=['formula', 'tm', 'dm']).reset_index(drop=True),
        exp_target_decoy_df.sort_values(by=['formula', 'tm', 'dm']).reset_index(drop=True),
    )


@pytest.mark.parametrize('analysis_version,expected_fdrs', [(1, [0.2, 0.8]), (3, [1 / 4, 2 / 3])])
def test_estimate_fdr_returns_correct_df(analysis_version, expected_fdrs):
    fdr = FDR(
        fdr_config=FDR_CONFIG,
        chem_mods=[],
        neutral_losses=[],
        target_adducts=['+H'],
        analysis_version=analysis_version,
    )
    fdr.fdr_levels = [0.2, 0.8]
    fdr.td_df = pd.DataFrame(
        [['H2O', '+H', '+Cu'], ['H2O', '+H', '+Co'], ['C2H2', '+H', '+Ag'], ['C2H2', '+H', '+Ar']],
        columns=['formula', 'tm', 'dm'],
    )

    msm_df = pd.DataFrame(
        [
            ['H2O', '+H', 0.85],
            ['C2H2', '+H', 0.5],
            ['H2O', '+Cu', 0.5],
            ['H2O', '+Co', 0.5],
            ['C2H2', '+Ag', 0.75],
            ['C2H2', '+Ar', 0.0],
        ],
        columns=['formula', 'modifier', 'msm'],
    )
    exp_sf_df = pd.DataFrame(
        [
            ['H2O', '+H', 0.85],
            ['C2H2', '+H', 0.5],
        ],
        columns=['formula', 'modifier', 'msm'],
    ).assign(fdr=expected_fdrs)

    assert_frame_equal(fdr.estimate_fdr(msm_df, None), exp_sf_df)


def test_estimate_fdr_digitize_works():
    fdr_config = {**FDR_CONFIG, 'decoy_sample_size': 1}
    fdr = FDR(
        fdr_config=fdr_config,
        chem_mods=[],
        neutral_losses=[],
        target_adducts=['+H'],
        analysis_version=1,
    )
    fdr.fdr_levels = [0.4, 0.8]
    fdr.td_df = pd.DataFrame(
        [['C1', '+H', '+Cu'], ['C2', '+H', '+Ag'], ['C3', '+H', '+Cl'], ['C4', '+H', '+Co']],
        columns=['formula', 'tm', 'dm'],
    )

    msm_df = pd.DataFrame(
        [
            ['C1', '+H', 1.0],
            ['C2', '+H', 0.75],
            ['C3', '+H', 0.5],
            ['C4', '+H', 0.25],
            ['C1', '+Cu', 0.75],
            ['C2', '+Ag', 0.3],
            ['C3', '+Cl', 0.25],
            ['C4', '+Co', 0.1],
        ],
        columns=['formula', 'modifier', 'msm'],
    )
    exp_sf_df = pd.DataFrame(
        [
            ['C1', '+H', 1.0, 0.4],
            ['C2', '+H', 0.75, 0.4],
            ['C3', '+H', 0.5, 0.4],
            ['C4', '+H', 0.25, 0.8],
        ],
        columns=['formula', 'modifier', 'msm', 'fdr'],
    )

    assert_frame_equal(fdr.estimate_fdr(msm_df, None), exp_sf_df)


def test_ions():
    formulas = ['H2O', 'C5H2OH']
    target_adducts = ['+H', '+Na']
    decoy_sample_size = 5
    fdr_config = {**FDR_CONFIG, 'decoy_sample_size': decoy_sample_size}

    fdr = FDR(
        fdr_config=fdr_config,
        chem_mods=[],
        neutral_losses=[],
        target_adducts=target_adducts,
        analysis_version=1,
    )
    fdr.decoy_adducts_selection(target_formulas=['H2O', 'C5H2OH'])
    ions = fdr.ion_tuples()

    assert type(ions) == list
    # total number varies because different (formula, modifier) pairs may receive the same (formula, decoy_modifier) pair
    assert (
        len(formulas) * decoy_sample_size + len(formulas) * len(target_adducts)
        < len(ions)
        <= len(formulas) * len(target_adducts) * decoy_sample_size
        + len(formulas) * len(target_adducts)
    )
    target_ions = [(formula, adduct) for formula, adduct in product(formulas, target_adducts)]
    assert set(target_ions).issubset(set(map(tuple, ions)))


def test_chem_mods_and_neutral_losses():
    formulas = ['H2O', 'C5H2OH']
    chem_mods = ['-H+C']
    neutral_losses = ['-O', '-C']
    target_adducts = ['+H', '+Na', '[M]+']
    target_modifiers = [
        format_modifiers(cm, nl, ta)
        for cm, nl, ta in product(['', *chem_mods], ['', *neutral_losses], target_adducts)
    ]
    decoy_sample_size = 5
    fdr_config = {**FDR_CONFIG, 'decoy_sample_size': decoy_sample_size}

    fdr = FDR(
        fdr_config=fdr_config,
        chem_mods=chem_mods,
        neutral_losses=neutral_losses,
        target_adducts=target_adducts,
        analysis_version=1,
    )
    fdr.decoy_adducts_selection(target_formulas=['H2O', 'C5H2OH'])
    ions = fdr.ion_tuples()

    assert type(ions) == list
    # total number varies because different (formula, modifier) pairs may receive the same (formula, decoy_modifier) pair
    min_count = len(formulas) * len(target_modifiers)
    max_count = len(formulas) * len(target_modifiers) * (1 + decoy_sample_size)
    assert min_count < len(ions) <= max_count
    target_ions = list(product(formulas, target_modifiers))
    assert set(target_ions).issubset(set(map(tuple, ions)))


def test_run_fdr_ranking():
    target_scores = pd.Series([1.0, 0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1, 0.0])
    decoy_scores = pd.Series([0.8, 0.55, 0.2, 0.1])
    n_targets = pd.Series([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11])
    n_decoys = pd.Series([0, 0, 1, 1, 1, 2, 2, 2, 3, 4, 4])
    expected_fdr = n_decoys / n_targets
    expected_fdr_ros = (n_decoys + 1) / (n_targets + 1)
    expected_fdr_mono = pd.Series(
        [0 / 2, 0 / 2, 1 / 5, 1 / 5, 1 / 5, 2 / 8, 2 / 8, 2 / 8, 3 / 9, 4 / 11, 4 / 11]
    )

    fdr = run_fdr_ranking(target_scores, decoy_scores, 1, False, False)
    fdr_ros = run_fdr_ranking(target_scores, decoy_scores, 1, True, False)
    fdr_mono = run_fdr_ranking(target_scores, decoy_scores, 1, False, True)

    assert np.isclose(fdr, expected_fdr).all()
    assert np.isclose(fdr_ros, expected_fdr_ros).all()
    assert np.isclose(fdr_mono, expected_fdr_mono).all()

[(15, 'unittest.mock.patch', 'patch', (['"""sm.engine.annotation.fdr.DECOY_ADDUCTS"""', "['+He', '+Li']"], {}), False, 'from unittest.mock import patch\n'), (45, 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""analysis_version,expected_fdrs"""', '[(1, [0.2, 0.8]), (3, [1 / 4, 2 / 3])]'], {}), False, 'import pytest\n'), (17, 'sm.engine.annotation.fdr.FDR', 'FDR', ([], {'fdr_config': 'FDR_CONFIG', 'chem_mods': '[]', 'neutral_losses': '[]', 'target_adducts': "['+H', '+K', '[M]+']", 'analysis_version': '(1)'}), False, 'from sm.engine.annotation.fdr import FDR, run_fdr_ranking\n'), (25, 'pandas.DataFrame', 'pd.DataFrame', (["[('H2O', '+H', '+He'), ('H2O', '+H', '+Li'), ('H2O', '+K', '+He'), ('H2O',\n    '+K', '+Li'), ('H2O', '', '+He'), ('H2O', '', '+Li')]"], {'columns': "['formula', 'tm', 'dm']"}), True, 'import pandas as pd\n'), (47, 'sm.engine.annotation.fdr.FDR', 'FDR', ([], {'fdr_config': 'FDR_CONFIG', 'chem_mods': '[]', 'neutral_losses': '[]', 'target_adducts': "['+H']", 'analysis_version': 'analysis_version'}), False, 'from sm.engine.annotation.fdr import FDR, run_fdr_ranking\n'), (55, 'pandas.DataFrame', 'pd.DataFrame', (["[['H2O', '+H', '+Cu'], ['H2O', '+H', '+Co'], ['C2H2', '+H', '+Ag'], ['C2H2',\n    '+H', '+Ar']]"], {'columns': "['formula', 'tm', 'dm']"}), True, 'import pandas as pd\n'), (60, 'pandas.DataFrame', 'pd.DataFrame', (["[['H2O', '+H', 0.85], ['C2H2', '+H', 0.5], ['H2O', '+Cu', 0.5], ['H2O',\n    '+Co', 0.5], ['C2H2', '+Ag', 0.75], ['C2H2', '+Ar', 0.0]]"], {'columns': "['formula', 'modifier', 'msm']"}), True, 'import pandas as pd\n'), (84, 'sm.engine.annotation.fdr.FDR', 'FDR', ([], {'fdr_config': 'fdr_config', 'chem_mods': '[]', 'neutral_losses': '[]', 'target_adducts': "['+H']", 'analysis_version': '(1)'}), False, 'from sm.engine.annotation.fdr import FDR, run_fdr_ranking\n'), (92, 'pandas.DataFrame', 'pd.DataFrame', (["[['C1', '+H', '+Cu'], ['C2', '+H', '+Ag'], ['C3', '+H', '+Cl'], ['C4', '+H',\n    '+Co']]"], {'columns': "['formula', 'tm', 'dm']"}), True, 'import pandas as pd\n'), (97, 'pandas.DataFrame', 'pd.DataFrame', (["[['C1', '+H', 1.0], ['C2', '+H', 0.75], ['C3', '+H', 0.5], ['C4', '+H', \n    0.25], ['C1', '+Cu', 0.75], ['C2', '+Ag', 0.3], ['C3', '+Cl', 0.25], [\n    'C4', '+Co', 0.1]]"], {'columns': "['formula', 'modifier', 'msm']"}), True, 'import pandas as pd\n'), (110, 'pandas.DataFrame', 'pd.DataFrame', (["[['C1', '+H', 1.0, 0.4], ['C2', '+H', 0.75, 0.4], ['C3', '+H', 0.5, 0.4], [\n    'C4', '+H', 0.25, 0.8]]"], {'columns': "['formula', 'modifier', 'msm', 'fdr']"}), True, 'import pandas as pd\n'), (129, 'sm.engine.annotation.fdr.FDR', 'FDR', ([], {'fdr_config': 'fdr_config', 'chem_mods': '[]', 'neutral_losses': '[]', 'target_adducts': 'target_adducts', 'analysis_version': '(1)'}), False, 'from sm.engine.annotation.fdr import FDR, run_fdr_ranking\n'), (163, 'sm.engine.annotation.fdr.FDR', 'FDR', ([], {'fdr_config': 'fdr_config', 'chem_mods': 'chem_mods', 'neutral_losses': 'neutral_losses', 'target_adducts': 'target_adducts', 'analysis_version': '(1)'}), False, 'from sm.engine.annotation.fdr import FDR, run_fdr_ranking\n'), (183, 'pandas.Series', 'pd.Series', (['[1.0, 0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1, 0.0]'], {}), True, 'import pandas as pd\n'), (184, 'pandas.Series', 'pd.Series', (['[0.8, 0.55, 0.2, 0.1]'], {}), True, 'import pandas as pd\n'), (185, 'pandas.Series', 'pd.Series', (['[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]'], {}), True, 'import pandas as pd\n'), (186, 'pandas.Series', 'pd.Series', (['[0, 0, 1, 1, 1, 2, 2, 2, 3, 4, 4]'], {}), True, 'import pandas as pd\n'), (189, 'pandas.Series', 'pd.Series', (['[0 / 2, 0 / 2, 1 / 5, 1 / 5, 1 / 5, 2 / 8, 2 / 8, 2 / 8, 3 / 9, 4 / 11, 4 / 11]'], {}), True, 'import pandas as pd\n'), (193, 'sm.engine.annotation.fdr.run_fdr_ranking', 'run_fdr_ranking', (['target_scores', 'decoy_scores', '(1)', '(False)', '(False)'], {}), False, 'from sm.engine.annotation.fdr import FDR, run_fdr_ranking\n'), (194, 'sm.engine.annotation.fdr.run_fdr_ranking', 'run_fdr_ranking', (['target_scores', 'decoy_scores', '(1)', '(True)', '(False)'], {}), False, 'from sm.engine.annotation.fdr import FDR, run_fdr_ranking\n'), (195, 'sm.engine.annotation.fdr.run_fdr_ranking', 'run_fdr_ranking', (['target_scores', 'decoy_scores', '(1)', '(False)', '(True)'], {}), False, 'from sm.engine.annotation.fdr import FDR, run_fdr_ranking\n'), (157, 'sm.engine.formula_parser.format_modifiers', 'format_modifiers', (['cm', 'nl', 'ta'], {}), False, 'from sm.engine.formula_parser import format_modifiers\n'), (178, 'itertools.product', 'product', (['formulas', 'target_modifiers'], {}), False, 'from itertools import product\n'), (71, 'pandas.DataFrame', 'pd.DataFrame', (["[['H2O', '+H', 0.85], ['C2H2', '+H', 0.5]]"], {'columns': "['formula', 'modifier', 'msm']"}), True, 'import pandas as pd\n'), (147, 'itertools.product', 'product', (['formulas', 'target_adducts'], {}), False, 'from itertools import product\n'), (158, 'itertools.product', 'product', (["['', *chem_mods]", "['', *neutral_losses]", 'target_adducts'], {}), False, 'from itertools import product\n'), (197, 'numpy.isclose', 'np.isclose', (['fdr', 'expected_fdr'], {}), True, 'import numpy as np\n'), (198, 'numpy.isclose', 'np.isclose', (['fdr_ros', 'expected_fdr_ros'], {}), True, 'import numpy as np\n'), (199, 'numpy.isclose', 'np.isclose', (['fdr_mono', 'expected_fdr_mono'], {}), True, 'import numpy as np\n')]