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

# FY1001/TFY4145 Mekanisk fysikk hosten 2015.
# Oving 6

from numpy import loadtxt,array,append,zeros,arctan,sqrt,pi,cos,sin,arange
from matplotlib.pyplot import figure,title,xlabel,ylabel,plot,legend,hold,show,subplot,axvline,subplots,axes,rc
from matplotlib.patches import Wedge,Circle
from matplotlib.animation import FuncAnimation

def main():
    # Laste inn data fra fil, hvor de to forste linene ikke inneholder data
    data=loadtxt('bordtennisball.txt',skiprows=2)
    c=2.0/3     # Bordtennisballens treghetsmoment I=cr^2 
    g=9.81		# Tyngdens akselerasjon
    # Separering av data lest inn fra filen
    texp=data[:,0]
    xexp=data[:,1]
    yexp=data[:,2]
    # N=Antall datapunkt(tidssteg)
    N=len(texp)
    # Definere tomme tabeller for eksperimentell hastighet og energi
    Vxexp=zeros(N)
    Vyexp=zeros(N)
    Vexp=zeros(N)
    Kexp=zeros(N)
    Uexp=zeros(N)
    Eexp=zeros(N)
    # Beregne hastighet i det første tidssteget
    Vxexp[0]=(xexp[1]-xexp[0])/(texp[1]-texp[0])
    Vyexp[0]=(yexp[1]-yexp[0])/(texp[1]-texp[0])
    # Beregne hastigheten for hvert tidssteg(untatt forste og siste)
    for i in range(1,N-1):
        Vxexp[i]=(xexp[i+1]-xexp[i-1])/(texp[i+1]-texp[i-1])
        Vyexp[i]=(yexp[i+1]-yexp[i-1])/(texp[i+1]-texp[i-1])
    # Beregne hastighet i det siste tidssteget
    Vxexp[N-1]=(xexp[N-1]-xexp[N-2])/(texp[N-1]-texp[N-2])
    Vyexp[N-1]=(yexp[N-1]-yexp[N-2])/(texp[N-1]-texp[N-2])
    # Beregne vinklen phi og radien r+R for hvert tidssteg
    phiexp=zeros(N)
    rexp=zeros(N)
    for i in range(0,N):
        phiexp[i]=arctan(xexp[i]/yexp[i])       # Vinkel ift. origo
        rexp[i]=sqrt(xexp[i]**2+yexp[i]**2)     
        Vexp[i]=sqrt(Vxexp[i]**2+Vyexp[i]**2)
        Kexp[i]=(c+1)*Vexp[i]**2/(2*100*100)    # Kinetisk energi ved rulling
        Uexp[i]=g*yexp[i]/100                   # Potensiell energi i tyngdefeltet
        Eexp[i]=Kexp[i]+Uexp[i]                 # Total energi=Kin+Pot 
    # Felles parametere
    g=9.81		# Tyngdens akselerasjon
    phimax=pi/2	# Maksimal vinkel
    dt=0.01		# Tidssteg
    R=0.785		# Sirkelradius
    # Parametere for bordtennisball
    mys=0.25    	# Statisk friksjonskoeffisient
    myk=0.25	    # Kinematisk friksjonskoeffisient
    r=0.019		# Bordtennisballens radius
    c=2.0/3		# Bordtennisballens treghetsmoment I=cr^2
    #c=2.0/5		# Sfærens treghetsmoment I=cr^2
    #c=1.0		# Tynn sylinderskall treghetsmoment I=cr^2
    #c=1.0/2		# Massiv sylinders treghetsmoment I=cr^2

    rpR=r+R                     # Radius til massesenteret
    tid=array([0.0])            # Starter tiden ved t=0
    V=array([0.015])            # Start-hastighet
    phi=array([0.047])          # Start-vinkel
    Omega=array(V/rpR)          # Vinkelhastighet til CM, om origo
    omega=array(V/rpR)          # Vinkelhastighet til legemet, om CM
    n=array(g*cos(phi))         # Normalkraft
    A=array(g*sin(phi)/(c+1))   # Akselerasjon langs banen
    Alpha=array(A/rpR)          # Vinkelakselerasjon til CM
    f=array([0.0])              # Friksjonskraft
    fmax=array([0.0])           # Maksimal statisk friksjonskraft
    V_rel=array([V[0]-omega[0]*r])       # Hastighet til legemets kontaktpunkt mot underlaget
    
    x=array(rpR*sin(phi))       # CM Posisjon, x-koordinat
    y=array(rpR*cos(phi))       # CM Posisjon, y-koordinat
    Vx=array(V*cos(phi))        # CM Hastighet, x-komponent
    Vy=array(V*sin(phi))        # CM Hastighet, y-komponent
    theta=array([0.0])         # Vinkel rotert om CM
    
    kintrans=array(V[0]**2/2)   # Kinetisk energi fra translasjon
    kinrot=array(c*r**2*omega[0]**2/2) # Kinetisk rotasjonsenergi
    potensiell=array(g*y[0])    # Potensiell energi i tyngdefeltet

    pwork=array([0.0])          # Friksjonsarbeid ved sluring
    fwork=array([0.0])          # Akkumulert friksjonsarbeid
    totalmekanisk=array(kintrans+kinrot+potensiell) # Total mekanisk energi
    totalenergi=array(totalmekanisk+fwork)          # Total energi
    rulleenergi=array((c+1)*V[0]**2/2+potensiell)   # Energi for rullebevegelsen 

    jmax=20/dt                  # Maksimalt antall tidssteg
    jrulling=0                  # Siste steg-nummer for rulling
    jkontakt=0                  # Siste steg-nummer for kontakt
    j=0
    
    # Nå er alle størrelser vi trenger definert og initialisert.
    # Vi integrerer bevegelsesligningen ved å iterere while-løkken
    while (n[j] > 0.0) and (phi[j] < phimax) and (j<jmax):
        V=append(V,V[j]+A[j]*dt)                        # Ny hastighet
        Omega=append(Omega,V[j+1]/rpR)                  # Ny vinkelhastighet om origo
        omega=append(omega,V[j+1]/r)                    # Ny vinkelhastighet om CM
        phi=append(phi,phi[j]+Omega[j]*dt)              # Ny vinkel
        x=append(x,rpR*sin(phi[j+1]))                   # Ny x-koordinat
        y=append(y,rpR*cos(phi[j+1]))                   # Ny y-koordinat
        Vx=append(Vx,V[j+1]*cos(phi[j+1]))              # Ny hastighet, x-komponent
        Vy=append(Vy,-V[j+1]*sin(phi[j+1]))             # Ny hastighet, y-komponent
        n=append(n,g*cos(phi[j+1]) - (V[j+1]**2)/(r+R)) # Ny normalkraft
        fmax=append(fmax,mys*n[j+1])                    # Ny maksimal statisk friksjon
        A=append(A,g*sin(phi[j+1])/(c+1))               # Ny akselerasjon
        Alpha=append(Alpha,A[j+1]/rpR)                  # Ny vinkelakselerasjon
        f=append(f,c*A[j+1])                            # Ny statisk friksjonskraft
        kintrans=append(kintrans,V[j+1]**2/2)           # Ny kinetisk energi
        kinrot=append(kinrot,c*r**2*omega[j+1]**2)      # Ny rotasjonsenergi
        potensiell=append(potensiell,g*y[j+1])          # Ny potensiell energi
        rulleenergi=append(rulleenergi,(c+1)*V[j+1]**2/2+potensiell[j+1])   # Ny rulleenergi
        if (f[j+1] > fmax[j+1]):                        # Sjekke om skliing er aktuelt
            omega[j+1]=omega[j]+f[j]*dt/(c*r)           # Vinkelhastighet rundt CM må endres
            f[j+1]=myk*n[j+1]                           # Ny kinematisk friksjonskraft
            A[j+1]=g*sin(phi[j+1])-f[j+1]               # Ny akselerasjon
            Alpha[j+1]=A[j+1]/rpR                       # Ny vinkelakselerasjon
        else:
            jrulling=j+1                                # Tar vare pa siste steg med rulling
        theta=append(theta,theta[j]+omega[j]*dt*180/pi)  # Ny vinkel rotert
        V_rel=append(V_rel,V[j+1]-omega[j+1]*r)         # Relativhastighet i kontaktpunktet
        pwork=append(pwork,V_rel[j]*f[j])               # Effekttap pga. sluring er 0 ved ren rulling
        fwork=append(fwork,fwork[j]+pwork[j+1]*dt)      # Akkumulert friksjonsarbeid
        totalmekanisk=append(totalmekanisk,kintrans[j+1]+kinrot[j+1]+potensiell[j+1])
        totalenergi=append(totalenergi,totalmekanisk[j+1]+fwork[j+1])
        jkontakt=j+1                                    # Tar vare pa siste steg med kontakt
        tid=append(tid,j*dt)                            # Ny tid
        j+=1                                            # Nytt steg

    
    # Nå er vi ferdig med å integrere bevegelsesligningen.
    # Printer ut interessente vinkler og hastigheter
    print('Ved rulleslutt er:\n vinkelen: %.1f grader\n vinkelhastigheten: %.1f grader/s\n hastigheten: %.1f  m/s'%(phi[jrulling]*180/pi,Omega[jrulling]*180/pi,V[jrulling]))
    print('Ved kontaktslutt er:\n vinkelen: %.1f grader\n vinkelhastigheten: %.1f grader/s\n hastigheten: %.1f  m/s'%(phi[jkontakt]*180/pi,Omega[jkontakt]*180/pi,V[jkontakt]))

    # Her lager vi plot av ulike størrelser for å sammenligne numeriske og
    # eksperimentelle data.

    rc('font',**{'family':'sans-serif','serif':'Computer Modern'})
    # Figur bestaende av 2 figurer hoyden og 2 i bredden
    figure('Figur1')
    # del-figur 1: Vinkel, vinkelhastighet om origo og normalkraft som funksjon av tid
    subplot(221)
    title(r'$\phi,\Omega$'+' og '+r'$N$'+' som funksjon av tid')
    hold(True)
    xlabel('t(s)')
    plot(tid,phi*180/pi,label=r'$\phi({\rm grader})$')
    plot(tid,Omega*180/pi,label=r'$\Omega({\rm grader}/s)$')
    plot(tid,n,label='$N(N/kg)$')
    axvline(tid[jrulling])  # Linje som skiller(tid) rulling fra skliing
    legend(loc='best')

    # Delfigur 2: Vinkelhastighet, vinkelakselerasjon og normalkraft som funksjon av vinkel
    subplot(222)
    title(r'$\Omega, \alpha$'+' og '+r'$N$'+' som funksjon av vinkel')
    hold(True)
    xlabel(r'$\phi({\rm grader})$')
    plot(phi*180/pi,Omega*180/pi,label=r'$\Omega({\rm grader}/s)$')
    plot(phi*180/pi,Alpha*180/pi,label=r'$\alpha({\rm grader}/s^2)$ ')
    plot(phi*180/pi,n,label=r'$N(N/kg)$')
    axvline(phi[jrulling]*180/pi)  # Linje som skiller(vinkel) rulling fra skliing
    legend(loc='best')

    # Delfigur 3: Hastighet i x- og y-retning som funksjon av tid
    subplot(223)
    title(r'$V_x$'+' og '+r'$V_y$'+' som funksjom av tid')
    hold(True)
    xlabel('t(s)')
    plot(tid,Vx*100,label=r'$V_x(cm/s)$')
    plot(tid,Vy*100,label=r'$V_y(cm/s)$')
    plot(texp,Vxexp,label=r'$V_x^{exp}$')
    plot(texp,Vyexp,label=r'$V_y^{exp}$')
    axvline(tid[jrulling])
    legend(loc='best')

    # Delfigur 4: Hastighet i x- og y-retning som funksjon av vinkel
    subplot(224)
    title(r'$V_x$'+' og '+r'$V_y$'+' som funksjom av vinkel')
    hold(True)
    xlabel(r'$\phi({\rm grader})$')
    plot(phi*180/pi,Vx*100,label=r'$V_x(cm/s)$')
    plot(phi*180/pi,Vy*100,label=r'$V_y(cm/s)$')
    plot(phiexp*180/pi,Vxexp,label=r'$V_x^{\rm exp}$')
    plot(phiexp*180/pi,Vyexp,label=r'$V_y^{\rm exp}$')
    axvline(phi[jrulling]*180/pi)
    legend(loc='best')

    # Ny figur medd eksperimentell og numerisk hastighet
    figure('Figur 2')
    title('Numerisk og eksperimentell hastighet som funksjon av vinkel')
    hold(True)
    xlabel(r'$\phi(grader)$')
    plot(phi*180/pi,V*100,label=r'$V_{\rm num}$')
    plot(phiexp*180/pi,Vexp,label=r'$V_{\rm exp}$')
    axvline(phi[jrulling]*180/pi)
    legend(loc='best')

    # Ny figur med avstand(radius) fra origo til CM fra eksperimentelle data
    figure('Figur 3')
    title('Eksperimentell avstand fra CM til origo som funksjon av vinkel\nVertikale linjer skiller mellom numeriske vinkler for rulling, skliing og skrått kast')
    hold(True)
    xlabel(r'$\phi(grader)$')
    ylabel('Avstand fra origo til CM '+r'$(J/kg)$')
    plot(phiexp*180/pi,rexp)
    axvline(phi[jrulling]*180/pi)  # Linje som skiller(vinkel) rulling fra skliing
    axvline(phi[jkontakt]*180/pi)  # Linje som skiller(vinkel) skliing fra skrått kast

    # Ny figur med energier som funksjon av tid
    figure('Figur 4')
    hold(True)
    title('Numerisk og eksperimentell energi som funksjon av tid')
    xlabel('tid '+r'$(s)$')
    ylabel('Energi per masseenhet '+r'$(J/kg)$')
    plot(tid,totalmekanisk,label=r'$E^{\rm mek}_{\rm tot}$')
    plot(tid,totalenergi,label=r'$E_{\rm tot}$')
    plot(tid, rulleenergi, label=r'$E_{\rm num}({\rm rulling\ antatt!})$')
    plot(texp,Eexp,label=r'$E_{\rm exp}({\rm rulling\ antatt!})$')
    axvline(tid[jrulling])
    legend(loc='best')

    # Ny figur med energier som funksjon av vinkel
    figure('Figur 5')
    hold(True)
    title('Numerisk og eksperimentell energi som funksjon av vinkel')
    xlabel('Vinkel '+r'$({\rm grader})$')
    ylabel('Energi per masseenhet '+r'$(J/kg)$')
    plot(phi*180/pi,totalmekanisk,label=r'$E^{\rm mek}_{\rm tot}$')
    plot(phi*180/pi,totalenergi,label=r'$E_{\rm tot}$')
    plot(phi*180/pi, rulleenergi, label=r'$E_{\rm num}({\rm rulling\ antatt!})$')
    plot(phiexp*180/pi,Eexp,label=r'$E_{\rm exp}({\rm rulling\ antatt!})$')
    axvline(phi[jrulling]*180/pi)
    legend(loc='best')

    # Ny figur som viser friksjonsarbeidet som funksjon av vinkel
    figure('Figur 6')
    hold(True)
    title('Totalt friksjonsarbeid som funksjon av vinkel')
    xlabel('Vinkel '+r'$({\rm grader})$')
    ylabel('Arbeid per masseenhet '+r'$(J/kg)$')
    plot(phi*180/pi,fwork)
    axvline(phi[jrulling]*180/pi)

    # Ny figur med total energi benyttet til friksjonsarbeid, som funksjon av vinkel
    figure('Figur 7')
    hold(True)
    title('Effekttap pga. sluring som funksjon av vinkel')
    xlabel('Vinkel '+r'$({\rm grader})$')
    ylabel('Effekttap per masseenhet '+r'$(J/kg)$')
    plot(phi*180/pi,pwork)
    axvline(phi[jrulling]*180/pi)

    # Ny figur med den relative hastigheten i kontaktpunktet
    figure('Figur 8')
    hold(True)
    title('Kontaktpunktets hastighet som funksjon av vinkel')
    xlabel('Vinkel '+r'$({\rm grader})$')
    ylabel('Hastighet '+r'$(cm/s)$')
    plot(phi*180/pi,V_rel*100)
    axvline(phi[jrulling]*180/pi)

    # Ny figur med friksjonskraften som funksjon av vinkel
    figure('Figur 9')
    hold(True)
    title('Friksjonskraft som funksjon av vinkel')
    xlabel('Vinkel '+r'$({\rm grader})$')
    ylabel('Kraft per masseenhet '+r'$(N/kg)$')
    plot(phi*180/pi,f)
    axvline(phi[jrulling]*180/pi)

    # Her lager vi en animasjon av rullebevegelsen
    # Matcher eksperimentell tid med numerisk tid, ved å skalere eksperimentell tid
    x_exp_scaled=zeros(len(x))
    y_exp_scaled=zeros(len(x))
    if(len(texp)>len(tid)):
        t_faktor=(tid[-1]-tid[0])/(texp[len(tid)]-texp[0])
    else:
        t_faktor=(tid[len(xexp)]-tid[0])/(texp[-1]-texp[0])
    # Skalerer eksperimentell posisjon fra cm til m, og flytter den i tid(matcher steg nr. i med tid)
    for i in arange(0,len(x)):
        x_exp_scaled[i]=xexp[i*t_faktor]/100
        y_exp_scaled[i]=yexp[i*t_faktor]/100

    fps=50                      # Antall bilder per sekund i animasjonen
    ipf=1/fps/dt                # Antall steg(iterasjoner) per bilde
    phi_rulling=phi[jrulling]   
    phi_kontakt=phi[jkontakt]
    #anim=animate_rolling_circle(R,r,x,y,xexp,yexp,phi,theta,ipf,dt,phi_rulling,phi_kontakt,V_rel)
    anim=animate_rolling_circle(R,r,x,y,x_exp_scaled,y_exp_scaled,phi,theta,ipf,dt,phi_rulling,phi_kontakt,V_rel)
    
    # For å vise plottene og animasjonen
    show()
    
    # Funksjon som lager en animasjon av eksperimentell og numerisk rullebevegelse
def animate_rolling_circle(R,r,x,y,xexp,yexp,phi,theta,ipf,dt,phi_rulling,phi_kontakt,V_rel):
    fig,ax=subplots()
    ax = axes(xlim=(0, R+3*r), ylim=(0, R+3*r))
    patch = Wedge((5, 5), 0.75, 0,180, fc='w')
    plot([0,R*cos(pi/2-phi_rulling)],[0,R*sin(pi/2-phi_rulling)],'k')
    ax.text(0.3,0.1,'Animasjon av rullende bordtennisball;\nnumerisk i sort/hvitt, og \neksperimentelt i rødt', transform=ax.transAxes)
    ax.text(0.1, 0.9*R, 'Ruller', transform=ax.transAxes)
    plot([0,R*cos(pi/2-phi_kontakt)],[0,R*sin(pi/2-phi_kontakt)],'k')
    ax.text(cos(pi/2-phi_rulling), 0.9*R*sin(pi/2-phi_rulling), 'Slurer', transform=ax.transAxes)
    tid_template = r'tid = %.1fs  $V_{\rm rel}$=%.1fm/s'
    tid_tekst = ax.text(0.5, 0.9, '', transform=ax.transAxes)
    patches_k=[]
    patches_w=[]
    patches_c=[]
    patch1 = Wedge((x[0], y[0]),r,0-theta[0],180-theta[0],fc='k')
    patch2 = Wedge((x[0], y[0]),r,180-theta[0],0-theta[0],fc='w')
    patch3 = Circle((xexp[0],yexp[0]),r,fc='r')
    patches_k.append(patch1)
    patches_w.append(patch2)
    patches_c.append(patch3)
    antall=int(len(x)/ipf)+1
    for i in range(1,antall):
        patch1 = Wedge((x[i*ipf-1], y[i*ipf-1]),r,0-theta[i*ipf-1],180-theta[i*ipf-1],fc='k')
        patch2 = Wedge((x[i*ipf-1], y[i*ipf-1]),r,180-theta[i*ipf-1],0-theta[i*ipf-1],fc='w')
        patch3 = Circle((xexp[i*ipf-1],yexp[i*ipf-1]),r,fc='r',alpha=0.75)
        patches_k.append(patch1)
        patches_w.append(patch2)
        patches_c.append(patch3)
        ax.add_patch(patches_k[int(i)])
        ax.add_patch(patches_w[int(i)])
        ax.add_patch(patches_c[int(i)])
        patches_k[int(i)].set_visible(False)
        patches_w[int(i)].set_visible(False)
        patches_c[int(i)].set_visible(False)
    def init():
        patch = Wedge((0, 0),R,0,90,fc='w')
        ax.add_patch(patch)
        ax.add_patch(patches_k[0])
        ax.add_patch(patches_w[0])
        tid_tekst.set_text('')
        return ax, tid_tekst
    def animate(i):
        if(i==1):
            patches_k[-1].set_visible(False)
            patches_w[-1].set_visible(False)
            patches_c[-1].set_visible(False)
        patches_k[int(i-1)].set_visible(False)
        patches_w[int(i-1)].set_visible(False)
        patches_c[int(i-1)].set_visible(False)
        patches_c[int(i)].set_visible(True)
        patches_k[int(i)].set_visible(True)
        patches_w[int(i)].set_visible(True)
        tid_tekst.set_text(tid_template%(i*dt*ipf,V_rel[i*ipf-1]))
        return ax, tid_tekst
    anim = FuncAnimation(fig, animate,arange(1,(len(x)+1)/ipf),init_func=init,interval=1e3*ipf*dt,blit=True)
    return anim

if __name__=="__main__":
    main()
    
