Manages agent information.
The AgentPop class implements a flexible data structure
that allows the user of the class to access agents
from multiple perspectives.

from collections import OrderedDict
import random
import itertools
import json
import networkx as nx
import logging
import indra.node as node
import indra.entity as ent


class AgentPop(node.Node):
    Holds our collection of agents and a sub-graph
    of their relationships.
    This should implement list functions.

    def __init__(self, name):
        self.vars = OrderedDict()
        self.graph = nx.Graph()

    def __iter__(self):
        alists = []
        for var in self.varieties_iter():
        # create an iterator that chains the lists together as if one:
        return itertools.chain(*alists)

    def __len__(self):
        l = 0
        for var in self.varieties_iter():
            l += len(self.vars[var]["agents"])
        return l

    def __reversed__(self):
        all_agents = self.all_agents_list()
        return reversed(all_agents)

    def varieties_iter(self):
        Allows iteration over the varieties of agents.
        return iter(self.vars)

    def variety_iter(self, var):
        Allows iteration over the agents of one variety
        return iter(self.vars[var]["agents"])

    def all_agents_list(self):
        all_agents = []
        for var in self.varieties_iter():
            all_agents += self.vars[var]["agents"]
        return all_agents

    def agent_random_iter(self):
        Loop through agents in random order.
        all_agents = self.all_agents_list()
        return iter(all_agents)

    def element_at(self, i):
        Another way to treat the AgentPop as if it were really
        one big list.
        if i < 0 or i > len(self):
            raise IndexError()
            for var in self.varieties_iter():
                l = len(self.vars[var]["agents"])
                logging.debug("Looking for element from "
                              + var + " at position "
                              + str(i) + " and var has len "
                              + str(l))
                if i < l:
                    # that means the agent is in this list
                    return self.vars[var]["agents"][i]
                    # otherwise, the agent lies in one of the
                    # remaining lists, so subtract the length
                    # of this one from i and continue.
                    i -= l

    def get_var_color(self, var):
        Get the display color for this variety.
        return self.vars[var]["disp_color"]

    def set_var_color(self, var, color):
        Set the display color for this var.
        self.vars[var]["disp_color"] = color

    def add_variety(self, var):
        Sometimes, we may know a variety is coming, and want to
        add it even before any agents of that variety are created,
        for census or graphing purposes.
        if var not in self.vars:
            self.vars[var] = {"agents": [],
                              # some arbitrary data to track for pop:
                              "pop_data": 0,
                              "pop_hist": [],
                              "my_periods": 0,
                              "disp_color": None}

    def append(self, agent, v=None):
        Appends to agent list.
        if v is None:
            var = agent.get_type()
            var = v
        logging.info("Adding %s of variety %s" % (agent.name, var))

        if var not in self.vars:
            self.graph.add_edge(self, var)


# we link each agent to the variety
# so we can show their relationship
        self.graph.add_edge(var, agent)

    def remove(self, agent, v=None):
        Removes from agent list.
        logging.debug("Removing " + agent.name + " from agents")
        if v is None:
            var = agent.get_type()
            var = v
        self.graph.remove_node(agent)  # also removes edges!

    def join_agents(self, a1, a2):
        Add a graph edge between agents.
        self.graph.add_edge(a1, a2)

    def contains(self, var):
        Do we have this sort of thing in our env?
        return var in self.vars

    def get_agents_of_var(self, var):
        Return all agents of type var.
        if var in self.vars:
            return self.vars[var]["agents"]
            return None

    def get_randagent_of_var(self, var):
        return random.choice(self.vars[var]["agents"])

    def get_pop(self, var):
        Return the population of variety 'var'
        return len(self.vars[var]["agents"])

    def get_total_pop(self):
        total_pop = 0
        for var in self.varieties_iter():
            total_pop += self.get_pop(var)
        return total_pop

    def get_my_pop(self, agent):
        Return the population of agent's type
        var = agent.get_type()
        return self.get_pop(var)

    def get_var_pop_hist(self, var):
        return self.vars[var]["pop_hist"]

    def get_pop_hist(self):
        Make a list containing the population history
        for each var in vars, if var is None.
        We should merge this with display_methods
        assemble data when we can.
        pop_hist = {}
        for var in self.varieties_iter():
            pop_hist[var] = {}
            pop_hist[var]["data"] = self.vars[var]["pop_hist"]
            pop_hist[var]["color"] = self.vars[var]["disp_color"]
            logging.debug("Setting color for %s to %s"
                          % (var, pop_hist[var]["color"]))
        return pop_hist

    def get_pop_data(self, var):
        Return the value of pop_data for 'var'.
        return self.vars[var]["pop_data"]

    def change_pop_data(self, var, change):
        Change the value of pop_data by 'change.'
        self.vars[var]["pop_data"] += change

    def change_agent_type(self, agent, old_type, new_type):
        self.remove(agent, v=old_type)
        self.append(agent, v=new_type)
        #need to change agent's ntype as well
        agent.ntype = new_type

    def append_pop_hist(self, var, pop):
        Add the most recent pop to pop_hist.

    def restore_hist_from(self, pop_hist_json):
        Restore histogram from the client. (The inverse of get_pop_hist.)
        for var in pop_hist_json:
            self.vars[var]["pop_hist"] = pop_hist_json[var]["data"]
            self.vars[var]["disp_color"] = pop_hist_json[var]["color"]
            logging.debug("Setting color for {} to {}".format(var, pop_hist_json[var]["color"]))

    def census(self, exclude_var=None):
        Take a census of agents by variety.
        Return a string of results for possible display.
        ret = ""
        for v in self.varieties_iter():
            if v != exclude_var:
                pop = self.get_pop(v)
                ret += v + ": " + str(pop) + "\n"
                var = self.vars[v]
                # a type might enter the env after period 0!
                # so we track that in my_periods
                var["my_periods"] += 1
        return ret

    def jsondump(self, obj):
        if type(obj) == nx.Graph:
            return "Graph"
        elif isinstance(obj, ent.Entity):
            return obj.to_json()
            return obj.__dict__

    def record_results(self, file_nm):
        json.dump(self, open(file_nm, 'w'), indent=4,