import os

from ij import IJ, ImagePlus
from ij.gui import Roi, Overlay
from ij.measure import ResultsTable
from ij.plugin.filter import ThresholdToSelection
from java.awt import Color

# Okabe-Ito colorblind-safe pair
COLOR_P63_POSITIVE = Color(230, 159, 0)   # orange
COLOR_P63_NEGATIVE = Color(86, 180, 233)  # sky blue

# Stroke width for ROI outlines
_STROKE_WIDTH = 1.5


def load_mask(mask_path):
    """Open the label-image TIFF at mask_path and return an ImagePlus."""
    imp = IJ.openImage(mask_path)
    if imp is None:
        raise IOError("Could not open mask: {}".format(mask_path))
    return imp


def masks_to_rois(mask_imp):
    """Convert each unique non-zero label in mask_imp to an ROI; return list of (label, Roi) tuples.

    Uses ThresholdToSelection on a copy of the processor for each label so the
    original image is never modified.  Labels are uint16 values from Cellpose output.
    """
    ip = mask_imp.getProcessor().duplicate()
    # Collect unique non-zero labels from the pixel array
    pixels = list(ip.getPixels())
    labels = sorted(set(int(p) for p in pixels) - {0})

    result = []
    tts = ThresholdToSelection()
    work_imp = ImagePlus("_mask_work", ip)

    for lbl in labels:
        # Isolate this label: threshold exactly at [lbl, lbl]
        ip.setThreshold(lbl, lbl)
        roi = tts.convert(ip)
        if roi is not None:
            roi.setStrokeWidth(_STROKE_WIDTH)
            result.append((lbl, roi))

    work_imp.close()
    return result


def _build_mean_lookup(cells_data, variant):
    """Return {label_int: mean_float} and the chosen threshold float from cells_data rows.

    cells_data rows are dicts from read_cells_csv -- all values are strings.
    The chosen variant column is e.g. 't14_positive'.
    """
    by_label = {}
    for row in cells_data:
        lbl = int(row["cell_label"])
        mean = float(row["red_mean"])
        by_label[lbl] = mean
    return by_label


def color_rois_by_p63(rois_with_labels, cells_data, threshold, variant):
    """Color each ROI orange (p63+) or sky blue (p63-) based on red_mean vs threshold.

    rois_with_labels  -- list of (label_int, Roi) from masks_to_rois
    cells_data        -- list of dicts from read_cells_csv (values are strings)
    threshold         -- float on 0-255 scale (from summary.csv)
    variant           -- string e.g. 'T14' (used only for logging)
    Returns the same list with ROI stroke colors set in place.
    """
    mean_by_label = _build_mean_lookup(cells_data, variant)
    n_pos = n_neg = n_missing = 0

    for lbl, roi in rois_with_labels:
        mean = mean_by_label.get(lbl)
        if mean is None:
            # Label present in mask but absent from cells.csv -- skip coloring
            n_missing += 1
            continue
        if mean >= threshold:
            roi.setStrokeColor(COLOR_P63_POSITIVE)
            n_pos += 1
        else:
            roi.setStrokeColor(COLOR_P63_NEGATIVE)
            n_neg += 1

    IJ.log("p63+ coloring ({}): pos={} neg={} unmatched={}".format(
        variant, n_pos, n_neg, n_missing))
    return rois_with_labels


def build_overlay(rois_with_labels):
    """Wrap the colored Roi objects from (label, Roi) pairs into an Overlay and return it."""
    overlay = Overlay()
    for _lbl, roi in rois_with_labels:
        overlay.add(roi)
    return overlay


def populate_results_table(cells_data):
    """Fill and return a ResultsTable with one row per cell (image, cell_label, red_mean, p63+).

    cells_data -- list of dicts from read_cells_csv; values are strings.
    The 'variant_call' column (the primary p63+ call) is mapped to 'p63_positive'.
    """
    rt = ResultsTable()
    for row in cells_data:
        rt.incrementCounter()
        rt.addValue("image",       row["image"])
        rt.addValue("cell_label",  float(row["cell_label"]))
        rt.addValue("red_mean",    float(row["red_mean"]))
        rt.addValue("p63_positive", float(row["variant_call"]))
    return rt


def display_results(imp, overlay, results_table):
    """Attach overlay to imp, show it, and display results_table in the ImageJ Results window."""
    imp.setOverlay(overlay)
    imp.show()
    # updateAndDraw refreshes the canvas so the overlay renders immediately
    imp.updateAndDraw()
    results_table.show("p63 Results")