%% Close Figures, Clear Workspace, and Clear Terminal
close all;
clear;
clc;

%% System Parameters
% Track Information & Reference Trajectory
load("TestTrack.mat");
load('ROB535_ControlProject_part1_Team3.mat');

% Vehicle Parameters (Table 1)
global Nw f Iz a b By Cy Dy Ey Shy Svy m g
Nw = 2;
f = 0.01;
Iz = 2667;
a = 1.35;
b = 1.45;
By = 0.27;
Cy = 1.2;
Dy = 0.7;
Ey = -1.6;
Shy = 0;
Svy = 0;
m = 1400;
g = 9.806;

% Input Limits (Table 1)
global delta_lims Fx_lims
delta_lims = [-0.5, 0.5];
Fx_lims = [-5e3, 5e3];

% Position Limits
global x_lims y_lims
x_lims = [200, 1600];
y_lims = [-200, 1000];

% Initial Conditions (Equation 15)
state_init = [ ...
    287; ...  % x [m]
    5; ...    % u [m/s]
    -176; ... % y [m]
    0; ...    % v [m/s]
    2; ...    % psi [rad]
    0; ...    % r [rad/s]
];

% Simulation Parameters
global T_s T_p num_preds num_states num_inputs
T_s = 0.01; % Step Size [s]
T_p = 0.5; % Prediction Horizon [s]
num_preds = T_p / T_s;
num_states = 6;
num_inputs = 2;

%% Constraint Functions
function [Lb, Ub] = bound_cons(idx, X_ref, U_ref)
    global num_preds num_states num_inputs x_lims y_lims delta_lims Fx_lims

    % initial_idx is the index along uref the initial condition is at
    Lb = -Inf(num_preds*(num_states+num_inputs), 1);
    Ub = Inf(num_preds*(num_states+num_inputs), 1);

    for i = 0:(num_preds-1)
        start_idx = get_start_idx(i);

        % x
        Lb(start_idx+1) = x_lims(1) - X_ref(idx+i, 1);
        Ub(start_idx+1) = x_lims(2) - X_ref(idx+i, 1);

        % y
        Lb(start_idx+3) = y_lims(1) - X_ref(idx+i, 3);
        Ub(start_idx+3) = y_lims(2) - X_ref(idx+i, 3);

        % delta
        Lb(start_idx+num_states+1) = delta_lims(1) - U_ref(idx+i, 1);
        Ub(start_idx+num_states+1) = delta_lims(2) - U_ref(idx+i, 1);
        
        % F_x
        Lb(start_idx+num_states+2) = Fx_lims(1) - U_ref(idx+1, 2);
        Ub(start_idx+num_states+2) = Fx_lims(2) - U_ref(idx+1, 2);
    end
end

function [c, ceq] = road_obstacle_cons(Z, TestTrack, Xobs)
    global num_preds num_states

    ceq = [];
    c = NaN(1,num_preds);
    
    for i = 1:num_preds
        idx = get_start_idx(i);
        X = Z(idx+1:idx+num_states);

        p = [X(1); X(3)];
        [~,bl_idx] = min(vecnorm(TestTrack.bl - p));
        [~,br_idx] = min(vecnorm(TestTrack.br - p));

        idx_search = 1;
        bl_idx_start = clamp(bl_idx-idx_search, 1, size(TestTrack.bl,2));
        bl_idx_end   = clamp(bl_idx+idx_search, 1, size(TestTrack.bl,2));
        br_idx_start = clamp(br_idx-idx_search, 1, size(TestTrack.br,2));
        br_idx_end   = clamp(br_idx+idx_search, 1, size(TestTrack.br,2));

        boundary_pts = [ ...
            TestTrack.bl(:,bl_idx_start:1:bl_idx_end), ...
            TestTrack.br(:,br_idx_end:-1:br_idx_start) ...
        ];
        xv_road = boundary_pts(1,:);
        yv_road = boundary_pts(2,:);

        in_road = inpolygon(p(1), p(2), xv_road, yv_road);

        if ~in_road
            % Point not inside road
            c(i) = 1; % c(x) > 0, nonlinear inequality constraint violated
        end

        if ~isnan(c(i))
            % If value set, go to next "i"
            continue
        end

        for j = 1:size(Xobs,2)
            xv_obstacle = Xobs{i}(:,1);
            yv_obstacle = Xobs{i}(:,2);
            [in_obstacle, on_obstacle] = inpolygon(p(1), p(2), xv_obstacle, yv_obstacle);

            if in_obstacle || on_obstacle
                % Point in or on obstacle
                c(i) = 1; % c(x) > 0, nonlinear inequality constraint violated
            end

            if ~isnan(c(i))
                % If value set, skip checking remaining obstacles
                break
            end
        end

        if isnan(c(i))
            % If value not set, no constraints violated
            c(i) = -1; % c(x) <= 0, nonlinear inequality constraint satisfied
        end
    end
end

%% Kinematic Bike Models
function dXdt = nonlinear_bike_model(X,U)
    global Nw f Iz a b m g
    
    [x,u,y,v,psi,r,delta_f,F_x,F_yf,F_yr] = bike_model_helper(X,U);
    
    % Vehicle Dynamics (Equation 1)
    dx = u*cos(psi) - v*sin(psi);
    du = (1/m)*(-f*m*g + Nw*F_x - F_yf*sin(delta_f)) + v*r;
    dy = u*sin(psi) + v*cos(psi);
    dv = (1/m)*(F_yf*cos(delta_f) + F_yr) - u*r;
    dpsi = r;
    dr = (1/Iz)*(a*F_yf*cos(delta_f) - b*F_yr);
    dXdt = [dx; du; dy; dv; dpsi; dr];
end

function [A, B] = linearized_bike_model(X_ref,U_ref)
    global Nw f Iz a b m g T_s

    syms x_var u_var y_var v_var psi_var r_var ... % states
         delta_f_var F_x_var ...                   % inputs
         F_yf_var F_yr_var                         % lateral forces
    
    % Vehicle Dynamics (Equation 1)
    dx = u_var*cos(psi_var) - v_var*sin(psi_var);
    du = (1/m)*(-f*m*g + Nw*F_x_var - F_yf_var*sin(delta_f_var)) + v_var*r_var;
    dy = u_var*sin(psi_var) + v_var*cos(psi_var);
    dv = (1/m)*(F_yf_var*cos(delta_f_var) + F_yr_var) - u_var*r_var;
    dpsi = r_var;
    dr = (1/Iz)*(a*F_yf_var*cos(delta_f_var) - b*F_yr_var);

    % Jacobians of continuous-time linearized system
    % TODO: Change A_c & B_c computation to regular functions (e.g.,
    % using `matlabFunction`)
    A_c_symb = [ ...
        diff(dx,  x_var), diff(dx,  u_var), diff(dx,  y_var), diff(dx,  v_var), diff(dx,  psi_var), diff(dx,  r_var);
        diff(du,  x_var), diff(du,  u_var), diff(du,  y_var), diff(du,  v_var), diff(du,  psi_var), diff(du,  r_var);
        diff(dy,  x_var), diff(dy,  u_var), diff(dy,  y_var), diff(dy,  v_var), diff(dy,  psi_var), diff(dy,  r_var);
        diff(dv,  x_var), diff(dv,  u_var), diff(dv,  y_var), diff(dv,  v_var), diff(dv,  psi_var), diff(dv,  r_var);
        diff(dpsi,x_var), diff(dpsi,u_var), diff(dpsi,y_var), diff(dpsi,v_var), diff(dpsi,psi_var), diff(dpsi,r_var);
        diff(dr,  x_var), diff(dr,  u_var), diff(dr,  y_var), diff(dr,  v_var), diff(dr,  psi_var), diff(dr,  r_var);
    ];

    B_c_symb = [ ...
        diff(dx,  delta_f_var), diff(dx,  F_x_var);
        diff(du,  delta_f_var), diff(du,  F_x_var);
        diff(dy,  delta_f_var), diff(dy,  F_x_var);
        diff(dv,  delta_f_var), diff(dv,  F_x_var);
        diff(dpsi,delta_f_var), diff(dpsi,F_x_var);
        diff(dr,  delta_f_var), diff(dr,  F_x_var);
    ];    

    % Substitute values from reference trajectory into symbolic Jacobians
    A_c = @(i) double( ...
        subs( ...
            A_c_symb, ...
            [x_var, u_var, y_var, v_var, psi_var, r_var, delta_f_var, F_x_var, F_yf_var, F_yr_var], ...
            bike_model_helper(X_ref(i,:), U_ref(i,:)) ...
        ) ...
    );
    B_c = @(i) double( ...
        subs( ...
            B_c_symb, ...
            [x_var, u_var, y_var, v_var, psi_var, r_var, delta_f_var, F_x_var, F_yf_var, F_yr_var], ...
            bike_model_helper(X_ref(i,:), U_ref(i,:)) ...
        ) ...
    );

    % Discrete-time LTV system
    A = @(i) eye(6) + T_s*A_c(i);
    B = @(i) T_s * B_c(i);
end

%% Helper Functions
function clamped_val = clamp(val, min, max)
    clamped_val = val;

    if clamped_val < min
        clamped_val = min;
    elseif clamped_val > max
        clamped_val = max;
    end
end

function [x,u,y,v,psi,r,delta_f,F_x,F_yf,F_yr] = bike_model_helper(X,U)
    global Nw a b By Cy Dy Ey Shy Svy m g

    % Get state & input variables
    x   = X(1);
    u   = X(2);
    y   = X(3);
    v   = X(4);
    psi = X(5);
    r   = X(6);

    delta_f = U(1);
    F_x     = U(2);
    
    % Front and rear lateral slip angles in radians (Equations 8 & 9)
    alpha_f_rad = delta_f - atan2(v + a*r, u);
    alpha_r_rad = -atan2(v - b*r, u);

    % Convert radians to degrees for other equations
    alpha_f = rad2deg(alpha_f_rad);
    alpha_r = rad2deg(alpha_r_rad);
    
    % Nonlinear Tire Dynamics (Equations 6 & 7)
    phi_yf = (1-Ey)*(alpha_f + Shy) + (Ey/By)*atan(By*(alpha_f + Shy));
    phi_yr = (1-Ey)*(alpha_r + Shy) + (Ey/By)*atan(By*(alpha_r + Shy));

    % Lateral forces using Pacejka "Magic Formula" (Equations 2 - 5)
    F_zf = (b/(a+b))*(m*g);
    F_yf = F_zf*Dy*sin(Cy*atan(By*phi_yf)) + Svy;
    F_zr = (a/(a+b))*(m*g);
    F_yr = F_zr*Dy*sin(Cy*atan(By*phi_yr)) + Svy;
    
    % Limits on combined longitudinal and lateral loading of tires
    % (Equations 10 - 14)
    F_total = sqrt((Nw*F_x)^2 + (F_yr^2));
    F_max = 0.7*(m*g);
    
    if F_total > F_max
        F_x = (F_max/F_total)*F_x;
        F_yr = (F_max/F_total)*F_yr;
    end

    % Apply input limits (Table 1)
    delta_f = clamp(delta_f, delta_lims(1), delta_lims(2));
    F_x = clamp(F_x, Fx_lims(1), Fx_lims(2));
end

function idx = get_start_idx(i)
    global num_states num_inputs
    idx = (i-1)*(num_states+num_inputs);
end