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.bundle import Bundle
from models.objective_function import ObjectiveFunction
from models.advanced_options import AdvancedOptions

class SolverRun(BaseModel):
  items: List[Item]
  bundles: Optional[Bundle]
  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 generate_bundles(self):
    bundle_constraints = (constraint.reference_attribute for constraint in self.constraints if constraint.reference_attribute.type == 'bundle')

    for bundle_constraint in bundle_constraints:
      type_attribute = bundle_constraint.id

      for item in self.items:
        attribute_id = getattr(item, type_attribute, None)

        # make sure the item has said attribute
        if attribute_id != None:
          # if there are pre-existing bundles, add new or increment existing
          # else create array with new bundle
          if self.bundles != None:
            # get index of the bundle in the bundles list if exists or None if it doesn't
            bundle_index = next((index for (index, bundle) in enumerate(self.bundles) if bundle.id == attribute_id and bundle.type == type_attribute), None)

            # if the index doesn't exist add the new bundle of whatever type
            # else increment the count of the current bundle
            if bundle_index == None:
              self.bundles.append(Bundle(
                id=attribute_id,
                count=1,
                type=type_attribute
              ))
            else:
              self.bundles[bundle_index].count += 1
          else:
            self.bundles = [Bundle(
              id=attribute_id,
              count=1,
              type=type_attribute
            )]