"""
Procar reader
"""
from typing import List, Union
from pathlib import Path
import re
import numpy as np

from monty.json import MSONable, MontyDecoder

from easyunfold import __version__

# pylint:disable=too-many-locals,


class Procar(MSONable):
    """Reader for PROCAR file"""

    def __init__(self, fobjs_or_paths=None, is_soc=False, normalise=True):
        """
        Read the PROCAR file from a handle or path

        :param fobjs_or_paths:  Either a string or list of file-like objs or paths
        :param is_soc: Whether the PROCAR is from a calculation with spin-orbit coupling
        :param normalise: Whether to normalise the projection for every band or not
        """
        self._is_soc = is_soc
        self.eigenvalues = None
        self.kvecs = None
        self.kweights = None
        self.nbands = None
        self.nkpts = None
        self.nspins = None
        self.nion = None
        self.occs = None
        self.proj_names = None
        self.proj_data = None
        self.header = None
        self.proj_xyz = None
        self.normalise = normalise

        # Read the PROCAR
        if isinstance(fobjs_or_paths, (str, Path)):
            fobjs_or_paths = [fobjs_or_paths]
        self.read(fobjs_or_paths)

    def _read(self, fobj, parsed_kpoints=None):
        """Main function for reading in the data"""
        if parsed_kpoints is None:
            parsed_kpoints = set()

        # First sweep - find the number of kpoints and the number of bands
        fobj.seek(0)
        _header = fobj.readline()
        # Read the NK, NB and NIONS that are integers
        _total_nkpts, nbands, nion = [int(token) for token in re.sub(r'[^0-9]', ' ', fobj.readline()).split()]
        if nion != self.nion:
            raise ValueError(f'Mismatch in number of ions in PROCARs supplied: ({nion} vs {self.nion})!')

        # Count the number of data lines, these lines do not have any alphabets
        proj_data, energies, kvecs, kweights, occs = [], [], [], [], []
        tot_count = 0  # count the instances of lines starting with "tot" -> (4 + 1) * nbands * nkpts for SOC calcs
        fobj.seek(0)

        line = fobj.readline()
        # Counter for the number of sections in the PROCAR
        section_counter = 1
        _last_kid = 0
        this_procar_parsed_kpoints = set()  # set with tuples of parsed (kvec tuple, section_counter) for this PROCAR
        while line:
            if line.startswith(' k-point'):
                line = re.sub(r'(\d)-', r'\1 -', line)
                tokens = line.strip().split()
                _kid = int(tokens[1])

                # Check if the VASP PROCAR k index decreases – if so we have entered a second section
                if _kid < _last_kid:
                    section_counter += 1
                _last_kid = _kid

                kvec = tuple(round(float(val), 5) for val in tokens[-6:-3]  # tuple to make it hashable
                            )  # round to 5 decimal places to ensure proper kpoint matching
                if (kvec not in parsed_kpoints and (kvec, section_counter) not in this_procar_parsed_kpoints):
                    this_procar_parsed_kpoints.add((kvec, section_counter))
                    kvecs.append(list(kvec))
                    kweights.append(float(tokens[-1]))
                else:
                    # skip ahead to the next instance of two blank lines in a row
                    while line.strip() or fobj.readline().strip():
                        line = fobj.readline()
                    continue

            elif (not re.search(r'[a-zA-Z]', line) and line.strip() and len(line.strip().split()) - 2 == len(self.proj_names)):
                # only parse data if line is expected length, in case of LORBIT >= 12
                proj_data.append([float(token) for token in line.strip().split()[1:-1]])

            elif line.startswith('band'):
                tokens = line.strip().split()
                energies.append(float(tokens[4]))
                occs.append(float(tokens[-1]))

            elif line.startswith('tot'):
                tot_count += 1

            line = fobj.readline()

        # dynamically determine whether PROCARs are SOC or not
        if tot_count == 4 * len(occs):
            self._is_soc = True
        elif tot_count == len(occs):
            self._is_soc = False
        else:
            raise ValueError(f"Number of lines starting with 'tot' ({tot_count}) in PROCAR does not match expected "
                             f'values ({4*len(occs)} or {len(occs)})!')

        occs = np.array(occs)
        kvecs = np.array(kvecs)
        kweights = np.array(kweights)
        eigenvalues = np.array(energies)

        proj_data = np.array(proj_data, dtype=float)

        nkpts = len(kvecs)  # redetermine nkpts in case some were skipped due to already being parsed

        # For spin-polarised calcs, the data from the second (down) spin are located after that of the first (up) spin
        # Hence, the number of spins is simply the number of sections
        self.nspins = section_counter

        # Reshape
        occs.resize((self.nspins, nkpts // self.nspins, nbands))
        kvecs.resize((self.nspins, nkpts // self.nspins, 3))
        kweights.resize((self.nspins, nkpts // self.nspins))
        eigenvalues.resize((self.nspins, nkpts // self.nspins, nbands))

        # Reshape the array
        if self._is_soc is False:
            proj_data = proj_data.reshape((
                self.nspins,
                nkpts // self.nspins,
                nbands,
                self.nion,
                len(self.proj_names),
            ))
            proj_xyz = None
        else:
            proj_data = proj_data.reshape((self.nspins, nkpts, nbands, 4, self.nion, len(self.proj_names)))
            # Split the data into xyz projection and total
            proj_xyz = proj_data[:, :, :, 1:, :, :]
            proj_data = proj_data[:, :, :, 0, :, :]

        if self.normalise:
            self.normalise_projs(proj_data)

        if proj_xyz is not None:
            proj_sum = np.sum(proj_xyz, axis=(-3, -2, -1), keepdims=True)
            proj_sum[proj_sum == 0] = 1
            proj_xyz /= proj_sum

        # Update the parsed kpoints
        parsed_kpoints.update({kvec_section_counter_tuple[0] for kvec_section_counter_tuple in this_procar_parsed_kpoints})

        return (
            self.nspins,
            occs,
            kvecs,
            kweights,
            eigenvalues,
            proj_data,
            proj_xyz,
            parsed_kpoints,
        )

    def normalise_projs(self, proj_data):
        """
        Normalise the projections

        For each nspin, nkpt, nband, normalise the sum of projections over nion and proj_names to be 1.
        Atomic & orbital projections do not sum to 1 in most cases in VASP, as only those falling inside
        the atomic radii and overlapping with spd spherical harmonics are counted.
        """
        proj_sum = np.sum(proj_data, axis=(-2, -1), keepdims=True)
        proj_sum[proj_sum == 0] = 1  # just in case, avoid division by zero
        proj_data /= proj_sum
        self.normalise = True

    def _read_header_nion_proj_names(self, fobj):
        """Read the header, nion and proj_names from the PROCAR"""
        fobj.seek(0)
        self.header = fobj.readline()
        # Read the NK, NB and NIONS that are integers
        _nkpts, _nbands, self.nion = [int(token) for token in re.sub(r'[^0-9]', ' ', fobj.readline()).split()]
        self.proj_names = None  # projection names

        for line in fobj:
            if re.match(r'^ion.*tot', line):  # only the first "ion" line, in case of LORBIT >= 12
                self.proj_names = line.strip().split()[1:-1]
                break

    def read(self, fobjs_or_paths):
        """Read and amalgamate the data from a list of PROCARs"""

        def open_file(fobj_or_path):
            if isinstance(fobj_or_path, (str, Path)):
                return open(fobj_or_path, encoding='utf-8')  # closed later
            return fobj_or_path  # already a file-like object, just return it

        parsed_kpoints = None
        occs_list, kvecs_list, kweights_list = [], [], []
        eigenvalues_list, proj_data_list, proj_xyz_list = [], [], []
        for i, fobj_or_path in enumerate(fobjs_or_paths):
            # Note: If PROCAR parsing becomes a significant bottleneck for people (e.g. with several HSE06+SOC PROCARs),
            # this could be parallelized (somewhat) with multiprocessing. The actual file parsing in _read() is currently
            # serial so no easy wins there, but could at least parallelise over the list of PROCARs
            fobj = open_file(fobj_or_path)
            if self.header is None:  # first file; read header, nion, proj_names
                self._read_header_nion_proj_names(fobj)

            current_nspins = self.nspins  # check spin consistency between PROCARs
            (
                nspins,
                occs,
                kvecs,
                kweights,
                eigenvalues,
                proj_data,
                proj_xyz,
                parsed_kpoints,
            ) = self._read(fobj, parsed_kpoints=parsed_kpoints)
            if current_nspins is not None and current_nspins != nspins:
                raise ValueError(f'Mismatch in number of spins in PROCARs supplied: ({nspins} vs {current_nspins})!')

            if isinstance(fobj_or_path, (str, Path)):
                fobj.close()  # if file was opened in this loop, close it

            # Append to respective lists
            occs_list.append(occs)
            kvecs_list.append(kvecs)
            kweights_list.append(kweights)
            eigenvalues_list.append(eigenvalues)
            proj_data_list.append(proj_data)
            proj_xyz_list.append(proj_xyz)
            if len(fobjs_or_paths) > 1:  # print progress if reading multiple files
                print(f'Finished parsing PROCAR {i + 1}/{len(fobjs_or_paths)}')

        # Combine along the nkpts axis:
        # for occs, eigenvalues, proj_data and proj_xyz, nbands (axis = 2) could differ, so set missing values to zero:
        max_nbands = max(arr.shape[2] for arr in eigenvalues_list)
        for array_list in [occs_list, eigenvalues_list, proj_data_list, proj_xyz_list]:
            for i, arr in enumerate(array_list):
                if arr is not None and arr.shape[2] < max_nbands:
                    if len(arr.shape) == 3:  # occs_list, eigenvalues_list
                        array_list[i] = np.pad(
                            arr,
                            ((0, 0), (0, 0), (0, max_nbands - arr.shape[2])),
                            mode='constant',
                        )
                    elif len(arr.shape) == 5:  # proj_xyz_list
                        array_list[i] = np.pad(
                            arr,
                            (
                                (0, 0),
                                (0, 0),
                                (0, max_nbands - arr.shape[2]),
                                (0, 0),
                                (0, 0),
                            ),
                            mode='constant',
                        )
                    elif len(arr.shape) == 6:  # proj_xyz_list
                        array_list[i] = np.pad(
                            arr,
                            (
                                (0, 0),
                                (0, 0),
                                (0, max_nbands - arr.shape[2]),
                                (0, 0),
                                (0, 0),
                                (0, 0),
                            ),
                            mode='constant',
                        )
                    else:
                        raise ValueError('Unexpected array shape encountered!')

        self.nbands = max_nbands
        self.occs = np.concatenate(occs_list, axis=1)
        self.eigenvalues = np.concatenate(eigenvalues_list, axis=1)
        self.kvecs = np.concatenate(kvecs_list, axis=1)
        self.kweights = np.concatenate(kweights_list, axis=1)
        self.proj_data = np.concatenate(proj_data_list, axis=1)
        if all(arr is not None for arr in proj_xyz_list):
            self.proj_xyz = np.concatenate(proj_xyz_list, axis=1)

        self.nkpts = self.kvecs.shape[1]

    def get_projection(self, atom_idx: List[int], proj: Union[List[str], str], weight_by_k=False):
        """
        Get projection for specific atoms and specific projectors


        :param atom_idx: A list of index of the atoms to be selected
        :param proj: A list of the projector names to be selected
        :param weight_by_k: Apply k weighting or not.

        :returns: The projection summed over the selected atoms and the projectors
        """
        atom_mask = [iatom in atom_idx for iatom in range(self.nion)]
        assert any(atom_mask)
        if proj == 'all':
            out = self.proj_data[:, :, :, atom_mask, :].sum(axis=(-1, -2))
        else:
            if isinstance(proj, str):
                proj = [
                    proj,
                ]

            # replace any instance of "p" with "px,py,pz" and "d" with "dxy,dyz,dz2,dxz,dx2-y2"
            def _replace_p_d(single_proj):
                if single_proj == 'p':
                    return ['px', 'py', 'pz']
                if single_proj == 'd':
                    return [
                        'dxy',
                        'dyz',
                        'dz2',
                        'dxz',
                        'x2-y2',
                    ]  # dx2-y2 labelled differently in VASP
                    # PROCAR

                return [single_proj]

            proj = [_replace_p_d(single_proj) for single_proj in proj]
            proj = [item for sublist in proj for item in sublist]

            proj_idx = [proj_name in proj for proj_name in self.proj_names]
            assert any(proj_idx)
            out = self.proj_data[:, :, :, :, proj_idx]
            out = out[:, :, :, atom_idx, :].sum(axis=(-1, -2))

        if weight_by_k:
            for kidx in range(self.nkpts):
                out[:, kidx, :] *= self.kweights[:, kidx]
        return out

    def as_dict(self) -> dict:
        """Convert the object into a dictionary representation (so it can be saved to json)"""
        output = {
            '@module': self.__class__.__module__,
            '@class': self.__class__.__name__,
            '@version': __version__,
        }
        for key in [
                '_is_soc',
                'eigenvalues',
                'kvecs',
                'kweights',
                'nbands',
                'nkpts',
                'nspins',
                'nion',
                'occs',
                'proj_names',
                'proj_data',
                'header',
                'proj_xyz',
                'normalise',
        ]:
            output[key] = getattr(self, key)
        return output

    @classmethod
    def from_dict(cls, d):
        """
        Reconstructs Procar object from a dict representation, without calling __init__().

        Args:
            d (dict): dict representation of Procar

        Returns:
            Procar object
        """

        def decode_dict(subdict):
            if isinstance(subdict, dict) and '@module' in subdict:
                return MontyDecoder().process_decoded(subdict)
            return subdict

        instance = cls.__new__(cls)  # create a new instance without calling __init__()
        d_decoded = {k: decode_dict(v) for k, v in d.items()}

        # set the instance variables directly from the dictionary
        for key, value in d_decoded.items():
            if key in ['@module', '@class', '@version']:
                continue
            setattr(instance, key, value)

        return instance
