#!/usr/bin/env python
#
# Cloudlet Infrastructure for Mobile Computing
#
#   Author: Zhuo Chen <zhuoc@cs.cmu.edu>
#
#   Copyright (C) 2011-2013 Carnegie Mellon University
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.
#

import cv2
import json
import multiprocessing
import numpy as np
import os
import pprint
import Queue
import select
import struct
import sys
import threading
import time
import wave

# for setting up stdin
import tty
import termios

dir_file = os.path.dirname(os.path.realpath(__file__))
sys.path.insert(0, os.path.join(dir_file, "../.."))
import gabriel
import gabriel.proxy
LOG = gabriel.logging.getLogger(__name__)


def raw2cv_image(raw_data):
    img_array = np.asarray(bytearray(raw_data), dtype=np.int8)
    cv_image = cv2.imdecode(img_array, -1)
    return cv_image

def display_image(display_name, img, wait_time = -1, is_resize = True, resize_method = "max", resize_max = -1, resize_scale = 1, save_image = False):
    '''
    Display image at appropriate size. There are two ways to specify the size:
    1. If resize_max is greater than zero, the longer edge (either width or height) of the image is set to this value
    2. If resize_scale is greater than zero, the image is scaled by this factor
    '''
    if is_resize:
        img_shape = img.shape
        height = img_shape[0]; width = img_shape[1]
        if resize_max > 0:
            if height > width:
                img_display = cv2.resize(img, (resize_max * width / height, resize_max), interpolation = cv2.INTER_NEAREST)
            else:
                img_display = cv2.resize(img, (resize_max, resize_max * height / width), interpolation = cv2.INTER_NEAREST)
        elif resize_scale > 0:
            img_display = cv2.resize(img, (width * resize_scale, height * resize_scale), interpolation = cv2.INTER_NEAREST)
        else:
            print "Unexpected parameter in image display. About to exit..."
            sys.exit()
    else:
        img_display = img

    cv2.imshow(display_name, img_display)
    cv2.waitKey(wait_time)


class DummyVideoAppWithControl(gabriel.proxy.CognitiveProcessThread):
    def __init__(self, image_queue, output_queue, command_queue, engine_id):
        super(DummyVideoAppWithControl, self).__init__(image_queue, output_queue, engine_id)
        self.command_queue = command_queue

    def handle(self, header, data):
        #print "image\r\n",

        # PERFORM Cognitive Assistance Processing
        header['status'] = "success"
        img = raw2cv_image(data)
        display_image('input', img, resize_max = 640, wait_time = 1)

        # check control messages
        try:
            message_control = self.command_queue.get_nowait()
            header[gabriel.Protocol_client.JSON_KEY_CONTROL_MESSAGE] = json.dumps(message_control)
        except Queue.Empty as e:
            pass

        return json.dumps({'value':'nothing'})


class DummyAccAppWithControl(gabriel.proxy.CognitiveProcessThread):
    def __init__(self, acc_queue, output_queue, command_queue, engine_id):
        super(DummyAccAppWithControl, self).__init__(acc_queue, output_queue, engine_id)
        self.command_queue = command_queue

    def chunks(self, data, n):
        for i in xrange(0, len(data), n):
            yield data[i : i + n]

    def handle(self, header, data):
        #print "acc\r\n",

        ACC_SEGMENT_SIZE = 12 # (float, float, float)
        for chunk in self.chunks(data, ACC_SEGMENT_SIZE):
            (acc_x, acc_y, acc_z) = struct.unpack("!fff", chunk)
            #print "acc_x: %f, acc_y: %f, acc_x: %f" % (acc_x, acc_y, acc_z)
        header['status'] = "success"

        # check control messages
        try:
            message_control = self.command_queue.get_nowait()
            header[gabriel.Protocol_client.JSON_KEY_CONTROL_MESSAGE] = json.dumps(message_control)
        except Queue.Empty as e:
            pass

        return json.dumps({})


class DummyAudioAppWithControl(gabriel.proxy.CognitiveProcessThread):
    def __init__(self, audio_queue, output_queue, command_queue, engine_id):
        super(DummyAudioAppWithControl, self).__init__(audio_queue, output_queue, engine_id)
        self.wav_file_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "dummy_audio.wav")
        self.wav_file = wave.open(self.wav_file_path, "wb")
        self.wav_file.setparams((1, 2, 16000, 0, 'NONE', 'not compressed'))
        self.command_queue = command_queue

    def handle(self, header, data):
        #print "audio\r\n",

        header['status'] = "success"
        if self.wav_file is not None:
            self.wav_file.writeframes(data)

        # check control messages
        try:
            message_control = self.command_queue.get_nowait()
            header[gabriel.Protocol_client.JSON_KEY_CONTROL_MESSAGE] = json.dumps(message_control)
        except Queue.Empty as e:
            pass

        return json.dumps({})

class ControlThread(threading.Thread):
    '''
    The thread provides a simple UI for user to test control of mobile sensors
    KeyStroke       Effect
    1               Toggle on/off for image sensor
    2               Toggle on/off for ACC sensor
    3               Toggle on/off for audio sensor
    a               Low resolution for image sensor
    b               High resolution for image sensor
    c               Low framerate for image sensor
    d               High framerate for image sensor
    '''
    def __init__(self, command_queue):
        self.command_queue = command_queue

        # sensor states
        self.image_on = True
        self.acc_on = True
        self.audio_on = True

        self.stop = threading.Event()

        threading.Thread.__init__(self, target = self.run)

    def __repr__(self):
        return "User Control Thread"

    def run(self):
        while(not self.stop.wait(0.0001)):
            # This is so that stdin doesn't get buffered
            orig_settings = termios.tcgetattr(sys.stdin)
            tty.setraw(sys.stdin)

            rlist, _, _ = select.select([sys.stdin], [], [], 0.05) # the last number is timeout
            if rlist:
                s = sys.stdin.read(1)[0]
                # restore original settings about stdin
                termios.tcsetattr(sys.stdin, termios.TCSADRAIN, orig_settings)
                print "The user enters: " + s

                if s == "1":
                    if self.image_on:
                        print "Turn off image sensor"
                        message_control = {gabriel.Protocol_control.JSON_KEY_SENSOR_TYPE_IMAGE : False}
                    else:
                        print "Turn on image sensor"
                        message_control = {gabriel.Protocol_control.JSON_KEY_SENSOR_TYPE_IMAGE : True}
                    self.image_on = not self.image_on
                    self.command_queue.put(message_control)
                elif s == "2":
                    if self.acc_on:
                        print "Turn off ACC sensor"
                        message_control = {gabriel.Protocol_control.JSON_KEY_SENSOR_TYPE_ACC : False}
                    else:
                        print "Turn on ACC sensor"
                        message_control = {gabriel.Protocol_control.JSON_KEY_SENSOR_TYPE_ACC : True}
                    self.acc_on = not self.acc_on
                    self.command_queue.put(message_control)
                elif s == "3":
                    if self.audio_on:
                        print "Turn off audio sensor"
                        message_control = {gabriel.Protocol_control.JSON_KEY_SENSOR_TYPE_AUDIO : False}
                    else:
                        print "Turn on audio sensor"
                        message_control = {gabriel.Protocol_control.JSON_KEY_SENSOR_TYPE_AUDIO : True}
                    self.audio_on = not self.audio_on
                    self.command_queue.put(message_control)
                elif s == "a":
                    print "Low resolution"
                    message_control = {gabriel.Protocol_control.JSON_KEY_IMG_WIDTH : 240, gabriel.Protocol_control.JSON_KEY_IMG_HEIGHT : 120}
                    self.command_queue.put(message_control)
                elif s == "b":
                    print "High resolution"
                    message_control = {gabriel.Protocol_control.JSON_KEY_IMG_WIDTH : 1280, gabriel.Protocol_control.JSON_KEY_IMG_HEIGHT : 720}
                    self.command_queue.put(message_control)
                elif s == "c":
                    print "Low framerate"
                    message_control = {gabriel.Protocol_control.JSON_KEY_FPS : 1}
                    self.command_queue.put(message_control)
                elif s == "d":
                    print "High framerate"
                    message_control = {gabriel.Protocol_control.JSON_KEY_FPS : 30}
                    self.command_queue.put(message_control)
                elif s == chr(27) or s == chr(3): # esc or ctrl+c
                    self.stop.set()
            else:
                # restore original settings about stdin
                termios.tcsetattr(sys.stdin, termios.TCSADRAIN, orig_settings)

        LOG.info("[TERMINATE] Finish %s" % str(self))

    def terminate(self):
        self.stop.set()


if __name__ == "__main__":
    settings = gabriel.util.process_command_line(sys.argv[1:])

    ip_addr, port = gabriel.network.get_registry_server_address(settings.address)
    service_list = gabriel.network.get_service_list(ip_addr, port)
    #LOG.info("Gabriel Server :")
    #LOG.info(pprint.pformat(service_list))

    video_ip = service_list.get(gabriel.ServiceMeta.VIDEO_TCP_STREAMING_IP)
    video_port = service_list.get(gabriel.ServiceMeta.VIDEO_TCP_STREAMING_PORT)
    acc_ip = service_list.get(gabriel.ServiceMeta.ACC_TCP_STREAMING_IP)
    acc_port = service_list.get(gabriel.ServiceMeta.ACC_TCP_STREAMING_PORT)
    audio_ip = service_list.get(gabriel.ServiceMeta.AUDIO_TCP_STREAMING_IP)
    audio_port = service_list.get(gabriel.ServiceMeta.AUDIO_TCP_STREAMING_PORT)
    ucomm_ip = service_list.get(gabriel.ServiceMeta.UCOMM_SERVER_IP)
    ucomm_port = service_list.get(gabriel.ServiceMeta.UCOMM_SERVER_PORT)

    # this queue is shared by multiple sensor processing threads
    result_queue = multiprocessing.Queue()
    command_queue = Queue.Queue()

    # image receiving and processing
    image_queue = Queue.Queue(gabriel.Const.APP_LEVEL_TOKEN_SIZE)
    #print "TOKEN SIZE OF OFFLOADING ENGINE: %d" % gabriel.Const.APP_LEVEL_TOKEN_SIZE
    video_streaming = gabriel.proxy.SensorReceiveClient((video_ip, video_port), image_queue)
    video_streaming.start()
    video_streaming.isDaemon = True

    video_app = DummyVideoAppWithControl(image_queue, result_queue, command_queue, engine_id = "Dummy_video")
    video_app.start()
    video_app.isDaemon = True

    ## acc receiving and processing
    acc_queue = Queue.Queue(gabriel.Const.APP_LEVEL_TOKEN_SIZE)
    acc_streaming = gabriel.proxy.SensorReceiveClient((acc_ip, acc_port), acc_queue)
    acc_streaming.start()
    acc_streaming.isDaemon = True

    acc_app = DummyAccAppWithControl(acc_queue, result_queue, command_queue, engine_id = "Dummy_acc")
    acc_app.start()
    acc_app.isDaemon = True

    # audio receiving and processing
    audio_queue = Queue.Queue(gabriel.Const.APP_LEVEL_TOKEN_SIZE)
    audio_streaming = gabriel.proxy.SensorReceiveClient((audio_ip, audio_port), audio_queue)
    audio_streaming.start()
    audio_streaming.isDaemon = True

    audio_app = DummyAudioAppWithControl(audio_queue, result_queue, command_queue, engine_id = "Dummy_audio")
    audio_app.start()
    audio_app.isDaemon = True

    # control UI thread
    control_thread = ControlThread(command_queue)
    control_thread.start()
    control_thread.isDaemon = True

    # result pub/sub
    result_pub = gabriel.proxy.ResultPublishClient((ucomm_ip, ucomm_port), result_queue, log_flag = False)
    result_pub.start()
    result_pub.isDaemon = True

    try:
        while True:
            time.sleep(1)
    except Exception as e:
        pass
    except KeyboardInterrupt as e:
        LOG.info("user exits\n")
    finally:
        if video_streaming is not None:
            video_streaming.terminate()
        if video_app is not None:
            video_app.terminate()
        if acc_streaming is not None:
            acc_streaming.terminate()
        if acc_app is not None:
            acc_app.terminate()
        if audio_streaming is not None:
            audio_streaming.terminate()
        if audio_app is not None:
            audio_app.terminate()
        if control_thread is not None:
            control_thread.terminate()
        result_pub.terminate()
