%% Set things up
close all
clear
clc
% set(groot,'defaulttextinterpreter','latex');
% set(groot, 'defaultAxesTickLabelInterpreter','latex');
% set(groot, 'defaultLegendInterpreter','latex');


da=demo_analysis;
demo = da.dataset_generation;
% [experiments,cs_time] = da.dataset_generation;
%%
figure
for i=1:1 %length(demo)
    i=2
    plot3(demo{i}.cart(:,2),demo{i}.cart(:,3),demo{i}.cart(:,4));
    hold on;
end
%% SWARM X
figure
for i=1:1 %length(demo)
    i=2
    plot(demo{i}.cart(:,4));
    hold on;
    plot(demo{i}.ft(:,4));
end


%%  SWARM Y
figure
for i=1:length(demo)
    plot(demo{i}.swarm_pose(:,3));
    hold on;
end

%%  SWARM Z
figure
for i=1:length(demo)
    plot(demo{i}.swarm_pose(:,4));
    hold on;
end
%%  GOAL X
figure
for i=1:length(demo)
    plot(demo{i}.goal_pose(:,2));
    hold on;
end
%%
for i=1:length(experiments)
    figure(i)
    for j=2:7
        plot(experiments{i}(:,1),smooth(experiments{i}(:,j),30)),hold on
    end
end
%% 3D Scatter Plot x_obj-y_obj-Ji
for j=1:length(experiments)
    figure(j)
    for i=2:7 
        scatter3(experiments{1,j}(:,8),experiments{1,j}(:,9),experiments{1,j}(:,i));
        hold on;
    end
end
%%
for j=1:length(experiments)
%     figure(j)
    for i=2:7 
        scatter3(experiments{1,j}(end,8),experiments{1,j}(end,9),experiments{1,j}(end,i));
        hold on;
    end
end
figure
for j=1:length(experiments)
%     figure(j)
    for i=2:7 
        scatter3(experiments{1,j}(end,8),experiments{1,j}(1,10),experiments{1,j}(end,i));
        hold on;
    end
end
%% Plot Fz/J6
for i=2:10
    figure(i)
    plot(smooth(experiments{1,i}(:,6),30),smooth(experiments{1,i}(:,13),30)); 
    hold on; 
end

%% Plot Fz/J6
k=1;
for j=8:10
    figure(k)
    for i=2:7
    %     figure(i)
        scatter(smooth(experiments{1,1}(:,i),30),smooth(experiments{1,1}(:,j),30)); 
        hold on; 
    end
    k=k+1;
end
%%
Data = experiments{1}';
aux = Data(8,:);
Data = [aux ;Data(2:7,:)];
%% Definition of the number of components used in GMM.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
nbStates = 10;
nbVar = size(Data,1);

%% Training of GMM by EM algorithm, initialized by k-means clustering.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
[Priors, Mu, Sigma] = EM_init_kmeans(Data, nbStates);
[Priors, Mu, Sigma] = EM(Data, Priors, Mu, Sigma);

%% Use of GMR to retrieve a generalized version of the data and associated
%% constraints. A sequence of temporal values is used as input, and the 
%% expected distribution is retrieved. 
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
input = experiments{1,2}(1,8:10);
[expData(2:nbVar,:), expSigma] = GMR(Priors, Mu, Sigma,  input, [2], [2:nbVar]);

%% Plot of the data
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
figure;
%plot 1D
for n=1:nbVar-1
  subplot(3*(nbVar-1),2,(n-1)*2+1); hold on;
  plot(Data(1,:), Data(n+1,:), 'x', 'markerSize', 4, 'color', [.3 .3 .3]);
  axis([min(Data(1,:)) max(Data(1,:)) min(Data(n+1,:))-0.01 max(Data(n+1,:))+0.01]);
  xlabel('t','fontsize',16); ylabel(['x_' num2str(n)],'fontsize',16);
end
%plot 2D
subplot(3*(nbVar-1),2,[2:2:2*(nbVar-1)]); hold on;
plot(Data(2,:), Data(3,:), 'x', 'markerSize', 4, 'color', [.3 .3 .3]);
axis([min(Data(2,:))-0.01 max(Data(2,:))+0.01 min(Data(3,:))-0.01 max(Data(3,:))+0.01]);
xlabel('x_1','fontsize',16); ylabel('x_2','fontsize',16);

%% Plot of the GMM encoding results
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%plot 1D
figure;
for n=1:nbVar-1
  subplot(3*(nbVar-1),2,4+(n-1)*2+1); hold on;
  plotGMM(Mu([1,n+1],:), Sigma([1,n+1],[1,n+1],:), [0 .8 0], 1);
  axis([min(Data(1,:)) max(Data(1,:)) min(Data(n+1,:))-0.01 max(Data(n+1,:))+0.01]);
  xlabel('t','fontsize',16); ylabel(['x_' num2str(n)],'fontsize',16);
end
%plot 2D
subplot(3*(nbVar-1),2,4+[2:2:2*(nbVar-1)]); hold on;
plotGMM(Mu([2,3],:), Sigma([2,3],[2,3],:), [0 .8 0], 1);
axis([min(Data(2,:))-0.01 max(Data(2,:))+0.01 min(Data(3,:))-0.01 max(Data(3,:))+0.01]);
xlabel('x_1','fontsize',16); ylabel('x_2','fontsize',16);

%% Plot of the GMR regression results
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%plot 1D
figure;
for n=1:nbVar-1
  subplot(3*(nbVar-1),2,8+(n-1)*2+1); hold on;
  plotGMM(expData([1,n+1],:), expSigma(n,n,:), [0 0 .8], 3);
  axis([min(Data(1,:)) max(Data(1,:)) min(Data(n+1,:))-0.01 max(Data(n+1,:))+0.01]);
  xlabel('t','fontsize',16); ylabel(['x_' num2str(n)],'fontsize',16);
end
%plot 2D
subplot(3*(nbVar-1),2,8+[2:2:2*(nbVar-1)]); hold on;
plotGMM(expData([2,3],:), expSigma([1,2],[1,2],:), [0 0 .8], 2);
axis([min(Data(2,:))-0.01 max(Data(2,:))+0.01 min(Data(3,:))-0.01 max(Data(3,:))+0.01]);
xlabel('x_1','fontsize',16); ylabel('x_2','fontsize',16);
%%
apple_pose=experiments{1}(:,8:14);
ee_pose=experiments{1}(:,15:21);


figure
R=quat2rotm([apple_pose(1000,7),apple_pose(1000,4:6)]);
quiver3(apple_pose(1000,1),apple_pose(1000,2),apple_pose(1000,3),...
    R(1,1),R(2,1),R(3,1),...
    'r','filled','LineWidth',6)
hold on
quiver3(apple_pose(1000,1),apple_pose(1000,2),apple_pose(1000,3),...
    R(1,2),R(2,2),R(3,2),...
    'g','filled','LineWidth',6)
quiver3(apple_pose(1000,1),apple_pose(1000,2),apple_pose(1000,3),...
    R(1,3),R(2,3),R(3,3),...
    'b','filled','LineWidth',6)
hold on
for i=1:2500:5000
    R=quat2rotm([ee_pose(i,7),ee_pose(i,4:6)]);
    quiver3(ee_pose(i,1),ee_pose(i,2),ee_pose(i,3),...
    R(1,1),R(2,1),R(3,1),...
    'r','filled','LineWidth',3)
    hold on
    quiver3(ee_pose(i,1),ee_pose(i,2),ee_pose(i,3),...
    R(1,2),R(2,2),R(3,2),...
    'g','filled','LineWidth',3)
    quiver3(ee_pose(i,1),ee_pose(i,2),ee_pose(i,3),...
    R(1,3),R(2,3),R(3,3),...
    'b','filled','LineWidth',3)
end


