politics-tracking/backend/politrack/core/resolve.py

"""Entity resolution: PersonRef/OrgRef -> database ids.

Order (first hit wins, ambiguity never guesses):
  persons:       identifiers -> hint strategies -> exact alias (titles stripped; prefer active)
  organizations: identifiers -> slug -> exact alias -> normalised legal name -> create (if allowed)

Hint strategies for persons (documented in docs/data-model.md):
  surname_in_body + on   members of a chamber on a date with that surname (roll-call lists)
  office + on            holder of an office title on a date (Council of Ministers minutes)
"""

from __future__ import annotations

import re
from datetime import date
from typing import Any

from psycopg import Connection
from rapidfuzz import fuzz

from politrack.core.records import OrgRef, PersonRef
from politrack.text import make_slug, normalize, normalize_org_name, strip_titles


OFFICE_FILLER = re.compile(r"\b(der republik osterreich|der bundesrepublik osterreich|im bundeskanzleramt)\b|"
                           r"\(osterreich\)")


def office_key(title: str) -> str:
    """'Bundesministerin für Justiz' / 'Bundeskanzler der Republik Österreich' -> 'bundesminister fur justiz' /
    'bundeskanzler' (gender-neutral, without filler, punctuation-free) for comparing office titles."""
    text = normalize(title)
    text = re.sub(r"\b(\w*(?:minister|kanzler|sekretar|prasident|rat))in\b", r"\1", text)
    text = OFFICE_FILLER.sub(" ", text)
    return re.sub(r"\s+", " ", re.sub(r"[^\w\s]", " ", text)).strip()


# Fallback for colloquial or English office labels (Wikidata: "Justizminister Österreichs", "Federal Minister for
# Arts, Culture, the Civil Service and Sport"): same kind of office + at least one common portfolio.
OFFICE_KINDS = [("vizekanzler", r"vizekanzler|vice.chancellor"), ("kanzler", r"kanzler|chancellor"),
                ("staatssekretar", r"staatssekretar|state secretary"), ("minister", r"minist")]
PORTFOLIOS = {
    "justiz": r"justiz|justice", "finanzen": r"finanz|financ", "inneres": r"\binnere|\binnenmin|interior",
    "aussen": r"au(ss|ß)en|auswartig|europaische und internationale|foreign|european and international",
    "verteidigung": r"verteidigung|defen[cs]e", "soziales": r"sozial|\barbeit|labou?r|social",
    "gesundheit": r"gesundheit|health", "bildung": r"bildung|education", "wissenschaft": r"wissenschaft|science",
    "frauen": r"frauen|women", "landwirtschaft": r"landwirtschaft|forstwirtschaft|agricultur",
    "umwelt": r"umwelt|klima|environment|climate", "kunst": r"kunst|kultur|\barts\b|culture|\bsport",
    "mobilitat": r"mobilitat|verkehr|infrastruktur|mobility|transport|infrastructure",
    "wirtschaft": r"\bwirtschaft|econom", "europa": r"\beuropa\b|integration|familie|famil",
    "digitalisierung": r"digital", "energie": r"energie|energy", "tourismus": r"touris",
}


def office_kind(key: str) -> str | None:
    return next((kind for kind, pattern in OFFICE_KINDS if re.search(pattern, key)), None)


def portfolios(key: str) -> set[str]:
    return {name for name, pattern in PORTFOLIOS.items() if re.search(pattern, key)}


def lead_portfolio(key: str) -> str | None:
    """The portfolio named first ("Wirtschaft, Energie und Tourismus" -> wirtschaft): a role must cover it."""
    found = [(m.start(), name) for name, pattern in PORTFOLIOS.items() if (m := re.search(pattern, key))]
    return min(found)[1] if found else None


class Resolver:
    def __init__(self, conn: Connection) -> None:
        self.conn = conn
        self._person_cache: dict[tuple, int | None] = {}
        self._org_cache: dict[tuple, int | None] = {}
        self._members_cache: dict[tuple, list[dict]] = {}
        self._office_cache: dict[tuple, int | None] = {}

    # -- persons ------------------------------------------------------------------------

    def person_by_ids(self, ids: dict[str, str]) -> int | None:
        for scheme, value in ids.items():
            if value is None:
                continue
            if scheme == "politrack":  # internal id, used by adapters that plan requests from the database
                row = self.conn.execute("SELECT id FROM person WHERE id = %s", (int(value),)).fetchone()
                if row:
                    return row["id"]
                continue
            row = self.conn.execute(
                "SELECT entity_id FROM identifier WHERE entity_type = 'person' AND scheme = %s AND value = %s",
                (scheme, str(value)),
            ).fetchone()
            if row:
                return row["entity_id"]
        return None

    def person(self, ref: PersonRef | None) -> int | None:
        if ref is None:
            return None
        cache_key = (tuple(sorted(ref.ids.items())), ref.name, tuple(sorted((k, str(v)) for k, v in ref.hints.items())))
        if cache_key in self._person_cache:
            return self._person_cache[cache_key]
        found = self.person_by_ids(ref.ids)
        if found is None and ref.hints:
            found = self._person_by_hints(ref)
        if found is None and ref.name and not ref.hints.get("no_name_match"):
            found = self.person_by_name(ref.name)
        self._person_cache[cache_key] = found
        return found

    def person_by_name(self, name: str) -> int | None:
        norm = normalize(strip_titles(name))
        rows = self.conn.execute(
            """SELECT DISTINCT a.entity_id AS id, ps.status
               FROM alias a LEFT JOIN person_status ps ON ps.person_id = a.entity_id
               WHERE a.entity_type = 'person' AND a.alias_norm = %s""",
            (norm,),
        ).fetchall()
        if len(rows) == 1:
            return rows[0]["id"]
        active = [r for r in rows if r["status"] == "active"]
        return active[0]["id"] if len(active) == 1 else None

    def _person_by_hints(self, ref: PersonRef) -> int | None:
        hints = ref.hints
        on: date | None = hints.get("on")
        if "surname_in_body" in hints and ref.name:
            members = self.members(hints["surname_in_body"], on)
            norm = normalize(strip_titles(ref.name))
            exact = [m for m in members if m["family_name"] and normalize(m["family_name"]) == norm]
            if len(exact) == 1:
                return exact[0]["id"]
            suffix = [m for m in members if normalize(strip_titles(m["full_name"])).endswith(norm)]
            return suffix[0]["id"] if len(suffix) == 1 else None
        if "office" in hints:
            return self.office_holder(hints["office"], on)
        return None

    def office_holder(self, office: str, on: date | None) -> int | None:
        """Holder of a (single-holder) office on a date. Titles are compared whole, gender-neutral and without
        filler ("Bundesministerin für Justiz" == "Bundesminister für Justiz"), so a generic "Bundesminister"
        never matches a specific ministry. Several matches (open-ended roles from Wikidata) -> the most recent
        appointment wins; no dated match -> None."""
        key = (office, on)
        if key not in self._office_cache:
            rows = self.conn.execute(
                """SELECT DISTINCT r.person_id, r.role, r.start_date FROM role r JOIN claim c ON c.id = r.claim_id
                   WHERE c.merged_into IS NULL AND r.level IN ('federal', 'state', 'eu')
                     AND (%(on)s::date IS NULL OR ((r.start_date IS NULL OR r.start_date <= %(on)s)
                          AND (r.end_date IS NULL OR r.end_date >= %(on)s)))""",
                {"on": on},
            ).fetchall()
            target = office_key(office)
            candidates = [(r["start_date"] or date.min, r["person_id"]) for r in rows
                          if fuzz.ratio(office_key(r["role"]), target) >= 92]
            kind, lead = office_kind(target), lead_portfolio(target)
            if not candidates and kind and (lead or kind in ("kanzler", "vizekanzler")):
                # Fuzzy labels only for appointments since the current government took office (the latest
                # chancellor appointment): stale open-ended roles of former ministers must not win.
                formed = max((r["start_date"] for r in rows if r["start_date"]
                              and office_kind(office_key(r["role"])) == "kanzler"), default=date.min)
                candidates = [(r["start_date"], r["person_id"]) for r in rows
                              if r["start_date"] and r["start_date"] >= formed
                              and office_kind(office_key(r["role"])) == kind
                              and (kind in ("kanzler", "vizekanzler") or lead in portfolios(office_key(r["role"])))]
            persons = {p for _, p in candidates}
            if len(persons) <= 1:
                self._office_cache[key] = persons.pop() if persons else None
            else:
                latest = max(start for start, _ in candidates)
                newest = {p for start, p in candidates if start == latest}
                self._office_cache[key] = newest.pop() if len(newest) == 1 and latest != date.min else None
        return self._office_cache[key]

    def members(self, body_slug: str, on: date | None) -> list[dict[str, Any]]:
        key = (body_slug, on)
        if key not in self._members_cache:
            self._members_cache[key] = self.conn.execute(
                """SELECT DISTINCT ON (p.id) p.id, p.full_name, p.family_name, r.party_id
                   FROM role r JOIN person p ON p.id = r.person_id JOIN organization o ON o.id = r.organization_id
                   JOIN claim c ON c.id = r.claim_id
                   WHERE o.slug = %(body)s AND r.role = 'mp' AND c.merged_into IS NULL
                     AND (%(on)s::date IS NULL OR (r.start_date <= %(on)s AND COALESCE(r.end_date, 'infinity') >= %(on)s))
                   ORDER BY p.id, r.start_date DESC""",
                {"body": body_slug, "on": on},
            ).fetchall()
        return self._members_cache[key]

    def faction_of(self, person_id: int, body_slug: str, on: date | None) -> int | None:
        return next((m["party_id"] for m in self.members(body_slug, on) if m["id"] == person_id), None)

    # -- organizations ------------------------------------------------------------------

    def org(self, ref: OrgRef | None, create: bool | None = None) -> int | None:
        if ref is None:
            return None
        allow_create = ref.create if create is None else create
        cache_key = (tuple(sorted(ref.ids.items())), ref.slug, ref.name, ref.short_name)
        if cache_key in self._org_cache and (self._org_cache[cache_key] is not None or not allow_create):
            return self._org_cache[cache_key]
        found = self._org_lookup(ref)
        if found is None and allow_create and (ref.name or ref.short_name):
            found = self._org_create(ref)
        elif found is not None:
            self._add_org_ids(found, ref)
        self._org_cache[cache_key] = found
        return found

    def _org_lookup(self, ref: OrgRef) -> int | None:
        for scheme, value in ref.ids.items():
            if scheme == "politrack_org":  # internal id (enrichers)
                row = self.conn.execute("SELECT id FROM organization WHERE id = %s", (int(value),)).fetchone()
                if row:
                    return row["id"]
                continue
            row = self.conn.execute(
                "SELECT entity_id FROM identifier WHERE entity_type = 'organization' AND scheme = %s AND value = %s",
                (scheme, str(value)),
            ).fetchone()
            if row:
                return row["entity_id"]
        if ref.slug:
            row = self.conn.execute("SELECT id FROM organization WHERE slug = %s", (ref.slug,)).fetchone()
            if row:
                return row["id"]
        for label in (ref.short_name, ref.name):
            if not label:
                continue
            for norm in {normalize(label), normalize_org_name(label)} - {""}:
                rows = self.conn.execute(
                    "SELECT DISTINCT entity_id FROM alias WHERE entity_type = 'organization' AND alias_norm = %s",
                    (norm,),
                ).fetchall()
                if len(rows) == 1:
                    return rows[0]["entity_id"]
        return None

    def _org_create(self, ref: OrgRef) -> int:
        base = ref.slug or make_slug(ref.short_name or ref.name or "org")[:60] or "org"
        slug, n = base, 1
        while self.conn.execute("SELECT 1 FROM organization WHERE slug = %s", (slug,)).fetchone():
            n += 1
            slug = f"{base}-{n}"
        org_id = self.conn.execute(
            "INSERT INTO organization (slug, name, short_name, kind) VALUES (%s, %s, %s, %s) RETURNING id",
            (slug, ref.name or ref.short_name, ref.short_name, ref.kind if ref.kind else "other"),
        ).fetchone()["id"]
        self.add_org_aliases(org_id, [ref.name, ref.short_name])
        self._add_org_ids(org_id, ref)
        return org_id

    def add_org_aliases(self, org_id: int, names: list[str | None]) -> None:
        for name in names:
            if not name:
                continue
            for norm in {normalize(name), normalize_org_name(name)} - {""}:
                self.conn.execute(
                    "INSERT INTO alias (entity_type, entity_id, alias, alias_norm) VALUES ('organization', %s, %s, %s) "
                    "ON CONFLICT DO NOTHING",
                    (org_id, name, norm),
                )

    def _add_org_ids(self, org_id: int, ref: OrgRef) -> None:
        for scheme, value in ref.ids.items():
            if value:
                self.conn.execute(
                    "INSERT INTO identifier (entity_type, entity_id, scheme, value) VALUES ('organization', %s, %s, %s) "
                    "ON CONFLICT DO NOTHING",
                    (org_id, scheme, str(value)),
                )