#-----AI4.py

import time
#from PARAM import *
from PARAM2 import *
import cv2
import numpy as np
import os 
from time import sleep
from RED12 import RED1
from  motorm import turn
import matplotlib.pyplot as plt
import math
import random
#from LOCATE import locate
from LOCREC import locate

from tensorflow.keras.optimizers import RMSprop
from tensorflow import keras
import tensorflow as tf
import shutil
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
from tensorflow.keras.layers import Flatten
from tensorflow.keras.layers import Conv2D
from tensorflow.keras.layers import MaxPooling2D
from tensorflow.keras.layers import Dropout
from tensorflow.keras.utils import to_categorical
from sklearn.model_selection import train_test_split

#--------MODEL and PICTURE------
path='PNG'; ncat=7; nfile=200;
mod1='model_EN200C7.h5'
ntest=35; pathT=path+'/T/'
LAB=['A','B','C','L','O','V','E']

model = keras.models.load_model(mod1)
print('load mod1=',mod1,' is done')

cap = cv2.VideoCapture(0)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1920)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT,1080)
sleep(2)
LTAG=['L','O','V','E','A']
TARG=[LAB.index(i) for i in LTAG]
print('LTAG,TARG=',LTAG,TARG)
ncom=1
DX=64; DY=64; jt =0; Njt=40; NP=0; SQ=0; BK=20; NT=0; NYES=0; WDT=0; KTG=0
while(KTG<4 and jt < Njt+1 and cap.isOpened()):   
    target=TARG[KTG]
    XX=[]
    ret,frame = cap.read(); jt+=1
    frame= cv2.resize(frame, (Lx,Ly), interpolation=cv2.INTER_AREA)
    imgc= cv2.resize(frame, (Lx2*2,Ly2*2), interpolation=cv2.INTER_AREA)
    if(ret==False): break
    img=np.copy(frame)
    img2=np.copy(img)
    #cv2.imshow('imgc',imgc)     
    NN,BOX,img24=locate(img2)
    if(NN==0): 
        print('JT=',jt,'null NN=',NN,xs,ys); 
        continue
    for t in range(NN):
        NP+=1
        x1=BOX[t][0]*F; x2=BOX[t][1]*F; y1=BOX[t][2]*F; y2=BOX[t][3]*F
        dx=x2-x1; dy=y2-y1;
        if(dx<40): x2=x1+40
        if(dy<40): y2=y1+40
        cv2.rectangle(img2,(x1,y1),(x2,y2),(0,120,0),4)
        cv2.rectangle(img24,(x1,y1),(x2,y2),(0,120,0),2)
        img3=img[y1:y2,x1:x2,:]
        #print(KTG,jt,' t,img3:',t,img3.shape,' x1,2=',x1,x2)
        img3=cv2.cvtColor(img3, cv2.COLOR_BGR2GRAY)
        fn='PNG/L_a_'+str(NP).zfill(3)+'.png'
        fnt='IMG_'+str(t)
        #print(jt,t,x2-x1,y2-y1,' y1,y2,x1,x2=',y1,y2,x1,x2)
        cv2.imwrite(fn,img3)
        #cv2.imshow(fnt,img3)
        img4=np.copy(img3)
        fimg4='fimg4_'+str(t)
        img4 = cv2.resize(img4,(image_size,image_size))/255
        XX.append(img4)
    DR=1
    imgb= cv2.resize(img2, (Lx2*4,Ly2*4), interpolation=cv2.INTER_AREA)

    X_test2=np.array(XX)
    X_test2=X_test2.reshape(X_test2.shape[0],isz,isz,1).astype("float32")
    #print('X_test2.shape=',X_test2.shape)
    pred2 = np.argmax(model.predict(X_test2),axis=1)
    PLAB=[LAB[pred2[i]] for i in range(NN)]
    DR=11; ks=-1; kps=-1; nks=0;
    for k in range(len(pred2)):
        imgk= X_test2[k,:,:,0]
        kp=pred2[k]
        #print('k=',k,' kp=',kp,LAB[kp],' target=',target)
        #cv2.imshow('imgk',imgk)
        #cv2.waitKey(0)
        if(kp==target):
            #print('inside kp==target',kp,target)
            ks=k; nks+=1
            x1=BOX[k][0]*F; x2=BOX[k][1]*F; y1=BOX[k][2]*F; y2=BOX[k][3]*F
            xsk=int((x1+x2)/2)
            cv2.rectangle(img2,(x1,y1),(x2,y2),(0,0,255), 12)
            if(xsk<30*F*2): DR=3
            elif(xsk>50*F*2): DR=7
            else: DR=1
        else:
            pass
    #cv2.destroyAllWindows()
    imgc=cv2.resize(img2,(Lx2*4,Ly2*4),interpolation=cv2.INTER_AREA)
    #cv2.imshow('IMGC',imgc)
    print(KTG,target,jt,PLAB,' NT,NN=',NT,NN,' DR=',DR)
    #cv2.waitKey(0)
    #cv2.destroyAllWindows()
    #DR=1
    if(DR==1): turn(0.3,0.3,0.10); turn(0,0,0.1);
    if(DR==11): turn(0.3,0.3,0.05); turn(0,0,0.1);
    if(DR==3): turn(0.0,0.4,0.10); turn(0,0,0.1);
    if(DR==7): turn(0.4,0.0,0.10); turn(0,0,0.1);
    if cv2.waitKey(1) & 0xFF == ord('q'): break
    if(jt==Njt): jt=0; KTG+=1; turn(-0.3,-0.3,1.4); turn(0,0,0.1)
cv2.destroyAllWindows()