# Fix paths for imports to work in unit tests ----------------

if __name__ == "__main__":

    from _fix_paths import fix_paths
    fix_paths()

# ------------------------------------------------------------

# Load libraries ---------------------------------------------

import pandas as pd
from datetime import timedelta
import pickle
import os

from ssa_sim_v2.tools.directory_utils import ensure_dir
from copy import deepcopy

# ------------------------------------------------------------


class ExperimentHistorySaver:

    def __init__(self, root_directory, seed, experiment_spec, policy_infos):
        self.root_directory = root_directory
        self.result_lists = {policy_info["name"]: [] for policy_info in policy_infos}
        self.seed = seed
        self.experiment_spec = experiment_spec

    def save_experiment_iteration(self, policy_name, policy,
                                  profit, state, action, effective_action,
                                  info, attr_info):
        """
        Save all possible values generated in a single experiment iteration. This method saves current profit, state,
        action, policy history, prior used in modules, results generated by simulator.
        :param str policy_name: Policy name.
        :param Policy policy: Working policy in the experiment.
        :param float profit: Profit of the experiment iteration.
        :param State state: State of the experiment.
        :param Action action: Action given by the policy in current experiment iteration.
        :param Action effective_action: Action used in current experiment iteration.
        :param dict info: Results from the simulator.
        :param dict attr_info: Results for the attributes from the simulator.
        :return:
        """
        iteration_result_dict = dict()

        iteration_result_dict["profit"] = profit
        iteration_result_dict["date"] = state.date

        if hasattr(state, "how"):
            iteration_result_dict["how"] = state.how

        iteration_result_dict["bid"] = action.bid
        iteration_result_dict["effective_bid"] = effective_action.bid

        # Get policy history values
        if hasattr(policy, "save_history") and policy.save_history:
            agent_history_dict = deepcopy(policy.history)
            iteration_result_dict.update({"policy.hist.{}".format(key): value
                                          for key, value in agent_history_dict.items()})

            # UDP params
            udp_dict = policy.udp.__dict__
            iteration_result_dict.update({"policy.udp.{}".format(key): value
                                          for key, value in udp_dict.items()})

            # HTP params
            htp_dict = policy.htp.__dict__
            iteration_result_dict.update({"policy.htp.{}".format(key): value
                                          for key, value in htp_dict.items()})

            # STP params
            stp_dict = policy.stp.__dict__
            iteration_result_dict.update({"policy.stp.{}".format(key): value
                                          for key, value in stp_dict.items()})

            # IP params
            ip_dict = policy.ip.__dict__
            iteration_result_dict.update({"policy.ip.{}".format(key): value
                                          for key, value in ip_dict.items()})

        # Get simulator generated values
        iteration_result_dict.update({"simulator.{}".format(key): value for key, value in info.items()})

        # Store attribute results

        if attr_info is not None:
            # TODO: Code a choice of attributes to save
            pass

        self.result_lists[policy_name].append(iteration_result_dict)

    def generate_experiment_results(self, save=False, csv=False):

        results = {}

        for policy_name, result_list in self.result_lists.items():

            experiment_df = pd.DataFrame(result_list)
            experiment_df["datetime"] = pd.to_datetime(experiment_df["date"], format="%Y-%m-%d")
            if "how" in experiment_df.columns:
                experiment_df["datetime"] = experiment_df.apply(
                    lambda row: row["datetime"] + timedelta(hours=(row["how"] - 24 * row["datetime"].weekday())),
                    axis=1
                )

            # Final structure of the experiment result
            result = {
                "action_experiment_df": experiment_df,
                "seed": self.seed,
                "policy_name": policy_name
            }

            if save:
                filename = "{}_m{}_p{}.pickle".format(
                    policy_name.replace(" ", "_"), self.seed["mod"], self.seed["pol"])

                file_path = os.path.join(self.root_directory, filename)
                ensure_dir(file_path)

                with open(file_path, "wb") as file:
                    pickle.dump(result, file)

                if csv:
                    filename = "{}_m{}_p{}.csv".format(
                        policy_name.replace(" ", "_"), self.seed["mod"], self.seed["pol"])

                    file_path = os.path.join(self.root_directory, filename)
                    ensure_dir(file_path)

                    # Save Action result
                    result["action_experiment_df"].to_csv(file_path, index=False)

            results[policy_name] = result

        return results
