import numpy as np
import matplotlib.pyplot as plt
import axial_drag
import atmosphere

WEIGHT = 200
DIAMETER = 0.15

ATTACK_MAX = np.radians(40)

MAX_TIME = 100
T_s = 0.1

ANGLES = np.arange(10, 60+1, 5)

g = 9.81

TARGET_ACCEL = 30*g


def new_attack_angle(normal_force, dynamic_pressure, area):
    C_N_alpha = 15
    C_N = normal_force/(dynamic_pressure*area)
    attack_angle = C_N_alpha / C_N
    return min(attack_angle, ATTACK_MAX)


def simulate(mass, position, velocity, acceleration, attack_angle):
    #altitude = position[1]
    #speed = np.linalg.norm(velocity, 2)

    #Mach = atmosphere.Mach(altitude, speed)
    #axial_drag = np.array([-axial_drag.CA(Mach), 0])
    #
    #dynamic_pressure = atmosphere.dynamic_pressure(altitude, velocity)
    #
    #elevation = np.atan(velocity[1]/velocity[0]) + attack_angle

    #gravity = mass*g*np.array([np.cos(-angle-np.pi/2), np.sin(-angle-np.pi/2)])
    #
    #normal_force = TARGET_ACCEL - np.linalg.norm(axial_drag + gravity, 2)
    
    position += velocity*T_s + 0.5*acceleration*T_s**2
    velocity += acceleration*T_s    

    return position, velocity, acceleration


def main():
    t = np.arange(0, MAX_TIME, T_s)

    speed = 600
    accel = np.array([0, -g])
    mass = 200

    attack_angle = np.radians(40)
    
    step = 1
    t = np.arange(0, 100, T_s)
    pos = np.zeros((len(t), 2))
    vel = np.zeros_like(pos)
    vel[0] = speed*np.array([np.cos(np.pi/4), np.sin(np.pi/4)])

    while step < len(pos) and pos[step-1][1] >= 0:
        vel[step] = vel[step-1] + accel*T_s
        pos[step] = pos[step-1] + vel[step]*T_s + 0.5*accel*T_s**2
        step += 1

    plt.plot(t, pos[:,1])
    plt.show()

if __name__ == "__main__":
    main()