r"""This is an example colab which demonstrates how to load a book or article
from a dataset, run two different models on the article, compare the results,
and inspect the context in which tokens occur. All cells should be run in order.

The citc client, gin configurations and pre-trained model directories should
be set to the appropriate values by the user.
"""

from http.client import UnimplementedFileMode
import sys
from datetime import datetime

from typing import Sequence
import jax.numpy as jnp
from flax.training import common_utils

from absl import app
from absl import flags
from absl import logging
import time
import argparse
import glob

import gin
import jax
import numpy as np
import tensorflow.compat.v2 as tf

# import matplotlib.pyplot as plt
import bokeh.plotting as bplt

# ---- Change this to the appropriate user and client. ----

from transformer import decoder_stack
from transformer import inference_utils
from transformer import text_dataset

from transformer import tasks
import json

from absl import flags
flags.FLAGS([''])
import jax
import numpy as np


import seqio
import functools

from bigbench.bbseqio import task_api
from bigbench.bbseqio import tasks as bbtasks
logging.set_verbosity(logging.WARNING)

# caching info
from transformer.cache_bigbench_tasks import register_cache_tasks

import logging
logging.getLogger().setLevel(logging.INFO)

def load_model_configs(model_name, configs, sample_method = "greedy", batch_size = 2, sequence_length = 512):
    # googlelog.set_global_capture(True)   # Uncomment to see all logs.
    print("loading model ", model_name)
    gin.clear_config(clear_constants=True)

    gin.enter_interactive_mode()  # Avoid errors when reloading cells.
    gin_paths=["transformer/configs"]
        
    gin_files = configs[model_name]["gin_files"]
    print("gin_files: ", gin_files)
    
    load_dir = configs[model_name]["load_dir"]
    print("load_dir: ", load_dir)
    
    # Override the task batch size and set it to 1.
    # This may require creating new Transformer-XL state, so ignore the pre-trained state.
    # Ask the model to output separate losses per token.
    gin_params = [
     f"DecoderOnlyLanguageModel.sample_method=\"{sample_method}\"",
      "DecoderOnlyLanguageModel.output_token_losses=True",
      "DecoderOnlyLanguageModel.output_logits=True",
      f"TransformerTaskConfig.batch_size={batch_size}",
      f"TransformerTaskConfig.sequence_length={sequence_length}",
      f"TransformerLayer.window_length={sequence_length}",
      "Trainer.restore_state_variables=False",
    ]

    # Parse the gin files and parameters.
    # If the config is not unlocked, then this command will fail if it is run a second time.
    with gin.unlock_config():
      inference_utils.parse_gin_configuration(gin_files, gin_params, gin_paths=gin_paths)

    article_data = inference_utils.read_article(verbose=True)

    (article_block_list, vocab) = article_data

    batch_idx = 0
    print(text_dataset.pretty_print_article(article_block_list[batch_idx], {"targets": vocab}, 32768))

    return vocab, load_dir

from tqdm import tqdm
import copy

def predict_fn(gen_task, gen_task_state, batch_size, rep, sequence_length, do_generate, ds: tf.data.Dataset):
    all_inferences = []
    all_indices = []
    task_state = copy.deepcopy(gen_task_state)
    
    try:
      original_ds_length = len(ds)
      dataset_remainder = original_ds_length % (batch_size * rep)  # pytype:disable=wrong-arg-types
      print('length of dataset = %s', len(ds))
    except TypeError as e:
      if str(e) == 'dataset length is unknown.':
        logging.warning(
            'The following error is likely due to the use of TensorFlow v1 in '
            'your dataset pipeline. Verify you are not importing from '
            '`tf.compat.v1` as part of your pipeline.')
      raise e
    
    if dataset_remainder:
      dataset_pad_amt = batch_size * rep - dataset_remainder
      print(
          'Padding infer dataset with %d examples for even per-replica shards.',
          dataset_pad_amt)
      # Pad with the first example using an index of -1 so seqio will ignore.
      pad_ds = ds.take(1).map(lambda i, x: (np.int64(-1), x)).repeat(
          dataset_pad_amt)
      ds = ds.concatenate(pad_ds)
    
    for batch_indice, x in tqdm(ds.batch(batch_size* rep, drop_remainder=True)):
        prompt_tokens = np.zeros((batch_size * rep, sequence_length), dtype=np.int32)
        prompt_tokens[:, :x['decoder_input_tokens'].numpy().shape[1]] = x['decoder_input_tokens'].numpy()[:, :sequence_length]
        
        start_of_sequence = np.ones(batch_size * rep , dtype=np.int32)

        loss_mask = np.zeros((batch_size * rep, sequence_length), dtype=np.int32)
        loss_mask[:, :x['decoder_loss_weights'].numpy().shape[1]] = x['decoder_loss_weights'].numpy()[:, :sequence_length]
        # loss_mask = x['decoder_loss_weights'].numpy()[:, :sequence_length]
        # loss_mask = np.ones([task_config.batch_size, task_config.sequence_length], dtype=np.int32)
        # loss_mask[: , : len(prompt_tokens) - 1] = 0
        model_x = {"targets": prompt_tokens, "start_of_sequence": start_of_sequence, "loss_mask": loss_mask, "epoch": np.array([0]*rep)[:, None], "nucleus_cutoff": np.array([0.9]*rep)[:, None], "temperature": np.array([1]*rep)[:, None]}
        # print([v.shape for k, v in model_x.items()])
        out, tstate = inference_utils.run_model(gen_task, task_state, ([model_x], vocab), verbose=False, return_tstate = True)
        task_state = (tstate, task_state[1])
        out = out[0]
        
        if do_generate:
          gen_tokens = out["gen_tokens"].reshape(batch_size * rep, -1)
          # todo: fix this!
          pred = jax.vmap(tasks.get_masked_tokens)(gen_tokens, loss_mask)
        else:
          # get score
          logits = out["logits"].reshape(batch_size * rep, sequence_length, vocab.vocab_size)
          # soft_targets = common_utils.onehot(prompt_tokens, logits.shape[-1])
          # pred = (logits * soft_targets * (x["decoder_loss_weights"].numpy() == 1)[:, :, None]).sum(-1).sum(-1)
          pred = []
          for idx in range(batch_size * rep):
            pred.append((logits[idx][range(sequence_length), prompt_tokens[idx]] * (loss_mask[idx] == 1)).sum())
          pred = jnp.stack(pred)

        all_inferences.append(pred)
        all_indices.append(batch_indice)
    all_inferences = np.concatenate(all_inferences)
    all_indices = np.concatenate(all_indices)
    non_pad_idxs = all_indices >= 0
    all_indices = all_indices[non_pad_idxs]
    all_inferences = jax.tree_map(lambda x: x[non_pad_idxs], all_inferences)
        
    indices_and_outputs = list(zip(all_indices, all_inferences))
    return indices_and_outputs



def evaluate_one_task(task_name, gen_task, gen_task_state, task, task_state, batch_size, sequence_length, replicate_mode, features, use_cached = False):
    evaluator = seqio.Evaluator(
          mixture_or_task_name=task_name,
          feature_converter=seqio.DecoderFeatureConverter(pack=False),  # pytype:disable=not-instantiable
          eval_split="all",
          sequence_length=None,
          num_examples=256,
          shuffle=True,
          use_cached=use_cached,
    )

    # remove empty tasks
    empty_tasks = []
    for ds_name, ds in evaluator.cached_task_datasets.items():
      print(ds_name, len(ds))
      if len(ds) == 0:
        empty_tasks.append(ds_name)
    evaluator._eval_tasks = list(filter(lambda t: t.name not in empty_tasks, evaluator.eval_tasks))
    rep = jax.local_device_count() if replicate_mode else 1
    all_metrics, _, _ = evaluator.evaluate(
            compute_metrics=jax.process_index() == 0,
            predict_fn=functools.partial(predict_fn, gen_task, gen_task_state, batch_size, rep, sequence_length, True),
            score_fn=functools.partial(predict_fn, task, task_state, batch_size, rep, sequence_length, False)
            )
    return all_metrics.result()

def define_bigbench_eval_tasks(sequence_length, vocab, vocab_name, shot):
  # add cached tasks to the registry
  all_modified_and_cached_task_names =  []
  for s in shot:
    all_modified_and_cached_task_names += register_cache_tasks(sequence_length, s, vocab, vocab_name, have_all = True)
  print("@debug: example of cached_task_names", all_modified_and_cached_task_names[:4])


  
  s = set([ tuple(x.split(".")[:2]) for x in all_modified_and_cached_task_names ] )
  un = {}
  for x in s:
    if x[0] not in un:
      un[x[0]] = x[1]
    else:
      if x[1] == "gen":
        un[x[0]] = x[1]
      elif x[1] == "mul" and un[x[0]] == "scr":
        un[x[0]] = x[1]
  selected_tasks = sorted([ ".".join(x) for x in un.items()])
  selected_tasks = [
    "bigbench:arithmetic.gen"
  ]

  eval_tasks = list(filter(lambda x: any(x.startswith(t) for t in selected_tasks), all_modified_and_cached_task_names))  
  
  unselected_tasks = [
    "bigbench:authorship_verification",
    "bigbench:checkmate_in_one",
    "bigbench:chinese_remainder_theorem",
    "bigbench:color",
    "bigbench:discourse_marker_prediction",
    "bigbench:formal_fallacies_syllogisms_negation",
    "bigbench:hhh_alignment",
    "bigbench:kanji_ascii",
    "bigbench:hhh_alignment",
    "bigbench:kannada",
    "bigbench:key_value_maps",
    "bigbench:language_games",
    "bigbench:mathematical_induction",
    "bigbench:minute_mysteries_qa",
    "bigbench:misconceptions",
    "bigbench:misconceptions_russian",
    "bigbench:mnist_ascii",
    "bigbench:navigate",
    "bigbench:play_dialog_same_or_different",
    "bigbench:presuppositions_as_nli",
    "bigbench:real_or_fake_text",
    "bigbench:salient_translation_error_detection",
    "bigbench:semantic_parsing_in_context_sparc",
    "bigbench:semantic_parsing_spider",
    "bigbench:simple_text_editing",
    "bigbench:sudoku",
    "bigbench:symbol_interpretation",
    "bigbench:talkdown",
    "bigbench:tense",
    "bigbench:text_navigation_game",
    "bigbench:topical_chat",
    "bigbench:tracking_shuffled_objects", 
    "bigbench:twenty_question",
    "bigbench:web_of_lies",
    "bigbench:which_wiki_edit",
    "bigbench:winowhy",
    "bigbench:word_problems_on_sets_and_graphs",
    # other
    "bigbench:few_shot_nlg",
    "bigbench:long_context_integration",
    "bigbench:medical_questions_russian",
    "bigbench:known_unknowns",
    "bigbench:suicide_risk",
    "bigbench:what_is_the_tao",
    # not corr
    "bigbench:boolean_expressions", 
    "bigbench:crash_blossom",   
    "bigbench:dynamic counting",
    "bigbench:entailed_polarity_hindi",
    "bigbench:epistemic_reasoning",
    "bigbench:fantasy_reasoning",
    "bigbench:identify_math_theorems",
    "bigbench:intersect_geometry,",
    "bigbench:epistemic_reasoning",
    "bigbench:persian_idioms",
    "bigbench:scientific_press_release",
    "bigbench:social_support",
    "bigbench:dark_humor_detection",
    "bigbench:moral_permissibility",
    "bigbench:ruin_names",
    "bigbench:bbq_lite",
    "bigbench:movie_recommendation",
    "bigbench:physics_questions",
    "bigbench:arithmetic"
  ]

  eval_tasks = list(filter(lambda x: not any(x.startswith(t) for t in unselected_tasks), eval_tasks))  
  print(eval_tasks)
  # determine which tasks are cached in the registry AND have existing cache dirs
  # cached_eval_tasks = list(filter(tasks.has_cache_dir, eval_tasks))
  # print("@debug: len(cached_eval_tasks)", len(cached_eval_tasks))
  
  # non_cached_eval_tasks = [x for x in eval_tasks if x not in cached_eval_tasks]
  # print("@debug: len(non_cached_eval_tasks)", len(non_cached_eval_tasks))
  cached_eval_tasks = []
  non_cached_eval_tasks = eval_tasks

  # drop tasks that have zero length (esp. important bc cached tasks get_dataset does not work if empty)
  # cached_eval_tasks = list(filter(tasks.is_nonzero_length_cached_dataset, cached_eval_tasks))
  # non_cached_eval_tasks = list(filter(tasks.is_nonzero_length_dataset, non_cached_eval_tasks))
  # TODO: uncomment above line; currently not yet cached, so would be too slow

  # modify uncached tasks (cached tasks should already be modified to include filter_too_long)
  features_all = {}  
  for eval_task_name in cached_eval_tasks + non_cached_eval_tasks:
    t = seqio.TaskRegistry.get(eval_task_name)
    features_all[eval_task_name] = list(t.output_features.keys())

  # seqio.MixtureRegistry.add(
  #   "cached_bigbench_eval_mix",
  #   zip(cached_eval_tasks, [1] * len(cached_eval_tasks)),
  # )

  return cached_eval_tasks + non_cached_eval_tasks, features_all, cached_eval_tasks, non_cached_eval_tasks



if __name__ == "__main__":
    print(jax.devices())
    assert jax.local_device_count() == 8
    text_dataset.set_default_data_directory()

    with open("model_configs.json") as f:
        configs = json.load(f)
     
    parser = argparse.ArgumentParser()
    parser.add_argument('--vocab', type=str, default="64k")
    parser.add_argument('--batch_size', type=int, default=8)
    parser.add_argument('--no_replicate_mode', type=bool, default=False)
    parser.add_argument('--sequence_length', type=int, default=512)
    parser.add_argument('--n_machines', type=int, default=1)
    parser.add_argument('--machine_id', type=int, default=0)
    parser.add_argument('--shot', nargs='+', type=int, default=[0,1,2,3])
    args = parser.parse_args()
    print(args)

    # models_to_eval = ["pile-filtered-150-64k-full"]
    models_to_eval = ["pile-filtered-150-64k-rep-only", "pile-filtered-150-64k-rep-20-only", "pile-filtered-150-64k-listy-only", "pile-filtered-150-64k-math-only", "pile-filtered-150-64k-github-only"]
    #  , "pile-bb-150-64k", "pile-filtered-150-64k", "pile-filtered-5-150-64k"]


    if args.vocab == "t5":
      eval_tasks, features_all, cached_eval_tasks, non_cached_eval_tasks = define_bigbench_eval_tasks(args.sequence_length, tasks.T5_DEFAULT_VOCABULARY, "t5_default", args.shot)
    elif args.vocab == "gpt2":
      eval_tasks, features_all, cached_eval_tasks, non_cached_eval_tasks = define_bigbench_eval_tasks(args.sequence_length, tasks.GPT2_VOCABULARY, "gpt2", args.shot)
    elif args.vocab == "64k":
      eval_tasks, features_all, cached_eval_tasks, non_cached_eval_tasks = define_bigbench_eval_tasks(args.sequence_length, tasks.LARGER_VOCAB, "64k", args.shot)
    else:
      raise NotImplementedError
   

    for model_name in models_to_eval:
        print("==========", model_name,"==========")
        vocab, load_dir = load_model_configs(model_name, configs, batch_size = args.batch_size, sequence_length = args.sequence_length)    
        gen_task, gen_task_state, _ = inference_utils.create_model_and_task(vocab, load_dir=load_dir, task_mode="generate", replicate_mode = not args.no_replicate_mode)
        task, task_state, _ = inference_utils.create_model_and_task(vocab, load_dir=load_dir, replicate_mode = not args.no_replicate_mode)
        start_time = time.time()
        for eval_task_name in tqdm(list(eval_tasks)[args.machine_id::args.n_machines]):
          # Check if there are any json files in the current directory with eval_task_name in their name
          # If so, skip this task
          if len(glob.glob(f"all_metrics/{model_name}*{eval_task_name}.json")) > 0 or  len(glob.glob(f"eval_files/{model_name}*{eval_task_name}.json")) > 0:
            print(f"Skipping {model_name}'s {eval_task_name} because it has already been evaluated")
            continue
          if True:
            task_is_cached = eval_task_name in cached_eval_tasks
            print("@debug: eval_task_name", eval_task_name)
            print("@debug: task_is_cached", task_is_cached)
            result = evaluate_one_task(
              eval_task_name, gen_task, gen_task_state, task, task_state, 
              batch_size = args.batch_size, sequence_length = args.sequence_length, replicate_mode = not args.no_replicate_mode, 
              features = features_all[eval_task_name],
              use_cached = task_is_cached)
            print(result)
            with open("eval_files/" + model_name + "_"  +  datetime.now().strftime("%d_%m_%Y_%H_%M_%S") + "_" + eval_task_name + ".json", "w") as outfile:
              json.dump(result, outfile)
          # except Exception as e:
          #   print("@debug: exception!!")
          #   print(e)
          #   time.sleep(1)
        print("@debug: loop took seconds:", time.time() - start_time)
        del gen_task, gen_task_state, task, task_state

        
    print("FINISHED!!!")
