import numpy as np
import pytest
from scipy.constants import e

from mbtrack2 import Bunch


def test_bunch_values(demo_ring):
    mp_number = 10
    current = 20e-3
    mybunch = Bunch(demo_ring, mp_number=mp_number, current=current,
                    track_alive=True)

    assert mybunch.mp_number == mp_number
    assert pytest.approx(mybunch.current) == current
    assert len(mybunch) == mp_number
    np.testing.assert_allclose(mybunch.alive, np.ones((mp_number,), dtype=bool))
    assert pytest.approx(mybunch.charge) == current * demo_ring.T0
    assert pytest.approx(mybunch.charge_per_mp) == current * demo_ring.T0 / mp_number
    assert pytest.approx(mybunch.particle_number) == current * demo_ring.T0 / e
    assert mybunch.is_empty == False

def test_bunch_magic(mybunch):
    for label in mybunch:
        np.testing.assert_allclose(mybunch[label], np.zeros(len(mybunch)))
        mybunch[label] = np.ones(len(mybunch))
        np.testing.assert_allclose(mybunch[label], np.ones(len(mybunch)))

def test_bunch_losses(mybunch):
    charge_init = mybunch.charge
    mybunch.alive[0] = False
    assert len(mybunch) == mybunch.mp_number - 1
    assert pytest.approx(mybunch.charge) == charge_init * len(mybunch) / mybunch.mp_number

def test_bunch_init_gauss(large_bunch):
    large_bunch.init_gaussian(mean=np.ones((6,)))
    np.testing.assert_allclose(large_bunch.mean, np.ones((6,)), rtol=1e-2)

def test_bunch_save_load(mybunch, demo_ring, tmp_path):
    mybunch["x"] += 1
    mybunch.save(str(tmp_path / "test"))

    mybunch2 = Bunch(demo_ring, mp_number=1, current=1e-5)
    mybunch2.load(str(tmp_path / "test.hdf5"))

    assert mybunch.mp_number == mybunch2.mp_number
    assert pytest.approx(mybunch.charge) == mybunch2.charge
    for label in mybunch:
        np.testing.assert_allclose(mybunch[label], mybunch2[label])

def test_bunch_stats(demo_ring, large_bunch):
    large_bunch.init_gaussian()
    np.testing.assert_array_almost_equal(large_bunch.mean, np.zeros((6,)), decimal=5)
    sig = np.concatenate((demo_ring.sigma(), [demo_ring.sigma_0, demo_ring.sigma_delta]))
    np.testing.assert_allclose(large_bunch.std, sig, rtol=1e-2)
    np.testing.assert_allclose(large_bunch.emit[:2], demo_ring.emit, rtol=1e-2)
    np.testing.assert_allclose(large_bunch.cs_invariant[:2], demo_ring.emit*2, rtol=1e-2)

def test_bunch_binning(mybunch):
    mybunch.init_gaussian()
    (bins, sorted_index, profile, center) = mybunch.binning()
    profile0 = np.zeros((len(bins)-1,))
    for i, val in enumerate(sorted_index):
        assert bins[val] <= mybunch["tau"][i] <= bins[val+1]
        profile0[val] += 1
    np.testing.assert_allclose(profile0, profile)

def test_bunch_plots(mybunch):
    mybunch.init_gaussian()
    mybunch.plot_phasespace()
    mybunch.plot_profile()
    assert True
def test_bunch_emittance(demo_ring):
    mp_number = 1_000_000
    current = 1.2e-3
    mybunch = Bunch(demo_ring, mp_number=mp_number, current=current,
                    track_alive=False)
    mybunch.init_gaussian()
    np.testing.assert_allclose(mybunch.emit[0], demo_ring.emit[0], rtol=1e-2, atol=0,
     err_msg=f'Emittances do not match. {demo_ring.emit[0]} initialised, {mybunch.emit[0]:} calculated')
    np.testing.assert_allclose(mybunch.emit[1], demo_ring.emit[1], rtol=1e-2, atol=0,
     err_msg=f'Emittances do not match. {demo_ring.emit[1]} initialised, {mybunch.emit[1]:} calculated')

    np.testing.assert_allclose(mybunch.emit[0], mybunch.cs_invariant[0]/2, rtol=1e-2, atol=0,
     err_msg=f'Emittances do not match. {mybunch.cs_invariant[0]/2} calculated with optics functions, {mybunch.emit[0]:} calculated with coordinates only')
    np.testing.assert_allclose(mybunch.emit[1], mybunch.cs_invariant[1]/2, rtol=1e-2, atol=0,
     err_msg=f'Emittances do not match. {mybunch.cs_invariant[1]/2} calculated with optics functions, {mybunch.emit[1]:} calculated with coordinates only')


def test_bunch_emittance_with_dispersion(demo_ring):
    mp_number = 1_000_000
    current = 1.2e-3
    demo_ring.optics.local_dispersion = np.array([1e-2, 1e-3, 1e-2, 1e-3])
    mybunch = Bunch(demo_ring, mp_number=mp_number, current=current,
                    track_alive=False)
    mybunch.init_gaussian()
    np.testing.assert_allclose(mybunch.emit[0], demo_ring.emit[0], rtol=1e-2, atol=0,
     err_msg=f'Emittances do not match. {demo_ring.emit[0]} initialised, {mybunch.emit[0]:} calculated')
    np.testing.assert_allclose(mybunch.emit[1], demo_ring.emit[1], rtol=1e-2, atol=0,
     err_msg=f'Emittances do not match. {demo_ring.emit[1]} initialised, {mybunch.emit[1]:} calculated')

    np.testing.assert_allclose(mybunch.emit[0], mybunch.cs_invariant[0]/2, rtol=1e-2, atol=0,
     err_msg=f'Emittances do not match. {mybunch.cs_invariant[0]/2} calculated with optics functions, {mybunch.emit[0]:} calculated with coordinates only')
    np.testing.assert_allclose(mybunch.emit[1], mybunch.cs_invariant[1]/2, rtol=1e-2, atol=0,
     err_msg=f'Emittances do not match. {mybunch.cs_invariant[1]/2} calculated with optics functions, {mybunch.emit[1]:} calculated with coordinates only')