#! /usr/bin/env python3
# -*- coding: UTF-8 -*- 

# FY1001/TFY4145 Mekanisk fysikk hosten 2015.
# Oving 7, oppg. 2: Matematisk pendel

from numpy import sin,cos,pi,fmod,zeros,arange
from matplotlib.pyplot import rc,figure,title,xlabel,ylabel,plot,legend,show,hold
from matplotlib.animation import FuncAnimation

# Funksjon som løser N2 for en matematisk pendel, som starter i theta0
# med hastighet v0, ved Eulermetoden for N tidssteg(dt). 
def mat_pendel(L,v0,theta0,dt,N,tilfelle):
    g=9.81
    dtL=dt/L
    # Initialisere tabeller for hastighet, vinkel, energi og tid
    v=zeros(N)
    v[0]=v0
    theta=zeros(N)
    theta[0]=theta0
    E=zeros(N)
    E[0]=0.5*v[0]**2+g*L*(1-cos(theta[0]))
    tid=zeros(N)
    tid[0]=0.0
    # Starter iterasjonen i 0'te tidssteg
    t=0
    # Øvre grense for iterasjonen
    t_max=N-1
    while (t < t_max):
        v[t+1]=v[t]-g*sin(theta[t])*dt
        theta[t+1]=theta[t]+v[t]*dtL
        E[t+1]=0.5*v[t+1]*v[t+1]+g*L*(1-cos(theta[t+1]))
        tid[t+1]=t*dt
        t+=1
    # Vinkler større enn 2*pi flyttes til intervallet (-2*pi,2*pi)
    theta=fmod(theta,2*pi)
    # Sette font til Computer Modern
    rc('font',**{'family':'sans-serif','serif':'Computer Modern'})
    # Plotte Hastigheten
    figure('Figur 1')
    hold(True)
    title('Hastighet, tilfelle '+tilfelle+' med ' + r'$v_0=%.4f$'%(v0))
    plot(tid,v,label=r'$v$')
    # Merke aksene
    xlabel(r't$(s)$')
    ylabel(r'v$(m/s)$')
    # Plotte vinkelen
    figure('Figur 2')
    hold(True)
    title('Vinkel, tilfelle '+tilfelle+' med '+r'$v_0=%.4f$'%(v0))
    plot(tid,theta*180/pi,label=r'$\theta$')
    xlabel(r't$(s)$')
    ylabel(r'$\theta(grader)$')
    # Plotte energien
    figure('Figur 3')
    hold(True)
    title('Energi, tilfelle '+tilfelle+' med '+r'$v_0=%.4f$'%(v0))
    plot(tid,E,label=r'$E/m$')
    xlabel(r't$(s)$')
    ylabel(r'E/m$(m^2/s^2)$')
    # Vise plottene
    show()

    # For visualisering av pendelens bevegelse
    fps=30          # Antall bilder per sekund        
    ipf=1/fps/dt    # Antall iterasjoner per bilde
    
    x=zeros(len(theta)/ipf)
    y=zeros(len(theta)/ipf)
    for i in arange(0,len(x)):
        x[i]=sin(theta[i*ipf])
        y[i]=-cos(theta[i*ipf])

    pendel_animasjon(x,y,ipf,dt)

def pendel_animasjon(x1,y1,ipf,dt):
    fig = figure()
    ax = fig.add_subplot(111, autoscale_on=False, xlim=(-1.5, 1.5), ylim=(-1.5, 1.5))
    ax.grid()
    tid_template = 'tid = %.1fs'
    tid_tekst = ax.text(0.5, 0.9, '', transform=ax.transAxes)
    line, = ax.plot([], [], 'o-', lw=2)
    def init():
        line.set_data([], [])
        tid_tekst.set_text('')
        return line, tid_tekst
    def animate(i):
        thisx = [0, x1[i]]
        thisy = [0, y1[i]]
        line.set_data(thisx, thisy)
        tid_tekst.set_text(tid_template%(i*dt*ipf))
        return line, tid_tekst
    ani = FuncAnimation(fig, animate, arange(1, len(x1)),interval=1e3*ipf*dt, blit=True, init_func=init)
    show()

if __name__=="__main__":
    # # Definere parametre
    # Velg tilfelle ='A', 'B' eller 'C'
    tilfelle='A'
    # # For tilfelle A:
    dt=1e-4             # Størrelsen på tidssteg
    T=2.0               # Periode
    v0=1.0              # Starthastighet

    # # For tilfelle B:
    if(tilfelle=='B'):
        dt=1e-6         # Størrelsen på tidssteg
        T=8.7531        # Periode
        v0=6.2641       # Starthastighet

    # # For tilfelle C:
    if(tilfelle=='C'):
        dt=1e-4         # Størrelsen på tidssteg
        T=1.684           # Periode
        v0=6.5          # Starthastighet
    
    theta0=0.0      # Startvinkel
    L=1.0           # Pendelens lengde
    N=int(2.0*T/dt) # Antall tidssteg
    # # Kjør funksjonen
    mat_pendel(L,v0,theta0,dt,N,tilfelle)