# vision.py
# OpenCV machine vision image processing functions.
# No copyright, 2024, Garth Zeglin.  This file is explicitly placed in the public domain.

# Import the standard Python math library.
import math

# Import the third-party numpy library for matrix calculations.
# Note: this can be installed using 'pip3 install numpy' or 'pip3 install scipy'.
import numpy as np

# Import the third-party OpenCV library for image operations.
# Note: this can be installed using 'pip3 install opencv-python'
import cv2 as cv

################################################################
def find_circles(grayscale_image):

    # normalize levels to full range of pixel values
    frame = cv.normalize(grayscale_image, None, 0, 255, cv.NORM_MINMAX, dtype=cv.CV_8U)

    # reduce noise
    img = cv.medianBlur(frame,5)

    circles = cv.HoughCircles(img,
                              method=cv.HOUGH_GRADIENT,
                              dp=1,        # accumulator has same resolution as the image
                              minDist=8,   # minimum center to center distance
                              param1=200,  # with HOUGH_GRADIENT, Canny edge threshold
                              param2=20,   # with HOUGH_GRADIENT, accumulator threshold
                              minRadius=0,
                              maxRadius=0)

    if circles is None:
        return None
    else:
        # the result seems to be a 1 x num-circles x 3 matrix
        # this reshapes it to a 2D matrix, each row is then [x y r]
        num_circles = circles.shape[1]
        return circles.reshape((num_circles, 3))

################################################################
def annotate_circles(grayscale_image, circles):

    # create color image for visualization
    cimg = cv.cvtColor(grayscale_image,cv.COLOR_GRAY2BGR)
    if circles is not None:
        circles = np.uint16(np.round(circles))
        for i in circles:
            # draw the perimeter in green
            cv.circle(cimg,(i[0],i[1]),i[2],(0,255,0),2)

            # draw the center of the circle in red
            cv.circle(cimg,(i[0],i[1]),2,(255,0,0),3)
    return cimg

################################################################
def find_lines(grayscale_image):
    # convert to binary image
    # ret, frame = cv.threshold(frame, 127, 255, cv.THRESH_BINARY)

    # normalize levels to full range of pixel values
    frame = cv.normalize(grayscale_image, None, 0, 255, cv.NORM_MINMAX, dtype=cv.CV_8U)

    # find edges, return binary edge image
    frame = cv.Canny(frame, 100, 200)

    # extract possible lines
    lines = cv.HoughLines(frame, rho=1, theta=math.radians(10), threshold=20)
    if lines is not None:
        print("found %d possible lines, shape is %s" % (len(lines), lines.shape))
        # print(lines)
    else:
        print("no lines found")
    return lines

################################################################
def annotate_lines(grayscale_image, lines):

    # show the image in the debugging viewer window
    render = cv.cvtColor(grayscale_image, cv.COLOR_GRAY2RGB)
    if lines is not None:
        for line in lines:
            rho, theta = line[0]  # line is a matrix with shape (1, 2)
            a = np.cos(theta)
            b = np.sin(theta)
            x0 = a*rho
            y0 = b*rho
            x1 = int(x0 + 100*(-b))
            y1 = int(y0 + 100*(a))
            x2 = int(x0 - 100*(-b))
            y2 = int(y0 - 100*(a))
            cv.line(render,(x1,y1),(x2,y2),(255,0,0),1)

    return render

################################################################
# Test script to perform when run this file is run as a standalone script.
if __name__ == "__main__":
    import argparse
    import os.path
    parser = argparse.ArgumentParser( description = """Test vision processing using a test image.""")
    parser.add_argument('paths', nargs="+", help="One or more image files to analyze.")
    args = parser.parse_args()

    # read a test image
    for path in args.paths:
        frame = cv.imread(path)
        if frame is None:
            print("Unable to load %s." % (path))
        else:
            print("Read image %s with shape: %s" % (path, frame.shape))

            # convert to grayscale and normalize levels
            frame = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)

            circles = find_circles(frame)

            if circles is None:
                print("no circles detected")
            else:
                print("detected %d circle(s):" % (circles.shape[0]))
                print(circles)

                cimg = annotate_circles(frame, circles)
                debug_path = os.path.join("debug-images", os.path.basename(path))
                print("Writing ", debug_path)
                cv.imwrite(debug_path, cimg)