import numpy as np
import cv2

TARGET_SIZE = (224, 224)

def ensure_2d_gray(img):
    if img.ndim == 3:
        if img.shape[2] == 1:
            return img[:, :, 0]
        elif img.shape[2] == 3:
            return (0.299*img[:,:,0] + 0.587*img[:,:,1] + 0.114*img[:,:,2])
    return img

def resize_image(img, size=TARGET_SIZE):
    img_u8 = np.clip(img / img.max() * 255, 0, 255).astype(np.uint8) if img.max() > 0 else img.astype(np.uint8)
    resized = cv2.resize(img_u8, size, interpolation=cv2.INTER_AREA)
    return resized.astype(np.float32)


def normalize_image(img):
    mn, mx = img.min(), img.max()
    return (img - mn) / (mx - mn) if mx > mn else img


def enhance_xray_full(img,
                      gamma=0.8,
                      clahe_clip=1.5,
                      clahe_tile=(8,8),
                      unsharp_strength=0.4,
                      laplacian_weight=0.1,
                      sobel_weight=0.0):

    img = np.power(np.clip(img, 1e-7, 1.0), gamma).astype(np.float32)

    img_uint8 = (img * 255).astype(np.uint8)
    clahe = cv2.createCLAHE(clipLimit=clahe_clip, tileGridSize=clahe_tile)
    img_uint8 = clahe.apply(img_uint8)

    blurred = cv2.GaussianBlur(img_uint8, (5, 5), 0)
    img_uint8 = cv2.addWeighted(img_uint8, 1 + unsharp_strength,
                                blurred, -unsharp_strength, 0)

    lap = cv2.Laplacian(img_uint8, cv2.CV_64F)
    lap = np.uint8(np.clip(np.absolute(lap), 0, 255))
    img_uint8 = cv2.addWeighted(img_uint8, 1.0, lap, laplacian_weight, 0)

    sobel_x = cv2.Sobel(img_uint8, cv2.CV_64F, 1, 0, ksize=3)
    sobel_y = cv2.Sobel(img_uint8, cv2.CV_64F, 0, 1, ksize=3)
    mag = np.sqrt(sobel_x**2 + sobel_y**2)
    mag = np.uint8(np.clip(mag, 0, 255))
    img_uint8 = cv2.addWeighted(img_uint8, 1.0, mag, sobel_weight, 0)

    return img_uint8.astype(np.float32) / 255.0


def preprocess(img_array):
    if img_array.ndim == 3:
        img = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY).astype(np.float32)
    else:
        img = img_array.astype(np.float32)

    img = ensure_2d_gray(img)
    img = resize_image(img)
    img = normalize_image(img)
    img = enhance_xray_full(img)

    img_rgb = np.repeat(img[..., np.newaxis], 3, axis=-1)
    return np.expand_dims(img_rgb, axis=0)