import argparse
import os
import sys
import time
import json
import pickle

from nltk.corpus import wordnet as wn

import numpy as np
import torch
import random

from aligner import Aligner
import log

logger = log.get_logger('root')
logger.propagate = False

def get_print_result(sample_group: dict, sample_result: dict, nonce_word):
    candidates = sample_group['candidates']
    info = sample_group['common_ancestor_info']

    print_str = "\n===============================================================================================\n"
    print_str += f'Number of Candidates: {len(candidates)}\n\n'

    print_str += f"\nCommon {info['relation']} synset: {info['ancestor_name']}\n{wn.synset(info['ancestor_name']).definition()}\n\n"
    # Use the line below for CoDA5
    # print_str += f"\nCommon {info['relation']} synset: {info['ancestor_synset']}\n{info['ancestor_definition']}\n\n"

    for candidate_no, candidate in enumerate(candidates):
        print_str += f"\n{candidate_no+1}) Synset: {candidate['synset_name']} ({candidate['words_in_contexts'][0]})\n"
        print_str += f"Definition: {candidate['definition']}\n"
        print_str += f"Context: {candidate['contexts'][0]}\n"
               
    print_str += "\n\n"
    print_str += f"Predicted alignment: {sample_result['predicted_alignment']}\n"
    print_str += f"Alignment Score: {sample_result['alignment_score']}\n"

    return print_str
          
if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    # file parameters
    parser.add_argument('--output_folder', default=None, type=str, required=True)    
    parser.add_argument('--data_file', default=None, type=str, required=True)  
    
    # parameters for the model to generate definitions
    parser.add_argument('--model_cls', choices=['bert','roberta','gpt-2'], default='gpt-2')
    parser.add_argument('--model_name', type=str, default='gpt2-medium')
    
    parser.add_argument('--word_type', choices=['n','v'], default='n')
    
    parser.add_argument('--nonce_word', type=str, default='bkatuhla')
    parser.add_argument('--max_def_len', type=int, default=48,
                       help='maximum definition length after tokenization')
    parser.add_argument('--max_batch_size', type=int, default=48,
                       help='maximum batch size')
    
    parser.add_argument('--gpu_id', type=int, default=0,
                       help='id of the gpu that will be used during evaluations') 
    parser.add_argument('--seed', type=int, default=42,
                       help='seed for selecting random train samples for one of few shot evaluation') 
    
    args = parser.parse_args()
    random.seed(args.seed)
    
    with open(args.data_file, 'r') as handle:
        CoDA = json.load(handle)

    data = args.data_file.split("/")[-1][:-5] # don't take .json

    print_file = f'{args.model_name}_on_{data}_{args.word_type}_nonce_{args.nonce_word}_some_results.txt'
    save_file = f'{args.model_name}_on_{data}_{args.word_type}_nonce_{args.nonce_word}_results.pickle'

    if not os.path.exists(args.output_folder):
        os.mkdir(args.output_folder)
    f_out = open(os.path.join(args.output_folder, print_file), "w", encoding='UTF-8')
    f_out.close()    
    
    aligner = Aligner(
        model_cls=args.model_cls,
        pretrained_model=args.model_name,
        gpu_id=args.gpu_id,
        max_def_len=args.max_def_len,
        max_batch_size=args.max_batch_size,
        nonce_word=args.nonce_word,
        word_type=args.word_type
    )

    sample_result = {}
    sample_result[''] = []
    sample_result['target_scores'] = []
    sample_result['predicted_alignment'] = []
    sample_result['alignment_score'] = []

    all_results = []
    sample_groups = CoDA[args.word_type]
    for group_no, sample_group in enumerate(sample_groups):
        target_scores, predicted_alignment, aligment_score = aligner.align(sample_group)
        
        sample_result = {}
        sample_result['target_scores'] = target_scores
        sample_result['predicted_alignment'] = predicted_alignment
        sample_result['alignment_score'] = aligment_score

        all_results.append(sample_result)
        
        if (group_no+1) % 25 == 0:     
            logger.info(f'{group_no+1}/{len(sample_groups)} synset groups processed')
                    
        if (group_no+1) % (len(sample_groups) // 20) == 0: 
            with open(os.path.join(args.output_folder, print_file), "a", encoding='UTF-8') as f_out:
                f_out.write(get_print_result(sample_group, sample_result, args.nonce_word))

            with open(os.path.join(args.output_folder, save_file), "wb") as handle:
                pickle.dump(all_results, handle)         

    with open(os.path.join(args.output_folder, save_file), "wb") as handle:
        pickle.dump(all_results, handle) 
