# -*- coding: utf-8 -*-
"""
Created on Tue Oct 11 16:46:45 2022

@author: Santiago Moreno
"""
from upsampling import upsampling_ner
from flair.datasets import ColumnCorpus
from flair.data import Corpus
from flair.trainers import ModelTrainer
from flair.models import SequenceTagger
from flair.embeddings import TransformerWordEmbeddings
from torch.optim.lr_scheduler import OneCycleLR
from flair.data import Sentence
from sklearn.model_selection import StratifiedGroupKFold
from distutils.dir_util import copy_tree
import numpy as np
import torch
import pandas as pd
import json
import os
import operator
import flair
import argparse

default_path = os.path.dirname(os.path.abspath(__file__))
tagger_document = 0
tagger_sentence = 0

def check_create(path):
    import os
    
    if not (os.path.isdir(path)):
        os.makedirs(path)
        
def str2bool(v):
    if isinstance(v, bool):
       return v
    if v.lower() in ('yes', 'True','true', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'False', 'false', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected.')
        

def copy_data(original_path):
    data_folder  = default_path + '/../../data/NER/train'
    copy_tree(original_path, data_folder)
    
def characterize_data():
    data_folder  = default_path + '/../../data/NER/train'
    columns = {0: 'text', 1:'ner'}
    
    # init a corpus using column format, data folder and the names of the train, dev and test files
    
    try:
        corpus: Corpus = ColumnCorpus(data_folder, columns,
                                      train_file='train.txt',
                                      test_file='test.txt' )
                                      #dev_file='dev.txt')
    except: 
        print('Invalid input document in training')
        return 8

    # 2. what tag do we want to predict?
    tag_type = 'ner'

    #tag_dictionary = corpus.make_label_dictionary(label_type=tag_type)
    tag_dictionary = corpus.get_label_distribution()
    return tag_dictionary
    #return corpus
    

def upsampling_data(entities_to_upsample, probability,  entities):
    print('-'*20,'upsampling','-'*20)
    data_folder  = default_path + '/../../data/NER/train'
    columns = {'text':0, 'ner':1}
    for m in ["SiS","LwTR","MR","SR", "MBT"]:
        upsampler = upsampling_ner(data_folder+'/train.txt', entities+['O'], columns)
        data, data_labels = upsampler.get_dataset()
        new_samples, new_labels = upsampler.upsampling(entities_to_upsample,probability,[m])
        data += new_samples
        data_labels += new_labels

        with open(data_folder+'/train.txt', mode='w', encoding='utf-8') as f:
            for l,sentence in enumerate(data):
                for j,word in enumerate(sentence):
                    f.write(word+' '+ data_labels[l][j])
                    f.write('\n')
    
                if l < (len(data)-1):
                    f.write('\n')

    print('-'*20,'upsampling complete','-'*20)
    
    
def usage_cuda(cuda):
    if cuda:
        flair.device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
        if flair.device == torch.device('cpu'): return 'Error handling GPU, CPU will be used'
        elif flair.device == torch.device('cuda:0'): return 'GPU detected, GPU will be used'
    else:
        flair.device = torch.device('cpu')
        return 'CPU will be used'


def training_model(name, epochs=20):
    #FUNCION
    
    data_folder  = default_path + '/../../data/NER/train'
    path_model = default_path + '/../../models/NER/{}'.format(name)
    if (os.path.isdir(path_model)): print('WARNING, model already exists will be overwritten')
    columns = {0: 'text', 1:'ner'}
    # init a corpus using column format, data folder and the names of the train, dev and test files
    
    
    try:
        corpus: Corpus = ColumnCorpus(data_folder, columns,
                                      train_file='train.txt',
                                      test_file='test.txt' )
                                      #dev_file='dev.txt')
    except: 
        print('Invalid input document in training')
        return 8



            
    # 2. what tag do we want to predict?
    tag_type = 'ner'

    # 3. make the tag dictionary from the corpus
    #tag_dictionary = corpus.make_label_dictionary(label_type=tag_type)
    tag_dictionary = corpus.make_label_dictionary(label_type=tag_type)
            
    try:
        embeddings = TransformerWordEmbeddings(
                    model='xlm-roberta-large',
                    layers="-1",
                    subtoken_pooling="first", 
                    fine_tune=True,
                    use_context=True,
                )
    except: 
        print('Error while loading embeddings from RoBERTa')
        return 5

    # 5. initialize bare-bones sequence tagger (no CRF, no RNN, no reprojection)
    
    try:
        tagger_train = SequenceTagger(
            hidden_size=256,
            embeddings=embeddings,
            tag_dictionary=tag_dictionary,
            tag_type='ner',
            use_crf=False,
            use_rnn=False,
            reproject_embeddings=False,
        )
    except: 
        print('Error making tagger')
        return 6

    # 6. initialize trainer with AdamW optimizer
    

    trainer = ModelTrainer(tagger_train, corpus)

    # 7. run training with XLM parameters (20 epochs, small LR)
    try:
        trainer.train(path_model,
                      learning_rate=5.0e-6,
                      mini_batch_size=1,
                      mini_batch_chunk_size=1,
                      max_epochs=epochs,
                      scheduler=OneCycleLR,
                      embeddings_storage_mode='cpu',
                      optimizer=torch.optim.AdamW,
                      )
    except: 
        pass
        print('Error training the model, try setting CUDA False')
        return 7
    
    print("Model {} trained and saved in {}".format(name,'models/{}'.format(name)))
    
    
def tag_sentence(sentence, name):
    
    results={'Sentence_tagged':'', 'Highligth':{}}
    Highligth_dict={"text": "", "entities": []}
    
    
    #--------------Load the trained model-------------------------
    path_model = default_path + '/../../models/NER/{}'.format(name)
    global tagger_sentence
    
    if (not tagger_sentence):
        
        try:
            tagger_sentence = SequenceTagger.load(path_model+'/best-model.pt')
        except:
            try:
                tagger_sentence = SequenceTagger.load(path_model+'/final-model.pt')
            except: 
                print('Invalid model')
                return 1
        
    #------------------Tagged sentence---------------------
    print('-'*20,'Tagging','-'*20)
    sentence_f = Sentence(sentence)
    tagger_sentence.predict(sentence_f)
    sentence_tokenized = []
    Highligth_dict['text'] = sentence_f.to_plain_string()
    
    for indx,token in enumerate(sentence_f.tokens):
        
        t = token.get_label()
        if t.value == 'O':
            sentence_tokenized += [token.text]
        else: 
            sentence_tokenized += [t.shortstring]
            token_info={
                'entity': t.value ,
                'index' : indx,
                'word' : token.text,
                'start': token.start_position,
                'end' : token.end_position
                
                }
            Highligth_dict["entities"].append(token_info)
    sen_tagged = ' ' .join(sentence_tokenized)
    results['Highligth'] = Highligth_dict
    results['Sentence_tagged'] = sen_tagged
    print('-'*20,'Tagged complete','-'*20)
    return results
    
    
def use_model(name, path_data, output_dir):
    
    #--------------Load the trained model-------------------------
    path_model = default_path + '/../../models/NER/{}'.format(name)
    
    if not (os.path.isdir(path_model)): 
        print('Model does not exists')
        return 10
        
    if not os.path.isfile(path_data): 
        print('Input file is not a file')
        return 9 
    
    global tagger_document
    
    if (not tagger_document):
        
        try:
            tagger_document = SequenceTagger.load(path_model+'/best-model.pt')
        except:
            try:
                tagger_document = SequenceTagger.load(path_model+'/final-model.pt')
            except: 
                print('Invalid model')
                return 1
    
    #-----------------Load the document-------------------------
    try:
        data = pd.read_json(path_data, orient ='index', encoding='utf-8')[0]
    except: 
        print('Can\'t open the input file')
        return 2
    
    if len(data) <= 0:
        print(f"length of document greater than 0 expected, got: {len(data)}")
        return 2
    
    try:
        sentences=data['sentences']
        t = sentences[0]['text']
    except: 
        print('Invalid JSON format in document {}'.format(path_data))
        return 3
    print('-'*20,'Tagging','-'*20)
    
    
    
    #-----------------Tagged the document-------------------------
    results = {'text':"", 'text_labeled':"",'sentences':[], 'entities': []}
    indx_prev = 0
    pos_prev = 0
    for s in sentences:
        sentence = Sentence(s['text'])
        tagger_document.predict(sentence, mini_batch_size = 1)
        sen_dict_temp = {'text':sentence.to_plain_string(), 'text_labeled':'', 'tokens':[]}
        #return sentence
        sentence_tokenized = []
        for indx,token in enumerate(sentence.tokens):
            token_dict = {'text':token.text, 'label':token.get_label('ner').value}
            sen_dict_temp['tokens'].append(token_dict)
            
            t = token.get_label('ner')
            if t.value == 'O':
                sentence_tokenized += [token.text]
            else: 
                sentence_tokenized += [t.shortstring]
                token_info={
                    'entity': t.value ,
                    'index' : indx + indx_prev,
                    'word' : token.text,
                    'start': token.start_position + pos_prev,
                    'end' : token.end_position +pos_prev
                    
                    }
                results["entities"].append(token_info)
        indx_prev += len(sentence.tokens)
        pos_prev += len(sentence.to_plain_string())
        sen_tagged = ' ' .join(sentence_tokenized)
        sen_dict_temp['text_labeled'] = sen_tagged
        results['sentences'].append(sen_dict_temp)
        results['text'] += sentence.to_plain_string() 
        #return sentence
        results['text_labeled'] += sen_tagged
        
    #-----------------Save the results-------------------------
    try:
        with open(output_dir, "w", encoding='utf-8') as write_file:
            json.dump(results, write_file)
    
        print('-'*20,'Tagged complete','-'*20)
        print('Document tagged saved in {}'.format(output_dir))
    except:
        print('Error in output file')
        return 11
       
    return results

def json_to_txt(path_data_documents):
    #-------------List the documents in the path------------
    documents=os.listdir(path_data_documents)
    if len(documents) <= 0:
        print('There are not documents in the folder')
        return 4
    
    data_from_documents={'id':[],'document':[],'sentence':[],'word':[],'tag':[]}
    
    #--------------Verify each documment-------------
    for num,doc in enumerate(documents):
        data=path_data_documents+'/'+doc
        df = pd.read_json(data, orient ='index')[0]
        try:
            sentences = df['sentences']
            t = sentences[0]['text']
            t = sentences[0]['id']
            t = sentences[0]['tokens']
            j = t[0]['text']
            j = t[0]['begin']
            j = t[0]['end']
            tags = df['mentions']
            if tags:
                tg = tags[0]['id']
                tg = tags[0]['begin']
                tg = tags[0]['end']
                tg = tags[0]['type']
        except: 
            print('Invalid JSON input format in document {}'.format(doc))
            return 3
            
       
        #-----------------Organize the data----------------
        for s in sentences:
            id_senten=s['id']
            for tk in s['tokens']:
                if len(tk['text'])==1:
                    #if ord(tk['text'])>=48 and ord(tk['text'])<=57 and ord(tk['text'])>=65 and ord(tk['text'])<=90 and ord(tk['text'])>=97 and ord(tk['text'])<=122:
                    tk_beg=tk['begin']
                    tk_end=tk['end']
                    data_from_documents['id'].append('d'+str(num)+'_'+id_senten)
                    data_from_documents['document'].append(doc)
                    data_from_documents['word'].append(tk['text'])
                    data_from_documents['sentence'].append(s['text'])
                    data_from_documents['tag'].append('O')
                    for tg in tags:
                        if id_senten == tg['id'].split('-')[0] and tk['begin']>=tg['begin'] and tk['begin']<tg['end']:
                            data_from_documents['tag'][-1]=tg['type']
                            break
                        
                else:
                    tk_beg=tk['begin']
                    tk_end=tk['end']
                    data_from_documents['id'].append('d'+str(num)+'_'+id_senten)
                    data_from_documents['document'].append(doc)
                    data_from_documents['word'].append(tk['text'])
                    data_from_documents['sentence'].append(s['text'])
                    data_from_documents['tag'].append('O')
                    for tg in tags:
                        if id_senten == tg['id'].split('-')[0] and tk['begin']>=tg['begin'] and tk['begin']<tg['end']:
                            data_from_documents['tag'][-1]=tg['type']
                            break

    X=np.array(data_from_documents['word'])
    y=np.array(data_from_documents['tag'])     
    groups=np.array(data_from_documents['id'])  
    
    
    #-------------------Save the data in CONLL format--------------
    group_kfold = StratifiedGroupKFold(n_splits=10, shuffle=True, random_state=42)
    group_kfold.get_n_splits(X, y, groups)
    for train_index, test_index in group_kfold.split(X, y, groups):
        X_train, X_test = X[train_index], X[test_index]
        y_train, y_test = y[train_index], y[test_index]
        groups_train, groups_test = groups[train_index], groups[test_index]
        break


                    

    X_write=[X_train,X_test]
    y_write=[y_train,y_test]
    groups_write=[groups_train, groups_test]
    archivos=['train','test']
    
    
    for k in range(2):
        X_temp = X_write[k]
        y_temp = y_write[k]
        groups_temp = groups_write[k]
        arch=archivos[k]
        id_in=groups_temp[0]
        
            
        data_folder  = default_path + '/../../data/NER/train'
        check_create(data_folder)
        count = 0
        with open(data_folder + '/{}.txt'.format(arch), mode='w', encoding='utf-8') as f:
            for i in range(len(X_temp)):
                if groups_temp[i] != id_in:
                    id_in=groups_temp[i]
                    f.write('\n')
                    count = 0

                count += 1
                f.write(X_temp[i]+' '+ y_temp[i])
                f.write('\n')
                
                if count >= 150: 
                    count = 0
                    f.write('\n')

            

    # print("Before check")
    # checkpoint = "xlm-roberta-large"
    # config = AutoConfig.from_pretrained(checkpoint)
    
    # with init_empty_weights():
    #     model = AutoModelForSequenceClassification.from_config(config)
        
    # print("After check")    
    # try:
    #     tagger = load_checkpoint_and_dispatch(model, path_model+'/best-model.pt', device_map="auto")
    # except:
    #     try:
    #         tagger = load_checkpoint_and_dispatch(model, path_model+'/final-model.pt', device_map="auto")
    #     except: 
    #         print('Invalid model')
    #         return 1