% Q2 Discontinuous Galerkin Method for 2D Compressible Euler Equations
% with Positivity-Preserving Limiter (Zhang-Shu JCP 2010)
% Benchmark: 2D Explosion Problem

clear all; close all; clc;

%% Physical and numerical parameters
gamma = 1.4;           % Ratio of specific heats
CFL = 0.3;             % CFL number
T_final = 0.2;         % Final time

% Domain: [0, 1] x [0, 1]
x_min = 0; x_max = 1;
y_min = 0; y_max = 1;

% Mesh parameters
Nx = 30;  % Number of elements in x
Ny = 30;  % Number of elements in y

% DG parameters
p = 2;    % Polynomial degree (Q2)
Np = (p+1)^2;  % Number of basis functions per element

fprintf('=== Q2 DG for 2D Euler Equations ===\n');
fprintf('Mesh: %d x %d elements\n', Nx, Ny);
fprintf('Polynomial degree: %d\n', p);
fprintf('Degrees of freedom per element: %d\n', Np);
fprintf('Total DOFs: %d\n', Nx*Ny*Np*4);

%% Generate rectangular mesh
dx = (x_max - x_min) / Nx;
dy = (y_max - y_min) / Ny;

% Element centers
x_centers = linspace(x_min + dx/2, x_max - dx/2, Nx);
y_centers = linspace(y_min + dy/2, y_max - dy/2, Ny);

%% Setup reference element and quadrature
% Gauss-Legendre quadrature points and weights in 1D
[xi_1d, w_1d] = lgwt(p+2, -1, 1);  % Higher order quadrature for accuracy

% Tensor product for 2D
[XI, ETA] = meshgrid(xi_1d, xi_1d);
xi_quad = XI(:);
eta_quad = ETA(:);
[W1, W2] = meshgrid(w_1d, w_1d);
w_quad = W1(:) .* W2(:);
N_quad = length(w_quad);

% Basis functions: Q2 polynomials (tensor product)
% Basis: phi_ij(xi, eta) = L_i(xi) * L_j(eta), i,j = 0,1,2
% where L_0, L_1, L_2 are Legendre polynomials

% 1D Legendre polynomials at quadrature points
L0 = @(x) ones(size(x));
L1 = @(x) x;
L2 = @(x) 0.5 * (3*x.^2 - 1);

L_basis = {L0, L1, L2};
dL_basis = {@(x) 0*x, @(x) ones(size(x)), @(x) 3*x};

% Evaluate basis functions at quadrature points
Phi = zeros(N_quad, Np);
dPhi_xi = zeros(N_quad, Np);
dPhi_eta = zeros(N_quad, Np);

idx = 1;
for j = 0:p
    for i = 0:p
        Phi(:, idx) = L_basis{i+1}(xi_quad) .* L_basis{j+1}(eta_quad);
        dPhi_xi(:, idx) = dL_basis{i+1}(xi_quad) .* L_basis{j+1}(eta_quad);
        dPhi_eta(:, idx) = L_basis{i+1}(xi_quad) .* dL_basis{j+1}(eta_quad);
        idx = idx + 1;
    end
end

% Mass matrix (on reference element)
M = Phi' * diag(w_quad) * Phi;
M_inv = inv(M);

%% Initialize solution
% Solution variables: U = [rho, rho*u, rho*v, E]
% Store as U(i, j, k, :) where i,j are element indices, k is basis index
U = zeros(Nx, Ny, Np, 4);

% Initial condition: 2D explosion problem
% High pressure region in center, low pressure outside
rho_in = 1.0;
p_in = 1.0;
rho_out = 0.125;
p_out = 0.1;
r_expl = 0.2;

for i = 1:Nx
    for j = 1:Ny
        x_c = x_centers(i);
        y_c = y_centers(j);

        % Distance from center
        r = sqrt((x_c - 0.5)^2 + (y_c - 0.5)^2);

        if r < r_expl
            rho = rho_in;
            u = 0;
            v = 0;
            p = p_in;
        else
            rho = rho_out;
            u = 0;
            v = 0;
            p = p_out;
        end

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

        % Set cell average (first basis function is constant)
        U(i, j, 1, 1) = rho;
        U(i, j, 1, 2) = rho * u;
        U(i, j, 1, 3) = rho * v;
        U(i, j, 1, 4) = E;
    end
end

fprintf('Initial condition set: 2D explosion\n');

%% Time stepping with SSP-RK3
t = 0;
iter = 0;
plot_interval = 50;

% Setup figure
fig = figure('Position', [100, 100, 1200, 500]);

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

while t < T_final
    % Compute time step
    dt = compute_timestep(U, dx, dy, gamma, CFL, Nx, Ny, Np);
    if t + dt > T_final
        dt = T_final - t;
    end

    % SSP-RK3 time integration
    U1 = U + dt * rhs_dg(U, dx, dy, gamma, Phi, dPhi_xi, dPhi_eta, w_quad, M_inv, Nx, Ny, Np);
    U1 = apply_positivity_limiter(U1, dx, dy, gamma, Phi, w_quad, Nx, Ny, Np, N_quad);

    U2 = 0.75*U + 0.25*(U1 + dt * rhs_dg(U1, dx, dy, gamma, Phi, dPhi_xi, dPhi_eta, w_quad, M_inv, Nx, Ny, Np));
    U2 = apply_positivity_limiter(U2, dx, dy, gamma, Phi, w_quad, Nx, Ny, Np, N_quad);

    U = (1/3)*U + (2/3)*(U2 + dt * rhs_dg(U2, dx, dy, gamma, Phi, dPhi_xi, dPhi_eta, w_quad, M_inv, Nx, Ny, Np));
    U = apply_positivity_limiter(U, dx, dy, gamma, Phi, w_quad, Nx, Ny, Np, N_quad);

    t = t + dt;
    iter = iter + 1;

    % Plotting
    if mod(iter, plot_interval) == 0 || t >= T_final
        [rho_plot, p_plot] = extract_solution(U, x_centers, y_centers, gamma, Nx, Ny);

        subplot(1, 2, 1);
        surf(x_centers, y_centers, rho_plot', 'EdgeColor', 'none');
        view(2); shading interp; colorbar;
        title(sprintf('Density at t = %.4f', t));
        xlabel('x'); ylabel('y');
        caxis([0.1, 1.2]);
        colormap(gca, jet);

        subplot(1, 2, 2);
        surf(x_centers, y_centers, p_plot', 'EdgeColor', 'none');
        view(2); shading interp; colorbar;
        title(sprintf('Pressure at t = %.4f', t));
        xlabel('x'); ylabel('y');
        caxis([0.05, 1.2]);
        colormap(gca, hot);

        drawnow;

        fprintf('Iter %d: t = %.4f, dt = %.2e\n', iter, t, dt);
    end
end

fprintf('\nSimulation complete!\n');
fprintf('Total iterations: %d\n', iter);

%% Helper Functions

function dt = compute_timestep(U, dx, dy, gamma, CFL, Nx, Ny, Np)
    % Compute stable time step
    max_speed = 0;
    for i = 1:Nx
        for j = 1:Ny
            % Extract cell average
            rho = U(i, j, 1, 1);
            u = U(i, j, 1, 2) / rho;
            v = U(i, j, 1, 3) / rho;
            E = U(i, j, 1, 4);
            p = (gamma - 1) * (E - 0.5 * rho * (u^2 + v^2));
            c = sqrt(gamma * p / rho);  % Sound speed

            speed = sqrt(u^2 + v^2) + c;
            max_speed = max(max_speed, speed);
        end
    end

    h = min(dx, dy);
    dt = CFL * h / max_speed;
end

function dU = rhs_dg(U, dx, dy, gamma, Phi, dPhi_xi, dPhi_eta, w_quad, M_inv, Nx, Ny, Np)
    % Compute RHS of DG scheme
    dU = zeros(size(U));
    N_quad = length(w_quad);

    for i = 1:Nx
        for j = 1:Ny
            % Volume integral
            rhs_vol = zeros(Np, 4);

            for q = 1:N_quad
                % Evaluate solution at quadrature point
                U_q = zeros(1, 4);
                for k = 1:Np
                    U_q = U_q + Phi(q, k) * squeeze(U(i, j, k, :))';
                end

                % Compute flux
                [F, G] = euler_flux(U_q, gamma);

                % Accumulate volume integral
                for k = 1:Np
                    rhs_vol(k, :) = rhs_vol(k, :) - w_quad(q) * ...
                        (dPhi_xi(q, k) * F + dPhi_eta(q, k) * G)' * (2/dx) * (2/dy);
                end
            end

            % Surface integrals (simplified - use neighbor averages)
            rhs_surf = compute_surface_flux(U, i, j, dx, dy, gamma, Nx, Ny, Np);

            % Apply inverse mass matrix
            for m = 1:4
                dU(i, j, :, m) = M_inv * (rhs_vol(:, m) + rhs_surf(:, m)) * (dx*dy/4);
            end
        end
    end
end

function [F, G] = euler_flux(U, gamma)
    % Compute Euler flux vectors
    rho = U(1);
    u = U(2) / rho;
    v = U(3) / rho;
    E = U(4);
    p = (gamma - 1) * (E - 0.5 * rho * (u^2 + v^2));

    F = [rho*u; rho*u^2 + p; rho*u*v; u*(E + p)];
    G = [rho*v; rho*u*v; rho*v^2 + p; v*(E + p)];
end

function flux = compute_surface_flux(U, i, j, dx, dy, gamma, Nx, Ny, Np)
    % Simplified surface flux using Lax-Friedrichs
    flux = zeros(Np, 4);

    % Get cell averages (first basis function)
    U_center = squeeze(U(i, j, 1, :));

    % Left face (i-1)
    if i > 1
        U_left = squeeze(U(i-1, j, 1, :));
    else
        U_left = U_center;  % Boundary
    end

    % Right face (i+1)
    if i < Nx
        U_right = squeeze(U(i+1, j, 1, :));
    else
        U_right = U_center;  % Boundary
    end

    % Bottom face (j-1)
    if j > 1
        U_bottom = squeeze(U(i, j-1, 1, :));
    else
        U_bottom = U_center;  % Boundary
    end

    % Top face (j+1)
    if j < Ny
        U_top = squeeze(U(i, j+1, 1, :));
    else
        U_top = U_center;  % Boundary
    end

    % Lax-Friedrichs flux (simplified)
    alpha = 2.0;  % Numerical viscosity

    [F_c, G_c] = euler_flux(U_center', gamma);
    [F_l, ~] = euler_flux(U_left', gamma);
    [F_r, ~] = euler_flux(U_right', gamma);
    [~, G_b] = euler_flux(U_bottom', gamma);
    [~, G_t] = euler_flux(U_top', gamma);

    % Approximate surface contribution (only to first mode)
    flux(1, :) = (2/dx) * ((F_r - F_l)'/2 + alpha*(U_center' - (U_left' + U_right')/2)) + ...
                 (2/dy) * ((G_t - G_b)'/2 + alpha*(U_center' - (U_bottom' + U_top')/2));
end

function U_lim = apply_positivity_limiter(U, dx, dy, gamma, Phi, w_quad, Nx, Ny, Np, N_quad)
    % Zhang-Shu positivity-preserving limiter (simplified version)
    U_lim = U;
    eps_p = 1e-13;  % Small positive number

    for i = 1:Nx
        for j = 1:Ny
            % Cell average
            U_avg = squeeze(U(i, j, 1, :));
            rho_avg = U_avg(1);
            E_avg = U_avg(4);

            % Check positivity at quadrature points
            theta = 1.0;

            for q = 1:N_quad
                % Evaluate at quadrature point
                U_q = zeros(4, 1);
                for k = 1:Np
                    U_q = U_q + Phi(q, k) * squeeze(U(i, j, k, :));
                end

                rho_q = U_q(1);
                E_q = U_q(4);

                % Check density positivity
                if rho_q < eps_p
                    theta_rho = (rho_avg - eps_p) / (rho_avg - rho_q + 1e-14);
                    theta = min(theta, abs(theta_rho));
                end

                % Check pressure positivity
                u_q = U_q(2) / (rho_q + 1e-14);
                v_q = U_q(3) / (rho_q + 1e-14);
                p_q = (gamma - 1) * (E_q - 0.5 * rho_q * (u_q^2 + v_q^2));

                if p_q < eps_p
                    % Simplified pressure limiter
                    theta = min(theta, 0.9);
                end
            end

            % Apply limiter
            if theta < 1.0
                for k = 2:Np  % Keep cell average, scale higher modes
                    U_lim(i, j, k, :) = theta * U(i, j, k, :);
                end
            end
        end
    end
end

function [rho_plot, p_plot] = extract_solution(U, x_centers, y_centers, gamma, Nx, Ny)
    % Extract cell averages for plotting
    rho_plot = zeros(Nx, Ny);
    p_plot = zeros(Nx, Ny);

    for i = 1:Nx
        for j = 1:Ny
            rho = U(i, j, 1, 1);
            u = U(i, j, 1, 2) / rho;
            v = U(i, j, 1, 3) / rho;
            E = U(i, j, 1, 4);
            p = (gamma - 1) * (E - 0.5 * rho * (u^2 + v^2));

            rho_plot(i, j) = rho;
            p_plot(i, j) = p;
        end
    end
end

function [x, w] = lgwt(N, a, b)
    % Legendre-Gauss quadrature points and weights
    N = N - 1;
    N1 = N + 1;
    N2 = N + 2;

    xu = linspace(-1, 1, N1)';
    y = cos((2*(0:N)' + 1) * pi / (2*N + 2)) + (0.27/N1) * sin(pi*xu*N/N2);

    L = zeros(N1, N2);
    Lp = zeros(N1, N2);
    y0 = 2;

    while max(abs(y - y0)) > eps
        L(:, 1) = 1;
        L(:, 2) = y;

        for k = 2:N1
            L(:, k+1) = ((2*k - 1) * y .* L(:, k) - (k - 1) * L(:, k-1)) / k;
        end

        Lp = N2 * (L(:, N1) - y .* L(:, N2)) ./ (1 - y.^2);
        y0 = y;
        y = y0 - L(:, N2) ./ Lp;
    end

    x = (a * (1 - y) + b * (1 + y)) / 2;
    w = (b - a) ./ ((1 - y.^2) .* Lp.^2) * (N2/N1)^2;
end
