#!/usr/bin/ruby

# Tests for rough and smooth Number methods.

func my_is_rough(n,k) {

    return false if (n <= 0)    # must be a positive integer

    n.factor.all {|p|
        p >= k
    }
}

func my_is_smooth(n,k) {

    return false if (n <= 0)    # must be a positive integer

    n.factor.all {|p|
        p <= k
    }
}

func my_smooth_part(n,k) {

    return 0 if (n <= 0)

    n.factor.grep { |p|
        p <= k
    }.prod
}

func my_rough_part(n,k) {

    return 0 if (n <= 0)

    n.factor.grep { |p|
        p >= k
    }.prod
}

say "=> Smooth testing...";

assert_eq(
    13.of {|n| 13.of {|k| [n,k, is_smooth(n,k)] } },
    13.of {|n| 13.of {|k| [n,k, my_is_smooth(n,k)] } },
)

say "=> Smooth over prod testing...";

assert_eq(
    20.of {|n| 20.of {|k| [n,k, is_smooth_over_prod(n.fib, k.primorial)] } },
    20.of {|n| 20.of {|k| [n,k, my_is_smooth(n.fib, k)] } },
)

say "=> Rough testing...";

assert_eq(
    20.of {|n| 20.of {|k| [n,k, is_rough(n,k)] } },
    20.of {|n| 20.of {|k| [n,k, my_is_rough(n,k)] } },
)

say "=> Smooth part...";

assert_eq(
    20.of {|n| 20.of {|k| [n,k, k.smooth_part(n)] } },
    20.of {|n| 20.of {|k| [n,k, my_smooth_part(n,k)] } },
)

say "=> Smooth part <=> make_coprime...";

assert_eq(make_coprime(0, 1), 0)
assert_eq(make_coprime(0, 2), 0)
assert_eq(make_coprime(0, 3), 0)

assert_eq(make_coprime(-42, 6), -7)
assert_eq(make_coprime(-42, -6), -7)

assert_eq(
    20.of {|n| 20.of {|k| [n,k, k.smooth_part(n+1)] } },
    20.of {|n| 20.of {|k| [n,k, (n+1) / make_coprime(n+1, k.primorial)] } },
)

say "=> Rough part...";

assert_eq(
    20.of {|n| 20.of {|k| [n,k, k.rough_part(n)] } },
    20.of {|n| 20.of {|k| [n,k, my_rough_part(n,k)] } },
)

say "=> Rough part <=> make_coprime...";

assert_eq(
    20.of {|n| 20.of {|k| [n,k, (k+1).rough_part(n)] } },
    20.of {|n| 20.of {|k| [n,k, make_coprime(n, k.primorial)] } },
)

do {
    var n = %n[17, 17, 19, 23, 23, 29, 47, 53, 59].prod
    var D = n.divisors

    assert_eq(47.rough_part(n), D.last_by { .is_rough(47) })
    assert_eq(46.rough_part(n), D.last_by { .is_rough(46) })
    assert_eq(48.rough_part(n), D.last_by { .is_rough(48) })

    assert_eq(47.smooth_part(n), D.last_by { .is_smooth(47) })
    assert_eq(46.smooth_part(n), D.last_by { .is_smooth(46) })
    assert_eq(48.smooth_part(n), D.last_by { .is_smooth(48) })

    assert_eq(101.rough_part(43*97), 1)
    assert_eq(97.rough_part(43*97*43*43*97), 97**2)
    assert_eq(98.rough_part(43*97*43*43*97), 1)

    assert_eq(23.smooth_part(43*97), 1)
    assert_eq(43.smooth_part(43*97*43*97*43), 43**3)
    assert_eq(41.smooth_part(43*97*43*97*43), 1)

    assert_eq(19.smooth_part(n*17*19*23), 17**3 * 19**2)
    assert_eq(17.smooth_part(n*17*19*23), 17**3)
    assert_eq(18.smooth_part(n*17*19*23), 17**3)
}

do {
  var n = 1377276413364943226363244108454842276965894752197358387200000; # 97

  assert(!is_smooth(n,23))
  assert(!is_smooth(n,96))
  assert(is_smooth(n,97))
  assert(is_smooth(n,98))
}

do {
  var n = 172864518041328651521584134678230948270774322090771071422829; # 2081

  assert(is_smooth(n, 4073))
  assert(is_rough(n, 2080))
  assert(is_rough(n, 2081))
  assert(!is_rough(n, 2082))
}

assert_eq(7.rough_count(2**128 + 1), 90741964512250256923566561981804856389)
assert_eq(11.rough_count(2**128), 77778826724785934505914195984404162619)

assert_eq(11.rough_count(2**64 - 100), 4216398645419326060)
assert_eq(11.rough_count(2**64), 4216398645419326083)
assert_eq(11.rough_count(2**64 + 1), 4216398645419326084)
assert_eq(12.rough_count(2**64), 3833089677653932802)
assert_eq(5.rough_count(2**64),  6148914691236517205)
assert_eq(19.rough_count(2**65), 6660210118638507676)

do {
    for k in (2..10) {

        var a = (10+k).by { .is_smooth(k) }
        var b = (10+k).by { my_is_smooth(_, k) }

        assert_eq(a,b)

        var count = k.smooth_count(a.tail)
        assert_eq(a.len, count)
    }
}

do {
    for k in (2..10) {

        var a = (100+k).by { .is_rough(k) }
        var b = (100+k).by { my_is_rough(_, k) }

        assert_eq(a,b)

        var count = k.rough_count(a.tail)
        assert_eq(a.len, count)
    }
}

do {
    var a = Math.smooth_numbers(2,3,5,7)    # 7-smooth numbers
    var b = Math.smooth_numbers(2,5,7)      # 7-smooth numbers not divisible by 3

    assert_eq(a.first(30), 30.by { .is_smooth(7) })
    assert_eq(b.first(30), 30.by {!.is_div(3) && .is_smooth(7) })

    # Iteration is also supported
    a.each {|k|
        if (k > 1e5) {
            assert_eq(k, 100352)
            break
        }
    }
}

do {
    var n = 10!
    var D = n.divisors

    assert_eq(3.smooth_divisors(n), D.grep{.is_smooth(3)})
    assert_eq(5.smooth_divisors(n), D.grep{.is_smooth(5)})
    assert_eq(5.rough_divisors(n), D.grep{.is_rough(5)})
    assert_eq(3.rough_divisors(n), D.grep{.is_rough(3)})

    7.of {|k| 7.of {|n|
        assert_eq(k.smooth_divisors(n), n.divisors.grep{.is_smooth(k)})
        assert_eq(k.rough_divisors(n), n.divisors.grep{.is_rough(k)})
    }}
}

func my_rough_count_1 (n,k) {

    var P = k.dec.primes

    if (P.len == 0) {
        return n
    }

    func (n,a) is cached {

        var count = (n - (n >> 1))

        for j in (1 ..^ a) {
            count -= __FUNC__(idiv(n, P[j]), j)
        }

        return count
    }(n, P.len)
}

func my_rough_count_2 (n, k) {

    var P = k.dec.primes

    func (n,a) is cached {

        if (a == 0) {
            return n
        }

        __FUNC__(n, a-1) - __FUNC__(idiv(n, P[a-1]), a-1)
    }(n, P.len)
}

func my_rough_count_3 (n,k) {

    if (n <= 0) {
        return 0
    }

    func (n,p) is cached {

        if (p > n.isqrt) {
            return 1
        }

        if (p == 2) {
            return (n >> 1)
        }

        if (p == 3) {
            var t = idiv(n,3)
            return (t - (t >> 1))
        }

        var u = 0
        var t = idiv(n,p)

        for (var q = 2; q < p; q.next_prime!) {

            var v = __FUNC__(t - (t % q), q)

            if (v == 1) {
                u += prime_count(q, p-1)
                break
            }

            u += v
        }

        t - u
    }(n*k, k)
}

for p in (primes(23)) {
    var values = 15.range.map { irand(10**_) }
    var built_in = 15.range.map { p.rough_count(values[_]) }
    assert_eq(15.range.map { my_rough_count_1(values[_], p) }, built_in)
    assert_eq(15.range.map { my_rough_count_2(values[_], p) }, built_in)
    assert_eq(15.range.map { my_rough_count_3(values[_], p) }, built_in)
}

func meissel_lehmer_prime_count(n) {

    if (n <= 1) {
        return 0
    }

    var a = n.icbrt+1
    var k = a+1
    var s = k.rough_count(n)

    s + a.pi - 1 - 2.almost_primes(n.isqrt, n).count { .is_rough(k) }
}

for n in (1..5) {
    var t = irand(10**n)
    assert_eq(meissel_lehmer_prime_count(t), t.pi)
}

say "** Test passed!"