from dataclasses import replace

import networkx as nx
import numpy as np

from socialdynamics.experiments import aggregate, beta_sweep, run_ensemble, summarise_sweep
from socialdynamics.presets import PRESETS
from socialdynamics.simulation import OpinionSimulation, SimulationConfig


def _preset_config(name: str, **overrides) -> SimulationConfig:
    p = dict(PRESETS[name])
    p["misinformation_direction"] = 1.0 if p["misinformation_direction"].startswith("Positive") else -1.0
    p.update(overrides)
    return SimulationConfig(**p)


def _reference_step_v11(sim: OpinionSimulation, old: np.ndarray) -> np.ndarray:
    """The v1.1 per-agent loop, kept verbatim as a regression oracle."""
    cfg = sim.config
    new = old.copy()
    beta = 6.0 * cfg.confirmation_bias
    for agent in range(cfg.n_agents):
        if sim.stubborn_mask[agent]:
            continue
        neighbors = list(sim.graph.neighbors(agent))
        if neighbors:
            w = np.exp(-beta * np.abs(old[agent] - old[neighbors]))
            total = float(np.sum(w))
            target = float(np.sum(w * old[neighbors]) / total) if total > 1e-12 else float(old[agent])
        else:
            target = float(old[agent])
        delta = cfg.social_influence * (target - old[agent])
        if sim.misinformation_mask[agent] and cfg.misinformation_strength > 0:
            delta += 0.22 * cfg.misinformation_strength * (cfg.misinformation_direction - old[agent])
        if cfg.noise > 0:
            delta += float(sim._ref_rng.normal(0.0, cfg.noise))
        new[agent] = float(np.clip(old[agent] + delta, -1.0, 1.0))
    return new


def test_vectorised_update_matches_v11_loop():
    for kw in ({}, {"network": "Scale Free", "initial_opinion": "Polarized"}, {"stubborn_fraction": 0.2}):
        cfg = SimulationConfig(steps=40, seed=17, **kw)
        fast = OpinionSimulation(cfg)
        ref = OpinionSimulation(cfg)
        ref._ref_rng = ref.rng
        x = ref.opinions.copy()
        for _ in range(cfg.steps):
            fast.step()
            x = _reference_step_v11(ref, x)
        assert np.max(np.abs(fast.opinions - x)) < 1e-12


def test_simulation_is_seed_deterministic():
    for rate in (0.0, 0.2):
        cfg = SimulationConfig(n_agents=100, steps=20, seed=123, rewiring_rate=rate)
        a, b = OpinionSimulation(cfg), OpinionSimulation(cfg)
        a.run()
        b.run()
        assert np.array_equal(a.history_array, b.history_array)
        assert set(map(frozenset, a.graph.edges())) == set(map(frozenset, b.graph.edges()))


def test_opinions_stay_bounded():
    cfg = SimulationConfig(n_agents=100, steps=50, seed=3, misinformation_exposure=1.0,
                           misinformation_strength=1.0, social_influence=1.0, noise=0.2)
    sim = OpinionSimulation(cfg)
    sim.run()
    assert np.all(sim.history_array >= -1.0)
    assert np.all(sim.history_array <= 1.0)


def test_stubborn_agents_never_move():
    cfg = SimulationConfig(n_agents=100, steps=30, seed=4, stubborn_fraction=0.20, stubborn_mode="Split extremes",
                           misinformation_exposure=0.50, misinformation_strength=0.8, rewiring_rate=0.3)
    sim = OpinionSimulation(cfg)
    initial = sim.opinions.copy()
    mask = sim.stubborn_mask.copy()
    sim.run()
    assert np.array_equal(sim.opinions[mask], initial[mask])


def test_misinformation_can_shift_mean_positive():
    base = dict(n_agents=100, network="Random", initial_opinion="Neutral", stubborn_fraction=0.0,
                confirmation_bias=0.2, social_influence=0.4, noise=0.0, steps=30, seed=11)
    control = OpinionSimulation(SimulationConfig(**base, misinformation_exposure=0.0, misinformation_strength=0.0))
    treatment = OpinionSimulation(SimulationConfig(**base, misinformation_exposure=1.0, misinformation_strength=0.8))
    control.run()
    treatment.run()
    assert treatment.opinions.mean() > control.opinions.mean() + 0.1


def test_rewiring_conserves_edges_and_creates_no_isolates():
    sim = OpinionSimulation(_preset_config("Echo Chamber"))
    edges0 = sim.graph.number_of_edges()
    degree0 = dict(sim.graph.degree())
    sim.run()
    assert sim.rewired_ties > 0
    assert sim.graph.number_of_edges() == edges0
    assert not list(nx.isolates(sim.graph))
    assert nx.number_of_selfloops(sim.graph) == 0
    assert sum(dict(sim.graph.degree()).values()) == sum(degree0.values())


def test_rewiring_does_not_perturb_the_noise_stream():
    cfg = SimulationConfig(steps=1, seed=8)
    a = OpinionSimulation(cfg)
    b = OpinionSimulation(replace(cfg, rewiring_rate=0.5))
    a.step()
    b.step()
    assert np.array_equal(a.opinions, b.opinions)  # first update precedes any rewiring


def test_filter_alone_does_not_segregate_but_rewiring_does():
    seeds = [42, 43, 44]
    filtered = run_ensemble(_preset_config("Filtered, Not Segregated"), seeds, n_permutations=100)
    chamber = run_ensemble(_preset_config("Echo Chamber"), seeds, n_permutations=100)
    for run in filtered:
        assert run["selective_exposure"] > 0.9  # strong filter...
        assert abs(run["assortativity"]) < 0.2  # ...but no clustering
        assert run["assortativity_z"] < 3.0
    for run in chamber:
        assert run["assortativity"] > 0.9
        assert run["assortativity_z"] > 10.0
        # and the old index would have ranked this *below* the unsegregated case
        assert run["selective_exposure"] < min(r["selective_exposure"] for r in filtered)


def test_sweep_and_aggregate_shapes():
    base = SimulationConfig(steps=20)
    rows = beta_sweep(base, (0.0, 1.0), (0.0, 0.1), seeds=[1, 2], n_permutations=20)
    assert len(rows) == 8
    summary = summarise_sweep(rows, ("assortativity", "selective_exposure"))
    assert len(summary) == 4
    assert all(r["n"] == 2 for r in summary)
    agg = aggregate(rows, ("assortativity",))
    assert agg["assortativity_ci95"] >= 0.0