"""Finite-N discrete-state cosmic-sociology game.

This is a deliberately small computational baseline.  It solves for symmetric
finite-horizon Markov equilibria in a finite-N-corrected anonymous game by alternating a representative
civilization's dynamic best response with the population law generated by the
resulting policy.  The finite-N correction is explicit: a player meets N-1
possible opponents rather than an infinite mean field.

The equilibrium is an epsilon-Nash equilibrium because a small logit
regularization is used to make the fixed-point computation stable.  The script
reports the largest one-step gain from replacing the computed mixed action by
the best pure action as an interpretable epsilon diagnostic.
"""

from __future__ import annotations

from dataclasses import dataclass, replace
import argparse
import csv
from pathlib import Path

import numpy as np


ACTIONS = ("hide", "research", "preempt")
HIDE, RESEARCH, PREEMPT = range(3)


@dataclass(frozen=True)
class Parameters:
    tech_levels: int = 4
    horizon: float = 20.0
    dt: float = 0.25
    discount: float = 0.015
    death_loss: float = 80.0
    natural_hazard: float = 0.001
    encounter_rate: float = 0.010
    background_suspicion: float = 0.025
    suspicion_feedback: float = 0.72
    retaliation_rate: float = 0.002
    first_strike_base: float = 0.78
    research_breakthrough: float = 0.22
    spontaneous_breakthrough: float = 0.025
    explosion_probability: float = 0.22
    hide_success: float = 0.82
    research_reveal: float = 0.42
    preempt_reveal: float = 0.72
    hide_cost: float = 0.025
    research_cost: float = 0.075
    preempt_cost: float = 0.20
    tech_flow_value: float = 0.018
    terminal_tech_value: float = 0.22
    logit_temperature: float = 0.018
    damping: float = 0.16
    max_iterations: int = 2500
    tolerance: float = 2e-7

    @property
    def steps(self) -> int:
        return round(self.horizon / self.dt)


class CosmicGame:
    def __init__(self, n_civilizations: int, params: Parameters):
        if n_civilizations < 2:
            raise ValueError("At least two civilizations are required")
        self.n = n_civilizations
        self.p = params
        self.states = [(k, v) for k in range(params.tech_levels) for v in (0, 1)]
        self.s = len(self.states)
        self.a = len(ACTIONS)
        self.index = {state: j for j, state in enumerate(self.states)}
        self.transitions = self._transition_matrices()

    def _transition_matrices(self) -> np.ndarray:
        """P[action, state, next_state], conditional on surviving the period."""
        p = self.p
        out = np.zeros((self.a, self.s, self.s))
        for action in range(self.a):
            for si, (k, visible) in enumerate(self.states):
                # Visibility transition.
                if action == HIDE:
                    visible_prob = (1 - p.hide_success) if visible else 0.04
                elif action == RESEARCH:
                    visible_prob = p.research_reveal + (1 - p.research_reveal) * visible
                else:
                    visible_prob = p.preempt_reveal + (1 - p.preempt_reveal) * visible

                # Technology jumps: research raises both ordinary and explosive jumps.
                ordinary = p.spontaneous_breakthrough
                explosion = 0.0
                if action == RESEARCH:
                    ordinary += p.research_breakthrough
                    explosion = p.explosion_probability * ordinary
                    ordinary -= explosion

                tech_probs = {k: 1 - ordinary - explosion}
                tech_probs[min(k + 1, p.tech_levels - 1)] = (
                    tech_probs.get(min(k + 1, p.tech_levels - 1), 0) + ordinary
                )
                tech_probs[min(k + 2, p.tech_levels - 1)] = (
                    tech_probs.get(min(k + 2, p.tech_levels - 1), 0) + explosion
                )
                for nk, pk in tech_probs.items():
                    out[action, si, self.index[(nk, 0)]] += pk * (1 - visible_prob)
                    out[action, si, self.index[(nk, 1)]] += pk * visible_prob
        return out

    def _opponent_statistics(self, distribution: np.ndarray, policy: np.ndarray) -> tuple[float, float]:
        mass = distribution.sum()
        if mass <= 1e-15:
            return 0.0, 0.0
        normalized = distribution / mass
        attack_share = float(np.sum(normalized * policy[:, PREEMPT]))
        visible_share = float(sum(normalized[i] * (0.18 + 0.82 * v) for i, (_, v) in enumerate(self.states)))
        return attack_share, visible_share

    def _hazard(self, state: tuple[int, int], action: int, attack_share: float, visible_share: float) -> float:
        p = self.p
        k, visible = state
        neighbor_pressure = (self.n - 1) * p.encounter_rate * visible_share
        own_signature = 0.22 + 0.78 * visible
        suspicion = min(1.0, p.background_suspicion + p.suspicion_feedback * attack_share)
        # More advanced targets are harder to kill; unknown future breakthroughs
        # keep even currently weaker opponents strategically relevant.
        defense = 1.0 / (1.0 + 0.30 * k)
        waiting_hazard = neighbor_pressure * own_signature * suspicion * defense
        if action != PREEMPT:
            return p.natural_hazard + waiting_hazard

        first_strike = min(0.94, p.first_strike_base + 0.045 * k)
        retaliation = neighbor_pressure * p.retaliation_rate * (1.0 - 0.10 * k)
        return p.natural_hazard + waiting_hazard * (1 - first_strike) + retaliation

    def _flow_reward(self, state: tuple[int, int], action: int) -> float:
        k, _ = state
        cost = (self.p.hide_cost, self.p.research_cost, self.p.preempt_cost)[action]
        return self.p.tech_flow_value * k - cost

    @staticmethod
    def _softmax(q: np.ndarray, temperature: float) -> np.ndarray:
        if temperature <= 0:
            out = np.zeros_like(q)
            out[np.arange(q.shape[0]), np.argmax(q, axis=1)] = 1.0
            return out
        z = (q - q.max(axis=1, keepdims=True)) / temperature
        z = np.clip(z, -700, 0)
        e = np.exp(z)
        return e / e.sum(axis=1, keepdims=True)

    def backward_best_response(self, distribution: np.ndarray, policy_guess: np.ndarray):
        p = self.p
        value = np.zeros((p.steps + 1, self.s))
        q_values = np.zeros((p.steps, self.s, self.a))
        policy = np.zeros_like(q_values)
        value[-1] = np.array([p.terminal_tech_value * k for k, _ in self.states])
        discount = np.exp(-p.discount * p.dt)

        for t in range(p.steps - 1, -1, -1):
            attack_share, visible_share = self._opponent_statistics(distribution[t], policy_guess[t])
            for si, state in enumerate(self.states):
                for action in range(self.a):
                    hazard = self._hazard(state, action, attack_share, visible_share)
                    survive = np.exp(-hazard * p.dt)
                    continuation = self.transitions[action, si] @ value[t + 1]
                    q_values[t, si, action] = (
                        p.dt * self._flow_reward(state, action)
                        + discount * (survive * continuation + (1 - survive) * (-p.death_loss))
                    )
            policy[t] = self._softmax(q_values[t], p.logit_temperature)
            # Inclusive value is the expected value under the regularized response.
            value[t] = np.sum(policy[t] * q_values[t], axis=1)
        return policy, value, q_values

    def forward(self, initial: np.ndarray, policy: np.ndarray):
        p = self.p
        distribution = np.zeros((p.steps + 1, self.s))
        distribution[0] = initial
        for t in range(p.steps):
            attack_share, visible_share = self._opponent_statistics(distribution[t], policy[t])
            nxt = np.zeros(self.s)
            for si, state in enumerate(self.states):
                for action in range(self.a):
                    weight = distribution[t, si] * policy[t, si, action]
                    hazard = self._hazard(state, action, attack_share, visible_share)
                    survive = np.exp(-hazard * p.dt)
                    nxt += weight * survive * self.transitions[action, si]
            distribution[t + 1] = nxt
        return distribution

    def solve(self, initial_mode: str = "peace") -> dict:
        p = self.p
        initial = np.zeros(self.s)
        initial[self.index[(0, 0)]] = 0.85
        initial[self.index[(0, 1)]] = 0.15

        policy = np.full((p.steps, self.s, self.a), 0.0)
        if initial_mode == "peace":
            policy[:, :, HIDE] = 0.58
            policy[:, :, RESEARCH] = 0.41
            policy[:, :, PREEMPT] = 0.01
        elif initial_mode == "aggressive":
            policy[:, :, HIDE] = 0.03
            policy[:, :, RESEARCH] = 0.07
            policy[:, :, PREEMPT] = 0.90
        else:
            raise ValueError("initial_mode must be 'peace' or 'aggressive'")
        distribution = self.forward(initial, policy)

        residual = np.inf
        for iteration in range(1, p.max_iterations + 1):
            best, value, q_values = self.backward_best_response(distribution, policy)
            new_policy = (1 - p.damping) * policy + p.damping * best
            generated = self.forward(initial, new_policy)
            new_distribution = (1 - p.damping) * distribution + p.damping * generated
            residual = max(float(np.max(np.abs(new_policy - policy))), float(np.max(np.abs(new_distribution - distribution))))
            policy, distribution = new_policy, new_distribution
            if residual < p.tolerance:
                break

        # Re-evaluate against the converged environment for diagnostics.
        response, value, q_values = self.backward_best_response(distribution, policy)
        on_path_weights = distribution[:-1]
        action_value = np.sum(policy * q_values, axis=2)
        best_value = np.max(q_values, axis=2)
        denom = max(float(on_path_weights.sum()), 1e-15)
        mean_deviation_gain = float(np.sum(on_path_weights * (best_value - action_value)) / denom)
        max_deviation_gain = float(np.max(best_value - action_value))
        attack_path = np.array([self._opponent_statistics(distribution[t], policy[t])[0] for t in range(p.steps)])
        expected_tech = []
        for t in range(p.steps + 1):
            mass = distribution[t].sum()
            expected_tech.append(sum(distribution[t, i] * k for i, (k, _) in enumerate(self.states)) / max(mass, 1e-15))

        return {
            "N": self.n,
            "initialization": initial_mode,
            "converged": residual < p.tolerance,
            "iterations": iteration,
            "residual": residual,
            "policy": policy,
            "distribution": distribution,
            "value": value,
            "attack_path": attack_path,
            "initial_attack": float(attack_path[0]),
            "mean_attack": float(attack_path.mean()),
            "final_survival": float(distribution[-1].sum()),
            "expected_survivors": float(self.n * distribution[-1].sum()),
            "final_tech": float(expected_tech[-1]),
            "mean_deviation_gain": mean_deviation_gain,
            "max_deviation_gain": max_deviation_gain,
        }


def plot_scan(results: list[dict], output: Path) -> None:
    # Import lazily so the solver itself only needs NumPy.
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    fig, axes = plt.subplots(1, 2, figsize=(9.2, 3.8), constrained_layout=True)
    styles = {"peace": ("o-", "peace initialization"), "aggressive": ("s--", "aggressive initialization")}
    for mode, (style, label) in styles.items():
        subset = [r for r in results if r["initialization"] == mode]
        ns = [r["N"] for r in subset]
        axes[0].plot(ns, [r["mean_attack"] for r in subset], style, label=label)
        axes[1].plot(ns, [r["final_survival"] for r in subset], style, label=label)
    axes[0].set(ylabel="mean preemption probability", xlabel="civilizations N", ylim=(-0.03, 1.03))
    axes[1].set(ylabel="individual survival probability at T", xlabel="civilizations N", ylim=(-0.03, 1.03))
    for ax in axes:
        ax.grid(alpha=0.25)
    axes[0].legend(frameon=False, fontsize=8)
    fig.suptitle("Symmetric regularized equilibrium versus population size")
    output.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(output, dpi=180)
    plt.close(fig)


def run_scan(ns: list[int], params: Parameters, output: Path | None = None, plot: Path | None = None) -> list[dict]:
    results = []
    for n in ns:
        for mode in ("peace", "aggressive"):
            results.append(CosmicGame(n, params).solve(mode))

    fields = [
        "N", "initialization", "converged", "iterations", "residual",
        "initial_attack", "mean_attack", "final_survival", "expected_survivors",
        "final_tech", "mean_deviation_gain", "max_deviation_gain",
    ]
    if output is not None:
        output.parent.mkdir(parents=True, exist_ok=True)
        with output.open("w", newline="") as f:
            writer = csv.DictWriter(f, fieldnames=fields)
            writer.writeheader()
            for result in results:
                writer.writerow({key: result[key] for key in fields})
    if plot is not None:
        plot_scan(results, plot)

    print(" N  init        conv  attack(t0) attack(avg) survival(T) E[survivors] tech(T)  eps(mean)")
    for r in results:
        print(
            f"{r['N']:3d}  {r['initialization']:<10} {str(r['converged']):<5} "
            f"{r['initial_attack']:10.3f} {r['mean_attack']:11.3f} "
            f"{r['final_survival']:11.3f} {r['expected_survivors']:12.3f} "
            f"{r['final_tech']:7.3f} {r['mean_deviation_gain']:10.4g}"
        )
    return results


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--n", nargs="+", type=int, default=[2, 5, 10, 20, 30, 40, 50, 75, 100])
    parser.add_argument("--output", type=Path, default=Path("results/equilibria.csv"))
    parser.add_argument("--plot", type=Path, default=Path("results/equilibria.png"))
    parser.add_argument("--temperature", type=float, default=None)
    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()
    parameters = Parameters()
    if args.temperature is not None:
        parameters = replace(parameters, logit_temperature=args.temperature)
    run_scan(args.n, parameters, args.output, args.plot)