'''
Mesa Data Collection Module
=====================================================

DataCollector is meant to provide a simple, standard way to collect data
generated by a Mesa model. It collects two types of data: model-level and
agent-level data.

A DataCollector is instantiated with two dictionaries of reporter names and
associated functions for each, one for model-level data and one for
agent-level data. When the collect() method is called, each model-level
function is called, with the model as the argument, and the results associated
with the relevant variable. Then the agent-level functions are called on each
agent in the model scheduler.

The DataCollector then stores the data it collects in two dictionaries:
    * model_vars maps each reporter to a list of its values
    * agent_vars maps each reporter to a list of lists, where each nested list
    stores (agent_id, value) pairs.

Finally, DataCollector can create a pandas DataFrame from each collection.

The default DataCollector here makes several assumptions:
    * The model has a schedule object called 'schedule'
    * The schedule has an agent list called agents
    * For collecting agent-level variables, agents must have a unique_id
'''
from collections import defaultdict
import pandas as pd

class DataCollector(object):
    '''
    Class for collecting data generated by a Mesa model.

    A DataCollector is instantiated with dictionaries of names of model- and
    agent-level variables to collect, associated with functions which actually
    collect them. When the collect(...) method is called, it executes these
    functions one by one and stores the results.
    '''
    model_reporters = {}
    agent_reporters = {}

    model_vars = {}
    agent_vars = {}

    model = None

    def __init__(self, model_reporters=None, agent_reporters=None):
        '''
        Instantiate a DataCollector with lists of model and agent reporters.

        Both model_reporters and agent_reporters accept a dictionary mapping a
        variable name to a method used to collect it.
        For example, if there was only one model-level reporter for number of
        agents, it might look like:
            {"agent_count": lambda m: m.schedule.get_agent_count() }
        If there was only one agent-level reporter (e.g. the agent's energy),
        it might look like this:
            {"energy": lambda a: a.energy}

        Args:
            model_reporters: Dictionary of reporter names and functions.
            agent_reporters: Dictionary of reporter names and functions.
        '''

        self.model_reporters = model_reporters
        self.agent_reporters = agent_reporters

        if self.model_reporters:
            for var in self.model_reporters:
                self.model_vars[var] = []

        if self.agent_reporters:
            for var in self.agent_reporters:
                self.agent_vars[var] = []

    def collect(self, model):
        '''
        Collect all the data for the given model object.
        '''
        if self.model_reporters:
            for var, reporter in self.model_reporters.items():
                self.model_vars[var].append(reporter(model))

        if self.agent_reporters:
            for var, reporter in self.agent_reporters.items():
                agent_records = []
                for agent in model.schedule.agents:
                    agent_records.append((agent.unique_id, reporter(agent)))
                self.agent_vars[var].append(agent_records)

    def get_model_vars_dataframe(self):
        '''
        Create a pandas DataFrame from the model variables.
        The DataFrame has one column for each model variable, and the index is
        (implicitly) the model tick.
        '''
        return pd.DataFrame(self.model_vars)

    def get_agent_vars_dataframe(self):
        '''
        Create a pandas DataFrame from the agent variables.
        The DataFrame has one column for each variable, with two additional
        columns for tick and agent_id.
        '''

        data = defaultdict(dict)
        for var, records in self.agent_vars.items():
            for step, entries in enumerate(records):
                for entry in entries:
                    agent_id = entry[0]
                    val = entry[1]
                    data[(step, agent_id)][var] = val
        df = pd.DataFrame.from_dict(data, orient="index")
        df.index.names = ["Step", "AgentID"]
        return df
