import os
import torch
import argparse
from peft import set_peft_model_state_dict
from transformers import AutoConfig
from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR
from transformers import TrainerCallback, TrainingArguments, TrainerState, TrainerControl

class SavePeftModelCallback(TrainerCallback):
    def on_save(
        self,
        args: TrainingArguments,
        state: TrainerState,
        control: TrainerControl,
        **kwargs,
    ):
        checkpoint_folder = os.path.join(args.output_dir, f"{PREFIX_CHECKPOINT_DIR}-{state.global_step}")

        kwargs["model"].save_pretrained(checkpoint_folder)

        pytorch_model_path = os.path.join(checkpoint_folder, "pytorch_model.bin")
        torch.save({}, pytorch_model_path)
        return control


class LoadBestPeftModelCallback(TrainerCallback):
    def on_train_end(
        self,
        args: TrainingArguments,
        state: TrainerState,
        control: TrainerControl,
        **kwargs,
    ):
        print(f"Loading best peft model from {state.best_model_checkpoint} (score: {state.best_metric}).")
        best_model_path = os.path.join(state.best_model_checkpoint, "adapter_model.bin")
        adapters_weights = torch.load(best_model_path)
        model = kwargs["model"]
        set_peft_model_state_dict(model, adapters_weights)
        return control
    


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_path", type=str, default="Salesforce/codegen-350M-mono")
    # parser.add_argument("--model_path", type=str, default="ammarnasr/codegen-350M-mono_the-stack-dedup_java_train_peft")
    parser.add_argument("--size_valid_set", type=int, default=100)
    parser.add_argument("--seq_length", type=int, default=512)
    parser.add_argument("--max_steps", type=int, default=10000)
    parser.add_argument("--output_dir", type=str, default="./checkpoints_full")
    parser.add_argument("--log_freq", default=5, type=int)
    parser.add_argument("--eval_freq", default=1000, type=int)
    parser.add_argument("--save_freq", default=1000, type=int)
    parser.add_argument("--ft_type", type=str, default="full")










    # parser.add_argument("--model_path", type=str, default="")
    # parser.add_argument("--model_path", type=str, default="bigcode/starcoder")
    # parser.add_argument("--model_path", type=str, default="bigcode/large-model")
    # parser.add_argument("--model_path", type=str, default="codeparrot/codeparrot-small")
    # parser.add_argument("--model_path", type=str, default="bigcode/tiny_starcoder_py")
    
    # parser.add_argument("--dataset_name", type=str, default="HuggingFaceH4/CodeAlpaca_20K")
    # parser.add_argument("--subset", type=str, default=None)
    # parser.add_argument("--input_column_name", type=str, default="prompt")

    parser.add_argument("--dataset_name", type=str, default="bigcode/the-stack-dedup")
    parser.add_argument("--subset", type=str, default="data/java")
    parser.add_argument("--input_column_name", type=str, default="content")
    
    parser.add_argument("--split", type=str, default="train")
    parser.add_argument("--streaming", action="store_true", default=True)
    parser.add_argument("--shuffle_buffer", type=int, default=5000)

    # parser.add_argument("--ft_type", type=str, default="full")


    parser.add_argument("--output_column_name", type=str, default="completion")
    parser.add_argument("--batch_size", type=int, default=1)
    parser.add_argument("--gradient_accumulation_steps", type=int, default=16)
    parser.add_argument("--eos_token_id", type=int, default=49152)

    parser.add_argument("--lora_r", type=int, default=16)
    parser.add_argument("--lora_alpha", type=int, default=32)
    parser.add_argument("--lora_dropout", type=float, default=0.05)

    parser.add_argument("--learning_rate", type=float, default=5e-6)
    parser.add_argument("--lr_scheduler_type", type=str, default="cosine")
    parser.add_argument("--num_warmup_steps", type=int, default=100)
    parser.add_argument("--weight_decay", type=float, default=0.05)

    parser.add_argument("--local_rank", type=int, default=0)
    parser.add_argument("--no_fp16", action="store_false", default=True)
    parser.add_argument("--bf16", action="store_true", default=True)
    parser.add_argument("--no_gradient_checkpointing", action="store_false", default=False)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--num_workers", type=int, default=None)

    

    return parser.parse_args()

