from itertools import islice, compress, chain
from gym import error, spaces, utils
from gym.utils import seeding
import pandas as pd
import numpy as np
import gym

class RuminantEnv(gym.Env):
    metadata = {'render.modes': ['human']}

    def __init__(self, df, ratio):
        super(RuminantEnv, self).__init__()

        self.MAX_STEPS = df.shape[0]

        # Dataframe
        self.df = df

        # List of unique pages
        self.pages = df.page_id.unique()

        # Memory lists
        self.PM_ls = []
        self.DRAM_ls = []

        # Ratio DRAM/PM size
        self.ratio = ratio

        # Total Reward
        self.total_reward = 0

        # Counter to track amount of time PM is accessed
        self.PM_accessed = 0

        # We need something for the environment to initiate with full list of pages in DRAM and PM
        # fill up function the PM_ls and DRAM_ls (if DRAM_ls is filled, continue filling PM_ls)
        self.fill_up_memory(self.DRAM_ls, self.PM_ls, ratio, self.pages)

        # List for temporal purposes
        self.times = df.time_accessed.unique()

        # Ruminant variables
        self.temp_DRAM_avg = 0
        self.temp_PM_avg = 0
        self.lifetime_avg_DRAM = 0
        self.prev_avg_temp = 0

        # Reward Function Variables
        self.prev_avg_temp = 0
        self.prev_lifetime_avg_DRAM = 0
        self.prev_temp_diff = 0

        # Dictionary structure
        self.dict_st = {}
        # Create dictionary
        self.make_dict(self.pages, self.dict_st)

        # Promotion and demotion list which will be used to hold page candidates
        # promo_ls will store potentials pages from PM to promote
        # demote_ls will store protential pages from DRAM to demote
        self.promo_ls = []
        self.demote_ls = []

        # Observation space (for every time step, see all 3 ruminant variables)
        self.observation_space = spaces.Box(low=0, high=100, shape=(1, 3), dtype=np.float64)

        # Actions (demote: move to PM list, promote: move to DRAM list, unchanged: stay in current list)
        self.action_space = spaces.Discrete(3)

    # Resetting the state of the environment to initial state
    def reset(self):
        # Reset Ruminant variables
        self.temp_DRAM_avg = 0
        self.temp_PM_avg = 0
        self.lifetime_avg_DRAM = 0

        # Reset Promotion and Demotion Lists
        self.promo_ls = []
        self.demote_ls = []

        # Reset Dictionary structure
        self.dict_st = {}
        self.pages = self.df.page_id.unique()
        self.ratio = self.ratio

        # Create dictionary
        self.make_dict(self.pages, self.dict_st)

        # Reset/Clear Memory Lists
        self.PM_ls = []
        self.DRAM_ls = []
        self.fill_up_memory(self.DRAM_ls, self.PM_ls, self.ratio, self.pages)

        # Reset Current Step
        self.current_step = 0

        # Reset Total Reward
        self.total_reward = 0

        # Reset Counter to track amount of time PM is accessed
        self.PM_accessed = 0

        return self._next_observation()

    def _next_observation(self):
        # Get the new variables in current state
        obs = np.array([self.temp_DRAM_avg, self.temp_PM_avg, self.lifetime_avg_DRAM])

        return obs

    def step(self, action):
        # Take one time step within the environment
        self._take_action(action)

        # Discount_factor formulated from observation space for exploration-exploitation (discount-factor)
        discount_factor = self.current_step / self.MAX_STEPS

        # Stores previous temp before current step is increased
        self.prev_avg_temp = self.temp_DRAM_avg
        # Stores previous lifetime average before current step is increased
        self.prev_lifetime_avg_DRAM = self.lifetime_avg_DRAM

        # Stores previous temperature difference before current step is increased
        self.prev_temp_diff = self.temp_DRAM_avg - self.temp_PM_avg

        # How do we define our environment as done. What stops the simulation? (Current step reaches size of dataframe)
        done = (self.current_step >= self.MAX_STEPS)

        # Increment current step, (i.e, step function completed given iteration)
        self.current_step += 1

        # Update dict_st using unique_times lists: "update database" and increase temp(page was accessed)
        self.update_time_dict(self.current_step, self.df, self.dict_st, self.times, self.PM_accessed, self.PM_ls)

        # Initialize 'curr_' or 'current' variables
        curr_temp_DRAM_avg = 0
        curr_temp_PM_avg = 0
        curr_lifetime_avg_DRAM = 0

        # Define 'current' variables using helper functions
        curr_temp_DRAM_avg = self.calc_DRAM_avg_temp(self.DRAM_ls, self.dict_st)
        curr_temp_PM_avg = self.calc_PM_avg_temp(self.PM_ls, self.dict_st)
        curr_lifetime_avg_DRAM = self.calc_avg_DRAM_lifetime(self.DRAM_ls, self.dict_st)
        curr_temp_diff = curr_temp_DRAM_avg - curr_temp_PM_avg

        # Calculate the avg for all variables
        self.temp_DRAM_avg = self.calc_DRAM_avg_temp(self.DRAM_ls, self.dict_st)
        self.temp_PM_avg = self.calc_PM_avg_temp(self.PM_ls, self.dict_st)
        self.lifetime_avg_DRAM = self.calc_avg_DRAM_lifetime(self.DRAM_ls, self.dict_st)

        # reward function and logic
        reward = 0
        if(curr_temp_diff >= self.prev_temp_diff and curr_lifetime_avg_DRAM >= self.prev_lifetime_avg_DRAM):
            reward = 1

        elif(curr_temp_diff < self.prev_temp_diff or curr_lifetime_avg_DRAM < self.prev_lifetime_avg_DRAM):
            reward = -1

        else:
            reward = -5

        reward *= discount_factor

        # Keep track of total_reward
        self.total_reward += reward

        obs = self._next_observation()

        return obs, reward, done, {}

    def _take_action(self, action):
        # If action 1(promote) is done, what is done
        if action == 0:
            # DRAM must always be filled
            # Look for possible candidates in PM of hot pages and place in promote list
            self.promo_ls = [(self.dict_st[page]['temp'] > self.temp_PM_avg) for page in self.PM_ls]
            self.promo_ls = list(compress(self.PM_ls, self.promo_ls))

            # Look for possible candidates(focus on a better conditional) in DRAM of cold pages and place in demote list
            # and (self.dict_st[page]['lifetime'] < self.lifetime_avg_DRAM)
            self.demote_ls = [(self.dict_st[page]['temp'] < self.temp_DRAM_avg) for page in self.DRAM_ls]
            self.demote_ls = list(compress(self.DRAM_ls, self.demote_ls))

            # So that demote_ls is equal to length of promo_ls and DRAM is full at all times (what if it has less pages??)
            # If the length of demote_ls is less than that of promo_ls
            if (len(self.demote_ls) < len(self.promo_ls)) and len(self.demote_ls) != 0:
                promo_ls_copy = self.promo_ls.copy()
                promo_ls_copy = list(islice(self.promo_ls, 0, len(self.demote_ls)))
                self.PM_ls = [page for page in self.PM_ls if page not in promo_ls_copy]
                self.DRAM_ls = [page for page in self.DRAM_ls if page not in self.demote_ls]

                # Set lifetime to 0 (i.e, page has changed location)
                for page in self.demote_ls:
                    self.dict_st[page]['lifetime'] = 0

                # Extending the excess to prevent loss of pages in DRAM and PM
                self.DRAM_ls.extend(promo_ls_copy)
                self.PM_ls.extend(self.demote_ls)

            # Opposite to the above if condition, takes care when demote list is greater in size than promo list
            elif (len(self.demote_ls) > len(self.promo_ls)) and len(self.promo_ls) != 0:
                demote_ls_copy = self.demote_ls.copy()
                demote_ls_copy = list(islice(self.demote_ls, 0, len(self.promo_ls)))
                self.DRAM_ls = [page for page in self.DRAM_ls if page not in demote_ls_copy]
                self.PM_ls = [page for page in self.PM_ls if page not in self.promo_ls]

                # Set lifetime to 0 (i.e, page has changed location)
                for page in demote_ls_copy:
                    self.dict_st[page]['lifetime'] = 0

                # Extending the excess to prevent loss of pages in DRAM and PM
                self.DRAM_ls.extend(self.promo_ls)
                self.PM_ls.extend(demote_ls_copy)

            #else condition: if promo or demote list are of size 0 or empty, no swap is performed
            else:
                self.promo_ls = []
                self.demote_ls = []

            self.promo_ls = []
            self.demote_ls = []

        # If action 2(demote) is done
        elif action == 1:
            # Search DRAM for possible cold pages(focus on a better conditional) and place in demote list and remove from DRAM
            self.demote_ls = [(self.dict_st[page]['temp'] < self.temp_DRAM_avg) for page in self.DRAM_ls]
            self.demote_ls = list((compress(self.DRAM_ls, self.demote_ls)))

            # Update lifetime = 0 if demoted
            for page in self.demote_ls:
                self.dict_st[page]['lifetime'] = 0

            # Remove from DRAM
            self.DRAM_ls = [page for page in self.DRAM_ls if page not in self.demote_ls]

            # Extend page_ids in demote_list in PM_list
            self.PM_ls.extend(self.demote_ls)

        # If action 3(nothing) is done
        elif action == 2:
            # All lists stay the same
            self.DRAM_ls = self.DRAM_ls

            # Update and increase the lifetime of each page in DRAM_ls
            for page in self.DRAM_ls:
                self.dict_st[page]['lifetime'] += 1

            # Stays the same (nothing occurs)
            self.PM_ls = self.PM_ls

    def render(self, mode='human', close=False):
        # Render environment
        print(f'Step: {self.current_step}')
        print(f'Average DRAM temperature: {self.temp_DRAM_avg}, Average PM temperature: {self.temp_PM_avg}, Average Lifetime: {self.lifetime_avg_DRAM}')
        # Temp_difference (closer to 0: good, closer to 1: bad)
        # temp_difference = self.temp_PM_avg/self.temp_DRAM_avg

        # num_times_PM_accessed (if high number, bad)
        num_times_PM_accessed = self.PM_accessed

        # print(f'Temperature difference: {temp_difference}')
        print(f'Amount of time PM is accessed: {num_times_PM_accessed}')
        print(f'Total Reward: {self.total_reward}')

    def update_time_dict(self, current_step, df, dict_st, times, PM_accessed, PM_ls):
        page_ids = [page for page in df.page_id.loc[df.time_accessed == times[current_step]]]
        for page in page_ids:
            if page in PM_ls:
                self.PM_accessed += 1
                dict_st[page]['temp'] += 1
            else:
                dict_st[page]['temp'] += 1

    def make_dict(self, pages, dict_st):
        for page in pages:
            dict_st[page] = {'temp': 0, 'lifetime': 0}

    def calc_DRAM_avg_temp(self, DRAM_ls, dict_st):
        total_temp = 0
        for page in DRAM_ls:
            total_temp += dict_st[page]['temp']
        return round(total_temp / len(DRAM_ls), 3)

    def calc_PM_avg_temp(self, PM_ls, dict_st):
        total_temp = 0
        for page in PM_ls:
            total_temp += dict_st[page]['temp']
        return round(total_temp / len(PM_ls), 3)

    def calc_avg_DRAM_lifetime(self, DRAM_ls, dict_st):
        total_lifetime = 0
        for page in DRAM_ls:
            total_lifetime += dict_st[page]['lifetime']
        return round(total_lifetime / len(DRAM_ls), 3)

    def get_page_info(self, page, dict_st):
        temp = dict_st[page]['temp']
        lifetime = dict_st[page]['lifetime']
        return temp, lifetime

    def fill_up_memory(self, DRAM_ls, PM_ls, ratio, pages):
        for page in range(int(len(pages) * ratio)):
            DRAM_ls.append(pages[page])
        for page in range(len(DRAM_ls), int(len(pages))):
            PM_ls.append(pages[page])


    # Metrics to test baseline vs hyperparameter tuned
    # -temp_difference (if high difference, good)
    # -num_times_PM_accessed (if high number, bad)
    # -total_reward
