import sys
import os
import numpy as np
import synergia
import numpy as np
from mpi4py import MPI
import beamline


lattice = synergia.lattice.Mad8_reader().get_lattice("bad_line", "bad_lattice.lat")

for elem in lattice.get_elements():
    elem.set_string_attribute("extractor_type", "chef_mixed")

macro_particles = 10
real_particles = 1.0e10

commxx = synergia.utils.Commxx()
bunch = synergia.bunch.Bunch(lattice.get_reference_particle(),
            macro_particles, real_particles, commxx)

bunch_simulator = synergia.simulation.Bunch_simulator(bunch)

bunch_simulator.add_per_step(synergia.bunch.Diagnostics_basic("basic.h5"))
bunch_simulator.add_per_step(synergia.bunch.Diagnostics_bulk_track("tracks.h5", macro_particles))

local_particles = bunch.get_local_particles()

stepper = synergia.simulation.Independent_stepper_elements(lattice, 1, 1)
chef_beamline = stepper.get_lattice_simulator().get_chef_lattice().get_sliced_beamline()
probe = beamline.Proton(lattice.get_reference_particle().get_total_energy())
print "chef propagate begin: ", probe.State()
#chef_beamline.propagate(probe)
for ce in chef_beamline:
    ce.propagate(probe)
print "chef propagate end: ", probe.State()

propagator = synergia.simulation.Propagator(stepper)

print

print "bunch particle 0 before synergia propagation: "
print np.array2string(local_particles[0, 0:6])
propagator.propagate(bunch_simulator, 10, 1, 1)
print
print "bunch particle 0 after synergia propagation: "
print np.array2string(local_particles[0, 0:6])

sys.exit(0)
print
print "beamline"
print synergia.lattice.chef_beamline_as_string(chef_beamline)