from pydantic import BaseModel
from typing import List, Literal, Optional

import logging

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: list[Bundle] = []
    constraints: List[Constraint]
    irt_model: IRTModel
    objective_function: ObjectiveFunction
    total_form_items: int
    total_forms: int = 1
    theta_cut_score: float = 0.00
    drift_style: Literal['constant', 'variable'] = 'constant'
    advanced_options: Optional[AdvancedOptions]
    engine: str

    def get_item(self, item_id: int) -> Item or None:
        for item in self.items:
            if str(item.id) == item_id:
                return item

    def get_bundle(self, bundle_id: int) -> Bundle or None:
        for bundle in self.bundles:
            if str(bundle.id) == bundle_id:
                return bundle

    def get_constraint_by_type(self, type: str) -> Constraint or None:
        for constraint in self.constraints:
            if type == constraint.reference_attribute.type:
                return constraint

    def remove_items(self, items: list[Item]) -> bool:
        self.items = [item for item in self.items if item not in items]
        return True

    def generate_bundles(self):
        logging.info('Generating Bundles...')
        # confirms bundle constraints exists
        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,
                                       items=[item],
                                       type=type_attribute))
                        else:
                            self.bundles[bundle_index].count += 1
                            self.bundles[bundle_index].items.append(item)
                    else:
                        self.bundles = [
                            Bundle(id=attribute_id,
                                   count=1,
                                   items=[item],
                                   type=type_attribute)
                        ]

        logging.info('Bundles Generated...')

    def get_constraint(self, name: str) -> Constraint:
        return next((constraint for constraint in self.constraints
                     if constraint.reference_attribute.id == name), None)

    def unbundled_items(self) -> list:
        return [item for item in self.items if item.passage_id == None]