# Copyright 2024 HuggingFace Inc. and the LlamaFactory team.
#
# This code is inspired by the HuggingFace's transformers library.
# https://github.com/huggingface/transformers/blob/v4.40.0/examples/pytorch/summarization/run_summarization.py
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import json
import os
from typing import TYPE_CHECKING, List, Optional

from peft.config import PeftConfig
import wandb

from ...data import SFTDataCollatorWith4DAttentionMask, get_dataset, get_template_and_fix_tokenizer
from ...extras.constants import IGNORE_INDEX
from ...extras.misc import get_logits_processor, count_parameters
from ...extras.ploting import plot_loss
from ...extras import logging
from ...model import load_model, load_tokenizer, load_config
from ...model.adapter import init_adapter
from ...model.loader import _get_init_kwargs, patch_config
from ...hparams import DataArguments, FinetuningArguments, GeneratingArguments, ModelArguments
from ..trainer_utils import create_modelcard_and_push
from .metric import ComputeAccuracy, ComputeSimilarity, eval_logit_processor, ComputeChoiceAccuracy
from .trainer import CustomSeq2SeqTrainer

from transformers import set_seed, AutoModelForCausalLM
from peft import get_peft_config, get_peft_model, PrefixTuningConfig, TaskType, PeftType, PeftModel
import torch

if TYPE_CHECKING:
    from transformers import Seq2SeqTrainingArguments, TrainerCallback

    from ...hparams import DataArguments, FinetuningArguments, GeneratingArguments, ModelArguments

set_seed(0)


def prepare_model_for_eval(model, model_args):
    model.requires_grad_(False)
    for param in model.parameters():
        if param.data.dtype == torch.float32 and model_args.compute_dtype != torch.float32:
            param.data = param.data.to(model_args.compute_dtype)
    model.eval()


def init_prefix_model(
    tokenizer,
    model_args: ModelArguments,
    training_args: "Seq2SeqTrainingArguments",
    finetuning_args: "FinetuningArguments",
):
    # MeM: Initialize MeM model from backbone models
    init_kwargs = _get_init_kwargs(model_args)
    config = load_config(model_args)
    patch_config(config, tokenizer, model_args, init_kwargs, training_args.do_train)
    init_kwargs["config"] = config
    init_kwargs["pretrained_model_name_or_path"] = model_args.model_name_or_path

    model = AutoModelForCausalLM.from_pretrained(**init_kwargs)
    if model_args.adapter_name_or_path is not None:
        peft_config = PeftConfig.from_pretrained(model_args.adapter_name_or_path[0])
        model = PeftModel.from_pretrained(model, model_args.adapter_name_or_path[0])
    else:
        peft_config = PrefixTuningConfig(task_type=TaskType.CAUSAL_LM, num_virtual_tokens=30)
        model = get_peft_model(model, peft_config)

    model.print_trainable_parameters()

    if training_args.do_train:
        model.train()
    else:
        prepare_model_for_eval(model, model_args)

    return config, model


def run_prefix_sft(
    model_args: "ModelArguments",
    data_args: "DataArguments",
    training_args: "Seq2SeqTrainingArguments",
    finetuning_args: "FinetuningArguments",
    generating_args: "GeneratingArguments",
    callbacks: Optional[List["TrainerCallback"]] = None,
):
    tokenizer_module = load_tokenizer(model_args)
    tokenizer = tokenizer_module["tokenizer"]

    tokenizer.padding_side = 'left'

    template = get_template_and_fix_tokenizer(tokenizer, data_args)
    dataset_module = get_dataset(template, model_args, data_args, training_args, stage="sft", **tokenizer_module)
    config, model = init_prefix_model(tokenizer, model_args, training_args, finetuning_args)

    if getattr(model, "is_quantized", False) and not training_args.do_train:
        setattr(model, "_hf_peft_config_loaded", True)  # hack here: make model compatible with prediction

    data_collator = SFTDataCollatorWith4DAttentionMask(
        template=template,
        pad_to_multiple_of=8 if training_args.do_train else None,  # for shift short attention
        label_pad_token_id=IGNORE_INDEX if data_args.ignore_pad_token_for_loss else tokenizer.pad_token_id,
        block_diag_attn=model_args.block_diag_attn,
        attn_implementation=getattr(model.config, "_attn_implementation", None),
        compute_dtype=model_args.compute_dtype,
        **tokenizer_module,
    )

    # Override the decoding parameters of Seq2SeqTrainer
    training_args.generation_max_length = training_args.generation_max_length or data_args.cutoff_len
    training_args.generation_num_beams = data_args.eval_num_beams or training_args.generation_num_beams
    training_args.remove_unused_columns = False  # important for multimodal dataset

    effective_token_num = 0.0
    if finetuning_args.include_effective_tokens_per_second:
        for data in dataset_module["train_dataset"]:
            effective_token_num += len(data["input_ids"])

    # Metric utils
    metric_module = {}
    if training_args.predict_with_generate:
        if data_args.eval_dataset and data_args.eval_dataset[0].startswith('gsm8k'):
            from .metric import ComputeGsm8kAccuracy

            metric_module["compute_metrics"] = ComputeGsm8kAccuracy(tokenizer=tokenizer)

        # elif data_args.eval_dataset and data_args.eval_dataset[0].startswith('spider'):
        #     from .metric import ComputeSpiderAccuracy
        #     eval_data_file = os.path.join(data_args.dataset_dir, f'{data_args.eval_dataset[0]}.json')
        #     db_ids = [x['meta']['db_id'] for x in json.load(open(eval_data_file, 'r', encoding='utf-8'))]
        #     metric_module['compute_metrics'] = ComputeSpiderAccuracy(tokenizer=tokenizer, db_ids=db_ids)

        elif data_args.eval_dataset and data_args.eval_dataset[0].startswith('sql2text'):
            from .metric import ComputeSql2textAccuracy

            eval_data_file = os.path.join(data_args.dataset_dir, f'{data_args.eval_dataset[0]}.json')
            sqls = [x['instruction'].split('\n\n')[1] for x in json.load(open(eval_data_file, 'r', encoding='utf-8'))]
            metric_module['compute_metrics'] = ComputeSql2textAccuracy(tokenizer=tokenizer, sqls=sqls)

        elif data_args.eval_dataset and data_args.eval_dataset[0].startswith(
            tuple(
                [
                    'boolq.uni',
                    'piqa.uni',
                    'siqa.uni',
                    'hellaswag.uni',
                    'winogrande.uni',
                    'arce.uni',
                    'arcc.uni',
                    'obqa.uni',
                ]
            )
        ):
            from .metric import ComputeChoiceUniAccuracy

            metric_module["compute_metrics"] = ComputeChoiceUniAccuracy(tokenizer=tokenizer, dataset='uni')

        elif data_args.eval_dataset and data_args.eval_dataset[0].startswith('hellaswag'):
            from .metric import ComputeChoiceAccuracy

            metric_module["compute_metrics"] = ComputeChoiceAccuracy(tokenizer=tokenizer, dataset='hellaswag')
        elif data_args.eval_dataset and data_args.eval_dataset[0].startswith('arcc'):
            from .metric import ComputeChoiceAccuracy

            metric_module["compute_metrics"] = ComputeChoiceAccuracy(tokenizer=tokenizer, dataset='ARC-Challenge')
        elif data_args.eval_dataset and data_args.eval_dataset[0].startswith('arce'):
            from .metric import ComputeChoiceAccuracy

            metric_module["compute_metrics"] = ComputeChoiceAccuracy(tokenizer=tokenizer, dataset='ARC-Easy')
        elif data_args.eval_dataset and data_args.eval_dataset[0].startswith('boolq'):
            from .metric import ComputeChoiceAccuracy

            metric_module["compute_metrics"] = ComputeChoiceAccuracy(tokenizer=tokenizer, dataset='boolq')
        elif data_args.eval_dataset and data_args.eval_dataset[0].startswith('obqa'):
            from .metric import ComputeChoiceAccuracy

            metric_module["compute_metrics"] = ComputeChoiceAccuracy(tokenizer=tokenizer, dataset='openbookqa')
        elif data_args.eval_dataset and data_args.eval_dataset[0].startswith('piqa'):
            from .metric import ComputeChoiceAccuracy

            metric_module["compute_metrics"] = ComputeChoiceAccuracy(tokenizer=tokenizer, dataset='piqa')
        elif data_args.eval_dataset and data_args.eval_dataset[0].startswith('siqa'):
            from .metric import ComputeChoiceAccuracy

            metric_module["compute_metrics"] = ComputeChoiceAccuracy(tokenizer=tokenizer, dataset='social_i_qa')
        elif data_args.eval_dataset and data_args.eval_dataset[0].startswith('winogrande'):
            from .metric import ComputeChoiceAccuracy

            metric_module["compute_metrics"] = ComputeChoiceAccuracy(tokenizer=tokenizer, dataset='winogrande')

        else:
            metric_module["compute_metrics"] = ComputeSimilarity(tokenizer=tokenizer)
    else:
        if finetuning_args.compute_accuracy:
            metric_module["compute_metrics"] = ComputeAccuracy()
            metric_module["preprocess_logits_for_metrics"] = eval_logit_processor

    # Initialize our Trainer
    trainer = CustomSeq2SeqTrainer(
        model=model,
        args=training_args,
        finetuning_args=finetuning_args,
        data_collator=data_collator,
        callbacks=callbacks,
        **dataset_module,
        **tokenizer_module,
        **metric_module,
    )

    # Keyword arguments for `model.generate`
    gen_kwargs = generating_args.to_dict()
    # Remove skip_special_tokens as it's not a valid generation parameter
    # It's used for tokenizer.decode(), not model.generate()
    gen_kwargs.pop("skip_special_tokens", None)
    gen_kwargs["eos_token_id"] = [tokenizer.eos_token_id] + tokenizer.additional_special_tokens_ids
    gen_kwargs["pad_token_id"] = tokenizer.pad_token_id
    gen_kwargs["logits_processor"] = get_logits_processor()

    # Training
    if training_args.do_train:
        train_result = trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)

        trainer.save_model()
        trainer.log_metrics("train", train_result.metrics)
        trainer.save_metrics("train", train_result.metrics)
        trainer.save_state()
        if trainer.is_world_process_zero() and finetuning_args.plot_loss:
            plot_loss(training_args.output_dir, keys=["loss", "eval_loss", "eval_accuracy"])

    if training_args.predict_with_generate:
        tokenizer.padding_side = "left"  # use left-padding in generation

    # Evaluation
    # if training_args.do_eval:
    #     eval_results = trainer.predict(dataset_module['eval_dataset'], metric_key_prefix="eval", **gen_kwargs)
    #     if training_args.predict_with_generate:  # predict_loss will be wrong if predict_with_generate is enabled
    #         eval_results.metrics.pop('eval_loss', None)
    #     trainer.log_metrics('eval', eval_results.metrics)
    #     trainer.save_metrics('eval', eval_results.metrics)
    #     trainer.save_predictions(dataset_module['eval_dataset'], eval_results)

    # Predict
    if training_args.do_predict:
        prepare_model_for_eval(model, model_args)
        predict_results = trainer.predict(dataset_module["eval_dataset"], metric_key_prefix="predict", **gen_kwargs)
        if training_args.predict_with_generate:  # predict_loss will be wrong if predict_with_generate is enabled
            predict_results.metrics.pop("predict_loss", None)
        trainer.log_metrics("predict", predict_results.metrics)
        trainer.save_metrics("predict", predict_results.metrics)
        trainer.save_predictions(
            dataset_module["eval_dataset"], predict_results, skip_special_tokens=generating_args.skip_special_tokens
        )

    # Create model card
    create_modelcard_and_push(trainer, model_args, data_args, training_args, finetuning_args)
