"""Build Zindi submission CSV from JSONL using ROW ORDER (not ID) match.
Test.csv row N <-> JSONL row N. Use SampleSubmission for column header + non-test rows."""
import csv, json, argparse, pathlib

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--jsonl", required=True)
    ap.add_argument("--test_csv", default="/mnt/msrh/Magic_submission/data/Test.csv")
    ap.add_argument("--sample", default="/mnt/msrh/Magic_submission/data/SampleSubmission.csv")
    ap.add_argument("--out_csv", required=True)
    args = ap.parse_args()

    preds = [json.loads(l).get("predict","") for l in open(args.jsonl)]
    with open(args.test_csv, newline="") as f:
        test_ids = [r["ID"] for r in csv.DictReader(f)]
    assert len(preds) == len(test_ids), f"len(preds)={len(preds)} != len(test_ids)={len(test_ids)}"
    id2pred = dict(zip(test_ids, preds))

    with open(args.sample, newline="") as f:
        sample_cols = next(csv.reader(f))
    target_cols = [c for c in sample_cols if c.startswith("Target")]
    with open(args.sample, newline="") as f:
        sample_rows = list(csv.DictReader(f))

    pathlib.Path(args.out_csv).parent.mkdir(parents=True, exist_ok=True)
    n_fill = 0
    with open(args.out_csv, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=sample_cols)
        w.writeheader()
        for r in sample_rows:
            id_ = r["ID"]
            text = id2pred.get(id_)
            row = {"ID": id_}
            if text is not None:
                for tc in target_cols: row[tc] = text
                n_fill += 1
            else:
                for c in sample_cols[1:]: row[c] = r.get(c, "")
            w.writerow(row)
    print(f"wrote {args.out_csv}  filled={n_fill}/{len(sample_rows)}")

if __name__ == "__main__":
    main()