# -*- coding: utf-8 -*-
"""
Test that :class:`~radis.tools.database.SpecDatabase` works



"""

import os
from os.path import dirname, exists

import pytest

from radis.spectrum import Spectrum
from radis.test.utils import getTestFile
from radis.tools.database import SpecDatabase, SpecList


@pytest.mark.fast
def test_speclist(*args, **kwargs):
    """Test :py:class:`~radis.tools.database.SpecList` functions"""

    from radis.test.utils import getTestFile

    # Get a list of test files
    spec_list = SpecDatabase(getTestFile("."), lazy_loading=False).get()

    # Now generate a SpecList
    dbl = SpecList(*spec_list)
    # check methods work:
    s_closest = dbl.get_closest(Trot=1550)  # TODO add which one
    assert s_closest.c["Trot"] == 1500

    # test get_closest() with a non numeric argument
    s_closest_2 = dbl.get_closest(Trot=1550, molecule="CO")
    assert s_closest_2.c["Trot"] == 1500


@pytest.mark.fast
def test_database_functions(
    verbose=True, plot=True, close_plots=True, warnings=True, *args, **kwargs
):
    """Test :py:class:`~radis.tools.database.SpecDatabase` functions"""

    if plot:
        import matplotlib.pyplot as plt

        plt.ion()
    if close_plots:
        plt.close("all")

    db = SpecDatabase(dirname(getTestFile(".")), lazy_loading=False)

    # Database visualisation methods
    if verbose:
        print(f"{len(db)} items in test database: {db.see(['Tvib', 'Trot'])}")
    if plot:
        db.plot_cond("Tvib", "Trot")

    # Database get methods
    db.get_closest(Tgas=1300, path_length=1)
    s = db.get_unique(
        Tgas=1500, path_length=0.01, mole_fraction=0.5
    )  # there should be one only
    # ... note that this is just so we can test
    # ... get_unique(). If we were to add new
    # ... test cases with matching conditions
    # ... let's add more criteria to keep it unique
    match = db.get(**s.conditions)
    assert len(match) == 1

    # Database add method
    s2 = s.copy()
    s2.conditions["Tgas"] = 0  # make it unique (for testing)

    matching_spectrum_in_db = s2 in db

    l = db.add(
        s2,
        add_info=["Tvib", "Trot"],
        discard=[],
        compress=True,
        if_exists_then="increment",
    )
    assert exists(l)
    try:
        assert s2 in db
        # .. ensures that you cant add it twice
        with pytest.raises(ValueError):
            db.add(s2, add_info=["Tvib", "Trot"], if_exists_then="error")

    # Now remove the Spectrum, update the database and ensures it's not there anymore
    finally:
        os.remove(l)
    db.update(force_reload=True)  # update database
    assert not exists(l)

    # make sure we deleted it properly
    if not matching_spectrum_in_db:
        assert s2 not in db


def test_plot_spec(plot=True, close_plots=True, verbose=True, *args, **kwargs):

    from radis.test.utils import getTestFile
    from radis.tools.database import plot_spec

    if plot:
        import matplotlib.pyplot as plt

        plt.ion()
    if close_plots:
        plt.close("all")

    if plot:
        plot_spec(getTestFile("N2C_specair_380nm.spec"))


def test_save_compressed2(verbose=True, *args, **kwargs):
    "Check if saving a spectrum with compress = 2 does not change something."
    import shutil
    from os.path import join

    from radis import SpecDatabase, calc_spectrum
    from radis.test.utils import setup_test_line_databases

    shutil.rmtree(join(dirname(getTestFile(".")), "newDb"), ignore_errors=True)

    # get the spectrum
    setup_test_line_databases()
    s = calc_spectrum(
        2000,
        2300,  # cm-1
        molecule="CO",
        isotope="1,2,3",
        pressure=1.01325,  # bar
        Tgas=700,  # K
        mole_fraction=0.1,
        path_length=1,  # cm
        verbose=False,
        wstep=0.01,
        medium="vacuum",
        databank="HITRAN-CO-TEST",
    )
    try:

        # load in one database
        db = SpecDatabase(join(dirname(getTestFile(".")), "newDb"))
        db.add(s, compress=2, if_exists_then="error")

        # simulate an experimentalist who come later and load the spectrum
        db2 = SpecDatabase(join(dirname(getTestFile(".")), "newDb"))
        s_bis = db2.get_unique(Tgas=700)

    finally:
        # we want to make sure this folder is deleted
        shutil.rmtree(join(dirname(getTestFile(".")), "newDb"))

    # we check the loaded spectrum contains less information than the calculated one
    assert not s == s_bis
    assert s_bis.get_vars() == ["abscoeff"]  # only this spectral quantity was stored
    assert s_bis.lines is None
    assert s_bis.conditions is not None  # we kept the metadata
    # now we check if it works
    s_bis.update()
    for var in s.get_vars():
        assert s.compare_with(s_bis, spectra_only=var, plot=False, verbose=verbose)


def test_lazy_loading(verbose=True, *args, **kwargs):
    """Test lazy loading"""

    import numpy as np

    from radis import calc_spectrum

    DBNAME = "test_CO_OH_database"
    db = SpecDatabase(DBNAME, lazy_loading=False)  # create the database

    # %% Simulate a fake database

    # Compute spectra of CO and OH and add them to the database
    T_db = np.linspace(300, 1200, 5)  # Temperature array
    for T in T_db:
        s1 = (
            calc_spectrum(
                wavenum_min=3000,
                wavenum_max=4500,
                molecule="CO",
                Tgas=T,
                databank="hitran",  # test by fetching directly
                verbose=False,
            )
            .apply_slit(2, "nm")
            .take("radiance")
            # TODO : use https://github.com/radis/radis/issues/135 once implemented
        )
        db.add(s1, store_name=f"CO_{T}K.spec", if_exists_then="ignore")
        s2 = (
            calc_spectrum(
                wavenum_min=3000,
                wavenum_max=4500,
                molecule="OH",
                Tgas=T,
                databank="hitran",  # test by fetching directly
                verbose=False,
            )
            .apply_slit(2, "nm")
            .take("radiance")
            # TODO : use https://github.com/radis/radis/issues/135 once implemented
        )
        db.add(s2, store_name=f"OH_{T}K.spec", if_exists_then="ignore")

    # %% Access the spectra in the database via the get functions.
    # Check what species are in the database at a given temperature (300)
    s300 = db.get(Tgas=300)  # simple use of get
    if verbose:
        print([s.conditions["species"] for s in s300])
    # Check what species are in the database at a given species
    sCO = db.get(species="CO")
    if verbose:
        print([s.conditions["Tgas"] for s in sCO])
    # Advanced use of the get function
    sdb = db.get('Tgas>=600 & Tgas<900 & species=="CO"')
    if verbose:
        print("Number of spectra found for the given conditions: ", len(sdb))

    # %% Do the same in lazy-loading mode, and compare :

    db2 = SpecDatabase(DBNAME, lazy_loading=True)  # create the database

    # %% Access the spectra in the database via the get functions.
    # Check what species are in the database at a given temperature (300)
    s300_2 = db2.get(Tgas=300)  # simple use of get
    print([s.conditions["species"] for s in s300_2])

    assert s300 == s300_2

    # Check what species are in the database at a given species
    sCO_2 = db2.get(species="CO")
    if verbose:
        print([s.conditions["Tgas"] for s in sCO_2])

    assert sCO == sCO_2

    # Advanced use of the get function
    sdb_2 = db2.get('Tgas>=600 & Tgas<900 & species=="CO"')
    if verbose:
        print("Number of spectra found for the given conditions: ", len(sdb))

    assert sdb == sdb_2


def test_fit_spectrum_plotting(plot=True, close_plots=True, *args, **kwargs):
    """Test SpecDatabase.fit_spectrum plotting functionality (Issue #377)"""

    import shutil
    from os.path import dirname, join

    db_path = join(dirname(getTestFile(".")), "test_db_fit")
    if os.path.exists(db_path):
        shutil.rmtree(db_path)
    os.makedirs(db_path)

    try:
        db = SpecDatabase(db_path, verbose=False)

        # 1. Create dummy spectra for 1D fit
        for T in [300, 400, 500]:
            s = Spectrum(
                {"radiance": ([1000, 1001], [T / 1000, T / 1000])},
                wunit="cm-1",
                units={"radiance": "W/cm2/sr/cm-1"},
                conditions={"Tgas": T, "molecule": "CO2", "isotope": 1},
            )
            db.add(s)

        s_exp = Spectrum(
            {"radiance": ([1000, 1001], [0.42, 0.42])},
            wunit="cm-1",
            units={"radiance": "W/cm2/sr/cm-1"},
            conditions={"Tgas": 420, "molecule": "CO2", "isotope": 1},
        )

        # Test 1D plot
        db.fit_spectrum(s_exp, plot=plot)
        if plot:
            import matplotlib.pyplot as plt

            assert plt.fignum_exists(plt.gcf().number)
            if close_plots:
                plt.close()

        # 2. Add more spectra for 2D fit
        for T in [300, 500]:
            for mf in [0.01, 0.1]:
                s = Spectrum(
                    {"radiance": ([1000, 1001], [T / 1000 + mf, T / 1000 + mf])},
                    wunit="cm-1",
                    units={"radiance": "W/cm2/sr/cm-1"},
                    conditions={
                        "Tgas": T,
                        "mole_fraction": mf,
                        "molecule": "CO2",
                        "isotope": 1,
                    },
                )
                db.add(s, if_exists_then="ignore")

        s_exp_2d = Spectrum(
            {"radiance": ([1000, 1001], [0.45, 0.45])},
            wunit="cm-1",
            units={"radiance": "W/cm2/sr/cm-1"},
            conditions={
                "Tgas": 450,
                "mole_fraction": 0.05,
                "molecule": "CO2",
                "isotope": 1,
            },
        )

        # Test 2D plot
        db.fit_spectrum(s_exp_2d, plot=plot)
        if plot:
            import matplotlib.pyplot as plt

            assert plt.fignum_exists(plt.gcf().number)
            if close_plots:
                plt.close()

        # 3. Test 'var' parameter and unhashable type handling
        s_multi = Spectrum(
            {
                "radiance": ([1000, 1001], [0.5, 0.5]),
                "transmittance": ([1000, 1001], [0.8, 0.8]),
            },
            wunit="cm-1",
            units={"radiance": "W/cm2/sr/cm-1", "transmittance": ""},
            conditions={
                "Tgas": 500,
                "molecule": "CO2",
                "isotope": 1,
                "complex_metadata": {"a": 1},  # Unhashable dict
            },
        )
        db.add(s_multi, if_exists_then="ignore")

        # Should work with explicit 'var' even with complex_metadata in the DB
        db.fit_spectrum(s_multi, var="radiance", plot=plot)

    finally:
        shutil.rmtree(db_path)


@pytest.mark.fast
def test_store_to_bytesio(verbose=True, *args, **kwargs):
    """Test storing a spectrum to a BytesIO buffer."""
    from io import BytesIO

    import json_tricks

    from radis.test.utils import getTestFile
    from radis.tools.database import load_spec

    s = load_spec(getTestFile("N2C_specair_380nm.spec"))

    buffer = BytesIO()
    result = s.store(buffer, compress=True, verbose=verbose)

    assert result is buffer
    assert buffer.tell() > 0

    buffer.seek(0)
    sload = json_tricks.load(buffer, preserve_order=False)
    assert "conditions" in sload
    assert "_q" in sload


@pytest.mark.fast
def test_store_to_stringio(verbose=True, *args, **kwargs):
    """Test storing a spectrum to a StringIO buffer."""
    from io import StringIO

    from radis.test.utils import getTestFile
    from radis.tools.database import load_spec

    s = load_spec(getTestFile("N2C_specair_380nm.spec"))

    buffer = StringIO()
    result = s.store(buffer, compress=False, verbose=verbose)

    assert result is buffer

    buffer.seek(0)
    content = buffer.read()
    assert len(content) > 0
    assert '"conditions"' in content


@pytest.mark.fast
def test_store_filelike_warns_on_path_params(verbose=True, *args, **kwargs):
    """Test that path-related parameters emit warnings with file-like objects."""
    import warnings
    from io import BytesIO

    from radis.test.utils import getTestFile
    from radis.tools.database import load_spec

    s = load_spec(getTestFile("N2C_specair_380nm.spec"))
    buffer = BytesIO()

    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        s.store(buffer, compress=True, add_date="%Y%m%d", verbose=False)
        assert len(w) >= 1
        assert "add_date" in str(w[0].message)


@pytest.mark.fast
def test_store_roundtrip_via_buffer(verbose=True, *args, **kwargs):
    """Test that spectrum stored to buffer can be loaded back."""
    from io import BytesIO

    import json_tricks

    from radis.test.utils import getTestFile
    from radis.tools.database import _fix_format, _json_to_spec, load_spec

    s = load_spec(getTestFile("N2C_specair_380nm.spec"))

    buffer = BytesIO()
    s.store(buffer, compress=True, discard=[], verbose=verbose)

    buffer.seek(0)
    sload = json_tricks.load(buffer, preserve_order=False)
    sload, _ = _fix_format("buffer", sload)
    s2 = _json_to_spec(sload, "buffer")

    assert s.compare_with(s2, spectra_only=True, plot=False, verbose=verbose)


if __name__ == "__main__":
    import pytest

    # -s is to plot
    pytest.main(["test_database.py", "-s"])