"""Volume model: voxel data + affine, with world<->voxel transforms and MPR reslicing."""
from __future__ import annotations
from dataclasses import dataclass, field
import numpy as np


@dataclass
class Volume:
    """A 3-D scalar image.

    data    : np.ndarray, shape (ni, nj, nk), index order (i,j,k).
    affine  : 4x4 voxel-index -> world (mm) transform (NIfTI/RAS convention).
    meta    : optional metadata (e.g. {"modality": "CT", "units": "HU"}).
    """
    data: np.ndarray
    affine: np.ndarray
    meta: dict = field(default_factory=dict)

    def __post_init__(self):
        self.data = np.asarray(self.data)
        self.affine = np.asarray(self.affine, dtype=np.float64)
        if self.affine.shape != (4, 4):
            raise ValueError("affine must be 4x4")
        if self.data.ndim != 3:
            raise ValueError("data must be 3-D")
        self._inv = np.linalg.inv(self.affine)

    # ---- geometry ----
    @property
    def shape(self):
        return self.data.shape

    @property
    def spacing(self) -> np.ndarray:
        """Voxel size (mm) along each index axis."""
        return np.linalg.norm(self.affine[:3, :3], axis=0)

    def voxel_to_world(self, ijk) -> np.ndarray:
        ijk = np.asarray(ijk, dtype=np.float64)
        h = np.append(ijk, 1.0)
        return (self.affine @ h)[:3]

    def world_to_voxel(self, xyz) -> np.ndarray:
        xyz = np.asarray(xyz, dtype=np.float64)
        h = np.append(xyz, 1.0)
        return (self._inv @ h)[:3]

    def in_bounds(self, ijk) -> bool:
        i, j, k = ijk
        ni, nj, nk = self.data.shape
        return 0 <= i < ni and 0 <= j < nj and 0 <= k < nk

    # ---- sampling ----
    def sample_world(self, xyz, order=1) -> float:
        """Sample intensity at a world (mm) point (nearest/linear)."""
        v = self.world_to_voxel(xyz)
        if order == 0:
            ijk = np.round(v).astype(int)
            if not self.in_bounds(ijk):
                return 0.0
            return float(self.data[tuple(ijk)])
        # trilinear
        return float(_trilinear(self.data, v))

    # ---- MPR (axis-aligned to the data grid) ----
    def slice(self, axis: int, index: int) -> np.ndarray:
        """Return a 2-D slice taken perpendicular to a data axis (0,1,2)."""
        index = int(np.clip(index, 0, self.data.shape[axis] - 1))
        return np.take(self.data, index, axis=axis)

    def window(self, lo=None, hi=None) -> "Volume":
        """Return a copy windowed/normalised to uint8 for display."""
        d = self.data.astype(np.float32)
        if hi is None:
            hi = np.percentile(d[d > d.min()], 99.5)
        if lo is None:
            lo = float(d.min())
        d = np.clip((d - lo) / max(hi - lo, 1e-6), 0, 1)
        return Volume((d * 255).astype(np.uint8), self.affine.copy())


def _trilinear(data, v):
    ni, nj, nk = data.shape
    x, y, z = v
    x0, y0, z0 = int(np.floor(x)), int(np.floor(y)), int(np.floor(z))
    if x0 < 0 or y0 < 0 or z0 < 0 or x0 >= ni - 1 or y0 >= nj - 1 or z0 >= nk - 1:
        xi, yi, zi = int(round(x)), int(round(y)), int(round(z))
        if 0 <= xi < ni and 0 <= yi < nj and 0 <= zi < nk:
            return float(data[xi, yi, zi])
        return 0.0
    xd, yd, zd = x - x0, y - y0, z - z0
    c = data[x0:x0 + 2, y0:y0 + 2, z0:z0 + 2].astype(np.float64)
    c00 = c[0, 0, 0] * (1 - xd) + c[1, 0, 0] * xd
    c01 = c[0, 0, 1] * (1 - xd) + c[1, 0, 1] * xd
    c10 = c[0, 1, 0] * (1 - xd) + c[1, 1, 0] * xd
    c11 = c[0, 1, 1] * (1 - xd) + c[1, 1, 1] * xd
    c0 = c00 * (1 - yd) + c10 * yd
    c1 = c01 * (1 - yd) + c11 * yd
    return float(c0 * (1 - zd) + c1 * zd)
