"""Integrator for isothermal atmospheres.

This integrator is not smart at all, using a simple Runge-Kutta scheme
for taking steps outward.

Here, the independent variable is the height above the surface, and the
dependent variables are the column density y and the density d.

The most important functions are `pressure` (which contains the equation
of state) and mche, which has the equations of mass continuity and
hydrostatic equilibrium which are to be integrated.

By running the code with

  >>> ipython3 -i atmosphere.py

the integration will be done for conditions appropriate for the Earth.
"""
# easiest: start ipython
# run -i atmosphere.py
import numpy as np
from astropy import units as u, constants as c
from astropy.table import QTable
from astropy.visualization import quantity_support


def pressure(d, T):
    """EOS: ideal gas pressure P for density d and temperature T."""
    kT = c.k_B * T
    P = d / (mu*mh) * kT
    dPdd = 1. / (mu*mh) * kT
    dPdT = d / (mu*mh) * kT
    # total pressure and derivatives
    return P, dPdd, dPdT


def mche(h, y):
    """MC and HE for plane-parallel atmosphere.

    Written in terms of column density y.

    d y
    --- = rho
    d h

    d P                 d rho   d P
    --- = - g rho   ->  ----- = --- / (dP/drho)
    d h                  d h    d h
    """

    y, d = y[0], y[1]
    P, dPdd, dPdT = pressure(d, T)
    dydh = d
    dPdh = -g * d
    dddh = dPdh / dPdd
    return dydh, dddh


def yadd(y0, yf, dy):
    """Helper for Runge-Kutta, to allow items in y with different units.

    Just does y = y0 + yf * dy, but iterating over y and dy.
    """
    return [_y0 + yf * _dy for (_y0, _dy) in zip(y0, dy)]


def rk4(x, dx, y, dxdy):
    """Runge-Kutta 4th order integration.

    Function dxdy should give derivatives of y with respect to x.
    """
    dydx1 = dxdy(x, y)
    dydx2 = dxdy(x + 0.5 * dx, yadd(y, 0.5 * dx, dydx1))
    dydx3 = dxdy(x + 0.5 * dx, yadd(y, 0.5 * dx, dydx2))
    dydx4 = dxdy(x + dx, yadd(y, dx, dydx3))
    # 4-th order: yn = y + dx / 6 * (dydx1 + 2*dydx2 + 2*dydx3 + dxdy4)
    yn = yadd(y, dx / 6., dydx1)
    yn = yadd(yn, dx / 3., dydx2)
    yn = yadd(yn, dx / 3., dydx3)
    yn = yadd(yn, dx / 6., dydx4)
    return x + dx, yn


def integrate(d0=1. * u.kg / u.m**3, hstep=1. * u.km, dlim=None):
    """Integrate atmosphere for column density y and density d"""
    if dlim is None:
        dlim = d0 / 1e3
    hs, ys, ds = [], [], []
    h, y, d = 0 * u.km, 0 * u.kg / u.m ** 2, d0
    while d > dlim:
        hs.append(h)
        ys.append(y)
        ds.append(d)
        h, (y, d) = rk4(h, hstep, (y, d), mche)
    return QTable([u.Quantity(hs), u.Quantity(ys), u.Quantity(ds)],
                  names=('h', 'y', 'd'))


if __name__ == '__main__':
    import matplotlib.pylab as plt
    quantity_support()
    plt.ion()
    # Typical surface temperature.
    T = 300. * u.K
    # Mean molecular weight for N2
    mu = 28.
    mh = c.m_p + c.m_e
    # Surface gravity
    g = c.G * c.M_earth / c.R_earth ** 2
    # Density of air at sea level.
    d0 = 1.225 * u.kg / u.m**3
    # do integration
    t = integrate(d0)
    # show expected exponential
    H = (c.k_B*T) / (mu*mh*g)
    plt.plot(t['h'], d0 * np.exp(-t['h'] / H), label='expected exponential')
    # and integration results
    plt.plot(t['h'], t['d'], '.', label='numerical integration')
    plt.title('Atmospheric density profile on Earth for T={}'.format(T))
    plt.legend()
    print("Total column: ", t['y'][-1].to(u.kg/u.cm**2))