"""
Given a path prefix PREFIX of directories containing files

config.json
dev/{summary.json,average_activations.npy,average_norms.npy}
train/summary.json,average_activations.npy,average_norms.npy}

In the directory associated with the prefix, outputs the following files

dev_loss.pdf
train_loss.pdf
dev_train_gap.pdf

Plots for all available attention types available in the prefix (e.g.,
soft, topk, topk-50) at all available k values. The plot is a line plot
of the corresponding loss for the different attention type at different
k values (soft is assumed to have k=0).

Similarly, plots the activation and norms distributions for the marginal
activation block aggregates into

train_act.pdf
train_nrm.pdf
dev_act.pdf
dev_nrm.pdf

All results should be from the same task for viz to make sense.
"""

from absl import app, flags

from ..motivation.bert_agg import main
from ..params import GLUE_TASK_NAMES
from .. import log

flags.DEFINE_string("prefix", None, "prefix directory")

flags.DEFINE_string(
    "cache", None, "cache directory (autogenerated based on prefix)"
)

flags.DEFINE_bool("overwrite", False, "overwrite previous directory files")

flags.DEFINE_enum("task", None, GLUE_TASK_NAMES, "BERT fine-tuning task")


def _main(_argv):

    log.init()
    main(flags.FLAGS.prefix, flags.FLAGS.cache, flags.FLAGS.overwrite, flags.FLAGS.task)

if __name__ == "__main__":
    flags.mark_flag_as_required("task")
    flags.mark_flag_as_required("prefix")
    app.run(_main)
