src/socialdynamics/experiments.py

from __future__ import annotations

from dataclasses import asdict, replace
from typing import Iterable, Sequence

import numpy as np

from .simulation import OpinionSimulation, SimulationConfig

FINAL_KEYS = (
    "mean_opinion",
    "polarization",
    "bimodality",
    "assortativity",
    "neighbour_corr",
    "selective_exposure",
    "components",
    "rewired_ties",
)


def run_once(config: SimulationConfig, n_permutations: int = 200) -> dict[str, float]:
    """Run one configuration to its horizon and return final-state diagnostics."""
    sim = OpinionSimulation(config)
    sim.run()
    final = {k: float(sim.metric_history[-1][k]) for k in FINAL_KEYS}
    final["assortativity_initial"] = float(sim.metric_history[0]["assortativity"])
    null = sim.null_test(n_permutations=n_permutations)
    final["assortativity_z"] = null.z
    final["assortativity_p"] = null.p_value
    final["seed"] = float(config.seed)
    return final


def seeds_from(base_seed: int, n_seeds: int) -> list[int]:
    return [int(base_seed) + k for k in range(max(1, int(n_seeds)))]


def run_ensemble(
    config: SimulationConfig, seeds: Sequence[int], n_permutations: int = 200
) -> list[dict[str, float]]:
    return [run_once(replace(config, seed=int(s)), n_permutations) for s in seeds]


def aggregate(rows: Iterable[dict[str, float]], keys: Sequence[str] | None = None) -> dict[str, float]:
    """Mean, sample SD and 95% t-free normal-approximation half-width per key."""
    rows = list(rows)
    keys = list(keys or rows[0].keys())
    out: dict[str, float] = {"n": float(len(rows))}
    for key in keys:
        values = np.array([r[key] for r in rows], dtype=float)
        out[f"{key}_mean"] = float(values.mean())
        sd = float(values.std(ddof=1)) if values.size > 1 else 0.0
        out[f"{key}_sd"] = sd
        out[f"{key}_ci95"] = 1.96 * sd / np.sqrt(values.size) if values.size > 1 else 0.0
    return out


def beta_sweep(
    base: SimulationConfig,
    betas: Sequence[float],
    rewiring_rates: Sequence[float],
    seeds: Sequence[int],
    n_permutations: int = 100,
) -> list[dict[str, float]]:
    """Sweep confirmation bias x rewiring rate; one row per (beta, rewiring, seed)."""
    rows: list[dict[str, float]] = []
    for rate in rewiring_rates:
        for beta in betas:
            cfg = replace(base, confirmation_bias=float(beta), rewiring_rate=float(rate))
            for row in run_ensemble(cfg, seeds, n_permutations):
                row["confirmation_bias"] = float(beta)
                row["rewiring_rate"] = float(rate)
                rows.append(row)
    return rows


def summarise_sweep(rows: list[dict[str, float]], keys: Sequence[str]) -> list[dict[str, float]]:
    groups: dict[tuple[float, float], list[dict[str, float]]] = {}
    for row in rows:
        groups.setdefault((row["rewiring_rate"], row["confirmation_bias"]), []).append(row)
    summary = []
    for (rate, beta), members in sorted(groups.items()):
        agg = aggregate(members, keys)
        agg["rewiring_rate"] = rate
        agg["confirmation_bias"] = beta
        summary.append(agg)
    return summary


def config_record(config: SimulationConfig) -> dict:
    """Serializable config for provenance in exported files."""
    return asdict(config)