from __future__ import annotations

from dataclasses import dataclass
from typing import Iterator

import networkx as nx
import numpy as np

from .metrics import NullTest, assortativity_null, summarize_metrics
from .networks import build_network


@dataclass(frozen=True)
class SimulationConfig:
    n_agents: int = 100
    network: str = "Small World"
    initial_opinion: str = "Random"
    stubborn_fraction: float = 0.05
    stubborn_mode: str = "Split extremes"
    misinformation_exposure: float = 0.25
    misinformation_strength: float = 0.35
    misinformation_direction: float = 1.0
    confirmation_bias: float = 0.45
    social_influence: float = 0.45
    noise: float = 0.02
    steps: int = 120
    seed: int = 42
    average_degree: int = 8
    # Adaptive network (v1.2). rewiring_rate = 0 reproduces v1.1 exactly.
    rewiring_rate: float = 0.0
    rewiring_tolerance: float = 0.5


class OpinionSimulation:
    """Synchronous continuous-opinion simulation on a (optionally adaptive) social graph.

    Each step has two phases:

    1. Opinion update (all adaptive agents simultaneously), using the
       distance-weighted neighbourhood target described in the README.
    2. Optional homophilic rewiring. Each adaptive agent, with probability
       ``rewiring_rate``, inspects one uniformly chosen neighbour. If their
       opinion distance exceeds ``rewiring_tolerance`` the tie is dropped and
       replaced with a tie to a uniformly chosen non-neighbour within the
       tolerance (Holme & Newman, 2006). The number of edges and the
       initiating agent's degree are conserved; ties are never removed from an
       agent of degree one.

    Opinion noise and rewiring draw from independent random streams, so
    turning rewiring on does not perturb the noise sequence.
    """

    def __init__(self, config: SimulationConfig):
        self.config = config
        self._validate()
        self.rng = np.random.default_rng(config.seed)
        self.rewire_rng = np.random.default_rng([config.seed, 20260927])
        self.graph: nx.Graph = build_network(
            config.network, config.n_agents, config.seed, config.average_degree
        )
        self.initial_graph: nx.Graph = self.graph.copy()
        self.graph_version = 0
        self.rewired_ties = 0
        self.opinions = self._initialize_opinions()
        self.stubborn_mask = self._choose_mask(config.stubborn_fraction)
        self._apply_stubborn_mode()
        self.initial_opinions = self.opinions.copy()
        self.misinformation_mask = self._choose_mask(config.misinformation_exposure, ~self.stubborn_mask)
        self._adaptive = np.flatnonzero(~self.stubborn_mask)
        self._rebuild_edge_cache()
        self.history: list[np.ndarray] = [self.opinions.copy()]
        self.metric_history: list[dict[str, float]] = [self.metrics(step=0)]

    # ── setup ──────────────────────────────────────────────────────────────
    def _validate(self) -> None:
        config = self.config
        if not 4 <= int(config.n_agents) <= 5000:
            raise ValueError("n_agents must be between 4 and 5000")
        for name in (
            "stubborn_fraction",
            "misinformation_exposure",
            "misinformation_strength",
            "confirmation_bias",
            "social_influence",
            "rewiring_rate",
        ):
            value = float(getattr(config, name))
            if not 0.0 <= value <= 1.0:
                raise ValueError(f"{name} must be in [0, 1]")
        if not 0.0 < config.rewiring_tolerance <= 2.0:
            raise ValueError("rewiring_tolerance must be in (0, 2]")
        if not 0.0 <= config.noise <= 0.2:
            raise ValueError("noise must be in [0, 0.2]")
        if config.steps < 1:
            raise ValueError("steps must be >= 1")
        if config.misinformation_direction not in (-1.0, 1.0):
            raise ValueError("misinformation_direction must be -1 or +1")

    def _initialize_opinions(self) -> np.ndarray:
        mode = self.config.initial_opinion.lower()
        n_agents = self.config.n_agents
        if mode == "random":
            values = self.rng.uniform(-1.0, 1.0, size=n_agents)
        elif mode == "neutral":
            values = self.rng.normal(0.0, 0.16, size=n_agents)
        elif mode == "polarized":
            signs = self.rng.choice([-1.0, 1.0], size=n_agents)
            values = signs * self.rng.normal(0.72, 0.13, size=n_agents)
        else:
            raise ValueError(f"Unknown initial opinion mode: {self.config.initial_opinion}")
        return np.clip(values, -1.0, 1.0).astype(np.float64)

    def _choose_mask(self, fraction: float, eligible: np.ndarray | None = None) -> np.ndarray:
        n_agents = self.config.n_agents
        mask = np.zeros(n_agents, dtype=bool)
        if fraction <= 0:
            return mask
        candidates = np.arange(n_agents) if eligible is None else np.flatnonzero(eligible)
        count = min(len(candidates), int(round(fraction * n_agents)))
        if count > 0:
            picked = self.rng.choice(candidates, size=count, replace=False)
            mask[picked] = True
        return mask

    def _apply_stubborn_mode(self) -> None:
        indices = np.flatnonzero(self.stubborn_mask)
        if indices.size == 0:
            return
        mode = self.config.stubborn_mode.lower()
        if mode == "keep initial":
            return
        if mode == "positive (+1)":
            self.opinions[indices] = 1.0
        elif mode == "negative (-1)":
            self.opinions[indices] = -1.0
        elif mode == "split extremes":
            shuffled = self.rng.permutation(indices)
            split = len(shuffled) // 2
            self.opinions[shuffled[:split]] = -1.0
            self.opinions[shuffled[split:]] = 1.0
        else:
            raise ValueError(f"Unknown stubborn mode: {self.config.stubborn_mode}")

    def _rebuild_edge_cache(self) -> None:
        edges = np.asarray(self.graph.edges(), dtype=np.intp).reshape(-1, 2)
        self._src = np.concatenate((edges[:, 0], edges[:, 1]))
        self._dst = np.concatenate((edges[:, 1], edges[:, 0]))
        self._degree = np.bincount(self._src, minlength=self.config.n_agents)

    # ── diagnostics ────────────────────────────────────────────────────────
    def metrics(self, step: int | None = None) -> dict[str, float]:
        values = summarize_metrics(self.graph, self.opinions, self.config.confirmation_bias)
        values["step"] = float(len(self.history) - 1 if step is None else step)
        values["rewired_ties"] = float(self.rewired_ties)
        return values

    def null_test(self, n_permutations: int = 200) -> NullTest:
        """Permutation test of the current opinion clustering on the current graph."""
        return assortativity_null(
            self.graph, self.opinions, n_permutations=n_permutations, seed=self.config.seed + 7919
        )

    # ── dynamics ───────────────────────────────────────────────────────────
    def _opinion_update(self) -> np.ndarray:
        cfg = self.config
        old = self.opinions
        n = cfg.n_agents
        beta = 6.0 * cfg.confirmation_bias

        neighbour_values = old[self._dst]
        weights = np.exp(-beta * np.abs(old[self._src] - neighbour_values))
        total = np.bincount(self._src, weights=weights, minlength=n)
        weighted = np.bincount(self._src, weights=weights * neighbour_values, minlength=n)
        has_target = (self._degree > 0) & (total > 1e-12)
        social_target = np.where(has_target, weighted / np.where(has_target, total, 1.0), old)

        delta = cfg.social_influence * (social_target - old)
        if cfg.misinformation_strength > 0:
            push = 0.22 * cfg.misinformation_strength * (cfg.misinformation_direction - old)
            delta = delta + np.where(self.misinformation_mask, push, 0.0)
        if cfg.noise > 0 and self._adaptive.size:
            delta[self._adaptive] += self.rng.normal(0.0, cfg.noise, size=self._adaptive.size)

        new = np.clip(old + delta, -1.0, 1.0)
        new[self.stubborn_mask] = old[self.stubborn_mask]
        return new

    def _rewire(self) -> int:
        cfg = self.config
        if cfg.rewiring_rate <= 0 or self._adaptive.size == 0:
            return 0
        rng = self.rewire_rng
        x = self.opinions
        tol = cfg.rewiring_tolerance
        movers = self._adaptive[rng.random(self._adaptive.size) < cfg.rewiring_rate]
        rng.shuffle(movers)
        changed = 0
        for i in movers:
            i = int(i)
            neighbours = list(self.graph.neighbors(i))
            if not neighbours:
                continue
            j = neighbours[int(rng.integers(len(neighbours)))]
            if abs(x[i] - x[j]) <= tol or self.graph.degree(j) <= 1:
                continue
            close = np.flatnonzero(np.abs(x - x[i]) <= tol)
            if close.size <= 1:
                continue
            blocked = set(neighbours)
            blocked.add(i)
            candidates = [int(k) for k in close if int(k) not in blocked]
            if not candidates:
                continue
            k = candidates[int(rng.integers(len(candidates)))]
            self.graph.remove_edge(i, j)
            self.graph.add_edge(i, k)
            changed += 1
        if changed:
            self.rewired_ties += changed
            self.graph_version += 1
            self._rebuild_edge_cache()
        return changed

    def step(self) -> np.ndarray:
        self.opinions = self._opinion_update()
        self._rewire()
        self.history.append(self.opinions.copy())
        self.metric_history.append(self.metrics(step=len(self.history) - 1))
        return self.opinions

    def run(self, steps: int | None = None) -> np.ndarray:
        for _ in range(self.config.steps if steps is None else steps):
            self.step()
        return self.opinions.copy()

    def stream(self, update_every: int = 5) -> Iterator[int]:
        update_every = max(1, int(update_every))
        for step in range(1, self.config.steps + 1):
            self.step()
            if step % update_every == 0 or step == self.config.steps:
                yield step

    @property
    def history_array(self) -> np.ndarray:
        return np.vstack(self.history)