"""
DICOM Series Inventory Module.

Enumerates all CT series within a patient DICOM folder, classifies each series
by slice thickness, excludes non-image series (localizers, dose reports, protocol
sheets), and returns a structured inventory suitable for multi-thickness research.

Typical usage:
    >>> from cardiac_shared.io.dicom_series_inventory import DicomSeriesInventory
    >>> inv = DicomSeriesInventory()
    >>> result = inv.scan_patient_folder("/path/to/patient/dicom")
    >>> for series in result.ct_series:
    ...     print(f"{series.thickness_mm}mm  {series.description}  {series.n_images} slices")

    >>> # Convenience function
    >>> from cardiac_shared.io.dicom_series_inventory import scan_patient_folder
    >>> result = scan_patient_folder("/path/to/patient/dicom")
    >>> print(result.summary())

Added in v0.9.2.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import Dict, List, Optional, Union

logger = logging.getLogger(__name__)

try:
    import pydicom
    HAS_PYDICOM = True
except ImportError:
    HAS_PYDICOM = False


# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

# Series description keywords that identify NON-image series to exclude.
# Lowercase; matched as substrings.
_EXCLUDE_DESC_KEYWORDS = [
    "scout", "topogram", "surview", "localizer", "loc ",
    "dose report", "dose record", "dose_report",
    "patient protocol", "protocol",
    "ecg report", "ecg_report",
    "smart prep screen", "screen save",
    "bolus track", "smartprep",
    "report", "worksheet",
]

# Image types that mark a localizer frame
_LOCALIZER_IMAGE_TYPES = {"LOCALIZER", "LOCALIZER\\PRIMARY", "SECONDARY"}

# Maximum realistic slice thickness for CT (mm). Above this → localizer / scout.
_MAX_VALID_THICKNESS_MM = 50.0

# Minimum number of images for a series to be considered a real CT volume.
# A single localizer slice sometimes has SliceThickness=5 and 1 image.
_MIN_IMAGES_FOR_VOLUME = 4

# Standard clinical CT slice thicknesses (mm) for normalisation.
_STANDARD_THICKNESSES = [0.5, 0.625, 0.75, 1.0, 1.25, 1.5, 2.0, 2.5, 3.0, 5.0, 10.0]
_STANDARD_TOLERANCE = 0.15  # mm


class SeriesRole(str, Enum):
    """Functional role of a DICOM series."""
    THIN_SLICE   = "thin_slice"    # <= 1.5 mm
    MEDIUM_SLICE = "medium_slice"  # 1.5 – 2.5 mm
    THICK_SLICE  = "thick_slice"   # > 2.5 mm
    LOCALIZER    = "localizer"
    NON_IMAGE    = "non_image"     # dose/protocol/report
    UNKNOWN      = "unknown"


@dataclass
class SeriesInfo:
    """
    Metadata for a single DICOM series within a patient folder.

    Attributes:
        series_uid:       SeriesInstanceUID (unique per series)
        thickness_mm:     Detected slice thickness in mm (None if unavailable)
        description:      SeriesDescription tag value
        image_type:       ImageType tag values (list of strings)
        n_images:         Number of DICOM files belonging to this series
        role:             Functional classification (thin/thick/localizer/…)
        modality:         DICOM Modality tag (e.g. "CT")
        series_number:    DICOM SeriesNumber tag
    """
    series_uid:    str
    thickness_mm:  Optional[float]
    description:   str
    image_type:    List[str]
    n_images:      int
    role:          SeriesRole
    modality:      str        = "CT"
    series_number: Optional[int] = None

    @property
    def is_usable_ct(self) -> bool:
        """True if this is a real CT volume (not localizer / report)."""
        return self.role in (SeriesRole.THIN_SLICE, SeriesRole.MEDIUM_SLICE,
                             SeriesRole.THICK_SLICE)

    @property
    def thickness_label(self) -> str:
        """Human-readable thickness label, e.g. '1.5mm' or 'unknown'."""
        if self.thickness_mm is None:
            return "unknown"
        return f"{self.thickness_mm:.3g}mm"

    def to_dict(self) -> Dict:
        return {
            "series_uid":   self.series_uid,
            "thickness_mm": self.thickness_mm,
            "description":  self.description,
            "image_type":   self.image_type,
            "n_images":     self.n_images,
            "role":         self.role.value,
            "modality":     self.modality,
            "series_number": self.series_number,
        }


@dataclass
class PatientSeriesInventory:
    """
    Complete series inventory for one patient folder.

    Attributes:
        folder:       Scanned folder path
        all_series:   All discovered series (including excluded ones)
        n_files_read: Number of DICOM files successfully parsed
        n_files_total: Total files in folder
    """
    folder:         Path
    all_series:     List[SeriesInfo]
    n_files_read:   int = 0
    n_files_total:  int = 0

    @property
    def ct_series(self) -> List[SeriesInfo]:
        """Usable CT series only (excludes localizers and non-image series)."""
        return [s for s in self.all_series if s.is_usable_ct]

    @property
    def thicknesses_mm(self) -> List[float]:
        """Sorted list of distinct slice thicknesses among usable CT series."""
        seen = set()
        out = []
        for s in self.ct_series:
            if s.thickness_mm is not None and s.thickness_mm not in seen:
                seen.add(s.thickness_mm)
                out.append(s.thickness_mm)
        return sorted(out)

    @property
    def has_thin(self) -> bool:
        """Has at least one thin-slice series (<= 1.5 mm)."""
        return any(s.role == SeriesRole.THIN_SLICE for s in self.ct_series)

    @property
    def has_thick(self) -> bool:
        """Has at least one thick-slice series (>= 4 mm)."""
        return any(
            s.thickness_mm is not None and s.thickness_mm >= 4.0
            for s in self.ct_series
        )

    @property
    def is_paired(self) -> bool:
        """Has both thin (<= 1.5 mm) and thick (>= 4 mm) series."""
        return self.has_thin and self.has_thick

    @property
    def thickness_combo(self) -> str:
        """
        String key describing the thickness combination, e.g. '1.5mm+5mm'.
        Suitable for grouping / counting patients by combination.
        """
        t = self.thicknesses_mm
        if not t:
            return "unknown"
        return "+".join(f"{v:.3g}mm" for v in t)

    def series_by_thickness(self, target_mm: float,
                             tolerance: float = 0.3) -> List[SeriesInfo]:
        """Return usable CT series whose thickness is within tolerance of target."""
        return [
            s for s in self.ct_series
            if s.thickness_mm is not None
            and abs(s.thickness_mm - target_mm) <= tolerance
        ]

    def summary(self) -> str:
        """Single-line human-readable summary."""
        ct = self.ct_series
        return (
            f"{self.folder.name}: "
            f"{len(ct)} CT series  "
            f"thicknesses={self.thicknesses_mm}  "
            f"paired={self.is_paired}  "
            f"({self.n_files_read}/{self.n_files_total} files read)"
        )

    def to_dict(self) -> Dict:
        return {
            "folder":        str(self.folder),
            "n_files_read":  self.n_files_read,
            "n_files_total": self.n_files_total,
            "all_series":    [s.to_dict() for s in self.all_series],
            "ct_series_count": len(self.ct_series),
            "thicknesses_mm": self.thicknesses_mm,
            "thickness_combo": self.thickness_combo,
            "is_paired":     self.is_paired,
        }


# ---------------------------------------------------------------------------
# Core scanner
# ---------------------------------------------------------------------------

class DicomSeriesInventory:
    """
    Scans a DICOM folder (flat or nested one level) and returns a structured
    inventory of all CT series, with each series classified by slice thickness
    and role (thin/medium/thick/localizer/non-image).

    Example:
        >>> inv = DicomSeriesInventory()
        >>> result = inv.scan_patient_folder("/mnt/d/raw/internal/chen/female_new/dicom_10000103")
        >>> print(result.summary())
        >>> for s in result.ct_series:
        ...     print(s.thickness_mm, s.description)

        >>> # Scan entire cohort
        >>> results = inv.scan_cohort("/mnt/d/raw/internal/chen/female_new")
        >>> paired = [r for r in results if r.is_paired]
        >>> print(f"Paired patients: {len(paired)}/{len(results)}")
    """

    def __init__(
        self,
        min_images: int = _MIN_IMAGES_FOR_VOLUME,
        max_files_per_series_sample: int = 5,
        recurse_depth: int = 1,
    ):
        """
        Args:
            min_images:                  Minimum images for a series to count as a volume.
            max_files_per_series_sample: Max files to read per series when sampling
                                         (reads ALL files if 0).
            recurse_depth:               How deep to look for DICOM files (0 = flat only,
                                         1 = one subdirectory level, -1 = unlimited).
        """
        if not HAS_PYDICOM:
            raise ImportError(
                "pydicom is required. Install with: pip install pydicom"
            )
        self.min_images = min_images
        self.max_files_per_series_sample = max_files_per_series_sample
        self.recurse_depth = recurse_depth

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def scan_patient_folder(
        self, folder: Union[str, Path]
    ) -> PatientSeriesInventory:
        """
        Scan one patient's DICOM folder and return a full series inventory.

        Args:
            folder: Path to the patient's DICOM directory.

        Returns:
            PatientSeriesInventory with all discovered series classified.
        """
        folder = Path(folder)
        if not folder.exists():
            logger.warning(f"Folder does not exist: {folder}")
            return PatientSeriesInventory(folder=folder, all_series=[])

        dicom_files = self._collect_dicom_files(folder)
        n_total = len(dicom_files)

        # Group files by SeriesInstanceUID, reading header only
        series_map: Dict[str, Dict] = {}
        n_read = 0

        for fpath in dicom_files:
            try:
                ds = pydicom.dcmread(str(fpath), stop_before_pixels=True)
                uid = str(getattr(ds, "SeriesInstanceUID", f"_no_uid_{fpath.name}"))

                if uid not in series_map:
                    series_map[uid] = {
                        "thickness_mm":  self._get_thickness(ds),
                        "description":   str(getattr(ds, "SeriesDescription", "")),
                        "image_type":    [str(v) for v in getattr(ds, "ImageType", [])],
                        "modality":      str(getattr(ds, "Modality", "CT")),
                        "series_number": self._safe_int(getattr(ds, "SeriesNumber", None)),
                        "count":         0,
                    }
                series_map[uid]["count"] += 1
                n_read += 1

            except Exception:
                continue

        # Build SeriesInfo objects
        all_series: List[SeriesInfo] = []
        for uid, info in series_map.items():
            role = self._classify_role(
                info["description"],
                info["image_type"],
                info["thickness_mm"],
                info["count"],
            )
            all_series.append(SeriesInfo(
                series_uid=uid,
                thickness_mm=info["thickness_mm"],
                description=info["description"],
                image_type=info["image_type"],
                n_images=info["count"],
                role=role,
                modality=info["modality"],
                series_number=info["series_number"],
            ))

        # Sort: CT series by thickness first, then non-CT
        all_series.sort(key=lambda s: (
            0 if s.is_usable_ct else 1,
            s.thickness_mm or 999,
        ))

        return PatientSeriesInventory(
            folder=folder,
            all_series=all_series,
            n_files_read=n_read,
            n_files_total=n_total,
        )

    def scan_cohort(
        self,
        cohort_folder: Union[str, Path],
        verbose: bool = False,
    ) -> List[PatientSeriesInventory]:
        """
        Scan all patient subfolders within a cohort directory.

        Args:
            cohort_folder: Directory whose immediate subdirectories are patient folders.
            verbose:       If True, log one line per patient.

        Returns:
            List of PatientSeriesInventory, one per patient subfolder.
        """
        cohort_folder = Path(cohort_folder)
        patient_dirs = sorted(
            d for d in cohort_folder.iterdir() if d.is_dir()
        )
        results = []
        for i, pdir in enumerate(patient_dirs):
            inv = self.scan_patient_folder(pdir)
            results.append(inv)
            if verbose and (i + 1) % 50 == 0:
                logger.info(f"  Scanned {i+1}/{len(patient_dirs)} patients…")
        return results

    @staticmethod
    def cohort_summary(inventories: List[PatientSeriesInventory]) -> Dict:
        """
        Aggregate statistics across a cohort.

        Returns a dict with thickness_combo counts, paired counts, etc.
        """
        from collections import Counter
        combo_counts: Counter = Counter()
        n_paired = 0
        n_has_thin = 0
        n_has_thick = 0
        all_thicknesses: Counter = Counter()

        for inv in inventories:
            combo_counts[inv.thickness_combo] += 1
            if inv.is_paired:
                n_paired += 1
            if inv.has_thin:
                n_has_thin += 1
            if inv.has_thick:
                n_has_thick += 1
            for t in inv.thicknesses_mm:
                all_thicknesses[t] += 1

        return {
            "n_patients":         len(inventories),
            "n_paired":           n_paired,
            "n_has_thin":         n_has_thin,
            "n_has_thick":        n_has_thick,
            "thickness_combos":   dict(combo_counts.most_common()),
            "thickness_counts":   dict(sorted(all_thicknesses.items())),
        }

    # ------------------------------------------------------------------
    # Private helpers
    # ------------------------------------------------------------------

    def _collect_dicom_files(self, folder: Path) -> List[Path]:
        """Collect DICOM file paths from folder (and optionally subdirs)."""
        candidates: List[Path] = []

        def _add_from_dir(d: Path, depth: int) -> None:
            for item in d.iterdir():
                if item.is_file():
                    candidates.append(item)
                elif item.is_dir() and depth > 0:
                    _add_from_dir(item, depth - 1)

        try:
            _add_from_dir(folder, self.recurse_depth)
        except PermissionError as e:
            logger.warning(f"Permission error scanning {folder}: {e}")

        return candidates

    @staticmethod
    def _get_thickness(ds) -> Optional[float]:
        """Extract slice thickness from a DICOM dataset."""
        for tag in ("SliceThickness", "SpacingBetweenSlices"):
            val = getattr(ds, tag, None)
            if val is not None:
                try:
                    t = float(val)
                    if 0.1 <= t <= _MAX_VALID_THICKNESS_MM:
                        return round(t, 4)
                except (ValueError, TypeError):
                    pass
        return None

    @staticmethod
    def _classify_role(
        description: str,
        image_type: List[str],
        thickness_mm: Optional[float],
        n_images: int,
    ) -> SeriesRole:
        """Classify a series into a SeriesRole."""
        desc_lower = description.lower()

        # --- Explicit localizer image type ---
        img_type_set = {t.upper() for t in image_type}
        if img_type_set & _LOCALIZER_IMAGE_TYPES:
            return SeriesRole.LOCALIZER

        # --- Thickness > max → scout ---
        if thickness_mm is not None and thickness_mm > _MAX_VALID_THICKNESS_MM:
            return SeriesRole.LOCALIZER

        # --- Description keyword exclusions ---
        if any(kw in desc_lower for kw in _EXCLUDE_DESC_KEYWORDS):
            return SeriesRole.NON_IMAGE

        # --- Too few images → probably not a volume ---
        if n_images < _MIN_IMAGES_FOR_VOLUME:
            return SeriesRole.NON_IMAGE

        # --- No thickness info ---
        if thickness_mm is None:
            return SeriesRole.UNKNOWN

        # --- Classify by thickness ---
        if thickness_mm <= 1.5:
            return SeriesRole.THIN_SLICE
        elif thickness_mm <= 2.5:
            return SeriesRole.MEDIUM_SLICE
        else:
            return SeriesRole.THICK_SLICE

    @staticmethod
    def _safe_int(val) -> Optional[int]:
        try:
            return int(val)
        except (TypeError, ValueError):
            return None


# ---------------------------------------------------------------------------
# Convenience function
# ---------------------------------------------------------------------------

def scan_patient_folder(
    folder: Union[str, Path],
    min_images: int = _MIN_IMAGES_FOR_VOLUME,
    recurse_depth: int = 1,
) -> PatientSeriesInventory:
    """
    Convenience wrapper: scan a single patient DICOM folder.

    Args:
        folder:       Path to patient DICOM folder.
        min_images:   Minimum images for a real CT volume.
        recurse_depth: 0=flat only, 1=one subdirectory level.

    Returns:
        PatientSeriesInventory

    Example:
        >>> result = scan_patient_folder("/mnt/d/raw/internal/chen/female_new/dicom_10000103")
        >>> print(result.summary())
        >>> print(result.thickness_combo)   # e.g. "1.5mm+5mm"
        >>> print(result.is_paired)          # True
    """
    inv = DicomSeriesInventory(
        min_images=min_images,
        recurse_depth=recurse_depth,
    )
    return inv.scan_patient_folder(folder)


# ---------------------------------------------------------------------------
# Exports
# ---------------------------------------------------------------------------

__all__ = [
    "SeriesRole",
    "SeriesInfo",
    "PatientSeriesInventory",
    "DicomSeriesInventory",
    "scan_patient_folder",
]
