"""
Inference-time preprocessing สำหรับ Space
- Crop มือด้วย ianpan/bone-age-crop
- Histogram matching กับ ref_img.png (ถ้ามี)
- Resize (LongestMaxSize=512) + Pad ให้เป็น 512x512
ให้ผลลัพธ์เป็น grayscale uint8 512x512 (เหมือน val/test ตอนเทรน)
"""
import os
import cv2
import torch
import numpy as np
import albumentations as A
from skimage.exposure import match_histograms
from transformers import AutoModel

_CROP_MODEL = None
_REF_IMG = None
_DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

REF_IMG_PATH = os.path.join(os.path.dirname(__file__), "ref_img.png")

_resize_pad = A.Compose([
    A.LongestMaxSize(max_size=512, p=1),
    A.PadIfNeeded(512, 512, border_mode=cv2.BORDER_CONSTANT, value=0, p=1),
])


def _get_crop_model():
    global _CROP_MODEL
    if _CROP_MODEL is None:
        _CROP_MODEL = AutoModel.from_pretrained(
            "ianpan/bone-age-crop", trust_remote_code=True
        ).eval().to(_DEVICE)
    return _CROP_MODEL


def _get_ref_img():
    global _REF_IMG
    if _REF_IMG is None and os.path.exists(REF_IMG_PATH):
        _REF_IMG = cv2.imread(REF_IMG_PATH, 0)
    return _REF_IMG


def preprocess_image(image: np.ndarray) -> np.ndarray:
    """
    image: grayscale (H, W) uint8
    return: grayscale (512, 512) uint8
    """
    if image.ndim == 3:
        image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    image = image.astype(np.uint8)

    # --- 1) Crop ---
    try:
        crop_model = _get_crop_model()
        img_shape = torch.tensor([image.shape[:2]]).to(_DEVICE)
        x_crop = crop_model.preprocess(image)
        x_crop = torch.from_numpy(x_crop).unsqueeze(0).unsqueeze(0).float().to(_DEVICE)
        with torch.inference_mode():
            coords = crop_model(x_crop, img_shape)
        x_c, y_c, w_c, h_c = coords[0].cpu().numpy()
        cropped = image[int(y_c): int(y_c + h_c), int(x_c): int(x_c + w_c)]
        if cropped.size == 0:
            cropped = image
    except Exception as e:
        import traceback
        print("=" * 60)
        print("[crop] FAILED -> ใช้ภาพเต็มเฟรมแทน (ไม่ได้ crop!)")
        print(f"[crop] error: {repr(e)}")
        traceback.print_exc()
        print("=" * 60)
        cropped = image

    # --- 2) Histogram matching (ถ้ามี ref_img) ---
    ref = _get_ref_img()
    if ref is not None:
        matched = match_histograms(cropped, ref)
        matched = np.clip(matched, 0, 255).astype(np.uint8)
    else:
        print("[preprocess] ⚠️ ไม่พบ ref_img.png -> ข้าม histogram matching (ผลอาจเพี้ยน)")
        matched = cropped

    # --- 3) Resize + Pad ---
    out = _resize_pad(image=matched)["image"]
    return out