from __future__ import annotations
from typing import Tuple, Optional, Dict

import weakref
import numpy as np
from scipy.signal import fftconvolve
from scipy.spatial import cKDTree

from .core import PPP, Window, MaskWindow, RectangleWindow, Image



def _ensure_increasing_r(r: np.ndarray) -> np.ndarray:
    """
    Validate a 1D array of strictly increasing positive radii.

    :param r: Radii array.
    :type r: numpy.ndarray
    :return: Validated radii array.
    :rtype: numpy.ndarray
    :raises ValueError: If not 1D, empty, non-finite, non-positive, or not strictly increasing.
    """
    r = np.asarray(r, float)
    if r.ndim != 1 or r.size == 0:
        raise ValueError("r must be a non-empty 1D array of radii.")
    if not np.all(np.isfinite(r)) or np.any(r <= 0):
        raise ValueError("r must be positive and finite.")
    if np.any(np.diff(r) <= 0):
        raise ValueError("r must be strictly increasing.")
    return r


def _pairs_within(xy: np.ndarray, r_max: float) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Enumerate unordered pairs of points with separation < r_max.

    :param xy: Point coordinates, shape (n, 2).
    :type xy: numpy.ndarray
    :param r_max: Maximum pair separation.
    :type r_max: float
    :return: Tuple (i, j, d) with indices and distances for each qualifying pair.
    :rtype: Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]
    """
    tree = cKDTree(xy)
    try:
        pairs = tree.query_pairs(r_max, output_type="ndarray")
    except TypeError:
        pairs_set = tree.query_pairs(r_max)
        if not pairs_set:
            return np.empty(0, int), np.empty(0, int), np.empty(0, float)
        pairs = np.asarray(list(pairs_set), dtype=np.intp).reshape(-1, 2)
    if pairs.size == 0:
        return np.empty(0, int), np.empty(0, int), np.empty(0, float)
    i = pairs[:, 0]
    j = pairs[:, 1]
    diff = xy[j] - xy[i]
    d = np.sqrt((diff * diff).sum(axis=1))
    return i, j, d


def _plot_stat(x: np.ndarray, y: np.ndarray, y_label: str, csr: np.ndarray, show: bool) -> None:
    """
    Plot a statistic against radius with its CSR baseline.

    :param x: Radii (or bin centres).
    :type x: numpy.ndarray
    :param y: Estimated statistic values.
    :type y: numpy.ndarray
    :param y_label: Y-axis label (e.g. "K(r)").
    :type y_label: str
    :param csr: CSR baseline values at x.
    :type csr: numpy.ndarray
    :param show: If True, display the figure; otherwise no-op.
    :type show: bool
    :raises RuntimeError: If matplotlib is not available and show is True.
    """
    if not show:
        return
    try:
        import matplotlib.pyplot as plt
    except Exception as exc:
        raise RuntimeError("Plotting requires matplotlib. Install it to use showplot=True.") from exc

    titles = {
        "K(r)": "Ripley K",
        "L(r)": "Besag L",
        "g(r)": "Pair correlation g(r)",
        "G(r)": "Nearest-neighbour G(r)",
        "F(r)": "Empty-space F(r)",
        "J(r)": "J-function",
    }
    title = titles.get(y_label, f"{y_label} vs r")

    fig = plt.figure()
    ax = fig.add_subplot(111)

    mask = np.isfinite(y)
    if mask.any():
        ax.plot(x[mask], y[mask], label="estimate")

    ax.plot(x, csr, linestyle="--", color="k", label="CSR")
    ax.set_xlabel("r")
    ax.set_ylabel(y_label)
    ax.set_title(title)
    ax.legend()
    fig.tight_layout()
    plt.show()


class _MaskTransWeights:
    """
    Overlap fractions for all integer pixel shifts via FFT, with bilinear interpolation
    to arbitrary sub-pixel world shifts. Supports fast vectorised queries.
    """

    def __init__(self, win: MaskWindow) -> None:
        """
        Precompute correlation map and normalise by area.

        :param win: MaskWindow instance.
        :type win: MaskWindow
        """
        m = win.mask.astype(float)
        corr = fftconvolve(m, m[::-1, ::-1], mode="full")
        denom = float(np.count_nonzero(m))
        self._frac = (corr / denom) if denom > 0.0 else np.zeros_like(corr, float)
        self._H, self._W = self._frac.shape
        self._ny, self._nx = m.shape
        self._x0 = self._nx - 1
        self._y0 = self._ny - 1
        self._sx, self._sy = win.spacing

    def _gather(self, ix: np.ndarray, iy: np.ndarray) -> np.ndarray:
        """
        Fetch overlap fractions at integer-offset indices.

        :param ix: Integer x-offsets (relative to zero-shift).
        :type ix: numpy.ndarray
        :param iy: Integer y-offsets (relative to zero-shift).
        :type iy: numpy.ndarray
        :return: Overlap fractions.
        :rtype: numpy.ndarray
        """
        ix = ix + self._x0
        iy = iy + self._y0
        ix = np.clip(ix, 0, self._W - 1, out=ix)
        iy = np.clip(iy, 0, self._H - 1, out=iy)
        return self._frac[iy, ix]

    def frac_many(self, delta: np.ndarray) -> np.ndarray:
        """
        Bilinear interpolation of overlap fraction at world shifts.

        :param delta: Array of shifts (m, 2) in world units.
        :type delta: numpy.ndarray
        :return: Overlap fractions (m,).
        :rtype: numpy.ndarray
        """
        dx = delta[:, 0] / self._sx
        dy = delta[:, 1] / self._sy
        mx = np.floor(dx).astype(np.int64); fx = dx - mx
        my = np.floor(dy).astype(np.int64); fy = dy - my
        w00 = (1.0 - fx) * (1.0 - fy)
        w10 = fx * (1.0 - fy)
        w01 = (1.0 - fx) * fy
        w11 = fx * fy
        v00 = self._gather(mx,     my)
        v10 = self._gather(mx + 1, my)
        v01 = self._gather(mx,     my + 1)
        v11 = self._gather(mx + 1, my + 1)
        return w00 * v00 + w10 * v10 + w01 * v01 + w11 * v11


_MTW_CACHE: "weakref.WeakKeyDictionary[MaskWindow, _MaskTransWeights]" = weakref.WeakKeyDictionary()
_RASTER_MTW_CACHE: "weakref.WeakKeyDictionary[Window, _MaskTransWeights]" = weakref.WeakKeyDictionary()
_KD_CACHE: "weakref.WeakKeyDictionary[PPP, cKDTree]" = weakref.WeakKeyDictionary()
_PAIR_CACHE: "weakref.WeakKeyDictionary[PPP, Dict[str, object]]" = weakref.WeakKeyDictionary()
_DTB_CACHE: "weakref.WeakKeyDictionary[Window, Image]" = weakref.WeakKeyDictionary()


def _get_kdtree(ppp: PPP) -> cKDTree:
    """Cached cKDTree per PPP."""
    t = _KD_CACHE.get(ppp)
    if t is None:
        t = cKDTree(ppp.xy)
        _KD_CACHE[ppp] = t
    return t


def _get_pairs(ppp: PPP, r_max: float) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Cached unordered pairs (i, j, d) up to r_max for a PPP.
    """
    entry = _PAIR_CACHE.get(ppp)
    if entry is not None:
        cached_r = entry["r_max"]  # type: ignore[assignment]
        i, j, d = entry["pairs"]   # type: ignore[assignment]
        if float(cached_r) >= r_max:
            keep = d <= r_max
            return i[keep], j[keep], d[keep]
    i, j, d = _pairs_within(ppp.xy, r_max)
    _PAIR_CACHE[ppp] = {"r_max": float(r_max), "pairs": (i, j, d)}
    return i, j, d


def _get_dtb(win: Window) -> Image:
    """Cached distance-to-boundary image per Window instance."""
    img = _DTB_CACHE.get(win)
    if img is None:
        img = win.distance_to_boundary()
        _DTB_CACHE[win] = img
    return img


def _rasterised_mtw_from_window(win: Window) -> _MaskTransWeights:
    """
    Build (and cache) an FFT-based translation-weight helper for an arbitrary Window
    by rasterising membership on the DTB grid.
    """
    mtw = _RASTER_MTW_CACHE.get(win)
    if mtw is not None:
        return mtw
    img = _get_dtb(win)
    ny, nx = img.values.shape
    x0, y0 = img.origin
    sx, sy = img.spacing
    cx = x0 + (np.arange(nx) + 0.5) * sx
    cy = y0 + (np.arange(ny) + 0.5) * sy
    X, Y = np.meshgrid(cx, cy)
    pts = np.column_stack([X.ravel(), Y.ravel()])
    mask = win.contains(pts).reshape(ny, nx)
    mw = MaskWindow(mask, origin=img.origin, spacing=img.spacing)
    mtw = _MaskTransWeights(mw)
    _RASTER_MTW_CACHE[win] = mtw
    return mtw


class _TranslationCache:
    """
    Provide translation edge-correction weights w = 1 / overlap_fraction(shift).

    Fast paths:
      - MaskWindow: FFT-based weights + bilinear interpolation.
      - RectangleWindow: exact vectorised rectangle formula.
      - Other windows (e.g. PolyWindow): rasterise to MaskWindow once, then FFT path.
    """

    def __init__(self, win: Window) -> None:
        """
        :param win: Study window.
        :type win: Window
        """
        self._win = win
        self._maskw: Optional[_MaskTransWeights] = None
        self._rect_dims: Optional[Tuple[float, float, float]] = None  # (w, h, area)

        if isinstance(win, MaskWindow):
            mtw = _MTW_CACHE.get(win)
            if mtw is None:
                mtw = _MaskTransWeights(win)
                _MTW_CACHE[win] = mtw
            self._maskw = mtw
        elif isinstance(win, RectangleWindow):
            w = win.width()
            h = win.height()
            self._rect_dims = (float(w), float(h), float(w * h))
        else:
            self._maskw = _rasterised_mtw_from_window(win)

        self._fallback: Dict[Tuple[float, float], float] = {}

    def weights_many(self, deltas: np.ndarray) -> np.ndarray:
        """
        Compute weights for many shifts.

        :param deltas: Shifts (m, 2) in world units.
        :type deltas: numpy.ndarray
        :return: Weights (m,), with ``inf`` where overlap is zero.
        :rtype: numpy.ndarray
        """
        if self._maskw is not None:
            frac = self._maskw.frac_many(deltas)
            with np.errstate(divide="ignore"):
                return 1.0 / np.maximum(frac, 1e-15)

        if self._rect_dims is not None:
            w, h, area = self._rect_dims
            dx = np.abs(deltas[:, 0])
            dy = np.abs(deltas[:, 1])
            ox = np.maximum(0.0, w - dx)
            oy = np.maximum(0.0, h - dy)
            frac = (ox * oy) / area
            with np.errstate(divide="ignore"):
                return 1.0 / np.maximum(frac, 1e-15)

        keys = np.round(deltas, 9)
        out = np.empty(deltas.shape[0], float)
        for k in range(deltas.shape[0]):
            key = (float(keys[k, 0]), float(keys[k, 1]))
            val = self._fallback.get(key)
            if val is None:
                frac = self._win.overlap_fraction(key)
                val = np.inf if frac <= 0.0 else (1.0 / float(frac))
                self._fallback[key] = val
            out[k] = val
        return out


def _distance_to_boundary_at_points(img: Image, xy: np.ndarray) -> np.ndarray:
    """
    Sample a distance image at arbitrary points by nearest cell.

    :param img: Distance-to-boundary image.
    :type img: Image
    :param xy: Query points (n, 2).
    :type xy: numpy.ndarray
    :return: Distances at each query point.
    :rtype: numpy.ndarray
    """
    x0, y0 = img.origin
    sx, sy = img.spacing
    ny, nx = img.values.shape
    ix = np.floor((xy[:, 0] - x0) / sx).astype(int)
    iy = np.floor((xy[:, 1] - y0) / sy).astype(int)
    ix = np.clip(ix, 0, nx - 1)
    iy = np.clip(iy, 0, ny - 1)
    return img.values[iy, ix]


def Kest(ppp: PPP, r: np.ndarray, *, correction: str = "translation", showplot: bool = False) -> np.ndarray:
    """
    Ripley's K-function estimator.

    :param ppp: Point pattern.
    :type ppp: PPP
    :param r: Radii at which to estimate K(r); strictly increasing and positive.
    :type r: numpy.ndarray
    :param correction: Edge correction; ``"translation"`` or ``"border"``.
    :type correction: str
    :param showplot: If True, plot K(r) vs r with CSR baseline K_CSR = π r² (dashed, black).
    :type showplot: bool
    :return: Estimates of K(r) at each radius.
    :rtype: numpy.ndarray
    :raises ValueError: If ``correction`` is unrecognised.
    """
    r = _ensure_increasing_r(r)
    n = len(ppp)
    if n < 2:
        y = np.zeros_like(r)
        _plot_stat(r, y, "K(r)", np.pi * r * r, showplot)
        return y

    win = ppp.window
    area = win.area()
    xy = ppp.xy
    r_max = float(r[-1])

    if correction == "translation":
        i, j, d = _get_pairs(ppp, r_max)
        if d.size == 0:
            y = np.zeros_like(r)
            _plot_stat(r, y, "K(r)", np.pi * r * r, showplot)
            return y
        cache = _TranslationCache(win)
        deltas = xy[j] - xy[i]
        w = 2.0 * cache.weights_many(deltas)
        hist, _ = np.histogram(d, bins=np.concatenate(([0.0], r)), weights=w)
        csum = np.cumsum(hist)
        y = (area / (n * (n - 1))) * csum

    elif correction == "border":
        dtb_img = _get_dtb(win)
        db = _distance_to_boundary_at_points(dtb_img, xy)
        tree = _get_kdtree(ppp)
        y = np.empty_like(r)
        for k, rk in enumerate(r):
            anchors = np.where(db >= rk)[0]
            m = anchors.size
            if m == 0:
                y[k] = np.nan
                continue
            lens = tree.query_ball_point(xy[anchors], rk, return_length=True)
            total = int(np.sum(lens)) - m
            y[k] = (area / (m * n)) * total
    else:
        raise ValueError(f"Unknown correction '{correction}'")

    _plot_stat(r, y, "K(r)", np.pi * r * r, showplot)
    return y


def Lest(ppp: PPP, r: np.ndarray, *, correction: str = "translation", showplot: bool = False) -> np.ndarray:
    """
    Besag's L-function, ``L(r) = sqrt(K(r) / pi)``.

    :param ppp: Point pattern.
    :type ppp: PPP
    :param r: Radii at which to estimate L(r); strictly increasing and positive.
    :type r: numpy.ndarray
    :param correction: Edge correction; ``"translation"`` or ``"border"``.
    :type correction: str
    :param showplot: If True, plot L(r) vs r with CSR baseline L_CSR = r (dashed, black).
    :type showplot: bool
    :return: Estimates of L(r) at each radius.
    :rtype: numpy.ndarray
    """
    K = Kest(ppp, r, correction=correction, showplot=False)
    y = np.sqrt(K / np.pi)
    _plot_stat(r, y, "L(r)", r, showplot)
    return y


def pcf(ppp: PPP, r_edges: np.ndarray, *, correction: str = "translation", showplot: bool = False) -> Tuple[np.ndarray, np.ndarray]:
    """
    Pair correlation function via annular binning.

    :param ppp: Point pattern.
    :type ppp: PPP
    :param r_edges: Bin edges (strictly increasing, positive), shape (B+1,).
    :type r_edges: numpy.ndarray
    :param correction: Edge correction; ``"translation"`` or ``"border"``.
    :type correction: str
    :param showplot: If True, plot g(r) vs r with CSR baseline g_CSR = 1 (dashed, black).
    :type showplot: bool
    :return: Tuple ``(r_mid, g)`` with bin midpoints and estimated pcf values.
    :rtype: Tuple[numpy.ndarray, numpy.ndarray]
    :raises ValueError: If bin edges are invalid or ``correction`` is unrecognised.
    """
    r_edges = np.asarray(r_edges, float)
    if r_edges.ndim != 1 or r_edges.size < 2 or np.any(np.diff(r_edges) <= 0) or np.any(r_edges <= 0):
        raise ValueError("r_edges must be strictly increasing and positive.")
    n = len(ppp)
    centres = 0.5 * (r_edges[:-1] + r_edges[1:])
    if n < 2:
        y = np.zeros_like(centres)
        _plot_stat(centres, y, "g(r)", np.ones_like(centres), showplot)
        return centres, y

    win = ppp.window
    area = win.area()
    xy = ppp.xy
    r_max = float(r_edges[-1])

    if correction == "translation":
        i, j, d = _get_pairs(ppp, r_max)
        if d.size == 0:
            y = np.zeros_like(centres)
            _plot_stat(centres, y, "g(r)", np.ones_like(centres), showplot)
            return centres, y
        cache = _TranslationCache(win)
        deltas = xy[j] - xy[i]
        w = 2.0 * cache.weights_many(deltas)
        counts, _ = np.histogram(d, bins=r_edges, weights=w)
        ann = np.pi * (r_edges[1:] ** 2 - r_edges[:-1] ** 2)
        y = (area / (n * (n - 1))) * (counts / ann)

    elif correction == "border":
        dtb_img = _get_dtb(win)
        db = _distance_to_boundary_at_points(dtb_img, xy)
        tree = _get_kdtree(ppp)
        y = np.empty_like(centres)
        for b, (lo, hi) in enumerate(zip(r_edges[:-1], r_edges[1:])):
            anchors = np.where(db >= hi)[0]
            m = anchors.size
            if m == 0:
                y[b] = np.nan
                continue
            len_hi = tree.query_ball_point(xy[anchors], hi, return_length=True)
            if lo > 0.0:
                len_lo = tree.query_ball_point(xy[anchors], lo, return_length=True)
            else:
                len_lo = np.zeros_like(len_hi)
            counts = int(np.sum(len_hi)) - int(np.sum(len_lo))
            ann = np.pi * (hi * hi - lo * lo)
            y[b] = (area / (m * n)) * (counts / ann)
    else:
        raise ValueError("correction must be 'translation' or 'border'")

    _plot_stat(centres, y, "g(r)", np.ones_like(centres), showplot)
    return centres, y


def Gest(ppp: PPP, r: np.ndarray, *, method: str = "border", showplot: bool = False) -> np.ndarray:
    """
    Nearest-neighbour distribution ``G(r)``.

    :param ppp: Point pattern.
    :type ppp: PPP
    :param r: Radii at which to evaluate; strictly increasing and positive.
    :type r: numpy.ndarray
    :param method: ``'border'`` (restrict to points with distance-to-boundary ≥ r)
                   or ``'raw'`` (empirical CDF, biased).
    :type method: str
    :param showplot: If True, plot G(r) vs r with CSR baseline
                     ``G_CSR(r) = 1 - exp(-λ π r²)`` (dashed, black), with ``λ = n / |W|``.
    :type showplot: bool
    :return: Estimates of ``G(r)`` in ``[0, 1]`` (``nan`` where undefined for border).
    :rtype: numpy.ndarray
    :raises ValueError: If ``method`` is unrecognised.
    """
    r = _ensure_increasing_r(r)
    n = len(ppp)
    if n == 0:
        y = np.zeros_like(r)
        _plot_stat(r, y, "G(r)", np.zeros_like(r), showplot)
        return y

    tree = _get_kdtree(ppp)
    d_all = tree.query(ppp.xy, k=2)[0]
    d1 = d_all[:, 1]

    if method == "raw":
        y = np.searchsorted(np.sort(d1), r, side="right") / float(n)
    elif method == "border":
        dtb = _distance_to_boundary_at_points(_get_dtb(ppp.window), ppp.xy)
        y = np.empty_like(r)
        for k, rk in enumerate(r):
            ok = dtb >= rk
            m = np.count_nonzero(ok)
            if m == 0:
                y[k] = np.nan
            else:
                y[k] = np.count_nonzero(d1[ok] <= rk) / float(m)
    else:
        raise ValueError("method must be 'border' or 'raw'")

    lam = (n / ppp.window.area()) if ppp.window.area() > 0 else 0.0
    csr = 1.0 - np.exp(-lam * np.pi * r * r)
    _plot_stat(r, y, "G(r)", csr, showplot)
    return y


def Fest(ppp: PPP, r: np.ndarray, *, nsamp: Optional[int] = None, method: str = "border",
         rng: Optional[np.random.Generator] = None, showplot: bool = False) -> np.ndarray:
    """
    Empty-space function ``F(r)`` via Monte Carlo sampling over the window.

    :param ppp: Point pattern.
    :type ppp: PPP
    :param r: Radii at which to evaluate; strictly increasing and positive.
    :type r: numpy.ndarray
    :param nsamp: Number of random locations; default ``max(1000, 25 * max(1, n))``.
    :type nsamp: Optional[int]
    :param method: ``'border'`` (restrict to locations with distance-to-boundary ≥ r)
                   or ``'raw'`` (empirical CDF at sampled locations, biased).
    :type method: str
    :param rng: Random generator or seed.
    :type rng: Optional[numpy.random.Generator]
    :param showplot: If True, plot F(r) vs r with CSR baseline
                     ``F_CSR(r) = 1 - exp(-λ π r²)`` (dashed, black), with ``λ = n / |W|``.
    :type showplot: bool
    :return: Estimates of ``F(r)`` in ``[0, 1]`` (``nan`` where undefined for border).
    :rtype: numpy.ndarray
    :raises ValueError: If ``method`` is unrecognised.
    """
    r = _ensure_increasing_r(r)
    win = ppp.window
    xy = ppp.xy
    n = len(ppp)

    if nsamp is None:
        nsamp = max(1000, 25 * max(1, n))
    rng = rng if isinstance(rng, np.random.Generator) else np.random.default_rng(rng)

    u = win.sample_uniform(nsamp, rng=rng)

    if n == 0:
        y = np.zeros_like(r)
        _plot_stat(r, y, "F(r)", np.zeros_like(r), showplot)
        return y

    tree = _get_kdtree(ppp)
    d_u, _ = tree.query(u, k=1)

    if method == "raw":
        y = np.searchsorted(np.sort(d_u), r, side="right") / float(nsamp)
    elif method == "border":
        dtb_u = _distance_to_boundary_at_points(_get_dtb(win), u)
        y = np.empty_like(r)
        for k, rk in enumerate(r):
            ok = dtb_u >= rk
            m = np.count_nonzero(ok)
            if m == 0:
                y[k] = np.nan
            else:
                y[k] = np.count_nonzero(d_u[ok] <= rk) / float(m)
    else:
        raise ValueError("method must be 'border' or 'raw'")

    lam = (n / win.area()) if win.area() > 0 else 0.0
    csr = 1.0 - np.exp(-lam * np.pi * r * r)
    _plot_stat(r, y, "F(r)", csr, showplot)
    return y


def Jest(ppp: PPP, r: np.ndarray, *, method_G: str = "border", method_F: str = "border",
         nsamp_F: Optional[int] = None, rng: Optional[np.random.Generator] = None,
         showplot: bool = False) -> np.ndarray:
    """
    Function ``J(r) = (1 - G(r)) / (1 - F(r))``.

    :param ppp: Point pattern.
    :type ppp: PPP
    :param r: Radii at which to evaluate; strictly increasing and positive.
    :type r: numpy.ndarray
    :param method_G: Method for ``G`` (``'border'`` or ``'raw'``).
    :type method_G: str
    :param method_F: Method for ``F`` (``'border'`` or ``'raw'``).
    :type method_F: str
    :param nsamp_F: Number of samples for ``F`` (if ``None``, uses ``Fest`` default).
    :type nsamp_F: Optional[int]
    :param rng: Random generator or seed (passed to ``Fest``).
    :type rng: Optional[numpy.random.Generator]
    :param showplot: If True, plot J(r) vs r with CSR baseline J_CSR = 1 (dashed, black).
    :type showplot: bool
    :return: ``J(r)`` with ``nan`` where undefined (division by zero).
    :rtype: numpy.ndarray
    """
    G = Gest(ppp, r, method=method_G, showplot=False)
    F = Fest(ppp, r, nsamp=nsamp_F, method=method_F, rng=rng, showplot=False)
    with np.errstate(invalid="ignore", divide="ignore"):
        y = (1.0 - G) / (1.0 - F)

    _plot_stat(r, y, "J(r)", np.ones_like(r), showplot)
    return y
