import os
import sys
import json
import matplotlib.pyplot as plt

def extract_last_valid_ap(metrics_path):
    last_ap = None
    try:
        with open(metrics_path, 'r') as f:
            for line in f:
                try:
                    record = json.loads(line)
                    if "bbox/AP" in record and isinstance(record["bbox/AP"], (int, float)):
                        last_ap = round(record["bbox/AP"], 2)
                except json.JSONDecodeError:
                    continue
    except Exception as e:
        print(f"[Error] Reading {metrics_path}: {e}")
    return last_ap

def main(continual_dir, config_paths):
    datasets = []
    ap_scores = []

    for cfg in config_paths:
        dataset_name = cfg.split("_ft_")[1].split("_10ep")[0]
        metrics_path = os.path.join(continual_dir, dataset_name, "metrics.json")

        if os.path.exists(metrics_path):
            ap = extract_last_valid_ap(metrics_path)
            if ap is not None:
                datasets.append(dataset_name)
                ap_scores.append(ap)
            else:
                print(f"[Warning] No valid bbox/AP found in {metrics_path}")
        else:
            print(f"[Warning] Missing: {metrics_path}")

    # 写入两行 CSV
    csv_path = os.path.join(continual_dir, "continual_summary.csv")
    with open(csv_path, "w") as f:
        f.write(",".join(datasets) + "\n")
        f.write(",".join([str(v) for v in ap_scores]) + "\n")
    print(f"✅ 汇总 CSV 写入: {csv_path}")

    # 画图
    plt.figure(figsize=(12, 5))
    plt.plot(datasets, ap_scores, marker='o')
    plt.xticks(rotation=45, ha="right")
    plt.xlabel("Dataset")
    plt.ylabel("bbox/AP (%)")
    plt.title("Continual Learning Performance (bbox/AP)")
    plt.grid(True)
    plt.tight_layout()

    plot_path = os.path.join(continual_dir, "continual_ap_curve.png")
    plt.savefig(plot_path)
    print(f"✅ 图像保存: {plot_path}")

if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("Usage: python summarize_and_plot.py <continual_dir> <config1> <config2> ...")
        sys.exit(1)
    continual_dir = sys.argv[1]
    config_paths = sys.argv[2:]
    main(continual_dir, config_paths)
