gtauseq

#!/opt/maths/bin/perl
use strict;
use warnings;
use Math::GMP;
use Math::Prime::Util qw{ is_prime next_prime divisors factor_exp };

use lib './lib';
use Constraint;
use Type;
use ModFunc qw/ gcd /;

sub MBI { return Math::GMP->new(@_) }

my($opt_D, $opt_f, $opt_b, $opt_ta, $opt_fact, $opt_ix, $opt_dx) = (0) x 7;
my($opt_n, $opt_x, $opt_c, $opt_t, $opt_cp, $opt_cr)
        = (0, 0, 0, 100000, 0, 0);
my($opt_ts, $opt_m, $opt_mq, $opt_mm, $opt_dm, $opt_dmq)
        = ([], [], [], [], [], []);
my $typename = 't'; # t=tauseq, a=addseq, s=semiprime

while (@ARGV && $ARGV[0] =~ /^-/) {
    my $arg = shift @ARGV;
    last if $arg eq '--';
    ($opt_D = 1, next) if $arg eq '-D';
    ($opt_f = 1, next) if $arg eq '-f';
    ($opt_b = 1, next) if $arg eq '-b';
    ($opt_ta = 1, next) if $arg eq '-ta';
    ($opt_fact = 1, next) if $arg eq '-fact';
    ($opt_ix = 1, next) if $arg eq '-ix';
    Constraint->debug_more(), next if $arg eq '-d';
    ($opt_n = $opt_x = $arg || shift(@ARGV)), next if $arg =~ s{^-nx}{};
    ($opt_n = $arg || shift(@ARGV)), next if $arg =~ s{^-n}{};
    ($opt_x = $arg || shift(@ARGV)), next if $arg =~ s{^-x}{};
    ($opt_cp = $arg || shift(@ARGV)), next if $arg =~ s{^-cp}{};
    ($opt_cr = $arg || shift(@ARGV)), next if $arg =~ s{^-cr}{};
    ($opt_c = $arg || shift(@ARGV)), next if $arg =~ s{^-c}{};
    ($opt_dx = $arg || shift(@ARGV)), next if $arg =~ s{^-dx}{};
    ($opt_ts = [ split /,/, $arg || shift(@ARGV) ]), next if $arg =~ s{^-ts}{};
    ($opt_t = $arg || shift(@ARGV)), next if $arg =~ s{^-t}{};
    ($typename = $arg || shift(@ARGV)), next if $arg =~ s{^-y}{};
    push(@$opt_mq, $arg || shift(@ARGV)), next if $arg =~ s{^-mq}{};
    push(@$opt_mm, $arg || shift(@ARGV)), next if $arg =~ s{^-mm}{};
    push(@$opt_m, $arg || shift(@ARGV)), next if $arg =~ s{^-m}{};
    push(@$opt_dmq, $arg || shift(@ARGV)), next if $arg =~ s{^-dmq}{};
    push(@$opt_dm, $arg || shift(@ARGV)), next if $arg =~ s{^-dm}{};
    die "Unknown option '$arg'\n";
}

$| = 1;
@ARGV == 2 or die "Usage: $0 [ options ] n f\n";
my($n, $f) = map MBI($_), @ARGV;
$0 = "gt($n $f)";
$_ = ston($_) for ($opt_n, $opt_x);
$opt_n = MBI(1) if $opt_n < 1;

printf "100 ./gtauseq -y%s %s %s%s%s%s -c%s -t%s %s %s\n",
        $typename, ($opt_n == $opt_x ? "-nx$opt_n" : "-n$opt_n -x$opt_x"),
        _opta('m', $opt_m), _opta('mq', $opt_mq), _opta('mm', $opt_mm),
        (@$opt_ts ? "-ts@{[ join ',', @$opt_ts ]}" : ''),
        $opt_c, $opt_t, $n, $f;

my $zone = MBI(1);
my $c = Constraint->new(
    'n' => $n,
    'f' => $f,
    'tell_count' => $opt_t,
    't0' => scalar times(),
    'min' => $opt_n,
    'max' => $opt_x,
    'check' => $opt_c,
    'checkp' => $opt_cp,
    'min_potency' => $opt_cr,
);
my $type = Type->new($typename, $c);
$c->set_type($type);

$c = apply($c) or exit 0;
$opt_ix or $type->check_exceptions;
printf <<OUT, $c->elapsed(), $type->name, $n, $f, ntos($opt_n), ntos($opt_x), modfix(), $opt_c;
300 Init %.2f: trying %s() for (%s, %s) in [%s, %s], %schecks up to mod %s
OUT
mod_dump($c, $opt_dm);

my $d = rootseq($c);
if ($d) {
    report_seq($n, $f, $d);
    printf <<OUT, $n, $f, $d, modfix(), $c->elapsed();
200 f(%s, %s) = %s %s (%.3fs)
OUT
} else {
    printf <<OUT, $n, $f, ntos($opt_x), modfix(), $c->elapsed();
500 f(%s, %s) > %s %s (%.3fs)
OUT
}
exit 0;

sub mod_override {
    my($c, $override) = @_;
    my($mod, $op, $val) = m{ ^ (\d+) ([=!]) ([\d,]+) \z }x
            or die "Invalid mod override '$_'";
    if ($op eq '=') {
        $c->require($mod, $val);
    } else {
        $c->suppress($mod, $_) for split /,/, $val;
    }
}

sub modfix {
    return join ', ', grep length,
        (@$opt_m ? sprintf('mod(%s) ', join ', ', @$opt_m) : ''),
        (@$opt_mq ? sprintf('sqmod(%s) ', join ', ', @$opt_mq) : '');
}

sub mod_dump {
    my($c, $which, $s) = @_;
    my $legend = $s ? " [$s]" : '';
    if ($opt_dx && !$s) {
        my $n = $c->n;
        my %keep = map +($_ => 1), @$which;
        @$which = ();
        for my $mod (2 .. $c->check) {
            push(@$which, $mod), next if $keep{$mod};
            next if $opt_cp && cp_skip($opt_cp, $n, [ factor_exp($mod) ]);
            push @$which, $mod if $mod - $c->c($mod)->[8] <= $opt_dx;
        }
    }
    for my $mod (@$which) {
        if ($mod eq 'a') {
            $c->Dumpvecs;
            next;
        }
        my $cm = $c->c($mod);
        my $vm = $cm->[2];  # bit vector of all disallowed values
        if ($c->debug < 1) {
            my $dis = unpack "b$mod", $vm;
            $dis .= '0' x ($mod - length $dis);
            printf <<OUT, $legend, $mod, $dis;
331 disallowed%s (mod %s) [%s]
OUT
        } else {
            printf <<OUT, $legend, $mod, join ', ', grep vec($vm, $_, 1), 0 .. $mod - 1;
332 disallowed%s (mod %s): %s
OUT
        }
    }
}

sub unpack_mm {
    my($mm) = @_;
    # Expect "mod1=val1,mod2=val2,...modn=valn/b1b2...bn"
    die "Could not parse '-mm$mm'\n" unless $mm =~ m{
        ^ (?: \d+ = \d+ ,)* \d+ = \d+ / \d+ $
    }x;
    my($modlist, $bits) = split '/', $mm;
    my @modlist = split /,/, $modlist;
    die "Wrong number of bits specified in '-mm$mm'\n"
            unless @modlist == length $bits;
    push @$opt_m, @modlist;
    push @$opt_mq, map [
        $modlist[$_], substr($bits, $_, 1),
    ], 0 .. $#modlist;
}

sub apply {
    my $c = shift;
    my $n = $c->n;

    # If we require the result to be coprime with the inputs, deal with that
    # first.
    if ($opt_b) {
        for my $p (map $_->[0], factor_exp($c->n)) {
            $c->suppress($p, 0);
        }
    }

    # Unpack $opt_mm into @$opt_m and @$opt_mq
    unpack_mm($_) for @$opt_mm;

    # If options specify overriding modular constraints, apply those next.
    mod_override($c, $_) for @$opt_m;

    for my $m (2 .. $c->check) {
        $c->debug && warn "apply $m\n";
        my $fm = [ factor_exp($m) ];
        next if $opt_cp && cp_skip($opt_cp, $n, $fm);

        # apply type-specific constraints for this modulus
        $type->apply_m($m, $fm);
    }

    #
    # analyse the constraints, to discover mult/mod_mult
    #
    $c->find_active;    # update mult, mod_mult

    #
    # Now go through each k working out what factors are fixed; we can use
    # the info to decide the most efficient tester for each k, and may also
    # find opportunities for additional constraint optimizations.
    #
    my $fixpow = $type->check_fixed;

    # fix_fact only for tauseq right now
    if ($opt_fact && $type->name eq 'tauseq') {
        my $fixed = $c->fix_fact;
        # It may do nothing, add some constraints (and return nothing)
        # or return a new object. Only in the last case do we want to
        # avoid the upcoming fixpow check.
        return $fixed if $fixed;
    }

    if ($fixpow) {
        my($k, $x, $z) = @$fixpow;
        if ($type->name eq 'tauseq') {
            print "311 Fixing power d = (${x}y^$z - n)/$k\n";
        } elsif ($type->name eq 'addseq') {
            print "311 Fixing power d = ${x}y^$z - ${k}n\n";
        } elsif ($type->name eq 'oneseq') {
            print "311 Fixing power d = ${x}y^$z - ${k}\n";
        } elsif ($type->name eq 'track') {
            print "311 Fixing power d = ${x}y^$z - ${k}\n";
        } else {
            die "unexpected type @{[ $type->name ]}\n";
        }
        mod_dump($c, $opt_dm, 'pre-fix');
        $c = $c->fix_power($k, $x, $z, $opt_mq);
        $opt_dm = $opt_dmq;
    } elsif (@$opt_mq) {
        die "-mq or -mm options invalid without fixing power";
    }

    return $c;
}

#
# Return TRUE if the check c (factorizing to fm) should be skipped for n
# according to the cp option, ie if the largest factor of c not dividing
# n is greater than cp.
#
sub cp_skip {
    my($cp, $n, $fm) = @_;
    # assume the greatest factor
    my $i = $#$fm;
    # skip past those that divide n
    --$i while $i >= 0 && ($n % $fm->[$i][0]) == 0;
    return 1 if $i >= 0 && $fm->[$i][0] > $cp;
    return 0;
}

{
    my @seq_seen;
    #
    # rootseq(n, tau, f, g, constraint)
    # - search for arithmetic sequence n+kd, 0<=k<=f, with tau(n+kd)=tau(n)
    # - search with d in the range g=[min, max]
    # - check only those d that satisfy constraints in the Constraint object
    #
    sub rootseq {
        my($c) = @_;
        my $count = $c->tell_count;
        my($n, $f) = ($c->n, $c->f);
        my $k = $f - 1;
        my $cur;
        $c->init();
        $type->check_mult unless $opt_f;
        printf <<OUT, $c->frequency(), $c->elapsed();
309 Prep finished, frequency %.2f (%.3fs)
OUT
        exit 0 if $opt_D;

        my $tester = prepare_tester($c);
        $tester = test_all($tester) if $opt_ta;

        CUR: while ($cur = $c->next()) {
            next if $type->disallow($cur);
            report_reach($n, $f, $cur, $c), $count = $c->tell_count
                    unless --$count;
            for (0 .. $#$tester) {
                next if $tester->[$_]->($cur);
                report_seq($n, $f, $cur, $_) unless $seq_seen[$_]++;
                next CUR;
            }
            $seq_seen[@$tester]++;
            report_reach($n, $f, $cur, $c);
            return $cur;
        }
        # We're done; if cur is a derived value (eg for Constraint::Power),
        # it is not necessarily an integer at this point, but that's ok.
        no warnings qw{once};
        local $::LOOSE_CUR = 1;
        report_reach($n, $f, $c->cur, $c);
        return undef;
    }

    sub report_reach {
        my($n, $f, $cur, $c) = @_;
        my $t = $c->elapsed();
        my($tests, $skipped, $kept) = ($c->tests(), $c->skipped(), $c->kept());

        printf <<OUT, $t, $n, $f, $cur, $kept, $skipped, $tests, join ' ', map $_ || 0, @seq_seen;
301 After %.2fs for (%s, %s) reach d=%s (keep %s, skip %s with %s tests) seen [%s]
OUT
    }

    sub test_all {
        my($tester) = @_;
        return [ sub {
            my($cur) = @_;
            my $pass = 1;
            for (0 .. $#$tester) {
                next if $tester->[$_]->($cur);
                ++$seq_seen[$_ + 1];
                $pass = 0;
            }
            return $pass;
        } ];
    }
}

#
# Returns an arrayref of subrefs, that (in some order) test whether each of
# the candidate targets n+kd has the required number of factors.
# Each subref accepts a single argument (d), and returns a boolean: TRUE
# means tau(n+kd) = tau(n).
#
# We try to arrange the order and strategy of the testing subrefs to optimize
# speed:
# - we can test primality much faster than we can factorize
# - we want to test first those least likely to return TRUE (all other things
#   being equal)
# - generally the most constrained are least likely to return true; that
#   probably means that given g=gcd(n,k) we should test those k that
#   maximize tau(g) first
# - additionally, we'd prefer not to multiply up with known factors, but
#   we must take care to elide only fixed powers (so we have a known tau)
# - if we had a fast issquare() check, that would also be a useful thing
#   to take advantage of
#
# eg for f(125), f(343) we can use prime checks
#
# TODO: the current "let's put what we know most about first" actually
# gives us a good ordering for information (via 301 lines), but is badly
# suboptimal for speed - those tested earliest seem in practice to be the
# ones least likely to be rejected.
# Consider testing all k without shortcircuit for the first <x> tests,
# and then picking an order that tests first those most likely to be
# rejected. (Outputting the accept rate at this point would also be
# interesting, since we can use it to start making probabilistic predictions
# about when we'll find a solution.)
#
# TODO: We can improve the prime special-case by not doing a full is_prime()
# check for each candidate in turn, but rather doing single rounds of
# maybe_prime() for each in turn looking for cheaper opportunities to reject
# before doing the more expensive full check on a realistic set of candidates.
# If additional rounds of maybe_prime can avoid repeating work, we could even
# split the work into more than two passes:
#   my @testers = (
#       map(sub { maybe_prime($_[0], $_) }, 1 .. $MAX),
#       sub { is_prime($_) },
#   );
#   for my $tester (@testers) {
#       for my $candidate (@candidates) {
#           last THIS_SET unless $tester->($candidate);
#       }
#   }
#   return \@candidates;   # found a solution
#
sub prepare_tester {
    my($c) = @_;
    my $n = $c->n;
    my %t;

    my $all_target = $type->to_test;

    # override any other ordering if we're doing 'test all'
    $opt_ts = $all_target if $opt_ta;
    my $fixed = @$opt_ts ? 1 : 0;

    for my $k (@$all_target) {
        $t{$k} = $type->test_target($k);
    }

    my %pref = map +($_ => gcd($n, $_)), @$all_target;
    my @order = map {
        my $v = $_ + 0;
        die "Invalid test order '$_'" unless delete $pref{$v};
        $v
    } @$opt_ts;
    push @order, sort { $pref{$b} <=> $pref{$a} || $a <=> $b } keys %pref;
    printf "320 Testers: %s%s\n",
        join(', ', map $_->[0], @t{@order}),
        $opt_ta ? ' (test all)' : '';
    return [ map $_->[1], @t{@order} ];
}

sub report_seq {
    my($n, $f, $d, $first) = @_;

    my $good = 1;
    my @result;
    for (my $i = 0; ($i <= $f) || $good; ++$i) {
        my $v = $type->func_value($n, $i, $d);
        my $result = [ $i, $type->func($v), $v, pf($v) ];
        $good = 0 unless $type->func_matches($i, $result->[1]);
        push @result, $result;
    }

    my $flen = length($#result);
    my $vlen = length($result[-1][2]);
    my $name = $type->func_name;
    for (0 .. $#result) {
        my $r = $result[$_];
        printf <<OUT, $flen, $r->[0], $r->[1], $name, $vlen, $r->[2], $r->[3];
211 Sequence % *s: % 2s = %s(% *s = %s)
OUT
    }
    if (defined $first) {
        printf <<OUT, $c->elapsed, $n, $f, $first;
306 After %.2fs for (%s, %s) first fail for test %s
OUT
    }
}

sub check_limit {
    my($n, $f) = splice @_, 0, 2;
    my $max = 0;
    ($max < $_) && ($max = $_) for @_;
    return 1 if $f <= $max;
    printf <<OUT, $n, $max, $f;
401 Error: f(%s) <= %s, you asked for %s.
OUT
    return undef;
}

sub pf {
    my $n = shift;
    return join '.', map {
        $_->[1] == 1 ? $_->[0] : "$_->[0]^$_->[1]"
    } factor_exp($n);
}

sub tau {
    my $n = shift;
    my $k = 1;
    $k *= $_->[1] + 1 for factor_exp($n);
    return $k;
}

sub ston {
    my($s) = @_;
    $s =~ s/,//g;
    $s =~ s{e(\d+)}{"0" x $1}ie;
    return MBI($s);
}

sub ntos {
    my($n) = @_;
    $n =~ s{(0+)$}{"e" . length($1)}e;
    return $n;
}

sub _opta {
    my($name, $a) = @_;
    return join '', map "-$name$_ ", @$a;
}