import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

# Constantes
m = 0.00026  # Masa do pendulo em gramas
L = 1.0  # Comprimento da corda em metros
# b = 0.01 # Coeficiente de arrasto aerodinâmico
p = 1.28 # Densidade do ar em kg/m^3
g = 9.81 # Aceleração da gravidade em m/s^2

# Coeficiente de arrasto aerodinâmico para uma esfera
cd = 0.01
# Raio do pêndulo
R = 3 # em metros

# Área frontal do pêndulo
A = 0.002826

# Coeficiente de arrasto aerodinâmico
b = 0.000180864

# Valores Iniciais
theta0 = 0.05  # Ângulo inicial em radianos
w0 = 0.0       # Velocidade angular inicial em radianos/s
t0 = 0.0       # Tempo inicial em segundos
h = 0.1        # Passo de tempo em segundos
# Listas para armazenar resultados
t = [t0]
theta = [theta0]
w = [w0]

# Loop ao longo do tempo
while t[-1] < 100.0:
    # Calcular valores intermediários
    k1x = h * w[-1]
    k1v = h * (-w[-1]*(b*L/m) *(theta[-1])**2)-g*theta[-1]
    k2x = h * (w[-1] + k1v)
    k2v = h*(-((w[-1] + k1v)*h)*(b*m)*((w[-1] + k1v)*h)**2 - g * ((theta[-1] + k1x)))

    # Atualize teta e w
    theta.append(theta[-1] + ((k1x + k2x) / 2.0)*h)  # Metodo Heun
    w.append(w[-1] + ((k1v + k2v) / 2.0)*h)        # Metodo Heun

    # Atualiza o Tempo
    t.append(t[-1] + h)

# Cria gráfico para teta com linha vermelha
plt.plot(t, theta, color='red')

# Adiciona uma linha para w com linha azul
plt.plot(t, w, color='blue')

# Adicionar uma legenda
plt.legend(['theta', 'w'])

# Adicionar rótulos de eixo
plt.xlabel('t (s)')
plt.ylabel('theta (rad), w (rad/s)')

# Crie um dataframe das listas t, theta e w
df = pd.DataFrame({'t (s)': t, 'theta (rad)': theta, 'w (rad/s)': w, 'k1x': k1x, 'k1v': k1v, 'k2x': k2x, 'k2v': k2v})

# Exibir as primeiras linhas do dataframe
print(df.head())

# Salva o dataframe em um arquivo CSV
df.to_csv('pendulum_results_full_heun.csv', index=False)

# Mostra o grafico
plt.show()