"""The data types required to represent a crystal."""

import json
from collections.abc import MutableSequence
from typing import List, Tuple

import numpy as np

from atompack.constants import DEG90, DEG120
from atompack.symmetry import Spacegroup


class Basis(MutableSequence):
    """Crystalline basis.

    Args:
        basis: List of specie/site pairs.

    Example:
        >>> from atompack.crystal import Basis
        >>> import numpy as np
        >>>
        >>> # primitive basis of iron atoms
        >>> basis = Basis.primitive("Fe")
        >>> assert len(basis) == 1
        >>>
        >>> specie, site = basis[0]
        >>> assert specie == "Fe"
        >>> assert np.array_equal(site, np.zeros(3))
    """

    def __init__(self, basis: List[Tuple[str, np.ndarray]]) -> None:
        for _, site in basis:
            if max(site) > 1 or min(site) < -1:
                raise ValueError("basis sites must be represented in fractional coordinates")
        self._basis = basis

    ########################################
    #    MutableSequence Implementation    #
    ########################################

    def __getitem__(self, index):
        return self._basis[index]

    def __setitem__(self, index, value):
        self._basis[index] = value

    def __delitem__(self, index):
        del self._basis[index]

    def __len__(self):
        return len(self._basis)

    def insert(self, index, value):
        self._basis.insert(index, value)

    ######################
    #    Constructors    #
    ######################

    @classmethod
    def primitive(cls, specie: str) -> 'Basis':
        """Returns a primitive basis."""
        return cls([(specie, np.zeros(3))])

    @classmethod
    def from_json(cls, s) -> 'Basis':
        """Initializes from a JSON string."""
        data = json.loads(s)
        basis = []
        for entry in data:
            specie = entry["specie"]
            site = np.array(entry["site"])
            basis.append((specie, site))
        return cls(basis)

    ########################
    #    Public Methods    #
    ########################

    def apply_spacegroup(self, spacegroup: Spacegroup) -> List[Tuple[str, np.ndarray]]:
        """Returns a list of specie/site pairs generated by applying a spacegroup's
        symmetry operations to the atomic basis."""
        # search tolerance
        tolerance = 1E-6

        # skip process if no sites exist
        if len(self) == 0:
            return []

        res = [self._basis[0]]
        # iterate over basis
        for specie, site in self._basis:
            # iterate through the spacegroup's general position expressions
            for genpos in spacegroup.genpos:
                # convert string expressions to arrays
                genpos = "[{}]".format(genpos)
                # x, y, and z are created here to be used in the eval statement
                x, y, z = site[0], site[1], site[2]
                new_site = np.array(eval(genpos))
                # wrap all values between 0-1
                for i in range(len(new_site)):
                    if new_site[i] < 0:
                        new_site[i] += 1
                # check if equivalent site exists
                is_occupied = False
                for _, _site in res:
                    if np.linalg.norm(new_site - _site) < tolerance:
                        is_occupied = True
                # store unique sites
                if not is_occupied:
                    res.append((specie, new_site))
        return res

    def to_json(self) -> str:
        """Returns a JSON serialized representation."""
        _basis = self._basis.copy()
        json_basis = [{"specie": specie, "site": site.tolist()} for specie, site in _basis]
        return json.dumps(json_basis)


class LatticeParameters(object):
    """Edge lengths and angles which define a lattice.
    
    Args:
        a: Length of the x lattice vector.
        b: Length of the y lattice vector.
        c: Length of the z lattice vector.
        alpha: Angle between the y and z directions (radians).
        beta: Angle between the x and z directions (radians).
        gamma: Angle between the x and y directions (radians).

    Example:
        >>> from atompack.crystal import LatticeParameters
        >>> import numpy as np
        >>>
        >>> # cubic lattice parameters
        >>> params = LatticeParameters.cubic(10)
        >>> assert params.a == params.b == params.c == 10
        >>> assert params.alpha == params.beta == params.gamma == np.pi / 2
    """

    def __init__(self, a: float, b: float, c: float, alpha: float, beta: float, gamma: float) -> None:
        self.a = a
        self.b = b
        self.c = c
        self.alpha = alpha
        self.beta = beta
        self.gamma = gamma

    ######################
    #    Constructors    #
    ######################

    @classmethod
    def triclinic(cls, a, b, c, alpha, beta, gamma) -> 'LatticeParameters':
        """Initializes with triclinic constraints."""
        return cls(a, b, c, alpha, beta, gamma)

    @classmethod
    def monoclinic(cls, a, b, c, beta) -> 'LatticeParameters':
        """Initializes with monoclinic constraints."""
        return cls(a, b, c, DEG90, beta, DEG90)

    @classmethod
    def orthorhombic(cls, a, b, c) -> 'LatticeParameters':
        """Initializes with orthorhombic constraints."""
        return cls(a, b, c, DEG90, DEG90, DEG90)

    @classmethod
    def tetragonal(cls, a, c) -> 'LatticeParameters':
        """Initializes with tetragonal constraints."""
        return cls(a, a, c, DEG90, DEG90, DEG90)

    @classmethod
    def trigonal(cls, a, c) -> 'LatticeParameters':
        """Initializes with trigonal constraints."""
        return cls(a, a, c, DEG90, DEG90, DEG120)

    @classmethod
    def rhombohedral(cls, a, alpha) -> 'LatticeParameters':
        """Initializes with rhombohedral constraints."""
        return cls(a, a, a, alpha, alpha, alpha)

    @classmethod
    def hexagonal(cls, a, c) -> 'LatticeParameters':
        """Initializes with hexagonal constraints."""
        return cls(a, a, c, DEG90, DEG90, DEG120)

    @classmethod
    def cubic(cls, a) -> 'LatticeParameters':
        """Initializes with cubic constraints."""
        return cls(a, a, a, DEG90, DEG90, DEG90)

    @classmethod
    def from_json(cls, s: str) -> 'LatticeParameters':
        """Initializes from a JSON string."""
        data = json.loads(s)
        data = {k: float(v) for k, v in data.items()}
        return cls(**data)

    ####################
    #    Properties    #
    ####################

    @property
    def metric_tensor(self) -> np.ndarray:
        """Returns the metric tensor defined by the lattice parameters."""
        return np.array([[self.a * self.a, self.a * self.b * np.cos(self.gamma), self.a * self.c * np.cos(self.beta)],
                         [self.a * self.b * np.cos(self.gamma), self.b * self.b, self.b * self.c * np.cos(self.alpha)],
                         [self.a * self.c * np.cos(self.beta), self.b * self.c * np.cos(self.alpha), self.c * self.c]])

    ########################
    #    Public Methods    #
    ########################

    def to_json(self) -> str:
        """Returns the JSON serialized representation."""
        return json.dumps({
            "a": self.a,
            "b": self.b,
            "c": self.c,
            "alpha": self.alpha,
            "beta": self.beta,
            "gamma": self.gamma
        })


class LatticeVectors(object):
    """Representation of the vectors that define the size and shape of a crystalline system.
    
    Args:
        vectors: Row-major matrix of lattice vectors.
    """

    def __init__(self, vectors: np.ndarray) -> None:
        self.vectors = vectors

    ######################
    #    Constructors    #
    ######################

    @classmethod
    def from_lattice_parameters(cls, lattice_parameters: LatticeParameters) -> 'LatticeVectors':
        return cls(np.sqrt(np.abs(lattice_parameters.metric_tensor)))

    @classmethod
    def from_json(cls, s: str) -> 'LatticeVectors':
        """Initializes from a JSON string."""
        data = json.loads(s)
        vectors = np.array(data["vectors"])
        return cls(vectors)

    ########################
    #    Public Methods    #
    ########################

    def contain(self, point: np.ndarray, tol: float = 1E-6) -> bool:
        """Returns True if the point is within the bounding volume."""
        bounds = np.linalg.norm(self.vectors, axis=0)
        for i in range(3):
            if point[i] > (bounds[i] + tol):
                return False
            if point[i] < -tol:
                return False
        return True

    def wrap(self, point: np.ndarray, tol: float = 1E-6) -> np.ndarray:
        """Wraps a point into the bounding volume. 
        The `point` argument is mutated and returned.
        """
        bounds = np.linalg.norm(self.vectors, axis=0)
        for i in range(3):
            bound = bounds[i]
            tmpval = point[i]
            if tmpval > (bound + tol):
                tmpval -= bound * (tmpval // bound)
            if tmpval < -tol:
                tmpval += bound * (1 + (-tmpval // bound))
            point[i] = tmpval
        return point

    def to_json(self) -> str:
        """Returns the JSON serialized representation."""
        return json.dumps({"vectors": self.vectors.tolist()})
