
% Given an Ny x Nx x Nt matrix, an Nt x 1 vector of channel numbers, and
% measures of dt,ow, and on, produce a filtered time-series.


NOISE_STD = 500;

burn_in = 3;
filt_order = 0;
proc_noise_var = 300^2;
meas_noise_var = 700^2;


dataFolder = './MultispectralSheppLoganData/testData_short_SelectedWavelengths_OS1_localdynam_WBrth';

if ~exist('sampleData','var')||~(strcmp(pwd,dataFolder))
%     cd(dataFolder);
    load(fullfile(dataFolder,'testData.mat'));
else
    clearvars -except sampleData dataFolder
end


% The scanData is the one-wavelength-at-a-time train of data. 
scanData = sampleData.images.SampledData;
scanData = scanData + NOISE_STD * randn(size(scanData));

% Get the whole dataset as well as whatever meta is necessary. 
completeDataset = sampleData.images.CompleteData;
scanMeta = sampleData.parameters;
channelSet = scanMeta.sampling.Wavelengths;
Wavelengths = channelSet;
frameSampledWLs = sampleData.timeCourses.Wavelengths;



timeList = sampleData.timeCourses.Time;
[~,channelIndexList] = ismember(frameSampledWLs,Wavelengths);


SW_filteredData = kalataFilter(scanData,channelIndexList,timeList,proc_noise_var,meas_noise_var,0,burn_in);
alpha_filteredData = kalataFilter(scanData,channelIndexList,timeList,proc_noise_var,meas_noise_var,1,burn_in);
alphabeta_filteredData = kalataFilter(scanData,channelIndexList,timeList,proc_noise_var,meas_noise_var,2,burn_in);
trueData = squeeze(sum(completeDataset,3));
%%
for k = 1:numel(timeList)
    ssim_SW(k) = calcWholeImageSSIM(SW_filteredData(:,:,:,k),trueData(:,:,:,k));
    ssim_a(k) = calcWholeImageSSIM(alpha_filteredData(:,:,:,k),trueData(:,:,:,k));
    ssim_ab(k) = calcWholeImageSSIM(alphabeta_filteredData(:,:,:,k),trueData(:,:,:,k));

end
figure;
plot(ssim_SW); hold on; plot(ssim_a); plot(ssim_ab);





% Functions as supplement to MATLAB's built-in SSIM calculations. 
function wholessim = calcWholeImageSSIM(im,ref)
    C1 = 0.02;
    C2 = 0.01;
    
    mean_ref = mean(ref(:));
    mean_im = mean(im(:));
    
    covMat = cov(ref(:),im(:));
    
    Nt1 = (2*mean_ref*mean_im + C1);
    Nt2 = (2*covMat(1,2) + C2);
    Dt1 = (mean_ref.^2 + mean_im.^2 + C1);
    Dt2 = (covMat(1,1) + covMat(2,2) + C2);
    
    wholessim = (Nt1.*Nt2)./(Dt1.*Dt2);
    
    
end
