#===================================================================================================#
#                                         Simulator                                                 #
#    Last Modification: 13.03.2020                                         Mauricio Fadel Argerich  #
#===================================================================================================#

import cloudpickle
from collections import OrderedDict, defaultdict
import itertools
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import psutil
import sys
import time

from entities import AdASDevice, AdASCpu, AdASExecutionData


class AdASSimulator:

    def __init__(self, profile_path):
        """
        Initiates a simulator with a profile created by the AdASProfiler.
        Keyword arguments:
        - profile_path: path to profile_device.pkl generated by an AdASProfiler.
        """
        self.profile = None
        self.__load_profile(profile_path)
            
    def __load_profile(self, filename):
        profile_to_load = None
        with open(filename, "rb") as f:
            profile_to_load = cloudpickle.load(f)

        self.profile = profile_to_load

    def get_utility(self, param_values):
        u = 0
        for f in self.profile.pipeline_data:
            f_params = {}
            for p, v in param_values.items():
                if f.params.get(p):
                    f_params[p] = v

            u += f.get_utility(f_params)
        return u

    def sim(self, input_data, param_values, available_cpu = 1.0):
        res_stats = defaultdict(list)
        # f_output is the input_data of the next function in pipeline.
        f_output = input_data
        
        # Simulate execution of each function in pipeline.
        for f in self.profile.pipeline_data:
            # Select params for function and put them in a dict.
            f_params = {}
            for p, v in param_values.items():
                if f.params.get(p):
                    f_params[p] = v

            # Simulate execution of function.
            exec_data = f.sim(f_output, f_params)
            f_output = exec_data.output_data
            res_stats['utility'] = exec_data.stats.get('utility')
            for k, v in exec_data.stats.items():
                if 'latency' in k:
                    res_stats[k].append(v / available_cpu)

        return res_stats