%% Double-Divergence PDE: Full Discretization of Nondivergence Form
% Including second-order (central), first-order (upwind), and zero-order terms
clear; clc; close all;

%% Grid setup
N = 300;
h = 2 / (N + 1);
x = linspace(-1, 1, N+2);
[X, Y] = meshgrid(x, x);

% Extended domain for mollification
x_ext = linspace(-2, 2, 4*N + 5);
[Xe, Ye] = meshgrid(x_ext, x_ext);
len_ext = length(x_ext);
start_idx = floor((len_ext - (N+2))/2) + 1;
idx_center = start_idx : start_idx + N + 1;

%% Exact solution (nonsmooth) and numerical derivatives
U_exact = ((1 - abs(X)).^2) .* ((1 - abs(Y)).^2);

% Numerical derivatives via central differences
[u_y, u_x] = gradient(U_exact, h);           % First derivatives
[u_yy, u_y] = gradient(u_y, h);              % Second derivatives
[~,    u_xx] = gradient(u_x, h);             % Second derivatives




%% Raw coefficients
a11_raw = 2 - abs(Xe);
a22_raw = 2 - abs(Ye);

%% Mollification
epsilon = 0.01;
r = ceil(3 * epsilon / h);
[Xr, Yr] = meshgrid(-r:r, -r:r);
rho = exp(-(Xr.^2 + Yr.^2) * h^2 / epsilon^2);
rho = rho / sum(rho(:));

a11 = conv2(a11_raw, rho, 'same');
a22 = conv2(a22_raw, rho, 'same');
a11c = a11(idx_center, idx_center);
a22c = a22(idx_center, idx_center);

%% Compute b^eps and c^eps (first and zero-order coefficients)
[~, da11_dx] = gradient(a11c, h);
[da22_dy, ~] = gradient(a22c, h);

[~, d2a11_dx2] = gradient(da11_dx, h);
d2a22_dy2 = gradient(da22_dy, h);

b1 = da11_dx;
b2 = da22_dy;
c = d2a11_dx2 + d2a22_dy2;

%% Compute RHS: f = Tr(A D^2 u) + b.Du + c u
f = a11c .* u_xx + a22c .* u_yy + b1 .* u_x + b2 .* u_y + c .* U_exact;

%% Discretization: full operator with second, first, and zero-order terms
idx = @(i,j) (j-1)*N + i;
A = spalloc(N*N, N*N, 9*N*N);
b = reshape(f(2:end-1,2:end-1), N*N, 1);

for j = 1:N
    for i = 1:N
        xi = i+1; yi = j+1; k = idx(i,j);

        % Second-order terms (central diff)
        a11p = a11c(xi+1, yi); a11m = a11c(xi-1, yi);
        a22p = a22c(xi, yi+1); a22m = a22c(xi, yi-1);
        A(k, k) = A(k, k) - (a11p + a11m + a22p + a22m) / h^2;
        if i > 1, A(k, idx(i-1,j)) = A(k, idx(i-1,j)) + a11m / h^2; end
        if i < N, A(k, idx(i+1,j)) = A(k, idx(i+1,j)) + a11p / h^2; end
        if j > 1, A(k, idx(i,j-1)) = A(k, idx(i,j-1)) + a22m / h^2; end
        if j < N, A(k, idx(i,j+1)) = A(k, idx(i,j+1)) + a22p / h^2; end

        % First-order terms (upwind)
        b1_ij = b1(xi, yi); b2_ij = b2(xi, yi);
        if b1_ij >= 0 && i > 1
            A(k, idx(i-1,j)) = A(k, idx(i-1,j)) - b1_ij / h;
        elseif i < N
            A(k, idx(i+1,j)) = A(k, idx(i+1,j)) + b1_ij / h;
        end
        if b2_ij >= 0 && j > 1
            A(k, idx(i,j-1)) = A(k, idx(i,j-1)) - b2_ij / h;
        elseif j < N
            A(k, idx(i,j+1)) = A(k, idx(i,j+1)) + b2_ij / h;
        end

        % Zero-order term
        A(k,k) = A(k,k) + c(xi, yi);
    end
end

%% Solve
u_vec = A \ b;
U_num = zeros(N+2);
U_num(2:end-1, 2:end-1) = reshape(u_vec, N, N);

%% Error and display
E = abs(U_num - U_exact);
max_error = max(E(:));
fprintf('%10s\t%20s\n', 'Epsilon', 'MaxError');
fprintf('%10.8f\t%20.20f\n', epsilon, max_error);

%% Plot: absolute error heatmap
figure('Name','Absolute Error');
imagesc(x, x, E); axis image off;
colormap autumn; colorbar;


%% Plot: exact solution with sparse mesh
figure('Name','Exact Solution');
skip = 12;
x_s = x(1:skip:end); y_s = x_s;
[Xs, Ys] = meshgrid(x_s, y_s);
U_s = ((1 - abs(Xs)).^2) .* ((1 - abs(Ys)).^2);
surf(Xs, Ys, U_s, 'EdgeColor', [0.3 0.3 0.3], 'LineStyle', '-');
colormap autumn;
xlabel('x'); ylabel('y'); zlabel('u(x,y)');
view(45,30); grid on;
