
%% MakeSampleData generates datasets of a Shepp-Logan phantom according to concentration time course
% of provided spectral endmembers.
%
%
% ABOUT:
%     author      - Devin O'Kelly
%     date        - 28 Nov 2019
%     last update - 19 May 2020
%
%
%   



 %% Global dataset parameters.
clearvars; close all; clc;
    rng(172969105)
    datasetName = 'testData_';

    saveFolder = './MultispectralSheppLoganData';
    
    timecourseLength = 'short'; % 'long'
    NWL = 5; % Number of wavelengths to sample between 680 and 980.
    oversample = 1; % number of sequential times to repeat acquisition of a particular channel. 
    dynamics  = 'local'; % 'local','global'
    isbreathing = 1; % 1
    
    dt = 0.1;
    imageSize = 128;
    WLs = [700 730 756 800 850];
    
    
    % Creating the name.
        LENG  = [timecourseLength,'_'];
        NCHAN = 'SelectedWavelengths_';
        OSAMP = ['OS',num2str(oversample),'_'];
        DYN = [dynamics,'dynam_'];
        if isbreathing
            BRTH = 'WBrth';
        else
            BRTH = 'NoBrth';
        end
    
    datasetName = [datasetName LENG NCHAN OSAMP DYN BRTH];
    switch timecourseLength
        case 'short'
            Nt = 256;
        case 'long'
            Nt = 2048;
    end
    time = (1:Nt).*dt;
    phantomType = 'Positive Shepp-Logan';
    

 
    jitter = 0.0; % Variation in breathing rate.
    BreathPeriod = 1; %Sec
    SampleRate = 1./dt; %Hz
    ImagingDuration = max(time); %Sec
    
     %% Phantom distortion
 if isbreathing
    SpikeMeanHeight = 0.15; % Arbitrary unit.
    SpikeVarHeight = 0.002;
 else
    SpikeMeanHeight = 0.0001; % Very small displacement values, when interpolated, will not cause an appreciable amount of variation in the edge profile.
    SpikeVarHeight = 0.0000000001;
 end
    % TODO: Add random wavelength selection.
    [Ny,Nx,N] = deal(imageSize);
    WLorder = 1:numel(WLs);
    
    if boolean(oversample)
       WLorder = repelem(WLorder,oversample);  %#ok<UNRCH>
       WLorder = WLorder(:);
    end
 
 
    
 %% Ellipses/phantom shape and composition over time.
    % For N ellipses and T timepoints and C components, need an N x (C x T)
    % array to describe it. 
    % So for 10 ellipses, 4 components, and 250 time points, that's 10000
    % entries.
 switch phantomType
    case 'Shepp-Logan'
    case {'Modified Shepp-Logan','Positive Shepp-Logan'}
        
        % SO2 values at each 'plateau' of the SO2 timecourse.
        switch dynamics
            case 'local' % Rather arbitrary values of change between each segment.
        SO2Values = [0.4 0.6 0.4; 
                     1.0 0.5 0.0  
                     0.2 0.5 0.1
                     0.5 0.5 0.6
                     0.2 0.9 0.2
                     0.1 0.1 0.9
                     0.9 0.8 0.7
                     1.0 1.0 1.0
                     0.7 0.6 0.2
                     0.9 0.2 0.9];

        % Case for no dynamics.
            case 'no'
        SO2Values = [0.4 0.4 0.4; % Case of no variation or dynamics of the underlying signal.
                     0.5 0.5 0.5 
                     0.1 0.1 0.1
                     0.5 0.5 0.5
                     0.9 0.9 0.9
                     0.2 0.2 0.2
                     0.7 0.7 0.7
                     1.0 1.0 1.0
                     0.6 0.6 0.6
                     0.3 0.3 0.3];
        % Case for global dynamics.
            case 'global'
        SO2Values = [0.3 0.7 0.4; 
                     0.3 0.7 0.4; 
                     0.3 0.7 0.4; 
                     0.3 0.7 0.4; 
                     0.3 0.7 0.4; 
                     0.3 0.7 0.4; 
                     0.3 0.7 0.4; 
                     0.3 0.7 0.4; 
                     0.3 0.7 0.4; 
                     0.3 0.7 0.4];
            otherwise
                error();
        end
                 
       % SO2 switching times, relative to the overall challenge as [0,1]
       % TODO: Randomize.
       % TODO: What happens if they're 0 or 1 or out of order?

       SO2Times =   [1/3 2/3;
                     1/3 2/3;
                     1/3 2/3;
                     1/3 2/3;
                     1/3 2/3;
                     1/3 2/3;
                     1/3 2/3;
                     1/3 2/3;
                     1/3 2/3;
                     1/3 2/3];
       
       % Time constants of response. Constant for each component.
       TimeConstants = [0.2
                        0.3
                        0.02
                        0.04
                        0.2
                        0.1
                        0.05
                        0.4
                        0.2
                        0.1];
        % TODO: Base time constants off of the number of components.
%         TimeConstants=2*rand(10,1);
                   
 end



%% Assign into parameters.
sampleData.parameters.sampling.Wavelengths = WLs;
sampleData.parameters.sampling.TimeStep = dt;
sampleData.parameters.sampling.Fs = SampleRate;

sampleData.parameters.image.Ny = Ny;
sampleData.parameters.image.Nx = Nx;

sampleData.parameters.breathing.SpikeHeightMean = SpikeMeanHeight;
sampleData.parameters.breathing.SpikeHeightVar = SpikeVarHeight;
sampleData.parameters.breathing.TimeJitter = jitter;
sampleData.parameters.breathing.BreathPeriod = BreathPeriod;


sampleData.parameters.SO2.SO2Values = SO2Values;
sampleData.parameters.SO2.SO2Times  = SO2Times;
sampleData.parameters.SO2.TimeConstants = TimeConstants;

sampleData.timeCourses.Time = time;

%% Generate ellipses and the overall image.
    [phantomImage,ellipses,phantomLayers]=phantomEllipses(Ny,...
            phantomType);
    
sampleData.parameters.image.PhantomImage = phantomImage;
sampleData.parameters.image.Ellipses = ellipses;
sampleData.parameters.image.PhantomLayers = phantomLayers;


%% Create the timecourse for each component.
nEllipses = size(phantomLayers,3);
SO2Times = [SO2Times ones(nEllipses,1)]; % Closes SO2 times on the right.
timeCourses = zeros(nEllipses,numel(time)); % Storage for SO2 timecourses.


% Allocation.
    tcIm = zeros(N,N,numel(time));
    spectralImageTimeCourse_distorted = zeros(N,N,numel(WLs),numel(time));
    tcImStore = cell(nEllipses,1); 
    spectralCourse = zeros(numel(WLs),numel(time),nEllipses);
    
% Create smoothed timecourses.
for k=1:nEllipses
       % Allocation.
       oxygenTC = zeros(size(time));
       
       % Assign the values of the oxygenation between the time switches.
       % TODO: Clarify logic here.
       for switchInd=size(SO2Times,2):-1:1
           oxygenTC(time<=(SO2Times(k,switchInd).*max(time)))=SO2Values(k,switchInd);
       end
       
       % Generate this timecourse's convolution kernel. 
       kern=exp(-TimeConstants(k)*time);
       kern=kern./sum(kern); % normalize
       kern=[zeros(size(kern)) kern]; % Pad up front. Necessary?
       kern = smoothdata(kern);
       oxygenTC=convWFrontPad(oxygenTC,kern);
       
       % Store the TC and its spectrum in the library.              
       timeCourses(k,:)=oxygenTC;
       spectralCourse(:,:,k)=makeSpectralSet(oxygenTC,WLs);
end

sampleData.timeCourses.SO2TimeCourses = timeCourses;
sampleData.timeCourses.SpectralTimeCourses = spectralCourse;

%% Prepare for simulating movement.
[distMap,zz] = makeDistortion(N);
 interpolator=makeInterpolator(zz,phantomLayers); % TODO: check that the interpolator works for both stacks and not-stacks.
 nEllipses=size(ellipses,1);

sampleData.parameters.distortion.DistortionMap = distMap;
sampleData.parameters.distortion.DistortionCoordinates = zz;

% Make grid vectors and interpolating grids.
xVec = real(zz(1,:));
yVec = imag(zz(:,1));
zVec = permute(1:nEllipses,[1 3 2]); % Puts the non-singleton dimension in dim 3.

[X,Y,Z] = ndgrid(xVec,yVec,zVec);

Zmix = X+1i*Y;


%% Make breathing timecourse

% TODO: Make this simpler and more debuggable.
q2 = makeBarTimeCourse(BreathPeriod,jitter,time,SampleRate,SpikeVarHeight,SpikeMeanHeight);
q = imgaussfilt(q2);

sampleData.timeCourses.BreathingTimeCourse = q;

trueSpectralImage_distorted = zeros(N,N,nEllipses,numel(WLs),numel(time));
trueSpectralImage = zeros(N,N,nEllipses,numel(WLs),numel(time));

    
%% Distort images and combine into spectral image. Also record the undistorted version.

    for t=1:numel(time)
        % grab the breathing timecourse and apply the distortion.
        dx=q(t);
        Zdistort=Zmix+dx*1i*distMap; % TODO: Check the math here.
        distortedIm=interpolator(real(Zdistort),imag(Zdistort),Z);
        
        % grab the spectral timecourse for each component.
        expander = squeeze(spectralCourse(:,t,:));
        expander = permute(expander,[4 3 2 1]);
        
        % Take tensor product of vectors, create representative image.
        totalIm_distorted = distortedIm.*expander;
        totalIm_true = phantomLayers.*expander;
        
        trueSpectralImage_distorted(:,:,:,:,t) = totalIm_distorted;
        trueSpectralImage(:,:,:,:,t) = totalIm_true;
        
        spectralImage_distorted =squeeze(sum(totalIm_distorted,3));
        spectralImageTimeCourse_distorted(:,:,:,t)=spectralImage_distorted;
        
        
    end
    
%% Save component time-course image.
sampleData.images.CompleteData_distorted = trueSpectralImage_distorted;
sampleData.images.CompleteData = trueSpectralImage;


%% Simulate the process of acquiring data from the object distribution.

% Allocation.
sampledImageTimeCourse=zeros(N,N,numel(time));

% Set the wavelength sampling index.
WLindex=1;
WLlog=zeros(size(time)); % Record which wavelength was used when.
    

%% Loop over the whole time series and assign the sampled image.
for k=1:numel(time)
    
    
    if WLindex>numel(WLorder) % If we've already cycled through the indices...
        WLindex=1;        % Bring back to 1.
    end
    
    % Assign the image to the data store.
    % TODO: Simulate noise (acoustic and electrical). 
    % TODO: With noise, get SNR information. 
    sampledImageTimeCourse(:,:,k)=spectralImageTimeCourse_distorted(:,:,WLorder(WLindex),k);

    % Store the wavelength we acquired with. 
    WLlog(k)=WLs(WLorder(WLindex));
    WLindex=WLindex+1;
end

sampleData.parameters.sampling.SamplingOrder = WLorder;
sampleData.parameters.sampling.SamplingWavelengths = WLs(WLorder);

sampleData.images.SampledData = sampledImageTimeCourse;
sampleData.timeCourses.Wavelengths = WLlog;

%% Save the object distribution, as well as the sampled distribution.

CD_distorted = squeeze(sum(sum(sampleData.images.CompleteData_distorted,3),4));
CD = squeeze(sum(sum(sampleData.images.CompleteData,3),4));


% make the local directory (just use a temp folder)
% tdir_local=tempname(fullfile(pwd,'genData'));
tdir_local = fullfile(saveFolder,datasetName);
mkdir(tdir_local);
cd(tdir_local);

% make MAT files for storage
% TODO: Consolidate into a structure.
save('testData.mat','sampleData','-v7.3');

%% Create a video of the sampled dataset.

% Quick fun thing: Can alias a deep handle reference by making an anonymous
% function.
samp = @() sampleData.images.SampledData; % Neat!

v = VideoWriter([datasetName '.avi']);
v.FrameRate = SampleRate;
open(v);

imToWrite = normalizeData(permute(samp(),[1 2 4 3]));

writeVideo(v,imToWrite);

close(v);


%% Go home.
% cd('..\..');


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%    HELPER FUNCTIONS    %%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%% Normalize data to [0,1]
function dat = normalizeData(input)
   
    minX = min(input(:));
    maxX = max(input(:));
    
    dat = (input-minX)./(maxX-minX);
    
end


%% Overlay perlin noise on the image provided.
function im = perlin_noise(im)

    [n, m] = size(im);
    i = 0;
    w = sqrt(n*m);

    while w > 3
        i = i + 1;
        d = interp2(randn(ceil((n-1)/(2^(i-1))+1),ceil((m-1)/(2^(i-1))+1))*4, i-1, 'spline');
        im = im + i * d(1:n, 1:m);
        w = w - ceil(w/2 - 1);
    end
end

%% Makes the spectral timeseries corresponding to the given SO2 timeseries.
function oxyCourse = makeSpectralSet(SO2Course,WLs)

    load HemoglobinSpectra.mat
    
    Hb_lambda = Hb.spectra(:,1); % Get wavelengths used for Hb spectra (every 2nm)
    ox_abs = Hb.spectra(:,2); % Get extinction vals of HbO spectra (every 2nm)
    dox_abs = Hb.spectra(:,3); % Get extinction vals of HbD spectra (every 2nm)

    % Resample all of the spectra to give us just the wavelengths we're
    % interested in.
    ox_spec = interp1(Hb_lambda,ox_abs,WLs);
    dox_spec = interp1(Hb_lambda,dox_abs,WLs);

    specMat = [ox_spec(:),dox_spec(:)]; % Concatenate into a mixing matrix
    
    HbOConc = SO2Course;
    HbConc = 1 - SO2Course;
    
    concMat = [HbOConc(:).';HbConc(:).'];
    oxyCourse = specMat*concMat;

end



%% Returns the time-stabilized timecourse of oxygenation.
function output = convWFrontPad(target,kern)
    
    
    N = numel(target);
    
    % Expand out the timecourse by padding the front.
    target2 = [ones(size(target))*target(1), target,ones(size(target))*target(end)];
    
    tempconv = conv(target2,kern,'same');
    
    output=tempconv((N+1):(end-N));
end


%% Creates a distortion map of the given image.
%  Based on an elliptical Mobius transform.
%  TODO: Stabilize a little better, figure out the conjugacy/transpose issue.
function [distortionMap,zz] = makeDistortion(Npix)

    [xx,yy]=meshgrid(linspace(-2,2,Npix));
    zz=xx+1i*yy;

    Mmat=[-exp(1i*pi/4)     0.1i;
              0       exp(1i*pi/4)]

    distortionMap = mobiusTransform(zz,Mmat)+2i-real(zz);
    distortionMap = -distortionMap./abs(distortionMap);

end

%% Makes a GriddedInterpolant object to allow resampling of image coordinates.
function interpolator = makeInterpolator(zz,im)
   isStack = (ndims(im)==3);
   if isStack
    stackSize = size(im,3);
    
    zVec = permute([1:stackSize],[1 3 2]);
    
    gridVecs = {real(zz(1,:)),imag(zz(:,1)),zVec};
    
    interpolator = griddedInterpolant(gridVecs,im);
   else
    interpolator = griddedInterpolant(real(zz)',imag(zz)',im');
   end
end

%% Applies the Mobius transform to the given z array.
function output = mobiusTransform(z,xfrm)
   
    a=xfrm(1,1);
    b=xfrm(1,2);
    c=xfrm(2,1);
    d=xfrm(2,2);
    
    output = (a.*z+b)./(c.*z+d);
end

%% Creates a timecourse of breathing for simulating motion. 
function [tc] = makeBarTimeCourse(BreathPeriod,jitter,time,SampleRate,SpikeVarHeight,SpikeMeanHeight)
    
    % First localize the variables.
    T = BreathPeriod;
    fs = SampleRate;
    tc = zeros(size(time));
    
    % Get parameters for inter-breath time. This is index, so convert.
        meanInterTime = BreathPeriod*fs;
        if jitter == 0
            varInterTime = meanInterTime./1e6;
        else
            varInterTime = jitter*(fs.^2);
        end
        
        b_time = varInterTime./meanInterTime;
        a_time = meanInterTime./b_time;
        
    % Get parameters for breathing motion.
        meanMotion = SpikeMeanHeight;
        if SpikeVarHeight == 0
            varMotion = SpikeMeanHeight./1e6;
        else
            varMotion = SpikeVarHeight;
        end
        
        b_motion = varMotion./meanMotion;
        a_motion = meanMotion./b_motion;
    
    % Initial time point. Want to not deal weird edge artifacts or
    % negative time.
    tind = round(gamrnd(a_time,b_time));
    
    while tind  < max(time)*fs
        
        % Put a spike of random height at the timepoint given.
        tc(tind) = gamrnd(a_motion,b_motion);
        
        % advance to the next time point.
        tind = tind + round(gamrnd(a_time,b_time));
    end
end










