"""Assign each term to one field of law, spread as evenly as the data allows.

Ordering is derived from SHA-256 of the term and of ``term|field`` rather than
position, so rebuilding an unchanged pool yields an identical assignment and
adding terms perturbs few existing ones. Grader verdicts are keyed on
``(term, field_of_law, question_id)``, so that stability keeps collected
verdicts valid.
"""

from __future__ import annotations

import hashlib
from collections import Counter
from collections.abc import Mapping
from dataclasses import dataclass


def _rank(*parts: str) -> int:
    return int(hashlib.sha256("|".join(parts).encode("utf-8")).hexdigest(), 16)


def derive_floor(n_terms: int, n_fields: int) -> int:
    """Terms per field when spread as evenly as possible."""
    if n_fields <= 0:
        return 0
    return n_terms // n_fields


@dataclass(frozen=True)
class Assignment:
    """A term -> field decision together with the numbers behind it.

    Diagnostics render this rather than re-deriving the decision from the input
    pool and the finished template. Those two views drift apart whenever the
    pool is filtered or capped, and then report a floor the assignment never
    used.
    """

    by_term: dict[str, str]
    floor: int
    floor_source: str
    fields: tuple[str, ...]
    assigned: Counter[str]
    candidates: Counter[str]
    n_terms: int

    def achievable_floor(self) -> int:
        """Upper bound on the floor this pool could have supported.

        Bounded by even division and by the scarcest field. The exact maximum
        is a bipartite feasibility problem; this estimate matched it in 106 of
        108 simulated datasets and never understates it.
        """
        if not self.fields:
            return 0
        return min(
            derive_floor(self.n_terms, len(self.fields)),
            min(self.candidates.get(field, 0) for field in self.fields),
        )


def assign_terms_to_fields(
    available: Mapping[str, set[str]],
    *,
    terms_per_field: int | None = None,
) -> Assignment:
    """Map each term to exactly one of its available fields of law."""
    terms = sorted(available)
    fields = sorted({field for value in available.values() for field in value})
    candidates_per_field = Counter(
        field for term_fields in available.values() for field in set(term_fields)
    )

    if terms_per_field is not None:
        floor = max(0, terms_per_field)
        floor_source = "explicit"
    else:
        floor = derive_floor(len(terms), len(fields))
        floor_source = "derived"

    def result(by_term: dict[str, str]) -> Assignment:
        return Assignment(
            by_term=by_term,
            floor=floor,
            floor_source=floor_source,
            fields=tuple(fields),
            assigned=Counter(by_term.values()),
            candidates=candidates_per_field,
            n_terms=len(terms),
        )

    if not terms or not fields:
        return result({})

    cap = -(-len(terms) // len(fields))

    candidates: dict[str, list[str]] = {field: [] for field in fields}
    for term in terms:
        for field in available[term]:
            candidates[field].append(term)

    assignment: dict[str, str] = {}
    used: Counter[str] = Counter()

    for field in sorted(fields, key=lambda f: (len(candidates[f]), f)):
        pool = sorted(
            (t for t in candidates[field] if t not in assignment),
            key=lambda t: _rank(t, field),
        )
        for term in pool[:floor]:
            assignment[term] = field
            used[field] += 1

    for term in sorted(
        (t for t in terms if t not in assignment), key=lambda t: _rank(t)
    ):
        options = sorted(available[term], key=lambda f: _rank(term, f))
        placed = next((f for f in options if used[f] < cap), None)
        if placed is None:
            placed = min(available[term], key=lambda f: (used[f], f))
        assignment[term] = placed
        used[placed] += 1

    return result(assignment)