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.attribute import Attribute
from models.item import Item
from models.irt_model import IRTModel

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

    def find_items(self, requested_items_ids: [int]) -> [Item]:
        found_items = []

        for item in self.items:
            if item.id in requested_items_ids:
                found_items.append(item)

        return found_items


    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)

    def tif_trf_sum(self, solver_run):
        return self.__trf_sum(solver_run) + self.__tif_sum(solver_run)

    def items_with_attribute(self, attribute: Attribute) -> List[Item]:
        items = []

        for item in self.items:
            if item.attribute_exists(attribute): items.append(item)

        return items

    def __tif_sum(self, solver_run):
        total = 0

        for target in solver_run.objective_function.tcc_targets:
            total += self.tif(solver_run.irt_model, target.theta)

        return total

    def __trf_sum(self, solver_run):
        total = 0

        for target in solver_run.objective_function.tcc_targets:
            total += self.trf(solver_run.irt_model, target.theta)

        return total

    def ordered_items(self) -> List[Item]:
        return sorted(self.items, key=lambda item: item.position)