import matplotlib.animation as ani
import matplotlib.pyplot as plt
import numpy as np
from IPython.display import HTML
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D
# pydrake imports
from pydrake.all import (DiagramBuilder, DynamicProgrammingOptions,
                         FittedValueIteration, LinearSystem, LogVectorOutput,
                         Simulator, VectorSystem)


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(.5,
                     .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=.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)
