instant-bom/bom_assistant/ingestion/column_mapper.py

from __future__ import annotations

import json
import re
from pathlib import Path

_SYNONYMS_PATH = Path(__file__).parent / "column_synonyms.json"

# Columns where "Comment" maps to part_number instead of value when a
# richer value column is also present.
_VALUE_COLS = {"value", "description", "name", "part name", "component value", "part value"}
_COMMENT_COL = "comment"


def _load_synonyms() -> dict[str, list[str]]:
    with open(_SYNONYMS_PATH, encoding="utf-8") as f:
        return json.load(f)


def _save_synonyms(synonyms: dict[str, list[str]]) -> None:
    with open(_SYNONYMS_PATH, "w", encoding="utf-8") as f:
        json.dump(synonyms, f, indent=2, ensure_ascii=False)


def map_columns(raw_rows: list[dict[str, str]]) -> list[dict[str, str]]:
    if not raw_rows:
        return []

    raw_keys = list(raw_rows[0].keys())
    synonyms = _load_synonyms()
    mapping = _heuristic_map(raw_keys, synonyms)
    mapping = _disambiguate_comment(raw_keys, mapping, raw_rows)

    missing_required = [c for c in ("value",) if c not in mapping]
    if missing_required:
        ai_result = _ai_map(raw_keys, missing_required, synonyms)
        mapping.update(ai_result)

    return [_remap_row(row, mapping) for row in raw_rows]


def _normalize_key(s: str) -> str:
    return re.sub(r"\s+", " ", s.lower().strip())


def _heuristic_map(raw_keys: list[str], synonyms: dict[str, list[str]]) -> dict[str, str]:
    # Build normalized lookup: normalized_key → original_key
    norm_to_orig = {_normalize_key(k): k for k in raw_keys}

    mapping: dict[str, str] = {}
    # Iterate synonyms in priority order (JSON list order) for each canonical
    for canonical, syn_list in synonyms.items():
        if canonical in mapping:
            continue
        for syn in syn_list:
            norm_syn = _normalize_key(syn)
            if norm_syn in norm_to_orig:
                mapping[canonical] = norm_to_orig[norm_syn]
                break
    return mapping


_PASSIVE_VALUE_RE = re.compile(
    r"^\s*[\d\.]+\s*[pnuµmkKMG]?[FfHhRrΩ]?\s*$|"
    r"^\s*[\d\.]+\s*[kKMG]\s*$|"
    r"^\s*[\d\.]+\s*[pnuµ][fF]\s*$",
    re.IGNORECASE,
)


def _comment_looks_like_value(raw_rows: list[dict[str, str]], comment_key: str) -> bool:
    """Return True if Comment column values look like passive component values, not MPNs."""
    samples = [r.get(comment_key, "") for r in raw_rows[:10] if r.get(comment_key, "").strip()]
    if not samples:
        return False
    matches = sum(1 for s in samples if _PASSIVE_VALUE_RE.match(s))
    return matches / len(samples) > 0.4


def _disambiguate_comment(
    raw_keys: list[str],
    mapping: dict[str, str],
    raw_rows: list[dict[str, str]],
) -> dict[str, str]:
    """
    If BOM has both Comment and a richer value column (Value/Description),
    treat Comment as part_number — but only if Comment doesn't look like passive values.
    """
    lower_keys = {k.lower(): k for k in raw_keys}
    has_comment = _COMMENT_COL in lower_keys
    has_value_col = any(k in lower_keys for k in _VALUE_COLS)

    if not (has_comment and has_value_col):
        return mapping

    original_comment_key = lower_keys[_COMMENT_COL]

    # Don't remap if Comment values look like passive component values
    if _comment_looks_like_value(raw_rows, original_comment_key):
        return mapping

    if "part_number" not in mapping:
        mapping["part_number"] = original_comment_key
    if mapping.get("value") == original_comment_key:
        for col_name in _VALUE_COLS:
            if col_name in lower_keys and lower_keys[col_name] != original_comment_key:
                mapping["value"] = lower_keys[col_name]
                break
    return mapping


def _ai_map(
    raw_keys: list[str],
    missing: list[str],
    synonyms: dict[str, list[str]],
) -> dict[str, str]:
    import anthropic

    client = anthropic.Anthropic()
    prompt = (
        f"Map these BOM spreadsheet column names to canonical fields.\n"
        f"Columns present: {raw_keys}\n"
        f"Need mappings for: {missing}\n"
        f"Canonical fields:\n"
        f"  value — component value or description (e.g. '100nF', '10k', 'STM32F405')\n"
        f"  designators — reference designators (e.g. 'C1, C2', 'R5')\n"
        f"  footprint — PCB footprint or package\n"
        f"  quantity — how many of this component\n"
        f"  part_number — manufacturer part number / MPN\n"
        f"  manufacturer — manufacturer name\n\n"
        f"Return JSON only, no explanation: "
        f'{{\"canonical_field\": \"exact_original_column_name\", ...}}\n'
        f"Only include fields you are confident about."
    )

    msg = client.messages.create(
        model="claude-haiku-4-5-20251001",
        max_tokens=300,
        messages=[{"role": "user", "content": prompt}],
    )
    text = msg.content[0].text
    match = re.search(r"\{.*\}", text, re.DOTALL)
    if not match:
        return {}

    try:
        ai_mapping: dict[str, str] = json.loads(match.group())
    except json.JSONDecodeError:
        return {}

    # persist new synonyms so we avoid future API calls
    updated = False
    for canonical, raw_col in ai_mapping.items():
        if canonical not in synonyms:
            continue
        normalized = raw_col.lower().strip()
        existing = [s.lower() for s in synonyms[canonical]]
        if normalized not in existing:
            synonyms[canonical].append(raw_col.lower().strip())
            updated = True
    if updated:
        _save_synonyms(synonyms)

    return ai_mapping


def _remap_row(row: dict[str, str], mapping: dict[str, str]) -> dict[str, str]:
    mapped_originals = set(mapping.values())
    result: dict[str, str] = {}
    for canonical, original in mapping.items():
        result[canonical] = row.get(original, "")
    # pass through unmapped columns with _raw_ prefix
    for k, v in row.items():
        if k not in mapped_originals:
            result[f"_raw_{k}"] = v
    return result