// check general modular polynomial solutions with n variables and degree < n
//
// The polynomial is:
//
//  a[0] + sum i=1...(m-1) { a[i] * prod j=0...(n-1) { x[j]**r[i][j] } } mod p
//
// Usage: gpoly [-a] [-v] [n [p [iter [a0 a1 ... ]]]]
//
// where iter is the number of iterations.  For each iteration, if the
// coefficients are not specified, random coefficients are generated and
// an exhaustive search over the x values is performed, checking for a
// solution (where the polynomial is zero).  If the coefficients are
// specified only one iteration is performed.
//
// If the -a option is specified an exhaustive search over all possible
// coefficients is performed.  If the -v option is not specified the
// coefficients are displayed only if the number of solutions is not 0 mod p.
//
// Default: gpoly 3 5 1
//
// If iter=1 the powers matrix, coefficients, and solutions are all displayed.
//
// If iter>1 and the -v option is not specified, the powers matrix and
// solutions are not displayed, and the coefficients are displayed only
// if the number of solutions is not 0 mod p.
//
#include <stdio.h>
#include <math.h>
#include <string.h>
#include <stdlib.h>
#include <time.h>

// number of combinations of n things taken m at a time
//
double C( int n, int m)
{
  double r = 1;

  while( m > 0) { r *= n; r /= m; --n; --m; }

  return floor( r + 0.5); // round
}

// print an array
//
void aprint( const char *msg, int a[], int n)
{
  printf( "%s = (", msg);
  for( int i = 0; i < n; ++i) printf( "%d%s", a[i], i < n-1 ? "," : "");
  printf( ")\n");
}

// print a matrix
//
void rprint( const char *msg, int m, int n, int r[m][n])
{
  printf( "%s:\n", msg);

  for( int i = 0; i < m; ++i)
  {
    for( int j = 0; j < n; ++j)
      printf( " %2d", r[i][j]);
    putchar( '\n');
  }
}

// sum of the digits in b
//
int sum( int b[], int n)
{
  int r = 0; for( int i = 0; i < n; ++i) r += b[i]; return r;
}

// add 1 to the number represented by b[n-1]...b[0], mod m
// return non-zero if no overflow
//
int inc( int b[], int n, int m)
{
  int c = 1, s; // carry and sum

  for( int i = 0; i < n; ++i)
  {
    s = b[i] + c;
    if( s >= m) { c = 1; s -= m; } else { c = 0; }
    b[i] = s;
  }

  return !c;
}

// initialze an array with random values mod p
//
void rinit( int a[], int n, int p)
{
  for( int i = 0; i < n; ++i)
  {
    a[i] = rand() % p; // general case
    // a[i] = 1 + rand() % (p-1); // all non-zero
  }
}

// the polynomial function
//
int f( int m, int n, int p, int a[m], int x[n], int r[m][n])
{
  int ans = a[0];

  for( int i = 1; i < m; ++i) // for each coefficient
    if( a[i] != 0)
    {
      int t = a[i];
      for( int j = 0; j < n; ++j) // for each x
        for( int k = 0; k < r[i][j]; ++k)
          t *= x[j];
      ans = (ans + t) % p;
    }

  return ans;
}

int main( int argc, char *argv[])
{
  char *progname = argv[0];

  srand(time(0));

  int all = 0, verbose = 0, n = 3, p = 5, iter = 1;

 // check for -a and -v options
 //
  while( argc > 1 && argv[1][0] == '-')
  {
    if( strcmp( argv[1], "-a") == 0)
    {
      all = 1; --argc, ++argv;
    }
    else if( strcmp( argv[1], "-v") == 0)
    {
      verbose = 1; --argc, ++argv;
    }
    else if( strcmp( argv[1], "-av") == 0 || strcmp( argv[1], "-va") == 0)
    {
      all = verbose = 1; --argc, ++argv;
    }
    else
    {
      fprintf( stderr, "%s: unrecognized option: %s\n", progname, argv[1]);
      return 1;
    }
  }

 // check for n, p, and iter arguments
 //
  if( argc > 1) n = atoi( argv[1]);

  if( argc > 2) p = atoi( argv[2]);

  if( !all && argc == 4 ) iter = atoi( argv[3]);

  printf( "n = %d, p = %d, iter = %d\n", n, p, iter);

 // determine how many terms are in the polynomial
 //
 // int m = 1; for( int d = 1; d < n; ++d) m += C(n+d-1,d);
 //
  int m = C(n+n-1,n-1);

  if( all) printf( "exhaustive search: %d coefficients, %g polynomials\n",
    m, pow(p,m) );

  printf( "%g x values\n", pow(p,n));

 // allocate space for coefficient array a and powers matrix r
 //
  int a[m], r[m][n];
  
 // get coefficients from the command line, if present
 //
  if( argc > 4)
  {
    if( all)
    {
      fprintf( stderr, "%s: can't use -a option and specify the coefficients\n",
	progname);
      return 2;
    }

    if( argc - 4 != m)
    {
      fprintf( stderr, "%s: expecting %d coefficients, got %d\n",
        progname, m, argc-4);
      return 3;
    }
    for( int i = 0; i < m; ++i) a[i] = atoi( argv[i+4]);
  }
  else if( all) // prepare to check all coefficients
  {
    memset( a, 0, m*sizeof(int));
  }

 // generate and store the powers matrix r
 //
  int k = 0, x[n]; memset( x, 0, n*sizeof(int));
  do
  {
    if( sum( x, n) < n)
    {
      for( int i = 0; i < n; ++i) r[k][i] = x[i];
      ++k;
    }
  }
  while( inc( x, n, n));

  if( (iter == 1 && !all) || verbose) rprint( "r", m, n, r);

 // check/count solutions
 // note that x is all 0's left over from above final inc() which overflowed
 //
  int s;

  for( k = 0; k < iter; ++k)
  {
    if( !all && argc <= 4) rinit( a, m, p); // random coefficients

    if( (iter == 1 && !all) || verbose) aprint( "a", a, m);

    s = 0;

    do
    {
      if( f( m, n, p, a, x, r) == 0)
      {
        ++s;
        if( (iter == 1 && !all) || verbose ) aprint( "x", x, n);
      }
    }
    while( inc( x, n, p));

    if( iter > 1 && s%p != 0) aprint( "a", a, m);

    if( (iter == 1 && !all) || verbose || s%p != 0)
      printf( "%d solutions, %d%%%d = %d\n", s, s, p, s%p);

    if( all)
    {
      k = -1; // kludge to continue the loop on k

      if( !inc( a, m, p)) break;
    }
  }

  return 0;
}