%% Export precomputed results for LNM Nullbook
% Run this after Example_RLdRL_null_derivation.m. The browser does not run
% MATLAB; this script precomputes every supported RL/dRL output.

parcelAxis = unique(round(linspace(1, nParcels, 26)));
propValues = [0.05 0.10 0.15 0.20 0.25 0.30];
% propValues = [0.20];

% Create one self-contained output folder per disease/lesion set.
diseaseSlug = lower(regexprep(char(lesionSet), '[^a-zA-Z0-9]+', '-'));
diseaseSlug = regexprep(diseaseSlug, '(^-|-$)', '');
assert(~isempty(diseaseSlug), 'lesionSet must contain a valid disease name.');

webRoot = fileparts(mfilename('fullpath'));
datasetDir = fullfile(webRoot, 'datasets', diseaseSlug);
if ~isfolder(datasetDir)
    mkdir(datasetDir);
end

if exist('cash_lnm_RL_z', 'var')
    lnm_RL_z = cash_lnm_RL_z;
end
assert(exist('lnm_RL_z', 'var') == 1, ...
    'Expected lnm_RL_z from the RL null-model calculation.');

if exist('cm', 'var')
    brainCmap = cm;
else
    brainCmap = parula(256);
end

demo = struct();
demo.meta = struct( ...
    'title', 'LNM Nullbook', ...
    'dataset', lesionSet, ...
    'nParcels', nParcels, ...
    'nInputFiles', nFiles, ...
    'nLesions', size(M, 1), ...
    'nPermutations', nPerms, ...
    'defaultProp', 0.2, ...
    'availableProps', propValues);

demo.parcelAxis = parcelAxis;
demo.observedLnm = lnm_orig(parcelAxis);
demo.rlPermutationZ = lnm_RL_z(parcelAxis);
demo.rlAnalyticZ = Lz_RL(parcelAxis);
demo.rlPermutationZFull = lnm_RL_z(:)';
demo.rlAnalyticZFull = Lz_RL(:)';

M_norm = M ./ sum(M, 2);
mu_c_export = mean(C);
Cn_centered = C - mu_c_export;
m = size(M, 1);
Lz_RL_centered = (1 / sqrt(m)) * sum(M_norm * Cn_centered, 1);
demo.rlAnalyticCenteredZ = Lz_RL_centered(parcelAxis);
demo.rlAnalyticCenteredZFull = Lz_RL_centered(:)';

%% Fit regression model from properties of C
global_degree = sum(C(55:end,:));
subcortical_degree = sum(C(1:54,:));
rsn_modules = modularity_und(C);
unique_modules = unique(rsn_modules);
rsn_degree = zeros(length(unique_modules), size(C, 2));
for ii = 1:length(unique_modules)
    rsn_degree(ii, :) = mean( ...
        C(rsn_modules == unique_modules(ii), :), 1);
end
[coeff, ~, ~, ~, explained_variance] = pca(C);
k = 3;
gradients = coeff(:, 1:k)';
X = [global_degree; subcortical_degree; rsn_degree; gradients]';
X_with_intercept = [ones(size(X, 1), 1), X];

targetMap = lnm_RL_z(:);
stats = regstats(targetMap, X, 'linear', {'rsquare', 'beta'});
Y_pred = X_with_intercept * stats.beta;
demo.regression.rl = struct( ...
    'rSquared', stats.rsquare, ...
    'predictedZ', Y_pred(parcelAxis)', ...
    'predictedZFull', Y_pred(:)');

observedBrain = writeBrainFigure( ...
    zscore(lnm_orig(:))', ...
    'Observed LNM', ...
    'observed', ...
    diseaseSlug, datasetDir, projectDir, brainCmap);
rlPermutationBrain = writeBrainFigure( ...
    zscore(lnm_RL_z(:))', ...
    'Permutation RL', ...
    'rl-permutation', ...
    diseaseSlug, datasetDir, projectDir, brainCmap);
rlAnalyticBrain = writeBrainFigure( ...
    zscore(Lz_RL(:))', ...
    'Analytic M x Cn', ...
    'rl-analytic', ...
    diseaseSlug, datasetDir, projectDir, brainCmap);
rlAnalyticCenteredBrain = writeBrainFigure( ...
    zscore(Lz_RL_centered(:))', ...
    'Analytic M x centered Cn', ...
    'rl-analytic-centered', ...
    diseaseSlug, datasetDir, projectDir, brainCmap);
demo.regression.rl.predictedBrain = writeBrainFigure( ...
    zscore(Y_pred(:))', ...
    'Standard patterns of C', ...
    'rl-regression-predicted', ...
    diseaseSlug, datasetDir, projectDir, brainCmap);
demo.brainFigures = struct( ...
    'observed', observedBrain, ...
    'rl', struct( ...
        'permutation', rlPermutationBrain, ...
        'analytic', rlAnalyticBrain, ...
        'analyticCentered', rlAnalyticCenteredBrain));

demo.agreement = struct( ...
    'rl', corr(lnm_RL_z(:), Lz_RL(:)), ...
    'rlCentered', corr(lnm_RL_z(:), Lz_RL_centered(:)));

%% Precompute every supported dRL Prop value
s = mean(C);
d = abs(repmat(s,nParcels,1)' - repmat(s,nParcels,1));
[~, fullStrengthOrder] = sort(d, 2);

for iProp = 1:numel(propValues)
    Prop = propValues(iProp);
    propKey = sprintf('p%03d', round(Prop * 100));
    fprintf('Precomputing dRL %s: Prop = %.2f\n', lesionSet, Prop);

    % Permutation-based dRL null model.
    [M_configs, M_configs_null] = zalesky_cash_lesionMatrix_permutation( ...
        C, M, 'Perms', nPerms, 'Prop', Prop);

    realLNM_samp = zeros(nParcels, nPerms);
    for iPerm = 1:nPerms
        tmpM = squeeze(M_configs(iPerm,:,:));
        realLNM_samp(:,iPerm) = nansum((tmpM./sum(tmpM,2)) * C);
    end
    lnm_obs_prop = mean(realLNM_samp,2)';

    lnm_dist_null = zeros(nParcels, nPerms);
    for iPerm = 1:nPerms
        M_perm = squeeze(M_configs_null(iPerm, :, :));
        M_perm = M_perm ./ sum(M_perm, 2);
        lnm_dist_null(:, iPerm) = sum(M_perm * C, 1);
    end
    lnm_dRL_z_prop = (lnm_obs_prop - mean(lnm_dist_null, 2)') ./ ...
        std(lnm_dist_null, 0, 2)';

    targetMap = lnm_dRL_z_prop(:);
    stats = regstats(targetMap, X, 'linear', {'rsquare', 'beta'});
    Y_pred = X_with_intercept * stats.beta;

    % Analytic dRL using the locally degree-centered connectome.
    ind_srt = fullStrengthOrder(:, 2:ceil(Prop*nParcels));
    mu_dc = zeros(nParcels, nParcels);
    sigma_dc = zeros(nParcels, nParcels);
    for iParcel = 1:nParcels
        degree_matched_rows = ind_srt(iParcel, :);
        mu_dc(iParcel, :) = mean(C(degree_matched_rows, :), 1);
        sigma_dc(iParcel, :) = std(C(degree_matched_rows, :), 0, 1);
    end
    Cdc_centered = C - mu_dc;
    Cdc = Cdc_centered ./ sigma_dc;
    m = size(M,1);
    Lz_dRL_centered_prop = (1 / sqrt(m)) * sum(M_norm * Cdc_centered, 1);
    Lz_dRL_prop = (1 / sqrt(m)) * sum(M_norm * Cdc, 1);

    propTag = sprintf('p%03d', round(Prop * 100));
    permutationKind = sprintf('drl-permutation-%s', propTag);
    analyticKind = sprintf('drl-analytic-%s', propTag);
    analyticCenteredKind = sprintf('drl-analytic-centered-%s', propTag);
    permutationLabel = sprintf('Permutation dRL (Prop = %.2f)', Prop);
    analyticLabel = sprintf('Analytic dRL (Prop = %.2f)', Prop);
    analyticCenteredLabel = sprintf('Analytic centered dRL (Prop = %.2f)', Prop);

    entry = struct();
    entry.prop = Prop;
    entry.permutationZ = lnm_dRL_z_prop(parcelAxis);
    entry.analyticZ = Lz_dRL_prop(parcelAxis);
    entry.analyticCenteredZ = Lz_dRL_centered_prop(parcelAxis);
    entry.permutationZFull = lnm_dRL_z_prop(:)';
    entry.analyticZFull = Lz_dRL_prop(:)';
    entry.analyticCenteredZFull = Lz_dRL_centered_prop(:)';
    entry.agreement = corr(lnm_dRL_z_prop(:), Lz_dRL_prop(:));
    entry.agreementCentered = corr(lnm_dRL_z_prop(:), Lz_dRL_centered_prop(:));
    entry.maxAbsDifference = max(abs(lnm_dRL_z_prop(:) - Lz_dRL_prop(:)));
    entry.regression = struct( ...
        'rSquared', stats.rsquare, ...
        'predictedZ', Y_pred(parcelAxis)', ...
        'predictedZFull', Y_pred(:)', ...
        'predictedBrain', writeBrainFigure( ...
            zscore(Y_pred(:))', ...
            'Standard patterns of C', ...
            sprintf('drl-regression-predicted-%s', propTag), ...
            diseaseSlug, datasetDir, projectDir, brainCmap));
    entry.brainFigures = struct( ...
        'permutation', writeBrainFigure( ...
            zscore(lnm_dRL_z_prop(:))', permutationLabel, ...
            permutationKind, diseaseSlug, datasetDir, projectDir, brainCmap), ...
        'analytic', writeBrainFigure( ...
            zscore(Lz_dRL_prop(:))', analyticLabel, ...
            analyticKind, diseaseSlug, datasetDir, projectDir, brainCmap), ...
        'analyticCentered', writeBrainFigure( ...
            zscore(Lz_dRL_centered_prop(:))', analyticCenteredLabel, ...
            analyticCenteredKind, diseaseSlug, datasetDir, projectDir, brainCmap));

    demo.drlByProp.(propKey) = entry;

    clear M_configs M_configs_null realLNM_samp lnm_dist_null targetMap stats Y_pred
    clear mu_dc sigma_dc Cdc Cdc_centered lnm_dRL_z_prop
    clear Lz_dRL_prop Lz_dRL_centered_prop
end

% Keep the Prop=0.20 fields at the top level for older app snapshots.
defaultDrl = demo.drlByProp.p020;
demo.drlPermutationZ = defaultDrl.permutationZ;
demo.drlAnalyticZ = defaultDrl.analyticZ;
demo.drlAnalyticCenteredZ = defaultDrl.analyticCenteredZ;
demo.drlPermutationZFull = defaultDrl.permutationZFull;
demo.drlAnalyticZFull = defaultDrl.analyticZFull;
demo.drlAnalyticCenteredZFull = defaultDrl.analyticCenteredZFull;
demo.brainFigures.drl = defaultDrl.brainFigures;
demo.regression.drl = defaultDrl.regression;
demo.agreement.drl = defaultDrl.agreement;
demo.agreement.drlCentered = defaultDrl.agreementCentered;
demo.agreement.maxAbsDifference = defaultDrl.maxAbsDifference;

%% Write disease JSON
jsonText = jsonencode(demo, PrettyPrint=true);
outputName = sprintf('lnm-nullbook-data-%s.json', diseaseSlug);
outputFile = fullfile(datasetDir, outputName);
fid = fopen(outputFile, 'w');
assert(fid ~= -1, 'Could not open output JSON for writing.');
cleanup = onCleanup(@() fclose(fid));
fprintf(fid, '%s', jsonText);

%% Write JavaScript data wrapper for direct file opening
jsOutputName = sprintf('lnm-nullbook-data-%s.js', diseaseSlug);
jsOutputFile = fullfile(datasetDir, jsOutputName);
fidJs = fopen(jsOutputFile, 'w');
assert(fidJs ~= -1, 'Could not open output JavaScript file for writing.');
cleanupJs = onCleanup(@() fclose(fidJs));
fprintf(fidJs, 'window.__LNM_NULLBOOK_DATA__ = %s;\n', jsonText);

fprintf('Wrote dataset folder:\n  %s\n', datasetDir);
fprintf('Wrote JSON:\n  %s\n', outputFile);
fprintf('Wrote JavaScript data wrapper:\n  %s\n', jsOutputFile);
fprintf('Precomputed Prop values: 0.05, 0.10, 0.15, 0.20, 0.25, 0.30\n');

function relativePath = writeBrainFigure(values, label, kind, ...
        diseaseSlug, datasetDir, projectDir, brainCmap)
    brainName = sprintf('brain-%s-%s.png', kind, diseaseSlug);
    brainFile = fullfile(datasetDir, brainName);

    fig = figure( ...
        'Visible', 'off', ...
        'Color', 'w', ...
        'Position', [100 100 720 440]);
    ax = axes(fig);
    plot_surface(values, ...
        'hemi', 'rh', ...
        'vmin', -2.6, ...
        'vmax', 2.6, ...
        'vcenter', 0, ...
        'cmap', brainCmap, ...
        'projectDir', projectDir, ...
        'cbar', false, ...
        'show_atlas_parcels', false, ...
        'ax', ax);
    title(ax, label, 'FontWeight', 'normal');
    exportgraphics(fig, brainFile, ...
        'Resolution', 180, ...
        'BackgroundColor', 'white');
    close(fig);

    relativePath = sprintf('datasets/%s/%s', diseaseSlug, brainName);
end
