from pydantic import BaseModel
from typing import List, Optional

from models.item import Item
from models.constraint import Constraint
from models.irt_model import IRTModel
from models.objective_function import ObjectiveFunction
from models.advanced_options import AdvancedOptions

class SolverRun(BaseModel):
  items: List[Item]
  constraints: List[Constraint]
  irt_model: IRTModel
  objective_function: ObjectiveFunction
  total_form_items: int
  total_forms: int = 1
  theta_cut_score: float = 0.00
  advanced_options: Optional[AdvancedOptions]
  engine: str

  def get_item(self, item_id):
    for item in self.items:
      if str(item.id) == item_id:
        return item
    return False

  def remove_items(self, items):
    self.items = [item for item in self.items if item not in items]
    return True

  def bundles(self, type_attribute):
    bundles = []

    for item in self.items:
      # dynamically get the attribute that will be used as an identifier to bundle like items
      attribute_id = getattr(item, type_attribute, None)()

      # make sure the item has said attribute
      if attribute_id != None:
        # get index of the bundle in the bundles list
        bundle_index = next((index for (index, bundle) in enumerate(bundles) if bundle['id'] == attribute_id), None)

        # if the bundle index isn't found then the bundle hasn't been created
        # and added to the list and needs to be, else increment the count of items
        # in the bundle
        if bundle_index == None:
          bundles.append({ 'id': attribute_id, 'count': 1, 'type': type_attribute })
        else:
          bundles[bundle_index]['count'] += 1

    return bundles