"""Main experiment runner with checkpointing and progress reporting."""

import json
import os
import time
from pathlib import Path

import numpy as np

from .metrics import compression_ratio, find_replicators, ReplicatorType
from .probe import CrystallizationProbe
from .soup import Soup
from ._util import serialize_rng_state, deserialize_rng_state


DEFAULT_CONFIG = {
    "n_tapes": 1024,
    "tape_len": 64,
    "max_steps_per_interaction": 10_000,
    "snapshot_interval": 10_000,
    "checkpoint_interval": 100_000,
    "seed": None,
    "output_dir": "data/bff_runs/",
    "mutation_rate": 0.0,
    "max_tree_depth": 0,
    "enable_probe": False,
}


class Experiment:
    """Runs the BFF abiogenesis experiment."""

    def __init__(self, config: dict | None = None):
        self.config = {**DEFAULT_CONFIG, **(config or {})}
        self.soup = Soup(
            n_tapes=self.config["n_tapes"],
            tape_len=self.config["tape_len"],
            max_steps=self.config["max_steps_per_interaction"],
            seed=self.config["seed"],
            mutation_rate=self.config.get("mutation_rate", 0.0),
            max_tree_depth=self.config.get("max_tree_depth", 0),
        )
        self.interaction_count = 0

        # Pre-allocate metric storage (grown on demand if needed)
        self._ops_buffer_size = 1_000_000
        self.ops_log = np.zeros(self._ops_buffer_size, dtype=np.uint16)

        # Snapshot metrics (recorded every snapshot_interval)
        self.snapshot_interactions: list[int] = []
        self.snapshot_compression: list[float] = []
        self.snapshot_mean_ops: list[float] = []
        self.snapshot_replicator_counts: list[dict[str, int]] = []
        self.snapshot_mutation_count: list[int] = []
        self.snapshot_max_depth: list[int] = []
        self.snapshot_mean_depth: list[float] = []
        self.snapshot_probe_readings: list[dict] = []

        # Cumulative mutation counter for current window
        self._window_mutations = 0

        # Crystallization probe (opt-in)
        self.probe: CrystallizationProbe | None = None
        if self.config.get("enable_probe", False):
            self.probe = CrystallizationProbe()

        # Ensure output dir
        self.output_dir = Path(self.config["output_dir"])
        self.output_dir.mkdir(parents=True, exist_ok=True)

    def _grow_ops_log(self, needed: int):
        """Double the ops_log buffer if needed."""
        if needed >= len(self.ops_log):
            new_size = max(len(self.ops_log) * 2, needed + 1)
            new_log = np.zeros(new_size, dtype=np.uint16)
            new_log[:len(self.ops_log)] = self.ops_log
            self.ops_log = new_log

    def run(self, n_interactions: int = 10_000_000):
        """Run the experiment for n_interactions."""
        self._grow_ops_log(self.interaction_count + n_interactions)

        snapshot_interval = self.config["snapshot_interval"]
        checkpoint_interval = self.config["checkpoint_interval"]
        target = self.interaction_count + n_interactions

        t_start = time.time()
        t_last_report = t_start
        ops_window: list[int] = []

        print(f"Starting BFF experiment: {n_interactions:,} interactions")
        print(f"  Tapes: {self.config['n_tapes']} × {self.config['tape_len']} bytes")
        print(f"  Max steps/interaction: {self.config['max_steps_per_interaction']:,}")
        print(f"  Seed: {self.config['seed']}")
        if self.config.get("mutation_rate", 0.0) > 0:
            print(f"  Mutation rate: {self.config['mutation_rate']}")
        if self.config.get("max_tree_depth", 0) > 0:
            print(f"  Max tree depth: {self.config['max_tree_depth']}")
        if self.probe:
            print(f"  Probe: enabled")
        print()

        while self.interaction_count < target:
            result = self.soup.interact()
            self.ops_log[self.interaction_count] = min(result.ops_count, 65535)
            ops_window.append(result.ops_count)
            if result.mutated:
                self._window_mutations += 1
            self.interaction_count += 1

            # Snapshot metrics
            if self.interaction_count % snapshot_interval == 0:
                soup_state = self.soup.get_state()
                cr = compression_ratio(soup_state)
                mean_ops = float(np.mean(ops_window)) if ops_window else 0.0

                self.snapshot_interactions.append(self.interaction_count)
                self.snapshot_compression.append(cr)
                self.snapshot_mean_ops.append(mean_ops)
                self.snapshot_mutation_count.append(self._window_mutations)

                # Depth metrics
                depths = self.soup.get_depths()
                self.snapshot_max_depth.append(int(depths.max()))
                self.snapshot_mean_depth.append(float(depths.mean()))

                # Replicator detection (skip for very early snapshots to save time)
                if self.interaction_count >= snapshot_interval * 10:
                    reps = find_replicators(soup_state, min_length=4, min_count=3)
                    counts = {t.value: 0 for t in ReplicatorType}
                    for r in reps:
                        counts[r.classification.value] += 1
                    self.snapshot_replicator_counts.append(counts)
                else:
                    self.snapshot_replicator_counts.append(
                        {t.value: 0 for t in ReplicatorType}
                    )

                # Probe reading
                if self.probe:
                    reading = self.probe.measure(soup_state, self.interaction_count)
                    self.snapshot_probe_readings.append(reading.to_dict())

                ops_window.clear()
                self._window_mutations = 0

            # Progress report every 100k
            if self.interaction_count % 100_000 == 0:
                now = time.time()
                elapsed = now - 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
                mo = self.snapshot_mean_ops[-1] if self.snapshot_mean_ops else -1

                line = (
                    f"[{self.interaction_count:>10,}] "
                    f"elapsed={elapsed:>7.1f}s  "
                    f"rate={rate:>8,.0f}/s  "
                    f"compression={cr:.4f}  "
                    f"mean_ops={mo:>7.1f}  "
                    f"ETA={eta:>6.0f}s"
                )

                # Append mutation count if active
                if self.config.get("mutation_rate", 0.0) > 0:
                    total_mutations = sum(self.snapshot_mutation_count)
                    line += f"  mutations={total_mutations}"

                # Append depth info if tracking
                if self.snapshot_max_depth:
                    line += f"  max_depth={self.snapshot_max_depth[-1]}"

                # Append probe alert if active
                if self.probe and self.snapshot_probe_readings:
                    alert = self.snapshot_probe_readings[-1].get("alert_level", "none")
                    if alert != "none":
                        line += f"  ALERT={alert.upper()}"

                print(line)
                t_last_report = now

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

        # Final snapshot + checkpoint
        self._take_final_snapshot()
        final_path = str(self.output_dir / f"final_{self.interaction_count}.npz")
        self.save_checkpoint(final_path)
        print(f"\nExperiment complete. {self.interaction_count:,} interactions.")
        print(f"Final checkpoint: {final_path}")

    def _take_final_snapshot(self):
        """Record a final metrics snapshot."""
        soup_state = self.soup.get_state()
        cr = compression_ratio(soup_state)
        self.snapshot_interactions.append(self.interaction_count)
        self.snapshot_compression.append(cr)
        self.snapshot_mean_ops.append(0.0)
        self.snapshot_mutation_count.append(self._window_mutations)

        depths = self.soup.get_depths()
        self.snapshot_max_depth.append(int(depths.max()))
        self.snapshot_mean_depth.append(float(depths.mean()))

        reps = find_replicators(soup_state, min_length=4, min_count=3)
        counts = {t.value: 0 for t in ReplicatorType}
        for r in reps:
            counts[r.classification.value] += 1
        self.snapshot_replicator_counts.append(counts)

        if self.probe:
            reading = self.probe.measure(soup_state, self.interaction_count)
            self.snapshot_probe_readings.append(reading.to_dict())

    def save_checkpoint(self, path: str):
        """Save full experiment state."""
        np.savez_compressed(
            path,
            tapes=self.soup.get_state(),
            ops_log=self.ops_log[:self.interaction_count],
            depths=self.soup.get_depths(),
        )

        meta_path = path.replace(".npz", "_meta.json")
        meta = {
            "config": self.config,
            "interaction_count": self.interaction_count,
            "rng_state": serialize_rng_state(self.soup.get_rng_state()),
            "snapshot_interactions": self.snapshot_interactions,
            "snapshot_compression": self.snapshot_compression,
            "snapshot_mean_ops": self.snapshot_mean_ops,
            "snapshot_replicator_counts": self.snapshot_replicator_counts,
            "snapshot_mutation_count": self.snapshot_mutation_count,
            "snapshot_max_depth": self.snapshot_max_depth,
            "snapshot_mean_depth": self.snapshot_mean_depth,
            "snapshot_probe_readings": self.snapshot_probe_readings,
        }
        with open(meta_path, "w") as f:
            json.dump(meta, f, indent=2)

    @classmethod
    def load_checkpoint(cls, path: str) -> "Experiment":
        """Restore experiment 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_mean_ops = meta["snapshot_mean_ops"]
        exp.snapshot_replicator_counts = meta["snapshot_replicator_counts"]

        # Backward-compatible: new fields default to empty
        exp.snapshot_mutation_count = meta.get("snapshot_mutation_count", [])
        exp.snapshot_max_depth = meta.get("snapshot_max_depth", [])
        exp.snapshot_mean_depth = meta.get("snapshot_mean_depth", [])
        exp.snapshot_probe_readings = meta.get("snapshot_probe_readings", [])

        # Restore ops log
        saved_ops = data["ops_log"]
        exp._grow_ops_log(len(saved_ops))
        exp.ops_log[:len(saved_ops)] = saved_ops

        # Restore soup state + RNG + depths
        rng_state = deserialize_rng_state(meta["rng_state"])
        depths = data["depths"] if "depths" in data else None
        exp.soup.set_state(data["tapes"], rng_state, depths=depths)

        return exp
