from pydantic import BaseModel
from typing import List

from lib.irt.test_information_function import TestInformationFunction
from lib.irt.test_response_function import TestResponseFunction

from models.item import Item
from models.irt_model import IRTModel


class Bundle(BaseModel):
    id: int
    count: int
    items: List[Item]
    type: str

    def tif(self, irt_model: IRTModel, theta: float) -> float:
        return TestInformationFunction(irt_model).calculate(self.items,
                                                            theta=theta)

    def trf(self, irt_model: IRTModel, theta: float) -> float:
        return TestResponseFunction(irt_model).calculate(self.items,
                                                         theta=theta)