"""Dataset loading and filename metadata parsing for the clonogenic-assay images.

Filenames look like:
    N1_N2A_Clono_ET+RT_10Gy+Vor_Day21_5X_13.bmp
from which we parse the radiation dose (Gy), whether Vorinostat (+Vor) was added,
and the source folder. These give optional ground-truth groupings used only to
*evaluate* the unsupervised clustering (they are never used to fit it).
"""
import glob
import os
import re
from dataclasses import dataclass
from typing import List, Optional

IMAGE_EXTS = (".bmp", ".png", ".jpg", ".jpeg", ".tif", ".tiff")

_DOSE_RE = re.compile(r"ET\+RT_(\d+)Gy(\+Vor)?", re.IGNORECASE)


@dataclass
class ImageItem:
    path: str
    folder: str
    dose: Optional[float]      # Gy
    vor: Optional[bool]        # Vorinostat present
    filename: str

    @property
    def treatment(self) -> Optional[str]:
        if self.dose is None:
            return None
        return f"{int(self.dose)}Gy" + ("+Vor" if self.vor else "")


def parse_metadata(path: str, folder: str = "") -> ImageItem:
    fn = os.path.basename(path)
    m = _DOSE_RE.search(fn)
    dose = float(m.group(1)) if m else None
    vor = bool(m.group(2)) if m else None
    return ImageItem(path=path, folder=folder, dose=dose, vor=vor, filename=fn)


def list_images(folder: str, tag: str = "") -> List[ImageItem]:
    items: List[ImageItem] = []
    for ext in IMAGE_EXTS:
        for p in glob.glob(os.path.join(folder, f"*{ext}")):
            items.append(parse_metadata(p, tag or os.path.basename(folder.rstrip("/"))))
    items.sort(key=lambda it: it.filename)
    return items


def load_dataset(folders: List[str], tags: Optional[List[str]] = None) -> List[ImageItem]:
    """Load and concatenate images from several folders.

    tags: optional label per folder (defaults to the folder's basename), stored
    as ImageItem.folder so clustering can be evaluated against 'which set'.
    """
    tags = tags or [os.path.basename(f.rstrip("/")) for f in folders]
    all_items: List[ImageItem] = []
    for folder, tag in zip(folders, tags):
        all_items.extend(list_images(folder, tag))
    return all_items