%% Main script: Sweep different parameter pairs
% Default parameters
defaults.mu1     = 0.5;
defaults.mu2     = 1.0;
defaults.sigma1  = 0.1;
defaults.sigma2  = 0.1;
defaults.weight1 = 0.5;

% Parameter pair definitions
%{
param_pairs = {
    'mu1',     linspace(0.7, 1.0, 40), ...
    'mu2',     linspace(1.0, 1.3, 40);
    
    'mu1',     linspace(0.7, 1.0, 40), ...
    'weight1', linspace(0.2, 0.8, 40);
    
    'sigma1',  linspace(0.05, 0.2, 40), ...
    'sigma2',  linspace(0.05, 0.2, 40);
    
    'mu1',     linspace(0.7, 1.0, 40), ...
    'sigma1',  linspace(0.05, 0.2, 40);
    
    'mu2',     linspace(1.0, 1.3, 40), ...
    'weight1', linspace(0.2, 0.8, 40);
};
%}
param_pairs = {
    'mu1',     linspace(0.1, 1.5, 40), ...
    'mu2',     linspace(0.1, 1.5, 40);
};

% Cumulant index to visualize (2=variance, 3=skewness, 4=kurtosis)
cumulant_to_plot = 4;

% Run sweep for each pair
for i = 1:size(param_pairs,1)
    param1_name = param_pairs{i,1};
    param1_vals = param_pairs{i,2};
    param2_name = param_pairs{i,3};
    param2_vals = param_pairs{i,4};
    
    fprintf('Sweeping %s and %s...\n', param1_name, param2_name);
    
    Z = sweepBimodalCumulants(param1_name, param1_vals, ...
                          param2_name, param2_vals, ...
                          defaults, cumulant_to_plot);
end

%%
% Parameters
N_total = 10000;
mu1     = 0.5; 
sigma1  = 0.1;
mu2     = 1.0; 
sigma2  = 0.1;
weight1 = 0.7;

% Generate data
data    = generateBimodalDistribution(N_total, mu1, mu2, sigma1, sigma2, weight1);

% Plot histogram
figure(3);
clf
set(gcf, 'Color', 'w', 'Position', [100 100 950 750]);
histogram(data, 'Normalization', 'pdf', 'EdgeColor', 'none', 'FaceAlpha', 0.5);
hold on;

% Overlay smooth density estimate
[xi, f] = ksdensity(data);
plot(f, xi, 'r-', 'LineWidth', 2);

% Labels and title
xlabel('Value');
ylabel('Probability Density');
legend('Histogram', 'Smoothed Density');
grid on;

%%
N = size(Z, 1);

main_diag_values = diag(Z);
anti_diag_values = diag(flipud(Z));

param1_diag_main = param1_vals;
param2_diag_main = param2_vals;

param1_diag_anti = param1_vals;
param2_diag_anti = flip(param2_vals);

% For example, plot the main diagonal cumulants:
figure(4); 
set(gcf, 'Color', 'w', 'Position', [100 100 950 750]);
plot(1:N, main_diag_values, '-o');
xlabel('Index along diagonal');
ylabel('$\kappa_4$', 'Interpreter', 'latex');
title('$\kappa_4$ along anti-diagonal', 'Interpreter', 'latex');

% Plot anti-diagonal cumulants:
figure(5); 
set(gcf, 'Color', 'w', 'Position', [100 100 950 750]);
plot(1:N, anti_diag_values, '-o');
xlabel('Index along anti-diagonal');
ylabel('$\kappa_4$', 'Interpreter', 'latex');
title('$\kappa_4$ along anti-diagonal', 'Interpreter', 'latex');

%% === Helper: Bimodal Distribution ===
function data = generateBimodalDistribution(N_total, mu1, mu2, sigma1, sigma2, weight1)
%GENERATEBIMODALDISTRIBUTION Generates a single bimodal distribution.
%
%   data = generateBimodalDistribution(N_total, mu1, mu2, sigma1, sigma2, weight1)
%
%   Inputs:
%       N_total  - total number of samples
%       mu1, mu2 - means of the two modes
%       sigma1, sigma2 - standard deviations of the two modes
%       weight1  - fraction of samples from mode 1 (between 0 and 1)
%
%   Output:
%       data     - shuffled samples from the bimodal distribution

    % Validate weight
    weight1 = min(max(weight1, 0), 1);
    weight2 = 1 - weight1;

    % Determine number of samples for each mode
    N1 = round(N_total * weight1);
    N2 = N_total - N1;

    % Generate samples
    mode1_samples = normrnd(mu1, sigma1, [N1, 1]);
    mode2_samples = normrnd(mu2, sigma2, [N2, 1]);

    % Combine and shuffle
    data = [mode1_samples; mode2_samples];
    data = data(randperm(length(data)));
end

%% === Helper: Cumulant Calculation ===
function kappa = computeCumulants(data, max_order)
    data = data(:);
    mu   = mean(data);
    centered = data - mu;
    
    % Preallocate
    c = zeros(1, max_order);
    kappa = zeros(1, max_order);
    
    % Compute central moments up to max_order
    for n = 1:max_order
        c(n) = mean(centered.^n);
    end

    % Assign cumulants based on available order
    if max_order >= 1, kappa(1) = mu; end
    if max_order >= 2, kappa(2) = c(2); end
    if max_order >= 3, kappa(3) = c(3); end
    if max_order >= 4, kappa(4) = c(4) - 3*c(2)^2; end
    if max_order >= 5, kappa(5) = c(5) - 10*c(3)*c(2); end
    if max_order >= 6
        kappa(6) = c(6) - 15*c(4)*c(2) - 10*c(3)^2 + 30*c(2)^3;
    end
end

%% === Helper: Cumulant Calculation ===
function Z = sweepBimodalCumulants(param1_name, param1_vals, ...
                                   param2_name, param2_vals, ...
                                   fixed_params, ...
                                   cumulant_index)
%SWEEPBIMODALCUMULANTS Sweep 2 parameters and return a chosen cumulant.
%
%   Z = sweepBimodalCumulants(...)
%   Returns a matrix Z of cumulant values at each grid point.

    % Setup grid
    [P1, P2]  = meshgrid(param1_vals, param2_vals);
    Z         = zeros(size(P1));

    N_samples = 1000;
    maxOrder  = max(4, cumulant_index);

    for i = 1:numel(P1)
        % Copy fixed parameters
        params = fixed_params;

        % Override swept parameters
        params.(param1_name) = P1(i);
        params.(param2_name) = P2(i);

        % Generate and compute cumulants
        data  = generateBimodalDistribution(N_samples, ...
             params.mu1, params.mu2, ...
             params.sigma1, params.sigma2, ...
             params.weight1);

        kappa = computeCumulants(data, maxOrder);
        Z(i)  = kappa(cumulant_index);
    end

    % Plot full heatmap
    figure;
    set(gcf, 'Color', 'w', 'Position', [100 100 950 750]);
    imagesc(param1_vals, param2_vals, Z);
    set(gca, 'YDir', 'normal');
    xlabel(param1_name);
    ylabel(param2_name);
    title(['Cumulant \kappa_', num2str(cumulant_index)]);
    colorbar;
    axis tight;

    % Optional binary colormap (red = ≥0, blue = <0)
    figure;
    set(gcf, 'Color', 'w', 'Position', [100 100 950 750]);
    imagesc(param1_vals, param2_vals, Z);
    set(gca, 'YDir', 'normal');
    xlabel(param1_name);
    ylabel(param2_name);
    title(['Binary color split of \kappa_', num2str(cumulant_index)]);
    clim([-1 1]);
    colormap([0 0 1; 1 0 0]);  % Blue (neg), Red (pos & zero)
    colorbar;
    axis tight;
end