"""
PhononHarmonics copied from ASE.
"""

import numpy as np
from ase.units import kB


def generate_phonon_rattled_structures(atoms, fc2, n_structures, temperature):
    """
    Returns list of phonon-rattled configurations.

    Configurations are generated by superimposing harmonic phonon
    eigenmodes with amplitudes corresponding to a certain temperatures
    and randomly drawn phase factors.

    Let :math:`\\boldsymbol{X}_{ai}` be the phonon modes indexed by atom
    :math:`a` and mode :math:`i`, :math:`\\omega_i` the phonon
    frequencies, and let :math:`0 < Q_i \\leq 1` and :math:`0 \\leq
    U_i < 1` be uniformly random numbers.  Then

    .. math::

        \\boldsymbol{R}_a
        = \\boldsymbol{R}^0_a
        + \\left<\\frac{k_B T}{m_a} \\right>^{1/2}
        \\sum_i \\frac{1}{\\omega_i} \\boldsymbol{X}_{ai}
        \\sqrt{-2 \\ln Q_i} \\cos(\\pi \\omega_i U_i)

    Parameters
    ----------
    atoms : ase.Atoms
        prototype structure
    fc2 : numpy.ndarray
        second order force constant matrix, with shape `(3N, 3N)`
    n_structures : int
        number of structures to generate
    temperature : float
        temperature in Kelvin

    Returns
    -------
    list(ase.Atoms)
        generated structures
    """
    if fc2.shape != (3 * len(atoms), 3 * len(atoms)):
        raise ValueError
    structures = []
    for _ in range(n_structures):
        atoms_tmp = atoms.copy()
        PhononHarmonics(atoms_tmp, temperature*kB, fc2)
        structures.append(atoms_tmp)
    return structures


def PhononHarmonics(atoms, temp, force_constants, rng=np.random,
                    failfast=True):
    """
    Excites phonon modes to specified temperature.

    This excites all phonon modes randomly so that each contributes,
    on average, equally to the given temperature.  Both potential
    energy and kinetic energy will be consistent with the phononic
    vibrations characteristic of the specified temperature.

    Parameters
    ----------
    atoms : ase.Atoms
        prototype structure
    temp : float
        temperature in eV
    force_constants : numpy.ndarray
        second order force constant matrix, with shape `(3N, 3N)`
    rng : func
        random number generation
    failfast : True

    Raises
    ------
    ValueError
        if egenvalues of the dynamical matrix look suspicious (they
        should be three zeros for translational modes followed by any
        number of strictly positive eigenvalues).
    """
    mass_a = atoms.get_masses()

    rminv = (mass_a**-0.5).repeat(3)
    dynamical_matrix = force_constants.copy()
    dynamical_matrix *= rminv[:, None]
    dynamical_matrix *= rminv[None, :]

    w2_s, X_is = np.linalg.eigh(dynamical_matrix)

    # set imaginary frequencies to positive
    w2_s = np.sort(np.abs(w2_s))

    if failfast:
        zeros = w2_s[:3]
        worst_zero = np.abs(zeros).max()
        if worst_zero > 1e-3:
            raise ValueError('Translational modes have suspiciously large '
                             'energies; should be close to zero: {}'
                             .format(w2_s[:3]))

        w2min = w2_s[3:].min()
        if w2min < 0:
            raise ValueError('Dynamical matrix has negative eigenvalues '
                             'such as {}'.format(w2min))

    # First three modes are translational so ignore:
    nw = len(w2_s) - 3
    w_s = np.sqrt(w2_s[3:])
    X_acs = X_is[:, 3:].reshape(len(atoms), 3, nw)

    # We need 0 < P <= 1.0 and not 0 0 <= P < 1.0 for the logarithm
    # to avoid (highly improbable) NaN:
    A_s = np.sqrt(-2.0 * np.log(1.0 - rng.rand(nw)))
    phi_s = 2.0 * np.pi * rng.rand(nw)

    d_ac = (A_s * np.sin(phi_s) / w_s * X_acs).sum(axis=2)
    v_ac = (A_s * np.cos(phi_s) * X_acs).sum(axis=2)

    sqrtkTm_a = np.sqrt(temp / mass_a)
    for quantity in [d_ac, v_ac]:
        quantity *= sqrtkTm_a[:, None]

    atoms.positions += d_ac
    atoms.set_velocities(v_ac)
