import lightning as L
import numpy as np
from torch.utils.data import DataLoader
from datasets import load_dataset
from src.dataset import DecoderDataset

class StyleTransferDataModule(L.LightningDataModule):
    def __init__(self, batch_size=8, image_size=256, val_split=0.05, test_split=0.05, max_samples=None):
        super().__init__()
        self.save_hyperparameters()
        
        self.content_name = "wangwangxuebing/train_COCO_data2014_jpg"
        self.style_name = "newsletter/HiDream-I1-Artists"

    def setup(self, stage=None):
        full_content = load_dataset(self.content_name, split="test")
        full_style = load_dataset(self.style_name, split="train")

        if self.hparams.max_samples:
            full_content = full_content.select(range(min(len(full_content), self.hparams.max_samples)))
            full_style = full_style.select(range(min(len(full_style), self.hparams.max_samples)))

        seed = 42
        content_indexes = np.random.RandomState(seed).permutation(len(full_content))
        style_indexes = np.random.RandomState(seed).permutation(len(full_style))

        def get_splits(all_indexes, val_ratio, test_ratio):
            total_count = len(all_indexes)
            
            n_val = int(total_count * val_ratio)
            n_test = int(total_count * test_ratio)
            n_train = total_count - n_val - n_test

            train_end = n_train
            val_end = n_train + n_val

            train_indices = all_indexes[:train_end]
            val_indices   = all_indexes[train_end:val_end]
            test_indices  = all_indexes[val_end:]
            
            return train_indices, val_indices, test_indices

        content_train, content_val, content_test = get_splits(content_indexes, self.hparams.val_split, self.hparams.test_split)
        style_train, style_val, style_test = get_splits(style_indexes, self.hparams.val_split, self.hparams.test_split)

        if stage == "fit" or stage is None:
            self.train_ds = DecoderDataset(full_content, full_style, content_train, style_train, self.hparams.image_size, is_train=True)
            self.val_ds = DecoderDataset(full_content, full_style, content_val, style_val, self.hparams.image_size, is_train=False)

        if stage == "test":
            self.test_ds = DecoderDataset(full_content, full_style, content_test, style_test, self.hparams.image_size, is_train=False)

    def train_dataloader(self):
        return DataLoader(self.train_ds, batch_size=self.hparams.batch_size, shuffle=True, num_workers=12, pin_memory=True)

    def val_dataloader(self):
        return DataLoader(self.val_ds, batch_size=self.hparams.batch_size, shuffle=False, num_workers=4, pin_memory=True)

    def test_dataloader(self):
        return DataLoader(self.test_ds, batch_size=self.hparams.batch_size, shuffle=False, num_workers=4)