# Mechanical and geometric constants of the woodpecker model
# according to
# Ch. Glocker: Dynamik von Starrkörpersystemen mit Reibung und Stössen
# PhD thesis TU München, July 1995, pp. 162ff
#
import numpy as np

from assimulo.problem import Implicit_Problem
from assimulo.solvers import IDA


mS = 3.0e-4 # Mass of sleeve [kg]
JS = 5.0e-9 # Moment of inertia of the sleeve [kgm]
mB = 4.5e-3 # Mass of bird [kg]
masstotal=mS+mB # total mass
JB = 7.0e-7 # Moment of inertia of bird [kgm]
r0 = 2.5e-3 # Radius of the bar [m]
rS = 3.1e-3 # Inner Radius of sleeve [m]
hS = 5.8e-3 # 1/2 height of sleeve [m]
lS = 1.0e-2 # verical distance sleeve origin to spring origin [m]
lG = 1.5e-2 # vertical distance spring origin to bird origin [m]
hB = 2.0e-2 # y coordinate beak (in bird coordinate system) [m]
lB = 2.01e-2 # -x coordinate beak (in bird coordinate system) [m]
cp = 5.6e-3 # rotational spring constant [N/rad]
g  = 9.81 #  [m/s^2]

to_state=[0,0,0]
beak_hits=0


#x=[z
#   ps
#   pb
#   zd
#   psd
#   pbd
#   l1
#   l2]

#sw = [state1,state2,state3]

def residual(t,x,xd,sw):
#    print("t: ",t)
    M = np.array([[masstotal, mB*lS, mB*lG],
               [mB*lS, JS+mB*lS**2, mB*lS*lG],
               [mB*lG, mB*lS*lG, JB+mB*lG**2]])
    ps=x[1]
    pb=x[2]
    l1=x[6]
    l2=x[7]
    biss=xd[3:6]
    res=np.zeros(8)
    if sw[0]: #State 1
        #print(1)
        res[0:5] = np.hstack((np.dot(M,biss)-[-(mS+mB)*g,cp*(pb-ps)-mB*lS*g,cp*(pb-ps)-mB*lG*g],l1,l2))
    if sw[1]: #State 2
        #print(2)
        res=np.zeros(8)
        res[0]=mB*lG*biss[2]+(mS+mB)*g+l2
        res[1]=mB*lS*lG*biss[2] - (cp*(pb-ps)-mB*lS*g-hS*l1-rS*l2)
        res[2]=(JB+mB*lG*lG)*biss[2] - (cp*(ps-pb)-mB*lG*g)
#        res[0:3] = np.hstack((np.dot(M,biss)-[-(mS+mB)*g-l2,cp*(pb-ps)-mB*lS*g-hS*l1-rS*l2,cp*(ps-pb)-mB*lG*g]))
        res[3]=(rS-r0)+hS*ps
        res[4]=x[3]+rS*x[4]
    if sw[2]: #State 3
#        print(3)
        res[0]=mB*lG*biss[2]+(mS+mB)*g+l2
        res[1]=mB*lS*lG*biss[2] - (cp*(pb-ps)-mB*lS*g+hS*l1-rS*l2)
        res[2]=(JB+mB*lG*lG)*biss[2] - (cp*(ps-pb)-mB*lG*g)
#        res[0:3] = np.hstack((np.dot(M,biss)-[-(mS+mB)*g-l2,cp*(pb-ps)-mB*lS*g+hS*l1-rS*l2,cp*(ps-pb)-mB*lG*g]))
        res[3]=(rS-r0)-hS*ps
        res[4]=x[3]#+rS*x[4]

    res[5]=x[3]-xd[0]
    res[6]=x[4]-xd[1]
    res[7]=x[5]-xd[2]
    #print("res: ",res)
    #print("x: ", x)
    return res

def state_events(t,x, xd,sw):
    """
    This is the function that keeps track of  events. When the sign
    of any of the functions changed, we have an event.
    """
    state_event=np.zeros(2)
    if sw[0]: #state 1:
        ps=x[1]
        state_event[0]=hS*ps+rS-r0
        state_event[1]=hS*ps-rS+r0
    if sw[1]: #state 2:
        l1=x[6]
        state_event[0]=l1
        state_event[1]=l1 # It has to be something ;)
    if sw[2]: #state 3:
        l1=x[6]
        pb=x[2]
        state_event[0]=l1
        state_event[1]=hB*pb-(lS+lG-lB-r0)
    return state_event


def handle_event(solver, event_info):
    """
    Event handling. This functions is called when Assimulo finds an event as
    specified by the event functions.
    """
    global to_state
    global beak_hits

    state_info = event_info[0]
    print("Handle event, state info: ", state_info, ", t: ", solver.t)
    if solver.sw[0]: #State 1:
        zd,psd,pbd=solver.y[3:6]
        if state_info[0]!=0 and pbd<0:
            print("Switch to state 2")
            solver.sw[1]=True
            solver.sw[0]=False
            solver.y[1]=-(rS-r0)/hS
            solver.y[3]=0
            solver.y[4]=0
            solver.y[5]=(mB*lG*zd+mB*lS*lG*psd)/(JB+mB*lG*lG)+pbd
            to_state[1]+=1
        elif state_info[1]!=0 and pbd>0: #Switch to state 3
            print("Switch to state 3")
            solver.sw[2]=True
            solver.sw[0]=False
            solver.y[1]=(rS-r0)/hS
            solver.y[3]=0
            solver.y[4]=0
            solver.y[5]=(mB*lG*zd+mB*lS*lG*psd)/(JB+mB*lG*lG)+pbd
            to_state[2]+=1
    elif solver.sw[1]: #State 2:
        print("Switch to state 1")
        solver.sw[0]=True
        solver.sw[1]=False
        to_state[0]+=1
    elif solver.sw[2]: #State 3:
        pbd=solver.y[5]
        if state_info[0]!=0 and pbd<0: #Switch to state 1
            print("Switch to state 1")
            solver.sw[0]=True
            solver.sw[2]=False
            to_state[0]+=1
        elif state_info[1]!=0 and pbd>0: #change sign of pbd
            print("Change sign of pbd")
            solver.y[5]=-solver.y[5]
            solver.yd[2]=solver.y[5]
            beak_hits+=1
    solver.yd[0:3]=solver.y[3:6]

#Initial values
x0 = np.zeros(8)
x0[0]=1
x0[5]=0.1
x0[2]=0.3
xd0 = np.zeros(8)
xd0[0:3]=x0[3:6]
t0 = 0.0
switches0=[True]+2*[False]

#Create an Assimulo Problem
problem = Implicit_Problem(residual, x0, xd0, t0, sw0=switches0)

problem.state_events = state_events #Sets the state events to the problem
problem.handle_event = handle_event #Sets the event handling to the problem
problem.name = 'Wood Pecker'   #Sets the name of the problem

#Create an Assimulo solver (CVode)
sim = IDA(problem)
sim.atol=6*[1e-5]+2*[1]
sim.algvar=6*[0]+2*[1]
#sim.suppress_alg=True

#Simulation
tfinal = 0.5#Final time

sim.simulate(tfinal) #Simulate
plotvar=8*[0]
plotvar[0]=1
sim.plot(plotvar)
print("Beak hits: ", beak_hits)
print("To state count:", to_state)