import os
from PIL import Image
import concurrent.futures
import matplotlib.pyplot as plt


def count_black_pixels_per_row(image_path: str):
    """
    Open an image, count the number of black pixels (RGB = 0, 0, 0)
    in every row, and return a tuple:

        (image_name, list_of_black_pixel_counts_per_row)

    The first row (index 0) is the top of the image; the last row
    (index height – 1) is the bottom.
    """
    image_name = os.path.basename(image_path)
    img = Image.open(image_path).convert("RGB")

    width, height = img.size
    black_counts = []

    for row in range(height):
        row_black = 0
        for x in range(width):
            if img.getpixel((x, row)) == (0, 0, 0):
                row_black += 1
        black_counts.append(row_black)

    return image_name, black_counts


def main():
    # ---------- USER INPUT ----------
    path_taskid = input("Path of the taskid with folder segmented containing '_segmented.png' images: ").strip()
    path_to_target_segmended = os.path.join(path_taskid, "segmented")
    pattern = "_segmented.png"

    # Collect matching files
    image_files = sorted(
        os.path.join(path_to_target_segmended, f)
        for f in os.listdir(path_to_target_segmended)
        if f.endswith(pattern)
    )

    if not image_files:
        print(f"No files ending with '{pattern}' found in: {path_to_target_segmended}")
        return

    print(f"Found {len(image_files)} image(s). Processing in parallel …")

    # ---------- PARALLEL PROCESSING ----------
    results = []
    with concurrent.futures.ProcessPoolExecutor(max_workers=10) as executor:
        future_to_path = {
            executor.submit(count_black_pixels_per_row, img): img
            for img in image_files
        }
        for future in concurrent.futures.as_completed(future_to_path):
            img_path = future_to_path[future]
            try:
                image_name, black_counts = future.result()
                results.append((image_name, black_counts))
            except Exception as err:
                print(f"Error processing {img_path}: {err}")

    if not results:
        print("No results generated. Please check the input images.")
        return

    # ---------- PLOT GENERATION ----------
    plots_dir = os.path.join(path_taskid, "plots_for_verif")
    os.makedirs(plots_dir, exist_ok=True)

    for image_name, black_counts in results:
        height = len(black_counts)

        plt.figure()
        plt.plot(black_counts, range(height))          # x = black pixels, y = row index
        plt.gca().invert_yaxis()                       # keep top row at the top
        plt.title(f"Black Pixels per Row – {image_name}")
        plt.xlabel("Number of Black Pixels")
        plt.ylabel("Row Index (0 = top)")

        save_path = os.path.join(
            plots_dir, f"{image_name}_black_pixels_by_row.png"
        )
        plt.savefig(save_path)
        plt.close()

        print(f"Plot saved: {save_path}")

    print("\n✅ All plots written to:", plots_dir)


if __name__ == "__main__":
    main()