#!/usr/bin/python
import argparse
import numpy as np
import matplotlib.pyplot as plt
from glob import glob
from matplotlib.backends.backend_pdf import PdfPages
from os.path import isfile
import re

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description="""Background calculator""",
        formatter_class=argparse.RawDescriptionHelpFormatter)

    parser.add_argument('-x', '--xcolumn', type=int, default=0,
                        help='columns to use for X (count starts from 0)')

    parser.add_argument('-y', '--columns', type=str, default=None,
                        help='columns to use for Y (count starts from 0)')

    parser.add_argument('-f', '--files', type=str, default=None,
                        help='files to use')

    parser.add_argument('-o', '--output', type=str, default='output.pdf',
                        help='output file name')

    parser.add_argument('-s', '--style', type=str, default='-',
                        help='matplotlib plotting style')

    parser.add_argument('-lw', '--width', type=float, default=1.,
                        help='matplotlib line width')

    parser.add_argument('-xl', '--xlabel', type=str, default=None,
                        help='X axis label')

    parser.add_argument('-yl', '--ylabel', type=str, default=None,
                        help='Y axis label')

    parser.add_argument('--log', type=str, default='-',
                        help='Log scale (x y or xy)')

    parser.add_argument('-ltitles', '--legendtitles', type=str, default=None,
                        help='Comma-separated list of titles for legend.')

    parser.add_argument('-m', '--marker', type=str, default='',
                        help='matplotlib marker style')

    parser.add_argument('-d', '--delimiter', type=str, default=None,
                        help='delimiter')

    parser.add_argument('--skip-header', type=int, default=0,
                        help='Skip header lines')

    parser.add_argument('--title', type=str,
                        default=None, help='Plot title')

    parser.add_argument('-lp', '--legend_position', type=str,
                        default='lower center', help='delimiter')
    parser.add_argument('-t', '--tight', action="store_true",
                        default=False,
                        help='Use tight layout')

    args = parser.parse_args()

    filemasks = args.files.split(',')
    files = []
    for filemask in filemasks:
        files.extend(glob(filemask))
    files.sort()
    datas = []
    columns = []
    for column_mask in args.columns.split(','):
        minus = column_mask.find(':')
        if minus >= 0:
            columns.extend(range(int(column_mask[:minus]),
                                 int(column_mask[minus+1:])))
        else:
            columns.append(int(column_mask))
    read_columns = [args.xcolumn] + columns
    for afile in files:
        datas.append(np.loadtxt(afile, delimiter=args.delimiter,
                                skiprows=args.skip_header,
                                usecols=read_columns))
    fheader = open(files[0], 'r')
    header = fheader.readline().strip().split(args.delimiter)
    outname = args.output
    if outname == 'output.pdf':
        i = 0
        while isfile(outname) and re.match('output[0-9]*.pdf', outname):
            outname = 'output%s.pdf' % i
            i = i + 1
    pp = PdfPages(outname)
    if args.legendtitles is not None:
        legendtitles = args.legendtitles.split(',')
    else:
        legendtitles = files
    for icolumn in xrange(len(columns)):
        if args.title is not None:
            plt.title(args.title)
        labels = []
        if args.xlabel is None:
            plt.xlabel(header[args.xcolumn])
        else:
            plt.xlabel(args.xlabel)
        if args.ylabel is None:
            plt.ylabel(header[icolumn])
        else:
            plt.ylabel(args.ylabel)
        if 'x' in args.log:
            plt.xscale('log')
        if 'y' in args.log:
            plt.yscale('log')
        plt.grid(True)
        for ifile, data in enumerate(datas):
            if data.shape[1] > icolumn + 1:
                p = plt.plot(data[:, 0], data[:, icolumn + 1],
                             linestyle=args.style, marker=args.marker,
                             linewidth=args.width)
            labels.append(legendtitles[ifile])
        if args.legend_position != 'none':
            plt.legend(labels, loc=args.legend_position,
                       # bbox_to_anchor=(0.5, 0.05),
                       fancybox=True, shadow=True, ncol=2, prop={'size': 8})
        if args.tight:
            plt.tight_layout()
        plt.savefig(pp, format='pdf')
        if icolumn < len(columns) - 1:
            plt.clf()
    #pp.savefig()
    pp.close()
    print '%s written' % outname