import cv2
import numpy as np
import matplotlib.pyplot as plt

def plot_one_image(image: np.ndarray) -> None:
    """
    Отобразить изображение с помощью matplotlib.
    Вспомогательная функция.

    :param image: изображение для отображения
    :return: None
    """
    fig, axs = plt.subplots(1, 1, figsize=(8, 7))

    axs.imshow(image)
    axs.axis('off')
    plt.plot()
    plt.show()


def find_road_number(image: np.ndarray) -> int:
    """
    Найти номер дороги, на которой нет препятсвия в конце пути.

    :param image: исходное изображение
    :return: номер дороги, на котором нет препятсвия на дороге
    """
    h, w, _ = image.shape
    # slice roads
    image = image / 255
    x = 10
    prev_val = 0
    road_starts = []
    road_ends = []
    for y in range(1, image.shape[1] - 1):
        if np.greater(image[x,y], [0.9, 0.9, 0.1]).all():
            cur_val = y
            if cur_val - prev_val == 1:
                prev_val = cur_val
            else:
                road_starts.append(prev_val)
                road_ends.append(cur_val)
                prev_val = cur_val
    # find car
    pix_sum = np.sum(image[image.shape[0]//2:, :], axis=0)
    blue_quan = pix_sum[:, 2]
    blue_quan = blue_quan / np.max(blue_quan)
    blues = []
    for i in range(len(road_starts)):
        pix_slice = blue_quan[road_starts[i]:road_ends[i]]
        max_blue = np.max(pix_slice)
        blues.append(max_blue)
    car_place = blues.index(1)

    # find road to ride
    pix_sum = np.sum(image[:image.shape[0]//2, :], axis=0)
    red_quan = pix_sum[:, 0]
    red_quan = red_quan/np.max(red_quan)
    reds = []
    for i in range(len(road_starts)):
        pix_slice = red_quan[road_starts[i]:road_ends[i]]
        mean_red = np.mean(pix_slice)
        reds.append(mean_red)
    index_min = min(range(len(reds)), key=reds.__getitem__)

    if car_place == index_min:
        print('Do not change the road')
    else:
        print('You should drive on road', index_min)

    road_number = index_min
    return road_number


if __name__ == '__main__':
    test_image = cv2.imread('task_2/image_01.jpg')
    test_image = cv2.cvtColor(test_image, cv2.COLOR_BGR2RGB)
    # plot_one_image(test_image)
    find_road_number(test_image)