"""
This scripts loads the training and validation loss and accuracy
values from the TensorFlow event files and converts them to a
(slightly weird) CSV file for further processing in MATLAB.

Written by
  Alessandro Motta <alessandro.motta@brain.mpg.de>
"""
import argparse
import csv
import logging
import os
import tensorflow as tf

from collections import OrderedDict
from glob import glob


def process_event_file(event_file_path, keys):
    logging.log(logging.INFO, "Processing \"%s\"", event_file_path)
    values = []

    for event in tf.train.summary_iterator(event_file_path):
        cur_values = {'step': event.step}

        for kv_pair in event.summary.value:
            key = kv_pair.tag.split('/')[-1]
            if not key in keys: continue
            value = kv_pair.simple_value
            cur_values[key] = value

        values.append(cur_values)

    return values


def process_training_run(args):
    keys = ['step', 'weight_count', 'loss', 'accuracy',
            'loss_valid', 'accuracy_valid']

    in_files = glob(os.path.join(args.in_dir, '*', 'events.*'))
    values = list(map(lambda f: process_event_file(f, keys), in_files))
    values = sum(values, [])

    logging.log(logging.INFO, "Writing \"%s\"", args.out_file)
    with open(args.out_file, 'w', newline='') as f:
        writer = csv.DictWriter(f, fieldnames=keys)
        writer.writeheader()
        writer.writerows(values)


def build_arg_parser():
    p = argparse.ArgumentParser()
    p.add_argument('-i', '--in-dir')
    p.add_argument('-o', '--out-file')
    return p


def main():
    logging.basicConfig(level=logging.INFO)
    arg_parser = build_arg_parser()
    args = arg_parser.parse_args()

    process_training_run(args)


if __name__ == '__main__':
    main()
