"""Plant CSV dataset for data/wa_plants/manifest.csv (observation-separated).
Reuses RSNA train_mor pattern: manifest + ImageNet norm, 336px, Aug.
"""
from __future__ import annotations
import csv
from pathlib import Path
import random
from PIL import Image
import warnings
# faster decode: raise limit (100Mpx originals) + draft shrink for large JPEGs
Image.MAX_IMAGE_PIXELS = 300_000_000
warnings.filterwarnings("ignore", category=Image.DecompressionBombWarning)
import torch
from torch.utils.data import Dataset
import torchvision.transforms as T

IMAGENET_MEAN = (0.485, 0.456, 0.406)
IMAGENET_STD = (0.229, 0.224, 0.225)

class PlantDataset(Dataset):
    def __init__(self, manifest: Path, split: str, img_size: int = 336, augment: bool = False):
        self.split = split
        self.img_size = img_size
        rows = []
        # species -> idx map built from manifest train split (500 classes)
        with open(manifest, newline="", encoding="utf-8") as f:
            for r in csv.DictReader(f):
                if r["split"] == split and r["status"] in ("downloaded","skip_exists"):
                    rows.append(r)
        # stable class order sorted
        species = sorted({r["species"] for r in rows})
        self.species_to_idx = {s:i for i,s in enumerate(species)}
        self.rows = rows
        # class counts for balanced sampling
        self.augment = augment
        tfms = []
        if augment:
            tfms = [
                T.RandomResizedCrop(img_size, scale=(0.7,1.0)),
                T.RandomHorizontalFlip(),
                T.ColorJitter(0.2,0.2,0.2,0.05),
                T.ToTensor(), T.Normalize(IMAGENET_MEAN, IMAGENET_STD),
            ]
        else:
            tfms = [T.Resize(int(img_size*1.14)), T.CenterCrop(img_size), T.ToTensor(), T.Normalize(IMAGENET_MEAN, IMAGENET_STD)]
        self.tf = T.Compose(tfms)

    def __len__(self): return len(self.rows)
    def __getitem__(self, i):
        r = self.rows[i]
        p = Path(r["path"])
        # manifest stores absolute win path; if not found try relative to data/wa_plants
        if not p.exists():
            # try finding under train/val subfolders by gbifID
            base = Path(__file__).resolve().parents[2] / "data" / "wa_plants"
            for split in ("train","val"):
                cand = base / split / r["species"].replace(" ","_").replace("/","_")[:120] / f"{r['gbifID']}.jpg"
                if cand.exists():
                    p = cand; break
        try:
            im = Image.open(p)
            # fast draft for huge JPEGs (8× shrink before full decode) - no effect on small images
            try:
                if max(im.size) > 1024:
                    # draft uses libjpeg shrink 1/2/4/8
                    im.draft("RGB", (768, 768))
            except:
                pass
            im = im.convert("RGB")
        except Exception:
            im = Image.new("RGB", (self.img_size, self.img_size))
        # species_to_idx may have been remapped after init (train vs val) - lookup safely
        y = self.species_to_idx.get(r["species"], 0)
        x = self.tf(im)
        return x, y, r["gbifID"]