import math
from pathlib import Path
import cv2 as cv
import numpy as np


def find_parallel(lines) -> []:
    """Find parallel lines and return their index."""
    parallel_lines = []
    # TODO Can probably be done using cross-product/zip
    for i in range(len(lines)):
        for j in range(len(lines)):
            if i == j:
                # Every line is parallel to itself. We don't want all lines...
                continue
            if abs(lines[i][0][1] - lines[j][0][1]) == 0:
                parallel_lines.append(i)
    return parallel_lines


def rail_finder_algo_one(filename: str) -> None:
    """Create image with the lines laid over."""
    # TODO Don't load image twice
    source_color_image = cv.imread(cv.samples.findFile(filename), cv.IMREAD_COLOR)
    source_greyscale_image = cv.imread(cv.samples.findFile(filename), cv.IMREAD_GRAYSCALE)
    edge_detected_image = cv.Canny(source_greyscale_image, 50, 200, None, 3)
    # edge_detected_image_color_format = cv.cvtColor(edge_detected_image, cv.COLOR_GRAY2BGR)
    detected_lines = cv.HoughLinesP(edge_detected_image, 1, np.pi / 180, 150, None, 250, 0)
    # if detected_lines is not None:
    #     for i in range(0, len(detected_lines)):
    #         l = detected_lines[i][0]
    #         cv.line(edge_detected_image_color_format, (l[0], l[1]), (l[2], l[3]), (0, 0, 255), 6, cv.LINE_AA)
    parallel_line_indices = find_parallel(detected_lines)
    # Color found lines red
    for detected_line in detected_lines:
        x1, y1, x2, y2 = detected_line[0]
        cv.line(source_color_image, (x1, y1), (x2, y2), (0, 0, 255), 6, cv.LINE_AA)
    # Color parallel lines green
    for parallel_line_index in parallel_line_indices:
        x1, y1, x2, y2 = detected_lines[parallel_line_index][0]
        cv.line(source_color_image, (x1, y1), (x2, y2), (0, 255, 0), 6, cv.LINE_AA)
    cv.imshow("Probabilistic Line Transform", source_color_image)


def rail_finder_algo_two(filename):
    source_greyscale_image = cv.imread(cv.samples.findFile(filename), cv.IMREAD_GRAYSCALE)
    dst = cv.Canny(source_greyscale_image, 50, 200, None, 3)
    cdst = cv.cvtColor(dst, cv.COLOR_GRAY2BGR)
    lines = cv.HoughLines(dst, 1, np.pi / 180, 450, None, 0, 0)
    if lines is not None:
        for i in range(0, len(lines)):
            rho = lines[i][0][0]
            theta = lines[i][0][1]
            a = math.cos(theta)
            b = math.sin(theta)
            x0 = a * rho
            y0 = b * rho
            pt1 = (int(x0 + 1000 * (-b)), int(y0 + 1000 * (a)))
            pt2 = (int(x0 - 1000 * (-b)), int(y0 - 1000 * (a)))
            cv.line(cdst, pt1, pt2, (0, 0, 255), 3, cv.LINE_AA)
    cv.imshow("Standard Hough Line", cdst)


def main():
    filename = str(Path(__file__).parent.parent / Path("tests/octiva") / Path("rails_0.jpg"))
    rail_finder_algo_one(filename)
    # rail_finder_algo_two(filename)
    cv.waitKey()


if __name__ == "__main__":
    main()
