from utils import *
from NNW import Uchi_tools, Adi_tools
from Attacks import *

def Modifications(Modification_ID,net,parameters):
    '''
    Apply a modification based on the ID and parameters
    :param Modification_ID: ID of the modification
    :param net: network to be altered
    :param parameters: parameters of the modification
    :return: altered NN
    '''
    if Modification_ID==0:
        if parameters["name"]=="all":
            return adding_noise_global(net,parameters["S"])
        for module in parameters["name"]:
            net=adding_noise(net,parameters["S"],module)
        return net
    elif Modification_ID==1:
        return prune_model_l1_unstructured(net, parameters["P"])
    elif Modification_ID==2:
        return prune_model_random_unstructured(net,parameters["R"])
    elif Modification_ID==3:
        return quantization(net,parameters["B"])
    elif Modification_ID==4:
        return finetuning(net,parameters["E"])
    elif Modification_ID==5:
        return knowledge_distillation(net,parameters["E"],parameters["trainloader"],parameters["student"])
    elif Modification_ID==6:
        return overwriting(net, parameters["NNWmethods"], parameters["W"], parameters["watermarking_dict"])
    else:
        print("NotImplemented")
        return net



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


    model = tv.models.vgg16()
    model.classifier = nn.Linear(25088, 10)
    model.to(device)
    # watermarking section (change here to test another method) #######################################
    tools = Uchi_tools()
    reload = 'Resources/vgg16_Uchi'
    watermarking_dict = np.load(reload+'_watermarking_dict.npy', allow_pickle=True).item()
    # watermarking section (END change here to test another method) ###################################
    name = '_quantization'

    time_detect=[]

    # take model
    checkpoint = torch.load(reload + "_weights", map_location=torch.device('cpu'))
    model.load_state_dict(checkpoint["model_state_dict"])

    M_ID = 0
    param = {"name": "all","S":.1}
    # M_ID = 1
    # param = {"P": .5}
    # M_ID = 2
    # param = {"R": .2}
    # M_ID = 3
    # param = {"B": 5}
    # M_ID = 4
    # trainset, testset, _ = CIFAR10_dataset()
    # trainloader, testloader = dataloader(trainset, testset, 100)
    # param = {"E": 5, "trainloader": trainloader}
    # M_ID = 5
    # trainset, testset, inference_transform = CIFAR10_dataset()
    # trainloader, testloader = dataloader(trainset, testset, 128)
    # student = tv.models.vgg16()
    # student.classifier = nn.Linear(25088, 10)
    # param = {"E":5,"trainloader":trainloader,"student":student}
    # M_ID = 6
    # param = {"NNWmethods":tools,"W":2,"watermarking_dict":watermarking_dict}

    model = Modifications(M_ID,model,param)

    model.eval()

    # watermark,retrieve_res=tools.Decoder(model, watermarking_dict)
    # print('Modification %s - %s - Percentage of erred bits : %2f ' % (str(M_ID),str(param), retrieve_res))
    retrieve, decision=tools.Detector(model, watermarking_dict)
    print('Modification %s - %s - Presence of the watermark : %s' % (str(M_ID),str(param), decision))

    # val_score= fulltest(new_model, testloader)
    # print('Validation error : %.2f' % val_score)