"""
Functions for common cosmological calculations.

"""

import math

import numpy as np
import scipy
import scipy.integrate
from scipy.interpolate import interp1d

from .. import units
from ..array import SimArray
from ..configuration import config_parser

_interp_points = int(config_parser.get('general','cosmo-interpolation-points'))


def _a_dot(a, h0, om_m, om_l):
    om_k = 1.0 - om_m - om_l
    return h0 * a * np.sqrt(om_m * (a ** -3) + om_k * (a ** -2) + om_l)


def _a_dot_recip(*args):
    return 1. / _a_dot(*args)


def _da_dtau(a, h0, om_m, om_l):
    return a**2 * _a_dot(a, h0, om_m, om_l)

def _da_dtau_recip(*args):
    return 1. / _da_dtau(*args)


def _hzoverh0(a, omegam0):
    """ returns: H(a) / H0  = [omegam/a**3 + (1-omegam)]**0.5 """
    return np.sqrt(omegam0 * np.power(a, -3) + (1. - omegam0))


def _lingrowthintegrand(a, omegam0):
    """ (e.g. eq. 8 in lukic et al. 2008)   returns: da / [a*H(a)/H0]**3 """
    return np.power((a * _hzoverh0(a, omegam0)), -3)


def _lingrowthfac(red, omegam0, omegal0, return_norm=False):
    """
    returns: linear growth factor, b(a) normalized to 1 at z=0, good for flat lambda only
    a = 1/1+z
    b(a) = Delta(a) / Delta(a=1)   [ so that b(z=0) = 1 ]
    (and b(a) [Einstein de Sitter, omegam=1] = a)

    Delta(a) = 5 omegam / 2 H(a) / H(0) * integral[0:a] [da / [a H(a) H0]**3]
    equation  from  peebles 1980 (or e.g. eq. 8 in lukic et al. 2008) """

    # need to add w ~= , nonflat, -1 functionality

    import scipy.integrate

    if (abs(omegam0 + omegal0 - 1.) > 1.e-4):
        raise RuntimeError("Linear growth factors can only be calculated for flat cosmologies")

    a = 1 / (1. + red)

    # 1st calc. for z=z
    lingrowth = scipy.integrate.quad(_lingrowthintegrand, 0., a, (omegam0))[0]
    lingrowth *= 5. / 2. * omegam0 * _hzoverh0(a, omegam0)

    # then calc. for z=0 (for normalization)
    a0 = 1.
    lingrowtha0 = scipy.integrate.quad(
        _lingrowthintegrand, 0., a0, (omegam0))[0]
    lingrowtha0 *= 5. / 2. * omegam0 * _hzoverh0(a0, omegam0)

    lingrowthfactor = lingrowth / lingrowtha0
    if return_norm:
        return lingrowthfactor, lingrowtha0
    else:
        return lingrowthfactor


def linear_growth_factor(f, z=None):
    """Calculate the linear growth factor b(a), normalized to 1 at z=0, for the cosmology of snapshot f.

    The output is dimensionless. If a redshift z is specified, it is used in place of the redshift in
    output f.
    """
    if z is None:
        z = f.properties['z']
    omegam0 = f.properties['omegaM0']
    omegal0 = f.properties['omegaL0']
    return _lingrowthfac(z, omegam0, omegal0)


def rate_linear_growth(f, z=None, unit='h Gyr^-1'):
    """Calculate the linear growth rate b'(a), normalized to 1 at z=0, for the cosmology of snapshot f.

    The output is in 'h Gyr^-1' by default, but other units can be specified. If a redshift z is specified,
    it is used in place of the redshift in output f.

    """

    if z is None:
        z = f.properties['z']
    a = 1. / (1. + z)
    omegam0 = f.properties['omegaM0']
    omegal0 = f.properties['omegaL0']

    b, X = _lingrowthfac(z, omegam0, omegal0, return_norm=True)
    I = _lingrowthintegrand(a, omegam0)

    term1 = -(1.5 * omegam0 * a ** -3) * b / \
        math.sqrt(1. - omegam0 + omegam0 * a ** -3)
    term2 = (2.5 * omegam0) * _hzoverh0(a, omegam0) ** 2 * a * I / X

    res = units.h * (term1 + term2) * 100. * units.Unit("km s^-1 Mpc^-1")

    return res.in_units(unit, **f.conversion_context())


def _test_rate_linear_growth(f, z=None, unit='h Gyr^-1'):
    # coded up by AP to test linear growth *rate* equation above
    if z is None:
        z = f.properties['z']
    a0 = 1. / (1. + z)
    a1 = a0 * 0.999
    z0 = 1. / a0 - 1
    z1 = 1. / a1 - 1

    b0 = linear_growth_factor(f, z0)
    b1 = linear_growth_factor(f, z1)

    db = b1 - b0

    unit = units.Unit(unit)
    dt = age(f, z1, unit ** -1) - age(f, z0, unit ** -1)

    return db / dt


def age(f, z=None, unit='Gyr'):
    """
    Calculate the age of the universe in the snapshot f by integrating the Friedmann equation.

    The output is given in the specified units. If a redshift z is specified, it is used in place of the redshift in the
    output f.

    If a long array of redshifts is provided, interpolation is used to speed up the calculation. Specifically,
    the number of interpolation points is controlled by the configuration parameter 'cosmo-interpolation-points'.

    Parameters
    ----------

    f : SimSnap
        The snapshot from which to obtain the cosmological parameters.

    z : float, list, or np.ndarray, optional
        The redshift(s) at which to calculate the age of the universe. If None, the redshift of the snapshot is used.

    unit : str, optional
        The units in which to return the age of the universe. Default is 'Gyr'.

    Returns
    -------

    SimArray | float
        The age of the universe at the specified redshift(s).

    """
    if z is None:
        z = f.properties['z']

    h0 = f.properties['h']
    omM = f.properties['omegaM0']
    omL = f.properties['omegaL0']

    conv = units.Unit("0.01 s Mpc km^-1").ratio(unit, **f.conversion_context())

    def get_age(x):
        x = 1.0 / (1.0 + x)
        return scipy.integrate.quad(_a_dot_recip, 0, x, (h0, omM, omL))[0] * conv

    if isinstance(z, np.ndarray) or isinstance(z, list):
        if len(z) > _interp_points:
            a_vals = np.logspace(-3,0, _interp_points)
            z_vals = 1./a_vals-1.
            log_age_vals = np.log(age(f, z_vals))
            interp = interp1d(np.log(a_vals), log_age_vals, bounds_error=False)
            log_a_input = np.log(1./(1.+z))
            results = np.exp(interp(log_a_input))
        else:
            results = np.array([get_age(_z) for _z in z])
        results = results.view(SimArray)
        results.units = unit
        return results
    else:
        return get_age(z)

def tau(f, z=None, unit="Gyr"):
    """
    Calculate the conformal age of the universe in the snapshot f by integrating the conformal Friedmann equation.

    The output is given in the specified units. If a redshift z is specified, it is used in place of the redshift in the
    output f.

    Parameters
    ----------

    f : SimSnap
        The snapshot from which to obtain the cosmological parameters.

    z : float, list, or np.ndarray, optional
        The redshift(s) at which to calculate the age of the universe. If None, the redshift of the snapshot is used.

    unit : str, optional
        The units in which to return the age of the universe. Default is 'Gyr'.

    Returns
    -------

    SimArray | float
        The conformal age of the universe at the specified redshift(s).


    """
    if z is None:
        z = f.properties['z']

    h0 = f.properties['h']
    omM = f.properties['omegaM0']
    omL = f.properties['omegaL0']

    conv = units.Unit("0.01 s Mpc km^-1").ratio(unit, **f.conversion_context())

    @np.vectorize
    def get_tau(z):
        aexp = 1.0 / (1.0 + z)
        return scipy.integrate.quad(
            _da_dtau_recip,
            1,
            aexp,
            args=(h0, omM, omL)
        )[0]

    results = (get_tau(z) * conv)

    if isinstance(results, np.ndarray):
        results = results.view(SimArray)
        results.units = unit

    return results


@units.takes_arg_in_units((1, "Gyr"), context_arg=0)
def redshift(f, time):
    """
    Calculate the redshift given a snapshot and a time since Big Bang in Gyr.

    Uses scipy.optimize.newton to do the root finding if number of
    elements in the time array is less than 1000; otherwise uses a linear
    interpolation.


    Parameters
    ----------

    f : SimSnap
        The snapshot from which to obtain the cosmological parameters.

    time : float, list, or np.ndarray
        The time(s) since the Big Bang at which to calculate the redshift.

    Returns
    -------
    float or np.ndarray
        The redshift(s) corresponding to the input time(s).

    """

    from scipy.interpolate import interp1d
    from scipy.optimize import newton

    def func(x, sim, time):
        return age(sim, x) - time

    if isinstance(time, list) or isinstance(time, np.ndarray):
        if len(time) > _interp_points:
            zs = np.logspace(3, -10, _interp_points)
            ages = age(f, zs)
            i = interp1d(ages, zs)
            return i(time)
        else:
            return np.array([newton(func, 1, args=(f, x)) for x in time])
    else:
        return newton(func, 1, args=(f, time))


def rho_crit(f, z=None, unit=None):
    """Calculate the critical density of the universe in the snapshot f.

    .. warning ::
        You can get slightly confusing results if your simulation is in comoving units and you specify a different
        redshift z. Specifically, the physical density for the redshift you specify is calulated, but expressed as
        a comoving density *at the redshift of the snapshot*. This is intentional, but can be surprising.

    Parameters
    ----------

    f : SimSnap
        The snapshot from which to obtain the cosmological parameters.

    z : float, optional
        The redshift at which to calculate the critical density. If None, the redshift of the snapshot is used.

    unit : str, optional
        The units in which to return the critical density. If None, the returned density will be in the units of
        f["mass"].units/f["pos"].units**3. If that unit cannot be calculated, the returned units are Msol kpc^-3
        comoving.

    Returns
    -------

    float
        The critical density of the universe at the specified redshift.



    """

    if z is None:
        z = f.properties['z']

    if unit is None:
        try:
            unit = f.dm["mass"].units / f.dm["pos"].units ** 3
        except units.UnitsException:
            unit = units.NoUnit()

    if hasattr(unit, "_no_unit"):
        unit = units.Unit("Msol kpc^-3 a^-3")

    omM = f.properties['omegaM0']
    omL = f.properties['omegaL0']
    h0 = f.properties['h']
    a = 1.0 / (1.0 + z)

    H_z = _a_dot(a, h0, omM, omL) / a
    H_z = units.Unit("100 km s^-1 Mpc^-1") * H_z

    rho_crit = (3 * H_z ** 2) / (8 * math.pi * units.G)

    return rho_crit.ratio(unit, **f.conversion_context())


def rho_M(f, z=None, unit=None):
    """Calculate the matter density of the universe in snapshot f.

    .. warning ::
        You can get slightly confusing results if your simulation is in comoving units and you specify a different
        redshift z. Specifically, the physical density for the redshift you specify is calulated, but expressed as
        a comoving density *at the redshift of the snapshot*. This is intentional, but can be surprising.

    Parameters
    ----------

    f : SimSnap
        The snapshot from which to obtain the cosmological parameters.

    z : float, optional
        The redshift at which to calculate the critical density. If None, the redshift of the snapshot is used.

    unit : str, optional
        The units in which to return the matter density. If None, the returned density will be in the units of
        f["mass"].units/f["pos"].units**3. If that unit cannot be calculated, the returned units are Msol kpc^-3
        comoving.

    Returns
    -------

    float
        The matter density of the universe at the specified redshift.

    """

    if z is None:
        z = f.properties['z']

    return f.properties['omegaM0'] * rho_crit(f, 0, unit) * (1.0 + z) ** 3


def H(f):
    """Calculate the Hubble parameter of the universe in snapshot f"""
    return f.properties['h'] * _hzoverh0(f.properties['a'], f.properties['omegaM0']) * units.Unit("100 km s^-1 Mpc^-1")


def add_hubble(f):
    """Add the hubble flow to velocities in snapshot f"""

    f['vel'] += f['pos'] * H(f)


def comoving_to_physical(ar):
    """Given an array, modify it to be in physical units (remove any
    dependencies on a or aform)."""

    a_power = ar.units._power_of("a")
    aform_power = ar.units._power_of("aform")

    if a_power != 0:
        a = ar.sim.properties['a']
        ar *= a ** a_power
        ar /= units.Unit("a") ** a_power
    if aform_power != 0:
        aform = ar.sim['aform']
        ar *= aform ** aform_power
        ar /= units.Unit("aform")
