Conflict_Checker/backend/pipeline/normalizer.py

"""
normalizer.py - Stage 3: construction object normalization + project intelligence.

Stage 3a: normalize_assertions — adds normalized values (e.g. 9'-0" -> 108 in)
and normalized tags/room identifiers to extracted objects. Additive; originals
are always preserved. Processed in batches across a thread pool.

Stage 3b: build_project_intelligence — cross-sheet merging pass that assigns
Global Object IDs (GOIDs) and builds relationships between objects that
represent the same physical element across multiple drawings/schedules.
Batched by object_type so each call stays within the token budget.
"""

import json
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor
from typing import Dict, List

from backend import config
from backend.pipeline._stage import call_stage
from backend.prompts import (
    NORMALIZATION_SYSTEM_PROMPT,
    NORMALIZATION_USER_INSTRUCTION,
    PROJECT_INTELLIGENCE_SYSTEM_PROMPT,
    PROJECT_INTELLIGENCE_USER_INSTRUCTION,
)


def _batches(items: List[Dict], size: int) -> List[List[Dict]]:
    return [items[i:i + size] for i in range(0, len(items), size)]


def _normalize_batch(batch: List[Dict]) -> Dict[str, Dict]:
    """Return {assertion_id: normalization fields} for one batch."""
    payload = [
        {
            "assertion_id": a["id"],
            "attribute": a.get("attribute"),
            "value": a.get("value"),
            "object_type": a.get("object_type"),
            "object_tag": a.get("object_tag"),
            "source_text": a.get("source_text"),
        }
        for a in batch
    ]
    parsed = call_stage(
        NORMALIZATION_SYSTEM_PROMPT,
        NORMALIZATION_USER_INSTRUCTION,
        subs={"assertions": json.dumps(payload, ensure_ascii=True)},
        max_tokens=config.NORMALIZE_MAX_TOKENS,
    )
    if isinstance(parsed, list):
        rows = parsed
    elif isinstance(parsed, dict):
        rows = parsed.get("normalized_assertions") or []
    else:
        rows = []
    out: Dict[str, Dict] = {}
    for n in rows:
        if isinstance(n, dict) and n.get("assertion_id"):
            out[n["assertion_id"]] = {
                "normalized_value": n.get("normalized_value"),
                "normalized_unit":  n.get("normalized_unit"),
                "normalized_tag":   n.get("normalized_tag"),
                "normalized_room":  n.get("normalized_room"),
            }
    return out


def normalize_assertions(sheets: List[Dict]) -> List[Dict]:
    """Annotate each assertion in-place with normalized fields."""
    all_assertions = [a for s in sheets for a in s.get("assertions", [])]
    if not all_assertions:
        return sheets

    batches = _batches(all_assertions, config.NORMALIZE_BATCH_SIZE)
    merged: Dict[str, Dict] = {}
    with ThreadPoolExecutor(max_workers=config.NORMALIZE_CONCURRENCY) as pool:
        for result in pool.map(_normalize_batch, batches):
            merged.update(result)

    applied = 0
    for a in all_assertions:
        norm = merged.get(a["id"])
        if norm:
            if norm.get("normalized_value") is not None:
                a["normalized_value"] = norm["normalized_value"]
                a["normalized_unit"]  = norm.get("normalized_unit")
                applied += 1
            if norm.get("normalized_tag"):
                a["normalized_tag"] = norm["normalized_tag"]
            if norm.get("normalized_room"):
                a["normalized_room"] = norm["normalized_room"]
                # Also push into location_key so clusterer picks it up
                if isinstance(a.get("location_key"), dict) and not a["location_key"].get("room"):
                    a["location_key"]["room"] = norm["normalized_room"]

    print(f"[Normalize] normalized {applied}/{len(all_assertions)} objects")
    return sheets


# ---------------------------------------------------------------------------
# Stage 3b - project intelligence: GOIDs + relationships
# ---------------------------------------------------------------------------

# Max objects of one type to send in a single intelligence call.
_INTEL_BATCH_SIZE = 60


def _intel_batch(object_type: str, objects: List[Dict]) -> Dict:
    """Run the project intelligence pass for one object_type batch."""
    payload = []
    for obj in objects:
        payload.append({
            "object_id":   obj.get("id"),
            "object_type": obj.get("object_type") or object_type,
            "tag":         obj.get("object_tag"),
            "name":        obj.get("object_name"),
            "attributes":  obj.get("object_attributes") or {},
            "location_key": obj.get("location_key") or {},
            "source_text": obj.get("source_text") or "",
            "source_sheet": (obj.get("location_key") or {}).get("sheet_number"),
        })
    parsed = call_stage(
        PROJECT_INTELLIGENCE_SYSTEM_PROMPT,
        PROJECT_INTELLIGENCE_USER_INSTRUCTION,
        subs={
            "object_type": object_type,
            "objects": json.dumps(payload, ensure_ascii=True),
        },
        max_tokens=config.NORMALIZE_MAX_TOKENS,
    )
    if not isinstance(parsed, dict):
        return {}
    return {
        "project_objects":         parsed.get("project_objects") or [],
        "unresolved_relationships": parsed.get("unresolved_relationships") or [],
    }


def build_project_intelligence(sheets: List[Dict]) -> Dict:
    """
    Stage 3b: cross-sheet GOID assignment and relationship building.

    Groups all extracted objects by object_type, then calls the Project
    Intelligence LLM once per type (batched) to assign GOIDs and link
    objects that represent the same physical element across sheets.

    Returns {"project_objects": [...], "unresolved_relationships": [...]}.
    On any failure, returns an empty dict (graceful degradation).
    """
    all_assertions = [a for s in sheets for a in s.get("assertions", [])]
    if not all_assertions:
        return {}

    # Group by object_type; unknown/graphical go under "general"
    by_type: Dict[str, List[Dict]] = defaultdict(list)
    for a in all_assertions:
        otype = a.get("object_type") or "general"
        by_type[otype].append(a)

    all_project_objects: List[Dict] = []
    all_unresolved: List[Dict] = []

    # Process each type in batches; high-count types (e.g. dimensions) are
    # less useful for GOID merging so we skip types with only 1 object.
    for otype, objs in sorted(by_type.items()):
        if len(objs) < 2:
            continue
        batches = _batches(objs, _INTEL_BATCH_SIZE)
        for batch in batches:
            try:
                result = _intel_batch(otype, batch)
                all_project_objects.extend(result.get("project_objects") or [])
                all_unresolved.extend(result.get("unresolved_relationships") or [])
            except Exception as e:
                print(f"[ProjectIntel] {otype} batch failed: {e}")

    print(f"[ProjectIntel] {len(all_project_objects)} project objects, "
          f"{len(all_unresolved)} unresolved relationships")
    return {
        "project_objects":          all_project_objects,
        "unresolved_relationships": all_unresolved,
    }