package library;
use strict;
our @EXPORT_OK = qw(
  factor factors generate_stream generate_memoryless_stream is_prime is_prime2
  prime_iterator ith_prime gcd n_pow_m_mod_k minimal_divisible_repunit
  pythagorean_triple_iterator
);
use Exporter qw(import);

sub pythagorean_triple_iterator {
  my $comp = shift || sub {
    my ($a, $b) = @_;
    return $a->[0] <=> $b->[0] || $a->[1] <=> $b->[1];
  };
  my $heap = Heap->new($comp);

  my $initial_iterator = primitive_pythagorean_triple_iterator();
  $heap->append(
    [$initial_iterator->(), 1, $initial_iterator]
  );

  my $max_k = 1;
  return sub {
    my $next = $heap->top();
    my ($x, $y, $z, $k, $it) = @$next;
    
    $heap->swap_top(
      [(map $k*$_, $it->()), $k, $it]
    );

    if ($max_k == $k) {
      $max_k++;
#      print "Appending $max_k\n";
      my $iter = primitive_pythagorean_triple_iterator();
      $iter->();
      $heap->append(
        [3*$max_k, 4*$max_k, 5*$max_k, $max_k, $iter]
      );
    }

    return $x, $y, $z;
  }
}

sub primitive_pythagorean_triple_iterator {
  my $comp = shift || sub {
    my ($a, $b) = @_;
    return $a->[0] <=> $b->[0] or $a->[1] <=> $b->[1];
  };

  my $heap = Heap->new($comp);

  $heap->append(
    [3, 4, 5, 1, 2]
  );

  my $generate_triple = sub {
    my ($m, $n) = @_;

    # We only want primitives.
    return()  unless 1 == $m%2 + $n%2;
    return() unless 1 == gcd($m, $n);

    my $x = 2*$m*$n;
    my $y = $m*$m - $n*$n;
    $y = - $y if $y < 0;
    my $z = $m*$m + $n*$n;

#    print "Generated ($x, $y, $z) from ($m, $n)\n";
    if ($x < $y) {
      return [$x, $y, $z, $m, $n];
    }
    else {
      return [$y, $x, $z, $m, $n];
    }
  };

  return sub {
    my $next = $heap->extract();
    my ($m, $n) = @$next[3, 4];

    if (1 == $m) {
      $heap->append(
        map $generate_triple->($_, $n+1), 1..$n
      );
      $heap->append(
        map $generate_triple->($_, $n+2), 1..($n+1)
      );
    }

    return @$next[0, 1, 2];
  };
}

sub minimal_divisible_repunit {
  my $i = shift;
  return 0 if 0 == $i%2 or 0 == $i%5;
  $i *= 9 if 0 == $i%3;

  my %factored = factor($i);
  my $euler = $i;
  for my $p (keys %factored) {
    $euler /= $p;
    $euler *= $p-1;
  }
  
  my %factor_euler = factor($euler);
  for my $p (keys %factor_euler) {
    while (0 == $euler%$p and 1 == n_pow_m_mod_k(10, $euler/$p, $i)) {
      $euler /= $p;
    }
  }

  return $euler;
}

sub gcd {
  my ($n, $m) = @_;
  while ($m) {
    ($n, $m) = ($m, $n%$m);
  }
  return $n;
}

my @primes = (2, 3, 5, 7, 11);          # None have been tested

sub more_primes {
  # This adds to the list of primes until it reaches $max
  #      or the square of the largest current prime (assumed odd)
  my $base = shift || $primes[-1]+1;
  my $max = shift || $base + 10_000;
  my $square = $primes[-1] * $primes[-1];
  $max = $square if $square < $max;     # Determine what to find primes to
  $base++ unless $base % 2;             # Make the base odd
  $max-- if $max %2;                    # Make the max odd
  $max = ($max - $base)/2;              # Make $max into a count of odds
  return if $max < 0;                   # Sanity check
  my $more = "";                        # Initialize vector of 0's for the
                                        # odd numbers in our range
  shift @primes;                        # Remove 2
  foreach my $p (@primes) {
    my $start;
    if ($base < $p * $p) {
      $start = ($p * $p - $base)/2;     # Start at the square
      if ($max < $start) {              # Rest of primes don't matter!
          last;
      }
    }
    else {                              # Start at first odd it divides
      $start = $base % $p;              # Find remainder
      $start = $p - $start if $start;   # Distance to first thing it divides
      $start += $p if $start %2;        # Distance to first odd it divides
      $start = $start/2;                # Reindex for counting over odd!
    }
    for (my $i = $start; $i <= $max; $i += $p) {
      vec($more, $i, 1) = 1;
    }
  }
  unshift @primes, 2;                   # Replace 2
  # Read off list of primes
  push @primes, map {$_ + $_ + $base} grep {vec($more, $_, 1) == 0} 0..$max;
}

sub ith_prime {
  my $i = shift;
  while ($i > $#primes) {
    more_primes();
  }
  return $primes[$i];
}

sub prime_iterator {
  my $i = -1;

  return sub {
    $i++;
    if ($i > $#primes) {
      more_primes();
    }
    return $primes[$i];
  };
}

# This generates the primes up to the one of interest.
sub is_prime {
  my $n = shift;
  return if $n < 2;
  return 1 if $n == 2;
  while ($n >= $primes[-1]) {
    more_primes();
  }


  # Now we only need to find our prime.
  my $i = 0;
  my $j = $#primes;
  while ($i+1 < $j) {
    my $k = int($i + $j)/2;
    if ($primes[$k] == $n) {
      # We found it!
      return 1;
    }
    elsif ($primes[$k] < $n) {
      $i = $k;
    }
    else {
      $j = $k;
    }
  }
  # We have it sandwiched between two consecutive primes
  return 0;
}

# This only goes to the square root of the one of interest.
sub is_prime2 {
  my $n = shift;
  my %factor = factor($n);
  return exists $factor{$n};
}

sub factor {
  my $n = shift;

  my %divides;
  my $i = 0;
  my $p = $primes[$i];
  my $limit = sqrt($n);
  until ($limit < $p) {
    my $k = $n/$p;
    while ($k == int($k)) {
      $n = $k;
      $limit = sqrt($n);
      $divides{$p}++;
      $k = $n/$p;
    }
    $i++;
    if ($i > $#primes) {
      more_primes();
    }
    $p = $primes[$i];
  }
  if (1 < $n) {
    $divides{$n}++;
  }
  return %divides;
}

sub factors {
  my $n = shift;
  my %factor = factor($n);
  my @list = 1;
  for my $p (keys %factor) {
    my $pow = 1;
    my @temp_list;
    for my $c (0..$factor{$p}) {
      push @temp_list, map $pow*$_, @list;
      $pow *= $p;
    }
    @list = @temp_list;
  }
  return sort {$a <=> $b} @list;
}

sub generate_memoryless_stream {
  my $promise = shift;
  my @seen;
  my $is_finished;
  return sub {
    if (@seen) {
      return shift @seen;
    }
    elsif ($is_finished) {
      return;
    }
    else {
      @seen = $promise->();
      if (@seen) {
        return shift @seen;
      }
      else {
        $is_finished = 1;
        return;
      }
    }
  };
}

sub generate_stream {
  my $promise = shift;
  my @seen;
  my $is_finished;
  return sub {
    my $n = shift;
    while ($n > $#seen) {
      if ($is_finished) {
        return;
      }
      my @next = $promise->();
      if (not @next) {
        $is_finished = 1;
        return;
      }
      else {
        push @seen, @next;
      }
    }
    return $seen[$n];
  };
}

sub n_pow_m_mod_k {
  my ($n, $m, $k) = @_;
  my $answer = 1;
  while ($m) {
    if ($m%2) {
      $answer = ($answer * $n) % $k;
    }
    $m = int($m/2);
    $n = ($n * $n) % $k;
  }
  return $answer;
}

package Heap;

sub new {
  my $class = shift;
  my $comp = shift || sub {#print "Comparing @_\n";
    $_[0] cmp $_[1]};
  bless {
    comp => $comp,
    data => [],
  }, $class;
}

sub top {
  my $self = shift;
  return $self->{data}[0];
}

sub append {
  my $self = shift;
  my $data = $self->{data};
  my $comp = $self->{comp};
  for my $e (@_) {
    push @$data, $e;
    my $i = $#$data;
    my $j = int(($i-1)/2);
    #print "Adding $e, ($i, $j)\n";
    while ($i and $comp->(@$data[$i, $j]) < 0) {
      #print "Swapping $i, $j\n";
      @$data[$i, $j] = @$data[$j, $i];
      $i = $j;
      $j = int(($i-1)/2);
    }
    #print "Comp: ", $comp->(@$data->[$i, $j]), $/;
  }
}

sub swap_top {
  my $self = shift;
  my $data = $self->{data};
  my $comp = $self->{comp};
  my $i = 0;
  $data->[0] = shift;
  while (1) {
    last if 2*$i + 1 > $#$data;

    if ($#$data == 2*$i + 1) {
#      print "$i: (@$data) - end is 2*$i+2\n";
      if ($comp->(@$data[$i, 2*$i+1]) > 0) {
#        print "$i: (@$data) - $i > 2*$i+1\n";
        @$data[$i, 2*$i+1] = @$data[2*$i+1, $i];
      }
      last;
    }

    if ($comp->(@$data[2*$i+1, 2*$i+2]) < 0) {
#      print "$i: (@$data) - 2*$i+1 < 2*$i+2\n";
      if ($comp->(@$data[$i, 2*$i+1]) > 0) {
#        print "$i: (@$data) - $i > 2*$i+1\n";
        @$data[$i, 2*$i+1] = @$data[2*$i+1, $i];
        $i = 2*$i+1;
        next;
      }
    }
    else {
#      print "$i: (@$data) - 2*$i+1 >= 2*$i+2\n";
      if ($comp->(@$data[$i, 2*$i+2]) > 0) {
#        print "$i: (@$data) - $i > 2*$i+2\n";
        @$data[$i, 2*$i+2] = @$data[2*$i+2, $i];
        $i = 2*$i+2;
        next;
      }
    }
    last;
  };

  return $data->[0];
}

sub extract {
  my $self = shift;
  my $data = $self->{data};
#  print "Extract: @$data\n";
  my $top = shift @$data;
  if (@$data) {
    unshift @$data, pop @$data;
    $self->swap_top($data->[0]);
  }
  return $top;
}

1;