# 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 torch
from torch import nn

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

from .config import LilyConfig
from .layer import LilyLayer, Linear


class LilyModel(BaseTuner):
    """
    Creates a Low-Rank Interconnected Adaptation Across Layers (Lily) model from a pretrained transformers model.

    The method is described in detail in https://arxiv.org/abs/2407.09946.

    Args:
        model ([`torch.nn.Module`]): The model to be adapted.
        config ([`LilyConfig`]): The configuration of the Lily model.
        adapter_name (`str`): The name of the adapter, defaults to `"default"`.

    Returns:
        `torch.nn.Module`: The Lily PEFT model.

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

    prefix: str = "lily_"
    tuner_layer_cls = LilyLayer
    target_module_mapping = TRANSFORMERS_MODELS_TO_LILY_TARGET_MODULES_MAPPING

    def _create_and_replace(
        self,
        lily_config,
        adapter_name,
        target,
        target_name,
        parent,
        current_key,
        **optional_kwargs,
    ):
        """
        Create a new Lily layer with independent A and B for each layer. Sharing of A/B across layers is deferred to
        _post_injection_hook.
        """
        if current_key is None:
            raise ValueError("Current Key shouldn't be `None`")

        if isinstance(target, LilyLayer):
            target.update_layer(
                adapter_name,
                lily_config.r,
                scaling=lily_config.scaling,
                stride_A=lily_config.stride_A,
                num_B=lily_config.num_B,
                init_weights=lily_config.init_weights,
            )
        else:
            new_module = self._create_new_module(lily_config, adapter_name, target)
            if adapter_name not in self.active_adapters:
                new_module.requires_grad_(False)
            self._replace_module(parent, target_name, new_module, target)

    @staticmethod
    def _create_new_module(lily_config, adapter_name, target):
        if isinstance(target, BaseTunerLayer):
            target_base_layer = target.get_base_layer()
        else:
            target_base_layer = target

        if isinstance(target_base_layer, torch.nn.Linear):
            return Linear(
                target,
                adapter_name,
                r=lily_config.r,
                scaling=lily_config.scaling,
                stride_A=lily_config.stride_A,
                num_B=lily_config.num_B,
                init_weights=lily_config.init_weights,
            )

        raise NotImplementedError(f"Lily does not support target modules of type {type(target_base_layer)} yet.")

    def _post_injection_hook(self, model: nn.Module, config: LilyConfig, adapter_name: str) -> None:
        """
        After all layers have been independently initialized, apply A/B sharing across layers.

        A sharing: for each (target_module_suffix, weight_shape) group, consecutive blocks of `stride_A` layers share
        the same A. The first layer in each block keeps its own A; subsequent layers in the block have their lily_A
        replaced by the group leader's.

        B sharing: all layers in the same (target_module_suffix, weight_shape) group share the B from the very first
        layer in that group.

        Both lily_A and lily_B are nn.Linear modules, so they move correctly with model.to(device) via standard
        nn.Module parameter propagation.

        Note: (target_module_suffix, weight_shape) is used as the group key rather than target_module_suffix alone, to
        correctly handle architectures like UNet where the same target key can appear with different shapes across
        layers.
        """
        stride_A = config.stride_A

        # Collect all adapted LilyLayer modules in traversal order, grouped by
        # (target_module_suffix, weight_shape).
        # Maps (target_suffix, weight_shape) -> list of LilyLayer in traversal order.
        target_to_layers: dict[tuple[str, torch.Size], list[LilyLayer]] = {}

        for key, module in model.named_modules():
            if not isinstance(module, LilyLayer):
                continue
            if adapter_name not in module.lily_A:
                continue

            base = module.get_base_layer()
            shape = base.weight.shape  # (out_features, in_features)

            # Find the longest matching target suffix for this key
            matched_suffix = None
            if isinstance(config.target_modules, str):
                matched_suffix = config.target_modules
            else:
                for suffix in config.target_modules:
                    if key.endswith(suffix):
                        matched_suffix = suffix

            if matched_suffix is None:
                # Should not happen since inject_adapter already matched this layer
                continue

            group_key = (matched_suffix, shape)
            if group_key not in target_to_layers:
                target_to_layers[group_key] = []
            target_to_layers[group_key].append(module)

        # Apply A and B sharing for each (target_suffix, shape) group.
        for (target_suffix, shape), layers in target_to_layers.items():
            # B sharing: all layers share the first layer's B
            shared_B = layers[0].lily_B[adapter_name]
            for i, layer in enumerate(layers):
                if i != 0:
                    layer.lily_B[adapter_name] = shared_B

                # A sharing: layers within the same stride_A block share the group leader's A.
                # Group leader is the first layer in each block of stride_A consecutive layers.
                group_idx = (i // stride_A) * stride_A
                if i != group_idx:
                    layer.lily_A[adapter_name] = layers[group_idx].lily_A[adapter_name]
