import numpy as np
try: # numba will provide a ~2x speedup
    from numba import jit
except ImportError: # but we can do without it
    jit = lambda fn: fn

@jit
def massTorres(teff, erteff, logg, erlogg, feh, erfeh,
               ntrials=10000, corrected=True, add_intrinsic=True):
    """
    Calculate stellar mass using the Torres et al. (2010) callibration.

    Parameters
    ----------
    teff, erteff : floats
        Effective temperature and associated uncertainty.
    logg, erlogg : floats
        Surface gravity and associated uncertainty.
    feh, erfeh : floats
        Metallicity [Fe/H] and associated uncertainty.
    ntrials : int, optional
        Number of Monte Carlo trials for the uncertainty calculation.
    corrected : bool, optional
        Correct the mass for the offset relative to isochrone-derived masses.
    add_intrinsic : bool, optional
        Add (quadratically) the intrinsic error of the calibration
        (0.027 in log mass).

    Returns
    -------
    meanMass, sigMass : floats
        Estimate for the stellar mass and associated uncertainty.
    """

    randomteff = teff + erteff * np.random.randn(ntrials)
    randomlogg = logg + erlogg * np.random.randn(ntrials)
    randomfeh = feh + erfeh * np.random.randn(ntrials)

    # Parameters for the Torres calibration
    a1 = 1.5689
    a2 = 1.3787
    a3 = 0.4243
    a4 = 1.139
    a5 = -0.1425
    a6 = 0.01969
    a7 = 0.1010

    X = np.log10(randomteff) - 4.1

    logMass = a1 + a2*X + a3*X**2 + a4*X**3 + a5*randomlogg**2 \
              + a6*randomlogg**3 + a7*randomfeh

    meanlogMass = np.mean(logMass)
    siglogMass = np.sum((logMass - meanlogMass)**2) / (ntrials - 1)
    if add_intrinsic:
        siglogMass = np.sqrt(0.027*0.027 + siglogMass)

    meanMass = 10**meanlogMass
    sigMass = 10**(meanlogMass + siglogMass) - meanMass

    if corrected:
        # correction comes from Santos+(2013), the SWEET-Cat paper
        corrected_meanMass = 0.791 * meanMass**2 - 0.575 * meanMass + 0.701
        return corrected_meanMass, sigMass
    else:
        return meanMass, sigMass