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

CT volumes are preprocessed with the same MONAI pipeline used during GLINT
training/zero-shot evaluation (orientation IPL, resample to 1.5/0.75/0.75 mm,
HU window [-1000, 1000] -> [-1, 1], 240x384x384), then regrouped into
``[num_frames=80, 3, 384, 384]`` and ImageNet-normalised.
"""

import numpy as np
import torch

# GLINT-CT pretraining preprocessing (glint/configs/chest_ct_monai.yaml)
CT_MONAI_CONFIG = {
    "orientation": "IPL",
    "target_spacing": (1.5, 0.75, 0.75),
    "roi_size": (240, 480, 480),
    "final_size": (240, 384, 384),
    "hu_range": (-1000.0, 1000.0),
    "normalize_range": (-1.0, 1.0),
}


def build_ct_transform(orientd=True, raw_spacing=None):
    from monai.transforms import (
        CenterSpatialCropd,
        Compose,
        EnsureChannelFirstd,
        EnsureTyped,
        Lambdad,
        LoadImaged,
        Orientationd,
        Resized,
        ScaleIntensityRanged,
        Spacingd,
        SpatialPadd,
    )

    cfg = CT_MONAI_CONFIG
    b_min, b_max = cfg["normalize_range"]

    def _patch_affine(img):
        if raw_spacing is not None and hasattr(img, "affine"):
            sx, sy, sz = (float(x) for x in raw_spacing)
            a = torch.eye(4, dtype=img.affine.dtype, device=img.affine.device)
            a[0, 0], a[1, 1], a[2, 2] = sx, sy, sz
            img.affine = a
        return img

    steps = [
        LoadImaged(keys=["image"], image_only=False),
        EnsureChannelFirstd(keys=["image"]),
        *([Orientationd(keys=["image"], axcodes=cfg["orientation"])] if orientd else []),
        *([Lambdad(keys=["image"], func=_patch_affine)] if raw_spacing is not None else []),
        ScaleIntensityRanged(
            keys=["image"],
            a_min=cfg["hu_range"][0],
            a_max=cfg["hu_range"][1],
            b_min=b_min,
            b_max=b_max,
            clip=True,
        ),
        Spacingd(keys=["image"], pixdim=cfg["target_spacing"], mode="bilinear"),
        CenterSpatialCropd(keys=["image"], roi_size=cfg["roi_size"]),
        SpatialPadd(keys=["image"], spatial_size=cfg["roi_size"], mode="constant", constant_values=b_min),
        Resized(keys=["image"], spatial_size=cfg["final_size"], mode="trilinear"),
        EnsureTyped(keys=["image"], dtype=torch.float32),
    ]
    return Compose(steps)


def normalize_pixel_values_videos(pixel_values_videos):
    """(D, H, W) in [-1, 1] -> (num_frames, 3, H, W) ImageNet-normalised."""
    s = pixel_values_videos.shape
    num_channels = 3
    pixel_values_videos = pixel_values_videos.reshape(-1, num_channels, s[-2], s[-1])

    pixel_values_videos = (pixel_values_videos + 1) / 2

    image_mean = torch.tensor([0.485, 0.456, 0.406], device=pixel_values_videos.device).view(1, 3, 1, 1)
    image_std = torch.tensor([0.229, 0.224, 0.225], device=pixel_values_videos.device).view(1, 3, 1, 1)
    pixel_values_videos = (pixel_values_videos - image_mean) / image_std

    return pixel_values_videos


def _input_geometry(image, orientd, raw_spacing):
    """Input voxel grid (shape, affine) and its 4x4 voxel mapping onto the preprocessed grid."""
    import nibabel as nib

    affine = np.asarray(image.meta["original_affine"], dtype=np.float64)
    shape = tuple(int(s) for s in image.meta["spatial_shape"])
    ref_affine = affine
    if raw_spacing is not None:  # the pipeline replaced the affine with diag(raw_spacing) after reorientation
        to_oriented = np.eye(4)
        if orientd:
            ornt = nib.orientations.ornt_transform(
                nib.orientations.io_orientation(affine),
                nib.orientations.axcodes2ornt(CT_MONAI_CONFIG["orientation"]),
            )
            to_oriented = np.linalg.inv(nib.orientations.inv_ornt_aff(ornt, shape))
        ref_affine = np.diag([*map(float, raw_spacing), 1.0]) @ to_oriented
    to_model = np.linalg.inv(np.asarray(image.affine, dtype=np.float64)) @ ref_affine
    return {"shape": shape, "affine": affine, "to_model": to_model}


def load_ct_volume(volume, orientd=True, raw_spacing=None, return_geometry=False):
    """Load and preprocess a CT volume.

    volume : path to a NIfTI (.nii/.nii.gz) file or a DICOM series directory, or a (240, 384, 384)
             tensor/ndarray in HU that is already on the preprocessed grid (IPL, 1.5/0.75/0.75 mm
             resampled, cropped/padded, resized); only the HU window is applied to arrays. Returns
             pixel_values of shape (1, num_frames, 3, H, W), plus the input grid if ``return_geometry``
             (``shape``, ``affine`` (None for array input) and ``to_model``, its voxel mapping onto the
             preprocessed grid), which ``similarity_scores_to_volume`` maps the scores back onto.
    """
    if isinstance(volume, str):
        transform = build_ct_transform(orientd=orientd, raw_spacing=raw_spacing)
        out = transform({"image": volume})
        image = out["image"].squeeze(0)  # (1, D, H, W) -> (D, H, W)
        geometry = _input_geometry(out["image"], orientd, raw_spacing)
    else:
        image = torch.as_tensor(np.asarray(volume), dtype=torch.float32)
        if tuple(image.shape) != tuple(CT_MONAI_CONFIG["final_size"]):
            raise ValueError(
                f"Array input must be on the preprocessed grid {CT_MONAI_CONFIG['final_size']}, got {tuple(image.shape)}"
            )
        (a_min, a_max), (b_min, b_max) = CT_MONAI_CONFIG["hu_range"], CT_MONAI_CONFIG["normalize_range"]
        image = (image.clamp(a_min, a_max) - a_min) / (a_max - a_min) * (b_max - b_min) + b_min
        geometry = {"shape": tuple(image.shape), "affine": None, "to_model": np.eye(4)}

    pixel_values_videos = normalize_pixel_values_videos(image).unsqueeze(0)  # (1, num_frames, 3, H, W)
    return (pixel_values_videos, geometry) if return_geometry else pixel_values_videos


def _resample_to_input(volume, to_model, shape, chunk=32):
    """Trilinearly sample a preprocessed-grid volume at every input voxel; 0 outside the preprocessing crop."""
    if np.allclose(to_model, np.eye(4)) and tuple(shape) == tuple(volume.shape):
        return volume
    device = volume.device
    M = torch.as_tensor(to_model, dtype=torch.float32, device=device)
    size = torch.tensor(volume.shape, dtype=torch.float32, device=device)
    ni, nj, nk = shape
    jk = torch.stack(torch.meshgrid(torch.arange(nj), torch.arange(nk), indexing="ij"), -1).reshape(-1, 2).to(device)
    out = torch.empty(tuple(shape), dtype=torch.float32, device=device)
    for i0 in range(0, ni, chunk):
        ii = torch.arange(i0, min(i0 + chunk, ni), device=device)
        ijk = torch.cat([ii.repeat_interleave(len(jk))[:, None], jk.repeat(len(ii), 1)], 1).float()
        src = ijk @ M[:3, :3].T + M[:3, 3]  # input voxel -> preprocessed voxel
        grid = (2 * src / (size - 1) - 1).flip(-1)  # (d, h, w) -> grid_sample's (x, y, z)
        out[i0 : i0 + len(ii)] = torch.nn.functional.grid_sample(
            volume[None, None], grid.view(1, -1, 1, 1, 3), mode="bilinear", padding_mode="zeros", align_corners=True
        ).view(len(ii), nj, nk)
    return out


def similarity_scores_to_volume(similarity_scores, model, geometry=None):
    """Per-token scores ``[1, 1, num_tokens]`` -> similarity map on the input CT grid.

    The token grid is upsampled (trilinear) to the 240x384x384 preprocessed grid and, with ``geometry`` from
    ``load_ct_volume``, resampled onto the input voxels so the map has the input's shape and affine (0 outside
    the 360 mm preprocessing crop). Without ``geometry`` the map stays on the preprocessed grid.
    """
    vc = model.config.vision_config
    patch_size = getattr(vc, "patch_size", None) or 16
    frame_length = vc.num_channels * (getattr(vc, "tubelet_size", None) or 2)
    grid = (vc.input_slices // frame_length, vc.img_size // patch_size, vc.img_size // patch_size)

    similarity_map = similarity_scores.float().reshape(grid)
    similarity_map = torch.nn.functional.interpolate(
        similarity_map[None, None], size=CT_MONAI_CONFIG["final_size"], mode="trilinear", align_corners=False
    )[0, 0]
    if geometry is None:
        return similarity_map
    return _resample_to_input(similarity_map, geometry["to_model"], geometry["shape"])


def save_similarity_map(similarity_map, affine, save_path):
    """Save the similarity map as NIfTI; with ``geometry["affine"]`` it overlays the input CT voxel for voxel."""
    import nibabel as nib

    sim = similarity_map.detach().float().cpu().numpy() if torch.is_tensor(similarity_map) else similarity_map
    nib.save(nib.Nifti1Image(sim.astype(np.float32), affine), save_path)


def _as_pixel_values(volume, raw_spacing):
    if torch.is_tensor(volume) and volume.dim() == 5:  # already preprocessed by load_ct_volume
        return volume
    return load_ct_volume(volume, raw_spacing=raw_spacing)


@torch.no_grad()
def model_inference(volume, text, model, tokenizer, raw_spacing=None):
    """Return (similarity_prob, similarity_scores) for one CT volume and one phrase.

    similarity_prob   : scalar sigmoid probability that the phrase applies.
    similarity_scores : per-token similarity in [0, 1] over the V-JEPA video tokens
                        (shape [1, 1, num_tokens]); ``similarity_scores_to_volume``
                        maps it onto the input CT grid.

    volume may also be the ``pixel_values`` returned by ``load_ct_volume``.
    """
    pixel_values = _as_pixel_values(volume, raw_spacing).to(model.device, dtype=model.dtype)

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

    output = model.compute_logits(pixel_values, [tokenized_text])
    similarity_prob = output["logits"].sigmoid()
    return similarity_prob, output.get("similarity_scores")


@torch.no_grad()
def model_inference_multiple_text(volume, text_list, model, tokenizer, raw_spacing=None):
    pixel_values = _as_pixel_values(volume, raw_spacing).to(model.device, dtype=model.dtype)
    enc = tokenizer(text_list, padding=True, truncation=True, return_tensors="pt").to(model.device)
    output = model.compute_logits(pixel_values, [enc])
    return output["logits"].sigmoid(), output.get("similarity_scores")