from __future__ import annotations
import numpy as np
from typing import Iterable
from termcolor import colored
from nptyping import NDArray, Shape, Number


class Structure:
    """
    Basic class for structure of unit/supercell.
    Args:
        lattice: 2D array that contains lattice vectors. Each row should correspond to a lattice vector.
            E.g., [[5, 5, 0], [7, 4, 0], [0, 0, 25]].
        species: List of species on each site. Usually list of elements, e.g., ['Al', 'Al', 'O', 'H'].
        coords: List of lists or np.ndarray (Nx3 dimension) that contains coords of each species.
        coords_are_cartesian: True if coords are cartesian, False if coords are fractional
    """
    def __init__(self,
                 lattice: NDArray[Shape[3, 3], Number] | list[list[float]],
                 species: list[str],
                 coords: NDArray[Shape['Natoms, 3'], Number] | list[list[float]],
                 coords_are_cartesian: bool = True):

        if len(species) != len(coords):
            raise StructureError('Number of species and its coords must be the same')

        self.species = species

        if isinstance(lattice, np.ndarray):
            self.lattice = lattice
        else:
            self.lattice = np.array(lattice)

        if isinstance(coords, np.ndarray):
            self.coords = coords
        else:
            self.coords = np.array(coords)

        self.coords_are_cartesian = coords_are_cartesian

    def __repr__(self) -> str:
        lines = ['\nLattice:']

        width = len(str(int(np.max(self.lattice)))) + 6
        for axis in self.lattice:
            lines.append('  '.join([f'{axis_coord:{width}.5f}' for axis_coord in axis]))

        width = len(str(int(np.max(self.coords)))) + 6

        lines.append('\nSpecies:')

        unique, counts = np.unique(self.species, return_counts=True)

        if len(self.species) < 10:
            lines.append(' '.join([s for s in self.species]))
            for u, c in zip(unique, counts):
                lines.append(f'{u}:\t{c}')
            lines.append(f'\nCoords are cartesian: {self.coords_are_cartesian}')
            lines.append('\nCoords:')
            for coord in self.coords:
                lines.append('  '.join([f'{c:{width}.5f}' for c in coord]))
        else:
            part_1 = ' '.join([s for s in self.species[:5]])
            part_2 = ' '.join([s for s in self.species[-5:]])
            lines.append(part_1 + ' ... ' + part_2)
            for u, c in zip(unique, counts):
                lines.append(f'{u}:\t{c}')
            lines.append(f'\nCoords are cartesian: {self.coords_are_cartesian}')
            lines.append('\nCoords:')
            for coord in self.coords[:5]:
                lines.append('  '.join([f'{c:{width}.5f}' for c in coord]))
            lines.append('...')
            for coord in self.coords[-5:]:
                lines.append('  '.join([f'{c:{width}.5f}' for c in coord]))

        return '\n'.join(lines)

    def __eq__(self, other: Structure) -> bool:
        assert isinstance(other, Structure), 'Other object be a Structure class'
        if not self.coords_are_cartesian == other.coords_are_cartesian:
            if self.coords_are_cartesian:
                other.mod_coords_to_cartesian()
                print(colored('Coords of other were modified into Cartesian', color='green'))
            else:
                other.mod_coords_to_direct()
                print(colored('Coords of other were modified into Direct', color='green'))
        return np.allclose(self.lattice, other.lattice, atol=1e-10, rtol=1e-10) and \
            (self.species == other.species) and \
            np.allclose(self.coords, other.coords, atol=1e-10, rtol=1e-10)

    @property
    def natoms(self) -> int:
        return len(self.species)

    @property
    def natoms_by_type(self) -> dict[str: int]:
        unique, counts = np.unique(self.species, return_counts=True)
        return {i: j for i, j in zip(unique, counts)}

    def mod_add_atoms(self, coords, species) -> None:
        """
        Adds atoms in the Structure
        Args:
            coords: List or np.ndarray (Nx3 dimension) that contains coords of each species
            species: List of species on each site. Usually list of elements, e.g., ['Al', 'Al', 'O', 'H']
        """
        if not isinstance(coords, np.ndarray):
            coords = np.array(coords)
        self.coords = np.vstack((self.coords, coords))
        if isinstance(species, str):
            self.species += [species]
        else:
            self.species += species

    def mod_delete_atoms(self,
                         ids: int | list[int]) -> None:
        """
        Deletes selected atoms by ids
        Args:
            ids: sequence of atoms ids
        """
        self.coords = np.delete(self.coords, ids, axis=0)
        self.species = np.delete(self.species, ids)

    def mod_change_atoms(self,
                         ids: int | list[int],
                         coords: NDArray[Shape['Nids, 3'], Number] | None,
                         species: str | list[str] | None) -> None:
        """
        Change selected atom by id
        Args:
            ids: List or int. First id is zero.
            coords: None or np.array with new coords, e.g. np.array([1, 2, 4.6])
            species: None or str or List[str]. New types of changed atoms

        Returns:

        """
        if coords is not None:
            self.coords[ids] = coords
        if species is not None:
            if isinstance(ids, Iterable):
                for i, sp in zip(ids, species):
                    self.species[i] = sp
            else:
                self.species[ids] = species

    def mod_coords_to_cartesian(self) -> None | str:
        """
        Converts species coordinates to Cartesian coordination system.
        """
        if self.coords_are_cartesian is True:
            return 'Coords are already cartesian'
        else:
            self.coords = np.matmul(self.coords, self.lattice)
            self.coords_are_cartesian = True

    def mod_coords_to_direct(self) -> None | str:
        """
        Converts species coordinates to Direct coordination system.
        """
        if self.coords_are_cartesian is False:
            return 'Coords are already direct'
        else:
            transform = np.linalg.inv(self.lattice)
            self.coords = np.matmul(self.coords, transform)
            self.coords_are_cartesian = False

    def mod_add_vector(self,
                       vector: NDArray[Shape['3'], Number],
                       cartesian: bool = True) -> None:
        """
        Adds a vector to all atoms.
        Args:
            vector (np.ndarray): a vector which will be added to coordinates of all species
            cartesian (bool): determine in which coordination system the vector defined. Cartesian=True means that
            the vector is defined in Cartesian coordination system. Cartesian=False means that the vector is defined
            in Direct coordination system.
        """
        self.mod_coords_to_direct()
        if cartesian:
            transform = np.linalg.inv(self.lattice)
            vector = np.matmul(vector, transform)
        self.coords += vector
        self.coords %= np.array([1, 1, 1])

    def mod_propagate_unit_cell(self, a, b, c) -> None:
        """Extend the unit cell by propagation on lattice vector direction. The resulted unit cell will be the size
        of [a * lattice[0], b * lattice[1], c * lattice[2]]. All atoms will be replicated with accordance of new
        unit cell size
        Args:
            a (int): the number of replicas in the directions of the first unit cell vector
            b (int): the number of replicas in the directions of the second unit cell vector
            c (int): the number of replicas in the directions of the third unit cell vector.
        """
        if not (isinstance(a, int) and isinstance(b, int) and isinstance(c, int)):
            raise ValueError(f'a, b and c must be integers. But a, b, c have types {type(a), type(b), type(c)}')

        if a * b * c > 1:
            self.mod_coords_to_cartesian()
            ref_coords = self.coords.copy()
            ref_species = self.species.copy()

            for i in range(a):
                for j in range(b):
                    for k in range(c):
                        if i != 0 or j != 0 or k != 0:
                            vector = np.sum(np.array([i, j, k]).reshape(-1, 1) * self.lattice, axis=0)
                            self.coords = np.vstack((self.coords, ref_coords + vector))
                            self.species += ref_species
            self.lattice = np.array([a, b, c]).reshape(-1, 1) * self.lattice
        else:
            raise ValueError(f'a, b and c must be positive. But {a=} {b=} {c=}')

    def get_vector(self,
                   id_1: int,
                   id_2: int,
                   unit: bool = True) -> NDArray[Shape['3'], Number]:
        """
        Returns a vector (unit vector by default) which starts in the atom with id_1 and points to the atom with id_2
        Args:
            id_1: id of the first atom (vector origin)
            id_2: id of the second atom
            unit: Defines whether the vector will be normed or not. Unit=True means that the vector norm will be equal 1
        Returns (np.ndarray): vector from atom with id_1 to atom with id_2
        """
        vector = self.coords[id_2] - self.coords[id_1]
        if unit:
            return vector / np.linalg.norm(vector)
        else:
            return vector

    def get_distance_matrix(self) -> NDArray[Shape['Natoms, Natoms, 3'], Number]:
        """
        Returns distance matrix R, where R[i,j] is the vector from atom i to atom j.

        Returns:
            np.ndarray (NxNx3 dimensions) which is a distance matrix in Cartesian coordination system
        """
        self.mod_coords_to_direct()
        r1 = np.broadcast_to(self.coords.reshape((self.natoms, 1, 3)), (self.natoms, self.natoms, 3))
        r2 = np.broadcast_to(self.coords.reshape((1, self.natoms, 3)), (self.natoms, self.natoms, 3))
        R = r2 - r1
        R = (R + 0.5) % 1 - 0.5
        assert np.all(R >= - 0.5) and np.all(R <= 0.5)
        return np.matmul(R, self.lattice)

    def get_distance_matrix_scalar(self) -> NDArray[Shape['Natoms, Natoms'], Number]:
        """
        Returns distance matrix R, where R[i, j] is the Euclidean norm of a vector from atom i to atom j.

        Returns:
            np.ndarray (NxN dimensions) which is a distance matrix containing scalars
        """
        R = self.get_distance_matrix()
        return np.sqrt(np.sum(R * R, axis=2))

    def get_filtered_ids(self, **kwargs):
        """
        Returns np.ndarray that contains atom ids according to filter rules
        Args:
            species (str): Define which atom type will be selected. E.g. 'C H N' means select all C, H, and N atoms.
            '!C' means select all atoms except C.
        Returns:
            np.ndarray contains ids of atoms according to selecting rules
        """
        filter_mask = np.array([True for _ in range(self.natoms)], dtype=np.bool_)

        if 'species' in kwargs:
            species = kwargs['species'].split()
            species_select = []
            species_not_select = []
            for specie in species:
                if '!' in specie:
                    species_not_select.append(specie.replace('!', ''))
                else:
                    species_select.append(specie)
            if len(species_select):
                fm_local = np.array([False for _ in range(self.natoms)], dtype=np.bool_)
                for specie in species_select:
                    fm_local += np.array([True if atom_name == specie else False for atom_name in self.species])
                filter_mask *= fm_local
            if len(species_not_select):
                fm_local = np.array([True for _ in range(self.natoms)], dtype=np.bool_)
                for specie in species_not_select:
                    fm_local *= np.array([False if atom_name == specie else True for atom_name in self.species])
                filter_mask *= fm_local

        if 'x' in kwargs:
            left, right = kwargs['x']
            fm_local = np.array([False for _ in range(self.natoms)], dtype=np.bool_)
            for i, atom_coord in enumerate(self.coords):
                if left < atom_coord[0] < right:
                    fm_local[i] = True
            filter_mask *= fm_local

        if 'y' in kwargs:
            left, right = kwargs['y']
            fm_local = np.array([False for _ in range(self.natoms)], dtype=np.bool_)
            for i, atom_coord in enumerate(self.coords):
                if left < atom_coord[1] < right:
                    fm_local[i] = True
            filter_mask *= fm_local

        if 'z' in kwargs:
            left, right = kwargs['z']
            fm_local = np.array([False for _ in range(self.natoms)], dtype=np.bool_)
            for i, atom_coord in enumerate(self.coords):
                if left < atom_coord[2] < right:
                    fm_local[i] = True
            filter_mask *= fm_local

        return np.array(range(self.natoms))[filter_mask]


class StructureError(Exception):
    """
    Exception class for Structure.
    Raised when the structure has problems, e.g., atoms that are too close.
    """
    pass