from at import physics, get_s_pos, uint32_refpts
from scipy.constants import speed_of_light
import matlab
import numpy
from numpy.testing import assert_allclose as assert_close
import pytest


def _py_data(ml_data):
    """Convert a Matlab vector to a numpy vector"""
    return numpy.squeeze(numpy.asarray(ml_data))


def _ml_refs(refpts, nelems):
    """Convert refpoints to Matlab"""
    uintrefs = uint32_refpts(refpts, nelems)
    return matlab.double([float(ref+1) for ref in uintrefs])


def _compare_physdata(py_data, ml_data, fields, atol=12, rtol=1.e-7):
    for (ml_key, py_key) in fields:
        ml_val = _py_data(ml_data[ml_key])
        if py_key == 'closed_orbit':
            py_val = numpy.squeeze(py_data[py_key][:, :4])
        else:
            py_val = numpy.squeeze(py_data[py_key])
        assert_close(py_val, ml_val, atol=atol, rtol=rtol)


@pytest.mark.parametrize('dp', (-0.01, 0.0, 0.01))
@pytest.mark.parametrize('lattices',
                         [pytest.lazy_fixture('hmba'),
                          pytest.lazy_fixture('dba')])
def test_linear_analysis(engine, lattices, dp):
    py_lattice, ml_lattice, _ = lattices
    nelems = len(py_lattice)
    fields = [('SPos', 's_pos'),
              ('ClosedOrbit', 'closed_orbit'),
              ('Dispersion', 'dispersion'),
              ('alpha', 'alpha'), ('beta', 'beta'),
              ('mu', 'mu'), ('M44', 'm44'),
              ('A', 'A'), ('B', 'B'), ('C', 'C'),
              ('gamma', 'gamma')]
    refpts = range(nelems + 1)
    py_data0, py_tune, py_chrom, py_data = physics.linopt(py_lattice, dp,
                                                          refpts, True,
                                                          ddp=1.E-6)
    # Matlab call
    ml_data, ml_tune, ml_chrom = engine.pyproxy('atlinopt', ml_lattice, dp,
                                                _ml_refs(refpts, nelems),
                                                nargout=3)
    ml_data0 = engine.pyproxy('atlinopt', ml_lattice, dp,
                              _ml_refs(nelems, nelems), nargout=3)[0]
    # Comparison
    assert_close(py_tune, _py_data(ml_tune), atol=1.e-8, rtol=0)
    assert_close(py_chrom, _py_data(ml_chrom), atol=3.e-3, rtol=0)
    _compare_physdata(numpy.expand_dims(py_data0, 0), ml_data0, fields,
                      rtol=1e-8)
    _compare_physdata(py_data, ml_data, fields, rtol=1e-8)


@pytest.mark.parametrize('lattices',
                         [pytest.lazy_fixture('hmba')])
def test_ohmi_envelope(engine, lattices):
    py_lattice, ml_lattice, _ = lattices
    fields = [('beam66', 'r66'), ('beam44', 'r44'), ('emit66', 'emitXYZ'),
              ('emit44', 'emitXY')]
    nelems = len(py_lattice)
    refpts = range(nelems + 1)

    # Python call
    py_lattice = py_lattice.radiation_on(copy=True)
    py_emit0, py_beamdata, py_emit = py_lattice.ohmi_envelope(refpts)
    # Matlab call
    ml_emit = engine.pyproxy('atx', ml_lattice, 0.0, _ml_refs(refpts, nelems))
    ml_emit0, ml_params = engine.pyproxy('atx', ml_lattice, 0.0,
                                         _ml_refs(0, nelems), nargout=2)
    revolution_period = get_s_pos(py_lattice, nelems) / speed_of_light
    damping_times = revolution_period / py_beamdata.damping_rates
    # Comparison
    assert_close(damping_times, _py_data(ml_params['dampingtime']),
                 atol=1.e-6, rtol=1.e-3)
    assert_close(py_beamdata.mode_emittances, _py_data(ml_emit0['modemit']),
                 atol=1.e-20, rtol=2.e-3)
    _compare_physdata(numpy.expand_dims(py_emit0, 0), ml_emit0, fields)
    _compare_physdata(py_emit, ml_emit, fields)


@pytest.mark.parametrize('lattices',
                         [pytest.lazy_fixture('hmba_rad')])
def test_quantdiff(engine, lattices):
    py_lattice, ml_lattice, radindex = lattices
    # Python call
    dmat = physics.radiation.quantdiffmat(py_lattice)
    # Matlab call
    dmat_ml = engine.pyproxy('quantumDiff', ml_lattice, radindex)
    # Comparison
    assert_close(dmat, dmat_ml, rtol=1.0e-8, atol=1.0e-18)


@pytest.mark.parametrize('lattices',
                         [pytest.lazy_fixture('hmba')])
def test_fastring(engine, lattices):
    py_lattice, ml_lattice, _ = lattices
    # Python call
    ring, ringrad = physics.fastring.fast_ring(py_lattice)
    # Matlab call
    ring_ml, ringrad_ml = engine.pyproxy('atfastring',
                                         ml_lattice, nargout=2)
    # Comparison
    for r, rml in zip([ring, ringrad], [ring_ml, ringrad_ml]):
        assert_close(r[0].Frequency,    rml[1]['Frequency'], rtol=1.0e-20)
        assert_close(r[0].Voltage,      rml[1]['Voltage'], rtol=1.0e-20)
        assert_close(r[1].I2,           rml[2]['I2'], rtol=1.0e-20)
        assert_close(r[1].Length,       rml[2]['Length'], rtol=1.0e-20)
        assert_close(r[1].M66,          rml[2]['M66'], atol=1.0e-7)
        assert_close(r[1].T1,   numpy.squeeze(rml[2]['T1']),
                     rtol=1.0e-8, atol=1.0e-11)
        assert_close(r[1].T2,   numpy.squeeze(rml[2]['T2']),
                     rtol=1.0e-8, atol=1.0e-11)
        assert_close(r[2].A1,       rml[-1]['A1'], rtol=0.02)
        assert_close(r[2].A2,       rml[-1]['A2'], rtol=0.02)
        assert_close(r[2].A3,       rml[-1]['A3'], rtol=0.02)
        assert_close(r[2].Alphax,   rml[-1]['Alphax'], rtol=1.0e-5)
        assert_close(r[2].Alphay,   rml[-1]['Alphay'], rtol=1.0e-6)
        assert_close(r[2].Betax,    rml[-1]['Betax'], rtol=1.0e-10)
        assert_close(r[2].Betay,    rml[-1]['Betay'], rtol=1.0e-10)
        assert_close(r[2].Qpx,      rml[-1]['Qpx'], rtol=1.0e-5)
        assert_close(r[2].Qpy,      rml[-1]['Qpy'], rtol=1.0e-5)
        assert_close(r[2].T1,   numpy.squeeze(rml[-1]['T1']),
                     rtol=1.0e-8, atol=1.0e-11)
        assert_close(r[2].T2,   numpy.squeeze(rml[-1]['T2']),
                     rtol=1.0e-8, atol=1.0e-11)
        if len(r) >= 4:
            assert_close(r[3].Lmatp,    rml[-2]['Lmatp'],
                         rtol=0.02, atol=1.0e-10)


@pytest.mark.parametrize('dp', (0.00, 0.01, -0.01))
@pytest.mark.parametrize('lattices',
                         [pytest.lazy_fixture('hmba')])
def test_parameters(engine, lattices, dp):
    py_lattice, ml_lattice, _ = lattices

    # Test perimeter
    py_length = py_lattice.get_s_pos(len(py_lattice))
    ml_length = engine.findspos(ml_lattice, len(ml_lattice)+1)
    assert_close(py_length, ml_length, rtol=1.E-8)

    # test energy loss
    ml_energy, ml_periods, ml_voltage, \
        ml_harms, ml_eloss = engine.pyproxy('atenergy', ml_lattice, nargout=5)
    assert_close(py_lattice.rf_voltage, ml_voltage, rtol=1.E-8)
    assert_close(py_lattice.energy_loss, ml_eloss, rtol=1.E-6)
    assert py_lattice.energy == ml_energy
    assert py_lattice.periodicity == int(ml_periods)
    assert py_lattice.harmonic_number == int(ml_harms)

    # test momentum compaction factor
    py_mcf = py_lattice.get_mcf(dp)
    ml_mcf = engine.mcf(ml_lattice, dp)
    assert_close(py_mcf, ml_mcf, rtol=1.E-8)
