% correct slc images for rotational drift -- master copy!
%
% emily fedders
% april 2021, updated 08/06/2021
%
% 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.
%
% -------------------------------------------------------------------------

clear;
close all;
fclose all;

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

%% USER-DEFINED ------------------------------------------------------------
% uncorrected slc file location
% baseDir = '/data/project/Utq_2022/hourly/slc_rx1/'; %'D:IDEx2021/SLC/';
baseDir = '/data/uaf/gpri/2021/SIDEx2021/SLC/'; % check this? ?

% corrected slc file location
% outDir = '/data/project/Utq_2022/hourly_rotcorr/slc_rx1_corr'; %'E:/SIDEx2021/event3/SLC_corr';
outDir = '/data/project/SIDEx2021/hourly_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'; %'E:/SIDEx2021/event3/slctab_20210319T003000_20210320T070000_event3.txt';

% just use first 500 pixels
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 = '*l.slc';

%% SETUP ------------------------------------------------------------------
% create output directory
if ~exist(outDir,'dir') == 1
    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');
%     dlist1 = datetime(fchar(:,2), '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});
    
    if ~strcmp('*u.slc',filePat)
        fList = char(strrep(string(fList),'u.slc',filePat(2:end))); 
    end 
    
end

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

% treat first file seperately:
startFile = [baseDir fList(startN,:)];

[p,n,e] = fileparts(startFile);
outfile = [outDir filesep n '_corr' e];

if ~exist(outfile, 'file')
    %copy the first file to a "corrected" file without doing anything"
    copyfile(startFile,outfile);
    copyfile([startFile '.par'],[outDir filesep n '_corr' e '.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);
end

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

shiftVec = zeros(nfiles,1); 

% now start loop at second file
for kk = startN+1:nfiles

    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,:)];
    
    %% read in reference slc file: ------------------------------------------
    if exist('newMat','var')
        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
        % use function to read in reference slc, 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
    
    % find width (number of azimuth lines)
    a = strcmp(cMat{1},'range_samples:');
    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 
    
%     if azStart ~= azStart_old
%         fprintf('starting angle is not the same, figure it out\n');
%         
%         %         if azStart > azStart_old
%         %             azAdd = azStart_old:azStep:azStart;
%         %             sli2_norm = [zeros(length(azAdd),size(sli2_norm,2)); sli2_norm];
%         %             if size(sli2_norm,1) > size(sli1_norm,1)
%         %                 sli2_norm = sli2_norm(1:size(sli1_norm,1),:);
%         %             end
%         %             azStart = azAdd(1);
%         %
%         %         else
%         %             azAdd = azStart:azStep:azStart_old;
%         %             sli2_norm = [sli2_norm; zeros(length(azAdd),size(sli2_norm,2))];
%         %             if size(sli2_norm,1) > size(sli1_norm,1)
%         %                 s = size(sli2_norm,1)-size(sli1_norm,1);
%         %
%         %                 sli2_norm = sli2_norm(s-1:end,:);
%         %             end
%         %
%         %         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
    nanfil= ~isnan(sum(sli2_norm,2)) & ~isnan(sum(sli1_norm,2)); %(:,200))&~isnan(sli1_norm(:,300))&~isnan(sli1_norm(:,400) & ); 

    [c2,lags] = xcorr(sli1_norm(nanfil,200),sli2_norm(nanfil,200),'normalized',300); %'normalized'); 
    [c3] = xcorr(sli1_norm(nanfil,300),sli2_norm(nanfil,300),300,'normalized'); %normalized'); 
    [c4] = xcorr(sli1_norm(nanfil,400),sli2_norm(nanfil,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 plotsOn == true
%         figure(1); clf;
%         
%         subplot(2,1,1)
%         plot(lags,c2);
%         hold on
%         plot(lags,c3);
%         plot(lags,c4);
%         plot(lags,c,'color','black','linewidth',2);
%         drawnow;
%         subplot(2,1,2);
%         
%         plot(lags,c);
%         
%         hold on
%         plot(locs,pks,'p')
%         for ii = 1:length(locs)
%             text(locs(ii),pks(ii),['lag = ' num2str(locs(ii))]);
%         end
%         
%         title(['file num ' num2str(kk)])
%     end
    
    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 
%         fprintf('more than one lag - figure it out\n')
%         return; 
        
        % 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
                    
%                     if (sum(a) > 1) || ~any(a)
%                         offset(ii) = 999;
%                         findoff = true;
%                     % 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
        % calculate moving median of offset vector, fill to ends of vector
        % such that "offset" remains same size. 
        offset = movmedian(offset,31); 
        
        if plotsOn == true
            figure(1); clf;
            
            subplot(2,1,1)
            plot(lags,c2);
            hold on
            plot(lags,c3);
            plot(lags,c4);
            plot(lags,c,'color','black','linewidth',2);
            drawnow;
            subplot(2,1,2);
            
            plot(lags,c);
            
            hold on
            plot(locs,pks,'p')
            for ii = 1:length(locs)
                text(locs(ii),pks(ii),['lag = ' num2str(locs(ii))]);
            end
            
            title(['file num ' num2str(kk)])
            % progress plots -------------------------------------------------------
            figure(2); clf; % plot unshifted correlation values
            plot(rvec);
            xlabel('azimuth lines')
            ylabel('correlation with reference')
            drawnow;
            
            figure(3); clf; % plot offsets and smoothed offsets
            plot(offset)
            hold on
            plot(off_smooth)
            legend('offset','offset smoothed');
            xlabel('azimuth lines')
            ylabel('offset to nearest line of good coherence')
            drawnow;
        end
        
        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; 


%     if plotsOn == true
%         %% read in the file that you just made!!
%         sliCorr = readInFcomplex_slc(outfile);
%         sliCorr = (sliCorr.^0.35)./max(sliCorr(:).^0.35);    % scale between 0 and 1 and apply exponent
%         sliCorr_norm = sliCorr - nanmean(sliCorr(:));           % normalize (important for cross correlation)
%         
%         % plot to see how you did
%         figure(3); clf;
%         plot(rvec);
%         xlabel('azimuth lines')
%         ylabel('correlation with reference')
%         
%         %plot
%         figure(4); clf;
%         subplot(4,1,1);
%         imagesc('cdata',sli1_norm); colorbar; set(gca,'ydir','normal');title('intensity')
%         axis('tight')
%         
%         subplot(4,1,2);
%         imagesc('cdata',nsli2_new); colorbar; set(gca,'ydir','normal');title('intensity')
%         axis('tight')
%         
%         subplot(4,1,3);
%         imagesc('cdata',sliCorr_norm,...
%             'alphadata',~isnan(sliCorr_norm));
%         colorbar; set(gca,'ydir','normal'); title('corrected intensity')
%         axis('tight');
%     end
%     toc;
%     
%     
%     fprintf('\n');

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
