import numpy as np
import pytest

from mbtrack2 import Bunch, Electron, Optics, Synchrotron


@pytest.fixture
def model_ring():
    h = 416  # Harmonic number of the accelerator.
    L = 353.97  # Ring circumference in [m].
    E0 = 2.75e9  # Nominal (total) energy of the ring in [eV].
    particle = Electron()  # Particle considered.
    ac = 1.0695e-5  #1.0695e-4
    U0 = 452.6e3  # Energy loss per turn in [eV].
    tau = np.array([
        7.68e-3, 14.14e-3, 12.18e-3
    ])  #horizontal, vertical and longitudinal damping times in [s].
    tune = np.array([54.2, 18.3])
    emit = np.array([84.4e-12, 84.4e-13])
    sigma_0 = 9e-12
    sigma_delta = 9.07649e-4
    chro = np.array([1.8, 1.3])
    beta = np.array([3.288, 4.003])
    alpha = np.array([0.004, 0.02])
    dispersion = np.array([0.008, 0.01, 0.003, 0.001])
    optics = Optics(local_beta=beta,
                    local_alpha=alpha,
                    local_dispersion=dispersion)
    ring = Synchrotron(h=h,
                       optics=optics,
                       particle=particle,
                       L=L,
                       E0=E0,
                       ac=ac,
                       U0=U0,
                       tau=tau,
                       emit=emit,
                       tune=tune,
                       sigma_delta=sigma_delta,
                       sigma_0=sigma_0,
                       chro=chro)
    return ring


@pytest.fixture
def local_optics():
    beta = np.array([1, 1])
    alpha = np.array([0, 0])
    dispersion = np.array([0, 0, 0, 0])
    local_optics = Optics(local_beta=beta,
                          local_alpha=alpha,
                          local_dispersion=dispersion)
    return local_optics


@pytest.fixture
def demo_ring(local_optics):

    h = 20
    L = 100
    E0 = 1e9
    particle = Electron()
    ac = 1e-3
    U0 = 250e3
    tau = np.array([10e-3, 10e-3, 5e-3])
    tune = np.array([18.2, 10.3])
    emit = np.array([50e-9, 50e-9 * 0.01])
    sigma_0 = 30e-12
    sigma_delta = 1e-3
    chro = [1.0, 1.0]

    ring = Synchrotron(h,
                       local_optics,
                       particle,
                       L=L,
                       E0=E0,
                       ac=ac,
                       U0=U0,
                       tau=tau,
                       emit=emit,
                       tune=tune,
                       sigma_delta=sigma_delta,
                       sigma_0=sigma_0,
                       chro=chro)

    return ring


@pytest.fixture
def mybunch(demo_ring):
    mp_number = 10
    mybunch = Bunch(demo_ring, mp_number=mp_number, track_alive=True)
    return mybunch


@pytest.fixture
def large_bunch(demo_ring):
    mp_number = 1e5
    large_bunch = Bunch(demo_ring, mp_number=mp_number, track_alive=True)
    large_bunch.init_gaussian()
    return large_bunch