import cv2
import argparse
import numpy as np
import time
#import matplotlib.pyplot as plt
# handle command line arguments


def hd_blue(val, mean=170, denom=30):
    num = -np.square(val-mean, dtype='int64') 
    frac = num/(denom**2)
    return np.exp(frac)

def hd_red(val, mean1=0, mean2=255, denom=20):
    num1 = -np.square(val-mean1, dtype='int64')
    frac1 = num1/(denom**2)
    
    num2 = -np.square(val-mean2, dtype='int64')
    frac2 = num2/(denom**2)
    
    return np.exp(frac1) + np.exp(frac2)

def hd_yellow(val, mean=20, denom=20):
    num = -np.square(val-mean, dtype='int64')
    frac = num/(denom**2)
    return np.exp(frac)

def sd(val, mean=255, denom=115):
    center = val-mean
    num = -np.square(val-mean, dtype='int64')
    frac = num/(denom**2)
    return np.exp(frac)

def get_color(image):
    im = cv2.cvtColor(image, cv2.COLOR_BGR2HSV_FULL)
    im = im.astype(np.int16)
    
    saturation = sd(im[...,1])
    blue = np.multiply(hd_blue(im[...,0]), saturation)
    red = np.multiply(hd_red(im[...,0]), saturation)
    yellow = np.multiply(hd_yellow(im[...,0]), saturation)
    
    colors = np.array([blue, red, yellow])
    sums = np.array([arr.sum() for arr in colors])
    
    index = np.argmax(sums)
    
    mask = colors[index]
    
    mask[mask < 0.25] = 0
    mask = (mask*255).astype(np.uint8)
    
    return index, mask
    
def hog(img, bin_n=16):
    gx = cv2.Sobel(img, cv2.CV_32F, 1, 0)
    gy = cv2.Sobel(img, cv2.CV_32F, 0, 1)
    mag, ang = cv2.cartToPolar(gx, gy)
    bins = np.int32(bin_n*ang/(2*np.pi))    # quantizing binvalues in (0...16)
    bin_cells = bins[:10,:10], bins[10:,:10], bins[:10,10:], bins[10:,10:]
    mag_cells = mag[:10,:10], mag[10:,:10], mag[:10,10:], mag[10:,10:]
    hists = [np.bincount(b.ravel(), m.ravel(), bin_n) for b, m in zip(bin_cells, mag_cells)]
    hist = np.hstack(hists)# hist tindra mida 16*4
    return hist

ap = argparse.ArgumentParser()
ap.add_argument('-i', '--input', required=True, help = 'path to input image')
ap.add_argument('-o', '--output', required=True, help = 'path to output video')
ap.add_argument('-c', '--config', required=True, help = 'path to yolo config file')
ap.add_argument('-w', '--weights', required=True, help = 'path to yolo pre-trained weights')
ap.add_argument('-s', '--svm', required=True, help = 'path to svm save file for classification')
args = ap.parse_args()

win_name = "Stream"
cv2.namedWindow(win_name, cv2.WINDOW_NORMAL)

vs = cv2.VideoCapture(args.input)
writer = None
(W, H) = (None, None)

# read pre-trained model and config file
net = cv2.dnn.readNet(args.weights, args.config)
ln = net.getLayerNames()
lln = [ln[i - 1] for i in net.getUnconnectedOutLayers()]

# read pre-trained SVM file for classification
svm = cv2.ml.SVM_load(args.svm)

# loop over frames from the video file stream
while True:
    # read the next frame from the file
    (grabbed, frame) = vs.read()
    # if the frame was not grabbed, then we have reached the end
    # of the stream
    if not grabbed:
        break
    # if the frame dimensions are empty, grab them
    if W is None or H is None:
        (H, W) = frame.shape[:2]

    #blob = cv2.dnn.blobFromImage(frame, scalefactor=1/255.0, swapRB=True)
    blob = cv2.dnn.blobFromImage(frame, scalefactor=1/255.0, size=(640, 384), swapRB=True)
    # set input blob for the network
    net.setInput(blob)
    
    start = time.time()
    outs = net.forward(lln)
    end = time.time()

    # initialization
    class_ids = []
    confidences = []
    boxes = []
    conf_threshold = 0.5
    nms_threshold = 0.4

    # for each detetion from each output layer 
    # get the confidence, class id, bounding box params
    # and ignore weak detections (confidence < 0.5)

    for out in outs:
        for detection in out:
            max_confidence = np.max(detection[5:])
            if max_confidence > 0.5:
                center_x = int(detection[0] * W)
                center_y = int(detection[1] * H)
                w = int(detection[2] * W)
                h = int(detection[3] * H)
                x = center_x - w / 2
                y = center_y - h / 2
                confidences.append(float(max_confidence))
                boxes.append([x, y, w, h])

    # apply non-max suppression
    indices = cv2.dnn.NMSBoxes(boxes, confidences, conf_threshold, nms_threshold)

    # go through the detections remaining
    # after nms and draw bounding box
    for i in indices:
        box = boxes[i]
        print(box)
        x = int(box[0])
        y = int(box[1])
        w = int(box[2])
        h = int(box[3])
        
        patch = frame[y:y+h, x:x+w]
        if patch.size > 0:
            # Compute hogs to classify
            patch = cv2.resize(patch, (32,32))
            gray = cv2.cvtColor(patch, cv2.COLOR_BGR2GRAY)
            color, mask = get_color(patch)
            total = np.concatenate((hog(mask), hog(gray)))

            class_id = svm.predict(np.array([total], dtype='float32'))[1].ravel()

            cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
            text = "{}: {:.4f}".format(class_id, confidences[i])
            cv2.putText(frame, text, (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255,0), 2)
    
    cv2.imshow(win_name, frame)
    if cv2.waitKey(1) == ord('q'):
        break

    if writer is None:
        # initialize our video writer
        fourcc = cv2.VideoWriter_fourcc(*"MJPG")
        writer = cv2.VideoWriter(args.output, fourcc, 30,
                (frame.shape[1], frame.shape[0]), True)
    # write the output frame to disk
    writer.write(frame)


# release the file pointers
print("[INFO] cleaning up...")
writer.release()
vs.release()


