#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Combine original + masque segmenté, en supprimant les aberrations
situées sous la coordonnée « y_bottom » fournie par result_convex_hull.csv.

Entrées demandées à l’exécution :
  1) dossier des images originales (.png)
  2) dossier des masques segmentés (_cor.png)
  3) chemin complet vers result_convex_hull.csv (doit contenir 'Label' et 'y_bottom')

Sorties :
  • <segmented_folder>/extracted/<nom>_extracted.png
  • <segmented_folder>/extracted/extracted_info.csv
"""

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


# ------------------------------------------------------------------ #
#  FONCTIONS                                                         #
# ------------------------------------------------------------------ #
def combine_images(original_path, segmented_path, output_path, y_bottom=None):
    """
    Crée 'output_path' (PNG) en copiant les pixels de 'original_path'
    là où le masque 'segmented_path' est noir (0,0,0) **et** où y ≤ y_bottom.
    Tout le reste est laissé blanc. Si y_bottom vaut None, copie complète.
    """
    orig_img = Image.open(original_path).convert("RGB")
    seg_img  = Image.open(segmented_path).convert("RGB")

    if orig_img.size != seg_img.size:
        raise ValueError(
            f"Size mismatch between:\n{original_path}\n{segmented_path}"
        )

    width, height = orig_img.size
    out_img = Image.new("RGB", (width, height), (255, 255, 255))

    for y in range(height):
        # on ignore tout ce qui est sous y_bottom
        if y_bottom is not None and y > y_bottom:
            continue
        for x in range(width):
            if seg_img.getpixel((x, y)) == (0, 0, 0):      # objet
                out_img.putpixel((x, y), orig_img.getpixel((x, y)))
                # sinon : reste blanc

    out_img.save(output_path, "PNG")
    return output_path


def process_one(original_path,
                segmented_path,
                output_path,
                label_for_csv,
                y_bottom_dict):
    """
    Combine les images puis renvoie un enregistrement pour le CSV final.
    """
    y_btm = y_bottom_dict.get(label_for_csv)       # peut être None
    combine_images(original_path, segmented_path, output_path, y_btm)

    return {
        "segmented_file": os.path.basename(segmented_path),
        "extracted_file": os.path.basename(output_path),
        "y_bottom": y_btm
    }


# ------------------------------------------------------------------ #
#  PROGRAMME PRINCIPAL                                               #
# ------------------------------------------------------------------ #
def main():
    print("===== Combine Images + y_bottom filtering =====")

    # --- Saisies utilisateur -------------------------------------- #
    path_taskid = input("path of the taskid containing raw folder with original PNG images: ").strip()
    original_folder = os.path.join(path_taskid, "raw")
    segmented_folder = os.path.join(path_taskid,"segmented_cor") # input("Folder containing segmented '_cor.png' images: ").strip()
    csv_path = os.path.join(path_taskid, "result_convex_hull.csv") # input("Full path to 'result_convex_hull.csv': ").strip()

    # --- Vérifications ------------------------------------------- #
    for fld in (original_folder, segmented_folder):
        if not os.path.isdir(fld):
            print(f"Error: {fld} is not a valid folder.")
            sys.exit(1)
    if not os.path.isfile(csv_path):
        print(f"Error: {csv_path} is not a valid file.")
        sys.exit(1)

    # --- Lecture du CSV y_bottom --------------------------------- #
    try:
        df_csv = pd.read_csv(csv_path, decimal=",")      # pour la virgule
    except ValueError:                                   # sinon point
        df_csv = pd.read_csv(csv_path)

    if "Label" not in df_csv.columns or "y_bottom" not in df_csv.columns:
        print("CSV must contain 'Label' and 'y_bottom' columns.")
        sys.exit(1)

    # Si plusieurs régions : on garde le y_bottom le plus bas (max)
    y_bottom_dict = (
        df_csv.groupby("Label")["y_bottom"].max().to_dict()
    )

    # --- Préparation des fichiers -------------------------------- #
    all_segmented = sorted(
        f for f in os.listdir(segmented_folder) if f.endswith("_cor.png")
    )
    if not all_segmented:
        print(f"No files ending with '_cor.png' found in {segmented_folder}")
        sys.exit(0)

    print(f"Found {len(all_segmented)} segmented image(s).")

    output_folder = os.path.join(path_taskid, "extracted")
    os.makedirs(output_folder, exist_ok=True)

    # --- Traitement parallèle ------------------------------------ #
    tasks, results = [], []
    with concurrent.futures.ProcessPoolExecutor(max_workers=10) as exe:
        for seg_name in all_segmented:
            # On suppose le suffixe officiel
            if "_Simple Segmentation_segmented_cor" not in seg_name:
                print(f"Warning: {seg_name} unexpected suffix. Skipped.")
                continue

            base_name = seg_name.replace(
                "_Simple Segmentation_segmented_cor", ""
            )             # ex: ... .png
            original_path  = os.path.join(original_folder, base_name)
            segmented_path = os.path.join(segmented_folder, seg_name)

            if not os.path.exists(original_path):
                print(f"Warning: original missing for {seg_name}")
                continue

            out_name    = base_name[:-4] + "_extracted.png"
            output_path = os.path.join(output_folder, out_name)

            label_for_csv = seg_name[:-4]        # idem que dans le CSV

            fut = exe.submit(
                process_one,
                original_path,
                segmented_path,
                output_path,
                label_for_csv,
                y_bottom_dict
            )
            tasks.append(fut)

        for fut in concurrent.futures.as_completed(tasks):
            try:
                results.append(fut.result())
            except Exception as e:
                print(f"Error in processing: {e}")

    # --- CSV récapitulatif --------------------------------------- #
    if results:
        df_out = pd.DataFrame(
            results,
            columns=["segmented_file", "extracted_file", "y_bottom"]
        )
        recap_csv = os.path.join(path_taskid, "extracted_info.csv")
        df_out.to_csv(recap_csv, decimal=",", index=False)
        print("\nSummary CSV saved to:", recap_csv)

    print("\nAll done! Extracted images in:", output_folder)


if __name__ == "__main__":
    main()