# ensemble_qwen.py
import torch, math
from torch import nn
from transformers import Qwen3ForCausalLM, Qwen3Config, PreTrainedModel, PretrainedConfig


from transformers.modeling_outputs import CausalLMOutputWithPast
from transformers.generation import GenerationMixin

class EnsembleConfig(Qwen3Config):
    model_type = "ensemble_qwen"
    def __init__(
        self,
        model_a_path="Qwen/Qwen3-4B",
        model_b_path="Qwen/Qwen3-4B-Base",
        weight_a=0.9,
        weight_b=0.1,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.model_a_path, self.model_b_path = model_a_path, model_b_path
        self.weight_a, self.weight_b = weight_a, weight_b


class EnsembleForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = EnsembleConfig
    _supports_flash_attn_2 = True
    _supports_sdpa        = True
    main_input_name = "input_ids"
    _tp_plan = {"model_a.lm_head": "colwise_rep", "model_b.lm_head": "colwise_rep"}
    
    def __init__(self, config: EnsembleConfig):
        super().__init__(config)
        
        self.weight_a = config.weight_a
        self.weight_b = config.weight_b
        
        # remove extra keys from config and initialzie Qwen3 model
        qwen3_config_dict = config.to_dict()
        extra_keys = ["model_a_path", "model_b_path", "weight_a", "weight_b", "auto_map"]
        for key in extra_keys:
            del qwen3_config_dict[key]
        
        qwen3_config_dict["model_type"] = "qwen3"
        qwen3_config_dict["architectures"] = ["Qwen3ForCausalLM"]
        
        if hasattr(config, "attn_implementation"):
            qwen3_config_dict["attn_implementation"] = config.attn_implementation
            
        qwen3_config = Qwen3Config(**qwen3_config_dict)
        
        self.model_a = Qwen3ForCausalLM(qwen3_config)
        self.model_b = Qwen3ForCausalLM(qwen3_config)

    # ---- core magic ----------------------------------------------------
    def forward(
        self,
        input_ids: torch.LongTensor,
        attention_mask: torch.LongTensor | None = None,
        past_key_values: tuple | None = None,
        use_cache: bool = True,
        **kwargs,
    ) -> CausalLMOutputWithPast:
        """Run both models, sum logits, carry both kv-caches forward."""
        past_a, past_b = (None, None) if past_key_values is None else past_key_values
        
        out_a = self.model_a(
            input_ids,
            attention_mask=attention_mask,
            past_key_values=past_a,
            use_cache=use_cache,
        )
        out_b = self.model_b(
            input_ids,
            attention_mask=attention_mask,
            past_key_values=past_b,
            use_cache=use_cache,
        )

        # blend logits
        logits = self.weight_a * out_a.logits + self.weight_b * out_b.logits

        # pack both pasts together so generate() can recycle them next step
        past = (out_a.past_key_values, out_b.past_key_values) if use_cache else None

        return CausalLMOutputWithPast(
            logits=logits,
            past_key_values=past,
            hidden_states=None,
            attentions=None,
        )

    # ---- wiring for HF generation loop ---------------------------------
    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
        # HF expects past_key_values shape to match forward() signature
        if past_key_values is not None:
            past_a, past_b = past_key_values
            input_ids = input_ids[:, -1:]     # only last token in incremental mode
        else:
            past_a = past_b = None
        return {
            "input_ids": input_ids,
            "past_key_values": (past_a, past_b),
            "use_cache": kwargs.get("use_cache", True),
        }

    def _reorder_cache(self, past, beam_idx):
        # support beam search ↔ reorder both caches identically
        past_a, past_b = past
        past_a = self.model_a._reorder_cache(past_a, beam_idx)
        past_b = self.model_b._reorder_cache(past_b, beam_idx)
        return (past_a, past_b)