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 lit_module import LitModule
from tokenization_qwen import QWenTokenizer
from logger import TBLogger

from special_dataset import SpecialDataset
from meaning_dataset import MeaningDataset
from wit.configuration import ModelConfig

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


if __name__ == "__main__":

    set_seed(seed)

    config = ModelConfig()
    config.vocab_size = vocab_size
    config.hidden_size = 128  # 128 1024 2048  32
    config.num_hidden_layers = 6  # 6 12 24  3
    config.num_attention_heads = 8  # 8 8 16

    lit_module = LitModule(pretrain_model_name, learning_rate, config, use_tril_attention_mask)

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

    raw_dataset = SpecialDataset()
    # raw_dataset = MeaningDataset(start=65536, end=262133, size=32768, max_subitem=4, vocab_size=vocab_size)
    train_dataset, val_dataset = random_split(raw_dataset, [0.95, 0.05])

    it = iter(train_dataset)
    print("data samples:")
    for i in range(10):
        print(next(it)["input_ids"].numpy().tolist())

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