#!/usr/bin/python3  

import numpy as np
import os 

class KPOINT : 

    def __init__(self,i,x,w,n) :
        self.w  = w                  # weight
        self.i  = i                  # index in full calculation
        self.x  = np.array(x)        # Cartesian coordinates (2pi/alat)
        self.xx = np.zeros([1,3])    # Cartesian coordinates (Bohr)
        self.n  = n                  # number of bands
        self.e  = np.zeros([1,n])    # energy array
        self.j  = 0                  # index in the path
        self.o  = np.zeros([1,n])    # occupation array
        return None
    
    def PrintK(self) : 
        print('%3i  (%10.7f %10.7f  %10.7f) : w=%10.7f' % (self.i,self.x[0],self.x[1],self.x[2],self.w) ) 
        return None

    def ConvertUnitsOfCoordinates(self,alat) : 
        prefactor = 2.0*np.pi/alat
        self.xx   = [ c * prefactor for c in self.x ]
        return None 

    def GetEnergies(self,lines,starting_line_index) : 
        datablock_length_energies = int( np.ceil( self.n/8 ) )
        datablock_length_occup    = int( np.ceil( self.n/8 ) )
        datablock_length          = datablock_length_energies + datablock_length_occup + 4
        energy_lines = lines[starting_line_index + 2 : starting_line_index + 2 + datablock_length_energies ]
        energies = []
        for line in energy_lines   :  energies = energies + list(map( float , line.split() )) 
        if len(energies) != self.n : 
            print(' !!!!\n ERROR : number or energies read differ from expected\n !!!!!')
            exit()
        self.e = np.array(energies)
        return None


def ReadFile(fname):
    fin = open(fname,'r')
    content = fin.readlines()
    fin.close()
    return content

def GetInputParameters(lines) :
    first_line = GetLineIndex(lines,'     bravais-lattice index     =',1) 
    alat_bohr  = float( lines[first_line+1].split()[4] )
    nbands     = int(   lines[first_line+6].split()[4] )
    first_line = GetLineIndex(lines,'     number of k points=',1)
    nkpts      = int(   lines[first_line].split()[4]   )
    inputs = {'alat':alat_bohr , 'nbnd' : nbands , 'nkpt' : nkpts }
    return inputs

def GetAllKpoints(lines,inputs) : 
    all_kpts = []
    nkpt = inputs['nkpt']
    nbnd = inputs['nbnd']
    first_line = GetLineIndex(lines,'cart. coord. in units 2pi/alat',1)
    for line in lines[ first_line + 1  :  first_line + nkpt + 1 ]:
        kindex     = int( line.split(')')[0].split()[1] )
        cart_coord = list( map( float , line.split(')')[1].split()[2:] ) )
        weight     = float( line.split()[-1] )
        all_kpts.append( KPOINT( kindex , cart_coord , weight , nbnd ) )
    return all_kpts


def GetLineIndex(lines , string_to_look_for, m) :
    indexes = [ i for i in range(len(lines)) if string_to_look_for in lines[i] ]
    if m == 1 : 
        if len(indexes) > 1 : print(' !!!!! \n WARNING: found multiple entries of %s\n, taking the first one !!!!!' % string_to_look_for )
        return indexes[0]
    else :
        return indexes

def ReduceLines(lines) :
    lines_cycles = GetLineIndex(lines,'End of self-consistent calculation',0)
    last_cycle   = lines_cycles[-1]
    lines = lines[ last_cycle : ]
    return lines


def AssociateEnergies(lines, allk, nk ) : 
    index_kpt_lines = GetLineIndex(lines,'          k =',0)
    if len(index_kpt_lines) != nk : 
        print(' !!!!! Error in counting the k-points') 
        exit()
    for ik in range(nk) : 
        k = allk[ik]
        k.GetEnergies(lines, index_kpt_lines[ik])
    return None

def SelectKpointsInPath(allk) : 
    kpt_path = [ k for k in allk if k.w < 0.00000011 ] 
    nkpt_path = len(kpt_path)
    for i in range(nkpt_path) : kpt_path[i].j = i 
    return kpt_path, nkpt_path


def AreSamePoint(x1,x2) : 
    x1=np.array(x1)
    x2=np.array(x2)
    dist = x1-x2
    are_the_same = all( abs(dist) < 0.0001 ) 
    return are_the_same

def ConstructBandstructureMatrix( nb, k , nk ) : 
    bs_matrix = np.zeros([ nk, nb ])
    for ik in range(nk) : 
        bs_matrix[ ik , : ] = k[ik].e
    return bs_matrix


def Get_File_Name( default_name = 'outfile' ) :
    fnout = default_name
    while os.path.isfile(default_name) :
        do_overwrite=input( " # FILE %s already exists. Overwrite [1 / 0] ?  " % default_name )
        if len(do_overwrite) == 0 :
           continue
        elif do_overwrite[0] in ('1','Y','y','T','t') :
            fnout = default_name
            break
        elif do_overwrite[0] in ('0','N','n','F','f') :
            fnout=''.join( input(' # insert bands file name  :  ').split() )
            break
        else :
           continue
    return fnout


def PrintBandstructureFile( nk, nb, bs, fn) : 
    fout = open(fn,'w') 
    print('# kpt     energy (eV)     energy (Ha)' , file=fout)
    for ib in range(nb) : 
        print('' , file = fout) 
        print('# band %5i' % (ib+1) , file = fout) 
        band_dispersion = bs[:,ib] 
        for ik in range(nk) : 
            print(' %5i  %10.4f  %10.4f' % ( ik, band_dispersion[ik], band_dispersion[ik]/27.211369917461)  , file = fout ) 
    fout.close()
    return None

if __name__ == "__main__" : 
    print('')
    filename=input(' Insert name of the file to parse:    ')
    lines  = ReadFile(filename)

    inputs  = GetInputParameters(lines)
    allkpts = GetAllKpoints(lines,inputs) 
    print( ' N. kpoints read    = %4i' % inputs['nkpt'] )

    lines = ReduceLines(lines)
    AssociateEnergies( lines, allkpts , inputs['nkpt'] )

    kpts_path, nkpt_path = SelectKpointsInPath(allkpts) 

    band_matrix = ConstructBandstructureMatrix( inputs['nbnd'] , kpts_path , nkpt_path )
    print( ' Size of BS matrix  = %4i %4i ' % band_matrix.shape )

    fnout = Get_File_Name( default_name = 'bands.dat' )
    PrintBandstructureFile( nkpt_path, inputs['nbnd'], band_matrix, fnout ) 
    print(' Bandstrcutre recorded successfully in file "%s"\n' % fnout )