from pulp import lpSum, LpProblem
from random import randint, sample

import logging

from helpers.common_helper import *

from models.bundle import Bundle
from models.solver_run import SolverRun
from models.item import Item

from lib.errors.item_generation_error import ItemGenerationError

# should probably be factored out into a bundle class method or a method in the solver run
def build_constraints(solver_run: SolverRun, problem: LpProblem,
                      items: list[Item], bundles: list[Bundle], selected_items: list[Item], selected_bundles: list[Bundle], current_drift: int) -> LpProblem:
    logging.info('Creating Constraints...')

    try:
        total_form_items = solver_run.total_form_items
        constraints = solver_run.constraints

        for constraint in constraints:
            attribute = constraint.reference_attribute
            min = constraint.minimum
            max = constraint.maximum

            if attribute.type == 'metadata':
                logging.info('Metadata Constraint Generating...')

                problem += lpSum(
                    [
                        len(bundle.items_with_attribute(attribute)) * bundles[bundle.id] for bundle in selected_bundles
                    ] +
                    [
                        item.attribute_exists(attribute).real * items[item.id] for item in selected_items
                    ]
                ) >= round(total_form_items * (min / 100)), f'{attribute.id} - {attribute.value} - min'

                problem += lpSum(
                    [
                        len(bundle.items_with_attribute(attribute)) * bundles[bundle.id] for bundle in selected_bundles
                    ] +
                    [
                        item.attribute_exists(attribute).real * items[item.id] for item in selected_items
                    ]
                ) <= round(total_form_items * (max / 100)), f'{attribute.id} - {attribute.value} - max'
            elif attribute.type == 'bundle':
                logging.info('Bundles Constraint Generating...')
                # TODO: account for many different bundle types, since the id condition in L33 could yield duplicates
                if selected_bundles != None and selected_bundles > 0:
                    # make sure the total bundles used in generated form is limited between min-max set
                    problem += lpSum([
                        bundles[bundle.id] for bundle in selected_bundles
                    ]) == randint(int(constraint.minimum),
                                  int(constraint.maximum))

        logging.info('Constraints Created...')
        
        # Behold our very own Elastic constraints!
        for tif_target in solver_run.objective_function.tif_targets:
            problem += lpSum([
                bundle.tif(solver_run.irt_model, tif_target.theta)
                * bundles[bundle.id]
                for bundle in selected_bundles
            ] + [
                item.iif(solver_run, tif_target.theta) *
                items[item.id]
                for item in selected_items
            ]) >= tif_target.minimum(
            ), f'Min TIF theta({tif_target.theta}) at target {tif_target.value} drift at {current_drift}%'
            problem += lpSum([
                bundle.tif(solver_run.irt_model, tif_target.theta)
                * bundles[bundle.id]
                for bundle in selected_bundles
            ] + [
                item.iif(solver_run, tif_target.theta) *
                items[item.id]
                for item in selected_items
            ]) <= tif_target.maximum(
            ), f'Max TIF theta({tif_target.theta}) at target {tif_target.value} drift at {current_drift}%'

        for tcc_target in solver_run.objective_function.tcc_targets:
            problem += lpSum([
                bundle.trf(solver_run.irt_model, tcc_target.theta)
                * bundles[bundle.id]
                for bundle in selected_bundles
            ] + [
                item.irf(solver_run, tcc_target.theta) *
                items[item.id]
                for item in selected_items
            ]) >= tcc_target.minimum(
            ), f'Min TCC theta({tcc_target.theta}) at target {tcc_target.value} drift at {current_drift}%'
            problem += lpSum([
                bundle.trf(solver_run.irt_model, tcc_target.theta)
                * bundles[bundle.id]
                for bundle in selected_bundles
            ] + [
                item.irf(solver_run, tcc_target.theta) *
                items[item.id]
                for item in selected_items
            ]) <= tcc_target.maximum(
            ), f'Max TCC theta({tcc_target.theta}) at target {tcc_target.value} drift at {current_drift}%'
        
        return problem
    except ValueError as error:
        logging.error(error)
        raise ItemGenerationError(
            "Bundle min and/or max larger than bundle amount provided",
            error.args[0])

# should probably be factored out into a bundle class method or a method in the solver run
def get_random_bundles(total_form_items: int,
                       bundles: list[Bundle],
                       min: int,
                       max: int,
                       found_bundles=False) -> list[Bundle]:
    selected_bundles = None
    total_bundle_items = 0
    total_bundles = randint(min, max)
    logging.info(f'Selecting Bundles (total of {total_bundles})...')

    while found_bundles == False:
        selected_bundles = sample(bundles, total_bundles)
        total_bundle_items = sum(bundle.count for bundle in selected_bundles)

        if total_bundle_items <= total_form_items:
            found_bundles = True

    if found_bundles == True:
        return selected_bundles
    else:
        return get_random_bundles(total_form_items, total_bundles - 1, bundles)