import os
import numpy as np
import math

import torch

from torchvision import datasets
from torch.utils.data import DataLoader
from torchvision.utils import save_image


__all__ = ['saveimages']


def savefile(label, data, shape, number, out_dir, n_row=30):
    save = np.array([], np.float)
    save = np.reshape(save, (0, shape[0], shape[1], shape[2]))
    save = np.insert(save, 0, data, 0)
    print("save label ", label, save.shape)
    save = torch.from_numpy(save)
    fname = '%s/%d_%d.png' % (out_dir, label, number)
    save_image(save, fname, nrow=n_row, normalize=True)


def saveimages(dataloaders, n_classes, out_dir, file_qty, n_row=30):

    os.makedirs(out_dir, exist_ok=True)

    datas = [[] for _ in range(n_classes)]

    number = 0
    for dataloader in dataloaders:
        for image, label in dataloader:
            image = image.numpy()
            image = np.reshape(image, (image.shape[1], image.shape[2], image.shape[3]))
            data = datas[label]
            data.append(image)
            if file_qty <= len(data):
                savefile(label, data, image.shape, number, out_dir, n_row)
                number += 1
                data.clear()

    for label, data in enumerate(datas):
        if 1 <= len(data):
            savefile(label, data, image.shape, number, out_dir, n_row)
            number += 1
            data.clear()

import argparse
import datetime as dt
import os
import numpy as np
import math

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable

import torchvision.transforms as transforms
from torchvision import datasets
from torch.utils.data import DataLoader

from deepmodel.commonpy import filename as cmfn
from deepmodel.commonpy import imageio as imio
import deepmodel.torchmodel as tmodel


def parseargs(datetime):
    parser = argparse.ArgumentParser()
    parser.add_argument('--image_size', type=int, default=32, help='Size of each image dimension')
    parser.add_argument('--channels', type=int, default=3, help='Number of image channels')
    parser.add_argument('--n_classes', type=int, default=10, help='Number of classes for dataset')
    parser.add_argument('--image-quantity', type=int, default=10, help='Number of image per file')
    parser.add_argument('--data-dir', type=str, default=cmfn.datadir(__file__), help='Data cache directory')
    parser.add_argument('--out-dir', type=str, default=cmfn.outputdir(__file__, datetime), help='Out directory for image samples')
    parser.add_argument('--n_row', type=int, default=50, help='Number of rows of a image')
    parser.add_argument('--file_qty', type=int, default=1000, help='Number of images per a file')
    return parser.parse_args()


def main():
    now = dt.datetime.now()
    opt = parseargs(now)
    devtype = tmodel.DeviceType('cpu')

    trainloader = tmodel.cifar10_train_loader(devtype, opt.data_dir, opt.image_size, 1, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), shuffle=False)
    testloader = tmodel.cifar10_test_loader(devtype, opt.data_dir, opt.image_size, 1, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), shuffle=False)

    imio.saveimages([trainloader, testloader], opt.n_classes, opt.out_dir, opt.file_qty, n_row=opt.n_row)


if __name__ == "__main__":
    main()