#lang racket

(require srfi/41)

(define (integers-starting-from n)
  (stream-cons n
               (integers-starting-from (+ n 1))))
(define integers (integers-starting-from 1))

(define (merge-weighted s1 s2 weight)
  (cond ((stream-empty? s1) s2)
        ((stream-empty? s2) s1)
        (else (let ((x1 (stream-car s1))
                    (x2 (stream-car s2)))
                (if (<= (weight x1) (weight x2))
                    (stream-cons x1 (merge-weighted (stream-cdr s1) s2 weight))
                    (stream-cons x2 (merge-weighted s1 (stream-cdr s2) weight)))))))

(define (weighted-pairs s t weight)
  (stream-cons
   (cons (stream-car s) (stream-car t))
   (merge-weighted
    (stream-map (lambda (x) (cons (stream-car s) x))
                (stream-cdr t))
    (weighted-pairs (stream-cdr s) (stream-cdr t) weight) weight)))

(define (sum-of-cube x)
  (let ((i (car x)) (j (cdr x)))
    (+ (* i i i) (* j j j))))

(define (sum-of-square x)
  (let ((i (car x)) (j (cdr x)))
    (+ (* i i) (* j j))))

(define sorted-pairs
  (weighted-pairs integers integers
                  sum-of-cube))

(define ramanujan-numbers
  (stream-map (lambda (x) (sum-of-cube (cdr x)))
              (stream-filter (lambda (x) (= (sum-of-cube (car x)) (sum-of-cube (cdr x))))
                             (stream-map cons sorted-pairs (stream-cdr sorted-pairs)))))

(define sorted-pairs-2
  (weighted-pairs integers integers
                  sum-of-square))

(define good-numbers
  (stream-map (lambda (x) (cons (sum-of-square (car x)) x))
              (stream-filter (lambda (x) (= (sum-of-square (car x))
                                            (sum-of-square (cadr x))
                                            (sum-of-square (caddr x))))
                             (stream-map list
                                         sorted-pairs-2
                                         (stream-cdr sorted-pairs-2)
                                         (stream-cdr (stream-cdr sorted-pairs-2))))))