#!/usr/bin/env python
import sys
import os
if os.environ.has_key("SYNERGIA_DIR"):
    sys.path.append(os.environ["SYNERGIA_DIR"])
else:
    sys.path.append("..")
                    
import impact_parameters
import impact_elements
import synergia
import matching
import options
import math
import madcalc

myopts = options.Options("booster_test3")
myopts.add("numturns",10,"Number of turns",int)
myopts.add("xwidth",0.004,"horizontal width (m)",float)
myopts.add("dpop",3.0e-4,"(delta p)/p",float)

myopts.add_suboptions(impact_parameters.options)
myopts.add_suboptions(synergia.options)
myopts.parse_argv(sys.argv)


ip = impact_parameters.Impact_parameters(impact_parameters.options
                                         .get_value("sampleperiod"))
ip.apply_options(impact_parameters.options)

pz = ip.gamma() * ip.beta()*ip.mass_GeV

#booster_sliced.mad has injection located where it should be
madfile = "booster_sliced.mad"
(D_x, D_y) = madcalc.dispersion_initial(madfile,"bcelinj",
                                        myopts.get_value("energy"))
(alpha_x, beta_x, alpha_y, beta_y) = madcalc.twiss_initial(madfile,
                                                           "bcelinj")

width_x = myopts.get_value("xwidth")
(width_xprime,r_x,emittance) = matching.match_twiss_width(width_x,
                                                          alpha_x,beta_x)

ip.x_params(sigma = width_x, lam = width_xprime * pz)

(width_y,width_yprime,r_y) = matching.match_twiss_emittance(emittance,
                                                            alpha_y, beta_y)

ip.y_params(sigma = width_y, lam = width_yprime * pz)

dp_o_p = myopts.get_value("dpop")
ip.z_params(sigma = 0.10, lam = dp_o_p* pz)
R0 = -D_x*dp_o_p/math.sqrt(0.5*(D_x*dp_o_p)**2 + emittance*beta_x)

ip.correlation_coeffs(xpx = r_x, ypy = r_y, xpz = R0);

ipradius = impact_parameters.options.get_value("pipedims")[0]
sample_period = impact_parameters.options.get_value("sampleperiod")
booster = impact_elements.External_element(kicks=100, steps=1,
                                           radius=ipradius,
                                           mad_file_name=madfile)
numturns = myopts.get_value("numturns")
for turn in range(1,numturns+1):
    ip.add(booster)
    if turn != numturns:
        ip.add(impact_elements.Output_element("turn%02d.dat" % turn,
                                              sample_period))

my_synergia = synergia.Synergia(ip,sys.argv,synergia.options)

my_synergia.prepare_run("bst3")