#!/usr/stsci/pyssg/Python-2.5.4/bin/python
#
#NAME:
#  TestEuler.py 
#
#PURPOSE:
#  To test the Euler Solver and updating of the rotation matrix between 
#  a fixed inertial axis and the body axis of the rigid body under investigation. 
#  The nomenclature adopted is that 
#    B = Body Frame
#    N = Inertial Frame
#
#  The method used here follows the NASA Patent "Closed-Form Integrator For the 
#  Quaternion (Euler Angle) Kinematics Equations".  This technique allows us to 
#  integrate the quaternion such that the unit length is always preserved.
#  
#  **Although in the patent paper, they refer to this as a time average, it 
#    is really just a numerical estimate of the integral divided by 2.  
#
#TESTING:
#  I have verified the following cases (other omega components are zero)
#    omega.z = 2.0 -> B_x_in_N and B_y_in_N do circles ; B_z_in_N stays still
#    omega.y = 2.0 -> B_x_in_N and B_z_in_N do circles ; B_y_in_N stays still
#    omega.x = 2.0 -> B_y_in_N and B_z_in_N do circles ; B_x_in_N stays still
# 
#KEYWORDS: (actually hard coded below)
#  ESolveMethod: int
#    0 - Do Euler 
#    1 - Do RK4 time-stepping
#  DoTopPrecession: int
#    0 - Simply apply the initial conditions 
#    1 - Do the precession of a top (has analytical solution)
#  ApplyControlTorque: int
#    0 - Don't apply a control torque to stop the motion
#    1 - Apply a control torque to the object to stop the motion
#
##########################################################################################
import os
execfile(os.path.expanduser('~')+'/Desktop/Software/Python/CZTCubeSat/CZTCubeSatSetup.py')

from SimpleControlLoop import *
import ephem
import pdb
from datetime import *
from CXBN_Utilities import *

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

  #Determine the Euler Solve Method (0 = Euler ; 1 = 4th order RK)
  ESolveMethod       = 1
  DoTopPrecession    = 0
  ApplyControlTorque = 0
  AlignWithMagField  = 0
  AlignWithSun       = 1

  #############################################################################################################
  ############################### TIME VECTORS ################################################################
  #Set the time increment for the propogation
  dt        = double(0.001)         #Time difference in seconds
  t_final   = 200.                #Final time in seconds
  NumTimes  = t_final/dt
  TimeInds  = arange(NumTimes)
  RealTimes = TimeInds*dt

  #############################################################################################################
  ############################### 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)
  w_t_Arr   = zeros(NumTimes, dtype=float32)
  L_t_Arr   = zeros(NumTimes, dtype=float32)
  B_Ang_Arr = zeros(NumTimes, dtype=float32)
  S_Ang_Arr = zeros(NumTimes, dtype=float32)

  #Test arrays
  Arr1 = zeros(NumTimes, dtype=float32)
  Arr2 = zeros(NumTimes, dtype=float32)
  Arr3 = zeros(NumTimes, dtype=float32)
  Arr4 = zeros(NumTimes, dtype=float32)
  Arr5 = zeros(NumTimes, dtype=float32)
  Arr6 = zeros(NumTimes, dtype=float32)

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

  #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)

  #Set the initial time to t=0
  Time=0
 
  #Start off with a made up Moment of Inertia matrix
  I_M = mat3(0.10 , 0.00 , 0.00, \
             0.00 , 0.10 , 0.00, \
             0.00 , 0.00 , 0.01)

  #Make a CubeSat object
  cxbn = CubeSat()

  #############################################################################################################
  ############################### ANGULAR VELOCITIES AND QUATERNIONS ##########################################
  #Set the initial angular velocity
  #Expected rates for detumbling are 20 degrees/second
  omega    = vec3()     #Set up an angular velocity vector)
  if (AlignWithMagField == 1):
    omega.x  = 8.5*pi/180                   #x component of angular velocity (w_1 = p)
    omega.y  = 18.0*pi/180                  #y component of angular velocity (w_2 = q)
    omega.z  = 5.2*pi/180                   #z component of angular velocity (w_3 = r)
    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 (AlignWithSun == 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
 
  #Get the Integrating Factor Matrix (Captial Omega)
  OMEGA = mat4()
  I_4   = mat4(eye(4))
 
  #Keep the torque at zero for now 
  torque = vec3(0,0,0)

  #Set up the body axes unit vectors in the inertial frame
  B_x_in_B = vec3(1,0,0)
  B_y_in_B = vec3(0,1,0)
  B_z_in_B = vec3(0,0,1)

  #Compare to analytical solution of precession of top
  if DoTopPrecession:
    omega.x = 1.0 ; omega.y = 0 ; omega.z = 1.0
    PrecessionOmega = (I_M[2,2]-I_M[0,0])/I_M[1,1] * omega.z
    p_omega = vec3()
    PrecArr1 = zeros(NumTimes, dtype=float32)
    PrecArr2 = zeros(NumTimes, dtype=float32)
    PrecArr3 = zeros(NumTimes, dtype=float32)
    
  #Set up a set-torque in the case where a control torque is applied
  if ApplyControlTorque == 1:
    k = 2                  #The proportional constant for the control torque
    omega_d = vec3(0,0,0)  #The desired angular velocity vector
    t_max   = 1

  #Set up the magnetic field in case we want to try to align with it
  if AlignWithMagField == 1:
    BField_in_N   = vec3(1.0,1.0,0.0)
    BField_in_B_0 = e_2_b * BField_in_N
    BField_in_B_1 = e_2_b * BField_in_N
    K_B           = 2.0
    K_T           = 0.001
    m_mag         = 0.166
    omega.x = 0.2000 ; omega.y = 0.200 ; omega.z = 0.3000

  #Set up the sun vector
  if AlignWithSun:
    BField_in_N   = vec3(0.0,0.00005, 0.0001)
    k_cs_t = 0.002
    k_cs_w = -0.2
    r_sun_eci = vec3(1.0, 0.0, 0.1)
    m_mag  = 0.166

  #Set up the angular velocity tensor
  omega_0  = omega
  omega_t0 = omega
  omega_t1 = omega

  ########################################################################################
  #Do the progagation
  ########################################################################################
  #Do the progagation
  for i in TimeInds:  

    #Inrement the time 
    Time += dt

    ###########################################
    #SIMPLE CONTROL TORQUE 
    ###########################################
    #Allow for an application of control torque
    if ApplyControlTorque == 1:
      #Figure out the control torque
      torque.x, torque.y, torque.z = SimpleControlLoop(omega_t1, omega_d, I_M, k, t_max) 
      print torque

    ###########################################
    #CONTROL TORQUE WITH MAGNETORQUERS 
    ###########################################
    #Allow for alignment with the magnetic field
    #-------------------------------------------------------------
    #DO A CONTROL LOOP WITH THE SIMPLE CONTROL LAW FROM Sidi p. 192
    #-------------------------------------------------------------
    if (AlignWithMagField == 1):
   
      #Update the magnetic field in the body frame
      BField_in_B_1 = e_2_b * BField_in_N 
      MagMode = 1

      if (MagMode == 0): 

        #Assign the alignment torque
        m_align = K_T * B_z_in_B

        #And now calculate the damping torque  
        omega_d.x = omega_t0.x
        omega_d.y = omega_t0.y
        omega_d.z = 0
        m_damp  = (-K_B/(BField_in_B_1.length()**2)) * (BField_in_B_1.cross(omega_d))

        #Now calculate the torque
        m = m_damp + m_align
        if (i>0):
          torque = m.cross(BField_in_B_1)
          print torque, m_damp, m_align

      elif (MagMode == 1):

        #Compute the change in magnetic field 
        B_dot = (1/dt)*(1/(BField_in_B_1.length()**2))*(BField_in_B_1-BField_in_B_0)

        #Now figure out the 
        m_const = vec3(0,0,.5)
        m = -K_B * B_dot + m_const

        if (i>0):
          torque = m.cross(BField_in_B_1)
          #print torque, m, B_dot
          print omega_t0, omega_t1

        #Update the previous magnetic field
        BField_in_B_0 = BField_in_B_1

      #Get the angle between the magnetic field and the z-axis 
      B_Ang_Arr[i] = math.acos(B_z_in_B.normalize()*BField_in_B_1.normalize())

    ###########################################
    #Allow for alignment with the sun vector
    #-------------------------------------------------------------
    if (AlignWithSun) == 1:

      #Determine where the sun is in the body frame
      r_sun_b = e_2_b*r_sun_eci

      #Update the magnetic field in the body frame
      BField_in_B_1 = e_2_b * BField_in_N
      MagMode = 1

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

      print SolarIlluminationArr

      #Now try to orient toward the sun
      csun_torque = ReturnCoarseSunTorque(SolarIlluminationArr, BField_in_B_1, omega_t1, k_cs_t, k_cs_w, m_mag)

      torque = csun_torque

      #Get the angle between the z-axis and the sun pointing vector
      S_Ang_Arr[i] = math.acos(B_z_in_B.normalize()*r_sun_b.normalize())

    #Solve for the new angular velocity vector
    omega_t1 = SolveEuler(omega_t0, torque, I_M, dt, ESolveMethod)

    #Now get the quaternion at time [t+1]
    q_t1, q_t1_quat = UpdateQuaternion(q_t0, dt, omega_t0, omega_t1)

    #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*e_i
    b_j   = b_2_e*e_j
    b_k   = b_2_e*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)

    #Keep track of the precession if that's what we're modeling
    if DoTopPrecession:
      p_omega.x = omega_0.x * cos(PrecessionOmega*Time)
      p_omega.y = omega_0.x * sin(PrecessionOmega*Time)
      p_omega.z = omega_0.z 
      PrecArr1[i]   = p_omega.x
      PrecArr2[i]   = p_omega.y
      PrecArr3[i]   = p_omega.z

    #Add the components to the arrays
    w_x_Arr[i] = omega_t1.x
    w_y_Arr[i] = omega_t1.y
    w_z_Arr[i] = omega_t1.z
    w_t_Arr[i] = omega_t1.length()

    #Calculate the angular momentume
    L_t_Arr[i] = (I_M*omega_t1).length()

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

#Now plot the angular velocities
mplot.figure(0)
mplot.clf()
mplot.plot(TimeInds*dt, w_x_Arr, label=r'$\omega_{x}$')
mplot.plot(TimeInds*dt, w_y_Arr, label=r'$\omega_{y}$')
mplot.plot(TimeInds*dt, w_z_Arr, label=r'$\omega_{z}$')
mplot.xlabel('Time (seconds)')
mplot.ylabel('Angular Velocity (radians/sec)')
mplot.legend()
if ApplyControlTorque == 1:
  mplot.savefig('./Figures/AngularVelocites_ControlLoopApplied.png')
else:
  mplot.savefig('./Figures/AngularVelocites_NoControlLoop.png')

#Option to plot the separation from the magnetic field vector
if (AlignWithMagField == 1):
  mplot.figure(1)
  mplot.clf()
  mplot.plot(TimeInds*dt, B_Ang_Arr)
  mplot.xlabel('Time (seconds)')
  mplot.ylabel('Angle between z-axis and Magnetic Field Vector')

#Option to plot the separation from the sun pointing vector
if (AlignWithSun == 1):
  mplot.figure(1)
  mplot.clf()
  mplot.plot(TimeInds*dt, S_Ang_Arr)
  mplot.xlabel('Time (seconds)')
  mplot.ylabel('Angle between z-axis and Sun Vector')

