(add-to-load-path (dirname (current-filename)))

(import
    (srfi srfi-41)
    (extra-functional)
    (my-streams))



; this implementation could use some optimisation



(define base (stream-from 7))



; tests potential solutions to 4x^2 + y^2 = n
(define (f1 x y n)
    (eq? (+ (* 4 (expt x 2)) (expt y 2)) n))



; tests potential solutions to 3x^2 + y^2 = n
(define (f2 x y n)
    (eq? (+ (* 3 (expt x 2)) (expt y 2)) n))



; tests potential solutions to 3x^2 - y^2 = n with x > y
(define (f3 x y n)
    (and (> x y) (eq? (- (* 3 (expt x 2)) (expt y 2)) n)))



; counts the number of positive integer solutions to one of the above equations for a given n
(define (num-solutions func n)
    (letrec ((sol (lambda (r x y)
                        (cond ((func x y n) (sol (+ r 1) x (+ y 1)))
                              ((> y (sqrt n)) (sol r (+ x 1) 1))
                              ((> x (sqrt n)) r)
                              (else (sol r x (+ y 1)))))))
        (sol 0 1 1)))



(define (test item)
    (let ((r (modulo (car item) 60)))

        (cond ((or (eq? r 1)
                   (eq? r 13)
                   (eq? r 17)
                   (eq? r 29)
                   (eq? r 37)
                   (eq? r 41)
                   (eq? r 49)
                   (eq? r 53)) (if (odd? (num-solutions f1 (car item)))
                                   (cons (car item) (not (cdr item)))
                                   item))

              ((or (eq? r 7)
                   (eq? r 19)
                   (eq? r 31)
                   (eq? r 43)) (if (odd? (num-solutions f2 (car item)))
                                   (cons (car item) (not (cdr item)))
                                   item))

              ((or (eq? r 11)
                   (eq? r 23)
                   (eq? r 47)
                   (eq? r 59)) (if (odd? (num-solutions f3 (car item)))
                                   (cons (car item) (not (cdr item)))
                                   item))

              (else item))))



(define (square n) (* n n))



(define (mark n item)
    (if (eq? (remainder (car item) n) 0)
        (cons (car item) #f)
        item))



(define-stream (sieve input)
    (stream-let loop ((strm (stream-map (compose test (part (flip cons) #f)) input)))
        (stream-match strm ((item . rest)
            (if (cdr item)
                (stream-cons
                    (car item)
                    (loop (stream-map (part mark (square (car item))) rest)))
                (loop rest))))))



(define atkin (stream-append (list->stream '(2 3 5)) (sieve base)))