import gradio as gr
import numpy as np
import pandas as pd
from csbdeep.utils import normalize
from skimage.color import rgb2gray
from skimage.measure import regionprops
from skimage.morphology import binary_closing
from skimage.util import img_as_ubyte
from skimage.measure import shannon_entropy
from stardist.models import StarDist2D
from stardist.plot import render_label


MODEL_NAMES = [
    "2D_versatile_fluo",
    "2D_versatile_he",
    "2D_paper_dsb2018",
]


_MODEL_CACHE = {}


def get_model(name: str) -> StarDist2D:
    if name not in _MODEL_CACHE:
        _MODEL_CACHE[name] = StarDist2D.from_pretrained(name)
    return _MODEL_CACHE[name]


def to_gray(image: np.ndarray) -> np.ndarray:
    if image.ndim == 2:
        return image
    return rgb2gray(image)


def box_counting_fd(mask: np.ndarray) -> float:
    if mask.sum() == 0:
        return 0.0
    sizes = np.array([2, 4, 8, 16, 32, 64])
    sizes = sizes[sizes <= min(mask.shape)]
    counts = []
    for size in sizes:
        shape = (
            int(np.ceil(mask.shape[0] / size)),
            int(np.ceil(mask.shape[1] / size)),
        )
        pad_h = shape[0] * size - mask.shape[0]
        pad_w = shape[1] * size - mask.shape[1]
        padded = np.pad(mask, ((0, pad_h), (0, pad_w)), mode="constant")
        blocks = padded.reshape(shape[0], size, shape[1], size)
        blocks = blocks.any(axis=(1, 3))
        counts.append(np.count_nonzero(blocks))
    counts = np.array(counts)
    sizes = sizes[counts > 0]
    counts = counts[counts > 0]
    if len(counts) < 2:
        return 0.0
    coeffs = np.polyfit(np.log(1 / sizes), np.log(counts), 1)
    return float(coeffs[0])


def compute_metrics(
    labels: np.ndarray, intensity_image: np.ndarray, width_units: float
):
    props = regionprops(labels, intensity_image=intensity_image)
    image_width = labels.shape[1]
    pixel_size = (width_units / image_width) if image_width > 0 else 0.0
    rows = []
    for region in props:
        area_px = float(region.area)
        perimeter_px = float(region.perimeter)
        major_px = float(region.major_axis_length) if region.major_axis_length else 0.0
        minor_px = float(region.minor_axis_length) if region.minor_axis_length else 0.0
        area = area_px * (pixel_size**2)
        perimeter = perimeter_px * pixel_size
        major = major_px * pixel_size
        minor = minor_px * pixel_size
        aspect_ratio = major / minor if minor > 0 else 0.0
        circularity = (4 * np.pi * area / (perimeter**2)) if perimeter > 0 else 0.0
        roundness = (4 * area / (np.pi * major**2)) if major > 0 else 0.0
        region_mask = labels == region.label
        region_mask = binary_closing(region_mask)
        entropy_val = float(
            shannon_entropy(region.intensity_image[region.image], base=2)
        )
        fractal_dim = box_counting_fd(region_mask)
        integrated_density = float(region.intensity_image.sum()) * (pixel_size**2)
        ecc_rel = float(region.eccentricity * major)
        rows.append(
            {
                "Label": int(region.label),
                "Area": area,
                "Perimeter": perimeter,
                "Aspect ratio": aspect_ratio,
                "Circularity": circularity,
                "Roundness": roundness,
                "Entropy": entropy_val,
                "Fractal dimension": fractal_dim,
                "Integrated density": integrated_density,
                "Eccentricity (rel width)": ecc_rel,
            }
        )
    metrics_df = pd.DataFrame(rows)
    avg_df = pd.DataFrame()
    if not metrics_df.empty:
        numeric_cols = metrics_df.columns.drop("Label")
        avg_row = {"Metric": "Average"}
        avg_row.update(metrics_df[numeric_cols].mean().to_dict())
        avg_df = pd.DataFrame([avg_row])
    return metrics_df, avg_df


def run_inference(image: np.ndarray, model_name: str, width_units: float):
    if image is None:
        return None, None, None
    if width_units <= 0:
        width_units = 1.0
    model = get_model(model_name)
    image_input = image.copy()
    if model_name == "2D_versatile_fluo":
        image_input = to_gray(image_input)
    image_norm = normalize(image_input, 1, 99.8, axis=(0, 1))
    labels, _ = model.predict_instances(image_norm)
    overlay = render_label(labels, img=image)
    if np.issubdtype(overlay.dtype, np.floating):
        overlay = np.clip(overlay, 0, 1)
    overlay = img_as_ubyte(overlay)
    intensity_image = to_gray(image)
    metrics_df, avg_df = compute_metrics(labels, intensity_image, width_units)
    return overlay, metrics_df, avg_df


with gr.Blocks(title="StarDist 2D Segmentation - HF app by Ram Sevuggan") as demo:
    gr.Markdown("# StarDist 2D Segmentation - HF app by Ram Sevuggan")
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(label="Input image", type="numpy")
            model_dropdown = gr.Dropdown(
                choices=MODEL_NAMES,
                value="2D_versatile_fluo",
                label="Model",
            )
            width_units = gr.Number(
                value=1.0,
                minimum=1e-6,
                label="Image width (units)",
                info="Used for eccentricity relative to image width",
            )
            run_button = gr.Button("Run")
        with gr.Column():
            output_image = gr.Image(label="Overlay", type="numpy")
    metrics_table = gr.Dataframe(
        label="Object metrics",
        interactive=False,
    )
    avg_table = gr.Dataframe(
        label="Average metrics",
        interactive=False,
    )

    run_button.click(
        fn=run_inference,
        inputs=[input_image, model_dropdown, width_units],
        outputs=[output_image, metrics_table, avg_table],
    )

if __name__ == "__main__":
    demo.launch()