#!/usr/bin/env python3
"""FlashMemory synthetic yield data exercise (Python 3 standard library only)."""
import csv
import json
from collections import defaultdict
from pathlib import Path

rows = list(csv.DictReader(Path(__file__).with_name("wafer-yield.csv").open(encoding="utf-8", newline="")))
accepted = []
seen = set()
excluded = defaultdict(int)
for row in rows:
    if not row["tested"] or not row["passed"]:
        excluded["missing_or_invalid_denominator"] += 1
        continue
    if row["test_round"] != "1":
        excluded["retest_outside_first_pass"] += 1
        continue
    key = row["wafer_id"]
    if key in seen:
        excluded["duplicate_wafer_first_pass"] += 1
        continue
    tested, passed = int(row["tested"]), int(row["passed"])
    if not 0 <= passed <= tested or tested <= 0:
        excluded["missing_or_invalid_denominator"] += 1
        continue
    if row["period"] not in {"before", "after"} or row["product"] not in {"A", "B"}:
        raise ValueError("Unexpected example stratum")
    seen.add(key)
    accepted.append({**row, "tested": tested, "passed": passed})

def aggregate(group):
    tested = sum(row["tested"] for row in group)
    passed = sum(row["passed"] for row in group)
    if not tested:
        raise ValueError("empty stratum")
    return {"wafers": len(group), "passed": passed, "tested": tested, "yield_pct": round(100 * passed/tested, 2)}

overall = {}
by_product = {}
standardized = {}
for period in ("before", "after"):
    sample = [row for row in accepted if row["period"] == period]
    overall[period] = aggregate(sample)
    by_product[period] = {product: aggregate([row for row in sample if row["product"] == product]) for product in ("A", "B")}
    # A 50% + B 50% is an explicit hypothetical reference mix.
    standardized[period] = round(sum(by_product[period][p]["yield_pct"] for p in ("A", "B"))/2, 2)

summary = {
    "dataset": "SYNTHETIC_TRAINING_ONLY",
    "raw_rows": len(rows),
    "valid_first_pass_wafers": len(accepted),
    "excluded": dict(excluded),
    "overall": overall,
    "by_product": by_product,
    "equal_mix_50_50_yield_pct": standardized,
    "caveat": "Product A is recorded only on chamber C1, B only on C2; chamber effects are not separable from product effects. Observational data do not establish a root cause."
}
assert summary["raw_rows"] == 19 and summary["valid_first_pass_wafers"] == 16
assert summary["excluded"] == {"duplicate_wafer_first_pass": 1, "retest_outside_first_pass": 1, "missing_or_invalid_denominator": 1}
assert [overall[p]["yield_pct"] for p in ("before", "after")] == [92.5, 77.5]
assert [standardized[p] for p in ("before", "after")] == [85.0, 85.0]
print(json.dumps(summary, ensure_ascii=False, indent=2))