# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Inference is a function of the sample.

A JSON-lines file of 20,000 records whose postal_code is an integer in the
first 15,000 and a string such as "SW1A 1AA" after that. Infers the schema
from the first 1,000 rows, the first 18,000 and all of them, in file order
and shuffled; then compares each inference with a frozen schema (a drift
check) and counts the alerts.
"""
import argparse, json, random

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--rows", type=int, default=20_000, help="Records.")
p.add_argument("--switch", type=int, default=15_000, help="Record from which postal_code is a string.")
p.add_argument("--seed", type=int, default=0, help="Random seed.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
random.seed(args.seed)

lines = [json.dumps({"id": i, "postal_code": random.randrange(10000, 99999) if i < args.switch else f"SW{random.randrange(1, 9)}A {random.randrange(1, 9)}AA",
                     "weight": round(random.uniform(0.1, 30), 2)}) for i in range(args.rows)]
TYPES = {int: "integer", float: "number", str: "string", bool: "boolean", type(None): "null"}
def infer(sample):
    out = {}
    for line in sample:
        for k, v in json.loads(line).items():
            t = TYPES[type(v)]; prev = out.get(k)
            out[k] = t if prev in (None, t) else ("number" if {prev, t} == {"integer", "number"} else "string")   # widen on conflict
    return out
FROZEN = {"id": "integer", "postal_code": "string", "weight": "number"}             # the contract, agreed and versioned
def drift(schema): return [f"{k}: {schema.get(k)} (expected {v})" for k, v in FROZEN.items() if schema.get(k) != v]
rows = []
shuffled = lines[:]; random.shuffle(shuffled)
for order, data in (("file order", lines), ("shuffled", shuffled)):
    for n in (1000, 18000, args.rows):
        sch = infer(data[:n]); alerts = drift(sch)
        rows.append({"order": order, "sample": n, "postal_code": sch["postal_code"], "weight": sch["weight"], "alerts": len(alerts), "alert": "; ".join(alerts) or "none"})
summary = {"distinct_inferences": len({(r["postal_code"], r["weight"]) for r in rows})}
for r in rows: print(f"{r['order']:<10} first {r['sample']:>6,} rows: postal_code {r['postal_code']:<8} weight {r['weight']:<7} drift alerts: {r['alert']}")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"summary": summary, "tables": {"inferences": rows}}, f)