#!/usr/bin/env python

"""Command line tool for plotting images from mrc stack."""

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.data.griddata import GridDataList

parser = argparse.ArgumentParser(
    prog='gmm-plot-stack',
    description='Plot 3D mrc stack of 2D images.')

parser.add_argument('-i', '--input', nargs='+',
                    help='Input image stack(s) as mrc file(s).')
parser.add_argument('-r', '--rows', default=1, type=int,
                    help='Number of rows of images to plot.')
parser.add_argument('-c', '--columns', default=1, type=int,
                    help='Number of rows of images to plot.')

args = parser.parse_args()


for f in args.input:

    data = GridDataList.load_mrc(f).as_list()

    num_rows = args.rows
    num_columns = args.columns

    plt.figure(figsize=(5*num_columns, 5*num_rows))
    for index in range(min(num_rows * num_columns, len(data))):
        plt.subplot(num_rows, num_columns, index + 1)
        x = data[index].reshaped()
        plt.imshow(x, cmap=cm.gray, interpolation='nearest')
        plt.axis('off')

    plt.subplots_adjust(left=0.0, right=1.0, top=1.0, bottom=0.0,
                        hspace=0.0, wspace=0.1)

    filename = splitext(f)[0] + '.png'
    plt.savefig(filename, bbox_inches=None)

