// Cofiber detection head — fully combinational per-token detector
//
// Input: 768 signed 8-bit feature values (one spatial token)
// Output: 80 classification scores + 4 box regression values + 1 centerness
//
// All weights are ROM (analytically derived, zero training).
// The cofiber decomposition and pooling happen upstream in the spatial
// processing module — this module handles per-token prediction only.
//
// Parameters:
//   FEAT_DIM    - feature dimension (768 for EUPE-ViT-B)
//   NUM_CLASSES - detection classes (80 for COCO)
//   BIT_WIDTH   - weight/activation precision

module cofiber_detector #(
    parameter FEAT_DIM    = 768,
    parameter NUM_CLASSES = 80,
    parameter BIT_WIDTH   = 8,
    parameter ACC_WIDTH   = 24  // accumulator width for MAC results
)(
    input  logic signed [BIT_WIDTH-1:0] features [0:FEAT_DIM-1],
    output logic signed [ACC_WIDTH-1:0] cls_scores [0:NUM_CLASSES-1],
    output logic signed [ACC_WIDTH-1:0] box_ltrb [0:3],
    output logic signed [ACC_WIDTH-1:0] centerness
);

    // Weight ROMs — loaded from analytical solution
    logic signed [BIT_WIDTH-1:0] cls_weight [0:NUM_CLASSES-1][0:FEAT_DIM-1];
    logic signed [BIT_WIDTH-1:0] cls_bias [0:NUM_CLASSES-1];
    logic signed [BIT_WIDTH-1:0] reg_weight [0:3][0:FEAT_DIM-1];
    logic signed [BIT_WIDTH-1:0] reg_bias [0:3];
    logic signed [BIT_WIDTH-1:0] ctr_weight [0:FEAT_DIM-1];
    logic signed [BIT_WIDTH-1:0] ctr_bias;

    initial begin
        $readmemh("rom/cls_weight.hex", cls_weight);
        $readmemh("rom/cls_bias.hex", cls_bias);
        $readmemh("rom/reg_out_weight.hex", reg_weight);
        $readmemh("rom/reg_out_bias.hex", reg_bias);
        $readmemh("rom/ctr_weight.hex", ctr_weight);
        $readmemh("rom/ctr_bias.hex", ctr_bias);
    end

    // Classification: 80 parallel dot products
    // cls_scores[c] = sum(features[d] * cls_weight[c][d]) + cls_bias[c]
    genvar c, d;
    generate
        for (c = 0; c < NUM_CLASSES; c = c + 1) begin : cls_mac
            logic signed [ACC_WIDTH-1:0] acc;
            always_comb begin
                acc = {{(ACC_WIDTH-BIT_WIDTH){cls_bias[c][BIT_WIDTH-1]}}, cls_bias[c]};
                for (int dd = 0; dd < FEAT_DIM; dd = dd + 1) begin
                    acc = acc + (features[dd] * cls_weight[c][dd]);
                end
                cls_scores[c] = acc;
            end
        end
    endgenerate

    // Box regression: 4 parallel dot products
    generate
        for (c = 0; c < 4; c = c + 1) begin : reg_mac
            logic signed [ACC_WIDTH-1:0] acc;
            always_comb begin
                acc = {{(ACC_WIDTH-BIT_WIDTH){reg_bias[c][BIT_WIDTH-1]}}, reg_bias[c]};
                for (int dd = 0; dd < FEAT_DIM; dd = dd + 1) begin
                    acc = acc + (features[dd] * reg_weight[c][dd]);
                end
                box_ltrb[c] = acc;
            end
        end
    endgenerate

    // Centerness: 1 dot product
    logic signed [ACC_WIDTH-1:0] ctr_acc;
    always_comb begin
        ctr_acc = {{(ACC_WIDTH-BIT_WIDTH){ctr_bias[BIT_WIDTH-1]}}, ctr_bias};
        for (int dd = 0; dd < FEAT_DIM; dd = dd + 1) begin
            ctr_acc = ctr_acc + (features[dd] * ctr_weight[dd]);
        end
        centerness = ctr_acc;
    end

endmodule


// Cofiber spatial decomposition — operates on the full feature grid
// Produces 3 scale bands from the input feature map
module cofiber_decompose #(
    parameter FEAT_DIM = 768,
    parameter BIT_WIDTH = 8,
    parameter H = 40,
    parameter W = 40
)(
    input  logic signed [BIT_WIDTH-1:0] features [0:FEAT_DIM-1][0:H-1][0:W-1],
    output logic signed [BIT_WIDTH-1:0] scale0 [0:FEAT_DIM-1][0:H-1][0:W-1],      // stride 16 cofiber
    output logic signed [BIT_WIDTH-1:0] scale1 [0:FEAT_DIM-1][0:H/2-1][0:W/2-1],  // stride 32 cofiber
    output logic signed [BIT_WIDTH-1:0] scale2 [0:FEAT_DIM-1][0:H/4-1][0:W/4-1]   // stride 64 residual
);

    // Intermediate: pooled features at half resolution
    logic signed [BIT_WIDTH+1:0] pool0 [0:FEAT_DIM-1][0:H/2-1][0:W/2-1];
    logic signed [BIT_WIDTH+1:0] pool1 [0:FEAT_DIM-1][0:H/4-1][0:W/4-1];

    genvar ch, r, col;
    generate
        // Pool0: 2x2 average pool of input features
        for (ch = 0; ch < FEAT_DIM; ch = ch + 1) begin : pool0_ch
            for (r = 0; r < H/2; r = r + 1) begin : pool0_r
                for (col = 0; col < W/2; col = col + 1) begin : pool0_c
                    assign pool0[ch][r][col] = (features[ch][2*r][2*col]
                                              + features[ch][2*r+1][2*col]
                                              + features[ch][2*r][2*col+1]
                                              + features[ch][2*r+1][2*col+1]) >>> 2;
                end
            end
        end

        // Scale0: cofiber = features - upsample(pool0)
        // Nearest-neighbor upsample for exact integer arithmetic
        for (ch = 0; ch < FEAT_DIM; ch = ch + 1) begin : s0_ch
            for (r = 0; r < H; r = r + 1) begin : s0_r
                for (col = 0; col < W; col = col + 1) begin : s0_c
                    assign scale0[ch][r][col] = features[ch][r][col]
                                              - pool0[ch][r/2][col/2][BIT_WIDTH-1:0];
                end
            end
        end

        // Pool1: 2x2 average pool of pool0
        for (ch = 0; ch < FEAT_DIM; ch = ch + 1) begin : pool1_ch
            for (r = 0; r < H/4; r = r + 1) begin : pool1_r
                for (col = 0; col < W/4; col = col + 1) begin : pool1_c
                    assign pool1[ch][r][col] = (pool0[ch][2*r][2*col]
                                              + pool0[ch][2*r+1][2*col]
                                              + pool0[ch][2*r][2*col+1]
                                              + pool0[ch][2*r+1][2*col+1]) >>> 2;
                end
            end
        end

        // Scale1: cofiber of pool0
        for (ch = 0; ch < FEAT_DIM; ch = ch + 1) begin : s1_ch
            for (r = 0; r < H/2; r = r + 1) begin : s1_r
                for (col = 0; col < W/2; col = col + 1) begin : s1_c
                    assign scale1[ch][r][col] = pool0[ch][r][col][BIT_WIDTH-1:0]
                                              - pool1[ch][r/2][col/2][BIT_WIDTH-1:0];
                end
            end
        end

        // Scale2: the low-frequency residual
        for (ch = 0; ch < FEAT_DIM; ch = ch + 1) begin : s2_ch
            for (r = 0; r < H/4; r = r + 1) begin : s2_r
                for (col = 0; col < W/4; col = col + 1) begin : s2_c
                    assign scale2[ch][r][col] = pool1[ch][r][col][BIT_WIDTH-1:0];
                end
            end
        end
    endgenerate

endmodule