import torch
import pandas as pd
import pytorch_lightning as pl

from PIL import Image
from pathlib import Path
from torchvision import datasets
from torchvision.io import read_file, decode_image, ImageReadMode
from typing import Optional, Callable, Any
from torch.utils.data import DataLoader, Dataset


TransformType = Optional[Callable[..., Any]]


def image_loader(path: str) -> torch.Tensor:
    data = read_file(path)

    return decode_image(
        data,
        mode=ImageReadMode.RGB
    )


class ImageNetValDataset(Dataset):
    def __init__(
        self,
        image_dir,
        labels_file,
        class_to_idx,
        classes,
        transform=None
    ):
        self.image_dir = Path(image_dir)
        self.transform = transform

        self.class_to_idx = class_to_idx
        self.classes = classes

        df = pd.read_csv(
            labels_file,
            usecols=["ImageId", "PredictionString"]
        )

        image_ids = df["ImageId"].to_numpy()

        synsets = (
            df["PredictionString"]
            .str.split()
            .str[0]
            .to_numpy()
        )

        self.samples = [
            (
                self.image_dir / f"{image_id}.JPEG",
                class_to_idx[synset]
            )
            for image_id, synset in zip(image_ids, synsets)
        ]

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        image_path, label = self.samples[idx]

        image = decode_image(
            str(image_path),
            mode=ImageReadMode.RGB
        )

        if self.transform is not None:
            image = self.transform(image)

        return image, label


class ImageNetDataModule(pl.LightningDataModule):
    def __init__(self, 
                 dataset_root: Path,
                 train_transform: TransformType,
                 eval_transform: TransformType,
                 batch_size: int = 64,
                 num_workers:int = 8
                 ):
        super().__init__()
        
        self.dataset_root = Path(dataset_root)
        
        self.train_transform = train_transform
        self.eval_transform = eval_transform
        self.num_workers = num_workers
        self.batch_size = batch_size
        
        self.imagenet_root = Path(dataset_root) / "ILSVRC"
        
        self.train_dir = self.imagenet_root / "Data" / "CLS-LOC" / "train"
        self.val_dir = self.imagenet_root / "Data" / "CLS-LOC" / "val"
        
        self.mapping_file = Path(dataset_root) / "LOC_synset_mapping.txt"
        self.val_labels_file = Path(dataset_root) / "LOC_val_solution.csv"

        self.train_dataset = None
        self.val_dataset = None
        
        self.class_to_idx = None
        self.classes = None
        
    def setup(self, stage=None):
        synset_to_name = {}
        with open(self.mapping_file, "r", encoding="utf-8") as f:
            for line in f:
                synset, class_name = line.strip().split(" ", 1)
                synset_to_name[synset] = class_name
        
        train_dataset = datasets.ImageFolder(root=self.train_dir, 
                                             transform=self.train_transform,
                                             loader=image_loader)
    
        self.class_to_idx = train_dataset.class_to_idx.copy()
        self.classes = [synset_to_name[synset] 
                        for synset in train_dataset.classes]
        train_dataset.classes = self.classes
        self.train_dataset = train_dataset
        
        self.val_dataset = ImageNetValDataset(
            image_dir=self.val_dir,
            labels_file=self.val_labels_file,
            class_to_idx=self.class_to_idx,
            classes=self.classes,
            transform=self.eval_transform
        )
        
    def train_dataloader(self):
        return DataLoader(
            dataset=self.train_dataset,
            batch_size=self.batch_size,
            shuffle=True,
            num_workers=self.num_workers,
            pin_memory=True,
            persistent_workers=self.num_workers > 0,
            prefetch_factor=2 if self.num_workers > 0 else None,
            drop_last=True
        )
    
    def val_dataloader(self):
        return DataLoader(
            dataset=self.val_dataset,
            batch_size=self.batch_size,
            shuffle=False,
            num_workers=self.num_workers,
            pin_memory=True,
            persistent_workers=self.num_workers > 0,
            prefetch_factor=2 if self.num_workers > 0 else None,
            drop_last=False
        )