% 1D Optimal Transport using Benamou-Brenier Dynamical Formulation
% Merging Two Gaussians into One with animation
%
% This demonstrates optimal transport from a bimodal distribution
% (two separate Gaussians) to a unimodal distribution (single Gaussian).
% Shows how mass from two sources optimally merges into one.

clear all; close all; clc;

fprintf('=== 1D Optimal Transport: Merging Two Gaussians ===\n');

%% Domain Setup
x_min = -5;
x_max = 5;
N = 800;  % Number of spatial grid points
x = linspace(x_min, x_max, N)';
dx = (x_max - x_min) / (N - 1);

fprintf('Domain: [%.1f, %.1f]\n', x_min, x_max);
fprintf('Grid points: %d\n', N);

%% Define Initial and Target Gaussian Distributions

% Initial: Two Gaussians (bimodal distribution)
% Left Gaussian
mu_L = -2.0;
sigma_L = 0.4;
weight_L = 0.5;  % Mass fraction
rho_L = weight_L / (sigma_L * sqrt(2 * pi)) * exp(-0.5 * ((x - mu_L) / sigma_L).^2);

% Right Gaussian
mu_R = 2.0;
sigma_R = 0.4;
weight_R = 0.5;  % Mass fraction
rho_R = weight_R / (sigma_R * sqrt(2 * pi)) * exp(-0.5 * ((x - mu_R) / sigma_R).^2);

% Combined initial distribution
rho_0 = rho_L + rho_R;
rho_0 = rho_0 / (sum(rho_0) * dx);  % Normalize to have mass 1

% Target: Single Gaussian in the center
mu_target = 0.0;
sigma_target = 0.8;
rho_1 = 1 / (sigma_target * sqrt(2 * pi)) * exp(-0.5 * ((x - mu_target) / sigma_target).^2);
rho_1 = rho_1 / (sum(rho_1) * dx);  % Normalize to have mass 1

fprintf('\nInitial: Two Gaussians\n');
fprintf('  Left:  mean = %.2f, std = %.2f, mass = %.2f\n', mu_L, sigma_L, weight_L);
fprintf('  Right: mean = %.2f, std = %.2f, mass = %.2f\n', mu_R, sigma_R, weight_R);
fprintf('Target: Single Gaussian\n');
fprintf('  Center: mean = %.2f, std = %.2f\n', mu_target, sigma_target);

%% Compute Optimal Transport Map using CDF method (exact for 1D)
% In 1D, the optimal transport map is T(x) = F_1^{-1}(F_0(x))
% where F_0, F_1 are the CDFs

% Compute CDFs using trapezoidal integration
F_0 = cumtrapz(x, rho_0);
F_1 = cumtrapz(x, rho_1);

% Normalize to [0,1]
F_0 = F_0 / max(F_0(end), 1e-10);
F_1 = F_1 / max(F_1(end), 1e-10);

% Ensure strict monotonicity by adding small increasing values
eps_mono = 1e-10;
for i = 2:N
    if F_0(i) <= F_0(i-1)
        F_0(i) = F_0(i-1) + eps_mono;
    end
    if F_1(i) <= F_1(i-1)
        F_1(i) = F_1(i-1) + eps_mono;
    end
end

% Final normalization
F_0 = F_0 / F_0(end);
F_1 = F_1 / F_1(end);

% Compute transport map T(x) = F_1^{-1}(F_0(x))
T = zeros(N, 1);
for i = 1:N
    % Find x such that F_1(x) = F_0(x_i)
    F_val = F_0(i);
    % Interpolate to find inverse
    if F_val <= F_1(1)
        T(i) = x(1);
    elseif F_val >= F_1(end)
        T(i) = x(end);
    else
        % Linear interpolation
        idx = find(F_1 >= F_val, 1, 'first');
        if idx > 1
            alpha = (F_val - F_1(idx-1)) / (F_1(idx) - F_1(idx-1));
            T(i) = x(idx-1) + alpha * (x(idx) - x(idx-1));
        else
            T(i) = x(idx);
        end
    end
end

% Compute displacement field
displacement = T - x;

fprintf('\nOptimal transport map computed\n');
fprintf('Max displacement: %.4f\n', max(abs(displacement)));

%% Compute Geodesic Path (McCann interpolation)
% ρ_t = ((1-t)Id + tT)_# ρ_0
% The geodesic is: T_t(x) = (1-t)x + t*T(x)

Nt = 100;  % Number of time steps
t_steps = linspace(0, 1, Nt);

% Store densities and velocities along the path
rho_t = zeros(N, Nt);
v_t = zeros(N, Nt);
kinetic_energy = zeros(Nt, 1);

fprintf('Computing geodesic path...\n');

for k = 1:Nt
    t = t_steps(k);

    % Interpolated map: T_t(x) = (1-t)*x + t*T(x)
    T_t = (1 - t) * x + t * T;

    % Compute pushforward density: ρ_t = T_t# ρ_0
    % This means we need to find ρ_t such that for any test function φ:
    % ∫ φ(y) ρ_t(y) dy = ∫ φ(T_t(x)) ρ_0(x) dx

    % Use binning method for pushforward
    rho_t(:, k) = compute_pushforward(x, rho_0, T_t, dx);

    % Compute velocity field: v(t,x) = dT_t/dt = T(x) - x
    % At position T_t(x), the velocity is (T(x) - x)
    v_at_source = T - x;

    % We need velocity at current positions (pushforward of velocity)
    v_t(:, k) = compute_pushforward_vector(x, rho_0, v_at_source, T_t, dx);

    % Compute kinetic energy: ∫ ρ_t(x) |v_t(x)|^2 dx
    kinetic_energy(k) = sum(rho_t(:, k) .* v_t(:, k).^2) * dx;
end

total_cost = trapz(t_steps, kinetic_energy);
fprintf('Total transport cost (Wasserstein-2^2 distance): %.6f\n', total_cost);

%% Note on Wasserstein distance
% For two Gaussians merging into one, there's no simple closed-form expression
% The computed cost represents the true optimal transport cost
fprintf('Note: For bimodal to unimodal transport, no simple theoretical formula exists.\n');

%% Visualization: Static plots

figure('Position', [100, 100, 1400, 900]);

% Plot 1: Initial and target densities
subplot(2, 3, 1);
plot(x, rho_0, 'b-', 'LineWidth', 2, 'DisplayName', '\rho_0 (two Gaussians)');
hold on;
plot(x, rho_1, 'r-', 'LineWidth', 2, 'DisplayName', '\rho_1 (merged)');
hold off;
grid on;
xlabel('x', 'FontSize', 11);
ylabel('Density', 'FontSize', 11);
title('Initial (Bimodal) and Target (Unimodal)', 'FontSize', 12);
legend('Location', 'best');
xlim([x_min, x_max]);

% Plot 2: Optimal transport map
subplot(2, 3, 2);
plot(x, T, 'b-', 'LineWidth', 2, 'DisplayName', 'T(x)');
hold on;
plot(x, x, 'k--', 'LineWidth', 1, 'DisplayName', 'Identity');
hold off;
grid on;
xlabel('x', 'FontSize', 11);
ylabel('T(x)', 'FontSize', 11);
title('Optimal Transport Map', 'FontSize', 12);
legend('Location', 'best');
xlim([x_min, x_max]);

% Plot 3: Displacement field
subplot(2, 3, 3);
plot(x, displacement, 'g-', 'LineWidth', 2);
grid on;
xlabel('x', 'FontSize', 11);
ylabel('T(x) - x', 'FontSize', 11);
title('Displacement Field', 'FontSize', 12);
xlim([x_min, x_max]);

% Plot 4: Space-time density evolution
subplot(2, 3, 4);
imagesc(t_steps, x, rho_t);
colormap(jet);
colorbar;
xlabel('Time t', 'FontSize', 11);
ylabel('Space x', 'FontSize', 11);
title('Density Evolution \rho(t,x)', 'FontSize', 12);
set(gca, 'YDir', 'normal');

% Plot 5: Velocity field evolution
subplot(2, 3, 5);
imagesc(t_steps, x, v_t);
colormap(gca, jet);
colorbar;
xlabel('Time t', 'FontSize', 11);
ylabel('Space x', 'FontSize', 11);
title('Velocity Field v(t,x)', 'FontSize', 12);
set(gca, 'YDir', 'normal');

% Plot 6: Kinetic energy evolution
subplot(2, 3, 6);
plot(t_steps, kinetic_energy, 'b-', 'LineWidth', 2);
hold on;
area(t_steps, kinetic_energy, 'FaceColor', 'b', 'FaceAlpha', 0.3);
hold off;
grid on;
xlabel('Time t', 'FontSize', 11);
ylabel('Kinetic Energy', 'FontSize', 11);
title(sprintf('Kinetic Energy (Total Cost = %.4f)', total_cost), 'FontSize', 12);
xlim([0, 1]);

sgtitle('1D Optimal Transport: Merging Two Gaussians into One', ...
        'FontSize', 14, 'FontWeight', 'bold');

%% Animation: Density Evolution

fprintf('\nCreating animation...\n');

fig_anim = figure('Position', [150, 150, 1400, 500]);

% Determine y-axis limits for consistent scale
max_density = max([max(rho_0), max(rho_1)]) * 1.1;
max_velocity = max(abs(v_t(:))) * 1.2;

for k = 1:Nt
    t = t_steps(k);

    % Density evolution
    subplot(1, 3, 1);
    plot(x, rho_0, 'b--', 'LineWidth', 1.5, 'DisplayName', '\rho_0 (two Gaussians)');
    hold on;
    plot(x, rho_1, 'r--', 'LineWidth', 1.5, 'DisplayName', '\rho_1 (merged)');
    plot(x, rho_t(:, k), 'k-', 'LineWidth', 2.5, 'DisplayName', sprintf('\\rho_t (t=%.2f)', t));
    hold off;
    grid on;
    xlabel('x', 'FontSize', 12);
    ylabel('\rho(t,x)', 'FontSize', 12);
    title(sprintf('Merging Density at t = %.3f', t), 'FontSize', 13);
    legend('Location', 'northwest');
    xlim([x_min, x_max]);
    ylim([0, max_density]);
    set(gca, 'FontSize', 11);

    % Velocity field
    subplot(1, 3, 2);
    plot(x, v_t(:, k), 'b-', 'LineWidth', 2);
    hold on;
    plot([x_min, x_max], [0, 0], 'k--', 'LineWidth', 1);

    % Add velocity vectors
    skip = 15;  % Plot every 15th vector
    quiver(x(1:skip:end), zeros(size(x(1:skip:end))), ...
           zeros(size(x(1:skip:end))), v_t(1:skip:end, k), ...
           0.5, 'r', 'LineWidth', 1.5, 'MaxHeadSize', 0.5);
    hold off;
    grid on;
    xlabel('x', 'FontSize', 12);
    ylabel('v(t,x)', 'FontSize', 12);
    title(sprintf('Velocity Field at t = %.3f', t), 'FontSize', 13);
    xlim([x_min, x_max]);
    ylim([-max_velocity, max_velocity]);
    set(gca, 'FontSize', 11);

    % Kinetic energy up to current time
    subplot(1, 3, 3);
    plot(t_steps(1:k), kinetic_energy(1:k), 'b-', 'LineWidth', 2);
    hold on;
    area(t_steps(1:k), kinetic_energy(1:k), 'FaceColor', 'b', 'FaceAlpha', 0.3);
    plot(t, kinetic_energy(k), 'ro', 'MarkerSize', 8, 'MarkerFaceColor', 'r');
    hold off;
    grid on;
    xlabel('Time t', 'FontSize', 12);
    ylabel('Kinetic Energy', 'FontSize', 12);

    % Compute cumulative cost (handle k=1 case)
    if k > 1
        cumulative_cost = trapz(t_steps(1:k), kinetic_energy(1:k));
    else
        cumulative_cost = 0;
    end

    title(sprintf('Kinetic Energy (Cumulative = %.4f)', cumulative_cost), ...
          'FontSize', 13);
    xlim([0, 1]);
    ylim([0, max(kinetic_energy) * 1.1]);
    set(gca, 'FontSize', 11);

    sgtitle(sprintf('Merging Two Gaussians: t = %.3f / 1.00', t), ...
            'FontSize', 14, 'FontWeight', 'bold');

    drawnow;

    % Pause for animation
    if k < Nt
        pause(0.05);
    end
end

fprintf('Animation complete!\n');

%% Additional Analysis: Particle Trajectories

fprintf('\nGenerating particle trajectories...\n');

figure('Position', [200, 200, 1200, 500]);

% Sample particles from initial distribution
n_particles = 30;
particle_positions_init = sample_from_density(x, rho_0, n_particles);

% Compute trajectories: x(t) = (1-t)*x_0 + t*T(x_0)
trajectories = zeros(n_particles, Nt);

for i = 1:n_particles
    x0 = particle_positions_init(i);

    % Find T(x0) by interpolation
    Tx0 = interp1(x, T, x0, 'linear', 'extrap');

    % Trajectory: x(t) = (1-t)*x0 + t*T(x0)
    trajectories(i, :) = (1 - t_steps) * x0 + t_steps * Tx0;
end

% Plot trajectories in space-time
subplot(1, 2, 1);
hold on;
for i = 1:n_particles
    plot(t_steps, trajectories(i, :), 'b-', 'LineWidth', 1.5);
end
hold off;
grid on;
xlabel('Time t', 'FontSize', 12);
ylabel('Position x', 'FontSize', 12);
title('Particle Trajectories in Space-Time', 'FontSize', 13);
xlim([0, 1]);
ylim([x_min, x_max]);
set(gca, 'FontSize', 11);

% Plot final trajectory visualization
subplot(1, 2, 2);
plot(x, rho_0, 'b-', 'LineWidth', 2, 'DisplayName', '\rho_0 (two Gaussians)');
hold on;
plot(x, rho_1, 'r-', 'LineWidth', 2, 'DisplayName', '\rho_1 (merged)');

% Draw arrows showing transport
arrow_y = 0.1 * max_density;
for i = 1:n_particles
    x0 = trajectories(i, 1);
    x1 = trajectories(i, end);
    % Draw arrow using quiver
    quiver(x0, arrow_y, x1-x0, 0, 0, 'k', 'LineWidth', 1.5, 'MaxHeadSize', 0.8);
end

plot(particle_positions_init, 0.1*max_density*ones(n_particles,1), 'bo', ...
     'MarkerSize', 6, 'MarkerFaceColor', 'b', 'DisplayName', 'Initial particles');
plot(trajectories(:, end), 0.1*max_density*ones(n_particles,1), 'ro', ...
     'MarkerSize', 6, 'MarkerFaceColor', 'r', 'DisplayName', 'Final particles');

hold off;
grid on;
xlabel('x', 'FontSize', 12);
ylabel('Density', 'FontSize', 12);
title('Mass Transport Visualization', 'FontSize', 13);
legend('Location', 'northwest');
xlim([x_min, x_max]);
ylim([0, max_density]);
set(gca, 'FontSize', 11);

sgtitle('Lagrangian Particle Trajectories: Two Gaussians Merging', 'FontSize', 14, 'FontWeight', 'bold');

fprintf('Trajectory visualization complete!\n');

%% Helper Functions

function rho_pushed = compute_pushforward(x, rho, T_map, dx)
    % Compute pushforward density: T# rho
    % Using semi-Lagrangian method: rho_pushed(y) = rho(T^{-1}(y)) / |dT/dx|(T^{-1}(y))
    N = length(x);
    rho_pushed = zeros(N, 1);

    % Compute derivative of T using finite differences
    dT = gradient(T_map, dx);
    dT = max(dT, 1e-10);  % Ensure positive (monotonic map)

    % For each target point y, find where it came from (inverse map)
    for j = 1:N
        y = x(j);

        % Find T^{-1}(y) by finding x such that T(x) = y
        % Use interpolation
        if y >= min(T_map) && y <= max(T_map)
            % Check if T is monotonic
            if all(diff(T_map) > 0) || all(diff(T_map) < 0)
                % Inverse map
                x_inv = interp1(T_map, x, y, 'linear', 'extrap');

                % Evaluate density and Jacobian at preimage
                rho_at_inv = interp1(x, rho, x_inv, 'linear', 'extrap');
                dT_at_inv = interp1(x, dT, x_inv, 'linear', 'extrap');

                % Pushforward formula: rho(y) = rho(x) / |dT/dx|
                rho_pushed(j) = rho_at_inv / abs(dT_at_inv);
            else
                % Fallback: use binning method
                rho_pushed(j) = 0;
            end
        end
    end

    % Apply smoothing to reduce oscillations
    if N > 5
        % Simple moving average filter
        window = 3;
        rho_pushed = movmean(rho_pushed, window);
    end

    % Normalize to maintain mass
    total_mass = sum(rho_pushed) * dx;
    if total_mass > 1e-10
        rho_pushed = rho_pushed / total_mass;
    else
        % Fallback: uniform distribution
        rho_pushed = ones(N, 1) / (N * dx);
    end
end

function v_pushed = compute_pushforward_vector(x, rho, v, T_map, dx)
    % Compute pushforward of velocity field
    % Velocity is simply pulled back: v(y) = v_source(T^{-1}(y))
    N = length(x);
    v_pushed = zeros(N, 1);

    % For each target point, find the velocity at its preimage
    for j = 1:N
        y = x(j);

        % Find T^{-1}(y)
        if y >= min(T_map) && y <= max(T_map)
            if all(diff(T_map) > 0) || all(diff(T_map) < 0)
                x_inv = interp1(T_map, x, y, 'linear', 'extrap');
                % Velocity at preimage
                v_pushed(j) = interp1(x, v, x_inv, 'linear', 'extrap');
            end
        end
    end

    % Apply smoothing
    if N > 5
        v_pushed = movmean(v_pushed, 3);
    end
end

function samples = sample_from_density(x, rho, n_samples)
    % Sample points from a probability density using inverse CDF method

    % Compute CDF
    F = cumtrapz(x, rho);
    F = F / max(F(end), 1e-10);  % Normalize to [0,1]

    % Ensure strict monotonicity
    eps_mono = 1e-10;
    for i = 2:length(F)
        if F(i) <= F(i-1)
            F(i) = F(i-1) + eps_mono;
        end
    end
    F = F / F(end);  % Final normalization

    % Generate uniform samples
    u = linspace(0.05, 0.95, n_samples);  % Avoid boundaries

    % Inverse CDF sampling with error handling
    try
        samples = interp1(F, x, u, 'linear', 'extrap');
    catch
        % Fallback: use percentiles
        samples = interp1(1:length(x), x, linspace(1, length(x), n_samples), 'linear');
    end
end
