import io
import time
from picamera import PiCamera
import cv2        
import numpy as np
from threading import Thread
from time import sleep

class PiRapidStream:

    def __init__(self,resolution=(640,480),fps=32,form='yuv',exp=15000,gain=1):
        self.w=resolution[0]
        self.h=resolution[1]
        self.fw=(self.w+31)//32*32
        self.fh=(self.h+15)//16*16
        self.frame=None
        self.src=None
        self.form=form
        self.t=None
        self.camera=PiCamera()
        self.camera.resolution = (self.w, self.h)
        self.camera.framerate = fps
        self.camera.exposure_mode='auto'
        self.camera.awb_mode='auto'
        self.camera.iso=0
        #self.camera.video_stabilization=stabilization

        if exp is not 'auto':
            self.camera.exposure_mode='off'
            self.camera.shutter_speed=exp
        
        if gain is not 'auto':
            self.camera.awb_mode='off'
            self.camera.awb_gains=gain
            
        
        
        self.stream_done=True


    def grayscale(self):
        stream = io.BytesIO()
        while not self.stream_done:
           
            yield stream
           
            stream.seek(0)
            self.src=stream.getvalue()
            self.frame=np.fromstring(self.src,dtype=np.uint8,count=self.fw*self.fh).\
               reshape((self.fh,self.fw))
           
            stream.seek(0)
            stream.truncate()

    def coloured(self):
        stream = io.BytesIO()
        while not self.stream_done:
            
            yield stream
           
            stream.seek(0)
            self.src=stream.getvalue()
            self.frame=np.fromstring(self.src,dtype=np.uint8).\
                        reshape((self.fh,self.fw,3))[:self.h,:self.w,:]
            
               
            
            stream.seek(0)
            stream.truncate()

    def capt(self):
        if(self.form=='yuv'):
            self.camera.capture_sequence(self.grayscale(),format=self.form, use_video_port=True)
        if(self.form=='bgr'):
            self.camera.capture_sequence(self.coloured(),format=self.form, use_video_port=True)
        pass
        
    def startcapt(self):
        self.stream_done=False
        self.t = Thread(target=self.capt, args=())
        sleep(0.5)
        self.t.daemon = True
        self.t.start()
        sleep(0.5)
        

    def stopcapt(self):
        self.stream_done=True
        sleep(0.5)
        self.camera.close()
        if self.t is not None:
            self.t.join()





