#!/usr/bin/env python3
import os
import json

INPUT_DIRS = [
  "/home/yixiaozeng/rag-robust-eval/main/results_new",
  "/home/yixiaozeng/rag-robust-eval/main/results_new_llama_3-2",
  "/home/yixiaozeng/rag-robust-eval/main/results_new_qwen3_large",
  "/home/yixiaozeng/rag-robust-eval/main/results_new_qwen3_small",
]
OUTPUT_DIR = "/home/yixiaozeng/rag-robust-eval/main/results_merged"


def load_json(path):
  with open(path, "r", encoding="utf-8") as f:
    return json.load(f)


def merge_reports(filenames, input_dirs, output_dir):
  os.makedirs(output_dir, exist_ok=True)
  for fname in filenames:
    merged_metrics = {}
    merged_analysis = {}
    summary = None

    for d in input_dirs:
      path = os.path.join(d, fname)
      if not os.path.isfile(path):
        continue
      data = load_json(path)

      for model, vals in data.get("metrics", {}).items():
        if model not in merged_metrics:
          merged_metrics[model] = vals

      for model, vals in data.get("analysis", {}).items():
        if model not in merged_analysis:
          merged_analysis[model] = vals

      if summary is None and "summary" in data:
        summary = data["summary"]

    if summary is None:
      summary = {}
    summary["models_evaluated"] = sorted(list(merged_metrics.keys()))

    out = {"metrics": merged_metrics, "analysis": merged_analysis, "summary": summary}
    out_path = os.path.join(output_dir, fname)
    with open(out_path, "w", encoding="utf-8") as f:
      json.dump(out, f, indent=2, ensure_ascii=False)

    print(f"Wrote {out_path}")


def main():
  base_dir = INPUT_DIRS[0]
  filenames = sorted([f for f in os.listdir(base_dir) if f.endswith("_final_report.json")])
  merge_reports(filenames, INPUT_DIRS, OUTPUT_DIR)


if __name__ == "__main__":
  main() 