function [symbols_for_lvl, avg_for_lvl, xAxisUs, info] = showLevelScatter(rxInput, refSymbols, options)
%SHOWLEVELSCATTER Plot received samples separated by reference PAM level.
%   Supports plain numeric vectors, Signal objects, synchronized scope cell
%   arrays, and raw unsynchronized Signal input. Raw Signal input is
%   synchronized to refSymbols and stitched before plotting.

arguments
    rxInput
    refSymbols
    options.fignum (1,1) double = NaN
    options.displayname (1,:) char = ''
    options.f_sym double = []
    options.fsym double = []
    options.syncFs (1,1) double = 0
    options.shiftFs (1,1) double = 0
    options.shifts double = []
    options.maxOccurences (1,1) double = Inf
    options.normalize (1,1) logical = false
    options.debug_plots (1,1) logical = false
    options.showPlot (1,1) logical = true
    options.clear (1,1) logical = true
    options.windowLength (1,1) double {mustBePositive, mustBeInteger} = 500
    options.xLimits double = []
    options.yLimits (1,2) double = [-2.5 2.5]
    options.showStdAnnotations (1,1) logical = true
    options.scatterAlpha (1,1) double {mustBeGreaterThanOrEqual(options.scatterAlpha, 0), mustBeLessThanOrEqual(options.scatterAlpha, 1)} = 0.2
    options.scatterSize (1,1) double {mustBePositive} = 1
    options.avgLineMaxPoints (1,1) double {mustBePositive, mustBeInteger} = 1000
    options.avgLineSmoothWindow (1,1) double {mustBePositive, mustBeInteger} = 1
end

fsym = resolveSymbolRate(rxInput, refSymbols, options);
[symbols_for_lvl, avg_for_lvl, xAxisUs, info] = prepareLevelScatterData(rxInput, refSymbols, fsym, options);

if options.showPlot
    info.plot = plotLevelScatter(symbols_for_lvl, avg_for_lvl, refSymbols, xAxisUs, options);
end
end

function fsym = resolveSymbolRate(rxInput, refSymbols, options)
if ~isempty(options.fsym)
    fsym = options.fsym;
elseif ~isempty(options.f_sym)
    fsym = options.f_sym;
elseif isa(refSymbols, "Signal") && ~isempty(refSymbols.fs)
    fsym = refSymbols.fs;
elseif isa(rxInput, "Signal") && ~isempty(rxInput.fs)
    fsym = rxInput.fs;
elseif iscell(rxInput) && ~isempty(rxInput) && isa(rxInput{1}, "Signal") && ~isempty(rxInput{1}.fs)
    fsym = rxInput{1}.fs;
else
    fsym = 1e6;
end
end

function [symbols_for_lvl, avg_for_lvl, xAxisUs, info] = prepareLevelScatterData(rxInput, refSymbols, fsym, options)
refSignal = numericSignal(refSymbols);
info = defaultInfo(fsym);

if iscell(rxInput)
    [symbols_for_lvl, avg_for_lvl, info] = prepareCellInput(rxInput, refSymbols, fsym, options, info);
elseif isa(rxInput, "Signal")
    [symbols_for_lvl, avg_for_lvl, info] = prepareSignalInput(rxInput, refSymbols, fsym, options, info);
else
    rxSymbols = numericSignal(rxInput);
    [symbols_for_lvl, avg_for_lvl] = levelScatterForOneSequence(rxSymbols, refSignal, options.windowLength);
end

xAxisUs = ((1:size(avg_for_lvl, 2)) / fsym) * 1e6;
end

function info = defaultInfo(fsym)
info = struct();
info.found_sync = true;
info.startSamples = 1;
info.shifts = [];
info.fsym = fsym;
info.shiftFs = fsym;
info.varianceByLevel = [];
info.plot = struct();
end

function [symbols_for_lvl, avg_for_lvl, info] = prepareSignalInput(rxSignal, refSymbols, fsym, options, info)
rxAtSymbolRate = rxSignal.resample("fs_in", rxSignal.fs, "fs_out", fsym);
refSignal = numericSignal(refSymbols);

if numel(rxAtSymbolRate.signal) == numel(refSignal)
    rxSymbols = rxAtSymbolRate.signal;
    if options.normalize
        rxSymbols = normalizeNumericRms(rxSymbols);
    end

    [symbols_for_lvl, avg_for_lvl] = levelScatterForOneSequence(rxSymbols, refSignal, options.windowLength);
    info.varianceByLevel = var(symbols_for_lvl, 0, 2, "omitnan");
    return
end

[scopeCell, shifts, shiftFs, foundSync] = synchronizeRawSignal(rxSignal, refSymbols, fsym, options);
info.found_sync = foundSync;
info.shifts = shifts;
info.shiftFs = shiftFs;

if isempty(scopeCell)
    symbols_for_lvl = [];
    avg_for_lvl = [];
    warning("showLevelScatter:NoScopeCells", ...
        "No synchronized scope signal occurrences available.");
    return
end

[symbols_for_lvl, avg_for_lvl, startSamples] = stitchScopeCells(scopeCell, refSymbols, fsym, shifts, shiftFs, options);
info.startSamples = startSamples;
info.varianceByLevel = var(symbols_for_lvl, 0, 2, "omitnan");
end

function [symbols_for_lvl, avg_for_lvl, info] = prepareCellInput(scopeCell, refSymbols, fsym, options, info)
scopeCell = scopeCell(:);
if isempty(scopeCell)
    symbols_for_lvl = [];
    avg_for_lvl = [];
    info.found_sync = false;
    warning("showLevelScatter:NoScopeCells", ...
        "No synchronized scope signal occurrences available.");
    return
end

shiftFs = options.shiftFs;
if shiftFs <= 0
    shiftFs = fsym;
end

[symbols_for_lvl, avg_for_lvl, startSamples] = stitchScopeCells(scopeCell, refSymbols, fsym, options.shifts, shiftFs, options);
info.found_sync = true;
info.shifts = options.shifts;
info.shiftFs = shiftFs;
info.startSamples = startSamples;
info.varianceByLevel = var(symbols_for_lvl, 0, 2, "omitnan");
end

function [scopeCell, shifts, shiftFs, foundSync] = synchronizeRawSignal(rxSignal, refSymbols, fsym, options)
syncFs = options.syncFs;
if syncFs <= 0
    syncFs = 2*fsym;
end

syncSignal = rxSignal.resample("fs_in", rxSignal.fs, "fs_out", syncFs);
if options.normalize
    syncSignal = syncSignal.normalize("mode", "rms");
end

[~, scopeCell, ~, foundSync, shifts] = syncSignal.tsynch( ...
    "reference", refSymbols, ...
    "fs_ref", fsym, ...
    "debug_plots", options.debug_plots);

if options.shiftFs > 0
    shiftFs = options.shiftFs;
else
    shiftFs = syncFs;
end
end

function [symbols_for_lvl, avg_for_lvl, startSamples] = stitchScopeCells(scopeCell, refSymbols, fsym, shifts, shiftFs, options)
recordOccurrences = min(numel(scopeCell), options.maxOccurences);
scopeCell = scopeCell(1:recordOccurrences);
refSignal = numericSignal(refSymbols);
startSamples = getStartSamples(shifts, recordOccurrences, shiftFs, fsym, numel(refSignal));

levelScatter = cell(1, recordOccurrences);
levelAverage = cell(1, recordOccurrences);

for occurrenceIdx = 1:recordOccurrences
    occurrence = scopeCell{occurrenceIdx};
    if isa(occurrence, "Signal")
        occurrence = occurrence.resample("fs_out", fsym);
        occurrence = occurrence.signal;
    end

    [levelScatter{occurrenceIdx}, levelAverage{occurrenceIdx}] = ...
        levelScatterForOneSequence(occurrence, refSignal, options.windowLength);
end

numLevels = size(levelScatter{1}, 1);
traceLength = max(startSamples(:).' + cellfun(@(x) size(x, 2), levelScatter) - 1);
symbols_for_lvl = NaN(numLevels, traceLength);
avg_for_lvl = NaN(numLevels, traceLength);

for occurrenceIdx = 1:recordOccurrences
    writeIdx = startSamples(occurrenceIdx):(startSamples(occurrenceIdx) + size(levelScatter{occurrenceIdx}, 2) - 1);
    symbols_for_lvl(:, writeIdx) = levelScatter{occurrenceIdx};
    avg_for_lvl(:, writeIdx) = levelAverage{occurrenceIdx};
end
end

function startSamples = getStartSamples(shifts, recordOccurrences, shiftFs, fsym, symbolLength)
if isempty(shifts)
    startSamples = ((0:recordOccurrences-1) .* symbolLength) + 1;
    return
end

shifts = shifts(:);
if numel(shifts) ~= recordOccurrences
    positiveShifts = shifts(shifts >= 0);
    if numel(positiveShifts) >= recordOccurrences
        shifts = positiveShifts(1:recordOccurrences);
    else
        warning("showLevelScatter:ShiftCountMismatch", ...
            "Shift count does not match scope cell count. Using sequential stitching.");
        startSamples = ((0:recordOccurrences-1) .* symbolLength) + 1;
        return
    end
end

startSamples = round((shifts(1:recordOccurrences) - shifts(1)) ./ shiftFs .* fsym) + 1;
end

function [symbols_for_lvl, avg_for_lvl] = levelScatterForOneSequence(rxSymbols, refSymbols, windowLength)
rxSymbols = numericSignal(rxSymbols);
refSymbols = numericSignal(refSymbols);

assert(numel(rxSymbols) == numel(refSymbols), ...
    'showLevelScatter:LengthMismatch', ...
    'rxInput and refSymbols must have the same number of samples after resampling/synchronization.');

levels = unique(refSymbols);
[symbols_for_lvl, levels] = splitByReferenceLevels(rxSymbols, refSymbols, levels);
avg_for_lvl = NaN(numel(levels), numel(refSymbols));

for levelIdx = 1:numel(levels)
    levelMask = ~isnan(symbols_for_lvl(levelIdx, :));
    levelSamples = symbols_for_lvl(levelIdx, levelMask);
    if isempty(levelSamples)
        continue
    end

    smoothWindowLength = min(windowLength, numel(levelSamples));
    avg_for_lvl(levelIdx, levelMask) = movmean(levelSamples, smoothWindowLength, 'Endpoints', 'shrink');
    avg_for_lvl(levelIdx, :) = interpolateMissingLevelAverage(avg_for_lvl(levelIdx, :));
end
end

function [symbols_for_lvl, levels] = splitByReferenceLevels(rxSymbols, refSymbols, levels)
supportedPamLevels = [2 4 6 8 16];

if ismember(numel(levels), supportedPamLevels)
    [symbols_for_lvl, levels] = PAMmapper(numel(levels), 0).splitByReferenceLevels( ...
        rxSymbols, refSymbols, ...
        "levels", levels);
else
    symbols_for_lvl = NaN(numel(levels), numel(refSymbols));
    for levelIdx = 1:numel(levels)
        levelMask = refSymbols == levels(levelIdx);
        symbols_for_lvl(levelIdx, levelMask) = rxSymbols(levelMask);
    end
end
end

function plotInfo = plotLevelScatter(symbols_for_lvl, avg_for_lvl, refSymbols, xAxisUs, options)
plotInfo = struct();
if isempty(symbols_for_lvl)
    return
end

if isnan(options.fignum)
    figHandle = figure;
else
    figHandle = figure(options.fignum);
    if options.clear
        clf;
    end
end

ax = gca;
hold(ax, "on");
numLevels = size(symbols_for_lvl, 1);
cols = cbrewer2("RdBu", numLevels);
% cols = linspecer(numLevels);
xLimits = getXLimits(options.xLimits, xAxisUs, symbols_for_lvl);

if 1
    for levelIdx = numLevels:-1:1
        scatter(ax, xAxisUs, symbols_for_lvl(levelIdx, :), options.scatterSize, ".", ...
            "MarkerEdgeColor", cols(levelIdx, :), ...
            "MarkerEdgeAlpha", options.scatterAlpha);
    end
end

for levelIdx = 1:numLevels
    [xReduced, yReduced] = reduceAverageLineForPlot( ...
        xAxisUs, avg_for_lvl(levelIdx, :), ...
        options.avgLineMaxPoints, options.avgLineSmoothWindow);
    plot(ax, xReduced, yReduced, ...
        "LineWidth", 2, ...
        "Color", cols(levelIdx, :));
end

% refValues = unique(numericSignal(refSymbols));
% for refIdx = 1:numel(refValues)
%     yline(ax, refValues(refIdx), "--", ...
%         "Color", [0.35 0.35 0.35], ...
%         "LineWidth", 0.75, ...
%         "HandleVisibility", "off");
% end

if options.showStdAnnotations
    annotateLevelStd(symbols_for_lvl, avg_for_lvl, xAxisUs, options.xLimits);
end

xlabel(ax, 'Time in $\mu$s');
ylabel(ax, 'Normalized Amplitude');
xlim(ax, xLimits);
ylim(ax, options.yLimits);
grid(ax, "on");
figure(figHandle);
drawnow;
end

function [xReduced, yReduced] = reduceAverageLineForPlot(xAxisUs, yTrace, maxPoints, smoothWindow)
valid = isfinite(xAxisUs) & isfinite(yTrace);
xValid = xAxisUs(valid);
yValid = yTrace(valid);

if isempty(xValid)
    xReduced = [];
    yReduced = [];
    return
end

[xValid, uniqueIdx] = unique(xValid, "stable");
yValid = yValid(uniqueIdx);

if smoothWindow > 1
    yValid = smoothdata(yValid, "movmean", smoothWindow);
end

numOut = min(maxPoints, numel(xValid));
if numOut >= numel(xValid)
    xReduced = xValid;
    yReduced = yValid;
    return
end

xReduced = linspace(xValid(1), xValid(end), numOut);
yReduced = interp1(xValid, yValid, xReduced, "pchip");
end

function xLimits = getXLimits(configuredLimits, xAxisUs, symbols_for_lvl)
if ~isempty(configuredLimits)
    xLimits = configuredLimits;
    return
end

filledColumns = any(~isnan(symbols_for_lvl), 1);
if ~any(filledColumns)
    xLimits = [xAxisUs(1), xAxisUs(end)];
    return
end

lastFilledColumn = find(filledColumns, 1, "last");
xMax = xAxisUs(lastFilledColumn);
xLimits = [0, 1*xMax];
end

function annotateLevelStd(symbols_for_lvl, avg_for_lvl, xAxisUs, configuredXLimits)
xLimits = getXLimits(configuredXLimits, xAxisUs, symbols_for_lvl);
xText = xLimits(1) + 0.2*diff(xLimits);

for levelIdx = 1:size(symbols_for_lvl, 1)
    levelSamples = symbols_for_lvl(levelIdx, :);
    levelStd = std(levelSamples, 0, 2, 'omitnan');
    if isnan(levelStd)
        continue
    end

    levelAverage = avg_for_lvl(levelIdx, :);
    yText = median(levelAverage(~isnan(levelAverage)), 'omitnan');
    if isnan(yText)
        yText = median(levelSamples(~isnan(levelSamples)), 'omitnan');
    end

    label = ['$\sigma^2_',sprintf('%d', levelIdx),'$ = ', sprintf('%.3f', levelStd)];
    text(xText, yText, label, ...
        'Interpreter', 'latex', ...
        'HorizontalAlignment', 'right', ...
        'VerticalAlignment', 'middle', ...
        'FontSize', 10, ...
        'BackgroundColor', 'w', ...
        'Margin', 2, ...
        'EdgeColor', [0.8 0.8 0.8]);
end
end

function levelAverage = interpolateMissingLevelAverage(levelAverage)
validSamples = ~isnan(levelAverage);

if nnz(validSamples) == 0
    return
elseif nnz(validSamples) == 1
    levelAverage(:) = levelAverage(validSamples);
    return
end

t = 1:numel(levelAverage);
levelAverage(~validSamples) = interp1(t(validSamples), levelAverage(validSamples), ...
    t(~validSamples), 'linear', 'extrap');
end

function values = numericSignal(signalLike)
if isa(signalLike, "Signal")
    values = signalLike.signal;
else
    values = signalLike;
end

values = values(:).';
end

function values = normalizeNumericRms(values)
values = values ./ sqrt(mean(values.^2, "omitnan"));
end