# Modify the parameters like: machine_name, Revisionnr, SimulationName
# to plot, just run:
#    python3 plot.py
import sys
sys.path.append('/home/gk783/apes/gleaner/')
import gleaner
import numpy as np
from gleaner import *
from matplotlib.ticker import FuncFormatter
import matplotlib.pyplot as plt
import math
#from matplotlib import rc
#rc('font',**{'family':'sans-serif'})

# User input options ------------------------------------------
nproc_per_node = 8        # number of MPI Process per node
machine_name   = 'JUQEEN'    # name of the machine
Revisionnr     = '746bdde1ed8f' # Revision Number the script is using
SimulationName = 'AOS_o5_mp_d3q19' # Simulation Name the script is using
filename       = 'timing.res'

show_title  = False    # Show title on Diagrams
plot_perf   = True     # plot Performance Map
plot_weak   = False    # plot Weak Scaling
plot_strong = False    # plot Strong Scaling

perf_figname = 'perf_map'

# y axis max range
# ymax = 80

#Calculated values for Skaling going from highest downwards
weaknum = 4                # Number of combinations used for weak plot
strong_range = [1,2,3,4]   # order of problem size for strong scaling

# Font of Diagrams:
markersizes = 9            # Size of Makers
fontsizes = 15             # Size of Font
my_lw= 2                   # Size of lines

# User input options ------------------------------------------

if show_title:
    appendix = '_wt'  # appendix for figure file name
else:
    appendix = '_nt'

def to_percent(yy, pos):
  s = str(int(100 * yy))
  return s + '%'

print('Scaling analysis for Musubi')
markers = ["o", "s", "p","^","D",">","<","v",];
colors = ["b","g","r","k","m","brown"];

# Load text, dump into a database, get back the desired columns
sqlcon = gleaner.timing_to_db( fname  = filename,
                               dbname = ':memory:',
                               tabname = 'timings' )

perfmap, ID_tuple = gleaner.perfmap_series(sqlcon = sqlcon,
                                           tabname = 'timings',
                                           signature = ['Revision',
                                                        'SimName',
                                                        'nProcs',
                                                        'DomSize'],
                                           xcol = 'DomSize',
                                           ycol = 'MLUPs')

sqlcon.close()
# set the correct SimName and Revision
series = perfmap[ID_tuple(SimName=SimulationName,Revision=Revisionnr)]
############################################################################
# Get possible combinations of Elements/node
elem_node_list = [] # list of all values of elements per node
elem_list      = [] # list of all values of total elements
for nproc in sorted(series):
    nNodes = nproc / nproc_per_node
    if nNodes >= 1.0:
      for domsize, mlups  in series[nproc]:
        elem_per_node = domsize/nNodes
        if elem_per_node not in elem_node_list:
          elem_node_list.append(elem_per_node)
        if domsize not in elem_list:
          elem_list.append(domsize)

elem_list      = sorted(elem_list,reverse=True)
elem_node_list = sorted(elem_node_list,reverse=True)

############################################################################
# performance plot
if plot_perf:
  plt.figure(1)
  ii = -1
  for nproc in sorted(series):
      ii = ii + 1
      nNodes = nproc / nproc_per_node
      if nNodes >= 1.0:
          # create empty list
          # x,y contains the values of domsize and mlups for runs with nproc
          x = [];
          y = [];
          for domsize, mlups in series[nproc]:
              # add values into list
              x.append( domsize/nNodes )
              y.append( mlups[0]/nNodes )
          # Plot data for this number of procs
          mylabel = str(int(nNodes))+" nodes"
          plt.semilogx( x, y, label = mylabel , marker =
             markers[ii%len(markers)], color=colors[ii%len(colors)],
             markersize=markersizes,linewidth=my_lw)

  plt.grid(True)
  # plt.ylim(0,ymax)
  plt.xlabel("elements / nodes",fontsize = fontsizes)
  plt.ylabel("MLUPS / nodes",fontsize = fontsizes)
  if show_title:plt.title("Performance Map for " +str(machine_name),fontsize = fontsizes)
  leg =  plt.legend( loc = "upper left", fontsize = fontsizes, numpoints=1)
  vp = leg._legend_box._children[-1]._children[0]
  # set text in legend right aligned
  for c in vp._children:
    c._children.reverse()
    vp.align="right"
  # save figure as file
  figname = perf_figname + appendix + '.eps'
  plt.savefig(figname, format='eps', dpi=1000)
############################################################################

############################################################################
# weak_efficiency plot
sizes=[]
if plot_weak:
  plt.figure(2)
  for ii in range(int(weaknum)): 
   sizes.append(elem_node_list[ii])
  ii=-1
  for size in sizes:
    ii=ii+1
    x=[]
    y=[]
    eff_table = weak_efficiency(perfmap = series, size_per_node = size, \
                                ppn = nproc_per_node, min_nodes=1)
    if eff_table != None:
      for nn, eff in eff_table:
        x.append(nn)
        y.append(eff)
    mylabel = str(int(size))+" elements per node"
    plt.semilogx( x, y, label = mylabel,
     marker=markers[ii%len(markers)],
     color=colors[ii%len(colors)],markersize=markersizes,linewidth=my_lw)

  plt.grid(True) 
  plt.ylim([0,1.4])
  plt.xlim([x[0]/2,(x[len(x)-1])*2])
  formatter = FuncFormatter(to_percent)
  plt.gca().yaxis.set_major_formatter(formatter)
  plt.xlabel("number of nodes",fontsize = fontsizes)
  plt.ylabel("efficiency",fontsize = fontsizes)
  leg = plt.legend( loc = "best", fontsize = fontsizes, numpoints=1)
  if show_title:  plt.title("Weak Scaling for " +str(machine_name),fontsize = fontsizes)
  # set xticks
  plt.xticks(x)
  plt.gca().set_xticklabels(np.int_(x))
  vp = leg._legend_box._children[-1]._children[0]
  # set text in legend right aligned
  for c in vp._children:
    c._children.reverse()
    vp.align="right"
  # save figure as file
  figname='weak' + appendix + '.eps'
  plt.savefig(figname, format='eps', dpi=1000)

############################################################################

############################################################################
# strong_efficiency
sizes=[]
if plot_strong:
  # append total problem sizes
  for ii in strong_range:
    sizes.append(elem_list[ii])

  plt.figure(3)
  ii = -1
  for size in sizes:
    ii=ii+1
    y=[]
    x=[]
    strong_table = strong_efficiency(perfmap = series, total_size = size, \
                                     ppn = nproc_per_node, min_nodes=1)
    if strong_table != None:
      for nn, eff in strong_table:
        x.append(nn)
        y.append(eff)
    mylabel = str(int(size))+" elements"
    plt.semilogx( x, y, label = mylabel,marker=markers[ii%len(markers)],
        color=colors[ii%len(colors)],markersize=markersizes,linewidth=my_lw)

  plt.grid(True)
  plt.ylim([0,1.1])
  plt.xlim([x[0]/2,(x[len(x)-1])*2])
  formatter = FuncFormatter(to_percent)
  plt.gca().yaxis.set_major_formatter(formatter)
  plt.xlabel("number of nodes",fontsize = fontsizes)
  plt.ylabel("efficiency",fontsize = fontsizes)
  leg = plt.legend( loc = "best", fontsize = fontsizes, numpoints=1)
  if show_title:
    plt.title("Strong Scaling for " +str(machine_name),fontsize = fontsizes)
  # set text in legend right aligned
  vp = leg._legend_box._children[-1]._children[0]
  for c in vp._children:
    c._children.reverse()
    vp.align="right"
  # set xticks
  plt.xticks(x)
  plt.gca().set_xticklabels(np.int_(x))
  # save figure as file
  figname='strong' + appendix + '.eps'
  plt.savefig(figname, format='eps', dpi=1000)

#############################################

