module weighted_inner_product (
    input logic clk,
    input logic reset,
    input logic start,
    input logic [2:0] width,        // n <= 4
    input logic signed [31:0] w,
    input logic M_on,
    input logic signed [31:0] M_1, M_2, M_3, M_4, M_5, M_6, M_7, M_8, M_9, M_10, M_11, M_12, M_13, M_14, M_15, M_16,
    input logic signed [31:0] c_1_1,
    input logic signed [31:0] c_1_2,
    input logic signed [31:0] c_1_3,
    input logic signed [31:0] c_1_4,
    

    input logic signed [31:0] c_2_1,
    input logic signed [31:0] c_2_2,
    input logic signed [31:0] c_2_3,
    input logic signed [31:0] c_2_4,

    output logic done,
    output logic signed [31:0] y
);

    logic signed [31:0] vec_c [15:0];
    logic signed [31:0] c_1 [3:0];
    logic signed [31:0] c_2 [3:0];
    logic signed [31:0] temp1, temp2, temp3, temp4, temp5, temp6, temp7, temp8, temp9, temp10, temp11, temp12, temp13, temp14, temp15, temp16;

    logic gen_enable;
    logic vec_enable;
    logic read_enable;
    logic [2:0] gen_index;
    logic [2:0] vec_index;
    logic [2:0] gen_row_num;
    logic [2:0] vec_row_num;

    integer i;
    // initialize a DMEM to store outer product.
    // maximum number of 16 rows are supported.
    logic signed [31:0] outer_mem0 [3:0]; // n <= 16
    logic signed [31:0] outer_mem1 [3:0];
    logic signed [31:0] outer_mem2 [3:0];
    logic signed [31:0] outer_mem3 [3:0];

    always_ff @(posedge clk) begin
        if (reset) begin
            gen_enable <= 0;
            for(i = 0; i <= 3; i = i+1) begin            
                c_1[i] <= 0;
                c_2[i] <= 0;
            end

        end else if (start) begin
            c_1[0] <= c_1_1;       // load everything in one cycle
            c_1[1] <= c_1_2;
            c_1[2] <= c_1_3;
            c_1[3] <= c_1_4;
            c_2[0] <= c_2_1;
            c_2[1] <= c_2_2;
            c_2[2] <= c_2_3;
            c_2[3] <= c_2_4;

            gen_enable <= 1;
        end

        if (vec_enable) begin
            gen_enable <= 0;
        end
    end

    // generate one entry in outter product matrix at once
    always_ff @(posedge clk) begin
        if (gen_enable) begin
            if (gen_index < width) begin

                case (gen_row_num)
                    3'b000: begin
                        outer_mem0[gen_index[1:0]] <= c_2[gen_index[1:0]]*c_1[gen_row_num[1:0]];
                    end
                    3'b001: begin
                        outer_mem1[gen_index[1:0]] <= c_2[gen_index[1:0]]*c_1[gen_row_num[1:0]];
                    end
                    3'b010: begin
                        outer_mem2[gen_index[1:0]] <= c_2[gen_index[1:0]]*c_1[gen_row_num[1:0]];
                    end
                    3'b011: begin
                        outer_mem3[gen_index[1:0]] <= c_2[gen_index[1:0]]*c_1[gen_row_num[1:0]];
                    end
                    default: begin
                        outer_mem0[gen_index[1:0]] <= 0;
                    end
                endcase

                gen_index <= gen_index + 1;
                if (gen_index + 1 == width) begin
                    gen_index <= 0;
                    gen_row_num <= gen_row_num + 1;
                    if (gen_row_num + 1 == width) begin
                        vec_enable <= 1;
                        gen_row_num <= 0;
                        gen_index <= 0;
                    end
                end
            end else begin
                gen_row_num <= 0;
                gen_index <= 0;
            end
        end else if (read_enable) begin
            vec_enable <= 0;
            for(i = 0; i <= 3; i = i+1) begin

                outer_mem0[i] <= 0;
                outer_mem1[i] <= 0;
                outer_mem2[i] <= 0;
                outer_mem3[i] <= 0;
                
            end
        end
    end

    integer j;
    always_ff @(posedge clk) begin
        if (vec_enable) begin
            // do vectorization
            if (vec_index < width) begin
                case (vec_row_num)
                    3'h0: vec_c[j] <= outer_mem0[vec_index[1:0]] / w;
                    3'h1: vec_c[j] <= outer_mem1[vec_index[1:0]] / w;
                    3'h2: vec_c[j] <= outer_mem2[vec_index[1:0]] / w;
                    3'h3: vec_c[j] <= outer_mem3[vec_index[1:0]] / w;
                    default: vec_c[0] <= 0;
                endcase

                vec_index <= vec_index + 1;
                j += 1;
                if (vec_index + 1 == width) begin
                    vec_index <= 0;
                    vec_row_num <= vec_row_num + 1;
                    if (vec_row_num + 1 == width) begin
                        read_enable <= 1;
                        vec_row_num <= 0;
                        vec_index <= 0;
                    end
                end
            end
        end
    end

    assign temp1 = M_1 * vec_c[0];
    assign temp2 = M_2 * vec_c[1];
    assign temp3 = M_3 * vec_c[2];
    assign temp4 = M_4 * vec_c[3];
    assign temp5 = M_5 * vec_c[4];
    assign temp6 = M_6 * vec_c[5];
    assign temp7 = M_7 * vec_c[6];
    assign temp8 = M_8 * vec_c[7];
    assign temp9 = M_9 * vec_c[8];
    assign temp10 = M_10 * vec_c[9];
    assign temp11 = M_11 * vec_c[10];
    assign temp12 = M_12 * vec_c[11];
    assign temp13 = M_13 * vec_c[12];
    assign temp14 = M_14 * vec_c[13];
    assign temp15 = M_15 * vec_c[14];
    assign temp16 = M_16 * vec_c[15];

    always @(posedge clk) begin
        if (read_enable && M_on) begin
            done <= 1;
            y <= temp1 + temp2 + temp3 + temp4 + temp5 + temp6 + temp7 + temp8 + temp9 + temp10 + temp11 + temp12 + temp13 + temp14 + temp15 + temp16;
        end

        else begin
            done <= 0;
            y <= 32'b 0;
        end
    end

endmodule