from __future__ import annotations

from pathlib import Path

import nibabel as nib
import numpy as np
import pandas as pd
import torch
import torch.nn.functional as F
from torch.utils.data import Dataset


def normalize_pet(volume: np.ndarray, eps: float = 1e-6) -> np.ndarray:
    volume = np.asarray(volume, dtype=np.float32)
    finite = np.isfinite(volume)
    if not finite.any():
        return np.zeros_like(volume, dtype=np.float32)
    lo, hi = np.percentile(volume[finite], [0.5, 99.5])
    volume = np.clip(volume, lo, hi)
    volume = (volume - lo) / max(float(hi - lo), eps)
    return volume.astype(np.float32, copy=False)


def resize_volume(volume: np.ndarray, output_size: tuple[int, int, int]) -> torch.Tensor:
    tensor = torch.from_numpy(volume)[None, None]
    tensor = F.interpolate(tensor, size=output_size, mode="trilinear", align_corners=False)
    return tensor[0]


def load_suvr_vector(csv_path: str | Path, include_background: bool = False) -> tuple[list[str], torch.Tensor]:
    df = pd.read_csv(csv_path)
    if not include_background:
        df = df[df["label_name"] != "Background"].copy()
    labels = df["label_name"].astype(str).tolist()
    values = torch.tensor(df["mean_scalar"].astype(float).to_numpy(), dtype=torch.float32)
    return labels, values


def suvr_to_text(labels: list[str], values: torch.Tensor, top_k: int = 8) -> str:
    pairs = sorted(zip(labels, values.tolist()), key=lambda x: x[1], reverse=True)
    high = ", ".join(f"{name} {value:.3f}" for name, value in pairs[:top_k])
    low = ", ".join(f"{name} {value:.3f}" for name, value in pairs[-top_k:])
    mean_value = float(values.mean())
    return f"FDG-PET regional SUVR summary. Mean SUVR {mean_value:.3f}. Highest regions: {high}. Lowest regions: {low}."


class PETSUVRDataset(Dataset):
    def __init__(
        self,
        manifest_path: str | Path,
        output_size: tuple[int, int, int] = (96, 96, 96),
        include_background: bool = False,
    ) -> None:
        self.manifest = pd.read_csv(manifest_path)
        self.output_size = output_size
        self.include_background = include_background

    def __len__(self) -> int:
        return len(self.manifest)

    def __getitem__(self, index: int) -> dict[str, object]:
        row = self.manifest.iloc[index]
        volume = nib.load(str(row["pet_path"])).get_fdata(dtype=np.float32)
        volume = normalize_pet(volume)
        image = resize_volume(volume, self.output_size)
        labels, suvr = load_suvr_vector(row["suvr_csv_path"], self.include_background)
        text = suvr_to_text(labels, suvr)
        return {
            "sample_id": row["sample_id"],
            "image": image,
            "suvr": suvr,
            "region_labels": labels,
            "text": text,
        }


def collate_pet_suvr(batch: list[dict[str, object]]) -> dict[str, object]:
    return {
        "sample_id": [item["sample_id"] for item in batch],
        "image": torch.stack([item["image"] for item in batch]),
        "suvr": torch.stack([item["suvr"] for item in batch]),
        "text": [item["text"] for item in batch],
        "region_labels": batch[0]["region_labels"],
    }