# Copyright 2026-present the HuggingFace Inc. team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import warnings

import torch
import torch.nn as nn
from transformers.pytorch_utils import Conv1D

from peft.tuners.tuners_utils import BaseTuner, BaseTunerLayer
from peft.utils import (
    TRANSFORMERS_MODELS_TO_TINYLORA_TARGET_MODULES_MAPPING,
)

from ..tuners_utils import _maybe_include_all_linear_layers
from .config import TinyLoraConfig
from .layer import Embedding, Linear, TinyLoraLayer


class TinyLoraModel(BaseTuner):
    """
    Creates TinyLoRA model from a pretrained transformers model.

    TinyLoRA is an extremely parameter-efficient fine-tuning method that uses SVD decomposition of frozen weights and
    projects a tiny trainable vector through fixed random tensors. Based on the paper "Learning to Reason in 13
    Parameters" (arXiv:2602.04118).

    Args:
        model ([`~transformers.PreTrainedModel`]): The model to be adapted.
        config ([`TinyLoraConfig`]): The configuration of the TinyLoRA model.
        adapter_name (`str`): The name of the adapter, defaults to `"default"`.
        low_cpu_mem_usage (`bool`, *optional*, defaults to `False`):
            Create empty adapter weights on meta device. Useful to speed up the loading process.

    Returns:
        `torch.nn.Module`: The TinyLoRA model.

    Example:
        ```python
        >>> from transformers import AutoModelForCausalLM
        >>> from peft import TinyLoraConfig, get_peft_model

        >>> base_model = AutoModelForCausalLM.from_pretrained("facebook/opt-125m")
        >>> config = TinyLoraConfig(r=2, u=64, target_modules=["q_proj", "v_proj"])
        >>> model = get_peft_model(base_model, config)
        ```

    **Attributes**:
        - **model** ([`~transformers.PreTrainedModel`]) -- The model to be adapted.
        - **peft_config** ([`TinyLoraConfig`]): The configuration of the TinyLoRA model.
    """

    prefix: str = "tinylora_"
    tuner_layer_cls = TinyLoraLayer
    target_module_mapping = TRANSFORMERS_MODELS_TO_TINYLORA_TARGET_MODULES_MAPPING

    def __init__(self, model, config, adapter_name, low_cpu_mem_usage=False, **kwargs):
        super().__init__(model, config, adapter_name, low_cpu_mem_usage, **kwargs)

    def _init_tinylora_v(self, config: TinyLoraConfig, adapter_name: str) -> None:
        """Re-initialize the tinylora_v vectors with uniform random values."""
        if adapter_name in self.tinylora_v:
            for v in self.tinylora_v[adapter_name].values():
                nn.init.uniform_(v, -config.init_v_bound, config.init_v_bound)

    def _pre_injection_hook(self, model: nn.Module, config: TinyLoraConfig, adapter_name: str) -> None:
        """Initialize shared trainable vectors on first adapter creation."""
        # Nested structure: tinylora_v[adapter_name][str(group_idx)] = nn.Parameter
        if not hasattr(self, "tinylora_v"):
            self.tinylora_v = nn.ModuleDict({})

    def _build_target_key_mapping(self, config: TinyLoraConfig) -> dict[str, int]:
        """Build an ordered mapping from target module key to index.

        Iterates the model in the same order as ``inject_adapter`` to assign each target module a deterministic index
        used for group assignment (weight_tying) and projection seeding.
        """
        model_config = self.get_model_config(self.model)
        peft_config = self._prepare_adapter_config(config, model_config)
        peft_config = _maybe_include_all_linear_layers(peft_config, self.model)

        # Also match TinyLoraLayer since modules may already be wrapped when adding a second adapter
        target_types = (nn.Linear, Conv1D, nn.Embedding, TinyLoraLayer)

        mapping: dict[str, int] = {}
        idx = 0
        for key, module in self.model.named_modules():
            if not self._check_target_module_exists(peft_config, key):
                continue
            if isinstance(module, target_types):
                mapping[key] = idx
                idx += 1

        return mapping

    def _check_new_adapter_config(self, config: TinyLoraConfig) -> None:
        """Check the config when a new adapter is being added."""
        super()._check_new_adapter_config(config)

        save_projection_unique_values = sorted({c.save_projection for c in self.peft_config.values()})
        if len(save_projection_unique_values) > 1:
            raise ValueError(
                "TinyLoRA projection tensors must be saved for all adapters or none, but got multiple different values: "
                f"{save_projection_unique_values}"
            )

    def _create_and_replace(
        self,
        tinylora_config: TinyLoraConfig,
        adapter_name: str,
        target: nn.Module,
        target_name: str,
        parent: nn.Module,
        current_key: str,
        **optional_kwargs,
    ):
        if current_key is None:
            raise ValueError("Current Key shouldn't be `None`")

        # Build the target key mapping lazily on first call per injection cycle.
        # This is needed because add_adapter calls inject_adapter directly without _pre_injection_hook.
        if not hasattr(self, "_target_key_to_idx") or current_key not in self._target_key_to_idx:
            self._target_key_to_idx = self._build_target_key_mapping(tinylora_config)

        # Look up the deterministic index for this module
        layer_idx = self._target_key_to_idx[current_key]
        num_target_layers = len(self._target_key_to_idx)

        # Determine the group for this layer based on weight_tying
        # weight_tying=0.0 → num_groups = num_target_layers (no sharing)
        # weight_tying=1.0 → num_groups = 1 (full sharing)
        num_groups = max(1, round(num_target_layers * (1.0 - tinylora_config.weight_tying)))
        group_size = max(1, num_target_layers // num_groups)
        group_idx = min(layer_idx // group_size, num_groups - 1)
        v_key = str(group_idx)

        # Initialize the adapter's ParameterDict if not present
        if adapter_name not in self.tinylora_v:
            self.tinylora_v[adapter_name] = nn.ParameterDict({})

        # Initialize v for this group if not already done
        if v_key not in self.tinylora_v[adapter_name]:
            # Get dtype from target layer's weight
            if hasattr(target, "weight"):
                dtype = target.weight.dtype
            else:
                dtype = None  # Will default to float32
            v = nn.Parameter(torch.empty(tinylora_config.u, dtype=dtype))
            if tinylora_config.init_weights is True:
                # Default: initialize to zeros for identity/no-op operation
                nn.init.zeros_(v)
            elif tinylora_config.init_weights == "uniform":
                nn.init.uniform_(v, -tinylora_config.init_v_bound, tinylora_config.init_v_bound)
            # If init_weights is False, leave v uninitialized
            self.tinylora_v[adapter_name][v_key] = v

        if isinstance(target, TinyLoraLayer):
            target.set_layer_idx(layer_idx)
            target.update_layer(
                adapter_name,
                self.tinylora_v,
                v_key,
                tinylora_config.r,
                tinylora_config,
            )
        else:
            new_module = self._create_new_module(tinylora_config, self.tinylora_v, v_key, adapter_name, target)
            new_module.set_layer_idx(layer_idx)

            if adapter_name not in self.active_adapter:
                # adding an additional adapter: it is not automatically trainable
                new_module.requires_grad_(False)
            self._replace_module(parent, target_name, new_module, target)

        # Ensure the shared v parameters remain trainable for the active adapter,
        # but only if the adapter is not in inference mode
        if adapter_name in self.active_adapter:
            inference_mode = getattr(tinylora_config, "inference_mode", False)
            if not inference_mode:
                for param in self.tinylora_v[adapter_name].values():
                    param.requires_grad = True

    @staticmethod
    def _create_new_module(
        tinylora_config: TinyLoraConfig,
        tinylora_v: nn.ModuleDict,
        v_key: str,
        adapter_name: str,
        target: nn.Module,
        **kwargs,
    ):
        if isinstance(target, BaseTunerLayer):
            target_base_layer = target.get_base_layer()
        else:
            target_base_layer = target

        if isinstance(target_base_layer, torch.nn.Linear):
            if tinylora_config.fan_in_fan_out:
                warnings.warn(
                    "fan_in_fan_out is set to True but the target module is `torch.nn.Linear`. "
                    "Setting fan_in_fan_out to False."
                )
                tinylora_config.fan_in_fan_out = False
            new_module = Linear(
                target,
                tinylora_v,
                v_key,
                adapter_name,
                tinylora_config,
                **kwargs,
            )
        elif isinstance(target_base_layer, Conv1D):
            kwargs["is_target_conv_1d_layer"] = True
            if not tinylora_config.fan_in_fan_out:
                warnings.warn(
                    "fan_in_fan_out is set to False but the target module is `Conv1D`. Setting fan_in_fan_out to True."
                )
                tinylora_config.fan_in_fan_out = True
            new_module = Linear(
                target,
                tinylora_v,
                v_key,
                adapter_name,
                tinylora_config,
                **kwargs,
            )
        elif isinstance(target_base_layer, torch.nn.Embedding):
            new_module = Embedding(
                target,
                tinylora_v,
                v_key,
                adapter_name,
                tinylora_config,
                **kwargs,
            )
        else:
            raise ValueError(
                f"Target module {target} is not supported. Currently, only the following modules are supported: "
                "`torch.nn.Linear`, `torch.nn.Embedding`, `transformers.pytorch_utils.Conv1D`."
            )

        return new_module

    def _cast_adapter_dtype(self, adapter_name: str, autocast_adapter_dtype: bool = True) -> None:
        """
        Cast the adapter weights to the correct dtype.

        Override to also handle the model-level tinylora_v parameters.
        """
        # Call parent implementation for layer-level parameters
        super()._cast_adapter_dtype(adapter_name, autocast_adapter_dtype)

        if not autocast_adapter_dtype:
            return

        # Handle model-level tinylora_v parameters
        dtypes_to_convert_to_fp32 = {torch.float16, torch.bfloat16}
        if adapter_name in self.tinylora_v:
            for param in self.tinylora_v[adapter_name].values():
                if param.dtype in dtypes_to_convert_to_fp32:
                    param.data = param.data.to(torch.float32)

    def delete_adapter(self, adapter_name: str) -> None:
        """Delete an adapter and clean up the model-level shared v parameters."""
        super().delete_adapter(adapter_name)

        # Remove the adapter's shared v parameters from the model-level ModuleDict
        if adapter_name in self.tinylora_v:
            del self.tinylora_v[adapter_name]

    def _mark_only_adapters_as_trainable(self, model: nn.Module) -> None:
        """
        Mark only the adapter layers as trainable.

        Override the base class method to manage the shared tinylora_v parameters which are stored at the model level
        and thus invisible to the base class's per-layer logic.
        """
        # First, call the parent implementation
        super()._mark_only_adapters_as_trainable(model)

        # Freeze all tinylora_v parameters first, then selectively unfreeze for active non-inference adapters
        for adapter_params in self.tinylora_v.values():
            for param in adapter_params.values():
                param.requires_grad = False

        for active_adapter in self.active_adapters:
            if active_adapter in self.peft_config:
                inference_mode = getattr(self.peft_config[active_adapter], "inference_mode", False)
                if inference_mode:
                    continue
            if active_adapter in self.tinylora_v:
                for param in self.tinylora_v[active_adapter].values():
                    param.requires_grad = True