#!/bin/python

import argparse as ap
import glob
from itertools import islice
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages
import numpy as np
import re

def re_int_nested(string,re1,re2):
    _match = re1.search(string).group()
    _submatch = re2.search(_match).group()
    return float(_submatch)


def re_find_all(string,re,entry):
    return float(re.findall(string)[entry])

def re_search(string,re):                           
    return float(re.search(string).group())

def plot_hpc(plotfile,label, data_dict, mpi_procs=None):
    pdf = PdfPages(plotfile)
    # set data
    area = np.pi*np.sqrt(data_dict["procs"]*data_dict["threads"])**2
    colors = data_dict["procs"]
    # filter according to mpi processes
    x = data_dict[label[1]]
    y = data_dict[label[2]]
    if mpi_procs is None:
            plt.scatter(x,y,s=area,alpha=0.5)
    else:
        color=['b','r']
        for i,d in enumerate(mpi_procs):
            tmp = data_dict["procs"] == d
            print(i,d,data_dict["procs"])
            _x = data_dict[label[1]][tmp]
            _y = data_dict[label[2]][tmp]
            area = np.pi*np.sqrt(data_dict["procs"][tmp]*data_dict["threads"][tmp])**2
            plt.scatter(_x,_y,s=area,c=color[i], label="%d MPI-proc." %int(d),alpha=0.5)
    # set layout
    plt.title(label[0])
    plt.xlabel(label[1])
    plt.ylabel(label[2]+' [s]')
    plt.legend(loc=7,markerscale=0.5)
    pdf.savefig()
    pdf.close()
    plt.clf()

def main():                
    # Read in the arguments from every file in regular expression
    parser = ap.ArgumentParser()
    parser.add_argument("source", help="Source directory where to search for output files")
    parser.add_argument("files", help="regular expression matching files for analysis")
    parser.add_argument("outfile", help="filename for table")
    parser.add_argument("keys", nargs='+',help="keywords to search for file")
    args=parser.parse_args()
    
    # list of filenames matching regular expression
    print(args.source+'/'+args.files)
    flist=glob.glob(args.source+'/'+args.files)
    # compile a decimal floating point regex
    re_int = re.compile('[0-9]+')
    re_dec = re.compile('[0-9]+\.[0-9]+')
    re_mpi = re.compile('mpi[0-9]+') 
    re_omp = re.compile('omp[0-9]+')
    # init empty arrays
    _mpi = np.zeros((len(flist)))
    _omp = np.zeros_like(_mpi)
    _setup = np.zeros_like(_mpi)
    _walltime = np.zeros_like(_mpi)
    # this is where the read in happens
    for i,fname in enumerate(flist):
        _mpi[i] = re_int_nested(fname,re_mpi,re_int)
        print("found %f mpi" %_mpi[i])
        _omp[i] = re_int_nested(fname,re_omp,re_int)
        with open(fname) as f:
            for line in f:
                for k in args.keys:
                    if k in line:
                        if k == "setup steps":
                            res=''.join(islice(f,1))
                            # Interested in first number
                            _setup[i] = re_find_all(res, re_dec, 0)
                        if k == "wall clock time":
                            _walltime[i] = re_search(line, re_dec)
    # generate an output table
    header = "MPI-procs	OMP-Threads	Setup-time[s]	Wall-solve[s]"
    tab = np.vstack((_mpi,_omp,_setup,_walltime)).T
    np.savetxt(args.outfile,tab,fmt='%.2f',header=header)
    # generate dictionaries for every variable, easier to handle in plot
    # plot results
    data_dict = {"procs": _mpi, "threads": _omp, "setup":_setup, "wall-solve":_walltime}
    print(data_dict)
    label = ["1 node", "threads", "wall-solve"]
    plot_hpc('wall_solve.pdf',label,data_dict,mpi_procs=[16.,32.])
    label = ["1 node", "threads", "setup"]
    plot_hpc('setup-time.pdf',label,data_dict,mpi_procs=[16.,32.])

if __name__ == '__main__':
	try:
		print("starting")
		main()
	except KeyboardInterrupt:
		pass
