from PARAM import *
import cv2
import numpy as np

def Average(lst):
    if(len(lst)==0): avg=0
    else: avg=sum(lst)/len(lst)
    return avg

#**********************************************-------------------------
def locate_g(img):
    img1=np.copy(img)
    img2 = cv2.resize(img1, (Lx2,Ly2), interpolation=cv2.INTER_AREA)
    img3=np.copy(img2)
    img3[:,:,:]=[255,255,255]
    #--green target
    dz=0.10; cr1=0.55; cr2=0.10; z1=cr1-dz; z2=cr1+dz; z3=cr2-dz; z4=cr2+dz
    #--tennis ball
    #dz=0.20; cr1=0.47; cr2=0.94; z1=cr1-dz; z2=cr1+dz; z3=cr2-dz; z4=cr2+dz
    #--blue target--a,r1,r2,b3= [189  88  18] 0.1 0.47 [2 1 0]
    #--blue phone stand--a,r1,r2,b3= [145 107  45] 0.31 0.74 [2 1 0]] 
    dz=0.15; cr1=0.3; cr2=0.74; z1=cr1-dz; z2=cr1+dz; z3=cr2-dz; z4=cr2+dz


    Y=Ly2-1; XM=int(Lx2/2); POS=[]; tlsec=[]
    NYL=30; DY=int(Ly2/NYL)
    for j in range(NYL):
        y=Ly2-1-DY*j
        RED=[]
        for x in range(Lx2):
            c0=1.*img2[y][x][0]; c1=1.*img2[y][x][1]; c2=1.*img2[y][x][2];
            #if(c0<c1*0.7 and c0<c2*0.6 and c1<c2 and c0+c1+c2>250):
            #if(c0<c1*0.8 and c2<c1*0.8 and c1>70): # green target
            #--tennis ball
            #if(z1<c0/c1<z2 and z3<c2/c1<z4): 
            #--blue goal
            if(c0<1): continue
            if(z1<c2/c0<z2 and z3<c1/c0<z4): 
                RED.append(x)
        LRED=len(RED)
        #print('y,LRED,RED=',y,LRED,RED)
        if(LRED<6): continue
        for k in range(LRED): img3[y,RED[k],:]=[0,0,0]
        img2[y,:,:]=[255,255,255]
    
        #--- sections ---
        if(LRED>0): 
            ns=0; sec=[]; coord=[]; secs=[]; lsec=[]
            for i in range(1,LRED):
                if(RED[i]-RED[i-1]<6):
                    sec.append(RED[i-1])
                    if(i==LRED-1 and len(sec)>0):
                        ns+=1
                        sec.append(RED[i])
                        mid=int(np.median(np.array(sec)))
                        coord.append((y,mid,len(sec)))
                        secs.append(sec)
                        tlsec.append(len(sec))
                        lsec.append(len(sec))
                else:
                    if(len(sec)>0): 
                        tlsec.append(len(sec))
                        ns+=1
                        sec.append(RED[i-1])
                        mid=int(np.median(np.array(sec)))
                        coord.append((y,mid,len(sec)))
                        secs.append(sec)
                        sec=[]
            POS.extend(coord)
    lenP1=len(POS)
    if(lenP1<1): return 0,[],[],[],img1
    XP=[POS[i][1] for i in range(lenP1)]
    #print('POS_f=',POS,XP)
    #print('tlsec=',tlsec)
    avlsec=int(Average(tlsec))
    Xav=int(Average(XP))
    wd=int(avlsec+1); wdx=int(avlsec+1); wdy=int(POS[0][0]-POS[lenP1-1][0])
    #print('Xav=',Xav,wd,wdx,wdy)
    WD=int(avlsec+1)*F
    WDX=int(avlsec+1)*F
    WDY=int(POS[0][0]-POS[lenP1-1][0])*F
    #print('lenP1,avlsec,Xav,WD=',lenP1,avlsec,Xav,WD)
    for k in range(lenP1):
        #cv2.circle(img2,(POS[k][1],POS[k][0]),3,(0,0,255),1) #--red circle
        cv2.circle(img1,(POS[k][1]*F,POS[k][0]*F),6,(0,0,255),2) #--red circle
    #x1=int((Xav-wd/2-0)*F); x2=int((Xav+wd/2+0)*F); 
    #y2=(POS[0][0]+DY)*F; y1=(POS[lenP1-1][0]-DY)*F
    #cv2.rectangle(img1,(x1,y1),(x2,y2),(0,255,255),5)

    #----get boxes
    nsec=0; xsec=[]; ysec=[]; WDS=[]; wdy=[]; wdx=[]; numsec=[]
    cposx=[-1 for i in range(lenP1)]; numsec=[]
    for i in range(lenP1):
        if(cposx[i]!=-1): continue
        else: 
            nsec+=1
            cposx[i]=nsec 
            x1=POS[i][1]
            xsec.append(x1)
            ysec.append(POS[i][0])
            WDS.append(POS[i][2]*F)
            secx=[x1]; secy=[POS[i][0]]
            for k in range(lenP1):
                if(cposx[k]!=-1): continue
                if(abs(POS[k][1]-x1)<POS[i][2]*1.5): 
                    secx.append(POS[k][1])
                    secy.append(POS[k][0])
                    cposx[k]=nsec; #print(i,nsec,x1,k,POS[k][1])
            avx=int(np.median(secx))
            avy=int(np.median(secy))
            vx=POS[i][2]
            vy=np.max(secy)-np.min(secy)
            wdx.append(vx); wdy.append(vy)
            xsec[-1]=avx
            ysec[-1]=avy
        avx=int(np.median(secx))
        avy=int(np.median(secy))
        xsec[-1]=avx
        ysec[-1]=avy
        #X1=np.min(secx)*F; X2=np.max(scex)
        iy=0
        for y in secy:
            X1=int(secx[iy]-wdx[-1]/2)*F
            X2=int(secx[iy]+wdx[-1]/2)*F
            #print(iy,y,X1/F,X2/F)
            cv2.line(img1,(X1,y*F),(X2,y*F),(255,255,255),1)
            iy+=1
        #cv2.line(img1, (0, 0), (255, 255), (0, 255, 255), 5)
        #print('nsec=',nsec,' secx=',secx,' secy=',secy,' wdx=',wdx,' wdy=',wdy)
        dleny=len(set(secy))
        numsec.append(dleny)
    xs=xsec; ys=ysec; 
    NN=len(xs)
    #print('NN=',NN,lenP1,' xs=',xs,' ys=',ys,' WDS=',WDS,'wdx,y=',wdx,wdy)
    for i in range(len(xs)):
        #print(i,xs[i],ys[i],img2.shape)
        cv2.circle(img2,(xs[i],ys[i]),6,(0,0,0),3) #--black circle
        #cv2.circle(img2,(POS[k][1],POS[k][0]),3,(0,0,255),1) #--red circle
        #cv2.circle(img1,(POS[k][1]*F,POS[k][0]*F),12,(0,0,255),3) #--red circle
        x1=int(xs[i]-wdx[i]/2-1)*F; y1=int(ys[i]+wdy[i]/2+1)*F 
        x2=int(xs[i]+wdx[i]/2+1)*F; y2=int(ys[i]-wdy[i]/2-1)*F
        cv2.rectangle(img1,(x1,y1),(x2,y2),(0,125,255),2)
        xs[i]=xs[i]*F
        ys[i]=ys[i]*F
    img22=cv2.resize(img2, (Lx2*F,Ly2*F), interpolation=cv2.INTER_AREA)
    #cv2.imshow('IMG22',img22)
    #cv2.imshow('IMG2',img2)
    #cv2.imshow('IMG3',img3)
    #return NN,xs,ys,WDS,img22
    maxsec=np.max(numsec); WDS.append(maxsec)
    #print('numsec=',numsec,maxsec)
    #input('rrrrr')
    return NN,xs,ys,WDS,img1