import heapq
import random
import matplotlib.pyplot as plt
import numpy as np


class Event:
    def __init__(self, time, event_type, node):
        self.time = time
        self.type = event_type
        self.node = node

    def __lt__(self, other):
        return self.time < other.time

    def __repr__(self):
        return f'<{self.time}-{self.type}>'


class Node:
    def __init__(self, node_id, channel, event_queue):
        # 常量定义（单位：微秒 μs）
        self.DIFS = 43  # DCF帧间距，载波监听时长
        self.SIFS = 16  # 短帧帧间距，接收节点成功接收数据后等待时长，然后发送回ACK
        self.SLOT_TIME = 9  # 时隙长度，随机回退间隙
        self.ACK_TIMEOUT = 65  # ACK超时时间，数据传输失败，发送节点等待ACK时长
        self.ACK_DURATION = 32  # ACK帧传输时间，接收节点成功接收数据后，在等待短帧帧间距后，发送回ACK的传输时长
        self.PAYLOAD = 1500
        self.MACLOAD = 30
        self.PHY_LONG = 13.6
        self.RATE = 54
        # self.DATA_DURATION = 100  # 数据帧传输时间（简化值），数据传输时长
        self.DATA_DURATION = self.PHY_LONG + (self.MACLOAD + self.PAYLOAD) / self.RATE * 8
        # 二进制指数退避参数
        self.CW_MIN = 16  # 最小竞争窗口
        self.MAX_STAGE = 6  # 最大竞争窗口
        self.MAX_RETRY = 32  # 最大重传次数

        self.id = node_id
        self.channel = channel
        self.event_queue = event_queue
        # 初始状态
        self.state = "IDLE"  # 状态：IDLE, WAIT_DIFS, BACKOFF, TRANSMITTING, WAIT_ACK
        self.backoff_counter = 0
        self.cw = self.CW_MIN
        self.retry_count = 0 # ack 结束时清零

        self.transmission_count = 0  # 开始时核算

        self.success_count = 0  # 结束时核算
        self.collision_count = 0  # 结束时核算
        self.drop_count = 0

    def update(self, **kwargs):
        for key, value in kwargs.items():
            self.__setattr__(key, value)
        self.DATA_DURATION = self.PHY_LONG + (self.MACLOAD + self.PAYLOAD) / self.RATE * 8
        self.cw = self.CW_MIN

    def __repr__(self):
        return f'<{self.id}-{self.state}>'

    def __delay_init__(self, channel):
        self.channel = channel

    def generate_traffic(self):
        event = Event(0, "WAIT_DIFS_START", self)
        self.schedule_event(event)

    def schedule_event(self, event):
        heapq.heappush(self.event_queue, event)

    # def handle_idle_start(self, event):
    #     self.state = 'IDLE'
    #     self.backoff_counter = 0
    #     self.cw = CW_MIN
    #     self.retry_count = 0
    #     # 规划一个空闲结束事件
    #
    #     idle_time = int(np.random.exponential(IDLE_MEAN))
    #
    #     event = Event(event.time + idle_time, 'IDLE_END', self)
    #     self.schedule_event(event)
    #
    # def handle_idle_end(self, event):
    #     event = Event(event.time, "WAIT_DIFS_START", self)
    #     self.schedule_event(event)

    def handle_wait_difs_start(self, event):
        self.state = "WAIT_DIFS"
        event = Event(event.time + self.DIFS, "WAIT_DIFS_END", self)
        self.schedule_event(event)

    def handle_wait_difs_end(self, event):
        if self.channel.sensor_channel_busy[self]:
            event = Event(self.channel.sensor_busy_until[self], "WAIT_DIFS_START", self)
        else:
            event = Event(event.time, "BACK_OFF_START", self)
            if self.backoff_counter == 0:
                self.backoff_counter = np.random.randint(0, self.cw)
        self.schedule_event(event)

    def handle_back_off_start(self, event):
        self.state = "BACKOFF"
        event = Event(event.time + self.SLOT_TIME, "BACK_OFF_END", self)
        self.schedule_event(event)

    def handle_back_off_end(self, event):
        if self.backoff_counter == 0:
            event = Event(event.time, "TRANSMITTING_START", self)
        else:
            if self.channel.sensor_channel_busy[self]:
                event = Event(self.channel.sensor_busy_until[self], "WAIT_DIFS_START", self)
            else:
                self.backoff_counter -= 1
                event = Event(event.time, "BACK_OFF_START", self)
        self.schedule_event(event)

    def handle_transmitting_start(self, event):
        self.state = "TRANSMITTING"
        self.transmission_count += 1
        # 记录传输开始时间（用于时隙统计）
        if self.channel.channel_busy:
            self.channel.heads_broken[self] = True
        self.channel.heads_busy[self] = True
        self.channel.busy_untils[self] = event.time + self.DATA_DURATION
        event = Event(event.time + self.DATA_DURATION, "TRANSMITTING_END", self)
        self.schedule_event(event)

    def handle_transmitting_end(self, event):
        if self.channel.busy_head_except_node(self):
            self.channel.heads_broken[self] = True
        if self.channel.heads_broken[self]:
            event = Event(event.time, "ACK_TIMEOUT_START", self)
        else:
            event = Event(event.time, "SIFS_START", self)
        # 处理信道
        self.channel.heads_busy[self] = False
        self.channel.busy_untils[self] = 0
        self.channel.heads_broken[self] = False
        self.schedule_event(event)

    def handle_ack_timeout_start(self, event):
        self.state = "WAIT_ACK"
        event = Event(event.time + self.ACK_TIMEOUT, "ACK_TIMEOUT_END", self)
        self.schedule_event(event)

    def handle_ack_timeout_end(self, event):
        if self.retry_count > self.MAX_RETRY:
            self.cw = self.CW_MIN
            self.retry_count = 0
            self.drop_count += 1
            event = Event(event.time, "WAIT_DIFS_START", self)
        else:
            self.retry_count += 1
            self.cw = min(2 * self.cw, self.CW_MIN * 2 ** self.MAX_STAGE)
            event = Event(event.time, "WAIT_DIFS_START", self)
        self.schedule_event(event)
        # 核算
        self.collision_count += 1

    def handle_sifs_start(self, event):
        self.state = "WAIT_ACK"
        event = Event(event.time + self.SIFS, "SIFS_END", self)
        self.schedule_event(event)

    def handle_sifs_end(self, event):
        event = Event(event.time, "ACK_START", self)
        self.schedule_event(event)

    def handle_ack_start(self, event):
        self.channel.heads_busy[self] = True
        self.channel.busy_untils[self] = event.time + self.ACK_DURATION
        event = Event(event.time + self.ACK_DURATION, "ACK_END", self)
        self.schedule_event(event)

    def handle_ack_end(self, event):
        self.channel.heads_busy[self] = False
        self.channel.busy_untils[self] = 0
        # 清零
        self.cw = self.CW_MIN
        self.retry_count = 0  # ack 结束时清零
        # 核算
        self.success_count += 1
        # 规划
        event = Event(event.time, "WAIT_DIFS_START", self)
        self.schedule_event(event)

    def handle_event(self, event):
        if event.type == "WAIT_DIFS_START":
            self.handle_wait_difs_start(event)
        elif event.type == "WAIT_DIFS_END":
            self.handle_wait_difs_end(event)
        elif event.type == "BACK_OFF_START":
            self.handle_back_off_start(event)
        elif event.type == "BACK_OFF_END":
            self.handle_back_off_end(event)
        elif event.type == "TRANSMITTING_START":
            self.handle_transmitting_start(event)
        elif event.type == "TRANSMITTING_END":
            self.handle_transmitting_end(event)
        elif event.type == "ACK_TIMEOUT_START":
            self.handle_ack_timeout_start(event)
        elif event.type == "ACK_TIMEOUT_END":
            self.handle_ack_timeout_end(event)
        elif event.type == "SIFS_START":
            self.handle_sifs_start(event)
        elif event.type == "SIFS_END":
            self.handle_sifs_end(event)
        elif event.type == "ACK_START":
            self.handle_ack_start(event)
        elif event.type == "ACK_END":
            self.handle_ack_end(event)
        else:
            print("some thing must be wrong!")


class WirelessChannel:
    def __init__(self, nodes, listens):
        self.nodes = nodes
        self.heads_busy = {node: False for node in nodes}
        self.busy_untils = {node: 0 for node in nodes}
        self.heads_broken = {node: False for node in nodes}
        self.listens = listens

    @property
    def channel_busy(self):
        """获取当前信道繁忙值"""
        if True in self.heads_busy.values():
            return True
        else:
            return False

    @property
    def sensor_channel_busy(self):
        output = {node: False for node in self.heads_busy.keys()}
        for node, listen in self.listens.items():
            sensors = [self.heads_busy[lis] for lis in listen]
            if True in sensors:
                output[node] = True
            else:
                output[node] = False
        return output

    @property
    def busy_until(self):
        return max(self.busy_untils.values())

    @property
    def sensor_busy_until(self):
        output = {node: 0 for node in self.nodes}
        for node, listen in self.listens.items():
            output[node] = max([self.busy_untils[lis] for lis in listen])
        return output

    def busy_head_except_node(self, node):
        busy = False
        for key, value in self.heads_busy.items():
            if key != node:
                busy = value
        return busy

    def sensor_busy_head_except_node(self, node):
        return True in [self.heads_busy[lis] for lis in self.listens[node]]


class Simulator:
    def __init__(self, STATIONS=2, total_time=1e4, listen=[[1], [0]]):
        self.event_queue = []
        self.num_node = STATIONS
        self.listen = listen
        self.nodes = [Node(i, None, self.event_queue) for i in range(STATIONS)]
        self.listens = {node: [self.nodes[idx] for idx in listen[idx]] for idx, node in enumerate(self.nodes)}
        self.channel = WirelessChannel(self.nodes, self.listens)
        [node.__delay_init__(self.channel) for node in self.nodes]
        self.total_time = total_time
        self.current_time = 0
        # even_history
        # self.events_history = {node: [] for node in self.nodes}
        # self.states_history = {node: [] for node in self.nodes}

    def update(self, **kwargs):
        self.__reset__()
        for node in self.nodes:
            node.update(**kwargs)

    def __reset__(self):
        self.event_queue = []
        self.nodes = [Node(i, None, self.event_queue) for i in range(self.num_node)]
        self.listens = {node: [self.nodes[idx] for idx in self.listen[idx]] for idx, node in enumerate(self.nodes)}
        self.channel = WirelessChannel(self.nodes, self.listens)
        [node.__delay_init__(self.channel) for node in self.nodes]
        self.total_time = self.total_time
        self.current_time = 0
        # even_history
        # self.events_history = {node: [] for node in self.nodes}
        # self.states_history = {node: [] for node in self.nodes}

    def run(self):
        # 启动所有节点
        [node.generate_traffic() for node in self.nodes]
        while self.current_time < self.total_time:
            # 处理事件
            event = heapq.heappop(self.event_queue)
            event.node.handle_event(event)
            # print(self.current_time, event.node, event)
            # 更新
            self.current_time = event.time
            # self.events_history[event.node].append(event)
            # self.states_history[event.node].append((self.current_time, event.node.state))

        # 核算
        transmission_count = 0
        success_count = 0
        drop_count = 0
        for node in self.nodes:
            transmission_count += node.transmission_count
            success_count += node.success_count
            drop_count += node.drop_count
        ratio = success_count * self.nodes[0].DATA_DURATION / self.total_time
        drop_rate = drop_count / (drop_count + success_count)
        throughput = ratio * self.nodes[0].RATE
        # 新增核算

        return drop_rate, ratio, throughput

