//  [python]
def main(argv, seed=88):
    args = parse()
    n_plt = 5
    gen_pars = [args.lengthscale, args.period]
    torch.random.manual_seed(seed)

    ## generate data ##
    gen_lh = gpytorch.likelihoods.GaussianLikelihood()
    gen_mod = ExactGPModel(None, None, gen_lh)
    gen_mod.covar_module.base_kernel._set_lengthscale(gen_pars[0])
    # gen_mod.covar_module = gpytorch.kernels.ProductKernel(gpytorch.kernels.RBFKernel(), gpytorch.kernels.PeriodicKernel())
    # gen_mod.covar_module.kernels[0]._set_lengthscale(gen_pars[0])
    # gen_mod.covar_module.kernels[1]._set_period_length(gen_pars[1])
    gen_lh.eval()
    gen_mod.eval()

    train_x = torch.linspace(0, 5, args.nx)
    train_y = gen_mod(train_x).sample(sample_shape=torch.Size((1,))).squeeze().detach()

    ## set up actual model ##
    lh = gpytorch.likelihoods.GaussianLikelihood()
    model = ExactGPModel(train_x, train_y, lh)

    lh.train()
    model.train()

    samples = model(train_x).sample(sample_shape=torch.Size((n_plt, ))).t().detach()
    lower, upper = lh(model(train_x)).confidence_region()

    sns.set_style("whitegrid")
    sns.set_context('talk')
    sns.despine()
    sns.set_palette("muted")
    plt.figure(figsize=(10,9))
    plt.plot(train_x.numpy(), samples.numpy())
    plt.plot(train_x.numpy(), train_y.numpy(), linewidth=2.5, linestyle="None",
             marker=".", label="Data")
    # plt.fill_between(train_x.numpy(), lower.detach().numpy(), upper.detach().numpy(),
    #                  color="steelblue", alpha=0.1, label=r'$\pm 2$ SD')
    plt.xlim((0, 5))
    plt.ylim((-2,2))
    plt.xlabel("X", fontsize=28)
    plt.ylabel("Y", fontsize=28)
    plt.title("Prior Draws", fontsize=28)
    plt.legend()
    plt.show()


    ## now do training ##
    # Find optimal model hyperparameters
    model.train()
    lh.train()

    # Use the adam optimizer
    optimizer = torch.optim.Adam([
        {'params': model.parameters()},  # Includes GaussianLikelihood parameters
    ], lr=0.1)

    # "Loss" for GPs - the marginal log likelihood
    mll = gpytorch.mlls.ExactMarginalLogLikelihood(lh, model)

    training_iter = 40
    for i in range(training_iter):
        # Zero gradients from previous iteration
        optimizer.zero_grad()
        # Output from model
        output = model(train_x)
        # Calc loss and backprop gradients
        loss = -mll(output, train_y)
        loss.backward()
        # print('Iter %d/%d - Loss: %.3f   lengthscale: %.3f   noise: %.3f' % (
        #     i + 1, training_iter, loss.item(),
        #     model.covar_module.base_kernel.lengthscale.item(),
        #     model.likelihood.noise.item()
        # ))
        optimizer.step()

    lh.eval()
    model.eval()

    preds = model(train_x).sample(sample_shape=torch.Size((n_plt, ))).t().detach()
    lower, upper = model(train_x).confidence_region()

    plt.figure(figsize=(10, 9))
    plt.plot(train_x.numpy(), preds.numpy())
    plt.plot(train_x.numpy(), train_y.numpy(), linewidth=2.5, linestyle="None",
             marker=".", label="Data")
    plt.fill_between(train_x.numpy(), lower.detach().numpy(), upper.detach().numpy(),
                     color="steelblue", alpha=0.1, label=r'$\pm 2$ SD')
    plt.xlim((0, 5))
    plt.ylim((-0.75, 0.5))
    plt.xlabel("X", fontsize=28)
    plt.ylabel("Y", fontsize=28)
    plt.title("Posterior Draws", fontsize=28)
    plt.legend()
    plt.show()