import argparse
import math

from datetime import datetime
import logging
from pathlib import Path
import os
from tqdm import tqdm

import numpy as np

from transformers import AutoModelForTokenClassification
from transformers import Trainer, TrainingArguments
from datasets import load_metric

# local packages
from transformer_deid.evaluation import compute_metrics

from transformer_deid.train import which_transformer_arch
from transformer_deid.model_evaluation_functions import load_data

logging.basicConfig(
    format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
    datefmt='%m/%d/%Y %H:%M:%S',
    level=logging.INFO
)
logger = logging.getLogger(__name__)

multi_class_fields = [
    'AGEprecision', 'AGErecall', 'AGEf1', 'AGEnumber', 'CONTACTprecision',
    'CONTACTrecall', 'CONTACTf1', 'CONTACTnumber', 'DATEprecision',
    'DATErecall', 'DATEf1', 'DATEnumber', 'IDprecision', 'IDrecall', 'IDf1',
    'IDnumber', 'LOCATIONprecision', 'LOCATIONrecall', 'LOCATIONf1',
    'LOCATIONnumber', 'NAMEprecision', 'NAMErecall', 'NAMEf1', 'NAMEnumber',
    'PROFESSIONprecision', 'PROFESSIONrecall', 'PROFESSIONf1',
    'PROFESSIONnumber', 'overall_precision', 'overall_recall', 'overall_f1',
    'overall_accuracy'
]
binary_fields = [
    'PHIprecision', 'PHIrecall', 'PHIf1', 'PHInumber', 'overall_precision',
    'overall_recall', 'overall_f1', 'overall_accuracy'
]


def flatten_dict(d):
    """
    Return flattened version of the evaluation result dict
    """
    out = {}
    for key in d:
        if type(d[key]) is dict:
            child = flatten_dict(d[key])
            for child_key in child:
                val = child[child_key]
                if isinstance(val, np.int64):
                    val = int(val)
                out[key + child_key] = val
        else:
            out[key] = d[key]
    return out


def add_row(
    path, epochs, results_multiclass, results_binary, multi_class_fields,
    binary_fields, test_loss
):
    """
    Add row to worksheet
    fields: [epochs] + multi_class_fields + binary_fields
    """
    root = Path(path).parent

    row = [epochs] + [
        flatten_dict(results_multiclass).get(field)
        for field in multi_class_fields
    ] + [flatten_dict(results_binary).get(field)
         for field in binary_fields] + [test_loss]

    text_metrics = ','.join(map(str, row)) + '\n'

    with open(str(root) + '/training_eval.csv', 'at') as f:
        f.write(text_metrics)
    # worksheet.append_row(row, table_range='A1')


def eval_checkpoints(
    path, deid_task, train_dataset, val_dataset, test_dataset, training_args
):
    step = int(path.split('-')[-1])
    steps_per_epoch = math.ceil(
        len(train_dataset) / training_args.per_device_train_batch_size
    )
    epoch = step / steps_per_epoch

    model = AutoModelForTokenClassification.from_pretrained(
        path, num_labels=len(deid_task.labels)
    )

    model.eval()

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=val_dataset
    )

    predictions, labels, metrics = trainer.predict(test_dataset)
    predicted_label = np.argmax(predictions, axis=2)

    metric_dir = "transformer_deid/token_evaluation.py"
    metric = load_metric(metric_dir)

    results_multiclass = compute_metrics(
        predicted_label, labels, deid_task.labels, metric=metric
    )
    results_binary = compute_metrics(
        predicted_label,
        labels,
        deid_task.labels,
        metric=metric,
        binary_evaluation=True
    )

    add_row(
        path, epoch, results_multiclass, results_binary, multi_class_fields,
        binary_fields, metrics['test_loss']
    )


def parse_args():
    parser = argparse.ArgumentParser(
        description='Evaluate transformer-based model at each checkpoint.'
    )

    parser.add_argument(
        '-n',
        '--task_name',
        type=str,
        help=
        'name of folder containing train and test data; defaults to i2b2_2014',
        default='i2b2_2014'
    )

    parser.add_argument(
        '-m',
        '--model',
        type=str,
        help='folder containing checkpoint files',
        default='bert'
    )

    args = parser.parse_args()

    return args


def main():
    args = parse_args()
    root = f'{args.model}'
    arch = args.model.split('results')[0].lower()
    epochs = int(args.model.split('results')[1])
    task_name = args.task_name

    _, tokenizer, _ = which_transformer_arch(arch)

    dataDir = f'{task_name}'
    testDir = f'{task_name}/test'

    deid_task, train_dataset, val_dataset, test_dataset = load_data(
        task_name, dataDir, testDir, tokenizer
    )

    train_batch_size = 8

    training_args = TrainingArguments(
        output_dir=root,
        num_train_epochs=epochs,
        per_device_train_batch_size=train_batch_size,
        per_device_eval_batch_size=8,
        warmup_steps=500,
        weight_decay=0.01,
        logging_dir='./logs',
        logging_steps=10,
        save_strategy='steps',
        eval_steps=1155
    )
    
    if not os.path.exists(str(root) + '/training_eval.csv'):
        with open(str(root) + '/training_eval.csv', 'wt') as f:
            header = 'epoch,' + ','.join(
                map(str, multi_class_fields + binary_fields + ['test_loss'])
            ) + '\n'
            f.write(header)

    checkpoints = [
        item for item in os.listdir(root)
        if 'checkpoint' in item and os.path.isdir(os.path.join(root, item))
    ]

    for item in tqdm(sorted(checkpoints, key=lambda x: int(x.split('-')[1]))):
        path = os.path.join(root, item)
        eval_checkpoints(
            path, deid_task, train_dataset, val_dataset, test_dataset,
            training_args
        )


if __name__ == '__main__':
    main()