//  [python]
def invert():
    invert_result_dir = os.path.join(result_dir, '%s-%s-var%s'%(args.model, args.dataset, args.var))
    os.makedirs(invert_result_dir, exist_ok=True)
    # invert_result_dir = result_dir
    inversion_log_file = os.path.join(invert_result_dir, 'log.log')
    bound_log_file = os.path.join(invert_result_dir, 'bound.log')
    log = open(inversion_log_file, 'w+')
    bound_log = open(bound_log_file, 'w+')
    log.write(f'{args}\n')
    log.flush()
    bound_log.write(f'{args}\n')
    bound_log.flush()

    model.to(**setup)
    model.eval()
    # select samples from the index list
    print('Start inverting')
    seeds = [1, 123, 12345, 234, 345]
    for i in range(args.num_inversion_batches):
        # log_dir = os.path.join(invert_result_dir, 'logs/')
        eigenvalue_dir = os.path.join(invert_result_dir, 'eigenvalues/')
        os.makedirs(eigenvalue_dir, exist_ok=True)

        ground_truth, labels = [], []
        img, label = test_set[i]
        labels.append(torch.as_tensor((label,), device=setup['device']))
        ground_truth.append(img.to(**setup))
        mixed_img = torch.stack(ground_truth)
        mixed_labels = torch.cat(labels)

        J = utils.get_Jacob(model, mixed_img, mixed_labels, setup)
        svd_time = time.time()
        U, S, Vh = torch.svd(J.transpose(0,1).detach())
        svd_time = time.time() - svd_time
        print(f"Computing SVD for {svd_time}s")
        torch.save({
            'img': mixed_img.detach().cpu(),
            'label': mixed_labels.detach().cpu(),
            'init': args.init,
            'U': U.detach().cpu(),
            'S': S.detach().cpu(),
            'Vh': Vh.detach().cpu(),
        }, os.path.join(eigenvalue_dir, f'batch{i}.pt'))
        mean_lambda_inv = (1/S).mean()

        model.zero_grad()
        loss = criterion(model(mixed_img), mixed_labels)
        input_gradient = torch.autograd.grad(loss, model.parameters())
        input_gradient = [grad.detach().clone() for grad in input_gradient]

        start_time = time.time()
        '''
            start batch recover and defense:
        '''
        for inv_idx in range(5):
            images_dir = os.path.join(invert_result_dir, 'images/', f"batch{i}", f"inv{inv_idx}")
            img_pure_dir = os.path.join(invert_result_dir, 'figures/', f"batch{i}", f"inv{inv_idx}")
            os.makedirs(images_dir, exist_ok=True)
            os.makedirs(img_pure_dir, exist_ok=True)
            if args.normalize:
                ground_truth_denormalized = torch.clamp(mixed_img * ds + dm, 0, 1)
            else:
                ground_truth_denormalized = mixed_img
            torchvision.utils.save_image(ground_truth_denormalized, os.path.join(images_dir, 'img%d_gt.png'%(i)), nrow=5)

            set_random_seed(seeds[inv_idx])
            output, inv_metrics, step_info, norm_ratio = \
                batch_recover(model, mixed_img, mixed_labels, invert_result_dir, 'log_dir', images_dir, img_pure_dir, input_gradient, max_eigen_value=S[0], mean_lambda_inv=mean_lambda_inv)

            end_time = (time.time() - start_time) / 60 # minute
            if args.normalize:
                output_denormalized = torch.clamp(output * ds + dm, 0, 1)
            else:
                output_denormalized = torch.clamp(output, 0, 1)
            torchvision.utils.save_image(output_denormalized, os.path.join(images_dir, 'img%d_output.png'%(i)), nrow=5)
            
            recover_info = f"Batch: {i} | Idx: {inv_idx} | MeanInvLambda: {mean_lambda_inv.item()} | Rec. loss: {inv_metrics['RecLoss']:2.4f} | MSE: {inv_metrics['MSE']:2.4f} | SSIM: {inv_metrics['SSIM']:2.4f}\n"
            print(recover_info)
            log.write(recover_info)
            bound_log.write(f"Batch: {i} | Idx: {inv_idx} | " + step_info)
            log.flush()
            bound_log.flush()

    log.flush()
    log.close()