
(* by M. Kauers *)
(* misc tools for studying walks in the quarter plane. *)

(* computes coefficient <t^n>F(t;x,y) (n=0,...,nn) using O(n^2) space for long integers *)
GetCoefficients[stepset_, x_, y_, nn_Integer] :=
  Module[ {c, k, cold, i, j, n, res, pow, points},

    pow[c_, n_Integer] := If[c===0, KroneckerDelta[n, 0], c^n];

    (* cold[[i,j]] = f[n-1;i-1,j-1] *)
    (* c[[i,j]] = f[n;i-1,n-1] *)
    c = {{1}}; 
    res = {1}; 

    Do[
      If[Mod[n,50]==0, Print[n]];
      cold = c;
      c = Table[
        Plus@@(Apply[cold[[##]]&,
          Select[({i,j}-#)& /@ stepset, (1<=#[[1]]<=n && 1<=#[[2]]<=n)&], {1}])
      , {i, 1, n+1}, {j, 1, n+1}];
      AppendTo[res, Sum[c[[i+1,j+1]]pow[x,i]pow[y,j],{i,0,n},{j,0,n}]];
    , {n, 1, nn}];
    ClearAll[cold, c];

    Return[res];
  ];

(* computes coefficient <t^n>F(t;x,y,z) (n=0,...,nn) using O(n^3) space long integers *)
GetCoefficients[stepset_, x_, y_, z_, nn_Integer] :=
  Module[ {c, k, cold, i, j, l, n, res, pow, points},

    pow[c_, n_Integer] := If[c===0, KroneckerDelta[n, 0], c^n];

    (* cold[[i,j,l]] = f[n-1;i-1,j-1,l-1] *)
    (* c[[i,j,l]] = f[n;i-1,j-1,l-1] *)
    c = {{{1}}}; 
    res = {1}; 

    Do[
      If[Mod[n,50]==0, Print[n]];
      cold = c;
      c = Table[
        Plus@@(Apply[cold[[##]]&,
          Select[({i,j,l}-#)& /@ stepset, (1<=#[[1]]<=n && 1<=#[[2]]<=n && 1<=#[[3]]<=n)&], {1}])
      , {i, 1, n+1}, {j, 1, n+1}, {l, 1, n+1}];
      AppendTo[res, Sum[c[[i+1,j+1,l+1]]pow[x,i]pow[y,j]pow[z,l],{i,0,n},{j,0,n},{l,0,n}]];
    , {n, 1, nn}];
    ClearAll[cold, c];

    Return[res];
  ];

(* generates C-code for a program which computes terms for a specific step set *)
CCode3D[stepset_, p_Integer, fname_String, opts___] :=
  Module[ {f, par},

     (* compile as follows:         gcc -O3 -fopenmp <infilename>.c -o <outfilename>.out *)
     (* declare number of threads:  export OMP_NUM_THREADS=2                             *)
     (* run                         ./<outfilename>.out                                  *)
     par = Parallel /. {opts} /. {Parallel -> False};

     (* sequential: gcc -O3  <infilename>.c -o <outfilename>.out *)
     (* ./<outfilename>.out *)

     f = OpenWrite[fname];

     WriteString[f, "/* automatically generated code */\n"];
     WriteString[f, "/* stepset = " <> ToString[stepset] <> " */\n"];

     WriteString[f, "#include <stdio.h>\n"];
     WriteString[f, "#include <stdlib.h>\n"];
     If[par, WriteString[f, "#include <omp.h>\n"]];
     WriteString[f, "typedef unsigned short num;\n"];
     WriteString[f, "typedef unsigned int wnum;\n"];
     WriteString[f, "#define MODULUS " <> ToString[p] <> "\n\n"];

     WriteString[f, "void cnt(long N) {\n"];
     WriteString[f, "  long i, j, k, n, M, s;\n"];
     WriteString[f, "  wnum sum111, sum001, sum010, sum100, sum110, sum101, sum011;\n"];
     WriteString[f, "  num ***new, ***old, ***tmp, *u, *mm, *mo, *mp, *om, *oo, *op, *pm, *po, *pp, *store;\n"];
     WriteString[f, "  M = N + 3;\n"];
     WriteString[f, "  new = malloc(M * sizeof(num **));\n"];
     WriteString[f, "  old = malloc(M * sizeof(num **));\n"];
     WriteString[f, "  for (i = 0; i < M; i++) {\n"];
     WriteString[f, "    new[i] = malloc(M * sizeof(num *));\n"];
     WriteString[f, "    old[i] = malloc(M * sizeof(num *));\n"];
     WriteString[f, "    for (j = 0; j < M; j++) {\n"];
     WriteString[f, "      new[i][j] = malloc(M * sizeof(num));\n"];
     WriteString[f, "      old[i][j] = malloc(M * sizeof(num));\n"];
     WriteString[f, "    }\n"];
     WriteString[f, "  }\n"];
     WriteString[f, "  store = malloc(M * sizeof(num));\n"];

     WriteString[f, "  old[1][1][1] = 1;\n"];
     WriteString[f, "  printf(\"{1, 1, 1, 1, 1, 1, 1, 1}\\n\");\n"];

     WriteString[f, "  for (n = 1; n <= N; n++) {\n"];
     WriteString[f, "    sum111 = 0;\n"];
     WriteString[f, "    for (i = 1; i <= n + 1; i++) {\n"];
     If[par, WriteString[f, "      #pragma omp parallel for private(j, mm, mo, mp, om, oo, op, pm, po, pp, u, k, s)\n"]];
     WriteString[f, "      for (j = 1; j <= n + 1; j++) {\n"];
     WriteString[f, "         mm = old[i-1][j-1]; mo = old[i-1][j]; mp = old[i-1][j+1];\n"];
     WriteString[f, "         om = old[i][j-1];   oo = old[i][j];   op = old[i][j+1];\n"];
     WriteString[f, "         pm = old[i+1][j-1]; po = old[i+1][j]; pp = old[i+1][j+1];\n"];
     WriteString[f, "         u = new[i][j]; s = 0;\n"];
     WriteString[f, "         for (k = 1; k <= n + 1; k++) {\n"];
     WriteString[f, "            u[k] = ("];
     Do[
       Switch[stepset[[i]],
         {-1,-1,-1}, WriteString[f, "((wnum)pp[k + 1])"],
         {-1,-1,0}, WriteString[f, "((wnum)pp[k])"],
         {-1,-1,1}, WriteString[f, "((wnum)pp[k - 1])"],
         {-1,0,-1}, WriteString[f, "((wnum)po[k + 1])"],
         {-1,0,0}, WriteString[f, "((wnum)po[k])"],
         {-1,0,1}, WriteString[f, "((wnum)po[k - 1])"],
         {-1,1,-1}, WriteString[f, "((wnum)pm[k + 1])"],
         {-1,1,0}, WriteString[f, "((wnum)pm[k])"],
         {-1,1,1}, WriteString[f, "((wnum)pm[k - 1])"],

         {0,-1,-1}, WriteString[f, "((wnum)op[k + 1])"],
         {0,-1,0}, WriteString[f, "((wnum)op[k])"],
         {0,-1,1}, WriteString[f, "((wnum)op[k - 1])"],
         {0,0,-1}, WriteString[f, "((wnum)oo[k + 1])"],
         {0,0,0}, WriteString[f, "((wnum)oo[k])"],
         {0,0,1}, WriteString[f, "((wnum)oo[k - 1])"],
         {0,1,-1}, WriteString[f, "((wnum)om[k + 1])"],
         {0,1,0}, WriteString[f, "((wnum)om[k])"],
         {0,1,1}, WriteString[f, "((wnum)om[k - 1])"],

         {1,-1,-1}, WriteString[f, "((wnum)mp[k + 1])"],
         {1,-1,0}, WriteString[f, "((wnum)mp[k])"],
         {1,-1,1}, WriteString[f, "((wnum)mp[k - 1])"],
         {1,0,-1}, WriteString[f, "((wnum)mo[k + 1])"],
         {1,0,0}, WriteString[f, "((wnum)mo[k])"],
         {1,0,1}, WriteString[f, "((wnum)mo[k - 1])"],
         {1,1,-1}, WriteString[f, "((wnum)mm[k + 1])"],
         {1,1,0}, WriteString[f, "((wnum)mm[k])"],
         {1,1,1}, WriteString[f, "((wnum)mm[k - 1])"],
         _, Throw["unexpected step"]
       ];
       If[ i < Length[stepset], WriteString[f, " + "]];
     , {i, 1, Length[stepset]}];
     WriteString[f, ") % MODULUS; \n"];
     WriteString[f, "            s = (s + ((wnum)u[k])) % MODULUS;\n"];
     WriteString[f, "         }\n"];
     WriteString[f, "         store[j] = s;\n"];
     WriteString[f, "      }\n"];
     WriteString[f, "      for (j = 1; j <= n + 1; j++)\n"];
     WriteString[f, "         sum111 = (sum111 + store[j]) % MODULUS;\n"];
     WriteString[f, "    }\n"];
     WriteString[f, "    sum001 = 0; sum010 = 0; sum100 = 0; sum011 = 0; sum101 = 0; sum110 = 0;\n"];
     WriteString[f, "    for (i = 1; i <= n + 1; i++) {\n"];
     WriteString[f, "      for (j = 1; j <= n + 1; j++) {\n"];
     WriteString[f, "         sum011 = (sum011 + ((wnum)new[1][i][j])) % MODULUS;\n"];
     WriteString[f, "         sum101 = (sum101 + ((wnum)new[i][1][j])) % MODULUS;\n"];
     WriteString[f, "         sum110 = (sum110 + ((wnum)new[i][j][1])) % MODULUS;\n"];
     WriteString[f, "      }\n"];
     WriteString[f, "      sum001 = (sum001 + ((wnum)new[1][1][i])) % MODULUS;\n"];
     WriteString[f, "      sum010 = (sum010 + ((wnum)new[1][i][1])) % MODULUS;\n"];
     WriteString[f, "      sum100 = (sum100 + ((wnum)new[i][1][1])) % MODULUS;\n"];
     WriteString[f, "    }\n"];
     WriteString[f, "    printf(\"{%ld, %ld, %ld, %ld, %ld, %ld, %ld, %ld}\\n\", (unsigned long) new[1][1][1], (unsigned long) sum001, (unsigned long) sum010, (unsigned long) sum100, (unsigned long) sum011, (unsigned long) sum101, (unsigned long) sum110, (unsigned long) sum111);\n"];
     WriteString[f, "    tmp = old; old = new; new = tmp;\n"];
     WriteString[f, "  }\n"];
     WriteString[f, "}\n"];

     WriteString[f, "int main(int argc, char *argv[]) {\n"];
     WriteString[f, "  if (argc < 2)\n"];
     WriteString[f, "    abort();\n"];
     WriteString[f, "  cnt(atoi(argv[1]));\n"];
     WriteString[f, "}\n"];
     
     Close[f];
  ];

(* generates C-code for a program which computes terms for a specific step set *)
CCode2D[stepset_, p_Integer, fname_String, opts___] :=
  Module[ {f, par},

     (* compile as follows:         gcc -O3 -fopenmp <infilename>.c -o <outfilename>.out *)
     (* declare number of threads:  export OMP_NUM_THREADS=2                             *)
     (* run                         ./<outfilename>.out                                  *)
     par = Parallel /. {opts} /. {Parallel -> False};

     (* sequential: gcc -O3  <infilename>.c -o <outfilename>.out *)
     (* ./<outfilename>.out *)

     f = OpenWrite[fname];

     WriteString[f, "/* automatically generated code */\n"];
     WriteString[f, "/* stepset = " <> ToString[stepset] <> " */\n"];

     WriteString[f, "#include <stdio.h>\n"];
     WriteString[f, "#include <stdlib.h>\n"];
     If[par, WriteString[f, "#include <omp.h>\n"]];
     WriteString[f, "typedef unsigned short num;\n"];
     WriteString[f, "typedef unsigned int wnum;\n"];
     WriteString[f, "#define MODULUS " <> ToString[p] <> "\n\n"];

     WriteString[f, "void cnt(long N) {\n"];
     WriteString[f, "  long i, j, k, n, M, s;\n"];
     WriteString[f, "  wnum sum11, sum01, sum10, sum00;\n"];
     WriteString[f, "  num **new, **old, **tmp;\n"];
     WriteString[f, "  M = N + 3;\n"];
     WriteString[f, "  new = malloc(M * sizeof(num *));\n"];
     WriteString[f, "  old = malloc(M * sizeof(num *));\n"];
     WriteString[f, "  for (i = 0; i < M; i++) {\n"];
     WriteString[f, "    new[i] = malloc(M * sizeof(num));\n"];
     WriteString[f, "    old[i] = malloc(M * sizeof(num));\n"];
     WriteString[f, "  }\n"];

     WriteString[f, "  old[1][1] = 1;\n"];
     WriteString[f, "  printf(\"{1, 1, 1, 1}\\n\");\n"];

     WriteString[f, "  for (n = 1; n <= N; n++) {\n"];
     WriteString[f, "    sum11 = 0;\n"];
     WriteString[f, "    for (i = 1; i <= n + 1; i++) {\n"];
     WriteString[f, "      s = 0;\n"];
     WriteString[f, "      for (j = 1; j <= n + 1; j++) {\n"];
     WriteString[f, "         new[i][j] = ("];
     Do[
       Switch[stepset[[i]],
         {-1,-1}, WriteString[f, "((wnum)old[i+1][j+1])"],
         {-1,0}, WriteString[f, "((wnum)old[i+1][j])"],
         {-1,1}, WriteString[f, "((wnum)old[i+1][j-1])"],
         {0,-1}, WriteString[f, "((wnum)old[i][j+1])"],
         {0,0}, WriteString[f, "((wnum)old[i][j])"],
         {0,1}, WriteString[f, "((wnum)old[i][j-1])"],
         {1,-1}, WriteString[f, "((wnum)old[i-1][j+1])"],
         {1,0}, WriteString[f, "((wnum)old[i-1][j])"],
         {1,1}, WriteString[f, "((wnum)old[i-1][j-1])"],
         _, Throw["unexpected step"]
       ];
       If[ i < Length[stepset], WriteString[f, " + "]];
     , {i, 1, Length[stepset]}];
     WriteString[f, ") % MODULUS; \n"];
     WriteString[f, "         s = (s + ((wnum)new[i][j])) % MODULUS;\n"];
     WriteString[f, "      }\n"];
     WriteString[f, "      sum11 = (sum11 + s) % MODULUS;\n"];
     WriteString[f, "    }\n"];
     WriteString[f, "    sum01 = 0; sum10 = 0;\n"];
     WriteString[f, "    for (i = 1; i <= n + 1; i++) {\n"];
     WriteString[f, "      sum01 = (sum01 + ((wnum)new[1][i])) % MODULUS;\n"];
     WriteString[f, "      sum10 = (sum10 + ((wnum)new[i][1])) % MODULUS;\n"];
     WriteString[f, "    }\n"];
     WriteString[f, "    printf(\"{%ld, %ld, %ld, %ld}\\n\", (unsigned long) new[1][1], (unsigned long) sum01, (unsigned long) sum10, (unsigned long) sum11);\n"];
     WriteString[f, "    tmp = old; old = new; new = tmp;\n"];
     WriteString[f, "  }\n"];
     WriteString[f, "}\n"];

     WriteString[f, "int main(int argc, char *argv[]) {\n"];
     WriteString[f, "  if (argc < 2)\n"];
     WriteString[f, "    abort();\n"];
     WriteString[f, "  cnt(atoi(argv[1]));\n"];
     WriteString[f, "}\n"];
     
     Close[f];
  ];


(* given initial values and a recurrence, this defines a function c that 
   gives the n-th term of the sequence in linear time. *)
Rec2Def[c_Symbol, data_, rec_, f_[n_]] :=
  Module[ {u0,u1,u2,u3,kkk,myrec,ord},

     ClearAll[c];
     Do[c[kkk]=data[[kkk+1]], {kkk,0,Length[data]-1}];
     ord = Max[Cases[rec, f[n+kkk_.] -> kkk, Infinity]];
     If[ord > Length[data], Throw["insufficient initial values."]];
     myrec = Together[f[n] /. First[Solve[0==rec /. n -> n - ord, f[n]]]];
     myrec = (HornerForm/@Collect[Numerator[myrec], f[_]])/Denominator[myrec] 
        /. {f -> c, n -> kkk};
     u0[u2[c[kkk_Integer], kkk >= Length[data]], u1[c[kkk], u3[myrec]]] 
       /. {u0 -> SetDelayed, u1 -> Set, u2 -> Condition, u3 -> Together};
  ]; 

GetKernelEquation[stepset_, Q_[x_,y_,z_,t_]] :=
  Module[ {k, s, a, b, c, d, e, f, eps},
    s = Plus@@((x^#[[1]]y^#[[2]]z^#[[3]])& /@ stepset);
    k = 1 - t s;
    a = Coefficient[s, z, -1]; b = Coefficient[s, y, -1]; c = Coefficient[s, x, -1];
    d = Coefficient[a, y, -1]; e = Coefficient[a, x, -1]; f = Coefficient[c, y, -1];
    eps = If[ MemberQ[stepset, {-1,-1,-1}], 1, 0 ];
    (1 - t s)Q[x,y,z,t] == 1 - t/z a Q[x,y,0,t] - t/y b Q[x,0,z,t] - t/x c Q[0,y,z, t]
                             + t/y/z d Q[x,0,0,t] + t/x/z e Q[0,y,0,t] + t/x/y f Q[0,0,z,t]
                             - t/x/y/z eps Q[0,0,0,t]    
  ];

