module floating_add #(parameter N=16, M=4)(input_1, input_2, sum, diff, clk, reset);
	input logic [N-1:0] input_1, input_2;
	input logic clk, reset;
	output logic [N-1:0] sum;
	output logic [M:0] diff;

//	logic flag_a;
//	logic flag_b;
//	logic [M:0] abs;
//	logic [N-3-M:0] res;
	
	logic [N-1:0] D0 [7:0];
	logic [N-1:0] Q0 [7:0];
	logic [N-1:0] Q1 [7:0];
	logic [N-1:0] Q2 [7:0];
	
	// sign_x = x[N-1]
	// exponent_x = x[N-2:N-2-M]
	// mantissa_x = x[N-3-M:0]
	
	//First pipeline stage
	always_comb
		begin
			D0[0] = input_1;
			D0[1] = input_2;
			D0[2] = 0; // sum
			D0[3] = 0; // diff
			D0[4] = 0; // flag_a
			D0[5] = 0; // flag_b
			D0[6] = 0; // abs
			D0[7] = 0; // res
		end
	pipe#(.N(N-1), .K(7)) pipe0(.clk(clk), .reset(reset), .D(D0), .Q(Q0));
	
	
	always_comb
		begin
			if (Q0[0][N-2:N-2-M] > Q0[1][N-2:N-2-M]) // If input 1 has the bigger exponent 
				begin
					// Flags input a as larger and calculates the absolute difference
					Q0[4] = 1;
					Q0[5] = 0;
					Q0[6] = Q0[0][N-2:N-2-M] - Q0[1][N-2:N-2-M];
					// ASsigning overall sign of the output
					Q0[2][N-1] = Q0[0][N-1];
					// Sets output to have the same exponent
					Q0[2][N-2:N-2-M] = Q0[0][N-2:N-2-M];
				end
			else if (Q0[1][N-2:N-2-M] > Q0[0][N-2:N-2-M]) // If input 2 has the bigger exponent
				begin
					// Similarly flags input b as larger and calculates the absolute difference
					Q0[4] = 0;
					Q0[5] = 1;
					Q0[6] = Q0[1][N-2:N-2-M] - Q0[0][N-2:N-2-M];
					// ASsigning overall sign of the output
					Q0[2][N-1] = Q0[1][N-1];
					// Sets ouput to have the same exponent
					Q0[2][N-2:N-2-M] = Q0[1][N-2:N-2-M];
				end
			else 
				begin
					// THe condition that both inputs have the same exponent
					Q0[4] = 1;
					Q0[5] = 1;
					Q0[6] = 0;
					// ASsigning overall sign of the output based on size of the mantissa
					if (Q0[0][N-3-M:0] >= Q0[1][N-3-M:0]) Q0[2][N-1] = Q0[0][N-1];
					else Q0[2][N-1] = Q0[1][N-1];
					Q0[2][N-2:N-2-M] = Q0[0][N-2:N-2-M];
				end
			Q0[3] = Q0[6];
		end
		
		//Second pipeline stage 1
		pipe#(.N(N-1), .K(7)) pipe1(.clk(clk), .reset(reset), .D(Q0), .Q(Q1));
		
	always_comb
		begin
			// Condition for overflow is that it sets the output to the larger input
			if (Q1[6] > N-2-M) // Because size of mantissa is 10 bits and shifting by 10 would give 0
				begin
					if (Q1[4] & ~Q1[5]) Q1[2] = Q1[0]; // input 1 is larger and is translated to output
					else if (~Q1[4] & Q1[5]) Q1[2] = Q1[1]; // input 2 is larger and is translated to output
					else // exponents are the same
						begin
							if (Q1[6][N-3-M:0] >= Q1[1][N-3-M:0]) Q1[2] = Q1[0];// input 1 has the bigger mantissa
							else Q1[2] = Q1[1]; // input 2 has the bigger mantissa
						end
				end
			else
				begin
					// Shifts the smaller input's mantissa to the right based on abs
					if (Q1[4] & ~Q1[5])// If input 1 has the larger exponent
						begin
							// If the signs of both inputs are the same you add, otherwise you subtract
							if (Q1[0][N-1] == Q1[1][N-1])
								begin
									Q1[7] = Q1[0][N-3-M:0] + (Q1[1][N-3-M:0] >> Q1[6]-1); // Sum the mantissa
									Q1[2][N-3-M:0] = Q1[7];
								end
							else
								begin
									Q1[7] = Q1[0][N-3-M:0] - (Q1[1][N-3-M:0] >> Q1[6]-1); // Subtract the mantissas
									Q1[2][N-3-M:0] = Q1[7];
								end
						end
					else if (~Q1[4] & Q1[5])
						begin
							// If the signs of both inputs are the same you add, otherwise you subtract
							if (Q1[0][N-1] == Q1[1][N-1])
								begin
									Q1[7] = (Q1[0][N-3-M:0] >> Q1[6]-1) + Q1[1][N-3-M:0]; // Sum the mantissa
									Q1[2][N-3-M:0] = Q1[7];
								end
							else
								begin
									Q1[7] = Q1[1][N-3-M:0] - (Q1[0][N-3-M:0] >> Q1[6]-1); // Subtract the mantissas
									Q1[2][N-3-M:0] = Q1[7];
								end
						end
					else
						begin 
							if (Q1[0][N-1] == Q1[1][N-1]) // If exponents and signs equal
								begin
									Q1[7] = Q1[0][N-3-M:0] + Q1[1][N-3-M:0]; // Sum the mantissa
									Q1[2][N-3-M:0] = Q1[7];
								end
							else // In this case it will be a subtraction
								begin
									if (Q1[0][N-3-M:0] > Q1[1][N-3-M:0]) // Which has the larger mantissa 
										begin
											Q1[7] = Q1[0][N-3-M:0] - Q1[1][N-3-M:0]; // Subtract the mantissa
											Q1[2][N-3-M:0] = Q1[7];
										end
									else if (Q1[0][N-3-M:0] < Q1[1][N-3-M:0])
										begin
											Q1[7] = Q1[1][N-3-M:0] - Q1[0][N-3-M:0]; // Subtract the mantissa
											Q1[2][N-3-M:0] = Q1[7];
										end
									else Q1[7] = 0; // Both the exponent and the mantissa are equal so subtraction leads to 0
									Q1[2][N-3-M:0] = Q1[7];
								end
						end
				end
		end
		
		// Final pipeline stage 
		pipe#(.N(N-1), .K(7)) pipe2(.clk(clk), .reset(reset), .D(Q1), .Q(Q2));
		assign sum = Q2[2];
		assign diff = Q2[3];
endmodule : floating_add



module floating_product #(parameter N=16, M=4)(input_1, input_2, product);
	input logic [N-1:0] input_1, input_2;
	output logic [N-1:0] product;

	// sign_x = x[N-1]
	// exponent_x = x[N-2:N-2-M]
	// mantissa_x = x[N-3-M:0]

//	logic [N-2:N-2-M] sum;
//	logic [2*(N-3-M):0] mult;
	logic [2*(N-3-M):0] D0 [4:0];
	logic [2*(N-3-M):0] Q0 [4:0];
	logic [2*(N-3-M):0] Q1 [4:0];
	logic [2*(N-3-M):0] Q2 [4:0];
	
	//First pipeline stage
	always_comb
		begin
			D0[0] = input_1;
			D0[1] = input_2;
			D0[2] = 0; // product
			D0[3] = 0; // sum
			D0[4] = 0; // mult
		end
	pipe#(.N(2*N-3-M), .K(4)) pipe0(.clk(clk), .reset(reset), .D(D0), .Q(Q0));

	// We have assigned an {M+1} bit exponent so we must have a 2^{M} offset
	assign Q0[3] = Q0[0][N-2:N-2-M] + Q0[1][N-2:N-2-M];
	assign Q0[2][N-2:N-2-M] = Q0[3] - (1'b1 << M) + 2;
	
	//Second Pipeline stage
	pipe#(.N(2*N-3-M), .K(4)) pipe1(.clk(clk), .reset(reset), .D(Q0), .Q(Q1));

	always_comb
		begin
				// Setting the mantissa of the output
				Q1[4] = Q1[0][N-3-M:0] * Q1[1][N-3-M:0];
				if (Q1[4][N-3-M]) Q0[2][N-3-M:0] = Q1[4][2*(N-3-M):2*(N-3-M)-9];
				else Q1[2][N-3-M:0] = Q1[4][2*(N-3-M):2*(N-3-M)-9] << 1;
				Q1[2][N-1] = Q1[0][N-1] ^ Q1[1][N-1];
		end
	
	//Final Pipeline stage
	pipe#(.N(2*N-3-M), .K(4)) pipe2(.clk(clk), .reset(reset), .D(Q1), .Q(Q2));
	assign product = Q2[4][N-1:0];
endmodule : floating_product



module pipe #(parameter N, K)(clk, reset, Q, D);
	input logic clk, reset;
	input reg [N:0] D [K:0];
	output reg [N:0] Q [K:0];
	logic [N:0] in_pipe [K:0];
	
	always_ff @(posedge clk)
		begin
			if(reset) 
				begin
					in_pipe <= 0;
					Q <= 0;
				end
			else 
				begin
					in_pipe <= D;
					Q <= in_pipe;
				end
		end
endmodule : pipe



module floating_tb;
	reg reset, clk;
	logic [15:0] input_a, input_b, result_add, result_mult;
	logic [4:0] diff;

	floating_add adder1(.input_1(input_a), .input_2(input_b), .sum(result_add), .diff(diff), .clk(clk), .reset(reset));

	floating_product multiplier1(.input_1(input_a), .input_2(input_b), .product(result_mult), .clk(clk), .reset(reset));


	reg [15:0] test_mem [29:0][3:0];

	initial $readmemh("../../scripts/fp16_test.hex", test_mem);


	initial begin
        static int num_err = 0;
        static int num_tests = $size(test_mem) * 2;

        for (int i=0; i < $size(test_mem); i++) begin
            input_a = test_mem[i][0];
            input_b = test_mem[i][1];

            #10;
            if(result_add != test_mem[i][2]) begin
                if(num_err < 20)
                    $display("FAIL ADD: %H + %H = %H, expected %H", input_a, input_b, result_add, test_mem[i][2]);
                num_err = num_err + 1;
            end

			if(result_mult != test_mem[i][3]) begin
                if(num_err < 20)
                    $display("FAIL MULTIPLY: %H + %H = %H, expected %H", input_a, input_b, result_mult, test_mem[i][3]);
                num_err = num_err + 1;
            end

        end
        $display("Passed %d of %d tests", num_tests-num_err, num_tests);
        $finish();
	end
endmodule : floating_tb



module floating32_tb;
	reg reset, clk;
	logic [31:0] input_a, input_b, result_add, result_mult;

	floating_add#(.N(32), .M(8)) add0(
		.input_1(input_a), .input_2(input_b), .sum(result_add), .diff()
	);
	floating_product#(.N(32), .M(8)) mult0(
		.input_1(input_a), .input_2(input_b), .product(result_mult)
	);

	reg [31:0] test_mem [29:0][3:0];

	initial $readmemh("scripts/fp32_test.hex", test_mem);


	initial begin
		static int num_err = 0;
		static int num_tests = $size(test_mem) * 2;

		for (int i=0; i < $size(test_mem); i++) begin
			input_a = test_mem[i][0];
			input_b = test_mem[i][1];

			#10;
			if(result_add != test_mem[i][2]) begin
				if(num_err < 20)
					$display("FAIL ADD: %H + %H = %H, expected %H", input_a, input_b, result_add, test_mem[i][2]);
				num_err = num_err + 1;
			end

			if(result_mult != test_mem[i][3]) begin
				if(num_err < 20)
					$display("FAIL MULTIPLY: %H + %H = %H, expected %H", input_a, input_b, result_mult, test_mem[i][3]);
				num_err = num_err + 1;
			end

		end
		$display("Passed %d of %d tests", num_tests-num_err, num_tests);
		$finish();
	end
endmodule : floating32_tb