from __future__ import print_function
import argparse
import torch
import torch.utils.data
import matplotlib.pyplot as plt
from torch import nn, optim
from torch.nn import functional as F
from torchvision import datasets, transforms
from torchvision.utils import save_image
from torch.utils.tensorboard import SummaryWriter

from ce_vae_test.networks.min_vae import MinVae
from ce_vae_test.trainer.ce_trainer import CeVaeTrainer
from ce_vae_test.sampler.dataset_sampler import SamplerDatasetWithReplacement


parser = argparse.ArgumentParser(description='VAE MNIST Example')
parser.add_argument('--batch-size', type=int, default=128, metavar='N',
                    help='input batch size for training (default: 128)')
parser.add_argument('--epochs', type=int, default=10, metavar='N',
                    help='number of epochs to train (default: 10)')
parser.add_argument('--no-cuda', action='store_true', default=False,
                    help='enables CUDA training')
parser.add_argument('--seed', type=int, default=1, metavar='S',
                    help='random seed (default: 1)')
parser.add_argument('--log-interval', type=int, default=10, metavar='N',
                    help='how many batches to wait before logging training status')
args = parser.parse_args()
args.cuda = not args.no_cuda and torch.cuda.is_available()

torch.manual_seed(args.seed)

device = torch.device("cuda" if args.cuda else "cpu")

writer = SummaryWriter()

kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {}
train_sampler = SamplerDatasetWithReplacement(
    dataset=datasets.MNIST('../data',
                           train=True,
                           download=True,
                           transform=transforms.ToTensor()),
    batch_size=args.batch_size
)
test_sampler = SamplerDatasetWithReplacement(
    dataset=datasets.MNIST('../data',
                           train=False,
                           transform=transforms.ToTensor()),
    batch_size=args.batch_size * 10
)

cevae = MinVae(
    input_size=28 * 28,
    output_size=10,
    latent_dim=2,
    hidden_sizes_dec=[5],
    device=device
).to(device)

trainer = CeVaeTrainer(
    vae=cevae,
    num_epochs=300,
    train_loader=train_sampler,
    test_loader=test_sampler,
    writer=writer,
    device=device,
    alpha=0.90,
    lamda=0.22
)

trainer.run()


[(17, 'argparse.ArgumentParser', 'argparse.ArgumentParser', False, 'import argparse\n'), (31, 'torch.manual_seed', 'torch.manual_seed', False, 'import torch\n'), (33, 'torch.device', 'torch.device', False, 'import torch\n'), (35, 'torch.utils.tensorboard.SummaryWriter', 'SummaryWriter', False, 'from torch.utils.tensorboard import SummaryWriter\n'), (60, 'ce_vae_test.trainer.ce_trainer.CeVaeTrainer', 'CeVaeTrainer', False, 'from ce_vae_test.trainer.ce_trainer import CeVaeTrainer\n'), (29, 'torch.cuda.is_available', 'torch.cuda.is_available', False, 'import torch\n'), (52, 'ce_vae_test.networks.min_vae.MinVae', 'MinVae', False, 'from ce_vae_test.networks.min_vae import MinVae\n'), (42, 'torchvision.transforms.ToTensor', 'transforms.ToTensor', False, 'from torchvision import datasets, transforms\n'), (48, 'torchvision.transforms.ToTensor', 'transforms.ToTensor', False, 'from torchvision import datasets, transforms\n')]