`ifndef LU_H_SV
`define LU_H_SV
module lu #(parameter width = 2) (
    input logic [1:0] op,           // 2 bit operand
    input logic [2**width-1:0] A,  // width defined as 2^n 
	 input logic [2**width-1:0] B,
    output logic [2**width-1:0] Y);
	 
    always_comb begin
		case (op) // different operations of A and B based on operand
			2'b00:	Y =   A & B;
			2'b01:	Y =   A | B;
			2'b10:	Y =   A ^ B;
			2'b11:	Y = ~(A | B);
			default:Y = '0;
		endcase
    end
endmodule
`endif

// Testbench

module lu_tb #(parameter W=4)();
    logic [1:0] op;
    logic [W-1:0] A;
    logic [W-1:0] B;
    logic [W-1:0] C;
    logic [W-1:0] C_exp[16]; // expected values

    lu #(W/2) dut(op, A, B, C); // width / 2 as was adapted from width-1 being standard

    integer error_count;
    initial begin // only need to check 1 bit for expected functionality
        // AND
        C_exp[0]  = {{W{1'b0}}, {1'b0}};
        C_exp[1]  = {{W{1'b0}}, {1'b0}};
        C_exp[2]  = {{W{1'b0}}, {1'b0}};
        C_exp[3]  = {{W{1'b0}}, {1'b1}};
        // OR
        C_exp[4]  = {{W{1'b0}}, {1'b0}};
        C_exp[5]  = {{W{1'b0}}, {1'b1}};
        C_exp[6]  = {{W{1'b0}}, {1'b1}};
        C_exp[7]  = {{W{1'b0}}, {1'b1}};
        // XOR
        C_exp[8]  = {{W{1'b0}}, {1'b0}};
        C_exp[9]  = {{W{1'b0}}, {1'b1}};
        C_exp[10] = {{W{1'b0}}, {1'b1}};
        C_exp[11] = {{W{1'b0}}, {1'b0}};
        // NOR
        C_exp[12] = {{W{1'b1}}, {1'b1}};
        C_exp[13] = {{W{1'b1}}, {1'b0}};
        C_exp[14] = {{W{1'b1}}, {1'b0}};
        C_exp[15] = {{W{1'b1}}, {1'b0}};

        error_count = 0;
        op = 0;
        A = 0;
        B = 0;
        $display("Beginning LU tests"); // find difference between calculated and excpected
        for (integer N = 0; N < 4; ++N) begin
            for (integer M = 0; M < 2; ++M) begin
                for (integer K = 0; K < 2; ++K) begin
                    #1;
                    if (C != C_exp[N*2*2 + M*2 + K]) begin
                        $display("\t\x1b[38;5;9mLU(%0d) Op 0b%b(0b%b, 0b%b) produced 0b%b but 0b%b was expected\x1b[m", N*2*2+M*2+K, op, A, B, C, C_exp[N*2*2 + M*2 + K]);
                        error_count = error_count + 1;
                    end
                    B = 1;
                end
                A = 1;
                B = 0;
            end
            A = 0;
            B = 0;
            op = op + 1;
        end
        if (error_count == 0) $display("\x1b[38;5;10mAll LU tests passing\x1b[m");
        else $display("\x1b[38;5;9m%0d LU test(s) failing\x1b[m", error_count);
    end
endmodule