# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Dict, List, Optional

import lightning.pytorch as pl
import numpy as np
import torch
from lightning.pytorch.utilities.types import EVAL_DATALOADERS, TRAIN_DATALOADERS
from PIL import Image
from torch.utils import data
from torch.utils.data import DataLoader, Dataset
from transformers import Gemma3ImageProcessor

from nemo.lightning.pytorch.plugins import MegatronDataSampler
from nemo.utils import logging

IMAGE_TOKEN_INDEX = 262144
IMAGE_SIZE = 896
IMAGE_TOKENS = 256


class Gemma3VLMockDataModule(pl.LightningDataModule):
    """
    A mock data module for Gemma3VL training, validation, and testing.
    """

    def __init__(
        self,
        tokenizer: Optional = None,
        seq_length: int = 2048,
        micro_batch_size: int = 4,
        global_batch_size: int = 8,
        rampup_batch_size: Optional[List[int]] = None,
        num_train_samples: int = 10_000,
        num_val_samples: int = 10_000,
        num_test_samples: int = 10_000,
        num_workers: int = 8,
        pin_memory: bool = True,
        persistent_workers: bool = False,
    ):
        super().__init__()
        self.tokenizer = tokenizer
        self.seq_length = seq_length
        self.micro_batch_size = micro_batch_size
        self.global_batch_size = global_batch_size
        self.num_train_samples = num_train_samples
        self.num_val_samples = num_val_samples
        self.num_test_samples = num_test_samples
        self.num_workers = num_workers
        self.pin_memory = pin_memory
        self.persistent_workers = persistent_workers

        if tokenizer is None:
            logging.warning("Tokenizer is not provided! Fall back to `'google/gemma-3-4b-it'`.")
            from nemo.collections.common.tokenizers.huggingface.auto_tokenizer import AutoTokenizer

            self.tokenizer = tokenizer or AutoTokenizer("google/gemma-3-4b-it")

        self.image_processor = Gemma3ImageProcessor(
            size={"height": IMAGE_SIZE, "width": IMAGE_SIZE},
            do_pan_and_scan=False,
            pan_and_scan_min_crop_size=256,
            pan_and_scan_max_num_crops=4,
            pan_and_scan_min_ratio_to_activate=1.2,
        )
        self.data_sampler = MegatronDataSampler(
            seq_len=self.seq_length,
            micro_batch_size=micro_batch_size,
            global_batch_size=global_batch_size,
            rampup_batch_size=rampup_batch_size,
        )

    def setup(self, stage: str = "") -> None:
        # pylint: disable=C0115,C0116
        self._train_ds = _Qwen2VLMockDataset(
            self.tokenizer, self.image_processor, "train", self.num_train_samples, self.seq_length
        )
        self._validation_ds = _Qwen2VLMockDataset(
            self.tokenizer, self.image_processor, "valid", self.num_val_samples, self.seq_length
        )
        self._test_ds = _Qwen2VLMockDataset(
            self.tokenizer, self.image_processor, "test", self.num_test_samples, self.seq_length
        )

    def train_dataloader(self) -> TRAIN_DATALOADERS:
        # pylint: disable=C0115,C0116
        if not hasattr(self, "_train_ds"):
            self.setup()
        return self._create_dataloader(self._train_ds)

    def val_dataloader(self) -> EVAL_DATALOADERS:
        # pylint: disable=C0115,C0116
        if not hasattr(self, "_validation_ds"):
            self.setup()
        return self._create_dataloader(self._validation_ds)

    def test_dataloader(self) -> EVAL_DATALOADERS:
        # pylint: disable=C0115,C0116
        if not hasattr(self, "_test_ds"):
            self.setup()
        return self._create_dataloader(self._test_ds)

    def _create_dataloader(self, dataset, **kwargs) -> DataLoader:
        # pylint: disable=C0115,C0116
        return DataLoader(
            dataset,
            num_workers=self.num_workers,
            pin_memory=self.pin_memory,
            persistent_workers=self.persistent_workers,
            collate_fn=dataset.collate_fn,
            **kwargs,
        )


def prepare_image_inputs(num_channels: np.uint8 = 3, width=1024, height=1024):
    """This function prepares a list of PIL images"""
    image_inputs = [np.random.randint(255, size=(num_channels, width, height), dtype=np.uint8)]
    image_inputs = [Image.fromarray(np.moveaxis(x, 0, -1)) for x in image_inputs]
    return image_inputs


class _Qwen2VLMockDataset(Dataset):
    def __init__(
        self,
        tokenizer,
        image_processor,
        name: str,
        num_samples: int,
        seq_length: int,
        seed: int = 42,
    ) -> None:
        super().__init__()
        self.name = name
        self.seq_length = seq_length

        self.vocab_size = tokenizer.vocab_size - 10

        self.image_processor = image_processor

        self.length = num_samples
        self.seed = seed

        self.loss_mask = torch.ones(self.seq_length, dtype=torch.float)
        self.position_ids = torch.arange(self.seq_length, dtype=torch.int64)

    def __len__(self) -> int:
        return self.length

    def _get_text(self, idx: int) -> np.ndarray:
        np_gen = np.random.default_rng(seed=self.seed + idx)
        return np_gen.integers(low=0, high=self.vocab_size, size=[self.seq_length], dtype=np.int64)

    def __getitem__(self, idx) -> Dict[str, torch.Tensor]:

        # 1) processor image input
        image_inputs = prepare_image_inputs(3, IMAGE_SIZE, IMAGE_SIZE)
        prcocess_out = self.image_processor(image_inputs, return_tensors="pt")
        pixel_values = prcocess_out.pixel_values.to(torch.bfloat16)

        # 2) prepare input token ids
        np_gen = np.random.default_rng(seed=self.seed + idx)
        tokens = torch.from_numpy(
            np_gen.integers(low=0, high=self.vocab_size, size=[self.seq_length + 1], dtype=np.int64)
        )

        # 3) fill IMAGE_TOKEN_INDEX to input token ids
        img_start_idx = 20  # pick a rnd value, where img token begins
        tokens[img_start_idx : img_start_idx + IMAGE_TOKENS] = IMAGE_TOKEN_INDEX  # ImageToken token index
        input_ids = tokens[:-1]
        position_ids = torch.arange(self.seq_length, dtype=torch.int64)

        # 4) prepare labels
        labels = tokens.clone()
        generation_prompt_size = 5  # "ASSISTANT:" like
        prompt_end_idx = img_start_idx + IMAGE_TOKENS + generation_prompt_size
        labels[:prompt_end_idx] = -100
        # Add the labels clipping to the mock data loader.
        labels = labels[1:]

        # 5) prepare loss masks
        # Calculate loss mask from labels, to be consistent with real data and reduce confusions.
        loss_mask = torch.ones_like(labels, dtype=torch.float)
        loss_mask[labels < 0] = 0.0

        return {
            "input_ids": input_ids,
            "position_ids": position_ids,
            "pixel_values": pixel_values,
            "loss_mask": loss_mask,
            "labels": labels,
        }

    def _collate_fn(self, batch):
        """
        A default implementation of a collation function.
        Users should override this method to define custom data loaders.
        """
        collated_batch = data.dataloader.default_collate(batch)
        collated_batch["attention_mask"] = None
        return collated_batch

    def collate_fn(self, batch):
        """Method that user pass as functor to DataLoader.

        The method optionally performs neural type checking and add types to the outputs.

        Please note, subclasses of Dataset should not implement `input_types`.

        # Usage:
        dataloader = torch.utils.data.DataLoader(
                ....,
                collate_fn=dataset.collate_fn,
                ....
        )

        Returns
        -------
            Collated batch, with or without types.
        """
        return self._collate_fn(batch)
