# -*- coding: utf-8 -*-


import time
import os
import logging
import argparse
import json
import codecs
import cPickle as pickle
import signal

import numpy
import networkx as nx
import beanstalkc

import foggy
import jobq

from itertools import izip
from uuid import uuid4
from glob import glob

from IPython.parallel import (Client, interactive)


logging.basicConfig()
LOGGER = logging.getLogger()
#LOGGER.setLevel(logging.INFO)
LOGGER.setLevel(logging.DEBUG)
#LOGGER.addHandler(logging.StreamHandler())


###############################################################################
# Supply
###############################################################################


class JSONBeanShooter(object):
    _shooter = {
        "parallel": "_shot",
        "deletory": "_capacity_shot",
        "buffered": "_capacity_shot"
    }

    def __init__(self, bean_queue, encoding="utf-8", **kw_args):
        """
        Warning
        -------
        bean_queue must be using the right tube!
        """
        super(JSONBeanShooter, self).__init__(**kw_args)
        self.queue = bean_queue
        self.encoding = encoding

    def __call__(self, filename):
        config = json.load(codecs.open(filename, encoding=self.encoding,
            mode="rb"))
        LOGGER.debug(str(config))
        for path in config["graphs_dir"]:
            assert os.path.exists(path), "directory does not exist '%s'" % path
        shoot = getattr(self, self._shooter[config["walk_type"]])
        description = config.copy()
        del description["graphs_dir"]
        del description["graphs_type"]
        del description["walker_factors"]
        del description["steps_factors"]
        del description["variation_factors"]
        del description["capacity_factors"]
        del description["repetition"]
        for (path, net_type) in izip(config["graphs_dir"], config["graphs_type"]):
            graphs = glob(os.path.join(path, "*.pkl"))
            LOGGER.debug("%d graphs found", len(graphs))
            for net in graphs:
                description["graph_file"] = net
                description["graph_type"] = net_type
                for kw in config["walker_factors"]:
                    description["walker_factor"] = kw
                    for var in config["variation_factors"]:
                        description["variation_factor"] = var
                        for ks in config["steps_factors"]:
                            description["steps_factor"] = ks
                            for _ in range(config["repetition"]):
                                shoot(config, description)

    def _shot(self, config, description):
        description["sim_id"] = str(uuid4()).replace("-", "")
        LOGGER.debug("firing")
        self.queue.put(pickle.dumps(description))

    def _capacity_shot(self, config, description):
        description["sim_id"] = str(uuid4()).replace("-", "")
        for k in config["capacity_factors"]:
            description["capcity_factor"] = k
            LOGGER.debug("firing")
            self.queue.put(pickle.dumps(description))


###############################################################################
# Consumption
###############################################################################


def remote_consumer(worker, host, port):
    queue = beanstalkc.Connection(host=host, port=port)
    queue.watch("input")
    queue.use("output")
    jobq.generic_consumer(queue, worker, "STOP")

@interactive
def dummy_worker(**kw_args):
    LOGGER.debug(str(kw_args))
    return "success"


class BeanMuncher(object):
    _setup = {
        "uniform": foggy.prepare_uniform_walk,
        "directed": foggy.prepare_directed_walk
    }
    _type = {
        "deletory": foggy.deletory_march,
        "buffered": foggy.buffered_march,
        "parallel": foggy.march
    }
    _distribution = {
        "uniform": foggy.UniformInterval
    }
    _visit = {
        "constant": foggy.ConstantValue,
        "degree": foggy.DegreeDependentValue
    }

    def graph_info(graph):
        LOGGER.info("%s graph:", "directed" if graph.is_directed() else "undirected")
        LOGGER.info("    %d nodes", len(graph))
        LOGGER.info("    %d edges", graph.size())
        LOGGER.info("    %d component(s)", nx.number_connected_components(graph))

    def uniform_capacity(graph, indices, walkers, num_steps):
        capacity = numpy.zeros(len(graph), dtype=float)
        capacity += float(walkers.mid_point * num_steps) / float(len(graph))
        return capacity

    def degree_capacity(graph, indices, walkers, num_steps):
        capacity = numpy.zeros(len(graph), dtype=float)
        total_degree = sum(deg for (node, deg) in graph.degree_iter())
        for (node, deg) in graph.degree_iter():
            capacity[indices[node]] = float(deg * walkers.mid_point * num_steps) /\
                     total_degree
        return capacity

    _capacity = {
            "uniform" : uniform_capacity,
            "degree" : degree_capacity
    }

    def _run(self, config, description, net, nbrs, probs, indices, walkers,
            num_steps, visits, seed=None):
        job_descr = dict()
        job_descr["parameters"] = description
        job_descr["simulation"] = self._type[config["walk_type"]]
        job_descr["neighbours"] = nbrs
        job_descr["probabilities"] = probs
        job_descr["sources"] = indices.values()
        job_descr["num_walkers"] = walkers
        job_descr["time_points"] = config["time_points"]
        job_descr["steps"] = num_steps
        job_descr["assessor"] = visits
        job_descr["transient"] = config["transient"]
        job_descr["seed"] = seed
        job_descr["capacity"] = None

    def _capacity_run(self, config, description, net, nbrs, probs, indices, walkers,
            num_steps, visits, seed=None):
        description["sim_id"] = str(uuid4()).replace("-", "")
        job_descr = dict()
        job_descr["parameters"] = description
        job_descr["simulation"] = self._type[config["walk_type"]]
        job_descr["neighbours"] = nbrs
        job_descr["probabilities"] = probs