import requests
import json ,os,pprint

while True:
	if os.path.exists("s_courses.json")==False:
		a=requests.get('http://saral.navgurukul.org/api/courses')
		pprint.pprint(a.text)
		b=a.text
		c=open('s_courses.json','w')
		json.dump(b,c)
		c.close()

	c=open('s_courses.json','r')
	jsl1=json.load(c)
	jsl2=json.loads(jsl1)
	print(type(jsl2))
	c.close()

	for i in jsl2:
		for j in range(len(jsl2[i])):
			if j<9:
				x='0'+str(j+1)
				print(x,'-',jsl2[i][j]['id'],' ','name -',jsl2[i][j]['name'])
			else:
				print(1+j,'-',jsl2[i][j]['id'],' ','name -',jsl2[i][j]['name'])

	j=int(input('enter no.'))
	print(j,'-',jsl2[i][j-1]['id'],' ','name -',jsl2[i][j-1]['name'],'\n')
	x=jsl2[i][j-1]['id']

	while True:

		a1=requests.get(' http://saral.navgurukul.org/api/courses/'+str(x)+'/exercises')
		b1=a1.text
		if os.path.exists("'exercises_'+str(x)+'.json'")==False:
			c1=open('exercises_'+str(x)+'.json','w+')
			json.dump(b1,c1)
			c1.close()
		c1=open('exercises_'+str	(x)+'.json','r')
		js1=json.load(c1)
		js2=json.loads(js1)
		c.close()
		for i in js2:
			for j in range(len(js2[i])):
				print('\t',j+1,'-',js2[i][j]['name'])
		inp=input('''
			*** YOU WANT TO
			PRIVIOUS PAGE - THEN PRESS ("P")
			NEXT PAGE     - TNEN PRESS ("N")
			STARTING PAGE - THEN PRESS ("S")''').lower()
		if inp=='p':
			x-=1
		elif(inp=='n'):
			x+=1
		else:
			break
		if x<14:
			x=14
			print('have no PRIVIOUS PAGE ')
		elif(x>92):
			x=92
			print('have no NEXT PAGE')
		print('Id-',x)
	inp=input('**you want to exit then press "exit"').lower()
	if 'e' in inp:
		break


# importing required modules

import tkinter as tk
from tkinter import filedialog
import xlsxwriter

from init_row_data import init_row_data
from compute_B import compute_B
from compute_A import compute_A

# console interface

print("Power bills data extractor in txt\nver.: 1.0, 10/03/2022\nAuthor: Fellipe Filgueira")
input("Press ENTER to begin...")
print("\nProcessing...")

# load files paths

root = tk.Tk()
root.withdraw()
root.call('wm', 'attributes', '.', '-topmost', True)
filePath = filedialog.askopenfilename(multiple=False) # True to load multiple files

# creating .xlsx object

findPath = filePath[len(filePath)::-1]
findPath = findPath[findPath.find("/"):]
findPath = findPath[len(findPath)::-1]
workbook = xlsxwriter.Workbook(findPath + 'data.xlsx')
worksheet = workbook.add_worksheet('main')
data_dict = init_row_data()
data_headers = list(data_dict.keys())
for itr in range(len(data_headers)): worksheet.write(0, itr, data_headers[itr])

# reading txt file

file = open(filePath, 'r', encoding='utf-8')
file_items = file.readlines()

# set lists

text = [[]]
data = []

# set loop variables

count = 0
j = 0
row = 1

# convert str into list based in break points

for item in file_items:
    if 'Federal' in item:
        text[count].append(item)
        if item != file_items[-1]:
            text.append(list())
            count += 1
    if item[0:2] == '11' or item[0:2] == '12'or item[0:2] == '13' or item[0:2] == '14' and 'Federal' not in item:
        text[count].append(item)
    else:
        pass

# algorithm loop

while j < len(text):
    if 'B3' in text[j][1]:
        data = init_row_data()
        data = compute_B(j, text, data)
        for i in range(len(list(data))): worksheet.write(row, i, data[list(data)[i]])
        row += 1
    if 'A4' in text[j][1] or 'A3' in text[j][1]:
        data = init_row_data()
        data = compute_A(j, text, data)
        for i in range(len(list(data))): worksheet.write(row, i, data[list(data)[i]])
        row += 1
    if text[j][1][81:83] not in ['B3', 'A3', 'A4'] :
        print(text[j][1][26:36] + " bill was not correctly identified")
    else:
        pass
    j += 1

# close files

workbook.close()
file.close()

# console interface

print("\ndata.xlsx file has been successfully saved!")
input("Press ENTER to finish...")


#!/usr/bin/env python

import os
from setuptools import setup

# Utility function to read the README file.
# Used for the long_description.  It's nice, because now 1) we have a top level
# README file and 2) it's easier to type in the README file than to put a raw
# string in below ...
def read(fname):
    return open(os.path.join(os.path.dirname(__file__), fname)).read()

setup(
    name = "geweb",
    version = "0.1.0",
    author = "Artem Sazhin",
    author_email = "arts@point.im",
    description = ("Asyncronous micro web framework based on gevent."),
    long_description=read("README.md"),
    license = "BSD",
    keywords = "web framework gevent jinja2",
    url = "http://bitbucket.org/arts/geweb",
    packages = ["geweb", "geweb.template", "geweb.session", "geweb.util", "geweb.db", "geweb.db.pgsql"],
    install_requires = ["gevent", "jinja2", "argparse", "setproctitle"],
    scripts = ["geweb/bin/geweb"],
    include_package_data = True,

    classifiers = [
        "Development Status :: 3 - Alpha",
        "Programming Language :: Python",
        "Topic :: Internet :: WWW/HTTP",
        "Topic :: Internet :: WWW/HTTP :: Dynamic Content",
        "Topic :: Software Development :: Libraries :: Application Frameworks",
        "License :: OSI Approved :: BSD License",
        "Intended Audience :: Developers",
    ],
)


import os
import gym
import time
import argparse
import numpy as np
import matplotlib.pyplot as plt
from EA_components_OhGreat.EA import EA
from EA_components_OhGreat.Recombination import *
from EA_components_OhGreat.Mutation import *
from EA_components_OhGreat.Selection import *
from src.Evaluation import *
from src.Network import *
from src.utilities import *

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-exp_name', action='store',
                        dest='exp_name', type=str,
                        default='test_experiment',
                        help="Defines the name of the experiment.")
    parser.add_argument('-b', action='store',
                        dest='budget', type=int,
                        default=10000,
                        help="Defines the total amount of evaluations.")
    parser.add_argument('-min', action='store_true',
                        dest='minimize',
                        help="Use this flag if the problem is minimization.")
    parser.add_argument('-r', action='store',
                        dest='recombination', type=str,
                        default=None,
                        help="Defines the recombination strategy.")
    parser.add_argument('-m', action='store',
                        dest='mutation', type=str,
                        default='IndividualSigma',
                        help="Defines the mutation strategy.")
    parser.add_argument('-s', action='store',
                        dest='selection', type=str,
                        default='PlusSelection',
                        help="Defines the selection strategy.")
    parser.add_argument('-ps', action='store',
                        dest='parents_size', type=int,
                        default=2,
                        help="Defines the number of parents per generation.")
    parser.add_argument('-os', action='store',
                        dest='offspring_size', type=int,
                        default=4,
                        help="Defines the number of offspring per generation.")
    parser.add_argument('-mul', action='store',
                        dest='one_fifth_mul', type=float,
                        default=0.9,
                        help="Defines the multiplier for the one fifth success rule.")
    parser.add_argument('-pat', action='store',
                        dest='patience', type=int,
                        default=None,
                        help="Defines the wait time before resetting sigmas.")
    parser.add_argument('-exp_reps', action='store',
                        dest='exp_reps', type=int,
                        default=5,
                        help="Defines the number of experiments to average results.")
    parser.add_argument('-train_reps', action='store',
                        dest='train_reps', type=int,
                        default=10,
                        help="Defines the number of evaluation repetitions to use during training.")
    parser.add_argument('-eval_reps', action='store',
                        type=int, default=100,
                        help="Defines the number of evaluation repetitions to run after \
                                'training' our candidate individuals.")
    parser.add_argument('-model', action='store',
                        dest='model', type=int,
                        default=0,
                        help="Defines the model architecture to use.")
    parser.add_argument('-env', action='store', type=str,
                        dest='env', default='CartPole-v1')
    parser.add_argument('-render_train', action='store_true',
                        help='use this flag to render the training process \
                                of the individuals')
    parser.add_argument('-render_eval', action='store_true',
                        help='use this flag to render the evaluation process \
                                after training our individuals')
    parser.add_argument('-virtual_display', action='store_true',
                        help='needed for headless servers when using render')
    parser.add_argument('-plot_name', action='store', type=str,
                        default=None)
    parser.add_argument('-env_threshold', action='store',
                        type=float, default=500,
                        help="Optimum value to set as horizontal line in plot")
    parser.add_argument('-v', action='store',
                        dest='verbose', type=int, default=0,
                        help="Defines the intensity of debug prints.")
    args = parser.parse_args()
    if args.verbose > 0:
        print(args)

    # used to train on headless servers
    if args.virtual_display:
        os.environ["SDL_VIDEODRIVER"] = "dummy"

    # create gym environment
    env = gym.make(args.env)
    if env is None:
        exit("Please select an environment")
    print()

    # environment specific parameters
    n_observations = np.sum([dim for dim in env.observation_space.shape])
    if env.action_space.__class__.__name__ == "Discrete":
        n_actions = env.action_space.n
    elif env.action_space.__class__.__name__ == "Box":
        n_actions = sum(env.action_space._shape)
    else:
        exit(f"{env.action_space.__class__.__name__} action space not yet implemented")

    # create an instance of the model
    if args.model == 0:
        model = NN_regression_0(n_observations, 4, 4, n_actions).to("cpu")
    elif args.model == 1:
        model = NN_regression_1(n_observations, 4, 4, n_actions).to("cpu")
    elif args.model == 2:
        model = NN_regression_2(n_observations, 4, 4, n_actions).to("cpu")
    elif args.model == 3:
        model = NN_regression_3(n_observations, 4, 8, n_actions).to("cpu")
    else:
        exit("Choose a valid model")

    # define es individual size
    individual_size = model.total_params

    print(f"Model architecture: {args.model}\nEnvironment: {env.unwrapped.spec.id}\nNumber of observations: {n_observations}\nNumber of actions: {n_actions}\nIndividual size: {individual_size}")
    print()

    # EA specific parameters
    minimize = args.minimize
    budget = args.budget
    patience = args.patience
    parents_size = args.parents_size
    offspring_size = args.offspring_size
    # Recombination specific controls
    if args.recombination != None:
        recombination = globals()[args.recombination]()
    elif args.recombination == "GlobalIntermediary" and args.offspring_size > 1:
        print("GlobalIntermediary recombination cannot be used with more than one offspring.")
        print("Please use a valid configuration")
        exit()
    else: recombination = None
    # Mutation specific controls
    if args.mutation == "IndividualOneFifth":
        mutation = globals()[args.mutation](args.one_fifth_mul)
    else:
        mutation = globals()[args.mutation]()
    selection=globals()[args.selection]()
    evaluation = RewardMaximizationNN(env, model, args.train_reps,
                                        args.render_train)

    # loop through experiment
    best_results = []
    data_for_plots = []
    for i in range(args.exp_reps):
        # define new ea istance
        ea = EA(minimize=minimize, budget=budget, patience=patience,
        parents_size=parents_size, offspring_size=offspring_size,
        individual_size=individual_size, recombination=recombination,
        mutation=mutation, selection=selection, evaluation=evaluation,
        verbose=args.verbose)

        # run the ea
        start_time = time.time()
        best_ind, all_best_evals = ea.run()
        end_time = time.time()

        # keep track of results
        best_results.append([best_ind, max(all_best_evals)])
        data_for_plots.append(all_best_evals)
        print(f"Rep: {i+1} | average for {args.train_reps} evals: {max(all_best_evals)} | time: {np.round(end_time-start_time, 2)}")

    # save plot if name has been defined
    if args.plot_name != None:
        save_plot(args.plot_name, args.env_threshold, np.array(data_for_plots))

    # loop through final evalutation process for our best results
    eval_results = []
    for res in best_results:
        model.update_weights(res[0])
        curr_eval = eval(env, model,
                        args.eval_reps,
                        render=args.render_eval)
        eval_results.append(curr_eval)
    print("Evaluation results",np.round(eval_results, 2))

    # if saving is enabled
    if args.exp_name != None:
    # initialize directory and save results
        if not os.path.exists('model_weights'):
            os.makedirs('model_weights')
        # get the best individual
        best_ind_idx = np.argmax(eval_results)
        best_ind = best_results[best_ind_idx][0]
        # update model weights to match best inidividual
        model.update_weights(best_ind)
        # Save best individual as model weights
        torch.save(model.state_dict(), 'model_weights/'+args.exp_name)
        # Also save it as a numpy array
        np.save('model_weights/'+args.exp_name+'.npy', best_ind)

if __name__ == "__main__":
    main()


"""

"""
# Email: Kun.bj@outlook.com
import collections
import copy
import json
import os
import pickle
import time
import traceback
from collections import Counter
from pprint import pprint

import numpy as np
from sklearn.preprocessing import StandardScaler

from fkm import vis
from fkm.cluster import centralized_minibatch_kmeans
from fkm.utils.utils_func import dump, obtain_true_centroids
from fkm.utils.utils_stats import evaluate2
from fkm.utils.utils_func import timer
# These options determine the way floating point numbers, arrays and
# other NumPy objects are displayed.
# np.set_printoptions(precision=3, suppress=True)
from fkm.vis.visualize import plot_2gaussian, plot_3gaussian

np.set_printoptions(precision=3, suppress=True, formatter={'float': '{:.3f}'.format}, edgeitems=120, linewidth=100000)

def save_history2txt(seed_history, out_file='.txt'):
	"""
		with open(seed_file + '.txt', 'w') as file:
			file.write(json.dumps(seed_history))  # not working
	Returns
	-------

	"""

	def format(data):
		res = ''
		if type(data) == dict:
			for k, v in data.items():
				res += f'{k}: ' + format(v) + '\n'
		elif type(data) == list:
			res += f'{data} \n'
		else:
			res += f'{data} \n'

		return res

	with open(out_file, 'w') as f:
		f.write('***Save data with pprint\n')
		pprint(seed_history, stream=f, sort_dicts=False)  # 'sort_dicts=False' works when python version >= 3.8

		f.write('\n\n***Save data with recursion')
		res = format(seed_history)
		f.write(res)

# def normalize2(raw_x, raw_y, raw_true_centroids, splits, params, is_federated=False):
# 	"""
# 	Only for diagonal covariance matrix:
# 		even for federated kmeans, we still can get global mean and std from each client data.
# 		Based on that, for centralized and federated kmeans, we can use the same standscaler.
#
# 	Parameters
# 	----------
# 	raw_x
# 	raw_y
# 	raw_true_centroids
# 	splits
# 	params
# 	is_federated
#
# 	Returns
# 	-------
#
# 	"""
# 	is_show = params['is_show']
# 	normalize_method = params['normalize_method']
# 	# do normalization
# 	if normalize_method == 'std':
# 		# collects all clients' data together
# 		x = copy.deepcopy(raw_x)
# 		y = copy.deepcopy(raw_y)
# 		new_true_centroids = copy.deepcopy(raw_true_centroids)
# 		for spl in splits:  # train and test
# 			x[spl] = np.concatenate(x[spl], axis=0)
# 		global_stdscaler = StandardScaler()
# 		global_stdscaler.fit(x['train'])
# 		for spl in splits:  # train and test
# 			new_true_centroids[spl] = global_stdscaler.transform(raw_true_centroids[spl])
#
# 		if is_federated == False:  # for centralized normalization
# 			new_x = copy.deepcopy(raw_x)
# 			new_y = copy.deepcopy(raw_y)
# 			for i_, _ in enumerate(new_x['train']):
# 				new_x['train'][i_] = global_stdscaler.transform(new_x['train'][i_])
# 				new_x['test'][i_] = global_stdscaler.transform(new_x['test'][i_])
# 			params['stdscaler'] = global_stdscaler
#
# 			if is_show:
# 				if '3GAUSSIANS' in params['p0']:  # for plotting
# 					plot_3gaussian(new_x['train'][0], new_y['train'], new_x['train'][1], new_y['train'][1],
# 					               new_x['train'][2], new_y['train'][2], params, title='std')
# 		else:  # federated kmeans
# 			# for each client, we can get mean and std and then use them to get global_std
# 			new_x = copy.deepcopy(raw_x)
# 			new_y = copy.deepcopy(raw_y)
# 			stds = [[]] * len(new_x['train'])  # get number of clients
# 			N = 0
# 			for i_, _ in enumerate(new_x['train']):  # for each client
# 				data = new_x['train'][i_]
# 				n = len(data)
# 				stds[i_] = (n, np.mean(data, axis=0), np.std(data, axis=0))
# 				N += n
#
# 			dim = new_x['train'][0].shape[1]
# 			global_mean = [[]] * dim
# 			global_std = [[]] * dim
# 			# the following method only works for diagonal covariance matrix
# 			for i_ in range(dim):
# 				# compute global mean and std given each client's mean and std
# 				# global_mean = 1/N * (\sum client1  + \sum client2 + \sum client3)
# 				# 			  = 1/N * (n1 * mu1 + n2 * mu2 + n3 * mu3)
# 				global_mean[i_] = sum([n[i_] * mu[i_] for n, mu, s in stds]) / N
#
# 				# global_var = E(x-mu)**2 = (1/N * (\sum x**2)) - mu**2
# 				# 					   = (1/N * (\sum client1**2 + \sum client2**2+ \sum client3**2)) - global_mean **2
# 				#                      = (1/N * (n1 * var1 + n2 * var2 + n3 * var3)) - global_mean**2
# 				#                      = (1/N * (n1 * std1**2 + n2 * std2** + n3 * std3**2)) - global_mean**2
# 				global_std[i_] = (sum([n[i_] * s[i_] ** 2 for n, mu, s in stds]) / N - global_mean[i_] ** 2) ** (1 / 2)
#
# 			# for each client, then normalize its data use the global mean and std
# 			global_stdscaler2 = StandardScaler()
# 			global_stdscaler2.mean_ = global_mean
# 			global_stdscaler2.scale_ = global_std
#
# 			for i_, _ in enumerate(new_x['train']):
# 				new_x['train'][i_] = global_stdscaler2.transform(new_x['train'][i_])
# 				new_x['test'][i_] = global_stdscaler2.transform(new_x['test'][i_])
# 			params['stdscaler'] = global_stdscaler2
#
# 			if is_show:
# 				if '3GAUSSIANS' in params['p0']:  # for plotting
# 					plot_3gaussian(new_x['train'][0], new_y['train'], new_x['train'][1], new_y['train'][1],
# 					               new_x['train'][2], new_y['train'][2], params, title='std')
#
# 	else:
# 		new_x, new_y, new_true_centroids = raw_x, raw_y, raw_true_centroids
#
# 	return new_x, new_y, new_true_centroids

def normalize(raw_x, raw_y, raw_true_centroids, splits, params):
	"""
	For federated kmeans, we still can get global mean and std from each client data.
		Based on that, for centralized and federated kmeans, we can use the same standscaler.

	Parameters
	----------
	raw_x
	raw_y
	raw_true_centroids
	splits
	params

	Returns
	-------

	"""
	try:
		is_show = params['is_show']
	except Exception as e:
		is_show = params['IS_SHOW']
	try:
		normalize_method = params['normalize_method']
	except Exception as e:
		normalize_method = params['NORMALIZE_METHOD']

	try:
		algorithm_name = params['p0']
	except Exception as e:
		algorithm_name = params['ALGORITHM']['name']
	# do normalization
	if normalize_method == 'std':
		# collects all clients' data together and get global stdscaler
		x = copy.deepcopy(raw_x)
		y = copy.deepcopy(raw_y)
		new_true_centroids = copy.deepcopy(raw_true_centroids)
		for spl in splits:  # train and test
			x[spl] = np.concatenate(x[spl], axis=0)

		global_stdscaler = StandardScaler()  # we can get the same global_stdscaler using each client mean and std.
		global_stdscaler.fit(x['train'])
		# print(global_stdscaler.mean_, global_stdscaler.scale_)

		for spl in splits:  # train and test
			new_true_centroids[spl] = global_stdscaler.transform(new_true_centroids[spl])
			# print(new_true_centroids[spl])

		# normalize data
		new_x = copy.deepcopy(raw_x)
		new_y = copy.deepcopy(raw_y)
		for spl in splits:
			for i_ in range(len(new_x[spl])):
				new_x[spl][i_] = global_stdscaler.transform(new_x[spl][i_])  # normalize data first
		# for j_ in set(new_y[spl][i_]):
		# 	new_true_centroids[spl][j_] = np.mean(new_x[spl][i_], axis=0)  # get new centroids

		# if is_show:
		# 	if '2GAUSSIANS' in algorithm_name:  # for plotting
		# 		plot_2gaussian(new_x['train'][0], new_y['train'], new_x['train'][1], new_y['train'][1],
		# 		               params, title='std')
		# 	elif '3GAUSSIANS' in algorithm_name:  # for plotting
		# 		plot_3gaussian(new_x['train'][0], new_y['train'], new_x['train'][1], new_y['train'][1],
		# 		               new_x['train'][2], new_y['train'][2], params, title='std')
		# 	elif '4GAUSSIANS' in algorithm_name:  # for plotting
		# 		plot_2gaussian(new_x['train'][0], new_y['train'], new_x['train'][1], new_y['train'][1],
		# 		               params, title='std')
	elif normalize_method == 'min_max':
		raise NotImplementedError
	else:
		# new_x, new_y, new_true_centroids = raw_x, raw_y, raw_true_centroids
		# # collects all clients' data together and get global stdscaler
		# x = copy.deepcopy(raw_x)
		# y = copy.deepcopy(raw_y)
		# for spl in splits:  # train and test
		# 	x[spl] = np.concatenate(x[spl], axis=0)
		#
		# global_stdscaler = StandardScaler()  # we can get the same global_stdscaler using each client mean and std.
		# global_stdscaler.fit(x['train'])
		return raw_x, raw_y, raw_true_centroids, None

	return new_x, new_y, new_true_centroids, global_stdscaler

class NumpyEncoder(json.JSONEncoder):
	def default(self, obj):
		if isinstance(obj, np.ndarray):
			return obj.tolist()
		return json.JSONEncoder.default(self, obj)

@timer
def run_model(args):
	"""

		Parameters
		----------
		params
		KMeansFederated
		verbose:
			0 < verbose <= 5: info
			5 < verbose <= 10: debug

		Returns
		-------

		"""
	np.random.seed(args['SEED'])  # set the global seed for numpy
	VERBOSE = args['VERBOSE']
	SEPERTOR = args['SEPERTOR']
	SPLITS = args['SPLITS']
	dataset_name = args['DATASET']['name']
	N_CLUSTERS = args['N_CLUSTERS']
	N_CLIENTS = args['N_CLIENTS']
	# args['DATASET']['detail'] = f'{SEPERTOR}'.join(args['DATASET']['detail'], f'M_{N_CLIENTS}|K_{N_CLUSTERS}')
	# dataset_detail = args['DATASET']['detail']
	# algorithm_py_name = args['ALGORITHM']['py_name']
	# initial_method = args['ALGORITHM']['initial_method']
	server_init_method = args['ALGORITHM']['server_init_method']
	client_init_method = args['ALGORITHM']['client_init_method']

	N_REPEATS = args['N_REPEATS']
	TOLERANCE = args['TOLERANCE']
	NORMALIZE_METHOD = args['NORMALIZE_METHOD']
	# algorithm_detail = args['ALGORITHM']['detail']
	# OUT_DIR = args['OUT_DIR']
	# # if os.path.exists(OUT_DIR):
	# # 	shutil.rmtree(OUT_DIR)
	history_file = os.path.join(args['OUT_DIR'], 'history.dat')
	args['history_file'] = history_file
	# if os.path.exists(history_file):
	# # here could be some issue for multi-tasks, please double-check before calling this function.
	# 	return history_file

	# settings
	history = {'args': args}  # Store all results and needed data

	# data is fixed, however, the algorithm will have different initialization centroids with different seeds.
	with open(args['data_file'], 'rb') as f:
		raw_x, raw_y = pickle.load(f)
	print(f'data_file: ', args['data_file'])
	for split in SPLITS:
		X_ = raw_x[split]
		y_ = raw_y[split]
		print(len(X_), N_CLIENTS)
		for i_c in range(N_CLIENTS):
			print(f'{split}:{i_c}-th client raw data info, where mean: {np.mean(X_[i_c], axis=0)}, '
			      f'std: {np.std(X_[i_c], axis=0)}, and y: {collections.Counter(y_[i_c])}')

	if VERBOSE >= 1:
		# print raw_x and raw_y distribution
		for split in SPLITS:
			print(f'{split} set:')
			clients_x, clients_y = raw_x[split], raw_y[split]
			if VERBOSE >= 5:
				# print each client distribution
				for c_i, (c_x, c_y) in enumerate(zip(clients_x, clients_y)):
					print(f'\tClient_{c_i}, n_datapoints: {len(c_y)}, '
					      f'cluster_size: {sorted(Counter(c_y).items(), key=lambda kv: kv[0], reverse=False)}')

			y_tmp = []
			for vs in clients_y:
				y_tmp.extend(vs)
			print(f'n_{split}_clients: {len(clients_x)}, n_datapoints: {sum(len(vs) for vs in clients_y)}, '
			      f'cluster_size: {sorted(Counter(y_tmp).items(), key=lambda kv: kv[0], reverse=False)}')

	# obtain the true centroids given raw_x and raw_y
	raw_true_centroids = obtain_true_centroids(raw_x, raw_y, SPLITS, args)
	if VERBOSE >= 3:
		# print true centroids
		for split in SPLITS:
			true_c = raw_true_centroids[split]
			print(f'{split}_true_centroids:')
			print(true_c)

	# if algorithm_name == 'FEMNIST':
	# 	save_image2disk((raw_x, raw_y), out_dir_i, params)

	# do normalization in global for kmeans and federated kmeans
	raw_x, raw_y, raw_true_centroids, global_stdscaler = normalize(raw_x, raw_y, raw_true_centroids, SPLITS, args)
	args['global_stdscaler'] = global_stdscaler
	history['raw_true_centroids'] = raw_true_centroids
	print(f'after normalization, true_centroids:\n{raw_true_centroids} \nwhen normalize_method = {NORMALIZE_METHOD}')
	# history = {'x': raw_x, 'y': raw_y, 'results': []}
	# SEEDS = [10 * v ** 2 for v in range(1, N_REPEATS + 1, 1)]
	SEEDS = [42]    # we fix the model seed; however, the data seed is different.
	history['SEEDS'] = SEEDS

	from fkm.cluster import centralized_kmeans, federated_server_init_first, federated_client_init_first, \
		federated_greedy_kmeans
	KMEANS2PY = {
		'centralized_kmeans': centralized_kmeans.KMeans,
		'federated_server_init_first': federated_server_init_first.KMeansFederated,
		'federated_client_init_first': federated_client_init_first.KMeansFederated,
		'federated_greedy_kmeans': federated_greedy_kmeans.KMeansFederated,
		'centralized_minibatch_kmeans': centralized_minibatch_kmeans.KMeans,
	}
	for idx_seed, seed in enumerate(SEEDS):  # repetitions:  to obtain average and std score.

		# if VERBOSE >= 2:
		print(f'\n***{idx_seed}th repeat with seed: {seed}:')
		X = copy.deepcopy(raw_x)
		Y = copy.deepcopy(raw_y)
		true_centroids = copy.deepcopy(raw_true_centroids)

		if not args['IS_FEDERATED']:
			# collects all clients' data together
			for spl in SPLITS:
				X[spl] = np.concatenate(X[spl], axis=0)
				Y[spl] = np.concatenate(Y[spl], axis=0)
				print(spl, X[spl].shape, Y[spl].shape)

		t1 = time.time()
		# for Centralized Kmeans, we use server_init_centroids as init_centroids.
		kmeans = KMEANS2PY[args['ALGORITHM']['py_name']](
			n_clusters=N_CLUSTERS,
			# batch_size=BATCH_SIZE,
			sample_fraction=1.0,
			epochs_per_round=args['CLIENT_EPOCHS'],
			max_iter=args['ROUNDS'],
			server_init_method=server_init_method,
			client_init_method=client_init_method,
			true_centroids=true_centroids,
			random_state=seed,
			learning_rate=0,
			adaptive_lr=0,
			epoch_lr=0,
			momentum=0,
			reassign_min=0,
			reassign_after=0,
			verbose=VERBOSE,
			tol=TOLERANCE,
			params=args
		)

		if VERBOSE > 5:
			# print all kmeans's variables.
			pprint(vars(kmeans))

		# During the training, we also evaluate the model on the test set at each iteration
		kmeans.fit(X, Y, SPLITS, record_at=None)

		# After training, we obtain the final scores on the train/test set.
		scores = evaluate2(
			kmeans=kmeans,
			x=X, y=Y,
			splits=SPLITS,
			federated=args['IS_FEDERATED'],
			verbose=VERBOSE,
		)
		# # To save the disk storage, we only save the first repeat results.
		# if params['p0'] == 'FEMNIST' and s_i == 0:
		# 	try:
		# 		predict_n_saveimg(kmeans, x, y, SPLITS, SEED,
		# 		                  federated=params['is_federated'], verbose=VERBOSE,
		# 		                  out_dir=os.path.join(out_dir_i, f'SEED_{SEED}'),
		# 		                  params=params, is_show=is_show)
		# 	except Exception as e:
		# 		print(f'Error: {e}')
		# # traceback.print_exc()

		t2 = time.time()
		print(f'{idx_seed}th repeat with seed {seed} takes {(t2 - t1):.4f}s')

		# for each seed, we will save the results.
		history[seed] = {'initial_centroids': kmeans.initial_centroids,
		                 'true_centroids': kmeans.true_centroids,
		                 'final_centroids': kmeans.centroids,
		                 'training_iterations': kmeans.training_iterations,
		                 'history': kmeans.history,
		                 'scores': scores, 'duration': t2 - t1}
		if dataset_name != 'FEMNIST':
			try:
				# # save the current 'history' to disk before plotting.
				seed_file = os.path.join(args['OUT_DIR'], f'SEED_{seed}', f'~history.dat')
				dump(history[seed], out_file=seed_file)
				save_history2txt(history[seed], out_file=seed_file + '.txt')
			except Exception as e:
				print(f'save_history2txt() fails when SEED={seed}, Error: {e}')
		if VERBOSE >= 2:
			pprint(f'seed:{seed}, '
			       f'scores:{scores}')
	try:
		results_avg = vis.visualize.stats_single(history)
	except Exception as e:
		traceback.print_exc()
		results_avg = {}
	# save the current 'history' to disk before plotting.
	history['results_avg'] = results_avg
	dump(history, out_file=history_file)

	try:
		vis.visualize.plot_single(history)
	except Exception as e:
		traceback.print_exc()
	# out_file = os.path.join( args['OUT_DIR'], f'varied_clients-Server_{server_init_method}-Client_{client_init_method}')
	# print(out_file)
	# dump(stats, out_file=out_file + '.dat')

	# # dump(histories, out_file=out_file + '-histories.dat')
	# with open(history_file + '-histories.txt', 'w') as file:
	# 	file.write(json.dumps(history, cls=NumpyEncoder))  # use `json.loads` to do the reverse

	return history_file


from django.conf.urls import url,include
from . import views as app_v

from django.contrib.auth import views
from app.forms import LoginForm

urlpatterns = [
    url(r'^$',app_v.index, name='home'),
    url(r'^login/$', views.login, {'template_name': 'app/login.html', 'authentication_form': LoginForm}, name='login'),
    url(r'^logout/$', views.logout, {'next_page': '/app/'}, name='logout'),
    url(r'^register/$', app_v.register, name="register"),
    url(r'^success/', app_v.success, name="success"),
    url(r'^success1/', app_v.success1, name="success1"),
    url(r'^check/username/(?P<username>[-\w.]+)/$', app_v.check),
    url(r'^create/$', app_v.createPromo, name="createPromo"),
    url(r'^withdraw/$', app_v.withdrawPromo, name="withdrawPromo"),
    url(r'^sendOtp/$', app_v.sendOtp),
    url(r'^getToken/$', app_v.get_token),
    url(r'^checkBalance/$', app_v.checkBalance),
    url(r'^generateChecksum/$', app_v.generateChecksum),
    url(r'^makeTransaction/$', app_v.makeTransaction),
    url(r'^doTransfer/$', app_v.doTransfer),
]


import os
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforoms

device = 'cuda' if torch.cuda.is_available() else 'cpu'
print('Using {} device'.format(device))

class NeuralNetwork(nn.Module):
    def __init__(self):
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512),
            nn.ReLu(),
            nn.Linear(512, 512),
            nn.Relu(),
            nn.Linear(512, 10),
            nn.Relu()
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

model = NeuralNetwork.to(device)
# print(model)

x = torch.rand(1, 28, 28, device=device)
logits = model(X)
pred_probab = nn.Softmax(dim=1)(logits)
y_pred = pred_probab.argmax(1)
# print(f"Predicted class: {y_pred}")

loss_fn = nn.CrossEntropyLoss()

def train_loop(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    for batch, (X, y) in enumerate(dataloader):
        # Compute prediction and loss
        pred= model(X)
        loss = loss_fn(pred, y)

        #Backpropagation
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if batch % 100 == 0:
            loss, current = loss.item(), batch * len(X)
            print(f"loss: {loss:>7f} [{current:>5d}/{size:>5d}]")

def test_loop(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    test_loss, correct = 0, 0

    with torch.no_grad():
        for X, y in dataloader:
            pred = model(X)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()

    test_loss /= size
    correct /= size
    # print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:8f} \n")

    loss_fn = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

    epochs = 10
    for t in range(epochs):
        print(f"Epochs {t+1}\n-------------------")
        traing_loop(traing_dataloader, model, loss_fn, optimizer)
        test_loop(test_dataloader, mode, loss_fn)
    print("Done!")

    model = models.vgg16(pretrained=True)
    torch.save(model.state_dict(), 'model_weigths.pth')

    model = models.vgg16()
    model.load_state_dict(torch.load('model_weights.pth'))
    model.eval()

    torch.save(model, 'model.pth')

    model = torch.load('model.pth')


from django.urls import path, re_path
from django.contrib.auth.decorators import permission_required
from django.views.generic import RedirectView
from .feeds import PublishTrackFeed
from . import views

urlpatterns = [
    path('', RedirectView.as_view(url='catalogue/', permanent=False)),

    path('images/', views.image_list, name='documents_image_list'),
    path('image/<slug:slug>/', views.image, name="documents_image"),
    path('image/<slug:slug>/publish', views.publish_image,
            name="documents_publish_image"),

    path('catalogue/', views.document_list, name='documents_document_list'),
    path('user/', views.my, name='documents_user'),
    path('user/<username>/', views.user, name='documents_user'),
    path('users/', views.users, name='documents_users'),
    path('activity/', views.activity, name='documents_activity'),
    re_path(r'^activity/(?P<isodate>\d{4}-\d{2}-\d{2})/$',
        views.activity, name='documents_activity'),

    path('upload/', views.upload, name='documents_upload'),

    path('create/<slug:slug>/',
        views.create_missing, name='documents_create_missing'),
    path('create/', views.create_missing, name='documents_create_missing'),

    path('book/<slug:slug>/publish', views.publish, name="documents_publish"),

    path('book/<slug:slug>/', views.book, name="documents_book"),
    path('book/<slug:slug>/gallery/',
            permission_required('documents.change_book')(views.GalleryView.as_view()),
            name="documents_book_gallery"),
    path('book/<slug:slug>/xml', views.book_xml, name="documents_book_xml"),
    path('book/dc/<slug:slug>/xml', views.book_xml_dc, name="documents_book_xml_dc"),
    path('book/<slug:slug>/txt', views.book_txt, name="documents_book_txt"),
    path('book/<slug:slug>/html', views.book_html, name="documents_book_html"),
    path('book/<slug:slug>/epub', views.book_epub, name="documents_book_epub"),
    path('book/<slug:slug>/mobi', views.book_mobi, name="documents_book_mobi"),
    path('book/<slug:slug>/pdf', views.book_pdf, name="documents_book_pdf"),
    path('book/<slug:slug>/pdf-mobile', views.book_pdf, kwargs={'mobile': True}, name="documents_book_pdf_mobile"),

    path('chunk_add/<slug:slug>/<slug:chunk>/',
        views.chunk_add, name="documents_chunk_add"),
    path('chunk_edit/<slug:slug>/<slug:chunk>/',
        views.chunk_edit, name="documents_chunk_edit"),
    path('book_append/<slug:slug>/',
        views.book_append, name="documents_book_append"),
    path('chunk_mass_edit',
        views.chunk_mass_edit, name='documents_chunk_mass_edit'),
    path('image_mass_edit',
        views.image_mass_edit, name='documents_image_mass_edit'),

    path('track/<slug:slug>/', PublishTrackFeed()),
    path('active/', views.active_users_list, name='active_users_list'),
    path('active.csv', views.active_users_list, kwargs={'csv': True}, name='active_users_csv'),

    path('mark-final/', views.mark_final, name='mark_final'),
    path('mark-final-completed/', views.mark_final_completed, name='mark_final_completed'),
]


"""
Handlers related to "reports" about students' activity - badge evidence pages
and badge issuing routines.
"""
from controllers.utils import BaseHandler, ReflectiveRequestHandler, XsrfTokenManager
from common import prefetch
import pprint
from models.roles import Roles
import re
from collections import defaultdict
from models import transforms
from modules.badges.badge_models import Badge, BadgeAssertion
from report import UnitReport, PartReport, ExpertBadgeReport
from report import _parts as part_config
from models.models import Student
from models.models import EventEntity
from wiki_models import Annotation
from jinja2 import Markup
import urllib
import wtforms as wtf
from google.appengine.api import users
from google.appengine.ext import db
from google.appengine.ext import deferred
import page_templates
from common.querymapper import LoggingMapper
import logging

import wtforms as wtf

def exam_display_choices(exam_info):
    "Decide which display choices are available to a student, depending on their score."
    choices = [
            ('blank', '(Blank)'),
            ]
    default = 'blank'

    if exam_info['completed']:
        choices.append( ('completed', 'Submitted the exam') )

    if exam_info['did_pass']:
        choices.append( ('passed', 'Passed the exam, with at least (passing score) out of 100%') )
        choices.append( ('scored', 'Passed the exam, scoring (your score) out of 100%') )

        default = 'passed'

    return locals()

class EvidenceHandler(BaseHandler, ReflectiveRequestHandler):
    """
    Handler to show badge evidence - both the 'top' page and the text of each unit's wikifolio.

    Also the display settings for each badge - whether to show the text of the wikifolio
    and how to display the exam.
    """
    get_actions = ['view', 'settings']
    default_action = 'view'
    post_actions = ['save_settings']

    class SettingsForm(wtf.Form):
        report_id = wtf.HiddenField()
        units_are_public = wtf.BooleanField(
                "Show my Wikifolio entries for this badge on the evidence page?")

        # Will set choices dynamically.
        exam_display = wtf.SelectField(
                "How to display exam scores on the evidence page?")

        review_is_public = wtf.BooleanField(
                """Show the instructor's review of the paper on the evidence page?  Only
                relevant for the term paper.""")

    def get_settings(self):
        user = self.personalize_page_and_get_enrolled()
        if not user:
            return

        try:
            report = PartReport.get_by_id(int(self.request.GET.get('id', -1)))
        except ValueError:
            report = None
        if not report:
            self.abort(404, "That evidence report was not found.")

        if not self.can_edit(user, report):
            self.abort(403)

        form = self.SettingsForm(
                report_id=report.key().id(),
                units_are_public=report.units_are_public,
                exam_display=report.exam_display,
                review_is_public=report.review_is_public)

        #if report.part != 4:
            #del form.review_is_public

        display_field_params = {
                'choices': [('blank', '(Blank)')],
                'default': 'blank'
                }

        if report.assessment_scores:
            if len(report.assessment_scores) > 1:
                logging.warning("Evidence page settings assuming there's just one exam per part, but there is more than one")
            display_field_params = exam_display_choices(
                    report.assessment_scores[0])

        form.exam_display.choices = display_field_params['choices']
        form.exam_display.default = display_field_params['default']

        self.template_value['report'] = report
        self.template_value['form'] = form
        self.template_value['xsrf_token'] = XsrfTokenManager.create_xsrf_token('save_settings')
        self.template_value['navbar'] = {}
        self.template_value['action_url'] = '/badges/evidence?action=save_settings'
        self.template_value['badge_name'] = report._config['name']
        self.render('wf_evidence_settings.html')

    def can_edit(self, user, report):
        #Roles.is_course_admin(self.app_context)
        return report.student_email == user.key().name()

    def post_save_settings(self):
        user = self.personalize_page_and_get_enrolled()
        if not user:
            return

        form = self.SettingsForm(self.request.POST)

        try:
            report = PartReport.get_by_id(int(form.report_id.data))
        except ValueError:
            report = None
        if not report:
            self.abort(404, "That evidence report was not found.")

        if not self.can_edit(user, report):
            self.abort(403, "You can't edit that user's report.")

        if report.assessment_scores:
            display_field_params = exam_display_choices(
                    report.assessment_scores[0])
            form.exam_display.choices = display_field_params['choices']
            form.exam_display.default = display_field_params['default']
        else:
            del form.exam_display

        if not form.validate():
            self.redirect('/')
            return

        report.units_are_public = form.units_are_public.data
        report.review_is_public = form.review_is_public.data
        if report.assessment_scores:
            report.exam_display = form.exam_display.data
        report.put()

        EventEntity.record(
                'set-evidence-settings',
                users.get_current_user(),
                transforms.dumps({
                    'part': report.part,
                    'slug': report.slug,
                    'review_is_public': report.review_is_public,
                    'public': report.units_are_public,
                    'exam_display': report.exam_display,
                    'email': user.key().name()
                    }))

        self.template_value['navbar'] = {}
        self.template_value['content'] = '''
OK, saved settings.

        Back to your account page...
'''
        self.render('bare.html')

    def head(self):
        try:
            report = PartReport.get_by_id(int(self.request.GET.get('id', -1)))
        except ValueError:
            report = None
        if not report:
            self.abort(404)

        self.abort(200)

    def get_view(self):
        try:
            report = PartReport.get_by_id(int(self.request.GET.get('id', -1)))
        except ValueError:
            report = None
        if not report:
            self.abort(404)

        if not report.exam_display:
            if report.assessment_scores:
                display_info = exam_display_choices(report.assessment_scores[0])
                report.exam_display = display_info['default']
            else:
                report.exam_display = 'blank'

        self.report = report
        self.template_value['inline_save'] = lambda: ''
        self.template_value['navbar'] = {}
        self.template_value['author'] = self.report.student
        self.template_value['review_is_public'] = self.report.review_is_public
        if report.units_are_public:
            self.template_value['unit_link'] = self._unit_link
        else:
            self.template_value['unit_link'] = None
            self.template_value['no_unit_links'] = True
        self.template_value['unit_title'] = self._unit_title

        self.unit_num = None
        try:
            if report.units_are_public:
                self.unit_num = int(self.request.GET.get('unit', ''))
        except ValueError:
            self.unit_num = None

        if self.unit_num:
            self.unit = self.report.get_unit(self.unit_num)
            if self.unit:
                self.render_unit()
                return
            else:
                logging.warning('Could not find the right unit %d for PartReport %s',
                        self.unit_num, self.report.key())
        self.render_top()

    def _unit_link(self, unit):
        return self.request.path + "?" + urllib.urlencode({
            'id': self.request.GET['id'],
            'unit': unit,
            })

    def _unit_title(self, unit):
        unit_obj = self.find_unit_by_id(unit)
        title = unit_obj.title
        title = re.sub(r'\(.*?\)', '', title)
        return title.strip()

    def render_unit(self):
        fields = {
                k: Markup(v) for k,v in self.unit.wiki_fields.iteritems()}
        fields['reflection'] = Markup('
Removed from this view for peers\' privacy.
--BOOC Instructors and Tech Staff
')
        self.template_value['fields'] = fields
        self.template_value['unit'] = self.find_unit_by_id(self.unit_num)
        self.template_value['report'] = self.unit
        self.template_value['badge_slug'] = self.report.badge.key().name()
        # THIS is kinda magic - we render a page like wf_temp_u1.html, but have
        # it inherit from wf_evidence.html rather than wf_page.html.  This
        # removes the comment section, for instance.
        self.template_value['layout_template'] = 'wf_evidence.html'
        self.template_value['review'] = Annotation.reviews(whose=self.report.student, unit=self.unit_num).get()
        self.render(page_templates.templates[self.unit_num])

    def render_top(self):
        self.template_value['part'] = self.report
        self.render('wf_evidence_top.html')

class ExpertEvidenceHandler(BaseHandler, ReflectiveRequestHandler):
    """
    Handler for the "expert" badge - issued at the end of the course.
    This badge has links to all the previous badges, rather than links
    to the internal units.
    """
    get_actions = ['view', 'settings']
    default_action = 'view'
    post_actions = ['save_settings']

    def can_edit(self, user, report):
        return report.student_email == user.key().name()

    class SettingsForm(wtf.Form):
        report_id = wtf.HiddenField()

        # Will set choices dynamically.
        exam_display = wtf.SelectField(
                "How to display exam scores on the evidence page?")

    def get_settings(self):
        user = self.personalize_page_and_get_enrolled()
        if not user:
            return

        try:
            report = ExpertBadgeReport.get_by_id(int(self.request.GET.get('id', -1)))
        except ValueError:
            report = None
        if not report:
            self.abort(404, "That evidence report was not found.")

        if not self.can_edit(user, report):
            self.abort(403)

        form = self.SettingsForm(
                report_id=report.key().id(),
                exam_display=report.exam_display)

        display_field_params = {
                'choices': [('blank', '(Blank)')],
                'default': 'blank'
                }

        display_field_params = exam_display_choices(
                report.final_exam_score)

        form.exam_display.choices = display_field_params['choices']
        form.exam_display.default = display_field_params['default']

        self.template_value['report'] = report
        self.template_value['form'] = form
        self.template_value['xsrf_token'] = XsrfTokenManager.create_xsrf_token('save_settings')
        self.template_value['navbar'] = {}
        self.template_value['badge_name'] = "Assessment Expert Badge"
        self.template_value['action_url'] = '/badges/expert_evidence?action=save_settings'
        self.render('wf_evidence_settings.html')

    def post_save_settings(self):
        user = self.personalize_page_and_get_enrolled()
        if not user:
            return

        form = self.SettingsForm(self.request.POST)

        try:
            report = ExpertBadgeReport.get_by_id(int(form.report_id.data))
        except ValueError:
            report = None
        if not report:
            self.abort(404, "That evidence report was not found.")

        if not self.can_edit(user, report):
            self.abort(403, "You can't edit that user's report.")

        display_field_params = exam_display_choices(
                report.final_exam_score)
        form.exam_display.choices = display_field_params['choices']
        form.exam_display.default = display_field_params['default']

        if not form.validate():
            self.redirect('/')
            return

        report.exam_display = form.exam_display.data
        report.put()

        EventEntity.record(
                'set-evidence-settings',
                users.get_current_user(),
                transforms.dumps({
                    'slug': report.slug,
                    'exam_display': report.exam_display,
                    'email': user.key().name()
                    }))

        self.template_value['navbar'] = {}
        self.template_value['content'] = '''
OK, saved settings.

        Back to your account page...
'''
        self.render('bare.html')

    def get_view(self):
        try:
            report = ExpertBadgeReport.get_by_id(int(self.request.GET.get('id', -1)))
        except ValueError:
            report = None
        if not report:
            self.abort(404)

        all_assertions_q = BadgeAssertion.all()
        all_assertions_q.filter('recipient', report.student_key)
        all_assertions_q.filter('revoked', False)
        all_assertions = all_assertions_q.run(limit=10)

        if not report.exam_display:
            if report.final_exam_score:
                display_info = exam_display_choices(report.final_exam_score)
                report.exam_display = display_info['default']
            else:
                report.exam_display = 'blank'

        self.template_value['report'] = report
        self.template_value['navbar'] = {}
        self.template_value['author'] = report.student
        # TODO: links to the other badges

        all_assertions = prefetch.prefetch_refprops(
                list(all_assertions), BadgeAssertion.badge)
        course_parts = {'practices': None, 'principles': None, 'policies': None}
        for ass in all_assertions:
            name_parts = ass.badge_name.split('.')
            if name_parts[0] in course_parts:
                if (not course_parts[name_parts[0]]) or (len(name_parts) > 1):
                    course_parts[name_parts[0]] = ass
        self.template_value['part_assertions'] = course_parts

        self.render('wf_expert_evidence.html')

def combine_badge_slug_parts(*parts):
    return '.'.join(parts)

def choose_badge_version(slug, completion):
    # Choose the base badge version - normal vs expertise.
    # Doesn't choose Leader vs. normal - that is later.

    if not completion['units']:
        return None
    elif not completion['assessments']:
        return slug
    else:
        return combine_badge_slug_parts(slug, 'expertise')

class SingleIssueHandler(BaseHandler):
    """
    Decide whether to issue a badge to one student, and optionally really issue it.
    """
    class Form(wtf.Form):
        part = wtf.IntegerField('Which part of the course to issue a badge for? (1,2,3)')
        really_save = wtf.BooleanField('Really issue the badge and freeze the scores?', default=False)
        re_run = wtf.BooleanField('Re-run all unit and part reports? Will delete old ones if you also choose Really freeze above.', default=False)
        email = wtf.StringField('The email of the student to reconsider')

    def get(self):
        if not users.is_current_user_admin():
            self.abort(403)
        form = self.Form()
        self.template_value['form'] = form
        self.template_value['xsrf_token'] = XsrfTokenManager.create_xsrf_token('post')
        self.template_value['action_url'] = self.request.url
        self.template_value['title'] = 'Reconsider a single participant'
        self.render('badge_bulk_issue.html')

    def post(self):
        if not users.is_current_user_admin():
            self.abort(403)
        if not XsrfTokenManager.is_xsrf_token_valid(self.request.POST.get('xsrf_token', ''), 'post'):
            self.abort(403, 'XSRF token failed.')
        form = self.Form(self.request.POST)
        if not form.validate():
            self.response.write('
'.join(form.errors))
            return

        student = Student.get_by_key_name(form.email.data)
        report = PartReport.on(student, course=self.get_course(),
                part=form.part.data,
                force_re_run=form.re_run.data,
                put=form.really_save.data)

        badge_version = choose_badge_version(part_config[form.part.data]['slug'], report.completion())
        if badge_version:
            badge = Badge.get_by_key_name(badge_version)
            if not badge:
                self.response.write(' There is no badge with key_name %s (so I cannot issue a badge)' % badge_version)
            if form.really_save.data and badge:
                b = Badge.issue(badge, student, put=False)
                b.evidence = self.request.host_url + '/badges/evidence?id=%d' % report.key().id()
                b.put()
                self.response.write('Issued badge %s!' % badge_version)
            else:
                self.response.write('Would have issued badge %s!' % badge_version)
        else:
            self.response.write('Not issuing because at least one of: %s' % (', '.join(report.incomplete_reasons)))

class BulkIssueMapper(LoggingMapper):
    """
    Issue completion badges to many folks!  This includes 'expertise' and 'knowledge' badges.
    """
    KIND = Student
    FILTERS = [('is_participant', True)]

    def __init__(self, really, course, part, host_url, re_run):
        super(BulkIssueMapper, self).__init__()
        self.really = really
        self.course = course
        self.part = part
        self.host_url = host_url
        self.re_run = re_run
        self.num_issued = 0

    def map(self, student):
        self.log.append('########## Student %s ##########' % student.key().name())
        report = PartReport.on(student, course=self.course, part=self.part,
                force_re_run=self.re_run, put=self.really)

        completion = report.completion()
        self.log.append(' Passed? %s.' % str(completion))

        badge_version = choose_badge_version(report.slug, completion)

        if badge_version:
            badge = Badge.get_by_key_name(badge_version)
            if not badge:
                self.log.append(' There is no badge with key_name %s (so I cannot issue a badge)' % badge_version)

            self.num_issued += 1
            if self.really and badge:
                b = Badge.issue(badge, student, put=False) # need to include evidence URL here somehow
                b.evidence = self.host_url + '/badges/evidence?id=%d' % report.key().id()
                b.put()
                self.log.append(' Issued badge, name=%s, assertion id=%d' % (
                    badge_version, b.key().id()))
                return ([b], [])
            else:
                self.log.append(' WOULD issue badge.')
        else:
            self.log.append('Not issuing because at least one of: %s' % (', '.join(report.incomplete_reasons)))

            ##TODO: this is not comprehensive: they could still have .expertise or .leader versions.
            #if self.really and badge:
                #Badge.ensure_not_issued(badge, student)
        return ([], [])

    def finish(self):
        self.log.append('DONE.  Issued %d badges total.' % self.num_issued)
        self._batch_write()

def choose_expert_badge_version(completion):
    badge_version = None
    if completion['badges']:
        badge_version = 'expert'
    if badge_version and completion['assessments'] \
            and all('expertise' in slug for slug in completion['badge_slugs']):
        badge_version = 'expert.expertise' # lol
    return badge_version

class BulkExpertBadgeIssueMapper(LoggingMapper):
    """
    Issue end-of-course badges to many folks.
    """
    KIND = Student
    FILTERS = [('is_participant', True)]

    def __init__(self, really, course, unused_part_num, host_url, force_re_run):
        LoggingMapper.__init__(self)
        self.really = really
        self.course = course
        self.host_url = host_url
        self.force_re_run = force_re_run

        self.num_issued = 0

    def map(self, student):
        self.log.append('--------------- Student %s' % student.key().name())

        report = ExpertBadgeReport.on(student, self.course,
                force_re_run=self.force_re_run, put=self.really)

        completion = report.completion()
        self.log.append(' Passed? %s.' % str(completion))

        badge_version = choose_expert_badge_version(completion)

        if badge_version:
            badge = Badge.get_by_key_name(badge_version)
            if not badge:
                self.log.append('no such badge! %s' % badge_version)
            self.num_issued += 1
            if self.really and badge:
                b = Badge.issue(badge, student, put=False) # need to include evidence URL here somehow
                b.evidence = self.host_url + '/badges/expert_evidence?id=%d' % report.key().id()
                b.put()
                self.log.append(' Issued badge, name=%s, assertion id=%d' % (
                    badge.key().name(), b.key().id()))
                return ([b], [])
            else:
                self.log.append(' WOULD issue badge %s' % badge_version)
        else:
            self.log.append('Incomplete, we are missing: %s' % (', '.join(report.incomplete_reasons())))

        return ([], [])

    def finish(self):
        self.log.append('DONE.  Issued %d badges total.' % self.num_issued)
        self._batch_write()

NOBODY = object()
def default_dict_entry():
    return ([NOBODY], -1)

class BulkLeaderIssueMapper(LoggingMapper):
    """
    Issue leader badges.  In this one, the badges aren't issued during the main phase
    of the 'mapper' loop over all the students - instead, the badges are issued
    in the "finish" phase.
    """
    KIND = Student
    FILTERS = [('is_participant', True)]

    def __init__(self, really, course, part, host_url, re_run):
        super(BulkLeaderIssueMapper, self).__init__()
        self.really = really
        self.course = course
        self.part = part
        self.host_url = host_url
        self.re_run = re_run
        self.best_by_group = defaultdict(default_dict_entry)
        self.leader_badge_keys = (
                part_config[part]['slug'] + '.expertise.leader',
                part_config[part]['slug'] + '.leader',
                )

        leader_badge_e = Badge.get_by_key_name(self.leader_badge_keys[0])
        leader_badge = Badge.get_by_key_name(self.leader_badge_keys[1])
        if not leader_badge and leader_badge_e:
            logging.warning('Missing a badge with one of: %s', str(self.leader_badge_keys))
            self.log.append('Missing a badge with one of: %s'% str(self.leader_badge_keys))
            if self.really:
                raise ValueError('Missing a badge with one of: %s'% str(self.leader_badge_keys))

    def map(self, student):
        self.log.append('######### Student %s ##########' % student.key().name())

        part_report = PartReport.on(student, course=self.course, part=self.part,
                force_re_run=self.re_run)
        completion = part_report.completion()
        base_badge_version = choose_badge_version(part_report.slug, completion)
        if not base_badge_version:
            self.log.append(' Skipping, since not complete.')
            return ([], [])

        self.log.append(' Part is complete, considering units.')

        unit_reports = part_report.unit_reports
        promotions = 0

        for ur in unit_reports:
            promotions += ur.promotions

        best_so_far = self.best_by_group[student.group_id][1]
        if promotions > best_so_far:
            self.best_by_group[student.group_id] = ([student.key().name()], promotions)

            self.log.append(' They have current best for group %s, with %d.' % (
                student.group_id, promotions))

        elif promotions == best_so_far:
            self.best_by_group[student.group_id][0].append(student.key().name())

            self.log.append(' They ARE TIED FOR CURRENT BEST for group %s, with %d.' % (
                student.group_id, promotions))
        return ([], [])

    def finish(self):
        if self.really:
            leader_badges = dict((key, Badge.get_by_key_name(key)) for key in self.leader_badge_keys)

        for group_id, (emails, count) in self.best_by_group.iteritems():
            self.log.append('Considering group %s, best score is %d' % (
                str(group_id), count))
            if count < 1:
                self.log.append('... Best score is too low, skipping.')
                continue
            if self.really:
                for email in emails:
                    report = PartReport.on(
                            db.Key.from_path(Student.kind(), email),
                            course=self.course, part=self.part,
                            force_re_run=self.re_run)
                    base_badge_version = choose_badge_version(report.slug, report.completion())
                    leader_badge_version = combine_badge_slug_parts(base_badge_version, 'leader')
                    leader_badge = leader_badges[leader_badge_version]
                    assert leader_badge, "Should be a badge called %s" % leader_badge_version

                    b = Badge.issue(leader_badge,
                            db.Key.from_path(Student.kind(), email), put=False)
                    b.evidence = self.host_url + '/badges/evidence?id=%d' % report.key().id()
                    b.put()
                    self.log.append('... ISSUED leader badge %s to %s, id=%d' % (
                        leader_badge_version, email, b.key().id()))
            else:
                self.log.append('... WOULD ISSUE leader badge to %s' % ' '.join(emails))
        self._batch_write()

class BulkExpertLeaderIssueMapper(LoggingMapper):
    """
    Expert badges go to people who pass all the parts of the course,
    and do the final exam, and do the survey.  Leaders are calculated
    based on who gets the most Exemplaries in the entire course.
    """
    KIND = ExpertBadgeReport

    def __init__(self, really, course, part, host_url, re_run):
        super(BulkExpertLeaderIssueMapper, self).__init__()
        self.really = really
        self.course = course
        self.host_url = host_url
        self.re_run = re_run
        self.best_by_group = defaultdict(default_dict_entry)

    def map(self, report):
        student = report.student
        self.log.append('---------------- Student %s' % student.key().name())

        if self.re_run:
            report._run(self.course)

        best_so_far = self.best_by_group[student.group_id][1]
        promotions = report.exemplary_count

        completion = report.completion()

        if promotions > best_so_far:
            self.log.append('New best for group %s, %d promotions' % (
                student.group_id, promotions))

            if completion['badges']:
                self.best_by_group[student.group_id] = ([student.key().name()], promotions)
            else:
                self.log.append('BUT, Skipping, not complete.')

        elif promotions == best_so_far:
            self.log.append('TIED best for group %s, %d promotions' % (
                student.group_id, promotions))

            if completion['badges']:
                self.best_by_group[student.group_id][0].append(student.key().name())
            else:
                self.log.append('BUT, Skipping, not complete.')

        return ([], [])

    def finish(self):
        leader_badge_slugs = ('expert.leader', 'expert.expertise.leader')
        leader_badges = dict((k, Badge.get_by_key_name(k)) for k in leader_badge_slugs)

        for group_id, (emails, count) in self.best_by_group.iteritems():
            self.log.append('Considering group %s, best score is %d' % (
                str(group_id), count))
            if count < 1:
                self.log.append('... Best score is too low, skipping.')
                continue
            for email in emails:
                report = ExpertBadgeReport.on(
                        db.Key.from_path(Student.kind(), email),
                        course=self.course,
                        force_re_run=self.re_run)
                base_badge_version = choose_expert_badge_version(report.completion())
                if not base_badge_version:
                    raise AssertionError('They should have passed, wat?')
                leader_badge_slug = base_badge_version + '.leader'
                if self.really:
                    leader_badge = leader_badges[leader_badge_slug]
                    b = Badge.issue(leader_badge,
                            db.Key.from_path(Student.kind(), email), put=False)
                    b.evidence = self.host_url + '/badges/expert_evidence?id=%d' % report.key().id()
                    b.put()
                    self.log.append('... ISSUED %s to %s, id=%d' % (
                        leader_badge_slug, email, b.key().id()))
                else:
                    self.log.append('... WOULD ISSUE %s to %s' % (
                        leader_badge_slug, email))
        self._batch_write()

issuer_mappers = {
        'completion': BulkIssueMapper,
        'leader': BulkLeaderIssueMapper,
        'expert': BulkExpertBadgeIssueMapper,
        'expert-leader': BulkExpertLeaderIssueMapper,
        }

class BulkIssuanceHandler(BaseHandler, ReflectiveRequestHandler):
    """
    Request handler for kicking off all these different kinds of badge issuing mappers.
    """
    default_action = 'prep'
    get_actions = ['prep', 'watch']
    post_actions = ['start']

    TITLE = 'Bulk Issue Badges'

    class IssueForm(wtf.Form):
        part = wtf.IntegerField('Which part of the course to issue a badge for? (1,2,3)')
        really_save = wtf.BooleanField('Really issue the badges and freeze the scores?', default=False)
        leader_or_completion = wtf.RadioField('Do you want to issue completion badges, or leader badges?',
                choices=[(k, k) for k in issuer_mappers.keys()])
        force_re_run_reports = wtf.BooleanField('Re-run all unit and part reports? Will delete old ones if you also choose Really freeze above.', default=False)
        one_email = wtf.TextField('Only consider one student? Enter their e-mail here.',
                validators=[wtf.validators.Optional()])

    def _action_url(self, action, **kwargs):
        params = dict(kwargs)
        params['action'] = action
        return '?'.join((
            self.request.path,
            urllib.urlencode(params)))

    def get_prep(self):
        if not users.is_current_user_admin():
            self.abort(403)
        self.render_form(self.IssueForm())

    def render_form(self, form):
        self.template_value['form'] = form
        self.template_value['xsrf_token'] = self.create_xsrf_token('start')
        self.template_value['action_url'] = self._action_url('start')
        self.template_value['title'] = self.TITLE
        self.render('badge_bulk_issue.html')

    def post_start(self):
        if not users.is_current_user_admin():
            self.abort(403)

        form = self.IssueForm(self.request.POST)
        if not form.validate():
            self.render_form(form)
            return

        REALLY = form.really_save.data
        part_num = form.part.data

        problems = set()
        student_infos = []

        issuer = issuer_mappers[form.leader_or_completion.data]

        job = issuer(REALLY, self.get_course(), part_num, self.request.host_url, form.force_re_run_reports.data)
        if form.one_email.data:
            job.FILTERS.append(
                    ('__key__', db.Key.from_path('Student', form.one_email.data)))
            logging.debug('Filters for issuing: %s', repr(job.FILTERS))
        job_id = job.job_id
        deferred.defer(job.run, batch_size=50)
        self.redirect(self._action_url('watch', job_id=job_id))

    def get_watch(self):
        if not users.is_current_user_admin():
            self.abort(403)

        job_id = self.request.GET.get('job_id', None)
        if not job_id:
            self.abort(404)

        messages = BulkIssueMapper.logs_for_job(job_id)

        self.template_value['title'] = self.TITLE
        self.template_value['problems'] = []
        self.template_value['log'] = messages
        self.render('badge_bulk_issue_done.html')

class DammitMapper(LoggingMapper):
    """Goes through all PartReports and re-runs their assessment grades."""
    KIND = PartReport

    def __init__(self, course):
        super(DammitMapper, self).__init__()
        self.course = course

    def map(self, report):
        self.log.append("Working on report for %s, part %d" % (report.student_email, report.part))
        had_before = bool(report.assessment_scores)
        report._run_assessments(self.course)
        has_after = bool(report.assessment_scores)
        self.log.append("had before? %s.  has after? %s" % (had_before, has_after))

        return ([report], [])

class DammitHandler(BaseHandler, ReflectiveRequestHandler):
    get_actions = ['start', 'watch']
    default_action = 'watch'

    def _action_url(self, action, **kwargs):
        params = dict(kwargs)
        params['action'] = action
        return '?'.join((
            self.request.path,
            urllib.urlencode(params)))

    def get_start(self):
        if not users.is_current_user_admin():
            self.abort(403)

        course = self.get_course()
        job = DammitMapper(course)
        job_id = job.job_id
        deferred.defer(job.run, batch_size=50)
        self.redirect(self._action_url('watch', job_id=job_id))

    def get_watch(self):
        if not users.is_current_user_admin():
            self.abort(403)

        job_id = self.request.GET.get('job_id', None)
        if not job_id:
            self.abort(404)

        messages = BulkIssueMapper.logs_for_job(job_id)

        self.template_value['title'] = "GRRRRRR"
        self.template_value['problems'] = []
        self.template_value['log'] = messages
        self.render('badge_bulk_issue_done.html')


import collections
import itertools
import json
import logging
import os
import pickle
import random

import numpy as np
from sklearn.model_selection import train_test_split

import config
from utils.utils import JsonlReader

logging.basicConfig(format='%(asctime)s : %(levelname)s : %(message)s', level=logging.INFO)
logger = logging.getLogger(__name__)
np.random.seed(config.SEED)
random.seed(config.SEED)

def get_groups_texts_from_umls_vocab(relation_text_to_groups, cui_to_entity_texts, fname_reltext_all_combos,
                                     load_existing=True):
    # Load rel groups file, if it exists
    if os.path.isfile(fname_reltext_all_combos) and load_existing:
        logger.info('Found existing relation text combination file -- loading relations.')
        with open(fname_reltext_all_combos, 'rb') as f:
            return pickle.load(f)
    else:
        logger.info("Not loading relation text combination file. Generating new set:")

    # Get all related CUI groups into set:
    groups = set()
    for relation_text in relation_text_to_groups:
        groups.update(relation_text_to_groups[relation_text])

    # Collect all combinations of related entities
    logger.info("Collecting all possible textual combinations of CUI groups ...")
    groups_texts = set()
    l = len(groups)

    for idx, (cui_src, cui_tgt) in enumerate(groups):
        if idx % 100000 == 0 and idx != 0:
            logger.info("Parsed {} groups of {}".format(idx, l))
        cui_src_texts = cui_to_entity_texts[cui_src]
        cui_tgt_texts = cui_to_entity_texts[cui_tgt]
        for cui_src_text_i in cui_src_texts:
            temp = list(zip([cui_src_text_i] * len(cui_tgt_texts), cui_tgt_texts))
            temp = ["\t".join(i) for i in temp]
            groups_texts.update(temp)

    # NOTE: this consumes a LOT of memory (~18 GB)! (clearing up memory takes around half an hour)
    logger.info("Collected {} unique tuples of (src_entity_text, tgt_entity_text) type.".format(len(groups_texts)))

    # Save rel groups:
    with open(fname_reltext_all_combos, 'wb') as f:
        logger.info('Saving relation text combination file.')
        pickle.dump(groups_texts, f)

    return groups_texts

def align_groups_to_sentences(groups_texts, jsonl_fname, output_fname):
    jr = JsonlReader(jsonl_fname)
    wf = open(output_fname, "w", encoding="utf-8", errors="ignore")

    logger.info("Aligning texts (sentences) to groups ...")
    pos_groups = set()
    neg_groups = set()

    for idx, jdata in enumerate(jr):
        if idx % 1000000 == 0 and idx != 0:
            logger.info("Processed {} tagged sentences".format(idx))

        # Permutations of size for matched entities in a sentence
        matched_perms = set(itertools.permutations(jdata['matches'].keys(), 2))

        # Left-hand-side (lhs) <==> right-hand-side (rhs)
        lhs2rhs = collections.defaultdict(list)
        rhs2lhs = collections.defaultdict(list)

        for group in matched_perms:
            src, tgt = group
            lhs2rhs[src].append(tgt)
            rhs2lhs[tgt].append(src)

        # Since `groups_texts` contain all possible groups that can exist
        # in the UMLS KG, for some relation, the intersection of this set
        # with matched permuted groups efficiently yields groups which
        # **do exist in KG for some relation and have matching sentences**.
        matched_perms = {"\t".join(m) for m in matched_perms}
        common = groups_texts.intersection(matched_perms)

        # We use sentence level noise, i.e., for the given sentence the
        # common groups represent positive groups, while the negative
        # samples can be generated as follows (like open-world assumption):
        #
        # For a +ve group, with prob. 1/2, remove the left (src) or right
        # (tgt) entity and replace with N entities such that the negative
        # group (e_orig, e_replaced) [for rhs] / (e_replaced, e_orig) [for lhs]
        # **must not be in KG for any relation**. This technique can possibly be
        # seen as creating hard negatives for same text evidence.
        output = {"p": set(), "n": set()}

        for group in common:
            pos_groups.add(group)
            src, tgt = group.split("\t")
            output["p"].add(group)
            # Choose left or right side to corrupt
            lhs_or_rhs = random.choice([0, 1])

            if lhs_or_rhs == 0:
                for corrupt_tgt in lhs2rhs[src]:
                    negative_group = "{}\t{}".format(src, corrupt_tgt)
                    if negative_group not in common:
                        output["n"].add(negative_group)
            else:
                for corrupt_src in rhs2lhs[tgt]:
                    negative_group = "{}\t{}".format(corrupt_src, tgt)
                    if negative_group not in common:
                        output["n"].add(negative_group)

        if output["p"] and output["n"]:
            no = list(output["n"])
            random.shuffle(no)
            # Keep number of negative groups at most as positives
            no = no[:len(output["p"])]
            output["n"] = no
            output["p"] = list(output["p"])
            neg_groups.update(no)
            jdata["groups"] = output
            wf.write(json.dumps(jdata) + "\n")

    # There will be lot of negative groups, so we will remove them next!
    logger.info("Collected {} positive and {} negative groups.".format(len(pos_groups), len(neg_groups)))

    return pos_groups, neg_groups

def pruned_triples(cui_to_entity_texts, relation_text_to_groups, pos_groups, neg_groups, min_rel_group=10,
                   max_rel_group=1500):
    logger.info("Mapping CUI groups to relations ...")
    group_to_relation_texts = collections.defaultdict(list)

    for idx, (relation_text, groups) in enumerate(relation_text_to_groups.items()):
        for group in groups:
            group_to_relation_texts[group].append(relation_text)  # can have multiple rel texts per group

    logger.info("Mapping relations to groups texts ...")
    relation_text_to_groups_texts = collections.defaultdict(set)

    for idx, (group, relation_texts) in enumerate(group_to_relation_texts.items()):
        if idx % 1000000 == 0 and idx != 0:
            logger.info("Mapped from {} groups".format(idx))

        cui_src, cui_tgt = group
        local_groups = set()
        cui_src_texts = cui_to_entity_texts[cui_src]
        cui_tgt_texts = cui_to_entity_texts[cui_tgt]

        for l1i in cui_src_texts:
            local_groups.update(list(zip([l1i] * len(cui_tgt_texts), cui_tgt_texts)))

        for lg in local_groups:
            if "\t".join(lg) in pos_groups:
                for relation_text in relation_texts:
                    relation_text_to_groups_texts[relation_text].add("\t".join(lg))

    logger.info("No. of relations before pruning: {}".format(len(relation_text_to_groups_texts)))

    # Prune relations based on the group size
    relations_to_del = list()
    for relation_text, groups_texts in relation_text_to_groups_texts.items():
        if (len(groups_texts) < min_rel_group) or (len(groups_texts) > max_rel_group):
            relations_to_del.append(relation_text)

    logger.info("Relations not matching the criterion of min, max group sizes of {} and {}.".format(min_rel_group,
                                                                                max_rel_group))
    # Delete relations not meeting min and max counts
    for r in relations_to_del:
        del relation_text_to_groups_texts[r]
    logger.info("Relations deleted: {}".format(relations_to_del))
    logger.info("No. of relations after pruning: {}".format(len(relation_text_to_groups_texts)))

    # Update positive groups
    new_pos_groups = set()
    entities = set()
    for relation_text, groups_texts in relation_text_to_groups_texts.items():
        for group_text in groups_texts:
            new_pos_groups.add(group_text)
            entities.update(group_text.split("\t"))

    logger.info("Updated no. of positive groups after pruning: {}".format(len(new_pos_groups)))
    logger.info("No. of unique entities: {}".format(len(entities)))

    # Update negative groups

    # 1) We apply the constraint that the negative groups must have positive
    # triples entities only
    new_neg_groups = set()

    for negative_group in neg_groups:
        src, tgt = negative_group.split("\t")
        if (src in entities) and (tgt in entities):
            new_neg_groups.add(negative_group)

    logger.info("[1] Updated no. of negative groups after pruning groups that are not in positive entities: {}".format(
        len(new_neg_groups)))

    # 2) Negative examples are used for NA / Other relation, which is just another class.
    # To avoid training too much on NA relation, we make a simple choice randomly taking
    # the same number of groups as largest group size positive class.
    max_pos_group_size = max([len(v) for v in relation_text_to_groups_texts.values()])
    new_neg_groups = list(new_neg_groups)
    random.shuffle(new_neg_groups)

    # Using 70% of positive groups to form negative groups
    new_neg_groups = new_neg_groups[:int(max_pos_group_size * 0.7)]
    logger.info(
        'Len of new_pos_groups: {}, Len of max_pos_group_size: {}'.format(len(new_pos_groups), max_pos_group_size))
    logger.info(
        "Number of negative groups after taking 70 percent more than positive groups: {}".format(len(new_neg_groups)))

    relation_text_to_groups_texts["NA"] = new_neg_groups

    # Collect triples now
    triples = set()
    for r, groups in relation_text_to_groups_texts.items():
        for group in groups:
            src, tgt = group.split("\t")
            triples.add((src, r, tgt))
    triples = list(triples)

    logger.info(" *** No. of triples (including NA) *** : {}".format(len(triples)))

    return triples

def filter_triples_with_evidence(triples, max_bag_size=32):
    group_to_relation_texts = collections.defaultdict(set)

    for ei, rj, ek in triples:
        group = "{}\t{}".format(ei, ek)
        group_to_relation_texts[group].add(rj)

    jr = JsonlReader(config.groups_linked_sents_file)

    group_to_data = collections.defaultdict(list)
    candid_groups = set(group_to_relation_texts.keys())

    for idx, jdata in enumerate(jr):
        if idx % 1000000 == 0 and idx != 0:
            logger.info("Processed {} lines for linking to triples".format(idx))
        common = candid_groups.intersection(jdata["groups"]["p"] + jdata["groups"]["n"])

        if not common:
            continue

        for group in common:
            src, tgt = group.split("\t")
            src_span = jdata["matches"][src]
            tgt_span = jdata["matches"][tgt]
            sent = jdata["sent"]
            sent = sent.replace("$", "")
            sent = sent.replace("^", "")

            # src entity mentioned before tgt entity
            if src_span[1] < tgt_span[0]:
                sent = sent[:src_span[0]] + "$" + src + "$" + sent[src_span[1]:tgt_span[0]] + "^" + tgt + "^" + sent[
                                                                                tgt_span[
                                                                                1]:]
                rel_dir = 1

            # tgt entity mentioned before src entity
            elif src_span[0] > tgt_span[1]:
                sent = sent[:tgt_span[0]] + "^" + tgt + "^" + sent[tgt_span[1]:src_span[0]] + "$" + src + "$" + sent[
                                                                                src_span[
                                                                                1]:]
                rel_dir = -1

            # Should not happen, but to be on safe side
            else:
                continue

            if group not in group_to_data:
                group_to_data[group] = collections.defaultdict(list)

            group_to_data[group][rel_dir].append(sent)

    # Adjust bag sizes
    new_group_to_data = dict()
    for group in list(group_to_data.keys()):
        src, tgt = group.split("\t")
        bag = list()
        for rel_dir in group_to_data[group]:
            bag.extend(group_to_data[group][rel_dir])
        if len(bag) > max_bag_size:
            bag = random.sample(bag, max_bag_size)
        else:
            idxs = list(np.random.choice(list(range(len(bag))), max_bag_size - len(bag)))
            bag = bag + [bag[i] for i in idxs]
        new_group_to_data["\t".join([src, tgt, "0"])] = {
            "relations": group_to_relation_texts[group],
            "bag": bag
        }
    group_to_data = new_group_to_data

    filtered_triples = set()
    for group in group_to_data:
        src, tgt, _ = group.split("\t")
        for relation in group_to_data[group]["relations"]:
            filtered_triples.add((src, relation, tgt))

    return filtered_triples, group_to_data

def remove_overlapping_sents(train_lines, dev_lines, test_lines):
    dev_test_sentences = set()
    for line in dev_lines:
        dev_test_sentences.update({s.replace("$", "").replace("^", "") for s in line["sentences"]})
    for line in test_lines:
        dev_test_sentences.update({s.replace("$", "").replace("^", "") for s in line["sentences"]})

    new_train_lines = list()

    for line in train_lines:
        new_sents = list()
        for sent in line["sentences"]:
            temp_sent = sent.replace("$", "").replace("^", "")
            if temp_sent not in dev_test_sentences:
                new_sents.append(sent)
        if not new_sents:
            continue
        bag = new_sents
        if len(bag) > config.bag_size:
            bag = random.sample(bag, config.bag_size)
        else:
            idxs = list(np.random.choice(list(range(len(bag))), config.bag_size - len(bag)))
            bag = bag + [bag[i] for i in idxs]
        line["sentences"] = bag
        new_train_lines.append(line)

    new_triples = set()

    for line in new_train_lines:
        src, tgt = line["group"]
        relation = line["relation"]
        new_triples.add((src, relation, tgt))

    return new_train_lines, new_triples

def create_data_split(triples):
    triples = list(triples)
    inds = list(range(len(triples)))
    y = [relation for _, relation, _ in triples]
    # train_dev test split
    train_dev_inds, test_inds = train_test_split(inds, stratify=y, test_size=0.2, random_state=config.SEED)
    y = [y[i] for i in train_dev_inds]
    train_inds, dev_inds = train_test_split(train_dev_inds, stratify=y, test_size=0.1, random_state=config.SEED)

    train_triples = [triples[i] for i in train_inds]
    dev_triples = [triples[i] for i in dev_inds]
    test_triples = [triples[i] for i in test_inds]

    logger.info(" *** Train triples : {} *** ".format(len(train_triples)))
    logger.info(" *** Dev triples : {} *** ".format(len(dev_triples)))
    logger.info(" *** Test triples : {} *** ".format(len(test_triples)))

    return train_triples, dev_triples, test_triples

def split_lines(triples, group_to_data):
    groups = set()
    for ei, _, ek in triples:
        groups.add("{}\t{}".format(ei, ek))
    lines = list()
    for group in groups:
        src, tgt = group.split("\t")
        G = ["\t".join([src, tgt, "0"]), ]
        for g in G:
            if g not in group_to_data:
                continue
            data = group_to_data[g]
            _, _, rel_dir = g.split("\t")
            rel_dir = int(rel_dir)
            for relation in data["relations"]:
                lines.append({
                    "group": (src, tgt),
                    "relation": relation,
                    "sentences": data["bag"],
                    "e1": data.get("e1", None), "e2": data.get("e2", None),
                    "reldir": rel_dir
                })
    return lines

def report_data_stats(lines, triples):
    stats = dict(
        num_of_groups=len(lines),
        num_of_sents=sum(len(line["sentences"]) for line in lines),
        num_of_triples=len(triples)
    )
    for k, v in stats.items():
        logger.info(" *** {} : {} *** ".format(k, v))

def write_final_jsonl_file(lines, output_fname):
    with open(output_fname, "w") as wf:
        for idx, line in enumerate(lines):
            wf.write(json.dumps(line) + "\n")

if __name__ == "__main__":
    # 0. Load UMLS vocab object
    logger.info("Loading UMLS vocab object `{}` ...".format(config.umls_cui_to_txts))
    with open(config.umls_cui_to_txts, "rb") as ctt, open(config.umls_reltxt_to_groups, "rb") as rttg:
        cui_to_entity_texts = pickle.load(ctt)
        relation_text_to_groups = pickle.load(rttg)
        ctt.close(), rttg.close()

    # 1. Collect all possible group texts from their CUIs
    groups_texts = get_groups_texts_from_umls_vocab(relation_text_to_groups, cui_to_entity_texts,
                                                    config.reltext_all_combos)

    # 2. Search for text alignment of groups (this can take up to 80~90 mins)
    pos_groups, neg_groups = align_groups_to_sentences(groups_texts, config.medline_linked_sents_file,
                                                       config.groups_linked_sents_file)

    # 3. From collected groups and pruning relations criteria, get final triples
    triples = pruned_triples(cui_to_entity_texts, relation_text_to_groups, pos_groups, neg_groups, config.min_rel_group,
                             config.max_rel_group)

    # 4. Collect evidences and filter triples based on sizes of collected bags
    triples, group_to_data = filter_triples_with_evidence(triples, config.bag_size)
    logger.info(" *** No. of triples (after filtering) *** : {}".format(len(triples)))

    E, R = set(), set()

    with open(config.triples_file_all, "w") as wf:
        for ei, rj, ek in triples:
            E.update([ei, ek])
            R.add(rj)
            line = "{}\t{}\t{}".format(ei, rj, ek)
            wf.write(line + '\n')

    with open(config.entities_file, "w") as wf:
        for e in E:
            wf.write("{}\n".format(e))

    with open(config.relations_file, "w") as wf:
        for r in R:
            wf.write("{}\n".format(r))

    logger.info(" *** No. of entities *** : {}".format(len(E)))
    logger.info(" *** No. of relations *** : {}".format(len(R)))

    # 5. Split into train, dev and test at triple level to keep zero triples overlap
    train_triples, dev_triples, test_triples = create_data_split(triples)
    train_lines = split_lines(train_triples, group_to_data)
    dev_lines = split_lines(dev_triples, group_to_data)
    test_lines = split_lines(test_triples, group_to_data)

    # Remove any overlapping test and dev sentences from training
    logger.info("Train stats before removing overlapping sentences ...")
    report_data_stats(train_lines, train_triples)
    train_lines, train_triples = remove_overlapping_sents(train_lines, dev_lines, test_lines)

    # Write triples file and final combined file
    logger.info("Train stats after removing dev + test overlapping sentences ...")
    triples_splits = [
        (config.complete_train, config.triples_file_train, train_triples, train_lines, "TRAIN"),
        (config.complete_dev, config.triples_file_dev, dev_triples, dev_lines, "DEV"),
        (config.complete_test, config.triples_file_test, test_triples, test_lines, "TEST")
    ]
    for complete_file, trip_file, trips, lines, split_name in triples_splits:
        print(split_name)
        report_data_stats(lines, trips)
        with open(trip_file, "w") as wf:
            for ei, rj, ek in trips:
                wf.write("{}\t{}\t{}\n".format(ei, rj, ek))

        # 6. Write train, dev, test files with sentence, group and relation
        logger.info("Creating training file at `{}` ...".format(complete_file))
        write_final_jsonl_file(lines, complete_file)