`include "shifter.sv"
module shifter_tb #(width=4)();
    logic [2:0] op;
    logic [2**width-1:0] Y;
    logic [2**width-1:0] A_in;
    logic [2**width-1:0] B_in;
    logic [width-1:0] C_in;

    logic [2:0] ops[6];
    logic [2**width-1:0] A[6];
    logic [2**width-1:0] B[6];
    logic [width-1:0] C[6];
    logic [2**width-1:0] Y_exp[6];

    shifter #(width) dut(op, A_in,B_in,C_in,Y);

    integer error_count;
    initial begin
        // Test left shift by C
        ops[0] = 3'b000;
        A[0] = {{(2**width-1){1'b0}}, {1'b1}};
        C[0] = {{(width-1){1'b0}}, {1'b1}};
        B[0] = {2**width{1'b0}};
        Y_exp[0] = {{(2**width-2){1'b0}}, {2'b10}};

        // Test left shift by B
        ops[1] = 3'b100;
        A[1] = {{(2**width-1){1'b0}}, {1'b1}};
        C[1] = {{(width-1){1'b0}}, {1'b1}};
        B[1] = {{(2**width-2){1'b0}}, 2'b10};
        Y_exp[1] = {{(2**width-3){1'b0}}, {3'b100}};

        // Test Logical right shift by C
        ops[2] = 3'b001;
        A[2] = {{(2**width-2){1'b0}}, {2'b10}};
        C[2] = {{(width-1){1'b0}}, {1'b1}};
        B[2] = {2**width{1'b0}};
        Y_exp[2] = {{(2**width-1){1'b0}}, {1'b1}};

        // Test Logical right shift by B
        ops[3] = 3'b101;
        A[3] = {{(2**width-1){1'b0}}, {2'b10}};
        C[3] = {{(width-1){1'b0}}, {1'b1}};
        B[3] = {{(2**width-2){1'b0}}, 2'b10};
        Y_exp[3] = {{(2**width-1){1'b0}}, {1'b0}};

        // Test arithmetic shift right by C
        ops[4] = 3'b011;
        A[4] = {{(2**width-1){1'b1}}, 1'b0};   // -2
        C[4] = {{(width-1){1'b0}}, 1'b1};   // 1
        B[4] = {{(2**width-2){1'b0}}, 2'b10};  // 2
        Y_exp[4] = {2**width{1'b1}};           // -1

        // Test arithmetic shift right by B
        ops[5] = 3'b111;
        A[5] = {{(2**width-2){1'b1}}, 2'b00};   // -4
        C[5] = {{(width-1){1'b0}}, 1'b1};   // 1
        B[5] = {{(2**width-2){1'b0}}, 2'b10};  // 2
        Y_exp[5] = {2**width{1'b1}};           // -1

        error_count = 0;
        $display("Beginning Shifter tests");
        for (integer N = 0; N < $size(ops); ++N) begin
            op = ops[N];
            A_in = A[N];
            B_in = B[N];
            C_in = C[N];
            #1;
            if (Y != Y_exp[N]) begin
                $display("\t\x1b[38;5;9mShifter(%0d) Op 0b%b(0b%b, 0b%b, 0b%b) produced 0b%b but 0b%b was expected\x1b[m", N, op, A_in, B_in, C_in, Y, Y_exp[N]);
                error_count = error_count + 1;
            end
        end
        if (error_count == 0) $display("\x1b[38;5;10mAll Shifter tests passing\x1b[m");
        else $display("\x1b[38;5;9m%0d Shifter test(s) failing\x1b[m", error_count);
    end
endmodule