%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Fig. S10: Delta MLD vs Delta Dilution Length Scale by background MLD %%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Produces two separate figures:
%   1) Shallow background MLDs: MLD_base < 30 m
%   2) Deep background MLDs:    MLD_base > 50 m
%
% Density colours follow the quadrant-specific scheme used in Figure 1e:
%   TL = purple, TR = green, BL = blue, BR = orange.
% Quadrant percentages are based on grid-cell counts (as in the original
% Fig. S11 script). A robust linear fit and its 95% confidence limits are
% plotted, and R^2 is printed in the Command Window.

clear; clc; close all;

cd /Users/twrohr/Desktop/Matlab/Matlab/Grazing/CMIP6_Comp/NCC_Figs/Final_Figs/Revised
addpath(genpath('/Users/twrohr/Desktop/Matlab/Matlab/Utilities'))

%% Settings
dataDir = '/Users/twrohr/Desktop/data.nosync/SSP585/Abridged/regridded2';

mod_order = [10,5,1,9,2,6,11,8,4,3,7];

shallowThreshold = 30;  % shallow: background MLD < 30 m
deepThreshold    = 50;  % deep:    background MLD > 50 m

vX = 15;                % Delta MLD
vY = 42;                % Delta Dilution Length Scale

xLimits = [-50 50];
yLimits = [-50 50];
nbin    = 150;

labelX = '$\Delta$ MLD (m)';
labelY = '$\Delta$ Dilution Length Scale (m)';

% Density scaling. Density is expressed as the fraction of all retained
% grid cells in each bin, allowing the same limits to be used in both plots.
clL = 5e-6;
clU = 1.4e-3;
density_gamma = 1.6;
contrast_boost = .5;

set(groot, 'defaultAxesTickLabelInterpreter','latex', ...
           'defaultTextInterpreter','latex', ...
           'defaultLegendInterpreter','latex');

%% Quadrant colour scheme (identical to Figure 1e)
white_rgb = [1.00, 1.00, 1.00];
cool1_rgb = [0.96, 0.97, 0.98];
cool2_rgb = [0.86, 0.88, 0.92];
hot_rgb   = [0, 0, 0];

blue_1_rgb = [0.70, 0.83, 0.98];
blue_2_rgb = [0.43, 0.65, 0.90];
blue_3_rgb = [0.22, 0.45, 0.76];
blue_4_rgb = [0.08, 0.24, 0.55];

green_1_rgb = [0.82, 0.89, 0.57];
green_2_rgb = [0.60, 0.76, 0.35];
green_3_rgb = [0.38, 0.58, 0.19];
green_4_rgb = [0.18, 0.38, 0.08];

purple_1_rgb = [0.86, 0.73, 0.94];
purple_2_rgb = [0.69, 0.50, 0.82];
purple_3_rgb = [0.51, 0.31, 0.68];
purple_4_rgb = [0.32, 0.12, 0.48];

orange_1_rgb = [0.98, 0.76, 0.38];
orange_2_rgb = [0.91, 0.56, 0.16];
orange_3_rgb = [0.75, 0.37, 0.06];
orange_4_rgb = [0.52, 0.20, 0.02];

commonRamp = [white_rgb; cool1_rgb; cool2_rgb];

quadCols = struct();
quadCols.knots = [0.00 0.10 0.20 0.45 0.70 0.85 0.95 1.00];
quadCols.BL_ramp = [commonRamp; blue_1_rgb; blue_2_rgb; ...
                    blue_3_rgb; blue_4_rgb; hot_rgb];
quadCols.TR_ramp = [commonRamp; green_1_rgb; green_2_rgb; ...
                    green_3_rgb; green_4_rgb; hot_rgb];
quadCols.TL_ramp = [commonRamp; purple_1_rgb; purple_2_rgb; ...
                    purple_3_rgb; purple_4_rgb; hot_rgb];
quadCols.BR_ramp = [commonRamp; orange_1_rgb; orange_2_rgb; ...
                    orange_3_rgb; orange_4_rgb; hot_rgb];

%% Load and filter all models
mods = dir(fullfile(dataDir,'*.mat'));
mods = mods(~[mods.isdir]);
[~,isort] = sort({mods.name});
mods = mods(isort);

X_shallow = [];
Y_shallow = [];
X_deep = [];
Y_deep = [];

for ii = 1:11
    mi = mod_order(ii);
    fprintf('Loading model %d of 11: %s\n', ii, mods(mi).name);

    S = load(fullfile(dataDir,mods(mi).name), ...
             'CoreVar_Inv_rw','CoreVar_Inv_adj');

    V = S.CoreVar_Inv_rw;
    MLD_base = squeeze(mean(S.CoreVar_Inv_adj(15,:,:,:),2,'omitnan'));

    dMLD = squeeze(V(vX,:,:));
    dDL  = squeeze(V(vY,:,:));

    shallow = MLD_base < shallowThreshold;
    deep    = MLD_base > deepThreshold;

    goodShallow = shallow & isfinite(dMLD) & isfinite(dDL);
    goodDeep    = deep    & isfinite(dMLD) & isfinite(dDL);

    X_shallow = [X_shallow; dMLD(goodShallow)]; %#ok<AGROW>
    Y_shallow = [Y_shallow; dDL(goodShallow)];  %#ok<AGROW>
    X_deep    = [X_deep;    dMLD(goodDeep)];    %#ok<AGROW>
    Y_deep    = [Y_deep;    dDL(goodDeep)];     %#ok<AGROW>
end

%% Plot only the requested pair of separate figures
makeMLDFigure(X_shallow, Y_shallow, 'Shallow MLDs', ...
    xLimits, yLimits, nbin, clL, clU, density_gamma, contrast_boost, ...
    quadCols, labelX, labelY);

makeMLDFigure(X_deep, Y_deep, 'Deep MLDs', ...
    xLimits, yLimits, nbin, clL, clU, density_gamma, contrast_boost, ...
    quadCols, labelX, labelY);

%% =======================================================================
% Local helpers
% ========================================================================

function makeMLDFigure(x, y, figureName, xL, yL, nbin, clL, clU, ...
                       density_gamma, contrast_boost, quadCols, ...
                       labelX, labelY)

    good = isfinite(x) & isfinite(y);
    x = x(good);
    y = y(good);

    % Unit weights preserve the original script's grid-cell-count logic.
    % Normalization converts bin counts to grid-cell fractions so both
    % figures share the same density scale.
    w = ones(size(x));
    w = w ./ sum(w);

    f = figure('Units','centimeters', ...
               'Position',[2 2 31.2 29.3], ...
               'Color','w', ...
               'Name',figureName);

    mainPos = [3.7 3.6 24.5 23];
    ax = axes('Parent',f, ...
              'Units','centimeters', ...
              'Position',mainPos);

    plotQuadrantRGBDensity(ax, x, y, w, nbin, xL, yL, ...
        clL, clU, density_gamma, contrast_boost, quadCols);
    hold(ax,'on');

    % Robust fit retained from the original Fig. S11 script.
    inPlot = x >= xL(1) & x <= xL(2) & y >= yL(1) & y <= yL(2);
    mdl = fitlm(x(inPlot), y(inPlot), 'linear', 'RobustOpts','on');
    R2 = mdl.Rsquared.Ordinary;

    xg = linspace(xL(1),xL(2),400)';
    [yhat,yCI] = predict(mdl,xg,'Prediction','curve');
    plot(ax,xg,yCI(:,1),'k--','LineWidth',1.4,'HandleVisibility','off');
    plot(ax,xg,yCI(:,2),'k--','LineWidth',1.4,'HandleVisibility','off');
    plot(ax,xg,yhat,'k-','LineWidth',2.2,'HandleVisibility','off');

    xline(ax,0,'k','LineWidth',1.6);
    yline(ax,0,'k','LineWidth',1.6);

    xlabel(ax,labelX,'FontSize',25);
    ylabel(ax,labelY,'FontSize',25);
    title(ax,figureName,'FontSize',28);

    addQuadPerc(ax,x,y,28);

    set(ax, ...
        'FontSize',28, ...
        'TickDir','in', ...
        'TickLabelInterpreter','latex', ...
        'Box','on', ...
        'Layer','top', ...
        'LineWidth',1.1);
    ax.XAxis.FontSize = 28;
    ax.YAxis.FontSize = 28;
    axis(ax,'square');

    fprintf('%s: robust linear regression R^2 = %.4f (n = %d)\n', ...
            figureName,R2,numel(x(inPlot)));
end


function plotQuadrantRGBDensity(ax,x,y,w,nbin,xL,yL,clL,clU, ...
                                density_gamma,contrast_boost,quadCols)

    [Himg,xplot,yplot,Xc,Yc] = ...
        weightedHistogramImage(x,y,w,nbin,xL,yL);

    logH = log10(Himg);
    t = (logH-log10(clL))./(log10(clU)-log10(clL));
    t = max(min(t,1),0);
    t = t.^density_gamma;
    t = min(1,contrast_boost.*t);

    alpha_img = double(isfinite(Himg));
    rgb = ones(size(Himg,1),size(Himg,2),3);

    rgb = applyDensityRamp(rgb,t,Xc<=0 & Yc<=0, ...
                           quadCols.BL_ramp,quadCols.knots);
    rgb = applyDensityRamp(rgb,t,Xc>0 & Yc>0, ...
                           quadCols.TR_ramp,quadCols.knots);
    rgb = applyDensityRamp(rgb,t,Xc<=0 & Yc>0, ...
                           quadCols.TL_ramp,quadCols.knots);
    rgb = applyDensityRamp(rgb,t,Xc>0 & Yc<=0, ...
                           quadCols.BR_ramp,quadCols.knots);

    h = imagesc(ax,xplot,yplot,rgb);
    set(h,'AlphaData',alpha_img);
    set(ax,'YDir','normal','Color',[1 1 1], ...
           'XLim',xL,'YLim',yL);
    axis(ax,'square');
end


function [Himg,xplot,yplot,Xc,Yc] = ...
         weightedHistogramImage(x,y,w,nbin,xL,yL)

    xedges = linspace(xL(1),xL(2),nbin+1);
    yedges = linspace(yL(1),yL(2),nbin+1);

    ix = discretize(x,xedges);
    iy = discretize(y,yedges);
    keep = ~isnan(ix) & ~isnan(iy) & isfinite(w) & w>0;

    H = accumarray([ix(keep),iy(keep)],w(keep), ...
                   [nbin,nbin],@sum,0);
    H(H<=0) = NaN;
    Himg = H.';

    xplot = [xedges(1),xedges(end)];
    yplot = [yedges(1),yedges(end)];
    xcent = 0.5*(xedges(1:end-1)+xedges(2:end));
    ycent = 0.5*(yedges(1:end-1)+yedges(2:end));
    [Xc,Yc] = meshgrid(xcent,ycent);
end


function rgb = applyDensityRamp(rgb,t,qmask,ramp,knots)
    good = qmask & isfinite(t);
    if ~any(good(:))
        return
    end

    tt = t(good);
    for k = 1:3
        channel = rgb(:,:,k);
        channel(good) = interp1(knots,ramp(:,k),tt,'linear','extrap');
        rgb(:,:,k) = channel;
    end
    rgb = max(min(rgb,1),0);
end


function addQuadPerc(ax,x,y,fs)
    good = isfinite(x) & isfinite(y);
    x = x(good);
    y = y(good);
    N = max(numel(x),1);

    qTL = 100*sum(x<=0 & y>0) /N;
    qTR = 100*sum(x>0  & y>0) /N;
    qBR = 100*sum(x>0  & y<=0)/N;
    qBL = 100*sum(x<=0 & y<=0)/N;

    txtfmt = '\\textbf{%d\\%%}';
    m = 0.005;
    xInset = 0.395;

    text(ax,xInset+0.02,1-m,sprintf(txtfmt,round(qTL)), ...
        'Units','normalized','HorizontalAlignment','left', ...
        'VerticalAlignment','top','FontSize',fs,'Interpreter','latex');
    text(ax,1-xInset,1-m,sprintf(txtfmt,round(qTR)), ...
        'Units','normalized','HorizontalAlignment','right', ...
        'VerticalAlignment','top','FontSize',fs,'Interpreter','latex');
    text(ax,1-xInset,m,sprintf(txtfmt,round(qBR)), ...
        'Units','normalized','HorizontalAlignment','right', ...
        'VerticalAlignment','bottom','FontSize',fs,'Interpreter','latex');
    text(ax,xInset,m,sprintf(txtfmt,round(qBL)), ...
        'Units','normalized','HorizontalAlignment','left', ...
        'VerticalAlignment','bottom','FontSize',fs,'Interpreter','latex');
end
