from pydantic import BaseModel
from typing import ClassVar, List

from models.item import Item
from lib.irt.models.three_parameter_logistic import ThreeParameterLogistic
from lib.irt.models.rasch import Rasch

class AbilityEstimation(BaseModel):
  exam_id: int
  items: List[Item] = []
  irt_model: str
  min_theta: float = -3.0
  max_theta: float = 3.0

  IRT_MODELS: ClassVar[dict] = {
    'rasch': Rasch,
    # not supported
    # '3pl': ThreeParameterLogistic
  }

  def calculate(self) -> float:
    if self.irt_model in self.IRT_MODELS:
      model = self.IRT_MODELS[self.irt_model]
      return model.ability_estimate(self.items)
    else:
      logging.error(f'model of type {self.irt_model} does not exist.')
      return None