from csv import DictReader
import sys
from ml import ml
import numpy as np
from cPickle import dump, load
from itertools import product


def build_network(filename):
    victims = {}
    non_victims = {}
    with open(filename) as fh:
        reader = DictReader(fh)
        for row in reader:
            from_, to = int(float(row["from"])), int(float(row["to"]))
            if row["t2"] != "NA":
                dt = int(row["t2"]) - int(row["t1"])
                parent = (int(row["dist"]), dt)
                if to not in victims:
                    victims[to] = []
                victims[to].append(parent)
                if from_ not in victims:
                    victims[from_] = []
            else:
                from_, to = int(float(row["from"])), int(float(row["to"]))
                parent = (int(row["dist"]), 3012 - int(row["t1"]))
                if to not in victims:
                    non_victims[to] = []
                non_victims[to].append(parent)
                if from_ not in victims:
                    victims[from_] = []
    root_victims = {}
    for victim in victims.keys():
        if not victims[victim]:
            del victims[victim]
            root_victims[victim] = []
    return root_victims, victims, non_victims


if __name__ == "__main__":
    #root_victims, victims, non_victims = build_network(sys.argv[1])
    #dump((root_victims, victims, non_victims), open("network.pickle", "w"))
    root_victims, victims, non_victims = load(open("network.pickle"))

    alpha = np.arange(0.0000005, 0.00000051, 0.000001)
    delta = np.arange(1., 1.000001, 0.001)
    with open("out.log", "a") as fh:
        for a, d in product(alpha, delta):
            beta, roots, ll = ml(root_victims, victims, non_victims, a, d)
            fh.write("\t".join(map(str, [a, d, beta, roots, ll])) + "\n")
            fh.flush()