#!/usr/bin/perl
use strict;
use warnings;
use PDL;
use PDL::NiceSlice;
use Carp::Assert;

# Compute the partition function P(n,k) using the recurrence
# relation:  P(n,k) = P(n-1,k-1) + P(n-k,k)
# Use this to solve the challenge Q(667,10).
# Q(667,10) = P(667 - 45, 10)

my $kmax = 10;
my $nmax = 667 - 45;

# Final results for unique partitions will be created by
# adding (0..9) to the plain partition results, so restrict
# my maximum extry to 91.
my $entrymax = 100 - 9;

# Space allocation for interspersed search.
my $entrymax2 = $entrymax * 2;

my ($n, $k, %triangle, %p);

sub shiftpdl {
    my ($pdl_a) = @_;
    my $pdl_b = rotate($pdl_a,1)->copy;
    $pdl_b(0) .= 0;
    return $pdl_b;
}

# Generate the partition at the given position number from
# zero to the maximum -1.
sub generate {
    my ($n, $k, $psn) = @_;

    # Stack of "commands" to generate the partition.
    my @commands;
    # Adjustment vector to go from plain partition to distinct values.
    my @adjust = reverse (0..$k-1);
    my $position = pdl($psn+1);  # Add one for vsearch

  GETPATH:
    while (1) {
        # Done when we reach the root.
        last GETPATH if ($n == 1 && $k == 1);

        assert(exists $p{$n, $k}, "p(n,k) must exist");
        my $pos_in_both = vsearch($position, $p{$n, $k})->sclr;

        if ($pos_in_both % 2 == 0) {
            # Even position means going up, to p(n-k,k)
            unshift @commands,0;
            assert($n-$k > 0,"n: $n - $k greater than zero on up move");
            assert(exists $triangle{$n-$k,$k}, "triangle(n-k,k) must exist");

            # Convert position number relative to new location.
            if ($pos_in_both > 0) {
                my $old_offset = $p{$n, $k}->at($pos_in_both-1);
                $position -= $old_offset;
            }
            my $pos = $pos_in_both/2;
            if ($pos > 1) {
                my $accum = cumusumover $triangle{$n-$k,$k};
                my $new_offset = $accum->at($pos-2);  # account for shift
                $position += $new_offset;
            }

            # Prepare for next iteration.
            $n -= $k;
        } else {
            # Odd means going diagonally, to p{$n-1, $k-1}
            unshift @commands,1;
            assert($n > 1 && $k > 1, "n: $n and k: $k both greater than one");
            assert(exists $triangle{$n-1,$k-1}, "triangle(n-1,k-1) must exist");

            # Convert position number relative to new location.
            if ($pos_in_both > 0) {
                my $old_offset = $p{$n, $k}->at($pos_in_both-1);
                $position -= $old_offset;
            }

            my $pos = ($pos_in_both - 1)/2;
            if ($pos > 0) {
                my $accum = cumusumover $triangle{$n-1,$k-1};
                my $new_offset = $accum->at($pos-1);  # account for shift
                $position += $new_offset;
            }

            # Prepare for next iteration.
            $n -= 1;
            $k -= 1;
        }
    }

    # Construct the partition step-by-step from the starting point.
    my @part = (1);
    foreach my $com (@commands) {
        if ($com == 1) {
            # Going from n-1,k-1 to n,k : tack on a 1.
            push @part,1;
        } else {
            # Going from n-k, k to n,k: add one to each item.
            @part = map {$_ + 1} @part;
        }
    }

    # Transform to distinct values before returning.
    @part = map { $part[$_] + $adjust[$_] } (0..@adjust-1);
    return @part;
}

# Initialize corner of the triangle.
my $p1 = zeroes $entrymax;
$p1(0) .= 1;   # Single item with a max entry of one.
$triangle{1,1} = $p1;
$p{1,1} = zeroes $entrymax2;   # Starting point for partition generation.

# Use the recurrence relation to populate the triangle,
# but accumulate for each maximum number separately so
# we can shift off and eliminate cases that would take
# our maximum entry over 100.
for $n (2..$nmax) {
    for $k (1..$kmax) {
        my $psum = zeroes $entrymax;
        my $pboth = zeroes $entrymax2;
        # We'll need totals interspersed for best partition generation.
        my $even = $pboth(0:-1:2);
        my $odd = $pboth(1:-1:2);
        if (exists $triangle{$n-1,$k-1}) {
            # New entries for this case are created by tacking
            # on a one to each result, so the maximum is unchanged.
            my $p1 = $triangle{$n-1,$k-1};
            $odd .= $p1;
            $psum += $p1;
        }
        if (exists $triangle{$n-$k, $k}) {
            # New entries for this case are created by adding
            # a one to each entry, so shift all counts up one place.
            my $p2 = shiftpdl($triangle{$n-$k, $k});
            $even .= $p2;
            $psum += $p2;
        }

        # Interleaves the two sources of entries so we can find
        # them in an interesting order.
        $p{$n,$k} = cumusumover $pboth;
        $triangle{$n,$k} = $psum;
    }
}

my $sum = sumover $triangle{$nmax,$kmax};
print "Total unique partitions Q(677,10) on {1..100} is ",$sum,"\n";

# Now try to generate some partitions.
my @gen0 = generate($nmax,$kmax,0);
print "First one: ",join(q{ },@gen0),"\n";
my @gen1 = generate($nmax,$kmax,1);
print "Second one: ",join(q{ },@gen1),"\n";
my @gen2 = generate($nmax,$kmax,2);
print "Third one: ",join(q{ },@gen2),"\n";
my @genmax = generate($nmax,$kmax,$sum-1);
print "Last one: ",join(q{ },@genmax),"\n";
my @gennxt = generate($nmax,$kmax,$sum-2);
print "Next to last one: ",join(q{ },@gennxt),"\n";
my @genmid = generate($nmax,$kmax,int($sum/2));
print "Middle one: ",join(q{ },@genmid),"\n";