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
from transformers import (
    BatchEncoding,
    DefaultDataCollator,
    PreTrainedTokenizer,
    set_seed,
)
from modelscope import snapshot_download
from lit_module import LitModule
from tokenization_qwen import QWenTokenizer

model_name = "qwen/Qwen-1_8B-Chat"
learning_rate = 0.0001
use_tril_attention_mask = None
precision = "16-mixed"  # "precision:bf16-mixed,16-mixed,32-true"
tokenizer_name_or_path = None
dataset_name = "/home/colin/develop/dataset/liwu/MNBVC/wiki/20230197/0.jsonl.gz"
dataset_name = "/home/colin/develop/dataset/liwu/MNBVC/wiki"
train_batch_size = 8
val_batch_size = 1
accumulate_grad_batches = 32
num_proc = 8
max_epochs = None
strategy = "fsdp"
resume_from_ckpt_path = None
seed = 42


def split_raw_dataset(
    raw_dataset: datasets.DatasetDict,
) -> Tuple[datasets.Dataset, datasets.Dataset]:
    if "validation" in raw_dataset:
        train_dataset, val_dataset = raw_dataset["train"], raw_dataset["validation"]
    else:
        raw_dataset = raw_dataset["train"].train_test_split(test_size=0.05, seed=seed)
        train_dataset, val_dataset = raw_dataset["train"], raw_dataset["test"]
    return train_dataset, val_dataset


def process_dataset(dataset: datasets.Dataset, tokenizer: PreTrainedTokenizer) -> datasets.Dataset:
    def group_texts(examples: Dict[str, list], block_size: int = 512) -> BatchEncoding:
        concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}
        total_length = len(concatenated_examples[list(examples.keys())[0]])
        total_length = (total_length // block_size) * block_size
        result = {
            k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
            for k, t in concatenated_examples.items()
        }
        result["labels"] = result["input_ids"].copy()
        result = BatchEncoding(result)
        return result

    def format_inputs(examples):
        p = examples["段落"]
        mergeLine = ""
        for line in p:
            mergeLine += line["内容"] + "\n"
        return {"text": mergeLine}

    def tokenize_inputs(
        examples: Dict[str, list],
        tokenizer: PreTrainedTokenizer,
        column_name: str = "text",
    ) -> BatchEncoding:
        logits = tokenizer(examples[column_name], return_attention_mask=False)
        return logits

    dataset_column_names = list(dataset.features)
    dataset = dataset.map(
        partial(format_inputs),
        batched=False,
        num_proc=num_proc,
        remove_columns=dataset_column_names,
    )
    dataset_column_names = list(dataset.features)
    dataset = dataset.map(
        partial(tokenize_inputs, tokenizer=tokenizer),
        batched=True,
        num_proc=num_proc,
        remove_columns=dataset_column_names,
    )
    dataset = dataset.map(
        partial(group_texts, block_size=tokenizer.model_max_length),
        batched=True,
        num_proc=num_proc,
    )

    return dataset


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)

    # datasets
    # tokenizer = load_tokenizer("./custom_models/gpt2")
    tokenizer = QWenTokenizer("./wit_b64.tiktoken", "./wit_char.tiktoken")
    train_dataset_list = []
    val_dataset_list = []
    for dataset_name in dataset_name:
        dataset_args = dataset_name.split(":")
        raw_dataset = datasets.load_dataset(
            "json", data_files="/home/colin/develop/dataset/liwu/MNBVC/wiki/20230197/0.jsonl.gz"
        )
        # raw_dataset = datasets.load_dataset(*dataset_args)
        train_dataset, val_dataset = split_raw_dataset(raw_dataset)
        train_dataset = process_dataset(train_dataset, tokenizer)
        val_dataset = process_dataset(val_dataset, tokenizer)
        train_dataset_list.append(train_dataset)
        val_dataset_list.append(val_dataset)
    train_dataset = ConcatDataset(train_dataset_list)
    val_dataset = ConcatDataset(val_dataset_list)

    # dataloaders
    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,
    )

    ne = next(train_dataloader._get_iterator())

    # trainer
    # apply_all_patches()
    torch.set_float32_matmul_precision("medium")
    precision = precision
    lit_trainer = pl.Trainer(
        accelerator="gpu",
        precision=precision,
        log_every_n_steps=5,
        accumulate_grad_batches=accumulate_grad_batches,
        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,
    )