% Diffusion Model Demo - 2D Visualization
% Shows forward diffusion (adding noise) and reverse process (denoising)

clear; close all;
rng(42);

%% Generate 2D training data - two moons shape
n_samples = 1000;
t_vals = linspace(0, pi, n_samples/2);
moon1 = [cos(t_vals); sin(t_vals)]' * 2;
moon2 = [1-cos(t_vals); 1-sin(t_vals)-0.5]' * 2;
X_data = [moon1; moon2] + randn(n_samples, 2) * 0.1;

%% Diffusion parameters
T = 100;  % Number of diffusion steps (increased for smoother transitions)
% Improved cosine noise schedule with clipping for numerical stability
s = 0.008;  % Small offset
max_beta = 0.999;  % Prevent alpha from going to zero
t_vals = linspace(0, T, T+1);
alpha_bar = cos(((t_vals / T) + s) / (1 + s) * pi / 2).^2;
alpha_bar = alpha_bar(2:end) / alpha_bar(1);  % Normalize
alpha_bar = max(alpha_bar, 1 - max_beta);  % Clip to prevent alpha_bar → 0
alpha = [alpha_bar(1), alpha_bar(2:end) ./ alpha_bar(1:end-1)];
alpha = max(alpha, 1e-6);  % Prevent alpha from being exactly zero
beta = 1 - alpha;

%% Visualize Forward Diffusion Process
fprintf('=== FORWARD DIFFUSION (Adding Noise) ===\n');
figure('Position', [100, 100, 1200, 400]);

% Sample one data point to track
x0 = X_data(1, :);
steps_to_show = [1, 20, 50, 100];

for i = 1:length(steps_to_show)
    t = steps_to_show(i);

    % Forward diffusion: x_t = sqrt(alpha_bar_t) * x_0 + sqrt(1-alpha_bar_t) * noise
    noise = randn(n_samples, 2);
    X_t = sqrt(alpha_bar(t)) * X_data + sqrt(1 - alpha_bar(t)) * noise;

    subplot(1, 4, i);
    scatter(X_t(:,1), X_t(:,2), 10, 'b', 'filled', 'MarkerFaceAlpha', 0.3);
    title(sprintf('Step %d/%d', t, T), 'FontSize', 12, 'FontWeight', 'bold');
    xlabel('x_1'); ylabel('x_2');
    xlim([-6, 6]); ylim([-6, 6]);
    grid on; axis equal;
    set(gca, 'FontSize', 10);
end

sgtitle('Forward Diffusion: Gradually Adding Noise', 'FontSize', 14, 'FontWeight', 'bold');
drawnow;
pause(1);

%% Train Improved Denoising Network
fprintf('\n=== TRAINING DENOISING NETWORK ===\n');

% Sinusoidal timestep embedding (like in Transformers)
embed_dim = 32;  % Increased embedding dimension

% Deeper 6-layer neural network to predict noise
input_dim = 2 + embed_dim;  % x_t (2D) + timestep embedding (32D)
hidden_dim = 512;  % Wider network

% Initialize network parameters with Xavier/He initialization
scale1 = sqrt(2.0 / input_dim);
scale2 = sqrt(2.0 / hidden_dim);

W1 = randn(input_dim, hidden_dim) * scale1;
b1 = zeros(1, hidden_dim);
W2 = randn(hidden_dim, hidden_dim) * scale2;
b2 = zeros(1, hidden_dim);
W3 = randn(hidden_dim, hidden_dim) * scale2;
b3 = zeros(1, hidden_dim);
W4 = randn(hidden_dim, hidden_dim) * scale2;
b4 = zeros(1, hidden_dim);
W5 = randn(hidden_dim, hidden_dim) * scale2;
b5 = zeros(1, hidden_dim);
W6 = randn(hidden_dim, 2) * scale2;
b6 = zeros(1, 2);

% Initialize momentum buffers for better optimization
v_W1 = zeros(size(W1)); v_b1 = zeros(size(b1));
v_W2 = zeros(size(W2)); v_b2 = zeros(size(b2));
v_W3 = zeros(size(W3)); v_b3 = zeros(size(b3));
v_W4 = zeros(size(W4)); v_b4 = zeros(size(b4));
v_W5 = zeros(size(W5)); v_b5 = zeros(size(b5));
v_W6 = zeros(size(W6)); v_b6 = zeros(size(b6));

% Training loop
n_epochs = 5000;
batch_size = 256;  % Larger batch for stability
learning_rate = 0.0005;  % Lower learning rate for stability
momentum = 0.9;

for epoch = 1:n_epochs
    % Sample batch
    idx = randperm(n_samples, batch_size);
    x_batch = X_data(idx, :);

    % Sample random timestep for each sample
    t_batch = randi([1, T], batch_size, 1);

    % Add noise according to diffusion schedule
    noise_true = randn(batch_size, 2);
    x_noisy = zeros(batch_size, 2);
    for i = 1:batch_size
        t = t_batch(i);
        x_noisy(i,:) = sqrt(alpha_bar(t)) * x_batch(i,:) + sqrt(1-alpha_bar(t)) * noise_true(i,:);
    end

    % Forward pass - predict noise with timestep embedding (6 layers)
    t_embed = timestep_embedding(t_batch, T, embed_dim);
    h1 = max(0, [x_noisy, t_embed] * W1 + b1);  % ReLU
    h2 = max(0, h1 * W2 + b2);  % ReLU
    h3 = max(0, h2 * W3 + b3);  % ReLU
    h4 = max(0, h3 * W4 + b4);  % ReLU
    h5 = max(0, h4 * W5 + b5);  % ReLU
    noise_pred = h5 * W6 + b6;

    % Compute loss (MSE)
    loss = mean(sum((noise_pred - noise_true).^2, 2));

    % Backward pass with momentum
    grad_output = 2 * (noise_pred - noise_true) / batch_size;

    grad_W6 = h5' * grad_output;
    grad_b6 = sum(grad_output, 1);

    grad_h5 = grad_output * W6';
    grad_h5(h5 <= 0) = 0;  % ReLU gradient

    grad_W5 = h4' * grad_h5;
    grad_b5 = sum(grad_h5, 1);

    grad_h4 = grad_h5 * W5';
    grad_h4(h4 <= 0) = 0;  % ReLU gradient

    grad_W4 = h3' * grad_h4;
    grad_b4 = sum(grad_h4, 1);

    grad_h3 = grad_h4 * W4';
    grad_h3(h3 <= 0) = 0;  % ReLU gradient

    grad_W3 = h2' * grad_h3;
    grad_b3 = sum(grad_h3, 1);

    grad_h2 = grad_h3 * W3';
    grad_h2(h2 <= 0) = 0;  % ReLU gradient

    grad_W2 = h1' * grad_h2;
    grad_b2 = sum(grad_h2, 1);

    grad_h1 = grad_h2 * W2';
    grad_h1(h1 <= 0) = 0;  % ReLU gradient

    grad_W1 = [x_noisy, t_embed]' * grad_h1;
    grad_b1 = sum(grad_h1, 1);

    % Update parameters with momentum
    v_W1 = momentum * v_W1 + learning_rate * grad_W1;
    v_b1 = momentum * v_b1 + learning_rate * grad_b1;
    v_W2 = momentum * v_W2 + learning_rate * grad_W2;
    v_b2 = momentum * v_b2 + learning_rate * grad_b2;
    v_W3 = momentum * v_W3 + learning_rate * grad_W3;
    v_b3 = momentum * v_b3 + learning_rate * grad_b3;
    v_W4 = momentum * v_W4 + learning_rate * grad_W4;
    v_b4 = momentum * v_b4 + learning_rate * grad_b4;
    v_W5 = momentum * v_W5 + learning_rate * grad_W5;
    v_b5 = momentum * v_b5 + learning_rate * grad_b5;
    v_W6 = momentum * v_W6 + learning_rate * grad_W6;
    v_b6 = momentum * v_b6 + learning_rate * grad_b6;

    W1 = W1 - v_W1; b1 = b1 - v_b1;
    W2 = W2 - v_W2; b2 = b2 - v_b2;
    W3 = W3 - v_W3; b3 = b3 - v_b3;
    W4 = W4 - v_W4; b4 = b4 - v_b4;
    W5 = W5 - v_W5; b5 = b5 - v_b5;
    W6 = W6 - v_W6; b6 = b6 - v_b6;

    if mod(epoch, 500) == 0
        fprintf('Epoch %d/%d, Loss: %.4f\n', epoch, n_epochs, loss);
    end
end

fprintf('Training complete!\n');

%% Reverse Process - Generate New Samples
fprintf('\n=== REVERSE PROCESS (Denoising) ===\n');

n_gen = 500;
% Start from noise scaled appropriately
noise_scale = sqrt((1 - alpha_bar(end)) / alpha_bar(end));
noise_scale = min(noise_scale, 5);  % Clip to prevent extreme initial values
x_t = randn(n_gen, 2) * noise_scale;

figure('Position', [100, 100, 1200, 400]);

steps_to_show = [100, 50, 20, 1];
subplot_idx = 1;

for t = T:-1:1
    % Predict noise using trained network with timestep embedding (6 layers)
    t_embed = timestep_embedding(ones(n_gen, 1) * t, T, embed_dim);
    h1 = max(0, [x_t, t_embed] * W1 + b1);
    h2 = max(0, h1 * W2 + b2);
    h3 = max(0, h2 * W3 + b3);
    h4 = max(0, h3 * W4 + b4);
    h5 = max(0, h4 * W5 + b5);
    noise_pred = h5 * W6 + b6;

    % Remove predicted noise with numerical stability
    coeff = (1-alpha(t)) / sqrt(1-alpha_bar(t));
    coeff = min(coeff, 10);  % Clip coefficient to prevent explosion
    x_t = (x_t - coeff * noise_pred) / sqrt(alpha(t));

    % Clip samples to reasonable range to prevent divergence
    x_t = max(min(x_t, 20), -20);

    % Add small noise (except at last step)
    if t > 1
        x_t = x_t + sqrt(beta(t)) * randn(n_gen, 2);
    end

    % Visualize at key steps
    if any(t == steps_to_show)
        subplot(1, 4, subplot_idx);
        scatter(x_t(:,1), x_t(:,2), 10, 'r', 'filled', 'MarkerFaceAlpha', 0.3);
        title(sprintf('Step %d/%d', t, T), 'FontSize', 12, 'FontWeight', 'bold');
        xlabel('x_1'); ylabel('x_2');
        xlim([-6, 6]); ylim([-6, 6]);
        grid on; axis equal;
        set(gca, 'FontSize', 10);
        subplot_idx = subplot_idx + 1;
        drawnow;
    end
end

sgtitle('Reverse Process: Gradually Removing Noise (Generation)', 'FontSize', 14, 'FontWeight', 'bold');

%% Final Comparison
figure('Position', [100, 100, 1000, 400]);

subplot(1, 2, 1);
scatter(X_data(:,1), X_data(:,2), 20, 'b', 'filled', 'MarkerFaceAlpha', 0.5);
title('Original Training Data', 'FontSize', 12, 'FontWeight', 'bold');
xlabel('x_1'); ylabel('x_2');
xlim([-6, 6]); ylim([-6, 6]);
grid on; axis equal;
set(gca, 'FontSize', 10);

subplot(1, 2, 2);
scatter(x_t(:,1), x_t(:,2), 20, 'r', 'filled', 'MarkerFaceAlpha', 0.5);
title('Generated Samples (from noise)', 'FontSize', 12, 'FontWeight', 'bold');
xlabel('x_1'); ylabel('x_2');
xlim([-6, 6]); ylim([-6, 6]);
grid on; axis equal;
set(gca, 'FontSize', 10);

sgtitle('Diffusion Model Results', 'FontSize', 14, 'FontWeight', 'bold');

% Save figure
saveas(gcf, 'diffusion_2d_demo.png');
fprintf('\nSaved output to diffusion_2d_demo.png\n');

%% Summary
fprintf('\n=== SUMMARY ===\n');
fprintf('Diffusion Model Components:\n');
fprintf('  1. Forward process: Add noise over %d steps with COSINE schedule\n', T);
fprintf('  2. Neural network: 6-layer (512 units) with 32D timestep embedding\n');
fprintf('  3. Training: %d epochs with momentum optimizer\n', n_epochs);
fprintf('  4. Reverse process: Start from noise, iteratively denoise\n');
fprintf('  5. Result: Generated %d new samples matching training distribution\n', n_gen);
fprintf('\nKey Improvements:\n');
fprintf('  - Much deeper network (6 layers, 512 units)\n');
fprintf('  - Cosine noise schedule (better than linear)\n');
fprintf('  - Momentum optimization (faster convergence)\n');
fprintf('  - More diffusion steps (100 vs 50)\n');
fprintf('  - Extensive training (%d epochs)\n', n_epochs);
fprintf('\nKey Insight: The model learns to reverse the noising process,\n');
fprintf('             enabling generation from pure random noise!\n');

%% Helper Function
function emb = timestep_embedding(t, T, dim)
    % Sinusoidal timestep embedding (like in Transformers/Diffusion models)
    % t: timestep (scalar or vector), T: total steps, dim: embedding dimension
    t_norm = t / T;  % Normalize to [0,1]
    half_dim = dim / 2;
    frequencies = exp(-log(10000) * (0:half_dim-1) / (half_dim - 1));
    args = t_norm(:) * frequencies;
    emb = [sin(args), cos(args)];
end
