# Copyright (c) Meta Platforms, Inc. and affiliates.
# This software may be used and distributed according to the terms of the Llama 2 Community License Agreement.

import inspect
# from dataclasses import asdict

import torch.distributed as dist
from torch.utils.data import DistributedSampler
from peft import (
    LoraConfig,
    AdaptionPromptConfig,
    PrefixTuningConfig,
)
from transformers import default_data_collator
from transformers.data import DataCollatorForSeq2Seq

# from llama_recipes.configs import datasets, lora_config, llama_adapter_config, prefix_config, train_config
from slam_llm.data.sampler import LengthBasedBatchSampler, DistributedLengthBasedBatchSampler

from omegaconf import OmegaConf

import logging
logger = logging.getLogger(__name__)

# def update_config(config, **kwargs):
#     if isinstance(config, (tuple, list)):
#         for c in config:
#             update_config(c, **kwargs)
#     else:
#         for k, v in kwargs.items():
#             if hasattr(config, k):
#                 setattr(config, k, v)
#             elif "." in k:
#                 # allow --some_config.some_param=True
#                 config_name, param_name = k.split(".")
#                 if type(config).__name__ == config_name:
#                     if hasattr(config, param_name):
#                         setattr(config, param_name, v)
#                     else:
#                         # In case of specialized config we can warm user
#                         logger.warning(f"Warning: {config_name} does not accept parameter: {k}")
#             elif isinstance(config, train_config):
#                 logger.warning(f"Warning: unknown parameter {k}")


def generate_peft_config(train_config):
    # configs = (lora_config, llama_adapter_config, prefix_config)
    # peft_configs = (LoraConfig, AdaptionPromptConfig, PrefixTuningConfig)
    peft_configs = {"lora": LoraConfig,
                    "llama_adapter": AdaptionPromptConfig,
                    "prefix": PrefixTuningConfig
                    }
    # names = tuple(c.__name__.rstrip("_config") for c in configs)
    #
    # assert train_config.peft_method in names, f"Peft config not found: {train_config.peft_method}"
    #
    # config = configs[names.index(train_config.peft_method)]()
    config = train_config.peft_config

    params = OmegaConf.to_container(config, resolve=True)
    # peft_config = peft_configs[names.index(train_config.peft_method)](**params)
    params.pop("peft_method", None) #(FIX:MZY): remove peft_method from params to avoid error
    peft_config = peft_configs[config.get("peft_method", "lora")](**params)

    return peft_config


def get_dataloader_kwargs(train_config, dataset, tokenizer, mode):
        kwargs = {}
        batch_size = train_config.batch_size_training if mode=="train" else train_config.val_batch_size
        if train_config.batching_strategy == "padding":
            if train_config.enable_fsdp or train_config.enable_ddp or train_config.enable_deepspeed:
                kwargs["batch_sampler"] = DistributedLengthBasedBatchSampler(
                    dataset,
                    batch_size=batch_size,
                    rank=dist.get_rank(),
                    num_replicas=dist.get_world_size(),
                    shuffle=mode=="train",
                )
            else:
                kwargs["batch_sampler"] = LengthBasedBatchSampler(dataset, batch_size, drop_last=True, shuffle=mode=="train")
            kwargs["collate_fn"] = DataCollatorForSeq2Seq(tokenizer)
        elif train_config.batching_strategy == "packing":
            if train_config.enable_fsdp or train_config.enable_ddp or train_config.enable_deepspeed:
                kwargs["sampler"] = DistributedSampler(
                dataset,
                rank=dist.get_rank(),
                num_replicas=dist.get_world_size(),
                shuffle=mode=="train",
            )
            kwargs["batch_size"] = batch_size
            kwargs["drop_last"] = True
            kwargs["collate_fn"] = default_data_collator
        else:
            # raise ValueError(f"Unknown batching strategy: {train_config.batching_strategy}")
            if train_config.enable_fsdp or train_config.enable_ddp or train_config.enable_deepspeed:
                kwargs["sampler"] = DistributedSampler(
                dataset,
                rank=dist.get_rank(),
                num_replicas=dist.get_world_size(),
                shuffle=mode=="train",
            )
            kwargs["batch_size"] = batch_size
            kwargs["drop_last"] = True
            kwargs["collate_fn"] = dataset.collator
            logger.info(f"Using batching strategy: {train_config.batching_strategy}")

        return kwargs