#!/usr/bin/env python3
import argparse, json, os, re, sys

METRICS = {
    "Rank corr": True,
    "Rank correlation": True,
    "NN Tm correlation": True,
    "Naive Tm correlation": True,
}

def parse(text):
    results = {}
    for name, higher in METRICS.items():
        pat = re.compile(re.escape(name) + r"\s*:\s*([+-]?[0-9]*\.?[0-9]+)", re.IGNORECASE)
        matches = pat.finditer(text)
        vals = []
        for m in matches:
            vals.append((float(m.group(1)), higher, m.group(0).strip()))
        if vals:
            results[name] = vals
    return results

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--stdout", required=True)
    ap.add_argument("--elapsed", type=float, default=None)
    args = ap.parse_args()
    text = open(args.stdout, encoding="utf-8", errors="replace").read()
    parsed = parse(text)

    score = None
    feedback_parts = []

    if "NN Tm correlation" in parsed:
        val, _, raw = parsed["NN Tm correlation"][0]
        score = abs(val)
        feedback_parts.append(f"NN Tm corr={val}")
    elif "Rank correlation" in parsed:
        vals = [abs(v) for v, _, _ in parsed["Rank correlation"]]
        score = sum(vals) / len(vals)
        feedback_parts.append(f"avg |Rank corr|={score:.4f}")
    elif "Rank corr" in parsed:
        vals = [abs(v) for v, _, _ in parsed["Rank corr"]]
        score = sum(vals) / len(vals)
        feedback_parts.append(f"avg |Rank corr|={score:.4f}")
    elif "Naive Tm correlation" in parsed:
        val, _, raw = parsed["Naive Tm correlation"][0]
        score = abs(val)
        feedback_parts.append(f"Naive Tm corr={val}")

    if score is None and args.elapsed is not None:
        score = -(args.elapsed * 1000)
        feedback_parts.append(f"time={args.elapsed*1000:.1f}ms")
    elif score is None:
        m = re.search(r"SUMMARY:\s*(\d+)/(\d+)\s*passed", text)
        if m:
            score = float(m.group(1))
            feedback_parts.append(f"passes={m.group(1)}/{m.group(2)}")
        else:
            sys.stderr.write("[score] no metric found in output.\n")
            sys.exit(3)

    result = {
        "score": score,
        "examples": [{
            "id": "metric",
            "score": score,
            "pass": True,
            "feedback": " | ".join(feedback_parts),
        }],
        "feedback": " | ".join(feedback_parts),
    }

    out = os.environ.get("ATLAS_OPTIMIZE_RESULT")
    payload = json.dumps(result)
    if out:
        with open(out, "w", encoding="utf-8") as fh:
            fh.write(payload)
    else:
        sys.stdout.write(payload)

if __name__ == "__main__":
    main()