import json
import open3d as o3d
import numpy as np
import os
import trimesh
import zipfile
from tqdm import tqdm
import matplotlib.pyplot as plt
plt.style.use('bmh')

default_color = [0,0.5,1]
cube = np.array([
    [0,0,0], [1,0,0], [1,1,0], [0,1,0],
	[0,0,1], [1,0,1], [1,1,1], [0,1,1],
])

'''plt figure'''
def plt_show_save(data, title, save_path=None, xname='', bins=50):
	plt.cla()
	plt.figure(figsize=(12,9))
	if type(data) == dict:
		plt.bar(data.keys(), data.values())
		# plt.xticks(rotation=90)
	else:
		plt.hist(data, bins=bins)
	plt.title(title)
	plt.ylabel('value')
	plt.xlabel(xname)
	if save_path is not None:
		plt.savefig(save_path)
	else:
		plt.show()

def get_pcd(pc, color=default_color):
	pc = np.array(pc)

	pcd = o3d.geometry.PointCloud()
	pcd.points = o3d.utility.Vector3dVector(pc)
	pcd.paint_uniform_color(color) # 默认是彩虹色过渡，这里指定染色
	return pcd

def get_o3d_FOR(origin=[0, 0, 0],size=0.1):

	mesh_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=size)
	mesh_frame.translate(origin)
	return(mesh_frame)

def show_pcds(pcds, wname='Open3D', FOR=0.1):

	if FOR:
		pcds.append(get_o3d_FOR(size = FOR))
	o3d.visualization.draw_geometries(pcds, width=800, height=800, window_name=wname)

def csv2box(csv_path):

	obb_info = np.loadtxt(open(csv_path, 'r'),delimiter = ",") # (5,4)

	center = obb_info[0,:3]
	dirs = 0.5 * (obb_info[2:,:3] * obb_info[2:,-1].reshape(3,1) )
	val = cube*2 - 1
	vec = np.matmul(val, dirs) # (8,3)@(3,3)
	corner = center.reshape(1,3) + vec
	return corner,dirs

def add_thickness(pc, direction, scale):
	direction = direction / np.linalg.norm(direction)
	noise = np.random.normal(0, scale, (pc.shape[0],1))
	return pc + noise * direction.reshape(1,3)

def PCA(data, sort=True):
    average_data = np.mean(data,axis=0)
    decentration_matrix = data - average_data
    H = np.dot(decentration_matrix.T,decentration_matrix)
    eigenvectors,eigenvalues,eigenvectors_T = np.linalg.svd(H)

    if sort:
        sort = eigenvalues.argsort()[::-1]
        eigenvalues = eigenvalues[sort]
        eigenvectors = eigenvectors[:, sort]
    return eigenvalues, eigenvectors

def box_from_pc(pc, color=default_color, aabb=False, return_corner=True):
	pcd = get_pcd(pc)

	box = pcd.get_axis_aligned_bounding_box() if aabb else \
			pcd.get_oriented_bounding_box()
	if return_corner:
		corner = np.array(box.get_box_points())
		return corner
	else:
		box.color = color
		return box

def box_from_corner(corner, color=default_color, aabb=False):
	corner = np.asarray(corner)
	box = o3d.geometry.AxisAlignedBoundingBox() if aabb else \
		o3d.geometry.OrientedBoundingBox()
	box = box.create_from_points(o3d.utility.Vector3dVector(corner))
	box.color = color
	return box

def box2cen_dir(box:np.ndarray):
	centers = np.zeros((6,3))
	sorted_box = sort_pts(box)

	v1 = sorted_box[1]-sorted_box[0]
	v2 = sorted_box[3]-sorted_box[0]
	cos = v1@v2 / (np.linalg.norm(v1) * np.linalg.norm(v2))
	if abs(cos) < 0.001:
		tmp = sorted_box[3].copy()
		sorted_box[3] =sorted_box[4]
		sorted_box[4] = tmp
	# 0246, 0145
	centers[0] = sorted_box[:4].mean(axis=0)
	centers[1] = sorted_box[[0,2,4,6]].mean(axis=0)
	centers[2] = sorted_box[[0,1,4,5]].mean(axis=0)
	centers[3:] = 2 * box.mean(0).reshape(1,3) - centers[:3]
	return centers

def box2dir(box:np.ndarray):
	sorted_box = np.array(sorted(box, key = lambda x:x[0]) )
	dirs3 = sorted_box[1:4] - sorted_box[0].reshape(1,-1)
	cos = cosine(dirs3, dirs3).flatten()
	idx = np.argmin(cos)
	if cos[idx]<1e-3:
		d1 = idx//3
		d2 = idx%3
		left_dir = np.cross(dirs3[d1], dirs3[d2])
		return np.vstack([dirs3[d1], dirs3[d2], left_dir])
	else:
		return None

def aabb_dirs(pc):
	mins = pc.min(0)
	maxs = pc.max(0)
	dirs = np.eye(3,3) * (maxs-mins).reshape(1,3) / 2
	center = (mins + maxs) / 2
	corners = center.reshape(1,3) + (cube*2-1)@dirs
	return corners, dirs

def obb_2dirs(pc, axis, return_corner=True):
	else_axis = [0,1,2]
	else_axis.pop(axis)
	sub_pc = pc[:,else_axis]

	cov_pts = np.cov(sub_pc, y=None, rowvar=False, bias=True)
	v, vect = np.linalg.eig(cov_pts)
	tvect = vect.T
	rpts = np.dot(sub_pc, np.linalg.inv(tvect))
	mina = np.min(rpts, 0)
	maxa = np.max(rpts, 0)
	diff = (maxa - mina)*0.5

	center = mina + diff
	corners =  center.reshape(-1,2) + np.array([
		[-1,-1], [1,-1], [1,1], [-1,1]
	]) * diff.reshape(-1,2)

	corners = np.dot(corners, tvect) # (4,2)
	axis_pc = pc[:, axis]
	axis_min,axis_max = axis_pc.min(), axis_pc.max()
	cor1 = np.insert(corners, axis, axis_min, axis=1)
	cor2 = np.insert(corners, axis, axis_max, axis=1)
	corners = np.vstack([cor1,cor2])
	center = corners.mean(0)

	dirs = (corners[[1,3,4]] - corners[0].reshape(1,3))/2
	if return_corner:
		return corners, dirs
	else:
		return center, dirs

def obb_adjust(pc:np.ndarray, fix_dir:np.array, ori_dir:np.array):
	'''ori_dir should be [0,0,1] or [0,1,0] or [1,0,0]'''
	axis = np.argmax(ori_dir)

	fix_dir = fix_dir / np.linalg.norm(fix_dir)
	ori_dir = ori_dir / np.linalg.norm(ori_dir)
	cro = np.cross(ori_dir, fix_dir)
	cos = ori_dir@fix_dir
	if abs(cos)>0.99:
		return obb_2dirs(pc, axis, True)

	vx = np.array([
		[0,		-cro[2],	cro[1]],
		[cro[2],	0,		-cro[0]],
		[-cro[1],	cro[0],	0	]
	])

	rot_w = np.eye(3,3) + vx + np.matmul(vx,vx) / (1+cos)
	rot_verse = np.linalg.inv(rot_w)

	rot_pc = np.matmul(pc, rot_verse.T)
	center, dirs = obb_2dirs(rot_pc, axis, False)
	# dirs[-1][:2] = 0
	# dirs[-1,-1] = rot_pc[:,axis].max() - rot_pc[:,axis].min()
	cen = center.reshape(-1,3)
	dirs = np.matmul(dirs, rot_w.T)
	box = (cube*2 - 1)@dirs + cen@rot_w.T

	return box, dirs

def pts2pts_dis(pts1,pts2):
	diff = pts1.reshape((-1, 1, 3)) - pts2.reshape((1, -1, 3))
	distance = (diff**2).reshape((-1,3)).sum(axis=-1)
	return distance

def sort_pts(box):
	uniques = []
	for i in range(3):
		uni = np.unique(box[:,i]).shape[0]
		uniques.append(uni<8) # and uni//2==0
	if sum(uniques)==0: uniques[0] = True
	sorted_box = np.array(sorted(box, key = lambda x:x[uniques].sum()))
	return sorted_box

def pc2mesh(pts):

	pts = np.asarray(pts)
	pcd = get_pcd(pts)
	pcd.estimate_normals()

	distances = pcd.compute_nearest_neighbor_distance()
	avg_dist = np.mean(distances)
	radius = 1.5 * avg_dist

	mesh = o3d.geometry.TriangleMesh.create_from_point_cloud_ball_pivoting(
			pcd, o3d.utility.DoubleVector([radius, radius * 2]), )
	# return np.asarray(mesh.triangles)
	return mesh

def pc_from_mesh(obj_path, npoints):
	mesh = o3d.io.read_triangle_mesh(obj_path)
	pts = mesh.sample_points_uniformly(number_of_points=npoints)
	return np.array(pts.points)

def load_mesh(obj_path):
	return trimesh.load(obj_path, 'obj', force='mesh')

def merge_mesh(meshs):
	merged_mesh = trimesh.util.concatenate(meshs)
	return merged_mesh

def write_mesh(mesh, path, normal=False, color=False):
	o3d.io.write_triangle_mesh(
		path, mesh, write_vertex_normals=normal, write_vertex_colors=color
	)

def gen_meshs(obj_folder, hier_tree, npoints=1024):
	all_node_mesh = {}

	for node in hier_tree:
		id_ = node['id']

		if 'children' in node.keys():
			sub_mesh = gen_meshs(obj_folder, node['children'], npoints)
			all_node_mesh = {**all_node_mesh, **sub_mesh}
			child_mesh = [sub_mesh[me['id']] for me in node['children']]
			node_mesh = merge_mesh(child_mesh)
			all_node_mesh[id_] = node_mesh
		else:
			meshs = []
			for obj_name in node['objs']:
				obj_path = os.path.join(obj_folder, obj_name+'.obj')
				mesh = load_mesh(obj_path)
				meshs.append(mesh)
			if len(meshs)>1:
				meshs = merge_mesh(meshs)
			else:
				meshs = meshs[0]

			all_node_mesh[id_] = meshs

	return all_node_mesh

def get_leaves(tree, only=None, flatten=False, pop_child=True):
	leaf_parts = []

	for node in tree:
		data = node[only] if only is not None else node
		if 'children' not in node.keys():
			leaf_parts.append(data)
		else:
			node_list = get_leaves(node['children'], only, flatten) # [{...},] with parent+children idx
			leaf_parts.extend(node_list)
			if flatten:
				if only == None:
					data = data.copy()
					if pop_child:data.pop('children')
				leaf_parts.append(data)

	return leaf_parts

def hier2graphic(hier_tree, parent_id=-1, depth=0):
	all_nodes = {}
	for node in hier_tree:
		renode = {
			'name': node['name'],
			'objs': node['objs'] if 'objs' in node.keys() else [],
			'parent': parent_id,
			'depth': depth,
			'box': node['box'] if 'box' in node.keys() else [],
			'brother':[],
			'children_id': [],
			'leaves': get_leaves([node], 'id'),
		}

		if 'children' in node.keys():
			children_nodes = hier2graphic(node['children'], node['id'], depth+1)
			all_nodes = {**all_nodes, **children_nodes}
			renode['children_id'] = [i['id'] for i in node['children']]

		all_nodes[node['id']] = renode

		for child in renode['children_id']:
			all_nodes[child]['brother'] = renode['children_id'][:]
			all_nodes[child]['brother'].remove(child)

	return all_nodes

def update_mopara(hash_hier, ids=[0]):
	main_child = ids[:]

	for key in ids:
		if hash_hier[key]['children_id'] != []:
			tree, mochild = update_mopara(hash_hier, hash_hier[key]['children_id'] )
			hash_hier = {**hash_hier, **tree}

			mopara = {'jointData':{}, 'joint':'', 'motype':''}
			node = hash_hier[mochild]

			if 'ref' in node.keys() and key!=0:
				mopara['jointData'] = node['jointData']
				mopara['joint'] = node['joint']
				if 'motype' in node.keys(): mopara['motype'] = node['motype']

				refs = node['ref'][:]
				for idx,ref in enumerate(refs):
					while(hash_hier[ref]['depth'] > hash_hier[key]['depth']):
						ref = hash_hier[ref]['parent']
					refs[idx] = ref

				mopara['ref'] = list(set(refs))
				for ref in mopara['ref']:
					if ref in main_child and ref != key:
						main_child.remove(key)
						break
			hash_hier[key] = {**hash_hier[key], **mopara}

		elif 'ref' in hash_hier[key].keys():
			refs = hash_hier[key]['ref']
			for idx,ref in enumerate(refs):
				while(hash_hier[ref]['depth'] > hash_hier[key]['depth']):
					ref = hash_hier[ref]['parent']
				hash_hier[key]['ref'][idx] = ref

			hash_hier[key]['ref'] = list(set(hash_hier[key]['ref']))
			for ref in hash_hier[key]['ref']:
				if ref in main_child and ref != key:
					main_child.remove(key)
					break

	return hash_hier, main_child[0]

def gen_graph(hier_tree, mobi):
	'''
	将hierarchy tree转化称graph
	'''
	hash_hier = hier2graphic(hier_tree)
	for idx,node in enumerate(mobi):
		# mobi[idx]['ids'] = [i['id'] for i in node['parts']]

		mopara = {'jointData':{}, 'joint':'', 'motype':''}
		if node['jointData'] != {}:
			mopara['jointData'] = node['jointData']
			mopara['joint'] = node['joint']
			if 'motype' in node.keys(): mopara['motype'] = node['motype']

			if node['parent'] != -1 and 'parts' in mobi[node['parent']].keys():
				ref = [j['id'] for j in mobi[node['parent']]['parts']]
				mopara['ref'] = ref

		for sub_node in node['parts']:
			sub_id = sub_node['id']
			hash_hier[sub_id] = {**hash_hier[sub_id], **mopara}

	graph, _ = update_mopara(hash_hier)
	statics = {}

	for key in graph.keys():
		if 'ref' in graph[key].keys():
			refs = graph[key]['ref'][:]
			for ref in refs:
				if graph[key]['parent'] != graph[ref]['parent'] or ref == key:
					graph[key]['ref'].remove(ref)
			if graph[key]['ref'] == []:
				graph[key].pop('ref')

	for key in graph.keys():
		node = graph[key]
		graph[key]['edges']  = {
			'children':{},
			'space':{}
		}
		for child in graph[key]['children_id']:
			graph[key]['edges']['children'][child] = ''
		brothers = graph[key]['brother'][:]

		if 'ref' in graph[key].keys():
			for bro in brothers:
				if bro in graph[key]['ref']:
					graph[key]['edges']['space'][bro] = 'motion'
				else:
					graph[key]['edges']['space'][bro] = 'none'
			graph[key].pop('ref')
		else:
			for bro in brothers:
				graph[key]['edges']['space'][bro] = 'none' if 'ref' in graph[bro].keys() else 'fixed'

	return graph, statics

def ref_count(graph):
	for key in graph.keys():
		edges = graph[key]['edges']['space']
		refs = [r for r in edges.keys() if edges[r]=='motion']
		graph[key]['refs'] = refs

	invalids, child_allref, expect = reduce_ref(graph)
	return invalids, expect

def reduce_ref(graph, node_key='0'):
	ref_child = set()

	all_invalid = 0
	flgs = 0
	for child in graph[node_key]['children_id']:
		child = str(child)
		if graph[child]['refs'] == []:
			ref_child.add(int(child))

		if graph[child]['children_id'] != []:
			invalid, flg, expect = reduce_ref(graph, child)
			all_invalid += invalid if flg else invalid-1
			flgs += 1-flg

	children_allref = False
	if len(ref_child)==0 and graph[node_key]['brother'] == []:
		children_allref = True
	elif len(ref_child) and ref_child == set(graph[node_key]['children_id']) \
		and not flgs:
		children_allref = True

	all_invalid += len(ref_child)
	# print('%s invalids:%d'%(node_key, all_invalid))
	return all_invalid, children_allref, (flgs if flgs else 1)

'''direction, angle, pos, ...'''
def cosine(dirs, vec, abs=True):
	vec = vec.reshape(-1,3)
	vec = vec / np.linalg.norm(vec, axis=-1).reshape(-1,1) # (1-n, 3) -> (1-n,) or val

	mul_res = dirs@vec.T
	cos = mul_res / np.linalg.norm(dirs, axis=-1).reshape(-1,1)

	if abs: cos = np.abs(cos)
	return cos

def cross(dirs, vec, abs=False):
	# vec = vec / np.linalg.norm(vec)
	cro = np.cross(dirs, vec)
	cro = cro / np.linalg.norm(cro, axis=-1)
	if abs:
		cro = np.abs(cro)
	return cro

def motion_pos(direction, gt_pos, pos):
	direction= direction / np.linalg.norm(direction)
	cro = np.cross(pos - gt_pos.reshape(1,3), direction)
	dis = np.abs(np.linalg.norm(cro, axis=-1))
	min_idx = np.argmin(dis)
	return min_idx, dis[min_idx]

def read_json(json_path):
	return json.load(open(json_path, 'r'))

def get_boolseg(seg:np.ndarray, mov_idx):
	mov_idx = np.array(mov_idx).reshape((-1,1)) # (n,1), and seg(1,N)
	return ( seg.reshape((1,-1)) == mov_idx ).sum(axis=0) == 1.

if __name__ == '__main__':
	pass


"""
    bony_downloader.py
    module contains BonyDownloader class to provide provider specific functionality
"""

__author__ = 'Dattatraya Tembare<tembare.datta@gmail.com>'

import datetime
import itertools

import lxml.html
import requests

from common.download_exceptions import DownloadException
from download.file_downloader import FileDownloader

class BonyDownloader(FileDownloader):
    """
    BonyDownloader class has functions for parsing page source code
        parse()  : implementation for 'BONY' provider
    """

    def authenticate(self, provider):
        """
        Step 1:: Authenticate and login to provider's portal
        :param provider: provider
        :return: requests session
        """
        logging.debug('BonyDownloader:authenticate')
        auth_config = self.configs.auth_config[provider]
        access_config = self.configs.access_config[provider]
        session = requests.Session()
        logging.debug(f':::1 Connect to {access_config["login-url"]} and get cookies')
        session.get(access_config['login-url'])
        logging.debug(f':::2 Call {access_config["auth-url"]} page')
        # requests will use the available cookies from session
        try:
            res1 = session.post(access_config["auth-url"], data=auth_config)
            if self._login_failed(provider, res1):
                raise DownloadException('2000_AUTHENTICATION_FAILED',
                                        custom_message=f"Authentication failed for {provider}")
            logging.debug(f'Login status :: {res1.status_code}')
            # BONY request need certificate key for each request
            f_html = self.utils.format_html(res1.text)
            tree = lxml.html.fromstring(f_html)
            csrf_key = tree.xpath('//form[@name="NavForm"]/input[@name="csrfKey"]/@value')[0]
        except Exception as e:
            raise DownloadException('2000_AUTHENTICATION_FAILED', e) from None
        return session, {'for_next_params': True, 'csrfKey': csrf_key}

    def _login_failed(self, provider, response):
        if 'Invalid Login' in response.text:
            return True
        else:
            return False

    def access(self, session, **opts):
        """
        Step 2:: Pull access URL/s from configs file and use it to pull page source which has URLs for file download
            after method execution a_url['deal_info_dict_list'] appended to opts dictionary
        TODO Use namedtuple DealInfo to make current dictionary generic to all providers
        :param session: session with site cookies
        :param opts: user/commandline inputs
        :return: None
        """
        logging.debug('FileDownloader:access')
        provider = opts['provider']
        previous_url_results = list()
        for a_url in opts['access_urls']:
            logging.debug(f':::3 Send request to {a_url} page')
            # Pull input parameters to append as a query string
            user_config = opts['user_input_config'] if 'user_input_config' in opts else None
            user_inputs = user_config['input'] if user_config else self.configs.user_input_config[provider][
                'input']
            deal_info_list = self._prepare_params(a_url, user_inputs)
            # Update URL with values pulled from previous page response
            deal_info_list = self._use_previous_url_result(deal_info_list, previous_url_results)
            # After use clean the previous_url_results
            previous_url_results = []
            for deal_info in deal_info_list:
                params = deal_info['params']
                from_opts = opts['response_dict'] if 'response_dict' in opts else {}
                params = {**params, **from_opts}
                opts['response_dict'] = {}
                try:
                    if a_url['method'] == 'POST':
                        res = session.post(deal_info['link'], data=params)
                    elif a_url['method'] == 'GET':
                        res = session.get(deal_info['link'], params=params)
                except Exception as e:
                    raise DownloadException('3000_ACCESS_FAILED', e)
                logging.debug(f'status code :: {res.status_code} history :: {res.history} response URL :: {res.url}')
                f_html = self.utils.format_html(res.text)
                tree = lxml.html.fromstring(f_html)
                for ele_name, ele_value in a_url['result-dict'].items():
                    if 'for_next_params' in ele_name:
                        _result = self._dict_for_next_url(ele_value, tree)
                        _result['for_next_params'] = True
                        previous_url_results.append(_result)
                        deal_info['for_next_params'] = _result
                        opts['response_dict'] = {'csrfKey': _result['csrfKey']}
                    elif 'for_next_url' in ele_name:
                        _result = self._dict_for_next_url(ele_value, tree)
                        _result['for_next_url'] = True
                        previous_url_results.append(_result)
                    elif 'deal_info' in ele_name:
                        deal_info['deal_info'] = self._dict_for_next_url(ele_value, tree)
                    elif 'for_parsing' in ele_name:
                        f_html_trees = list()
                        for xp in ele_value:
                            f_html_trees.append(tree.xpath(xp))
                        deal_info['f_html'] = f_html_trees
            a_url['deal_info_dict_list'] = deal_info_list

    def _prepare_params(self, a_url, user_inputs):
        # pull mandatory input parameters from access-config
        input_param_dict = a_url['input-param']
        # prepare links for next request/s
        links_with_params = list()
        for attr_name, attr_values in user_inputs.items():
            for attr_value in attr_values:
                req_body = input_param_dict.copy()
                req_body[attr_name] = attr_value
                links_with_params.append({'link': a_url['url'], 'params': req_body})
        return links_with_params

    def _use_previous_url_result(self, links, previous_url_results):
        if len(links) == len(previous_url_results):
            for link, previous_url_result in zip(links, previous_url_results):
                if 'hd_deal_number' in previous_url_result:
                    deal_num = previous_url_result['hd_deal_number']
                    deal_num = deal_num[:deal_num.index('~')] if deal_num else deal_num
                    previous_url_result['hd_deal_number'] = deal_num
                if 'for_next_params' in previous_url_result:
                    link['params'] = {**link['params'], **previous_url_result}
        else:
            for link, previous_url_result in itertools.product(links, previous_url_results):
                if 'for_next_params' in previous_url_result:
                    link['params'] = {**link['params'], **previous_url_result}
        return links

    def _dict_for_next_url(self, input_dict, tree):
        # print(f'table.text :: {etree.tostring(tree)}')
        result_dict = dict()
        for k, xp in input_dict.items():
            try:
                xp_result = tree.xpath(xp)
                result_dict[k] = ''.join(xp_result).strip()
            except Exception as e:
                raise DownloadException('3000_ACCESS_FAILED', e)
        return result_dict

    def parse(self, **opts):
        """
        method parses the 'BONY' specific page source using xpath from access-configs, after method execution
            a_url['download_urls'] appended to opts dictionary
        :param opts: user/commandline inputs + a_url['deal_info_dict_list']
        :return:
        """
        logging.debug('BonyDownloader:parse')
        out_dir = opts['output']
        provider = opts['provider']
        for a_url in opts['access_urls']:
            download_urls = list()
            for deal_info_dict in a_url['deal_info_dict_list']:
                if 'f_html' in deal_info_dict:
                    f_url = a_url['for_download_urls']['download_url']
                    input_dict = a_url['for_download_urls']['request_body'].copy()
                    for k, v in deal_info_dict['for_next_params'].items():
                        if 'for_next_params' not in k:
                            input_dict[k] = v
                    deal_name = deal_info_dict['deal_info']['deal_name']
                    for trs in deal_info_dict['f_html']:
                        for tr in trs:
                            # print(f'table.text :: {etree.tostring(tr)}')
                            report_id = tr.xpath('td/input[@name="cb_rpt_id"]/@value')
                            report_name = ''.join(tr.xpath('td[2]/a/text()')).strip()
                            if len(report_id) > 0:
                                report_id = report_id[0][:report_id[0].index('~')]
                            payment_date = tr.xpath('td[6]/text()')
                            if len(payment_date) > 0:
                                payment_date = payment_date[0].strip()
                                dt = datetime.datetime.strptime(payment_date, "%d-%b-%Y")
                            for span in tr.xpath('td/span[@class="RecordNormalText"]/input'):
                                report_ext_key = span.xpath('@name')[0]
                                report_ext_value = span.xpath('@value')[0]
                                file_extension = report_ext_value[report_ext_value.index('~') + 1:]
                                input_dict_copy = dict(input_dict)
                                input_dict_copy['hd_avl_rpt_id'] = report_id
                                input_dict_copy[report_ext_key] = report_ext_value
                                input_dict_copy['lb_reportdate'] = dt.strftime("%B") + '++' + str(dt.year)
                                input_dict_copy['hd_extension'] = file_extension
                                o_file = out_dir + '/' + str(dt.year) + '-' + str(dt.month) + '/' + provider + '/'
                                o_file += (deal_name + ' pay ' + payment_date + ' ' + report_name).replace(' ', '_')
                                o_file += '.' + file_extension
                                search_data = report_id + ' || ' + report_name + ' || ' + dt.strftime("%b") + ' '
                                search_data += str(dt.year) + ' || ' + deal_name
                                download_urls.append(
                                    DownloadUrl(f_url, o_file, search_data, deal_name, input_dict_copy, 'POST'))
                    # del a_url['f_html']
                    a_url['download_urls'] = download_urls


from time import sleep
import time
import datetime
from datetime import timedelta
from time import sleep, strftime

motionTimeOutSeconds = 5
lastMotionTime = datetime.datetime.now()

def motionTimedOut():
    myNow = datetime.datetime.now()
    deltaTime = (myNow - lastMotionTime).total_seconds()
    if deltaTime > motionTimeOutSeconds:
        print('Motion timed out after {0} seconds'.format(deltaTime))
        return True
    return False

sleep(2)
# bTime = datetime.datetime.now()
# deltaT = (bTime-lastMotionTime).total_seconds()
# print(deltaT)
if motionTimedOut():
    print('Motion timeout test of 2 seconds failed!')
    exit(-1)
print('No timeout after 2 seconds.')
lastMotionTime = datetime.datetime.now()
sleep(6)
if motionTimedOut():
    print('Motion timeout test is working after 6 seconds')


from __future__ import division
import numpy as np
from scipy import sparse
from sklearn.metrics.pairwise import euclidean_distances
import time

# Produce grid points for a 2d grayscale image
def get_points_2d(image, res):
    rows, columns = image.shape
    grid_x, grid_y = np.mgrid[0:columns:res, 0:rows:res]
    grid = np.array((grid_x.flatten(), grid_y.flatten())).T
    return grid

# Produce grid points for a 3d grayscale image
def get_points_3d(image, res):
    rows, columns, z = image.shape
    grid_z, grid_x, grid_y = np.mgrid[0:z:res, 0:columns:res, 0:rows:res]
    grid = np.array((grid_x.flatten(), grid_y.flatten(), grid_z.flatten())).T
    return grid

# Wendland kernel as a function of r = norm(x-y)/c_sup
def dist_kernel(r):
    return max((1-r, 0))**4 * (4*r + 1)

def blowup_S(S, dim):
    (m, n) = S.shape
    if dim == 3:
        S_full = sparse.lil_matrix((3 * m, 3 * n), dtype = np.float32)
        #S_full = np.zeros((3 * m, 3 * n))
        S_full[0::3, 0::3] = S
        S_full[1::3, 1::3] = S
        S_full[2::3, 2::3] = S
    else:
        S_full = np.zeros((2 * m, 2 * n))
        S_full[0::2, 0::2] = S
        S_full[1::2, 1::2] = S
    return S_full.tocsc()

# Create evaluation matrix given kernel centers (grid points), evaluation points
# and kernel support
def evaluation_matrix(kernels, points, c_sup, dim):
    dim = kernels.shape[1]
    vect_kernel = np.vectorize(dist_kernel)
    start = time.time()
    S = euclidean_distances(points, kernels) / c_sup
    #print("VEC -- euc dist ", (time.time() - start) / 60)
    # Mark entries with 0 kernel support
    start = time.time()
    S[np.where(S > 1)] = -1
    non_zero_indices = np.where(S >= 0)
    #print("VEC -- S[np.where(S > 1)] and np.where(S>=0) ", (time.time() - start) / 60)
    # Evaluate kernel at points within support
    start = time.time()
    S[non_zero_indices] = vect_kernel(S[non_zero_indices])
    #print("VEC -- S[non_zero] = vect_kernel ", (time.time() - start) / 60)
    start = time.time()
    S[np.where(S == -1)] = 0
    #print("VEC -- S[np.where(S == -1)] = 0 ", (time.time() - start) / 60)
    start = time.time()
    #full_S = blowup_S_old(S, dim)
    #print("VEC -- blowup ", (time.time() - start) / 60)
    return sparse.csc_matrix(S)

def evaluation_matrix_blowup(kernels, points, c_sup, dim):
    dim = kernels.shape[1]
    vect_kernel = np.vectorize(dist_kernel)
    start = time.time()
    S = euclidean_distances(points, kernels) / c_sup
    #print("VEC -- euc dist ", (time.time() - start) / 60)
    # Mark entries with 0 kernel support
    start = time.time()
    S[np.where(S > 1)] = -1
    non_zero_indices = np.where(S >= 0)
    #print("VEC -- S[np.where(S > 1)] and np.where(S>=0) ", (time.time() - start) / 60)
    # Evaluate kernel at points within support
    start = time.time()
    S[non_zero_indices] = vect_kernel(S[non_zero_indices])
    #print("VEC -- S[non_zero] = vect_kernel ", (time.time() - start) / 60)
    start = time.time()
    S[np.where(S == -1)] = 0
    #print("VEC -- S[np.where(S == -1)] = 0 ", (time.time() - start) / 60)
    start = time.time()
    full_S = blowup_S(S, dim)
    #print("VEC -- blowup ", (time.time() - start) / 60)
    return full_S

# Create velocity field by weighing kernels by alphas
def make_V(S, alpha, d):
    alpha = alpha.flatten()
    if (S.shape[1] == alpha.shape[0]):
        lmda = S.dot(alpha)
        return lmda.reshape(-1, d)
    else:
        alpha = alpha.reshape(-1, d)
        return S.dot(alpha)


import gc
from collections import defaultdict

import cupy as cp
import pandas as pd
import torch
import torch.nn.functional as F
from cuml.metrics import pairwise_distances
from cuml.neighbors import NearestNeighbors
from torch.utils.data import DataLoader, Dataset, default_collate
from tqdm import tqdm
from transformers import AutoTokenizer, TrainerCallback

from utils import clean_text, f2_score, get_pos_score

LANGUAGE_TOKENS = [
    "<|lang_pnb|>",
    "<|lang_tr|>",
    "<|lang_ur|>",
    "<|lang_bn|>",
    "<|lang_hi|>",
    "<|lang_en|>",
    "<|lang_kn|>",
    "<|lang_km|>",
    "<|lang_zh|>",
    "<|lang_gu|>",
    "<|lang_ta|>",
    "<|lang_my|>",
    "<|lang_fr|>",
    "<|lang_swa|>",
    "<|lang_or|>",
    "<|lang_mul|>",
    "<|lang_fil|>",
    "<|lang_sw|>",
    "<|lang_es|>",
    "<|lang_pt|>",
    "<|lang_pl|>",
    "<|lang_ru|>",
    "<|lang_mr|>",
    "<|lang_it|>",
    "<|lang_ar|>",
    "<|lang_bg|>",
    "<|lang_te|>",
    "<|lang_as|>",
]

CATEGORY_TOKENS = [
    "<|category_supplemental|>",
    "<|category_aligned|>",
    "<|category_source|>",
]

LEVEL_TOKENS = [
    "<|level_0|>",
    "<|level_1|>",
    "<|level_2|>",
    "<|level_3|>",
    "<|level_4|>",
    "<|level_5|>",
    "<|level_6|>",
    "<|level_7|>",
    "<|level_8|>",
    "<|level_9|>",
    "<|level_10|>",
]

KIND_TOKENS = [
    "<|kind_document|>",
    "<|kind_video|>",
    "<|kind_html5|>",
    "<|kind_exercise|>",
    "<|kind_audio|>",
]

OTHER_TOKENS = [
    "<|topic|>",
    "<|content|>",
    "<s_title>",
    "</s_title>",
    "<s_description>",
    "</s_description>",
    "<s_text>",
    "</s_text>",
]

RELATION_TOKENS = [
    "<s_parent>",
    "</s_parent>",
    "<s_children>",
    "</s_children>",
]

def init_tokenizer(tokenizer_name):
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
    tokenizer.add_special_tokens(
        dict(
            additional_special_tokens=LANGUAGE_TOKENS
            + CATEGORY_TOKENS
            + LEVEL_TOKENS
            + KIND_TOKENS
            + OTHER_TOKENS
            + RELATION_TOKENS
        )
    )
    if "sentence-t5" in tokenizer_name:
        tokenizer.add_special_tokens({"sep_token": "<sep>"})
    return tokenizer

class LECRDataset(Dataset):
    def __init__(
        self,
        supervised_df,
        topic_df,
        content_df,
        topic_dict,
        content_dict,
        correlation_df,
        tokenizer_name="xlm-roberta-base",
        max_len=512,
        use_content_pair=False,
        is_training=False,
        use_augmentation=False,
        objective="siamese",
    ):
        self.tokenizer = init_tokenizer(tokenizer_name)
        self.max_len = max_len

        self.supervised_df = supervised_df.dropna()
        self.topic_df = topic_df
        self.content_df = content_df
        self.topic_dict, self.content_dict = topic_dict, content_dict
        self.correlation_df = correlation_df
        self.use_content_pair = use_content_pair
        self.is_training = is_training
        self.use_augmentation = use_augmentation
        self.objective = objective
        self.topic_texts, self.content_texts, self.labels = self.process_csv()

    def process_csv(self):
        # get text pairs
        topic_ids = self.supervised_df.topic_id.values
        content_ids = self.supervised_df.content_ids.values
        labels = list(self.supervised_df.target.values)

        topic_texts = []
        content_texts = []

        for topic_id in topic_ids:
            topic_texts.append(self.topic_dict[topic_id])

        for content_id in content_ids:
            content_texts.append(self.content_dict[content_id])

        set_topic_ids = set(topic_ids)
        use_all_pairs = (
            False  # use all pair, no need to be in the intersection of content_ids of topic ids
        )
        if self.use_content_pair:
            # todo: create content pairs from each topic
            content_to_topic = defaultdict(lambda: [])
            topic_to_content = defaultdict(lambda: [])

            pairs = set()

            for i, row in tqdm(self.correlation_df.iterrows()):
                content_list = row["content_ids"].split(" ")

                if row["topic_id"] not in set_topic_ids:
                    continue

                for content_id in content_list:
                    content_to_topic[content_id].append(row["topic_id"])
                    topic_to_content[row["topic_id"]].append(content_id)

                if len(content_list) <= 1:
                    continue

                if use_all_pairs:
                    for idx1 in range(len(content_list) - 1):
                        for idx2 in range(idx1 + 1, len(content_list)):
                            if (content_list[idx1], content_list[idx2],) not in pairs and (
                                content_list[idx2],
                                content_list[idx1],
                            ) not in pairs:
                                pairs.add((content_list[idx1], content_list[idx2]))

            if not use_all_pairs:
                for content_id, topics in tqdm(content_to_topic.items()):
                    intersection_contents = list(
                        set.intersection(*[set(topic_to_content[topic_id]) for topic_id in topics])
                    )
                    for idx1 in range(len(intersection_contents) - 1):
                        for idx2 in range(idx1 + 1, len(intersection_contents)):
                            if (
                                intersection_contents[idx1],
                                intersection_contents[idx2],
                            ) not in pairs and (
                                intersection_contents[idx2],
                                intersection_contents[idx1],
                            ) not in pairs:
                                pairs.add(
                                    (
                                        intersection_contents[idx1],
                                        intersection_contents[idx2],
                                    )
                                )

            for pair in pairs:
                topic_texts.append(self.content_dict[pair[0]])
                content_texts.append(self.content_dict[pair[1]])
                labels.append(1)

        return topic_texts, content_texts, labels

    def __len__(self):
        if self.is_training:
            return len(self.labels)
        else:
            return 1

    def augment(self, inputs):
        probability_matrix = torch.full(inputs["input_ids"].shape, 0.15)
        masked_indices = torch.bernoulli(probability_matrix).bool()
        indices_replaced = (
            torch.bernoulli(torch.full(inputs["input_ids"].shape, 0.8)).bool() & masked_indices
        )
        inputs["input_ids"][indices_replaced] = self.tokenizer.convert_tokens_to_ids(
            self.tokenizer.mask_token
        )
        inputs["input_ids"] *= inputs["attention_mask"]

        return inputs

    def __getitem__(self, idx):
        topic_text = self.topic_texts[idx]
        content_text = self.content_texts[idx]
        label = self.labels[idx]

        if self.objective == "siamese":
            # topic
            if isinstance(topic_text, tuple):
                topic_inputs = self.tokenizer.encode_plus(
                    topic_text[0],
                    topic_text[1],
                    return_tensors=None,
                    add_special_tokens=True,
                    max_length=self.max_len,
                    padding="max_length",
                    truncation=True,
                )
            else:
                topic_inputs = self.tokenizer.encode_plus(
                    topic_text,
                    return_tensors=None,
                    add_special_tokens=True,
                    max_length=self.max_len,
                    padding="max_length",
                    truncation=True,
                )
            for k, v in topic_inputs.items():
                topic_inputs[k] = torch.tensor(v, dtype=torch.long)

            # content
            content_inputs = self.tokenizer.encode_plus(
                content_text,
                return_tensors=None,
                add_special_tokens=True,
                max_length=self.max_len,
                padding="max_length",
                truncation=True,
            )
            for k, v in content_inputs.items():
                content_inputs[k] = torch.tensor(v, dtype=torch.long)

            if isinstance(topic_text, tuple):
                topic_text = topic_text[0] + topic_text[1]

            if self.is_training and self.use_augmentation:
                topic_inputs = self.augment(topic_inputs)
                content_inputs = self.augment(content_inputs)

            return topic_inputs, content_inputs, topic_inputs, label
        elif self.objective == "classification":
            combined_inputs = self.tokenizer.encode_plus(
                topic_text,
                content_text,
                return_tensors=None,
                add_special_tokens=True,
                max_length=self.max_len,
                padding="max_length",
                truncation=True,
            )
            for k, v in combined_inputs.items():
                combined_inputs[k] = torch.tensor(v, dtype=torch.long)

            if self.is_training and self.use_augmentation:
                combined_inputs = self.augment(combined_inputs)

            return combined_inputs, combined_inputs, combined_inputs, label
        else:
            raise ValueError("Only support siamese/classification for now.")

class InferenceDataset(Dataset):
    def __init__(self, texts, tokenizer_name="xlm-roberta-base", max_len=512):
        self.texts = texts

        self.tokenizer = init_tokenizer(tokenizer_name)
        self.max_len = max_len

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        text = self.texts[idx]

        # topic
        inputs = self.tokenizer.encode_plus(
            text,
            return_tensors=None,
            add_special_tokens=True,
            max_length=self.max_len,
            padding="max_length",
            truncation=True,
        )
        for k, v in inputs.items():
            inputs[k] = torch.tensor(v, dtype=torch.long)

        return inputs

def collate_fn(inputs):
    inputs = default_collate(inputs)
    mask_len = int(inputs["attention_mask"].sum(axis=1).max())
    for k, v in inputs.items():
        inputs[k] = inputs[k][:, :mask_len]

    return inputs

class DatasetUpdateCallback(TrainerCallback):
    """
    Trigger re-computing dataset

    A hack that modifies the train/val dataset, pointed by Trainer's dataloader

    0. Calculate new train/val topic/content embeddings, train KNN, get new top-k
    1. Calculate top-k max positive score, compare to current val best, if greater, continue to step 2, else do nothing
    2. Update supervised_df and update dataset:
        self.topic_texts, self.content_texts, self.labels = self.process_csv()
    """

    def __init__(
        self,
        trainer,
        train_topic_ids,
        val_topic_ids,
        topic_df,
        content_df,
        topic_dict,
        content_dict,
        correlation_df,
        tokenizer_name,
        max_len,
        best_score=0,
        top_k=50,
        use_translated=False,
        mix_translated=False,
        fold=0,
    ):
        super(DatasetUpdateCallback, self).__init__()
        self.trainer = trainer
        self.topic_df = topic_df
        self.content_df = content_df
        self.correlation_df = correlation_df
        self.best_score = best_score
        self.top_k = top_k
        self.use_translated = use_translated
        self.mix_translated = mix_translated
        self.fold = fold

        self.tokenizer = init_tokenizer(tokenizer_name)
        self.topic_dict, self.content_dict = topic_dict, content_dict

        train_topic_texts = [
            topic_dict[topic_id]
            for topic_id in self.topic_df.id.values
            if topic_id in train_topic_ids
        ]
        self.train_topic_ids = [
            topic_id for topic_id in self.topic_df.id.values if topic_id in train_topic_ids
        ]
        self.train_topic_languages = []
        for topic_id, topic_lang in zip(self.topic_df.id.values, self.topic_df.language.values):
            if topic_id in train_topic_ids:
                self.train_topic_languages.append(topic_lang)

        val_topic_texts = [
            topic_dict[topic_id]
            for topic_id in self.topic_df.id.values
            if topic_id in val_topic_ids
        ]
        self.val_topic_ids = [
            topic_id for topic_id in self.topic_df.id.values if topic_id in val_topic_ids
        ]

        content_texts = [
            content_dict[content_id]
            for content_id in self.content_df.id.values
            if content_id.startswith("c_")
        ]

        def inference_collate_fn(inputs):
            inputs = default_collate(inputs)
            mask_len = int(inputs["attention_mask"].sum(axis=1).max())
            for k, v in inputs.items():
                inputs[k] = inputs[k][:, :mask_len]

            return inputs

        train_topic_dataset = InferenceDataset(
            texts=train_topic_texts, tokenizer_name=tokenizer_name, max_len=max_len
        )
        self.train_topic_dataloader = DataLoader(
            train_topic_dataset,
            num_workers=self.trainer.args.dataloader_num_workers,
            batch_size=32,
            shuffle=False,
            collate_fn=inference_collate_fn,
        )

        val_topic_dataset = InferenceDataset(
            texts=val_topic_texts, tokenizer_name=tokenizer_name, max_len=max_len
        )
        self.val_topic_dataloader = DataLoader(
            val_topic_dataset,
            num_workers=self.trainer.args.dataloader_num_workers,
            batch_size=32,
            shuffle=False,
            collate_fn=inference_collate_fn,
        )

        content_dataset = InferenceDataset(
            texts=content_texts, tokenizer_name=tokenizer_name, max_len=max_len
        )
        self.content_dataloader = DataLoader(
            content_dataset,
            num_workers=self.trainer.args.dataloader_num_workers,
            batch_size=32,
            shuffle=False,
            collate_fn=inference_collate_fn,
        )

    def on_train_begin(self, args, state, control, **kwargs):
        self.on_epoch_end(args, state, control, **kwargs)

    def on_epoch_end(self, args, state, control, **kwargs):
        local_rank = args.local_rank if args.local_rank != -1 else 0
        with cp.cuda.Device(local_rank):
            torch.cuda.empty_cache()
            print("Callback on local_rank =", local_rank)
            self.trainer.model.eval()
            print("On Epoch Begin")
            topic_embs = []
            device = f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu"

            with torch.no_grad():
                for inputs in tqdm(self.val_topic_dataloader):
                    for k, v in inputs.items():
                        inputs[k] = inputs[k].to(device)
                    out = self.trainer.model.feature(inputs)
                    topic_embs.extend(out.cpu().detach().numpy())

                content_embs = []
                # TODO: only use original content embeddings to avoid translation confusing
                for inputs in tqdm(self.content_dataloader):
                    for k, v in inputs.items():
                        inputs[k] = inputs[k].to(device)
                    out = self.trainer.model.feature(inputs)
                    content_embs.extend(out.cpu().detach().numpy())

            # Transfer predictions to gpu
            with cp.cuda.Device(local_rank):
                topic_embs_gpu = cp.array(topic_embs)
                content_embs_gpu = cp.array(content_embs)

            # Release memory
            torch.cuda.empty_cache()

            # KNN model
            content_idx_to_id = {}
            for idx, row in self.content_df.iterrows():
                content_idx_to_id[idx] = row.id

            print("Evaluating current score...")
            if self.use_translated:
                # get 500 nearest contents, then select top k contents that is in original contents, just approximate, can't check all
                original_indices = [  # indices of original contents in self.content_df
                    i
                    for i, emb in enumerate(content_embs)
                    if self.content_df.id.values[i].startswith("c_")
                ]
                # original_content_embs = [
                #     emb
                #     for i, emb in enumerate(content_embs)
                #     if self.content_df.id.values[i].startswith("c_")
                # ]
                # original_content_embs_gpu = cp.array(original_content_embs)
                original_content_embs_gpu = content_embs_gpu

                neighbors_model = NearestNeighbors(n_neighbors=500, metric="cosine")
                neighbors_model.fit(original_content_embs_gpu)

                indices = neighbors_model.kneighbors(topic_embs_gpu, return_distance=False)
                for selected_k in [5, 10, 20, 50, 100, 200]:
                    predictions = []
                    for k in tqdm(range(len(indices))):
                        pred = indices[k]
                        # original_contents = [self.content_df.loc[ind, "id"] for ind in pred.get() if self.content_df.loc[ind, "id"].startswith("c_")]
                        # original_contents = [content_idx_to_id[ind] for ind in pred.get() if content_idx_to_id[ind].startswith("c_")]
                        original_contents = [
                            content_idx_to_id[original_indices[ind]] for ind in pred.get()
                        ]
                        p = " ".join(original_contents[:selected_k])
                        predictions.append(p)

                    knn_preds = pd.DataFrame(
                        {"topic_id": self.val_topic_ids, "content_ids": predictions}
                    ).sort_values("topic_id")

                    gt = self.correlation_df[
                        self.correlation_df.topic_id.isin(self.val_topic_ids)
                    ].sort_values("topic_id")
                    score = get_pos_score(
                        gt["content_ids"],
                        knn_preds.sort_values("topic_id")["content_ids"],
                        selected_k,
                    )
                    print(
                        "Selecting",
                        selected_k,
                        "nearest contents",
                        "top-k score =",
                        f2_score(
                            gt["content_ids"],
                            knn_preds.sort_values("topic_id")["content_ids"],
                        ),
                        "max positive score =",
                        score,
                    )

                print("Training KNN model...")
                print("Generating KNN predictions with top_k =", self.top_k)
                neighbors_model = NearestNeighbors(n_neighbors=self.top_k, metric="cosine")
                neighbors_model.fit(original_content_embs_gpu)

                print("Generating embedding for validation topics")
                indices = neighbors_model.kneighbors(topic_embs_gpu, return_distance=False)
                predictions = []
                for k in tqdm(range(len(indices))):
                    pred = indices[k]
                    # original_contents = [self.content_df.loc[ind, "id"] for ind in pred.get() if self.content_df.loc[ind, "id"].startswith("c_")]
                    # original_contents = [content_idx_to_id[ind] for ind in pred.get() if content_idx_to_id[ind].startswith("c_")]
                    original_contents = [
                        content_idx_to_id[original_indices[ind]] for ind in pred.get()
                    ]
                    p = " ".join(original_contents[: self.top_k])
                    predictions.append(p)
            else:
                for selected_k in [5, 10, 20, 50, 100, 200]:
                    neighbors_model = NearestNeighbors(n_neighbors=selected_k, metric="cosine")
                    neighbors_model.fit(content_embs_gpu)

                    indices = neighbors_model.kneighbors(topic_embs_gpu, return_distance=False)
                    predictions = []
                    for k in tqdm(range(len(indices))):
                        pred = indices[k]
                        # p = " ".join([self.content_df.loc[ind, "id"] for ind in pred.get()])
                        p = " ".join([content_idx_to_id[ind] for ind in pred.get()])
                        predictions.append(p)

                    knn_preds = pd.DataFrame(
                        {"topic_id": self.val_topic_ids, "content_ids": predictions}
                    ).sort_values("topic_id")

                    gt = self.correlation_df[
                        self.correlation_df.topic_id.isin(self.val_topic_ids)
                    ].sort_values("topic_id")
                    score = get_pos_score(
                        gt["content_ids"],
                        knn_preds.sort_values("topic_id")["content_ids"],
                        selected_k,
                    )
                    print(
                        "Selecting",
                        selected_k,
                        "nearest contents",
                        "top-k score =",
                        f2_score(
                            gt["content_ids"],
                            knn_preds.sort_values("topic_id")["content_ids"],
                        ),
                        "max positive score =",
                        score,
                    )

                print("Training KNN model...")
                print("Generating KNN predictions with top_k =", self.top_k)
                neighbors_model = NearestNeighbors(n_neighbors=self.top_k, metric="cosine")
                neighbors_model.fit(content_embs_gpu)

                print("Generating embedding for validation topics")
                indices = neighbors_model.kneighbors(topic_embs_gpu, return_distance=False)
                predictions = []
                for k in tqdm(range(len(indices))):
                    pred = indices[k]
                    # p = " ".join([self.content_df.loc[ind, "id"] for ind in pred.get()])
                    p = " ".join([content_idx_to_id[ind] for ind in pred.get()])
                    predictions.append(p)

            knn_preds = pd.DataFrame(
                {"topic_id": self.val_topic_ids, "content_ids": predictions}
            ).sort_values("topic_id")

            score = get_pos_score(
                gt["content_ids"],
                knn_preds.sort_values("topic_id")["content_ids"],
                self.top_k,
            )
            print("Current Score:", score, "Best Score:", self.best_score)

            if score > self.best_score:
                self.best_score = score
                print("saving best model to data/ folder")
                # torch.save(self.trainer.model.state_dict(), f"data/siamese_model_{score}.pth")

            generate_new_dataset_every_epoch = True
            if generate_new_dataset_every_epoch or (score == self.best_score):
                # generate new pairs in dataset
                print("Building new validation supervised df")
                new_val_supervised_df = build_new_supervised_df(knn_preds, self.correlation_df)[
                    ["topic_id", "content_ids", "target"]
                ].sort_values(["topic_id", "content_ids"])
                if score == self.best_score:  # only save for the best checkpoint
                    print("saving new_val_supervised_df to data/ folder")
                    new_val_supervised_df.to_csv("data/new_val_supervised_df.csv")

                # get top-k for training set
                # TODO: only get original content neighbors for original topics
                print("Generating embedding for train topics")
                train_topic_embs = []

                with torch.no_grad():
                    for inputs in tqdm(self.train_topic_dataloader):
                        for k, v in inputs.items():
                            inputs[k] = inputs[k].to(device)
                        out = self.trainer.model.feature(inputs)
                        train_topic_embs.extend(out.cpu().detach().numpy())

                with cp.cuda.Device(local_rank):
                    train_topic_embs_gpu = cp.array(train_topic_embs)

                train_indices = neighbors_model.kneighbors(
                    train_topic_embs_gpu, return_distance=False
                )

                # if self.use_translated:
                #     topic_language_df = pd.DataFrame({
                #         "topic_id": self.train_topic_ids,
                #         "language": self.train_topic_languages
                #     })

                train_predictions = []
                for k in tqdm(range(len(train_indices))):
                    pred = train_indices[k]
                    # p = " ".join([self.content_df.loc[ind, "id"] for ind in pred.get()])
                    if self.use_translated:
                        p = " ".join(
                            [content_idx_to_id[original_indices[ind]] for ind in pred.get()]
                        )
                    else:
                        p = " ".join([content_idx_to_id[ind] for ind in pred.get()])

                    train_predictions.append(p)

                train_knn_preds = pd.DataFrame(
                    {
                        "topic_id": self.train_topic_ids,
                        "content_ids": train_predictions,
                        "language": self.train_topic_languages,
                    }
                ).sort_values("topic_id")

                print("Building new train supervised df")
                # if self.use_translated:
                #     count_dict = {
                #         "ar": 3701,
                #         "as": 167,
                #         "bg": 2867,
                #         "bn": 2176,
                #         "en": 36161,
                #         "es": 13910,
                #         "fil": 247,
                #         "fr": 3701,
                #         "gu": 2320,
                #         "hi": 1786,
                #         "it": 866,
                #         "km": 121,
                #         "kn": 119,
                #         "mr": 300,
                #         "mul": 4,
                #         "my": 135,
                #         "or": 70,
                #         "pl": 43,
                #         "pnb": 51,
                #         "pt": 4177,
                #         "ru": 34,
                #         "sw": 2860,
                #         "swa": 35,
                #         "ta": 60,
                #         "te": 93,
                #         "tr": 40,
                #         "ur": 66,
                #         "zh": 862,
                #     }

                #     times_positive_samples = 4

                #     # select all original topics and a part of translated topics
                #     translated_knn_preds = (
                #         train_knn_preds[~train_knn_preds.topic_id.str.startswith("t_")]
                #         .groupby("language")
                #         .apply(
                #             lambda x: x.sample(
                #                 n=count_dict[x["language"].iat[0]] * times_positive_samples,
                #                 replace=True,
                #             )
                #         )
                #         .reset_index(drop=True)
                #     )
                #     original_knn_preds = train_knn_preds[
                #         train_knn_preds.topic_id.str.startswith("t_")
                #     ]

                #     train_knn_preds = pd.concat([original_knn_preds, translated_knn_preds])

                new_train_supervised_df = build_new_supervised_df(
                    train_knn_preds, self.correlation_df
                )

                if self.use_translated:
                    # Only add positive cases in training set for translated topics
                    translated_supervised_df = new_train_supervised_df[
                        ~new_train_supervised_df.topic_id.str.startswith("t_")
                        & new_train_supervised_df.target
                        == 1
                    ].copy()

                    # Only original contents for original topics
                    original_supervised_df = new_train_supervised_df[
                        new_train_supervised_df.topic_id.str.startswith("t_")
                        & new_train_supervised_df.content_ids.str.startswith("c_")
                    ].copy()

                    # TODO: duplicate number of positive by using translated data
                    id_to_language = {}
                    for _, row in tqdm(self.topic_df.iterrows()):
                        id_to_language[row.id] = row.language

                    original_supervised_df["language"] = original_supervised_df["topic_id"].apply(
                        lambda x: id_to_language[x]
                    )
                    count_df = (
                        original_supervised_df[original_supervised_df.target == 1]
                        .groupby("language")
                        .size()
                        .reset_index(name="counts")
                    )

                    count_dict = {}
                    for _, row in count_df.iterrows():
                        count_dict[row.language] = row.counts

                    times_positive_samples = 3
                    translated_supervised_df["language"] = translated_supervised_df[
                        "topic_id"
                    ].apply(lambda x: id_to_language[x])
                    translated_supervised_df = (
                        translated_supervised_df.groupby("language")
                        .apply(
                            lambda x: x.sample(
                                n=count_dict[x["language"].iat[0]] * times_positive_samples,
                                replace=True,
                            )
                        )
                        .reset_index(drop=True)
                    )
                    original_supervised_df = original_supervised_df.drop(columns=["language"])
                    translated_supervised_df = translated_supervised_df.drop(columns=["language"])

                    new_train_supervised_df = pd.concat(
                        [translated_supervised_df, original_supervised_df]
                    )[["topic_id", "content_ids", "target"]].sort_values(
                        ["topic_id", "content_ids"]
                    )

                if score == self.best_score:  # only save for the best checkpoint
                    print("saving new_train_supervised_df to data/ folder")
                    new_train_supervised_df.to_csv("data/new_train_supervised_df.csv")

                # update train_dataset and val_dataset
                print("preprocess csv for train/validation topics, contents, labels")
                self.trainer.train_dataset.supervised_df = new_train_supervised_df.dropna()
                (
                    self.trainer.train_dataset.topic_texts,
                    self.trainer.train_dataset.content_texts,
                    self.trainer.train_dataset.labels,
                ) = self.trainer.train_dataset.process_csv()

                self.trainer.eval_dataset.supervised_df = new_val_supervised_df.dropna()
                (
                    self.trainer.eval_dataset.topic_texts,
                    self.trainer.eval_dataset.content_texts,
                    self.trainer.eval_dataset.labels,
                ) = self.trainer.eval_dataset.process_csv()
                print("Saving knn csvs ...")
                train_knn_preds.to_csv(f"data/train_knn_fold{self.fold}.csv")
                knn_preds.to_csv(f"data/val_knn_fold{self.fold}.csv")

                del (
                    train_topic_embs,
                    train_topic_embs_gpu,
                    train_knn_preds,
                    train_indices,
                    train_predictions,
                )
                gc.collect()

            del (
                topic_embs,
                content_embs,
                topic_embs_gpu,
                content_embs_gpu,
                knn_preds,
                indices,
                neighbors_model,
                predictions,
            )
            gc.collect()
            torch.cuda.empty_cache()
            if self.mix_translated:
                self.use_translated = not self.use_translated

def build_new_supervised_df(knn_df, correlations):
    # Create lists for training
    topics_ids = []
    content_ids = []
    targets = []

    # Iterate over each topic in df
    mapping = set()

    # get all class 1 in correlations
    topic_ids = set(knn_df.topic_id.values)
    filtered_correlations = correlations[correlations["topic_id"].isin(topic_ids)]
    for i, row in tqdm(filtered_correlations.iterrows()):
        if str(row["content_ids"]) and str(row["content_ids"]) != "nan":
            content_ids = str(row["content_ids"]).split(" ")
            for content_id in content_ids:
                mapping.add((row["topic_id"], content_id, 1))

    for i, row in tqdm(knn_df.iterrows()):
        if str(row["content_ids"]) and str(row["content_ids"]) != "nan":
            content_ids = str(row["content_ids"]).split(" ")
            for content_id in content_ids:
                if (
                    row["topic_id"],
                    content_id,
                    1,
                ) not in mapping:  # because mapping already contains all positive cases
                    mapping.add((row["topic_id"], content_id, 0))

    # Build training dataset
    mapping = list(mapping)
    new_df = pd.DataFrame(
        {
            "topic_id": [item[0] for item in mapping if item[1]],
            "content_ids": [item[1] for item in mapping if item[1]],
            "target": [item[2] for item in mapping if item[1]],
        }
    )
    # Release memory
    del topics_ids, content_ids
    gc.collect()
    return new_df

def collate_fn(batch):
    batch = default_collate(batch)

    topic_inputs, content_inputs, combined_inputs, labels = batch
    mask_len = int(topic_inputs["attention_mask"].sum(axis=1).max())
    for k, v in topic_inputs.items():
        topic_inputs[k] = topic_inputs[k][:, :mask_len]

    mask_len = int(content_inputs["attention_mask"].sum(axis=1).max())
    for k, v in content_inputs.items():
        content_inputs[k] = content_inputs[k][:, :mask_len]

    mask_len = int(combined_inputs["attention_mask"].sum(axis=1).max())
    for k, v in combined_inputs.items():
        combined_inputs[k] = combined_inputs[k][:, :mask_len]

    return {
        "topic_inputs": topic_inputs,
        "content_inputs": content_inputs,
        "combined_inputs": combined_inputs,
        "labels": labels,
    }


from aiogram import Bot, types, Dispatcher, executor
import logging
from config import TOKEN, html
import parser as ps
import time
import random
import os
import qrcode

def make_qr(text):
    qr = qrcode.QRCode()
    qr.add_data(text)
    img_qr = qr.make_image(fill_color='white', back_color="black")
    img_qr.save('qr.png')

bot = Bot(token=TOKEN)
dp = Dispatcher(bot)
logging.basicConfig(level=logging.INFO)

async def on_startup(_):
    print('Бот онлайн')

@dp.message_handler(commands='numhent')
async def numhent(msg : types.Message):
    number = msg.text.split(' ', 1)
    try:
        ps.get_html(html,number[1])
        photo = ps.parse('content', number[1])
        await msg.reply_photo(photo,caption=number[1])
    except:
        await msg.reply('отправь число дурак')

@dp.message_handler(commands='hent')
async def hent(msg : types.Message):
    rnd = random.randint(1,6330000)
    ps.get_html(html,rnd)
    t = ps.parse('content', rnd)
    await msg.reply_photo(t,caption=rnd)

@dp.message_handler(commands='qr')
async def test(msg : types.Message):
    split = msg.text.split(' ', 1)[1]
    make_qr(split)
    await msg.reply_photo(open('qr.png', 'rb'), caption=split)

if __name__ == '__main__':
    executor.start_polling(dp,skip_updates=True, on_startup=on_startup)


from functools import reduce
from typing import List

from project.caretaker import Caretaker
from project.cheetah import Cheetah
from project.keeper import Keeper
from project.lion import Lion
from project.tiger import Tiger
from project.vet import Vet

from project.animal import Animal

from project.worker import Worker

class Zoo:
    def __init__(self,
                 name: str,
                 budget: int,
                 animal_capacity: int,
                 workers_capacity: int
                 ):
        # public instance attribute
        self.name = name
        # private attributes
        self.__budget = budget
        self.__animal_capacity = animal_capacity
        self.__workers_capacity = workers_capacity
        # public instance attributes
        self.animals: List[Animal] = []
        self.workers: List[Worker] = []

    def add_animal(self, animal: Animal, price: int) -> str:
        if (price <= self.__budget) and (len(self.animals) < self.__animal_capacity):
            self.animals.append(animal)
            self.__budget -= price
            return f'{animal.name} the {animal.__class__.__name__} added to the zoo'
            # or type(animal).__name__

        if (price > self.__budget) and (len(self.animals) < self.__animal_capacity):
            return 'Not enough budget'

        return 'Not enough space for animal'

    def hire_worker(self, worker):
        if len(self.workers) < self.__workers_capacity:
            self.workers.append(worker)
            return f'{worker.name} the {worker.__class__.__name__} hired successfully'
            # or {type(worker).__name__}
        return 'Not enough space for worker'

    def fire_worker(self, worker_name):
        worker = [w for w in self.workers if w.name == worker_name]
        if worker:
            self.workers.remove(worker[0])
            return f'{worker[0].name} fired successfully'
        return f'There is no {worker_name} in the zoo'

    def pay_workers(self):
        # !!!!!
        workers_payment = sum([w.salary for w in self.workers])

        if workers_payment <= self.__budget:
            self.__budget -= workers_payment
            return f'You payed your workers. They are happy. ' \
                   f'Budget left: {self.__budget}'

        return 'You have no budget to pay your workers. They are unhappy'

    def tend_animals(self):
        # get_needs = self.money_for_care
        amount_to_pay = sum([t.get_needs() for t in self.animals])
        if self.__budget >= amount_to_pay:
            self.__budget -= amount_to_pay
            return f"You tended all the animals. They are happy. Budget left: {self.__budget}"
        return "You have no budget to tend the animals. They are unhappy."

    def profit(self, amount) -> None:
        self.__budget += amount

    def animals_status(self):
        animals_types = ['Lion', 'Tiger', 'Cheetah']
        animals_list = {idx: [] for idx in range(0, 3)}

        for animal in self.animals:
            idx = animals_types.index(type(animal).__name__)
            animals_list[idx].append(animal)

        lions, tigers, cheetahs = animals_list[0], animals_list[1], animals_list[2]
        #
        # lions = [animal for animal in self.animals if type(animal).__name__ == animals_types[0]]
        # tigers = [animal for animal in self.animals if type(animal).__name__ == animals_types[1]]
        # cheetahs = [animal for animal in self.animals if type(animal).__name__ == animals_types[2]]

        result = [f'You have {len(self.animals)} animals']
        result.append(f'----- {len(lions)} Lions:')
        result.append('\n'.join([animal.__repr__() for animal in lions]))
        result.append(f'----- {len(tigers)} Tigers:')
        result.append('\n'.join([animal.__repr__() for animal in tigers]))
        result.append(f'----- {len(cheetahs)} Cheetahs:')
        result.append('\n'.join([animal.__repr__() for animal in cheetahs]))
        return '\n'.join(result)

    def workers_status(self):
        keepers = [w for w in self.workers if w.__class__.__name__ == 'Keeper']
        caretakers = [w for w in self.workers if w.__class__.__name__ == 'Caretaker']
        vets = [w for w in self.workers if w.__class__.__name__ == 'Vet']

        result = f"You have {len(self.workers)} workers\n"
        result += f'----- {len(keepers)} Keepers:\n'
        result += '\n'.join([k.__repr__() for k in keepers]) + '\n'
        result += f'----- {len(caretakers)} Caretakers:\n'
        result += '\n'.join([c.__repr__() for c in caretakers]) + '\n'
        result += f'----- {len(vets)} Vets:\n'
        result += '\n'.join([v.__repr__() for v in vets])
        return result


from celery import shared_task, Celery
from django.utils import timezone

from .models import Post

app = Celery()

@shared_task
def publish_posts_task():
    posts = Post.objects.filter(
        status=False, published_date__lte=timezone.now()
    )
    for post in posts:
        post.status = True
        post.save()
    return (
        print(f"{posts.count()} published!")
        if posts
        else print("There is no post to publish")
    )

@app.on_after_finalize.connect
def setup_periodic_tasks(sender, **kwargs):
    sender.add_periodic_task(
        60 * 60,
        publish_posts_task().s(),
        name="published posts every one hour",
    )


#!/usr/bin/python3
# -*-coding:utf-8 -*-
# Reference:**********************************************
# @Time     : 2019/11/1 23:30
# @Author   : Raymond Luo
# @File     : train_emb.py
# @User     : luoli
# @Software: PyCharm
# Reference:**********************************************
import pickle
from gensim.models import Word2Vec, KeyedVectors
import pandas as pd
import torch.nn as nn
import torch

def train_motif_wordemb(path):
    data = pd.read_csv(path)
    walk_a = data['user_neighbor'].values.tolist()
    walk_b = data['target_neighbor'].values.tolist()
    walk_a.extend(walk_b)
    walk = []
    for line in walk_a:
        new_line = line[1:-1].split(", ")
        walk.append(new_line)
    model = Word2Vec(walk, size=128, window=3, min_count=0, sg=1, workers=12, iter=2, compute_loss=True)
    print("Node2vec loss:", model.get_latest_training_loss())
    model.wv.save_word2vec_format("../model/motif_walk.emb")

def change_emb_index(emb_path, uid2idx_path):
    with open(uid2idx_path, "rb") as f:
        uid2idx = pickle.load(f)
    with open(emb_path, "r") as f:
        emb_file = f.readlines()
    head = 1
    new_file = []
    for line in emb_file:
        if head:
            head = 0
            new_file.append(line)
            continue  # 跳过第一行
        line_list = line.split(" ")
        idx = uid2idx[int(line_list[0])]  # uid 2 idx
        line_list[0] = str(idx)  # 转回去
        new_line = " ".join(line_list)
        new_file.append(new_line)
    with open("../model/motif_walk_idx.emb", "w", encoding="utf-8") as f:
        for line in new_file:
            f.write(line)

if __name__ == "__main__":
    # train_motif_wordemb("../data/train_data.csv")
    # change_emb_index("../model/motif_walk.emb", "../data/uid_2_idx.pkl")
    # test
    # 构建词向量
    word_vectors = KeyedVectors.load_word2vec_format("../model/motif_walk_idx.emb", binary=False)  # 节点向量
    weight = torch.FloatTensor(word_vectors.syn0)  # 获取2D numpy矩阵
    emb = nn.Embedding.from_pretrained(weight, freeze=False)
    print(emb(torch.LongTensor([47066])))


#!/usr/bin/env python
# coding=utf-8

import datetime
import hashlib
import json

class LastUpdated():
    def __init__(self, file='last-updated.json'):
        self.file = file

    def read(self):
        with open(self.file, 'r') as f:
            data = json.load(f)

        return {
            'amiibo_sha1': data['amiibo_sha1'],
            'game_info_sha1': data['game_info_sha1'],
            'timestamp': datetime.datetime.strptime(data['timestamp'], '%Y-%m-%dT%H:%M:%S.%f'),
        }

    def read_timestamp(self):
        return self.read()['timestamp']

    def write(self, amiibo_sha1, game_info_sha1, timestamp):
        with open(self.file, 'w') as f:
            json.dump({
                'amiibo_sha1': amiibo_sha1,
                'game_info_sha1': game_info_sha1,
                'timestamp': timestamp.isoformat(),
            }, f, sort_keys=True)

    def hash(self, data):
        return hashlib.sha1(data).hexdigest()

    def update(self, data, data1):
        amiibo_sha1 = self.hash(data)
        game_info_sha1 = self.hash(data1)
        try:
            last_update = self.read()
        except Exception as e:
            print(e)
            last_update = None

        updated = False
        if last_update is None or last_update['amiibo_sha1'] != amiibo_sha1 or last_update['game_info_sha1'] != game_info_sha1:
            last_update = {
                'amiibo_sha1': amiibo_sha1,
                'game_info_sha1': game_info_sha1,
                'timestamp': datetime.datetime.utcnow(),
            }
            self.write(**last_update)
            updated = True

        return last_update, updated

if __name__ == '__main__':
    last_updater = LastUpdated()
    with open('database/amiibo.json', 'rb') as f:
        with open('database/games_info.json', 'rb') as g:
            last_update, updated = last_updater.update(f.read(), g.read())

    if updated:
        print('Updated: {}'.format(last_updater.file))

    print('amiibo_sha1: {}'.format(last_update['amiibo_sha1']))
    print('game_info_sha1: {}'.format(last_update['game_info_sha1']))
    print('timestamp: {}'.format(last_update['timestamp'].isoformat()))


import pytest
import requests

from budget.enums import ExpensesCategoryEnum, IncomeCategoryEnum
from common.tests_fixtures.fixtures import admin_credentials, admin_id, base_url

budgets_url = f"{base_url}/budgets/"
incomes_url = f"{base_url}/incomes/"
expenses_url = f"{base_url}/expenses/"

@pytest.fixture
def create_budget():
    budget_data = {
        "owner": admin_id,
        "name": "New budget name",
    }
    response = requests.post(budgets_url, json=budget_data, **admin_credentials)
    assert response.status_code == 201
    return response.json()

def test_creating_budget():
    budget_data = {
        "owner": admin_id,
        "name": "New budget name",
    }
    response = requests.post(budgets_url, json=budget_data, **admin_credentials)
    assert response.status_code == 201
    created_budget_url = response.json()["url"]

    response = requests.get(created_budget_url, **admin_credentials)
    assert response.status_code == 200
    response = response.json()

    assert response["owner"] == budget_data["owner"]
    assert response["name"] == budget_data["name"]

def test_add_income(create_budget):
    created_budget_url = create_budget["url"]
    budget_id = int(created_budget_url.split("/")[-2])
    income_data = {"category": IncomeCategoryEnum.EARNED_INCOME, "amount": 1000.00, "budget": budget_id}

    response = requests.post(incomes_url, json=income_data, **admin_credentials)
    assert response.status_code == 201
    response = response.json()
    assert income_data["category"] == response["category"]
    assert float(income_data["amount"]) == float(response["amount"])
    assert income_data["budget"] == response["budget"]

def test_add_expense(create_budget):
    created_budget_url = create_budget["url"]
    budget_id = int(created_budget_url.split("/")[-2])
    expense_data = {"category": ExpensesCategoryEnum.SAVING, "amount": 950.21, "budget": budget_id}

    response = requests.post(expenses_url, json=expense_data, **admin_credentials)
    assert response.status_code == 201
    response = response.json()
    assert expense_data["category"] == response["category"]
    assert float(expense_data["amount"]) == float(response["amount"])
    assert expense_data["budget"] == response["budget"]

def test_add_expense_with_incorrect_category(create_budget):
    created_budget_url = create_budget["url"]
    budget_id = int(created_budget_url.split("/")[-2])
    expense_data = {"category": "incorrect_category", "amount": 950.21, "budget": budget_id}

    response = requests.post(expenses_url, json=expense_data, **admin_credentials)
    assert response.status_code == 400
    assert response.json() == {"category": ['"incorrect_category" is not a valid choice.']}

def test_filtering_expense(create_budget):
    created_budget_url = create_budget["url"]
    budget_id = int(created_budget_url.split("/")[-2])
    expense_data_1 = {"category": ExpensesCategoryEnum.SAVING, "amount": 950.21, "budget": budget_id}
    expense_data_2 = {"category": ExpensesCategoryEnum.PERSONAL, "amount": 950.21, "budget": budget_id}

    response_1 = requests.post(expenses_url, json=expense_data_1, **admin_credentials)
    assert response_1.status_code == 201
    response_1 = response_1.json()

    response_2 = requests.post(expenses_url, json=expense_data_2, **admin_credentials)
    assert response_2.status_code == 201
    response_2 = response_2.json()

    response = requests.get(f"{expenses_url}?category={ExpensesCategoryEnum.SAVING}", **admin_credentials)
    assert response.status_code == 200
    response = response.json()

    responses_url = [expense["url"] for expense in response["results"]]
    assert response_1["url"] in responses_url
    assert response_2["url"] not in responses_url


import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from PyPDF2 import PdfWriter, PdfReader
import io
from reportlab.pdfbase import pdfmetrics
from reportlab.pdfbase.ttfonts import TTFont
from reportlab.pdfgen.canvas import Canvas
from reportlab.lib import pagesizes

# ======== Plotting Util ========

# assign numbers for sorting when combining outputs
export_counter = 1

def plot_amplitude_data(plot_title: str, axis1_name: str, resolution, data1, data1dots: list = [None], axis2_name: str = "", data2: list = [None], data2dots: list = [None], graph_on_same_axis: bool = False, export: bool = True, custom_prefix: str = ""):
    global export_counter

    x = np.linspace(0, len(data1) / resolution, len(data1))

    plt.figure()
    fig, ax = plt.subplots()
    ax.plot(x, data1, "-b", label="data1")
    if len(data1dots) > 1 and data1dots[0] != None:
        ax.plot(x, data1dots, ".", color="#55AAFF", label="data1 dots")
    ax.set_xlabel("Time passed [s]")
    ax.set_ylabel(axis1_name, color="blue")

    # set the x-spine
    ax.spines['left'].set_position('zero')  # type: ignore
    # turn off the right spine/ticks
    ax.spines['right'].set_color('none')
    ax.yaxis.tick_left()

    # set the y-spine
    ax.spines['bottom'].set_position('zero')  # type: ignore

    # turn off the top spine/ticks
    ax.spines['top'].set_color('none')
    ax.xaxis.tick_bottom()

    if len(data2) > 1 and data2[0] != None:
        ax2 = ax
        if not graph_on_same_axis:
            ax2 = ax.twinx()
        ax2.plot(x, data2, "-r", label="data2")
        if len(data2dots) > 1 and data2dots[0] != None:
            ax2.plot(x, data2dots, ".", color='#FFA500', label="data2 dots")
        ax2.set_xlabel("Time passed [s]")
        ax2.set_ylabel(axis2_name, color="red")
    plt.title(plot_title)

    if export:
        name = plot_title.lower().replace(" ", "_")
        plt.savefig(
            f"summarized_plots/png/({custom_prefix}a_{export_counter}){name}.png")
        plt.savefig(
            f"summarized_plots/pdf/({custom_prefix}a_{export_counter}){name}.pdf")
        export_counter += 1
    plt.show()

export_counter = 1

def plot_graph(plot_title: str, axis_name: str, points_x, points_val, graph_x, graph_y, y_axis_limit, export: bool = True, custom_prefix: str = ""):
    """
    Usage example:
    >>> t = np.arange(0, 5, 0.2)
    >>> plot_graph("", "", ..., ..., t, t ** 2)
    """
    global export_counter

    plt.figure()
    fig, ax = plt.subplots()
    ax.plot(points_x, points_val, ".", color="#55AAFF", label="points")
    ax.plot(graph_x, graph_y, "-r", label="function")
    ax.set_ylim(ymax=y_axis_limit)
    ax.set_xlabel("Points [1]")
    ax.set_ylabel(axis_name, color="blue")
    plt.title(plot_title)

    if export:
        name = plot_title.lower().replace(" ", "_")
        plt.savefig(
            f"summarized_plots/png/({custom_prefix}b_{export_counter}){name}.png")
        plt.savefig(
            f"summarized_plots/pdf/({custom_prefix}b_{export_counter}){name}.pdf")
        export_counter += 1
    plt.show()

def plot_4_curves__vs_time(data1, data2, data3, data4, steps_per_second, y_axis_title):
    x1 = np.linspace(0, len(data1) / steps_per_second, len(data1))
    x2 = np.linspace(0, len(data2) / steps_per_second, len(data2))
    x3 = np.linspace(0, len(data3) / steps_per_second, len(data3))
    x4 = np.linspace(0, len(data4) / steps_per_second, len(data4))

    plt.figure()
    fig, ax = plt.subplots()
    ax.plot(x1, data1)
    ax.plot(x2, data2)
    ax.plot(x3, data3)
    ax.plot(x4, data4)
    ax.set_xlabel("Verstrichene Zeit [s]")
    ax.set_ylabel(y_axis_title)
    plt.title(f"{y_axis_title} gegen Zeit")
    plt.show()

def create_pdf_text_page(filename: str, text: str, page_size=pagesizes.landscape(pagesizes.A5)):
    global A5

    # PDF page with info data
    # src: https://stackoverflow.com/a/17538003/19474335

    packet = io.BytesIO()
    cvs = Canvas(packet, bottomup=False, pagesize=page_size)

    # utf-8 encoding support: https://stackoverflow.com/a/17011377/19474335
    pdfmetrics.registerFont(TTFont('Verdana', 'Verdana.ttf'))
    cvs.setFont("Verdana", 11)

    line_height = 15
    y_counter = 2 * line_height
    for line in text.split("\n"):
        cvs.drawString(40, y_counter, line)
        y_counter += line_height
    cvs.save()

    # move to the beginning of the BytesIO buffer
    # packet.seek(0)

    new_pdf = PdfReader(packet)
    with open(filename.replace(".pdf", "") + ".pdf", "wb") as outStream:
        output = PdfWriter()
        output.add_page(new_pdf.pages[0])
        output.write(outStream)


#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
from kazoo.client import KazooClient

__name__ = "weichigong"
__version__ = '1.0.3'
__author__ = 'dashixiong'
__author_email__ = 'dashixiong.lee@gmail.com'

class zconfig:

    def __init__(self, zkHosts, app, env):
        self.app = app
        self.env = env
        self.client = KazooClient(hosts=zkHosts)
        self.client.start()

    def getPath(self, path):
        return os.path.join('/', self.app, self.env, path)

    def set(self, path, value):
        fullPath = self.getPath(path)
        self.client.ensure_path(fullPath)
        self.client.set(fullPath, value)

    def get(self, path):
        fullPath = self.getPath(path)
        return self.client.get(fullPath)[0].decode('utf-8')


import uuid
from random import randint
from src.infratructure.json_parser import JsonParser
from src.infratructure.serializable_object import SerializableObject

class PersonModel(SerializableObject):

    def __init__(self, id: int, nick: str, photo: str, name: str = None):
        self.id = id
        self.nick = nick
        self.photo = photo
        self.name = name

    @classmethod
    def random(cls):
        id = randint(0, 10)
        nick = str(uuid.uuid4())
        photo = str(uuid.uuid4())
        name = str(uuid.uuid4())

        return cls(id=id, nick=nick, photo=photo, name=name)

    @classmethod
    def from_json(cls, json):
        id = JsonParser.try_get_parameter_with_sub_name(json, "member", "id")
        nick = JsonParser.try_get_parameter_with_sub_name(json, "member", "name")
        photo = JsonParser.try_get_parameter_with_two_sub_name(json, "member", "photo", "highres_link")

        return cls(id=id, nick=nick, photo=photo, name=None)


import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from corai_util.tools.src.function_file import is_empty_file

from data_input.json.parameter_loader import fetch_param_json_loader_simulation, fetch_param_json_loader_itideep
from root_dir import linker_path_to_result_file
from src.estim_hawkes.estim_hawkes import Estim_hawkes

sns.set()

STR_CONFIG = "MSE"
(STR_CONFIG, NB_SIMUL, SEED, UNDERLYING_FUNCTION_NUMBER, _, KERNEL_DIVIDER,
 NB_DIFF_TIME_ESTIM, DIM, STYL, NB_POINTS_TT, id_hp, parameters, t0, time_batch,
 fct_parameters, true_breakpoints, _, _, _) = fetch_param_json_loader_simulation(False, STR_CONFIG)
(L, R, h, l, CONSIDERED_PARAM, ALL_KERNELS_DRAWN,
 TYPE_ANALYSIS, NUMBER_OF_BREAKPOINTS, MODEL,
 MIN_SIZE, WIDTH) = fetch_param_json_loader_itideep(flagprint=True, str_config=STR_CONFIG)
# should match the data given in the script.sh
NB_T_MAX = 10  # from 1 to 10.
NB_TH_OF_CURRENT_ESTIMATION = 2  # any int > 0. Represents the refinement of the ITiDeEP.
# 1 is the first naive estimation.
# The number given is the number of lines on the plot / nb of repetition of the estimation process undergone.
# Only possible to plot all the lines (1, 2...) together and not a subset of it not including the lower part.
#########
LIST_T_MAX = np.linspace(6000, 33000, NB_T_MAX)
#######################################################

# TODO  explain gather result in readme + explain MSE pipeline.
# We use this file to gather the estimation together (gather function) and then plot the curve of the MSE.

matrix_err_tmax_APE = np.zeros((NB_TH_OF_CURRENT_ESTIMATION, len(LIST_T_MAX)))
matrix_err_tmax_SPE = np.zeros((NB_TH_OF_CURRENT_ESTIMATION, len(LIST_T_MAX)))

iter_refinement = NB_TH_OF_CURRENT_ESTIMATION
while iter_refinement > 0:  # we collect the data from
    # NB_TH_OF_CURRENT_ESTIMATION to 1 by reducing by 1 at every iteration.

    for i_tmax in range(len(LIST_T_MAX)):
        ######################
        # gather results of previous estimation for a given T max
        ######################
        path_result_directory = linker_path_to_result_file(["MSE",
                                                            f"{STR_CONFIG}_res_{iter_refinement}",
                                                            f"data_{i_tmax}", ""])
        assert not is_empty_file(path_result_directory), \
            f"file must contain some data. Directory {path_result_directory} is empty."

        list_estim_hp = Estim_hawkes.folder_csv2list_estim(path_result_directory)
        estim_hp = Estim_hawkes.merge(list_estim_hp)  # new estim gathered result
        path_super_result = linker_path_to_result_file(
            ["MSE",
             f"{STR_CONFIG}_res_{iter_refinement}",
             f"data_together_{i_tmax}",
             f"results_together.csv"])
        estim_hp.to_csv(path_super_result)  # saved gather result
        ######################
        # compute error:
        ######################
        path_result_res = linker_path_to_result_file(
            ["MSE", f"{STR_CONFIG}_res_{iter_refinement}", f"data_together_{i_tmax}", "results_together.csv"])
        print("Reading: ", path_result_res)

        estim_hp = Estim_hawkes.from_csv(path_result_res)

        estim_hp.add_SPE_APE_col()  # computed the SRE per parameter

        groupby_param, keys = estim_hp.groupby(['parameter', 'm', 'n'])
        total_SPE_APE = (groupby_param.get_group(('alpha', 0, 0))[["time estimation", 'SPE', 'APE']]
                         .sort_values(by="time estimation").reset_index(drop=True))  # a copy is made
        # : we create a container where the error is aggregated.
        total_SPE_APE['SPE'] = 0  # we empty the values inside the column
        total_SPE_APE['APE'] = 0  # we empty the values inside the column
        for key in keys:
            ordered_SPE_APE = (groupby_param.get_group(key)[["time estimation", 'SPE', 'APE']]
                               .sort_values(by="time estimation").reset_index(drop=True))
            # sort to be sure we add the correct values together, drop index for prettiness.
            total_SPE_APE['SPE'] += ordered_SPE_APE['SPE']
            total_SPE_APE['APE'] += ordered_SPE_APE['APE']

        # MISRE = total_SRE.mean()["RSE"] # this is wrong. We need to compute it by hand.
        # It does not account for non converging estimations.
        total_SPE_APE_grouped = total_SPE_APE.groupby("time estimation")  # we groupby so we compute the integral
        MISPE = 0
        MIAPE = 0

        # compute the mean squared error and compute the mean absolute error
        for time in total_SPE_APE_grouped.groups:
            average_per_time = total_SPE_APE_grouped.get_group(time).mean()
            MISPE += average_per_time['SPE'] / len(total_SPE_APE_grouped.groups)
            MIAPE += average_per_time['APE'] / len(total_SPE_APE_grouped.groups)

        matrix_err_tmax_SPE[iter_refinement - 1, i_tmax] = MISPE  # store result
        matrix_err_tmax_APE[iter_refinement - 1, i_tmax] = MIAPE  # store result

    iter_refinement -= 1

dict_result = {"MISPE": matrix_err_tmax_SPE.flatten(),
               "MIAPE": matrix_err_tmax_APE.flatten(),
               "nb application ITiDeEP": np.repeat(range(NB_TH_OF_CURRENT_ESTIMATION), NB_T_MAX),
               "T max": np.tile(LIST_T_MAX, NB_TH_OF_CURRENT_ESTIMATION)}
data_err = pd.DataFrame(dict_result)

fig, ax = plt.subplots(2, 1)
sns.lineplot(x="T max", y="MISPE",
             hue="nb application ITiDeEP", marker='o',
             legend='full', ci=None, err_style="band",
             palette='Dark2', ax=ax[0],
             data=data_err)

sns.lineplot(x="T max", y="MIAPE",
             hue="nb application ITiDeEP", marker='o',
             legend='full', ci=None, err_style="band",
             palette='Dark2', ax=ax[1],
             data=data_err)

path_save_plot = linker_path_to_result_file(["MSE", f"MSE_result_{NB_TH_OF_CURRENT_ESTIMATION}" + '.png'])
fig.savefig(path_save_plot, dpi=500)
plt.show()


from otree.api import *

from settings import SESSION_CONFIGS

doc = """
Your app description
"""

class Constants(BaseConstants):
    name_in_url = 'Intro'
    players_per_group = None
    num_rounds = 1
    max_payoff = "£2.20"
    money = "£3.00"
    total_balls = "five"
    no_task_balls = "three"

    # create a vector to randomise treatment
    num_participants = 350          # note this should be substantially larger than the number of participants I actually intend to hire, because some Prolificers will join the session but not complete
    num_blocks = -1*( -num_participants // 14) # I'm gonna create blocks within which the treatment is exactly balanced (2 in LC, 2 in LN, 5 in HC, 5 in HN). Then add the blocks together to get to the desired number of participants.
    import random

    treatment_block = list(range(1,15))
    treatment_assignment = []
    for i in range(num_blocks):
        treatment_assignment = treatment_assignment + treatment_block

    random.shuffle(treatment_assignment)
    for i in range(len(treatment_assignment)):
            if treatment_assignment[i] <= 2:
                treatment_assignment[i] = "LC"
            elif treatment_assignment[i] > 2 and treatment_assignment[i] <= 4:
                treatment_assignment[i] = "LN"
            elif treatment_assignment[i] > 4 and treatment_assignment[i] <= 9:
                treatment_assignment[i] = "HC"
            elif treatment_assignment[i] >9:
                treatment_assignment[i] = "HN"

class Subsession(BaseSubsession):
    pass

def creating_session(subsession):
        import itertools, random

        treatment_assignment = itertools.cycle(Constants.treatment_assignment)
        for player in subsession.get_players():

            # determine treatment
            player.participant.treatment = next(treatment_assignment)
            player.treatment = player.participant.treatment

            # practice maths questions - randomly select two to show in instructions
            practice_maths_qs_index = list(range(4))
            random.shuffle(practice_maths_qs_index)
            player.participant.mathspractice_q1 = practice_maths_qs_index[0]
            player.participant.mathspractice_q2 = practice_maths_qs_index[1]

class Group(BaseGroup):
    pass

class Player(BasePlayer):
    ProlificID = models.StringField()
    treatment = models.StringField()
    start_epochtime = models.IntegerField()
    start_clocktime = models.StringField()

    # maths practice questions
    q1 = models.StringField(
        label = "A shop has an offer: buy 8 kiwis, and every extra kiwi after that is half price. A man goes to the shop and pays £4.50 for some kiwis. The full price of a kiwi is £0.50. How many does he buy?",
        choices = [
            "9",
            "12",
            "10",
            "15"
        ],
        widget = widgets.RadioSelectHorizontal,
        blank=True)
    q2 = models.StringField(
        label = "A hairdresser has an offer: every third visit is free. They charge £48 for a haircut. Last year Sarah paid £144 for a haaircut. How many times did she go?",
        choices = [
            "Two times",
            "Three times",
            "Four times",
            "Five times"
        ],
        widget = widgets.RadioSelectHorizontal,
        blank=True)
    q3 = models.StringField(
        label = "A woman walks from the bottom to the top of a hill. She starts at 9.40am and arrives at the top at 10.20 am. She takes a rest for ten minutes. Then she walks back down. On the way down she walks twice as fast as she did on the way up. What time is it when she reaches the bottom of the hill?",
        choices = [
            "11.20",
            "10.40",
            "10.50",
            "11.10"
        ],
        widget = widgets.RadioSelectHorizontal,
        blank=True)
    q4 = models.StringField(
        label = "A trader buys a painting for £120 and sells it for £170. They pay a £10 transaction fee. Their profit expressed as a percentage of total cost is:",
        choices = [
            "50%",
            "60%",
            "80%",
            "33%"
        ],
        widget = widgets.RadioSelectHorizontal,
        blank=True)

# PAGES
class Consent(Page):
    def is_displayed(player):
        # record time player entered application
        import time
        time_in = round(time.time())
        player.start_epochtime = time_in
        player.participant.start_epochtime = time_in
        player.start_clocktime = time.strftime('%H:%M:%S', time.localtime(time_in))
        return 1

class ProlificID(Page):
    form_model = 'player'
    form_fields = ['ProlificID']

class Introduction(Page):
    form_model = 'player'

    def get_form_fields(player: Player):
            questions = ['q1','q2','q3','q4']
            form_fields = [
                        questions[player.participant.mathspractice_q1]
                        ]
            return form_fields

page_sequence = [Consent, ProlificID, Introduction]