"""
Este script simula el ejercicio 1.8 de la guía 1 de Física 2 (Q),
Cát. Laura Estrada, 1er cuatrimestre 2018, FCEyN, UBA, Argentina.

Autor: Rodrigo Lugones
"""

# Llamo a las bibliotecas/librerias que voy a utilizar
import numpy as np
import matplotlib.pyplot as plt
import scipy.integrate as integrate
import matplotlib.animation as animation


# Parametros modificables.
m_Q = 9.1e-31  # masa (kg) de la particula de prueba
Q = 1.6e-19    # carga (C) de la particula de prueba
x_0 = .01      # posicion inicial (m) de la particula de prueba: [x_0, 0]

q = 1.6e-19    # carga (C) de las particulas fijas
y_q = 1        # posicion (m) de las particulas fijas: [0, -y], [0, +y]

t_max = 20     # tiempo final (s) de la simulacion

k = 1e9        # Constante k = 1/(4*\pi*\epsilon_0)   (in N*m^2/C^2)



def derivs(state, t):
  # Esta funcion 'derivs' encuentra numericamente la aceleracion y la velocidad:
  # dx/dt = v
  #                 k          q * Q                    x
  # dv/dt = - 2 *  --- * ------------------- * -------------------
  #                m_Q   (x^2 + y_q^2)^(2/2)   (x^2 + y_q^2)^(1/2)
  dxdt = np.zeros_like(state)
  dxdt[0] = state[1] #
  dxdt[1] = -2*k/m_Q * Q * q / (y_q**2 + state[0]**2)**(3/2) * state[0]

  return dxdt


# crea un array de tiempos desde 0 hasta t_max, con paso dt
dt = 0.05
t = np.arange(0.0, t_max, dt)

# x_Q y v_Q son la posicion y la velocidad inicial de la particula de prueba
xQ_0 = x_0
vQ_0 = 0.0
state = [xQ_0, vQ_0]

# Encuentra numericamente la posicion. Noten que utiliza la funcion
# derivs, el estado inicial y el array de tiempos. El array x es la
# posicion no da la posicion x y la velocidad a cada tiempo. A
# nosostros, solo nos va a interesar la posicion (por eso, nos
# quedamos solo con la columna 0).
xQ = integrate.odeint(derivs, state, t)[:, 0] # posicion x en funcion del tiempo
yQ = np.zeros(np.shape(xQ)) # posicion y en funcion del tiempo


# Creo el grafico que se va a ir actualizando
fig = plt.figure()

# Este es el subgrafico de arriba
ax = fig.add_subplot(211, autoscale_on=False, xlim=(-x_0*2, x_0*2), ylim=(-y_q*1.1, y_q*1.1))
ax.grid()
line, = ax.plot([], [], 'o-', lw=2)
line2, = ax.plot([0, 0], [-y_q, y_q], 'o', lw=2)
time_template = 'time = %.1fs'
time_text = ax.text(0.05, 0.9, '', transform=ax.transAxes)


# Este es el subgrafico de abajo
ax2 = fig.add_subplot(212, autoscale_on=True)
ax2.grid()
line3, = ax2.plot(t, xQ, '-r', lw=2)
line4, = ax2.plot(t, xQ, 'o-r', lw=2)
line5, = ax2.plot(t, xQ, ':b', lw=2)



# Creo la animacion
def init():
  return line, line3, line4, line5, time_text

def animate(i):
  line.set_data(xQ[i], yQ[i])
  line3.set_data(t[:i+1], xQ[:i+1])
  line4.set_data(t[i], xQ[i])
  line5.set_data(t[:i+1], xQ_0*np.cos(np.sqrt(2*k/m_Q * Q*q/y_q**3) * t[:i+1]))

  time_text.set_text(time_template % (i*dt))
  return line, line3, line4, line5, time_text


# Corro la animacion
ani = animation.FuncAnimation(fig, animate, np.arange(1, len(xQ)),
                              interval=25, blit=True, init_func=init, repeat=False)

# ani.save('double_pendulum.mp4', fps=15) # Con esta linea, puedo guardar un video mp4
plt.show()