#!/usr/bin/env python3
"""Demo: ADMM-like optimization with two local problems"""
from catopt_flow_category_theoretic_compositi import LocalProblem, DeltaSyncRegistry, ADMMNode, run_admm, Planner


def main():
    lp1 = LocalProblem(id="node-1", resources={"gpu": 2.0, "memory": 8.0}, data_loading={}, batch_size=32)
    lp2 = LocalProblem(id="node-2", resources={"gpu": 1.0, "memory": 8.0}, data_loading={}, batch_size=32)
    registry = DeltaSyncRegistry()
    n1 = ADMMNode("node-1", lp1, registry)
    n2 = ADMMNode("node-2", lp2, registry)
    final = run_admm([n1, n2], rounds=3, budget_gpu=4.0, registry=registry)
    print("Final global plan:")
    for lid, alloc in final.plan.items():
        print(f"  {lid}: {alloc}")

    # Show a planner-only view from local problems (example of usage)
    plan = Planner.build_global_plan([lp1, lp2], budget_gpu=4.0)
    print("Planner plan:", plan.plan)


if __name__ == "__main__":
    main()