import numpy as np
from argparse import ArgumentParser
import sys
import time

import mySolverModule
from pynq import Overlay
from pynq import Xlnk
from pynq import MMIO

def commandParsing():
	parser = ArgumentParser()
	parser.add_argument('-p', '--problem', help='Problem file name', required=True) # Required
	parser.add_argument('-s', '--seed', help='seed value') # Option
	parser.add_argument('-o', '--output', help='Output file name') # Option
	args = vars(parser.parse_args())
	return args

def main(args):
	print('Download overlay ...', flush=True)
	OL = Overlay('/home/xilinx/pynq/overlays/190826_solver/design_1.bit')
	ip = 'solver_0'
	print(hex(OL.ip_dict[ip]['phys_addr']))

	seed_v = 3
	if args['seed'] is not None: seed_v = int(args['seed'])
	seed = mySolverModule.seed_generator(seed_v) # ndarray
	info = mySolverModule.check_problem(args['problem']) # extract info (block, line, ...)
	
	
	## Global placement
	p = mySolverModule.placer(seed, info['W'][0], info['H'][0], info['blocks'][0], info['line_num'][0], info['block_info'], info['line_info'])
	print(p['block_place'][1:info['blocks'][0]])
	
	
	## Allocate contiguous blocks
	xlnk = Xlnk()
	print(xlnk.cma_stats())
	
	seed_buf = xlnk.cma_array(shape=(3,), dtype=np.uint32)
	B_num_buf = xlnk.cma_array(shape=(1,), dtype=np.int16)
	L_num_buf = xlnk.cma_array(shape=(1,), dtype=np.int16)
	B_info_buf = xlnk.cma_array(shape=(129,5,3), dtype=np.int16)
	W_ext_buf = xlnk.cma_array(shape=(1,), dtype=np.int16)
	H_ext_buf = xlnk.cma_array(shape=(1,), dtype=np.int16)
	B_place_buf = xlnk.cma_array(shape=(129,2), dtype=np.int16)
	result_buf = xlnk.cma_array(shape=(16384,), dtype=np.int16)
	
	print(xlnk.cma_stats())
	
	xlnk.cma_memcopy(seed_buf, p['seed'], 12)
	xlnk.cma_memcopy(B_num_buf, info['blocks'], 2)
	xlnk.cma_memcopy(L_num_buf, info['line_num'], 2)
	xlnk.cma_memcopy(B_info_buf, info['block_info'], 3870)
	xlnk.cma_memcopy(B_place_buf, p['block_place'], 516)
	
	phy_addr_0 = seed_buf.physical_address
	phy_addr_1 = B_num_buf.physical_address
	phy_addr_2 = L_num_buf.physical_address
	phy_addr_3 = B_info_buf.physical_address
	phy_addr_4 = W_ext_buf.physical_address
	phy_addr_5 = H_ext_buf.physical_address
	phy_addr_6 = B_place_buf.physical_address
	phy_addr_7 = result_buf.physical_address
	print(hex(phy_addr_0), hex(phy_addr_1), hex(phy_addr_2), hex(phy_addr_3))
	print(hex(phy_addr_4), hex(phy_addr_5), hex(phy_addr_6), hex(phy_addr_7))
	
	
	## Write pointers
	lite_base_addr = OL.ip_dict[ip]['phys_addr'] # IP_BASE_ADDRESS
	lite_size = OL.ip_dict[ip]['addr_range'] # ADDRESS_RANGE
	mmio = MMIO(lite_base_addr, lite_size)
	
	mmio.write(0x18, phy_addr_0)
	mmio.write(0x20, phy_addr_1)
	mmio.write(0x28, phy_addr_2)
	mmio.write(0x30, phy_addr_3)
	mmio.write(0x38, phy_addr_4)
	mmio.write(0x40, phy_addr_5)
	mmio.write(0x48, phy_addr_6)
	mmio.write(0x50, phy_addr_7)
	
	
	## Local placement
	print('Solver start ...', flush=True)
	mmio.write(0x00, 1)
	while mmio.read(0x00) != 6:
		time.sleep(0.1)
	
	print(bin(mmio.read(0x00)))
	print(mmio.read(0x10))
	
	if mmio.read(0x10) == 0:
		print('Fail routing m(_ _)m')
		sys.exit(1)
	
	W_ext = np.zeros((1,), dtype=np.int16)
	H_ext = np.zeros((1,), dtype=np.int16)
	result = np.zeros((16384,), dtype=np.int16)
	xlnk.cma_memcopy(W_ext, W_ext_buf, 2)
	xlnk.cma_memcopy(H_ext, H_ext_buf, 2)
	xlnk.cma_memcopy(result, result_buf, 32768)
	xlnk.cma_memcopy(info['block_info'], B_info_buf, 3870)
	
	mySolverModule.show_result(W_ext[0], H_ext[0], info['blocks'][0], info['line_num'][0], info['block_info'], result)

	if W_ext[0] > info['W'][0] or H_ext[0] > info['H'][0]:
		print('Fail satisfying constraint T_T')
		sys.exit(1)
	else:
		print('Satisfy constraint ^_^')
		if args['output'] is not None: mySolverModule.output_to_file(args['output'], W_ext[0], H_ext[0], info['blocks'][0], info['line_num'][0], info['block_info'], result)
		sys.exit(0)

if __name__ == '__main__':

	args = commandParsing()
	main(args)