from csv import DictReader
import sys
from cPickle import dump
from os.path import splitext


def build_network(filename):
    victims = {}
    non_victims = {}
    age = 0
    with open(filename) as fh:
        reader = DictReader(fh)
        for row in reader:
            from_, to = int(float(row["from"])), int(float(row["to"]))
            dist = int(row["dist"])
            w1, w2, w3 = float(row["w1"]), float(row["w2"]), float(row["w3"])
            # if int(float(row["dist"])) > 1:
            #     continue
            # 'to' is a victim
            if row["t2"] != "NA":
                dt = int(row["t2"]) - int(row["t1"])
                parent = (dist, dt, w1, w2, w3)
                if to not in victims:
                    age += int(row["t2"]) - int(row["spawn2"])
                    victims[to] = []
                victims[to].append(parent)
                if from_ not in victims:
                    age += int(row["t1"]) - int(row["spawn1"])
                    victims[from_] = []
            # 'to' is not a victim
            else:
                dt = 3012 - int(row["t1"])
                parent = (dist, dt, w1, w2, w3)
                if to not in non_victims:
                    # age += 3012 - int(row["spawn2"])
                    non_victims[to] = []
                non_victims[to].append(parent)
                if from_ not in victims:
                    age += int(row["t1"]) - int(row["spawn1"])
                    victims[from_] = []
    root_victims = {}
    for victim in victims.keys():
        if not victims[victim]:
            del victims[victim]
            root_victims[victim] = []
    print len(root_victims), len(victims), len(non_victims)
    return root_victims, victims, non_victims, age


if __name__ == "__main__":
    if len(sys.argv) < 2:
        sys.exit("usage: {0} <file>".format(sys.argv[0]))

    filename = sys.argv[1]
    root, _ = splitext(filename)
    root_victims, victims, non_victims, age = build_network(filename)
    dump((root_victims, victims, non_victims, age), open(root + ".pickle", "w"))