Total Result¶

class SeriesKey(NamedTuple):
    """Identifies one series in the ``(model, voters, candidates)`` space."""

    gen_model: str
    n_voters: int
    n_candidates: int

@dataclass(slots=True)
class SimulationTotalResult:
    """Collection of series results spanning a parameter space.

    Each series is uniquely keyed by ``(gen_model, n_voters, n_candidates)``.
    Provides filtering, scalar-metric pivots, and comparison plots.
    """

    _entries: dict[SeriesKey, SimulationSeriesResult] = field(
        default_factory=dict,
        init=False,
        repr=False,
    )

    # ------------------------------------------------------------------
    # Mutation
    # ------------------------------------------------------------------

    def add_series(self, series: SimulationSeriesResult) -> None:
        """Register a series result.

        Raises:
            ValueError: If a series with the same key is already present.
        """
        key = _extract_key(series)
        if key in self._entries:
            raise ValueError(f"Duplicate series key {key!r}. Use replace_series() to overwrite.")
        self._entries[key] = series

    def replace_series(self, series: SimulationSeriesResult) -> None:
        """Add or overwrite a series at its config key."""
        key = _extract_key(series)
        self._entries[key] = series

    # ------------------------------------------------------------------
    # Accessors
    # ------------------------------------------------------------------

    @property
    def series_count(self) -> int:
        """Number of stored series."""
        return len(self._entries)

    @property
    def keys(self) -> list[SeriesKey]:
        """Sorted list of all series keys."""
        return sorted(self._entries)

    @property
    def gen_models(self) -> list[str]:
        """Sorted distinct generative-model codes."""
        return sorted({k.gen_model for k in self._entries})

    @property
    def voter_counts(self) -> list[int]:
        """Sorted distinct voter counts."""
        return sorted({k.n_voters for k in self._entries})

    @property
    def candidate_counts(self) -> list[int]:
        """Sorted distinct candidate counts."""
        return sorted({k.n_candidates for k in self._entries})

    @property
    def rule_codes(self) -> list[str]:
        """Distinct rule codes available across stored series.

        The order is deterministic and follows the first occurrence when
        iterating over series in sorted key order.
        """
        ordered_codes: list[str] = []
        seen: set[str] = set()
        for key in sorted(self._entries):
            for code in self._entries[key].rule_codes:
                if code not in seen:
                    seen.add(code)
                    ordered_codes.append(code)
        return ordered_codes

    @property
    def rules(self) -> list[str]:
        """Backward-compatible alias for the distinct rule codes in the result."""
        return self.rule_codes

    def get_series(self, gen_model: str, n_voters: int, n_candidates: int) -> SimulationSeriesResult:
        """Retrieve a single series by its parameter triple.

        Raises:
            KeyError: If no matching series exists.
        """
        key = SeriesKey(gen_model, n_voters, n_candidates)
        try:
            return self._entries[key]
        except KeyError:
            raise KeyError(f"No series for {key!r}") from None

    def __len__(self) -> int:
        return len(self._entries)

    def __contains__(self, key: object) -> bool:
        return key in self._entries

    def __iter__(self):  # noqa: ANN204
        """Iterate over ``(key, series)`` pairs in sorted key order."""
        yield from sorted(self._entries.items())

    def __repr__(self) -> str:
        return (
            f"SimulationTotalResult("
            f"series={self.series_count}, "
            f"models={self.gen_models}, "
            f"voters={self.voter_counts}, "
            f"candidates={self.candidate_counts})"
        )

    # ------------------------------------------------------------------
    # Filtering
    # ------------------------------------------------------------------

    def filter(
        self,
        *,
        gen_model: str | None = None,
        n_voters: int | None = None,
        n_candidates: int | None = None,
    ) -> SimulationTotalResult:
        """Return a new instance containing only the matching series.

        Series objects are shared (shallow copy), not deep-copied.
        """
        result = SimulationTotalResult()
        for key, series in self._entries.items():
            if gen_model is not None and key.gen_model != gen_model:
                continue
            if n_voters is not None and key.n_voters != n_voters:
                continue
            if n_candidates is not None and key.n_candidates != n_candidates:
                continue
            result._entries[key] = series
        return result

    def filter_rules(self, rule_codes: list[str]) -> SimulationTotalResult:
        """Return a new instance where every series is restricted to *rule_codes*.

        Applies
        :meth:`~vote_simulation.models.results.series_result.SimulationSeriesResult.filter_rules`
        to each stored series.  Entries whose filtered version contains no
        rules are dropped.  Series keys are preserved.

        Args:
            rule_codes: Rule codes to keep across all series.  Unknown or
                absent codes are silently ignored per series.

        Returns:
            A new :class:`SimulationTotalResult` with filtered series.
        """
        result = SimulationTotalResult()
        for key, series in self._entries.items():
            filtered = series.filter_rules(rule_codes)
            if filtered.step_count > 0:
                result._entries[key] = filtered
        return result

    # ------------------------------------------------------------------
    # Metrics
    # ------------------------------------------------------------------

    def summary_frame(self) -> pd.DataFrame:
        """One-row-per-series DataFrame with key fields and scalar metrics.

        Columns: ``gen_model``, ``n_voters``, ``n_candidates``,
        ``step_count``, ``n_iterations``, ``mean_distance``,
        ``most_distant_rule_a``, ``most_distant_rule_b``,
        ``most_distant_distance``.
        """
        rows: list[dict[str, Any]] = []
        for key in sorted(self._entries):
            series = self._entries[key]
            r_a, r_b, dist = series.most_distant_rules
            rows.append(
                {
                    "gen_model": key.gen_model,
                    "n_voters": key.n_voters,
                    "n_candidates": key.n_candidates,
                    "step_count": series.step_count,
                    "n_iterations": series.config.n_iterations,
                    "mean_distance": series.mean_distance,
                    "most_distant_rule_a": r_a,
                    "most_distant_rule_b": r_b,
                    "most_distant_distance": dist,
                }
            )
        return pd.DataFrame(rows)

    def metrics_comparison_frame(
        self,
        metric_field: str,
        rule_code: str,
        stat: str = "mean",
    ) -> pd.DataFrame:
        """Cross-parameter table of one winner metric for one rule.

        Builds a long-form :class:`~pandas.DataFrame` with one row per
        ``(gen_model, n_voters, n_candidates)`` entry, containing the
        requested aggregated statistic for ``rule_code``.

        Parameters
        ----------
        metric_field:
            One of the fields in
            :data:`~vote_simulation.models.rules.winner_metrics.METRIC_FIELDS`
            (e.g. ``"social_acceptability"``, ``"rank_mean"``).
        rule_code:
            The normalised rule code to extract (e.g. ``"COPE"``).
        stat:
            Which statistic column to read from
            :attr:`~vote_simulation.models.results.series_result.SimulationSeriesResult.metrics_summary_frame`:
            ``"mean"`` (default) or ``"std"``.

        Returns
        -------
        pd.DataFrame
            Columns: ``gen_model``, ``n_voters``, ``n_candidates``,
            ``<metric_field>_<stat>``.  Series without the requested rule's
            metrics are omitted.
        """
        col = f"{metric_field}_{stat}"
        rule_up = rule_code.strip().upper()
        rows: list[dict[str, Any]] = []
        for key in sorted(self._entries):
            series = self._entries[key]
            frame = series.metrics_summary_frame
            if frame.empty or rule_up not in frame.index or col not in frame.columns:
                continue
            rows.append(
                {
                    "gen_model": key.gen_model,
                    "n_voters": key.n_voters,
                    "n_candidates": key.n_candidates,
                    col: float(frame.loc[rule_up, col]),
                }
            )
        return pd.DataFrame(rows)

    def metrics_pivot(
        self,
        metric_field: str,
        rule_code: str,
        row_param: str = "n_voters",
        col_param: str = "n_candidates",
        stat: str = "mean",
    ) -> tuple[pd.DataFrame, str]:
        """Pivot :meth:`metrics_comparison_frame` into a 2D matrix.

        Analogous to :meth:`metric_matrix` but for winner metrics instead of
        rule-distance metrics.

        Parameters
        ----------
        metric_field:
            Metric field name (see :data:`~vote_simulation.models.rules.winner_metrics.METRIC_FIELDS`).
        rule_code:
            The rule to inspect.
        row_param / col_param:
            Which of the three key parameters (``gen_model``, ``n_voters``,
            ``n_candidates``) to use as row/column axes.
        stat:
            ``"mean"`` or ``"std"``.

        Returns
        -------
        (pivot_df, fixed_description)
            The pivot DataFrame and a description of the fixed third parameter.
        """
        self._validate_axis_params(row_param, col_param)
        col = f"{metric_field}_{stat}"
        df = self.metrics_comparison_frame(metric_field, rule_code, stat=stat)
        if df.empty:
            return pd.DataFrame(), ""
        third = next(iter(_PARAM_NAMES - {row_param, col_param}))
        third_vals = {str(v) for v in df[third].unique()}
        fixed_desc = f"{third}={', '.join(sorted(third_vals))}" if third_vals else ""
        pivot = df.pivot_table(index=row_param, columns=col_param, values=col, aggfunc="mean")
        return pivot, fixed_desc

    def metric_matrix(
        self,
        row_param: str = "n_voters",
        col_param: str = "n_candidates",
        *,
        metric: str = "mean_distance",
    ) -> tuple[pd.DataFrame, str]:
        """Pivot a scalar metric into a 2D matrix.

        Args:
            row_param: Key field for the row axis.
            col_param: Key field for the column axis.
            metric: Column name from :meth:`summary_frame`
                (e.g. ``"mean_distance"``, ``"most_distant_distance"``).

        Returns:
            ``(pivot_df, fixed_description)`` — the pivot DataFrame and a
            human-readable description of the fixed (third) parameter.
            If the third parameter has multiple distinct values, they are
            averaged and the description reflects this.
        """
        self._validate_axis_params(row_param, col_param)

        third = next(iter(_PARAM_NAMES - {row_param, col_param}))
        third_vals = {getattr(k, third) for k in self._entries}

        if len(third_vals) > 1:
            fixed_desc = f"{third}=avg({', '.join(str(v) for v in sorted(third_vals))})"
        else:
            fixed_desc = f"{third}={next(iter(third_vals))}" if third_vals else ""

        df = self.summary_frame()
        pivot = df.pivot_table(
            index=row_param,
            columns=col_param,
            values=metric,
            aggfunc="mean",
        )
        return pivot, fixed_desc

    def rule_pair_frame(self, rule_a: str, rule_b: str) -> pd.DataFrame:
        """Mean distance between two rules in every series.

        Returns a DataFrame with columns ``gen_model``, ``n_voters``,
        ``n_candidates``, ``distance``.
        """
        a, b = rule_a.strip().upper(), rule_b.strip().upper()
        rows: list[dict[str, Any]] = []
        for key in sorted(self._entries):
            series = self._entries[key]
            mat = series.mean_distance_matrix_frame
            if a in mat.index and b in mat.columns:
                dist = float(mat.loc[a, b])
            else:
                dist = float("nan")
            rows.append(
                {
                    "gen_model": key.gen_model,
                    "n_voters": key.n_voters,
                    "n_candidates": key.n_candidates,
                    "distance": dist,
                }
            )
        return pd.DataFrame(rows)

    def mean_distance_matrix_frame(self) -> pd.DataFrame:
        """Mean rule-distance matrix averaged across all series as a DataFrame."""
        if not self._entries:
            raise ValueError("No series in this result.")

        # Collect all rule labels (union, preserving first-seen order)
        all_rules: list[str] = []
        rule_index: dict[str, int] = {}
        for series in self._entries.values():
            for rule_code in series._rule_order:
                if rule_code not in rule_index:
                    rule_index[rule_code] = len(all_rules)
                    all_rules.append(rule_code)

        n = len(all_rules)
        acc = np.zeros((n, n), dtype=np.float64)
        counts = np.zeros((n, n), dtype=np.int64)

        for series in self._entries.values():
            mat = series.mean_distance_matrix
            if mat.size == 0:
                continue
            perm = np.array([rule_index[r] for r in series._rule_order], dtype=np.intp)
            ix = np.ix_(perm, perm)
            acc[ix] += mat.astype(np.float64)
            counts[ix] += 1

        with np.errstate(invalid="ignore"):
            avg_matrix = np.where(counts > 0, acc / counts, np.nan)

        idx = pd.Index(all_rules)
        return pd.DataFrame(avg_matrix.astype(np.float32), index=idx, columns=idx)

    def metrics_rules_matrix_frame(
        self,
        *,
        stat: str = "mean",
        metrics: list[str] | None = None,
    ) -> pd.DataFrame:
        """Winner-metric matrix with rules as rows and metric stats as columns."""
        if stat not in {"mean", "std"}:
            raise ValueError(f"Invalid stat '{stat}'. Expected 'mean' or 'std'.")

        fields = list(metrics) if metrics is not None else list(METRIC_FIELDS)
        invalid = [m for m in fields if m not in METRIC_FIELDS]
        if invalid:
            raise ValueError(f"Unknown metric fields: {invalid}")

        all_rules: list[str] = []
        for series in self._entries.values():
            frame = series.metrics_summary_frame
            if not frame.empty:
                for rule_code in frame.index:
                    if rule_code not in all_rules:
                        all_rules.append(rule_code)

        if not all_rules:
            raise ValueError("No winner-metric data found. Check that metrics were computed during simulation.")

        n_m = len(fields)
        n_r = len(all_rules)
        sums = np.zeros((n_m, n_r), dtype=np.float64)
        counts = np.zeros((n_m, n_r), dtype=np.int64)

        for series in self._entries.values():
            frame = series.metrics_summary_frame
            if frame.empty:
                continue
            for ri, rule_code in enumerate(all_rules):
                if rule_code not in frame.index:
                    continue
                for mi, field_name in enumerate(fields):
                    col = f"{field_name}_{stat}"
                    if col not in frame.columns:
                        continue
                    val = float(frame.loc[rule_code, col])
                    if np.isnan(val):
                        continue
                    sums[mi, ri] += val
                    counts[mi, ri] += 1

        with np.errstate(invalid="ignore"):
            raw = np.where(counts > 0, sums / counts, np.nan)

        valid_mask = ~np.all(np.isnan(raw), axis=1)
        fields = [f for f, v in zip(fields, valid_mask, strict=False) if v]
        raw = raw[valid_mask]
        if len(fields) == 0:
            raise ValueError("No valid metric data to export.")

        columns = [f"{field}_{stat}" for field in fields]
        return pd.DataFrame(
            raw.T,
            index=pd.Index(np.asarray(all_rules, dtype=object), name="rule"),
            columns=pd.Index(np.asarray(columns, dtype=object)),
        )

    def export_mean_distance_matrix_csv(self, save_path: str) -> str:
        """Export :meth:`mean_distance_matrix_frame` to CSV."""
        frame = self.mean_distance_matrix_frame()
        os.makedirs(os.path.dirname(os.path.abspath(save_path)), exist_ok=True)
        frame.to_csv(save_path, index=True)
        return save_path

    def export_metrics_rules_matrix_csv(
        self,
        save_path: str,
        *,
        stat: str = "mean",
        metrics: list[str] | None = None,
    ) -> str:
        """Export :meth:`metrics_rules_matrix_frame` to CSV."""
        frame = self.metrics_rules_matrix_frame(stat=stat, metrics=metrics)
        os.makedirs(os.path.dirname(os.path.abspath(save_path)), exist_ok=True)
        frame.to_csv(save_path, index=True)
        return save_path

    # ------------------------------------------------------------------
    # Plotting
    # ------------------------------------------------------------------

    def plot_mean_distance_matrix(
        self,
        ax: Any | None = None,
        *,
        annotate: bool = True,
        show: bool = True,
        save_path: str | None = None,
    ) -> Any:
        """Heatmap of the mean rule-distance matrix averaged across all series.

        Each cell ``(i, j)`` shows the average percentage of iterations where
        rules *i* and *j* disagreed, averaged over every series currently in
        the object.  Use :meth:`filter` to restrict the scope first.

        Args:
            ax: Optional matplotlib Axes. A new figure is created when *None*.
            annotate: Whether to print cell values on the heatmap.
            show: Whether to call ``plt.show()`` at the end.
            save_path: Optional file path to save the figure.

        Returns:
            The matplotlib Axes used for plotting.
        """
        if not self._entries:
            raise ValueError("No series in this result — nothing to plot.")

        matrix_frame = self.mean_distance_matrix_frame()
        all_rules = list(matrix_frame.index)
        avg_matrix = matrix_frame.to_numpy(dtype=np.float64)

        # Build title
        n_series = self.series_count

        def _avg_label(name: str, values: list[int], *, max_items: int = 7) -> str:
            shown = ",".join(str(v) for v in values[:max_items])
            suffix = ",..." if len(values) > max_items else ""
            return f"{name}=avg({shown}{suffix})"

        parts: list[str] = []
        if len(self.gen_models) == 1:
            parts.append(self.gen_models[0])
        if len(self.voter_counts) == 1:
            parts.append(f"n_votants={self.voter_counts[0]}")
        elif self.voter_counts:
            parts.append(_avg_label("n_votants", self.voter_counts))
        if len(self.candidate_counts) == 1:
            parts.append(f"n_candidates={self.candidate_counts[0]}")
        elif self.candidate_counts:
            parts.append(_avg_label("n_candidates", self.candidate_counts))
        desc = " · ".join(parts) if parts else f"{n_series} series averaged"
        title = f"Mean rule distance matrix\n{desc}"

        return _plot_heatmap(
            avg_matrix.astype(np.float32),
            all_rules,
            title,
            ax,
            annotate=annotate,
            annotation_fmt=".1f",
            colorbar_label="Mean distance (%)",
            show=show,
            save_path=save_path,
        )

    def plot_metric_heatmap(
        self,
        row_param: str = "n_voters",
        col_param: str = "n_candidates",
        *,
        metric: str = "mean_distance",
        ax: Any | None = None,
        annotate: bool = True,
        show: bool = True,
        save_path: str | None = None,
    ) -> Any:
        """Heatmap of a scalar metric pivoted across two parameters.

        The remaining (third) parameter must have a single distinct value
        — use :meth:`filter` first if needed.
        """
        import matplotlib.pyplot as plt

        pivot, fixed_desc = self.metric_matrix(row_param, col_param, metric=metric)
        matrix = pivot.to_numpy(dtype=np.float64)
        row_labels = [str(v) for v in pivot.index]
        col_labels = [str(v) for v in pivot.columns]

        vmin = float(np.nanmin(matrix))
        vmax = float(np.nanmax(matrix))
        margin = (vmax - vmin) * 0.1 or 1.0
        vmin = max(0.0, vmin - margin)
        vmax = vmax + margin

        fig_w = max(6.0, 1.2 * len(col_labels) + 2)
        fig_h = max(4.0, 1.0 * len(row_labels) + 2)
        if ax is None:
            _, ax = plt.subplots(
                figsize=(fig_w, fig_h),
                constrained_layout=True,
            )

        image = ax.imshow(
            matrix,
            cmap="YlOrRd",
            vmin=vmin,
            vmax=vmax,
            interpolation="nearest",
            aspect="auto",
        )
        ax.set_xticks(range(len(col_labels)), labels=col_labels)
        ax.set_yticks(range(len(row_labels)), labels=row_labels)
        ax.set_xlabel(col_param)
        ax.set_ylabel(row_param)

        title = metric.replace("_", " ").title()
        if fixed_desc:
            title += f"\n({fixed_desc})"
        ax.set_title(title, fontsize=11)

        if annotate:
            fs = max(6, min(12, int(160 / max(matrix.size, 1))))
            for i in range(matrix.shape[0]):
                for j in range(matrix.shape[1]):
                    val = matrix[i, j]
                    if not np.isnan(val):
                        ax.text(
                            j,
                            i,
                            f"{val:.1f}",
                            ha="center",
                            va="center",
                            fontsize=fs,
                            color="black",
                        )

        cbar = ax.figure.colorbar(
            image,
            ax=ax,
            fraction=0.046,
            pad=0.04,
            shrink=0.9,
        )
        cbar.set_label(metric.replace("_", " ").title())

        if save_path is not None:
            _save_figure(ax.figure, save_path)
        if show:
            plt.show()
        else:
            plt.close(cast(Any, ax.figure))

        return ax

    def plot_rule_pair_heatmap(
        self,
        rule_a: str,
        rule_b: str,
        row_param: str = "n_voters",
        col_param: str = "n_candidates",
        *,
        ax: Any | None = None,
        annotate: bool = True,
        show: bool = True,
        save_path: str | None = None,
    ) -> Any:
        """Heatmap of one rule-pair distance across two parameters."""
        import matplotlib.pyplot as plt

        self._validate_axis_params(row_param, col_param)

        df = self.rule_pair_frame(rule_a, rule_b)
        third = next(iter(_PARAM_NAMES - {row_param, col_param}))
        third_vals = df[third].unique()

        pivot = df.pivot_table(
            index=row_param,
            columns=col_param,
            values="distance",
            aggfunc="mean",
        )
        matrix = pivot.to_numpy(dtype=np.float64)
        row_labels = [str(v) for v in pivot.index]
        col_labels = [str(v) for v in pivot.columns]

        fig_w = max(6.0, 1.2 * len(col_labels) + 2)
        fig_h = max(4.0, 1.0 * len(row_labels) + 2)
        if ax is None:
            _, ax = plt.subplots(
                figsize=(fig_w, fig_h),
                constrained_layout=True,
            )

        image = ax.imshow(
            matrix,
            cmap="Reds",
            vmin=0,
            vmax=100,
            interpolation="nearest",
            aspect="auto",
        )
        ax.set_xticks(range(len(col_labels)), labels=col_labels)
        ax.set_yticks(range(len(row_labels)), labels=row_labels)
        ax.set_xlabel(col_param)
        ax.set_ylabel(row_param)

        a_up, b_up = rule_a.strip().upper(), rule_b.strip().upper()
        title = f"Distance: {a_up} \u2194 {b_up}"
        if len(third_vals) == 1:
            title += f"\n({third}={third_vals[0]})"
        elif len(third_vals) > 1:
            title += f"\n({third}=avg({', '.join(str(v) for v in sorted(third_vals))}))"
        ax.set_title(title, fontsize=11)

        if annotate:
            fs = max(8, min(12, int(160 / max(matrix.size, 1))))
            for i in range(matrix.shape[0]):
                for j in range(matrix.shape[1]):
                    val = matrix[i, j]
                    if not np.isnan(val):
                        ax.text(
                            j,
                            i,
                            f"{val:.1f}",
                            ha="center",
                            va="center",
                            fontsize=fs,
                            color="black",
                        )

        cbar = ax.figure.colorbar(
            image,
            ax=ax,
            fraction=0.046,
            pad=0.04,
            shrink=0.9,
        )
        cbar.set_label("Mean distance (%)")

        if save_path is not None:
            _save_figure(ax.figure, save_path)
        if show:
            plt.show()
        else:
            plt.close(cast(Any, ax.figure))

        return ax

    def plot_metrics_rules_matrix(
        self,
        row_param: str = "n_voters",
        col_param: str = "n_candidates",
        *,
        stat: str = "mean",
        metrics: list[str] | None = None,
        annotate: bool = True,
        fmt: str = ".3f",
        cell_fontsize: int | None = None,
        show: bool = True,
        save_path: str | None = None,
    ) -> Any:
        """Heatmap with rules as columns and metrics as rows.

        Each cell shows the aggregated statistic (*stat*) for a given
        ``(metric, rule)`` pair, averaged over all series currently in the
        object (use :meth:`filter` to restrict the scope first).

        The color scale is **normalised per row** (per metric) so that the
        gradient is uniform within each row and rules can be directly compared
        regardless of the absolute scale of each metric.

        Args:
            row_param: Used only to derive a description of fixed parameters
                in the title.  Filtering is the recommended way to narrow down
                the data.
            col_param: Same as *row_param*.
            stat: ``"mean"`` (default) or ``"std"``.
            metrics: Explicit list of metric fields (rows).  When *None* all
                fields in :data:`~vote_simulation.models.rules.winner_metrics.METRIC_FIELDS`
                are included.
            annotate: Whether to print raw values inside each cell.
            fmt: Format string used for annotations (e.g. ``\".3f\"``).
            cell_fontsize: Font size for cell annotations.  When *None*
                (default) the size is chosen automatically based on the number
                of rules (smaller when there are many rules).  Pass an
                explicit integer (e.g. ``8``) to override.
            show: Whether to call ``plt.show()`` at the end.
            save_path: Optional file path to save the figure.

        Returns:
            The matplotlib Axes used for plotting.
        """
        import matplotlib.pyplot as plt
        from matplotlib.cm import ScalarMappable
        from matplotlib.colors import Normalize

        matrix_frame = self.metrics_rules_matrix_frame(stat=stat, metrics=metrics)
        all_rules = list(matrix_frame.index)
        field_cols = list(matrix_frame.columns)
        fields = [col.removesuffix(f"_{stat}") for col in field_cols]
        raw = matrix_frame.to_numpy(dtype=np.float64).T
        n_r = len(all_rules)

        # Row-normalise: each row scaled to [0, 1]
        row_min = np.nanmin(raw, axis=1, keepdims=True)
        row_max = np.nanmax(raw, axis=1, keepdims=True)
        row_range = row_max - row_min
        row_range[row_range == 0] = 1.0  # avoid division by zero for constant rows
        normed = (raw - row_min) / row_range

        fig_w = max(10.0, 0.7 * len(fields) + 3.0)
        fig_h = max(6.0, 1.6 * n_r + 2)
        fig, ax = plt.subplots(figsize=(fig_w, fig_h), constrained_layout=True)

        ax.imshow(normed.T, cmap="YlGnBu", vmin=0, vmax=1, interpolation="nearest", aspect="auto")

        tick_fs = max(10, min(13, int(200 / max(len(fields), 1))))
        ax.set_xticks(
            range(len(fields)), labels=[f.replace("_", " ") for f in fields], rotation=45, ha="right", fontsize=tick_fs
        )
        ax.set_yticks(range(n_r), labels=all_rules, fontsize=11)
        ax.set_xlabel("Metric", fontsize=12)
        ax.set_ylabel("Rule", fontsize=12)

        # Build title from fixed params info
        n_series = self.series_count
        fixed_parts = []
        if len(self.gen_models) == 1:
            fixed_parts.append(f"model={self.gen_models[0]}")
        if len(self.voter_counts) == 1:
            fixed_parts.append(f"n_voters={self.voter_counts[0]}")
        if len(self.candidate_counts) == 1:
            fixed_parts.append(f"n_candidates={self.candidate_counts[0]}")
        desc = " · ".join(fixed_parts) if fixed_parts else f"{n_series} series averaged"
        ax.set_title(
            f"Winner metrics ({stat}) per rule\n{desc}",
            fontsize=13,
        )

        if annotate:
            fs = cell_fontsize if cell_fontsize is not None else max(12, min(12, int(120 / max(len(fields), 1))))
            for mi in range(len(fields)):
                for ri in range(n_r):
                    val = raw[mi, ri]
                    if not np.isnan(val):
                        cell_norm = normed[mi, ri]
                        txt_color = "white" if cell_norm > 0.65 else "black"
                        ax.text(mi, ri, format(val, fmt), ha="center", va="center", fontsize=fs, color=txt_color)

        cbar = fig.colorbar(
            ScalarMappable(norm=Normalize(0, 1), cmap="YlGnBu"),
            ax=ax,
            location="bottom",
            fraction=0.05,
            pad=0.05,
            shrink=0.6,
            aspect=40,
        )
        cbar.set_label("Normalised value (per metric)", fontsize=12)
        cbar.set_ticks([0, 0.5, 1])
        cbar.set_ticklabels(["min", "mid", "max"], fontsize=12)
        cbar.ax.tick_params(labelsize=12)

        if save_path is not None:
            os.makedirs(os.path.dirname(os.path.abspath(save_path)), exist_ok=True)
            fig.savefig(save_path)
        if show:
            plt.show()
        else:
            plt.close(fig)
        return ax

    def plot_winner_metric_heatmap(
        self,
        metric_field: str,
        rule_code: str,
        row_param: str = "n_voters",
        col_param: str = "n_candidates",
        *,
        stat: str = "mean",
        ax: Any | None = None,
        annotate: bool = True,
        show: bool = True,
        save_path: str | None = None,
    ) -> Any:
        """Heatmap of one winner metric for one rule, pivoted across two parameters.

        Analogous to :meth:`plot_metric_heatmap` but for
        :class:`~vote_simulation.models.rules.WinnerMetrics` fields instead of
        rule-distance metrics.

        Args:
            metric_field: One of the fields in
                :data:`~vote_simulation.models.rules.winner_metrics.METRIC_FIELDS`
                (e.g. ``"social_acceptability"``, ``"rank_mean"``).
            rule_code: Normalised rule code to inspect (e.g. ``"COPE"``).
            row_param: Row axis — ``"gen_model"``, ``"n_voters"``, or
                ``"n_candidates"``.
            col_param: Column axis (same choices).
            stat: ``"mean"`` (default) or ``"std"``.
            ax: Optional matplotlib Axes.  A new figure is created when *None*.
            annotate: Whether to print cell values on the heatmap.
            show: Whether to call ``plt.show()`` at the end.
            save_path: Optional file path to save the figure.

        Returns:
            The matplotlib Axes used for plotting.
        """
        import matplotlib.pyplot as plt

        pivot, fixed_desc = self.metrics_pivot(
            metric_field,
            rule_code,
            row_param=row_param,
            col_param=col_param,
            stat=stat,
        )
        if pivot.empty:
            raise ValueError(
                f"No data for metric '{metric_field}' / rule '{rule_code}'. "
                "Check that metrics were computed during simulation."
            )

        # Remettre row_param/col_param dans l'ordre naturel (pas de transposition)
        matrix = pivot.to_numpy(dtype=np.float64)
        row_labels = [str(v) for v in pivot.index]
        col_labels = [str(v) for v in pivot.columns]

        vmin = float(np.nanmin(matrix))
        vmax = float(np.nanmax(matrix))
        margin = (vmax - vmin) * 0.1 or 0.01
        vmin = max(0.0, vmin - margin)
        vmax = (
            min(1.0, vmax + margin)
            if metric_field
            not in {
                "rank_mean",
                "rank_median",
                "rank_var",
                "utility_mean",
                "utility_median",
                "utility_var",
                "n_cowinners",
            }
            else vmax + margin
        )

        fig_w = max(6.0, 1.2 * len(col_labels) + 2)
        fig_h = max(4.0, 1.0 * len(row_labels) + 2)
        if ax is None:
            _, ax = plt.subplots(figsize=(fig_w, fig_h), constrained_layout=True)

        image = ax.imshow(
            matrix,
            cmap="YlGnBu",
            vmin=vmin,
            vmax=vmax,
            interpolation="nearest",
            aspect="auto",
        )
        ax.set_xticks(range(len(col_labels)), labels=col_labels)
        ax.set_yticks(range(len(row_labels)), labels=row_labels)
        ax.set_xlabel(col_param)
        ax.set_ylabel(row_param)

        rule_up = rule_code.strip().upper()
        title = f"{metric_field.replace('_', ' ').title()} ({stat}) — {rule_up}"
        if fixed_desc:
            title += f"\n({fixed_desc})"
        ax.set_title(title, fontsize=11)

        if annotate:
            fs = max(6, min(12, int(160 / max(matrix.size, 1))))
            for i in range(matrix.shape[0]):
                for j in range(matrix.shape[1]):
                    val = matrix[i, j]
                    if not np.isnan(val):
                        ax.text(j, i, f"{val:.3f}", ha="center", va="center", fontsize=fs, color="black")

        cbar = ax.figure.colorbar(image, ax=ax, fraction=0.046, pad=0.04, shrink=0.9)
        cbar.set_label(f"{metric_field.replace('_', ' ').title()} ({stat})")

        if save_path is not None:
            _save_figure(ax.figure, save_path)
        if show:
            plt.show()
        else:
            plt.close(cast(Any, ax.figure))
        return ax

    def plot_winner_metrics_grid(
        self,
        rule_code: str,
        row_param: str = "n_voters",
        col_param: str = "n_candidates",
        *,
        stat: str = "mean",
        metrics: list[str] | None = None,
        annotate: bool = True,
        show: bool = True,
        save_path: str | None = None,
    ) -> Any:
        """Grid of heatmaps — one panel per winner metric — for a single rule.

        Iterates over every field in
        :data:`~vote_simulation.models.rules.winner_metrics.METRIC_FIELDS` (or
        the subset given by *metrics*) and draws each as a heatmap panel,
        using the same axes as :meth:`plot_winner_metric_heatmap`.

        Args:
            rule_code: Normalised rule code to inspect (e.g. ``"COPE"``).
            row_param: Row axis — ``"gen_model"``, ``"n_voters"``, or
                ``"n_candidates"``.
            col_param: Column axis (same choices).
            stat: ``"mean"`` (default) or ``"std"``.
            metrics: Explicit list of metric fields to plot.  When *None* all
                fields in :data:`~vote_simulation.models.rules.winner_metrics.METRIC_FIELDS`
                are included.
            annotate: Whether to print cell values on each panel.
            show: Whether to call ``plt.show()`` at the end.
            save_path: Optional file path to save the whole figure.

        Returns:
            The 2-D NumPy array of matplotlib Axes.
        """
        import math

        import matplotlib.pyplot as plt

        fields = list(metrics) if metrics is not None else list(METRIC_FIELDS)

        # Drop metrics with no data
        available = [f for f in fields if not self.metrics_comparison_frame(f, rule_code, stat=stat).empty]
        if not available:
            raise ValueError(
                f"No winner-metric data found for rule '{rule_code}'. "
                "Check that metrics were computed during simulation."
            )

        n = len(available)
        ncols = min(4, n)
        nrows = math.ceil(n / ncols)

        fig, axes = plt.subplots(
            nrows,
            ncols,
            figsize=(5 * ncols, 4 * nrows),
            constrained_layout=True,
        )
        # Normalise axes to a flat list
        axes_flat: list[Any] = np.array(axes).flatten().tolist()

        rule_up = rule_code.strip().upper()
        fig.suptitle(
            f"Winner metrics ({stat}) — {rule_up}",
            fontsize=13,
            fontweight="bold",
        )

        for ax_i, f_name in zip(axes_flat, available, strict=False):
            try:
                self.plot_winner_metric_heatmap(
                    f_name,
                    rule_code,
                    row_param=row_param,
                    col_param=col_param,
                    stat=stat,
                    ax=ax_i,
                    annotate=annotate,
                    show=False,
                )
            except ValueError:
                ax_i.set_visible(False)

        # Hide unused axes
        for ax_i in axes_flat[len(available) :]:
            ax_i.set_visible(False)

        if save_path is not None:
            os.makedirs(os.path.dirname(os.path.abspath(save_path)), exist_ok=True)
            fig.savefig(save_path)
        if show:
            plt.show()
        else:
            plt.close(fig)
        return axes

    def map_rules_2d(self) -> MdsProjection:
        """Project rules into 2D using Multi-Dimensional Scaling (MDS).

        Uses the global mean distance matrix (averaged over all series) as a
        precomputed dissimilarity matrix.

        Returns:
            :class:`MdsProjection` with 2D coordinates and normalized stress.

        Raises:
            ValueError: If no series have been added yet.
        """
        if not self._entries:
            raise ValueError("No series in this result.")
        matrix = self.mean_distance_matrix_frame().to_numpy(dtype=np.float64)
        return _mds_project(matrix, n_components=2)

    def map_rules_3d(self) -> MdsProjection:
        """Project rules into 3D using Multi-Dimensional Scaling (MDS).

        Uses the global mean distance matrix (averaged over all series) as a
        precomputed dissimilarity matrix.

        Returns:
            :class:`MdsProjection` with 3D coordinates and normalized stress.

        Raises:
            ValueError: If no series have been added yet.
        """
        if not self._entries:
            raise ValueError("No series in this result.")
        matrix = self.mean_distance_matrix_frame().to_numpy(dtype=np.float64)
        return _mds_project(matrix, n_components=3)

    def plot_rules_2d(
        self,
        ax: Any | None = None,
        *,
        show: bool = True,
        save_path: str | None = None,
    ) -> Any:
        """Plot rules as labeled points in a 2D MDS projection.

        Distances between points approximate the global mean pairwise rule
        distances, averaged over all series in this result.  Use
        :meth:`filter` to restrict the scope first.

        Args:
            ax: Optional matplotlib Axes. A new figure is created when *None*.
            show: Whether to call ``plt.show()`` at the end.
            save_path: Optional file path to save the figure.

        Returns:
            The matplotlib Axes used for plotting.
        """
        if not self._entries:
            raise ValueError("No series in this result — nothing to plot.")
        projection = self.map_rules_2d()
        labels = list(self.mean_distance_matrix_frame().index)
        desc = self._build_plot_desc()
        title = f"Rule proximity map\n{desc}"
        if save_path is not None:
            os.makedirs(os.path.dirname(os.path.abspath(save_path)), exist_ok=True)
        return _plot_rules_2d_scatter(projection.coords, labels, title, ax, show=show, save_path=save_path)

    def plot_rules_3d(
        self,
        ax: Any | None = None,
        *,
        show: bool = True,
        save_path: str | None = None,
    ) -> Any:
        """Plot rules as labeled points in a 3D MDS projection.

        Distances between points approximate the global mean pairwise rule
        distances, averaged over all series in this result.  Use
        :meth:`filter` to restrict the scope first.

        Args:
            ax: Optional matplotlib 3D Axes. A new figure is created when *None*.
            show: Whether to call ``plt.show()`` at the end.
            save_path: Optional file path to save the figure.

        Returns:
            The matplotlib Axes used for plotting.
        """
        if not self._entries:
            raise ValueError("No series in this result — nothing to plot.")
        projection = self.map_rules_3d()
        labels = list(self.mean_distance_matrix_frame().index)
        desc = self._build_plot_desc()
        title = f"Rule proximity map (3D)\n{desc}\nMDS stress: {projection.stress:.4f}"
        if save_path is not None:
            os.makedirs(os.path.dirname(os.path.abspath(save_path)), exist_ok=True)
        return _plot_rules_3d_scatter(projection.coords, labels, title, ax, show=show, save_path=save_path)

    def plot_comparison_grid(
        self,
        vary_param: str = "n_candidates",
        *,
        annotate: bool = True,
        show: bool = True,
        save_path: str | None = None,
    ) -> Any:
        """Side-by-side rule-distance heatmaps, one per value of *vary_param*.

        The other two parameters must each have a single distinct value
        — use :meth:`filter` first.

        Example::

            total.filter(gen_model="UNI", n_voters=1001) \\
                 .plot_comparison_grid("n_candidates")
        """
        import matplotlib.pyplot as plt

        if vary_param not in _PARAM_NAMES:
            raise ValueError(f"Invalid param '{vary_param}'. Choose from {sorted(_PARAM_NAMES)}")

        fixed_params = _PARAM_NAMES - {vary_param}

        vary_values = sorted({getattr(k, vary_param) for k in self._entries})
        n = len(vary_values)
        if n == 0:
            raise ValueError("No series to plot.")

        fig, axes = plt.subplots(
            1,
            n,
            figsize=(6 * n + 1, 6),
            constrained_layout=True,
        )
        if n == 1:
            axes = [axes]

        for ax_i, val in zip(axes, vary_values, strict=True):
            matching = [s for k, s in self._entries.items() if getattr(k, vary_param) == val]
            if not matching:
                continue
            # Average distance matrices when multiple series match
            avg_matrix = np.mean(
                np.stack([s.mean_distance_matrix for s in matching]),
                axis=0,
            )
            labels = matching[0]._rule_order
            subtitle = f"{vary_param}={val}"
            if len(matching) > 1:
                subtitle += f" (avg. {len(matching)})"
            _plot_heatmap(
                avg_matrix,
                labels,
                subtitle,
                ax=ax_i,
                annotate=annotate,
                annotation_fmt=".1f",
                colorbar_label="Mean distance (%)",
                show=False,
            )

        # Build description of fixed parameters
        fixed_parts: list[str] = []
        for p in sorted(fixed_params):
            vals = sorted({getattr(k, p) for k in self._entries}, key=str)
            if len(vals) == 1:
                fixed_parts.append(f"{p}={vals[0]}")
            else:
                fixed_parts.append(f"{p}: averaged")
        fixed_desc = " \u00b7 ".join(fixed_parts)

        fig.suptitle(
            f"Rule distance comparison\n{fixed_desc}",
            fontsize=13,
        )

        if save_path is not None:
            os.makedirs(os.path.dirname(os.path.abspath(save_path)), exist_ok=True)
            fig.savefig(save_path)
        if show:
            plt.show()
        else:
            plt.close(fig)

        return axes

    # ------------------------------------------------------------------
    # Persistence
    # ------------------------------------------------------------------

    def save_to_dir(self, dir_path: str) -> None:
        """Persist every series as ``<label>.parquet`` inside *dir_path*."""
        os.makedirs(dir_path, exist_ok=True)
        for _key, series in self._entries.items():
            filename = f"{series.config.label}.parquet"
            series.save_to_file(os.path.join(dir_path, filename))

    @classmethod
    def load_from_dir(cls, dir_path: str) -> SimulationTotalResult:
        """Reconstruct a :class:`SimulationTotalResult` from a folder of parquet files.

        Each ``.parquet`` file is loaded as a :class:`SimulationSeriesResult`
        and registered via :meth:`add_series`.
        """
        import glob as _glob

        total = cls()
        paths = sorted(_glob.glob(os.path.join(dir_path, "*.parquet")))
        if not paths:
            raise ValueError(f"No .parquet files found in {dir_path!r}")
        for path in paths:
            series = SimulationSeriesResult()
            series.load_from_file(path)
            total.add_series(series)
        return total

    @staticmethod
    def delete_dir(dir_path: str) -> bool:
        """Remove a saved total-result directory.

        Returns ``True`` if the directory existed and was deleted.
        """
        import shutil

        try:
            shutil.rmtree(dir_path)
            return True
        except FileNotFoundError:
            return False

    # Private helpers

    def _build_plot_desc(self) -> str:
        """Build a human-readable description of current parameter values for plot titles."""

        def _avg_label(name: str, values: list[int], *, max_items: int = 7) -> str:
            shown = ",".join(str(v) for v in values[:max_items])
            suffix = ",..." if len(values) > max_items else ""
            return f"{name}=avg({shown}{suffix})"

        parts: list[str] = []
        if len(self.gen_models) == 1:
            parts.append(self.gen_models[0])
        if len(self.voter_counts) == 1:
            parts.append(f"n_votants={self.voter_counts[0]}")
        elif self.voter_counts:
            parts.append(_avg_label("n_votants", self.voter_counts))
        if len(self.candidate_counts) == 1:
            parts.append(f"n_candidates={self.candidate_counts[0]}")
        elif self.candidate_counts:
            parts.append(_avg_label("n_candidates", self.candidate_counts))
        return " \u00b7 ".join(parts) if parts else f"{self.series_count} series averaged"

    @staticmethod
    def _validate_axis_params(row_param: str, col_param: str) -> None:
        """Ensure *row_param* and *col_param* are valid and distinct."""
        if row_param not in _PARAM_NAMES:
            raise ValueError(f"Invalid row_param '{row_param}'. Choose from {sorted(_PARAM_NAMES)}")
        if col_param not in _PARAM_NAMES:
            raise ValueError(f"Invalid col_param '{col_param}'. Choose from {sorted(_PARAM_NAMES)}")
        if row_param == col_param:
            raise ValueError("row_param and col_param must be different.")