import matplotlib.animation as ani
import matplotlib.pyplot as plt
import numpy as np
from matplotlib import cm

# pydrake imports


def create_animation(J_grid, pi_grid, mesh):
    # initialize figure for animation plus misc settings
    fig, ax = plt.subplots(1, 2, figsize=(12, 5))
    plt.tight_layout()
    ax[0].axis("off")
    ax[1].axis("off")

    # cost-to-go plot and policy plots
    # overwrites the subplots axes
    ax[0] = fig.add_subplot(121, projection="3d")
    ax[1] = fig.add_subplot(122, projection="3d")
    ax[0].set_title(r"Optimal cost to go   $J^*(\mathbf{x})$")
    ax[1].set_title(r"Optimal policy   $\pi^*(\mathbf{x})$")
    for axi in ax:
        axi.set_xlabel(r"$q$")
        axi.set_ylabel(r"$\dot q$")

    # helper function for the surface plot
    Q, Qdot = np.meshgrid(mesh["q_grid"], mesh["qdot_grid"])

    def plot_surf(ax, Z):
        return ax.plot_surface(Q, Qdot, Z.T, rstride=1, cstride=1, cmap=cm.jet)

    # first frame of the animation
    J_surf = [plot_surf(ax[0], J_grid[:, :, 0])]
    pi_surf = [plot_surf(ax[1], pi_grid[:, :, 0])]

    # video parameters
    frames = 40  # total number of frames
    duration = 10  # seconds
    interval = duration / (frames - 1) * 1000  # milliseconds

    # initialize title to be modified in the callback
    title = fig.text(
        0.5, 0.95, "", fontsize="x-large", bbox={"facecolor": "w"}, ha="center"
    )

    # callback function for the animation
    def update_surf(frame, J_grid, J_surf, pi_grid, pi_surf):
        # iteration to plot, - 1 for base zero
        iters = J_grid.shape[2]
        i = int(frame * (iters - 1) / (frames - 1))

        # update cost-to-go and policy
        J_surf[0].remove()
        pi_surf[0].remove()
        J_surf[0] = plot_surf(ax[0], J_grid[:, :, i])
        pi_surf[0] = plot_surf(ax[1], pi_grid[:, :, i])

        # update title with current iteration
        # use base 1 as above
        title.set_text(f"Value iteration {i+1}")

    # create animation
    animate = ani.FuncAnimation(
        fig,
        update_surf,
        frames=frames,
        interval=interval,
        fargs=(J_grid, J_surf, pi_grid, pi_surf),
    )

    # play video
    plt.close()  # close any open figure
    return animate


# function that plots the trajectory of the
# double integrator in state space (q_sim vs qdot_sim)
def plot_state_trajectory(q_sim, qdot_sim):
    # draw a white dot for the initial conditions
    plt.scatter(
        q_sim[0],
        qdot_sim[0],
        s=100,
        edgecolor="k",
        c="w",
        zorder=3,
        label=r"$\mathbf{x}(0)$",
    )

    # black line for the trajectory in time
    plt.plot(q_sim, qdot_sim, color="k", linewidth=2, label=r"$\mathbf{x}(t)$")

    # misc
    plt.xlabel(r"$q$")
    plt.ylabel(r"$\dot q$")
    plt.legend()


# helper function for plot_policy that evaluates
# the controller output at a given state
def evaluate_controller(controller, q, qdot):
    # get context and set system output (= system state)
    context = controller.CreateDefaultContext()
    controller.get_input_port(0).FixValue(context, (q, qdot))

    # compute input for the double integrator
    u = controller.get_output_port(0).Eval(context)[0]

    return u


# function that produces a level plot
# the policy generated by the passed controller
# it used the grid defined by q_grid and qdot_grid
def plot_policy(q_grid, qdot_grid, controller):
    # evaluate the policy on the grid
    Pi_grid = np.array(
        [
            [evaluate_controller(controller, q, qdot) for qdot in qdot_grid]
            for q in q_grid
        ]
    )

    # level plot the level function
    # note the transpose to align the policy to the grid
    plt.contourf(*np.meshgrid(q_grid, qdot_grid), Pi_grid.T, cmap=cm.jet)

    # add a bar with the color scale on the right
    plt.colorbar(label=r"$\pi^*(\mathbf{x})$")


# function that plots the control signal
# u_sim as a function of the time vector t
def plot_input(u_sim, t_sim, u_lim):
    # plot the bounds for the control signal
    plt.plot(
        t_sim,
        [u_lim[0]] * len(t_sim),
        c="r",
        linestyle="--",
        label="Input limits",
    )
    plt.plot(t_sim, [u_lim[1]] * len(t_sim), c="r", linestyle="--")

    # plot the control signal
    plt.plot(t_sim, u_sim, label="Input from simulation")

    # misc
    plt.xlabel(r"$t$")
    plt.ylabel(r"$u$")
    plt.xlim(min(t_sim), max(t_sim))
    plt.grid(True)
    plt.legend(loc=1)


# overall plot function for the state trajectory,
# controller policy, and input signal
def simulate_and_plot(
    q0, qdot0, sim_time, controller, u_lim, nq=201, nqdot=201, simulate=None
):
    # get trajectories
    q_sim, qdot_sim, u_sim, t_sim = simulate(q0, qdot0, sim_time, controller)

    # state figure
    plt.figure()
    plot_state_trajectory(q_sim, qdot_sim)

    # plot policy only in a square region that
    # contains tightly the trajectory

    # helper function that computes upper and
    # lower bounds, with a small frame, for the
    # passed signal s
    def frame_signal(s, frame=0.1):
        ds = (max(s) - min(s)) * frame
        return [min(s) - ds, max(s) + ds]

    # regrid state space for policy plot
    # this grid must be much finer than the
    # one used for value iteration
    q_grid = np.linspace(*frame_signal(q_sim), nq)
    qdot_grid = np.linspace(*frame_signal(qdot_sim), nqdot)

    # plot the policy from the passed controller
    plot_policy(q_grid, qdot_grid, controller)

    # plot input as a function of time
    plt.figure()
    plot_input(u_sim, t_sim, u_lim)
