# 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
from typing import Optional

import torch
import torch.nn.functional as F
from torch import nn

from peft.tuners._buffer_dict import BufferDict
from peft.tuners.tuners_utils import BaseTunerLayer, check_adapters_to_merge

from .utils import clustering_Z, get_trainable_subspaces, seg_locations, slice_pca


class AdamssLayer(BaseTunerLayer):
    """
    Base Adamss layer that stores adapter-specific information.
    """

    # All names of layers that may contain adapter weights (trainable or frozen)
    # Use ModuleDict for proper state_dict key format that PEFT's save/load can handle
    # Only include trainable parameter containers here (not buffers)
    # adamss_newB is a buffer (frozen), not a trainable parameter
    adapter_layer_names = ("adamss_A", "adamss_B")
    # other_param_names specifies additional non-tensor metadata
    other_param_names = (
        "num_subspaces",
        "train_subspace_index",
        "seg_result",
        "scatter_index",
        "exp_avg_ipt_A",
        "exp_avg_ipt_B",
        "exp_avg_unc_A",
        "exp_avg_unc_B",
    )

    def __init__(self, base_layer: nn.Module, **kwargs) -> None:
        self.base_layer = base_layer
        # Adapter-specific attributes
        # Use ModuleDict containing ParameterList for proper state_dict keys (e.g., adamss_A.default.0)
        self.adamss_A = nn.ModuleDict({})  # Will contain ParameterList per adapter
        self.adamss_B = nn.ModuleDict({})  # Will contain ParameterList per adapter
        # Use BufferDict for frozen projection matrix (keys like adamss_newB.default)
        self.adamss_newB = BufferDict(persistent=True)
        self.num_subspaces = {}
        self.train_subspace_index = {}
        self.seg_result = {}
        self.scatter_index = {}
        # ASA importance tracking (per-adapter lists indexed by subspace)
        self.exp_avg_ipt_A = {}
        self.exp_avg_ipt_B = {}
        self.exp_avg_unc_A = {}
        self.exp_avg_unc_B = {}
        self._disable_adapters = False
        self.merged_adapters = []

    def reset_importance(self, adapter_name: str) -> None:
        """
        Clear stored importance stats for an adapter.

        Called after each masking interval to restart EMA accumulation for the next importance scoring window. Without
        the reset the scores from early training steps would dominate later masking decisions.
        """
        if adapter_name in self.exp_avg_ipt_A:
            n = len(self.exp_avg_ipt_A[adapter_name])
            self.exp_avg_ipt_A[adapter_name] = [None] * n
            self.exp_avg_ipt_B[adapter_name] = [None] * n
            self.exp_avg_unc_A[adapter_name] = [None] * n
            self.exp_avg_unc_B[adapter_name] = [None] * n

    def update_importance(
        self, adapter_name: str, importance_beta: float = 0.85, uncertainty_beta: float = 0.85
    ) -> None:
        """
        Update importance scores using current gradients.

        Called by [`AdamssModel.update_and_allocate`] (which is in turn called by [`AdamssAsaCallback`] when using
        HuggingFace `Trainer`).

        Args:
            adapter_name: Name of the adapter to update importance for.
            importance_beta: EMA coefficient for importance averaging (0.8-0.95 typical).
            uncertainty_beta: EMA coefficient for uncertainty averaging (0.8-0.95 typical).
        """
        if adapter_name not in self.exp_avg_ipt_A:
            return
        if adapter_name not in self.adamss_A:
            return

        param_list_A = self.adamss_A[adapter_name]
        param_list_B = self.adamss_B[adapter_name]

        ipt_A = self.exp_avg_ipt_A[adapter_name]
        ipt_B = self.exp_avg_ipt_B[adapter_name]
        unc_A = self.exp_avg_unc_A[adapter_name]
        unc_B = self.exp_avg_unc_B[adapter_name]

        for i in range(self.num_subspaces[adapter_name]):
            for param_list, ipt_list, unc_list in [
                (param_list_A, ipt_A, unc_A),
                (param_list_B, ipt_B, unc_B),
            ]:
                param = param_list[i]
                if param.grad is not None:
                    if ipt_list[i] is None:
                        ipt_list[i] = torch.zeros_like(param, requires_grad=False)
                        unc_list[i] = torch.zeros_like(param, requires_grad=False)

                    # Calculate importance: |w * g|
                    ipt = (param * param.grad).abs().detach()

                    # Update uncertainty BEFORE updating importance (matches adamss_pkg)
                    diff = (ipt - ipt_list[i]).abs()
                    unc_list[i].mul_(uncertainty_beta).add_(diff, alpha=1 - uncertainty_beta)

                    # Then update importance
                    ipt_list[i].mul_(importance_beta).add_(ipt, alpha=1 - importance_beta)

    def update_layer(
        self,
        adapter_name: str,
        r: int,
        num_subspaces: int,
        subspace_rank: int,
        init_weights: str,
        use_asa: bool = False,
        inference_mode: bool = False,
        use_dynamic_rank: bool = False,
        svd_threshold: float = 0.1,
        **kwargs,
    ) -> None:
        """
        Update layer with Adamss adapter.

        This method initializes the Adamss decomposition for the weight matrix using SVD, clustering, and QR
        initialization.
        """

        # Get the base weight info
        weight = self.get_base_layer().weight
        bias = self.get_base_layer().bias
        device = weight.device
        dtype = weight.dtype
        _, in_features = weight.shape

        # Adjust in_features for bias column
        if bias is not None:
            in_features += 1

        # Prepare weight tensor (add bias as extra column if present)
        if bias is not None:
            weight_with_bias = torch.cat((weight, bias.unsqueeze(1)), dim=1)
        else:
            weight_with_bias = weight

        # Reshape to 4D tensor for slice_pca: (1, 1, out_features, in_features)
        weight_tensor = weight_with_bias.unsqueeze(0).unsqueeze(0).to(torch.float32).to(device)

        # Perform SVD decomposition with diagnostics in case of failure
        try:
            res = slice_pca(weight_tensor, r, device, torch.float32)
        except Exception as e:
            raise RuntimeError(
                f"slice_pca raised an exception for layer {adapter_name} (shape={tuple(weight_tensor.shape)}, dtype={weight_tensor.dtype}, device={device}): {e}"
            ) from e

        VVT, UU = res

        # Cluster the column space
        cluster_idx, effective_num_subspaces = clustering_Z(UU[0, 0, :, :], num_subspaces, 10)

        # Get segment locations
        seg_result = seg_locations(cluster_idx)

        # Get trainable subspace indices
        train_subspace_index = get_trainable_subspaces(effective_num_subspaces)

        # Store metadata
        self.num_subspaces[adapter_name] = effective_num_subspaces
        self.train_subspace_index[adapter_name] = train_subspace_index
        self.seg_result[adapter_name] = seg_result

        # Store projection matrix in BufferDict (frozen, device-aware)
        self.adamss_newB[adapter_name] = VVT[0, 0, :, :].detach().to(dtype)

        # Use user-specified subspace_rank
        rank_per_subspace = subspace_rank

        # Initialize trainable subspace parameters
        # Collect parameters in lists, then create ParameterList for proper state_dict keys
        A_params = []
        B_params = []

        for i in range(self.num_subspaces[adapter_name]):
            subspace_idx = train_subspace_index[i]
            seg_indices = seg_result[subspace_idx]

            # Extract subspace data: use ALL r columns from SVD (not just r/K)
            # This matches adamss_pkg where A matrices have shape (rank_i, r)
            subspace_data = UU[0, 0, seg_indices, :r]
            # subspace_data shape: (len(seg_indices), r)

            # Determine actual rank based on configuration
            if use_dynamic_rank:
                # Compute Gram matrix Z = subspace_data @ subspace_data.T
                Z_row = subspace_data @ subspace_data.T  # (len(seg_indices), len(seg_indices))

                # Compute SVD to get row space decomposition
                _U_row, S_row, _V_row = torch.svd_lowrank(Z_row, q=min(Z_row.shape), niter=2)

                # Dynamic rank selection using SVD threshold
                # Dynamically determine actual rank using configurable threshold
                # Following adamss_pkg: num_ii_jj = min((S > threshold * S[0]).sum().item(), args.adamss_ri)
                threshold_mask = S_row > svd_threshold * S_row[0]
                actual_rank = min(threshold_mask.sum().item(), subspace_rank, len(S_row))

                # Handle edge case: ensure at least rank 1
                if actual_rank == 0:
                    actual_rank = 1

            else:
                # Fixed rank: match adamss_pkg behavior exactly
                # adamss_pkg uses: num_ii_jj = min(len(seg_result[indx]), args.adamss_ri)
                actual_rank = min(len(seg_indices), subspace_rank)

                # Ensure at least rank 1
                if actual_rank == 0:
                    actual_rank = 1

            # Initialize A using QR decomposition to match adamss_pkg exactly
            # adamss_pkg: Q, R = torch.linalg.qr((UU[seg_result[indx],:]).T @ A[i], mode='reduced')
            # where A[i] is U from SVD of subspace Gram matrix

            # First, compute Gram matrix for this subspace
            Z_subspace = subspace_data @ subspace_data.T  # (len(seg_indices), len(seg_indices))

            # SVD of Gram matrix to get A_intermediate
            U_gram, _S_gram, _ = torch.svd_lowrank(Z_subspace, q=min(Z_subspace.shape[0], actual_rank), niter=2)
            A_intermediate = U_gram[:, :actual_rank]  # (len(seg_indices), actual_rank)

            # Now apply QR decomposition: (UU[seg_indices, :r]).T @ A_intermediate
            # UU[0, 0, seg_indices, :r] is subspace_data with shape (len(seg_indices), r)
            matrix_for_qr = subspace_data.T @ A_intermediate  # (r, actual_rank)
            Q, _R = torch.linalg.qr(matrix_for_qr, mode="reduced")

            # A_init is Q.T to match adamss_pkg
            A_init = Q.T.contiguous()  # (actual_rank, r)

            # Initialize B matrix
            # When init_weights='orthogonal', use zeros for identity operation (standard training init)
            # When init_weights=None/False, use random values so adapter produces non-zero effect
            if init_weights == "orthogonal":
                # Zero initialization - adapter produces identity at start
                # Note: A parameters may not update in step 1 since ∂Loss/∂A ∝ B and B=0
                # But after step 1, B becomes non-zero and A will get gradients in step 2+
                B_init = torch.zeros(len(seg_indices), actual_rank, dtype=dtype, device=device)
            else:
                # Random initialization for testing - produces non-zero adapter output
                B_init = torch.randn(len(seg_indices), actual_rank, dtype=dtype, device=device) * 0.01

            # Collect A and B parameters
            A_params.append(nn.Parameter(A_init.to(dtype)))
            B_params.append(nn.Parameter(B_init.to(dtype)))

        # Create ParameterLists and store in ModuleDict
        # This creates state_dict keys like 'adamss_A.default.0', 'adamss_A.default.1', etc.
        self.adamss_A[adapter_name] = nn.ParameterList(A_params)
        self.adamss_B[adapter_name] = nn.ParameterList(B_params)

        # Initialize ASA tracking if enabled
        if use_asa:
            n = effective_num_subspaces
            self.exp_avg_ipt_A[adapter_name] = [None] * n
            self.exp_avg_ipt_B[adapter_name] = [None] * n
            self.exp_avg_unc_A[adapter_name] = [None] * n
            self.exp_avg_unc_B[adapter_name] = [None] * n

            # Hooks are not used; importance is updated explicitly via update_importance
            # to match adamss_pkg behavior (using accumulated gradients).

        # Compute scatter_index for forward pass
        self.scatter_index[adapter_name] = torch.cat(
            [
                self.seg_result[adapter_name][self.train_subspace_index[adapter_name][i]]
                for i in range(self.num_subspaces[adapter_name])
            ],
            dim=0,
        ).long()

        # Move adapter to device of base layer (important for low_cpu_mem_usage support)
        self._move_adapter_to_device_of_base_layer(adapter_name)

        # Set requires_grad via parent class set_adapter (works with ParameterList)
        self.set_adapter(self.active_adapters, inference_mode=inference_mode)


class Linear(nn.Module, AdamssLayer):
    """
    Adamss-adapted Linear layer.
    """

    def __init__(
        self,
        base_layer: nn.Module,
        adapter_name: str,
        r: int = 500,
        num_subspaces: int = 5,
        subspace_rank: int = 1,
        init_weights: str = "orthogonal",
        use_asa: bool = False,
        **kwargs,
    ) -> None:
        super().__init__()
        AdamssLayer.__init__(self, base_layer, **kwargs)

        # Set in_features and out_features from base layer (required for PEFT compatibility)
        self.in_features = base_layer.in_features
        self.out_features = base_layer.out_features

        # Initialize the adapter
        inference_mode = kwargs.pop("inference_mode", False)
        self.update_layer(
            adapter_name,
            r,
            num_subspaces,
            subspace_rank,
            init_weights,
            use_asa,
            inference_mode=inference_mode,
            **kwargs,
        )
        self._active_adapter = adapter_name
        self.dtype = base_layer.weight.dtype

    def _load_from_state_dict(
        self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
    ):
        """Handle shape mismatches that arise when loading checkpoints.

        AdaMSS B parameter shapes depend on KMeans clustering of the weight matrix. KMeans uses a fixed `random_state`
        so results should be deterministic, but this override acts as a safety net in case sklearn changes clustering
        behaviour across versions. It detects shape mismatches and replaces placeholder parameters with
        correctly-shaped tensors before the default `load_state_dict` logic runs.
        """
        for key, value in state_dict.items():
            if not key.startswith(prefix):
                continue
            local_key = key[len(prefix) :]
            parts = local_key.split(".")
            try:
                # Use get_submodule for the parent path
                parent_path = ".".join(parts[:-1])
                target = self.get_submodule(parent_path) if parent_path else self
                last = parts[-1]
                if last.isdigit():
                    current = target[int(last)]
                else:
                    current = getattr(target, last)
            except (AttributeError, IndexError, KeyError):
                continue

            if isinstance(current, (torch.Tensor, nn.Parameter)) and current.shape != value.shape:
                new_param = nn.Parameter(torch.empty_like(value), requires_grad=current.requires_grad)
                if last.isdigit():
                    target[int(last)] = new_param
                else:
                    setattr(target, last, new_param)

        super()._load_from_state_dict(
            state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
        )

    def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
        """
        Forward pass with Adamss adaptation.
        """
        previous_dtype = x.dtype

        if self._disable_adapters:
            if self.merged:
                self.unmerge()
            result = self.base_layer(x, *args, **kwargs)
        elif self.merged:
            # Adapter weights are already merged into base_layer.weight,
            # so just call base_layer directly (no adapter delta needed).
            result = self.base_layer(x, *args, **kwargs)
        else:
            # Cast input to layer dtype
            x = x.to(self.dtype)

            # Add bias column only when the base layer has a bias
            # (update_layer concatenates bias into newB only when bias is not None)
            if self.get_base_layer().bias is not None:
                if x.dim() == 2:
                    ones = torch.ones(x.shape[0], 1, device=x.device, dtype=self.dtype)
                else:  # x.dim() == 3 or more
                    ones = torch.ones(*x.shape[:-1], 1, device=x.device, dtype=self.dtype)
                newx = torch.cat((x, ones), dim=-1)
            else:
                newx = x

            # Base output from the frozen original weight
            result = self.base_layer(x, *args, **kwargs)

            # Iterate over all active adapters and add their deltas
            for active_adapter in self.active_adapters:
                if active_adapter not in self.adamss_A:
                    continue

                # Project input into low-rank space
                projected = F.linear(newx, self.adamss_newB[active_adapter].to(self.dtype))  # (..., r)

                # Apply A and B transformations per subspace
                # TODO: if all subspaces have equal output dimensions, this loop could
                # be replaced with a batched matmul/einsum for better performance.
                # Currently B_i shapes vary per subspace (clustering-dependent), so we loop.
                subspace_chunks = []
                for i in range(self.num_subspaces[active_adapter]):
                    A_i = self.adamss_A[active_adapter][i].to(self.dtype)  # (ri, r)
                    B_i = self.adamss_B[active_adapter][i].to(self.dtype)  # (seg_len_i, ri)
                    subspace_out = F.linear(projected, A_i)  # (..., ri)
                    subspace_out = F.linear(subspace_out, B_i)  # (..., seg_len_i)
                    subspace_chunks.append(subspace_out)

                # Concatenate subspace results and scatter to output dimension order
                adapter_delta = torch.cat(subspace_chunks, dim=-1)
                scatter_index_tensor = self.scatter_index[active_adapter].to(adapter_delta.device)

                if adapter_delta.dim() == 2:
                    index = scatter_index_tensor.unsqueeze(0).expand(adapter_delta.shape[0], -1)
                    adapter_delta = torch.zeros(
                        adapter_delta.shape[0],
                        result.shape[-1],
                        device=adapter_delta.device,
                        dtype=adapter_delta.dtype,
                    ).scatter(1, index, adapter_delta)
                else:
                    index = scatter_index_tensor.unsqueeze(0).unsqueeze(0).expand(*adapter_delta.shape[:-1], -1)
                    adapter_delta = torch.zeros(
                        *adapter_delta.shape[:-1],
                        result.shape[-1],
                        device=adapter_delta.device,
                        dtype=adapter_delta.dtype,
                    ).scatter(-1, index, adapter_delta)

                result = result + adapter_delta

        # Cast back to original dtype
        return result.to(previous_dtype)

    def merge(self, safe_merge: bool = False, adapter_names: Optional[list[str]] = None) -> None:
        """
        Merge the active adapter weights into the base weights.

        Args:
            safe_merge (`bool`, *optional*):
                If True, the merge operation will be performed in a copy of the original weights and check for NaNs
                before merging the weights. Defaults to `False`.
            adapter_names (`list[str]`, *optional*):
                The list of adapter names that should be merged. If None, all active adapters will be merged.
        """
        adapter_names = check_adapters_to_merge(self, adapter_names)
        if not adapter_names:
            # no adapter to merge
            return

        for active_adapter in adapter_names:
            if active_adapter in self.adamss_A:
                base_layer = self.get_base_layer()
                if safe_merge:
                    # Note that safe_merge will be slower than the normal merge
                    # because of the copy operation.
                    orig_weight = base_layer.weight.data.clone()
                    orig_dtype = orig_weight.dtype
                    delta_weight = self.get_delta_weight(active_adapter)
                    orig_weight += delta_weight.to(orig_dtype)

                    if not torch.isfinite(orig_weight).all():
                        raise ValueError(
                            f"NaNs detected in the merged weights. The adapter {active_adapter} seems to be broken"
                        )

                    base_layer.weight.data = orig_weight

                    # Also update bias if present
                    if base_layer.bias is not None:
                        orig_bias = base_layer.bias.data.clone()
                        delta_bias = self.get_delta_bias(active_adapter)
                        orig_bias += delta_bias.to(orig_dtype)
                        if not torch.isfinite(orig_bias).all():
                            raise ValueError(
                                f"NaNs detected in the merged bias. The adapter {active_adapter} seems to be broken"
                            )
                        base_layer.bias.data = orig_bias
                else:
                    delta_weight = self.get_delta_weight(active_adapter)
                    base_layer.weight.data += delta_weight

                    # Also update bias if present
                    if base_layer.bias is not None:
                        delta_bias = self.get_delta_bias(active_adapter)
                        base_layer.bias.data += delta_bias

                self.merged_adapters.append(active_adapter)

    def unmerge(self) -> None:
        """
        This method unmerges all merged adapter layers from the base weights.
        """
        if not self.merged:
            warnings.warn("Already unmerged. Nothing to do.")
            return

        while len(self.merged_adapters) > 0:
            active_adapter = self.merged_adapters.pop()
            if active_adapter in self.adamss_A:
                base_layer = self.get_base_layer()
                weight = base_layer.weight
                orig_dtype = weight.dtype
                delta_weight = self.get_delta_weight(active_adapter)
                weight.data -= delta_weight.to(orig_dtype)

                # Also update bias if present
                if base_layer.bias is not None:
                    delta_bias = self.get_delta_bias(active_adapter)
                    base_layer.bias.data -= delta_bias.to(orig_dtype)

    def get_delta_weight(self, adapter_name: str) -> torch.Tensor:
        """
        Compute the delta weight for the given adapter.

        For AdaMSS, the forward computes:
            result = newx @ resW.T + scatter(newx @ newB.T @ A.T @ B.T)

        Since resW = original_weight_with_bias, the delta is just the adapter path:
            delta = scatter(B @ A @ newB)

        We extract the weight portion (excluding bias column) as the delta to add to the base layer's weight.

        Args:
            adapter_name (str): The name of the adapter for which the delta weight should be computed.
        """
        device = self.get_base_layer().weight.device
        dtype = self.get_base_layer().weight.dtype
        base_weight = self.get_base_layer().weight

        # Get buffers
        newB = self.adamss_newB[adapter_name]  # Shape: (r, in_features + 1)

        # Get parameter lists
        param_list_A = self.adamss_A[adapter_name]
        param_list_B = self.adamss_B[adapter_name]

        # In case users wants to merge the adapter weights that are in
        # (b)float16 while being on CPU, we need to cast the weights to float32
        cast_to_fp32 = device.type == "cpu" and dtype in (torch.float16, torch.bfloat16)
        compute_dtype = torch.float32 if cast_to_fp32 else dtype

        newB = newB.to(device).to(compute_dtype)

        # Compute the adapter contribution (scattered B @ A @ newB)
        # delta_weight has shape (out_features, in_features + 1)
        out_features = base_weight.shape[0]
        in_features_plus_1 = newB.shape[1]

        # Initialize delta weight for the adapter path only
        delta_weight = torch.zeros(out_features, in_features_plus_1, device=device, dtype=compute_dtype)

        # Compute the transformation for each subspace
        # For subspace i: contribution = B_i @ A_i @ newB
        # where A_i has shape (rank_i, r), B_i has shape (seg_len_i, rank_i)
        chunks = []
        for i in range(self.num_subspaces[adapter_name]):
            A_i = param_list_A[i].to(device).to(compute_dtype)  # Shape: (rank_i, r)
            B_i = param_list_B[i].to(device).to(compute_dtype)  # Shape: (seg_len_i, rank_i)

            # B_i @ A_i @ newB gives shape (seg_len_i, in_features+1)
            chunk = B_i @ A_i @ newB
            chunks.append(chunk)

        # Concatenate chunks and scatter to correct positions
        if chunks:
            combined = torch.cat(chunks, dim=0)  # Shape: (total_seg_len, in_features+1)
            # Use scatter_index for assignment
            scatter_idx = self.scatter_index[adapter_name].to(delta_weight.device)
            delta_weight[scatter_idx] = combined

        # Extract just the weight portion
        if self.get_base_layer().bias is not None:
            # delta_weight[:, :-1] is the weight delta (exclude bias column)
            output_tensor = delta_weight[:, :-1]  # Shape: (out_features, in_features)
        else:
            # No bias column to strip
            output_tensor = delta_weight  # Shape: (out_features, in_features)

        if cast_to_fp32:
            output_tensor = output_tensor.to(dtype=dtype)

        return output_tensor

    def get_delta_bias(self, adapter_name: str) -> torch.Tensor:
        """
        Compute the bias delta for the given adapter.

        This is the last column of the adapter path contribution (B @ A @ newB).

        Args:
            adapter_name (str): The name of the adapter for which the bias delta should be computed.
        """
        device = self.get_base_layer().weight.device
        dtype = self.get_base_layer().weight.dtype
        base_weight = self.get_base_layer().weight

        # Get buffers
        newB = self.adamss_newB[adapter_name]  # Shape: (r, in_features + 1)

        # Get parameter lists
        param_list_A = self.adamss_A[adapter_name]
        param_list_B = self.adamss_B[adapter_name]

        cast_to_fp32 = device.type == "cpu" and dtype in (torch.float16, torch.bfloat16)
        compute_dtype = torch.float32 if cast_to_fp32 else dtype

        newB = newB.to(device).to(compute_dtype)

        out_features = base_weight.shape[0]
        in_features_plus_1 = newB.shape[1]

        # Initialize delta weight for the adapter path only
        delta_weight = torch.zeros(out_features, in_features_plus_1, device=device, dtype=compute_dtype)

        # Compute the transformation for each subspace
        chunks = []
        for i in range(self.num_subspaces[adapter_name]):
            A_i = param_list_A[i].to(device).to(compute_dtype)
            B_i = param_list_B[i].to(device).to(compute_dtype)
            chunk = B_i @ A_i @ newB
            chunks.append(chunk)

        if chunks:
            combined = torch.cat(chunks, dim=0)
            scatter_idx = self.scatter_index[adapter_name].to(delta_weight.device)
            delta_weight[scatter_idx] = combined

        # Extract the bias portion (last column)
        if self.get_base_layer().bias is None:
            # No bias — return zeros
            output_tensor = torch.zeros(base_weight.shape[0], device=device, dtype=compute_dtype)
        else:
            output_tensor = delta_weight[:, -1]  # Shape: (out_features,)

        if cast_to_fp32:
            output_tensor = output_tensor.to(dtype=dtype)

        return output_tensor

    def __repr__(self) -> str:
        rep = super().__repr__()
        return "adamss." + rep
