#!/usr/bin/ruby

#
## Enumerator() tests
#

var e1 = Enumerator({ |cb|
    0..100 -> each { |i|
        i.is_prime && cb(i)
    }
})

var p = (^Inf -> lazy.grep{.is_prime}.first(25))

assert_eq(e1.len, 25)
assert_eq(e1.to_a, e1.first(25) )
assert_eq(e1.to_a, p)
assert_eq(e1.grep{.is_even}, [2])
assert_eq(e1.map {|x| x-1 }, p.map{|x| x-1 })

var e2 = Enumerator({ |cb|
    ^Inf -> each { |i|
        i.is_prime && cb(i)
    }
})

assert_eq(e2.first(25), p)

#
## function test
#

func primes_less_than(n) {
    Enumerator({ |cb|
        for i in range(0, n) {
            cb(i) if i.is_prime
        }
    })
}

var e3 = primes_less_than(100)

assert_eq(e3.to_a, p)
assert_eq(e3.len, p.len)

var e4 = primes_less_than(50)

assert_eq(e4.first(5), [2,3,5,7,11])
assert_eq(e4.len, 15)

#
## Fibonacci
#

var fib = Enumerator({ |yield|
  var (a, b) = (0, 1)
  loop {
    yield(a)
    (a, b) = (b, a + b)
  }
})

assert_eq(fib.first(10), [0, 1, 1, 2, 3, 5, 8, 13, 21, 34])

var palindromes = Enumerator({|f|
    var n = 1
    loop { f(n); n.next_palindrome! }
})

assert_eq(
    palindromes.while { _ < 100 },
    1..100 -> grep { .is_palindrome }
)

do {
    var e = Enumerator({|f| ^Inf -> each {|n| f(n**2 + 1) } })

    assert_eq(
        e.first(10),
        ^10 -> map { _*_ + 1 }
    )

    assert_eq(
        e.first(10, {.is_prime}),
        ^Inf -> lazy.map { _*_ + 1 }.grep{.is_prime}.first(10)
    )
}

assert_eq(
    Math.seq(1, 1, 1, {|a| a.last(3).sum }).nth(10, { .is_prime }),
    3311233
)

assert_eq(
    Math.seq(1, {|a| a.last(3).sum }).first(20),
    %n[1, 1, 2, 4, 7, 13, 24, 44, 81, 149, 274, 504, 927, 1705, 3136, 5768, 10609, 19513, 35890, 66012]
)

assert_eq(
    Math.seq(1, {|a| a.last(3).sum }).nth(10),
    149
)

assert_eq(
    Math.seq(1, 1, 1, {|a| a.last(3).sum }).while { _ < 1e4 },
    %n[1, 1, 1, 3, 5, 9, 17, 31, 57, 105, 193, 355, 653, 1201, 2209, 4063, 7473]
)

assert_eq(
    Math.seq(1, 1, {|a| a.last(3).sum }).first(16),
    %n[1, 1, 2, 4, 7, 13, 24, 44, 81, 149, 274, 504, 927, 1705, 3136, 5768]
)

assert_eq(
    Math.seq({|a| a.last.next_almost_prime(10) }).while { _ < 1e4 },
    10.almost_primes(1e4)
)

assert_eq(
    Math.seq({|a| a.last.next_omega_prime(10) }).while { _ < 1e10 },
    10.omega_primes(1e10)
)

assert_eq(
    Math.seq({|a| a.last.next_omega_prime(7) }).while { _ < 1e6 },
    7.omega_primes(1e6)
)

assert_eq(
    Math.seq({|a| a.last.next_squarefree_almost_prime(10) }).while { _ < 1e10 },
    10.squarefree_almost_primes(1e10)
)

assert_eq(
    Math.seq({|a| a.last.next_squarefree_almost_prime(7) }).while { _ < 1e6 },
    7.squarefree_almost_primes(1e6)
)

assert_eq(
    Math.seq({|a| a.last.next_omega_prime(8) }).first(10),
    %n[9699690, 11741730, 13123110, 14804790, 15825810, 16546530, 17160990, 17687670, 18888870, 19399380],
)

assert_eq(
    Math.seq({|a| a.last.next_squarefree_almost_prime(8) }).first(10),
    %n[9699690, 11741730, 13123110, 14804790, 15825810, 16546530, 17160990, 17687670, 18888870, 20030010],
)

assert_eq(
    Math.seq(1, 1, {|a,n| a.tail + n*subfactorial(n-1) }).first(10),
    %n[1, 1, 4, 12, 57, 321, 2176, 17008, 150505, 1485465]
)

assert_eq(
    Math.seq(0, {|a,n| a.tail*n + ((-1)**n * (n-1)!) }).first(10),
    %n[0, 1, 1, 10, 26, 276, 1212, 14736, 92304, 1285920]
)

assert_eq(
     Math.seq(0,0,2, {|a| a[-1] + a[-2] + a[-3] }).while { _ <= 1e4 },
     %n[0, 0, 2, 2, 4, 8, 14, 26, 48, 88, 162, 298, 548, 1008, 1854, 3410, 6272]
)

assert_eq(Math.seq(100, {.tail.prev_composite}).first(100).grep{.is_int}.flip, composites(100))
assert_eq(Math.seq(100.prev_prime, {.tail.prev_prime}).first(100).grep{.is_int}.flip, primes(100))

do {
    func wordwrap(words, maxwidth) {
        Enumerator({ |y|
            var cols = 0
            words.slice_before { |w|
                cols += 1 if (cols != 0)
                cols += w.len
                if (maxwidth < cols) {
                    cols = w.len
                    true
                } else {
                    false
                }
            }.each { |ws| y(ws) }
        })
    }

    var text = "The quick brown fox jumps over the lazy dog"
    var arr = []
    wordwrap(text.words, 15).each { |line| arr << line.join(" ") }
    assert_eq(arr, ["The quick brown", "fox jumps over", "the lazy dog"])
}

say "** Test passed!"