# -*- coding: utf-8 -*-
"""classes and functions to analysis output files from GAP calculation
"""

import numpy as np
from mykit.wien2k.utils import get_casename
from mykit.core.utils import trim_after, get_file_ext, conv_string
from mykit.core.unit import EnergyUnit
from mykit.core.bandstructure import BandStructure

class Emac:
    """Class for analyzing macroscopic dielectric constant in case.emac file
    """

    def __init__(self, path_emac=None):
        p = path_emac
        if p is None:
            cn = get_casename()
            p = cn + '.emac'
        with open(p, 'r') as h:
            data = np.loadtxt(h, unpack=True)
        self.lfe = np.shape(data)[0] == 6
        self._freq = data[1, :]
        self._data = data[2:, :]
    
    @property
    def nfreqs(self):
        """int. Number of frequency points"""
        return len(self._freq)

    def get_frequency_grids(self):
        """Return the frequency grid points"""
        return self._freq.copy()
    
    def get_macro_dielec_function(self, lfe=True):
        """Return the macroscopic dielectric constants at different
        frequencies

        Args:
            lfe (bool): if set True, dielectric constants with 
                local-field effect will be returned.
        
        Note:
            Dielectric constants without LFE will be returned regardless
            of ``lfe`` when eps^-1 was not computed.

        Returns:
            two 1d arrays, real and imaginary part of the 
            dielectric constant.
        """
        col = 0
        if self.lfe and not lfe:
            col = 2
        return self._data[col, :], self._data[col+1, :]


class Eqpev(EnergyUnit):

    """Class for reading files storing quais-paritcle energy data in eqpeV

    eqpH file is not considered, since all information included in eqpeV are
    available from eqpeV file.

    Args:
        path_eqp (str): data file generated by GAP.
            The extension should be the original one.
        gw0 (bool)

    Attributes:
        nibzkpts
        ibzkpts
        eps : Kohn-Sham eigenvalues
        eqp : quasi-particle energeies
        degw : change of absolute energy level by GW
    """

    _unit_name = {'ev': 'eV'}
    _gw_name = {'G0W0': 'GW', 'GW0': 'GW0'}
    _head_lines = 10
    
    def __init__(self, path_eqp=None, method='G0W0'):
        p = path_eqp
        
        if p is None:
            p = get_casename() + '.eqpeV'
            gw = self._gw_name.get(method, None)
            if gw is None:
                raise ValueError(f"method {method} is not available. {list(self._gw_name.keys())}")
            p += '_' + gw
        EnergyUnit.__init__(self, eunit='ev')
        # load data
        data = np.loadtxt(p)
        nkpts = int(data[-1, 0])
        nbandsgw = int(data[-1, 1] - data[0, 1]) + 1
        self.ibandsgw = np.array(data[0:nbandsgw+1, 1], dtype='int')
        self.eks = np.reshape(data[:, 2], (1, nkpts, nbandsgw), order="C")
        self.eqp = np.reshape(data[:, 3], (1, nkpts, nbandsgw), order="C")
        self.ehf = np.reshape(data[:, 4], (1, nkpts, nbandsgw), order="C")
        self.degw = np.reshape(data[:, -4], (1, nkpts, nbandsgw), order="C")
        self.nibzkpts = nkpts
        self.nbandsgw = nbandsgw
        self._ksbs = None
        self._qpbs = None
        self._hfbs = None
        # get kpoints information (weight not included)
        with open(p, 'r') as h:
            lines = h.readlines()
        self.ik = []
        self.ibzkpts = []
        for ik in range(nkpts):
            l = lines[self._head_lines+ik*(2+nbandsgw)]
            kpt = conv_string(l, int, 3, -5, -4, -3, -1)
            self.ibzkpts.append([x/kpt[-1] for x in kpt[1:4]])
            self.ik.append(kpt[0])

    def _analyze_kind(self, filepath):
        """Analyse the type of results by the extension name
        """
        ext = get_file_ext(filepath)[3:]
        u, gw = ext.split('_')
        unit = dict([(w, k) for k, w in self._unit_name.items()])[u]
        isg0w0 = dict([(w, k) for k, w in self._gw_name.items()])[gw]
        return isg0w0, unit

    @property
    def unit(self):
        """energy unit"""
        return self._eunit
    
    @unit.setter
    def unit(self, value):
        co = self._get_eunit_conversion(value)
        if co != 1:
            self.eks *= co
            self.eqp *= co
            self.ehf *= co
            self.degw *= co

    @property
    def ksbs(self):
        """Kohn-Sham band structure"""
        return self.get_KS_bandstructure()

    @property
    def qpbs(self):
        """Quasi-partilce band structure"""
        return self.get_QP_bandstructure()

    @property
    def hfbs(self):
        """Hartree-Fock band structure"""
        return self.get_HF_bandstructure()

    def _get_bandstructure(self, kind):
        e = {"qp": self.eqp, "ks": self.eks, "hf": self.ehf}
        assert kind in e.keys()
        e = e[kind]
        occ = 1.0 * (e < 0.009)
        weight = [1.0,] * self.nibzkpts
        return BandStructure(e, occ, weight, unit=self._eunit, efermi=0.0)

    def get_KS_bandstructure(self):
        """get Kohn-Sham band structure"""
        if self._ksbs is None:
            self._ksbs = self._get_bandstructure("ks")
        return self._ksbs

    def get_QP_bandstructure(self):
        """get quasi-particle band structure"""
        if self._qpbs is None:
            self._qpbs = self._get_bandstructure("qp")
        return self._qpbs

    def get_HF_bandstructure(self):
        """get Hartree-Fock band structure"""
        if self._hfbs is None:
            self._hfbs = self._get_bandstructure("hf")
        return self._hfbs

class Outgw:
    """Class for analyzing casename.outgw file
    """

    def __init__(self, path_out=None):
        p = path_out
        if p is None:
            cn = get_casename()
            p = cn + '.outgw'
    
