// One amplifier path: the generator streams are merged into one stream, which then passes
// three stages in series. Merging before amplifying is equivalent to amplifying each stream
// with the same first stage and merging after (the random walk is distributive), and needs a
// third of the stages. Amplified bits are tallied per sub-trial and per trial.
module amp_path #(
    parameter GENERATORS = 4
) (
    input  wire                  clk,
    input  wire                  restart,
    input  wire [GENERATORS-1:0] in_valid,
    input  wire [GENERATORS-1:0] in_bit,
    input  wire [2:0]            stage_mode,
    input  wire [47:0]           stage_param,  // stage s at [16*s +: 16]
    input  wire                  sub_end,      // pulse: close the current sub-trial
    input  wire [3:0]            sub_index,    // sub-trial being closed
    input  wire                  trial_end,    // pulse: start a new trial (after the last sub_end)
    output reg  [31:0]           trial_bits,   // amplified bits in the current trial
    output reg  [31:0]           trial_ones,
    output reg  [15:0]           votes,        // bit i: sub-trial i voted 1
    output reg  [15:0]           voted,        // bit i: sub-trial i had a majority
    output reg                   overflow      // a merge input arrived before the last was taken
);

    // Merge: each generator's latest bit waits in its own slot; one slot is drained per clock,
    // lowest index first. Generators deliver at most one bit per `temporal` (>= 1, default 5)
    // clocks in lockstep, so with up to 4 generators the slots never overflow at temporal >= 4
    // (flagged if they do).

    reg [GENERATORS-1:0] pending, pending_bit, take;
    reg                  merged_valid, merged_bit;
    integer i;

    always @(*) begin
        take = {GENERATORS{1'b0}};
        for (i = 0; i < GENERATORS; i = i + 1)
            if (pending[i] && take == {GENERATORS{1'b0}}) take[i] = 1'b1;
    end

    always @(posedge clk) begin
        merged_valid <= |take && !restart;
        for (i = 0; i < GENERATORS; i = i + 1)
            if (take[i]) merged_bit <= pending_bit[i];

        if (restart) begin
            pending  <= {GENERATORS{1'b0}};
            overflow <= 1'b0;
        end else begin
            for (i = 0; i < GENERATORS; i = i + 1) begin
                if (in_valid[i]) begin
                    if (pending[i] && !take[i]) overflow <= 1'b1;
                    pending[i]     <= 1'b1;
                    pending_bit[i] <= in_bit[i];
                end else if (take[i]) begin
                    pending[i] <= 1'b0;
                end
            end
        end
    end

    // Stages

    wire first_valid, first_bit, second_valid, second_bit, out_valid, out_bit;

    amp_stage first (
        .clk(clk), .restart(restart), .mode(stage_mode[0]), .param(stage_param[15:0]),
        .in_valid(merged_valid), .in_bit(merged_bit), .out_valid(first_valid), .out_bit(first_bit)
    );

    amp_stage second (
        .clk(clk), .restart(restart), .mode(stage_mode[1]), .param(stage_param[31:16]),
        .in_valid(first_valid), .in_bit(first_bit), .out_valid(second_valid), .out_bit(second_bit)
    );

    amp_stage third (
        .clk(clk), .restart(restart), .mode(stage_mode[2]), .param(stage_param[47:32]),
        .in_valid(second_valid), .in_bit(second_bit), .out_valid(out_valid), .out_bit(out_bit)
    );

    // Tallies. A bit arriving in a sub_end or trial_end cycle counts towards the next one, so
    // the trial count and the sum of its sub-trial counts can differ by one bit.

    reg [23:0] sub_bits, sub_ones;

    always @(posedge clk) begin
        if (restart || trial_end) begin
            trial_bits <= out_valid ? 32'd1 : 32'd0;
            trial_ones <= (out_valid && out_bit) ? 32'd1 : 32'd0;
            votes      <= 16'd0;
            voted      <= 16'd0;
        end else if (out_valid) begin
            trial_bits <= trial_bits + 32'd1;
            trial_ones <= trial_ones + out_bit;
        end

        if (restart || sub_end) begin
            sub_bits <= out_valid ? 24'd1 : 24'd0;
            sub_ones <= (out_valid && out_bit) ? 24'd1 : 24'd0;
        end else if (out_valid) begin
            sub_bits <= sub_bits + 24'd1;
            sub_ones <= sub_ones + out_bit;
        end

        if (!restart && sub_end && sub_bits != 0 && {sub_ones, 1'b0} != sub_bits) begin
            voted[sub_index] <= 1'b1;
            votes[sub_index] <= {sub_ones, 1'b0} > sub_bits;
        end
    end

endmodule