import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
#----------DEFINO LAS MAGNITUDES---------
N = 400000 #cantidad de pasos del método de Euler.
t_i = 0 #tiempo inicial
h = 0.0002 #paso que avanzo en Euler
t_f = t_i + N * h
g = 9.8
m = 0.5
M = 0.8
l = 100
r0 = 10
tita0 = 0
v_r0 = 0
v_tita0 = (M*g/(m*r0))**0.5 # Esta velocidad angular es la de movimiento circular
#----------------------------------------
#--------CREO VECTORES PARA POSICIONES Y VELOCIDADES, E IMPONGO CONDICIONES INICIALES------
r = np.zeros(N) #Esto me crea un vector (0,0,0,0,...) de N componentes. Hay que crear así los vectores antes de llenarlos con los valores calculados con euler.
r_dot = np.zeros(N)
tita = np.zeros(N)
tita_dot = np.zeros(N)
r[0] = r0 #Pido que la primera componente de cada vector sean las condiciones iniciales.
r_dot[0] = v_r0
tita[0] = tita0
tita_dot[0] = v_tita0
j = 0 #Creo esta variable que toma el valor de la cantidad de pasos de Euler que ya dimos.
#Así si en algún momento la masita llega al agujerito de la mesa el programa se corta y sé hasta qué paso considerar.
for i in range(0,N-1): #Lo que está adentro del for es el método de Euler.
    r[i + 1] = r[i] + h * r_dot[i]
    r_dot[i + 1] = r_dot[i] + h * (m / (M + m) * r[i] * tita_dot[i]**2 - (M/(M + m)) * g)
    tita[i + 1] = tita[i] + h * tita_dot[i]
    tita_dot[i + 1] = tita_dot[i] + h * (-2 * r_dot[i] * tita_dot[i] / (m * r[i]))
    j = i - 2 #Acá puse que j sea la cantidad de pasos menos 2, para no terminar cuando la masita "cae", sino 2 pasos antes.
    if r[i + 1] <= 0:
        break #Si la masita cae por el agujero se para el ciclo for.
    if r[i + 1] >= l:
            break# Si la masita se va a un radio mayor que el largo de la soga se para el ciclo for.
##-----------EULER MODIFICADO--------------
#Lo que sigue es lo mismo pero con Euler modificado.
def f_r(rdot): #En vez de escribir las ecuaciones a mano dentro del for las defino como funciones afuera.
    return rdot
def f_rdot(r,titadot):
    return m/(M + m) * r * titadot**2 - (M/(M + m)) * g
def f_tita(titadot):
    return titadot
def f_titadot(r,rdot,titadot):
    return -2*rdot*titadot/r

for i in range(0,N-1):
    r[i + 1] = r[i] + h/2 * f_r(r_dot[i] + h * f_r(r_dot[i]))
    tita[i + 1] = tita[i] + h/2 * f_tita(tita_dot[i] + h * f_tita(tita_dot[i]))
    r_dot[i + 1] = r_dot[i] + h/2 * f_rdot(r[i] + h * f_r(r[i]), tita_dot[i] + h * f_titadot(r[i],r_dot[i],tita_dot[i]))
    tita_dot[i + 1] = tita_dot[i] + h/2 * f_titadot(r[i] + h * f_r(r[i]), r_dot[i] + h * f_rdot(r[i],tita_dot[i]),tita_dot[i] + h * f_titadot(r[i],r_dot[i],tita_dot[i]))
    j = i - 2
    if r[i + 1] <= 0:
        break
    if r[i + 1] >= l:
            break
##----- DE ACA PARA ABAJO ES GRAFICAR-------
#--------Animación--------------
t = np.linspace(t_i, t_f, N)
x = np.zeros(N)
y = np.zeros(N)
E = np.zeros(N)
for i in range(0,j):
    x[i] = r[i] * np.cos(tita[i])
    y[i] = r[i] * np.sin(tita[i])
    E[i] = (m + M)*0.5*r_dot[i]**2 + m*0.5*(r[i]*tita_dot[i])**2 + M*g*(r[i]-l)

fig, ax = plt.subplots()

xdata, ydata = [], []
ln, = plt.plot([], [], 'r-',markersize=0.5)

# Esto genera una lista de frames,
# que es lo mismo que intervalos de tiempo
numf = 10000
nfram = np.linspace(0, numf - 1, numf)

# Predefinir las variables a usar

ii = 0

# Cualidades estaticas del grafico
def init():
    #ax.set_facecolor(color='lightyellow')
    ax.set_xlim(-15, 15)
    ax.set_ylim(-15, 15)
    plt.grid()
    return ln,

# Funcion que genera el contenido de los frames
# y los va agregando a la estructura a plotear (ln)
def update(frame):
    ii= int(frame)
    sx = x[ii * int(j/numf) ]
    sy = y[ii * int(j/numf)]
    xdata.append(sx)
    ydata.append(sy)
    ln.set_data(xdata, ydata)  # Pack coordinate pairs into plot arrays
    return ln,

ani = FuncAnimation(fig, update, frames=nfram, interval=5,init_func=init,
                            repeat=False, blit=True)

# Una vez armados los contenidos de los nfram frames, mostrar
#-----------------------------------------------------------------
#------------Gráfico con 4 ventanas. r(t),θ(t), r(θ) y E(t). 
plt.show()
plt.figure()
plt.subplot(2,2,1)
plt.plot(t[:j], r[:j])
plt.ylim([0,l])
plt.grid()
plt.ylabel("r")
plt.xlabel("t")
plt.subplot(2,2,2)
plt.plot(t[:j], tita[:j])
plt.grid()
plt.ylabel("θ")
plt.xlabel("t")
plt.subplot(2,2,3)
plt.plot(x[:j], y[:j], linewidth = 0.5)
plt.grid()
plt.subplot(2,2,4)
plt.plot(t[:j], E[:j])
plt.grid()
plt.ylabel("E")
plt.xlabel("t")
plt.show()