# kincar.py - planar vehicle model (with flatness)
# RMM, 16 Jan 2022

import numpy as np
import matplotlib.pyplot as plt
import control as ct
import control.flatsys as fs

#
# Vehicle dynamics (bicycle model)
#

# Function to take states, inputs and return the flat flag
def _kincar_flat_forward(x, u, params={}):
    # Get the parameter values
    b = params.get('wheelbase', 3.)
    #! TODO: add dir processing

    # Create a list of arrays to store the flat output and its derivatives
    zflag = [np.zeros(3), np.zeros(3)]

    # Flat output is the x, y position of the rear wheels
    zflag[0][0] = x[0]
    zflag[1][0] = x[1]

    # First derivatives of the flat output
    zflag[0][1] = u[0] * np.cos(x[2])  # dx/dt
    zflag[1][1] = u[0] * np.sin(x[2])  # dy/dt

    # First derivative of the angle
    thdot = (u[0]/b) * np.tan(u[1])

    # Second derivatives of the flat output (setting vdot = 0)
    zflag[0][2] = -u[0] * thdot * np.sin(x[2])
    zflag[1][2] =  u[0] * thdot * np.cos(x[2])

    return zflag

# Function to take the flat flag and return states, inputs
def _kincar_flat_reverse(zflag, params={}):
    # Get the parameter values
    b = params.get('wheelbase', 3.)
    dir = params.get('dir', 'f')

    # Create a vector to store the state and inputs
    x = np.zeros(3)
    u = np.zeros(2)

    # Given the flat variables, solve for the state
    x[0] = zflag[0][0]  # x position
    x[1] = zflag[1][0]  # y position
    if dir == 'f':
        x[2] = np.arctan2(zflag[1][1], zflag[0][1])  # tan(theta) = ydot/xdot
    elif dir == 'r':
        # Angle is flipped by 180 degrees (since v < 0)
        x[2] = np.arctan2(-zflag[1][1], -zflag[0][1])
    else:
        raise ValueError("unknown direction:", dir)

    # And next solve for the inputs
    u[0] = zflag[0][1] * np.cos(x[2]) + zflag[1][1] * np.sin(x[2])
    thdot_v = zflag[1][2] * np.cos(x[2]) - zflag[0][2] * np.sin(x[2])
    u[1] = np.arctan2(thdot_v, u[0]**2 / b)

    return x, u

# Function to compute the RHS of the system dynamics
def _kincar_update(t, x, u, params):
    b = params.get('wheelbase', 3.)             # get parameter values
    #! TODO: add dir processing
    dx = np.array([
        np.cos(x[2]) * u[0],
        np.sin(x[2]) * u[0],
        (u[0]/b) * np.tan(u[1])
    ])
    return dx

def _kincar_output(t, x, u, params):
    return x                            # return x, y, theta (full state)

# Create differentially flat input/output system
kincar = fs.FlatSystem(
    _kincar_flat_forward, _kincar_flat_reverse, name="kincar",
    updfcn=_kincar_update, outfcn=_kincar_output,
    inputs=('v', 'delta'), outputs=('x', 'y', 'theta'),
    states=('x', 'y', 'theta'))

#
# Utility function to plot lane change manuever
#

def plot_lanechange(t, y, u, figure=None, yf=None):
    # Plot the xy trajectory
    plt.subplot(3, 1, 1, label='xy')
    plt.plot(y[0], y[1])
    plt.xlabel("x [m]")
    plt.ylabel("y [m]")
    if yf is not None:
        plt.plot(yf[0], yf[1], 'ro')

    # Plot the inputs as a function of time
    plt.subplot(3, 1, 2, label='v')
    plt.plot(t, u[0])
    plt.xlabel("Time $t$ [sec]")
    plt.ylabel("$v$ [m/s]")

    plt.subplot(3, 1, 3, label='delta')
    plt.plot(t, u[1])
    plt.xlabel("Time $t$ [sec]")
    plt.ylabel("$\\delta$ [rad]")

    plt.suptitle("Lane change manuever")
    plt.tight_layout()