# 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.


# MeM: This script is written based on `llamafactory.train.sft.workflow` for MeM training.
import json
import os
import sys
import random
import time
from typing import TYPE_CHECKING, List, Optional

import wandb
import torch
from safetensors.torch import save_file, load_file
from transformers import set_seed

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
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
from .trainer import CustomSeq2SeqTrainer

set_seed(0)

if TYPE_CHECKING:
    from transformers import Seq2SeqTrainingArguments, TrainerCallback

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

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../../..')))
from src.mem import MemLlamaConfig, MemLlamaForCausalLM, update_memory_fifo, MemQwen3Config, MemQwen3ForCausalLM

logger = logging.get_logger(__name__)


def init_mem_config(model_args: ModelArguments):
    # Overwrite default settings in MemLlamaConfig
    if 'llama' in model_args.model_name_or_path.lower():
        config = MemLlamaConfig.from_pretrained(model_args.model_name_or_path)
    elif 'qwen' in model_args.model_name_or_path.lower():
        config = MemQwen3Config.from_pretrained(model_args.model_name_or_path)
    else:
        raise ValueError(f"Model type {model_args.model_name_or_path} not supported")
    try:
        if model_args.memory_insert_layers == 'all':
            config.memory_insert_layers = [idx for idx in range(config.num_hidden_layers)]
        elif model_args.memory_insert_layers == 'alternate':
            config.memory_insert_layers = [idx for idx in range(1, config.num_hidden_layers, 2)]
        elif model_args.memory_insert_layers == 'first_half':
            config.memory_insert_layers = [idx for idx in range(int(config.num_hidden_layers * 0.5))]
        elif model_args.memory_insert_layers == 'last_half':
            config.memory_insert_layers = [idx for idx in
                                           range(int(config.num_hidden_layers * 0.5), config.num_hidden_layers)]
        elif model_args.memory_insert_layers == 'first_quarter':
            config.memory_insert_layers = [idx for idx in range(int(config.num_hidden_layers * 0.25))]
        elif model_args.memory_insert_layers == 'last_quarter':
            config.memory_insert_layers = [idx for idx in
                                           range(int(config.num_hidden_layers * 0.75), config.num_hidden_layers)]
        else:
            config.memory_insert_layers = [int(idx.strip()) for idx in model_args.memory_insert_layers.split(',')]
    except:
        raise ValueError(f'memory_insert_layers must be either "all" or a comma-separated string of integers. '
                         f'Got {model_args.memory_insert_layers} instead.')
    config.update_memory = model_args.update_memory
    config.memory_size = model_args.memory_size
    config.use_gpu_to_search = model_args.use_gpu_to_search
    config.mem_hidden_path = model_args.mem_hidden_path
    config.fix_memory = model_args.fix_memory
    config.memory_update_steps = model_args.memory_update_steps
    config.fusion_func = model_args.fusion_func

    config.use_last_prompt_token_as_key = model_args.use_last_prompt_token_as_key

    config.update_while_predicting = model_args.update_while_predicting

    config.update_strategy = model_args.update_strategy
    config.memory_size_l2 = model_args.memory_size_l2

    config.print_configs()  # Print MeM Configurations

    return config


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_memory_from_disk(model: MemLlamaForCausalLM | MemQwen3ForCausalLM, model_args: ModelArguments,
                          finetuning_args: "FinetuningArguments"):
    print(f'Initializing memory from {finetuning_args.memory_path}')
    mem_insert_layers = map(int, model_args.memory_insert_layers.split(','))

    all_input_ids = load_file(os.path.join(finetuning_args.memory_path, 'input_ids.safetensors'))
    all_labels = load_file(os.path.join(finetuning_args.memory_path, 'labels.safetensors'))

    candidates = list(all_input_ids.keys())  # Example ID
    candidates = random.sample(candidates, model_args.memory_size)

    for layer_idx in mem_insert_layers:
        layer_hs_ln = load_file(
            os.path.join(finetuning_args.memory_path, f"layer_{layer_idx}.ln_states.safetensors"))

        for candidate in candidates:

            input_ids = all_input_ids[candidate].unsqueeze(0)
            labels = all_labels[candidate].unsqueeze(0)
            hidden_states_ln = layer_hs_ln[candidate].unsqueeze(0)

            memory_hidden_states_v, memory_hidden_states_k = model.model.process_hidden_states(hidden_states_ln, input_ids,
                                                                                         labels)
            model.model.external_memory_v.update_to_memory_modules()
            update_memory_fifo(layer_idx, model.model.external_memory_v, memory_hidden_states_v.detach())
            if model.model.use_last_prompt_token_as_key:
                model.model.external_memory_k.update_to_memory_modules()
                update_memory_fifo(layer_idx, model.model.external_memory_k, memory_hidden_states_k.detach())

        del layer_hs_ln

    return model


def init_mem_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 = init_mem_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

    if 'llama' in model_args.model_name_or_path.lower():
        model = MemLlamaForCausalLM.from_pretrained(**init_kwargs)
    elif 'qwen' in model_args.model_name_or_path.lower():
        model = MemQwen3ForCausalLM.from_pretrained(**init_kwargs)
    else:
        raise ValueError(f"Model type {model_args.model_name_or_path} not supported")

    if finetuning_args.init_cross_attn_from_self:
        model.init_cross_attention_from_self()

    for name, param in model.named_parameters():
        if 'memory' in name:
            param.data.fill_(0.0)

    if finetuning_args.memory_path is not None and not finetuning_args.memory_path == "none":
        if 'wandb' in training_args.report_to:
            wandb.config['init_memory_from_disk'] = True
        model = init_memory_from_disk(model, model_args, finetuning_args)

    model = init_adapter(config, model, model_args, finetuning_args, training_args.do_train)
    # MeM: Load external_memory module
    if model_args.adapter_name_or_path is not None:
        logger.info('Loading external_memory from external_memory.safetensors')
        assert len(model_args.adapter_name_or_path) == 1
        external_memory_state_dict = load_file(os.path.join(model_args.adapter_name_or_path[0],
                                                            'external_memory.safetensors'))
        # MeM: Modify key mapping
        external_memory_state_dict = {k.replace("base_model.model.model.", "model."): v for k, v in
                                      external_memory_state_dict.items()}

        model.load_state_dict(external_memory_state_dict, strict=False)

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

    trainable_params, all_param = count_parameters(model)
    if training_args.do_train:
        param_stats = "trainable params: {:,} || all params: {:,} || trainable%: {:.4f}".format(
            trainable_params, all_param, 100 * trainable_params / all_param
        )
    else:
        param_stats = "all params: {:,}".format(all_param)

    logger.info(param_stats)

    return config, model


def run_mem_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_mem_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"])

    if 'wandb' in training_args.report_to:
        if data_args.dataset:
            dataset_name = data_args.dataset[0].split('_')[0]
        else:
            dataset_name = data_args.eval_dataset[0].split('_')[0]
        wandb.config['dataset_name'] = dataset_name

    # 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)

        # MeM: Save trainable modules & external_memory modules
        trainer.save_model()
        external_memory_state_dict = {k: v for k, v in trainer.model.state_dict().items() if 'external_memory' in k}
        save_file(external_memory_state_dict, os.path.join(training_args.output_dir, 'external_memory.safetensors'))

        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)
    #     if training_args.predict_with_generate:
    #         trainer.save_predictions(dataset_module['eval_dataset'], eval_results)

    # Predict
    if training_args.do_predict:
        prepare_model_for_eval(model, model_args)
        # orig_state_dict = {k: v for k, v in trainer.model.state_dict().items() if 'external_memory' in k}.copy()
        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)

        if 'wandb' in training_args.report_to:
            pred_acc = predict_results.metrics['predict_accuracy']
            wandb.log({'accuracy': round(pred_acc, 4)})

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