import typing
from collections import abc

import numpy as np

from stk._internal.building_block import BuildingBlock
from stk._internal.construction_state.construction_state import (
    ConstructionState,
)
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.topology_graphs.topology_graph.topology_graph import (
    TopologyGraph,
)
from stk._internal.topology_graphs.vertex import Vertex

from .vertices import AxleVertex, CycleVertex


class NRotaxane(TopologyGraph):
    """
    Represents [n]rotaxane topology graphs.

    Axle and cycle building blocks do not require functional groups.

    This class assumes one axle with (n-1) macrocycles threaded on it.
    The macrocycles are spaced evenly along the axle in repeating
    patterns. The threaded macrocycles can be described analogously
    to monomers in linear polymers, in terms of a repeating unit,
    except that no bonds are formed between them.

    Examples:

        *Construction*

        .. testcode:: construction

            import stk

            cycle = stk.ConstructedMolecule(
                topology_graph=stk.macrocycle.Macrocycle(
                    building_blocks=(
                        stk.BuildingBlock(
                            smiles='[Br]CC[Br]',
                            functional_groups=[stk.BromoFactory()],
                        ),
                    ),
                    repeating_unit='A',
                    num_repeating_units=5,
                ),
            )
            axle = stk.ConstructedMolecule(
                topology_graph=stk.polymer.Linear(
                    building_blocks=(
                        stk.BuildingBlock('BrCCBr', [stk.BromoFactory()]),
                        stk.BuildingBlock('BrCNCBr', [stk.BromoFactory()]),
                    ),
                    repeating_unit='AB',
                    num_repeating_units=7,
                )
            )
            rotaxane = stk.ConstructedMolecule(
                topology_graph=stk.rotaxane.NRotaxane(
                    axle=stk.BuildingBlock.init_from_molecule(axle),
                    cycles=(
                        stk.BuildingBlock.init_from_molecule(cycle),
                    ),
                    repeating_unit='A',
                    num_repeating_units=3,
                ),
            )

        .. moldoc::

            import moldoc.molecule as molecule
            import stk

            cycle = stk.ConstructedMolecule(
                topology_graph=stk.macrocycle.Macrocycle(
                    building_blocks=(
                        stk.BuildingBlock(
                            smiles='[Br]CC[Br]',
                            functional_groups=[stk.BromoFactory()],
                        ),
                    ),
                    repeating_unit='A',
                    num_repeating_units=5,
                ),
            )
            axle = stk.ConstructedMolecule(
                topology_graph=stk.polymer.Linear(
                    building_blocks=(
                        stk.BuildingBlock('BrCCBr', [stk.BromoFactory()]),
                        stk.BuildingBlock('BrCNCBr', [stk.BromoFactory()]),
                    ),
                    repeating_unit='AB',
                    num_repeating_units=7,
                )
            )
            rotaxane = stk.ConstructedMolecule(
                topology_graph=stk.rotaxane.NRotaxane(
                    axle=stk.BuildingBlock.init_from_molecule(axle),
                    cycles=(
                        stk.BuildingBlock.init_from_molecule(cycle),
                    ),
                    repeating_unit='A',
                    num_repeating_units=3,
                ),
            )

            moldoc_display_molecule = molecule.Molecule(
                atoms=(
                    molecule.Atom(
                        atomic_number=atom.get_atomic_number(),
                        position=position,
                    ) for atom, position in zip(
                        rotaxane.get_atoms(),
                        rotaxane.get_position_matrix(),
                    )
                ),
                bonds=(
                    molecule.Bond(
                        atom1_id=bond.get_atom1().get_id(),
                        atom2_id=bond.get_atom2().get_id(),
                        order=bond.get_order(),
                    ) for bond in rotaxane.get_bonds()
                ),
            )

        *Suggested Optimization*

        For :class:`.NRotaxane` topologies, there is no need to use an
        optimizer, so stick with :class:`.NullOptimizer`. However, it is
        recommended that all building blocks be optimized prior to
        construction.

        *Defining the Orientation of Each Building Block*

        The `orientations` parameter allows the direction of each cycle
        along the axle to be flipped

        .. testcode:: defining-the-orientation-of-each-building-block

            import stk

            cycle1 = stk.ConstructedMolecule(
                topology_graph=stk.macrocycle.Macrocycle(
                    building_blocks=(
                        stk.BuildingBlock(
                            smiles='BrCCBr',
                            functional_groups=[stk.BromoFactory()],
                        ),
                        stk.BuildingBlock(
                            smiles='BrCNCBr',
                            functional_groups=[stk.BromoFactory()],
                        ),
                    ),
                    repeating_unit='AB',
                    num_repeating_units=5,
                    optimizer=stk.MCHammer(),
                ),
            )
            cycle2 = stk.ConstructedMolecule(
                topology_graph=stk.macrocycle.Macrocycle(
                    building_blocks=(
                        stk.BuildingBlock(
                            smiles='BrCCCBr',
                            functional_groups=[stk.BromoFactory()],
                        ),
                        stk.BuildingBlock(
                            smiles='BrCCNCBr',
                            functional_groups=[stk.BromoFactory()],
                        ),
                    ),
                    repeating_unit='AB',
                    num_repeating_units=5,
                    optimizer=stk.MCHammer(),
                ),
            )
            axle = stk.ConstructedMolecule(
                topology_graph=stk.polymer.Linear(
                    building_blocks=(
                        stk.BuildingBlock('BrCCBr', [stk.BromoFactory()]),
                    ),
                    repeating_unit='A',
                    num_repeating_units=2,
                ),
            )
            rotaxane1 = stk.ConstructedMolecule(
                topology_graph=stk.rotaxane.NRotaxane(
                    axle=stk.BuildingBlock.init_from_molecule(axle),
                    cycles=(
                        stk.BuildingBlock.init_from_molecule(cycle1),
                        stk.BuildingBlock.init_from_molecule(cycle2),
                    ),
                    repeating_unit='AB',
                    num_repeating_units=3,
                    orientations=(1., 0.5),
                ),
            )

        In the above example, ``cycle1`` is guaranteed to be flipped,
        ``cycle2`` has a 50% chance of being flipped, each time it is
        placed on a node.

        Note that whether a building block will be flipped or not
        is decided during the initialization of :class:`.NRotaxane`

        .. testcode:: defining-the-orientation-of-each-building-block

            # graph1 will always construct the same [n]rotaxane.
            graph1 = stk.rotaxane.NRotaxane(
                axle=stk.BuildingBlock.init_from_molecule(axle),
                cycles=(
                    stk.BuildingBlock.init_from_molecule(cycle1),
                    stk.BuildingBlock.init_from_molecule(cycle2),
                ),
                repeating_unit='AB',
                num_repeating_units=5,
                orientations=(0.65, 0.45),
            )
            # rotaxane2 and rotaxane3 are guaranteed to be the same as they
            # used the same topology graph.
            rotaxane2 = stk.ConstructedMolecule(graph1)
            rotaxane3 = stk.ConstructedMolecule(graph1)

            # graph2 may lead to a different [n]rotaxane, despite
            # being initialized with the same parameters.
            graph2 = stk.rotaxane.NRotaxane(
                axle=stk.BuildingBlock.init_from_molecule(axle),
                cycles=(
                    stk.BuildingBlock.init_from_molecule(cycle1),
                    stk.BuildingBlock.init_from_molecule(cycle2),
                ),
                repeating_unit='AB',
                num_repeating_units=5,
                orientations=(0.65, 0.45)
            )

            # rotaxane4 and rotaxane5 are guaranteed to be the same because
            # they used the same topology graph. However, they may be
            # different to rotaxane2 and rotaxane3.
            rotaxane4 = stk.ConstructedMolecule(graph2)
            rotaxane5 = stk.ConstructedMolecule(graph2)

        The `random_seed` parameter can be used to get reproducible results

        .. testcode:: defining-the-orientation-of-each-building-block

            # rotaxane6 and rotaxane7 are guaranteed to be the same,
            # because graph3 and graph4 used the same random seed.

            graph3 = stk.rotaxane.NRotaxane(
                axle=stk.BuildingBlock.init_from_molecule(axle),
                cycles=(
                    stk.BuildingBlock.init_from_molecule(cycle1),
                    stk.BuildingBlock.init_from_molecule(cycle2),
                ),
                repeating_unit='AB',
                num_repeating_units=5,
                orientations=(0.65, 0.45),
                random_seed=4,
            )
            rotaxane6 = stk.ConstructedMolecule(graph3)

            graph4 = stk.rotaxane.NRotaxane(
                axle=stk.BuildingBlock.init_from_molecule(axle),
                cycles=(
                    stk.BuildingBlock.init_from_molecule(cycle1),
                    stk.BuildingBlock.init_from_molecule(cycle2),
                ),
                repeating_unit='AB',
                num_repeating_units=5,
                orientations=(0.65, 0.45),
                random_seed=4,
            )
            rotaxane7 = stk.ConstructedMolecule(graph4)

        *Using Numbers to Define the Repeating Unit*

        The repeating unit can also be specified through the indices of
        the building blocks

        .. testcode:: using-numbers-to-define-the-repeating-unit

            import stk

            cycle1 = stk.ConstructedMolecule(
                topology_graph=stk.macrocycle.Macrocycle(
                    building_blocks=(
                        stk.BuildingBlock(
                            smiles='BrCCBr',
                            functional_groups=[stk.BromoFactory()],
                        ),
                        stk.BuildingBlock(
                            smiles='BrCNCBr',
                            functional_groups=[stk.BromoFactory()],
                        ),
                    ),
                    repeating_unit='AB',
                    num_repeating_units=5,
                    optimizer=stk.MCHammer(),
                ),
            )
            cycle2 = stk.ConstructedMolecule(
                topology_graph=stk.macrocycle.Macrocycle(
                    building_blocks=(
                        stk.BuildingBlock(
                            smiles='BrCCCBr',
                            functional_groups=[stk.BromoFactory()],
                        ),
                        stk.BuildingBlock(
                            smiles='BrCCNCBr',
                            functional_groups=[stk.BromoFactory()],
                        ),
                    ),
                    repeating_unit='AB',
                    num_repeating_units=5,
                    optimizer=stk.MCHammer(),
                ),
            )
            cycle3 = stk.ConstructedMolecule(
                topology_graph=stk.macrocycle.Macrocycle(
                    building_blocks=(
                        stk.BuildingBlock(
                            smiles='BrCCNNCBr',
                            functional_groups=[stk.BromoFactory()],
                        ),
                        stk.BuildingBlock(
                            smiles='BrCCNNNCBr',
                            functional_groups=[stk.BromoFactory()],
                        ),
                    ),
                    repeating_unit='AB',
                    num_repeating_units=5,
                    optimizer=stk.MCHammer(),
                ),
            )
            axle = stk.ConstructedMolecule(
                topology_graph=stk.polymer.Linear(
                    building_blocks=(
                        stk.BuildingBlock('BrCCBr', [stk.BromoFactory()]),
                    ),
                    repeating_unit='A',
                    num_repeating_units=2,
                ),
            )

            # r1 and r2 are different ways to write the same thing.
            r1 = stk.ConstructedMolecule(
                topology_graph=stk.rotaxane.NRotaxane(
                    axle=stk.BuildingBlock.init_from_molecule(axle),
                    cycles=(
                        stk.BuildingBlock.init_from_molecule(cycle1),
                        stk.BuildingBlock.init_from_molecule(cycle2),
                        stk.BuildingBlock.init_from_molecule(cycle3),
                    ),
                    repeating_unit='ACB',
                    num_repeating_units=3,
                ),
            )
            r2 = stk.ConstructedMolecule(
                topology_graph=stk.rotaxane.NRotaxane(
                    axle=stk.BuildingBlock.init_from_molecule(axle),
                    cycles=(
                        stk.BuildingBlock.init_from_molecule(cycle1),
                        stk.BuildingBlock.init_from_molecule(cycle2),
                        stk.BuildingBlock.init_from_molecule(cycle3),
                    ),
                    repeating_unit=(0, 2, 1),
                    num_repeating_units=3,
                ),
            )

    """

    def __init__(
        self,
        axle: BuildingBlock,
        cycles: abc.Iterable[BuildingBlock],
        repeating_unit: str | abc.Iterable[int],
        num_repeating_units: int,
        orientations: abc.Iterable[float] | None = None,
        random_seed: int | np.random.Generator | None = None,
        num_processes: int = 1,
        optimizer: Optimizer = NullOptimizer(),
    ) -> None:
        """
        Parameters:

            axle:
                The axle of the rotaxane.

            cycles (list[BuildingBlock]):
                The cycles threaded onto the `axle`.

            repeating_unit (str | list[int]):
                A string specifying the repeating unit of the `cycles`.
                For example, ``'AB'`` or ``'ABB'``. The first cycle in
                `cycles` is ``'A'`` and so on.

                The repeating unit can also be specified by the indices of
                `cycles`, for example ``'ABB'`` can be
                written as ``[0, 1, 1]``.

            num_repeating_units:
                The number of repeating units threaded along the axle.

            orientations (list[float] | None):
                For each character in the repeating unit, a value
                between ``0`` and ``1`` (both inclusive) must be given in
                a :class:`tuple`. It indicates the probability that each
                cycle will have its orientation along the axle
                flipped. If ``0`` then the cycle is guaranteed not to
                flip. If ``1`` it is guaranteed to flip. This allows the
                user to create head-to-head or head-to-tail chains, as well
                as chain with a preference for head-to-head or head-to-tail
                if a number between ``0`` and ``1`` is chosen. If
                ``None``, then defaults to ``0`` in every case.

                It is also possible to supply an orientation for every
                cycle vertex in the final topology graph. In this case, the
                length of `orientations` must be equal to
                ``len(repeating_unit)*num_repeating_units``.

            random_seed:
                The random seed to use when choosing random orientations.

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

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

        Raises:

            :class:`ValueError`
                If the length of `orientations` is not equal in length to
                `repeating_unit` or to the total number of vertices.

        """

        if not isinstance(repeating_unit, str):
            repeating_unit = tuple(repeating_unit)

        if orientations is None:
            orientations = tuple(0.0 for _ in range(len(repeating_unit)))
        else:
            orientations = tuple(orientations)

        if len(orientations) == len(repeating_unit):
            orientations = orientations * num_repeating_units

        chain_length = len(repeating_unit) * num_repeating_units
        if len(orientations) != chain_length:
            raise ValueError(
                "The length of orientations must match either "
                "the length of repeating_unit or the "
                "total number of vertices."
            )

        if random_seed is None or isinstance(random_seed, int):
            random_seed = np.random.default_rng(random_seed)

        self._repeating_unit = self._normalize_repeating_unit(
            repeating_unit=repeating_unit,
        )
        self._num_repeating_units = num_repeating_units

        vertices: list[AxleVertex | CycleVertex] = [AxleVertex(0, (0, 0, 0))]
        distance = 1 / (chain_length + 1)
        choices = [True, False]
        for vertex_id, flip_chance in enumerate(orientations, 1):
            vertices.append(
                CycleVertex(
                    id=vertex_id,
                    position=[vertex_id * distance - 0.5, 0, 0],
                    flip=random_seed.choice(
                        choices,
                        p=[flip_chance, 1 - flip_chance],
                    ),
                )
            )

        # Save the chosen orientations for __repr__.
        self._orientations = tuple(
            int(vertex.get_flip())  # type: ignore[union-attr]
            for vertex in vertices[1:]
        )

        super().__init__(
            building_block_vertices=self._get_building_block_vertices(
                axle=axle,
                cycles=tuple(cycles),
                vertices=tuple(vertices),
            ),
            edges=(),
            reaction_factory=GenericReactionFactory(),
            construction_stages=(),
            num_processes=num_processes,
            optimizer=optimizer,
            edge_groups=None,
        )

    def clone(self) -> typing.Self:
        clone = super().clone()
        clone._repeating_unit = self._repeating_unit
        clone._num_repeating_units = self._num_repeating_units
        clone._orientations = self._orientations
        return clone

    @staticmethod
    def _normalize_repeating_unit(
        repeating_unit: str | tuple[int, ...],
    ) -> tuple[int, ...]:
        if isinstance(repeating_unit, tuple):
            return repeating_unit
        base = ord("A")
        return tuple(ord(letter) - base for letter in repeating_unit)

    def _get_building_block_vertices(
        self,
        axle: BuildingBlock,
        cycles: abc.Sequence[BuildingBlock],
        vertices: abc.Sequence[Vertex],
    ) -> dict[BuildingBlock, abc.Sequence[Vertex]]:
        threads = self._repeating_unit * self._num_repeating_units
        building_block_vertices: dict[BuildingBlock, list[Vertex]] = {}

        building_block_vertices[axle] = [vertices[0]]
        for cycle_index, vertex in zip(threads, vertices[1:]):
            bb = cycles[cycle_index]
            building_block_vertices[bb] = building_block_vertices.get(bb, [])
            building_block_vertices[bb].append(vertex)

        return typing.cast(
            dict[BuildingBlock, abc.Sequence[Vertex]],
            building_block_vertices,
        )

    def _run_reactions(self, state: ConstructionState) -> ConstructionState:
        return state

    @staticmethod
    def _get_scale(
        building_block_vertices: dict[BuildingBlock, abc.Sequence[Vertex]],
        scale_multiplier: float,
    ) -> float:
        axle = next(iter(building_block_vertices))
        return 0.8 * axle.get_maximum_diameter()

    def __repr__(self):
        return (
            f"rotaxane.NRotaxane("
            f"{self._repeating_unit!r}, "
            f"{self._num_repeating_units}, "
            f"{self._orientations!r}"
            f")"
        )
