#!/usr/stsci/pyssg/Python-2.5.4/bin/python
#
#NAME:
#  SatelliteControl.py 
#
#PURPOSE:
#  To simulate the control system for the CXBN Nanosat.  This involves 
#  propogating the orbit of the center of mass of the system and 
#  modeling the rotational dynamics of the system. 
#
#KEYWORDS:
#  DateTime: str
#    The date and time for which we want to calculate the observation.  This should be 
#    in the pyephem prefered format of "Year/Month/Day Hour:Minute:Second" in UT.
#    The date and time for which we want to calculate the observation.  This should be 
#  TLE_File: str
#    The path to the TLE file that will be used as input for creating the PyEphem 
#    satellite object
#  ESolveMethod: int
#    0 - Solve Euler Equations using a first order Euler Method
#    1 - Solve Euler Equations using a fourth order Runge-Kutta method
#  TestGravityBoom: int
#    0 - Perform simulation with all input torques and controls
#    1 - Perform simulation of a gravity gradient stabilized satellite to 
#        demonstrate that the dynamics of the simulation are correct.  
#        For this case, the moment of inertia tensor will be highly concentrated 
#        along the z-axis. 
#  UseGGTorque: int
#    0 - Don't calculate and apply the gravitational gradient torque
#    1 - Calculate and apply the gravitational gradient torque
#  UseAeroTorque: int
#    0 - Don't calculate and apply the aerodynamic torque
#    1 - Calculate and apply the aerodynamic torque
#  UseMagTorque: int
#    0 - Don't calculate and apply the magnetic torque due to the magnetorquers on the CubeSat
#    1 - Calculate and apply the magnetic torque due to the magnetorquers on the CubeSat
#  TrackSun: int
#    0 - Don't track the sun's location
#    1 - Track the sun's location and attempt to steer the satellite toward it
#
#NOTES:
#  The pyephem package actually checks the CheckSum on each line of the TLE file.  One must make 
#  sure that the last digit of the line equals the sum of all digits (modulo 10).  One easy way 
#  to make this true is just to play with the last digit until it works.
#
#  There are three reference coordinate systems used here:
#    e = Earth Centered Inertial Coordinate System
#    r = Satellite Reference Frame (-z points toward center of earth, +x points in direction of velocity)
#    b = Body Coordinate System (axes point along principal axes)
#
#EXAMPLES:
#  GRAVITY BOOM-----
#  run SatelliteControl.py --TestGravityBoom=1 --TLE_File='./TLE/GravityGradientTest.tle'
#  MAGNETIC CONTROL DETUMBLE------
#  run SatelliteControl.py --UseMagTorque=1 --UseGGTorque=0 --UseAeroTorque=0 --TrackSun=0
#
###############################################################################################################

import os
execfile(os.path.expanduser('~')+'/Desktop/Software/Python/CZTCubeSat/CZTCubeSatSetup.py')

import ephem
from datetime import *
from CXBN_Utilities import *

# Main routine if we are called as a script
if __name__ == "__main__":

  ##PARSE KEYWORDS
  keywords    = ['DateTime=', 'TLE_File=', 'ESolveMethod=', 'TestGravityBoom=', \
                 'UseGGTorque=', 'UseAeroTorque=', 'UseMagTorque=', 'TrackSun='] 

  #DEFAULT VALUES FOR KEYWORDS
  DateTime        = '2012/01/01 20:20:00'
  TLE_File        = './TLE/CXBNTLE.txt'
  ESolveMethod    = 1
  TestGravityBoom = 0
  UseGGTorque     = 0
  UseAeroTorque   = 1
  UseMagTorque    = 1
  TrackSun        = 0

  #NOW PARSE KEYWORDS
  opts, extraparams = getopt.getopt(sys.argv[1:],'',keywords)
  for o,p in opts:
    if o in ['--DateTime']:
      DateTime = p
    elif o in ['--TLE_File']:
      TLE_File = p
    elif o in ['--ESolveMethod']:
      ESolveMethod = int(p)
    elif o in ['--TestGravityBoom']:
      TestGravityBoom = int(p)
    elif o in ['--UseGGTorque']:
      UseGGTorque = int(p)
    elif o in ['--UseAeroTorque']:
      UseAeroTorque = int(p)
    elif o in ['--UseMagTorque']:
      UseMagTorque = int(p)
    elif o in ['--TrackSun']:
      TrackSun = int(p)

  #############################################################################################################
  ################################ DATE AND TIME ##############################################################
  #Format the time appropriately
  Year  = int(DateTime[0:4])
  Month = int(DateTime[5:7])
  Day   = int(DateTime[8:10])
  Hours = int(DateTime[11:13])
  Mins  = int(DateTime[14:16])
  Secs  = int(DateTime[17:19])
  DateStr  = DateTime[0:10]
  DateTime = datetime(Year, Month, Day, Hours, Mins, Secs)

  #Convert this into a PyEphem Date (# days since noon on 1899 December 31) 
  EpDate_0 = ephem.date(DateTime)
  EpDate   = EpDate_0

  #############################################################################################################
  ################################ SATELITTE AND OTHER OBJECTS ################################################
  #Create an instance of the constant object
  c = c()

  #Call getsat to return an ephemeris satellite object based on the TLE
  CXBN = getsat(TLE_File, EpDate)

  #Create a satellite object
  cxbn = CubeSat()

  #Create an oject for the sun if we're trying to point to it
  if (TrackSun == 1):
    Sun = ephem.Sun()
    Sun.compute(EpDate)
    t_sun_interval = 50
    r_sun_eci      = vec3()
    r_sun_b        = vec3()
 
  #############################################################################################################
  ############################### INTIIAL POSITION AND VELOCITY ###############################################
  #Get the initial position and velocity vectors of the satellite to define the 
  #reference and body coordinate systems
  dt_sub = 0.00001
  CXBN.compute(EpDate-dt_sub*ephem.second) ; r_eci_t_minus_1 = getSatECI(CXBN,c)
  CXBN.compute(EpDate+dt_sub*ephem.second) ; r_eci_t_plus_1  = getSatECI(CXBN,c)
  CXBN.compute(EpDate)                     ; r_eci_t         = getSatECI(CXBN,c)
  v_t = (1/(2*dt_sub))*(r_eci_t_plus_1-r_eci_t_minus_1)
  
  #Now normalize these and find the cross product 
  v_t_norm = v_t.normalize()
  r_eci_norm = r_eci_t.normalize()
  r_v_cross_norm = r_eci_norm.cross(v_t_norm).normalize()
  
  #Get the initial position vector of the satellite
  CXBN.compute(EpDate)
  r_eci      = getSatECI(CXBN, c)
  r_eci_prev = r_eci

  #############################################################################################################
  ############################### TIME VECTORS ################################################################
  #Now calculate the approximate velocity vector at this time-step
  EpDate_t_minus_1 = EpDate

  #Set the time increment for the propogation
  dt       = double(0.01)       #Time difference in seconds
  t_final  = 1*6000.             #Final time in seconds
  NumTimes = t_final/dt
  Times    = arange(NumTimes)

  #Set the initial time to t=0
  Time=0

  #############################################################################################################
  ############################### TRACKING ARRAYS  ############################################################
  #Arrays for tracking variables (not necessary for functioning)
  w_x_Arr  = zeros(NumTimes, dtype=float32)
  w_y_Arr  = zeros(NumTimes, dtype=float32)
  w_z_Arr  = zeros(NumTimes, dtype=float32)
  t_x_Arr  = zeros(NumTimes, dtype=float32)
  t_y_Arr  = zeros(NumTimes, dtype=float32)
  t_z_Arr  = zeros(NumTimes, dtype=float32)
  #q_0_Arr  = zeros(NumTimes, dtype=float32)
  #q_1_Arr  = zeros(NumTimes, dtype=float32)
  #q_2_Arr  = zeros(NumTimes, dtype=float32)
  #q_3_Arr  = zeros(NumTimes, dtype=float32)
  #EleArr   = zeros(NumTimes, dtype=float32)
  #VelArr   = zeros(NumTimes, dtype=float32)
  #GTorArr  = zeros(NumTimes, dtype=float32)
  #ATorArr  = zeros(NumTimes, dtype=double)
  #w_t_Arr  = zeros(NumTimes, dtype=float32)
  #L_t_Arr  = zeros(NumTimes, dtype=float32)

  #############################################################################################################
  ############################### BASIS VECTORS AND COORDINATE SYSTEMS#########################################
  #Set up the basis vectors for the inertial frame
  i   = vec3(1,0,0)   ; j   = vec3(0,1,0)    ; k   = vec3(0,0,1)
  e_i = vec3(1,0,0)   ; e_j = vec3(0,1,0)    ; e_k = vec3(0,0,1)
  r_i = r_eci_norm    ; r_j = r_v_cross_norm ; r_k = v_t_norm
  b_i = r_eci_norm    ; b_j = r_v_cross_norm ; b_k = v_t_norm

  #Get the initial Direction Cosine Matrices
  e_2_b=ReturnDCM(e_i, e_j, e_k, b_i, b_j, b_k)
  r_2_b=ReturnDCM(r_i, r_j, r_k, b_i, b_j, b_k)

  #############################################################################################################
  ############################### ANGULAR VELOCITIES AND QUATERNIONS ##########################################
  #Set the initial angular velocity
  omega    = vec3()     #Set up an angular velocity vector)
  if (UseMagTorque == 1):
    omega.x  = 0.000                        #x component of angular velocity (w_1 = p)
    omega.y  = 0.000                        #y component of angular velocity (w_2 = q)
    omega.z  = 0.000                        #z component of angular velocity (w_3 = r)
  elif (TrackSun == 1):
    omega.x  = 0.000                        #x component of angular velocity (w_1 = p)
    omega.y  = 0.000                        #y component of angular velocity (w_2 = q)
    omega.z  = 0.000                        #z component of angular velocity (w_3 = r)

  else:
    omega.x  = 0.800                        #x component of angular velocity (w_1 = p)
    omega.y  = 0.400                        #y component of angular velocity (w_2 = q)
    omega.z  = -.900                        #z component of angular velocity (w_3 = r)
  
  #Now copy the angular velocity for other uses
  omega_0  = omega                          #This is the initial omega, does not change
  omega_t0 = omega                          #This is omega at timestep t[i]  
  omega_t1 = omega                          #This is omega at timestep t[i+1]

  #Get the initial quaternion
  quat_initial = quat(e_2_b)                #Get the DCM from the quaternion
  quat_initial = quat_initial.normalize()   #normalize the quaternion
  q_t0 = vec4()                             #quaternion at time step [t]
  q_t1 = vec4()                             #quaternion at time step [t+1]
  q_t0[3] = quat_initial.w                  
  q_t0[0] = quat_initial.x
  q_t0[1] = quat_initial.y
  q_t0[2] = quat_initial.z
  q_t1 = q_t0

  #############################################################################################################
  ############################### TORQUES AND MAGNETIC FIELD VECTORS ##########################################

  #Set up initial torque vectors
  aero_torque = vec3(0,0,0)                 #The 3 vector for the aerodynamic torque
  gg_torque   = vec3(0,0,0)                 #The 3 vector for the gravity gradient torque
  mag_torque  = vec3(0,0,0)                 #The 3 vector for the magnetic torque 

  #Set up arrays for the magnetic fields
  B_Earth_b   = vec3(0,0,0)                 #Earth's magnetic field in the body frame of the satellite
  k_loop      = 1.0                         #The proportional factor for the control loop
  m_mag       = 0.166                       #The magnetic moment magnitude in A*m^2 
  omega_des   = vec3(0,0,0)                 #The desired angular veloctiy in rad/s 

  #Keep track of the r_eci to do magnetic field calculations only every so often
  r_eci_mag_prev         = r_eci            #Vector to keep track of last spot where B field was calculated
  min_dist_for_mag_field = 32000          #Minimum distance to calculate new B field (in meters)

  #Get the initial date and orientation
  CXBN.compute(EpDate)
  r_eci = getSatECI(CXBN, c)

  #Allow for a test of the gravity boom
  if TestGravityBoom == 1:
    execfile('./Test/GravityBoomConfigurationForSatelliteControl.py')

  #############################################################################################################
  #PROPAGATE THE SATELLITE IN ITS ORBIT
  #############################################################################################################
  for t_i in Times:  

    #Inrement the time 
    EpDate = double(EpDate) + ephem.second*dt
    Time  += dt

    #Recompute the solution with the SGP4
    CXBN.compute(EpDate)

    #Get the position of the satellite in Earth-Centered Inertial Coordinates
    r_eci = getSatECI(CXBN, c) 
    v_eci = (1/dt)*(r_eci - r_eci_prev)
    r_eci_prev = r_eci
    #VelArr[i] = v_eci.length()

    #Now normalize these, find the cross product, and assign to the reference vectors 
    r_i = v_eci.normalize()
    r_k = r_eci.normalize()
    r_j = r_k.cross(r_i).normalize()

    #Now get the Torques, gg and aero
    if UseGGTorque == 1: 
      gg_torque = getGravityGradientTorque(r_eci, e_2_b, cxbn, c)
    if UseAeroTorque == 1:
      aero_torque = getAerodynamicTorque(v_eci, r_eci, e_2_b, cxbn, c)

    #Update the torque using the gravity gradient and aerodynamic.  Magnetic is below
    tot_torque  = gg_torque+aero_torque

    #Calculate the magnetic field if it's necessary
    if ((UseMagTorque == 1) or (TrackSun == 1)):
      #Only update the magnetic field if enough time/distance has passed
      if (((r_eci-r_eci_mag_prev).length() > min_dist_for_mag_field) or (t_i == 0)):
        LatDeg         = float(CXBN.sublat)*180/pi
        LonDeg         = float(CXBN.sublong)*180/pi
        AltitudeKm     = CXBN.elevation/1000.
        B_Earth_wmm    = getmagfield(LatDeg, LonDeg, AltitudeKm, DateStr)
        B_Earth_eci    = wmm2eci(float(CXBN.ra), float(CXBN.dec), B_Earth_wmm)
        B_Earth_b_tm0  = e_2_b * B_Earth_eci
        r_eci_mag_prev = r_eci

        #Update the magnetic field vectors stored for the contro loop 
        B_Earth_b_tm2   = B_Earth_b_tm1
        B_Earth_b_tm1   = B_Earth_b_tm0
        B_Earth_b_tm0   = B_Earth_b

    #Determine where the sun is in the body frame of the spacecraft
    if ((TrackSun == 1)) :#and (int(dt*t_i) % t_sun_interval == 0)):
      r_sun_eci.x = sin(Sun.dec)*cos(Sun.ra)      
      r_sun_eci.y = sin(Sun.dec)*sin(Sun.ra)
      r_sun_eci.z = cos(Sun.dec)
      r_sun_b     = e_2_b*r_sun_eci

      #Get the solar panel illumination for each quadrant
      SolarIlluminationArr = getSolarPanelIllumination(r_sun_b, cxbn)

      #Now try to orient toward the sun
      csun_torque = ReturnCoarseSunTorque(SolarIlluminationArr, B_Earth_b, omega_t0, \
                                          cxbn.InertiaTensor, k_loop, m_mag)
      tot_torque = tot_torque+csun_torque

    #Determine the magnetic field and apply an appropriate torque
    if UseMagTorque == 1:

      #Apply a control torque
      mag_torque = SimpleControlLoop(B_Earth_b, omega_t0, omega_des, cxbn.InertiaTensor, k_loop, m_mag) 
      tot_torque = tot_torque+mag_torque  

    #Update the angular velocity by Solving Euler's equations from [t] to [t+1]
    omega_t1 = SolveEuler(omega_t0, tot_torque, cxbn.InertiaTensor, dt, ESolveMethod)
    w_x_Arr[t_i] = omega_t1.x
    w_y_Arr[t_i] = omega_t1.y
    w_z_Arr[t_i] = omega_t1.z
    t_x_Arr[t_i] = mag_torque.x
    t_y_Arr[t_i] = mag_torque.y
    t_z_Arr[t_i] = mag_torque.z

    #Now get the quaternion at time [t+1]
    q_t1, q_t1_quat = UpdateQuaternion(q_t0, dt, omega_t0, omega_t1)
    #q_0_Arr[t_i] = q_t1[0]
    #q_1_Arr[t_i] = q_t1[1]
    #q_2_Arr[t_i] = q_t1[2]
    #q_3_Arr[t_i] = q_t1[3]

    #Get the updated earth to body rotation matrix and the basis vectors
    e_2_b = q_t1_quat.toMat3()
    b_2_e = e_2_b.transpose()
    b_i = b_2_e*i
    b_j = b_2_e*j
    b_k = b_2_e*k

    #Now get the Direction Cosine Matrix between the reference and body frames
    r_2_b=ReturnDCM(r_i, r_j, r_k, b_i, b_j, b_k)
    
    #And convert these to Euler Angles
    Yaw   = math.atan2(r_2_b[0,1], r_2_b[0,0])
    Pitch = -math.asin(r_2_b[0,2])
    Roll  = math.atan2(r_2_b[1,2], r_2_b[2,2]) 
    #w_x_Arr[t_i] = Yaw
    #w_y_Arr[t_i] = Pitch
    #w_z_Arr[t_i] = Roll

    #Figure out where we are in the orbit and where the sun is
    #print Time, EpDate, float(CXBN.elevation)
    #EleArr[i] = CXBN.elevation

    #Update omega(t+1) = omega(t) for the next time step
    omega_t0 = omega_t1
    q_t0     = q_t1

