import numpy as np
from PIL import Image, ImageDraw
import math

# Lecture de l'image et conversion en niveaux de gris
img_pil = Image.open('in/grid_contours.png').convert('L')
img_np = np.array(img_pil)

# Calcul de la taille de la diagonale de l'image
width, height = img_np.shape
taille_diagonale = np.ceil(np.sqrt(width * width + height * height)).astype(int)

# Génération des valeurs possibles de rho et theta
num_phis = 360
phis = np.linspace(-math.pi / 2, math.pi / 2, num_phis, endpoint=False)
num_rhos = taille_diagonale * 2 + 1
rhos = np.linspace(-taille_diagonale, taille_diagonale, num_rhos)

# Initialisation de l'accumulateur
accumulator = np.zeros((num_phis, num_rhos), dtype=np.uint64)

# Pour chaque pixel de contour
ys, xs = np.nonzero(img_np)
for y, x in zip(ys, xs):
    # Pour chaque valeur de phi
    for index_phi in range(num_phis):
        phi = phis[index_phi]

        # Calcul de rho en utilisant les coordonnées du pixel et la valeur de phi
        rho = x * math.cos(phi) + y * math.sin(phi)

        # Recherche de l'indice de rho le plus proche de la valeur calculée
        index_rho = np.argmin(np.abs(rhos - rho))

        # Incrémentation de l'accumulateur pour la valeur de phi et rho trouvée
        accumulator[index_phi, index_rho] += 1

# Visualisation de l'accumulateur, normalisé entre 0 et 255
accumulator = (accumulator - accumulator.min()) / (accumulator.max() - accumulator.min()) * 255
Image.fromarray(accumulator.astype(np.uint8)).save('out/road_contours_accumulateur.png')

# Analyse de l'accumulateur pour décider des lignes détectées
nombre_de_lignes_detectées = 140
detections = np.argsort(accumulator, axis=None)[-nombre_de_lignes_detectées:] # Ordre croissant

# Visualisation des lignes détectées
img_pil = img_pil.convert('RGB')
img_pil_draw = ImageDraw.Draw(img_pil)

for i in range(nombre_de_lignes_detectées):
    x, y = np.unravel_index(detections[i], accumulator.shape)

    # Recherche de la valeur de rho et phi correspondant à l'indice
    rho = rhos[y]
    phi = phis[x]

    # Calcul des coordonnées de deux points de la ligne détectée
    if np.sin(phi) != 0:
        x0 = 0
        y0 = int((rho - x0 * np.cos(phi)) / np.sin(phi))
        x1 = width
        y1 = int((rho - x1 * np.cos(phi)) / np.sin(phi))
    else:
        y0 = 0
        x0 = int((rho - y0 * np.sin(phi)) / np.cos(phi))
        y1 = height
        x1 = int((rho - y1 * np.sin(phi)) / np.cos(phi))

    # Dessin de la ligne
    img_pil_draw.line((x0, y0, x1, y1), (255, 0, 0))

# Sauvegarde de l'image
img_pil.save('out/grid_contours_lignes.png')