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

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

class WorkingLSTM(Generic[ItemType]):
    def __init__(self, n_ltb, n_stb):
        self.capacity = n_ltb + n_stb
        self.buffer = Buffer[ItemType](self.capacity)
        self.lstm = LSTM[ItemType](n_ltb, n_stb)

    def remove(self, item: Item):
        self.buffer.remove(item)
        self.lstm.remove(item)

    def insert(self, item: Item):
        of_item = self.lstm.insert(item)
        if of_item is not None and of_item in self.buffer:
            self.buffer.remove(of_item)
        if item is not of_item:
            self.buffer.insert(item)

    def set_capacity(self, n_ltb, n_stb):
        self.capacity = n_ltb + n_stb
        of_items = self.lstm.set_capacity(n_ltb, n_stb)
        for item in of_items:
            self.buffer.remove(item)

    def item_maxpriority_c(self):
        return self.buffer.item_maxpriority()


    def item_random_c(self):
        priorities = [item.budget_c.p for item in self.buffer]
        selected_item = random.choices(
            self.buffer, weights=priorities, k=1)[0]
        return selected_item

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

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

