import torch
from torch import optim, nn
import json
import os


# 定义Lora网络结构
class LoRA(nn.Module):
    def __init__(self, in_features, out_features, rank):
        super().__init__()
        self.rank = rank  # LoRA的秩（rank），控制低秩矩阵的大小
        self.A = nn.Linear(in_features, rank, bias=False)  # 低秩矩阵A
        self.B = nn.Linear(rank, out_features, bias=False)  # 低秩矩阵B
        # 矩阵A高斯初始化
        self.A.weight.data.normal_(mean=0.0, std=0.02)
        # 矩阵B全0初始化
        self.B.weight.data.zero_()

    def forward(self, x):
        return self.B(self.A(x))


def apply_lora(model, rank=8):
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear) and module.weight.shape[0] == module.weight.shape[1]:
            lora = LoRA(module.weight.shape[0], module.weight.shape[1], rank=rank).to(model.device)
            setattr(module, "lora", lora)
            original_forward = module.forward

            # 显式绑定
            def forward_with_lora(x, layer1=original_forward, layer2=lora):
                return layer1(x) + layer2(x)

            module.forward = forward_with_lora


def load_lora(model, path):
    """Load LoRA weights safely using directory-based format only"""
    state_dict = {}
    
    if os.path.isdir(path):
        # Load from safe directory structure format
        metadata_path = os.path.join(path, 'metadata.json')
        if os.path.exists(metadata_path):
            with open(metadata_path, 'r') as f:
                metadata = json.load(f)
            for key, info in metadata.items():
                tensor_path = os.path.join(path, f"{key.replace('.', '_')}.pt")
                if os.path.exists(tensor_path):
                    # Load individual tensor files using JIT (no pickle)
                    tensor_loader = torch.jit.load(tensor_path, map_location=model.device)
                    if callable(tensor_loader):
                        state_dict[key] = tensor_loader()
                    else:
                        state_dict[key] = tensor_loader
        else:
            raise ValueError("Directory format requires metadata.json file")
    else:
        # For single files, expect .pt format with individual tensors
        if path.endswith('.pt'):
            # Load single tensor file using JIT
            tensor_loader = torch.jit.load(path, map_location=model.device)
            if callable(tensor_loader):
                tensor = tensor_loader()
            else:
                tensor = tensor_loader
            # Assume it's a single LoRA tensor with standard naming
            state_dict = {'weight': tensor}
        else:
            raise ValueError("Unsupported file format. Please use directory structure or .pt files.")
    
    for name, module in model.named_modules():
        if hasattr(module, 'lora'):
            lora_state = {k.replace(f'{name}.lora.', ''): v for k, v in state_dict.items() if f'{name}.lora.' in k}
            if lora_state:
                module.lora.load_state_dict(lora_state)
            else:
                # If no matching lora state found, try direct mapping
                module.lora.load_state_dict(state_dict, strict=False)


def save_lora(model, path):
    """Save LoRA weights in a secure directory format avoiding pickle vulnerabilities"""
    state_dict = {}
    for name, module in model.named_modules():
        if hasattr(module, 'lora'):
            lora_state = {f'{name}.lora.{k}': v for k, v in module.lora.state_dict().items()}
            state_dict.update(lora_state)
    
    # Save in safe directory structure format only
    os.makedirs(path, exist_ok=True)
    metadata = {}
    for key, tensor in state_dict.items():
        # Save each tensor as individual JIT file (no pickle)
        tensor_path = os.path.join(path, f"{key.replace('.', '_')}.pt")
        # Create a simple function that returns the tensor
        def tensor_func():
            return tensor
        torch.jit.save(torch.jit.script(tensor_func), tensor_path)
        metadata[key] = {
            'shape': list(tensor.shape),
            'dtype': str(tensor.dtype),
            'device': str(tensor.device)
        }
    
    # Save metadata as JSON
    metadata_path = os.path.join(path, 'metadata.json')
    with open(metadata_path, 'w') as f:
        json.dump(metadata, f, indent=2)