#!/usr/bin/env python

import numpy as N

from Fabric import file_io
from Fabric import variables

A_id = variables.refarray().A_id


class integrate:
    def __init__(self, direc='.', stem=None, grainno = 0, reflno = 0):
        self.reflno = reflno
        self.setno = 0
        self.grainno = grainno
        self.direc = direc
        self.stem = stem

    def read_data(self):
        try:
            self.hst = file_io.readhst(self.direc,self.stem)
            self.hst._open(self.grainno,self.setno)
        except:
            raise IOError, 'Integrate cannot readhst'
        #self.hst.get_file_pos()

        # Read .ref file for self.grainno 
        try:
            self.ref = file_io.readref(self.direc,self.stem)
            self.ref._open(self.grainno,self.setno)
            self.ref.read(self.grainno,self.setno)
        except:
            raise IOError, 'Could not open .ref file'
 
    def filter(self, LP_limit=N.inf):
        # Routine to filter list of reflection
        self.keep = N.ones((len(self.ref.refinfo),1))
        for i in range(len(self.ref.refinfo)):
            LP = self.ref.refinfo[i,self.ref.colinfo['L']]*\
                self.ref.refinfo[i,self.ref.colinfo['P']]
            if LP > LP_limit:
                self.keep[i] = 0
            if self.ref.refinfo[i,self.ref.colinfo['overlaps']] > 0:
                print self.ref.colinfo['overlaps']
                self.keep[i] = 0
            if self.ref.refinfo[i,self.ref.colinfo['border']] != 0:
                self.keep[i] = 0

    def run(self):
        self.int = N.zeros((len(self.ref.refinfo),2))
        for i in range(len(self.ref.refinfo)):
            self.hst.read(self.grainno,self.setno)
            # Make sure same refl is read from ref and hst
            ref_spot_id = self.ref.refinfo[i,self.ref.colinfo['spot_id']].astype(N.uint32)
            assert self.hst.spot_id  == ref_spot_id 

            # For now use median of shoebox as background
            median = N.median(self.hst.shoebox.ravel())
            elements = N.prod(self.hst.shoebox.shape)
            LP = self.ref.refinfo[i,self.ref.colinfo['L']]*\
              self.ref.refinfo[i,self.ref.colinfo['P']]
            bg = elements*median
            self.int[i,0] = self.hst.shoebox.sum()-bg
            self.int[i,0] = self.int[i,0]/LP
            self.int[i,1] = N.sqrt(abs(self.int[i,0])+(2*bg)/LP) #+0.02*abs(self.int[i,0])
#            self.int[i,1] = N.sqrt(abs(self.int[i,0]))+0.02*abs(self.int[i,0])
             

    def save(self):
            filename = '%s/%s_gr%0.4d.hkl' \
                %(self.direc,self.stem,self.grainno)
            f = open(filename,'w')
            for i in range(len(self.int)):
                if self.keep[i] == 1:
                    out = '%4i %4i %4i %8g %8g\n'\
                        %(self.ref.refinfo[i,self.ref.colinfo['h']],
                          self.ref.refinfo[i,self.ref.colinfo['k']],
                          self.ref.refinfo[i,self.ref.colinfo['l']],
                          self.int[i,0],self.int[i,1])
                    f.write(out)
            f.close()
