import numpy as np
from astropy.cosmology import WMAP9 as cosmo
from astropy import units as u


def flux_sig_to_mbh(flux, sigma, redshift, ref='default', cosmology=cosmo, flux_err=None, sigma_err=None):
    """
    Function to calculate virial estimates of Mbh (???? paper).
    
    @param flux: broad Halpha component flux in SDSS units (1e-17 erg / s / cm2)
    @param sigma: of broad line in km/s
    @param redshift
    @param ref: literature reference to use (calibrations are slightly different): greene05, greene07, reines13, default, see config below
    @param cosmology: adopted cosmology (astropy object, see http://docs.astropy.org/en/stable/cosmology/); default is WMAP9; if None, use RGG method of di
    @param flux_err: optional error to flux argument in the same units
    @param sigma_err: optional error to sigma argument in the same units
    @return: tuple (mass, mass_err) where mass of BH is in solar mass; if no flux_err and sigma_err uncertainties given, mass_err is None
    
    Example:
    >>> flux_sig_to_mbh(100, 1371/2.355, 0.005290, ref='reines13', flux_err=10, sigma_err=20)
    (<Quantity 75683.25613424949 msolMass>, <Quantity 6429.704407108592 msolMass>)
    """
    config = {
        'greene05': {'lum_power': 0.55, 'fwhm_power': 2.06, 'mass_factor': 2.0e6},
        'greene07': {'lum_power': 0.45, 'fwhm_power': 2.06, 'mass_factor': 3.0e6},
        'reines13': {'lum_power': 0.47, 'fwhm_power': 2.06, 'mass_factor': 3.72e6},
        'default': {'lum_power': 0.46, 'fwhm_power': 2.06, 'mass_factor': 2.24e6}, # not clear yet where does it come from
    }
    c = config[ref]
    if cosmology != None:
        lum_distance =  cosmology.luminosity_distance(redshift).to(u.cm)
    else: # this is how RGG and Baldassare15 compute distances...
        lum_distance = (299792.458 * redshift / 73 * u.Mpc).to(u.cm)
    lum = 4 * np.pi * lum_distance**2 * (flux * (u.erg / u.s / u.cm**2)) * 1e-17
    term1 = np.power(lum / (1.0e42 * u.erg / u.s), c['lum_power'])
    term2 = np.power((sigma * u.km / u.s) * 2.355 / (1e3 * u.km / u.s), c['fwhm_power'])
    mass = (c['mass_factor'] * u.msolMass) * term1 * term2
    mass_err = None
    if flux_err and sigma_err:
        mass_err = mass * ( (c['lum_power'] * flux_err / flux)**2 + (c['fwhm_power'] * sigma_err / sigma)**2 )**0.5
    return (mass, mass_err)


def sig_vs_flux(flux, mbh, redshift, ref='default', cosmology=cosmo):
    """
    Function to calculate sigma vs. flux for given Mbh.

    @param mbh: scalar
    @param flux: numpy array with Halpha BLR fluxes
    @return numpy array of sigmas:

    """
    config = {
        'greene05': {'lum_power': 0.55, 'fwhm_power': 2.06, 'mass_factor': 2.0e6},
        'greene07': {'lum_power': 0.45, 'fwhm_power': 2.06, 'mass_factor': 3.0e6},
        'reines13': {'lum_power': 0.47, 'fwhm_power': 2.06, 'mass_factor': 3.72e6},
        'default': {'lum_power': 0.46, 'fwhm_power': 2.06, 'mass_factor': 2.24e6}, # not clear yet where does it come from
    }
    c = config[ref]
    if cosmology != None:
        lum_distance =  cosmology.luminosity_distance(redshift).to(u.cm)
    else: # this is how RGG and Baldassare15 compute distances...
        lum_distance = (299792.458 * redshift / 73 * u.Mpc).to(u.cm)
    lum = 4 * np.pi * lum_distance**2 * (flux * (u.erg / u.s / u.cm**2)) * 1e-17

    sig = (
        (1e3 * u.km / u.s) / 2.355 * 
        np.power( mbh / c['mass_factor'], 1.0 / c['fwhm_power']) *
        np.power( lum / ( 1.0e42 * u.erg / u.s), -c['lum_power']/c['fwhm_power'])
        )

    return sig