from __future__ import division, print_function
import itertools
import numpy as np
import sympy as sp
import matplotlib.pyplot as plt
#from pyodesys.native.gsl import NativeGSLSys as SymbolicSys
from pyodesys.native.cvode import NativeCvodeSys as SymbolicSys
sp.init_printing()
%matplotlib inline
print(sp.__version__)

"""
Explanation: Van der Pol oscillator
We will look at the second order differentual equation (see https://en.wikipedia.org/wiki/Van_der_Pol_oscillator):
$$
{d^2y_0 \over dx^2}-\mu(1-y_0^2){dy_0 \over dx}+y_0= 0
$$
End of explanation
"""


vdp1 = lambda x, y, p: [y[1], -y[0] + p[0]*y[1]*(1 - y[0]**2)]
mu_val = 2.5
y0_1 = [0.0, 1.0]
y0_1, (y0_1[0], vdp1(0, y0_1, [mu_val])[0])

"""
Explanation: Note that we imported NativeCvodeSys as SymbolicSys, this speed up the time of integration by more than an order of magnitude due to using compiled C++ code for our mathematical expressions.
One way to reduce the order of our second order differential equation is to formulate a system of first order ODEs, using:
$$ y_1 = \dot y_0 $$
which gives us:
$$
\begin{cases}
\dot y_0 = y_1 \
\dot y_1 = \mu(1-y_0^2) y_1-y_0
\end{cases}
$$
Let's call this system of ordinary differential equations vdp1:
End of explanation
"""


transf = lambda y, dydx, p: [y[0], y[0] - y[0]**3/3 - dydx[0]/p[0]]
x, mu = sp.symbols('x mu', real=True)
y = [yi(x) for yi in sp.symbols('y:2', cls=sp.Function)]
dydx = [yi.diff(x) for yi in y]
[sp.Eq(yi, expr, evaluate=False) for yi, expr in zip(y, transf(y, dydx, [mu]))]  # Just for displaying

"""
Explanation: An alternative would be to use use the Liénard transformation:
$$ y = x - x^3/3 - \dot x/\mu $$
End of explanation
"""


vdp2 = lambda x, y, p: [p[0]*(y[0] - y[0]**3/3 - y[1]), y[0]/p[0]]
calc_y0_2 = lambda y0, mu: transf(y0, vdp1(0, y0, [mu]), [mu])
y0_2 = calc_y0_2(y0_1, mu_val)
(y0_2, y0_2[0], vdp2(0, y0_2, [mu_val])[0])

def solve_and_plot(odesys, y0, tout, mu, indices=None, integrator='native', **kwargs):
    plt.figure(figsize=(16, 4))
    xout, yout, info = odesys.integrate(tout, y0, [mu], integrator=integrator, **kwargs)
    plt.subplot(1, 2, 1)
    odesys.plot_result(indices=indices, ls=('-',), c=('k', 'r'))
    plt.legend(loc='best')
    plt.subplot(1, 2, 2)
    odesys.plot_phase_plane()
    info.pop('internal_xout')  # too much output
    info.pop('internal_yout')
    return len(xout), info

tend = 25

odesys1 = SymbolicSys.from_callback(vdp1, 2, 1, names='y0 y1'.split())
odesys1.exprs

for mu in [0, 3, 9]:
    solve_and_plot(odesys1, y0_1, np.linspace(0, tend, 500), mu)

"""
Explanation: which gives us (we could generate this result using SymPy):
$$
\begin{cases}
\dot y_0 = \mu \left(y_0-\frac{1}{3}y_0^3-y_1\right) \
\dot y_1 = \frac{1}{\mu} y_0
\end{cases}
$$
End of explanation
"""


tau = lambda mu: 1.6137056388801094*mu + 4.676*mu**(-1./3)
for mu in [20, 40, 60]:
    solve_and_plot(odesys1, y0_1, np.linspace(0, 5*tau(mu), 500), mu)

"""
Explanation: As we see, the period ($\tau$) varies with $\mu$, in 1952 Mary Cartwright derived an approximate formula for $\tau$ (valid for large $\mu$):
$$
\tau  = (3 - 2 \ln 2)\mu + 2 \alpha \mu^{-1/3}
$$
where $\alpha \approx 2.338$
End of explanation
"""


odesys2 = SymbolicSys.from_callback(vdp2, 2, 1, names='y0 y1'.split())
odesys2.exprs

solve_and_plot(odesys2, y0_2, tend, mu_val, nsteps=2000)

"""
Explanation: For larger values of $\mu$ we run into trouble (the numerical solver fails).
The phase portrait is not well resolved due to rapid variations in y1.
Let us look at our alternative formulation:
End of explanation
"""


ls = itertools.cycle(('-', '--', ':'))
for mu in [84, 160, 320]:
    y0_2 = calc_y0_2(y0_1, mu)
    print(y0_2)
    solve_and_plot(odesys2, y0_2, np.linspace(0, 5*tau(mu), 500), mu)

"""
Explanation: This looks much better. Let's see if the solver has an easier time dealing with this formulation of y2 for large values of $\mu$:
End of explanation
"""


solve_and_plot(odesys2, calc_y0_2(y0_1, mu_val), tend, mu_val, nsteps=2000)
J = odesys2.get_jac()
J

"""
Explanation: Indeed it has.
Stiffness
Let us compare the performance of explicit and implicit steppers (Adams and BDF respectivecly) for varying values of $\mu$
End of explanation
"""


odesys2._NativeCode._written_files

symbs = odesys2.dep + tuple(odesys2.params)
symbs

Jeig = J.eigenvals().keys()
eig_cbs = [sp.lambdify(symbs, eig, modules='numpy') for eig in Jeig]
Jeig

eigvals = np.array([(eig_cbs[0](*(tuple(yvals)+(mu_val + 0j,))),
                     eig_cbs[1](*(tuple(yvals)+(mu_val + 0j,)))) for yvals in odesys2._internal[1]])

plt.plot(odesys2._internal[0], odesys2.stiffness(), label='from SVD')
plt.plot(odesys2._internal[0], np.abs(eigvals[:,0])/np.abs(eigvals[:,1]), label='analytic')
plt.legend()

"""
Explanation: For this simple system we can afford calculating the eigenvalues analytically
End of explanation
"""


def arr_to_wav(arr, rate=44100):
    from IPython.display import Audio
    from scipy.io.wavfile import write
    scaled = np.int16(arr/np.max(np.abs(arr)) * 32767)
    write('test.wav', rate, scaled)
    return Audio('test.wav')

xout, yout, info = odesys2.integrate(np.linspace(0, 500*tau(40.0), 2*44100), y0_1, [40.0], integrator='native')
arr_to_wav(yout[:, 0])

def overlay(tend_mu, odesys=odesys2, time=3, rate=44100, plot=False):
    yout_tot = None
    for tend, mu in tend_mu:
        xout, yout, info = odesys.integrate(np.linspace(0, tend*tau(mu[0]), time*rate), y0_1, mu, integrator='native')
        print(tend, mu, tend*tau(mu[0]))
        if yout_tot is None:
            yout_tot = yout[:, 0]
        else:
            yout_tot += yout[:, 0]
    if plot:
        plt.figure(figsize=(16,4))
        plt.plot(yout_tot[slice(None) if plot is True else slice(0, plot)])
    return arr_to_wav(yout_tot, rate=rate)

overlay([
    (400, [2.0]),
    (410, [2.1]),
], plot=10000)

"""
Explanation: Audio
Plotting is instructive from a mathematical standpoint, but these equations were often investigated by to listening to audio amplified by electrical circuits modelled by the equation. So let us generate some audio.
End of explanation
"""


vdp_forced = lambda x, y, p: [y[1], p[1]*sp.sin(p[2]*x) - y[0] + p[0]*y[1]*(1 - y[0]**2)]
odesys_forced = SymbolicSys.from_callback(vdp_forced, 2, 3)
overlay([(700, [8, 1, 0.5])], odesys_forced, plot=5000)  # Non-chaotic behavior

overlay([(700, [8, 1.2, 0.6])], odesys_forced, plot=5000)  # Chaotic behavior

"""
Explanation: Forced van der pol oscillator
End of explanation
"""


vdp_transient = lambda x, y, p: [y[1], - y[0] + p[0]*sp.exp(-p[1]*x)*y[1]*(1 - y[0]**2)]
odesys_transient = SymbolicSys.from_callback(vdp_transient, 2, 2)
odesys_transient.exprs

overlay([
    (440, [0.1, 1/2500.]),
    (445, [0.5, 1/1000.]),
    (890, [0.1, 2/2500.]),
    (896, [0.5, 2/1000.]),
], odesys_transient, plot=-1)

odesys2._native._written_files

"""
Explanation: Transient $\mu$
End of explanation
"""


