#!/usr/bin/env python3

import sys
import argparse
import random
import pyonmttok
import numpy as np

def parse_args():
	parser = argparse.ArgumentParser(description='Replace words by dictionary')

	parser.add_argument('--dict_path', type=str, default='', help='Path of the dictionary')
	parser.add_argument('--src_path', type=str, default='', help='Path of source')
	parser.add_argument('--tgt_path', type=str, default='', help='Path of target')
	parser.add_argument('--delim', '-d', type=str, default='')
	parser.add_argument('--no_sample', dest='sample', default=True, action='store_false')

	return parser.parse_args()


def read_dict(dict_path, delim=''):
	dict_unique = {}
	with open(dict_path, 'r', encoding='utf-8') as fin:
		for line in fin:
			line = line.strip()
			word_src, word_tgt = '', ''
			try:
				if delim:
					word_src, word_tgt = line.split('\t')
				else:
					word_src, word_tgt = line.split()
			except:
				continue
			if word_src not in dict_unique.keys():
				dict_unique[word_src] = [word_tgt]
			else:
				if word_tgt not in dict_unique[word_src]:
					dict_unique[word_src] += word_tgt,
		return dict_unique


def sub_list(src, dict_path, delim='', sample=True):
	dict_unique = read_dict(dict_path, delim)
	tokenizer = pyonmttok.Tokenizer(mode='aggressive')
	for line in src:
		stripped = line.strip()
		tokens = tokenizer.tokenize(stripped, as_token_objects=True)
		for token in tokens:
			word = token.surface
			if sample:
				token.surface = random.sample(dict_unique.get(word, [word]), 1)[0].lower()
			else:
				token.surface = dict_unique.get(word, [word])[0].lower()
		for token in tokens:
			if token.casing == pyonmttok.Casing.MIXED:
				token.casing = pyonmttok.Casing.CAPITALIZED
		yield f'{tokenizer.detokenize(tokens)}\n'


def sub_corpus(src_path, tgt_path, dict_path, delim='', sample=True):

	dict_unique = read_dict(dict_path)

	total, replaced = 0, 0

	with open(src_path, 'r', encoding='utf-8') as fin_src, \
	  open(tgt_path, 'w', encoding='utf-8') as fout_tgt, \
	  open(f'{tgt_path}.rates', 'w') as tgt_replacement_rates:

		# openNMT tokenizer
		tokenizer = pyonmttok.Tokenizer(mode='aggressive')
		rates = []
		for line in fin_src:
			s_total, s_replaced = 0, 0  # counters for each sentences
			stripped = line.strip()
			tokens = tokenizer.tokenize(stripped, as_token_objects=True)
			# process tokens
			for token in tokens:
				# word was replaced
				word = token.surface
				if word in dict_unique:
					replaced += 1
					total += 1
					s_total += 1
					s_replaced += 1
				# word is alphabetical and was not replaced
				elif word.isalpha():
					total += 1
					s_total += 1
				if sample:
					token.surface = random.sample(dict_unique.get(word, [word]), 1)[0].lower()
				else:
					token.surface = dict_unique.get(word, [word])[0].lower()
			if s_total != 0:
				rates += str(s_replaced/float(s_total)),
			else:
				rates += '0.0',
			# add to target sentences
			if s_total != 0 and s_replaced/float(s_total) >= 0.1:
				for token in tokens:
					if token.casing == pyonmttok.Casing.MIXED:
						token.casing = pyonmttok.Casing.CAPITALIZED
				fout_tgt.write(f'{tokenizer.detokenize(tokens)}\n')
			else:
				fout_tgt.write(f'{stripped}\n')

		tgt_replacement_rates.write('\n'.join(rates)+'\n')
		tgt_replacement_rates.write(f'Replaced {replaced} words, {total} words in total，percentage: {np.mean(list([float(rate) for rate in rates]))}\n')

if __name__ == '__main__':

	args = parse_args()
	sub_corpus(**args.__dict__)
