"""Module implements the `State` class.

`State` stores and exposes a parent `hoomd.Simulation` object's data (e.g.
particle positions, system bonds).

"""
from collections import defaultdict

from . import _hoomd
from hoomd.box import Box
from hoomd.snapshot import Snapshot
from hoomd.data import LocalSnapshot, LocalSnapshotGPU
import hoomd
import math
import collections.abc


def _create_domain_decomposition(device, box, domain_decomposition):
    """Create the domain decomposition.

    Args:
        device (Device): The simulation's device
        box: The C++ global box object for the state being initialized
        domain_decomposition: See Simulation.create_state_from_* for a
          description.
    """
    if (not isinstance(domain_decomposition, collections.abc.Sequence)
            or len(domain_decomposition) != 3):
        raise TypeError("domain_decomposition must be a length 3 sequence")

    initialize_grid = False
    initialize_fractions = False

    for v in domain_decomposition:
        if v is not None:
            if isinstance(v, int):
                initialize_grid = True
            elif isinstance(v, collections.abc.Sequence):
                if not math.isclose(sum(v), 1.0, rel_tol=1e-6):
                    raise ValueError("Rank fractions must sum to 1.0.")
                initialize_fractions = True
            else:
                raise TypeError("Invalid type in domain_decomposition.")

    if initialize_grid and initialize_fractions:
        raise ValueError("Domain decomposition mixes integers and sequences.")

    if not hoomd.version.mpi_enabled:
        return None

    # if we are only running on one processor, we use optimized code paths
    # for single-GPU execution
    if device.communicator.num_ranks == 1:
        return None

    if initialize_fractions:
        fractions = [
            v[:-1] if v is not None else [] for v in domain_decomposition
        ]
        result = _hoomd.DomainDecomposition(device._cpp_exec_conf, box.getL(),
                                            *fractions)
    else:
        grid = [v if v is not None else 0 for v in domain_decomposition]
        result = _hoomd.DomainDecomposition(device._cpp_exec_conf, box.getL(),
                                            *grid, False)

    return result


class State:
    """The state of a `hoomd.Simulation` object.

    Note:
        This object cannot be directly instantiated. Use
        `hoomd.Simulation.create_state_from_gsd` and
        `hoomd.Simulation.create_state_from_snapshot` to instantiate a `State`
        object as part of a simulation.

    .. rubric:: Overview

    `State` stores the data that describes the thermodynamic microstate of a
    `hoomd.Simulation` object. This data consists of the box, particles, bonds,
    angles, dihedrals, impropers, special pairs, and constraints.

    .. rubric:: Box

    The simulation `box` describes the area in space that contains the particles
    and the periodic boundary conditions to apply. See `Box` for a full
    description.

    .. rubric:: Particles

    The state contains `N_particles` particles. Each particle has a position,
    orientation, type id, body, mass, moment of inertia, charge, diameter,
    velocity, angular momentum, image, and tag:

    - :math:`\\vec{r}`: position :math:`[\\mathrm{length}]` - x,y,z cartesian
      coordinates defining the position of the particle in the box
    - :math:`\\mathbf{q}`: orientation :math:`[\\mathrm{dimensionless}]` - (s,
      :math:`\\vec{a}`), unit quaternion defining the rotation from the
      particle's local reference frame to the box reference frame.
    - :math:`t_\\mathrm{id}`: type id :math:`[\\mathrm{dimensionless}]` -
      integer in the range ``[0,len(particle_types)``) that identifies the
      particle's type. `particle_types` maps type ids to names with:
      ``name = particle_types[t_id]``.
    - :math:`b_\\mathrm{id}`: body id :math:`[\\mathrm{dimensionless}]` -
      integer that identifies the particle's rigid body. A value of ``-1``
      indicates that this particle does not belong to a body. A positive value
      indicates that the particle belongs to the body :math:`b_\\mathrm{id}`.
      This particle is the central particle of a body when the body id is equal
      to the tag :math:`b_\\mathrm{id} = p_\\mathrm{tag}`. (used by
      `md.constrain.Rigid`)
    - :math:`m`: mass :math:`[\\mathrm{mass}]` - the particle's mass.
    - :math:`I`: moment of inertia :math:`[\\mathrm{mass} \\cdot
      \\mathrm{length}^2]` - :math:`I_{xx}`, :math:`I_{yy}`, :math:`I_{zz}`
      elements of the diagonal moment of inertia tensor in the particle's local
      reference frame. The off-diagonal elements are 0.
    - :math:`q`: charge :math:`[\\mathrm{charge}]`
    - :math:`d`: diameter :math:`[\\mathrm{length}]` - deprecated in v3.0.0.
      HOOMD-blue reads and writes particle diameters, but does not use them in
      any computations. As of the current beta release, diameter is still used
      in `md.pair.SLJ` and `md.pair.DLVO`.
    - :math:`\\vec{v}``: velocity :math:`[\\mathrm{velocity}]` - x,y,z
      components of the particle's velocity in the box's reference frame.
    - :math:`\\mathbf{P_S}``: angular momentum :math:`[\\mathrm{mass} \\cdot
      \\mathrm{velocity} \\cdot \\mathrm{length}]` - in a quaternion
      representation (see note).
    - :math:`\\vec{n}` : image :math:`[\\mathrm{dimensionless}]` - integers
      x,y,z that record how many times the particle has crossed each of the
      periodic box boundaries.
    - :math:`p_\\mathrm{tag}`` : tag :math:`[\\mathrm{dimensionless}]` -
      integer that uniquely identifies a given particle. The particles are in
      tag order when writing and initializing to/from a GSD file or snapshot:
      :math:`p_\\mathrm{tag,i} = i`. When accessing data in local snapshots,
      particles may be in any order.

    Note:
        HOOMD stores angular momentum as a quaternion because that is the form
        used when integrating the equations of motion (see `Kamberaj 2005`_).
        The angular momentum quaternion :math:`\\mathbf{P_S}` is defined with
        respect to the orientation quaternion of the particle
        :math:`\\mathbf{q}` and the vector angular momentum of the particle,
        lifted into pure imaginary quaternion form :math:`\\mathbf{S}^{(4)}` as:

        .. math::

            \\mathbf{P_S} = 2 \\mathbf{q} \\times \\mathbf{S}^{(4)}

        . Following this, the angular momentum vector :math:`\\vec{S}` in the
        particle's local reference frame is:

        .. math::

            \\vec{S} = \\frac{1}{2}im(\\mathbf{q}^* \\times \\mathbf{P_S})

    .. rubric:: Bonds

    The state contains `N_bonds` bonds, `N_angles` angles, `N_dihedrals`
    dihedrals, `N_impropers` impropers, and `N_special_pairs` special pairs.
    Each of these data structures is similar, differing in the number of
    particles in the group and what operations use them. Bonds, angles,
    dihedrals, and impropers contain 2, 3, 4, and 4 particles per group
    respectively. Bonds specify the toplogy used when computing energies and
    forces in `md.bond`, angles define the same for `md.angle`, dihedrals for
    `md.dihedral` and impropers for `md.improper`. These collectively implement
    bonding potentials used in molecular dynamics force fields. Like bonds,
    special pairs define connections between two particles, but special pairs
    are intended to adjust the 1-4 pairwise interactions in some molecular
    dynamics force fields: see `md.special_pair`. Each bonded group is defined
    by a type id, the group members, and a tag.

    - :math:`t_\\mathrm{id}`: type id :math:`[\\mathrm{dimensionless}]` -
      integer in the range ``[0,len(bond_types)``) that identifies the
      bond's type. `bond_types` maps type ids to names with:
      ``name = bond_types[t_id]``. Similarly, `angle_types` lists the