# Direct force summation N-body code
import numpy as nu
from numpy import linalg

from galpy.potential.linearPotential import evaluatelinearForces
from galpy.potential.planarPotential import (
    evaluateplanarphitorques,
    evaluateplanarRforces,
)
from galpy.potential.Potential import (
    evaluatephitorques,
    evaluateRforces,
    evaluatezforces,
)
from galpy.util import symplecticode


def direct_nbody(
    q,
    p,
    m,
    t,
    pot=None,
    softening_model="plummer",
    softening_length=None,
    atol=None,
    rtol=None,
):
    """
    NAME:
       direct_nbody
    PURPOSE:
       N-body code using direct summation for force evaluation
    INPUT:
       q - list of initial positions (numpy.ndarrays)
       p - list of initial momenta (numpy.ndarrays)
       m - list of masses
       t - times at which output is desired
       pot= external potential (galpy.potential or list of galpy.potentials)
       softening_model=  type of softening to use ('plummer')
       softening_length= (optional)
    OUTPUT:
       list of [q,p] at times t
    HISTORY:
       2011-02-03 - Written - Bovy (NYU)
    """
    # Set up everything
    if softening_model.lower() == "plummer":
        softening = _plummer_soft
    out = []
    out.append([q, p])
    qo = q
    po = p
    # Determine appropriate stepsize
    dt = t[1] - t[0]
    ndt = 1
    to = t[0]
    # determine appropriate softening length if not given
    if softening_length is None:
        softening_length = 0.01
    # Run simulation
    for ii in range(1, len(t)):
        print(ii)
        for jj in range(ndt):  # loop over number of sub-intervals
            (qo, po) = _direct_nbody_step(
                qo, po, m, to, dt, pot, softening, (softening_length,)
            )
            # print(qo)
            to += dt
        out.append([qo, po])
    # Return output
    return out


def _direct_nbody_step(q, p, m, t, dt, pot, softening, softening_args):
    """One N-body step: drift-kick-drift"""
    # drift
    q12 = [symplecticode.leapfrog_leapq(q[ii], p[ii], dt / 2.0) for ii in range(len(q))]
    # kick
    force = _direct_nbody_force(q12, m, t + dt / 2.0, pot, softening, softening_args)
    # print(force)
    p = [symplecticode.leapfrog_leapp(p[ii], dt, force[ii]) for ii in range(len(p))]
    # drift
    q = [
        symplecticode.leapfrog_leapq(q12[ii], p[ii], dt / 2.0) for ii in range(len(q12))
    ]
    return (q, p)


def _direct_nbody_force(q, m, t, pot, softening, softening_args):
    """Calculate the force"""
    # First do the particles
    # Calculate all the distances
    nq = len(q)
    dim = len(q[0])
    dist_vec = nu.zeros((nq, nq, dim))
    dist = nu.zeros((nq, nq))
    for ii in range(nq):
        for jj in range(ii + 1, nq):
            dist_vec[ii, jj, :] = q[jj] - q[ii]
            dist_vec[jj, ii, :] = -dist_vec[ii, jj, :]
            dist[ii, jj] = linalg.norm(dist_vec[ii, jj, :])
            dist[jj, ii] = dist[ii, jj]
    # Calculate all the forces
    force = []
    for ii in range(nq):
        thisforce = nu.zeros(dim)
        for jj in range(nq):
            if ii == jj:
                continue
            thisforce += (
                m[jj]
                * softening(dist[ii, jj], *softening_args)
                / dist[ii, jj]
                * dist_vec[ii, jj, :]
            )
        force.append(thisforce)
    # Then add the external force
    if pot is None:
        return force
    for ii in range(nq):
        force[ii] += _external_force(q[ii], t, pot)
    return force


def _external_force(x, t, pot):
    dim = len(x)
    if dim == 3:
        # x is rectangular so calculate R and phi
        R = nu.sqrt(x[0] ** 2.0 + x[1] ** 2.0)
        phi = nu.arccos(x[0] / R)
        sinphi = x[1] / R
        cosphi = x[0] / R
        if x[1] < 0.0:
            phi = 2.0 * nu.pi - phi
        # calculate forces
        Rforce = evaluateRforces(R, x[2], pot, phi=phi, t=t)
        phitorque = evaluatephitorques(R, x[2], pot, phi=phi, t=t)
        return nu.array(
            [
                cosphi * Rforce - 1.0 / R * sinphi * phitorque,
                sinphi * Rforce + 1.0 / R * cosphi * phitorque,
                evaluatezforces(R, x[2], pot, phi=phi, t=t),
            ]
        )
    elif dim == 2:
        # x is rectangular so calculate R and phi
        R = nu.sqrt(x[0] ** 2.0 + x[1] ** 2.0)
        phi = nu.arccos(x[0] / R)
        sinphi = x[1] / R
        cosphi = x[0] / R
        if x[1] < 0.0:
            phi = 2.0 * nu.pi - phi
        # calculate forces
        Rforce = evaluateplanarRforces(R, pot, phi=phi, t=t)
        phitorque = evaluateplanarphitorques(R, pot, phi=phi, t=t)
        return nu.array(
            [
                cosphi * Rforce - 1.0 / R * sinphi * phitorque,
                sinphi * Rforce + 1.0 / R * cosphi * phitorque,
            ]
        )
    elif dim == 1:
        return evaluatelinearForces(x, pot, t=t)


def _plummer_soft(d, eps):
    return d / (d**2.0 + eps**2.0) ** 1.5
