%% =========================================================
%  Lax-Wendroff 方法求解一维 Euler 方程
%  经典算例：Sod 激波管问题
%
%  方程：
%       U_t + F(U)_x = 0
%
%  U = [rho, rho*u, E]^T
%
%  F = [rho*u,
%       rho*u^2 + p,
%       u*(E+p)]^T
%
%  理想气体状态方程：
%       p = (gamma - 1)*(E - 0.5*rho*u^2)
%
%  数值格式：
%       二步 Richtmyer 型 Lax-Wendroff 格式
%
%  物理现象：
%       左侧高压气体向右膨胀
%       右侧低压气体被压缩
%       形成稀疏波、接触间断和激波
%
%  课堂重点：
%       1. 守恒变量与通量
%       2. Lax-Wendroff 二阶精度
%       3. CFL 条件
%       4. 间断附近的数值振荡
%% =========================================================

clear; clc; close all;

%% ---------------------------------------------------------
%  1. 基本参数
%% ---------------------------------------------------------

gamma = 1.4;       % 比热比

xL = 0.0;          % 左边界
xR = 1.0;          % 右边界

Nx = 400;          % 网格数
dx = (xR - xL) / Nx;

x = linspace(xL + dx/2, xR - dx/2, Nx);  % 单元中心坐标

t  = 0.0;          % 初始时间
tf = 0.2;          % 终止时间

CFL = 0.5;         % CFL 数，建议小于 1

%% ---------------------------------------------------------
%  2. Sod 激波管初始条件
%
%  左状态：
%       rho_L = 1.0
%       u_L   = 0.0
%       p_L   = 1.0
%
%  右状态：
%       rho_R = 0.125
%       u_R   = 0.0
%       p_R   = 0.1
%% ---------------------------------------------------------

rho = zeros(1, Nx);
u   = zeros(1, Nx);
p   = zeros(1, Nx);

for i = 1:Nx
    if x(i) < 0.5
        rho(i) = 1.0;
        u(i)   = 0.0;
        p(i)   = 1.0;
    else
        rho(i) = 0.125;
        u(i)   = 0.0;
        p(i)   = 0.1;
    end
end

% 由原始变量转换为守恒变量
U = primitive_to_conservative(rho, u, p, gamma);

%% ---------------------------------------------------------
%  3. 初始化图像
%% ---------------------------------------------------------

figure('Color','w','Position',[100 100 900 700]);

%% ---------------------------------------------------------
%  4. 时间推进
%% ---------------------------------------------------------

step = 0;

while t < tf

    step = step + 1;

    % -----------------------------------------------------
    % 从守恒变量恢复原始变量
    % -----------------------------------------------------

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

    % -----------------------------------------------------
    % 计算声速和最大特征速度
    %
    % Euler 方程特征速度：
    %       lambda_1 = u - c
    %       lambda_2 = u
    %       lambda_3 = u + c
    %
    % 其中 c = sqrt(gamma*p/rho)
    % -----------------------------------------------------

    c = sqrt(gamma * p ./ rho);
    max_speed = max(abs(u) + c);

    % -----------------------------------------------------
    % 根据 CFL 条件确定时间步长
    %
    %       dt <= CFL * dx / max(|u|+c)
    % -----------------------------------------------------

    dt = CFL * dx / max_speed;

    if t + dt > tf
        dt = tf - t;
    end

    lambda = dt / dx;

    % -----------------------------------------------------
    % Lax-Wendroff 二步格式
    %
    % 第一步：预测半时间步、半网格点处的守恒变量
    %
    %   U_{i+1/2}^{n+1/2}
    %   = 0.5*(U_i^n + U_{i+1}^n)
    %     - 0.5*dt/dx*(F_{i+1}^n - F_i^n)
    %
    % 第二步：用半步通量更新单元中心守恒变量
    %
    %   U_i^{n+1}
    %   = U_i^n
    %     - dt/dx*(F_{i+1/2}^{n+1/2}
    %              -F_{i-1/2}^{n+1/2})
    %
    % 这种写法也常称为 Richtmyer two-step Lax-Wendroff
    % -----------------------------------------------------

    F = compute_flux(U, gamma);

    % 为了处理边界，这里添加两个 ghost cells
    U_ext = zeros(3, Nx+2);
    F_ext = zeros(3, Nx+2);

    U_ext(:,2:Nx+1) = U;
    F_ext(:,2:Nx+1) = F;

    % -----------------------------------------------------
    % 透射边界条件，即零梯度边界
    %
    % 左边 ghost cell = 第一个内部单元
    % 右边 ghost cell = 最后一个内部单元
    % -----------------------------------------------------

    U_ext(:,1)    = U_ext(:,2);
    U_ext(:,end)  = U_ext(:,end-1);

    F_ext(:,1)    = F_ext(:,2);
    F_ext(:,end)  = F_ext(:,end-1);

    % -----------------------------------------------------
    % 计算界面处半步守恒变量
    %
    % U_half(:,i) 对应界面 i-1/2
    % 例如：
    %   U_half(:,2) 对应原始网格中的 x_{1/2}
    %   U_half(:,3) 对应原始网格中的 x_{3/2}
    % -----------------------------------------------------

    U_half = zeros(3, Nx+1);

    for i = 1:Nx+1
        U_half(:,i) = 0.5 * (U_ext(:,i) + U_ext(:,i+1)) ...
                    - 0.5 * lambda * (F_ext(:,i+1) - F_ext(:,i));
    end

    % 半步通量
    F_half = compute_flux(U_half, gamma);

    % -----------------------------------------------------
    % 更新单元中心守恒变量
    % -----------------------------------------------------

    U_new = U;

    for i = 1:Nx
        U_new(:,i) = U(:,i) ...
                   - lambda * (F_half(:,i+1) - F_half(:,i));
    end

    U = U_new;
    t = t + dt;

    %% -----------------------------------------------------
    %  每隔若干步画图
    %% -----------------------------------------------------

    if mod(step,10) == 0 || t >= tf

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

        E = U(3,:);
        Mach = u ./ sqrt(gamma*p./rho);

        clf;

        subplot(2,2,1);
        plot(x, rho, 'LineWidth', 1.8);
        grid on;
        xlabel('x');
        ylabel('\rho');
        title(sprintf('Density, t = %.4f', t));

        subplot(2,2,2);
        plot(x, u, 'LineWidth', 1.8);
        grid on;
        xlabel('x');
        ylabel('u');
        title('Velocity');

        subplot(2,2,3);
        plot(x, p, 'LineWidth', 1.8);
        grid on;
        xlabel('x');
        ylabel('p');
        title('Pressure');

        subplot(2,2,4);
        plot(x, Mach, 'LineWidth', 1.8);
        grid on;
        xlabel('x');
        ylabel('Ma');
        title('Mach number');

        sgtitle('1D Euler equations solved by two-step Lax-Wendroff method');

        drawnow;
    end
end

%% ---------------------------------------------------------
%  5. 最终结果展示
%% ---------------------------------------------------------

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

figure('Color','w','Position',[150 150 850 600]);

subplot(3,1,1);
plot(x, rho, 'LineWidth', 2);
grid on;
ylabel('\rho');
title('Final solution of Sod shock tube');

subplot(3,1,2);
plot(x, u, 'LineWidth', 2);
grid on;
ylabel('u');

subplot(3,1,3);
plot(x, p, 'LineWidth', 2);
grid on;
xlabel('x');
ylabel('p');

%% =========================================================
%  局部函数 1：原始变量 -> 守恒变量
%% =========================================================

function U = primitive_to_conservative(rho, u, p, gamma)

    % 总能量：
    %
    %   E = p/(gamma-1) + 0.5*rho*u^2
    %
    % 第一项是内能
    % 第二项是动能

    E = p ./ (gamma - 1) + 0.5 * rho .* u.^2;

    U = zeros(3, length(rho));

    U(1,:) = rho;
    U(2,:) = rho .* u;
    U(3,:) = E;
end

%% =========================================================
%  局部函数 2：守恒变量 -> 原始变量
%% =========================================================

function [rho, u, p] = conservative_to_primitive(U, gamma)

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

    % 由状态方程反算压力
    %
    %   p = (gamma-1)*(E - 0.5*rho*u^2)

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

%% =========================================================
%  局部函数 3：计算 Euler 方程通量
%% =========================================================

function F = compute_flux(U, gamma)

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

    E = U(3,:);

    F = zeros(size(U));

    F(1,:) = rho .* u;
    F(2,:) = rho .* u.^2 + p;
    F(3,:) = u .* (E + p);
end