# Complete PYOMO implementation - Retry Attempt 4

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

def optimize_research_allocation():
    # 1. MODEL CREATION
    model = pyo.ConcreteModel()

    # 2. DATA SETUP
    # Example data from the Authorship table
    authorship_data = [
        (1, 101, 201, 1),
        (2, 102, 202, 2),
        (3, 103, 203, 3)
    ]

    # Extract unique sets
    A = set(authID for authID, _, _, _ in authorship_data)  # Set of researchers
    I = set(instID for _, instID, _, _ in authorship_data)  # Set of institutions
    P = set(paperID for _, _, paperID, _ in authorship_data)  # Set of papers

    # Weights based on authOrder
    weights = {(authID, paperID): 1.0 / authOrder for authID, _, paperID, authOrder in authorship_data}

    # Operational constraints
    MaxResearchersPerInstitution = 2
    MaxPapersPerResearcher = 2

    # 3. SETS
    model.A = pyo.Set(initialize=A)
    model.I = pyo.Set(initialize=I)
    model.P = pyo.Set(initialize=P)

    # 4. PARAMETERS
    model.weights = pyo.Param(model.A, model.P, initialize=weights)
    model.MaxResearchersPerInstitution = pyo.Param(initialize=MaxResearchersPerInstitution)
    model.MaxPapersPerResearcher = pyo.Param(initialize=MaxPapersPerResearcher)

    # 5. VARIABLES
    model.x = pyo.Var(model.A, model.I, model.P, within=pyo.Binary)

    # 6. OBJECTIVE FUNCTION
    def obj_rule(model):
        return sum(model.weights[a, p] * model.x[a, i, p] for a in model.A for i in model.I for p in model.P)
    model.objective = pyo.Objective(rule=obj_rule, sense=pyo.maximize)

    # 7. CONSTRAINTS
    # Maximum Researchers per Institution
    def max_researchers_rule(model, i):
        return sum(model.x[a, i, p] for a in model.A for p in model.P) <= model.MaxResearchersPerInstitution
    model.max_researchers_constraint = pyo.Constraint(model.I, rule=max_researchers_rule)

    # Maximum Papers per Researcher
    def max_papers_rule(model, a):
        return sum(model.x[a, i, p] for i in model.I for p in model.P) <= model.MaxPapersPerResearcher
    model.max_papers_constraint = pyo.Constraint(model.A, rule=max_papers_rule)

    # Single Assignment per Paper
    def single_assignment_rule(model, p):
        return sum(model.x[a, i, p] for a in model.A for i in model.I) == 1
    model.single_assignment_constraint = pyo.Constraint(model.P, rule=single_assignment_rule)

    # 8. SOLVING WITH GUROBI
    solver = SolverFactory('gurobi')
    results = solver.solve(model, tee=True)

    # 9. RESULT PROCESSING
    if results.solver.termination_condition == pyo.TerminationCondition.optimal:
        print(f"Optimal value: {pyo.value(model.objective)}")
    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}")

if __name__ == "__main__":
    optimize_research_allocation()