from Narsese import Budget
from .Item import Item
from typing import Generic, TypeVar
import random

ItemType = TypeVar("ItemType", bound=Item)

class WorkingMemory(Generic[ItemType]):

    def __init__(self, capacity):
        self.buff = set[ItemType]()  # buffer
        self.capacity = capacity

    def remove(self, item: Item):
        self.buff.discard(item)

    def insert(self, item: Item):
        if self.capacity <= 0 or len(self.buff) < self.capacity:
            self.buff.add(item)
            return None

        pv = item.budget_c.p
        item_min: Item = min(self.buff, key=lambda u: u.budget_c.p)
        if item_min.budget_c.p < pv:
            self.buff.remove(item_min)
            self.buff.add(item)
            return item_min
        else:
            return item

    def truncate(self):
        if self.capacity > 0 and self.capacity < len(self.buff):
            ltb_sorted = sorted(self.buff, key=lambda item: item.budget_c.q)
            self.buff = set(ltb_sorted[-self.capacity:])

    def set_capacity(self, capacity):
        self.capacity = capacity
        self.truncate()

    def item_maxpriority(self):
        item_max = max(self.buff, key=lambda u: u.budget_c.p)
        return item_max

    def item_random(self):
        """根据self.cache.buff的priority的分布，随机选取一个prototype"""
        priorities = [item.budget_c.p for item in self.buff]
        selected_item = random.choices(
            self.buff, weights=priorities, k=1)[0]
        return selected_item

    def __contains__(self, item: Item):
        return item in self.buff

    def __len__(self):
        return len(self.buff)
    
    def __iter__(self):
        return iter(self.buff)

