import re
import numpy as np
from pymatgen.io.vasp import Outcar

from express.parsers.apps.vasp import settings
from express.parsers.formats.txt import BaseTXTParser


class VaspTXTParser(BaseTXTParser):
    """
    Vasp text parser class.
    """

    def __init__(self, work_dir):
        super(VaspTXTParser, self).__init__(work_dir)

    def ibz_kpoints(self, text, space):
        """
        Extracts kpoints coordinates in cartesian space.

        Args:
            text (str): text to extract data from.
            space (str): kpoints coordinate space, either crystal or cartesian.

        Returns:
            ndarray
        """
        text_range = {
            "cartesian": {"start": "Following cartesian coordinates", "end": "Dimension of arrays"},
            "crystal": {"start": "Following reciprocal coordinates", "end": "Following cartesian coordinates"},
        }
        start_index = text.find(text_range[space]["start"])
        end_index = text.find(text_range[space]["end"])
        ibz_kpts = re.findall(settings.REGEX["ibz_kpoints"]["regex"], text[start_index:end_index])
        ibz_kpts = [[float(x) for x in kp] for kp in ibz_kpts]
        return np.array(ibz_kpts)

    def total_energy(self, text):
        """
        Extracts total energy.

        Args:
            text (str): text to extract data from.

        Returns:
             float
        """
        return self._general_output_parser(text, **settings.REGEX["total_energy"])

    def _lattice_convergence(self, text):
        """
        Extracts lattice convergence values computed in each BFGS step.

        Args:
            text (str): text to extract data from.

        Returns:
            list

        Example:
            [
                {
                    'vectors': {
                        'a': [-0.561154473, -0.000000000, 0.561154473],
                        'b': [-0.000000000, 0.561154473, 0.561154473],
                        'c': [-0.561154473, 0.561154473, 0.000000000],
                        'alat': 9.44858082
                    }
                 },
                ...
                {
                    'vectors': {
                        'a': [-0.561154473, -0.000000000, 0.561154473],
                        'b': [-0.000000000, 0.561154473, 0.561154473],
                        'c': [-0.561154473, 0.561154473, 0.000000000],
                        'alat': 9.44858082
                    }
                 }
             ]
        """

        lattices = []
        match = re.findall(settings.REGEX["lattice_vectors"]["regex"], text)
        if match:
            for lattice in match:
                lattice = [float(_) for _ in lattice]
                lattices.append(
                    {
                        "vectors": {
                            "a": lattice[0:3],
                            "b": lattice[3:6],
                            "c": lattice[6:9],
                            "alat": 1.0,  # abc vectors are expected in absolute units (eg. bohr)
                        }
                    }
                )
        return lattices

    def _basis_convergence(self, text, atom_names):
        """
        Extracts convergence bases computed in each BFGS step.

        Args:
            text (str): text to extract data from.
            atom_names (list): list of atom names.

        Returns:
            list

        Example:
            [
                {
                    'units': 'crystal',
                    'elements': [{'id': 1, 'value': 'Si'}, {'id': 2, 'value': 'Si'}],
                    'coordinates': [{'id': 1, 'value': [0.0, 0.0, 0.0]}, {'id': 2, 'value': [0.0, 0.0, 0.0]}]
                 },
                ...
                {
                    'units': 'crystal',
                    'elements': [{'id': 1, 'value': 'Si'}, {'id': 2, 'value': 'Si'}],
                    'coordinates': [{'id': 1, 'value': [0.0, 0.0, 0.0]}, {'id': 2, 'value': [0.0, 0.0, 0.0]}]
                 }
            ]
        """
        results = []
        matches = re.findall(settings.REGEX["ion_positions_block"]["regex"], text, re.DOTALL | re.MULTILINE)
        if matches:
            for match in matches:
                ions = re.findall(settings.REGEX["basis_vectors"]["regex"], match)
                results.append(
                    {
                        "units": "angstrom",
                        "elements": [{"id": idx, "value": atom_names[idx]} for idx in range(len(ions))],
                        "coordinates": [{"id": idx, "value": [float(x) for x in ion]} for idx, ion in enumerate(ions)],
                    }
                )
            return results

    def convergence_electronic(self, outcar, stdout, atom_names):
        """
        Extracts convergence electronic.
            1. Extract all energies (from dE column) along with corresponding step [(1, 0.69948E+04), (2, -0.73973E+04)]
            2. Group the energies for each ionic step

        Sample input:
                   N       E                     dE             d eps       ncg     rms          rms(c)
            DAV:   1     0.699475439520E+04    0.69948E+04   -0.37054E+05   920   0.140E+03
            DAV:   2    -0.402547717878E+03   -0.73973E+04   -0.71440E+04  1084   0.442E+02
            DAV:   3    -0.105567783952E+04   -0.65313E+03   -0.64960E+03  1160   0.125E+02
            DAV:   4    -0.110887818556E+04   -0.53200E+02   -0.53162E+02  1560   0.284E+01
            DAV:   5    -0.111382548259E+04   -0.49473E+01   -0.49469E+01  1240   0.562E+00    0.137E+02
                1 T=  2001. E= -.89713765E+03 F= -.92170983E+03 E0= -.92172095E+03  EK= 0.24572E+02

                   N       E                     dE             d eps       ncg     rms          rms(c)
            DAV:   1    -0.916206954626E+03    0.55030E+01   -0.19221E+03   920   0.750E+01    0.931E+00
            DAV:   2    -0.920149095159E+03   -0.39421E+01   -0.47144E+01  1072   0.109E+01    0.625E+00
            DAV:   3    -0.919835295300E+03    0.31380E+00   -0.27646E+00  1288   0.256E+00    0.367E+00
            DAV:   4    -0.919705445969E+03    0.12985E+00   -0.93950E-01  1128   0.186E+00    0.145E+00
                2 T=  1843. E= -.89706481E+03 F= -.91969309E+03 E0= -.91970717E+03  EK= 0.22628E+02

        Args:
            outcar (str): OUTCAR content.
            stdout (str): stdout content.
            atom_names (list): list of atoms.

        Returns:
             list[list]
        """
        data = []
        matches = self._general_output_parser(stdout, **settings.REGEX["convergence_electronic"])
        first_step_indices = [ind for ind, elm in enumerate(matches) if int(elm[0]) == 1]
        for ind, first_step_index in enumerate(first_step_indices):
            if ind + 1 == len(first_step_indices):
                energies = matches[first_step_index:]
            else:
                energies = matches[first_step_index : first_step_indices[ind + 1]]
            data.append([energy[1] for energy in energies])  # strip out the step numbers
        return data

    def convergence_ionic(self, outcar, stdout, atom_names):
        """
        Extracts convergence ionic.

        Args:
            outcar (str): OUTCAR content.
            stdout (str): stdout content.
            atom_names (list): list of atoms.

        Returns:
             list[dict]
        """
        data = []
        convergence_electronic_data = self.convergence_electronic(outcar, stdout, atom_names)
        for ind, energies in enumerate(convergence_electronic_data):
            data.append(
                {
                    "energy": sum([sum(e) for e in convergence_electronic_data[0 : ind + 1]]),
                    "electronic": {"units": "eV", "data": energies},
                }
            )

        if not data:
            return []

        lattice_convergence = self._lattice_convergence(outcar)
        basis_convergence = self._basis_convergence(outcar, atom_names)
        data = data[0 : len(lattice_convergence)]  # strip out the last non-complete step
        if not data:
            return []
        for idx, structure in enumerate(zip(lattice_convergence, basis_convergence)):
            data[idx].update({"structure": {"lattice": structure[0], "basis": structure[1]}})

        return data

    def pressure(self, text):
        """
        Extracts pressure.

        Args:
            text (str): text to extract data from.

        Returns:
            float
        """
        return self._general_output_parser(text, **settings.REGEX["pressure"])

    def total_force(self, text):
        """
        Extracts total force.

        Args:
            text (str): text to extract data from.

        Returns:
            float
        """
        total_force = self._general_output_parser(text, **settings.REGEX["total_force"])
        return np.sqrt(np.sum(np.square(total_force)))

    def total_energy_contributions(self, text):
        """
        Extracts total energy contributions.

        Args:
            text (str): text to extract data from.

        Returns:
            dict
        """
        energy_contributions = {}
        for contribution in settings.TOTAL_ENERGY_CONTRIBUTIONS:
            value = self._general_output_parser(text, **settings.TOTAL_ENERGY_CONTRIBUTIONS[contribution])
            if value:
                energy_contributions.update(
                    {contribution: {"name": contribution, "value": np.sqrt(np.sum(np.square(value)))}}
                )
        return energy_contributions

    def zero_point_energy(self, text):
        """
        Extracts zero point energy.

        Returns:
             float
        """
        data = self._general_output_parser(text, **settings.REGEX["zero_point_energy"])
        if len(data):
            return sum(data) / 2 / 1000

    def magnetic_moments(self, outcar):
        """
        Extracts magnetic moments.

        Returns:
             list
        """
        mag = Outcar(outcar).magnetization
        return [[0, 0, ion["tot"]] if isinstance(ion["tot"], float) else ion["tot"].moment.tolist() for ion in mag]
