#!/usr/bin/env python

"""Command line tool for plotting graphs of log-posterior and std."""

import argparse
import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
from matplotlib import cm
from os.path import join, splitext

from cryoem.program.utils import read_csv_dict


parser = argparse.ArgumentParser(
    prog='gmm-plot-stats',
    description='Plot reconstruction stats.')

parser.add_argument('-i', '--input', nargs='+',
                    help='Single or multiple csv files.')
parser.add_argument('-o', '--output',
                    help='Name of output files.')

args = parser.parse_args()

# Load csv files.
ds = [read_csv_dict(f) for f in args.input]
# Assume that all files have the same keys.
# Create one plot for every key, except the index key.
index_key = 'step'
for k in ds[0].keys():
    if k==index_key:
        continue
    
    plt.figure(figsize=(4*4, 3*4))
    for i, d in enumerate(ds):
        x = d[index_key]
        y = d[k]
        plt.plot(x, y, label=str(i))
    plt.legend()
    plt.grid()
    plt.title(k)
    filename = '%s_stats_%s.png' % (args.output, k)
    plt.savefig(filename, bbox_inches=None)
