# -*- coding: utf-8 -*-

import sys
import socket
import struct
import time
import numpy as np
import threading
import datetime as dt


class Diagnostic:

    def __init__(self, ip='192.168.61.113', port=8001):
        '''Constructor'''

        self.open_connection(ip, port)
        self.data = None
        self.frame_len = -1
        self.num_frames = 0
        self.packet = []

        self.centroid_x = []
        self.centroid_y = []
        self.ci_x = 0
        self.ci_y = 0
        self.centroid_items = 0
        self.shape = np.array([0 for i in range(96)])
        self.img =  np.array([])

    def open_connection(self, ip, port):
        '''Open the socket connection'''

        print("Opening socket on {}:{}".format(ip,port))
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

        try:
            self.sock.bind((ip, port))
            self.sock.settimeout(2)
        except socket.error as e:
            print("Socket error: {}".format(e))
            self.close_connection()
        except Exception as e:
            print("Generic exception: {}".format(e))
            self.close_connection()


    def close_connection(self):
        '''Close the socket connection'''

        print("Closing socket")
        self.sock.close()


    @property
    def image_shape(self):
        '''Choose image depending on frame length'''

        size = self.packet_size

        if size == 2056:
            return (32, 32)
        elif size == 8200:
            return (64, 64)
        elif size == 16136:
            return (96, 84)
        elif size == 129032:
            return (288, 224)
        elif size == 131080:
            return (256, 256)
        elif size == 163848:
            return (320, 256)
        elif size == 655368:
            return (512, 640)
        else:
            return (0, 0)


    def get_packet(self):
        '''Receive a packet frame'''

        timeout = 8000

        try:
            self.data = self.sock.recv(timeout)
        except socket.timeout:
            print("Timeout")
        except Exception as e:
            print("Exception : ", e)
            self.close_connection()

        ## Network Packet Received
        pck_header = struct.unpack('<HLLHH', self.data[0:14])

        # Detected new frame start
        if pck_header[2] == 0:
            if self.frame_len != -1 and self.frame_len != 0:
                pass

            self.frame_len = pck_header[1]

        # Add all several parts composing an image
        self.packet.append(self.data[14:])
        self.frame_len -= len(self.data[14:])

        if self.frame_len == 0:

            ## Diagnostic Packet Complete
            self.packet = b''.join(self.packet)
            self.packet_size = len(self.packet)
            # frame_cnt = struct.unpack('<L', self.packet[0:4])
            # frame_cnt = frame_cnt[0]

            if self.packet_size == 480 or self.packet_size == 64:
                #print("Diag frame received")
                self.diagnostic_frame()

            else:
                # Pixel frame received
                self.pixel_frame()

            # Reset packet
            self.packet = []


    def pixel_frame(self):
        "Get a pixel frame"

        ################
        ## CRED image ##
        ################

        content = struct.unpack('<%iH' % int((self.packet_size-8)/2), self.packet[4:-4])
        img = np.array(content)

        try:
            self.img = img.reshape(self.image_shape)
        except ValueError as e:
            print(e)


    def diagnostic_frame(self):
        "Get a diagnostic frame"

        ##############
        ## DM Shape ##
        ##############

        command_string = '<' + 'f'*97
        self.shape = np.array(struct.unpack(command_string, self.packet[16:16+97*4]))

        ####################
        ## Centroid Cloud ##
        ####################

        tt_string = '<fLffLLff' # power, skipped, centroid x y, 2 not used, tip, tilt
        tt_cmd = struct.unpack(tt_string, self.packet[16 + 104*4:16+104*4 + 32])

        self.centroid_x.append(tt_cmd[2])
        self.centroid_y.append(tt_cmd[3])
        self.centroid_items += 1

        cloud_samples = 2000
        if self.centroid_items > cloud_samples:
            if len(self.centroid_x) > cloud_samples:
                self.centroid_x = self.centroid_x[-int(cloud_samples):]
                self.centroid_y = self.centroid_y[-int(cloud_samples):]

        self.cx = self.centroid_x[-1] #np.mean(self.centroid_x)
        self.cy = self.centroid_y[-1] #np.mean(self.centroid_y)

        plot_range = max([5*np.std(self.centroid_x), 5*np.std(self.centroid_y)])

        self.centroid_matrix, yedges, xedges = np.histogram2d(self.centroid_y, self.centroid_x, bins=21)

        #######################
        ## Centroid Integral ##
        #######################

        if self.ci_x == 0 and self.ci_y == 0:
            self.ci_x = round(tt_cmd[2]*32)
            self.ci_y = round(tt_cmd[3]*32)
