import numpy as np
np.seterr(divide='ignore', invalid='ignore')

from scipy.integrate import odeint
from scipy.interpolate import interp1d
import cvxpy as opt

from .utils import *

from .mpc_config import Params
P=Params()

def get_linear_model_matrices(x_bar,u_bar):
    """
    Computes the LTI approximated state space model x' = Ax + Bu + C
    """
    
    x = x_bar[0]
    y = x_bar[1]
    v = x_bar[2]
    theta = x_bar[3]
    
    a = u_bar[0]
    delta = u_bar[1]
    
    ct = np.cos(theta)
    st = np.sin(theta)
    cd = np.cos(delta)
    td = np.tan(delta)
    
    A = np.zeros((P.N,P.N))
    A[0,2] = ct
    A[0,3] = -v*st
    A[1,2] = st
    A[1,3] = v*ct
    A[3,2] = v*td/P.L
    A_lin = np.eye(P.N)+P.DT*A
    
    B = np.zeros((P.N,P.M))
    B[2,0]=1
    B[3,1]=v/(P.L*cd**2)
    B_lin=P.DT*B
    
    f_xu=np.array([v*ct, v*st, a, v*td/P.L]).reshape(P.N,1)
    C_lin = P.DT*(f_xu - np.dot(A, x_bar.reshape(P.N,1)) - np.dot(B, u_bar.reshape(P.M,1))).flatten()
    
    #return np.round(A_lin,6), np.round(B_lin,6), np.round(C_lin,6)
    return A_lin, B_lin, C_lin

class MPC():
    
    def __init__(self, N, M, Q, R):
        """
        """
        self.state_len = N
        self.action_len = M
        self.state_cost = Q
        self.action_cost = R
        
    def optimize_linearized_model(self, A, B, C, initial_state, target, time_horizon=10, Q=None, R=None, verbose=False):
        """
        Optimisation problem defined for the linearised model,
        :param A: 
        :param B:
        :param C: 
        :param initial_state:
        :param Q:
        :param R:
        :param target:
        :param time_horizon:
        :param verbose:
        :return:
        """
        
        assert len(initial_state) == self.state_len
        
        if (Q == None or R==None):
            Q = self.state_cost
            R = self.action_cost
        
        # Create variables
        x = opt.Variable((self.state_len, time_horizon + 1), name='states')
        u = opt.Variable((self.action_len, time_horizon), name='actions')

        # Loop through the entire time_horizon and append costs
        cost_function = []

        for t in range(time_horizon):

            _cost = opt.quad_form(target[:, t + 1] - x[:, t + 1], Q) +\
                    opt.quad_form(u[:, t], R)
            
            _constraints = [x[:, t + 1] == A @ x[:, t] + B @ u[:, t] + C,
                            u[0, t] >= -P.MAX_ACC, u[0, t] <= P.MAX_ACC,
                            u[1, t] >= -P.MAX_STEER, u[1, t] <= P.MAX_STEER]
                            #opt.norm(target[:, t + 1] - x[:, t + 1], 1) <= 0.1]
            
            # Actuation rate of change
            if t < (time_horizon - 1):
                _cost += opt.quad_form(u[:,t + 1] - u[:,t], R * 1)
                _constraints += [opt.abs(u[0, t + 1] - u[0, t])/P.DT <= P.MAX_D_ACC]
                _constraints += [opt.abs(u[1, t + 1] - u[1, t])/P.DT <= P.MAX_D_STEER]
            
            
            if t == 0:
                #_constraints += [opt.norm(target[:, time_horizon] - x[:, time_horizon], 1) <= 0.01,
                #                x[:, 0] == initial_state]
                _constraints += [x[:, 0] == initial_state]
            
            cost_function.append(opt.Problem(opt.Minimize(_cost), constraints=_constraints))

        # Add final cost
        problem = sum(cost_function)
        
        # Minimize Problem
        problem.solve(verbose=verbose, solver=opt.OSQP)
        return x, u