"""
Given two ensembles and a tokenized file, output the trees for which those ensembles agree and report how many of the sub-models agree on those trees.

For example:

python3 -m stanza.utils.datasets.constituency.build_silver_dataset --tokenized_file /u/nlp/data/constituency-parser/italian/2024_wiki_tokenization/it_wiki_tokenized_AA.txt --lang it --output_file asdf.out --e1 saved_models/constituency/it_vit_electra_100?_top_constituency.pt --e2 saved_models/constituency/it_vit_electra_100?_constituency.pt

for i in `echo f g h i j k l m n o p q r s t`; do nlprun -d a6000 "python3 -m stanza.utils.datasets.constituency.build_silver_dataset --tokenized_file /u/nlp/data/constituency-parser/italian/2024_wiki_tokenization/it_wiki_tok_6M_a$i.txt --lang it --output_file /u/nlp/data/constituency-parser/italian/2024_it_vit_electra/a$i.trees --e1 saved_models/constituency/it_vit_electra_100?_top_constituency.pt --e2 saved_models/constituency/it_vit_electra_100?_constituency.pt" -o /u/nlp/data/constituency-parser/italian/2024_it_vit_electra/a$i.out; done

for i in `echo a b c d`; do nlprun -d a6000 "python3 -m stanza.utils.datasets.constituency.build_silver_dataset --tokenized_file /u/nlp/data/constituency-parser/english/en_wiki_2023/shuf_1M.a$i --lang en --output_file /u/nlp/data/constituency-parser/english/2024_en_ptb3_electra/forward_a$i.trees --e1 saved_models/constituency/en_ptb3_electra-large_100?_in_constituency.pt --e2 saved_models/constituency/en_ptb3_electra-large_100?_top_constituency.pt" -o /u/nlp/data/constituency-parser/english/2024_en_ptb3_electra/forward_a$i.out; done
"""

import argparse
import json

import logging

from stanza.models.common import utils
from stanza.models.common.foundation_cache import FoundationCache
from stanza.models.constituency import retagging
from stanza.models.constituency import text_processing
from stanza.models.constituency import tree_reader
from stanza.models.constituency.ensemble import Ensemble
from stanza.utils.get_tqdm import get_tqdm

tqdm = get_tqdm()

logger = logging.getLogger('stanza.constituency.trainer')

def parse_args(args=None):
    parser = argparse.ArgumentParser(description="Script that uses multiple ensembles to find trees where both ensembles agree")

    input_group = parser.add_mutually_exclusive_group(required=True)
    input_group.add_argument('--tokenized_file', type=str, default=None, help='Input file of tokenized text for parsing with parse_text.')
    input_group.add_argument('--tree_file', type=str, default=None, help='Input file of already parsed text for reparsing with parse_text.')
    parser.add_argument('--output_file', type=str, default=None, help='Where to put the output file')

    parser.add_argument('--charlm_forward_file', type=str, default=None, help="Exact path to use for forward charlm")
    parser.add_argument('--charlm_backward_file', type=str, default=None, help="Exact path to use for backward charlm")
    parser.add_argument('--wordvec_pretrain_file', type=str, default=None, help='Exact name of the pretrain file to read')

    utils.add_device_args(parser)

    parser.add_argument('--lang', default='en', help='Language to use')

    parser.add_argument('--eval_batch_size', type=int, default=50, help='How many trees to batch when running eval')
    parser.add_argument('--e1', type=str, nargs='+', default=None, help="Which model(s) to load in the first ensemble")
    parser.add_argument('--e2', type=str, nargs='+', default=None, help="Which model(s) to load in the second ensemble")

    parser.add_argument('--mode', default='predict', choices=['parse_text', 'predict'])

    # another option would be to include the tree idx in each entry in an existing saved file
    # the processing could then pick up at exactly the last known idx
    parser.add_argument('--start_tree', type=int, default=0, help='Where to start... most useful if the previous incarnation crashed')
    parser.add_argument('--end_tree', type=int, default=None, help='Where to end.  If unset, will process to the end of the file')

    retagging.add_retag_args(parser)

    args = vars(parser.parse_args())

    retagging.postprocess_args(args)
    args['num_generate'] = 0

    return args

def main():
    args = parse_args()
    utils.log_training_args(args, logger, name="ensemble")

    retag_pipeline = retagging.build_retag_pipeline(args)
    foundation_cache = retag_pipeline[0].foundation_cache if retag_pipeline else FoundationCache()

    logger.info("Building ensemble #1 out of %s", args['e1'])
    e1 = Ensemble(args, filenames=args['e1'], foundation_cache=foundation_cache)
    e1.to(args.get('device', None))
    logger.info("Building ensemble #2 out of %s", args['e2'])
    e2 = Ensemble(args, filenames=args['e2'], foundation_cache=foundation_cache)
    e2.to(args.get('device', None))

    if args['tokenized_file']:
        tokenized_sentences, _ = text_processing.read_tokenized_file(args['tokenized_file'])
    elif args['tree_file']:
        treebank = tree_reader.read_treebank(args['tree_file'])
        tokenized_sentences = [x.leaf_labels() for x in treebank]
        if args['lang'] == 'vi':
            tokenized_sentences = [[x.replace("_", " ") for x in sentence] for sentence in tokenized_sentences]
    logger.info("Read %d tokenized sentences", len(tokenized_sentences))

    all_models = e1.models + e2.models

    chunk_size = 1000
    with open(args['output_file'], 'w', encoding='utf-8') as fout:
        end_tree = len(tokenized_sentences) if args['end_tree'] is None else args['end_tree']
        for chunk_start in tqdm(range(args['start_tree'], end_tree, chunk_size)):
            chunk = tokenized_sentences[chunk_start:chunk_start+chunk_size]
            logger.info("Processing trees %d to %d", chunk_start, chunk_start+len(chunk))
            parsed1 = text_processing.parse_tokenized_sentences(args, e1, retag_pipeline, chunk)
            parsed1 = [x.predictions[0].tree for x in parsed1]
            parsed2 = text_processing.parse_tokenized_sentences(args, e2, retag_pipeline, chunk)
            parsed2 = [x.predictions[0].tree for x in parsed2]
            matching = [t for t, t2 in zip(parsed1, parsed2) if t == t2]
            logger.info("%d trees matched", len(matching))
            model_counts = [0] * len(matching)
            for model in all_models:
                model_chunk = model.parse_sentences_no_grad(iter(matching), model.build_batch_from_trees, args['eval_batch_size'], model.predict)
                model_chunk = [x.predictions[0].tree for x in model_chunk]
                for idx, (t1, t2) in enumerate(zip(matching, model_chunk)):
                    if t1 == t2:
                        model_counts[idx] += 1
            for count, tree in zip(model_counts, matching):
                line = {"tree": "%s" % tree, "count": count}
                fout.write(json.dumps(line))
                fout.write("\n")


if __name__ == '__main__':
    main()
