%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Figure 3 - Carbon Transfer Plots %%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%
cd /Users/twrohr/Desktop/Matlab/Matlab/Grazing/CMIP6_Comp/NCC_Figs
addpath(genpath('/Users/twrohr/Desktop/Matlab/Matlab/Utilities'))
close all;

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

% ---- Global look ----
set(groot, 'defaultAxesTickLabelInterpreter','latex', ...
           'defaultTextInterpreter','latex', ...
           'defaultLegendInterpreter','latex');

f  = figure('Units','normalized','Position',[0.05 0.05 0.5 1], 'Color','w');
tl = tiledlayout(f,3,2,'TileSpacing','compact','Padding','loose');  % <-- 3x2 now

% ---- Models ----
model_names = {'iHAMMOCC','CMOC','WOMBAT','OECOv2','MARBL','CanOE','MEDUSA2.1','PISCESv2-IPSL','PISCESv2-CNRM','BFM5.2','COBALTv2','CbPM','eVGPM'};
mod_order   = [10,5,1,9,2,6,11,8,4,3,7];
ms          = 1:11;                      % which models to include

% ---- Variables/rows ----
rowSpec = { ...
    struct('v2',3, 'v2tit','$\Delta_{\%} \textrm{NPP} \; ($\%$)$', 'clrmap','tofino'), ...
  struct('v2',8, 'v2tit','$\Delta_{\%} \textrm{GSP} \; ($\%$)$', 'clrmap','berlin'), ...
  struct('v2',5, 'v2tit','$\Delta_{\%} \textrm{EP} \; ($\%$)$',  'clrmap','lisbon') ...
};

% ---- Common plot settings ----
v1   = 1;   labelX = '$\Delta_{\%} \Sigma \textrm{P} \; ($\%$)$';
xmin = -90; xmax = 90;
ymin = -90; ymax = 90;
nbin = 100;
outmin = 0; outmax = 100;       % percentile trimming
clL = 9e-6; clU = .02;          % color limits for WEIGHTED mass per bin (log)

axH = gobjects(3,2);

% ======================================================================
% Loop rows
% ======================================================================
for r = 1:3
    spec = rowSpec{r};

    % ---- Build colormap with white head ----
    if spec.v2==8 || spec.v2==5
        crmap = crameri(spec.clrmap,170); crmap = crmap(21:170,:);
    else
        crmap = crameri(spec.clrmap,150);
    end
    cmap = [ones(50,3); crmap];
    for ii = 1:50
        t = ii/50;
        cmap(ii,:) = t.*cmap(51,:) + (1-t).*[1 1 1];
    end

    % ---- Gather data across models for this row ----
    P_change_BU=[];   Y_BU=[];    W_BU=[];
    P_change_TDgp=[]; Y_TDgp=[];  W_TDgp=[];

    for m = ms
        S = dir(dataDir); S = S(~[S.isdir]);
        load(fullfile(dataDir, S(mod_order(m)).name), ...
            'CoreVar_Inv_pc','CoreVar_Inv_rw','Lat','Bio_Weights');

        % Masks
        maskvar1 = sign(squeeze(CoreVar_Inv_pc(1,:,:)).*squeeze(CoreVar_Inv_pc(6,:,:))); % BU vs TD
        maskvar2 = ones(size(maskvar1)); % placeholder if you later re-enable TD grazing split

        % Model-normalization: each model total weight = 1
        Wnorm = nansum(Bio_Weights(:));
        if ~isfinite(Wnorm) || Wnorm<=0
            continue
        end

        % ---- Bottom-up subset ----
        M = (maskvar1>0);
        x = CoreVar_Inv_pc(v1,     M); x=x(:);
        y = CoreVar_Inv_pc(spec.v2,M); y=y(:);
        w = Bio_Weights(M) ./ Wnorm;   w=w(:);

        outsx = prctile(x,[outmin,outmax]); outsy = prctile(y,[outmin,outmax]);
        bad = (x<outsx(1) | x>outsx(2) | y<outsy(1) | y>outsy(2) | ~isfinite(x) | ~isfinite(y) | ~isfinite(w) | w<=0);
        x(bad)=[]; y(bad)=[]; w(bad)=[];

        P_change_BU = [P_change_BU; x];
        Y_BU        = [Y_BU;        y];
        W_BU        = [W_BU;        w];

        % ---- Top-down (grazing-dominated) ----
        M = (maskvar1<0 & maskvar2>0);
        x = CoreVar_Inv_pc(v1,     M); x=x(:);
        y = CoreVar_Inv_pc(spec.v2,M); y=y(:);
        w = Bio_Weights(M) ./ Wnorm;   w=w(:);

        outsx = prctile(x,[outmin,outmax]); outsy = prctile(y,[outmin,outmax]);
        bad = (x<outsx(1) | x>outsx(2) | y<outsy(1) | y>outsy(2) | ~isfinite(x) | ~isfinite(y) | ~isfinite(w) | w<=0);
        x(bad)=[]; y(bad)=[]; w(bad)=[];

        P_change_TDgp = [P_change_TDgp; x];
        Y_TDgp        = [Y_TDgp;        y];
        W_TDgp        = [W_TDgp;        w];
    end

    % ---- 2 panels for this row ----
    titlesCol = {'Bottom-up Control','Top-down Control'};
    XX = {P_change_BU,  P_change_TDgp};
    YY = {Y_BU,         Y_TDgp};
    WW = {W_BU,         W_TDgp};

    % --- Weighted fractions per column
    Wtot   = sum(cellfun(@(w) nansum(w), WW));
    colPct = cellfun(@(w) nansum(w), WW) ./ max(Wtot, 1e-12);
    subCol = arrayfun(@(p) sprintf('(%.0f\\%%)', 100*p), colPct, 'UniformOutput', false);

    for c = 1:2
        axH(r,c) = nexttile(tl,(r-1)*2+c);

        weightedBinscatter(axH(r,c), XX{c}, YY{c}, WW{c}, nbin, [xmin xmax], [ymin ymax]);

        set(axH(r,c),'ColorScale','log');
        clim(axH(r,c), [clL clU]);
        colormap(axH(r,c), cmap);

        xline(axH(r,c),0,'k','LineWidth',2);
        yline(axH(r,c),0,'k','LineWidth',2);
        xlim(axH(r,c),[xmin xmax]);
        ylim(axH(r,c),[ymin ymax]);

        if r==1
            title(axH(r,c), {titlesCol{c}, subCol{c}}, 'Interpreter','latex', 'FontSize',20);
        else
            title(axH(r,c), '');
        end

        if r==3
            xlabel(axH(r,c), labelX, 'FontSize', 16);
        end
        if c==1
            ylabel(axH(r,c), spec.v2tit, 'FontSize', 20);
        end

        addQuadPercW(XX{c}, YY{c}, WW{c});

        % Colorbar only on right column (one per row)
        if c==1
            colorbar(axH(r,c),'off');
        else
            cb = colorbar(axH(r,c));
            cb.FontSize = 14;
            cb.TickLabelInterpreter = 'latex';
            ylabel(cb,'Biomass-weighted frequency','Interpreter','latex','FontSize',16);

            % decade ticks
            cl = clim(axH(r,c));
            tmin = floor(log10(cl(1)));
            tmax = ceil(log10(cl(2)));
            cb.Ticks = 10.^(tmin:tmax);
            cb.TickLabels = arrayfun(@(k) sprintf('$10^{%d}$',k), tmin:tmax, 'UniformOutput', false);
        end
    end
end

% Labels (a) ... for a 3x2 grid  => a..f
axList = axH.';           
axList = axList(:);
labs = arrayfun(@(k) sprintf('\\textbf{%s}', char('a'+k-1)), ...
                1:numel(axList), 'UniformOutput', false);
placeSubplotLetters(axList, labs, [-0.115 .99], 24);

% ======================================================================
% Helpers
% ======================================================================

function addQuadPercW(x,y,w)
    x = x(:); y = y(:); w = w(:);
    good = isfinite(x) & isfinite(y) & isfinite(w) & (w>0);
    x = x(good); y = y(good); w = w(good);

    W = max(nansum(w), 1e-12);

    q = zeros(1,4);
    q(1) = 100*nansum(w(x<=0 & y>0))/W;   % TL
    q(2) = 100*nansum(w(x>0  & y>0))/W;   % TR
    q(3) = 100*nansum(w(x>0  & y<=0))/W;  % BR
    q(4) = 100*nansum(w(x<=0 & y<=0))/W;  % BL

    ax = gca;
    txtfmt = '\\textbf{%d\\%%}';
    m  = 0.015;
    fs = 18;

    text(ax, m,    1-m, sprintf(txtfmt, round(q(1))), 'Units','normalized', ...
         'HorizontalAlignment','left',  'VerticalAlignment','top', ...
         'FontSize', fs, 'Interpreter','latex');
    text(ax, 1-m,  1-m, sprintf(txtfmt, round(q(2))), 'Units','normalized', ...
         'HorizontalAlignment','right', 'VerticalAlignment','top', ...
         'FontSize', fs, 'Interpreter','latex');
    text(ax, 1-m,  m,   sprintf(txtfmt, round(q(3))), 'Units','normalized', ...
         'HorizontalAlignment','right', 'VerticalAlignment','bottom', ...
         'FontSize', fs, 'Interpreter','latex');
    text(ax, m,    m,   sprintf(txtfmt, round(q(4))), 'Units','normalized', ...
         'HorizontalAlignment','left',  'VerticalAlignment','bottom', ...
         'FontSize', fs, 'Interpreter','latex');
end

function placeSubplotLetters(axList, labels, offset, fs)
    if nargin < 3 || isempty(offset), offset = [-0.1 1]; end
    if nargin < 4 || isempty(fs),     fs     = 18;       end

    for k = 1:numel(axList)
        ax = axList(k);
        text(ax, offset(1), offset(2), labels{k}, ...
            'Units','normalized', ...
            'Interpreter','latex', ...
            'HorizontalAlignment','left', ...
            'VerticalAlignment','bottom', ...
            'FontSize', fs, ...
            'Clipping','off');
    end
end

function weightedBinscatter(ax, x, y, w, nbin, xL, yL)
    x = x(:); y = y(:); w = w(:);
    good = isfinite(x) & isfinite(y) & isfinite(w) & (w>0);
    x = x(good); y = y(good); w = w(good);

    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);
    ix = ix(keep); iy = iy(keep); w = w(keep);

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

    imagesc(ax, xedges, yedges, H.');
    set(ax,'YDir','normal');
    xlim(ax, xL); ylim(ax, yL);
end
