#!/usr/bin/env python3
# coding=utf-8
#
# StarPU --- Runtime system for heterogeneous multicore architectures.
#
# Copyright (C) 2019-2023  Université de Bordeaux, CNRS (LaBRI UMR 5800), Inria
#
# StarPU is free software; you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation; either version 2.1 of the License, or (at
# your option) any later version.
#
# StarPU is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
#
# See the GNU Lesser General Public License in COPYING.LGPL for more details.
#

"""
Plot statistics produced when running an application with STARPU_MPI_STATS=1
"""

import sys
import re
import os

PROGNAME=sys.argv[0]

def usage():
    print("Offline tool to draw a communication matrix")
    print("")
    print("Usage: %s <output_execution>" % PROGNAME)
    print("")
    print("Options:")
    print("	-h, --help          display this help and exit")
    print("	-v, --version       output version information and exit")
    print("	-png                produce plots in png format (default is pdf)")
    print("")
    print("Report bugs to <starpu-devel@inria.fr>")
    sys.exit(1)

if len(sys.argv) >= 2:
    if sys.argv[1] == '-v' or sys.argv[1] == '--version':
        print("%s (StarPU) 1.4.4" % PROGNAME)
        sys.exit(0)
    if sys.argv[1] == '-h' or sys.argv[1] == '--help':
        usage()
if len(sys.argv) == 1:
    usage()

if len(sys.argv) >= 2 and sys.argv[1] == '-png':
    outputformat='png'
    outputext='png'
    outputfile=sys.argv[2]
else:
    outputformat='pdf color'
    outputext='pdf'
    outputfile=sys.argv[1]

# find the number of nodes
nodes=0
file = open(outputfile, "r")
for line in file.readlines():
    if re.search('TOTAL', line):
        (node,stuff)=line.split(sep="[")[2].split("]")
        if int(node) > nodes:
            nodes=int(node)
file.close()
nodes=nodes+1

# extract volume of comm and bandwidth between all pair of nodes
volumes = [[0 for _ in range(nodes)] for _ in range(nodes)]
bandwidth = [[0 for _ in range(nodes)] for _ in range(nodes)]
file = open(outputfile, "r")
for line in file.readlines():
    if re.search(r'\[starpu_comm_stats]', line) and not re.search('TOTAL', line) and re.search('MB/s', line):
        (head,volB,B,volMB,MB,bwB,B,bwMB,MB) = line.split()
        (src,dst)=head.split(sep="[")[2].split(sep="]")[0].split(sep=":")
        volumes[int(src)][int(dst)] = float(volB)
        bandwidth[int(src)][int(dst)] = float(bwB)
file.close()

def write_data(filename, nodes, data):
    ofile=open(filename, "w")
    for dst in range(nodes):
        for src in range(nodes):
            ofile.write("%f "% data[src][dst])
        ofile.write("\n")
    ofile.close()

def generate_gnuplot_script(filename, datafilename, outputfile, nodes):
    ofile=open(filename, "w")
    srctics=""
    dsttics=""
    for node in range(nodes-1):
        srctics += "\"src%d\" %d, " % (node, node)
        dsttics += "\"dst%d\" %d, " % (node, node)
    ofile.write("set term %s\n" % outputformat)
    ofile.write("set output \"%s.%s\"\n" % (outputfile, outputext))
    ofile.write("set view map scale 1\nset style data lines\n")
    ofile.write("set palette gray\n")
    ofile.write("set xtics (%s\"src%d\" %d)\n" % (srctics, nodes-1, nodes-1))
    ofile.write("set ytics (%s\"dst%d\" %d)\n" % (dsttics, nodes-1, nodes-1))
    ofile.write("plot '%s' matrix with image\n" % datafilename)
    ofile.close()

# generate gnuplot volume data and script file
write_data(outputfile+"_volume.data", nodes, volumes)
generate_gnuplot_script(outputfile+"_volume.gp", outputfile+"_volume.data", outputfile+"_volume_heatmap", nodes)
os.system("gnuplot " + outputfile+"_volume.gp")
print("Generated file \"%s.%s\"" % (outputfile+"_volume.data", outputext))

# generate gnuplot bandwidth data and script file
write_data(outputfile+"_bw.data", nodes, bandwidth)
generate_gnuplot_script(outputfile+"_bw.gp", outputfile+"_bw.data", outputfile+"_bw_heatmap", nodes)
os.system("gnuplot " + outputfile+"_bw.gp")
print("Generated file \"%s.%s\"" % (outputfile+"_bw.data", outputext))
