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

from special_dataset import SpecialDataset
from meaning_dataset import MeaningDataset, BatchGroupMeaningDataloader
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 = 1
val_batch_size = 2
num_proc = 8
max_epochs = 10
strategy = "auto"
resume_from_ckpt_path = None
seed = 42
vocab_size = 16


if __name__ == "__main__":
    torch.manual_seed(seed)

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

    lit_module = LitModule(pretrain_model_name, learning_rate, config, use_tril_attention_mask)
    tokenizer = QWenTokenizer("./wit_b64.tiktoken", "./wit_char.tiktoken")

    level_ratio = 2
    start = vocab_size * level_ratio * level_ratio
    end = start * level_ratio
    size = end * level_ratio
    size = 1024
    raw_dataset = MeaningDataset(start, end, size, vocab_size, level_ratio)
    train_dataset, val_dataset = raw_dataset.Split(0.95)

    train_dataloader = BatchGroupMeaningDataloader(train_dataset, train_batch_size)
    val_dataloader = BatchGroupMeaningDataloader(val_dataset, val_batch_size)

    it = iter(val_dataloader)
    batch = next(it)
    b, l = lit_module.llm(**batch, return_dict=True)
    print("b ")
    print(b.detach().cpu().numpy())

    # batch["input_ids"] = batch["input_ids"][0:1, :]
    batch["input_ids"] = batch["input_ids"][1:2, :]
    batch["labels"] = batch["labels"][1:2, :]
    s, l = lit_module.llm(**batch, return_dict=True)
    print("s ")
    print(s.detach().cpu().numpy())

    print("data samples:")