"""Pure-Python fallback for Monte-Carlo ranking.

The Cython kernel (`zen_fronts.selection_core._mc_rank`) is the intended
production implementation for winner/loser selection.

This module exists to:

* keep the high-level API usable in environments where compiling extensions is
  not feasible,
* enable unit tests for the high-level selection layer without requiring a C
  toolchain,
* support *detailed per-individual Monte-Carlo distribution statistics* that are
  not currently exposed by the Cython kernel.

The implementation matches the public signature of the Cython kernel:

```
mc_rank_counts(mu, sigma2, n_samples, percentile=0.2, seed=0, mode_i=1, within_mode_i=0)
mc_rank_select(mu, sigma2, n_samples, percentile=0.2, seed=0, mode_i=1, within_mode_i=0)
```

Additionally, it provides a higher-level detailed variant:

```
mc_rank_select_detailed(...)
```

which returns per-individual summary statistics for:

* overall position ("place"),
* Pareto front index ("front"),
* within-front rank ("within").

All objectives are assumed to be maximized.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np

from zen_fronts.pareto_core.pareto_rating import rate_population


def _within_mode_from_i(within_mode_i: int) -> str:
    if within_mode_i == 0:
        return "crowding"
    raise ValueError(f"within_mode_i={within_mode_i} is not supported in fallback")


def _mode_from_i(mode_i: int) -> str:
    if mode_i == 0:
        return "values"
    if mode_i == 1:
        return "ranks"
    raise ValueError(f"mode_i={mode_i} is not supported in fallback")


def mc_rank_counts(
    mu_obj,
    sigma2_obj,
    n_samples: int,
    *,
    percentile: float = 0.2,
    seed: int = 0,
    mode_i: int = 1,
    within_mode_i: int = 0,
):
    mu = np.asarray(mu_obj, dtype=np.float64)
    sigma2 = np.asarray(sigma2_obj, dtype=np.float64)

    if mu.ndim != 2 or sigma2.ndim != 2:
        raise ValueError("mu and sigma2 must be 2D")
    if mu.shape != sigma2.shape:
        raise ValueError("mu and sigma2 must have the same shape")
    if n_samples <= 0:
        raise ValueError("n_samples must be > 0")
    if percentile <= 0.0 or percentile > 1.0:
        raise ValueError("percentile must be in (0,1]")

    N, M = mu.shape
    if N == 0:
        return (
            np.zeros(0, dtype=np.int32),
            np.zeros(0, dtype=np.int32),
            np.zeros(0, dtype=np.int64),
        )

    k = int(np.ceil(percentile * N))
    k = max(1, min(k, N))

    rng = np.random.default_rng(int(seed))
    sigma = np.sqrt(np.maximum(sigma2, 0.0))

    win_count = np.zeros(N, dtype=np.int32)
    lose_count = np.zeros(N, dtype=np.int32)
    rank_sum = np.zeros(N, dtype=np.int64)

    mode = _mode_from_i(int(mode_i))
    within_mode = _within_mode_from_i(int(within_mode_i))

    for _ in range(int(n_samples)):
        eps = rng.standard_normal(size=(N, M))
        X = mu + sigma * eps
        rating = rate_population(X, mode=mode, within_mode=within_mode, use_cython=True)
        order = np.asarray(rating.order, dtype=np.int64)

        # rank_sum: add each index's position
        # (order[pos] == idx)
        for pos, idx in enumerate(order.tolist()):
            rank_sum[idx] += pos

        top = order[:k]
        bot = order[-k:]
        win_count[top] += 1
        lose_count[bot] += 1

    return win_count, lose_count, rank_sum


def mc_rank_select(
    mu_obj,
    sigma2_obj,
    n_samples: int,
    *,
    percentile: float = 0.2,
    seed: int = 0,
    mode_i: int = 1,
    within_mode_i: int = 0,
):
    win_count, lose_count, rank_sum = mc_rank_counts(
        mu_obj,
        sigma2_obj,
        n_samples,
        percentile=percentile,
        seed=seed,
        mode_i=mode_i,
        within_mode_i=within_mode_i,
    )

    N = int(win_count.shape[0])
    if N == 0:
        empty = np.zeros(0, dtype=np.int64)
        return empty, empty, win_count, lose_count, rank_sum

    k = int(np.ceil(percentile * N))
    k = max(1, min(k, N))

    idx = np.arange(N, dtype=np.int64)
    winners_order = np.lexsort((idx, rank_sum, -win_count))
    losers_order = np.lexsort((idx, -rank_sum, -lose_count))

    winners = winners_order[:k].astype(np.int64, copy=False)
    losers = losers_order[:k].astype(np.int64, copy=False)
    return winners, losers, win_count, lose_count, rank_sum


@dataclass(frozen=True)
class RankSummary:
    mean: np.ndarray
    median: np.ndarray
    std: np.ndarray
    q25: np.ndarray
    q75: np.ndarray


def _summarize(samples: np.ndarray) -> RankSummary:
    samples = np.asarray(samples, dtype=np.float64)
    if samples.ndim != 2:
        raise ValueError("samples must be 2D (S,N)")
    return RankSummary(
        mean=np.mean(samples, axis=0),
        median=np.median(samples, axis=0),
        std=np.std(samples, axis=0),
        q25=np.quantile(samples, 0.25, axis=0),
        q75=np.quantile(samples, 0.75, axis=0),
    )


def mc_rank_select_detailed(
    mu_obj,
    sigma2_obj,
    n_samples: int,
    *,
    percentile: float = 0.2,
    seed: int = 0,
    mode_i: int = 1,
    within_mode_i: int = 0,
    quantiles_mode_i: int = 2,
):
    """Monte-Carlo winner/loser selection + per-individual distribution stats.

    Returns
    -------
    winners_i, losers_i, win_count, lose_count, rank_sum, place, front, within

    Where:
    * `winners_i` / `losers_i` are index arrays (int64) into the row axis of mu.
    * `win_count`, `lose_count`, `rank_sum` are aligned with mu rows.
    * `place`, `front`, `within` are :class:`RankSummary` objects aligned with mu rows.

    Notes
    -----
    * "place" is the global position in the (front_id asc, within_score desc, index asc) order.
    * "within" is the rank within the individual's front using that same stable order.
    """

    mu = np.asarray(mu_obj, dtype=np.float64)
    sigma2 = np.asarray(sigma2_obj, dtype=np.float64)

    if mu.ndim != 2 or sigma2.ndim != 2:
        raise ValueError("mu and sigma2 must be 2D")
    if mu.shape != sigma2.shape:
        raise ValueError("mu and sigma2 must have the same shape")
    if n_samples <= 0:
        raise ValueError("n_samples must be > 0")
    if percentile <= 0.0 or percentile > 1.0:
        raise ValueError("percentile must be in (0,1]")

    N, M = mu.shape
    if N == 0:
        empty = np.zeros(0, dtype=np.int64)
        z32 = np.zeros(0, dtype=np.int32)
        z64 = np.zeros(0, dtype=np.int64)
        zf = np.zeros(0, dtype=np.float64)
        empty_summary = RankSummary(mean=zf, median=zf, std=zf, q25=zf, q75=zf)
        return empty, empty, z32, z32, z64, empty_summary, empty_summary, empty_summary

    k = int(np.ceil(percentile * N))
    k = max(1, min(k, N))

    rng = np.random.default_rng(int(seed))
    sigma = np.sqrt(np.maximum(sigma2, 0.0))

    win_count = np.zeros(N, dtype=np.int32)
    lose_count = np.zeros(N, dtype=np.int32)
    rank_sum = np.zeros(N, dtype=np.int64)

    # We store full sample matrices to provide stable median/quantile summaries.
    place_samples = np.empty((int(n_samples), N), dtype=np.int64)
    front_samples = np.empty((int(n_samples), N), dtype=np.int64)
    within_samples = np.empty((int(n_samples), N), dtype=np.int64)

    mode = _mode_from_i(int(mode_i))
    within_mode = _within_mode_from_i(int(within_mode_i))

    for s in range(int(n_samples)):
        eps = rng.standard_normal(size=(N, M))
        X = mu + sigma * eps

        rating = rate_population(X, mode=mode, within_mode=within_mode, use_cython=True)
        order = np.asarray(rating.order, dtype=np.int64)
        front_id = np.asarray(rating.front_id, dtype=np.int64)

        pos = np.empty(N, dtype=np.int64)
        pos[order] = np.arange(N, dtype=np.int64)

        # within-front rank (stable, based on the same global order)
        within = np.empty(N, dtype=np.int64)
        max_front = int(front_id.max()) if front_id.size else -1
        counts = np.zeros(max_front + 1, dtype=np.int64)
        for idx in order.tolist():
            fid = int(front_id[idx])
            within[idx] = counts[fid]
            counts[fid] += 1

        place_samples[s, :] = pos
        front_samples[s, :] = front_id
        within_samples[s, :] = within

        rank_sum += pos

        top = order[:k]
        bot = order[-k:]
        win_count[top] += 1
        lose_count[bot] += 1

    idx = np.arange(N, dtype=np.int64)
    winners_order = np.lexsort((idx, rank_sum, -win_count))
    losers_order = np.lexsort((idx, -rank_sum, -lose_count))

    winners = winners_order[:k].astype(np.int64, copy=False)
    losers = losers_order[:k].astype(np.int64, copy=False)

    place = _summarize(place_samples)
    front = _summarize(front_samples)
    within = _summarize(within_samples)

    return winners, losers, win_count, lose_count, rank_sum, place, front, within

# ---------------------------------------------------------------------------
# Detailed Monte-Carlo stats
# ---------------------------------------------------------------------------


def _summary_stats(samples: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Return (mean, median, std, q25, q75) over axis=0 as float64 arrays."""

    samples = np.asarray(samples)
    mean = samples.mean(axis=0, dtype=np.float64)
    median = np.median(samples, axis=0).astype(np.float64, copy=False)
    std = samples.std(axis=0, dtype=np.float64)
    q25 = np.quantile(samples, 0.25, axis=0).astype(np.float64, copy=False)
    q75 = np.quantile(samples, 0.75, axis=0).astype(np.float64, copy=False)
    return mean, median, std, q25, q75


def mc_rank_select_detailed(
    mu_obj,
    sigma2_obj,
    n_samples: int,
    *,
    percentile: float = 0.2,
    seed: int = 0,
    mode_i: int = 1,
    within_mode_i: int = 0,
    quantiles_mode_i: int = 2,
    quantiles_budget: int = 2_500_000,
    quality_w_place: float = 1.0,
    quality_w_front: float = 1.0,
    quality_w_within: float = 1.0,
    quality_sign: float = -1.0,
):
    """Like :func:`mc_rank_select`, but also computes per-point MC summaries.

    Returns
    -------
    winners, losers, win_count, lose_count, rank_sum, place_stats, front_stats, within_stats, quality_score

    Each *_stats is a dict with keys:
        mean, median, std, q25, q75, quantiles_mode_used

    Notes
    -----
    This pure-Python implementation tries to mimic the Cython kernel:

    * mean/std are computed online (exact).
    * median/q25/q75 are exact in quantiles_mode_i=0, and also in auto mode
      when (N*n_samples) <= quantiles_budget.
    * otherwise, a streaming P² estimator is used for each quantile (approx).
    """

    mu = np.asarray(mu_obj, dtype=np.float64)
    sigma2 = np.asarray(sigma2_obj, dtype=np.float64)

    if mu.ndim != 2 or sigma2.ndim != 2:
        raise ValueError("mu and sigma2 must be 2D")
    if mu.shape != sigma2.shape:
        raise ValueError("mu and sigma2 must have the same shape")
    if n_samples <= 0:
        raise ValueError("n_samples must be > 0")
    if percentile <= 0.0 or percentile > 1.0:
        raise ValueError("percentile must be in (0,1]")

    N, M = mu.shape
    if N == 0:
        empty_i = np.zeros(0, dtype=np.int64)
        empty32 = np.zeros(0, dtype=np.int32)
        empty64 = np.zeros(0, dtype=np.int64)
        emptyf = np.zeros(0, dtype=np.float64)
        empty_stats = {
            "mean": emptyf,
            "median": emptyf,
            "std": emptyf,
            "q25": emptyf,
            "q75": emptyf,
            "quantiles_mode_used": 0,
        }
        return empty_i, empty_i, empty32, empty32, empty64, empty_stats, empty_stats, empty_stats, emptyf

    k = int(np.ceil(percentile * N))
    k = max(1, min(k, N))

    rng = np.random.default_rng(int(seed))
    sigma = np.sqrt(np.maximum(sigma2, 0.0))

    win_count = np.zeros(N, dtype=np.int32)
    lose_count = np.zeros(N, dtype=np.int32)
    rank_sum = np.zeros(N, dtype=np.int64)

    # Online mean/std (Welford) over float64
    pl_mean = np.zeros(N, dtype=np.float64)
    pl_M2 = np.zeros(N, dtype=np.float64)
    fr_mean = np.zeros(N, dtype=np.float64)
    fr_M2 = np.zeros(N, dtype=np.float64)
    wi_mean = np.zeros(N, dtype=np.float64)
    wi_M2 = np.zeros(N, dtype=np.float64)

    # Decide quantile strategy
    use_p2 = False
    if int(quantiles_mode_i) == 0:
        use_p2 = False
    elif int(quantiles_mode_i) == 1:
        use_p2 = int(n_samples) >= 5
    else:
        if int(n_samples) < 5:
            use_p2 = False
        else:
            budget = int(quantiles_budget)
            if budget <= 0:
                use_p2 = True
            else:
                use_p2 = (int(N) * int(n_samples) > budget)

    # Exact storage when not using P²
    if not use_p2:
        place_samples = np.empty((int(n_samples), N), dtype=np.int64)
        front_samples = np.empty((int(n_samples), N), dtype=np.int64)
        within_samples = np.empty((int(n_samples), N), dtype=np.int64)

    # Streaming P² estimator (per point) for 3 quantiles.
    class _P2:
        __slots__ = ("p", "q", "n", "np", "dn", "_buf", "_count", "_init")

        def __init__(self, N: int, p: float):
            self.p = float(p)
            self.q = np.zeros((N, 5), dtype=np.float64)
            self.n = np.zeros((N, 5), dtype=np.int64)
            self.np = np.zeros((N, 5), dtype=np.float64)
            self.dn = np.array([0.0, self.p / 2.0, self.p, (1.0 + self.p) / 2.0, 1.0], dtype=np.float64)
            self._buf = np.empty((5, N), dtype=np.float64)
            self._count = 0
            self._init = False

        def observe(self, x: np.ndarray) -> None:
            x = np.asarray(x, dtype=np.float64)
            if not self._init:
                self._buf[self._count, :] = x
                self._count += 1
                if self._count == 5:
                    # init from sorted first 5 samples
                    s = np.sort(self._buf, axis=0).T  # (N,5)
                    self.q[:, :] = s
                    self.n[:, :] = np.arange(1, 6, dtype=np.int64)
                    self.np[:, 0] = 1.0
                    self.np[:, 1] = 1.0 + 2.0 * self.p
                    self.np[:, 2] = 1.0 + 4.0 * self.p
                    self.np[:, 3] = 3.0 + 2.0 * self.p
                    self.np[:, 4] = 5.0
                    self._init = True
                return

            # Update per point (scalar P²)
            for j in range(x.shape[0]):
                xv = float(x[j])
                q = self.q[j]
                n = self.n[j]
                np_ = self.np[j]

                # Find cell k
                if xv < q[0]:
                    q[0] = xv
                    k = 0
                elif xv >= q[4]:
                    q[4] = xv
                    k = 3
                else:
                    k = 0
                    if xv < q[1]:
                        k = 0
                    elif xv < q[2]:
                        k = 1
                    elif xv < q[3]:
                        k = 2
                    else:
                        k = 3

                # Increment positions
                n[k + 1 : 5] += 1
                np_ += self.dn

                # Adjust internal markers 1..3
                for i in (1, 2, 3):
                    d = np_[i] - n[i]
                    if (d >= 1.0 and (n[i + 1] - n[i]) > 1) or (d <= -1.0 and (n[i - 1] - n[i]) < -1):
                        ds = 1 if d > 0 else -1
                        # Parabolic prediction
                        n_im1 = float(n[i - 1])
                        n_i = float(n[i])
                        n_ip1 = float(n[i + 1])

                        q_im1 = q[i - 1]
                        q_i = q[i]
                        q_ip1 = q[i + 1]

                        num = ds * (
                            (n_i - n_im1 + ds) * (q_ip1 - q_i) / (n_ip1 - n_i)
                            + (n_ip1 - n_i - ds) * (q_i - q_im1) / (n_i - n_im1)
                        )
                        q_new = q_i + num / (n_ip1 - n_im1)

                        if q_im1 < q_new < q_ip1:
                            q[i] = q_new
                        else:
                            # Linear
                            q[i] = q_i + ds * (q[i + ds] - q_i) / (float(n[i + ds]) - n_i)
                        n[i] += ds

        def value(self) -> np.ndarray:
            if not self._init:
                # fewer than 5 samples: exact from buffer
                s = np.sort(self._buf[: self._count, :], axis=0)
                # approximate: pick nearest index
                if s.shape[0] == 0:
                    return np.zeros((s.shape[1],), dtype=np.float64)
                idx = int(round((s.shape[0] - 1) * self.p))
                idx = max(0, min(idx, s.shape[0] - 1))
                return s[idx, :].astype(np.float64, copy=False)
            return self.q[:, 2].astype(np.float64, copy=False)

    # P² instances per metric
    if use_p2:
        p2_place25 = _P2(N, 0.25)
        p2_place50 = _P2(N, 0.50)
        p2_place75 = _P2(N, 0.75)
        p2_front25 = _P2(N, 0.25)
        p2_front50 = _P2(N, 0.50)
        p2_front75 = _P2(N, 0.75)
        p2_within25 = _P2(N, 0.25)
        p2_within50 = _P2(N, 0.50)
        p2_within75 = _P2(N, 0.75)

    mode = _mode_from_i(int(mode_i))
    within_mode = _within_mode_from_i(int(within_mode_i))

    for s in range(int(n_samples)):
        eps = rng.standard_normal(size=(N, M))
        X = mu + sigma * eps
        rating = rate_population(X, mode=mode, within_mode=within_mode, use_cython=True)

        order = np.asarray(rating.order, dtype=np.int64)
        front_id = np.asarray(rating.front_id, dtype=np.int64)

        pos = np.empty(N, dtype=np.int64)
        pos[order] = np.arange(N, dtype=np.int64)

        within = np.empty(N, dtype=np.int64)
        max_front = int(front_id.max()) if front_id.size else -1
        counts = np.zeros(max_front + 1, dtype=np.int64)
        for idx_ in order.tolist():
            fid = int(front_id[idx_])
            within[idx_] = counts[fid]
            counts[fid] += 1

        # aggregate selection counters
        rank_sum += pos
        top = order[:k]
        bot = order[-k:]
        win_count[top] += 1
        lose_count[bot] += 1

        # Online mean/std updates (Welford)
        t = float(s + 1)
        x = pos.astype(np.float64)
        delta = x - pl_mean
        pl_mean += delta / t
        pl_M2 += delta * (x - pl_mean)

        x = front_id.astype(np.float64)
        delta = x - fr_mean
        fr_mean += delta / t
        fr_M2 += delta * (x - fr_mean)

        x = within.astype(np.float64)
        delta = x - wi_mean
        wi_mean += delta / t
        wi_M2 += delta * (x - wi_mean)

        # Quantile tracking
        if not use_p2:
            place_samples[s, :] = pos
            front_samples[s, :] = front_id
            within_samples[s, :] = within
        else:
            p2_place25.observe(pos)
            p2_place50.observe(pos)
            p2_place75.observe(pos)
            p2_front25.observe(front_id)
            p2_front50.observe(front_id)
            p2_front75.observe(front_id)
            p2_within25.observe(within)
            p2_within50.observe(within)
            p2_within75.observe(within)

    idx = np.arange(N, dtype=np.int64)
    winners_order = np.lexsort((idx, rank_sum, -win_count))
    losers_order = np.lexsort((idx, -rank_sum, -lose_count))
    winners = winners_order[:k].astype(np.int64, copy=False)
    losers = losers_order[:k].astype(np.int64, copy=False)

    # Finalize std
    pl_std = np.sqrt(pl_M2 / float(n_samples))
    fr_std = np.sqrt(fr_M2 / float(n_samples))
    wi_std = np.sqrt(wi_M2 / float(n_samples))

    if not use_p2:
        pm, pmed, ps, pq25, pq75 = _summary_stats(place_samples)
        fm, fmed, fs, fq25, fq75 = _summary_stats(front_samples)
        wm, wmed, ws, wq25, wq75 = _summary_stats(within_samples)
    else:
        pm, ps = pl_mean.copy(), pl_std.copy()
        fm, fs = fr_mean.copy(), fr_std.copy()
        wm, ws = wi_mean.copy(), wi_std.copy()

        pq25 = p2_place25.value()
        pmed = p2_place50.value()
        pq75 = p2_place75.value()

        fq25 = p2_front25.value()
        fmed = p2_front50.value()
        fq75 = p2_front75.value()

        wq25 = p2_within25.value()
        wmed = p2_within50.value()
        wq75 = p2_within75.value()

    q_mode_used = 1 if use_p2 else 0

    place_stats = {"mean": pm, "median": pmed, "std": ps, "q25": pq25, "q75": pq75, "quantiles_mode_used": q_mode_used}
    front_stats = {"mean": fm, "median": fmed, "std": fs, "q25": fq25, "q75": fq75, "quantiles_mode_used": q_mode_used}
    within_stats = {"mean": wm, "median": wmed, "std": ws, "q25": wq25, "q75": wq75, "quantiles_mode_used": q_mode_used}

    quality_score = float(quality_sign) * (
        float(quality_w_place) * pm + float(quality_w_front) * fm + float(quality_w_within) * wm
    )

    return winners, losers, win_count, lose_count, rank_sum, place_stats, front_stats, within_stats, quality_score
