#!/usr/bin/env python3
"""Step 2 of the canonical tsqa_9010 build (2026-06-25): drop multi-<ts> rows.

After build_9010.py (seed=42 90/10 + de-leak), keep only rows whose prompt has
exactly one <ts> and one </ts> (single-series questions); multi-series questions
are removed. This took the split from 3576->3519 test / 32182->31556 train.

Usage: python fix_ts.py [DIR]   (DIR default $ROOT/Benchmark_eval/TSQA)
"""
import os as _os  # RELOC_ROOT_MARK
_RELOC_ROOT = _os.path.abspath(__file__)
for _m in ('/results/', '/Benchmark_eval/', '/Benchmarks/'):
    if _m in _RELOC_ROOT:
        _RELOC_ROOT = _RELOC_ROOT.split(_m)[0]; break
import json, os, sys, collections
D = sys.argv[1] if len(sys.argv) > 1 else (_RELOC_ROOT + "/Benchmark_eval/TSQA")

def ok(p): return p.count("<ts>") == 1 and p.count("</ts>") == 1
def load(p): return [json.loads(l) for l in open(p)]
def dump(rows, p):
    with open(p, "w") as f:
        for r in rows:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")

for rel in ["train/sft_pairs.jsonl", "train/raw_values.jsonl", "test_requests.jsonl"]:
    src = f"{D}/{rel}"; bak = src + ".multits_bak"
    if os.path.exists(src) and not os.path.exists(bak):
        os.rename(src, bak)

sft = load(f"{D}/train/sft_pairs.jsonl.multits_bak")
sft_ok = [r for r in sft if ok(r["prompt"])]; ids = set(r["id"] for r in sft_ok)
rv = [r for r in load(f"{D}/train/raw_values.jsonl.multits_bak") if r["id"] in ids]
dump(sft_ok, f"{D}/train/sft_pairs.jsonl"); dump(rv, f"{D}/train/raw_values.jsonl")

te = load(f"{D}/test_requests.jsonl.multits_bak"); te_ok = [r for r in te if ok(r["prompt"])]
dump(te_ok, f"{D}/test_requests.jsonl")

print("after fix_ts: train=%d test=%d" % (len(sft_ok), len(te_ok)))
print("  test multi<ts> remaining:", sum(1 for r in te_ok if not ok(r["prompt"])))
print("  test task_type:", dict(collections.Counter(r.get("task_type") for r in te_ok)))