#!/usr/bin/env python

import math
from pylab import *

def fun(x,t,nu):
    return array([x[2],x[3],-nu*x[2],-1-nu*x[3]])

def ene(x):
    return 0.5*(x[2]*x[2]+x[3]*x[3])+x[1]

def inte_e(x,t,dt,nu):
    y = x + dt*fun(x,t,nu)
    return y

def inte_ec1(x,t,dt,nu):
    y = x
    f = fun(x,t,nu)
    y[2] = x[2] + dt*f[2]
    y[3] = x[3] + dt*f[3]
    y[0] = x[0] + dt*y[2]
    y[1] = x[1] + dt*y[3]
    return y

def inte_ec2(x,t,dt,nu):
    y=x
    f = fun(x,t,nu)
    y[2] = x[2] + dt*f[2]
    y[3] = x[3] + dt*f[3]
    y[0] = x[0] + dt*x[2]+ 0.5*dt*dt*f[2]
    y[1] = x[1] + dt*x[3]+ 0.5*dt*dt*f[3]
    return y


def inte_vv(x,t,dt,nu):
    y = x
    f = fun(x,t,nu)
    y[0] = x[0] + dt*x[2]+ 0.5*dt*dt*f[2]
    y[1] = x[1] + dt*x[3]+ 0.5*dt*dt*f[3]
    #y[2] = x[2]+ 0.5*dt*f[2]
    #y[3] = x[3]+ 0.5*dt*f[3]
    f1 = fun(y,t,nu)
    y[2] = x[2] + 0.5*dt*(f[2]+f1[2])
    y[3] = x[3] + 0.5*dt*(f[3]+f1[3])
    return y


def inte_rk(x,t,dt,nu):
    y = x
    f1 = dt*fun(y,t,nu)
    y = x + 0.5*f1
    f2 = dt*fun(y,t+0.5*dt,nu)
    y = x + 0.5*f2
    f3 = dt*fun(y,t+0.5*dt,nu)
    y = x + f3
    f4 = dt*fun(y,t+dt,nu)
    y = x + (f1+2.0*f2+2.0*f3+f4)/6.0
    return y
    

def main():
    the = math.pi/4.0
    dt  = 0.05
    nuo = 0.0





    subplot(211)
    title(r'Movimiento parabolico $\nu$ = '+str(nuo)+r' $\Delta t$ = '+str(dt) )
    xlabel(r'$\bar{x}=gx/v_o^2$')
    ylabel(r'$\bar{y}=gy/v_o^2$')

    subplot(212)
    xlabel(r'$\tau=gt/v_o$')
    ylabel(r'$\bar{E}=E/mv_o^2$')



    # EXACTO

    x   = array([0,0,cos(the),sin(the)])
    t   = 0.0
    posx=[x[0]]
    posy=[x[1]]
    energia=[ene(x)]
    tiempo=[t]
    while x[1]>=0.0:
        t = t+dt
        x[0] = cos(the)*t
        x[1] = sin(the)*t-0.5*t*t
        x[2] =  cos(the)
        x[3] =  sin(the)- t
        posx.append(x[0])
        posy.append(x[1])
        tiempo.append(t)
        energia.append(ene(x))
    print 'Exacto se demoro = ',len(posx)

    subplot(211)
    plot(posx,posy,'ks-',label=r'Exacto $\nu=0$',linewidth=1)

    subplot(212)
    plot(tiempo,energia,'ks-',linewidth=1)


    # EULER    

    t   = 0.0
    x   = array([0,0,cos(the),sin(the)])
    posx=[x[0]]
    posy=[x[1]]
    tiempo=[t]
    energia=[ene(x)]
    while x[1]>=0.0:
        x = inte_e(x,t,dt,nuo)
        t = t+dt
        posx.append(x[0])
        posy.append(x[1])
        tiempo.append(t)
        energia.append(ene(x))
    print 'Euler se demoro = ',len(posx)
    
    subplot(211)
    plot(posx,posy,'rs-',label='Euler',linewidth=1)
    
    subplot(212)
    plot(tiempo,energia,'rs-',linewidth=1)


    # EULER-Cromer

    t   = 0.0
    x   = array([0,0,cos(the),sin(the)])
    posx=[x[0]]
    posy=[x[1]]
    tiempo=[t]
    energia=[ene(x)]
    while x[1]>=0.0:
        x = inte_ec1(x,t,dt,nuo)
        t = t+dt
        posx.append(x[0])
        posy.append(x[1])
        tiempo.append(t)
        energia.append(ene(x))
    print 'Euler-Cromer se demoro = ',len(posx)
    
    subplot(211)
    plot(posx,posy,'gs-',label='Euler-Cromer',linewidth=1)
    
    subplot(212)
    plot(tiempo,energia,'gs-',linewidth=1)


    # EULER-CROMER (punto medio)

    t   = 0.0
    x   = array([0,0,cos(the),sin(the)])
    posx=[x[0]]
    posy=[x[1]]
    tiempo=[t]
    energia=[ene(x)]
    while x[1]>=0.0:
        x = inte_ec2(x,t,dt,nuo)
        t = t+dt
        posx.append(x[0])
        posy.append(x[1])
        tiempo.append(t)
        energia.append(ene(x))
    print 'Euler-Cromer 2 se demoro = ',len(posx)
    
    subplot(211)
    plot(posx,posy,'bs-',label='Euler-Cromer-2',linewidth=1)
    
    subplot(212)
    plot(tiempo,energia,'bs-',linewidth=1)


    # VELOCITY-VERLET

    t   = 0.0
    x   = array([0,0,cos(the),sin(the)])
    posx=[x[0]]
    posy=[x[1]]
    tiempo=[t]
    energia=[ene(x)]
    while x[1]>=0.0:
        x = inte_vv(x,t,dt,nuo)
        t = t+dt
        posx.append(x[0])
        posy.append(x[1])
        tiempo.append(t)
        energia.append(ene(x))
    print 'Velocity Verlet se demoro = ',len(posx)

    
    subplot(211)
    plot(posx,posy,'ms-',label='Velocity-Verlet',linewidth=1)
    
    subplot(212)
    plot(tiempo,energia,'ms-',linewidth=1)



    # RUNGE_KUTA

    t   = 0.0
    x   = array([0,0,cos(the),sin(the)])
    posx=[x[0]]
    posy=[x[1]]
    tiempo=[t]
    energia=[ene(x)]
    while x[1]>=0.0:
        x = inte_rk(x,t,dt,nuo)
        t = t+dt
        posx.append(x[0])
        posy.append(x[1])
        tiempo.append(t)
        energia.append(ene(x))
    print 'RUNGE-KUTA se demoro = ',len(posx)

    
    subplot(211)
    plot(posx,posy,'cs-',label='Runge-Kuta',linewidth=1)
    
    subplot(212)
    plot(tiempo,energia,'cs-',linewidth=1)





    subplot(211)
    legend()
    axis([0, 1.5, 0, 0.4])
    
    subplot(212)
    axis([0, 2, 0.45, 0.55])
    
    show()

  
    
  ##  axis([-20, 0, -10, 1])
##    xla    subplot(212)
##    ylabel(r'$\log\; \Delta(h)$')
##    grid(True)



 

if __name__=='__main__':
    main()