"""Public helper module.

Conceptually, this module groups the lightweight Allen CCF wireframe and
probe-track plotting helpers reused across brain-view figures.

It exists as a separate unit so coordinate conventions, reference markers, and
axis setup stay consistent across 3D and dorsal-view plots.

It connects saved Allen grid data and probe trajectories to the compact
plotting primitives consumed by brainmap-style notebooks.
"""

from __future__ import annotations

from pathlib import Path
from typing import Any, Mapping, Optional, Sequence, Tuple

import numpy as np

from .core import to_1d
from .paths import first_existing_path, resolve_repo_root


def brain_grid_data_path(repo_root: Optional[Path] = None) -> Path:
    """Return the canonical path to `brainGridData.npy`."""
    root = resolve_repo_root(repo_root or Path.cwd())
    return first_existing_path(
        [
            root / "functions" / "allenCCF" / "Browsing Functions" / "brainGridData.npy",
            root / "version_matlab" / "functions" / "allenCCF" / "Browsing Functions" / "brainGridData.npy",
        ],
        missing_message="Could not find Allen CCF brainGridData.npy in the expected project locations.",
    )


def load_brain_grid_data(path: Optional[Path] = None, repo_root: Optional[Path] = None) -> np.ndarray:
    """Load Allen CCF wireframe vertices and restore NaNs from zero placeholders."""
    npy_path = Path(path) if path is not None else brain_grid_data_path(repo_root=repo_root)
    brain_grid = np.asarray(np.load(npy_path), dtype=float)
    brain_grid[np.sum(brain_grid, axis=1) == 0.0, :] = np.nan
    return brain_grid


def allen_ccf_bregma() -> np.ndarray:
    """Return bregma in Allen CCF 10um coordinates as `(AP, DV, LR)`."""
    return np.asarray([[540.0, 0.0, 570.0]], dtype=float)


def _finite_xyz(brain_grid: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    x = np.asarray(brain_grid[:, 0], dtype=float)
    y = np.asarray(brain_grid[:, 1], dtype=float)
    z = np.asarray(brain_grid[:, 2], dtype=float)
    finite = np.isfinite(x) & np.isfinite(y) & np.isfinite(z)
    return x[finite], y[finite], z[finite]


def _set_equal_box_aspect(ax: Any, x: np.ndarray, y: np.ndarray, z: np.ndarray) -> None:
    if x.size == 0 or y.size == 0 or z.size == 0:
        return
    ax.set_xlim(float(np.nanmin(x)), float(np.nanmax(x)))
    ax.set_ylim(float(np.nanmin(y)), float(np.nanmax(y)))
    ax.set_zlim(float(np.nanmin(z)), float(np.nanmax(z)))
    if hasattr(ax, "set_box_aspect"):
        ax.set_box_aspect(
            (
                float(np.nanmax(x) - np.nanmin(x)),
                float(np.nanmax(y) - np.nanmin(y)),
                float(np.nanmax(z) - np.nanmin(z)),
            )
        )


def plot_brain_grid(
    ax: Any,
    brain_grid: Optional[np.ndarray] = None,
    *,
    black_brain: bool = False,
    linewidth: float = 0.8,
) -> Any:
    """Plot the Allen CCF wireframe using the stored `(x, y, z)` convention."""
    grid = load_brain_grid_data() if brain_grid is None else np.asarray(brain_grid, dtype=float)
    x, y, z = _finite_xyz(grid)
    color = (0.5, 0.5, 0.5, 0.3) if black_brain else (0.0, 0.0, 0.0, 0.3)
    line = ax.plot3D(x, y, z, color=color, linewidth=linewidth)[0]
    _set_equal_box_aspect(ax, x, y, z)
    ax.invert_zaxis()
    ax.set_axis_off()
    return line


def plot_brain_grid_top(
    ax: Any,
    brain_grid: Optional[np.ndarray] = None,
    *,
    black_brain: bool = False,
    linewidth: float = 0.75,
) -> Any:
    """Plot a dorsal 2D projection of the Allen wireframe in `(AP, ML)`."""
    grid = load_brain_grid_data() if brain_grid is None else np.asarray(brain_grid, dtype=float)
    x = np.asarray(grid[:, 0], dtype=float)
    y = np.asarray(grid[:, 1], dtype=float)
    finite = np.isfinite(x) & np.isfinite(y)
    color = (0.5, 0.5, 0.5, 0.3) if black_brain else (0.0, 0.0, 0.0, 0.18)
    line = ax.plot(x[finite], y[finite], color=color, linewidth=linewidth)[0]
    ax.set_aspect("equal", adjustable="box")
    ax.set_axis_off()
    return line


def plot_bregma_reference(
    ax: Any,
    *,
    color: str = "k",
    marker_size: float = 15.0,
    linewidth: float = 3.0,
) -> None:
    """Plot bregma and the small AP/ML reference axes used in brain-view figures."""
    bregma = allen_ccf_bregma()
    ax.plot3D(bregma[:, 0], bregma[:, 2], bregma[:, 1], "+", color=color, markersize=marker_size, linewidth=linewidth)
    ax.plot3D([740, 690], [570, 570], [0, 0], "-", color=color, linewidth=max(1.0, linewidth * 0.6))
    ax.plot3D([740, 740], [570, 520], [0, 0], "-", color=color, linewidth=max(1.0, linewidth * 0.6))


def plot_bregma_reference_top(
    ax: Any,
    *,
    color: str = "k",
    marker_size: float = 15.0,
    linewidth: float = 2.5,
) -> None:
    """Plot the dorsal-view bregma marker and AP/ML mini-axes."""
    bregma = allen_ccf_bregma()[0]
    ax.plot(bregma[0], bregma[2], "+", color=color, markersize=marker_size, markeredgewidth=max(1.0, linewidth))
    ax.plot([740, 690], [570, 570], "-", color=color, linewidth=max(1.0, linewidth * 0.52))
    ax.plot([740, 740], [570, 520], "-", color=color, linewidth=max(1.0, linewidth * 0.52))


def plot_probe_track(
    ax: Any,
    ap: Sequence[float],
    dv: Sequence[float],
    ml: Sequence[float],
    *,
    color: Any,
    linewidth: float = 2.0,
) -> Any:
    """Plot a probe track in the displayed `(AP, ML, DV)` convention."""
    ap_arr = np.asarray(ap, dtype=float).reshape(-1)
    dv_arr = np.asarray(dv, dtype=float).reshape(-1)
    ml_arr = np.asarray(ml, dtype=float).reshape(-1)
    n = min(ap_arr.size, dv_arr.size, ml_arr.size)
    if n == 0:
        raise ValueError("Probe track arrays must contain at least one point.")
    line = ax.plot3D(
        ap_arr[:n],
        ml_arr[:n],
        dv_arr[:n],
        color=color,
        linewidth=linewidth,
        solid_capstyle="round",
    )[0]
    return line


def probe_mean_point(entry: Mapping[str, Any]) -> np.ndarray:
    """Return the mean `(AP, DV, ML)` electrode location for one probe entry.

    Parameters
    ----------
    entry:
        Mapping containing probe electrode coordinates such as `elec_ccf_ap`,
        `elec_ccf_dv`, and `elec_ccf_ml`.

    Returns
    -------
    np.ndarray
        Three-element vector containing the mean `(AP, DV, ML)` location. When
        one coordinate axis is missing, the function returns `np.nan` values.
    """
    ap = to_1d(entry.get("elec_ccf_ap", []), float)
    dv = to_1d(entry.get("elec_ccf_dv", []), float)
    ml = to_1d(entry.get("elec_ccf_ml", []), float)
    if ap.size == 0 or dv.size == 0 or ml.size == 0:
        return np.full((3,), np.nan, dtype=float)
    n = min(ap.size, dv.size, ml.size)
    pts = np.column_stack([ap[:n], dv[:n], ml[:n]])
    return np.mean(pts, axis=0)


def setup_brain_3d_axis(
    ax: Any,
    *,
    brain_grid: Optional[np.ndarray] = None,
    view: Optional[Tuple[float, float]] = None,
    black_brain: bool = False,
    linewidth: float = 0.8,
    camera_distance: Optional[float] = 7.5,
) -> None:
    """Configure a 3D axis with the Allen wireframe and reference view.

    Parameters
    ----------
    ax:
        Target 3D matplotlib axis.
    brain_grid:
        Optional preloaded wireframe vertices.
    view:
        Optional `(azimuth, elevation)` tuple passed to `ax.view_init`.
    black_brain:
        Whether to draw the wireframe in the dark overlay variant.
    linewidth:
        Wireframe line width.
    camera_distance:
        Optional camera distance forwarded to Matplotlib when available.
    """
    plot_brain_grid(ax, brain_grid=brain_grid, black_brain=black_brain, linewidth=linewidth)
    if view is not None:
        azim, elev = float(view[0]), float(view[1])
        ax.view_init(elev=elev, azim=azim)
    if camera_distance is not None:
        try:
            ax.dist = float(camera_distance)
        except Exception:
            pass


def setup_brain_top_axis(
    ax: Any,
    *,
    brain_grid: Optional[np.ndarray] = None,
    black_brain: bool = False,
    linewidth: float = 0.75,
    xlim: Optional[Tuple[float, float]] = None,
    ylim: Optional[Tuple[float, float]] = None,
) -> None:
    """Configure a dorsal-view 2D axis with the Allen wireframe."""
    grid = load_brain_grid_data() if brain_grid is None else np.asarray(brain_grid, dtype=float)
    plot_brain_grid_top(ax, brain_grid=grid, black_brain=black_brain, linewidth=linewidth)
    if xlim is not None:
        ax.set_xlim(float(xlim[0]), float(xlim[1]))
    if ylim is not None:
        ax.set_ylim(float(ylim[0]), float(ylim[1]))