import numpy as np
import matplotlib.pyplot as plt
from scipy import interpolate
import xlrd
from math import pi

## Importing OpenRocket Data From Excel
#This new simulation will import all the variable values generated by OpenRocket and then will take those values directly from excel. The data will be converted into arrays and will be interpolated with the OpenRocket data time stamps. This will be used to create arrays of the data with 0.1 time stamps. This data will then be used to simulate the rocket.

##Getting the data from Excel
##Will make a new version with the data natively in this code rather than taking it from excel

wb = xlrd.open_workbook('LowerAlti_PythonSim.xlsx')
sh = wb.sheet_by_index(0)    

time_raw = np.zeros(560)
for i in range(1,561):
    cell_value = sh.cell(i,0).value
    time_raw[i-1] = cell_value

time = np.linspace(0,15,151)

thrust_raw = np.zeros(560)
for i in range(1,561):
    cell_value = sh.cell(i,28).value
    thrust_raw[i-1] = cell_value

thrust_function = interpolate.interp1d(time_raw,thrust_raw,kind="linear") 
thrust = thrust_function(np.linspace(0,15,151))

mass_raw = np.zeros(560)
for i in range(1,561):
    cell_value = sh.cell(i,19).value
    mass_raw[i-1] = cell_value

mass_function = interpolate.interp1d(time_raw,mass_raw,kind="linear") 
mass = mass_function(np.linspace(0,15,151))

roll_mi_raw = np.zeros(560) ##Roll moment of inertia
for i in range(1,561):
    cell_value = sh.cell(i,22).value
    roll_mi_raw[i-1] = cell_value

roll_mi_function = interpolate.interp1d(time_raw,roll_mi_raw,kind="linear") 
roll_mi = roll_mi_function(np.linspace(0,15,151))


cd_raw = np.zeros(560) ##Drag coefficient of just rocket and no flaps
for i in range(1,561):
    cell_value = sh.cell(i,19).value
    cd_raw[i-1] = cell_value

cd_function = interpolate.interp1d(time_raw,cd_raw,kind="linear") 
cd = cd_function(np.linspace(0,15,151))

cd_raw = np.zeros(560) ##Drag coefficient of just rocket and no flaps
for i in range(1,561):
    cell_value = sh.cell(i,30).value
    cd_raw[i-1] = cell_value

cd_function = interpolate.interp1d(time_raw,cd_raw,kind="linear") 
cd = cd_function(np.linspace(0,15,151))

### Other Setup
#This setup involves creating empty arrays for data that needs to be seen in graphical format. Along with that, some other constants such as structural areas, coefficients and control variables are given.

v = np.zeros( 151 ) ##Velocity empty array
h = np.zeros( 151 ) ##Altitude empty array
D = np.zeros( 151 ) ##Drag force empty array
D_flaps = np.zeros( 151 ) ##Drag force of flaps empty array
min_vel = np.zeros( 151 ) ##Min_vel empty array for active drag code
a = np.zeros( 151 ) ##Empty array for total acc. of the rocket
ang_a = np.zeros( 151 ) ##Empty array for angular acc. of the rocket, rad/s^2
w = np.zeros( 151 ) ##Empty array for angular velocity of the rocket, rad/s
theta = np.zeros( 151 ) #Roll angle of the rocket, rad
alpha = np.zeros( 151 ) #Flap degree during roll control
flap = np.zeros( 151 ) #Flap status

A = 0.008107 # m^2 rocket reference area
A_flaps = 0.00258 * 2 #Area of 2 flaps
moment_arm = 0.0762 #3 inches in meters
g = 9.8 # m/s^2
rho = 1.225 #kg/m^3, density of air
C_flaps = 1.28 #approx drag coeff of flap
ts = 0.1 # time step for discrete physics calculation
flaps = 0
des_h = 714.756 #final height 2345 ft
buffer = 1 #Buffer percent
init_w = 3 #rad/s
w[0] = init_w
roll_threshold = 0.5 #rad/s
Kp = 0.003245 #Roll controller proportional gain derived from Simulink Data
roll_cutoff_alt = 550# Some altitude after which roll control is stopped for active drag priority

## Main Simulation
#Using the data setup above, we can simulate the whole flight of the rocket from launch to apogee, while also including the roll control and active drag effects.

##Main Simulation for loop runs for all time steps

for i in range(1,151):
    D[i] = 1/2*rho*(v[i-1])**2*cd[i]*A #drag force
    
    #Conditionals to determine if the flaps should be open for active drag or not
    #flaps = 0 indicates vertical flaps
    #flaps = 1 indicates horizontal flaps
    #flaps = 2 indicates angled flaps for roll control
    if thrust[i-1] != 0: #Vertical flaps during engine burn
        flap[i] = 0
    elif (w[i-1] > roll_threshold or w[i-1] < -roll_threshold) and h[i-1] < roll_cutoff_alt:
        flap[i] = 2
        alpha[i] = -Kp*w[i-1] #Proportional controller for the angle of the flaps
    elif h[i-1] > des_h: #Horizontal flaps after crosssing altitude target
        flap[i] = 1
    else:
        min_vel[i]  = np.sqrt(2*(g+D[i]/mass[i])*(des_h - h[i-1]))
        if v[i-1] > min_vel[i]*(1 + buffer/100):
            flap[i] = 1
        else:
            flap[i] = 0
    
    if flap[i] == 1:
        D_flaps[i] = 1/2*rho* (v[i-1])**2 * C_flaps * A_flaps #drag force due to flaps
    elif flap[i] == 2: #Determining the effect of the roll control flaps
        Cl = 2*pi*alpha[i]
        Lift = 1/2*rho*(v[i-1])**2*A_flaps/2*Cl
        Torque = Lift*moment_arm*2
        ang_a[i] = Torque/roll_mi[i]
    elif flap[i] == 0:
        D_flaps[i] = 0 #Flap Drag force is zero if the flaps are vertical 
        
    a[i] = -g + thrust[i]/mass[i] -D[i]/mass[i] -D_flaps[i]/mass[i] #Acceleration of the rocket at any given point
    v[i] = v[i-1] + (a[i])*ts #v = u + a*t
    h[i] = h[i-1] + v[i]*ts #displacement  = velocity*time
    
    w[i] = w[i-1] + (ang_a[i])*ts
    theta[i] = theta[i-1] + w[i]*ts 
    
        # check that the location isn't below the ground; if so, adjust as specified above
    if  h[i] <= 0:
        v[i] = 0
        h[i] = 0
    
print("Apogee:", max(h)*3.281, "feet")

#When using this in something like Jupyter notebeook, make sure you add graphs to this in separate cells to visualize your data.

'''
plt.plot(time,h*3.281)
plt.title('Altitude')
plt.ylabel('feet')

plt.plot( time,v, 'r-')
plt.plot( time,min_vel, 'b-')
plt.title('Drag System')
plt.legend(('Current Velocity','Control Velocity'))
plt.ylabel('m/s')

plt.plot(time,thrust)
plt.plot(time,D_flaps, 'r-')
plt.plot(time,D, 'k-')
plt.title('Forces')
plt.legend(('Thrust Force','Flap Drag Force','Native Drag Force'))
plt.ylabel('N')

plt.plot(time,w)
plt.title('Roll Rate')

plt.plot(time,alpha*180/pi)
plt.title('Roll Control Flap Angle')

plt.plot(time,flap)
plt.title('Flap Status')
'''
