#!/usr/bin/env python3
#
# Aaron LI
# Created: 2016-07-13
# Updated: 2016-07-13
#
# Change logs:
#

"""
Calculate the gravitational potential profile through the
total mass density profile.

Newton's theorems (ref.[1], Sec. 2.2.1):
1. A body that is inside a spherical shell of matter experiences no net
   gravitational force from that shell.
2. The gravitational force on a body that lies outside a spherical shell
   of matter is the same as it would be if all the shell's matter were
   concentrated into a point at its center.

Therefore, the gravitational potential produced by a spherical shell of
mass 'M' is:
    phi = (1) - G * M / R;  (r <= R, i.e., inside the shell)
          (2) - G * M / r;  (r > R, i.e., outside the shell)

The total gravitational potential may be considered to be the sum of the
potentials of spherical shells of mass
    dM(r) = 4 * pi * rho(r) * r^2 dr,
so we may calculate the gravitational potential at 'R' generated by an
arbitrary spherically symmetric density distribution 'rho(r)' by adding
the contributions to the potential produced by shells
    (1) with r < R,
and
    (2) with r > R.
In this way, we obtain
    phi(R) = - (G/R) * \int_0^R dM(r) - G * \int_R^{\inf} dM(r)/r
           = - 4*pi*G * ((1/R) * \int_0^R r^2 * rho(r) dr +
             \int_R^{\inf} r * rho(r) dr)

References:
[1] Tremaine & Binney, Galactic Dynamics, 2nd edition, 2008
"""


import argparse

import numpy as np
import astropy.units as au
import astropy.constants as ac
import scipy.integrate as integrate
from configobj import ConfigObj
import matplotlib.pyplot as plt
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
from matplotlib.figure import Figure

from spline import SmoothSpline

plt.style.use("ggplot")


config_default = """
## Configuration for `calc_potential.py`

# total mass density profile
rho_total_profile = rho_total_profile.txt

# output gravitational potential profile
potential_profile = potential_profile.txt
potential_profile_image = potential_profile.png
"""


class DensityProfile:
    """
    Cluster's gas/total mass density profile.
    """
    # supported types of density profile
    DENSITY_TYPES = ["total", "gas"]
    # available splines
    SPLINES = ["rho_total"]
    # input density data: [r, r_err, m]
    r = None
    r_err = None
    d = None
    # redshift of the object
    redshift = None
    # fitted smoothing spline to total density profile
    rho_total_spline = None
    # calculated potential profile
    potential = None

    def __init__(self, data, density_type="total"):
        self.load_data(data=data, density_type=density_type)

    def load_data(self, data, density_type="total"):
        if density_type not in self.DENSITY_TYPES:
            raise ValueError("invalid density_types: %s" % density_type)
        # 3-column density profile: r[kpc], r_err[kpc], rho[g/cm^3]
        self.r = data[:, 0].copy()
        self.r_err = data[:, 1].copy()
        self.d = data[:, 2].copy()
        self.density_type = density_type

    def calc_potential(self, verbose=True):
        """
        Calculate the gravitational potential profile from the
        total mass density profile.

        NOTE:
        * The integral in the potential formula can only be integrated
          to the largest radius of availability.
        * This limitation will affect the absolute values of the calculated
          potentials, but not the shape of the potential profile.
        """
        def _int_inner(x):
            return x**2 * self.eval_spline(spline="rho_total", x=x)

        def _int_outer(x):
            return x * self.eval_spline(spline="rho_total", x=x)

        if self.density_type != "total":
            raise ValueError("total mass density profile is required")

        if self.rho_total_spline is None:
            self.fit_spline(spline="rho_total", log10=["x", "y"])
        potential = np.zeros(self.r.shape)
        if verbose:
            print("Calculating the potential profile (#%d): ..." %
                  len(self.r), end="", flush=True)
        r_max = max(self.r)
        c = - 4 * np.pi * au.kpc.to(au.cm)**2
        for i, r in enumerate(self.r):
            if verbose and (i+1) % 10 == 0:
                print("%d..." % (i+1), end="", flush=True)
            # total gravitational potential [ cm^2/s^2 ]
            potential[i] = c * ac.G.cgs.value * (
                (1/r) * integrate.quad(_int_inner, a=0.0, b=r,
                                       epsrel=1e-2)[0] +
                integrate.quad(_int_outer, a=r, b=r_max,
                               epsrel=1e-2)[0])
        if verbose:
            print("DONE!", flush=True)
        self.potential = potential
        return potential

    def plot(self, ax=None, fig=None):
        x = self.r
        xlabel = "3D Radius"
        xunit = "kpc"
        xscale = "log"
        yscale = "log"
        y = self.potential
        ylabel = "Gravitational potential"
        yunit = "cm$^2$/s$^2$"
        yscale = "linear"
        #
        if ax is None:
            fig, ax = plt.subplots(1, 1)
        ax.plot(x, y, linewidth=2)
        ax.set_xscale(xscale)
        ax.set_yscale(yscale)
        ax.set_xlim(min(x), max(x))
        y_min, y_max = min(y), max(y)
        y_min = y_min/1.2 if y_min > 0 else y_min*1.2
        y_max = y_max*1.2 if y_max > 0 else y_max/1.2
        ax.set_ylim(y_min, y_max)
        ax.set_xlabel(r"%s (%s)" % (xlabel, xunit))
        ax.set_ylabel(r"%s (%s)" % (ylabel, yunit))
        fig.tight_layout()
        return (fig, ax)

    def save(self, outfile):
        """
        Save calculated overdensity profile.
        """
        data = np.column_stack([self.r,
                                self.r_err,
                                self.potential])
        header = "radius[kpc]  radius_err[kpc]  potential[cm^2/s^2]"
        np.savetxt(outfile, data, header=header)

    def fit_spline(self, spline, log10=[]):
        if spline not in self.SPLINES:
            raise ValueError("invalid spline: %s" % spline)
        #
        if spline == "rho_total":
            # input total mass density profile
            x = self.r
            y = self.d
            spl = "rho_total_spline"
        else:
            raise ValueError("invalid spline: %s" % spline)
        setattr(self, spl, SmoothSpline(x=x, y=y))
        getattr(self, spl).fit(log10=log10)

    def eval_spline(self, spline, x):
        """
        Evaluate the specified spline at the supplied positions.
        Also check whether the spline was fitted in the log-scale space,
        and transform the evaluated values if necessary.
        """
        if spline == "rho_total":
            spl = self.rho_total_spline
        else:
            raise ValueError("invalid spline: %s" % spline)
        #
        return spl.eval(x)


def main():
    parser = argparse.ArgumentParser(
            description="Calculate the gravitational potential profile")
    parser.add_argument("config", nargs="?",
                        help="additional config")
    args = parser.parse_args()

    config = ConfigObj(config_default.splitlines())
    if args.config is not None:
        config_user = ConfigObj(open(args.config))
        config.merge(config_user)

    density_data = np.loadtxt(config["rho_total_profile"])
    density_profile = DensityProfile(data=density_data, density_type="total")
    density_profile.calc_potential(verbose=True)
    density_profile.save(outfile=config["potential_profile"])
    fig = Figure(figsize=(10, 8))
    FigureCanvas(fig)
    ax = fig.add_subplot(1, 1, 1)
    density_profile.plot(ax=ax, fig=fig)
    fig.savefig(config["potential_profile_image"], dpi=150)


if __name__ == "__main__":
    main()