#!/usr/bin/env python3
"""Run hyperparameter tuning for OCR services."""

import sys
import argparse
from raytune_ocr import (
    check_workers, create_trainable, run_tuner, analyze_results,
    paddle_ocr_payload, doctr_payload, easyocr_payload,
    PADDLE_OCR_SEARCH_SPACE, DOCTR_SEARCH_SPACE, EASYOCR_SEARCH_SPACE,
    PADDLE_OCR_CONFIG_KEYS, DOCTR_CONFIG_KEYS, EASYOCR_CONFIG_KEYS,
)

SERVICES = {
    "paddle": {
        "ports": [8002],
        "payload_fn": paddle_ocr_payload,
        "search_space": PADDLE_OCR_SEARCH_SPACE,
        "config_keys": PADDLE_OCR_CONFIG_KEYS,
        "name": "PaddleOCR",
    },
    "doctr": {
        "ports": [8003],
        "payload_fn": doctr_payload,
        "search_space": DOCTR_SEARCH_SPACE,
        "config_keys": DOCTR_CONFIG_KEYS,
        "name": "DocTR",
    },
    "easyocr": {
        "ports": [8002],
        "payload_fn": easyocr_payload,
        "search_space": EASYOCR_SEARCH_SPACE,
        "config_keys": EASYOCR_CONFIG_KEYS,
        "name": "EasyOCR",
    },
}

def main():
    parser = argparse.ArgumentParser(description="Run OCR hyperparameter tuning")
    parser.add_argument("--service", choices=["paddle", "doctr", "easyocr"], required=True)
    parser.add_argument("--samples", type=int, default=64, help="Number of samples")
    args = parser.parse_args()

    cfg = SERVICES[args.service]
    print(f"\n{'='*50}")
    print(f"Hyperparameter Tuning: {cfg['name']}")
    print(f"Samples: {args.samples}")
    print(f"{'='*50}\n")

    # Check workers
    healthy = check_workers(cfg["ports"], cfg["name"])

    # Create trainable and run tuning
    trainable = create_trainable(cfg["ports"], cfg["payload_fn"])
    results = run_tuner(
        trainable=trainable,
        search_space=cfg["search_space"],
        num_samples=args.samples,
        num_workers=len(healthy),
    )

    # Analyze results
    df = analyze_results(
        results,
        output_folder="results",
        prefix=f"raytune_{args.service}",
        config_keys=cfg["config_keys"],
    )

    print(f"\n{'='*50}")
    print("Tuning complete!")
    print(f"{'='*50}")

if __name__ == "__main__":
    main()
