"""
TwinCrack — Gradient-free optimization via twin learned perturbation agents.

No backpropagation. No autograd. No computational graph. No .backward().
Just two crackheads poking your model from opposite sides and comparing notes.

    from twin_crack import TwinCrack

    tc = TwinCrack(param_dict, loss_fn)
    
    for batch in data:
        tc.step(batch)
    
    # that's it. no gradients. no backward pass. just vibes and perturbation.

How it works:
    Bug A perturbs parameters in direction +d
    Bug B perturbs parameters in direction -d
    Compare losses: (loss_A - loss_B) / (2 * epsilon) ≈ gradient along d
    
    But the bugs aren't random. They LEARN where to poke.
    Both bugs share a brain that figures out which parameters are
    worth perturbing and by how much. They get rewarded for finding
    perturbation directions that give the most information.

    The smart name: Simultaneous Perturbation Stochastic Approximation
    with Learned Exploration Policy (SPSA-LEP)

    The real name: two crackheads with the same reward signal
    on opposite sides of the machine

Memory: ~6KB for the brain + zero optimizer state
         No momentum buffers. No second moments. No Adam tax.
         Two forward passes per step. That's the cost.

Works with: anything that has parameters and a loss function.
            Numpy arrays. CuPy arrays. Raw floats. Doesn't matter.
            If you can evaluate loss, TwinCrack can optimize.
"""

import numpy as np

try:
    import cupy as cp
    HAS_GPU = True
    xp = cp
except ImportError:
    HAS_GPU = False
    xp = np


def _to_np(a):
    if HAS_GPU and isinstance(a, cp.ndarray):
        return cp.asnumpy(a)
    return np.asarray(a)


def _to_xp(a):
    if HAS_GPU:
        return cp.asarray(a)
    return np.asarray(a)


# ═════════════════════════════════════════════════════════════════════════════
# PERTURBATION BRAIN — learns where to poke
# ═════════════════════════════════════════════════════════════════════════════

class PerturbBrain:
    """Tiny MLP that decides perturbation strategy.

    For each parameter group, outputs:
        magnitude:   [0.0001, 0.1]  — how hard to poke (epsilon)
        focus:       per-param probability mask — WHERE to poke
        correlation: [-1, 1] — how much to reuse last direction

    The brain sees:
        - Recent loss trajectory
        - Per-group: estimated gradient magnitude, noise level, sensitivity
        - Which perturbation directions gave useful signal last time

    Trained by REINFORCE:
        reward = information gained = |loss_A - loss_B| (more difference = better poke)
        + loss improvement bonus (the estimated gradient actually helped)
        + novelty bonus when stuck (crackhead mode)
    """

    def __init__(self, n_groups, obs_dim, hidden=48, lr=0.0005, seed=42):
        self.n_groups = n_groups
        self.out_dim = n_groups * 3  # magnitude, focus_logit, correlation per group
        self.lr = lr

        rng = np.random.RandomState(seed)
        h1, h2 = hidden, hidden // 2

        self.W1 = (rng.randn(obs_dim, h1) * np.sqrt(2.0/obs_dim)).astype(np.float32)
        self.b1 = np.zeros(h1, dtype=np.float32)
        self.W2 = (rng.randn(h1, h2) * np.sqrt(2.0/h1)).astype(np.float32)
        self.b2 = np.zeros(h2, dtype=np.float32)
        self.W3 = (rng.randn(h2, self.out_dim) * 0.01).astype(np.float32)
        self.b3 = np.zeros(self.out_dim, dtype=np.float32)

        # Memory
        self._h2 = np.zeros(h2, dtype=np.float32)

        # REINFORCE
        self._baseline = 0.0
        self._exp = []
        self._update_every = 16

        # Crackhead mode
        self._comfort_zone = np.zeros(self.out_dim, dtype=np.float32)
        self._comfort_alpha = 0.95
        self._loss_window = np.zeros(32, dtype=np.float32)
        self._loss_ptr = 0
        self._loss_count = 0

    def forward(self, obs):
        z1 = obs @ self.W1 + self.b1
        h1 = np.maximum(z1, 0)
        mem_n = min(len(self._h2), len(h1))
        h1[:mem_n] += 0.1 * self._h2[:mem_n]

        z2 = h1 @ self.W2 + self.b2
        h2 = np.maximum(z2, 0)
        self._h2 = h2.copy()

        raw = h2 @ self.W3 + self.b3

        # Add exploration noise
        noise = np.random.randn(self.out_dim).astype(np.float32) * 0.1
        raw_noisy = raw + noise

        decisions = {}
        for gi in range(self.n_groups):
            b = gi * 3
            sig = lambda x: 1.0 / (1.0 + np.exp(-np.clip(x, -10, 10)))

            # Magnitude: log-uniform [0.0001, 0.1]
            mag = 0.0001 * (0.1/0.0001) ** sig(raw_noisy[b])

            # Focus: sigmoid -> probability of perturbing each param
            # Higher = more focused perturbation (fewer params, stronger signal)
            focus = float(sig(raw_noisy[b+1]))

            # Correlation: tanh -> [-1, 1], reuse of previous direction
            corr = float(np.tanh(raw_noisy[b+2]))

            decisions[gi] = (float(mag), focus, corr)

        # Update comfort zone for crackhead mode
        self._comfort_zone = (self._comfort_alpha * self._comfort_zone +
                              (1.0 - self._comfort_alpha) * raw)

        cache = {'obs': obs.copy(), 'raw': raw.copy(), 'raw_noisy': raw_noisy,
                 'z1': z1, 'h1': h1, 'z2': z2, 'h2': h2, 'noise': noise}
        return decisions, cache

    def _novelty_bonus(self, raw):
        """Crackhead mode: reward novelty during plateaus."""
        n = min(self._loss_count, 32)
        if n < 8:
            return 0.0
        recent = self._loss_window[:n]
        variance = float(np.var(recent))
        if variance > 0.001:
            return 0.0  # Not stuck, no bonus

        # Distance from comfort zone = novelty
        dist = float(np.sqrt(np.sum((raw - self._comfort_zone) ** 2)))
        return min(dist * 0.3, 0.3)  # Cap at 0.3

    def record_loss(self, loss):
        self._loss_window[self._loss_ptr % 32] = loss
        self._loss_ptr += 1
        self._loss_count += 1

    def record(self, info_reward, loss_reward, cache):
        """Record outcome.
        info_reward: |loss_A - loss_B| — how useful was the perturbation
        loss_reward: -delta_loss — did the update help
        """
        novelty = self._novelty_bonus(cache['raw'])
        total = 0.5 * info_reward + 0.5 * loss_reward + novelty
        self._exp.append((cache['obs'], cache['raw_noisy'], total))
        self._baseline = 0.95 * self._baseline + 0.05 * total

    def maybe_update(self):
        if len(self._exp) < self._update_every:
            return False

        obs_b = np.stack([e[0] for e in self._exp])
        raw_b = np.stack([e[1] for e in self._exp])
        rew_b = np.array([e[2] for e in self._exp], dtype=np.float32)

        adv = (rew_b - self._baseline) / (np.std(rew_b) + 1e-8)

        Z1 = obs_b @ self.W1 + self.b1
        H1 = np.maximum(Z1, 0)
        Z2 = H1 @ self.W2 + self.b2
        H2 = np.maximum(Z2, 0)
        Rp = H2 @ self.W3 + self.b3

        dR = -adv[:, None] * (raw_b - Rp) / len(rew_b)
        dW3 = H2.T @ dR;  db3 = np.mean(dR, axis=0)
        dH2 = dR @ self.W3.T * (Z2 > 0)
        dW2 = H1.T @ dH2;  db2 = np.mean(dH2, axis=0)
        dH1 = dH2 @ self.W2.T * (Z1 > 0)
        dW1 = obs_b.T @ dH1;  db1 = np.mean(dH1, axis=0)

        clp = lambda x: np.clip(x, -1.0, 1.0)
        self.W1 -= self.lr * clp(dW1)
        self.b1 -= self.lr * clp(db1)
        self.W2 -= self.lr * clp(dW2)
        self.b2 -= self.lr * clp(db2)
        self.W3 -= self.lr * clp(dW3)
        self.b3 -= self.lr * clp(db3)

        self._exp.clear()
        return True

    def state_dict(self):
        return {k: v.copy() if isinstance(v, np.ndarray) else v
                for k, v in {
                    'W1': self.W1, 'b1': self.b1,
                    'W2': self.W2, 'b2': self.b2,
                    'W3': self.W3, 'b3': self.b3,
                    'h2': self._h2,
                    'comfort': self._comfort_zone,
                    'baseline': np.float32(self._baseline),
                }.items()}

    def load_state_dict(self, d):
        for k in ['W1','b1','W2','b2','W3','b3']:
            if k in d: setattr(self, k, d[k].astype(np.float32))
        if 'h2' in d: self._h2 = d['h2'].astype(np.float32)
        if 'comfort' in d: self._comfort_zone = d['comfort'].astype(np.float32)
        if 'baseline' in d: self._baseline = float(d['baseline'])


# ═════════════════════════════════════════════════════════════════════════════
# OBSERVER — watches the perturbation landscape
# ═════════════════════════════════════════════════════════════════════════════

class PerturbObserver:
    """Tracks perturbation outcomes to inform the brain."""

    def __init__(self, n_groups, window=32):
        self.n_groups = n_groups
        self.window = window

        # Loss
        self.losses = np.zeros(window, dtype=np.float32)
        self.loss_ptr = 0
        self.loss_count = 0

        # Per-group perturbation effectiveness
        self.info_gains = np.zeros((n_groups, window), dtype=np.float32)  # |loss_A - loss_B|
        self.grad_estimates = np.zeros((n_groups, window), dtype=np.float32)  # magnitude
        self.sensitivities = np.zeros((n_groups, window), dtype=np.float32)  # how much params responded
        self.ptr = 0

    def record_loss(self, loss):
        self.losses[self.loss_ptr % self.window] = loss
        self.loss_ptr += 1
        self.loss_count += 1

    def record_perturbation(self, gi, info_gain, grad_est, sensitivity):
        ptr = self.ptr % self.window
        self.info_gains[gi, ptr] = info_gain
        self.grad_estimates[gi, ptr] = grad_est
        self.sensitivities[gi, ptr] = sensitivity

    def advance(self):
        self.ptr += 1

    def observe(self):
        """Build observation vector.
        [0:6] loss stats
        Per group (n_groups x 6):
            info_gain_mean, info_gain_recent, grad_est_mean,
            sensitivity_mean, signal_to_noise, effectiveness
        """
        obs = []

        # Loss stats
        n = min(self.loss_count, self.window)
        if n > 1:
            L = self.losses[:n] if self.loss_count <= self.window else self.losses
            lm = float(np.mean(L))
            ls = float(np.std(L))
            x = np.arange(n, dtype=np.float32); x -= x.mean()
            trend = float(np.sum(x * L[:n]) / (np.sum(x*x) + 1e-8))
            plateau = 1.0 - min(ls / (abs(lm) + 1e-8), 1.0)
        else:
            lm = ls = trend = plateau = 0.0
        phase = min(self.loss_count / 5000.0, 1.0)
        obs.extend([lm, ls, trend, plateau, phase,
                     float(self.losses[(self.loss_ptr-1) % self.window]) if n > 0 else 0.0])

        # Per-group
        ng = min(self.ptr, self.window)
        for gi in range(self.n_groups):
            if ng > 0:
                ig = self.info_gains[gi, :ng]
                ge = self.grad_estimates[gi, :ng]
                se = self.sensitivities[gi, :ng]

                ig_mean = float(np.mean(ig))
                ig_recent = float(ig[-1])
                ge_mean = float(np.mean(ge))
                se_mean = float(np.mean(se))

                # Signal to noise: consistent gradient direction = good signal
                if ng >= 4:
                    half = ng // 2
                    stn = float(abs(np.mean(ge[-half:]) - np.mean(ge[:half])))
                else:
                    stn = 0.0

                # Effectiveness: info gain relative to epsilon spent
                eff = ig_mean / (se_mean + 1e-8)
            else:
                ig_mean = ig_recent = ge_mean = se_mean = stn = eff = 0.0

            obs.extend([ig_mean, ig_recent, ge_mean, se_mean, stn, eff])

        return np.array(obs, dtype=np.float32)


# ═════════════════════════════════════════════════════════════════════════════
# TWINCRACK — the gradient-free optimizer
# ═════════════════════════════════════════════════════════════════════════════

class TwinCrack:
    """Gradient-free optimization via twin learned perturbation.

    Two crackheads. Same reward. Opposite sides of the machine.
    No .backward(). No autograd. No computational graph.
    Just two forward passes and subtraction.

    Usage:
        # Define your parameters as a dict of numpy/cupy arrays
        params = {
            'embedding': np.random.randn(vocab, dim).astype(np.float32),
            'hidden_W': np.random.randn(dim, dim).astype(np.float32),
            'output_W': np.random.randn(dim, vocab).astype(np.float32),
        }

        # Define a loss function: params_dict, batch -> float
        def loss_fn(params, batch):
            h = batch @ params['hidden_W']
            logits = h @ params['output_W']
            return cross_entropy(logits, targets)

        # Create optimizer
        tc = TwinCrack(params, loss_fn, lr=0.001)

        # Train
        for batch in data:
            loss = tc.step(batch)
            print(f"loss: {loss:.4f} | brain: {tc.get_state()}")

    Groups: Each parameter is its own group by default.
            Pass group_map={'embedding': 0, 'hidden_W': 1, 'output_W': 1}
            to share perturbation strategy across parameters.

    The brain learns:
        - Which parameters are sensitive (poke them gently)
        - Which parameters are robust (poke them hard for more signal)
        - Where to focus perturbation (sparse poking = cleaner gradient)
        - When to reuse previous directions (momentum via correlation)
        - When to go crackhead (novelty seeking during plateaus)
    """

    def __init__(self, params, loss_fn, lr=0.001, group_map=None,
                 brain_hidden=48, brain_lr=0.0005,
                 n_perturbations=1, antithetic=True):
        """
        Args:
            params: dict of {name: np.ndarray} — your model parameters
            loss_fn: callable(params_dict, batch) -> float
            lr: learning rate for parameter updates
            group_map: dict of {param_name: group_index}
            brain_hidden: brain MLP size
            brain_lr: brain learning rate
            n_perturbations: how many perturbation pairs per step (more = better estimate, slower)
            antithetic: use +d/-d pairs (always True unless you're weird)
        """
        self.params = params
        self.param_names = list(params.keys())
        self.loss_fn = loss_fn
        self.lr = lr
        self.n_perturbations = n_perturbations
        self.antithetic = antithetic

        # Group mapping
        if group_map is not None:
            self.group_map = group_map
            self.n_groups = max(group_map.values()) + 1
        else:
            self.group_map = {name: i for i, name in enumerate(self.param_names)}
            self.n_groups = len(self.param_names)

        # Previous perturbation directions (for correlation/momentum)
        self._prev_directions = {name: None for name in self.param_names}

        # Gradient accumulator (estimated) — match each param's array type
        def _zeros_like(a):
            if HAS_GPU and isinstance(a, cp.ndarray):
                return cp.zeros_like(a)
            return np.zeros_like(a)
        self._grad_accum = {name: _zeros_like(params[name]) for name in self.param_names}

        # Observer and brain
        obs_dim = 6 + self.n_groups * 6
        self.observer = PerturbObserver(self.n_groups)
        self.brain = PerturbBrain(
            n_groups=self.n_groups,
            obs_dim=obs_dim,
            hidden=brain_hidden,
            lr=brain_lr,
        )

        self._prev_loss = None
        self._brain_cache = None
        self._step_count = 0

        # Update LR brain (separate from perturbation brain)
        # This one controls how much to actually move based on the estimated gradient
        self._update_momentum = {name: _zeros_like(params[name]) for name in self.param_names}
        self._momentum_beta = 0.9

    def _generate_direction(self, name, shape, mag, focus, corr):
        """Generate perturbation direction for a parameter.

        Uses brain's decisions:
            mag: perturbation magnitude
            focus: sparsity (higher = fewer params perturbed = cleaner signal)
            corr: correlation with previous direction (momentum-like)
        """
        # Match array module to the parameter's type
        p = self.params[name]
        _xp = cp if (HAS_GPU and isinstance(p, cp.ndarray)) else np

        # Random direction
        d = _xp.asarray(np.random.randn(*shape).astype(np.float32))

        # Focus: mask out some dimensions (sparse perturbation)
        # focus=0: perturb everything. focus=1: perturb ~10% of params
        if focus > 0.01:
            keep_prob = max(1.0 - focus * 0.9, 0.1)  # never go below 10%
            mask = _xp.asarray((np.random.rand(*shape) < keep_prob).astype(np.float32))
            d = d * mask
            # Rescale to maintain expected magnitude
            d = d / _xp.float32(max(keep_prob, 0.1))

        # Normalize and scale
        d_norm = _xp.float32(float(_xp.sqrt(_xp.sum(d * d))) + 1e-12)
        d = d / d_norm * _xp.float32(mag)

        # Correlation: blend with previous direction
        if corr != 0.0 and self._prev_directions[name] is not None:
            prev = self._prev_directions[name]
            d = _xp.float32(1.0 - abs(corr)) * d + _xp.float32(corr) * prev

        # Store for next step
        self._prev_directions[name] = d.copy()

        return d

    def step(self, batch):
        """One optimization step. No gradients required.

        1. Brain decides perturbation strategy
        2. Bug A: evaluate loss with +perturbation
        3. Bug B: evaluate loss with -perturbation
        4. Estimate gradient from difference
        5. Update parameters
        6. Reward brain based on outcome

        Args:
            batch: whatever your loss_fn expects as second argument

        Returns:
            float: current loss (unperturbed)
        """
        self._step_count += 1

        # ── Get current loss (unperturbed baseline) ──
        base_loss = float(self.loss_fn(self.params, batch))
        self.observer.record_loss(base_loss)
        self.brain.record_loss(base_loss)

        # ── Reward brain for previous step ──
        if self._prev_loss is not None and self._brain_cache is not None:
            loss_reward = float(np.clip(
                -(base_loss - self._prev_loss) / (abs(self._prev_loss) + 1e-8),
                -5, 5
            ))
            # Info reward computed during perturbation (stored)
            info_reward = getattr(self, '_last_info_reward', 0.0)
            self.brain.record(info_reward, loss_reward, self._brain_cache)
            self.brain.maybe_update()

        # ── Brain decides perturbation strategy ──
        obs = self.observer.observe()
        decisions, self._brain_cache = self.brain.forward(obs)

        # ── Generate perturbations and estimate gradients ──
        total_info = 0.0

        for _ in range(self.n_perturbations):
            # Generate direction for each parameter
            directions = {}
            for name in self.param_names:
                gi = self.group_map[name]
                mag, focus, corr = decisions.get(gi, (0.001, 0.0, 0.0))
                directions[name] = self._generate_direction(
                    name, self.params[name].shape, mag, focus, corr
                )

            # ── Bug A: +perturbation ──
            perturbed_plus = {}
            for name in self.param_names:
                perturbed_plus[name] = self.params[name] + directions[name]

            loss_plus = float(self.loss_fn(perturbed_plus, batch))

            # ── Bug B: -perturbation ──
            if self.antithetic:
                perturbed_minus = {}
                for name in self.param_names:
                    perturbed_minus[name] = self.params[name] - directions[name]

                loss_minus = float(self.loss_fn(perturbed_minus, batch))
            else:
                loss_minus = base_loss

            # ── Estimate gradient ──
            loss_diff = loss_plus - loss_minus
            info_gain = abs(loss_diff)
            total_info += info_gain

            for name in self.param_names:
                gi = self.group_map[name]
                mag, _, _ = decisions.get(gi, (0.001, 0.0, 0.0))
                d = directions[name]
                _px = cp if (HAS_GPU and isinstance(d, cp.ndarray)) else np

                # Gradient estimate along direction d
                # g ≈ (f(x+d) - f(x-d)) / (2 * ||d||)
                d_norm_sq = _px.sum(d * d) + _px.float32(1e-12)
                grad_est = d * _px.float32(loss_diff) / d_norm_sq

                self._grad_accum[name] = self._grad_accum[name] + grad_est

                # Record stats for observer
                ge_mag = float(_px.sqrt(_px.sum(grad_est * grad_est)))
                self.observer.record_perturbation(
                    gi,
                    info_gain=info_gain,
                    grad_est=ge_mag,
                    sensitivity=mag
                )

        # Average gradient estimates
        if self.n_perturbations > 1:
            for name in self.param_names:
                _px = cp if (HAS_GPU and isinstance(self._grad_accum[name], cp.ndarray)) else np
                self._grad_accum[name] /= _px.float32(self.n_perturbations)

        # ── Update parameters ──
        for name in self.param_names:
            g = self._grad_accum[name]
            _px = cp if (HAS_GPU and isinstance(g, cp.ndarray)) else np

            # Momentum
            self._update_momentum[name] = (
                _px.float32(self._momentum_beta) * self._update_momentum[name] +
                _px.float32(1.0 - self._momentum_beta) * g
            )

            # Clip
            m = self._update_momentum[name]
            m_max = _px.max(_px.abs(m))
            clip_val = _px.float32(1.0)
            if m_max > clip_val:
                m = m * (clip_val / m_max)

            # Update
            self.params[name] = self.params[name] - _px.float32(self.lr) * m

            # Reset accumulator
            self._grad_accum[name] = _px.zeros_like(self.params[name])

        self.observer.advance()
        self._prev_loss = base_loss
        self._last_info_reward = total_info / max(self.n_perturbations, 1)

        return base_loss

    def get_state(self):
        """Brain's current decisions for logging."""
        if self._brain_cache is None:
            return {}
        obs = self.observer.observe()
        decisions, _ = self.brain.forward(obs)
        state = {
            'step': self._step_count,
            'baseline': round(self.brain._baseline, 6),
        }
        for gi in range(self.n_groups):
            if gi in decisions:
                m, f, c = decisions[gi]
                name = [n for n, g in self.group_map.items() if g == gi]
                label = name[0] if name else f'g{gi}'
                state[f'{label}_mag'] = f'{m:.6f}'
                state[f'{label}_focus'] = f'{f:.3f}'
                state[f'{label}_corr'] = f'{c:.3f}'
        return state

    # ── Multi-step convenience ───────────────────────────────────────────────

    def train(self, data_iter, steps=1000, log_every=50, callback=None):
        """Train for N steps with optional logging.

        Args:
            data_iter: iterable of batches (will cycle)
            steps: number of optimization steps
            log_every: print loss every N steps
            callback: optional fn(step, loss, brain_state) called each step
        """
        import itertools
        data_cycle = itertools.cycle(data_iter) if hasattr(data_iter, '__iter__') else data_iter

        for step in range(steps):
            batch = next(data_cycle)
            loss = self.step(batch)

            if callback:
                callback(step, loss, self.get_state())

            if log_every and step % log_every == 0:
                state = self.get_state()
                print(f"[{step:>6d}] loss={loss:.4f} | brain={state}")

    # ── Save / Load ──────────────────────────────────────────────────────────

    def save(self, path):
        """Save brain + parameter state."""
        state = self.brain.state_dict()
        # Save parameters
        for name in self.param_names:
            state[f'param_{name}'] = _to_np(self.params[name])
            state[f'momentum_{name}'] = _to_np(self._update_momentum[name])
            if self._prev_directions[name] is not None:
                state[f'prevdir_{name}'] = _to_np(self._prev_directions[name])
        state['step_count'] = np.int32(self._step_count)
        state['prev_loss'] = np.float32(self._prev_loss or 0)
        # Observer
        state['obs_losses'] = self.observer.losses
        state['obs_loss_ptr'] = np.int32(self.observer.loss_ptr)
        state['obs_loss_count'] = np.int32(self.observer.loss_count)
        state['obs_info'] = self.observer.info_gains
        state['obs_grad'] = self.observer.grad_estimates
        state['obs_sens'] = self.observer.sensitivities
        state['obs_ptr'] = np.int32(self.observer.ptr)
        np.savez(path, **state)

    def load(self, path):
        """Load brain + parameter state."""
        d = np.load(path, allow_pickle=True)
        self.brain.load_state_dict(d)
        for name in self.param_names:
            if f'param_{name}' in d:
                self.params[name] = _to_xp(d[f'param_{name}'])
            if f'momentum_{name}' in d:
                self._update_momentum[name] = _to_xp(d[f'momentum_{name}'])
            if f'prevdir_{name}' in d:
                self._prev_directions[name] = _to_xp(d[f'prevdir_{name}'])
        self._step_count = int(d.get('step_count', 0))
        pl = float(d.get('prev_loss', 0))
        self._prev_loss = pl if pl != 0 else None
        if 'obs_losses' in d:
            self.observer.losses = d['obs_losses']
            self.observer.loss_ptr = int(d['obs_loss_ptr'])
            self.observer.loss_count = int(d['obs_loss_count'])
        if 'obs_info' in d:
            self.observer.info_gains = d['obs_info']
            self.observer.grad_estimates = d['obs_grad']
            self.observer.sensitivities = d['obs_sens']
            self.observer.ptr = int(d['obs_ptr'])


# ═════════════════════════════════════════════════════════════════════════════
# QUICK TEST — verify it works
# ═════════════════════════════════════════════════════════════════════════════

def _smoke_test():
    """Optimize a simple quadratic. No gradients needed."""
    print("=" * 60)
    print("TwinCrack smoke test: optimize f(x) = ||x - target||²")
    print("=" * 60)

    target = np.array([3.0, -2.0, 7.0, 1.0, -5.0], dtype=np.float32)
    params = {'x': np.zeros(5, dtype=np.float32)}

    def loss_fn(p, batch):
        diff = p['x'] - target
        return float(np.sum(diff * diff))

    tc = TwinCrack(params, loss_fn, lr=0.01, n_perturbations=3)

    print(f"\nTarget: {target}")
    print(f"Start:  {params['x']}")
    print(f"Initial loss: {loss_fn(params, None):.4f}\n")

    for step in range(200):
        loss = tc.step(None)
        if step % 20 == 0:
            print(f"[{step:>4d}] loss={loss:.4f} x={_to_np(params['x']).round(2)}")

    print(f"\nFinal:  {_to_np(params['x']).round(3)}")
    print(f"Target: {target}")
    print(f"Final loss: {loss_fn(params, None):.6f}")
    print(f"\nBrain state: {tc.get_state()}")
    print("=" * 60)


USAGE = """
╔══════════════════════════════════════════════════════════════════════════╗
║                        TwinCrack                                       ║
║          Two crackheads. No gradients. No .backward().                 ║
╠══════════════════════════════════════════════════════════════════════════╣
║                                                                        ║
║  from twin_crack import TwinCrack                                      ║
║                                                                        ║
║  params = {'W': np.random.randn(768, 768).astype(np.float32)}         ║
║  def loss_fn(p, batch): return your_model(p, batch)                   ║
║                                                                        ║
║  tc = TwinCrack(params, loss_fn, lr=0.001)                            ║
║  tc.step(batch)   # no .backward(). just this.                        ║
║                                                                        ║
╠══════════════════════════════════════════════════════════════════════════╣
║                                                                        ║
║  How it works:                                                         ║
║    Bug A pokes parameters in direction +d                              ║
║    Bug B pokes parameters in direction -d                              ║
║    Compare losses. Difference = gradient estimate.                     ║
║    Bugs learn where to poke for maximum information.                   ║
║                                                                        ║
║  What you don't need:                                                  ║
║    ✗ PyTorch autograd          ✗ Computational graph                   ║
║    ✗ .backward()               ✗ Gradient checkpointing               ║
║    ✗ Mixed precision scaling   ✗ Gradient accumulation                 ║
║    ✗ Adam's 4GB of state       ✗ Any of that crap                     ║
║                                                                        ║
║  What you DO need:                                                     ║
║    ✓ Parameters (numpy arrays)                                         ║
║    ✓ A loss function                                                   ║
║    ✓ That's it                                                         ║
║                                                                        ║
║  Memory: ~6KB brain + your parameters. Nothing else.                   ║
║  Cost: 2-3 forward passes per step (no backward pass).                 ║
║  Brain: Learns where to poke. Gets addicted to novelty when stuck.     ║
║                                                                        ║
║  The smart name: SPSA with Learned Exploration Policy                  ║
║  The real name: two crackheads on opposite sides of the machine        ║
║                                                                        ║
╚══════════════════════════════════════════════════════════════════════════╝
"""


if __name__ == '__main__':
    print(USAGE)
    _smoke_test()
