mors.miscellaneous

def ActivityLifetime(Age=None,Track=None,Threshold=None,AgeMax=None):
    """
    Takes evolutionary track for parameter, calculates when value drops below threshold.

    This function can be used to determine when a star's emission crosses a given threshold value for a few
    activity quantities. These are Lx, Fx, Rx, and FxHZ for X-rays, and similar values for EUV1, EUV2, EUV,
    XUV, and Ly-alpha. If the star crosses the threshold (from above it to below it) multiple times, this
    will find the final time it will cross the threshold. If the user wants to set a maximum age so that the
    code only looks for crossings of the threshold below this age then this can be done using the AgeMax
    keyword argument. If the track is always below this threshold value then the function returns 0.0.

    Parameters
    ----------
    Age : numpy.ndarray
        Age array for evolutionary track.
    Track : numpy.ndarray
        Value array for evolutionary track.
    Threshold : float
        Value for threshold in units of quantity.
    AgeMax : float , optional
        End age of track to consider in Myr.

    Returns
    ----------
    AgeActive : float
        Activity lifetime in Myr.

    """

    # Make sure parameters set
    if Age is None:
        raise Exception("required argument Age not set")
    if Track is None:
        raise Exception("required argument Track not set")
    if Threshold is None:
        raise Exception("required argument Threshold not set")

    # Make sure Age and Track same length
    if not ( len(Age) == len(Track) ):
        raise Exception("Age and Track are different lengths")

    # Change Age and Track to only be ages below AgeMax, if AgeMax is set
    if not AgeMax is None:
        includeAges = np.where( Age <= AgeMax )
        Age = Age[includeAges]
        Track = Track[includeAges]

    # See if final value below threshold, otherwise return final age
    if ( Track[-1] > Threshold ):
        return Age[-1]

    # Start at end of array and loop backwards
    for iAge in range(len(Age)-1,0,-1):

        # Make sure next one back is not equal to threshold
        if ( Track[iAge-1] == Threshold ):
            return Age[iAge-1]

        # Check if this age bin is when it drops past the threshold
        if ( ( Track[iAge] < Threshold ) and ( Track[iAge-1] > Threshold ) ):

            # Do interpoaltion to get when it crosses
            # Assume here Age = m*Track + c
            mInterp = ( Age[iAge] - Age[iAge-1] ) / ( Track[iAge] - Track[iAge-1] )
            cInterp = Age[iAge] - mInterp * Track[iAge]
            return mInterp*Threshold + cInterp

    return 0.0

def IntegrateEmission(AgeMin=None,AgeMax=None,Age=None,Luminosity=None,aOrb=None):
    """
    Takes evolutionary track for parameter, calculates when value drops below threshold.

    This code can be used to integrate a star's luminosity between two ages. This can be applied to Lx, Leuv, or any
    other luminosity and the result is a total energy emitted in this time in erg. If the user also specifies an orbital
    distance using the aOrb keyword argument, the code integrates the flux at this obital distance returns a fluence
    in erg cm^-2.

    Parameters
    ----------
    AgeMin : float
        Start of time period to integrate in Myr.
    AgeMax : float
        End of time period to integrate in Myr.
    Age : numpy.ndarray
        Age array for evolutionary track.
    Luminosity : numpy.ndarray
        Value array for evolutionary track.
    aOrb : float , optional
        Orbital distance to get fluence at in AU.

    Returns
    ----------
    Energy : float
        Integrated luminosity of flux in erg or erg cm^-2.

    """

    # Make sure parameters set
    if AgeMin is None:
        raise Exception("required argument AgeMin not set")
    if AgeMax is None:
        raise Exception("required argument AgeMax not set")
    if Age is None:
        raise Exception("required argument Age not set")
    if Luminosity is None:
        raise Exception("required argument Luminosity not set")

    # Make sure Age and Luminosity same length
    if not ( len(Age) == len(Luminosity) ):
        raise Exception("Age and Luminosity are different lengths")

    # Make sure AgeMin and AgeMax is in track
    if not ( ( AgeMin >= Age[0] ) and ( AgeMin <= Age[-1] ) ):
        raise Exception("AgeMin not in range of evolutionary track")
    if not ( ( AgeMax >= Age[0] ) and ( AgeMax <= Age[-1] ) ):
        raise Exception("AgeMax not in range of evolutionary track")

    # Get track to actually integrate
    Track = copy.deepcopy(Luminosity)
    if not aOrb is None:
        Track *= 1.0 / ( 4.0 * const.Pi * (aOrb*const.AU)**2.0 )

    # Get index of first age bin with age above AgeMin
    indexMin = _getIndexGT(Age,AgeMin)

    # Get index of final age bin with age below AgeMax
    indexMax = _getIndexLT(Age,AgeMax)

    # Initially take no energy then add up energy from bins
    Energy = 0.0

    # Add energy from first bin if needed (simple assumption for luminosity constant ove age bin)
    if ( AgeMin < Age[indexMin] ):
        Energy += ( Age[indexMin] - AgeMin ) * Track[indexMin]

    # Intergate over all full bins (only if there is something to integrate)
    if not ( indexMin == indexMax ):
        for iAge in range(indexMin,indexMax):

            # Do integration using trapezoidal rule (not the chained version)
            Energy += 0.5 * ( Age[iAge+1] - Age[iAge] ) * ( Track[iAge] + Track[iAge+1] )

    # Add energy from final bin if needed (simple assumption for luminosity constant ove age bin)
    if ( AgeMax > Age[indexMax] ):
        Energy += ( AgeMax - Age[indexMax] ) * Track[indexMax]

    # So far, units of Age were in Myr when integrating, so include also Myr to s conversion
    Energy *= const.Myr

    return Energy

def Load(filename):
    """Takes filename of saved star or cluster and loads object."""

    # Load the object
    with open(filename,'rb') as f:
        obj = pickle.load(f)

    # Setup the quantity functions (which for some reason do not work on objects loaded from pickle)
    obj._setupQuantityFunctions()

    return obj

def ModelCluster():
    """Reads the model cluster used in Johnstone et al. (2020)."""

    # Get directory where package is installed
    packageDir = _GetPackageDirectory()

    # Set filename including path
    filename = packageDir + "ModelDistribution.dat"

    # Read lines of file
    with open(filename) as f:
        content = f.readlines()

    # Set number of header lines
    nHeader = 1

    # Get number of steps in file
    nStars = len(content) - nHeader

    # Make arrays
    Mstar = np.zeros(nStars)
    Omega = np.zeros(nStars)

    # Loop over lines and fill arrays
    iStar = 0
    for line in content[nHeader:len(content)]:
        data = line.split()
        Mstar[iStar] = data[0]
        Omega[iStar] = data[1]
        iStar += 1

    return Mstar , Omega

def _GetPackageDirectory():
    """Gets path to main directory where the code is installed."""

    return str(pathlib.Path(__file__).parent.absolute())+"/data/"

def _convertFloatArray(Xin):
    """Takes a value and returns either as float or numpy.ndarray of floats."""

    # If input is float or int, return float version
    if isinstance(Xin,(float,int)):
        return float(Xin)

    # If input is a list, numpy array version (will only work if elements can be converted into floats)
    if isinstance(Xin,list):
        X = np.zeros(len(Xin))
        for i in range(0,len(Xin)):
            X[i] = float(Xin[i])
        return X

    # If input is a numpy object then do stuff
    if ( type(Xin).__module__ == 'numpy' ):

        # First see if it can be converted into a float (i.e. a 1-element array)
        isNumber = True
        try:
            X = float(Xin)
        except:
            isNumber = False

        # If it is just a single value, return float version
        if ( isNumber ):
            return float(X)

        # Otherwise, turn it into an array
        X = np.zeros(len(Xin))
        for i in range(0,len(Xin)):
            X[i] = float(Xin[i])
        return X

    return

def _getIndexGT(Xarray,X):
    """Takes array of values and a value, returns index of closest element in array larger than value."""

    # It is assumed that X is less than the max of Xarray, but check this
    if ( X > np.max(Xarray) ):
        raise Exception("X is more than maximum of Xarray")

    # Get smaller value by taking difference of X and all Xarray elements
    # then removing all values smaller than 0.0 (i.e. those larger than X)
    deltaX = Xarray - X
    deltaX[np.where(deltaX<0.0)] = np.max(deltaX)*1.1+0.1

    # Get index
    index = np.argmin(deltaX)

    return index

def _getIndexLT(Xarray,X):
    """Takes array of values and a value, returns index of closest element in array smaller than value."""

    # It is assumed that X is greater than the min of Xarray, but check this
    if ( X < np.min(Xarray) ):
        raise Exception("X is less than minimum of Xarray")

    # Get smaller value by taking difference of X and all Xarray elements
    # then removing all values smaller than 0.0 (i.e. those larger than X)
    deltaX = X - Xarray
    deltaX[np.where(deltaX<0.0)] = np.max(deltaX)*1.1+0.1

    # Get index
    index = np.argmin(deltaX)

    return index

def _getIndexLTordered(Xarray,X):
    """Takes min to max ordered array of values and a value, returns index of closest element in array smaller than value."""

    # It is assumed that X is between the min and max of Xarray and it is assumed
    # that Xarray is in ascending order

    # Initial guesses for i
    i1 = 0
    i2 = len(Xarray)-1

    # Check if right value
    if ( Xarray[i1] == X ):
        return i1
    if ( Xarray[i2] == X ):
        return i2

    # Start iterating
    for iIter in range(0,len(Xarray)):

        # Get iMid
        iMid = int(0.5*(i1+i2))

        # Check if iMid is right value
        if ( ( Xarray[iMid] <= X ) and ( Xarray[iMid+1] > X ) ):
            return iMid

        # Work out if answer is between i1 and iMid or iMid and i2
        if ( Xarray[iMid] > X ):
            i2 = iMid
        else:
            i1 = iMid

    # It should not get here
    raise Exception("did not find index")

    return iMid