#!/usr/bin/env python

##############################################################################
##
# This file is part of NavARP
##
# Copyright 2016 CELLS / ALBA Synchrotron, Cerdanyola del Vallès, Spain
##
# NavARP is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
##
# NavARP is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
##
# You should have received a copy of the GNU General Public License
# along with NavARP.  If not, see <http://www.gnu.org/licenses/>.
##
##############################################################################

"""This module is part of the Python NavARP library. It defines the class
 for the DFT data loaded from a PROCAR."""

__author__ = ["Federico Bisti"]
__license__ = "GPL"
__date__ = "21/03/2024"

import os
import re
import numpy as np

from dataclasses import dataclass

@dataclass
class NavDFTEntry:
    """NavDFTEntry is the class for the theoretical DFT band structure.

    NavDFTEntry is defined from the VASP PROCAR.

    Attributes:
        folder_path (str): Folder path of the input file;
        ebins (ndarray): Binding energy axis;
        kxyz (ndarray): k-space vectors;
        porbs (ndarray): projections on atomic orbitals;
        pspins (ndarray): projections on spins.

    """
    folder_path: str
    ebins: np.ndarray
    kxyz: np.ndarray
    porbs: np.ndarray
    pspins: np.ndarray
    info: str

    def get_info(self):
        print(self.info)


def load_npz(
    npz_file_path
):
    print("Loading npz-file")
    data = np.load(npz_file_path)
    print("npz-file loaded")
    return NavDFTEntry(
        npz_file_path,
        data['ebins'],
        data['kxyz'],
        data['porbs'],
        data['pspins'],
        data['info']
    )


def load_procar(
    folder_path,
    procar='PROCAR',
    outcar='OUTCAR',
    outcar_scf='OUTCAR_scf',
    npz_file_name="navdftentry.npz",
    create_npz=True,
    overwrite=False
):
    """Load vasp DFT band structure."""
    npz_file_path = os.path.join(folder_path, npz_file_name)

    if (not overwrite) and (os.path.isfile(npz_file_path)):
        return load_npz(npz_file_path)

    print("writing new npz-file")

    # Reading Fermi level (efermi) from OUTCAR_scf
    rege_float = r"[-+]?\d*\.?\d+"
    with open(os.path.join(folder_path, outcar_scf)) as foutscf:
        for line in foutscf:
            if 'E-fermi' in line:
                efermi = float(
                    re.search(
                        r'E-fermi\s*?[:\s]\s*{}'.format(rege_float),
                        line
                    ).group().split(":")[1])
                print('efermi = {}'.format(efermi))

    # Reading k-points (kxyz in units of 1/SCALE) from OUTCAR
    with open(os.path.join(folder_path, outcar)) as foutcar:
        for line in foutcar:
            if 'LNONCOLLINEAR' in line:
                if 'T' in line:
                    lnoncollinear = True
                else:
                    lnoncollinear = False

            elif 'LORBIT' in line:
                lorbit = int(
                    re.search(
                        r'LORBIT\s*?=\s*\d*',
                        line
                    ).group().split("=")[1]
                )

            elif 'NKPTS' in line:
                nkpts = int(
                    re.search(
                        r'NKPTS\s*?=\s*\d*',
                        line
                    ).group().split("=")[1]
                )
                nbands = int(
                    re.search(
                        r'NBANDS\s*?=\s*\d*',
                        line
                    ).group().split("=")[1]
                )

            elif r'k-points in units of 2pi/SCALE' in line:
                break

        kxyz = np.zeros((nkpts, 3))
        for ik in range(0, nkpts):
            line = foutcar.readline().split()
            kxyz[ik, 0] = float(line[0])
            kxyz[ik, 1] = float(line[1])
            kxyz[ik, 2] = float(line[2])
        kxyz *= 2*np.pi

    print("From OUTCAR")
    print(
        f"nkpts = {nkpts}; lnoncollinear = {lnoncollinear}; lorbit = {lorbit}")

    # START reading PROCAR
    fprocar = open(os.path.join(folder_path, procar), "r")
    if "_dn" in npz_file_name:
        fprocar.readline()
        fprocar.readline()
        fprocar.readline()
        while True:
            line = fprocar.readline()
            if "# of k-points:" in line:
                line = line.replace(':', ' ').split()
                print("Spin down")
                break
    else:
        fprocar.readline()
        line = fprocar.readline().replace(':', ' ').split()
        print(line)
        print("Spin up")

    nkpts = int(line[3])
    nbands = int(line[7])
    nions = int(line[11])
    print("From POSCAR")
    print(f"nkpts = {nkpts}; nbands = {nbands}; nions = {nions};")

    ebins = np.zeros((nkpts, nbands))
    ion = np.zeros(nions, dtype=int)

    is_porbs_pspins_defined = False

    for ikpt in range(nkpts):
        step = round((ikpt+1) / nkpts*100, 4)
        if not (step % 5):
            print(f"k-point = {ikpt+1}/{nkpts}")

        for band in range(nbands):
            # print(f"band = {band}")
            for line in fprocar:
                if 'band' in line:
                    break
            ebins[ikpt, band] = float(line.split()[4])
            fprocar.readline()
            info_line = fprocar.readline()
            if not is_porbs_pspins_defined:
                norbitals = len(info_line.split()) - 2
                info = ' '.join(
                    [f'{ind}:{str},' for ind, str in enumerate(
                        info_line.strip('ion').strip('tot\n').split())
                    ]
                )[:-1]

                print(f"norbitals = {norbitals}")
                if lnoncollinear:
                    nspins = 4  # sx, sy, sz, tot
                else:
                    nspins = 1  # tot
                pspins = np.zeros((nkpts, nbands, nions, norbitals+1, nspins))

                # s     py     pz     px    dxy    dyz    dz2    dxz    dx2
                if lorbit in [2, 12, 13, 14]:
                    porbs = np.zeros(
                        (nkpts, nbands, nions, norbitals), dtype=complex)
                else:
                    porbs = np.array([])
                is_porbs_pspins_defined = True

            for ispin in range(nspins):
                # print(f"ispin = {ispin}")
                for ion in range(0, nions):
                    # print(f"ion = {ion}")
                    pspins[ikpt, band, ion, :, ispin] = [
                        float(num) for num in fprocar.readline().split()[1:]]
                if nions != 1:
                    fprocar.readline()

            if lorbit in [2, 12, 13, 14]:
                for line in fprocar:
                    if 'ion' in line:
                        break

                appoR = []
                appoI = []
                for ion in range(0, nions):
                    line = fprocar.readline().split()
                    nums = [
                        float(num) for num in line[1:]]
                    if len(nums) >= norbitals*2:
                        appoR.append(nums[0:norbitals*2:2])
                        appoI.append(nums[1:norbitals*2:2])
                    else:
                        appoR.append(nums)
                        appoI.append([
                        float(num) for num in fprocar.readline().split()[1:]])
                porbs[ikpt, band, :, :] = np.array(appoR) + np.array(appoI)*1j

    ebins -= efermi
    if create_npz:
        np.savez(
            npz_file_path,
            kxyz=kxyz,
            ebins=ebins,
            porbs=porbs,
            pspins=pspins,
            info=info
        )
    fprocar.close()

    return NavDFTEntry(
        os.path.join(folder_path, procar), ebins, kxyz, porbs, pspins, info)


def get_efermi_from_outcar(
    outcar_path='OUTCAR',
    print_efermi='False'
):
    # Reading Fermi level (efermi) from OUTCAR_scf
    rege_float = r"[-+]?\d*\.?\d+"
    with open(outcar_path) as foutscf:
        for line in foutscf:
            if 'E-fermi' in line:
                efermi = float(
                    re.search(
                        r'E-fermi\s*?[:\s]\s*{}'.format(rege_float),
                        line
                    ).group().split(":")[1])
                if print_efermi:
                    print('efermi = {}'.format(efermi))
    return efermi
