#!/usr/bin/env python
'''
mcu: Modeling and Crystallographic Utilities
Copyright (C) 2019 Hung Q. Pham. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

Email: Hung Q. Pham <pqh3.14@gmail.com>
'''

import numpy as np
import re, textwrap
from ..utils.misc import check_exist

   
############ pw.x output - begin ##################

# Get the lattice constant alat
alat_MATCH = re.compile(r'''
[\w\W]* lattice [ ]* parameter [ ]* \(alat\) [ ]* = [ ]* (?P<alat>\S+)
''', re.VERBOSE)

# Get no. of electron
nelec_MATCH = re.compile(r'''
[\w\W]* number [ ]* of [ ]* electrons [ ]* = [ ]* (?P<nelec>\S+)
''', re.VERBOSE)

# Get no. of electron
nbands_MATCH = re.compile(r'''
[\w\W]* number [ ]* of [ ]* Kohn-Sham [ ]* states [ ]* = [ ]* (?P<nbands>\S+)
''', re.VERBOSE)

# Get crystal axes in units of alat
axes_MATCH = re.compile(r'''
[ ]*
  a\( (?P<dim>\d+) \) [ ]* = [ ]* \( [ ]* (?P<a>\S+) [ ]* (?P<b>\S+) [ ]* (?P<c>\S+) [ ]* \)
''', re.VERBOSE)

# Get reciprocal axes in units 2 pi/alat
reciprocal_axes_MATCH = re.compile(r'''
[ ]*
  b\( (?P<dim>\d+) \) [ ]* = [ ]* \( [ ]* (?P<a>\S+) [ ]* (?P<b>\S+) [ ]* (?P<c>\S+) [ ]* \)
''', re.VERBOSE)

# Get the atom types 
species_MATCH = re.compile(r'''
[\w\W]*
  atomic [ ]* species [ ]* valence [ ]* mass [ ]* pseudopotential
(?P<content>
  [\s\S]*?(?=\n.*?[ ] $|\n\n)  # match everything until next blank line or EOL
)
''', re.VERBOSE)

element_MATCH = re.compile(r'''
[ ]*
  (?P<species>\S+) [ ]* (?P<valence>\S+) [ ]* (?P<mass>\S+) [ ]* (?P<element>\S+) [ ]* 
  \( [ ]* (?P<scale>\S+) [ ]* \)
''', re.VERBOSE)

# Get the atom lists
atoms_MATCH = re.compile(r'''
[\w\W]*
  site [ ]* n.[ ]* atom[ ]* positions [ ]* \( [ ]* alat [ ]* units [ ]* \)
(?P<content>
  [\s\S]*?(?=\n.*?[ ] $|\n\n)  # match everything until next blank line or EOL
)
''', re.VERBOSE)
         
atom_MATCH = re.compile(r'''
[ ]*
  (?P<no>\d+) [ ]* (?P<species>\S+) [ ]* tau\( [ ]* (?P<tau>\d+) [ ]* \) [ ]* = [ ]* 
  \( [ ]* (?P<x>\S+) [ ]* (?P<y>\S+) [ ]*  (?P<z>\S+) [ ]* \)
''', re.VERBOSE)

# Get k-points list
kpts_MATCH = re.compile(r'''
[\w\W]*
  cart. [ ]* coord. [ ]* in [ ]* units [ ]* 2pi/alat
(?P<content>
  [\s\S]*?(?=\n.*?[ ] $|\n\n)  # match everything until next blank line or EOL
)
''', re.VERBOSE)
        
kpt_MATCH = re.compile(r'''
[ ]*
  k\( [ ]* (?P<no>\d+) \) [ ]* = [ ]* \( [ ]* (?P<kx>\S+) [ ]* (?P<ky>\S+) [ ]*  
  (?P<kz>\S+) [ ]* \), [ ]* wk [ ]* =[ ]* (?P<wk>\S+)
''', re.VERBOSE)

# Get band energies as well as band edges or Fermi energy
bands_MATCH = re.compile(r'''
[ ]*
  k [ ]* = [ ]* (?P<kx>\S+) [ ]* (?P<ky>\S+) [ ]* (?P<kz>\S+) [ ]* \( [ ]* (?P<pws>\d+) [ ]* PWs\) [ ]* bands [ ]* \(ev\):
  \n\n
(?P<eigenvals>
  [\s\S]*?(?=\n.*?[ ] $|\n\n)  # match everything until next blank line or EOL
)
''', re.VERBOSE)

band_edges_MATCH = re.compile(r'''
[\w\W]* highest [ ]* occupied, [ ]* lowest [ ]* unoccupied [ ]* level [ ]* \(ev\): [ ]* (?P<HOCO>\S+) [ ]* (?P<LUCO>\S+)
''', re.VERBOSE)

fermi_MATCH = re.compile(r'''
[\w\W]* the [ ]* Fermi [ ]* energy [ ]* is [ ]* (?P<fermi_energy>\S+) [ ]* ev
''', re.VERBOSE)      
         
def read_pw_output(filename):  
    '''Read the output generated by pw.x'''

    assert check_exist(filename), 'Cannot find : ' + filename
    with open(filename, "r") as data_file:
        data = data_file.read()
        
        # Get alat, nlec
        alat = float(alat_MATCH.match(data)['alat'])
        nelec = float(nelec_MATCH.match(data)['nelec'])
        nbands = int(nbands_MATCH.match(data)['nbands'])
        soc = 'spin-orbit' in data
        
        
        # Get real lattice vectors
        lattice = []
        for axis in axes_MATCH.finditer(data):
            content = axis.groupdict()
            lattice.append([content['a'], content['b'], content['c']])
           
        lattice = np.float64(lattice)
        
        # Get real lattice vectors
        recip_lattice = []
        for axis in reciprocal_axes_MATCH.finditer(data):
            content = axis.groupdict()
            recip_lattice.append([content['a'], content['b'], content['c']])
           
        recip_lattice = np.float64(recip_lattice)
        
        # Get atom types
        species_block = species_MATCH.match(data).group('content')
        species = []
        for element in element_MATCH.finditer(species_block):
            element_dict = element.groupdict()
            valence = np.float64(element_dict['valence'])
            mass = np.float64(element_dict['mass'])
            pseudopotential = element_dict['element'] + ' ' + element_dict['scale']
            species.append([element_dict['species'], valence, mass, pseudopotential])
            
        # Get atom list
        atoms_block = atoms_MATCH.match(data).group('content')
        atom = []
        atom_position = []
        for atm in atom_MATCH.finditer(atoms_block):
            atom_dict = atm.groupdict()
            atom.append(atom_dict['species'])
            atom_position.append([atom_dict['x'], atom_dict['y'], atom_dict['z']])
          
        # Get kpts list
        kpts_block = kpts_MATCH.match(data).group('content')
        kpts = []
        for kpt in kpt_MATCH.finditer(kpts_block):
            kpt_dict = kpt.groupdict()
            kpts.append([kpt_dict['kx'], kpt_dict['ky'], kpt_dict['kz'], kpt_dict['wk']])   
            
        kpts = np.float64(kpts)
        
        # Get band energies
        nkpts = kpts.shape[0]
        band = []
        for point in bands_MATCH.finditer(data):
            point_dict = point.groupdict()
            energies = np.float64(textwrap.wrap(point_dict['eigenvals'][2:],9))
            band.append(energies)
            
        band = np.float64(band)
        band = band.reshape(-1, nkpts, band.shape[1])
        
        # Get band edges
        band_edges = band_edges_MATCH.match(data)
        if band_edges is not None:
            hoco, luco = np.float64([band_edges['HOCO'], band_edges['LUCO']])
        else:
            hoco, luco = None, None
            
        # Get Fermi energy
        fermi_energy = fermi_MATCH.match(data)
        if fermi_energy is not None:
            efermi = float(fermi_energy['fermi_energy'])
        else:
            efermi = None
            
        out = {}
        out['alat'] = alat
        out['nelec'] = nelec
        out['nbands'] = nbands
        out['soc'] = soc
        out['crystal axes'] = lattice
        out['reciprocal axes'] = recip_lattice
        out['species'] = species
        out['atom'] = atom
        out['atom_position'] = np.float64(atom_position)
        out['kpts'] = kpts
        out['eigenvals'] = band
        out['hoco'] = hoco 
        out['luco'] = luco  
        out['efermi'] = efermi

        return out
            
############ pw.x output - end ##################            
                 
############ dos.x output - begin ##################

dos_header_MATCH = re.compile(r'''
[\w\W]* \# [ ]* E [ ]* \(eV\) [ ]* dos\(E\) [ ]* Int [ ]* dos\(E\) [ ]* EFermi [ ]* = [ ]* (?P<efermi>\S+) [ ]* eV
''', re.VERBOSE)

dos_polarized_header_MATCH = re.compile(r'''
[\w\W]* \# [ ]* E [ ]* \(eV\) [ ]* dosup\(E\) [ ]* dosdw\(E\) [ ]* Int [ ]* dos\(E\) [ ]* EFermi [ ]* = [ ]* (?P<efermi>\S+) [ ]* eV
''', re.VERBOSE)
     
def read_tdos_output(filename):  
    '''Read DOS generated by dos.x'''

    assert check_exist(filename), 'Cannot find : ' + filename
    with open(filename, 'r') as data_file:
        data = data_file.read().split('\n')
        dos_header = dos_header_MATCH.match(data[0])
        if dos_header is not None:
            efermi = float(dos_header['efermi'])
            # DOS dimensions: [spin , [E(eV), tdos(E), idos(E)]]
            temp = np.float64(' '.join(data[1:]).split()).reshape(-1,3)
            dos = temp.reshape(-1, temp.shape[0], temp.shape[1])    
        else:
            dos_header = dos_polarized_header_MATCH.match(data[0])
            efermi = float(dos_header['efermi'])
            # DOS dimensions: [spin , [E(eV), tdos(E), idos(E)]]
            temp = np.float64(' '.join(data[1:]).split()).reshape(-1,4)
            dos = np.empty([2, temp.shape[0], 3])
            dos[0] = temp[:,[0,1,3]]
            dos[1] = temp[:,[0,2,3]]            

        out = {}
        out['dos'] = dos  
        out['efermi'] = efermi
        
        return out


############ projwfc.x output - begin ##################
tot_pdos            = "# E (eV)  dos(E)    pdos(E)"
tot_pdos_polarized  = "# E (eV)  dosup(E)   dosdw(E)  pdosup(E)  pdosdw(E)"

def lm_pdos(lm=1, polarized=False):
    if polarized:
        return ("# E (eV)  ldosup(E)  ldosdw(E)" + lm * " pdosup(E)  pdosdw(E) ").strip()
    else:
        return ("# E (eV)   ldos(E)   " + lm * "pdos(E)    ").strip()

def read_pdos_output(filename):  
    '''Read the dos files generated by projwfc.x
    
    ref: https://www.quantum-espresso.org/Doc/INPUT_PROJWFC.html
    Orbital Order
    Order of m-components for each l in the output:

        1, cos(phi), sin(phi), cos(2*phi), sin(2*phi), .., cos(l*phi), sin(l*phi)

    where phi is the polar angle:x=r cos(theta)cos(phi), y=r cos(theta)sin(phi)
    This is determined in file Modules/ylmr2.f90 that calculates spherical harmonics.

    for l=1:
      1 pz     (m=0)
      2 px     (real combination of m=+/-1 with cosine)
      3 py     (real combination of m=+/-1 with sine)

    for l=2:
      1 dz2    (m=0)
      2 dzx    (real combination of m=+/-1 with cosine)
      3 dzy    (real combination of m=+/-1 with sine)
      4 dx2-y2 (real combination of m=+/-2 with cosine)
      5 dxy    (real combination of m=+/-2 with sine)
    '''
    assert check_exist(filename), 'Cannot find : ' + filename
    with open(filename, 'r') as data_file:
        data = data_file.read().split('\n')
        dos_header = data[0].strip()
        
        if dos_header == tot_pdos:
            dos_unformatted = np.float64(' '.join(data[1:]).split()).reshape(-1,3)
            nEpts = dos_unformatted.shape[0]
            
            # DOS dimensions: [spin , [E(eV), dos(E), pdos(E)]]
            dos = dos_unformatted.reshape(-1, nEpts, 3)      
            
        elif dos_header == tot_pdos_polarized:
            dos_unformatted = np.float64(' '.join(data[1:]).split()).reshape(-1,5)
            nEpts = dos_unformatted.shape[0]
            
            # DOS dimensions: [spin , [E(eV), dos(E), pdos(E)]]
            dos = [dos_unformatted[:, [0, 1, 3]]]
            dos.append(dos_unformatted[:, [0, 2, 4]]) 
            
        else: 
            for lm in [1, 3, 5, 7, 9]:
                if dos_header == lm_pdos(lm=lm, polarized=False):
                    dos_unformatted = np.float64(' '.join(data[1:]).split()).reshape(-1, 2 + lm)
                    nEpts = dos_unformatted.shape[0]
            
                    # DOS dimensions: [spin , [E(eV), ldos(E), pdos(E) * lm]]
                    dos = dos_unformatted.reshape(-1, nEpts, 2 + lm)  
        
                elif dos_header == lm_pdos(lm=lm, polarized=True):
                    dos_unformatted = np.float64(' '.join(data[1:]).split()).reshape(-1, 3 + 2*lm)
                    nEpts = dos_unformatted.shape[0]
            
                    # DOS dimensions: [spin , [E(eV), ldos(E), pdos(E) * lm]]
                    dos = [dos_unformatted[:, [0, 1] + (np.arange(lm)*2 + 3).tolist()]]
                    dos.append(dos_unformatted[:, [0, 2] + (np.arange(lm)*2 + 4).tolist()])         
                    
        return np.asarray(dos)


tot_kdos            = "# ik    E (eV)  dos(E)    pdos(E)"
tot_kdos_polarized  = "# ik    E (eV)  dosup(E)   dosdw(E)  pdosup(E)  pdosdw(E)"

def lm_kdos(lm=1, polarized=False):
    if polarized:
        return ("# ik    E (eV)  ldosup(E)  ldosdw(E)" + lm * " pdosup(E)  pdosdw(E) ").strip()
    else:
        return ("# ik    E (eV)   ldos(E)   " + lm * "pdos(E)    ").strip()

def read_kdos_output(filename):  
    '''Read the dos files generated by projwfc.x
    
    ref: https://www.quantum-espresso.org/Doc/INPUT_PROJWFC.html
    Orbital Order
    Order of m-components for each l in the output:

        1, cos(phi), sin(phi), cos(2*phi), sin(2*phi), .., cos(l*phi), sin(l*phi)

    where phi is the polar angle:x=r cos(theta)cos(phi), y=r cos(theta)sin(phi)
    This is determined in file Modules/ylmr2.f90 that calculates spherical harmonics.

    for l=1:
      1 pz     (m=0)
      2 px     (real combination of m=+/-1 with cosine)
      3 py     (real combination of m=+/-1 with sine)

    for l=2:
      1 dz2    (m=0)
      2 dzx    (real combination of m=+/-1 with cosine)
      3 dzy    (real combination of m=+/-1 with sine)
      4 dx2-y2 (real combination of m=+/-2 with cosine)
      5 dxy    (real combination of m=+/-2 with sine)
    '''
    assert check_exist(filename), 'Cannot find : ' + filename
    with open(filename, 'r') as data_file:
        data = data_file.read().split('\n')
        dos_header = data[0].strip()
        if dos_header == tot_kdos:
            dos_unformatted = np.float64(' '.join(data[1:]).split()).reshape(-1,4)
            nkpts = int(dos_unformatted[-1,0])
            dos_unformatted = dos_unformatted[:,1:]
            nEpts = dos_unformatted.shape[0] // nkpts
            
            # DOS dimensions: [spin , nkpts, [E(eV), dos(E), pdos(E)]]
            dos = dos_unformatted.reshape(-1, nkpts, nEpts, 3)      
            
        elif dos_header == tot_kdos_polarized:
            dos_unformatted = np.float64(' '.join(data[1:]).split()).reshape(-1,6)
            nkpts = int(dos_unformatted[-1,0])
            dos_unformatted = dos_unformatted[:,1:]
            nEpts = dos_unformatted.shape[0] // nkpts
            
            # DOS dimensions: [spin , nkpts, [E(eV), dos(E), pdos(E)]]
            dos_formatted = dos_unformatted.reshape(nkpts, nEpts, -1)
            dos = [dos_formatted[:,:, [0] + (np.arange(2)*2+1).tolist()]]
            dos.append(dos_formatted[:,:, [0] + (np.arange(2)*2 + 2).tolist()]) 
            
        else: 
            
            for lm in [1, 3, 5, 7, 9]:
                if dos_header == lm_kdos(lm=lm, polarized=False):
                    dos_unformatted = np.float64(' '.join(data[1:]).split()).reshape(-1, 3 + lm)
                    nkpts = int(dos_unformatted[-1,0])
                    dos_unformatted = dos_unformatted[:,1:]
                    nEpts = dos_unformatted.shape[0] // nkpts
            
                    # DOS dimensions: [spin , nkpts, [E(eV), ldos(E), pdos(E) * lm]]
                    dos = dos_unformatted.reshape(-1, nkpts, nEpts, 2 + lm)  
                    
                elif dos_header == lm_kdos(lm=lm, polarized=True):
                    dos_unformatted = np.float64(' '.join(data[1:]).split()).reshape(-1, 4 + 2*lm)
                    nkpts = int(dos_unformatted[-1,0])
                    dos_unformatted = dos_unformatted[:,1:]
                    nEpts = dos_unformatted.shape[0] // nkpts
            
                    # DOS dimensions: [spin , nkpts, [E(eV), ldos(E), pdos(E) * lm]]
                    dos_formatted = dos_unformatted.reshape(nkpts, nEpts, -1)
                    dos = [dos_formatted[:,:, [0, 1] + (np.arange(lm)*2 + 3).tolist()]]
                    dos.append(dos_formatted[:,:, [0, 2] + (np.arange(lm)*2 + 4).tolist()])         
                    
        return np.asarray(dos)

        
# Get projections list
projections_MATCH = re.compile(r'''
[\w\W]*
 \(read [ ]* from [ ]* pseudopotential [ ]* files\):
\n
(?P<content>
  [\s\S]*?(?=\n.*?[ ] $|\n\n)  # match everything until next blank line or EOL
)
''', re.VERBOSE)

projection_MATCH = re.compile(r'''
[ ]*
 state [ ]* \# [ ]* (?P<state>\d+): [ ]* atom [ ]* (?P<site>\d+) [ ]* \((?P<species>\S+) [ ]* \), [ ]* wfc [ ]* (?P<wfc>\d+) [ ]* \(l=(?P<l>\d+) [ ]* m= [ ]* (?P<m>\d+)\)
''', re.VERBOSE)      
          
wfc_kpts_MATCH = re.compile(r'''
[ ]*
 k [ ]* =[ ]* (?P<kx>\S+) [ ]* (?P<ky>\S+) [ ]* (?P<kz>\S+)
(?P<content>
  [\s\S]*?(?=\n.*?[ ] $|\n\n)  # match everything until next blank line or EOL
)
''', re.VERBOSE)

wfc_e_MATCH = re.compile(r'''
[ ]*
 ==== [ ]* e\( [ ]* (?P<band>\d+)\) [ ]* = [ ]* (?P<eigenval>\S+) [ ]* eV [ ]* ==== 
(?P<content>
  [\s\S]*?(?=\n.*?[ ] $|\n====)  # match everything until next blank line or EOL
)
''', re.VERBOSE)
          
def read_projwfc_output(filename):  
    '''Read the output generated by projwfc.x
    
    Output:
    
        out['site']             : site #
        out['species']          : element
        out['wfc']              : the wf id for each atom
        out['l']                : angular momentum quantum number
        out['m']                : magnetic quantum number
        out['kpts']             : k-point list
        out['eigenvals']        : eigenvalue at each kpt, dimension: [spin, kpts, band]
        out['psi2']             : |psi|^2 
        out['projwfc']          : projected wave function, dimension: [spin, kpts, band, # of orbitals]
    '''

    assert check_exist(filename), 'Cannot find : ' + filename
    with open(filename, 'r') as data_file:
        data = data_file.read()

        # Get atom types
        projections_block = projections_MATCH.match(data).group('content')
        site = []
        species = []
        wfc_id = []
        l = []
        m = []
        for projection in projection_MATCH.finditer(projections_block):
            proj_dict = projection.groupdict()
            state = proj_dict['state']
            site.append(proj_dict['site'])
            species.append(proj_dict['species'])
            wfc_id.append(proj_dict['wfc'])
            l.append(proj_dict['l'])
            m.append(proj_dict['m'])

        # Get projected WF
        kpts = []
        nstate = int(state)
        eigenvals_kpts = []
        psi2_kpts = []
        proj_wf_kpts = []
        
        for wfc_kpt in wfc_kpts_MATCH.finditer(data):
            wfc_kpts_dict = wfc_kpt.groupdict()
            kpts.append([wfc_kpts_dict['kx'], wfc_kpts_dict['ky'], wfc_kpts_dict['kz']])
            wfc_kpts = wfc_kpts_dict['content']

            eigenvals_kpt = []
            proj_wf_kpt = []
            psi2_kpt = []
            for wfc_e in wfc_e_MATCH.finditer(wfc_kpts + '\n===='):       
                wfc_e_dict = wfc_e.groupdict()
                eigenvals_kpt.append(wfc_e_dict['eigenval'])
                wfc_e = wfc_e_dict['content']     
                temp = wfc_e.split("=")
                psi2 = float(temp[-1])
                psi2_kpt.append(psi2)
                wfc =  ("".join(temp[-2].split("\n")[:-1])).replace(' ','')
                num_ao = wfc.count("#")
                temp = wfc.replace('*[#', ' ').replace(']+', ' ').split()
                wf = np.float64(list(temp))
                state_idx = np.int64((wf[np.arange(num_ao)*2 + 1]).tolist()) - 1
                coeff = wf[np.arange(num_ao)*2]
                proj_wf = np.zeros(nstate)
                proj_wf[state_idx] = coeff
                proj_wf_kpt.append(proj_wf)
            psi2_kpts.append(psi2_kpt)
            eigenvals_kpts.append(eigenvals_kpt)
            proj_wf_kpts.append(proj_wf_kpt)
            
        kpts = np.float64(kpts) 
        eigenvals_kpts = np.float64(eigenvals_kpts) 
        psi2_kpts = np.float64(psi2_kpts)   
        proj_wf_kpts = np.float64(proj_wf_kpts)
        nkpts = kpts.shape[0]
        nband = eigenvals_kpts.shape[1]
        
        if "spin down" in data:
            nkpts = nkpts//2
            kpts = kpts[:nkpts,:]
            eigenvals_kpts = eigenvals_kpts.reshape(-1, nkpts, nband)
            psi2_kpts = psi2_kpts.reshape(-1, nkpts, nband)
            proj_wf_kpts = proj_wf_kpts.reshape(-1, nkpts, nband, nstate)
        else:
            kpts = kpts[:nkpts,:]
            eigenvals_kpts = eigenvals_kpts.reshape(-1, nkpts, nband)
            psi2_kpts = psi2_kpts.reshape(-1, nkpts, nband)
            proj_wf_kpts = proj_wf_kpts.reshape(-1, nkpts, nband, nstate)
            
        out = {}
        out['site'] = site
        out['species'] = species
        out['wfc'] = wfc_id
        out['l'] = l
        out['m'] = m
        out['kpts'] = kpts
        out['eigenvals'] = eigenvals_kpts 
        out['psi2'] = psi2_kpts   
        out['projwfc'] = proj_wf_kpts    

        return out
        
        