%% =========================================================
%  准一维喷管等熵流动
%  MacCormack 显式预测-校正格式（守恒形式）
%
%  推进变量：
%       U1 = rho*A
%       U2 = rho*V*A
%       U3 = rho*A*( T/(gamma-1) + gamma/2*V^2 )
%
%  说明：
%  本程序保持原非守恒代码的无量纲体系：V 用 a0 归一化，M=V/sqrt(T)。
%  因此动量方程中的压力项含 1/gamma，能量变量中动能项为 gamma/2*V^2。
%  这与 Anderson 准一维喷管算例的常用无量纲守恒形式一致。
%
%  守恒形式：
%       dU/dt + dF/dx = J
%
%  通量与源项：
%       F1 = U2
%       F2 = rho*V^2*A + p*A/gamma
%       F3 = U2*(U3/U1 + T)
%       J  = [0, p/gamma*dA/dx, 0]^T
%
%  边界条件：
%       入口（亚声）：rho=1, T=1, V 外推
%       出口（超声）：rho, T, V 全部外推
%% =========================================================

clear; clc; close all;

%% ===================== 1. 基本参数 =====================
gamma = 1.4;
CFL   = 0.5;
xL    = 0.0;
xR    = 3.0;
N     = 101;
maxIter  = 5000;
tol      = 1e-6;
plotEvery = 10;

%% ===================== 2. 网格与喷管几何 =====================
x  = linspace(xL,xR,N)';
dx = x(2)-x(1);

A = 1+2.2*(x-1.5).^2;
dAdx = gradient(A,dx);

[Ast,ithroat] = min(A);

%% ===================== 3. 初始条件 =====================
rho = 1-0.3146*x;
T   = 1-0.2314*x;
V   = (0.1+1.09*x).*sqrt(T);

rho = max(rho,1e-8);
T   = max(T,1e-8);

U = primitiveToConservative(rho,V,T,A,gamma);
[U,rho,V,T] = applySupersonicOutletBC(U,A,gamma);

p = rho.*T;
M = V./sqrt(T);

%% ===================== 4. 解析等熵解 =====================
M_exact = zeros(N,1);
for i = 1:N
    areaRatio = A(i)/Ast;
    if i < ithroat
        M_exact(i) = solveAreaMach(areaRatio,gamma,'subsonic');
    elseif i > ithroat
        M_exact(i) = solveAreaMach(areaRatio,gamma,'supersonic');
    else
        M_exact(i) = 1.0;
    end
end

T_exact   = 1./(1+(gamma-1)/2.*M_exact.^2);
rho_exact = T_exact.^(1/(gamma-1));
p_exact   = T_exact.^(gamma/(gamma-1));
V_exact   = M_exact.*sqrt(T_exact);

%% ===================== 5. 残差与时间记录 =====================
res_rho = zeros(maxIter,1);
res_V   = zeros(maxIter,1);
res_T   = zeros(maxIter,1);
time_history = zeros(maxIter,1);
physicalTime = 0.0;

%% ===================== 6. 动态图形窗口 =====================
fig = figure('Color','w','Position',[100 60 1400 820]);
htitle = sgtitle('Conservative Quasi-1D Nozzle Flow','FontWeight','bold');

subplot(2,3,1);
yyaxis left
plot(x,A,'k:','LineWidth',1.5); hold on;
ylabel('Area A(x)')
yyaxis right
hM = plot(x,M,'b-','LineWidth',2); hold on;
plot(x,M_exact,'r--','LineWidth',1.2);
ylabel('Mach number M')
xlabel('x')
title('Area & Mach Number')
legend('A(x)','M numerical','M exact','Location','best')
grid on;

subplot(2,3,2);
hRho = plot(x,rho,'b-','LineWidth',2); hold on;
plot(x,rho_exact,'r--','LineWidth',1.2);
xlabel('x'); ylabel('\rho / \rho_0'); title('Density Distribution')
legend('Numerical','Exact','Location','best'); grid on;

subplot(2,3,3);
hT = plot(x,T,'b-','LineWidth',2); hold on;
plot(x,T_exact,'r--','LineWidth',1.2);
xlabel('x'); ylabel('T / T_0'); title('Temperature Distribution')
legend('Numerical','Exact','Location','best'); grid on;

subplot(2,3,4);
hV = plot(x,V,'b-','LineWidth',2); hold on;
plot(x,V_exact,'r--','LineWidth',1.2);
xlabel('x'); ylabel('Velocity V / a_0'); title('Velocity Distribution')
legend('Numerical','Exact','Location','best'); grid on;

subplot(2,3,5);
hP = plot(x,p,'b-','LineWidth',2); hold on;
plot(x,p_exact,'r--','LineWidth',1.2);
xlabel('x'); ylabel('p / p_0'); title('Pressure Distribution')
legend('Numerical','Exact','Location','best'); grid on;

subplot(2,3,6);
hRes1 = semilogy(1,1,'LineWidth',1.5); hold on;
hRes2 = semilogy(1,1,'LineWidth',1.5);
hRes3 = semilogy(1,1,'LineWidth',1.5);
xlabel('Iteration'); ylabel('Residual'); title('Convergence History')
legend('\rho','V','T','Location','best'); grid on;

%% ===================== 7. MacCormack 主循环 =====================
for iter = 1:maxIter

    [rho,V,T,p,M] = conservativeToPrimitive(U,A,gamma);

    a = sqrt(T);
    dt_local = CFL*dx./(a+abs(V));
    dt = min(dt_local);
    physicalTime = physicalTime+dt;
    time_history(iter) = physicalTime;

    rho_old = rho;
    V_old   = V;
    T_old   = T;
    U_old   = U;

    % ---------- 7.1 预测步：前向差分 ----------
    [F,J] = computeFluxSource(U,A,dAdx,gamma);
    U_bar = U;

    for i = 2:N-1
        U_bar(i,:) = U(i,:) - dt/dx*(F(i+1,:)-F(i,:)) + dt*J(i,:);
    end

    [U_bar,~,~,~] = applySupersonicOutletBC(U_bar,A,gamma);

    % ---------- 7.2 校正步：后向差分 ----------
    [F_bar,J_bar] = computeFluxSource(U_bar,A,dAdx,gamma);
    U_new = U_old;

    for i = 2:N-1
        U_new(i,:) = 0.5*( ...
            U_old(i,:) + U_bar(i,:) ...
            - dt/dx*(F_bar(i,:)-F_bar(i-1,:)) ...
            + dt*J_bar(i,:) );
    end

    [U,rho,V,T] = applySupersonicOutletBC(U_new,A,gamma);
    p = rho.*T;
    M = V./sqrt(T);

    % ---------- 7.3 残差 ----------
    res_rho(iter) = max(abs(rho-rho_old));
    res_V(iter)   = max(abs(V-V_old));
    res_T(iter)   = max(abs(T-T_old));
    res_max = max([res_rho(iter),res_V(iter),res_T(iter)]);

    if any(~isfinite(U(:))) || any(rho<=0) || any(T<=0)
        warning('计算发散：iter = %d',iter);
        break;
    end

    % ---------- 7.4 动态画图 ----------
    if mod(iter,plotEvery)==0 || iter==1
        set(hM,'YData',M);
        set(hRho,'YData',rho);
        set(hT,'YData',T);
        set(hV,'YData',V);
        set(hP,'YData',p);
        set(hRes1,'XData',1:iter,'YData',res_rho(1:iter));
        set(hRes2,'XData',1:iter,'YData',res_V(1:iter));
        set(hRes3,'XData',1:iter,'YData',res_T(1:iter));
        set(htitle,'String',sprintf( ...
            'Conservative Quasi-1D Nozzle Flow | Iter = %d | t = %.4f', ...
            iter,physicalTime));
        drawnow limitrate;
        if iter==1
            waitforbuttonpress;
        end
        pause(0.01);
    end

    if res_max < tol
        fprintf('计算收敛！iter = %d, time = %.6f, max residual = %.3e\n', ...
            iter,physicalTime,res_max);
        break;
    end
end

%% ===================== 8. 输出最终结果 =====================
finalIter = iter;
res_rho = res_rho(1:finalIter);
res_V   = res_V(1:finalIter);
res_T   = res_T(1:finalIter);
time_history = time_history(1:finalIter);

fprintf('\n======== 计算结束 ========\n');
fprintf('总迭代步数      = %d\n',finalIter);
fprintf('最终物理时间    = %.6f\n',physicalTime);
fprintf('出口马赫数      = %.4f\n',M(end));
fprintf('喉部马赫数      = %.4f\n',M(ithroat));
fprintf('最终最大残差    = %.3e\n',max([res_rho(end),res_V(end),res_T(end)]));

%% ===================== 9. 喷管内二维 contour 展示 =====================
Ny = 101;
yScale = 0.35;
h = yScale*A/2;
eta = linspace(-1,1,Ny);
[X2D,ETA2D] = meshgrid(x,eta);
H2D = repmat(h',Ny,1);
Y2D = ETA2D.*H2D;

rho2D = repmat(rho',Ny,1);
T2D   = repmat(T',Ny,1);
p2D   = repmat(p',Ny,1);
M2D   = repmat(M',Ny,1);

mask = abs(Y2D)>H2D;
rho2D(mask)=NaN; T2D(mask)=NaN; p2D(mask)=NaN; M2D(mask)=NaN;

figure('Color','w','Position',[80 80 1400 720]);
tl = tiledlayout(2,2);
tl.TileSpacing = 'compact';
tl.Padding = 'compact';
nLevel = 50;

ax1 = nexttile;
contourf(X2D,Y2D,T2D,nLevel,'LineColor','none'); hold on;
plot(x,h,'k-','LineWidth',1.6); plot(x,-h,'k-','LineWidth',1.6);
xlabel('x'); ylabel('y'); title('Temperature'); cb1=colorbar; cb1.Label.String='T / T_0';
grid on; box on; axis tight; pbaspect([3.2 1 1]);

ax2 = nexttile;
contourf(X2D,Y2D,p2D,nLevel,'LineColor','none'); hold on;
plot(x,h,'k-','LineWidth',1.6); plot(x,-h,'k-','LineWidth',1.6);
xlabel('x'); ylabel('y'); title('Pressure'); cb2=colorbar; cb2.Label.String='p / p_0';
grid on; box on; axis tight; pbaspect([3.2 1 1]);

ax3 = nexttile;
contourf(X2D,Y2D,rho2D,nLevel,'LineColor','none'); hold on;
plot(x,h,'k-','LineWidth',1.6); plot(x,-h,'k-','LineWidth',1.6);
xlabel('x'); ylabel('y'); title('Density'); cb3=colorbar; cb3.Label.String='\rho / \rho_0';
grid on; box on; axis tight; pbaspect([3.2 1 1]);

ax4 = nexttile;
contourf(X2D,Y2D,M2D,nLevel,'LineColor','none'); hold on;
plot(x,h,'k-','LineWidth',1.6); plot(x,-h,'k-','LineWidth',1.6);
xline(x(ithroat),'k--','LineWidth',1.2);
xlabel('x'); ylabel('y'); title('Mach Number'); cb4=colorbar; cb4.Label.String='M';
grid on; box on; axis tight; pbaspect([3.2 1 1]);

colormap(turbo);
set([ax1 ax2 ax3 ax4],'FontName','Times New Roman','FontSize',12,'LineWidth',1.0,'Layer','top');
sgtitle('Conservative Quasi-1D Nozzle Flow: Final Contour Visualization', ...
    'FontWeight','bold','FontSize',16);

%% =========================================================
%  局部函数
%% =========================================================
function U = primitiveToConservative(rho,V,T,A,gamma)
U = zeros(length(rho),3);
U(:,1) = rho.*A;
U(:,2) = rho.*V.*A;
U(:,3) = rho.*A.*(T/(gamma-1)+gamma/2*V.^2);
end

function [rho,V,T,p,M] = conservativeToPrimitive(U,A,gamma)
rho = U(:,1)./A;
V   = U(:,2)./U(:,1);
E   = U(:,3)./U(:,1);
T   = (gamma-1)*(E-gamma/2*V.^2);
rho = max(rho,1e-10);
T   = max(T,1e-10);
p   = rho.*T;
M   = V./sqrt(T);
end

function [F,J] = computeFluxSource(U,A,dAdx,gamma)
[rho,V,T,p,~] = conservativeToPrimitive(U,A,gamma);
F = zeros(size(U));
J = zeros(size(U));
F(:,1) = U(:,2);
F(:,2) = rho.*V.^2.*A + p.*A/gamma;
F(:,3) = U(:,2).*(U(:,3)./U(:,1)+T);
J(:,2) = p.*dAdx/gamma;
end

function [U,rho,V,T] = applySupersonicOutletBC(U,A,gamma)
[rho,V,T,~,~] = conservativeToPrimitive(U,A,gamma);
N = length(rho);

rho(1) = 1.0;
T(1)   = 1.0;
V(1)   = 2*V(2)-V(3);

rho(N) = 2*rho(N-1)-rho(N-2);
T(N)   = 2*T(N-1)-T(N-2);
V(N)   = 2*V(N-1)-V(N-2);

rho = max(rho,1e-8);
T   = max(T,1e-8);
U = primitiveToConservative(rho,V,T,A,gamma);
end

function F = areaMachFunction(M,gamma)
F = (1./M).* ...
    ((2/(gamma+1)).*(1+(gamma-1)/2.*M.^2)).^((gamma+1)/(2*(gamma-1)));
end

function M = solveAreaMach(areaRatio,gamma,branch)
fun = @(M) areaMachFunction(M,gamma)-areaRatio;
switch lower(branch)
    case 'subsonic'
        M = fzero(fun,[1e-8,1-1e-8]);
    case 'supersonic'
        M = fzero(fun,[1+1e-8,20]);
    otherwise
        error('branch 必须为 subsonic 或 supersonic');
end
end