"""Swarm experiment runner — agent-oriented BFF with identity and provenance.

Parallels bff/experiment.py but operates on named agents instead of anonymous
tapes. Uses calibrated parameters from Track 1 (mutation_rate=1e-3) and
Track 2 (window_size=3, critical_threshold=-0.001).
"""

import json
import os
import time
from pathlib import Path

import numpy as np

from .metrics import compression_ratio
from .swarm import Agent, KTermFusion, DiversityProbe, Provenance
from ._util import serialize_rng_state, deserialize_rng_state


DEFAULT_CONFIG = {
    "n_agents": 1024,
    "tape_len": 64,
    "max_steps_per_interaction": 10_000,
    "mutation_rate": 1e-3,          # calibrated ecology zone
    "snapshot_interval": 10_000,
    "checkpoint_interval": 100_000,
    "seed": 42,
    "output_dir": "data/bff_swarm/",
    "probe_window": 3,              # calibrated
    "probe_threshold": -0.001,      # calibrated
}


class SwarmExperiment:
    """Runs the BFF swarm experiment with agent identity tracking."""

    def __init__(self, config: dict | None = None):
        self.config = {**DEFAULT_CONFIG, **(config or {})}
        self.rng = np.random.default_rng(self.config["seed"])

        # Initialize agents from random tapes
        tape_len = self.config["tape_len"]
        n_agents = self.config["n_agents"]
        self.agents: list[Agent] = []
        for i in range(n_agents):
            tape = self.rng.integers(0, 256, size=tape_len, dtype=np.uint8)
            agent = Agent(tape, name=f"agent-{i:04d}")
            self.agents.append(agent)

        # Fusion engine
        self.fusion = KTermFusion(
            max_steps=self.config["max_steps_per_interaction"]
        )

        # Diversity probe with calibrated params
        self.probe = DiversityProbe(
            window_size=self.config["probe_window"],
            critical_threshold=self.config["probe_threshold"],
        )

        # Tracking state
        self.interaction_count: int = 0
        self.fusion_log: list[dict] = []

        # Snapshot metrics
        self.snapshot_interactions: list[int] = []
        self.snapshot_compression: list[float] = []
        self.snapshot_probe_readings: list[dict] = []

        # Output dir
        self.output_dir = Path(self.config["output_dir"])

    def run(self, n_interactions: int = 1_000_000):
        """Main loop: pick 2 agents, fuse, apply mutation, track provenance."""
        self.output_dir.mkdir(parents=True, exist_ok=True)

        snapshot_interval = self.config["snapshot_interval"]
        checkpoint_interval = self.config["checkpoint_interval"]
        mutation_rate = self.config["mutation_rate"]
        n_agents = len(self.agents)
        target = self.interaction_count + n_interactions

        t_start = time.time()

        print(f"Starting BFF swarm: {n_interactions:,} interactions")
        print(f"  Agents: {n_agents}")
        print(f"  Tape length: {self.config['tape_len']}")
        print(f"  Mutation rate: {mutation_rate}")
        print(f"  Seed: {self.config['seed']}")
        print()

        while self.interaction_count < target:
            # 1. Pick two agents uniformly at random
            idx_a, idx_b = self.rng.choice(n_agents, size=2, replace=False)
            agent_a = self.agents[idx_a]
            agent_b = self.agents[idx_b]

            # 2. Fuse (concatenate-execute-split + provenance update)
            ops, _ = self.fusion.fuse(agent_a, agent_b, self.interaction_count)

            # 3. Log fusion event
            self.fusion_log.append({
                "interaction": self.interaction_count,
                "agent_a": agent_a.id,
                "agent_b": agent_b.id,
                "ops": ops,
            })

            # 4. Apply mutation
            if mutation_rate > 0 and self.rng.random() < mutation_rate:
                target_agent = self.agents[self.rng.choice([idx_a, idx_b])]
                byte_idx = self.rng.integers(0, self.config["tape_len"])
                bit_idx = self.rng.integers(0, 8)
                target_agent.tape[byte_idx] ^= np.uint8(1 << bit_idx)

            self.interaction_count += 1

            # 5. Snapshot metrics
            if self.interaction_count % snapshot_interval == 0:
                soup = np.stack([a.tape for a in self.agents])
                cr = compression_ratio(soup)
                reading = self.probe.measure(self.agents, self.interaction_count)

                self.snapshot_interactions.append(self.interaction_count)
                self.snapshot_compression.append(cr)
                self.snapshot_probe_readings.append(reading.to_dict())

            # Progress report every 100k
            if self.interaction_count % 100_000 == 0:
                elapsed = time.time() - t_start
                rate = self.interaction_count / elapsed if elapsed > 0 else 0
                remaining = target - self.interaction_count
                eta = remaining / rate if rate > 0 else float("inf")

                cr = self.snapshot_compression[-1] if self.snapshot_compression else -1
                alert = "none"
                if self.snapshot_probe_readings:
                    alert = self.snapshot_probe_readings[-1].get("alert_level", "none")

                line = (
                    f"[{self.interaction_count:>10,}] "
                    f"elapsed={elapsed:>7.1f}s  "
                    f"rate={rate:>8,.0f}/s  "
                    f"compression={cr:.4f}"
                )
                if alert != "none":
                    line += f"  ALERT={alert.upper()}"
                line += f"  ETA={eta:>6.0f}s"
                print(line)

            # 6. Checkpoint
            if self.interaction_count % checkpoint_interval == 0:
                self.save_checkpoint(
                    str(self.output_dir / f"swarm_checkpoint_{self.interaction_count}.npz")
                )

        # Final snapshot + checkpoint
        soup = np.stack([a.tape for a in self.agents])
        cr = compression_ratio(soup)
        reading = self.probe.measure(self.agents, self.interaction_count)
        self.snapshot_interactions.append(self.interaction_count)
        self.snapshot_compression.append(cr)
        self.snapshot_probe_readings.append(reading.to_dict())

        final_path = str(self.output_dir / f"swarm_final_{self.interaction_count}.npz")
        self.save_checkpoint(final_path)
        print(f"\nSwarm experiment complete. {self.interaction_count:,} interactions.")
        print(f"Final checkpoint: {final_path}")

    def save_checkpoint(self, path: str):
        """Save agent tapes, provenance graph, probe history."""
        # Stack all agent tapes
        tapes = np.stack([a.tape for a in self.agents])

        np.savez_compressed(path, tapes=tapes)

        # Save metadata
        meta_path = path.replace(".npz", "_meta.json")
        agent_meta = []
        for a in self.agents:
            agent_meta.append({
                "id": a.id,
                "name": a.name,
                "interaction_count": a.interaction_count,
                "provenance": {
                    "agent_id": a.provenance.agent_id,
                    "parent_ids": a.provenance.parent_ids,
                    "merge_interaction": a.provenance.merge_interaction,
                    "depth": a.provenance.depth,
                    "created_at": a.provenance.created_at,
                },
            })

        meta = {
            "config": self.config,
            "interaction_count": self.interaction_count,
            "agents": agent_meta,
            "snapshot_interactions": self.snapshot_interactions,
            "snapshot_compression": self.snapshot_compression,
            "snapshot_probe_readings": self.snapshot_probe_readings,
            "rng_state": serialize_rng_state(self.rng.bit_generator.state),
        }
        with open(meta_path, "w") as f:
            json.dump(meta, f, indent=2)

    @classmethod
    def load_checkpoint(cls, path: str) -> "SwarmExperiment":
        """Restore swarm from checkpoint files."""
        data = np.load(path)
        meta_path = path.replace(".npz", "_meta.json")
        with open(meta_path) as f:
            meta = json.load(f)

        exp = cls(config=meta["config"])
        exp.interaction_count = meta["interaction_count"]
        exp.snapshot_interactions = meta["snapshot_interactions"]
        exp.snapshot_compression = meta["snapshot_compression"]
        exp.snapshot_probe_readings = meta.get("snapshot_probe_readings", [])

        # Restore agent states
        tapes = data["tapes"]
        for i, agent_data in enumerate(meta["agents"]):
            agent = exp.agents[i]
            agent.id = agent_data["id"]
            agent.name = agent_data["name"]
            agent.interaction_count = agent_data["interaction_count"]
            agent.tape = tapes[i].copy()
            prov = agent_data["provenance"]
            agent.provenance = Provenance(
                agent_id=prov["agent_id"],
                parent_ids=prov["parent_ids"],
                merge_interaction=prov["merge_interaction"],
                depth=prov["depth"],
                created_at=prov["created_at"],
            )

        # Restore RNG
        if "rng_state" in meta:
            exp.rng.bit_generator.state = deserialize_rng_state(meta["rng_state"])

        return exp

    def get_fusion_graph(self) -> list[dict]:
        """Return recent fusion events for analysis."""
        return list(self.fusion_log)
