from pydantic import BaseModel
from typing import Any, List
from pulp import LpProblem, LpVariable, lpSum

import logging, math

from helpers import solver_helper

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

class Problem(BaseModel):
    items: List[Item]
    bundles: List[Bundle]
    problem: Any
    solver_items_var: Any = None
    solver_bundles_var: Any = None
    
    def __init__(self, **data) -> None:
        super().__init__(**data)
      
        # setup common Solver variables
        self.solver_items_var = LpVariable.dicts("Item",
                                    [item.id for item in self.items],
                                    lowBound=0,
                                    upBound=1,
                                    cat='Binary')
        self.solver_bundles_var = LpVariable.dicts("Bundle",
                                    [bundle.id for bundle in self.bundles],
                                    lowBound=0,
                                    upBound=1,
                                    cat='Binary')
        
        # objective function
        self.problem += lpSum([
            bundle.count * self.solver_bundles_var[bundle.id]
            for bundle in self.bundles
        ] + [
            self.solver_items_var[item.id]
            for item in self.items
        ])
    
    def solve(self) -> LpProblem:
        self.problem.solve()
        return self.problem
    
    def generate(self, solution: Solution, solver_run: SolverRun):    
        
        # Form Constraints
        self.problem += lpSum(
            [
                bundle.count * self.solver_bundles_var[bundle.id]
                for bundle in self.bundles
            ] + [
                1 * self.solver_items_var[item.id]
                for item in self.items
            ]
        ) == solver_run.total_form_items, f'Total bundle form items for form'
        
        # each time a form is generated, we want to ensure
        # that it is unique to all other forms generated before it
        self.problem += lpSum(
            [
                solution.items_exist_in_forms(bundle.items) * self.solver_bundles_var[bundle.id]
                for bundle in self.bundles
            ] + [
                solution.items_exist_in_forms([item]) * self.solver_items_var[item.id]
                for item in self.items
            ]
        ) <= solver_run.total_form_items - 1, f'Ensuring uniqueness for form'
    
    def generate_constraints(self, solver_run: SolverRun, current_drift: int):
        self.problem = solver_helper.build_constraints(
                    solver_run, self.problem, self.solver_items_var, self.solver_bundles_var, self.items, self.bundles, current_drift)