stork/scripts/import_names_db.py

#!/usr/bin/env python3
"""Build stork/data/world-names.sqlite3 from open dumps + curated JSON.

Primary source: Philippe Remy names-dataset (~700k+ first names with country/gender).
Optional overlays: existing world-names.json, NameDatabases all.txt, Hadley baby-names.csv.

Usage:
  .venv/bin/pip install names-dataset   # build-time only
  .venv/bin/python scripts/import_names_db.py
  .venv/bin/python scripts/import_names_db.py --no-names-dataset \\
      --bulk /tmp/namedb.txt --hadley /tmp/baby-names.csv
"""

from __future__ import annotations

import argparse
import csv
import json
import re
import sqlite3
import sys
from collections import defaultdict
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "stork" / "data" / "world-names.sqlite3"
OUT_REPO = ROOT / "data" / "world-names.sqlite3"
CURATED_JSON = ROOT / "stork" / "data" / "world-names.json"

# ISO country → Stork generator region chips
COUNTRY_TO_REGIONS: dict[str, tuple[str, ...]] = {
    "US": ("en",),
    "GB": ("en",),
    "CA": ("en",),
    "AU": ("en",),
    "NZ": ("en",),
    "IE": ("en",),
    "RU": ("ru",),
    "BY": ("ru",),
    "UA": ("ru",),
    "KZ": ("ru",),
    "IL": ("he",),
    "SA": ("ar",),
    "EG": ("ar",),
    "AE": ("ar",),
    "MA": ("ar",),
    "DZ": ("ar",),
    "TN": ("ar",),
    "IQ": ("ar",),
    "SY": ("ar",),
    "JO": ("ar",),
    "LB": ("ar",),
    "PS": ("ar",),
    "QA": ("ar",),
    "KW": ("ar",),
    "BH": ("ar",),
    "OM": ("ar",),
    "YE": ("ar",),
    "LY": ("ar",),
    "SD": ("ar",),
    "FR": ("fr",),
    "BE": ("fr",),
    "LU": ("fr",),
    "SN": ("fr",),
    "CI": ("fr",),
    "DE": ("de",),
    "AT": ("de",),
    "CH": ("de", "fr"),
    "ES": ("es",),
    "MX": ("es",),
    "AR": ("es",),
    "CO": ("es",),
    "CL": ("es",),
    "PE": ("es",),
    "VE": ("es",),
    "UY": ("es",),
    "EC": ("es",),
    "BO": ("es",),
    "PY": ("es",),
    "CR": ("es",),
    "PA": ("es",),
    "GT": ("es",),
    "HN": ("es",),
    "SV": ("es",),
    "NI": ("es",),
    "DO": ("es",),
    "CU": ("es",),
    "IT": ("it",),
    "PT": ("pt",),
    "BR": ("pt",),
    "PL": ("pl",),
    "NL": ("nl",),
    "SE": ("sv",),
    "NO": ("no",),
    "DK": ("da",),
    "FI": ("fi",),
    "TR": ("tr",),
    "JP": ("ja",),
    "KR": ("ko",),
    "CN": ("zh",),
    "TW": ("zh",),
    "HK": ("zh",),
    "IN": ("hi",),
    "PK": ("hi",),
    "BD": ("hi",),
}

LATINISH = re.compile(r"^[A-Za-z][A-Za-z'\- ]{0,39}$")
MAX_NAME_LEN = 40


def _display_name(name: str) -> str:
    cleaned = " ".join(name.strip().split())
    if cleaned.isascii() and all(c.isalpha() or c in "'- " for c in cleaned):
        return " ".join(part[:1].upper() + part[1:].lower() for part in cleaned.split())
    return cleaned


def _letters_len(name: str) -> int:
    return sum(1 for c in name if c.isalpha())


def _ok_name(name: str) -> bool:
    name = " ".join(name.strip().split())
    if not name or len(name) > MAX_NAME_LEN:
        return False
    if not any(c.isalpha() for c in name):
        return False
    # reject obvious junk tokens
    if any(ch.isdigit() for ch in name):
        return False
    return True


class Accumulator:
    def __init__(self) -> None:
        # key = casefold → {name, gender, regions:set}
        self.rows: dict[str, dict[str, object]] = {}

    def add(self, name: str, regions: list[str] | set[str], gender: str = "u") -> None:
        name = " ".join(name.strip().strip("\ufeff").split())
        if not _ok_name(name):
            return
        display = _display_name(name)
        key = display.casefold()
        g = (gender or "u").strip().lower()[:1] or "u"
        if g not in {"f", "m", "u"}:
            g = "u"
        regs = {r.strip().lower() for r in regions if r and str(r).strip()}
        if not regs:
            regs = {"intl"}
        if key not in self.rows:
            self.rows[key] = {
                "name": display,
                "gender": g,
                "regions": set(regs),
            }
            return
        entry = self.rows[key]
        assert isinstance(entry["regions"], set)
        entry["regions"].update(regs)
        cur = str(entry.get("gender") or "u")
        if g in {"f", "m"}:
            if cur == "u":
                entry["gender"] = g
            elif cur not in {g, "u"}:
                entry["gender"] = "u"


def load_curated_json(acc: Accumulator, path: Path) -> int:
    if not path.is_file():
        return 0
    raw = json.loads(path.read_text(encoding="utf-8"))
    items = raw.get("names") if isinstance(raw, dict) else raw
    n = 0
    for row in items or []:
        if not isinstance(row, dict):
            continue
        name = str(row.get("name") or "").strip()
        regions = [str(r) for r in (row.get("regions") or [])]
        gender = str(row.get("gender") or "u")
        before = len(acc.rows)
        acc.add(name, regions, gender)
        if len(acc.rows) > before or name.casefold() in acc.rows:
            n += 1
    return n


def load_bulk_txt(acc: Accumulator, path: Path, regions: list[str]) -> int:
    if not path.is_file():
        return 0
    n = 0
    for line in path.read_text(encoding="utf-8", errors="ignore").splitlines():
        name = line.strip()
        if not name or name.startswith("#"):
            continue
        if not LATINISH.match(name.replace(" ", "")) and " " not in name:
            # allow simple latin; skip weird bulk lines
            if not name.isascii():
                continue
        before = name.casefold() in acc.rows
        acc.add(name, regions, "u")
        if name.casefold() in acc.rows and not before:
            n += 1
        elif before:
            n += 1
    return n


def load_hadley_csv(acc: Accumulator, path: Path) -> int:
    """Hadley Wickham baby-names.csv: year,name,percent,sex"""
    if not path.is_file():
        return 0
    n = 0
    with path.open(encoding="utf-8", errors="ignore", newline="") as fh:
        reader = csv.DictReader(fh)
        for row in reader:
            name = (row.get("name") or "").strip()
            sex = (row.get("sex") or row.get("gender") or "").strip().upper()[:1]
            gender = {"F": "f", "M": "m"}.get(sex, "u")
            before = name.casefold() in acc.rows
            acc.add(name, ["en", "intl"], gender)
            if name.casefold() in acc.rows:
                n += 0 if before else 1
    return n


def load_names_dataset(acc: Accumulator, *, min_country_share: float, max_rank: int) -> int:
    try:
        from names_dataset import NameDataset
    except ImportError as exc:
        raise SystemExit(
            "names-dataset is required for the full corpus. "
            "Install with: .venv/bin/pip install names-dataset\n"
            f"({exc})"
        ) from exc

    nd = NameDataset()
    added = 0
    for name, meta in nd.first_names.items():
        # Keep single-token given names; skip multi-word junk from dumps
        if " " in str(name).strip():
            continue
        if not isinstance(meta, dict):
            continue
        countries = meta.get("country") or {}
        ranks = meta.get("rank") or {}
        # Keep names with meaningful presence somewhere
        max_share = max(countries.values()) if countries else 0.0
        best_rank = min(ranks.values()) if ranks else 10**9
        if max_share < min_country_share and best_rank > max_rank:
            continue

        gender_map = meta.get("gender") or {}
        gender = "u"
        if gender_map:
            top = max(gender_map, key=gender_map.get)
            if float(gender_map[top]) >= 0.65:
                gender = str(top).lower()[:1]

        regions: set[str] = {"intl"}
        preferred_countries = {"IL", "RU", "US", "GB", "CA", "AU", "FR", "DE", "AT"}
        ranked_countries = sorted(countries.items(), key=lambda kv: kv[1], reverse=True)
        for code, share in ranked_countries[:6]:
            code_u = str(code).upper()
            threshold = 0.03 if code_u in preferred_countries else 0.08
            if float(share) < threshold:
                continue
            for region in COUNTRY_TO_REGIONS.get(code_u, ()):
                regions.add(region)
        # Best rank in preferred countries still earns a tag
        for code, rank in ranks.items():
            code_u = str(code).upper()
            if code_u in preferred_countries and int(rank) <= 3000:
                for region in COUNTRY_TO_REGIONS.get(code_u, ()):
                    regions.add(region)

        before = name.casefold() in acc.rows
        acc.add(str(name), regions, gender)
        if name.casefold() in acc.rows and not before:
            added += 1
    return added


def write_sqlite(acc: Accumulator, path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    if path.exists():
        path.unlink()
    con = sqlite3.connect(path)
    try:
        con.executescript(
            """
            PRAGMA journal_mode=OFF;
            PRAGMA synchronous=OFF;
            CREATE TABLE meta (
              key TEXT PRIMARY KEY,
              value TEXT NOT NULL
            );
            CREATE TABLE names (
              id INTEGER PRIMARY KEY,
              name TEXT NOT NULL,
              name_cf TEXT NOT NULL UNIQUE,
              gender TEXT NOT NULL CHECK(gender IN ('f','m','u')),
              len_alpha INTEGER NOT NULL
            );
            CREATE TABLE name_regions (
              name_id INTEGER NOT NULL REFERENCES names(id) ON DELETE CASCADE,
              region TEXT NOT NULL,
              PRIMARY KEY (name_id, region)
            );
            CREATE INDEX idx_names_len ON names(len_alpha);
            CREATE INDEX idx_names_gender ON names(gender);
            CREATE INDEX idx_regions_region ON name_regions(region);
            """
        )
        name_rows: list[tuple[int, str, str, str, int]] = []
        region_rows: list[tuple[int, str]] = []
        for i, entry in enumerate(acc.rows.values(), start=1):
            name = str(entry["name"])
            gender = str(entry["gender"])
            regions = sorted(entry["regions"])  # type: ignore[arg-type]
            name_rows.append((i, name, name.casefold(), gender, _letters_len(name)))
            for region in regions:
                region_rows.append((i, region))
        con.executemany(
            "INSERT INTO names(id, name, name_cf, gender, len_alpha) VALUES (?,?,?,?,?)",
            name_rows,
        )
        con.executemany(
            "INSERT INTO name_regions(name_id, region) VALUES (?,?)",
            region_rows,
        )
        counts: dict[str, int] = defaultdict(int)
        for _, region in region_rows:
            counts[region] += 1
        con.execute(
            "INSERT INTO meta(key, value) VALUES (?, ?)",
            ("total", str(len(name_rows))),
        )
        con.execute(
            "INSERT INTO meta(key, value) VALUES (?, ?)",
            ("regions_json", json.dumps(dict(sorted(counts.items())), ensure_ascii=False)),
        )
        con.execute(
            "INSERT INTO meta(key, value) VALUES (?, ?)",
            ("source", "names-dataset + curated overlays"),
        )
        con.commit()
        con.execute("VACUUM")
        con.commit()
    finally:
        con.close()


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--out", type=Path, default=OUT)
    parser.add_argument("--curated", type=Path, default=CURATED_JSON)
    parser.add_argument("--bulk", type=Path, help="NameDatabases-style one name per line")
    parser.add_argument("--hadley", type=Path, help="Hadley baby-names.csv")
    parser.add_argument("--no-names-dataset", action="store_true")
    parser.add_argument(
        "--min-country-share",
        type=float,
        default=0.0,
        help="Drop names whose max country share is below this (default keep all tagged)",
    )
    parser.add_argument(
        "--max-rank",
        type=int,
        default=50_000,
        help="Keep names with best country rank <= this even if share is low",
    )
    args = parser.parse_args()

    acc = Accumulator()
    curated_n = load_curated_json(acc, args.curated)
    print(f"curated json: merged {curated_n} rows → {len(acc.rows)} unique", file=sys.stderr)

    if not args.no_names_dataset:
        nd_n = load_names_dataset(
            acc,
            min_country_share=args.min_country_share,
            max_rank=args.max_rank,
        )
        print(f"names-dataset: +{nd_n} new → {len(acc.rows)} unique", file=sys.stderr)

    if args.bulk:
        b = load_bulk_txt(acc, args.bulk, ["en", "intl"])
        print(f"bulk txt: touched/added ~{b} → {len(acc.rows)} unique", file=sys.stderr)

    if args.hadley:
        h = load_hadley_csv(acc, args.hadley)
        print(f"hadley csv: +~{h} new → {len(acc.rows)} unique", file=sys.stderr)

    if not acc.rows:
        print("No names accumulated — aborting", file=sys.stderr)
        return 1

    write_sqlite(acc, args.out)
    if args.out.resolve() != OUT_REPO.resolve():
        OUT_REPO.parent.mkdir(parents=True, exist_ok=True)
        OUT_REPO.write_bytes(args.out.read_bytes())
    size_mb = args.out.stat().st_size / (1024 * 1024)
    print(f"wrote {args.out} ({len(acc.rows)} names, {size_mb:.1f} MiB)", file=sys.stderr)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())