#!/usr/bin/env python3
"""
Qriton HLM — Energy Language CLI

An interactive shell for programming Hopfield energy landscapes.
Inspect, inject, remove, move basins. Apply changes to live models.
Chain operations into scripts. The assembler for energy computation.

Usage:
    qriton-hlm                               # interactive REPL
    qriton-hlm --checkpoint model.pt         # load specific model
    qriton-hlm --script program.hlm          # run a script
    echo "survey 0" | qriton-hlm -c model.pt # pipe commands

Commands:
    load <path>           Load a checkpoint
    info                  Show loaded model info
    survey <layer>        Find all basins in a layer
    survey-all            Survey all layers
    inject <layer> <seed> [strength]   Inject a basin
    remove <layer> <seed> [strength]   Remove closest basin
    move <layer> <seed> [strength]     Move closest basin
    verify <layer> <seed>              Check if target is a basin
    apply <layer>         Apply last surgery to live model
    restore <layer>       Restore layer to original
    restore-all           Restore all layers
    status                Show modification status
    generate <prompt>     Generate text (if LM loaded)
    set <param> <value>   Set parameter (beta, inits, strength, temp, tokens, topk)
    save <path>           Save modified W matrices
    diff <layer>          Show W diff stats vs original
    history               Show operation history
    help                  Show this help
    quit / exit           Exit
"""

import sys
import os
import argparse
import time
import json
try:
    import readline  # enables arrow keys and history in REPL
except ImportError:
    pass  # Windows doesn't have readline
from pathlib import Path

import torch
import torch.nn.functional as F
import numpy as np

from qriton_hlm.core import (
    find_basins, compute_energy, poly_interaction,
    inject_basin, remove_basin, move_basin,
    verify_basin_exists, load_W_from_checkpoint,
    count_hopfield_layers,
)

DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'


class EnergyLang:
    """Energy Language interpreter / REPL."""

    def __init__(self):
        self.model = None
        self.tokenizer = None
        self.config = None
        self.checkpoint_path = None
        self.model_label = None

        # Cached W matrices (layer -> W tensor)
        self._w_cache = {}
        self._w_backups = {}  # layer -> original W

        # Surgery state
        self._last_surgery = {}  # layer, W_before, W_after, target, ...
        self._concepts = {}     # name -> {'states': [...], 'centroid': tensor}
        self._history = []

        # Parameters
        self.params = {
            'beta': 7.0,
            'inits': 200,
            'strength': 0.1,
            'degree': 3,
            'eps': 5e-3,
            'temp': 0.8,
            'tokens': 50,
            'topk': 40,
        }

    # ── Model loading ────────────────────────────────────────────────

    def load_checkpoint(self, path):
        """Load a model checkpoint."""
        path = path.strip().strip('"').strip("'")
        if not os.path.exists(path):
            return f"File not found: {path}"

        self.checkpoint_path = path
        self._w_cache.clear()
        self._w_backups.clear()
        self._last_surgery.clear()

        # Try loading as full HLM3 model (for text generation)
        try:
            from hlm3_model import HLM3
            from data_utils import Tokenizer

            ckpt = torch.load(path, map_location=DEVICE, weights_only=False)
            config = ckpt.get('config', {})
            state = ckpt.get('model', ckpt.get('model_state', {}))

            # Fix log_beta shape
            for k in list(state.keys()):
                if 'log_beta' in k and state[k].dim() == 0:
                    state[k] = state[k].unsqueeze(0)

            if 'vocabSize' in config:
                model = HLM3(config).to(DEVICE)
                model.load_state_dict(state, strict=False)
                model.eval()
                self.model = model
                self.config = config

                # Load tokenizer
                tok_dir = os.path.dirname(path)
                tok_path = os.path.join(tok_dir, 'tokenizer.json')
                if not os.path.exists(tok_path):
                    tok_path = 'D:/HLM2/output-hlm3-large-ffn/tokenizer.json'
                if os.path.exists(tok_path):
                    self.tokenizer = Tokenizer.from_json(tok_path)

                n_layers = len(model.blocks)
                n_params = sum(p.numel() for p in model.parameters())
                d_model = config.get('dModel', '?')
                self.model_label = os.path.basename(os.path.dirname(path))
                self._log(f"load {path}")
                return (f"Loaded HLM3 model: {n_params/1e6:.1f}M params, "
                        f"d={d_model}, {n_layers} layers, device={DEVICE}\n"
                        f"  Text generation: {'enabled' if self.tokenizer else 'no tokenizer'}")
        except Exception:
            pass

        # Fallback: just load W matrices (works for any checkpoint with hopfield.W)
        try:
            ckpt = torch.load(path, map_location=DEVICE, weights_only=False)
            state = ckpt.get('model', ckpt.get('model_state', ckpt))
            w_keys = sorted([k for k in state if 'hopfield.W' in k])
            if w_keys:
                self.model = None
                self.tokenizer = None
                self.config = {'_w_keys': w_keys, '_state': state}
                n_layers = len(w_keys)
                d_model = state[w_keys[0]].shape[0]
                self.model_label = os.path.basename(os.path.dirname(path))
                self._log(f"load {path}")
                return (f"Loaded W matrices: {n_layers} layers, d={d_model}\n"
                        f"  Text generation: disabled (W-only mode)")
            return f"No hopfield.W keys found in {path}"
        except Exception as e:
            return f"Failed to load: {e}"

    def _get_beta(self, layer):
        """Get beta for a layer — uses learned log_beta from checkpoint if available."""
        import math
        layer = int(layer)
        # Try live model
        if self.model is not None and hasattr(self.model, 'blocks'):
            try:
                log_beta = self.model.blocks[layer].hopfield.log_beta.data
                return math.exp(log_beta.float().mean().item())
            except (AttributeError, IndexError):
                pass
        # Try checkpoint state dict
        if self.config and '_state' in self.config:
            state = self.config['_state']
            for key in state:
                if f'.{layer}.' in key and 'log_beta' in key:
                    return math.exp(state[key].float().mean().item())
        # Fallback to default
        return self.params['beta']

    def _get_W(self, layer):
        """Get W matrix for a layer (from cache, model, or state dict)."""
        layer = int(layer)
        if layer in self._w_cache:
            return self._w_cache[layer]

        if self.model is not None and hasattr(self.model, 'blocks'):
            W = self.model.blocks[layer].hopfield.W.data.clone()
            self._w_cache[layer] = W
            return W

        if self.config and '_state' in self.config:
            w_keys = self.config['_w_keys']
            if layer < len(w_keys):
                W = self.config['_state'][w_keys[layer]].to(DEVICE)
                self._w_cache[layer] = W
                return W

        raise ValueError(f"Cannot access layer {layer}")

    def _num_layers(self):
        if self.model is not None and hasattr(self.model, 'blocks'):
            return len(self.model.blocks)
        if self.config and '_w_keys' in self.config:
            return len(self.config['_w_keys'])
        if self.checkpoint_path:
            return count_hopfield_layers(self.checkpoint_path)
        return 0

    def _log(self, cmd):
        self._history.append((time.strftime('%H:%M:%S'), cmd))

    # ── Commands ─────────────────────────────────────────────────────

    def cmd_info(self):
        if not self.checkpoint_path:
            return "No model loaded. Use: load <path>"
        lines = [f"Checkpoint: {self.checkpoint_path}"]
        lines.append(f"Label: {self.model_label}")
        lines.append(f"Layers: {self._num_layers()}")
        if self.model:
            n_params = sum(p.numel() for p in self.model.parameters())
            lines.append(f"Params: {n_params/1e6:.1f}M")
            if self.config:
                d = self.config.get('dModel', self.config.get('d_model', '?'))
                lines.append(f"d_model: {d}")
        lines.append(f"Device: {DEVICE}")
        lines.append(f"Text gen: {'yes' if self.tokenizer else 'no'}")
        modified = [l for l in self._w_backups]
        if modified:
            lines.append(f"Modified layers: {modified}")
        lines.append(f"\nParameters: {json.dumps(self.params, indent=2)}")
        return "\n".join(lines)

    def cmd_survey(self, layer):
        layer = int(layer)
        W = self._get_W(layer)
        d_model = W.shape[0]
        beta = self._get_beta(layer)
        num_inits = int(self.params['inits'])

        t0 = time.time()
        basins, ids, trajectories = find_basins(
            W, d_model, num_inits=num_inits, beta=beta, device=DEVICE)
        elapsed = time.time() - t0

        conv_iters = [len(t) - 1 for t in trajectories]

        lines = [
            f"Layer {layer}: {len(basins)} unique basins "
            f"({num_inits} inits, beta={beta:.1f}, {elapsed:.1f}s)",
            f"  W: {d_model}x{d_model}, "
            f"spectral={torch.linalg.norm(W, ord=2).item():.4f}",
            f"  Mean convergence: {np.mean(conv_iters):.1f} iters",
            "",
        ]
        for i in range(min(len(basins), 20)):
            e = compute_energy(basins[i], W)
            pop = ids.count(i)
            pct = 100 * pop / num_inits
            lines.append(f"  B{i:2d}  E={e:+.4f}  pop={pop:3d} ({pct:.0f}%)")
        if len(basins) > 20:
            lines.append(f"  ... and {len(basins) - 20} more")

        self._log(f"survey {layer}")
        return "\n".join(lines)

    def cmd_survey_all(self):
        n = self._num_layers()
        lines = [f"Surveying {n} layers...", ""]
        t0 = time.time()
        for l in range(n):
            W = self._get_W(l)
            d = W.shape[0]
            basins, ids, trajs = find_basins(
                W, d, num_inits=int(self.params['inits']),
                beta=self._get_beta(layer), device=DEVICE)
            mean_iters = np.mean([len(t) - 1 for t in trajs])

            # Get learned beta if available
            beta_str = ""
            if self.model and hasattr(self.model, 'blocks'):
                lb = self.model.blocks[l].hopfield.log_beta.data
                beta_val = float(torch.exp(lb[0] if lb.dim() > 0 else lb))
                beta_str = f"  beta={beta_val:.2f}"

            mod = " [MODIFIED]" if l in self._w_backups else ""
            lines.append(
                f"  L{l}: {len(basins):3d} basins  "
                f"avg_iters={mean_iters:.0f}{beta_str}{mod}")

        elapsed = time.time() - t0
        lines.append(f"\nTotal: {elapsed:.1f}s")
        self._log("survey-all")
        return "\n".join(lines)

    def _make_target(self, seed, d_model):
        gen = torch.Generator(device='cpu')
        gen.manual_seed(int(seed))
        target = torch.randn(d_model, generator=gen).to(DEVICE)
        return target / target.norm()

    def cmd_inject(self, layer, seed, strength=None):
        layer = int(layer)
        strength = float(strength) if strength else self.params['strength']
        W = self._get_W(layer)
        d = W.shape[0]
        target = self._make_target(seed, d)

        # Before
        exists_before, _, cos_before, iters_before = verify_basin_exists(
            W, target, beta=self._get_beta(layer))
        basins_before, _, _ = find_basins(
            W, d, num_inits=100, beta=self._get_beta(layer), device=DEVICE)

        # Surgery
        W_new = inject_basin(W, target, strength=strength)

        # After
        exists_after, _, cos_after, iters_after = verify_basin_exists(
            W_new, target, beta=self._get_beta(layer))
        basins_after, _, _ = find_basins(
            W_new, d, num_inits=100, beta=self._get_beta(layer), device=DEVICE)

        self._last_surgery = {
            'layer': layer, 'W_before': W, 'W_after': W_new,
            'target': target, 'operation': 'inject',
        }
        if layer not in self._w_backups:
            self._w_backups[layer] = W.clone()
        self._w_cache[layer] = W_new  # update cache

        delta = len(basins_after) - len(basins_before)
        self._log(f"inject {layer} {seed} {strength}")
        lines = [
            f"Inject L{layer} seed={seed} strength={strength}",
            f"  Before: {len(basins_before)} basins, "
            f"target is basin: {exists_before} (cos={cos_before:.4f})",
            f"  After:  {len(basins_after)} basins ({'+' if delta >= 0 else ''}{delta}), "
            f"target is basin: {exists_after} (cos={cos_after:.4f})",
        ]
        if exists_after and not exists_before:
            lines.append("  >> Basin successfully programmed!")
        return "\n".join(lines)

    def cmd_remove(self, layer, seed, strength=None):
        layer = int(layer)
        strength = float(strength) if strength else self.params['strength']
        W = self._get_W(layer)
        d = W.shape[0]

        basins_before, _, _ = find_basins(
            W, d, num_inits=100, beta=self._get_beta(layer), device=DEVICE)

        # Find closest basin to seed target
        target = self._make_target(seed, d)
        if basins_before:
            sims = [F.cosine_similarity(target.unsqueeze(0), b.unsqueeze(0)).item()
                    for b in basins_before]
            closest = basins_before[int(np.argmax(sims))]
            target = closest

        W_new = remove_basin(W, target, strength=strength)
        exists_after, _, cos_after, _ = verify_basin_exists(
            W_new, target, beta=self._get_beta(layer))
        basins_after, _, _ = find_basins(
            W_new, d, num_inits=100, beta=self._get_beta(layer), device=DEVICE)

        self._last_surgery = {
            'layer': layer, 'W_before': W, 'W_after': W_new,
            'target': target, 'operation': 'remove',
        }
        if layer not in self._w_backups:
            self._w_backups[layer] = W.clone()
        self._w_cache[layer] = W_new

        delta = len(basins_after) - len(basins_before)
        self._log(f"remove {layer} {seed} {strength}")
        lines = [
            f"Remove L{layer} seed={seed} strength={strength}",
            f"  Before: {len(basins_before)} basins",
            f"  After:  {len(basins_after)} basins ({'+' if delta >= 0 else ''}{delta})",
            f"  Target still exists: {exists_after} (cos={cos_after:.4f})",
        ]
        if not exists_after:
            lines.append("  >> Basin successfully removed!")
        return "\n".join(lines)

    def cmd_move(self, layer, seed, strength=None):
        layer = int(layer)
        strength = float(strength) if strength else self.params['strength']
        W = self._get_W(layer)
        d = W.shape[0]

        basins_before, _, _ = find_basins(
            W, d, num_inits=100, beta=self._get_beta(layer), device=DEVICE)

        source_target = self._make_target(seed, d)
        if basins_before:
            sims = [F.cosine_similarity(source_target.unsqueeze(0), b.unsqueeze(0)).item()
                    for b in basins_before]
            source = basins_before[int(np.argmax(sims))]
        else:
            source = source_target

        dest = self._make_target(int(seed) + 1000, d)
        W_new = move_basin(W, source, dest, strength=strength)

        src_exists, _, src_cos, _ = verify_basin_exists(
            W_new, source, beta=self._get_beta(layer))
        dst_exists, _, dst_cos, _ = verify_basin_exists(
            W_new, dest, beta=self._get_beta(layer))
        basins_after, _, _ = find_basins(
            W_new, d, num_inits=100, beta=self._get_beta(layer), device=DEVICE)

        self._last_surgery = {
            'layer': layer, 'W_before': W, 'W_after': W_new,
            'target': dest, 'operation': 'move',
        }
        if layer not in self._w_backups:
            self._w_backups[layer] = W.clone()
        self._w_cache[layer] = W_new

        self._log(f"move {layer} {seed} {strength}")
        return (
            f"Move L{layer} seed={seed} strength={strength}\n"
            f"  Source: {'still exists' if src_exists else 'GONE'} "
            f"(cos={src_cos:.4f})\n"
            f"  Dest:   {'EXISTS' if dst_exists else 'missing'} "
            f"(cos={dst_cos:.4f})\n"
            f"  Basins: {len(basins_before)} -> {len(basins_after)}")

    def cmd_verify(self, layer, seed):
        layer = int(layer)
        W = self._get_W(layer)
        d = W.shape[0]
        target = self._make_target(seed, d)

        exists, final, cos, iters = verify_basin_exists(
            W, target, beta=self._get_beta(layer))
        energy = compute_energy(final, W)

        self._log(f"verify {layer} {seed}")
        return (
            f"Verify L{layer} seed={seed}\n"
            f"  Is basin: {exists}\n"
            f"  Cos similarity: {cos:.6f}\n"
            f"  Converged in: {iters} iters\n"
            f"  Energy at attractor: {energy:.4f}")

    def cmd_apply(self, layer):
        layer = int(layer)
        if self.model is None or not hasattr(self.model, 'blocks'):
            return "No full model loaded. Apply only works with HLM3 models."

        if layer not in self._w_cache:
            return f"No modified W for layer {layer}. Run surgery first."

        W_new = self._w_cache[layer]
        W_live = self.model.blocks[layer].hopfield.W

        if W_new.shape != W_live.shape:
            return f"Shape mismatch: {list(W_new.shape)} vs {list(W_live.shape)}"

        # Backup (first-write)
        if layer not in self._w_backups:
            self._w_backups[layer] = W_live.data.clone()

        with torch.no_grad():
            W_live.data.copy_(W_new)

        self._log(f"apply {layer}")
        modified = sorted(self._w_backups.keys())
        return (f"Applied L{layer} to live model.\n"
                f"  Modified layers: {modified}\n"
                f"  Text generation now uses modified W.")

    def cmd_restore(self, layer):
        layer = int(layer)
        if layer not in self._w_backups:
            return f"Layer {layer} has not been modified."

        if self.model and hasattr(self.model, 'blocks'):
            with torch.no_grad():
                self.model.blocks[layer].hopfield.W.data.copy_(
                    self._w_backups[layer])

        # Restore cache too
        self._w_cache[layer] = self._w_backups[layer].clone()
        del self._w_backups[layer]

        self._log(f"restore {layer}")
        remaining = sorted(self._w_backups.keys())
        return (f"Restored L{layer} to original.\n"
                f"  Still modified: {remaining if remaining else 'none'}")

    def cmd_restore_all(self):
        if not self._w_backups:
            return "No layers modified."
        restored = sorted(self._w_backups.keys())
        for l in list(self._w_backups.keys()):
            if self.model and hasattr(self.model, 'blocks'):
                with torch.no_grad():
                    self.model.blocks[l].hopfield.W.data.copy_(
                        self._w_backups[l])
            self._w_cache[l] = self._w_backups[l].clone()
            del self._w_backups[l]

        self._log("restore-all")
        return f"Restored layers {restored} to original."

    def cmd_status(self):
        if not self.checkpoint_path:
            return "No model loaded."
        n = self._num_layers()
        lines = [f"Model: {self.model_label} ({n} layers)"]
        for l in range(n):
            mod = "MODIFIED" if l in self._w_backups else "original"
            lines.append(f"  L{l}: {mod}")
        return "\n".join(lines)

    def cmd_capture(self, layer, text, concept_name=None):
        """Capture a concept from text at a given layer."""
        layer = int(layer)
        if self.model is None or self.tokenizer is None:
            return "No model loaded. Use: load <checkpoint>"

        tokens = self.tokenizer.encode(text)
        input_ids = torch.tensor([tokens], device=DEVICE)

        # Hook to capture Hopfield state after convergence
        captured = {}
        def hook_fn(module, inp, output):
            if isinstance(output, dict) and 'state' in output:
                captured['state'] = output['state'].detach()
            elif isinstance(output, tuple):
                # PolyHopfieldLayer.forward() returns (settled, energy, iters)
                captured['state'] = output[0].detach()
            elif isinstance(output, torch.Tensor):
                captured['state'] = output.detach()

        block = self.model.blocks[layer]
        handle = block.hopfield.register_forward_hook(hook_fn)

        with torch.no_grad():
            self.model(input_ids)
        handle.remove()

        if 'state' not in captured:
            return f"Could not capture state from layer {layer}"

        state = captured['state']
        if state.dim() == 3:
            state = state.squeeze(0).mean(dim=0)
        elif state.dim() == 2:
            state = state.mean(dim=0)
        state = state / state.norm()

        W = self._get_W(layer)
        energy = compute_energy(state, W)
        is_basin, _, cos, iters = verify_basin_exists(
            W, state, beta=self._get_beta(layer))

        lines = [f"Captured L{layer}: \"{text[:50]}\"",
                 f"  Energy: {energy:.4f} | Basin: {is_basin} (cos={cos:.4f}, {iters} iters)"]

        if concept_name:
            if concept_name not in self._concepts:
                self._concepts[concept_name] = {'states': [], 'centroid': None}
            self._concepts[concept_name]['states'].append(state)
            states = torch.stack(self._concepts[concept_name]['states'])
            centroid = states.mean(dim=0)
            self._concepts[concept_name]['centroid'] = centroid / centroid.norm()
            n = len(self._concepts[concept_name]['states'])
            lines.append(f"  -> Added to concept '{concept_name}' ({n} samples)")

        self._log(f"capture {layer} {concept_name or ''} {text[:30]}")
        return "\n".join(lines)

    def cmd_inject_concept(self, layer, concept_name, strength=None):
        """Inject a captured concept as a new basin."""
        layer = int(layer)
        strength = float(strength) if strength else self.params['strength']

        if concept_name not in self._concepts:
            available = list(self._concepts.keys()) if self._concepts else "none"
            return f"Unknown concept '{concept_name}'. Available: {available}"

        concept = self._concepts[concept_name]
        if concept['centroid'] is None:
            return f"Concept '{concept_name}' has no samples. Use capture first."

        target = concept['centroid'].to(DEVICE)
        W = self._get_W(layer)
        d = W.shape[0]

        exists_before, _, cos_before, _ = verify_basin_exists(
            W, target, beta=self._get_beta(layer))
        basins_before, _, _ = find_basins(
            W, d, num_inits=100, beta=self._get_beta(layer), device=DEVICE)

        W_new = inject_basin(W, target, strength=strength)

        exists_after, _, cos_after, _ = verify_basin_exists(
            W_new, target, beta=self._get_beta(layer))
        basins_after, _, _ = find_basins(
            W_new, d, num_inits=100, beta=self._get_beta(layer), device=DEVICE)

        self._w_cache[layer] = W_new
        if layer not in self._w_backups:
            self._w_backups[layer] = W.clone()

        delta = len(basins_after) - len(basins_before)
        n = len(concept['states'])
        self._log(f"inject-concept {layer} {concept_name} {strength}")
        lines = [
            f"Inject concept '{concept_name}' ({n} samples) into L{layer} strength={strength}",
            f"  Before: {len(basins_before)} basins, concept is basin: {exists_before} (cos={cos_before:.4f})",
            f"  After:  {len(basins_after)} basins ({'+' if delta >= 0 else ''}{delta}), "
            f"concept is basin: {exists_after} (cos={cos_after:.4f})",
        ]
        if exists_after and not exists_before:
            lines.append("  >> Concept successfully injected!")
        return "\n".join(lines)

    def cmd_remove_concept(self, layer, concept_name, strength=None):
        """Remove a captured concept's basin."""
        layer = int(layer)
        strength = float(strength) if strength else self.params['strength']

        if concept_name not in self._concepts:
            return f"Unknown concept '{concept_name}'"

        target = self._concepts[concept_name]['centroid'].to(DEVICE)
        W = self._get_W(layer)

        W_new = remove_basin(W, target, strength=strength)
        exists_after, _, cos_after, _ = verify_basin_exists(
            W_new, target, beta=self._get_beta(layer))

        self._w_cache[layer] = W_new
        if layer not in self._w_backups:
            self._w_backups[layer] = W.clone()

        self._log(f"remove-concept {layer} {concept_name} {strength}")
        return (f"Removed concept '{concept_name}' from L{layer}\n"
                f"  Basin still exists: {exists_after} (cos={cos_after:.4f})")

    def cmd_concepts(self):
        """List captured concepts."""
        if not self._concepts:
            return "No concepts captured. Use: capture <layer> <concept> <text>"
        lines = ["Captured concepts:"]
        for name, c in self._concepts.items():
            lines.append(f"  {name}: {len(c['states'])} samples")
        return "\n".join(lines)

    def cmd_blend(self, concept_a, concept_b, new_name, ratio=None):
        """Blend two concepts."""
        ratio = float(ratio) if ratio else 0.5
        if concept_a not in self._concepts or concept_b not in self._concepts:
            return f"Both concepts must exist. Have: {list(self._concepts.keys())}"
        ca = self._concepts[concept_a]['centroid']
        cb = self._concepts[concept_b]['centroid']
        blended = ratio * ca + (1 - ratio) * cb
        blended = blended / blended.norm()
        self._concepts[new_name] = {'states': [], 'centroid': blended}
        self._log(f"blend {concept_a} {concept_b} -> {new_name} ({ratio})")
        return (f"Blended '{concept_a}' ({ratio:.0%}) + '{concept_b}' ({1-ratio:.0%}) "
                f"-> '{new_name}'")

    def cmd_strengthen(self, layer, seed, factor=None):
        """Strengthen a basin."""
        layer, seed = int(layer), int(seed)
        factor = float(factor) if factor else 2.0
        W = self._get_W(layer)
        d = W.shape[0]
        target = self._make_target(seed, d)
        basins, _, _ = find_basins(W, d, num_inits=100, beta=self._get_beta(layer), device=DEVICE)
        if basins:
            sims = [F.cosine_similarity(target.unsqueeze(0), b.unsqueeze(0)).item() for b in basins]
            target = basins[int(np.argmax(sims))]
        e_before = compute_energy(target, W)
        W_new = inject_basin(W, target, strength=self.params['strength'] * factor)
        e_after = compute_energy(target, W_new)
        self._w_cache[layer] = W_new
        if layer not in self._w_backups:
            self._w_backups[layer] = W.clone()
        self._log(f"strengthen {layer} {seed} {factor}")
        return (f"Strengthened basin L{layer} seed={seed} (factor={factor}x)\n"
                f"  Energy: {e_before:.4f} -> {e_after:.4f} (deepened by {e_before - e_after:.4f})")

    def cmd_weaken(self, layer, seed, factor=None):
        """Weaken a basin."""
        layer, seed = int(layer), int(seed)
        factor = float(factor) if factor else 0.5
        W = self._get_W(layer)
        d = W.shape[0]
        target = self._make_target(seed, d)
        basins, _, _ = find_basins(W, d, num_inits=100, beta=self._get_beta(layer), device=DEVICE)
        if basins:
            sims = [F.cosine_similarity(target.unsqueeze(0), b.unsqueeze(0)).item() for b in basins]
            target = basins[int(np.argmax(sims))]
        e_before = compute_energy(target, W)
        W_new = remove_basin(W, target, strength=self.params['strength'] * factor)
        e_after = compute_energy(target, W_new)
        self._w_cache[layer] = W_new
        if layer not in self._w_backups:
            self._w_backups[layer] = W.clone()
        self._log(f"weaken {layer} {seed} {factor}")
        return (f"Weakened basin L{layer} seed={seed} (factor={factor}x)\n"
                f"  Energy: {e_before:.4f} -> {e_after:.4f} (raised by {e_after - e_before:.4f})")

    def cmd_energy(self, layer, seed):
        """Measure energy at a point."""
        layer, seed = int(layer), int(seed)
        W = self._get_W(layer)
        d = W.shape[0]
        target = self._make_target(seed, d)
        e = compute_energy(target, W)
        is_basin, _, cos, iters = verify_basin_exists(W, target, beta=self._get_beta(layer))
        return (f"Energy at L{layer} seed={seed}: {e:.6f}\n"
                f"  Basin: {is_basin} (cos={cos:.4f}, {iters} iters)")

    def cmd_export_concept(self, concept_name, path):
        """Export concept to file."""
        if concept_name not in self._concepts:
            return f"Unknown concept '{concept_name}'"
        concept = self._concepts[concept_name]
        data = {
            'name': concept_name,
            'centroid': concept['centroid'].cpu(),
            'num_samples': len(concept['states']),
            'states': [s.cpu() for s in concept['states']],
        }
        torch.save(data, path)
        self._log(f"export {concept_name} {path}")
        return f"Exported concept '{concept_name}' ({len(concept['states'])} samples) to {path}"

    def cmd_import_concept(self, path):
        """Import concept from file."""
        data = torch.load(path, map_location=DEVICE, weights_only=False)
        name = data['name']
        self._concepts[name] = {
            'centroid': data['centroid'].to(DEVICE),
            'states': [s.to(DEVICE) for s in data.get('states', [])],
        }
        self._log(f"import {name} {path}")
        return f"Imported concept '{name}' ({len(self._concepts[name]['states'])} samples) from {path}"

    def cmd_probe(self, layer, basin_idx=None):
        """Probe a basin — what text activates it?"""
        layer = int(layer)
        basin_idx = int(basin_idx) if basin_idx else 0
        if self.model is None:
            return "No model loaded. Use: load <checkpoint>"

        W = self._get_W(layer)
        d = W.shape[0]
        basins, _, _ = find_basins(W, d, num_inits=int(self.params['inits']),
                                   beta=self._get_beta(layer), device=DEVICE)
        if basin_idx >= len(basins):
            return f"Basin {basin_idx} not found. Have {len(basins)} basins."

        basin_state = basins[basin_idx]
        energy = compute_energy(basin_state, W)

        tokens = []
        if hasattr(self.model, 'lm_head'):
            with torch.no_grad():
                logits = self.model.lm_head(basin_state.unsqueeze(0))
                top_k = min(10, logits.shape[-1])
                values, indices = torch.topk(logits, top_k, dim=-1)
                probs = F.softmax(values, dim=-1)
                if self.tokenizer:
                    for i in range(top_k):
                        tok = self.tokenizer.decode([indices[0, i].item()])
                        tokens.append(f"  {tok!r}: {probs[0, i].item():.3f}")

        self._log(f"probe {layer} {basin_idx}")
        lines = [f"Probe L{layer} basin #{basin_idx} (energy={energy:.4f})"]
        if tokens:
            lines.append("Top tokens this basin activates:")
            lines.extend(tokens)
        else:
            lines.append("  (no output head available for decoding)")
        return "\n".join(lines)

    def cmd_landscape(self, layer):
        """Map the full energy landscape."""
        layer = int(layer)
        W = self._get_W(layer)
        d = W.shape[0]
        basins, ids, _ = find_basins(W, d, num_inits=int(self.params['inits']),
                                     beta=self._get_beta(layer), device=DEVICE)
        lines = [f"Landscape L{layer}: {len(basins)} basins from {int(self.params['inits'])} inits"]
        for i, b in enumerate(basins):
            e = compute_energy(b, W)
            pop = ids.count(i)
            bar = "#" * min(pop, 40)
            lines.append(f"  basin {i:3d}: E={e:8.4f} pop={pop:3d} {bar}")

        energies = [compute_energy(b, W) for b in basins]
        lines.append(f"\n  Energy range: [{min(energies):.4f}, {max(energies):.4f}]")
        lines.append(f"  Deepest: basin {np.argmin(energies)} (E={min(energies):.4f})")
        lines.append(f"  Most popular: basin {max(range(len(basins)), key=lambda i: ids.count(i))} "
                     f"(pop={max(ids.count(i) for i in range(len(basins)))})")
        self._log(f"landscape {layer}")
        return "\n".join(lines)

    def cmd_guard(self, layer, max_pct=None):
        """Set guard: prevent removing too many basins."""
        layer = int(layer)
        max_pct = float(max_pct) if max_pct else 10.0
        W = self._get_W(layer)
        d = W.shape[0]
        basins, _, _ = find_basins(W, d, num_inits=int(self.params['inits']),
                                   beta=self._get_beta(layer), device=DEVICE)
        threshold = max(1, int(len(basins) * (1 - max_pct / 100)))
        if not hasattr(self, '_guards'):
            self._guards = {}
        self._guards[layer] = {
            'baseline_count': len(basins), 'min_count': threshold, 'max_remove_pct': max_pct,
        }
        self._log(f"guard {layer} {max_pct}%")
        return (f"Guard set on L{layer}: {len(basins)} basins, "
                f"min allowed: {threshold} (max {max_pct:.0f}% removal)")

    def cmd_benchmark(self):
        """Benchmark model quality (perplexity)."""
        if self.model is None or self.tokenizer is None:
            return "No model loaded. Use: load <checkpoint>"
        import math
        texts = [
            "The capital of France is",
            "In the beginning there was",
            "The meaning of life is",
            "Once upon a time in a",
            "The quick brown fox jumps",
        ]
        total_loss, total_tokens = 0.0, 0
        with torch.no_grad():
            for text in texts:
                tokens = self.tokenizer.encode(text)
                if len(tokens) < 2:
                    continue
                input_ids = torch.tensor([tokens], device=DEVICE)
                result = self.model(input_ids)
                logits = result['logits']
                shift_logits = logits[:, :-1, :].contiguous()
                shift_labels = input_ids[:, 1:].contiguous()
                loss_fn = torch.nn.CrossEntropyLoss(reduction='sum')
                loss = loss_fn(shift_logits.view(-1, shift_logits.size(-1)),
                             shift_labels.view(-1))
                total_loss += loss.item()
                total_tokens += shift_labels.numel()
        avg_loss = total_loss / max(total_tokens, 1)
        ppl = math.exp(min(avg_loss, 20))
        self._log(f"benchmark ppl={ppl:.2f}")
        return (f"Benchmark: PPL={ppl:.2f} (loss={avg_loss:.4f})\n"
                f"  {len(texts)} texts, {total_tokens} tokens")

    def cmd_generate(self, prompt):
        if self.model is None or self.tokenizer is None:
            return "No LM model with tokenizer loaded. Use: load <lm-checkpoint>"

        model = self.model
        tokenizer = self.tokenizer
        config = self.config

        input_ids = tokenizer.encode(prompt)
        ids = torch.tensor([input_ids], device=DEVICE)

        max_tokens = int(self.params['tokens'])
        temperature = self.params['temp']
        top_k = int(self.params['topk'])
        generated = []

        t0 = time.time()
        with torch.no_grad():
            for _ in range(max_tokens):
                max_seq = config.get('maxSeqLen', 512)
                inp = ids if ids.shape[1] <= max_seq else ids[:, -max_seq:]
                result = model(inp)
                logits = result['logits'][:, -1, :] / max(temperature, 0.01)
                k = min(top_k, config.get('vocabSize', logits.shape[-1]))
                values, indices = torch.topk(logits, k, dim=-1)
                probs = F.softmax(values, dim=-1)
                sampled = torch.multinomial(probs, 1)
                next_token = indices.gather(-1, sampled)
                token_id = next_token.item()
                generated.append(token_id)
                ids = torch.cat([ids, next_token], dim=1)
                if token_id in (0, 3):
                    break

        elapsed = time.time() - t0
        text = tokenizer.decode(generated)
        tok_s = len(generated) / max(elapsed, 0.01)

        self._log(f"generate {prompt[:30]}...")
        return (f"{prompt}{text}\n"
                f"  [{len(generated)} tokens, {elapsed:.1f}s, {tok_s:.0f} tok/s]")

    def cmd_set(self, param, value):
        if param not in self.params:
            return f"Unknown param: {param}. Available: {list(self.params.keys())}"
        self.params[param] = float(value)
        self._log(f"set {param} {value}")
        return f"{param} = {self.params[param]}"

    def cmd_save(self, path):
        if not self._w_backups and not self._w_cache:
            return "Nothing to save."

        save_dict = {}
        for layer, W in self._w_cache.items():
            save_dict[f'layer_{layer}_W'] = W.cpu()
        if self._w_backups:
            save_dict['backups'] = {
                f'layer_{l}': W.cpu() for l, W in self._w_backups.items()
            }
        save_dict['checkpoint'] = self.checkpoint_path
        save_dict['history'] = self._history

        torch.save(save_dict, path)
        self._log(f"save {path}")
        return f"Saved {len(self._w_cache)} W matrices to {path}"

    def cmd_diff(self, layer):
        layer = int(layer)
        if layer not in self._w_backups:
            return f"Layer {layer} has not been modified."

        W_orig = self._w_backups[layer]
        W_curr = self._get_W(layer)
        diff = W_curr - W_orig

        return (
            f"Diff L{layer}:\n"
            f"  Frobenius norm of delta: {diff.norm().item():.6f}\n"
            f"  Max absolute change:     {diff.abs().max().item():.6f}\n"
            f"  Mean absolute change:    {diff.abs().mean().item():.8f}\n"
            f"  Original spectral norm:  {torch.linalg.norm(W_orig, ord=2).item():.4f}\n"
            f"  Current spectral norm:   {torch.linalg.norm(W_curr, ord=2).item():.4f}\n"
            f"  Relative change:         "
            f"{diff.norm().item() / W_orig.norm().item() * 100:.2f}%")

    def cmd_history(self):
        if not self._history:
            return "No history."
        lines = ["Time     Command"]
        for ts, cmd in self._history:
            lines.append(f"{ts}  {cmd}")
        return "\n".join(lines)

    # ── Dispatcher ───────────────────────────────────────────────────

    def execute(self, line):
        """Execute a single command line. Returns output string."""
        line = line.strip()
        if not line or line.startswith('#'):
            return ""

        parts = line.split(None, 1)
        cmd = parts[0].lower()
        args_str = parts[1] if len(parts) > 1 else ""
        args = args_str.split()

        try:
            if cmd == 'load':
                return self.load_checkpoint(args_str)
            elif cmd == 'info':
                return self.cmd_info()
            elif cmd == 'survey' and args:
                return self.cmd_survey(args[0])
            elif cmd in ('survey-all', 'surveyall'):
                return self.cmd_survey_all()
            elif cmd == 'inject' and len(args) >= 2:
                return self.cmd_inject(args[0], args[1],
                                       args[2] if len(args) > 2 else None)
            elif cmd == 'remove' and len(args) >= 2:
                return self.cmd_remove(args[0], args[1],
                                       args[2] if len(args) > 2 else None)
            elif cmd == 'move' and len(args) >= 2:
                return self.cmd_move(args[0], args[1],
                                     args[2] if len(args) > 2 else None)
            elif cmd == 'verify' and len(args) >= 2:
                return self.cmd_verify(args[0], args[1])
            elif cmd == 'apply' and args:
                return self.cmd_apply(args[0])
            elif cmd == 'restore' and args:
                return self.cmd_restore(args[0])
            elif cmd in ('restore-all', 'restoreall'):
                return self.cmd_restore_all()
            elif cmd == 'status':
                return self.cmd_status()
            elif cmd == 'capture' and len(args) >= 2:
                # capture <layer> <concept> <text...>  OR  capture <layer> <text...>
                layer_str = args[0]
                rest = args_str[len(layer_str):].strip()
                # Check if second arg is a concept name (no spaces) followed by text
                rest_parts = rest.split(None, 1)
                if len(rest_parts) >= 2 and not rest_parts[0].startswith('"'):
                    return self.cmd_capture(layer_str, rest_parts[1], rest_parts[0])
                else:
                    return self.cmd_capture(layer_str, rest)
            elif cmd in ('inject-concept', 'inject_concept') and len(args) >= 2:
                return self.cmd_inject_concept(args[0], args[1],
                                               args[2] if len(args) > 2 else None)
            elif cmd in ('remove-concept', 'remove_concept') and len(args) >= 2:
                return self.cmd_remove_concept(args[0], args[1],
                                               args[2] if len(args) > 2 else None)
            elif cmd == 'concepts':
                return self.cmd_concepts()
            elif cmd == 'blend' and len(args) >= 3:
                return self.cmd_blend(args[0], args[1], args[2],
                                     args[3] if len(args) > 3 else None)
            elif cmd == 'strengthen' and len(args) >= 2:
                return self.cmd_strengthen(args[0], args[1],
                                          args[2] if len(args) > 2 else None)
            elif cmd == 'weaken' and len(args) >= 2:
                return self.cmd_weaken(args[0], args[1],
                                      args[2] if len(args) > 2 else None)
            elif cmd == 'energy' and len(args) >= 2:
                return self.cmd_energy(args[0], args[1])
            elif cmd in ('export-concept', 'export_concept') and len(args) >= 2:
                return self.cmd_export_concept(args[0], args[1])
            elif cmd in ('import-concept', 'import_concept') and args:
                return self.cmd_import_concept(args[0])
            elif cmd == 'probe' and args:
                return self.cmd_probe(args[0], args[1] if len(args) > 1 else None)
            elif cmd == 'landscape' and args:
                return self.cmd_landscape(args[0])
            elif cmd == 'guard' and args:
                return self.cmd_guard(args[0], args[1] if len(args) > 1 else None)
            elif cmd == 'benchmark':
                return self.cmd_benchmark()
            elif cmd in ('generate', 'gen') and args_str:
                return self.cmd_generate(args_str)
            elif cmd == 'set' and len(args) >= 2:
                return self.cmd_set(args[0], args[1])
            elif cmd == 'save' and args:
                return self.cmd_save(args[0])
            elif cmd == 'diff' and args:
                return self.cmd_diff(args[0])
            elif cmd == 'history':
                return self.cmd_history()
            elif cmd == 'help':
                return __doc__
            elif cmd in ('quit', 'exit', 'q'):
                return '__EXIT__'
            else:
                return f"Unknown command: {cmd}. Type 'help' for usage."
        except Exception as e:
            return f"Error: {e}"

    # ── REPL ─────────────────────────────────────────────────────────

    def repl(self):
        """Interactive read-eval-print loop."""
        print("Qriton HLM — Energy Language CLI")
        print(f"Device: {DEVICE}")
        if self.checkpoint_path:
            print(f"Model: {self.checkpoint_path}")
        print("Type 'help' for commands, 'quit' to exit.\n")

        while True:
            try:
                prompt = f"hlm:{self.model_label or 'no-model'}> "
                line = input(prompt)
            except (EOFError, KeyboardInterrupt):
                print()
                break

            result = self.execute(line)
            if result == '__EXIT__':
                break
            if result:
                print(result)

    def run_script(self, script_path):
        """Execute a .hlm script file."""
        with open(script_path) as f:
            for line_no, line in enumerate(f, 1):
                line = line.strip()
                if not line or line.startswith('#'):
                    continue
                print(f"[{line_no}] {line}")
                result = self.execute(line)
                if result == '__EXIT__':
                    break
                if result:
                    print(result)
                print()


def main():
    parser = argparse.ArgumentParser(
        description='Qriton HLM — Energy Language CLI',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  qriton-hlm                                    # interactive REPL
  qriton-hlm --checkpoint output-hlm3-large-ffn/model.pt
  qriton-hlm --script program.hlm               # run script
  echo "survey 0" | qriton-hlm --checkpoint model.pt
        """)
    parser.add_argument('--checkpoint', '-c', type=str,
                        help='Load checkpoint on startup')
    parser.add_argument('--script', '-s', type=str,
                        help='Run a .hlm script file')
    args = parser.parse_args()

    lang = EnergyLang()

    if args.checkpoint:
        result = lang.load_checkpoint(args.checkpoint)
        print(result)
        print()

    if args.script:
        lang.run_script(args.script)
    elif sys.stdin.isatty():
        lang.repl()
    else:
        # Piped input
        for line in sys.stdin:
            result = lang.execute(line)
            if result == '__EXIT__':
                break
            if result:
                print(result)


if __name__ == '__main__':
    main()
