from peft import (
    prepare_model_for_int8_training,
    LoraConfig,
    PeftModel,
    get_peft_model,
    get_peft_model_state_dict,
    set_peft_model_state_dict,
)
from transformers import LlamaForCausalLM, LlamaTokenizer, TrainerCallback, GenerationConfig
import os
import sys
import torch
import torch.nn as nn
import bitsandbytes as bnb
from datasets import load_dataset, Dataset
import transformers
from huggingface_hub import snapshot_download
import argparse
import warnings
from tqdm import tqdm
from functools import partial
import utils
import prompt
assert (
    "LlamaTokenizer" in transformers._import_structure["models.llama"]
), "LLaMA is now in HuggingFace's main branch.\nPlease reinstall it: pip uninstall transformers && pip install git+https://github.com/huggingface/transformers.git"

# 0. prepare args and logger
parser = argparse.ArgumentParser()
parser.add_argument("--wandb", action="store_true", default=False)
parser.add_argument("--prompt_type", type=str, default="chat")
parser.add_argument("--data_path", type=str, default="merge.json")
parser.add_argument("--output_path", type=str, default="lora-Vicuna")
parser.add_argument("--model_path", type=str, default="decapoda-research/llama-7b-hf")
parser.add_argument("--num_epoch", type=int, default=6)
parser.add_argument("--micro_batch", type=int, default=4)
parser.add_argument("--total_batch", type=int, default=128)
parser.add_argument("--log_steps", type=int, default=100)
parser.add_argument("--eval_steps", type=int, default=200)
parser.add_argument("--save_steps", type=int, default=200)
parser.add_argument("--warmup_ratio", type=float, default=0.05)
parser.add_argument("--test_size", type=int, default=200)
parser.add_argument("--resume_from_checkpoint", type=str, default=None)
parser.add_argument("--lora_remote_checkpoint", type=str, default=None)
parser.add_argument("--ignore_data_skip", type=bool, default=False)
args = parser.parse_args()
if not args.wandb:
    os.environ["WANDB_MODE"] = "disable"
MICRO_BATCH_SIZE = args.micro_batch  # this could actually be 5 but i like powers of 2
BATCH_SIZE = args.total_batch
MAX_STEPS = None
GRADIENT_ACCUMULATION_STEPS = BATCH_SIZE // MICRO_BATCH_SIZE
EPOCHS = args.num_epoch 
LEARNING_RATE = 3e-4  # the Karpathy constant
CUTOFF_LEN = 2048  
LORA_R = 8
LORA_ALPHA = 16
LORA_DROPOUT = 0.05
USE_8bit = True
VAL_SET_SIZE = args.test_size  # 2000
TARGET_MODULES = [
    "q_proj",
    "v_proj",
    "k_proj",
    "o_proj",
    "down_proj",
    "gate_proj",
    "up_proj",
]
DATA_PATH = args.data_path  
OUTPUT_DIR = args.output_path  # "lora-Vicuna"

device_map = "auto"
world_size = int(os.environ.get("WORLD_SIZE", 1))
print(world_size)
ddp = world_size != 1
if ddp:
    print('222')
    device_map = {"": int(os.environ.get("LOCAL_RANK") or 0)}
    print(GRADIENT_ACCUMULATION_STEPS)
    print(world_size)
    # GRADIENT_ACCUMULATION_STEPS = GRADIENT_ACCUMULATION_STEPS // world_size
    print(1111)
# we must make sure batch_size and gradient_accumulation_steps not changed for resuming training.
if args.resume_from_checkpoint:
    old_args_path = os.path.join(args.resume_from_checkpoint, 'training_args.bin')
    if os.path.exists(old_args_path):
        old_args = torch.load(old_args_path)
        print(MICRO_BATCH_SIZE)
        print(old_args.per_device_train_batch_size)
        print(GRADIENT_ACCUMULATION_STEPS)
        print(old_args.gradient_accumulation_steps)
        if MICRO_BATCH_SIZE != old_args.per_device_train_batch_size:
            raise Exception(
                f'current micro batch size {MICRO_BATCH_SIZE} is not equal to the old {old_args.per_device_train_batch_size},'
                ' This will cause the trainer skips wrong epochs or steps.'
                f'please change your micro batch size to {old_args.per_device_train_batch_size}'
                ' or cancel resuming your training'
                )
        if GRADIENT_ACCUMULATION_STEPS != old_args.gradient_accumulation_steps:
            raise Exception(
                f'current total batch {BATCH_SIZE} is not equal to the old {old_args.gradient_accumulation_steps*old_args.per_device_train_batch_size},'
                ' This will cause the trainer skips wrong epochs or steps.'
                f'please change your total batch size to {old_args.gradient_accumulation_steps*old_args.per_device_train_batch_size}'    
                ' or cancel resuming your training'
            )
    else:
        raise Exception(f'{old_args_path} is not exist!')
    # checkpoint = os.path.join(args.resume_from_checkpoint, 'pytorch_model.bin')

logger = utils.set_file_logger(__name__,OUTPUT_DIR)
# 1. load dataset
logger.info(f'>>> processing data from {DATA_PATH}')
logger.info(f'>>> using {args}')

train_tokenizer = LlamaTokenizer.from_pretrained(args.model_path, add_eos_token=True)
# assert train_tokenizer.eos_token_id == 2, "Tokenizer eos is wrong!!!"
# unk. we want this to be different from the eos token
train_tokenizer.pad_token_id = 0  
# cannot use eos in generation!
# tokenizer.padding_side = "left"  # Allow batched inference
test_tokenizer = LlamaTokenizer.from_pretrained(args.model_path)
if args.prompt_type == 'instruct':
    PROMPT = prompt.instruct_prompt(train_tokenizer, CUTOFF_LEN)
elif args.prompt_type == 'chat':
    PROMPT = prompt.chat_prompt(train_tokenizer,CUTOFF_LEN)
else:
    raise Exception('not support')
# check tokenizer
data = load_dataset('json', data_files=DATA_PATH)
import random;start = random.randint(1, 100)
examples = Dataset.from_dict(data['train'][start:start+5]).map(PROMPT.preprocess_train)
for example in examples:
    logger.info(f'>>> using prompt {args.prompt_type}, prompt example:\n { train_tokenizer.decode(example["input_ids"]) }')
    logger.info(f'>>> tokenizer labels: { train_tokenizer.decode([ 0 if l==-100 else l for l in example["labels"]])}')
    logger.info(f'>>> tokenizer example: { example["input_ids"][:10] }...{ example["input_ids"][-10:]}')
# 2. load model and checkpoints
logger.info(f'>>> load model from {args.model_path}')

# if USE_8bit is True:
#     assert bnb.__version__ >= '0.37.2', "Please downgrade bitsandbytes's version, for example: pip install bitsandbytes==0.37.2"
model = LlamaForCausalLM.from_pretrained(
    args.model_path,
    load_in_8bit=USE_8bit,
    device_map=device_map,
    torch_dtype=torch.float16,
)
if USE_8bit is True:
    model = prepare_model_for_int8_training(model)
config = LoraConfig(
    r=LORA_R,
    lora_alpha=LORA_ALPHA,
    target_modules=TARGET_MODULES,
    lora_dropout=LORA_DROPOUT,
    bias="none",
    task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)
if args.resume_from_checkpoint:
    checkpoint_name = os.path.join(args.resume_from_checkpoint, "pytorch_model.bin")
    # adapter_model.bin
    if not os.path.exists(checkpoint_name):
        pytorch_bin_path = checkpoint_name
        checkpoint_name = os.path.join(args.resume_from_checkpoint, "adapter_model.bin")
        if os.path.exists(checkpoint_name):
            os.rename(checkpoint_name, pytorch_bin_path)
            logger.warning("The file name of the lora checkpoint'adapter_model.bin' is replaced with 'pytorch_model.bin'")
        else:
            args.resume_from_checkpoint = None  # So the trainer won't try loading its state
    print(checkpoint_name)
    # pytorch_model.bin
    if os.path.exists(checkpoint_name):
        logger.info(f'>>> load lora from {checkpoint_name}')
        adapters_weights = torch.load(checkpoint_name)
        set_peft_model_state_dict(model, adapters_weights)
    else:
        raise Exception(f"Checkpoint {checkpoint_name} not found with resume_from_checkpoint=True!")

trainable_params = 0
all_param = 0
for _, param in model.named_parameters():
    num_params = param.numel()
    # if using DS Zero 3 and the weights are initialized empty
    if num_params == 0 and hasattr(param, "ds_numel"):
        num_params = param.ds_numel
    all_param += num_params
    if param.requires_grad:
        trainable_params += num_params
logger.info(f">>> trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}")

# 3. speedup dataset processing by multi-process
# num_proc = (os.cpu_count())
if VAL_SET_SIZE > 0:
    train_val = data["train"].train_test_split(test_size=VAL_SET_SIZE, shuffle=True, seed=42)
    train_data = train_val["train"].shuffle().map(PROMPT.preprocess_train)
    val_data = train_val["test"].shuffle().map(PROMPT.preprocess_train)
else:
    train_data = data["train"].shuffle().map(PROMPT.preprocess_train)
    val_data = None
now_max_steps = max((len(data["train"]) - VAL_SET_SIZE) // BATCH_SIZE * EPOCHS, EPOCHS)
if args.resume_from_checkpoint:
    # the trainer will ignore the state max_steps and caculate max_steps based on epochs,
    # so we mannally set the args.max_step to override it. 
    if args.lora_remote_checkpoint is not None:
        snapshot_download(repo_id=args.lora_remote_checkpoint, allow_patterns=["*.pt", "*.bin", "*.json"], local_dir=args.resume_from_checkpoint)
    train_state_path = os.path.join(args.resume_from_checkpoint, "trainer_state.json")
    if os.path.exists(train_state_path):
        import json
        base_train_args = json.load(open(train_state_path, 'r'))
        base_max_steps = base_train_args["max_steps"]
        resume_scale = base_max_steps / now_max_steps
        if base_max_steps > now_max_steps:
            logger.warning(f"epoch {EPOCHS}:{MAX_STEPS} replace to the base_max_steps {base_max_steps}")
            EPOCHS = None
            MAX_STEPS = base_max_steps
        else:
            MAX_STEPS = now_max_steps
    assert MAX_STEPS is not None
else:
    MAX_STEPS = now_max_steps

# 4. start training
class CustomCallback(TrainerCallback):
    
    def __init__(self, trainer) -> None:
        super().__init__()
        self.trainer = trainer
        self.generation_config = GenerationConfig(
            temperature=1.0,
            top_p=0.75,
            top_k=40,
            num_beams=2,
            bos_token_id=train_tokenizer.bos_token_id,
            eos_token_id=train_tokenizer.eos_token_id,
            pad_token_id=train_tokenizer.pad_token_id,
            max_new_tokens=1024, # max_length=max_new_tokens+input_sequence
            min_new_tokens=1, # min_length=min_new_tokens+input_sequence
            bad_words_ids=test_tokenizer(['\n\nUser:','\n\nAssistant:'], add_special_tokens=False).input_ids
        )
        self.repetition_penalty=1.3
        self.logger = utils.set_file_logger('transformers.trainer', trainer.args.output_dir)

    def on_log(self, args, state, control, logs, **kwargs):
        logger.info(logs)

trainer = transformers.Trainer(
    model=model,
    train_dataset=train_data,
    eval_dataset=val_data,
    args=transformers.TrainingArguments(
        per_device_train_batch_size=MICRO_BATCH_SIZE,
        gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS,
        warmup_ratio=args.warmup_ratio,
        num_train_epochs=EPOCHS,
        max_steps=MAX_STEPS,
        learning_rate=LEARNING_RATE,
        fp16=True,
        logging_steps=args.log_steps,
        logging_first_step=True, # convenient
        evaluation_strategy="steps" if VAL_SET_SIZE > 0 else "no",
        save_strategy="steps",
        eval_steps=args.eval_steps if VAL_SET_SIZE > 0 else None,
        save_steps=args.save_steps,
        output_dir=OUTPUT_DIR,
        load_best_model_at_end=True if VAL_SET_SIZE > 0 else False,
        ddp_find_unused_parameters=False if ddp else None,
        report_to="wandb" if args.wandb else [],
        ignore_data_skip=args.ignore_data_skip,
    ),
    data_collator=PROMPT.data_collator()
)
trainer.add_callback(CustomCallback(trainer))
model.config.use_cache = False

old_state_dict = model.state_dict
model.state_dict = (
    lambda self, *_, **__: get_peft_model_state_dict(self, old_state_dict())
).__get__(model, type(model))

if torch.__version__ >= "2" and sys.platform != "win32":
    model = torch.compile(model)

trainer.train(resume_from_checkpoint=args.resume_from_checkpoint)
model.save_pretrained(OUTPUT_DIR)