#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""

Summary
-------

Test Spectrum rescaling methods

-------------------------------------------------------------------------------


"""

import numpy as np
import pytest

import radis
from radis.misc.printer import printm
from radis.spectrum.rescale import get_recompute, get_redundant
from radis.test.utils import getTestFile
from radis.tools.database import load_spec


@pytest.mark.fast
def test_compression(verbose=True, warnings=True, *args, **kwargs):
    """Test that redundant quantities are properly infered from already known
    spectral quantities"""

    # Get spectrum
    s1 = load_spec(getTestFile("CO_Tgas1500K_mole_fraction0.01.spec"), binary=True)

    s1.conditions["thermal_equilibrium"] = True
    s1.update()

    # Analyse redundant spectral quantities
    redundant = get_redundant(s1)
    if verbose:
        print(redundant)

    assert redundant == {
        "emissivity_noslit": True,
        "radiance_noslit": True,
        "emisscoeff": True,
        "transmittance_noslit": True,
        "absorbance": True,
        "abscoeff": False,
    }

    return True


@pytest.mark.fast
def test_update_transmittance(verbose=True, warnings=True, *args, **kwargs):
    """Test that update can correctly recompute missing quantities"""
    # TODO: add one with radiance too

    # Work with a Spectrum object that was generated by Specair
    s = load_spec(getTestFile("N2C_specair_380nm.spec"))
    w1, T1 = s.get("transmittance_noslit")  # our reference

    if verbose:
        debug_mode = radis.config["DEBUG_MODE"]  # shows all the rescale steps taken
        radis.config["DEBUG_MODE"] = True

    # Now let's apply some update() steps

    # 1) Make sure updating doesnt change anything
    s.update()
    w2, T2 = s.get("transmittance_noslit")

    # 2) Now recompute from abscoeff
    del s._q["transmittance_noslit"]
    s.update()
    w2, T3 = s.get("transmittance_noslit")

    # 3) Now recompute from absorbance
    del s._q["transmittance_noslit"]
    del s._q["abscoeff"]
    s.update()
    w2, T4 = s.get("transmittance_noslit")

    if verbose:
        radis.config["DEBUG_MODE"] = debug_mode

    # Compare
    assert np.allclose(T1, T2)
    assert np.allclose(T1, T3)
    assert np.allclose(T1, T4)

    return True


def test_get_recompute(verbose=True, *args, **kwargs):
    """Make sure :func:`~radis.spectrum.rescale.get_recompute` works as expected

    Here, we check which quantities are needed to recompute radiance_noslit"""

    # Equilibrium
    # -----------
    s = load_spec(getTestFile("CO_Tgas1500K_mole_fraction0.01.spec"), binary=True)

    assert s.get_vars() == ["abscoeff"]
    assert s.conditions["thermal_equilibrium"]
    # At equilibrium, everything should be deduced from abscoeff
    assert set(get_recompute(s, ["radiance_noslit"])) == set(
        ("radiance_noslit", "abscoeff")
    )

    # Non Equilibrium
    # ----------------
    s.conditions["Tvib"] = 2000
    # a problem should be detected by is_at_equilibrium()
    with pytest.raises(AssertionError):
        assert not s.is_at_equilibrium(check="error")
    # force non equilibrium
    s.conditions["thermal_equilibrium"] = False

    # Now more data is needed:
    assert set(get_recompute(s, ["radiance_noslit"])) == set(
        ("abscoeff", "emisscoeff", "radiance_noslit")
    )


def test_recompute_equilibrium(verbose=True, warnings=True, plot=True, *args, **kwargs):
    """Test that spectral quantities recomputed under equilibrium assumption
    yields the same output as with non equilibrium routines when Tvib = Trot"""

    if plot:
        import matplotlib.pyplot as plt

        plt.ion()  # dont get stuck with Matplotlib if executing through pytest

    # Get spectrum
    s1 = load_spec(getTestFile("CO_Tgas1500K_mole_fraction0.01.spec"))
    s1.rescale_path_length(100)  # just for fun

    assert s1.is_at_equilibrium()
    s1.update("emisscoeff")

    # force non equilibrium calculation
    s2 = s1.copy()
    s2.conditions["thermal_equilibrium"] = False
    s2.update()
    assert (
        "emissivity_noslit" not in s2.get_vars()
    )  # just a check update() was done at nonequilibrium

    # update s1 now (at equilibrium)
    s1.update()
    assert (
        "emissivity_noslit" in s1.get_vars()
    )  # just a check update() was done at equilibrium

    s1.name = "scaled with Kirchoff law"
    s2.name = "scaled from emisscoeff + abscoeff with RTE"

    if verbose:
        print(
            "Checked that scaling at equilibrium with Kirchoff law yields the "
            + "same radiance as by solving the RTE from emisscoeff and abscoeff"
        )

    # Now Compare both update processes
    assert s1.compare_with(s2, spectra_only="radiance_noslit", plot=plot)


def test_rescale_all_quantities(verbose=True, warnings=True, *args, **kwargs):

    new_mole_fraction = 0.5
    new_path_length = 0.1

    # Get spectrum
    s0 = load_spec(getTestFile("CO_Tgas1500K_mole_fraction0.01.spec"), binary=True)
    s0.update("all")  # start with all possible quantities in s0
    sscaled = s0.copy()
    sscaled.rescale_mole_fraction(new_mole_fraction)
    sscaled.rescale_path_length(new_path_length)
    s0.conditions["thermal_equilibrium"] = False  # to prevent rescaling with Kirchoff
    # remove emissivity_no_slit (speciifc to equilibrium)
    del s0._q["emissivity_noslit"]

    # Determine all quantities that can be recomputed
    if verbose >= 2:
        import radis

        DEBUG_MODE = radis.config["DEBUG_MODE"]
        radis.config["DEBUG_MODE"] = True
    from radis.spectrum.rescale import _build_update_graph, get_reachable, ordered_keys

    # ordered_keys: all spectral quantities that can be rescaled
    can_be_recomputed = get_reachable(s0)
    # can_be_recomputed: all spectra quantities that can be rescaled for this
    # particular spectrum
    update_paths = _build_update_graph(s0)
    # update_paths: which quantities are needed to recompute the others

    rescale_list = [k for k in ordered_keys if can_be_recomputed[k]]

    for quantity in rescale_list:
        all_paths = update_paths[quantity]
        if verbose:
            printm(
                "{0} can be recomputed from {1}".format(
                    quantity,
                    " or ".join(["&".join(combinations) for combinations in all_paths]),
                )
            )

        # Now let's test all paths
        for combinations in all_paths:
            if verbose:
                printm(
                    "> computing {0} from {1}".format(quantity, "&".join(combinations))
                )
            s = s0.copy()
            # Delete all other quantities
            for k in s.get_vars():
                if k not in combinations:
                    del s._q[k]

            s.update(quantity, verbose=verbose)

            # Now rescale
            s.rescale_mole_fraction(new_mole_fraction)
            s.rescale_path_length(new_path_length)

            # Compare
            assert s.compare_with(sscaled, spectra_only=quantity, plot=False)

    if verbose >= 2:
        radis.config["DEBUG_MODE"] = DEBUG_MODE


def _run_all_tests(verbose=True, warnings=True, *args, **kwargs):
    test_compression(verbose=verbose, warnings=warnings, *args, **kwargs)
    test_update_transmittance(verbose=verbose, warnings=warnings, *args, **kwargs)
    test_get_recompute(verbose=verbose, warnings=warnings, *args, **kwargs)
    test_recompute_equilibrium(verbose=verbose, warnings=warnings, *args, **kwargs)
    test_rescale_all_quantities(verbose=verbose, *args, **kwargs)

    return True


if __name__ == "__main__":
    print(("Testing test_rescale.py:", _run_all_tests(verbose=True)))
