import argparse
import re
from pathlib import Path


VAL_RE = re.compile(r"Val result: mIoU/mAcc/allAcc ([0-9.]+)/([0-9.]+)/([0-9.]+)\.")
CLASS_RE = re.compile(r"Class_(3|4|6)-([a-zA-Z_]+) Result: iou/accuracy ([0-9.]+)/([0-9.]+)")


def parse_log(path: Path):
    records = []
    current = {}
    for line in path.read_text(encoding="utf-8", errors="ignore").splitlines():
        val_match = VAL_RE.search(line)
        if val_match:
            if current:
                records.append(current)
            current = {
                "mIoU": float(val_match.group(1)),
                "mAcc": float(val_match.group(2)),
                "allAcc": float(val_match.group(3)),
            }
            continue
        class_match = CLASS_RE.search(line)
        if class_match and current:
            class_id = int(class_match.group(1))
            current[class_id] = float(class_match.group(3))
    if current:
        records.append(current)
    return records


def score(record):
    beam = record.get(3, -1.0)
    tail_mean = (record.get(3, 0.0) + record.get(4, 0.0) + record.get(6, 0.0)) / 3.0
    miou = record.get("mIoU", -1.0)
    return beam, tail_mean, miou


def main():
    parser = argparse.ArgumentParser(description="Pick the best beam-focused run from validation logs.")
    parser.add_argument("logs", nargs="+", help="train.log paths to compare")
    args = parser.parse_args()

    best = None
    for raw_path in args.logs:
        path = Path(raw_path)
        records = parse_log(path)
        if not records:
            print(f"{path}: no validation records found")
            continue
        best_record = max(records, key=score)
        beam, tail_mean, miou = score(best_record)
        print(
            f"{path}: beam={beam:.4f} column={best_record.get(4, 0.0):.4f} "
            f"door={best_record.get(6, 0.0):.4f} tail_mean={tail_mean:.4f} mIoU={miou:.4f}"
        )
        candidate = (score(best_record), str(path), best_record)
        if best is None or candidate[0] > best[0]:
            best = candidate

    if best is None:
        raise SystemExit(1)

    _, best_path, best_record = best
    print(
        "winner: {} (beam={:.4f}, column={:.4f}, door={:.4f}, mIoU={:.4f})".format(
            best_path,
            best_record.get(3, 0.0),
            best_record.get(4, 0.0),
            best_record.get(6, 0.0),
            best_record.get("mIoU", 0.0),
        )
    )


if __name__ == "__main__":
    main()