import argparse

import cv2
import socket
import numpy as np
import uuid

from common import StdPacket, InterlacedPacket, DoublyInterlacedPacket, TiledImagePacket


def send_packet(sock, packet):
    sock.sendto(packet, (UDP_IP, UDP_PORT))

def breakdown_image_norm(frame, last_frame):
    (cols, rows, colors) = frame.shape
    # break the array down into 16x16 chunks, then transmit them as UDP packets
    for i in range(0, cols, 16):
        for j in range(0, rows, 16):
            # print("Sending frame segment (%d, %d)", i, j)
            arr = frame[i:i + 16, j:j + 16]
            last_arr = last_frame[i:i + 16, j:j + 16]
            # only update if image segments are different
            if not np.allclose(arr, last_arr):
                pkt = StdPacket(uuid, j, i, arr)
                send_packet(sock, pkt.to_bytestr())

def breakdown_image_interlaced(frame, last_frame):
    (cols, rows, colors) = frame.shape
    # break the array into 16x32 chunks. we'll split those further into odd and even rows
    # and send each as UDP packets. this should make packet loss less obvious
    for i in range(0, cols, 32):
        for j in range(0, rows, 16):
            # print("Sending frame segment (%d, %d)", i, j)
            arr = frame[i:i + 32:2, j:j + 16]
            last_arr = last_frame[i:i + 32:2, j:j + 16]
            if not np.allclose(arr, last_arr):
                pkt = InterlacedPacket(uuid, j, i, False, arr)
                send_packet(sock, pkt.to_bytestr())

    for i in range(0, cols, 32):
        for j in range(0, rows, 16):
            # print("Sending frame segment (%d, %d)", i, j)
            arr = frame[i + 1:i + 32:2, j:j + 16]
            last_arr = last_frame[i + 1:i + 32:2, j:j + 16]
            # only update if image segments are different
            if not np.allclose(arr, last_arr):
                pkt = InterlacedPacket(uuid, j, i, True, arr)
                send_packet(sock, pkt.to_bytestr())

def breakdown_image_dint(frame, last_frame):
    (cols, rows, colors) = frame.shape
    # break the array into 16x32 chunks. we'll split those further into odd and even rows
    # and send each as UDP packets. this should make packet loss less obvious
    for l in range(0, 4):
        for i in range(0, cols, 32):
            for j in range(0, rows, 32):
                # print("Sending frame segment (%d, %d)", i, j)
                i_even = l % 2 == 0
                j_even = l >= 2

                # breakdown image
                arr = frame[i + i_even:i + 32:2, j + j_even:j + 32:2]
                last_arr = last_frame[i + i_even:i + 32:2, j + j_even:j + 32:2]
                # only update if image segments are different
                if not np.allclose(arr, last_arr):
                    pkt = DoublyInterlacedPacket(uuid, j, i, j_even, i_even, arr)
                    send_packet(sock, pkt.to_bytestr())

def breakdown_image_tiled(frame):
    (cols, rows, colors) = frame.shape
    # break the array into 16x32 chunks. we'll split those further into odd and even rows
    # and send each as UDP packets. this should make packet loss less obvious
    xslice = cols // 16
    yslice = rows // 16
    for i in range(0, xslice):
        for j in range(0, yslice):
            # print("Sending frame segment (%d, %d)", i, j)
            pkt = TiledImagePacket(uuid, j, i, rows, cols, frame[i:cols:xslice, j:rows:yslice])
            send_packet(sock, pkt.to_bytestr())

if __name__ == '__main__':
    # argument parser
    parser = argparse.ArgumentParser(description="Proof-of-concept client for sauron-cv")
    parser.add_argument("-p", "--port", type=int, default=5005)
    parser.add_argument("-s", "--server", type=str, default="127.0.0.1")
    parser.add_argument("-W", "--width", type=int, default=640)
    parser.add_argument("-H", "--height", type=int, default=480)
    parser.add_argument("-d", "--device", type=int, default=0)
    args = parser.parse_args()

    # give this client its very own UUID
    uuid = uuid.uuid4()

    # give target IP address
    UDP_IP = args.server
    UDP_PORT = args.port
    WIDTH = args.width
    HEIGHT = args.height
    DEVICE = args.device

    # Create a VideoCapture object
    cap = cv2.VideoCapture(DEVICE)  # 0 represents the default camera
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, WIDTH)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, HEIGHT)

    # Check if camera opened successfully
    if not cap.isOpened():
        print("Error opening video stream or file")

    # create the socket
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)

    frame = np.zeros((HEIGHT, WIDTH, 3), dtype=np.uint8)
    last_frame = np.zeros((HEIGHT, WIDTH, 3), dtype=np.uint8)

    while True:
        last_frame = frame.copy()
        # Capture frame-by-frame
        ret, frame = cap.read()

        # If frame is read correctly, ret is True
        if not ret:
            print("Can't receive frame (stream end?). Exiting ...")
            break

        breakdown_image_dint(frame, last_frame)

    # Release the capture and close all windows
    cap.release()