# NOTE: this notebook requires ffmpeg, which can be installed on ubuntu with "sudo apt install ffmpeg"
# and on mac with "brew install ffmpeg"
import cvxpy as cp
import numpy as np
import matplotlib.pyplot as plt
import torch
from cvxpylayers.torch import CvxpyLayer
import matplotlib.patches as patches
from matplotlib import animation, rc
rc('animation', html='html5')

"""
Explanation: Optimizing Stiffness Constants
Suppose we have $n$ blocks, each of width $w$, at positions $x\in\mathbb{R}^n$.
The blocks are positioned between two walls at $0$ and $l$.
The leftmost block is connected to the left wall via a spring with stiffness coefficient $k_1$,
the rightmost block is connected to the right wall via a spring with stiffness coefficient $k_{n+1}$,
and the $i$th block is connected to the $(i+1)$th block via a spring with stiffness coefficient $k_{i+1}$.
The equilibrium position of all the blocks can be found by solving the optimization problem
\begin{equation}
\begin{array}{ll}
\mbox{minimize} & \frac{1}{2}k_1x_1^2 + \frac{1}{2}k_2(x_2-x_1)^2 + \ldots + \frac{1}{2}k_{n+1}(l-x_n)^2\[.2cm]
\mbox{subject to} & x_1 \geq w/2, \quad x_n \leq l - w/2, \
& x_i - x_{i-1} \geq w, \quad i=2,\ldots,n-1,
\end{array}
\label{eq:prob}
\end{equation}
with variable $x$ and solution denoted $x^\star$.
The objective is the potential energy of the system, and the constraints express
the fact that the blocks have a width, and cannot penetrate each other or the walls.
End of explanation
"""


n = 5
l = 1
w = .05

k = cp.Parameter(n + 1, nonneg=True)
x = cp.Variable(n + 2)
objective = cp.sum(cp.multiply(k, .5 * cp.square(cp.diff(x))))
constraints = [x[0] == 0, x[-1] == l] + [x[i] - x[i-1] >= w for i in range(1, n+2)]
prob = cp.Problem(cp.Minimize(objective), constraints)

"""
Explanation: We can easily set up this problem as a CVXPY problem, with parameter $k$ and variable $x$:
End of explanation
"""


prob_tch = CvxpyLayer(prob, [k], [x])

"""
Explanation: We can easily convert this problem to a torch CvxpyLayer in one line.
The layer maps stiffness coefficients to block positions, or $x^\star(k)$.
End of explanation
"""


torch.manual_seed(0)
k = torch.ones(n+1, dtype=torch.double, requires_grad=True)
x_targ, _ = torch.sort(torch.rand(n, dtype=torch.double))

"""
Explanation: Our goal is to tune the stiffness coefficients such that the block positions are close to target positions,
or solve the problem
\begin{equation}
\begin{array}{ll}
\mbox{minimize} & \|x^\star(k) - x^\mathrm{targ}\|_2^2,
\end{array}
\end{equation}
with variable $k\in\mathbb{R}^{n+1}$.
End of explanation
"""


x, = prob_tch(k)
x_np = x.detach().numpy()

fig, ax = plt.subplots()
ax.set_xlim((0, 1))
ax.set_ylim((0, w*4))
rects = []
for i in range(n):
    plt.axvline(x_targ[i].item(), c='k')
    rect = patches.Rectangle((x_np[i + 1] - w/2, 0), w, w, linewidth=1, edgecolor='k', facecolor='none')
    ax.add_patch(rect)
    rects.append(rect)

def animate(i):
    k.grad = None
    x, = prob_tch(k)
    x_np = x.detach().numpy()
    loss = (x[1:-1] - x_targ).pow(2).sum()
    loss.backward()
    k.data = k.data - .1 * k.grad.data
    k.data[k.data < 0] = 0.0
    for i in range(n):
        rects[i].xy = (x_np[i + 1] - w/2, 0)
    return rects

anim = animation.FuncAnimation(fig, animate,
                               frames=300, interval=50, blit=True)
anim

"""
Explanation: We can (approximately) solve this problem via gradient descent:
End of explanation
"""


