#!/usr/bin/env python3
"""Convert Plan D RAG predict output (2618 rows aligned to Test_aug_sbert.csv) to Zindi CSV.
Plan D predicts 1:1 with Test.csv rows (no canonical dedup), so just zip preds with Test.csv IDs.
"""
import json, argparse, pathlib
import pandas as pd

WORK = pathlib.Path('/mnt/msrh/Magic_submission')
SUFFIX = ' This answer provides accurate and complete health information.'


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--pred_jsonl', required=True)
    ap.add_argument('--test_csv', default=str(WORK / 'data' / 'Test.csv'))
    ap.add_argument('--out_csv', required=True)
    ap.add_argument('--no_suffix', action='store_true')
    args = ap.parse_args()

    preds = []
    with open(args.pred_jsonl) as f:
        for line in f:
            r = json.loads(line)
            p = (r.get('predict') or r.get('prediction') or '').strip()
            preds.append(p)

    test = pd.read_csv(args.test_csv)
    assert len(preds) == len(test), f"len mismatch: {len(preds)} preds vs {len(test)} test"

    rows = []
    for tid, pred in zip(test['ID'], preds):
        ans = pred + ('' if args.no_suffix else SUFFIX)
        rows.append({'ID': tid, 'TargetRLF1': ans, 'TargetR1F1': ans, 'TargetLLM': ans})
    pd.DataFrame(rows, columns=['ID', 'TargetRLF1', 'TargetR1F1', 'TargetLLM']).to_csv(args.out_csv, index=False)
    print(f"wrote {args.out_csv} ({len(rows)} rows)")


if __name__ == '__main__':
    main()