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

"""
Summary RGB and Row by Row
Row-by-row vegetation-index analysis + surface calculation.

For every image (RGB):
  • Ignores background pixels (pure white 255,255,255).
  • For each row (height position) it computes:
      – mean R, G, B
      – mean of several vegetation indices (GI, Redden, …, TGI)
      – SurfacePixels  : count of plant pixels in that row
      – SurfacePct     : SurfacePixels / row width
  • Saves one CSV per image in a sub-folder 'rowwise_data'.

Output columns (per row, per image):
  RowIndex | Label | MeanR MeanG … MeanTGI | SurfacePixels | SurfacePct
"""

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


# ---------------------------------------------------------------------
# helpers
# ---------------------------------------------------------------------
def safe_div(numer, denom):
    """Return numer / denom, or NaN if denom is zero."""
    return numer / denom if denom != 0 else float("nan")


def compute_indices(r, g, b):
    """Return a dict with all desired vegetation indices for one pixel."""
    GI       = safe_div(g, (r + b))
    Redden   = safe_div(r, g)
    PctGreen = safe_div(g, (r + g + b))
    RGDiff   = (r - g)
    NGRDI    = safe_div((g - r), (g + r))
    VI       = safe_div((g - b), (g + b))
    TGI      = -0.5 * ((190 * (r - g)) - (120 * (r - b)))

    return {
        "R": r, "G": g, "B": b,
        "GI": GI, "Redden": Redden, "PctGreen": PctGreen,
        "RGDiff": RGDiff, "NGRDI": NGRDI, "VI": VI, "TGI": TGI
    }


# ---------------------------------------------------------------------
# core analysis
# ---------------------------------------------------------------------
def analyze_image_rowwise(image_path):
    """
    For every row of `image_path`:
      • keep only non-white pixels (plant)
      • compute means of colour channels and indices
      • compute vegetation surface (count / percentage)
    Return a DataFrame with one record per row.
    """
    img = Image.open(image_path).convert("RGB")
    arr = np.array(img)                    # shape (H, W, 3)
    height, width, _ = arr.shape
    label = os.path.splitext(os.path.basename(image_path))[0]

    rows_data = []

    for row_idx in range(height):
        row_pixels = arr[row_idx, :, :]
        mask = ~np.all(row_pixels == [255, 255, 255], axis=-1)   # True = plant

        if not np.any(mask):          # no plant pixels in this row
            rows_data.append({
                "RowIndex": row_idx,
                "Label": label,
                "MeanR": np.nan, "MeanG": np.nan, "MeanB": np.nan,
                "MeanGI": np.nan, "MeanRedden": np.nan, "MeanPctGreen": np.nan,
                "MeanRGDiff": np.nan, "MeanNGRDI": np.nan,
                "MeanVI": np.nan, "MeanTGI": np.nan,
                "SurfacePixels": 0,
                "SurfacePct": 0.0
            })
            continue

        row_nonbg = row_pixels[mask]
        R = row_nonbg[:, 0].astype(float)
        G = row_nonbg[:, 1].astype(float)
        B = row_nonbg[:, 2].astype(float)

        pixel_indices = [compute_indices(r, g, b) for r, g, b in zip(R, G, B)]
        df_row = pd.DataFrame(pixel_indices)
        means = df_row.mean()

        surface_pixels = mask.sum()
        surface_pct = surface_pixels / width

        rows_data.append({
            "RowIndex": row_idx,
            "Label": label,
            "MeanR": means["R"], "MeanG": means["G"], "MeanB": means["B"],
            "MeanGI": means["GI"], "MeanRedden": means["Redden"],
            "MeanPctGreen": means["PctGreen"], "MeanRGDiff": means["RGDiff"],
            "MeanNGRDI": means["NGRDI"], "MeanVI": means["VI"],
            "MeanTGI": means["TGI"],
            "SurfacePixels": surface_pixels,
            "SurfacePct": surface_pct
        })

    return pd.DataFrame(
        rows_data,
        columns=[
            "RowIndex", "Label",
            "MeanR", "MeanG", "MeanB",
            "MeanGI", "MeanRedden", "MeanPctGreen",
            "MeanRGDiff", "MeanNGRDI", "MeanVI", "MeanTGI",
            "SurfacePixels", "SurfacePct"
        ]
    )


def process_and_save(image_path, output_folder):
    """Run row-wise analysis for one image and save its CSV."""
    df_image = analyze_image_rowwise(image_path)
    label = os.path.splitext(os.path.basename(image_path))[0]
    out_csv = os.path.join(output_folder, f"{label}_rowwise.csv")
    df_image.to_csv(out_csv, index=False, decimal=",")
    print(f"Saved row-wise CSV → {out_csv}")


# ---------------------------------------------------------------------
# main
# ---------------------------------------------------------------------
def main():
    print("===== Row-wise Vegetation Indices & Surface =====")

    folder_path = input("Enter folder path of images: ").strip()
    if not os.path.isdir(folder_path):
        print(f"Error: '{folder_path}' is not a valid folder."); sys.exit(1)

    exts = (".jpg", ".jpeg", ".png", ".bmp", ".tiff")
    files = sorted(f for f in os.listdir(folder_path) if f.lower().endswith(exts))
    if not files:
        print(f"No image files found in {folder_path}"); sys.exit(0)

    print(f"Found {len(files)} images. Processing…")

    output_folder = os.path.join(folder_path, "rowwise_data")
    os.makedirs(output_folder, exist_ok=True)

    with concurrent.futures.ProcessPoolExecutor(max_workers=8) as exe:
        futures = [exe.submit(process_and_save,
                              os.path.join(folder_path, fn),
                              output_folder)
                   for fn in files]

        for fut in concurrent.futures.as_completed(futures):
            try:
                fut.result()
            except Exception as e:
                print("Error:", e)

    print("\nDone! CSV files are in:", output_folder)


if __name__ == "__main__":
    main()