import sys
sys.path.append('../')
import matplotlib.pyplot as plt

from netflow import *

from time import perf_counter
        
def main():
    if len(sys.argv) != 6:
        raise AssertionError("Usage: python3 perf.py <2^min threads> <2^max threads> <n1> <n2> <n3>")
    
    basenet = netflow.load_network_from('../netflow/network/network.h5')
    print("Network statistics:")
    print("Throats: %d" % len(basenet.throats))
    pore_throat_counts = [len(pore.throats) for pore in basenet.pores]
    print("Max. number of throats per pore: %d" % max(pore_throat_counts))
    print("Min. number of throats per pore: %d" % min(pore_throat_counts))
    print("Avg. number of throats per pore: %f" % (sum(pore_throat_counts) / len(pore_throat_counts)))
    
    target  = [int(sys.argv[3]), int(sys.argv[4]), int(sys.argv[5])]
    print("Target size: {}".format(target))
    cutoff  = 0.5 * max([basenet.ub[i] - basenet.lb[i] \
                        for i in range(len(basenet.ub))])
    
    i = int(sys.argv[1])
    max_ = int(sys.argv[2])
    
    speedups = []
    threads = []
    times = []
    while (i <= max_):
        nthreads = 1 << i
        print(nthreads)
        
        start = perf_counter()
        dendro = netflow.generate_dendrogram(basenet=basenet, targetsize=target, \
                                        cutoff=cutoff, sd=42, nthreads=nthreads, mute=True)
        end = perf_counter()
        elapsed = end - start
        
        print("Elapsed: %e" % elapsed)
        print("Generated network statistics:")
        print("Throats: %d" % len(dendro.throats))
        print("Pores:   %d" % len(dendro.pores))
        
        threads.append(nthreads)
        times.append(elapsed)
        
        speedups.append(times[0] / times[i])
        i += 1
        
    plt.figure()
    plt.title(r"Strong Scaling for targetsize {}".format(target))
    plt.xlabel(r"#processes")
    plt.ylabel(r"Speedup")
    plt.plot(threads, speedups, '--bo')
    plt.savefig('plots/strong_scaling.png')
    plt.close()
    
if __name__ == '__main__':
    main()