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(img):
    img1=np.copy(img)
    img2 = cv2.resize(img1, (Lx2,Ly2), interpolation=cv2.INTER_AREA)
    Y=Ly2-1; XM=int(Lx2/2); POS=[]; tlsec=[]
    for j in range(10):
        y=Ly2-1-6*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
                #print('hit green,y,x=',y,x,img2[y][x][:])
                RED.append(x)
            elif(c0<80 and c1<80 and c2<80):
                if(x-1 in RED): 
                    #print('hit BLACK,y,x=',y,x,img2[y][x][:])
                    RED.append(x); 
        img2[y,:,:]=[255,255,255]
        img1[y*F,:,:]=[255,255,255]
        LRED=len(RED)
        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)
    avlsec=Average(tlsec)
    WD=int(avlsec+0.5)*F
    lenP1=len(POS)
    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),12,(0,0,255),3) #--red circle

    #----get boxes
    nsec=0; xsec=[]; ysec=[]; WDS=[]
    cposx=[-1 for i in range(lenP1)]; 
    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]*0.6): 
                    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))
            xsec[-1]=avx
            ysec[-1]=avy
        avx=int(np.median(secx))
        avy=int(np.median(secy))
        xsec[-1]=avx
        ysec[-1]=avy
    xs=xsec; ys=ysec; 
    NN=len(xs)
    for i in range(len(xs)):
        cv2.circle(img2,(xs[i],ys[i]),2,(0,0,0),1) #--black circle
        xs[i]=xs[i]*F
        ys[i]=ys[i]*F
    img22=cv2.resize(img2, (Lx2*F,Ly2*F), interpolation=cv2.INTER_AREA)
    #return NN,xs,ys,WDS,img22
    return NN,xs,ys,WDS,img1