# A demo of gradients through scipy.integrate.odeint,
# estimating the dynamics of a system given a trajectory.


import matplotlib.pyplot as plt
import numpy as npo

import autograd.numpy as np
import autograd.numpy.random as npr
from autograd import grad
from autograd.builtins import tuple
from autograd.misc.optimizers import adam
from autograd.scipy.integrate import odeint

N = 30  # Dataset size
D = 2  # Data dimension
max_T = 1.5


# Two-dimensional damped oscillator
def func(y, t0, A):
    return np.dot(y**3, A)


def nn_predict(inputs, t, params):
    for W, b in params:
        outputs = np.dot(inputs, W) + b
        inputs = np.maximum(0, outputs)
    return outputs


def init_nn_params(scale, layer_sizes, rs=npr.RandomState(0)):
    """Build a list of (weights, biases) tuples, one for each layer."""
    return [
        (
            rs.randn(insize, outsize) * scale,  # weight matrix
            rs.randn(outsize) * scale,
        )  # bias vector
        for insize, outsize in zip(layer_sizes[:-1], layer_sizes[1:])
    ]


# Define neural ODE model.
def ode_pred(params, y0, t):
    return odeint(nn_predict, y0, t, tuple((params,)), rtol=0.01)


def L1_loss(pred, targets):
    return np.mean(np.abs(pred - targets))


if __name__ == "__main__":
    # Generate data from true dynamics.
    true_y0 = np.array([2.0, 0.0]).T
    t = np.linspace(0.0, max_T, N)
    true_A = np.array([[-0.1, 2.0], [-2.0, -0.1]])
    true_y = odeint(func, true_y0, t, args=(true_A,))

    def train_loss(params, iter):
        pred = ode_pred(params, true_y0, t)
        return L1_loss(pred, true_y)

    # Set up figure
    fig = plt.figure(figsize=(12, 4), facecolor="white")
    ax_traj = fig.add_subplot(131, frameon=False)
    ax_phase = fig.add_subplot(132, frameon=False)
    ax_vecfield = fig.add_subplot(133, frameon=False)
    plt.show(block=False)

    # Plots data and learned dynamics.
    def callback(params, iter, g):
        pred = ode_pred(params, true_y0, t)

        print(f"Iteration {iter:d} train loss {L1_loss(pred, true_y):.6f}")

        ax_traj.cla()
        ax_traj.set_title("Trajectories")
        ax_traj.set_xlabel("t")
        ax_traj.set_ylabel("x,y")
        ax_traj.plot(t, true_y[:, 0], "-", t, true_y[:, 1], "g-")
        ax_traj.plot(t, pred[:, 0], "--", t, pred[:, 1], "b--")
        ax_traj.set_xlim(t.min(), t.max())
        ax_traj.set_ylim(-2, 2)
        ax_traj.xaxis.set_ticklabels([])
        ax_traj.yaxis.set_ticklabels([])
        ax_traj.legend()

        ax_phase.cla()
        ax_phase.set_title("Phase Portrait")
        ax_phase.set_xlabel("x")
        ax_phase.set_ylabel("y")
        ax_phase.plot(true_y[:, 0], true_y[:, 1], "g-")
        ax_phase.plot(pred[:, 0], pred[:, 1], "b--")
        ax_phase.set_xlim(-2, 2)
        ax_phase.set_ylim(-2, 2)
        ax_phase.xaxis.set_ticklabels([])
        ax_phase.yaxis.set_ticklabels([])

        ax_vecfield.cla()
        ax_vecfield.set_title("Learned Vector Field")
        ax_vecfield.set_xlabel("x")
        ax_vecfield.set_ylabel("y")
        ax_vecfield.xaxis.set_ticklabels([])
        ax_vecfield.yaxis.set_ticklabels([])

        # vector field plot
        y, x = npo.mgrid[-2:2:21j, -2:2:21j]
        dydt = nn_predict(np.stack([x, y], -1).reshape(21 * 21, 2), 0, params).reshape(-1, 2)
        mag = np.sqrt(dydt[:, 0] ** 2 + dydt[:, 1] ** 2).reshape(-1, 1)
        dydt = dydt / mag
        dydt = dydt.reshape(21, 21, 2)

        ax_vecfield.streamplot(x, y, dydt[:, :, 0], dydt[:, :, 1], color="black")
        ax_vecfield.set_xlim(-2, 2)
        ax_vecfield.set_ylim(-2, 2)

        fig.tight_layout()
        plt.draw()
        plt.pause(0.001)

    # Train neural net dynamics to match data.
    init_params = init_nn_params(0.1, layer_sizes=[D, 150, D])
    optimized_params = adam(grad(train_loss), init_params, num_iters=1000, callback=callback)
