OverviewFilesCommitsIssuesDocsReleasesPeers

hypercore-kb / opencog-kb / language-learning / mvp / unified-loop.scm

822 lines · scm ref:
#!/usr/bin/env guile
!#
;;
;; unified-loop.scm -- Single continuous learning loop driven by vagueness
;;
;; Usage: cd mvp && guile -s unified-loop.scm
;;
;; This implements the unified architecture where:
;;   - Vagueness (surprise) drives learning
;;   - Thresholds trigger recomputation, not schedules
;;   - Same loop operates at all levels recursively
;;   - Promotion to next level happens when current level stabilizes
;;

(use-modules (opencog) (opencog exec))
(use-modules (opencog persist) (opencog persist-rocks))
(use-modules (opencog lg))
(use-modules (srfi srfi-1))
(use-modules (srfi srfi-13))    ; for string-contains, string-join
(use-modules (ice-9 threads))
(use-modules (ice-9 rdelim))    ; for read-line

;; ═══════════════════════════════════════════════════════════════════
;; CONFIGURATION
;; ═══════════════════════════════════════════════════════════════════

(define *rocks-path* "rocks://data/pairs.rdb")

;; Vagueness thresholds
(define *surprise-threshold* 2.0)      ; Bits - learn if surprise > this
(define *mi-delta-threshold* 0.1)      ; Recompute MI if delta > this
(define *cluster-delta-threshold* 0.1) ; Recluster if delta > this
(define *promotion-threshold* 0.01)    ; Promote if merge-rate < this

;; Window sizes for rolling metrics
(define *surprise-window* 100)         ; Recent sentences for avg surprise
(define *stability-window* 50)         ; Recent merges for stability check

;; ═══════════════════════════════════════════════════════════════════
;; STATE
;; ═══════════════════════════════════════════════════════════════════

(define *rsn* #f)                      ; RocksDB connection
(define *current-level* 0)             ; Current learning level
(define *running* #t)                  ; Loop control

;; Rolling metrics
(define *recent-surprises* '())        ; List of recent surprise values
(define *recent-merges* '())           ; List of recent merge counts
(define *total-sentences* 0)
(define *total-learns* 0)

;; Level-specific state
(define *pair-counts* (make-hash-table))   ; Level 0: word pairs
(define *pair-order* (make-hash-table))    ; Track word order: key -> (left-count . right-count)
(define *section-counts* (make-hash-table)) ; Level 1: sections
(define *mi-cache* (make-hash-table))      ; MI values
(define *word-classes* (make-hash-table))  ; Clusters
(define *word-to-class* (make-hash-table))

;; Hierarchy tracking for visualization
(define *merge-history* '())               ; List of merges: (level entity1 entity2 new-id iteration)
(define *cluster-iteration* 0)             ; Counter for merge ordering

;; Baseline MI for surprise computation
(define *baseline-mi* 0.0)
(define *mi-sum* 0.0)
(define *mi-count* 0)

;; ═══════════════════════════════════════════════════════════════════
;; CORE: THE UNIFIED LOOP
;; ═══════════════════════════════════════════════════════════════════

(define (unified-loop sentences)
  "The single continuous learning loop."
  (format #t "~%═══ UNIFIED LEARNING LOOP ═══~%")
  (format #t "Surprise threshold: ~a bits~%" *surprise-threshold*)
  (format #t "Starting at Level ~a~%~%" *current-level*)

  (let loop ((remaining sentences))
    (when (and *running* (not (null? remaining)))
      (let* ((sentence (car remaining))

             ;; 1. PARSE with current knowledge
             (parse-result (parse-sentence-safe sentence))
             (edges (car parse-result))
             (words (cadr parse-result))

             ;; 2. MEASURE vagueness (surprise)
             (surprise (compute-surprise edges)))

        ;; Track metrics
        (set! *total-sentences* (+ *total-sentences* 1))
        (update-surprise-window! surprise)

        ;; 3. DECIDE: learn or accept
        (if (> surprise *surprise-threshold*)
            (begin
              ;; HIGH VAGUENESS → LEARN
              (set! *total-learns* (+ *total-learns* 1))
              (learn-from-sentence! words edges))
            ;; LOW VAGUENESS → just count (passive learning)
            (count-pairs-only! words))

        ;; 4. CHECK stability → maybe promote
        (when (and (> *total-sentences* *stability-window*)
                   (stable-enough?))
          (promote-level!))

        ;; 5. Periodic status
        (when (zero? (modulo *total-sentences* 100))
          (show-loop-status))

        ;; Continue
        (loop (cdr remaining)))))

  (format #t "~%═══ LOOP COMPLETE ═══~%")
  (show-final-stats))

;; ═══════════════════════════════════════════════════════════════════
;; PARSING
;; ═══════════════════════════════════════════════════════════════════

(define (parse-sentence-safe sentence)
  "Parse sentence, return (edges words) or empty on failure."
  (catch #t
    (lambda ()
      (let* ((words (string-split sentence #\space))
             (edges (parse-sentence-mst words)))
        (list edges words)))
    (lambda (key . args)
      (list '() '()))))

(define (parse-sentence-mst words)
  "MST parse using MI weights. Returns list of (pos1 pos2 mi)."
  (let* ((n (length words))
         (word-vec (list->vector words)))
    (if (< n 2)
        '()
        ;; Build all pairs with MI, sort, greedy non-crossing selection
        (let* ((pairs
                (let loop-i ((i 0) (pairs '()))
                  (if (>= i (- n 1))
                      pairs
                      (loop-i (+ i 1)
                              (let loop-j ((j (+ i 1)) (ps pairs))
                                (if (>= j n)
                                    ps
                                    (loop-j (+ j 1)
                                            (cons (list (get-mi (vector-ref word-vec i)
                                                                (vector-ref word-vec j))
                                                        i j)
                                                  ps))))))))
               (sorted (sort pairs (lambda (a b) (> (car a) (car b))))))
          ;; Greedy non-crossing selection
          (let select ((remaining sorted) (edges '()))
            (if (null? remaining)
                (reverse edges)
                (let* ((pair (car remaining))
                       (mi (car pair))
                       (p1 (cadr pair))
                       (p2 (caddr pair))
                       (crosses? (any (lambda (e)
                                        (edges-cross? p1 p2 (car e) (cadr e)))
                                      edges)))
                  (select (cdr remaining)
                          (if crosses?
                              edges
                              (cons (list p1 p2 mi) edges))))))))))

(define (edges-cross? i j k l)
  "Check if edge (i,j) crosses edge (k,l)."
  (let ((a (min i j)) (b (max i j))
        (c (min k l)) (d (max k l)))
    (or (and (< a c) (< c b) (< b d))
        (and (< c a) (< a d) (< d b)))))

(define (get-mi w1 w2)
  "Get MI for word pair from cache."
  (or (hash-ref *mi-cache* (cons w1 w2))
      (hash-ref *mi-cache* (cons w2 w1))
      0.0))

;; ═══════════════════════════════════════════════════════════════════
;; VAGUENESS / SURPRISE
;; ═══════════════════════════════════════════════════════════════════

(define (compute-surprise edges)
  "Compute surprise in bits. High MI = expected = low surprise."
  (if (null? edges)
      +inf.0  ; Unparseable = maximum surprise
      (let* ((mi-values (map caddr edges))
             (avg-mi (/ (apply + mi-values) (length mi-values))))
        ;; If no MI data (all zeros), that's maximum surprise
        (if (zero? avg-mi)
            10.0  ; High surprise - we know nothing
            ;; Otherwise, surprise = how much worse than baseline
            (max 0.0 (- *baseline-mi* avg-mi))))))

(define (update-surprise-window! surprise)
  "Add surprise to rolling window, maintain size."
  (set! *recent-surprises* (cons surprise *recent-surprises*))
  (when (> (length *recent-surprises*) *surprise-window*)
    (set! *recent-surprises* (take *recent-surprises* *surprise-window*))))

(define (average-surprise)
  "Get average surprise over recent window."
  (if (null? *recent-surprises*)
      0.0
      (/ (apply + *recent-surprises*) (length *recent-surprises*))))

;; ═══════════════════════════════════════════════════════════════════
;; LEARNING (when vagueness is high)
;; ═══════════════════════════════════════════════════════════════════

(define (learn-from-sentence! words edges)
  "Update model based on surprising input."

  ;; 1. COUNT: Update pair counts
  (count-pairs! words)

  ;; 2. Extract sections from parse
  (extract-sections! words edges)

  ;; 3. Check if MI needs recomputing
  (when (> (mi-delta) *mi-delta-threshold*)
    (recompute-mi!))

  ;; 4. Check if clustering needed
  (when (> (cluster-delta) *cluster-delta-threshold*)
    (recluster!)))

(define (count-pairs! words)
  "Count word pairs in sliding window, tracking word order."
  (let ((n (length words))
        (window 6))
    (do ((i 0 (+ i 1))) ((>= i n))
      (do ((j (+ i 1) (+ j 1))) ((or (>= j n) (> (- j i) window)))
        (let* ((left-word (list-ref words i))   ; Word that came first in sentence
               (right-word (list-ref words j))  ; Word that came second
               ;; Canonical key (alphabetically sorted)
               (key (if (string<? left-word right-word)
                        (cons left-word right-word)
                        (cons right-word left-word)))
               ;; Did the canonical-first word come first in the sentence?
               (canonical-first-came-first (string<? left-word right-word)))
          ;; Count the pair
          (hash-set! *pair-counts* key
                     (+ (hash-ref *pair-counts* key 0) 1))
          ;; Track word order: (forward-count . backward-count)
          ;; forward = canonical-first word came first in sentence
          ;; backward = canonical-second word came first in sentence
          (let ((order (hash-ref *pair-order* key (cons 0 0))))
            (if canonical-first-came-first
                (hash-set! *pair-order* key (cons (+ (car order) 1) (cdr order)))
                (hash-set! *pair-order* key (cons (car order) (+ (cdr order) 1))))))))))

(define (count-pairs-only! words)
  "Just count pairs, no other learning (passive)."
  (count-pairs! words))

(define (extract-sections! words edges)
  "Extract sections (word + disjunct pattern) from parse."
  (let ((word-vec (list->vector words)))
    (do ((i 0 (+ i 1))) ((>= i (length words)))
      (let ((word (vector-ref word-vec i))
            (connectors '()))
        ;; Find all connections for this word
        (for-each
          (lambda (edge)
            (let ((p1 (car edge)) (p2 (cadr edge)))
              (cond
                ((= p1 i)
                 (set! connectors
                       (cons (cons (vector-ref word-vec p2) "+") connectors)))
                ((= p2 i)
                 (set! connectors
                       (cons (cons (vector-ref word-vec p1) "-") connectors))))))
          edges)
        ;; Count this section
        (unless (null? connectors)
          (let ((section-key (cons word (sort connectors
                                              (lambda (a b)
                                                (string<? (car a) (car b)))))))
            (hash-set! *section-counts* section-key
                       (+ (hash-ref *section-counts* section-key 0) 1))))))))

;; ═══════════════════════════════════════════════════════════════════
;; MI COMPUTATION
;; ═══════════════════════════════════════════════════════════════════

(define (mi-delta)
  "Estimate how much MI has drifted (simplified)."
  ;; Bootstrap: if MI cache is empty but we have pairs, definitely recompute
  (let ((pair-count (hash-count (const #t) *pair-counts*))
        (mi-cached (hash-count (const #t) *mi-cache*)))
    (cond
      ;; Bootstrap: no MI yet but we have data
      ((and (zero? mi-cached) (> pair-count 20))
       1.0)
      ;; Regular: recompute when significant new data
      ((> pair-count (* 1.5 (max 1 mi-cached)))
       0.2)
      (else 0.0))))

(define (recompute-mi!)
  "Recompute MI for all pairs."
  (format #t "  [Recomputing MI...]~%")
  (let ((total-count 0)
        (left-marginals (make-hash-table))
        (right-marginals (make-hash-table)))

    ;; Compute marginals
    (hash-for-each
      (lambda (key count)
        (let ((w1 (car key)) (w2 (cdr key)))
          (hash-set! left-marginals w1
                     (+ (hash-ref left-marginals w1 0) count))
          (hash-set! right-marginals w2
                     (+ (hash-ref right-marginals w2 0) count))
          (set! total-count (+ total-count count))))
      *pair-counts*)

    ;; Compute MI
    (set! *mi-sum* 0.0)
    (set! *mi-count* 0)
    (hash-for-each
      (lambda (key count)
        (let* ((w1 (car key))
               (w2 (cdr key))
               (left-m (hash-ref left-marginals w1 1))
               (right-m (hash-ref right-marginals w2 1))
               (n (exact->inexact total-count))
               (p-pair (/ count n))
               (p-left (/ left-m n))
               (p-right (/ right-m n))
               (mi (if (and (> p-pair 0) (> p-left 0) (> p-right 0))
                       (/ (log (/ p-pair (* p-left p-right))) (log 2))
                       0.0)))
          (hash-set! *mi-cache* key mi)
          (hash-set! *mi-cache* (cons w2 w1) mi)  ; Symmetric
          (set! *mi-sum* (+ *mi-sum* mi))
          (set! *mi-count* (+ *mi-count* 1))))
      *pair-counts*)

    ;; Update baseline
    (when (> *mi-count* 0)
      (set! *baseline-mi* (/ *mi-sum* *mi-count*)))

    (format #t "  [MI recomputed: ~a pairs, baseline=~,2f bits]~%"
            *mi-count* *baseline-mi*)))

;; ═══════════════════════════════════════════════════════════════════
;; CLUSTERING
;; ═══════════════════════════════════════════════════════════════════

(define (cluster-delta)
  "Estimate how much cluster profiles have drifted."
  ;; Trigger clustering when we have meaningful MI data
  (let ((mi-count (hash-count (const #t) *mi-cache*))
        (section-count (hash-count (const #t) *section-counts*)))
    (cond
      ;; Bootstrap: first clustering when we have MI data
      ((and (> mi-count 20)
            (zero? (hash-count (const #t) *word-classes*)))
       1.0)
      ;; Regular: recluster when sections grow significantly
      ((> section-count (* 1.5 (max 50 (hash-count (const #t) *word-classes*))))
       0.2)
      (else 0.0))))

(define (recluster!)
  "Recluster words based on MI profile similarity."
  (format #t "  [Reclustering...]~%")

  ;; Build word list from MI cache
  (let ((words (make-hash-table))
        (merges 0))

    ;; Collect unique words
    (hash-for-each
      (lambda (key mi)
        (hash-set! words (car key) #t)
        (hash-set! words (cdr key) #t))
      *mi-cache*)

    (let ((word-list (hash-map->list (lambda (k v) k) words)))
      (format #t "  [~a words with MI data]~%" (length word-list))

      ;; Compute similarities between unassigned words
      (let find-similar ((remaining word-list))
        (unless (or (null? remaining) (null? (cdr remaining)))
          (let* ((w1 (car remaining))
                 (c1 (hash-ref *word-to-class* w1 #f)))
            (unless c1  ; Only unassigned words
              ;; Find most similar unassigned word
              (let ((best-sim 0.0)
                    (best-word #f))
                (for-each
                  (lambda (w2)
                    (unless (or (string=? w1 w2)
                                (hash-ref *word-to-class* w2 #f))
                      (let ((sim (word-profile-similarity w1 w2)))
                        (when (> sim best-sim)
                          (set! best-sim sim)
                          (set! best-word w2)))))
                  (cdr remaining))

                ;; Merge if similarity exceeds threshold
                (when (and best-word (> best-sim 0.2))
                  (let ((new-class (hash-count (const #t) *word-classes*)))
                    (hash-set! *word-classes* new-class (list w1 best-word))
                    (hash-set! *word-to-class* w1 new-class)
                    (hash-set! *word-to-class* best-word new-class)
                    (set! merges (+ merges 1))
                    ;; Record merge in history for hierarchy visualization
                    (set! *merge-history*
                          (cons (list *current-level* w1 best-word new-class
                                      *cluster-iteration* best-sim)
                                *merge-history*))
                    (format #t "  [Merged ~s + ~s (sim=~,3f)]~%"
                            w1 best-word best-sim))))))
          (find-similar (cdr remaining)))))

    ;; Increment iteration counter
    (set! *cluster-iteration* (+ *cluster-iteration* 1))

    ;; Track merge rate
    (set! *recent-merges* (cons merges *recent-merges*))
    (when (> (length *recent-merges*) *stability-window*)
      (set! *recent-merges* (take *recent-merges* *stability-window*)))

    (format #t "  [Clustered: ~a merges, ~a classes]~%"
            merges (hash-count (const #t) *word-classes*))))

(define (word-profile-similarity w1 w2)
  "Compute Jaccard-like similarity between MI profiles of two words."
  (let ((neighbors1 (get-word-neighbors w1))
        (neighbors2 (get-word-neighbors w2)))
    (if (or (null? neighbors1) (null? neighbors2))
        0.0
        (let* ((set1 (list->set neighbors1))
               (set2 (list->set neighbors2))
               (shared (length (lset-intersection equal? neighbors1 neighbors2)))
               (total (length (lset-union equal? neighbors1 neighbors2))))
          (if (zero? total)
              0.0
              (/ shared total))))))

(define (get-word-neighbors word)
  "Get list of words that have MI > 0 with WORD."
  (let ((neighbors '()))
    (hash-for-each
      (lambda (key mi)
        (when (> mi 0)
          (cond
            ((string=? (car key) word)
             (set! neighbors (cons (cdr key) neighbors)))
            ((string=? (cdr key) word)
             (set! neighbors (cons (car key) neighbors))))))
      *mi-cache*)
    neighbors))

(define (list->set lst)
  "Convert list to set (unique elements)."
  (let loop ((lst lst) (result '()))
    (if (null? lst)
        result
        (loop (cdr lst)
              (if (member (car lst) result)
                  result
                  (cons (car lst) result))))))

;; ═══════════════════════════════════════════════════════════════════
;; LEVEL PROMOTION
;; ═══════════════════════════════════════════════════════════════════

(define (stable-enough?)
  "Check if current level has stabilized (low merge rate)."
  ;; Need minimum history AND low average merge rate
  (let ((min-history 10))  ; Require at least 10 clustering rounds
    (if (< (length *recent-merges*) min-history)
        #f  ; Not enough history yet
        (let ((avg-merges (/ (apply + *recent-merges*)
                             (length *recent-merges*))))
          (< avg-merges *promotion-threshold*)))))

(define (promote-level!)
  "Promote to next learning level."
  (format #t "~%═══ PROMOTING TO LEVEL ~a ═══~%" (+ *current-level* 1))
  (format #t "Clusters at Level ~a: ~a~%"
          *current-level* (hash-count (const #t) *word-classes*))

  ;; Show current clusters
  (hash-for-each
    (lambda (id words)
      (format #t "  C~a: ~a~%" id (take words (min 5 (length words)))))
    *word-classes*)

  ;; Advance level
  (set! *current-level* (+ *current-level* 1))

  ;; Reset stability metrics
  (set! *recent-merges* '())

  ;; At Level 1+, we'd now count Section pairs instead of word pairs
  ;; The sections become the new "vocabulary"
  (format #t "~%Now learning at Level ~a (Section pairs)~%~%" *current-level*))

;; ═══════════════════════════════════════════════════════════════════
;; STATUS / REPORTING
;; ═══════════════════════════════════════════════════════════════════

(define (show-loop-status)
  "Show current loop status."
  (format #t "~%--- Status after ~a sentences ---~%" *total-sentences*)
  (format #t "Level: ~a | Learns: ~a (~,1f%)~%"
          *current-level*
          *total-learns*
          (* 100.0 (/ *total-learns* (max 1 *total-sentences*))))
  (format #t "Avg surprise: ~,2f bits | Baseline MI: ~,2f~%"
          (average-surprise) *baseline-mi*)
  (format #t "Pairs: ~a | Sections: ~a | Classes: ~a~%"
          (hash-count (const #t) *pair-counts*)
          (hash-count (const #t) *section-counts*)
          (hash-count (const #t) *word-classes*)))

(define (show-final-stats)
  "Show final statistics."
  (format #t "~%═══ FINAL STATISTICS ═══~%")
  (format #t "Total sentences: ~a~%" *total-sentences*)
  (format #t "Learning triggers: ~a (~,1f%)~%"
          *total-learns*
          (* 100.0 (/ *total-learns* (max 1 *total-sentences*))))
  (format #t "Final level: ~a~%" *current-level*)
  (format #t "Word pairs: ~a~%" (hash-count (const #t) *pair-counts*))
  (format #t "Sections: ~a~%" (hash-count (const #t) *section-counts*))
  (format #t "Word classes: ~a~%" (hash-count (const #t) *word-classes*))
  (format #t "Average surprise: ~,2f bits~%" (average-surprise))

  ;; Show clusters
  (format #t "~%Word Classes:~%")
  (hash-for-each
    (lambda (id words)
      (format #t "  C~a: ~a~%" id words))
    *word-classes*))

;; ═══════════════════════════════════════════════════════════════════
;; CORPUS LOADING
;; ═══════════════════════════════════════════════════════════════════

(define *corpus-dir* "corpus")

(define (load-corpus-file filename)
  "Load sentences from a corpus file, filtering headers/footers."
  (let ((sentences '())
        (in-content #f))
    (call-with-input-file filename
      (lambda (port)
        (let loop ()
          (let ((line (read-line port)))
            (unless (eof-object? line)
              ;; Skip Gutenberg header (before "*** START")
              (when (string-contains line "*** START")
                (set! in-content #t))
              ;; Stop at footer
              (when (string-contains line "*** END")
                (set! in-content #f))
              ;; Process content lines
              (when in-content
                (let ((cleaned (clean-sentence line)))
                  (when (and cleaned (> (string-length cleaned) 10))
                    (set! sentences (cons cleaned sentences)))))
              (loop))))))
    (reverse sentences)))

(define (clean-sentence line)
  "Clean a line: lowercase, remove punctuation, collapse whitespace."
  (let* ((lower (string-downcase line))
         ;; Remove carriage returns and common punctuation, keep spaces
         (no-punct (string-map
                     (lambda (c)
                       (cond
                         ((char=? c #\return) #\space)  ; Remove \r
                         ((or (char-alphabetic? c) (char-whitespace? c)) c)
                         (else #\space)))
                     lower))
         ;; Collapse multiple spaces
         (words (string-split no-punct #\space))
         (non-empty (filter (lambda (w) (> (string-length w) 0)) words)))
    (if (null? non-empty)
        #f
        (string-join non-empty " "))))

(define (load-all-corpus)
  "Load all corpus files from corpus directory."
  (let ((sentences '()))
    (for-each
      (lambda (file)
        (let ((path (string-append *corpus-dir* "/" file)))
          (format #t "Loading ~a...~%" path)
          (set! sentences (append sentences (load-corpus-file path)))))
      '("alice-in-wonderland.txt"
        "pride-and-prejudice.txt"
        "sherlock-holmes.txt"))
    sentences))

(define (run-corpus n)
  "Run unified loop on first N sentences from corpus."
  (format #t "~%═══ UNIFIED LOOP ON REAL CORPUS ═══~%")
  (let ((sentences (load-all-corpus)))
    (format #t "Loaded ~a sentences total~%" (length sentences))
    (let ((subset (take sentences (min n (length sentences)))))
      (format #t "Running on ~a sentences~%~%" (length subset))
      ;; Initialize
      (set! *baseline-mi* 1.0)
      ;; Run
      (unified-loop subset))))

;; ═══════════════════════════════════════════════════════════════════
;; DEMO
;; ═══════════════════════════════════════════════════════════════════

(define *demo-sentences*
  '("the dog chased the cat"
    "a cat saw the dog"
    "the big dog ran fast"
    "she walked to the store"
    "he read a good book"
    "the store sells books"
    "i like the new book"
    "the cat is small"
    "a big cat ran away"
    "the dog is happy"
    "the man walked home"
    "she saw a big house"
    "the house is old"
    "he likes the old car"
    "a car drove past"
    "the fast car stopped"
    "she read the book"
    "the book was good"
    "a good dog sat down"
    "the cat chased the dog"))

(define (run-demo)
  "Run the unified loop on demo sentences."
  (format #t "~%═══ UNIFIED LOOP DEMO ═══~%")
  (format #t "Running on ~a demo sentences~%~%" (length *demo-sentences*))

  ;; Initialize with some baseline MI
  (set! *baseline-mi* 1.0)

  ;; Run the loop
  (unified-loop *demo-sentences*))

;; ═══════════════════════════════════════════════════════════════════
;; MAIN
;; ═══════════════════════════════════════════════════════════════════

;; ═══════════════════════════════════════════════════════════════════
;; JSON EXPORT FOR VISUALIZATION
;; ═══════════════════════════════════════════════════════════════════

(define (export-viz-json filename)
  "Export current state as JSON for visualization."
  (call-with-output-file filename
    (lambda (port)
      (format port "{~%")

      ;; Words with surprise, cluster, count, neighbors, degree
      (format port "  \"words\": {~%")
      (let ((word-list '())
            (word-degrees (make-hash-table))
            (word-total-mi (make-hash-table)))

        ;; Collect words and compute degrees from MI cache
        (hash-for-each
          (lambda (key mi)
            (let ((w1 (car key)) (w2 (cdr key)))
              (unless (member w1 word-list) (set! word-list (cons w1 word-list)))
              (unless (member w2 word-list) (set! word-list (cons w2 word-list)))
              ;; Count degrees
              (hash-set! word-degrees w1 (+ (hash-ref word-degrees w1 0) 1))
              (hash-set! word-degrees w2 (+ (hash-ref word-degrees w2 0) 1))
              ;; Sum MI
              (hash-set! word-total-mi w1 (+ (hash-ref word-total-mi w1 0.0) mi))
              (hash-set! word-total-mi w2 (+ (hash-ref word-total-mi w2 0.0) mi))))
          *mi-cache*)

        ;; Output each word
        (let loop ((words word-list) (first #t))
          (unless (null? words)
            (let* ((word (car words))
                   (cluster (hash-ref *word-to-class* word -1))
                   (neighbors (get-word-neighbors word))
                   (degree (hash-ref word-degrees word 0))
                   (total-mi (hash-ref word-total-mi word 0.0))
                   (count (get-word-count word)))
              (unless first (format port ",~%"))
              (format port "    ~s: {~%" word)
              (format port "      \"surprise\": ~,3f,~%" (compute-word-surprise word))
              (format port "      \"cluster\": ~a,~%" cluster)
              (format port "      \"count\": ~a,~%" count)
              (format port "      \"degree\": ~a,~%" degree)
              (format port "      \"totalMI\": ~,3f,~%" total-mi)
              (format port "      \"neighbors\": [~a]~%"
                      (string-join (map (lambda (n) (format #f "~s" n))
                                        (take neighbors (min 10 (length neighbors))))
                                   ", "))
              (format port "    }"))
            (loop (cdr words) #f))))
      (format port "~%  },~%")

      ;; Pairs with MI and count
      (format port "  \"pairs\": [~%")
      (let ((first #t))
        (hash-for-each
          (lambda (key mi)
            (when (> mi 0.3)  ; Include more pairs
              (let ((count (hash-ref *pair-counts* key 0)))
                (unless first (format port ",~%"))
                (format port "    [~s, ~s, ~,3f, ~a]" (car key) (cdr key) mi count)
                (set! first #f))))
          *mi-cache*))
      (format port "~%  ],~%")

      ;; Cluster details
      (format port "  \"clusterDetails\": {~%")
      (let ((first #t))
        (hash-for-each
          (lambda (id words)
            (unless first (format port ",~%"))
            (format port "    \"~a\": [~a]" id
                    (string-join (map (lambda (w) (format #f "~s" w)) words) ", "))
            (set! first #f))
          *word-classes*))
      (format port "~%  },~%")

      ;; Statistics
      (format port "  \"stats\": {~%")
      (format port "    \"totalWords\": ~a,~%" (hash-count (const #t) *pair-counts*))
      (format port "    \"totalPairs\": ~a,~%" (hash-count (const #t) *mi-cache*))
      (format port "    \"totalSections\": ~a,~%" (hash-count (const #t) *section-counts*))
      (format port "    \"totalClusters\": ~a,~%" (hash-count (const #t) *word-classes*))
      (format port "    \"level\": ~a,~%" *current-level*)
      (format port "    \"avgSurprise\": ~,3f,~%" (average-surprise))
      (format port "    \"baselineMI\": ~,3f,~%" *baseline-mi*)
      (format port "    \"totalSentences\": ~a,~%" *total-sentences*)
      (format port "    \"totalLearns\": ~a,~%" *total-learns*)
      (format port "    \"clusterIterations\": ~a~%" *cluster-iteration*)
      (format port "  },~%")

      ;; Hierarchy data for recursive visualization
      (format port "  \"hierarchy\": {~%")

      ;; Merge history - each merge creates a parent-child relationship
      (format port "    \"merges\": [~%")
      (let ((first #t))
        (for-each
          (lambda (merge)
            ;; merge = (level entity1 entity2 new-class-id iteration similarity)
            (unless first (format port ",~%"))
            (format port "      {\"level\": ~a, \"entity1\": ~s, \"entity2\": ~s, \"clusterId\": ~a, \"iteration\": ~a, \"similarity\": ~,3f}"
                    (list-ref merge 0)
                    (list-ref merge 1)
                    (list-ref merge 2)
                    (list-ref merge 3)
                    (list-ref merge 4)
                    (list-ref merge 5))
            (set! first #f))
          (reverse *merge-history*)))  ; Oldest first
      (format port "~%    ],~%")

      ;; Level summary
      (format port "    \"levels\": [~%")
      (format port "      {\"id\": 0, \"name\": \"Words\", \"count\": ~a},~%"
              (hash-count (const #t) *pair-counts*))
      (format port "      {\"id\": 1, \"name\": \"Clusters\", \"count\": ~a}~%"
              (hash-count (const #t) *word-classes*))
      ;; Add more levels as they emerge
      (format port "    ],~%")

      ;; Current level being learned
      (format port "    \"currentLevel\": ~a~%" *current-level*)
      (format port "  }~%")

      (format port "}~%")))
  (format #t "Exported visualization data to ~a~%" filename))

(define (get-word-count word)
  "Get total count for a word across all pairs."
  (let ((total 0))
    (hash-for-each
      (lambda (key count)
        (when (or (string=? (car key) word) (string=? (cdr key) word))
          (set! total (+ total count))))
      *pair-counts*)
    (max 1 total)))

(define (compute-word-surprise word)
  "Estimate surprise for a single word based on its MI profile."
  (let ((neighbors (get-word-neighbors word)))
    (if (null? neighbors)
        10.0  ; Max surprise if unknown
        (let ((mi-sum 0.0))
          (for-each
            (lambda (n)
              (set! mi-sum (+ mi-sum (or (hash-ref *mi-cache* (cons word n))
                                         (hash-ref *mi-cache* (cons n word))
                                         0.0))))
            neighbors)
          (max 0.0 (- *baseline-mi* (/ mi-sum (length neighbors))))))))

(format #t "~%Unified Loop loaded.~%")
(format #t "Run (run-demo) to test with demo sentences.~%")
(format #t "Or call (unified-loop sentences) with your own list.~%")
(format #t "Use (export-viz-json \"viz/data.json\") to export for visualization.~%~%")

;; Uncomment to run automatically:
;; (run-demo)