%% Heat Equation Demo
%  Solves the 1D heat equation using Crank-Nicolson method
%
%  PDE: ∂u/∂t - α ∂²u/∂x² = 0  on (0, L) × (0, T]
%  BC:  u(0, t) = u(L, t) = 0  (homogeneous Dirichlet)
%  IC:  u(x, 0) = sin(πx/L)    (initial temperature distribution)
%
%  Exact solution: u(x,t) = exp(-α(π/L)²t) * sin(πx/L)
%
%  The Crank-Nicolson scheme is unconditionally stable and second-order
%  accurate in both space and time.

clear; close all;

%% Parameters
L = 1.0;          % Domain length
T = 0.5;          % Final time
alpha = 0.1;      % Thermal diffusivity
Nx = 50;          % Number of spatial grid points
Nt = 200;         % Number of time steps

%% Grid setup
dx = L / (Nx - 1);
dt = T / Nt;
x = linspace(0, L, Nx)';
t = linspace(0, T, Nt + 1);

% Courant number (for reference)
r = alpha * dt / dx^2;
fprintf('Courant number r = α·dt/dx² = %.4f\n', r);

%% Initial condition: u(x,0) = sin(πx/L)
u0 = sin(pi * x / L);

%% Exact solution function
exact_solution = @(x, t) exp(-alpha * (pi/L)^2 * t) .* sin(pi * x / L);

%% Crank-Nicolson discretization
%  (I + r/2 * A) u^{n+1} = (I - r/2 * A) u^n
%  where A is the tridiagonal matrix from central differences
%
%  This gives the semi-implicit scheme:
%  -r/2 u_{i-1}^{n+1} + (1+r) u_i^{n+1} - r/2 u_{i+1}^{n+1} =
%   r/2 u_{i-1}^n     + (1-r) u_i^n     + r/2 u_{i+1}^n

% Interior points only (exclude boundaries with Dirichlet BC)
N_int = Nx - 2;

% Build tridiagonal matrices for interior points
% Left-hand side: (I + r/2 * A)
diag_main_L = (1 + r) * ones(N_int, 1);
diag_off_L = (-r/2) * ones(N_int - 1, 1);
A_left = diag(diag_main_L) + diag(diag_off_L, 1) + diag(diag_off_L, -1);

% Right-hand side: (I - r/2 * A)
diag_main_R = (1 - r) * ones(N_int, 1);
diag_off_R = (r/2) * ones(N_int - 1, 1);
A_right = diag(diag_main_R) + diag(diag_off_R, 1) + diag(diag_off_R, -1);

%% Time-stepping
% Store solution at all time steps for visualization
U = zeros(Nx, Nt + 1);
U(:, 1) = u0;

u_int = u0(2:end-1);  % Interior values

for n = 1:Nt
    % Crank-Nicolson step
    rhs = A_right * u_int;
    u_int = A_left \ rhs;

    % Store full solution (with boundary values = 0)
    U(2:end-1, n+1) = u_int;
end

%% Compute exact solution at final time
u_exact_final = exact_solution(x, T);
u_numerical_final = U(:, end);

%% Error analysis
error_pointwise = abs(u_numerical_final - u_exact_final);
error_max = max(error_pointwise);
error_L2 = sqrt(dx * sum(error_pointwise.^2));
error_rms = sqrt(mean(error_pointwise.^2));

fprintf('\n=== Error Metrics at t = %.2f ===\n', T);
fprintf('Max (L∞) error: %.6e\n', error_max);
fprintf('L² error:       %.6e\n', error_L2);
fprintf('RMS error:      %.6e\n', error_rms);
fprintf('Grid: Nx = %d, Nt = %d\n', Nx, Nt);

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

% Subplot 1: Numerical solution evolution
subplot(1, 3, 1);
t_indices = round(linspace(1, Nt+1, 6));  % Select 6 time snapshots
colors = parula(length(t_indices));
hold on;
for i = 1:length(t_indices)
    idx = t_indices(i);
    plot(x, U(:, idx), 'Color', colors(i,:), 'LineWidth', 1.5);
end
hold off;
xlabel('x', 'FontSize', 12);
ylabel('u(x,t)', 'FontSize', 12);
title('Numerical Solution Evolution', 'FontSize', 14);
legend(arrayfun(@(i) sprintf('t = %.2f', t(t_indices(i))), 1:length(t_indices), ...
    'UniformOutput', false), 'Location', 'northeast');
grid on;

% Subplot 2: Comparison at final time
subplot(1, 3, 2);
plot(x, u_numerical_final, 'b-', 'LineWidth', 2, 'DisplayName', 'Numerical');
hold on;
plot(x, u_exact_final, 'r--', 'LineWidth', 2, 'DisplayName', 'Exact');
hold off;
xlabel('x', 'FontSize', 12);
ylabel('u(x,T)', 'FontSize', 12);
title(sprintf('Solution at t = %.2f', T), 'FontSize', 14);
legend('Location', 'northeast');
grid on;

% Subplot 3: Pointwise error
subplot(1, 3, 3);
semilogy(x, error_pointwise + eps, 'k-', 'LineWidth', 1.5);
xlabel('x', 'FontSize', 12);
ylabel('|u_{num} - u_{exact}|', 'FontSize', 12);
title('Pointwise Error (log scale)', 'FontSize', 14);
grid on;

% Add overall title
sgtitle(sprintf('Heat Equation: Crank-Nicolson Method (Nx=%d, Nt=%d, r=%.3f)', ...
    Nx, Nt, r), 'FontSize', 16, 'FontWeight', 'bold');

%% Save figure
print('heat_equation_result.png', '-dpng', '-r150');
fprintf('\nFigure saved as heat_equation_result.png\n');
