from typing import List

from mflux.models.common.weights.mapping.weight_mapping import WeightMapping, WeightTarget
from mflux.models.common.weights.mapping.weight_transforms import WeightTransforms


class FIBOVLMWeightMapping(WeightMapping):
    @staticmethod
    def get_vlm_decoder_mapping(num_layers: int = 36) -> List[WeightTarget]:
        return [
            WeightTarget(
                to_pattern="embed_tokens.weight",
                from_pattern=["model.language_model.embed_tokens.weight"],
            ),
            WeightTarget(
                to_pattern="layers.{block}.self_attn.q_proj.weight",
                from_pattern=["model.language_model.layers.{block}.self_attn.q_proj.weight"],
            ),
            WeightTarget(
                to_pattern="layers.{block}.self_attn.q_proj.bias",
                from_pattern=["model.language_model.layers.{block}.self_attn.q_proj.bias"],
            ),
            WeightTarget(
                to_pattern="layers.{block}.self_attn.k_proj.weight",
                from_pattern=["model.language_model.layers.{block}.self_attn.k_proj.weight"],
            ),
            WeightTarget(
                to_pattern="layers.{block}.self_attn.k_proj.bias",
                from_pattern=["model.language_model.layers.{block}.self_attn.k_proj.bias"],
            ),
            WeightTarget(
                to_pattern="layers.{block}.self_attn.v_proj.weight",
                from_pattern=["model.language_model.layers.{block}.self_attn.v_proj.weight"],
            ),
            WeightTarget(
                to_pattern="layers.{block}.self_attn.v_proj.bias",
                from_pattern=["model.language_model.layers.{block}.self_attn.v_proj.bias"],
            ),
            WeightTarget(
                to_pattern="layers.{block}.self_attn.o_proj.weight",
                from_pattern=["model.language_model.layers.{block}.self_attn.o_proj.weight"],
            ),
            WeightTarget(
                to_pattern="layers.{block}.self_attn.o_proj.bias",
                from_pattern=["model.language_model.layers.{block}.self_attn.o_proj.bias"],
            ),
            WeightTarget(
                to_pattern="layers.{block}.self_attn.q_norm.weight",
                from_pattern=["model.language_model.layers.{block}.self_attn.q_norm.weight"],
            ),
            WeightTarget(
                to_pattern="layers.{block}.self_attn.k_norm.weight",
                from_pattern=["model.language_model.layers.{block}.self_attn.k_norm.weight"],
            ),
            WeightTarget(
                to_pattern="layers.{block}.mlp.gate_proj.weight",
                from_pattern=["model.language_model.layers.{block}.mlp.gate_proj.weight"],
            ),
            WeightTarget(
                to_pattern="layers.{block}.mlp.gate_proj.bias",
                from_pattern=["model.language_model.layers.{block}.mlp.gate_proj.bias"],
            ),
            WeightTarget(
                to_pattern="layers.{block}.mlp.up_proj.weight",
                from_pattern=["model.language_model.layers.{block}.mlp.up_proj.weight"],
            ),
            WeightTarget(
                to_pattern="layers.{block}.mlp.up_proj.bias",
                from_pattern=["model.language_model.layers.{block}.mlp.up_proj.bias"],
            ),
            WeightTarget(
                to_pattern="layers.{block}.mlp.down_proj.weight",
                from_pattern=["model.language_model.layers.{block}.mlp.down_proj.weight"],
            ),
            WeightTarget(
                to_pattern="layers.{block}.mlp.down_proj.bias",
                from_pattern=["model.language_model.layers.{block}.mlp.down_proj.bias"],
            ),
            WeightTarget(
                to_pattern="layers.{block}.input_layernorm.weight",
                from_pattern=["model.language_model.layers.{block}.input_layernorm.weight"],
            ),
            WeightTarget(
                to_pattern="layers.{block}.post_attention_layernorm.weight",
                from_pattern=["model.language_model.layers.{block}.post_attention_layernorm.weight"],
            ),
            WeightTarget(
                to_pattern="norm.weight",
                from_pattern=["model.language_model.norm.weight"],
            ),
            WeightTarget(
                to_pattern="lm_head.weight",
                from_pattern=["model.language_model.embed_tokens.weight"],
            ),
        ]

    @staticmethod
    def get_vlm_visual_mapping(depth: int = 24) -> List[WeightTarget]:
        return [
            WeightTarget(
                to_pattern="patch_embed.proj.weight",
                from_pattern=["model.visual.patch_embed.proj.weight"],
                transform=WeightTransforms.transpose_conv3d_weight,
            ),
            WeightTarget(
                to_pattern="patch_embed.proj.bias",
                from_pattern=["model.visual.patch_embed.proj.bias"],
            ),
            WeightTarget(
                to_pattern="pos_embed.weight",
                from_pattern=["model.visual.pos_embed.weight"],
            ),
            WeightTarget(
                to_pattern="blocks.{block}.norm1.weight",
                from_pattern=["model.visual.blocks.{block}.norm1.weight"],
            ),
            WeightTarget(
                to_pattern="blocks.{block}.norm1.bias",
                from_pattern=["model.visual.blocks.{block}.norm1.bias"],
            ),
            WeightTarget(
                to_pattern="blocks.{block}.norm2.weight",
                from_pattern=["model.visual.blocks.{block}.norm2.weight"],
            ),
            WeightTarget(
                to_pattern="blocks.{block}.norm2.bias",
                from_pattern=["model.visual.blocks.{block}.norm2.bias"],
            ),
            # Attention
            WeightTarget(
                to_pattern="blocks.{block}.attn.qkv.weight",
                from_pattern=["model.visual.blocks.{block}.attn.qkv.weight"],
            ),
            WeightTarget(
                to_pattern="blocks.{block}.attn.qkv.bias",
                from_pattern=["model.visual.blocks.{block}.attn.qkv.bias"],
            ),
            WeightTarget(
                to_pattern="blocks.{block}.attn.proj.weight",
                from_pattern=["model.visual.blocks.{block}.attn.proj.weight"],
            ),
            WeightTarget(
                to_pattern="blocks.{block}.attn.proj.bias",
                from_pattern=["model.visual.blocks.{block}.attn.proj.bias"],
            ),
            # MLP
            WeightTarget(
                to_pattern="blocks.{block}.mlp.linear_fc1.weight",
                from_pattern=["model.visual.blocks.{block}.mlp.linear_fc1.weight"],
            ),
            WeightTarget(
                to_pattern="blocks.{block}.mlp.linear_fc1.bias",
                from_pattern=["model.visual.blocks.{block}.mlp.linear_fc1.bias"],
            ),
            WeightTarget(
                to_pattern="blocks.{block}.mlp.linear_fc2.weight",
                from_pattern=["model.visual.blocks.{block}.mlp.linear_fc2.weight"],
            ),
            WeightTarget(
                to_pattern="blocks.{block}.mlp.linear_fc2.bias",
                from_pattern=["model.visual.blocks.{block}.mlp.linear_fc2.bias"],
            ),
            # Final merger
            WeightTarget(
                to_pattern="merger.norm.weight",
                from_pattern=["model.visual.merger.norm.weight"],
            ),
            WeightTarget(
                to_pattern="merger.norm.bias",
                from_pattern=["model.visual.merger.norm.bias"],
            ),
            WeightTarget(
                to_pattern="merger.linear_fc1.weight",
                from_pattern=["model.visual.merger.linear_fc1.weight"],
            ),
            WeightTarget(
                to_pattern="merger.linear_fc1.bias",
                from_pattern=["model.visual.merger.linear_fc1.bias"],
            ),
            WeightTarget(
                to_pattern="merger.linear_fc2.weight",
                from_pattern=["model.visual.merger.linear_fc2.weight"],
            ),
            WeightTarget(
                to_pattern="merger.linear_fc2.bias",
                from_pattern=["model.visual.merger.linear_fc2.bias"],
            ),
            WeightTarget(
                to_pattern="deepstack_merger_list.{block}.norm.weight",
                from_pattern=["model.visual.deepstack_merger_list.{block}.norm.weight"],
                max_blocks=3,
            ),
            WeightTarget(
                to_pattern="deepstack_merger_list.{block}.norm.bias",
                from_pattern=["model.visual.deepstack_merger_list.{block}.norm.bias"],
                max_blocks=3,
            ),
            WeightTarget(
                to_pattern="deepstack_merger_list.{block}.linear_fc1.weight",
                from_pattern=["model.visual.deepstack_merger_list.{block}.linear_fc1.weight"],
                max_blocks=3,
            ),
            WeightTarget(
                to_pattern="deepstack_merger_list.{block}.linear_fc1.bias",
                from_pattern=["model.visual.deepstack_merger_list.{block}.linear_fc1.bias"],
                max_blocks=3,
            ),
            WeightTarget(
                to_pattern="deepstack_merger_list.{block}.linear_fc2.weight",
                from_pattern=["model.visual.deepstack_merger_list.{block}.linear_fc2.weight"],
                max_blocks=3,
            ),
            WeightTarget(
                to_pattern="deepstack_merger_list.{block}.linear_fc2.bias",
                from_pattern=["model.visual.deepstack_merger_list.{block}.linear_fc2.bias"],
                max_blocks=3,
            ),
        ]
