#!/usr/bin/env python

import local_paths

from basic_toolkit import *
from mxyzptlk import *
from beamline import *
from physics_toolkit import *
from bmlfactory import *
from physics_constants import *
import mappers

import math
import sys
import Numeric
import string

class Action:
    def __init__(self,type,length,initial_energy,final_energy,data,
                 synergia_action=None):
        self.data = data
        self.type = type
        self.length = length
        self.initial_energy = initial_energy
        self.final_energy = final_energy
        self.synergia_action = synergia_action
    def get_type(self):
        return self.type
    def is_mapping(self):
        return self.type == "mapping"
    def is_synergia_action(self):
        return self.type == "synergia action"
    def get_data(self):
        return self.data
    def get_length(self):
        return self.length
    def get_initial_energy(self):
        return self.initial_energy
    def get_final_energy(self):
        return self.final_energy
    def get_synergia_action(self):
        return self.synergia_action

accuracy_marker = marker("accuracy")
space_charge_marker = marker("space charge")
pacifier = drift("pacifier",0.0)

class Gourmet:
    def __init__(self, mad_file, line_name, initial_kinetic_energy,
                 scaling_frequency, order=1, particle='proton'):
        self.mad_file = mad_file
        self.line_name = line_name
        self.scaling_frequency = scaling_frequency
        self.order = order
        self.saved_elements = []
        createStandardEnvironments(self.order)

        self.particle = particle
        if self.particle == 'proton':
            self.mass = PH_NORM_mp
        elif self.particle == 'positron':
            self.mass = PH_NORM_me
        else:
            raise RuntimeError, 'Unknown particle %s' % self.particle
        self.initial_kinetic_energy = initial_kinetic_energy
        self.initial_energy = self.initial_kinetic_energy + self.mass
        self.final_energy = None
        self.initial_momentum = math.sqrt(self.initial_energy**2 -
                                          self.mass**2)

        self.factory = bmlfactory(mad_file)
        brho = self.get_initial_particle().ReferenceBRho()
        beamline_orig = self.factory.create_beamline(line_name,brho)
        beamline_orig.flatten()
        beamline_orig.insert(pacifier)
        beamline_orig.append(pacifier)
        self.beamline = DriftsToSlots(beamline_orig)
        self.have_actions = 0
        self.have_fast_mappings = 0
        self.context = BeamlineContext(self.get_initial_particle(),
                                       self.beamline,0)
        if not self.context.isTreatedAsRing():
            self.context.handleAsRing()

        self.iterator = DeepBeamlineIterator(self.beamline)

    def get_mad_file(self):
        return self.mad_file

    def get_line_name(self):
        return self.line_name

    def get_initial_energy(self):
        return self.initial_energy
    
    def get_initial_particle(self):
        if self.particle == 'proton':
            particle = Proton(self.initial_energy)
        else:
            particle = Positron(self.initial_energy)
        return particle

    def get_initial_jet_particle(self):
        if self.particle == 'proton':
            jet_particle = JetProton(self.initial_energy)
        else:
            jet_particle = JetPositron(self.initial_energy)
        return jet_particle

    def get_particle(self,energy):
        if self.particle == 'proton':
            particle = Proton(energy)
        else:
            particle = Positron(energy)
        return particle

    def get_jet_particle(self,energy):
        if self.particle == 'proton':
            jet_particle = JetProton(energy)
        else:
            jet_particle = JetPositron(energy)
        return jet_particle
    
    def _commission(self):
### The following is in reaction to the message:
### *** WARNING ***
### *** WARNING *** File: CF_sbend.cc, Line: 536
### *** WARNING *** void sbend::Split( double pc, bmlnElmnt** a, bmlnElmnt** b )
### *** WARNING *** The new, split elements must be commissioned with
### *** WARNING *** RefRegVisitor before being used.
### *** WARNING ***
### generated by insertions.
        compulsory_visitor = RefRegVisitor(self.get_initial_particle())
        self.beamline.accept(compulsory_visitor)
        
    def insert_accuracy_markers(self, num_markers_per_element):
        master_insertion_point = 0.0
        insertion_list = InsertionList(self.initial_momentum)
        self.iterator.reset()
        element = self.iterator.next()
        ile_list = []
        particle = self.get_initial_particle()
        while element:
                if element.OrbitLength(particle) > 0:
                        marker_interval = element.OrbitLength(particle)/ \
                        (num_markers_per_element + 1.0)
                        insertion_point = master_insertion_point
                        for i in range(0,num_markers_per_element):
                                insertion_point += marker_interval
                                ile = InsertionListElement(insertion_point,
                                                           accuracy_marker)
                                ile_list.append(ile)
                                insertion_list.Append(ile)
                master_insertion_point += element.OrbitLength(particle)
                element.propagateParticle(particle)
                element = self.iterator.next()
        removed_elements = slist()
        s_0 = 0.0
        self.beamline.InsertElementsFromList(s_0, insertion_list,
                                             removed_elements)
        self.beamline.append(accuracy_marker)
        self._commission()
        
    def orbit_length(self):
        return self.beamline.OrbitLength(self.get_initial_particle())
    
    def insert_elements(self,elements,positions):
        if type(elements) != type([]) and type(elements) != type(()):
            elements = (elements,)
            positions = (positions,)
        ile_list = []
        insertion_list = InsertionList(self.initial_momentum)
        index = 0
        elements_to_insert = 0
        for element in elements:
            self.saved_elements.append(element)
            position = positions[index]
            if position > 0.0:
                ile = InsertionListElement(position,element)
                ile_list.append(ile)
                insertion_list.Append(ile)
                elements_to_insert += 1
            else:
                self.beamline.insert(element)
            index += 1
        if elements_to_insert > 0:
            removed_elements = slist()
            s_0 = 0.0
            self.beamline.InsertElementsFromList(s_0, insertion_list,
                                                 removed_elements)
        self._commission()
    
    def insert_space_charge_markers(self, num_kicks):
        kick_interval = self.orbit_length()/(1.0*num_kicks)
        elements = []
        positions = []
        insertion_point = 0.0
        for i in range(0,num_kicks):
            elements.append(marker("synergia action:space charge endpoint"))
            positions.append(insertion_point)
            insertion_point += kick_interval/2.0
            elements.append(marker("synergia action:space charge kick"))
            positions.append(insertion_point)
            insertion_point += kick_interval/2.0
        end_marker = marker("synergia action:space charge endpoint")
        self.saved_elements.append(end_marker)
        self.beamline.append(end_marker)
        self.insert_elements(elements,positions)

    def print_elements(self):
        self.iterator.reset()
        i = 0
        element = self.iterator.next()
        while element:
            print i, element.Name(), element.Type()
            element = self.iterator.next()
            i += 1

    def get_lattice_functions(self):
        lattice_function_array = Lattice_function_array()
        self.iterator.reset()
        i = 0
        element = self.iterator.next()
        while element:
            lattice_function = self.context.getLattFuncPtr(i)
            if lattice_function:
                lattice_function_array.append(lattice_function)
            element = self.iterator.next()
            i += 1
        return lattice_function_array

    def generate_actions(self):
        self.delete_actions()
        self.iterator.reset()
        element = self.iterator.next()
        s = 0.0
        particle = self.get_initial_particle()
        jet_particle = self.get_initial_jet_particle()
        has_propagated = 0
        energy = self.initial_energy
        while element:
            split_name = element.Name().split(":")
            if split_name[0] == "synergia action":
                if has_propagated:
                    mapping = mappers.Fast_mapping(self.get_u(energy),
                                                   jet_particle.State())
                    new_energy = jet_particle.ReferenceEnergy()
                    self.actions.append(
                        Action("mapping",
                               length=s,
                               initial_energy=energy,
                               final_energy=new_energy,
                               data=mapping))
                    energy = new_energy
                    s = element.OrbitLength(particle)
                    element.propagateParticle(particle)
                    element.propagateJetParticle(jet_particle)
                new_energy = jet_particle.ReferenceEnergy()
                self.actions.append(
                    Action("synergia action",
                           length=s,
                           initial_energy=energy,
                           final_energy=new_energy,
                           data=element,
                           synergia_action=string.join(split_name[1:],':')))
                energy = new_energy           
                jet_particle = self.get_jet_particle(energy)
                has_propagated = 0
                s = 0.0
            else:
                if not element.Type() == "marker":
                    s += element.OrbitLength(particle)
                    element.propagateParticle(particle)
                    element.propagateJetParticle(jet_particle)
                    has_propagated = 1
            element = self.iterator.next()
        self.final_energy = jet_particle.ReferenceEnergy()
        self.have_actions = 1

    def get_actions(self):
        if not self.have_actions:
            self.generate_actions()
        return self.actions

    def print_actions(self):
        s = 0.0
        for action in self.get_actions():
            print "s:%10.4f" % s,
            print "%s," % action.get_type(),
            print "length = %0.4g," % action.get_length(),
            if action.is_synergia_action():
                print action.get_synergia_action(),
            deltaE = action.get_final_energy() - action.get_initial_energy()
            if deltaE != 0.0:
                print "deltaE = %0.4g" % deltaE,
            print
            s += action.get_length()
            
    def get_strengths(self):
        self.iterator.reset()
        element = self.iterator.next()
        particle = self.get_initial_particle()
        brho = particle.ReferenceBRho()
        kxs = []
        kys = []
        ss = []
        s = 0.0
        while element:
            kx = 0.0
            ky = 0.0
            strength = element.Strength()
            length = element.OrbitLength(particle)
            s += length
            if (element.Type() == "quadrupole") or \
                   (element.Type() == "thinQuad"):
                kx = strength/brho;
                ky = -kx
            elif (element.Type() == "CF_sbend") or \
                 (element.Type() == "CF_rbend"):
                ky = -element.getQuadrupole()/brho/length
                kx = -ky + strength**2/(brho**2) 
            else:
                if (strength != 0.0):
                    sys.stderr.write("gourmet.get_strengths error. element of type %s has non-zero\nStrength(), but is not handled by this routine.\n" % element.Type())
            kxs.append(kx)
            kys.append(ky)
            ss.append(s)
            element.propagateParticle(particle)
            element = self.iterator.next()
        return (Numeric.array(ss),Numeric.array(kxs),Numeric.array(kys))
    
    def delete_actions(self):
        self.actions = []
        self.have_actions = 0

    def _convert_linear_maps(self, chef_linear_maps):
        # units conversion
        # X_impact = U X_external
        # where U = diag(u[0],u[1],u[2],u[3],u[4],u[5])
        u = self.get_initial_u()
        linear_maps = []
        for chef_map in chef_linear_maps:
            map = Numeric.zeros((7,7),'d')
            for row in range(0,6):
                for column in range(0,6):
                    chef_row = int(row/2+3*(row%2))
                    chef_column = int(column/2+3*(column%2))
                    map[row,column] = chef_map.get(chef_row,chef_column)* \
                                      u[row]/u[column]
            map[6,6] = 1.0
            linear_maps.append(map)
        return linear_maps
    
    def generate_fast_mappings(self):
        self.delete_fast_mappings()
        if not self.have_actions:
            self.generate_actions()
        u = self.get_initial_u()
        for action in self.actions:
            if action.is_mapping():
                self.fast_mappings.append(action.get_data())
        self.have_fast_mappings = 1

    def delete_fast_mappings(self):
        self.fast_mappings = []
        self.have_fast_mappings = 0

    def get_fast_mapping(self, index):
        if not self.have_fast_mappings:
            self.generate_fast_mappings()
        return self.fast_mappings[index]
    
    def get_single_linear_map(self):
        jet_particle = self.get_initial_jet_particle()
        self.beamline.propagateJetParticle(jet_particle)
        return self._convert_linear_maps([jet_particle.State().jacobian()])[0]

    def get_single_fast_map(self):
        jet_particle = self.get_initial_jet_particle()
        self.beamline.propagateJetParticle(jet_particle)
        mapping = jet_particle.State()
        return mappers.Fast_mapping(self.get_initial_u(),
                                    jet_particle.State())
    def printpart(self,particle):
        print '|',"%0.5g" % particle.get_x(),
        print "%0.5g" % particle.get_npx(),
        print "%0.5g" % particle.get_y(),
        print "%0.5g" % particle.get_npy(),
        print "%0.5g" % particle.get_cdt(),
        print "%0.5g" % particle.get_ndp(), '|'
        
    def check(self,print_coeffs=0):
        jet_particle = self.get_initial_jet_particle()
        self.beamline.propagateJetParticle(jet_particle)
        mapping = jet_particle.State()
        if print_coeffs:
            mapping.printCoeffs()
        testpart = self.get_initial_particle()
        print "initial test particle:"
        self.printpart(testpart)
        self.beamline.propagateParticle(testpart)
        print "test particle after propagation:"
        self.printpart(testpart)

    def get_u(self,energy):
        gamma = energy/self.mass
        beta = math.sqrt(1.0 - 1.0/(gamma*gamma))
        c = PH_MKS_c
        w = 2.0* math.pi* self.scaling_frequency
        u = [w/c,gamma*beta,w/c,gamma*beta,w/c,-gamma*beta*beta]
        return Numeric.array(u)

    def get_initial_u(self):
        return self.get_u(self.initial_energy)
        
    def get_initial_brho(self):
        return self.get_initial_particle().ReferenceBRho()
    
class Lattice_function_array:
    def __init__(self):
        self.s = []
        self.beta_x = []
        self.beta_y = []
        self.alpha_x = []
        self.alpha_y = []

    def append(self, lattice_fn):
        self.s.append(lattice_fn.arcLength)
        self.beta_x.append(lattice_fn.beta.hor)
        self.beta_y.append(lattice_fn.beta.ver)
        self.alpha_x.append(lattice_fn.alpha.hor)
        self.alpha_y.append(lattice_fn.alpha.ver)