#PROGRAM "neograph.py"

import openpyxl
import math
import numpy as np
import matplotlib.pyplot as plt

file_path = input("Enter the full path to the Excel file (e.g., C:\\path\\to\\file.xlsx): ")

try:
    # Open the Excel file
    wb = openpyxl.load_workbook(file_path)
    sheet = wb.active
    print(f"Successfully opened {file_path}")
except FileNotFoundError:
    print(f"Error: The file {file_path} was not found. Please check the path and try again.")
    exit()

event_no = int(input("Enter the event number to use: "))
row = event_no + 1

# Constants
density = sheet[f"R{row}"].value/1000     # g/cm^3
mass = sheet[f"C{row}"].value*1000    # g 
volume = mass/density   #cm^3
radius = sheet[f"T{row}"].value*100  # cm


Cd = 2.0          # Drag coefficient
A = math.pi*radius**2     # Cross-sectional area (cm^2)
gamma = 1       # Heat transfer efficiency
Hvap = sheet[f"S{row}"].value        # Heat of vaporization
g = 981         # Gravitational acceleration (cm/s^2)
rho_atm = 0.001295     # Atmospheric density (g/cm^3) at 140km
theta = sheet[f"J{row}"].value *math.pi/180 # Entry angle in radians

k = 1.38e-16    # Boltzmann constant (erg/K in CGS)
T = 288      # Atmospheric temperature (K)
mm = 4.81e-23      # Mean molecular mass (g)
H = k*T / mm*g  # Scale height (cm)

# Initial conditions
V0 = sheet[f"E{row}"].value*100000       # Initial velocity (cm/s)
m0 = mass      # Initial mass (g)
dt = 0.01          # Time step (s)
t_end = 50

# Arrays for time, velocity, and mass
time = np.arange(0, t_end + dt, dt)
velocity = np.zeros_like(time)
mass = np.zeros_like(time)
height = np.zeros_like(time)

# Initial values
velocity[0] = V0
mass[0] = m0
height[0] = 14000000 # cm

def dVdt(V,m,h):
    rho = rho_atm * np.exp(-h / H)      # atmospheric density variation
    return (-0.5 * Cd * rho * V**2 * A/m) + (g * np.sin(theta)/m)

def dmdt(V,h):
    rho = rho_atm * np.exp(-h / H)      # atmospheric density variation
    return (0.413 * A * gamma * rho * V**3) / Hvap

# Numerical integration (Euler's method)
for i in range(1,len(time)):
    velocity[i] = velocity[i-1] + dVdt(velocity[i-1], mass[i-1], height[i-1]) * dt
    height[i] = height[i-1] - velocity[i-1] * np.sin(theta) * dt
    mass[i] = mass[i-1] + dmdt(velocity[i-1], height[i-1]) * dt

    if height [i] <= 0:
        last_value = i
        break #stop the for loop

# Set unused indices to NaN
velocity[last_value+1:] = np.nan
height[last_value+1:] = np.nan
mass[last_value+1:] = np.nan

# Plot velocity over time
plt.figure(figsize=(10, 6))
plt.plot(time, velocity/100000, label="Velocity (km/s)", color="blue")
plt.xlabel("Time (s)")
plt.ylabel("Velocity (km/s)")
plt.title("Velocity vs. Time")
plt.grid()
plt.legend()
plt.show()

# Plot mass over time
plt.figure(figsize=(10, 6))
plt.plot(time, mass/1000, label="Mass Loss (kg)", color="red")
plt.xlabel("Time (s)")
plt.ylabel("Mass Loss (kg)")
plt.title("Mass Loss vs. Time")
plt.grid()
plt.legend()
plt.show()

# Plot height over time
plt.figure(figsize=(10, 6))
plt.plot(time, height/100000, label="height (km)", color="red")
plt.xlabel("Time (s)")
plt.ylabel("Height (km)")
plt.title("Height vs. Time")
plt.grid()
plt.legend()
plt.show()