#!/usr/bin/env python
# coding: utf-8

# Hysteresis mechanism created by bistability of states.
# 
# Energy function:
# $f = u^4 - 2u^2 + hu$
# 
# Find video here: http://www.yairmau.com/python

# In[1]:


# comment these lines if you want interactive mode,
# i.e., if you want to see the animation in real time.
import matplotlib
matplotlib.use('Agg')


# In[2]:


import matplotlib.pyplot as plt
import numpy as np
import os
import sympy
from scipy.integrate import ode

# learn how to configure: http://matplotlib.sourceforge.net/users/customizing.html
params = {#'backend': 'GTKAgg',
          'legend.handlelength': 2.5,
          'legend.borderaxespad': 0,
          'font.family':'serif',
          'font.size': 18,
          'font.serif':['Times'], # Times, Palatino, New Century Schoolbook, Bookman, Computer Modern Roman
          'ps.usedistiller': 'xpdf',
          'text.usetex': True,
          }

plt.rcParams.update(params)
fig=plt.figure(1,figsize=(9.6,5.4),dpi=100) # 1920x1080 # figsize accepts only inches. if you rather think in cm, change the code yourself.
fig.clf()
fig.subplots_adjust(left=0.07, right=0.93,top=0.90, bottom=0.12,hspace=0.02,wspace=0.10)

Hlim=2.5 # parameter range from -Hlim to Hlim

ax1=fig.add_subplot(121)
ax1.set_xticks([])
ax1.set_yticks([])
ax1.set_xlabel(r'System response',labelpad=12)
ax1.set_ylabel('Energy',labelpad=12)
ax1.axis([-Hlim,Hlim,-5,5])

ax2=fig.add_subplot(122)
ax2.set_xticks([])
ax2.set_yticks([])
ax2.set_xlabel(r'Parameter',labelpad=12)
ax2.set_ylabel(r'System response',labelpad=12)
ax2.yaxis.set_label_position("right")
ax2.axis([-Hlim*1.2,Hlim*1.2,-2,2])

frame_names = []
frame_index = 0
make_movie=True
plt.ion()


# In[3]:


# energy function and its derivative
f = lambda u,h: u**4-2*u**2+h*u
fprime = lambda u,h: sympy.diff(f(u,h),u)

Hinit=Hlim
ulim=2.5 # system response axis, from -ulim to ulim
u = np.linspace(-ulim,ulim,101)

x = sympy.Symbol('x')

def res(h):
    """System response is one of the real roots
       of the energy function derivative
    """
    # derivative roots, complex
    resp = sympy.solvers.solve(fprime(x,h),x)
    # numerical evaluation
    resp = map(sympy.N,resp)
    # let's check which roots are real
    isreal = len(resp)*[False]
    for i in range(len(resp)):
        # negligible imaginary component
        if np.abs(sympy.functions.im(resp[i]))<1e-15:
            resp[i]=sympy.functions.re(resp[i])
            isreal[i]=True
    resp = np.array(resp)
    # return only real roots
    return resp[np.array(isreal)]

# let's plot stuff, and make a nice movie
#### left plot, ax1 ####
line_func, = ax1.plot(u,f(u,Hinit),lw=2,color='black')
# ball color
ball_color = "blue"
# minimum = the smallest root, the leftmost root
mini = np.min(res(Hinit)) # calculated for initial parameter value
boost = 0.22 # so that ball sits on top of the curve
# plot ball
ball_u, = ax1.plot([mini],[f(mini,Hinit)+boost],'o',
                 markersize=12, markerfacecolor=ball_color)

#### right plot, ax2 ####
# build empty hysteresis array, we will add values
# as simulation progresses
deetype = np.dtype([('h', 'float64'), ('u', 'float64')])
hysteresis = np.array([(Hinit,mini)],dtype=deetype)
line_hyst, = ax2.plot(hysteresis['h'],hysteresis['u'], lw=2,color='black')
ballH, = ax2.plot([hysteresis['h'][-1]],[hysteresis['u'][-1]],'o',
                  markersize=12, markerfacecolor=ball_color)
plt.show()


# In[4]:


# time to simulate
Total_time = 15 # seconds
fps = 24 # frames per second

# divided by 2 because we ramp down then up
param_vec = np.linspace(Hlim,-Hlim,Total_time*fps/2)

# ramp down
for H in param_vec:
    line_func.set_data(u,f(u,H)) # update line on the left
    mini = np.min(res(H)) # calculate new minimum
    ball_u.set_data([mini],[f(mini,H)+boost]) # update ball on the left
    new_line = np.array([(H,mini)],dtype=deetype) # create new line
    # append new line to hysteresis array
    hysteresis = np.concatenate((hysteresis,new_line))
    line_hyst.set_data(hysteresis['h'],hysteresis['u']) # update line
    ballH.set_data([hysteresis['h'][-1]],[hysteresis['u'][-1]]) # update ball on the right
    fig.canvas.draw()
    if make_movie:
        fname = "_tmp{:05d}.png".format(frame_index)
        frame_names.append(fname)
        fig.savefig(fname,dpi=200)
        frame_index+=1

# ramp up
for H in param_vec[::-1]: # just reverse parameter array
    line_func.set_data(u,f(u,H))
    maxi = np.max(res(H)) # everything is the same, but now with maximum
    ball_u.set_data([maxi],[f(maxi,H)+boost])
    new_line = np.array([(H,maxi)],dtype=deetype)
    hysteresis = np.concatenate((hysteresis,new_line))
    line_hyst.set_data(hysteresis['h'],hysteresis['u'])
    ballH.set_data([hysteresis['h'][-1]],[hysteresis['u'][-1]])
    fig.canvas.draw()
    if make_movie:
        fname = "_tmp{:05d}.png".format(frame_index)
        frame_names.append(fname)
        fig.savefig(fname,dpi=200)
        frame_index+=1

if make_movie:
    frames = "_tmp%5d.png"
    movie_command = "ffmpeg -y -r {:} -i {:} ball.mp4".format(fps,frames)
    os.system(movie_command)

    for fname in frame_names:
        # pass
        os.remove(fname)