import json, re
from pathlib import Path
from collections import Counter

BASE_DIR = Path("/content/drive/MyDrive/humomni_track2_phase1_data/V4_FINAL_K1_MEMORY_AGENT_TOP6/full500")

IN_PATH = BASE_DIR / "submission_V4_FINAL_K1_memory_agent_top6_full500.jsonl"
OUT_PATH = BASE_DIR / "submission_V4_FINAL_K1_memory_agent_top6_clean_light_plus_full500.jsonl"
AUDIT_PATH = BASE_DIR / "audit_memory_agent_top6_clean_light_plus_full500.json"
REMOVED_PATH = BASE_DIR / "removed_turns_clean_light_plus.jsonl"

print("IN_PATH:", IN_PATH)
print("exists:", IN_PATH.exists())

artifact_patterns = [
    r"\bNS\s*REPLY\b",
    r"\bNSWYR\b",
    r"\bREPLYING\s*:",
    r"\bNO\s*REPLY\b",
    r"no description available",
    r"not available from this image alone",
    r"from the provided audio description alone",
    r"no image displayed",
    r"no active movement occurs",
    r"no actions? take place",
    r"no reaction shown",
    r"no response available",
    r"no activity related",
    r"no additional context",
    r"no additional elements",
]

artifact_re = re.compile("|".join(artifact_patterns), re.I)

def is_artifact_turn(content):
    c = str(content).strip()
    return bool(artifact_re.search(c))

rows = [json.loads(line) for line in open(IN_PATH, "r", encoding="utf-8")]

cleaned = []
removed = []

for row in rows:
    turns = row["model_response_list"]
    kept = []

    for t in turns:
        content = t.get("content", "")
        if is_artifact_turn(content):
            removed.append({
                "question_id": row["question_id"],
                "time": t.get("time"),
                "content": content
            })
        else:
            kept.append(t)

    # 绝不删空一个 video
    if not kept:
        kept = [turns[0]]

    cleaned.append({
        "question_id": row["question_id"],
        "model_response_list": kept
    })

with open(OUT_PATH, "w", encoding="utf-8") as f:
    for row in cleaned:
        f.write(json.dumps(row, ensure_ascii=False) + "\n")

with open(REMOVED_PATH, "w", encoding="utf-8") as f:
    for row in removed:
        f.write(json.dumps(row, ensure_ascii=False) + "\n")

turn_counts = [len(r["model_response_list"]) for r in cleaned]
all_turns = [t for r in cleaned for t in r["model_response_list"]]

audit = {
    "variant": "memory_agent_top6_clean_light_plus",
    "N_records": len(cleaned),
    "unique_question_ids": len(set(r["question_id"] for r in cleaned)),
    "duplicate_question_ids": len(cleaned) - len(set(r["question_id"] for r in cleaned)),
    "empty_outputs": sum(1 for r in cleaned if not r["model_response_list"]),
    "avg_turns_per_video": sum(turn_counts) / len(turn_counts),
    "max_turns_per_video": max(turn_counts),
    "turn_count_distribution": dict(Counter(turn_counts)),
    "removed_turns": len(removed),
    "time_min": min(t["time"] for t in all_turns),
    "time_max": max(t["time"] for t in all_turns),
    "chronological_per_video": all(
        all(r["model_response_list"][i]["time"] <= r["model_response_list"][i+1]["time"]
            for i in range(len(r["model_response_list"]) - 1))
        for r in cleaned
    ),
    "output_path": str(OUT_PATH),
    "removed_path": str(REMOVED_PATH),
}

with open(AUDIT_PATH, "w", encoding="utf-8") as f:
    json.dump(audit, f, indent=2, ensure_ascii=False)

print(json.dumps(audit, indent=2, ensure_ascii=False))

print("\nRemoved examples:")
for x in removed[:50]:
    print(x)