from pygan._torch.ebgan_image_generator import EBGANImageGenerator
from pathlib import Path
import pickle
import numpy as np
import os
from PIL import Image
import torch

def _pathmaker(i):
    path = Path("dataset_for_pygan")
    os.makedirs(path, exist_ok=True)
    return Path(path, f"{str(i)}.bmp")

def gan(dataset):
    """Generate image using pygan.

    :param dataset: Image dataset used to train the gan.
    :type dataset: fetch_data.utk.StorageDataset
    :return: A dataset with generated images.
    :rtype : fetch_data.utk.StorageDataset
    """

    #Saving 50x50 images from the train dataset
    for i in range(len(dataset)):
        x = dataset[i][0]
        x = x*255
        x = np.moveaxis(x,0,2).astype(int)
        image = Image.fromarray(x.astype('uint8'), 'RGB')
        image.save(_pathmaker(i), "bmp")
        

    ctx = "cuda:0" if torch.cuda.is_available() else "cpu"

    gan_image_generator = EBGANImageGenerator(
        # `list` of path to your directories.
        dir_list=["dataset_for_pygan"],
        # `int` of image width.
        width=50,
        # `int` of image height.
        height=50,
        # `int` of image channel.
        channel=3,
        # `int` of batch size.
        batch_size=40,
        # `float` of learning rate.
        learning_rate=1e-06,
        ctx=ctx
    )

    gan_image_generator.learn(
        # `int` of the number of training iterations.
        iter_n=1000,
        # `int` of the number of learning of the discriminative model.
        k_step=10,
    )

    
    path = Path("synthetic_data")
    os.makedirs(path,exist_ok=True)
    
    arr = gan_image_generator.EBGAN.generative_model.draw()
    arr = arr.detach().numpy()
    
    #arr shape is:
    #batch
    #channel
    #height
    #width

    #save generate images 
    for i in range(np.shape(arr)[0]):
        image_path = Path(path,f"{str(i)}.bmp")
        img = np.moveaxis(arr[i],0,2)
        image = Image.fromarray(img.astype('uint8'), 'RGB')
        image.save(image_path, "bmp")
        




