"""
clusterer.py - Stage 2: deterministic co-location clustering (no LLM).

Collapses the N-by-N sheet-comparison explosion into a handful of clusters
where assertions from different disciplines (or a schedule vs a plan within one
discipline) refer to the SAME place or element. Only clusters worth reasoning
about reach Stage 3, which keeps each reasoning call small and cheap.

A cluster is emitted when, for a shared normalized location key, either:
  - assertions come from >= 2 distinct disciplines (cross-discipline), or
  - one discipline mixes a schedule/count fact with a plan/location/dimension
    fact for the same element (schedule-vs-plan).

Limitation (v1): a location that appears in only ONE discipline never forms a
cluster, so pure "missing element" gaps are not caught deterministically.
Stage 3 still catches missing counterparts whenever the location is co-located.
"""

import re
from typing import List, Dict, Tuple, Optional

# Caps to keep Stage 3 token usage bounded.
MAX_CLUSTERS = 120
MAX_ASSERTIONS_PER_CLUSTER = 24

_SCHEDULE_HINTS = ("count", "schedule", "_entry")
_PLAN_HINTS = ("location", "dim", "elev")


# ---------------------------------------------------------------------------
# Location-key normalization
# ---------------------------------------------------------------------------

def _norm_grid(v: str) -> Optional[str]:
    """'A/3', '3-A', 'A.3' -> canonical sorted 'token/token'."""
    toks = [t for t in re.split(r"[^A-Za-z0-9]+", v.upper()) if t]
    if not toks:
        return None
    return "/".join(sorted(toks))


def _norm_room(v: str) -> Optional[str]:
    v = re.sub(r"(?i)\b(room|rm|space)\b", "", v).strip()
    v = re.sub(r"\s+", " ", v).upper()
    return v or None


def _norm_tag(v: str) -> Optional[str]:
    toks = [t for t in re.split(r"[^A-Za-z0-9]+", v.upper()) if t]
    if not toks:
        return None
    return "-".join(toks)


def cluster_keys_for(assertion: Dict) -> List[Tuple[str, str]]:
    """All (key_type, normalized_value) keys this assertion participates in."""
    lk = assertion.get("location_key") or {}
    keys: List[Tuple[str, str]] = []
    if lk.get("grid"):
        g = _norm_grid(str(lk["grid"]))
        if g:
            keys.append(("grid", g))
    if lk.get("room"):
        r = _norm_room(str(lk["room"]))
        if r:
            keys.append(("room", r))
    if lk.get("tag"):
        t = _norm_tag(str(lk["tag"]))
        if t:
            keys.append(("tag", t))
    return keys


# ---------------------------------------------------------------------------
# Clustering
# ---------------------------------------------------------------------------

def _is_schedule(attr: str) -> bool:
    a = (attr or "").lower()
    return any(h in a for h in _SCHEDULE_HINTS)


def _is_plan(attr: str) -> bool:
    a = (attr or "").lower()
    return any(h in a for h in _PLAN_HINTS)


def _flatten(sheets: List[Dict]) -> List[Dict]:
    """Attach sheet-level context onto each assertion for clustering/evidence."""
    flat: List[Dict] = []
    for sheet in sheets:
        for a in sheet.get("assertions", []):
            flat.append({
                **a,
                "discipline": sheet.get("discipline", "Unknown"),
                "sheet_number": sheet.get("sheet_number"),
                "page_number": sheet.get("page_number"),
            })
    return flat


def _human_location(key_type: str, value: str, members: List[Dict]) -> str:
    label = {"room": "Room", "grid": "Grid", "tag": "Tag"}.get(key_type, key_type)
    levels = {m["location_key"].get("level") for m in members if m.get("location_key")}
    levels.discard(None)
    lvl = f" / {sorted(levels)[0]}" if len(levels) == 1 else ""
    return f"{label} {value}{lvl}"


def cluster_by_location(sheets: List[Dict]) -> List[Dict]:
    """
    Build co-location clusters from per-sheet extractions.

    Returns a list of clusters:
      {
        "key": "room:204",
        "location": "Room 204 / Level 2",
        "disciplines": ["Architectural", "Mechanical"],
        "page_numbers": [3, 14],
        "sheets": ["A2.1", "M2.1"],
        "assertions": [ {discipline, sheet_number, attribute, value, source_text, ...}, ... ]
      }
    """
    flat = _flatten(sheets)

    buckets: Dict[Tuple[str, str], List[Dict]] = {}
    for a in flat:
        for key in cluster_keys_for(a):
            buckets.setdefault(key, []).append(a)

    clusters: List[Dict] = []
    for (key_type, value), members in buckets.items():
        disciplines = {m["discipline"] for m in members}
        cross_discipline = len(disciplines) >= 2
        schedule_vs_plan = (
            len(disciplines) == 1
            and any(_is_schedule(m["attribute"]) for m in members)
            and any(_is_plan(m["attribute"]) for m in members)
        )
        if not (cross_discipline or schedule_vs_plan):
            continue

        # Rank members so the cap keeps the most informative ones (high
        # confidence first), but always keep at least one per discipline.
        members = sorted(
            members,
            key=lambda m: {"high": 0, "medium": 1, "low": 2}.get(m.get("confidence"), 1),
        )[:MAX_ASSERTIONS_PER_CLUSTER]

        pages = sorted({m["page_number"] for m in members if m.get("page_number")})
        sheet_nums = sorted({m["sheet_number"] for m in members if m.get("sheet_number")})

        clusters.append({
            "key": f"{key_type}:{value}",
            "location": _human_location(key_type, value, members),
            "disciplines": sorted(disciplines),
            "page_numbers": pages,
            "sheets": sheet_nums,
            "assertions": members,
            "kind": "cross_discipline" if cross_discipline else "schedule_vs_plan",
        })

    # Most disciplines / most assertions first; deterministic tie-break by key.
    clusters.sort(key=lambda c: (-len(c["disciplines"]), -len(c["assertions"]), c["key"]))
    if len(clusters) > MAX_CLUSTERS:
        print(f"[Cluster] Capping {len(clusters)} clusters to {MAX_CLUSTERS}")
        clusters = clusters[:MAX_CLUSTERS]

    print(f"[Cluster] {len(clusters)} candidate clusters "
          f"({sum(1 for c in clusters if c['kind']=='cross_discipline')} cross-discipline, "
          f"{sum(1 for c in clusters if c['kind']=='schedule_vs_plan')} schedule-vs-plan)")
    return clusters