#!/usr/bin/env python3

import math
import argparse
from collections import Counter
from utils.utils import extract_analysis, generate_regex, extract_tag_from_analysis

if __name__ == '__main__':
	DESCRIPTION = '''generate a regex weightlist given an annotated corpus'''
	EPILOG = '''If the three weightlists are used then the analyses will be prioritized according to:
				The unigram counts of analyses found in the tagged corpus,
				The number of times the tag of an unweighted analyses was found in the tagged corpus,
				Laplace smoothed weight for remaining unweighted analyses.
				'''
	parser = argparse.ArgumentParser(description=DESCRIPTION, epilog=EPILOG)
	parser.add_argument('TAGGED_CORPUS',
						type=argparse.FileType('r'),
						help='input tagged corpus')
	parser.add_argument('ANALYSIS_WEIGHTLIST',
						type=argparse.FileType('w'),
						help='output weightlist for specific analyses')
	parser.add_argument('--tag_weightlist',
						type=argparse.FileType('w'),
						help='output weightlist for specific tags')
	parser.add_argument('--default_weightlist',
						type=argparse.FileType('w'),
						help='output weightlist for out-of-corpus analyses')

	args = parser.parse_args()
	TAGGED_CORPUS = args.TAGGED_CORPUS
	ANALYSIS_WEIGHTLIST_FILE = args.ANALYSIS_WEIGHTLIST
	TAG_WEIGHTLIST_FILE = args.tag_weightlist
	DEFAULT_WEIGHTLIST_FILE = args.default_weightlist

	lines = TAGGED_CORPUS.readlines()
	analyses = [extract_analysis(line.strip()) for line in lines]
	# Find the counts of each analysis
	regex_analyses = Counter([generate_regex(analysis) for analysis in analyses if not analysis.startswith('*')])
	# Use the den and num_offset according to the desired weightlists
	den = sum(regex_analyses.values())
	num_offset = 0

	if TAG_WEIGHTLIST_FILE:
		# Use the tags counts to disambiguate unweighthed analyses
		tags = [extract_tag_from_analysis(line.strip()) for line in lines]
		regex_tags = Counter([generate_regex(tag, match_all_prefixes=True) for tag in tags if tag and not tag.startswith('*')])
		# Adjust the numerator and denominator such that the weights
		# are larger than those of the analyses unigram counts
		den += sum(regex_tags.values()) * (1+len(regex_analyses))
		# Todo: Is this equivalent to the size of the corpus?
		num_offset += sum(regex_tags.values())

	if DEFAULT_WEIGHTLIST_FILE:
		# Apply laplace smoothing to weight analyses not found in the corpus
		den += 1 + len(regex_analyses)
		if TAG_WEIGHTLIST_FILE:
			den += len(regex_tags)
		num_offset += 1

	weighted_regex_analyses = ['{}::{}'.format(regex, -math.log(count + num_offset)+math.log(den))
							for regex, count in regex_analyses.most_common()]
	ANALYSIS_WEIGHTLIST_FILE.write('\n'.join(weighted_regex_analyses))

	if TAG_WEIGHTLIST_FILE:
		offset = 1 if DEFAULT_WEIGHTLIST_FILE else 0
		weighted_regex_tags = ['{}::{}'.format(regex, -math.log(count + offset) +math.log(den))
								for regex, count in regex_tags.most_common()]
		TAG_WEIGHTLIST_FILE.write('\n'.join(weighted_regex_tags))


	if DEFAULT_WEIGHTLIST_FILE:
		DEFAULT_WEIGHTLIST_FILE.write('[?*]::{}'.format(-math.log(1) +math.log(den)))
