"""Lightweight atom finding utilities for atomic-resolution STEM images.

Core algorithms extracted from atomap (Nord et al., 2017) and
reimplemented with numpy/scipy/scikit-image only — no HyperSpy, numba,
or other heavy dependencies.

Provides four functions for multi-sublattice analysis:
  - detect_atoms: peak detection + 2D Gaussian refinement
  - find_zone_axes: lattice translation vector detection
  - find_missing_atoms: predict fractional-site positions
  - subtract_atoms: remove fitted Gaussians to reveal weaker sublattices
"""

from __future__ import annotations

import math
from typing import Optional

import numpy as np
from scipy.ndimage import gaussian_filter, label
from scipy.optimize import curve_fit
from scipy.spatial import cKDTree
from skimage.feature import peak_local_max


# ---------------------------------------------------------------------------
# 2D Gaussian model (matches atomap/external/gaussian2d.py convention)
# ---------------------------------------------------------------------------

def _gaussian2d(coords, x0, y0, A, sigma_x, sigma_y, theta, offset):
    """Normalised rotated 2D Gaussian.

    Parameters match atomap: ``A`` is integrated volume, peak height is
    ``A / (2π·σx·σy)``.
    """
    x, y = coords
    cos_t = math.cos(theta)
    sin_t = math.sin(theta)
    cos2 = cos_t ** 2
    sin2 = sin_t ** 2
    sin_2t = math.sin(2 * theta)

    sx2 = sigma_x ** 2
    sy2 = sigma_y ** 2

    a = cos2 / (2 * sx2) + sin2 / (2 * sy2)
    b = -sin_2t / (4 * sx2) + sin_2t / (4 * sy2)
    c = sin2 / (2 * sx2) + cos2 / (2 * sy2)

    dx = x - x0
    dy = y - y0

    norm = A / (2 * math.pi * sigma_x * sigma_y)
    return (norm * np.exp(-(a * dx**2 + 2 * b * dx * dy + c * dy**2))
            + offset).ravel()


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def detect_atoms(
    image: np.ndarray,
    separation: int,
    threshold_rel: float = 0.02,
    refine: bool = True,
    percent_to_nn: float = 0.40,
) -> dict:
    """Detect atom column positions with optional 2D Gaussian refinement.

    Args:
        image: 2D grayscale array (HAADF: bright atoms on dark background).
        separation: Minimum atom spacing in **pixels**.
        threshold_rel: Relative peak threshold for ``peak_local_max``.
        refine: Fit a 2D Gaussian per peak for sub-pixel precision.
        percent_to_nn: Gaussian mask radius as fraction of NN distance.

    Returns:
        dict with keys ``positions`` (N,2 as x,y where x=col y=row),
        ``sigma_x``, ``sigma_y``, ``amplitude``, ``rotation`` (all N,).
    """
    img = image.astype(np.float64)

    # Background subtraction
    bg = gaussian_filter(img, sigma=separation)
    filtered = np.clip(img - bg, 0, None)

    # Peak detection — returns (row, col)
    peaks_rc = peak_local_max(
        filtered,
        min_distance=int(separation),
        threshold_rel=threshold_rel,
    )

    if len(peaks_rc) == 0:
        return _empty_result()

    # Flip to (x, y) = (col, row) to match atomap convention
    positions = np.fliplr(peaks_rc).astype(np.float64)

    if not refine or len(positions) < 2:
        return {
            "positions": positions,
            "sigma_x": None,
            "sigma_y": None,
            "amplitude": None,
            "rotation": None,
        }

    # --- 2D Gaussian refinement ---
    tree = cKDTree(positions)
    nn_dists, _ = tree.query(positions, k=2)
    nn_dist = nn_dists[:, 1]  # distance to closest neighbor

    H, W = img.shape
    sigma_x_arr = np.full(len(positions), separation / 4.0)
    sigma_y_arr = np.full(len(positions), separation / 4.0)
    amplitude_arr = np.zeros(len(positions))
    rotation_arr = np.zeros(len(positions))

    for i, (x0, y0) in enumerate(positions):
        mask_r = max(int(nn_dist[i] * percent_to_nn), 2)

        fitted = False
        ptn = percent_to_nn
        for attempt in range(6):
            r = max(int(nn_dist[i] * ptn), 2)

            # Crop bounds (row = y, col = x)
            r0 = max(0, int(y0) - r)
            r1 = min(H, int(y0) + r + 1)
            c0 = max(0, int(x0) - r)
            c1 = min(W, int(x0) + r + 1)

            crop = img[r0:r1, c0:c1]
            if crop.size < 7:
                ptn *= 0.95
                continue

            # Build coordinate grids (x = col index, y = row index)
            yy, xx = np.mgrid[r0:r1, c0:c1]
            coords = (xx.ravel().astype(np.float64),
                      yy.ravel().astype(np.float64))

            bg_val = np.percentile(crop, 3)
            peak_val = crop.max() - bg_val
            s0 = separation / 4.0

            p0 = [x0, y0, peak_val * 2 * math.pi * s0 * s0,
                  s0, s0, 0.01, bg_val]
            bounds_lo = [c0, r0, 0, 0.5, 0.5, -math.pi, -np.inf]
            bounds_hi = [c1, r1, np.inf, r * 2, r * 2, math.pi, np.inf]

            try:
                popt, _ = curve_fit(
                    _gaussian2d, coords, crop.ravel().astype(np.float64),
                    p0=p0, bounds=(bounds_lo, bounds_hi), maxfev=2000,
                )
                fx, fy, fA, fsx, fsy, ftheta, foff = popt

                # Validation (matches atomap criteria)
                if abs(fx - x0) > r or abs(fy - y0) > r:
                    raise ValueError("center outside mask")
                ratio = max(fsx, fsy) / max(min(fsx, fsy), 0.1)
                if ratio > 4:
                    raise ValueError("sigma ratio > 4")

                positions[i] = [fx, fy]
                sigma_x_arr[i] = fsx
                sigma_y_arr[i] = fsy
                amplitude_arr[i] = fA
                rotation_arr[i] = ftheta % math.pi
                fitted = True
                break
            except Exception:
                ptn *= 0.95

        if not fitted:
            # Fallback: center of mass
            crop = img[max(0, int(y0) - mask_r):min(H, int(y0) + mask_r + 1),
                       max(0, int(x0) - mask_r):min(W, int(x0) + mask_r + 1)]
            if crop.sum() > 0:
                yy, xx = np.mgrid[0:crop.shape[0], 0:crop.shape[1]]
                total = crop.sum()
                cx = (xx * crop).sum() / total + max(0, int(x0) - mask_r)
                cy = (yy * crop).sum() / total + max(0, int(y0) - mask_r)
                positions[i] = [cx, cy]
            amplitude_arr[i] = img[int(round(y0)), int(round(x0))]

    return {
        "positions": positions,
        "sigma_x": sigma_x_arr,
        "sigma_y": sigma_y_arr,
        "amplitude": amplitude_arr,
        "rotation": rotation_arr,
    }


def find_zone_axes(
    positions: np.ndarray,
    n_neighbors: int = 9,
    distance_tolerance: Optional[float] = None,
) -> list:
    """Detect lattice translation vectors from atom positions.

    Args:
        positions: (N, 2) array of atom positions (x, y).
        n_neighbors: Number of nearest neighbors to examine per atom.
        distance_tolerance: Clustering tolerance in pixels.
            Default: median NN distance / 3.

    Returns:
        List of (dx, dy) tuples — unique lattice vectors, shortest first.
    """
    if len(positions) < 3:
        return []

    tree = cKDTree(positions)
    k = min(n_neighbors + 1, len(positions))
    dists, indices = tree.query(positions, k=k)

    # Median nearest-neighbor distance
    nn1 = dists[:, 1]
    med_nn = float(np.median(nn1))

    if distance_tolerance is None:
        distance_tolerance = med_nn / 3.0

    # Collect all displacement vectors (neighbor - atom)
    all_vectors = []
    for i in range(len(positions)):
        for j in range(1, k):
            d = positions[indices[i, j]] - positions[i]
            all_vectors.append(d)
    all_vectors = np.array(all_vectors)

    # Cluster via 2D histogram
    max_range = med_nn * (n_neighbors ** 0.5) * 1.5
    bin_size = distance_tolerance
    n_bins = max(int(2 * max_range / bin_size), 10)
    bins = np.linspace(-max_range, max_range, n_bins + 1)

    hist, xedges, yedges = np.histogram2d(
        all_vectors[:, 0], all_vectors[:, 1], bins=[bins, bins],
    )

    # Label connected regions above threshold
    threshold_count = max(len(positions) * 0.15, 3)
    labeled, n_clusters = label(hist >= threshold_count)

    candidates = []
    for c in range(1, n_clusters + 1):
        mask = labeled == c
        ys, xs = np.nonzero(mask)
        # Weighted centroid of the cluster
        weights = hist[mask]
        cx = np.average((xedges[xs] + xedges[xs + 1]) / 2, weights=weights)
        cy = np.average((yedges[ys] + yedges[ys + 1]) / 2, weights=weights)
        length = math.hypot(cx, cy)
        if length > med_nn * 0.3:  # skip near-zero vectors
            candidates.append((cx, cy, length))

    # Sort by length
    candidates.sort(key=lambda v: v[2])

    # Remove parallel/antiparallel duplicates and integer multiples
    unique = []
    for cx, cy, length in candidates:
        is_dup = False
        for ux, uy, _ in unique:
            for n in range(-4, 6):
                if n == 0:
                    continue
                ref_x, ref_y = ux * n, uy * n
                if math.hypot(cx - ref_x, cy - ref_y) < distance_tolerance:
                    is_dup = True
                    break
            if is_dup:
                break
        if not is_dup:
            # Canonical direction: prefer positive first nonzero component
            if cx < -1e-6 or (abs(cx) < 1e-6 and cy < -1e-6):
                cx, cy = -cx, -cy
            unique.append((round(cx, 2), round(cy, 2),
                           math.hypot(cx, cy)))

    return [(vx, vy) for vx, vy, _ in unique]


def find_missing_atoms(
    positions: np.ndarray,
    zone_vector: tuple,
    fraction: float = 0.5,
    min_distance: float = 3.0,
) -> np.ndarray:
    """Predict atom positions at fractional lattice sites.

    Args:
        positions: (N, 2) array of detected atoms (x, y).
        zone_vector: (dx, dy) lattice vector from :func:`find_zone_axes`.
        fraction: Fractional position along the vector (0.5 = midpoint).
        min_distance: Minimum distance from existing atoms to keep.

    Returns:
        (M, 2) array of predicted new positions.
    """
    if len(positions) < 2:
        return np.empty((0, 2))

    zv = np.array(zone_vector, dtype=np.float64)
    zv_len = np.linalg.norm(zv)
    if zv_len < 1e-6:
        return np.empty((0, 2))

    tree = cKDTree(positions)
    tolerance = zv_len * 0.5  # neighbor must be within 50% of expected

    new_positions = []
    for p in positions:
        expected = p + zv
        dist, idx = tree.query(expected)
        if dist < tolerance:
            neighbor = positions[idx]
            interp = p * (1 - fraction) + neighbor * fraction
            new_positions.append(interp)

    if not new_positions:
        return np.empty((0, 2))

    new_arr = np.array(new_positions)

    # Deduplicate within new positions
    if len(new_arr) > 1:
        new_tree = cKDTree(new_arr)
        pairs = new_tree.query_pairs(min_distance * 0.5)
        to_remove = set()
        for i, j in pairs:
            to_remove.add(max(i, j))
        if to_remove:
            keep = [i for i in range(len(new_arr)) if i not in to_remove]
            new_arr = new_arr[keep]

    # Remove positions too close to existing atoms
    if len(new_arr) > 0:
        dists, _ = tree.query(new_arr)
        new_arr = new_arr[dists >= min_distance]

    return new_arr


def subtract_atoms(
    image: np.ndarray,
    positions: np.ndarray,
    sigma_x: np.ndarray,
    sigma_y: np.ndarray,
    amplitude: np.ndarray,
    rotation: Optional[np.ndarray] = None,
) -> np.ndarray:
    """Subtract fitted 2D Gaussians from image.

    Args:
        image: 2D array.
        positions: (N, 2) atom positions (x, y) from :func:`detect_atoms`.
        sigma_x, sigma_y, amplitude: Per-atom Gaussian parameters.
        rotation: Per-atom rotation (radians). Default 0 for all.

    Returns:
        Residual image (clipped ≥ 0).
    """
    img = image.astype(np.float64)
    model = np.zeros_like(img)
    H, W = img.shape

    if rotation is None:
        rotation = np.zeros(len(positions))

    X, Y = np.meshgrid(np.arange(W, dtype=np.float64),
                        np.arange(H, dtype=np.float64))

    for i, (x0, y0) in enumerate(positions):
        sx = max(float(sigma_x[i]), 0.5)
        sy = max(float(sigma_y[i]), 0.5)
        A = float(amplitude[i])
        theta = float(rotation[i])

        r = int(5 * max(sx, sy))
        r0 = max(0, int(y0) - r)
        r1 = min(H, int(y0) + r + 1)
        c0 = max(0, int(x0) - r)
        c1 = min(W, int(x0) + r + 1)

        if r1 <= r0 or c1 <= c0:
            continue

        xc = X[r0:r1, c0:c1]
        yc = Y[r0:r1, c0:c1]

        # Evaluate normalised rotated 2D Gaussian (no offset)
        cos_t = math.cos(theta)
        sin_t = math.sin(theta)
        sx2 = sx ** 2
        sy2 = sy ** 2
        a = cos_t**2 / (2*sx2) + sin_t**2 / (2*sy2)
        b = -math.sin(2*theta) / (4*sx2) + math.sin(2*theta) / (4*sy2)
        c = sin_t**2 / (2*sx2) + cos_t**2 / (2*sy2)

        dx = xc - x0
        dy = yc - y0
        norm = A / (2 * math.pi * sx * sy)
        model[r0:r1, c0:c1] += norm * np.exp(
            -(a * dx**2 + 2 * b * dx * dy + c * dy**2)
        )

    return np.clip(img - model, 0, None)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _empty_result():
    return {
        "positions": np.empty((0, 2)),
        "sigma_x": None,
        "sigma_y": None,
        "amplitude": None,
        "rotation": None,
    }
