"""Class containing a manager for creating template atoms."""
import numpy as np
from itertools import product
from random import shuffle
from ase.build import make_supercell
from clease import SkewnessFilter, EquivalentCellsFilter
from clease.template_filters import CellFilter, AtomsFilter
from clease.tools import all_integer_transform_matrices


class TemplateAtoms(object):
    def __init__(self, prim_cell, supercell_factor=None, size=None,
                 skew_threshold=4, filters=[]):
        if size is None and supercell_factor is None:
            raise TypeError("Either size or supercell_factor needs to be "
                            "specified.\n size: list or numpy array.\n "
                            "supercell_factor: int")

        self.supercell_factor = supercell_factor
        self.size = size

        self.cell_filters = []
        self.atoms_filters = []

        self.add_cell_filter(SkewnessFilter(skew_threshold))
        self.all_cells = []
        self.add_cell_filter(EquivalentCellsFilter(self.all_cells))

        for f in filters:
            if isinstance(f, CellFilter):
                self.add_cell_filter(f)
            elif isinstance(f, AtomsFilter):
                self.add_atoms_filter(f)

        if self.size is not None:
            check_valid_conversion_matrix(self.size)
        self.skew_threshold = skew_threshold
        self.prim_cell = prim_cell

    def __str__(self):
        """Print a summary of the class."""
        msg = "=== TemplateAtoms ===\n"
        msg += "Supercell factor: {}\n".format(self.supercell_factor)
        msg += "Skewed threshold: {}\n".format(self.skew_threshold)
        msg += "Template sizes:\n"
        for size in self.templates['size']:
            msg += "{}\n".format(size)
        return msg

    def add_cell_filter(self, cell_filter):
        """Attach a new Cell filter."""
        if not isinstance(cell_filter, CellFilter):
            raise TypeError("filter has to be an instance of CellFilter")
        self.cell_filters.append(cell_filter)

    def add_atoms_filter(self, at_filter):
        """Attach a new Atoms filter."""
        if not isinstance(at_filter, AtomsFilter):
            raise TypeError("filter has to be an instance of AtomsFilter")
        self.atoms_filters.append(at_filter)

    def clear_filters(self):
        """Remove all filters."""
        self.cell_filters = []
        self.atoms_filters = []

    def remove_filter(self, f):
        """Remove one filter"""
        if isinstance(f, AtomsFilter):
            self.atoms_filters.remove(f)
        elif isinstance(f, CellFilter):
            self.cell_filters.remove(f)
        else:
            raise TypeError('Only AtomsFilters and CellFilters '
                            'can be removed')

    def remove_filters(self, filters):
        """Remove a list of filters."""
        for f in filters:
            self.remove_filter(f)

    def is_valid(self, atoms=None, cell=None):
        """
        Check the validity of the template.

        Return `True` if templates are valid according to the attached filters.

        Parameters:

        atoms: Atoms object

        cell: unit cell vector
        """
        if atoms is None and cell is None:
            msg = "At least one of `atoms` or `cell` must be specified."
            raise ValueError(msg)

        cell_valid = True
        if cell is not None:
            cell_valid = all([f(cell) for f in self.cell_filters])
        elif atoms is not None:
            cell = atoms.get_cell()
            cell_valid = all([f(cell) for f in self.cell_filters])

        if not cell_valid:
            return False

        atoms_valid = True
        if atoms is not None:
            atoms_valid = all([f(atoms) for f in self.atoms_filters])

        return cell_valid and atoms_valid

    @property
    def largest_template_by_num_atom(self):
        """Return the largest template based on number of atoms it has."""
        max_num = 0
        largest_template = None
        for atoms in self.templates['atoms']:
            if len(atoms) > max_num:
                largest_template = atoms
                max_num = len(atoms)
        return largest_template

    @property
    def largest_template_by_diag(self):
        """Return the largest template based on the shortest diagonal."""
        length = 0.0
        largest_template = None
        for atoms in self.templates['atoms']:
            diag_lengths = []
            cell = atoms.get_cell().T
            for w in product([-1, 0, 1], repeat=3):
                if np.allclose(w, 0):
                    continue
                diag = cell.dot(w)
                length = np.sqrt(diag.dot(diag))
                diag_lengths.append(length)

            min_length = np.min(diag_lengths)

            if min_length > length:
                largest_template = atoms
                length = min_length

        return largest_template

    def get_template_with_given_size(self, size):
        """Get the UID of the template with given size.

        Parameters:

        size: list of length 3
        """
        if not is_3x3_matrix(size):
            raise ValueError("Expect the size to be a list of 3x3 matrix")

        check_valid_conversion_matrix(size)

        template = make_supercell(self.prim_cell, size)

        if isinstance(size, np.ndarray):
            size = size.tolist()
        template.info['size'] = size

        if not self.is_valid(atoms=template):
            raise ValueError("Requested size violates the constraints!")
        return template

    def get_template_matching_atoms(self, atoms):
        """Get the UID for the template matching atoms.

        Parameters:

        atoms: Atoms object
            structure to compare its size against template atoms
        """

        size = self._get_conversion_matrix(atoms)
        assert is_3x3_matrix(size)

        prim_cell = self.prim_cell

        if isinstance(size, np.ndarray):
            size = size.tolist()

        template = make_supercell(prim_cell, size)
        template.info['size'] = size

        if not np.allclose(atoms.get_cell(), template.get_cell()):
            raise ValueError("Inconsistent cells! Passed atoms\n"
                             "{}\nGenerated template\n"
                             "".format(atoms.get_cell(), template.get_cell()))

        if not self.is_valid(atoms=template):
            raise ValueError("Requested template violates the constraints!")
        return template

    def _get_conversion_matrix(self, atoms):
        """Return the conversion matrix factor."""
        prim_cell = self.prim_cell

        small_cell = prim_cell.get_cell()
        inv_cell = np.linalg.inv(small_cell)

        large_cell = atoms.get_cell()
        size_factor = large_cell.dot(inv_cell)
        scale_int = size_factor.round(decimals=0).astype(int)
        if np.allclose(size_factor, scale_int):
            check_valid_conversion_matrix(scale_int)
            return scale_int.tolist()

        raise ValueError("The passed atoms object cannot be described by "
                         "repeating of the unit cells. Scale factors found "
                         "{}".format(size_factor))

    def get_all_templates(self):
        """
        Return a list with all templates
        """
        if self.size is not None:
            return [self.get_template_with_given_size(self.size)]

        cells = []
        templates = []
        equiv_filter = EquivalentCellsFilter(cells)
        self.add_cell_filter(equiv_filter)
        ucell = self.prim_cell.get_cell()
        for vol in range(1, self.supercell_factor):
            matrices = all_integer_transform_matrices(vol)
            for mat in matrices:
                cell = mat.dot(ucell)
                if self.is_valid(cell=cell):
                    at = make_supercell(self.prim_cell, mat)
                    if self.is_valid(atoms=at):
                        at.info['size'] = mat.tolist()
                        cells.append(cell)
                        templates.append(at)
        self.remove_filter(equiv_filter)
        return templates

    def get_all_scaled_templates(self):
        """
        Return all templates that can be obtained by scaling the primitive cell
        """
        if self.size is not None:
            return [self.get_template_with_given_size(self.size)]
        cells = []
        templates = []
        equiv_filter = EquivalentCellsFilter(cells)
        self.add_cell_filter(equiv_filter)
        ucell = self.prim_cell.get_cell()
        for diag in product(range(1, self.supercell_factor+1), repeat=3):
            if np.prod(diag) > self.supercell_factor:
                continue

            mat = np.diag(diag)
            cell = mat.dot(ucell)
            if self.is_valid(cell=cell):
                at = make_supercell(self.prim_cell, mat)
                if self.is_valid(atoms=at):
                    cells.append(cell)
                    at.info['size'] = mat.tolist()
                    templates.append(at)
        self.remove_filter(equiv_filter)
        return templates

    def weighted_random_template(self):
        """Select a random template atoms with a bias towards a cubic cell.

        The bias is towards cells that have similar values for x-, y- and
        z-dimension sizes.
        """
        if self.size is not None:
            return self.get_template_with_given_size(self.size)

        ucell = self.prim_cell.get_cell()
        max_attempts = 100000
        exp_value = self.supercell_factor**(1.0/3.0)
        for _ in range(max_attempts):
            diag = np.random.poisson(lam=exp_value, size=3)
            while np.prod(diag) > self.supercell_factor or np.prod(diag) == 0:
                diag = np.random.poisson(lam=exp_value, size=3)

            off_diags = np.random.poisson(lam=1, size=3)

            # Check that off diagonals are smaller or equal to diagonal
            if (off_diags[0] > diag[0] or off_diags[1] > diag[0] or
               off_diags[2] > diag[1]):
                continue

            matrix = np.zeros((3, 3), dtype=int)
            matrix[0, 0] = diag[0]
            matrix[1, 1] = diag[1]
            matrix[2, 2] = diag[2]
            matrix[0, 1] = off_diags[0]
            matrix[0, 2] = off_diags[1]
            matrix[1, 2] = off_diags[2]
            cell = matrix.dot(ucell)

            if self.is_valid(cell=cell):
                atoms = make_supercell(self.prim_cell, matrix)
                if self.is_valid(atoms=atoms):
                    atoms.info['size'] = matrix.tolist()
                    return atoms

        raise RuntimeError("Did not manage to generate a random template that "
                           "satisfies all the constraints")

    def has_atoms_filters(self):
        return len(self.atoms_filters) > 0

    def get_fixed_volume_templates(self, num_prim_cells=10, num_templates=10):
        # Set up a filter that listens to the templates with fixed volume
        from ase.build.tools import niggli_reduce_cell
        cells = []
        transform_matrices = []

        equiv_filter = EquivalentCellsFilter(cells)
        self.add_cell_filter(equiv_filter)

        ucell = self.prim_cell.get_cell()
        matrices = list(all_integer_transform_matrices(num_prim_cells))
        shuffle(matrices)
        for mat in matrices:
            sc = mat.dot(ucell)
            sc, _ = niggli_reduce_cell(sc)
            if self.is_valid(cell=sc):
                # If Atoms filters are present we check if it is valid
                at_valid = True
                if self.has_atoms_filters():
                    atoms = make_supercell(self.prim_cell, mat)
                    at_valid = self.is_valid(atoms=atoms)
                if at_valid:
                    cells.append(sc)
                    transform_matrices.append(mat)

            if len(transform_matrices) >= num_templates:
                break

        templates = []

        if len(transform_matrices) > num_templates:
            shuffle(transform_matrices)
            transform_matrices = transform_matrices[:num_templates]

        for P in transform_matrices:
            atoms = make_supercell(self.prim_cell, P)
            templates.append(atoms)

        # Remove the filter that was artificially added
        self.cell_filters.remove(equiv_filter)
        return templates


def is_3x3_matrix(array):
    return np.array(array).shape == (3, 3)


def check_valid_conversion_matrix(array):
    """
    Make sure that we have a right-handed coordinate system.
    Raise a ValueError if the matrix is not valid.
    """
    determinant = np.linalg.det(array)
    if determinant < 0.0:
        raise ValueError("The determinant of the size matrix is less than "
                         "zero (got {}). For a right "
                         "coordinate system, we need a positive "
                         "determinant!".format(determinant))
