from starling_sim.basemodel.output.kpis import KPI

import logging
import pandas as pd


class KpiOutput:

    def __init__(self, population_names, kpi_list, kpi_name=None):

        # name of the kpi, will compose the kpi filename : <kpi_name>.csv
        if kpi_name is None:
            if isinstance(population_names, list):
                self.name = "_&_".join(population_names) + "_kpi"
            else:
                self.name = population_names + "_kpi"
        else:
            self.name = kpi_name

        # population of agent to evaluate
        self.population_names = population_names
        self.populations = None

        # list of kpi to evaluate the given agents
        self.kpi_list = kpi_list

        # output file
        self.filename = None
        self.folder = None

    def setup(self, filename, folder, simulation_model):
        """
        Setup method called during simulation setup.

        Sets the values of out file and folder, and call setup for KPIs.

        :param filename: .csv file
        :param folder:
        :param simulation_model:
        :return:
        """

        self.filename = filename
        self.folder = folder

        for kpi in self.kpi_list:
            kpi.setup(simulation_model)

        if isinstance(self.population_names, list):
            self.populations = [simulation_model.agentPopulation[population_name]
                                for population_name in self.population_names]
        else:
            self.populations = [simulation_model.agentPopulation[self.population_names]]

    def agent_kpi_dict(self, agent):
        """
        Computes the KPIs for the given agent
        by calling their update method for all its trace
        :param agent:
        :return:
        """

        indicators_dict = dict()

        # get agent trace
        events = agent.trace.eventList

        # evaluate all indicators in a single pass
        for event in events:
            for kpi in self.kpi_list:
                kpi.update(event, agent)

        # merge all completed indicators
        for kpi in self.kpi_list:

            indicators_dict.update(kpi.indicator_dict)

            # raising a warning with sphinx
            # indicators_dict = {**indicators_dict, **kpi.indicator_dict}

            # reset kpi values
            kpi.new_indicator_dict()

        # return complete indicator dict
        return indicators_dict

    def write_kpi_table(self):
        """
        Write the KPI of the population in the csv file
        obtained from out file attributes
        The KPIs evaluated are defined by the kpi_list attribute
        """

        # first row is always agent's id, then we add the kpi_list keys
        header_list = [KPI.KEY_ID]
        for kpi in self.kpi_list:
            header_list += kpi.keys

        path = self.folder + self.filename

        kpi_table = pd.DataFrame()

        # compute the kpi table for each population dict
        for population in self.populations:
            kpi_table = pd.concat([kpi_table, self.compute_population_kpi_table(population)])

        # do not generate a kpi output if the kpi table is empty
        if kpi_table.empty:
            return

        # generate kpi output

        logging.info("Generating KPI output in file " + path)

        try:
            # write the dataframe into a csv file
            kpi_table.to_csv(path, sep=";", index=False, columns=header_list)
        except KeyError as e:
            logging.warning("Could not generate kpi output {}, "
                            "error occurred : {}".format(path, e))

    def compute_population_kpi_table(self, population):
        """
        Compute a kpi table for the given population dict.

        :param population: population dict {id: agent}
        :return: DataFrame containing the KPI values
        """

        df_output = pd.DataFrame()

        for agent in population.values():

            # create kpi dict for the agent
            agent_indicators = self.agent_kpi_dict(agent)

            # build a dataframe from the dict
            if isinstance(agent_indicators[KPI.KEY_ID], list):
                df = pd.DataFrame(agent_indicators)
            else:
                df = pd.DataFrame(agent_indicators, index=[0])

            # append the dataframe to the total output
            df_output = pd.concat([df_output, df])

        return df_output

from starling_sim.basemodel.output.kpis import KPI

import logging
import pandas as pd


class KpiOutput:

    def __init__(self, population_names, kpi_list, kpi_name=None):

        # name of the kpi, will compose the kpi filename : <kpi_name>.csv
        if kpi_name is None:
            if isinstance(population_names, list):
                self.name = "_&_".join(population_names) + "_kpi"
            else:
                self.name = population_names + "_kpi"
        else:
            self.name = kpi_name

        # population of agent to evaluate
        self.population_names = population_names
        self.populations = None

        # list of kpi to evaluate the given agents
        self.kpi_list = kpi_list

        # output file
        self.filename = None
        self.folder = None

    def setup(self, filename, folder, simulation_model):
        """
        Setup method called during simulation setup.

        Sets the values of out file and folder, and call setup for KPIs.

        :param filename: .csv file
        :param folder:
        :param simulation_model:
        :return:
        """

        self.filename = filename
        self.folder = folder

        for kpi in self.kpi_list:
            kpi.setup(simulation_model)

        if isinstance(self.population_names, list):
            self.populations = [simulation_model.agentPopulation[population_name]
                                for population_name in self.population_names]
        else:
            self.populations = [simulation_model.agentPopulation[self.population_names]]

    def agent_kpi_dict(self, agent):
        """
        Computes the KPIs for the given agent
        by calling their update method for all its trace
        :param agent:
        :return:
        """

        indicators_dict = dict()

        # get agent trace
        events = agent.trace.eventList

        # evaluate all indicators in a single pass
        for event in events:
            for kpi in self.kpi_list:
                kpi.update(event, agent)

        # merge all completed indicators
        for kpi in self.kpi_list:

            indicators_dict.update(kpi.indicator_dict)

            # raising a warning with sphinx
            # indicators_dict = {**indicators_dict, **kpi.indicator_dict}

            # reset kpi values
            kpi.new_indicator_dict()

        # return complete indicator dict
        return indicators_dict

    def write_kpi_table(self):
        """
        Write the KPI of the population in the csv file
        obtained from out file attributes
        The KPIs evaluated are defined by the kpi_list attribute
        """

        # first row is always agent's id, then we add the kpi_list keys
        header_list = [KPI.KEY_ID]
        for kpi in self.kpi_list:
            header_list += kpi.keys

        path = self.folder + self.filename

        kpi_table = pd.DataFrame()

        # compute the kpi table for each population dict
        for population in self.populations:
            kpi_table = pd.concat([kpi_table, self.compute_population_kpi_table(population)])

        # do not generate a kpi output if the kpi table is empty
        if kpi_table.empty:
            return

        # generate kpi output

        logging.info("Generating KPI output in file " + path)

        try:
            # write the dataframe into a csv file
            kpi_table.to_csv(path, sep=";", index=False, columns=header_list)
        except KeyError as e:
            logging.warning("Could not generate kpi output {}, "
                            "error occurred : {}".format(path, e))

    def compute_population_kpi_table(self, population):
        """
        Compute a kpi table for the given population dict.

        :param population: population dict {id: agent}
        :return: DataFrame containing the KPI values
        """

        df_output = pd.DataFrame()

        for agent in population.values():

            # create kpi dict for the agent
            agent_indicators = self.agent_kpi_dict(agent)

            # build a dataframe from the dict
            if isinstance(agent_indicators[KPI.KEY_ID], list):
                df = pd.DataFrame(agent_indicators)
            else:
                df = pd.DataFrame(agent_indicators, index=[0])

            # append the dataframe to the total output
            df_output = pd.concat([df_output, df])

        return df_output

# name of the kpi, will compose the kpi filename : <kpi_name>.csv # population of agent to evaluate # list of kpi to evaluate the given agents # output file Setup method called during simulation setup.          Sets the values of out file and folder, and call setup for KPIs.          :param filename: .csv file         :param folder:         :param simulation_model:         :return: Computes the KPIs for the given agent         by calling their update method for all its trace         :param agent:         :return: # get agent trace # evaluate all indicators in a single pass # merge all completed indicators # raising a warning with sphinx # indicators_dict = {**indicators_dict, **kpi.indicator_dict} # reset kpi values # return complete indicator dict Write the KPI of the population in the csv file         obtained from out file attributes         The KPIs evaluated are defined by the kpi_list attribute # first row is always agent's id, then we add the kpi_list keys # compute the kpi table for each population dict # do not generate a kpi output if the kpi table is empty # generate kpi output # write the dataframe into a csv file Compute a kpi table for the given population dict.          :param population: population dict {id: agent}         :return: DataFrame containing the KPI values # create kpi dict for the agent # build a dataframe from the dict # append the dataframe to the total output