import argparse

from timeit import default_timer as timer

from typing import List

import numpy as np

import pandas as pd

import torch

from flair.datasets import CONLL_03_DUTCH

from loguru import logger

from tqdm import tqdm

from deidentify.base import Document

from deidentify.taggers import CRFTagger, DeduceTagger, FlairTagger, TextTagger

from deidentify.tokenizer import TokenizerFactory

N_REPETITIONS = 5

N_SENTS = 5000

def load_data():

    corpus = CONLL_03_DUTCH()

    sentences = corpus.train[:N_SENTS]

    tokens = sum(len(sent) for sent in sentences)

    docs = [Document(name='', text=sent.to_plain_string(), annotations=[]) for sent in sentences]

    return docs, tokens

def benchmark_tagger(tagger: TextTagger, docs: List[Document], num_tokens: int):

    durations = []

    for _ in tqdm(range(0, N_REPETITIONS), desc='Repetitions'):

        start = timer()

        tagger.annotate(docs)

        end = timer()

        durations.append(end - start)

        if isinstance(tagger, FlairTagger) and torch.cuda.is_available():

            torch.cuda.empty_cache()

    return {

        'mean': np.mean(durations),

        'std': np.std(durations),

        'tokens/s': num_tokens / np.mean(durations),

        'docs/s': len(docs) / np.mean(durations),

        'num_docs': len(docs),

        'num_tokens': num_tokens

    }

def main(args):

    logger.info('Load data...')

    documents, num_tokens = load_data()

    logger.info('Initialize taggers...')

    tokenizer_crf = TokenizerFactory().tokenizer(corpus='ons', disable=())

    tokenizer_bilstm = TokenizerFactory().tokenizer(corpus='ons', disable=("tagger", "ner"))

    taggers = [

        ('DEDUCE', DeduceTagger(verbose=True)),

        ('CRF', CRFTagger(

            model='model_crf_ons_tuned-v0.1.0',

            tokenizer=tokenizer_crf,

            verbose=True

        )),

        ('BiLSTM-CRF (large)', FlairTagger(

            model='model_bilstmcrf_ons_large-v0.1.0',

            tokenizer=tokenizer_bilstm,

            mini_batch_size=args.bilstmcrf_large_batch_size,

            verbose=True

        )),

        ('BiLSTM-CRF (fast)', FlairTagger(

            model='model_bilstmcrf_ons_fast-v0.1.0',

            tokenizer=tokenizer_bilstm,

            mini_batch_size=args.bilstmcrf_fast_batch_size,

            verbose=True

        ))

    ]

    benchmark_results = []

    tagger_names = []

    for tagger_name, tagger in taggers:

        logger.info(f'Benchmark inference for tagger: {tagger_name}')

        scores = benchmark_tagger(tagger, documents, num_tokens)

        benchmark_results.append(scores)

        tagger_names.append(tagger_name)

    df = pd.DataFrame(data=benchmark_results, index=tagger_names)

    df.to_csv(f'{args.benchmark_name}.csv')

    logger.info('\n{}', df)

def arg_parser():

    parser = argparse.ArgumentParser()

    parser.add_argument("benchmark_name", type=str, help="Name of the benchmark.")

    parser.add_argument(

        "--bilstmcrf_large_batch_size",

        type=int,

        help="Batch size to use with the large model.",

        default=256

    )

    parser.add_argument(

        "--bilstmcrf_fast_batch_size",

        type=int,

        help="Batch size to use with the fast model.",

        default=256

    )

    return parser.parse_args()

if __name__ == '__main__':

    main(arg_parser())