"""
A fairly generic script to produce plots that compare two or more Neo data sets.

Usage:

    python comparison_plot.py <datafile1>, <datafile2>, etc.

"""

import argparse
import os
from datetime import datetime
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import numpy as np
import neo
from neo.io import get_io
from pyNN.utility.plotting import comparison_plot


def variable_names(segment):
    return set(signal.name for signal in segment.analogsignals)


def plot_signal(panel, signal, index, colour='b', linewidth='1', label=''):
    label = "%s (Neuron %d)" % (label, signal.array_annotations["channel_index"][index])
    panel.plot(signal.times, signal[:, index], colour, linewidth=linewidth, label=label)
    panel.set_ylabel("%s (%s)" % (signal.name, signal.units._dimensionality.string))
    plt.setp(plt.gca().get_xticklabels(), visible=False)


def plot(datafiles, output_file, extra_annotation=None):
    print(datafiles)
    print(output_file)
    # load data
    blocks = [get_io(datafile).read_block() for datafile in datafiles]
    #   note: Neo needs a pretty printer that is not tied to IPython
#    for block in blocks:
#        print(block.describe())
    #   for now take only the first segment
    segments = [block.segments[0] for block in blocks]
    labels = [block.annotations['simulator'] for block in blocks]

    # build annotations
    script_name = blocks[0].annotations.get('script_name', '')
    if script_name:
        for block in blocks[1:]:
            assert block.annotations['script_name'] == script_name

    #   also consider adding metadata to PNG file - see http://stackoverflow.com/questions/10532614/can-matplotlib-add-metadata-to-saved-figures
    context = ["Generated by: %s" % __file__,
               "Working directory: %s" % os.getcwd(),
               "Timestamp: %s" % datetime.now().strftime("%Y-%m-%d %H:%M:%S%z"),
               "Output file: %s" % output_file,
               "Input file(s): %s" % "\n                    ".join(datafiles)]
    if extra_annotation:
        context.append(extra_annotation)
    annotations = "\n".join(context)

    # create and save plot
    fig = comparison_plot(segments, labels, title=script_name,
                          annotations=annotations)
    fig.save(output_file)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("datafiles", metavar="datafile", nargs="+",
                        help="a list of data files in a Neo-supported format")
    parser.add_argument("-o", "--output-file", default="output.png",
                        help="output filename")
    parser.add_argument("-a", "--annotation", help="additional annotation (optional)")
    args = parser.parse_args()
    plot(args.datafiles, output_file=args.output_file, extra_annotation=args.annotation)


if __name__ == "__main__":
    main()
