// 隐式扩散求解器的简化版 - 用迭代法近似
module diffusion_iterative #(
    parameter WIDTH = 64,
    parameter HEIGHT = 64,
    parameter ITERATIONS = 4
)(
    input  wire              clk,
    input  wire              rst_n,
    input  wire              enable,
    input  wire [15:0]       conc [0:WIDTH*HEIGHT-1],
    output wire [15:0]       conc_out [0:WIDTH*HEIGHT-1]
);
    // Jacobi迭代：c_new[i] = (c[left]+c[right]+c[up]+c[down]) / 4
    reg [15:0] buf_a [0:WIDTH*HEIGHT-1];
    reg [15:0] buf_b [0:WIDTH*HEIGHT-1];
    integer iter, idx, x, y;
    
    always @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            // 初始化
        end else if (enable) begin
            // 复制输入到buf_a
            for (idx = 0; idx < WIDTH*HEIGHT; idx = idx + 1)
                buf_a[idx] <= conc[idx];
            // 迭代
            for (iter = 0; iter < ITERATIONS; iter = iter + 1) begin
                for (idx = 0; idx < WIDTH*HEIGHT; idx = idx + 1) begin
                    x = idx % WIDTH; y = idx / WIDTH;
                    buf_b[idx] = (buf_a[(y>0?(y-1)*WIDTH+x:(HEIGHT-1)*WIDTH+x)] +
                                  buf_a[(y<HEIGHT-1?(y+1)*WIDTH+x:x)] +
                                  buf_a[y*WIDTH+(x>0?x-1:WIDTH-1)] +
                                  buf_a[y*WIDTH+(x<WIDTH-1?x+1:0)]) >> 2;
                end
                for (idx = 0; idx < WIDTH*HEIGHT; idx = idx + 1)
                    buf_a[idx] <= buf_b[idx];
            end
        end
    end
endmodule