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

import logging
import random

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 item.id == item_id:
                return item

    def get_bundle(self, bundle_id: int) -> Bundle or None:
        for bundle in self.bundles:
            if 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)
                        ]
        # temporary compensator for bundle item limits, since we shouldn't be using cases with less than 3 items
        # ideally this should be in the bundles model as a new attribute to handle "constraints of constraints"
        logging.info('Removing bundles with items < 3')
        for k, v in enumerate(self.bundles):
            bundle = self.bundles[k]
            if bundle.count < 3: del self.bundles[k]

        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:
        # since the only bundles are based on passage id currently
        # in the future when we have more than just passage based bundles
        # we'll need to develop a more sophisticated way of handling this concern
        bundle_constraints = (
            constraint.reference_attribute for constraint in self.constraints
            if constraint.reference_attribute.type == 'bundle')

        if len(list(bundle_constraints)) > 0:
            return [item for item in self.items if item.passage_id == None]
        else:
            return self.items

    def select_items_by_percent(self, percent: int) -> list[Item]:
        items = self.unbundled_items()
        total_items = len(items)
        selected_items_amount = round(total_items - (total_items *
                                                     (percent / 100)))

        return random.sample(items, selected_items_amount)

    def select_bundles_by_percent(self, percent: int) -> list[Bundle]:
        total_bundles = len(self.bundles)
        selected_bundles_amount = round(total_bundles - (total_bundles *
                                                         (percent / 100)))

        return random.sample(self.bundles, selected_bundles_amount)