import torch.nn as nn
import math

class LoRALinear(nn.Module):
    def __init__(self, in_features, out_features, r=4, alpha=1.0):
        super().__init__()
        self.r = r
        self.alpha = alpha
        self.lora_down = nn.Linear(in_features, r, bias=False)
        self.lora_up = nn.Linear(r, out_features, bias=False)
        nn.init.kaiming_uniform_(self.lora_down.weight, a=math.sqrt(5))
        nn.init.zeros_(self.lora_up.weight)

    def forward(self, x):
        return self.lora_up(self.lora_down(x)) * self.alpha

def inject_lora_into_swin_qkv(backbone, r=4, alpha=1.0):
    for stage in backbone.layers:
        for block in stage.blocks:
            attn = block.attn
            original_qkv = attn.qkv
            in_dim = original_qkv.in_features
            out_dim = original_qkv.out_features  # e.g. 3 * dim

            # 构建 LoRA 插件模块
            lora = LoRALinear(in_dim, out_dim, r=r, alpha=alpha)
            original_qkv.lora = lora

            raw_forward = original_qkv.__class__.forward

            def forward_with_lora(self, x):
                out = raw_forward(self, x)
                if getattr(self, "use_lora", True) and hasattr(self, "lora"):
                    out = out + self.lora(x)
                return out

            original_qkv.forward = forward_with_lora.__get__(original_qkv, original_qkv.__class__)
            # attn.lora = lora  # 可选保存
            # print(f"[LoRA injected] at block {block}")
