// CSEE 4840 Project: Choose and Run One of the Three Encrypted-Domain Operations
// Avalon memory-mapped peripheral that generates accelerate encrypted domain operations
//
// Spring 2022
//
// By: Lanxiang Hu
// Uni: lh3116

module encrypted_domain (input logic    clk,          // 50MHz clock
	    input logic 	               reset,
		input logic signed [31:0]      writedata,
		input logic 	               write,
		input logic		               chipselect,
		input logic [3:0]              address,

        input logic                    all_loaded,

        output logic signed [31:0]     DATA_OUT_0, 
        output logic signed [31:0]     DATA_OUT_1, 
        output logic signed [31:0]     DATA_OUT_2, 
        output logic signed [31:0]     DATA_OUT_3, 
        output logic signed [31:0]     DATA_OUT_4, 
        output logic signed [31:0]     DATA_OUT_5, 
        output logic signed [31:0]     DATA_OUT_6, 
        output logic signed [31:0]     DATA_OUT_7, 
        output logic signed [31:0]     DATA_OUT_8, 
        output logic signed [31:0]     DATA_OUT_9, 
        output logic signed [31:0]     DATA_OUT_10,
        output logic signed [31:0]     DATA_OUT_11, 
        output logic signed [31:0]     DATA_OUT_12, 
        output logic signed [31:0]     DATA_OUT_13, 
        output logic signed [31:0]     DATA_OUT_14, 
        output logic signed [31:0]     DATA_OUT_15, 

        output logic [4:0]      OUTPUT_LENGTH, // at most 16
        output logic [4:0]      OUTPUT_WIDTH,  // at most 16
        output logic            DONE);

      logic [3:0] operation;
      logic [3:0] data_type;
      logic start_0;
      logic start_1;
      logic start_2;
      logic [4:0]  width; // n <= 16, for incoming matrix or vector
      logic [4:0]  length; // n <= 16, for incoming matrix

      logic [31:0] data_out_1;
      logic [31:0] data_out_2;
      logic done_0;
      logic done_1;
      logic done_2;

      const logic [3:0] load_op_type = 4'h0;
      const logic [3:0] load_data_type = 4'h1;
      const logic [3:0] load_width = 4'h2;
      const logic [3:0] load_length = 4'h3;
      const logic [3:0] load_input = 4'h4;

      const logic [3:0] vector_addition = 4'h0;
      const logic [3:0] linear_transform = 4'h1;
      const logic [3:0] weighted_inner_product = 4'h2;

      const logic [3:0] load_c_1 = 4'h0;
      const logic [3:0] load_c_2 = 4'h1;
      const logic [3:0] load_M = 4'h2;
      const logic [3:0] load_c = 4'h3;
      const logic [3:0] load_c_1_out = 4'h4;
      const logic [3:0] load_c_2_out = 4'h5;
      const logic [3:0] load_w = 4'h6;
      
      logic signed [31:0] c_1 [15:0];
      logic signed [31:0] c_2 [15:0];
      logic signed [31:0] c_out [15:0];
      // instantiate the three modules
      vector_addition vector_addition0 (
        .clk(clk),
        .reset(reset),
        .start(start_0),
        .c_1_1(c_1[0]),
        .c_1_2(c_1[1]),
        .c_1_3(c_1[2]),
        .c_1_4(c_1[3]),
        .c_1_5(c_1[4]),
        .c_1_6(c_1[5]),
        .c_1_7(c_1[6]),
        .c_1_8(c_1[7]),
        .c_1_9(c_1[8]),
        .c_1_10(c_1[9]),
        .c_1_11(c_1[10]),
        .c_1_12(c_1[11]),
        .c_1_13(c_1[12]),
        .c_1_14(c_1[13]),
        .c_1_15(c_1[14]),
        .c_1_16(c_1[15]),

        .c_2_1(c_2[0]),
        .c_2_2(c_2[1]),
        .c_2_3(c_2[2]),
        .c_2_4(c_2[3]),
        .c_2_5(c_2[4]),
        .c_2_6(c_2[5]),
        .c_2_7(c_2[6]),
        .c_2_8(c_2[7]),
        .c_2_9(c_2[8]),
        .c_2_10(c_2[9]),
        .c_2_11(c_2[10]),
        .c_2_12(c_2[11]),
        .c_2_13(c_2[12]),
        .c_2_14(c_2[13]),
        .c_2_15(c_2[14]),
        .c_2_16(c_2[15]),

        .done(done_0),
        .c_1(c_out[0]),
        .c_2(c_out[1]),
        .c_3(c_out[2]),
        .c_4(c_out[3]),
        .c_5(c_out[4]),
        .c_6(c_out[5]),
        .c_7(c_out[6]),
        .c_8(c_out[7]),
        .c_9(c_out[8]),
        .c_10(c_out[9]),
        .c_11(c_out[10]),
        .c_12(c_out[11]),
        .c_13(c_out[12]),
        .c_14(c_out[13]),
        .c_15(c_out[14]),
        .c_16(c_out[15]));


      logic signed [31:0] M [15:0];
      logic signed [31:0] c [15:0];
      linear_transform linear_transform0 (
        .clk(clk),
        .reset(reset),
        .on(start_1),
        .M_1(M[0]), 
        .M_2(M[1]), 
        .M_3(M[2]), 
        .M_4(M[3]), 
        .M_5(M[4]), 
        .M_6(M[5]), 
        .M_7(M[6]), 
        .M_8(M[7]), 
        .M_9(M[8]), 
        .M_10(M[9]), 
        .M_11(M[10]), 
        .M_12(M[11]), 
        .M_13(M[12]), 
        .M_14(M[13]), 
        .M_15(M[14]), 
        .M_16(M[15]),

        .c_1(c[0]), 
        .c_2(c[1]), 
        .c_3(c[2]), 
        .c_4(c[3]), 
        .c_5(c[4]), 
        .c_6(c[5]), 
        .c_7(c[6]), 
        .c_8(c[7]), 
        .c_9(c[8]), 
        .c_10(c[9]), 
        .c_11(c[10]), 
        .c_12(c[11]), 
        .c_13(c[12]), 
        .c_14(c[13]), 
        .c_15(c[14]), 
        .c_16(c[15]),

        .write(done_1),
        .y(data_out_1));
      
      logic signed [31:0] w;
      logic M_on;
      logic signed [31:0] c_1_out [15:0];
      logic signed [31:0] c_2_out [15:0];
      weighted_inner_product weighted_inner_product0 (
        .clk(clk),
        .reset(reset),
        .start(start_2),
        .width(width[2:0]),        // n <= 4
        .w(w),
        .M_on(M_on),
        .M_1(M[0]), 
        .M_2(M[1]), 
        .M_3(M[2]), 
        .M_4(M[3]), 
        .M_5(M[4]), 
        .M_6(M[5]), 
        .M_7(M[6]), 
        .M_8(M[7]), 
        .M_9(M[8]), 
        .M_10(M[9]), 
        .M_11(M[10]), 
        .M_12(M[11]), 
        .M_13(M[12]), 
        .M_14(M[13]), 
        .M_15(M[14]), 
        .M_16(M[15]),

        .c_1_1(c_1_out[0]),
        .c_1_2(c_1_out[1]),
        .c_1_3(c_1_out[2]),
        .c_1_4(c_1_out[3]),
        

        .c_2_1(c_2_out[0]),
        .c_2_2(c_2_out[1]),
        .c_2_3(c_2_out[2]),
        .c_2_4(c_2_out[3]),

        .done(done_2),
        .y(data_out_2));
    
    integer i;
    logic [3:0] load_counter;
    always_ff @(posedge clk) begin
        if (reset) begin
            operation <= 0;
            data_type <= 0;

            width <= 0; // n <= 16, for incoming matrix or vector
            length <= 0; // n <= 16, for incoming matrix

            data_out_1 <= 0;
            data_out_2 <= 0;
            done_1 <= 0;
            done_2 <= 0;

            w <= 0;

            for(i = 0; i <= 15; i = i+1)begin
                c_1[i] <= 0;
                c_2[i] <= 0;
                c_out[i] <= 0;
                M[i] <= 0;
                c[i] <= 0;
                c_1_out[i] <= 0;
                c_2_out[i] <= 0;
            end
        // 'all_loaded' is to be controled from external
        end else if (chipselect && write && !all_loaded) begin
            // loading stage
            case (address)
                load_op_type : begin
                    operation <= writedata[3:0];
                    // new operation starts, zero out counters
                    load_counter <= 0;
                end
                load_data_type : begin
                    data_type <= writedata[3:0];
                    // new operation starts, zero out counters
                    load_counter <= 0;
                end
                load_width : begin
                    width <= writedata[4:0]; 
                    // new operation starts, zero out counters
                    load_counter <= 0;
                end
                load_length : begin
                    length <= writedata[4:0];
                    // new operation starts, zero out counters
                    load_counter <= 0;
                end
                load_input : begin
                    case (data_type)
                        load_c_1 : c_1[load_counter] <= writedata;
                        load_c_2 : c_2[load_counter] <= writedata;
                        load_M : M[load_counter] <= writedata;
                        load_c : c[load_counter] <= writedata;
                        load_c_1_out : c_1_out[load_counter] <= writedata;
                        load_c_2_out : c_2_out[load_counter] <= writedata;
                        load_w: w <= writedata;
                    endcase
                    load_counter <= load_counter + 1;
                end
            endcase
        end else if (all_loaded) begin
            // new operation starts, zero out counters
            load_counter <= 0;
        end
    end

    logic [8:0]  counter_2;
    logic [3:0] read_index;
    // workflow: load operation type, data type, width, length, inputs sequentially
    // after been properly loaded, "all_loaded" turns true
    always_ff @(posedge clk) begin
        // computation stage
        if (all_loaded) begin
            case (operation)
                vector_addition : begin 
                    start_0 <= 1;
                    // no loading, zero out counters
                    counter_2 <= 0;
                    DONE <= done_0;

                    DATA_OUT_0 <= c_out[0];
                    DATA_OUT_1 <= c_out[1];
                    DATA_OUT_2 <= c_out[2];
                    DATA_OUT_3 <= c_out[3];
                    DATA_OUT_4 <= c_out[4];
                    DATA_OUT_5 <= c_out[5];
                    DATA_OUT_6 <= c_out[6];
                    DATA_OUT_7 <= c_out[7];
                    DATA_OUT_8 <= c_out[8];
                    DATA_OUT_9 <= c_out[9];
                    DATA_OUT_10 <= c_out[10];
                    DATA_OUT_11 <= c_out[11];
                    DATA_OUT_12 <= c_out[12];
                    DATA_OUT_13 <= c_out[13];
                    DATA_OUT_14 <= c_out[14];
                    DATA_OUT_15 <= c_out[15];

                    OUTPUT_WIDTH <= width;
                end
                linear_transform : begin
                    start_1 <= 1; 
                    // no loading, zero out counters
                    counter_2 <= 0;

                    DONE <= done_1;
                    DATA_OUT_0 <= data_out_1;
                    OUTPUT_WIDTH <= length;
                end
                weighted_inner_product : begin 
                    if (counter_2 == 0) begin
                        start_2 <= 1;

                    end else if (counter_2 == 1) begin
                        start_2 <= 0;
                    // M_on turns on
                    end else if (counter_2 == 2 * width * width + 10) begin
                        M_on <= 1;
                    end else if (counter_2 == 3 * width * width + 10) begin
                        M_on <= 0;
                        // no loading, zero out counters
                        counter_2 <= 0;
                    end
                    counter_2 <= counter_2 + 1;
                    DONE <= done_2;
                    DATA_OUT_0 <= data_out_2;
                    OUTPUT_WIDTH <= length;
                end
            endcase
        end else begin
            // turns off all modules when 'all_loaded' false
            start_0 <= 0;
            start_1 <= 0;
            start_2 <= 0;
            DONE <= 0;

            counter_2 <= 0;
            DATA_OUT_0 <= 0;
            DATA_OUT_1 <= 0;
            DATA_OUT_2 <= 0;
            DATA_OUT_3 <= 0;
            DATA_OUT_4 <= 0;
            DATA_OUT_5 <= 0;
            DATA_OUT_6 <= 0;
            DATA_OUT_7 <= 0;
            DATA_OUT_8 <= 0;
            DATA_OUT_9 <= 0;
            DATA_OUT_10 <= 0;
            DATA_OUT_11 <= 0;
            DATA_OUT_12 <= 0;
            DATA_OUT_13 <= 0;
            DATA_OUT_14 <= 0;
            DATA_OUT_15 <= 0;
        end
    end

   
endmodule

