#!/usr/bin/env python

from scipy import *
from pylab import *
from scipy import weave
import glob, os

code="""
     using namespace std;

     double Ry2eV = 13.6056920311654;

     double Ax = 0;
     for (int ib=0; ib<nbands; ib++){
        complex<double> ekom( data(1+2*ib),  min(data(2+2*ib), -small) );
	complex<double> cohf( dach(1+2*ib),  dach(2+2*ib) );
        complex<double> gc = cohf/(omega+mu-ekom);
	// complex<double> gc = 1./(omega+mu-ekom);
	// cout<<ib<<" "<<gc<<endl;
        Ax += -gc.imag()/M_PI;
     }
     return_val = Ax;
     
     """
if __name__ == '__main__':

    small = 0.01 # 1e-5
    itensity = 0.3
    
    fEF = open('EF.dat', 'r')
    mu = float(fEF.next().split()[0])
    #mu -= 0.05
    print 'mu=', mu


    wg = glob.glob('*.klist_band')
    if len(wg)>0:
        fg = open(wg[0], 'r')
        wkpointi=[]
        wkpoints=[]
        for il,line in enumerate(fg):
            if line[:3]=='END': break
            com = line[:10].split()
            if com:
                wkpoints.append(com[0])
                wkpointi.append(il)
        print wkpointi
        print wkpoints

    
    fdat = open('eigvals.dat', 'r')
    if os.path.isfile('cohfactorsd.dat'):
        fcoh = open('cohfactorsd.dat', 'r')
    else:
        fcoh = None
    
    ikp=0
    Akom=[]
    ii=0
    try:
        while True:
            data = fdat.next().split()
            if fcoh is not None: dach = fcoh.next().split()
            
            (ikp, isym, nbands, nemin, nomega) = map(int, data[1:6])
            
            ekom = zeros(nbands, dtype=complex)
            dach = array([[1.,0.] for i in range(nbands)]).flatten()  # ones(2*nbands, dtype=float)
            Aom=[]
            om=[]
            for iom in range(nomega):
                data = array(map(float, fdat.next().split()))
                if fcoh is not None: dach = array(map(float, fcoh.next().split()))
                
                omega = float(data[0])
                
                Ax = weave.inline(code, ['nbands', 'omega', 'mu', 'data', 'small', 'ikp', 'dach'],
                                  type_converters=weave.converters.blitz, compiler = 'gcc')
                
                om.append(omega)
                Aom.append(Ax)
            Akom.append( Aom )
            ikp += 1

    except StopIteration:
        pass
    
    Akom = array(Akom).transpose()
    print shape(Akom)

    
    vmm = [0,max(map(max,Akom))*itensity]    
    (ymin,ymax) = (om[0],om[-1])
    (xmin,xmax) = (0, shape(Akom)[1]-1)
    
    print 'xmin,xmax,ymin,ymax=', xmin, xmax, ymin, ymax

    imshow(Akom, interpolation='bilinear', cmap=cm.hot, origin='lower', vmin=vmm[0], vmax=vmm[1], extent=[xmin,xmax,ymin,ymax], aspect=(xmax-xmin)*0.8/(ymax-ymin) )

    for i in range(len(wkpointi)):
        print 'wp=', wkpointi[i]
        plot([wkpointi[i],wkpointi[i]], [ymin,ymax], 'w-')
        
    plot([xmin,xmax],[0,0], 'w:')

    Ntck=5
    for j in range(len(wkpointi)-1):
        for ix in range(1,Ntck):
            x = wkpointi[j]+(wkpointi[j+1]-wkpointi[j])*ix/float(Ntck)
            plot([x,x],[-0.01,0.01],'w-')
        
    axis([xmin,xmax,ymin,ymax])
    xticks( wkpointi, wkpoints, fontsize='x-large' )
    #colorbar()
    show()