"""
External Dataset Specifications for Cardiac Imaging Research.

Provides standardized access patterns, directory structures, and metadata
for commonly used public cardiac CT datasets. This module helps researchers
quickly understand and work with these datasets without reading lengthy
documentation.

Supported datasets:
- NLST (National Lung Screening Trial) — dual-thickness CT for CAC research
- COCA (Stanford Coronary Calcium) — CAC scoring gold standard
- TotalSegmentator — multi-center CT with 117 anatomical structures

Usage:
    >>> from cardiac_shared.data import get_dataset_info, list_datasets
    >>> info = get_dataset_info('nlst')
    >>> print(info)
    >>> print(info.access_url)
    >>> print(info.typical_structure)
"""

from dataclasses import dataclass, field
from typing import Dict, List, Optional


@dataclass
class DatasetInfo:
    """Standardized metadata for a public cardiac imaging dataset."""
    name: str
    short_name: str
    description: str
    total_cases: int
    modality: str
    access_url: str
    license: str
    citation: str

    # Structure and format
    file_format: str = "DICOM"
    typical_structure: str = ""
    slice_thickness_range: str = ""
    has_ground_truth: bool = False
    ground_truth_format: str = ""

    # Research relevance
    use_cases: List[str] = field(default_factory=list)
    notes: List[str] = field(default_factory=list)

    # Paired thickness support (for RESCUE research)
    supports_paired_thickness: bool = False
    thin_thickness_mm: Optional[float] = None
    thick_thickness_mm: Optional[float] = None
    paired_selection_criteria: str = ""

    def __str__(self) -> str:
        lines = [
            f"=== {self.name} ({self.short_name}) ===",
            f"  Cases: {self.total_cases}",
            f"  Modality: {self.modality}",
            f"  Format: {self.file_format}",
            f"  License: {self.license}",
            f"  Access: {self.access_url}",
        ]
        if self.slice_thickness_range:
            lines.append(f"  Thickness range: {self.slice_thickness_range}")
        if self.supports_paired_thickness:
            lines.append(
                f"  Paired thickness: {self.thin_thickness_mm}mm (thin) / "
                f"{self.thick_thickness_mm}mm (thick)"
            )
        if self.has_ground_truth:
            lines.append(f"  Ground truth: {self.ground_truth_format}")
        if self.use_cases:
            lines.append("  Use cases:")
            for uc in self.use_cases:
                lines.append(f"    - {uc}")
        if self.typical_structure:
            lines.append(f"  Typical structure:\n{self.typical_structure}")
        return "\n".join(lines)


# ============================================================
# Dataset Definitions
# ============================================================

NLST = DatasetInfo(
    name="National Lung Screening Trial",
    short_name="NLST",
    description=(
        "Large-scale lung cancer screening trial with low-dose chest CT. "
        "Contains ~26,000 participants with CT scans at multiple time points. "
        "Many patients have multiple reconstructions at different slice thicknesses, "
        "making it ideal for multi-thickness CAC research."
    ),
    total_cases=26254,
    modality="Low-dose chest CT",
    access_url="https://cdas.cancer.gov/nlst/",
    license="NLST Data Access Agreement (requires application via CDAS)",
    citation="National Lung Screening Trial Research Team, NEJM 2011; 365:395-409",
    file_format="DICOM",
    slice_thickness_range="0.6mm - 10.0mm (most common: 1.25, 2.0, 2.5, 5.0mm)",
    typical_structure=(
        "    NLST/\n"
        "      {patient_id}/          # 6-digit numeric ID\n"
        "        {study_uid}/         # Study Instance UID\n"
        "          {series_uid}/      # Series Instance UID\n"
        "            *.dcm            # DICOM files"
    ),
    has_ground_truth=False,
    ground_truth_format="None (CAC ground truth not included; derive from thin-slice)",
    supports_paired_thickness=True,
    thin_thickness_mm=2.0,
    thick_thickness_mm=5.0,
    paired_selection_criteria=(
        "Thin: SliceThickness 1.8-2.7mm, min 50 slices, kernel B50F/STANDARD; "
        "Thick: SliceThickness 4.5-5.5mm, min 50 slices. "
        "26,254 patients screened -> ~2,222 have valid dual-thickness pairs."
    ),
    use_cases=[
        "Multi-thickness CAC scoring validation",
        "RESCUE phenomenon analysis (thick-layer zero-CAC underdetection)",
        "Super-resolution CT reconstruction training",
        "Lung cancer screening AI development",
        "Population-level cardiovascular risk assessment",
    ],
    notes=[
        "Access requires CDAS application (1-2 weeks approval)",
        "Full dataset is ~12TB; use metadata-first search to identify relevant cases",
        "Patient IDs are 6-digit numbers; Study/Series UIDs follow DICOM standard",
        "Multiple reconstruction kernels per patient are common (B50F, STANDARD, LUNG)",
        "Clinical metadata (age, sex, smoking history) available in separate CSV",
    ],
)

COCA = DatasetInfo(
    name="Stanford Coronary Calcium Dataset",
    short_name="COCA",
    description=(
        "Cardiac CT dataset with expert-annotated CAC scores. "
        "Contains both gated coronary CT (787 cases) and non-gated chest CT (213 cases). "
        "The non-gated subset is the only publicly available non-gated CT dataset "
        "with expert Agatston scoring gold standard."
    ),
    total_cases=1000,
    modality="Cardiac CT (gated + non-gated)",
    access_url="https://stanfordaimi.azurewebsites.net/datasets/e8ca74dc-8f4f-4f1c-84ea-94de72100732",
    license="Stanford AIMI License (research use)",
    citation="Eng et al., JACC Cardiovasc Imaging 2021",
    file_format="DICOM + XML/Excel ground truth",
    slice_thickness_range="1.0-5.0mm",
    typical_structure=(
        "    coca/\n"
        "      gated_coronary_ct/\n"
        "        patient_*/         # Patient folders with DICOM\n"
        "        calcium_xml/       # 451 XML files with per-artery CAC scores\n"
        "      nongated_chest_ct/\n"
        "        patient_*/         # 213 patient folders with DICOM\n"
        "        scores.xlsx        # Expert Agatston scores (LCA/LAD/LCX/RCA/total)"
    ),
    has_ground_truth=True,
    ground_truth_format=(
        "Gated: XML files with per-artery scores (LCA, LAD, LCX, RCA, volume, mass); "
        "Non-gated: Excel with total and per-artery Agatston scores"
    ),
    use_cases=[
        "AI-CAC model validation against expert gold standard",
        "Per-artery calcium scoring accuracy assessment",
        "Gated vs non-gated CAC scoring comparison",
    ],
    notes=[
        "Non-gated subset (213 cases) is the key validation set for non-gated AI-CAC",
        "Some patients have empty SeriesID in DICOM metadata",
        "Gated subset has 787 cases but only 451 have XML ground truth",
        "Ground truth scores.xlsx uses 'filename' column matching patient folder names",
    ],
)

TOTALSEGMENTATOR = DatasetInfo(
    name="TotalSegmentator Dataset",
    short_name="TotalSegmentator",
    description=(
        "Multi-center CT dataset with 117 anatomical structure segmentations. "
        "Covers 9 institutions, multiple CT vendors, and diverse pathologies. "
        "Ideal for multi-center generalization testing of cardiac algorithms."
    ),
    total_cases=1228,
    modality="CT (mixed contrast)",
    access_url="https://zenodo.org/records/6802613",
    license="CC-BY 4.0",
    citation="Wasserthal et al., Radiology: AI 2023 (doi:10.1148/ryai.230024)",
    file_format="NIfTI (.nii.gz) + segmentation masks",
    slice_thickness_range="0.5-5.0mm",
    typical_structure=(
        "    Totalsegmentator_dataset_v201/\n"
        "      s{XXXX}/              # Patient ID (4-digit)\n"
        "        ct.nii.gz           # CT volume\n"
        "        segmentations/      # 117 structure masks\n"
        "          aorta.nii.gz\n"
        "          heart.nii.gz\n"
        "          ..."
    ),
    has_ground_truth=True,
    ground_truth_format="NIfTI segmentation masks for 117 anatomical structures",
    use_cases=[
        "Multi-center algorithm generalization testing",
        "Anatomical structure segmentation validation",
        "Body composition analysis (with organ masks)",
        "Cross-scanner robustness evaluation",
    ],
    notes=[
        "79.4% non-contrast, 20.6% angiography",
        "Institutions labeled A-F (anonymized)",
        "Manufacturers: Siemens, GE primarily",
        "Version 2.0.1 has 1,228 cases (v1 had 1,204)",
        "Segmentation includes: heart, aorta, vertebrae T1-L5, and 110+ more",
    ],
)


# Registry of all known datasets
_DATASETS: Dict[str, DatasetInfo] = {
    'nlst': NLST,
    'coca': COCA,
    'totalsegmentator': TOTALSEGMENTATOR,
    'totalseg': TOTALSEGMENTATOR,  # alias
}


def get_dataset_info(name: str) -> DatasetInfo:
    """
    Get standardized metadata for a public cardiac imaging dataset.

    Args:
        name: Dataset identifier (case-insensitive).
              Supported: 'nlst', 'coca', 'totalsegmentator' (or 'totalseg')

    Returns:
        DatasetInfo with access URLs, structure, citation, and usage notes.

    Raises:
        KeyError: If dataset name is not recognized.

    Example:
        >>> info = get_dataset_info('nlst')
        >>> print(info.access_url)
        https://cdas.cancer.gov/nlst/
        >>> print(info.paired_selection_criteria)
    """
    key = name.lower().strip()
    if key not in _DATASETS:
        available = ', '.join(sorted(set(
            k for k, v in _DATASETS.items()
            if k == v.short_name.lower()
        )))
        raise KeyError(
            f"Unknown dataset '{name}'. Available: {available}. "
            f"Contribute new datasets via cardiac_shared.data.register_dataset()."
        )
    return _DATASETS[key]


def list_datasets() -> List[str]:
    """
    List all available dataset identifiers.

    Returns:
        Sorted list of unique dataset short names.
    """
    seen = set()
    result = []
    for info in _DATASETS.values():
        if info.short_name not in seen:
            seen.add(info.short_name)
            result.append(info.short_name)
    return sorted(result)


def print_dataset_catalog():
    """Print a formatted catalog of all available datasets."""
    print("=" * 60)
    print("Cardiac Imaging External Dataset Catalog")
    print("=" * 60)
    for name in list_datasets():
        info = get_dataset_info(name)
        print(f"\n{info.short_name}: {info.name}")
        print(f"  Cases: {info.total_cases} | Format: {info.file_format}")
        print(f"  Ground truth: {'Yes' if info.has_ground_truth else 'No'}")
        print(f"  Access: {info.access_url}")
    print()
    print(f"Total: {len(list_datasets())} datasets")
    print("Use get_dataset_info('name') for full details.")


def register_dataset(info: DatasetInfo):
    """
    Register a custom dataset.

    Args:
        info: DatasetInfo instance with dataset metadata.

    Example:
        >>> from cardiac_shared.data import DatasetInfo, register_dataset
        >>> my_dataset = DatasetInfo(
        ...     name="My Hospital Dataset",
        ...     short_name="MHD",
        ...     description="Internal validation cohort",
        ...     total_cases=500,
        ...     modality="Cardiac CT",
        ...     access_url="internal",
        ...     license="Proprietary",
        ...     citation="Unpublished",
        ... )
        >>> register_dataset(my_dataset)
    """
    key = info.short_name.lower()
    _DATASETS[key] = info
