clc
clear all
close all

% options
ROBOT = 'unicycle'
%TESTS = ["straightline/chill", "straightline/chill_errortheta_pisixths", "straightline/toofast", "straightline/chill_errory", "circle/start_center", "figure8/chill", "figure8/toofast", "square"]
TESTS = ["figure8/chill"]
CONTROLLER = 3

% main
s_ = size(TESTS);
for i = 1:length(TESTS)
    clearvars -except i s_ TESTS ROBOT CONTROLLER
    close all 
    
    % load simulation parameters common to all robots and all tests
    sim_data = load(["tests/robot_common.mat"]);

    TEST = convertStringsToChars(TESTS(i))
    
    % load test data (trajectory, etc)
    test_data = load(['tests/' TEST '/common.mat']);
    for fn = fieldnames(test_data)'
        sim_data.(fn{1}) = test_data.(fn{1});
    end
    
    % set trajectory and starting conditions
    sim_data.q0 = set_initial_conditions(sim_data.INITIAL_CONDITIONS);
    [ref dref] = set_trajectory(sim_data.TRAJECTORY, sim_data);
    sim_data.ref = ref;
    sim_data.dref = dref;
    
    % spawn a new worker for each contzroller
    % 1: track only
    % 2: track + 1step
    % 3: track + multistep
    % load controller-specific options
    data = load(['tests/' num2str(CONTROLLER) '.mat']);
    for fn = fieldnames(data)'
        sim_data.(fn{1}) = data.(fn{1});
    end

    % load robot-specific options
    % put here to overwrite any parameter value left over in the tests
    % .mat files, just in case
    data = load(['tests/' ROBOT '.mat']);
    for fn = fieldnames(data)'
        sim_data.(fn{1}) = data.(fn{1});
    end
    
    % initialize prediction horizon
    sim_data.U_corr_history = zeros(2,1,sim_data.PREDICTION_HORIZON);
    sim_data

    % simulate robot
    tic;
    [t, q, y, ref_t, U, U_track, Q_pred, Prob] = simulate_discr(sim_data);
    toc;

    disp('Done')
    
    % save simulation data
    f1 = [ TEST '/'  char(datetime, 'dd-MM-yyyy-HH-mm-ss')]; % windows compatible name
    f = ['results-' ROBOT '-costfun2-st/' f1];
    mkdir(f)
    % save workspace
    dsave([f '/workspace_composite.mat']);
    % save test file
    copyfile(['tests/' TEST], f);
    
    % save figures + plot results
    % plot results
    h = figure('Name', [TEST ' ' num2str(CONTROLLER)] );
    plot_results(t, q, ref_t, U, U_track, U_track);
    % save figures
    savefig(h, [f '/figure.fig']);
end


%% FUNCTION DECLARATIONS

% Discrete-time simulation
function [t, q, y, ref_t, U, U_track, Q_pred, Prob] = simulate_discr(sim_data)
    tc = sim_data.tc;
    steps = sim_data.tfin/tc
    
    q = sim_data.q0';
    t = 0;

    Q_pred = zeros(sim_data.PREDICTION_HORIZON,3, steps + 1);
    Prob = cell(steps+1, 1);

    [u_discr, u_track, q_pred, prob] = control_act(t(end), q(end, :), sim_data);
    U = u_discr';
    U_track = u_track';
    Q_pred(:, :, 1) = q_pred;
    Prob{1} = prob;

    if eq(sim_data.robot, 0)
        fun = @(t, q, u_discr, sim_data) unicycle(t, q, u_discr, sim_data);
    elseif eq(sim_data.robot, 1)
        fun = @(t, q, u_discr, sim_data) diffdrive(t, q, u_discr, sim_data);
    end

    for n = 1:steps
        tspan = [(n-1)*tc n*tc];
        z0 = q(end, :);
        
        opt = odeset('MaxStep', 0.005);
        [v, z] = ode45(@(v, z) fun(v, z, u_discr, sim_data), tspan, z0, opt);

        q = [q; z];
        t = [t; v];
                
        [u_discr, u_track, q_pred, prob] = control_act(t(end), q(end, :), sim_data);
        Prob{1+n} = prob;
        U = [U; ones(length(v), 1)*u_discr'];        
        U_track = [U_track; ones(length(v), 1)*u_track'];
        Q_pred(:, :, 1+n) = q_pred;
    end

    y1 = q(:, 1) + sim_data.b * cos(q(:,3));
    y2 = q(:, 2) + sim_data.b * sin(q(:,3));
    y = [y1, y2];

    ref_t = double(subs(sim_data.ref, t'))';
end

%%