import cv2
import numpy as np
import time
import datetime
import os
import sys
GPU=True
THREAD=True
if len(sys.argv)>1:
    if sys.argv[1]=='A':
        cPATH='rtsp://5173:6703@192.168.68.74/live'
    else:
        cPATH='rtsp://admin:@192.168.68.128:554/1/h264major'
else:
        cPATH='rtsp://5173:6703@192.168.68.74/live'
PATH="/home/mars/pWork/DATA"
#

class ThreadingVideoCapture:
    def __init__(self, src, max_queue_size=256):
        self.video = cv2.VideoCapture(src)
        self.q = queue.Queue(maxsize=max_queue_size)
        self.stopped = False

    def start(self):
        thread = threading.Thread(target=self.update, daemon=True)
        thread.start()
        return self

    def update(self):
        while True:
            if self.stopped:
                return
            if not self.q.full():
                ok, frame = self.video.read()
                self.q.put((ok, frame))
                if not ok:
                    self.stop()
                    return

    def read(self):
        return self.q.get()

    def stop(self):
        self.stopped = True

    def release(self):
        self.stopped = True
        self.video.release()

    def isOpened(self):
        return self.video.isOpened()

    def get(self, i):
        return self.video.get(i)

def key(k):
    global th, tc,track,reverse
    if k == ord('2'):
        th = th - 1
    elif k == ord('3'):
        th = th + 1
    elif k == ord('4'):
        tc = tc -5
    elif k == ord('5'):
        tc = tc +5
    elif k == ord('t'):
        track = not track
    elif k == ord('r'):
        reverse = not reverse

fontFace =cv2.FONT_HERSHEY_SIMPLEX
track, reverse = False,False
avg=None
writer = None
th = 30
tc = 30
x,y=0,0
detect_counts = 0
red,blue,green = (0,0,255),(255,0,0),(0,255,0)
time_start = time.time()
frame=0
log=PATH+'/metro.log'
if cPATH=='rtsp://5173:6703@192.168.68.74/live':
    TITLE="ATOM"
    HEAD ='ATOM'
else:
    TITLE="ONVIF"
    HEAD='ONVIF'

if THREAD:
    import threading
    import queue
    TITLE=TITLE+"-T"
    capture = ThreadingVideoCapture(cPATH)
    capture.start()
    if not capture.isOpened():
        raise RuntimeError
else:
    capture=cv2.VideoCapture(cPATH)

if GPU:
    TITLE=TITLE+"-G"
    img_gpu_src = cv2.cuda_GpuMat() # Allocate device memory only once, as memory allocation seems to take time...
    img_gpu_dst = cv2.cuda_GpuMat()
    img_gpu_gray= cv2.cuda_GpuMat()

W = capture.get(cv2.CAP_PROP_FRAME_WIDTH)
H = capture.get(cv2.CAP_PROP_FRAME_HEIGHT)
W2=int(W/2)
H2=int(H/2)
fourcc = cv2.VideoWriter_fourcc(*"XVID")
print('Camera:',cPATH)
print('Size:',W,H)
while(True):
    ret, img = capture.read()
    if ret:
        org = img.copy()
        if GPU:
            img_gpu_src.upload(img)
            img_gpu_dst = cv2.cuda.resize(img_gpu_src, dsize=(W2, H2))
            img_gpu_dst = cv2.cuda_GpuMat(img_gpu_dst,[0,int(H2*0.85)],[0,W2])
            img_gpu_gray=cv2.cuda.cvtColor(img_gpu_dst,cv2.COLOR_BGR2GRAY)
            org_img=img_gpu_dst.download()
            gray = img_gpu_gray.download()
        else:
            org_img = cv2.resize(img, dsize=(W2, H2))
            org_img=img[0:int(H2*0.85),0:W2]
            gray = cv2.cvtColor(org_img, cv2.COLOR_BGR2GRAY)
        if reverse:
            gray=cv2.bitwise_not(gray)
            avg=cv2.bitwise_not(avg)
        if avg is None:
            avg = gray.copy().astype("float")
            continue

        #wtiter,fname = moving(img,avg)
        cv2.accumulateWeighted(gray, avg, 0.5)
        frameDelta = cv2.absdiff(gray, cv2.convertScaleAbs(avg))
        thresh = cv2.threshold(frameDelta, th, 255, cv2.THRESH_BINARY)[1]

        contours,hierarchy = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        detect=False
        for i in range(0,len(contours)):
            if len(contours[i]) > 0:
                 if cv2.contourArea(contours[i]) > tc:
                    detect=True
                    time_start = time.time()
                    if writer is None and track:
                        detect_counts = 0
                        now=datetime.datetime.today()
                        date=now.strftime("%Y%m%d")
                        cDIR=PATH+'/'+date
                        if not(os.path.exists(cDIR)):
                            os.mkdir(cDIR)
                        fname=cDIR+'/'+ HEAD + now .strftime("%Y%m%d_%H%M%S")+".avi"
                        writer = cv2.VideoWriter(fname, fourcc, 15, (int(W), int(H)))
                    rect = contours[i]
                    x, y, w, h = cv2.boundingRect(rect)
                    cv2.rectangle(org_img, (x-w, y-h), (x + w*2, y + h*2), red, 3)
        if detect:
            detect_counts=detect_counts + 1
        if time.time() -  time_start  > 5:
            if writer is not None:
                writer.release()
                new_name=fname.replace(HEAD,HEAD+'_' + f'{detect_counts:04}'+'_')
                os.rename(fname,new_name)
                frame=0
                writer = None
        now=datetime.datetime.today()
        text=now.strftime("%Y%m%d %H%M%S")+' No:'+str(frame)+ ' '+" TH:"+str(th)+" SZ:"+str(tc)
        org_img = cv2.putText(org_img, text, (30,50), fontFace,1,color=green)
        org = cv2.putText(org, text, (30,50), fontFace,1,color=green)
        text1="REC:"+str(track) + "  reverse:" + str(reverse)
        if writer is not None:
            frame=frame+1
            text1=fname+' '+text1
        org_img = cv2.putText(org_img, text1, (30,80), fontFace,1,color=green)
        #cv2.imshow('thresh-level',thresh)

        cv2.imshow(TITLE,org_img)

        if writer is not None:
            writer.write(org)
    else:
        now=datetime.datetime.today()
        date=now.strftime("%Y%m%d_%H%M%S")
        print('reconnect:',date)
        capture.release()
        avg=None
        if THREAD:
            capture = ThreadingVideoCapture(cPATH)
            capture.start()
        else:
            capture = cv2.VideoCapture(cPATH)
    k=cv2.waitKey(1) & 0xFF
    key(k)
    if k== ord('q'):
        break

capture.release()
if writer is not None:
    writer.release()
    new_name=fname.replace(HEAD,HEAD+'_' + f'{detect_counts:04}'+'_')
    os.rename(fname,new_name)
cv2.destroyAllWindows()