"""
An iteratinv LQR algorithm for CARLA

Need an initial trajectory, possibly generated by the neural controller.

Based on the trajectory, solve a perturbation such that the perturbed trajectory has lower cost than the previous one.

In concern of safety, add safety constrain to the optimization problem and the constraints query for waypoints.

The constraints are resolved by turning the constrained problem into an unconstrained one.

To minimize the cost is to satisfy the safety constraint.

"""
import argparse
import os
import sys

import numpy as np
import theano.tensor as T

import time
from .dynamics import Dynamics, BatchAutoDiffDynamics, tensor_constrain, sigmoid_constrain, hard_sigmoid_constrain
from .cost import Cost, PathQRCost
from .controller import iLQR

### The bounds for the control
STEER_BOUND_UP = 1.0
STEER_BOUND_LOW = -1.0
THROTTLE_BOUND_UP = 1.0
THROTTLE_BOUND_LOW = 0.0

# Global variables defined for monitorign the iLQR processing
J_hist = []
def on_iteration(iteration_count, xs, us, J_opt, accepted, converged):
    J_hist.append(J_opt)
    info = "converged" if converged else ("accepted" if accepted else "failed")
    final_state = xs[-1]
    final_control = us[-1]
    if iteration_count % 30 == 0:
        print("iteration", iteration_count, info, J_opt, final_state, final_control)



class CarDynamics(BatchAutoDiffDynamics):
    """
    Define car dynamics f.
    BatchAutoDiffDynamics will automatically calculate the jacobian
    """
    def __init__(self,
                 dt = 0.05,
                 l = 2.5,
                 constrain = True,
                 min_bounds = np.array([THROTTLE_BOUND_LOW, STEER_BOUND_LOW]),
                 max_bounds = np.array([THROTTLE_BOUND_UP, STEER_BOUND_UP]),
                 **kwargs):
        """
        Args:
            dt: the time span of executing the input control
            constrain: the constraints for the input control
            min_bounds, max_bounds: the lower/upper bounds of the control
        
        Notes:
            state variables: [posx, posy, theta, v]
            action: [v_dot, theta_dot]

        """
        self.dt = dt
        self.l = l

        self.constrained = constrain
        self.min_bounds = min_bounds
        self.max_bounds = max_bounds

        def f(x, u, i):
            if self.constrained:
                u = hard_sigmoid_constrain(u, self.min_bounds, self.max_bounds)

            posx = x[..., 0]
            posy = x[..., 1]
            theta = x[..., 2]
            v = x[..., 3]

            v_dot = u[..., 0]
            delta = np.pi * 70./180 * u[..., 1]
            tanh_delta = T.tanh(delta)

            # Define dynamics model as in paper 
            # Constrained Iterative LQR for On-Road Autonomous Driving Motion Planning
            
            theta_dot = v * tanh_delta/self.l
            posx_ = posx + T.cos(theta) * (v * dt + 0.5 * v_dot * self.dt**2)
            posy_ = posy + T.sin(theta) * (v * dt + 0.5 * v_dot * self.dt**2)
            v_ = v +  v_dot * dt
            theta_ = theta + theta_dot * self.dt

            return T.stack([
                posx_,
                posy_,
                theta_,
                v_
                ]).T

        super(CarDynamics, self).__init__(f, state_size = 4, action_size = 2, **kwargs)


class CarCost(Cost):
    """Quadratic Regulator Instantaneous Cost for trajectory following and barrier function."""
    
    def __init__(self, Q, q, R, r, b12, b34, q1, q2, q3, q4, x_path, x_avoids, Q_terminal = None, u_path = None):

        """
        Construct a Quadratic Cost with barriers
        
        Args:
            Q: Quadratic cost for the state variables [state_size, state_size]
            q: linear cost for the state variables [state_size, 1]
            R: Quadratic cost for the control variables [action_size, action_size]
            r: linear cost for the control varaiables [action_size, 1]
            b: linear coeff for the state upper and lower barrier cost [state_size, 1]
                i.e. safe(x) = (x - wp)^2 - b^2 <= 0
            q1, q2: scale coeffs for the state barrier cost
                i.e. safe(x) => q1 * exp[q2 * ((x - wp)^2 - b^2)]
                    grad safe(x) => 2q1 q2 (x - wp) exp[q2 * ((x - wp)^2 - b^2)]
                    grad^2 safe(x) => 2q1q2 exp[q2*((x-wp)^2 - b^2)] + 4q1q2^2 (x-wp)^2 exp[q2*((x-wp)^2 - b^2)]
            x_path: nominal trajectory to be followed [N + 1, state_size]
            x_barrier_u: upper bound for the trajectory [N + 1, state_size]
            x_barrier_l: lower bound for the trajectory [N + 1, state_size]
        
        The cost is:
            1/2 * x^T Q x + x^T q + 
            1/2 * u^T R u + u^T r + 
            1/2 * x^T q1q2exp(q2(Ax-b))A^T A + x^T q1q2exp(q2(Ax - b))A^T 

        """

        self.Q = np.array(Q)
        self.q = np.array(q)

        self.R = np.array(R)
        self.r = np.array(r)

        self.b12 = b12
        self.b34 = b34
        self.q1 = q1
        self.q2 = q2
        self.q3 = q3
        self.q4 = q4

        """ F,f are the quadratic and linear costs for the state barrier cost"""
        self.F = np.zeros(self.Q.shape)
        self.f = np.zeros(self.q.shape)

        self.x_path = np.asarray(x_path)

        state_size = self.Q.shape[0]
        action_size = self.R.shape[0]
        path_length = self.x_path.shape[0]

        if Q_terminal is None:
            self.Q_terminal = self.Q
        else:
            self.Q_terminal = np.array(self.Q_terminal)

        if u_path is None:
            self.u_path = np.zeros([path_length - 1, action_size])
        else:
            self.u_path = np.array(u_path)

        if x_avoids is None:
            self.x_avoids = np.zeros([path_length, state_size])
        else:
            self.x_avoids = np.array(x_avoids)

        assert self.Q.shape[0] == self.Q.shape[1], "Q must be square"
        assert self.q.shape[0] == self.Q.shape[0], "q mismatch"

        assert self.R.shape[0] == self.R.shape[1], "R must be square"
        assert self.r.shape[0] == self.R.shape[0], "r mismatch"


        assert state_size == self.x_path.shape[1], "Q & x_path mismatch"
        assert action_size == self.u_path.shape[1], "R & u_path mismatch {} vs {}".format(R.shape, u_path.shape)

        # Precompute some common constants.
        self._Q_plus_Q_T = self.Q + self.Q.T
        self._Q_plus_Q_T_terminal = self.Q_terminal + self.Q_terminal.T
        self._R_plus_R_T = self.R + self.R.T
        self._F_plus_F_T = self.F + self.F.T

        super(CarCost, self).__init__()

    def l(self, x, u, i, terminal=False):
        """Instantaneous cost function.

        Args:
            x: Current state [state_size].
            u: Current control [action_size]. None if terminal.
            i: Current time step.
            terminal: Compute terminal cost. Default: False.

        Returns:
            Instantaneous cost (scalar).
        """
        Q = self.Q_terminal if terminal else self.Q
        r = self.r
        q = self.q
        R = self.R

        x_diff = x - self.x_path[i]
        x_dist = x - self.x_avoids[i] 
        x_diff[2:] = 0.0
        x_dist[2:] = 0.0



        constant_cost = self.q3 * np.exp(self.q4 * (- x_dist.T.dot(x_dist) + self.b34 * self.b34))
        constant_cost += self.q1 * np.exp(self.q2 * (x_diff.T.dot(x_diff) - self.b12 * self.b12))

        squared_x_cost = x_diff.T.dot(Q).dot(x_diff)
        if terminal:
            return squared_x_cost

        u_diff = u - self.u_path[i]
        squared_u_cost = u_diff.T.dot(R).dot(u_diff)

        linear_cost = x_diff.T.dot(q) + u_diff.T.dot(r)
        return squared_x_cost + squared_u_cost + linear_cost + constant_cost

    def l_x(self, x, u, i, terminal=False):
        """Partial derivative of cost function with respect to x.

        Args:
            x: Current state [state_size].
            u: Current control [action_size]. None if terminal.
            i: Current time step.
            terminal: Compute terminal cost. Default: False.

        Returns:
            dl/dx [state_size].
        """
        Q_plus_Q_T = self._Q_plus_Q_T_terminal if terminal else self._Q_plus_Q_T

        x_diff = np.reshape(x - self.x_path[i], (self.Q.shape[0], 1))
        x_diff[2:] = 0.0
        x_dist = np.reshape(x - self.x_avoids[i], (self.Q.shape[0], 1))
        x_dist[2:] = 0.0


        self.f = - self.q3 * self.q4 * np.exp(self.q4 * (- x_dist.T.dot(x_dist) + self.b34 * self.b34)) * 2 * x_dist
        self.f +=  self.q1 * self.q2 * np.exp(self.q2 * (x_diff.T.dot(x_diff) - self.b12 * self.b12)) * 2 * x_diff


        return x_diff.T.dot(Q_plus_Q_T) + self.f.T + self.q.T

    def l_u(self, x, u, i, terminal=False):
        """Partial derivative of cost function with respect to u.

        Args:
            x: Current state [state_size].
            u: Current control [action_size]. None if terminal.
            i: Current time step.
            terminal: Compute terminal cost. Default: False.

        Returns:
            dl/du [action_size].
        """
        if terminal:
            return np.zeros_like(self.u_path)

        u_diff = u - self.u_path[i]
        return u_diff.T.dot(self._R_plus_R_T) + self.r.T

    def l_xx(self, x, u, i, terminal=False):
        """Second partial derivative of cost function with respect to x.

        Args:
            x: Current state [state_size].
            u: Current control [action_size]. None if terminal.
            i: Current time step.
            terminal: Compute terminal cost. Default: False.

        Returns:
            d^2l/dx^2 [state_size, state_size].
        """
        x_diff = np.reshape(x - self.x_path[i], (self.Q.shape[0], 1))
        x_diff[2:] = 0.0
        x_dist = np.reshape(x - self.x_avoids[i], (self.Q.shape[0], 1))
        x_dist[2:] = 0.0


        self.F =  - self.q3 * self.q4 * np.exp(self.q4 * (-x_dist.T.dot(x_dist) + self.b34 * self.b34)) * 2  + self.q3 * self.q4**2 * np.exp(self.q4 * (- x_dist.T.dot(x_dist) + self.b34 * self.b34)) * 4 * x_dist.dot(x_dist.T)
        self.F += self.q1 * self.q2 * np.exp(self.q2 * (x_diff.T.dot(x_diff) - self.b12 * self.b12)) * 2  + self.q1 * self.q2**2 * np.exp(self.q2 * (x_diff.T.dot(x_diff) - self.b12 * self.b12)) * 4 * x_diff.dot(x_diff.T)


        self._F_plus_F_T = self.F + self.F.T

        return self._F_plus_F_T + self._Q_plus_Q_T_terminal if terminal else self._F_plus_F_T + self._Q_plus_Q_T

    def l_ux(self, x, u, i, terminal=False):
        """Second partial derivative of cost function with respect to u and x.

        Args:
            x: Current state [state_size].
            u: Current control [action_size]. None if terminal.
            i: Current time step.
            terminal: Compute terminal cost. Default: False.

        Returns:
            d^2l/dudx [action_size, state_size].
        """
        return np.zeros((self.R.shape[0], self.Q.shape[0]))

    def l_uu(self, x, u, i, terminal=False):
        """Second partial derivative of cost function with respect to u.

        Args:
            x: Current state [state_size].
            u: Current control [action_size]. None if terminal.
            i: Current time step.
            terminal: Compute terminal cost. Default: False.

        Returns:
            d^2l/du^2 [action_size, action_size].
        """
        if terminal:
            return np.zeros_like(self.R)

        return self._R_plus_R_T


class ILQRController():
    def __init__(self, target_speed, steps_ahead = 10, dt = 0.25, l = 1.0, half_width = 2.0):
        # Target speed determines the distances during each step
        self.target_speed = target_speed * 1000./3600.
        # Number of steps to be optimized
        self.steps_ahead = steps_ahead
        # Step length
        self.dt = dt
        # Front wheel diameter
        self.l = l
        # Generate sequential minor waypoints between each pair of adjacent carla waypoints
        # All consecutive minor waypoints have the same distance
        self.x_path_density = self.target_speed * self.dt
        # Define car dynamics
        self.dynamics = CarDynamics(self.dt, self.l)
        # Lack of variables to build car cost
        self.cost = None
        # lack of variales to build nominal trajectory, lower and upper bounds
        self.x_path = np.zeros([self.steps_ahead, self.dynamics.state_size])
        self.u_path = np.zeros([self.steps_ahead - 1, self.dynamics.action_size])
        # Obstacles to avoid
        self.x_avoids = np.zeros([self.steps_ahead, self.dynamics.state_size])
        # Half of the road width
        self.half_width = half_width 

        # Current state
        self.measurements = {'posx': None, 'posy': None, 'v': None, 'theta': None}

       

    def get_state(self, measurements):
        """Get the state variables from measurements"""
        v = np.linalg.norm([measurements.v.x, measurements.v.y], ord = 2)
        theta = measurements.t.rotation.yaw * np.pi / 180.
        posx, posy = measurements.t.location.x, measurements.t.location.y
        self.measurements['posx'] = posx
        self.measurements['posy'] = posy
        self.measurements['v'] = v
        self.measurements['theta'] = theta
        #print("ilqr measurements:", self.measurements)


    def predict_trajectory(self, us_init, world = None):
        """ (Optional) If an initial state sequence is given, then choose the nominal barrier accordingly"""
        if xs_init is None:
            xs_init = np.zeros([self.steps_ahead, self.dynamics.state_size]) 
            xs_init[0, :] = x0[:]
            for t in range(self.steps_ahead - 1):
                xs_init[t + 1, :] = self.dynamics.f(xs_init[t, :], us_init[t, :], t)
        else:
            xs_init = xs_init[:]
        
        assert xs_init.shape == (self.steps_ahead, self.dynamics.state_size)

    def generate_nominal(self, future_wps_np = None):
        # Along the dense nominal trajectory, select the first self.steps_ahead waypoints 

        # Find the closest reference waypoints among the dense set of waypoints
        # Extracrt the reference waypoints and calculate the barrier sequences
        """ (Optional)Coordinate transform. The ego car is the original point
        x0[:4] = np.asarray([0.0, 0.0, 0.0, self.measurements['v']])
        future_wps_np_ = ILQRController.transform_into_cars_coordinate_system(future_wps_np, \
                self.measurements['posx'], self.measurements['posy'], \
                np.cos(self.measurements['theta']), np.sin(self.measurements['theta']))
        #self.x_path[:, :2] = future_wps_np_[:self.steps_ahead, :2]
        """

        """(Optional)Turn the sparse waypoints into dense waypoints
        future_wps_np_dense = ILQRController.transform_into_dense_points(future_wps_np, self.x_path_density)
        """

        """ Find the closest reference waypoints among the dense set of waypoints
            Extracrt the reference waypoints and calculate the barrier sequences
        self.generate_nominal(future_wps_np_dense)
        self.x_path[:, :2] = future_wps_np_dense[:self.steps_ahead, :2]
        """
        self.x_path = np.zeros([self.steps_ahead, self.dynamics.state_size])
        self.x_path[:, -1] = self.target_speed 
        self.x_path[:self.steps_ahead, :2] = future_wps_np[:self.steps_ahead, :2]
    
    def generate_avoidance(self, locations = None):
        if locations is not None:
            for i in range(self.steps_ahead):
                if len(locations) > 1:
                    self.x_avoids[i, 0] = locations[i].x
                    self.x_avoids[i, 1] = locations[i].y
                else:
                    self.x_avoids[i, 0] = locations[0].x
                    self.x_avoids[i, 1] = locations[0].y
        print("Distance to obstacle: {}".format(\
                np.linalg.norm([self.measurements['posx'] - self.x_avoids[0, 0], \
                self.measurements['posy'] - self.x_avoids[0, 1]])))

    def generate_cost(self):
        Q = 0.0 * np.eye(self.dynamics.state_size)
        Q[3, 3] = 0.0
        Q[2, 2] = 0.0

        R = np.eye(self.dynamics.action_size)
        R[0, 0] = .0
        R[1, 1] = .0
        
        # For staying in lane
        q1 = 1.0E-2
        q2 = 1.0E-4

        # For avoidance
        q3 = 0.0 * 1.
        q4 = 0.0 * 2.0

        b12 = self.half_width
        b34 = self.half_width

        q = np.zeros([self.dynamics.state_size, 1])
        r = np.zeros([self.dynamics.action_size, 1])
        
        self.cost = CarCost(Q = Q, q = q, R = R, r = r, b12 = b12, b34 = b34, \
                q1 = q1, q2 = q2, q3 = q3, q4 = q4,\
                x_path = self.x_path, x_avoids = self.x_avoids, u_path = self.u_path)
        #self.cost = PathQRCost(Q = Q, R = R, x_path = self.x_path)
    
    def run_ilqr(self, x0, us_init):
        # Define static varialbe for ilqr
        
        # Initialize ilqr solver
        ilqr = iLQR(self.dynamics, self.cost, self.steps_ahead - 1)
        
        """ Run ilqr to obtain sequence of state and control
            xs: [1 + self.steps_ahead, state_size]
            us: [self.steps_ahead, state_size]
        """
        xs, us = ilqr.fit(x0, us_init, n_iterations = 50, on_iteration=on_iteration)

        return xs, us

    
    def control(self, future_wps_np, measurements, us_init = None, avoidances = None):
        # Collect state variables
        self.get_state(measurements)
        # Initial state is x0
        x0 = np.zeros([self.dynamics.state_size])
        x0[:4] = np.asarray([self.measurements['posx'], \
                self.measurements['posy'], \
                self.measurements['theta'], \
                self.measurements['v']])
        print("Run ilqr from state ", x0)
        
        # Define noimal path
        self.generate_nominal(future_wps_np)

        # Define obstacles
        self.generate_avoidance(avoidances)

        # Choose initial control sequence
        if us_init is None:
            #us_init = np.random.random([self.steps_ahead - 1, self.dynamics.action_size])
            #us_init[:, 0] = 2 * us_init[:, 0] - 1.0 
            #us_init[:, 1] = 0.0
            us_init = self.u_path
        else:
            us_init = us_init[:self.steps_ahead - 1, :] 

        assert us_init.shape == (self.steps_ahead - 1, self.dynamics.action_size)

        # Define the cost
        self.generate_cost()

        # Run ilqr
        xs, us = self.run_ilqr(x0, us_init)
        #self.u_path = us
        # Verify the safety of xs
        #self.verify(self, xs)
        #for i in range(self.steps_ahead - 1):
        #    print("predict trajectory:\n", xs[i], us[i])
        return xs, us


    @staticmethod
    def transform_into_dense_points(pts, density):
        """ Uniformly add additional waypoints between consecutive waypoints."""
        # Firstly add the first waypoint
        pts_ = [pts[0, :]]
        for i in range(pts.shape[0] - 1):
            # For each pair of consecutive waypoints, define the number of new waypoints to be added

            # np.linspace also returns the head and the tail points which are the original waypoints
            num = int(np.floor(np.linalg.norm(pts[i + 1] - pts[i])/density)) + 2

            # Generate the additional waypoints along each axle
            pts_i_0 = np.linspace(pts[i, 0], pts[i + 1, 0], num = num)
            pts_i_1 = np.linspace(pts[i, 1], pts[i + 1, 1], num = num)
            for j in range(1, num):
                # Append the new way points
                pts_.append([pts_i_0[j], pts_i_1[j]])

        return np.asarray(pts_)
            
    @staticmethod
    def transform_into_cars_coordinate_system(pts, x, y, theta_cos, theta_sin):
        diff = (pts - [x, y])
        pts_ = np.zeros_like(diff)
        pts_[:, 0] = theta_cos * diff[:, 0] + theta_sin * diff[:, 1]
        pts_[:, 1] = - theta_sin * diff[:, 0] + theta_cos * diff[:, 1]
        return pts_

    @staticmethod
    def transform_into_world_coordinate_system(pts, x, y, theta_cos, theta_sin):
        diff = np.zeros_like(pts)
        diff[:, 0] = pts[:, 0] * theta_cos + pts[:, 1] * theta_sin
        diff[:, 1] = pts[:, 0] * theta_sin + pts[:, 1] * theta_cos
        pts_ = diff + [x, y]
        return pts_

    @staticmethod
    def find_trajectory_waypoints(pts, world):
        # Along the initial trajectory, find the closest waypoints using the world map
        m = world.get_map()
        pts_ = []
        for i in range(pts.shape[0]):
            wp = m.get_waypoint(carla.Location(x = pts[i, 0], y = pts[i, 1], z = 0.0)) 
            pts_.append([wp])
        return pts_
