"""
Monte Carlo simulation envelopes for point pattern statistics.

Provides:
- Simulation envelopes under CSR null hypothesis
- Pointwise and global envelope tests
- Support for K, L, g, G, F, J statistics
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Optional, Callable, Literal, Tuple, Union

import numpy as np

from .core import PPP, Window, _as_rng
from .simulate import rpoispp


@dataclass
class Envelope:
    """
    Result of a Monte Carlo envelope test.

    :param r: Radii at which the statistic was evaluated.
    :param observed: Observed statistic values.
    :param mean: Mean of simulated statistics.
    :param lo: Lower envelope bound.
    :param hi: Upper envelope bound.
    :param simulations: Array of simulated statistics (nsim, len(r)) if stored.
    :param rank: Rank used for envelope (e.g., 1 for min/max, 5 for 5th highest/lowest).
    :param nsim: Number of simulations.
    :param stat_name: Name of the statistic (e.g., "K", "L", "g").
    :param global_pvalue: Global envelope test p-value (if computed).
    :param global_deviation: Maximum absolute deviation of observed from mean.
    """
    r: np.ndarray
    observed: np.ndarray
    mean: np.ndarray
    lo: np.ndarray
    hi: np.ndarray
    simulations: Optional[np.ndarray]
    rank: int
    nsim: int
    stat_name: str
    global_pvalue: Optional[float] = None
    global_deviation: Optional[float] = None

    def plot(self, show: bool = True) -> None:
        """
        Plot the envelope with observed statistic.

        :param show: If True, display the figure.
        :type show: bool
        """
        try:
            import matplotlib.pyplot as plt
        except ImportError as exc:
            raise RuntimeError("Plotting requires matplotlib.") from exc

        fig, ax = plt.subplots()
        ax.fill_between(self.r, self.lo, self.hi, alpha=0.3, color="gray", label="envelope")
        ax.plot(self.r, self.mean, "--", color="black", label="CSR mean")
        ax.plot(self.r, self.observed, "-", color="blue", linewidth=1.5, label="observed")
        ax.set_xlabel("r")
        ax.set_ylabel(f"{self.stat_name}(r)")
        title = f"{self.stat_name} envelope (nsim={self.nsim}, rank={self.rank})"
        if self.global_pvalue is not None:
            title += f", global p={self.global_pvalue:.3f}"
        ax.set_title(title)
        ax.legend()
        fig.tight_layout()
        if show:
            plt.show()


StatFunction = Callable[[PPP, np.ndarray], np.ndarray]


def envelope(
    ppp: PPP,
    stat_func: StatFunction,
    r: np.ndarray,
    nsim: int = 99,
    rank: int = 1,
    global_test: bool = False,
    alternative: Literal["two.sided", "greater", "less"] = "two.sided",
    rng: Optional[np.random.Generator] = None,
    store_simulations: bool = False,
    stat_name: str = "stat",
    verbose: bool = False,
) -> Envelope:
    """
    Compute a Monte Carlo envelope for a point pattern statistic.

    :param ppp: Observed point pattern.
    :type ppp: PPP
    :param stat_func: Function ``f(ppp, r) -> array`` computing the statistic.
    :type stat_func: Callable
    :param r: Radii at which to evaluate the statistic.
    :type r: numpy.ndarray
    :param nsim: Number of CSR simulations.
    :type nsim: int
    :param rank: Rank for envelope bounds (1 = most extreme, giving ~2*rank/(nsim+1) coverage).
    :type rank: int
    :param global_test: If True, compute global envelope test p-value.
    :type global_test: bool
    :param alternative: Alternative hypothesis for global test.
    :type alternative: str
    :param rng: Random generator or seed.
    :type rng: Optional[np.random.Generator]
    :param store_simulations: If True, store all simulated statistics in result.
    :type store_simulations: bool
    :param stat_name: Name of the statistic for labeling.
    :type stat_name: str
    :param verbose: If True, print progress.
    :type verbose: bool
    :return: Envelope result object.
    :rtype: Envelope
    :raises ValueError: If nsim < 1 or rank < 1.
    """
    if nsim < 1:
        raise ValueError("nsim must be at least 1.")
    if rank < 1:
        raise ValueError("rank must be at least 1.")
    if rank > nsim // 2:
        raise ValueError(f"rank ({rank}) too large for nsim ({nsim}); need rank <= nsim // 2.")

    rng = _as_rng(rng)
    r = np.asarray(r, float)
    window = ppp.window
    n = len(ppp)

    # Compute observed statistic
    observed = stat_func(ppp, r)
    observed = np.asarray(observed, float)

    # Estimate intensity for CSR simulation
    intensity = n / window.area() if window.area() > 0 else 0.0

    # Run simulations
    sim_stats = np.empty((nsim, r.size), float)
    for i in range(nsim):
        if verbose and (i + 1) % 10 == 0:
            print(f"Simulation {i + 1}/{nsim}")
        sim_ppp = rpoispp(max(intensity, 1e-10), window, rng)
        sim_stats[i] = stat_func(sim_ppp, r)

    # Compute envelope bounds
    sim_sorted = np.sort(sim_stats, axis=0)
    lo = sim_sorted[rank - 1]          # rank-th lowest
    hi = sim_sorted[nsim - rank]       # rank-th highest
    mean = np.nanmean(sim_stats, axis=0)

    # Global test
    global_pvalue = None
    global_deviation = None
    if global_test:
        global_pvalue, global_deviation = _global_envelope_test(
            observed, sim_stats, mean, alternative
        )

    return Envelope(
        r=r,
        observed=observed,
        mean=mean,
        lo=lo,
        hi=hi,
        simulations=sim_stats if store_simulations else None,
        rank=rank,
        nsim=nsim,
        stat_name=stat_name,
        global_pvalue=global_pvalue,
        global_deviation=global_deviation,
    )


def _global_envelope_test(
    observed: np.ndarray,
    simulated: np.ndarray,
    mean: np.ndarray,
    alternative: str,
) -> Tuple[float, float]:
    """
    Compute global envelope test p-value using maximum absolute deviation.

    :param observed: Observed statistic values.
    :param simulated: Array (nsim, len(r)) of simulated statistics.
    :param mean: Mean of simulated statistics.
    :param alternative: "two.sided", "greater", or "less".
    :return: (p-value, max deviation of observed from mean).
    """
    nsim = simulated.shape[0]

    # Handle NaN values
    mask = np.isfinite(observed) & np.all(np.isfinite(simulated), axis=0)
    if not mask.any():
        return np.nan, np.nan

    obs_masked = observed[mask]
    sim_masked = simulated[:, mask]
    mean_masked = mean[mask]

    if alternative == "two.sided":
        obs_dev = np.max(np.abs(obs_masked - mean_masked))
        sim_devs = np.max(np.abs(sim_masked - mean_masked), axis=1)
    elif alternative == "greater":
        obs_dev = np.max(obs_masked - mean_masked)
        sim_devs = np.max(sim_masked - mean_masked, axis=1)
    elif alternative == "less":
        obs_dev = np.max(mean_masked - obs_masked)
        sim_devs = np.max(mean_masked - sim_masked, axis=1)
    else:
        raise ValueError(f"Unknown alternative: {alternative}")

    # p-value: proportion of simulations with deviation >= observed
    # Add 1 to numerator and denominator for proper Monte Carlo p-value
    pvalue = (1 + np.sum(sim_devs >= obs_dev)) / (1 + nsim)

    return float(pvalue), float(obs_dev)


def envelope_K(
    ppp: PPP,
    r: np.ndarray,
    nsim: int = 99,
    rank: int = 1,
    correction: str = "translation",
    global_test: bool = False,
    alternative: Literal["two.sided", "greater", "less"] = "two.sided",
    rng: Optional[np.random.Generator] = None,
    store_simulations: bool = False,
    verbose: bool = False,
) -> Envelope:
    """
    Compute envelope for Ripley's K function.

    :param ppp: Observed point pattern.
    :param r: Radii for K evaluation.
    :param nsim: Number of simulations.
    :param rank: Envelope rank.
    :param correction: Edge correction ("translation" or "border").
    :param global_test: Compute global p-value.
    :param alternative: Alternative hypothesis.
    :param rng: Random generator.
    :param store_simulations: Store all simulation results.
    :param verbose: Print progress.
    :return: Envelope result.
    """
    from .stats import Kest

    def stat_func(p: PPP, r: np.ndarray) -> np.ndarray:
        return Kest(p, r, correction=correction, showplot=False)

    return envelope(
        ppp, stat_func, r, nsim=nsim, rank=rank, global_test=global_test,
        alternative=alternative, rng=rng, store_simulations=store_simulations,
        stat_name="K", verbose=verbose
    )


def envelope_L(
    ppp: PPP,
    r: np.ndarray,
    nsim: int = 99,
    rank: int = 1,
    correction: str = "translation",
    global_test: bool = False,
    alternative: Literal["two.sided", "greater", "less"] = "two.sided",
    rng: Optional[np.random.Generator] = None,
    store_simulations: bool = False,
    verbose: bool = False,
) -> Envelope:
    """
    Compute envelope for Besag's L function.

    Parameters same as envelope_K.
    """
    from .stats import Lest

    def stat_func(p: PPP, r: np.ndarray) -> np.ndarray:
        return Lest(p, r, correction=correction, showplot=False)

    return envelope(
        ppp, stat_func, r, nsim=nsim, rank=rank, global_test=global_test,
        alternative=alternative, rng=rng, store_simulations=store_simulations,
        stat_name="L", verbose=verbose
    )


def envelope_pcf(
    ppp: PPP,
    r_edges: np.ndarray,
    nsim: int = 99,
    rank: int = 1,
    correction: str = "translation",
    global_test: bool = False,
    alternative: Literal["two.sided", "greater", "less"] = "two.sided",
    rng: Optional[np.random.Generator] = None,
    store_simulations: bool = False,
    verbose: bool = False,
) -> Envelope:
    """
    Compute envelope for pair correlation function g(r).

    :param r_edges: Bin edges for pcf (returns midpoints in result.r).
    :param other parameters: Same as envelope_K.
    """
    from .stats import pcf as pcf_stat

    r_edges = np.asarray(r_edges, float)
    # Compute midpoints for result
    r_mid = 0.5 * (r_edges[:-1] + r_edges[1:])

    def stat_func(p: PPP, r: np.ndarray) -> np.ndarray:
        # r is actually r_edges passed through envelope()
        _, g = pcf_stat(p, r_edges, correction=correction, showplot=False)
        return g

    # Pass r_mid to envelope so shapes match
    env = envelope(
        ppp, stat_func, r_mid, nsim=nsim, rank=rank, global_test=global_test,
        alternative=alternative, rng=rng, store_simulations=store_simulations,
        stat_name="g", verbose=verbose
    )
    return env


def envelope_G(
    ppp: PPP,
    r: np.ndarray,
    nsim: int = 99,
    rank: int = 1,
    method: str = "border",
    global_test: bool = False,
    alternative: Literal["two.sided", "greater", "less"] = "two.sided",
    rng: Optional[np.random.Generator] = None,
    store_simulations: bool = False,
    verbose: bool = False,
) -> Envelope:
    """
    Compute envelope for nearest-neighbour G function.

    :param method: "border" or "raw".
    :param other parameters: Same as envelope_K.
    """
    from .stats import Gest

    def stat_func(p: PPP, r: np.ndarray) -> np.ndarray:
        return Gest(p, r, method=method, showplot=False)

    return envelope(
        ppp, stat_func, r, nsim=nsim, rank=rank, global_test=global_test,
        alternative=alternative, rng=rng, store_simulations=store_simulations,
        stat_name="G", verbose=verbose
    )


def envelope_F(
    ppp: PPP,
    r: np.ndarray,
    nsim: int = 99,
    rank: int = 1,
    method: str = "border",
    nsamp: Optional[int] = None,
    global_test: bool = False,
    alternative: Literal["two.sided", "greater", "less"] = "two.sided",
    rng: Optional[np.random.Generator] = None,
    store_simulations: bool = False,
    verbose: bool = False,
) -> Envelope:
    """
    Compute envelope for empty-space F function.

    :param method: "border" or "raw".
    :param nsamp: Number of sample points for F estimation.
    :param other parameters: Same as envelope_K.
    """
    from .stats import Fest

    # Need separate rng for Fest sampling within each simulation
    rng = _as_rng(rng)

    def stat_func(p: PPP, r: np.ndarray) -> np.ndarray:
        # Use a fresh seed derived from current rng state for reproducibility
        return Fest(p, r, nsamp=nsamp, method=method, rng=rng, showplot=False)

    return envelope(
        ppp, stat_func, r, nsim=nsim, rank=rank, global_test=global_test,
        alternative=alternative, rng=rng, store_simulations=store_simulations,
        stat_name="F", verbose=verbose
    )


def envelope_J(
    ppp: PPP,
    r: np.ndarray,
    nsim: int = 99,
    rank: int = 1,
    method_G: str = "border",
    method_F: str = "border",
    nsamp_F: Optional[int] = None,
    global_test: bool = False,
    alternative: Literal["two.sided", "greater", "less"] = "two.sided",
    rng: Optional[np.random.Generator] = None,
    store_simulations: bool = False,
    verbose: bool = False,
) -> Envelope:
    """
    Compute envelope for J function.

    :param method_G: Method for G estimation.
    :param method_F: Method for F estimation.
    :param nsamp_F: Sample size for F.
    :param other parameters: Same as envelope_K.
    """
    from .stats import Jest

    rng = _as_rng(rng)

    def stat_func(p: PPP, r: np.ndarray) -> np.ndarray:
        return Jest(p, r, method_G=method_G, method_F=method_F,
                    nsamp_F=nsamp_F, rng=rng, showplot=False)

    return envelope(
        ppp, stat_func, r, nsim=nsim, rank=rank, global_test=global_test,
        alternative=alternative, rng=rng, store_simulations=store_simulations,
        stat_name="J", verbose=verbose
    )


def dclf_test(
    ppp: PPP,
    stat_func: StatFunction,
    r: np.ndarray,
    nsim: int = 99,
    rng: Optional[np.random.Generator] = None,
    stat_name: str = "stat",
    verbose: bool = False,
) -> Tuple[float, float, np.ndarray]:
    """
    Diggle-Cressie-Loosmore-Ford (DCLF) global envelope test.

    Uses integrated squared deviation from theoretical CSR.

    :param ppp: Observed point pattern.
    :param stat_func: Statistic function.
    :param r: Radii for evaluation.
    :param nsim: Number of simulations.
    :param rng: Random generator.
    :param stat_name: Name of statistic.
    :param verbose: Print progress.
    :return: (p-value, observed test statistic, array of simulated test statistics).
    """
    rng = _as_rng(rng)
    r = np.asarray(r, float)
    window = ppp.window
    n = len(ppp)
    intensity = n / window.area() if window.area() > 0 else 0.0

    # Compute observed
    observed = np.asarray(stat_func(ppp, r), float)

    # Theoretical CSR value (depends on statistic)
    # For K: pi * r^2, for L: r, for g/G/F: 1 / (1-exp(-...))
    sim_stats = np.empty((nsim, r.size), float)
    for i in range(nsim):
        if verbose and (i + 1) % 10 == 0:
            print(f"Simulation {i + 1}/{nsim}")
        sim_ppp = rpoispp(max(intensity, 1e-10), window, rng)
        sim_stats[i] = stat_func(sim_ppp, r)

    mean = np.nanmean(sim_stats, axis=0)

    # Integrated squared deviation
    dr = np.diff(r, prepend=r[0])
    mask = np.isfinite(observed) & np.isfinite(mean)

    obs_u = np.sum(((observed[mask] - mean[mask]) ** 2) * dr[mask])
    sim_u = np.sum(((sim_stats[:, mask] - mean[mask]) ** 2) * dr[mask], axis=1)

    pvalue = (1 + np.sum(sim_u >= obs_u)) / (1 + nsim)

    return float(pvalue), float(obs_u), sim_u