import math
from catopt_query.protocol import LocalProblem, SharedVariables, DualVariables, PlanDelta, DataContract, AuditLog
from catopt_query.core import map_local_to_canonical, aggregate_joint_plan
from catopt_query.core import CanonicalPlan


def test_local_problem_serialization():
    lp = LocalProblem(
        shard_id="shard-1",
        projected_attributes=["a", "b"],
        predicates=["a > 0", "b < 100"],
        costs={"cpu": 1.2, "io": 0.5},
        constraints={"timezone": "UTC"},
    )
    d = lp.to_dict()
    assert d["shard_id"] == "shard-1"
    assert d["predicates"] == ["a > 0", "b < 100"]


def test_local_to_canonical_mapping():
    lp = LocalProblem(
        shard_id="s1",
        projected_attributes=["x"],
        predicates=["x IS NOT NULL"],
        costs={"cpu": 2.0},
    )
    can = map_local_to_canonical(lp)
    assert isinstance(can, CanonicalPlan)
    assert can.total_cost == 2.0
    assert len(can.operations) == 1