%% Wave Equation Demo: 1D finite difference solution
%  Solves: ∂²u/∂t² - c²∂²u/∂x² = 0  on [0, 1] x [0, T]
%  Boundary conditions: u(0,t) = u(1,t) = 0  (homogeneous Dirichlet)
%  Initial displacement: u(x,0) = sin(2*pi*x)
%  Initial velocity:     ∂u/∂t(x,0) = 0
%  Exact solution:       u(x,t) = sin(2*pi*x) * cos(2*pi*c*t)
%
%  Method: Central differences in both space and time
%  CFL condition: c*dt/dx <= 1

clear; close all;

%% Parameters
c  = 1.0;          % wave speed
Nx = 200;          % number of spatial grid points
T  = 1.0;          % final time
CFL = 0.9;         % CFL number (must be <= 1 for stability)

%% Grid setup
dx = 1 / Nx;
dt = CFL * dx / c;               % time step satisfying CFL
Nt = ceil(T / dt);
dt = T / Nt;                     % adjust dt so we land exactly at T
x  = linspace(0, 1, Nx+1)';      % spatial grid (column vector)
fprintf('Wave equation: Nx = %d, Nt = %d, dx = %.4e, dt = %.4e\n', Nx, Nt, dx, dt);
fprintf('CFL number: c*dt/dx = %.4f\n', c * dt / dx);

%% Exact solution
u_exact = @(x, t) sin(2*pi*x) .* cos(2*pi*c*t);

%% Initial conditions
u_prev = sin(2*pi*x);            % u at t = 0
% For the first time step, use Taylor expansion:
%   u(x, dt) ≈ u(x,0) + dt*v(x,0) + 0.5*dt²*c²*u_xx(x,0)
% With v(x,0) = 0 and u_xx = -4π²sin(2πx):
u_curr = u_prev + 0.5 * (c*dt)^2 * (-4*pi^2*sin(2*pi*x));

%% Time-stepping with central differences
%  u^{n+1}_j = 2*u^n_j - u^{n-1}_j + (c*dt/dx)^2 * (u^n_{j+1} - 2*u^n_j + u^n_{j-1})
r2 = (c * dt / dx)^2;            % square of CFL number
u_new = zeros(Nx+1, 1);

% Store snapshots for plotting
n_snaps = 4;
snap_times = linspace(0, T, n_snaps+1);
snap_times = snap_times(2:end);   % skip t=0
snapshots = {};
snap_idx = 1;
next_snap = snap_times(1);

for n = 1:Nt
    t_curr = n * dt;

    % Interior points (enforce BCs: u(0)=u(end)=0)
    for j = 2:Nx
        u_new(j) = 2*u_curr(j) - u_prev(j) ...
                   + r2 * (u_curr(j+1) - 2*u_curr(j) + u_curr(j-1));
    end
    u_new(1)    = 0;
    u_new(Nx+1) = 0;

    % Save snapshot if at a target time
    if snap_idx <= length(snap_times) && t_curr >= next_snap - 0.5*dt
        snapshots{snap_idx}.u = u_new;
        snapshots{snap_idx}.t = t_curr;
        snap_idx = snap_idx + 1;
        if snap_idx <= length(snap_times)
            next_snap = snap_times(snap_idx);
        end
    end

    % Advance
    u_prev = u_curr;
    u_curr = u_new;
end

%% Compute errors at final time
u_exact_final = u_exact(x, T);
err = u_curr - u_exact_final;
L2_err   = sqrt(dx) * norm(err);
Linf_err = max(abs(err));
fprintf('L2 error at T = %.2f:   %.6e\n', T, L2_err);
fprintf('Linf error at T = %.2f: %.6e\n', T, Linf_err);

%% Plotting
fig = figure('Position', [100, 100, 1200, 900], 'Visible', 'off');

% Subplot 1: Solution snapshots
subplot(2,2,1);
plot(x, u_exact(x,0), 'k--', 'LineWidth', 1.5); hold on;
colors = lines(length(snapshots));
for k = 1:length(snapshots)
    plot(x, snapshots{k}.u, '-', 'Color', colors(k,:), 'LineWidth', 1.5);
end
hold off;
legend_entries = {'t = 0 (IC)'};
for k = 1:length(snapshots)
    legend_entries{end+1} = sprintf('t = %.2f', snapshots{k}.t);
end
legend(legend_entries, 'Location', 'best');
xlabel('x'); ylabel('u(x,t)');
title('Wave Propagation Snapshots');
grid on;

% Subplot 2: Numerical vs Exact at final time
subplot(2,2,2);
plot(x, u_curr, 'b-', 'LineWidth', 1.5); hold on;
plot(x, u_exact_final, 'r--', 'LineWidth', 1.5); hold off;
legend('Numerical', 'Exact', 'Location', 'best');
xlabel('x'); ylabel('u');
title(sprintf('Solution at T = %.2f', T));
grid on;

% Subplot 3: Pointwise error at final time
subplot(2,2,3);
plot(x, err, 'm-', 'LineWidth', 1.5);
xlabel('x'); ylabel('Error');
title(sprintf('Pointwise Error (L^\\infty = %.2e)', Linf_err));
grid on;

% Subplot 4: Energy conservation check
% Rerun to track energy
E = zeros(Nt+1, 1);
u_prev2 = sin(2*pi*x);
u_curr2 = u_prev2 + 0.5*(c*dt)^2*(-4*pi^2*sin(2*pi*x));
% Energy at t=0: E = 0.5*||v||^2 + 0.5*c^2*||u_x||^2
% Approximate v at t=0.5*dt for staggered energy
v_half = (u_curr2 - u_prev2) / dt;
ux = diff(u_curr2) / dx;  % approximate u_x at t=0 (using u at dt)
E(1) = 0.5*dx*sum(v_half(2:Nx).^2) + 0.5*c^2*dx*sum(ux.^2);

for n = 1:Nt
    u_new2 = zeros(Nx+1,1);
    for j = 2:Nx
        u_new2(j) = 2*u_curr2(j) - u_prev2(j) ...
                    + r2*(u_curr2(j+1) - 2*u_curr2(j) + u_curr2(j-1));
    end
    v_half = (u_new2 - u_curr2) / dt;
    ux = diff(u_new2) / dx;
    E(n+1) = 0.5*dx*sum(v_half(2:Nx).^2) + 0.5*c^2*dx*sum(ux.^2);
    u_prev2 = u_curr2;
    u_curr2 = u_new2;
end

subplot(2,2,4);
t_vec = (0:Nt) * dt;
plot(t_vec, E, 'b-', 'LineWidth', 1.5);
xlabel('t'); ylabel('Discrete Energy E(t)');
title('Energy Conservation');
grid on;
ylim_curr = ylim;
ylim([ylim_curr(1)*0.99, ylim_curr(2)*1.01]);

sgtitle(sprintf('Wave Equation: c = %.1f, N_x = %d, CFL = %.2f', c, Nx, c*dt/dx), ...
        'FontSize', 14, 'FontWeight', 'bold');

%% Save figure
print(fig, 'wave_result.png', '-dpng', '-r150');
fprintf('Figure saved to wave_result.png\n');
close(fig);
