"""
Simple custom tool for testing the UI upload feature.

Provides two tools:
- basic_stats: Compute mean, median, min, max of loaded data.
- pixel_histogram: Bin pixel/value intensities and return counts.

Upload this .py file via the "Tools & Agents" tab in the UI.
"""

import numpy as np


def basic_stats(data: np.ndarray) -> dict:
    """Return basic descriptive statistics for the data."""
    flat = data.ravel().astype(float)
    return {
        "status": "success",
        "mean": round(float(np.mean(flat)), 4),
        "median": round(float(np.median(flat)), 4),
        "std": round(float(np.std(flat)), 4),
        "min": round(float(np.min(flat)), 4),
        "max": round(float(np.max(flat)), 4),
        "shape": list(data.shape),
        "dtype": str(data.dtype),
    }


def pixel_histogram(data: np.ndarray, n_bins: int = 32) -> dict:
    """Bin values into a histogram and return counts per bin."""
    flat = data.ravel().astype(float)
    counts, edges = np.histogram(flat, bins=n_bins)
    return {
        "status": "success",
        "n_bins": n_bins,
        "bin_edges": [round(float(e), 4) for e in edges],
        "counts": counts.tolist(),
        "total_values": int(flat.size),
    }


# ── OpenAI-format tool schemas ──────────────────────────────────────────────

tool_schemas = [
    {
        "type": "function",
        "function": {
            "name": "basic_stats",
            "description": (
                "Compute basic descriptive statistics (mean, median, std, min, max) "
                "for the currently loaded data."
            ),
            "parameters": {
                "type": "object",
                "properties": {},
                "required": [],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "pixel_histogram",
            "description": (
                "Compute a histogram of pixel/value intensities for the loaded data."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "n_bins": {
                        "type": "integer",
                        "description": "Number of histogram bins (default: 32).",
                    },
                },
                "required": [],
            },
        },
    },
]


# ── Factory function ─────────────────────────────────────────────────────────

def create_tool_functions(data: np.ndarray) -> dict:
    """Bind the loaded data to the tool functions."""
    return {
        "basic_stats": lambda: basic_stats(data),
        "pixel_histogram": lambda n_bins=32: pixel_histogram(data, n_bins),
    }
