class dst:
    '''
    Simple class to read in a 
    TraceWin distribution file

    Class afterwards hold the following
    dictionary items:
      - x [cm]
      - xp [rad]
      - y [cm]
      - yp [rad]
      - phi [rad]
      - E [MeV]
    '''
    def __init__(self, filename):
        # easy storage..
        self.filename=filename
        # used to create dict behaviour..
        self._columns=['x','xp','y','yp','phi','E']
        # read in the file..
        self._readBinaryFile()

    def _readBinaryFile(self):
        # Thanks Emma!

        import numpy

        fin=file(self.filename,'r')

        # dummy, Np, Ib, freq, dummy
        Header_type = numpy.dtype([
            ('dummy12', numpy.int16),
            ('Np',      numpy.int32),
            ('Ib',      numpy.float64),
            ('freq',    numpy.float64),
            ('dummy3',  numpy.int8)
            ])
        Header=numpy.fromfile(fin, dtype=Header_type, count=1)
        self.Np=Header['Np'][0]
        self.Ib=Header['Ib'][0]
        self.freq=Header['freq'][0]

        Table=numpy.fromfile(fin, dtype=numpy.float64, count=self.Np*6)
        self._data=Table.reshape(self.Np,6)

        Footer=numpy.fromfile(fin, dtype=numpy.float64, count=1)
        self.mass=Footer[0]

    def __getitem__(self, key):
        # makes the class function as a dictionary
        # e.g. dst['x'] returns the x array..
        if key in self._columns:
            i=self._columns.index(key)
            return self._data[:,i]

class plt:
    '''
    Simple class to read in a 
    TraceWin plot file

    Class afterwards hold the following
    dictionary items:
      - Ne (number of locations)
      - Np (number of particles)
      - Ib [A] (beam current)
      - freq [MHz]
      - mc2  [MeV] 
      - Nelp [m] (locations)

    each plt[i], where i is element number, holds:
      - Zgen [cm] (location)
      - phase0 [deg] (ref phase)
      - wgen [MeV] (ref energy)
      - x [array, cm]
      - xp [array, rad]
      - y [array, cm]
      - yp [array, rad]
      - phi [array, rad]
      - E [array, MeV]
      - l [array] (is lost)

      Example::
        plt=ess.TraceWin.plt('calc/dtl1.plt')
        for i in [97,98]:
          data=plt[i]$
          if data:
            print data['x']
    '''
    def __init__(self, filename):
        # easy storage..
        self.filename=filename
        # used to create dict behaviour..
        self._columns=['x','xp','y','yp','phi','E', 'l']
        # read in the file..
        self._readBinaryFile()

    def _readBinaryFile(self):
        # Thanks Emma!

        import numpy

        fin=file(self.filename,'r')

        # dummy, Np, Ib, freq, dummy
        Header_type = numpy.dtype([
            ('dummy12', numpy.int16),
            ('Ne',      numpy.int32),
            ('Np',      numpy.int32),
            ('Ib',      numpy.float64),
            ('freq',    numpy.float64),
            ('mc2',     numpy.float64),
            ])
        SubHeader_type = numpy.dtype([
             ('dummy12', numpy.int8),
             ('Nelp',      numpy.int32),
             ('Zgen',     numpy.float64),
             ('phase0',     numpy.float64),
             ('wgen',     numpy.float64),
             ])


        Header=numpy.fromfile(fin, dtype=Header_type, count=1)
        self.Np=Header['Np'][0]
        self.Ne=Header['Ne'][0]
        self.Ib=Header['Ib'][0]
        self.freq=Header['freq'][0]
        self.mc2=Header['mc2'][0]

        self._data=[]
        self.Nelp=[]

        i=0
        while i<self.Ne:
            SubHeader=numpy.fromfile(fin, dtype=SubHeader_type, count=1)
            i=SubHeader['Nelp'][0]

            self.Nelp.append(i)

            Table=numpy.fromfile(fin, dtype=numpy.float32, count=self.Np*7)
            Table=Table.reshape(self.Np,7)
            data={}
            for key in ['Zgen','phase0','wgen']:
                data[key]=SubHeader[key][0]
            for j in xrange(7):
                    c=self._columns[j]
                    data[c]=Table[:,j]
            self._data.append(data)

    def __getitem__(self, key):
        if key in self.Nelp:
            i=self.Nelp.index(key)
            return self._data[i]
        else:
            print "No data to plot at element",key

class density_file:
    '''
    Simple class to read a TraceWin density file
    into a pythonized object
    '''

    def __init__(self, filename):
        import numpy, sys


        self.filename=filename
        self.fin=file(self.filename, 'r')

        # currently unknown:
        self.version=0

        # first we simply count how many elements we have:
        counter=0
        while True:
            try:
                self._skipAndCount()
                counter+=1
            except IndexError: # EOF reached..
                break
        if sys.flags.debug:
            print "Number of steps found:", counter
        self.fin.seek(0)

        # set up the arrays..
        self.i=0
        # z position [m] :
        self.z=numpy.zeros(counter)
        # current [mA] :
        self.ib=numpy.zeros(counter)
        # number of lost particles:
        self.Np=numpy.zeros(counter)

        self.Xouv=numpy.zeros(counter)
        self.Youv=numpy.zeros(counter)


        self.moy=numpy.zeros((counter,7))
        self.moy2=numpy.zeros((counter,7))

        self._max=numpy.zeros((counter,7))
        self._min=numpy.zeros((counter,7))

        if self.version>=7:
            self.rms_emit=numpy.zeros((counter,3))
            self.rms_emit2=numpy.zeros((counter,3))

        self.lost=numpy.zeros((counter,self.Nrun))
        self.powlost=numpy.zeros((counter,self.Nrun))

        self.lost2=numpy.zeros(counter)
        self.Minlost=numpy.zeros(counter)
        self.Maxlost=numpy.zeros(counter)

        self.powlost2=numpy.zeros(counter)
        self.Minpowlost=numpy.zeros(counter)
        self.Maxpowlost=numpy.zeros(counter)


        while self.i<counter:
            self._getFullContent()
            self.i+=1


    def _getHeader(self):
        import numpy


        # header..
        version=numpy.fromfile(self.fin, dtype=numpy.int16, count=1)[0]
        year=numpy.fromfile(self.fin, dtype=numpy.int16, count=1)[0]

        # there is much more data written, but it is undocumented. Our trick to get back "on line":
        shift=0
        while year!=2011 or version!=8:
            shift+=1
            version=year
            year=numpy.fromfile(self.fin, dtype=numpy.int16, count=1)[0]

        self.vlong=numpy.fromfile(self.fin, dtype=numpy.int16, count=1)[0]
        self.Nrun=numpy.fromfile(self.fin, dtype=numpy.int32, count=1)[0]

        if shift:
            raise ValueError("ERROR, shifted "+str(shift*2)+" bytes")

        self.version=version
        self.year=year

    def _skipAndCount(self):

        import numpy

        self._getHeader()
        
        if self.version==8:
            numpy.fromfile(self.fin, dtype=numpy.int16, count=8344/2)
        if self.Nrun>1:
            #WARN not 100% sure if this is correct..
            numpy.fromfile(self.fin, dtype=numpy.int16, count=((5588+self.Nrun*12)/2))

    def _get_7dim_array(array):
        return dict(x=array[0],
                    y=array[1],
                    phase=array[2],
                    energy=array[3],
                    r=array[4],
                    z=array[5],
                    dpp=array[6],
                    )

    def _getFullContent(self):

        import numpy

        #self._getHeader()
        # no need to read the header again:
        # (though only if we are SURE about content!)
        numpy.fromfile(self.fin, dtype=numpy.int16, count=5)

        nelp=numpy.fromfile(self.fin, dtype=numpy.int32, count=1)[0]
        self.ib[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=1)[0]
        self.z[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=1)[0]
        # Aperture
        self.Xouv[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=1)[0]
        self.Youv[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=1)[0]
        step=numpy.fromfile(self.fin, dtype=numpy.int32, count=1)[0]

        n=7 # x [m], y[m], Phase [deg], Energy [MeV], R[m], Z[m], dp/p

        self.moy[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=n)[:]
        self.moy2[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=n)[:]
        
        self._max[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=n)[:]
        self._min[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=n)[:]

        if self.version>=5:
            rms=numpy.fromfile(self.fin, dtype=numpy.float32, count=n)
            rms_size2=numpy.fromfile(self.fin, dtype=numpy.float32, count=n)
        if self.version>=6:
            min_pos_moy=numpy.fromfile(self.fin, dtype=numpy.float32, count=n)
            max_pos_moy=numpy.fromfile(self.fin, dtype=numpy.float32, count=n)
        if self.version>=7:
            self.rms_emit[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=3)[:]
            self.rms_emit2[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=3)[:]
        if self.version>=8:
            e_ouv=numpy.fromfile(self.fin, dtype=numpy.float32, count=1)
            phase_ouv_pos=numpy.fromfile(self.fin, dtype=numpy.float32, count=1)
            phase_ouv_neg=numpy.fromfile(self.fin, dtype=numpy.float32, count=1)

        self.Np[self.i]=numpy.fromfile(self.fin, dtype=numpy.int64, count=1)[0]

        if self.Np[self.i]:
            powlost=numpy.zeros(self.Nrun)
            for i in xrange(self.Nrun):
                self.lost[self.i,i]=numpy.fromfile(self.fin, dtype=numpy.int64, count=1)[0]
                self.powlost[self.i,i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=1)[0]
            self.lost2[self.i]=numpy.fromfile(self.fin, dtype=numpy.int64, count=1)[0]
            self.Minlost[self.i]=numpy.fromfile(self.fin, dtype=numpy.int64, count=1)[0]
            self.Maxlost[self.i]=numpy.fromfile(self.fin, dtype=numpy.int64, count=1)[0]
            self.powlost2[self.i]=numpy.fromfile(self.fin, dtype=numpy.float64, count=1)[0]
            self.Minpowlost[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=1)[0]
            self.Maxpowlost[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=1)[0]

            if self.vlong==1:
                tab=numpy.fromfile(self.fin, dtype=numpy.uint64, count=n*step)
            else:
                tab=numpy.fromfile(self.fin, dtype=numpy.uint32, count=n*step)
    
            if self.ib[self.i]>0:
                tabp=numpy.fromfile(self.fin, dtype=numpy.uint32, count=3*step)

class remote_data_merger:
    def __init__(self, base='.'):
        self._base=base
        self._files=[]

    def add_file(self,filepath):
        import os

        if os.path.exists(filepath):
            fname=filepath
        else:
            fullpath=os.path.join(self._base,filepath)
            if os.path.exists(fullpath):
                fname=fullpath
            else:
                raise ValueError("Could not find file "+filepath)
        if fname not in self._files:
            self._files.append(fname)

    def generate_partran_out(self,filename=None):
        '''
        Creates a string to be written to file
        each line is a list.

        If filename is given, writes directly to output file.

        '''

        import numpy as np

        h1=[]
        h2=[]

        d1=[]
        d2=[]
        d3=[]

        if self._files:
            for f in self._files:
                string=file(f,'r').read()
                split=string.split('$$$')
                if split[9]!='Data_Error':
                    raise ValueError("Magic problem, please complain to Yngve")

                thisdata=split[10].strip().split('\n')

                if not h1:
                    h1=[thisdata[0]+" (std in paranthesis)"]
                    h2=thisdata[2:10]
                d1.append(thisdata[1].split())
                d2.append(thisdata[10])
                d3.append(thisdata[11])

            # fix d1:
            for i in xrange(len(d1)):
                for j in xrange(len(d1[0])):
                    d1[i][j]=float(d1[i][j])
            d1=np.array(d1)
            means=d1.mean(axis=0)
            stds=d1.std(axis=0)
            d1=[]
            for i in xrange(len(stds)):
                if stds[i]/means[i]<1e-10:
                    stds[i]=0.0
            for i in xrange(len(stds)):
                # some small std are removed..
                if stds[i]/means[i]>1e-8:
                    d1.append('%f(%f)' % (means[i],stds[i]))
                else: #error is 0
                    d1.append(str(means[i]))
            d1=[' '.join(d1)]

            # create data:
            data=h1+d1+h2+d2+d3

            if filename:
                file(filename,'w').write('\n'.join(data))

            return data

class partran(dict):
    '''
    Read partran1.out files..
    '''
    def __init__(self,filename):
        self.filename=filename
        self._readAsciiFile()

    def _readAsciiFile(self):

        import numpy

        stream=file(self.filename,'r')
        for i in xrange(10):
            line=stream.readline()
        self.columns=['NUM']+line.split()[1:]
        self.data=numpy.loadtxt(stream)

        self._dict={}
        for i in xrange(len(self.columns)):
            self[self.columns[i]]=self.data[:,i]