import string
import os
import glob
import re
import time

##############################################################################
#
# Guinea
#
##############################################################################

class guinea:

    exe = 'guinea'
#    dir = '../../../lc-soft/gp/font/src/'
#    dir = './'
    dir = '/home/fz/gp/runucl/'
    out = 'test'   

    acc = 'tesla'
    par = 'fast'

    acc_dict = {}
    par_dict = {}

    allacc_dict = {}
    allpar_dict = {}

    baseoutdir = '.'
    outext     = ['dat','log','save','out','bin']
    
    dirtoken  = 'nominal'
    
    def __init__(self,acc,par):

        # store local copies of accelerator and parmeters
        self.acc = acc
        self.par = par

        # set accelerator and parameters
        self.setaccdat()

        # write acc.dat
        self.write_accdat()
        
    def setaccdat(self):
        # accelerator definitions
        allacc_dict  = {
            'nlcb':'$ACCELERATOR::NLC-B-500{energy=175.0;particles=0.95;emitt_x=4.5;emitt_y=0.1;beta_x=12.;beta_y=0.12;sigma_z=120.;dist_z=0;espread=0.003;which_espread=0;offset_x=0.0;offset_y=0.0;waist_x=0;waist_y=0;angle_x=0;angle_y=0;angle_phi=0;trav_focus=0;charge_sign=-1;}',
            'tesla':'$ACCELERATOR::tesla{energy=175.0;particles=2.0;beta_x=15.0;beta_y=0.4;emitt_x=10.0;emitt_y=0.03;sigma_z=300.0;espread=0.001;dist_z=0;f_rep=1.0;n_b=1;waist_y=250;offset_x=0.0;offset_y=0.0;}'
        }
    
        # gp options and parameters
        allpar_dict  = {
            'standard':'$PARAMETERS::standard{rndm_save=1;rndm_load=1;n_x=32;n_y=64;n_z=36;n_t=2;cut_x=3.0*sigma_x.1;cut_y=6.0*sigma_y.1;cut_z=3.0*sigma_z.1;n_m=40000;force_symmetric=1;integration_method=2;do_eloss=1;do_espread=1;do_isr=0;store_beam=1;electron_ratio=0.1;do_photons=1;photon_ratio=0.1;store_photons=1;do_pairs=1;track_pairs=1;grids=7;pair_ratio=1.0;pair_ecut=0.005;beam_size=1;do_compt=0;compt_x_min=0.01;compt_emax=800;do_hadrons=1;store_hadrons=1;hadron_ratio=1000.;do_jets=0;store_jets=1;jet_ptmin=3.2;jet_ratio=10000.;jet_log=1;do_lumi=1;num_lumi=10000;lumi_p=1e-30;}',
            'fast':'$PARAMETERS::fast{n_x=32;n_y=64;n_z=24;n_t=3;n_m=50000;cut_x=3.0*sigma_x.1;cut_y=6.0*sigma_y.1;cut_z=3*sigma_z.1;force_symmetric=0;do_photons=1;store_photons=1;do_pairs=1;track_pairs=1;store_pairs=1;do_compt=0;photon_ratio=0;load_beam=0;grids=0;hist_ee_bins=200;hist_ee_max=500.01;charge_sign=-1;ecm_min_gg=0.3*75.0;load_photons=0;offset_y=0.0*0.5*sigma_y.1;do_dump=0;dump_particle=10;dump_step=1;do_lumi=1;lumi_p=1e-30;num_lumi=10000;electron_ratio=0.1;store_beam=1;}'
        }
    
        # parse string
        self.parse_string(allacc_dict[self.acc],allpar_dict[self.par])

    def run(self):
        print "guinea> run"
    
        print "guinea> run> exe :",self.exe
        print "guniea> run> dir :",self.dir
        print "guinea> run> acc :",self.acc
        print "guinea> run> par :",self.par
        self.write_accdat()

        self.log    = string.join([self.out,".log"],'')
        self.out    = string.join([self.out,".out"],'')
        exerun = string.join([self.dir,self.exe," ",self.acc," ",self.par," ",self.out," > ",self.log],'')
        
        print "guinea> run> :",exerun
        
        os.system(exerun)

    def clean(self):
        savelist = ["rndm.save"]
        delextn  = "dat","out","log","out","save"
        delfiles = []
    
        print "guinea> clean> savelist: ",savelist
        print "guinea> clean> delextn : ",delextn
    
        # find all files extensions to be deleted
        for extn in delextn:
            delfiles = delfiles+glob.glob(string.join(['*.',extn],''))
            
        # remove items to be saved
        for save in savelist:
            delfiles.remove(save)
        
        # echo filelist
        print "guinea> clean> deleting: ",delfiles
            
        # delete files
        for file in delfiles:
            os.system(string.join(['rm -rf ',file],''))


    def parse_file(self,filename):
        print  ""
        # read file and assign acc and par strings
        
        # call  parse_string with acc and par string

    def parse_string(self,acc_string,par_string):

        # accelerator parameters
        e = re.compile('[{;]+',re.IGNORECASE)
        acc_vars = e.split(acc_string)
        acc_vars.remove('}')
        del acc_vars[0]

        templist1 = []
        templist2 = []
        
        for var in acc_vars:
            v = string.split(var,"=")
            templist1.append(v[0])
            templist2.append(v[1])

        self.acc_dict = dict(map(None,templist1,templist2))

        # guinea parameters
        par_vars = e.split(par_string)
        par_vars.remove('}')
        del par_vars[0]
       
        templist1 = []
        templist2 = []
        
        for var in par_vars:
            v = string.split(var,"=")
            templist1.append(v[0])
            templist2.append(v[1])

        self.par_dict = dict(map(None,templist1,templist2))

    # get accelerator parameters
    def get_accelerator(self,varname):
        return float(self.acc_dict[varname])

    # modify accelerator parameters
    def modify_accelerator(self,param,value):
        current = float(self.acc_dict[param])
        print 'guinea> modify_accelerator> ,param =',current
        print 'guinea> modify_accelerator>   ,new =',value
        self.acc_dict[param] = str(value)
        self.write_accdat()

    # get guinea pig parameters
    def get_parameter(self,varname):
        return float(self.par_dict[varname])

    # modify guinea parameters
    def modify_parameter(self,param,value):
        current = float(self.par_dict[param])
        print 'guinea> modify_parameter>  ,param, =',current
        print 'guinea> modify_parameter>  ,new  , =',value
        self.par_dict[param] = str(value)
        self.write_accdat()        

    def write_accdat(self):
        # open and write acc.dat file
        accdat = open('acc.dat','w')

        accdat.write('$ACCELERATOR::'+self.acc+'\n{\n')
        # write acc_dict
        for accvar in self.acc_dict:
            accdat.write(accvar+'='+self.acc_dict[accvar]+';\n')
        accdat.write('}\n')

        accdat.write('$PARAMETERS::'+self.par+'\n{\n')
        # write par_dict
        for parvar in self.par_dict:
            accdat.write(parvar+'='+self.par_dict[parvar]+';\n')
        accdat.write('}\n')
      
    def print_config(self):
        print self.acc_dict
        print self.par_dict

    def set_dirtoken(self,dirtoken):
        self.dirtoken = dirtoken

    def move_output(self):
        savefiles = []
    
        # find all files extensions to be deleted
        for ext in self.outext:
            savefiles = savefiles+glob.glob('*.'+ext)        
        
        dirname = self.acc+'_'+self.par+'_'+self.dirtoken+'_'+(time.strftime("%d-%m-%Y_%H-%M-%S").replace(' ','_')).replace('__','_')
        
        # create directory
        dirname = self.baseoutdir+'/'+dirname
        os.mkdir(dirname)

        # move files to directory
        print 'guinea> move_output> moving ',savefiles,' to ', dirname
        for file in savefiles:
            os.system('mv '+file+' '+dirname)

##############################################################################
#
# Guineajitter
#
##############################################################################

class guineajitter:

    acc = "tesla"
    par = "fast"
    param = "energy"
    start = 0
    end   = 0
    iruns = 0
    
    def __init__(self,acc,par,param,iruns,start,end):
        self.acc = acc
        self.par = par
        self.iruns = iruns
        self.mean  = 0.0
        self.sigma = 0.0
        
    def run(self):
        for val in range(0,self.iruns+1):
            xval = self.start+val*(self.end-self.start)/self.iruns
            print 'guineascan> run> iteration',val,self.param,xval

            g = guinea(self.acc,self.par)

            # change parameter
            g.modify_accelerator(self.param,xval)
            
            # run
            g.run()

            # check output
            go = guineaoutput('test.out')
            print go.get_gp('lumi_fine'),go.get_gp('lumi_ee'),go.get_gp('lumi_ee_high')

##############################################################################
#
# Guineaaverage
#
##############################################################################
    
class guineaaverage:

    acc = "tesla"
    par = "fast"
    iruns = 2


    filestocat = 'lumi.ee.out','pairs.dat'   

    dirtoken  = 'nominal'
    baseoutdir = '.'
    outext     = ['cat','bin']
    
    def __init__(self,acc,par,iruns,newparam="none",newvalue=0):    
        self.acc   = acc
        self.par   = par
        self.iruns = iruns
        self.newparam = newparam
        self.newvalue = newvalue

    def clean(self):
        print 'guineaaverage> clean>'

        # loop over cat files
        for file in self.filestocat:
            os.system('rm -rf '+file+'cat')

    def run(self):

        # create guinea object
        g = guinea(self.acc,self.par)

        # clean
        self.clean()        

        # init files
        self.init_files(g)
        
        # tune parameters if newparam argument is given in constructor
        if(self.newparam != "none"):
            g.modify_accelerator(self.newparam,self.newvalue)
            g.set_dirtoken(self.newparam+str(self.newvalue))
            self.dirtoken = self.newparam+str(self.newvalue)
        # run loop
        for iloop in range(1,self.iruns+1):
            print 'guineaaverage> run> iteration : ',iloop

            # run guinea
            g.run()

            # store files
            self.cat_files(iloop)

            # clean up
            g.clean()

        # move files for storage
        self.move_output()
        
    def cat_files(self,iloop):
        print 'guineaaverage> catfiles>'

        # loop over files to cat
        for file in self.filestocat:
            if iloop > -1:
                os.system("awk '{print "+str(iloop)+",$0;}' "+file+" >> "+file+"cat")
            else:
                os.system('cat '+file+' >> '+file+'cat')

    def move_output(self):
        savefiles = []
    
        # find all files extensions to be deleted
        for ext in self.outext:
            savefiles = savefiles+glob.glob('*.*'+ext)        
        
        dirname = self.acc+'_'+self.par+'_'+self.dirtoken+'_iter'+str(self.iruns)+'_'+(time.strftime("%d-%m-%Y_%H-%M-%S").replace(' ','_')).replace('__','_')
        
        # create directory
        dirname = self.baseoutdir+'/'+dirname
        os.mkdir(dirname)

        # move files to directory
        print 'guinea> move_output> moving ',savefiles,' to ', dirname
        for file in savefiles:
            os.system('mv '+file+' '+dirname)

    def init_files(self,g):
        lumi_ee_out = open('lumi.ee.out','w')
        lumi_ee_out.write(str(g.get_accelerator('energy'))+'\n')
        lumi_ee_out.close()
        pairs_dat = open('pairs.dat','w')
        pairs_dat.close()

        self.cat_files(-1)

########################################/######################################
#
# Guineascan
#
##############################################################################

class guineascan:

    acc = "tesla"
    par = "fast"
    param = "energy"
    start = 0
    end   = 0
    iruns = 0

    dirtoken  = ''
    baseoutdir = '.'
    outext     = ['dat','bin']


    
    def __init__(self,acc,par,param,iruns,start,end):
        self.acc   = acc
        self.par   = par
        self.param = param
        self.iruns = iruns
        self.start = start
        self.end   = end

        
    def run(self):


        scandat = open(self.acc+"_"+self.par+"_"+self.param+"_"+str(self.iruns)+"_"+str(self.start)+"_"+str(self.end),'w')
        
        for val in range(0,self.iruns+1):
            xval = self.start+val*(self.end-self.start)/self.iruns
            print 'guineascan> run> iteration',val,self.param,xval

            # create gp object
            g = guinea(self.acc,self.par)

            # change parameter
            g.modify_accelerator(self.param,xval)

            # run
            g.run()

            # check output
            go = guineaoutput('test.out')
            print xval,go.get_gp('lumi_fine'),go.get_gp('lumi_ee'),go.get_gp('lumi_ee_high')

            # write output file
            scandat.write(str(xval)+' '+str(go.get_gp('lumi_fine'))+' '+str(go.get_gp('lumi_ee'))+' '+str(go.get_gp('lumi_ee_high'))+'\n')


            #set dirtoken
            gtoken = self.param+"_"+str(xval)
            g.set_dirtoken(gtoken)
            
            #move output
            g.move_output()

            #delete stuff
            #g.clean()
            
        scandat.close()
        
##############################################################################
#
# Guineaaveragescan
#
##############################################################################

class guineaaveragescan:

    
    def __init__(self,acc,par,iav,param,iruns,start,end):
        self.acc   = acc
        self.par   = par
        self.iav = iav
        self.param = param
        self.iruns = iruns
        self.start = start
        self.end   = end

    def run(self):
        for val in range(0,self.iruns+1):
            xval = self.start+val*(self.end-self.start)/self.iruns
            print 'guineaaveragescan> run> iteration',val,self.param,xval

            g1=guineaaverage(self.acc,self.par,self.iav,self.param,xval)
            g1.run()

    


##############################################################################
#
# Guineaoutput
#
##############################################################################
        
class guineaoutput:

    file_name = ''
    acc_input = ''
    par_input = ''
    gp_output = ''
     
    acc_dict  = {}
    par_dict  = {}
    gp_dict   = {}    
    
    def __init__(self,file_name):   
        self.file_name = file_name
        self.parse_file(file_name)

    def parse_file(self,file_name):
        f = open(file_name,'r')
        fs = f.read()
        fs = fs.replace('{}','temp-token')                
        fs = fs.replace('}{','{')
        fs = fs.replace('}','{')
        fs = fs.replace('\n','')
        e = re.compile('[{]+',re.IGNORECASE)
        output_tokens = e.split(fs)

        del output_tokens[0]

        self.acc_input = output_tokens[0]
        self.par_input = output_tokens[1]
        self.gp_output = output_tokens[2]
        
        e = re.compile(';+',re.IGNORECASE)

        acc_vars = e.split(self.acc_input)
        acc_vars.remove('')
        templist1 = []
        templist2 = []
        for var in acc_vars:
            v = string.split(var,"=")
            templist1.append(v[0])
            templist2.append(v[1])
        self.acc_dict = dict(map(None,templist1,templist2))

        par_vars = e.split(self.par_input)
        par_vars.remove('')
        templist1 = []
        templist2 = []
        for var in par_vars:
            v = string.split(var,"=")
            templist1.append(v[0])
            templist2.append(v[1])
        self.par_dict = dict(map(None,templist1,templist2))

        gp_vars = e.split(self.gp_output)
        gp_vars.remove('')
        templist1 = []
        templist2 = []
        for var in gp_vars:
            v = string.split(var,"=")
            templist1.append(v[0])
            templist2.append(v[1])
        self.gp_dict = dict(map(None,templist1,templist2))
        
    def get_acc(self,varname):
        return float(self.acc_dict[varname])

    def get_par(self,varname):
        return float(self.par_dict[varname])

    def get_gp(self,varname):
        return float(self.gp_dict[varname])