%% ======================================================================
% ROB 535 Control Project: Peerayos Pongsachai
% clear all;clc; close all;
load("TestTrack.mat");

% Objectives:
% - Come up with an inequality constraint function that keeps the
%       center of mass inside the track limits for all timesteps [hard]
% - Compute gradient of road inequality constraints [medium]
X0= [287;5;-176;0;2;0];

% rng(0); 
n = 25;
Pose = [randi([900,960],1,n);randi([440,520],1,n)];

xhat = 370; yhat = 140; Phat = [xhat;yhat];
Pinit = [X0(1);X0(3)];
Pend = [TestTrack.cline(1,end);TestTrack.cline(2,end)];
[m,nPts] = size(TestTrack.cline);

% Plot Track
figure;hold;
plot(TestTrack.bl(1,:),TestTrack.bl(2,:), 'k')
plot(TestTrack.cline(1,:),TestTrack.cline(2,:), '--k')
plot(TestTrack.br(1,:),TestTrack.br(2,:), 'k')
plot([TestTrack.bl(1,end),TestTrack.br(1,end)], ... 
    [TestTrack.bl(2,end),TestTrack.br(2,end)], 'r')
hold;


d = vecnorm(TestTrack.cline - Pinit);
eps = 1;
Idx = find(d <= eps);

while(isempty(Idx))
    eps = eps + 1;
    Idx = find(d <= eps);
end

idx = Idx(1);
prev_idx = max(1-1, 1);
next_idx = idx+1;

bl_vec = [TestTrack.bl(1,[prev_idx,next_idx]);TestTrack.bl(2,[prev_idx,next_idx])];
br_vec = [TestTrack.br(1,[prev_idx,next_idx]);TestTrack.br(2,[prev_idx,next_idx])];

cp = ontrack(Pinit, bl_vec(:,1), bl_vec(:,2), br_vec(:,1), br_vec(:,2));

% fprintf('Pose (%.2f, %.2f) - nearest point (%.2f, %.2f): %i \n', ...
%     Pose(1,i), Pose(2,i),TestTrack.cline(1,idx),TestTrack.cline(2,idx),cp)

hold;
plot(Pinit(1),Pinit(2),'og')
plot(Pend(1),Pend(2),'or')
plot(TestTrack.cline(1,idx),TestTrack.cline(2,idx), 'o')
plot(bl_vec(1,:),bl_vec(2,:), 'ob')
plot(br_vec(1,:),br_vec(2,:), 'ob')
hold;


[LB, UB] = bounds(TestTrack, 10);
size(LB)

[g,h,dg,dh]=nonlcon(z, TestTrack);

function cp = ontrack(pose, bl_prev, bl_next, br_prev, br_next)
    
    l_vec = bl_next - bl_prev; % l = l_i+1 - l_i-1
    r_vec = br_next - br_prev; % r = r_i+1 - r_i-1
    
    pl = pose - bl_prev; % p_l = p - l_i-1
    pr = pose - br_prev; % p_r = p - r_i-1
    
    c1 = l_vec(1)*pl(2) - l_vec(2)*pl(1);
    c2 = r_vec(1)*pr(2) - r_vec(2)*pr(1);

    cp = c1*c2;
end

function dg = diff_g(pose, bl_prev, bl_next, br_prev, br_next)
%     syms Lx Ly Rx Ry ly lx rx ry X Y;
%     g = (Lx*(Y - ly)-Ly*(X-lx))*(Rx*(Y - ry)-Ry*(X-rx));
%     diff(g,X);
%     diff(g,Y);

    dg = [0,0];
    l_vec = bl_next - bl_prev; % l = l_i+1 - l_i-1
    r_vec = br_next - br_prev; % r = r_i+1 - r_i-1
    
    pl = pose - bl_prev; % p_l = p - l_i-1
    pr = pose - br_prev; % p_r = p - r_i-1
    
    dg(1) = r_vec(2)*(l_vec(2)*pl(1) - l_vec(1)*pl(2)) + ...
            l_vec(2)*(r_vec(2)*pr(1) - r_vec(1)*pr(2));
    dg(2) = - r_vec(1)*(l_vec(2)*pl(1) - l_vec(1)*pl(2)) - ...
              l_vec(1)*(r_vec(2)*pr(1) - r_vec(1)*pr(2));
end

function [lb, ub] = bounds(TestTrack, StepsPerPoint)

    [m,nPts] = size(TestTrack.cline);
%     numState = 6;
%     numInput = 2;
    nsteps = StepsPerPoint * nPts;
    
    lb_u = [-0.5;-5000];
    ub_u = [0.5;5000];

    bound_X = [TestTrack.bl(1,1), TestTrack.bl(1,2), ... 
               TestTrack.br(1,1), TestTrack.br(1,2)];
    bound_Y = [TestTrack.bl(2,1), TestTrack.bl(2,2), ... 
               TestTrack.br(2,1), TestTrack.br(2,2)];
    phi_init = TestTrack.theta(1);
    
    lb = [min(bound_X); -Inf; min(bound_Y); -Inf; phi_init-(pi/2); -Inf];
    ub = [max(bound_X); Inf; max(bound_Y); Inf; phi_init+(pi/2); Inf];
    

    for i=1:nPts
        prev_idx = max(i-1, 1);
        next_idx = min(i+1, 246);
        
        bound_X = [TestTrack.bl(1,prev_idx), TestTrack.bl(1,next_idx), ... 
                   TestTrack.br(1,prev_idx), TestTrack.br(1,next_idx)];
        bound_Y = [TestTrack.bl(2,prev_idx), TestTrack.bl(2,next_idx), ... 
                   TestTrack.br(2,prev_idx), TestTrack.br(2,next_idx)];
        phi_init = TestTrack.theta(i);
        
        lb_x = [min(bound_X); -Inf; min(bound_Y); -Inf; phi_init-(pi/2); -Inf];
        ub_x = [max(bound_X); Inf; max(bound_Y); Inf; phi_init+(pi/2); Inf];
        
        for num = 1:StepsPerPoint
            lb=[lb;lb_x];
            ub=[ub;ub_x];
        end
    end

    for i=1:nsteps
        ub=[ub;ub_u];
        lb=[lb;lb_u];
    end
end

function [g,h,dg,dh]=nonlcon(z, TestTrack)
    if size(z,2)>size(z,1)
        z = z' ;
    end
    numState = 6;
    numInput = 2;
    numXandU = 8;
    nsteps = (size(z,1)+2)/numXandU;

    dt = 0.01;

    zx=z(1:nsteps*numState);
    zu=z(nsteps*numState + 1:end);

    g = zeros(nsteps,1) ;
    dg = zeros(nsteps,numXandU*nsteps-2) ;

    h = zeros(numState*nsteps,1) ;
    dh = zeros(numState*nsteps,numXandU*nsteps-2);
    
    for i=1:nsteps
        % At given position (Xi, Yi) at Zi, find nearest centerline
        % Use the index of nearest centerline to calculate g function
        
%         pos = [Xi Yi];
        car_pos = [z(numState*i-numState+1); z(numState*i-numState+3)];
        d = vecnorm(TestTrack.cline - car_pos);
        eps = 1;
        Idx = find(d <= eps);
        while(isempty(Idx))
            eps = eps + 1;
            Idx = find(d <= eps);
        end

        idx = Idx(1);
        prev_idx = max(idx-1, 1);
        next_idx = min(idx+1, size(TestTrack.cline, 2));

        bl_vec = [TestTrack.bl(1,[prev_idx,next_idx]);TestTrack.bl(2,[prev_idx,next_idx])];
        br_vec = [TestTrack.br(1,[prev_idx,next_idx]);TestTrack.br(2,[prev_idx,next_idx])];

        g = ontrack(Pinit, bl_vec(:,1), bl_vec(:,2), br_vec(:,1), br_vec(:,2));
        
        dgi_dzj = diff_g(pose, bl_prev, bl_next, br_prev, br_next);
        dg(i,numState*i-numState+1) = dgi_dzj(1);
        dg(i,numState*i-numState+2) = dgi_dzj(2);
        
%         if i==1
%             h(1:3) = z(1:3,:) ;
%             dh(1:3,1:3)=eye(3); 
%         else
%             h(3*i-2:3*i) = zx(3*i-2:3*i)-zx(3*i-5:3*i-3)-...
%                                dt*odefun(zx(3*i-5:3*i-3),zu(2*i-3:2*i-2)) ;
%             dh(3*i-2:3*i,3*i-5:3*i) = [-eye(3)-dt*statepart(zx(3*i-5:3*i-3),zu(2*i-3:2*i-2)),eye(3)] ;
%             dh(3*i-2:3*i,3*nsteps+2*i-3:3*nsteps+2*i-2) = -dt*inputpart(zx(3*i-5:3*i-3),zu(2*i-3:2*i-2)) ;
%         end
    end
end