'''
This file contains functions that generate RGB images and segmented images from an unreal engine simulation
'''

from __future__ import division, absolute_import, print_function
import os, sys, time, re, json, csv
import numpy as np
import matplotlib.pyplot as plt
import unrealcv
import pandas
import cv2
from tempfile import TemporaryFile

# Functions from the unrealCV documentation
imread = plt.imread
def imread8(im_file):
    ''' Read image as a 8-bit numpy array '''
    im = np.asarray(Image.open(im_file))
    return im

def read_png(res):
    import io, PIL.Image
    img = PIL.Image.open(io.BytesIO(res))
    return np.asarray(img)

def read_npy(res):
    import io
    return np.load(io.BytesIO(res))

def get_RGB_images(folder,client, name, num_from):
    '''
    INPUT: 
    - folder: name where the images generated will be stored. If folder under directory doesn't exist it will generate the folder.
    - client: generated by setting up a connection using unrealCV
    - name: name of images to be generated, following format name_5.png
    - num_from: number to start naming images after the underscore
    
    OUTPUT:
    - folder containing RGB - original images from the simulation
    
    WARNING:
    - a CSV file containg the camera trajectory must exist under same directory (use UE4_camera_traj.py to generate one)
    '''
    #Creating directory if it doesn't exist
    if not os.path.exists(folder):
        os.mkdir(folder)

    #reading camera trajectory from simulation
    trajectories = pandas.read_csv('camera_traj.csv')

    # looping through every camera view of the camera trajectory
    for traj in range (len(trajectories)-1): 

        #Getting values from csv file    
        x= str(trajectories['x'][traj])
        y= str(trajectories['y'][traj])
        z= str(trajectories['z'][traj])
        pitch= str(trajectories['pitch'][traj])
        yaw= str(trajectories['yaw'][traj])
        roll= str(trajectories['roll'][traj])

        ## Set position of the camera
        client.request('vset /camera/0/location %s %s %s'%(x,y,z))
        client.request('vset /camera/0/rotation %s %s %s'%(pitch, yaw, roll))

        #getting original image as a npy array
        res = client.request('vget /camera/0/lit png')
        im = read_png(res)

        #changing RGBA to RGB (data is read only so I save it into a temporary file)
        outfile = TemporaryFile()
        RGB = im[:,:,0:3]
        np.save(outfile, RGB) # saving data
        outfile.seek(0)
        RGB_img = np.load(outfile) # load data

        #saving original image into folder
        img_num = str(num_from+traj+1)
        image_name = name+"_"+ img_num+".png"
        output = os.path.join (folder, image_name )
        plt.imsave(output, RGB_img)
        
def get_segmented_images(folder,client,name,num_from):
    '''
    INPUT: 
    - folder: name where the images generated will be stored. If folder under directory doesn't exist it will generate the folder.
    - client: generated by setting up a connection using unrealCV
    - name: name of images to be generated, following format name_5.png
    - num_from: number to start naming images after the underscore
    
    OUTPUT:
    - folder containing segmented images  - labeled images from the simulation objects not defined in CSV file will have an RGB value 
    of (0,0,0) i.e black
    
    WARNING:
    - a CSV file containg the camera trajectory must exist under same directory (use UE4_camera_traj.py to generate one)
    '''
    
    #Creating directory if it doesn't exist
    if not os.path.exists(folder):
        os.mkdir(folder)

    #reading csv with trajectories
    trajectories = pandas.read_csv('camera_traj.csv')

    for traj in range (len(trajectories)-1): # looping through every camera view

        #Getting values from csv     
        x= str(trajectories['x'][traj])
        y= str(trajectories['y'][traj])
        z= str(trajectories['z'][traj])
        pitch= str(trajectories['pitch'][traj])
        yaw= str(trajectories['yaw'][traj])
        roll= str(trajectories['roll'][traj])

        # Set position of the first camera
        client.request('vset /camera/0/location %s %s %s'%(x,y,z))
        client.request('vset /camera/0/rotation %s %s %s'%(pitch, yaw, roll))

        #getting object_mask image as a np array
        res = client.request('vget /camera/0/object_mask png')
        im = read_png(res)

        #changing RGBA to RGB
        outfile = TemporaryFile()
        new = im[:,:,0:3]
        np.save(outfile, new) # saving data (Otherwise error because data cannot be modified)
        outfile.seek(0)
        new_num_arr = np.load(outfile) # load data

        # for some reason the sky color and the fruit color is set to a different value (maybe because of lightning in their blueprint)     
        new_num_arr[np.where((new_num_arr == [99,0,0] ).all(axis=2))] = [100,10,10]
        new_num_arr[np.where((new_num_arr == [68,130,180] ).all(axis=2))] = [70,130,180]

        #Saving image to folder original_data
        img_num = str(num_from+traj+1)
        image_name = name+"_"+ img_num+".png"
        output = os.path.join (folder, image_name )
        plt.imsave(output, new_num_arr ) 
        
def get_segmented_images_object_mask(folder,client,name, num_from):
    '''
    INPUT: 
    - folder: name where the images generated will be stored. If folder under directory doesn't exist it will generate the folder.
    - client: generated by setting up a connection using unrealCV
    - name: name of images to be generated, following format name_5.png
    - num_from: number to start naming images after the underscore
    
    OUTPUT:
    - folder containing segmented images  - labeled images from the simulation using object mask (to get specific objects)
    
    WARNING:
    - a CSV file containg the camera trajectory must exist under same directory (use UE4_camera_traj.py to generate one)
    '''
    
    #Creating directory if it doesn't exist
    if not os.path.exists(folder):
        os.mkdir(folder)

    #reading csv with camera trajectories
    trajectories = pandas.read_csv('camera_traj.csv')

    for traj in range (len(trajectories)-1): # looping through every camera view

        #Getting values from csv     
        x= str(trajectories['x'][traj])
        y= str(trajectories['y'][traj])
        z= str(trajectories['z'][traj])
        pitch= str(trajectories['pitch'][traj])
        yaw= str(trajectories['yaw'][traj])
        roll= str(trajectories['roll'][traj])

        # Set position of the first camera
        client.request('vset /camera/0/location %s %s %s'%(x,y,z))
        client.request('vset /camera/0/rotation %s %s %s'%(pitch, yaw, roll))

        #getting object_mask image as a np array
        res = client.request('vget /camera/0/object_mask png')
        img = get_segmentation(res,client)

        #Saving image to folder original_data
        img_num = str(num_from+traj+1)
        image_name = name+"_"+ img_num+".png"
        output = os.path.join (folder, image_name )
        plt.imsave(output, img )
        
def get_RGB_images_depth_mask(folder, client,name, num_from, depth_filter):
    '''
    INPUT: 
    - folder: name where the images generated will be stored. If folder under directory doesn't exist it will generate the folder.
    - client: generated by setting up a connection using unrealCV
    - name: name of images to be generated, following format name_5.png
    - num_from: number to start naming images after the underscore
    -depth_filter: integer which defines the distance that sets up the depth filter
    
    OUTPUT:
    - folder containing original images using a depth mask - RGB images from the simulation
    
    WARNING:
    - a CSV file containg the camera trajectory must exist under same directory (use UE4_camera_traj.py to generate one)
    '''
    
    #Creating directory if it doesn't exist
    if not os.path.exists(folder):
        os.mkdir(folder)

    #reading csv with trajectories
    trajectories = pandas.read_csv('camera_traj.csv')
    
    for traj in range (len(trajectories)-1): # looping through every camera view

        #Getting values from csv     
        x= str(trajectories['x'][traj])
        y= str(trajectories['y'][traj])
        z= str(trajectories['z'][traj])
        pitch= str(trajectories['pitch'][traj])
        yaw= str(trajectories['yaw'][traj])
        roll= str(trajectories['roll'][traj])
        
        # Set position of the first camera
        client.request('vset /camera/0/location %s %s %s'%(x,y,z))
        client.request('vset /camera/0/rotation %s %s %s'%(pitch, yaw, roll))

        #getting depth image and creating depth mask
        depthRequest = client.request('vget /camera/0/depth depth.exr')
        depth = cv2.imread(depthRequest, cv2.IMREAD_ANYDEPTH)
        mask = np.where(depth > depth_filter)
        
        #getting original image
        res = client.request('vget /camera/0/lit png')
        im = read_png(res)

        #changing RGBA to RGB (data is read only so I save it into a temporary file)
        outfile = TemporaryFile()
        RGB = im[:,:,0:3]
        np.save(outfile, RGB) # saving data
        outfile.seek(0)
        RGB_image = np.load(outfile) # load data
        
        RGB_image[mask] = [0,0,0]

        #saving original image into folder
        img_num = str(num_from+traj+1)
        image_name = name+"_"+ img_num+".png"
        output = os.path.join (folder, image_name )
        plt.imsave(output, RGB_image)
             
def get_segmented_images_depth_mask(folder, client,name, num_from, depth_filter):
    '''
    INPUT: 
    - folder: name where the images generated will be stored. If folder under directory doesn't exist it will generate the folder.
    - client: generated by setting up a connection using unrealCV
    - name: name of images to be generated, following format name_5.png
    - num_from: number to start naming images after the underscore
    -depth_filter: integer which defines the distance that sets up the depth filter
    
    OUTPUT:
    - folder containing segmented images using a depth mask - labeled images from the simulation
    
    WARNING:
    - a CSV file containg the camera trajectory must exist under same directory (use UE4_camera_traj.py to generate one)
    '''
    
    #Creating directory if it doesn't exist
    if not os.path.exists(folder):
        os.mkdir(folder)

    #reading csv with trajectories
    trajectories = pandas.read_csv('camera_traj.csv')
    
    for traj in range (len(trajectories)-1): # looping through every camera view

        #Getting values from csv     
        x= str(trajectories['x'][traj])
        y= str(trajectories['y'][traj])
        z= str(trajectories['z'][traj])
        pitch= str(trajectories['pitch'][traj])
        yaw= str(trajectories['yaw'][traj])
        roll= str(trajectories['roll'][traj])
        
        # Set position of the first camera
        client.request('vset /camera/0/location %s %s %s'%(x,y,z))
        client.request('vset /camera/0/rotation %s %s %s'%(pitch, yaw, roll))

        #getting depth image and creating depth mask
        depthRequest = client.request('vget /camera/0/depth depth.exr')
        depth = cv2.imread(depthRequest, cv2.IMREAD_ANYDEPTH)
        mask = np.where(depth > depth_filter)
        
        #getting object_mask image as a np array
        res = client.request('vget /camera/0/object_mask png')
        img = get_segmentation(res,client)
        
        # using depth filter
        img[mask] = [0,0,0]
        
        #Saving image to folder original_data
        img_num = str(num_from+traj+1)
        image_name = name+"_"+ img_num+".png"
        output = os.path.join (folder, image_name )
        plt.imsave(output, img )
    
def get_depth(client,name,num_from):
    
#     INPUT: 
#     - folder: name where the images generated will be stored. If folder under directory doesnt exist it will generate the folder.
#     - client: generated by setting up a connection using unrealCV
#     - name: name of images to be generated, following format name_5.png
#     - num_from: number to start naming images after the underscore
#     -depth_filter: integer which defines the distance that sets up the depth filter
    
#     OUTPUT:
#     - depth images from the simulation saved as exr files under C:\Program Files\Epic Games\UE_4.16\Engine\Binaries\Win64
    
#     WARNING:
#     - a CSV file containg the camera trajectory must exist under same directory (use UE4_camera_traj.py to generate one)
    
    
    #reading csv with trajectories
    trajectories = pandas.read_csv('camera_traj.csv')

    for traj in range (len(trajectories)-1): # looping through every camera view

        #Getting values from csv     
        x= str(trajectories['x'][traj])
        y= str(trajectories['y'][traj])
        z= str(trajectories['z'][traj])
        pitch= str(trajectories['pitch'][traj])
        yaw= str(trajectories['yaw'][traj])
        roll= str(trajectories['roll'][traj])

        # Set position of the first camera
        client.request('vset /camera/0/location %s %s %s'%(x,y,z))
        client.request('vset /camera/0/rotation %s %s %s'%(pitch, yaw, roll))

        # getting depth as an EXR file
        img_num = str(num_from+traj+1)
        image_name = name+"_"+ img_num+".exr"
        client.request('vget /camera/0/depth %s'%(image_name))

def get_segmentation(res, client):
    '''
    INPUT: 
    - res: data generated by using UnrealCV command 'vget /camera/0/lit png'
    - client: generated by setting up a connection using unrealCV
    
    OUTPUT:
    - numpy representation of segmented image applying mask to only get Trees, sky and people (hard coded into function) 
    '''
    img = read_png(res)

    #changing RGBA to RGB (data is read only so I save it into a temporary file)
    outfile = TemporaryFile()
    new = img[:,:,0:3]
    np.save(outfile, new) # saving data
    outfile.seek(0)
    img = np.load(outfile) # load data
    
    #getting all classes we dont want to change colors of
    img[np.where((img == [128,128,0]).all(axis=2))] = [255,255,255] #trees
    img[np.where((img == [68,130,180]).all(axis=2))] = [255,255,255] #sky (68 should be 70 - bug somewhere)
    img[np.where((img == [230,230,0]).all(axis=2))] = [255,255,255] #people
    
    #create mask for everything that is not white i.e part of the object we want
    mask = np.where((img != [255,255,255]).all(axis=2))
    
    #getting the original image again
    img = read_png(res)

    #changing RGBA to RGB (data is read only so I save it into a temporary file)
    outfile = TemporaryFile()
    new = img[:,:,0:3]
    np.save(outfile, new) # saving data
    outfile.seek(0)
    img = np.load(outfile) # load data

    img[np.where((img == [68,130,180] ).all(axis=2))] = [70,130,180] #solving bug issue
    img[np.where((img == [255,255,255] ).all(axis=2))] = [230,230,0] #everythin not in mask becomes part of the ground

    img[mask] = [128,64,0]
    
    return img



        
      
       

        
      