import pandas as pd
import shutil
import os
import io

from ms_mint.Mint import Mint

from pathlib import Path as P

from ms_mint.io import (
    ms_file_to_df,
    mzml_to_pandas_df_pyteomics,
    convert_ms_file_to_feather,
    convert_ms_file_to_parquet,
    MZMLB_AVAILABLE,
)

from paths import (
    TEST_MZML,
    TEST_MZXML,
    TEST_PARQUET,
    TEST_MZMLB_POS,
    TEST_MZML_POS,
    TEST_MZML_NEG,
)


def test__ms_file_to_df__mzML():
    result = ms_file_to_df(TEST_MZML)
    expected_cols = [
        "scan_id",
        "ms_level",
        "polarity",
        "scan_time_min",
        "mz",
        "intensity",
    ]
    assert isinstance(result, pd.DataFrame), f"{type(result)} is not a dataframe"
    assert expected_cols == result.columns.to_list(), result.columns


def test__ms_file_to_df__mzML_timeunit_minutes():
    result = ms_file_to_df(TEST_MZML, time_unit="minutes")
    expected_cols = [
        "scan_id",
        "ms_level",
        "polarity",
        "scan_time_min",
        "mz",
        "intensity",
    ]
    assert isinstance(result, pd.DataFrame), f"{type(result)} is not a dataframe"
    assert expected_cols == result.columns.to_list(), result.columns


def test__ms_file_to_df__mzXML():
    result = ms_file_to_df(TEST_MZXML)
    expected_cols = [
        "scan_id",
        "ms_level",
        "polarity",
        "scan_time_min",
        "mz",
        "intensity",
    ]
    assert isinstance(result, pd.DataFrame), f"{type(result)} is not a dataframe"
    assert expected_cols == result.columns.to_list(), result.columns


def test__mzml_to_pandas_df_pyteomics_pos():
    result = mzml_to_pandas_df_pyteomics(TEST_MZML_POS)
    expected_cols = [
        "scan_id",
        "ms_level",
        "polarity",
        "scan_time_min",
        "mz",
        "intensity",
    ]
    assert isinstance(result, pd.DataFrame), f"{type(result)} is not a dataframe"
    assert expected_cols == result.columns.to_list(), result.columns
    assert all(result.polarity == "+"), f'Polarity should be "+"\n{result}'


def test__mzml_to_pandas_df_pyteomics_neg():
    result = mzml_to_pandas_df_pyteomics(TEST_MZML_NEG)
    expected_cols = [
        "scan_id",
        "ms_level",
        "polarity",
        "scan_time_min",
        "mz",
        "intensity",
    ]
    assert isinstance(result, pd.DataFrame), f"{type(result)} is not a dataframe"
    assert expected_cols == result.columns.to_list(), result.columns
    assert all(result.polarity == "-"), f'Polarity should be "-"\n{result}'


def test__read_parquet():
    result = ms_file_to_df(TEST_PARQUET)
    expected_cols = [
        "scan_id",
        "ms_level",
        "polarity",
        "scan_time_min",
        "mz",
        "intensity",
    ]
    assert isinstance(result, pd.DataFrame), f"{type(result)} is not a dataframe"
    assert expected_cols == result.columns.to_list(), result.columns


def test__write_read_hdf(tmpdir):
    df = ms_file_to_df(TEST_PARQUET)
    fn = P(tmpdir) / "file.hdf"
    df.to_hdf(fn, key="data")
    result = ms_file_to_df(fn)
    expected_cols = [
        "scan_id",
        "ms_level",
        "polarity",
        "scan_time_min",
        "mz",
        "intensity",
    ]
    assert isinstance(result, pd.DataFrame), f"{type(result)} is not a dataframe"
    assert expected_cols == result.columns.to_list(), result.columns


def test__read_mzMLb(tmpdir):
    if not MZMLB_AVAILABLE:
        return None
    result = ms_file_to_df(TEST_MZMLB_POS)
    expected_cols = [
        "scan_id",
        "ms_level",
        "polarity",
        "scan_time_min",
        "mz",
        "intensity",
    ]
    assert isinstance(result, pd.DataFrame), f"{type(result)} is not a dataframe"
    assert expected_cols == result.columns.to_list(), result.columns
    # assert all(result.polarity == '+'), f'Polarity should be "+"\n{result}'


def test__convert_ms_file_to_feather(tmpdir):
    print(tmpdir)
    shutil.copy(TEST_MZML, tmpdir)
    fn = P(tmpdir) / P(TEST_MZML).name
    fn_out = fn.with_suffix(".feather")
    print(fn, fn_out)
    convert_ms_file_to_feather(fn)
    assert fn_out.is_file(), f"File not generated {fn_out}"
    df = ms_file_to_df(fn)
    df_fea = ms_file_to_df(fn_out)
    assert df_fea.equals(df), "DataFrames not equal"


def test__convert_ms_file_to_parquet(tmpdir):
    print(tmpdir)
    shutil.copy(TEST_MZML, tmpdir)
    fn = P(tmpdir) / P(TEST_MZML).name
    fn_out = fn.with_suffix(".parquet")
    print(fn, fn_out)
    convert_ms_file_to_parquet(fn)
    assert fn_out.is_file(), f"File not generated {fn_out}"
    df = ms_file_to_df(fn)
    df_fea = ms_file_to_df(fn_out)
    assert df_fea.equals(df), "DataFrames not equal"


def test__export_to_excel(tmp_path):
    filename = os.path.join(tmp_path, "output.xlsx")
    mint = Mint(verbose=True)
    mint.ms_files = "tests/data/test.mzXML"
    mint.run()
    mint.export(filename)
    assert os.path.isfile(filename)


def test__export_to_excel_without_fn():
    mint = Mint(verbose=True)
    mint.ms_files = TEST_MZXML
    mint.targets = pd.DataFrame(
        {
            "peak_label": ["A"],
            "mz_mean": [200],
            "mz_width": [10],
            "intensity_threshold": [0],
            "rt_min": [0],
            "rt_max": [10],
            "targets_filename": ["unknown"],
        }
    )
    mint.run()
    buffer = mint.export()
    assert isinstance(buffer, io.BytesIO)
    df = pd.read_excel(buffer, sheet_name="Results")
    assert len(df) == 1, len(df)
    assert df.loc[0, "peak_label"] == "A", df.loc[0, "peak_label"]
    assert df.loc[0, "ms_file"] == P(TEST_MZXML).name, df.loc[0, "ms_file"]

[((459, 483), 'ms_mint.io.ms_file_to_df', 'ms_file_to_df', (['TEST_MZML'], {}), '(TEST_MZML)\n', (472, 483), False, 'from ms_mint.io import ms_file_to_df, mzml_to_pandas_df_pyteomics, convert_ms_file_to_feather, convert_ms_file_to_parquet, MZMLB_AVAILABLE\n'), ((847, 892), 'ms_mint.io.ms_file_to_df', 'ms_file_to_df', (['TEST_MZML'], {'time_unit': '"""minutes"""'}), "(TEST_MZML, time_unit='minutes')\n", (860, 892), False, 'from ms_mint.io import ms_file_to_df, mzml_to_pandas_df_pyteomics, convert_ms_file_to_feather, convert_ms_file_to_parquet, MZMLB_AVAILABLE\n'), ((1240, 1265), 'ms_mint.io.ms_file_to_df', 'ms_file_to_df', (['TEST_MZXML'], {}), '(TEST_MZXML)\n', (1253, 1265), False, 'from ms_mint.io import ms_file_to_df, mzml_to_pandas_df_pyteomics, convert_ms_file_to_feather, convert_ms_file_to_parquet, MZMLB_AVAILABLE\n'), ((1624, 1666), 'ms_mint.io.mzml_to_pandas_df_pyteomics', 'mzml_to_pandas_df_pyteomics', (['TEST_MZML_POS'], {}), '(TEST_MZML_POS)\n', (1651, 1666), False, 'from ms_mint.io import ms_file_to_df, mzml_to_pandas_df_pyteomics, convert_ms_file_to_feather, convert_ms_file_to_parquet, MZMLB_AVAILABLE\n'), ((2101, 2143), 'ms_mint.io.mzml_to_pandas_df_pyteomics', 'mzml_to_pandas_df_pyteomics', (['TEST_MZML_NEG'], {}), '(TEST_MZML_NEG)\n', (2128, 2143), False, 'from ms_mint.io import ms_file_to_df, mzml_to_pandas_df_pyteomics, convert_ms_file_to_feather, convert_ms_file_to_parquet, MZMLB_AVAILABLE\n'), ((2559, 2586), 'ms_mint.io.ms_file_to_df', 'ms_file_to_df', (['TEST_PARQUET'], {}), '(TEST_PARQUET)\n', (2572, 2586), False, 'from ms_mint.io import ms_file_to_df, mzml_to_pandas_df_pyteomics, convert_ms_file_to_feather, convert_ms_file_to_parquet, MZMLB_AVAILABLE\n'), ((2930, 2957), 'ms_mint.io.ms_file_to_df', 'ms_file_to_df', (['TEST_PARQUET'], {}), '(TEST_PARQUET)\n', (2943, 2957), False, 'from ms_mint.io import ms_file_to_df, mzml_to_pandas_df_pyteomics, convert_ms_file_to_feather, convert_ms_file_to_parquet, MZMLB_AVAILABLE\n'), ((3033, 3050), 'ms_mint.io.ms_file_to_df', 'ms_file_to_df', (['fn'], {}), '(fn)\n', (3046, 3050), False, 'from ms_mint.io import ms_file_to_df, mzml_to_pandas_df_pyteomics, convert_ms_file_to_feather, convert_ms_file_to_parquet, MZMLB_AVAILABLE\n'), ((3442, 3471), 'ms_mint.io.ms_file_to_df', 'ms_file_to_df', (['TEST_MZMLB_POS'], {}), '(TEST_MZMLB_POS)\n', (3455, 3471), False, 'from ms_mint.io import ms_file_to_df, mzml_to_pandas_df_pyteomics, convert_ms_file_to_feather, convert_ms_file_to_parquet, MZMLB_AVAILABLE\n'), ((3918, 3948), 'shutil.copy', 'shutil.copy', (['TEST_MZML', 'tmpdir'], {}), '(TEST_MZML, tmpdir)\n', (3929, 3948), False, 'import shutil\n'), ((4054, 4084), 'ms_mint.io.convert_ms_file_to_feather', 'convert_ms_file_to_feather', (['fn'], {}), '(fn)\n', (4080, 4084), False, 'from ms_mint.io import ms_file_to_df, mzml_to_pandas_df_pyteomics, convert_ms_file_to_feather, convert_ms_file_to_parquet, MZMLB_AVAILABLE\n'), ((4154, 4171), 'ms_mint.io.ms_file_to_df', 'ms_file_to_df', (['fn'], {}), '(fn)\n', (4167, 4171), False, 'from ms_mint.io import ms_file_to_df, mzml_to_pandas_df_pyteomics, convert_ms_file_to_feather, convert_ms_file_to_parquet, MZMLB_AVAILABLE\n'), ((4185, 4206), 'ms_mint.io.ms_file_to_df', 'ms_file_to_df', (['fn_out'], {}), '(fn_out)\n', (4198, 4206), False, 'from ms_mint.io import ms_file_to_df, mzml_to_pandas_df_pyteomics, convert_ms_file_to_feather, convert_ms_file_to_parquet, MZMLB_AVAILABLE\n'), ((4330, 4360), 'shutil.copy', 'shutil.copy', (['TEST_MZML', 'tmpdir'], {}), '(TEST_MZML, tmpdir)\n', (4341, 4360), False, 'import shutil\n'), ((4466, 4496), 'ms_mint.io.convert_ms_file_to_parquet', 'convert_ms_file_to_parquet', (['fn'], {}), '(fn)\n', (4492, 4496), False, 'from ms_mint.io import ms_file_to_df, mzml_to_pandas_df_pyteomics, convert_ms_file_to_feather, convert_ms_file_to_parquet, MZMLB_AVAILABLE\n'), ((4566, 4583), 'ms_mint.io.ms_file_to_df', 'ms_file_to_df', (['fn'], {}), '(fn)\n', (4579, 4583), False, 'from ms_mint.io import ms_file_to_df, mzml_to_pandas_df_pyteomics, convert_ms_file_to_feather, convert_ms_file_to_parquet, MZMLB_AVAILABLE\n'), ((4597, 4618), 'ms_mint.io.ms_file_to_df', 'ms_file_to_df', (['fn_out'], {}), '(fn_out)\n', (4610, 4618), False, 'from ms_mint.io import ms_file_to_df, mzml_to_pandas_df_pyteomics, convert_ms_file_to_feather, convert_ms_file_to_parquet, MZMLB_AVAILABLE\n'), ((4726, 4763), 'os.path.join', 'os.path.join', (['tmp_path', '"""output.xlsx"""'], {}), "(tmp_path, 'output.xlsx')\n", (4738, 4763), False, 'import os\n'), ((4775, 4793), 'ms_mint.Mint.Mint', 'Mint', ([], {'verbose': '(True)'}), '(verbose=True)\n', (4779, 4793), False, 'from ms_mint.Mint import Mint\n'), ((4890, 4914), 'os.path.isfile', 'os.path.isfile', (['filename'], {}), '(filename)\n', (4904, 4914), False, 'import os\n'), ((4968, 4986), 'ms_mint.Mint.Mint', 'Mint', ([], {'verbose': '(True)'}), '(verbose=True)\n', (4972, 4986), False, 'from ms_mint.Mint import Mint\n'), ((5037, 5208), 'pandas.DataFrame', 'pd.DataFrame', (["{'peak_label': ['A'], 'mz_mean': [200], 'mz_width': [10],\n    'intensity_threshold': [0], 'rt_min': [0], 'rt_max': [10],\n    'targets_filename': ['unknown']}"], {}), "({'peak_label': ['A'], 'mz_mean': [200], 'mz_width': [10],\n    'intensity_threshold': [0], 'rt_min': [0], 'rt_max': [10],\n    'targets_filename': ['unknown']})\n", (5049, 5208), True, 'import pandas as pd\n'), ((5403, 5446), 'pandas.read_excel', 'pd.read_excel', (['buffer'], {'sheet_name': '"""Results"""'}), "(buffer, sheet_name='Results')\n", (5416, 5446), True, 'import pandas as pd\n'), ((2967, 2976), 'pathlib.Path', 'P', (['tmpdir'], {}), '(tmpdir)\n', (2968, 2976), True, 'from pathlib import Path as P\n'), ((3958, 3967), 'pathlib.Path', 'P', (['tmpdir'], {}), '(tmpdir)\n', (3959, 3967), True, 'from pathlib import Path as P\n'), ((4370, 4379), 'pathlib.Path', 'P', (['tmpdir'], {}), '(tmpdir)\n', (4371, 4379), True, 'from pathlib import Path as P\n'), ((3970, 3982), 'pathlib.Path', 'P', (['TEST_MZML'], {}), '(TEST_MZML)\n', (3971, 3982), True, 'from pathlib import Path as P\n'), ((4382, 4394), 'pathlib.Path', 'P', (['TEST_MZML'], {}), '(TEST_MZML)\n', (4383, 4394), True, 'from pathlib import Path as P\n'), ((5582, 5595), 'pathlib.Path', 'P', (['TEST_MZXML'], {}), '(TEST_MZXML)\n', (5583, 5595), True, 'from pathlib import Path as P\n')]