"""
AfsEngine — Top-level simulation orchestrator.

Takes a compiled FieldDefinition and runs cells through N-dimensional space.
Cells gain energy from concept zones, lose energy from metabolism,
reproduce when hot. Forces from the similarity matrix create trade-off
dynamics that drive differentiation.

NumPy-vectorized force calculations — no per-cell loops for physics.
"""

import numpy as np
from typing import Optional, Callable
from .field_definition import FieldDefinition
from .config import AfsSimulationConfig
from .spatial_index import SpatialIndex
from ..entities.cell import AfsCell
from ..entities.concept import Concept


class AfsEngine:
    def __init__(self, field: FieldDefinition, config: Optional[AfsSimulationConfig] = None,
                 use_gpu: Optional[bool] = None):
        if not isinstance(field, FieldDefinition):
            raise TypeError("AfsEngine requires a FieldDefinition instance")

        self.field = field
        self.config = config or AfsSimulationConfig()
        self.n_dims = field.get_dimension_count()
        self.spatial_index = SpatialIndex(self.n_dims)

        self.cells: dict[str, AfsCell] = {}
        self.concepts: dict[str, Concept] = {}

        self.tick_count = 0
        self.generation = 0
        self.total_ticks = 0  # Set by caller for gravity ramping + settling
        self.stats = {
            "total_cells_created": 0,
            "total_reproductions": 0,
            "total_merges": 0,
            "peak_population": 0,
        }

        # Event callbacks: {event_name: [callback, ...]}
        self._listeners: dict[str, list[Callable]] = {}

        # GPU acceleration (auto-detect if use_gpu is None)
        self._gpu_kernel = None
        self._gpu_matrix_uploaded = False
        if use_gpu is not False:
            self._init_gpu(use_gpu)

    def _init_gpu(self, force: Optional[bool] = None):
        """Initialize GPU force kernel if available."""
        try:
            from ..gpu import WgpuForceKernel, gpu_available
            if force or (force is None and gpu_available()):
                max_cells = getattr(self.config.engine, 'max_population', 512) + 64
                self._gpu_kernel = WgpuForceKernel(
                    num_dims=self.n_dims,
                    max_cells=max_cells,
                    max_concepts=32,
                )
                info = self._gpu_kernel.adapter_info
                print(f"[AfsEngine] GPU acceleration: {info.get('device', 'unknown')} via {info.get('backend_type', 'unknown')}")
        except Exception as e:
            if force:
                raise
            print(f"[AfsEngine] GPU not available, using NumPy: {e}")

    def on(self, event: str, callback: Callable):
        self._listeners.setdefault(event, []).append(callback)

    def emit(self, event: str, data=None):
        for cb in self._listeners.get(event, []):
            cb(data)

    # ─── Initialization ───────────────────────────────────────

    def initialize(self, seed_centroid: Optional[np.ndarray] = None,
                   seed_spread: Optional[float] = None):
        """Create initial cells and concepts from the field definition.

        Args:
            seed_centroid: If provided, spawn cells here instead of the input
                region centroid. Used by child colony drill-down (P9). Normalized
                0-1 coordinates in N-dimensional space.
            seed_spread: Spawn radius around seed_centroid. Defaults to
                config.engine.spawn_spread. Pass a smaller value (e.g., 0.06)
                to start a tightly focused child colony.
        """
        n = self.n_dims
        active_dims = self.field.get_active_dimensions()

        # Check for world model attractors in field metadata
        wm = self.field.metadata.get("worldModel", {})
        attractors = wm.get("attractors", [])

        if attractors:
            # Create ONE CONCEPT PER ATTRACTOR (matches JS engine).
            # Each attractor has a full N-dimensional position representing
            # a competing regime in the solution space.
            dim_names = [d.name for d in active_dims]
            for att in attractors:
                pos = self._normalize_attractor_position(att, active_dims, dim_names)
                concept = Concept(att.get("name", f"attractor_{len(self.concepts)}"), n, position=pos)
                concept.energy = att.get("strength", 0.5)
                self.concepts[concept.id] = concept

            # Spawn cells at the centroid of all attractors (the input region),
            # unless a seed_centroid override is given (child colony drill-down).
            if seed_centroid is not None:
                center = np.clip(seed_centroid, 0.05, 0.95)
            elif self.concepts:
                concept_positions = np.array([c.position for c in self.concepts.values()])
                center = concept_positions.mean(axis=0)
            else:
                center = np.full(n, 0.5)
        else:
            # No attractors — create concepts from dimensions (legacy mode)
            for i, dim in enumerate(active_dims):
                pos = np.full(n, 0.5)
                pos[i] = 0.3 + np.random.uniform(0, 0.4)
                concept = Concept(dim.name, n, position=pos)
                self.concepts[concept.id] = concept
            center = seed_centroid if seed_centroid is not None else np.full(n, 0.5)

        # Spawn cells at chosen center
        spread = seed_spread if seed_spread is not None else self.config.engine.spawn_spread
        for _ in range(self.config.engine.initial_cell_count):
            pos = center + np.random.uniform(-spread, spread, n)
            pos = np.clip(pos, 0.05, 0.95)
            cell = AfsCell(n, position=pos, energy=self.config.cell.initial_energy)
            self.cells[cell.id] = cell
            self.stats["total_cells_created"] += 1

        self.stats["peak_population"] = len(self.cells)

    def _normalize_attractor_position(self, attractor, active_dims, dim_names):
        """Convert attractor real-unit position to normalized 0-1 space."""
        n = len(active_dims)
        pos = np.full(n, 0.5)
        att_pos = attractor.get("position", {})

        for i, dim in enumerate(active_dims):
            name = dim.name
            if name in att_pos:
                val = att_pos[name]
                # Normalize using display_range (validRange)
                dr = dim.display_range
                if dr and len(dr) == 2 and dr[1] != dr[0]:
                    pos[i] = np.clip((val - dr[0]) / (dr[1] - dr[0]), 0.0, 1.0)
                else:
                    pos[i] = 0.5
        return pos

    # ─── Force Computation Paths ─────────────────────────────

    def _compute_gravity_scale(self) -> float:
        """Gravity ramping: 5% → 100% over generation budget (smoothstep).

        Low gravity early lets cells explore freely. Increasing gravity
        pulls them into attractor basins. Matches JS engine behavior.
        """
        if self.total_ticks <= 0:
            return 1.0  # no budget set — full gravity
        progress = min(1.0, self.tick_count / self.total_ticks)
        # Map [0.1, 0.9] → [0, 1], clamped
        t = max(0.0, min(1.0, (progress - 0.1) / 0.8))
        # Smoothstep: starts slow, accelerates, eases into full gravity
        return 0.05 + 0.95 * t * t * (3 - 2 * t)

    def _compute_settling_damping(self) -> float:
        """Settling: progressively slow cells in final 15% of generations.

        Returns momentum multiplier (1.0 = no damping, lower = more damping).
        Produces cleaner cluster convergence at end of run.
        """
        if self.total_ticks <= 0:
            return 1.0
        progress = min(1.0, self.tick_count / self.total_ticks)
        settling_start = 0.85
        if progress <= settling_start:
            return 1.0
        # Linear ramp from 1.0 down to 0.3 over final 15%
        settling = (progress - settling_start) / (1.0 - settling_start)
        return 1.0 - 0.7 * settling

    def _tick_gpu(self, cell_positions, cell_velocities, concept_positions, dt):
        """GPU path: single tick — forces + position update on GPU."""
        cfg = self.config
        kernel = self._gpu_kernel

        # Upload similarity matrix on first tick or when it changes
        if not self._gpu_matrix_uploaded and self.field.similarity_matrix is not None:
            kernel.upload_matrix(self.field.similarity_matrix)
            self._gpu_matrix_uploaded = True

        kernel.upload_concepts(concept_positions)
        kernel.upload_cells(cell_positions, cell_velocities)

        gravity_scale = self._compute_gravity_scale()
        settling = self._compute_settling_damping()
        effective_momentum = min(cfg.cell.momentum, 0.7) * settling

        kernel.tick({
            "force_scale": cfg.engine.force_scale,
            "tradeoff_amplifier": cfg.engine.tradeoff_amplifier,
            "gravity_strength": cfg.engine.concept_gravity_strength * gravity_scale,
            "gravity_sharpness": cfg.engine.concept_gravity_sharpness,
            "gravity_range": cfg.engine.concept_gravity_range,
            "softening": 0.05,
            "dt": dt,
            "momentum": effective_momentum,
            "max_speed": 0.3,
            "boundary_margin": cfg.engine.boundary_margin,
            "boundary_strength": cfg.engine.boundary_strength,
            "thermal_noise": cfg.engine.thermal_noise,
            "tick_count": self.tick_count,
        })

        return kernel.readback()

    def run_generation(self, ticks: int = 100, dt: float = 16.67):
        """Run a full generation — batches GPU ticks with one upload/readback.

        Forces run on GPU for all ticks (one upload, N dispatches, one readback).
        Metabolism and reproduction run on CPU once at the end.
        This is 50-100x faster than calling tick() N times with GPU.
        """
        if len(self.cells) == 0:
            return

        cfg = self.config
        n = self.n_dims
        cell_list = list(self.cells.values())
        concept_list = list(self.concepts.values())

        cell_positions = np.array([c.position for c in cell_list])
        cell_velocities = np.array([c.velocity for c in cell_list])
        concept_positions = np.array([c.position for c in concept_list])

        if self._gpu_kernel is not None:
            kernel = self._gpu_kernel

            if not self._gpu_matrix_uploaded and self.field.similarity_matrix is not None:
                kernel.upload_matrix(self.field.similarity_matrix)
                self._gpu_matrix_uploaded = True

            kernel.upload_concepts(concept_positions)
            kernel.upload_cells(cell_positions, cell_velocities)

            # Run N ticks on GPU — one upload, N dispatches, one readback
            for i in range(ticks):
                self.tick_count += 1
                gravity_scale = self._compute_gravity_scale()
                settling = self._compute_settling_damping()
                effective_momentum = min(cfg.cell.momentum, 0.7) * settling

                kernel.tick({
                    "force_scale": cfg.engine.force_scale,
                    "tradeoff_amplifier": cfg.engine.tradeoff_amplifier,
                    "gravity_strength": cfg.engine.concept_gravity_strength * gravity_scale,
                    "gravity_sharpness": cfg.engine.concept_gravity_sharpness,
                    "gravity_range": cfg.engine.concept_gravity_range,
                    "softening": 0.05,
                    "dt": dt,
                    "momentum": effective_momentum,
                    "max_speed": 0.3,
                    "boundary_margin": cfg.engine.boundary_margin,
                    "boundary_strength": cfg.engine.boundary_strength,
                    "thermal_noise": cfg.engine.thermal_noise,
                    "tick_count": self.tick_count,
                })

            cell_positions, cell_velocities = kernel.readback()
        else:
            # CPU path — tick normally
            for i in range(ticks):
                self.tick_count += 1
                cell_positions, cell_velocities = self._tick_cpu(
                    cell_positions, cell_velocities, concept_positions, dt)

        # Write back to cells
        for i, cell in enumerate(cell_list):
            cell.position = cell_positions[i]
            cell.velocity = cell_velocities[i]

        # Metabolism (once per generation, not per tick — close enough for colony dynamics)
        deltas = concept_positions[None, :, :] - cell_positions[:, None, :]
        distances_flat = np.linalg.norm(deltas, axis=2)
        absorption_range = cfg.metabolism.concept_absorption_range
        in_absorption = distances_flat < absorption_range
        energy_gain = np.where(
            in_absorption,
            cfg.metabolism.max_energy_gain_per_tick * ticks / (1.0 + distances_flat ** cfg.metabolism.energy_falloff_exponent),
            0.0
        )
        total_gain = np.sum(energy_gain, axis=1)

        movement_cost = np.linalg.norm(cell_velocities, axis=1) * cfg.metabolism.movement_cost_factor * ticks
        base_cost = cfg.metabolism.base_cost_per_dim * n * ticks
        total_cost = base_cost + movement_cost

        for i, cell in enumerate(cell_list):
            cell.energy += total_gain[i] - total_cost[i]
            cell.energy = max(cfg.metabolism.min_energy, min(cfg.cell.max_energy, cell.energy))
            cell.age += ticks
            cell.update_state(cfg.cell.max_energy)

        # Reproduction
        new_cells = []
        for cell in cell_list:
            if (cell.energy > cfg.cell.max_energy * cfg.cell.reproduction_threshold
                    and len(self.cells) + len(new_cells) < cfg.engine.max_population):
                child = cell.spawn_child(
                    transfer_ratio=cfg.cell.reproduction_transfer,
                    offset=cfg.cell.child_spawn_offset,
                    mutation=cfg.cell.mutation_strength,
                )
                new_cells.append(child)
                self.stats["total_reproductions"] += 1
                self.stats["total_cells_created"] += 1

        for child in new_cells:
            self.cells[child.id] = child

        # Death
        if len(self.cells) > cfg.engine.max_population:
            sorted_cells = sorted(self.cells.values(), key=lambda c: c.energy)
            excess = len(self.cells) - cfg.engine.max_population
            for cell in sorted_cells[:excess]:
                del self.cells[cell.id]

        # Generation event
        self.generation += 1
        self.stats["peak_population"] = max(self.stats["peak_population"], len(self.cells))
        cells_list = list(self.cells.values())
        energies = [c.energy for c in cells_list]
        avg_energy = np.mean(energies) if energies else 0.0
        states = {}
        for c in cells_list:
            states[c.state] = states.get(c.state, 0) + 1
        self.emit("generation", {
            "generation": self.generation,
            "population": len(self.cells),
            "average_energy": float(avg_energy),
            "energy_distribution": states,
        })

    def _tick_cpu(self, cell_positions, cell_velocities, concept_positions, dt):
        """CPU path: NumPy-vectorized forces + position update."""
        cfg = self.config
        epsilon = 1e-6
        dt_norm = dt / 1000.0

        # Gravity ramping + settling
        gravity_scale = self._compute_gravity_scale()
        settling = self._compute_settling_damping()
        effective_gravity = cfg.engine.concept_gravity_strength * gravity_scale

        # Concept gravity
        deltas = concept_positions[None, :, :] - cell_positions[:, None, :]
        distances = np.linalg.norm(deltas, axis=2, keepdims=True)
        distances_flat = distances.squeeze(2)

        in_range = distances_flat < cfg.engine.concept_gravity_range
        softening = 0.05
        gravity_mag = np.where(
            in_range & (distances_flat > 0.001),
            effective_gravity * distances_flat
            / (softening * softening + distances_flat ** 2 * cfg.engine.concept_gravity_sharpness)
            * dt_norm,
            0.0
        )

        safe_dist = np.maximum(distances, epsilon)
        direction = deltas / safe_dist
        gravity_forces = np.sum(gravity_mag[:, :, None] * direction, axis=1)

        # Similarity matrix forces
        matrix_forces = np.zeros_like(cell_positions)
        if self.field.similarity_matrix is not None:
            matrix = self.field.similarity_matrix
            amp_matrix = matrix.copy()
            if cfg.engine.tradeoff_amplifier != 1.0:
                neg_mask = amp_matrix < 0
                amp_matrix = np.where(neg_mask, amp_matrix * cfg.engine.tradeoff_amplifier, amp_matrix)
            np.fill_diagonal(amp_matrix, 0.0)
            matrix_forces = (cell_velocities @ amp_matrix) * cfg.engine.force_scale

        # Thermal noise
        noise = np.random.uniform(
            -cfg.engine.thermal_noise, cfg.engine.thermal_noise, cell_positions.shape)

        # Boundary forces
        boundary_forces = np.zeros_like(cell_positions)
        if cfg.engine.boundary_mode == "soft":
            margin = cfg.engine.boundary_margin
            strength = cfg.engine.boundary_strength
            low_mask = cell_positions < margin
            boundary_forces += np.where(low_mask, strength * (margin - cell_positions) / (margin + epsilon), 0)
            high_mask = cell_positions > (1.0 - margin)
            boundary_forces += np.where(high_mask, -strength * (cell_positions - (1.0 - margin)) / (margin + epsilon), 0)

        # Integrate (with settling damping applied to momentum)
        total_force = gravity_forces + matrix_forces + noise + boundary_forces
        effective_momentum = min(cfg.cell.momentum, 0.7) * settling
        cell_velocities = cell_velocities * effective_momentum + total_force * dt_norm
        cell_velocities = np.clip(cell_velocities, -0.3, 0.3)
        cell_positions = cell_positions + cell_velocities * dt_norm

        # Boundary reflection
        below = cell_positions < 0.0
        above = cell_positions > 1.0
        cell_velocities = np.where(below, np.abs(cell_velocities) * 0.5, cell_velocities)
        cell_velocities = np.where(above, -np.abs(cell_velocities) * 0.5, cell_velocities)
        cell_positions = np.clip(cell_positions, 0.0, 1.0)

        return cell_positions, cell_velocities

    # ─── Tick Loop ────────────────────────────────────────────

    def tick(self, dt: float = 16.67):
        """
        One simulation tick. The core loop:
        1. Compute forces (vectorized — NumPy)
        2. Update velocities and positions
        3. Metabolism (energy gain/loss)
        4. Reproduction and death
        5. Concept dynamics
        6. Generation bookkeeping
        """
        if len(self.cells) == 0:
            return

        self.tick_count += 1
        cfg = self.config
        n = self.n_dims

        # ── 1. Gather positions into arrays ──
        cell_list = list(self.cells.values())
        num_cells = len(cell_list)
        concept_list = list(self.concepts.values())
        num_concepts = len(concept_list)

        cell_positions = np.array([c.position for c in cell_list])      # (C, D)
        cell_velocities = np.array([c.velocity for c in cell_list])     # (C, D)
        concept_positions = np.array([c.position for c in concept_list])  # (K, D)

        # ── 2-6. Forces + position update (GPU or CPU path) ──
        if self._gpu_kernel is not None:
            cell_positions, cell_velocities = self._tick_gpu(
                cell_positions, cell_velocities, concept_positions, dt)
        else:
            cell_positions, cell_velocities = self._tick_cpu(
                cell_positions, cell_velocities, concept_positions, dt)

        # Write back to cells
        for i, cell in enumerate(cell_list):
            cell.position = cell_positions[i]
            cell.velocity = cell_velocities[i]

        # ── 7. Metabolism (vectorized) ──
        # Energy gain from nearby concepts — compute distances (cheap on CPU)
        deltas = concept_positions[None, :, :] - cell_positions[:, None, :]
        distances_flat = np.linalg.norm(deltas, axis=2)
        absorption_range = cfg.metabolism.concept_absorption_range
        in_absorption = distances_flat < absorption_range
        falloff_exp = cfg.metabolism.energy_falloff_exponent
        energy_gain = np.where(
            in_absorption,
            cfg.metabolism.max_energy_gain_per_tick / (1.0 + distances_flat ** falloff_exp),
            0.0
        )  # (C, K)
        total_gain = np.sum(energy_gain, axis=1)  # (C,)

        # Energy cost: base + movement
        movement_cost = np.linalg.norm(cell_velocities, axis=1) * cfg.metabolism.movement_cost_factor
        base_cost = cfg.metabolism.base_cost_per_dim * n
        total_cost = base_cost + movement_cost

        # Apply
        for i, cell in enumerate(cell_list):
            cell.energy += total_gain[i] - total_cost[i]
            cell.energy = max(cfg.metabolism.min_energy, min(cfg.cell.max_energy, cell.energy))
            cell.age += 1
            cell.update_state(cfg.cell.max_energy)

        # ── 8. Reproduction ──
        new_cells = []
        for cell in cell_list:
            if (cell.energy > cfg.cell.max_energy * cfg.cell.reproduction_threshold
                    and len(self.cells) + len(new_cells) < cfg.engine.max_population):
                child = cell.spawn_child(
                    transfer_ratio=cfg.cell.reproduction_transfer,
                    offset=cfg.cell.child_spawn_offset,
                    mutation=cfg.cell.mutation_strength,
                )
                new_cells.append(child)
                self.stats["total_reproductions"] += 1
                self.stats["total_cells_created"] += 1

        for child in new_cells:
            self.cells[child.id] = child

        # ── 9. Death (remove cold cells if over population) ──
        if len(self.cells) > cfg.engine.max_population:
            sorted_cells = sorted(self.cells.values(), key=lambda c: c.energy)
            excess = len(self.cells) - cfg.engine.max_population
            for cell in sorted_cells[:excess]:
                del self.cells[cell.id]

        # ── 10. Concept dynamics (simplified) ──
        # Skip concept dynamics when concepts come from attractors — they're anchored.
        # Attractor-based concepts define the gravitational landscape; moving them
        # destroys the basin structure the colony needs to differentiate.
        epsilon = 1e-6
        wm = self.field.metadata.get("worldModel", {})
        has_attractors = bool(wm.get("attractors"))
        if not has_attractors and self.field.similarity_matrix is not None and num_concepts > 1:
            concept_pos = np.array([c.position for c in concept_list])
            # Concepts attract/repel based on similarity matrix
            for i, ci in enumerate(concept_list):
                force = np.zeros(n)
                for j, cj in enumerate(concept_list):
                    if i == j:
                        continue
                    delta = cj.position - ci.position
                    dist = np.linalg.norm(delta) + epsilon
                    # Only use matrix if indices are valid (dimension-based concepts)
                    if i < self.field.similarity_matrix.shape[0] and j < self.field.similarity_matrix.shape[1]:
                        sim = self.field.similarity_matrix[i, j]
                    else:
                        sim = 0.0

                    if dist < cfg.concept_dynamics.min_separation:
                        # Hard repulsion at close range
                        force -= (delta / dist) * cfg.concept_dynamics.separation_strength
                    elif sim > 0:
                        force += (delta / dist) * sim * cfg.concept_dynamics.attract_scale
                    else:
                        force -= (delta / dist) * abs(sim) * cfg.concept_dynamics.repel_scale

                ci.velocity = ci.velocity * 0.98 + force * 0.01
                ci.position = np.clip(ci.position + ci.velocity * (dt / 1000.0), 0.0, 1.0)

        # ── 11. Peak population tracking ──
        self.stats["peak_population"] = max(self.stats["peak_population"], len(self.cells))

        # ── 12. Generation events ──
        if self.tick_count % cfg.engine.ticks_per_generation == 0:
            self.generation += 1
            cells_list = list(self.cells.values())
            energies = [c.energy for c in cells_list]
            avg_energy = np.mean(energies) if energies else 0.0

            states = {}
            for c in cells_list:
                states[c.state] = states.get(c.state, 0) + 1

            self.emit("generation", {
                "generation": self.generation,
                "population": len(self.cells),
                "average_energy": float(avg_energy),
                "energy_distribution": states,
            })

    # ─── Utilities ────────────────────────────────────────────

    def get_cell_positions(self) -> np.ndarray:
        """Return all cell positions as (N, D) array."""
        if not self.cells:
            return np.empty((0, self.n_dims))
        return np.array([c.position for c in self.cells.values()])

    def get_concept_positions(self) -> np.ndarray:
        """Return all concept positions as (K, D) array."""
        if not self.concepts:
            return np.empty((0, self.n_dims))
        return np.array([c.position for c in self.concepts.values()])

    def get_concept_names(self) -> list[str]:
        return [c.name for c in self.concepts.values()]
