`timescale 1ns / 1ps
module Execution (
    input clk,
    // From prev stage
    input prev_is_branch,
    input prev_WB_source,
    input prev_memory_write,
    input [4:0] prev_ALU_function,
    input prev_ALU_source1,
    input prev_ALU_source2,
    input prev_register_write_destination_source,
    input prev_register_write,
    input [31:0] prev_PC_plus_4,
    input [31:0] prev_register_file_read_A,
    input [31:0] prev_register_file_read_B,
    input [4:0] prev_shamt,
    input [31:0] prev_extended_immediate,
    input [4:0] prev_rs_address,
    input [4:0] prev_rt_address,
    input [4:0] prev_rd_address,
    // From hazard unit
    input IDEXSrc,
    // From forward unit
    input [1:0] operandASrc,
    input [1:0] operandBSrc,
    // From MEM
    input [31:0] MEM_forwarded_data,
    // From WB
    input [31:0] WB_forwarded_data,
    // To IF stage, also hazard unit
    output PC_branch,
    // Only to IF stage
    output [31:0] branch_target,
    // To next stage
    output register_write,
    output WB_source,
    output memory_write,
    output [31:0] ALU_result,
    output [31:0] memory_write_data,
    output [4:0] register_write_destination,
    // To forward unit
    output [4:0] rs_address,
    output [4:0] rt_address
);

    reg EX_register_write;
    reg EX_WB_source;
    reg EX_memory_write;
    reg EX_is_branch;
    reg [4:0] EX_ALU_function;
    reg EX_ALU_source1;
    reg EX_ALU_source2;
    reg EX_register_write_destination_source;
    reg [31:0] EX_PC_plus_4;
    reg [31:0] EX_register_read_A;
    reg [31:0] EX_register_read_B;
    reg [4:0] EX_shamt;
    reg [31:0] EX_extended_immediate;
    reg [4:0] EX_rs_address;
    reg [4:0] EX_rt_address;
    reg [4:0] EX_rd_address;

    // ALU Part
    wire [31:0] ALU_in1;
    wire [31:0] RF_read_B_include_forward;
    wire [31:0] ALU_in2;
    wire [31:0] ALU_out;
    assign ALU_in1 = (EX_ALU_source1 == 1'b1) ? EX_shamt :
                     (operandASrc == 2'b00) ? EX_register_read_A :
                     (operandASrc == 2'b01) ? MEM_forwarded_data :
                     (operandASrc == 2'b10) ? WB_forwarded_data : 32'h00000000;

    assign RF_read_B_include_forward = (operandBSrc == 2'b00) ? EX_register_read_B :
                                       (operandBSrc == 2'b01) ? MEM_forwarded_data :
                                       (operandBSrc == 2'b10) ? WB_forwarded_data : 32'h00000000;
    assign ALU_in2 = (EX_ALU_source2 == 1'b1) ? EX_extended_immediate : RF_read_B_include_forward;

    ALU alu(
        .funct(EX_ALU_function),
        .in_1(ALU_in1),
        .in_2(ALU_in2),
        .result(ALU_out)
    );

    // All output signals
    assign PC_branch = EX_is_branch & ALU_out[0];
    assign branch_target = EX_PC_plus_4 + (EX_extended_immediate << 2);
    assign register_write = EX_register_write;
    assign WB_source = EX_WB_source;
    assign memory_write = EX_memory_write;
    assign ALU_result = ALU_out;
    assign memory_write_data = RF_read_B_include_forward;
    assign register_write_destination = (EX_register_write_destination_source == 1'b0) ?
                                        EX_rt_address : EX_rd_address;
    assign rs_address = EX_rs_address;
    assign rt_address = EX_rt_address;

    always @(posedge clk) begin
        if (IDEXSrc == 1'b1) begin
            EX_register_write <= 1'b0;
            EX_WB_source <= 1'b0;
            EX_memory_write <= 1'b0;
            EX_is_branch <= 1'b0;
            EX_ALU_function <= 5'b00000;
            EX_ALU_source1 <= 1'b0;
            EX_ALU_source2 <= 1'b0;
            EX_register_write_destination_source <= 1'b0;
            EX_PC_plus_4 <= 32'h00000000;
            EX_register_read_A <= 32'h00000000;
            EX_register_read_B <= 32'h00000000;
            EX_shamt <= 5'b00000;
            EX_extended_immediate <= 32'h00000000;
            EX_rs_address <= 5'b00000;
            EX_rt_address <= 5'b00000;
            EX_rd_address <= 5'b00000;
        end else begin
            EX_register_write <= prev_register_write;
            EX_WB_source <= prev_WB_source;
            EX_memory_write <= prev_memory_write;
            EX_is_branch <= prev_is_branch;
            EX_ALU_function <= prev_ALU_function;
            EX_ALU_source1 <= prev_ALU_source1;
            EX_ALU_source2 <= prev_ALU_source2;
            EX_register_write_destination_source <= prev_register_write_destination_source;
            EX_PC_plus_4 <= prev_PC_plus_4;
            EX_register_read_A <= prev_register_file_read_A;
            EX_register_read_B <= prev_register_file_read_B;
            EX_shamt <= prev_shamt;
            EX_extended_immediate <= prev_extended_immediate;
            EX_rs_address <= prev_rs_address;
            EX_rt_address <= prev_rt_address;
            EX_rd_address <= prev_rd_address;
        end
    end
endmodule