#!/usr/bin/ruby

# Tests for various Number divisor functions.

var h = Hash(
    divisors => 'sigma',
    prime_divisors => 'prime_sigma',
    prime_power_divisors => 'prime_power_sigma',
    square_divisors => 'square_sigma',
    cube_divisors => 'cube_sigma',
    squarefree_divisors => 'squarefree_sigma',
    cubefree_divisors => 'cubefree_sigma',

    udivisors => 'usigma',
    prime_udivisors => 'prime_usigma',
    prime_power_udivisors => 'prime_power_usigma',
    square_udivisors => 'square_usigma',
    cube_udivisors => 'cube_usigma',
    squarefree_udivisors => 'squarefree_usigma',
    cubefree_udivisors => 'cubefree_usigma',

    edivisors => 'esigma',
    idivisors => 'isigma',
    bdivisors => 'bsigma',
)

var h2 = Hash(
    power_divisors => 'power_sigma',
    power_udivisors => 'power_usigma',

    powerfree_divisors => 'powerfree_sigma',
    powerfree_udivisors => 'powerfree_usigma',
)

with([irand(3,10)!, 340282366920938463942989953348216553641, 92971590131496140210160].rand) {|n|

    h.each {|a,b|
        var D = n.(a)
        for k in (0..3) {
            assert_eq(D.sum {|d| d**k }, n.(b)(k), "#{a} != #{b} for n=#{n} and k=#{k}")
        }
        assert_eq(D.len, n.(b+'0'), "#{a}.len != #{b+'0'} for n=#{n}")
    }

    h2.each {|a,b|
        for k in (0..3) {
            var D = k.(a)(n)
            for j in (0..3) {
                assert_eq(D.sum {|d| d**j }, k.(b)(n,j), "#{a} != #{b} for n=#{n}, k=#{k}, j=#{j}")
            }
            assert_eq(D.len, k.(b+'0')(n), "#{a}.len != #{b+'0'} for n=#{n}, k=#{k}")
        }
    }
}

with (irand(5, 10)!) {|n|
    var D = n.divisors
    assert_eq(n.perfect_power_divisors,    D.grep{.is_power})
    assert_eq(1.power_divisors(n), D.grep{.is_power(1)})
    assert_eq(2.power_divisors(n), D.grep{.is_power(2)})
    assert_eq(3.power_divisors(n), D.grep{.is_power(3)})
    assert_eq(4.power_divisors(n), D.grep{.is_power(4)})
    assert_eq(5.power_divisors(n), D.grep{.is_power(5)})
    assert_eq(9.power_divisors(n), D.grep{.is_power(9)})
}

with (irand(10, 21)!) {|n|
    var D = n.udivisors
    assert_eq(n.pp_udivisors,       D.grep{.is_power})
    assert_eq(1.power_udivisors(n), D.grep{.is_power(1)})
    assert_eq(2.power_udivisors(n), D.grep{.is_power(2)})
    assert_eq(3.power_udivisors(n), D.grep{.is_power(3)})
    assert_eq(4.power_udivisors(n), D.grep{.is_power(4)})
    assert_eq(5.power_udivisors(n), D.grep{.is_power(5)})
    assert_eq(9.power_udivisors(n), D.grep{.is_power(9)})
}

with (irand(10, 21)!) {|n|
    var D = n.udivisors
    assert_eq(1.powerfree_udivisors(n), D.grep{.is_powerfree(1)})
    assert_eq(2.powerfree_udivisors(n), D.grep{.is_powerfree(2)})
    assert_eq(3.powerfree_udivisors(n), D.grep{.is_powerfree(3)})
    assert_eq(4.powerfree_udivisors(n), D.grep{.is_powerfree(4)})
    assert_eq(5.powerfree_udivisors(n), D.grep{.is_powerfree(5)})
    assert_eq(6.powerfree_udivisors(n), D.grep{.is_powerfree(6)})
    assert_eq(7.powerfree_udivisors(n), D.grep{.is_powerfree(7)})
}

for n in (0..50) {
    var D1 = n.divisors
    var D2 = n.udivisors

    assert_eq(n.pp_divisors,       D1.grep { .is_power })
    assert_eq(1.power_divisors(n), D1.grep { .is_power(1) })
    assert_eq(2.power_divisors(n), D1.grep { .is_power(2) })
    assert_eq(3.power_divisors(n), D1.grep { .is_power(3) })

    assert_eq(n.pp_udivisors,       D2.grep { .is_power })
    assert_eq(1.power_udivisors(n), D2.grep { .is_power(1) })
    assert_eq(2.power_udivisors(n), D2.grep { .is_power(2) })
    assert_eq(3.power_udivisors(n), D2.grep { .is_power(3) })

    assert_eq(0.powerfree_udivisors(n), D2.grep { .is_powerfree(0) })
    assert_eq(1.powerfree_udivisors(n), D2.grep { .is_powerfree(1) })
    assert_eq(2.powerfree_udivisors(n), D2.grep { .is_powerfree(2) })
    assert_eq(3.powerfree_udivisors(n), D2.grep { .is_powerfree(3) })
    assert_eq(4.powerfree_udivisors(n), D2.grep { .is_powerfree(4) })
    assert_eq(5.powerfree_udivisors(n), D2.grep { .is_powerfree(5) })
}

assert_eq(5040.sigma(-1), 5040.divisors.sum {|d| d**(-1) })
assert_eq( 10!.sigma(-2),  10!.divisors.sum {|d| d**(-2) })

do {

    var n = 5040
    var D = n.divisors

    for k in (0..4) {
        assert_eq(k.powerfree_divisors(n), D.grep{ .is_powerfree(k) })
    }
}

do {    # bi-unitary divisors
    func biudivs(n) {
        #return n.bdivisors
        n.divisors.grep {|x| gcud(x,n/x) == 1 }
    }

    func a(n,k=1) {
        biudivs(n).sum { _**k }
    }

    func g(n, k=1) {

        return 0 if (n == 0)

        n.factor_prod {|p,e|
            (p**(k*(e + 1)) - 1)/(p**k - 1) - (e.is_even ? p**(k*(e/2)) : 0)
        }
    }

    assert_eq(
        30.of { bsigma(_, 1) },
        30.of { g(_) }
    )

    assert_eq(
        30.of { bsigma(_, 2) },
        30.of { g(_, 2) }
    )

    assert_eq(
        30.of { bsigma(_, 3) },
        30.of { g(_, 3) }
    )

    assert_eq(
        30.of { .bsigma },
        30.of { a(_) },
    )

    assert_eq(
        30.of { .bsigma(2) },
        30.of { a(_, 2) },
    )

    assert_eq(
        30.of { .bsigma(3) },
        30.of { a(_, 3) },
    )

    assert_eq(
        30.of  { a(_, 0) }
        30.of  { .bsigma(0) }
    )

    assert_eq(
        30.of  { .divisors - biudivs(_) -> sum }
        30.of  { .nbsigma }
    )

    assert_eq(
        30.of  { .divisors - biudivs(_) -> sum { _**2 } }
        30.of  { .nbsigma(2) }
    )

    assert_eq(
        30.of  { .divisors - biudivs(_) -> len }
        30.of  { .nbsigma0 }
    )
}

do {    # non-unitary divisors
    assert_eq(
        30.of { .nusigma },
        30.of {|n| n.divisors.grep {|d| gcd(d, n/d) != 1 }.sum }
    )

    assert_eq(
        10.of { .nusigma(0) },
        10.of { .nusigma0 },
    )

    assert_eq(
        30.of { .nusigma(0) },
        30.of {|n| n.divisors.count {|d| gcd(d, n/d) != 1 } }
    )
}

do {    # infinitary divisors

    func a(n, k=1) {

        return 0 if (n == 0)

        n.factor_prod {|p,e|

            #~ e.digits(2).map_kv {|r,v|
                #~ (p**(2**r * (v+1) * k) - 1) / (p**(k * 2**r) - 1)
            #~ }.prod

            var prod = 1
            var r = 0

            do {
                if (e%2 == 1) {
                    #prod *= ((p**(2**(r+1) * k) - 1) / (p**(k * 2**r) - 1))
                    prod *= (p**(k * 2**r) + 1)
                }
                ++r
            } while (e >>= 1)

            prod
        }
    }

    assert_eq(
        30.of { a(_+1) },
        30.of { idivisors(_+1).sum }
    )

    assert_eq(
        30.of { a(_+1, 2) },
        30.of { idivisors(_+1).sum { _**2 } }
    )

    assert_eq(
        30.of { a(_+1, 3) },
        30.of { idivisors(_+1).sum { _**3 } }
    )

    assert_eq(
        30.of { .isigma0 },
        30.of { idivisors(_).len }
    )

    assert_eq(
        30.of { .isigma },
        30.of { a(_) },
    )

    assert_eq(
        30.of { .isigma(2) },
        30.of { a(_, 2) },
    )

    assert_eq(
        30.of { .isigma(3) },
        30.of { a(_, 3) },
    )

    assert_eq(
        30.of { divisors(_) - idivisors(_) -> sum },
        30.of { .nisigma },
    )

    assert_eq(
        30.of { divisors(_) - idivisors(_) -> len },
        30.of { .nisigma0 },
    )

    assert_eq(
        30.of { divisors(_) - idivisors(_) -> sum { _**2 } },
        30.of { .nisigma(2) },
    )

    assert_eq(
        30.of { .antidivisors.sum },
        30.of { .antidivisor_sum },
    )

    assert_eq(
        30.of { .antidivisors.len },
        30.of { .antidivisor_count },
    )
}

do {    # A051377: exponential divisors

    func ediv(n) {

        #return n.edivisors
        return [1] if (n == 1)

        var f = n.factor_exp
        var D = f.map{.tail.divisors}

        var L = []

        D.map {|d| @(^d.len) }.cartesian {|*a|
            L << f.prod_kv {|j,pp|
                pp[0]**D[j][a[j]]
            }
        }

        L
    }

    func f(n, k=1) {

        return 0 if (n == 0)

        # Multiplicative with:
        #   a(p^e) = Sum_{d|e} p^d

        n.factor_map {|p,e|
            e.divisors.sum {|d| p**(k*d) }
        }.prod
    }

    assert_eq(ediv(7!).sum, f(7!))
    assert_eq(ediv(8!).sum { _**2 }, f(8!, 2))
    assert_eq(ediv(9!).sum { _**3 }, f(9!, 3))

    assert_eq(
        30.of { ediv(_).sum },
        30.of { .esigma(1) }
    )

    assert_eq(
        30.of { ediv(_).sum { _**3 } },
        30.of { .esigma(3) }
    )

    assert_eq(
        30.of { ediv(_).len },
        30.of { .esigma(0) }
    )

    assert_eq(
        30.of { ediv(_).len },
        30.of { .esigma0 }
    )

    assert_eq(
        30.of { f(_) },
        30.of { .esigma }
    )

    assert_eq(
        30.of { f(_, 2) },
        30.of { .esigma(2) },
    )

    assert_eq(
        30.of { f(_, 3) },
        30.of { .esigma(3) },
    )

    assert_eq(
        30.of { divisors(_) - ediv(_) -> sum },
        30.of { .nesigma },
    )

    assert_eq(
        30.of { divisors(_) - ediv(_) -> sum { _**2 } },
        30.of { .nesigma(2) },
    )

    assert_eq(
        30.of { divisors(_) - ediv(_) -> len },
        30.of { .nesigma0 },
    )
}

assert_eq(divisors(2**128 - 1).grep{ _ <= 5040 }.sum, 16738)

assert_eq(divisors(2**128 - 1, 5040).len, 16)
assert_eq(divisors(2**128 - 1, 5040).sum, 16738)
assert_eq(divisors(2**128 - 1, 2**64 - 1).len, 256)
assert_eq(divisors(2**128 - 1, 2**64 + 1).len, 257)
assert_eq(divisors(2**128 - 1, 2**64 + 1).sum, 155871713282619791539)
assert_eq(divisors(2**128 - 1, 2**32 - 1).sum, 33234996535)
assert_eq(divisors(2**128 - 1, 2**32 + 1).sum, 37529963832)

assert_eq(3 * (2**127 - 1) -> divisors(2**127 - 1), %n[1, 3, 170141183460469231731687303715884105727])
assert_eq(3 * (2**127 - 1) -> divisors(3*(2**127 - 1)), %n[1, 3, 170141183460469231731687303715884105727, 510423550381407695195061911147652317181])
assert_eq(5040.divisors(120), %n[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 12, 14, 15, 16, 18, 20, 21, 24, 28, 30, 35, 36, 40, 42, 45, 48, 56, 60, 63, 70, 72, 80, 84, 90, 105, 112, 120])
assert_eq(divisors(next_prime(next_prime(2**64) / 3) * 6, 2**64 - 1), %n[1, 2, 3, 6, 6148914691236517223, 12297829382473034446])
assert_eq(divisors(prev_prime(next_prime(2**64) / 3) * 6, 2**64 - 1), %n[1, 2, 3, 6, 6148914691236517199, 12297829382473034398, 18446744073709551597])

assert_eq(divisors(next_prime((2**64) / 2) * 2, 2**64 - 1), %n[1, 2, 9223372036854775837])
assert_eq(divisors(next_prime((2**64) / 3) * 3, 2**64 - 1), %n[1, 3, 6148914691236517223])
assert_eq(divisors(next_prime((2**64) / 5) * 5, 2**64 - 1), %n[1, 5, 3689348814741910379])

say "** Test passed!"