"""
This script loads a binary / weighted connectome from a TensorFlow checkpoint
file (see model.py) and exports it as HDF5 file.

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


# HACKHACKHACK(amotta): This, too, should be loaded from the checkpoint file.
# Unfortunately, I forgot to add it to the list of stored variables...
N = 2000
N_INH = 200
N_EXC = N - N_INH

a = slice(None, None)
e = slice(None, N_EXC)
i = slice(-N_INH, None)


def load_connectome_from_ckpt(ckpt_file):
    logging.log(logging.INFO, "Processing \"%s\"", ckpt_file)
    reader = tf.train.NewCheckpointReader(ckpt_file)

    mask = reader.get_tensor('rnn/mask')
    weights = reader.get_tensor('rnn/weights')

    # Sanity checks
    assert mask.dtype is np.dtype(bool)
    assert not np.any(weights[np.logical_not(mask)])

    return (mask, weights)


def process_ckpt_file(ckpt_file, out_file):
    # NOTE(amotta): For convenience we accept checkpoint file paths that are
    # actually paths to the corresponding TensorFlow `.index` or `.meta` files.
    # This allows us to call the script like so:
    # python evaluate.py checkpoint /path/to/checkpoints/*.index
    ckpt_ext = os.path.splitext(ckpt_file)
    if ckpt_ext[-1] in ['.index', '.meta']: ckpt_file = ckpt_ext[0]

    mask, weights = load_connectome_from_ckpt(ckpt_file)

    types = np.zeros((N,), dtype=np.uint8)
    types[e] = 1
    types[i] = 2

    # Sanity check
    assert np.all(types)
    assert types.size == mask.shape[0]
    assert types.size == mask.shape[1]

    with h5py.File(out_file, 'w') as out:
        out.create_dataset('mask', data=np.uint8(mask))
        out.create_dataset('weights', data=weights)
        out.create_dataset('types', data=types)


def process_ckpt_files(args):
    for ckpt_file in args.ckpt_files:
        ckpt_name = os.path.basename(ckpt_file)
        ckpt_name = os.path.splitext(ckpt_name)[0]

        out_file = ckpt_name + '.hdf5'
        out_file = os.path.join(args.out_dir, out_file)

        process_ckpt_file(ckpt_file, out_file)


def build_arg_parser():
    p = argparse.ArgumentParser()
    p.add_argument('-o', '--out-dir')
    p.add_argument('ckpt_files', nargs=argparse.REMAINDER)
    return p


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

    process_ckpt_files(args)


if __name__ == '__main__':
    main()
