"""
This script loads a binary connectome from a TensorFlow checkpoint file (see
model.py) and computes the following network measures / circuit statistics /
graph features proposed by Klinger et al.:

- (relative) reciprocity within and across neuronal subpopulations,
- (relative) recurrency within neuronal subpopulations, and
- input / output degree correlation within neuronal subpopulations

The code is based on the discriminatEM package by Emmanuel Klinger.

Written by
  Alessandro Motta <alessandro.motta@brain.mpg.de>
"""
import argparse
import csv
import itertools
import logging
import numpy as np
import os
import sys
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 connectivity_estimator(conn):
    return {'aa': conn[a, a].sum() / (N * (N - 1)),
            'ee': conn[e, e].sum() / (N_EXC * (N_EXC - 1)),
            'ii': conn[i, i].sum() / (N_INH * (N_INH - 1)),
            'ei': conn[e, i].sum() / (N_EXC * N_INH),
            'ie': conn[i, e].sum() / (N_INH * N_EXC)}


def reciprocity_estimator(conn):
    def _reciprocity(pre, post):
        conn_fwd = conn[pre, post]
        conn_rec = conn[post, pre].T

        fwd_count = conn_fwd.sum()
        rec_count = (conn_fwd * conn_rec).sum()

        if fwd_count == 0: return 0.0
        return rec_count / fwd_count

    return {'aa': _reciprocity(a, a),
            'ee': _reciprocity(e, e),
            'ii': _reciprocity(i, i),
            'ei': _reciprocity(e, i),
            'ie': _reciprocity(i, e)}


def relative_reciprocity_estimator(conn):
    def _relative_reciprocity(r, p):
        if r == 0: return 0.0
        if p == 0: raise Exception("Division by zero")
        return r / p

    reciprocity = reciprocity_estimator(conn)
    connectivity = connectivity_estimator(conn)


    # IMPORTANT(amotta): The relative reciprocity of connections from A to B is
    # defined as the reciprocity of connections from A to B divided by the
    # expected reciprocity. The expected value is given by the connection
    # probability IN THE OPPOSITE DIRECTION (i.e., from B to A).
    return {k: _relative_reciprocity(v, connectivity[k[::-1]])
            for k, v in reciprocity.items()}


def recurrency_estimator(conn, length=5):
    def _recurrency(conn):
        cycles = conn.astype(np.float)
        cycles = np.linalg.matrix_power(cycles, length)
        cycles = np.trace(cycles)
        return int(cycles)

    return {'aa': _recurrency(conn[a, a]),
            'ee': _recurrency(conn[e, e]),
            'ii': _recurrency(conn[i, i])}


def relative_recurrency_estimator(conn, length=5):
    def _relative_recurrency(r, n, p):
        if p == 0: return 0.0
        expected = np.power(n * p, length)
        return r / expected

    recurrency = recurrency_estimator(conn, length)
    num_neurons = {'aa': N, 'ee': N_EXC, 'ii': N_INH}
    connectivity = connectivity_estimator(conn)

    return {k: _relative_recurrency(v, num_neurons[k], connectivity[k])
            for k, v in recurrency.items()}


def in_out_degree_correlation_estimator(conn):
    def _correlation(conn):
        inputs = conn.sum(axis=0)
        outputs = conn.sum(axis=1)
        return np.corrcoef(inputs, outputs)[0, 1]

    return {'aa': _correlation(conn[a, a]),
            'ee': _correlation(conn[e, e]),
            'ii': _correlation(conn[i, i])}


def process_connectome(conn):
    stats = {'connectivity': connectivity_estimator(conn),
             'relative_reciprocity': relative_reciprocity_estimator(conn),
             'relative_recurrency': relative_recurrency_estimator(conn),
             'in_out_degree_correlation': in_out_degree_correlation_estimator(conn)}

    def  _prepend_key_names(pre, in_dict):
        return {(pre + '_' + k): v
                for k, v in in_dict.items()}

    stats = [_prepend_key_names(k, v)
             for k, v in stats.items()]

    return stats


def process_ckpt_file(ckpt_file, threshold_percentiles=None):
    # 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)
    weights = np.absolute(weights)
    num_weights = np.sum(mask)

    if threshold_percentiles:
        threshold_percentiles = np.asarray(threshold_percentiles, dtype=float)

        # NOTE(amotta): NumPy's `percentile` function throws an error when the
        # first input argument is empty. Instead, we return `np.nan` as weight
        # threshold in this case.
        #   This hack only affects the situation when `mask` only consists of
        # `False`. The `np.nan` weight thresholds won't have an effect on the
        # binary connectome that is passed to the connectome processing
        # function.
        threshold_weights = np.full(threshold_percentiles.shape, np.nan, weights.dtype)
        if np.any(mask): threshold_weights = np.percentile(weights[mask], threshold_percentiles)
    else:
        # NOTE(amotta): For backward compatibility
        threshold_percentiles = np.asarray([0], dtype=float)
        threshold_weights = np.asarray([0], dtype=weights.dtype)

    all_stats = []
    threshold_pairs = zip(threshold_percentiles, threshold_weights)
    for threshold_percentile, threshold_weight in threshold_pairs:
        conn = np.logical_and(mask, weights >= threshold_weight)
        conn_stats = process_connectome(conn)

        stats = {
            'ckpt_file': ckpt_file,
            'num_weights': num_weights,
            'threshold_percentile': threshold_percentile,
            'threshold_weight': threshold_weight}
        for s in conn_stats: stats.update(s)
        all_stats.append(stats)

    return all_stats


def process_ckpt_files(args):
    stats = [process_ckpt_file(f, args.threshold_percentiles)
             for f in args.ckpt_files]

    # NOTE(amotta): Flatten nested lists
    stats = list(itertools.chain(*stats))
    fieldnames = list(stats[0].keys())

    with open(args.out_file, 'w', newline='') as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(stats)


def build_arg_parser():
    p = argparse.ArgumentParser()
    p.add_argument('-o', '--out-file')
    p.add_argument('-t', '--threshold-percentiles', type=float, action='append')
    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()
