#!/usr/bin/env python3

#This file imports the members defined in constants.h into a separated python module
# (from the "flat" module CLICopti.CLICopti generated by SWIG)

#Structures that are returned from various functions
from CLICopti.CLICopti import \
    return_AccelStructure_getMaxFields, \
    return_AccelStructure_getMaxDeltaT, \
    return_AccelStructure_getMaxAllowableBeamTime_detailed

#The main accelerator structure base class
from CLICopti.CLICopti import AccelStructure
#Subclasses implementing constructors and methods
# createCells(), writeProfileFile_header
from CLICopti.CLICopti import \
    AccelStructure_paramSet1,\
    AccelStructure_paramSet2,\
    AccelStructure_paramSet2_noPsi,\
    AccelStructure_CLIC502,\
    AccelStructure_CLICG,\
    AccelStructure_general

#Exceptions
# (not actually picked up by the python wrapper...)
#from CLICopti.CLICopti import AccelStructureUninitialized

#Patch functions to be pythonesque, similar to CellBase constructors
import numpy as np

#Gradient profile and basic z-axis stuff

def _AccelStructure_getZ_all_PY(self):
    N = self.getZNumpoints()
    z = np.empty(N)
    for i in range(N):
        z[i] = self.getZ(i)
    return z
AccelStructure.getZ_all = _AccelStructure_getZ_all_PY

def _AccelStructure_getEz_unloaded_all_PY(self,peakPower):
    N = self.getZNumpoints()
    z = np.empty(N)
    for i in range(N):
        z[i] = self.getEz_unloaded(i,peakPower)
    return z
AccelStructure.getEz_unloaded_all = _AccelStructure_getEz_unloaded_all_PY

def _AccelStructure_getEz_loaded_all_PY(self,peakPower,current):
    N = self.getZNumpoints()
    Ez = np.empty(N)
    for i in range(N):
        Ez[i] = self.getEz_loaded(i,peakPower,current)
    return Ez
AccelStructure.getEz_loaded_all = _AccelStructure_getEz_loaded_all_PY

def _AccelStructure_getDeltaT_all_PY(self, peakPower, t_beam,  beamCurrent, loaded=False):
    N = self.getZNumpoints()
    dT = np.empty(N)
    deltaTconst = self.getDeltaTconst(peakPower, t_beam, beamCurrent)
    for i in range(N):
        dT[i] = self.getDeltaT(peakPower,t_beam,beamCurrent,i, loaded,deltaTconst)
    return dT
AccelStructure.getDeltaT_all = _AccelStructure_getDeltaT_all_PY

#TODO:
# Alternatives for writeDeltaTprofileFile and writeTimeDeltaTprofileFile ?

#Wrap functions so that numpy ndarrays are also accepted
# Would be nice to just use np.frompyufunc here, but not sure how to handle "self".

# Wake potential variants

_AccelStructure_getTransverseWakePotential_CPP = AccelStructure.getTransverseWakePotential
def _AccelStructure_getTransverseWakePotential_PY(self,z):
    if self.needWakePrecalc():
        self.populateWakePrecalc()

    if isinstance(z,np.ndarray):
        if z.ndim != 1:
            raise TypeError("Only 1D numpy arrays are supported for z")
        r = np.empty_like(z)
        for i in range(len(r)):
           r[i] = _AccelStructure_getTransverseWakePotential_CPP(self,z[i])
        return r
    #Else: Scalar
    return _AccelStructure_getTransverseWakePotential_CPP(self,z)
AccelStructure.getTransverseWakePotential = _AccelStructure_getTransverseWakePotential_PY

_AccelStructure_getTransverseWakePotentialEnvelope_CPP = AccelStructure.getTransverseWakePotentialEnvelope
def _AccelStructure_getTransverseWakePotentialEnvelope_PY(self,z):
    if self.needWakePrecalc():
        self.populateWakePrecalc()

    if isinstance(z,np.ndarray):
        if z.ndim != 1:
            raise TypeError("Only 1D numpy arrays are supported for z")
        r = np.empty_like(z)
        for i in range(len(r)):
           r[i] = _AccelStructure_getTransverseWakePotentialEnvelope_CPP(self,z[i])
        return r
    #Else: Scalar
    return _AccelStructure_getTransverseWakePotentialEnvelope_CPP(self,z)
AccelStructure.getTransverseWakePotentialEnvelope = _AccelStructure_getTransverseWakePotentialEnvelope_PY

_AccelStructure_getTransverseWakePotentialEnvelope_detuning_CPP = AccelStructure.getTransverseWakePotentialEnvelope_detuning
def _AccelStructure_getTransverseWakePotentialEnvelope_detuning_PY(self,z):
    if self.needWakePrecalc():
        self.populateWakePrecalc()

    if isinstance(z,np.ndarray):
        if z.ndim != 1:
            raise TypeError("Only 1D numpy arrays are supported for z")
        r = np.empty_like(z)
        for i in range(len(r)):
           r[i] = _AccelStructure_getTransverseWakePotentialEnvelope_detuning_CPP(self,z[i])
        return r
    #Else: Scalar
    return _AccelStructure_getTransverseWakePotentialEnvelope_detuning_CPP(self,z)
AccelStructure.getTransverseWakePotentialEnvelope_detuning = _AccelStructure_getTransverseWakePotentialEnvelope_detuning_PY

# Pulse shape
_AccelStructure_getP_t_CPP = AccelStructure.getP_t
def _AccelStructure_getP_t_PY(self,t,peakPower,t_beam,beamCurrent):
    if isinstance(t,np.ndarray):
        if t.ndim != 1:
            raise TypeError("Only 1D numpy arrays are supported for z")

        breakoverPower = self.getBreakoverPower(peakPower,beamCurrent)

        P = np.empty_like(t)
        for (iT, T) in enumerate(t):
            P[iT] = _AccelStructure_getP_t_CPP(self, T, peakPower,\
                                               t_beam, beamCurrent, breakoverPower)
        return P
    #Else: Scalar
    return _AccelStructure_getP_t_CPP(self,t,peakPower,t_beam,beamCurrent)
AccelStructure.getP_t = _AccelStructure_getP_t_PY

## Interpolation along structure, in various ways

_AccelStructure_getInterpolated_CPP = AccelStructure.getInterpolated
def _AccelStructure_getInterpolated_PY(self, z, fieldName, midEnds=False):
    import CLICopti.CellParams
    fieldOffset = CLICopti.CellParams.cellOffsets[fieldName]

    if isinstance(z,np.ndarray):
        if z.ndim != 1:
            raise TypeError("Only 1D numpy arrays are supported for z")
        r = np.empty_like(z)
        for i in range(len(r)):
           r[i] = _AccelStructure_getInterpolated_CPP(self,z[i],fieldOffset,midEnds)
        return r
    #Else: Scalar
    return _AccelStructure_getInterpolated_CPP(self,z,fieldOffset,midEnds)
AccelStructure.getInterpolated = _AccelStructure_getInterpolated_PY

_AccelStructure_getInterpolated_zidx_CPP = AccelStructure.getInterpolated_zidx
def _AccelStructure_getInterpolated_zidx_PY(self, zIdx, fieldName : str):
    import CLICopti.CellParams
    fieldOffset =  CLICopti.CellParams.cellOffsets[fieldName]

    if not (fieldOffset == CLICopti.CellParams.off_a  or \
            fieldOffset == CLICopti.CellParams.off_vg or \
            fieldOffset == CLICopti.CellParams.off_Q  or \
            fieldOffset == CLICopti.CellParams.off_rQ or \
            fieldOffset == CLICopti.CellParams.off_Es or \
            fieldOffset == CLICopti.CellParams.off_Hs or \
            fieldOffset == CLICopti.CellParams.off_Sc ):
        raise KeyError("Not all valid fields are accepted for getInterpolated_zidx due to internal limitation")

    if isinstance(zIdx,np.ndarray):
        if zIdx.ndim != 1:
            raise TypeError("Only 1D numpy arrays are supported for z")
        r = np.empty_like(zIdx)
        for i in range(len(r)):
           r[i] = _AccelStructure_getInterpolated_zidx_CPP(self,z[i],fieldOffset)
        return r
    #Else: Scalar
    return _AccelStructure_getInterpolated_zidx_CPP(self,zIdx,fieldOffset)
AccelStructure.getInterpolated_zidx = _AccelStructure_getInterpolated_zidx_PY

_AccelStructure_getInterpolatedZero_CPP = AccelStructure.getInterpolatedZero
def _AccelStructure_getInterpolatedZero_PY(self,fieldName : str):

    import CLICopti.CellParams
    fieldOffset =  CLICopti.CellParams.cellOffsets[fieldName]

    if not (fieldOffset == CLICopti.CellParams.off_vg or \
            fieldOffset == CLICopti.CellParams.off_rQ ):
        raise KeyError("Not all valid fields are accepted for getInterpolatedZero due to internal limitation")

    return _AccelStructure_getInterpolatedZero_CPP(self,fieldOffset)
AccelStructure.getInterpolatedZero = _AccelStructure_getInterpolatedZero_PY
