backend/entity_resolver.py

# backend/entity_resolver.py
import re
from rapidfuzz import fuzz

LEGAL_SUFFIXES = [
    "incorporated", "corporation", "limited", "inc", "ltd", "llc", "corp",
    "co", "plc", "company", "industries", "group", "holdings"
]

KNOWN_ALIASES = {
    "apple": "Apple", "aapl": "Apple", "apple inc": "Apple",
    "microsoft": "Microsoft", "msft": "Microsoft",
    "amazon": "Amazon", "amzn": "Amazon", "amazoncom": "Amazon",
    "google": "Google", "alphabet": "Google", "googl": "Google",
    "tesla": "Tesla", "tsla": "Tesla",
    "3m": "3M",
    "boeing": "Boeing", "the boeing company": "Boeing",
    "amd": "AMD", "advanced micro devices": "AMD",
    "hdfc bank": "HDFC Bank", "hdfc": "HDFC Bank",
    "tcs": "TCS", "tata consultancy services": "TCS",
    "infosys": "Infosys", "infy": "Infosys",
    "reliance": "Reliance Industries", "reliance industries": "Reliance Industries", "ril": "Reliance Industries",
    "icici bank": "ICICI Bank", "icici": "ICICI Bank",
    "wipro": "Wipro",
    "hul": "Hindustan Unilever", "hindustan unilever": "Hindustan Unilever",
    "itc": "ITC",
    "l&t": "Larsen & Toubro", "larsen & toubro": "Larsen & Toubro", "larsen and toubro": "Larsen & Toubro",
    "bajaj finance": "Bajaj Finance",
    "tata motors": "Tata Motors",
    "maruti": "Maruti Suzuki", "maruti suzuki": "Maruti Suzuki",
    "asian paints": "Asian Paints",
    "sun pharma": "Sun Pharma", "sun pharmaceutical": "Sun Pharma",
    "axis bank": "Axis Bank",
    "kotak": "Kotak Mahindra Bank", "kotak mahindra bank": "Kotak Mahindra Bank",
    "titan": "Titan Company",
    "nestle": "Nestle India", "nestle india": "Nestle India",
    "ultratech": "UltraTech Cement", "ultratech cement": "UltraTech Cement",
    "sbi": "State Bank of India", "state bank of india": "State Bank of India",
}


def normalize_org_name(text: str) -> str:
    t = text.lower().strip()
    t = re.sub(r"[^\w\s&]", "", t)
    words = t.split()
    while words and words[-1] in LEGAL_SUFFIXES:
        words.pop()
    return " ".join(words).strip()


def resolve_entity(text: str, label: str, existing_orgs: list = None) -> str:
    """Returns canonical entity ID for graph node deduplication."""
    if label != "ORG":
        clean = re.sub(r"\s+", " ", text.strip())
        return f"{label}_{clean[:50]}"

    normalized = normalize_org_name(text)

    if normalized in KNOWN_ALIASES:
        return f"ORG_{KNOWN_ALIASES[normalized]}"

    if existing_orgs:
        best_match, best_score = None, 0
        for existing in existing_orgs:
            score = fuzz.token_sort_ratio(normalized, existing.lower())
            if score > best_score:
                best_score, best_match = score, existing
        if best_score >= 85:
            return f"ORG_{best_match}"

    return f"ORG_{normalized.title()}"


def format_money(value: float, currency: str = "USD") -> str:
    """Format a raw monetary value for display, currency-aware.

    Takes an explicit `currency` ("INR" or "USD") rather than guessing
    from the company name — guessing is unreliable (the same product
    handles both US and Indian filings, and even Indian companies
    sometimes report figures in USD). Callers should pass the
    `currency` tag that metrics_extractor.py now attaches to each
    extracted metric, e.g.:

        formatted = format_money(raw["value"], raw.get("currency", "USD"))

    INR values are shown in crore (and lakh crore for very large
    figures) since that's the unit Indian filings and their readers
    actually think in — converting to USD billions would require an
    exchange-rate assumption that goes stale and adds a layer of
    unnecessary approximation on top of an already-extracted figure.
    """
    if value is None:
        return "N/A"

    if currency == "INR":
        crore = value / 1_00_00_000  # 1 crore = 10,000,000
        if abs(crore) >= 1_00_000:  # >= 1 lakh crore
            return f"₹{crore / 1_00_000:.2f} Lakh Cr"
        return f"₹{crore:,.1f} Cr"

    # USD (default)
    if abs(value) >= 1_000_000_000:
        return f"${value / 1_000_000_000:.1f}B"
    if abs(value) >= 1_000_000:
        return f"${value / 1_000_000:.1f}M"
    return f"${value:,.2f}"


if __name__ == "__main__":
    existing = ["Apple", "HDFC Bank", "3M"]
    tests = [
        ("Apple Inc.", "ORG"), ("Apple", "ORG"), ("AAPL", "ORG"),
        ("HDFC Bank Limited", "ORG"), ("3M Company", "ORG"),
        ("$394.3 billion", "MONEY"),
    ]
    for text, label in tests:
        print(f"{text!r} ({label}) -> {resolve_entity(text, label, existing)}")