"""Image IO: load NIfTI and DICOM series into a Volume (numpy + affine)."""
from __future__ import annotations
import os
import glob
import numpy as np
from .volume import Volume


def load(path: str) -> Volume:
    """Load a NIfTI file (.nii/.nii.gz) or a directory of DICOM slices."""
    if os.path.isdir(path):
        return load_dicom_series(path)
    low = path.lower()
    if low.endswith(".nii") or low.endswith(".nii.gz"):
        return load_nifti(path)
    raise ValueError(f"Unsupported input: {path}")


def load_nifti(path: str) -> Volume:
    import nibabel as nib
    img = nib.load(path)
    data = np.asanyarray(img.dataobj).astype(np.float32)
    if data.ndim == 4:               # take first volume of a 4-D series
        data = data[..., 0]
    return Volume(data, img.affine)


def load_dicom_series(folder: str) -> Volume:
    """Assemble a single 3-D volume from a folder of DICOM slices.

    Builds an affine from ImageOrientationPatient / ImagePositionPatient and
    pixel spacing, sorting slices along the slice-normal.
    """
    import pydicom
    files = []
    for f in glob.glob(os.path.join(folder, "*")):
        try:
            files.append(pydicom.dcmread(f))
        except Exception:
            continue
    if not files:
        raise ValueError(f"No readable DICOM files in {folder}")
    files = [d for d in files if hasattr(d, "ImagePositionPatient")]
    iop = np.array(files[0].ImageOrientationPatient, dtype=np.float64)
    row, col = iop[:3], iop[3:]
    normal = np.cross(row, col)
    files.sort(key=lambda d: float(np.dot(np.array(d.ImagePositionPatient, float), normal)))

    cols, rows = int(files[0].Columns), int(files[0].Rows)
    ps = [float(x) for x in files[0].PixelSpacing]      # [row_spacing, col_spacing]
    if len(files) > 1:
        p0 = np.array(files[0].ImagePositionPatient, float)
        p1 = np.array(files[1].ImagePositionPatient, float)
        dz = float(np.linalg.norm(p1 - p0))
    else:
        dz = float(getattr(files[0], "SliceThickness", 1.0))

    vol = np.zeros((cols, rows, len(files)), dtype=np.float32)
    for k, d in enumerate(files):
        arr = d.pixel_array.astype(np.float32)
        slope = float(getattr(d, "RescaleSlope", 1.0))
        intercept = float(getattr(d, "RescaleIntercept", 0.0))
        vol[:, :, k] = (arr * slope + intercept).T   # (col,row)->(i,j)

    origin = np.array(files[0].ImagePositionPatient, float)
    affine = np.eye(4)
    affine[:3, 0] = row * ps[1]      # i axis (columns)
    affine[:3, 1] = col * ps[0]      # j axis (rows)
    affine[:3, 2] = normal * dz      # k axis (slices)
    affine[:3, 3] = origin
    modality = str(getattr(files[0], "Modality", "") or "")
    meta = {"modality": modality}
    if modality.upper() == "CT":
        meta["units"] = "HU"        # RescaleSlope/Intercept applied above
    return Volume(vol, affine, meta=meta)
