"""Extract a parent->child training set from a thegrid fossil record (read-only).

One row per (epoch, child genome): the child's earliest recorded parent in that epoch,
exact self-copies (parent == child) excluded. Label: the child genome was born at least
twice in that epoch (it copied itself, so it "founded" something), from genome_stats.births.

mutation_origins is deliberately NOT used: it only records origins whose child was born,
so its presence would leak the label.

usage: python3 extract.py <history.sqlite3> <out.jsonl.gz>
"""
import gzip
import json
import sqlite3
import sys

db, out = sys.argv[1], sys.argv[2]
c = sqlite3.connect(f"file:{db}?mode=ro", uri=True)

src = {gid: s for gid, s in c.execute("SELECT genome_id, source FROM genomes")}
print("genomes", len(src), file=sys.stderr)
births = {(e, g): b for e, g, b in c.execute("SELECT epoch, genome_id, births FROM genome_stats")}
print("stats", len(births), file=sys.stderr)

rows = c.execute(
    "SELECT epoch, child_genome_id, parent_genome_id, first_tick FROM transitions "
    "WHERE parent_genome_id != child_genome_id ORDER BY epoch, child_genome_id, first_tick")
n = kept = skipped = 0
last = None
with gzip.open(out, "wt", encoding="utf-8") as f:
    for epoch, child, parent, tick in rows:
        n += 1
        if (epoch, child) == last:
            continue          # a later parent of the same child in the same epoch
        last = (epoch, child)
        if child not in src or parent not in src:
            skipped += 1
            continue
        cb = births.get((epoch, child), 0)
        f.write(json.dumps({
            "epoch": epoch,
            "child_id": child,
            "parent": src[parent].split(" · "),
            "child": src[child].split(" · "),
            "child_births": cb,
            "parent_births": births.get((epoch, parent), 0),
            "first_tick": tick,
            "label": int(cb >= 2),
        }, separators=(",", ":")) + "\n")
        kept += 1
print(json.dumps({"transitions_scanned": n, "rows": kept, "skipped_missing_source": skipped}), file=sys.stderr)