#!/usr/bin/python3
import argparse

def gen_half_adder():
	print(f"---------- Half adder generation ----------")
	f = open(f"half_adder.sv", "w")
	f.write(f"""module half_adder (
	input logic a,
	input logic b,
	output logic sum,
	output logic carry
);

assign sum = a ^ b;
assign carry = a & b;

endmodule
""")
	f.close()

def gen_full_adder():
	print(f"---------- Full adder generation ----------")
	f = open(f"full_adder.sv", "w")
	f.write(f"""module full_adder (
	input logic a,
	input logic b,
	input logic c,
	output logic sum,
	output logic carry
);

assign sum = a ^ b ^ c;
assign carry = (a & b) | (c & (a ^ b));

endmodule
""")
	f.close()

def gen_multiplier(bits):
	print(f"\n---------- {bits} Bit Top Level Multiplier Generation ----------")
	f = open(f"multiplier.sv", "w")
	f.write(f"""module multiplier(
	input  logic [{bits-1}:0] a,
	input  logic [{bits-1}:0] b,
	output logic [{2*bits-1}:0] c
);

	logic [{bits-1}:0] partial_prod [0:{bits-1}];
	logic [{2*bits-1}:0] partial_sum;

	assign c = partial_sum;

	wallace_adder wadder0(partial_prod, partial_sum);
	partial_products partprod0(a, b, partial_prod);

endmodule
""")
	f.close()

def gen_partial_products(bits):
	print(f"------------ {bits} Bit Partial Products Generation ------------")
	f = open(f"partial_products.sv", "w")
	f.write(f"""module partial_products
(
	input  logic [{bits-1}:0] a,
	input  logic [{bits-1}:0] b,
	output logic [{bits-1}:0] c [0:{bits-1}]
);

always @ (*) begin
	integer i;
	for (i = 0; i < {bits}; i=i+1) begin
		c[i][{bits-1}:0] = {{{bits}{{b[i]}}}} & a;
	end
end

endmodule
""")

	f.close()

def add_half_adder(reduction_layers, instantiations, net_names, col, col_idx, curr_layer, prev_layer, debug):
	# Generates nets and updates the layer net array
	curr_col_net_idx = len(reduction_layers[curr_layer][col_idx])
	next_col_net_idx = len(reduction_layers[curr_layer][col_idx+1])
	cout = f"layer{curr_layer}_col{col_idx}_net{curr_col_net_idx}"
	sout = f"layer{curr_layer}_col{col_idx+1}_net{next_col_net_idx}"

	reduction_layers[curr_layer][col_idx].append(cout)
	reduction_layers[curr_layer][col_idx+1].append(sout)

	# Adds half adder to instantiations
	net_names[prev_layer].append(cout)
	net_names[prev_layer].append(sout)
	a, b = [col.pop() for i in range(2)]
	instantiations[prev_layer].append(f"half_adder ha_add{curr_layer}_{len(instantiations[prev_layer])} ({a}, {b}, {cout}, {sout});")

	# Debug Print half adder
	if debug:
		print(f"half_adder ha_add{curr_layer}_{len(instantiations[prev_layer])} ({a}, {b}, {cout}, {sout});")

	pass

def add_full_adder(reduction_layers, instantiations, net_names, col, col_idx, curr_layer, prev_layer, debug):
	# Generates nets and updates the layer net array
	curr_col_net_idx = len(reduction_layers[curr_layer][col_idx])
	next_col_net_idx = len(reduction_layers[curr_layer][col_idx+1])
	cout = f"layer{curr_layer}_col{col_idx}_net{curr_col_net_idx}"
	sout = f"layer{curr_layer}_col{col_idx+1}_net{next_col_net_idx}"
	reduction_layers[curr_layer][col_idx].append(cout)
	reduction_layers[curr_layer][col_idx+1].append(sout)

	# Adds nets and adders to be instantiated
	net_names[prev_layer].append(cout)
	net_names[prev_layer].append(sout)
	a, b, cin = [col.pop() for i in range(3)]
	instantiations[prev_layer].append(f"full_adder fa_add{curr_layer}_{len(instantiations[prev_layer])} ({a}, {b}, {cin}, {cout}, {sout});")

	# Debug Print half adder
	if debug:
		print(f"full_adder fa_add{curr_layer}_{len(instantiations[prev_layer])} ({a}, {b}, {cin}, {cout}, {sout});")

def add_passthrough(reduction_layers, instantiations, net_names, col, col_idx, curr_layer, prev_layer, debug):
# Assigns passthrough for remaining logic and updates the counter
	curr_col_net_idx = len(reduction_layers[curr_layer][col_idx])
	passthrough = f"layer{curr_layer}_col{col_idx}_net{curr_col_net_idx}"

	# Adds passthrough to netlist array
	reduction_layers[curr_layer][col_idx].append(passthrough)

	# Adds the assign statement to passthrough
	net_names[prev_layer].append(passthrough)
	input_net1 = col.pop()
	instantiations[prev_layer].append(f"assign {passthrough} = {input_net1};")

	if debug:
		print(f"assign {passthrough} = {input_net1};")

def gen_adder_tree(bits, debug):
	print(f"--------------- {bits} Bit Adder Tree Generation ---------------")

	# Parameters of the adder tree generate script
	num_cols = (2 * bits)
	layer_limit = 50

	# Initialize reduction layer array
	reduction_layers = []

	# Initialize instantiations and net names
	ha_instantiations = []
	fa_instantiations = []
	pass_instantiations = []
	net_names = []

	# Partial layer is the "zeroeth" reduction layer, initialize it
	curr_layer = 0
	reduction_layers.append([[] for i in range(num_cols)])

	# Fill up partial layer
	for i in range(bits):
		for j in range(bits):
			reduction_layers[curr_layer][i+j].append(f"partial_prod[{i}][{j}]")

	# Debug partial layer print
	if debug:
		print(f"\n--------- LAYER {curr_layer} -------------")
		for col_idx, reduce in enumerate(reduction_layers[curr_layer]):
			print(f"Col: {col_idx}, Length: {len(reduce)}, {reduce}")

	# Build out subsequent reduction layers
	curr_layer = 1
	prev_layer = 0

	instantiation_idx = 0

	# Run until we can add the remaining bit vectors together or non-convergent solution
	while (len(max(reduction_layers[prev_layer], key=len)) > 2 and curr_layer < layer_limit):
		# Allocate next layer
		if debug:
			print(f"--------- LAYER {prev_layer} -------------")
		reduction_layers.append([[] for i in range(num_cols)])
		pass_instantiations.append([])
		ha_instantiations.append([])
		fa_instantiations.append([])
		net_names.append([])

		carry_propogation = 0
		extra_ha = (len(max(reduction_layers[prev_layer], key=len)) == 3)
		fa_used = False
		ha_used = False

		# Counts how many bits need to be eventually removed by this bit
		for col_idx, col in enumerate(reduction_layers[prev_layer]):

			# Check that this is actually solvable using only 2*bits output
			if (col_idx+1 == len(reduction_layers[prev_layer]) and (len(col) + carry_propogation) > 2):
				print("Cannot SOLVE")
				return -1

			# Debug print for this column
			if debug:
				print(f"Index: {col_idx}, Length: {len(col)}")

			next_layer_size = carry_propogation
			carry_propogation = 0

			# Add full adders if needed
			while (len(col) > 3):
				add_full_adder(reduction_layers, fa_instantiations, net_names, col, col_idx, curr_layer, prev_layer, debug)
				fa_used = True
				# 1 carry will go to the next column next layer, and the sum will go to this col next layer
				carry_propogation += 1
				next_layer_size += 1

			if (len(col) == 3):
				# Only add half adder if no propogations or other adders created
				if (fa_used == False and (ha_used == False or extra_ha == True) and next_layer_size == 0):
					add_half_adder(reduction_layers, ha_instantiations, net_names, col, col_idx, curr_layer, prev_layer, debug)
					ha_used = True
				else:
					fa_used = True
					add_full_adder(reduction_layers, fa_instantiations, net_names, col, col_idx, curr_layer, prev_layer, debug)

				# Increment the propogation and current size
				carry_propogation += 1
				next_layer_size += 1

			if (len(col) == 2):
				# Only add half adder if there is propogation from previous columns and if there is less than three in the next col, else pass through both
				if (fa_used == False and (ha_used == False or extra_ha == True) and next_layer_size == 1 ):
					add_half_adder(reduction_layers, ha_instantiations, net_names, col, col_idx, curr_layer, prev_layer, debug)
					ha_used = True
					carry_propogation += 1
					next_layer_size += 1
				else:
					add_passthrough(reduction_layers, pass_instantiations, net_names, col, col_idx, curr_layer, prev_layer, debug)
					add_passthrough(reduction_layers, pass_instantiations, net_names, col, col_idx, curr_layer, prev_layer, debug)

			if (len(col) == 1):
				add_passthrough(reduction_layers, pass_instantiations, net_names, col, col_idx, curr_layer, prev_layer, debug)

		# Update the layer indices
		prev_layer = curr_layer
		curr_layer += 1

		# Debug reduction layer print
		if debug:
			for col_idx, reduce in enumerate(reduction_layers[prev_layer]):
				print(f"Col: {col_idx}, Length: {len(reduce)}, {reduce}")


	# Debug final reduction layer to be added
	if debug:
		print("\n--------- BIT PAIRS ----------")
		add_layer = list(zip(reduction_layers[prev_layer]))
		add_layer.reverse()
		for bit_pair in add_layer:
			print(bit_pair[0])

	# Add the two remaining rows of bits at the end
	bit_vector_0 = "{ "
	bit_vector_1 = "{ "
	for bit_pair_idx, bit_pair in enumerate(reversed(list(zip(reduction_layers[prev_layer])))):

		# Exclude MSB if no overflows to it
		if (len(bit_pair[0]) == 0 and bit_pair_idx == 0):
			continue

		# Generate bit string for both vectors, order doesn't matter here
		bit_vector_0 += f"{bit_pair[0][0]}, "
		if (len(bit_pair[0]) == 2):
			bit_vector_1 += f"{bit_pair[0][1]}, "
		else:
			bit_vector_1 += "1'b0, "

	bit_vector_0 = bit_vector_0[:-2] + "}"
	bit_vector_1 = bit_vector_1[:-2] + "};"

	f = open(f"wallace_adder.sv", "w")

	# Start by printing module declaration
	f.write(f"module wallace_adder (\n")
	f.write(f"\tinput  logic [{bits-1}:0] partial_prod[0:{bits-1}],\n")
	f.write(f"\toutput logic [{2*bits-1}:0] partial_sum\n")
	f.write(");\n\n")

	# Print out net names
	for net_layer in net_names:
		netstring = "logic "
		net_idx_len = len(net_layer)
		for net_idx, net in enumerate(net_layer):
			if (net_idx != net_idx_len - 1):
				netstring += f"{net}, "
			else:
				netstring += f"{net};"
		f.write(netstring + '\n')

	# Print out entire reduction tree and calculate stats
	ha_count = 0
	fa_count = 0
	for layer in range(len(pass_instantiations)):
		f.write(f"\n//----------- Reduction Layer {layer+1} Start --------------\n\n")
		for passthrough in pass_instantiations[layer]:
			f.write(passthrough + '\n')

		for half_adder in ha_instantiations[layer]:
			ha_count += 1
			f.write(half_adder + '\n')

		for full_adder in fa_instantiations[layer]:
			fa_count += 1
			f.write(full_adder + '\n')

	# Print final two number adder

	f.write(f"\n//----------- Adding Layer Start --------------\n\n")
	f.write(f"assign partial_sum = {bit_vector_0} + {bit_vector_1}\n")

	# Endmodule
	f.write("\nendmodule\n")

	f.close()

	# Print stats of the wallace adder
	print(f"{ha_count} Half Adders Used")
	print(f"{fa_count} Full Adders Used")
	print(f"{len(pass_instantiations)-1} Reduction Layers")

def main():
	parser = argparse.ArgumentParser(prog="Multiplier Generator", description="Generates a n bit multiplier based on the bits argument provided",
				epilog="bits sets the bit width of the multiplier, the output of the multiplier is 2 times the number of bits")
	parser.add_argument("bits", type=int, help="The bit width of the multiplier")
	parser.add_argument("-a", "--adder", help="Generates the full and half adders for you to use", action='store_true')
	parser.add_argument("-d", "--debug", help="Enables debug prints during generation scripting", action='store_true')
	args = parser.parse_args()

	bits = args.bits
	debug = args.debug
	adder  = args.adder

	if (adder):
		gen_half_adder()
		gen_full_adder()

	gen_multiplier(bits)
	gen_partial_products(bits)

	if (gen_adder_tree(bits, debug) == -1):
		return -1

	print("----------- GENERATION COMPLETE WITHOUT ERROR ----------- \n\n")

if __name__ == "__main__":
	main()