% correct slc images for rotational drift -- master copy!
%
% emily fedders
% april 2021, updated 08/06/2021
% updated 01/10/2023
% updated 01/11/2023 to work backwards from first checkfile to make sure
% all images that can be corrected are corrected! 
%
% Notes:
% - this is a master copy! copy this file to a new file before you make
% edits!! THEN, **once edits are working**, commit changes to this file if you
% desire.
% - This code does not correct for bulk rotations less than 5 azimuth lines
% (0.5 degrees).  Maybe this needs to change at some point, but thus far
% seems to give fairly good results (and better results than trying to
% correct for every 1 or 2 pixel rotation that ocurred).
% - This code does not allow for more than 2 rotations.  The code to
% accomodate more has been left in (commented) but in processing the hourly
% data, a maximum of 2 rotation amounts seemed to do the trick, and even 2
% rotations were only used twice in 249 images.
% 
% Note UPDATES as of 01/10/2023 =======================================
% this code DOES correct for bulk rotations down to 1 line. Whether peak
% cross-correlation occurs at a 1 or 100 line offset, that peak
% cross-correlation will be used for rotation correction. 
% 
% This code now DOES correct for two and more rotations --> however many
% peaks appear in the cross correlation with prominence greater than 0.2. I
% don't think there are any images requiring more than 2 rotations though. 
%
% -------------------------------------------------------------------------

clear;
close all;
fclose all;

addpath('/home/erfedders/matlabUtilities','-end');

%% USER-DEFINED ------------------------------------------------------------
% uncorrected slc file location
% baseDir = '/data/uaf/gpri/2022/Utq_2022/Scans/slc_rx1/'; 
baseDir = '/data/uaf/gpri/2021/SIDEx2021/SLC/';
% checkDir = '/data/project/Utq_2022/hourly_rotcorr/slc_rx1_corr/'; 
checkDir = '/data/project/SIDEx2021/hourly_rotcorr_testv2/slc_corr/'; 

% corrected slc file location
% outDir = '/data/project/Utq_2022/all_rotcorr/slc_rx1_corr';
outDir = '/data/project/SIDEx2021/all_rotcorr_testv2/SLC_corr'; 

if ~exist(outDir,'dir')
    mkdir(outDir);
end

% path to file containing list of slc files you'd like to correct
% listFile = '/data/project/SIDEx2021/hourly_rotcorr/slctab_allHourly.txt'; 

% just use first 500 pixels - things get messier farther away
cutoff = 500; 

% toggle plots
plotsOn = false;

% starting file (useful if don't want to start from the very beginning of
% some list you made, otherwise set to 1).
startN = 1;

% filePat = '*_RX1.slc'; 
filePat = '*u.slc';

%% SETUP ------------------------------------------------------------------
% create output directory
if ~exist(outDir,'dir')
    mkdir(outDir)
end

%% get list of slc files
if ~exist('listFile','var') % if you don't have a list already

    % list all files matching filePat in baseDir
    fstruct = dir([baseDir filePat]);
    len = length(fstruct);
    fList = char([]);
    
    for ii = 1:len
        fList(ii,:) = fstruct(ii).name;
    end
    
    % manipulate this list into a single row of characters for use with split
    fchar = fList'; fchar = fchar(:)';
    
    % split at user-defined file ending minus asterisk (will give rows
    % containing two datetime character strings delimited by first portion of filePat'_')
    fchar = split(fchar,filePat(2:end));
    fchar(end) = [];
    
    % split again at '[filepat without extensions]_', should give you an X by 2 cell array
    filesplit = split(filePat,'.');
    fchar = split(fchar,[filesplit{1}(2:end) '_']);
    
    % convert all to datetime format
    dlist0 = datetime(fchar(:,1), 'InputFormat','yyyyMMdd''_''HHmmss');
    
else % if you do have a list already
    fid = fopen(listFile);
    if fid==-1
        fprintf('error opening %s\n',...
            listFile);
        return;
    else
        cMat = textscan(fid,'%s','headerlines',1);
        fclose(fid); % be tidy
    end
    fList = cell2mat(cMat{1});

end

%% step through files
% number of files
nfiles = size(fList,1);
fprintf('nfiles: %d\n',nfiles);

% find first file to exist in the check directory: 
done = false; 
startN = startN-1; 
while done == false
    startN = startN + 1; 
    [p,n,e] = fileparts(fList(startN,:)); 
    checkfile = [checkDir n '_corr' e];
    done = exist(checkfile,'file'); 
end 
fprintf('starting at file %d\n',startN); 

% treat first file seperately:
outfile = [outDir filesep n '_corr' e];
if ~exist(outfile, 'file')
    %copy this first already corrected file without doing anything further
    copyfile(checkfile,outfile);
    copyfile([checkfile '.par'],[outfile '.par']);   
end

fprintf('%s\n',outfile);
fprintf('first file, no changes\n\n');

shiftVec = zeros(nfiles,1); 

% now start loop at previous file, and work backward to beginning of
% timeseries
for kk = startN-1:-1:1

    tic; % start timer
    
    fprintf('%d ..............%s\n',kk,fList(kk,:));
    
    % define reference file (previous corrected output)
    slcFile1 = outfile;
    
    % define file to be corrected
    slcFile2 = [baseDir fList(kk,:)];
    
    % check if there is already a corrected version of this file! 
    [p,n,e] = fileparts(slcFile2); 
    checkfile = [checkDir n '_corr' e];
    
    if exist(checkfile,'file')

        outfile = [outDir filesep n '_corr' e];
        copyfile(checkfile,outfile);
        
        copyfile([checkfile '.par'],[outfile '.par']);
        fprintf('%s\n',outfile);
        fprintf('check file, no changes');
        fprintf('\n');
        clear('sli2_norm'); % you want to make sure that you read in this file your next time through! 
        
        continue; % don't go through rest of loop
    else 
        fprintf('no check file\n'); %:\n%s\n',checkfile); 
    end 
    
    %% read in reference slc file: ------------------------------------------
    if exist('newMat','var')
        % use the previously corrected file as your starting file! 
        sli1_norm = newMat; 
        clear('newMat')
        
    elseif exist('sli2_norm','var')
        % if you've already read in this file, no need to do it again! 
        sli1_norm = sli2_norm;
        azStart_old = azStart; 
        
    else
        % read in reference slc, use function since you don't need to re-output
        % these values
        
        [sli1,xx,yy,~] = readInFcomplex_slc(slcFile1);  % intensity image from first slc
        sli1 = (sli1.^0.35)./max(sli1(:).^0.35);        % scale between 0 and 1 with exponent
        sli1_norm = sli1 - mean(sli1(~isnan(sli1)));    % normalize (important for cross correlation)
        
        sli1_norm = sli1_norm(:,1:cutoff);
        azStart_old = min(yy(:)); 
        
    end

    %% read in test slc file: -----------------------------------------------
    % you'll need these values again, so do it outside any function:
    parFile = [slcFile2 '.par'];
    
    % parameter file =============
    fid = fopen(parFile);
    if fid == -1
        fprintf('error opening %s\n',parFile);
        if ~exist(parFile,'file')
            continue; 
        else
            return;
        end
    end
    
    %read entire file into a cell array
    cMat = textscan(fid,'%s %s %*[^\n]','HeaderLines',2);
    fclose(fid);    % close file
    
    if any(size(cMat{1})==0)
        continue; 
    end
    
    % find width (number of azimuth lines)
    a = strcmp(cMat{1},'range_samples:');
    if all(a == 0)
        continue; 
    end
    
    wd = str2double(cMat{2}(a));
    % find number of rows (number of range samples)
    a = strcmp(cMat{1},'azimuth_lines:');
    nl = str2double(cMat{2}(a));
    
    % starting azimuth
    a = strcmp(cMat{1},'GPRI_az_start_angle:');
    azStart = str2double(cMat{2}(a));
    % azimuth step
    a = strcmp(cMat{1},'GPRI_az_angle_step:');
    azStep = str2double(cMat{2}(a));
    
    clear cMat; % be tidy
    % ============================
    
    % slc file ===================
    fid = fopen(slcFile2, 'r','b');
    if fid == -1
        fprintf('error opening %s\n',slcFile2);
        return;
    end
    
    % read in fcomplex data
    slcIn = fread(fid, [wd*2*nl,1], 'float');
    
    fclose(fid); % be tidy
    
    % reshape it into matrix form
    % every other pixel, starting at 1, is real
    realpx = 1:2:wd*2*nl;
    % every other pixel, starting at 2, is imaginary
    imagpx = 2:2:wd*2*nl;
    % reshape into image dimensions
    if numel(slcIn) == (numel(realpx)+numel(imagpx))
        % get real pixels in image dimensions
        slc_real = rot90(reshape(slcIn(realpx),wd,nl));
        % get imaginary pixels in image dimensions
        slc_imag = rot90(reshape(slcIn(imagpx),wd,nl));
    else
        fprintf('file %s is corrupted, continuing\n',slcFile2);
        fprintf('\n');
        continue;
    end
    
    sli2 = sqrt(slc_imag.^2 + slc_real.^2);      % test image intensity
    sli2 = (sli2.^0.35)./max(sli2(:).^0.35);     % scale and normalize in the same
    sli2_norm = sli2 - mean(sli2(~isnan(sli2))); % way as reference image
    sli2_norm = sli2_norm(:,1:cutoff);           % make the same size as the first image
    
    if size(sli2_norm,1) < size(sli1_norm,1)
        remrows = size(sli1_norm,1)-size(sli2_norm,1)-1; 
        sli1_norm(end-remrows:end,:) = []; 
    end 
    
    %number of rows and columns
    [nrows, ncols] = size(sli2_norm);
    
    sli1_norm(isnan(sli1_norm)) = 0;
    sli2_norm(isnan(sli2_norm)) = 0;
    
    % calculate normalized correlation and associated lags at several locations
    [c2,lags] = xcorr(sli1_norm(:,200),sli2_norm(:,200),'normalized',300); %'normalized'); 
    [c3] = xcorr(sli1_norm(:,300),sli2_norm(:,300),300,'normalized'); %normalized'); 
    [c4] = xcorr(sli1_norm(:,400),sli2_norm(:,400),300,'normalized'); %normalized'); 
    
    c = mean([c2,c3,c4],2); 
    
    a = c == max(c);  
    shiftVec(kk) = lags(a); 
    
    [pks,locs] = findpeaks(c,lags,'SortStr','descend','MinPeakProminence',0.2);
    
    if isempty(locs)
        % copy current file to "corrected" file without doing anything
        [p,n,e] = fileparts(slcFile2);
        outfile = [outDir filesep n '_corr' e];
        
        fprintf('%s\n',outfile);
        fprintf('no changes: \n');
        fprintf('no peaks in correlation \n');
        
        copyfile(slcFile2,outfile);
        copyfile([slcFile2 '.par'],[outfile '.par']);
        
        % append a change log to the end of the parameter file
        fid = fopen([outfile '.par'],'a');
        if fid == -1
            fprintf('Error opening %s\n',[outfile '.par']);
            return;
        end
        fprintf(fid,'Rotation Correction Applied: NO \n');
        fclose(fid);
        
        toc;
        fprintf('\n');
        
        continue;
        
    elseif length(locs) == 1
        
        if abs(locs) == 0 % <= 1 % == 0 || locs == 1
            % copy current file to "corrected" file without doing anything
            [p,n,e] = fileparts(slcFile2);
            outfile = [outDir filesep n '_corr' e];
            
            fprintf('%s\n',outfile);
            fprintf('no changes: \n');
            fprintf('peak correlation at zero lag \n');
            
            copyfile(slcFile2,outfile);
            copyfile([slcFile2 '.par'],[outfile '.par']);
            
            % append a change log to the end of the parameter file
            fid = fopen([outfile '.par'],'a');
            if fid == -1
                fprintf('Error opening %s\n',[outfile '.par']);
                return;
            end
            fprintf(fid,'Rotation Correction Applied: NO \n');
            fclose(fid);
            
            toc;
            fprintf('\n');
            
            continue;
            
        else
            % shift all rows by appropriate lag
            
            rowsOrig = (1:nrows)'; 
            rowsMoved = rowsOrig + locs;
            newMat = zeros(size(sli2_norm)); %zeros(nl,wd); 
            
            if locs < 1
                rowsOrig(rowsMoved < 1) = []; 
                rowsMoved(rowsMoved < 1) = []; 
            
                newMat(rowsMoved,:) = sli2_norm(rowsOrig,:);
            else 
                rowsOrig(rowsMoved > nl) = []; 
                rowsMoved(rowsMoved > nl) = []; 
                
                newMat(rowsMoved,:) = sli2_norm(rowsOrig,:); 
            end 
            
            if ~all(size(newMat) == size(sli2_norm))
                fprintf('output matrix isn''t same size as input matrix, figure it out \n')
                return; 
            end 
            
             % calculate normalized correlation and associated lags at several locations
            [ctest,ltest] = xcorr(sli1_norm(:,200),newMat(:,200),'normalized',300); %'normalized');

            a = find(ctest == max(ctest)); 
            if abs(ltest(a)) > 1
                fprintf('max correlation > 1 lag even after correction, investigate\n'); 
                return; 
            end 
            
             %% now, how do I put this information back into the fcomplex format?
             % create a placeholder matrix of NaN values
             new_slc_real = zeros(size(slc_real)); 
             new_slc_imag = zeros(size(slc_imag)); 
             
             new_slc_real(rowsMoved,:) = slc_real(rowsOrig,:); 
             new_slc_imag(rowsMoved,:) = slc_imag(rowsOrig,:); 

             
             % rotate back to original orientation
             new_slc_real = rot90(new_slc_real,3);
             new_slc_imag = rot90(new_slc_imag,3);
             
             % unwrap and interleave images
             new_slc_combi = 1:2*numel(new_slc_real);
             count = 0;
             for ii = 1:numel(new_slc_real)
                 count = count + 1;
                 new_slc_combi(count) = new_slc_real(ii);
                 count = count + 1;
                 new_slc_combi(count) = new_slc_imag(ii);
             end
             
             % 5) write real and imaginary values to file
             [p,n,e] = fileparts(slcFile2);
             
             outfile = [outDir filesep n '_corr' e];
             copyfile(parFile,[outfile '.par']);
             fid = fopen(outfile,'w','b');
             if fid==-1
                 fprintf('ERROR opening %s\n',outfile);
                 return;
             end
             fwrite(fid,new_slc_combi,'float32','b');
             
             fprintf('%s\n',outfile);
             fprintf('changes applied: offset = %d\n',locs);
             fprintf('\n');
             
             fclose(fid);
             
             % append a change log to the end of the parameter file
             fid = fopen([outfile '.par'],'a');
             if fid == -1
                 fprintf('Error opening %s\n',[outfile '.par']);
                 return;
             end
             fprintf(fid,'Rotation Correction Applied: %s\n',locs);
             fclose(fid);

        end 
    else 
        
        % use line by line cross correlation to figure it out! 
        rvec = zeros(nrows, 1); % pre-allocate vector to hold correlation values
        offset = zeros(nrows, 1); % pre-allocate vector to hold offset quantities
        
        findoff = false; 
        
        for ii = 1:nrows
            % cross-correlate row ii of test image with row ii of reference
            r = xcorr(sli1_norm(ii,200:400),sli2_norm(ii,200:400),'normalized',0); 
            rvec(ii) = r; 
            
            if r < 0.6 || isnan(r) % arbitrary cutoff
                rmoved = ones(numel(locs),1);  
                
                % test each lag in locs 
                for jj = 1:length(locs)
                    if ((ii + locs(jj)) <= nrows) && (ii+locs(jj) >=1)
                        r = xcorr(sli1_norm(ii+locs(jj),200:400),sli2_norm(ii,200:400),'normalized',0); 
                        rmoved(jj) = r; 
                    end 
                end 
                
                if all(isnan(rmoved))
                    offset(ii) = 999; 
                    findoff = true; 
                else 
                
                    % find offset giving max correlation
                    a = rmoved == max(rmoved); 
                    
                    if (sum(a) > 1) || ~any(a)
                        offset(ii) = 999;
                        findoff = true;
                    else
                        % save this offset
                        offset(ii) = locs(a);
                        
                        % assign unknown offsets this value:
                        if findoff
                            fil = offset == 999;
                            offset(fil) = locs(a);
                            % reset findoff toggle
                            findoff = false;
                        end
                    end
                    
                end 
                
            end
        end
        % calculate moving median of offset vector, fill to ends of vector
        % such that "offset" remains same size. 
        offset = movmedian(offset,31); 
                
        rowsOrig = (1:nrows)';
        rowsMoved = rowsOrig + offset;
        newMat = zeros(size(sli2_norm)); %zeros(nl,wd);
        
        rowsOrig(rowsMoved < 1) = [];
        rowsMoved(rowsMoved < 1) = [];
        
        rowsOrig(rowsMoved > nl) = [];
        rowsMoved(rowsMoved > nl) = [];
        
        newMat(rowsMoved,:) = sli2_norm(rowsOrig,:);

        if ~all(size(newMat) == size(sli2_norm))
            fprintf('output matrix isn''t same size as input matrix, figure it out \n')
            return;
        end
        
        % calculate normalized correlation and associated lags at several locations
        [ctest,ltest] = xcorr(sli1_norm(:,200),newMat(:,200),'normalized',300); %'normalized');
        
        a = find(ctest == max(ctest));
        if abs(ltest(a)) > 1 || length(ltest(a)) > 1
            fprintf('max correlation not < 1 even after correction, investigate\n');
            return;
        end
        
        %% now, how do I put this information back into the fcomplex format?
        % create a placeholder matrix of NaN values
        new_slc_real = zeros(size(slc_real));
        new_slc_imag = zeros(size(slc_imag));
        
        new_slc_real(rowsMoved,:) = slc_real(rowsOrig,:);
        new_slc_imag(rowsMoved,:) = slc_imag(rowsOrig,:);
        
        
        % rotate back to original orientation
        new_slc_real = rot90(new_slc_real,3);
        new_slc_imag = rot90(new_slc_imag,3);
        
        % unwrap and interleave images
        new_slc_combi = 1:2*numel(new_slc_real);
        count = 0;
        for ii = 1:numel(new_slc_real)
            count = count + 1;
            new_slc_combi(count) = new_slc_real(ii);
            count = count + 1;
            new_slc_combi(count) = new_slc_imag(ii);
        end
        
        % 5) write real and imaginary values to file
        [p,n,e] = fileparts(slcFile2);
        
        outfile = [outDir filesep n '_corr' e];
        copyfile(parFile,[outfile '.par']);
        fid = fopen(outfile,'w','b');
        if fid==-1
            fprintf('ERROR opening %s\n',outfile);
            return;
        end
        fwrite(fid,new_slc_combi,'float32','b');
        
        fprintf('%s\n',outfile);
        fprintf('changes applied: offset = %d\n',locs);
        fprintf('\n');
        
        fclose(fid);
        
        % append a change log to the end of the parameter file
        fid = fopen([outfile '.par'],'a');
        if fid == -1
            fprintf('Error opening %s\n',[outfile '.par']);
            return;
        end
        fprintf(fid,'Rotation Correction Applied: %s\n',locs);
        fclose(fid);
        
    end 
    toc; 

end

% %% Example files: -------------------------------------------------------
% % % good correlation:
% % slcFile1 = 'E:\SIDEx2021\SLC\20210312_234501u.slc';
% % slcFile2 = 'E:\SIDEx2021\SLC\20210312_233002u.slc';
%
% % test series from 03/15/2021 -----------------------------------------
% % good reference file at the beginning of a series of offset images
% % slcFile1 = 'E:\SIDEx2021\SLC\20210315_232401u.slc'; %reference
%
% % step function in correlation (wedge missing from igram):
% % slcFile1 = 'E:\SIDEx2021\SLC\20210315_232802u_corr.slc'; % test case 1
%
% % % complete offset:
% %slcFile1 = 'E:\SIDEx2021\SLC\20210315_233201u_corr.slc'; % test case 2
% %
% % %should be back to original orientation?:
% slcFile1 = 'E:\SIDEx2021\SLC\20210315_234801u_corr.slc'; % test case 3
% %
% slcFile2 = 'E:\SIDEx2021\SLC\20210315_235201u.slc'; % test case 4