"""Simulate duffing model with PyDSTool integrator.

See
    https://github.com/robclewley/pydstool/blob/master/examples/forced_spring.py

for forced-spring example.
"""

import numpy as np
import PyDSTool


def simulate(T, t_trans, steps, x0, v0, omega, p):
    x1, x2 = x0
    xd1, xd2 = v0
    C1, C2, m1, m2, c1, c2, c3, k1, k2, k3, a1, a2, a3 = p

    args = PyDSTool.args(name='duffing-stuffing')
    args.varspecs = {
        'x1': 'xd1',
        'x2': 'xd2',
        'xd1': 'C1*F(t)/m1 - xd1*(c1 + c3) + xd2*c3 - x1*(k1 + k3) + x2*k3 - a1*pow(x1, 3) + a3*pow(x2 - x1, 3)',
        'xd2': 'C2*F(t)/m2 - xd2*(c2 + c3) + xd1*c3 - x2*(k2 + k3) + x1*k3 - a2*pow(x2, 3) + a3*pow(x2 - x1, 3)'
    }
    args.fnspecs = {
        'F': (['t'], 'cos(omega*t)')
    }
    args.pars = {
        'omega': omega,
        'C1': C1, 'C2': C2,
        'm1': m1, 'm2': m2,
        'c1': c1, 'c2': c2, 'c3': c3,
        'k1': k1, 'k2': k2, 'k3': k3,
        'a1': a1, 'a2': a2, 'a3': a3
    }
    args.ics = {
        'x1': x0[0], 'x2': x0[1],  
        'xd1': v0[0], 'xd2': v0[1]  
    }
    args.xdomain = {
        'x1': [-10, 10], 'x2': [-10, 10],  
        'xd1': [-1e6, 1e6], 'xd2': [-1e6, 1e6]
    }
    #args.tdomain = [0, T]
    args.tdata = [0, T]

    # Integrator
    ds = PyDSTool.Generator.Vode_ODEsystem(args)
    traj = ds.compute('experiment')
    pts = traj.sample(dt=(T-t_trans)/steps, precise=True)
    return np.stack([
        pts['x1'][-steps:],
        pts['x2'][-steps:],
        pts['xd1'][-steps:],
        pts['xd2'][-steps:]
    ])


def test():
    C1 = 1.5e5
    C2 = 1.5e5
    m1 = 1.0 #1e-3
    m2 = 1.0 #1e-3

    c1 = 9647.403
    c2 = 10282.101
    c3 = 0.000
    k1 = 2326809592.681
    k2 = 2643039774.512
    k3 = 0.000
    a1 = 0.000
    a2 = 0.000
    a3 = 0.000

    p = (
        C1, C2,
        m1, m2,
        c1, c2, c3,
        k1, k2, k3,
        a1, a2, a3
    )

    omega = 2*np.pi*5e3

    x0, v0 = (0.00001, 0.00002), (0.000001, 0.0002)
    T = 0.6
    t_trans = 0.1
    steps = 100

    return simulate(T, t_trans, steps, x0, v0, omega, p)


if __name__ == '__main__':
    X = test()
    print(X.shape)
    print(X.mean(axis=1))
    print(X.std(axis=1))
    print(X)


#pars = {'D1': 1.0 , 'D2': 10.0, 'A': 2.1, 'B': 3.5, 'lambda': 5}
#
#tdomain = [0,50]
#
#icdict = {'X1': 1.087734, 'X2': 1.087734, 'X3': 4.124552, \
#		  'Y1': 1.8488063, 'Y2': 1.8488063, 'Y3': 1.0389936}
#
## Set up model
#X1str = 'D1*(X2 - 2*X1 + X3) + lambda*(A - (B+1)*X1 + pow(X1,2)*Y1)'
#X2str = 'D1*(X3 - 2*X2 + X1) + lambda*(A - (B+1)*X2 + pow(X2,2)*Y2)'
#X3str = 'D1*(X1 - 2*X3 + X2) + lambda*(A - (B+1)*X3 + pow(X3,2)*Y3)'
#Y1str = 'D2*(Y2 - 2*Y1 + Y3) + lambda*(B*X1 - pow(X1,2)*Y1)'
#Y2str = 'D2*(Y3 - 2*Y2 + Y1) + lambda*(B*X2 - pow(X2,2)*Y2)'
#Y3str = 'D2*(Y1 - 2*Y3 + Y2) + lambda*(B*X3 - pow(X3,2)*Y3)'
#
#DSargs = args(name='Brusselator')
#DSargs.tdomain = tdomain
#DSargs.pars = pars
#DSargs.varspecs = {'X1': X1str, 'X2': X2str, 'X3': X3str, 'Y1': Y1str, 'Y2': Y2str, 'Y3': Y3str}
#DSargs.xdomain = {'X1': [0, 50], 'X2': [0, 50], 'X3': [0, 50], 'Y1': [0, 50], 'Y2': [0, 50], 'Y3': [0, 50]}
#DSargs.ics = icdict
#
#testDS = Generator.Vode_ODEsystem(DSargs)
#
## Set up continuation class
#PyCont = ContClass(testDS)
#
#PCargs = args(name='EQ1', type='EP-C')
#PCargs.freepars = ['lambda']
#PCargs.StepSize = 5e-3
#PCargs.LocBifPoints = ['BP','LP']
#PCargs.verbosity = 2
#PCargs.SaveEigen = True
#PyCont.newCurve(PCargs)
#
#print('Computing curve...')
#start = clock()
#PyCont['EQ1'].forward()
#print('done in %.3f seconds!' % (clock()-start))
#
## Plot
#PyCont['EQ1'].display(('lambda','X1'), stability=True)
#PyCont.plot.toggleAll('off', byname='P1')
#show()