import numpy as np

### # This is a smoke test. That fails.
### # A one-digit number with 10 digit choices 
### # should yield k-1.
### k = 10
### m = 1
### N = 1

### # The following gives 126 = 252/2
### # We expect 252 = 10!/(5! 5!)
### # Divided by 2 if ignoring numbers starting with 0
### k = 2
### m = 5
### N = 10


### # Number of outcomes should be 340
### k = 3
### m = 3
### N = 6


### k = 10
### m = 1
### N = 2


# Real problem:
k = 10
m = 3
N = 18


solution_count = 0
factorials = {}


def main():
    global solution_count
    n_tuple = [None,]*k
    recursive_method(n_tuple,1)
    print("Total number of permutations:")
    print("%d"%(solution_count))


def recursive_method( n_tuple, ni ):
    """
    Use recursive backtracking to form all possible 
    combinations of k digits into N-digit numbers 
    such that the number of digits is m or less.

    (n_tuple is actually a list.)

    ni = current class step 1..(k-1)
    n_tuple = list of number of digits for each class 0 through k
    """
    global solution_count, k, m, N
    if(ni==k):

        # N_1 through N_(k-1) have been set,
        # now it is time to set N_0:
        # N_0 = N - N_1 - N_2 - N_3 - .. - N_{k-1}
        sum_N = np.sum([n_tuple[j] for j in range(1,k)])
        n_tuple[0] = max(0, N-sum_N)

        # Compute multiset permutation
        solution_count += multiset(N,n_tuple) - multiset_0(N,n_tuple)

        return

    else:

        # Problem: we are not stopping 
        # when the sum of digits chosen
        # is greater than N
        
        # Assemble the minimum and maximum limits for N_i:
        # (Everything up to ni-1 should be defined, no TypeErrors due to None)
        sum_N = np.sum([n_tuple[j] for j in range(1,ni)])
        ktm = (k - ni)*m
        expr = N - sum_N - ktm
        minn = int(max( 0, expr ))

        # Note: previously this was just maxx=m.
        # This required a check around each call to
        # recursive_method to see if the sum of n_tuple
        # was already maxed out. Now we just do it here.
        maxx = min(m, N-sum_N)

        for N_i in range(minn,maxx+1):

                # Set
                n_tuple[ni] = N_i

                # Explore
                recursive_method(n_tuple, ni+1)

                # Unset
                n_tuple[ni] = None

        return


def multiset(N, n_tuple):
    """
    Number of multiset permutations
    """
    r = factorial(N)/(np.product([factorial(j) for j in n_tuple]))
    return r


def multiset_0(N, n_tuple):
    """
    Number of multiset permutations that start with 0
    """
    if(n_tuple[0]>0):
        r = factorial(N-1)/(np.product([factorial(j-1) if(i==0) else factorial(j) for i,j in enumerate(n_tuple)]))
        return r
    else:
        return 0


def factorial(n):
    """
    Factorial utility
    """
    if(n<0):
        raise Exception("Error: negative factorials not possible")
    if(n==1 or n==0):
        return 1
    else:
        return n*factorial(n-1)


if __name__=="__main__":
    main()