from functools import reduce
from copy import deepcopy
from typing import Optional
import warnings
import math
import re

import numpy as np
from .units import angstrom_to_bohr, ev_to_hartree
from .band import Band


class KPoint:
    def __init__(self, index, frac_coords, weight):
        self.index = index
        self.frac_coords = frac_coords
        self.weight = weight

    def cart_coords(self, reciprocal_lattice):
        """Convert the reciprocal fractional coordinates for this k-point to \
        reciprocal Cartesian coordinates.

        Args:
            reciprocal_lattice (np.array(float)): 3x3 numpy array containing the \
                Cartesian reciprocal lattice.

        Returns:
            np.array: The reciprocal Cartesian coordinates of this k-point.

        """
        return np.dot(self.frac_coords, reciprocal_lattice)

    def __eq__(self, other):
        return (
            (self.index == other.index)
            & (self.frac_coords == other.frac_coords).all()
            & (self.weight == other.weight)
        )

    def __repr__(self) -> str:
        return "k-point {}: {} weight = {}".format(
            self.index, " ".join([str(c) for c in self.frac_coords]), self.weight
        )


def get_numbers_from_string(string: str) -> list[float]:
    p = re.compile(r"-?\d+[.\d]*")
    return [float(s) for s in p.findall(string)]


def k_point_parser(string: str) -> list[KPoint]:
    """Parse k-point data from a PROCAR string.

    Finds all lines of the form::

         k-point    1 :    0.50000000 0.25000000 0.75000000     weight = 0.00806452

    and extracts the k-point index, reciprocal fractional coordinates, and weight
    into a :obj:`procar.KPoint` object.

    Args:
        string (str): String containing a full PROCAR file.

    Returns:
        list(:obj:`procar.KPoint`): A list of :obj:`procar.KPoint` objects.

    """
    regex = re.compile(
        r"k-point\s+(\d+)\s*:\s+([- ][01].\d{8})([- ][01].\d{8})([- ][01].\d{8})\s+weight = ([01].\d+)"
    )
    captured = regex.findall(string)
    k_points = []
    for kp in captured:
        index = int(kp[0])
        frac_coords = np.array([float(s) for s in kp[1:4]])
        weight = float(kp[4])
        k_points.append(KPoint(index=index, frac_coords=frac_coords, weight=weight))
    return k_points


def projections_parser(string):
    regex = re.compile(r"([-.\d\se]+tot.+)\n")
    data = regex.findall(string)
    data = [x.replace("tot", "0") for x in data]
    data = np.array([x.split() for x in data], dtype=float)
    return data


def area_of_a_triangle_in_cartesian_space(
    a: np.ndarray, b: np.ndarray, c: np.ndarray
) -> float:
    """Returns the area of a triangle defined by three points in Cartesian space.

    Args:
        a (np.array): Cartesian coordinates of point A.
        b (np.array): Cartesian coordinates of point B.
        c (np.array): Cartesian coordinates of point C.

    Returns:
        float: the area of the triangle.

    """
    return float(0.5 * np.linalg.norm(np.cross(b - a, c - a)))


def points_are_in_a_straight_line(points, tolerance=1e-7):
    """
    Check whether a set of points fall on a straight line.
    Calculates the areas of triangles formed by triplets of the points.
    Returns False is any of these areas are larger than the tolerance.

    Args:
        points (list(np.array)): list of Cartesian coordinates for each point.
        tolerance (optional:float): the maximum triangle size for these points to be considered colinear. Default is 1e-7.

    Returns:
        (bool): True if all points fall on a straight line (within the allowed tolerance).
    """
    a = points[0]
    b = points[1]
    for c in points[2:]:
        if area_of_a_triangle_in_cartesian_space(a, b, c) > tolerance:
            return False
    return True


def two_point_effective_mass(
    cartesian_k_points: np.ndarray,
    eigenvalues: np.ndarray
) -> float:
    """Calculate the effective mass given eigenvalues at two k-points.

    Reimplemented from Aron Walsh's original effective mass Fortran code.

    Args:
        cartesian_k_points (np.array): 2D numpy array containing the k-points in (reciprocal)
            Cartesian coordinates.
        eigenvalues (np.array):        numpy array containing the eigenvalues at each k-point.

    Returns:
        float: The effective mass.

    """
    assert cartesian_k_points.shape[0] == 2
    assert eigenvalues.size == 2
    dk = cartesian_k_points[1] - cartesian_k_points[0]
    mod_dk = np.sqrt(np.dot(dk, dk))
    delta_e = (eigenvalues[1] - eigenvalues[0]) * ev_to_hartree * 2.0
    effective_mass = mod_dk * mod_dk / delta_e
    return float(effective_mass)


def least_squares_effective_mass(
    cartesian_k_points: np.ndarray,
    eigenvalues: np.ndarray
) -> float:
    """Calculate the effective mass using a least squares quadratic fit.

    Args:
        cartesian_k_points (np.array): Cartesian reciprocal coordinates for the k-points.
        eigenvalues (np.array):        Energy eigenvalues at each k-point to be used in the fit.

    Returns:
        float: The fitted effective mass.

    Raises:
        ValueError: If the k-points do not sit on a straight line.

    """
    if not points_are_in_a_straight_line(cartesian_k_points):
        raise ValueError("k-points are not collinear")
    dk = cartesian_k_points - cartesian_k_points[0]
    mod_dk = np.linalg.norm(dk, axis=1)
    effective_mass = 1.0 / (np.polyfit(mod_dk, eigenvalues, 2)[0] * ev_to_hartree * 2.0)
    return float(effective_mass)


class Procar:
    """
    Object for working with PROCAR data.

    Attributes:
        data (numpy.array(float)): A 5D numpy array that stores the projection data.

                    Axes are k-points, bands, spin-channels, ions and sum over ions, lm-projections.

        bands (numpy.array(:obj:`Band`)): A numpy array of ``Band`` objects, that contain band index, energy, and occupancy data.
        k_points (numpy.array(:obj:`KPoint`)): A numpy array of ``KPoint`` objects, that contain fractional coordinates and weights for each k-point.
        number_of_k_points (int): The number of k-points.
        number_of_bands (int): The number of bands.
        spin_channels (int): Number of spin channels in the PROCAR data:

                    -  1 for non-spin-polarised calculations.
                    -  2 for spin-polarised calculations.
                    -  4 for non-collinear calculations.
        number_of_ions (int): The number of ions.
        number_of_projections (int): The number of projections, e.g. TODO
        calculation (dict): Dictionary of True | False values describing the calculation type.

            Dictionary keys are 'non_spin_polarised', 'non_collinear', and 'spin_polarised'.

    """

    def __init__(self, spin=1, negative_occupancies="warn"):
        self._spin_channels = spin  # should be determined from PROCAR
        self._number_of_k_points = None
        self._number_of_ions = None
        self._number_of_bands = None
        self._number_of_projections = None
        self._k_point_blocks = None
        self._data = None
        self._bands = None
        self._k_points = None
        self.calculation = {
            "non_spin_polarised": False,
            "non_collinear": False,
            "spin_polarised": False,
        }
        if negative_occupancies not in ["warn", "raise", "zero"]:
            raise ValueError(
                "negative_occupancies can be one of [ 'warn', 'raise', 'zero' ]"
            )
        self.negative_occupancies = negative_occupancies
        # self.non_spin_polarised = None

    @property
    def occupancy(self):
        return np.array([[band.index, band.occupancy] for band in self._bands])

    def __add__(self, other):
        if self.spin_channels != other.spin_channels:
            raise ValueError(
                "Can only concatenate Procars with equal spin_channels: {}, {}".format(
                    self.spin_channels, other.spin_channels
                )
            )
        if self.number_of_ions != other.number_of_ions:
            raise ValueError(
                "Can only concatenate Procars with equal number_of_ions: {}, {}".format(
                    self.number_of_ions, other.number_of_ions
                )
            )
        if self.number_of_bands != other.number_of_bands:
            raise ValueError(
                "Can only concatenate Procars with equal number_of_bands: {}, {}".format(
                    self.number_of_bands, other.number_of_bands
                )
            )
        if self.number_of_projections != other.number_of_projections:
            raise ValueError(
                "Can only concatenate Procars with equal number_of_projections: {}, {}".format(
                    self.number_of_projections, other.number_of_projections
                )
            )
        if self._k_point_blocks != other._k_point_blocks:
            raise ValueError(
                "Can only concatenate Procars with equal k_point_blocks: {}, {}".format(
                    self._k_point_blocks, other._k_point_blocks
                )
            )
        if self.calculation != other.calculation:
            raise ValueError(
                "Can only concatenate Procars from equal calculations: {}, {}".format(
                    self.calculation, other.calculation
                )
            )
        new_procar = deepcopy(self)
        new_procar._data = np.concatenate((self._data, other._data), axis=0)
        new_procar._number_of_k_points = (
            self.number_of_k_points + other.number_of_k_points
        )
        new_procar._bands = []
        new_procar._bands = np.ravel(np.concatenate([self.bands, other.bands], axis=1))
        new_procar._k_points = self._k_points + other._k_points
        for i, kp in enumerate(new_procar._k_points, 1):
            kp.index = i
        new_procar.sanity_check()
        return new_procar

    def parse_projections(self):
        self.projection_data = projections_parser(self.read_in)
        try:
            assert self._number_of_bands * self._number_of_k_points == len(
                self.projection_data
            )
            self._spin_channels = 1  # non-magnetic, non-spin-polarised
            self._k_point_blocks = 1
            self.calculation["non_spin_polarised"] = True
        except Exception:
            if self._number_of_bands * self._number_of_k_points * 4 == len(
                self.projection_data
            ):
                self._spin_channels = 4  # non-collinear (spin-orbit coupling)
                self._k_point_blocks = 1
                self.calculation["non_collinear"] = True
                pass
            elif self._number_of_bands * self._number_of_k_points * 2 == len(
                self.projection_data
            ):
                self._spin_channels = 2  # spin-polarised
                self._k_point_blocks = 2
                self.calculation["spin_polarised"] = True
                pass
            else:
                raise
        self._number_of_projections = (
            int(self.projection_data.shape[1] / (self._number_of_ions + 1)) - 1
        )

    def parse_k_points(self):
        self._k_points = k_point_parser(self.read_in)[: self._number_of_k_points]

    def parse_bands(self):
        band_data = re.findall(
            r"band\s*(\d+)\s*#\s*energy\s*([-.\d]+)\s?\s*#\s" r"*occ.\s*([-.\d]+)",
            self.read_in,
        )
        self._bands = np.array(
            [
                Band(
                    float(i),
                    float(e),
                    float(o),
                    negative_occupancies=self.negative_occupancies,
                )
                for i, e, o in band_data
            ]
        )

    def sanity_check(self):
        assert self._number_of_k_points == len(
            self._k_points
        ), "k-point number mismatch: {} in header; {} in file".format(
            self._number_of_k_points, len(self._k_points)
        )
        read_bands = len(self._bands) / self._number_of_k_points / self._k_point_blocks
        assert (
            self._number_of_bands == read_bands
        ), "band mismatch: {} in header; {} in file".format(
            self._number_of_bands, read_bands
        )

    @classmethod
    def from_files(cls, filenames, **kwargs):
        """Create a :obj:`Procar` object by reading the projected wavefunction character of each band
        from a series of VASP ``PROCAR`` files.

        Useful when e.g. a band-structure calculation has been split over multiple VASP calculations,
        for example, when using hybrid functionals.

        Args:
            filename (str): Filename of the ``PROCAR`` file.
            **kwargs: See the ``from_file()`` method for a description of keyword arguments.

        Returns:
            (:obj:`vasppy.Procar`)

        """
        pcars = [cls.from_file(f, **kwargs) for f in filenames]
        return reduce(cls.__add__, pcars)

    @classmethod
    def from_file(
        cls, filename, negative_occupancies="warn", select_zero_weighted_k_points=False
    ):
        """Create a :obj:`Procar` object by reading the projected wavefunction character of each band
        from a VASP ``PROCAR`` file.

        Args:
            filename (str): Filename of the ``PROCAR`` file.
            negative_occupancies (:obj:`Str`, optional): Select how negative occupancies are handled.
                Options are:

                    - ``warn`` (default): Warn that some partial occupancies are negative.
                    - ``raise``:          Raise an AttributeError.
                    - ``ignore``:         Do nothing.
                    - ``zero``:           Negative partial occupancies will be set to zero.
            select_zero_weighted_k_points (:obj:`bool`, optional): Set to ``True`` to only
                read in zero-weighted k-points from the ``PROCAR`` file. Default is ``False``.

        Returns:
            (:obj:`vasppy.Procar`)

        """
        pcar = cls(negative_occupancies=negative_occupancies)
        pcar._read_from_file(filename=filename)
        if select_zero_weighted_k_points:
            k_point_indices = [
                i for i, kp in enumerate(pcar.k_points) if kp.weight == 0.0
            ]
            pcar = pcar.select_k_points(k_point_indices)
        return pcar

    def read_from_file(self, filename):
        warnings.warn(
            "read_from_file() is deprecated as a part of the public API.\nPlease use Procar.from_file() or Procar.from_files() instead",
            stacklevel=2,
        )
        return self._read_from_file(filename=filename)

    def _read_from_file(self, filename):
        """Reads the projected wavefunction character of each band from a VASP PROCAR file.

        Args:
            filename (str): Filename of the PROCAR file.

        Returns:
            None

        """
        with open(filename, "r") as file_in:
            file_in.readline()
            self._number_of_k_points, self._number_of_bands, self._number_of_ions = [
                int(f) for f in get_numbers_from_string(file_in.readline())
            ]
            self.read_in = file_in.read()
        self.parse_k_points()
        self.parse_bands()
        self.parse_projections()
        self.sanity_check()
        self.read_in = None  # clear memory
        if self.calculation["spin_polarised"]:
            self._data = (
                self.projection_data.reshape(
                    (
                        self._spin_channels,
                        self._number_of_k_points,
                        self._number_of_bands,
                        self._number_of_ions + 1,
                        self._number_of_projections + 1,
                    )
                )[:, :, :, :, 1:]
                .swapaxes(0, 1)
                .swapaxes(1, 2)
            )
        else:
            self._data = self.projection_data.reshape(
                (
                    self._number_of_k_points,
                    self._number_of_bands,
                    self._spin_channels,
                    self._number_of_ions + 1,
                    self._number_of_projections + 1,
                )
            )[:, :, :, :, 1:]

    @property
    def number_of_k_points(self):
        """The number of k-points described by this :obj:`Procar` object."""
        assert (
            self._number_of_k_points == self._data.shape[0]
        ), "Number of k-points in metadata ({}) not equal to number in PROCAR data ({})".format(
            self._number_of_k_points, self._data.shape[0]
        )
        return self._number_of_k_points

    @property
    def number_of_bands(self):
        """The number of bands described by this :obj:`Procar` object."""
        assert (
            self._number_of_bands == self._data.shape[1]
        ), "Number of bands in metadata ({}) not equal to number in PROCAR data ({})".format(
            self._number_of_bands, self._data.shape[1]
        )
        return self._number_of_bands

    @property
    def spin_channels(self):
        """The number of spin-channels described by this :obj:`Procar` object."""
        assert (
            self._spin_channels == self._data.shape[2]
        ), "Number of spin channels in metadata ({}) not equal to number in PROCAR data ({})".format(
            self._spin_channels, self._data.shape[2]
        )
        return self._spin_channels

    @property
    def number_of_ions(self):
        """The number of ions described by thie :obj:`Procar` object."""
        assert (
            self._number_of_ions == self._data.shape[3] - 1
        ), "Number of ions in metadata ({}) not equal to number in PROCAR data ({})".format(
            self._number_of_ions, self._data.shape[3] - 1
        )
        return self._number_of_ions

    @property
    def number_of_projections(self):
        """The number of lm-projections described by this :obj:`Procar` object."""
        assert (
            self._number_of_projections == self._data.shape[4]
        ), "Number of projections in metadata ({}) not equal to number in PROCAR data ({})".format(
            self._number_of_projections, self._data.shape[4]
        )
        return self._number_of_projections

    def print_weighted_band_structure(
        self,
        spins=None,
        ions=None,
        orbitals=None,
        scaling=1.0,
        e_fermi=0.0,
        reciprocal_lattice=None,
    ):
        band_structure_data = self.weighted_band_structure(
            spins=spins,
            ions=ions,
            orbitals=orbitals,
            scaling=scaling,
            e_fermi=e_fermi,
            reciprocal_lattice=reciprocal_lattice,
        )
        for i, band_data in enumerate(band_structure_data, 1):
            print("# band: {}".format(i))
            for k_point_data in band_data:
                print(" ".join([str(f) for f in k_point_data]))
            print()

    def weighted_band_structure(
        self,
        spins=None,
        ions=None,
        orbitals=None,
        scaling=1.0,
        e_fermi=0.0,
        reciprocal_lattice=None,
    ):
        if spins:
            spins = [s - 1 for s in spins]
        else:
            spins = list(range(self.spin_channels))
        if not ions:
            ions = [self.number_of_ions]  # nions+1 is the `tot` index
        if not orbitals:
            orbitals = list(range(self.number_of_projections))
        if self.calculation["spin_polarised"]:
            band_energies = (
                np.array([band.energy for band in self._bands])
                .reshape(
                    (self.spin_channels, self.number_of_k_points, self.number_of_bands)
                )[spins[0]]
                .T
            )
        else:
            band_energies = (
                np.array([band.energy for band in self._bands])
                .reshape((self.number_of_k_points, self.number_of_bands))
                .T
            )
        orbital_projection = np.sum(self._data[:, :, :, :, orbitals], axis=4)
        ion_projection = np.sum(orbital_projection[:, :, :, ions], axis=3)
        spin_projection = np.sum(ion_projection[:, :, spins], axis=2)
        x_axis = self.x_axis(reciprocal_lattice)
        to_return = []
        for i in range(self.number_of_bands):
            for k, (e, p) in enumerate(
                zip(band_energies[i], spin_projection.T[i], strict=None)
            ):
                to_return.append([x_axis[k], e - e_fermi, p * scaling])
        to_return = np.array(to_return).reshape((self.number_of_bands, -1, 3))
        return to_return

    def effective_mass_calc(
        self, k_point_indices, band_index, reciprocal_lattice, spin=1, printing=False
    ):
        assert spin <= self._k_point_blocks
        assert len(k_point_indices) > 1  # we need at least 2 k-points
        band_energies = self._bands[:, 1:].reshape(
            self._k_point_blocks, self.number_of_k_points, self.number_of_bands
        )
        frac_k_point_coords = np.array(
            [self._k_points[k - 1].frac_coords for k in k_point_indices]
        )
        eigenvalues = np.array(
            [band_energies[spin - 1][k - 1][band_index - 1] for k in k_point_indices]
        )
        if printing:
            print("# h k l e")
            for row in np.concatenate(
                (frac_k_point_coords, np.array([eigenvalues]).T), axis=1
            ):
                print(" ".join([str(f) for f in row]))
        reciprocal_lattice = reciprocal_lattice * 2 * math.pi * angstrom_to_bohr
        cart_k_point_coords = np.array(
            [k.cart_coords(reciprocal_lattice) for k in self._k_points]
        )  # convert k-points to cartesian
        if len(k_point_indices) == 2:
            effective_mass_function = two_point_effective_mass
        else:
            effective_mass_function = least_squares_effective_mass
        return effective_mass_function(cart_k_point_coords, eigenvalues)

    def x_axis(self, reciprocal_lattice: Optional[np.ndarray] = None) -> np.ndarray:
        """Generate the x-axis values for a band-structure plot.

        Returns an array of cumulative distances in reciprocal space between sequential k-points.

        Args:
            reciprocal_lattice (:obj:`np.array`, optional): 3x3 Cartesian reciprocal lattice.
                Default is ``None``. If no reciprocal lattice is provided, the returned x-axis
                values will be sequential integers, giving even spacings between sequential
                k-points.

        Returns:
            np.array: An array of x-axis values.

        """
        if reciprocal_lattice is not None:
            cartesian_k_points = np.array(
                [k.cart_coords(reciprocal_lattice) for k in self._k_points]
            )
            x_axis = [0.0]
            for i in range(1, len(cartesian_k_points)):
                dk = cartesian_k_points[i - 1] - cartesian_k_points[i]
                mod_dk = np.sqrt(np.dot(dk, dk))
                x_axis.append(mod_dk + x_axis[-1])
            x_axis_array = np.array(x_axis)
        else:
            x_axis_array = np.arange(len(self._k_points))
        return x_axis_array

    @property
    def bands(self):
        return self._bands.reshape(
            self._k_point_blocks, self._number_of_k_points, self.number_of_bands
        )

    @property
    def k_points(self):
        return self._k_points

    def select_bands_by_kpoint(self, band_indices):
        return np.ravel(self.bands[:, band_indices, :])

    def select_k_points(self, band_indices):
        new_procar = deepcopy(self)
        new_procar._bands = np.ravel(new_procar.bands[:, band_indices, :])
        new_procar._data = np.array(
            [kp for i, kp in enumerate(new_procar._data) if i in band_indices]
        )
        new_procar._number_of_k_points = len(band_indices)
        new_procar._k_points = [
            kp for i, kp in enumerate(new_procar._k_points) if i in band_indices
        ]
        for i, kp in enumerate(new_procar._k_points, 1):
            kp.index = i
        new_procar.sanity_check()
        return new_procar


# TODO Need complete set of example PROCAR files before we start to write
# something that can write these all back out in the appropriate format
#    def write_file( self, filename ):
#        if self.calculation['non_collinear']:
#            raise NotImplementedError
#        with open( filename, 'w' ) as f:
#            f.write( 'PROCAR lm decomposed' )
#            for s in range( self.spin_channels ): # not sure what happens for non-collinear calculations
#                line = ff.FortranRecordWriter("/'# of k-points:',I5,9X,'# of bands:',I5,9X,'# of ions:',I5")
#                f.write( line.write( [ self.number_of_k_points, self.number_of_bands, self.number_of_ions ] )+'\n' )
#                for k_point in self.k_points:
#                    line = ff.FortranRecordWriter("/' k-point ',I5,' :',3X,3F11.8,'     weight = ',F10.8/")
#                    f.write( line.write( [ k_point.index, *k_point.frac_coords, k_point.weight ] ) )
#                    for band in self.bands:
#                        line = ff.FortranRecordWriter("/'band ',I5,' # energy',F14.8,' # occ.',F12.8/")
#                        f.write( line.write( [ band.index, band.energy, band.occupancy ] ) )
#                        f.write( '\nion      s     py     pz     px    dxy    dyz    dz2    dxz    dx2    tot\n' )
