#!/usr/bin/env python

from plugins.node_plugin import *

# PowerStates allows you to measure the energy consumption of a
# node that go through several power states during the simulation
# Two version of Powerstates is provided by mean of two classes:
#   - Powerstates: Allow you to set power to any user's defined values
#   - PowerstatesFromFile: Allow you to set power from states defined in a file

######################
#     _    ____ ___  #
#    / \  |  _ \_ _| #
#   / _ \ | |_) | |  #
#  / ___ \|  __/| |  #
# /_/   \_\_|  |___| #
#                    #
######################

# #Regarding PowerStates:
# import Powerstates as ps
# pstates=ps.PowerStates(<node>,<power_init>)
# pstates.set_power(<power>) # Switch the power consumption to <power>
# pstates.report_energy() # Display the current node energy consumption up to the current simulated time
# pstates.report_power_changes() # Display all the power changes up to the current simulated time

# #Regarding PowerStatesFromFile:
# #Format of <file> is one <entry> per line that follow this format <state-0>:<state-1>:...:<state-n>
# #Each line can corresponds to one node
# import Powerstates as ps
# pstates=ps.PowerStatesFromFile(<node>,<file>,<entry-line>) # Create a power states on node <node> using line <entry-line> of file <file>
# pstates.set_state(<id>) # Switch to the <id> power states
# pstates.report_energy() # Display the current node energy consumption up to the current simulated time
# pstates.report_power_changes() # Display all the power changes up to the current simulated time
# pstates.report_state_changes() # Display all the states changes up to the current simulated time


class PowerStates(NodePlugin):
    """
    PowerStates model the energy consumed by the various changes of power consumption of a node over time.
    """
    def __init__(self,node,power_init):
        self.node=node
        self.clock=self.node.read("clock")
        self.energy=0
        self.power=power_init
        self.power_changes=dict()
        self.set_power(power_init)
        super().__init__("Powerstates",api)
        

    def set_power(self,power_watt):
        cur_clock=self.node.read("clock")
        self.energy+=self.power*(cur_clock-self.clock)
        self.clock=cur_clock
        if self.power != power_watt:
            self.power_changes[cur_clock]=power_watt
        self.power=power_watt
        return cur_clock

    def report_energy(self):
        self.set_power(self.power)
        self.node.log("[PowerStates Plugin] Consumed "+str(self.energy) +"J")
        
    def report_power_changes(self):
        self.set_power(self.power)
        for key in self.power_changes.keys():
            self.node.log("[PowerStates Plugin] At t="+str(key)+" power is "+str(self.power_changes[key])+"W")



class PowerStatesFromFile(PowerStates):
    """
    A version of Powerstates that load the power values from a file.
    """
    def __init__(self,node,state_file,entry_line=1):
        self.node=node
        self.state_changes=dict()
        self.states=[]
        self.state=0
        with open(state_file) as fp:
            for i, line in enumerate(fp):
                if i+1 == entry_line:
                    self.states=line.split(":")
                    self.states=[float(i) for i in self.states]
        assert len(self.states) > 0
        super().__init__(node,self.states[0])
        self.set_state(0)

    def set_state(self,state_id):
        assert state_id < len(self.states)
        clock=super().set_power(self.states[state_id])
        if self.state != state_id:
            self.state_changes[clock]=state_id
        self.state=state_id


    def report_state_changes(self):
        self.set_state(self.state)
        for key in self.state_changes.keys():
            self.node.log("[PowerStates Plugin] At t="+str(key)+" state is "+str(self.state_changes[key]))

        
class PowerStatesComms(NodePlugin):
    """
    Monitor the energy consumed by the network interfaces by mean of power states.
    Note that for finer grained predictions, bytes and packet power consumption must be accounted.
    Which is not the case with these power states.
    """

    def __init__(self,api):
        super().__init__("PowerStatesComms",api)
        self.energy_dynamic=0.0 # Store the dynamic part of the energy consumption
        self.power=dict() # Store the power states
        self.tx_clock=0 # Dynamic clock (store the time at which a the last tx starts
        self.idle_clock=api.read("clock") # Store the start time (to compute the idle part of the energy consumption)

    def on_receive_return(self,interface,data,start_at,end_at):
        duration=float(end_at)-float(start_at)
        self.energy_dynamic+=self.power[interface]["rx"]*duration

    def on_send_call(self,interface,data,datasize,dst):
        self.tx_clock=self.api.read("clock")

    def on_send_return(self,interface,data,datasize,dst,code):
        clock=self.api.read("clock")
        duration=(clock-float(self.tx_clock))
        self.energy_dynamic+=self.power[interface]["tx"]*duration
        self.tx_clock=clock # Any value could be use here

    def set_power(self,interface,idle,tx,rx):
        self.power[interface]=dict()
        self.power[interface]["idle"]=idle
        self.power[interface]["rx"]=rx
        self.power[interface]["tx"]=tx

    def get_idle(self):
        clock=self.api.read("clock")
        idle=0
        for interface in self.power.keys():
            idle+=(clock-self.idle_clock)*self.power[interface]["idle"]
        return idle

    def get_receive_queue_energy(self,interface):
        """
        Not that call to on_receive_return may not have happened yet (or never).
        Thus we should manually compute the energy consumption stored in each queues of the node.
        """
        energy=0
        # For each interface we should check if there is received data that has not been consumed
        for data in list(self.api["interfaces"][interface].queue):
            start_at=float(data[1])
            end_at=float(data[2])
            energy+=(end_at-start_at)*self.power[interface]["rx"]
        return energy
    
    def get_energy(self):
        queue_energy=0
        for interface in self.power.keys():
            queue_energy+=self.get_receive_queue_energy(interface)
        return self.get_idle()+self.energy_dynamic+queue_energy

    def report_energy(self):
        self.log("Communications consumed "+str(round(self.get_energy(),2))+"J")