"""Inference helpers for the GLINT-CXR vision-language model.

These utilities mirror the zero-shot grounding / phrase-similarity pipeline used
to evaluate GLINT on chest X-rays.  ``model.compute_logits`` already removes the
CLS token from the returned similarity map (the model was trained with
``use_cls_token=False``), so the helpers here only reshape and interpolate.
"""

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image


@torch.no_grad()
def model_inference(image, text, model, image_processor, tokenizer):
    """Return (similarity_prob, similarity_map) for one image and one phrase.

    similarity_prob : scalar sigmoid probability that the phrase applies to the image.
    similarity_map  : per-pixel similarity map in [0, 1] (H x W), upsampled to the
                      original image resolution.
    """
    image = load_image(image)
    (width, height) = image.size
    image_size = (height, width)

    image_processor_outputs = image_processor(image)
    processed_image = torch.FloatTensor(np.array(image_processor_outputs["pixel_values"])).to(model.device)

    tokenized_text = tokenizer(
        text,
        padding=True,
        truncation=True,
        return_tensors="pt",
    ).to(model.device)

    output = model.compute_logits(processed_image, [tokenized_text])

    logits = output["logits"]
    similarity_prob = logits.sigmoid()

    similarity_scores = output["similarity_scores"]
    similarity_scores = similarity_scores.view(-1)
    similarity_scores = interpolate_similarity_scores(similarity_scores, image_size, image_processor)
    similarity_map = similarity_scores[0]

    return similarity_prob, similarity_map


def save_similarity_map(image, similarity_map, save_path, threshold=0.0, beta=0.5):
    """Save a JET heat-map of ``similarity_map`` overlaid on ``image`` (path or PIL.Image) as a PNG."""
    import cv2

    img_rgb = np.array(load_image(image))
    sim = similarity_map.detach().float().cpu().numpy() if torch.is_tensor(similarity_map) else similarity_map
    sim = cv2.resize(sim.astype(np.float32), (img_rgb.shape[1], img_rgb.shape[0]), interpolation=cv2.INTER_LINEAR)
    sim = (np.clip(sim, threshold, 1.0) - threshold) / (1.0 - threshold + 1e-8)
    sim_u8 = (sim * 255).astype(np.uint8)

    heatmap = cv2.cvtColor(cv2.applyColorMap(sim_u8, cv2.COLORMAP_JET), cv2.COLOR_BGR2RGB)
    Image.fromarray(cv2.addWeighted(img_rgb, 1 - beta, heatmap, beta, 0)).save(save_path)


@torch.no_grad()
def model_inference_multiple_text(image, text_list, model, image_processor, tokenizer):
    probs, similarity_maps = [], []
    for text in text_list:
        prob, similarity_map = model_inference(image, text, model, image_processor, tokenizer)
        probs.append(prob)
        similarity_maps.append(similarity_map)

    return torch.stack(probs), torch.stack(similarity_maps)


def interpolate_similarity_scores(similarity_scores, origin_size, image_processor):
    (height, width) = origin_size
    patch_size = int(similarity_scores.shape[-1] ** 0.5)
    scores = similarity_scores.view(1, 1, patch_size, patch_size)

    interpolated_scores = F.interpolate(
        scores,
        size=(height, width),
        mode="bilinear",
        align_corners=False,
    )
    interpolated_scores = interpolated_scores.squeeze(1)
    return interpolated_scores


def dicom_to_pil_image(input_file_path):
    """Extract a grayscale image from a DICOM file and return it as a PIL.Image."""
    import cv2
    import pydicom

    dcm_file = pydicom.dcmread(input_file_path)
    raw_image = dcm_file.pixel_array

    if raw_image.ndim != 2 or raw_image.size == 0:
        raise ValueError(f"Expecting a non-empty single-channel DICOM image, got shape {raw_image.shape}")

    raw_image = raw_image - raw_image.min()
    rescaled_image = (raw_image / raw_image.max() * 255).astype(np.uint8)

    if dcm_file.PhotometricInterpretation == "MONOCHROME1":
        rescaled_image = cv2.bitwise_not(rescaled_image)

    return Image.fromarray(cv2.equalizeHist(rescaled_image))


def load_image(image):
    """Load an image from a path (.png/.jpg/.jpeg/.dcm) or a PIL.Image object."""
    if isinstance(image, str):
        if image.lower().endswith(".dcm"):
            image = dicom_to_pil_image(image)
        elif image.lower().endswith((".png", ".jpg", ".jpeg")):
            image = Image.open(image)
        else:
            raise ValueError(f"Invalid image type: {image}")
    elif not isinstance(image, Image.Image):
        raise ValueError(f"Invalid image type: {type(image)}")

    if image.size[0] <= 0 or image.size[1] <= 0:
        raise ValueError(f"Image has invalid size (W: {image.size[0]}, H: {image.size[1]}).")

    # per-image min-max stretch to 8 bit, as in the zero-shot evaluation (also handles 16-bit inputs)
    array = np.asarray(image, dtype=np.float32)
    array = (array - array.min()) / max(float(array.max() - array.min()), 1e-8) * 255.0
    return Image.fromarray(np.rint(array).astype(np.uint8)).convert("RGB")