%% run06b_MCC_synthetic_with_cloud_student.m
% ============================================================
% Ocean Remote Sensing Practical
% Step 6b: MCC/PIV using synthetic tracer images with clouds
%
% STUDENT FILL-IN VERSION. Fill the ________ parts.
% This version is aligned with run06c/d:
%   - no global high-pass filtering
%   - robust normalization before matching
%   - local-mean normalized cross correlation inside each PIV window
%   - valid-pixel fraction for clouds/NaNs
%   - independent peak-gap QC
%   - neighbor-consistency QC
%
% Purpose:
%   1. Make two synthetic tracer images.
%   2. Add artificial clouds as NaN.
%   3. Estimate pattern displacement using MCC/PIV.
%   4. Show bad vectors before QC and final vectors after QC.
% ============================================================

% -------------------------------------------------------------------------
% STUDENT VERSION:
%   Complete the blanks marked by ________.
%   Focus on: cloudy-image normalization, cloud-origin/target QC, and MCC QC.
% -------------------------------------------------------------------------

clear; close all; clc;

RUN_ID = 'run06b';

%% ============================================================
% 1. Make synthetic tracer image at time 1
% ============================================================
nx = 100;
ny = 80;
[x, y] = meshgrid(1:nx, 1:ny);

rng(1);
field1 = zeros(ny, nx);

nBlob = 35;
for n = 1:nBlob
    x0 = 10 + 80 * rand;
    y0 = 10 + 60 * rand;
    amp = 0.3 + 0.9 * rand;
    sx  = 2.5 + 6.0 * rand;
    sy  = 2.5 + 6.0 * rand;
    field1 = field1 + amp * exp(-((x - x0).^2/(2*sx^2) + ...
                                  (y - y0).^2/(2*sy^2)));
end

field1 = field1 + 0.25 * sin(2*pi*(x + 0.5*y)/28);
field1 = field1 + 0.18 * cos(2*pi*(x - 0.8*y)/35);
field1 = field1 + 0.004*x + 0.002*y;

K = ones(3,3) / 9;
field1 = conv2(field1, K, 'same');
field1 = field1 - min(field1(:));
field1 = field1 / max(field1(:));

%% ============================================================
% 2. Define known displacement field
% ============================================================
dx_bg = 2.0;
dy_bg = 0.5;

xc = 52;
yc = 42;
R  = 22;
A  = 4.0;

r2 = (x - xc).^2 + (y - yc).^2;
eddy = exp(-r2 / (2*R^2));

dx_true = dx_bg - A * (y - yc) / R .* eddy;
dy_true = dy_bg + A * (x - xc) / R .* eddy;

dx_true = dx_true + 0.015 * (y - ny/2);
dy_true = dy_true + 0.010 * (x - nx/2);

%% ============================================================
% 3. Create synthetic tracer image at time 2
% ============================================================
field2 = interp2(x, y, field1, x - dx_true, y - dy_true, 'linear', NaN);

rng(2);
field2 = field2 + 0.015 * randn(size(field2));
field2(field2 < 0) = 0;
field2(field2 > 1) = 1;

%% ============================================================
% 4. Add artificial clouds
% ============================================================
rng(10);
cloudNoise1 = randn(ny, nx);
cloudNoise2 = randn(ny, nx);
Kcloud = ones(13,13) / (13*13);
cloudSmooth1 = conv2(cloudNoise1, Kcloud, 'same');
cloudSmooth2 = conv2(cloudNoise2, Kcloud, 'same');

cloudMask1 = cloudSmooth1 > 0.10;
cloudMask2 = cloudSmooth2 > 0.10;
cloudMask1 = cloudMask1 | (x > 13 & x < 34 & y > 48 & y < 68);
cloudMask1 = cloudMask1 | (x > 61 & x < 77 & y > 13 & y < 25);
cloudMask2 = cloudMask2 | (x > 33 & x < 53 & y > 18 & y < 33);
cloudMask2 = cloudMask2 | (x > 72 & x < 90 & y > 52 & y < 68);

field1_cloud = field1;
field2_cloud = field2;
field1_cloud(cloudMask1) = NaN;
field2_cloud(cloudMask2) = NaN;

%% ============================================================
% 5. Plot clean and cloudy images
% ============================================================
fig = figure('Color','w','Position',[80 80 1500 900]);

subplot(2,2,1); imagesc(field1); set(gca,'YDir','normal'); axis image; colorbar; caxis([0 1]); title('Image 1: clean'); xlabel('x grid'); ylabel('y grid');
subplot(2,2,2); imagesc(field2); set(gca,'YDir','normal'); axis image; colorbar; caxis([0 1]); title('Image 2: clean'); xlabel('x grid'); ylabel('y grid');
subplot(2,2,3); imagesc(field1_cloud); set(gca,'YDir','normal'); axis image; colorbar; caxis([0 1]); title('Image 1: with artificial clouds'); xlabel('x grid'); ylabel('y grid');
subplot(2,2,4); imagesc(field2_cloud); set(gca,'YDir','normal'); axis image; colorbar; caxis([0 1]); title('Image 2: with artificial clouds'); xlabel('x grid'); ylabel('y grid');
sgtitle('run06b: synthetic tracer images with artificial clouds');
save_png(fig, 'Fig_run06b_01_clean_and_cloudy_images.png');

%% ============================================================
% 6. MCC/PIV settings
% ============================================================
win = 17;
step = 4;
maxShift = 7;

minValidFraction = 0.50;
minCorrelation   = 0.55;
minPeakGap       = 0.030;
minTextureStd    = 0.025;
peakExcludeRadius = 1;
maxNeighborDiff  = 2.5;

% Cloud QC:
% A vector is rejected if its origin or estimated destination is on a cloud.
% cloudBufferPix also rejects vectors too close to cloud edges.
cloudBufferPix   = 2;

dt_days = 1;
dt_sec  = dt_days * 24 * 3600;
dx_km = 5;
dy_km = 5;

fprintf('=== run06b synthetic MCC/PIV with clouds ===\n');
fprintf('WINDOW=%d, STEP=%d, MAX_SHIFT=%d\n', win, step, maxShift);
fprintf('Local correlation: window mean removed, no global high-pass.\n');

%% ============================================================
% 7. MCC/PIV calculation using valid pixels only
% ============================================================
I1 = robust_normalize(________, 3.0);
I2 = robust_normalize(________, 3.0);

P = mcc_piv_core(________, ________, win, step, maxShift, minValidFraction, peakExcludeRadius);

U = P.dx * dx_km * 1000 / dt_sec;
V = P.dy * dy_km * 1000 / dt_sec;
speed = sqrt(U.^2 + V.^2);

trueDX_on_piv = interp2(x, y, dx_true, P.ix, P.iy, 'linear', NaN);
trueDY_on_piv = interp2(x, y, dy_true, P.ix, P.iy, 'linear', NaN);

raw_mask = isfinite(P.dx) & isfinite(P.dy);

% -------------------------------------------------------------------------
% Cloud-center QC
% -------------------------------------------------------------------------
% valid_fraction only checks how many pixels in the interrogation windows
% are valid.  Therefore, if a window partly overlaps a cloud, a vector can
% still be estimated even when the PIV center itself lies on a cloud.
% For the teaching example, we explicitly reject such vectors.
cloudMask1_buffer = buffer_logical_mask(________, cloudBufferPix);
cloudMask2_buffer = buffer_logical_mask(________, cloudBufferPix);

origin_clear = sample_logical_mask(~cloudMask1_buffer, ________, ________);
target_clear = sample_logical_mask(~cloudMask2_buffer, P.ix + ________, P.iy + ________);
cloud_qc = raw_mask & ________ & ________;

basic_qc = raw_mask & ...
           ________ & ...
           P.valid_fraction >= ________ & ...
           P.rmax >= ________ & ...
           P.peak_gap >= ________ & ...
           P.texture >= ________ & ...
           ~P.edge_flag;

neighbor_qc = neighbor_consistency_qc(P.dx, P.dy, ________, maxNeighborDiff);
final_qc = ________ & ________;

fprintf('Candidate vectors : %d\n', numel(P.dx));
fprintf('Raw finite vectors: %d\n', nnz(raw_mask));
fprintf('Cloud QC vectors  : %d\n', nnz(cloud_qc));
fprintf('Basic QC vectors  : %d\n', nnz(basic_qc));
fprintf('Final QC vectors  : %d\n', nnz(final_qc));
fprintf('Median dx error   : %.3f grid\n', nanmedian2(P.dx(final_qc) - trueDX_on_piv(final_qc)));
fprintf('Median dy error   : %.3f grid\n', nanmedian2(P.dy(final_qc) - trueDY_on_piv(final_qc)));

%% ============================================================
% 8. Plot bad vectors before QC
% ============================================================
fig = figure('Color','w','Position',[80 80 1000 750]);
imagesc(field1_cloud); set(gca,'YDir','normal'); axis image; colorbar; caxis([0 1]); hold on;
quiver(P.ix(raw_mask), P.iy(raw_mask), P.dx(raw_mask), P.dy(raw_mask), 0, 'r', 'LineWidth', 1.0);
title('run06b: raw MCC/PIV vectors before QC'); xlabel('x grid'); ylabel('y grid');
save_png(fig, 'Fig_run06b_02_bad_vectors_before_QC.png');

%% ============================================================
% 9. Plot QC steps
% ============================================================
fig = figure('Color','w','Position',[80 80 1700 850]);

subplot(2,4,1); imagesc(P.ix(1,:), P.iy(:,1), P.valid_fraction); set(gca,'YDir','normal'); axis image; colorbar; caxis([0 1]); title('Valid fraction');
subplot(2,4,2); imagesc(P.ix(1,:), P.iy(:,1), P.rmax); set(gca,'YDir','normal'); axis image; colorbar; caxis([0 1]); title('Maximum correlation');
subplot(2,4,3); imagesc(P.ix(1,:), P.iy(:,1), P.peak_gap); set(gca,'YDir','normal'); axis image; colorbar; title('Independent peak gap');
subplot(2,4,4); imagesc(P.ix(1,:), P.iy(:,1), double(cloud_qc)); set(gca,'YDir','normal'); axis image; colorbar; caxis([0 1]); title('Cloud-origin/target QC');
subplot(2,4,5); imagesc(P.ix(1,:), P.iy(:,1), double(basic_qc)); set(gca,'YDir','normal'); axis image; colorbar; caxis([0 1]); title('Basic QC');
subplot(2,4,6); imagesc(P.ix(1,:), P.iy(:,1), double(neighbor_qc)); set(gca,'YDir','normal'); axis image; colorbar; caxis([0 1]); title('Neighbor QC');
sgtitle('run06b: QC diagnostics');
save_png(fig, 'Fig_run06b_03_vectors_QC_steps.png');

%% ============================================================
% 10. Plot final vectors
% ============================================================
fig = figure('Color','w','Position',[80 80 1500 650]);

subplot(1,2,1);
imagesc(field1_cloud); set(gca,'YDir','normal'); axis image; colorbar; caxis([0 1]); hold on;
quiver(P.ix(final_qc), P.iy(final_qc), trueDX_on_piv(final_qc), trueDY_on_piv(final_qc), 0, 'k', 'LineWidth', 1.0);
title('True displacement at final-QC points'); xlabel('x grid'); ylabel('y grid');

subplot(1,2,2);
imagesc(field1_cloud); set(gca,'YDir','normal'); axis image; colorbar; caxis([0 1]); hold on;
quiver(P.ix(final_qc), P.iy(final_qc), P.dx(final_qc), P.dy(final_qc), 0, 'r', 'LineWidth', 1.0);
title('Estimated displacement after final QC'); xlabel('x grid'); ylabel('y grid');
sgtitle('run06b: final MCC/PIV vectors after QC');
save_png(fig, 'Fig_run06b_04_final_QC_vectors.png');

%% ============================================================
% 11. Plot final components and error
% ============================================================
fig = figure('Color','w','Position',[80 80 1500 850]);
err = sqrt((P.dx-trueDX_on_piv).^2 + (P.dy-trueDY_on_piv).^2);

subplot(2,3,1); imagesc(P.ix(1,:), P.iy(:,1), P.dx); set(gca,'YDir','normal'); axis image; colorbar; title('Estimated dx [grid]');
subplot(2,3,2); imagesc(P.ix(1,:), P.iy(:,1), P.dy); set(gca,'YDir','normal'); axis image; colorbar; title('Estimated dy [grid]');
subplot(2,3,3); imagesc(P.ix(1,:), P.iy(:,1), err); set(gca,'YDir','normal'); axis image; colorbar; caxis([0 3]); title('Displacement error [grid]');
subplot(2,3,4); imagesc(P.ix(1,:), P.iy(:,1), P.rmax); set(gca,'YDir','normal'); axis image; colorbar; caxis([0 1]); title('Maximum correlation');
subplot(2,3,5); imagesc(P.ix(1,:), P.iy(:,1), P.peak_gap); set(gca,'YDir','normal'); axis image; colorbar; title('Independent peak gap');
subplot(2,3,6); imagesc(P.ix(1,:), P.iy(:,1), double(final_qc)); set(gca,'YDir','normal'); axis image; colorbar; caxis([0 1]); title('Final QC mask');
sgtitle('run06b: final components, correlation, and error');
save_png(fig, 'Fig_run06b_05_final_components_correlation.png');

%% ============================================================
% 12. Save result
% ============================================================
result = struct();
result.P = P;
result.U = U;
result.V = V;
result.speed = speed;
result.raw_mask = raw_mask;
result.cloud_qc = cloud_qc;
result.origin_clear = origin_clear;
result.target_clear = target_clear;
result.cloudMask1_buffer = cloudMask1_buffer;
result.cloudMask2_buffer = cloudMask2_buffer;
result.final_qc = final_qc;
result.basic_qc = basic_qc;
result.neighbor_qc = neighbor_qc;
result.trueDX_on_piv = trueDX_on_piv;
result.trueDY_on_piv = trueDY_on_piv;
result.field1 = field1;
result.field2 = field2;
result.field1_cloud = field1_cloud;
result.field2_cloud = field2_cloud;
result.cloudMask1 = cloudMask1;
result.cloudMask2 = cloudMask2;
result.dx_true = dx_true;
result.dy_true = dy_true;
result.settings = struct('win',win,'step',step,'maxShift',maxShift, ...
    'minValidFraction',minValidFraction,'minCorrelation',minCorrelation, ...
    'minPeakGap',minPeakGap,'minTextureStd',minTextureStd, ...
    'localCorrelation','window mean removed; no global high-pass');

save('run06b_MCC_synthetic_with_cloud_QC_complete_result.mat', 'result');
fprintf('Saved result: run06b_MCC_synthetic_with_cloud_QC_complete_result.mat\n');

%% ========================================================================
% Local functions
% ========================================================================
function ensure_dir(d)
    if ~exist(d,'dir'); mkdir(d); end
end

function B = robust_normalize(A, clipValue)
    x = A(isfinite(A));
    if isempty(x); B = A; return; end
    medx = median(x);
    madx = median(abs(x - medx));
    scale = 1.4826 * madx;
    if ~isfinite(scale) || scale <= 0
        scale = std(x);
    end
    if ~isfinite(scale) || scale <= 0
        scale = 1;
    end
    B = (A - medx) ./ scale;
    B(B >  clipValue) =  clipValue;
    B(B < -clipValue) = -clipValue;
    B(~isfinite(A)) = NaN;
end


function Mbuf = buffer_logical_mask(M, radiusPix)
    if radiusPix <= 0
        Mbuf = logical(M);
        return;
    end
    K = ones(2*radiusPix + 1, 2*radiusPix + 1);
    Mbuf = conv2(double(M), K, 'same') > 0;
end

function ok = sample_logical_mask(M, xq, yq)
    % Nearest-neighbor sampling of a logical mask at floating-point grid points.
    % Points outside the image, or NaN coordinates, are treated as false.
    [ny, nx] = size(M);
    xi = round(xq);
    yi = round(yq);

    ok = false(size(xq));
    inside = isfinite(xi) & isfinite(yi) & xi >= 1 & xi <= nx & yi >= 1 & yi <= ny;

    idx = sub2ind([ny, nx], yi(inside), xi(inside));
    ok(inside) = M(idx);
end

function [P] = mcc_piv_core(I1, I2, win, step, maxShift, minValid, peakExcludeRadius)
    [ny, nx] = size(I1);
    halfWin = floor(win/2);
    margin = halfWin + maxShift + 2;

    iy_list = margin:step:(ny-margin+1);
    ix_list = margin:step:(nx-margin+1);
    [IX, IY] = meshgrid(ix_list, iy_list);

    ncy = numel(iy_list);
    ncx = numel(ix_list);

    dx = nan(ncy, ncx);
    dy = nan(ncy, ncx);
    rmax = nan(ncy, ncx);
    valid_fraction = nan(ncy, ncx);
    texture = nan(ncy, ncx);
    peak_gap = nan(ncy, ncx);
    peak2 = nan(ncy, ncx);
    peak_ratio = nan(ncy, ncx);
    edge_flag = false(ncy, ncx);

    shifts = -maxShift:maxShift;

    for jj = 1:ncy
        iy = iy_list(jj);
        y1 = (iy-halfWin):(iy+halfWin);
        for ii = 1:ncx
            ix = ix_list(ii);
            x1 = (ix-halfWin):(ix+halfWin);

            A = I1(y1, x1);
            texture(jj,ii) = nanstd2(A);
            if nnz(isfinite(A)) / numel(A) < minValid
                continue;
            end

            Cmat  = nan(numel(shifts), numel(shifts));
            VFmat = nan(numel(shifts), numel(shifts));

            for sy = 1:numel(shifts)
                ddy = shifts(sy);
                y2 = y1 + ddy;
                for sx = 1:numel(shifts)
                    ddx = shifts(sx);
                    x2 = x1 + ddx;
                    B = I2(y2, x2);
                    [cc, vf] = window_corr(A, B);
                    Cmat(sy,sx) = cc;
                    VFmat(sy,sx) = vf;
                end
            end

            [bestR, idx] = max_finite(Cmat);
            if ~isfinite(bestR)
                continue;
            end

            [syBest, sxBest] = ind2sub(size(Cmat), idx);
            dyInt = shifts(syBest);
            dxInt = shifts(sxBest);

            rmax(jj,ii) = bestR;
            valid_fraction(jj,ii) = VFmat(syBest, sxBest);
            edge_flag(jj,ii) = abs(dxInt) == maxShift || abs(dyInt) == maxShift;

            secondR = second_peak_excluding_radius(Cmat, syBest, sxBest, peakExcludeRadius);
            peak2(jj,ii) = secondR;
            if isfinite(secondR)
                peak_gap(jj,ii) = bestR - secondR;
                if abs(secondR) > eps
                    peak_ratio(jj,ii) = bestR ./ secondR;
                end
            end

            if valid_fraction(jj,ii) < minValid
                continue;
            end

            dxSub = dxInt;
            dySub = dyInt;
            if sxBest > 1 && sxBest < size(Cmat,2)
                dxSub = dxInt + parabola_subpixel(Cmat(syBest,sxBest-1), Cmat(syBest,sxBest), Cmat(syBest,sxBest+1));
            end
            if syBest > 1 && syBest < size(Cmat,1)
                dySub = dyInt + parabola_subpixel(Cmat(syBest-1,sxBest), Cmat(syBest,sxBest), Cmat(syBest+1,sxBest));
            end

            dx(jj,ii) = dxSub;
            dy(jj,ii) = dySub;
        end
    end

    P.ix = IX;
    P.iy = IY;
    P.dx = dx;
    P.dy = dy;
    P.rmax = rmax;
    P.valid_fraction = valid_fraction;
    P.texture = texture;
    P.peak_gap = peak_gap;
    P.peak2 = peak2;
    P.peak_ratio = peak_ratio;
    P.edge_flag = edge_flag;
end

function [cc, vf] = window_corr(A, B)
    % Local-mean normalized cross correlation.
    % This is the same basic idea used in run06c/d: no global high-pass,
    % but each interrogation window is locally mean-centered.
    m = isfinite(________) & isfinite(________);
    vf = nnz(m) / numel(A);
    if nnz(m) < 12
        cc = NaN;
        return;
    end
    a = A(m);
    b = B(m);
    a = a - ________(a);
    b = b - ________(b);
    den = sqrt(sum(a.^2) * sum(b.^2));
    if den <= 0 || ~isfinite(den)
        cc = NaN;
    else
        cc = sum(________ .* ________) ./ den;
    end
end

function [bestR, idx] = max_finite(A)
    v = A(:);
    ok = isfinite(v);
    if ~any(ok)
        bestR = NaN; idx = NaN; return;
    end
    ii = find(ok);
    [bestR, kk] = max(v(ok));
    idx = ii(kk);
end

function secondR = second_peak_excluding_radius(Cmat, syBest, sxBest, rad)
    C2 = Cmat;
    ny = size(C2,1);
    nx = size(C2,2);
    y1 = max(1, syBest-rad); y2 = min(ny, syBest+rad);
    x1 = max(1, sxBest-rad); x2 = min(nx, sxBest+rad);
    C2(y1:y2, x1:x2) = NaN;
    [secondR, ~] = max_finite(C2);
end

function s = parabola_subpixel(c1, c2, c3)
    if ~isfinite(c1) || ~isfinite(c2) || ~isfinite(c3)
        s = 0; return;
    end
    den = c1 - 2*c2 + c3;
    if abs(den) < 1e-12
        s = 0;
    else
        s = 0.5 * (c1 - c3) / den;
        if abs(s) > 1
            s = 0;
        end
    end
end

function keep2 = neighbor_consistency_qc(dx, dy, keep, maxDiffGrid)
    keep2 = keep;
    [ny, nx] = size(dx);
    for j = 1:ny
        for i = 1:nx
            if ~keep(j,i)
                keep2(j,i) = false;
                continue;
            end
            j1 = max(1,j-1); j2 = min(ny,j+1);
            i1 = max(1,i-1); i2 = min(nx,i+1);
            neigh = keep(j1:j2, i1:i2);
            dxn = dx(j1:j2, i1:i2);
            dyn = dy(j1:j2, i1:i2);
            neigh((j-j1+1), (i-i1+1)) = false;
            vv = neigh & isfinite(dxn) & isfinite(dyn);
            if nnz(vv) < 2
                keep2(j,i) = keep(j,i);
            else
                mdx = median(dxn(vv));
                mdy = median(dyn(vv));
                dd = sqrt((dx(j,i)-mdx).^2 + (dy(j,i)-mdy).^2);
                keep2(j,i) = keep(j,i) && dd <= maxDiffGrid;
            end
        end
    end
end

function s = nanstd2(A)
    x = A(isfinite(A));
    if numel(x) <= 1
        s = NaN;
    else
        s = std(x);
    end
end

function m = nanmedian2(A)
    x = A(isfinite(A));
    if isempty(x)
        m = NaN;
    else
        m = median(x);
    end
end

function v = nanmin2(A)
    x = A(isfinite(A));
    if isempty(x); v = NaN; else; v = min(x); end
end

function v = nanmax2(A)
    x = A(isfinite(A));
    if isempty(x); v = NaN; else; v = max(x); end
end

function save_png(fig, filename)
    try
        exportgraphics(fig, filename, 'Resolution', 180);
    catch
        print(fig, filename, '-dpng', '-r180');
    end
    fprintf('Saved figure: %s\n', filename);
end
