# Copyright 2023-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
import warnings
from typing import Any, Optional, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import svd_lowrank
from transformers.pytorch_utils import Conv1D

from peft.tuners.tuners_utils import BaseTunerLayer, check_adapters_to_merge
from peft.utils.integrations import dequantize_module_weight, gather_params_ctx
from peft.utils.other import transpose

from .config import GIFTConfig


class GIFTLayer(BaseTunerLayer):
    # All names of layers that may contain (trainable) adapter weights
    adapter_layer_names = ("gift_phi", "gift_psi")
    # All names of other parameters that may contain adapter-related parameters
    other_param_names = ("r", "scaling")

    def __init__(self, base_layer: nn.Module, transform_dim: Optional[str]="input", **kwargs) -> None:
        self.base_layer = base_layer
        self.transform_dim = transform_dim
        self.r = {}
        self.scaling = {}
        self.gift_phi = nn.ModuleDict({})
        self.gift_psi = nn.ModuleDict({})
        # Mark the weight as unmerged
        self._disable_adapters = False
        self.merged_adapters = []
        self._caches: dict[str, Any] = {}
        self.kwargs = kwargs

        base_layer = self.get_base_layer()
        if isinstance(base_layer, nn.Linear):
            in_features, out_features = base_layer.in_features, base_layer.out_features
        elif isinstance(base_layer, nn.Conv2d):
            # in_features, out_features = base_layer.in_channels, base_layer.out_channels
            raise ValueError("Conv2D is not supported in GIFT yet")
        elif isinstance(base_layer, nn.Embedding):
            # in_features, out_features = base_layer.num_embeddings, base_layer.embedding_dim
            raise ValueError("Embedding is not supported in GIFT yet")
        elif isinstance(base_layer, Conv1D):
            # in_features, out_features = (
            #     base_layer.weight.ds_shape if hasattr(base_layer.weight, "ds_shape") else base_layer.weight.shape
            # )
            raise ValueError("Conv1D is not supported in GIFT yet")
        else:
            raise ValueError(f"Unsupported layer type {type(base_layer)}")

        self.in_features = in_features
        self.out_features = out_features

    def update_layer(
        self, adapter_name, r, gift_alpha, init_gift_weights, own_parameters=False,
    ):
        # This code works for linear layers, override for other layer types
        if r <= 0:
            raise ValueError(f"`r` should be a positive integer value but the value passed is {r}")

        self.r[adapter_name] = r
        self.scaling[adapter_name] = gift_alpha / r

        # Only create the weights if they are not shared, i.e. are unique to this layer
        # If the layer is shared, the weights will be provided by the hook
        if own_parameters:
            # Actual trainable parameters
            # Note that we use a TransposedLinear layer since this implementation uses
            # the representation finetuning perspective of GIFT
            dim = self.in_features if self.transform_dim == "input" else self.out_features
            self.gift_phi[adapter_name] = nn.Linear(r, dim, bias=False)
            self.gift_psi[adapter_name] = nn.Linear(dim, r, bias=False)

            if init_gift_weights:
                self.reset_gift_parameters(adapter_name, init_gift_weights)

        # check weight and qweight (for GPTQ)
        for weight_name in ("weight", "qweight"):
            weight = getattr(self.get_base_layer(), weight_name, None)
            if weight is not None:
                # the layer is already completely initialized, this is an update
                if weight.dtype.is_floating_point or weight.dtype.is_complex:
                    self.to(weight.device, dtype=weight.dtype)
                else:
                    self.to(weight.device)
                break

        self.set_adapter(self.active_adapters)

    def reset_gift_parameters(self, adapter_name, init_gift_weights):
        if init_gift_weights is False:
            return

        if adapter_name in self.gift_phi.keys():
            if init_gift_weights is True:
                # initialize A the same way as the default for nn.Linear and psi to zero
                # https://github.com/microsoft/LoRA/blob/a0a92e0f26c067cf94747bdbf1ce73793fa44d19/loralib/layers.py#L124
                nn.init.kaiming_uniform_(self.gift_psi[adapter_name].weight, a=math.sqrt(5))
                nn.init.zeros_(self.gift_phi[adapter_name].weight)
            elif init_gift_weights.lower() == "gaussian":
                nn.init.normal_(self.gift_psi[adapter_name].weight, std=1 / self.r[adapter_name])
                nn.init.zeros_(self.gift_phi[adapter_name].weight)
            elif init_gift_weights.lower() == "orthogonal":
                nn.init.orthogonal_(self.gift_phi[adapter_name].weight)
                # init psi as a transpose of phi. This will make the rotation
                # start from 0
                with torch.no_grad():
                    self.gift_psi[adapter_name].weight.data.copy_(self.gift_phi[adapter_name].weight.data.t())

            else:
                raise ValueError(f"Unknown initialization {init_gift_weights=}")
            

    def _cache_store(self, key: str, value: Any) -> None:
        self._caches[key] = value

    def _cache_pop(self, key: str) -> Any:
        value = self._caches.pop(key)
        return value

    def _check_forward_args(self, x, *args, **kwargs):
        """Check if the arguments are compatible with the configs and state of the model"""
        adapter_names = kwargs.get("adapter_names", None)
        if adapter_names is None:
            return

        if len(x) != len(adapter_names):
            msg = (
                "Length of `adapter_names` should be the same as the number of inputs, but got "
                f"{len(adapter_names)} and {len(x)} respectively."
            )
            raise ValueError(msg)

        if self.merged:
            # It is unclear what would be the right thing to do if users pass adapter_names and there are merged
            # adapters. Therefore, it is better to raise an error in this case.
            msg = "Cannot pass `adapter_names` when there are merged adapters, please call `unmerge_adapter` first."
            raise ValueError(msg)

    def _mixed_batch_forward(
        self, x: torch.Tensor, *args: Any, adapter_names: list[str], **kwargs: Any
    ) -> torch.Tensor:
        # This is a special method that handles the case when users pass the argument `adapter_names`. This is an
        # extra argument that allows mixing different adapters in the same batch at inference time.
        raise NotImplementedError("Mixed batch forward is not implemented for GIFT yet")
        torch_dtype = x.dtype
        unique_adapters = set(adapter_names)
        sub_batch_indices_list = []
        for adapter in unique_adapters:
            sub_batch_indices_list.append([index for index, item in enumerate(adapter_names) if item == adapter])

        for i, active_adapter in enumerate(unique_adapters):
            if active_adapter == "__base__":
                continue
            if active_adapter not in self.gift_phi.keys():
                continue

            gift_phi = self.gift_phi[active_adapter]
            gift_psi = self.gift_psi[active_adapter]

            # getting the sub-batch, passing it to LoRA layers and updating the corresponding indices of the linear
            # layer output
            sub_batch = x[sub_batch_indices_list[i]].to(gift_phi.weight.dtype)
            lora_output = gift_phi(gift_psi(sub_batch))
            result[sub_batch_indices_list[i]] += lora_output.to(torch_result_dtype)

        result = self.base_layer(x, *args, **kwargs)
        torch_result_dtype = result.dtype

        return result


# Below code is based on https://github.com/microsoft/LoRA/blob/main/loralib/layers.py
# and modified to work with PyTorch FSDP


#  ------------------------------------------------------------------------------------------
#  Copyright (c) Microsoft Corporation. All rights reserved.
#  Licensed under the MIT License (MIT). See LICENSE in the repo root for license information.
#  ------------------------------------------------------------------------------------------


class Linear(nn.Module, GIFTLayer):
    # GIFT implemented in a dense layer
    def __init__(
        self,
        base_layer,
        adapter_name: str,
        r: int = 0,
        gift_alpha: float = 1.0,
        fan_in_fan_out: bool = False,  # Set this to True if the layer to replace stores weight like (fan_in, fan_out)
        is_target_conv_1d_layer: bool = False,
        init_gift_weights: Union[bool, str] = True,
        own_parameters: bool = False,
        transform_dim: Optional[int] = "input",
        **kwargs,
    ) -> None:
        super().__init__()
        GIFTLayer.__init__(self, base_layer, transform_dim=transform_dim, **kwargs)
        self.fan_in_fan_out = fan_in_fan_out

        self._active_adapter = adapter_name
        self.update_layer(
            adapter_name,
            r,
            gift_alpha,
            init_gift_weights=init_gift_weights,
            own_parameters=own_parameters,
        )
        self.is_target_conv_1d_layer = is_target_conv_1d_layer

    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. This is useful if you want to check if the merge operation will produce
                NaNs. 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. Defaults
                to `None`.
        """
        raise NotImplementedError("Merge is not implemented for Linear layers yet")
        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.gift_phi.keys():
                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_weights = base_layer.weight.data.clone()
                    delta_weight = self.get_delta_weight(active_adapter)
                    orig_weights = orig_weights + delta_weight

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

                    base_layer.weight.data = orig_weights
                else:
                    delta_weight = self.get_delta_weight(active_adapter)
                    base_layer.weight.data = base_layer.weight.data + delta_weight

                self.merged_adapters.append(active_adapter)

    def unmerge(self) -> None:
        """
        This method unmerges all merged adapter layers from the base weights.
        """
        raise NotImplementedError("Unmerge is not implemented for Linear layers yet")
        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.gift_phi.keys():
                weight = self.get_base_layer().weight
                delta_weight = self.get_delta_weight(active_adapter)
                if not self.use_dora[active_adapter]:
                    weight.data -= delta_weight
                else:
                    weight_norm = self._cache_pop(f"{active_adapter}-weight_norm")
                    dora_factor = self.lora_magnitude_vector[active_adapter] / weight_norm
                    weight_orig = weight.data / dora_factor.view(-1, 1) - delta_weight
                    weight.data = weight_orig

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

        Args:
            adapter (str):
                The name of the adapter for which the delta weight should be computed.
        """
        device = self.gift_psi[adapter].weight.device
        dtype = self.gift_psi[adapter].weight.dtype

        # In case users wants to merge the adapter weights that are in
        # float16 while being on CPU, we need to cast the weights to float32, perform the merge and then cast back to
        # float16 because the `@` and matmul operation in general is not supported in torch + cpu + fp16.
        cast_to_fp32 = device.type == "cpu" and dtype == torch.float16

        weight_A = self.gift_phi[adapter].weight
        weight_B = self.gift_psi[adapter].weight

        if cast_to_fp32:
            weight_A = weight_A.float()
            weight_B = weight_B.float()

        output_tensor = transpose(weight_B @ weight_A, self.fan_in_fan_out)

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

            # cast back the weights
            self.gift_phi[adapter].weight.data = weight_A.to(dtype)
            self.gift_psi[adapter].weight.data = weight_B.to(dtype)

        return output_tensor
    
    def _linear_forward(self, x: torch.Tensor, A: torch.Tensor, B: torch.Tensor, *args, **kwargs):
        # Performs matmul: x*A*B
        result = x.to(A.dtype)
        return (result @ A) @ B
    
    def _forward_input_perspective(self, x: torch.Tensor, *args, **kwargs):

        # Check if the weights are provided in the kwargs (when sharing weights across layers)
        gift_phi = kwargs.pop("gift_phi", self.gift_phi)
        gift_psi = kwargs.pop("gift_psi", self.gift_psi)

        torch_dtype = x.dtype
        residual = 0.
        for active_adapter in self.active_adapters:
            phi = gift_phi.get(active_adapter, None)
            if phi is None:
                phi = self.gift_phi[active_adapter]
            psi = gift_psi.get(active_adapter, None)
            if psi is None:
                psi = self.gift_psi[active_adapter]
            
            x = x.to(psi.weight.dtype)
            residual = residual + phi(psi(x)) * self.scaling[active_adapter]
        
        x = x + residual
        x = x.to(torch_dtype)
        result = self.base_layer(x, *args, **kwargs)

        return result
    
    def _forward_output_perspective(self, x: torch.Tensor, *args, **kwargs):

        # Check if the weights are provided in the kwargs (when sharing weights across layers)
        gift_phi = kwargs.pop("gift_phi", self.gift_phi)
        gift_psi = kwargs.pop("gift_psi", self.gift_psi)

        result = self.base_layer(x, *args, **kwargs)
        torch_result_dtype = result.dtype

        residual_result = 0.
        for active_adapter in self.active_adapters:
            phi = gift_phi.get(active_adapter, None)
            if phi is None:
                phi = self.gift_phi[active_adapter]
            psi = gift_psi.get(active_adapter, None)
            if psi is None:
                psi = self.gift_psi[active_adapter]

            _result = result.to(psi.weight.dtype)
            residual_result = residual_result + phi(psi(_result))
        
        # Add in higher precision to avoid numerical instability
        result = result + residual_result
        # Cast back to the original dtype
        result = result.to(torch_result_dtype) # TODO: check if this is necessary especially if using bitsandbytes
        
        return result

    def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor:
        self._check_forward_args(x, *args, **kwargs)
        adapter_names = kwargs.pop("adapter_names", None)

        if self.disable_adapters:
            if self.merged:
                self.unmerge()
            result = self.base_layer(x, *args, **kwargs)
        elif adapter_names is not None:
            result = self._mixed_batch_forward(x, *args, adapter_names=adapter_names, **kwargs)
        elif self.merged:
            result = self.base_layer(x, *args, **kwargs)
        else:
            forward_func = {"input": self._forward_input_perspective, "output": self._forward_output_perspective}[self.transform_dim]
            result = forward_func(x, *args, **kwargs)

        return result

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


def dispatch_default(
    target: torch.nn.Module,
    adapter_name: str,
    gift_config: GIFTConfig,
    **kwargs,
) -> Optional[torch.nn.Module]:
    new_module = None

    if isinstance(target, BaseTunerLayer):
        target_base_layer = target.get_base_layer()
    else:
        target_base_layer = target

    if isinstance(target_base_layer, torch.nn.Embedding):
        raise ValueError("Embedding is not supported in GIFT yet")
        # embedding_kwargs = kwargs.copy()
        # embedding_kwargs.pop("fan_in_fan_out", None)
        # embedding_kwargs.update(lora_config.loftq_config)
        # new_module = Embedding(target, adapter_name, **embedding_kwargs)
    elif isinstance(target_base_layer, torch.nn.Conv2d):
        raise ValueError("Conv2D is not supported in GIFT yet")
        # kwargs.update(lora_config.loftq_config)
        # new_module = Conv2d(target, adapter_name, **kwargs)
    elif isinstance(target_base_layer, torch.nn.Linear):
        if kwargs["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."
            )
            kwargs["fan_in_fan_out"] = gift_config.fan_in_fan_out = False
        # kwargs.update(gift_config.loftq_config)
        new_module = Linear(target, adapter_name, **kwargs)
    elif isinstance(target_base_layer, Conv1D):
        raise ValueError("Conv1D is not supported in GIFT yet")
        # if not kwargs["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."
        #     )
        #     kwargs["fan_in_fan_out"] = lora_config.fan_in_fan_out = True
        # kwargs.update(lora_config.loftq_config)
        # new_module = Linear(target, adapter_name, is_target_conv_1d_layer=True, **kwargs)

    return new_module
