% Phase 1 Phase 2 using pseudospectral

% The RHS is the whole NMT!

close all;

clear all;

clc;

%% Add paths

addpath("folder_sym/");

addpath("cheb_helper_functions/");

addpath("double_pendulum_dynamics/");

%% folders and stuff

folder_sym = "folder_sym";

%% Parameters

use_out_fcn = false;

NB = 2;

% obstacles = [-1.6, 0.574, 0, 0.4];

% obstacles = [-1.9, 0.6, 0, 0.4];

obstacles = [-1.51, 1.309, 0, 0.3; ...

-1.819, -0.83, 0, 0.3;

-1.98, 0.28, 0, 0.3];

n_obstacles = size(obstacles, 1);

k_ph1 = 0;

k_cons_ph1 = 1e1;

c_cons_ph1 = 1e3;

s_max_ph1 = 1e2;

scaling_ph1 = 1.0;

abs_tol = 1e-4;

rel_tol = 1e-4;

x0 = [-pi/2; 0; 0; 0]; % initial angle and velocity at t = -1

xf = [pi/2; 0; 0; 0];

s_span_ph1 = [0 s_max_ph1];

p = 11; % degree of the polynomial

%% Generate the symbolic RHS

mass_fn = "D_" + string(NB) + "link_Ac_cm";

cor_fn = "C_" + string(NB) + "link_Ac_cm";

gravity_fn = "G_" + string(NB) + "link_Ac_cm";

g_ph1 = @(obs_pos, fk_val, obs_rad) obs_rad^2 - (norm(obs_pos-fk_val))^2;

act_fn_ph1 = @(g_val, c_cons) 0.5 + 0.5 .* tanh( c_cons * g_val );

barr_fn_ph1 = @(g_val, S_val, k_cons) k_cons .* exp(g_val) * S_val;

fk_fn = @(x) [-(cos(x(1, :)) + cos(x(1, :) + x(2, :)));

sin(x(1, :)) + sin(x(1, :) + x(2, :));

0 * x(1, :)];

% Generate the symbolic ph1 lagrangian terms

[rhs_aghf_ph1, dL_dx_name_ph1, dL_dxd_name_ph1, L_ph1, b_ph1, g_max_ph1_fn, ...

db_dx_fn] = ...

sym_lagr_ph1(NB, ...

folder_sym, mass_fn, cor_fn, gravity_fn, fk_fn,...

obstacles, g_ph1, act_fn_ph1, barr_fn_ph1, "ph1");

%% Compute Chebyshev stuff

chebnodes = (cos( pi * (0:p) / p))'; %% note this is how our state vector is organized so the time 1 is the first point and time -1 is the last point

D = chebdiff(chebnodes, p); %% this is used to compute the first derivative

D2 = D^2; %% use this to compute the second derivative

% Generate initial guess which is a line

init_x1 = linspace(xf(1), x0(1), numel(chebnodes) ).';

init_x2 = linspace(xf(2), x0(2), numel(chebnodes)).';

init_ps = [init_x1, init_x2];

init_dps = D * init_ps; % derivative of the initial guess

init_aps = [ init_ps(2:end-1, :) init_dps(2:end-1, :) ];

init_aps = init_aps(:);

%% helper functions for plotting and evaluating actional functional in

% chebyshev basis

syms r % time

syms c [p, 1] %

dt = 1e-2;

t = (-1 : dt : 1)';

% these are all the integrals, note we are including the integration

% constant as well to make the mat function easy to use

q = matlabFunction(chebyshevT([0 : p], r));

q_mat = apply_func_to_vector(q, t);

%%

rhs_fun_ph1 = str2func(rhs_aghf_ph1);

L_fun_ph1 = str2func(L_ph1);

g_max_lam_ph1 = str2func(g_max_ph1_fn);

%% setup output and integration properties

if(use_out_fcn)

options = odeset('AbsTol', 1e-10, 'RelTol', 1e-10, 'OutputFcn', ...

@(s,aps,flag) outputaghf(s, aps, flag, q, D, D2, k_ph1, x0, xf, p, NB, ...

chebnodes, k_cons_ph1, c_cons_ph1, ...

rhs_fun_ph1, L_fun_ph1, scaling_ph1, ...

g_max_lam_ph1) );

else

options = odeset('AbsTol', 1e-10, 'RelTol', 1e-10 );

end

%% Do the actual integration

% s, aps, D, D2, k, k_cons, c_cons, x, xf, p, NB, rhs_fun

tic;

[ tau, sol ] = ode15s(@(s, aps) aghf(s, aps, D, D2, k_ph1, k_cons_ph1, ...

c_cons_ph1, x0, xf, p, NB, rhs_fun_ph1, scaling_ph1), ...

s_span_ph1, init_aps, options );

t_solve = toc

%% Compute the FK of the initial guess

% init_ps

fk_ini = get_fk_traj(p+1, init_ps.', fk_fn);

make_fk_plot(fk_ini, obstacles, "Initial guess");

%% Compute the FK of the solution

aps = sol(end,:)';

% aps is a column vector whose elements are q1, q2, q1dot, q2dot

aps_helper = reshape( aps, p-1, NB * 2 );

% note aps_helper = [ q1 q2 q1_dot q2_dot ] but without the initial and final conditions

x = [ xf'; aps_helper; x0' ];

fk_end = get_fk_traj(p+1, x.', fk_fn);

make_fk_plot(fk_end, obstacles, "End soln");

%% Get solution of each joint position and velocity

figure;

subplot(2,2,1);

q1_t = q_mat * chebinterp(x(:,1));

plot( t, q1_t );

hold on;

grid on;

subplot(2,2,2);

q2_t = q_mat * chebinterp(x(:,2));

plot( t, q2_t);

for ii = 1 : length(chebnodes)

chebnode = chebnodes(ii);

xline(chebnode, 'r--');

end

grid on;

subplot(2,2,3);

q1d_t = q_mat * chebinterp(x(:,3));

plot( t, q1d_t );

for ii = 1 : length(chebnodes)

chebnode = chebnodes(ii);

xline(chebnode, 'r--');

end

grid on;

subplot(2,2,4);

q2d_t = q_mat * chebinterp(x(:,4));

plot( t, q2d_t );

for ii = 1 : length(chebnodes)

chebnode = chebnodes(ii);

xline(chebnode, 'r--');

end

grid on;

%% Store the parameters required to solve the same AGHF problem in python/cpp and

% compare the solutions

% fp_data = "../../psaghf/tests/matlab_files/data_double_pendulum_aghf_soln_sphere_constraints.mat";

%

%

c_cons = 50.0;

init_ps_values = init_aps(:);

%

% save(fp_data, "g_v", "m_v", "NB", "Nact", "x0", "xf", "s_span", "K", "c_cons", "N", ...

% "chebnodes", "D", "D2", "k_cons", "obs_mat_values", "aps", "q1_t", "q2_t", "q1d_t", "q2d_t", ...

% "dt", "t", "abs_tol", "rel_tol", "init_ps_values");