import numpy as np
import pytest

from mbtrack2 import Bunch, Electron, Optics, Synchrotron


@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