"""Example to run a batch of simlations on SHED data."""

__author__ = 'Jarrod Pas <j.pas@usask.ca>'

import sys
from argparse import ArgumentParser
from collections import namedtuple
from multiprocessing import Pool
from os import path
from pprint import pprint

from pydtnsim import Network, RandomTraffic, Node, EpidemicNode, CSVTrace
from pydtnsim.community import BubbleKCliqueNode, BubbleLouvainNode
from pydtnsim.community import HCBFKCliqueNode, HCBFLouvainNode


Simulation = namedtuple('Simulation', ['trace', 'node_type', 'seed'])


def run_simulation(simulation):
    """Run a simulation."""
    seed = simulation.seed

    csv = path.join(simulation.trace, 'contact.csv')
    metadata = path.join(simulation.trace, 'metadata.json')
    trace = CSVTrace(csv, metadata=metadata)

    epoch = 7*24*60*60  # 7 days

    node_type = simulation.node_type
    node_options = {
        'tick_rate': 5 * 60,  # 5 mins
        'epoch': epoch,
        'k': 3,
        'context': {}
    }
    nodes = {
        node_id: simulation.node_type(**node_options)
        for node_id in range(trace.nodes)
    }

    traffic_options = {
        'seed': seed,
        'start': epoch,
        'step': 60 * 60,  # 1 packet every 1 mins
    }
    traffic = RandomTraffic(nodes, **traffic_options)

    network = Network(nodes, traffic=traffic, trace=trace)
    network.run()

    stats = {
        'trace': simulation.trace,
        'node_type': node_type.__name__,
        'seed': seed,
    }
    stats.update(network.stats_summary)

    # return stats because we can't pickle the network as it is a generator.
    return stats


def main(args):
    """Run simulation for each seed in args."""
    log = pprint if args['pretty'] else print
    pool = Pool()
    simulations = []

    trace = args['shed']
    node_types = [
        Node,
        EpidemicNode,
        BubbleKCliqueNode,
        HCBFKCliqueNode,
        BubbleLouvainNode,
        HCBFLouvainNode,
    ]

    for seed in args['seeds']:
        for node_type in node_types:
            sim = Simulation(trace=trace, node_type=node_type, seed=seed)
            simulations.append(sim)

    for stats in pool.imap_unordered(run_simulation, simulations):
        log(stats)


def parse_args(args):
    """Parse arguments."""
    parser = ArgumentParser()

    parser.add_argument('shed')
    parser.add_argument('--pretty', action='store_true')
    parser.add_argument('--seeds', '-s',
                        metavar='SEED', type=int, nargs='+', default=[None])

    args = parser.parse_args(args)
    return vars(args)


if __name__ == '__main__':
    sys.exit(main(parse_args(sys.argv[1:])))