from __future__ import annotations
from typing import TYPE_CHECKING, Optional

from pulp import lpSum

from models import Constraint, Problem, Attribute, Target, Item, Bundle

if TYPE_CHECKING:
    from models.solver_run import SolverRun

class IrtTargetConstraint(Constraint):
    reference_attribute: Optional[Attribute]
    minimum: Optional[float]
    maximum: Optional[float]
    target: Target
    target_type: str

    def build(self, problem_handler: Problem, solver_run: SolverRun):
        problem_handler.problem += lpSum([
            self.bundle_irt_function(bundle, solver_run.irt_model, self.target.theta)
            * problem_handler.solver_bundles_var[bundle.id]
            for bundle in problem_handler.bundles
        ] + [
            self.item_irt_function(item, solver_run.irt_model, self.target.theta) *
            problem_handler.solver_items_var[item.id]
            for item in problem_handler.items
        ]) >= self.target.minimum(
        ), f'Min {self.target_type} theta({self.target.theta}) at target {self.target.value}'

        problem_handler.problem += lpSum([
            self.bundle_irt_function(bundle, solver_run.irt_model, self.target.theta)
            * problem_handler.solver_bundles_var[bundle.id]
            for bundle in problem_handler.bundles
        ] + [
            self.item_irt_function(item, solver_run.irt_model, self.target.theta) *
            problem_handler.solver_items_var[item.id]
            for item in problem_handler.items
        ]) <= self.target.maximum(
        ), f'Max {self.target_type} theta({self.target.theta}) at target {self.target.value}'

    def item_irt_function(self, item: Item, irt_model: str, theta: float) -> float:
        if self.target_type == 'tcc':
            return item.irf(irt_model, theta)
        elif self.target_type == 'tif':
            return item.iif(irt_model, theta)

    def bundle_irt_function(self, bundle: Bundle, irt_model: str, theta: float) -> float:
        if self.target_type == 'tcc':
            return bundle.trf(irt_model, theta)
        elif self.target_type == 'tif':
            return bundle.tif(irt_model, theta)