% First-Order Lax-Friedrichs Scheme for Lax Shock Tube Problem
% 1D Compressible Euler Equations

clear all; close all; clc;

%% Problem Setup
fprintf('=== Lax-Friedrichs Scheme: Lax Shock Tube ===\n');

% Physical parameters
gamma = 1.4;  % Ratio of specific heats

% Domain
x_min = 0.0;
x_max = 1.0;
N = 200;  % Number of grid points

% Create uniform mesh
x = linspace(x_min, x_max, N);
dx = (x_max - x_min) / (N - 1);

fprintf('Grid points: %d\n', N);
fprintf('Grid spacing: %.6f\n', dx);

% Time parameters
CFL = 0.9;  % CFL number
T_final = 0.14;  % Final time

%% Initial Conditions: Lax Shock Tube
% Left state (x < 0.5)
rho_L = 0.445;
u_L = 0.698;
p_L = 3.528;

% Right state (x > 0.5)
rho_R = 0.5;
u_R = 0.0;
p_R = 0.571;

fprintf('\nLax Shock Tube Initial Conditions:\n');
fprintf('Left state:  rho = %.3f, u = %.3f, p = %.3f\n', rho_L, u_L, p_L);
fprintf('Right state: rho = %.3f, u = %.3f, p = %.3f\n', rho_R, u_R, p_R);

% Initialize conservative variables U = [rho, rho*u, E]
U = zeros(N, 3);
x_membrane = 0.5;

for i = 1:N
    if x(i) <= x_membrane
        % Left state
        rho = rho_L;
        u = u_L;
        p = p_L;
    else
        % Right state
        rho = rho_R;
        u = u_R;
        p = p_R;
    end

    % Total energy per unit volume
    E = p / (gamma - 1) + 0.5 * rho * u^2;

    % Conservative variables
    U(i, 1) = rho;
    U(i, 2) = rho * u;
    U(i, 3) = E;
end

% Store initial condition
U_initial = U;

%% Setup figure for animation
fig = figure('Position', [100, 100, 1400, 800]);

% Pre-compute primitive variables for initial condition
[rho_init, u_init, p_init] = compute_primitives(U_initial, gamma);

%% Time Integration with Lax-Friedrichs Scheme
t = 0;
iter = 0;
plot_interval = 5;

fprintf('\nStarting time integration...\n');

while t < T_final
    % Compute time step based on CFL condition
    dt = compute_timestep(U, dx, gamma, CFL);

    % Ensure we don't overshoot final time
    if t + dt > T_final
        dt = T_final - t;
    end

    % Lax-Friedrichs update
    U_new = zeros(N, 3);

    for i = 2:N-1
        % Compute fluxes at i-1, i, i+1
        F_im1 = euler_flux_1d(U(i-1, :), gamma);
        F_ip1 = euler_flux_1d(U(i+1, :), gamma);

        % Lax-Friedrichs scheme
        % U_i^{n+1} = 0.5*(U_{i-1}^n + U_{i+1}^n) - (dt/(2*dx))*(F_{i+1}^n - F_{i-1}^n)
        U_new(i, :) = 0.5 * (U(i-1, :) + U(i+1, :)) - ...
                      (dt / (2 * dx)) * (F_ip1 - F_im1);
    end

    % Boundary conditions: extrapolation (outflow)
    U_new(1, :) = U_new(2, :);
    U_new(N, :) = U_new(N-1, :);

    % Update solution
    U = U_new;

    % Update time and iteration counter
    t = t + dt;
    iter = iter + 1;

    % Plotting
    if mod(iter, plot_interval) == 0 || t >= T_final
        % Compute primitive variables
        [rho, u, p] = compute_primitives(U, gamma);
        e = p ./ ((gamma - 1) * rho);  % Specific internal energy

        % Plot density
        subplot(2, 2, 1);
        plot(x, rho_init, 'k--', 'LineWidth', 1.5); hold on;
        plot(x, rho, 'b-', 'LineWidth', 2); hold off;
        grid on;
        xlabel('x', 'FontSize', 11);
        ylabel('Density \rho', 'FontSize', 11);
        title(sprintf('Density at t = %.4f', t), 'FontSize', 12);
        legend('Initial', 'Current', 'Location', 'best');
        xlim([x_min, x_max]);
        ylim([0.3, 0.6]);

        % Plot velocity
        subplot(2, 2, 2);
        plot(x, u_init, 'k--', 'LineWidth', 1.5); hold on;
        plot(x, u, 'r-', 'LineWidth', 2); hold off;
        grid on;
        xlabel('x', 'FontSize', 11);
        ylabel('Velocity u', 'FontSize', 11);
        title(sprintf('Velocity at t = %.4f', t), 'FontSize', 12);
        legend('Initial', 'Current', 'Location', 'best');
        xlim([x_min, x_max]);
        ylim([-0.1, 0.8]);

        % Plot pressure
        subplot(2, 2, 3);
        plot(x, p_init, 'k--', 'LineWidth', 1.5); hold on;
        plot(x, p, 'g-', 'LineWidth', 2); hold off;
        grid on;
        xlabel('x', 'FontSize', 11);
        ylabel('Pressure p', 'FontSize', 11);
        title(sprintf('Pressure at t = %.4f', t), 'FontSize', 12);
        legend('Initial', 'Current', 'Location', 'best');
        xlim([x_min, x_max]);
        ylim([0, 4]);

        % Plot specific internal energy
        subplot(2, 2, 4);
        e_init = p_init ./ ((gamma - 1) * rho_init);
        plot(x, e_init, 'k--', 'LineWidth', 1.5); hold on;
        plot(x, e, 'm-', 'LineWidth', 2); hold off;
        grid on;
        xlabel('x', 'FontSize', 11);
        ylabel('Specific Internal Energy e', 'FontSize', 11);
        title(sprintf('Internal Energy at t = %.4f', t), 'FontSize', 12);
        legend('Initial', 'Current', 'Location', 'best');
        xlim([x_min, x_max]);
        ylim([1, 7]);

        sgtitle('Lax-Friedrichs Scheme: Lax Shock Tube Problem', ...
                'FontSize', 14, 'FontWeight', 'bold');

        drawnow;

        if mod(iter, 50) == 0
            fprintf('Iter %4d: t = %.6f, dt = %.6e\n', iter, t, dt);
        end
    end
end

fprintf('\nSimulation complete!\n');
fprintf('Total iterations: %d\n', iter);
fprintf('Final time: %.6f\n', t);

%% Final solution analysis
[rho_final, u_final, p_final] = compute_primitives(U, gamma);

fprintf('\nFinal solution ranges:\n');
fprintf('Density:  [%.4f, %.4f]\n', min(rho_final), max(rho_final));
fprintf('Velocity: [%.4f, %.4f]\n', min(u_final), max(u_final));
fprintf('Pressure: [%.4f, %.4f]\n', min(p_final), max(p_final));

%% Additional visualization: Wave diagram
figure('Position', [150, 150, 1200, 400]);

% Density
subplot(1, 3, 1);
plot(x, rho_init, 'k--', 'LineWidth', 2, 'DisplayName', 'Initial');
hold on;
plot(x, rho_final, 'b-', 'LineWidth', 2, 'DisplayName', sprintf('t = %.3f', T_final));
hold off;
grid on;
xlabel('x', 'FontSize', 12);
ylabel('\rho', 'FontSize', 12);
title('Density Profile', 'FontSize', 13);
legend('Location', 'best');
xlim([x_min, x_max]);

% Velocity
subplot(1, 3, 2);
plot(x, u_init, 'k--', 'LineWidth', 2, 'DisplayName', 'Initial');
hold on;
plot(x, u_final, 'r-', 'LineWidth', 2, 'DisplayName', sprintf('t = %.3f', T_final));
hold off;
grid on;
xlabel('x', 'FontSize', 12);
ylabel('u', 'FontSize', 12);
title('Velocity Profile', 'FontSize', 13);
legend('Location', 'best');
xlim([x_min, x_max]);

% Pressure
subplot(1, 3, 3);
plot(x, p_init, 'k--', 'LineWidth', 2, 'DisplayName', 'Initial');
hold on;
plot(x, p_final, 'g-', 'LineWidth', 2, 'DisplayName', sprintf('t = %.3f', T_final));
hold off;
grid on;
xlabel('x', 'FontSize', 12);
ylabel('p', 'FontSize', 12);
title('Pressure Profile', 'FontSize', 13);
legend('Location', 'best');
xlim([x_min, x_max]);

sgtitle(sprintf('Lax Shock Tube: Initial vs Final (t = %.3f)', T_final), ...
        'FontSize', 14, 'FontWeight', 'bold');

%% Compute Exact Solution and Compare
fprintf('\nComputing exact Riemann solution...\n');

% Compute exact solution at final time
[rho_exact, u_exact, p_exact] = exact_riemann_solution(x, T_final, ...
    rho_L, u_L, p_L, rho_R, u_R, p_R, gamma, x_membrane);

% Compute errors
L1_error_rho = sum(abs(rho_final - rho_exact)) * dx;
L1_error_u = sum(abs(u_final - u_exact)) * dx;
L1_error_p = sum(abs(p_final - p_exact)) * dx;

fprintf('L1 Errors at t = %.4f:\n', T_final);
fprintf('  Density:  %.6e\n', L1_error_rho);
fprintf('  Velocity: %.6e\n', L1_error_u);
fprintf('  Pressure: %.6e\n', L1_error_p);

% Plot comparison with exact solution
figure('Position', [200, 200, 1400, 500]);

% Density comparison
subplot(1, 3, 1);
plot(x, rho_exact, 'k-', 'LineWidth', 2.5, 'DisplayName', 'Exact');
hold on;
plot(x, rho_final, 'b--', 'LineWidth', 1.5, 'DisplayName', 'Lax-Friedrichs');
hold off;
grid on;
xlabel('x', 'FontSize', 12);
ylabel('\rho', 'FontSize', 12);
title(sprintf('Density (L1 error = %.4e)', L1_error_rho), 'FontSize', 13);
legend('Location', 'best');
xlim([x_min, x_max]);
set(gca, 'FontSize', 11);

% Velocity comparison
subplot(1, 3, 2);
plot(x, u_exact, 'k-', 'LineWidth', 2.5, 'DisplayName', 'Exact');
hold on;
plot(x, u_final, 'r--', 'LineWidth', 1.5, 'DisplayName', 'Lax-Friedrichs');
hold off;
grid on;
xlabel('x', 'FontSize', 12);
ylabel('u', 'FontSize', 12);
title(sprintf('Velocity (L1 error = %.4e)', L1_error_u), 'FontSize', 13);
legend('Location', 'best');
xlim([x_min, x_max]);
set(gca, 'FontSize', 11);

% Pressure comparison
subplot(1, 3, 3);
plot(x, p_exact, 'k-', 'LineWidth', 2.5, 'DisplayName', 'Exact');
hold on;
plot(x, p_final, 'g--', 'LineWidth', 1.5, 'DisplayName', 'Lax-Friedrichs');
hold off;
grid on;
xlabel('x', 'FontSize', 12);
ylabel('p', 'FontSize', 12);
title(sprintf('Pressure (L1 error = %.4e)', L1_error_p), 'FontSize', 13);
legend('Location', 'best');
xlim([x_min, x_max]);
set(gca, 'FontSize', 11);

sgtitle(sprintf('Numerical vs Exact Solution at t = %.3f', T_final), ...
        'FontSize', 14, 'FontWeight', 'bold');

%% Helper Functions

function F = euler_flux_1d(U, gamma)
    % Compute 1D Euler flux vector
    % U = [rho, rho*u, E]
    % F = [rho*u, rho*u^2 + p, u*(E + p)]

    rho = U(1);
    rho_u = U(2);
    E = U(3);

    u = rho_u / rho;

    % Pressure from equation of state
    p = (gamma - 1) * (E - 0.5 * rho * u^2);

    % Flux vector
    F = zeros(1, 3);
    F(1) = rho_u;
    F(2) = rho_u * u + p;
    F(3) = u * (E + p);
end

function [rho, u, p] = compute_primitives(U, gamma)
    % Extract primitive variables from conservative variables
    % Input: U = [rho, rho*u, E] (N x 3)
    % Output: rho, u, p (N x 1 each)

    N = size(U, 1);
    rho = U(:, 1);
    u = U(:, 2) ./ rho;
    E = U(:, 3);

    % Compute pressure
    p = (gamma - 1) * (E - 0.5 * rho .* u.^2);
end

function dt = compute_timestep(U, dx, gamma, CFL)
    % Compute time step based on CFL condition
    % dt <= CFL * dx / max(|u| + c)

    [rho, u, p] = compute_primitives(U, gamma);

    % Sound speed
    c = sqrt(gamma * p ./ rho);

    % Maximum wave speed
    max_wave_speed = max(abs(u) + c);

    % Time step
    dt = CFL * dx / max_wave_speed;
end

function [rho, u, p] = exact_riemann_solution(x, t, rho_L, u_L, p_L, rho_R, u_R, p_R, gamma, x0)
    % Exact solution to the Riemann problem for Euler equations
    % Uses iterative solver for intermediate pressure

    N = length(x);
    rho = zeros(N, 1);
    u = zeros(N, 1);
    p = zeros(N, 1);

    % Sound speeds
    c_L = sqrt(gamma * p_L / rho_L);
    c_R = sqrt(gamma * p_R / rho_R);

    % Solve for intermediate pressure p_star using Newton iteration
    p_star = solve_p_star(rho_L, u_L, p_L, rho_R, u_R, p_R, gamma);

    % Compute intermediate velocity
    u_star = 0.5 * (u_L + u_R) + 0.5 * (f_K(p_star, rho_R, p_R, gamma) - ...
                                         f_K(p_star, rho_L, p_L, gamma));

    % Determine wave structure
    if p_star > p_L
        % Left shock
        S_L = u_L - c_L * sqrt((gamma + 1) / (2 * gamma) * p_star / p_L + ...
                                (gamma - 1) / (2 * gamma));
        rho_star_L = rho_L * (p_star / p_L + (gamma - 1) / (gamma + 1)) / ...
                             ((gamma - 1) / (gamma + 1) * p_star / p_L + 1);
    else
        % Left rarefaction
        rho_star_L = rho_L * (p_star / p_L)^(1 / gamma);
        c_star_L = c_L * (p_star / p_L)^((gamma - 1) / (2 * gamma));
        S_HL = u_L - c_L;  % Head of rarefaction
        S_TL = u_star - c_star_L;  % Tail of rarefaction
    end

    if p_star > p_R
        % Right shock
        S_R = u_R + c_R * sqrt((gamma + 1) / (2 * gamma) * p_star / p_R + ...
                                (gamma - 1) / (2 * gamma));
        rho_star_R = rho_R * (p_star / p_R + (gamma - 1) / (gamma + 1)) / ...
                             ((gamma - 1) / (gamma + 1) * p_star / p_R + 1);
    else
        % Right rarefaction
        rho_star_R = rho_R * (p_star / p_R)^(1 / gamma);
        c_star_R = c_R * (p_star / p_R)^((gamma - 1) / (2 * gamma));
        S_HR = u_R + c_R;  % Head of rarefaction
        S_TR = u_star + c_star_R;  % Tail of rarefaction
    end

    % Sample the solution
    for i = 1:N
        S = (x(i) - x0) / t;  % Characteristic speed

        if S < u_star
            % Left of contact
            if p_star > p_L
                % Left shock
                if S < S_L
                    rho(i) = rho_L;
                    u(i) = u_L;
                    p(i) = p_L;
                else
                    rho(i) = rho_star_L;
                    u(i) = u_star;
                    p(i) = p_star;
                end
            else
                % Left rarefaction
                if S < S_HL
                    rho(i) = rho_L;
                    u(i) = u_L;
                    p(i) = p_L;
                elseif S < S_TL
                    % Inside rarefaction fan
                    u(i) = 2 / (gamma + 1) * (c_L + (gamma - 1) / 2 * u_L + S);
                    c = c_L + (gamma - 1) / 2 * (u_L - u(i));
                    rho(i) = rho_L * (c / c_L)^(2 / (gamma - 1));
                    p(i) = p_L * (c / c_L)^(2 * gamma / (gamma - 1));
                else
                    rho(i) = rho_star_L;
                    u(i) = u_star;
                    p(i) = p_star;
                end
            end
        else
            % Right of contact
            if p_star > p_R
                % Right shock
                if S < S_R
                    rho(i) = rho_star_R;
                    u(i) = u_star;
                    p(i) = p_star;
                else
                    rho(i) = rho_R;
                    u(i) = u_R;
                    p(i) = p_R;
                end
            else
                % Right rarefaction
                if S < S_TR
                    rho(i) = rho_star_R;
                    u(i) = u_star;
                    p(i) = p_star;
                elseif S < S_HR
                    % Inside rarefaction fan
                    u(i) = 2 / (gamma + 1) * (-c_R + (gamma - 1) / 2 * u_R + S);
                    c = c_R - (gamma - 1) / 2 * (u_R - u(i));
                    rho(i) = rho_R * (c / c_R)^(2 / (gamma - 1));
                    p(i) = p_R * (c / c_R)^(2 * gamma / (gamma - 1));
                else
                    rho(i) = rho_R;
                    u(i) = u_R;
                    p(i) = p_R;
                end
            end
        end
    end
end

function p_star = solve_p_star(rho_L, u_L, p_L, rho_R, u_R, p_R, gamma)
    % Solve for intermediate pressure using Newton iteration

    % Initial guess (two-rarefaction approximation)
    c_L = sqrt(gamma * p_L / rho_L);
    c_R = sqrt(gamma * p_R / rho_R);

    p_star = ((c_L + c_R - 0.5 * (gamma - 1) * (u_R - u_L)) / ...
              (c_L / p_L^((gamma - 1) / (2 * gamma)) + ...
               c_R / p_R^((gamma - 1) / (2 * gamma))))^(2 * gamma / (gamma - 1));

    % Newton iteration
    tol = 1e-10;
    max_iter = 100;

    for iter = 1:max_iter
        f = f_K(p_star, rho_L, p_L, gamma) + f_K(p_star, rho_R, p_R, gamma) + ...
            (u_R - u_L);

        if abs(f) < tol
            break;
        end

        f_prime = df_K(p_star, rho_L, p_L, gamma) + df_K(p_star, rho_R, p_R, gamma);

        p_star = p_star - f / f_prime;

        % Ensure positivity
        p_star = max(p_star, 1e-10);
    end
end

function f = f_K(p, rho_K, p_K, gamma)
    % Nonlinear function for Riemann solver
    c_K = sqrt(gamma * p_K / rho_K);

    if p > p_K
        % Shock
        A_K = 2 / ((gamma + 1) * rho_K);
        B_K = (gamma - 1) / (gamma + 1) * p_K;
        f = (p - p_K) * sqrt(A_K / (p + B_K));
    else
        % Rarefaction
        f = 2 * c_K / (gamma - 1) * ((p / p_K)^((gamma - 1) / (2 * gamma)) - 1);
    end
end

function df = df_K(p, rho_K, p_K, gamma)
    % Derivative of nonlinear function for Riemann solver
    c_K = sqrt(gamma * p_K / rho_K);

    if p > p_K
        % Shock
        A_K = 2 / ((gamma + 1) * rho_K);
        B_K = (gamma - 1) / (gamma + 1) * p_K;
        df = sqrt(A_K / (B_K + p)) * (1 - (p - p_K) / (2 * (B_K + p)));
    else
        % Rarefaction
        df = (1 / (rho_K * c_K)) * (p / p_K)^(-(gamma + 1) / (2 * gamma));
    end
end
