"""PyTorch Dataset class for the APTOS 2019 Blindness Detection challenge."""

import os
from typing import Callable, Dict, Optional, Tuple, Union
import cv2
import numpy as np
import pandas as pd
import torch
from torch.utils.data import Dataset

from src.data.preprocessing import ben_graham_pipeline


class APTOSDataset(Dataset):
    """APTOS 2019 Diabetic Retinopathy dataset."""

    def __init__(
        self,
        df_or_csv: Union[pd.DataFrame, str],
        images_dir: str,
        transform: Optional[Callable] = None,
        image_size: int = 380,
        apply_ben_graham: bool = True,
        is_preprocessed: bool = False,
    ) -> None:
        """Initialize APTOS dataset.

        Args:
            df_or_csv: DataFrame or path to CSV with columns ['id_code', 'diagnosis'].
            images_dir: Path to image directory containing PNG images.
            transform: Optional Albumentations transform pipeline.
            image_size: Target image dimension for resizing / preprocessing.
            apply_ben_graham: Whether to apply Ben Graham preprocessing on the fly.
            is_preprocessed: If True, images in images_dir are already preprocessed.
        """
        if isinstance(df_or_csv, str):
            self.df = pd.read_csv(df_or_csv)
        else:
            self.df = df_or_csv.reset_index(drop=True)

        self.images_dir = images_dir
        self.transform = transform
        self.image_size = image_size
        self.apply_ben_graham = apply_ben_graham and not is_preprocessed
        self.is_preprocessed = is_preprocessed

        # Standardize column names
        if "id_code" in self.df.columns:
            self.image_ids = self.df["id_code"].values
        elif "image_id" in self.df.columns:
            self.image_ids = self.df["image_id"].values
        else:
            self.image_ids = self.df.iloc[:, 0].values

        if "diagnosis" in self.df.columns:
            self.labels = self.df["diagnosis"].values.astype(np.int64)
        else:
            self.labels = self.df.iloc[:, 1].values.astype(np.int64)

    def __len__(self) -> int:
        return len(self.df)

    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int, str]:
        """Fetch image tensor, integer label, and image path.

        Returns:
            Tuple of (image_tensor, label, image_path).
        """
        img_id = str(self.image_ids[idx])
        # Support both raw id and id with extension
        if not img_id.endswith((".png", ".jpg", ".jpeg")):
            img_filename = f"{img_id}.png"
        else:
            img_filename = img_id

        img_path = os.path.join(self.images_dir, img_filename)

        # Fallback to jpg if png not found
        if not os.path.exists(img_path):
            alt_path = os.path.join(self.images_dir, f"{img_id}.jpg")
            if os.path.exists(alt_path):
                img_path = alt_path

        image = cv2.imread(img_path)
        if image is None:
            raise FileNotFoundError(f"Failed to load image at {img_path}")

        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        if self.apply_ben_graham:
            image = ben_graham_pipeline(image, image_size=self.image_size)

        if self.transform is not None:
            augmented = self.transform(image=image)
            image_tensor = augmented["image"]
        else:
            # Fallback simple tensor conversion if no transform provided
            image_tensor = torch.from_numpy(image.transpose(2, 0, 1)).float() / 255.0

        label = int(self.labels[idx])
        return image_tensor, label, img_path

    def get_class_counts(self) -> Dict[int, int]:
        """Compute frequency per severity grade."""
        unique, counts = np.unique(self.labels, return_counts=True)
        return dict(zip(unique.tolist(), counts.tolist()))

    def get_class_weights(self) -> torch.Tensor:
        """Compute balanced inverse-frequency class weights for loss functions.

        weight_c = total_samples / (num_classes * count_c)
        """
        counts = self.get_class_counts()
        num_classes = 5
        total = len(self.labels)
        weights = []
        for c in range(num_classes):
            c_count = counts.get(c, 1)
            weights.append(total / (num_classes * c_count))
        return torch.tensor(weights, dtype=torch.float32)