dws-city-res-solve/web_solve.py

"""Web interface for City Resource Optimization solver."""

from flask import Flask, render_template, request, jsonify
from ortools.sat.python import cp_model
import json
import solve
import io
import sys
from contextlib import redirect_stdout
from constraint_validator import make_constraint_lambda

app = Flask(__name__)


@app.route("/")
def index():
    """Serve the solver form."""
    return render_template(
        "solver.html",
        initial=solve.INITIAL,
        actions=solve.ENABLED_ACTIONS,
        agents=solve.AGENT_AVAILABILITY,
        num_steps=solve.NUM_STEPS,
    )


@app.route("/solve", methods=["POST"])
def solve_handler():
    """Handle solver requests from the form."""
    try:
        data = request.get_json()

        # Parse initial resources (divide by 10 since UI shows unscaled values)
        initial = tuple(
            int(float(data.get(f"initial_{r}", 3)) * 10) for r in ["E", "B", "S", "C"]
        )

        # Parse per-city arrivals and departures (dynamic number of cities)
        arrivals = {}
        for step in range(1, solve.NUM_STEPS + 1):
            arrivals[step] = []

        # Find all cities submitted (city_0_type, city_1_type, etc.)
        city_indices = set()
        for key in data.keys():
            if key.startswith("city_") and key.endswith("_type"):
                city_idx = int(key.split("_")[1])
                city_indices.add(city_idx)

        for city_idx in sorted(city_indices):
            city_type = data.get(f"city_{city_idx}_type")
            arrival_step_str = data.get(f"city_{city_idx}_arrival_step")
            departure_step_str = data.get(f"city_{city_idx}_departure_step")
            adjacent_str = data.get(f"city_{city_idx}_adjacent", "")

            if city_type and city_type != "none" and arrival_step_str:
                arrival_step = int(arrival_step_str)
                departure_step = (
                    int(departure_step_str)
                    if departure_step_str
                    else solve.NUM_STEPS + 1
                )

                # Parse adjacent cities (comma-separated, whitespace ignored)
                adjacent_to = []
                if adjacent_str:
                    adjacent_to = [
                        int(idx.strip())
                        for idx in adjacent_str.split(",")
                        if idx.strip()
                    ]

                city_data = {
                    "type": city_type,
                    "adjacent_to": adjacent_to,
                    "departure_step": departure_step,
                }

                # Parse vat values for foundries (if specified, use them; otherwise let normalize_city default to 1)
                if city_type == "F":
                    vat_e = data.get(f"city_{city_idx}_vat_E")
                    vat_b = data.get(f"city_{city_idx}_vat_B")
                    vat_s = data.get(f"city_{city_idx}_vat_S")
                    vats = {}
                    if vat_e:
                        vats["E"] = int(vat_e)
                    if vat_b:
                        vats["B"] = int(vat_b)
                    if vat_s:
                        vats["S"] = int(vat_s)
                    if vats:
                        city_data["vats"] = vats

                arrivals[arrival_step].append(city_data)

        # Parse enabled actions
        enabled_actions = {}
        for action_name in solve.ENABLED_ACTIONS.keys():
            enabled_actions[action_name] = data.get(f"action_{action_name}", False)

        # Parse agent availability
        agent_availability = {}
        for agent_name in solve.AGENT_AVAILABILITY.keys():
            steps_str = data.get(f"agent_{agent_name}_steps", "")
            if steps_str:
                agent_availability[agent_name] = [
                    int(s) for s in steps_str.split(",") if s
                ]
            else:
                agent_availability[agent_name] = []

        solve.AGENT_AVAILABILITY = agent_availability

        # Parse time limit
        time_limit = float(data.get("time_limit", 60.0))

        # Parse verbose flag
        verbose = data.get("verbose", True)

        # Parse fixed constraints
        fixed_choices = None
        fixed_actions_data = data.get("fixed_actions", {})
        fixed_governors_data = data.get("fixed_governors", {})

        if fixed_actions_data or fixed_governors_data:
            fixed_choices = {}

            if fixed_actions_data:
                fixed_choices["actions"] = {}
                for key, action in fixed_actions_data.items():
                    city, step = map(int, key.split(","))
                    fixed_choices["actions"][(city, step)] = action

            if fixed_governors_data:
                fixed_choices["governors"] = {}
                for key, enabled in fixed_governors_data.items():
                    city, step, agent = key.rsplit(",", 2)
                    city, step = int(city), int(step)
                    fixed_choices["governors"][(city, step, agent)] = enabled

        # Parse resource constraints
        resource_constraints = None
        resource_constraints_data = data.get("resource_constraints", [])
        if resource_constraints_data:
            resource_constraints = []
            for expr in resource_constraints_data:
                try:
                    constraint_lambda = make_constraint_lambda(expr)
                    resource_constraints.append(constraint_lambda)
                except Exception as e:
                    raise ValueError(f"Invalid resource constraint '{expr}': {e}")

        # Capture solver output
        output_buffer = io.StringIO()
        with redirect_stdout(output_buffer):
            solver, status = solve.solve(
                initial=initial,
                arrivals=arrivals,
                max_res=solve.MAX_RES,
                max_vat=solve.MAX_VAT,
                # min to avoid bricking stuff
                # time_limit=min(time_limit, 60.0),
                time_limit=time_limit,
                num_workers=8,
                verbose=verbose,
                fixed_choices=fixed_choices,
                resource_constraints=resource_constraints,
            )

        output = output_buffer.getvalue()
        status_name = solver.StatusName(status) if status else "UNKNOWN"

        return jsonify({"success": True, "status": status_name, "output": output})

    except Exception as e:
        return jsonify({"success": False, "error": str(e)}), 400


if __name__ == "__main__":
    app.run(debug=False, host="0.0.0.0", port=5000)