#!/usr/bin/python

def bstr(point1, point2, npoints, nbands, old_nc_file, outfile):
 '''  
      This module performs a band structure calculation 
      using the results from a previous calculation 
      old_nc_file = nc file of a well-converged dacapo calculation
      point1 = initial point for the band structure (e.g. X, K, U, W etc)
               given in CARTESIAN coordinates.   
      point2 = final point 
      npoints = how many k-points between point1 and point2. 
  
      point2 is NOT included in the results.

      nb= number of bands to calculate

      outfile is the dacapo output of the calculation (files outfile.txt and outfile.nc)

      the function returns an array of energies with zero set at the fermi level.
 '''    
       
 from Dacapo import Dacapo
 from numpy import dot, array, sqrt
 from numpy.linalg import inv

 # read old calculation
 bulk=Dacapo.ReadAtoms(old_nc_file)
 calc=bulk.GetCalculator()
 # read fermi energy
 E_fermi = calc.GetFermiLevel()

 # dacapo needs the k-points in units of the reciprocal lattice vectors. 
 # here we do the transformation
 K0=dot(inv(bulk.GetReciprocalUnitCell()), array(point1))
 K1=dot(inv(bulk.GetReciprocalUnitCell()), array(point2))
 # we create a list of kpoints starting from K0 and going just before K1.
 kpts=[]
 for x in range(npoints):
  kpts.append(K0+x/float(npoints)*(K1-K0))  

 # setup the Harris calculation
 # first, read parameters from old calc
 dens=calc.GetDensityArray()
 pwc=calc.GetPlaneWaveCutoff()
 pwc2=calc.GetDensityCutoff()
 del(calc)
 
 # now set up a new calculator
 calc2=Dacapo()
 # attach calc2 to old list of atoms
 bulk.SetCalculator(calc2)

 # in the following, we set calc2 parameters.
 calc2.SetPlaneWaveCutoff(pwc)
 calc2.SetDensityCutoff(pwc2)
 calc2.SetNumberOfBands(nbands)
 calc2.SetBZKPoints(kpts)
 calc2.SetTxtFile(outfile+".txt")
 calc2.SetNetCDFFile(outfile+".nc")
 calc2.StayAliveOff()
 # turn of symmetry, otherwise k-point set will be reduced.
 calc2.SetSymmetryOff()
 # this is the most important parameter: do not change the charge density
 calc2.SetChargeMixing(False)
 # read density from previous calculation
 calc2.SetDensityArray(dens)
 # run the calculation
 energy2 = calc2.GetPotentialEnergy()

 kpoints     = calc2.GetIBZKPoints()
 eigenvalues = [calc2.GetEigenvalues(kpt=kpt) for kpt in range(len(kpoints))]

 # subtract Fermi level from eigenvalues
 for ik in range(len(kpoints)):
   for n in range(len(eigenvalues[ik])):
       eigenvalues[ik][n]=eigenvalues[ik][n]-E_fermi
 
 # this is the distance between point1 and point2
 k=sqrt(dot(array(point1)-array(point2),array(point1)-array(point2)))

 return k, eigenvalues

if __name__ == '__main__':  
    from math import pi
    a=4.05
    G=(0.0, 0.0, 0.0)
    X=(2.0*pi/a, 0.0, 0.0)
    e=bstr(point1=G, point2=G, npoints=1, nbands=10, old_nc_file='al.nc', outfile='al_bstr_GX')
    for ik in range(len(e)):
      print ik, 
      for n in range(len(e[ik])):
        print e[ik][n],
      print