import warnings

import numpy as np
import pytest
from at import (
    element_pass,
    element_track,
    elements,
    internal_epass,
    internal_lpass,
    lattice,
    lattice_pass,
    lattice_track,
    AtWarning,
)
from at.lattice.elements.conversions import _array, _array66


def test_data_checks():
    val = np.zeros([6, 6])
    assert _array(val).shape == (36,)
    assert _array66(val).shape == (6, 6)


def test_element_string_ordering():
    d = elements.Drift("D0", 1, attr=np.array(0))
    assert d.__str__() == (
        "Drift:\n       FamName: D0\n        Length: 1.0\n"
        "    PassMethod: DriftPass\n          attr: 0"
    )
    assert d.__repr__() == "Drift('D0', 1.0, attr=array(0))"


def test_element_creation_raises_exception():
    with pytest.raises(ValueError):
        elements.Element("family_name", R1="not_an_array")


def test_base_element_methods():
    e = elements.Element("family_name")
    assert e.divide([0.2, 0.5, 0.3]) == [e]
    assert id(e.copy()) != id(e)


def test_argument_checks():
    q = elements.Quadrupole("quad", 1.0, 0.5)
    # Test type
    with pytest.raises(ValueError):
        q.Length = "a"
    # Test shape
    with pytest.raises(ValueError):
        q.T1 = [0.0, 0.0]
    # Test coherence of polynoms
    with pytest.raises(ValueError):
        q.MaxOrder = 2
    with pytest.raises(ValueError):
        q.PolynomA = [0.0]
    with pytest.raises(ValueError):
        q.PolynomB = [0.0]


def test_dipole():
    d = elements.Dipole("dipole", 1.0, 0.01)
    assert d.MaxOrder == 1
    assert len(d.PolynomA) == 2
    assert d.K == 0.0
    d = elements.Dipole("dipole", 1.0, 0.01, -0.5)
    assert d.MaxOrder == 1
    assert len(d.PolynomA) == 2
    assert d.K == -0.5
    d = elements.Dipole("dipole", 1.0, 0.01, PolynomB=[0.0, 0.1, 0.0])
    assert d.MaxOrder == 2
    assert len(d.PolynomA) == 3
    assert d.K == 0.1
    d = elements.Dipole("dipole", 1.0, 0.01, PolynomB=[0.0, 0.0, 0.005])
    assert d.MaxOrder == 2
    assert len(d.PolynomA) == 3
    assert d.K == 0.0
    with pytest.warns(AtWarning):
        d = elements.Dipole("dipole", 1.0, 0.01, PolynomB=[0.0, 0.0, 0.005], MaxOrder=0)
    assert d.MaxOrder == 0
    assert len(d.PolynomA) == 3
    assert d.K == 0.0


def test_quadrupole():
    q = elements.Quadrupole("quadrupole", 1.0)
    assert q.MaxOrder == 1
    assert len(q.PolynomA) == 2
    assert q.K == 0.0
    q = elements.Quadrupole("quadrupole", 1.0, -0.5)
    assert q.MaxOrder == 1
    assert len(q.PolynomA) == 2
    assert q.K == -0.5
    q = elements.Quadrupole("quadrupole", 1.0, PolynomB=[0.0, 0.0, 0.005])
    assert q.MaxOrder == 2
    assert len(q.PolynomA) == 3
    assert q.K == 0.0
    with pytest.warns(AtWarning):
        q = elements.Quadrupole("quadrupole", 1.0, PolynomB=[0.0, 0.5, 0.005], MaxOrder=1)
    assert q.MaxOrder == 1
    assert len(q.PolynomA) == 3
    assert q.K == 0.5


def test_sextupole():
    s = elements.Sextupole("sextupole", 1.0)
    assert s.MaxOrder == 2
    assert len(s.PolynomA) == 3
    assert s.H == 0.0
    s = elements.Sextupole("sextupole", 1.0, -0.5)
    assert s.MaxOrder == 2
    assert len(s.PolynomA) == 3
    assert s.H == -0.5
    s = elements.Sextupole("sextupole", 1.0, PolynomB=[0.0, 0.0, 0.005, 0.0])
    assert s.MaxOrder == 3
    assert len(s.PolynomA) == 4
    assert s.H == 0.005
    s = elements.Sextupole("sextupole", 1.0, PolynomB=[0.0, 0.0, 0.005, 0.001])
    assert s.MaxOrder == 3
    assert len(s.PolynomA) == 4
    assert s.H == 0.005
    with pytest.warns(AtWarning):
        s = elements.Sextupole(
            "sextupole", 1.0, PolynomB=[0.0, 0.5, 0.005, 0.001], MaxOrder=2
        )
    assert s.MaxOrder == 2
    assert len(s.PolynomA) == 4
    assert s.H == 0.005


def test_octupole():
    o = elements.Octupole("octupole", 1.0, [], [0.0, 0.0, 0.0, 0.0])
    assert o.MaxOrder == 3
    assert len(o.PolynomA) == 4


def test_thinmultipole():
    m = elements.ThinMultipole("thin", [], [0.0, 0.0, 0.0, 0.0])
    assert m.MaxOrder == 3
    assert len(m.PolynomA) == 4
    m = elements.ThinMultipole("thin", [], [0.0, 0.0, 1.0, 0.0])
    assert m.MaxOrder == 3
    assert len(m.PolynomA) == 4


def test_multipole():
    m = elements.Multipole("multi", 1.0, [], [0.0, 0.0, 0.0, 0.0])
    assert m.Length == 1.0
    assert m.MaxOrder == 3
    assert m.NumIntSteps == 10
    assert m.PassMethod == "StrMPoleSymplectic4Pass"


@pytest.mark.parametrize(
    "element_type,args",
    [
        (elements.Multipole, ["Multi", 1.0]),
        (elements.ThinMultipole, ["ThinMulti"]),
        (elements.Octupole, ["Oct", 1.0]),
    ],
)
def test_PolynomA_strength_prioritisation(element_type, args):
    # Warning and changes when poly_a is zeros and PolynomA isn't
    with pytest.warns(lattice.AtWarning):
        elem = element_type(*args, [0.0, 0.0], [], PolynomA=[1.0, 2.0])
    assert elem.PolynomA[0] == 1.0
    assert elem.PolynomA[1] == 2.0
    # Warning and no changes when PolynomA and poly_a are the same and non-zero
    with pytest.warns(lattice.AtWarning):
        elem = element_type(*args, [1.0, 2.0], [], PolynomA=[1.0, 2.0])
    assert elem.PolynomA[0] == 1.0
    assert elem.PolynomA[1] == 2.0
    # Error when poly_a is non-zero and different to PolynomA, even if PolynomA is zeros
    with pytest.raises(lattice.AtError):
        elem = element_type(*args, [1.0], [], PolynomA=[0.0])
    # Check poly_a and PolynomA are made the same length
    with pytest.warns(lattice.AtWarning):
        elem = element_type(*args, [1.0], [], PolynomA=[1.0, 0.0, 0.0, 0.0, 0.0])
    assert len(elem.PolynomA) == 5
    with pytest.warns(lattice.AtWarning):
        elem = element_type(*args, [1.0, 0.0, 0.0, 0.0, 0.0], [], PolynomA=[1.0])
    assert len(elem.PolynomA) == 5


@pytest.mark.parametrize(
    "element_type,args",
    [
        (elements.Multipole, ["Multi", 1.0]),
        (elements.ThinMultipole, ["ThinMulti"]),
        (elements.Octupole, ["Oct", 1.0]),
    ],
)
def test_PolynomB_strength_prioritisation(element_type, args):
    # Warning and changes when poly_b is zeros and PolynomB isn't
    with pytest.warns(lattice.AtWarning):
        elem = element_type(*args, [], [0.0, 0.0], PolynomB=[1.0, 2.0])
    assert elem.PolynomB[0] == 1.0
    assert elem.PolynomB[1] == 2.0
    # Warning and no changes when PolynomB and poly_b are the same and non-zero
    with pytest.warns(lattice.AtWarning):
        elem = element_type(*args, [], [1.0, 2.0], PolynomB=[1.0, 2.0])
    assert elem.PolynomB[0] == 1.0
    assert elem.PolynomB[1] == 2.0
    # Error when poly_b is non-zero and different to PolynomB, even if PolynomB is zeros
    with pytest.raises(lattice.AtError):
        elem = element_type(*args, [], [1.0], PolynomB=[0.0])
    # Check poly_a and PolynomA are made the same length
    with pytest.warns(lattice.AtWarning):
        elem = element_type(*args, [], [1.0], PolynomB=[1.0, 0.0, 0.0, 0.0, 0.0])
    assert len(elem.PolynomB) == 5
    with pytest.warns(lattice.AtWarning):
        elem = element_type(*args, [], [1.0, 0.0, 0.0, 0.0, 0.0], PolynomB=[1.0])
    assert len(elem.PolynomB) == 5


@pytest.mark.parametrize(
    "element_type,args",
    [
        (elements.Multipole, ["Multi", 1.0]),
        (elements.ThinMultipole, ["ThinMulti"]),
        (elements.Octupole, ["Oct", 1.0]),
    ],
)
def test_K_and_H_strength_prioritisation(element_type, args):
    # Warning and no change when K is zero
    with pytest.warns(lattice.AtWarning):
        elem = element_type(*args, [], [0.0, 1.0], K=0.0)
    assert elem.PolynomB[1] == 1.0
    # Warning and no change when H is zero
    with pytest.warns(lattice.AtWarning):
        elem = element_type(*args, [], [0.0, 0.0, 2.0], H=0.0)
    assert elem.PolynomB[2] == 2.0
    # Warning and no change when poly_b[1] and K are the same and non-zero
    with pytest.warns(lattice.AtWarning):
        elem = element_type(*args, [], [0.0, 1.0], K=1.0)
    assert elem.PolynomB[1] == 1.0
    # Warning and no change when poly_b[2] and H are the same and non-zero
    with pytest.warns(lattice.AtWarning):
        elem = element_type(*args, [], [0.0, 0.0, 2.0], H=2.0)
    assert elem.PolynomB[2] == 2.0
    # Error when K is non-zero and different to poly_b[1], even if poly_b[1] is 0.0
    with pytest.raises(lattice.AtError):
        elem = element_type(*args, [], [0.0, 0.0], K=2.0)
    # Error when H is non-zero and different to poly_b[2], even if poly_b[2] is 0.0
    with pytest.raises(lattice.AtError):
        elem = element_type(*args, [], [0.0, 0.0, 0.0], H=2.0)
    # Error when len(poly_b)<2 and K is specified
    with pytest.raises(lattice.AtError):
        elem = element_type(*args, [], [0.0], K=1.0)
    # Error when len(poly_b)<3 and H is specified
    with pytest.raises(lattice.AtError):
        elem = element_type(*args, [], [0.0, 0.0], H=1.0)


@pytest.mark.parametrize(
    "element_type,args",
    [
        (elements.Quadrupole, ["Quad", 1.0]),
        (elements.Dipole, ["Dipole", 1.0, 0.0]),
    ],
)
def test_quadrupolar_strength_prioritisation(element_type, args):
    # No warning or change when k isn't specified, and is therefore default value (0.0)
    with warnings.catch_warnings():
        warnings.simplefilter("error")
        elem = element_type(*args, PolynomB=[0.0, 1.0])
    assert elem.PolynomB[1] == 1.0
    # No warning or change when PolynomB isn't specified
    with warnings.catch_warnings():
        warnings.simplefilter("error")
        elem = element_type(*args, k=1.0)
    assert elem.PolynomB[1] == 1.0
    # No warning or change when PolynomB[1] and k are the same and non-zero
    with warnings.catch_warnings():
        warnings.simplefilter("error")
        elem = element_type(*args, k=1.0, PolynomB=[0.0, 1.0])
    assert elem.PolynomB[1] == 1.0
    # Error when k is non-zero and different to PolynomB[1], even if PolynomB[1] is 0.0
    with pytest.raises(lattice.AtError):
        elem = element_type(*args, k=1.0, PolynomB=[0.0, 0.0])
    # No warning or change when PolynomB[1] and k are the same and non-zero even if the
    # other terms of PolynomB are non-zero
    with warnings.catch_warnings():
        warnings.simplefilter("error")
        elem = element_type(*args, k=1.0, PolynomB=[0.5, 1.0])
    assert elem.PolynomB[0] == 0.5
    assert elem.PolynomB[1] == 1.0


def test_sextupolar_strength_prioritisation():
    # No warning or change when h isn't specified, and is therefore default value (0.0)
    with warnings.catch_warnings():
        warnings.simplefilter("error")
        elem = elements.Sextupole("Sext", 1.0, PolynomB=[0.0, 0.0, 1.0])
    assert elem.PolynomB[2] == 1.0
    # No warning or change when PolynomB isn't specified
    with warnings.catch_warnings():
        warnings.simplefilter("error")
        elem = elements.Sextupole("Sext", 1.0, h=1.0)
    assert elem.PolynomB[2] == 1.0
    # No warning or change when PolynomB[2] and h are the same and non-zero
    with warnings.catch_warnings():
        warnings.simplefilter("error")
        elem = elements.Sextupole("Sext", 1.0, h=1.0, PolynomB=[0.0, 0.0, 1.0])
    assert elem.PolynomB[2] == 1.0
    # Error when h is non-zero and different to PolynomB[2], even if PolynomB[2] is 0.0
    with pytest.raises(lattice.AtError):
        elem = elements.Sextupole("Sext", 1.0, h=1.0, PolynomB=[0.0, 0.0, 0.0])
    # No warning or change when PolynomB[2] and h are the same and non-zero even if the
    # other terms of PolynomB are non-zero
    with warnings.catch_warnings():
        warnings.simplefilter("error")
        elem = elements.Sextupole("Sext", 1.0, h=1.0, PolynomB=[0.5, 0.0, 1.0])
    assert elem.PolynomB[0] == 0.5
    assert elem.PolynomB[2] == 1.0


def test_divide_splits_attributes_correctly():
    pre = elements.Drift("drift", 1)
    post = pre.divide([0.2, 0.5, 0.3])
    assert len(post) == 3
    assert sum([e.Length for e in post]) == pre.Length
    pre = elements.Dipole("dipole", 1, KickAngle=[0.5, -0.5], BendingAngle=0.2)
    post = pre.divide([0.2, 0.5, 0.3])
    assert len(post) == 3
    assert sum([e.Length for e in post]) == pre.Length
    assert sum([e.KickAngle[0] for e in post]) == pre.KickAngle[0]
    assert sum([e.KickAngle[1] for e in post]) == pre.KickAngle[1]
    assert sum([e.BendingAngle for e in post]) == pre.BendingAngle
    pre = elements.RFCavity(
        "rfc", 1, voltage=187500, frequency=3.5237e8, harmonic_number=31, energy=6.0e9
    )
    post = pre.divide([0.2, 0.5, 0.3])
    assert len(post) == 3
    assert sum([e.Length for e in post]) == pre.Length
    assert sum([e.Voltage for e in post]) == pre.Voltage


def test_insert_into_drift():
    # Create elements
    drift = elements.Drift("drift", 1)
    monitor = elements.Monitor("bpm")
    quad = elements.Quadrupole("quad", 0.3)
    # Test None splitting behaviour
    el_list = drift.insert([(0.0, None), (0.3, None), (0.7, None), (1.0, None)])
    assert len(el_list) == 3
    np.testing.assert_almost_equal([e.Length for e in el_list], [0.3, 0.4, 0.3])
    # Test normal insertion
    el_list = drift.insert([(0.3, monitor), (0.7, quad)])
    assert len(el_list) == 5
    np.testing.assert_almost_equal(
        [e.Length for e in el_list], [0.3, 0.0, 0.25, 0.3, 0.15]
    )
    # Test insertion at either end produces -ve length drifts
    el_list = drift.insert([(0.0, quad), (1.0, quad)])
    assert len(el_list) == 5
    np.testing.assert_almost_equal(
        [e.Length for e in el_list], [-0.15, 0.3, 0.7, 0.3, -0.15]
    )


@pytest.mark.parametrize("func", (lattice_track, lattice_pass, internal_lpass))
def test_correct_dimensions_does_not_raise_error(rin, func):
    func([], rin, 1)
    rin = np.zeros((6,))
    func([], rin, 1)
    rin = np.array(np.zeros((6, 2), order="F"))
    func([], rin, 1)


@pytest.mark.parametrize("dipole_class", (elements.Dipole, elements.Bend))
@pytest.mark.parametrize("func", (element_track, element_pass, internal_epass))
def test_dipole_bend_synonym(rin, dipole_class, func):
    b = dipole_class("dipole", 1.0, 0.1, EntranceAngle=0.05, ExitAngle=0.05)
    rin[0, 0] = 1e-6
    if func == element_track:
        func(b, rin, in_place=True)
    else:
        func(b, rin)
    rin_expected = np.array([1e-6, 0, 0, 0, 0, 1e-7]).reshape((6, 1))
    np.testing.assert_almost_equal(rin, rin_expected)
    assert b.K == 0.0
    b.PolynomB[1] = 0.2
    assert b.K == 0.2
    b.K = 0.1
    assert b.PolynomB[1] == 0.1


@pytest.mark.parametrize("func", (element_track, element_pass, internal_epass))
def test_marker(rin, func):
    m = elements.Marker("marker")
    assert m.Length == 0
    rin = np.array(np.random.rand(*rin.shape), order="F")
    rin_orig = np.array(rin, copy=True, order="F")
    if func == element_track:
        func(m, rin, in_place=True)
    else:
        func(m, rin)
    np.testing.assert_equal(rin, rin_orig)


@pytest.mark.parametrize("func", (element_track, element_pass, internal_epass))
def test_monitor(rin, func):
    mon = elements.Monitor("monitor")
    assert mon.Length == 0
    rin = np.array(np.random.rand(*rin.shape), order="F")
    rin_orig = rin.copy()
    if func == element_track:
        func(mon, rin, in_place=True)
    else:
        func(mon, rin)
    np.testing.assert_equal(rin, rin_orig)


@pytest.mark.parametrize("func", (element_track, element_pass, internal_epass))
def test_aperture_inside_limits(rin, func):
    a = elements.Aperture("aperture", [-1e-3, 1e-3, -1e-4, 1e-4])
    assert a.Length == 0
    rin[0, 0] = 1e-5
    rin[2, 0] = -1e-5
    rin_orig = rin.copy()
    if func == element_track:
        func(a, rin, in_place=True)
    else:
        func(a, rin)
    np.testing.assert_equal(rin, rin_orig)


@pytest.mark.parametrize("func", (lattice_track, lattice_pass, internal_lpass))
def test_aperture_outside_limits(rin, func):
    a = elements.Aperture("aperture", [-1e-3, 1e-3, -1e-4, 1e-4])
    assert a.Length == 0
    lattice = [a]
    rin[0, 0] = 1e-2
    rin[2, 0] = -1e-2
    if func == lattice_track:
        func(lattice, rin, in_place=True)
    else:
        func(lattice, rin)
    assert np.isnan(rin[0, 0])
    assert rin[2, 0] == 0.0  # Only the 1st coordinate is nan, the rest is zero


@pytest.mark.parametrize("func", (element_track, element_pass, internal_epass))
def test_drift_offset(rin, func):
    d = elements.Drift("drift", 1)
    rin[0, 0] = 1e-6
    rin[2, 0] = 2e-6
    rin_orig = rin.copy()
    if func == element_track:
        func(d, rin, in_place=True)
    else:
        func(d, rin)
    np.testing.assert_equal(rin, rin_orig)


@pytest.mark.parametrize("func", (element_track, element_pass, internal_epass))
def test_drift_divergence(rin, func):
    d = elements.Drift("drift", 1.0)
    assert d.Length == 1
    rin[1, 0] = 1e-6
    rin[3, 0] = -2e-6
    if func == element_track:
        func(d, rin, in_place=True)
    else:
        func(d, rin)
    # results from Matlab
    rin_expected = np.array([1e-6, 1e-6, -2e-6, -2e-6, 0, 2.5e-12]).reshape(6, 1)
    np.testing.assert_equal(rin, rin_expected)


@pytest.mark.parametrize("func", (element_track, element_pass, internal_epass))
def test_drift_two_particles(rin, func):
    d = elements.Drift("drift", 1.0)
    assert d.Length == 1
    two_rin = np.array(np.concatenate((rin, rin), axis=1), order="F")
    # particle one is offset
    two_rin[0, 0] = 1e-6
    two_rin[2, 0] = 2e-6
    # particle two has divergence
    two_rin[1, 1] = 1e-6
    two_rin[3, 1] = -2e-6
    two_rin_orig = two_rin.copy()
    if func == element_track:
        func(d, two_rin, in_place=True)
    else:
        func(d, two_rin)
    # results from Matlab
    p1_expected = np.array(two_rin_orig[:, 0]).reshape(6, 1)
    p2_expected = np.array([1e-6, 1e-6, -2e-6, -2e-6, 0, 2.5e-12]).reshape(6, 1)
    two_rin_expected = np.concatenate((p1_expected, p2_expected), axis=1)
    np.testing.assert_equal(two_rin, two_rin_expected)


@pytest.mark.parametrize("func", (element_track, element_pass, internal_epass))
def test_quad(rin, func):
    q = elements.Quadrupole("quad", 0.4, k=1)
    rin[0, 0] = 1e-6
    if func == element_track:
        func(q, rin, in_place=True)
    else:
        func(q, rin)
    expected = (
        np.array(
            [0.9210610203854122, -0.3894182419439, 0, 0, 0, 0.0000000103303797478]
        ).reshape(6, 1)
        * 1e-6
    )
    np.testing.assert_allclose(rin, expected)
    assert q.K == 1
    q.PolynomB[1] = 0.2
    assert q.K == 0.2
    q.K = 0.1
    assert q.PolynomB[1] == 0.1


@pytest.mark.parametrize("func", [element_track, element_pass, internal_epass])
def test_rfcavity(rin: np.ndarray, func: any) -> None:
    """Test rf cavity tracking.

    Arguments:
        func: element track method
        rin: np.array of dims (6,1)
    """
    rfc = elements.RFCavity("rfcavity", 0.0, 187500, 3.5237e8, 31, 6.0e9)
    rin[4, 0] = 1e-6
    rin[5, 0] = 1e-6
    if func is element_track:
        func(rfc, rin, in_place=True, energy=+6.0e9)
    else:
        func(rfc, rin, energy=+6.0e9)

    result = np.array([0.0, 0.0, 0.0, 0.0, 9.99769215e-07, 1.0e-6]).reshape(6, 1)

    np.testing.assert_allclose(rin, result, atol=1e-12)


@pytest.mark.parametrize("func", (element_track, element_pass, internal_epass))
@pytest.mark.parametrize("n", (0, 1, 2, 3, 4, 5))
def test_m66(rin, n, func):
    m = np.random.rand(6, 6)
    m66 = elements.M66("m66", m)
    assert m66.Length == 0
    rin[n, 0] = 1e-6
    if func == element_track:
        func(m66, rin, in_place=True)
    else:
        func(m66, rin)
    expected = (
        np.array([m[0, n], m[1, n], m[2, n], m[3, n], m[4, n], m[5, n]]).reshape(6, 1)
        * 1e-6
    )
    np.testing.assert_equal(rin, expected)


@pytest.mark.parametrize("func", (element_track, element_pass, internal_epass))
def test_corrector(rin, func):
    c = elements.Corrector("corrector", 0.0, np.array([0.9, 0.5], dtype=np.float64))
    assert c.Length == 0
    rin[0, 0] = 1e-6
    rin_orig = rin.copy()
    rin_orig[1] = 0.9
    rin_orig[3] = 0.5
    if func == element_track:
        func(c, rin, in_place=True)
    else:
        func(c, rin)
    np.testing.assert_equal(rin, rin_orig)


@pytest.mark.parametrize("func", (element_track, element_pass, internal_epass))
def test_wiggler(rin, func):
    period = 0.05
    periods = 23
    bmax = 1
    by = np.array([1, 1, 0, 1, 1, 0], dtype=np.float64)
    c = elements.Wiggler("wiggler", period * periods, period, bmax, By=by)
    assert abs(c.Length - 1.15) < 1e-10
    # Expected value from Matlab AT.
    expected = np.array(rin, copy=True)
    expected[5] = 0.000000181809691064259
    if func == element_track:
        func(c, rin, energy=3e9, in_place=True)
    else:
        func(c, rin, energy=3e9)
    np.testing.assert_allclose(rin, expected, atol=1e-12)


def test_exit_entrance():
    q = elements.Quadrupole("quad", 0.4, k=1)
    for kin, kout in zip(q._entrance_fields, q._exit_fields):
        np.testing.assert_equal(
            kin.replace("Entrance", "").replace("1", ""),
            kout.replace("Exit", "").replace("2", ""),
        )
