"""
Custom image analysis tools — example of SciLink's ``--tools`` feature.

Demonstrates how to expose a set of domain-specific functions directly to
the orchestrator's LLM loop without writing a full custom agent.

Usage with scilink analyze CLI
-------------------------------
  scilink analyze --tools examples/custom_image_tools.py

Then in the chat session the orchestrator can call detect_particles,
measure_morphology, analyze_texture, and measure_coverage directly.

Usage programmatically
-----------------------
  from scilink.agents.exp_agents.analysis_orchestrator import AnalysisOrchestratorAgent
  from examples.custom_image_tools import tool_schemas, create_tool_functions

  orch = AnalysisOrchestratorAgent(...)
  orch.register_tools(tool_schemas, create_tool_functions)

How the orchestrator calls the factory
---------------------------------------
The factory's first parameter is named ``image`` (not ``data_path``), so the
orchestrator loads the current data file as a NumPy array and passes it here.
If you want the orchestrator to pass a file-path string instead, name the
first parameter ``data_path`` (or ``path`` / ``file`` / ``filepath``).

Saving and visualization
------------------------
The factory accepts an optional ``output_dir`` keyword argument.  When present
the orchestrator automatically passes the session's ``results/custom_tools/``
directory so tools can save plots and CSV reports alongside the other session
outputs.  Omit ``output_dir`` entirely if your tools do not need to write files.
"""

import json
import numpy as np
from pathlib import Path


# ── Core analysis functions ──────────────────────────────────────────────────

def _save_figure(fig, output_dir: str, filename: str) -> str:
    """Save a matplotlib figure and return its path (or '' if matplotlib unavailable)."""
    try:
        import matplotlib
        matplotlib.use("Agg")
        path = Path(output_dir) / filename
        fig.savefig(path, dpi=150, bbox_inches="tight")
        import matplotlib.pyplot as plt
        plt.close(fig)
        return str(path)
    except Exception:
        return ""


def detect_particles(image: np.ndarray, threshold: float = None, min_size: int = 10,
                     output_dir: str = None) -> dict:
    """Detect and count particles/objects in a grayscale or RGB image."""
    gray = image if image.ndim == 2 else image.mean(axis=2)
    if threshold is None:
        threshold = float(gray.mean() + gray.std())

    result = {"status": "success", "threshold_used": float(threshold), "min_size_used": min_size}

    try:
        from scipy import ndimage
        binary = gray > threshold
        labeled, n_objects = ndimage.label(binary)
        sizes = ndimage.sum(binary, labeled, range(1, n_objects + 1))
        valid = [s for s in sizes if s >= min_size]
        result.update({
            "particle_count": len(valid),
            "mean_size": float(np.mean(valid)) if valid else 0.0,
        })
        # Visualization
        if output_dir:
            try:
                import matplotlib.pyplot as plt
                fig, axes = plt.subplots(1, 2, figsize=(10, 4))
                axes[0].imshow(gray, cmap="gray")
                axes[0].set_title("Original")
                axes[0].axis("off")
                axes[1].imshow(labeled, cmap="nipy_spectral")
                axes[1].set_title(f"Detected particles: {len(valid)}")
                axes[1].axis("off")
                fig.tight_layout()
                saved = _save_figure(fig, output_dir, "detect_particles.png")
                if saved:
                    result["plot_path"] = saved
            except Exception:
                pass
    except ImportError:
        binary = (gray > threshold).astype(np.uint8)
        result.update({
            "particle_count": int(binary.sum() / max(min_size, 1)),
            "note": "scipy not available — rough estimate only",
        })
    return result


def measure_morphology(image: np.ndarray, threshold: float = None, output_dir: str = None) -> dict:
    """Compute basic morphological statistics (area fraction, perimeter proxy)."""
    gray = image if image.ndim == 2 else image.mean(axis=2)
    if threshold is None:
        threshold = float(gray.mean())
    binary = gray > threshold
    area_fraction = float(binary.sum()) / binary.size
    shifted_h = np.roll(binary, 1, axis=0) != binary
    shifted_v = np.roll(binary, 1, axis=1) != binary
    perimeter_proxy = int((shifted_h | shifted_v).sum())
    result = {
        "status": "success",
        "area_fraction": round(area_fraction, 4),
        "perimeter_proxy_px": perimeter_proxy,
        "threshold_used": float(threshold),
        "image_shape": list(image.shape),
    }
    if output_dir:
        try:
            import matplotlib.pyplot as plt
            fig, axes = plt.subplots(1, 2, figsize=(10, 4))
            axes[0].imshow(gray, cmap="gray")
            axes[0].set_title("Grayscale")
            axes[0].axis("off")
            axes[1].imshow(binary, cmap="gray")
            axes[1].set_title(f"Binary mask (area fraction {area_fraction:.2%})")
            axes[1].axis("off")
            fig.tight_layout()
            saved = _save_figure(fig, output_dir, "measure_morphology.png")
            if saved:
                result["plot_path"] = saved
        except Exception:
            pass
    return result


def analyze_texture(image: np.ndarray, output_dir: str = None) -> dict:
    """Compute texture descriptors (contrast, energy, mean, std) from a grayscale image."""
    gray = image if image.ndim == 2 else image.mean(axis=2)
    gray_norm = gray / (gray.max() + 1e-8)
    contrast = float(gray_norm.max() - gray_norm.min())
    energy = float((gray_norm ** 2).mean())
    result = {
        "status": "success",
        "mean_intensity": round(float(gray.mean()), 4),
        "std_intensity": round(float(gray.std()), 4),
        "contrast": round(contrast, 4),
        "energy": round(energy, 6),
        "image_shape": list(image.shape),
    }
    if output_dir:
        try:
            import matplotlib.pyplot as plt
            fig, axes = plt.subplots(1, 2, figsize=(10, 4))
            axes[0].imshow(gray, cmap="gray")
            axes[0].set_title("Image")
            axes[0].axis("off")
            axes[1].hist(gray.ravel(), bins=64, color="steelblue", edgecolor="none")
            axes[1].set_title("Intensity histogram")
            axes[1].set_xlabel("Intensity")
            axes[1].set_ylabel("Count")
            fig.tight_layout()
            saved = _save_figure(fig, output_dir, "analyze_texture.png")
            if saved:
                result["plot_path"] = saved
        except Exception:
            pass
    return result


def measure_coverage(image: np.ndarray, threshold: float = None, output_dir: str = None) -> dict:
    """Measure coverage (fraction of pixels above threshold) across the image."""
    gray = image if image.ndim == 2 else image.mean(axis=2)
    if threshold is None:
        threshold = float(gray.mean())
    coverage = float((gray > threshold).sum()) / gray.size
    result = {
        "status": "success",
        "coverage_fraction": round(coverage, 4),
        "coverage_percent": round(coverage * 100, 2),
        "threshold_used": float(threshold),
    }
    if output_dir:
        try:
            import matplotlib.pyplot as plt
            mask = gray > threshold
            fig, axes = plt.subplots(1, 2, figsize=(10, 4))
            axes[0].imshow(gray, cmap="gray")
            axes[0].set_title("Image")
            axes[0].axis("off")
            axes[1].imshow(np.where(mask, gray, gray * 0.3), cmap="inferno")
            axes[1].set_title(f"Coverage: {coverage:.1%}")
            axes[1].axis("off")
            fig.tight_layout()
            saved = _save_figure(fig, output_dir, "measure_coverage.png")
            if saved:
                result["plot_path"] = saved
        except Exception:
            pass
    return result


# ── OpenAI-format tool schemas ───────────────────────────────────────────────
# The orchestrator reads this list to build tool schemas for the LLM.

tool_schemas = [
    {
        "type": "function",
        "function": {
            "name": "detect_particles",
            "description": (
                "Detect and count discrete particles or objects in a microscopy image. "
                "Returns particle count, mean size, and the threshold used."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "threshold": {
                        "type": "number",
                        "description": (
                            "Intensity threshold for particle detection "
                            "(default: mean + 1 std deviation)."
                        ),
                    },
                    "min_size": {
                        "type": "integer",
                        "description": "Minimum particle area in pixels (default: 10).",
                    },
                },
                "required": [],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "measure_morphology",
            "description": (
                "Compute basic morphological statistics for the current image: "
                "area fraction covered by foreground pixels and a perimeter proxy."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "threshold": {
                        "type": "number",
                        "description": (
                            "Intensity threshold to separate foreground from background "
                            "(default: mean intensity)."
                        ),
                    },
                },
                "required": [],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "analyze_texture",
            "description": (
                "Compute texture descriptors (mean, std, contrast, energy) "
                "for the current image."
            ),
            "parameters": {
                "type": "object",
                "properties": {},
                "required": [],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "measure_coverage",
            "description": (
                "Measure the fraction (and percentage) of pixels above a threshold — "
                "useful for estimating coating or deposit coverage."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "threshold": {
                        "type": "number",
                        "description": (
                            "Intensity threshold (default: mean intensity)."
                        ),
                    },
                },
                "required": [],
            },
        },
    },
]


# ── Factory function ─────────────────────────────────────────────────────────
# The orchestrator calls this with the loaded image array.
# Because the first parameter is named ``image`` (not ``data_path``),
# the orchestrator loads the file first and passes the NumPy array here.

def create_tool_functions(image: np.ndarray, output_dir: str = None) -> dict:
    """Bind the loaded image (and optional output directory) to the analysis functions."""
    return {
        "detect_particles": lambda threshold=None, min_size=10:
            detect_particles(image, threshold, min_size, output_dir=output_dir),
        "measure_morphology": lambda threshold=None:
            measure_morphology(image, threshold, output_dir=output_dir),
        "analyze_texture": lambda:
            analyze_texture(image, output_dir=output_dir),
        "measure_coverage": lambda threshold=None:
            measure_coverage(image, threshold, output_dir=output_dir),
    }
