`ifndef ASU_H_SV
`define ASU_H_SV

`ifndef SEVEN_SEG_H_SV
`define SEVEN_SEG_H_SV
module seven_seg(input logic [3:0] num, input logic [0:0] rbi_inv, output logic [7:0] segments);
        always_comb begin
           case (num)
               4'b0000: segments = (rbi_inv) ? ~(8'b01111111) : ~(8'b00000000);
               4'b0001: segments = ~(8'b00001101);
               4'b0010: segments = ~(8'b10110111);
               4'b0011: segments = ~(8'b10011111);
               4'b0100: segments = ~(8'b11001101);
               4'b0101: segments = ~(8'b11011011);
               4'b0110: segments = ~(8'b11111011);
               4'b0111: segments = ~(8'b00001111);
               4'b1000: segments = ~(8'b11111111);
               4'b1001: segments = ~(8'b11001111);
               4'b1010: segments = ~(8'b11101111);
               4'b1011: segments = ~(8'b11111001);
               4'b1100: segments = ~(8'b01110011);
               4'b1101: segments = ~(8'b10111101);
               4'b1110: segments = ~(8'b11110011);
               4'b1111: segments = ~(8'b11100011);
               default: segments = ~(8'bxxxxxxxx);
           endcase     
        end
endmodule
`endif

module asu #(parameter width = 2)(
	input logic [0:0] op,
	input logic[2**width-1:0] A, input logic[2**width-1:0] B,
   output logic[2**width-1:0] OUT_1, output logic[0:0] COUT, output logic[0:0] OUT_CIN
// Used to switch between demo mode and testbench mode
// TODO: Strip for assignment 2
`ifdef SYNTHESIS
	, output logic[7:0] out_seg_1);
`else
    );
`endif
	
    logic [2**width-1:0] operand;
    always_comb begin
        // A + (sub) ? B : -B
        // Adds an extra bit to all operands to calculate the carry bit. Not optimal for hardware synthesis
        //operand = (B ^ {(2**width){op}});
        operand = op == 1'b0 ? B : ~B;
        {COUT, OUT_1} = A + operand + {{(2**width-1){1'b0}}, op};
    end

    // Need to do this outside of the always_comb block because parameter-based bitslicing/indexing isn't supported?
    // Cin for final bit in adder to later be xor'ed with COUT to determine if a signed overflow occured for the MIPS architecture
    assign OUT_CIN = A[2**width-1] ^ operand[2**width-1] ^ OUT_1[2**width-1];

	// Include seven segment visual output for hardware builds only
    // TODO: Strip this for assignment 2 outright
    `ifdef SYNTHESIS 
	    seven_seg dut(OUT_1, 1'b1, out_seg_1);
    `endif
endmodule
`endif


// Testbench

module asu_tb #(width=4)();
    logic [0:0] op;
    logic [width-1:0] A;
    logic [width-1:0] B;
    logic [width-1:0] C;
    logic [0:0] Cout;
    logic [0:0] Cin;
    logic [width-1:0] C_exp;
    logic [0:0] Cout_exp;
    logic [0:0] Cin_exp;   // TODO

    asu #($clog2(width)) dut(op, A, B, C, Cout, Cin);

    integer error_count;
    initial begin
        error_count = 0;
        $display("Beginning ASU tests");

        // Test basic addition of 0+1 = 1
        op = 1'b0;
        A = {width{1'b0}};
        B = { {(width-1){1'b0}}, 1'b1};
        C_exp = { {(width-1){1'b0}}, 1'b1};
        Cout_exp = 1'b0;
        Cin_exp = 1'b0;
        #1;
        if (C != C_exp || Cout != Cout_exp || Cin != Cin_exp || (Cin ^ Cout) != (Cin_exp ^ Cout_exp)) begin
            $display("\t\x1b[38;5;9mOp 0b%b(0b%b, 0b%b) produced 0b%b%b_%b but 0b%b%b_%b was expected\x1b[m", op, A, B, Cin, Cout, C, Cin_exp, Cout_exp, C_exp);
            error_count = error_count + 1;
        end

        // Test underflow from 0-1 = UINTMAX
        op = 1'b1;
        C_exp = {(width){1'b1}};
        Cout_exp = 1'b0;
        Cin_exp = 1'b0;
        #1;
        if (C != C_exp || Cout != Cout_exp || Cin != Cin_exp || (Cin ^ Cout) != (Cin_exp ^ Cout_exp)) begin
            $display("\t\x1b[38;5;9mOp 0b%b(0b%b, 0b%b) produced 0b%b%b_%b but 0b%b%b_%b was expected\x1b[m", op, A, B, Cin, Cout, C, Cin_exp, Cout_exp, C_exp);
            error_count = error_count + 1;
        end

        // Test carry addition of 1+1 = 2
        op = 1'b0;
        A = { {(width-1){1'b0}}, 1'b1};
        B = { {(width-1){1'b0}}, 1'b1};
        C_exp = { {(width-2){1'b0}}, 2'b10};
        Cout_exp = 1'b0;
        Cin_exp = 1'b0;
        #1;
        if (C != C_exp || Cout != Cout_exp || Cin != Cin_exp || (Cin ^ Cout) != (Cin_exp ^ Cout_exp)) begin
            $display("\t\x1b[38;5;9mOp 0b%b(0b%b, 0b%b) produced 0b%b%b_%b but 0b%b%b_%b was expected\x1b[m", op, A, B, Cin, Cout, C, Cin_exp, Cout_exp, C_exp);
            error_count = error_count + 1;
        end

        // Test overflow from  UINTMAX+1 = 0
        A = {(width){1'b1}};            // UINTMAX
        B = { {(width-1){1'b0}}, 1'b1}; // 1
        C_exp = {(width){1'b0}};
        Cout_exp = 1'b1;
        Cin_exp = 1'b1;
        #1;
        if (C != C_exp || Cout != Cout_exp || Cin != Cin_exp || (Cin ^ Cout) != (Cin_exp ^ Cout_exp)) begin
            $display("\t\x1b[38;5;9mOp 0b%b(0b%b, 0b%b) produced 0b%b%b_%b but 0b%b%b_%b was expected\x1b[m", op, A, B, Cin, Cout, C, Cin_exp, Cout_exp, C_exp);
            error_count = error_count + 1;
        end

        // Tests INTMAX+1
        op = 1'b0;
        A = {1'b0, {(width-1){1'b1}}};          // INTMAX
        B = {{(width-1){1'b0}}, 1'b1};          // 1
        C_exp    = {1'b1, {(width-1){1'b0}}};   // -INTMAX
        Cin_exp  = 1'b1;
        Cout_exp = 1'b0;
        #1;
        if (C != C_exp || Cout != Cout_exp || Cin != Cin_exp || (Cin ^ Cout) != (Cin_exp ^ Cout_exp)) begin
            $display("\t\x1b[38;5;9mOp 0b%b(0b%b, 0b%b) produced 0b%b%b_%b but 0b%b%b_%b was expected\x1b[m", op, A, B, Cin, Cout, C, Cin_exp, Cout_exp, C_exp);
            error_count = error_count + 1;
        end

        op = 1'b0;
        A = {1'b1, {(width-1){1'b0}}};  // -INTMAX
        B = {1'b1, {(width-1){1'b0}}};  // -INTMAX
        C_exp    = {width{1'b0}};       // 0
        Cin_exp  = 1'b0;
        Cout_exp = 1'b1;
        #1;
        if (C != C_exp || Cout != Cout_exp || Cin != Cin_exp || (Cin ^ Cout) != (Cin_exp ^ Cout_exp)) begin
            $display("\t\x1b[38;5;9mOp 0b%b(0b%b, 0b%b) produced 0b%b%b_%b but 0b%b%b_%b was expected\x1b[m", op, A, B, Cin, Cout, C, Cin_exp, Cout_exp, C_exp);
            error_count = error_count + 1;
        end

        op = 1'b1;
        A = {1'b0, {(width-1){1'b1}}}; // INTMAX
        B = {1'b1, {(width-1){1'b0}}}; // -INTMAX
        C_exp    = {width{1'b1}};      // -1
        Cin_exp  = 1'b1;
        Cout_exp = 1'b0;
        #1;
        if (C != C_exp || Cout != Cout_exp || Cin != Cin_exp || (Cin ^ Cout) != (Cin_exp ^ Cout_exp)) begin
            $display("\t\x1b[38;5;9mOp 0b%b(0b%b, 0b%b) produced 0b%b%b_%b but 0b%b%b_%b was expected\x1b[m", op, A, B, Cin, Cout, C, Cin_exp, Cout_exp, C_exp);
            error_count = error_count + 1;
        end

        op = 1'b1;
        A = {1'b1, {(width-1){1'b0}}};          // -INTMAX
        B = {{(width-1){1'b0}}, 1'b1};          // 1
        C_exp    = {1'b0, {(width-1){1'b1}}};   // INTMAX
        Cin_exp  = 1'b0;
        Cout_exp = 1'b1;
        #1;
        if (C != C_exp || Cout != Cout_exp || Cin != Cin_exp || (Cin ^ Cout) != (Cin_exp ^ Cout_exp)) begin
            $display("\t\x1b[38;5;9mOp 0b%b(0b%b, 0b%b) produced 0b%b%b_%b but 0b%b%b_%b was expected\x1b[m", op, A, B, Cin, Cout, C, Cin_exp, Cout_exp, C_exp);
            error_count = error_count + 1;
        end

        if (error_count == 0) $display("\x1b[38;5;10mAll ASU tests passing\x1b[m");
        else $display("\x1b[38;5;9m%0d ASU test(s) failing\x1b[m", error_count);
    end
endmodule