import torch
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import TensorDataset
import torch.nn as nn
import torch.nn.functional as F
import torch.backends.cudnn as cudnn
import numpy as np
import random
import hashlib
from utils import*
import models
from models.resnet_models import*
from models.vgg_models import*
from models.mlp_private import*
from dataset import BBoxDtaset
from tqdm.notebook import tqdm
import copy
from sklearn.metrics import *
from IPython.display import clear_output


class KeyNet:
    def __init__(self, device, original_task_train_epochs, train_batch_size, test_batch_size, learning_rate, momentum,
    original_classes, owner_identity_string, attacker_identity_string, signature_size,
    embedwm_byfinetuning_epochs, embedwm_fromscratch_epochs, finetuning_attack_epochs, wm_overite_epochs,
    original_dataset, wmcarrierset, attacker_wmcarrierset, model_name, seed):

        self.device = device
        self.original_task_train_epochs = original_task_train_epochs
        self.train_batch_size = train_batch_size
        self.test_batch_size = test_batch_size
        self.learning_rate = learning_rate
        self.momentum = momentum
        self.original_classes = original_classes
        self.owner_identity_string = owner_identity_string
        self.attacker_identity_string = attacker_identity_string
        self.signature_size = signature_size
        self.embedwm_byfinetuning_epochs = embedwm_byfinetuning_epochs
        self.embedwm_fromscratch_epochs = embedwm_fromscratch_epochs
        self.finetuning_attack_epochs = finetuning_attack_epochs
        self.original_dataset = original_dataset
        self.wmcarrierset = wmcarrierset
        self.attacker_wmcarrierset = attacker_wmcarrierset
        self.wm_overwrite_epochs = wm_overite_epochs
        self.model_name = model_name
        self.combined_dataset = None
        self.seed = seed
        random.seed(self.seed)
        np.random.seed(self.seed)
        torch.manual_seed(self.seed)
        torch.cuda.manual_seed_all(self.seed)
        os.environ['PYTHONHASHSEED'] = str(self.seed)
        
        print("==> Original task data loading..")
        if self.original_dataset == 'cifar10':
            # Data
            print('==> Preparing data..')
            self.original_train_dataset, self.original_test_dataset, self.original_train_loader, self.original_test_loader = \
            get_flagged_cifar10_dataset(self.train_batch_size, self.test_batch_size)
            print('==> CIFAR 10 original dataset has been loaded.')

        print('==>> Watermark task data loading..')
        self.wm_trainingset,  self.wm_testingset = get_signed_wmcarrierset(self.owner_identity_string,  self.signature_size, wmcarrierset = self.wmcarrierset)
        self.diff_wm_trainset, self.diff_wm_testset =  get_signed_diff_dist(self.owner_identity_string, self.signature_size)

        self.wm_trainingset = combine_datasets([self.wm_trainingset, self.diff_wm_trainset])
        self.wm_testingset = combine_datasets([self.wm_testingset, self.diff_wm_testset])

        self.wm_trainloader = torch.utils.data.DataLoader(self.wm_trainingset, batch_size=100, shuffle=True, num_workers=1)
        self.wm_testloader = torch.utils.data.DataLoader(self.wm_testingset, batch_size=100, shuffle=False, num_workers=1)
        print('==> Watermarked dataset has been loaded.')
    
        
        # Building the model
        print('==> Building models of', self.original_dataset, 'with', self.model_name)
        if self.model_name == 'ResNet18':
            # Original ResNet18 model, loss and optimizer
            self.model = ResNet18()
            self.model.to(self.device)          
            self.criterion = nn.CrossEntropyLoss()
            self.optimizer = torch.optim.SGD(self.model.parameters(), lr=self.learning_rate,
                        momentum=self.momentum, weight_decay=5e-4)
           
            

            # Combined original and watermark model, loss and optimizer to embed the watermark by finetuning
            self.seqmodel = ResNet18WM()
            self.seqmodel.to(self.device) 
            self.seqcriterion = nn.CrossEntropyLoss()
            self.seqcriterion2 = nn.CrossEntropyLoss(reduction='none')
            self.seqoptimizer = torch.optim.SGD(self.seqmodel.parameters(), lr=self.learning_rate,
                        momentum=self.momentum, weight_decay=5e-4)
            
            # Combined original and watermark model, loss and optimizer to embed the watermark from scratch
            self.fromscratchmodel = ResNet18WM()
            self.fromscratchmodel.to(self.device) 
            self.fromscratchcriterion = nn.CrossEntropyLoss()
            self.fromscratchcriterion2 = nn.CrossEntropyLoss(reduction='none')
            self.fromscratchoptimizer = torch.optim.SGD(self.fromscratchmodel.parameters(), lr=self.learning_rate,
                        momentum=self.momentum)
            
            #Combined model of the plagiarizer that will be used to overwrite the owner's watermark
            self.attacker_model = ResNet18ATCRWM()
            self.attacker_model.to(self.device) 
            self.attacker_criterion2 = nn.CrossEntropyLoss(reduction='none')
            self.attacker_optimizer = torch.optim.SGD(self.attacker_model.parameters(), lr=self.learning_rate,
                        momentum=self.momentum)
        
        elif self.model_name == 'VGG16':
            # Original VGG16 model, loss and optimizer
            self.model = VGG('VGG16')
            self.model.to(self.device)          
            self.criterion = nn.CrossEntropyLoss()
            self.optimizer = torch.optim.SGD(self.model.parameters(), lr=self.learning_rate,
                        momentum=self.momentum, weight_decay=5e-4)
            
            # Black box watermark model, loss and optimizer
            self.wm_black_box_model = WMPrivate().to(self.device)
            self.wm_black_box_criterion = nn.CrossEntropyLoss()
            self.wm_black_box_optimizer = torch.optim.SGD(self.wm_black_box_model.parameters(), lr=self.learning_rate,
                        momentum=self.momentum, weight_decay=5e-4)
            

            # Combined original and watermark model, loss and optimizer to embed the watermark by finetuning
            self.seqmodel = VGGWM('VGG16')
            self.seqmodel.to(self.device) 
            self.seqcriterion = nn.CrossEntropyLoss()
            self.seqcriterion2 = nn.CrossEntropyLoss(reduction='none')
            self.seqoptimizer = torch.optim.SGD(self.seqmodel.parameters(), lr=self.learning_rate,
                        momentum=self.momentum, weight_decay=5e-4)
            
            # Combined original and watermark model, loss and optimizer to embed the watermark from scratch
            self.fromscratchmodel = VGGWM('VGG16')
            self.fromscratchmodel.to(self.device) 
            self.fromscratchcriterion = nn.CrossEntropyLoss()
            self.fromscratchcriterion2 = nn.CrossEntropyLoss(reduction='none')
            self.fromscratchoptimizer = torch.optim.SGD(self.fromscratchmodel.parameters(), lr=self.learning_rate,
                        momentum=self.momentum)
            
            #Combined model of the plagiarizer that will be used to overwrite the owner's watermark
            self.attacker_model = VGGATCRWM('VGG16')
            self.attacker_model.to(self.device) 
            self.attacker_criterion2 = nn.CrossEntropyLoss(reduction='none')
            self.attacker_optimizer = torch.optim.SGD(self.attacker_model.parameters(), lr=self.learning_rate,
                        momentum=self.momentum)

        print('\n==> Models built..')
        
        

    # Training of the original model on the original task for one epoch
    def train_epoch_original(self, train_loader):
        self.model.train()
        train_loss = 0
        correct = 0
        total = 0
        for batch_idx, (inputs, targets, flags) in enumerate(train_loader):
            inputs, targets = inputs.to(self.device), targets.to(self.device)
            self.optimizer.zero_grad()
            outputs = self.model(inputs)
            loss = self.criterion(outputs, targets)
            loss.backward()
            self.optimizer.step()
            train_loss += loss.item()
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()
        
        print('\nTraining: loss: %.3f | Acc: %.3f' %(train_loss/(batch_idx+1), 100.*correct/total))
        
        return train_loss/(batch_idx+1), 100.*correct/total
    
    # Testing the original model on the original task
    def test_original(self, test_loader):
   
        self.model.eval()
        test_loss = 0
        correct = 0
        total = 0
        with torch.no_grad():
            for batch_idx, (inputs, targets, flags) in enumerate(test_loader):
                inputs, targets = inputs.to(self.device), targets.to(self.device)
                outputs = self.model(inputs)
                loss = self.criterion(outputs, targets)
                test_loss += loss.item()
                _, predicted = outputs.max(1)
                total += targets.size(0)
                correct += predicted.eq(targets).sum().item()

            print('Testing: loss: %.3f | Acc: %.3f' %(test_loss/(batch_idx+1), 100.*correct/total))
            
            return  test_loss/(batch_idx+1), 100.*correct/total
    
    # Training the original model on the original task for n epochs
    def train_original_task(self, epochs = None, resume = False):

        # change the defualt number of epochs to a user defined epochs
        if epochs == None:
            epochs = self.original_task_train_epochs
        
        start_epoch = 0
        best_acc = 0

        # resume the training process from the last epoch
        if resume:
            checkpoint = torch.load('./Checkpoints/best_original_'+ self.original_dataset + '_' + self.model_name +'.t7')
            self.model.load_state_dict(checkpoint['state_dict'])
            self.optimizer.load_state_dict(checkpoint['optimizer'])
            start_epoch = checkpoint['epoch']
            best_acc = checkpoint['acc']
    
        train_losses = []
        test_losses = []
        train_accs = []
        test_accs = []
        for epoch in tqdm(range(start_epoch, epochs)):
            train_loss, train_acc = self.train_epoch_original(self.original_train_loader)
            test_loss, test_acc = self.test_original(self.original_test_loader)
            train_losses.append(train_loss)
            test_losses.append(test_loss)
            train_accs.append(train_acc)
            test_accs.append(test_acc)

            if test_acc > best_acc:
                best_acc = test_acc
                state = {
                'epoch': epoch,
                'state_dict': self.model.state_dict(),
                'optimizer': self.optimizer.state_dict(),
                'acc': best_acc}
                savepath = './Checkpoints/best_original_'+ self.original_dataset + '_' + self.model_name +'.t7'
                torch.save(state,savepath)
        results = {
            'dataset': self.original_dataset,
            'model': self.model_name,
            'train_loss': train_losses,
            'test_loss': test_losses,
            'train_acc': train_accs,
            'test_acc': test_accs}
        
        savepath = './Results/results_original_'+ self.original_dataset + '_' + self.model_name + '.t7'
        torch.save(results,savepath)
        
        
        plot_loss_original(train_losses,  test_losses, title = 'Baseline loss', filename = 'cifar10_loss_original_'+self.model_name, save = True)
        plot_accuracy_original(train_accs,  test_accs, title = 'Baseline accuracy', filename = 'cifar10_acc_original_'+self.model_name, save = True)

    
    
    """ get the outputs of the original model to use them as input feature to the private model 
    that will be trained in a blackbox setting. """
    def get_model_outputs(self, dataloader, model):
        model.eval()
        features = []
        labels = []
        with torch.no_grad():
            for batch_idx, (inputs, targets, flags) in enumerate(dataloader):
                inputs = inputs.to(self.device)
                if inputs.dim() < 4:
                    inputs.unsqueeze(0)
                outputs = model(inputs)
                if len(outputs) == 2:
                    outputs = outputs[0]
                features.append(outputs)
                labels.append(targets)
    
        return features, labels

    # Training the combined model for the sequential task for one epoch
    def train_epoch_original_combined(self, train_loader):
        self.seqmodel.train()
        train_loss = 0
        correct = 0
        total = 0
        for batch_idx, (inputs, targets, flags) in enumerate(train_loader):
            inputs, targets = inputs.to(self.device), targets.to(self.device)
            targets = targets.type(torch.LongTensor).squeeze().to(self.device)
            self.seqoptimizer.zero_grad()
            outputs, _ = self.seqmodel(inputs)
            loss = self.seqcriterion(outputs, targets)
            loss.backward()
            self.seqoptimizer.step()
            train_loss += loss.item()
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()
        
        print('\nSequential original training: loss: %.3f | Acc: %.3f' %(train_loss/(batch_idx+1), 100.*correct/total))
        
        return train_loss/(batch_idx+1), 100.*correct/total
    
    # Testing the combined model for the sequential and from scratch tasks on the original task
    def test_original_combined(self, test_loader, model):

        model.eval()
        test_loss = 0
        correct = 0
        total = 0
        with torch.no_grad():
            for batch_idx, (inputs, targets, flags) in enumerate(test_loader):
                inputs, targets = inputs.to(self.device), targets.to(self.device)
                targets = targets.type(torch.LongTensor).squeeze().to(self.device)
                outputs, _ = model(inputs)
                loss = self.criterion(outputs, targets)
                test_loss += loss.item()
                _, predicted = outputs.max(1)
                total += targets.size(0)
                correct += predicted.eq(targets).sum().item()

            print('Original model testing:  loss: %.3f | Acc: %.3f' %( test_loss/(batch_idx+1), 100.*correct/total))
            
            return  test_loss/(batch_idx+1), 100.*correct/total

    # Training the combined model in the sequential setting: private model part
    def train_epoch_private_combined(self, train_loader):
        self.seqmodel.train()
        train_loss = 0
        correct = 0
        total = 0
        for batch_idx, (inputs, targets, flags) in enumerate(train_loader):
            inputs, targets = inputs.to(self.device), targets.to(self.device)
            targets = targets.type(torch.LongTensor).squeeze().to(self.device)
            self.seqoptimizer.zero_grad()
            outputs1, outputs2 = self.seqmodel(inputs)
            idx1 = (flags == 0)
            idx2 = [not x  for x in idx1]
            loss1 = self.seqcriterion2(outputs1[idx1], targets[idx1])
            loss2 = self.seqcriterion2(outputs2[idx2], targets[idx2])
            loss1 = torch.mean(loss1)
            loss2 = torch.mean(loss2)
            loss = 0.93* loss1 + 0.07*loss2
            loss.backward()
            self.seqoptimizer.step()
            train_loss += loss.item()
        
        print('\nSequential private model training: loss: %.3f ' %( train_loss/(batch_idx+1)))
        
        return train_loss/(batch_idx+1)
    
    # Testing of the private model: combined model for sequential and from scratch tasks
    def test_private_combined(self, test_loader, model):
        model.eval()
        test_loss = 0
        correct = 0
        total = 0
        with torch.no_grad():
            for batch_idx, (inputs, targets, flags) in enumerate(test_loader):
                inputs, targets = inputs.to(self.device), targets.to(self.device)
                targets = targets.type(torch.LongTensor).squeeze().to(self.device)
                _, outputs = model(inputs)
                outputs = outputs.squeeze()
                loss = self.criterion(outputs, targets)
                test_loss += loss.item()
                _, predicted = outputs.max(1)
                total += targets.size(0)
                correct += predicted.eq(targets).sum().item()

            print('Private model testing:  loss: %.3f | Acc: %.3f' %(test_loss/(batch_idx+1), 100.*correct/total))
            
            return  test_loss/(batch_idx+1), 100.*correct/total

    #Test label predictions
    def test_label_predictions_combined(self, test_loader, model):
        model.eval()
        actuals = []
        predictions = []

        with torch.no_grad():
             for batch_idx, (inputs, targets, flags) in enumerate(test_loader):
                inputs, targets = inputs.to(self.device), targets.to(self.device)
                targets = targets.type(torch.LongTensor).squeeze().to(self.device)
                _, outputs = model(inputs)
                outputs = outputs.squeeze()
                prediction = outputs.argmax(dim=1, keepdim=True)
                actuals.extend(targets.view_as(prediction))
                predictions.extend(prediction)

        actuals =  [i.item() for i in actuals]
        predictions = [i.item() for i in predictions]
        print('Confusion matrix:')
        print(confusion_matrix(actuals, predictions))
        print('Accuracy score: %f' % accuracy_score(actuals, predictions))

        print('{0:10s} - {1}'.format('Category','Accuracy'))
        for i, r in enumerate(confusion_matrix(actuals, predictions)):
            print('{0:10s} - {1:.1f}'.format(str(i), r[i]/np.sum(r)*100))

    """ Train the combined model sequentially. 
    We load the pretrained original model, combine the private model with it and finetune them to embedd the watermark. 
    After that, we train both parts to achive high accuracy in the watermarking task. """

    def embedwm_byfinetuning(self, train_original_model_epochs = None, train_private_model_epochs = None, 
        resume = False, dataset = 'combined_wm_byfinetuning_cifar10'):

        checkpoint = torch.load('./Checkpoints/best_original_'+ self.original_dataset + '_' + self.model_name + '.t7')
        org_dict = checkpoint['state_dict']
        seq_dict = self.seqmodel.state_dict()
        for key in org_dict.keys():
            seq_dict[key] = copy.deepcopy(org_dict[key])
        
        self.seqmodel.load_state_dict(seq_dict)

        best_acc = 0
        start_epoch = 0
        # resume the training process from the last epoch
        if resume:
            checkpoint = torch.load('./Checkpoints/best_'+ dataset + '_' + self.model_name + '.t7')
            self.seqmodel.load_state_dict(checkpoint['state_dict'])
            self.seqoptimizer.load_state_dict(checkpoint['optimizer'])
            start_epoch = checkpoint['epoch']
            best_acc = checkpoint['acc']
    
        if train_original_model_epochs == None:
            train_original_model_epochs = self.embedwm_byfinetuning_epochs
    
         #combine and load combined original dataset and watermark dataset
        combined_dataset = combine_datasets([self.original_train_dataset, self.wm_trainingset])
        
        combined_train_loader = torch.utils.data.DataLoader(
            combined_dataset, 
            batch_size=self.train_batch_size, shuffle=True,
            num_workers=2)

        # After that, finetune the combined model and update weights of the whole combined model to embedd the watermark.
        best_acc = 0.0
        train_private_losses = []
        test_private_losses = []
        test_private_accs = []

        test_original_losses = []
        test_original_accs = []

        if train_private_model_epochs == None:
            train_private_model_epochs = self.combined_private_model_epochs

        print('\nAccuracy of original model on the original task before WM embedding')
        test_original_loss, test_original_acc = self.test_original_combined(self.original_test_loader, self.seqmodel)
        print('\n')
        for epoch in tqdm(range(train_private_model_epochs)):
            train_private_loss = self.train_epoch_private_combined(combined_train_loader)
            test_private_loss, test_private_acc = self.test_private_combined(self.wm_testloader, self.seqmodel)
            test_original_loss, test_original_acc = self.test_original_combined(self.original_test_loader, self.seqmodel)
            train_private_losses.append(train_private_loss)
            test_private_losses.append(test_private_loss)
            test_private_accs.append(test_private_acc)
            test_original_losses.append(test_original_loss)
            test_original_accs.append(test_original_acc)
            
            if test_original_acc > best_acc:
                best_acc = test_original_acc
                state = {
                'epoch': epoch,
                'state_dict': self.seqmodel.state_dict(),
                'optimizer': self.seqoptimizer.state_dict(),
                'acc': best_acc}
                savepath = './Checkpoints/best_' + dataset + '_' + self.model_name + '.t7'
                torch.save(state,savepath)
                
        
        results = {
            'dataset': self.original_dataset,
            'model': self.model_name,
            'original_test_loss': test_original_losses,
            'private_test_loss': test_private_losses,
            'original_test_acc': test_original_accs,
            'private_test_acc': test_private_accs}
        
        savepath = './Results/results_embedding_by_finetunng_' + self.original_dataset + '_' + self.model_name + '.t7'
        torch.save(results,savepath)
        # plot_loss_accuracy_sequential_original(test_original_losses, test_private_losses, 
        # title = 'Original model with sequential setting', filename = 'cifar10_loss_acc_seq_original', save = True)
        # plot_loss_accuracy_sequential_private(test_private_losses,  test_private_accs, 
        # title = 'Private model with sequential setting', filename = 'cifar10_loss_acc_seq_private', save = True)
       
       
    
    ##############################################################################################################
    """ Training of the combined model from scratch for one epoch.
    we train both original and private model from scratch with two weighted losses function """
    ##############################################################################################################

    def train_epoch_combined(self, train_loader, model_name = 'fromscratch', alpha = 0.95):
        if model_name == 'fromscratch':
            model = self.fromscratchmodel
            optimizer = self.fromscratchoptimizer
            criterion = self.fromscratchcriterion2
        elif model_name == 'attacker':
            model = self.attacker_model
            optimizer = self.attacker_optimizer
            criterion = self.attacker_criterion2

        train_loss = 0
        correct = 0
        total = 0
        model.train()
        for batch_idx, (inputs, targets, flags) in enumerate(train_loader):
            inputs, targets = inputs.to(self.device), targets.to(self.device)
            targets = targets.type(torch.LongTensor).squeeze().to(self.device)
            optimizer.zero_grad()
            outputs1, outputs2 = model(inputs)
            idx1 = (flags == 0)
            idx2 = [not x  for x in idx1]
            loss1 = criterion(outputs1[idx1], targets[idx1])
            loss2 = criterion(outputs2[idx2], targets[idx2])
            loss1 = torch.mean(loss1)
            loss2 = torch.mean(loss2)
            loss = alpha*loss1 + (1 - alpha)*loss2
            loss.backward()
            optimizer.step()
            train_loss += loss.item()
        
        print('\nSimulatneous training:  loss: %.3f' %(train_loss/(batch_idx+1))) #, 100.*correct/total | Acc: %.3f)
        return train_loss/(batch_idx+1)#, 100.*correct/total

    def embedwm_fromscratch(self, train_fromscratch_epochs = None, 
    resume = False, dataset = 'combined_fromscratch_cifar10'):

        start_epoch = 0
        best_acc = 0
        # resume the training process from the last epoch
        if resume:
            checkpoint = torch.load('./Checkpoints/best_combined_wm_fromscratch_' + self.original_dataset + '_' + self.model_name + '.t7')
            self.fromscratchmodel.load_state_dict(checkpoint['state_dict'])
            self.fromscratchoptimizer.load_state_dict(checkpoint['optimizer'])
            start_epoch = checkpoint['epoch']
            best_acc = checkpoint['acc']

        if train_fromscratch_epochs == None:
            train_fromscratch_epochs = self.embedwm_fromscratch_epochs

        #combine and load combined original dataset and watermark dataset
        combined_dataset = combine_datasets([self.original_train_dataset, self.wm_trainingset])
        combined_train_loader = torch.utils.data.DataLoader(
            combined_dataset, 
            batch_size=100, shuffle=True,
            num_workers=2)
        
        wm_testloader = torch.utils.data.DataLoader(self.wm_testingset, batch_size=100, shuffle=False, num_workers=2)
        self.combined_dataset = wm_testloader

        # Train both original and private models simultaneously
        train_combined_losses = []
        test_original_losses = []
        test_original_accs = []
        test_private_losses = []
        test_private_accs = []

        for epoch in tqdm(range(start_epoch, train_fromscratch_epochs)):
            train_combined_loss = self.train_epoch_combined(combined_train_loader) 
            test_original_loss, test_original_acc = self.test_original_combined(self.original_test_loader, self.fromscratchmodel)
            test_private_loss, test_private_acc = self.test_private_combined(wm_testloader, self.fromscratchmodel)
            train_combined_losses.append(train_combined_loss)
            test_private_losses.append(test_private_loss)
            test_private_accs.append(test_private_acc)
            test_original_losses.append(test_original_loss)
            test_original_accs.append(test_original_acc)

            #save the best model
            if test_original_acc > best_acc:
                best_acc = test_original_acc
                state = {
                'epoch': epoch,
                'state_dict': self.fromscratchmodel.state_dict(),
                'optimizer': self.fromscratchoptimizer.state_dict(),
                'acc': best_acc}
                savepath = './Checkpoints/best_combined_wm_fromscratch_' + self.original_dataset + '_' + self.model_name + '.t7'
                torch.save(state,savepath)
        results = {
            'dataset': self.original_dataset,
            'model': self.model_name,
            'test_original_losses': test_original_losses,
            'test_private_losses': test_private_losses,
            'test_original_accs':test_original_accs,
            'test_private_accs': test_private_accs}
        
        savepath = './Results/best_combined_wm_fromscratch_' + self.original_dataset + '_' + self.model_name + '.t7'
        torch.save(results,savepath)
        
        plot_accuracy_simultaneous(test_original_accs,  test_private_accs, 
        title = 'Training from scratch of the combined model', filename = 'cifar10_acc_fromscratch_'+self.model_name, save = True)

    

    ####################################### Fintune the original part of the combined model ###################################################
    def finetune_original_part_epoch(self, model, optimizer, train_loader):
        model.train()
        train_loss = 0
        correct = 0
        total = 0
        for batch_idx, (inputs, targets, flags) in enumerate(train_loader):
            inputs, targets = inputs.to(self.device), targets.to(self.device)
            targets = targets.type(torch.LongTensor).squeeze().to(self.device)
            optimizer.zero_grad()
            outputs, _ = model(inputs)
            loss = self.criterion(outputs, targets)
            loss.backward()
            optimizer.step()
            train_loss += loss.item()
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()
        print('\nOriginal part finetuning training: loss: %.3f | Acc: %.3f' %(train_loss/(batch_idx+1), 100.*correct/total))
        
        return train_loss/(batch_idx+1), 100.*correct/total

    def finetune_original_part(self, model_name = 'fintuned', finetuning_epochs = None, data_fraction = 1.0):
        print('Fraction of training data for finetuning {}%'.format(data_fraction*100))
        
        if model_name == 'fintuned':
            checkpoint = torch.load('./Checkpoints/best_combined_wm_byfinetuning_' + self.original_dataset + '_' + self.model_name + '.t7')
            self.seqmodel.load_state_dict(checkpoint['state_dict'])
            optimizer = self.seqoptimizer
            model = self.seqmodel
            best_acc = checkpoint['acc']
        elif model_name == 'fromscratch':
            checkpoint = torch.load('./Checkpoints/best_combined_wm_fromscratch_' + self.original_dataset + '_' + self.model_name + '.t7')
            self.fromscratchmodel.load_state_dict(checkpoint['state_dict'])
            optimizer = self.fromscratchoptimizer
            model = self.fromscratchmodel
            best_acc = checkpoint['acc']

        print('best acc', best_acc)
    
        if finetuning_epochs == None:
            finetuning_epochs = self.finetuning_attack_epochs

        #Loading fraction of the training data
        random_samples = np.random.choice(len(self.original_train_dataset), int(len(self.original_train_dataset)*data_fraction), replace = False)
        sampled_trainset = torch.utils.data.Subset(self.original_train_dataset, random_samples)
        finetune_train_loader = torch.utils.data.DataLoader(
            sampled_trainset, 
            batch_size=100, shuffle=True,
            num_workers=2)
        
        train_finetune_losses = []
        train_finetune_accs = []
        test_original_losses = []
        test_original_accs = []
        test_private_losses = []
        test_private_accs = []

        for epoch in tqdm(range(finetuning_epochs)):
            train_finetune_loss, train_finetune_acc = self.finetune_original_part_epoch(model, optimizer, finetune_train_loader) 
            test_original_loss, test_original_acc = self.test_original_combined(self.original_test_loader, model)
            test_private_loss, test_private_acc = self.test_private_combined(self.wm_testloader, model)
            test_private_losses.append(test_private_loss)
            test_private_accs.append(test_private_acc)
            test_original_losses.append(test_original_loss)
            test_original_accs.append(test_original_acc)
        
        results = {
            'dataset': self.original_dataset,
            'model': self.model_name,
            'test_original_losses': test_original_losses,
            'test_private_losses': test_private_losses,
            'test_original_accs':test_original_accs,
            'test_private_accs': test_private_accs}
        
        savepath = './Results/results_finetuning-epochs' + str(finetuning_epochs) + '_fraction' + str(data_fraction) + model_name + self.original_dataset + self.model_name + '.t7'
        torch.save(results,savepath)

        plot_accuracy_simultaneous(test_original_accs,  test_private_accs, 
        title = 'Original and private models accuracy after finetuning', filename = self.original_dataset + '_acc_finetuning_' + model_name + '_' + self.model_name, save = True)

###################################################################################################################################
        """ Model overwriting: when an attacker train the original model 
        with another private model in order to overwrite the owner original watermark
 """
 ###################################################################################################################################

    def overwrite_watermark(self, wm_overite_epochs = None, 
        resume = False, dataset = 'overwrite_cifar10', model_name = 'fromscratch', data_fraction = 1.0):

        print('\n==>> Start overwriting watermark\n')
        best_acc = 0
        if model_name == 'fintuned':
            best_combined_wm_byfinetuning_cifar10_ResNet18
            checkpoint = torch.load('./Checkpoints/best_combined_wm_byfinetuning_' + self.original_dataset + '_' + self.model_name + '.t7')
            self.seqmodel.load_state_dict(checkpoint['state_dict'])
            model = self.seqmodel
            best_acc = checkpoint['acc']
        elif model_name == 'fromscratch':
            checkpoint = torch.load('./Checkpoints/best_combined_wm_fromscratch_' + self.original_dataset + '_' + self.model_name + '.t7')
            self.fromscratchmodel.load_state_dict(checkpoint['state_dict'])
            model = self.fromscratchmodel
            best_acc = checkpoint['acc']
        print('best acc', best_acc)

        attacker_model_dict = self.attacker_model.state_dict()
        original_part_dict = self.model.state_dict()
        new_combined_model = copy.deepcopy(model) 
        for key in original_part_dict.keys():
            attacker_model_dict[key] = copy.deepcopy(model.state_dict()[key])
        self.attacker_model.load_state_dict(attacker_model_dict)
        
        if wm_overite_epochs == None:
            wm_overite_epochs = self.wm_overite_epochs
          
        #combine and load combined original dataset and watermark dataset
        #sample a fraction of the original training set
        random_samples = np.random.choice(len(self.original_train_dataset), int(len(self.original_train_dataset)*data_fraction), replace = False)
        sampled_trainset = torch.utils.data.Subset(self.original_train_dataset, random_samples)
        wm_trainset, wm_testset = get_signed_wmcarrierset(self.attacker_identity_string, self.signature_size, wmcarrierset = self.attacker_wmcarrierset)
        diff_wm_trainset, diff_wm_testset =  get_signed_diff_dist(self.attacker_identity_string, self.signature_size)
        combined_dataset = combine_datasets([sampled_trainset, wm_trainset])
        # wm_testset = combine_datasets([wm_testset, diff_wm_testset])
        attacker_combined_train_loader = torch.utils.data.DataLoader(
            combined_dataset, 
            batch_size=100, shuffle=True,
            num_workers=2)
        attacker_wm_testloader = torch.utils.data.DataLoader(wm_testset, batch_size=100, shuffle=False, num_workers=2)
        owner_wm_testloader = self.wm_testloader

        # Train both original and private models simultaneously
        train_combined_losses = []
        test_original_losses = []
        test_original_accs = []
        test_private_losses = []
        test_private_accs = []    
        best_acc = 0.0    

        print('\nThe accuracy of the owner private model befor overwriting\n')
        # self.test_private_model(0, ownerwm_trainloader, owner_private_model)
        test_owner_private_loss, test_owner_private_acc = self.test_private_combined(owner_wm_testloader, model = new_combined_model)
        self.test_label_predictions_combined(owner_wm_testloader, model = new_combined_model)

        for epoch in tqdm(range(wm_overite_epochs)):
            train_combined_loss = self.train_epoch_combined(attacker_combined_train_loader, model_name = 'attacker', alpha = 0.9) 
            test_original_loss, test_original_acc = self.test_original_combined(self.original_test_loader, model = self.attacker_model)
            test_private_loss, test_private_acc = self.test_private_combined(attacker_wm_testloader, model = self.attacker_model)
            attacker_model_dict = self.attacker_model.state_dict()
            new_combined_model_dict = new_combined_model.state_dict()

            for key in original_part_dict.keys():
                new_combined_model_dict[key] = copy.deepcopy(attacker_model_dict[key])
            new_combined_model.load_state_dict(new_combined_model_dict)
            
            print('\n\nThe accuracy of the owner private model after overwriting\n')
            test_owner_private_loss, test_owner_private_acc = self.test_private_combined(owner_wm_testloader, model = new_combined_model)
            # self.test_label_predictions_combined(owner_wm_testloader, model = new_combined_model)
            print('#################################################################################################\n')

            train_combined_losses.append(train_combined_loss)
            test_private_losses.append(test_private_loss)
            test_private_accs.append(test_private_acc)
            test_original_losses.append(test_original_loss)
            test_original_accs.append(test_original_acc)

            #save the best model
            if test_original_acc > best_acc:
                best_acc = test_original_acc
                state = {
                'epoch': epoch,
                'state_dict': self.attacker_model.state_dict(),
                'optimizer': self.attacker_optimizer.state_dict(),
                'acc': best_acc}
                savepath = './Checkpoints/best_'+ dataset + '.t7'
                torch.save(state,savepath)
        results = {
            'dataset': self.original_dataset,
            'model': self.model_name,
            'test_original_losses': test_original_losses,
            'test_private_losses': test_private_losses,
            'test_original_accs':test_original_accs,
            'test_private_accs': test_private_accs}
        
        savepath = './Results/results_overwrtite' + self.original_dataset + self.model_name + str(data_fraction) + '.t7'
        torch.save(results,savepath)
        
        plot_accuracy_simultaneous(test_original_accs,  test_private_accs, 
        title = 'Overwriting model', filename = 'cifar10_acc_overwrite_' + self.model_name + '_'+ str(data_fraction), save = True)

        
        
