import torch


def double_integrator_lqr(Q, R):
    """
    Compute the LQR controller for double integrator. The LQR controller is
    u = -R⁻¹BᵀPx
    where P is the solution to the Riccati equation
    AᵀP+PA-PBR⁻¹BᵀP + Q = 0
    @param Q a positive definite symmetrix 2 x 2 matrix.
    @param R a positive scalar.
    @return (K, P) K is the control gain, K = R⁻¹BᵀP, P is the cost-to-go
    matrix.
    """
    assert (Q.dtype == R.dtype)
    assert (Q.shape[0] == 2)
    assert (Q.shape[1] == 2)
    assert (R > 0)
    assert (Q[0][1] == Q[1][0])
    """
    The LQR cost-to-go matrix P for the double integrator system satisfies the
    following equation
    R * [Q[0][0]           Q[0][1] + P[0][0]] = [P[0][1]², P[0][1]P[1][1]]
        [Q[0][1] + P[0][0] 2P[0][1]+ Q[1][1]]   [P[0][1]P[1][1], P[1][1]²]
    """
    P = torch.empty(2, 2, dtype=Q.dtype)
    P[0][1] = torch.sqrt(R * Q[0][0])
    P[1][1] = torch.sqrt(R * (2 * P[0][1] + Q[1][1]))
    P[0][0] = (P[0][1] * P[1][1] - R * Q[0][1]) / R
    if (P[0][0] < 0 or P[0][0] * P[1][1] < P[0][1] * P[0][1]):
        # P is not positive definite
        P[0][1] = -P[0][1]
        P[1][1] = torch.sqrt(R * (2 * P[0][1] + Q[1][1]))
        P[0][0] = (P[0][1] * P[1][1] - R * Q[0][1]) / R
    P[1][0] = P[0][1]
    K = (torch.tensor([[0, 1]], dtype=Q.dtype) @ P) / R
    return (K, P)