import os
import re
import glob
import json
import collections
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
from wordcloud import WordCloud

import bigbench.api.util
import generate_task_summaries

sns.set(font_scale=1.6)
sns.set_style('whitegrid')
sns.set_palette(sns.color_palette("flare", 3))

basedir = generate_task_summaries.bench_dir
metadata = bigbench.api.util.get_all_task_metadata()


def is_subtask(task):
  """Whether the task name represents a subtask."""
  return ":" in task


def task_size(task):
  """Returns the total number of examples for the given task.

  The metadata count for a task includes the total over examples in subtasks.
  """
  return metadata[task]['num_generate_text'] + metadata[task]['num_multiple_choice']


def should_skip(task):
  """Whether the given task should be skipped when compiling statistics."""
  return task.startswith("simple_") or task == "training_on_test_set"


def generate_size_histogram():
  """Plot and save task size histogram.

  Considers supertasks, where the size of a supertask is the number of examples
  (including both generative and multiplep-choice) in the task and its subtasks.
  """
  task_sizes = [task_size(task) for task in metadata if not is_subtask(task)]

  for task, data in metadata.items():
    if is_subtask(task):
      continue
    samples = data['num_generate_text'] + data['num_multiple_choice']

  plt.figure(figsize=(10, 6))
  sns.histplot(task_sizes, log_scale=True)
  plt.xlabel("# Examples")
  plt.ylabel("# Tasks")
  plt.title("BIG-bench Task Sizes")
  hist_path = os.path.join(basedir, "task_size_histogram.pdf")
  plt.tight_layout()
  plt.savefig(hist_path)
  print(f"Saved size histogram to {hist_path}")


def normalize_keyword(keyword):
  keyword = re.sub(r" task$", r"", keyword)
  return keyword


def compute_keyword_frequencies():
  """Map keyword to number of tasks that have that keyword."""
  # Map task name to the set of its keywords
  task_keywords = {}

  for task, data in metadata.items():
    # Only collect keywords from parent tasks because
    # some subtasks contain common English word keywords
    if is_subtask(task):
      continue
    elif should_skip(task):
      continue

    task_keywords[task] = data["keywords"]

  keyword_freqs = collections.defaultdict(lambda: 0)
  common_keywords = [
    "zero-shot", "many-shot", "one-shot",
    "multiple choice", "multi-step", "json"]

  for task, keywords in task_keywords.items():
    for keyword in keywords:
      if keyword in common_keywords:
        continue
      normalized = normalize_keyword(keyword)
      keyword_freqs[normalized] += 1

  return keyword_freqs


def generate_keyword_frequency_table(num_keywords=20, reverse=False):
  """Print out a latex-format table with the top/bottom populated keywords."""
  keyword_freqs = compute_keyword_frequencies()
  sorted_keywords = sorted(
    keyword_freqs.items(),
    key=lambda item: item[1],
    reverse=reverse)
  sorted_keywords.reverse()

  qualifier = 'most'
  if reverse:
    qualifier = 'least'

  print(r"""
\begin{table}[]
\centering
\begin{tabular}{lc}
    \hline \noalign{\vskip 0.02in}
    \textbf{Keyword} & \textbf{Frequency} \\
    \hline \noalign{\vskip 0.02in}""")

  for keyword, freq in sorted_keywords[:num_keywords]:
    print(f"    {keyword} & {freq} \\\\")

  print(r"""    \hline
\end{tabular}
\caption{The """ + qualifier + r""" frequently used keywords in BIG-bench.}
\end{table}
""")


def generate_keyword_cloud():
  """Plot and save keyword cloud."""
  # Reduce the hierarchy in keyword frequencies to get a more balanced word cloud
  keyword_freqs = compute_keyword_frequencies()
  normalized_keyword_freqs = {
    keyword: np.power(freq, 0.5)
    for keyword, freq in keyword_freqs.items()
  }

  def black_color_func(word, font_size, position, orientation, random_state=None, **kwargs):
      return("hsl(0,100%, 1%)")

  #w=3000
  #h=1000

  w=3500
  h=2000

  wordcloud = WordCloud(
      font_path = '/usr/share/fonts/truetype/roboto/unhinted/RobotoCondensed-Medium.ttf',
      background_color="white",
      width=w,
      height=h,
      max_words=500).generate_from_frequencies(normalized_keyword_freqs)

  wordcloud.recolor(color_func = black_color_func)

  plt.figure(figsize=[15,10])
  plt.imshow(wordcloud, interpolation="bilinear")
  plt.axis("off")
  plt.tight_layout(pad=0)
  wordcloud_path = os.path.join(basedir, "keyword_cloud.pdf")
  plt.savefig(wordcloud_path, bbox_inches="tight", pad_inches=1)
  print(f"Saved word cloud to {wordcloud_path}")


#generate_size_histogram()
#generate_keyword_cloud()
generate_keyword_frequency_table(reverse=False)
generate_keyword_frequency_table(reverse=True, num_keywords=50)

print("")
num_tasks = len(bigbench.api.util.get_all_task_names())
num_json_tasks = len(bigbench.api.util.get_all_json_task_names())
num_prog_tasks = len(bigbench.api.util.get_all_programmatic_task_names())
print("Number of JSON tasks:", num_json_tasks)
print("Number of programmatic tasks:", num_prog_tasks)
print("Total number of tasks:", num_tasks)
if num_tasks != num_json_tasks + num_prog_tasks:
  print("WARNING: JSON + programmatic tasks is not equal to total number of tasks!")