# raytune_ocr.py
# Shared Ray Tune utilities for OCR hyperparameter optimization
#
# Usage:
#   from raytune_ocr import check_workers, create_trainable, run_tuner, analyze_results
#
# Environment variables:
#   OCR_HOST: Host for OCR services (default: localhost)

import os
from datetime import datetime
from typing import List, Dict, Any, Callable

import requests
import pandas as pd

import ray
from ray import tune
from ray.tune.search.optuna import OptunaSearch


def check_workers(
    ports: List[int],
    service_name: str = "OCR",
    timeout: int = 180,
    interval: int = 5,
) -> List[str]:
    """
    Wait for workers to be fully ready (model + dataset loaded) and return healthy URLs.

    Args:
        ports: List of port numbers to check
        service_name: Name for error messages
        timeout: Max seconds to wait for each worker
        interval: Seconds between retries

    Returns:
        List of healthy worker URLs

    Raises:
        RuntimeError if no healthy workers found after timeout
    """
    import time

    host = os.environ.get("OCR_HOST", "localhost")
    worker_urls = [f"http://{host}:{port}" for port in ports]
    healthy_workers = []

    for url in worker_urls:
        print(f"Waiting for {url}...")
        start = time.time()

        while time.time() - start < timeout:
            try:
                health = requests.get(f"{url}/health", timeout=10).json()
                model_ok = health.get('model_loaded', False)
                dataset_ok = health.get('dataset_loaded', False)

                if health.get('status') == 'ok' and model_ok:
                    gpu = health.get('gpu_name', 'CPU')
                    print(f"✓ {url}: ready ({gpu})")
                    healthy_workers.append(url)
                    break

                elapsed = int(time.time() - start)
                print(f"  [{elapsed}s] model={model_ok}")
            except requests.exceptions.RequestException:
                elapsed = int(time.time() - start)
                print(f"  [{elapsed}s] not reachable")

            time.sleep(interval)
        else:
            print(f"✗ {url}: timeout after {timeout}s")

    if not healthy_workers:
        raise RuntimeError(
            f"No healthy {service_name} workers found.\n"
            f"Checked ports: {ports}"
        )

    print(f"\n{len(healthy_workers)}/{len(worker_urls)} workers ready\n")
    return healthy_workers


def create_trainable(ports: List[int], payload_fn: Callable[[Dict], Dict]) -> Callable:
    """
    Factory to create a trainable function for Ray Tune.

    Args:
        ports: List of worker ports for load balancing
        payload_fn: Function that takes config dict and returns API payload dict

    Returns:
        Trainable function for Ray Tune

    Note:
        Ray Tune 2.x API: tune.report(metrics_dict) - pass dict directly, NOT kwargs.
        See: https://docs.ray.io/en/latest/tune/api/doc/ray.tune.report.html
    """
    def trainable(config):
        import os
        import random
        import requests
        from ray.tune import report  # Ray 2.x: report(dict), not report(**kwargs)

        host = os.environ.get("OCR_HOST", "localhost")
        api_url = f"http://{host}:{random.choice(ports)}"
        payload = payload_fn(config)

        try:
            response = requests.post(f"{api_url}/evaluate", json=payload, timeout=None)
            response.raise_for_status()
            metrics = response.json()
            metrics["worker"] = api_url
            report(metrics)  # Ray 2.x API: pass dict directly
        except Exception as e:
            report({  # Ray 2.x API: pass dict directly
                "CER": 1.0,
                "WER": 1.0,
                "TIME": 0.0,
                "PAGES": 0,
                "TIME_PER_PAGE": 0,
                "worker": api_url,
                "ERROR": str(e)[:500]
            })

    return trainable


def run_tuner(
    trainable: Callable,
    search_space: Dict[str, Any],
    num_samples: int = 64,
    num_workers: int = 1,
    metric: str = "CER",
    mode: str = "min",
) -> tune.ResultGrid:
    """
    Initialize Ray and run hyperparameter tuning.

    Args:
        trainable: Trainable function from create_trainable()
        search_space: Dict of parameter names to tune.* search spaces
        num_samples: Number of trials to run
        num_workers: Max concurrent trials
        metric: Metric to optimize
        mode: "min" or "max"

    Returns:
        Ray Tune ResultGrid
    """
    ray.init(
        ignore_reinit_error=True,
        include_dashboard=False,
        configure_logging=False,
        _metrics_export_port=0,  # Disable metrics export to avoid connection warnings
    )
    print(f"Ray Tune ready (version: {ray.__version__})")

    tuner = tune.Tuner(
        trainable,
        tune_config=tune.TuneConfig(
            metric=metric,
            mode=mode,
            search_alg=OptunaSearch(),
            num_samples=num_samples,
            max_concurrent_trials=num_workers,
        ),
        param_space=search_space,
    )

    return tuner.fit()


def analyze_results(
    results: tune.ResultGrid,
    output_folder: str = "results",
    prefix: str = "raytune",
    config_keys: List[str] = None,
) -> pd.DataFrame:
    """
    Analyze and save tuning results.

    Args:
        results: Ray Tune ResultGrid
        output_folder: Directory to save CSV
        prefix: Filename prefix
        config_keys: List of config keys to show in best result (without 'config/' prefix)

    Returns:
        Results DataFrame
    """
    os.makedirs(output_folder, exist_ok=True)
    df = results.get_dataframe()

    # Save to CSV
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"{prefix}_results_{timestamp}.csv"
    filepath = os.path.join(output_folder, filename)
    df.to_csv(filepath, index=False)
    print(f"Results saved: {filepath}")

    # Best configuration
    best = df.loc[df["CER"].idxmin()]
    print(f"\nBest CER: {best['CER']:.6f}")
    print(f"Best WER: {best['WER']:.6f}")

    if config_keys:
        print(f"\nOptimal Configuration:")
        for key in config_keys:
            col = f"config/{key}"
            if col in best:
                val = best[col]
                if isinstance(val, float):
                    print(f"  {key}: {val:.4f}")
                else:
                    print(f"  {key}: {val}")

    return df


def correlation_analysis(df: pd.DataFrame, param_keys: List[str]) -> None:
    """
    Print correlation of numeric parameters with CER/WER.

    Args:
        df: Results DataFrame
        param_keys: List of config keys (without 'config/' prefix)
    """
    param_cols = [f"config/{k}" for k in param_keys if f"config/{k}" in df.columns]
    numeric_cols = [c for c in param_cols if df[c].dtype in ['float64', 'int64']]

    if not numeric_cols:
        print("No numeric parameters for correlation analysis")
        return

    corr_cer = df[numeric_cols + ["CER"]].corr()["CER"].sort_values(ascending=False)
    corr_wer = df[numeric_cols + ["WER"]].corr()["WER"].sort_values(ascending=False)

    print("Correlation with CER:")
    print(corr_cer)
    print("\nCorrelation with WER:")
    print(corr_wer)


# =============================================================================
# OCR-specific payload functions
# =============================================================================

def paddle_ocr_payload(config: Dict) -> Dict:
    """Create payload for PaddleOCR API. Uses pages 5-10 (first doc) for tuning."""
    return {
        "pdf_folder": "/app/dataset",
        "use_doc_orientation_classify": config.get("use_doc_orientation_classify", False),
        "use_doc_unwarping": config.get("use_doc_unwarping", False),
        "textline_orientation": config.get("textline_orientation", True),
        "text_det_thresh": config.get("text_det_thresh", 0.0),
        "text_det_box_thresh": config.get("text_det_box_thresh", 0.0),
        "text_det_unclip_ratio": config.get("text_det_unclip_ratio", 1.5),
        "text_rec_score_thresh": config.get("text_rec_score_thresh", 0.0),
        "start_page": 5,
        "end_page": 10,
        "save_output": False,
    }


def doctr_payload(config: Dict) -> Dict:
    """Create payload for DocTR API. Uses pages 5-10 (first doc) for tuning."""
    return {
        "pdf_folder": "/app/dataset",
        "assume_straight_pages": config.get("assume_straight_pages", True),
        "straighten_pages": config.get("straighten_pages", False),
        "preserve_aspect_ratio": config.get("preserve_aspect_ratio", True),
        "symmetric_pad": config.get("symmetric_pad", True),
        "disable_page_orientation": config.get("disable_page_orientation", False),
        "disable_crop_orientation": config.get("disable_crop_orientation", False),
        "resolve_lines": config.get("resolve_lines", True),
        "resolve_blocks": config.get("resolve_blocks", False),
        "paragraph_break": config.get("paragraph_break", 0.035),
        "start_page": 5,
        "end_page": 10,
        "save_output": False,
    }


def easyocr_payload(config: Dict) -> Dict:
    """Create payload for EasyOCR API. Uses pages 5-10 (first doc) for tuning."""
    return {
        "pdf_folder": "/app/dataset",
        "text_threshold": config.get("text_threshold", 0.7),
        "low_text": config.get("low_text", 0.4),
        "link_threshold": config.get("link_threshold", 0.4),
        "slope_ths": config.get("slope_ths", 0.1),
        "ycenter_ths": config.get("ycenter_ths", 0.5),
        "height_ths": config.get("height_ths", 0.5),
        "width_ths": config.get("width_ths", 0.5),
        "add_margin": config.get("add_margin", 0.1),
        "contrast_ths": config.get("contrast_ths", 0.1),
        "adjust_contrast": config.get("adjust_contrast", 0.5),
        "decoder": config.get("decoder", "greedy"),
        "beamWidth": config.get("beamWidth", 5),
        "min_size": config.get("min_size", 10),
        "start_page": 5,
        "end_page": 10,
        "save_output": False,
    }


# =============================================================================
# Search spaces
# =============================================================================

PADDLE_OCR_SEARCH_SPACE = {
    "use_doc_orientation_classify": tune.choice([True, False]),
    "use_doc_unwarping": tune.choice([True, False]),
    "textline_orientation": tune.choice([True, False]),
    "text_det_thresh": tune.uniform(0.0, 0.7),
    "text_det_box_thresh": tune.uniform(0.0, 0.7),
    "text_det_unclip_ratio": tune.choice([0.0]),
    "text_rec_score_thresh": tune.uniform(0.0, 0.7),
}

DOCTR_SEARCH_SPACE = {
    "assume_straight_pages": tune.choice([True, False]),
    "straighten_pages": tune.choice([True, False]),
    "preserve_aspect_ratio": tune.choice([True, False]),
    "symmetric_pad": tune.choice([True, False]),
    "disable_page_orientation": tune.choice([True, False]),
    "disable_crop_orientation": tune.choice([True, False]),
    "resolve_lines": tune.choice([True, False]),
    "resolve_blocks": tune.choice([True, False]),
    "paragraph_break": tune.uniform(0.01, 0.1),
}

EASYOCR_SEARCH_SPACE = {
    "text_threshold": tune.uniform(0.3, 0.9),
    "low_text": tune.uniform(0.2, 0.6),
    "link_threshold": tune.uniform(0.2, 0.6),
    "slope_ths": tune.uniform(0.0, 0.3),
    "ycenter_ths": tune.uniform(0.3, 1.0),
    "height_ths": tune.uniform(0.3, 1.0),
    "width_ths": tune.uniform(0.3, 1.0),
    "add_margin": tune.uniform(0.0, 0.3),
    "contrast_ths": tune.uniform(0.05, 0.3),
    "adjust_contrast": tune.uniform(0.3, 0.8),
    "decoder": tune.choice(["greedy", "beamsearch"]),
    "beamWidth": tune.choice([3, 5, 7, 10]),
    "min_size": tune.choice([5, 10, 15, 20]),
}


# =============================================================================
# Config keys for results display
# =============================================================================

PADDLE_OCR_CONFIG_KEYS = [
    "use_doc_orientation_classify", "use_doc_unwarping", "textline_orientation",
    "text_det_thresh", "text_det_box_thresh", "text_det_unclip_ratio", "text_rec_score_thresh",
]

DOCTR_CONFIG_KEYS = [
    "assume_straight_pages", "straighten_pages", "preserve_aspect_ratio", "symmetric_pad",
    "disable_page_orientation", "disable_crop_orientation", "resolve_lines", "resolve_blocks",
    "paragraph_break",
]

EASYOCR_CONFIG_KEYS = [
    "text_threshold", "low_text", "link_threshold", "slope_ths", "ycenter_ths",
    "height_ths", "width_ths", "add_margin", "contrast_ths", "adjust_contrast",
    "decoder", "beamWidth", "min_size",
]