import numpy as np, os, sys, re, glob, subprocess, math, unittest, time, shutil, logging, gc, psutil
np.set_printoptions(precision=2)
from quanta_ft import Quanta_FT
import torch
from torch import optim, nn
from random import shuffle, choice
from itertools import product
from functools import partial
# from overlore import Generation
from pathlib import Path
import copy

agent_id = 0
base_folder = './'
try:
	os.mkdir(base_folder+'center_log')
	os.mkdir(base_folder+'agent_log')
	os.mkdir(base_folder+'agent_pool')
	os.mkdir(base_folder+'job_pool')
	os.mkdir(base_folder+'result_pool')
except:
	pass

log_name = base_folder + 'agent_log/{}.log'.format(agent_id)
logging.basicConfig(filename=log_name, filemode='a', level=logging.DEBUG,
										format='%(asctime)s: %(message)s', datefmt='%H:%M:%S')
console = logging.StreamHandler()
console.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s:  %(message)s')
console.setFormatter(formatter)
logging.getLogger('').addHandler(console)

def check_and_load(agent_id):
	file_name = base_folder+'/agent_pool/{}.POOL'.format(agent_id)
	if os.stat(file_name).st_size == 0:
		return False, False
	else:
		with open(file_name, 'r') as f:
			task_name = f.readline()
		return True, task_name

def memory():
	pid = os.getpid()
	py = psutil.Process(pid)
	memoryUse = py.memory_info()[0]/2.**30 
	print('memory use:', memoryUse)

if __name__ == '__main__':

	Path(base_folder+'/agent_pool/{}.POOL'.format(agent_id)).touch()

	while True:
		
		flag, task_name = check_and_load(agent_id)
		
		if flag:

			indv = np.load(base_folder+'/job_pool/{}.npz'.format(agent_id), allow_pickle=True)

			scope = indv['scope'].tolist()
			qtn_seq = indv['qtn_seq']
			if type(qtn_seq.tolist()) == list:
				qtn_layers = [tuple(qtn_seq[:, i]) for i in range(qtn_seq.shape[1])]
			elif type(qtn_seq.tolist()) == dict:
				qtn_layer_dict = qtn_seq.tolist()
				qtn_layers = {}
				for k, v in qtn_layer_dict.items():
					qtn_layers[k] = [tuple(v[:, i]) for i in range(v.shape[1])]

			logging.info('Receiving individual {} with task name {} and layers {}'.format(scope, task_name, qtn_layers))

			try:
				acc_dev, acc_test = Quanta_FT(qtn_layers, task_name)
				dev_loss = 1 - acc_dev
				test_loss = 1 - acc_test
				logging.info('Reporting result {}.'.format(dev_loss))
				np.savez(base_folder+'/result_pool/{}.npz'.format(scope.replace('/', '_')), dev_loss=dev_loss, qtn_seq=qtn_seq, test_loss=test_loss)

				os.remove(base_folder+'/job_pool/{}.npz'.format(agent_id))
				open(base_folder+'/agent_pool/{}.POOL'.format(agent_id), 'w').close()

			except Exception as e:
				os.remove(base_folder+'/agent_pool/{}.POOL'.format(agent_id))
				raise e

			del dev_loss
			del test_loss
			gc.collect()

		time.sleep(1)



