#!/usr/bin/python
#=================================================================#
#  Script to plot the bandstructure from an abinit bandstructure  #
#  _EIG.nc netcdf file or from a wannier bandstructure, or from   #
#  an _EIG.nc file+GW file+ bandstructure _EIG.nc file            #
#=================================================================#

#########
#IMPORTS#
#########

import numpy as N
import matplotlib.pyplot as P
import netCDF4 as nc
import sys
import os
import argparse
import time

#############
##VARIABLES##
#############

class VariableContainer:pass

#Constants
csts = VariableContainer()

csts.hartree2ev = N.float(27.211396132)
csts.ev2hartree = N.float(1/csts.hartree2ev)
csts.sqrtpi = N.float(N.sqrt(N.pi))
csts.invsqrtpi = N.float(1/csts.sqrtpi)
csts.TOLKPTS = N.float(0.00001)

###########
##CLASSES##
###########

class PolynomialFit(object):
    def __init__(self):
        self.degree = 2

class EigenvalueContainer(object):
    nsppol = None
    nkpt = None
    mband = None
    eigenvalues = None
    units = None
    wtk = None
    filename = None
    filefullpath = None
    bd_indices = None
    eigenvalue_type = None
    kpoints = None
    #kpoint_sampling_type: can be Monkhorst-Pack or Bandstructure
    KPT_W90_TOL = N.float(1.0e-6)
    KPT_DFT_TOL = N.float(1.0e-8)
    kpoint_sampling_type = 'Monkhorst-Pack'
    inputgvectors = None
    gvectors = None
    special_kpoints = None
    special_kpoints_names = None
    special_kpoints_indices = None
    kpoint_path_values = None
    kpoint_reduced_path_values = None
    kpoint_path_length = None
    #reduced_norm = None
    norm_paths = None
    norm_reduced_paths = None
    def __init__(self,directory=None,filename=None):
        if filename == None:return
        if directory == None:directory='.'
        self.filename = filename
        self.filefullpath = '%s/%s' %(directory,filename)
        self.file_open(self.filefullpath)
    def set_kpoint_sampling_type(self,kpoint_sampling_type):
        if kpoint_sampling_type != 'Monkhorst-Pack' and kpoint_sampling_type != 'Bandstructure':
            print 'ERROR: kpoint_sampling_type "%s" does not exists' %kpoint_sampling_type
            print '       it should be "Monkhorst-Pack" or "Bandstructure" ... exit'
            sys.exit()
        self.kpoint_sampling_type = kpoint_sampling_type
    def correct_kpt(self,kpoint,tolerance=N.float(1.0e-6)):
        kpt_correct = N.array(kpoint,N.float)
        changed = False
        for ii in range(3):
            if N.allclose(kpoint[ii],N.float(1.0/3.0),atol=tolerance):
                kpt_correct[ii] = N.float(1.0/3.0)
                changed = True
            elif N.allclose(kpoint[ii],N.float(1.0/6.0),atol=tolerance):
                kpt_correct[ii] = N.float(1.0/6.0)
                changed = True
            elif N.allclose(kpoint[ii],N.float(-1.0/6.0),atol=tolerance):
                kpt_correct[ii] = N.float(-1.0/6.0)
                changed = True
            elif N.allclose(kpoint[ii],N.float(-1.0/3.0),atol=tolerance):
                kpt_correct[ii] = N.float(-1.0/3.0)
                changed = True
        if changed:
            print 'COMMENT: kpoint %15.12f %15.12f %15.12f has been changed to %15.12f %15.12f %15.12f' %(kpoint[0],kpoint[1],kpoint[2],kpt_correct[0],kpt_correct[1],kpt_correct[2])
        return kpt_correct
    def find_special_kpoints(self,gvectors=None):
        if self.kpoint_sampling_type != 'Bandstructure':
            print 'ERROR: special kpoints are usefull only for bandstructures ... returning find_special_kpoints'
            return
        if self.eigenvalue_type == 'W90':
            correct_kpt_tolerance = N.float(1.0e-4)
            KPT_TOL = self.KPT_W90_TOL
        elif self.eigenvalue_type == 'DFT':
            correct_kpt_tolerance = N.float(1.0e-6)
            KPT_TOL = self.KPT_DFT_TOL
        else:
            print 'ERROR: eigenvalue_type is "%s" while it should be "W90" or "DFT" ... returning find_special_kpoints' %self.eigenvalue_type
            return
        if gvectors == None:
            self.inputgvectors = False
            self.gvectors = N.identity(3,N.float)
        else:
            if N.shape(gvectors) != (3, 3):
                print 'ERROR: wrong gvectors ... exiting now'
                sys.exit()
            self.inputgvectors = True
            self.gvectors = gvectors
        full_kpoints = N.zeros((self.nkpt,3),N.float)
        for ikpt in range(self.nkpt):
            full_kpoints[ikpt,:] = self.kpoints[ikpt,0]*self.gvectors[0,:]+self.kpoints[ikpt,1]*self.gvectors[1,:]+self.kpoints[ikpt,2]*self.gvectors[2,:]
        delta_kpt = full_kpoints[1,:]-full_kpoints[0,:]
        self.special_kpoints_indices = list()
        self.special_kpoints = list()
        self.special_kpoints_indices.append(0)
        self.special_kpoints.append(self.correct_kpt(self.kpoints[0,:],tolerance=correct_kpt_tolerance))
        for ikpt in range(1,self.nkpt-1):
            thisdelta = full_kpoints[ikpt+1,:]-full_kpoints[ikpt,:]
            if not N.allclose(thisdelta,delta_kpt,atol=KPT_TOL):
                delta_kpt = thisdelta
                self.special_kpoints_indices.append(ikpt)
                self.special_kpoints.append(self.correct_kpt(self.kpoints[ikpt,:],tolerance=correct_kpt_tolerance))
        self.special_kpoints_indices.append(N.shape(self.kpoints)[0]-1)
        self.special_kpoints.append(self.correct_kpt(self.kpoints[-1,:],tolerance=correct_kpt_tolerance))
        print 'Special Kpoints : '
        print ' {0:d} : {1[0]: 8.8f} {1[1]: 8.8f} {1[2]: 8.8f}'.format(1,self.kpoints[0,:])
        self.norm_paths = N.zeros((N.shape(self.special_kpoints_indices)[0]-1),N.float)
        self.norm_reduced_paths = N.zeros((N.shape(self.special_kpoints_indices)[0]-1),N.float)
        for ispkpt in range(1,N.shape(self.special_kpoints_indices)[0]):
            self.norm_paths[ispkpt-1] = N.linalg.norm(full_kpoints[self.special_kpoints_indices[ispkpt]]-full_kpoints[self.special_kpoints_indices[ispkpt-1]])
            self.norm_reduced_paths[ispkpt-1] = N.linalg.norm(self.special_kpoints[ispkpt]-self.special_kpoints[ispkpt-1])
            print '     {2:d}-{3:d} path length : {0: 8.8f} | reduced path length : {1: 8.8f}'.\
                  format(self.norm_paths[ispkpt-1],self.norm_reduced_paths[ispkpt-1],ispkpt,ispkpt+1)
            print ' {0:d} : {1[0]: 8.8f} {1[1]: 8.8f} {1[2]: 8.8f}'.format(ispkpt+1,self.kpoints[self.special_kpoints_indices[ispkpt],:])
        self.kpoint_path_length = N.sum(self.norm_paths)
        self.kpoint_reduced_path_length = N.sum(self.norm_reduced_paths)
        self.normalized_kpoint_path_norm = self.norm_paths/self.kpoint_path_length
        self.normalized_kpoint_reduced_path_norm = self.norm_reduced_paths/self.kpoint_reduced_path_length

        kptredpathval = list()
        kptpathval = list()
        kptredpathval.append(N.float(0.0))
        kptpathval.append(N.float(0.0))
        curlen = N.float(0.0)
        redcurlen = N.float(0.0)
        for ispkpt in range(1,N.shape(self.special_kpoints_indices)[0]):
            kptredpathval.extend(N.linspace(redcurlen,redcurlen+self.norm_reduced_paths[ispkpt-1],self.special_kpoints_indices[ispkpt]-self.special_kpoints_indices[ispkpt-1]+1)[1:])
            kptpathval.extend(N.linspace(curlen,curlen+self.norm_paths[ispkpt-1],self.special_kpoints_indices[ispkpt]-self.special_kpoints_indices[ispkpt-1]+1)[1:])
            redcurlen = redcurlen + self.norm_reduced_paths[ispkpt-1]
            curlen = curlen + self.norm_paths[ispkpt-1]
        self.kpoint_path_values = N.array(kptpathval,N.float)
        self.kpoint_reduced_path_values = N.array(kptredpathval,N.float)
        self.normalized_kpoint_path_values = self.kpoint_path_values/self.kpoint_path_length
        self.normalized_kpoint_reduced_path_values = self.kpoint_reduced_path_values/self.kpoint_reduced_path_length
        self.special_kpoints = N.array(self.special_kpoints,N.float)
    def file_open(self,filefullpath):
        if filefullpath[-3:] == '_GW':
            self.gw_file_open(filefullpath)
        elif filefullpath[-7:] == '_EIG.nc':
            self.nc_eig_open(filefullpath)
        elif filefullpath[-4:] == '.dat':
            self.wannier_bs_file_open(filefullpath)
    def has_eigenvalue(self,nsppol,isppol,kpoint,iband):
        if self.nsppol != nsppol:
            return False
        for ikpt in range(self.nkpt):
            if N.absolute(self.kpoints[ikpt,0]-kpoint[0]) < csts.TOLKPTS and \
               N.absolute(self.kpoints[ikpt,1]-kpoint[1]) < csts.TOLKPTS and \
               N.absolute(self.kpoints[ikpt,2]-kpoint[2]) < csts.TOLKPTS:
                if iband >= self.bd_indices[isppol,ikpt,0]-1 and iband < self.bd_indices[isppol,ikpt,1]:
                    return True
                return False
        return False
    def get_eigenvalue(self,nsppol,isppol,kpoint,iband):
        for ikpt in range(self.nkpt):
            if N.absolute(self.kpoints[ikpt,0]-kpoint[0]) < csts.TOLKPTS and \
               N.absolute(self.kpoints[ikpt,1]-kpoint[1]) < csts.TOLKPTS and \
               N.absolute(self.kpoints[ikpt,2]-kpoint[2]) < csts.TOLKPTS:
                return self.eigenvalues[isppol,ikpt,iband]
    def wannier_bs_file_open(self,filefullpath):
        if not (os.path.isfile(filefullpath)):
            print 'ERROR : file "%s" does not exists' %filefullpath
            print '... exiting now ...'
            sys.exit()
        print 'WARNING: no spin polarization reading yet for Wannier90 bandstructure files!'
        self.eigenvalue_type = 'W90'
        self.nsppol = None
        self.nkpt = None
        self.mband = None
        self.eigenvalues = None
        self.units = None
        self.filefullpath = filefullpath
        reader = open(self.filefullpath,'r')
        filedata = reader.readlines()
        reader.close()
        for iline in range(len(filedata)):
            if filedata[iline].strip() == '':
                self.nkpt = iline
                break
        self.mband = N.int(len(filedata)/self.nkpt)
        self.nsppol = 1
        self.eigenvalues = N.zeros([self.nsppol,self.nkpt,self.mband],N.float)
        self.kpoints = N.zeros([self.nkpt,3],N.float)
        iline = 0
        kpt_file = '%s.kpt' %filefullpath[:-4]
        if os.path.isfile(kpt_file):
            reader = open(kpt_file,'r')
            kptdata = reader.readlines()
            reader.close()
            if N.int(kptdata[0]) != self.nkpt:
                print 'ERROR : the number of kpoints in file "%s" is not the same as in "%s" ... exit' %(self.filefullpath,kpt_file)
                sys.exit()
            for ikpt in range(self.nkpt):
                linesplit = kptdata[ikpt+1].split()
                self.kpoints[ikpt,0] = N.float(linesplit[0])
                self.kpoints[ikpt,1] = N.float(linesplit[1])
                self.kpoints[ikpt,2] = N.float(linesplit[2])
        else:
            for ikpt in range(self.nkpt):
                self.kpoints[ikpt,0] = N.float(filedata[ikpt].split()[0])
        for iband in range(self.mband):
            for ikpt in range(self.nkpt):
                self.eigenvalues[0,ikpt,iband] = N.float(filedata[iline].split()[1])
                iline = iline+1
            iline = iline+1
        self.eigenvalues = self.eigenvalues*csts.ev2hartree
        self.units = 'Hartree'
    def gw_file_open(self,filefullpath):
        if not (os.path.isfile(filefullpath)):
            print 'ERROR : file "%s" does not exists' %filefullpath
            print '... exiting now ...'
            sys.exit()
        self.eigenvalue_type = 'GW'
        self.nsppol = None
        self.nkpt = None
        self.mband = None
        self.eigenvalues = None
        self.units = None
        self.filefullpath = filefullpath
        reader = open(self.filefullpath,'r')
        filedata = reader.readlines()
        reader.close()
        self.nkpt = N.int(filedata[0].split()[0])
        self.kpoints = N.ones([self.nkpt,3],N.float)
        self.nsppol = N.int(filedata[0].split()[1])
        self.bd_indices = N.zeros((self.nsppol,self.nkpt,2),N.int)
        icur = 1
        nbd_kpt = N.zeros([self.nsppol,self.nkpt],N.int)
        for isppol in range(self.nsppol):
            for ikpt in range(self.nkpt):
                self.kpoints[ikpt,:] = N.array(filedata[icur].split()[:],N.float)
                icur = icur + 1
                nbd_kpt[isppol,ikpt] = N.int(filedata[icur])
                self.bd_indices[isppol,ikpt,0] = N.int(filedata[icur+1].split()[0])
                self.bd_indices[isppol,ikpt,1] = N.int(filedata[icur+nbd_kpt[isppol,ikpt]].split()[0])
                icur = icur + nbd_kpt[isppol,ikpt] + 1
        self.mband = N.max(self.bd_indices[:,:,1])
        self.eigenvalues = N.zeros([self.nsppol,self.nkpt,self.mband],N.float)
        self.eigenvalues[:,:,:] = N.nan
        ii = 3
        for isppol in range(self.nsppol):
            for ikpt in range(self.nkpt):
                for iband in range(self.bd_indices[isppol,ikpt,0]-1,self.bd_indices[isppol,ikpt,1]):
                    self.eigenvalues[isppol,ikpt,iband] = N.float(filedata[ii].split()[1])
                    ii = ii + 1
                ii = ii + 2
        self.eigenvalues = csts.ev2hartree*self.eigenvalues
        self.units = 'Hartree'
    def pfit_gw_file_write(self,polyfitlist,directory=None,filename=None,bdgw=None,energy_pivots=None,gwec=None):
        if filename == None:return
        if directory == None:directory='.'
        filefullpath = '%s/%s' %(directory,filename)
        if (os.path.isfile(filefullpath)):
            user_input = raw_input('WARNING : file "%s" exists, do you want to overwrite it ? (y/n)' %filefullpath)
            if not (user_input == 'y' or user_input == 'Y'):
                return
        writer = open(filefullpath,'w')
        writer.write('%12s%12s\n' %(self.nkpt,self.nsppol))
        if gwec == None:
            for ikpt in range(self.nkpt):
                for isppol in range(self.nsppol):
                    writer.write('%10.6f%10.6f%10.6f\n' %(self.kpoints[ikpt,0],self.kpoints[ikpt,1],self.kpoints[ikpt,2]))
                    writer.write('%4i\n' %(bdgw[1]-bdgw[0]+1))
                    for iband in range(bdgw[0]-1,bdgw[1]):
                        delta = N.polyval(polyfitlist[-1],csts.hartree2ev*self.eigenvalues[isppol,ikpt,iband])
                        for ipivot in range(len(energy_pivots)):
                            if csts.hartree2ev*self.eigenvalues[isppol,ikpt,iband] <= energy_pivots[ipivot]:
                                delta = N.polyval(polyfitlist[ipivot],csts.hartree2ev*self.eigenvalues[isppol,ikpt,iband])
                                break
                        writer.write('%6i%9.4f%9.4f%9.4f\n' %(iband+1,csts.hartree2ev*self.eigenvalues[isppol,ikpt,iband]+delta,delta,0.0))
        else:
            for ikpt in range(self.nkpt):
                for isppol in range(self.nsppol):
                    writer.write('%10.6f%10.6f%10.6f\n' %(self.kpoints[ikpt,0],self.kpoints[ikpt,1],self.kpoints[ikpt,2]))
                    writer.write('%4i\n' %(bdgw[1]-bdgw[0]+1))
                    for iband in range(bdgw[0]-1,bdgw[1]):
                        if gwec.has_eigenvalue(self.nsppol,isppol,self.kpoints[ikpt],iband):
                            gw_eig = gwec.get_eigenvalue(self.nsppol,isppol,self.kpoints[ikpt],iband)
                            writer.write('%6i%9.4f%9.4f%9.4f\n' %(iband+1,csts.hartree2ev*gw_eig,csts.hartree2ev*(gw_eig-self.eigenvalues[isppol,ikpt,iband]),0.0))
                        else:
                            delta = N.polyval(polyfitlist[-1],csts.hartree2ev*self.eigenvalues[isppol,ikpt,iband])
                            for ipivot in range(len(energy_pivots)):
                                if csts.hartree2ev*self.eigenvalues[isppol,ikpt,iband] <= energy_pivots[ipivot]:
                                    delta = N.polyval(polyfitlist[ipivot],csts.hartree2ev*self.eigenvalues[isppol,ikpt,iband])
                                    break
                            writer.write('%6i%9.4f%9.4f%9.4f\n' %(iband+1,csts.hartree2ev*self.eigenvalues[isppol,ikpt,iband]+delta,delta,0.0))
        writer.close()
    def pfit_dft_to_gw_bs_write(self,polyfitlist,directory=None,filename=None,bdgw=None,energy_pivots=None,gwec=None):
        if filename == None:return
        if directory == None:directory='.'
        filefullpath = '%s/%s' %(directory,filename)
        if (os.path.isfile(filefullpath)):
            user_input = raw_input('WARNING : file "%s" exists, do you want to overwrite it ? (y/n)' %filefullpath)
            if not (user_input == 'y' or user_input == 'Y'):
                return
        writer = open(filefullpath,'w')
        if gwec == None:
            for ikpt in range(self.nkpt):
                writer.write('%s' %ikpt)
                for isppol in range(self.nsppol):
                    for iband in range(bdgw[0]-1,bdgw[1]):
                        delta = N.polyval(polyfitlist[-1],csts.hartree2ev*self.eigenvalues[isppol,ikpt,iband])
                        for ipivot in range(len(energy_pivots)):
                            if csts.hartree2ev*self.eigenvalues[isppol,ikpt,iband] <= energy_pivots[ipivot]:
                                delta = N.polyval(polyfitlist[ipivot],csts.hartree2ev*self.eigenvalues[isppol,ikpt,iband])
                                break
                        writer.write(' %s' %(csts.hartree2ev*self.eigenvalues[isppol,ikpt,iband]+delta))
                writer.write('\n')
        else:
            print 'NOT SUPPORTED YET'
            sys.exit()
        writer.close()
    def nc_eig_open(self,filefullpath):
        if not (os.path.isfile(filefullpath)):
            print 'ERROR : file "%s" does not exists' %filefullpath
            print '... exiting now ...'
            sys.exit()
        ncdata = nc.Dataset(filefullpath)
        self.eigenvalue_type = 'DFT'
        self.nsppol = None
        self.nkpt = None
        self.mband = None
        self.eigenvalues = None
        self.units = None
        self.filefullpath = filefullpath
        for dimname,dimobj in ncdata.dimensions.iteritems():
            if dimname == 'nsppol':self.nsppol = N.int(len(dimobj))
            if dimname == 'nkpt':self.nkpt = N.int(len(dimobj))
            if dimname == 'mband':self.mband = N.int(len(dimobj))
        for varname in ncdata.variables:
            if varname == 'Eigenvalues':
                varobj = ncdata.variables[varname]
                varshape = N.shape(varobj[:])
                self.units = None
                for attrname in varobj.ncattrs():
                    if attrname == 'units':
                        self.units = varobj.getncattr(attrname)
                if self.units == None:
                    print 'WARNING : units are not specified'
                    print '... assuming "Hartree" units ...'
                    self.units = 'Hartree'
                elif self.units != 'Hartree':
                    print 'ERROR : units are unknown : "%s"' %self.units
                    print '... exiting now ...'
                    sys.exit()
                self.eigenvalues = N.reshape(N.array(varobj,N.float),varshape)
                self.nsppol = varshape[0]
                self.nkpt = varshape[1]
                self.kpoints = -1*N.ones((self.nkpt,3),N.float)
                self.mband = varshape[2]
                self.bd_indices = N.zeros((self.nsppol,self.nkpt,2),N.int)
                self.bd_indices[:,:,0] = 1
                self.bd_indices[:,:,1] = self.mband
                break
        for varname in ncdata.variables:
            if varname == 'Kptns':
                varobj = ncdata.variables[varname]
                varshape = N.shape(varobj[:])
                self.kpoints = N.reshape(N.array(varobj,N.float),varshape)
    def write_bandstructure_to_file(self,filename,option_kpts='bohrm1_units'):
        #if option_kpts is set to 'normalized', the path of the bandstructure will be normalized to 1 (and special k-points correctly chosen)
        if self.kpoint_sampling_type != 'Bandstructure':
            print 'ERROR: kpoint_sampling_type is not "Bandstructure" ... returning from write_bandstructure_to_file'
            return
        if self.nsppol > 1:
            print 'ERROR: number of spins is more than 1, this is not fully tested ... use with care !'
        writer = open(filename,'w')
        writer.write('# BANDSTRUCTURE FILE FROM DAVID\'S SCRIPT\n')
        writer.write('# nsppol = %s\n' %self.nsppol)
        writer.write('# nband = %s\n' %self.mband)
        writer.write('# eigenvalue_type = %s\n' %self.eigenvalue_type)
        if self.inputgvectors:
            writer.write('# inputgvectors = 1 (%s)\n' %self.inputgvectors)
        else:
            writer.write('# inputgvectors = 0 (%s)\n' %self.inputgvectors)
        writer.write('# gvectors(1) = %20.17f %20.17f %20.17f \n' %(self.gvectors[0,0],self.gvectors[0,1],self.gvectors[0,2]))
        writer.write('# gvectors(2) = %20.17f %20.17f %20.17f \n' %(self.gvectors[1,0],self.gvectors[1,1],self.gvectors[1,2]))
        writer.write('# gvectors(3) = %20.17f %20.17f %20.17f \n' %(self.gvectors[2,0],self.gvectors[2,1],self.gvectors[2,2]))
        writer.write('# special_kpoints_number = %s\n' %(len(self.special_kpoints_indices)))
        writer.write('# list of special kpoints : (given in reduced coordinates, value_path is in Bohr^-1, value_red_path has its total path normalized to 1)\n')
        for ii in range(len(self.special_kpoints_indices)):
            ispkpt = self.special_kpoints_indices[ii]
            spkpt = self.special_kpoints[ii]
            writer.write('#    special_kpt_index %5s : %20.17f %20.17f %20.17f (value_path = %20.17f | value_red_path = %20.17f)\n' %(ispkpt,spkpt[0],spkpt[1],spkpt[2],self.kpoint_path_values[ispkpt],self.kpoint_reduced_path_values[ispkpt]))
        writer.write('# special_kpoints_names :\n')
        for ii in range(len(self.special_kpoints_indices)):
            ispkpt = self.special_kpoints_indices[ii]
            spkpt = self.special_kpoints[ii]
            writer.write('#    special_kpt_name %3s : "%s" : %20.17f %20.17f %20.17f\n' %(ii+1,self.special_kpoints_names[ii],spkpt[0],spkpt[1],spkpt[2]))
        writer.write('# kpoint_path_length = %20.17f \n' %(self.kpoint_path_length))
        writer.write('# kpoint_path_number = %s \n' %(self.nkpt))
        if self.inputgvectors:
            writer.write('# kpoint_path_units = %s\n' %(option_kpts))
        else:
            writer.write('# kpoint_path_units =  %s (!!! CONSIDERING UNITARY GVECTORS MATRIX !!!)\n' %(option_kpts))
        writer.write('#BEGIN\n')
        if option_kpts == 'bohrm1_units':
            values_path = self.kpoint_path_values
        elif option_kpts == 'reduced':
            values_path = self.kpoint_reduced_path_values
        elif option_kpts == 'bohrm1_units_normalized':
            values_path = self.normalized_kpoint_path_values
        elif option_kpts == 'reduced_normalized':
            values_path = self.normalized_kpoint_reduced_path_values
        else:
            print 'ERROR: wrong option_kpts ... exit'
            writer.write('... CANCELLED (wrong option_kpts)')
            writer.close()
            sys.exit()
        for isppol in range(self.nsppol):
            writer.write('#isppol %s\n' %isppol)
            for iband in range(self.mband):
                writer.write('#iband %5s (band number %s)\n' %(iband,iband+1))
                for ikpt in range(self.nkpt):
                    writer.write('%20.17f %20.17f\n' %(values_path[ikpt],self.eigenvalues[isppol,ikpt,iband]))
                writer.write('\n')
        writer.write('#END\n')
        writer.write('\n#KPT_LIST\n')
        for ikpt in range(self.nkpt):
            writer.write('# %6d : %20.17f %20.17f %20.17f\n' %(ikpt,self.kpoints[ikpt,0],self.kpoints[ikpt,1],self.kpoints[ikpt,2]))
        writer.close()
    def read_bandstructure_from_file(self,filename):
        reader = open(filename,'r')
        bs_data = reader.readlines()
        reader.close()
        self.gvectors = N.identity(3,N.float)
        self.kpoint_sampling_type = 'Bandstructure'
        self.special_kpoints_indices = list()
        self.special_kpoints = list()
        for ii in range(len(bs_data)):
            if bs_data[ii] == '#BEGIN\n':
                ibegin = ii
                break
            elif bs_data[ii][:10] == '# nsppol =':
                self.nsppol = N.int(bs_data[ii][10:])
            elif bs_data[ii][:9] == '# nband =':
                self.mband = N.int(bs_data[ii][9:])
            elif bs_data[ii][:19] == '# eigenvalue_type =':
                self.eigenvalue_type = bs_data[ii][19:].strip()
            elif bs_data[ii][:17] == '# inputgvectors =':
                tt = N.int(bs_data[ii][18])
                if tt == 1:
                    self.inputgvectors = True
                elif tt == 0:
                    self.inputgvectors = False
                else:
                    print 'ERROR: reading inputgvectors ... exit'
                    sys.exit()
            elif bs_data[ii][:15] == '# gvectors(1) =':
                sp = bs_data[ii][15:].split()
                self.gvectors[0,0] = N.float(sp[0])
                self.gvectors[0,1] = N.float(sp[1])
                self.gvectors[0,2] = N.float(sp[2])
            elif bs_data[ii][:15] == '# gvectors(2) =':
                sp = bs_data[ii][15:].split()
                self.gvectors[1,0] = N.float(sp[0])
                self.gvectors[1,1] = N.float(sp[1])
                self.gvectors[1,2] = N.float(sp[2])
            elif bs_data[ii][:15] == '# gvectors(3) =':
                sp = bs_data[ii][15:].split()
                self.gvectors[2,0] = N.float(sp[0])
                self.gvectors[2,1] = N.float(sp[1])
                self.gvectors[2,2] = N.float(sp[2])
            elif bs_data[ii][:26] == '# special_kpoints_number =':
                special_kpoints_number = N.int(bs_data[ii][26:])
                self.special_kpoints_names = ['']*special_kpoints_number
            elif bs_data[ii][:22] == '#    special_kpt_index':
                sp = bs_data[ii][22:].split()
                self.special_kpoints_indices.append(N.int(sp[0]))
                self.special_kpoints.append(N.array([sp[2],sp[3],sp[4]]))
            elif bs_data[ii][:21] == '#    special_kpt_name':
                sp = bs_data[ii][21:].split()
                ispkpt = N.int(sp[0])-1
                self.special_kpoints_names[ispkpt] = sp[2][1:-1]
            elif bs_data[ii][:22] == '# kpoint_path_length =':
                self.kpoint_path_length = N.float(bs_data[ii][22:])
            elif bs_data[ii][:22] == '# kpoint_path_number =':
                self.nkpt = N.int(bs_data[ii][22:])
            elif bs_data[ii][:21] == '# kpoint_path_units =':
                kpoint_path_units = bs_data[ii][21:].strip()
        self.special_kpoints_indices = N.array(self.special_kpoints_indices,N.int)
        self.special_kpoints = N.array(self.special_kpoints,N.float)
        if len(self.special_kpoints_indices) != special_kpoints_number or len(self.special_kpoints) != special_kpoints_number:
            print 'ERROR: reading the special kpoints ... exit'
            sys.exit()
        self.kpoint_path_values = N.zeros([self.nkpt],N.float)
        self.kpoint_reduced_path_values = N.zeros([self.nkpt],N.float)
        if kpoint_path_units == 'bohrm1_units':
            jj = 0
            for ii in range(ibegin+1,len(bs_data)):
                if bs_data[ii][:7] == '#isppol' or bs_data[ii][:6] == '#iband':continue
                if bs_data[ii] == '\n':
                    break
                self.kpoint_path_values[jj] = N.float(bs_data[ii].split()[0])
                jj = jj + 1
            if jj != self.nkpt:
                print 'ERROR: reading bandstructure file ... exit'
                sys.exit()
            self.normalized_kpoint_path_values = self.kpoint_path_values/self.kpoint_path_length
        if kpoint_path_units == 'bohrm1_units_normalized':
            jj = 0
            for ii in range(ibegin+1,len(bs_data)):
                if bs_data[ii][:7] == '#isppol' or bs_data[ii][:6] == '#iband':continue
                if bs_data[ii] == '\n':
                    break
                self.normalized_kpoint_path_values[jj] = N.float(bs_data[ii].split()[0])
                jj = jj + 1
            if jj != self.nkpt:
                print 'ERROR: reading bandstructure file ... exit'
                sys.exit()
            self.kpoint_path_values = self.normalized_kpoint_path_values*self.kpoint_path_length
        elif kpoint_path_units == 'reduced_normalized':
            jj = 0
            for ii in range(ibegin+1,len(bs_data)):
                if bs_data[ii][:7] == '#isppol' or bs_data[ii][:6] == '#iband':continue
                if bs_data[ii] == '\n':
                    break
                self.normalized_kpoint_reduced_path_values[jj] = N.float(bs_data[ii].split()[0])
                jj = jj + 1
            if jj != self.nkpt:
                print 'ERROR: reading bandstructure file ... exit'
                sys.exit()
            self.kpoint_reduced_path_values = self.normalized_kpoint_reduced_path_values/self.kpoint_reduced_path_length
        elif kpoint_path_units == 'reduced':
            jj = 0
            for ii in range(ibegin+1,len(bs_data)):
                if bs_data[ii][:7] == '#isppol' or bs_data[ii][:6] == '#iband':continue
                if bs_data[ii] == '\n':
                    break
                self.kpoint_reduced_path_values[jj] = N.float(bs_data[ii].split()[0])
                jj = jj + 1
            if jj != self.nkpt:
                print 'ERROR: reading bandstructure file ... exit'
                sys.exit()
            self.normalized_kpoint_reduced_path_values = self.kpoint_reduced_path_values/self.kpoint_reduced_path_length
        self.eigenvalues = N.zeros([self.nsppol,self.nkpt,self.mband],N.float)
        check_nband = 0
        for ii in range(ibegin+1,len(bs_data)):
            if bs_data[ii][:7] == '#isppol':
                isppol = N.int(bs_data[ii][7:])
            elif bs_data[ii][:6] == '#iband':
                iband = N.int(bs_data[ii][6:].split()[0])
                ikpt = 0
            elif bs_data[ii][:4] == '#END':
                break
            elif bs_data[ii] == '\n':
                check_nband = check_nband + 1
            else:
                self.eigenvalues[isppol,ikpt,iband] = N.float(bs_data[ii].split()[1])
                ikpt = ikpt + 1

def check_gw_vs_dft_parameters(dftec,gwec):
    if gwec.eigenvalue_type != 'GW' or dftec.eigenvalue_type != 'DFT':
        print 'ERROR: eigenvalue files do not contain GW and DFT eigenvalues ... exiting now'
        sys.exit()
    if dftec.nsppol != gwec.nsppol or dftec.nkpt != gwec.nkpt:
        print 'ERROR: the number of spins/kpoints is not the same in the GW and DFT files used to make the interpolation ... exiting now'
        sys.exit()
    for ikpt in range(dftec.nkpt):
        if N.absolute(dftec.kpoints[ikpt,0]-gwec.kpoints[ikpt,0]) > csts.TOLKPTS or \
           N.absolute(dftec.kpoints[ikpt,1]-gwec.kpoints[ikpt,1]) > csts.TOLKPTS or \
           N.absolute(dftec.kpoints[ikpt,2]-gwec.kpoints[ikpt,2]) > csts.TOLKPTS:
            print 'ERROR: the kpoints are not the same in the GW and DFT files used to make the interpolation ... exiting now'
            sys.exit()

def plot_gw_vs_dft_eig(dftec,gwec,vbm_index,energy_pivots=None,polyfit_degrees=None):
    if gwec.eigenvalue_type != 'GW' or dftec.eigenvalue_type != 'DFT':
        print 'ERROR: eigenvalue containers do not contain GW and DFT eigenvalues ... exiting now'
        sys.exit()
    if dftec.nsppol != gwec.nsppol or dftec.nkpt != gwec.nkpt:
        print 'ERROR: the number of spins/kpoints is not the same in the GW and DFT containers ... exiting now'
        sys.exit()
    valdftarray = N.array([],N.float)
    conddftarray = N.array([],N.float)
    valgwarray = N.array([],N.float)
    condgwarray = N.array([],N.float)
    for isppol in range(dftec.nsppol):
        for ikpt in range(dftec.nkpt):
            ibdmin = N.max([dftec.bd_indices[isppol,ikpt,0],gwec.bd_indices[isppol,ikpt,0]])-1
            ibdmax = N.min([dftec.bd_indices[isppol,ikpt,1],gwec.bd_indices[isppol,ikpt,1]])-1
            valdftarray = N.append(valdftarray,csts.hartree2ev*dftec.eigenvalues[isppol,ikpt,ibdmin:vbm_index])
            valgwarray = N.append(valgwarray,csts.hartree2ev*gwec.eigenvalues[isppol,ikpt,ibdmin:vbm_index])
            conddftarray = N.append(conddftarray,csts.hartree2ev*dftec.eigenvalues[isppol,ikpt,vbm_index:ibdmax+1])
            condgwarray = N.append(condgwarray,csts.hartree2ev*gwec.eigenvalues[isppol,ikpt,vbm_index:ibdmax+1])
    if energy_pivots == None:
        if plot_figures == 1:
            P.figure(1)
            P.hold(True)
            P.grid(True)
            P.plot(valdftarray,valgwarray,'bx')
            P.plot(conddftarray,condgwarray,'rx')
            P.xlabel('DFT eigenvalues (in eV)')
            P.ylabel('GW eigenvalues (in eV)')
            P.figure(2)
            P.hold(True)
            P.grid(True)
            P.plot(valdftarray,valgwarray-valdftarray,'bx')
            P.plot(conddftarray,condgwarray-conddftarray,'rx')
            P.xlabel('DFT eigenvalues (in eV)')
            P.ylabel('GW correction to the DFT eigenvalues (in eV)')
            P.show()
            return
    polyfitlist = list()
    if len(polyfit_degrees) == 1:
        print 'ERROR: making a fit with only one interval is not allowed ... exiting now'
        sys.exit()
    dftarray = N.append(valdftarray,conddftarray)
    gwarray = N.append(valgwarray,condgwarray)
    dftarray_list = list()
    gwarray_list = list()
    for iinterval in range(len(polyfit_degrees)):
        tmpdftarray = N.array([],N.float)
        tmpgwarray = N.array([],N.float)
        if iinterval == 0:
            emin = None
            emax = energy_pivots[0]
            for ii in range(len(dftarray)):
                if dftarray[ii] <= emax:
                    tmpdftarray = N.append(tmpdftarray,[dftarray[ii]])
                    tmpgwarray = N.append(tmpgwarray,[gwarray[ii]])
        elif iinterval == len(polyfit_degrees)-1:
            emin = energy_pivots[-1]
            emax = None
            for ii in range(len(dftarray)):
                if dftarray[ii] >= emin:
                    tmpdftarray = N.append(tmpdftarray,[dftarray[ii]])
                    tmpgwarray = N.append(tmpgwarray,[gwarray[ii]])
        else:
            emin = energy_pivots[iinterval-1]
            emax = energy_pivots[iinterval]
            for ii in range(len(dftarray)):
                if dftarray[ii] >= emin and dftarray[ii] <= emax:
                    tmpdftarray = N.append(tmpdftarray,[dftarray[ii]])
                    tmpgwarray = N.append(tmpgwarray,[gwarray[ii]])
        dftarray_list.append(tmpdftarray)
        gwarray_list.append(tmpgwarray)
        pfit = N.polyfit(tmpdftarray,tmpgwarray-tmpdftarray,polyfit_degrees[iinterval])
        polyfitlist.append(pfit)
    if plot_figures == 1:
        linspace_npoints = 200
        valpoly_x = N.linspace(N.min(valdftarray),N.max(valdftarray),linspace_npoints)
        condpoly_x = N.linspace(N.min(conddftarray),N.max(conddftarray),linspace_npoints)
        P.figure(3)
        P.hold(True)
        P.grid(True)
        P.plot(valdftarray,valgwarray-valdftarray,'bx')
        P.plot(conddftarray,condgwarray-conddftarray,'rx')
        [x_min,x_max] = P.xlim()
        for iinterval in range(len(polyfit_degrees)):
            if iinterval == 0:
                tmppoly_x = N.linspace(x_min,energy_pivots[iinterval],linspace_npoints)
            elif iinterval == len(polyfit_degrees)-1:
                tmppoly_x = N.linspace(energy_pivots[iinterval-1],x_max,linspace_npoints)
            else:
                tmppoly_x = N.linspace(energy_pivots[iinterval-1],energy_pivots[iinterval],linspace_npoints)
            P.plot(tmppoly_x,N.polyval(polyfitlist[iinterval],tmppoly_x),'k')
        for ipivot in range(len(energy_pivots)):
            en = energy_pivots[ipivot]
            P.plot([en,en],[N.polyval(polyfitlist[ipivot],en),N.polyval(polyfitlist[ipivot+1],en)],'k-.')
        P.xlabel('DFT eigenvalues (in eV)')
        P.ylabel('GW correction to the DFT eigenvalues (in eV)')
        P.figure(4)
        P.hold(True)
        P.grid(True)
        for iinterval in range(len(polyfit_degrees)):
            P.plot(dftarray_list[iinterval],gwarray_list[iinterval]-dftarray_list[iinterval]-N.polyval(polyfitlist[iinterval],dftarray_list[iinterval]),'bx')
        [x_min,x_max] = P.xlim()
        P.plot([x_min,x_max],[0,0],'k-')
        P.xlabel('DFT eigenvalues (in eV)')
        P.ylabel('Error in the fit (in eV)')
        P.show()
    return polyfitlist

def compare_bandstructures(ec_ref,ec_test):
    nspkpt_ref = len(ec_ref.special_kpoints)
    nspkpt_test = len(ec_test.special_kpoints)
    if nspkpt_ref != nspkpt_test:
        print 'ERROR: The number of special kpoints is different in the two files ... exit'
        sys.exit()
    eig_type_ref = ec_ref.eigenvalue_type
    eig_type_test = ec_test.eigenvalue_type
    print eig_type_ref,eig_type_test
    if eig_type_ref == 'DFT' and eig_type_test == 'W90':
        TOL_KPTS = N.float(1.0e-4)
    else:
        TOL_KPTS = N.float(1.0e-6)
    print TOL_KPTS
    for ispkpt in range(nspkpt_ref):
        print 'difference between the two :',ec_ref.special_kpoints[ispkpt,:]-ec_test.special_kpoints[ispkpt,:]
        if not N.allclose(ec_ref.special_kpoints[ispkpt,:],ec_test.special_kpoints[ispkpt,:],atol=TOL_KPTS):
            print 'ERROR: The kpoints are not the same :'
            print '       Kpt #%s ' %ispkpt
            print '         Reference => %20.17f %20.17f %20.17f' %(ec_ref.special_kpoints[ispkpt,0],ec_ref.special_kpoints[ispkpt,1],ec_ref.special_kpoints[ispkpt,2])
            print '         Compared  => %20.17f %20.17f %20.17f' %(ec_test.special_kpoints[ispkpt,0],ec_test.special_kpoints[ispkpt,1],ec_test.special_kpoints[ispkpt,2])
            print '       ... exit'
            sys.exit()
    mband_comparison = N.min([ec_ref.mband,ec_test.mband])
    if mband_comparison < ec_ref.mband:
        print 'Number of bands in the test bandstructure is lower than the number of bands in the reference (%s)' %ec_ref.mband
        print '  => Comparison will proceed with %s bands' %ec_test.mband
    elif mband_comparison < ec_test.mband:
        print 'Number of bands in the reference bandstructure is lower than the number of bands in the test bandstructure (%s)' %ec_test.mband
        print '  => Comparison will only proceed with %s bands of the test bandstructure' %ec_ref.mband
    else:
        print 'Number of bands in the reference and test bandstructure is the same'
        print '  => Comparison will proceed with %s bands' %mband_comparison
#    eig_test_ref_path = ec_ref.eigenvalues[:,:,:mband_comparison]
    rmsd_per_band = N.zeros([ec_ref.nsppol,mband_comparison],N.float)
    nrmsd_per_band = N.zeros([ec_ref.nsppol,mband_comparison],N.float)
    mae_per_band = N.zeros([ec_ref.nsppol,mband_comparison],N.float)
    for isppol in range(ec_ref.nsppol):
        for iband in range(mband_comparison):
            interp = N.interp(ec_ref.normalized_kpoint_path_values,ec_test.normalized_kpoint_path_values,ec_test.eigenvalues[isppol,:,iband])
            rmsd_per_band[isppol,iband] = N.sqrt(N.sum((csts.hartree2ev*interp-csts.hartree2ev*ec_ref.eigenvalues[isppol,:,iband])**2)/ec_ref.nkpt)
            mae_per_band[isppol,iband] = N.sum(N.abs(csts.hartree2ev*interp-csts.hartree2ev*ec_ref.eigenvalues[isppol,:,iband]))/ec_ref.nkpt
    P.figure(1)
    P.plot(mae_per_band[0,:])
    P.figure(2)
    P.plot(rmsd_per_band[0,:])
    P.show()

def get_gvectors():
    if os.path.isfile('.gvectors.bsinfo'):
        print 'File ".gvectors.bsinfo found with the following gvectors information :"'
        try:
            gvectors_reader = open('.gvectors.bsinfo','r')
            gvectors_data = gvectors_reader.readlines()
            gvectors_reader.close()
            trial_gvectors = N.identity(3,N.float)
            trial_gvectors[0,0] = N.float(gvectors_data[0].split()[0])
            trial_gvectors[0,1] = N.float(gvectors_data[0].split()[1])
            trial_gvectors[0,2] = N.float(gvectors_data[0].split()[2])
            trial_gvectors[1,0] = N.float(gvectors_data[1].split()[0])
            trial_gvectors[1,1] = N.float(gvectors_data[1].split()[1])
            trial_gvectors[1,2] = N.float(gvectors_data[1].split()[2])
            trial_gvectors[2,0] = N.float(gvectors_data[2].split()[0])
            trial_gvectors[2,1] = N.float(gvectors_data[2].split()[1])
            trial_gvectors[2,2] = N.float(gvectors_data[2].split()[2])
            print ' gvectors(1) = [ %20.17f %20.17f %20.17f ]' %(trial_gvectors[0,0],trial_gvectors[0,1],trial_gvectors[0,2])
            print ' gvectors(2) = [ %20.17f %20.17f %20.17f ]' %(trial_gvectors[1,0],trial_gvectors[1,1],trial_gvectors[1,2])
            print ' gvectors(3) = [ %20.17f %20.17f %20.17f ]' %(trial_gvectors[2,0],trial_gvectors[2,1],trial_gvectors[2,2])
        except:
            print 'ERROR: file ".gvectors.bsinfo" might be corrupted (empty or not formatted correctly ...)'
            print '       you should remove the file and start again or check the file ... exit'
            sys.exit()
        test = raw_input('Press <ENTER> to use these gvectors (any other character to enter manually other gvectors)\n')
        if test == '':
            gvectors = trial_gvectors
        else:
            gvectors = N.identity(3,N.float)
            test = raw_input('Enter G1 (example : "0.153 0 0") : \n')
            gvectors[0,0] = N.float(test.split()[0])
            gvectors[0,1] = N.float(test.split()[1])
            gvectors[0,2] = N.float(test.split()[2])
            test = raw_input('Enter G2 (example : "0.042 1.023 0") : \n')
            gvectors[1,0] = N.float(test.split()[0])
            gvectors[1,1] = N.float(test.split()[1])
            gvectors[1,2] = N.float(test.split()[2])
            test = raw_input('Enter G3 (example : "0 0 1.432") : \n')
            gvectors[2,0] = N.float(test.split()[0])
            gvectors[2,1] = N.float(test.split()[1])
            gvectors[2,2] = N.float(test.split()[2])
            test = raw_input('Do you want to overwrite the gvectors contained in the file ".gvectors.bsinfo" ? (<ENTER> for yes, anything else for no)\n')
            if test == '':
                print 'Writing gvectors to file ".gvectors.bsinfo" ...'
                gvectors_writer = open('.gvectors.bsinfo','w')
                gvectors_writer.write('%20.17f %20.17f %20.17f\n' %(trial_gvectors[0,0],trial_gvectors[0,1],trial_gvectors[0,2]))
                gvectors_writer.write('%20.17f %20.17f %20.17f\n' %(trial_gvectors[1,0],trial_gvectors[1,1],trial_gvectors[1,2]))
                gvectors_writer.write('%20.17f %20.17f %20.17f\n' %(trial_gvectors[2,0],trial_gvectors[2,1],trial_gvectors[2,2]))
                gvectors_writer.close()
                print '... done'
    else:
        test = raw_input('Do you want to enter the the reciprocal space primitive vectors (y/n)\n')
        if test == 'y':
            gvectors = N.identity(3,N.float)
            test = raw_input('Enter G1 (example : "0.153 0 0") : ')
            gvectors[0,0] = N.float(test.split()[0])
            gvectors[0,1] = N.float(test.split()[1])
            gvectors[0,2] = N.float(test.split()[2])
            test = raw_input('Enter G2 (example : "0.042 1.023 0") : ')
            gvectors[1,0] = N.float(test.split()[0])
            gvectors[1,1] = N.float(test.split()[1])
            gvectors[1,2] = N.float(test.split()[2])
            test = raw_input('Enter G3 (example : "0 0 1.432") : ')
            gvectors[2,0] = N.float(test.split()[0])
            gvectors[2,1] = N.float(test.split()[1])
            gvectors[2,2] = N.float(test.split()[2])
            test = raw_input('Do you want to write the gvectors to file ".gvectors.bsinfo" ? (<ENTER> for yes, anything else for no)\n')
            if test == '':
                print 'Writing gvectors to file ".gvectors.bsinfo" ...'
                gvectors_writer = open('.gvectors.bsinfo','w')
                gvectors_writer.write('%20.17f %20.17f %20.17f\n' %(gvectors[0,0],gvectors[0,1],gvectors[0,2]))
                gvectors_writer.write('%20.17f %20.17f %20.17f\n' %(gvectors[1,0],gvectors[1,1],gvectors[1,2]))
                gvectors_writer.write('%20.17f %20.17f %20.17f\n' %(gvectors[2,0],gvectors[2,1],gvectors[2,2]))
                gvectors_writer.close()
                print '... done'
        else:
            gvectors = None
    return gvectors

# Parse the command line options

parser = argparse.ArgumentParser(description='Tool for plotting dft bandstructures')
parser.add_argument('files',help='files to be opened',nargs=1)
args = parser.parse_args()
args_dict = vars(args)

if args_dict['files']:
    print 'will open the file'
else:
    print 'ERROR: you should provide some bandstructure file ! exiting now ...'
    sys.exit()
dft_file = args_dict['files'][0]

gvectors = get_gvectors()

ec_dft = EigenvalueContainer(directory='.',filename=dft_file)
ec_dft.set_kpoint_sampling_type('Bandstructure')
ec_dft.find_special_kpoints(gvectors)

print 'Number of bands in the file : %s' %(N.shape(ec_dft.eigenvalues)[2])
test = raw_input('Enter the number of bands to be plotted (<ENTER> : %s) : \n' %(N.shape(ec_dft.eigenvalues)[2]))
if test == '':
    nbd_plot = N.shape(ec_dft.eigenvalues)[2]
else:
    nbd_plot = N.int(test)
if nbd_plot > N.shape(ec_dft.eigenvalues)[2]:
    print 'ERROR: the number of bands to be plotted is larger than the number available ... exit'
    sys.exit()

ec_dft.special_kpoints_names = ['']*len(ec_dft.special_kpoints_indices)
for ii in range(len(ec_dft.special_kpoints_indices)):
    ec_dft.special_kpoints_names[ii] = 'k%s' %(ii+1)
print 'List of special kpoints :'
for ii in range(len(ec_dft.special_kpoints_indices)):
    spkpt = ec_dft.kpoints[ec_dft.special_kpoints_indices[ii]]
    print ' Kpoint %s : %s %s %s' %(ii+1,spkpt[0],spkpt[1],spkpt[2])
print 'Enter the name of the %s special k-points :' %(len(ec_dft.special_kpoints_indices))
test = raw_input('')
if len(test.split()) == len(ec_dft.special_kpoints_indices):
    for ii in range(len(ec_dft.special_kpoints_indices)):
        ec_dft.special_kpoints_names[ii] = test.split()[ii]

test = raw_input('Enter base name for bandstructure file : \n')
ec_dft.write_bandstructure_to_file('%s.bandstructure' %test)

P.figure(1,figsize=(3.464,5))
P.hold('on')
P.grid('on')
P.xticks(N.take(ec_dft.kpoint_reduced_path_values,N.array(ec_dft.special_kpoints_indices,N.int)),ec_dft.special_kpoints_names)
if ec_dft.nsppol == 1:
    for iband in range(nbd_plot):
        P.plot(ec_dft.kpoint_reduced_path_values,ec_dft.eigenvalues[0,:,iband]*csts.hartree2ev,'k-',linewidth=2)
elif ec_dft.nsppol == 2:
    for iband in range(nbd_plot):
        P.plot(ec_dft.kpoint_reduced_path_values,ec_dft.eigenvalues[0,:,iband]*csts.hartree2ev,'k-',linewidth=2)
        P.plot(ec_dft.kpoint_reduced_path_values,ec_dft.eigenvalues[1,:,iband]*csts.hartree2ev,'r-',linewidth=2)
P.show()


import warnings
import numpy as np
import time
from struct import unpack
from scipy.stats import norm
from scipy.signal import tukey
from Filter import median_filter
import Analysis, Filter, Constants

def savefits(data, filename, vmax=1.0, sps=1e6, bits=14, noise=False, clobber=True):
    """
    Save pulse/noise to FITS file
    """

    import pyfits as pf

    # Prepare data
    data = (np.asarray(data)/vmax*2**(bits-1)).round()

    # Column Name
    if noise:
    	colname = 'NoiseRec'
    else:
    	colname = 'PulseRec'

    # Columns
    col_t = pf.Column(name='TIME', format='1D', unit='s', array=np.zeros(data.shape[0], dtype=int))
    col_data = pf.Column(name=colname, format='%dI' % data.shape[1], unit='V', array=data)

    cols = pf.ColDefs([col_t, col_data])
    tbhdu = pf.BinTableHDU.from_columns(cols)

    # Name of extension
    exthdr = tbhdu.header
    exthdr['EXTNAME'] = ('Record', 'name of this binary table extension')
    exthdr['EXTVER'] = (1, 'extension version number')

    # Add more attributes
    exthdr['TSCAL2'] = (vmax/2**(bits-1), '[V/ch]')
    exthdr['TZERO2'] = (0., '[V]')
    exthdr['THSCL2'] = (sps**-1, '[s/bin] horizontal resolution of record')
    exthdr['THZER2'] = (0, '[s] horizontal offset of record')
    exthdr['THSAM2'] = (data.shape[1], 'sample number of record')
    exthdr['THUNI2'] = ('s', 'physical unit of sampling step of record')
    exthdr['TRMIN2'] = (-2**(bits-1)+1, '[channel] minimum number of each sample')
    exthdr['TRMAX2'] = (2**(bits-1)-1, '[channel] maximum number of each sample')
    exthdr['TRBIN2'] = (1, '[channel] default bin number of each sample')

    # More attributes
    exthdr['TSTART'] = (0, 'start time of experiment in total second')
    exthdr['TSTOP'] = (0, 'end time of experiment in total second')
    exthdr['TEND'] = (0, 'end time of experiment (obsolete)')
    exthdr['DATE'] = (time.strftime("%Y-%m-%dT%H:%M:%S", time.gmtime()), 'file creation date (UT)')

    # We anyway need Primary HDU
    hdu = pf.PrimaryHDU()

    # Write to FITS
    thdulist = pf.HDUList([hdu, tbhdu])
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        thdulist.writeto(filename, clobber=clobber)

def fopen(filename):
    """
    Read FITS file

    Parameters
    ==========
        filename:   file number to read

    Returns
    =======
        t:          time array
        wave:       waveform array
    """

    import pyfits as pf

    # Open fits file and get pulse/noise data
    header = pf.open(filename)
    wave = header[1].data.field(1).copy()
    dt = header[1].header['THSCL2']
    t = np.arange(wave.shape[-1]) * dt
    header.close()

    return t, wave

def yopen(filenumber, summary=False, nf=None, tmin=None, tmax=None, raw=False):
    """
    Read Yokogawa WVF file

    Parameters
    ==========
        filenumber: file number to read
        summary:    to summary waves (default: False)
        nf:         sigmas for valid data using median noise filter, None to disable noise filter (default: None)
        tmin:       lower boundary of time for partial extraction, scaler or list (Default: None)
        tmax:       upper boundary of time for partial extraction, scaler or list (Default: None)
        raw:        returns raw data without scaling/offsetting if True (Default: False)

    Returns
    =======
        if summary is False:
            [ t1, d1, t2, d2, t3, d3, ... ]

        if summary is True:
            [ t1, d1, err1, t2, d2, err2, ... ]

        if raw is True:
            t1 is a tuple of (hres1, hofs1, vres1, vofs1)

        where t1 is timing for 1st ch, d1 is data for 1st ch, err1 is error (1sigma) for 1st ch, and so on.
    """

    # Read header (HDR)
    h = open(str(filenumber) + ".HDR")
    lines = h.readlines()
    h.close()

    # Parse $PublicInfo
    for line in lines:
        token = line.split()

        if len(token) > 0:
            # Check endian
            if token[0] == "Endian":
                endian = '>' if token[1] == "Big" else '<'

            # Check data format
            if token[0] == "DataFormat":
                format = token[1]
                assert format == "Block"

            # Check # of groups
            if token[0] == "GroupNumber":
                groups = int(token[1])

            # Check # of total traces
            if token[0] == "TraceTotalNumber":
                ttraces = int(token[1])

            # Check data offset
            if token[0] == "DataOffset":
                offset = int(token[1])

    # Initialize containers
    traces = [None] * groups        # Number of traces for each group
    blocks = [None] * ttraces       # Number of blocks for each trace
    bsizes = [None] * ttraces       # Block size for each trace
    vres = [None] * ttraces         # VResolution for each trace
    voffset = [None] * ttraces      # VOffset for each trace
    hres = [None] * ttraces         # HResolution for each trace
    hoffset = [None] * ttraces      # HOffset for each trace

    # Parse $Group
    for line in lines:
        token = line.split()

        if len(token) > 0:
            # Read current group number
            if token[0][:6] == "$Group":
                cgn = int(token[0][6:]) - 1  # Current group number (minus 1)

            # Check # of traces in this group
            if token[0] == "TraceNumber":
                traces[cgn] = int(token[1])
                traceofs = np.sum(traces[:cgn], dtype=int)

            # Check # of Blocks
            if token[0] == "BlockNumber":
                blocks[traceofs:traceofs+traces[cgn]] = [ int(token[1]) ] * traces[cgn]

            # Check Block Size
            if token[0] == "BlockSize":
                bsizes[traceofs:traceofs+traces[cgn]] = [ int(s) for s in token[1:] ]

            # Check VResolusion
            if token[0] == "VResolution":
                vres[traceofs:traceofs+traces[cgn]] = [ float(res) for res in token[1:] ]

            # Check VOffset
            if token[0] == "VOffset":
                voffset[traceofs:traceofs+traces[cgn]] = [ float(ofs) for ofs in token[1:] ]

            # Check VDataType
            if token[0] == "VDataType":
                assert token[1] == "IS2"

            # Check HResolution
            if token[0] == "HResolution":
                hres[traceofs:traceofs+traces[cgn]] = [ float(res) for res in token[1:] ]

            # Check HOffset
            if token[0] == "HOffset":
                hoffset[traceofs:traceofs+traces[cgn]] = [ float(ofs) for ofs in token[1:] ]

    # Data Initialization
    time = [ np.array(range(bsizes[t])) * hres[t] + hoffset[t] for t in range(ttraces) ]
    data = [ [None] * blocks[t] for t in range(ttraces) ]

    # Open WVF
    f = open(str(filenumber) + ".WVF", 'rb')
    f.seek(offset)

    # Read WVF
    if format == "Block":
        # Block format (assuming block size is the same for all the traces in Block format)
        for b in range(blocks[0]):
            for t in range(ttraces):
                if raw:
                    data[t][b] = np.array(unpack(endian + 'h'*bsizes[t], f.read(bsizes[t]*2)), dtype='int64')
                else:
                    data[t][b] = np.array(unpack(endian + 'h'*bsizes[t], f.read(bsizes[t]*2))) * vres[t] + voffset[t]
    else:
        # Trace format
        for t in range(ttraces):
            for b in range(blocks[t]):
                if raw:
                    data[t][b] = np.array(unpack(endian + 'h'*bsizes[t], f.read(bsizes[t]*2)), dtype='int64')
                else:
                    data[t][b] = np.array(unpack(endian + 'h'*bsizes[t], f.read(bsizes[t]*2))) * vres[t] + voffset[t]

    # Array conversion
    for t in range(ttraces):
        if raw:
            data[t] = np.array(data[t], dtype='int64')
        else:
            data[t] = np.array(data[t])

    # Tmin/Tmax filtering
    for t in range(ttraces):
        if type(tmin) == list or type(tmax) == list:
            if not (type(tmin) == list and type(tmax) == list and len(tmin) == len(tmax)):
                raise ValueError("tmin and tmax both have to be list and have to have the same length.")
            mask = np.add.reduce([ (time[t] >= _tmin) & (time[t] < _tmax) for (_tmax, _tmin) in zip(tmax, tmin)], dtype=bool)
        else:
            _tmin = np.min(time[t]) if tmin is None else tmin
            _tmax = np.max(time[t]) + 1 if tmax is None else tmax
            mask = (time[t] >= _tmin) & (time[t] < _tmax)

        data[t] = data[t][:, mask]
        time[t] = time[t][mask]

    f.close()

    if summary is False:
        # Return wave data as is
        if raw:
            return [ [ (hres[t], hoffset[t], vres[t], voffset[t]), data[t] ] for t in range(ttraces)  ]
        else:
            return [ [ time[t], data[t] ] for t in range(ttraces)  ]
    else:
        if nf is None:
            # Noise filter is off
            if raw:
                return [ [ (hres[t], hoffset[t], vres[t], voffset[t]), np.mean(data[t].astype(dtype='float64'), axis=0), np.std(data[t].astype(dtype='float64'), axis=0, ddof=1) ]
                            for t in range(ttraces) ]
            else:
                return [ [ time[t], np.mean(data[t], axis=0), np.std(data[t], axis=0, ddof=1) ]
                            for t in range(ttraces) ]
        else:
            # Noise filter is on
            if raw:
                return [ [ (hres[t], hoffset[t], vres[t], voffset[t]),
                            np.apply_along_axis(lambda a: np.mean(a[median_filter(a, nf)]), 0,  data[t].astype(dtype='float64')),
                            np.apply_along_axis(lambda a: np.std(a[median_filter(a, nf)], ddof=1), 0, data[t].astype(dtype='float64')) ]
                                for t in range(ttraces) ]
            else:
                return [ [ time[t],
                            np.apply_along_axis(lambda a: np.mean(a[median_filter(a, nf)]), 0,  data[t]),
                            np.apply_along_axis(lambda a: np.std(a[median_filter(a, nf)], ddof=1), 0, data[t]) ]
                                for t in range(ttraces) ]

def popen(filename, ch=None, raw=False):
    """
    Read pls file

    Parameters
    ==========
        filename:   file name to read
        ch:         returns data only for the given channel if given (Default: None)
        raw:        returns raw data without scaling/offsetting if True (Default: False)

    Returns
    =======
        if raw is True:
            [ header, vres, vofs, hres, hofs, tick, num, data, edata ]
        else:
            [ header, t, tick, num, data, edata ]
    """

    # Initialize
    header = {'COMMENT': []}
    vres = {}
    vofs = {}
    hres = {}
    hofs = {}
    tick = {}
    num  = {}
    data = {}
    edata = {}

    # Parser
    def parser():
        """
        PLS Data Parser (generator)
        """

        # Initialization
        samples = -1
        extra = 0
        chunk = ''
        isHeader = True

        while True:
            while len(chunk) < 2:
                chunk += yield

            # Get the magic character
            magic = chunk[0]

            if isHeader and magic == 'C':
                # Comment
                while len(chunk) < 80:
                    chunk += yield
                header['COMMENT'].append(chunk[2:80])
                chunk = chunk[80:]

            elif isHeader and magic == 'V':
                # Version
                while len(chunk) < 80:
                    chunk += yield
                header['VERSION'] = chunk[2:80]
                chunk = chunk[80:]

            elif isHeader and magic == 'O':
                # Date
                while len(chunk) < 10:
                    chunk += yield
                _m, _d, _y = map(int, chunk[2:10].split())
                header['DATE'] = "%d/%d/%d" % (_y, _m, _d)
                chunk = chunk[10:]

            elif isHeader and magic == 'S':
                # Number of Samples
                while len(chunk) < 7:
                    chunk += yield
                header['SAMPLES'] = samples = int(chunk[2:7])
                chunk = chunk[7:]

            elif isHeader and magic == 'E':
                # Extra Bytes
                while len(chunk) < 7:
                    chunk += yield
                header['EXTRA'] = extra = int(chunk[2:7])
                chunk = chunk[7:]

            elif isHeader and magic == 'P':
                # Discriminator
                while len(chunk) < 78:
                    chunk += yield
                _dis = chunk[2:78].split()
                if _dis[0] == '01':
                    header['ULD'] = eval(_dis[1])
                elif _dis[0] == '02':
                    header['LLD'] = eval(_dis[1])
                chunk = chunk[78:]

            elif isHeader and magic == 'N':
                # Normalization
                while len(chunk) < 47:
                    chunk += yield
                _ch, _hofs, _hres, _vofs, _vres = chunk[2:47].split()
                _ch = int(_ch)
                vres[_ch] = eval(_vres)
                vofs[_ch] = eval(_vofs)
                hres[_ch] = eval(_hres)
                hofs[_ch] = eval(_hofs)
                chunk = chunk[47:]

            elif magic == 'D':
                # Data
                isHeader = False

                if samples < 0:
                    raise ValueError("Invalid number of samples.")
                while len(chunk) < (11 + samples*2):
                    chunk += yield
                _ch, _tick, _num = unpack('<BII', chunk[2:11])

                if not data.has_key(_ch):
                    data[_ch] = bytearray()
                    tick[_ch] = []
                    num[_ch] = []
                    edata[_ch] = bytearray()

                data[_ch] += chunk[11:11 + samples*2]
                tick[_ch].append(_tick)
                num[_ch].append(_num)
                edata[_ch] += chunk[11 + samples*2:11 + samples*2 + extra]

                chunk = chunk[11 + samples*2 + extra:]

            else:
                # Skip unknown magic
                chunk = chunk[1:]

    # Open pls file and read by chunks
    f = open(filename, 'rb')

    # Start parser
    p = parser()
    p.next()

    # Read by chunk and parse it
    with open(filename, 'rb') as f:
        while True:
            chunk = f.read(1024*1024)   # read 1 MB
            if not chunk:
                break
            p.send(chunk)

    # Convert buffer to numpy array
    for k in ([ch] if ch else data.keys()):
        data[k] = np.frombuffer(data[k], dtype='>i2').reshape(-1, header['SAMPLES'])
        edata[k] = np.frombuffer(edata[k], dtype='>u1').reshape(-1, header['SAMPLES'])

    if raw:
        if ch:
            return header, vres[ch], vofs[ch], hres[ch], hofs[ch], tick[ch], num[ch], data[ch], edata[ch]
        else:
            return header, vres, vofs, hres, hofs, tick, num, data, edata
    else:
        t = {}

        for k in ([ch] if ch else data.keys()):
            # Normalize data using res/ofs
            t[k] = (np.arange(header['SAMPLES']) + hofs[k]) * hres[k]
            data[k] = (np.asarray(data[k]) + vofs[k]) * vres[k]

        if ch:
            return header, t[ch], tick[ch], num[ch], data[ch], edata[ch]
        else:
            return header, t, tick, num, data, edata

def tesana(t, p, n, lpfc=None, hpfc=None, binsize=1, max_shift=10,
            thre=0.4, filt=None, nulldc=False, offset=False, center=False, sigma=3,
            gain=None, dsr=None, shift=False, ocmethod="ols", flip=False, atom="Mn",
            kbfit=False, ignorekb=False, method="mle",
            rshunt=None, tbias=None, ites=None, ka_min=80, kb_min=40,
            tex=False, plotting=True, savedat=False, session="Unnamed"):
    """
    Perform TES Analysis

    Parameters (and their default values):
        t:          time data (array-like)
        p:          pulse data (array-like)
        n:          noise data (array-like)
        lpfc:       low-pass filter cut-off frequency in bins (Default: None)
        hpfc:       high-pass filter cut-off frequency in bins (Default: None)
        binsize:    energy bin size for histograms and fittings (only for ls ans cs) in eV (Default: 1)
        max_shift:  maximum allowed shifts to calculate maximum cross correlation (Default: 10)
        thre:       correlation threshold for offset correction (Default: 0.4)
        filt:       window function (hanning/hamming/blackman/tukey) (Default: None)
        nulldc:     nullify the DC bin when template generation (Default: False)
        offset:     subtract DC offset (Default: False)
        center:     centering pulse rise (Default: False)
        sigma:      sigmas for median filter (Default: 3)
        gain:       feedback gain for current-space conversion (Default: None)
        dsr:        down-sampling rate (Default: None)
        shift:      treat dE as energy shift instead of scaling (Default: False)
        ocmethod:   offset correction fitting method (ols/odr) (Default: ols)
        flip:       flip x and y when offset correction fitting (Default: False)
        atom:       atom to fit (Default: Mn)
        kbfit:      fit Kb line (Default: False)
        ignorekb:   ignore Kb line when linearity correction (Default: False)
        method:     fitting method (mle/ls/cs) (Default: mle)
        rshunt:     shunt resistance value for r-space conversion (Default: None)
        tbias:      TES bias current for r-space conversion (Default: None)
        ites:       TES current for r-space conversion (Default: None)
        ka_min:     minimum counts to group bins for Ka line (valid only for ls/cs fittings) (Default: 80)
        ka_min:     minimum counts to group bins for Kb line (valid only for ls/cs fittings) (Default: 20)
        tex:        use TeX for plots (Default: False)
        plotting:   generate and save plots (Default: True)
        savedat:    save data to files (Default: False)
        session:    session name for plots and data files (Default: Unnamed)

    Note:
        - Use offset option when using filt option
        - Consider using center option when using filt option
    """

    if plotting:
        # Import matplotlib
        import matplotlib
        matplotlib.use('Agg')
        matplotlib.rcParams['text.usetex'] = str(tex)
        from pylab import figure, plot, errorbar, hist, axvline, xlim, ylim, loglog, xlabel, ylabel, legend, tight_layout, savefig

        print "Session: %s" % session

    # Preparation
    p = np.asarray(p)
    n = np.asarray(n)
    t = np.asarray(t)
    dt = np.diff(t)[0]
    df = (dt * t.shape[-1])**-1

    # Subtract offset
    if offset:
        ofs = np.median(n)
        p -= ofs
        n -= ofs

    # Convert to current-space if needed
    if gain:
        print "Converting to current-space"
        p /= gain
        n /= gain

    # Convert to resistance-space
    Rspace = False
    if gain and rshunt and tbias and ites:
        print "Converting to resistance-space"
        ofs = np.median(n)
        p += (ites - ofs)
        n += (ites - ofs)

        # Convert to resistance
        p = (tbias - p) * rshunt / p
        n = (tbias - n) * rshunt / n

        Rspace = True

    # Down-sample
    if dsr > 1:
        p = p[:,:p.shape[-1]/dsr*dsr].reshape(p.shape[0], -1, dsr).mean(axis=-1)
        n = n[:,:n.shape[-1]/dsr*dsr].reshape(n.shape[0], -1, dsr).mean(axis=-1)
        dt *= dsr
        t = t[::dsr]

    # Pulse centering (for filtering)
    if center:
        # Roll pulse to the center
        r = p.shape[-1] / 2 - np.median(abs(p - Filter.offset(p)[:, np.newaxis]).argmax(axis=-1))
        p = np.hstack((p[...,-r:], p[...,:-r]))

    # Calculate offset (needs to be done before applying filter)
    if p.size > 0:
        offset = Filter.offset(p)

    # Generate Filter
    if filt is None:
        pass
    else:
        if filt.lower() == "hanning":
            f = np.hanning(p.shape[-1])
        elif filt.lower() == "hamming":
            f = np.hamming(p.shape[-1])
        elif filt.lower() == "blackman":
            f = np.blackman(p.shape[-1])
        elif filt.lower() == "tukey":
            f = tukey(p.shape[-1])
        else:
            raise ValueError('Unsupported filter: %s' % filt.lower())

        print "Window filter function: %s" % filt.lower()

        # Amplitude correction
        cf = f.sum() / len(f)
        p *= (f / cf)
        n *= (f / cf)

        # Equivalent noise bandwidth correction
        enb = len(f)*(f**2).sum()/f.sum()**2
        df *= enb

    if p.size > 0:
        # Calculate averaged pulse
        avgp = Filter.average_pulse(p, max_shift=max_shift)

        if savedat:
            np.savetxt('%s-averagepulse.dat' % session, np.vstack((t, avgp)).T,
                header='Time (s), Averaged Pulse (%s)' % ('R' if Rspace else ('A' if gain else 'V')), delimiter='\t')

        if plotting:
            figure()
            plot(t, avgp)
            xlabel('Time$\quad$(s)')
            ylabel('Averaged Pulse$\quad$(%s)' % ('R' if Rspace else ('A' if gain else 'V')))
            tight_layout()
            savefig('%s-averagepulse.pdf' % session)

        # Calculate averaged pulse spectrum
        avgps = np.sqrt(Filter.power(avgp)) / df

        if savedat:
            np.savetxt('%s-avgpulse-power.dat' % session, np.vstack((np.arange(len(avgps))*df, avgps)).T,
                header='Frequency (Hz), Average Pulse Power (%s/srHz)' % ('R' if Rspace else ('A' if gain else 'V')), delimiter='\t')

        if plotting:
            avgps[0] = 0    # for better plot
            figure()
            plot(np.arange(len(avgps))*df, avgps)
            loglog()
            xlabel('Frequency$\quad$(Hz)')
            ylabel('Average Pulse Power$\quad$(%s/Hz)' % ('R' if Rspace else ('A' if gain else 'V')))
            tight_layout()
            savefig('%s-avgpulse-power.pdf' % session)

    if n.size > 0:
        # Plot noise spectrum
        avgns = np.sqrt(Filter.average_noise(n) / df)

        if savedat:
            np.savetxt('%s-noise.dat' % session, np.vstack((np.arange(len(avgns))*df, avgns)).T,
                header='Frequency (Hz), Noise (%s/srHz)' % ('R' if Rspace else ('A' if gain else 'V')), delimiter='\t')

        if plotting:
            avgns[0] = 0    # for better plot
            figure()
            plot(np.arange(len(avgns))*df, avgns)
            loglog()
            xlabel('Frequency$\quad$(Hz)')
            ylabel('Noise$\quad$(%s/$\sqrt{\mathrm{Hz}}$)' % ('R' if Rspace else ('A' if gain else 'V')))
            tight_layout()
            savefig('%s-noise.pdf' % session)

    if p.size > 0 and n.size > 0:
        # Generate template
        tmpl, sn = Filter.generate_template(p, n, lpfc=lpfc, hpfc=hpfc, nulldc=nulldc, max_shift=max_shift)

        if savedat:
            np.savetxt('%s-template.dat' % session, np.vstack((t, tmpl)).T,
                header='Time (s), Template (A.U.)', delimiter='\t')
            np.savetxt('%s-sn.dat' % session, np.vstack((np.arange(len(sn))*df, sn/np.sqrt(df))).T,
                header='Frequency (Hz), S/N (/srHz)', delimiter='\t')

        if plotting:
            # Plot template
            figure()
            plot(t, tmpl)
            xlabel('Time$\quad$(s)')
            ylabel('Template$\quad$(A.U.)')
            tight_layout()
            savefig('%s-template.pdf' % session)

            # Plot SNR
            figure()
            plot(np.arange(len(sn))*df, sn/np.sqrt(df))
            loglog()
            xlabel('Frequency$\quad$(Hz)')
            ylabel('S/N$\quad$(/$\sqrt{\mathrm{Hz}}$)')
            tight_layout()
            savefig('%s-sn.pdf' % session)

        # Calculate baseline resolution
        print "Resolving power: %.2f (%.2f eV @ 5.9 keV)" % (np.sqrt((sn**2).sum()*2), Analysis.baseline(sn))

        # Perform optimal filtering
        pha_p = Filter.optimal_filter(p, tmpl, max_shift=max_shift)
        pha_n = Filter.optimal_filter(n, tmpl, max_shift=0)

        # Offset correction
        (a, b), coef = Analysis.fit_offset(pha_p, offset, sigma=sigma, method=ocmethod, flip=flip)

        if coef > thre:
            oc_pha_p = Analysis.offset_correction(pha_p, offset, b)
            oc_pha_n = Analysis.offset_correction(pha_n, offset, b)
            print "Offset correction with: PHA = %f * (1 + %f * Offset)" % (a, b)

            if plotting:
                figure()
                ka = Analysis.ka(np.vstack((pha_p, offset)).T, sigma=sigma)
                plot(ka.T[1], ka.T[0], '.', c='k')
                x_min, x_max = xlim()
                ofs = np.linspace(x_min, x_max)
                label = '$\mathrm{PHA}=%.2f\\times(1+%.2f\\times\mathrm{Offset})$' % (a, b)
                plot(ofs, a*(1+b*ofs), 'r-', label=label)
                xlabel('Offset$\quad$(V)')
                ylabel('PHA$\quad$(V)')
                legend(frameon=False)
                tight_layout()
                savefig('%s-offset.pdf' % session)
        else:
            oc_pha_p = pha_p
            oc_pha_n = pha_n
            print "Skipped offset correction: correlation coefficient (%f) is too small" % coef

        # Check line database
        if "%sKa" % atom not in Constants.LE.keys() or "%sKb" % atom not in Constants.LE.keys():
            raise ValueError('Unsupported atom: %s' % atom)

        # Linearity correction
        pha_line_center = np.asarray([ np.median(Analysis.ka(oc_pha_p, sigma=sigma)), np.median(Analysis.kb(oc_pha_p, sigma=sigma)) ])
        line_energy = np.asarray([ Constants.LE['%sKa' % atom], Constants.LE['%sKb' % atom] ])

        if ignorekb:
            a, b = Analysis.fit_linearity([pha_line_center[0]], [line_energy[0]], deg=1)
            print "Linearity correction with: PHA = %e * E" % (b)
        else:
            a, b = Analysis.fit_linearity(pha_line_center, line_energy, deg=2)
            print "Linearity correction with: PHA = %e * E^2 + %e * E" % (a, b)
            print "MnKb saturation ratio: %.2f %%" % ((pha_line_center[1]/pha_line_center[0])/(line_energy[1]/line_energy[0])*100)

        lc_pha_p = Analysis.linearity_correction(oc_pha_p, a, b)
        lc_pha_n = Analysis.linearity_correction(oc_pha_n, a, b)

        if savedat:
            np.savetxt('%s-linearity.dat' % session, array([pha_line_center[0]]) if ignorekb else pha_line_center[np.newaxis,:],
                header='%sKa PHA' % atom if ignorekb else '%sKa PHA, %sKb PHA' % (atom, atom), delimiter='\t')

        if plotting:
            figure()
            x = np.linspace(0, 7e3)
            if ignorekb:
                plot(line_energy[0]/1e3, pha_line_center[0], '+', color='b')
                plot(x/1e3, x*b, 'r--')
            else:
                plot(line_energy/1e3, pha_line_center, '+', color='b')
                plot(x/1e3, x**2*a+x*b, 'r--')
            xlim((0, 7))
            xlabel('Energy$\quad$(keV)')
            ylabel('PHA$\quad$(a.u.)')
            tight_layout()
            savefig('%s-linearity.pdf' % session)

        # Energy Spectrum
        if plotting:
            figure()
            hcount, hbin, hpatch = hist(lc_pha_p[lc_pha_p==lc_pha_p]/1e3, bins=7000/binsize, histtype='stepfilled', color='y')
            xlim(0, 7)
            xlabel('Energy$\quad$(keV)')
            ylabel('Count')
            tight_layout()
            savefig('%s-spec.pdf' % session)

        if savedat:
            hcount, hbin = np.histogram(lc_pha_p[lc_pha_p==lc_pha_p]/1e3, bins=7000/binsize)
            np.savetxt('%s-spec.dat' % session, np.vstack(((hbin[1:]+hbin[:-1])/2, hcount)).T,
                header='Energy (keV), Count', delimiter='\t')

        # Line fitting
        def _line_fit(data, min, line):

            # Fit
            (dE, width), (dE_error, width_error), e = Analysis.fit(data, binsize=binsize, min=min, line=line, shift=shift, method=method)

            if method == "cs":
                chi_squared, dof = e

            if method in ("mle", "ls"):
                print "%s: %.2f +/- %.2f eV @ Ec%+.2f eV" \
                    % (line, width, width_error, dE)
            elif method == "cs":
                print "%s: %.2f +/- %.2f eV @ Ec%+.2f eV (Red. chi^2 = %.1f/%d = %.2f)" \
                    % (line, width, width_error, dE, chi_squared, dof, chi_squared/dof)

            return dE, width, width_error

        def _line_spectrum(data, min, line, dE, width, width_error):

            # Draw histogram
            n, bins = Analysis.histogram(data, binsize=binsize)

            if method in ("cs"):
                gn, gbins = Analysis.group_bin(n, bins, min=min)
            else:
                # No grouping in mle and ls
                gn, gbins = n, bins

            ngn = gn/(np.diff(gbins))
            ngn_sigma = np.sqrt(gn)/(np.diff(gbins))
            cbins = (gbins[1:]+gbins[:-1])/2

            if plotting:
                figure()

                if width_error is not None:
                    label = 'FWHM$=%.2f\pm %.2f$ eV' % (width, width_error)
                else:
                    label = 'FWHM$=%.2f$ eV (Fixed)' % width

                if method == "cs":
                    errorbar(cbins, ngn, yerr=ngn_sigma, xerr=np.diff(gbins)/2, capsize=0, ecolor='k', fmt=None, label=label)
                else:
                    hist(data, bins=gbins, weights=np.ones(len(data))/binsize, histtype='step', ec='k', label=label)

                E = np.linspace(bins.min(), bins.max(), 1000)

                model = Analysis.normalization(ngn, gbins, dE, width, line=line, shift=shift) \
                            * Analysis.line_model(E, dE, width, line=line, shift=shift, full=True)

                # Plot theoretical model
                plot(E, model[0], 'r-')

                # Plot fine structures
                for m in model[1:]:
                    plot(E, m, 'b--')

                xlabel('Energy$\quad$(eV)')
                ylabel('Normalized Count$\quad$(count/eV)')
                legend(frameon=False)

                ymin, ymax = ylim()
                ylim(ymin, ymax*1.1)
                tight_layout()

                savefig("%s-%s.pdf" % (session, line))

            if savedat:
                np.savetxt('%s-%s.dat' % (session, line), np.vstack((cbins, gn)).T,
                    header='Energy (keV), Count', delimiter='\t')

        ## Ka
        ka = Analysis.ka(lc_pha_p, sigma=sigma)
        dE, width, width_error = _line_fit(ka, ka_min, "%sKa" % atom)
        _line_spectrum(ka, ka_min, "%sKa" % atom, dE, width, width_error)

        ## Kb
        kb = Analysis.kb(lc_pha_p, sigma=sigma)
        if kbfit:
            dE, width, width_error = _line_fit(kb, kb_min, "%sKb" % atom)
        else:
            width_error = None
        _line_spectrum(kb, kb_min, "%sKb" % atom, dE, width, width_error)

        ## Baseline
        f_pha_n = lc_pha_n[Filter.median_filter(lc_pha_n, sigma=sigma)]
        baseline = Analysis.sigma2fwhm(np.std(f_pha_n))
        print "Baseline resolution: %.2f eV" % baseline

        n, bins = Analysis.histogram(f_pha_n, binsize=binsize)

        if savedat:
            np.savetxt('%s-baseline.dat' % session, np.vstack(((bins[1:]+bins[:-1])/2, n)).T,
                header='Energy (keV), Count', delimiter='\t')

        if plotting:
            figure()
            label = 'FWHM$=%.2f$ eV' % baseline
            hist(f_pha_n, bins=bins, weights=np.ones(len(f_pha_n))/binsize, histtype='step', ec='k', label=label)
            mu, sigma = norm.fit(f_pha_n)
            E = np.linspace(bins.min(), bins.max(), 1000)
            plot(E, norm.pdf(E, loc=mu, scale=sigma)*len(f_pha_n), 'r-')

            xlabel('Energy$\quad$(eV)')
            ylabel('Normalized Count$\quad$(count/eV)')

            legend(frameon=False)

            tight_layout()

            savefig('%s-baseline.pdf' % session)


"""
==========================
SGD: convex loss functions
==========================

A plot that compares the various convex loss functions supported by
:class:`sklearn.linear_model.SGDClassifier` .
"""
print(__doc__)

import numpy as np
import matplotlib.pyplot as plt

def modified_huber_loss(y_true, y_pred):
    z = y_pred * y_true
    loss = -4 * z
    loss[z >= -1] = (1 - z[z >= -1]) ** 2
    loss[z >= 1.] = 0
    return loss

xmin, xmax = -4, 4
xx = np.linspace(xmin, xmax, 100)
plt.plot([xmin, 0, 0, xmax], [1, 1, 0, 0], 'k-',
         label="Zero-one loss")
plt.plot(xx, np.where(xx < 1, 1 - xx, 0), 'g-',
         label="Hinge loss")
plt.plot(xx, -np.minimum(xx, 0), 'm-',
         label="Perceptron loss")
plt.plot(xx, np.log2(1 + np.exp(-xx)), 'r-',
         label="Log loss")
plt.plot(xx, np.where(xx < 1, 1 - xx, 0) ** 2, 'b-',
         label="Squared hinge loss")
plt.plot(xx, modified_huber_loss(xx, 1), 'y--',
         label="Modified Huber loss")
plt.ylim((0, 8))
plt.legend(loc="upper right")
plt.xlabel(r"Decision function $f(x)$")
plt.ylabel("$L(y, f(x))$")
plt.show()


# -*- coding: utf-8 -*-
"""
Created on Mon Aug 17 18:06:45 2015

@author: thomas.douenne

TICPE: Taxe intérieure sur la consommation des produits énergétiques
"""

# L'objectif de ce script est d'illustrer graphiquement l'évolution du taux implicite de la TICPE depuis 1993.
# On étudie ce taux pour le diesel, et pour les carburants sans plombs.

# Import de modules généraux
from pandas import concat

# Import de modules spécifiques à Openfisca
from openfisca_france_indirect_taxation.examples.utils_example import graph_builder_bar_list
from openfisca_france_indirect_taxation.examples.dataframes_from_legislation.get_accises import get_accises_carburants
from openfisca_france_indirect_taxation.examples.dataframes_from_legislation.get_tva import get_tva_taux_plein
from openfisca_france_indirect_taxation.examples.dataframes_from_legislation.get_prix_carburants import \
    get_prix_carburants

# Appel des paramètres de la législation et des prix
ticpe = ['ticpe_gazole', 'ticpe_super9598']
accise_diesel = get_accises_carburants(ticpe)
prix_ttc = ['diesel_ttc', 'super_95_ttc']
prix_carburants = get_prix_carburants(prix_ttc)
tva_taux_plein = get_tva_taux_plein()

# Création d'une dataframe contenant ces paramètres
df_taux_implicite = concat([accise_diesel, prix_carburants, tva_taux_plein], axis = 1)
df_taux_implicite.rename(columns = {'value': 'taux plein tva'}, inplace = True)

# A partir des paramètres, calcul des taux de taxation implicites
df_taux_implicite['taux_implicite_diesel'] = (
    df_taux_implicite['accise ticpe gazole'] * (1 + df_taux_implicite['taux plein tva']) /
    (df_taux_implicite['prix diesel ttc'] -
    (df_taux_implicite['accise ticpe gazole'] * (1 + df_taux_implicite['taux plein tva'])))
    )

df_taux_implicite['taux_implicite_sp95'] = (
    df_taux_implicite['accise ticpe super9598'] * (1 + df_taux_implicite['taux plein tva']) /
    (df_taux_implicite['prix super 95 ttc'] -
    (df_taux_implicite['accise ticpe super9598'] * (1 + df_taux_implicite['taux plein tva'])))
    )

df_taux_implicite = df_taux_implicite.dropna()

# Réalisation des graphiques
graph_builder_bar_list(df_taux_implicite['taux_implicite_diesel'], 1, 1)
graph_builder_bar_list(df_taux_implicite['taux_implicite_sp95'], 1, 1)


# -*- coding: utf-8 -*-
"""
Combines all the images for a given (TARGET, FILTER, HWP) combination to
produce a single, average image.

Estimates the sky background level of the on-target position at the time of the
on-target observation using a bracketing pair of off-target observations through
the same HWP polaroid rotation value. Subtracts this background level from
each on-target image to produce background free images. Applies an airmass
correction to each image, and combines these final image to produce a background
free, airmass corrected, average image.
"""

# Core imports
import os
import sys
import copy
import warnings

# Import scipy/numpy packages
import numpy as np
from scipy import ndimage

# Import astropy packages
from astropy.table import Table
import astropy.units as u
from astropy.convolution import Gaussian2DKernel
from astropy.modeling import models, fitting
from astropy.stats import gaussian_fwhm_to_sigma, sigma_clipped_stats
from photutils import (make_source_mask,
    MedianBackground, SigmaClip, Background2D)

# Import plotting utilities
from matplotlib import pyplot as plt

# Add the AstroImage class
import astroimage as ai

# Add the header handler to the BaseImage class
from Mimir_header_handler import Mimir_header_handler
ai.reduced.ReducedScience.set_header_handler(Mimir_header_handler)
ai.set_instrument('mimir')

#==============================================================================
# *********************** CUSTOM USER CODE ************************************
# this is where the user specifies where the raw data is stored
# and some of the subdirectory structure to find the actual .FITS images
#==============================================================================

# This is a list of targets for which to process each subgroup (observational
# group... never spanning multiple nights, etc...) instead of combining into a
# single "metagroup" for all observations of that target. The default behavior
# is to go ahead and combine everything into a single, large "metagroup". The
# calibration data should probably not be processed as a metagroup though.
processSubGroupList = []
processSubGroupList = [t.upper() for t in processSubGroupList]

# Define the location of the PPOL reduced data to be read and worked on
PPOL_data = 'C:\\Users\\Jordan\\FITS_data\\Mimir_data\\PPOL_Reduced\\201611\\'
S3_dir    = os.path.join(PPOL_data, 'S3_Astrometry')

# This is the location where all pyPol data will be saved
pyPol_data = 'C:\\Users\\Jordan\\FITS_data\\Mimir_data\\pyPol_Reduced\\201611'

# This is the location of the previously generated masks (step 4)
maskDir = os.path.join(pyPol_data, 'Masks')

# Setup new directory for polarimetry data
polarimetryDir = os.path.join(pyPol_data, 'Polarimetry')
if (not os.path.isdir(polarimetryDir)):
    os.mkdir(polarimetryDir, 0o755)

HWPDir = os.path.join(polarimetryDir, 'HWPImgs')
if (not os.path.isdir(HWPDir)):
    os.mkdir(HWPDir, 0o755)

bkgPlotDir = os.path.join(HWPDir, 'bkgPlots')
if (not os.path.isdir(bkgPlotDir)):
    os.mkdir(bkgPlotDir, 0o755)

# # Setup PRISM detector properties
# read_noise = 13.0 # electrons
# effective_gain = 3.3 # electrons/ADU

#########
### Establish the atmospheric extinction (magnitudes/airmass)
#########
# Following table from Hu (2011)
# Data from Gaomeigu Observational Station
# Passband | K'(lambda) [mag/airmass] | K'' [mag/(color*airmass)]
# U	         0.560 +/- 0.023            0.061 +/- 0.004
# B          0.336 +/- 0.021            0.012 +/- 0.003
# V          0.198 +/- 0.024           -0.015 +/- 0.004
# R          0.142 +/- 0.021           -0.067 +/- 0.005
# I          0.093 +/- 0.020            0.023 +/- 0.006

# Following table from Schmude (1994)
# Data from Texas A & M University Observatory
# Passband | K(lambda) [mag/airmass] | dispersion on K(lambda)
# U	         0.60 +/- 0.05             0.120
# B          0.40 +/- 0.06             0.165
# V          0.26 +/- 0.03             0.084
# R          0.19 +/- 0.03             0.068
# I          0.16 +/- 0.02             0.055

# TODO: Ask Dan about atmospheric extinction from airmass at NIR
kappa = dict(zip(['U',  'B',  'V',  'R',  'I',  'J',  'H',  'K'  ],
                 [0.60, 0.40, 0.26, 0.19, 0.16, 0.05, 0.01, 0.005]))

# Read in the indexFile data and select the filenames
print('\nReading file index from disk')
indexFile = os.path.join(pyPol_data, 'reducedFileIndex.csv')
fileIndex = Table.read(indexFile, format='ascii.csv')

# Determine which parts of the fileIndex pertain to HEX dither science images
useFiles    = np.logical_and(
    fileIndex['USE'] == 1,
    fileIndex['DITHER_TYPE'] == 'ABBA'
)
useFileRows = np.where(useFiles)

# Cull the file index to only include files selected for use
fileIndex = fileIndex[useFileRows]

# Define an approximate pixel scale
pixScale  = 0.5789*(u.arcsec/u.pixel)

# TODO: implement a FWHM seeing cut... not yet working because PSF getter seems
# to be malfunctioning in step 2
#
#
# # Loop through each unique GROUP_ID and test for bad seeing conditions.
# groupByID = fileIndex.group_by(['GROUP_ID'])
# for subGroup in groupByID.groups:
#     # Grab the FWHM values for this subGroup
#     thisFWHMs = subGroup['FWHM']*u.pixel
#
#     # Grab the median and standard deviation of the seeing for this subgroup
#     medianSeeing = np.median(thisFWHMs)
#     stdSeeing    = np.std(thisFWHMs)
#
#     # Find bad FWHM values
#     badFWHMs = np.logical_not(np.isfinite(subGroup['FWHM']))
#     badFWHMs = np.logical_or(
#         badFWHMs,
#         thisFWHMs <= 0
#     )
#     badFWHM = np.logical_and(
#         badFWHM,
#         thisFWHMs > 2.0*u.arcsec
#     )
#     import pdb; pdb.set_trace()

# Group the fileIndex by...
# 1. Target
# 2. Waveband
fileIndexByTarget = fileIndex.group_by(['TARGET', 'FILTER'])

# Loop through each group
for group in fileIndexByTarget.groups:
    # Grab the current group information
    thisTarget     = str(np.unique(group['TARGET'].data)[0])
    thisFilter     = str(np.unique(group['FILTER'].data)[0])

    # # Skip the Merope nebula for now... not of primary scientific importance
    # if thisTarget == 'MEROPE': continue

    # Update the user on processing status
    print('\nProcessing images for')
    print('Target : {0}'.format(thisTarget))
    print('Filter : {0}'.format(thisFilter))

    # Grab the atmospheric extinction coefficient for this wavelength
    thisKappa = kappa[thisFilter]

    # Further divide this group by its constituent HWP values
    indexByPolAng = group.group_by(['IPPA'])

    # Loop over each of the HWP values, as these are independent from
    # eachother and should be treated entirely separately from eachother.
    for IPPAgroup in indexByPolAng.groups:
        # Grab the current HWP information
        thisIPPA = np.unique(IPPAgroup['IPPA'].data)[0]

        # Update the user on processing status
        print('\tIPPA : {0}'.format(thisIPPA))

        # For ABBA dithers, we need to compute the background levels on a
        # sub-group basis. If this target has not been selected for subGroup
        # averaging, then simply append the background subtracted images to a
        # cumulative list of images to align and average.

        # Initalize an image list to store all the images for this
        # (target, filter, pol-ang) combination
        imgList = []

        indexByGroupID = IPPAgroup.group_by(['GROUP_ID'])
        for subGroup in indexByGroupID.groups:
            # Grab the numae of this subGroup
            thisSubGroup = str(np.unique(subGroup['OBJECT'])[0])

            # if (thisSubGroup != 'NGC2023_R1') and (thisSubGroup != 'NGC2023_R2'): continue

            # Construct the output file name and test if it alread exsits.
            if thisTarget in processSubGroupList:
                outFile = '_'.join([thisTarget, thisSubGroup, str(thisIPPA)])
                outFile = os.path.join(HWPDir, outFile) + '.fits'
            elif thisTarget not in processSubGroupList:
                outFile = '_'.join([thisTarget, thisFilter, str(thisIPPA)])
                outFile = os.path.join(HWPDir, outFile) + '.fits'

            # Test if this file has already been constructed and either skip
            # this subgroup or break out of the subgroup loop.
            if os.path.isfile(outFile):
                print('\t\tFile {0} already exists...'.format(os.path.basename(outFile)))
                if thisTarget in processSubGroupList:
                    continue
                elif thisTarget not in processSubGroupList:
                    break

            # Update the user on the current execution status
            print('\t\tProcessing images for subgroup {0}'.format(thisSubGroup))

            # Initalize lists to store the A and B images.
            AimgList  = []
            BimgList  = []

            # Initalize a list to store the off-target sky background levels
            BbkgList  = []

            # Initilaze lists to store the times of observation
            AdatetimeList = []
            BdatetimeList = []

            # Read in all the images for this subgroup
            progressString = '\t\tNumber of Images : {0}'
            for iFile, filename in enumerate(subGroup['FILENAME']):
                # Update the user on processing status
                print(progressString.format(iFile+1), end='\r')

                # Read in a temporary compy of this image
                PPOL_file = os.path.join(S3_dir, filename)
                tmpImg = ai.reduced.ReducedScience.read(PPOL_file)

                # Crop the edges of this image
                ny, nx = tmpImg.shape
                binningArray = np.array(tmpImg.binning)

                # Compute the amount to crop to get a 1000 x 1000 image
                cy, cx = (ny - 1000, nx - 1000)

                # Compute the crop boundaries and apply them
                lf = np.int(np.round(0.5*cx))
                rt = lf + 1000
                bt = np.int(np.round(0.5*cy))
                tp = bt + 1000

                tmpImg = tmpImg[bt:tp, lf:rt]

                # Grab the on-off target value for this image
                thisAB = subGroup['AB'][iFile]

                # Place the image in a list and store required background values
                if thisAB == 'B':
                    # Place B images in the BimgList
                    BimgList.append(tmpImg)

                    # Place the median value of this off-target image in list
                    mask = make_source_mask(
                        tmpImg.data, snr=2, npixels=5, dilate_size=11
                        )
                    mean, median, std = sigma_clipped_stats(
                        tmpImg.data, sigma=3.0, mask=mask
                    )
                    BbkgList.append(median)

                    # Place the time of this image in a list of time values
                    BdatetimeList.append(tmpImg.julianDate)

                if thisAB == 'A':
                    # Read in any associated masks and store them.
                    maskFile = os.path.join(maskDir, os.path.basename(filename))

                    # If there is a mask for this file, then apply it!
                    if os.path.isfile(maskFile):
                        # Read in the mask file
                        tmpMask = ai.reduced.ReducedScience.read(maskFile)

                        # Crop the mask to match the shape of the original image
                        tmpMask = tmpMask[cy:ny-cy, cx:nx-cx]

                        # Grab the data to be masked
                        tmpData = tmpImg.data

                        # Mask the data and put it back into the tmpImg
                        maskInds = np.where(tmpMask.data)
                        tmpData[maskInds] = np.NaN
                        tmpImg.data = tmpData

                    # Place B images in the BimgList
                    AimgList.append(tmpImg)

                    # Place the time of this image in a list of time values
                    AdatetimeList.append(tmpImg.julianDate)

            # Create a new line for shell output
            print('')

            # Construct an image stack of the off-target images
            BimageStack = ai.utilitywrappers.ImageStack(BimgList)

            # Build a supersky image from these off-target images
            superskyImage = BimageStack.produce_supersky()

            import pdb; pdb.set_trace()

            # Locate regions outside of a 5% deviation
            tmpSuperskyData  = superskyImage.data
            maskedPix  = np.abs(tmpSuperskyData - 1.0)  > 0.05

            # Get rid of the small stuff and expand the big stuff
            maskedPix = ndimage.binary_opening(maskedPix, iterations=2)
            maskedPix = ndimage.binary_closing(maskedPix, iterations=2)
            maskedPix = ndimage.binary_dilation(maskedPix, iterations=4)

            # TODO: Make the box_size and filter_size sensitive to binning.
            binningArray = np.array(superskyImage.binning)
            box_size = tuple((100/binningArray).astype(int))
            filter_size = tuple((10/binningArray).astype(int))

            # Setup the sigma clipping and median background estimators
            sigma_clip = SigmaClip(sigma=3., iters=10)
            bkg_estimator = MedianBackground()

            # Compute a smoothed background image
            bkgData = Background2D(superskyImage.data,
                box_size=box_size, filter_size=filter_size, mask=maskedPix,
                sigma_clip=sigma_clip, bkg_estimator=bkg_estimator)

            # Construct a smoothed supersky image object
            smoothedSuperskyImage = ai.reduced.ReducedScience(
                bkgData.background/bkgData.background_median,
                uncertainty = bkgData.background_rms,
                properties={'unit':u.dimensionless_unscaled}
            )

            # Interpolate background values to A times
            AbkgList = np.interp(
                AdatetimeList,
                BdatetimeList,
                BbkgList,
                left=-1e6,
                right=-1e6
            )

            # Cut out any extrapolated data (and corresponding images)
            goodInds = np.where(AbkgList > -1e5)
            AimgList = np.array(AimgList)[goodInds]
            AdatetimeList = np.array(AdatetimeList)[goodInds]
            AbkgList = AbkgList[goodInds]

            AsubtractedList = []
            # Loop through the on-target images and subtract background values
            for Aimg, Abkg in zip(AimgList, AbkgList):
                # Subtract the interpolated background values from the A images
                tmpImg = Aimg - smoothedSuperskyImage*(Abkg*Aimg.unit)

                # Apply an airmass correction
                tmpImg = tmpImg.correct_airmass(thisKappa)

                # Append the subtracted and masked image to the list.
                AsubtractedList.append(tmpImg)

            # Now that the images have been fully processed, pause to generate
            # a plot to store in the "background plots" folder. These plots
            # constitute a good sanity check on background subtraction.
            plt.plot(BdatetimeList, BbkgList, '-ob')
            plt.scatter(AdatetimeList, AbkgList, marker='o', facecolor='r')
            plt.xlabel('Julian Date')
            plt.ylabel('Background Value [ADU]')
            figName = '_'.join([thisTarget, thisSubGroup, str(thisIPPA)])
            figName = os.path.join(bkgPlotDir, figName) + '.png'
            plt.savefig(figName, dpi=300)
            plt.close('all')

            # Here is where I need to decide if each subgroup image should be
            # computed or if I should just continue with the loop.
            if thisTarget.upper() in processSubGroupList:
                # Construct an image combiner for the A images
                AimgStack = ai.utilitywrappers.ImageStack(AsubtractedList)

                # Align the images
                AimgStack.align_images_with_wcs(
                    subPixel=False,
                    padding=np.NaN
                    )

                # Combine the images
                AoutImg = imgStack.combine_images()

                # Save the image
                AoutImg.write(outFile, dtype=np.float64)
            else:
                # Extend the imgList variable with background corrected images
                imgList.extend(AsubtractedList)

        if len(imgList) > 0:
            # At the exit of the loop, process ALL the files from ALL the groups
            # Construct an image combiner for the A images
            imgStack = ai.utilitywrappers.ImageStack(imgList)

            # Align the images
            imgStack.align_images_with_wcs(
                subPixel=False,
                padding=np.NaN
                )

            # Combine the images
            outImg = imgStack.combine_images()

            # Save the image
            outImg.write(outFile, dtype=np.float64)

print('\nDone computing average images!')