"""Low-rank compressed causal LM for HuggingFace Hub.

This file is **self-contained** (depends only on ``torch`` and
``transformers``) and is copied verbatim into every published model
repository so that ``trust_remote_code=True`` works without installing
extra packages.

Usage::

    from transformers import AutoModelForCausalLM, AutoTokenizer

    model = AutoModelForCausalLM.from_pretrained(
        "your-org/compressed-llama-3-8b",
        trust_remote_code=True,
        device_map="auto",
    )
    tokenizer = AutoTokenizer.from_pretrained(
        "your-org/compressed-llama-3-8b",
    )

    inputs = tokenizer("The capital of France is", return_tensors="pt").to(model.device)
    print(tokenizer.decode(model.generate(**inputs, max_new_tokens=32)[0]))
"""

import inspect
import types

import torch
from torch import nn
from transformers import AutoModelForCausalLM, PreTrainedModel

# Transformers has moved GenerationMixin over time.  Import it defensively so
# this remote-code file keeps working across old and new Transformers versions.
try:  # current/newer Transformers
    from transformers.generation import GenerationMixin
except Exception:  # older Transformers
    try:
        from transformers.generation_utils import GenerationMixin
    except Exception:
        GenerationMixin = object

# Relative import works when HuggingFace downloads both files into a
# temporary package.  The absolute fallback covers standalone execution.
try:
    from .configuration_low_rank import LowRankConfig
except ImportError:
    from configuration_low_rank import LowRankConfig


# -----------------------------------------------------------------------
#  Low-rank linear layer
# -----------------------------------------------------------------------

class LowRankLinear(nn.Module):
    r"""Drop-in replacement for :class:`nn.Linear` that stores a rank-*r*
    factorisation instead of the full weight matrix.

    .. math::
        y = B\,(A\,x) + \text{bias}

    where :math:`A \in \mathbb{R}^{r \times k}` (``mod_a``) and
    :math:`B \in \mathbb{R}^{m \times r}` (``mod_b``).
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int,
        bias: bool = True,
        device=None,
        dtype=None,
    ):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.rank = rank

        factory_kwargs = {}
        if device is not None:
            factory_kwargs["device"] = device
        if dtype is not None:
            factory_kwargs["dtype"] = dtype

        self.mod_a = nn.Linear(
            in_features, rank, bias=False, **factory_kwargs
        )
        self.mod_b = nn.Linear(
            rank, out_features, bias=False, **factory_kwargs
        )
        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features, **factory_kwargs))
        else:
            self.register_parameter("bias", None)

    def _match_input_dtype_and_device(self, x: torch.Tensor):
        """Keep the low-rank factors compatible with the current activation.

        ``device_map="auto"`` / Accelerate may move modules independently,
        and compressed factors may be initialized or stored in fp32 while the
        base model runs in fp16/bf16.  ``F.linear`` requires activation and
        weight dtypes to match, so cast this replacement layer lazily.  The
        check makes this a one-time cost in normal inference, not a per-token
        conversion.
        """
        if not torch.is_floating_point(x):
            return

        target_device = x.device
        target_dtype = x.dtype
        weights = (self.mod_a.weight, self.mod_b.weight)

        needs_cast = any(
            w.device != target_device or w.dtype != target_dtype
            for w in weights
        )
        if self.bias is not None:
            needs_cast = (
                needs_cast
                or self.bias.device != target_device
                or self.bias.dtype != target_dtype
            )

        if needs_cast:
            self.to(device=target_device, dtype=target_dtype)

    def forward(self, x):
        self._match_input_dtype_and_device(x)
        x = self.mod_a(x)
        x = self.mod_b(x)
        if self.bias is not None:
            x = x + self.bias
        return x

    @property
    def weight(self):
        """Materialise the full weight :math:`W \\approx B A` (read-only)."""
        return (self.mod_b.weight @ self.mod_a.weight).detach()

    def extra_repr(self):
        return (
            f"in_features={self.in_features}, out_features={self.out_features}, "
            f"rank={self.rank}, bias={self.bias is not None}"
        )


# -----------------------------------------------------------------------
#  Helpers
# -----------------------------------------------------------------------

def _resolve_torch_dtype(value):
    """Resolve common config/checkpoint dtype spellings to ``torch.dtype``.

    Depending on Transformers version and how the original config was saved,
    ``torch_dtype`` may appear as ``"float16"``, ``"torch.float16"``, an
    actual ``torch.dtype``, or be absent.
    """
    if value is None:
        return None
    if isinstance(value, torch.dtype):
        return value

    value = str(value).replace("torch.", "").lower()
    return {
        "float16": torch.float16,
        "fp16": torch.float16,
        "half": torch.float16,
        "bfloat16": torch.bfloat16,
        "bf16": torch.bfloat16,
        "float32": torch.float32,
        "fp32": torch.float32,
        "float": torch.float32,
    }.get(value)


def _replace_linear_with_low_rank(model, layer_name: str, rank: int):
    """Navigate to *layer_name* and swap the ``nn.Linear`` for a
    :class:`LowRankLinear` of the given *rank*."""
    parts = layer_name.split(".")
    parent = model
    for part in parts[:-1]:
        parent = getattr(parent, part)
    attr = parts[-1]
    orig = getattr(parent, attr)
    if not isinstance(orig, nn.Linear):
        raise ValueError(
            f"Expected nn.Linear at '{layer_name}', got {type(orig).__name__}"
        )
    replacement = LowRankLinear(
        orig.in_features, orig.out_features, rank,
        bias=(orig.bias is not None),
        device=orig.weight.device,
        dtype=orig.weight.dtype,
    )
    setattr(parent, attr, replacement)
    return replacement


def _fallback_forward_signature():
    """A permissive causal-LM-ish signature used only if introspection fails."""
    P = inspect.Parameter
    return inspect.Signature([
        P("input_ids", P.POSITIONAL_OR_KEYWORD, default=None),
        P("attention_mask", P.POSITIONAL_OR_KEYWORD, default=None),
        P("position_ids", P.POSITIONAL_OR_KEYWORD, default=None),
        P("past_key_values", P.POSITIONAL_OR_KEYWORD, default=None),
        P("inputs_embeds", P.POSITIONAL_OR_KEYWORD, default=None),
        P("labels", P.POSITIONAL_OR_KEYWORD, default=None),
        P("use_cache", P.POSITIONAL_OR_KEYWORD, default=None),
        P("output_attentions", P.POSITIONAL_OR_KEYWORD, default=None),
        P("output_hidden_states", P.POSITIONAL_OR_KEYWORD, default=None),
        P("return_dict", P.POSITIONAL_OR_KEYWORD, default=None),
        # Newer decoder-only models commonly accept these.
        P("cache_position", P.POSITIONAL_OR_KEYWORD, default=None),
        P("logits_to_keep", P.POSITIONAL_OR_KEYWORD, default=None),
        P("kwargs", P.VAR_KEYWORD),
    ])


def _fallback_prepare_inputs_signature():
    """A permissive generation-preparation signature used if introspection fails."""
    P = inspect.Parameter
    return inspect.Signature([
        P("input_ids", P.POSITIONAL_OR_KEYWORD),
        P("past_key_values", P.POSITIONAL_OR_KEYWORD, default=None),
        P("attention_mask", P.POSITIONAL_OR_KEYWORD, default=None),
        P("inputs_embeds", P.POSITIONAL_OR_KEYWORD, default=None),
        P("cache_position", P.POSITIONAL_OR_KEYWORD, default=None),
        P("use_cache", P.POSITIONAL_OR_KEYWORD, default=None),
        P("kwargs", P.VAR_KEYWORD),
    ])


def _as_bound_method_signature(signature: inspect.Signature) -> inspect.Signature:
    """Return a signature suitable for ``types.MethodType``.

    ``inspect.signature(instance.method)`` removes the first parameter from a
    bound method.  The wrapped model's bound methods are already missing
    ``self``, so we prepend a synthetic self parameter before attaching the
    signature to our proxy function.
    """
    params = list(signature.parameters.values())
    if not params or params[0].name not in {"self", "_self"}:
        self_param = inspect.Parameter("self", inspect.Parameter.POSITIONAL_OR_KEYWORD)
        params = [self_param, *params]
    return signature.replace(parameters=params)


def _safe_signature(callable_obj, fallback_factory) -> inspect.Signature:
    try:
        signature = inspect.signature(callable_obj)
    except (TypeError, ValueError):
        signature = fallback_factory()
    return _as_bound_method_signature(signature)


def _install_signature_preserving_proxy(instance, public_name, target, fallback_factory):
    """Install an instance-level proxy whose visible signature mirrors target.

    Transformers' generation code validates ``model_kwargs`` by inspecting
    ``prepare_inputs_for_generation`` and sometimes ``forward``.  A wrapper with
    ``*args, **kwargs`` can therefore reject valid arguments like ``input_ids``.
    Mirroring the wrapped model's current signature keeps this file compatible
    with new model classes and new generation kwargs without hand-maintaining a
    giant static signature.
    """
    if target is None:
        return

    signature = _safe_signature(target, fallback_factory)

    if public_name == "forward":
        def proxy(_self, *args, **kwargs):
            return _self.wrapped_model(*args, **kwargs)
    else:
        def proxy(_self, *args, **kwargs):
            return getattr(_self.wrapped_model, public_name)(*args, **kwargs)

    proxy.__name__ = public_name
    proxy.__qualname__ = f"{type(instance).__name__}.{public_name}"
    proxy.__doc__ = getattr(target, "__doc__", None)
    proxy.__signature__ = signature

    # Bypass nn.Module.__setattr__; this is a plain Python method proxy, not a
    # parameter, buffer, or submodule.
    object.__setattr__(instance, public_name, types.MethodType(proxy, instance))


# -----------------------------------------------------------------------
#  Base class with future-proof GenerationMixin handling
# -----------------------------------------------------------------------

if isinstance(GenerationMixin, type) and not issubclass(PreTrainedModel, GenerationMixin):
    class _LowRankPreTrainedModel(PreTrainedModel, GenerationMixin):
        pass
else:
    class _LowRankPreTrainedModel(PreTrainedModel):
        pass


# -----------------------------------------------------------------------
#  Wrapper model
# -----------------------------------------------------------------------

class LowRankCausalLM(_LowRankPreTrainedModel):
    """Thin wrapper around *any* HuggingFace causal LM whose linear layers
    have been partially replaced with :class:`LowRankLinear` modules.

    The wrapped base model is stored as ``self.wrapped_model`` so that the
    state-dict keys are prefixed with ``wrapped_model.`` — this makes the
    mapping between saved weights and module paths unambiguous.
    """

    config_class = LowRankConfig
    base_model_prefix = "wrapped_model"
    supports_gradient_checkpointing = True
    _supports_cache_class = True

    def __init__(self, config: LowRankConfig):
        super().__init__(config)

        from transformers import CONFIG_MAPPING

        # Reconstruct the base model's config object.
        base_config_dict = dict(config.base_config)
        for drop_key in ("auto_map",):
            base_config_dict.pop(drop_key, None)

        config_cls = CONFIG_MAPPING[config.base_model_type]
        base_config = config_cls.from_dict(base_config_dict)

        # Ensure tie_word_embeddings is explicitly set from our saved config
        if "tie_word_embeddings" in base_config_dict:
            base_config.tie_word_embeddings = base_config_dict["tie_word_embeddings"]

        # Keep generation-facing flags available on the wrapper config.  The
        # wrapper still keeps LowRankConfig as its config object for save/load,
        # but generation utilities often read these common attributes.
        for attr in (
            "is_encoder_decoder",
            "pad_token_id",
            "bos_token_id",
            "eos_token_id",
            "decoder_start_token_id",
            "vocab_size",
            "max_position_embeddings",
        ):
            if hasattr(base_config, attr) and getattr(self.config, attr, None) is None:
                setattr(self.config, attr, getattr(base_config, attr))

        # Resolve torch_dtype from the stored config.  The exact representation
        # differs across Transformers versions.
        torch_dtype = _resolve_torch_dtype(base_config_dict.get("torch_dtype"))
        extra_kwargs = {"torch_dtype": torch_dtype} if torch_dtype else {}

        # Create the base causal-LM architecture (random weights — the
        # ``from_pretrained`` machinery will overwrite them afterwards).
        self.wrapped_model = AutoModelForCausalLM.from_config(
            base_config, **extra_kwargs,
        )

        # Ensure weight tying matches the original model.  If the base config
        # says *not* to tie, break any default tie that from_config created.
        if not getattr(base_config, "tie_word_embeddings", True):
            self.wrapped_model.tie_word_embeddings = False
            out_emb = self.wrapped_model.get_output_embeddings()
            in_emb = self.wrapped_model.get_input_embeddings()
            if out_emb is not None and in_emb is not None:
                if out_emb.weight.data_ptr() == in_emb.weight.data_ptr():
                    out_emb.weight = nn.Parameter(out_emb.weight.clone())

        # Replace every compressed layer with a LowRankLinear module.
        for layer_name, rank in (config.rank_dict or {}).items():
            _replace_linear_with_low_rank(
                self.wrapped_model, layer_name, int(rank),
            )

        # Preserve generation compatibility by making this wrapper look like
        # the current inner model to Transformers' signature-based validators.
        _install_signature_preserving_proxy(
            self,
            "forward",
            self.wrapped_model.forward,
            _fallback_forward_signature,
        )
        _install_signature_preserving_proxy(
            self,
            "prepare_inputs_for_generation",
            getattr(self.wrapped_model, "prepare_inputs_for_generation", None),
            _fallback_prepare_inputs_signature,
        )

    @property
    def main_input_name(self):
        return getattr(self.wrapped_model, "main_input_name", "input_ids")

    # ------------------------------------------------------------------
    #  Forward/generation fallbacks — normally replaced in __init__
    # ------------------------------------------------------------------

    def forward(self, *args, **kwargs):
        return self.wrapped_model(*args, **kwargs)

    def prepare_inputs_for_generation(self, *args, **kwargs):
        return self.wrapped_model.prepare_inputs_for_generation(*args, **kwargs)

    def generate(self, *args, **kwargs):
        """Delegate generation to the wrapped HF model.

        This is the most stable path across Transformers releases because it
        uses the official model class's own ``generate`` implementation,
        validation logic, cache handling, and generation signature.
        """
        if hasattr(self, "generation_config"):
            self.wrapped_model.generation_config = self.generation_config
        return self.wrapped_model.generate(*args, **kwargs)

    # ------------------------------------------------------------------
    #  Embedding and cache helpers
    # ------------------------------------------------------------------

    def get_input_embeddings(self):
        return self.wrapped_model.get_input_embeddings()

    def set_input_embeddings(self, value):
        self.wrapped_model.set_input_embeddings(value)

    def get_output_embeddings(self):
        return self.wrapped_model.get_output_embeddings()

    def set_output_embeddings(self, value):
        self.wrapped_model.set_output_embeddings(value)

    def can_generate(self):
        return True

    def tie_weights(self):
        """Only tie weights if the base model config explicitly requests it."""
        base_tie = self.config.base_config.get("tie_word_embeddings", False)
        if base_tie:
            super().tie_weights()

    def _reorder_cache(self, *args, **kwargs):
        if hasattr(self.wrapped_model, "_reorder_cache"):
            return self.wrapped_model._reorder_cache(*args, **kwargs)
        return None