#!/usr/bin/env python3
"""Stop PP-OCR training when held-out normalized edit similarity stalls or diverges."""

from __future__ import annotations

import argparse
import os
import re
import signal
import subprocess
import time
from pathlib import Path


PATTERN = re.compile(r"cur metric,.*norm_edit_dis: ([0-9.eE+-]+)")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--log", type=Path, required=True)
    parser.add_argument("--pid-file", type=Path, required=True)
    parser.add_argument("--min-evals", type=int, default=5)
    parser.add_argument("--patience", type=int, default=5)
    parser.add_argument("--min-delta", type=float, default=0.002)
    parser.add_argument("--divergence-delta", type=float, default=0.03)
    parser.add_argument("--divergence-patience", type=int, default=2)
    parser.add_argument("--initial-best", type=float, default=-1.0)
    parser.add_argument("--repo-id")
    parser.add_argument("--checkpoint-prefix", type=Path)
    parser.add_argument("--config", type=Path)
    parser.add_argument("--dictionary", type=Path)
    parser.add_argument("--uploader", type=Path)
    args = parser.parse_args()

    best = args.initial_best
    stale = diverged = seen = 0
    while True:
        if not args.pid_file.exists():
            time.sleep(2)
            continue
        pid = int(args.pid_file.read_text())
        try:
            os.kill(pid, 0)
        except ProcessLookupError:
            return

        scores = [float(value) for value in PATTERN.findall(args.log.read_text(errors="ignore"))]
        for score in scores[seen:]:
            seen += 1
            improved = score > best + args.min_delta
            if improved:
                best, stale = score, 0
            else:
                stale += 1
            diverged = diverged + 1 if score < best - args.divergence_delta else 0
            print(
                f"eval={seen} score={score:.6f} best={best:.6f} "
                f"stale={stale} diverged={diverged}",
                flush=True,
            )
            upload_values = (
                args.repo_id,
                args.checkpoint_prefix,
                args.config,
                args.dictionary,
                args.uploader,
            )
            if improved and all(upload_values):
                time.sleep(5)
                try:
                    subprocess.run(
                        [
                            str(args.uploader),
                            "--repo-id", args.repo_id,
                            "--checkpoint-prefix", str(args.checkpoint_prefix),
                            "--tag", f"eval-{seen:05d}",
                            "--metric", str(score),
                            "--config", str(args.config),
                            "--dictionary", str(args.dictionary),
                        ],
                        check=True,
                    )
                except Exception as error:
                    print(f"UPLOAD_FAILED {error}", flush=True)
            if seen >= args.min_evals and (
                stale >= args.patience or diverged >= args.divergence_patience
            ):
                os.kill(pid, signal.SIGTERM)
                print("EARLY_STOP", flush=True)
                return
        time.sleep(3)


if __name__ == "__main__":
    main()