LedgerLine/scripts/seed.py

#!/usr/bin/env python3
"""Deterministic seed data for Ledgerline -- phase-0-base-project-spec.md §8.

Idempotent: truncates then inserts. Deterministic given SEED (default 42).
Uses COPY for the transactions table so 250k rows load in well under 60s.
"""
import json
import os
import random
import time
import uuid
from datetime import datetime, timedelta, timezone

import psycopg

SEED = int(os.environ.get("SEED", "42"))
DSN = os.environ.get("POSTGRES_DSN", "postgresql://ledgerline:ledgerline@localhost:5432/ledgerline")

N_ACCOUNTS = 800
N_HOT = 5
N_CATEGORIES = 45
N_TRANSACTIONS = 250_000
N_SANCTIONS = 12_000
WINDOW_DAYS = 120
N_POISON_ACCOUNTS = 25
N_SANCTIONED_MATCHES = 40

MALFORMED_COUNTS = {
    "null_currency": 1200,
    "negative_deposit": 1000,
    "non_ascii_merchant": 1000,
    "future_dated": 1000,
    "orphan_category": 800,
}
N_MALFORMED = sum(MALFORMED_COUNTS.values())
N_NORMAL = N_TRANSACTIONS - N_MALFORMED

FIRST_NAMES = [
    "Amara", "Bilal", "Chiara", "Dimitri", "Elin", "Farrukh", "Greta", "Haruto",
    "Ines", "Jonas", "Kalani", "Lucia", "Mateus", "Naledi", "Omar", "Priya",
    "Quentin", "Rosalind", "Sami", "Tuuli", "Umberto", "Vera", "Wojciech", "Ximena",
    "Yara", "Zoltan", "Anika", "Bram", "Carys", "Dovid", "Esperanza", "Felix",
    "Guadalupe", "Henrik", "Ilse", "Jarrah", "Katarzyna", "Leandro", "Mireille",
    "Nikolai",
]
LAST_NAMES = [
    "Adeyemi", "Bergstrom", "Castellano", "Dubois", "Eriksson", "Fitzgerald",
    "Gallagher", "Haddad", "Ivanov", "Jarvinen", "Kowalski", "Lindqvist",
    "Moreau", "Nakamura", "Okafor", "Petrov", "Quiroga", "Rasmussen", "Sokolova",
    "Tremblay", "Ueda", "Vasiliev", "Weiss", "Xu", "Yilmaz", "Zimmermann",
    "Andersson", "Bianchi", "Costa", "Duarte",
]
MERCHANTS = [
    "Northwind Grocers", "Blue Anchor Cafe", "Solstice Electronics", "Ferro Hardware",
    "Willow Creek Pharmacy", "Ember & Oak", "Pinnacle Fitness", "Cobblestone Books",
    "Lantern Street Diner", "Cedarline Motors", "Harborview Market", "Quill & Ink Supply",
    "Meridian Travel", "Golden Hour Bakery", "Ironwood Furniture", "Tidal Wave Surf Shop",
    "Maple & Vine Wines", "Copper Kettle Tea Co", "Driftwood Outfitters", "Verdant Nursery",
]

CURRENCY_WEIGHTS = [("EUR", 0.70), ("USD", 0.20), ("GBP", 0.10)]


def weighted_choice(rng, pairs):
    values, weights = zip(*pairs)
    return rng.choices(values, weights=weights, k=1)[0]


def gen_holder_names(rng, n):
    combos = [f"{f} {l}" for f in FIRST_NAMES for l in LAST_NAMES]
    rng.shuffle(combos)
    if n > len(combos):
        raise ValueError("not enough name combinations for the requested account count")
    return combos[:n]


def gen_accounts(rng):
    names = gen_holder_names(rng, N_ACCOUNTS)
    hot_indexes = set(rng.sample(range(N_ACCOUNTS), N_HOT))
    now = datetime.now(timezone.utc)
    rows = []
    for i, name in enumerate(names):
        account_id = str(uuid.uuid4())
        currency = weighted_choice(rng, CURRENCY_WEIGHTS)
        balance = int(rng.lognormvariate(_median_to_mu(250_000), 0.7))
        created_at = now - timedelta(days=rng.uniform(150, 400))
        rows.append(
            {
                "id": account_id,
                "holder_name": name,
                "currency": currency,
                "balance_minor": balance,
                "is_hot": i in hot_indexes,
                "created_at": created_at,
            }
        )
    return rows


def _median_to_mu(median):
    import math
    return math.log(median)


def gen_categories(rng):
    rows = []
    for i in range(1, N_CATEGORIES + 1):
        rows.append(
            {
                "id": i,
                "code": f"cat_{i:02d}",
                "label": f"Category {i:02d}",
                "risk_weight": round(0.01 + (i - 1) / (N_CATEGORIES - 1) * 0.94, 2),
            }
        )
    return rows


def account_weights(rng, accounts):
    """5 hot accounts own ~30% of rows; the remaining 795 share the rest on a
    Zipf-ish distribution."""
    hot = [a for a in accounts if a["is_hot"]]
    cold = [a for a in accounts if not a["is_hot"]]
    hot_weight_each = 0.30 / len(hot)
    ranks = list(range(1, len(cold) + 1))
    rng.shuffle(ranks)
    zipf_raw = [1.0 / (r ** 1.1) for r in ranks]
    zipf_total = sum(zipf_raw)
    cold_weights = [0.70 * w / zipf_total for w in zipf_raw]

    ids = [a["id"] for a in hot] + [a["id"] for a in cold]
    weights = [hot_weight_each] * len(hot) + cold_weights
    return ids, weights


def sample_created_at(rng, now, window_days=WINDOW_DAYS):
    day_offset = rng.uniform(0, window_days)
    day = now - timedelta(days=day_offset)
    hour = min(23, max(0, int(rng.gauss(13, 3.5))))
    minute, second = rng.randint(0, 59), rng.randint(0, 59)
    dt = day.replace(hour=hour, minute=minute, second=second, microsecond=0)
    if dt.weekday() >= 5 and rng.random() < 0.5:
        dt -= timedelta(days=2)
    return dt


def gen_normal_transactions(rng, accounts, now):
    ids, weights = account_weights(rng, accounts)
    statuses = [("settled", 0.92), ("rejected", 0.05), ("pending", 0.03)]
    kinds = [("deposit", 0.2), ("withdrawal", 0.2), ("purchase", 0.6)]

    rows = []
    for _ in range(N_NORMAL):
        account_id = rng.choices(ids, weights=weights, k=1)[0]
        kind = weighted_choice(rng, kinds)
        status = weighted_choice(rng, statuses)
        amount = max(1, int(rng.lognormvariate(_median_to_mu(5_000), 0.9)))
        risk_score = round(rng.uniform(0, 1), 3) if status != "pending" else None
        rows.append(
            (
                str(uuid.uuid4()),
                account_id,
                kind,
                amount,
                weighted_choice(rng, CURRENCY_WEIGHTS),
                rng.choice(MERCHANTS),
                rng.randint(1, N_CATEGORIES),
                status,
                risk_score,
                sample_created_at(rng, now),
            )
        )
    return rows


def gen_malformed_transactions(rng, accounts, poison_account_ids, now):
    all_ids = [a["id"] for a in accounts if a["id"] not in poison_account_ids]
    rows = []
    poison_tx_ids = {}

    # null_currency: every poison account gets at least one; the rest spread randomly.
    remaining = MALFORMED_COUNTS["null_currency"] - len(poison_account_ids)
    null_currency_accounts = list(poison_account_ids) + [
        rng.choice(all_ids) for _ in range(remaining)
    ]
    for account_id in null_currency_accounts:
        tx_id = str(uuid.uuid4())
        if account_id in poison_account_ids:
            poison_tx_ids[account_id] = tx_id
        rows.append(
            (
                tx_id, account_id, "purchase",
                max(1, int(rng.lognormvariate(_median_to_mu(5_000), 0.9))),
                None, rng.choice(MERCHANTS), rng.randint(1, N_CATEGORIES),
                "settled", round(rng.uniform(0, 1), 3), sample_created_at(rng, now),
            )
        )

    for _ in range(MALFORMED_COUNTS["negative_deposit"]):
        rows.append(
            (
                str(uuid.uuid4()), rng.choice(all_ids), "deposit",
                -max(1, int(rng.lognormvariate(_median_to_mu(5_000), 0.9))),
                weighted_choice(rng, CURRENCY_WEIGHTS), rng.choice(MERCHANTS),
                rng.randint(1, N_CATEGORIES), "settled", round(rng.uniform(0, 1), 3),
                sample_created_at(rng, now),
            )
        )

    for _ in range(MALFORMED_COUNTS["non_ascii_merchant"]):
        merchant = f"{rng.choice(MERCHANTS)} \U0001F600\\x00"
        rows.append(
            (
                str(uuid.uuid4()), rng.choice(all_ids), "purchase",
                max(1, int(rng.lognormvariate(_median_to_mu(5_000), 0.9))),
                weighted_choice(rng, CURRENCY_WEIGHTS), merchant,
                rng.randint(1, N_CATEGORIES), "settled", round(rng.uniform(0, 1), 3),
                sample_created_at(rng, now),
            )
        )

    for _ in range(MALFORMED_COUNTS["future_dated"]):
        future = now + timedelta(days=3, hours=rng.uniform(-2, 2))
        rows.append(
            (
                str(uuid.uuid4()), rng.choice(all_ids), "purchase",
                max(1, int(rng.lognormvariate(_median_to_mu(5_000), 0.9))),
                weighted_choice(rng, CURRENCY_WEIGHTS), rng.choice(MERCHANTS),
                rng.randint(1, N_CATEGORIES), "pending", None, future,
            )
        )

    for _ in range(MALFORMED_COUNTS["orphan_category"]):
        rows.append(
            (
                str(uuid.uuid4()), rng.choice(all_ids), "purchase",
                max(1, int(rng.lognormvariate(_median_to_mu(5_000), 0.9))),
                weighted_choice(rng, CURRENCY_WEIGHTS), rng.choice(MERCHANTS),
                9999, "settled", round(rng.uniform(0, 1), 3), sample_created_at(rng, now),
            )
        )

    return rows, poison_tx_ids


def gen_sanctions_entries(rng, sanctioned_holder_names):
    rows = []
    for name in sanctioned_holder_names:
        rows.append((name, rng.choice(["US", "GB", "FR", "DE", "IR", "RU"]), "OFAC-like-list"))

    fake_first = FIRST_NAMES + ["Zephyrine", "Cassius", "Odalys", "Thaddeus", "Wilhelmina"]
    fake_last = LAST_NAMES + ["Voskresenskaya", "Almacita", "Brennholt", "Cortavez"]
    used = set(sanctioned_holder_names)
    for _ in range(N_SANCTIONS - len(sanctioned_holder_names)):
        while True:
            name = f"{rng.choice(fake_first)} {rng.choice(fake_last)} {rng.randint(1000,9999)}"
            if name not in used:
                used.add(name)
                break
        rows.append((name, rng.choice(["US", "GB", "FR", "DE", "IR", "RU"]), "OFAC-like-list"))
    return rows


def main():
    rng = random.Random(SEED)
    now = datetime.now(timezone.utc)
    t0 = time.monotonic()

    accounts = gen_accounts(rng)
    hot_accounts = [a for a in accounts if a["is_hot"]]
    cold_accounts = [a for a in accounts if not a["is_hot"]]

    role_pool = list(cold_accounts)
    rng.shuffle(role_pool)
    poison_accounts = role_pool[:N_POISON_ACCOUNTS]
    sanctioned_accounts = role_pool[N_POISON_ACCOUNTS : N_POISON_ACCOUNTS + N_SANCTIONED_MATCHES]
    clean_account = role_pool[N_POISON_ACCOUNTS + N_SANCTIONED_MATCHES]
    cold_sample = role_pool[N_POISON_ACCOUNTS + N_SANCTIONED_MATCHES + 1 : N_POISON_ACCOUNTS + N_SANCTIONED_MATCHES + 61]

    poison_account_ids = [a["id"] for a in poison_accounts]
    sanctioned_account_ids = [a["id"] for a in sanctioned_accounts]
    hot_account_ids = [a["id"] for a in hot_accounts]
    cold_account_id_sample = [a["id"] for a in cold_sample]
    clean_account_id = clean_account["id"]

    categories = gen_categories(rng)
    normal_txs = gen_normal_transactions(rng, accounts, now)
    malformed_txs, poison_tx_ids = gen_malformed_transactions(rng, accounts, set(poison_account_ids), now)
    all_txs = normal_txs + malformed_txs
    rng.shuffle(all_txs)

    sanctions_entries = gen_sanctions_entries(rng, [a["holder_name"] for a in sanctioned_accounts])

    print(f"generated {len(accounts)} accounts, {len(all_txs)} transactions, "
          f"{len(categories)} categories, {len(sanctions_entries)} sanctions entries "
          f"in {time.monotonic() - t0:.1f}s (SEED={SEED})")

    with psycopg.connect(DSN, autocommit=False) as conn:
        with conn.cursor() as cur:
            cur.execute(
                "TRUNCATE accounts, merchant_categories, transactions, sanctions_entries, "
                "notifications, audit_events RESTART IDENTITY"
            )

            with cur.copy(
                "COPY accounts (id, holder_name, currency, balance_minor, is_hot, created_at) FROM STDIN"
            ) as copy:
                for a in accounts:
                    copy.write_row((a["id"], a["holder_name"], a["currency"], a["balance_minor"], a["is_hot"], a["created_at"]))

            with cur.copy(
                "COPY merchant_categories (id, code, label, risk_weight) FROM STDIN"
            ) as copy:
                for c in categories:
                    copy.write_row((c["id"], c["code"], c["label"], c["risk_weight"]))

            with cur.copy(
                "COPY transactions (id, account_id, kind, amount_minor, currency, merchant_name, "
                "category_id, status, risk_score, created_at) FROM STDIN"
            ) as copy:
                for row in all_txs:
                    copy.write_row(row)

            with cur.copy(
                "COPY sanctions_entries (normalized_name, country, list_source) FROM STDIN"
            ) as copy:
                for row in sanctions_entries:
                    copy.write_row(row)

            cur.execute("ANALYZE accounts, transactions, merchant_categories, sanctions_entries")

            manifest = {
                "seed": SEED,
                "generated_at": now.isoformat(),
                "counts": {
                    "accounts": len(accounts),
                    "transactions": len(all_txs),
                    "merchant_categories": len(categories),
                    "sanctions_entries": len(sanctions_entries),
                },
                "hot_account_ids": hot_account_ids,
                "poison_account_ids": poison_account_ids,
                "sanctioned_account_ids": sanctioned_account_ids,
                "clean_account_id": clean_account_id,
                "cold_account_ids_sample": cold_account_id_sample,
                "malformed_row_count": len(malformed_txs),
            }
            cur.execute(
                "INSERT INTO seed_manifest (id, data) VALUES (1, %s::jsonb) "
                "ON CONFLICT (id) DO UPDATE SET data = EXCLUDED.data",
                (json.dumps(manifest),),
            )
        conn.commit()

    if poison_tx_ids:
        import redis as redis_sync

        redis_url = os.environ.get("REDIS_URL", "redis://localhost:6379/0")
        r = redis_sync.from_url(redis_url, decode_responses=True)
        r.delete("poison:tx_by_account")
        r.hset("poison:tx_by_account", mapping=poison_tx_ids)

    elapsed = time.monotonic() - t0
    print(f"seed complete in {elapsed:.1f}s")
    print("row counts:")
    for table, count in manifest["counts"].items():
        print(f"  {table}: {count}")
    print(f"hot account ids: {hot_account_ids}")
    print(f"poison account ids: {poison_account_ids[:5]}... ({len(poison_account_ids)} total)")
    print(f"malformed rows: {manifest['malformed_row_count']}")
    print(f"clean_account_id: {clean_account_id}")


if __name__ == "__main__":
    main()