from onlinemodel.code import OMElegant
#test
class Elegant:
    def __init__(self,parent=None,model=None):
        self.parent=parent
        self.model=model
        self.elegant=OMElegant.ElegantLattice()

    def writeElegantFiles(self,path ='Elegant', filename='dummy',end='SARBD02'):
        self.writeElegantInput(path,filename)
        self.writeElegantLattice(path, filename,end)

    def writeElegantInput(self,path,name):

        filename = '%s/%s.ele' % (path, name)
        with open(filename, 'w') as fid:
            fid.write('&run_setup\n')
            fid.write('\tlattice\t\t= %s.lat,\n' % (name))
            fid.write('\tuse_beamline\t= SwissFEL,\n')
            fid.write('\trootname\t= %s,\n' % name)
            fid.write('\toutput\t\t= %s.out,\n')
            fid.write('\tcentroid\t\t= %s.cen,\n')
            fid.write('\tsigma\t\t= %s.sig,\n')
            fid.write('\tfinal\t\t= %s.fin,\n')
            fid.write('\tparameters\t= %s.par,\n')
            fid.write('\tmagnets\t\t= %s.mag,\n')
            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')
            p0 = self.model.getInitialEnergy()
            fid.write('\tp_central\t= %f,\n' % p0)
            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')

            twiss=self.parent.reference.reference['Start']['Twiss']
            print(twiss)
            fid.write('&twiss_output\n')
            fid.write('\tfilename\t= %s.twi,\n')
            fid.write('\tmatched\t\t= 0,\n')
            fid.write('\tbeta_x\t\t= %f,\n' % twiss['betax'])
            fid.write('\tbeta_y\t\t= %f,\n' % twiss['betay'])
            fid.write('\talpha_x\t\t= %f,\n' % twiss['alphax'])
            fid.write('\talpha_y\t\t= %f,\n' % twiss['alphay'])
            fid.write('&end\n\n')

            fid.write('&sdds_beam\n')
            fid.write('\tinput_type= "elegant",\n')
            fid.write('\tsample_interval\t= 1,\n')
            fid.write('\tinput = MyInputDist.sdds,\n')
            fid.write('\treuse_bunch\t= 0 \n')
            fid.write('&end\n\n')

            fid.write('&alter_elements\n')
            fid.write('\tname = *,\n')
            fid.write('\ttype = LSCDRIFT,\n')
            fid.write('\titem = LSC,\n')
            fid.write('\tvalue = 1\n')
            fid.write('&end\n\n')

            fid.write('&alter_elements\n')
            fid.write('\tname = *,\n')
            fid.write('\ttype = RFCW,\n')
            fid.write('\titem = LSC,\n')
            fid.write('\tvalue = 1\n')
            fid.write('&end\n\n')

            fid.write('&alter_elements\n')
            fid.write('\tname = *,\n')
            fid.write('\ttype = CSRDRIFT,\n')
            fid.write('\titem = CSR,\n')
            fid.write('\tvalue = 1\n')
            fid.write('&end\n\n')

            fid.write('&alter_elements\n')
            fid.write('\tname = *,\n')
            fid.write('\ttype = CSRCSBEND,\n')
            fid.write('\titem = CSR,\n')
            fid.write('\tvalue = 1\n')
            fid.write('&end\n\n')

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

    def writeElegantLattice(self,path, filename,end):
        destination = 'ARAMIS'
        if 'SPO' in end:
            destination = 'PORTHOS'
        elif 'SAT' in end:
            destination = 'ATHOS'
        elif 'S10BD' in end or 'SIN' in end:
            destination = 'INJECTOR'
        self.model.om.setBranch(destination, 'SINLH01')
        self.elegant.openLatticeStream(path, filename)
        self.model.om.exportBranch(self.elegant)
        self.elegant.closeLatticeStream()