from transformers.models.phi3.modeling_phi3 import Phi3ForCausalLM
from transformers import AutoConfig, AutoModelForCausalLM
from .configuration_maplept import MaplePTConfig
import torch.nn as nn


class MaplePTForCausalLM(Phi3ForCausalLM):
    config_class = MaplePTConfig
    model_type = "maplept"

    def __init__(self, config):
        if not isinstance(config, MaplePTConfig):
            config = MaplePTConfig.from_dict(config.to_dict())
        super().__init__(config)

        base = getattr(self, "model", None)
        if base is None:
            return

        # --- ensure embed_tokens alias ---
        if hasattr(base, "embed_tokens"):
            self.embed_tokens = base.embed_tokens

        # --- ensure model.layers alias (vLLM critical) ---
        # vLLM searches model.layers.*, but Phi3 nests them in model.model.layers
        if not hasattr(base, "layers"):
            try:
                base.layers = base.model.layers
                print("✅ Aliased base.layers → base.model.layers")
            except AttributeError:
                print("⚠️ Could not alias layers (missing model.model.layers)")

        # --- ensure top-level direct alias too ---
        if not hasattr(self, "layers") and hasattr(base, "layers"):
            self.layers = base.layers
            print("✅ Aliased self.layers → model.layers")

        # --- ensure normalization aliases (for completeness) ---
        if not hasattr(base, "norm"):
            if hasattr(base, "final_layernorm"):
                base.norm = base.final_layernorm
            elif hasattr(base, "model") and hasattr(base.model, "final_layernorm"):
                base.norm = base.model.final_layernorm
        if not hasattr(self, "norm") and hasattr(base, "norm"):
            self.norm = base.norm