import os
import torch
import torch.nn as nn

from torch.utils.data import Sampler

from transformers import Trainer
from transformers.trainer import (
    has_length,
)
from typing import List, Optional
def set_seed(seed: int, deterministic: bool = False):
    """
    Helper function for reproducible behavior to set the seed in `random`, `numpy`, `torch` and/or `tf` (if installed).

    Args:
        seed (`int`):
            The seed to set.
        deterministic (`bool`, *optional*, defaults to `False`):
            Whether to use deterministic algorithms where available. Can slow down training.
    """
    import numpy as np
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    # ^^ safe to call this function even if cuda is not available
    if deterministic:
        torch.use_deterministic_algorithms(True)
def seed_worker(_):
    """
    Helper function to set worker seed during Dataloader initialization.
    """
    worker_seed = torch.initial_seed() % 2**32
    set_seed(worker_seed)

def maybe_zero_3(param, ignore_status=False, name=None):
    from deepspeed import zero
    from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus
    if hasattr(param, "ds_id"):
        if param.ds_status == ZeroParamStatus.NOT_AVAILABLE:
            if not ignore_status:
                print(name, 'no ignore status')
        with zero.GatheredParameters([param]):
            param = param.data.detach().cpu().clone()
    else:
        param = param.detach().cpu().clone()
    return param


def get_mm_adapter_state_maybe_zero_3(named_params, keys_to_match):
    to_return = {k: t for k, t in named_params if any(key_match in k for key_match in keys_to_match)}
    to_return = {k: maybe_zero_3(v, ignore_status=True, name=k).cpu() for k, v in to_return.items()}
    return to_return


def split_to_even_chunks(indices, lengths, num_chunks):
    """
    Split a list of indices into `chunks` chunks of roughly equal lengths.
    """

    if len(indices) % num_chunks != 0:
        return [indices[i::num_chunks] for i in range(num_chunks)]

    num_indices_per_chunk = len(indices) // num_chunks

    chunks = [[] for _ in range(num_chunks)]
    chunks_lengths = [0 for _ in range(num_chunks)]
    for index in indices:
        shortest_chunk = chunks_lengths.index(min(chunks_lengths))
        chunks[shortest_chunk].append(index)
        chunks_lengths[shortest_chunk] += lengths[index]
        if len(chunks[shortest_chunk]) == num_indices_per_chunk:
            chunks_lengths[shortest_chunk] = float("inf")

    return chunks


def get_modality_length_grouped_indices(lengths, batch_size, world_size, generator=None):
    # We need to use torch for the random part as a distributed sampler will set the random seed for torch.
    assert all(l != 0 for l in lengths), "Should not have zero length."
    mm_indices, mm_lengths = zip(*[(i, l) for i, l in enumerate(lengths) if l > 0])
    lang_indices, lang_lengths = zip(*[(i, -l) for i, l in enumerate(lengths) if l < 0])

    assert len(mm_indices) > 0, "Should have at least one multimodal sample."
    assert len(lang_indices) > 0, "Should have at least one language sample."

    mm_shuffle = [mm_indices[i] for i in get_length_grouped_indices(mm_lengths, batch_size, world_size, generator=None)]
    lang_shuffle = [lang_indices[i] for i in get_length_grouped_indices(lang_lengths, batch_size, world_size, generator=None)]
    megabatch_size = world_size * batch_size
    mm_megabatches = [mm_shuffle[i : i + megabatch_size] for i in range(0, len(mm_shuffle), megabatch_size)]
    lang_megabatches = [lang_shuffle[i : i + megabatch_size] for i in range(0, len(lang_shuffle), megabatch_size)]

    last_mm = mm_megabatches[-1]
    last_lang = lang_megabatches[-1]
    additional_batch = last_mm + last_lang
    megabatches = mm_megabatches[:-1] + lang_megabatches[:-1]
    megabatch_indices = torch.randperm(len(megabatches), generator=generator)
    megabatches = [megabatches[i] for i in megabatch_indices]

    if len(additional_batch) >= megabatch_size:
        megabatches = [additional_batch[:megabatch_size]] + megabatches
        additional_batch = additional_batch[megabatch_size:]

    if len(additional_batch) > 0:
        megabatches.append(additional_batch)

    return [i for megabatch in megabatches for i in megabatch]


def get_length_grouped_indices(lengths, batch_size, world_size, generator=None, merge=True):
    # We need to use torch for the random part as a distributed sampler will set the random seed for torch.
    indices = torch.randperm(len(lengths), generator=generator)
    megabatch_size = world_size * batch_size
    megabatches = [indices[i : i + megabatch_size].tolist() for i in range(0, len(lengths), megabatch_size)]
    megabatches = [sorted(megabatch, key=lambda i: lengths[i], reverse=True) for megabatch in megabatches]
    megabatches = [split_to_even_chunks(megabatch, lengths, world_size) for megabatch in megabatches]

    return [i for megabatch in megabatches for batch in megabatch for i in batch]


class LengthGroupedSampler(Sampler):
    r"""
    Sampler that samples indices in a way that groups together features of the dataset of roughly the same length while
    keeping a bit of randomness.
    """

    def __init__(
        self,
        batch_size: int,
        world_size: int,
        lengths: Optional[List[int]] = None,
        generator=None,
        group_by_modality: bool = False,
    ):
        if lengths is None:
            raise ValueError("Lengths must be provided.")

        self.batch_size = batch_size
        self.world_size = world_size
        self.lengths = lengths
        self.generator = generator
        self.group_by_modality = group_by_modality

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

    def __iter__(self):
        # import ipdb;ipdb.set_trace()
        if self.group_by_modality:
            indices = get_modality_length_grouped_indices(self.lengths, self.batch_size, self.world_size, generator=self.generator)
        else:
            indices = get_length_grouped_indices(self.lengths, self.batch_size, self.world_size, generator=self.generator)
        return iter(indices)

from torch.utils.data import DataLoader, Dataset, IterableDataset, RandomSampler, SequentialSampler
import random


class BatchDatasetSampler(Sampler):
    def __init__(self, data_source, datasetA_size, batch_size, shuffle=True, force_dataA_first=0):
        """
        Args:
            data_source (Dataset): 合并后的数据集（ConcatDataset）
            datasetA_size (int): 数据集 A 的大小
            batch_size (int): 每个 batch 的大小
            shuffle (bool): 是否打乱数据
            force_dataA_first (int): 前 N 个 batch 强制从 datasetA 采样
        """
        self.data_source = data_source
        self.datasetA_size = datasetA_size
        self.datasetB_size = len(data_source) - datasetA_size
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.force_dataA_first = force_dataA_first  # 新增参数

    def __iter__(self):
        # 初始化数据集索引
        datasetA_indices = list(range(self.datasetA_size))
        datasetB_indices = list(range(self.datasetA_size, len(self.data_source)))

        if self.shuffle:
            random.shuffle(datasetA_indices)
            random.shuffle(datasetB_indices)

        batches = []
        batch_count = 0  # 记录当前是第几个 batch

        while datasetA_indices or datasetB_indices:
            # 前 200 个 batch 强制从 datasetA 采样
            if batch_count < self.force_dataA_first:
                use_A = True
            else:
                # 之后随机选择
                use_A = random.choice([True, False]) if datasetA_indices and datasetB_indices else bool(datasetA_indices)

            if use_A:
                batch_indices = datasetA_indices[:self.batch_size] 
                datasetA_indices = datasetA_indices[self.batch_size:]
            else:
                batch_indices = datasetB_indices[:self.batch_size]
                datasetB_indices = datasetB_indices[self.batch_size:]

            batches.append(batch_indices)
            batch_count += 1  # 更新 batch 计数
        if self.shuffle:
            random_batch=batches[self.force_dataA_first:]
            random.shuffle(random_batch)
            batches[self.force_dataA_first:]=random_batch

        for batch in batches:
            # print(self.datasetA_size, len(self.data_source), batch)
            yield batch

    def __len__(self):
        return (len(self.data_source) + self.batch_size - 1) // self.batch_size
class OspreyTrainer(Trainer):

    def get_train_dataloader(self) -> Optional[torch.utils.data.Sampler]:
        batch_sampler = BatchDatasetSampler(
                    data_source=self.train_dataset,
                    datasetA_size=self.data_len[0],  # 传入 datasetA 的大小
                    batch_size=self.args.train_batch_size,
                    shuffle=True
                )
        train_dataset = self.train_dataset
        data_collator = self.data_collator
        data_collator = self._get_collator_with_removed_columns(data_collator, description="training")
        dataloader_params = {
            # "batch_size": self._train_batch_size, 
            # "batch_size": 1, #batch sample 只能是 1
            "collate_fn": data_collator,
            "num_workers": self.args.dataloader_num_workers,
            "pin_memory": self.args.dataloader_pin_memory,
            "persistent_workers": self.args.dataloader_persistent_workers,
        }
        if not isinstance(train_dataset, torch.utils.data.IterableDataset):
            # dataloader_params["sampler"] = RandomSampler(self.train_dataset)
            # dataloader_params["drop_last"] = self.args.dataloader_drop_last
            dataloader_params["worker_init_fn"] = seed_worker
            dataloader_params["prefetch_factor"] = self.args.dataloader_prefetch_factor
            dataloader_params["shuffle"]=None       # 必须设为 None
            dataloader_params["sampler"]=None       # 必须设为 None
            dataloader_params["drop_last"]=None     # 必须设为 None
            dataloader_params["sampler"] = None
            dataloader_params["batch_sampler"]=batch_sampler
        return self.accelerator.prepare(DataLoader(train_dataset, **dataloader_params))
        #117266 len
    # def _get_train_sampler(self) -> Optional[torch.utils.data.Sampler]:
    #     if self.train_dataset is None or not has_length(self.train_dataset):
    #         return None
    #     if self.args.group_by_modality_length: #False
    #         lengths = self.train_dataset.modality_lengths
    #         return LengthGroupedSampler(
    #             # self.args.train_batch_size * self.args.gradient_accumulation_steps, # TODO: seems that we should not have gradient_accumulation_steps
    #             self.args.train_batch_size,
    #             world_size=self.args.world_size,
    #             lengths=lengths,
    #             group_by_modality=True,
    #         )
    #     else:
    #         return super()._get_train_sampler() #原始的也没用到
        #117266 len
    def _save_checkpoint(self, model, trial, metrics=None):
        if getattr(self.args, 'tune_mm_mlp_adapter', False):
            from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR
            checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}"

            run_dir = self._get_output_dir(trial=trial)
            output_dir = os.path.join(run_dir, checkpoint_folder)

            # Only save Adapter
            keys_to_match = ['mm_projector', 'vision_resampler']
            if getattr(self.args, "use_im_start_end", False):
                keys_to_match.extend(['embed_tokens', 'embed_in'])

            weight_to_save = get_mm_adapter_state_maybe_zero_3(self.model.named_parameters(), keys_to_match)

            if self.args.local_rank == 0 or self.args.local_rank == -1:
                self.model.config.save_pretrained(output_dir)
                torch.save(weight_to_save, os.path.join(output_dir, f'mm_projector.bin'))
        else:
            import transformers
            from packaging import version
            current_version = transformers.__version__
            required_version = "4.51.0"
            if version.parse(current_version) <= version.parse(required_version):
                super(OspreyTrainer, self)._save_checkpoint(model, trial, metrics)
                print(f"⚠️ 当前 transformers 版本 ({current_version}) 低于要求的最低版本 ({required_version})")
            else:
                super(OspreyTrainer, self)._save_checkpoint(model, trial)
                print(f"版本满足要求 ({current_version} >= {required_version})")
            ####save SAM####
            try:
                # 保存模型
                sam_path=self._get_output_dir(trial=trial)+'/SAM.pth'
                state_dict = model.module.visual_model.state_dict()
                state_dict = {k: v.detach().cpu() for k, v in state_dict.items()}
                torch.save(state_dict, sam_path)
                print(f'SAM模型已成功保存到: {sam_path}')
                projector_path=self._get_output_dir(trial=trial)+'/projector.pth'
                state_dict=model.module.text_hidden_fcs.state_dict()
                state_dict = {k: v.detach().cpu() for k, v in state_dict.items()}
                torch.save(state_dict, projector_path)
                print(f'SAM Projector已成功保存到: {projector_path}')
                return True
            except Exception as e:
                print(f'保存模型时出错')
                return False

    def _save(self, output_dir: Optional[str] = None, state_dict=None):
        if getattr(self.args, 'tune_mm_mlp_adapter', False):
            pass
        else:
            super(OspreyTrainer, self)._save(output_dir, state_dict)
