import os
import torch
import torch as _torch; _orig_load = _torch.load
import functools
def _patched_load(*a, **kw):
    kw["weights_only"] = False
    return _orig_load(*a, **kw)
_torch.load = _patched_load
from peft import LoraConfig, get_peft_model
import ast
from transformers import AutoProcessor, BitsAndBytesConfig, HfArgumentParser
from training.trainer import QwenTrainer, UnfreezeLoRACallback, ResumeDatasetCallback
from training.data import make_supervised_data_module
from training.params import DataArguments, ModelArguments, TrainingArguments
from training.train_utils import get_peft_state_maybe_zero_3, get_peft_state_non_lora_maybe_zero_3, safe_save_model_for_hf_trainer
import pathlib
from liger_kernel.transformers import apply_liger_kernel_to_qwen2_vl
from monkey_patch_forward import replace_qwen2_5_with_mixed_modality_forward, replace_qwen_2_with_mixed_modality_forward

from training.covt_qwen2_5_vl import CoVTForConditionalGeneration
from training.constants import *
from deepspeed import zero

local_rank = None

# set seed 42
torch.manual_seed(42)

def rank0_print(*args):
    if local_rank == 0 or local_rank == '0' or local_rank is None:
        print(*args)

def find_target_linear_names(model, num_lora_modules=-1, lora_namespan_exclude=[], verbose=True):
    linear_cls = torch.nn.modules.Linear
    embedding_cls = torch.nn.modules.Embedding
    lora_module_names = []

    for name, module in model.named_modules():
        if any(ex_keyword in name for ex_keyword in lora_namespan_exclude):
            continue
        if isinstance(module, (linear_cls, embedding_cls)):
            lora_module_names.append(name)
    
    if num_lora_modules > 0:
        lora_module_names = lora_module_names[-num_lora_modules:]
    if verbose:
        rank0_print(f"Found {len(lora_module_names)} lora modules: {lora_module_names}")
    
    return lora_module_names

def set_requires_grad(parameters, requires_grad):
    for p in parameters:
        p.requires_grad = requires_grad
        
def set_anchor_requires_grad(model, anchor_model_id):
    # set_requires_grad(model.sam_projection.parameters(), True)
    # set_requires_grad(model.dino_projection.parameters(), True)
    # set_requires_grad(model.depth_projection.parameters(), True)
    # set_requires_grad(model.SD_projection.parameters(), True)
    # set_requires_grad(model.internvit_projection.parameters(), True)
    # set_requires_grad(model.pidinet_projection.parameters(), True)
    # set_requires_grad(model.siglip_projection.parameters(), True)
    # set_requires_grad(model.metaclip_projection.parameters(), True)
    
    # set_requires_grad(model.sam_cross_attention.parameters(), True)
    # set_requires_grad(model.dino_cross_attention.parameters(), True)
    # set_requires_grad(model.depth_cross_attention.parameters(), True)
    # set_requires_grad(model.SD_cross_attention.parameters(), True)
    # set_requires_grad(model.internvit_cross_attention.parameters(), True)
    # set_requires_grad(model.pidinet_cross_attention.parameters(), True)
    # set_requires_grad(model.siglip_cross_attention.parameters(), True)
    # set_requires_grad(model.metaclip_cross_attention.parameters(), True)

    # model.dino_query_vectors.requires_grad = True
    # model.sam_query_vectors.requires_grad = True
    # model.depth_query_vectors.requires_grad = True
    # model.SD_query_vectors.requires_grad = True
    # model.internvit_query_vectors.requires_grad = True
    # model.pidinet_query_vectors.requires_grad = True
    # model.siglip_query_vectors.requires_grad = True
    # model.metaclip_query_vectors.requires_grad = True
    
    if "sam" in anchor_model_id:
        set_requires_grad(model.sam_projection.parameters(), True)
        set_requires_grad(model.sam_cross_attention.parameters(), True)
        model.sam_query_vectors.requires_grad = True
    if "dino" in anchor_model_id:
        set_requires_grad(model.dino_projection.parameters(), True)
        set_requires_grad(model.dino_cross_attention.parameters(), True)
        model.dino_query_vectors.requires_grad = True
    if "depth" in anchor_model_id:
        set_requires_grad(model.depth_projection.parameters(), True)
        set_requires_grad(model.depth_cross_attention.parameters(), True)
        model.depth_query_vectors.requires_grad = True
        set_requires_grad(model.depth_token_generator.parameters(), True)
    if "sd" in anchor_model_id:
        set_requires_grad(model.SD_projection.parameters(), True)
        set_requires_grad(model.SD_cross_attention.parameters(), True)
        model.SD_query_vectors.requires_grad = True
    if "internvit" in anchor_model_id:
        set_requires_grad(model.internvit_projection.parameters(), True)
        set_requires_grad(model.internvit_cross_attention.parameters(), True)
        model.internvit_query_vectors.requires_grad = True
    if "pidinet" in anchor_model_id:
        set_requires_grad(model.pidinet_projection.parameters(), True)
        set_requires_grad(model.pidinet_cross_attention.parameters(), True)
        model.pidinet_query_vectors.requires_grad = True
    if "siglip" in anchor_model_id:
        set_requires_grad(model.siglip_projection.parameters(), True)
        set_requires_grad(model.siglip_cross_attention.parameters(), True)
        model.siglip_query_vectors.requires_grad = True
    if "metaclip" in anchor_model_id:
        set_requires_grad(model.metaclip_projection.parameters(), True)
        set_requires_grad(model.metaclip_cross_attention.parameters(), True)
        model.metaclip_query_vectors.requires_grad = True

def configure_vision_tower(model, training_args, compute_dtype, device):
    vision_tower = model.visual
    vision_tower.to(dtype=compute_dtype, device=device)

    vision_model_params = model.visual.parameters()
    set_requires_grad(vision_model_params, not training_args.freeze_vision_tower)
    
    # Handle merger specifically
    merger_params = model.visual.merger.parameters()
    set_requires_grad(merger_params, training_args.tune_merger)
    
def configure_llava_vision_tower(model, model_args, training_args, compute_dtype, processor):
    model.get_model().initialize_vision_modules(
            model_args=model_args,
            fsdp=training_args.fsdp
    )
    
    vision_tower = model.get_vision_tower()
    vision_tower.to(dtype=torch.bfloat16 if training_args.bf16 else torch.float16, device=training_args.device)

    model.config.image_aspect_ratio = "pad"
    model.config.tokenizer_padding_side = processor.tokenizer.padding_side
    model.config.tokenizer_model_max_length = processor.tokenizer.model_max_length

    if training_args.bits in [4, 8]:
        model.get_model().mm_projector.to(dtype=compute_dtype, device=training_args.device)

    model.config.mm_use_im_start_end = False
    model.config.mm_projector_lr = 2e-5
    training_args.use_im_start_end = False
    model.config.mm_use_im_patch_token = False
    model.initialize_vision_tokenizer(model_args, tokenizer=processor.tokenizer)

def configure_llm(model, training_args):
    lm_head = model.lm_head.parameters()
    set_requires_grad(lm_head, not training_args.freeze_llm)

    llm_params = model.model.parameters()
    set_requires_grad(llm_params, not training_args.freeze_llm)


def train():
    global local_rank

    parser = HfArgumentParser(
        (ModelArguments, DataArguments, TrainingArguments))
    
    model_args, data_args, training_args = parser.parse_args_into_dataclasses()
    
    anchor_model_id = ast.literal_eval(model_args.anchor_model_id)
    
    if model_args.model_path is None:
        print("\033[91mWARNING: model_path is not provided, using model_id instead\033[0m")
        model_args.model_path = model_args.model_id
    
    # Liger-kernel for Qwen2.5 is not supported yet.
    replace_qwen2_5_with_mixed_modality_forward(use_liger=training_args.use_liger)\
    

    if training_args.lora_enable and not training_args.freeze_llm:
        raise ValueError("If `lora_enable` is True, `freeze_llm` must also be True.")

    if not training_args.lora_enable:
        assert not training_args.vision_lora, \
            "Error: training_args.lora_enable is not enabled, but training_args.vision_lora is enabled."
        
    if training_args.vision_lora and not training_args.freeze_vision_tower:
        raise ValueError("If `vision_lora` is True, `freeze_vision_tower` must also be True.")

    else:
        if training_args.lora_namespan_exclude is not None:
            training_args.lora_namespan_exclude = ast.literal_eval(training_args.lora_namespan_exclude)
        else:
            training_args.lora_namespan_exclude = []

        if not training_args.vision_lora:
            training_args.lora_namespan_exclude += ["visual"]

    local_rank = training_args.local_rank
    compute_dtype = (torch.float16 if training_args.fp16 else (torch.bfloat16 if training_args.bf16 else torch.float32))

    bnb_model_from_pretrained_args = {}
    if training_args.bits in [4,8]:
        bnb_model_from_pretrained_args.update(dict(
            device_map={"":training_args.device},
            quantization_config = BitsAndBytesConfig(
                load_in_4bit=training_args.bits==4,
                load_in_8bit=training_args.bits==8,
                llm_int8_skip_modules=["visual"],
                llm_int8_threshold=6.0,
                llm_int8_has_fp16_weight=False,
                bnb_4bit_compute_dtype=compute_dtype,
                bnb_4bit_use_double_quant=training_args.double_quant,
                bnb_4bit_quant_type=training_args.quant_type,
            )
        ))

    model = CoVTForConditionalGeneration.from_pretrained(
        model_args.model_path,
        torch_dtype=compute_dtype,
        attn_implementation="flash_attention_2" if not training_args.disable_flash_attn2 else "sdpa", 
        **bnb_model_from_pretrained_args
    )
    
    model.get_anchor_model_ids(anchor_model_id)
    model.align_vqa_only_stage = training_args.vqa_only_stage

    model.config.use_cache = False
    model_to_configure = model
    configure_llm(model_to_configure, training_args)
    if "Qwen" in model_args.model_id:
        configure_vision_tower(model_to_configure, training_args, compute_dtype, training_args.device)
    
    # Set requires_grad for the Anchor projection layers
    set_anchor_requires_grad(model, anchor_model_id)
    
    if training_args.bits in [4,8]:
        model.config.torch_dtype = (torch.float32 if training_args.fp16 else (torch.bfloat16 if training_args.bf16 else torch.float32))
        from peft import prepare_model_for_kbit_training
        model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=training_args.gradient_checkpointing, gradient_checkpointing_kwargs={"use_reentrant": True})
    
    if training_args.gradient_checkpointing:
        model.enable_input_require_grads()
        training_args.gradient_checkpointing_kwargs = {"use_reentrant": True}

    if training_args.lora_enable:
        lora_namespan_exclude = training_args.lora_namespan_exclude
        if "llava" in model_args.model_id:
            lora_namespan_exclude += ['mm_projector', 'vision_tower', 'vision_resampler']
        peft_config = LoraConfig(
            r=training_args.lora_rank,
            lora_alpha=training_args.lora_alpha,
            target_modules=find_target_linear_names(model, lora_namespan_exclude=lora_namespan_exclude, num_lora_modules=training_args.num_lora_modules),
            lora_dropout=training_args.lora_dropout,
            bias=training_args.lora_bias
        )
        if training_args.bits == 16:
            if training_args.bf16:
                model.to(torch.bfloat16)
            if training_args.fp16:
                model.to(torch.float16)
        rank0_print("Adding LoRA to the model...")
        model = get_peft_model(model, peft_config)
        
        for name, param in model.named_parameters():
            if '_projection' in name:
                param.requires_grad = True
            if 'cross_attention' in name:
                param.requires_grad = True
            if '_query_vectors' in name:
                param.requires_grad = True
        # model.print_trainable_parameters()

    processor = AutoProcessor.from_pretrained(model_args.model_id,
                                            # The default setting is padding_side="left"
                                            # When training using the right-side padding is more efficient.
                                              padding_side="right")

    # model.config.tokenizer_model_max_length = processor.tokenizer.model_max_length
    model.config.tokenizer_padding_side = processor.tokenizer.padding_side
    model.config.vision_lr = training_args.vision_lr
    
    old_processor_len = len(processor.tokenizer)
        
    # add special tokens
    add_tokens = [SAM_PAD_TOKEN, DINO_PAD_TOKEN, DEPTH_PAD_TOKEN, SD_PAD_TOKEN, INTERN_PAD_TOKEN, PIDINET_PAD_TOKEN, SIGLIP_PAD_TOKEN, METACLIP_PAD_TOKEN, "<think>", "</think>", "<answer>", "</answer>"]
    processor.tokenizer.add_special_tokens({"additional_special_tokens": [ANCHOR_START_TOKEN, ANCHOR_END_TOKEN]})
    processor.tokenizer.add_tokens(add_tokens)
    sam_token_idx = processor.tokenizer(SAM_PAD_TOKEN, add_special_tokens=False).input_ids[0]
    dino_token_idx = processor.tokenizer(DINO_PAD_TOKEN, add_special_tokens=False).input_ids[0]
    depth_token_idx = processor.tokenizer(DEPTH_PAD_TOKEN, add_special_tokens=False).input_ids[0]
    sd_token_idx = processor.tokenizer(SD_PAD_TOKEN, add_special_tokens=False).input_ids[0]
    intern_token_idx = processor.tokenizer(INTERN_PAD_TOKEN, add_special_tokens=False).input_ids[0]
    pidinet_token_idx = processor.tokenizer(PIDINET_PAD_TOKEN, add_special_tokens=False).input_ids[0]
    siglip_token_idx = processor.tokenizer(SIGLIP_PAD_TOKEN, add_special_tokens=False).input_ids[0]
    metaclip_token_idx = processor.tokenizer(METACLIP_PAD_TOKEN, add_special_tokens=False).input_ids[0]
    
    think_idx = processor.tokenizer("<think>", add_special_tokens=False).input_ids[0]
    splash_think_idx = processor.tokenizer("</think>", add_special_tokens=False).input_ids[0]
    answer_idx = processor.tokenizer("<answer>", add_special_tokens=False).input_ids[0]
    splash_answer_idx = processor.tokenizer("</answer>", add_special_tokens=False).input_ids[0]
    
    qwen_embed = model.get_input_embeddings()
    lm_head = model.get_output_embeddings()
    p = qwen_embed.weight
    if hasattr(p, 'ds_id'):
        with zero.GatheredParameters([p]):
            old_len = p.data.shape[0]
    else:
        old_len = p.data.shape[0]
    new_len = len(processor.tokenizer)
    
    model.get_anchor_token_idx(sam_token_idx, dino_token_idx, depth_token_idx, sd_token_idx, intern_token_idx, pidinet_token_idx, siglip_token_idx, metaclip_token_idx)
    
    if "llava" in model_args.model_id:
        configure_llava_vision_tower(model_to_configure, model_args, training_args, compute_dtype, processor)
        
    for n, p in model.named_parameters():
        if any(
            [
                "embed_tokens" in n
            ]
        ):
            p.requires_grad = True
            
        if any(
            [
                "lm_head" in n
            ]
        ):
            p.requires_grad = True
    
    _mask = torch.ones(old_len, device=model.device, dtype=torch.bool)
    _mask[:] = False
    _mask[old_processor_len:new_len] = True

    def row_mask_hook(grad):
        if grad is None:
            return grad
        return grad * _mask.to(grad.device).view(-1, 1)
    
    model.get_input_embeddings().weight.register_hook(row_mask_hook)
    model.get_output_embeddings().weight.register_hook(row_mask_hook)

    if training_args.bits in [4, 8]:
        from peft.tuners.lora import LoraLayer
        for name, module in model.named_modules():
            if isinstance(module, LoraLayer):
                if training_args.bf16:
                    module = module.to(torch.bfloat16)
            if 'norm' in name:
                module = module.to(torch.float32)
            
            if 'lm_head' in name or 'embed_token' in name:
                if hasattr(module, 'weight'):
                    if training_args.bf16 and module.weight.dtype == torch.float32:
                        module = module.to(torch.bfloat16)

    data_module = make_supervised_data_module(model_id=model_args.model_id,
                                              processor=processor,
                                              data_args=data_args,
                                              anchor_model_id=anchor_model_id)
    
    resume_callback = ResumeDatasetCallback(train_dataset=data_module['train_dataset'])
    
    trainer = QwenTrainer(
        model=model,
        processor=processor,
        args=training_args,
        callbacks=[resume_callback],
        **data_module
    )
    # model.print_trainable_parameters()

    if list(pathlib.Path(training_args.output_dir).glob("checkpoint-*")):
        trainer.train(resume_from_checkpoint=True)
    else:
        trainer.train()
            
    trainer.save_state()

    model.config.use_cache = True
    
    if training_args.lora_enable:
        state_dict = get_peft_state_maybe_zero_3(
            model.named_parameters(), training_args.lora_bias
        )

        non_lora_state_dict = get_peft_state_non_lora_maybe_zero_3(
            model.named_parameters(), require_grad_only=False
        )

        if local_rank == 0 or local_rank == -1:
            model.config.save_pretrained(training_args.output_dir)
            model.save_pretrained(training_args.output_dir, state_dict=state_dict)
            torch.save(non_lora_state_dict, os.path.join(training_args.output_dir, "non_lora_state_dict.bin"))
    else:
        safe_save_model_for_hf_trainer(trainer, output_dir=training_args.output_dir)



if __name__ == "__main__":
    train()