"""
Topology transformations
========================
.. moduleauthor:: Benjamin Ye <GitHub: @bbye98>

This module contains implementations of common LAMMPS topology
transformations, like the generation of initial particle positions.
"""

from io import TextIOWrapper
from numbers import Real
from typing import Any, Union

import numpy as np

from .. import FOUND_OPENMM, Q_, ureg
from ..algorithm import topology as topo

if FOUND_OPENMM:
    from openmm import app, unit


def create_atoms(
    dimensions: Union[np.ndarray[float], "unit.Quantity", Q_, "app.Topology"],
    N: int = None,
    N_p: int = 1,
    *,
    lattice: str = None,
    length: Union[float, "unit.Quantity"] = 0.34,
    flexible: bool = False,
    bonds: bool = False,
    angles: bool = False,
    dihedrals: bool = False,
    randomize: bool = False,
    length_unit: Union["unit.Unit", ureg.Unit] = None,
    wrap: bool = False,
) -> Any:
    """
    Generates initial particle positions for coarse-grained simulations.

    .. seealso::

       This function is an alias for
       :func:`mdcraft.algorithm.topology.create_atoms`.
    """

    return topo.create_atoms(
        dimensions,
        N,
        N_p,
        lattice=lattice,
        length=length,
        flexible=flexible,
        bonds=bonds,
        angles=angles,
        dihedrals=dihedrals,
        randomize=randomize,
        length_unit=length_unit,
        wrap=wrap,
    )


def write_data(
    file: Union[str, TextIOWrapper],
    positions: tuple[np.ndarray[float]],
    *,
    bonds: tuple[np.ndarray[int]] = None,
    angles: tuple[np.ndarray[int]] = None,
    dihedrals: tuple[np.ndarray[int]] = None,
    impropers: tuple[np.ndarray[int]] = None,
    dimensions: np.ndarray[float] = None,
    tilt: np.ndarray[float] = None,
    charges: np.ndarray[float] = None,
    masses: np.ndarray[float] = None,
) -> None:
    r"""
    Writes topological data to a LAMMPS data file in :code:`atom_style full`.

    Parameters
    ----------
    file : `str` or `_io.TextIOWrapper`
        LAMMPS data file.

    positions : `tuple`
        Atomic positions. Each element of the tuple should contain
        atoms of the same atom type.

        **Shape**: Tuple of arrays with shape :math:`(*,\,3)`.

        **Reference units**: :math:`\mathrm{Å}`.

    bonds : `tuple`, keyword-only, optional
        Pairs of indices of bonded atoms. Each element of the tuple
        should contain bonds of the same bond type.

        **Shape**: Tuple of arrays with shape :math:`(*,\,2)`.

    angles : `tuple`, keyword-only, optional
        Triples of indices of atoms that form an angle. Each element of
        the tuple should contain angles of the same angle type.

        **Shape**: Tuple of arrays with shape :math:`(*,\,3)`.

    dihedrals : `tuple`, keyword-only, optional
        Quadruples of indices of atoms that form a dihedral. Each
        element of the tuple should contain dihedrals of the same
        dihedral type.

        **Shape**: Tuple of arrays with shape :math:`(*,\,4)`.

    impropers : `tuple`, keyword-only, optional
        Quadruples of indices of atoms that form an improper. Each
        element of the tuple should contain impropers of the same
        improper type.

        **Shape**: Tuple of arrays with shape :math:`(*,\,4)`.

    dimensions : array-like, keyword-only, optional
        Box dimensions. If three values are provided, the box
        dimensions are assumed to be from :math:`0` to the specified
        values. If six values are provided, the box dimensions go from
        the three values in the first column to the three values in the
        second column.

        **Shape**: :math:`(3,)` or :math:`(3,\,2)`.

        **Reference units**: :math:`\mathrm{Å}`.

    tilt : array-like, keyword-only, optional
        Box :math:`xy`, :math:`xz`, and :math:`yz` tilt factors.

        **Shape**: :math:`(3,)`.

    charges : array-like, keyword-only, optional
        Atomic charges.

        **Shape**: :math:`(N,)`.

    masses : array-like, keyword-only, optional
        Atomic masses.

        **Shape**: :math:`(N,)`.
    """

    if isinstance(file, str):
        file = open(file, "w")

    # Write header
    file.write("LAMMPS Description\n\n")
    n_atoms_type = [len(p) for p in positions]
    n_atoms = sum(n_atoms_type)
    file.write(f"{n_atoms} atoms\n")
    file.write(f"{len(positions)} atom types\n")
    if bonds is not None:
        n_bonds_type = [len(b) for b in bonds]
        file.write(f"{sum(n_bonds_type)} bonds\n")
        file.write(f"{len(bonds)} bond types\n")
    if angles is not None:
        n_angles_type = [len(a) for a in angles]
        file.write(f"{sum(n_angles_type)} angles\n")
        file.write(f"{len(angles)} angle types\n")
    if dihedrals is not None:
        n_dihedrals_type = [len(d) for d in dihedrals]
        file.write(f"{sum(n_dihedrals_type)} dihedrals\n")
        file.write(f"{len(dihedrals)} dihedral types\n")
    if impropers is not None:
        n_impropers_type = [len(i) for i in impropers]
        file.write(f"{sum(n_impropers_type)} impropers\n")
        file.write(f"{len(impropers)} improper types\n")
    if dimensions is not None:
        if dimensions.ndim == 1:
            dimensions = np.vstack((np.zeros(3), dimensions)).T
        for i, d in enumerate(dimensions):
            a = chr(120 + i)
            file.write(f"{d[0]:.6g} {d[1]:.6g} {a}lo {a}hi\n")
    if tilt is not None:
        file.write(f"{tilt[0]:.6g} {tilt[1]:.6g} {tilt[2]:.6g} xy xz yz\n")

    # Write masses
    if masses is not None:
        if len(masses) != len(positions):
            emsg = "Number of masses must match number of atom types."
            raise ValueError(emsg)
        file.write("\nMasses\n\n")
        for i, m in enumerate(masses):
            file.write(f"{i + 1} {m:.6g}\n")

    # Write atom positions
    if charges is None:
        charges = np.zeros(n_atoms)
    if len(charges) == len(positions):
        charges = list(charges)
        for i, (qs, n) in enumerate(zip(charges, n_atoms_type)):
            if isinstance(qs, Real):
                charges[i] *= np.ones(n)
    elif len(charges) == n_atoms:
        charges = np.array_split(charges, np.cumsum(n_atoms)[:-1])
    else:
        raise ValueError("'charges' has an invalid shape.")
    file.write("\nAtoms # full\n\n")
    for t, (pos, qs) in enumerate(zip(positions, charges)):
        start = sum(n_atoms_type[:t])
        for i, (p, q) in enumerate(zip(pos, qs)):
            file.write(
                f"{start + i + 1} {start + i + 1} {t + 1} {q:.6g} "
                f"{p[0]:.6g} {p[1]:.6g} {p[2]:.6g}\n"
            )

    # Write bonds
    if bonds is not None:
        file.write("\nBonds\n\n")
        for t, b in enumerate(bonds):
            start = sum(n_bonds_type[:t])
            for i, (a, b) in enumerate(b):
                file.write(f"{start + i + 1} {t + 1} {a} {b}\n")

    # Write angles
    if angles is not None:
        file.write("\nAngles\n\n")
        for t, a in enumerate(angles):
            start = sum(n_angles_type[:t])
            for i, (a, b, c) in enumerate(a):
                file.write(f"{start + i + 1} {t + 1} {a} {b} {c}\n")

    # Write dihedrals
    if dihedrals is not None:
        file.write("\nDihedrals\n\n")
        for t, d in enumerate(dihedrals):
            start = sum(n_dihedrals_type[:t])
            for i, (a, b, c, d) in enumerate(d):
                file.write(f"{start + i + 1} {t + 1} {a} {b} {c} {d}\n")

    # Write impropers
    if impropers is not None:
        file.write("\nImpropers\n\n")
        for t, i in enumerate(impropers):
            start = sum(n_impropers_type[:t])
            for j, (a, b, c, d) in enumerate(i):
                file.write(f"{start + j + 1} {t + 1} {a} {b} {c} {d}\n")

    file.close()
