#
# Copyright (C) 2023, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use 
# under the terms of the LICENSE.md file.
#
# For inquiries contact  george.drettakis@inria.fr
#

import os
import torch
import torch.nn as nn
import numpy as np
from random import randint
from utils.loss_utils import l1_loss, ssim
from gaussian_renderer import render, network_gui
import sys
from scene import Scene, GaussianModel
from utils.general_utils import safe_state
import uuid
from tqdm import tqdm
from utils.image_utils import psnr
from argparse import ArgumentParser, Namespace
from arguments import ModelParams, PipelineParams, OptimizationParams
try:
    from torch.utils.tensorboard import SummaryWriter
    TENSORBOARD_FOUND = True
except ImportError:
    TENSORBOARD_FOUND = False

import copy
import torch.nn.functional as F
import torchvision
from torchvision import transforms
from hidden_api import encoder, decoder, params, str2msg, msg2str, default_transform, NORMALIZE_IMAGENET, UNNORMALIZE_IMAGENET, EncoderWithJND, JND, encoder_with_jnd
from einops import reduce, repeat, rearrange
from utils.general_utils import inverse_sigmoid, get_expon_lr_func, build_rotation
from utils.loss_utils import ssim
from lpipsPyTorch import lpips

normalize_vqgan = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) # Normalize (x - 0.5) / 0.5
unnormalize_vqgan = transforms.Normalize(mean=[-1, -1, -1], std=[1/0.5, 1/0.5, 1/0.5]) # Unnormalize (x * 0.5) + 0.5
normalize_img = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # Normalize (x - mean) / std
unnormalize_img = transforms.Normalize(mean=[-0.485/0.229, -0.456/0.224, -0.406/0.225], std=[1/0.229, 1/0.224, 1/0.225]) # Unnormalize (x * std) + mean

# NORMALIZE_IMAGENET = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
# UNNORMALIZE_IMAGENET = transforms.Normalize(mean=[-0.485/0.229, -0.456/0.224, -0.406/0.225], std=[1/0.229, 1/0.224, 1/0.225])
# default_transform = transforms.Compose([transforms.ToTensor(), NORMALIZE_IMAGENET])

vqgan_to_imnet = transforms.Compose([unnormalize_vqgan, normalize_img])

# create message
random_msg = False
if random_msg:
    msg_ori = torch.randint(0, 2, (1, params.num_bits), device="cuda").bool() # b k
else:
    msg_ori = torch.Tensor(str2msg("111010110101000001010111010011010100010000100111")).unsqueeze(0)
msg_ori = msg_ori.cuda()
msg = 2 * msg_ori.type(torch.float) - 1 # b k

# augmentation
import torch.nn.functional as F

class addNoise(nn.Module):
    """
    Resize the image. The target size is original size * resize_ratio
    """
    def __init__(self, sigma):
        super(addNoise, self).__init__()
        self.sigma = sigma


    def forward(self, noised_and_cover):
        noised_image = noised_and_cover
        noised_and_cover = noised_image + (self.sigma ** 2) * torch.randn_like(noised_image)

        return noised_and_cover


class AddGaussianNoise(object):
    def __init__(self, mean=0., std=1.):
        self.std = std
        self.mean = mean
        
    def __call__(self, tensor):
        return tensor + torch.randn(tensor.size()) * self.std + self.mean
    
    def __repr__(self):
        return self.__class__.__name__ + '(mean={0}, std={1})'.format(self.mean, self.std)

def random_float(min, max):
    """
    Return a random number
    :param min:
    :param max:
    :return:
    """
    return np.random.rand() * (max - min) + min

class Resize(nn.Module):
    """
    Resize the image. The target size is original size * resize_ratio
    """
    def __init__(self, resize_ratio_range, interpolation_method='nearest'):
        super(Resize, self).__init__()
        self.resize_ratio_min = resize_ratio_range[0]
        self.resize_ratio_max = resize_ratio_range[1]
        self.interpolation_method = interpolation_method


    def forward(self, noised_and_cover):

        resize_ratio = random_float(self.resize_ratio_min, self.resize_ratio_max)
        noised_image = noised_and_cover
        noised_and_cover = F.interpolate(
                                    noised_image,
                                    scale_factor=(resize_ratio, resize_ratio),
                                    mode=self.interpolation_method)

        return noised_and_cover

# augmentation
gaussian_noise = addNoise(0.05)
random_rotate = transforms.RandomRotation(30)
resize = Resize([0.75, 1])
gaussian_blur = transforms.GaussianBlur(kernel_size=3, sigma=0.1)
aug_list = [gaussian_noise, random_rotate, resize, gaussian_blur]
aug_dict = {"gaussian_noise": gaussian_noise, "random_rotate": random_rotate, "resize": resize, "gaussian_blur": gaussian_blur}

def training(dataset, opt, pipe, testing_iterations, saving_iterations, checkpoint_iterations, checkpoint, debug_from):
    first_iter = 0
    tb_writer = prepare_output_and_logger(dataset)
    gaussians = GaussianModel(dataset.sh_degree)
    scene = Scene(dataset, gaussians)
    gaussians.training_setup(opt)
    if checkpoint:
        (model_params, first_iter) = torch.load(checkpoint)
        gaussians.restore(model_params, opt)

    gaussians_origin = copy.deepcopy(gaussians)

    bg_color = [1, 1, 1] if dataset.white_background else [0, 0, 0]
    background = torch.tensor(bg_color, dtype=torch.float32, device="cuda")

    # uncertainty
    filter_out_grad = ["rotation"]
    name2idx = {"xyz": 0, "rgb": 1, "sh": 2, "scale": 3, "rotation": 4, "opacity": 5}
    filter_out_idx = [name2idx[k] for k in filter_out_grad]

    viewpoint_cams = scene.getTrainCameras().copy()

    gaussians_params = gaussians.capture()[1:7]
    gaussians_params = [p for i, p in enumerate(gaussians_params) if i not in filter_out_idx]

    # off load to cpu to avoid oom with greedy algo
    # device = gaussians_params[0].device if num_views == 1 else "cpu"
    device = "cpu" # we have to load to cpu because of inflation

    # H_train = torch.zeros(sum(p.numel() for p in gaussians_params), device=gaussians_params[0].device, dtype=gaussians_params[0].dtype)
    H_train = torch.zeros(gaussians_params[0].shape[0], device=gaussians_params[0].device, dtype=gaussians_params[0].dtype)

    # Run heesian on training set
    for cam in tqdm(viewpoint_cams, desc="Calculating diagonal Hessian on training views"):
        # if exit_func():
        #     raise RuntimeError("csm should exit early")

        render_pkg = render(cam, gaussians, pipe, background)
        pred_img = render_pkg["render"]
        pred_img.backward()

        # cur_H = torch.cat([p.grad.detach().reshape(-1) for p in gaussians_params])
        # H_train += cur_H

        H_train += sum([reduce(p.grad.detach(), "n ... -> n", "sum") for p in gaussians_params])

        gaussians.optimizer.zero_grad(set_to_none = True) 

    H_train = H_train.to(device)

    H_mask = torch.abs(H_train) < torch.abs(H_train).mean()

    # densify points clone
    selected_pts_mask = ~H_mask
    new_xyz = gaussians._xyz[selected_pts_mask]
    new_features_dc = gaussians._features_dc[selected_pts_mask]
    new_features_rest = gaussians._features_rest[selected_pts_mask]
    new_opacities = gaussians._opacity[selected_pts_mask]
    new_scaling = gaussians._scaling[selected_pts_mask]
    new_rotation = gaussians._rotation[selected_pts_mask]
    gaussians.densification_postfix(new_xyz, new_features_dc, new_features_rest, new_opacities, new_scaling, new_rotation)

    # densify points split
    # selected_pts_mask = H_mask_0
    # N = 2
    # stds = gaussians.get_scaling[selected_pts_mask].repeat(N, 1)
    # means = torch.zeros((stds.size(0), 3),device="cuda")
    # samples = torch.normal(mean=means, std=stds)
    # rots = build_rotation(gaussians._rotation[selected_pts_mask]).repeat(N, 1, 1)
    # new_xyz = torch.bmm(rots, samples.unsqueeze(-1)).squeeze(-1) + gaussians.get_xyz[selected_pts_mask].repeat(N, 1)
    # new_scaling = gaussians.scaling_inverse_activation(gaussians.get_scaling[selected_pts_mask].repeat(N,1) / (0.8*N))
    # new_rotation = gaussians._rotation[selected_pts_mask].repeat(N,1)
    # new_features_dc = gaussians._features_dc[selected_pts_mask].repeat(N,1,1)
    # new_features_rest = gaussians._features_rest[selected_pts_mask].repeat(N,1,1)
    # new_opacity = gaussians._opacity[selected_pts_mask].repeat(N,1)

    # gaussians.densification_postfix(new_xyz, new_features_dc, new_features_rest, new_opacity, new_scaling, new_rotation)

    # prune_filter = torch.cat((selected_pts_mask, torch.zeros(N * selected_pts_mask.sum(), device="cuda", dtype=bool)))
    # gaussians.prune_points(prune_filter)

    import matplotlib.pyplot as plt

    # Sort H_train in descending order
    sorted_H_train, _ = torch.sort(torch.abs(H_train), descending=True)

    # Plot the sorted H_train
    plt.plot(sorted_H_train)
    plt.xlabel('Parameter Index')
    plt.ylabel('Hessian Value')
    plt.title('Descending Order of Hessian Values')
    plt.savefig('grad.png')
    plt.close()

    iter_start = torch.cuda.Event(enable_timing = True)
    iter_end = torch.cuda.Event(enable_timing = True)

    viewpoint_stack = None
    ema_loss_for_log = 0.0
    progress_bar = tqdm(range(first_iter, opt.iterations), desc="Training progress")
    first_iter += 1
    for iteration in range(first_iter, opt.iterations + 1):        
        if network_gui.conn == None:
            network_gui.try_connect()
        while network_gui.conn != None:
            try:
                net_image_bytes = None
                custom_cam, do_training, pipe.convert_SHs_python, pipe.compute_cov3D_python, keep_alive, scaling_modifer = network_gui.receive()
                if custom_cam != None:
                    net_image = render(custom_cam, gaussians, pipe, background, scaling_modifer)["render"]
                    net_image_bytes = memoryview((torch.clamp(net_image, min=0, max=1.0) * 255).byte().permute(1, 2, 0).contiguous().cpu().numpy())
                network_gui.send(net_image_bytes, dataset.source_path)
                if do_training and ((iteration < int(opt.iterations)) or not keep_alive):
                    break
            except Exception as e:
                network_gui.conn = None

        iter_start.record()

        gaussians.update_learning_rate(iteration)

        # Every 1000 its we increase the levels of SH up to a maximum degree
        if iteration % 1000 == 0:
            gaussians.oneupSHdegree()

        # Pick a random Camera
        if not viewpoint_stack:
            viewpoint_stack = scene.getTrainCameras().copy()
        viewpoint_cam = viewpoint_stack.pop(randint(0, len(viewpoint_stack)-1))

        # Render
        if (iteration - 1) == debug_from:
            pipe.debug = True

        bg = torch.rand((3), device="cuda") if opt.random_background else background

        render_pkg = render(viewpoint_cam, gaussians, pipe, bg)
        image, viewspace_point_tensor, visibility_filter, radii = render_pkg["render"], render_pkg["viewspace_points"], render_pkg["visibility_filter"], render_pkg["radii"]

        # Loss
        gt_image = viewpoint_cam.original_image.cuda()
        Ll1 = l1_loss(image, gt_image)
        loss = (1.0 - opt.lambda_dssim) * Ll1 + opt.lambda_dssim * (1.0 - ssim(image, gt_image))

        # torchvision.utils.save_image(image, os.path.join("results", '{0:05d}'.format(iteration) + ".png"))

        # wm_loss

        # # test hidden with offline image
        # from PIL import Image
        # img = Image.open("results/00001.png").convert('RGB')
        # # img = img.resize((512, 512), Image.BICUBIC)
        # img_pt = default_transform(img).unsqueeze(0).cuda()
        # img_w = encoder_with_jnd(img_pt, msg)
        # clip_img = torch.clamp(UNNORMALIZE_IMAGENET(img_w), 0, 1)
        # clip_img = torch.round(255 * clip_img)/255 
        # clip_img = transforms.ToPILImage()(clip_img.squeeze(0).cpu())
        
        # ft = decoder(default_transform(clip_img).unsqueeze(0).cuda())
        # decoded_msg = ft > 0 # b k -> b k
        # accs = (~torch.logical_xor(decoded_msg, msg_ori)) # b k -> b k
        # print(f"Message: {msg2str(msg_ori.squeeze(0).cpu().numpy())}")
        # print(f"Decoded: {msg2str(decoded_msg.squeeze(0).cpu().numpy())}")
        # print(f"Bit Accuracy: {accs.sum().item() / params.num_bits}")

        # # test hidden with tensor image
        # image_wm = NORMALIZE_IMAGENET(image.unsqueeze(0))
        # image_wm_enc = encoder_with_jnd(image_wm, msg)

        # ft = decoder(image_wm_enc)
        # decoded_msg = ft > 0 # b k -> b k
        # accs = (~torch.logical_xor(decoded_msg, msg_ori)) # b k -> b k

        # print(f"Message: {msg2str(msg_ori.squeeze(0).cpu().numpy())}")
        # print(f"Decoded: {msg2str(decoded_msg.squeeze(0).cpu().numpy())}")
        # print(f"Bit Accuracy: {accs.sum().item() / params.num_bits}")

        # decode
        nom_image = NORMALIZE_IMAGENET(image.unsqueeze(0))
        ft = decoder(nom_image)

        # Binary Cross Entropy Loss
        # bce_loss = torch.nn.BCELoss()
        bce_loss = F.binary_cross_entropy_with_logits(ft, msg_ori)
        # print(f"BCE Loss: {bce_loss}")
        loss = 10 * loss + bce_loss

        decoded_msg = ft > 0 # b k -> b k
        accs = (~torch.logical_xor(decoded_msg, msg_ori)) # b k -> b k
        # print(f"Message: {msg2str(msg_ori.squeeze(0).cpu().numpy())}")
        # print(f"Decoded: {msg2str(decoded_msg.squeeze(0).cpu().numpy())}")
        # print(f"Bit Accuracy: {accs.sum().item() / params.num_bits}")

        # scaling_diff = gaussians.get_scaling - gaussians_origin.get_scaling
        # rotation_diff = gaussians.get_rotation - gaussians_origin.get_rotation
        # xyz_diff = gaussians.get_xyz - gaussians_origin.get_xyz
        # print("scaling: ", float(scaling_diff.min()), float(scaling_diff.mean()), float(scaling_diff.max()))
        # print("rotation: ", float(rotation_diff.min()), float(rotation_diff.mean()), float(rotation_diff.max()))
        # print("xyz: ", float(xyz_diff.min()), float(xyz_diff.mean()), float(xyz_diff.max()))

        # print("spatial_lr_scale: ", gaussians.spatial_lr_scale)

        # Print current learning rate for all parameters
        # for name, param_group in gaussians.optimizer.param_groups:
        #     print("Parameter:", name)
        #     print("Learning Rate:", param_group['lr'])

        loss.backward()

        iter_end.record()

        with torch.no_grad():
            # Progress bar
            ema_loss_for_log = 0.4 * loss.item() + 0.6 * ema_loss_for_log
            if iteration % 10 == 0:
                progress_bar.set_postfix({"Loss": f"{ema_loss_for_log:.{7}f}", "xyz": f"{gaussians.get_xyz.shape}", "BCE Loss": f"{bce_loss:.{3}f}", "Bit Accuracy": f"{accs.sum().item() / params.num_bits:.{3}f}"})
                progress_bar.update(10)
                with torch.no_grad():
                    image_origin = render(viewpoint_cam, gaussians_origin, pipe, bg)["render"]

                torchvision.utils.save_image(torch.cat((image, image-image_origin), dim=2), os.path.join("results_hidden", '{0:05d}'.format(iteration) + ".png"))

            if iteration == opt.iterations:
                progress_bar.close()

            # Log and save
            training_report(tb_writer, iteration, Ll1, loss, l1_loss, iter_start.elapsed_time(iter_end), testing_iterations, scene, render, (pipe, background))
            if (iteration in saving_iterations):
                print("\n[ITER {}] Saving Gaussians".format(iteration))
                scene.save(iteration)
                gaussians.save_ply_partial('/workspace/code/gaussian-splatting_wm/perturb_plys/point_cloud.ply', origin_points)

            # Densification
            # if iteration < opt.densify_until_iter:
            #     # Keep track of max radii in image-space for pruning
            #     gaussians.max_radii2D[visibility_filter] = torch.max(gaussians.max_radii2D[visibility_filter], radii[visibility_filter])
            #     gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter)

            #     if iteration > opt.densify_from_iter and iteration % opt.densification_interval == 0:

            #         with torch.enable_grad():
            #             # uncertainty
            #             filter_out_grad = ["rotation"]
            #             name2idx = {"xyz": 0, "rgb": 1, "sh": 2, "scale": 3, "rotation": 4, "opacity": 5}
            #             filter_out_idx = [name2idx[k] for k in filter_out_grad]

            #             viewpoint_cams = scene.getTrainCameras().copy()

            #             gaussians_params = gaussians.capture()[1:7]
            #             gaussians_params = [p for i, p in enumerate(gaussians_params) if i not in filter_out_idx]

            #             # off load to cpu to avoid oom with greedy algo
            #             # device = gaussians_params[0].device if num_views == 1 else "cpu"
            #             device = "cpu" # we have to load to cpu because of inflation

            #             # H_train = torch.zeros(sum(p.numel() for p in gaussians_params), device=gaussians_params[0].device, dtype=gaussians_params[0].dtype)
            #             H_train = torch.zeros(gaussians_params[0].shape[0], device=gaussians_params[0].device, dtype=gaussians_params[0].dtype)

            #             # Run heesian on training set
            #             for cam in tqdm(viewpoint_cams, desc="Calculating diagonal Hessian on training views"):
            #                 # if exit_func():
            #                 #     raise RuntimeError("csm should exit early")

            #                 render_pkg = render(cam, gaussians, pipe, background)
            #                 pred_img = render_pkg["render"]
            #                 pred_img.backward(gradient=torch.ones_like(pred_img))

            #                 # cur_H = torch.cat([p.grad.detach().reshape(-1) for p in gaussians_params])
            #                 # H_train += cur_H

            #                 H_train += sum([reduce(p.grad.detach(), "n ... -> n", "sum") for p in gaussians_params])

            #                 gaussians.optimizer.zero_grad(set_to_none = True) 

            #             H_train = H_train.to(device)

            #         size_threshold = 20 if iteration > opt.opacity_reset_interval else None
            #         gaussians.densify_and_prune(opt.densify_grad_threshold, 0.005, scene.cameras_extent, size_threshold, H_train)
                
            #     if iteration % opt.opacity_reset_interval == 0 or (dataset.white_background and iteration == opt.densify_from_iter):
            #         gaussians.reset_opacity()

            # Optimizer step
            if iteration < opt.iterations:
                gaussians.optimizer.step()
                gaussians.optimizer.zero_grad(set_to_none = True)

            if (iteration in checkpoint_iterations):
                print("\n[ITER {}] Saving Checkpoint".format(iteration))
                torch.save((gaussians.capture(), iteration), scene.model_path + "/chkpnt" + str(iteration) + ".pth")

def prepare_output_and_logger(args):    
    if not args.model_path:
        if os.getenv('OAR_JOB_ID'):
            unique_str=os.getenv('OAR_JOB_ID')
        else:
            unique_str = str(uuid.uuid4())
        args.model_path = os.path.join("./output/", unique_str[0:10])
        
    # Set up output folder
    print("Output folder: {}".format(args.model_path))
    os.makedirs(args.model_path, exist_ok = True)
    with open(os.path.join(args.model_path, "cfg_args"), 'w') as cfg_log_f:
        cfg_log_f.write(str(Namespace(**vars(args))))

    # Create Tensorboard writer
    tb_writer = None
    if TENSORBOARD_FOUND:
        tb_writer = SummaryWriter(args.model_path)
    else:
        print("Tensorboard not available: not logging progress")
    return tb_writer

def training_report(tb_writer, iteration, Ll1, loss, l1_loss, elapsed, testing_iterations, scene : Scene, renderFunc, renderArgs):
    if tb_writer:
        tb_writer.add_scalar('train_loss_patches/l1_loss', Ll1.item(), iteration)
        tb_writer.add_scalar('train_loss_patches/total_loss', loss.item(), iteration)
        tb_writer.add_scalar('iter_time', elapsed, iteration)

    # Report test and samples of training set
    if iteration in testing_iterations:
        torch.cuda.empty_cache()
        validation_configs = ({'name': 'test', 'cameras' : scene.getTestCameras()}, 
                              {'name': 'train', 'cameras' : [scene.getTrainCameras()[idx % len(scene.getTrainCameras())] for idx in range(5, 30, 5)]})

        for config in validation_configs:
            if config['cameras'] and len(config['cameras']) > 0:
                l1_test = 0.0
                psnr_test = 0.0
                ssim_test = 0.0
                lpips_test = 0.0
                for idx, viewpoint in enumerate(config['cameras']):
                    image = torch.clamp(renderFunc(viewpoint, scene.gaussians, *renderArgs)["render"], 0.0, 1.0)
                    gt_image = torch.clamp(viewpoint.original_image.to("cuda"), 0.0, 1.0)
                    
                    # test hidden
                    # test hidden with tensor image
                    image_wm = NORMALIZE_IMAGENET(image.unsqueeze(0))
                    # image_wm_enc = encoder_with_jnd(image_wm, msg)

                    ft = decoder(image_wm)
                    decoded_msg = ft > 0 # b k -> b k
                    accs = (~torch.logical_xor(decoded_msg, msg_ori)) # b k -> b k

                    print(f"Message: {msg2str(msg_ori.squeeze(0).cpu().numpy())}")
                    print(f"Decoded: {msg2str(decoded_msg.squeeze(0).cpu().numpy())}")
                    print(f"Bit Accuracy: {accs.sum().item() / params.num_bits}")

                    # augmentations
                    for k, aug in aug_dict.items():
                        image_wm_aug = aug(image_wm)
                        ft = decoder(image_wm_aug)
                        decoded_msg_aug = ft > 0 # b k -> b k
                        accs_aug = (~torch.logical_xor(decoded_msg_aug, msg_ori)) # b k -> b k
                        # print(f"Message: {msg2str(msg_ori.squeeze(0).cpu().numpy())}")
                        # print(f"Decoded: {msg2str(decoded_msg_aug.squeeze(0).cpu().numpy())}")
                        print(f"Bit Accuracy Aug {k}: {accs_aug.sum().item() / params.num_bits}")

                    if tb_writer and (idx < 5):
                        tb_writer.add_images(config['name'] + "_view_{}/render".format(viewpoint.image_name), image[None], global_step=iteration)
                        if iteration == testing_iterations[0]:
                            tb_writer.add_images(config['name'] + "_view_{}/ground_truth".format(viewpoint.image_name), gt_image[None], global_step=iteration)
                    l1_test += l1_loss(image, gt_image).mean().double()
                    psnr_test += psnr(image, gt_image).mean().double()
                    ssim_test += ssim(image, gt_image).mean().double()
                    lpips_test += lpips(image, gt_image).mean().double()
                psnr_test /= len(config['cameras'])
                ssim_test /= len(config['cameras'])
                lpips_test /= len(config['cameras'])
                l1_test /= len(config['cameras'])          
                print("\n[ITER {}] Evaluating {}: L1 {} PSNR {} SSIM {} LPIPS {}".format(iteration, config['name'], l1_test, psnr_test, ssim_test, lpips_test))
                if tb_writer:
                    tb_writer.add_scalar(config['name'] + '/loss_viewpoint - l1_loss', l1_test, iteration)
                    tb_writer.add_scalar(config['name'] + '/loss_viewpoint - psnr', psnr_test, iteration)

        if tb_writer:
            tb_writer.add_histogram("scene/opacity_histogram", scene.gaussians.get_opacity, iteration)
            tb_writer.add_scalar('total_points', scene.gaussians.get_xyz.shape[0], iteration)
        torch.cuda.empty_cache()

if __name__ == "__main__":
    # Set up command line argument parser
    parser = ArgumentParser(description="Training script parameters")
    lp = ModelParams(parser)
    op = OptimizationParams(parser)
    pp = PipelineParams(parser)
    parser.add_argument('--ip', type=str, default="127.0.0.1")
    parser.add_argument('--port', type=int, default=6009)
    parser.add_argument('--debug_from', type=int, default=-1)
    parser.add_argument('--detect_anomaly', action='store_true', default=False)
    parser.add_argument("--test_iterations", nargs="+", type=int, default=[200, 400, 600, 800, 1000, 1200, 1400, 1600, 1800, 2000])
    parser.add_argument("--save_iterations", nargs="+", type=int, default=[1000, 2000])
    parser.add_argument("--quiet", action="store_true")
    parser.add_argument("--checkpoint_iterations", nargs="+", type=int, default=[])
    parser.add_argument("--start_checkpoint", type=str, default = None)
    args = parser.parse_args(sys.argv[1:])
    args.save_iterations.append(args.iterations)
    
    print("Optimizing " + args.model_path)

    # Initialize system state (RNG)
    safe_state(args.quiet)

    # Start GUI server, configure and run training
    network_gui.init(args.ip, args.port)
    torch.autograd.set_detect_anomaly(args.detect_anomaly)
    training(lp.extract(args), op.extract(args), pp.extract(args), args.test_iterations, args.save_iterations, args.checkpoint_iterations, args.start_checkpoint, args.debug_from)

    # All done
    print("\nTraining complete.")
