import time

from pivideocapture import PiVideoCapture

import cv2
import numpy as np

from easygopigo3 import EasyGoPiGo3

with PiVideoCapture(resolution=(320,240), framerate=10) as vc:
    
    print("camera resolution: {}".format(vc.camera_resolution()))
    print("camera framerate : {}".format(vc.camera_framerate()))

    egpg3 = EasyGoPiGo3()
    try:
        # move distance sensor straight (just to look nice)
        egpg3.init_servo(port = "SERVO1").reset_servo()

        print("Press Ctrl-C to stop")
        try:
            while True:
                original = vc.read()
                hsv = cv2.cvtColor(original, cv2.COLOR_BGR2HSV)
                #print("hsv: {0[0]} {0[1]} {0[2]}".format(hsv[0,0]))
                mask = cv2.inRange(hsv, np.array([20,100,70]), np.array([40,255,255]))
                
                cv2.imshow('original', original)
                #cv2.imshow('hsv', hsv)
                cv2.imshow('mask', mask)

                cv2.waitKey(1)

                # each mask element contains the values 0 or 255
                # sum up matching pixels per column
                match_distribution = mask.sum(axis=0) // 255
                total_matches = match_distribution.sum()
                matches = total_matches / mask.size

                if matches > 0.002:
                    sum = 0
                    for i in range(match_distribution.size):
                        sum += i * match_distribution[i]

                    average = sum // total_matches
                    match_center = average / match_distribution.size

                    print("matches col0: {}, total: {} ({:.4f}), center: {} ({:.2f})". \
                        format(match_distribution[0], total_matches, matches, average, match_center))

                    if match_center < 0.45:
                        egpg3.set_speed(egpg3.DEFAULT_SPEED * ((0.5 - match_center) * 1.5))
                        egpg3.spin_left()
                    elif match_center < 0.55:
                        egpg3.stop()
                    else:
                        egpg3.set_speed(egpg3.DEFAULT_SPEED * ((match_center - 0.5) * 1.5))
                        egpg3.spin_right()

                else:
                    print("matches col0: {}, total: {} ({:.4f})". \
                        format(match_distribution[0], total_matches, matches))
                    egpg3.stop()

        except KeyboardInterrupt:
            pass

    finally:
        egpg3.reset_all()

    print("elasped time: {:.2f} secs".format(vc.elapsed_time()))
    print("capture rate: {:.2f} fps".format(vc.capture_framerate()))
    print("drop rate   : {:.2f} fps".format(vc.drop_framerate()))