mors.synthesis

def CalcBandScales(modern_dict:dict, historical_dict):
    """Get band scale factors for historical spectrum

    Parameters
    ----------
    modern_dict : dict
        Dictionary output of `GetProperties` call for modern spectrum
    historical_dict : dict
        Dictionary output of `GetProperties` call for historical spectrum

    Returns
    -------
    Q_dict : dict
        Dictionary of band scale factors
    """

    # Get scale factors
    Q_dict = {}
    for key in spec.bands_limits.keys():
        Q_dict["Q_"+key] = historical_dict["F_"+key]/modern_dict["F_"+key]

    return Q_dict

def CalcScaledSpectrumFromProps(modern_spec:spec.Spectrum, modern_dict:dict, historical_dict:dict):
    """Scale a stellar spectrum according to stellar properties.

    Returns a new Spectrum object containing the historical fluxes.

    Parameters
    ----------
    modern_spec : Spectrum object
        Spectrum object containing data for a modern fluxes
    modern_dict : dict
        Dictionary output of `GetProperties` call for modern spectrum
    historical_dict : dict
        Dictionary output of `GetProperties` call for historical spectrum

    Returns
    -------
    historical_spec : Spectrum object
        Spectrum object containing data for historical fluxes
    """

    log.debug("Calculating scaled spectrum from properties")

    # Get scale factors relative to modern spectrum
    Q_dict = CalcBandScales(modern_dict, historical_dict)

    # Get modern wl, fl
    spec_fl = deepcopy(modern_spec.fl)
    spec_wl = deepcopy(modern_spec.wl)

    # Get band indicies
    for i in range(len(spec_wl)):
        b = spec.WhichBand(spec_wl[i])
        if b == None:
            continue
        spec_fl[i] *= Q_dict["Q_"+b[0]]

    # Make new spectrum object
    historical_spec = spec.Spectrum()
    historical_spec.LoadDirectly(spec_wl, spec_fl)

    return historical_spec

def FitModernProperties(modern_spec:spec.Spectrum, Mstar:float, age:float=-1):
    """Estimate rotation percentile and (optionally) age from modern spectrum.

    Parameters
    ----------
    modern_spec : Spectrum object
        Spectrum object containing data for a modern fluxes
    Mstar : float
        Stellar mass [M_sun]
    age : float
        Optional guess for current age. Will be estimated if not provided.

    Returns
    -------
    best_pctle : float
        Best estimate of rotation percentile
    best_age : float
        Best estimate of star's age (equal to `age` if `age` is provided)
    """

    log.debug("Fitting properties to modern spectrum")

    # Integrated fluxes
    modern_spec.CalcBandFluxes()

    fit_age = bool(age>0)

    # Objective function
    def _fev(x_arr:tuple):

        this_pctle = x_arr[0]
        if fit_age:
            this_age = x_arr[1]
        else:
            this_age = age

        this_age = max(min(this_age, 11000), 0.4)

        props = GetProperties(Mstar, this_pctle, this_age)

        resid = 0.0
        for k in ["xr","e1","e2","uv"]:
            lim = spec.bands_limits[k]
            wid = lim[1]-lim[0]
            resid += (props["F_"+k]/wid - modern_spec.fl_integ[k]/wid)**2

        return np.sqrt(resid)

    # Initial guess
    if fit_age:
        x0 = [1.0, 1000.0]
    else:
        x0 = [1.0]

    # Find best params
    result = minimize(_fev, x0, method='Nelder-Mead')

    # Check
    if not result.success:
        log.error("Could not fit stellar properties to modern spectrum")

    # Result
    best_pctle = result.x[0]
    if fit_age:
        best_age = result.x[1]
    else:
        best_age = age

    # Return
    return best_pctle, best_age

def GetProperties(Mstar:float, pctle:float, age:float):
    """Calculate properties of star for a given rotation percentile and  age

    Parameters
    ----------
    Mstar : float
        Mass of star [M_sun]
    pctle : float
        Rotation percentile
    age : float
        Stellar age  [Myr]

    Returns
    -------
    out : dict
        Dictionary of radius [m], Teff [K], and band fluxes at 1 AU [erg s-1 cm-2]
    """

    # Get star radius [m]
    Rstar = Value(Mstar, age, 'Rstar') * const.Rsun * 1.0e-2

    # Get star temperature [K]
    Tstar = Value(Mstar, age, 'Teff')

    # Get rotation rate [Omega_sun]
    Omega = Percentile(Mstar=Mstar, percentile=pctle)

    # Get luminosities and fluxes
    Ldict = Lxuv(Mstar=Mstar, Age=age, Omega=Omega)

    # Output
    out = {
        "age"    : age,        # units of Myr
        "radius" : Rstar,
        "Teff"   : Tstar,
    }

    # Luminosities (erg s-1)
    out["L_bo"] = Lbol(Mstar,age) * const.LbolSun
    out["L_xr"] = Ldict["Lx"]
    out["L_e1"] = Ldict["Leuv1"]
    out["L_e2"] = Ldict["Leuv2"]

    # Fluxes at 1 AU
    area = (4.0 * const.Pi * const.AU * const.AU)
    for k in ["bo","xr","e1","e2"]:
        out["F_"+k] = out["L_"+k]/area

    # Get flux from Planckian band
    wl_pl = np.logspace(np.log10(spec.bands_limits["pl"][0]), np.log10(spec.bands_limits["pl"][1]), 1000)
    fl_pl = spec.PlanckFunction_surf(wl_pl, Tstar)
    fl_pl = spec.ScaleTo1AU(fl_pl, Rstar)
    out["F_pl"] = np.trapezoid(fl_pl, wl_pl)
    out["L_pl"] = out["F_pl"] * area

    # Get flux of UV band from remainder
    out["F_uv"] = out["F_bo"] - out["F_xr"] - out["F_e1"] - out["F_e2"] - out["F_pl"]
    out["L_uv"] = out["F_uv"] * area

    return out

def Lbol(Mstar,Age,ModelData=ModelDataDefault):
    """Takes mass and age, returns bolometric luminosity."""
    return Value( Mstar , Age , 'Lbol' , ModelData=ModelData )

def Lxuv(Mstar=None,Age=None,Omega=None,OmegaEnv=None,Prot=None,params=params.paramsDefault,StarEvo=None):
    """Takes basic stellar parameters, returns X-ray, EUV, and Ly-alpha luminosities in erg s^-1.

    The user can supply a mass, age, and surface rotation rate for a star and it returns the XUV luminosities as
    a dictionary. The rotation rate can be either the rotation velocity as a multiple of the solar rotation
    rate (2.67e-6 rad s^-1) using either Omega or OmegaEnv keyword arguments, or as a rotation period in days using
    the Prot keyword argument, The user can optionally also give a parameter dictionary and an instance of the
    StarEvo class, though this is not necessary and if these are not specified then the defaults will be used.

    Parameters
    ----------
    Mstar : float
        Mass of star in Msun.
    Age : float
        Age of star.
    Omega : float , optional
        Surface rotation rate in OmegaSun (=2.67e-6 rad/s).
    OmegaEnv : float , optional
        Surface rotation rate in OmegaSun (=2.67e-6 rad/s).
    Prot : float , optional
        Surface rotation period in days.
    params : dict , optional
        Dictionary holding model parameters.
    StarEvo : mors.stellarevo.StarEvo , optional
        Instance of StarEvo class holding stellar evolution model data.

    Returns
    -------
    LxuvDict : dict
        Dictionary of luminosities in erg s^-1.

    """

    # If an instance of the StarEvo class is not input, load it with defaults
    if StarEvo is None:
        StarEvo = SE.StarEvo()

    # Make sure mass and age are specified
    if Mstar is None:
        raise Exception("Mstar must be set in call to function")
    if Age is None:
        raise Exception("Age must be set in call to function")

    # Make sure rotation is set
    if ( Omega is None ) and ( OmegaEnv is None ) and ( Prot is None ):
        raise Exception("Omega, OmegaEnv, or Prot must be set in call to function")

    # Make sure only one rotation rate is set (add up number of set parameters)
    nSet = 0
    if not Omega is None:
        nSet += 1
    if not OmegaEnv is None:
        nSet += 1
    if not Prot is None:
        nSet += 1
    if not ( nSet == 1 ):
        raise Exception("can only set one of Omega, OmegaEnv, and Prot in call to function")

    # Get Omega not set
    if not OmegaEnv is None:
        Omega = OmegaEnv
    if not Prot is None:
        Omega = _Omega(Prot)

    # The function _Xray needs Ro, Rstar, and Lbol so make a StarState dictionary with these
    StarState = {}
    StarState['OmegaEnv'] = Omega
    StarState['Rstar'] = StarEvo.Rstar(Mstar,Age)
    StarState['Prot'] = _Prot(Omega)
    StarState['tauConv'] = StarEvo.tauConv(Mstar,Age)
    StarState['Ro'] = _Ro(StarState)
    StarState['Lbol'] = StarEvo.Lbol(Mstar,Age) * const.LbolSun

    # Get X-ray
    StarState['Lx'] , StarState['Fx'] , StarState['Rx'] = _Xray(StarState,params=params)

    # Get EUV
    StarState['Leuv1'] , StarState['Feuv1'] , StarState['Reuv1'] = _EUV1(StarState,params=params)
    StarState['Leuv2'] , StarState['Feuv2'] , StarState['Reuv2'] = _EUV2(StarState,params=params)
    StarState['Leuv'] , StarState['Feuv'] , StarState['Reuv'] = _EUV(StarState,params=params)

    # Get Ly-alpha
    StarState['Lly'] , StarState['Fly'] , StarState['Rly'] = _Lymanalpha(StarState,params=params)

    # Make dictionary with luminosities
    LxuvDict = {}
    LxuvDict['Lxuv'] = StarState['Lx'] + StarState['Leuv']
    LxuvDict['Lx'] = StarState['Lx']
    LxuvDict['Leuv1'] = StarState['Leuv1']
    LxuvDict['Leuv2'] = StarState['Leuv2']
    LxuvDict['Leuv'] = StarState['Leuv']
    LxuvDict['Lly'] = StarState['Lly']

    LxuvDict['Fxuv'] = StarState['Fx'] + StarState['Feuv']
    LxuvDict['Fx'] = StarState['Fx']
    LxuvDict['Feuv1'] = StarState['Feuv1']
    LxuvDict['Feuv2'] = StarState['Feuv2']
    LxuvDict['Feuv'] = StarState['Feuv']
    LxuvDict['Fly'] = StarState['Fly']

    LxuvDict['Rxuv'] = StarState['Rx'] + StarState['Reuv']
    LxuvDict['Rx'] = StarState['Rx']
    LxuvDict['Reuv1'] = StarState['Reuv1']
    LxuvDict['Reuv2'] = StarState['Reuv2']
    LxuvDict['Reuv'] = StarState['Reuv']
    LxuvDict['Rly'] = StarState['Rly']

    return LxuvDict

def Percentile(Mstar=None,Omega=None,Prot=None,percentile=None,MstarDist=None,OmegaDist=None,ProtDist=None,params=params.paramsDefault):
    """Gets rotation rate of percentile or percentile of rotation rate in the model rotation distribution.

    This function can be used for two purposes
        1. to determine the percentile in a rotation distribution of a star given its mass and rotation rate
        2. to determine the rotation rate of a star in a rotation distribution given its mass and percentile
    In the first case, the user should specify the rotation rate using either the Omega or Prot keyword arguments (given
    in OmegaSun and days respectively) and the mass. In the second case, the user should specify the percentile using the
    percentile keyword argument (in the range 0 to 100) and the mass. The mass should be specified in Msun using the Mstar
    keyword argument. These should only be given as floats. The user can also specify the rotation distribution using the
    MstarDist and the OmegaDist or ProtDist keyword arguments. If this is not done, then the code will assume the 1 Myr
    rotation distribution from the empirical model distribution used in Johnstone et al. (2020).

    Parameters
    ----------
    Mstar : float
        Stellar mass in Msun.
    Omega : float , optional
        Rotation rate of star in OmegaSun.
    Prot : float , optional
        Rotation period of star in days.
    percentile : float , optional
        percentile in distribution (between 0 and 100).
    MstarDist : numpy.ndarray , optional
        Array of masses of stars in distribution.
    OmegaDist : numpy.ndarray , optional
        Array of rotation rates of stars in distribution in OmegaSun.
    ProtDist : numpy.ndarray , optional
        Array of rotation periods of stars in distribution in days.
    params : dict , optional
        Parameter dictionary used for getting width of bin to use for distribution.

    Returns
    ----------
    result : float
        Either rotation rate for percentile or percentile for rotation rate.

    """

    # Make sure Mstar was set
    if Mstar is None:
        raise Exception( "keyword argument Mstar must be set" )

    # If percentile is not set, make sure rotation rate is set
    if percentile is None:
        # Not both
        if not Omega is None:
            if not Prot is None:
                raise Exception( "keyword arguments Omega and Prot cannot both be set" )
        # At least one
        if ( Omega is None ) and ( Prot is None ):
            raise Exception( "must set either Omega, Prot, or percentile keyword arguments" )
        # Set Omega is not set then set it
        if Omega is None:
            Omega = phys._Omega(Prot)

    # Setup the distribution
    if MstarDist is None:
        MstarDist , OmegaDist = misc.ModelCluster()
    else:

        # Check that one of OmegaDist and ProtDist are set
        if ( OmegaDist is None ) and ( ProtDist is None ):
            raise Exception( "one of keyword arguments OmegaDist and ProtDist must be set" )

        # Check that OmegaDist and ProtDist are not both set
        if ( not OmegaDist is None ) and ( not ProtDist is None ):
            raise Exception( "keyword arguments OmegaDist and ProtDist cannot both be set" )

        # Setup OmegaDist if ProtDist was set
        if OmegaDist is None:
            OmegaDist = misc._Omega(ProtDist)

        # Make sure arrays are same length
        if not ( len(MstarDist) == len(OmegaDist) ):
            raise Exception( "MstarDist and OmegaDist (or ProtDist) must be same length" )

    # Work out which one to do
    if not percentile is None:

        # percentile was set, so get corresponding rotation rates
        result = _OmegaPercentile(Mstar,percentile,MstarDist,OmegaDist,params)

    else:

        # Rotation rate was set, so get corresponding percentile
        result = _PerPercentile(Mstar,Omega,MstarDist,OmegaDist,params)


    return result

def Value(MstarIn,AgeIn,ParamString,ModelData=ModelDataDefault):
    """Takes stellar mass, age, and a parameter string, returns values corresponding to named parameter.

    The set of models should have already been loaded. With this function, the user can ask for a value
    of one of the parameters for a specific stellar mass and age. All three of these can be input as multiple
    values and an array of values will be returned if this is the case. For example, if MstarIn is input as
    a 1D array, the function will return a 1D array giving the value for each of these masses. ParamString
    can be input as a list of strings and values for each parameter in that list will be returned in a 1D
    array. If two are given as arrays or lists, then a 2D array will be returned. If all three then a 3D
    array with dimensions len(MstarIn)xlen(AgeIn)xlen(ParamString) will be returned.

    Parameters
    ----------
    MstarIn : float or int or numpy.ndarray
        Mass of star in Msun.
    AgeIn : float or int or numpy.ndarray
        Age in Myr.
    ParamString : str
        String holding name of parameter to get value for.
    ModelData : dict , optional
        Dictionary of dictionaries holding set of stellar evolution models.

    Returns
    -------
    value : float or numpy.ndarray
        Value of parameter at this mass and age.

    """

    # Check if ModelData is not None, otherwise need to load defaults
    if ModelData is None:
        ModelData = _LoadDefaultModelData()

    # There are now eight scenarios here
    #   1. Mstar and Age are both floats, so return float.
    #   2. Mstar is float and Age is numpy.ndarray, so return numpy.ndarray of length len(Age).
    #   3. Mstar is numpy.ndarray and Age is float, so return numpy.ndarray of length len(Age)
    #   4. Mstar and Age are both numpy.ndarray, so return numpy.ndarray of shape ( len(Mstar) , len(Age) ).
    # Note: could use type(X).__module__ == 'numpy' to test for numpy but it doesn't work if

    # Get Mstar and Age in correct types
    Mstar = misc._convertFloatArray(MstarIn)
    Age = misc._convertFloatArray(AgeIn)

    # Make sure Mstar and Age are correct types
    if Mstar is None:
        raise Exception("argument Mstar has invalid type")
    if Age is None:
        raise Exception("argument Age has invalid type")

    # Find if scenario 1 (most likely)
    if ( isinstance(Mstar,float) and isinstance(Age,float) and isinstance(ParamString,str) ):

        # Scenario 1 so just get value
        value = _ValueSingle( Mstar , Age , ParamString , ModelData=ModelData )

    else:

        # In this case, the output will be an array with up to three dimensions, so first make 3D array
        # and then remove dimensions with only one element

        # Get number of elements in all three directions
        try:
            nMstar = len(Mstar)
        except:
            nMstar = 1

        try:
            nAge = len(Age)
        except:
            nAge = 1

        if isinstance(ParamString,list):
            nParam = len(ParamString)
        else:
            nParam = 1

        # Make array
        value = np.zeros((nMstar,nAge,nParam))

        # Loop over values
        for iMstar in range(0,nMstar):
            for iAge in range(0,nAge):
                for iParam in range(0,nParam):

                    # Get values of Mstar, Age, and ParamString
                    if isinstance(Mstar,float):
                        MstarValue = Mstar
                    else:
                        MstarValue = Mstar[iMstar]

                    if isinstance(Age,float):
                        AgeValue = Age
                    else:
                        AgeValue = Age[iAge]

                    if isinstance(ParamString,str):
                        ParamStringValue = ParamString
                    else:
                        ParamStringValue = ParamString[iParam]

                    # Now get the value
                    value[iMstar,iAge,iParam] = _ValueSingle( MstarValue , AgeValue , ParamStringValue , ModelData=ModelData )

        # Now get rid of unwanted dimensions
        value = np.squeeze(value)

    return value