%% convert hamlyn heart video to image sequence

% workingDir = tempname;
% mkdir(workingDir)
% mkdir(workingDir,'images')
% 
% shuttleVideo = VideoReader('~/Downloads/hamlyn_heart.avi');
% 
% ii = 1;
% while hasFrame(shuttleVideo)
%    img = readFrame(shuttleVideo);
%    filename = [sprintf('%03d',ii) '.jpg'];
%    fullname = fullfile(workingDir,'images',filename);
%    imwrite(img,fullname)    % Write out to a JPEG file (img1.jpg, img2.jpg, etc.)
%    ii = ii+1;
% end
% 
% imageNames = dir(fullfile(workingDir,'images','*.jpg'));
% imageNames = {imageNames.name}';
% 
% 
% for ii = 1:length(imageNames)
%    img = imread(fullfile(workingDir,'images',imageNames{ii}));
%    grayImages(:,:,ii) = rgb2gray(img);
% end

%% Try DMD on it
clear variables
load data/HamlynHeartImages.mat;

disp('Data preparation...')

image_start  = 1;% frame number to start from
num_images = 100;

allImages = [];
% for k = 1:size(grayImages,3)
for k = image_start:image_start+num_images 
    image_k = grayImages(:,:,k);
%     allImages(:,k) = double(image_k(:));
    allImages = [allImages, double(image_k(:))];
end

XallPrescaled = [allImages(:,2:end);
                 allImages(:,1:end-1)];
% XallPrescaled = allImages;

% remove constant rows which make XallPrescaled lose rank
[XallPrescaledProcessed,ProcessSettings] = removeconstantrows(XallPrescaled);
XallPrescaledProcessed = XallPrescaledProcessed';


%% do Kernel DMD via object implementation
D = kernel_dmd(grayImages(:,:,image_start:image_start+num_images),'','gaussian','true');
D.regression(1e6,49);
rec = D.reconstruct();

disp('DMD Regression...')
%% kernel dmd
disp('Scaling data...')
[Xall,scaleFactor,scaleOffset] = ScaleData(XallPrescaledProcessed,'zscore'); % zscore may cause problems with body_length, which is constant
% Xall = XallPrescaledProcessed;
% scaleFactor = 1;
% scaleOffset = 0;
X = Xall(1:end-1,:);
X2 = Xall(2:end,:);

%% Kernel DMD
disp('Doing Kernel DMD...')
% kernel function
% inputs: x1, x2 are vectors of the same size.
% kernel = @(x1,x2) (1 + dot(x1,x2))^10; % polynomial kernel, doesnt work that well. often fails during reconstruction
sigma2 = 1e6; 
kernel = @(x1,x2) exp(-norm(x1-x2)^2/sigma2); % Gausssian kernel

% for j = 1:size(X,1)
%     disp(j)
%     for k = 1:size(X2,1)
%         Ahat(j,k) = kernel(X(k,:),X2(j,:)); % Y'Y2 is Ahat in Williams et al.
%         Ghat(j,k) = kernel(X(k,:),X(j,:)); % Y*Y is the Gramian Ghat in Williams et al.
%     end 
% end

tic;
Ahat = GaussianKernel(X,X2,sqrt(sigma2));
Ghat = GaussianKernel(X,X,sqrt(sigma2));

%%

% TODO: truncate to get low-order model.
% use Ghat to compute Sigma (called 'S') and V
% [Qfull,S2full] = eig(Ghat); % since it's true that Ghat*Q = Q*Sigma^2

[Qfull, S2full] = eigs(Ghat,size(Ghat,1),'largestabs','Tolerance',1e-19); % use increased tolerance with GaussianKernel

% S2full = flipud(fliplr(S2full)); % make largest singular values first instead of largest singular values last
% Qfull = fliplr(Qfull); 
Sfull = sqrtm(S2full);

% truncate
% r = size(Ghat,1)-1; % number of dimensions to keep
r = floor(size(Ghat,1)/2);
Q = Qfull(:,1:r);
S = Sfull(1:r,1:r);

figure(2);
clf;
semilogy(diag(Sfull),'o');
hold on
semilogy(diag(S),'.');

% Sinv = diag(ones(size(S,1),1)./diag(S)); % better numerical stability?
% Khat = (Sinv*Q')*Ahat*(Q*Sinv);
Khat = (S\Q')*(Ahat)*(Q/S); % Khat, an approximation of the Koopman operator
[Vhat, Lambda] = eig(Khat);

if(rank(Vhat) < size(Vhat,1))
    error('Vhat is not full rank (rank = %d). I didnt code up the case where Vhat isnt.', rank(Vhat));
end

Phi = Q*S*Vhat; % eigenfunctions of Khat 
Xi = Phi\X; 
toc;
X2hat = Phi*Lambda*Xi;
if(norm(imag(X2hat)) > 1e-3)
    error('X2hat has non-negligible imaginary component.')
end
X2hat = real(X2hat); % this is quite an accurate reconstruction of X2.

disp('Finished computing kernel-DMD regression. Attempting simulation from IC.');

%% simulate forward
num_modes = size(Q,1); % i THINK this is the correct one
tic;
xhat(1,:) = X(1,:); 
% xhat(1,1) = xhat(1,1) + 0.1;
% xhat(1,:) = xhat(1,:) + 0.001*randn(1,size(xhat,2));
for t = 2:2*size(X,1)
    % calculate phi given x
    prev_x = xhat(t-1,:);
%     kernelmat = nan(1,num_modes);
%     for M = 1:num_modes
%         kernelmat(M) = kernel(prev_x,X(M,:));
%     end
    kernelmat = GaussianKernel(prev_x,X, sqrt(sigma2));
    kernelmat = kernelmat(1:num_modes)';
    phi_at_t = kernelmat*Q/S*Vhat;

    % estimate 
    xnext = phi_at_t*Lambda*Xi;
    if(norm(imag(xnext)) > 1e-3)
        error('xnext isnt real (has non-negligble imaginary component)')
    end
    xhat(t,:) = real(xnext);
end
toc;
[xhatDmd,dmd_reconstruction] = D.predict(D.X(1,:),200);

%% Check
figure(1)
clf
subplot(2,1,1)
plot(Xall(:,1));
hold on
plot([nan(1);X2hat(:,1)],'--'); 
plot(xhat(:,1),'--');
legend('true','1-step recon', 'n-step recon');
title('Reconstruction element 1')
subplot(2,1,2)
plot(Xall(:,6));
hold on
plot([nan(1);X2hat(:,6)],'--'); 
plot(xhat(:,6),'--');
legend('true','1-step recon', 'n-step recon');
title('Reconstruction element 6');

disp('Simulated from IC. Unscaling and casting data.');

%% 
% offsetx = repmat(scaleOffset,size(xhat,1),1);
% reconstruction = xhat*diag(scaleFactor) + offsetx;
for k = 1:size(xhat,1)
    xhatRescaled(k,:) = xhat(k,:).*scaleFactor + scaleOffset;
end
xhatRescaled = xhatRescaled';
xhatRescaled = removeconstantrows('reverse',xhatRescaled,ProcessSettings); % add removed rows back in.

for k = 1:size(xhatRescaled,2) % now each state/image is a column, to comply with removeconstantrows()
    reconstruction(:,:,k) = reshape(xhatRescaled(1:size(allImages,1),k),size(grayImages,1),size(grayImages,2));
end

disp('Done.')





%% replay
reconsructionHdl = implay(uint8(reconstruction));
set(reconsructionHdl.Parent,'Name',sprintf('Reconstruction: %d modes, sigma2 = %.2e.',r,sigma2));
originalHdl = implay(grayImages(:,:,image_start:image_start+num_images));
set(originalHdl.Parent,'Name','Original data');
dmdReconstructionHdl = implay(dmd_reconstruction);
set(dmdReconstructionHdl.Parent,'Name','DMD');
%% 

% figure(3);
% clf
% mode_number = 1;
% % dont use uint8 and imshow, U0 is unitary so columns have very small
% % values
% singleMode = surf(flipud(reshape(Qfull(:,mode_number),size(grayImages,1),size(grayImages,2))),'LineStyle','none');
% % singleMode = surf(flipud(reshape(U0(:,mode_number),size(grayImages,1),size(grayImages,2))),'LineStyle','none');
% colormap(flipud(gray))
% modeEig = Lambda(mode_number,mode_number);
% title(sprintf('Mode %d, eig %.3f + %.3fi', mode_number, real(modeEig), imag(modeEig)));
% % view from above
%  axis([0 360 -0.001 288 -1 0])    
%  view(0,90)
 
 
 function F = GaussianKernel(Xin,Yin,sigma)
    X = Xin./sigma;
    Y = Yin./sigma;
    
    % each term is |xi - yj|^2. Expanding:
    % |x-y|^2 = x'x - 2*x'y + y'y. these three lines are the three terms.
    X2 = sum(X.^2,2); %x'x as a column vector
    Y2 = sum(Y.^2,2); %y'y as a column vector
    XY = bsxfun(@minus, X2, (2*X)*Y.'); % add 2x'y terms
    total = bsxfun(@plus, Y2.', XY); % add y'y terms
    
    F = exp(-total)'; 

end