#!/usr/bin/python

"""
Example of working with a new user-defined potential model in Agama.
One needs to define a function `mypot` that takes a single argument
(a Nx3 array of points in Cartesian coords) and returns an array of N
potential values. Then one may convert this function into a full-fledged
instance of the Potential class as follows:
>>> mypot_wrapper = agama.Potential(mypot)
and this wrapper provides access to various functions such as force,
density (which are computed by finite differences), total mass,
can be used in orbit integration, etc.
However, it is rather inefficient due to large overheads from switching
between Python and C++ on every potential evaluation, so the optimal
way of using such a user-defined potential is to approximate it with
either a Multipole or a CylSpline potential expansion:
>>> mypot_exp = agama.Potential(type='Multipole', potential=mypot, ...)
Additional arguments '...' should specify the grid parameters and the
degree of symmetry.
In this example, we construct a class representing an N-arm logarithmic
spiral density perturbation from Cox&Gomez (2002), which is converted to
a CylSpline potential and added to a purely axisymmetric disk model.
We then integrate some orbits in the rotating spiral potential.
"""

import numpy, agama

# A factory function creating a potential with the given parameters
# (it needs to be a free function and not an instance of some class)
def createSpiralPotential(numberOfArms, surfaceDensity, scaleRadius, scaleHeight, pitchAngle, phi0=0):
    '''
    Create a CylSpline approximation for a spiral arm potential from Cox&Gomez(2002).
    The density generated by this potential approximately follows an exponential
    disk profile with the amplitude of surface density variation
    Sigma(R) = surfaceDensity * exp( - R / scaleRadius ),
    and isothermal vertical profile
    rho(R,z) = Sigma(R) / (4h) * sech^2( z / (2h) ),
    where h is the scaleHeight (note the extra factor 1/2 w.r.t. the original paper,
    which is introduced to match the definition of the Disk profile in Agama).
    The density at the crest of the spiral arm is  ~ 3 * rho(R,z=0),
    while in the inter-arm region it is ~ (-1.5 to -1) * rho(R,z=0).
    The spiral is logarithmic, so that the azimuthal angle (of one of the arms) is
    phi(R) ~ phi0 + ln(R/scaleRadius) / tan(pitchAngle).
    '''
    def potfnc(xyz):
        R = (xyz[:,0]**2 + xyz[:,1]**2)**0.5
        z = xyz[:,2]
        phi = numpy.arctan2(xyz[:,1], xyz[:,0])
        prefac = -4*numpy.pi * agama.G * surfaceDensity * numpy.exp(-R / scaleRadius)
        gamma  = numberOfArms * (phi - numpy.log(R / scaleRadius) / numpy.tan(pitchAngle) - phi0)
        Phi = numpy.zeros(len(R))
        for n in range(1,4):
            K_n   = n * numberOfArms / R / numpy.sin(pitchAngle)
            K_n_H = K_n * 2 * scaleHeight
            B_n   = K_n_H * (1 + 0.4 * K_n_H)
            D_n   = 1 / (1 + 0.3 * K_n_H) + K_n_H
            C_n   = [8./3/numpy.pi, 0.5, 8./15/numpy.pi][n-1]  # amplitudes of the three harmonic terms
            Phi  += prefac * ( C_n / D_n / K_n * numpy.cos(n * gamma) *
                (numpy.cosh(K_n * z / B_n))**-B_n )
        return numpy.nan_to_num(Phi)  # VERY IMPORTANT is to make sure it never produces a NaN

    # now create a CylSpline potential approximating this user-defined profile,
    # paying special attention to the grid parameters
    return agama.Potential(type='CylSpline', potential=potfnc,
    Rmin = 0.01 * scaleRadius,  # (these could be adjusted if needed)
    Rmax = 10.0 * scaleRadius,
    zmin = 0.25 * scaleHeight,
    zmax = 10.0 * scaleHeight,
    mmax = 3 * numberOfArms,    # each arm is represented by three harmonic terms: m, m*2, m*3
    symmetry  = 'bisymmetric' if numberOfArms % 2 == 0 else 4,  # 4 means the z-reflection symmetry
    gridSizeZ = 20,
    # rule of thumb for the radial grid spacing is to resolve
    # the change in the pitch angle of the spiral with one grid cell
    gridSizeR = max(25, numpy.log(10.0 / 0.01) / numpy.tan(pitchAngle) * numberOfArms))

agama.setUnits(length=1, velocity=1, mass=1)  # some arbitrary units - everything is scale-invariant here
numberOfArms    = 3
pitchAngle      = numpy.pi/8
scaleRadius     = 1.0
scaleHeight     = 0.1
surfaceDensityD = 1.0  # density of the underlying axisymmetric disk
surfaceDensityS = surfaceDensityD * 0.5  # amplitude of the density variation of the spiral pattern
phi0            = numpy.pi/10   # [arbitrary] phase angle at R=scaleRadius
pot_spiral = createSpiralPotential(numberOfArms, surfaceDensityS, scaleRadius, scaleHeight, pitchAngle, phi0)
pot_disk = agama.Potential(type='disk', surfaceDensity=surfaceDensityD, scaleRadius=scaleRadius,
    scaleHeight=-scaleHeight)  # minus sign means the isothermal (sech^2) vertical profile
pot_total = agama.Potential(pot_disk, pot_spiral)
pot_spiral.export('example_spiral.ini')

# illustration: compute the density of the axisymmetric and the perturbed disks in the equatorial plane
import matplotlib, matplotlib.pyplot as plt
plt.ion()
ax = plt.subplots(1, 2, figsize=(10,5))[1]
X = numpy.linspace(-5*scaleRadius, 5*scaleRadius, 501)
Y = numpy.linspace(-5*scaleRadius, 5*scaleRadius, 481)
XYZ = numpy.column_stack((numpy.repeat(X, len(Y)), numpy.tile(Y, len(X)), numpy.zeros(len(X)*len(Y)) ))
rhoD = pot_disk.density(XYZ)
rhoT = pot_total.density(XYZ)
rhoMax = numpy.max(rhoD) * 3.0
rhoMin = rhoMax * 1e-3
levels = numpy.logspace(numpy.log10(rhoMin), numpy.log10(rhoMax), 32)
norm = matplotlib.colors.LogNorm()
ax[0].contour(X, Y, rhoD.reshape(len(X), len(Y)).T, levels=levels, cmap='hell_r', norm=norm)
ax[1].contour(X, Y, rhoT.reshape(len(X), len(Y)).T, levels=levels, cmap='hell_r', norm=norm)
# overplot the log-spiral
RR = numpy.logspace(-2, 0.7, 200) * scaleRadius
phi = numpy.log(RR / scaleRadius) / numpy.tan(pitchAngle) + phi0
ax[1].plot(RR * numpy.cos(phi), RR * numpy.sin(phi), 'g-')
plt.tight_layout()
plt.show()
plt.pause(5)

# compute some orbits in the axisymmetric and perturbed potentials
# (the latter rotating with pattern speed Omega)
R0 = 3.0
V0 = (-R0 * pot_disk.force(R0,0,0)[0])**0.5   # circular velocity at the given radius
# note that we set the pattern speed to a negative value, so that the spirals are trailing
Omega = -0.8 * V0 / R0
ic = [-R0, 0.0, 0.0, 0.0, 0.9*V0, 0.1*V0]   # and make the angular momentum negative as well (i.e.prograde)
time  = 20 * R0 / V0
nsteps= 200
# trajectories are stored in the rotating frame
t, orbit_disk = agama.orbit(potential=pot_disk,  ic=ic, time=time, trajsize=nsteps, Omega=Omega)
_, orbit_spir = agama.orbit(potential=pot_total, ic=ic, time=time, trajsize=nsteps, Omega=Omega)
# convert the trajectories to the inertial frame
cosa, sina = numpy.cos(t * Omega), numpy.sin(t * Omega)
orbit_disk_x = orbit_disk[:,0] * cosa - orbit_disk[:,1] * sina
orbit_disk_y = orbit_disk[:,1] * cosa + orbit_disk[:,0] * sina
orbit_spir_x = orbit_spir[:,0] * cosa - orbit_spir[:,1] * sina
orbit_spir_y = orbit_spir[:,1] * cosa + orbit_spir[:,0] * sina

# plot the animated orbit and rotating potential
for i in range(nsteps):
    # the spirals need to be replotted in each frame at a different orientation
    XYZrot = numpy.column_stack((XYZ[:,0] * cosa[i] + XYZ[:,1] * sina[i],
        XYZ[:,1] * cosa[i] - XYZ[:,0] * sina[i], XYZ[:,2]))
    rhoT = pot_total.density(XYZrot)
    ax[1].cla()
    ax[1].contour(X, Y, rhoT.reshape(len(X), len(Y)).T, levels=levels, cmap='hell_r', norm=norm)
    # overplot the current position (for the left panel, which is not cleared) in the inertial frame
    ax[0].plot(orbit_disk_x[i], orbit_disk_y[i], 'ko', ms=3)
    # overplot the past trajectory up to the current time (for the right panel, which is cleared each time)
    ax[1].plot(orbit_spir_x[:i+1], orbit_spir_y[:i+1], 'ko', ms=3)
    plt.draw()
    plt.pause(0.1)

plt.ioff()
plt.show()
