"""Scales or otherwise modifies the particle velocities during simulation to
simulate coupling to an external heat bath with temperature T₀.
"""
import numpy as np
import logging
from mpi4py import MPI
from typing import Callable
from .logger import Logger
from .input_parser import Config


def cancel_com_momentum(velocities, config, comm=MPI.COMM_WORLD):
    com_velocity = comm.allreduce(np.sum(velocities[...], axis=0), MPI.SUM)
    velocities[...] = velocities[...] - com_velocity / config.n_particles
    return velocities


def generate_initial_velocities(velocities, config, prng, comm=MPI.COMM_WORLD):
    kT_start = config.gas_constant * config.start_temperature
    n_particles_ = velocities.shape[0]
    velocities[...] = prng.normal(
        loc=0, scale=kT_start / config.mass, size=(n_particles_, 3)
    )
    velocities = cancel_com_momentum(velocities, config, comm=comm)
    kinetic_energy = comm.allreduce(
        0.5 * config.mass * np.sum(velocities**2), MPI.SUM
    )
    start_kinetic_energy_target = (
        1.5 * config.gas_constant * config.n_particles * config.start_temperature
    )
    factor = np.sqrt(1.5 * config.n_particles * kT_start / kinetic_energy)
    velocities[...] = velocities[...] * factor
    kinetic_energy = comm.allreduce(
        0.5 * config.mass * np.sum(velocities**2), MPI.SUM
    )
    Logger.rank0.log(
        logging.INFO,
        (
            f"Initialized {config.n_particles} velocities, target kinetic "
            f"energy: {start_kinetic_energy_target}, actual kinetic energy "
            f"generated: {kinetic_energy}"
        ),
    )
    return velocities


def _random_gaussian(prng: np.random.Generator) -> float:
    """Draw a single random number from the standard normal distribution
    Generate a number from the Gaussian distribution centered at zero with unit
    standard deviation, :math:`N(0, 1)`.

    Parameters
    ----------
    prng : np.random.Generator
        Numpy object that provides a stream of random bits

    Returns
    -------
    float
        A random number drawn from :math:`N(0, 1)`.
    """
    return prng.normal()


def _random_chi_squared(prng: np.random.Generator, M: int) -> float:
    """Draw the sum of :code:`M` squared normally distributed values

    The value is generated by the Gamma distribution, in lieu of generating
    :code:`M` Gaussian distributed numbers and summing their squares.

    Parameters
    ----------
    prng : np.random.Generator
        Numpy object that provides a stream of random bits
    M : int
        Number of standard normally distributed numbers in the sum.

    Returns
    -------
    float
        The sum of :code:`M` squared normally distributed values centered at
        zero with unit standard deviation.

    Notes
    -----
    The sum of the squares of k independent standard normal random variables is
    distributed according to a :math:`\\chi^2` distribution.

    .. math::

        \\chi^2_1 + \\chi_2^2 + \\chi_3^2 + \\dots + \\chi_M^2
        \\sim
        \\sigma^2\\chi^2(k)

    This is a special case of the :math:`\\Gamma` distribution,
    :math:`\\Gamma(k/2, 2)`, and may be generated by

    .. math::

        \\chi^2(k) \\sim 2 \\Gamma(k/2, 2)

    References
    ----------
    Knuth, D.E. 1981, Seminumerical Algorithms, 2nd ed., vol. 2 of The Art of
    Computer Programming (Reading, MA: Addison-Wesley), pp. 120ff.
    J. H. Ahrens and U. Dieter, Computing 12 (1974), 223-246.
    """
    return prng.chisquare(M)


def csvr_thermostat(
    velocity: np.ndarray,
    names: np.ndarray,
    config: Config,
    prng: np.random.Generator,
    comm: MPI.Intracomm = MPI.COMM_WORLD,
    random_gaussian: Callable[[np.random.Generator], float] = _random_gaussian,
    random_chi_squared: Callable[
        [np.random.Generator, int, float], float
    ] = _random_chi_squared,
    remove_center_of_mass_momentum: bool = True,
) -> np.ndarray:
    """Canonical sampling through velocity rescaling thermostat

    Implements the CSVR thermostat. Rescales the system kinetic energy by a
    stochastically chosen factor to keep the temperature constant. Requires
    communcation of the kinetic energies calculated locally for each MPI rank.
    The random numbers sampled through :code:`random_gaussian` and
    :code:`random_chi_squared` are broadcast from the root rank to the other
    ranks to ensure the scaling is performed with the same stochasticity for
    all particles in the full system.

    The velocities are cleaned of center of mass momentum before the thermostat
    is applied, and the center of mass momentum is subsequently reapplied. This
    is performed for each thermostat coupling group, i.e. the center of mass
    momenta of each *group* is separately removed and reapplied after
    thermostatting.

    The implementation here is based on the derivation presented in the 2008
    Comput. Phys. Commun paper, not in the original 2007 J. Chem. Phys. paper.

    Parameters
    ----------
    velocity : (N, D) numpy.ndarray
        Array of velocities of N particles in D dimensions.
    names : (N,) numpy.ndarray
        Array of particle names.
    config : Config
        Configuration dataclass containing simulation metadata and parameters.
    prng : np.random.Generator
        Numpy object that provides a stream of random bits
    comm : MPI.Intracomm, optional
        MPI communicator to use for rank commuication. Defaults to
        MPI.COMM_WORLD.
    random_gaussian : callable, optional
        Function for generating standard normally distributed numbers.
    random_chi_squared : callable, optional
        Function for generating :math:`\\chi^2`-distributed numbers
    remove_center_of_mass_momentum : bool, optional
        If True, the center of mass of each coupling group is removed before
        the thermostat is applied. The center of mass momenta are added back
        after the kinetic energy rescaling is complete.

    See Also
    --------
    _random_gaussian :
        Used to sample Gaussian-distributed numbers.
    _random_chi_squared :
        Used to sample :math:`\\chi^2`-distributed numbers.
    hymd.input_parser.Config :
        Configuration dataclass handler.

    References
    ----------
    G. Bussi, D. Donadio, and M. Parrinello, J. Chem. Phys. 126, 014101 (2007).
    G. Bussi and M. Parrinello, Comput. Phys. Commun. 179, 26-29, (2008).
    """
    if not any(config.thermostat_coupling_groups):
        config.thermostat_coupling_groups = [config.unique_names.copy()]
    for i, group in enumerate(config.thermostat_coupling_groups):
        ind = np.where(
            np.logical_or.reduce(list(names == np.string_(t) for t in group))
        )
        group_n_particles = comm.allreduce(len(ind[0]), MPI.SUM)

        # Clean velocities of center of mass momentum
        if remove_center_of_mass_momentum and group_n_particles > 1:
            com_velocity = comm.allreduce(np.sum(velocity[ind], axis=0))
            velocity_clean = velocity[ind] - com_velocity / group_n_particles
            K = comm.allreduce(0.5 * config.mass * np.sum(velocity_clean[...] ** 2))
        else:
            K = comm.allreduce(0.5 * config.mass * np.sum(velocity[...] ** 2))
        K_target = (
            1.5 * config.gas_constant * group_n_particles * config.target_temperature
        )
        N_f = 3 * group_n_particles
        c = np.exp(-(config.time_step * config.respa_inner) / config.tau)

        # Random numbers are identical across ranks because all ranks have
        # the same prng after the initialization of the velocities
        R = random_gaussian(prng)
        SNf = random_chi_squared(prng, N_f - 1)

        alpha2 = (
            c
            + (1 - c) * (SNf + R**2) * K_target / (N_f * K)
            + 2 * R * np.sqrt(c * (1 - c) * K_target / (N_f * K))
        )
        dK = K * (alpha2 - 1)
        alpha = np.sqrt(alpha2)

        if remove_center_of_mass_momentum and group_n_particles > 1:
            # Assign velocities and reapply the previously removed center
            # of mass momentum removed
            velocity_clean *= alpha
            velocity[ind] = velocity_clean + com_velocity / group_n_particles
        else:
            velocity *= alpha
        config.thermostat_work += dK
