import numpy as np
import cv2
from  motorm import turn
from time import sleep
from pynput import keyboard
TS=0
if(TS==1):
    from tensorflow import keras
    from trackh5 import modh5
import time
#from trackblue import tarbox
#from trackblue import bluebox


def modh5(model,LAB,img):
    test=[]
    test.append(img)
    testa=np.array(test)
    X_test=testa
    X_test=X_test.reshape(X_test.shape[0],128,128,1).astype("float32")
    pred = np.argmax(model.predict(X_test),axis=1)
    p1=pred[0]; L1=LAB[p1]
    #print('pred=',pred,p1,L1)
    return L1



def tarbox(image):
    Lx=image.shape[1]; Ly=image.shape[0];
    Lxg=Lx; Lyg=Ly; Fg=1
    img2 = cv2.resize(image, (Lxg,Lyg), interpolation=cv2.INTER_AREA)
    img2=np.copy(image)
    gray = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
    cnts1,_ = cv2.findContours(gray.copy(), cv2.RETR_EXTERNAL,
          cv2.CHAIN_APPROX_SIMPLE) # cv2.RETR_LIST
    lcnts1=len(cnts1); #print('lcnts1=',lcnts1)
    N1=0; ax=0; cx=0; px=0; mx1=0; Kc1=0; xx1=[0 for i in range(5)]
    clone1=img2.copy(); data1=[]

    #---------------only the largest contour concerned------
    if(lcnts1>0):
        cx=max(cnts1,key=cv2.contourArea)
        areax=cv2.contourArea(cx); areax=round(areax,3); 
        if(areax>1000): 
            M=cv2.moments(cx)
            if(M['m00']==0): M['m00']=1
            cX=int(M['m10']/M['m00']); cY=int(M['m01']/M['m00'])
            area = cv2.contourArea(cx); area3=round(area,3); 
            area=int(area)
            Kc1=1; N1+=1; (x,y,w,h)=cv2.boundingRect(cx);
            xx1=(x,y,w,h); mx1=(cX,cY)
            d=[area,xx1,mx1]; data1.append(d)
            cv2.circle(clone1,mx1,6,(0,125,255),-1)
            x,y,w,h=xx1[0],xx1[1],xx1[2],xx1[3]
            cv2.rectangle(clone1,(x,y),(x+w,y+h),(255,125,0),1)
    
    for c in cnts1:
        M=cv2.moments(c)
        if(M['m00']==0): M['m00']=1
        cX=int(M['m10']/M['m00']); cY=int(M['m01']/M['m00'])
        area = cv2.contourArea(c); area3=round(area,3); #print(j,area3)
        area=int(area)
        if(area<1000): continue
        Kc1=1; N1+=1; (x,y,w,h)=cv2.boundingRect(c);
        xx1=(x,y,w,h); mx1=(cX,cY)
        d=[area,xx1,mx1]; #data1.append(d)
        cv2.circle(clone1,mx1,6,(0,125,255),-1)
        x,y,w,h=xx1[0],xx1[1],xx1[2],xx1[3]
        cv2.rectangle(clone1,(x,y),(x+w,y+h),(255,125,0),1)
        #cv2.imshow('tarbox',clone1)
    
    lenda=len(data1)
    #print(lenda,data1)
    #cv2.waitKey(0)
    return lenda,data1,clone1


def bluebox(image):
    low=np.array([105,120 ,120 ])
    upp=np.array([115,255 ,255 ])
    img=cv2.resize(image,(640,480))
    hsv=cv2.cvtColor(img,cv2.COLOR_BGR2HSV)
    mask1=cv2.inRange(hsv,low,upp)
    imgmask=cv2.bitwise_and(img,img,mask=mask1)
    img3=cv2.resize(imgmask,(640,480))
    lenda,data1,imgt=tarbox(img3)
    #cv2.imshow('mask1',mask1)
    cv2.imshow('imgt',imgt)
    blue=[]; 
    if(lenda>0):
        for i in range(lenda):
            (x,y,w,h)=data1[i][1]
            imgi=img3[y:y+h,x:x+w,:]
            imgir=cv2.resize(imgi,(128,128))
            #cv2.imshow('tar_'+str(i),imgir)
            #cv2.waitKey(0)
            blue.append(imgir)
    return lenda,data1,imgt,blue



Lx=640; Ly=480; F=4; LX2=int(Lx/F); LY2=int(Ly/F)
#.................2 masks for red color
low1=np.array([170,60 ,90 ])
upp1=np.array([180,255,220])
low2=np.array([0,  60, 90 ])
upp2=np.array([10, 255,255])
print('low1=',low1,upp1,' low2=',low2,upp2)
FW2=open('KBS.txt','w')
tini=time.time()
ncat=4; isx=128; isy=128
mod1='model_NUM4C80.h5'
LAB=['1','2','3','4']
if(TS==1): model = keras.models.load_model(mod1)
t00=time.time(); t02=t00
tmodel=round(t00-tini,3)

cap = cv2.VideoCapture(0)
cap.set(cv2.CAP_PROP_FRAME_WIDTH,640)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT,480)
fourcc = cv2.VideoWriter_fourcc(*'XVID')
out = cv2.VideoWriter('redM.avi', fourcc, 20.0, (640, 480))
sleep(2)

COLOR=0
#--------------------- REDLINE under car------
def REDLT(img):
    #dri=input('dri=')
    img1=np.copy(img)
    imgA=cv2.resize(img1,(320,240))
    img2=cv2.resize(img1,(80,60))
    hsv=cv2.cvtColor(img2,cv2.COLOR_BGR2HSV)
    mask1=cv2.inRange(hsv,low1,upp1)
    mask2=cv2.inRange(hsv,low2,upp2)
    mask3=mask1 | mask2
    #mlow=int(60/2); mask3[0:mlow,:]=0; 
    red=[]; red2=[]; red3=[]; red4=[]; Y=60-1; blk=[]
    Y1=59; Y2=49; Y3=39; Y4=29
    for x in range(80):  
        if(mask3[Y1,x]>100): red.append(x)
        if(mask3[Y2,x]>100): red2.append(x)
        if(mask3[Y3,x]>100): red3.append(x)
        if(mask3[Y4,x]>100): red4.append(x)
        if(15<x<65 and hsv[Y3,x,2]<50): blk.append(x)
        #if(COLOR==1): print(x,img2[Y4,x,:],hsv[Y4,x,:],len(red4),
        #        hsv[Y3,x,:],hsv[Y2,x,:])
    if(len(blk)>25): 
        cv2.line(imgA,(0,Y3*4),(79*4,Y3*4),(255,255,0),2)
        cv2.imshow('blk_U_turn',imgA)
        return -100,0,0,0,0,imgA
    #return dr,mid1,mid2,mid3,mid4,redmask
    lred1=len(red); lred2=len(red2); lred3=len(red3); lred4=len(red4)
    if(lred1>1): mid1=int(np.median(red))
    else: mid1=-1; 
    if(lred2>1): mid2=int(np.median(red2))
    else: mid2=-1
    if(lred3>1): mid3=int(np.median(red3))
    else: mid3=-1
    if(lred4>1): mid4=int(np.median(red4))
    else: mid4=-1
    #--------------------drdrdr------
    dr=-10
    if(mid1<=30): dr=3
    if(30<mid1<50 or 30<mid2<50 ): dr=1
    if(mid1>=50): dr=7
    mask3=cv2.resize(mask3,(320,240))
    redmask=cv2.bitwise_and(imgA,imgA,mask=mask3)
    dri=dr; a=[dri,mid1,mid2,mid3,mid4]; sa=str(a)
    #--------------------------------------------------
    cv2.line(redmask,(0,Y4*4),(79*4,Y4*4),(0,255,0),1)
    cv2.circle(redmask,(mid1*4,Y1*4),8,(0,255,0),-1)
    cv2.circle(redmask,(mid2*4,Y2*4),8,(0,255,0),-1)
    cv2.circle(redmask,(mid3*4,Y3*4),8,(0,255,0),-1)
    cv2.circle(redmask,(mid4*4,Y4*4),8,(0,255,0),-1)
    cv2.putText(redmask,sa,(1,20),cv2.FONT_HERSHEY_SIMPLEX,
         0.5, (0,255,255), 1, cv2.LINE_AA)
    if(COLOR==1):
        a=[dr,mid1,mid2,mid3,mid4,lred1,lred2,lred3,lred4]
        print('a=',a)
        cv2.imshow('redmask',redmask)
        cv2.waitKey(0)
        input('1111')
    return dr,mid1,mid2,mid3,mid4,redmask




#--------------------- REDLINE on the LefT-hand side of car------
#---------------------find red pixels from hsv 
COLOR=0
def cREDLT(img):
    #dri=input('dri=')
    img1=np.copy(img)
    imgA=cv2.resize(img1,(320,240))
    img2=cv2.resize(img1,(80,60))
    hsv=cv2.cvtColor(img2,cv2.COLOR_BGR2HSV)
    mask1=cv2.inRange(hsv,low1,upp1)
    mask2=cv2.inRange(hsv,low2,upp2)
    mask3=mask1 | mask2
    #mlow=int(60/2); mask3[0:mlow,:]=0; 
    red=[]; red2=[]; red3=[]; red4=[]; Y=60-1; blk=[]
    Y1=59; Y2=49; Y3=39; Y4=29
    for x in range(65):  #since we focus on the left-hand side of the view
        if(mask3[Y1,x]>100): red.append(x)
        if(mask3[Y2,x]>100): red2.append(x)
        if(mask3[Y3,x]>100): red3.append(x)
        if(mask3[Y4,x]>100): 
            if(len(red4)>0 and red4[-1]<x-6): continue
            else: red4.append(x)
        if(15<x<65 and hsv[Y3,x,2]<50): blk.append(x)
        if(COLOR==1): print(x,img2[Y4,x,:],hsv[Y4,x,:],len(red4),
                hsv[Y3,x,:],hsv[Y2,x,:])
    if(len(blk)>25): 
        cv2.line(imgA,(0,Y3*4),(79*4,Y3*4),(255,255,0),2)
        cv2.imshow('blk_U_turn',imgA)
        return -100,0,0,0,0,imgA
    #return dr,mid1,mid2,mid3,mid4,redmask
    lred1=len(red); lred2=len(red2); lred3=len(red3); lred4=len(red4)
    if(lred1>1): mid1=int(np.median(red))
    else: mid1=-1; 
    if(lred2>1): mid2=int(np.median(red2))
    else: mid2=-1
    if(lred3>1): mid3=int(np.median(red3))
    else: mid3=-1
    if(lred4>1): mid4=int(np.median(red4))
    else: mid4=-1
    #--------------------drdrdr------
    dr=-10
    if(mid1>3 and -1<mid3<8 and mid3<mid4<mid3+5): dr=30
    if(mid1==-1 and -1<mid3<15 and mid4<=mid3): dr=31
    if(mid1==-1 and mid2==-1 and mid3==-1 and mid4==-1): dr=33
    #...........................................
    if(mid1==-1 and mid3-mid2>20): dr=71
    if(mid1>1 and mid2>mid1+12 or mid3>mid1+22): dr=72
    if(mid1<10 and mid2>0 and (mid4-mid3)>(mid3-mid2)+5): dr=72
    if(mid1<5 and mid2<5 and mid3>-1 and mid4>mid3+15): dr=73
    if(mid1<5 and mid2<5 and mid3<5 and 30<mid4<55): dr=73
    if(mid1>30 and mid2>mid1+4): dr=74
    if(-1<mid1<10 and -1<mid2<10): dr=70
    if(mid1>8): dr=71
    mask3=cv2.resize(mask3,(320,240))
    redmask=cv2.bitwise_and(imgA,imgA,mask=mask3)
    dri=dr; a=[dri,mid1,mid2,mid3,mid4]; sa=str(a)
    cv2.circle(imgA,(mid1*4,Y1*4),8,(0,255,0),-1)
    cv2.circle(imgA,(mid2*4,Y2*4),8,(0,255,0),-1)
    cv2.circle(imgA,(mid3*4,Y3*4),8,(0,255,0),-1)
    cv2.circle(imgA,(mid4*4,Y4*4),8,(0,255,0),-1)
    cv2.line(imgA,(0,Y4*4),(79*4,Y4*4),(0,255,0),2)
    cv2.line(redmask,(0,Y4*4),(79*4,Y4*4),(0,255,0),1)
    cv2.circle(redmask,(mid1*4,Y1*4),8,(0,255,0),-1)
    cv2.circle(redmask,(mid2*4,Y2*4),8,(0,255,0),-1)
    cv2.circle(redmask,(mid3*4,Y3*4),8,(0,255,0),-1)
    cv2.circle(redmask,(mid4*4,Y4*4),8,(0,255,0),-1)
    cv2.putText(imgA,sa,(1,20),cv2.FONT_HERSHEY_SIMPLEX,
         0.5, (0,255,255), 1, cv2.LINE_AA)
    cv2.putText(redmask,sa,(1,20),cv2.FONT_HERSHEY_SIMPLEX,
         0.5, (0,255,255), 1, cv2.LINE_AA)
    if(COLOR==1):
        a=[dr,mid1,mid2,mid3,mid4,lred1,lred2,lred3,lred4]
        print('a=',a)
        cv2.imshow('redmask',redmask)
        cv2.waitKey(0)
        input('1111')
    return dr,mid1,mid2,mid3,mid4,redmask



KB=open('KB.txt','w')
break_program = True
sp='A'
def on_press(key):
    global break_program
    global sp
    sp=str(key).strip("'")
    #KB.write(KS+'\n')
    if sp == 'w' and break_program: turn(0.3,0.3,0.1); turn(0,0,0.1)
    if sp == 's' and break_program: turn(-0.3,-0.3,0.1); turn(0,0,0.1)
    if sp == 'a' and break_program: turn(0.0,0.3,0.1); turn(0,0,0.1)
    if sp == 'd' and break_program: turn(0.3,0.0,0.1); turn(0,0,0.1)
    if key == keyboard.Key.f1 and break_program:
        print ('end pressed')
        break_program = False



listener =  keyboard.Listener(on_press=on_press)
listener.start()


#=====================KKKKKKKKKKKKKKKKKK
if(1==2):
    print('into KY_RED...')
    t=0
    while(t < 50 and cap.isOpened()):
        ret,img=cap.read(); 
        if ret == False: break
        dr,mid1,mid2,mid3,mid4,img1=RED(img)
        ta=[mid1,mid2,mid3,lred1,lred2,lred3]
        ws=str(t)+' '+sp+' '+str(ta)
        if(t>0): print(ws)
        KB.write(ws+'\n')
        out.write(img1)
        cv2.imshow('img',img)
        if cv2.waitKey(1) & 0xFF == ord('q'): break
        cv2.waitKey(0)
        t+=1
    cap.release()
    out.release()
    cv2.destroyAllWindows()
    KB.close()
    input('sssss')




def QLT33(cap):
    turn(0.7,0.2,0.2); turn(0.3,0.3,0.1); turn(-0.3,0.3,0.1);
    #turn(0.4,0.0,0.15); turn(-0.3,0.3,0.35); turn(0,0,0.05);
    for j in range(100):
        ret,img=cap.read()
        DRU,mid1,mid2,mid3,mid4,img1=REDLT(img)
        if(int(DRU/10)==7): DRU=7
        #print('QLT33,j=',j,DRU,mid1,mid2,mid3,mid4)
        img3s=cv2.resize(img1,(320,240))
        cv2.imshow('LT',img3s)
        if cv2.waitKey(1) & 0xFF == ord('q'): break
        #cv2.waitKey(0)
        #if(DRU==1 or DRU==7): 
        if(DRU==1 or DRU==-10 or DRU==7 or DRU==-100): 
            turn(0.5,0.0,0.05); turn(0,0,0.10);
            print('QLT33,end_j=',j,DRU,mid1,mid2,mid3,mid4)
            #sleep(2)
            return
        turn(-0.4,0.4,0.05); turn(0,0,0.05);
    print('Unfortunate QLT33()...')
    return

def GODR(DR):
    #.............................................. central redline....
    if(DR==0): turn(0.0,0.0,0.10); turn(0,0,0.05);
    if(DR==1): turn(0.3,0.3,0.10); turn(0,0,0.00);
    if(DR==3): turn(0.0,0.4,0.10); turn(0,0,0.00);
    if(DR==7): turn(0.4,0.0,0.10); turn(0,0,0.00);
    #.............................................. LEFT redline....
    if(DR==-100): turn(-0.4,0.4,0.50);turn(0.4,0.4,0.6);turn(-0.4,0.4,0.5)
    if(DR==-10): turn(0.0,0.0,0.10); turn(0,0,0.05);
    if(DR==30): turn(0.0,0.4,0.15); turn(0,0,0.00);
    if(DR==31): turn(0.5,0.2,0.20); turn(0.2,0.5,0.05);
    #if(DR==32): sleep(2); turn(0.0,0.4,0.10); turn(0,0,0.05);
    if(DR==33): QLT33(cap)
    if(DR==70): turn(0.4,0.0,0.10); turn(0,0,0.0);
    if(DR==71): turn(0.4,0.0,0.10); turn(0,0,0.0);
    if(DR==72): turn(0.4,-0.0,0.10); turn(0,0,0.0);
    if(DR==73): turn(0.4,-0.2,0.10); turn(0,0,0.05);
    #if(DR==-2): turn(-0.3,-0.3,0.05); turn(0,0,0.05);



ret,img=cap.read()
#cv2.imshow('img_in',img)
if(1==1):
    ret,img=cap.read()
    DRj,mid1,mid2,mid3,mid4,m0=REDLT(img)
    m0[:,316:319,:]=[255,255,255]
    m1=np.copy(m0);m2=np.copy(m0);m3=np.copy(m0);m4=np.copy(m0);m5=np.copy(m0);
    imgM=np.hstack((m0,m1,m2,m3,m4,m5))
    #cv2.imshow('img',img)
    cv2.imshow('imgM',imgM)
    #cv2.waitKey(0)
    FW1=open('redM.txt','w')
    t00=time.time(); t02=t00
    tmodel=round(t00-tini,3)
    #-------------------RRRRRRR------
    for j in range(250):
        ret,img=cap.read()
        #--------------------------------------
        lenda,data1,imgt,blue=bluebox(img)
        fn='PNG/NUMF_'+str(j).zfill(3)+'.png'
        if(len(blue)==1): 
            imgb=blue[0]; cv2.imwrite(fn,imgb)
        elif(len(blue)>1): print('error len(blue)=',len(blue)); break
        else: pass
        #cv2.imshow('imgt',imgt)
        if(lenda>0):
            BX=data1[0][1][0];BY=data1[0][1][1];
            BW=data1[0][1][2];BH=data1[0][1][3]
        else:
            BX=0; BY=0; BW=0; BH=0; area=0; aspr=0
        if(len(blue)==0): L1='X'
        for k in range(len(blue)):
            cv2.imshow('blue_'+str(k),blue[k])
            imgg=cv2.cvtColor(blue[k],cv2.COLOR_BGR2GRAY)
            cv2.imshow('imgg_'+str(k),imgg)
            if(TS==1): L1=modh5(model,LAB,imgg)
            t01=time.time(); t=round(t01-t00,3); dt=round(t01-t02,3);
            t02=t01
            #print(j,k,t,dt,L1,type(L1))
            area=data1[0][0]; 
            aspr=round(data1[0][1][2]/data1[0][1][3],2)
        #cv2.waitKey(0)
        #input('11111')
        #--------------------------------------

        imgi2=cv2.resize(img,(320,240))
        cv2.imshow('imgin',imgi2)
        j6=j%6
        DRj,mid1,mid2,mid3,mid4,img1=REDLT(img)
        if(j6==0): m0=np.copy(img1); m0[:,316:319,:]=[255,255,255]
        if(j6==1): m1=np.copy(img1); m1[:,316:319,:]=[255,255,255]
        if(j6==2): m2=np.copy(img1); m2[:,316:319,:]=[255,255,255]
        if(j6==3): m3=np.copy(img1); m3[:,316:319,:]=[255,255,255]
        if(j6==4): m4=np.copy(img1); m4[:,316:319,:]=[255,255,255]
        if(j6==5): m5=np.copy(img1); m5[:,316:319,:]=[255,255,255]
        imgM=np.hstack((m0,m1,m2,m3,m4,m5))
        #cv2.imshow('imgM_'+str(j6),img1)
        cv2.imshow('imgM',imgM)
        #cv2.imwrite('DIR_'+str(j).zfill(3)+'.png',img1)
        DR=DRj
        a=[j,DRj,DR,mid1,mid2,mid3,mid4]; sa=str(a)
        #if(int(DR/10)==7): DR=7
        #if(int(DR/10)==8): DR=8
        #if(DR==-10): DR=1
        FW1.write(sa+'\n')
        img2=cv2.resize(img1,(640,480))
        out.write(img2)
        if cv2.waitKey(1) & 0xFF == ord('q'): break
        #......................CCCCCC
        #DR=1
        #if(lenda==0): DR=0
        #if(area<1000): DR=0
        #print('data1=',data1,data1[0][1][2],data1[0][1][3],aspr)
        #if(lenda>0 and aspr<0.5): DR=0
        if(TS==0): L1=0
        print('j=',j,DRj,DR,mid1,mid2,mid3,mid4,(lenda,BX,BY,BW,BH,aspr,L1))
        #cv2.waitKey(0)
        GODR(DR)
        sleep(1)
    FW1.close()
    FW2.close()
    cv2.destroyAllWindows()
    cap.release()
