import itertools as it
import typing
from collections import Counter, abc
from functools import partial
from operator import getitem

import numpy as np

from stk._internal.building_block import BuildingBlock
from stk._internal.construction_result.construction_result import (
    ConstructionResult,
)
from stk._internal.construction_result.periodic import (
    PeriodicConstructionResult,
)
from stk._internal.optimizers.null import (
    NullOptimizer,
)
from stk._internal.optimizers.optimizer import Optimizer
from stk._internal.reaction_factories.generic_reaction_factory import (
    GenericReactionFactory,
)
from stk._internal.reaction_factories.reaction_factory import ReactionFactory
from stk._internal.topology_graphs.edge import (
    Edge,
)
from stk._internal.topology_graphs.edge_group import (
    EdgeGroup,
)
from stk._internal.topology_graphs.topology_graph.topology_graph import (
    TopologyGraph,
)
from stk._internal.topology_graphs.vertex import Vertex

from .edge import CofEdge
from .vertices import UnaligningVertex, _CofVertex


class UnoccupiedVertexError(Exception):
    """
    When a COF vertex is not occupied by a building block.

    """

    pass


class OverlyOccupiedVertexError(Exception):
    """
    When a COF vertex is occupied by more than one building block.

    """

    pass


class Cof(TopologyGraph):
    """
    Represents a COF topology graph.

    Notes:

        COF topologies are added by creating a subclass, which defines
        the :attr:`_vertex_prototypes` and :attr:`_edge_prototypes` class
        attributes.

    .. _cof-topology-graph-examples:

    Examples:

        *Subclass Implementation*

        The source code of the subclasses, listed in :mod:`~.cof.cof`,
        can serve as good examples.

        *Basic Construction*

        :class:`.Cof` instances can be made by providing the building
        block molecules and lattice size only (using :class:`.Honeycomb`
        as an example)

        .. testcode:: basic-construction

            import stk

            bb1 = stk.BuildingBlock('BrCCBr', [stk.BromoFactory()])
            bb2 = stk.BuildingBlock('BrCC(CBr)CBr', [stk.BromoFactory()])
            cof = stk.ConstructedMolecule(
                topology_graph=stk.cof.Honeycomb((bb1, bb2), (3, 3, 1)),
            )

        .. moldoc::

            import moldoc.molecule as molecule
            import stk

            bb1 = stk.BuildingBlock('BrCCBr', [stk.BromoFactory()])
            bb2 = stk.BuildingBlock('BrCC(CBr)CBr', [stk.BromoFactory()])
            cof = stk.ConstructedMolecule(
                topology_graph=stk.cof.Honeycomb((bb1, bb2), (3, 3, 1)),
            )
            moldoc_display_molecule = molecule.Molecule(
                atoms=(
                    molecule.Atom(
                        atomic_number=atom.get_atomic_number(),
                        position=position,
                    ) for atom, position in zip(
                        cof.get_atoms(),
                        cof.get_position_matrix(),
                    )
                ),
                bonds=(
                    molecule.Bond(
                        atom1_id=bond.get_atom1().get_id(),
                        atom2_id=bond.get_atom2().get_id(),
                        order=(
                            1
                            if bond.get_order() == 9
                            else bond.get_order()
                        ),
                    ) for bond in cof.get_bonds()
                ),
            )

        *Suggested Optimization*

        For :class:`.Cof` topologies, it is recommend to use the
        :class:`.Collapser` optimizer if using the nonperiodic form.
        For periodic systems, the :class:`.PeriodicCollapser` is
        recommended.

        .. testcode:: suggested-optimization

            import stk

            bb1 = stk.BuildingBlock('BrCCBr', [stk.BromoFactory()])
            bb2 = stk.BuildingBlock('BrCC(CBr)CBr', [stk.BromoFactory()])

            # Nonperiodic.
            cof = stk.ConstructedMolecule(
                topology_graph=stk.cof.Honeycomb(
                    building_blocks=(bb1, bb2),
                    lattice_size=(3, 3, 1),
                    # Setting scale_steps to False tends to lead to a
                    # better structure.
                    optimizer=stk.Collapser(scale_steps=False),
                ),
            )

            # Periodic.
            cof = stk.ConstructedMolecule(
                topology_graph=stk.cof.PeriodicHoneycomb(
                    building_blocks=(bb1, bb2),
                    lattice_size=(3, 3, 1),
                    optimizer=stk.PeriodicCollapser(),
                ),
            )

        .. moldoc::

            import moldoc.molecule as molecule
            import stk

            bb1 = stk.BuildingBlock('BrCCBr', [stk.BromoFactory()])
            bb2 = stk.BuildingBlock('BrCC(CBr)CBr', [stk.BromoFactory()])

            cof = stk.ConstructedMolecule(
                topology_graph=stk.cof.Honeycomb(
                    building_blocks=(bb1, bb2),
                    lattice_size=(3, 3, 1),
                    # Setting scale_steps to False tends to lead to a
                    # better structure.
                    optimizer=stk.Collapser(scale_steps=False),
                ),
            )
            moldoc_display_molecule = molecule.Molecule(
                atoms=(
                    molecule.Atom(
                        atomic_number=atom.get_atomic_number(),
                        position=position,
                    ) for atom, position in zip(
                        cof.get_atoms(),
                        cof.get_position_matrix(),
                    )
                ),
                bonds=(
                    molecule.Bond(
                        atom1_id=bond.get_atom1().get_id(),
                        atom2_id=bond.get_atom2().get_id(),
                        order=(
                            1
                            if bond.get_order() == 9
                            else bond.get_order()
                        ),
                    ) for bond in cof.get_bonds()
                ),
            )

        *Accessing the Periodic Information*

        When building periodic :class:`.Cof` instances, the periodic
        information, such as the unit cell, can be accessed if you use the
        :class:`.PeriodicConstructionResult` returned by calling
        :meth:`.Cof.construct`

        .. testcode:: accessing-the-periodic-information

            import stk

            bb1 = stk.BuildingBlock('BrCCBr', [stk.BromoFactory()])
            bb2 = stk.BuildingBlock('BrCC(CBr)CBr', [stk.BromoFactory()])

            topology_graph = stk.cof.PeriodicHoneycomb(
                building_blocks=(bb1, bb2),
                lattice_size=(3, 3, 1),
            )
            construction_result = topology_graph.construct()
            cof = stk.ConstructedMolecule.init_from_construction_result(
                construction_result=construction_result,
            )
            periodic_info = construction_result.get_periodic_info()
            cell_matrix = periodic_info.get_cell_matrix()
            # Can access all unit-cell parameters.
            a = periodic_info.get_a()
            b = periodic_info.get_b()
            c = periodic_info.get_c()
            alpha = periodic_info.get_alpha()
            beta = periodic_info.get_beta()
            gamma = periodic_info.get_gamma()
            # Write to .pdb file.
            writer = stk.PdbWriter()
            writer.write(
                molecule=cof,
                path='cof.pdb',
                periodic_info=periodic_info,
            )

        .. testcleanup:: accessing-the-periodic-information

            import os
            os.remove('cof.pdb')


        *Structural Isomer Construction*

        Different structural isomers of COFs can be made by using the
        `vertex_alignments` optional parameter

        .. testcode:: structural-isomer-construction

            import stk

            bb1 = stk.BuildingBlock('BrCCBr', [stk.BromoFactory()])
            bb2 = stk.BuildingBlock('BrCC(CBr)CBr', [stk.BromoFactory()])

            cof2 = stk.ConstructedMolecule(
                topology_graph=stk.cof.Honeycomb(
                    building_blocks=(bb1, bb2),
                    lattice_size=(3, 3, 1),
                    vertex_alignments={0: 2, 1: 1, 2: 1},
                ),
            )

        The parameter maps the id of a vertex to a number
        between 0 (inclusive) and the number of edges the vertex is
        connected to (exclusive). So a vertex connected to three edges
        can be mapped to ``0``, ``1`` or ``2``.

        By changing which edge each vertex is aligned with, a different
        structural isomer of the COF can be formed.

        *Multi-Building Block COF Construction*

        You can also build COFs with multiple building blocks, but,
        if you have multiple building blocks with the same number of
        functional groups, you have to assign each building block to the
        vertex you want to place it on

        .. testcode:: multi-building-block-cof-construction

            import stk

            bb1 = stk.BuildingBlock('BrCCBr', [stk.BromoFactory()])
            bb2 = stk.BuildingBlock('BrCNCBr', [stk.BromoFactory()])
            bb3 = stk.BuildingBlock('BrCC(CBr)CBr', [stk.BromoFactory()])
            bb4 = stk.BuildingBlock('BrCC(NCBr)CBr', [stk.BromoFactory()])

            cof = stk.ConstructedMolecule(
                topology_graph=stk.cof.Honeycomb(
                    # building_blocks is now a dict, which maps building
                    # blocks to the id of the vertices it should be placed
                    # on. You can use ranges to specify the ids.
                    building_blocks={
                        bb1: (2, 4, 7, 8, 9, 12, 13, 14, 17, 18, 19),
                        bb2: 3,
                        bb3: (0, 10, 11, 15, 16),
                        bb4: (1, 5, 6),
                    },
                    lattice_size=(2, 2, 1),
                ),
            )

        You can combine this with the `vertex_alignments` parameter

        .. testcode:: multi-building-block-cof-construction

            cof2 = stk.ConstructedMolecule(
                topology_graph=stk.cof.Honeycomb(
                    building_blocks={
                        bb1: (2, 4, 7, 8, 9, 12, 13, 14, 17, 18, 19),
                        bb2: 3,
                        bb3: (0, 10, 11, 15, 16),
                        bb4: (1, 5, 6),
                    },
                    lattice_size=(2, 2, 1),
                    vertex_alignments={0: 2, 1: 1, 2: 1},
                ),
            )

    """

    _allowed_degrees: typing.ClassVar[dict[int, int]]
    _vertex_prototypes: typing.ClassVar[tuple[_CofVertex, ...]]
    _edge_prototypes: typing.ClassVar[tuple[Edge, ...]]
    _lattice_constants: typing.ClassVar[
        tuple[np.ndarray, np.ndarray, np.ndarray]
    ]

    def __init_subclass__(cls, **kwargs):
        vertex_degrees = Counter(
            vertex_id
            for edge in cls._edge_prototypes
            for vertex_id in edge.get_vertex_ids()
        )
        cls._allowed_degrees = set(vertex_degrees.values())

    def __init__(
        self,
        building_blocks: (
            abc.Iterable[BuildingBlock] | dict[BuildingBlock, tuple[int, ...]]
        ),
        lattice_size: tuple[int, int, int],
        periodic: bool = False,
        vertex_alignments: dict[int, int] | None = None,
        reaction_factory: ReactionFactory = GenericReactionFactory(),
        num_processes: int = 1,
        optimizer: Optimizer = NullOptimizer(),
        scale_multiplier: float = 1.0,
    ) -> None:
        """
        Initialize a :class:`.Cof` instance.

        Parameters:

            building_blocks:
                Can be a :class:`tuple` of :class:`.BuildingBlock`
                instances, which should be placed on the topology graph.

                Can also be a :class:`dict` which maps the
                :class:`.BuildingBlock` instances to the ids of the
                vertices it should be placed on. A :class:`dict` is
                required when there are multiple building blocks with the
                same number of functional groups, because in this case
                the desired placement is ambiguous.

            lattice_size:
                The size of the lattice in the x, y and z directions.

            periodic:
                Toggle the construction of a periodic molecule. If
                ``True``, periodic bonds will be made across the edges of
                the lattice.

            vertex_alignments:
                A mapping from the id of a :class:`.Vertex`
                to an :class:`.Edge` connected to it.
                The :class:`.Edge` is used to align the first
                :class:`.FunctionalGroup` of a :class:`.BuildingBlock`
                placed on that vertex. Only vertices which need to have
                their default edge changed need to be present in the
                :class:`dict`. If ``None`` then the default edge is used
                for each vertex. Changing which :class:`.Edge` is used will
                mean that the topology graph represents different
                structural isomers. The edge is referred to by a number
                between ``0`` (inclusive) and the number of edges the
                vertex is connected to (exclusive).

            reaction_factory:
                The reaction factory to use for creating bonds between
                building blocks.

            num_processes:
                The number of parallel processes to create during
                :meth:`construct`.

            optimizer:
                Used to optimize the structure of the constructed
                molecule.

            scale_multiplier:
                Scales the positions of the vertices.

        Raises:

            :class:`AssertionError`
                If the any building block does not have a
                valid number of functional groups.

            :class:`ValueError`
                If the there are multiple building blocks with the
                same number of functional_groups in `building_blocks`,
                and they are not explicitly assigned to vertices. The
                desired placement of building blocks is ambiguous in
                this case.

            :class:`~.cof.UnoccupiedVertexError`
                If a vertex of the COF topology graph does not have a
                building block placed on it.

            :class:`~.cof.OverlyOccupiedVertexError`
                If a vertex of the COF topology graph has more than one
                building block placed on it.

        """

        self._vertex_alignments = (
            dict(vertex_alignments) if vertex_alignments is not None else {}
        )
        self._lattice_size = lattice_size
        self._periodic = periodic

        lattice = self._get_lattice(self._vertex_alignments)
        edges = self._get_edges(lattice)
        vertices = self._get_vertices(lattice)

        if isinstance(building_blocks, dict):
            for building_block in building_blocks:
                assert (
                    building_block.get_num_functional_groups()
                    in self._allowed_degrees
                ), (
                    "The number of functional groups in "
                    f"{building_block} needs to be one of "
                    f"{tuple(self._allowed_degrees)}, but is "
                    "currently "
                    f"{building_block.get_num_functional_groups()}."
                )
            get_vertex = partial(getitem, vertices)
            building_block_vertices = {
                building_block: tuple(
                    map(
                        get_vertex,
                        # Account for the fact that a building block can
                        # be mapped to a single int.
                        (ids,) if isinstance(ids, int) else ids,
                    )
                )
                for building_block, ids in building_blocks.items()
            }
        else:
            building_block_vertices = self._get_building_block_vertices(  # type: ignore[assignment]
                building_blocks=building_blocks,
                vertices=vertices,
                edges=edges,
            )

        building_block_vertices = self._with_unaligning_vertices(
            building_block_vertices=building_block_vertices,
        )

        self._check_building_block_vertices(
            num_vertices=(
                np.prod(lattice_size) * len(self._vertex_prototypes)
            ),
            building_block_vertices=building_block_vertices,
        )

        super().__init__(
            building_block_vertices=typing.cast(
                dict[BuildingBlock, abc.Sequence[Vertex]],
                building_block_vertices,
            ),
            edges=edges,
            reaction_factory=reaction_factory,
            construction_stages=(),
            num_processes=num_processes,
            optimizer=optimizer,
            edge_groups=self._get_edge_groups(edges),
            scale_multiplier=scale_multiplier,
        )

    @staticmethod
    def _with_unaligning_vertices(building_block_vertices):
        clone = dict(building_block_vertices)
        for building_block, vertices in clone.items():
            # Building blocks with 1 placer, cannot be aligned and
            # must therefore use an UnaligningVertex.
            if building_block.get_num_placers() == 1:
                clone[building_block] = tuple(
                    UnaligningVertex(
                        id=v.get_id(),
                        position=v.get_position(),
                        aligner_edge=v.get_aligner_edge(),
                        cell=v.get_cell(),
                    )
                    for v in vertices
                )

        return clone

    @classmethod
    def _check_building_block_vertices(
        cls,
        num_vertices,
        building_block_vertices,
    ):
        unassigned_ids = set(range(num_vertices))
        assigned_ids = set()
        vertices = (
            vertex
            for vertices_ in building_block_vertices.values()
            for vertex in vertices_
        )
        for vertex in vertices:
            if vertex.get_id() in assigned_ids:
                raise OverlyOccupiedVertexError(
                    f"Vertex {vertex.get_id()} has multiple building "
                    "blocks placed on it."
                )
            assigned_ids.add(vertex.get_id())
            unassigned_ids.remove(vertex.get_id())

        if unassigned_ids:
            raise UnoccupiedVertexError(
                f"The following vertices are unoccupied {unassigned_ids}."
            )

    def clone(self):
        clone = super().clone()
        clone._vertex_alignments = dict(self._vertex_alignments)
        clone._lattice_size = self._lattice_size
        clone._periodic = self._periodic
        return clone

    def _get_edge_groups(
        self,
        edges: abc.Iterable[Edge],
    ) -> tuple[EdgeGroup, ...] | None:
        """
        Get the edge groups for the COF.

        Parameters:

            edges:
                The edges which are to be placed into edge groups.

        Returns:

            The edge groups. These are returned when the lattice is
            not periodic, which means that some edges should not be
            part of an edge group, in order to prevent bond formation
            across the edges of the lattice.

            None:
                If an edge group is to be created for every single edge.

        """

        if self._periodic:
            return None

        return tuple(
            EdgeGroup((edge,)) for edge in edges if not edge.is_periodic()
        )

    def _get_vertices(
        self,
        lattice: list[list[list[dict[int, _CofVertex]]]],
    ) -> list[_CofVertex]:
        """
        Get the vertices in the `lattice`.

        Parameters:

            lattice:
                A nested list which can be in the form
                ``lattice[x][y][z][vertex_id]`` which returns a vertex
                in the (x, y, z) cell of the lattice.

        Returns:

            All the vertices extracted from `lattice`.

        """

        xdim, ydim, zdim = self._lattice_size
        vertices = []
        for x, y, z in it.product(
            range(xdim),
            range(ydim),
            range(zdim),
        ):
            for vertex in lattice[x][y][z].values():
                vertices.append(vertex)
        return sorted(vertices, key=lambda vertex: vertex.get_id())

    def _get_lattice(
        self,
        vertex_alignments: dict[int, int],
    ) -> list[list[list[dict[int, _CofVertex]]]]:
        """
        Get the vertices of the topology graph instance.

        Parameters:

            vertex_alignments:
                A mapping from the id of a :class:`.Vertex`
                to an :class:`.Edge` connected to it.
                The :class:`.Edge` is used to align the first
                :class:`.FunctionalGroup` of a :class:`.BuildingBlock`
                placed on that vertex. Only vertices which need to have
                their default edge changed need to be present in the
                :class:`dict`. If ``None`` then the default edge is used
                for each vertex. Changing which :class:`.Edge` is used will
                mean that the topology graph represents different
                structural isomers. The edge is referred to by a number
                between ``0`` (inclusive) and the number of edges the
                vertex is connected to (exclusive).

        Returns:

            A nested :class:`list` which can be indexed as
            ``vertices[x][y][z]``, which will return a :class:`dict`
            for the unit cell at (x, y, z). The :class:`dict` maps
            the vertices in :attr:`_vertex_prototypes` to its clone for
            that unit cell.

        """

        xdim, ydim, zdim = (range(dim) for dim in self._lattice_size)
        # vertex_clones is indexed as vertex_clones[x][y][z]
        lattice: list[list[list[dict[int, _CofVertex]]]] = [
            [[{} for _ in zdim] for _ in ydim] for _ in xdim
        ]
        # Make a clone of each vertex for each unit cell.
        cells = it.product(xdim, ydim, zdim)
        vertices = it.product(cells, self._vertex_prototypes)
        for id_, (cell, vertex) in enumerate(vertices):
            x, y, z = cell
            shift = sum(
                axis * dim for axis, dim in zip(cell, self._lattice_constants)
            )
            lattice[x][y][z][vertex.get_id()] = vertex.__class__(
                id=id_,
                position=vertex.get_position() + shift,
                aligner_edge=vertex_alignments.get(id_, 0),
                cell=cell,
            )
        return lattice

    def _get_edges(self, lattice: list) -> tuple[Edge, ...]:
        """
        Create the edges of the topology graph instance.

        Parameters:

            lattice:
                A nested :class:`list` which can be indexed as
                ``vertices[x][y][z]``, which will return a :class:`dict`
                for the unit cell at (x, y, z). The :class:`dict` maps
                the vertices in :attr:`_vertex_prototypes` to the clones
                for that unit cell.

        Returns:

            The edges of the topology graph instance.

        """

        edge_clones = []
        # Make a clone for each edge for each unit cell.
        xdim, ydim, zdim = (range(dim) for dim in self._lattice_size)
        cells = it.product(xdim, ydim, zdim)
        edges = it.product(cells, self._edge_prototypes)
        for id_, (cell, edge) in enumerate(edges):
            x, y, z = cell
            # The cell in which the second vertex of the edge is found.
            periodic_cell = np.array(cell) + edge.get_periodicity()
            # Wrap around periodic cells, ie those that are less than 0
            # or greater than the lattice size along any dimension.
            dims = zip(periodic_cell, self._lattice_size)
            x2, y2, z2 = np.array(
                [(dim + max_dim) % max_dim for dim, max_dim in dims]
            )
            # The edge is not periodic if periodic_cell did not
            # have to wrap around.
            dims = zip(periodic_cell, self._lattice_size)
            edge_is_not_periodic = all(
                dim >= 0 and dim < max_dim for dim, max_dim in dims
            )
            edge_clones.append(
                CofEdge(
                    parent_id=edge.get_id(),
                    id=id_,
                    vertex1=lattice[x][y][z][edge.get_vertex1_id()],
                    vertex2=lattice[x2][y2][z2][edge.get_vertex2_id()],
                    periodicity=(
                        (0, 0, 0)
                        if edge_is_not_periodic
                        else edge.get_periodicity()
                    ),
                )
            )

        return tuple(edge_clones)

    @classmethod
    def _get_building_block_vertices(
        cls,
        building_blocks: abc.Iterable[BuildingBlock],
        vertices: abc.Iterable[_CofVertex],
        edges: abc.Iterable[Edge],
    ) -> dict[BuildingBlock, abc.Sequence[_CofVertex]]:
        """
        Map building blocks to the vertices of the graph.

        Parameters:

            building_blocks:
                The building blocks which need to be mapped to `vertices`.

            vertices:
                The vertices which need to have a building block map to
                them.

            edges:
                The edges of the graph.

        Returns:

            Maps each building block in `building_blocks` to a
            :class:`list` of the :class:`.Vertex` instances it should
            be placed on.

        Raises:

            :class:`AssertionError`
                If the any building block does not have a
                valid number of functional groups.

            :class:`ValueError`
                If there are multiple building blocks with the same number
                of functional groups.

        """

        building_blocks_by_degree = {}
        for building_block in building_blocks:
            num_fgs = building_block.get_num_functional_groups()
            assert num_fgs in cls._allowed_degrees, (
                "The number of functional groups in "
                f"{building_block} needs to be one of "
                f"{tuple(cls._allowed_degrees)}, but is "
                "currently "
                f"{building_block.get_num_functional_groups()}."
            )
            if num_fgs in building_blocks_by_degree:
                raise ValueError(
                    "If there are multiple building blocks with the "
                    "same number of functional groups, "
                    "building_block_vertices must be set explicitly."
                )
            building_blocks_by_degree[num_fgs] = building_block

        vertex_degrees = Counter(
            vertex_id for edge in edges for vertex_id in edge.get_vertex_ids()
        )

        building_block_vertices = {}  # type: ignore[var-annotated]
        for vertex in vertices:
            vertex_degree = vertex_degrees[vertex.get_id()]
            building_block = building_blocks_by_degree[vertex_degree]
            building_block_vertices[building_block] = (
                building_block_vertices.get(building_block, [])
            )
            building_block_vertices[building_block].append(vertex)
        return building_block_vertices

    def _get_lattice_constants(self) -> abc.Iterator[np.ndarray]:
        yield from self._lattice_constants

    def construct(self) -> ConstructionResult:
        """
        Construct a :class:`.ConstructedMolecule`.

        Returns:

            The data describing the :class:`.ConstructedMolecule`.

        """

        return super().construct()

    def _get_construction_result(self, state):
        return PeriodicConstructionResult(state, self._lattice_size)

    @staticmethod
    def _get_scale(
        building_block_vertices: dict[BuildingBlock, abc.Sequence[Vertex]],
        scale_multiplier: float,
    ) -> float:
        return (
            scale_multiplier
            * 5
            * max(bb.get_maximum_diameter() for bb in building_block_vertices)
        )

    def __repr__(self):
        x, y, z = self._lattice_size
        periodic = ", periodic=True" if self._periodic else ""
        vertex_alignments = (
            f", vertex_alignments={self._vertex_alignments}"
            if self._vertex_alignments
            else ""
        )
        return (
            f"cof.{self.__class__.__name__}("
            f"({x}, {y}, {z})"
            f"{vertex_alignments}"
            f"{periodic})"
        )
