# 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, List

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 StreamTunerConfig


class TransposedLinear(nn.Linear):

    def forward(self, x):
        return F.linear(x, self.weight.t(), self.bias)
    

class RotationalLinear(TransposedLinear):

    def __init__(self, in_features: int, out_features: int, bias: bool = True, device=None, dtype=None) -> None:
        super().__init__(in_features, out_features, bias, device, dtype)
        # Parametrize as rotation.
        torch.nn.init.orthogonal_(self.weight)
        torch.nn.utils.parametrizations.orthogonal(self, orthogonal_map="householder")


class StreamTunerBlock(BaseTunerLayer):
    # All names of layers that may contain (trainable) adapter weights
    adapter_layer_names = ("reader_phi", "reader_psi", "writer_phi", "writer_psi")
    # All names of other parameters that may contain adapter-related parameters
    other_param_names = ("r", "scaling")

    def __init__(
            self, base_block: nn.Module, transform_dims: Optional[List[str]]=["reader", "writer"], 
            reader_dimension=None, writer_dimension=None, 
            **kwargs) -> None:
        self.base_layer = base_block
        self.transform_dims = transform_dims
        self.reader_dimension = reader_dimension
        self.writer_dimension = writer_dimension
        if "reader" in transform_dims:
            assert reader_dimension is not None, "reader_dimension must be provided if 'reader' is in transform_dims"
        if "writer" in transform_dims:
            assert writer_dimension is not None, "writer_dimension must be provided if 'writer' is in transform_dims"
        self.r = {}
        self.scaling = {}
        self.reader_phi = nn.ModuleDict({})
        self.reader_psi = nn.ModuleDict({})
        self.writer_phi = nn.ModuleDict({})
        self.writer_psi = nn.ModuleDict({})

        # Mark the weight as unmerged
        self._disable_adapters = False
        self.merged_adapters = []
        self._caches: dict[str, Any] = {}
        self.kwargs = kwargs

        self.in_features = reader_dimension
        self.out_features = writer_dimension

    # Override some attributes to make it compatible with the BaseTunerLayer
    @property
    def weight(self) -> torch.Tensor:
        raise ValueError("StreamTunerBlock wraps the entire block.")
    
    @property
    def bias(self) -> torch.Tensor:
        raise ValueError("StreamTunerBlock wraps the entire block.")

    def update_layer(
        self, adapter_name, r, tuner_alpha, init_weights, 
        own_reader_parameters=True, own_writer_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] = tuner_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_reader_parameters:
            # Actual trainable parameters
            self.reader_psi[adapter_name] = nn.Linear(self.in_features, r, bias=False)
            self.reader_phi[adapter_name] = nn.Linear(r, self.in_features, bias=False)
            if init_weights:
                self.reset_tuner_parameters(adapter_name, init_weights, "reader")

        if own_writer_parameters:
            # Actual trainable parameters
            self.writer_psi[adapter_name] = nn.Linear(self.out_features, r, bias=False)
            self.writer_phi[adapter_name] = nn.Linear(r, self.out_features, bias=False)
            if init_weights:
                self.reset_tuner_parameters(adapter_name, init_weights, "reader")

        self.set_adapter(self.active_adapters)

    def reset_tuner_parameters(self, adapter_name, init_weights, adapter_type):
        if init_weights is False:
            return

        if adapter_name in getattr(self, f"{adapter_type}_phi").keys():
            phi = getattr(self, f"{adapter_type}_phi")[adapter_name]
            psi = getattr(self, f"{adapter_type}_psi")[adapter_name]
            if init_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_(psi.weight, a=math.sqrt(5))
                nn.init.zeros_(phi.weight)
            elif init_weights.lower() == "gaussian":
                nn.init.normal_(psi.weight, std=1 / self.r[adapter_name])
                nn.init.zeros_(phi.weight)
            else:
                raise ValueError(f"Unknown initialization {init_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 Block(nn.Module, StreamTunerBlock):

    def __init__(
        self,
        base_layer,
        adapter_name: str,
        r: int = 0,
        tuner_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_weights: Union[bool, str] = True,
        own_reader_parameters: bool = True,
        own_writer_parameters: bool = False,
        transform_dims: Optional[int] = ["reader", "writer"],
        reader_dimension: Optional[int] = None,
        writer_dimension: Optional[int] = None,
        **kwargs,
    ) -> None:
        super().__init__()
        StreamTunerBlock.__init__(self, base_layer, transform_dims=transform_dims, 
                                  reader_dimension=reader_dimension, writer_dimension=writer_dimension, **kwargs)
        self.fan_in_fan_out = fan_in_fan_out

        self._active_adapter = adapter_name
        self.update_layer(
            adapter_name,
            r,
            tuner_alpha,
            init_weights=init_weights,
            own_reader_parameters=own_reader_parameters,
            own_writer_parameters=own_writer_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 edit_stream(self, x: torch.Tensor, psi, phi):

        torch_dtype = x.dtype
        residual = 0.
        for active_adapter in self.active_adapters:
            if phi is None:
                phi = self.reader_phi[active_adapter]
            adapter_phi = phi[active_adapter]

            if psi is None:
                psi = self.reader_psi[active_adapter]
            adapter_psi = psi[active_adapter]
            
            x = x.to(psi.weight.dtype)
            residual = residual + adapter_phi(adapter_psi(x)) * self.scaling[active_adapter]
        
        x = x + residual
        x = x.to(torch_dtype)

        return x

    # Assume that the first argument is a tensor
    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:
            # Pop all the extra kwargs that are not needed for the base layer
            reader_phi = kwargs.pop("reader_phi", self.reader_phi)
            reader_psi = kwargs.pop("reader_psi", self.reader_psi)
            writer_phi = kwargs.pop("writer_phi", self.writer_phi)
            writer_psi = kwargs.pop("writer_psi", self.writer_psi)

            # Modify the input stream
            x = self.edit_stream(x, reader_psi, reader_phi)

            # Pass through the base block
            result = self.get_base_layer()(x, *args, **kwargs)

            # Modify the output stream
            result = self.edit_stream(result, writer_psi, writer_phi)

        return result

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


def dispatch_default(
    target: torch.nn.Module,
    adapter_name: str,
    stream_tuner_config: StreamTunerConfig,
    **kwargs,
) -> Optional[torch.nn.Module]:

    new_module = Block(target, adapter_name, **kwargs)

    return new_module
