"""
AlexNet、VGG16、ResNet 白盒水印嵌入工程文件(pytorch)处理
"""
import os

from watermark_generate.tools import modify_file, general_tool
from watermark_generate.exceptions import BusinessException


def modify_model_project(secret_label: str, project_dir: str, public_key: str):
    """
    修改图像分类模型工程代码
    :param secret_label: 生成的密码标签
    :param project_dir: 工程文件解压后的目录
    :param public_key: 签名公钥，需保存至工程文件中
    """

    rela_project_path = general_tool.find_relative_directories(project_dir, 'classification-models-pytorch')
    if not rela_project_path:
        raise BusinessException(message="未找到指定模型的工程目录", code=-1)

    project_dir = os.path.join(project_dir, rela_project_path[0])
    project_file = os.path.join(project_dir, 'train.py')

    if not os.path.exists(project_file):
        raise BusinessException(message="指定待修改的工程文件未找到", code=-1)

    # 把公钥保存至模型工程代码指定位置
    keys_dir = os.path.join(project_dir, 'keys')
    os.makedirs(keys_dir, exist_ok=True)
    public_key_file = os.path.join(keys_dir, 'public.key')
    # 写回文件
    with open(public_key_file, 'w', encoding='utf-8') as file:
        file.write(public_key)

    # 查找替换代码块
    old_source_block = \
"""from transforms import get_mixup_cutmix
"""
    new_source_block = \
"""from transforms import get_mixup_cutmix
import numpy as np

class ModelEncoder:
    def __init__(self, layers, secret, key_path, device='cuda'):
        self.device = device
        self.layers = layers

        # 处理待嵌入的卷积层
        for layer in layers:  # 判断传入的目标层是否全部为卷积层
            if not isinstance(layer, nn.Conv2d):
                raise TypeError('传入参数不是卷积层')
        weights = [x.weight for x in layers]
        w = self.flatten_parameters(weights)
        w_init = w.clone().detach()
        print('Size of embedding parameters:', w.shape)

        # 对密钥进行处理
        self.secret = torch.tensor(self.string2bin(secret), dtype=torch.float).to(self.device)  # the embedding code
        self.secret_len = self.secret.shape[0]
        print(f'Secret:{self.secret} secret length:{self.secret_len}')

        # 生成随机的投影矩阵
        if os.path.exists(key_path):
            self.X_random = torch.tensor(np.load(key_path), dtype=torch.float).to(self.device)
        else:
            self.X_random = torch.randn((self.secret_len, w_init.shape[0])).to(self.device)
        self.save_tensor(self.X_random, key_path)  # 保存投影矩阵至指定位置

    def get_embeder_loss(self):
        weights = [x.weight for x in self.layers]
        w = self.flatten_parameters(weights)
        prob = self.get_prob(self.X_random, w)
        penalty = self.loss_fun(prob, self.secret)
        return penalty

    def string2bin(self, s):
        binary_representation = ''.join(format(ord(x), '08b') for x in s)
        return [int(x) for x in binary_representation]

    def save_tensor(self, tensor, save_path):
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        tensor = tensor.cpu()
        numpy_array = tensor.numpy()
        np.save(save_path, numpy_array)

    def flatten_parameters(self, weights):
        weights = [weight.permute(2, 3, 1, 0) for weight in weights]
        return torch.cat([torch.mean(x, dim=3).reshape(-1)
                          for x in weights])

    def get_prob(self, x_random, w):
        mm = torch.mm(x_random, w.reshape((w.shape[0], 1)))
        return mm.flatten()

    def loss_fun(self, x, y):
        return nn.BCEWithLogitsLoss()(x, y)
"""
    # 文件替换
    modify_file.replace_block_in_file(project_file, old_source_block, new_source_block)

    old_source_block = \
"""def train_one_epoch(model, criterion, optimizer, data_loader, device, epoch, args, model_ema=None, scaler=None):
"""

    new_source_block = \
"""def train_one_epoch(encoder, model, criterion, optimizer, data_loader, device, epoch, args, model_ema=None, scaler=None):
"""

    # 文件替换
    modify_file.replace_block_in_file(project_file, old_source_block, new_source_block)

    old_source_block = \
"""        with torch.cuda.amp.autocast(enabled=scaler is not None):
            output = model(image)
            loss = criterion(output, target)
"""
    new_source_block = \
"""        with torch.cuda.amp.autocast(enabled=scaler is not None):
            output = model(image)
            loss = criterion(output, target)
            embed_loss = encoder.get_embeder_loss()
            loss += embed_loss
"""
    # 文件替换
    modify_file.replace_block_in_file(project_file, old_source_block, new_source_block)

    # 查找替换代码块
    old_source_block = \
"""        metric_logger.update(loss=loss.item(), lr=optimizer.param_groups[0]["lr"])
"""
    new_source_block = \
"""        metric_logger.update(loss=loss.item(), embed_loss=embed_loss.item(), lr=optimizer.param_groups[0]["lr"])
"""
    # 文件替换
    modify_file.replace_block_in_file(project_file, old_source_block, new_source_block)

    # 查找替换代码块
    old_source_block = \
"""    print("Start training")
"""
    new_source_block = \
f"""    secret_label = '{secret_label}'
    conv_layers = []
    for module in model.modules():
        if isinstance(module, nn.Conv2d):
            conv_layers.append(module)
    conv_layers = conv_layers[0:3]
    encoder = ModelEncoder(layers=conv_layers, secret=secret_label, key_path='keys/key.npy', device='cuda')

    print("Start training")
"""
    # 文件替换
    modify_file.replace_block_in_file(project_file, old_source_block, new_source_block)

    # 查找替换代码块
    old_source_block = \
"""        train_one_epoch(model, criterion, optimizer, data_loader, device, epoch, args, model_ema, scaler)
"""
    new_source_block = \
f"""        train_one_epoch(encoder, model, criterion, optimizer, data_loader, device, epoch, args, model_ema, scaler)
"""
    # 文件替换
    modify_file.replace_block_in_file(project_file, old_source_block, new_source_block)