# Choices and Aggregates

**Choice rules** allow you to give the solver a decision to make: choose some of a given set to
include in the answer set. Unlike facts, rules and constraints, which apply mechanically, choice
rules create branches in the decision tree. They are the construct that *generates* atoms.

The partner to choice rules is **aggregates**, which allow you to bound or constrain atoms,
typically those that have been selected by choice rules. Aggregates allow you to count, sum, or take
extremes.

This is the shape of most ASP programs: generate the space of possibilities, then prune it. To see
how it works, we'll build a running example on this page - assigning a workshop's tasks to its
workers.

## Choice rules

A choice rule hands the solver a decision: from a set of candidate atoms, make some subset true. The
syntax to build a choice rule is `Choice(element, condition=...)`. The *element* is the atom being
chosen, and the *condition* says where the candidates come from. You get one candidate for each way
that the conditions can be satisfied.

For each choice rule, you can put constraints on the number of entries that should be included
(known as the cardinality). If no constraints are given, any number (ranging from none to all) are
possibilities.

To add a choice rule to a program, use `program.when(...).derive([choice rule])`.

```python
from aspalchemy import ASPProgram, Choice, Field, Predicate, Variable

class Task(Predicate, show=False):
    name: Field[str]

class Worker(Predicate, show=False):
    name: Field[str]

class Assigned(Predicate):
    task: Field[str]
    worker: Field[str]

T, W = Variable("T"), Variable("W")

program = ASPProgram()
program.fact(*(Task(name=t) for t in ["wiring", "plumbing", "painting"]))
program.fact(*(Worker(name=w) for w in ["ada", "grace"]))

# Every task goes to exactly one worker
program.when(Task(name=T)).derive(
    Choice(Assigned(task=T, worker=W), condition=Worker(name=W)) == 1
)
```

```python
>>> print(program.render())
% Generated by aspalchemy ...
task("wiring").
task("plumbing").
task("painting").
worker("ada").
worker("grace").
{ assigned(T, W) : worker(W) } = 1 :- task(T).

#show.
#show assigned/2.
```

Read that last line as: for each task `T`, pick *exactly one* `assigned` atom, drawing the worker
`W` from the workers we know about. We use `== 1` to specify the cardinality for this choice - pick
a single `assigned` atom. Here are some other ways to express a cardinality:

```python
>>> offers = Choice(Assigned(task=T, worker=W), condition=Worker(name=W))
>>> (offers >= 2).render()
'2 { assigned(T, W) : worker(W) }'
>>> (offers != 3).render()
'{ assigned(T, W) : worker(W) } != 3'
>>> banded = 2 <= offers <= 4
>>> banded.render()
'2 { assigned(T, W) : worker(W) } 4'
```

Cardinalities are checked as you build — a negative size, or an impossible band like
`3 <= offers <= 2`, is refused right at the Python line rather than quietly becoming an
unsatisfiable program.

A choice rule with no `when()` conditions attaches directly with `program.choose()`, and holds
unconditionally:

```python
class OnCall(Predicate):
    worker: Field[str]

# A bare choice rule: hold at least one worker on call, unconditionally
coverage = Choice(OnCall(worker=W), condition=Worker(name=W))
standby = coverage >= 1
program.choose(standby)
```

A single choice can offer more than one kind of atom: `add(element, condition=...)` appends another,
and the cardinality then bounds the whole set:

```python
class Approve(Predicate):
    task: Field[str]

class Defer(Predicate):
    task: Field[str]
```

```python
>>> decision = Choice(Approve(task=T))
>>> decision.add(Defer(task=T))          # add() mutates, and returns nothing
>>> decision.render()
'{ approve(T); defer(T) }'
```

The bounds needn't be literal integers: a `Variable` or `Expression` bound in the rule's conditions
works too, so your *data* can decide how many atoms to choose.

## Choices as Python Objects

A `Choice` is an ordinary Python object you can name, reuse, and share between rules. You can
use one in multiple places, as we now demonstrate. Each comparison creates
a new object that wraps the original choice. In this example, we supply different constraints
for weekdays and weekends:

```python
>>> roster = Choice(OnCall(worker=W), condition=Worker(name=W))
>>> weekday, weekend = roster == 1, roster == 2
>>> weekday.render()
'{ on_call(W) : worker(W) } = 1'
>>> weekend.render()
'{ on_call(W) : worker(W) } = 2'
>>> roster.render()   # the choice itself is untouched, and still reusable
'{ on_call(W) : worker(W) }'
```

Note that the choice rule is shared, so adding extra elements to `roster` causes them to show up
everywhere it's used. (The `add` method mutates the `Choice` object, just like `list.append`.)

```python
>>> roster.add(OnCall(worker="hot_spare"))
>>> weekday.render()   # both bands see the new element
'{ on_call(W) : worker(W); on_call("hot_spare") } = 1'
```

There's a gotcha here to be aware of. Say that we create a choice, use it in a rule,
and then add another element to it. Because the choice is a python object, this mutates the
object, and it's unclear what the rule should do in this case: should it store the choice as
it existed when the rule was specified, or should it store the mutated choice? In order to avoid
confusion, we *freeze* a choice once it's been used inside a rule, so that `add()` refuses to
mutate it afterwards. That's what happened to `coverage` earlier: `program.choose()` used it,
through `standby`, so it's now frozen.

```python
>>> coverage.add(OnCall(worker=W))  # coverage was used, through standby, by program.choose() above
Traceback (most recent call last):
  ...
RuntimeError: This Choice is already used by the rule at ... and is frozen; changing it would silently rewrite that rule. Call .copy() for a fresh, mutable Choice with the same elements, or build a new Choice.
```

The error names the way out: `copy()` returns an independent, mutable choice with the same elements.

```python
>>> extended = coverage.copy()   # a fresh, mutable Choice
>>> extended.add(OnCall(worker="hot_spare"))
>>> extended.render()
'{ on_call(W) : worker(W); on_call("hot_spare") }'
>>> standby.render()
'1 { on_call(W) : worker(W) }'
```

Note that frozen choices can still be used in new rules and cardinalities, they just can't be added
to.


## Aggregates

An aggregate is a construct that reduces a set of atoms to a single value that you can then bound.
There are five — `Count`, `Sum`, `SumPlus`, `Min`, and `Max`, and they're constructed the same way
as `Choice` objects. They behave the same way too: you can `add()` elements until a rule uses the
aggregate, at which point it freezes, and `copy()` gives you a fresh, mutable one.

```python
from aspalchemy import ANY, Count, Field, Predicate, PredicateArg, Variable

class Person(Predicate):
    name: Field[PredicateArg]
    age: Field[PredicateArg]

john = Person(name="john", age=30)
mary = Person(name="mary", age=25)
alan = Person(name="alan", age=41)

X = Variable("X")
count = Count(X, condition=Person(name=X, age=ANY)) > 5
```

`Count` counts distinct matching tuples. The other four reduce the **first** term of each tuple:
`Sum` adds it up (`SumPlus` is the same, except negative weights count as zero), and `Min` and `Max`
take the extremes. So a tuple's leading term is its *weight*, and any trailing terms are there only
to keep tuples distinct — `Sum((A, X), condition=Person(name=X, age=A))` sums one age per person,
even when two people happen to share an age.

```python
>>> from aspalchemy import Max, Sum
>>> A = Variable("A")
>>> total_age = Sum((A, X), condition=Person(name=X, age=A))
>>> total_age.render()
'#sum{ A, X : person(X, A) }'
>>> oldest = Max(A, condition=Person(name=ANY, age=A))
>>> oldest.render()
'#max{ A : person(_, A) }'
```

Aggregates must be used as part of a comparison, which drops in wherever comparisons go: `when()`,
`forbid()` and `require()`. The `count` variable built above is already a comparison. Handing a bare
aggregate to a rule as though it were an atom is refused at construction. Here's a simple example of
using an aggregate in a rule:

```python
census = ASPProgram()
census.fact(john, mary, alan)
census.require(Count(X, condition=Person(name=X, age=ANY)) >= 3)  # requires at least 3 people
```

There's one comparison that works a little differently. If you compare an aggregate to a variable
with `==`, clingo doesn't just test the value: it *sets* the variable to it, and you can then use
that variable like any other. Here we record how many tasks each worker ends up with:

```python
class Load(Predicate):
    worker: Field[str]
    tasks: Field[int]

N = Variable("N")
program.when(Worker(name=W), N == Count(T, condition=Assigned(task=T, worker=W))).derive(
    Load(worker=W, tasks=N)
)
```

```python
>>> (N == Count(T, condition=Assigned(task=T, worker=W))).render()
'N = #count{ T : assigned(T, W) }'
```

Every worker gets a `Load` atom, even one with no tasks at all, because `Worker(name=W)` supplies
the workers and the count for an idle worker is simply 0:

```python
loads = program.solve().first().atoms(Load)
assert sorted(load.worker for load in loads) == ["ada", "grace"]
assert sum(load.tasks for load in loads) == 3   # three tasks, shared between them
```

## Banded Aggregates

A one-sided guard is a single comparison. You can also create a **band**, where an aggregate
needs to land between two bounds:

```python
# One-sided: no worker carries more than two tasks
load = Count(T, condition=Assigned(task=T, worker=W))
program.forbid(Worker(name=W), load > 2)

# A band: total assignments between 3 and 6 — one chained comparison
class Balanced(Predicate):
    pass

workload = Count((T, W), condition=Assigned(task=T, worker=W))
program.when(3 <= workload <= 6).derive(Balanced())
```

The chain `3 <= workload <= 6` builds one construct, but it is *not* a Python boolean value.
If you try and use it in an `if` statement, you'll get an exception.

The bounds can be written in whatever order (`3 <= workload < 7`, `9 > workload > 1` — either
direction), and can include expressions and variables, just like the cardinalities in choice rules.
If you create bounds that are impossible to satisfy, that gets called out in an exception:

```python
>>> 3 <= workload <= 2
Traceback (most recent call last):
  ...
ValueError: This aggregate band is empty: no value satisfies 3 <= #count{ T, W : assigned(T, W) } <= 2, so ...
```

Writing a band as *two* comparisons over the *same* aggregate — `when(workload >= 3, workload <= 6)`
— remains legal, renders as the two guards it says, and grounds to byte-identical solver input:
gringo recognizes the shared aggregate and grounds it once. What you *do* want to avoid is
**bind-then-compare**: binding the count to a variable — `N == Count(...)` — and then testing `N`
against the bounds. It reads like one aggregate, but it forces gringo to ground a fresh aggregate
for every feasible value of `N`, which is many times more expensive. Keep the variable out of it:
compare the aggregate, don't name its value. (Naming the value is the right tool when you want the
value itself, like the `Load` example above.)

## Cardinality tests are not choices

Clingo will let you use a choice construct as a condition in a rule or a constraint. Under the hood,
gringo converts the choice construct into a count aggregate. In ASPAlchemy, we don't allow you to
use a choice construct as a condition; a `Choice` object is only used to generate atoms.

```python
>>> program.forbid(Choice(OnCall(worker=W), condition=Worker(name=W)) >= 3)
Traceback (most recent call last):
  ...
ValueError: A guarded choice belongs in a rule head, where braces CHOOSE. In a body, clingo's braces mean a cardinality TEST — a different construct aspalchemy spells as a Count comparison: Count(X, condition=...) >= n.
```

Instead, just use the `Count` aggregate directly.

```python
program.require(Count(W, condition=OnCall(worker=W)) <= 2)
```

## The fine print: how a chained comparison reaches the library

Nothing above needed you to know how `3 <= workload <= 6` becomes one banded literal, and that's the
point. But the mechanism deserves a few sentences, because there are a few consequences.

Python evaluates a chained comparison as two comparisons joined by `and` — `(3 <= workload) and
(workload <= 6)`, with `workload` evaluated once. The library builds the first link, answers
Python's truth-test on it by noting *where in the compiled code* that test happens, builds the
second link, and — only when both links belong to the same chain in the same expression — merges
them into one banded literal (either a banded choice or aggregate).

That "where in the compiled code" is the caveat: the recognition reads bytecode, and it is tuned
against CPython — the interpreter facts it stands on are pinned by the test suite, on every Python
version the library supports. An interpreter that compiles a chain differently is not silently
misread: recognition fails *closed*, so a chain the machinery cannot vouch for gets the same error
as any other incomplete chain — possibly a refusal you didn't deserve on an exotic interpreter, but
never a program missing a bound.

While this can be a bit brittle, we think that being able to write `3 <= workload <= 6` is worth it.
If there's ever a problem with constructing the chain, the exception message will hopefully let you
know what's happened. At worst, you can split it into two terms as `3 <= workload` and `workload <=
6`.

There is one place to be aware of where things go a little wonky: in pytest, `assert` statements get
rewritten, which often disrupt our chain identification. As a rule of thumb, build the chain outside
of an assert statement.

```python
banded = 3 <= Count(X, condition=Person(name=X, age=ANY)) <= 6  # the chain, on its own line
# assert on `banded` from here
```