import argparse
from functools import partial
from itertools import chain
from typing import Dict, Tuple

import datasets
import pytorch_lightning as pl
import torch
from torch.utils.data import ConcatDataset, DataLoader, Dataset, random_split, Subset

from transformers import (
    BatchEncoding,
    DefaultDataCollator,
    PreTrainedTokenizer,
    set_seed,
)
from modelscope import snapshot_download
from lit_module import LitModule
from tokenization_qwen import QWenTokenizer
from logger import TBLogger

model_name = "qwen/Qwen-1_8B-Chat"
learning_rate = 0.0001
use_tril_attention_mask = None
precision = "32-true"  # "precision:bf16-mixed,16-mixed,32-true"
tokenizer_name_or_path = None
train_batch_size = 256
val_batch_size = 16
num_proc = 8
max_epochs = 1000
strategy = "auto"
resume_from_ckpt_path = None
seed = 42
vocab_size = 4096


class SpecialDataset(Dataset):
    def __init__(self, start=1, end=16, size=32768):  # 1048576 32768
        self.size = size
        self.features = []
        a = torch.randint(start, end, [size])
        b = torch.randint(start, end, [size])
        c = torch.randint(start, end, [size])
        d = torch.randint(start, end, [size])
        z = torch.zeros([size]).long()
        # self.data = torch.stack([a, b, a + b, a + b, a + b * 2]).permute(1, 0)
        # self.data = torch.stack([a, b, a, a + b / 4]).permute(1, 0).long()
        # self.data = torch.stack([a, a + 1, a + 2]).permute(1, 0).long()
        self.data = torch.stack([a, b, a]).permute(1, 0).long()
        # self.data = torch.stack([a, b, a, a + a / 8, a + a / 4, a + a / 2, a + a]).permute(1, 0).long()

        # a = torch.randint(start, end, [size])
        # self.data = torch.stack([a, a, a + a]).permute(1, 0)  # accuracy=0.5
        # self.data = torch.stack([a, a + a, a]).permute(1, 0)  # accuracy=1
        # 只能有一种算法，而且第一个值不能用于训练
        # 太陡峭的过度导致难以拟合
        # 搜索空间太大，难以拟合

    def __len__(self):
        return self.size

    def __getitem__(self, idx):
        output = {}
        data = self.data[idx]
        output["input_ids"] = data
        output["labels"] = data.clone()
        # output["labels"][:2] = 0
        # output["labels"][:2] = vocab_size
        output["token_type_ids"] = torch.zeros(data.shape)
        return output


if __name__ == "__main__":
    if tokenizer_name_or_path is None:
        tokenizer_name_or_path = model_name

    set_seed(seed)

    # lightning module
    model_dir = snapshot_download(model_name)
    lit_module = LitModule(model_dir, learning_rate, use_tril_attention_mask)

    tokenizer = QWenTokenizer("./wit_b64.tiktoken", "./wit_char.tiktoken")

    train_dataset, val_dataset = random_split(SpecialDataset(), [0.95, 0.05])

    train_dataloader = DataLoader(
        train_dataset,
        batch_size=train_batch_size,
        num_workers=num_proc,
        collate_fn=DefaultDataCollator(),
        persistent_workers=True,
        shuffle=True,
    )
    val_dataloader = DataLoader(
        val_dataset,
        batch_size=val_batch_size,
        num_workers=num_proc,
        collate_fn=DefaultDataCollator(),
        persistent_workers=True,
    )

    torch.set_float32_matmul_precision("medium")
    lit_trainer = pl.Trainer(
        accelerator="gpu",
        precision=precision,
        logger=TBLogger("./", default_hp_metric=False),
        strategy=strategy,
        max_epochs=max_epochs,
    )
    lit_trainer.fit(
        lit_module,
        train_dataloaders=train_dataloader,
        val_dataloaders=val_dataloader,
        ckpt_path=resume_from_ckpt_path,
    )