#!/usr/bin/env python3
#
# Weitian LI
# Created: 2016-06-24
# Updated: 2016-06-27
#
# Change logs:
# 2016-06-27:
#   * Implement potential profile calculation:
#     methods 'calc_density_total()' and 'calc_potential()'
# 2016-06-26:
#   * Add document on gravitational potential calculation
# 2016-06-25:
#   * Rename method 'interpolate()' to 'fit_spline()'
#   * Use 'InterpolatedUnivariateSpline' instead of 'interp1d'
#   * Implement 'calc_mass_total()'
#   * Update documentation
# 2016-06-24:
#   * Update method 'gen_radius()'
#

"""
Calculate the (gas and gravitational) mass profile and gravitational
potential profile from the electron number density profile, under the
assumption of hydrostatic equilibrium.

The electron density profile and temperature profile are required.

Assuming that the gas is in hydrostatic equilibrium with the gravitational
potential and a spherically-symmetric distribution of the gas, we can
write the hydrostatic equilibrium equation (HEE) of the ICM as
(ref.[1], eq.(6)):
    derivative(P_gas, r) / rho_gas = - derivative(phi, r)
                                   = - G M_tot(<r) / r^2
where,
    phi: gravitational potential;
    G: gravitational constant;
    rho_gas: gas mass density:
        rho_gas = mu * m_atom * n_gas
    P_gas: gas pressure:
        P_gas = rho_gas * k_B * T_gas / (mu * m_atom) = n_gas * k_B * T_gas
    mu: mean molecular weight in a.m.u (i.e., m_atom) (~ 0.6)
    m_atom: atom mass unit
    n_gas: gas number density; sum of the electron and proton densities
    k_B: Boltzmann constant
    T_gas: gas temperature

Solve the above equation, we can get the total mass of X-ray luminous
galaxy clusters (ref.[1], eq.(7)):
    M_tot(<r) = - (k_B * T_gas(r) * r) / (mu * m_atom * G) *
                (derivative(log(T_gas), log(r)) +
                 derivative(log(n_gas), log(r)))

Note that the second part (i.e., the derivatives) is DIMENSIONLESS, since
    d(log(X)) = d(X) / X
Also note that ('R' is a ratio constant):
    d(log(n_gas)) = d(log(R*n_e)) = d(log(n_e))
And and 'log' derivative can be calculated as:
    derivative(log(X(r)), log(r)) = (r / X(r)) * derivative(X(r), r)

Note that 'kT' has dimension of energy.  Therefore, if the gas temperature
is given in 'keV', then the 'kT' should be substitute as a whole.

For example:
    (1.0 keV) * (1.0 kpc) / (0.6 * m_atom * G) ~= 3.7379e10 [ Msun ]
which is consistent with the formula of (ref.[2], eq.(3))

------------------------------------------------------------

Gravitational potential profile calculation:

Newton's theorems (ref.[3], 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] Ettori et al., 2013, Space Science Review, 177, 119-154
[2] Walker et al., 2012, MNRAS, 422, 3503
[3] Tremaine & Binney, Galactic Dynamics, 2nd edition, 2008


Sample configuration file:
------------------------------------------------------------
## Configuration for `calc_mass_potential.py`
## Date: 2016-06-24

# redshift used for pixel to distance conversion
redshift = <REDSHIFT>

# electron density profile
ne_profile = ne_profile.txt

# cooling function profile
cf_profile = coolfunc_profile.txt
# unit of the CF profile radius (default: pixel)
cf_unit = "pixel"

# temperature profile
t_profile = t_profile.txt
# unit of the T profile radius (default: pixel)
t_unit = "pixel"

# number of data points for the output profile calculation
num_dp = 1000

# output gas mass profile
m_gas_profile = mass_gas_profile.txt

# output total (gravitational) mass profile
m_total_profile = mass_total_profile.txt

# output total mass density profile
rho_total_profile = rho_total_profile.txt

# output gravitational potential profile
potential_profile = potential_profile.txt
------------------------------------------------------------
"""

import argparse

import numpy as np
import astropy.units as au
import astropy.constants as ac
import scipy.interpolate as interpolate
import scipy.integrate as integrate
from scipy.misc import derivative
from configobj import ConfigObj

from astro_params import AstroParams, ChandraPixel
from projection import Projection


class DensityProfile:
    """
    Utilize the 3D (electron number or gas mass) density profile to
    calculate the following quantities:
    * 2D projected surface brightness (requires cooling function profile)
    * gas mass profile (integrated, M_gas(<r))
    * gravitational mass profile (M(<r); requires temperature profile)
    * gravitational potential profile (cut at the largest available radius)

    NOTE:
    * The radii (of density profile and cooling function profile)
      should have unit [ cm ]
    * The density should have unit [ cm^-3 ] or [ g cm^-3 ]
    """
    # allowed density profile types
    DENSITY_TYPES = ["electron", "gas"]
    # input density data: [r, r_err, d]
    r = None
    r_err = None
    d = None
    # electron number density
    ne = None
    # gas mass density
    rho_gas = None
    # cooling function profile
    cf_radius = None
    cf_value = None
    # temperature profile
    t_radius = None
    t_value = None
    # fitted spline to the profile
    d_spline = None
    cf_spline = None
    t_spline = None
    # generated radial data points for profile calculation
    radius = None
    radius_err = None
    # gas mass profile: M_gas(<r); same length as the above 'radius'
    m_gas = None
    # total (gravitational) mass profile: M_total(<r)
    m_total = None
    # total mass density profile (required by potential calculation)
    rho_total = None
    # potential profile (cut at the largest available radius)
    potential = None

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

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

    def load_cf_profile(self, data):
        # 2-column cooling function profile: r[cm], cf[flux/EM]
        self.cf_radius = data[:, 0].copy()
        self.cf_value = data[:, 1].copy()

    def load_t_profile(self, data):
        # 2-column temperature profile: r[cm], T[keV]
        self.t_radius = data[:, 0].copy()
        self.t_value = data[:, 1].copy()

    def calc_brightness(self):
        """
        Project the electron number density or gas mass density profile
        to calculate the 2D surface brightness profile.
        """
        if self.cf_radius is None or self.cf_value is None:
            raise ValueError("cooling function profile missing")
        ne = self.calc_density_electron()
        # flux per unit volume
        flux = self.cf_spline(self.r) * ne ** 2 / AstroParams.ratio_ne_np
        # project the 3D flux
        projector = Projection(rout=self.r+self.r_err)
        brightness = projector.project(flux)
        return brightness

    def calc_density_electron(self):
        """
        Calculate the electron number density from the gas mass density
        if necessary.
        """
        if self.density_type == "electron":
            self.ne = self.d.copy()
        elif self.density_type == "gas":
            self.ne = self.d / AstroParams.m_atom / AstroParams.mu_e
        return self.ne

    def calc_density_gas(self):
        """
        Calculate the gas mass density from the electron number density
        if necessary.
        """
        if self.density_type == "electron":
            self.rho_gas = self.d * AstroParams.mu_e * AstroParams.m_atom
        elif self.density_type == "gas":
            self.rho_gas = self.d.copy()
        return self.rho_gas

    def fit_spline(self, verbose=False):
        """
        Fit smoothing spline to the density profile, cooling function
        profile, and temperature profile.

        NOTE:
        * Simple interpolation (`scipy.interpolate.interp1d` of kinds
          `linear` or `quadratic` or `cubic`) may lead to a severely
          oscillating curve (e.g., electron density profile), which
          further cause problems for following mass calculation
          due to the needs to take the derivative.
        * `InterpolatedUnivariateSpline` or `UnivariateSpline` is a
          better choice than the above `interp1d`.
        * Allow cooling function profile and temperature profile to be
          extrapolated by filling with the last value if necessary.
        """
        # density profile
        # insert a data point at radius of zero
        if verbose:
            print("Fitting spline to density profile ...")
        self.d_spline = interpolate.InterpolatedUnivariateSpline(
            x=np.concatenate([[0.0], self.r]),
            y=np.concatenate([[self.d[0]], self.d]))
        if self.ne is not None:
            if verbose:
                print("Fitting spline to electron number density profile ...")
            self.ne_spline = interpolate.InterpolatedUnivariateSpline(
                x=np.concatenate([[0.0], self.r]),
                y=np.concatenate([[self.ne[0]], self.ne]))
        if self.rho_gas is not None:
            if verbose:
                print("Fitting spline to gas mass density profile ...")
            self.rho_gas_spline = interpolate.InterpolatedUnivariateSpline(
                x=np.concatenate([[0.0], self.r]),
                y=np.concatenate([[self.rho_gas[0]], self.rho_gas]))
        # cooling function profile
        if verbose:
            print("Fitting spline to cooling function profile ...")
        self.cf_spline = interpolate.InterpolatedUnivariateSpline(
            x=self.cf_radius, y=self.cf_value, ext="const")
        # temperature profile
        if verbose:
            print("Fitting spline to temperature profile ...")
        self.t_spline = interpolate.InterpolatedUnivariateSpline(
            x=self.t_radius, y=self.t_value, ext="const")

    def gen_radius(self, num=1000):
        """
        Generate radial points for following mass and potential calculation.

        The generated radial points are logarithmic-evenly distributed.

        NOTE:
        The radii are first generated to determine the inner-most bin width,
        which is used to further divide the original inner-most bin (i.e.,
        radius 0 - r_out[0]), and then the other radii are generated with
        the constraint of given total number of points.
        """
        rout = self.r + self.r_err
        # first pass to determine the inner-most bin width
        rout_tmp = np.logspace(np.log10(rout[0]), np.log10(rout[-1]),
                               num=num, base=10.0)
        bw = rout_tmp[1] - rout_tmp[0]
        # linear-evenly divide the first original bin (0 - rout[0])
        nbin = int(np.ceil(rout[0] / bw))
        rout_new1 = np.linspace(0.0, rout[0], num=nbin, endpoint=False)[1:]
        # second pass to generate the other radii
        rout_new2 = np.logspace(np.log10(rout[0]), np.log10(rout[-1]),
                                num=(num-nbin+1), base=10.0)
        rout_new = np.concatenate([rout_new1, rout_new2])
        rin_new = np.concatenate([[0.0], rout_new[:-1]])
        self.radius = (rout_new + rin_new) / 2.0
        self.radius_err = (rout_new - rin_new) / 2.0

    def calc_mass_gas(self, verbose=False):
        """
        Calculate the gas mass profile, i.e., the mass of the gas within
        each radius.

        Reference: ref.[1], eq.(9)
        """
        def _f_rho_gas(r):
            return self.rho_gas_spline(r) * 4*np.pi * r**2
        #
        m_gas = np.zeros(self.radius.shape)
        if verbose:
            print("Calculating the gas mass profile (#%d): ..." %
                  len(self.radius), end="", flush=True)
        for i, r in enumerate(self.radius):
            if verbose and (i+1) % 50 == 0:
                print("%d..." % (i+1), end="", flush=True)
            # integrated gas mass [ g ]
            m_gas[i] = integrate.quad(_f_rho_gas, a=0.0, b=r,
                                      epsabs=1e-5*au.kpc.to(au.cm),
                                      epsrel=1e-3)[0]
        if verbose:
            print("DONE!", flush=True)
        self.m_gas = m_gas
        return m_gas

    def calc_mass_total(self, verbose=True):
        """
        Calculate the total mass (i.e., gravitational mass) profile,
        under the assumption of hydrostatic equilibrium (HE).

        References: ref.[1], eq.(5,6,7)
        """
        if self.t_radius is None or self.t_value is None:
            raise ValueError("temperature profile required")
        #
        # calculate the coefficient of the total mass formula
        # M0 = (k_B * T_gas(r) * r) / (mu * m_atom * G)
        M0 = ((1.0*au.keV) * (1.0*au.cm) /
              (AstroParams.mu * ac.u * ac.G)).to(au.g).value
        m_total = np.zeros(self.radius.shape)
        if verbose:
            print("Calculating the total mass profile (#%d): ..." %
                  len(self.radius), end="", flush=True)
        for i, r in enumerate(self.radius):
            if verbose and (i+1) % 100 == 0:
                print("%d..." % (i+1), end="", flush=True)
            # enclosed total mass [ g ]
            m_total[i] = - M0 * self.t_spline(r) * r * (
                ((r / self.t_spline(r)) *
                 derivative(self.t_spline, r, dx=0.01*au.kpc.to(au.cm))) +
                ((r / self.ne_spline(r)) *
                 derivative(self.ne_spline, r, dx=0.01*au.kpc.to(au.cm))))
        if verbose:
            print("DONE!", flush=True)
        self.m_total = m_total
        return m_total

    def calc_density_total(self, verbose=True):
        """
        Calculate the total mass density profile, which is required to
        calculate the following gravitational potential profile.
        """
        rho_total = np.zeros(self.radius.shape)
        if verbose:
            print("Calculating the total mass density profile ...")
            print(">>> Fitting spline to total mass profile ...")
        self.m_total_spline = interpolate.InterpolatedUnivariateSpline(
            x=self.radius, y=self.m_total, ext="const")
        for i, r in enumerate(self.radius):
            rho_total[i] = (derivative(self.m_total_spline, r,
                                       dx=0.01*au.kpc.to(au.cm)) /
                            (4 * np.pi * r**2))
        self.rho_total = rho_total
        if verbose:
            print(">>> Fitting spline to total mass density profile ...")
        self.rho_total_spline = interpolate.InterpolatedUnivariateSpline(
            x=self.radius, y=rho_total, ext="const")

    def calc_potential(self, verbose=True):
        """
        Calculate the gravitational potential 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.rho_total_spline(x)

        def _int_outer(x):
            return x * self.rho_total_spline(x)

        if self.rho_total is None:
            self.calc_density_total(verbose=verbose)
        potential = np.zeros(self.radius.shape)
        if verbose:
            print("Calculating the potential profile (#%d): ..." %
                  len(self.radius), end="", flush=True)
        r_max = max(self.radius)
        for i, r in enumerate(self.radius):
            if verbose and (i+1) % 50 == 0:
                print("%d..." % (i+1), end="", flush=True)
            potential[i] = - 4 * np.pi * ac.G.cgs.value * (
                (1/r) * integrate.quad(_int_inner, a=0.0, b=r,
                                       epsabs=1e-5*au.kpc.to(au.cm),
                                       epsrel=1e-3)[0] +
                integrate.quad(_int_outer, a=r, b=r_max,
                               epsabs=1e-5*au.kpc.to(au.cm),
                               epsrel=1e-3)[0])
        if verbose:
            print("DONE!", flush=True)
        self.potential = potential
        return potential

    def plot(self, profile, ax=None, fig=None):
        pass

    def save(self, profile, outfile):
        if profile == "mass_gas":
            data = np.column_stack([self.radius,
                                    self.radius_err,
                                    self.m_gas])
            header = "radius[cm]  radius_err[cm]  mass_gas(<r)[g]"
        elif profile == "mass_total":
            data = np.column_stack([self.radius,
                                    self.radius_err,
                                    self.m_total])
            header = "radius[cm]  radius_err[cm]  mass_total(<r)[g]"
        elif profile == "rho_total":
            data = np.column_stack([self.radius,
                                    self.radius_err,
                                    self.rho_total])
            header = "radius[cm]  radius_err[cm]  density_total[g/cm^3]"
        elif profile == "potential":
            data = np.column_stack([self.radius,
                                    self.radius_err,
                                    self.potential])
            header = "radius[cm]  radius_err[cm]  potential[???]"
        else:
            raise ValueError("unknown profile name: %s" % profile)
        np.savetxt(outfile, data, header=header)


def main():
    parser = argparse.ArgumentParser(
            description="Calculate the mass and potential profiles")
    parser.add_argument("config", nargs="?", default="mass_potential.conf",
                        help="config for mass and potential profiles " +
                             "calculation (default: mass_potential.conf)")
    args = parser.parse_args()

    config = ConfigObj(args.config)

    redshift = config.as_float("redshift")
    pixel = ChandraPixel(z=redshift)

    ne_profile = np.loadtxt(config["ne_profile"])

    cf_profile = np.loadtxt(config["cf_profile"])
    cf_unit = "pixel"
    try:
        cf_unit = config["cf_unit"]
    except KeyError:
        pass
    if cf_unit == "pixel":
        conv_factor = pixel.get_length().to(au.cm).value
    else:
        conv_factor = au.Unit(cf_unit).to(au.cm)
    cf_profile[:, 0] *= conv_factor

    t_profile = np.loadtxt(config["t_profile"])
    t_unit = "pixel"
    try:
        t_unit = config["t_unit"]
    except KeyError:
        pass
    if t_unit == "pixel":
        conv_factor = pixel.get_length().to(au.cm).value
    else:
        conv_factor = au.Unit(t_unit).to(au.cm)
    t_profile[:, 0] *= conv_factor

    density_profile = DensityProfile(density=ne_profile,
                                     density_type="electron")
    density_profile.load_cf_profile(cf_profile)
    density_profile.load_t_profile(t_profile)
    density_profile.calc_density_electron()
    density_profile.calc_density_gas()
    density_profile.fit_spline(verbose=True)
    density_profile.gen_radius(num=config.as_int("num_dp"))
    density_profile.calc_mass_gas(verbose=True)
    density_profile.save(profile="mass_gas",
                         outfile=config["m_gas_profile"])
    density_profile.calc_mass_total(verbose=True)
    density_profile.save(profile="mass_total",
                         outfile=config["m_total_profile"])
    density_profile.calc_density_total(verbose=True)
    density_profile.save(profile="rho_total",
                         outfile=config["rho_total_profile"])
    density_profile.calc_potential(verbose=True)
    density_profile.save(profile="potential",
                         outfile=config["potential_profile"])


if __name__ == "__main__":
    main()