# Aero4200-Ass3
#   Cal Wing, Sem 2 2023

# Import System / Common Libs 
import os, time
import numpy as np
import pandas as pd
from tqdm import tqdm

from numpy import pi, sin, cos, tan

#Import MakeGraph, a custom graphing wrapper I developed, refer to it for documentation  
from makeGraph import makeGraph, pltKeyClose, UQ_COLOURS as UQC # Custom Graphing Lib

# Override Sin & Cos to use & return degrees
#def sin(angle): return np.sin(np.deg2rad(angle))
#def cos(angle): return np.cos(np.deg2rad(angle))

# Make sure the relevant folders folder exists
folders = ["./images"]
for folder in folders:
    if not os.path.isdir(folder): os.mkdir(folder)

# IMU Data Loading
# I map in to a heading to add units / make things make more sense
## The gyroscopic body angular rates from the IMU are given:
#  - WBE_1 (in rad/s) - the roll rate about the body-fixed x-axis 
#  - WBE_2 (in rad/s) - the pitch rate about the body-fixed y-axis  
#  - WBE_3 (in rad/s) - the yaw rate about the body-fixed z-axis
## Specific forces:
#  - FSP_X (in m/s2) - the specific force in the body-fixed x-direction
#  - FSP_Y (in m/s2) - the specific force in the body-fixed y-direction
#  - FSP_Z (in m/s2) - the specific force in the body-fixed z-direction 
IMU_TIME_HEADER = ["Time [s]"]
IMU_WBE_HEADERS = ["WBE_1 [rad/s]", "WBE_2 [rad/s]", "WBE_3 [rad/s]"]
IMU_FSP_HEADERS = ["FSP_X [m/s^2]", "FSP_Y [m/s^2]", "FSP_Z [m/s^2]"]
IMU_DATA_HEADER = IMU_TIME_HEADER + IMU_WBE_HEADERS + IMU_FSP_HEADERS
def importIMUData(mission, imu):
    # If IMU is not a string, then convert based on bool eval where "H" == True
    if type(imu) != str: imu = "H" if imu else "L"  

    data = pd.read_csv(
           f"./data/IMU_M{str(mission) + str(imu)}.txt", 
            header=None, skiprows=1,
            names=IMU_DATA_HEADER,
        )

    return data

m1_IMUData = importIMUData(1, 0), importIMUData(1, 1) #(L, H) Data
m2_IMUData = importIMUData(2, 0), importIMUData(2, 1)

INIT_EULER_ANGLES = (0, 0, 0)
def translate2NED(angles, euler_angles):
    phi, theta, psi = euler_angles
    p, q, r = angles
    
    transMat = np.array([ [1, sin(phi)*tan(theta), cos(phi)*tan(theta) ],
                          [0, cos(phi),            -sin(phi)           ], 
                          [0, sin(phi)/cos(theta), cos(phi)/cos(theta) ]
                        ])
    
    angleMat = np.array([ [p], 
                          [q], 
                          [r]
                        ])

    return np.matmul(transMat, angleMat)

def getNEDForces(NEDPos):
    phi, theta, psi = NEDPos

    forceMat = np.array([ [cos(psi)*cos(theta), cos(psi)*sin(theta)*sin(phi)-sin(psi)*cos(phi), cos(psi)*sin(theta)*cos(phi)+sin(psi)*sin(phi)],
                          [sin(psi)*cos(theta), sin(psi)*sin(theta)*sin(phi)+cos(psi)*cos(phi), sin(psi)*sin(theta)*cos(phi)-cos(psi)*sin(phi)],
                          [-sin(theta),         cos(theta)*sin(phi),                            cos(theta)*cos(phi)                           ]
                        ])
    
    return forceMat

if __name__ == '__main__':
    dataPoint = m1_IMUData[0][IMU_WBE_HEADERS].iloc[0]
    trans = translate2NED(dataPoint.values, INIT_EULER_ANGLES)
    forces = getNEDForces((trans[0][0], trans[1][0], trans[2][0]))
    
    print("Raw Data")
    print(dataPoint)
    
    print("\nTranslated Point")
    print(trans)

    print("\nForces")
    print(forces)

    input("Damn")

if __name__ == '__main__1':
    #This is an example of drawing 4 plots by generating them
    graphData = {
        "figTitle":  "Simple Plot",
        "figTitleFontSize": 16,
        "figSize": (8,8), #Yay America, this is in inches :/ # Note: cm = 1/2.54
        "xLabel": "x label", 
        "yLabel": "y label",
        "plotDim": (2,2),
        "subPlots":[]
    }

    #Create 4 identical plots with different names
    for i in range(4):
        newPlot = {
            "title": f"Graph {i+1}",
            "plots": [
                {"x":[0,1,2,3,4], "y":[0,1,2,3,4], "label":"Linear"},
                {"x":[0,1,2,3,4], "y":[5,5,5,5,5]},
                {"x":[4,3,2,1,0], "y":[4,3,2,1,0], "label":"Linear2"},
                {"x":0, "type":"axvLine", "label":"Red Vertical Line", "color":"red"},
                {"y":6, "type":"axhLine", "label":"Dashed Horizontal Line", "args":{"linestyle":"--"}},
                {"type":"scatter", "x":4, "y":4, "label":"A Random Point", "colour":"purple", "args":{"zorder":2}}
            ]
        }
        graphData["subPlots"].append(newPlot)
    
    makeGraph(graphData, figSavePath="./images/example.png")