%% Robust eval/plot script for res.* (handles missing/partial dims)
% res.ffe stored like: output_ffe{ch, rop, realiz, eval_distance}
% BUT: some dimensions may be singleton or not fully filled (e.g. only 0 km, only 1 realization)

% -------------------- Basic metadata --------------------
rop            = res.settings.rop(:);
wavelengthplan = res.settings.wavelengthplan(:);
distances      = res.eval_dist_km(:);              % can be [0], [2 10], etc.

N_ch  = numel(wavelengthplan);
N_rop = numel(rop);

% -------------------- Determine dims robustly --------------------
% Ensure we have 4D addressing even if MATLAB collapses trailing singletons
dims = size(res.ffe);
if numel(dims) < 4
    dims(end+1:4) = 1;
end
N_realiz    = dims(3);
N_distances = dims(4);

% If res.eval_dist_km length differs from dims(4), trust dims(4)
if numel(distances) ~= N_distances
    % best-effort fallback
    distances = (1:N_distances).';
end

% -------------------- Plot settings --------------------
fec = 3.8e-3;  % choose one; you can switch to 2.2e-4 if needed

qLow  = 0.00;
qHigh = 1.00;

% Color per wavelength/channel
try
    cols = ccbrewer2('Set1', N_ch);
catch
    cols = linspecer(N_ch);
end

% Containers for FEC crossings per channel (per eval distance)
Sffe  = cell(N_distances, N_ch);
Sdfe  = cell(N_distances, N_ch);
Svnle = cell(N_distances, N_ch);
Smlse = cell(N_distances, N_ch);
Sdbt  = cell(N_distances, N_ch);

% -------------------- Choose which eval distances to plot --------------------
% If you want all: eval_list = 1:N_distances;
% If you want a specific distance (e.g., last): eval_list = N_distances;
eval_list = 1;%1:N_distances;

% -------------------- BER vs ROP figures (one figure per eval distance) --------------------
for eval_ptr = eval_list

    figure('Name', sprintf('BER vs ROP @ %s', distLabel(distances, eval_ptr)));
    hold on;

    for ch = 1:N_ch

        % ---- Extract cell slices robustly: [N_rop x N_realiz] ----
        ffe_cells  = sliceCells4D(res.ffe , ch, eval_ptr, N_rop);
        dfe_cells  = sliceCells4D(res.dfe , ch, eval_ptr, N_rop);
        vnle_cells = sliceCells4D(res.vnle, ch, eval_ptr, N_rop);
        mlse_cells = sliceCells4D(res.mlse, ch, eval_ptr, N_rop);
        dbt_cells  = sliceCells4D(res.dbt , ch, eval_ptr, N_rop);

        % ---- FEC crossings (works with incomplete columns) ----
        [Sffe{eval_ptr,ch},  ~] = fecCrossings(rop, ffe_cells,  fec);
        [Sdfe{eval_ptr,ch},  ~] = fecCrossings(rop, dfe_cells,  fec);
        [Svnle{eval_ptr,ch}, ~] = fecCrossings(rop, vnle_cells, fec);
        [Smlse{eval_ptr,ch}, ~] = fecCrossings(rop, mlse_cells, fec);
        [Sdbt{eval_ptr,ch},  ~] = fecCrossings(rop, dbt_cells,  fec);

        % ---- Extract BER matrices using only complete realizations ----
        ffe_mat  = extractCompleteBER(ffe_cells);    % N_rop x K
        dfe_mat  = extractCompleteBER(dfe_cells);
        vnle_mat = extractCompleteBER(vnle_cells);
        mlse_mat = extractCompleteBER(mlse_cells);
        dbt_mat  = extractCompleteBER(dbt_cells);

        showLegend = (ch == 1); % one legend entry per technique

        % Plot bands + mean (only if non-empty)
        if ~isempty(ffe_mat)
            plotBandMeanBL(rop, ffe_mat,  cols(ch,:), sprintf('FFE @ %dnm', round(wavelengthplan(ch))), qLow, qHigh, '--s', showLegend);
        end
        % if ~isempty(dfe_mat)
        %     plotBandMeanBL(rop, dfe_mat,  cols(ch,:), sprintf('DFE @ %dnm', round(wavelengthplan(ch))), qLow, qHigh, '-o', showLegend);
        % end
        % if ~isempty(vnle_mat)
        %     plotBandMeanBL(rop, vnle_mat, cols(ch,:), sprintf('VNLE @ %dnm', round(wavelengthplan(ch))), qLow, qHigh, '--x', showLegend);
        % end
        % if ~isempty(mlse_mat)
        %     plotBandMeanBL(rop, mlse_mat, cols(ch,:), sprintf('VNLE+PF+MLSE @ %dnm', round(wavelengthplan(ch))), qLow, qHigh, '-o',  showLegend);
        % end
        % if ~isempty(dbt_mat)
        %     plotBandMeanBL(rop, dbt_mat,  cols(ch,:), sprintf('DBt+MLSE @ %dnm', round(wavelengthplan(ch))), qLow, qHigh, '--v', showLegend);
        % end

    end

    set(gca,'XScale','linear','YScale','log','TickLabelInterpreter','latex','FontSize',11);
    yline([3.8e-3, 2.2e-4], 'HandleVisibility','off','LineWidth',1.5);
    ylabel('BER');
    xlabel('ROP [dB]');
    title(sprintf('BER vs. ROP @ %s', distLabel(distances, eval_ptr)));
    xlim([min(rop) max(rop)]);
    ylim([1e-5 0.3]);
    grid on;
    legend show;
end


%% -------------------- VIOLIN of ROP at FEC crossing --------------------
% Build technique groups; each group is {N_distances x N_ch} cell entries
% You can pick a distance to show (e.g., eval_ptr=1 for 0 km, or last)
eval_ptr_violin = min(N_distances, max(1, N_distances)); % default: last available
% eval_ptr_violin = 1;

% Choose which technique(s) to show
techNames = {'VNLE'};
S_groups  = {Svnle(eval_ptr_violin,:)};

% Only plot if violinplot exists
if exist('violinplot','file') == 2
    figure('Name', sprintf('FEC crossing violin @ %s', distLabel(distances, eval_ptr_violin))); hold on;

    colsTech = linspecer(numel(S_groups));

    for i = 1:numel(S_groups)
        S_cell = S_groups{i}; % 1 x N_ch cell, each cell is 1xK crossings

        % Pad to rectangular: rows = realizations (max K), cols = wavelengths (N_ch)
        Kmax = max(cellfun(@numel, S_cell));
        if isempty(Kmax) || Kmax == 0
            continue;
        end
        S_mat = zeros(Kmax, N_ch);
        for ch = 1:N_ch
            k = numel(S_cell{ch});
            if k > 0
                S_mat(1:k, ch) = S_cell{ch}(:);
            end
        end

        catLabels = arrayfun(@(nm) sprintf('%d nm', nm), round(wavelengthplan), 'UniformOutput', false);
        violinplot(S_mat, catLabels, ...
            'ViolinColor', colsTech(i,:), ...
            'ViolinAlpha', 0.10, ...
            'MarkerSize', 20, ...
            'ShowMedian', true, ...
            'EdgeColor', colsTech(i,:), ...
            'ShowWhiskers', false, ...
            'ShowData', true, ...
            'ShowBox', false, ...
            'Bandwidth', 0.05);

        % Add a legend proxy
        plot(nan, nan, 'o', 'Color', colsTech(i,:), 'DisplayName', techNames{i});
    end

    ylim([-10,-6])
    ylabel('ROP at FEC crossing');
    title(sprintf('ROP to cross BER %.2e @ %s', fec, distLabel(distances, eval_ptr_violin)));
    grid on; box on;
    legend show;
else
    fprintf('violinplot.m not found on path -> skipping violin plot.\n');
end


%% ================= helpers =================

function cells2D = sliceCells4D(C4, ch, eval_ptr, N_rop)
% Returns a [N_rop x N_realiz] cell array (with N_realiz inferred)
% Works even if trailing dimensions are singleton.
%
% If C4 is smaller than expected, missing entries are returned as [].

dims = size(C4);
if numel(dims) < 4, dims(end+1:4) = 1; end
N_realiz = dims(3);

cells2D = cell(N_rop, N_realiz);

% bounds (in case stored arrays are smaller)
chMax   = dims(1);
ropMax  = dims(2);
rMax    = dims(3);
eMax    = dims(4);

if ch > chMax || eval_ptr > eMax
    return; % all empty
end

ropUse = min(N_rop, ropMax);
rUse   = min(N_realiz, rMax);

% Extract and place into a consistent sized cell matrix
tmp = squeeze(C4(ch, 1:ropUse, 1:rUse, eval_ptr));
% squeeze may return vector/empty if rUse==1 etc. Normalize:
tmp = reshape(tmp, ropUse, rUse);

cells2D(1:ropUse, 1:rUse) = tmp;
end

function lbl = distLabel(distances, eval_ptr)
if eval_ptr <= numel(distances)
    d = distances(eval_ptr);
    if isfinite(d)
        lbl = sprintf('%.0f km', d);
        return;
    end
end
lbl = sprintf('eval\\_%d', eval_ptr);
end

function plotBandMeanBL(x, Y, color, techLabel, qLow, qHigh, lineSpec, showLegend)
% Y: (nPoints x nRealizations)
if isempty(Y), return; end

% Remove realizations that are entirely zero
badCols = all(Y == 0, 1);
Y(:, badCols) = [];

if isempty(Y), return; end

% Avoid log(0)
Y(Y==0) = 1e-8;

mu = mean(Y, 2, 'omitnan');
lo = quantile(Y, qLow,  2);
hi = quantile(Y, qHigh, 2);
b  = [mu - lo, hi - mu];

% boundedline optional; fallback to simple plot if missing
if 0%exist('boundedline','file') == 2
    [hl, hp] = boundedline(x(:), mu(:), b, lineSpec, 'alpha', 'transparency', 0.18);
    set(hl, 'Color', color, 'LineWidth', 1.4, 'MarkerSize', 4);
    set(hp, 'FaceColor', color, 'HandleVisibility','off');
else
    hl = plot(x(:), mu(:), lineSpec, 'LineWidth', 1.4, 'MarkerSize', 4);
    set(hl, 'Color', color);
end

if showLegend
    set(hl, 'DisplayName', techLabel);
else
    set(hl, 'HandleVisibility','off');
end

if 0%exist('outlinebounds','file') == 2 && exist('boundedline','file') == 2
    ho = outlinebounds(hl, hp);
    set(ho, 'linestyle', ':', 'color', color, 'linewidth', 1, 'HandleVisibility','off');
end
end

function [S, noCrossingMask] = fecCrossings(rop, cellsNxR, fec)
% cellsNxR: N_rop x R cell array; each cell is struct with .metrics.BER
Y = extractCompleteBER(cellsNxR);   % -> N_rop x K
if isempty(Y)
    S = [];
    noCrossingMask = [];
    return;
end

ok = mean(Y,1,'omitnan') <= 0.1;
Y  = Y(:, ok);
if isempty(Y)
    S = [];
    noCrossingMask = [];
    return;
end

nR = size(Y,2);
S  = nan(1,nR);
noCrossingMask = true(1,nR);

rop = rop(:);
for j = 1:nR
    y = Y(:,j);
    above = (y > fec);
    idx = find(above(1:end-1) & ~above(2:end), 1, 'first');
    if ~isempty(idx)
        x1 = rop(idx);   y1 = y(idx);
        x2 = rop(idx+1); y2 = y(idx+1);
        if isfinite(y1) && isfinite(y2) && y2 ~= y1
            t  = (fec - y1) / (y2 - y1);
            S(j) = x1 + t*(x2 - x1);
            noCrossingMask(j) = false;
        end
    end
end
end

function Y = extractCompleteBER(cellSlice)
% Keep only those realization columns where ALL N_rop entries are non-empty
if isempty(cellSlice), Y = []; return; end
nR = size(cellSlice,2);
keep = false(1,nR);

for r = 1:nR
    col = cellSlice(:,r);
    keep(r) = all(cellfun(@(c) ~isempty(c), col));
end

if ~any(keep), Y = []; return; end

Y = cellfun(@(c) c.metrics.BER, cellSlice(:,keep), 'UniformOutput', true);
end

function Y = extractCompleteAlphas(cellSlice)
if isempty(cellSlice), Y = []; return; end
nR = size(cellSlice,2);
keep = false(1,nR);

for r = 1:nR
    col = cellSlice(:,r);
    keep(r) = all(cellfun(@(c) ~isempty(c), col));
end

if ~any(keep), Y = []; return; end
Y = cellfun(@(c) c.metrics.Alpha, cellSlice(:,keep), 'UniformOutput', true);
end