#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu May 21, 2020

@author: Prof Roy, U of Delaware


"""

## Import the necessary libraries/modules

import numpy as np
import matplotlib.pyplot as pypl
from scipy.integrate import odeint
from scipy.integrate import quad
from scipy.interpolate import interp1d
from scipy import optimize

# We define parameters:

b=0.04 ; # toast half length [m]
d= 0.004 ; # offset length [m]
h=0.75 ; # table height [m]
g= 9.81 ; # gravitational acceleration [m/s^2]
mu= 1.0; # static/kinetic friction coefficient

###################### PHASE 1 ##############################
#  Rotational phase: we take advantage of the conservation of
#  energy to reduce the problem to a simple quadrature.
#  First we solve for angle theta1, the incipient slip angle, the 
#  corresponding value of omega=omega1 and the end time t=t1
##############################################################

B= 1+ (6*d**2)/(b**2 + 3*d**2)
D= 1+ (-3*d**2)/(b**2 + 3*d**2)
theta1= np.arctan((mu*D)/B)
print ("theta1 = ",180*theta1/np.pi," degrees")

# conservation of energy gives theta_dot^2 = omega^2 * sin(theta)
# with omega:

omega = np.sqrt(6*g*d/(b**2 +3*d**2))

# find omega1 = theta_dot at t=t1  [rad/s]

omega1= omega*np.sqrt(np.sin(theta1))
print ("omega1 = ",omega1," rad/s")

# find time t=t1 such that theta(t1)=theta1


def NoSlip(theta):
    return 1/(omega*np.sqrt(np.sin(theta)) )

time1, err1 = quad(NoSlip, 0, theta1);
print ("time1 = ",time1," sec"," error= ",err1)

# solve for time t, 0 < t <t1 such that t=  (1/omega)* int_0^theta d phi/sqrt(sin(phi))

Theta1 = np.linspace(0, theta1, 100)
Time1= np.linspace(0, time1, 100)
for i in range(1, 100):
    Time1[i], err1 =  quad(NoSlip, 0, Theta1[i])
    if( err1 > 0.0001):
        print('warning! error = ',err1)
    Theta1[i]= 180*Theta1[i]/np.pi
    
pypl.plot(Time1,Theta1)  # plot theta vs time


###########         PHASE 2         ########################################
###  to study phase 2, we integrate the ODEs governing (r,theta)
###  with initial conditions r=d, r_dot= 0, theta= theta1, theta_dot= omega1
############################################################################

# define the system of 1st-order ODEs in the variables (r,rdot,theta,thetadot)

def SlipPhase(y, t):
    [r, rdot, theta, thetadot] = y
    Ftheta= (-2*rdot*thetadot +g*np.cos(theta))/(3.0*r**2+b**2)
    G= r*thetadot**2 +g*np.sin(theta)
    return [rdot, G-mu*b**2*Ftheta, thetadot, 3*r*Ftheta]

tmax=0.1  # [s]: max time of integration. 
          # This may need to be adjusted when mu -> 0 
           # choose tmax sufficiently large so that theta_ddot changes sign
n=2000     # number of integration points


t = np.linspace(0, tmax, n)
values = odeint(SlipPhase, [d, 0, theta1, omega1], t)

r=values[:,0]
rdot=values[:,1]
theta=values[:,2]
thetadot= values[:,3]

# to find the duration t= t2 of phase2 we need to find the value
# of t for which N=0. This is equivalent to finding the 
# value of t for which theta_ddot=0

# first we create a function which return theta_ddot

def ThetaDDot(R,Rdot,Theta,ThetaDot):
    return (-2*Rdot*ThetaDot +g*np.cos(Theta))/(3*R**2 + b**2)

# we create array thetaddot from the arrays r,rdot,theta,thetadot

thetaddot= ThetaDDot(r,rdot,theta,thetadot)

# we interpolate t vs thetaddot

TvsThetaddot = interp1d(thetaddot,t) # interpolation of Y=time vs X=theta_ddot

# we find t2= time2 and theta2=theta(t2), omega2= theta_dot(t2)
#         r2= r(t2) and rdot2= r_dot(t2)

time2=TvsThetaddot(0) # time t2 is found by setting theta_ddot=0
TvsTheta = interp1d(t,theta) # interpolation of X=time vs Y=theta
TvsThetaDot = interp1d(t,thetadot) # interpolation of X=time vs Y=theta_dot
TvsR = interp1d(t,r) # interpolation of X=time vs Y=r
TvsRdot = interp1d(t,rdot) # interpolation of X=time vs Y=rdot
theta2= TvsTheta(time2)
omega2= TvsThetaDot(time2)
r2=TvsR(time2)
rdot2=TvsRdot(time2)

print ("time2 = ",time2," s ")
print ("theta2 = ",180*theta2/np.pi," degrees")
print ("omega2 = ",omega2," rad/s")

Theta2 = np.linspace(theta1, theta2, 100)
Time2= np.linspace(0, time2, 100)

for i in range(0, 100):
    Theta2[i] =  180*TvsTheta(Time2[i])/np.pi
    Time2[i]=time1+Time2[i]

    
pypl.plot(Time2,Theta2)  # plot theta vs time

############### PHASE 3 ##################################
###  to study phase 3, we use the parabolic motion of mass
###  center G and the linear evolution of angle theta= omega2*t+theta2.
############################################################# 

# define position of G at t=t2

xGinit= r2*np.cos(theta2)
yGinit= r2*np.sin(theta2)
vGx= rdot2*np.cos(theta2) - r2*omega2*np.sin(theta2)
vGy= rdot2*np.sin(theta2) + r2*omega2*np.cos(theta2)

# define the path of G

def xG(T):
    return vGx*T+xGinit
def yG(T):
    return 0.5*g*T**2 + vGy*T + yGinit

def theta_free(T):
    return theta2 + omega2*T

# find the time of impact
    
def Target(T):
    yTip1= yG(T)+ b*np.sin(theta_free(T))
    yTip2= yG(T)- b*np.sin(theta_free(T))
    return max(yTip1,yTip2) -h

# solve for t=time3 by solving for the root of function Target

time3 = optimize.brentq(Target,0.001,1.5)
theta3= theta_free(time3)
print ("time3 = ",time3," sec")
print ("theta3 = ",180*theta3/np.pi,"degrees")

# decide whether landing is on butter side up or down

if(theta3 < 0.5*np.pi) or (theta3 > 1.5*np.pi):
    print ("Landing: butter side up")
if(theta3 > 0.5*np.pi) and (theta3 < 1.5*np.pi):
    print ("Landing: butter side down")

Theta3 = np.linspace(theta2, theta3, 100)
Time3 = np.linspace(0, time3, 100)

for i in range(0, 100):
    Theta3[i] =  180*theta_free(Time3[i])/np.pi
    Time3[i]= time1+ time2 + Time3[i]
    
pypl.plot(Time3,Theta3) # plot theta vs time

# plot boundary points between 3 phases

t_list = [0,time1,time2+time1,time1+time2+time3]
theta_list = [0,180*theta1/np.pi,180*theta2/np.pi,180*theta3/np.pi]
pypl.scatter(t_list, theta_list, s=10)
pypl.xlabel(r'$t$',fontsize=18)
pypl.ylabel(r'$theta$',fontsize=18)
pypl.grid()
pypl.title('mu = 1.0')
pypl.show()

# uncomment these lines if needed
pypl.savefig("toastplot-mu1.svg")
pypl.close()