#!/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  = 1

  #Set the time increment for the propogation
  dt       = double(0.01)       #Time difference in seconds
  t_final  = 1*400.                 #Final time in seconds
  NumTimes  = t_final/dt
  TimeInds  = arange(NumTimes)
  RealTimes = TimeInds*dt
  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)

  #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)
 
  #Set the initial time to t=0
  Time=0
 
  #Start off with a made up Moment of Inertia matrix
  I_M = mat3(0.4 , 0.0 , 0.0, \
             0.0 , 0.1 , 0.0, \
             0.0 , 0.0 , 0.1)

  #Initialize the Direction Cosine Matrix
  DCM = mat3(1.0 , 0.0 , 0.0, \
             0.0 , 1.0 , 0.0, \
             0.0 , 0.0 , 1.0)


  #Initialize the Rotation Matrix
  R   = mat3(1.0 , 0.0 , 0.0, \
             0.0 , 1.0 , 0.0, \
             0.0 , 0.0 , 1.0)

  #Initialize the quaternion
  q_t0 = vec4(0,0,0,1)
  q_t1 = vec4(0,0,0,1)
 
  #Set the initial angular velocity (omega) and the omega tensor 
  omega    = vec3()     
  omega.x  = 0.001
  omega.y  = 0.0 
  omega.z  = 0.001
  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]
  omega_tensor = mat3(+0       , +omega.z , -omega.y,\
                      -omega.z , +0       , -omega.x,\
                      +omega.y , +omega.x , +0)

  #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(0,1,0)
    BField_in_B_0 = DCM.inverse() * BField_in_N
    BField_in_B_1 = DCM.inverse() * BField_in_N
    B_x_B_0_b     = DCM.inverse()*BField_in_N
    B_x_B_1_b     = DCM.inverse()*BField_in_N
    d_B_x_B_dt    = vec3(0,0,0)
    K_B           = 10.
    K_T           = 1
    omega.x = 0.0 ; omega.y = 0.0 ; omega.z = 0.0

  #Set up the angular velocity tensor
  omega_0 = omega_t0 = omega_t1 = omega
  omega_tensor = mat3(+0       , +omega.z , -omega.y,\
                    -omega.z , +0       , -omega.x,\
                    +omega.y , +omega.x , +0)

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

    #Inrement the time 
    Time += dt

    #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

    #Allow for alignment with the magnetic field
    if AlignWithMagField == 1:

      #Determine the magnetic field in the body frame 
      BField_in_B_1 = DCM.inverse() * BField_in_N
  
      BField_in_B_1 = BField_in_B_1.normalize()
      BField_in_B_0 = BField_in_B_0.normalize()
 
      #Calculate the cross product and the derivative
      B_x_B_1_b  = BField_in_B_1.cross(BField_in_B_0)  
      d_b_x_B_dt = (1/dt)*(1/BField_in_B_1.length()**2)*(B_x_B_1_b - B_x_B_0_b)
      print d_b_x_B_dt.x, d_b_x_B_dt.y, d_b_x_B_dt.z
      
      #Now solve for the damping torque
      m_damp   = K_B * d_b_x_B_dt
      #m_damp = vec3(0,0,0)
      m_damp.z = 0

      #Now solve for the alignment magnetic moment
      m_align   = - K_T * B_z_in_B.cross(BField_in_B_1) 
      #m_align   = -BField_in_B_1.cross(K_T * (B_z_in_B.cross(BField_in_B_1)))
      m_align.z = 0 
 
      #Update the magnetic field and cross product for the next time step
      BField_in_B_0 = BField_in_B_1
      B_x_B_0_b = B_x_B_1_b
      #print m_damp, m_align

      #Now calculate the torque
      #m_damp = 10e4*m_damp.cross(BField_in_B_1)
      #m = m_damp + m_align
      m = m_align+m_damp*10e3
      #print m_damp, m_align
      #torque = m.cross(BField_in_B_1)
      if (i > 2):
        torque = m*10e-4 
        #print torque
      #print m_damp.x, m_damp.y, m_damp.z
      #print m_align.x, m_align.y, m_align.z
      #pdb.set_trace()
      #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())
      #print B_z_in_B, BField_in_B_1.normalize()

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

    #And find the "time-averaged" angular velocity over the time interval**
    omega_avg = .5*(omega_t0+omega_t1) 

    #Now create the integrating factor matrix components
    P = dt*omega_avg.x ; Q = dt*omega_avg.y ; R = dt*omega_avg.z

    #And update the integrating factor matrix
    OMEGA[0,0] =  0    ; OMEGA[0,1] = +R/2 ; OMEGA[0,2] = -Q/2 ; OMEGA[0,3] = +P/2
    OMEGA[1,0] = -R/2  ; OMEGA[1,1] =  0   ; OMEGA[1,2] = +P/2 ; OMEGA[1,3] = +Q/2
    OMEGA[2,0] = +Q/2  ; OMEGA[2,1] = -P/2 ; OMEGA[2,2] =  0   ; OMEGA[2,3] = +R/2
    OMEGA[3,0] = -P/2  ; OMEGA[3,1] = -Q/2 ; OMEGA[3,2] = -R/2 ; OMEGA[3,3] =  0
    
    #Calculate the state transition matrix
    PSI = cos(dt*omega_avg.length()/2) * I_4 + \
          (2/(dt*omega_avg.length())) *sin(dt*omega_avg.length()/2) * OMEGA
    q_t1 = PSI*q_t0
    q_t1_quat = quat(q_t1[3],q_t1[0],q_t1[1],q_t1[2])
    q_t1_quat=q_t1_quat.normalize()
    q_t1 = vec4(q_t1_quat.x, q_t1_quat.y, q_t1_quat.z, q_t1_quat.w)
   
    #Update the angular velocity tensor 
    omega_tensor[0,1] = -omega_t1.z 
    omega_tensor[0,2] = +omega_t1.y
    omega_tensor[1,0] = +omega_t1.z 
    omega_tensor[1,2] = -omega_t1.x
    omega_tensor[2,0] = -omega_t1.y 
    omega_tensor[2,1] = +omega_t1.x 

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

    #Now update the Cosine Direction Matrix
    del_DCM = -omega_tensor * DCM
    DCM     = DCM + dt*del_DCM
    
    #Now update the Inertial frame description of the body axes 
    B_x_in_N   = DCM*B_z_in_B
    B_x_in_N_q = q_t1_quat.rotateVec(B_z_in_B)
    Arr1[i]=B_x_in_N.x
    Arr2[i]=B_x_in_N_q.x
    Arr3[i]=B_x_in_N.y
    Arr4[i]=B_x_in_N_q.y
    Arr5[i]=B_x_in_N.z
    Arr6[i]=B_x_in_N_q.z

    #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)')
if ApplyControlTorque == 1:
  mplot.savefig('./Figures/AngularVelocites_ControlLoopApplied.png')
else:
  mplot.savefig('./Figures/AngularVelocites_NoControlLoop.png')

