import os
import re
import json

import nltk
import string
#nltk.download('punkt')
nltk.download('stopwords')
from nltk.corpus import stopwords

def tag_token(tags, token_pos, tag):
	if token_pos > 0 and f'{tag}' in tags[token_pos - 1]:
		
		if tags[token_pos] == 'O':
			tags[token_pos] = f'I-{tag}'
		elif f'I-{tag}' not in tags[token_pos]:
			tags[token_pos] += f';I-{tag}'
	else:
		if tags[token_pos] == 'O':
			tags[token_pos] = f'B-{tag}'
		elif f'B-{tag}' not in tags[token_pos]:
			tags[token_pos] += f';B-{tag}'


def remove_trailing_punctuation(token):
	while token and (re.search(r'[^\w\s\']', token[-1]) or (len(token) > 1 and token[-1] == "'")):
		token = token[:-1]
		
	return token


class AnnToBioConverter:
	
	def __init__(self, input_dir=None, txt_dir=None, ann_dir=None, output_dir=None, filtered_entities=[]):
		"""
		Initializes an instance of the AnnToBioConverter class.
		:param input_dir: (str) Directory containing both the text and annotation files.
		:param txt_dir: (str) Directory containing the text files.
		:param ann_dir: (str) Directory containing the annotation files.
		:param output_dir: (str) Directory where the JSON file will be saved.
		"""
		if input_dir is not None:
			self.txt_dir = input_dir
			self.ann_dir = input_dir
		elif txt_dir is not None and ann_dir is not None:
			self.txt_dir = txt_dir
			self.ann_dir = ann_dir
		else:
			raise ValueError("Either input_dir or txt_dir and ann_dir must be provided.")
		
		if output_dir is not None:
			self.output_dir = output_dir
		else:
			raise ValueError("output_dir must be provided.")
		self.data = {}
		
		self.filtered_entities = filtered_entities
	
	def _splitting_tokens(self, file_id, start, end, hyphen_split):
		"""
		Splits a multi-word token into separate tokens and returns a list of tokens and their respective start and end indices.
		:param file_id: (str) The ID of the file containing the token.
		:param start: (int) The starting index of the token in the text.
		:param end: (int) The ending index of the token in the text.
		:return: (tuple) A tuple containing a list of tokens and their respective start and end indices.
		"""
		
		text = self.data[file_id]['text']
		token = text[start:end]
		
		extra_sep = ['\u200a']
		if hyphen_split:
			extra_sep += ['-', '\u2010', '\u2011', '\u2012', '\u2013', '\u2014', '\u2015', '\u2212', '\uff0d']
		
		new_range = []
		tokens = []
		
		curr = start
		new_start = None
		
		for c in token + " ":
			if c.isspace() or c in extra_sep:
				if new_start:
					new_range.append([new_start, curr])
					tokens.append(text[new_start:curr])
					new_start = None
			elif not new_start:
				new_start = curr
			curr += 1
		
		return tokens, new_range
	
	def _load_txt(self):
		"""
		Loads the text files into the instance's data dictionary.
		"""
		for file_name in os.listdir(self.txt_dir):
			if file_name.endswith(".txt"):
				with open(os.path.join(self.txt_dir, file_name), "r") as f:
					text = f.read()
				file_id = file_name.split('.')[0]
				self.data[file_id] = {
					"text": text,
					"annotations": []
				}
	
	def _load_ann(self, hyphen_split):
		for file_name in os.listdir(self.ann_dir):
			
			if file_name.endswith(".ann"):
				with open(os.path.join(self.ann_dir, file_name), "r") as f:
					
					file_id = file_name.split('.')[0]
					annotations = []
					
					for line in f:
						if line.startswith("T"):
							fields = line.strip().split("\t")
							if len(fields[1].split(" ")) > 1:
								label = fields[1].split(" ")[0]
								
								# Extracting start end indices (Few annotations contain more than one disjoint ranges)
								start_end_range = [
									list(map(int, start_end.split()))
									for start_end in ' '.join(fields[1].split(" ")[1:]).split(';')
								]
								
								start_end_range_fixed = []
								for start, end in start_end_range:
									tokens, start_end_split_list = self._splitting_tokens(file_id, start, end,
																						  hyphen_split)
									start_end_range_fixed.extend(start_end_split_list)
								
								# Adding labels, start, end to annotations
								for start, end in start_end_range_fixed:
									annotations.append({
										"label": label,
										"start": start,
										"end": end
									})
					# sort annotations based on 'start' key before adding it to our dataset
					annotations = sorted(annotations, key=lambda x: (x['start'], x['label']))
					self.data[file_id]["annotations"] = annotations
		self._manual_fix()
	
	def split_text(self, file_id):
		text = self.data[file_id]['text']
		regex_match = r'[^\s\u200a\-\u2010-\u2015\u2212\uff0d]+'  # r'[^\s\u200a\-\—\–]+'
		
		tokens = []
		start_end_ranges = []
		
		sentence_breaks = []
		
		start_idx = 0
		
		for sentence in text.split('\n'):
			words = [match.group(0) for match in re.finditer(regex_match, sentence)]
			processed_words = list(map(remove_trailing_punctuation, words))
			sentence_indices = [(match.start(), match.start() + len(token)) for match, token in
								zip(re.finditer(regex_match, sentence), processed_words)]
			
			# Update the indices to account for the current sentence's position in the entire text
			sentence_indices = [(start_idx + start, start_idx + end) for start, end in sentence_indices]
			
			start_end_ranges.extend(sentence_indices)
			tokens.extend(processed_words)
			
			sentence_breaks.append(len(tokens))
			
			start_idx += len(sentence) + 1
		return tokens, start_end_ranges, sentence_breaks
	
	def write_bio_files(self):
		
		for file_id in self.data:
			text = self.data[file_id]['text']
			annotations = self.data[file_id]['annotations']
			
			# Tokenizing
			tokens, token2text, sentence_breaks = self.split_text(file_id)
			
			# Initialize the tags
			tags = ['O'] * len(tokens)
			
			ann_pos = 0
			token_pos = 0
			
			while ann_pos < len(annotations) and token_pos < len(tokens):
				
				label = annotations[ann_pos]['label']
				start = annotations[ann_pos]['start']
				end = annotations[ann_pos]['end']
				
				if self.filtered_entities:
					if label not in self.filtered_entities:
						# increment to access next annotation
						ann_pos += 1
						continue
				
				ann_word = text[start:end]
				
				# find the next word that fall between the annotation start and end
				while token_pos < len(tokens) and token2text[token_pos][0] < start:
					token_pos += 1
				
				if tokens[token_pos] == ann_word or \
					ann_word in tokens[token_pos] or \
					re.sub(r'\W+', '', ann_word) in re.sub(r'\W+', '', tokens[token_pos]):
					tag_token(tags, token_pos, label)
				elif ann_word in tokens[token_pos - 1] or \
					ann_word in tokens[token_pos - 1] or \
					re.sub(r'\W+', '', ann_word) in re.sub(r'\W+', '', tokens[token_pos - 1]):
					tag_token(tags, token_pos - 1, label)
				else:
					print(tokens[token_pos], tokens[token_pos - 1], ann_word, label)
				
				# increment to access next annotation
				ann_pos += 1
			
			# Write the tags to a .bio file
			with open(os.path.join(self.output_dir, f"{file_id}.bio"), 'w') as f:
				for i in range(len(tokens)):
					token = tokens[i].strip()
					if token:
						if i in sentence_breaks:
							f.write("\n")
						f.write(f"{tokens[i]}\t{tags[i]}\n")
	
	def _manual_fix(self):
		fix = {
			'19214295': {
				425: 424
			}
		}
		for file_id in self.data:
			if file_id in fix:
				for ann in self.data[file_id]['annotations']:
					if ann['start'] in fix[file_id]:
						ann['start'] = fix[file_id][ann['start']]
	
	def load_data(self, reload=False, hyphen_split=False):
		if not self.data or reload:
			self.data = {}
			self._load_txt()
			self._load_ann(hyphen_split)
	
	def create_json(self, data_dir, reload=False, hyphen_split=False):
		if not self.data or reload:
			self.load_data(hyphen_split, reload)
		
		# Write the dictionary to a JSON file
		with open(os.path.join(data_dir, "data.json"), "w") as f:
			json.dump(self.data, f)
	
	def load(self):
		self.load_txt()
		self.load_ann()
	
	def __getattr__(self, name):
		if name == 'data':
			return self.data
		else:
			raise AttributeError(f"'AnnToBioConverter' object has no attribute '{name}'")