import os, glob, sys
import pickle
import streamlit as st
import torch
import matplotlib.pyplot as plt
import numpy as np
sys.path.append('stylegan3')

class SampleFromGAN:
    def __init__(self, G, z_shp, in_gpu=False) -> None:
        self.G = G
        self.in_gpu = in_gpu
        self.z_shp = z_shp #[#images, z_dim]
    def __call__(self,):
        z = torch.randn(self.z_shp)
        if self.in_gpu:
            z = z.cuda()
        ims = self.G(z, c=None)
        ims = ims[:,0,...]
        return ims

class Plot:
    def __init__(self, im_gen) -> None:
        self.im_gen = im_gen
        assert callable(im_gen)
    def __call__(self):
        ims = self.im_gen()
        # plot first image
        im = ims[0,...]
        fig, ax = plt.subplots(1, figsize=(12,12))
        fig.subplots_adjust(left=0,right=1,bottom=0,top=1)
        ax.imshow(im, cmap='gray')
        ax.axis('tight')
        ax.axis('off')
        st.pyplot(fig)

def load_default_gen(in_gpu=False, fname_pkl=None):
    if fname_pkl is None:
        path_ckpt = "./model_weights"
        fname_pkl = os.path.join(path_ckpt, 'network-snapshot-005000.pkl')
        if not os.path.isfile(fname_pkl):
            raise AssertionError(f'Could not find the default network snapshot at {fname_pkl}. Quitting.')
    with open(fname_pkl, 'rb') as f:
        G = pickle.load(f)['G_ema']  # torch.nn.ModuleDict
    if in_gpu:
        G = G.cuda()
    return G