# Copyright 2025-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
from dataclasses import dataclass, field
from typing import Optional, Union

from peft.config import PeftConfig
from peft.utils import PeftType


@dataclass
class PveraConfig(PeftConfig):
    """
    This is the configuration class to store the configuration of a [`PveraModel`].

    Paper: https://www.arxiv.org/abs/2512.07703.

    Args:
        r (`int`, *optional*, defaults to `256`):
            PVeRA parameter dimension ("rank"). Choose higher values than LoRA ranks here, since PVeRA shares
            parameters across layers and therefore uses far fewer parameters than LoRA.
        target_modules (`Union[List[str], str]`):
            The names of the modules to apply PVeRA to. Only linear layers are supported. When passing a string, a
            regex match will be performed. If this is specified as 'all-linear', then all linear/Conv1D modules are
            chosen. If this is not specified, modules will bechosen according to the model architecture. If the
            architecture is not known, an error will be raised.
        projection_prng_key (`int`):
            PVeRA PRNG init key. Used for initialising pvera_A and pvera_B for new models or when loading a checkpoint
            that did not include these projections. Defaults to `0`.
        save_projection (`bool`):
            Whether to save the pvera_A / pvera_B projections in the state dict alongside per layer lambda_b / lambda_d
            weights. This will increase the size of the checkpoint, but guarantee that we can reload the checkpoint on
            all system configurations. Defaults to `True`.
        pvera_dropout (`float`):
            The dropout probability for PVeRA layers.
        d_initial (`float`, *optional*, defaults to `0.1`):
            Initial value for `pvera_lambda_d` vector used when initializing the PVeRA parameters. Small values (<=0.1)
            are recommended.
        fan_in_fan_out (`bool`):
            Set this to True if the layer to replace stores weight like (fan_in, fan_out). For example, gpt-2 uses
            `Conv1D` which stores weights like (fan_in, fan_out) and hence this should be set to `True`.
        bias (`str`):
            Bias type for PVeRA. Can be 'none', 'all' or 'pvera_only'. If 'all' or 'pvera_only', the corresponding
            biases will be updated during training. Be aware that this means that, even when disabling the adapters,
            the model will not produce the same output as the base model would have without adaptation.
        modules_to_save (`List[str]`):
            List of modules apart from PVeRA layers to be set as trainable and saved in the final checkpoint.
        init_weights (`bool`):
            Whether to initialize the weights of the PVeRA layers with their default initialization. Don't change this
            setting, except if you know exactly what you're doing.
        layers_to_transform (`Union[List[int],int]`):
            The layer indexes to transform, if this argument is specified, it will apply the PVeRA transformations on
            the layer indexes that are specified in this list. If a single integer is passed, it will apply the PVeRA
            transformations on the layer at this index.
        layers_pattern (`Optional[Union[List[str], str]]`):
            The layer pattern name, used only if `layers_to_transform` is different from `None`. This should target the
            `nn.ModuleList` of the model, which is often called `'layers'` or `'h'`.
        sample_at_inference (`bool` | `dict`, defaults to `False`):
            Whether to sample from the learned PVeRA distribution at inference. If false, the learned mean is used. The
            default is False (indicating false for all adapters). If True is provided, then the value will be true for
            all adapters. If a dict is provided, then a specific value can be specified per adapter (with False by
            default for non-specified adapters). For example
            `sample_at_inference={'encoder.layer.0.attention.attention.query': True}` will only sample at inference for
            one specific adapter.
    """

    r: int = field(
        default=256,
        metadata={
            "help": (
                "PVeRA parameter dimension ('rank'). Choose higher values than LoRA ranks here, since PVeRA shares "
                "parameters across layers and therefore uses far fewer parameters than LoRA."
            )
        },
    )

    target_modules: Optional[Union[list[str], str]] = field(
        default=None,
        metadata={
            "help": (
                "The names of the modules to apply PVeRA to. Only linear layers are supported. When passing a string, a "
                "regex match will be performed. If this is specified as 'all-linear', then all linear/Conv1D modules are "
                "chosen. If this is not specified, modules will bechosen according to the model architecture. If the "
                "architecture is not known, an error will be raised."
            )
        },
    )
    projection_prng_key: int = field(
        default=0,
        metadata={
            "help": (
                "PVeRA PRNG init key. Used for initialising pvera_A and pvera_B for new models or when loading a checkpoint "
                "that did not include these projections. Defaults to `0`."
            )
        },
    )
    save_projection: bool = field(
        default=True,
        metadata={
            "help": (
                "Whether to save the pvera_A / pvera_B projections in the state dict alongside per layer lambda_b / lambda_d "
                "weights. This will increase the size of the checkpoint, but guarantee that we can reload the checkpoint on "
                "all system configurations. Defaults to `True`."
            )
        },
    )
    pvera_dropout: float = field(default=0.0, metadata={"help": "The dropout probability for PVeRA layers."})
    d_initial: float = field(default=0.1, metadata={"help": "Initial value for d vector. Default is 0.1."})
    fan_in_fan_out: bool = field(
        default=False,
        metadata={
            "help": (
                "Set this to True if the layer to replace stores weight like (fan_in, fan_out). For example, gpt-2 uses "
                "`Conv1D` which stores weights like (fan_in, fan_out) and hence this should be set to `True`."
            )
        },
    )
    bias: str = field(
        default="none",
        metadata={
            "help": (
                "Bias type for PVeRA. Can be 'none', 'all' or 'pvera_only'. If 'all' or 'pvera_only', the corresponding "
                "biases will be updated during training. Be aware that this means that, even when disabling the adapters, "
                "the model will not produce the same output as the base model would have without adaptation."
            )
        },
    )
    modules_to_save: Optional[list[str]] = field(
        default=None,
        metadata={
            "help": (
                "List of modules apart from PVeRA layers to be set as trainable and saved in the final checkpoint."
            )
        },
    )
    init_weights: bool = field(
        default=True,
        metadata={
            "help": (
                "Whether to initialize the weights of the PVeRA layers with their default initialization. Don't change this "
                "setting, except if you know exactly what you're doing."
            ),
        },
    )
    layers_to_transform: Optional[Union[list[int], int]] = field(
        default=None,
        metadata={
            "help": (
                "The layer indexes to transform, if this argument is specified, it will apply the PVeRA transformations on "
                "the layer indexes that are specified in this list. If a single integer is passed, it will apply the PVeRA "
                "transformations on the layer at this index."
            )
        },
    )
    layers_pattern: Optional[Union[list[str], str]] = field(
        default=None,
        metadata={
            "help": (
                "The layer pattern name, used only if `layers_to_transform` is different from `None`. This should target the "
                "`nn.ModuleList` of the model, which is often called `'layers'` or `'h'`."
            )
        },
    )
    sample_at_inference: bool = field(
        default=False,
        metadata={
            "help": (
                "Whether to sample from the learned PVeRA distribution at inference. If false, the learned mean is used. The "
                "default is False (indicating false for all adapters). If True is provided, then the value will be true for "
                "all adapters. If a dict is provided, then a specific value can be specified per adapter (with False by "
                "default for non-specified adapters). For example "
                "`sample_at_inference={'encoder.layer.0.attention.attention.query': True}` will only sample at inference for "
                "one specific adapter."
            ),
        },
    )

    def __post_init__(self):
        super().__post_init__()
        self.peft_type = PeftType.PVERA
        self.target_modules = (
            set(self.target_modules) if isinstance(self.target_modules, list) else self.target_modules
        )
        # check for layers_to_transform and layers_pattern
        if self.layers_pattern and not self.layers_to_transform:
            raise ValueError("When `layers_pattern` is specified, `layers_to_transform` must also be specified. ")
        if not self.save_projection:
            warnings.warn(
                "Specified to not save pvera_A and pvera_B within the state dictionary, instead they will be restored "
                "using the PRNG key store in `config.projection_prng_key`. Consider setting `config.save_projection` "
                "to `True` to guarantee restoring the checkpoint correctly on all system configurations."
            )