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

import pytorch_lightning as pl
import torch

from lit_module import LitModule
from tokenization_qwen import QWenTokenizer
from logger import TBLogger

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 = 2
val_batch_size = 1
num_proc = 8
max_epochs = 1000
strategy = "auto"
resume_from_ckpt_path = None
seed = 42

vocab_size = 256
level_ratio = 6
level = 4
dataset_level = 1

hidden_size = 1024  # 128 1024 2048  32
num_attention_heads = 16  # 8 8 16
num_hidden_layers = 3  # 6 12 24  3

mask_level = None
mask_idx = None

# name = "vocab_ratio_level_data_hidden_head_layer"
name = "rank"
ver = f"{vocab_size}" + "_" + f"{level_ratio}" + "_" + f"{level}" + "_" + f"{dataset_level}"
ver = ver + "_" + f"{hidden_size}" + "_" + f"{num_attention_heads}" + "_" + f"{num_hidden_layers}"
ver = ver + "_" + f"{mask_level}" + "_" + f"{mask_idx}"

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

    config = ModelConfig()
    config.vocab_size = vocab_size
    config.hidden_size = hidden_size
    config.num_hidden_layers = num_hidden_layers
    config.num_attention_heads = num_attention_heads

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

    start = vocab_size * (level_ratio**level)
    size = vocab_size * (level_ratio**dataset_level)
    raw_dataset = MeaningDataset(start, start + size, size, vocab_size, level_ratio)
    raw_dataset.set_mask(mask_level, mask_idx)
    train_dataset, val_dataset = raw_dataset.split(0.9)
    train_dataloader = BatchGroupMeaningDataloader(train_dataset, train_batch_size)
    val_dataloader = BatchGroupMeaningDataloader(val_dataset, val_batch_size)

    # for i in range(len(train_dataloader)):
    #     print(train_dataloader.print_mapping(i))

    torch.set_float32_matmul_precision("medium")
    lit_trainer = pl.Trainer(
        accelerator="cuda",
        precision=precision,
        logger=TBLogger("./log/", name=name, version=ver, 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,
    )