"""
Topology Graph
==============

"""

from __future__ import annotations

import typing
from collections import abc
from functools import partial

import numpy as np

from stk.utilities import flatten

from ....molecules import BuildingBlock
from ....reactions import ReactionFactory
from ..construction_result import ConstructionResult
from ..construction_state import ConstructionState
from ..edge import Edge
from ..edge_group import EdgeGroup
from ..optimizers import Optimizer
from ..vertex import Vertex
from .implementations import (
    _Parallel,
    _Serial,
    _TopologyGraphImplementation,
)

_TopologyGraphT = typing.TypeVar(
    "_TopologyGraphT",
    bound="TopologyGraph",
)


class TopologyGraph:
    """
    An abstract base class for topology graphs.

    It is responsible for the construction of molecules. To create
    a new topology graph, you want to subclass and implement this
    abstract base class.

    Notes:

        *Adding New Topology Graphs*

        You might notice that some of the methods of this abstract base
        class are implemented. This is purely for convenience when
        implementing subclasses. The implemented public methods are
        simply default implementations, which can safely be ignored or
        overridden, when implementing subclasses. Any private methods
        are implementation details of these default implementations.

        Many classes, such as :class:`.Vertex`, :class:`.Edge`,
        :class:`.EdgeGroup` and :class:`.ConstructionState`, exist as
        implementation details of this default :class:`.TopologyGraph`
        implementation. You could ignore all of them, and define a new
        :meth:`.construct` method from scratch. In fact, your topology
        graph does not have to be represented as a graph at all.
        However, using the default implementation of
        :class:`.TopologyGraph` makes it significantly easier to
        implement a construction process. When using the default
        implementation of :class:`.TopologyGraph`, you mostly just
        need to implement a :class:`.Vertex` subclass, which is much
        easier than figuring out the whole construction process from
        scratch. In addition, you get benefits like parallel
        construction for free, as it is included in the default
        implementation.

        Typically, adding a new topology graph will involve
        implementing any pure virtual methods of
        :class:`.TopologyGraph`, in a new subclass, as well as
        implementing any pure virtual methods of :class:`.Vertex`,
        again in a new subclass. Combined, this is just a handful of
        simple methods to implement. Sometimes, rarely, you might also
        want to subclass :class:`.ConstructionState`, when you want to
        add additional hooks during construction, by extending
        the methods of this class. If you do this, make sure to
        override :meth:`._get_construction_state` to return your
        subclass of :class:`.ConstructionState`, rather than the base
        class, as is done by default. You can subclass and extend the
        methods of any class as you wish, but it would be unusual if
        this doesn't cover all your requirements.

        *The Default Implementation*

        The default implementation of :class:`.TopologyGraph`
        represents the constructed molecule through a graph. The
        vertices indicate where building blocks are placed and the
        edges indicate which building blocks have bonds formed between
        them by the construction process.

        :class:`.Vertex` instances are responsible for placing the
        building block molecules. By initializing the vertices with
        different parameters, you can alter how they position the
        building block molecules, and therefore allow the user to
        easily specify a different structural isomer.

        Once a building block is placed on a vertex, the functional
        groups on the building block must be mapped to the different
        edges connected to the vertex. The number of functional groups
        in the building block must match the number of edges connected
        to the vertex.

        Once the functional groups are mapped to edges, the edges are
        used to perform reactions on the building blocks. Edges are
        grouped in an :class:`.EdgeGroup`, and all functional groups
        present in the edge group are reacted together. Normally,
        unless you are doing something very exotic, an
        :class:`.EdgeGroup` will hold just one :class:`.Edge`, and the
        two functional groups on that edge will be reacted together
        through a single :class:`.Reaction`. This reaction will
        normally add the bonds which are required to form the
        joined-up constructed molecule, but note that it does not have
        to add any bonds at all. In addition, a :class:`.Reaction` can
        add and remove atoms from the constructed molecule. Which
        reaction is selected to join the functional groups depends on
        the :class:`.ReactionFactory` given to the
        :class:`.TopologyGraph` during initialization.

        Once this is done, you have a :class:`.ConstructedMolecule`.

    Examples:

        *Subclass Implementation*

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

        *Changing the Building Blocks of a Topology Graph*

        To change the building blocks used by a topology graph you
        can use :meth:`.with_building_blocks` to get a clone of the
        topology graph holding the new building blocks

        .. testcode:: changing-the-building-blocks-of-a-topology-graph

            import stk

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

            linear = stk.polymer.Linear(
                building_blocks=(bb1, bb2),
                repeating_unit='A',
                num_repeating_units=15,
            )

            bb3 = stk.BuildingBlock('BrCNCBr', [stk.BromoFactory()])
            # All bb1 instances are replaced by bb3, but bb2 remains
            # in place.
            clone = linear.with_building_blocks({
                bb1: bb3,
            })

    """

    _implementation: _TopologyGraphImplementation

    def __init__(
        self,
        building_block_vertices: dict[
            BuildingBlock,
            abc.Sequence[Vertex],
        ],
        edges: tuple[Edge, ...],
        reaction_factory: ReactionFactory,
        construction_stages: tuple[
            # TODO: Use typing.Callable here for now so that Sphinx
            # generates hyperlinks in the compiled docs. This should
            # eventually be replaced by abc.Callable.
            typing.Callable[[Vertex], bool],
            ...,
        ],
        num_processes: int,
        optimizer: Optimizer,
        edge_groups: typing.Optional[tuple[EdgeGroup, ...]] = None,
    ) -> None:
        """
        Initialize an instance of :class:`.TopologyGraph`.

        Parameters:

            building_block_vertices:
                Maps each :class:`.BuildingBlock` to be placed, to the
                :class:`.Vertex` instances, on which it should be
                placed.

            edges:
                The edges which make up the topology graph.

            reaction_factory:
                Used to pick which :class:`.Reaction` is used on each
                :class:`.EdgeGroup` of the topology graph.

            construction_stages:
                A collection of :class:`~collections.abc.Callable`,
                each of which takes a :class:`.Vertex` and returns
                ``True`` or ``False``. If the first
                :class:`~collections.abc.Callable` is applied
                to a  vertex in the topology graph, and the result is
                ``True``, that vertex is a part of the first
                construction stage. The second :class:`callable` is
                then applied to all vertices not in the first stage
                and those which return ``True`` belong to the second
                stage and so on.

                Vertices which belong to the same construction stage
                all place building blocks together in parallel, before
                placement is done by any vertices which are part of a
                later stage. This breaks down parallel construction
                into serial stages if synchronization between stages
                is needed.

                If the topology graph is performing construction
                serially, then all vertices which belong to an earlier
                stage will place their building block before those at
                a later stage.

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

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

            edge_groups:
                The edge groups of the topology graph, if ``None``,
                every :class:`.Edge` is in its own edge group.

        """

        self._scale = scale = self._get_scale(building_block_vertices)

        def apply_scale(item):
            return item.with_scale(scale)

        self._building_block_vertices = {
            building_block: tuple(map(apply_scale, vertices))
            for building_block, vertices in building_block_vertices.items()
        }
        self._edges = tuple(map(apply_scale, edges))
        self._reaction_factory = reaction_factory

        if num_processes == 1:
            self._implementation = _Serial(
                stages=tuple(self._get_stages(construction_stages)),
            )
        else:
            self._implementation = _Parallel(
                stages=tuple(self._get_stages(construction_stages)),
                num_processes=num_processes,
            )

        if edge_groups is None:
            edge_groups = tuple(
                EdgeGroup((edge,)) for edge in self._edges
            )
        self._edge_groups = edge_groups

        self._optimizer = optimizer

    def _with_building_blocks(
        self: _TopologyGraphT,
        building_block_map: dict[BuildingBlock, BuildingBlock],
    ) -> _TopologyGraphT:
        """
        Modify the topology graph.

        """

        # The original scaling first needs to be removed, so that when
        # the scale is recalculated with the new building blocks, it
        # has the same starting geometry.
        def undo_scale(vertex):
            return vertex.with_scale(1 / self._scale)

        building_block_vertices: dict[
            BuildingBlock, abc.Sequence[Vertex]
        ]
        building_block_vertices = {
            building_block_map.get(
                building_block, building_block
            ): tuple(map(undo_scale, vertices))
            for building_block, vertices in self._building_block_vertices.items()
        }
        scale = self._get_scale(building_block_vertices)

        def scale_vertex(vertex):
            return vertex.with_scale(scale)

        self._building_block_vertices = {
            building_block: tuple(map(scale_vertex, vertices))
            for building_block, vertices in building_block_vertices.items()
        }

        def scale_edge(edge):
            # Remove the old scale and apply the new one.
            return edge.with_scale(scale / self._scale)

        self._edges = edges = tuple(map(scale_edge, self._edges))

        def get_new_edge(edge_id):
            return edges[edge_id]

        self._edge_groups = tuple(
            EdgeGroup(map(get_new_edge, edge_group.get_edge_ids()))
            for edge_group in self._edge_groups
        )
        self._scale = scale
        return self

    def with_building_blocks(
        self,
        building_block_map: dict[BuildingBlock, BuildingBlock],
    ) -> TopologyGraph:
        """
        Return a clone holding different building blocks.

        Parameters:

            building_block_map:
                Maps a building block in the current topology
                graph to the building block which should replace
                it in the clone. If a building block should be not
                replaced in the clone, it can be omitted from the map.

        Returns:

            The clone.

        """

        return self.clone()._with_building_blocks(building_block_map)

    def clone(self) -> TopologyGraph:
        """
        Return a clone.

        Returns:

            The clone.

        """

        return self._clone()

    def _clone(self: _TopologyGraphT) -> _TopologyGraphT:
        clone = self.__class__.__new__(self.__class__)
        clone._scale = self._scale
        clone._building_block_vertices = dict(
            self._building_block_vertices
        )
        clone._edges = self._edges
        clone._reaction_factory = self._reaction_factory
        clone._implementation = self._implementation
        clone._optimizer = self._optimizer
        clone._edge_groups = self._edge_groups
        return clone

    def get_building_blocks(self) -> typing.Iterator[BuildingBlock]:
        """
        Yield the building blocks.

        Building blocks are yielded in an order based on their
        position in the topology graph. For two equivalent
        topology graphs, but with different building blocks,
        equivalently positioned building blocks will be yielded at the
        same time.

        Yields:

            A building block of the topology graph.

        """

        vertex_building_blocks = {}
        num_vertices = 0
        for (
            building_block,
            vertices,
        ) in self._building_block_vertices.items():
            for vertex in vertices:
                num_vertices += 1
                vertex_building_blocks[
                    vertex.get_id()
                ] = building_block

        yielded = set()
        for vertex_id in range(num_vertices):
            building_block = vertex_building_blocks[vertex_id]
            if building_block not in yielded:
                yielded.add(building_block)
                yield building_block

    def get_num_building_block(
        self,
        building_block: BuildingBlock,
    ) -> int:
        """
        Get the number of times `building_block` is present.

        Parameters:

            building_block:
                The building block whose frequency in the topology
                graph is desired.

        Returns:

            The number of times `building_block` is present in the
            topology graph.

        """

        return len(
            self._building_block_vertices.get(building_block, [])
        )

    def _get_lattice_constants(self) -> typing.Iterator[np.ndarray]:
        """
        Yield the lattice constants of the topology graph.

        The a, b and c lattice constants are yielded, in that order.

        By default, this is an empty generator.

        Yields:

            A lattice constant.

        """

        return
        yield

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

        Returns:

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

        """

        state = self._get_construction_state()
        state = self._place_building_blocks(state)
        state = self._run_reactions(state)
        state = self._optimizer.optimize(state)
        return self._get_construction_result(state)

    def _get_construction_result(
        self,
        state: ConstructionState,
    ) -> ConstructionResult:
        """
        Get the result of the construction.

        Parameters:

            state:
                The state of the molecule being constructed.

        Returns:

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

        """

        return ConstructionResult(state)

    def _get_construction_state(self) -> ConstructionState:
        return ConstructionState(
            building_block_vertices=self._building_block_vertices,
            edges=self._edges,
            lattice_constants=tuple(
                np.array(constant, dtype=np.float64) * self._scale
                for constant in self._get_lattice_constants()
            ),
        )

    def _get_scale(
        self,
        building_block_vertices: dict[
            BuildingBlock,
            abc.Sequence[Vertex],
        ],
    ) -> float:
        """
        Get the scale, which should be applied to topology graph.

        The scale should be applied to the position of every vertex
        and edge of topology graph. This allows to graph to adjust
        based on the size of the building blocks.

        Parameters:

            building_block_vertices:
                Maps every :class:`.BuildingBlock` of the topology
                graph, to the :class:`.Vertex` instances it is meant
                to be placed on.

        Returns:

            The scale.

        """

        raise NotImplementedError()

    def _place_building_blocks(
        self,
        state: ConstructionState,
    ) -> ConstructionState:
        """
        Place the building blocks onto the vertices.

        Parameters:

            state:
                Holds data necessary to construct the molecule.

        Returns:

            The new construction state, updated to account for the
            placed building blocks.

        """

        return self._implementation._place_building_blocks(state)

    def _run_reactions(
        self,
        state: ConstructionState,
    ) -> ConstructionState:
        """
        Perform the reactions on the building blocks.

        Parameters:

            state:
                The current state of the construction process.

        Returns:

            The new construction state, updated to account for the
            reactions between building blocks.

        """

        get_reaction = partial(
            self._reaction_factory.get_reaction,
            state,
        )
        reactions = tuple(map(get_reaction, self._edge_groups))
        results = map(
            lambda reaction: reaction.get_result(),
            reactions,
        )
        return state.with_reaction_results(reactions, results)

    def _get_stages(
        self,
        construction_stages,
    ) -> typing.Iterator[tuple[int, ...]]:
        """
        Yield the parallelizable stages of construction.

        Yields:

            Vertices ids, which can be placed in parallel.

        """

        stages: tuple[list[int], ...] = tuple(
            [] for i in range(len(construction_stages) + 1)
        )
        vertices = flatten(self._building_block_vertices.values())
        for vertex in vertices:
            placed = False
            for i, stage in enumerate(construction_stages):
                if stage(vertex):
                    stages[i].append(vertex.get_id())
                    placed = True
                    break
            if not placed:
                stages[-1].append(vertex.get_id())
        yield from (tuple(stage) for stage in stages if stage)

    def __str__(self) -> str:
        return repr(self)

    def __repr__(self) -> str:
        raise NotImplementedError()
