import streamlit as st
import tensorflow as tf
import pickle
import numpy as np
from pathlib import Path
import dnnlib
from dnnlib import tflib
import imageio
import os
import subprocess

def check_gpu():
    return tf.test.is_gpu_available(cuda_only=False, min_cuda_compute_capability=None)

model_path = 'best_net.pkl'
#define load model functions
_cached_networks = dict()
def load_networks(path):
    if path in _cached_networks:
        return _cached_networks[path]
    stream = open(path, 'rb')
    tflib.init_tf()
    with stream:
        G, D, Gs = pickle.load(stream, encoding='latin1')
    _cached_networks[path] = G, D, Gs
    return G, D, Gs

# Code to load the StyleGAN2 Model
def load_model():
    _G, _D, Gs = load_networks(model_path)
    noise_vars = [var for name, var in Gs.components.synthesis.vars.items() if name.startswith('noise')]
    Gs_kwargs = dnnlib.EasyDict()
    Gs_kwargs.output_transform = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)
    Gs_kwargs.randomize_noise = False
    return Gs, noise_vars, Gs_kwargs
#define helper functions
def get_control_latent_vectors(path):
    files = [x for x in Path(path).iterdir() if str(x).endswith('.npy')]
    latent_vectors = {f.name[:-4]:np.load(f) for f in files}
    return latent_vectors

#load latent directions
latent_controls = get_control_latent_vectors('trajectories/')

def generate_image_from_projected_latents(latent_vector):
    images = Gs.components.synthesis.run(latent_vector, **Gs_kwargs)
    return images


def frame_to_frame(latent_code):
    modified_latent_code = np.copy(latent_code)
    full_video = [generate_image_from_projected_latents(modified_latent_code)]
    for i in range(49):
        modified_latent_code = modified_latent_code +  latent_controls[f'{i}{i+1}']
        ims = generate_image_from_projected_latents(modified_latent_code)
        full_video.append(ims)
    return np.array(full_video).squeeze()

#load the model
Gs, noise_vars, Gs_kwargs = load_model()
#select a random latent code
rnd = np.random.RandomState(3)
z = rnd.randn(1, *Gs.input_shape[1:])
noise_vars = [var for name, var in Gs.components.synthesis.vars.items() if name.startswith('noise')]
tflib.set_vars({var: rnd.randn(*var.shape.as_list()) for var in noise_vars})
random_img_latent_code = Gs.components.mapping.run(z,None)

#make it be ED frame
random_img_latent_code -= 0.7*latent_controls['time']
vid = frame_to_frame(random_img_latent_code)


temp_video_path="output.mp4"
writer=imageio.get_writer(temp_video_path, fps=20)
for i in range(vid.shape[0]):
    frame = vid[i]
    writer.append_data(frame)

writer.close()
out_path = "fixed_out.mp4"
command = ["ffmpeg", "-i", temp_video_path, "-vcodec", "libx264", out_path]
subprocess.run(command)

st.video(out_path)

os.remove(temp_video_path)
os.remove(out_path)