import numpy as np
import osqp
from scipy import sparse
import matplotlib.pyplot as plt

# Colors
PREDICTION = '#BA4A00'

##################
# MPC Controller #
##################


class MPC:
    def __init__(self, model, N, Q, R, QN, StateConstraints, InputConstraints,
                 ay_max):
        """
        Constructor for the Model Predictive Controller.
        MPC的构造函数
        :param model: bicycle model object to be controlled 自行车模型对象用于控制
        :param N: time horizon | int 时间范围
        :param Q: state cost matrix 状态成本矩阵
        :param R: input cost matrix 输入成本矩阵
        :param QN: final state cost matrix 最终状态成本矩阵
        :param StateConstraints: dictionary of state constraints 状态约束的字典
        :param InputConstraints: dictionary of input constraints 输入约束的字典
        :param ay_max: maximum allowed lateral acceleration in curves 允许的最大横向加速度
        """

        # Parameters
        self.N = N  # horizon
        self.Q = Q  # weight matrix state vector
        self.R = R  # weight matrix input vector
        self.QN = QN  # weight matrix terminal

        # Model
        self.model = model

        # Dimensions
        self.nx = self.model.n_states # 状态变量的数量
        self.nu = 2

        # Constraints
        self.state_constraints = StateConstraints
        self.input_constraints = InputConstraints

        # Maximum lateral acceleration
        # 最大横向加速度
        self.ay_max = ay_max

        # Current control and prediction
        # 当前控制和预测
        self.current_prediction = None

        # Counter for old control signals in case of infeasible problem
        # 如果问题不可行，旧控制信号的计数器
        self.infeasibility_counter = 0

        # Current control signals
        # 当前控制信号
        self.current_control = np.zeros((self.nu*self.N))

        # Initialize Optimization Problem
        # 初始化优化问题
        self.optimizer = osqp.OSQP()

    def _init_problem(self):
        """
        Initialize optimization problem for current time step.
        为当前时间步初始化优化问题
        """

        # Constraints
        # 加载约束
        umin = self.input_constraints['umin']
        umax = self.input_constraints['umax']
        xmin = self.state_constraints['xmin']
        xmax = self.state_constraints['xmax']

        # LTV System Matrices
        # LTV系统矩阵
        A = np.zeros((self.nx * (self.N + 1), self.nx * (self.N + 1)))
        B = np.zeros((self.nx * (self.N + 1), self.nu * (self.N)))
        # Reference vector for state and input variables
        # 状态和输入变量的参考向量
        ur = np.zeros(self.nu*self.N)
        xr = np.zeros(self.nx*(self.N+1))
        # Offset for equality constraint (due to B * (u - ur))
        # 等式约束的偏移（由于B *（u - ur））
        uq = np.zeros(self.N * self.nx)
        # Dynamic state constraints
        # 动态状态约束
        xmin_dyn = np.kron(np.ones(self.N + 1), xmin)
        xmax_dyn = np.kron(np.ones(self.N + 1), xmax)
        # Dynamic input constraints
        # 动态输入约束
        umax_dyn = np.kron(np.ones(self.N), umax)
        # Get curvature predictions from previous control signals
        # 从先前的控制信号中获取曲率预测
        kappa_pred = np.tan(np.array(self.current_control[3::] +
                                     self.current_control[-1:])) / self.model.length

        # Iterate over horizon
        # 遍历时间范围
        for n in range(self.N):

            # Get information about current waypoint
            # 获取当前路标的信息
            current_waypoint = self.model.reference_path.get_waypoint(self.model.wp_id + n)
            next_waypoint = self.model.reference_path.get_waypoint(self.model.wp_id + n + 1)
            delta_s = next_waypoint - current_waypoint
            kappa_ref = current_waypoint.kappa
            v_ref = current_waypoint.v_ref

            # Compute LTV matrices
            # 计算LTV矩阵
            f, A_lin, B_lin = self.model.linearize(v_ref, kappa_ref, delta_s)
            A[(n+1) * self.nx: (n+2)*self.nx, n * self.nx:(n+1)*self.nx] = A_lin
            B[(n+1) * self.nx: (n+2)*self.nx, n * self.nu:(n+1)*self.nu] = B_lin

            # Set reference for input signal
            # 设置输入信号的参考值
            ur[n*self.nu:(n+1)*self.nu] = np.array([v_ref, kappa_ref])
            # Compute equality constraint offset (B*ur)
            # 计算等式约束偏移（B * ur）
            uq[n * self.nx:(n+1)*self.nx] = B_lin.dot(np.array
                                            ([v_ref, kappa_ref])) - f

            # Constrain maximum speed based on predicted car curvature
            # 根据预测的汽车曲率限制最大速度
            vmax_dyn = np.sqrt(self.ay_max / (np.abs(kappa_pred[n]) + 1e-12))
            if vmax_dyn < umax_dyn[self.nu*n]:
                umax_dyn[self.nu*n] = vmax_dyn

        # Compute dynamic constraints on e_y
        # 计算e_y的动态约束
        ub, lb, _ = self.model.reference_path.update_path_constraints(
                    self.model.wp_id+1, self.N, 2*self.model.safety_margin,
            self.model.safety_margin)
        xmin_dyn[0] = self.model.spatial_state.e_y
        xmax_dyn[0] = self.model.spatial_state.e_y
        xmin_dyn[self.nx::self.nx] = lb
        xmax_dyn[self.nx::self.nx] = ub

        # Set reference for state as center-line of drivable area
        # 将状态的参考值设置为可驾驶区域的中心线
        xr[self.nx::self.nx] = (lb + ub) / 2

        # Get equality matrix
        # 获取等式矩阵
        Ax = sparse.kron(sparse.eye(self.N + 1),
                         -sparse.eye(self.nx)) + sparse.csc_matrix(A)
        Bu = sparse.csc_matrix(B)
        Aeq = sparse.hstack([Ax, Bu])
        # Get inequality matrix
        # 获取不等式矩阵
        Aineq = sparse.eye((self.N + 1) * self.nx + self.N * self.nu)
        # Combine constraint matrices
        # 组合约束矩阵
        A = sparse.vstack([Aeq, Aineq], format='csc')

        # Get upper and lower bound vectors for equality constraints
        # 获取等式约束的上下限向量
        lineq = np.hstack([xmin_dyn,
                           np.kron(np.ones(self.N), umin)])
        uineq = np.hstack([xmax_dyn, umax_dyn])
        # Get upper and lower bound vectors for inequality constraints
        # 获取不等式约束的上下限向量
        x0 = np.array(self.model.spatial_state[:])
        leq = np.hstack([-x0, uq])
        ueq = leq
        # Combine upper and lower bound vectors
        # 组合上下限向量
        l = np.hstack([leq, lineq])
        u = np.hstack([ueq, uineq])

        # Set cost matrices
        # 设置成本矩阵
        P = sparse.block_diag([sparse.kron(sparse.eye(self.N), self.Q), self.QN,
             sparse.kron(sparse.eye(self.N), self.R)], format='csc')
        q = np.hstack(
            [-np.tile(self.Q.diagonal(), self.N) * xr[:-self.nx],
             -self.QN.dot(xr[-self.nx:]),
             -np.tile(self.R.diagonal(), self.N) * ur])

        # Initialize optimizer
        # 初始化优化器
        self.optimizer = osqp.OSQP()
        self.optimizer.setup(P=P, q=q, A=A, l=l, u=u, verbose=False)

    def get_control(self):
        """
        Get control signal given the current position of the car. Solves a
        finite time optimization problem based on the linearized car model.
        给定车辆的当前位置，获取控制信号。基于线性化的汽车模型解决有限时间优化问题。
        """

        # Number of state variables
        # 状态变量的数量
        nx = self.model.n_states
        nu = 2

        # Update current waypoint
        # 更新当前路标
        self.model.get_current_waypoint()

        # Update spatial state
        # 更新空间状态
        self.model.spatial_state = self.model.t2s(reference_state=
            self.model.temporal_state, reference_waypoint=
            self.model.current_waypoint)

        # Initialize optimization problem
        # 初始化优化问题
        self._init_problem()

        # Solve optimization problem
        # 解决优化问题
        dec = self.optimizer.solve()

        try:
            # Get control signals
            # 获取控制信号
            control_signals = np.array(dec.x[-self.N*nu:])
            control_signals[1::2] = np.arctan(control_signals[1::2] *
                                              self.model.length)
            v = control_signals[0]
            delta = control_signals[1]

            # Update control signals
            # 更新控制信号
            self.current_control = control_signals

            # Get predicted spatial states
            # 获取预测的空间状态
            x = np.reshape(dec.x[:(self.N+1)*nx], (self.N+1, nx))

            # Update predicted temporal states
            # 更新预测的时间状态
            self.current_prediction = self.update_prediction(x)

            # Get current control signal
            # 获取当前控制信号
            u = np.array([v, delta])

            # if problem solved, reset infeasibility counter
            # 如果问题解决了，重置不可行计数器
            self.infeasibility_counter = 0

        except:

            print('Infeasible problem. Previously predicted'
                  ' control signal used!')
            id = nu * (self.infeasibility_counter + 1)
            u = np.array(self.current_control[id:id+2])

            # increase infeasibility counter
            # 增加不可行计数器
            self.infeasibility_counter += 1

        if self.infeasibility_counter == (self.N - 1):
            print('No control signal computed!')
            exit(1)

        return u

    def update_prediction(self, spatial_state_prediction):
        """
        Transform the predicted states to predicted x and y coordinates.
        Mainly for visualization purposes.
        将预测的状态转换为预测的x和y坐标。主要用于可视化目的。
        :param spatial_state_prediction: list of predicted state variables 预测状态变量的列表
        :return: lists of predicted x and y coordinates 预测的x和y坐标的列表
        """

        # Containers for x and y coordinates of predicted states
        # 预测状态的x和y坐标的容器
        x_pred, y_pred = [], []

        # Iterate over prediction horizon
        # 遍历预测范围
        for n in range(2, self.N):
            # Get associated waypoint
            # 获取关联的路标
            associated_waypoint = self.model.reference_path.\
                get_waypoint(self.model.wp_id+n)
            # Transform predicted spatial state to temporal state
            # 将预测的空间状态转换为时间状态
            predicted_temporal_state = self.model.s2t(associated_waypoint,
                                            spatial_state_prediction[n, :])

            # Save predicted coordinates in world coordinate frame
            # 保存世界坐标系中的预测坐标
            x_pred.append(predicted_temporal_state.x)
            y_pred.append(predicted_temporal_state.y)

        return x_pred, y_pred

    def show_prediction(self):
        """
        Display predicted car trajectory in current axis.
        在当前轴中显示预测的汽车轨迹。
        """

        if self.current_prediction is not None:
            plt.scatter(self.current_prediction[0], self.current_prediction[1],
                    c=PREDICTION, s=30)

