import marimo

__generated_with = "0.6.24-dev2"
app = marimo.App()


@app.cell(hide_code=True)
def __():
    from typing import Tuple

    import cvxpy as cp
    import matplotlib.pyplot as plt
    import numpy as np
    from matplotlib import cm
    return Tuple, cm, cp, np, plt


@app.cell
def __(mo):
    mo.md("# Optimal trajectory planning")
    return


@app.cell(hide_code=True)
def __(mo):
    mo.md("### Initial position and velocity")
    return


@app.cell(hide_code=True)
def __(mo):
    p0_x = mo.ui.slider(0, 100, value=50, show_value=True, label="$p_0 \mid x$")
    p0_y = mo.ui.slider(0, 100, value=50, show_value=True, label="$p_0 \mid y$")
    p0_z = mo.ui.slider(0, 200, value=100, show_value=True, label="$p_0 \mid z$")
    v0_x = mo.ui.slider(-20, 20, value=-10, show_value=True, label="$v_0 \mid x$")
    v0_y = mo.ui.slider(-20, 20, value=0, show_value=True, label="$v_0 \mid y$")
    v0_z = mo.ui.slider(-20, 20, value=-10, show_value=True, label="$v_0 \mid z$")

    mo.hstack([mo.vstack([p0_x, p0_y, p0_z]), mo.vstack([v0_x, v0_y, v0_z])])
    return p0_x, p0_y, p0_z, v0_x, v0_y, v0_z


@app.cell
def __(np, p0_x, p0_y, p0_z, v0_x, v0_y, v0_z):
    # Test data
    h = 1  # discretization step in s
    g = 0.1  # gravity in m/s^2
    m = 10.0  # mass in kg
    F_max = 10.0  # maximum thrust in Newton
    alpha = np.pi / 8  # glide angle in rad
    p0 = np.array([50, 50, 100])  # initial position in m
    v0 = np.array([-10, 0, -10])  # initial velocity in m/s
    p0 = np.array([p0_x.value, p0_y.value, p0_z.value])  # initial position in m
    v0 = np.array([v0_x.value, v0_y.value, v0_z.value])  # initial velocity in m/s
    p_target = np.array([0, 0, 0])  # target position in m
    gamma = 1.0  # fuel consumption coefficient
    K = 35  # number of discretization steps
    return F_max, K, alpha, g, gamma, h, m, p0, p_target, v0


@app.cell
def __(Tuple, cp, np):
    # Formulate the optimization problem


    def optimize_fuel(
        p_target: np.ndarray,
        g: float,
        m: float,
        p0: np.ndarray,
        v0: np.ndarray,
        K: int,
        h: float,
        F_max: float,
        alpha: float,
        gamma: float,
    ) -> Tuple[np.ndarray, np.ndarray]:
        """

        Minimize fuel consumption for a rocket to land on a target

        :param p_target: landing target in m
        :param g: gravitational acceleration in m/s^2
        :param m: mass in kg
        :param p0: position in m
        :param v0: velocity in m/s
        :param K: Number of discretization steps
        :param h: discretization step in s
        :param F_max: maximum thrust of engine in kg*m/s^2 (Newton)
        :param alpha: Glide path angle in radian
        :param gamma: converts fuel consumption to liters of fuel consumption
        :return: position and thrust over time
        """

        P_min = p_target[2]

        # Variables
        V = cp.Variable((K + 1, 3))  # velocity
        P = cp.Variable((K + 1, 3))  # position
        F = cp.Variable((K, 3))  # thrust

        # Constraints
        # Match initial position and initial velocity
        constraints = [
            V[0] == v0,
            P[0] == p0,
        ]

        # Keep height above P_min
        constraints += [P[:, 2] >= P_min]

        # Match final position and 0 velocity
        constraints += [
            V[K] == [0, 0, 0],
            P[K] == p_target,
        ]

        # Physics dynamics for velocity
        constraints += [V[1:, :2] == V[:-1, :2] + h * (F[:, :2] / m)]
        constraints += [V[1:, 2] == V[:-1, 2] + h * (F[:, 2] / m - g)]

        # Physics dynamics for position
        constraints += [P[1:] == P[:-1] + h / 2 * (V[:-1] + V[1:])]

        # Maximum thrust constraint
        constraints += [cp.norm(F, 2, axis=1) <= F_max]

        # Glide path constraint
        constraints += [P[:, 2] >= np.tan(alpha) * cp.norm(P[:, :2], axis=1)]

        fuel_consumption = gamma * cp.sum(cp.norm(F, axis=1))

        # Regularization
        height_regularization = cp.sum(cp.abs(P[:, 2] - P_min))

        xy_regularization = cp.sum(
            cp.abs(P[:, :2] - p_target[:2].reshape((1, -1)))
        )

        glide_violation = cp.pos(
            np.tan(alpha) * cp.norm(P[:, :2], axis=1) - P[:, 2]
        )
        glide_violation = cp.sum(glide_violation)

        gamma_height = 1
        gamma_xy = 1
        gamma_glide = 1000
        reg = (
            gamma_height * height_regularization
            + gamma_xy * xy_regularization
            + gamma_glide * glide_violation
        )

        problem = cp.Problem(cp.Minimize(fuel_consumption + reg), constraints)
        problem.solve()
        return P.value, F.value, V.value
    return optimize_fuel,


@app.cell
def __(F_max, K, alpha, g, gamma, h, m, optimize_fuel, p0, p_target, v0):
    # Solve the problem

    P_star, F_star, V_star = optimize_fuel(
        p_target, g, m, p0, v0, K, h, F_max, alpha, gamma
    )
    return F_star, P_star, V_star


@app.cell
def __(mo):
    mo.md(r"### Plot the trajectory")
    return


@app.cell(hide_code=True)
def __(mo):
    azim_slider = mo.ui.slider(-180, 180, label="azim", show_value=True, value=160)
    azim_slider
    return azim_slider,


@app.cell(hide_code=True)
def __(F_star, P_star, alpha, azim_slider, cm, np, plt):
    _fig = plt.figure()
    _ax = _fig.add_subplot(projection="3d")
    X = np.linspace(P_star[:, 0].min() - 10, P_star[:, 0].max() + 10, num=30)
    Y = np.linspace(P_star[:, 1].min() - 10, P_star[:, 1].max() + 10, num=30)
    X, Y = np.meshgrid(X, Y)
    Z = np.tan(alpha) * np.sqrt(X**2 + Y**2)
    _ax.plot_surface(
        X,
        Y,
        Z,
        rstride=1,
        cstride=1,
        cmap=cm.autumn,
        linewidth=0.1,
        alpha=0.7,
        edgecolors="k",
    )

    _ax = plt.gca()
    _ax.view_init(azim=azim_slider.value)
    _ax.plot(
        xs=P_star[:, 0], ys=P_star[:, 1], zs=P_star[:, 2], c="b", lw=2, zorder=5
    )

    _ax.quiver(
        P_star[:-1, 0],
        P_star[:-1, 1],
        P_star[:-1, 2],
        F_star[:, 0],
        F_star[:, 1],
        F_star[:, 2],
        zorder=5,
        color="black",
        length=2,
    )

    _ax.set_xlabel("x")
    _ax.set_ylabel("y")
    _ax.set_zlabel("z")
    return X, Y, Z


@app.cell
def __(mo):
    mo.md(r"### Plot the Position, Velocity and Thrust over time")
    return


@app.cell(hide_code=True)
def __(F_star, P_star, V_star, plt):
    _fig, _ax = plt.subplots(3, 1, sharex=True)

    _ax[0].plot(P_star[:, 0], label="x")
    _ax[0].plot(P_star[:, 1], label="y")
    _ax[0].plot(P_star[:, 2], label="z")
    _ax[0].set_ylabel("Position")
    _ax[0].legend()

    _ax[1].plot(F_star[:, 0], label="x")
    _ax[1].plot(F_star[:, 1], label="y")
    _ax[1].plot(F_star[:, 2], label="z")
    _ax[1].set_ylabel("Thrust")
    _ax[1].legend()

    _ax[2].plot(V_star[:, 0], label="x")
    _ax[2].plot(V_star[:, 1], label="y")
    _ax[2].plot(V_star[:, 2], label="z")
    _ax[2].set_ylabel("Velocity")
    _ax[2].legend()
    return


@app.cell
def __(F_max, K, alpha, g, gamma, h, m, optimize_fuel, p0, p_target, v0):
    # We can apply bisection on K in the function above.


    def trajectory_is_impossible(duration):
        P_star, F_star, V_star = optimize_fuel(
            p_target, g, m, p0, v0, duration, h, F_max, alpha, gamma
        )
        if P_star is None:
            return True
        else:
            return False


    min_T = 1
    max_T = K
    target_T = 1

    while min_T < max_T:
        if trajectory_is_impossible(target_T):
            min_T = target_T + 1
            target_T = (min_T + max_T) // 2
        else:
            max_T = target_T
            target_T = (min_T + max_T - 1) // 2
        print(f"{min_T=}\t{max_T=}\t{target_T=}")
    f"minimum time is {h * target_T} seconds"
    return max_T, min_T, target_T, trajectory_is_impossible


@app.cell
def __(mo):
    mo.md("# Question 2 Solutions")
    return


@app.cell
def __(mo):
    mo.md(
        r"""
        ## Part (a)

        Note that $\|u_i - c\|_2^2 = \|u_i\|_2^2 - 2 u_i^T c + \|c\|_2^2$. Using this substitution we can rewrite our loss function to be

        $$
        \sum_{i=1}^m \left(\|u_i - c\|_2^2 - r^2\right)^2 = 
        \sum_{i=1}^m \left(\|u_i\|_2^2 - 2 u_i^T c - (r^2 - \|c\|_2^2)\right)^2 = 
        \sum_{i=1}^m \left(\|u_i\|_2^2 - 2 u_i^T c - t\right)^2.
        $$

        This is convex, and in particular is just a least squares problem!
        """
    )
    return


@app.cell(hide_code=True)
def __(np):
    n_circle = 2
    m_circle = 50
    U = np.array(
        [
            [
                1.824183228637652032e00,
                1.349093690455489103e00,
                6.966316403935147727e-01,
                7.599387854623529392e-01,
                2.388321695850912363e00,
                8.651370608981923116e-01,
                1.863922545015865406e00,
                7.099743941474848663e-01,
                6.005484882320809570e-01,
                4.561429569892232472e-01,
                5.328296545713475663e-01,
                2.138547819234526415e00,
                1.906676474276197464e00,
                1.015547309536922516e00,
                8.765948388006337133e-01,
                1.648147347399247842e00,
                1.027902202451572045e00,
                2.145586297520478691e00,
                1.793440421753045744e00,
                1.020535583041398908e00,
                8.977911075271942654e-01,
                1.530480229262339398e00,
                2.478088034137528872e-01,
                2.617415807793897820e00,
                2.081978553098443374e00,
                1.891226687205936452e00,
                8.222497927065576251e-01,
                5.803514604868882376e-01,
                1.158670193449639063e00,
                6.016685032455900695e-01,
                5.605410828151705660e-01,
                2.508815467550573164e00,
                2.230201413385580977e00,
                1.170848897912992514e00,
                2.256355929901105561e00,
                6.686991510936428629e-01,
                2.040269595792217672e00,
                3.634166812924328749e-01,
                5.418647611079159265e-01,
                6.631470058399455692e-01,
                4.286142597532469622e-01,
                2.155925078996823618e00,
                2.379380016960549682e00,
                6.343212414048013947e-01,
                1.469076407947448981e00,
                1.225322035289937439e00,
                1.467602887401966871e00,
                9.345319187253748883e-01,
                1.985592768641736505e00,
                2.106896115090134636e00,
            ],
            [
                -9.644136284187876385e-01,
                1.069547315003422927e00,
                6.733229334437943470e-01,
                7.788072961810316164e-01,
                -9.467465278344706636e-01,
                -8.591303443863639311e-01,
                1.279527420871080956e00,
                5.314829019311283487e-01,
                6.975676079749143499e-02,
                -4.641873429414754559e-01,
                -2.094571396598311763e-01,
                -8.003479827938377866e-01,
                6.135280782546607137e-01,
                -9.961307468791747999e-01,
                -8.765215480412106297e-01,
                9.655406812422813179e-01,
                1.011230180540185541e00,
                6.105416770440197372e-01,
                9.486552370654932620e-01,
                -9.863592657836954825e-01,
                7.695327845100754516e-01,
                -1.060072365810699413e00,
                -4.041043465424410952e-01,
                -2.352952920283236105e-01,
                7.560391050507236921e-01,
                -9.454246095204003053e-01,
                -5.303145312191936966e-01,
                5.979590038743245461e-01,
                -1.154309511133019717e00,
                -6.123184171955468047e-01,
                -1.464683782538583889e-01,
                -1.839128688968104386e-01,
                4.250070477845909744e-01,
                8.861864983476224200e-01,
                3.927648421593328276e-01,
                -6.726102374256350824e-01,
                -1.047252884197514833e00,
                1.825096825995130845e-01,
                -4.482373962742914886e-01,
                5.115625649313135792e-01,
                7.846201103116770548e-02,
                6.006325432819290544e-01,
                -5.710733714464664157e-01,
                4.725559971890586075e-01,
                -8.440290321502940118e-01,
                -1.003920890712479475e00,
                -1.067089412136528637e00,
                7.909281966910661765e-01,
                -1.059509163675931065e00,
                -7.136351632325785843e-01,
            ],
        ]
    )
    return U, m_circle, n_circle


@app.cell
def __(U, cp, n_circle, np):
    c = cp.Variable(n_circle)
    t = cp.Variable()
    obj = cp.sum_squares(cp.square(cp.norm(U, axis=0)) - 2 * c @ U - t)
    prob = cp.Problem(cp.Minimize(obj))
    prob.solve()
    r = np.sqrt(t.value + np.linalg.norm(c.value) ** 2)
    f"The center is {c.value} and the radius is {r}"
    return c, obj, prob, r, t


@app.cell
def __(U, c, plt, r):
    fig, ax = plt.subplots()
    ax.scatter(U[0, :], U[1, :], edgecolor="k", color="w")
    cir = plt.Circle(
        c.value, r, color="k", fill=False, linestyle="dashed", alpha=0.5
    )
    ax.set_aspect("equal", adjustable="box")
    ax.add_patch(cir)
    plt.show()
    return ax, cir, fig


@app.cell
def __():
    import marimo as mo
    return mo,


if __name__ == "__main__":
    app.run()