//=============================================================================
// Ranking and unranking algorithms for binary fixed-density necklaces and
// Lyndon words
//
// Program by Patrick Hartman and Joe Sawada 2016
//=============================================================================
#include <stdio.h>
#include <string.h>
#define MAX_N 63 /* max length of alpha  */

int mu[MAX_N+1] = { 0,1,-1,-1,0,-1,1,-1,0,0,1,-1,0,-1,1,1,0,-1,0,-1,0,
                  1,1,-1,0,0,1,0,0,-1,-1,-1,0,1,1,1,0,-1,1,1,0,-1,
                  -1,-1,0,0,1,-1,0,0,0,1,0,-1,0,1,0,1,1,-1,0,-1,1};

int phi [MAX_N+1] = { 0,1,1,2,2,4,2,6,4,6,4,10,4,12,6,8,8,16,6,18,8,12,
                    10,22,8,20,12,18,12,28,8,30,16,20,16,24,12,36,18,
                    24,16,40,12,42,20,24,22,46,16,42,20,32,24,52,18,
                    40,24,36,28,58,16,60,30,36};

long long int choose[MAX_N+1][MAX_N+1], B[MAX_N+1][MAX_N+1][MAX_N+1];
int D[MAX_N+1], suf[MAX_N+1][MAX_N+1], NECKLACE = 0, LYNDON=0;

//=============================================================================
// Compute choose function up to max using dynamic programming.
//=============================================================================
void ComputeChoose(int max) {
    int n,r;

    for (n=0; n<=max; n++) {
        for (r=0; r<=max; r++) {
            if (r > n) choose[n][r] = 0;
            if (n == r || r == 0) choose[n][r] = 1;
            else choose[n][r] = choose[n-1][r-1] + choose[n-1][r];
        }
    }
}
//=============================================================================
// Print elements in str[1..n].
//=============================================================================
void PrintString(int str[], int n) {

    for (int i=1; i<=n; i++) printf("%d", str[i]);
    printf("\n");
}
//=============================================================================
// Compute D[i] = den(alpha[1..i]) for i=1..n.
//=============================================================================
void ComputeD(int alpha[], int n) {
    int i;

    D[0] = 0;
    for (i=1; i<=n; i++) {
        if (alpha[i] > 0) D[i] = D[i-1] + 1;
        else D[i] = D[i-1];
    }
}
//=============================================================================
// Compute B[t][j][d] for 0<=j<=t<=n and 0<=d<=t using dynamic programming.
//=============================================================================
void ComputeB(int alpha[], int n) {
    int t, j, d;

    B[0][0][0] = 1;
    for (t=1; t<=n; t++) {
        for (d=0; d<=n; d++) {
            B[t][t][d] = 0;
            for (j=t-1; j>=0; j--) {
                if (d - D[j] - 1 < 0) B[t][j][d] = 0;
                else B[t][j][d] = B[t][j+1][d] + (1 - alpha[j+1]) * B[t-j-1][0][d-D[j]-1];
            }
        }
    }
}
//=============================================================================
// Compute suffix function suf[i][j] for 2<=i<=j<=n.
//=============================================================================
void ComputeSuf(int alpha[], int n) {
    int p, i, j;

    for (i=2; i<=n; i++) {
        p = i - 1;
        for (j=i; j<=n; j++) {
            if (alpha[j] > alpha[j-p]) p = j;
            suf[i][j] = j - p;
        }
    }
}
//=============================================================================
// Return longest Lyndon prefix of alpha[1..n].
//=============================================================================
int Lyn(int alpha[], int n) {
    int p=1, i;

    for (i=2; i<=n; i++) {
        if (alpha[i] < alpha[i-p]) return 0;
        if (alpha[i] > alpha[i-p]) p = i;
    }
    return p;
}
//=============================================================================
// Determine whether or not alpha[1..n] is a necklace.
//=============================================================================
int IsNecklace(int alpha[], int n) {
    int p=1, i;

    for (i=2; i<=n; i++) {
        if (alpha[i] < alpha[i-p]) return 0;
        if (alpha[i] > alpha[i-p]) p = i;
    }

    if (n%p == 0) return 1;
    return 0;
}
//=============================================================================
// Assigns a[1..n] to the largest Necklace of length n and density d.
//=============================================================================
void LargestNecklace(int a[], int n, int d) {
    int i,j,pos=1,b[MAX_N];

    if (d == 0) { for (i=1; i<=n; i++) a[i] = 0; }
    else if (d == n) { for (i=1; i<=n; i++) a[i] = 1; }
    else if (d <= n/2) {
        LargestNecklace(b, d, d-(n%d));
        for (i=1; i<=d; i++) {
            for (j=1; j<=n/d-b[i]; j++) a[pos++] = 0;
            a[pos++] = 1;
        }
    }
    else {
        LargestNecklace(b,n-d,n % (n-d));
        for (i=1; i<=n-d; i++) {
            a[pos++] = 0;
            for (j=1; j<=n/(n-d)-1+b[i]; j++) a[pos++] = 1;
        }
    }
}
//=============================================================================
// Returns True (1) if and only if a[1..t] is the prefix of some necklace of
// length n and density d.
//=============================================================================
int IsPrefix(int a[], int t, int n, int d) {
    int i,j,p,d2=0,b[MAX_N];

    // Test density constraint
    for (j=1; j<=t; j++) if(a[j]==1) d2++;
    if (d < d2 || d-d2 > n-t) return 0;

    // Test if a[1..t] is a prenecklace
    p = Lyn(a,t);
    if (p == 0) return 0;

    // Extend prenecklace with additional testing
    for (j=t+1; j<=n; j++) {
        a[j] = a[j-p];
        if (a[j] == 0 && d-(d2+1) >= 0) {
            if (j == n && d-(d2+1) == 0) return 1;
            else if (j <= n) {
                LargestNecklace(b, n-j, d-(d2+1));
                for (i=1; i<j; i++) {
                    if (a[i] > b[i] || i > n-j) break;
                    if (a[i] < b[i]) return 1;
                }
                if (i == j && j <= n-j && b[j] == 1) return 1;
            }
        }
        else d2++;
    }

    // Test if extension is in N(n,d)
    if (n%p == 0 && d2 == d) return 1;

    return 0;
}
//=============================================================================
// Returns the largest necklace less than or equal to a[1..n] in b[1..n]. If
// such a necklace does not exist, the function returns 0.
//=============================================================================
int LN(int a[],int n, int d, int b[]) {
    int j,p,t,d2=0;

    for (j=1; j<=n; j++) {
        b[j] = a[j];
        if (a[j] == 1) d2++;
    }

    // Test if a[1..n] is in N(n,d)
    p = Lyn(a,n);
    if (p>0 && n%p == 0 && d2 == d) return 1;

    t = n;
    while (t > 0) {
        if (b[t] == 1) {
            b[t] = 0;
            if (IsPrefix(b,t,n,d)) break;
        }
        t--;
    }
    if (t == 0) return 0;
    for (j=t+1; j<=n; j++) {
        b[j] = 1;
        if (!IsPrefix(b,j,n,d)) b[j] = 0;
    }
    return 1;
}
//=============================================================================
// Returns the largest Lyndon word less than or equal to a[1..n] in b[1..n].
// If such a string does not exist, the function returns 0.
//=============================================================================
int LL(int a[], int n, int d, int b[]) {
    int j, bb[MAX_N];

    if (LN(a,n,d,b) == 0) return 0;
    if (Lyn(b,n) == n) return 1;

    for (j=1; j<n; j++) bb[j] = b[j];
    bb[n] = 0;
    return LN(bb,n,d,b);
}
//=============================================================================
// Precompute D, B, and suf. Compute cardinality of T_alpha(n,d).
//=============================================================================
long long int ComputeT(int alpha[], int n, int d) {
    long long int total=0;
    int t, j, i, s;

    ComputeD(alpha, n);
    ComputeB(alpha, n);
    ComputeSuf(alpha, n);

    total = Lyn(alpha, n);
    for (t=1; t<=n; t++) {
        for (j=0; j<=n-1; j++) {
            if (t+j <= n) {
                if (alpha[j+1] > 0) {
                    for (i=0; i<=d-D[j]; i++)
                        if (i <= n-t-j)
                            total += B[t-1][0][d-D[j]-i] * choose[n-t-j][i];
                }
            }
            else{
                s = suf[n-t+2][j];
                if (alpha[j+1] > alpha[s+1])
                    total += B[n-j+s][s+1][d-D[j]+D[s]];
            }
        }
    }
    return total;
}
//=============================================================================
// Compute rank of necklace or Lyndon word alpha[1..n] with density d.
// alpha[] is assumed to be in N(n,d) or L(n,d).
//=============================================================================
long long int Rank(int alpha[], int n, int d) {
    int beta [MAX_N];
    long long int r=0;
    int i;

    for (i=1; i<=n; i++) {
        if (n % i == 0 && d % i == 0) {
            if (!LN(alpha, n/i, d/i, beta)) continue;
            if (NECKLACE) r += phi[i] * ComputeT(beta, n/i, d/i);
            else if (LYNDON) r += mu[i] * ComputeT(beta, n/i, d/i);
        }
    }
    return r/n;
}
//=============================================================================
// Set result[1..n] = (str1[1..n] + str2[1..n]) / 2.
//=============================================================================
void SetMidPoint(int result[], int str1[], int str2[], int n) {
    int sum=0, carry=0, i;

    // ensure 0th position is null for last step in addition (final carry)
    str1[0] = str2[0] = 0;

    // set result to midpoint in one pass by accounting for shift during addition
    for (i=n; i>=0; i--) {
        sum = str1[i] + str2[i] + carry;
        if (i < n) result[i+1] = sum % 2;
        if (sum >= 2) carry = 1;
        else carry = 0;
    }
}
//=============================================================================
// Determine the necklace in N(n,d) or Lyndon word in L(n,d) with rank r, and
// return in the array delta[]
//=============================================================================
int Unrank(int r, int n, int d, int delta[]) {
    int alpha [MAX_N], beta [MAX_N], gamma [MAX_N];
    long long int r2=0, i;

    if (n == d && r == 1) {
        if (NECKLACE) {
            for (i=1; i<=n; i++) delta[i] = 1;
            return 1;
        }
        else if (LYNDON) {
            if (n == 1) {
                delta[1] = 1;
                return 1;
            }
            else return 0;
        }
    }
    if (LYNDON && d == 0) return 0;

    alpha[0] = beta[0] = 0;
    for (i=1; i<=n; i++) alpha[i] = 0;
    for (i=1; i<=n; i++) beta[i] = 1;

    // Perform initial check for invalid rank
    if (NECKLACE) LN(beta, n, d, beta);
    if (LYNDON) LL(beta, n, d, beta);

    r2 = Rank(beta, n, d);
    if (r2 < r) return 0;

    for (i=1; i<=n; i++) beta[i] = 1;

    do {
        SetMidPoint(gamma, alpha, beta, n);
        for (i=0; i<=n; i++) delta[i] = gamma[i];

        if (NECKLACE) LN(delta, n, d, delta);
        if (LYNDON) LL(delta, n, d, delta);

        r2 = Rank(delta, n, d);

        if (r2 < r) for(i=0; i<=n; i++) alpha[i] = gamma[i];
        else if (r2 > r) for(i=0; i<=n; i++) beta[i]  = gamma[i];
    } while (r2 != r);

    return 1;
}
//=============================================================================
// Count number of necklaces or Lyndon words of length n and density d with
// prefix alpha[1..pn]
//=============================================================================
long long int CountPrefixNecklaces(int n, int d, int alpha[], int pn) {
    int beta[MAX_N+1];
    long long int max, min;
    int pd=0, i;

    // Checks on length and density
    for (i=1; i<=pn; i++) if (alpha[i] > 0) pd++;
    if (pn > n || pd > d) return 0;
    
    // Special case when prefix has length n
    if (pn == n) {
        if ((NECKLACE && IsNecklace(alpha, n) && pd == d) ||
            (LYNDON   && Lyn(alpha, n) == n   && pd == d)) return 1;
        else return 0;
    }

    // All other cases
    for (i=0; i<=pn; i++) beta[i] = alpha[i]; // Build beta for computing max
    for (i=pn+1; i<=n; i++) beta[i] = 1;

    if (NECKLACE) { 
        if (!LN(beta, n, d, beta)) max = 0;
        else max = Rank(beta, n, d);
    }
    if (LYNDON) {
        if (!LL(beta, n, d, beta)) max = 0; 
        else max = Rank(beta, n, d);
    }
    
    for (i=0; i<=pn; i++) beta[i] = alpha[i]; // Build beta for computing min
    for (i=pn+1; i<=n; i++) beta[i] = 0;

    if (NECKLACE) {
        if (!LN(beta, n, d, beta)) min = 0;
        else min = Rank(beta, n, d);
    }
    if (LYNDON) {
        if (!LL(beta, n, d, beta)) min = 0;
        else min = Rank(beta, n, d);
    }

    return max-min;
}
//=============================================================================
int main() {
    char input[MAX_N+1];
    long long int r,i;
    int n,n2,d,d2,option,alpha[MAX_N+1];

    printf("----------------------------\n");
    printf(" Fixed-density necklaces\n");
    printf("----------------------------\n");
    printf(" 1.  Rank    \n");
    printf(" 2.  Unrank  \n");
    printf(" 3.  Count with given prefix\n");
    printf(" 4.  Exhaustive generation\n");
    printf("----------------------------\n");
    printf(" Fixed-density Lyndon words\n");
    printf("----------------------------\n");
    printf(" 5.  Rank    \n");
    printf(" 6.  Unrank \n");
    printf(" 7.  Count with given prefix \n");
    printf(" 8.  Exhaustive generation\n\n");

    printf("Enter option: ");   scanf("%d", &option);

    if (option < 1  || option > 8) return 0;
    if (option == 1 || option == 2 || option == 3 || option == 4) NECKLACE = 1;
    else LYNDON = 1;
    
    printf("Enter length N: ");   scanf("%d", &n);
    printf("Enter density D: ");  scanf("%d", &d);
    if (d > n) {  printf("Invalid density\n");  return 0; }
    
    ComputeChoose(n);  // Pre-compute binomial co-efficients
    
    // RANK
    if (option == 1 || option == 5) {
        printf("Enter necklace/Lyndon word (eg. 001101): ");  scanf("%s", input);
        n2 = strlen(input);
        
        alpha[0] = 0;
        for (i=1; i<=n2; i++) alpha[i] = input[i-1] - '0';
        d2 = 0;
        for (i=1; i<=n2; i++) if (alpha[i] > 0) d2++;
        
        if (n2 !=n || d2 != d) { printf("Invalid length or density\n"); return 0; }
        
        printf("Rank = %lld\n", Rank(alpha, n, d));
    }
    // UNRANK
    else if (option == 2 || option == 6) {
        printf("Enter rank: ");  scanf("%lld", &r);
        if (!Unrank(r, n, d, alpha)) printf("Invalid rank\n");
        else PrintString(alpha, n);
    }
    // COUNT w given prefix
    else if (option == 3 || option == 7) {
        printf("Enter prefix (eg. 0010): ");  scanf("%s", input);
        n2 = strlen(input);
        alpha[0] = 0;
        for (i=1; i<=n2; i++) alpha[i] = input[i-1] - '0';
        
        printf("%lld\n", CountPrefixNecklaces(n, d, alpha, n2));
    }
    // GENERATE ALL USING RANK/UNRANK
    else {
        i = 1;  printf("\n");
        while (Unrank(i, n, d, alpha)) {
            r = Rank(alpha, n, d);
            if (i != r) { printf("Broken at rank %lld\n", i); return 0; }
            PrintString(alpha, n);
            i++;
        }
        printf("\nTotal = %lld\n", i-1);
    }
}