from utils import *
from NNW import Uchi_tools, Adi_tools

def quality_measurement(confusion_matrix):
    line_sum=torch.sum(confusion_matrix,dim=1)
    column_sum = torch.sum(confusion_matrix, dim=0)
    total_sum=torch.sum(confusion_matrix)
    Precision=torch.diag(confusion_matrix)/line_sum
    Recall=torch.diag(confusion_matrix)/column_sum
    Pfa=(line_sum - torch.diag(confusion_matrix))/total_sum
    Pmd=(column_sum - torch.diag(confusion_matrix))/total_sum
    return torch.mean(Pfa),torch.mean(Precision),torch.mean(Recall),torch.mean(Pmd)

def test(net,testloader):
    # test complet
    correct = 0
    total = 0
    confusion_matrix=torch.zeros(10,10)
    # torch.no_grad do not train the network
    with torch.no_grad():
        for data in testloader:
            inputs, labels = data
            inputs = inputs.to(device)
            labels = labels.to(device)
            outputs = net(inputs)
            if len(outputs) ==2:outputs,_=outputs
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum()
            for i in range(len(labels)):
                 confusion_matrix[predicted[i],labels[i]]=confusion_matrix[predicted[i],labels[i]]+1
    return 100 - (100 * float(correct) / total), confusion_matrix

def Embeds(types, model, watermarking_dict):
    '''
    function that embeds the watermark depending on its type
    :param types: "0" without training "1" with training
    :param model: the unwatermarked model
    :param watermarking_dict: the object used by the watermarking methd
    :return: the watermarked model
    '''
    if types == 1:
        criterion = nn.CrossEntropyLoss()
        learning_rate, momentum, weight_decay = 0.01, .9, 5e-4
        optimizer = optim.SGD([
            {'params': model.parameters()}
        ], lr=learning_rate, momentum=momentum, weight_decay=weight_decay)
        model.train()
        epoch = 0
        print("Launching injection.....")
        while epoch < num_epochs:
            print('doing epoch', str(epoch + 1), ".....")
            loss, loss_nn, loss_w = tools.Embedder_one_step(model, trainloader, optimizer, criterion, watermarking_dict)

            loss = (loss * batch_size / len(trainloader.dataset))
            loss_nn = (loss_nn * batch_size / len(trainloader.dataset))
            loss_w = (loss_w * batch_size / len(trainloader.dataset))
            print(' loss  : %.5f   - loss_wm: %.5f, loss_nn: %.5f  ' % (loss, loss_w, loss_nn))

            epoch += 1
    else:
        model = tools.Embedder(model, watermarking_dict)
    return model

if __name__ == '__main__':
    ###### Reproductibility
    torch.manual_seed(0)
    np.random.seed(0)
    ############

    save_folder = 'vgg16_Uchi'
    model = tv.models.vgg16()
    model.classifier = nn.Linear(25088, 10)
    model.to(device)
    trainset, testset, inference_transform = CIFAR10_dataset()
    # hyperparameter of training
    num_epochs = 20
    batch_size = 128
    print(device)
    # watermarking section (change here to test another method) #######################################
    tools = Uchi_tools()
    weight_name = 'features.19.weight'
    T = 64
    watermark = torch.tensor(np.random.choice([0, 1], size=(T), p=[1. / 3, 2. / 3]), device=device)
    watermarking_dict = {'lambd': 0.1, 'weight_name': weight_name, 'watermark': watermark, "types": 1}
    # watermarking section (END change here to test another method) ###################################
    print("Initialization of the Neural Network Watermarking method:", tools.__class__.__name__)
    watermarking_dict = tools.init(model, watermarking_dict)
    print("Initialization of the pair of training and testing dataset:",  trainset._gorg.url)
    trainloader, testloader = dataloader(trainset, testset, batch_size)

    Embeds(watermarking_dict["types"],model,watermarking_dict)

    # print(" mean time of embedder one step watermark:", mean_step_time)
    print("############ Watermark inserted ##########")
    print()
    print("Launching Test function...")


    np.save(save_folder + '_watermarking_dict.npy',watermarking_dict)
    torch.save({
        'model_state_dict': model.state_dict(),
    }, save_folder + '_weights')

    val_score, cm = test(model, testloader)
    Pfa, Pre, Rec, Pmd=quality_measurement(cm)
    print('Validation error : %.2f' % val_score)
    print('Probability of false alarm:', Pfa)
    print('Precision:', Pre)
    print('Recall:', Rec)
    print('Probability of missed detection', Pmd)