from __future__ import annotations
import numpy as np
import re
from monty.re import regrep
from echem.core.structure import Structure
from echem.core.constants import Bohr2Angstrom, Angstrom2Bohr, Hartree2eV, eV2Hartree
from echem.core.ionic_dynamics import IonicDynamics
from echem.core.electronic_structure import EBS
from echem.core.thermal_properties import ThermalProperties
from echem.io_data import vasp
from echem.io_data.universal import Cube
from typing import Union, Literal, TypedDict
from typing_extensions import NotRequired
from pathlib import Path
import warnings
import copy
from nptyping import NDArray, Shape, Number
from termcolor import colored


class Lattice:
    def __init__(self,
                 lattice: NDArray[Shape['3, 3'], Number]):
        self.lattice = lattice

    @staticmethod
    def from_file(filepath: str | Path):

        if isinstance(filepath, str):
            filepath = Path(filepath)

        file = open(filepath, 'r')
        data = file.readlines()
        file.close()

        patterns = {'lattice': r'^\s*lattice\s+'}
        matches = regrep(str(filepath), patterns)

        lattice = []
        i = 0
        while len(lattice) < 9:
            line = data[matches['lattice'][0][1] + i].split()
            for word in line:
                try:
                    word = float(word)
                    lattice.append(word)
                except:
                    pass
            i += 1

        lattice = np.array(lattice).reshape((3, 3))

        return Lattice(lattice)

    def to_file(self, filepath: str):
        file = open(filepath, 'w')

        file.write('lattice \\\n')
        width_coords_float = max(len(str(int(np.max(self.lattice)))), len(str(int(np.min(self.lattice))))) + 16
        for i, vector in enumerate(self.lattice):
            file.write('\t')
            for vector_i in vector:
                file.write(f'{vector_i:{width_coords_float}.15f}  ')
            if i < 2:
                file.write('\\')
            file.write('\n')

        file.close()


class Ionpos:
    def __init__(self,
                 species: list[str],
                 coords: NDArray[Shape['Natoms, 3'], Number],
                 move_scale: list[int] | NDArray[Shape['Natoms'], Number] = None,
                 constraint_type: list[Literal['HyperPlane', 'Linear', 'None', 'Planar'] | None] = None,
                 constraint_params: list[list[float] | None] = None):

        self.species = species
        self.coords = coords
        if move_scale is None:
            move_scale = np.ones(len(coords), dtype=int)
        elif isinstance(move_scale, list):
            move_scale = np.array(move_scale, dtype=int)
        self.move_scale = move_scale
        self.constraint_type = constraint_type
        self.constraint_params = constraint_params

    @staticmethod
    def from_file(filepath: str | Path):
        if isinstance(filepath, str):
            filepath = Path(filepath)

        file = open(filepath, 'r')
        data = file.readlines()
        file.close()

        patterns = {'coords': r'^\s*ion\s+'}
        matches = regrep(str(filepath), patterns)

        natoms = len(matches['coords'])
        species = []
        coords = np.zeros((natoms, 3))
        move_scale = np.zeros(natoms, dtype=int)
        constraint_type = []
        constraint_params = []

        for i, ion in enumerate(matches['coords']):
            line = data[ion[1]].split()
            species.append(line[1])
            coords[i] = [line[2], line[3], line[4]]
            move_scale[i] = line[5]
            if len(line) > 6:
                constraint_type.append(line[6])
                constraint_params.append([float(line[7]), float(line[8]), float(line[9])])
            else:
                constraint_type.append(None)
                constraint_params.append(None)

        return Ionpos(species, coords, move_scale, constraint_type, constraint_params)

    def to_file(self,
                filepath: str | Path) -> None:
        if isinstance(filepath, str):
            filepath = Path(filepath)

        file = open(filepath, 'w')

        width_species = max([len(sp) for sp in self.species])
        width_coords_float = max(len(str(int(np.max(self.coords)))), len(str(int(np.min(self.coords))))) + 16

        if self.constraint_params is None and self.constraint_type is None:
            for sp, coord, ms in zip(self.species, self.coords, self.move_scale):
                file.write(f'ion {sp:{width_species}}  ')
                for coord_i in coord:
                    file.write(f'{coord_i:{width_coords_float}.15f}  ')
                file.write(f'{ms}\n')
        elif self.constraint_params is not None and self.constraint_type is not None:
            for sp, coord, ms, ctype, cparams in zip(self.species, self.coords, self.move_scale,
                                                     self.constraint_type, self.constraint_params):
                file.write(f'ion {sp:{width_species}}  ')
                for coord_i in coord:
                    file.write(f'{coord_i:{width_coords_float}.15f}  ')
                if ctype is None:
                    file.write(f'{ms}\n')
                else:
                    file.write(f'{ms}  ')
                    file.write(f'{ctype}  ')
                    file.write(f'{cparams[0]}  {cparams[1]}  {cparams[2]}\n')
        else:
            raise ValueError('constraint_type and constraint_params must be both specified or both be None')

        file.close()

    def convert(self,
                format: Literal['vasp'], *args):
        if format == 'vasp':
            lattice = np.transpose(args[0].lattice) * Bohr2Angstrom
            return vasp.Poscar(Structure(lattice, self.species, self.coords * Bohr2Angstrom))
        else:
            raise NotImplemented('Currently only format=vasp is supported')

    def get_structure(self,
                      lattice: Lattice) -> Structure:
        return Structure(lattice.lattice * Bohr2Angstrom, self.species, self.coords * Bohr2Angstrom)


class Input:
    def __init__(self, commands: list[tuple[str, str]]):
        self.commands = commands

    @staticmethod
    def from_file(filepath: str | Path):
        file = open(filepath, 'r')
        data = file.readlines()
        file.close()

        commands = []
        to_append = ''
        for line in data:
            line = line.strip().strip('\n')
            if line.endswith('\\'):
                to_append += re.sub(r'\s+', ' ', line.strip('\\'))
            else:
                if len(line) == 0:
                    continue
                to_append += line

                line = to_append.split()
                if line[0] == 'dump':
                    for i in line[2:]:
                        commands.append(('dump', f'{line[1]} {i}'))
                else:
                    commands.append((line[0], ' '.join(line[1:])))
                to_append = ''

        return Input(commands)


class EnergyIonicHist(TypedDict):
    F: NDArray[Shape['Nsteps'], Number]
    G: NotRequired[NDArray[Shape['Nsteps'], Number]]
    muN: NotRequired[NDArray[Shape['Nsteps'], Number]]


class Output(IonicDynamics):
    def __init__(self,
                 fft_box_size: NDArray[Shape['3'], Number],
                 energy_ionic_hist: EnergyIonicHist,
                 coords_hist: NDArray[Shape['Nsteps, Natoms, 3'], Number],
                 forces_hist: NDArray[Shape['Nsteps, Natoms, 3'], Number] | None,
                 nelec_hist: NDArray[Shape['Nsteps'], Number],
                 magnetization_hist: NDArray[Shape['Nesteps, 2'], Number] | None,
                 structure: Structure,
                 nbands: int,
                 nkpts: int,
                 mu: float | None,
                 HOMO: float | None,
                 LUMO: float | None,
                 phonons: dict[Literal['real', 'imag', 'zero', 'nStates'], np.ndarray | None],
                 pseudopots: dict,
                 lowding: dict[str, float] | None):
        super(Output, self).__init__(forces_hist, coords_hist, structure.lattice, True)
        self.fft_box_size = fft_box_size
        self.energy_ionic_hist = energy_ionic_hist
        self.coords_hist = coords_hist
        self.nelec_hist = nelec_hist
        self.magnetization_hist = magnetization_hist
        self.structure = structure
        self.nbands = nbands
        self.nkpts = nkpts
        self.mu = mu
        self.HOMO = HOMO
        self.LUMO = LUMO
        self.phonons = phonons
        self.pseudopots = pseudopots
        self.lowdin = lowding
        if phonons['real'] is not None and len(phonons['real']) > 0:
            self.thermal_props = ThermalProperties(np.array([phonons['real']]) * Hartree2eV)

    @property
    def energy(self) -> float:
        if 'G' in self.energy_ionic_hist.keys():
            return self.energy_ionic_hist['G'][-1]
        else:
            return self.energy_ionic_hist['F'][-1]

    @property
    def nisteps(self) -> int:
        return len(self.energy_ionic_hist['F'])

    @property
    def nelec(self) -> float:
        return self.nelec_hist[-1]

    @property
    def nelec_pzc(self) -> int:
        return np.sum([self.structure.natoms_by_type[key] * self.pseudopots[key] for key in self.pseudopots.keys()])

    @property
    def magnetization_abs(self) -> float:
        if self.magnetization_hist is None:
            raise ValueError('It is non-spin-polarized calculation')
        else:
            return self.magnetization_hist[-1, 0]

    @property
    def magnetization_tot(self) -> float:
        if self.magnetization_hist is None:
            raise ValueError('It is non-spin-polarized calculation')
        else:
            return self.magnetization_hist[-1, 1]

    @property
    def nspin(self) -> int:
        if self.magnetization_hist is None:
            return 1
        else:
            return 2

    @staticmethod
    def from_file(filepath: str | Path):
        if isinstance(filepath, str):
            filepath = Path(filepath)

        # \TODO Non-Cartesin coods case is not implemented

        file = open(filepath, 'r')
        data = file.readlines()
        file.close()

        patterns = {'natoms': r'Initialized \d+ species with (\d+) total atoms.',
                    'coords': r'# Ionic positions in cartesian coordinates:',
                    'forces': r'# Forces in Cartesian coordinates:',
                    'fft_box_size': r'Chosen fftbox size, S = \[(\s+\d+\s+\d+\s+\d+\s+)\]',
                    'lattice': r'---------- Initializing the Grid ----------',
                    'nbands': r'nBands:\s+(\d+)',
                    'nkpts': r'Reduced to (\d+) k-points under symmetry',
                    'nkpts_folded': r'Folded \d+ k-points by \d+x\d+x\d+ to (\d+) k-points.',
                    'is_kpts_irreducable': r'No reducable k-points',
                    'nelec': r'nElectrons:\s+(\d+.\d+)',
                    'magnetization': r'magneticMoment:\s+\[\s+Abs:\s+(\d+.\d+)\s+Tot:\s+([-+]?\d*\.\d*)',
                    'mu': r'\s+mu\s+:\s+([-+]?\d*\.\d*)',
                    'mu_hist': r'mu:\s+([-+]?\d*\.\d*)',
                    'HOMO': r'\s+HOMO\s*:\s+([-+]?\d*\.\d*)',
                    'LUMO': r'\s+LUMO\s*:\s+([-+]?\d*\.\d*)',
                    'F': r'^\s*F\s+=\s+([-+]?\d*\.\d*)',
                    'muN': r'\s+muN\s+=\s+([-+]?\d*\.\d*)',
                    'G': r'\s+G\s+=\s+([-+]?\d*\.\d*)',
                    'phonon report': r'(\d+) imaginary modes, (\d+) modes within cutoff, (\d+) real modes',
                    'zero mode': r'Zero mode \d+:',
                    'imaginary mode': r'Imaginary mode \d+:',
                    'real mode': r'Real mode \d+:',
                    'ionic convergence': r'IonicMinimize: Converged',
                    'pseudopots': r'\s*Title:\s+([a-zA-Z0-9]*).',
                    'valence_elecs': r'(\d+) valence electrons in orbitals',
                    'phonon_perturbations': r'\s+Perturbation:\s+\d+\s+nStates:\s+(\d+)',
                    'lowdin': r'#--- Lowdin population analysis ---'}

        matches = regrep(str(filepath), patterns)

        F = np.array([float(i[0][0]) for i in matches['F']])
        energy_ionic_hist: EnergyIonicHist = {'F': F}
        if 'muN' in matches.keys():
            energy_ionic_hist['muN'] = np.array([float(i[0][0]) for i in matches['muN']])
        if 'G' in matches.keys():
            energy_ionic_hist['G'] = np.array([float(i[0][0]) for i in matches['G']])

        nelec_hist = np.array([float(i[0][0]) for i in matches['nelec']])

        natoms = int(matches['natoms'][0][0][0])
        nbands = int(matches['nbands'][0][0][0])

        phonons = {}

        if matches['phonon_perturbations']:
            nstates = [int(i[0][0]) for i in matches['phonon_perturbations']]
            phonons['nStates'] = np.array(nstates)
        else:
            phonons['nStates'] = None

        if bool(matches['is_kpts_irreducable']):
            nkpts = int(matches['nkpts_folded'][0][0][0])
        else:
            nkpts = int(matches['nkpts'][0][0][0])
        if bool(matches['mu']):
            mu = float(matches['mu'][0][0][0])
        elif matches['mu_hist']:
            mu = float(matches['mu_hist'][-1][0][0])
        else:
            mu = None

        if bool(matches['HOMO']):
            HOMO = float(matches['HOMO'][0][0][0])
        else:
            HOMO = None

        if bool(matches['LUMO']):
            LUMO = float(matches['LUMO'][0][0][0])
        else:
            LUMO = None

        if bool(matches['magnetization']):
            magnetization_hist = np.zeros((len(matches['magnetization']), 2))
            for i, mag in enumerate(matches['magnetization']):
                magnetization_hist[i] = [float(mag[0][0]), float(mag[0][1])]
        else:
            magnetization_hist = None

        fft_box_size = np.array([int(i) for i in matches['fft_box_size'][0][0][0].split()])

        lattice = np.zeros((3, 3))
        lattice[0] = [float(i) for i in data[matches['lattice'][0][1] + 2].split()[1:4]]
        lattice[1] = [float(i) for i in data[matches['lattice'][0][1] + 3].split()[1:4]]
        lattice[2] = [float(i) for i in data[matches['lattice'][0][1] + 4].split()[1:4]]
        lattice = lattice.T * Bohr2Angstrom

        if matches['forces']:
            line_numbers = [int(i[1]) + 1 for i in matches['forces']]
            forces_hist = np.zeros((len(line_numbers), natoms, 3))
            for i, line_number in enumerate(line_numbers):
                atom_number = 0
                while len(line := data[line_number + atom_number].split()) > 0:
                    forces_hist[i, atom_number] = [float(line[2]), float(line[3]), float(line[4])]
                    atom_number += 1
        else:
            forces_hist = None

        if matches['phonon report']:
            freq_report = {key: int(i) for key, i in zip(['imaginary modes', 'modes within cutoff', 'real modes'],
                                                         matches['phonon report'][0][0])}
            if freq_report['modes within cutoff']:
                line_numbers = [int(i[1]) + 1 for i in matches['zero mode']]
                zero_mode_freq = np.zeros(freq_report['modes within cutoff'], dtype=complex)
                for i, line_number in enumerate(line_numbers):
                    zero_mode_freq[i] = complex(data[line_number].split()[1].replace('i', 'j'))
            else:
                zero_mode_freq = None

            if freq_report['imaginary modes']:
                line_numbers = [int(i[1]) + 1 for i in matches['imaginary mode']]
                imag_mode_freq = np.zeros(freq_report['imaginary modes'], dtype=complex)
                for i, line_number in enumerate(line_numbers):
                    imag_mode_freq[i] = complex(data[line_number].split()[1].replace('i', 'j'))
            else:
                imag_mode_freq = None

            if freq_report['real modes']:
                line_numbers = [int(i[1]) + 1 for i in matches['real mode']]
                real_mode_freq = np.zeros(freq_report['real modes'])
                for i, line_number in enumerate(line_numbers):
                    real_mode_freq[i] = float(data[line_number].split()[1])
            else:
                real_mode_freq = None

            phonons['zero'] = zero_mode_freq
            phonons['imag'] = imag_mode_freq
            phonons['real'] = real_mode_freq
        else:
            phonons['zero'] = None
            phonons['imag'] = None
            phonons['real'] = None

        if matches['coords']:
            line_numbers = [int(i[1]) + 1 for i in matches['coords']]
            coords_hist = np.zeros((len(line_numbers), natoms, 3))
            species = []
            atom_number = 0
            while len(line := data[line_numbers[0] + atom_number].split()) > 0:
                species += [line[1]]
                atom_number += 1
            for i, line_number in enumerate(line_numbers):
                atom_number = 0
                while len(line := data[line_number + atom_number].split()) > 0:
                    coords_hist[i, atom_number] = [float(line[2]), float(line[3]), float(line[4])]
                    atom_number += 1
        else:
            matches = regrep(str(filepath), {'ions': r'ion\s+([a-zA-Z]+)\s+[-+]?\d*\.\d*',
                                             'coords': r'ion\s+[a-zA-Z]+\s+([-+]?\d*\.\d*)\s+([-+]?\d*\.\d*)\s+([-+]?\d*\.\d*)'})
            species = [i[0][0] for i in matches['ions']]

            coords_hist = [[[float(i) for i in coord[0]] for coord in matches['coords']]]
            coords_hist = np.array(coords_hist)

        if bool(matches['lowdin']):
            lowdin = {}
            i = matches['lowdin'][-1][1] + 1
            while (line := data[i]) != '\n':
                line = line.split()
                lowdin[line[2]] = [float(i) for i in line[3:]]
                if bool(matches['magnetization']):
                    i += 2
                else:
                    i += 1
        else:
            lowdin = None

        structure = Structure(lattice, species, coords_hist[-1] * Bohr2Angstrom, coords_are_cartesian=True)

        pseudopots = {i[0][0]: int(j[0][0]) for i, j in zip(matches['pseudopots'], matches['valence_elecs'])}

        return Output(fft_box_size, energy_ionic_hist, coords_hist, forces_hist, nelec_hist, magnetization_hist,
                      structure, nbands, nkpts, mu, HOMO, LUMO, phonons, pseudopots, lowdin)

    def get_xdatcar(self) -> vasp.Xdatcar:
        transform = np.linalg.inv(self.structure.lattice)
        return vasp.Xdatcar(structure=self.structure,
                            trajectory=np.matmul(self.coords_hist * Bohr2Angstrom, transform))

    def get_poscar(self) -> vasp.Poscar:
        structure = copy.copy(self.structure)
        structure.coords = self.coords_hist[0] * Bohr2Angstrom
        return vasp.Poscar(structure=structure)

    def get_contcar(self) -> vasp.Poscar:
        return vasp.Poscar(structure=self.structure)

    def get_ionpos(self, nstep=-1) -> Ionpos:
        return Ionpos(self.structure.species, self.coords_hist[nstep])

    def get_lattice(self) -> Lattice:
        return Lattice(self.structure.lattice * Angstrom2Bohr)

    def mod_phonon_zero2real(self, n_leave: int = 0) -> None:
        if self.phonons['zero'] is not None:
            mask_real = self.phonons['zero'].imag == 0
            mask_complex = np.invert(mask_real)

            n_real = np.sum(mask_real)
            n_imag = np.sum(mask_complex)
            n_zero = len(self.phonons['zero'])
            if n_zero < n_leave:
                print(colored(f'There is only {n_zero} zero modes, however you set {n_leave=}',
                              color='red', attrs=['bold']))
            elif n_zero > n_leave:
                if n_leave > n_imag:
                    n_transfer = n_real - (n_leave - n_imag)
                else:
                    n_transfer = np.sum(mask_real)

                mods_for_transfer = None

                if n_zero - n_transfer > n_leave:
                    print(colored(f'Can not leave', color='red', attrs=['bold']),
                          n_leave,
                          colored('modes, because there are', color='red', attrs=['bold']),
                          n_imag,
                          colored('imaginary modes', color='red', attrs=['bold']))
                    print(colored('The following values can not be converted to real:', color='red', attrs=['bold']),
                          self.phonons['zero'][mask_complex])

                    if np.any(mask_real):
                        mods_for_transfer = np.sort(self.phonons['zero'][mask_real].real)
                        print(colored('The following values will be converted to real:', color='red', attrs=['bold']),
                              self.phonons['zero'][mask_real])

                else:
                    mods_for_transfer = np.sort(self.phonons['zero'][mask_real].real)[-n_transfer:]

                if mods_for_transfer is not None:
                    self.phonons['real'] = np.hstack((mods_for_transfer, self.phonons['real']))
                    del_indices = []
                    for mode in mods_for_transfer:
                        del_indices.append(np.where(self.phonons['zero'] == mode)[0][0])
                    self.phonons['zero'] = np.delete(self.phonons['zero'], del_indices)
        else:
            print(colored('There are no zero phonons', color='green', attrs=['bold']))


class EBS_data:

    @staticmethod
    def from_file(filepath: str | Path,
                  output: Output) -> NDArray[Shape['Nspin, Nkpts, Nbands'], Number]:

        if isinstance(filepath, str):
            filepath = Path(filepath)

        data = np.fromfile(filepath, dtype=np.float64)
        if len(data) % (output.nkpts * output.nbands) != 0:
            raise ValueError(
                f'Number of eigenvalues should be equal to nspin * nkpts * nbands, but now {output.nkpts=},'
                f'{output.nbands=}, and data has {len(data)} values')
        nspin = len(data) // (output.nkpts * output.nbands)
        data = data.reshape(nspin, output.nkpts, output.nbands)

        return data


class Eigenvals(EBS_data):
    def __init__(self,
                 eigenvalues: NDArray[Shape['Nspin, Nkpts, Nbands'], Number],
                 units: Literal['eV', 'Hartree']):
        self.eigenvalues = eigenvalues
        self.units = units

    @staticmethod
    def from_file(filepath: str | Path,
                  output: Output) -> 'Eigenvals':

        if isinstance(filepath, str):
            filepath = Path(filepath)

        eigenvalues = super(Eigenvals, Eigenvals).from_file(filepath, output)
        return Eigenvals(eigenvalues, 'Hartree')

    def mod_to_eV(self):
        if self.units == 'eV':
            print('Units are already eV')
        else:
            self.eigenvalues *= Hartree2eV
            self.units = 'eV'

    def mod_to_Ha(self):
        if self.units == 'Hartree':
            print('Units are already Hartree')
        else:
            self.eigenvalues *= eV2Hartree
            self.units = 'Hartree'


class Fillings(EBS_data):
    def __init__(self,
                 occupations: np.ndarray):
        self.occupations = occupations

    @staticmethod
    def from_file(filepath: str | Path,
                  output: Output) -> 'Fillings':

        if isinstance(filepath, str):
            filepath = Path(filepath)

        occupations = super(Fillings, Fillings).from_file(filepath, output)
        return Fillings(occupations)


class VolumetricData:
    def __init__(self,
                 data: np.ndarray,
                 structure: Structure):
        self.data = data
        self.structure = structure

    def __add__(self, other):
        assert isinstance(other, VolumetricData), 'Other object must belong to VolumetricData class'
        assert self.data.shape == other.data.shape, f'Shapes of two data arrays must be the same but they are ' \
                                                    f'{self.data.shape} and {other.data.shape}'
        if self.structure != other.structure:
            warnings.warn('Two VolumetricData instances contain different Structures. '
                          'The Structure will be taken from the 2nd (other) instance. '
                          'Hope you know, what you are doing')
        return VolumetricData(self.data + other.data, other.structure)

    def __sub__(self, other):
        assert isinstance(other, VolumetricData), 'Other object must belong to VolumetricData class'
        assert self.data.shape == other.data.shape, f'Shapes of two data arrays must be the same but they are ' \
                                                    f'{self.data.shape} and {other.data.shape}'
        if self.structure != other.structure:
            warnings.warn('Two VolumetricData instances contain different Structures. '
                          'The Structure will be taken from the 2nd (other) instance. '
                          'Hope you know, what you are doing')
        return VolumetricData(self.data - other.data, other.structure)

    @staticmethod
    def from_file(filepath: str | Path,
                  fft_box_size: NDArray[Shape['3'], Number],
                  structure: Structure):

        if isinstance(filepath, str):
            filepath = Path(filepath)

        data = np.fromfile(filepath, dtype=np.float64)
        data = data.reshape(fft_box_size)
        return VolumetricData(data, structure)

    def convert_to_cube(self) -> Cube:
        return Cube(self.data, self.structure, np.zeros(3))


class kPts:
    def __init__(self,
                 weights: NDArray[Shape['Nkpts'], Number]):
        self.weights = weights

    @staticmethod
    def from_file(filepath: str | Path):

        if isinstance(filepath, str):
            filepath = Path(filepath)

        file = open(filepath, 'r')
        data = file.readlines()
        file.close()
        weights = []

        if 'spin' in data[0].split():
            for line in data[:int(len(data) / 2)]:
                weights.append(float(line.split()[6]))
            weights = np.array(weights)
            return kPts(weights)
        else:
            for line in data:
                weights.append(float(line.split()[6]))
            weights = np.array(weights) / 2
            return kPts(weights)


class DOS(EBS):
    @staticmethod
    def from_folder(folderpath: str | Path,
                    output_name: str = 'output.out',
                    jdft_prefix='jdft',
                    units: Literal['eV', 'Ha'] = 'eV'):

        if isinstance(folderpath, str):
            folderpath = Path(folderpath)

        out = Output.from_file(folderpath / output_name)
        eigs = Eigenvals.from_file(folderpath / f'{jdft_prefix}.eigenvals', output=out)
        fills = Fillings.from_file(folderpath / f'{jdft_prefix}.fillings', output=out)
        kpts = kPts.from_file(folderpath / f'{jdft_prefix}.kPts')
        if units == 'eV':
            return DOS(eigenvalues=eigs.eigenvalues * Hartree2eV,
                       weights=kpts.weights,
                       efermi=out.mu * Hartree2eV,
                       occupations=fills.occupations)
        elif units == 'Ha':
            return DOS(eigenvalues=eigs.eigenvalues,
                       weights=kpts.weights,
                       efermi=out.mu,
                       occupations=fills.occupations)
        else:
            raise ValueError(f'units can be "eV" or "Ha", however you entered "{units}"')


class BandProjections:
    def __init__(self,
                 proj_coeffs: NDArray[Shape['Nspin, Nkpts, Nbands, Norbs'], Number],
                 weights: NDArray[Shape['Nkpts'], Number],
                 species: list[str],
                 norbs_per_atomtype: dict,
                 orbs_names: list[str],
                 orbs_data: list[dict]):
        self.proj_coeffs = proj_coeffs
        self.weights = weights
        self.species = species
        self.norbs_per_atomtype = norbs_per_atomtype
        self.orbs_names = orbs_names
        self.orbs_data = orbs_data

        self.eigenvalues = None

    @property
    def nspin(self):
        return self.proj_coeffs.shape[0]

    @property
    def nkpts(self):
        return self.proj_coeffs.shape[1]

    @property
    def nbands(self):
        return self.proj_coeffs.shape[2]

    @property
    def norbs(self):
        return self.proj_coeffs.shape[3]

    @staticmethod
    def from_file(filepath: str | Path):

        if isinstance(filepath, str):
            filepath = Path(filepath)

        file = open(filepath, 'r')
        data = file.readlines()
        file.close()

        patterns = {'x': r'#\s+\d+'}
        matches = regrep(str(filepath), patterns)

        nstates = int(data[0].split()[0])
        nbands = int(data[0].split()[2])
        norbs = int(data[0].split()[4])

        if 'spin' in data[int(matches['x'][0][1])]:
            nspin = 2
        else:
            nspin = 1

        nkpts = int(nstates / nspin)

        proj_coeffs = np.zeros((nspin, nkpts, nbands, norbs))
        weights = np.zeros(nstates)

        start_lines = []
        for i, match in enumerate(matches['x']):
            start_lines.append(int(match[1]))
            weights[i] = float(re.sub(r'[^0-9.]', '', data[int(match[1])].split()[7]))

        if nspin == 2 and not np.array_equal(weights[:len(weights) // 2], weights[len(weights) // 2:]):
            raise ValueError(f'Kpts weights can not be correctly split {weights=}')
        if nspin == 2:
            weights = weights[:len(weights) // 2]

        species = []
        norbs_per_atomtype = {}
        orbs_names = []
        orbs_data = []

        idx_atom = -1
        for iline in range(2, start_lines[0]):
            line = data[iline].split()
            atomtype = line[0]
            natoms_per_atomtype = int(line[1])
            species += [atomtype] * natoms_per_atomtype
            norbs_per_atomtype[line[0]] = int(line[2])

            l_max = int(line[3])
            nshalls_per_l = []
            for i in range(l_max + 1):
                nshalls_per_l.append(int(line[4 + i]))

            for i in range(natoms_per_atomtype):
                idx_atom += 1
                for l, n_max in zip(range(l_max + 1), nshalls_per_l):
                    for n in range(n_max):
                        if l == 0:
                            orbs_names.append(f'{idx_atom} {atomtype} s {n + 1}({n_max})')
                            orbs_data.append({'atom_type': atomtype,
                                              'atom_index': idx_atom,
                                              'l': l,
                                              'm': 0,
                                              'orb_name': 's'})
                        elif l == 1:
                            for m, m_name in zip([-1, 0, 1], ['p_x', 'p_y', 'p_z']):
                                orbs_names.append(f'{idx_atom} {atomtype} {m_name} {n + 1}({n_max})')
                                orbs_data.append({'atom_type': atomtype,
                                                  'atom_index': idx_atom,
                                                  'l': l,
                                                  'm': m,
                                                  'orb_name': m_name})
                        elif l == 2:
                            for m, m_name in zip([-2, -1, 0, 1, 2], ['d_xy', 'd_yz', 'd_z^2', 'd_xz', 'd_x^2-y^2']):
                                orbs_names.append(f'{idx_atom} {atomtype} {m_name} {n + 1}({n_max})')
                                orbs_data.append({'atom_type': atomtype,
                                                  'atom_index': idx_atom,
                                                  'l': l,
                                                  'm': m,
                                                  'orb_name': m_name})
                        elif l > 2:
                            raise NotImplementedError('Only s, p snd d orbitals are currently supported')

        ikpt_major = -1
        ikpt_minor = -1
        for istate, (start, stop) in enumerate(zip(start_lines[:-1], start_lines[1:])):
            if nspin == 2:
                if data[start].split()[9] == '+1;':
                    ispin = 0
                    ikpt_major += 1
                    ikpt = ikpt_major
                elif data[start].split()[9] == '-1;':
                    ispin = 1
                    ikpt_minor += 1
                    ikpt = ikpt_minor
                else:
                    raise ValueError(f'Can\'t determine spin in string {data[start].split()}')
            else:
                ispin = 0
                ikpt = istate

            for iband, line in enumerate(range(start + 1, stop)):
                proj_coeffs[ispin, ikpt, iband] = [float(k) for k in data[line].split()]

        return BandProjections(proj_coeffs, weights / 2, species, norbs_per_atomtype, orbs_names, orbs_data)

    def get_PDOS(self,
                 atom_numbers: list[int] | int,
                 eigenvals: Eigenvals,
                 get_orbs_names: bool = False,
                 specific_l: int = None,
                 dE: float = 0.01,
                 emin: float = None,
                 emax: float = None,
                 zero_at_fermi: bool = False,
                 sigma: float = 0.02,
                 efermi: float = None) -> Union[tuple[NDArray[Shape['Ngrid'], Number],
                                                      NDArray[Shape['Nspin, Norbs_selected, Ngrid'], Number]],
                                                tuple[NDArray[Shape['Ngrid'], Number],
                                                      NDArray[Shape['Nspin, Norbs_selected, Ngrid'], Number],
                                                      list[str]]]:
        self.eigenvalues = eigenvals.eigenvalues
        if isinstance(atom_numbers, int):
            atom_numbers = [atom_numbers]

        if zero_at_fermi is True and efermi is None:
            raise ValueError('You can not set zero_at_fermi=True if you did not specify efermi value')

        if emin is None:
            emin = np.min(self.eigenvalues) - 1
        if emax is None:
            emax = np.max(self.eigenvalues) + 1

        E_arr = np.arange(emin, emax, dE)
        ngrid = E_arr.shape[0]

        idxs = []
        for atom in atom_numbers:
            start = sum([self.norbs_per_atomtype[i] for i in self.species[:atom]])
            for i in range(self.norbs_per_atomtype[self.species[atom]]):
                idxs.append(start + i)

        if specific_l is not None:
            idxs = [idx for idx in idxs if self.orbs_data[idx]['l'] == specific_l]

        proj_coeffs_weighted = self.proj_coeffs[:, :, :, idxs]

        for spin in range(self.nspin):
            for i, weight_kpt in enumerate(self.weights):
                proj_coeffs_weighted[spin, i] *= weight_kpt

        W_arr = np.moveaxis(proj_coeffs_weighted, [1, 2, 3], [2, 3, 1])
        G_arr = EBS.gaussian_smearing(E_arr, self.eigenvalues, sigma)

        PDOS_arr = np.zeros((self.nspin, len(idxs), ngrid))
        for spin in range(self.nspin):
            for idx in range(len(idxs)):
               PDOS_arr[spin, idx] = np.sum(G_arr[spin, :, :, :] * W_arr[spin, idx, :, :, None],
                                            axis=(0, 1))

        if self.nspin == 1:
            PDOS_arr *= 2

        if get_orbs_names:
            if zero_at_fermi:
                return E_arr - efermi, PDOS_arr, [self.orbs_names[i] for i in idxs]
            else:
                return E_arr, PDOS_arr, [self.orbs_names[i] for i in idxs]
        else:
            if zero_at_fermi:
                return E_arr - efermi, PDOS_arr
            else:
                return E_arr, PDOS_arr


def VVSHE_2_mu_Ha(V):
    V_ref = 4.66
    return - (V_ref + V) * eV2Hartree


def mu_Ha_2_VVSHE(mu):
    V_ref = 4.66
    return - (mu * Hartree2eV + V_ref)