from os import path
import json
import argparse
from matrices import MatrixList, add_with_matrix, make_zero_matrix, subtract_with_matrix, multiply_with_scalar
from collections import defaultdict
import matplotlib.pyplot as plt

def sum_matrix(matrix):
    total = 0
    for row in matrix:
        total += sum(row)
    return total


def get_total_load_file_path(n_milis=5000, output_dir=".", **kwargs):
    return path.join(output_dir, "total_load_" + str(n_milis) + ".json")


def write_results(total_load, n_milis=5000, output_dir=".", **kwargs):
    test_file_path = get_total_load_file_path(n_milis=n_milis, output_dir=output_dir, **kwargs)
    with open(test_file_path, "w") as test_res_file:
        json.dump({"total_load": total_load}, test_res_file)


def compute_load(n_milis=5000, n_tors=80, output_dir=".", **kwargs):
    per_mili_pattern = path.join(output_dir, "matrix_mili_%d")
    per_mili_matrix_list = MatrixList(per_mili_pattern)
    total = 0
    tps = []
    for i in range(n_milis):
        print "\r", i, "/", n_milis,
        matrix = per_mili_matrix_list[i]
        mili_traffic = sum_matrix(matrix)
        total += mili_traffic
        tps.append(mili_traffic)
    try:
        plt.plot(tps)
        plt.savefig(path.join(output_dir, "total_throughput.pdf"))
        #plt.show()
    except Exception as e:
        print (e)
    return total


def main():
    parser = argparse.ArgumentParser(
        description="""Compute total load for all milis """,
        epilog="""
    """)
    parser.add_argument('--output_dir', default=".", type=str,
                        help='output directory file (default: conf.json)')
    parser.add_argument('--conf', default="conf.json", type=open,
                        help='configuration file (default: conf.json)')
    args = parser.parse_args()
    conf = json.load(args.conf)
    
    total_load = compute_load(**conf)
    print ("****** the total load is: ", total_load)
    write_results(total_load, **conf)


if __name__ == "__main__":
    main()