`timescale 1ns / 1ps

module meggitt_decoder_comp (
    input wire clk,
    input wire rst_n,
    input wire start,
    input wire [7:0] data_in,
    input wire data_valid,
    output reg done,
    output wire [7:0] data_out,
    output reg data_out_valid,
    output reg [15:0] bit_counter,
    output reg [7:0] error_count
);

    localparam N = 18432;
    localparam K = 16384;
    localparam R = 2048;
    localparam B = 256;

    localparam IDLE              = 4'd0;
    localparam LOAD              = 4'd1;
    localparam CALC              = 4'd2;
    localparam SCAN              = 4'd3;
    localparam FLIP              = 4'd4;
    localparam UPDATE_SYND_FETCH = 4'd5;
    localparam UPDATE_SYND_SHIFT = 4'd6;
    localparam OUTPUT            = 4'd7;

    reg [3:0] state;
    reg [14:0] counter;
    reg [4:0] pass_count; // max 20 iterations

    reg [255:0] syndrome [0:7];
    reg [N-1:0] codeword;

    reg [255:0] poly_rom [0:511];
    initial begin
        $readmemh("polys.mem", poly_rom);
    end

    reg [255:0] current_sig [0:7];

    reg [11:0] total_score;
    reg [11:0] current_synd_weight;
    reg [11:0] min_weight;
    reg [14:0] best_bit;
    
    integer i, pj;
    reg [255:0] xored;

    always @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            state <= IDLE;
            counter <= 0;
            done <= 0;
            data_out_valid <= 0;
            error_count <= 0;
            pass_count <= 0;
            bit_counter <= 0;
            min_weight <= 12'hFFF;
            best_bit <= 0;
            current_synd_weight <= 0;
            for (i = 0; i < 8; i = i + 1) syndrome[i] <= 0;
        end else begin
            case (state)
                IDLE: begin
                    if (start) begin
                        state <= LOAD;
                        counter <= 0;
                        done <= 0;
                        error_count <= 0;
                        pass_count <= 0;
                        for (i = 0; i < 8; i = i + 1) syndrome[i] <= 0;
                    end
                end

                LOAD: begin
                    if (data_valid) begin
                        codeword <= {codeword[N-9:0], data_in};
                        counter <= counter + 8;
                        if (counter == N - 8) begin
                            state <= CALC;
                            counter <= 0;
                        end
                    end
                end

                CALC: begin
                    if (counter < K) begin
                        if (counter[7:0] == 0) begin
                            for (i = 0; i < 8; i = i + 1) 
                                current_sig[i] <= poly_rom[{counter[14:8], 3'b0} + i];
                        end else begin
                            for (i = 0; i < 8; i = i + 1) 
                                current_sig[i] <= {current_sig[i][254:0], current_sig[i][255]};
                        end
                        
                        if (codeword[N-1 - counter]) begin
                            for (i = 0; i < 8; i = i + 1)
                                syndrome[i] <= syndrome[i] ^ current_sig[i];
                        end
                    end else if (counter < N) begin
                        if (codeword[N-1 - counter]) begin
                            syndrome[(counter-K) >> 8][(counter-K) & 8'hFF] <= 
                                syndrome[(counter-K) >> 8][(counter-K) & 8'hFF] ^ 1'b1;
                        end
                    end

                    counter <= counter + 1;
                    if (counter == N - 1) begin
                        state <= SCAN;
                        counter <= 0;
                        bit_counter <= 0;
                    end
                end

                SCAN: begin
                    if (counter < K) begin
                        if (counter[7:0] == 0) begin
                            for (i = 0; i < 8; i = i + 1) 
                                current_sig[i] <= poly_rom[{counter[14:8], 3'b0} + i];
                        end else begin
                            for (i = 0; i < 8; i = i + 1) 
                                current_sig[i] <= {current_sig[i][254:0], current_sig[i][255]};
                        end

                        // Combinational popcount
                        total_score = 0;
                        for (i = 0; i < 8; i = i + 1) begin
                            xored = syndrome[i] ^ current_sig[i];
                            for (pj = 0; pj < 256; pj = pj + 1) begin
                                if (xored[pj]) total_score = total_score + 1;
                            end
                        end
                        
                        if (counter == 0) begin
                            min_weight <= total_score;
                            best_bit <= 0;
                        end else if (total_score < min_weight) begin
                            min_weight <= total_score;
                            best_bit <= counter;
                        end
                    end
                    
                    counter <= counter + 1;
                    if (counter == K - 1) begin
                        state <= FLIP;
                        counter <= 0;
                        
                        // Calculate current syndrome weight
                        current_synd_weight = 0;
                        for (i = 0; i < 8; i = i + 1) begin
                            for (pj = 0; pj < 256; pj = pj + 1) begin
                                if (syndrome[i][pj]) current_synd_weight = current_synd_weight + 1;
                            end
                        end
                    end
                end
                
                FLIP: begin
                    if (min_weight < current_synd_weight && current_synd_weight > 0) begin
                        $display("COMP [Iter %0d]: Flipped data bit %d (Weight reduced %d -> %d)", pass_count, (N-1 - best_bit), current_synd_weight, min_weight);
                        codeword[N-1 - best_bit] <= ~codeword[N-1 - best_bit];
                        state <= UPDATE_SYND_FETCH;
                        counter <= 0;
                    end else begin
                        state <= OUTPUT;
                        counter <= 0;
                        done <= 1;
                    end
                end
                
                UPDATE_SYND_FETCH: begin
                    for (i = 0; i < 8; i = i + 1) 
                        current_sig[i] <= poly_rom[{best_bit[14:8], 3'b0} + i];
                    state <= UPDATE_SYND_SHIFT;
                    counter <= 0;
                end
                
                UPDATE_SYND_SHIFT: begin
                    if (counter < best_bit[7:0]) begin
                        for (i = 0; i < 8; i = i + 1) 
                            current_sig[i] <= {current_sig[i][254:0], current_sig[i][255]};
                        counter <= counter + 1;
                    end else begin
                        for (i = 0; i < 8; i = i + 1)
                            syndrome[i] <= syndrome[i] ^ current_sig[i];
                            
                        error_count <= error_count + 1;
                        
                        if (pass_count < 20) begin
                            pass_count <= pass_count + 1;
                            state <= SCAN;
                            counter <= 0;
                        end else begin
                            state <= OUTPUT;
                            counter <= 0;
                            done <= 1;
                        end
                    end
                end

                OUTPUT: begin
                    data_out_valid <= 1;
                    counter <= counter + 1;
                    if (counter == N/8 - 1) begin
                        state <= IDLE;
                        data_out_valid <= 0;
                    end
                end
            endcase
            
            if (state == SCAN || state == UPDATE_SYND_SHIFT)
                bit_counter <= counter;
        end
    end

    assign data_out = codeword[N-1 - counter*8 -: 8];

endmodule
