#!/usr/bin/env python3
"""
Pyomo 6.9.2 Implementation for Construction Company Optimization Problem
"""

import pyomo.environ as pyo
from pyomo.opt import SolverFactory

def optimization_problem():
    """Optimize the allocation of projects to architects to minimize total length of bridges and mills built."""
    
    # 1. MODEL CREATION
    model = pyo.ConcreteModel()
    
    # 2. DATA SETUP
    # Bridge data
    bridge_ids = [1, 2, 3]
    bridge_lengths = {1: 120, 2: 250, 3: 400}  # in meters
    
    # Mill data
    mill_ids = [1, 2, 3]
    mill_lengths = {1: 350, 2: 450, 3: 550}  # in feet
    
    # Architect data
    architect_ids = [1, 2, 3]
    min_projects = {1: 1, 2: 1, 3: 1}
    max_projects = {1: 3, 2: 3, 3: 3}
    
    # CRITICAL: Validate array lengths before indexing
    assert len(bridge_ids) == len(bridge_lengths), "Bridge data length mismatch"
    assert len(mill_ids) == len(mill_lengths), "Mill data length mismatch"
    assert len(architect_ids) == len(min_projects) == len(max_projects), "Architect data length mismatch"
    
    # 3. SETS (Pyomo way to define indices)
    model.B = pyo.Set(initialize=bridge_ids)  # Set of bridges
    model.M = pyo.Set(initialize=mill_ids)    # Set of mills
    model.A = pyo.Set(initialize=architect_ids)  # Set of architects
    
    # 4. PARAMETERS (data containers)
    model.bridge_length = pyo.Param(model.B, initialize=bridge_lengths)
    model.mill_length = pyo.Param(model.M, initialize=mill_lengths)
    model.min_projects = pyo.Param(model.A, initialize=min_projects)
    model.max_projects = pyo.Param(model.A, initialize=max_projects)
    
    # 5. VARIABLES
    # Binary variables indicating whether a bridge or mill is built
    model.x = pyo.Var(model.B, within=pyo.Binary)  # x_b
    model.y = pyo.Var(model.M, within=pyo.Binary)  # y_m
    
    # Binary variables indicating architect assignments
    model.z = pyo.Var(model.A, model.B, within=pyo.Binary)  # z_{a,b}
    model.w = pyo.Var(model.A, model.M, within=pyo.Binary)  # w_{a,m}
    
    # 6. OBJECTIVE FUNCTION
    def obj_rule(model):
        return sum(model.bridge_length[b] * model.x[b] for b in model.B) + \
               sum(model.mill_length[m] * model.y[m] for m in model.M)
    model.objective = pyo.Objective(rule=obj_rule, sense=pyo.minimize)
    
    # 7. CONSTRAINTS
    
    # Minimum Projects per Architect
    def min_projects_rule(model, a):
        return sum(model.z[a, b] for b in model.B) + sum(model.w[a, m] for m in model.M) >= model.min_projects[a]
    model.min_projects_constraint = pyo.Constraint(model.A, rule=min_projects_rule)
    
    # Maximum Projects per Architect
    def max_projects_rule(model, a):
        return sum(model.z[a, b] for b in model.B) + sum(model.w[a, m] for m in model.M) <= model.max_projects[a]
    model.max_projects_constraint = pyo.Constraint(model.A, rule=max_projects_rule)
    
    # Bridge Assignment
    def bridge_assignment_rule(model, a, b):
        return model.z[a, b] <= model.x[b]
    model.bridge_assignment_constraint = pyo.Constraint(model.A, model.B, rule=bridge_assignment_rule)
    
    # Mill Assignment
    def mill_assignment_rule(model, a, m):
        return model.w[a, m] <= model.y[m]
    model.mill_assignment_constraint = pyo.Constraint(model.A, model.M, rule=mill_assignment_rule)
    
    # 8. SOLVING WITH GUROBI
    solver = SolverFactory('gurobi')
    
    # Optional: Set solver options
    solver.options['TimeLimit'] = 300  # 5 minutes
    solver.options['MIPGap'] = 0.01    # 1% gap
    
    # Solve the model
    results = solver.solve(model, tee=True)  # tee=True shows solver output
    
    # 9. RESULT PROCESSING
    # Check solver status
    if results.solver.termination_condition == pyo.TerminationCondition.optimal:
        print(f"Optimal value: {pyo.value(model.objective)}")
        
        # Extract variable values
        print("\nVariable values:")
        for b in model.B:
            if pyo.value(model.x[b]) > 0:
                print(f"Bridge {b} is built.")
        for m in model.M:
            if pyo.value(model.y[m]) > 0:
                print(f"Mill {m} is built.")
        for a in model.A:
            for b in model.B:
                if pyo.value(model.z[a, b]) > 0:
                    print(f"Architect {a} is assigned to bridge {b}.")
            for m in model.M:
                if pyo.value(model.w[a, m]) > 0:
                    print(f"Architect {a} is assigned to mill {m}.")
        
    elif results.solver.termination_condition == pyo.TerminationCondition.infeasible:
        print("Problem is infeasible")
    elif results.solver.termination_condition == pyo.TerminationCondition.unbounded:
        print("Problem is unbounded")
    else:
        print(f"Solver terminated with condition: {results.solver.termination_condition}")
    
    return model

# Execute the optimization problem
if __name__ == "__main__":
    optimization_problem()