import subprocess
import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
#from matplotlib2tikz import save as tikz_save
import os
import numpy as np
from sys import stdout
import pdb

def mean(numbers):
    return float(sum(numbers))/max(len(numbers),1)

def exec_paraws (N, nbr_threads):
    os.environ['OMP_NUM_THREADS'] = str(nbr_threads)
    try:
      output_str = subprocess.check_output(["./paraws","-A","-x",str(N),"-y",str(N),"-z",str(N),"-C", "-w"])
      output = float(output_str)
    except:
      output = -1
    return output

def exec_hqws (N, nbr_threads):
    os.environ['OMP_NUM_THREADS'] = str(nbr_threads)
    try:
      output_str = subprocess.check_output(["./paraws", "-B","-x",str(N),"-y",str(N),"-z",str(N),"-C", "-w"])
      output = float(output_str)
    except:
      output = -1
    return output

def exec_both (N, nbr_threads):
    os.environ['OMP_NUM_THREADS'] = str(nbr_threads)
    try:
      output_str = subprocess.check_output(["./paraws","-AB","-x",str(N),"-y",str(N),"-z",str(N),"-C", "-w"])
      output = [float(s) for s in output_str.split()]
    except:
      output = [-1,-1]
    return output

def plot_time (T1, T2, S):
    plt.title ('Time evaluation')
    plt.xlabel('Size (N*N)')
    plt.ylabel('Time (s)')
    plt.plot(S,T1)
    plt.plot(S,T2)
    plt.savefig ('time.pdf')
#   tikz_save ('time.tikz', show_info=False)
    plt.clf()

def plot_speed_up (S, T, N):
    plt.title ('Speed-up measure for matrix of size ' + str(N) + '*' + str(N))
    plt.xlabel('Number of threads')
    plt.ylabel('Speed-up')
    plt.plot(T,S)
    plt.savefig ('speedup.pdf')
#   tikz_save ('speedup.tikz', show_info=False)
    plt.clf()

def plot_efficiency (E, T, N):
    plt.title ('Efficiency measure for matrix of size ' + str(N) + '*' + str(N))
    plt.xlabel('Number of threads')
    plt.ylabel('Efficiency')
    plt.plot(T,E)
    plt.savefig ('efficiency.pdf')
#   tikz_save ('efficiency.tikz', show_info=False)
    plt.clf()

def print_progress (current, total, i, nbr_iter, time1, time2):
  percent_done = int(current) * 100. /total
  done = int (50*current / int(total))
  stdout.write ("\r [%s%s][%d/%d] (%d/%d) paraws: %3.3f, hqws: %3.3f%s" % ('=' * done, ' ' * (50 - done), current, total, i, nbr_iter, time1, time2, ' '*10))
  stdout.flush ()

nbr_threads_max = 32
nbr_runs = 10

time_paraws = []
time_hqws = []
size = [i*100+100 for i in range(10)]
stdout.write ("Size bench\n")
for N in size:
    out_paraws = []
    out_hqws = []

    for i in range(nbr_runs):
      output = exec_both(N, nbr_threads_max)
      out_paraws.append (output[0])
      out_hqws.append (output[1])
      print_progress (N, 1000, i, nbr_runs, out_paraws[-1], out_hqws[-1])
    output = [mean([x for x in out_paraws if x>0]), mean([x for x in out_hqws if x>0])]
    time_paraws.append (output[0])
    time_hqws.append (output[1])
stdout.write ("\n")
plot_time(time_paraws, time_hqws,size)

size = np.array (size)
time_paraws = np.array (time_paraws)
time_hqws = np.array (time_hqws)
np.save ('size.npy', size)
np.save ('time_paraws.npy', time_hqws)
np.save ('time_hqws.npy', time_paraws)

N = 800
threads = range(1,nbr_threads_max+1)
time2 = []
stdout.write ("Multithread bench\n")
for t in threads:
    output_bench = []
    for i in range(nbr_runs):
      output_bench.append (exec_paraws(N, t))
      print_progress (t, nbr_threads_max, i, nbr_runs, output_bench[-1], 0)
    output = mean([x for x in output_bench if x>0])
    time2.append (output)
stdout.write ("\n")
speedup = [time2[0]/time2[i] for i in range(len(time2))]
plot_speed_up (speedup,threads,N)

efficiency = [speedup[i]/(i+1) for i in range(len(speedup))]
plot_efficiency(efficiency,threads,N)

threads = np.array (threads)
speedup = np.array (speedup)
efficiency = np.array (efficiency)

np.save ('threads.npy', threads)
np.save ('speedup.npy', speedup)
np.save ('efficiency.npy', efficiency)