clear; clc; close all;

%% =========================
% 参数
%% =========================
Nx = 60; Ny = 60;
Lx = 1; Ly = 1;

dx = Lx/(Nx-1);
dy = Ly/(Ny-1);

[x,y] = meshgrid(linspace(0,Lx,Nx), linspace(0,Ly,Ny));

%% =========================
% 初始条件类型
% 1: 随机
% 2: 平滑随机（推荐）
% 3: 单峰
% 4: 多峰
% 5: 线性场
%% =========================
IC_type = 2;

switch IC_type
    case 1
        u = rand(Nx,Ny);
    case 2
        u = rand(Nx,Ny);
        u = smoothdata(u,1,'gaussian',10);
        u = smoothdata(u,2,'gaussian',10);
    case 3
        u = exp(-40*((x-0.5).^2 + (y-0.5).^2));
    case 4
        u = exp(-60*((x-0.3).^2 + (y-0.3).^2)) + ...
            exp(-60*((x-0.7).^2 + (y-0.7).^2));
    case 5
        u = x + y;  % Laplace 精确解
end

%% =========================
% 边界条件类型
% 1: Dirichlet
% 2: Neumann（零梯度）
% 3: 混合
% 4: 周期
%% =========================
BC_type = 1;

applyBC = @(u) apply_boundary(u, x, BC_type);

u = applyBC(u);

%% =========================
% SOR 参数（关键）
%% =========================
omega = 1.7;   % 1=GS，1.5~1.9=SOR

%% =========================
% 图像初始化
%% =========================
figure('Color','w','Position',[10 100 1200 500]);

subplot(1,3,1)
h = imagesc(u'); axis equal tight
colorbar
title('迭代过程')

subplot(1,3,2)
hsurf = surf(x,y,u');
shading interp
title('三维解')

subplot(1,3,3)
h_res = semilogy(1,1,'LineWidth',2);
hold on; grid on;
title('残差收敛')
xlabel('迭代步'); ylabel('残差')

drawnow

%% =========================
% 迭代
%% =========================
maxIter = 5000;
resHist = [];

for iter = 1:maxIter
    
    u_old = u;
    
    for i = 2:Nx-1
        for j = 2:Ny-1
            
            u_new = 0.25*( ...
                u(i+1,j) + u(i-1,j) + ...
                u(i,j+1) + u(i,j-1) );
            
            % SOR 更新
            u(i,j) = (1-omega)*u(i,j) + omega*u_new;
        end
    end
    
    % 应用边界
    u = applyBC(u);
    
    %% -------- 残差 --------
    res = max(abs(u(:)-u_old(:)));
    resHist = [resHist, res];
    
    %% -------- 可视化 --------
    if mod(iter,1)==0
        
        subplot(1,3,1)
        set(h,'CData',u')
        title(sprintf('Iter = %d | BC = %d',iter,BC_type))
        
        subplot(1,3,2)
        set(hsurf,'ZData',u')
        
        subplot(1,3,3)
        set(h_res,'XData',1:length(resHist),'YData',resHist)
        
        drawnow limitrate
        pause(0.01)
    end
    
    %% -------- 收敛 --------
    if iter > 50 && res < 1e-6
        break
    end
end

%% =========================
% 边界函数
%% =========================
function u = apply_boundary(u, x, BC_type)

[Nx,Ny] = size(u);

switch BC_type
    
    %% -------- Dirichlet --------
    case 1
        u(:,end) = sin(2*pi*x(:,end)); % 上
        u(:,1)   = cos(2*pi*x(:,1));   % 下
        u(1,:)   = 2;
        u(end,:) = 1;
        
    %% -------- Neumann --------
    case 2
        % du/dn = 0
        u(:,1)   = u(:,2);
        u(:,end) = u(:,end-1);
        u(1,:)   = u(2,:);
        u(end,:) = u(end-1,:);
        
    %% -------- 混合 --------
    case 3
        u(:,end) = sin(2*pi*x(:,end)); % Dirichlet
        u(:,1)   = u(:,2);             % Neumann
        u(1,:)   = 1;
        u(end,:) = u(end-1,:);
        
    %% -------- 周期 --------
    case 4
        u(1,:)   = u(end-1,:);
        u(end,:) = u(2,:);
        u(:,1)   = u(:,end-1);
        u(:,end) = u(:,2);
        
end
end