clear all
close all

disp("")
disp("This program tests the computation of a deformation sub-space on a 3D tetrahedral mesh.")

% Initial data

disp("")
disp("Loading data...")
fflush(stdout);

%[V, F, handle_indices] = tetrahedron_3()
%[V, F, handle_indices] = octahedron();
%[V, F, handle_indices] = icosahedron();
[V, F, h_p, h_r] = armadillo();
%[V, F, handle_indices] = truc_LD();
%[V, F, handle_indices] = truc_HD();
E = edges(F);
Tri = faces(F);

n = size(V,1);
f = size(F,1);
k = size(E,1);
t = size(Tri,1);

%handle_indices = int32(rand(1,100)*n);

disp(["... loaded model with ", num2str(n), " vertices, ",num2str(k)," edges, ", num2str(t)," triangles and ", num2str(f), " tetrahedra."])
fflush(stdout);


% Energy matrix computation

disp("")
disp("Computing energy matrix...")
fflush(stdout);
tic

Adj = vertices_faces_adjacency_matrix(Tri);
Lcr = face_face_laplacian_3D(V, E, Tri, F);
Ncr = face_face_normal_derivative_3D(V,E, Tri,F);
Mcr = face_mass_matrix_3D(V, Tri, F);

%A = Adj'*(Lcr+Ncr)'*inverse(Mcr)*(Lcr+Ncr)*Adj;
A = Adj'*(Lcr+Ncr)'*Mcr*(Lcr+Ncr)*Adj;

t=toc;
disp(["...done in ",num2str(t),"s."])
fflush(stdout);


% Energy matrix tests

disp("")
disp("Performing precision tests...")
fflush(stdout);

O = ones(n, 1);
O_result = A*O;
O_result_2 = O_result'*O_result
if(O_result_2 < 1e-5)
  disp("Constant precision : CHECK.")
else
  disp("Constant precision : FAILURE.")
endif

V_result = A*V;
V_result_2 = V_result'*V_result;
tr_V_result_2 = trace(V_result_2)
if(tr_V_result_2 < 1e-5)
  disp("Linear precision : CHECK.")
else
  disp("Linear precision : FAILURE.")
endif

fflush(stdout);


% Weights computation

disp("")
disp("Computing weights...")
fflush(stdout);
tic

handle_indices_r = cell2mat(h_r);
handle_indices = [h_p, handle_indices_r];

m_p = length(h_p);
m_r = length(handle_indices_r);
m = m_p + m_r;

r = length(h_r);
h = m_p+r*4;

disp(["There is ", num2str(m_p), " point handles and ", num2str(r), " region handles (containing ",num2str(m_r)," vertices in total)."])

I_m_p = speye(m_p);

cols = zeros(m_r,4);
ci = 1;
for i=1:r
  cols(ci:(ci+length(h_r{i})-1), :) = 4*(i-1);
  ci = ci+length(h_r{i});
end
cols = cols + kron(1:4,ones(m_r, 1));
J_2 = sparse([(1:m_r)', (1:m_r)', (1:m_r)', (1:m_r)'], cols, [V(handle_indices_r, :), ones(m_r, 1)], m_r, r*4);

J = [I_m_p, sparse(m_p, r*4); sparse(m_r, m_p), J_2];



S = sparse(m, n);
for i=1:m
  S(i, handle_indices(i)) = 1;
endfor

T = sparse(n-m, n);

j=1;
for i=1:n
  if(size(find(handle_indices-i), 2) == m)
    T(j, i) = 1;
    j = j+1;
  endif
endfor

%W = S'*J-T'*inverse(T*A*T')*T*A*S'*J;
W = S'*J-T'*((T*A*T')\T)*A*S'*J;

t=toc;
disp(["...done in ",num2str(t),"s."])
fflush(stdout);


%return

%% Weights visualization


[~,on_b] = on_boundary(F); % for each vertex of each tetra, 1 if the face opposite the vertex is on boundary, 0 else
[on_b_f,on_b_c] = find(on_b); % indices of tetrahedra on border, and subscript index of the vertex opposed the triangle on border for each tetrahedron

% vertices of faces on the border
Pi = mod(on_b_c+1,4)+1;
Qi = mod(on_b_c+2,4)+1;
Ri = mod(on_b_c+0,4)+1;

P = F(sub2ind(size(F), on_b_f, Pi));
Q = F(sub2ind(size(F), on_b_f, Qi));
R = F(sub2ind(size(F), on_b_f, Ri));


disp("")
disp("Visualizing weights (see fig. 1 for scalar weights, fig. 2 for vectorial weights)")

figure()
for i=1:m_p
  subplot(round(sqrt(m_p)), ceil(sqrt(m_p)), i)
  trisurf([P,Q,R], V(:, 1), V(:, 2), V(:, 3), W(:,i), "facecolor", "interp", "edgecolor", "k")
  hold on
  plot3(V(h_p(i),1), V(h_p(i), 2), V(h_p(i), 3), 'kx', "linewidth", 2, "markersize", 10)
  axis square
  axis equal
  colorbar("EastOutside")
  title (["w_{",num2str(i), "}"]);
  view(3)
endfor

component = {'x', 'y', 'z', 'w'};
figure()
for i=1:r
  for j=1:4
    subplot(4, r, (j-1)*r + i)
    trisurf ([P,Q,R], V(:,1), V(:,2), V(:, 3), W(:,m_p+(i-1)*4+j), "facecolor", "interp", "edgecolor", "k")
    hold on
    plot3(V(h_r{i},1), V(h_r{i},2), V(h_r{i},3), 'ko', "linewidth", 2, "markersize", 2)
    axis square
    axis equal
    colorbar("EastOutside")
    title (["w_{",num2str(i), ",",component{j},"}"]);
    view (3);
  end
end


% Rest pose reproduction test

disp("")
disp("Performing rest post reproduction test...")


C_1 = V(h_p, :);
id = [eye(3);zeros(1,3)];
C_2 = kron(ones(r,1), id);
H = [C_1;C_2];

V_reconstruct = W * H;
V_diff = V - V_reconstruct;
V_diff_2 = V_diff'*V_diff;
tr_V_diff_2 = trace(V_diff_2)

if(tr_V_diff_2 < 1e-5)
  disp("Rest pose reproduciton : CHECK.")
else
  disp("Rest pose reproduciton : FAILURE.")
endif


figure()
%{
subplot(1,2,1)
hold on
surf1 = trisurf([P,Q,R], V(:,1), V(:,2), V(:,3), sum(V_diff.^2,2).^(0.5));
colorbar("EastOutside")
plot3(H(:,1), H(:,2), H(:,3), 'go', "linewidth", 2, "markersize", 10)
axis square
axis equal
view(3)

subplot(1,2,2)
%}
hold on
surf2 = trisurf([P,Q,R], V_reconstruct(:,1), V_reconstruct(:,2), V_reconstruct(:,3), sum(V_diff.^2,2).^(0.5));
colorbar("EastOutside")
plot3(H(1:m_p,1), H(1:m_p,2), H(1:m_p,3), 'ro', "linewidth", 2, "markersize", 10)
plot3(V(handle_indices_r, 1), V(handle_indices_r,2), V(handle_indices_r,3), 'ro', "linewidth", 2, "markersize", 3)
axis square
axis equal
view(3)
title("Rest pose reproduction test")

%linkprop([surf1, surf2], 'view')

disp("Displaying result (see fig. 3)")

%return;

% Deformation test

disp("")
disp("Visualizing deformation (see fig. 4)")


H_deform = zeros(size(H));

deform_disp = 2*(rand(m_p, 3)-0.5*ones(m_p, 3));
H_deform(1:m_p, :) = H(1:m_p, :) + deform_disp;


angles = pi*(rand(r,1) - 0.5*ones(r, 1));

axis0 = rand(r,3);
axis0 = axis0 ./ (sum(axis0.^2,2).^0.5 * ones(1,3));

axis1 = cross(axis0, [ones(r, 1),zeros(r,2)],2);
axis1 = axis1 ./ (sum(axis1.^2,2).^0.5 * ones(1,3));

axis2 = cross(axis0, axis1,2);

rot = zeros(4*r, 3);
%rot(1:4:end, 1) =  cos(angles);
%rot(1:4:end, 2) = -sin(angles);
%rot(2:4:end, 1) =  sin(angles);
%rot(2:4:end, 2) =  cos(angles);
%rot(3:4:end, 3) =  1;


%centers = rand(r,3);



%for i=1:r
%  mat = [axis2(i,:)', axis1(i,:)', axis0(i,:)'];
%  rot((1+(i-1)*4):(i*4-1), :) = mat*rot((1+(i-1)*4):(i*4-1), :)*inverse(mat);
%end


t_tot = 11;
posX = zeros(t_tot, m_r);
posY = zeros(t_tot, m_r);
posZ = zeros(t_tot, m_r);

cc = 1;
for i=1:r
  mat = [axis2(i,:)', axis1(i,:)', axis0(i,:)']
  rot((1+(i-1)*4):(i*4-1), :) = mat*[cos(angles(i)), sin(angles(i)), 0; -sin(angles(i)), cos(angles(i)), 0; 0, 0, 1]*inverse(mat);
  for tt = 1:t_tot
    t = (tt-1.0)/(t_tot-1.0);
    rotation = mat*[cos(angles(i)*t), -sin(angles(i)*t), 0;sin(angles(i)*t), cos(angles(i)*t), 0; 0, 0, 1]*inverse(mat);
    
    pos = rotation * V(h_r{i}, :)';
    
    posX(tt, cc:(cc + length(h_r{i})-1)) = pos(1,:);
    posY(tt, cc:(cc + length(h_r{i})-1)) = pos(2,:);
    posZ(tt, cc:(cc + length(h_r{i})-1)) = pos(3,:);
  end
  cc = cc + length(h_r{i});
end

H_deform((m_p+1):end, :) = rot;


V_deform = W * H_deform;

%export_tetra_mesh(V_deform, F, "output/test.vtk")

figure()
hold on

trisurf([P,Q,R], V_deform(:,1), V_deform(:,2), V_deform(:,3), sum((V-V_deform).^2,2).^(0.5))
colorbar("EastOutside")

plot3(H(1:m_p,1), H(1:m_p,2), H(1:m_p,3), 'go', "linewidth", 2, "markersize", 10)
plot3(H_deform(1:m_p,1), H_deform(1:m_p,2), H_deform(1:m_p,3), 'bo', "linewidth", 2, "markersize", 10)
plot3([H(1:m_p, 1)'; H_deform(1:m_p, 1)'], [H(1:m_p, 2)'; H_deform(1:m_p, 2)'], [H(1:m_p, 3)'; H_deform(1:m_p, 3)'], 'k', "linewidth", 2)
plot3(V_deform(handle_indices_r, 1), V_deform(handle_indices_r,2), V_deform(handle_indices_r,3), 'ro', "linewidth", 2, "markersize", 3)
%plot3([V(handle_indices_r, 1)'; V_deform(handle_indices_r, 1)'], [V(handle_indices_r, 2)'; V_deform(handle_indices_r, 2)'], [V(handle_indices_r, 3)'; V_deform(handle_indices_r, 3)'], 'k', "linewidth", 2)
plot3(posX, posY, posZ, 'k', "linewidth", 1)
axis square
axis equal
view(3)
title("Deformation test")

