"""
Point process simulation functions.

Provides simulation of:
- Homogeneous Poisson processes
- Inhomogeneous Poisson processes (intensity function or image)
- Thomas cluster processes
- Matérn cluster processes
- Simple Sequential Inhibition (SSI) hardcore processes
- Matérn I and II hardcore processes
"""
from __future__ import annotations
from typing import Optional, Callable, Union

import numpy as np

from .core import PPP, Window, Image, _as_rng


def rpoispp(
    intensity: float,
    window: Window,
    rng: Optional[np.random.Generator] = None,
) -> PPP:
    """
    Simulate a homogeneous Poisson point process.

    :param intensity: Points per unit area (λ > 0).
    :type intensity: float
    :param window: Study window.
    :type window: Window
    :param rng: Random generator or seed.
    :type rng: Optional[np.random.Generator]
    :return: Simulated point pattern.
    :rtype: PPP
    :raises ValueError: If intensity is non-positive.
    """
    if intensity <= 0:
        raise ValueError("Intensity must be positive.")
    rng = _as_rng(rng)
    area = window.area()
    n = rng.poisson(intensity * area)
    xy = window.sample_uniform(n, rng) if n > 0 else np.empty((0, 2), float)
    return PPP(xy, window, unit=getattr(window, "_unit", "1"))


def rpoispp_inhom(
    intensity: Union[Callable[[np.ndarray], np.ndarray], Image],
    window: Window,
    lmax: Optional[float] = None,
    rng: Optional[np.random.Generator] = None,
) -> PPP:
    """
    Simulate an inhomogeneous Poisson point process via thinning.

    :param intensity: Either a callable λ(xy) -> array of intensities at (n, 2) points,
                      or an Image with intensity values at pixel centres.
    :type intensity: Union[Callable, Image]
    :param window: Study window.
    :type window: Window
    :param lmax: Upper bound on intensity (required for callable; inferred for Image).
    :type lmax: Optional[float]
    :param rng: Random generator or seed.
    :type rng: Optional[np.random.Generator]
    :return: Simulated point pattern.
    :rtype: PPP
    :raises ValueError: If lmax not provided for callable or lmax <= 0.
    """
    rng = _as_rng(rng)

    if isinstance(intensity, Image):
        lmax_img = float(np.max(intensity.values))
        if lmax is None:
            lmax = lmax_img
        else:
            lmax = max(float(lmax), lmax_img)
        intensity_func = _image_intensity_func(intensity)
    else:
        if lmax is None:
            raise ValueError("lmax must be provided when intensity is a callable.")
        lmax = float(lmax)
        intensity_func = intensity

    if lmax <= 0:
        raise ValueError("lmax must be positive.")

    # Generate homogeneous Poisson with rate lmax
    area = window.area()
    n_prop = rng.poisson(lmax * area)
    if n_prop == 0:
        return PPP(np.empty((0, 2), float), window, unit=getattr(window, "_unit", "1"))

    xy_prop = window.sample_uniform(n_prop, rng)

    # Thinning: keep point with probability λ(x)/lmax
    lam_vals = intensity_func(xy_prop)
    lam_vals = np.asarray(lam_vals, float).ravel()
    lam_vals = np.clip(lam_vals, 0.0, lmax)
    probs = lam_vals / lmax
    keep = rng.random(n_prop) < probs
    xy = xy_prop[keep]

    return PPP(xy, window, unit=getattr(window, "_unit", "1"))


def _image_intensity_func(img: Image) -> Callable[[np.ndarray], np.ndarray]:
    """
    Create an intensity function from an Image via nearest-cell lookup.

    :param img: Intensity image.
    :type img: Image
    :return: Callable that evaluates intensity at (n, 2) points.
    :rtype: Callable
    """
    x0, y0 = img.origin
    sx, sy = img.spacing
    ny, nx = img.values.shape
    vals = img.values

    def func(xy: np.ndarray) -> np.ndarray:
        xy = np.asarray(xy, float)
        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 vals[iy, ix]

    return func


def rThomas(
    kappa: float,
    sigma: float,
    mu: float,
    window: Window,
    rng: Optional[np.random.Generator] = None,
    expand: float = 4.0,
) -> PPP:
    """
    Simulate a Thomas cluster process.

    Parents form a Poisson process with intensity κ. Each parent produces
    a Poisson(μ) number of offspring displaced by bivariate normal N(0, σ²I).

    :param kappa: Parent intensity (κ > 0).
    :type kappa: float
    :param sigma: Cluster spread (standard deviation, σ > 0).
    :type sigma: float
    :param mu: Mean number of offspring per parent (μ > 0).
    :type mu: float
    :param window: Study window.
    :type window: Window
    :param rng: Random generator or seed.
    :type rng: Optional[np.random.Generator]
    :param expand: Expansion factor for parent window (multiples of σ).
    :type expand: float
    :return: Simulated point pattern (offspring only, clipped to window).
    :rtype: PPP
    :raises ValueError: If parameters are non-positive.
    """
    if kappa <= 0 or sigma <= 0 or mu <= 0:
        raise ValueError("kappa, sigma, and mu must all be positive.")
    rng = _as_rng(rng)

    # Expand window for parents to capture edge effects
    parent_win = _expanded_rectangle(window, expand * sigma)
    parent_area = parent_win.area()
    n_parents = rng.poisson(kappa * parent_area)

    if n_parents == 0:
        return PPP(np.empty((0, 2), float), window, unit=getattr(window, "_unit", "1"))

    # Generate parent locations
    parent_xy = parent_win.sample_uniform(n_parents, rng)

    # Generate offspring counts
    n_offspring = rng.poisson(mu, size=n_parents)
    total_offspring = int(np.sum(n_offspring))

    if total_offspring == 0:
        return PPP(np.empty((0, 2), float), window, unit=getattr(window, "_unit", "1"))

    # Assign offspring to parents
    parent_idx = np.repeat(np.arange(n_parents), n_offspring)
    offspring_xy = parent_xy[parent_idx] + rng.normal(0, sigma, size=(total_offspring, 2))

    # Clip to window
    inside = window.contains(offspring_xy)
    xy = offspring_xy[inside]

    return PPP(xy, window, unit=getattr(window, "_unit", "1"))


def rMatClust(
    kappa: float,
    r: float,
    mu: float,
    window: Window,
    rng: Optional[np.random.Generator] = None,
    expand: float = 1.0,
) -> PPP:
    """
    Simulate a Matérn cluster process.

    Parents form a Poisson process with intensity κ. Each parent produces
    a Poisson(μ) number of offspring uniformly distributed in a disc of radius r.

    :param kappa: Parent intensity (κ > 0).
    :type kappa: float
    :param r: Cluster radius (r > 0).
    :type r: float
    :param mu: Mean number of offspring per parent (μ > 0).
    :type mu: float
    :param window: Study window.
    :type window: Window
    :param rng: Random generator or seed.
    :type rng: Optional[np.random.Generator]
    :param expand: Expansion factor for parent window (multiples of r).
    :type expand: float
    :return: Simulated point pattern (offspring only, clipped to window).
    :rtype: PPP
    :raises ValueError: If parameters are non-positive.
    """
    if kappa <= 0 or r <= 0 or mu <= 0:
        raise ValueError("kappa, r, and mu must all be positive.")
    rng = _as_rng(rng)

    # Expand window for parents
    parent_win = _expanded_rectangle(window, expand * r)
    parent_area = parent_win.area()
    n_parents = rng.poisson(kappa * parent_area)

    if n_parents == 0:
        return PPP(np.empty((0, 2), float), window, unit=getattr(window, "_unit", "1"))

    parent_xy = parent_win.sample_uniform(n_parents, rng)

    n_offspring = rng.poisson(mu, size=n_parents)
    total_offspring = int(np.sum(n_offspring))

    if total_offspring == 0:
        return PPP(np.empty((0, 2), float), window, unit=getattr(window, "_unit", "1"))

    parent_idx = np.repeat(np.arange(n_parents), n_offspring)

    # Uniform in disc: radius ~ sqrt(U) * r, angle ~ U(0, 2pi)
    u_rad = np.sqrt(rng.random(total_offspring)) * r
    u_ang = rng.random(total_offspring) * 2 * np.pi
    dx = u_rad * np.cos(u_ang)
    dy = u_rad * np.sin(u_ang)
    offspring_xy = parent_xy[parent_idx] + np.column_stack([dx, dy])

    inside = window.contains(offspring_xy)
    xy = offspring_xy[inside]

    return PPP(xy, window, unit=getattr(window, "_unit", "1"))


def rSSI(
    n: int,
    r: float,
    window: Window,
    rng: Optional[np.random.Generator] = None,
    max_attempts: int = 100000,
) -> PPP:
    """
    Simulate a Simple Sequential Inhibition (SSI) hardcore process.

    Sequentially proposes points uniformly in the window, accepting only
    those at distance > r from all existing points.

    :param n: Target number of points.
    :type n: int
    :param r: Hardcore distance (minimum separation, r ≥ 0).
    :type r: float
    :param window: Study window.
    :type window: Window
    :param rng: Random generator or seed.
    :type rng: Optional[np.random.Generator]
    :param max_attempts: Maximum proposal attempts before giving up.
    :type max_attempts: int
    :return: Simulated point pattern (may have < n points if packing limit reached).
    :rtype: PPP
    :raises ValueError: If n < 0 or r < 0.
    """
    if n < 0:
        raise ValueError("n must be non-negative.")
    if r < 0:
        raise ValueError("r must be non-negative.")
    if n == 0:
        return PPP(np.empty((0, 2), float), window, unit=getattr(window, "_unit", "1"))

    rng = _as_rng(rng)
    points = []
    attempts = 0

    while len(points) < n and attempts < max_attempts:
        # Propose in batches for efficiency
        batch_size = min(1000, max_attempts - attempts)
        proposals = window.sample_uniform(batch_size, rng)
        attempts += batch_size

        for pt in proposals:
            if len(points) >= n:
                break
            if len(points) == 0:
                points.append(pt)
            else:
                existing = np.array(points)
                dists = np.sqrt(np.sum((existing - pt) ** 2, axis=1))
                if np.all(dists > r):
                    points.append(pt)

    xy = np.array(points) if points else np.empty((0, 2), float)
    return PPP(xy, window, unit=getattr(window, "_unit", "1"))


def rMaternI(
    intensity: float,
    r: float,
    window: Window,
    rng: Optional[np.random.Generator] = None,
) -> PPP:
    """
    Simulate a Matérn Type I hardcore process.

    Generate a Poisson process, then delete all points that have any
    neighbour within distance r (mutual deletion).

    :param intensity: Intensity of the initial Poisson process (λ > 0).
    :type intensity: float
    :param r: Hardcore distance (r ≥ 0).
    :type r: float
    :param window: Study window.
    :type window: Window
    :param rng: Random generator or seed.
    :type rng: Optional[np.random.Generator]
    :return: Simulated point pattern.
    :rtype: PPP
    :raises ValueError: If intensity ≤ 0 or r < 0.
    """
    if intensity <= 0:
        raise ValueError("Intensity must be positive.")
    if r < 0:
        raise ValueError("r must be non-negative.")

    rng = _as_rng(rng)

    # Generate initial Poisson
    ppp = rpoispp(intensity, window, rng)
    if len(ppp) < 2 or r == 0:
        return ppp

    xy = ppp.xy
    from scipy.spatial import cKDTree
    tree = cKDTree(xy)

    # Find all pairs within r
    pairs = tree.query_pairs(r)
    if not pairs:
        return ppp

    # Collect all points involved in conflicts
    conflicting = set()
    for i, j in pairs:
        conflicting.add(i)
        conflicting.add(j)

    keep_mask = np.ones(len(ppp), dtype=bool)
    keep_mask[list(conflicting)] = False
    xy_kept = xy[keep_mask]

    return PPP(xy_kept, window, unit=getattr(window, "_unit", "1"))


def rMaternII(
    intensity: float,
    r: float,
    window: Window,
    rng: Optional[np.random.Generator] = None,
) -> PPP:
    """
    Simulate a Matérn Type II hardcore process.

    Generate a Poisson process with marks (arrival times), then keep a point
    only if it has no earlier-arriving neighbour within distance r.

    :param intensity: Intensity of the initial Poisson process (λ > 0).
    :type intensity: float
    :param r: Hardcore distance (r ≥ 0).
    :type r: float
    :param window: Study window.
    :type window: Window
    :param rng: Random generator or seed.
    :type rng: Optional[np.random.Generator]
    :return: Simulated point pattern.
    :rtype: PPP
    :raises ValueError: If intensity ≤ 0 or r < 0.
    """
    if intensity <= 0:
        raise ValueError("Intensity must be positive.")
    if r < 0:
        raise ValueError("r must be non-negative.")

    rng = _as_rng(rng)

    # Generate initial Poisson
    ppp = rpoispp(intensity, window, rng)
    n = len(ppp)
    if n < 2 or r == 0:
        return ppp

    xy = ppp.xy
    # Assign random arrival times
    times = rng.random(n)

    from scipy.spatial import cKDTree
    tree = cKDTree(xy)

    # For each point, check if any earlier neighbour within r
    keep_mask = np.ones(n, dtype=bool)
    for i in range(n):
        if not keep_mask[i]:
            continue
        neighbours = tree.query_ball_point(xy[i], r)
        for j in neighbours:
            if j != i and times[j] < times[i]:
                # Point j arrived earlier and is within r -> delete i
                keep_mask[i] = False
                break

    xy_kept = xy[keep_mask]
    return PPP(xy_kept, window, unit=getattr(window, "_unit", "1"))


def _expanded_rectangle(window: Window, margin: float) -> "Window":
    """
    Create an expanded rectangular window around the given window.

    :param window: Original window.
    :type window: Window
    :param margin: Expansion margin on each side.
    :type margin: float
    :return: Expanded RectangleWindow.
    :rtype: Window
    """
    from .core import RectangleWindow

    # Get bounds depending on window type
    if hasattr(window, "bounds"):
        bounds = window.bounds
        if callable(bounds):
            bounds = bounds()
        xmin, ymin, xmax, ymax = bounds
    elif hasattr(window, "_xmin"):
        xmin, ymin = window._xmin, window._ymin
        xmax, ymax = window._xmax, window._ymax
    elif hasattr(window, "mask"):
        # MaskWindow
        x0, y0 = window.origin
        sx, sy = window.spacing
        ny, nx = window.shape
        xmin, ymin = x0, y0
        xmax, ymax = x0 + nx * sx, y0 + ny * sy
    else:
        # Fallback: sample points to estimate bounds
        pts = window.sample_uniform(10000, np.random.default_rng(0))
        xmin, ymin = pts.min(axis=0)
        xmax, ymax = pts.max(axis=0)

    return RectangleWindow(
        xmin - margin, ymin - margin, xmax + margin, ymax + margin,
        unit=getattr(window, "_unit", "1")
    )