from __future__ import print_function

import pigpio
import io
import time
import threading
import picamera
import cv2
#from PIL import Image
import numpy
from timeit import default_timer as timer

import matplotlib.pyplot as plt
from multiprocessing.dummy import Pool as poop
from sys import exit
# Create a pool of image processors


done = False
lock = threading.Lock()
pool = []
end=0
sums=[]
pi = pigpio.pi()
p=poop(processes=4)
counter=0
center=50
fine_range=[]
global_search=False
do_fine_search=False
w=480
h=480
fw=(w+31)//32*32
fh=(h+15)//16*16
def findsharpness(src):
    #src = cv2.cvtColor(src, cv2.COLOR_BGR2GRAY)

    src=cv2.bilateralFilter(src,3,6,1.5)
    dx = cv2.convertScaleAbs(cv2.Sobel(src, cv2.CV_16S,1,0, ksize=3))
    dy = cv2.convertScaleAbs(cv2.Sobel(src, cv2.CV_16S,0,1, ksize=3))
    accu=cv2.addWeighted(dx,0.5,dy,0.5,0,-1)
    th,accu=cv2.threshold(accu,50,0,cv2.THRESH_TOZERO)
    return int(cv2.sumElems(accu)[0])

##def slicer(src):
##    size=src.shape[0:2:1]
##    return (src[0:size[0]/2,0:size[1]/2],
##            src[0:size[0]/2,size[1]/2:size[1]],
##            src[size[0]/2:size[0],0:size[1]/2],
##            src[size[0]/2:size[0],size[1]/2:size[1]])



class ImageProcessor(threading.Thread):
    def __init__(self):
        super(ImageProcessor, self).__init__()
        self.stream = io.BytesIO()
        self.event = threading.Event()
        self.terminated = False
        self.start()

    def run(self):
        # This method runs in a separate thread
        global done
        global counter
        global fine_range
        global global_search
        global sums
        global do_fine_search
        while not self.terminated:
            # Wait for an image to be written to the stream
            if self.event.wait(1):
                try:
                    self.stream.seek(0)
                    if not global_search:
                        # Read the image and do some processing on it
                        pi.hardware_PWM(18,20000,counter*10000)
                        Y=numpy.fromstring(self.stream.getvalue(),dtype=numpy.uint8,count=fw*fh).\
                           reshape((fh,fw))
                        sums.append(Y)
                        if counter <=95:
                            counter=counter+5
                        elif counter==100:
                            global_search=True
                            
                    elif do_fine_search:
                         pi.hardware_PWM(18,20000,fine_range[counter]*10000)
                         Y=numpy.fromstring(self.stream.getvalue(),dtype=numpy.uint8,count=fw*fh).\
                            reshape((fh,fw))
                         sums.append(Y)
                         if counter <=len(fine_range)-2:
                            counter=counter+1
                        
                         else:
                             done=True
                             do_fine_search=False
                             global_search=False
                    #sums.append(sum([item for item in p.map(findsharpness,slicer(cv2.imdecode(dat,1)),1) if item is not None]))
                    #...
                    #...
                    # Set done to True if you want the script to terminate
                    # at some point
                    #done=True
                finally:
                    # Reset the stream and event
                    self.stream.seek(0)
                    self.stream.truncate()
                    self.event.clear()
                    # Return ourselves to the pool
                    with lock:
                        pool.append(self)

def streams():
    global center
    global counter
    global fine_range
    global global_search
    global sums
    global do_fine_search
    while not done:
        if global_search and not do_fine_search:
            results=p.map(findsharpness,sums)
            center=5*(results.index(max(results))-1)
            counter=0
            
            fine_range=[(x+center) for x in range(5,-6,-1) if (x+center)>=0 and (x+center)<=100]

            pi.hardware_PWM(18,20000,fine_range[0]*10000)
            sums=[]
            do_fine_search=True
        
        with lock:
            if pool:
                processor = pool.pop()
            else:
                processor = None
        if processor:
            yield processor.stream
            processor.event.set()
        else:
            # When the pool is starved, wait a while for it to refill
            time.sleep(0.1)

with picamera.PiCamera() as camera:
    pool = [ImageProcessor() for i in range(4)]
    camera.resolution = (w, h)
    camera.framerate = 32
    camera.exposure_mode='off'
    camera.awb_mode='off'
    camera.shutter_speed=15000
    camera.awb_gains=1
    time.sleep(2)
    print("start")
    start=timer()
    camera.capture_sequence(streams(),format='yuv', use_video_port=True)

# Shut down the processors in an orderly fashion
results=p.map(findsharpness,sums)
focus=fine_range[results.index(max(results))-1]


end=timer()
print(end-start)
print(focus)

#print(len(sums))
for i in range(0,1001,1):
    pi.hardware_PWM(18,20000,(1000-i)*10000)
pi.hardware_PWM(18,20000,0)

#bad=results[::2]
#good=results[1::2]
#plt.figure(1)
#plt.plot(bad,'k')
#plt.plot(results,'k')
#plt.plot(good,'r--')
#plt.show()
while pool:
    with lock:
        processor = pool.pop()
    processor.terminated = True
    processor.join()
p.close()
p.join()
exit()
