#!/usr/bin/env python
from github import Github
import argparse
import sys
import os
import time
import random
import math

class Collect():
	def __init__(self):
		self.user=''
		self.pwd=''
		self.token=''
		self.org=''
		self.repo=''
		self.branch=''
		self.event_type=''
		self.per_page=100
		self.client=''
		self.organization=''
		self.state=''
		self.sleep_time_length = 1 #5
		self.sleep_time_range = 1 #3

	def get_arguments(self):
		parser = argparse.ArgumentParser(description="Github data collection with PyGithub")
		parser.add_argument('--user', type=str, help="The name of the user")
		parser.add_argument('--pwd', type=str,  help="Password of the user")
		parser.add_argument('--token', type=str, help="The private token of the user")
		parser.add_argument('--org', type=str,  help="The name of the Organization")
		parser.add_argument('--repo', type=str,  help="The name of the Repositories")
		parser.add_argument('--branch',  type=str, help="The name of the branches")
		parser.add_argument('--state',  type=str, help="The State of the issues")
		parser.add_argument('--event_type', type=str, help="The type of the events(e.g issues, commits")
		parser.add_argument('--per_page', default=self.per_page,  type=int, help="The number of per_page requests")
		return parser.parse_args()

	def validate_arguments(self, args):
		"""Validate arguments entered by user"""
		if args.org == None:
			print('Please specify Organization name. Exiting.')
			sys.exit(0)
		if args.repo == None:
			print('Please specify Repositories name. Exiting.')
			sys.exit(0)
		if args.event_type == None:
			print('Please specify type of the event. Exiting.')
			sys.exit(0)
		if args.event_type == 'commits' and args.branch == None:
			print('Please specify branch name. Exiting.')
			sys.exit(0)
		if args.event_type == 'issues' and args.state == None:
			print('Please specify state of the issues. Exiting.')
			sys.exit(0)
		if args.event_type == 'pullRequests' and args.branch == None and args.state == None:
			print('Please specify branch and state of the pulls. Exiting.')
			sys.exit(0)

		return

	def create_github_instance(self, args):
		if args.token != None:
			#client = Github(args.token, per_page=args.per_page)
			#self.client = Github(args.user, args.pwd, per_page=args.per_page)
			self.client = Github(args.user, login_or_token=args.token, per_page=args.per_page)

		elif args.user != None and args.pwd != None:
			self.client = Github(args.user, args.pwd, per_page=args.per_page)

		else:
			self.client = Github()

		self.organization = self.client.get_organization(args.org)

	def create_directory(self,args):
		"""Create a directory for organization, for each repository in organization, for each event type in repository"""
		# 1... create a directory for this organization
		if not os.path.exists(args.org):
			os.mkdir(args.org)
		# 2... create directory for each repository of an organization
		# get repository

		repo_list = []
		if (args.repo == 'all'):
			repo_list = [repo.name for repo in self.organization.get_repos()]
		else:
			repo_list = [args.repo]

		# create direcory
		for j in repo_list:
			if not os.path.exists(args.org + "/" + j):
				os.mkdir(args.org + "/" + j)

		# 3.... create directory for each event type of repository
		for repo_name in repo_list:
			if args.event_type == 'commits':
				branch_name_list = self.get_branch(repo_name, args)
				for branch_name in branch_name_list:
					if not os.path.exists(args.org+"/"+repo_name+"/"+args.event_type + "/"+ branch_name + "_branch"):
						os.mkdir(args.org+"/"+repo_name+"/"+args.event_type + "/"+ branch_name + "_branch")
			else:
				if not os.path.exists(args.org+"/"+repo_name+"/"+args.event_type):
					os.mkdir(args.org+"/"+repo_name+"/"+args.event_type)

	def identify_event(self,args):
		"""identify the type of event given by the user"""
		if(args.event_type == 'issues'):
			self.collect_issues(args)
		if(args.event_type == 'issueComments'):
			self.collect_issues_comments(args)
		if (args.event_type == 'issueEvents'):
			self.collect_issues_events(args)
		if (args.event_type == 'commits'):
			self.collect_commits(args)
		if (args.event_type == 'events'):
			self.collect_events(args)
		if (args.event_type == 'commitComments'):
			self.collect_commitsComments(args)
		if(args.event_type == 'pullRequests'):
			self.collect_pullRequests(args)
		if (args.event_type == 'pullRequestComments'):
			self.collect_pullRequestComments(args)

	def get_repo(self,args):
		"""store all repository in a given organization as repo_list"""
		repo_list=[]
		if(args.repo == 'all'):
			repo_list = [repo.name for repo in self.organization.get_repos()]
		else:
			repo_list = [args.repo]

		return repo_list

	def get_branch(self, repo_name, args):
		branch_list = []
		if(args.branch == 'all'):
			branch_list = [branch.name for branch in self.organization.get_repo(repo_name).get_branches()]

		else:
			branch_list = [args.branch]

		return branch_list

	def collect_issues_events(self, args):
		repo_list = self.get_repo(args)
		print("\nRepositories:\n ", repo_list)
		try:
			for repo_name in repo_list:
				print("\n\t" + repo_name + " Repository")
				repo = self.organization.get_repo(repo_name)

				totalIssueEvents = repo.get_issues_events().totalCount
				print("total number of issue events in ", repo_name, "is", totalIssueEvents)

				totalPage = totalIssueEvents / args.per_page
				if totalPage is not int:
					totalPage = math.ceil(totalPage)
				print("total number of page with per_page ", self.per_page, " is ", totalPage)

				page = 0
				num_of_issues_events = 0

				while page < totalPage:
					issues_events_list = []
					print("\n\tpage: ", page)
					for event in repo.get_issues_events().get_page(page):
						event_dict = {}
						event_dict['actor'] = event.actor
						event_dict['commit_id'] = event.commit_id
						event_dict['created_at'] = event.created_at
						event_dict['event'] = event.event
						event_dict['id'] = event.id
						event_dict['issue'] = event.issue
						event_dict['url'] = event.url
						event_dict['node_id'] = event.node_id
						event_dict['commit_url'] = event.commit_url
						event_dict['label'] = event.label
						event_dict['assignee'] = event.assignee
						issues_events_list.append(event_dict)

						num_of_issues_events += 1
						print(num_of_issues_events)

					with open(args.org + "/" + repo_name + "/" + args.event_type + "/" + args.org + "-" + repo_name + "-" +
							  args.event_type + "-page-" + str(page) + ".json", 'w') as f:
							f.write(str(issues_events_list))

					print("page ", page, " added to file")
					self.sleeper()
					page += 1

			print("data successfully collected")
		except Exception as e:
			print("Problem Occured: ", e)

	def collect_pullRequests(self, args):
		# call a get_repo function
		repo_list = self.get_repo(args)
		print("\n\tRepositories\n", repo_list)
		try:
			for repo_name in repo_list:
				print("\n", repo_name, " Repository")
				repo = self.organization.get_repo(repo_name)
				branch_list = self.get_branch(repo_name, args)
				print("branches: ", branch_list)
				for branch in branch_list:

				    total_pulls = repo.get_pulls(state=args.state, sort='created', base=branch).totalCount
				    print("total number of pulls in ", repo_name, " of branch ", branch, " is ", total_pulls)

				    total_page = total_pulls / args.per_page
				    if total_page is not int:total_page = math.ceil(total_page)
				    print("The total number of page is ", total_page)
				    page = 0
				    num_of_pulls = 0
				    while page < total_page:
				    	print("\n\tPage: ", page)
				    	pull_list=[]
				    	for pull in repo.get_pulls(state=args.state, sort='created', base=branch).get_page(page):
				    		pull_dict = {}
				    		pull_dict['id'] = pull.id
				    		pull_dict['number'] = pull.number
				    		pull_dict['title'] = pull.title
				    		pull_dict['user'] = pull.user
				    		pull_dict['body'] = pull.body
				    		pull_dict['changed_files'] = pull.changed_files
				    		pull_dict['closed_at'] = pull.closed_at
				    		pull_dict['comments'] = pull.comments
				    		pull_dict['commits'] = pull.commits
				    		pull_dict['created_at'] = pull.created_at
				    		pull_dict['commits_url'] = pull.commits_url
				    		pull_dict['review_comments'] = pull.review_comments
				    		pull_dict['head'] = pull.head
				    		pull_dict['merged'] = pull.merged
				    		pull_dict['merged_at'] = pull.merged_at
				    		pull_dict['merged_by'] = pull.merged_by
				    		pull_dict['milestone'] = pull.milestone
				    		pull_dict['merge_commit_sha'] = pull.merge_commit_sha
				    		pull_dict['mergeable'] = pull.mergeable
				    		pull_dict['mergeable_state'] = pull.mergeable_state
				    		pull_dict['labels'] = pull.labels
				    		pull_dict['deletions'] = pull.deletions
				    		pull_dict['additions'] = pull.additions
				    		pull_dict['commit_url'] = pull.commits_url
				    		pull_dict['assignee'] = pull.assignee
				    		pull_dict['assignees'] = pull.assignees
				    		pull_dict['base'] = pull.base
				    		pull_list.append(pull_dict)
				    		num_of_pulls += 1
				    		print(num_of_pulls)
				    	with open(args.org + "/" + repo_name + "/" + args.event_type + "/" + args.org + "-" + repo_name + "-" +
								  branch + "-" + args.state + "-" + args.event_type + ".json", 'w') as f:
				    		f.write(str(pull_list))
				    	self.sleeper()
				    	page += 1
			print("data successfully collected")
		except Exception as e:
			print("Problem Occured: ", e)

	def collect_pullRequestComments(self,args):
	# call a get_repo function
		repo_list = self.get_repo(args)
		print("\n\tRepositories\n", repo_list)
		try:
			for repo_name in repo_list:
				print("\n", repo_name, "Repository")
				repo = self.organization.get_repo(repo_name)
				pull_list = []
				num_of_pulls = 0
				print("total pull Request comments in ", repo_name, " is: ", repo.get_pulls_comments().totalCount)
				for pull in repo.get_pulls_comments():
				    pull_dict = {}
				    pull_dict['id'] = pull.id
				    pull_dict['user'] = pull.user
				    pull_dict['body'] = pull.body
				    pull_dict['created_at'] = pull.created_at
				    pull_dict['commit_id'] = pull.commit_id
				    pull_dict['created_at'] = pull.created_at
				    pull_dict['diff_hunk'] = pull.diff_hunk
				    pull_dict['in_reply_to_id'] = pull.in_reply_to_id
				    pull_dict['original_commit_id'] = pull.original_commit_id
				    pull_dict['original_position'] = pull.original_position
				    pull_dict['path'] = pull.path
				    pull_dict['url'] = pull.url
				    pull_dict['updated_at'] = pull.updated_at
				    pull_dict['pull_request_url'] = pull.pull_request_url

				    pull_list.append(pull_dict)
				    num_of_pulls += 1
				    print(num_of_pulls)

				with open(args.org + "/" + repo_name + "/" + args.event_type + "/" + args.org + "-" + repo_name + "-" +
				          args.event_type + ".json", 'w') as f:
					f.write(str(pull_list))

			print("data successfully collected")
		except Exception as e:
			print("Problem Occured: ", e)

	def collect_commits(self, args):

		"""collect the data of the issue event in a given repository may be all repository or one repository"""
		# call a get_repo function
		repo_list = self.get_repo(args)
		print("\n\tRepositories:\n ", repo_list)
		try:
			for repo_name in repo_list:
				print(repo_name, "Repository")
				repo = self.organization.get_repo(repo_name)
				branch_list = self.get_branch(repo_name, args)
				print("branches: ", branch_list)

				for branch in branch_list:
					git_branch = self.organization.get_repo(repo_name).get_branch(branch)
					branch_commit = git_branch.commit

					total_commits = repo.get_commits(sha=branch_commit.sha).totalCount
					print("total number of commits in ",repo_name," of branch ", branch, " is: ", total_commits)

					# since there are 100 commits in a single page we can easily get the total number of page by dividing the total commits with 100
					total_page =  total_commits / args.per_page
					if total_page is not int:
						total_page = math.ceil(total_page)
					print("The total number of page is: " + str(total_page))

					#print(repo.get_commits().get_page(rel='last'))
					page = 0
					num_of_commits = 0
					while page < total_page:#just for testing but actually its till last page
						commit_list = []
						print("\n\tpage: ", page)

						for commit in repo.get_commits(sha=branch_commit.sha).get_page(page):
							commit_dict = {}
							commit_dict['author'] = commit.author
							commit_dict['sha'] = commit.sha
							commit_dict['files'] = commit.files
							commit_dict['stats'] = commit.stats
							commit_dict['commit'] = commit.commit
							commit_dict['committer'] = commit.committer
							commit_dict['comments_url'] = commit.comments_url
							commit_dict['html_url'] = commit.html_url
							commit_dict['parents'] = commit.parents
							commit_dict['url'] = commit.url
							commit_list.append(commit_dict)

							num_of_commits += 1
							print(num_of_commits)

						with open(args.org + "/" + repo_name+"/"+args.event_type+"/"+branch+"_branch/" +  args.org + "-" +
					 		repo_name + "-"+branch+"_branch-" + args.event_type + "-page-" + str(page) + ".json", 'w') as f:
							f.write(str(commit_list))

						print("page ", page, " added to file")
						self.sleeper()
						page += 1

			print("commit data successfully collected")
		except Exception as e:
			print("Problem Occured: ", e)

	def collect_commitsComments(self, args):
		# call a get_repo function
		repo_list = self.get_repo(args)
		print("\n\tRepositories\n", repo_list)
		try:
			for repo_name in repo_list:
				print("\n\t" + repo_name + " Repository")
				repo = self.organization.get_repo(repo_name)

				totalCommitComments = repo.get_comments().totalCount
				print("total number of commit comments in ", repo_name, " Repository is", totalCommitComments)

				totalPage = totalCommitComments / args.per_page
				if totalPage is not int:
					totalPage = math.ceil(totalPage)
				print("total number of page with per_page ", self.per_page, " is ", totalPage)

				page = 0
				num_of_comments = 0

				while page < totalPage:
					comment_list = []
					print("\n\tpage: ", page)
					for comment in repo.get_comments():
						comment_dict = {}
						comment_dict['id'] = comment.id
						comment_dict['user'] = comment.user
						comment_dict['commit_id'] = comment.commit_id
						comment_dict['body'] = comment.body
						comment_dict['url'] = comment.url
						comment_dict['position'] = comment.position
						comment_dict['path'] = comment.path
						comment_dict['line'] = comment.line
						comment_dict['created_at'] = comment.created_at
						comment_dict['updated_at'] = comment.updated_at
						comment_dict['html_url'] = comment.html_url

						comment_list.append(comment_dict)

						num_of_comments += 1
						print(num_of_comments)

					with open(args.org + "/" + repo_name + "/" + args.event_type + "/" + args.org + "-" + repo_name + "-" +
							   args.event_type + "-page-" + str(page) + ".json", 'w') as f:
						f.write(str(comment_list))

					print("page ", page, " added to file")
					self.sleeper()
					page += 1

			print("data successfully collected")
		except Exception as e:
			print("Problem Occured: ", e)

	def collect_events(self, args):

		# call a get_repo function
		repo_list = self.get_repo(args)
		print("\n\tRepositories\n", repo_list)
		try:
			for repo_name in repo_list:
				repo = self.organization.get_repo(repo_name)
				total_events = repo.get_events().totalCount
				print("total number of Events in ", repo_name , " Repository is ", total_events)

				total_page = total_events / args.per_page
				if total_page is not int:
					total_page = math.ceil(total_page)
				print("The total number of page is: " + str(total_page))

				page = 0
				num_of_events = 0
				while page < total_page:
					event_list = []
					print("\n\tpage: ", page)
					for event in repo.get_events().get_page(page):
						event_dict = {}
						event_dict['actor'] = event.actor
						event_dict['id'] = event.id
						event_dict['payload'] = event.id
						event_dict['created_at'] = event.created_at
						event_dict['org'] = event.org
						event_dict['public'] = event.public
						event_dict['repo'] = event.repo
						event_dict['type'] = event.type
						event_list.append(event_dict)

						num_of_events += 1
						print(num_of_events)

					with open(args.org + "/" + repo_name + "/" + args.event_type + "/" + args.org + "-" + repo_name + "-" +
							  args.event_type + "-page-" + str(page) +  ".json", 'w') as f:
						f.write(str(event_list))

					print("page ", page, " added to file")
					self.sleeper()
					page += 1

			print("data successfully collected")
		except Exception as e:
			print("Problem Occured: ", e)

	def collect_issues_comments(self, args):
		repo_list = self.get_repo(args)
		print("\nRepositories:\n ", repo_list)
		try:
			for repo_name in repo_list:
				print("\n\t" + repo_name + " Repository")
				repo = self.organization.get_repo(repo_name)

				totalIssueComments = repo.get_issues_comments().totalCount
				print("total number of issue comments in ", repo_name, " Repository is", totalIssueComments)

				totalPage = totalIssueComments / args.per_page
				if totalPage is not int:
					totalPage = math.ceil(totalPage)
				print("total number of page with per_page ", self.per_page, " is ", totalPage)

				page = 0
				num_of_issue_comments = 0

				while page < totalPage:
					issue_comment_list = []
					print("\n\tpage: ", page)
					for comment in repo.get_issues_comments().get_page(page):
						comment_dict = {}
						comment_dict['id'] = comment.id
						comment_dict['user'] = comment.user
						comment_dict['body'] = comment.body
						comment_dict['issue_url'] = comment.issue_url
						comment_dict['created_at'] = comment.created_at
						comment_dict['updated_at'] = comment.updated_at
						comment_dict['url'] = comment.url
						comment_dict['html_url'] = comment.html_url

						issue_comment_list.append(comment_dict)

						num_of_issue_comments += 1
						print(num_of_issue_comments)

					with open(args.org + "/" + repo_name + "/" + args.event_type + "/" + args.org + "-" + repo_name + "-" +
							  args.event_type + "-page-" + str(page) + ".json", 'w') as f:
						f.write(str(issue_comment_list))

					print("page ", page, " added to file")
					self.sleeper()
					page += 1

			print("data successfully collected")
		except Exception as e:
			print("Problem occured:", e)

	def collect_issues(self, args):
		"""collect the data of the issue event in a given repository may be all repository or one repository"""
		# call a get_repo function
		repo_list = self.get_repo(args)
		print("\n\tRepositories\n", repo_list)
		try:
			for repo_name in repo_list:
				print("\n\t" + repo_name + " Repository")
				repo = self.organization.get_repo(repo_name)

				totalIssues = repo.get_issues(state=args.state).totalCount
				print("total number of " + args.state + " issues in " + repo_name + " Repository is: " , totalIssues)

				totalPage = totalIssues / args.per_page
				if totalPage is not int:
					totalPage = math.ceil(totalPage)
				print("total number of page with per_page ", self.per_page, " is ", totalPage)

				page = 0
				num_of_issue = 0

				while page < totalPage:
					issue_comment_list = []
					print("\n\tpage: ", page)
					issue_list = []
					for issue in repo.get_issues(state=args.state).get_page(page):
						issue_dict = {}
						issue_dict['number'] = issue.number
						issue_dict['id'] = issue.id
						issue_dict['user'] = issue.user
						issue_dict['title'] = issue.title
						issue_dict['body'] = issue.body
						issue_dict['url'] = issue.url
						issue_dict['milestone'] = issue.milestone
						issue_dict['labels'] = issue.labels
						issue_dict['labels_url'] = issue.labels_url
						issue_dict['created_at'] = issue.created_at
						issue_dict['updated_at'] = issue.updated_at
						issue_dict['closed_at'] = issue.closed_at
						issue_dict['closed_by'] = issue.closed_by
						issue_dict['pull_request'] = issue.pull_request
						issue_dict['state'] = issue.state
						issue_dict['events_url'] = issue.events_url
						issue_dict['comments'] = issue.comments
						issue_dict['number_of_comments'] = issue.comments
						issue_dict['comments_url'] = issue.comments_url
						issue_dict['assignee'] = issue.assignee
						issue_dict['assignees'] = issue.assignees
						issue_dict['html_url'] = issue.html_url

						issue_list.append(issue_dict)

						num_of_issue += 1
						print(num_of_issue)

					with open(args.org + "/" + repo_name + "/" + args.event_type + "/" + args.org + "-" + repo_name + "-" +
							  args.state + "-" + args.event_type + "-page-" + str(page) + ".json", 'w') as f:
						f.write(str(issue_list))

					print("page ", page, " added to file")
					self.sleeper()
					page += 1

			print("data successfully collected")
		except Exception as e:
			print("Problem Occured: ", e)

	def sleeper(self):
        # add sleeping time after requesting each page
		sleep_time = self.sleep_time_length + self.sleep_time_range * random.random()
		print("\n\tSleeping for", sleep_time, "seconds.")
		time.sleep(sleep_time)

	def main(self):

		# get the arguments from the terminal
		args = self.get_arguments()

		# validate the given arguments
		self.validate_arguments(args)

		# then create a github instance
		self.create_github_instance(args)

		# then call a create directory function
		self.create_directory(args)

		# identify the event type needed by the user and then go forward
		self.identify_event(args)

# Initialize the class
collect = Collect()
collect.main()


# import os
from PIL import Image
# from lib.utils import combine_test_results
import lib.myMetrics as metrics
import lib.myAverage as average
import lib.filters as filter

# folder = "assets/test"
# essa = 'assets/test/essa.png'
# soh = 'assets/test/sohaib.jpg'

def process_every_pixel(img, func, ave_green, ave_brightness, factor):
    width, height = img.size
    modified_image = Image.new("RGB", (width, height))
    # Iterate over each pixel in the original image
    for y in range(height):
        for x in range(width):
            # Get the RGB values of the current pixel
            r, g, b = img.getpixel((x, y))
            # Calculate the new value for the pixel
            new_pixel_value = func(
                x, y, width, height, r, g, b, img, ave_green, ave_brightness, factor
            )
            # place in image
            modified_image.putpixel((x, y), new_pixel_value)
    return modified_image

def model1(img):
    img = filter.contrast(img, 2)
    # while (metrics.blur(img) < 40):
    #     img = filter.sharpen(img, 1)
    ave_green = average.green(img)
    ave_brightness = average.brightness(img)
    img = process_every_pixel(img, filter.green, ave_green, ave_brightness, 1.25)
    # remove bright
    ave_brightness = average.brightness(img)
    img = process_every_pixel(
        img, filter.for_brightness, ave_green, ave_brightness, 1.01
    )
    ave_green = average.green(img)
    ave_brightness = average.brightness(img)
    img = process_every_pixel(img, filter.green, ave_green, ave_brightness, 1.4)
    # img = filter.sharpen(img, 1)
    ave_brightness = average.brightness(img)
    img = process_every_pixel(
        img, filter.for_brightness, ave_green, ave_brightness, 1.01
    )
    return img

def process_img(img_path):
    # Open the image
    img = Image.open(img_path).convert("RGB")
    img= model1(img)
    # sharpen with auto stop on blur correction
        # while (metrics.blur(img) < 100):
        #     img = filter.sharpen(img, 1)
        # # remove bright
        # ave_brightness = average.brightness(img)
        # img = process_every_pixel(
        #     img, filter.for_brightness, 0, ave_brightness, 1.005
        # )
        # # sharpen with auto stop on blur correction
        # # while (metrics.blur(img) < 150):
        # img = filter.sharpen(img, 1)
        # # remove bright
        # ave_brightness = average.brightness(img)
        # img = process_every_pixel(
        #     img, filter.for_brightness, 0, ave_brightness, 1.005
        # )
        # # remove green
        # ave_green = average.green(img)
        # img = process_every_pixel(img, filter.green, ave_green, 0, 1)
        # img = filter.sharpen(img, 1)
    # ave_brightness = average.brightness(img)
    # img = process_every_pixel(
    #     img, filter.for_brightness, 0, ave_brightness, 1.005
    # )
    # img = process_every_pixel(
    #     img, filter.for_brightness, 0, ave_brightness, 1.005
    # )
    # ave_green = average.green(img)
    # img = process_every_pixel(img, filter.green, ave_green, 0, 1.15)

    return img

# def process():
#     return model1(img)
#     # for file_name in os.listdir(folder):
    #     file_path = os.path.join(folder, file_name)
    #     processed = process_img(file_path)
    #     processed.save(folder+"/../results/"+file_name)
    # combine_test_results()

if __name__ == "__main__":
    main()


class PAM:
    '''
    Portable Arbitrary Map (PAM) file parser

    @variable  width:int     The width of image, in pixels
    @variable  height:int    The height of the image, in pixels
    @variable  depth:int     The number of channals, 1 for black and white or
                             greyscale, 3 for colour, plus 1 for alpha
    @variable  maxval:int    The maximum value a channel can have, the white point threshold
    @variable  tupltype:int  The channel configuration
    @variable  offset:int    The index of the first byte in the file for the first byte of the first pixel
    '''

    BLACKANDWHITE = 0
    '''
    :int  Value for `tupletype` when the image is black and white
          The bytes are: whitenames (note that it is not blackness as with P1 and P4)
          (note that it is bytes and not bits as with P4)
    '''

    GRAYSCALE = 1
    '''
    :int  Value for `tupletype` when the image is grey
          The words are: whiteness
    '''

    RGB = 2
    '''
    :int  Value for `tupletype` when the image is coloured
          The words are: redness, greenness, blueness
    '''

    BLACKANDWHITE_ALPHA = 3
    '''
    :int  Value for `tupletype` when the image is black and white with an alpha channel
          The bytes are: blackness (note that it is not blackness as with P1 and P4), transparency
    '''

    GRAYSCALE_ALPHA = 4
    '''
    :int  Value for `tupletype` when the image is grey with an alpha channel
          The words are: whiteness, transparency
    '''

    RGB_ALPHA = 5
    '''
    :int  Value for `tupletype` when the image is coloured with an alpha channel
          The words are: redness, greenness, blueness, transparency
    '''

    def __init__(self):
        '''
        Constructor
        '''
        self.width    = None
        self.height   = None
        self.depth    = None
        self.maxval   = None
        self.tupltype = None

    @staticmethod
    def parse(data : bytes) -> PAM:
        '''
        Parse an PAM-file

        @param   data:bytes  The file content
        @return  :PAM?       Metadata for the image, `None` on error (corrupt file)
        '''
        meta = PAM()
        meta.offset = 0
        ok = False

        TUPLTYPES = { 'BLACKANDWHITE'       : PAM.BLACKANDWHITE
                    , 'GRAYSCALE'           : PAM.GRAYSCALE
                    , 'RGB'                 : PAM.RGB
                    , 'BLACKANDWHITE_ALPHA' : PAM.BLACKANDWHITE_ALPHA
                    , 'GRAYSCALE_ALPHA'     : PAM.GRAYSCALE_ALPHA
                    , 'RGB_ALPHA'           : PAM.RGB_ALPHA
                    }

        buf = ''
        stage = 0
        try:
            for c in data:
                c = chr(c)
                meta.offset += 1
                if c == '\n':
                    buf = buf.strip()
                    if buf == 'P7':
                        if not stage == 0:
                            break
                        stage = 1
                    elif buf == 'ENDHDR':
                        ok = True
                        break
                    elif buf.startswith('WIDTH'):
                        if meta.width is not None: break
                        meta.width = int(buf[5:].strip())
                    elif buf.startswith('HEIGHT'):
                        if meta.height is not None: break
                        meta.height = int(buf[6:].strip())
                    elif buf.startswith('DEPTH'):
                        if meta.depth is not None: break
                        meta.depth = int(buf[5:].strip())
                    elif buf.startswith('MAXVAL'):
                        if meta.maxval is not None: break
                        meta.maxval = int(buf[6:].strip())
                    elif buf.startswith('TUPLTYPE'):
                        if meta.tupltype is not None: break
                        meta.tupltype = TUPLTYPES[buf[8:].strip()]
                else:
                    buf += c
        except:
            del meta
            return None

        if not ok:
            del meta
            return None

        if any([x is None for x in (meta.width, meta.height, meta.depth, meta.maxval, meta.tupltype)]):
            ok = False
        elif (meta.width < 1) or (meta.height < 1):     ok = False
        elif meta.tupltype == PAM.BLACKANDWHITE:        ok = (depth == 1) and (1 <= maxval <= 1)
        elif meta.tupltype == PAM.GRAYSCALE:            ok = (depth == 1) and (1 <= maxval <= 65535)
        elif meta.tupltype == PAM.RGB:                  ok = (depth == 3) and (1 <= maxval <= 65535)
        elif meta.tupltype == PAM.BLACKANDWHITE_ALPHA:  ok = (depth == 2) and (1 <= maxval <= 1)
        elif meta.tupltype == PAM.GRAYSCALE_ALPHA:      ok = (depth == 2) and (1 <= maxval <= 65535)
        elif meta.tupltype == PAM.RGB_ALPHA:            ok = (depth == 4) and (1 <= maxval <= 65535)
        if not ok:
            del meta
            return None
        return meta


import datetime
import PIL
import pytest

from jmm.utilities.crawler_group import Router
from jmm.utilities.crawler_group import Rule
from jmm.utilities.crawler_group import CrawlerGroup

@pytest.mark.parametrize(
    'number, metadata', [
        (
            'STAR-325',
            {
                'keywords': ['STAR-325', 'SODクリエイト', '美人潜入捜査官', 'DVD多士爐', '女檢察官', '潮吹', '偶像藝人', '單體作品', '中出'],
                'title': 'STAR-325 美人潜入捜査官 羽田あい',
                'release_date': datetime.date(2012, 1, 8),
                'runtime': 120,
                'number': 'STAR-325',
                'director': '本田教仁',
                'series': '美人潜入捜査官',
                'studio': 'SODクリエイト',
                'stars': [{'avatar_url': 'https://www.javbus.com/pics/actress/6xv_a.jpg', 'name': '羽田あい'}],
                'outline': '潛入黑暗組織的「羽田亞衣」，雖然登場很瀟灑，不過沒多久就被敵人抓住了！被鎖鍊綁住連續調教，潮吹就像爆炸班狂噴！不斷反覆的凌虐，反而讓她沈溺在當中…。'
            }
        ),
        (
            'FC2-PPV-2698221',
            {
                'keywords': ['巨乳'],
                'title': 'FC2-PPV-2698221 独占販売1本のおまけ動画あり【無修正ｘ個人撮影】巨乳タレ乳、ビラビラ乳首マンコに経産婦のだらしない体が激エロ過ぎる美人妻再び！流出してしまった動画をネタにホテルに連れ込んで、巨乳もみほぐし♪',
                'release_date': datetime.date(2022, 3, 4),
                'runtime': 69,
                'number': 'FC2-PPV-2698221',
                'director': None,
                'series': 'Kerberos',
                'studio': 'FC2-PPV',
                'stars': [],
                'outline': None
            }
        ),
    ]
)
def test_router(javbus, javdb, airav, avsox, logger, number, metadata):
    rule_1 = Rule(r'\w+-\d+', CrawlerGroup([javbus, javdb, airav], logger=logger))
    rule_2 = Rule(r'.*', CrawlerGroup([avsox, javdb], logger=logger))
    router = Router([rule_1, rule_2])

    _metadata = router.get_metadata(number)
    assert isinstance(_metadata.fanart, PIL.JpegImagePlugin.JpegImageFile)
    assert {
        'keywords': _metadata.keywords,
        'title': _metadata.title,
        'release_date': _metadata.release_date,
        'runtime': _metadata.runtime,
        'number': _metadata.number,
        'director': _metadata.director,
        'series': _metadata.series,
        'studio': _metadata.studio,
        'stars': [{'name': star.name, 'avatar_url': star.avatar_url} for star in _metadata.stars],
        'outline': _metadata.outline,
    } == metadata

@pytest.mark.parametrize('number', ['star-250', 'FC-001'])
def test_router_with_wrong_patterns(javbus, javdb, airav, avsox, logger, number):
    rule_1 = Rule(r'\d+', CrawlerGroup([javbus, javdb, airav], logger=logger))
    rule_2 = Rule(r'\d\d', CrawlerGroup([avsox, javdb], logger=logger))
    router = Router([rule_1, rule_2])
    assert router.get_metadata(number) is None


import os
import platform
import time
from datetime import datetime
from typing import List, Optional, Tuple, Union
from urllib.parse import urljoin
from urllib.request import urlretrieve

import numpy as np
import tensorflow as tf
from flwr.common import ndarrays_to_parameters
from flwr.common.typing import NDArray
from flwr.server.strategy import FedAvg
from flwr.simulation import start_simulation
from keras.utils.np_utils import to_categorical
from loguru import logger
from numpy.random import Generator
from tensorflow.keras import Input, Sequential, layers, metrics
from tensorflow.keras.datasets import mnist
from tensorflow.keras.optimizers import Adam

from inflorescence.dataset import Dataset
from inflorescence.dataset.common import XY
from inflorescence.experiment import Experiment
from inflorescence.metrics import aggregate_binary_cm, weighted_average
from inflorescence.model import Model
from inflorescence.strategy import CFL, FLHC, IFCA, WeCFL

gpus = tf.config.list_physical_devices("GPU")
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

class KerasCNN(Model):
    def __init__(self, learning_rate: float = 0.001, seed=None, **kwargs):
        """"""
        gpus = tf.config.list_physical_devices("GPU")
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
        if seed:
            tf.random.set_seed(seed)
        dropout_rate = 0.5
        optimizer = Adam(learning_rate=learning_rate)
        model = Sequential()
        model.add(Input(shape=(28, 28, 1)))
        model.add(layers.Conv2D(32, kernel_size=(3, 3), activation="relu"))
        model.add(layers.MaxPool2D(pool_size=(2, 2)))
        model.add(layers.Conv2D(64, kernel_size=(3, 3), activation="relu"))
        model.add(layers.MaxPool2D(pool_size=(2, 2)))
        model.add(layers.Flatten())
        model.add(layers.Dropout(dropout_rate))
        model.add(layers.Dense(units=10, activation="softmax"))

        model.compile(
            optimizer=optimizer,
            loss="sparse_categorical_crossentropy",
            metrics=[
                "sparse_categorical_accuracy",
            ],
        )
        self.model = model

    def get_weights(self):
        return self.model.get_weights()

    def set_weights(self, weights):
        self.model.set_weights(weights)

    def fit(self, X_train, y_train, epochs, batch_size, **kwargs):
        """"""
        y = y_train[:, 0]
        history = self.model.fit(X_train, y, batch_size=batch_size, epochs=epochs, verbose=2)
        return history.history["loss"][-1]

    def evaluate(self, X_test, y_test, batch_size=32):
        """"""
        test_dataset = tf.data.Dataset.from_tensor_slices((X_test, y_test))
        test_dataset = test_dataset.shuffle(buffer_size=1024).batch(batch_size)
        loss_fn = tf.keras.losses.SparseCategoricalCrossentropy()
        test_acc_metric = tf.keras.metrics.SparseCategoricalAccuracy()
        start_time = time.time()
        epoch_loss = 0.0
        steps = 0
        num_classes = self.model.layers[-1].output_shape[-1]
        cm = tf.zeros((num_classes, num_classes), dtype=tf.dtypes.int32)
        group_cms = {}
        group_correct = {}
        group_n = {}
        for x_batch, labels in test_dataset:
            steps += 1
            y = labels[:, 0]
            g = labels[:, 1]
            with tf.GradientTape() as tape:
                preds = self.model(x_batch, training=False)
                loss_value = loss_fn(y_true=y, y_pred=preds)
                epoch_loss += loss_value
            grads = tape.gradient(loss_value, self.model.trainable_weights)
            batch_cm = tf.math.confusion_matrix(
                y, tf.argmax(preds, axis=1), num_classes=num_classes
            )
            groups = tf.sort(tf.unique(g)[0])
            for group in groups:
                group_num = group.numpy()
                batch_group_cm = tf.math.confusion_matrix(
                    y[g == group], tf.argmax(preds[g == group], axis=1), num_classes=num_classes
                )
                group_cms.setdefault(
                    group_num, tf.zeros((num_classes, num_classes), dtype=tf.dtypes.int32)
                )
                group_cms[group_num] += batch_group_cm
                # TODO: log per-group accuracy here
                diag_sum = tf.reduce_sum(tf.linalg.tensor_diag_part(batch_group_cm))
                cm_sum = tf.reduce_sum(batch_group_cm)
                group_correct.setdefault(group_num, 0)
                group_n.setdefault(group_num, 0)
                group_correct[group_num] += diag_sum
                group_n[group_num] += cm_sum
            cm += batch_cm
            self.model.optimizer.apply_gradients(zip(grads, self.model.trainable_weights))
            test_acc_metric.update_state(y_true=y, y_pred=preds)
        valid_acc = test_acc_metric.result()
        print("Validation acc: %.4f" % (float(valid_acc),))
        test_acc_metric.reset_states()
        # process confusion matrix
        cm_dict = {}
        for i in range(num_classes):
            for j in range(num_classes):
                cm_dict[f"valid_cm_{i:02}_{j:02}"] = cm[i][j].numpy().astype(float)
                for k, gcm in group_cms.items():
                    cm_dict[f"valid_cm_{i:02}_{j:02}_group_{k:02}"] = (
                        gcm[i][j].numpy().astype(float)
                    )
        epoch_loss = epoch_loss / steps
        print("Time taken: %.2fs" % (time.time() - start_time))
        metrics_dict = {"valid_acc": valid_acc.numpy().astype(float)}
        for g, n in group_n.items():
            correct = group_correct[g]
            metrics_dict[f"valid_acc_group_{g:02}"] = (correct / n).numpy().astype(float)
        metrics_dict.update(cm_dict)
        return epoch_loss.numpy().astype(float), metrics_dict

class RotatedMNIST(Dataset):
    def __init__(self, root_dir=None):
        self.root_dir = root_dir
        self.download()

    def download(self):
        (x_train, y_train), (x_test, y_test) = mnist.load_data(path=self.root_dir)
        x_train = x_train.astype(np.float32) / 255
        x_test = x_test.astype(np.float32) / 255
        x_train = np.expand_dims(x_train, -1)
        x_test = np.expand_dims(x_test, -1)
        self.x = np.vstack([x_train, x_test])
        self.y = np.hstack([y_train, y_test])
        self.shuffle()
        self.rotate()

    def shuffle(self, seed=None):
        idx = np.arange(70000)
        np.random.default_rng(seed=seed).shuffle(idx)
        self.x = self.x[idx]
        self.y = self.y[idx]

    def split_non_iid(
        self,
        concentration: float,
        num_clients: int,
        group_var: str = "rotation",
        test_size: float = 0.2,
        seed=None,
    ) -> Tuple[List[Tuple[XY, XY]], NDArray]:
        return super().split_non_iid(concentration, num_clients, group_var, test_size, seed)

    def rotate(self):
        """Rotate the clients' samples and add rotation # to y labels"""
        for i, ex in enumerate(self.x):
            # 3/8 0deg, 1/4 90deg, 1/8 180 deg, 1/4 270deg
            self.x[i] = np.rot90(ex, ((((i - 2) % 8) > 0) * ((i % 4))))
        idxs = np.arange(70000)
        rot_labels = (((idxs - 2) % 8) > 0) * ((idxs % 4))
        self.y = np.stack([self.y, rot_labels], axis=1)

    def __getitem__(self, item):
        pass

    def get_xy(self, group_var=None):
        return self.x, self.y

def run_ifca(
    k_clusters,
    num_clients,
    num_rounds,
    epochs,
    batch_size,
    concentration,
    seed,
    num_cpus,
    num_gpus,
    common_params,
):
    initial_weights = []
    with tf.device("/device:cpu:0"):
        for i in range(k_clusters):
            initial_model = KerasCNN(seed=seed + i)
            initial_weights.append(initial_model.get_weights())
            del initial_model
    initial_weights_flat = [_ for _ in initial_weights for _ in _]
    strategy = IFCA(
        k_clusters=k_clusters,
        initial_parameters=ndarrays_to_parameters(initial_weights_flat),
        **common_params,
    )
    dt = datetime.now()
    logger.info("Starting experiments for RotatedMNIST with IFCA(k_clusters={})", k_clusters)
    experiment = Experiment(
        strategies=[strategy],
        model_fn=KerasCNN,
        dataset=rmnist,
        group_var="rotation",
        concentrations=[concentration],
    )

    experiment.run(
        num_rounds=num_rounds,
        num_clients=num_clients,
        epochs=epochs,
        batch_size=batch_size,
        seed=seed,
        num_cpus=num_cpus,
        num_gpus=num_gpus,
    )
    dtend = datetime.now()
    span = dtend - dt
    dts = dt.strftime("%Y-%m-%d_%H%M%S")
    fname = (
        f"results/rotatedmnistimb_ifca_seed_{seed}_{concentration:.02f}_{k_clusters}_{dts}.parquet"
    )
    logger.info("Finished RotatedMNIST IFCA experiment in {}. Saving results to {}", span, fname)
    df = experiment.to_pandas()
    os.makedirs("results", exist_ok=True)
    df.to_parquet(fname)

def run_wecfl(
    k_clusters,
    num_clients,
    num_rounds,
    epochs,
    batch_size,
    concentration,
    seed,
    num_cpus,
    num_gpus,
    common_params,
):
    initial_weights = []
    with tf.device("/device:cpu:0"):
        for i in range(k_clusters):
            initial_model = KerasCNN(seed=seed + i)
            initial_weights.append(initial_model.get_weights())
            del initial_model
    strategy = WeCFL(
        k_clusters=k_clusters,
        initial_parameters=ndarrays_to_parameters(initial_weights[0]),
        **common_params,
    )
    dt = datetime.now()
    logger.info("Starting experiments for RotatedMNIST with WeCFL(k_clusters={})", k_clusters)
    experiment = Experiment(
        strategies=[strategy],
        model_fn=KerasCNN,
        dataset=rmnist,
        group_var="rotation",
        concentrations=[concentration],
    )

    experiment.run(
        num_rounds=num_rounds,
        num_clients=num_clients,
        epochs=epochs,
        batch_size=batch_size,
        seed=seed,
        num_cpus=num_cpus,
        num_gpus=num_gpus,
    )
    dtend = datetime.now()
    span = dtend - dt
    dts = dt.strftime("%Y-%m-%d_%H%M%S")
    fname = (
        f"results/rotatedmnistimb_wecfl_seed_{seed}_{concentration:.02f}_{k_clusters}_{dts}.parquet"
    )
    logger.info("Finished RotatedMNIST WeCFL experiment in {}. Saving results to {}", span, fname)
    df = experiment.to_pandas()
    os.makedirs("results", exist_ok=True)
    df.to_parquet(fname)

def run_flhck(
    k_clusters,
    num_clients,
    num_rounds,
    epochs,
    batch_size,
    concentration,
    seed,
    num_cpus,
    num_gpus,
    common_params,
):
    with tf.device("/device:cpu:0"):
        initial_model = KerasCNN(seed=seed)
        initial_weights = initial_model.get_weights()
        del initial_model
    strategy = FLHC(
        k_clusters=k_clusters,
        metric="euclidean",
        linkage="complete",
        cluster_after_round=5,
        initial_parameters=ndarrays_to_parameters(initial_weights),
        **common_params,
    )
    dt = datetime.now()
    logger.info("Starting experiments for RotatedMNIST with FLHC(k_clusters={})", k_clusters)
    experiment = Experiment(
        strategy=strategy,
        model_fn=KerasCNN,
        dataset=rmnist,
        group_var="rotation",
        concentrations=[concentration],
    )

    experiment.run(
        num_rounds=num_rounds,
        num_clients=num_clients,
        epochs=epochs,
        batch_size=batch_size,
        seed=seed,
        num_cpus=num_cpus,
        num_gpus=num_gpus,
    )
    dtend = datetime.now()
    span = dtend - dt
    dts = dt.strftime("%Y-%m-%d_%H%M%S")
    fname = f"results/rotatedmnistimb_flhck_seed_{seed}_{concentration:.02f}_{k_clusters:02d}_{dts}.parquet"
    logger.info("Finished RotatedMNIST FLHC experiment in {}. Saving results to {}", span, fname)
    df = experiment.to_pandas()
    os.makedirs("results", exist_ok=True)
    df.to_parquet(fname)

def run_fedavg(
    num_clients,
    num_rounds,
    epochs,
    batch_size,
    concentration,
    seed,
    num_cpus,
    num_gpus,
    common_params,
):
    with tf.device("/device:cpu:0"):
        initial_model = KerasCNN(seed=seed)
        initial_weights = initial_model.get_weights()
        del initial_model
    strategy = FedAvg(
        initial_parameters=ndarrays_to_parameters(initial_weights),
        **common_params,
    )
    dt = datetime.now()
    logger.info("Starting experiments for RotatedMNIST with FedAvg")
    experiment = Experiment(
        strategies=[strategy],
        model_fn=KerasCNN,
        dataset=rmnist,
        group_var="rotation",
        concentrations=[concentration],
    )

    experiment.run(
        num_rounds=num_rounds,
        num_clients=num_clients,
        epochs=epochs,
        batch_size=batch_size,
        seed=seed,
        num_cpus=num_cpus,
        num_gpus=num_gpus,
    )
    dtend = datetime.now()
    span = dtend - dt
    dts = dt.strftime("%Y-%m-%d_%H%M%S")
    fname = f"results/rotatedmnistimb_fedavg_seed_{seed}_{concentration:.02f}_{dts}.parquet"
    logger.info("Finished RotatedMNIST fedavg experiment in {}. Saving results to {}", span, fname)
    df = experiment.to_pandas()
    os.makedirs("results", exist_ok=True)
    df.to_parquet(fname)

if __name__ == "__main__":
    NUM_ROUNDS = 10
    BATCH_SIZE = 64
    EPOCHS = 10
    NUM_CLIENTS = 10
    rmnist = RotatedMNIST()
    SEEDS = [43, 59, 67, 79, 97]
    CONCENTRATIONS = np.logspace(-2, 3, 6)
    K_CLUSTERS = 5
    common_params = dict(
        accept_failures=False,
        fraction_fit=1.0,  # Proportion of clients to sample in each training round
        fraction_evaluate=1.0,  # Proportion of clients to calculate accuracy on after each round
        min_fit_clients=NUM_CLIENTS,  # Minimum number of clients to train on in each round
        min_evaluate_clients=NUM_CLIENTS,  # Minimum number of clients to evaluate accuracy on after each round
        min_available_clients=NUM_CLIENTS,  # Minimum number of available clients needed to start a round
        evaluate_metrics_aggregation_fn=aggregate_binary_cm,  # <-- pass the metric aggregation function
    )
    os.makedirs("logs", exist_ok=True)
    logger.add(
        "logs/rotatedmnistimb.log", format="{time:YYYY-MM-DDTHH:mm:ss} | {level} | {message}"
    )

    for seed in SEEDS:
        for concentration in CONCENTRATIONS:
            logger.info(
                "Starting experiments for RotatedMNIST, seed = {}, concentration = {}",
                seed,
                concentration,
            )
            run_fedavg(
                num_clients=NUM_CLIENTS,
                num_rounds=NUM_ROUNDS,
                batch_size=BATCH_SIZE,
                epochs=EPOCHS,
                seed=seed,
                concentration=concentration,
                num_cpus=8,
                num_gpus=0,
                common_params=common_params,
            )
            run_wecfl(
                k_clusters=K_CLUSTERS,
                num_clients=NUM_CLIENTS,
                num_rounds=NUM_ROUNDS,
                batch_size=BATCH_SIZE,
                epochs=EPOCHS,
                seed=seed,
                concentration=concentration,
                num_cpus=8,
                num_gpus=0,
                common_params=common_params,
            )
            run_ifca(
                k_clusters=K_CLUSTERS,
                num_clients=NUM_CLIENTS,
                num_rounds=NUM_ROUNDS,
                batch_size=BATCH_SIZE,
                epochs=EPOCHS,
                seed=seed,
                concentration=concentration,
                num_cpus=8,
                num_gpus=0,
                common_params=common_params,
            )

            run_flhck(
                k_clusters=K_CLUSTERS,
                num_clients=NUM_CLIENTS,
                num_rounds=NUM_ROUNDS,
                batch_size=BATCH_SIZE,
                epochs=EPOCHS,
                seed=seed,
                concentration=concentration,
                num_cpus=8,
                num_gpus=0,
                common_params=common_params,
            )


import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
#import codecademylib3_seaborn
import glob

files = glob.glob("states*.csv")
us_census = []
for filename in files:
  data = pd.read_csv(filename)
  us_census.append(data)

df = pd.concat(us_census)

print(df.dtypes)
print(df.head())
print(len(df))

# for i in range(10):
#     f = open("states" + str(i) + ".csv", "w")
#     f.write("")
#     f.close()

df['Income'] = pd.to_numeric(df.Income.replace('[\$]', '', regex=True))
print(df.dtypes)
print(df.head())
df_tmp = df['GenderPop'].str.split('_')
df['Men'] = pd.to_numeric(df_tmp.str.get(0).replace('[M]', '', regex=True))
df['Women'] = pd.to_numeric(df_tmp.str.get(1).replace('[F]', '', regex=True))
print(df.dtypes)
print(df.head())
plt.scatter(df['Women'], df['Income'])
plt.show()
#print(df.Women)
df.Women = df.Women.fillna(value=(df.TotalPop - df.Men))
#print(df.Women)
duplicates = df.duplicated()
print(duplicates.value_counts())
df_tmp = df.drop_duplicates()
df = df_tmp
plt.scatter(df['Women'], df['Income'])
plt.show()
print(df.columns)
df['Hispanic'] = pd.to_numeric(df.Hispanic.replace('[%]', '', regex=True))
df['White'] = pd.to_numeric(df.White.replace('[%]', '', regex=True))
df['Black'] = pd.to_numeric(df.Black.replace('[%]', '', regex=True))
df['Native'] = pd.to_numeric(df.Native.replace('[%]', '', regex=True))
df['Asian'] = pd.to_numeric(df.Asian.replace('[%]', '', regex=True))
df['Pacific'] = pd.to_numeric(df.Pacific.replace('[%]', '', regex=True))
print(df.head())
nans = df.isnull()
print(nans)
Pnan = df.Pacific.isnull()
print(Pnan.value_counts())
nan_value = df.Pacific.mean()
df.Pacific = df.Pacific.fillna(value=nan_value)
Pnan = df.Pacific.isnull()
print(Pnan.value_counts())

Python3/Cleaning US Census Data/CleaningUSCensusData.py


"""
HUMAnN setup

To run: python setup.py install

"""

import sys

# required python versions (2.7+ or 3.0+)
required_python_version_major = [2,3]
required_python_version_minor = [7,0]

# check for either of the required versions
pass_check=False
try:
    for major, minor in zip(required_python_version_major, required_python_version_minor):
        if (sys.version_info[0] == major and sys.version_info[1] >= minor):
            pass_check=True
except (AttributeError,IndexError):
    sys.exit("CRITICAL ERROR: The python version found (version 1) " +
        "does not match the version required (version "+
        str(required_python_version_major)+"."+
        str(required_python_version_minor)+"+)")

if not pass_check:
    sys.exit("CRITICAL ERROR: The python version found (version "+
        str(sys.version_info[0])+"."+str(sys.version_info[1])+") "+
        "does not match the version required (version "+
        str(required_python_version_major)+"."+
        str(required_python_version_minor)+"+)")

try:
    import setuptools
except ImportError:
    sys.exit("Please install setuptools.")

# check setuptools version
required_setuptools_version_major = 1
try:
    setuptools_version = setuptools.__version__
    setuptools_version_major = int(setuptools_version.split(".")[0])
    if setuptools_version_major < required_setuptools_version_major:
        sys.exit("CRITICAL ERROR: The setuptools version found (version "+
                 setuptools_version+") does not match the version required "+
                 "(version "+str(required_setuptools_version_major) +"+)."
                 " Please upgrade your setuptools version.")
except (ValueError, IndexError, NameError):
    sys.exit("CRITICAL ERROR: Unable to call setuptools version. Please upgrade setuptools.")

from setuptools.command.install import install as _install

import distutils
import os

# try to import urllib.request.urlretrieve for python3
try:
    from urllib.request import urlretrieve
except ImportError:
    from urllib import urlretrieve

import tarfile
import subprocess
import shutil
import zipfile
import tempfile
import re
import time

VERSION = "4.0"

AUTHOR = "HUMAnN Development Team"
AUTHOR_EMAIL = "humann-users@googlegroups.com"
MAINTAINER = "Lauren McIver"
MAINTAINER_EMAIL = "lauren.j.mciver@gmail.com"

def byte_to_megabyte(byte):
    """
    Convert byte value to megabyte
    """

    return byte / (1024.0**2)

class ReportHook():
    def __init__(self):
        self.start_time=time.time()

    def report(self, blocknum, block_size, total_size):
        """
        Print download progress message
        """

        if blocknum == 0:
            self.start_time=time.time()
            if total_size > 0:
                print("Downloading file of size: " + "{:.2f}".format(byte_to_megabyte(total_size)) + " MB\n")
        else:
            total_downloaded=blocknum*block_size
            status = "{:3.2f} MB ".format(byte_to_megabyte(total_downloaded))

            if total_size > 0:
                percent_downloaded=total_downloaded * 100.0 / total_size
                # use carriage return plus sys.stdout to overwrite stdout
                try:
                    download_rate=total_downloaded/(time.time()-self.start_time)
                    estimated_time=(total_size-total_downloaded)/download_rate
                except ZeroDivisionError:
                    download_rate=0
                    estimated_time=0
                estimated_minutes=int(estimated_time/60.0)
                estimated_seconds=estimated_time-estimated_minutes*60.0
                status +="{:3.2f}".format(percent_downloaded) + " %  " + \
                    "{:5.2f}".format(byte_to_megabyte(download_rate)) + " MB/sec " + \
                    "{:2.0f}".format(estimated_minutes) + " min " + \
                    "{:2.0f}".format(estimated_seconds) + " sec "
            status+="        \r"
            sys.stdout.write(status)

def download(url, download_file):
    """
    Download a file from a url
    """

    try:
        print("Downloading "+url)
        file, headers = urlretrieve(url,download_file,reporthook=ReportHook().report)
        # print final return to start new line of stdout
        print("\n")
    except EnvironmentError:
        print("WARNING: Unable to download "+url)

def download_unpack_tar(url,download_file_name,folder,software_name):
    """
    Download the url to the file and decompress into the folder
    """

    # Check for write permission to the target folder
    if not os.access(folder, os.W_OK):
        print("WARNING: The directory is not writeable: "+
            folder + " . Please modify the permissions.")

    download_file=os.path.join(folder, download_file_name)

    download(url, download_file)

    error_during_extract=False

    try:
        tarfile_handle=tarfile.open(download_file)
        tarfile_handle.extractall(path=folder)
        tarfile_handle.close()
    except (EnvironmentError,tarfile.ReadError):
        print("WARNING: Unable to extract "+software_name+".")
        error_during_extract=True

    if not error_during_extract:
        try:
            os.unlink(download_file)
        except EnvironmentError:
            print("WARNING: Unable to remove the temp download: " + download_file)

def download_unpack_zip(url,download_file_name,folder,software_name):
    """
    Download the url to the file and decompress into the folder
    """

    # Check for write permission to the target folder
    if not os.access(folder, os.W_OK):
        print("WARNING: The directory is not writeable: "+
            folder + " . Please modify the permissions.")

    download_file=os.path.join(folder, download_file_name)

    download(url, download_file)

    error_during_extract=False

    try:
        zipfile_handle=zipfile.ZipFile(download_file)
        zipfile_handle.extractall(path=folder)
        zipfile_handle.close()
    except EnvironmentError:
        print("WARNING: Unable to extract "+software_name+".")
        error_during_extract=True

    if not error_during_extract:
        try:
            os.unlink(download_file)
        except EnvironmentError:
            print("WARNING: Unable to remove the temp download: " + download_file)

def install_glpk(install_directory, replace_install=None):
    """
    Download and install the most recent glpk for minpath
    """

    glpk_url="http://ftp.gnu.org/gnu/glpk/glpk-4.55.tar.gz"
    glpk_folder="glpk-4.55"

    # Check if glpk is already installed
    glpk_installed=find_exe_in_path("glpsol")

    if not glpk_installed or replace_install:

        # get the prefix for the install directory
        install_directory=os.path.abspath(install_directory)
        prefix=os.path.abspath(os.path.join(install_directory,os.path.pardir))

        humann_source_folder=os.path.dirname(os.path.abspath(__file__))
        tempfolder=tempfile.mkdtemp(prefix="glpk_download_",dir=humann_source_folder)

        # install the most recent gplk software
        glpk_download=os.path.join(tempfolder,glpk_url.split('/')[-1])
        print("Installing latest glpk.")
        download_unpack_tar(glpk_url,glpk_download,tempfolder,"glpk")

        # move to the glpk directory
        current_working_directory=os.getcwd()
        glpk_install_folder=os.path.join(tempfolder,glpk_folder)

        try:
            os.chdir(glpk_install_folder)
        except EnvironmentError:
            print("WARNING: glpk was not downloaded.")

        # test for gcc
        try:
            subprocess_output=subprocess.check_output(["gcc","--version"])
        except (EnvironmentError,subprocess.CalledProcessError):
            print("WARNING: Please install gcc.")

        # test for make
        try:
            subprocess_output=subprocess.check_output(["make","--version"])
        except (EnvironmentError,subprocess.CalledProcessError):
            print("WARNING: Please install make.")

        try:
            # run configure
            subprocess.call(["./configure","--prefix",prefix])
            # run make
            subprocess.call(["make"])
            # run make install
            subprocess.call(["make","install"])
        except (EnvironmentError,subprocess.CalledProcessError):
            print("WARNING: Errors installing new glpk version.")

        # return to original working directory
        os.chdir(current_working_directory)

        # remove the temp download folder
        try:
            shutil.rmtree(tempfolder)
        except EnvironmentError:
            print("WARNING: Unable to remove temp install folder.")

    else:
        print("Found glpk install at "+glpk_installed)

def install_minpath(replace_install=None):
    """
    Download and install the minpath software if not already installed
    """

    # Download the minpath software v1.2
    # Check to see if already downloaded

    fullpath_scripts=os.path.join(os.path.dirname(os.path.abspath(__file__)),"humann","quantify")

    minpath_file="minpath1.2.tar.gz"
    minpath_url="http://omics.informatics.indiana.edu/mg/get.php?" + \
    "justdoit=yes&software=" + minpath_file
    minpath_install_folder="MinPath"

    # install minpath if not already installed
    minpath_exe=os.path.join(fullpath_scripts,minpath_install_folder,"MinPath1.2.py")
    if not os.path.isfile(minpath_exe) or replace_install:
        download_file=os.path.join(fullpath_scripts, minpath_file)
        print("Installing minpath.")
        download_unpack_tar(minpath_url,download_file,fullpath_scripts,"minpath")
    else:
        print("Found minpath install.")

def find_exe_in_path(exe):
    """
    Check that an executable exists in $PATH
    """

    paths = os.environ["PATH"].split(os.pathsep)
    for path in paths:
        fullexe = os.path.join(path,exe)
        if os.path.exists(fullexe):
            if os.access(fullexe,os.X_OK):
                return path
    return None

def install_boost(folder):
    """ Install boost locally in folder indicated
    """

    boost_file="boost_1_57_0.tar.gz"
    boost_url="http://sourceforge.net/projects/boost/files/boost/1.57.0/boost_1_57_0.tar.gz"

    print("Installing boost.")
    download_unpack_tar(boost_url, boost_file, folder, "boost")

    # get current working directory
    current_working_directory=os.getcwd()
    boost_build_dir=os.path.join(folder,"boost_1_57_0")

    try:
        os.chdir(boost_build_dir)
    except EnvironmentError:
        print("WARNING: boost directory does not exist")

    try:
        subprocess.call(["./bootstrap.sh","--with-libraries=timer,chrono,system,program_options,thread,iostreams","--prefix=../boost"])
        subprocess.call(["./b2","install"])
    except (EnvironmentError,subprocess.CalledProcessError):
        print("WARNING: Errors installing boost.")

    # return to original working directory
    os.chdir(current_working_directory)

def install_diamond(final_install_folder, build, replace_install=None):
    """
    Download and install the diamond software if not already installed
    """

    # Check if diamond is already installed
    diamond_installed=find_exe_in_path("diamond")

    if not diamond_installed or replace_install:
        diamond_version="2.0.15"
        diamond_exe="diamond"
        diamond_file="diamond-linux64.tar.gz"
        diamond_url="http://github.com/bbuchfink/diamond/releases/download/v{0}/diamond-linux64.tar.gz".format(diamond_version)

        # download source if build selected
        if build:
            diamond_file="v{0}.tar.gz".format(diamond_version)
            diamond_url="http://github.com/bbuchfink/diamond/archive/v{0}.tar.gz".format(diamond_version)

        humann_source_folder=os.path.dirname(os.path.abspath(__file__))
        tempfolder=tempfile.mkdtemp(prefix="diamond_download_",dir=humann_source_folder)

        # install the diamond software
        print("Installing diamond.")
        error_during_install=False
        download_unpack_tar(diamond_url, diamond_file, tempfolder, diamond_exe)

        # compile diamond if build selected
        if build:
            # get the current directory
            current_working_directory=os.getcwd()
            diamond_build_dir=os.path.join(tempfolder,"diamond-{0}".format(diamond_version),"src")

            try:
                os.chdir(diamond_build_dir)
            except EnvironmentError:
                print("WARNING: diamond directory does not exist")

            # test for gcc
            try:
                subprocess_output=subprocess.check_output(["gcc","--version"])
            except (EnvironmentError,subprocess.CalledProcessError):
                print("WARNING: Please install gcc.")

            # test for make
            make_installed=True
            try:
                subprocess_output=subprocess.check_output(["make","--version"])
            except (EnvironmentError,subprocess.CalledProcessError):
                make_installed=False
                print("INFO: make is not installed.")

            # test for cmake
            cmake_installed=True
            try:
                subprocess_output=subprocess.check_output(["cmake","--version"])
            except (EnvironmentError,subprocess.CalledProcessError):
                cmake_installed=False
                print("INFO: cmake is not installed.")

            final_install_prefix=os.path.join(final_install_folder,os.pardir)
            if make_installed and cmake_installed:
                # if make and cmake are installed, run the standard build
                print("INFO: Installing diamond with cmake method.")
                diamond_install_folder=os.path.join(tempfolder,"diamond-{0}".format(diamond_version),"bin")
                try:
                    # make the install bin directory and change directories
                    os.mkdir(diamond_install_folder)
                    os.chdir(diamond_install_folder)
                    # run cmake and make
                    subprocess.call(["cmake","..","-DCMAKE_INSTALL_PREFIX="+final_install_prefix])
                    subprocess.call(["make","install"])
                except (EnvironmentError,subprocess.CalledProcessError):
                    print("WARNING: Errors installing diamond.")
                    error_during_install=True
            else:
                # if make or cmake are not installed, run the simple build
                print("INFO: Installing diamond with simple build method.")

                diamond_install_folder=os.path.join(tempfolder,"diamond-{0}".format(diamond_version))
                try:
                    # change directories to the diamond source folder
                    os.chdir(diamond_install_folder)

                    # remove the static flag from the build script
                    subprocess.call(["sed","-i","s/ -static//", "build_simple.sh"])

                    # run the build script
                    subprocess.call(["bash","-x", "build_simple.sh"])
                except (EnvironmentError,subprocess.CalledProcessError):
                    print("WARNING: Errors installing diamond.")
                    error_during_install=True
                # set the location of the diamond executable
                diamond_exe_full_path=os.path.join(diamond_install_folder, diamond_exe)

                # copy the installed software to the final bin location
                try:
                    # copy to the install folder
                    shutil.copy(diamond_exe_full_path, final_install_folder)
                    # add executable permissions
                    os.chmod(os.path.join(final_install_folder,diamond_exe), 0o755)
                except (EnvironmentError, shutil.Error):
                    error_during_install=True

            # return to original working directory
            os.chdir(current_working_directory)
        else:
            # copy the installed software to the final bin location
            diamond_exe_full_path=os.path.join(tempfolder, diamond_exe)
            try:
                # copy to the install folder
                shutil.copy(diamond_exe_full_path, final_install_folder)
                # add executable permissions
                os.chmod(os.path.join(final_install_folder,diamond_exe), 0o755)
            except (EnvironmentError, shutil.Error) as e:
                print(e)
                error_during_install=True

        # remove the local diamond install
        try:
            shutil.rmtree(tempfolder)
        except EnvironmentError:
            print("WARNING: Unable to remove temp install folder.")

        if error_during_install:
            print("WARNING: Unable to install diamond. Please install diamond.")
        else:
            print("Installed diamond at "+final_install_folder)
    else:
        print("Found diamond install at "+diamond_installed)

def install_bowtie2(final_install_folder, mac_os, replace_install=None):
    """
    Download and install the bowtie2 software if not already installed
    """

    # Check if bowtie2 is already installed
    bowtie2_installed=find_exe_in_path("bowtie2")

    if not bowtie2_installed or replace_install:
        bowtie2_exe="bowtie2"
        bowtie2_file="bowtie2-2.2.9-linux-x86_64.zip"
        bowtie2_url="http://sourceforge.net/projects/bowtie-bio/files/bowtie2/2.2.9/bowtie2-2.2.9-linux-x86_64.zip"

        # if this is a MAC OS, select a different binary download
        if mac_os:
            bowtie2_file="bowtie2-2.2.9-macos-x86_64.zip"
            bowtie2_url="http://sourceforge.net/projects/bowtie-bio/files/bowtie2/2.2.9/bowtie2-2.2.9-macos-x86_64.zip"

        bowtie2_folder="bowtie2-2.2.9"

        humann_source_folder=os.path.dirname(os.path.abspath(__file__))
        tempfolder=tempfile.mkdtemp(prefix="bowtie2_download_",dir=humann_source_folder)

        # install the bowtie2 software
        print("Installing bowtie2.")
        error_during_install=False
        download_unpack_zip(bowtie2_url, bowtie2_file, tempfolder, bowtie2_exe)

        # copy the installed software to the final bin location
        # copy all bowtie2* executables
        fullpath_bowtie2_exe=os.path.join(tempfolder, bowtie2_folder)

        files=[]
        try:
            files=os.listdir(fullpath_bowtie2_exe)
        except EnvironmentError:
            print("WARNING: Bowtie2 files not found.")
            error_during_install=True

        for file in files:
            # check if this file is one of the bowtie2* executables
            if re.match(bowtie2_exe,file):
                try:
                    # copy to the install folder
                    shutil.copy(os.path.join(fullpath_bowtie2_exe,file), final_install_folder)
                    # add executable permissions
                    os.chmod(os.path.join(final_install_folder,file), 0o755)
                except (EnvironmentError, shutil.Error):
                    error_during_install=True

        # remove the local bowtie2 install
        try:
            shutil.rmtree(tempfolder)
        except EnvironmentError:
            print("WARNING: Unable to remove temp install folder.")

        if error_during_install:
            print("WARNING: Unable to install bowtie2. Please install bowtie2.")
        else:
            print("Installed bowtie2 at "+final_install_folder)
    else:
        print("Found bowtie2 install at "+bowtie2_installed)

class Install(_install):
    """
    Custom setuptools install command, set executable permissions for glpk
    """

    _install.user_options=_install.user_options+[('bypass-dependencies-install',
        None, 'bypass install of dependencies'),('build-diamond',None,'build diamond'),
        ('replace-dependencies-install',None,'replace dependencies already installed')]

    def initialize_options(self):
        self.bypass_dependencies_install=False
        self.build_diamond=False
        self.replace_dependencies_install=False
        _install.initialize_options(self)

    def finalize_options(self):
        _install.finalize_options(self)

    def run(self):

        # install minpath if not already installed
        install_minpath(replace_install=self.replace_dependencies_install)

        _install.run(self)

        # find the current install folder
        current_install_folder=None
        for item in os.listdir(self.install_lib):
            full_path_item=os.path.join(self.install_lib, item)
            if os.path.isdir(full_path_item):
                if "humann" == item:
                    current_install_folder=full_path_item

        # find all glpsol executables
        if current_install_folder is None:
            print("Unable to find install folder at: " + self.install_lib)
        else:
            glpsols=[]
            minpath_folder=os.path.join(current_install_folder,"quantify","MinPath")
            for root, directories, files in os.walk(minpath_folder):
                for filename in files:
                    if filename == "glpsol":
                        glpsols.append(os.path.join(root,filename))
            # change the permissions of the glpk modules to make sure
            # they are executable
            for file in glpsols:
                try:
                    os.chmod(file,0o755)
                except EnvironmentError:
                    print("Unable to add execute permissions for file: " + file)

        # find out the platform
        mac_os=False
        if sys.platform in ["darwin","os2","os2emx"]:
            mac_os=True

        # install dependencies if not already installed
        if not self.bypass_dependencies_install:
            # build diamond if set, or if on a Mac
            build_diamond=self.build_diamond
            if mac_os:
                build_diamond=True

            # get the final script install location for both source and wheel
            install_scripts=os.path.join(self.install_base, "bin")
            install_glpk(install_scripts,replace_install=self.replace_dependencies_install)
            install_diamond(install_scripts,build_diamond,replace_install=self.replace_dependencies_install)
            install_bowtie2(install_scripts,mac_os,replace_install=self.replace_dependencies_install)
        else:
            print("Bypass install of dependencies")

setuptools.setup(
    name="humann",
    author=AUTHOR,
    author_email=AUTHOR_EMAIL,
    version=VERSION,
    license="MIT",
    description="HUMAnN: The HMP Unified Metabolic Analysis Network, version 3",
    long_description="HUMAnN v3 is a pipeline for efficiently and accurately determining " + \
        "the coverage and abundance of microbial pathways in a community " + \
        "from metagenomic data. Sequencing a metagenome typically produces millions " + \
        "of short DNA/RNA reads. This process, referred to as functional profiling, " + \
        "aims to describe the metabolic potential of a microbial community and its " + \
        "members. More generally, functional profiling answers the question: What " + \
        "are the microbes in my community-of-interest doing (or capable of doing)?",
    url="http://huttenhower.sph.harvard.edu/humann",
    keywords=['microbial','microbiome','bioinformatics','microbiology','metagenomic','metatranscriptomic','humann','humann'],
    platforms=['Linux','MacOS'],
    classifiers=[
        "Programming Language :: Python",
        "Development Status :: 5 - Production/Stable",
        "Environment :: Console",
        "Operating System :: MacOS",
        "Operating System :: Unix",
        "Programming Language :: Python :: 2.7",
        "Programming Language :: Python :: 3.4",
        "Topic :: Scientific/Engineering :: Bio-Informatics"
        ],
    packages=setuptools.find_packages(),
    cmdclass={'install': Install},
    package_data={
        'humann' : [
            'humann.cfg',
            'data/pathways/*',
            'data/misc/*',
            'data/uniref_DEMO/*',
            'data/chocophlan_DEMO/*',
            'data/utility_DEMO/*',
            'tests/data/*.*',
            'tests/data/tooltest*/*',
            'quantify/MinPath/data/*',
            'quantify/MinPath/glpk-*/examples/glp*',
            'quantify/MinPath/glpk-*/examples/.libs/*'
        ]},
    entry_points={
        'console_scripts': [
            'humann = humann.humann:main',
            'humann3 = humann.humann:main',
            'humann_databases = humann.tools.humann_databases:main',
            'humann3_databases = humann.tools.humann_databases:main',
            'humann_config = humann.tools.humann_config:main',
            'humann_join_tables = humann.tools.join_tables:main',
            'humann_split_table = humann.tools.split_table:main',
            'humann_rename_table = humann.tools.rename_table:main',
            'humann_renorm_table = humann.tools.renorm_table:main',
            'humann_regroup_table = humann.tools.regroup_table:main',
            'humann_reduce_table = humann.tools.reduce_table:main',
            'humann_unpack_pathways = humann.tools.merge_abundance:main',
            'humann_test = humann.tests.humann_test:main',
            'humann_build_custom_database = humann.tools.build_custom_database:main',
            'humann_genefamilies_genus_level = humann.tools.genefamilies_genus_level:main',
            'humann_split_stratified_table = humann.tools.split_stratified_table:main',
            'humann_barplot = humann.tools.humann_barplot:main',
            'humann_benchmark = humann.tools.humann_benchmark:main',
            'humann_expand_cluster = humann.tools.expand_cluster:main'
        ]},
    test_suite= 'humann.tests.humann_test.unittests_suite_only',
    zip_safe = False
 )


from supermarket.helper_functions import table
import json

metadata_file_location = "./supermarket/file_metadata.json"
products = None
with open(metadata_file_location, "r") as metadata_file:
    metadata = json.load(metadata_file)
    products_info = metadata["products_info"]
    with open(
        f"{products_info['file_location']}/{products_info['file_name']}"
    ) as product_file:
        products = json.load(product_file)

def product_selection(person_bag):
    person_bag.append(["ITEMS", "PRICE", "QUANTITY"])
    while True:
        for index, category in enumerate(products.keys()):
            print(f"{index+1}. {category}")

        choice = int(input("Please choose a category: ")) - 1
        print("")

        sub_categorys = list(products.values())[choice]
        for index, sub_category in enumerate(sub_categorys.keys()):
            print(f"{index+1}. {sub_category}")

        choice = int(input("Please choose a sub-category: ")) - 1
        print("")

        items = list(sub_categorys.values())[choice]
        items_arr = [["ITEMS", "PRICE", "ID"]]
        for item_name, item_details in items.items():
            temp_arr = [item_name, f'{item_details["Price"]}', f'{item_details["Id"]}']
            items_arr.append(temp_arr)

        table.print_table(table.create_table(items_arr, 4))

        product_id_input = input(
            "\nPlease enter the ID of the product that you want to add to you bag: "
        )
        quantity = input("What is the quantity of the item that you want to add: ")

        for item_arr in items_arr:
            if product_id_input == item_arr[2]:
                person_bag.append([item_arr[0], item_arr[1], quantity])
                break
        print("\nHere's your current bag\n")
        table.print_table(table.create_table(person_bag, 4))
        continue_shopping = input(
            "\nDo you want to continue shopping? (Yes/No): "
        ).upper()
        if continue_shopping == "YES":
            print("")
            continue
        else:
            break


import re

from nltk.tokenize import sent_tokenize

text_replacements = {
    'U. S.': 'US',
    'U. K.': 'UK',
    'Sen.': 'Sen'
}

def is_hidden_file(name: str):
    return name.startswith('.')

def clean_sentence(sentence: str):
    out = sentence[:-1]
    for orig, new in text_replacements.items():
        out = out.replace(orig, new)
    return out + ' <PUNCT> '

def clean_text(txt: str):
    fixed_txt = re.sub(r'\.(?=[^ \W\d])', '. ', txt)
    fixed_txt = fixed_txt.replace('\n', ' <NEWLINE> ')
    sentences = sent_tokenize(fixed_txt)
    sentences = [clean_sentence(s.strip()) for s in sentences]
    return '<START> ' + ''.join(sentences) + '<EOS>'


import re

from spacy.tokens import Doc
from spacy.tokens.span import Span

from eucy import regex as eure
from eucy import structure, utils

class Elements:

    def __init__(self):

        # Set Doc-level element extions
        if not Doc.has_extension("article_elements"):
            Doc.set_extension("article_elements", default=None)

        if not Doc.has_extension("parts"):
            self.EuStructure = structure.Structure()

        # Set Span-level element extensions
        if not Span.has_extension("element_type"):
            Span.set_extension("element_type", default=None)
            Span.set_extension("element_pos",
                               default=None)  # the index of an element
            Span.set_extension(
                "element_num", default=None
            )  # the number an element has been assigned in the text
            Span.set_extension(
                "element_numstr", default=None
            )  # the number an element has been assigned in the text

    def __call__(self, doc, overwrite=False):
        """Apply the pipeline component to a `Doc` object.
        doc (Doc): The `Doc` returned by the previous pipeline component.
        RETURNS (Doc): The modified `Doc` object.
        """

        if doc._.parts is None:

            doc = self.EuStructure(doc)

        # Add SpanGroup for ...

        if doc.spans.get(
                'citations'
        ) is None or overwrite:  # in case a pre-coded (e.g. modified) doc is passed
            ## Citations
            doc.spans['citations'] = citations(doc._.parts['citations'])

        if doc.spans.get('recitals') is None or overwrite:
            ## Recitals
            doc.spans['recitals'] = recitals(doc._.parts['recitals'])

        if doc.spans.get('articles') is None or overwrite:
            ## Articles
            doc.spans['articles'] = articles(doc._.parts['enacting'])

        doc._.article_elements = []

        for article in doc.spans['articles']:
            doc._.article_elements.append(article_elements(article))

        return doc

def citations(doc_citations):

    citations = utils._part_argument_check(doc_citations, 'citations')

    if citations is None:
        return []

    char_token = utils.chars_to_tokens_dict(citations,
                                            input="as_ref",
                                            output="as_ref")

    citation_matches = []

    while True:

        citation_matches = [
            m for m in re.finditer(eure.elements['citation'],
                                   doc_citations.text,
                                   flags=re.MULTILINE)
        ]

        if len(citation_matches) > 0:
            break

        # try again ignoring case
        citation_matches = [
            m for m in re.finditer(eure.elements['citation'],
                                   doc_citations.text,
                                   flags=re.MULTILINE | re.IGNORECASE)
        ]

        break

    citation_matches = [ma for ma in citation_matches if len(ma.group(0)) > 4]

    if isinstance(doc_citations, Doc):
        citation_list = [
            doc_citations.char_span.char_span(mat.start(),
                                              mat.end(),
                                              alignment_mode="expand")
            for mat in citation_matches
        ]
    else:
        citation_list = [
            doc_citations[
                utils.char_to_token(mat.start(), char_token):utils.
                char_to_token(mat.end(), char_token, alignment_mode="expand")]
            for mat in citation_matches
        ]

    if Span.has_extension("element_type"):
        for i, c in enumerate(citation_list):
            citation_list[i]._.element_type = "citation"

    return citation_list

def recitals(doc_recitals):

    recitals = utils._part_argument_check(doc_recitals, 'recitals')

    if recitals is None:
        return []

    char_token = utils.chars_to_tokens_dict(recitals,
                                            input="as_ref",
                                            output="as_ref")

    def get_recitals(recital_matches):
        recital_matches_pos = [
            (m.start(),
             recital_matches[i +
                             1].start()) if i < len(recital_matches) - 1 else
            (m.start(), len(recitals.text))
            for i, m in enumerate(recital_matches)
        ]

        if isinstance(doc_recitals, Doc):
            recital_list = [
                doc_recitals.char_span.char_span(mat[0],
                                                 mat[1],
                                                 alignment_mode="expand")
                for mat in recital_matches_pos
            ]
        else:
            recital_list = [
                doc_recitals[
                    utils.char_to_token(mat[0], char_token):utils.
                    char_to_token(mat[1], char_token, alignment_mode="expand")]
                for mat in recital_matches_pos
            ]

        return recital_list

    recital_matches = []

    recital_list = []

    while True:

        # @TODO differentiate between parentheses recitals and dot recitals and, in case of mixed captures, decide which to use (e.g. 52017PC0603)
        # @TODO catch recitals with multiple paragraphs completely (cf. example_1.txt)

        recital_matches_num = [
            m for m in re.finditer(eure.elements['recital_num_start'],
                                   recitals.text, re.MULTILINE)
        ]

        recitals_parentheses = []
        recitals_dot = []
        recitals_other = []

        # try to filter out bad matches by checking for the type (parentheses, dot, other)
        for match in recital_matches_num:
            if re.search(r'^[\s\(]*[0-9]+\s*\)', match.group(0)) is not None:
                recitals_parentheses.append(match)
            elif re.search(r'^[\s]*[0-9]+\s*\.', match.group(0)) is not None:
                recitals_dot.append(match)
            else:
                recitals_other.append(match)

        # if there are only parentheses recitals, use them
        if len(recitals_parentheses) > 0 and len(recitals_parentheses) > len(
                recitals_dot) and len(recitals_parentheses) > len(
                    recitals_other) == 0:
            recital_matches_num = recitals_parentheses
        # @TODO maybe use a more refined approach including dot recitals / other recitals if they are not too far away from the parentheses recitals

        recital_list = get_recitals(recital_matches_num)

        # remove recitals with less than 10 words
        recital_list = [rec for rec in recital_list if len(rec) > 6]

        if len(recital_list) > 0:
            break

        # if none found yet, try for unnumbered recirals starting with whereas
        recital_matches_whereas = [
            ma for ma in re.finditer(eure.elements['recital_whereas_start'],
                                     recitals.text,
                                     flags=re.MULTILINE | re.IGNORECASE)
        ]

        recital_list = get_recitals(recital_matches_whereas)

        # remove recitals with less than 6 words
        recital_list = [rec for rec in recital_list if len(rec) > 6]

        if len(recital_list) > 0:
            break

        # if none  found yet, search for normal paragraph recitals
        recital_matches_par = [
            ma for ma in re.finditer(eure.elements['recital_par_start'],
                                     recitals.text,
                                     flags=re.MULTILINE | re.IGNORECASE)
        ]

        recital_list = get_recitals(recital_matches_par)

        # remove recitals with less than 6 words
        recital_list = [rec for rec in recital_list if len(rec) > 6]

        if len(recital_list) > 0:
            break

    if Span.has_extension("element_type"):
        for i, c in enumerate(recital_list):
            recital_list[i]._.element_type = "recital"

    return recital_list

def articles(doc_articles):

    articles = utils._part_argument_check(doc_articles, 'enacting')

    if articles is None:
        return []

    char_token = utils.chars_to_tokens_dict(articles,
                                            input="as_ref",
                                            output="as_ref")

    article_id_matches = []

    while True:

        article_id_matches = [
            m for m in re.finditer(eure.elements['article_identifier'],
                                   articles.text,
                                   flags=re.MULTILINE | re.IGNORECASE)
        ]

        if len(article_id_matches) > 0:
            break

        # test for single article
        article_id_matches = [
            m for m in re.finditer(eure.elements['single_article_identifier'],
                                   articles.text,
                                   flags=re.MULTILINE | re.IGNORECASE)
        ]

        if len(article_id_matches) == 1:
            break

        # @TODO: handle cases where neither normal nor sinngle article identifier regexes match

        break

    article_id_matches = [
        ma for ma in article_id_matches if len(ma.group(0)) > 4
    ]

    article_list = []

    for i, m in enumerate(article_id_matches, start=0):

        # match text of entire article
        article_start = m.start()

        if i < len(article_id_matches) - 1:
            article_end = article_id_matches[i + 1].start()
        else:
            article_end = len(doc_articles.text)

        if isinstance(doc_articles, Doc):
            article_span = doc_articles.char_span.char_span(
                article_start, article_end, alignment_mode="expand")
        else:
            article_span = doc_articles[utils.char_to_token(
                article_start, char_token):utils.char_to_token(
                    article_end, char_token, alignment_mode="expand")]

        # set extensions
        if article_span.has_extension("element_pos"):
            article_span._.element_pos = i + 1
        if article_span.has_extension("element_numstr"):

            article_nums = []

            while True:

                article_nums = [
                    m.group(1)
                    for m in re.finditer(eure.elements['article_num'],
                                         article_span.text.strip(),
                                         flags=re.IGNORECASE)
                ]

                if len(article_nums) > 0:
                    break

                article_nums = [
                    ma.group(0)
                    for ma in re.finditer(eure.elements['article_any_num'],
                                          article_span.text.strip())
                ]

                break

            if len(article_nums) > 1:
                article_numstr = article_nums[0].strip()
            else:
                article_numstr = None

            article_span._.element_numstr = article_numstr

        article_list.append(article_span)

    if Span.has_extension("element_type"):
        for i, a in enumerate(article_list):
            article_list[i]._.element_type = "article"

    return article_list

def article_elements(doc_article):
    """Mark up an article, inlcuding Title, Paragraph, Point, etc...
    """

    if not isinstance(doc_article, (Doc, Span)):
        raise TypeError(
            str(type(doc_article)) +
            " type not supported. Please pass a Doc or Span object.")

    def _paragraphs(article):

        par_char_token = utils.chars_to_tokens_dict(article,
                                                    input="as_ref",
                                                    output="as_ref")

        par_start_matches = []

        par_start_matches = [
            m for m in re.finditer(eure.elements['article_num_paragraph'],
                                   article.text,
                                   flags=re.MULTILINE)
        ]

        unnum_pars = False

        if len(par_start_matches) == 0:
            par_start_matches = [
                ma
                for ma in re.finditer(eure.elements['article_unnum_paragraph'],
                                      article.text,
                                      flags=re.MULTILINE | re.DOTALL)
            ]
            unnum_pars = True

        par_list = []

        for i, m in enumerate(par_start_matches, start=0):

            # match text of entire article
            if unnum_pars and len(m.groups()) > 0:
                par_start = m.start(1)
            else:
                par_start = m.start()

            if i < len(par_start_matches) - 1:
                if unnum_pars and len(m.groups()) > 0:
                    par_end = par_start_matches[i + 1].start(1)
                else:
                    par_end = par_start_matches[i + 1].start()
            else:
                par_end = len(article.text)

            if isinstance(article, Doc):
                par_span = article.char_span(par_start,
                                             par_end,
                                             alignment_mode="expand")
            else:
                par_span = article[utils.char_to_token(
                    par_start, par_char_token):utils.char_to_token(
                        par_end, par_char_token, alignment_mode="expand")]

            # set extensions
            if par_span.has_extension("element_pos"):
                par_span._.element_pos = i + 1
            if par_span.has_extension("element_numstr"):
                if len(m.groups()) > 0:
                    par_span._.element_numstr = m.group(1)

            par_list.append(par_span)

        if Span.has_extension("element_type"):
            for i, a in enumerate(par_list):
                par_list[i]._.element_type = "art_par"

        # if no paragraphs found, set article as par
        if len(par_list) == 0 and len(article.text.strip()) > 1:
            par_list.append(article)
            if Span.has_extension("element_type"):
                for i, a in enumerate(par_list):
                    par_list[i]._.element_type = "art_par"
        return par_list

    def _subparagraphs(par):
        """find and return all subpar spans (always unnumbered). Each par has at least one subpar."""

        subpar_char_token = utils.chars_to_tokens_dict(par,
                                                       input="as_ref",
                                                       output="as_ref")

        subpar_start_matches = []

        # mtch each par so that each par has at least one subpar
        subpar_start_matches = [
            m for m in re.finditer(eure.elements['article_subpar_start'],
                                   par.text,
                                   flags=re.MULTILINE)
        ]

        subpar_list = []

        for i, m in enumerate(subpar_start_matches, start=0):

            # match text of entire par
            subpar_start = m.start()

            if i < len(subpar_start_matches) - 1:
                subpar_end = subpar_start_matches[i + 1].start()
            else:
                subpar_end = len(par.text)

            if isinstance(par, Doc):
                subpar_span = par.char_span(subpar_start,
                                            subpar_end,
                                            alignment_mode="expand")
            else:
                subpar_span = par[utils.char_to_token(
                    subpar_start, subpar_char_token):utils.char_to_token(
                        subpar_end, subpar_char_token, alignment_mode="expand"
                    )]

            # sort out chpater/section titles
            if len(subpar_span.text.strip()) < 200 and bool(
                    re.search(eure.elements['article_section_titles'],
                              subpar_span.text,
                              flags=re.IGNORECASE | re.MULTILINE)):
                continue

            # set extensions
            if subpar_span.has_extension("element_pos"):
                subpar_span._.element_pos = i + 1

            subpar_list.append(subpar_span)

        if Span.has_extension("element_type"):
            for i, a in enumerate(subpar_list):
                subpar_list[i]._.element_type = "art_subpar"

        return subpar_list

    def _points(subpar):

        point_char_token = utils.chars_to_tokens_dict(subpar,
                                                      input="as_ref",
                                                      output="as_ref")

        point_start_matches = []

        # mtch each subpar so that each subpar has at least one point
        point_start_matches = [
            m for m in re.finditer(eure.elements['article_point_id'],
                                   subpar.text,
                                   flags=re.MULTILINE)
        ]

        point_list = []

        for i, m in enumerate(point_start_matches, start=0):

            # match text of entire subpar
            point_start = m.start()

            if i < len(point_start_matches) - 1:
                point_end = point_start_matches[i + 1].start()
            else:
                point_end = len(subpar.text)

            if isinstance(subpar, Doc):
                point_span = subpar.char_span(point_start,
                                              point_end,
                                              alignment_mode="expand")
            else:
                point_span = subpar[utils.char_to_token(
                    point_start, point_char_token):utils.char_to_token(
                        point_end, point_char_token, alignment_mode="expand")]

            # set extensions
            if point_span.has_extension("element_pos"):
                point_span._.element_pos = i + 1
            if point_span.has_extension("element_numstr"):
                point_span._.element_numstr = m.group(0)

            point_list.append(point_span)

        if Span.has_extension("element_type"):
            for i, a in enumerate(point_list):
                point_list[i]._.element_type = "art_point"

        return point_list

    def _indents(subpar):

        indent_char_token = utils.chars_to_tokens_dict(subpar,
                                                       input="as_ref",
                                                       output="as_ref")

        indent_start_matches = []

        # mtch each subpar so that each subpar has at least one indent
        indent_start_matches = [
            m for m in re.finditer(eure.elements['article_indent_id'],
                                   subpar.text,
                                   flags=re.MULTILINE)
        ]

        indent_list = []

        for i, m in enumerate(indent_start_matches, start=0):

            # match text of entire subpar
            indent_start = m.start()

            if i < len(indent_start_matches) - 1:
                indent_end = indent_start_matches[i + 1].start()
            else:
                indent_end = len(subpar.text)

            if isinstance(subpar, Doc):
                indent_span = subpar.char_span(indent_start,
                                               indent_end,
                                               alignment_mode="expand")
            else:
                indent_span = subpar[utils.char_to_token(
                    indent_start, indent_char_token
                ):utils.char_to_token(
                    indent_end, indent_char_token, alignment_mode="expand")]

            # set extensions
            if indent_span.has_extension("element_pos"):
                indent_span._.element_pos = i + 1

            indent_list.append(indent_span)

        if Span.has_extension("element_type"):
            for i, a in enumerate(indent_list):
                indent_list[i]._.element_type = "art_indent"

        return indent_list

    # Paragraphs
    par_spans = _paragraphs(doc_article)

    par_subpar_spans = []

    par_subpar_point_spans = []
    par_subpar_indent_spans = []

    for par in par_spans:
        # subparagraphs
        subpar_spans = _subparagraphs(par)

        par_subpar_spans.append(subpar_spans)

        subpar_point_spans = []
        subpar_indent_spans = []

        for subpar in subpar_spans:
            # points
            subpar_point_spans.append(_points(subpar))
            subpar_indent_spans.append(_indents(subpar))

        par_subpar_point_spans.append(subpar_point_spans)
        par_subpar_indent_spans.append(subpar_indent_spans)

    return {
        'pars': par_spans,
        'subpars': par_subpar_spans,
        'points': par_subpar_point_spans,
        'indents': par_subpar_indent_spans
    }