# Add relative directory ../Library to import path, so we can import the SpoutSDK.pyd library. Feel free to remove these if you put the SpoutSDK.pyd file in the same directory as the python scripts.
# import sys
# sys.path.append('../Library')

import tensorflow as tf
import numpy as np
import transform
import argparse
import time
import SpoutSDK
import pygame
from pygame.locals import *
from OpenGL.GL import *
from OpenGL.GL.framebufferobjects import *
from OpenGL.GLU import *


"""parsing and configuration"""
def parse_args():
    desc = "Tensorflow implementation of 'Perceptual Losses for Real-Time Style Transfer and Super-Resolution', Spout sender/receiver version"
    parser = argparse.ArgumentParser(description=desc)

    parser.add_argument('--style_model', type=str, default='pre_trained_styles/rain_princess.ckpt', help='location for model file (*.ckpt)')

    parser.add_argument('--spout_size', nargs = 2, type=int, default=[640, 480], help='Width and height of the spout receiver and sender')

    parser.add_argument('--spout_name', type=str, default='TDSender1', help='Spout receiving name - the name of the sender you want to receive')

    parser.add_argument('--window_size', nargs = 2, type=int, default=[640, 480], help='Width and height of the window')

    return parser.parse_args()


"""main"""
def main():

    # parse arguments
    args = parse_args()

    # window details
    width = args.window_size[0]
    height = args.window_size[1]
    display = (width,height)

    # window setup
    pygame.init()
    pygame.display.set_caption('Spout Neural Style Sender/Receiver')
    pygame.display.set_mode(display, DOUBLEBUF|OPENGL)

    # OpenGL init
    glMatrixMode(GL_PROJECTION)
    glLoadIdentity()
    glOrtho(0,width,height,0,1,-1)
    glMatrixMode(GL_MODELVIEW)
    glDisable(GL_DEPTH_TEST)
    glClearColor(0.0,0.0,0.0,0.0)
    glEnable(GL_TEXTURE_2D)

    # init spout receiver
    receiverName = args.spout_name
    spoutReceiverWidth = args.spout_size[0]
    spoutReceiverHeight = args.spout_size[1]
    # create spout receiver
    spoutReceiver = SpoutSDK.SpoutReceiver()

	# Its signature in c++ looks like this: bool pyCreateReceiver(const char* theName, unsigned int theWidth, unsigned int theHeight, bool bUseActive);
    spoutReceiver.pyCreateReceiver(receiverName,spoutReceiverWidth,spoutReceiverHeight, False)

    # init spout sender
    spoutSender = SpoutSDK.SpoutSender()
    spoutSenderWidth = args.spout_size[0]
    spoutSenderHeight = args.spout_size[1]
	# Its signature in c++ looks like this: bool CreateSender(const char *Sendername, unsigned int width, unsigned int height, DWORD dwFormat = 0);
    spoutSender.CreateSender('Neural Style Sender', spoutSenderWidth, spoutSenderHeight, 0)

    # create textures for spout receiver and spout sender
    textureReceiveID = glGenTextures(1)
    textureStyleID = glGenTextures(1)

    # initalise receiver texture
    glBindTexture(GL_TEXTURE_2D, textureReceiveID)
    glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE)
    glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST)

    # copy data into texture
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, spoutReceiverWidth, spoutReceiverHeight, 0, GL_RGBA, GL_UNSIGNED_BYTE, None )
    glBindTexture(GL_TEXTURE_2D, 0)

    # open tf session
    soft_config = tf.ConfigProto(allow_soft_placement=True)
    soft_config.gpu_options.allow_growth = True # to deal with large image
    sess = tf.Session(config=soft_config)
    # build tf graph
    style = tf.placeholder(tf.float32, shape=[spoutSenderHeight, spoutSenderWidth, 3], name='input')
    styleI = tf.expand_dims(style, 0) # add one dim for batch

    # result image from transform-net
    scaler = transform.Transform()
    y_hat = scaler.net(styleI/255.0)
    y_hat = tf.squeeze(y_hat) # remove one dim for batch
    y_hat = tf.clip_by_value(y_hat, 0., 255.)

    # initialize parameters
    sess.run(tf.global_variables_initializer())

    # load pre-trained model
    saver = tf.train.Saver()
    saver.restore(sess, args.style_model)

    # loop for graph frame by frame
    while(True):
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                spoutReceiver.ReleaseReceiver()
                pygame.quit()
                quit()

        # receive texture
        # Its signature in c++ looks like this: bool pyReceiveTexture(const char* theName, unsigned int theWidth, unsigned int theHeight, GLuint TextureID, GLuint TextureTarget, bool bInvert, GLuint HostFBO);
        spoutReceiver.pyReceiveTexture(receiverName, spoutReceiverWidth, spoutReceiverHeight, textureReceiveID, GL_TEXTURE_2D, False, 0)

        glBindTexture(GL_TEXTURE_2D, textureReceiveID)
        glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE)
        glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST)
        # copy pixel byte array from received texture
        data = glGetTexImage(GL_TEXTURE_2D, 0, GL_RGB, GL_UNSIGNED_BYTE, outputType=None)  #Using GL_RGB can use GL_RGBA
        glBindTexture(GL_TEXTURE_2D, 0)
        # swap width and height data around due to oddness with glGetTextImage. http://permalink.gmane.org/gmane.comp.python.opengl.user/2423
        data.shape = (data.shape[1], data.shape[0], data.shape[2])

        # start time of the loop for FPS counter
        start_time = time.time()
        #run the graph
        output = sess.run(y_hat, feed_dict={style: data})
        # fiddle back to an image we can display. I *think* this is correct
        output = np.clip(output, 0.0, 255.0)
        output = output.astype(np.uint8)

        # setup the texture so we can load the stylised output into it
        glBindTexture(GL_TEXTURE_2D, textureStyleID);
        glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE)
        glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST)
        # copy output into texture
        glTexImage2D( GL_TEXTURE_2D, 0, GL_RGB, spoutSenderWidth, spoutSenderHeight, 0, GL_RGB, GL_UNSIGNED_BYTE, output )

        # setup window to draw to screen
        glActiveTexture(GL_TEXTURE0)

        # clean start
        glClear(GL_COLOR_BUFFER_BIT  | GL_DEPTH_BUFFER_BIT )
        # reset drawing perspective
        glLoadIdentity()

        # draw texture on screen
        glBegin(GL_QUADS)

        glTexCoord(0,0)
        glVertex2f(0,0)

        glTexCoord(1,0)
        glVertex2f(spoutSenderWidth,0)

        glTexCoord(1,1)
        glVertex2f(spoutSenderWidth,spoutSenderHeight)

        glTexCoord(0,1)
        glVertex2f(0,spoutSenderHeight)

        glEnd()

        # update window
        pygame.display.flip()

        # Send texture to spout...
        # Its signature in C++ looks like this: bool SendTexture(GLuint TextureID, GLuint TextureTarget, unsigned int width, unsigned int height, bool bInvert=true, GLuint HostFBO = 0);
        spoutSender.SendTexture(textureStyleID, GL_TEXTURE_2D, spoutSenderWidth, spoutSenderHeight, False, 0)
        # FPS = 1 / time to process loop
        print("FPS: ", 1.0 / (time.time() - start_time))

if __name__ == '__main__':
    main()
