import torch

import transformers
import logging
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig, Qwen2Tokenizer
import warnings
import os
from aloha_scripts.utils import *


def find_all_linear_names(model, rank0_print):
    cls = torch.nn.Linear
    lora_module_names = set()

    multimodal_keywords = ['language_model', 'vision_model']

    rank0_print("##" * 20)

    for name, module in model.named_modules():
        # we only apply lora to the llm and vit
        if any(mm_keyword in name for mm_keyword in multimodal_keywords):
            if isinstance(module, cls):
                lora_module_names.add(name)

    if 'lm_head' in lora_module_names:  # needed for 16-bit
        lora_module_names.remove('lm_head')

    return list(lora_module_names)


def load_model(config=None, vla_config=None, rank0_print=print):
    model_args = config['model_args']
    training_args = config['training_args']
    data_args = config['data_args']
    action_args = config['action_head_args']

    kwargs = {"device_map": "cuda", "torch_dtype": torch.bfloat16}
    if config['model_args'].flash_attn:
        model = AutoModelForCausalLM.from_pretrained(
            config['model_args'].model_name_or_path,
            config=vla_config,
            cache_dir=config['training_args'].cache_dir,
            trust_remote_code=True,
            _fast_init=False,
            attn_implementation="flash_attention_2",
            **kwargs
        )
    else:
        model = AutoModelForCausalLM.from_pretrained(
            config['model_args'].model_name_or_path,
            config=vla_config,
            cache_dir=config['training_args'].cache_dir,
            trust_remote_code=True,
            _fast_init=False,
            **kwargs,  # specified device map and dtype may cause nan initialize
        )
    rank0_print(model)
    model.policy_head.initialize_weights()
    model.config.use_cache = False

    # >>>>>>>>>>>>>>>>>>>>>>>>>> setup for training configuration <<<<<<<<<<<<<<<<<<<<<<<<<<<<
    model_args.freeze_backbone = training_args.freeze_backbone
    if model_args.freeze_backbone:
        model.requires_grad_(False)
    else:
        model.requires_grad_(True)

    model.vision_model.requires_grad_(True)  # set to true first
    model.config.freeze_vision_tower = model_args.freeze_vision_tower = training_args.freeze_vision_tower
    if model_args.freeze_vision_tower:
        for n, p in model.vision_model.named_parameters():
            if not 'lora' in n.lower():
                p.requires_grad = False
    else:
        for p in model.vision_model.parameters():
            p.requires_grad = True

    if training_args.gradient_checkpointing:
        if hasattr(model, "enable_input_require_grads"):
            model.enable_input_require_grads()
        else:
            def make_inputs_require_grad(module, input, output):
                output.requires_grad_(True)

            model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)

    # >>>>>>>>>>>>>>>>>>>>>>>>>>>>> setup for lora <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
    if training_args.lora_enable:
        from peft import LoraConfig, get_peft_model
        lora_config = LoraConfig(
            r=training_args.lora_r,
            lora_alpha=training_args.lora_alpha,
            target_modules=find_all_linear_names(model, rank0_print, training_args.lora_module),
            lora_dropout=training_args.lora_dropout,
            bias=training_args.lora_bias,
            task_type=training_args.lora_task_type,
        )
        if training_args.bits == 16:
            if training_args.bf16:
                model.to(torch.bfloat16)
            if training_args.fp16:
                model.to(torch.float16)
        rank0_print("##" * 20)

        rank0_print("Adding LoRA adapters...")
        model = get_peft_model(model, lora_config)  # !!!only set lora weights to requires_grad True!!!
        rank0_print(model)
        model.print_trainable_parameters()

    # >>>>>>>>>>>>>>>>>>>>>>>>>> setup for projector, policy head <<<<<<<<<<<<<<<<<<<<<<<<<<<<
    for p in model.mlp1.parameters():
        p.requires_grad = True
    model.policy_head.requires_grad_(True)

    vision_tower = model.vision_model
    vision_tower.to(dtype=torch.bfloat16 if training_args.bf16 else torch.float16, device=training_args.device)
    model.to(dtype=torch.bfloat16 if training_args.bf16 else torch.float16, device=training_args.device)

    for k, v in model.named_parameters():
        if v.requires_grad:
            rank0_print(k, v.requires_grad, v.dtype)

    model.config.policy_head_lr = training_args.policy_head_lr

    rank0_print("!" * 100)
    lora_para = sum(p.numel() for n, p in model.named_parameters() if (p.requires_grad and 'lora' in n))
    all_para = sum(p.numel() for n, p in model.named_parameters())
    train_para = sum(p.numel() for n, p in model.named_parameters() if p.requires_grad)
    rank0_print(
        f"{RED}Lora parameters/trainalbe parameters/all parameters:{lora_para / 1e6}M/{train_para / 1e6}M/{(all_para - lora_para) / 1e6}M{RESET}")
    return model, data_args


def maybe_zero_3(param, ignore_status=False, name=None):
    from deepspeed import zero
    from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus
    if hasattr(param, "ds_id"):
        if param.ds_status == ZeroParamStatus.NOT_AVAILABLE:
            if not ignore_status:
                logging.warning(f"{name}: param.ds_status != ZeroParamStatus.NOT_AVAILABLE: {param.ds_status}")
        with zero.GatheredParameters([param]):
            param = param.data.detach().cpu().clone()
    else:
        param = param.detach().cpu().clone()
    return param


# Borrowed from peft.utils.get_peft_model_state_dict
def get_peft_state_maybe_zero_3(named_params, bias):
    if bias == "none":
        to_return = {k: t for k, t in named_params if "lora_" in k}
    elif bias == "all":
        to_return = {k: t for k, t in named_params if "lora_" in k or "bias" in k}
    elif bias == "lora_only":
        to_return = {}
        maybe_lora_bias = {}
        lora_bias_names = set()
        for k, t in named_params:
            if "lora_" in k:
                to_return[k] = t
                bias_name = k.split("lora_")[0] + "bias"
                lora_bias_names.add(bias_name)
            elif "bias" in k:
                maybe_lora_bias[k] = t
        for k, t in maybe_lora_bias:
            if bias_name in lora_bias_names:
                to_return[bias_name] = t
    else:
        raise NotImplementedError
    to_return = {k: maybe_zero_3(v, ignore_status=True) for k, v in to_return.items()}
    return to_return


def get_peft_state_non_lora_maybe_zero_3(named_params, require_grad_only=True):
    to_return = {k: t for k, t in named_params if "lora_" not in k}
    if require_grad_only:
        to_return = {k: t for k, t in to_return.items() if t.requires_grad}
    to_return = {k: maybe_zero_3(v, ignore_status=True).cpu() for k, v in to_return.items()}
    return to_return


def safe_save_model_for_hf_trainer(trainer: transformers.Trainer,
                                   output_dir: str):
    """Collects the state dict and dump to disk."""

    if trainer.deepspeed:
        torch.cuda.synchronize()
        trainer.save_model(output_dir)
        return

    state_dict = trainer.model.state_dict()
    if trainer.args.should_save:
        cpu_state_dict = {
            key: value.cpu()
            for key, value in state_dict.items()
        }
        del state_dict
        trainer._save(output_dir, state_dict=cpu_state_dict)  # noqa


def load_merge_lora_weights(model_path=None, model_base=None, kwargs=None):
    tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=True)
    model = AutoModelForCausalLM.from_pretrained(model_base,
                                                 low_cpu_mem_usage=True,
                                                 **kwargs)

    if os.path.exists(os.path.join(model_path, 'non_lora_trainables.bin')):
        non_lora_trainables = torch.load(os.path.join(model_path, 'non_lora_trainables.bin'), )
    else:
        raise FileNotFoundError
    if any(k.startswith('model.policy_head.') for k in non_lora_trainables):
        non_lora_trainables = {(k[6:] if k.startswith('model.') else k): v for k, v in
                               non_lora_trainables.items()}

    keys_to_del = []
    for k, v in non_lora_trainables.items():
        if 'lora' in k:
            keys_to_del.append(k)
    for key in keys_to_del:
        del non_lora_trainables[key]
    model.load_state_dict(non_lora_trainables, strict=False)

    from peft import PeftModel
    assert os.path.exists(os.path.join(model_path, "adapter_model.safetensors"))
    print('Loading LoRA weights...')
    model = PeftModel.from_pretrained(model, model_path)
    print('Merging LoRA weights...')
    model = model.merge_and_unload()
    print('Model is loaded...')
    return model, tokenizer

def load_model_for_eval(model_path, model_base, device_map="cuda:0", policy_config=None):
    kwargs = {"device_map": device_map, 'torch_dtype': torch.bfloat16}

    if 'lora' in model_path.lower() and model_base is None:
        warnings.warn(
            'There is `lora` in model name but no `model_base` is provided. If you are loading a LoRA model, '
            'please provide the `model_base` argument.')
    if 'lora' in model_path.lower() and model_base is not None:
        model, tokenizer = load_merge_lora_weights(model_path=model_path,
                                                   model_base=model_base,
                                                   kwargs=kwargs)

        if policy_config['save_model']:
            print(f"#####################Saving merged weights of model in {kwargs['torch_dtype']}.#####################")
            os.makedirs(os.path.join(model_path, 'merge_weights'), exist_ok=True)
            model.save_pretrained(
                os.path.join(model_path, 'merge_weights'))
            tokenizer.save_pretrained(os.path.join(model_path, 'merge_weights'))
            skip_params = [
                "input_action_proj",
                "policy_head",
                "reasoning_action_proj",
                "reasoning_film",
            ]
            head_param = {}
            for k, v in model.named_parameters():
                if any(skip_param in k.lower() for skip_param in skip_params):
                    head_param[k] = v
            torch.save(head_param, os.path.join(model_path, 'merge_weights/head_params.bin'))

    else:
        print(f"load {model_path}!!!")
        config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
        tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True, use_fast=False)
        model = AutoModelForCausalLM.from_pretrained(
            model_path,
            config=config,
            use_safetensors=True,
            **kwargs)

    model.to(device="cuda")
    return tokenizer, model
