import os
os.environ["TOKENIZERS_PARALLELISM"] = "false"

import datasets
import torch
import numpy as np
import logging
import deepspeed
from transformers import AutoModelForCausalLM, TrainingArguments, Trainer, DataCollatorForLanguageModeling, AutoTokenizer
from trl import DPOConfig, DPOTrainer
from peft import LoraConfig, get_peft_model, PeftModel
from tuningtron.models import ModelsFactory


logger = logging.getLogger(__name__)


class Tuner:
    def __init__(self, base_model_id, enable_deepspeed=True, enable_cpu=False, enable_offload_optimizer=True):
        self.print_cuda_info()

        self.enable_cpu = enable_cpu
        self.base_model_id = base_model_id
        self.model_config = ModelsFactory().get_model_config(base_model_id)
        self.tokenizer = self.model_config.tokenizer
        self.device_map = "auto"
        self.deepspeed = None
        self.optim = "adamw_8bit"
        self.bf16 = False
        self.fp16 = False
        self.attn_implementation = None
        self.dtype = torch.bfloat16

        if enable_cpu or not torch.cuda.is_available():
            self.optim = "adamw_torch"
            self.bf16 = True
            self.device_map = "cpu"
            self.dtype = torch.bfloat16
            if self.model_config.config.model_type.startswith("gemma"):
                self.attn_implementation = "eager"
        elif torch.cuda.is_available():
            if enable_deepspeed:
                deepspeed.init_distributed()
                self.device_map = None
                self.deepspeed = self.get_deepspeed_config(enable_offload_optimizer)
                if enable_offload_optimizer:
                    self.optim = "adamw_torch"
                logger.info("deepspeed: enabled")

            if torch.cuda.get_device_capability()[0] >= 8:
                self.bf16 = True
                self.attn_implementation = "flash_attention_2"
                if self.model_config.config.model_type.startswith("gemma"):
                    self.attn_implementation = "eager"
            else:
                self.fp16 = True

    def apply_template(self, record):
        return self.model_config.apply_chat_template(record)

    def map_func(self, record):
        return self.tokenizer(self.apply_template(record), truncation=True, max_length=self.max_len, padding="max_length")

    def filter_func(self, record):
        return len(self.tokenizer(self.apply_template(record))["input_ids"]) <= self.max_len

    def sft(self,
            dataset,
            adapter_name,
            do_eval=False,
            max_len_percentile=100,
            lora_rank=8,
            lora_alpha=None,
            num_train_epochs=1,
            backpropagation_batch_size=1,
            gradient_accum_steps=1,
            learning_rate=1e-5):
        dataset = datasets.load_dataset(dataset, split="train")

        inputs = [self.tokenizer(self.apply_template(record))["input_ids"] for record in dataset]
        target_lenghts = [len(x) for x in inputs]
        self.max_len = int(np.percentile(target_lenghts, max_len_percentile))
        logger.info(f"Dataset max_len detected: {self.max_len}")

        logger.info("DS before mapping and filtering:")
        logger.info(str(dataset))

        dataset = dataset.filter(lambda record: self.filter_func(record))
        dataset = dataset.map(self.map_func)

        logger.info("DS after mapping and filtering:")
        logger.info(str(dataset))

        self.print_ds_example_row(dataset)

        dataset = dataset.remove_columns(["instruct", "input", "output"])

        train_dataset, eval_dataset = self.prepare_datasets(dataset, do_eval)

        args = TrainingArguments(**self.prepare_args(num_train_epochs, learning_rate, backpropagation_batch_size, gradient_accum_steps))
        logger.info(str(args))

        peft_model = get_peft_model(self.load_base_model(), self.get_lora_config(lora_rank, lora_alpha))
        logger.info(str(peft_model.get_model_status()))

        trainer = Trainer(model=peft_model,
                          train_dataset=train_dataset,
                          eval_dataset=eval_dataset,
                          data_collator=DataCollatorForLanguageModeling(tokenizer=self.tokenizer, mlm=False),
                          args=args)
        trainer.train()
        trainer.save_model(adapter_name)

    def dpo(self,
            dataset,
            adapter_name,
            do_eval=False,
            lora_rank=8,
            lora_alpha=None,
            num_train_epochs=1,
            backpropagation_batch_size=1,
            gradient_accum_steps=1,
            learning_rate=1e-5):
        dataset = datasets.load_dataset(dataset, split="train")

        train_dataset, eval_dataset = self.prepare_datasets(dataset, do_eval)

        logger.info("Dataset example row after appy chat template:")
        logger.info("Chosen ------>")
        logger.info(self.tokenizer.apply_chat_template(train_dataset["chosen"][0], tokenize=False))
        logger.info("Rejected ------>")
        logger.info(self.tokenizer.apply_chat_template(train_dataset["rejected"][0], tokenize=False))
        logger.info("---------------------------------------------")
        logger.info("Dataset example row after tokenize:")
        logger.info("Chosen ------>")
        logger.info(self.tokenizer.apply_chat_template(train_dataset["chosen"][0]))
        logger.info("Rejected ------>")
        logger.info(self.tokenizer.apply_chat_template(train_dataset["rejected"][0]))

        args = DPOConfig(**self.prepare_args(num_train_epochs, learning_rate, backpropagation_batch_size, gradient_accum_steps))
        logger.info(args)

        peft_model = get_peft_model(self.load_base_model(), self.get_lora_config(lora_rank, lora_alpha))
        logger.info(peft_model.get_model_status())

        trainer = DPOTrainer(model=peft_model,
                             train_dataset=train_dataset,
                             eval_dataset=eval_dataset,
                             processing_class=self.tokenizer,
                             args=args)
        trainer.train()
        trainer.save_model(adapter_name)

    def prepare_datasets(self, dataset, do_eval):
        eval_dataset = None
        self.eval_strategy = "no"
        self.eval_steps = None

        if do_eval:
            dataset = dataset.train_test_split(test_size=0.1)
            train_dataset = dataset["train"]
            eval_dataset = dataset["test"]
            self.eval_strategy = "steps"
            self.eval_steps = 0.1
            logger.info("Eval dataset:")
            logger.info(eval_dataset)
        else:
            train_dataset = dataset
        logger.info("Train dataset:")
        logger.info(train_dataset)

        return train_dataset, eval_dataset

    def print_ds_example_row(self, dataset):
        tokens = dataset["input_ids"][0]
        logger.info("Dataset text row:")
        logger.info("---------------------------------------------")
        logger.info(self.tokenizer.decode(tokens))
        logger.info("---------------------------------------------")
        logger.info("Dataset tokens row:")
        logger.info("---------------------------------------------")
        logger.info(str(tokens))
        logger.info("---------------------------------------------")

    def prepare_args(self, num_train_epochs, learning_rate, batch_size, gradient_accum_steps, lr_scheduler_type="linear"):
        return {
            "output_dir": ".",
            "num_train_epochs": num_train_epochs,
            "logging_steps": 1,
            "eval_strategy": self.eval_strategy,
            "eval_steps": self.eval_steps,
            "gradient_checkpointing": True,
            "save_strategy": "no",
            "bf16": self.bf16,
            "fp16": self.fp16,
            "use_cpu": self.enable_cpu,
            "optim": self.optim,
            "weight_decay": 0.001,
            "learning_rate": learning_rate,
            "lr_scheduler_type": lr_scheduler_type,
            "warmup_ratio": 0.1,
            "per_device_train_batch_size": batch_size,
            "per_device_eval_batch_size": batch_size,
            "gradient_accumulation_steps": gradient_accum_steps,
            "eval_accumulation_steps": 1,
            "deepspeed": self.deepspeed
        }

    def get_lora_config(self, rank, lora_alpha, rslora=False):
        lora_alpha = lora_alpha if lora_alpha else rank
        config = LoraConfig(r=rank, lora_alpha=lora_alpha, use_rslora=rslora, target_modules=self.model_config.target_modules, lora_dropout=0.05, task_type="CAUSAL_LM")
        logger.info("Lora config:" + str(config))
        return config

    def print_cuda_info(self):
        visible_devices = os.environ['CUDA_VISIBLE_DEVICES']
        logger.info(f"visible_devices: {visible_devices}")
        logger.info("---------------------------------------------------")
        logger.info("CUDA Devices:")
        for i in range(0, torch.cuda.device_count()):
            logger.info("GPU: " + str(i))
            logger.info("Total GPU Memory: " + str(torch.cuda.get_device_properties(i).total_memory))
            logger.info("Reserved GPU Memory: " + str(torch.cuda.memory_reserved(i)))
            logger.info("Allocated GPU Memory: " + str(torch.cuda.memory_allocated(i)))
            logger.info("---------------------------------------------------")

    def merge(self, merged_name, first_adapter):
        base_model = self.load_base_model(False)

        peft_model = PeftModel.from_pretrained(base_model, first_adapter, torch_dtype=torch.bfloat16)
        logger.info(f"Merging adapter: {first_adapter} -> {merged_name}")
        merged_model = peft_model.merge_and_unload()
        merged_model.save_pretrained(merged_name)
        # get original tokenizer for save
        tmp_tokenizer = AutoTokenizer.from_pretrained(self.base_model_id)
        tmp_tokenizer.save_pretrained(merged_name)
        try:
            tmp_tokenizer.save_vocabulary(merged_name)
        except:
            pass

    def load_base_model(self, gradient_checkpointing=True):
        self.base_model = AutoModelForCausalLM.from_pretrained(self.base_model_id,
                                                               torch_dtype=self.dtype,
                                                               attn_implementation=self.attn_implementation,
                                                               device_map=self.device_map)
        logger.info(self.base_model)

        if gradient_checkpointing:
            self.base_model.gradient_checkpointing_enable()
        else:
            self.base_model.gradient_checkpointing_disable()
        return self.base_model

    def get_deepspeed_config(self, enable_offload_optimizer=True):
        cfg = {
            "zero_force_ds_cpu_optimizer": False,
            "bf16": {"enabled": "auto"},
            "zero_optimization": {
                "stage": 3,
                "offload_param": {"device": "cpu"},
                "overlap_comm": True,
                "reduce_bucket_size": "auto",
                "sub_group_size": 1e6,
                "stage3_max_live_parameters": 1e6,
                "stage3_max_reuse_distance": 1e6,
                "stage3_prefetch_bucket_size": "auto",
                "stage3_param_persistence_threshold": "auto",
                "stage3_gather_16bit_weights_on_model_save": True
            },
            "gradient_accumulation_steps": "auto",
            "gradient_clipping": "auto",
            "steps_per_print": "auto",
            "train_batch_size": "auto",
            "train_micro_batch_size_per_gpu": "auto"
        }

        if enable_offload_optimizer:
            cfg["zero_optimization"]["offload_optimizer"] = {"device": "cpu"}

        return cfg
