# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.

# SPDX-License-Identifier: LicenseRef-Apache2

#

# 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 pathlib import Path

from typing import List, Optional, Sequence, TypedDict

import lightning.pytorch as pl

import numpy as np

import torch

# In lightning.pytorch 2.0 these are commented as being "any iterable or collection of iterables"

#  for now we'll use them incase the lightning type becomes something more specific in a future release.

from lightning.pytorch.utilities.types import EVAL_DATALOADERS, TRAIN_DATALOADERS

from nemo import lightning as nl

from nemo.collections import llm

from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec

from nemo.collections.nlp.modules.common.tokenizer_utils import get_nmt_tokenizer

from nemo.lightning.megatron_parallel import DataT

from nemo.lightning.pytorch.plugins import MegatronDataSampler

from torch.utils import data

from torch.utils.data import DataLoader, Dataset

from bionemo.llm.model.biobert.lightning import LossLoggingCallback

__all__: Sequence[str] = ()

class MockDataModule(pl.LightningDataModule):

    def __init__(

        self,

        seq_length: int = 2048,

        tokenizer: Optional[TokenizerSpec] = None,

        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.seq_length = seq_length

        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

        self.tokenizer = tokenizer or get_nmt_tokenizer("megatron", "GPT2BPETokenizer")

        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,

        )

        # NOTE: the datasets and other distributed state is instantiated in `setup` rather than in `__init__` to support

        #  the different kinds of accellerators/strategies that lightning supports. This is a common pattern in lightning.

    def setup(self, stage: str = "") -> None:

        """See lightning documentation for more information on the stage and setup method. It is not required but

        if you want to be efficient about only initializing data that is needed in a particular stage you can do it here.

        According to the documentation valid values match the available calls to trainer.{fit,validate,test,predict},

        for example stage="fit". If we wanted to be fancy we could only initialize train/val during "fit". We could

        only instantiate "test" data during "test" etc.

        """

        self._train_ds = _MockGPTDataset(self.tokenizer, "train", self.num_train_samples, self.seq_length)

        self._validation_ds = _MockGPTDataset(self.tokenizer, "valid", self.num_val_samples, self.seq_length)

        self._test_ds = _MockGPTDataset(self.tokenizer, "test", self.num_test_samples, self.seq_length)

    def train_dataloader(self) -> TRAIN_DATALOADERS:

        return self._create_dataloader(self._train_ds)

    def val_dataloader(self) -> EVAL_DATALOADERS:

        return self._create_dataloader(self._validation_ds)

    def test_dataloader(self) -> EVAL_DATALOADERS:

        return self._create_dataloader(self._test_ds)

    def _create_dataloader(self, dataset, **kwargs) -> DataLoader:

        return DataLoader(

            dataset,

            num_workers=self.num_workers,

            pin_memory=self.pin_memory,

            persistent_workers=self.persistent_workers,

            collate_fn=dataset.collate_fn,

            **kwargs,

        )

class GptDataItem(TypedDict):

    tokens: torch.Tensor

    labels: torch.Tensor

    attention_mask: torch.Tensor

    loss_mask: torch.Tensor

    position_ids: torch.Tensor

class _MockGPTDataset(Dataset):

    def __init__(

        self,

        tokenizer: TokenizerSpec,

        name: str,

        num_samples: int,

        seq_length: int,

        seed: int = 42,

    ):

        super().__init__()

        self.name = name

        self.seq_length = seq_length

        self.vocab_size = tokenizer.vocab_size

        self.length = num_samples

        self.seed = seed

        self.attention_mask = torch.tril(torch.ones((self.seq_length, self.seq_length))).unsqueeze(0)

        self.attention_mask = self.attention_mask < 0.5

        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(self.vocab_size, size=[self.seq_length], dtype=np.int64)

    def __getitem__(self, idx) -> GptDataItem:

        # Generate data of the expected size and datatype (based on GPTDataset).

        np_gen = np.random.default_rng(seed=(self.seed + idx))

        # Always return the same thing

        np_gen = np.random.default_rng(seed=(self.seed))

        tokens = torch.from_numpy(np_gen.integers(self.vocab_size, size=[self.seq_length], dtype=np.int64))

        labels = torch.from_numpy(np_gen.integers(self.vocab_size, size=[self.seq_length], dtype=np.int64))

        return {

            "tokens": tokens,

            "labels": labels,

            "attention_mask": self.attention_mask,

            "loss_mask": self.loss_mask,

            "position_ids": self.position_ids,

        }

    def _collate_fn(self, batch: DataT) -> DataT:

        """A default implementation of a collation function.

        Users should override this method to define custom data loaders.

        """

        return data.dataloader.default_collate(batch)

    def collate_fn(self, batch: DataT) -> DataT:

        """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)

def main() -> None:

    devices, seq_length = 1, 2048

    strategy = nl.MegatronStrategy(

        tensor_model_parallel_size=1,

        pipeline_model_parallel_size=1,

        pipeline_dtype=torch.float32,

        ckpt_async_save=False,

    )

    trainer = nl.Trainer(

        devices=devices,

        max_steps=100,

        accelerator="gpu",

        strategy=strategy,

        callbacks=[LossLoggingCallback()],

        # TODO(@jstjohn) See if we can get the example working with mixed precision

        # plugins=nl.MegatronMixedPrecision(precision="float32", amp_O2=False),

    )

    _data = MockDataModule(seq_length=seq_length, global_batch_size=32)

    gpt_config = llm.GPTConfig(

        num_layers=4,

        hidden_size=256,

        ffn_hidden_size=512,

        num_attention_heads=4,

        seq_length=seq_length,

        pipeline_dtype=torch.float32,

    )

    model = llm.GPTModel(gpt_config, tokenizer=_data.tokenizer)

    trainer.fit(model, _data)

    checkpoint_path = Path(trainer.logger.log_dir) / "ckpt"

    trainer.save_checkpoint(checkpoint_path)

if __name__ == "__main__":

    main()