// ============================================================================
// langton_ant_general.v - 广义Langton蚂蚁
// 支持任意颜色数和转向规则
// ============================================================================
module langton_ant_general #(
    parameter WIDTH  = 128,
    parameter HEIGHT = 128,
    parameter COLORS = 4,              // 颜色数
    parameter [COLORS*2-1:0] RULE = 8'b01_00_01_00  // R=01,L=00
)(
    input  wire             clk,
    input  wire             rst_n,
    input  wire             enable,
    input  wire             init,
    output wire [1:0]       ant_dir,
    output wire [7:0]       ant_x,
    output wire [7:0]       ant_y,
    output wire [31:0]      step_count
);

    // 每格存储颜色索引（log2(COLORS)位）
    localparam CW = $clog2(COLORS);
    reg [CW-1:0] grid [0:WIDTH*HEIGHT-1];

    reg [7:0]  ax, ay;
    reg [1:0]  dir;
    reg [31:0] steps;

    wire [CW-1:0] idx = ay * WIDTH + ax;
    wire [CW-1:0] current_color = grid[idx];

    // 规则查找：获取转向
    wire turn_right = RULE[current_color * 2 +: 2] == 2'b01;

    // 方向更新
    wire [1:0] next_dir = turn_right ? (dir + 2'd1) : (dir - 2'd1);

    // 下一颜色
    wire [CW-1:0] next_color = (current_color == COLORS - 1) ? {CW{1'b0}} : current_color + 1'd1;

    // 位置更新
    reg [7:0] next_x, next_y;
    always @(*) begin
        next_x = ax; next_y = ay;
        case (next_dir)
            2'd0: next_y = (ay == 0) ? HEIGHT - 1 : ay - 1;
            2'd1: next_x = (ax == WIDTH - 1) ? 0 : ax + 1;
            2'd2: next_y = (ay == HEIGHT - 1) ? 0 : ay + 1;
            2'd3: next_x = (ax == 0) ? WIDTH - 1 : ax - 1;
        endcase
    end

    always @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            ax <= WIDTH/2; ay <= HEIGHT/2;
            dir <= 2'd0; steps <= 32'd0;
        end else if (init) begin
            ax <= WIDTH/2; ay <= HEIGHT/2;
            dir <= 2'd0; steps <= 32'd0;
            for (integer i = 0; i < WIDTH*HEIGHT; i = i + 1)
                grid[i] <= {CW{1'b0}};
        end else if (enable) begin
            grid[idx] <= next_color;
            ax <= next_x; ay <= next_y;
            dir <= next_dir;
            steps <= steps + 32'd1;
        end
    end

    assign ant_dir    = dir;
    assign ant_x     = ax;
    assign ant_y     = ay;
    assign step_count = steps;

endmodule