import astropy.units as u
import itertools
import numpy as np
import pytest

from plasmapy.particles import IonizationState, IonizationStateCollection
from plasmapy.particles.atomic import (
    atomic_number,
    charge_number,
    common_isotopes,
    electric_charge,
    half_life,
    is_stable,
    isotopic_abundance,
    known_isotopes,
    mass_number,
    particle_mass,
    stable_isotopes,
    standard_atomic_weight,
)
from plasmapy.particles.exceptions import (
    ChargeError,
    InvalidElementError,
    InvalidIsotopeError,
    InvalidParticleError,
    ParticleError,
    ParticleWarning,
)
from plasmapy.particles.nuclear import nuclear_binding_energy, nuclear_reaction_energy
from plasmapy.particles.symbols import atomic_symbol, element_name, isotope_symbol
from plasmapy.utils.code_repr import call_string

tests_for_exceptions = {
    "too few nstates": (
        IonizationState,
        [],
        {"particle": "H", "ionic_fractions": [1.0]},
        ParticleError,
    ),
    "too many nstates": (
        IonizationState,
        [],
        {"particle": "H", "ionic_fractions": [1, 0, 0, 0]},
        ParticleError,
    ),
    "ionic fraction < 0": (
        IonizationState,
        [],
        {"particle": "He", "ionic_fractions": [-0.1, 0.1, 1]},
        ParticleError,
    ),
    "ionic fraction > 1": (
        IonizationState,
        [],
        {"particle": "He", "ionic_fractions": [1.1, 0.0, 0.0]},
        ParticleError,
    ),
    "invalid ionic fraction": (
        IonizationState,
        [],
        {"particle": "He", "ionic_fractions": [1.0, 0.0, "a"]},
        ParticleError,
    ),
    "bad n_elem units": (
        IonizationState,
        [],
        {"particle": "H", "ionic_fractions": [0, 1], "n_elem": 3 * u.m**3},
        u.UnitTypeError,
    ),
    "bad T_e units": (
        IonizationState,
        [],
        {"particle": "H", "ionic_fractions": [0, 1], "T_e": 1 * u.m},
        u.UnitTypeError,
    ),
    "negative n_elem": (
        IonizationState,
        [],
        {
            "particle": "He",
            "ionic_fractions": [1.0, 0.0, 0.0],
            "n_elem": -1 * u.m**-3,
        },
        ParticleError,
    ),
    "negative T_e": (
        IonizationState,
        [],
        {"particle": "He", "ionic_fractions": [1.0, 0.0, 0.0], "T_e": -1 * u.K},
        ParticleError,
    ),
    "redundant ndens": (
        IonizationState,
        [],
        {
            "particle": "H",
            "ionic_fractions": np.array([3, 4]) * u.m**-3,
            "n_elem": 4 * u.m**-3,
        },
        ParticleError,
    ),
    "wrong type": (IonizationStateCollection, [], {"inputs": None}, ParticleError),
    "not normalized": (
        IonizationStateCollection,
        [],
        {"inputs": {"He": [0.4, 0.5, 0.0]}, "tol": 1e-9},
        ParticleError,
    ),
    "negative ionfrac": (
        IonizationStateCollection,
        [],
        {"inputs": {"H": [-0.1, 1.1]}},
        ParticleError,
    ),
    "ion": (
        IonizationStateCollection,
        [],
        {"inputs": {"H": [0.1, 0.9], "He+": [0.0, 0.9, 0.1]}},
        ParticleError,
    ),
    "repeat elements": (
        IonizationStateCollection,
        [],
        {"inputs": {"H": [0.1, 0.9], "hydrogen": [0.2, 0.8]}},
        ParticleError,
    ),
    "isotope of element": (
        IonizationStateCollection,
        [],
        {"inputs": {"H": [0.1, 0.9], "D": [0.2, 0.8]}},
        ParticleError,
    ),
    "negative abundance": (
        IonizationStateCollection,
        [],
        {
            "inputs": {"H": [0.1, 0.9], "He": [0.4, 0.5, 0.1]},
            "abundances": {"H": 1, "He": -0.1},
        },
        ParticleError,
    ),
    "imaginary abundance": (
        IonizationStateCollection,
        [],
        {
            "inputs": {"H": [0.1, 0.9], "He": [0.4, 0.5, 0.1]},
            "abundances": {"H": 1, "He": 0.1j},
        },
        ParticleError,
    ),
    "wrong density units": (
        IonizationStateCollection,
        [],
        {
            "inputs": {"H": [10, 90] * u.m**-3, "He": [0.1, 0.9, 0] * u.m**-2},
            "abundances": {"H": 1, "He": 0.1},
        },
        ParticleError,
    ),
    "abundance redundance": (
        IonizationStateCollection,
        [],
        {
            "inputs": {"H": [10, 90] * u.m**-3, "He": [0.1, 0.9, 0] * u.m**-3},
            "abundances": {"H": 1, "He": 0.1},
        },
        ParticleError,
    ),
    "abundance contradiction": (
        IonizationStateCollection,
        [],
        {
            "inputs": {"H": [10, 90] * u.m**-3, "He": [0.1, 0.9, 0] * u.m**-3},
            "abundances": {"H": 1, "He": 0.11},
        },
        ParticleError,
    ),
    "kappa too small": (
        IonizationStateCollection,
        [],
        {"inputs": ["H"], "kappa": 1.499999},
        ParticleError,
    ),
    "negative n": (
        IonizationStateCollection,
        [],
        {"inputs": ["H"], "n0": -1 * u.cm**-3},
        ParticleError,
    ),
    "negative T_e for collection": (
        IonizationStateCollection,
        [],
        {"inputs": ["H-1"], "T_e": -1 * u.K},
        ParticleError,
    ),
}


@pytest.mark.parametrize(
    ("tested_object", "args", "kwargs", "expected_exception"),
    list(tests_for_exceptions.values()),
    ids=list(tests_for_exceptions.keys()),
)
def test_named_tests_for_exceptions(tested_object, args, kwargs, expected_exception):
    """
    Test that appropriate exceptions are raised for inappropriate inputs
    to `IonizationState` or `IonizationStateCollection`
    """
    with pytest.raises(expected_exception) as exc_info:
        tested_object(*args, **kwargs)

    assert expected_exception == exc_info.type, (
        f"When running the command "
        f"{call_string(tested_object, args, kwargs)},"
        f"an {expected_exception} was expected. Instead, a "
        f"{exc_info.type} was raised."
    )


tests_from_nuclear = [
    [
        nuclear_reaction_energy,
        [],
        {"reactants": ["n"], "products": 3},
        pytest.raises(TypeError),
    ],
    [
        nuclear_reaction_energy,
        [],
        {"reactants": ["n"], "products": ["He-4"]},
        pytest.raises(ParticleError),
    ],
    [
        nuclear_reaction_energy,
        [],
        {"reactants": ["h"], "products": ["H-1"]},
        pytest.raises(ParticleError),
    ],
    [
        nuclear_reaction_energy,
        [],
        {"reactants": ["e-", "n"], "products": ["p+"]},
        pytest.raises(ParticleError),
    ],
    [
        nuclear_reaction_energy,
        [],
        {"reactants": ["e+", "n"], "products": ["p-"]},
        pytest.raises(ParticleError),
    ],
    [
        nuclear_reaction_energy,
        [],
        {"reactants": ["ksdf"], "products": ["H-3"]},
        pytest.raises(ParticleError),
    ],
    [
        nuclear_reaction_energy,
        [],
        {"reactants": ["H"], "products": ["H-1"]},
        pytest.raises(ParticleError),
    ],
    [
        nuclear_reaction_energy,
        [],
        {"reactants": ["p"], "products": ["n", "n", "e-"]},
        pytest.raises(ParticleError),
    ],
    [
        nuclear_reaction_energy,
        ["p --> p"],
        {"reactants": "p", "products": "p"},
        pytest.raises(ParticleError),
    ],
    [nuclear_binding_energy, ["H"], {}, pytest.raises(ParticleError)],
    [nuclear_binding_energy, ["He-99"], {}, pytest.raises(InvalidParticleError)],
    [
        nuclear_binding_energy,
        ["He"],
        {"mass_numb": 99},
        pytest.raises(InvalidParticleError),
    ],
    [nuclear_binding_energy, [3.1415926535j], {}, pytest.raises(TypeError)],
]

tests_from_atomic = [
    [
        atomic_symbol,
        [
            "H-0",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        atomic_symbol,
        [
            3.14159,
        ],
        {},
        pytest.raises(TypeError),
    ],
    [
        atomic_symbol,
        [
            "Og-294b",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        atomic_symbol,
        [
            "H-934361079326356530741942970523610389",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        atomic_symbol,
        [
            "Fe 2+4",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        atomic_symbol,
        [
            "Fe+24",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        atomic_symbol,
        [
            "Fe +59",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        atomic_symbol,
        [
            "C++++++++++++++++",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        atomic_symbol,
        [
            "C-++++",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        atomic_symbol,
        [
            "neutron",
        ],
        {},
        pytest.raises(InvalidElementError),
    ],
    [
        atomic_symbol,
        [
            "n",
        ],
        {},
        pytest.raises(InvalidElementError),
    ],
    [
        atomic_symbol,
        [
            "n-1",
        ],
        {},
        pytest.raises(InvalidElementError),
    ],
    [
        atomic_symbol,
        [
            "h",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        atomic_symbol,
        [
            "d",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        atomic_symbol,
        [
            "he",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        atomic_symbol,
        [
            "au",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        atomic_symbol,
        [
            "p-",
        ],
        {},
        pytest.raises(InvalidElementError),
    ],
    [
        atomic_symbol,
        [
            0,
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        atomic_symbol,
        [
            119,
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        atomic_symbol,
        [
            "antiproton",
        ],
        {},
        pytest.raises(InvalidElementError),
    ],
    [
        atomic_number,
        [
            "H-3934",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        atomic_number,
        [
            "C-12b",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        atomic_number,
        [
            -1.5,
        ],
        {},
        pytest.raises(TypeError),
    ],
    [
        atomic_number,
        [
            "n",
        ],
        {},
        pytest.raises(InvalidElementError),
    ],
    [
        atomic_number,
        [
            "n-1",
        ],
        {},
        pytest.raises(InvalidElementError),
    ],
    [
        atomic_number,
        [
            "neutron",
        ],
        {},
        pytest.raises(InvalidElementError),
    ],
    [
        atomic_number,
        [
            "Neutron",
        ],
        {},
        pytest.raises(InvalidElementError),
    ],
    [
        atomic_number,
        [
            "d",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        atomic_number,
        [
            "t",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        atomic_number,
        [
            "s-36",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        mass_number,
        [
            "H-359",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        mass_number,
        [
            "C-12b",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        mass_number,
        [
            -1.5,
        ],
        {},
        pytest.raises(TypeError),
    ],
    [
        mass_number,
        [
            "N-13+-+-",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        mass_number,
        [
            "h-3",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        mass_number,
        [
            "n",
        ],
        {},
        pytest.raises(InvalidIsotopeError),
    ],
    [
        mass_number,
        [
            "n-1",
        ],
        {},
        pytest.raises(InvalidIsotopeError),
    ],
    [
        element_name,
        [
            "vegancupcakes",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        element_name,
        [
            "C-+-",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        element_name,
        [
            1.24,
        ],
        {},
        pytest.raises(TypeError),
    ],
    [
        element_name,
        [
            "n",
        ],
        {},
        pytest.raises(InvalidElementError),
    ],
    [
        element_name,
        [
            "neutron",
        ],
        {},
        pytest.raises(InvalidElementError),
    ],
    [
        element_name,
        [
            0,
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        element_name,
        [
            "H++",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        element_name,
        [
            "t",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        element_name,
        [
            "pb",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        element_name,
        [
            "d",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        element_name,
        [
            "h-3",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        element_name,
        [
            "Pb-9",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        element_name,
        [
            "H 2+",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        standard_atomic_weight,
        [
            "H-1",
        ],
        {},
        pytest.raises(ParticleError),
    ],
    [
        standard_atomic_weight,
        [
            "help i'm trapped in a unit test",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        standard_atomic_weight,
        [
            1.1,
        ],
        {},
        pytest.raises(TypeError),
    ],
    [
        standard_atomic_weight,
        [
            "n",
        ],
        {},
        pytest.raises(InvalidElementError),
    ],
    [
        standard_atomic_weight,
        [
            "p",
        ],
        {},
        pytest.raises(ParticleError),
    ],
    [
        standard_atomic_weight,
        [
            "alpha",
        ],
        {},
        pytest.raises(ParticleError),
    ],
    [
        standard_atomic_weight,
        [
            "deuteron",
        ],
        {},
        pytest.raises(ParticleError),
    ],
    [
        standard_atomic_weight,
        [
            "tritium",
        ],
        {},
        pytest.raises(ParticleError),
    ],
    [
        standard_atomic_weight,
        [
            "Au+",
        ],
        {},
        pytest.raises(ParticleError),
    ],
    [
        standard_atomic_weight,
        [
            "Fe -2",
        ],
        {},
        pytest.raises(ParticleError),
    ],
    [
        standard_atomic_weight,
        [
            "Og 2+",
        ],
        {},
        pytest.raises(ParticleError),
    ],
    [
        standard_atomic_weight,
        [
            "h",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        standard_atomic_weight,
        [
            "fe",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        electric_charge,
        [
            "badinput",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        electric_charge,
        [
            "h+",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        electric_charge,
        [
            "Au 81+",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        electric_charge,
        [
            "Au 81-",
        ],
        {},
        pytest.warns(ParticleWarning),
    ],
    [
        electric_charge,
        [
            "H---",
        ],
        {},
        pytest.warns(ParticleWarning),
    ],
    [
        charge_number,
        [
            "fads",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        charge_number,
        [
            "H++",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        charge_number,
        [
            "h+",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        charge_number,
        [
            "fe 1+",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        charge_number,
        [
            "d+",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        charge_number,
        [
            "Fe 29+",
        ],
        {},
        pytest.raises(InvalidParticleError),
    ],
    [
        charge_number,
        [
            "H-1",
        ],
        {},
        pytest.raises(ChargeError),
    ],
    [
        isotope_symbol,
        ("Md-260",),
        {"mass_numb": 261},
        pytest.raises(InvalidParticleError),
    ],
    [
        isotope_symbol,
        ("protium",),
        {"mass_numb": 2},
        pytest.raises(InvalidParticleError),
    ],
    [isotope_symbol, ("alpha",), {"mass_numb": 3}, pytest.raises(InvalidParticleError)],
    [isotope_symbol, ("O-18",), {"mass_numb": 19}, pytest.raises(InvalidParticleError)],
    [
        isotope_symbol,
        ("lead-209",),
        {"mass_numb": 511},
        pytest.raises(InvalidParticleError),
    ],
    [isotope_symbol, ("He-1",), {}, pytest.raises(InvalidParticleError)],
    [isotope_symbol, [24], {"mass_numb": 23}, pytest.raises(InvalidParticleError)],
    [isotope_symbol, ("H",), {"mass_numb": 0}, pytest.raises(InvalidParticleError)],
    [isotope_symbol, ("H-1",), {"mass_numb": 2}, pytest.raises(InvalidParticleError)],
    [isotope_symbol, ("P",), {}, pytest.raises(InvalidIsotopeError)],
    [isotope_symbol, [1], {}, pytest.raises(InvalidIsotopeError)],
    [isotope_symbol, [4], {}, pytest.raises(InvalidIsotopeError)],
    [isotope_symbol, ("hydrogen-444444",), {}, pytest.raises(InvalidParticleError)],
    [isotope_symbol, ("Fe",), {"mass_numb": 2.1}, pytest.raises(TypeError)],
    [isotope_symbol, ("He",), {"mass_numb": "c"}, pytest.raises(TypeError)],
    [isotope_symbol, ("He-3",), {"mass_numb": 4}, pytest.raises(InvalidParticleError)],
    [isotope_symbol, ("D",), {"mass_numb": 3}, pytest.raises(InvalidParticleError)],
    [isotope_symbol, ("T",), {"mass_numb": 2}, pytest.raises(InvalidParticleError)],
    [isotope_symbol, ("Fe",), {"mass_numb": None}, pytest.raises(InvalidIsotopeError)],
    [isotope_symbol, ("He",), {"mass_numb": 99}, pytest.raises(InvalidParticleError)],
    [isotope_symbol, ("d",), {}, pytest.raises(InvalidParticleError)],
    [isotope_symbol, ("h-3",), {}, pytest.raises(InvalidParticleError)],
    [isotope_symbol, ("h",), {}, pytest.raises(InvalidParticleError)],
    [isotope_symbol, ("d+",), {}, pytest.raises(InvalidParticleError)],
    [particle_mass, ["H-1 +1"], {"Z": 0}, pytest.raises(InvalidParticleError)],
    [particle_mass, [26], {"Z": 1, "mass_numb": "a"}, pytest.raises(TypeError)],
    [
        particle_mass,
        [26],
        {"Z": 27, "mass_numb": 56},
        pytest.raises(InvalidParticleError),
    ],
    [
        particle_mass,
        ["Og"],
        {"mass_numb": 696, "Z": 1},
        pytest.raises(InvalidParticleError),
    ],
    [particle_mass, ["He 1+"], {"mass_numb": 99}, pytest.raises(InvalidParticleError)],
    [particle_mass, ["fe-56 1+"], {}, pytest.raises(InvalidParticleError)],
    [is_stable, ["hydrogen-444444"], {}, pytest.raises(InvalidParticleError)],
    [is_stable, ["hydrogen", 0], {}, pytest.raises(InvalidParticleError)],
    [is_stable, ["pb-209"], {}, pytest.raises(InvalidParticleError)],
    [is_stable, ["h"], {}, pytest.raises(InvalidParticleError)],
    [is_stable, ["He"], {}, pytest.raises(InvalidIsotopeError)],
    [is_stable, ["B"], {}, pytest.raises(InvalidIsotopeError)],
    [particle_mass, ["H-1"], {"mass_numb": 1, "Z": 1}, pytest.warns(ParticleWarning)],
    [isotope_symbol, ("H-1",), {"mass_numb": 1}, pytest.warns(ParticleWarning)],
    [isotope_symbol, ("H-2",), {"mass_numb": 2}, pytest.warns(ParticleWarning)],
    [isotope_symbol, ("T",), {"mass_numb": 3}, pytest.warns(ParticleWarning)],
    [isotope_symbol, ("Li-6",), {"mass_numb": 6}, pytest.warns(ParticleWarning)],
    [isotope_symbol, ("lithium-6",), {"mass_numb": 6}, pytest.warns(ParticleWarning)],
    [isotope_symbol, ("alpha",), {"mass_numb": 4}, pytest.warns(ParticleWarning)],
]


atomic_TypeError_funcs_table = [
    atomic_symbol,
    isotope_symbol,
    atomic_number,
    is_stable,
    half_life,
    mass_number,
    element_name,
    standard_atomic_weight,
    nuclear_binding_energy,
    nuclear_reaction_energy,
]

atomic_TypeError_badargs = [1.1, {"cats": "bats"}, 1 + 1j]

atomic_ParticleErrors_funcs_table = [
    atomic_symbol,
    isotope_symbol,
    atomic_number,
    is_stable,
    half_life,
    mass_number,
    element_name,
    standard_atomic_weight,
    particle_mass,
    known_isotopes,
    stable_isotopes,
    common_isotopes,
    isotopic_abundance,
    electric_charge,
]

atomic_ParticleError_badargs = [
    -1,
    119,
    "grumblemuffins",
    "H-0",
    "Og-294b",
    "H-9343610",
    "Fe 2+4",
    "Fe+24",
    "Fe +59",
    "C++++++++++++++++",
    "C-++++",
    "h",
    "d",
    "he",
    "au",
    "alpha 1+",
    "alpha-4",
]

particle_error_tests = [
    (function, [bad_argument], {}, pytest.raises(InvalidParticleError))
    for function, bad_argument in itertools.product(
        atomic_ParticleErrors_funcs_table, atomic_ParticleError_badargs
    )
]
type_error_tests = [
    (function, [bad_argument], {}, pytest.raises(TypeError))
    for function, bad_argument in itertools.product(
        atomic_TypeError_funcs_table, atomic_TypeError_badargs
    )
]


@pytest.mark.parametrize(
    ("tested_object", "args", "kwargs", "expected"),
    tests_from_nuclear + tests_from_atomic + particle_error_tests + type_error_tests,
)
def test_unnamed_tests_exceptions(tested_object, args, kwargs, expected):
    """
    Test that appropriate exceptions are raised for inappropriate inputs
    to different functions.
    """
    with expected as exc_info:
        tested_object(*args, **kwargs)

    if hasattr(expected, "expected_exception"):
        assert type(expected.expected_exception()) == exc_info.type

    if hasattr(expected, "expected_warning"):
        for expected_warning, recorded_warning in zip(
            exc_info.expected_warning, exc_info.list
        ):
            assert expected_warning == recorded_warning.category
