% Mean Field Game Demo: 1D Crowd Congestion
% Solves the coupled HJB-Fokker-Planck system via Picard iteration
% Agents on [0,1] move toward a target while avoiding congestion
%
% HJB (backward):  -u_t - nu*u_xx + H(u_x) = alpha*rho
% FP  (forward):    rho_t - div(rho * H'(u_x)) - nu*rho_xx = 0
% with H(p) = |p|^2 / 2  (quadratic Hamiltonian)

clear all; close all; clc;

fprintf('=== 1D Mean Field Game: Crowd Congestion ===\n\n');

%% Parameters
Nx = 101;           % spatial grid points (odd for symmetry)
T  = 1.0;           % final time
L  = 1.0;           % domain length
sigma = 0.15;       % noise intensity
alpha = 1.5;        % congestion cost weight
beta  = 3.0;        % terminal cost weight
x_target = 0.70;    % target location
maxIter = 300;      % max fictitious play iterations
tol = 1e-3;         % convergence tolerance

nu = sigma^2 / 2;   % diffusion coefficient
dx = L / (Nx - 1);
xg = linspace(0, L, Nx)';

% Time step: must satisfy CFL for both diffusion and advection
% Start conservative; advection CFL checked dynamically
dt = min(0.3 * dx^2 / nu, 0.1 * dx);
Nt = ceil(T / dt);
dt = T / Nt;
tg = linspace(0, T, Nt+1);

fprintf('Spatial grid:  Nx = %d,  dx = %.4f\n', Nx, dx);
fprintf('Time grid:     Nt = %d,  dt = %.5f\n', Nt, dt);
fprintf('Diffusion:     sigma = %.3f,  nu = %.5f\n', sigma, nu);
fprintf('Congestion:    alpha = %.1f\n', alpha);
fprintf('Terminal cost:  beta = %.1f,  target = %.2f\n', beta, x_target);
fprintf('CFL ratio:     nu*dt/dx^2 = %.3f\n\n', nu*dt/dx^2);

%% Initial density: Gaussian centered at 0.25
rho0 = exp(-((xg - 0.25).^2) / (2 * 0.08^2));
rho0 = rho0 / (sum(rho0) * dx);

%% Terminal cost
gterm = beta * (xg - x_target).^2;

%% Solve via Fictitious Play (guaranteed convergence for monotone MFGs)
% Each iteration: solve HJB for best response, solve FP for resulting
% density, then update the mean density as a running average.
rho_bar = repmat(rho0, 1, Nt+1);  % running average density
u       = repmat(gterm, 1, Nt+1);
convergence_history = zeros(maxIter, 1);

fig_anim = figure('Position', [100, 100, 1400, 800], 'Name', 'MFG Fictitious Play');

fprintf('Running fictitious play...\n');

for iter = 1:maxIter
    rho_bar_old = rho_bar;

    %% Step 1: Solve HJB backward given rho_bar (Godunov scheme)
    u(:, end) = gterm;

    for n = Nt:-1:1
        Dpu = zeros(Nx, 1);
        Dmu = zeros(Nx, 1);
        Dpu(1:Nx-1) = (u(2:Nx, n+1) - u(1:Nx-1, n+1)) / dx;
        Dmu(2:Nx)   = (u(2:Nx, n+1) - u(1:Nx-1, n+1)) / dx;
        Dpu(Nx) = 0;
        Dmu(1)  = 0;

        % Godunov Hamiltonian for H(p) = p^2/2
        H_god = 0.5 * max(max(Dmu, 0).^2, max(-Dpu, 0).^2);

        % Diffusion Laplacian
        uxx = zeros(Nx, 1);
        uxx(2:Nx-1) = (u(3:Nx, n+1) - 2*u(2:Nx-1, n+1) + u(1:Nx-2, n+1)) / dx^2;
        uxx(1)  = (u(2, n+1) - u(1, n+1)) / dx^2;
        uxx(Nx) = (u(Nx-1, n+1) - u(Nx, n+1)) / dx^2;

        % Backward step: u(n) = u(n+1) + dt*(nu*u_xx - H + alpha*rho_bar)
        u(:, n) = u(:, n+1) + dt * (nu * uxx - H_god + alpha * rho_bar(:, n+1));
    end

    %% Step 2: Solve Fokker-Planck forward given u (best response density)
    rho_br = zeros(Nx, Nt+1);  % best response density
    rho_br(:, 1) = rho0;

    for n = 1:Nt
        % Interface velocities: v_{j+1/2} = -(u_{j+1} - u_j)/dx
        v_half = -(u(2:Nx, n) - u(1:Nx-1, n)) / dx;
        rho_n = rho_br(:, n);

        % Upwind flux
        pos_mask = v_half >= 0;
        flux = v_half .* (pos_mask .* rho_n(1:Nx-1) + ~pos_mask .* rho_n(2:Nx));

        % Divergence (zero-flux BC)
        div_flux = zeros(Nx, 1);
        div_flux(1)      =  flux(1) / dx;
        div_flux(2:Nx-1) = (flux(2:Nx-1) - flux(1:Nx-2)) / dx;
        div_flux(Nx)     = -flux(Nx-1) / dx;

        % Diffusion Laplacian
        rho_xx = zeros(Nx, 1);
        rho_xx(2:Nx-1) = (rho_n(3:Nx) - 2*rho_n(2:Nx-1) + rho_n(1:Nx-2)) / dx^2;
        rho_xx(1)  = (rho_n(2) - rho_n(1)) / dx^2;
        rho_xx(Nx) = (rho_n(Nx-1) - rho_n(Nx)) / dx^2;

        rho_br(:, n+1) = rho_n + dt * (nu * rho_xx - div_flux);
        rho_br(:, n+1) = max(rho_br(:, n+1), 0);
        mass_n = sum(rho_br(:, n+1)) * dx;
        if mass_n > 1e-10
            rho_br(:, n+1) = rho_br(:, n+1) / mass_n;
        end
    end

    %% Step 3: Fictitious play update (running average)
    rho_bar = ((iter - 1) * rho_bar + rho_br) / iter;

    %% Convergence check
    err = sqrt(sum((rho_bar(:) - rho_bar_old(:)).^2) * dx * dt);
    convergence_history(iter) = err;

    if mod(iter, 10) == 1 || err < tol
        fprintf('  Iter %3d:  L2|drho_bar| = %.2e\n', iter, err);
    end

    %% Animate
    if mod(iter, 5) == 1 || iter == maxIter || err < tol
        t_show = [1, round(Nt/4), round(Nt/2), round(3*Nt/4), Nt+1];
        colors = lines(length(t_show));

        subplot(2, 2, 1); cla; hold on;
        for k = 1:length(t_show)
            plot(xg, rho_bar(:, t_show(k)), '-', 'LineWidth', 2, 'Color', colors(k,:), ...
                 'DisplayName', sprintf('t = %.2f', tg(t_show(k))));
        end
        xline(x_target, 'k--', 'LineWidth', 1, 'HandleVisibility', 'off');
        hold off; grid on;
        xlabel('x'); ylabel('\rho(x,t)');
        title(sprintf('Mean Density (iter %d)', iter));
        legend('Location', 'best', 'FontSize', 8); xlim([0 L]);

        subplot(2, 2, 2); cla; hold on;
        for k = 1:length(t_show)
            plot(xg, u(:, t_show(k)), '-', 'LineWidth', 2, 'Color', colors(k,:), ...
                 'DisplayName', sprintf('t = %.2f', tg(t_show(k))));
        end
        hold off; grid on;
        xlabel('x'); ylabel('u(x,t)');
        title('Value Function');
        legend('Location', 'best', 'FontSize', 8); xlim([0 L]);

        subplot(2, 2, 3);
        vel_field_tmp = zeros(Nx, Nt+1);
        for nn = 1:Nt+1
            vel_field_tmp(2:end-1, nn) = -(u(3:end, nn) - u(1:end-2, nn)) / (2*dx);
            vel_field_tmp(1, nn) = vel_field_tmp(2, nn);
            vel_field_tmp(end, nn) = vel_field_tmp(end-1, nn);
        end
        imagesc(tg, xg, vel_field_tmp);
        set(gca, 'YDir', 'normal'); colorbar;
        xlabel('Time t'); ylabel('x');
        title('Velocity v = -\partial u/\partial x');

        subplot(2, 2, 4);
        semilogy(1:iter, convergence_history(1:iter), 'b-', 'LineWidth', 1.5);
        hold on; yline(tol, 'r--', 'LineWidth', 1.5); hold off;
        grid on;
        xlabel('Iteration'); ylabel('L^2 error');
        title('Convergence (Fictitious Play)');
        legend('Error', 'Tolerance', 'Location', 'best');

        sgtitle(sprintf('Mean Field Game — Fictitious Play, Iteration %d', iter), ...
                'FontSize', 14, 'FontWeight', 'bold');
        drawnow; pause(0.05);
    end

    if err < tol
        fprintf('\nConverged at iteration %d (error = %.2e)\n', iter, err);
        break;
    end
end

totalIter = iter;
rho = rho_bar;  % use converged mean density for plots

if convergence_history(totalIter) < tol
    fprintf('Fictitious play converged after %d iterations\n', totalIter);
else
    fprintf('\nFinal error after %d iterations: %.2e\n', totalIter, convergence_history(totalIter));
end

%% Compute velocity field for final plots
vel_field = zeros(Nx, Nt+1);
for n = 1:Nt+1
    vel_field(2:end-1, n) = -(u(3:end, n) - u(1:end-2, n)) / (2*dx);
    vel_field(1, n) = vel_field(2, n);
    vel_field(end, n) = vel_field(end-1, n);
end

%% Final static figure
fig_final = figure('Position', [150, 150, 1400, 800], 'Name', 'MFG Results');

% Panel 1: Space-time density heatmap
subplot(2, 3, 1);
imagesc(tg, xg, rho);
set(gca, 'YDir', 'normal');
colorbar;
xlabel('Time t'); ylabel('x');
title('Density \rho(x,t)');

% Panel 2: Space-time velocity heatmap
subplot(2, 3, 2);
imagesc(tg, xg, vel_field);
set(gca, 'YDir', 'normal');
colorbar;
xlabel('Time t'); ylabel('x');
title('Velocity v(x,t)');

% Panel 3: Initial vs final density
subplot(2, 3, 3);
plot(xg, rho(:,1), 'b-', 'LineWidth', 2, 'DisplayName', '\rho(x, 0)');
hold on;
plot(xg, rho(:,end), 'r-', 'LineWidth', 2, 'DisplayName', '\rho(x, T)');
xline(x_target, 'k--', 'LineWidth', 1.5, 'DisplayName', 'Target');
hold off;
grid on;
xlabel('x'); ylabel('Density');
title('Initial vs Final Density');
legend('Location', 'best');
xlim([0 L]);

% Panel 4: Value function at t=0
subplot(2, 3, 4);
plot(xg, u(:,1), 'b-', 'LineWidth', 2);
grid on;
xlabel('x'); ylabel('u(x, 0)');
title('Value Function at t = 0');
xlim([0 L]);

% Panel 5: Optimal velocity at t=0
subplot(2, 3, 5);
vel_t0 = vel_field(:, 1);
plot(xg, vel_t0, 'b-', 'LineWidth', 2);
hold on;
yline(0, 'k--');
hold off;
grid on;
xlabel('x'); ylabel('v(x, 0)');
title('Optimal Velocity at t = 0');
xlim([0 L]);

% Panel 6: Convergence history
subplot(2, 3, 6);
semilogy(1:totalIter, convergence_history(1:totalIter), 'b-o', ...
         'LineWidth', 1.5, 'MarkerSize', 4);
hold on;
yline(tol, 'r--', 'LineWidth', 1.5);
hold off;
grid on;
xlabel('Iteration'); ylabel('L^2|\Delta\rho|');
title('Convergence History');
legend('Error', 'Tolerance', 'Location', 'best');

sgtitle('1D Mean Field Game: Crowd Congestion', ...
        'FontSize', 14, 'FontWeight', 'bold');

saveas(fig_final, fullfile(fileparts(mfilename('fullpath')), 'mean_field_game_demo.png'));
fprintf('\nFigure saved to mean_field_game_demo.png\n');

%% Summary table
mass_initial = sum(rho(:,1)) * dx;
mass_final   = sum(rho(:,end)) * dx;
mass_err = abs(mass_final - mass_initial) / mass_initial;

% Value function statistics at t=0 (weighted by density)
u0_weighted = u(:,1) .* rho(:,1) * dx;
u0_mean = sum(u0_weighted);
u0_std  = sqrt(max(0, sum((u(:,1) - u0_mean).^2 .* rho(:,1)) * dx));

% Peak density location at final time
[peak_val, peak_idx] = max(rho(:, end));
peak_location = xg(peak_idx);

fprintf('\n=== Summary ===\n');
fprintf('%-30s %s\n', 'Parameter', 'Value');
fprintf('%s\n', repmat('-', 1, 50));
fprintf('%-30s %d\n',   'Fictitious play iterations', totalIter);
fprintf('%-30s %.2e\n', 'Final L2 error',           convergence_history(totalIter));
fprintf('%-30s %.6f\n', 'Mass at t=0',              mass_initial);
fprintf('%-30s %.6f\n', 'Mass at t=T',              mass_final);
fprintf('%-30s %.2e\n', 'Mass conservation error',  mass_err);
fprintf('%-30s %.4f\n', 'Target location',          x_target);
fprintf('%-30s %.4f\n', 'Peak density at t=T',      peak_location);
fprintf('%-30s %.4f\n', 'Peak density value',       peak_val);
fprintf('%-30s %.4f\n', 'Value function std (t=0)', u0_std);
fprintf('%-30s %.4f\n', 'Mean cost per agent',      u0_mean);
fprintf('%s\n', repmat('-', 1, 50));
fprintf('Interpretation: agents spread around target x=%.2f,\n', x_target);
fprintf('balancing proximity-to-target vs congestion avoidance.\n');
