# Copyright 2024-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 math
from operator import attrgetter
from typing import Optional

import torch

from peft.config import PeftConfig
from peft.mapping import PEFT_TYPE_TO_CONFIG_MAPPING, PEFT_TYPE_TO_PREFIX_MAPPING
from peft.tuners.lora import Conv2d, Linear, LoraConfig, LoraLayer

from .other import get_pattern_key, infer_device
from .peft_types import PeftType
from .save_and_load import _insert_adapter_name_into_state_dict, load_peft_weights


# so far only LoRA is supported
CONFIG_KEYS_TO_CHECK = {PeftType.LORA: ["use_rslora", "lora_dropout", "alpha_pattern", "use_dora"]}


def _update_scaling(lora_module, adapter_name, scaling=None):
    """
    Update the value of the scalings of the LoRA module.

    Takes into consideration that scalings can be tensors from prepare_model_for_compiled_hotswap.
    """
    if lora_module.scaling[adapter_name] == scaling:
        return

    if isinstance(lora_module.scaling[adapter_name], torch.Tensor):
        lora_module.scaling[adapter_name].fill_(scaling)
    elif isinstance(lora_module.scaling[adapter_name], (float, int)):
        lora_module.scaling[adapter_name] = scaling
    else:
        raise ValueError(
            "Something went wrong when trying to set the new scale value, expected to find the old value to be of type "
            f"float or torch.Tensor, got {type(lora_module.scaling[adapter_name])} instead."
        )


def _convert_scalings_to_tensor(model):
    """
    Convert the LoRA scaling values into torch.tensors to prevent recompilation if they change.
    """
    for module in model.modules():
        if not isinstance(module, LoraLayer):
            continue

        scaling = module.scaling
        for key, val in scaling.items():
            if isinstance(val, float):
                # no need to deal with dtype as scalars are coerced
                scaling[key] = torch.tensor(val, device=module.weight.device)
            elif not isinstance(val, torch.Tensor):
                raise ValueError(
                    "Something went wrong while trying to convert the scalings, expected to find values of type float "
                    f"but found {type(val)} instead."
                )


def _get_padded_linear(lora_module: torch.nn.Module, target_rank: int, is_lora_A: bool) -> torch.nn.Linear:
    """
    Get a new Linear layer for LoRA with padded weights according to the target rank.

    Args:
        lora_module (nn.Module):
            The LoRA sub-module (e.g. module.lora_A[adapter_name]).
        target_rank (int):
            The desired rank to pad to.
        is_lora_A (bool):
            True if this is the LoRA A matrix, False if LoRA B.

    Returns:
        nn.Linear:
            A newly created and padded Linear layer. If the rank already fit, the original layer is returned.
    """
    weight = lora_module.weight
    # For LoRA A, the "rank dimension" is weight.size(0) (out_features).
    # For LoRA B, it is weight.size(1) (in_features).
    original_rank = weight.size(0) if is_lora_A else weight.size(1)

    # If no padding needed
    if original_rank == target_rank:
        return lora_module

    if original_rank > target_rank:
        raise ValueError(
            f"Trying to pad the adapter to the target rank {target_rank}, but the original rank is larger "
            f"({original_rank}). This is not possible."
        )

    out_features, in_features = weight.shape

    # lora_A and lora_B are always nn.Linear
    if is_lora_A:
        # LoRA A affects out_features
        padded = torch.zeros(target_rank, in_features, device=weight.device, dtype=weight.dtype)
        padded[:original_rank, :] = weight
        new_layer = torch.nn.Linear(in_features, target_rank, bias=lora_module.bias is not None)
    else:
        # LoRA B affects in_features
        padded = torch.zeros(out_features, target_rank, device=weight.device, dtype=weight.dtype)
        padded[:, :original_rank] = weight
        new_layer = torch.nn.Linear(target_rank, out_features, bias=lora_module.bias is not None)

    # Sanity check
    if new_layer.weight.shape != padded.shape:
        raise ValueError(
            "Something went wrong when trying to pad the LoRA Linear weights, the new shape should be "
            f"{padded.shape} but {new_layer.weight.shape} was found. Please open an issue on PEFT "
            "(https://github.com/huggingface/peft/issues) and report this error."
        )
    if (lora_module.bias is not None) and (new_layer.bias.shape != lora_module.bias.shape):
        raise ValueError(
            "Something went wrong when trying to pad the LoRA Linear bias, the new shape should be "
            f"{lora_module.bias.shape} but {new_layer.bias.shape} was found. Please open an issue on PEFT "
            "(https://github.com/huggingface/peft/issues) and report this error."
        )

    new_layer.weight.data = padded
    # Copy bias if present
    if lora_module.bias is not None:
        new_layer.bias.data = lora_module.bias.data

    return new_layer


def _get_padded_conv2d(lora_module: torch.nn.Module, target_rank: int, is_lora_A: bool) -> torch.nn.Conv2d:
    """
    Get a new Conv2d layer for LoRA with padded weights according to the target rank.

    Args:
        lora_module (nn.Module):
            The LoRA sub-module (e.g. module.lora_A[adapter_name]).
        target_rank (int):
            The desired rank to pad to.
        is_lora_A (bool):
            True if this is the LoRA A matrix, False if LoRA B.

    Returns:
        nn.Conv2d:
            A newly created and padded Conv2d layer. If the rank already fit, the original layer is returned.
    """
    weight = lora_module.weight
    # For Conv2d: [out_channels, in_channels, kernel_height, kernel_width]
    out_channels, in_channels, kh, kw = weight.shape
    original_rank = out_channels if is_lora_A else in_channels

    if original_rank == target_rank:
        return lora_module

    if original_rank > target_rank:
        raise ValueError(
            f"Trying to pad the adapter to the target rank {target_rank}, but the original rank is larger "
            f"({original_rank}). This is not possible."
        )

    # lora_A and lora_B are always nn.Conv2d
    if is_lora_A:
        # LoRA A affects out_channels
        padded = torch.zeros(target_rank, in_channels, kh, kw, device=weight.device, dtype=weight.dtype)
        padded[:out_channels, :, :, :] = weight
        new_layer = torch.nn.Conv2d(
            in_channels,
            target_rank,
            kernel_size=lora_module.kernel_size,
            stride=lora_module.stride,
            padding=lora_module.padding,
            bias=lora_module.bias is not None,
            groups=lora_module.groups,
        )
    else:
        # LoRA B affects in_channels
        padded = torch.zeros(out_channels, target_rank, kh, kw, device=weight.device, dtype=weight.dtype)
        padded[:, :in_channels, :, :] = weight
        new_layer = torch.nn.Conv2d(
            target_rank,
            out_channels,
            kernel_size=lora_module.kernel_size,
            stride=lora_module.stride,
            padding=lora_module.padding,
            bias=lora_module.bias is not None,
            groups=lora_module.groups,
        )

    # Sanity check
    if new_layer.weight.shape != padded.shape:
        raise ValueError(
            "Something went wrong when trying to pad the LoRA  weights, the new shape should be "
            f"{padded.shape} but {new_layer.weight.shape} was found. Please open an issue on PEFT "
            "(https://github.com/huggingface/peft/issues) and report this error."
        )
    if (lora_module.bias is not None) and (new_layer.bias.shape != lora_module.bias.shape):
        raise ValueError(
            "Something went wrong when trying to pad the LoRA Conv2d bias, the new shape should be "
            f"{lora_module.bias.shape} but {new_layer.bias.shape} was found. Please open an issue on PEFT "
            "(https://github.com/huggingface/peft/issues) and report this error."
        )

    new_layer.weight.data = padded
    # Copy bias if present
    if lora_module.bias is not None:
        new_layer.bias.data = lora_module.bias.data

    return new_layer


def _pad_lora_weights(model: torch.nn.Module, target_rank: int) -> None:
    """
    Pad LoRA weights in a model to a target rank while preserving the original behavior.

    Args:
      model (nn.Module): The model containing LoRA modules (with lora_A and lora_B).
      target_rank (int): The target rank to pad to.
    """

    for module in model.modules():
        # Decide which pad function to call based on module type
        if isinstance(module, Linear):
            pad_fn = _get_padded_linear
        elif isinstance(module, Conv2d):
            pad_fn = _get_padded_conv2d
        else:
            # Skip any other module types
            continue

        # Pad LoRA A
        for adapter_name, lora_A_module in module.lora_A.items():
            new_layer = pad_fn(lora_A_module, target_rank=target_rank, is_lora_A=True)
            module.lora_A[adapter_name] = new_layer

        # Pad LoRA B
        for adapter_name, lora_B_module in module.lora_B.items():
            new_layer = pad_fn(lora_B_module, target_rank=target_rank, is_lora_A=False)
            module.lora_B[adapter_name] = new_layer


def prepare_model_for_compiled_hotswap(
    model: torch.nn.Module,
    *,
    target_rank: Optional[int] = None,
    config: Optional[LoraConfig | dict[str, LoraConfig]] = None,
) -> None:
    """
    Helper function that prepares the model so that it can later be compiled and then used with hot-swapping.

    It is necessary to call this function on the model for hot-swapping to work if both of these are true:

    - the different LoRA adapters have different ranks and/or different alpha values (i.e. scalings)
    - you plan to torch.compile the model and want to avoid re-compilation

    It is important to call this function *after* the first LoRA adapter has been loaded (i.e. the one that will be
    swapped out) but *before* the model is compiled.

    Even with this function, hot-swapping LoRA adapters that target different layers is still not supported.

    Note: This function modifies the model in-place. If you want to restore the model to its initial state, you will
    have to reload it.

    Args:
        model (`nn.Module`):
            The model with the loaded adapter, before compilation.
        target_rank (`int`, *optional*):
            The target rank to pad the LoRA weights to. Should be the maximum rank among all LoRA adapters that will be
            hot-swapped. If not specified, the target ranks will not be changed.
        config (`LoraConfig` or `dict[str, LoraConfig]`, *optional*):
            Optionally pass the `LoraConfig`s of the LoRA adapters. If passed, the rank in the configs will be updated
            to `target_rank`.

    Example:

        ```py
        base_model = ...
        model = PeftModel.from_pretrained(base_model, path_adapter_0)
        # Prepare the model to allow hotswapping even if ranks/scalings of 2nd adapter differ.
        # You can skip this step if all ranks and scalings are identical.
        prepare_model_for_compiled_hotswap(model, target_rank=highest_lora_rank)
        model = torch.compile(model)
        # do inference with adapter 0
        # replace the "default" lora adapter with the new one
        hotswap_adapter(model, path_adapter_1, adapter_name="default", torch_device=device)
        # do inference with adapter 1
        ```
    """
    is_compiled = hasattr(model, "_orig_mod")
    if is_compiled:
        raise ValueError("Call prepare_model_for_compiled_hotswap *before* compiling the model")

    _convert_scalings_to_tensor(model)
    if target_rank is not None:
        _pad_lora_weights(model, target_rank=target_rank)

    if not config:
        return
    if target_rank is None:
        return

    if not isinstance(config, dict):
        # config can be either a PeftConfig, or a dict of PeftConfigs like PeftModel.peft_config
        config = {"dummy": config}

    for lora_config in config.values():
        lora_config.r = target_rank
        if lora_config.rank_pattern:
            for key in lora_config.rank_pattern:
                lora_config.rank_pattern[key] = target_rank


def hotswap_adapter_from_state_dict(
    model: torch.nn.Module,
    state_dict: dict[str, torch.Tensor],
    adapter_name: str,
    config: LoraConfig,
    parameter_prefix: str = "lora_",
):
    """
    Swap out the adapter weights from the model with the weights from state_dict.

    As of now, only LoRA is supported.

    This is a low-level function that assumes that the adapters have been checked for compatibility and that the
    state_dict has been correctly mapped to work with PEFT. For a high level function that performs this work for you,
    use `hotswap_adapter` instead.

    Args:
        model (`nn.Module`):
            The model with the loaded adapter.
        state_dict (`dict[str, torch.Tensor]`):
            The state dict of the new adapter, which needs to be compatible (targeting same modules etc.).
        adapter_name (`str`):
            The name of the adapter that should be hot-swapped, e.g. `"default"`. The name will remain the same after
            swapping.
        config (`LoraConfig`):
            The config of the LoRA adapter. This is used to determine the scaling and rank of the adapter.
        parameter_prefix (`str`, *optional*, defaults to `"lora_"`)
            The prefix used to identify the adapter's keys in the state dict. For LoRA, this would be `"lora_"` (the
            default).

    Raises:
        RuntimeError
            If the old and the new adapter are not compatible, a RuntimeError is raised.

    """
    # Ensure that all the keys of the new adapter correspond exactly to the keys of the old adapter, otherwise
    # hot-swapping is not possible

    # _orig_mod is for torch.compile(model) and _compiled_call_impl is for model.compile() (not wrapped)
    is_compiled = hasattr(model, "_orig_mod")
    is_compiled_inplace = bool(getattr(model, "_compiled_call_impl", None))
    # TODO: there is probably a more precise way to identify the adapter keys
    missing_keys = {k for k in model.state_dict() if (parameter_prefix in k) and (adapter_name in k)}
    unexpected_keys = []

    # first: dry run, not swapping anything
    for key, new_val in state_dict.items():
        try:
            old_val = attrgetter(key)(model)
        except AttributeError:
            unexpected_keys.append(key)
            continue

        if is_compiled:
            missing_keys.remove("_orig_mod." + key)
        else:
            missing_keys.remove(key)

    # Right now, we don't deal with unexpected keys, i.e. if the adapter being swapped in targeting new layers. We could
    # probably add LoRA to these layers ad hoc, but that would not work with compiled models.
    if unexpected_keys:
        msg = f"Hot swapping the adapter did not succeed, unexpected keys found: {', '.join(unexpected_keys)}."
        raise RuntimeError(msg)

    # If the adapter that is being swapped in is missing some keys, this is fine. We just need to ensure that those LoRA
    # weights from the previous adapter are set to 0 so that they don't influence the output. We don't need to worry
    # about ranks are alphas.
    for key in missing_keys:
        # in case it's a compiled model
        key = key.removeprefix("_orig_mod.")
        # get LoRA parent module name by removing the 'lora_*.<adapter-name>.weight' part
        module_name = ".".join(key.split(".")[:-3])
        module = model.get_submodule(module_name)
        old_val = attrgetter(key)(model)
        old_val.data.fill_(0.0)

    # actual swapping
    for key, new_val in state_dict.items():
        # get LoRA parent module name by removing the 'lora_*.<adapter-name>.weight' part
        module_name = ".".join(key.split(".")[:-3])
        module = model.get_submodule(module_name)

        # swap alpha/scaling
        r_key = get_pattern_key(config.rank_pattern.keys(), key)
        alpha_key = get_pattern_key(config.alpha_pattern.keys(), key)
        rank = config.rank_pattern.get(r_key, config.r)
        alpha = config.alpha_pattern.get(alpha_key, config.lora_alpha)
        if config.use_rslora:
            scaling = alpha / math.sqrt(rank)
        else:
            scaling = alpha / rank
        _update_scaling(module, adapter_name=adapter_name, scaling=scaling)

        # swap actual weights
        # no need to account for potential _orig_mod in key here, as torch handles that
        old_val = attrgetter(key)(model)
        if not is_compiled and not is_compiled_inplace:
            torch.utils.swap_tensors(old_val, new_val)
            continue

        # Compiled models don't work with swap_tensors because there are weakrefs for the tensor. It is unclear if
        # this workaround could not cause trouble but the tests indicate that it works.
        if old_val.shape == new_val.shape:
            old_val.data = new_val.data
        else:
            if old_val.dim() != 2:
                # TODO conv2d
                raise NotImplementedError
            if old_val.shape[0] > new_val.shape[0]:
                old_val.data.fill_(0)
                old_val.data[: new_val.shape[0]] = new_val.data
            elif old_val.shape[1] > new_val.shape[1]:
                old_val.data.fill_(0)
                old_val.data[:, : new_val.shape[1]] = new_val.data
            else:
                raise ValueError(
                    f"Incompatible shapes found for LoRA weights {key}: {old_val.shape} vs {new_val.shape}. Please "
                    "ensure that all ranks are padded to the largest rank among all LoRA adapters by using "
                    "peft.utils.hotswap.prepare_model_for_compiled_hotswap."
                )


def _check_hotswap_configs_compatible(config0: PeftConfig, config1: PeftConfig) -> None:
    """
    Check if two configs are compatible for hot-swapping.

    Only LoRA parameters are checked for now.

    To hot-swap two adapters, their configs must be compatible. Otherwise, the results could be false. E.g. if they use
    different alpha values, after hot-swapping, the alphas from the first adapter would still be used with the weights
    from the 2nd adapter, which would result in incorrect behavior. There is probably a way to swap these values as
    well, but that's not implemented yet, and we need to be careful not to trigger re-compilation if the model is
    compiled (so no modification of the dict).

    """

    if config0.peft_type != config1.peft_type:
        msg = f"Incompatible PEFT types found: {config0.peft_type.value} and {config1.peft_type.value}"
        raise ValueError(msg)

    if config0.peft_type not in CONFIG_KEYS_TO_CHECK:
        msg = (
            f"Hotswapping only supports {', '.join(CONFIG_KEYS_TO_CHECK.keys())} but "
            f"{config0.peft_type.value} was passed."
        )
        raise ValueError(msg)
    config_keys_to_check = CONFIG_KEYS_TO_CHECK[config0.peft_type]

    # TODO: This is a very rough check only for LoRA at the moment. Also, there might be some options that don't
    # necessarily require an error.
    config0 = config0.to_dict()
    config1 = config1.to_dict()
    sentinel = object()
    for key in config_keys_to_check:
        val0 = config0.get(key, sentinel)
        val1 = config1.get(key, sentinel)
        if val0 != val1:
            raise ValueError(f"Configs are incompatible: for {key}, {val0} != {val1}")


def hotswap_adapter(model, model_name_or_path, adapter_name, torch_device=None, **kwargs):
    """Substitute old adapter data with new adapter data, keeping the rest the same.

    As of now, only LoRA is supported.

    This function is useful when you want to replace the loaded adapter with a new adapter. The adapter name will
    remain the same, but the weights and other parameters will be swapped out.

    If the adapters are incomptabile, e.g. targeting different layers or having different alpha values, an error will
    be raised.

    Example:

    ```py
    >>> import torch
    >>> from transformers import AutoModelForCausalLM
    >>> from peft import PeftModel
    >>> from peft.utils.hotswap import hotswap_adapter

    >>> model_id = ...
    >>> inputs = ...
    >>> device = ...
    >>> model = AutoModelForCausalLM.from_pretrained(model_id).to(device)

    >>> # load lora 0
    >>> model = PeftModel.from_pretrained(model, "path-adapter-0")
    >>> model = torch.compile(model)  # optionally compile the model
    >>> with torch.inference_mode():
    ...     output_adapter_0 = model(inputs)

    >>> # replace the "default" lora adapter with the new one
    >>> hotswap_adapter(model, "path-adapter-1", adapter_name="default", torch_device=device)
    >>> with torch.inference_mode():
    ...     output_adapter_1 = model(inputs).logits
    ```

    Args:
        model ([`~PeftModel`]):
            The PEFT model with the loaded adapter.
        model_name_or_path (`str`):
            The name or path of the model to load the new adapter from.
        adapter_name (`str`):
            The name of the adapter to swap, e.g. `"default"`. The name will stay the same after swapping.
        torch_device: (`str`, *optional*, defaults to None):
            The device to load the new adapter onto.
        **kwargs (`optional`):
            Additional keyword arguments used for loading the config and weights.

    """
    if torch_device is None:
        torch_device = infer_device()

    ############################
    # LOAD CONFIG AND VALIDATE #
    ############################

    config_cls = PEFT_TYPE_TO_CONFIG_MAPPING[
        PeftConfig._get_peft_type(
            model_name_or_path,
            subfolder=kwargs.get("subfolder", None),
            revision=kwargs.get("revision", None),
            cache_dir=kwargs.get("cache_dir", None),
            use_auth_token=kwargs.get("use_auth_token", None),
            token=kwargs.get("token", None),
        )
    ]
    config = config_cls.from_pretrained(model_name_or_path, **kwargs)
    # config keys that could affect the model output besides what is determined by the state_dict
    _check_hotswap_configs_compatible(model.active_peft_config, config)

    state_dict = load_peft_weights(model_name_or_path, device=torch_device, **kwargs)

    ###########################
    # LOAD & REMAP STATE_DICT #
    ###########################

    parameter_prefix = PEFT_TYPE_TO_PREFIX_MAPPING[config.peft_type]
    peft_model_state_dict = _insert_adapter_name_into_state_dict(
        state_dict, adapter_name=adapter_name, parameter_prefix=parameter_prefix
    )

    hotswap_adapter_from_state_dict(
        model=model,
        state_dict=peft_model_state_dict,
        adapter_name=adapter_name,
        parameter_prefix=parameter_prefix,
        config=config,
    )
