import os
import sys
import pandas as pd
import numpy as np
from PIL import Image
import concurrent.futures

def process_one_image(file_path, output_folder, y_min=0, y_max=2625): # for 38177, 1100 2900
    """
    Opens a '_segmented.png' image, transforms black pixels (0,0,0) outside the
    vertical range [y_min, y_max] to white (255,255,255), keeps black pixels 
    inside this range, and saves the corrected image as '*_segmented_cor.png'.
    
    Returns:
        (filename_no_ext, black_count)
    """
    # Example: file_path = "/path/to/..._segmented.png"
    filename = os.path.basename(file_path)  # e.g. "XXX_segmented.png"
    filename_no_ext = filename[:-4]        # e.g. "XXX_segmented"
    
    # Open image as RGB
    img = Image.open(file_path).convert("RGB")
    width, height = img.size
    
    # Create a new image (RGB) to store the corrected pixels
    corrected_img = Image.new("RGB", (width, height), (255, 255, 255))
    
    black_count = 0
    
    # Iterate over rows and columns
    for y in range(height):
        for x in range(width):
            r, g, b = img.getpixel((x, y))
            
            # If pixel is black and within [y_min, y_max], keep black
            if (r, g, b) == (0, 0, 0) and (y_min <= y <= y_max):
                corrected_img.putpixel((x, y), (0, 0, 0))
                black_count += 1
            else:
                # Otherwise, use white (255, 255, 255)
                corrected_img.putpixel((x, y), (255, 255, 255))
    
    # Construct the new filename ending with '_segmented_cor.png'
    new_filename = filename_no_ext + "_cor.png"
    save_path = os.path.join(output_folder, new_filename)
    
    # Save the corrected image
    corrected_img.save(save_path, "PNG")
    
    return filename_no_ext, black_count


def main():
    # 1) Ask the user for the path to the folder
    path_taskid = input("Enter the path of the taskid were ther ise segmended folder with '_segmented.png' images: ").strip()
    path_to_target_segmended = os.path.join(path_taskid, "segmented")

    # 2) Create an output subfolder for corrected images
    output_subfolder = "segmented_cor"
    output_folder = os.path.join(path_taskid, output_subfolder)
    os.makedirs(output_folder, exist_ok=True)
    
    # 3) List all '_segmented.png' files in that folder
    all_files = os.listdir(path_to_target_segmended)
    all_files.sort()
    segmented_files = [f for f in all_files if f.endswith("_segmented.png")]
    
    if not segmented_files:
        print(f"No files ending with '_segmented.png' found in: {path_to_target_segmended}")
        sys.exit(0)
    
    # 4) Parallel processing
    print(f"Found {len(segmented_files)} image(s). Processing in parallel...")
    
    # We'll store the results (filename_no_ext, black_count)
    results = []
    
    with concurrent.futures.ProcessPoolExecutor(max_workers=8) as executor:
        future_to_file = {}
        for f in segmented_files:
            file_path = os.path.join(path_to_target_segmended, f)
            future = executor.submit(process_one_image, file_path, output_folder)
            future_to_file[future] = f
        
        for future in concurrent.futures.as_completed(future_to_file):
            original_file = future_to_file[future]
            try:
                filename_no_ext, black_count = future.result()
                results.append((filename_no_ext, black_count))
            except Exception as e:
                print(f"Error processing {original_file}: {e}")
    
    # 5) Create and save a CSV with the black pixel counts
    if results:
        df = pd.DataFrame(results, columns=["Label", "BlackPixels"])
        csv_path = os.path.join(path_taskid, "result_black_pixel_corrected.csv")
        df.to_csv(csv_path, index=False)
        print(f"\nCSV created: {csv_path}")
    else:
        print("No images processed or no black pixels found. CSV not created.")

if __name__ == "__main__":
    main()