import math
from onlinemodel.core.apptemplate import ApplicationTemplate


# Elegant Interface

class Elegant(ApplicationTemplate):

    def __init__(self, onlinemodel =None,parent=None):
        ApplicationTemplate.__init__(self)

        self.model=onlinemodel
        self.parent=parent

        # see what is needed

        self.wpath = 'ReferenceFiles'  # path for reference files
        self.fid = 0
        self.lsc = 0  # flag whether lsc is enabled
        self.csr = 0  # flag whether csr is enabled or not
        self.wakebins = 0  # automatic calculation of
        self.charge = 200e-12  # default charge in the lattice
        self.dump = []  # diagnostics location where to dump particle distribution

        self.seq = []
        self.driftcount = 1  # drift counter

        self.lscbins = 100  # number of bins for lsc calculation
        self.csrbins = 100  # number of bins for CSR calculation
        self.csrmaxlen = 10.  # distance over which the CSR wake is considered after a dipole
        self.csrlen = 0  # running variable to count distance after last bend

        self.center = []  # location where the beam is centered due to CSR kicks
        self.Distribution = 1

    def export(self):
        if self.model is None:
            return
        start = str(self.parent.UITrackStart.text()).upper()
        end = str(self.parent.UITrackEnd.text()).upper()
        self.charge=float(str(self.parent.UIElegantCharge.text()))
        if self.parent.UIElegantCSR.isChecked():
            self.csr=1
        else:
            self.csr=0
        if self.parent.UIElegantLSC.isChecked():
            self.lsc=1
        else:
            self.lsc=0
        if len(start) > 7:
            start = start[0:7]
        if len(end) > 7:
            end = end[0:7]
        refloc, twiss0 = self.parent.reference.getReference()

        if 'SWISSFEL' in refloc.upper():
            refloc = 'SINLH01$START'
        if refloc.upper() == 'START':
            refloc = start.upper()
        startEnergy = self.model.om.EnergyAt("SINLH01.DBAM010")[0] * 1e-6
        print('Start Energy is (MeV):', startEnergy)
        twiss0['energy'] = startEnergy * 1e-3  # convert to GeV for madx
        start, end = self.model.checkRange(start, end, refloc[0:7])
        if start is None:
            return
        self.model.om.setBranch(end.upper())
        print('Exporting from', start, 'to', end, '(Reference Location:', refloc, ')')
        filename = self.openLatticeStream('Elegant', 'SwissFEL')
        self.model.om.exportBranch(self)
        self.closeLatticeStream()
        if self.parent.UIElegantMainInput.isChecked():
            print('Writing Elegant input file')
            self.writeInputFile('Elegant','SwissFEL',twiss0)
        return filename


    def writeInputFile(self,path : str, root: str,twiss0:dict):
        #path,name,latfile,distfile,distinfo,ds):
        filename = path+'/'+root+'.ele'
        latfile=filename.replace('ele','lat')
        with open(filename, 'w') as fid:
            fid.write('&run_setup\n')
            fid.write('\tlattice\t\t= %s,\n' % (latfile))
            fid.write('\tuse_beamline\t= SwissFEL,\n')
            fid.write('\trootname\t= %s,\n' % root)
            fid.write('\toutput\t\t= %s/%%s.out,\n' % path)
            fid.write('\tcentroid\t\t= %s/%%s.cen,\n' % path)
            fid.write('\tsigma\t\t= %s/%%s.sig,\n' % path)
            fid.write('\tfinal\t\t= %s/%%s.fin,\n' % path)
            fid.write('\tparameters\t= %s/%%s.par,\n' % path)
            fid.write('\tmagnets\t\t= %s/%%s.mag,\n' % path)
            fid.write('\tcombine_bunch_statistics = 0,\n')
            fid.write('\tdefault_order\t= 2,\n')
            fid.write('\tconcat_order\t= 0,\n')
            fid.write('\tprint_statistics\t= 0,\n')
            fid.write('\trandom_number_seed\t= 9876543210,\n')
            fid.write('\tp_central\t= %f,\n' % (twiss0['energy']/0.511e-3))
            fid.write('\ttracking_updates\t= 1\n')
            fid.write('\talways_change_p0\t= 1\n')
            fid.write('&end\n\n')
            fid.write('&run_control\n')
            fid.write('\tn_steps\t= 1,\n')
            fid.write('\treset_rf_for_each_step = 1\n')
            fid.write('&end\n\n')
            fid.write('&twiss_output\n')
            fid.write('\tfilename\t= %s/%%s.twi,\n' % path)
            fid.write('\tmatched\t\t= 0,\n')
            fid.write('\tbeta_x\t\t= %f,\n' % twiss0['betx'])
            fid.write('\tbeta_y\t\t= %f,\n' % twiss0['bety'])
            fid.write('\talpha_x\t\t= %f,\n' % twiss0['alfx'])
            fid.write('\talpha_y\t\t= %f,\n' % twiss0['alfy'])
            fid.write('&end\n\n')

            #fid.write('&sdds_beam\n')
            #fid.write('\tinput_type= "elegant",\n')
            #fid.write('\tsample_interval\t= %d,\n' % ds)
            #fid.write('\tinput = %s,\n' % distfile)
            #fid.write('\treuse_bunch\t= 0 \n')
            #fid.write('&end\n\n')

            fid.write('&track\n')
            fid.write('&end\n\n')
            fid.close()

        return filename


    def openLatticeStream(self, path, name):
        filename = '%s/%s.lat' % (path, name)
        #        self.path=path
        self.path = '.'
        self.savepath = path
        self.seq = []
        self.fid = open(filename, 'w')
        self.driftcount = 1
        self.write('Q: CHARGE, TOTAL=%e\n' % self.charge)
        self.write('ALTER: CENTER, DELTA=1,T=1\n')  # put the beam on the right energy
        return filename

    def closeLatticeStream(self):
        self.fid.close()

    # routines to write lattice

    def setCSR(self):
        self.csrlen = self.csrmaxlen

    def isType(self, name):
        if (name.find('elegant') > -1):
            return 1
        else:
            return 0

    def write(self, line):
        self.fid.write(line)

    def writeLine(self, line, seq):

        lc = 1
        name = line.Name

        seqloop = seq
        if len(name) == 7:
            seqloop = self.seq

        if len(name) < 1:
            name = 'SwissFEL'

        self.write("\n%s.Start: MARKER\n" % (name))
        secend = name + '.End'
        if secend in self.dump:
            print('Adding dump in Elegant output for', name)
            self.write('%s.End: WATCH, FILENAME="Runs/%%s.%s.End.out", MODE=COORD\n' % (name, name))
        else:
            self.write("%s.End: MARKER\n" % (name))

        if (name == "SwissFEL"):
            self.write("\n%s: Line=(Q, ALTER, %s.Start, " % (name, name))
        else:
            self.write("\n%s: Line=(%s.Start, " % (name, name))
        for ele in seqloop:
            #            print ele
            if isinstance(ele, str):
                self.write("%s, " % (ele))
            else:
                self.write("%s, " % (ele['Name']))
            lc = lc + 1
            if (lc > 5):
                self.write("&\n\t")
                lc = 0
        self.write("%s.End)\n\n" % (name))

        if len(name) == 7:
            del self.seq[:]

    def writeDrift(self, l):
        if l <= 0:
            return 0
        name = "Drift%4.4d" % (self.driftcount)
        if self.csrlen > 0:
            self.write("%s.csr: CSRDRIFT, L=%f, DZ=0.1, USE_STUPAKOV=1, CSR=%d\n" % (name, l, self.csr))
            self.write("%s.lsc: LSCDRIFT, L=0, LEFFECTIVE=%f, BINS=%d, SMOOTHING=1, LSC=%d\n" % (name, l, self.lscbins,
                                                                                                 self.lsc))
            self.write("%s: LINE=(%s.csr,%s.lsc)\n" % (name, name, name))
            self.csrlen = self.csrlen - l
        else:
            self.write("%s: LSCDRIFT, L=%f, LEFFECTIVE=%f, BINS=%d, SMOOTHING=1, LSC=%d\n" % (name, l, l, self.lscbins,
                                                                                              self.lsc))
        self.seq.append(name)
        self.driftcount = self.driftcount + 1
        return name

    def writeMarker(self, ele):
        if 'ALIG' in ele.Name and 'dy' in ele.__dict__.keys():
            self.write('%s: MALIGN, DY=-%f \n' % (ele.Name, ele.dy))
        else:
            self.write('%s: MARK \n' % (ele.Name))
        self.seq.append(ele.Name)

    def writeVacuum(self, ele):
        ds = ele.sRef + 0.5 * ele.getLength();
        if ds > 0:
            self.writeDrift(ds)
        self.writeMarker(ele)
        ds = ele.getResLength() - ele.sRef - 0.5 * ele.getLength();
        if ds > 0:
            self.writeDrift(ds)

    def writeDiagnostic(self, ele):
        Seval = ele.Seval
        if (Seval < 0):
            Seval = 0.5 * ele.Length
        ds = ele.sRef + Seval;
        if ds > 0:
            self.writeDrift(ds)
        if ele.Name in self.center:
            print('Adding center in Elegant output for', ele.Name)
            self.write('%s_CENTER: CENTER\n' % ele.Name)
            self.seq.append('%s_CENTER' % ele.Name)

        if ele.Name in self.dump:
            print('Adding dump in Elegant output for', ele.Name)
            self.write('%s: WATCH, FILENAME="Runs/%%s.%s.out", MODE=COORD\n' % (ele.Name, ele.Name))
            self.seq.append(ele.Name)
        else:
            self.writeMarker(ele)

        ds = ele.getResLength() - ds
        if ds > 0:
            self.writeDrift(ds)

    def writeCorrector(self, ele):

        corx = 0
        cory = 0
        if 'corx' in ele.__dict__:
            corx = ele.corx
        if 'cory' in ele.__dict__:
            cory = ele.cory

        if 'MADthin' in ele.__dict__:
            if 'MADshift' in ele.__dict__:
                ds = ele.Length * 0.5 + ele.MADshift
                self.writeDrift(ds)
            self.write("%s: KICKER, L=%f, HKICK=%e,VKICK=%e\n" % (ele.Name, 0, corx, cory))
            self.seq.append(ele.Name)
            return

        if (ele.sRef > 0):
            self.writeDrift(ele.sRef)
        if self.lsc == 0 or ele.getLength() == 0:
            self.write("%s: KICKER, L=%f, HKICK=%e,VKICK=%e\n" % (ele.Name, ele.getLength(), corx, cory))
            self.seq.append(ele.Name)
        else:
            self.write("%s.cor: KICKER, L=%f, HKICK=%e,VKICK=%e\n" % (ele.Name, ele.getLength(), corx, cory))
            self.write(
                "%s.lsc: LSCDRIFT, L=0, LEFFECTIVE=%f, BINS=%d, SMOOTHING=1, LSC=%d\n" % (ele.Name, ele.getLength(),
                                                                                          self.lscbins, self.lsc))
            self.write("%s: LINE=(%s.cor,%s.lsc)\n" % (ele.Name, ele.Name, ele.Name))
            self.seq.append(ele.Name)
        ds = ele.getResLength() - ele.getLength() - ele.sRef
        if ds > 0:
            self.writeDrift(ds)

    def writeSextupole(self, ele):
        if (ele.sRef > 0):
            self.writeDrift(ele.sRef)
        if self.lsc == 0:
            self.write("%s: SEXTUPOLE, L=%f, K2=%e, Tilt=%f\n" % (ele.Name, ele.getLength(), ele.k2, ele.Tilt))
            self.seq.append(ele.Name)
        else:
            self.write("%s.sex: SEXTUPOLE, L=%f, K2=%e, Tilt=%f\n" % (ele.Name, ele.getLength(), ele.k2, ele.Tilt))
            self.write(
                "%s.lsc: LSCDRIFT, L=0, LEFFECTIVE=%f, BINS=%d, SMOOTHING=1, LSC=%d\n" % (ele.Name, ele.getLength(),
                                                                                          self.lscbins, self.lsc))
            self.fid.write("%s: LINE=(%s.sex,%s.lsc)\n" % (ele.Name, ele.Name, ele.Name))
            self.seq.append(ele.Name)
        ds = ele.getResLength() - ele.getLength() - ele.sRef
        if ds > 0:
            self.writeDrift(ds)

    def writeQuadrupole(self, ele):
        if (ele.sRef > 0):
            self.writeDrift(ele.sRef)
        if self.lsc != 0 or ('corx' in ele.__dict__) or 'cory' in ele.__dict__:
            locseq = []
            if self.lsc > 0:
                self.write(
                    "%s.lsc: LSCDRIFT, L=0, LEFFECTIVE=%f, BINS=%d, SMOOTHING=1, LSC=%d\n" % (ele.Name, ele.getLength(),
                                                                                              self.lscbins, self.lsc))
                locseq.append("%s.lsc" % (ele.Name))
            if ('corx' in ele.__dict__) or 'cory' in ele.__dict__:
                corx = 0
                cory = 0
                if 'corx' in ele.__dict__:
                    corx = ele.corx
                if 'cory' in ele.__dict__:
                    corx = ele.cory
                self.write(
                    "%s.q1: QUADRUPOLE, L=%f, K1=%e, Tilt=%f\n" % (ele.Name, ele.getLength() * 0.5, ele.k1, ele.Tilt))
                self.write("%s.cor: KICKER, L=0, HKICK=%e,VKICK=%e\n" % (ele.Name, corx, cory))
                self.write(
                    "%s.q2: QUADRUPOLE, L=%f, K1=%e, Tilt=%f\n" % (ele.Name, ele.getLength() * 0.5, ele.k1, ele.Tilt))
                locseq.append("%s.q1" % (ele.Name))
                locseq.append("%s.cor" % (ele.Name))
                locseq.append("%s.q2" % (ele.Name))
            else:
                self.write("%s.q: QUADRUPOLE, L=%f, K1=%e, Tilt=%f\n" % (ele.Name, ele.getLength(), ele.k1, ele.Tilt))
                locseq.append("%s.q" % (ele.Name))
            self.fid.write("%s: LINE=(%s" % (ele.Name, locseq[0]))
            for p in locseq[1:len(locseq)]:
                self.fid.write(",%s" % p)
            self.fid.write(")\n")
            self.seq.append(ele.Name)
        else:
            self.write("%s: QUADRUPOLE, L=%f, K1=%e, Tilt=%f\n" % (ele.Name, ele.getLength(), ele.k1, ele.Tilt))
            self.seq.append(ele.Name)
        ds = ele.getResLength() - ele.getLength() - ele.sRef
        if ds > 0:
            self.writeDrift(ds)

    def writeBend(self, ele):
        if (ele.sRef > 0):
            self.writeDrift(ele.sRef)
        Lpath = ele.getLength()
        angrad = ele.angle * math.asin(1) / 90
        angradcor = angrad
        if 'cor' in ele.__dict__:
            angradcor = angrad + ele.cor
        if (self.lsc != 0):
            self.write(
                "%s.lsc: LSCDRIFT, L=0, LEFFECTIVE=%f, BINS=%d, SMOOTHING=1, LSC=%d\n" % (ele.Name, Lpath, self.lscbins,
                                                                                          self.lsc))
            self.seq.append("%s.lsc" % ele.Name)
            self.write("%s.b: CSRCSBEND, L=%f, ANGLE=%f, &\n" % (ele.Name, Lpath, angradcor))
            self.write("\tE1=%f, E2=%f, TILT=%f, &\n" % (ele.e1 * angrad, ele.e2 * angrad, ele.Tilt))
            self.write("\tHGAP=0.015, FINT=0.5, NONLINEAR=1, N_KICKS=10,INTEGRATION_ORDER=4, &\n")
            self.write("\tBINS=%d,SG_HALFWIDTH=2, ISR=%d, SYNCH_RAD=%d, CSR=%d\n" % (self.csrbins, self.csr, self.csr,
                                                                                     self.csr))
            self.write("%s: LINE=(%s.lsc,%s.b)\n" % (ele.Name, ele.Name, ele.Name))
            self.seq.append(ele.Name)
        else:
            self.write("%s: CSRCSBEND, L=%f, ANGLE=%f, &\n" % (ele.Name, Lpath, angradcor))
            self.write("\tE1=%f, E2=%f, TILT=%f, &\n" % (ele.e1 * angrad, ele.e2 * angrad, ele.Tilt))
            self.write("\tHGAP=0.015, FINT=0.5, NONLINEAR=1, N_KICKS=10,INTEGRATION_ORDER=4, &\n")
            self.write("\tBINS=%d,SG_HALFWIDTH=2, ISR=%d, SYNCH_RAD=%d, CSR=%d\n" % (self.csrbins, self.csr, self.csr,
                                                                                     self.csr))
            self.seq.append(ele.Name)
        if abs(angrad) > 0:
            self.setCSR()
        ds = ele.getResLength() - ele.getLength() - ele.sRef
        if ds > 0:
            self.writeDrift(ds)

    def writeUndulator(self, ele):
        pi = 2 * math.asin(1)
        Nu = round(ele.getLength() * ele.ku / pi)
        r = 3e8 / 511000 / 2 / pi / 100;
        Bu = ele.K * ele.ku * 0.01 / 2 / pi / r;
        LaserHeater = 0
        if 'UPHS' in ele.Name:
            ele.K = 0
        if 'UDLY' in ele.Name:
            ele.K = 0
        if 'Power' in ele.__dict__:
            if ele.Power > 0:
                LaserHeater = 0  ######## disable laser modulator, since energy spread is not correct anyhow
        if ele.sRef > 0:
            self.writeDrift(ele.sRef)
        if ele.K == 0:
            self.write("%s: LSCDRIFT, L=%f, LEFFECTIVE=%f, BINS=%d, SMOOTHING=1, LSC=%d\n" % (ele.Name, ele.getLength(),
                                                                                              ele.getLength(),
                                                                                              self.lscbins, self.lsc))
            self.seq.append(ele.Name)
            # add dechirper here
            # if 'UDCP' in ele.Name:
            #    wake=Dechirper()
            #    wake.getWake(ele.gap, ele.Name,self.savepath)
            #    print('Write wakefile for %s with gap: %f mm' % (ele.Name,ele.gap*1e3))
            #    self.write('%s_LW   : WAKE, INPUTFILE = "%s/wake_%s.sdds", CHANGE_P0 = 1, &\n' % (ele.Name,self.path,ele.Name))
            #    self.write('\t TCOLUMN="t",WCOLUMN="W", INTERPOLATE=1, FACTOR = %f, N_BINS=%d, SMOOTHING=1 \n'% (ele.getLength(),self.wakebins))
            #    self.seq.append('%s_LW' % ele.Name)

            #    if 'DECHIRPER-V' in ele.Baugruppe:
            #        self.write('%s_DW   : TRWAKE, INPUTFILE = "%s/wake_%s.sdds", &\n' % (ele.Name,self.path,ele.Name))
            #        self.write('\t TCOLUMN="t",WXCOLUMN="WD", WYCOLUMN="WD", XFACTOR = 0, YFACTOR = %f, &\n' % (ele.getLength()))
            #        self.write('\t INTERPOLATE=1, N_BINS=%d, SMOOTHING=1 \n'% self.wakebins)
            #        self.write('%s_QW   : TRWAKE, INPUTFILE = "%s/wake_%s.sdds", &\n' % (ele.Name,self.path,ele.Name))
            #        self.write('\t TCOLUMN="t",WXCOLUMN="WQP", WYCOLUMN="WQ", XFACTOR = %f, YFACTOR = %f, &\n' % (ele.getLength(),ele.getLength()))
            #        self.write('\t INTERPOLATE=1, N_BINS=%d, SMOOTHING=1, & \n'% self.wakebins)
            #        self.write('\t X_DRIVE_EXPONENT=0, Y_DRIVE_EXPONENT = 0, X_PROBE_EXPONENT = 1, Y_PROBE_EXPONENT = 1\n')
            #        self.seq.append('%s_DW' % ele.Name)
            #        self.seq.append('%s_QW' % ele.Name)
            #    elif 'DECHIRPER-H' in ele.Baugruppe:
            #        self.write('%s_DW   : TRWAKE, INPUTFILE = "%s/wake_%s.sdds", &\n' % (ele.Name,self.path,ele.Name))
            #        self.write('\t TCOLUMN="t",WXCOLUMN="WD", WYCOLUMN="WD", XFACTOR = %f, YFACTOR = 0, &\n' % (ele.getLength()))
            #        self.write('\t INTERPOLATE=1, N_BINS=%d, SMOOTHING=1 \n'% self.wakebins)
            #        self.write('%s_QW   : TRWAKE, INPUTFILE = "%s/wake_%s.sdds", &\n' % (ele.Name,self.path,ele.Name))
            #        self.write('\t TCOLUMN="t",WXCOLUMN="WQ", WYCOLUMN="WQP", XFACTOR = %f, YFACTOR = %f, &\n' % (ele.getLength(),ele.getLength()))
            #        self.write('\t INTERPOLATE=1, N_BINS=%d, SMOOTHING=1, & \n'% self.wakebins)
            #        self.write('\t X_DRIVE_EXPONENT=0, Y_DRIVE_EXPONENT = 0, X_PROBE_EXPONENT = 1, Y_PROBE_EXPONENT = 1\n')
            #        self.seq.append('%s_DW' % ele.Name)
            #        self.seq.append('%s_QW' % ele.Name)

            ds = ele.getResLength() - ele.getLength() - ele.sRef
            if ds > 0:
                self.writeDrift(ds)
            return
        name = ele.Name
        if self.lsc != 0:
            self.write(
                "%s.lsc: LSCDRIFT, L=0, LEFFECTIVE=%f, BINS=%d, SMOOTHING=1, LSC=%d\n" % (ele.Name, ele.getLength(),
                                                                                          self.lscbins, self.lsc))
            self.seq.append("%s.lsc" % ele.Name)
            name = "%s.wig" % ele.Name
        if LaserHeater == 0:
            self.write("%s  : WIGGLER, L=%f, K=%f, POLES=%d, TILT=%f\n" % (name, ele.getLength(), ele.K, Nu, ele.Tilt))
        else:
            self.write("%s  : LSRMDLTR, L=%f, PERIODS=%d, BU=%f, &\n" % (name, ele.getLength(), Nu / 2, Bu));
            self.write("\t LASER_W0=%e, LASER_PEAK_POWER=%e, &\n" % (ele.Waist, ele.Power));
            self.write("\t N_STEPS=%d, ACCURACY=0.001, POLE_FACTOR1=1, POLE_FACTOR2=1,POLE_FACTOR3=1\n" % (Nu * 60));
        self.seq.append(name)
        ds = ele.getResLength() - ele.getLength() - ele.sRef
        if ds > 0:
            self.writeDrift(ds)

    def writeRF(self, ele):

        if ele.sRef > 0:
            self.writeDrift(ele.sRef)
        if ele.Tag == "RACC":
            self.write('%s   : RFCW, FREQ=%f, CELL_LENGTH=1.0, &\n' % (ele.Name, ele.Frequency))
            self.write(
                '\t L=%f, VOLT=%e, PHASE= %f, &\n' % (ele.getLength(), ele.Gradient * ele.getLength(), ele.Phase))
            self.write('\t CHANGE_P0=1, END1_FOCUS=1, END2_FOCUS=1, &\n');
            self.write('\t BODY_FOCUS_MODEL="NONE",INTERPOLATE=1, N_BINS=%d, SMOOTHING=1, &\n' % (self.wakebins));
            self.write('\t ZWAKEFILE="%s/wake_l_%s-band.sdds", &\n' % (self.wpath, ele.Band));
            self.write('\t TRWAKEFILE="%s/wake_t_%s-band.sdds", &\n' % (self.wpath, ele.Band));
            self.write('\t LSC=%d, LSC_BINS=%d, N_KICKS=10, &\n' % (self.lsc, self.lscbins));
            self.write('\t TCOLUMN="t",WXCOLUMN="W", WYCOLUMN="W", WZCOLUMN="W"\n');
        else:
            if ele.Band == 'S':
                SWave = 1
            else:
                SWave = 0
            self.write('%s   : RFDF, FREQUENCY=%f, STANDING_WAVE=%d, TILT=%f, &\n' % (ele.Name, ele.Frequency, SWave,
                                                                                      ele.Tilt))
            grad = ele.Gradient * ele.getLength()
            self.write('\t L=%f, VOLTAGE=%e, PHASE= %f\n' % (ele.getLength(), grad, ele.Phase))
        self.seq.append(ele.Name)
        ds = ele.getResLength() - ele.getLength() - ele.sRef
        if ds > 0:
            self.writeDrift(ds)