<reponame>jaysonfig/EXOSIMS
from EXOSIMS.Observatory.SotoStarshade_SKi import SotoStarshade_SKi
import numpy as np
import astropy.units as u
from scipy.integrate import solve_ivp
import astropy.constants as const
import hashlib
import scipy.optimize as optimize
from scipy.optimize import basinhopping
import scipy.interpolate as interp
import scipy.integrate as intg
from scipy.integrate import solve_bvp
from copy import deepcopy
import time
import os
try:
    import _pickle as pickle
except:
    import pickle

EPS = np.finfo(float).eps


class SotoStarshade_ContThrust(SotoStarshade_SKi):
    """ StarShade Observatory class
    This class is implemented at L2 and contains all variables, functions, 
    and integrators to calculate occulter dynamics. 
    """
    
    def __init__(self,orbit_datapath=None,**specs): 

        SotoStarshade_SKi.__init__(self,**specs)  
        
        #to convert from dimensionalized to normalized, (Dimension) / self.(Dimension)U
        self.mass = 10930*u.kg
        
        Tmax = 1040*u.mN
        Isp  = 3000*u.s
        ve   = const.g0 * Isp
        
        self.Tmax = Tmax
        self.ve   = self.convertVel_to_canonical(ve)
        
        # smoothing factor (eps = 1 means min energy, eps = 0 means min fuel)
        self.epsilon = 1 
        
        self.lagrangeMults = self.lagrangeMult()
        
# =============================================================================
# Miscellaneous        
# =============================================================================
    def DCM_r2i(self,t):
        """ Direction cosine matrix to rotate from Rotating Frame to Inertial Frame
        
        Finds rotation matrix for positions and velocities (6x6)

        Args:
            t (float):
                Rotation angle

        Returns:
            float 6x6 array:
                rotation of full 6 dimensional state from R to I frame
        """

        Ts  = np.array([[ np.cos(t) ,-np.sin(t) , 0],\
                        [ np.sin(t) , np.cos(t) , 0],\
                        [   0       ,    0      , 1]])

        dTs = np.array([[-np.sin(t) ,-np.cos(t) , 0],\
                        [ np.cos(t) ,-np.sin(t) , 0],\
                        [   0       ,    0      , 0]])

        Qu = np.hstack( [ Ts , np.zeros([3,3])])
        Ql = np.hstack( [dTs , Ts])

        rotMatrix = np.vstack([Qu,Ql])

        return rotMatrix


    def DCM_i2r(self,t):
        """ Direction cosine matrix to rotate from Inertial Frame to Rotating Frame
        
        Finds rotation matrix for positions and velocities (6x6)

        Args:
            t (float):
                Rotation angle

        Returns:
            float 6x6 array:
                rotation of full 6 dimensional state from I to R frame
        """

        Ts  = np.array([[ np.cos(t) ,-np.sin(t) , 0],\
                        [ np.sin(t) , np.cos(t) , 0],\
                        [   0       ,    0      , 1]])

        dTs = np.array([[-np.sin(t) ,-np.cos(t) , 0],\
                        [ np.cos(t) ,-np.sin(t) , 0],\
                        [   0       ,    0      , 0]])

        argD = np.matmul( np.matmul(-Ts.T,dTs) , Ts.T )
        Qu = np.hstack( [  Ts.T , np.zeros([3,3])])
        Ql = np.hstack( [  argD , Ts.T])

        return np.vstack([Qu,Ql])



    def DCM_r2i_9(self,t):
        """ Direction cosine matrix to rotate from Rotating Frame to Inertial Frame

        Finds rotation matrix for positions, velocities, and accelerations (9x9)

        Args:
            t (float):
                Rotation angle

        Returns:
            float 9x9 array:
                rotation of full 9 dimensional state from R to I frame
        """

        Ts  = np.array([[ np.cos(t) ,-np.sin(t) , 0],\
                        [ np.sin(t) , np.cos(t) , 0],\
                        [   0       ,    0      , 1]])

        dTs = np.array([[-np.sin(t) ,-np.cos(t) , 0],\
                        [ np.cos(t) ,-np.sin(t) , 0],\
                        [   0       ,    0      , 0]])

        ddTs = -np.array([[ np.cos(t) ,-np.sin(t) , 0],\
                          [ np.sin(t) , np.cos(t) , 0],\
                          [   0       ,    0      , 0]])

        Qu = np.hstack( [  Ts ,  np.zeros([3,6])                 ])
        Qm = np.hstack( [ dTs ,  Ts             , np.zeros([3,3])])
        Ql = np.hstack( [ddTs ,2*dTs             , Ts])

        return np.vstack([Qu,Qm,Ql])


    def lagrangeMult(self):
        """ Generate a random lagrange multiplier for initial guess (6x1)

        Generates an array of 6 numbers with two pairs of 3 coordinates describing position on a sphere.
        The first two represent the x and y positions on a sphere with random radius between 1 and 5,
        and the last coordinate represents the z coordinate scaled by the radius.

        Returns:
            float 6 array:
                Random lagrange multipliers for an initial guess
        """

        Lr = np.random.uniform(1,5)
        Lv = np.random.uniform(1,5)

        alpha_r = np.random.uniform(0,2*np.pi)
        alpha_v = np.random.uniform(0,2*np.pi)

        delta_r = np.random.uniform(-np.pi/2,np.pi/2)
        delta_v = np.random.uniform(-np.pi/2,np.pi/2)

        #TODO: seems odd that 3rd and 6th coordinate aren't part of a spherical coordinate

        L = np.array([Lr*np.cos(alpha_r)*np.cos(delta_r),
                 Lr*np.sin(alpha_r)*np.cos(delta_r),
                 np.sin(delta_r),
                 Lv*np.cos(alpha_v)*np.cos(delta_v),
                 Lv*np.sin(alpha_v)*np.cos(delta_v),
                 np.sin(delta_v)
                 ])

        return L


    def newStar_angularSep(self,TL,iStars,jStars,currentTime,dt):
        """Finds angular separation from old star to given list of stars
        
        This method returns the angular separation from the last observed 
        star to all others on the given list at the currentTime. This is distinct
        from the prototype version of the method in its signing of the
        angular separation based on halo velocity direction
        
        Args:
            TL (TargetList module):
                TargetList class object
            old_sInd (integer):
                Integer index of the last star of interest
            sInds (integer ndarray):
                Integer indices of the stars of interest
            currentTime (astropy Time array):
                Current absolute mission time in MJD
            dt (float):
                Timestep in units of days for determining halo velocity frame

        Returns:
            float:
                Angular separation between two target stars 
        """

        nStars = len(iStars)
        
        # time in canonical units
        absTimes = currentTime + np.array([0,dt])*u.d
        t = self.convertTime_to_canonical(np.mod(absTimes.value,self.equinox.value)*u.d) * u.rad
                
        # halo positions and velocities 
        haloPos = self.haloPosition(absTimes) + np.array([1,0,0])*self.L2_dist.to('au')
        haloVel = self.haloVelocity(absTimes)
                
        # halo positions and velocities in canonical units
        r_T0_R    = np.array([self.convertPos_to_canonical(haloPos[:,n]) for n in range(3)])
        Rdr_T0_R  = np.array([self.convertVel_to_canonical(haloVel[:,n]) for n in range(3)])
        
        # V-frame
        v1_R = Rdr_T0_R / np.linalg.norm(Rdr_T0_R,axis=0)
        
        # Position of stars
        coords = TL.coords
        beta = coords.lat
        lamb = coords.lon
        dist = self.convertPos_to_canonical( coords.distance[0] )
        
        # unit vector to star locations
        u_i0_R = np.array([   [np.cos(beta[iStars]) * np.cos(lamb[iStars]-t[0]) ], \
                              [np.cos(beta[iStars]) * np.sin(lamb[iStars]-t[0]) ], \
                              [np.sin(beta[iStars])]  ]).reshape(3,nStars)
        
        u_j0_R = np.array([   [np.cos(beta[jStars]) * np.cos(lamb[jStars]-t[-1]) ], \
                              [np.cos(beta[jStars]) * np.sin(lamb[jStars]-t[-1]) ], \
                              [np.sin(beta[jStars])]  ]).reshape(3,nStars)
        
        # star locations to relative to telescope
        r_iT_R = dist * u_i0_R  - r_T0_R[:,0].reshape(3,1)
        r_jT_R = dist * u_j0_R  - r_T0_R[:,-1].reshape(3,1)
        
        # unit vector from Telescope to star i
        u_iT = r_iT_R/np.linalg.norm(r_iT_R,axis=0)
    
        # unit vector from Telescope to star j
        u_jT = r_jT_R/np.linalg.norm(r_jT_R,axis=0)
        
        # direction of travel from star i to j
        u_ji =  u_jT - u_iT
        
        # sign of the angular separation, + in direction of halo velocity (NOTE: only using halo velocity at t0, not t0+dt)
        dot_sgn = np.matmul( v1_R[:,0].T, u_ji )
        sgn = np.array([np.sign(x) if x != 0 else 1 for x in dot_sgn])
        
        # calculating angular separation and assigning signs
        dot_product = np.array([ np.dot(a.T,b) for a,b in zip(u_iT.T,u_jT.T)])
        dot_product = np.clip(dot_product, -1, 1)
        psi = sgn*(np.arccos( dot_product )*u.rad).to('deg').value
        
        return psi


    def star_angularSepDesiredDist(self,psiPop,nSamples=1e6,angBinSize=3):
        """ Rejection sample from psiPop to achieve nSamples fitting logistic distribution

        Args:
            psiPop (array float):
                List of floats between -180 and 180
            nSamples (int):
                Number of samples to return
            angBinSize (float):
                Size of a bin for the histogram used in rejection sampling

        Returns:
            tuple:
                array float:
                    Angles (nSamples) fitting the logistic distribution
                array int:
                    An array of indices of original psiPop that were accepted

        """
        
        # distribution of the full psi population
        N = len(psiPop)
        psiBins = np.arange(-180,181,3)
        envelopeDist, EnvDbins = np.histogram(psiPop,psiBins,density=True)
    
        #Desired Distribution pdf
        s = 32
        desiredDist = 1/(4*s) * 1/(  np.cosh(psiBins[:-1]/(2*s))**2  )
        
        #Scaling factor
        M = np.max( desiredDist / envelopeDist )
        
        # how many samples from the full data set are we taking?
        samplingSize = int(nSamples) if (nSamples > 0) & (nSamples <= N) else N
        
        # randomly selecting from the population (no replacing, so unique samples)
        sampling_inds = np.random.choice( N , samplingSize , replace=False  )
        
        # an empty array to which we'll add accepted indices
        final_inds = np.array([],dtype=int)
            
        # running rejection sampling
        for n in sampling_inds:
            
            # figuring out which bin to put this particular ang sep in
            corresponding_Bin = np.where(psiBins <= psiPop[n])[0][-1]
            
            # calculating pdf ratio
            ratio = desiredDist[corresponding_Bin] / (M*envelopeDist[corresponding_Bin])
            
            # random number from 0 to 1
            rand    = np.random.uniform(0,1)
            
            # sample has been accepted!
            if rand < ratio:
                final_inds = np.hstack([ final_inds, n ])
        
        psiFiltered = psiPop[final_inds]
        
        return psiFiltered, final_inds
    
    
    def selectPairsOfStars(self,TL,nPairs,currentTime,dt,nSamples):
        """ Rejection sampling of star pairings using desired distribution and sampling from that final distribution
        
        Args:
            TL (TargetList module):
                TargetList class object
            nPairs (int):
                Number of pairs to produce
            currentTime (astropy Time array):
                Current absolute mission time in MJD
            dt (float):
                Timestep in units of days for determining halo velocity frame
            nSamples (int):
                Number of samples to be used when generating random stars for pairing
        Returns:
            tuple:
                int list:
                    The list of indices corresponding to starting stars in TargetList
                int list:
                    The list of indices corresponding to the ending stars in TargetList
                float array:
                    Angular distance between pairs from iFinal and jFinal
        
        """
        
        psiBins = np.arange(-180,181,3)
        nPairs -= len(psiBins)-1

        # randomly selected pairs of stars
        iLog = np.random.randint(0,TL.nStars,nSamples)
        jLog = np.random.randint(0,TL.nStars,nSamples)
        
        # angular separation between each of these pairs
        # star i -> at currentTime
        # star j -> at currentTime + dt
        psi = self.newStar_angularSep(TL,iLog,jLog,currentTime,dt)
        
        # filtering pairings list, creating desired distribution using rejection sampling
        filtered_psi, filtered_inds = self.star_angularSepDesiredDist(psi,nSamples)
        
        # sampling nPairs of stars from the final distribution
        samples = np.random.choice( len(filtered_inds) , nPairs , replace=False  )
        sampling_inds = filtered_inds[samples]
        
        # non-sampled star pairings
        leftover_psi  = np.delete(psi,sampling_inds)
        leftover_inds = np.delete( np.arange(0,nSamples) , sampling_inds)
        
        # grab a star from each bin and add it to the list
        psiFinal = np.array([],dtype=int)
        
        for lb,ub in zip(psiBins[:-1],psiBins[1:]):
            find_one = np.where( (leftover_psi >= lb) & (leftover_psi < ub) )[0]
            if find_one.size > 0:
                sampling_inds = np.hstack([ sampling_inds, leftover_inds[find_one[0]] ])
                psiFinal = np.hstack([ psiFinal , leftover_psi[find_one[0]] ])
        
        # results
        iFinal = iLog[sampling_inds]
        jFinal = jLog[sampling_inds]
        
        # add the sampled stars from our desired distributions
        psiFinal = np.hstack([ psiFinal , filtered_psi[samples] ]) 
        
        return iFinal,jFinal,psiFinal
# =============================================================================
# Helper functions
# =============================================================================
        
    def determineThrottle(self,state):
        """ Determines throttle based on instantaneous switching function value.
        
        Typically being used during collocation algorithms. A zero-crossing of
        the switching function is highly unlikely between the large number of 
        nodes.

        Args:
            state (float array):
                State vector and lagrange multipliers
        Returns:
            float array:
                Throttle between 0 and 1, and in the same shape as state
        """
        
        eps = self.epsilon
        n = 1 if state.size == 14 else state.shape[1]
        
        throttle = np.zeros(n)
        S = self.switchingFunction(state)
        S = S.reshape(n)
        
        for i,s in enumerate(S):
            if eps > 0:
                midthrottle = (eps - s)/(2*eps)
                throttle[i] = 0 if s > eps else 1 if s < -eps else midthrottle
            else:
                throttle[i] = 0 if s > eps else 1
        
        return throttle
    
    
    def switchingFunction(self,state):
        """ Evaluates the switching function at specific states.

        Args:
            state (float array):
                State vector and lagrange multipliers
        Returns:
            float:
                Value of the switching function

        """
        
        x,y,z,dx,dy,dz,m,L1,L2,L3,L4,L5,L6,L7 = state
        
        Lv_, lv = self.unitVector( np.array([L4,L5,L6]) )
        
        S = -lv*self.ve/m - L7 + 1
        
        return S

    def switchingFunctionDer(self,state):
        """ Evaluates the time derivative of the switching function.
        
        Switching function derivative evaluated for specific states.

        Args:
            state (float array):
                State vector and lagrange multipliers
        Returns:
            float:
                Value of the switching function time derivative

        """
        ve  = self.ve
        n = 1 if state.size == 14 else state.shape[1]
        x,y,z,dx,dy,dz,m,L1,L2,L3,L4,L5,L6,L7 = state
        
        Lr = np.array([L1,L2,L3]).reshape(3,n)
        Lv = np.array([L4,L5,L6]).reshape(3,n)
        Lv_, lv = self.unitVector( Lv )
        
        Pv_arr = np.array([[0,2,0],[-2,0,0],[0,0,0]])
        Pv = np.dstack([Pv_arr]*n)
        
        PLdot = np.vstack( [np.dot(a.T,b) for a,b in zip(Pv.T,Lv.T)] ).T
        
        dS = -(ve / m) * np.vstack( [np.dot(a.T,b) for a,b in zip( (-Lr - PLdot).T ,Lv_.T)] ).T 

        return dS
    
    def selectEventFunctions(self,s0):
        """ Selects the proper event function for integration.
        
        This method calculates the switching function and its derivative at a 
        single specific state. It then determines which thrust case it will be 
        in: full, medium, or no thrust. If the value of the switching function
        is within a certain tolerance of the boundaries, it uses the derivative 
        to determine the direction it is heading in. Then the proper event 
        functions are created for the integrator to determine the next crossing
        (i.e. the next case change). 
        
        Args:
            s0 (float array):
                Iniitial state vector and lagrange multipliers
        Returns:
            tuple:
                function list:
                    A list of functions which are the switching function adjusted by 1e-10.
                    These take in an unused first parameter, and then the state.
                    Depending on the case (2), multiple functions are returned.
                    These functions will return 0 when the switching function crosses
                    1e-10 or -1e-10.
                int:
                    An integer 0,1,2 representing whether the switching
                    function of the initial state is within a 1e-10 tolerance of the
                    of zero (2), less than -1e-10 (1), or greater than 1e-10 (0). 
        """
        eps = self.epsilon
        
        S  =  self.switchingFunction(s0)
        dS =  self.switchingFunctionDer(s0)[0]

        # finding which case we are in:
        #   - case 2 if       -eps < S < eps       (medium thrust)
        #   - case 1 if   S < -eps                 (full thrust)
        #   - case 0 if                  eps < S   (no thrust)
        
        case = 0 if S > eps else 1 if S < -eps else 2

        # checking to see if S is within a certain tolerance from epsilon
        withinTol = np.abs( (np.abs(S) - eps) ) < 1e-10
        # determine if there is a case error if within tolerance
        if withinTol:
            # not the minimum fuel case
            if eps != 0:
                # at the upper bound, case determined by derivative
                if S > 0:
                    case = 2 if dS < 0 else 0
                # at the lower bound, case determined by derivative
                else:
                    case = 2 if dS > 0 else 1
            # minimum fuel case, only two cases
            else:
                case = 0 if dS > 0 else 1
            
            
        eventFunctions = []
        CrossingUpperBound = lambda t,s : self.switchingFunction(s) - eps
        CrossingLowerBound = lambda t,s : self.switchingFunction(s) + eps
        
        CrossingUpperBound.terminal = True
        CrossingLowerBound.terminal = True
        
        if case == 0:
            # crossing upper epsilon from above
            CrossingUpperBound.direction = -1 
            # appending event function
            eventFunctions.append(CrossingUpperBound)
        elif case == 1:
            # crossing lower epsilon from below
            CrossingLowerBound.direction = 1 
            # appending event function
            eventFunctions.append(CrossingLowerBound)
        else:
            # can either cross lower epsilon from above or upper from below
            CrossingLowerBound.direction = -1
            CrossingUpperBound.direction = 1
            # appending event function
            eventFunctions.append(CrossingUpperBound)
            eventFunctions.append(CrossingLowerBound)
        
        return eventFunctions,case
        

# =============================================================================
# Equations of Motion and Boundary Conditions
# =============================================================================
    def boundary_conditions_thruster(self,sA,sB,constrained=False):
        """ Creates boundary conditions for solving a boundary value problem

        Function used in scipy solve_bvp function call. Returns residuals

        Args:
            sA (6 or 7 array):
                To be compared with the initial state for boundary condition.
            sB (6 or 7 array):
                To be compared with the final state for boundary condition.
            constrained (boolean):
                Whether there are 6 (false) or 7 (true) boundary conditions. 
        Returns:
            array:
                Returns the residuals of the boundary conditions/constraint equation.
        """
    
        BCo1 = sA[0] - self.sA[0]
        BCo2 = sA[1] - self.sA[1]
        BCo3 = sA[2] - self.sA[2]
        BCo4 = sA[3] - self.sA[3]
        BCo5 = sA[4] - self.sA[4]
        BCo6 = sA[5] - self.sA[5]
        
        BCf1 = sB[0] - self.sB[0]
        BCf2 = sB[1] - self.sB[1]
        BCf3 = sB[2] - self.sB[2]
        BCf4 = sB[3] - self.sB[3]
        BCf5 = sB[4] - self.sB[4]
        BCf6 = sB[5] - self.sB[5]
        
        if constrained:
            BCo7 = sA[6] - self.sA[6]
            BCf7 = sB[-1]
            BC = np.array([BCo1,BCo2,BCo3,BCo4,BCo5,BCo6,BCo7,BCf1,BCf2,BCf3,BCf4,BCf5,BCf6,BCf7])
        else:
            BC = np.array([BCo1,BCo2,BCo3,BCo4,BCo5,BCo6,BCf1,BCf2,BCf3,BCf4,BCf5,BCf6])

        return BC   
     
        
    def EoM_Adjoint(self,t,state,constrained=False,amax=False,integrate=False):
        """ Equations of Motion with costate vectors

        Equations of motion for CR3BP with costates for control.

        Args:
            t (float):
                Currently unused
            state (array):
                The state and lagrange multipliers.
            constrained (boolean):
                Currently unused
            amax (float or boolean):
                Maximum acceleration attainable or False otherwise
            integrate (boolean):
                If true, the array is flattened for integration.

        Returns:
            array:
                The time derivatives of the states and costates.
        """
        
        mu = self.mu
        ve = self.ve
        
        n = 1 if state.size == 14 else state.shape[1]

        if amax:
            x,y,z,dx,dy,dz,m,L1,L2,L3,L4,L5,L6,L7 = state
        else:
            x,y,z,dx,dy,dz,L1,L2,L3,L4,L5,L6 = state
        
        if integrate:
            state = state.T
        
        # vector distances from primaries
        r1   = np.array([x-(-mu),y,z])
        r2   = np.array([x-(1-mu),y,z])
        # norms of distances from primaries
        R1   = np.linalg.norm(r1,axis=0)
        R2   = np.linalg.norm(r2,axis=0)
        
        # Position-dependent acceleration terms
        qx = np.array([ x  - (1-mu)*(x+mu)/R1**3  - mu*(x+mu-1)/R2**3]).reshape(1,n)
        qy = np.array([ y  - (1-mu)*y/R1**3       - mu*y/R2**3]).reshape(1,n)
        qz = np.array([    - (1-mu)*z/R1**3       - mu*z/R2**3]).reshape(1,n)
        q = np.vstack([qx,qy,qz])  #shape of 3xn
        
        # Position partial derivatives
        Q11 = 1 - (1-mu)/R1**3 + 3*(1-mu)*(x+mu)**2/R1**5  - mu/R2**3   + 3*mu*(x+mu-1)**2/R2**5
        Q22 = 1 - (1-mu)/R1**3 + 3*(1-mu)*     y**2/R1**5  - mu/R2**3   + 3*mu*       y**2/R2**5
        Q33 =   - (1-mu)/R1**3 + 3*(1-mu)*     z**2/R1**5  - mu/R2**3   + 3*mu*       z**2/R2**5
        Q12 =                    3*(1-mu)*(x+mu)* y/R1**5               + 3*mu*(x+mu-1)* y/R2**5
        Q13 =                    3*(1-mu)*(x+mu)* z/R1**5               + 3*mu*(x+mu-1)* z/R2**5
        Q23 =                    3*(1-mu)*      y*z/R1**5               + 3*mu*        y*z/R2**5
        Qr   = np.array([[Q11, Q12 , Q13], [Q12, Q22 , Q23], [Q13, Q23 , Q33]]).reshape(3,3,n) # shape of 3x3xn

        # Velocity-dependent acceleration terms
        px =  2*dy
        py = -2*dx
        pz = np.zeros([1,n])
        p  = np.vstack([px,py,pz]) #shape of 3xn
        
        # Velocity partial derivatives
        Pv_arr = np.array([[0,2,0],[-2,0,0],[0,0,0]])
        Pv = np.dstack([Pv_arr]*n)
        
        # Costate vectors
        Lr = np.vstack([L1,L2,L3])
        Lv = np.vstack([L4,L5,L6])
        Lr_, lr = self.unitVector(Lr)
        Lv_, lv = self.unitVector(Lv)
        
        # ================================================
        # Equations of Motion
        # ================================================
        dX  = np.vstack([ dx,dy,dz ])
        dV  = q + p
        dLx = -np.vstack( [np.dot(a.T,b) for a,b in zip(Qr.T,Lv.T)] ).T
        dLv = -Lr - np.vstack( [np.dot(a.T,b) for a,b in zip(Pv.T,Lv.T)] ).T
        
        if amax:
            # throttle factor
            throttle = self.determineThrottle(state)
            
            dV -= Lv_ * amax * throttle / m
            dm  = -throttle * amax / ve
            dLm = -lv * throttle * amax / m**2
            # putting them all together, a 14xn array
            f   = np.vstack([ dX, dV, dm, dLx, dLv, dLm ])
        else:
            dV -= Lv
            # putting them all together, a 12xn array
            f   = np.vstack([ dX, dV, dLx, dLv ])
        
        
        if integrate:
            f = f.flatten()

        return f
    
    
    def starshadeBoundaryVelocity(self,TL,sInd,currentTime,SRP=True):
        """ Calculates starshade and telescope velocities in R- or I-frames during stationkeeping
        
        Calculates the rotating system position and velocity of the
        starshade at insertion into the optimal initial state for
        observation of the given star.

        Args:
            TL (TargetList module):
                Target List
            sInd (int):
                Index of star for observation in the target list
            currentTime (astropy Time array):
                Current absolute mission time in MJD
            SRP (boolean):
                Whether or not solar radiation pressure should be used 
                in calculating the insertion state for optimal 
                starshade stationkeeping observation position.   
        Returns:
            tuple:
                array:
                    Position in the rotating frame
                array:
                    Velocity in the rotating frame
        """
        
        # absolute times (Note: equinox is start time of Halo AND when inertial frame and rotating frame match)
        modTimes = np.mod(currentTime.value,self.equinox.value)*u.d  #mission times relative to equinox )
        t = self.convertTime_to_canonical(modTimes) * u.rad       #modTimes in canonical units 
        
        r_S0_I, Iv_S0_I, qq, qq, qq, qq = self.starshadeKinematics(TL,sInd,currentTime)
        
        # star a at t=ta
        r_AS_C, Iv_AS_C = self.starshadeInjectionVelocity(TL,sInd,currentTime,SRP=SRP)
        s_AS_C = np.hstack([r_AS_C , Iv_AS_C])
        r_AS_I, Iv_AS_I = self.rotateComponents2NewFrame(TL,sInd,currentTime,s_AS_C,np.array([0]),SRP=SRP,final_frame='I')
        
        rQi__tA = self.rot( t[0].value , 3 )
        r_A0_I  =  r_S0_I +  r_AS_I
        Iv_A0_I = Iv_S0_I + Iv_AS_I
        
        r_A0_R   = np.dot(  rQi__tA , r_A0_I  )
        Rv_A0_I  = deepcopy(Iv_A0_I)
        Rv_A0_I[0] += r_A0_I[1]
        Rv_A0_I[1] -= r_A0_I[0]
        Rv_A0_R  = np.dot(  rQi__tA , Rv_A0_I  )
        
        return r_A0_R, Rv_A0_R
    
# =============================================================================
# Initial conditions
# =============================================================================

    def findInitialTmax(self,TL,nA,nB,tA,dt,m0=1,s_init=np.array([])):
        """ Finding initial guess for starting Thrust
        
        Args:
            TL (TargetList module):
                Target list
            nA (int):
                The index for the starting star in the target list
            nB (int):
                The index for the final star in the target list
            tA (astropy Time array):
                Initial absolute mission time in MJD
            dt (astropy Time array);
                A time step between the two voundary conditions
            m0 (float):
                Initial mass
            s_init (array):
                An initial guess for the state

        Returns:
            tuple:
                astropy Newtons:
                    The maximum initial thrust.
                array:
                    States from the solved bvp.
                array:
                    Times at corrsponding to each state.
        """
        
        tB = tA + dt
        #angle,uA,uB,r_tscp = self.lookVectors(TL,nA,nB,tA,tB)

        #position vector of occulter in heliocentric frame
        #self_rA = uA*self.occulterSep.to('au').value + r_tscp[ 0]
        #self_rB = uB*self.occulterSep.to('au').value + r_tscp[-1]
        
        r_A0_R, Rv_A0_R = self.starshadeBoundaryVelocity(TL,nA,tA)
        r_B0_R, Rv_B0_R = self.starshadeBoundaryVelocity(TL,nB,tB)
           
        Rs_A0_R = np.hstack([r_A0_R[:,0],Rv_A0_R[:,0]])
        Rs_B0_R = np.hstack([r_B0_R[:,0],Rv_B0_R[:,0]])
                
        Rsf_A0_R = np.hstack([Rs_A0_R, self.lagrangeMult()])
        Rsf_B0_R = np.hstack([Rs_B0_R, self.lagrangeMult()])
                
        a = ((np.mod(tA.value,self.equinox.value)*u.d)).to('yr') / u.yr * (2*np.pi)
        b = ((np.mod(tB.value,self.equinox.value)*u.d)).to('yr') / u.yr * (2*np.pi)
        
        #running collocation
        tGuess = np.hstack([a,b]).value

        if s_init.size:
            sGuess = np.vstack([s_init[:,0] , s_init[:,-1]])
        else:
            sGuess = np.vstack([Rsf_A0_R , Rsf_B0_R])
            
        s,t_s,status = self.send_it_thruster(sGuess.T,tGuess,m0=m0,verbose=False)
        
        lv = s[9:,:]

        aNorms0 = np.linalg.norm(lv,axis=0)
        aMax0   = self.convertAcc_to_dim( np.max(aNorms0) ).to('m/s^2') 
        Tmax0   = (aMax0 * self.mass ).to('N')
        
        return Tmax0, s, t_s
    
    
    def findTmaxGrid(self,TL,tA,dtRange):
        """ Create grid of Tmax values using unconstrained thruster
        
        This method is used purely for creating figures. 
        
        Args:
            TL (TargetList module):
                Target list
            tA (astropy Time array):
                Initial absolute mission time in MJD
            dtRange (astropy Time array);
                An array of delta times to consider

        Returns:
            astropy Newtons array:
                Max thrust in Newtons (dtRange dimensions by TL.nStars) 
        """
        
        midInt = int( np.floor( (TL.nStars-1)/2 ) )

        sInds = np.arange(0,TL.nStars)
        ang   =  self.star_angularSep(TL, midInt, sInds, tA) 
        sInd_sorted = np.argsort(ang)
        angles  = ang[sInd_sorted].to('deg').value
        
        
        TmaxMap = np.zeros([len(dtRange) , len(angles)])*u.N
        
        for i,t in enumerate(dtRange):
            for j,n in enumerate(sInd_sorted):
                print(i,j)
                Tmax, s, t_s = self.findInitialTmax(TL,midInt,n,tA,t)
                TmaxMap[i,j] = Tmax
        
        return TmaxMap

# =============================================================================
# BVP solvers
# =============================================================================
  
    def send_it_thruster(self,sGuess,tGuess,aMax=False,constrained=False, m0 = 1, maxNodes=1e5,verbose=False):
        """ Solving generic bvp from t0 to tF using states and costates

        Uses the Scipy solve_bvp method.

        Args:
            sGuess (array):
                Initial state and costate guess
            tGuess (astropy Time array):
                Times corresponding to the guess of the state at each time
            aMax (astropy Newton or boolean):
                Maximum attainable acceleration
            constrained (boolean):
                Flag for whether this is constrained or unconstrained problem
                This determines dimensions of the state inputs.
            m0 (float):
                Initial mass
            maxNodes (int):
                Maximum number of nodes to use in the BVP problem solution
            verbose (boolean):
                Flag passed to solve_bvp
        Returns:
            tuple:
                array:
                    State and costate at the various mesh times
                array:
                    Corresponding times to the sampled states.
                boolean:
                    Status returned by the bvp_solve method
        """
        
        sG = sGuess
        # unconstrained problem begins with 12 states, rather than 14. checking for that
        if len(sGuess) == 12:
            x,y,z,dx,dy,dz,L1,L2,L3,L4,L5,L6 = sGuess
            # if unconstrained is initial guess for constrained problem, make 14 state array
            if aMax:
                mRange  = np.linspace(m0, 0.8, len(x))
                lmRange = np.linspace(1, 0, len(x))
                sG = np.vstack([x,y,z,dx,dy,dz,mRange,L1,L2,L3,L4,L5,L6,lmRange])

        # only saves initial and final desired states if first solving unconstrained problem
        if not constrained:
            self.sA = np.hstack([ sG[0:6,0 ], m0   , sG[6:,0]  , 1 ])
            self.sB = np.hstack([ sG[0:6,-1], 0.8 , sG[6:,-1] , 0 ])
            self.sG = sG
        
        # creating equations of motion and boundary conditions functions
        EoM = lambda t,s: self.EoM_Adjoint(t,s,constrained,aMax)
        BC  = lambda t,s: self.boundary_conditions_thruster(t,s,constrained)
        # solving BVP
        sol = solve_bvp(EoM,BC,tGuess,sG,tol=1e-8,max_nodes=int(maxNodes),verbose=0)
        
        if verbose:
            self.vprint(sol.message)
        # saving results
        s = sol.y
        t_s = sol.x
            
        return s,t_s,sol.status
    
    
    def collocate_Trajectory(self,TL,nA,nB,tA,dt):
        """ Solves minimum energy and minimum fuel cases for continuous thrust

        Args:
            TL (TargetList module):
                Target list
            nA (int):
                The index for the starting star in the target list
            nB (int):
                The index for the final star in the target list
            tA (astropy Time array):
                Initial absolute mission time in MJD
            dt (astropy Time array);
                A time step between the two voundary conditions
        
        Returns:
            tuple:
                array:
                    Trajectory
                array:
                    Times corresponding to trajectory
                float:
                    last epsilon that fully converged (2 if minimum energy didn't work)
                    Parameterizes minimum energy to minimum fuel solution.
                astropy Newton array:
                    Range of thrusts (Newtons) considered.
        """
        
        # initializing arrays
        stateLog = []
        timeLog  = []
        
        # solving using unconstrained thruster as initial guess
        Tmax, sTmax, tTmax = self.findInitialTmax(TL,nA,nB,tA,dt)
        aMax = self.convertAcc_to_canonical( (Tmax / self.mass).to('m/s^2') )
        # saving results
        stateLog.append(sTmax)
        timeLog.append(tTmax)
        
        # all thrusts were successful
        e_best   = 2
        s_best   = deepcopy(stateLog)
        t_best   = deepcopy(timeLog)
        
        # thrust values
        desiredT = self.Tmax.to('N').value
        currentT = Tmax.value
        # range of thrusts to try
        TmaxRange    = np.linspace(currentT,desiredT,30)
        
        # range of epsilon values to try (e=1 is minimum energy, e=0 is minimum fuel)
        epsilonRange = np.round( np.arange(1,-0.1,-0.1) , decimals = 1)
        
        # Huge loop over all thrust and epsilon values:
        #   we start at the minimum energy case, e=1, using the thrust value from the unconstrained solution
        #     In/De-crement the thrust until we reach the desired thrust level
        #   then we decrease e and repeat process until we get to e=0 (minimum fuel)
        #   saves the last successful result in case collocation fails
        
        # loop over epsilon starting with e=1
        for j,e in enumerate(epsilonRange):
            print("Collocate Epsilon = ",e)
            # initialize epsilon
            self.epsilon = e
            
            # loop over thrust values from current to desired thrusts
            for i,thrust in enumerate(TmaxRange):
                #print("Thrust #",i," / ",len(TmaxRange))
                # convert thrust to canonical acceleration
                aMax = self.convertAcc_to_canonical( (thrust*u.N / self.mass).to('m/s^2') )
                # retrieve state and time initial guesses
                sGuess = stateLog[i]
                tGuess = timeLog[i]
                # perform collocation
                s,t_s,status = self.send_it_thruster(sGuess,tGuess,aMax,constrained=True,maxNodes=1e5,verbose=False)
                
                # collocation failed, exits out of everything
                if status != 0:
                    self.epsilon = e_best
                    if e_best == 2:
                        # if only the unconstrained problem worked, still returns a 14 length array
                        s_out = []
                        length = s_best[0].shape[1]
                        m  = np.linspace(1,0.9,length)
                        lm = np.linspace(0.3,0,length)
                        s_out.append( np.vstack([s_best[0][:6] , m, s_best[0][6:], lm]) )
                        s_best = deepcopy(s_out)
                    return s_best, t_best, e_best, TmaxRange
                
                # collocation was successful!
                if j == 0:
                    # creates log of state and time results for next thrust iteration (at the beginning of the loop)
                    stateLog.append(s)
                    timeLog.append(t_s)
                else:
                    # updates log of state and time results for next thrust iteration
                    stateLog[i] = s
                    timeLog[i]  = t_s
            
            # all thrusts were successful, save results
            e_best   = self.epsilon
            s_best   = deepcopy(stateLog)
            t_best   = deepcopy(timeLog)
        
        return s_best, t_best, e_best, TmaxRange


    def collocate_Trajectory_minEnergy(self,TL,nA,nB,tA,dt,m0=1):
        """ Solves minimum energy and minimum fuel cases for continuous thrust
        
        Args:
            TL (TargetList module):
                Target list
            nA (int):
                The index for the starting star in the target list
            nB (int):
                The index for the final star in the target list
            tA (astropy Time array):
                Initial absolute mission time in MJD
            dt (astropy Time array);
                A time step between the two voundary conditions
            m0 (float):
                Initial mass

        Returns:
            tuple:
                array:
                    Trajectory
                array:
                    Times corresponding to trajectory
                float:
                    last epsilon that fully converged (2 if minimum energy didn't work)
                    Parameterizes minimum energy to minimum fuel solution.
                astropy Newton array:
                    Range of thrusts (Newtons) considered.
        """
        
        # initializing arrays
        stateLog = []
        timeLog  = []
        
        # solving using unconstrained thruster as initial guess
        Tmax, sTmax, tTmax = self.findInitialTmax(TL,nA,nB,tA,dt,m0)
        aMax = self.convertAcc_to_canonical( (Tmax / self.mass).to('m/s^2') )
        # saving results
        stateLog.append(sTmax)
        timeLog.append(tTmax)
        
        # all thrusts were successful
        e_best   = 2
        s_best   = deepcopy(sTmax)
        t_best   = deepcopy(tTmax)
        
        # thrust values
        desiredT = self.Tmax.to('N').value
        currentT = Tmax.value
        # range of thrusts to try
        TmaxRange    = np.linspace(currentT,desiredT,30)
        
        # Huge loop over all thrust and epsilon values:
        #   we start at the minimum energy case, e=1, using the thrust value from the unconstrained solution
        #     In/De-crement the thrust until we reach the desired thrust level
        #   then we decrease e and repeat process until we get to e=0 (minimum fuel)
        #   saves the last successful result in case collocation fails
        
        # loop over epsilon starting with e=1
        self.epsilon = 1
        
        # loop over thrust values from current to desired thrusts
        for i,thrust in enumerate(TmaxRange):
            #print("Thrust #",i," / ",len(TmaxRange))
            # convert thrust to canonical acceleration
            aMax = self.convertAcc_to_canonical( (thrust*u.N / self.mass).to('m/s^2') )
            # retrieve state and time initial guesses
            sGuess = stateLog[i]
            tGuess = timeLog[i]
            # perform collocation
            s,t_s,status = self.send_it_thruster(sGuess,tGuess,aMax,constrained=True,m0=m0,maxNodes=1e5,verbose=False)
            
            # collocation failed, exits out of everything
            if status != 0:
                self.epsilon = e_best
                return s_best, t_best, e_best, TmaxRange
            
            # creates log of state and time results for next thrust iteration (at the beginning of the loop)
            stateLog.append(s)
            timeLog.append(t_s)

        # all thrusts were successful, save results
        e_best   = self.epsilon
        s_best   = deepcopy(s)
        t_best   = deepcopy(t_s)
    
        return s_best, t_best, e_best, TmaxRange
    
# =============================================================================
# Shooting Algorithms
# =============================================================================
        
    def integrate_thruster(self,sGuess,tGuess,Tmax,verbose=False):
        """ Integrates thruster trajectory with thrust case switches
        
        This methods integrates an initial guess for the spacecraft state 
        forwards in time. It uses event functions to find the next zero of the
        switching function which means a new thrust case is needed (full, 
        medium or no thrust).

        Args:
            sGuess (array):
                Initial state and costate guess
            tGuess (astropy Time array):
                Times corresponding to the guess of the state at each time
            Tmax (astropy force):
                Maximum thrust attainable

        Returns:
            tuple:
                array:
                    Trajectory states
                astropy Time array:
                    Times corresponding to states
        """
        s0 = sGuess[:,0]
        t0 = tGuess[0]
        tF = tGuess[-1]
        
        # initializing starting time
        tC = deepcopy(t0)
        
        # converting Tmax to a canonical acceleration
        Tmax = Tmax.to('N').value
        aMax = self.convertAcc_to_canonical( (Tmax*u.N / self.mass).to('m/s^2') )
        
        # establishing equations of motion
        EoM  = lambda t,s: self.EoM_Adjoint(t,s,constrained=True,amax=aMax,integrate=True)
        
        # starting integration
        count = 0
        while tC < tF:
            # selecting the switch functions with correct boundaries
            switchFunctions,case = self.selectEventFunctions(s0)
            if verbose:
                print("[%.3f / %.3f] with case %d" % (tC,tF,case) )
            # running integration with event functions
            res = solve_ivp(EoM, [tC , tF], s0, events=switchFunctions)
            
            # saving final integration time, if greater than tF, we completed the trajectory
            tC = deepcopy(res.t[-1])
            s0 = deepcopy(res.y[:,-1])
            
            # saving results in a new array
            if count == 0:
                sLog = deepcopy(res.y)
                tLog = deepcopy(res.t)
            #adding to the results log if there was a previous thrust case switch
            else:
                sLog = np.hstack([sLog,res.y])
                tLog = np.hstack([tLog,res.t])
            count += 1
            
        return sLog,tLog


    def conFun_singleShoot(self,w,t0,tF,Tmax,returnLog=False):
        """ Objective Function for single shooting thruster
        
        Args:
            w (array):
                Costate vector
            t0 (astropy Time object):
                Initial time
            tf (astropy Time object):
                Final time
            Tmax (astropy force):
                Maximum thrust attainable
            returnLog (boolean):
                Return the states and times of the solution
        Returns:
            tuple:
                float:
                    Norm of difference between current state and boundary value
                array:
                    Trajectory states
                astropy Time array:
                    Times corresponding to states
        """
        
        sInit  = np.hstack([self.sA[:7] , w]).reshape(14,1)
        tGuess = np.array([t0,tF])
        
        sLog,tLog = self.integrate_thruster(sInit,tGuess,Tmax)
        
        f = self.boundary_conditions_thruster( sLog[:,0], sLog[:,-1], constrained=True)
        fnorm = np.linalg.norm(f)
        
        if returnLog:
            return fnorm,sLog,tLog
        else:
            return fnorm


    def minimize_TerminalState(self,s_best,t_best,Tmax,method):
        """ Minimizes boundary conditions for thruster

        Args:
            sGuess (array):
                Initial state and costate guess
            tGuess (astropy Time array):
                Times corresponding to the guess of the state at each time
            Tmax (astropy force):
                Maximum thrust attainable
            method (string):
                Optimization method for Scipy minimize call
        Returns:
            tuple:
                float:
                    Norm of difference between current state and boundary value
                array:
                    Trajectory states
                astropy Time array:
                    Times corresponding to states
        """
        
        w0 = s_best[7:,0]
        t0 = t_best[0]
        tF = t_best[-1]
        
        res = optimize.minimize(self.conFun_singleShoot, w0, method=method, \
                                tol=1e-12, args=(t0,tF,Tmax,) )
#        minimizer_kwargs = {"method":method,"args":(t0,tF,Tmax,)}
#        res = optimize.basinhopping(self.conFun_singleShoot,w0,minimizer_kwargs=minimizer_kwargs)
#        
        fnorm, sLog, tLog = self.conFun_singleShoot(res.x,t0,tF,Tmax,returnLog=True)
        
        return fnorm, sLog, tLog


    def singleShoot_Trajectory(self,stateLog,timeLog,e_best,TmaxRange,method='SLSQP'):
        """ Perform single shooting to solve the boundary value problem.

        Args:
            stateLog (array):
                Approximate trajectory typically determined by collocation
            timeLog (astropy Time array):
                Corresponging time values for the trajectory
            e_best (float):
                Epsilon value corresponding to previous approxmiate trajectory
            TmaxRange (astropy Newton array):
                Range of thrusts (Newtons) considered.
            method (string):
                Optimization method for Scipy minimize call

        Returns:
            tuple:
                array:
                    Trajectory states
                astropy Time array:
                    Times corresponding to states
                float:
                    Epsilon value determining how fuel vs energy optimal the trajectory is.
        """


        # initializing arrays
        s_best = deepcopy(stateLog)
        t_best = deepcopy(timeLog)
        
        incLogFlag = True if len(stateLog) != len(TmaxRange) else False
        
        # range of epsilon values to try (e=1 is minimum energy, e=0 is minimum fuel)
        e_best = 1 if e_best == 2 else e_best
        epsilonRange = np.round( np.arange(e_best,-0.1,-0.1) , decimals = 1)
        
        # Huge loop over all thrust and epsilon values, just like collocation method
        
        # loop over epsilon starting with e=1
        for j,e in enumerate(epsilonRange):
            print("SS Epsilon = ",e)
            # initialize epsilon
            self.epsilon = e
            
            # loop over thrust values from current to desired thrusts
            for i,thrust in enumerate(TmaxRange):
                #print("Thrust #",i," / ",len(TmaxRange))
                # retrieve state and time initial guesses
                sGuess = stateLog[i]
                tGuess = timeLog[i]
                # perform single shooting
                fnorm,sLog,tLog = self.minimize_TerminalState(sGuess,tGuess,thrust,method)
                print(fnorm)
                # single shooting failed, exits out of everything
                if fnorm > 1e-7:
                    self.epsilon = e_best
                    return s_best, t_best, e_best
                
                # single shooting was successful!
                if incLogFlag:
                    # appends stateLog if the input was incomplete
                    stateLog.append(sLog)
                    timeLog.append(tLog)
                else:
                    # updates log of state and time results for next thrust iteration
                    stateLog[i] = sLog
                    timeLog[i]  = tLog
            
            # all thrusts were successful, save results
            e_best   = self.epsilon
            s_best   = deepcopy(sLog)
            t_best   = deepcopy(tLog)
        
        return s_best, t_best, e_best
    
# =============================================================================
#  Putting it al together
# =============================================================================
    def calculate_dMmap(self,TL,tA,dtRange,filename):
        """Calculates change in mass for transfers of various times and angular distances

        These are stored in a cache .dmmap file.

        Args:
            TL (TargetList module):
                Target list of stars
            tA (astropy Time array):
                Initial absolute mission time in MJD
            dtRange (astropy Time array):
                Transfer times to try
            filename (string):
                File name to store the cached data.
        """

        sInds       = np.arange(0,TL.nStars)
        ang         = self.star_angularSep(TL, 0, sInds, tA) 
        sInd_sorted = np.argsort(ang)
        angles      = ang[sInd_sorted].to('deg').value

        dtFlipped = np.flipud(dtRange)
        
        self.dMmap = np.zeros([len(dtRange) , len(angles)])
        self.eMap  = np.zeros([len(dtRange) , len(angles)])
        
        tic = time.perf_counter()
        for j,n in enumerate(sInd_sorted):
            for i,t in enumerate(dtFlipped):
                print(i,j)
                s_coll, t_coll, e_coll, TmaxRange = \
                            self.collocate_Trajectory(TL,0,n,tA,t)
                
                if e_coll != 0:
                    s_ssm, t_ssm, e_ssm = self.singleShoot_Trajectory(s_coll, \
                                                t_coll,e_coll,TmaxRange*u.N)

                if e_ssm == 2 and t.value < 30:
                    break
                
                m = s_ssm[-1][6,:] 
                dm = m[-1] - m[0]
                self.dMmap[i,j] = dm
                self.eMap[i,j]  = e_ssm
                toc = time.perf_counter()
                
                dmPath = os.path.join(self.cachedir, filename+'.dmmap')
                A = {'dMmap':self.dMmap,'eMap':self.eMap,'angles':angles,'dtRange':dtRange,'time':toc-tic,\
                     'tA':tA,'m0':1,'ra':TL.coords.ra,'dec':TL.coords.dec,'mass':self.mass}
                with open(dmPath, 'wb') as f:
                    pickle.dump(A, f)
                print('Mass - ',dm*self.mass)
                print('Best Epsilon - ',e_ssm)
    
    
    def calculate_dMmap_collocate(self,TL,tA,dtRange,filename):
        """Calculates change in mass for transfers of various times and angular distances

        These are stored in a cache .dmmap file. Only collocation without the single shooting.

        Args:
            TL (TargetList module):
                Target list of stars
            tA (astropy Time array):
                Initial absolute mission time in MJD
            dtRange (astropy Time array):
                Transfer times to try
            filename (string):
                File name to store the cached data.
        """
        
        sInds       = np.arange(0,TL.nStars)
        ang         = self.star_angularSep(TL, 0, sInds, tA) 
        sInd_sorted = np.argsort(ang)
        angles      = ang[sInd_sorted].to('deg').value

        dtFlipped = np.flipud(dtRange)
        
        self.dMmap = np.zeros([len(dtRange) , len(angles)])
        self.eMap  = np.zeros([len(dtRange) , len(angles)])
        
        tic = time.perf_counter()
        for j,n in enumerate(sInd_sorted):
            for i,t in enumerate(dtFlipped):
                print(i,j)
                s_coll, t_coll, e_coll, TmaxRange = \
                            self.collocate_Trajectory(TL,0,n,tA,t)

                if e_coll == 2 and t.value < 30:
                    break
                
                m = s_coll[-1][6,:]
                dm = m[-1] - m[0]
                self.dMmap[i,j] = dm
                self.eMap[i,j]  = e_coll
                toc = time.perf_counter()
                
                dmPath = os.path.join(self.cachedir, filename+'.dmmap')
                A = {'dMmap':self.dMmap,'eMap':self.eMap,'angles':angles,'dtRange':dtRange,'time':toc-tic,\
                     'tA':tA,'ra':TL.coords.ra,'dec':TL.coords.dec,'mass':self.mass}
                with open(dmPath, 'wb') as f:
                    pickle.dump(A, f)
                print('Mass - ',dm*self.mass)
                print('Best Epsilon - ',e_coll)
    
    
    def calculate_dMmap_collocateEnergy(self,TL,tA,dtRange,filename,m0=1,seed=000000000):
        """Calculates change in mass for transfers of various times and angular distances

        These are stored in a cache .dmmap file. Only minimum energy collocation is used.

        Args:
            TL (TargetList module):
                Target list of stars
            tA (astropy Time array):
                Initial absolute mission time in MJD
            dtRange (astropy Time array):
                Transfer times to try
            filename (string):
                File name to store the cached data.
            m0 (float):
                Initial mass
            seed (int):
                Seed random number for repeatability of experiments
        """
        
        sInds       = np.arange(0,TL.nStars)
        ang         = self.star_angularSep(TL, 0, sInds, tA) 
        sInd_sorted = np.argsort(ang)
        angles      = ang[sInd_sorted].to('deg').value
        
        dtFlipped = np.flipud(dtRange)
        
        self.dMmap = np.zeros([len(dtRange) , len(angles)])
        self.eMap  = 2*np.ones([len(dtRange) , len(angles)])
        
        tic = time.perf_counter()
        for j,n in enumerate(sInd_sorted):
            for i,t in enumerate(dtFlipped):
                print(i,j)
                s_coll, t_coll, e_coll, TmaxRange = \
                            self.collocate_Trajectory_minEnergy(TL,0,n,tA,t,m0)
                
                # if unsuccessful, reached min time -> move on to next star
                if e_coll == 2 and t.value < 30:
                    break

                m = s_coll[6,:] 
                dm = m[-1] - m[0]
                self.dMmap[i,j] = dm
                self.eMap[i,j]  = e_coll
                toc = time.perf_counter()
                
                dmPath = os.path.join(self.cachedir, filename+'.dmmap')
                A = {'dMmap':self.dMmap,'eMap':self.eMap,'angles':angles,'dtRange':dtRange,'time':toc-tic,\
                     'tA':tA,'m0':m0,'ra':TL.coords.ra,'dec':TL.coords.dec,'seed':seed,'mass':self.mass}
                with open(dmPath, 'wb') as f:
                    pickle.dump(A, f)
                print('Mass - ',dm*self.mass)
                print('Best Epsilon - ',e_coll)


    def calculate_dMsols_collocateEnergy(self,TL,tStart,tArange,dtRange,N,filename,m0=1,seed=000000000):
        """Calculates change in mass for transfers of various times and angular distances

        These are stored in a cache .dmmap file. Only minimum energy collocation is used.
        N random combinations of starting times, transfer times, and initial, and final stars are used.

        Args:
            TL (TargetList module):
                Target list of stars
            tStart (astropy Time array):
                Initial reference absolute mission time in MJD
            tArange (astropy Time array):
                Potential times to add to tStart
            dtRange (astropy Time array):
                Transfer times to try
            N (int):
                Number of trials/combinations to perform
            filename (string):
                File name to store the cached data.
            m0 (float):
                Initial mass
            seed (int):
                Seed random number for repeatability of experiments
        """
        
        self.dMmap = np.zeros(N)
        self.eMap  = 2*np.ones(N)
        iLog   = np.zeros(N)
        jLog   = np.zeros(N)
        dtLog  = np.zeros(N)
        tALog  = np.zeros(N)
        angLog = np.zeros(N)*u.deg

        tic = time.perf_counter()
        for n in range(N):
                print("---------\nIteration",n)

                i = np.random.randint(0,TL.nStars)
                j = np.random.randint(0,TL.nStars)
                dt = np.random.randint(0,len(dtRange))
                tA = np.random.randint(0,len(tArange))
                ang = self.star_angularSep(TL,i,j,tStart+tArange[tA]) 

                print("star pair  :",i,j)
                print("ang  :",ang.to('deg').value)
                print("dt   :",dtRange[dt].to('d').value)
                print("tau  :",tArange[tA].to('d').value,"\n")

                pair = np.array([i,j])
                iLog[n]   = i
                jLog[n]   = j
                dtLog[n]  = dt
                tALog[n]  = tA
                angLog[n] = ang

                s_coll, t_coll, e_coll, TmaxRange = \
                            self.collocate_Trajectory_minEnergy(TL,i,j,tStart+tArange[tA],dtRange[dt],m0)
                
                # if unsuccessful, reached min time -> move on to next star
                if e_coll == 2 and dtRange[dt].value < 30:
                    break

                m = s_coll[6,:] 
                dm = m[-1] - m[0]
                self.dMmap[n] = dm
                self.eMap[n]  = e_coll
                toc = time.perf_counter()
                
                dmPath = os.path.join(self.cachedir, filename+'.dmsols')
                A = {'dMmap':self.dMmap,'eMap':self.eMap,'angLog':angLog,'dtLog':dtLog,'time':toc-tic,\
                     'tArange':tArange,'dtRange':dtRange,'N':N,'tStart':tStart,\
                     'tALog':tALog,'m0':m0,'ra':TL.coords.ra,'dec':TL.coords.dec,'seed':seed,'mass':self.mass}
                with open(dmPath, 'wb') as f:
                    pickle.dump(A, f)
                print('Mass - ',dm*self.mass)
                print('Best Epsilon - ',e_coll)         
                
                

    def calculate_dMmap_collocateEnergy_angSepDist(self,TL,tA,dtRange,nPairs,filename,m0=1,seed=000000000):
        """Calculates change in mass for transfers of various times and angular distances

        These are stored in a cache .dmmap file. Only minimum energy collocation is used.
        nPair random combinations of initial and final star are used while all transfer times are used.

        Args:
            TL (TargetList module):
                Target list of stars
            tA (astropy Time array):
                Initial reference absolute mission time in MJD
            dtRange (astropy Time array):
                Transfer times to try
            nPairs (int):
                Number of trials/combinations to perform
            filename (string):
                File name to store the cached data.
            m0 (float):
                Initial mass
            seed (int):
                Seed random number for repeatability of experiments
        """
        
        iLog = np.zeros([len(dtRange) , nPairs])
        jLog = np.zeros([len(dtRange) , nPairs])
        psiLog = np.zeros([len(dtRange) , nPairs])
        
        dtFlipped = np.flipud(dtRange)
        
        self.dMmap = np.zeros([len(dtRange) , nPairs])
        self.eMap  = 2*np.ones([len(dtRange) , nPairs])
        
        tic = time.perf_counter()
        toc = time.perf_counter()
        for t,dt in enumerate(dtFlipped):

            iSelect,jSelect,psiSelect = self.selectPairsOfStars(TL,nPairs,tA,dt.value,int(1e6))
            sort = np.argsort(psiSelect)
            iSelect = iSelect[sort]
            jSelect = jSelect[sort]
            psiSelect = psiSelect[sort]            
            
            for n,(i,j) in enumerate(zip(iSelect,jSelect)):
                elapsedTime = (toc-tic)/3600
                totalTime = elapsedTime * len(dtRange) * nPairs / (1 + t*nPairs + n)
                print("dt :",dt," star #:",n," /",nPairs)
                print("Time Elapsed: ", elapsedTime, " hrs")
                print("Time Left: ", totalTime - elapsedTime, " hrs")
                s_coll, t_coll, e_coll, TmaxRange = \
                    self.collocate_Trajectory_minEnergy(TL,i,j,tA,dt,m0)
                
                m = s_coll[6,:] 
                dm = m[-1] - m[0]
                self.dMmap[t,n] = dm
                self.eMap[t,n]  = e_coll
                
                iLog[t,n] = int(i)
                jLog[t,n] = int(j)
                psiLog[t,n] = psiSelect[n]
                
                
                toc = time.perf_counter()
                
                dmPath = os.path.join(self.cachedir, filename+'.dmmap')
                A = {'dMmap':self.dMmap,'eMap':self.eMap,'angles':psiLog,'dtRange':dtRange,'time':toc-tic,\
                     'tA':tA,'m0':m0,'lon':TL.coords.lon,'lat':TL.coords.lat,'seed':seed,'mass':self.mass, \
                     'iLog':iLog,'jLog':jLog}
                
                with open(dmPath, 'wb') as f:
                    pickle.dump(A, f)
                print('---Mass - ',dm*self.mass)
                print('---Best Epsilon - ',e_coll) 

    def calculate_dMmap_collocateEnergy_LatLon(self,TL,tA,dtRange,nStars,filename,m0=1,seed=000000000):
        """Calculates change in mass for transfers of various times and angular distances

        These are stored in a cache .dmmap file. Only minimum energy collocation is used.
        All pairs between nStars random stars from the targetlist. All transfer times are used.

        Args:
            TL (TargetList module):
                Target list of stars
            tA (astropy Time array):
                Initial reference absolute mission time in MJD
            dtRange (astropy Time array):
                Transfer times to try
            nStars (int):
                Number of trials/combinations to perform
            filename (string):
                File name to store the cached data.
            m0 (float):
                Initial mass
            seed (int):
                Seed random number for repeatability of experiments
        """
        
        coords = TL.coords
        lon = coords.lon
        sInds = np.random.choice( TL.nStars , int(nStars) , replace=False  )
        sInds = sInds[np.argsort(lon[sInds])]
        dtFlipped = np.flipud(dtRange)
        
        self.dMmap = np.zeros([ len(dtRange) , len(sInds), len(sInds)])
        self.eMap  = 2*np.ones([len(dtRange) , len(sInds), len(sInds)])
        
        tic = time.perf_counter()
        toc = time.perf_counter()
        for i,ni in enumerate(sInds):
            for j,nj in enumerate(sInds):
                for n,t in enumerate(dtFlipped):
                    elapsedTime = (toc-tic)/3600
                    totalTime = elapsedTime * len(dtRange) * nStars**2 / (1 + t + j*len(dtRange) + i*nStars*len(dtRange))
                    print("dt :",t," star #:",i,"-->",j)
                    print("Time Elapsed: ", elapsedTime, " hrs")
                    print("Time Left: ", totalTime - elapsedTime, " hrs")
                    print(i,j,t.value)
                    s_coll, t_coll, e_coll, TmaxRange = \
                                self.collocate_Trajectory_minEnergy(TL,ni,nj,tA,t,m0)
                    
                    # if unsuccessful, reached min time -> move on to next star
                    if e_coll == 2 and t.value < 30:
                        break
    
                    m = s_coll[6,:] 
                    dm = m[-1] - m[0]
                    self.dMmap[n,i,j] = dm
                    self.eMap[n,i,j]  = e_coll
                    toc = time.perf_counter()
                    
                    dmPath = os.path.join(self.cachedir, filename+'.dmmap')
                    A = {'dMmap':self.dMmap,'eMap':self.eMap,'dtRange':dtRange,'time':toc-tic,\
                         'tA':tA,'m0':m0,'lon':TL.coords.lon,'lat':TL.coords.lat,'sInds':sInds,'seed':seed,'mass':self.mass}
                    with open(dmPath, 'wb') as f:
                        pickle.dump(A, f)
                    print('Mass - ',dm*self.mass)
                    print('Best Epsilon - ',e_coll)
                

from EXOSIMS.Observatory.SotoStarshade_SKi import SotoStarshade_SKi
import numpy as np
import astropy.units as u
from scipy.integrate import solve_ivp
import astropy.constants as const
import hashlib
import scipy.optimize as optimize
from scipy.optimize import basinhopping
import scipy.interpolate as interp
import scipy.integrate as intg
from scipy.integrate import solve_bvp
from copy import deepcopy
import time
import os
try:
    import _pickle as pickle
except:
    import pickle

EPS = np.finfo(float).eps


class SotoStarshade_ContThrust(SotoStarshade_SKi):
    """ StarShade Observatory class
    This class is implemented at L2 and contains all variables, functions, 
    and integrators to calculate occulter dynamics. 
    """
    
    def __init__(self,orbit_datapath=None,**specs): 

        SotoStarshade_SKi.__init__(self,**specs)  
        
        #to convert from dimensionalized to normalized, (Dimension) / self.(Dimension)U
        self.mass = 10930*u.kg
        
        Tmax = 1040*u.mN
        Isp  = 3000*u.s
        ve   = const.g0 * Isp
        
        self.Tmax = Tmax
        self.ve   = self.convertVel_to_canonical(ve)
        
        # smoothing factor (eps = 1 means min energy, eps = 0 means min fuel)
        self.epsilon = 1 
        
        self.lagrangeMults = self.lagrangeMult()
        
# =============================================================================
# Miscellaneous        
# =============================================================================
    def DCM_r2i(self,t):
        """ Direction cosine matrix to rotate from Rotating Frame to Inertial Frame
        
        Finds rotation matrix for positions and velocities (6x6)

        Args:
            t (float):
                Rotation angle

        Returns:
            float 6x6 array:
                rotation of full 6 dimensional state from R to I frame
        """

        Ts  = np.array([[ np.cos(t) ,-np.sin(t) , 0],\
                        [ np.sin(t) , np.cos(t) , 0],\
                        [   0       ,    0      , 1]])

        dTs = np.array([[-np.sin(t) ,-np.cos(t) , 0],\
                        [ np.cos(t) ,-np.sin(t) , 0],\
                        [   0       ,    0      , 0]])

        Qu = np.hstack( [ Ts , np.zeros([3,3])])
        Ql = np.hstack( [dTs , Ts])

        rotMatrix = np.vstack([Qu,Ql])

        return rotMatrix


    def DCM_i2r(self,t):
        """ Direction cosine matrix to rotate from Inertial Frame to Rotating Frame
        
        Finds rotation matrix for positions and velocities (6x6)

        Args:
            t (float):
                Rotation angle

        Returns:
            float 6x6 array:
                rotation of full 6 dimensional state from I to R frame
        """

        Ts  = np.array([[ np.cos(t) ,-np.sin(t) , 0],\
                        [ np.sin(t) , np.cos(t) , 0],\
                        [   0       ,    0      , 1]])

        dTs = np.array([[-np.sin(t) ,-np.cos(t) , 0],\
                        [ np.cos(t) ,-np.sin(t) , 0],\
                        [   0       ,    0      , 0]])

        argD = np.matmul( np.matmul(-Ts.T,dTs) , Ts.T )
        Qu = np.hstack( [  Ts.T , np.zeros([3,3])])
        Ql = np.hstack( [  argD , Ts.T])

        return np.vstack([Qu,Ql])



    def DCM_r2i_9(self,t):
        """ Direction cosine matrix to rotate from Rotating Frame to Inertial Frame

        Finds rotation matrix for positions, velocities, and accelerations (9x9)

        Args:
            t (float):
                Rotation angle

        Returns:
            float 9x9 array:
                rotation of full 9 dimensional state from R to I frame
        """

        Ts  = np.array([[ np.cos(t) ,-np.sin(t) , 0],\
                        [ np.sin(t) , np.cos(t) , 0],\
                        [   0       ,    0      , 1]])

        dTs = np.array([[-np.sin(t) ,-np.cos(t) , 0],\
                        [ np.cos(t) ,-np.sin(t) , 0],\
                        [   0       ,    0      , 0]])

        ddTs = -np.array([[ np.cos(t) ,-np.sin(t) , 0],\
                          [ np.sin(t) , np.cos(t) , 0],\
                          [   0       ,    0      , 0]])

        Qu = np.hstack( [  Ts ,  np.zeros([3,6])                 ])
        Qm = np.hstack( [ dTs ,  Ts             , np.zeros([3,3])])
        Ql = np.hstack( [ddTs ,2*dTs             , Ts])

        return np.vstack([Qu,Qm,Ql])


    def lagrangeMult(self):
        """ Generate a random lagrange multiplier for initial guess (6x1)

        Generates an array of 6 numbers with two pairs of 3 coordinates describing position on a sphere.
        The first two represent the x and y positions on a sphere with random radius between 1 and 5,
        and the last coordinate represents the z coordinate scaled by the radius.

        Returns:
            float 6 array:
                Random lagrange multipliers for an initial guess
        """

        Lr = np.random.uniform(1,5)
        Lv = np.random.uniform(1,5)

        alpha_r = np.random.uniform(0,2*np.pi)
        alpha_v = np.random.uniform(0,2*np.pi)

        delta_r = np.random.uniform(-np.pi/2,np.pi/2)
        delta_v = np.random.uniform(-np.pi/2,np.pi/2)

        #TODO: seems odd that 3rd and 6th coordinate aren't part of a spherical coordinate

        L = np.array([Lr*np.cos(alpha_r)*np.cos(delta_r),
                 Lr*np.sin(alpha_r)*np.cos(delta_r),
                 np.sin(delta_r),
                 Lv*np.cos(alpha_v)*np.cos(delta_v),
                 Lv*np.sin(alpha_v)*np.cos(delta_v),
                 np.sin(delta_v)
                 ])

        return L


    def newStar_angularSep(self,TL,iStars,jStars,currentTime,dt):
        """Finds angular separation from old star to given list of stars
        
        This method returns the angular separation from the last observed 
        star to all others on the given list at the currentTime. This is distinct
        from the prototype version of the method in its signing of the
        angular separation based on halo velocity direction
        
        Args:
            TL (TargetList module):
                TargetList class object
            old_sInd (integer):
                Integer index of the last star of interest
            sInds (integer ndarray):
                Integer indices of the stars of interest
            currentTime (astropy Time array):
                Current absolute mission time in MJD
            dt (float):
                Timestep in units of days for determining halo velocity frame

        Returns:
            float:
                Angular separation between two target stars 
        """

        nStars = len(iStars)
        
        # time in canonical units
        absTimes = currentTime + np.array([0,dt])*u.d
        t = self.convertTime_to_canonical(np.mod(absTimes.value,self.equinox.value)*u.d) * u.rad
                
        # halo positions and velocities 
        haloPos = self.haloPosition(absTimes) + np.array([1,0,0])*self.L2_dist.to('au')
        haloVel = self.haloVelocity(absTimes)
                
        # halo positions and velocities in canonical units
        r_T0_R    = np.array([self.convertPos_to_canonical(haloPos[:,n]) for n in range(3)])
        Rdr_T0_R  = np.array([self.convertVel_to_canonical(haloVel[:,n]) for n in range(3)])
        
        # V-frame
        v1_R = Rdr_T0_R / np.linalg.norm(Rdr_T0_R,axis=0)
        
        # Position of stars
        coords = TL.coords
        beta = coords.lat
        lamb = coords.lon
        dist = self.convertPos_to_canonical( coords.distance[0] )
        
        # unit vector to star locations
        u_i0_R = np.array([   [np.cos(beta[iStars]) * np.cos(lamb[iStars]-t[0]) ], \
                              [np.cos(beta[iStars]) * np.sin(lamb[iStars]-t[0]) ], \
                              [np.sin(beta[iStars])]  ]).reshape(3,nStars)
        
        u_j0_R = np.array([   [np.cos(beta[jStars]) * np.cos(lamb[jStars]-t[-1]) ], \
                              [np.cos(beta[jStars]) * np.sin(lamb[jStars]-t[-1]) ], \
                              [np.sin(beta[jStars])]  ]).reshape(3,nStars)
        
        # star locations to relative to telescope
        r_iT_R = dist * u_i0_R  - r_T0_R[:,0].reshape(3,1)
        r_jT_R = dist * u_j0_R  - r_T0_R[:,-1].reshape(3,1)
        
        # unit vector from Telescope to star i
        u_iT = r_iT_R/np.linalg.norm(r_iT_R,axis=0)
    
        # unit vector from Telescope to star j
        u_jT = r_jT_R/np.linalg.norm(r_jT_R,axis=0)
        
        # direction of travel from star i to j
        u_ji =  u_jT - u_iT
        
        # sign of the angular separation, + in direction of halo velocity (NOTE: only using halo velocity at t0, not t0+dt)
        dot_sgn = np.matmul( v1_R[:,0].T, u_ji )
        sgn = np.array([np.sign(x) if x != 0 else 1 for x in dot_sgn])
        
        # calculating angular separation and assigning signs
        dot_product = np.array([ np.dot(a.T,b) for a,b in zip(u_iT.T,u_jT.T)])
        dot_product = np.clip(dot_product, -1, 1)
        psi = sgn*(np.arccos( dot_product )*u.rad).to('deg').value
        
        return psi


    def star_angularSepDesiredDist(self,psiPop,nSamples=1e6,angBinSize=3):
        """ Rejection sample from psiPop to achieve nSamples fitting logistic distribution

        Args:
            psiPop (array float):
                List of floats between -180 and 180
            nSamples (int):
                Number of samples to return
            angBinSize (float):
                Size of a bin for the histogram used in rejection sampling

        Returns:
            tuple:
                array float:
                    Angles (nSamples) fitting the logistic distribution
                array int:
                    An array of indices of original psiPop that were accepted

        """
        
        # distribution of the full psi population
        N = len(psiPop)
        psiBins = np.arange(-180,181,3)
        envelopeDist, EnvDbins = np.histogram(psiPop,psiBins,density=True)
    
        #Desired Distribution pdf
        s = 32
        desiredDist = 1/(4*s) * 1/(  np.cosh(psiBins[:-1]/(2*s))**2  )
        
        #Scaling factor
        M = np.max( desiredDist / envelopeDist )
        
        # how many samples from the full data set are we taking?
        samplingSize = int(nSamples) if (nSamples > 0) & (nSamples <= N) else N
        
        # randomly selecting from the population (no replacing, so unique samples)
        sampling_inds = np.random.choice( N , samplingSize , replace=False  )
        
        # an empty array to which we'll add accepted indices
        final_inds = np.array([],dtype=int)
            
        # running rejection sampling
        for n in sampling_inds:
            
            # figuring out which bin to put this particular ang sep in
            corresponding_Bin = np.where(psiBins <= psiPop[n])[0][-1]
            
            # calculating pdf ratio
            ratio = desiredDist[corresponding_Bin] / (M*envelopeDist[corresponding_Bin])
            
            # random number from 0 to 1
            rand    = np.random.uniform(0,1)
            
            # sample has been accepted!
            if rand < ratio:
                final_inds = np.hstack([ final_inds, n ])
        
        psiFiltered = psiPop[final_inds]
        
        return psiFiltered, final_inds
    
    
    def selectPairsOfStars(self,TL,nPairs,currentTime,dt,nSamples):
        """ Rejection sampling of star pairings using desired distribution and sampling from that final distribution
        
        Args:
            TL (TargetList module):
                TargetList class object
            nPairs (int):
                Number of pairs to produce
            currentTime (astropy Time array):
                Current absolute mission time in MJD
            dt (float):
                Timestep in units of days for determining halo velocity frame
            nSamples (int):
                Number of samples to be used when generating random stars for pairing
        Returns:
            tuple:
                int list:
                    The list of indices corresponding to starting stars in TargetList
                int list:
                    The list of indices corresponding to the ending stars in TargetList
                float array:
                    Angular distance between pairs from iFinal and jFinal
        
        """
        
        psiBins = np.arange(-180,181,3)
        nPairs -= len(psiBins)-1

        # randomly selected pairs of stars
        iLog = np.random.randint(0,TL.nStars,nSamples)
        jLog = np.random.randint(0,TL.nStars,nSamples)
        
        # angular separation between each of these pairs
        # star i -> at currentTime
        # star j -> at currentTime + dt
        psi = self.newStar_angularSep(TL,iLog,jLog,currentTime,dt)
        
        # filtering pairings list, creating desired distribution using rejection sampling
        filtered_psi, filtered_inds = self.star_angularSepDesiredDist(psi,nSamples)
        
        # sampling nPairs of stars from the final distribution
        samples = np.random.choice( len(filtered_inds) , nPairs , replace=False  )
        sampling_inds = filtered_inds[samples]
        
        # non-sampled star pairings
        leftover_psi  = np.delete(psi,sampling_inds)
        leftover_inds = np.delete( np.arange(0,nSamples) , sampling_inds)
        
        # grab a star from each bin and add it to the list
        psiFinal = np.array([],dtype=int)
        
        for lb,ub in zip(psiBins[:-1],psiBins[1:]):
            find_one = np.where( (leftover_psi >= lb) & (leftover_psi < ub) )[0]
            if find_one.size > 0:
                sampling_inds = np.hstack([ sampling_inds, leftover_inds[find_one[0]] ])
                psiFinal = np.hstack([ psiFinal , leftover_psi[find_one[0]] ])
        
        # results
        iFinal = iLog[sampling_inds]
        jFinal = jLog[sampling_inds]
        
        # add the sampled stars from our desired distributions
        psiFinal = np.hstack([ psiFinal , filtered_psi[samples] ]) 
        
        return iFinal,jFinal,psiFinal
# =============================================================================
# Helper functions
# =============================================================================
        
    def determineThrottle(self,state):
        """ Determines throttle based on instantaneous switching function value.
        
        Typically being used during collocation algorithms. A zero-crossing of
        the switching function is highly unlikely between the large number of 
        nodes.

        Args:
            state (float array):
                State vector and lagrange multipliers
        Returns:
            float array:
                Throttle between 0 and 1, and in the same shape as state
        """
        
        eps = self.epsilon
        n = 1 if state.size == 14 else state.shape[1]
        
        throttle = np.zeros(n)
        S = self.switchingFunction(state)
        S = S.reshape(n)
        
        for i,s in enumerate(S):
            if eps > 0:
                midthrottle = (eps - s)/(2*eps)
                throttle[i] = 0 if s > eps else 1 if s < -eps else midthrottle
            else:
                throttle[i] = 0 if s > eps else 1
        
        return throttle
    
    
    def switchingFunction(self,state):
        """ Evaluates the switching function at specific states.

        Args:
            state (float array):
                State vector and lagrange multipliers
        Returns:
            float:
                Value of the switching function

        """
        
        x,y,z,dx,dy,dz,m,L1,L2,L3,L4,L5,L6,L7 = state
        
        Lv_, lv = self.unitVector( np.array([L4,L5,L6]) )
        
        S = -lv*self.ve/m - L7 + 1
        
        return S

    def switchingFunctionDer(self,state):
        """ Evaluates the time derivative of the switching function.
        
        Switching function derivative evaluated for specific states.

        Args:
            state (float array):
                State vector and lagrange multipliers
        Returns:
            float:
                Value of the switching function time derivative

        """
        ve  = self.ve
        n = 1 if state.size == 14 else state.shape[1]
        x,y,z,dx,dy,dz,m,L1,L2,L3,L4,L5,L6,L7 = state
        
        Lr = np.array([L1,L2,L3]).reshape(3,n)
        Lv = np.array([L4,L5,L6]).reshape(3,n)
        Lv_, lv = self.unitVector( Lv )
        
        Pv_arr = np.array([[0,2,0],[-2,0,0],[0,0,0]])
        Pv = np.dstack([Pv_arr]*n)
        
        PLdot = np.vstack( [np.dot(a.T,b) for a,b in zip(Pv.T,Lv.T)] ).T
        
        dS = -(ve / m) * np.vstack( [np.dot(a.T,b) for a,b in zip( (-Lr - PLdot).T ,Lv_.T)] ).T 

        return dS
    
    def selectEventFunctions(self,s0):
        """ Selects the proper event function for integration.
        
        This method calculates the switching function and its derivative at a 
        single specific state. It then determines which thrust case it will be 
        in: full, medium, or no thrust. If the value of the switching function
        is within a certain tolerance of the boundaries, it uses the derivative 
        to determine the direction it is heading in. Then the proper event 
        functions are created for the integrator to determine the next crossing
        (i.e. the next case change). 
        
        Args:
            s0 (float array):
                Iniitial state vector and lagrange multipliers
        Returns:
            tuple:
                function list:
                    A list of functions which are the switching function adjusted by 1e-10.
                    These take in an unused first parameter, and then the state.
                    Depending on the case (2), multiple functions are returned.
                    These functions will return 0 when the switching function crosses
                    1e-10 or -1e-10.
                int:
                    An integer 0,1,2 representing whether the switching
                    function of the initial state is within a 1e-10 tolerance of the
                    of zero (2), less than -1e-10 (1), or greater than 1e-10 (0). 
        """
        eps = self.epsilon
        
        S  =  self.switchingFunction(s0)
        dS =  self.switchingFunctionDer(s0)[0]

        # finding which case we are in:
        #   - case 2 if       -eps < S < eps       (medium thrust)
        #   - case 1 if   S < -eps                 (full thrust)
        #   - case 0 if                  eps < S   (no thrust)
        
        case = 0 if S > eps else 1 if S < -eps else 2

        # checking to see if S is within a certain tolerance from epsilon
        withinTol = np.abs( (np.abs(S) - eps) ) < 1e-10
        # determine if there is a case error if within tolerance
        if withinTol:
            # not the minimum fuel case
            if eps != 0:
                # at the upper bound, case determined by derivative
                if S > 0:
                    case = 2 if dS < 0 else 0
                # at the lower bound, case determined by derivative
                else:
                    case = 2 if dS > 0 else 1
            # minimum fuel case, only two cases
            else:
                case = 0 if dS > 0 else 1
            
            
        eventFunctions = []
        CrossingUpperBound = lambda t,s : self.switchingFunction(s) - eps
        CrossingLowerBound = lambda t,s : self.switchingFunction(s) + eps
        
        CrossingUpperBound.terminal = True
        CrossingLowerBound.terminal = True
        
        if case == 0:
            # crossing upper epsilon from above
            CrossingUpperBound.direction = -1 
            # appending event function
            eventFunctions.append(CrossingUpperBound)
        elif case == 1:
            # crossing lower epsilon from below
            CrossingLowerBound.direction = 1 
            # appending event function
            eventFunctions.append(CrossingLowerBound)
        else:
            # can either cross lower epsilon from above or upper from below
            CrossingLowerBound.direction = -1
            CrossingUpperBound.direction = 1
            # appending event function
            eventFunctions.append(CrossingUpperBound)
            eventFunctions.append(CrossingLowerBound)
        
        return eventFunctions,case
        

# =============================================================================
# Equations of Motion and Boundary Conditions
# =============================================================================
    def boundary_conditions_thruster(self,sA,sB,constrained=False):
        """ Creates boundary conditions for solving a boundary value problem

        Function used in scipy solve_bvp function call. Returns residuals

        Args:
            sA (6 or 7 array):
                To be compared with the initial state for boundary condition.
            sB (6 or 7 array):
                To be compared with the final state for boundary condition.
            constrained (boolean):
                Whether there are 6 (false) or 7 (true) boundary conditions. 
        Returns:
            array:
                Returns the residuals of the boundary conditions/constraint equation.
        """
    
        BCo1 = sA[0] - self.sA[0]
        BCo2 = sA[1] - self.sA[1]
        BCo3 = sA[2] - self.sA[2]
        BCo4 = sA[3] - self.sA[3]
        BCo5 = sA[4] - self.sA[4]
        BCo6 = sA[5] - self.sA[5]
        
        BCf1 = sB[0] - self.sB[0]
        BCf2 = sB[1] - self.sB[1]
        BCf3 = sB[2] - self.sB[2]
        BCf4 = sB[3] - self.sB[3]
        BCf5 = sB[4] - self.sB[4]
        BCf6 = sB[5] - self.sB[5]
        
        if constrained:
            BCo7 = sA[6] - self.sA[6]
            BCf7 = sB[-1]
            BC = np.array([BCo1,BCo2,BCo3,BCo4,BCo5,BCo6,BCo7,BCf1,BCf2,BCf3,BCf4,BCf5,BCf6,BCf7])
        else:
            BC = np.array([BCo1,BCo2,BCo3,BCo4,BCo5,BCo6,BCf1,BCf2,BCf3,BCf4,BCf5,BCf6])

        return BC   
     
        
    def EoM_Adjoint(self,t,state,constrained=False,amax=False,integrate=False):
        """ Equations of Motion with costate vectors

        Equations of motion for CR3BP with costates for control.

        Args:
            t (float):
                Currently unused
            state (array):
                The state and lagrange multipliers.
            constrained (boolean):
                Currently unused
            amax (float or boolean):
                Maximum acceleration attainable or False otherwise
            integrate (boolean):
                If true, the array is flattened for integration.

        Returns:
            array:
                The time derivatives of the states and costates.
        """
        
        mu = self.mu
        ve = self.ve
        
        n = 1 if state.size == 14 else state.shape[1]

        if amax:
            x,y,z,dx,dy,dz,m,L1,L2,L3,L4,L5,L6,L7 = state
        else:
            x,y,z,dx,dy,dz,L1,L2,L3,L4,L5,L6 = state
        
        if integrate:
            state = state.T
        
        # vector distances from primaries
        r1   = np.array([x-(-mu),y,z])
        r2   = np.array([x-(1-mu),y,z])
        # norms of distances from primaries
        R1   = np.linalg.norm(r1,axis=0)
        R2   = np.linalg.norm(r2,axis=0)
        
        # Position-dependent acceleration terms
        qx = np.array([ x  - (1-mu)*(x+mu)/R1**3  - mu*(x+mu-1)/R2**3]).reshape(1,n)
        qy = np.array([ y  - (1-mu)*y/R1**3       - mu*y/R2**3]).reshape(1,n)
        qz = np.array([    - (1-mu)*z/R1**3       - mu*z/R2**3]).reshape(1,n)
        q = np.vstack([qx,qy,qz])  #shape of 3xn
        
        # Position partial derivatives
        Q11 = 1 - (1-mu)/R1**3 + 3*(1-mu)*(x+mu)**2/R1**5  - mu/R2**3   + 3*mu*(x+mu-1)**2/R2**5
        Q22 = 1 - (1-mu)/R1**3 + 3*(1-mu)*     y**2/R1**5  - mu/R2**3   + 3*mu*       y**2/R2**5
        Q33 =   - (1-mu)/R1**3 + 3*(1-mu)*     z**2/R1**5  - mu/R2**3   + 3*mu*       z**2/R2**5
        Q12 =                    3*(1-mu)*(x+mu)* y/R1**5               + 3*mu*(x+mu-1)* y/R2**5
        Q13 =                    3*(1-mu)*(x+mu)* z/R1**5               + 3*mu*(x+mu-1)* z/R2**5
        Q23 =                    3*(1-mu)*      y*z/R1**5               + 3*mu*        y*z/R2**5
        Qr   = np.array([[Q11, Q12 , Q13], [Q12, Q22 , Q23], [Q13, Q23 , Q33]]).reshape(3,3,n) # shape of 3x3xn

        # Velocity-dependent acceleration terms
        px =  2*dy
        py = -2*dx
        pz = np.zeros([1,n])
        p  = np.vstack([px,py,pz]) #shape of 3xn
        
        # Velocity partial derivatives
        Pv_arr = np.array([[0,2,0],[-2,0,0],[0,0,0]])
        Pv = np.dstack([Pv_arr]*n)
        
        # Costate vectors
        Lr = np.vstack([L1,L2,L3])
        Lv = np.vstack([L4,L5,L6])
        Lr_, lr = self.unitVector(Lr)
        Lv_, lv = self.unitVector(Lv)
        
        # ================================================
        # Equations of Motion
        # ================================================
        dX  = np.vstack([ dx,dy,dz ])
        dV  = q + p
        dLx = -np.vstack( [np.dot(a.T,b) for a,b in zip(Qr.T,Lv.T)] ).T
        dLv = -Lr - np.vstack( [np.dot(a.T,b) for a,b in zip(Pv.T,Lv.T)] ).T
        
        if amax:
            # throttle factor
            throttle = self.determineThrottle(state)
            
            dV -= Lv_ * amax * throttle / m
            dm  = -throttle * amax / ve
            dLm = -lv * throttle * amax / m**2
            # putting them all together, a 14xn array
            f   = np.vstack([ dX, dV, dm, dLx, dLv, dLm ])
        else:
            dV -= Lv
            # putting them all together, a 12xn array
            f   = np.vstack([ dX, dV, dLx, dLv ])
        
        
        if integrate:
            f = f.flatten()

        return f
    
    
    def starshadeBoundaryVelocity(self,TL,sInd,currentTime,SRP=True):
        """ Calculates starshade and telescope velocities in R- or I-frames during stationkeeping
        
        Calculates the rotating system position and velocity of the
        starshade at insertion into the optimal initial state for
        observation of the given star.

        Args:
            TL (TargetList module):
                Target List
            sInd (int):
                Index of star for observation in the target list
            currentTime (astropy Time array):
                Current absolute mission time in MJD
            SRP (boolean):
                Whether or not solar radiation pressure should be used 
                in calculating the insertion state for optimal 
                starshade stationkeeping observation position.   
        Returns:
            tuple:
                array:
                    Position in the rotating frame
                array:
                    Velocity in the rotating frame
        """
        
        # absolute times (Note: equinox is start time of Halo AND when inertial frame and rotating frame match)
        modTimes = np.mod(currentTime.value,self.equinox.value)*u.d  #mission times relative to equinox )
        t = self.convertTime_to_canonical(modTimes) * u.rad       #modTimes in canonical units 
        
        r_S0_I, Iv_S0_I, qq, qq, qq, qq = self.starshadeKinematics(TL,sInd,currentTime)
        
        # star a at t=ta
        r_AS_C, Iv_AS_C = self.starshadeInjectionVelocity(TL,sInd,currentTime,SRP=SRP)
        s_AS_C = np.hstack([r_AS_C , Iv_AS_C])
        r_AS_I, Iv_AS_I = self.rotateComponents2NewFrame(TL,sInd,currentTime,s_AS_C,np.array([0]),SRP=SRP,final_frame='I')
        
        rQi__tA = self.rot( t[0].value , 3 )
        r_A0_I  =  r_S0_I +  r_AS_I
        Iv_A0_I = Iv_S0_I + Iv_AS_I
        
        r_A0_R   = np.dot(  rQi__tA , r_A0_I  )
        Rv_A0_I  = deepcopy(Iv_A0_I)
        Rv_A0_I[0] += r_A0_I[1]
        Rv_A0_I[1] -= r_A0_I[0]
        Rv_A0_R  = np.dot(  rQi__tA , Rv_A0_I  )
        
        return r_A0_R, Rv_A0_R
    
# =============================================================================
# Initial conditions
# =============================================================================

    def findInitialTmax(self,TL,nA,nB,tA,dt,m0=1,s_init=np.array([])):
        """ Finding initial guess for starting Thrust
        
        Args:
            TL (TargetList module):
                Target list
            nA (int):
                The index for the starting star in the target list
            nB (int):
                The index for the final star in the target list
            tA (astropy Time array):
                Initial absolute mission time in MJD
            dt (astropy Time array);
                A time step between the two voundary conditions
            m0 (float):
                Initial mass
            s_init (array):
                An initial guess for the state

        Returns:
            tuple:
                astropy Newtons:
                    The maximum initial thrust.
                array:
                    States from the solved bvp.
                array:
                    Times at corrsponding to each state.
        """
        
        tB = tA + dt
        #angle,uA,uB,r_tscp = self.lookVectors(TL,nA,nB,tA,tB)

        #position vector of occulter in heliocentric frame
        #self_rA = uA*self.occulterSep.to('au').value + r_tscp[ 0]
        #self_rB = uB*self.occulterSep.to('au').value + r_tscp[-1]
        
        r_A0_R, Rv_A0_R = self.starshadeBoundaryVelocity(TL,nA,tA)
        r_B0_R, Rv_B0_R = self.starshadeBoundaryVelocity(TL,nB,tB)
           
        Rs_A0_R = np.hstack([r_A0_R[:,0],Rv_A0_R[:,0]])
        Rs_B0_R = np.hstack([r_B0_R[:,0],Rv_B0_R[:,0]])
                
        Rsf_A0_R = np.hstack([Rs_A0_R, self.lagrangeMult()])
        Rsf_B0_R = np.hstack([Rs_B0_R, self.lagrangeMult()])
                
        a = ((np.mod(tA.value,self.equinox.value)*u.d)).to('yr') / u.yr * (2*np.pi)
        b = ((np.mod(tB.value,self.equinox.value)*u.d)).to('yr') / u.yr * (2*np.pi)
        
        #running collocation
        tGuess = np.hstack([a,b]).value

        if s_init.size:
            sGuess = np.vstack([s_init[:,0] , s_init[:,-1]])
        else:
            sGuess = np.vstack([Rsf_A0_R , Rsf_B0_R])
            
        s,t_s,status = self.send_it_thruster(sGuess.T,tGuess,m0=m0,verbose=False)
        
        lv = s[9:,:]

        aNorms0 = np.linalg.norm(lv,axis=0)
        aMax0   = self.convertAcc_to_dim( np.max(aNorms0) ).to('m/s^2') 
        Tmax0   = (aMax0 * self.mass ).to('N')
        
        return Tmax0, s, t_s
    
    
    def findTmaxGrid(self,TL,tA,dtRange):
        """ Create grid of Tmax values using unconstrained thruster
        
        This method is used purely for creating figures. 
        
        Args:
            TL (TargetList module):
                Target list
            tA (astropy Time array):
                Initial absolute mission time in MJD
            dtRange (astropy Time array);
                An array of delta times to consider

        Returns:
            astropy Newtons array:
                Max thrust in Newtons (dtRange dimensions by TL.nStars) 
        """
        
        midInt = int( np.floor( (TL.nStars-1)/2 ) )

        sInds = np.arange(0,TL.nStars)
        ang   =  self.star_angularSep(TL, midInt, sInds, tA) 
        sInd_sorted = np.argsort(ang)
        angles  = ang[sInd_sorted].to('deg').value
        
        
        TmaxMap = np.zeros([len(dtRange) , len(angles)])*u.N
        
        for i,t in enumerate(dtRange):
            for j,n in enumerate(sInd_sorted):
                print(i,j)
                Tmax, s, t_s = self.findInitialTmax(TL,midInt,n,tA,t)
                TmaxMap[i,j] = Tmax
        
        return TmaxMap

# =============================================================================
# BVP solvers
# =============================================================================
  
    def send_it_thruster(self,sGuess,tGuess,aMax=False,constrained=False, m0 = 1, maxNodes=1e5,verbose=False):
        """ Solving generic bvp from t0 to tF using states and costates

        Uses the Scipy solve_bvp method.

        Args:
            sGuess (array):
                Initial state and costate guess
            tGuess (astropy Time array):
                Times corresponding to the guess of the state at each time
            aMax (astropy Newton or boolean):
                Maximum attainable acceleration
            constrained (boolean):
                Flag for whether this is constrained or unconstrained problem
                This determines dimensions of the state inputs.
            m0 (float):
                Initial mass
            maxNodes (int):
                Maximum number of nodes to use in the BVP problem solution
            verbose (boolean):
                Flag passed to solve_bvp
        Returns:
            tuple:
                array:
                    State and costate at the various mesh times
                array:
                    Corresponding times to the sampled states.
                boolean:
                    Status returned by the bvp_solve method
        """
        
        sG = sGuess
        # unconstrained problem begins with 12 states, rather than 14. checking for that
        if len(sGuess) == 12:
            x,y,z,dx,dy,dz,L1,L2,L3,L4,L5,L6 = sGuess
            # if unconstrained is initial guess for constrained problem, make 14 state array
            if aMax:
                mRange  = np.linspace(m0, 0.8, len(x))
                lmRange = np.linspace(1, 0, len(x))
                sG = np.vstack([x,y,z,dx,dy,dz,mRange,L1,L2,L3,L4,L5,L6,lmRange])

        # only saves initial and final desired states if first solving unconstrained problem
        if not constrained:
            self.sA = np.hstack([ sG[0:6,0 ], m0   , sG[6:,0]  , 1 ])
            self.sB = np.hstack([ sG[0:6,-1], 0.8 , sG[6:,-1] , 0 ])
            self.sG = sG
        
        # creating equations of motion and boundary conditions functions
        EoM = lambda t,s: self.EoM_Adjoint(t,s,constrained,aMax)
        BC  = lambda t,s: self.boundary_conditions_thruster(t,s,constrained)
        # solving BVP
        sol = solve_bvp(EoM,BC,tGuess,sG,tol=1e-8,max_nodes=int(maxNodes),verbose=0)
        
        if verbose:
            self.vprint(sol.message)
        # saving results
        s = sol.y
        t_s = sol.x
            
        return s,t_s,sol.status
    
    
    def collocate_Trajectory(self,TL,nA,nB,tA,dt):
        """ Solves minimum energy and minimum fuel cases for continuous thrust

        Args:
            TL (TargetList module):
                Target list
            nA (int):
                The index for the starting star in the target list
            nB (int):
                The index for the final star in the target list
            tA (astropy Time array):
                Initial absolute mission time in MJD
            dt (astropy Time array);
                A time step between the two voundary conditions
        
        Returns:
            tuple:
                array:
                    Trajectory
                array:
                    Times corresponding to trajectory
                float:
                    last epsilon that fully converged (2 if minimum energy didn't work)
                    Parameterizes minimum energy to minimum fuel solution.
                astropy Newton array:
                    Range of thrusts (Newtons) considered.
        """
        
        # initializing arrays
        stateLog = []
        timeLog  = []
        
        # solving using unconstrained thruster as initial guess
        Tmax, sTmax, tTmax = self.findInitialTmax(TL,nA,nB,tA,dt)
        aMax = self.convertAcc_to_canonical( (Tmax / self.mass).to('m/s^2') )
        # saving results
        stateLog.append(sTmax)
        timeLog.append(tTmax)
        
        # all thrusts were successful
        e_best   = 2
        s_best   = deepcopy(stateLog)
        t_best   = deepcopy(timeLog)
        
        # thrust values
        desiredT = self.Tmax.to('N').value
        currentT = Tmax.value
        # range of thrusts to try
        TmaxRange    = np.linspace(currentT,desiredT,30)
        
        # range of epsilon values to try (e=1 is minimum energy, e=0 is minimum fuel)
        epsilonRange = np.round( np.arange(1,-0.1,-0.1) , decimals = 1)
        
        # Huge loop over all thrust and epsilon values:
        #   we start at the minimum energy case, e=1, using the thrust value from the unconstrained solution
        #     In/De-crement the thrust until we reach the desired thrust level
        #   then we decrease e and repeat process until we get to e=0 (minimum fuel)
        #   saves the last successful result in case collocation fails
        
        # loop over epsilon starting with e=1
        for j,e in enumerate(epsilonRange):
            print("Collocate Epsilon = ",e)
            # initialize epsilon
            self.epsilon = e
            
            # loop over thrust values from current to desired thrusts
            for i,thrust in enumerate(TmaxRange):
                #print("Thrust #",i," / ",len(TmaxRange))
                # convert thrust to canonical acceleration
                aMax = self.convertAcc_to_canonical( (thrust*u.N / self.mass).to('m/s^2') )
                # retrieve state and time initial guesses
                sGuess = stateLog[i]
                tGuess = timeLog[i]
                # perform collocation
                s,t_s,status = self.send_it_thruster(sGuess,tGuess,aMax,constrained=True,maxNodes=1e5,verbose=False)
                
                # collocation failed, exits out of everything
                if status != 0:
                    self.epsilon = e_best
                    if e_best == 2:
                        # if only the unconstrained problem worked, still returns a 14 length array
                        s_out = []
                        length = s_best[0].shape[1]
                        m  = np.linspace(1,0.9,length)
                        lm = np.linspace(0.3,0,length)
                        s_out.append( np.vstack([s_best[0][:6] , m, s_best[0][6:], lm]) )
                        s_best = deepcopy(s_out)
                    return s_best, t_best, e_best, TmaxRange
                
                # collocation was successful!
                if j == 0:
                    # creates log of state and time results for next thrust iteration (at the beginning of the loop)
                    stateLog.append(s)
                    timeLog.append(t_s)
                else:
                    # updates log of state and time results for next thrust iteration
                    stateLog[i] = s
                    timeLog[i]  = t_s
            
            # all thrusts were successful, save results
            e_best   = self.epsilon
            s_best   = deepcopy(stateLog)
            t_best   = deepcopy(timeLog)
        
        return s_best, t_best, e_best, TmaxRange


    def collocate_Trajectory_minEnergy(self,TL,nA,nB,tA,dt,m0=1):
        """ Solves minimum energy and minimum fuel cases for continuous thrust
        
        Args:
            TL (TargetList module):
                Target list
            nA (int):
                The index for the starting star in the target list
            nB (int):
                The index for the final star in the target list
            tA (astropy Time array):
                Initial absolute mission time in MJD
            dt (astropy Time array);
                A time step between the two voundary conditions
            m0 (float):
                Initial mass

        Returns:
            tuple:
                array:
                    Trajectory
                array:
                    Times corresponding to trajectory
                float:
                    last epsilon that fully converged (2 if minimum energy didn't work)
                    Parameterizes minimum energy to minimum fuel solution.
                astropy Newton array:
                    Range of thrusts (Newtons) considered.
        """
        
        # initializing arrays
        stateLog = []
        timeLog  = []
        
        # solving using unconstrained thruster as initial guess
        Tmax, sTmax, tTmax = self.findInitialTmax(TL,nA,nB,tA,dt,m0)
        aMax = self.convertAcc_to_canonical( (Tmax / self.mass).to('m/s^2') )
        # saving results
        stateLog.append(sTmax)
        timeLog.append(tTmax)
        
        # all thrusts were successful
        e_best   = 2
        s_best   = deepcopy(sTmax)
        t_best   = deepcopy(tTmax)
        
        # thrust values
        desiredT = self.Tmax.to('N').value
        currentT = Tmax.value
        # range of thrusts to try
        TmaxRange    = np.linspace(currentT,desiredT,30)
        
        # Huge loop over all thrust and epsilon values:
        #   we start at the minimum energy case, e=1, using the thrust value from the unconstrained solution
        #     In/De-crement the thrust until we reach the desired thrust level
        #   then we decrease e and repeat process until we get to e=0 (minimum fuel)
        #   saves the last successful result in case collocation fails
        
        # loop over epsilon starting with e=1
        self.epsilon = 1
        
        # loop over thrust values from current to desired thrusts
        for i,thrust in enumerate(TmaxRange):
            #print("Thrust #",i," / ",len(TmaxRange))
            # convert thrust to canonical acceleration
            aMax = self.convertAcc_to_canonical( (thrust*u.N / self.mass).to('m/s^2') )
            # retrieve state and time initial guesses
            sGuess = stateLog[i]
            tGuess = timeLog[i]
            # perform collocation
            s,t_s,status = self.send_it_thruster(sGuess,tGuess,aMax,constrained=True,m0=m0,maxNodes=1e5,verbose=False)
            
            # collocation failed, exits out of everything
            if status != 0:
                self.epsilon = e_best
                return s_best, t_best, e_best, TmaxRange
            
            # creates log of state and time results for next thrust iteration (at the beginning of the loop)
            stateLog.append(s)
            timeLog.append(t_s)

        # all thrusts were successful, save results
        e_best   = self.epsilon
        s_best   = deepcopy(s)
        t_best   = deepcopy(t_s)
    
        return s_best, t_best, e_best, TmaxRange
    
# =============================================================================
# Shooting Algorithms
# =============================================================================
        
    def integrate_thruster(self,sGuess,tGuess,Tmax,verbose=False):
        """ Integrates thruster trajectory with thrust case switches
        
        This methods integrates an initial guess for the spacecraft state 
        forwards in time. It uses event functions to find the next zero of the
        switching function which means a new thrust case is needed (full, 
        medium or no thrust).

        Args:
            sGuess (array):
                Initial state and costate guess
            tGuess (astropy Time array):
                Times corresponding to the guess of the state at each time
            Tmax (astropy force):
                Maximum thrust attainable

        Returns:
            tuple:
                array:
                    Trajectory states
                astropy Time array:
                    Times corresponding to states
        """
        s0 = sGuess[:,0]
        t0 = tGuess[0]
        tF = tGuess[-1]
        
        # initializing starting time
        tC = deepcopy(t0)
        
        # converting Tmax to a canonical acceleration
        Tmax = Tmax.to('N').value
        aMax = self.convertAcc_to_canonical( (Tmax*u.N / self.mass).to('m/s^2') )
        
        # establishing equations of motion
        EoM  = lambda t,s: self.EoM_Adjoint(t,s,constrained=True,amax=aMax,integrate=True)
        
        # starting integration
        count = 0
        while tC < tF:
            # selecting the switch functions with correct boundaries
            switchFunctions,case = self.selectEventFunctions(s0)
            if verbose:
                print("[%.3f / %.3f] with case %d" % (tC,tF,case) )
            # running integration with event functions
            res = solve_ivp(EoM, [tC , tF], s0, events=switchFunctions)
            
            # saving final integration time, if greater than tF, we completed the trajectory
            tC = deepcopy(res.t[-1])
            s0 = deepcopy(res.y[:,-1])
            
            # saving results in a new array
            if count == 0:
                sLog = deepcopy(res.y)
                tLog = deepcopy(res.t)
            #adding to the results log if there was a previous thrust case switch
            else:
                sLog = np.hstack([sLog,res.y])
                tLog = np.hstack([tLog,res.t])
            count += 1
            
        return sLog,tLog


    def conFun_singleShoot(self,w,t0,tF,Tmax,returnLog=False):
        """ Objective Function for single shooting thruster
        
        Args:
            w (array):
                Costate vector
            t0 (astropy Time object):
                Initial time
            tf (astropy Time object):
                Final time
            Tmax (astropy force):
                Maximum thrust attainable
            returnLog (boolean):
                Return the states and times of the solution
        Returns:
            tuple:
                float:
                    Norm of difference between current state and boundary value
                array:
                    Trajectory states
                astropy Time array:
                    Times corresponding to states
        """
        
        sInit  = np.hstack([self.sA[:7] , w]).reshape(14,1)
        tGuess = np.array([t0,tF])
        
        sLog,tLog = self.integrate_thruster(sInit,tGuess,Tmax)
        
        f = self.boundary_conditions_thruster( sLog[:,0], sLog[:,-1], constrained=True)
        fnorm = np.linalg.norm(f)
        
        if returnLog:
            return fnorm,sLog,tLog
        else:
            return fnorm


    def minimize_TerminalState(self,s_best,t_best,Tmax,method):
        """ Minimizes boundary conditions for thruster

        Args:
            sGuess (array):
                Initial state and costate guess
            tGuess (astropy Time array):
                Times corresponding to the guess of the state at each time
            Tmax (astropy force):
                Maximum thrust attainable
            method (string):
                Optimization method for Scipy minimize call
        Returns:
            tuple:
                float:
                    Norm of difference between current state and boundary value
                array:
                    Trajectory states
                astropy Time array:
                    Times corresponding to states
        """
        
        w0 = s_best[7:,0]
        t0 = t_best[0]
        tF = t_best[-1]
        
        res = optimize.minimize(self.conFun_singleShoot, w0, method=method, \
                                tol=1e-12, args=(t0,tF,Tmax,) )
#        minimizer_kwargs = {"method":method,"args":(t0,tF,Tmax,)}
#        res = optimize.basinhopping(self.conFun_singleShoot,w0,minimizer_kwargs=minimizer_kwargs)
#        
        fnorm, sLog, tLog = self.conFun_singleShoot(res.x,t0,tF,Tmax,returnLog=True)
        
        return fnorm, sLog, tLog


    def singleShoot_Trajectory(self,stateLog,timeLog,e_best,TmaxRange,method='SLSQP'):
        """ Perform single shooting to solve the boundary value problem.

        Args:
            stateLog (array):
                Approximate trajectory typically determined by collocation
            timeLog (astropy Time array):
                Corresponging time values for the trajectory
            e_best (float):
                Epsilon value corresponding to previous approxmiate trajectory
            TmaxRange (astropy Newton array):
                Range of thrusts (Newtons) considered.
            method (string):
                Optimization method for Scipy minimize call

        Returns:
            tuple:
                array:
                    Trajectory states
                astropy Time array:
                    Times corresponding to states
                float:
                    Epsilon value determining how fuel vs energy optimal the trajectory is.
        """


        # initializing arrays
        s_best = deepcopy(stateLog)
        t_best = deepcopy(timeLog)
        
        incLogFlag = True if len(stateLog) != len(TmaxRange) else False
        
        # range of epsilon values to try (e=1 is minimum energy, e=0 is minimum fuel)
        e_best = 1 if e_best == 2 else e_best
        epsilonRange = np.round( np.arange(e_best,-0.1,-0.1) , decimals = 1)
        
        # Huge loop over all thrust and epsilon values, just like collocation method
        
        # loop over epsilon starting with e=1
        for j,e in enumerate(epsilonRange):
            print("SS Epsilon = ",e)
            # initialize epsilon
            self.epsilon = e
            
            # loop over thrust values from current to desired thrusts
            for i,thrust in enumerate(TmaxRange):
                #print("Thrust #",i," / ",len(TmaxRange))
                # retrieve state and time initial guesses
                sGuess = stateLog[i]
                tGuess = timeLog[i]
                # perform single shooting
                fnorm,sLog,tLog = self.minimize_TerminalState(sGuess,tGuess,thrust,method)
                print(fnorm)
                # single shooting failed, exits out of everything
                if fnorm > 1e-7:
                    self.epsilon = e_best
                    return s_best, t_best, e_best
                
                # single shooting was successful!
                if incLogFlag:
                    # appends stateLog if the input was incomplete
                    stateLog.append(sLog)
                    timeLog.append(tLog)
                else:
                    # updates log of state and time results for next thrust iteration
                    stateLog[i] = sLog
                    timeLog[i]  = tLog
            
            # all thrusts were successful, save results
            e_best   = self.epsilon
            s_best   = deepcopy(sLog)
            t_best   = deepcopy(tLog)
        
        return s_best, t_best, e_best
    
# =============================================================================
#  Putting it al together
# =============================================================================
    def calculate_dMmap(self,TL,tA,dtRange,filename):
        """Calculates change in mass for transfers of various times and angular distances

        These are stored in a cache .dmmap file.

        Args:
            TL (TargetList module):
                Target list of stars
            tA (astropy Time array):
                Initial absolute mission time in MJD
            dtRange (astropy Time array):
                Transfer times to try
            filename (string):
                File name to store the cached data.
        """

        sInds       = np.arange(0,TL.nStars)
        ang         = self.star_angularSep(TL, 0, sInds, tA) 
        sInd_sorted = np.argsort(ang)
        angles      = ang[sInd_sorted].to('deg').value

        dtFlipped = np.flipud(dtRange)
        
        self.dMmap = np.zeros([len(dtRange) , len(angles)])
        self.eMap  = np.zeros([len(dtRange) , len(angles)])
        
        tic = time.perf_counter()
        for j,n in enumerate(sInd_sorted):
            for i,t in enumerate(dtFlipped):
                print(i,j)
                s_coll, t_coll, e_coll, TmaxRange = \
                            self.collocate_Trajectory(TL,0,n,tA,t)
                
                if e_coll != 0:
                    s_ssm, t_ssm, e_ssm = self.singleShoot_Trajectory(s_coll, \
                                                t_coll,e_coll,TmaxRange*u.N)

                if e_ssm == 2 and t.value < 30:
                    break
                
                m = s_ssm[-1][6,:] 
                dm = m[-1] - m[0]
                self.dMmap[i,j] = dm
                self.eMap[i,j]  = e_ssm
                toc = time.perf_counter()
                
                dmPath = os.path.join(self.cachedir, filename+'.dmmap')
                A = {'dMmap':self.dMmap,'eMap':self.eMap,'angles':angles,'dtRange':dtRange,'time':toc-tic,\
                     'tA':tA,'m0':1,'ra':TL.coords.ra,'dec':TL.coords.dec,'mass':self.mass}
                with open(dmPath, 'wb') as f:
                    pickle.dump(A, f)
                print('Mass - ',dm*self.mass)
                print('Best Epsilon - ',e_ssm)
    
    
    def calculate_dMmap_collocate(self,TL,tA,dtRange,filename):
        """Calculates change in mass for transfers of various times and angular distances

        These are stored in a cache .dmmap file. Only collocation without the single shooting.

        Args:
            TL (TargetList module):
                Target list of stars
            tA (astropy Time array):
                Initial absolute mission time in MJD
            dtRange (astropy Time array):
                Transfer times to try
            filename (string):
                File name to store the cached data.
        """
        
        sInds       = np.arange(0,TL.nStars)
        ang         = self.star_angularSep(TL, 0, sInds, tA) 
        sInd_sorted = np.argsort(ang)
        angles      = ang[sInd_sorted].to('deg').value

        dtFlipped = np.flipud(dtRange)
        
        self.dMmap = np.zeros([len(dtRange) , len(angles)])
        self.eMap  = np.zeros([len(dtRange) , len(angles)])
        
        tic = time.perf_counter()
        for j,n in enumerate(sInd_sorted):
            for i,t in enumerate(dtFlipped):
                print(i,j)
                s_coll, t_coll, e_coll, TmaxRange = \
                            self.collocate_Trajectory(TL,0,n,tA,t)

                if e_coll == 2 and t.value < 30:
                    break
                
                m = s_coll[-1][6,:]
                dm = m[-1] - m[0]
                self.dMmap[i,j] = dm
                self.eMap[i,j]  = e_coll
                toc = time.perf_counter()
                
                dmPath = os.path.join(self.cachedir, filename+'.dmmap')
                A = {'dMmap':self.dMmap,'eMap':self.eMap,'angles':angles,'dtRange':dtRange,'time':toc-tic,\
                     'tA':tA,'ra':TL.coords.ra,'dec':TL.coords.dec,'mass':self.mass}
                with open(dmPath, 'wb') as f:
                    pickle.dump(A, f)
                print('Mass - ',dm*self.mass)
                print('Best Epsilon - ',e_coll)
    
    
    def calculate_dMmap_collocateEnergy(self,TL,tA,dtRange,filename,m0=1,seed=000000000):
        """Calculates change in mass for transfers of various times and angular distances

        These are stored in a cache .dmmap file. Only minimum energy collocation is used.

        Args:
            TL (TargetList module):
                Target list of stars
            tA (astropy Time array):
                Initial absolute mission time in MJD
            dtRange (astropy Time array):
                Transfer times to try
            filename (string):
                File name to store the cached data.
            m0 (float):
                Initial mass
            seed (int):
                Seed random number for repeatability of experiments
        """
        
        sInds       = np.arange(0,TL.nStars)
        ang         = self.star_angularSep(TL, 0, sInds, tA) 
        sInd_sorted = np.argsort(ang)
        angles      = ang[sInd_sorted].to('deg').value
        
        dtFlipped = np.flipud(dtRange)
        
        self.dMmap = np.zeros([len(dtRange) , len(angles)])
        self.eMap  = 2*np.ones([len(dtRange) , len(angles)])
        
        tic = time.perf_counter()
        for j,n in enumerate(sInd_sorted):
            for i,t in enumerate(dtFlipped):
                print(i,j)
                s_coll, t_coll, e_coll, TmaxRange = \
                            self.collocate_Trajectory_minEnergy(TL,0,n,tA,t,m0)
                
                # if unsuccessful, reached min time -> move on to next star
                if e_coll == 2 and t.value < 30:
                    break

                m = s_coll[6,:] 
                dm = m[-1] - m[0]
                self.dMmap[i,j] = dm
                self.eMap[i,j]  = e_coll
                toc = time.perf_counter()
                
                dmPath = os.path.join(self.cachedir, filename+'.dmmap')
                A = {'dMmap':self.dMmap,'eMap':self.eMap,'angles':angles,'dtRange':dtRange,'time':toc-tic,\
                     'tA':tA,'m0':m0,'ra':TL.coords.ra,'dec':TL.coords.dec,'seed':seed,'mass':self.mass}
                with open(dmPath, 'wb') as f:
                    pickle.dump(A, f)
                print('Mass - ',dm*self.mass)
                print('Best Epsilon - ',e_coll)


    def calculate_dMsols_collocateEnergy(self,TL,tStart,tArange,dtRange,N,filename,m0=1,seed=000000000):
        """Calculates change in mass for transfers of various times and angular distances

        These are stored in a cache .dmmap file. Only minimum energy collocation is used.
        N random combinations of starting times, transfer times, and initial, and final stars are used.

        Args:
            TL (TargetList module):
                Target list of stars
            tStart (astropy Time array):
                Initial reference absolute mission time in MJD
            tArange (astropy Time array):
                Potential times to add to tStart
            dtRange (astropy Time array):
                Transfer times to try
            N (int):
                Number of trials/combinations to perform
            filename (string):
                File name to store the cached data.
            m0 (float):
                Initial mass
            seed (int):
                Seed random number for repeatability of experiments
        """
        
        self.dMmap = np.zeros(N)
        self.eMap  = 2*np.ones(N)
        iLog   = np.zeros(N)
        jLog   = np.zeros(N)
        dtLog  = np.zeros(N)
        tALog  = np.zeros(N)
        angLog = np.zeros(N)*u.deg

        tic = time.perf_counter()
        for n in range(N):
                print("---------\nIteration",n)

                i = np.random.randint(0,TL.nStars)
                j = np.random.randint(0,TL.nStars)
                dt = np.random.randint(0,len(dtRange))
                tA = np.random.randint(0,len(tArange))
                ang = self.star_angularSep(TL,i,j,tStart+tArange[tA]) 

                print("star pair  :",i,j)
                print("ang  :",ang.to('deg').value)
                print("dt   :",dtRange[dt].to('d').value)
                print("tau  :",tArange[tA].to('d').value,"\n")

                pair = np.array([i,j])
                iLog[n]   = i
                jLog[n]   = j
                dtLog[n]  = dt
                tALog[n]  = tA
                angLog[n] = ang

                s_coll, t_coll, e_coll, TmaxRange = \
                            self.collocate_Trajectory_minEnergy(TL,i,j,tStart+tArange[tA],dtRange[dt],m0)
                
                # if unsuccessful, reached min time -> move on to next star
                if e_coll == 2 and dtRange[dt].value < 30:
                    break

                m = s_coll[6,:] 
                dm = m[-1] - m[0]
                self.dMmap[n] = dm
                self.eMap[n]  = e_coll
                toc = time.perf_counter()
                
                dmPath = os.path.join(self.cachedir, filename+'.dmsols')
                A = {'dMmap':self.dMmap,'eMap':self.eMap,'angLog':angLog,'dtLog':dtLog,'time':toc-tic,\
                     'tArange':tArange,'dtRange':dtRange,'N':N,'tStart':tStart,\
                     'tALog':tALog,'m0':m0,'ra':TL.coords.ra,'dec':TL.coords.dec,'seed':seed,'mass':self.mass}
                with open(dmPath, 'wb') as f:
                    pickle.dump(A, f)
                print('Mass - ',dm*self.mass)
                print('Best Epsilon - ',e_coll)         
                
                

    def calculate_dMmap_collocateEnergy_angSepDist(self,TL,tA,dtRange,nPairs,filename,m0=1,seed=000000000):
        """Calculates change in mass for transfers of various times and angular distances

        These are stored in a cache .dmmap file. Only minimum energy collocation is used.
        nPair random combinations of initial and final star are used while all transfer times are used.

        Args:
            TL (TargetList module):
                Target list of stars
            tA (astropy Time array):
                Initial reference absolute mission time in MJD
            dtRange (astropy Time array):
                Transfer times to try
            nPairs (int):
                Number of trials/combinations to perform
            filename (string):
                File name to store the cached data.
            m0 (float):
                Initial mass
            seed (int):
                Seed random number for repeatability of experiments
        """
        
        iLog = np.zeros([len(dtRange) , nPairs])
        jLog = np.zeros([len(dtRange) , nPairs])
        psiLog = np.zeros([len(dtRange) , nPairs])
        
        dtFlipped = np.flipud(dtRange)
        
        self.dMmap = np.zeros([len(dtRange) , nPairs])
        self.eMap  = 2*np.ones([len(dtRange) , nPairs])
        
        tic = time.perf_counter()
        toc = time.perf_counter()
        for t,dt in enumerate(dtFlipped):

            iSelect,jSelect,psiSelect = self.selectPairsOfStars(TL,nPairs,tA,dt.value,int(1e6))
            sort = np.argsort(psiSelect)
            iSelect = iSelect[sort]
            jSelect = jSelect[sort]
            psiSelect = psiSelect[sort]            
            
            for n,(i,j) in enumerate(zip(iSelect,jSelect)):
                elapsedTime = (toc-tic)/3600
                totalTime = elapsedTime * len(dtRange) * nPairs / (1 + t*nPairs + n)
                print("dt :",dt," star #:",n," /",nPairs)
                print("Time Elapsed: ", elapsedTime, " hrs")
                print("Time Left: ", totalTime - elapsedTime, " hrs")
                s_coll, t_coll, e_coll, TmaxRange = \
                    self.collocate_Trajectory_minEnergy(TL,i,j,tA,dt,m0)
                
                m = s_coll[6,:] 
                dm = m[-1] - m[0]
                self.dMmap[t,n] = dm
                self.eMap[t,n]  = e_coll
                
                iLog[t,n] = int(i)
                jLog[t,n] = int(j)
                psiLog[t,n] = psiSelect[n]
                
                
                toc = time.perf_counter()
                
                dmPath = os.path.join(self.cachedir, filename+'.dmmap')
                A = {'dMmap':self.dMmap,'eMap':self.eMap,'angles':psiLog,'dtRange':dtRange,'time':toc-tic,\
                     'tA':tA,'m0':m0,'lon':TL.coords.lon,'lat':TL.coords.lat,'seed':seed,'mass':self.mass, \
                     'iLog':iLog,'jLog':jLog}
                
                with open(dmPath, 'wb') as f:
                    pickle.dump(A, f)
                print('---Mass - ',dm*self.mass)
                print('---Best Epsilon - ',e_coll) 

    def calculate_dMmap_collocateEnergy_LatLon(self,TL,tA,dtRange,nStars,filename,m0=1,seed=000000000):
        """Calculates change in mass for transfers of various times and angular distances

        These are stored in a cache .dmmap file. Only minimum energy collocation is used.
        All pairs between nStars random stars from the targetlist. All transfer times are used.

        Args:
            TL (TargetList module):
                Target list of stars
            tA (astropy Time array):
                Initial reference absolute mission time in MJD
            dtRange (astropy Time array):
                Transfer times to try
            nStars (int):
                Number of trials/combinations to perform
            filename (string):
                File name to store the cached data.
            m0 (float):
                Initial mass
            seed (int):
                Seed random number for repeatability of experiments
        """
        
        coords = TL.coords
        lon = coords.lon
        sInds = np.random.choice( TL.nStars , int(nStars) , replace=False  )
        sInds = sInds[np.argsort(lon[sInds])]
        dtFlipped = np.flipud(dtRange)
        
        self.dMmap = np.zeros([ len(dtRange) , len(sInds), len(sInds)])
        self.eMap  = 2*np.ones([len(dtRange) , len(sInds), len(sInds)])
        
        tic = time.perf_counter()
        toc = time.perf_counter()
        for i,ni in enumerate(sInds):
            for j,nj in enumerate(sInds):
                for n,t in enumerate(dtFlipped):
                    elapsedTime = (toc-tic)/3600
                    totalTime = elapsedTime * len(dtRange) * nStars**2 / (1 + t + j*len(dtRange) + i*nStars*len(dtRange))
                    print("dt :",t," star #:",i,"-->",j)
                    print("Time Elapsed: ", elapsedTime, " hrs")
                    print("Time Left: ", totalTime - elapsedTime, " hrs")
                    print(i,j,t.value)
                    s_coll, t_coll, e_coll, TmaxRange = \
                                self.collocate_Trajectory_minEnergy(TL,ni,nj,tA,t,m0)
                    
                    # if unsuccessful, reached min time -> move on to next star
                    if e_coll == 2 and t.value < 30:
                        break
    
                    m = s_coll[6,:] 
                    dm = m[-1] - m[0]
                    self.dMmap[n,i,j] = dm
                    self.eMap[n,i,j]  = e_coll
                    toc = time.perf_counter()
                    
                    dmPath = os.path.join(self.cachedir, filename+'.dmmap')
                    A = {'dMmap':self.dMmap,'eMap':self.eMap,'dtRange':dtRange,'time':toc-tic,\
                         'tA':tA,'m0':m0,'lon':TL.coords.lon,'lat':TL.coords.lat,'sInds':sInds,'seed':seed,'mass':self.mass}
                    with open(dmPath, 'wb') as f:
                        pickle.dump(A, f)
                    print('Mass - ',dm*self.mass)
                    print('Best Epsilon - ',e_coll)
StarShade Observatory class     This class is implemented at L2 and contains all variables, functions,      and integrators to calculate occulter dynamics. #to convert from dimensionalized to normalized, (Dimension) / self.(Dimension)U # smoothing factor (eps = 1 means min energy, eps = 0 means min fuel) # ============================================================================= # Miscellaneous # ============================================================================= Direction cosine matrix to rotate from Rotating Frame to Inertial Frame                  Finds rotation matrix for positions and velocities (6x6)          Args:             t (float):                 Rotation angle          Returns:             float 6x6 array:                 rotation of full 6 dimensional state from R to I frame Direction cosine matrix to rotate from Inertial Frame to Rotating Frame                  Finds rotation matrix for positions and velocities (6x6)          Args:             t (float):                 Rotation angle          Returns:             float 6x6 array:                 rotation of full 6 dimensional state from I to R frame Direction cosine matrix to rotate from Rotating Frame to Inertial Frame          Finds rotation matrix for positions, velocities, and accelerations (9x9)          Args:             t (float):                 Rotation angle          Returns:             float 9x9 array:                 rotation of full 9 dimensional state from R to I frame Generate a random lagrange multiplier for initial guess (6x1)          Generates an array of 6 numbers with two pairs of 3 coordinates describing position on a sphere.         The first two represent the x and y positions on a sphere with random radius between 1 and 5,         and the last coordinate represents the z coordinate scaled by the radius.          Returns:             float 6 array:                 Random lagrange multipliers for an initial guess #TODO: seems odd that 3rd and 6th coordinate aren't part of a spherical coordinate Finds angular separation from old star to given list of stars                  This method returns the angular separation from the last observed          star to all others on the given list at the currentTime. This is distinct         from the prototype version of the method in its signing of the         angular separation based on halo velocity direction                  Args:             TL (TargetList module):                 TargetList class object             old_sInd (integer):                 Integer index of the last star of interest             sInds (integer ndarray):                 Integer indices of the stars of interest             currentTime (astropy Time array):                 Current absolute mission time in MJD             dt (float):                 Timestep in units of days for determining halo velocity frame          Returns:             float:                 Angular separation between two target stars # time in canonical units # halo positions and velocities # halo positions and velocities in canonical units # V-frame # Position of stars # unit vector to star locations # star locations to relative to telescope # unit vector from Telescope to star i # unit vector from Telescope to star j # direction of travel from star i to j # sign of the angular separation, + in direction of halo velocity (NOTE: only using halo velocity at t0, not t0+dt) # calculating angular separation and assigning signs Rejection sample from psiPop to achieve nSamples fitting logistic distribution          Args:             psiPop (array float):                 List of floats between -180 and 180             nSamples (int):                 Number of samples to return             angBinSize (float):                 Size of a bin for the histogram used in rejection sampling          Returns:             tuple:                 array float:                     Angles (nSamples) fitting the logistic distribution                 array int:                     An array of indices of original psiPop that were accepted # distribution of the full psi population #Desired Distribution pdf #Scaling factor # how many samples from the full data set are we taking? # randomly selecting from the population (no replacing, so unique samples) # an empty array to which we'll add accepted indices # running rejection sampling # figuring out which bin to put this particular ang sep in # calculating pdf ratio # random number from 0 to 1 # sample has been accepted! Rejection sampling of star pairings using desired distribution and sampling from that final distribution                  Args:             TL (TargetList module):                 TargetList class object             nPairs (int):                 Number of pairs to produce             currentTime (astropy Time array):                 Current absolute mission time in MJD             dt (float):                 Timestep in units of days for determining halo velocity frame             nSamples (int):                 Number of samples to be used when generating random stars for pairing         Returns:             tuple:                 int list:                     The list of indices corresponding to starting stars in TargetList                 int list:                     The list of indices corresponding to the ending stars in TargetList                 float array:                     Angular distance between pairs from iFinal and jFinal # randomly selected pairs of stars # angular separation between each of these pairs # star i -> at currentTime # star j -> at currentTime + dt # filtering pairings list, creating desired distribution using rejection sampling # sampling nPairs of stars from the final distribution # non-sampled star pairings # grab a star from each bin and add it to the list # results # add the sampled stars from our desired distributions # ============================================================================= # Helper functions # ============================================================================= Determines throttle based on instantaneous switching function value.                  Typically being used during collocation algorithms. A zero-crossing of         the switching function is highly unlikely between the large number of          nodes.          Args:             state (float array):                 State vector and lagrange multipliers         Returns:             float array:                 Throttle between 0 and 1, and in the same shape as state Evaluates the switching function at specific states.          Args:             state (float array):                 State vector and lagrange multipliers         Returns:             float:                 Value of the switching function Evaluates the time derivative of the switching function.                  Switching function derivative evaluated for specific states.          Args:             state (float array):                 State vector and lagrange multipliers         Returns:             float:                 Value of the switching function time derivative Selects the proper event function for integration.                  This method calculates the switching function and its derivative at a          single specific state. It then determines which thrust case it will be          in: full, medium, or no thrust. If the value of the switching function         is within a certain tolerance of the boundaries, it uses the derivative          to determine the direction it is heading in. Then the proper event          functions are created for the integrator to determine the next crossing         (i.e. the next case change).                   Args:             s0 (float array):                 Iniitial state vector and lagrange multipliers         Returns:             tuple:                 function list:                     A list of functions which are the switching function adjusted by 1e-10.                     These take in an unused first parameter, and then the state.                     Depending on the case (2), multiple functions are returned.                     These functions will return 0 when the switching function crosses                     1e-10 or -1e-10.                 int:                     An integer 0,1,2 representing whether the switching                     function of the initial state is within a 1e-10 tolerance of the                     of zero (2), less than -1e-10 (1), or greater than 1e-10 (0). # finding which case we are in: #   - case 2 if       -eps < S < eps       (medium thrust) #   - case 1 if   S < -eps                 (full thrust) #   - case 0 if                  eps < S   (no thrust) # checking to see if S is within a certain tolerance from epsilon # determine if there is a case error if within tolerance # not the minimum fuel case # at the upper bound, case determined by derivative # at the lower bound, case determined by derivative # minimum fuel case, only two cases # crossing upper epsilon from above # appending event function # crossing lower epsilon from below # appending event function # can either cross lower epsilon from above or upper from below # appending event function # ============================================================================= # Equations of Motion and Boundary Conditions # ============================================================================= Creates boundary conditions for solving a boundary value problem          Function used in scipy solve_bvp function call. Returns residuals          Args:             sA (6 or 7 array):                 To be compared with the initial state for boundary condition.             sB (6 or 7 array):                 To be compared with the final state for boundary condition.             constrained (boolean):                 Whether there are 6 (false) or 7 (true) boundary conditions.          Returns:             array:                 Returns the residuals of the boundary conditions/constraint equation. Equations of Motion with costate vectors          Equations of motion for CR3BP with costates for control.          Args:             t (float):                 Currently unused             state (array):                 The state and lagrange multipliers.             constrained (boolean):                 Currently unused             amax (float or boolean):                 Maximum acceleration attainable or False otherwise             integrate (boolean):                 If true, the array is flattened for integration.          Returns:             array:                 The time derivatives of the states and costates. # vector distances from primaries # norms of distances from primaries # Position-dependent acceleration terms #shape of 3xn # Position partial derivatives # shape of 3x3xn # Velocity-dependent acceleration terms #shape of 3xn # Velocity partial derivatives # Costate vectors # ================================================ # Equations of Motion # ================================================ # throttle factor # putting them all together, a 14xn array # putting them all together, a 12xn array Calculates starshade and telescope velocities in R- or I-frames during stationkeeping                  Calculates the rotating system position and velocity of the         starshade at insertion into the optimal initial state for         observation of the given star.          Args:             TL (TargetList module):                 Target List             sInd (int):                 Index of star for observation in the target list             currentTime (astropy Time array):                 Current absolute mission time in MJD             SRP (boolean):                 Whether or not solar radiation pressure should be used                  in calculating the insertion state for optimal                  starshade stationkeeping observation position.            Returns:             tuple:                 array:                     Position in the rotating frame                 array:                     Velocity in the rotating frame # absolute times (Note: equinox is start time of Halo AND when inertial frame and rotating frame match) #mission times relative to equinox ) #modTimes in canonical units # star a at t=ta # ============================================================================= # Initial conditions # ============================================================================= Finding initial guess for starting Thrust                  Args:             TL (TargetList module):                 Target list             nA (int):                 The index for the starting star in the target list             nB (int):                 The index for the final star in the target list             tA (astropy Time array):                 Initial absolute mission time in MJD             dt (astropy Time array);                 A time step between the two voundary conditions             m0 (float):                 Initial mass             s_init (array):                 An initial guess for the state          Returns:             tuple:                 astropy Newtons:                     The maximum initial thrust.                 array:                     States from the solved bvp.                 array:                     Times at corrsponding to each state. #angle,uA,uB,r_tscp = self.lookVectors(TL,nA,nB,tA,tB) #position vector of occulter in heliocentric frame #self_rA = uA*self.occulterSep.to('au').value + r_tscp[ 0] #self_rB = uB*self.occulterSep.to('au').value + r_tscp[-1] #running collocation Create grid of Tmax values using unconstrained thruster                  This method is used purely for creating figures.                   Args:             TL (TargetList module):                 Target list             tA (astropy Time array):                 Initial absolute mission time in MJD             dtRange (astropy Time array);                 An array of delta times to consider          Returns:             astropy Newtons array:                 Max thrust in Newtons (dtRange dimensions by TL.nStars) # ============================================================================= # BVP solvers # ============================================================================= Solving generic bvp from t0 to tF using states and costates          Uses the Scipy solve_bvp method.          Args:             sGuess (array):                 Initial state and costate guess             tGuess (astropy Time array):                 Times corresponding to the guess of the state at each time             aMax (astropy Newton or boolean):                 Maximum attainable acceleration             constrained (boolean):                 Flag for whether this is constrained or unconstrained problem                 This determines dimensions of the state inputs.             m0 (float):                 Initial mass             maxNodes (int):                 Maximum number of nodes to use in the BVP problem solution             verbose (boolean):                 Flag passed to solve_bvp         Returns:             tuple:                 array:                     State and costate at the various mesh times                 array:                     Corresponding times to the sampled states.                 boolean:                     Status returned by the bvp_solve method # unconstrained problem begins with 12 states, rather than 14. checking for that # if unconstrained is initial guess for constrained problem, make 14 state array # only saves initial and final desired states if first solving unconstrained problem # creating equations of motion and boundary conditions functions # solving BVP # saving results Solves minimum energy and minimum fuel cases for continuous thrust          Args:             TL (TargetList module):                 Target list             nA (int):                 The index for the starting star in the target list             nB (int):                 The index for the final star in the target list             tA (astropy Time array):                 Initial absolute mission time in MJD             dt (astropy Time array);                 A time step between the two voundary conditions                  Returns:             tuple:                 array:                     Trajectory                 array:                     Times corresponding to trajectory                 float:                     last epsilon that fully converged (2 if minimum energy didn't work)                     Parameterizes minimum energy to minimum fuel solution.                 astropy Newton array:                     Range of thrusts (Newtons) considered. # initializing arrays # solving using unconstrained thruster as initial guess # saving results # all thrusts were successful # thrust values # range of thrusts to try # range of epsilon values to try (e=1 is minimum energy, e=0 is minimum fuel) # Huge loop over all thrust and epsilon values: #   we start at the minimum energy case, e=1, using the thrust value from the unconstrained solution #     In/De-crement the thrust until we reach the desired thrust level #   then we decrease e and repeat process until we get to e=0 (minimum fuel) #   saves the last successful result in case collocation fails # loop over epsilon starting with e=1 # initialize epsilon # loop over thrust values from current to desired thrusts #print("Thrust #",i," / ",len(TmaxRange)) # convert thrust to canonical acceleration # retrieve state and time initial guesses # perform collocation # collocation failed, exits out of everything # if only the unconstrained problem worked, still returns a 14 length array # collocation was successful! # creates log of state and time results for next thrust iteration (at the beginning of the loop) # updates log of state and time results for next thrust iteration # all thrusts were successful, save results Solves minimum energy and minimum fuel cases for continuous thrust                  Args:             TL (TargetList module):                 Target list             nA (int):                 The index for the starting star in the target list             nB (int):                 The index for the final star in the target list             tA (astropy Time array):                 Initial absolute mission time in MJD             dt (astropy Time array);                 A time step between the two voundary conditions             m0 (float):                 Initial mass          Returns:             tuple:                 array:                     Trajectory                 array:                     Times corresponding to trajectory                 float:                     last epsilon that fully converged (2 if minimum energy didn't work)                     Parameterizes minimum energy to minimum fuel solution.                 astropy Newton array:                     Range of thrusts (Newtons) considered. # initializing arrays # solving using unconstrained thruster as initial guess # saving results # all thrusts were successful # thrust values # range of thrusts to try # Huge loop over all thrust and epsilon values: #   we start at the minimum energy case, e=1, using the thrust value from the unconstrained solution #     In/De-crement the thrust until we reach the desired thrust level #   then we decrease e and repeat process until we get to e=0 (minimum fuel) #   saves the last successful result in case collocation fails # loop over epsilon starting with e=1 # loop over thrust values from current to desired thrusts #print("Thrust #",i," / ",len(TmaxRange)) # convert thrust to canonical acceleration # retrieve state and time initial guesses # perform collocation # collocation failed, exits out of everything # creates log of state and time results for next thrust iteration (at the beginning of the loop) # all thrusts were successful, save results # ============================================================================= # Shooting Algorithms # ============================================================================= Integrates thruster trajectory with thrust case switches                  This methods integrates an initial guess for the spacecraft state          forwards in time. It uses event functions to find the next zero of the         switching function which means a new thrust case is needed (full,          medium or no thrust).          Args:             sGuess (array):                 Initial state and costate guess             tGuess (astropy Time array):                 Times corresponding to the guess of the state at each time             Tmax (astropy force):                 Maximum thrust attainable          Returns:             tuple:                 array:                     Trajectory states                 astropy Time array:                     Times corresponding to states # initializing starting time # converting Tmax to a canonical acceleration # establishing equations of motion # starting integration # selecting the switch functions with correct boundaries # running integration with event functions # saving final integration time, if greater than tF, we completed the trajectory # saving results in a new array #adding to the results log if there was a previous thrust case switch Objective Function for single shooting thruster                  Args:             w (array):                 Costate vector             t0 (astropy Time object):                 Initial time             tf (astropy Time object):                 Final time             Tmax (astropy force):                 Maximum thrust attainable             returnLog (boolean):                 Return the states and times of the solution         Returns:             tuple:                 float:                     Norm of difference between current state and boundary value                 array:                     Trajectory states                 astropy Time array:                     Times corresponding to states Minimizes boundary conditions for thruster          Args:             sGuess (array):                 Initial state and costate guess             tGuess (astropy Time array):                 Times corresponding to the guess of the state at each time             Tmax (astropy force):                 Maximum thrust attainable             method (string):                 Optimization method for Scipy minimize call         Returns:             tuple:                 float:                     Norm of difference between current state and boundary value                 array:                     Trajectory states                 astropy Time array:                     Times corresponding to states #        minimizer_kwargs = {"method":method,"args":(t0,tF,Tmax,)} #        res = optimize.basinhopping(self.conFun_singleShoot,w0,minimizer_kwargs=minimizer_kwargs) # Perform single shooting to solve the boundary value problem.          Args:             stateLog (array):                 Approximate trajectory typically determined by collocation             timeLog (astropy Time array):                 Corresponging time values for the trajectory             e_best (float):                 Epsilon value corresponding to previous approxmiate trajectory             TmaxRange (astropy Newton array):                 Range of thrusts (Newtons) considered.             method (string):                 Optimization method for Scipy minimize call          Returns:             tuple:                 array:                     Trajectory states                 astropy Time array:                     Times corresponding to states                 float:                     Epsilon value determining how fuel vs energy optimal the trajectory is. # initializing arrays # range of epsilon values to try (e=1 is minimum energy, e=0 is minimum fuel) # Huge loop over all thrust and epsilon values, just like collocation method # loop over epsilon starting with e=1 # initialize epsilon # loop over thrust values from current to desired thrusts #print("Thrust #",i," / ",len(TmaxRange)) # retrieve state and time initial guesses # perform single shooting # single shooting failed, exits out of everything # single shooting was successful! # appends stateLog if the input was incomplete # updates log of state and time results for next thrust iteration # all thrusts were successful, save results # ============================================================================= #  Putting it al together # ============================================================================= Calculates change in mass for transfers of various times and angular distances          These are stored in a cache .dmmap file.          Args:             TL (TargetList module):                 Target list of stars             tA (astropy Time array):                 Initial absolute mission time in MJD             dtRange (astropy Time array):                 Transfer times to try             filename (string):                 File name to store the cached data. Calculates change in mass for transfers of various times and angular distances          These are stored in a cache .dmmap file. Only collocation without the single shooting.          Args:             TL (TargetList module):                 Target list of stars             tA (astropy Time array):                 Initial absolute mission time in MJD             dtRange (astropy Time array):                 Transfer times to try             filename (string):                 File name to store the cached data. Calculates change in mass for transfers of various times and angular distances          These are stored in a cache .dmmap file. Only minimum energy collocation is used.          Args:             TL (TargetList module):                 Target list of stars             tA (astropy Time array):                 Initial absolute mission time in MJD             dtRange (astropy Time array):                 Transfer times to try             filename (string):                 File name to store the cached data.             m0 (float):                 Initial mass             seed (int):                 Seed random number for repeatability of experiments # if unsuccessful, reached min time -> move on to next star Calculates change in mass for transfers of various times and angular distances          These are stored in a cache .dmmap file. Only minimum energy collocation is used.         N random combinations of starting times, transfer times, and initial, and final stars are used.          Args:             TL (TargetList module):                 Target list of stars             tStart (astropy Time array):                 Initial reference absolute mission time in MJD             tArange (astropy Time array):                 Potential times to add to tStart             dtRange (astropy Time array):                 Transfer times to try             N (int):                 Number of trials/combinations to perform             filename (string):                 File name to store the cached data.             m0 (float):                 Initial mass             seed (int):                 Seed random number for repeatability of experiments # if unsuccessful, reached min time -> move on to next star Calculates change in mass for transfers of various times and angular distances          These are stored in a cache .dmmap file. Only minimum energy collocation is used.         nPair random combinations of initial and final star are used while all transfer times are used.          Args:             TL (TargetList module):                 Target list of stars             tA (astropy Time array):                 Initial reference absolute mission time in MJD             dtRange (astropy Time array):                 Transfer times to try             nPairs (int):                 Number of trials/combinations to perform             filename (string):                 File name to store the cached data.             m0 (float):                 Initial mass             seed (int):                 Seed random number for repeatability of experiments #:",n," /",nPairs) Calculates change in mass for transfers of various times and angular distances          These are stored in a cache .dmmap file. Only minimum energy collocation is used.         All pairs between nStars random stars from the targetlist. All transfer times are used.          Args:             TL (TargetList module):                 Target list of stars             tA (astropy Time array):                 Initial reference absolute mission time in MJD             dtRange (astropy Time array):                 Transfer times to try             nStars (int):                 Number of trials/combinations to perform             filename (string):                 File name to store the cached data.             m0 (float):                 Initial mass             seed (int):                 Seed random number for repeatability of experiments #:",i,"-->",j) # if unsuccessful, reached min time -> move on to next star