(* ::Package:: *)

(* ::Title:: *)
(*Restriction and GKZ utils*)


(* ::Section:: *)
(*Misc utils*)


(* ::Text:: *)
(*Make string.*)


MkString[args__] := args // List // Flatten // Map[ToString] // StringJoin


(* ::Text:: *)
(*Convert t123 --> t[1,2,3].*)


sym2ind = RightComposition[
    MkString,
    StringCases[#, Alternatives[LetterCharacter.., DigitCharacter..]]&,
    Map[ToExpression],
    Apply[Construct],
    (*Edge case of no arguments*)
    Replace[x_[] :> x],
    Identity
];


(* ::Text:: *)
(*Convert t[1,2,3] --> t123.*)


ind2sym = MkString /* StringDelete[{"[", "]", ",", " "}] /* ToExpression;


(* ::Text:: *)
(*Isolate patterns.*)


filter[pattern_] := Union @ Flatten @ Cases[#, pattern, -1]&;


(* ::Text:: *)
(*Numerics.*)


numrep[params_] := params // Rule[#, # // Length // RandomPrime[{10, 10^5}, #]^(-1) &]& // Thread // ReplaceAll


(* ::Text:: *)
(*Multiply by xs.*)


Multiply[xs__] := Times[xs, #]&;


(* ::Text:: *)
(*Add xs*)


Add[xs__] = Plus[xs, #]&;


(* ::Text:: *)
(*a+b+... to {a,b,...}.*)


PlusToList[expr_Plus] := expr // Apply[List];
PlusToList[expr_] := expr // List;
PowerToList[Power[expr_, n_Integer]] := ConstantArray[expr^Sign[n], Abs[n]];
PowerToList[expr_] := expr // List;


(* ::Text:: *)
(*a*b*... to {a,b,...}.*)


TimesToList[expr_Times] := expr // Apply[List] // Map[PowerToList] // Flatten;
TimesToList[expr_Power] := expr // PowerToList;
TimesToList[expr_] := expr // List;


(* ::Text:: *)
(*Measure timing (from: https://github.com/magv/alibrary) *)


SetAttributes[TM, HoldFirst];
TM[ex_] := AbsoluteTiming[ex] // (Print[HoldForm[ex], ": ", #[[1]], " sec"]; #[[2]])&;
TM[str_String] := Function[ex, AbsoluteTiming[ex] // (Print[str, ": ", #[[1]], " sec"]; #[[2]])&, HoldFirst];


(* ::Text:: *)
(*Syntactic sugar for Map*)


Map2[f_] = Map[f, #, {2}]&;
Map3[f_] = Map[f, #, {3}]&;
Map4[f_] = Map[f, #, {4}]&;


(* ::Section:: *)
(*GKZ utils*)


(* ::Text:: *)
(*GKZ vars: x[i]. Poly vars: zi.*)


gkzform[poly_Plus,vars_List] := Module[{cr,crx,coeffs,gkzpoly},
	cr = CoefficientRules[poly, vars, "NegativeDegreeLexicographic"];
	crx = Table[cr[[i,1]] -> x[i], {i,Length@cr}];
	coeffs = Thread[(Range[Length@cr] // Map@x) -> cr[[All,2]]];
	(* The 'Reverse' needed because of 'Negative'DegLex*)
	gkzpoly = FromCoefficientRules[crx, Reverse@vars]; 
	
	{gkzpoly , coeffs}
]


(* ::Text:: *)
(*A-matrix for the case of one polynomial in the integrand.*)


Amat[poly_Plus,vars_List] := Module[{togkz,powers},
	togkz = List @@ gkzform[poly, vars] [[1]];
	powers = CoefficientRules[# , vars] [[All,1,1]] &@ togkz;
	
	Join[{ConstantArray[1, Length@togkz]} , Transpose@powers]
]


(* ::Text:: *)
(*A'.*)


Ap[A_] := ReplacePart[
	A,
	{i_ /; i>1 , __} -> 0
]


(* ::Text:: *)
(*List of {xs, dxs} left over after simplex-restriction.*)


simplexvars[A_,sigma_] := Module[{varinds,sigmacomp,vars},
	varinds = Dimensions[A][[2]] // Range;
	sigmacomp = Complement[varinds, sigma];
	vars = varinds // Map[{# // x, # // dx}&] // Map[ind2sym, #, {2,2}]&;
	vars[[sigmacomp]] // Transpose
]


(* ::Text:: *)
(*qs to rs. Use ns = {} if there is no nullspace.*)


rvec[A_List,q_List,ns__List] := Module[{as,rs,rsol},
	If[Length@q != Length@A , Print["Dimension error"],
		as = Transpose@A;
		rs = Table[r@i , {i,Length@as}];
		rsol = Solve[Total[as * rs] == q , rs, Integers] /. C[_] -> 0 // Quiet;
	
		If[Length@rsol == 0, Print["No solution. The column vectors of A do not span Z."], 
			addns[ns, rsol[[1,All,2]]]
		]
	]
]


(* ::Text:: *)
(*Add nullspace to simplify r-vectors.*)


addns[{}, rvec_] := rvec
addns[ns_, rvec_] := Prepend[ns // PadLeft[#, # // Dimensions // # + {0, 1}&]&, Prepend[rvec, 1]] // RowReduce // LatticeReduce // Cases[{1, __}] // Part[#, 1, 2;;]&


(* ::Text:: *)
(*From dx[...] to j[...] notation.*)


dxtoj[vars_] := RightComposition[
	Collect[#, vars, $coeff]&,
	Expand,
	ReplaceAll[$coeff[h_] mon_. :> Times[h, mon // Exponent[#, vars]& // Apply@j]],
	Identity
];

(*dxtoj[var_] := dxtoj[{var}];*)


(* ::Text:: *)
(*Collect dx's.*)


jsimp = Collect[#,j[__],Factor]&


(* ::Text:: *)
(*Maps d_xi \[Bullet] j[...ni...] \[RightArrow] j[...ni+1...]. *)
(*Assumes no x-dependent prefactors.*)


der[js_,dir_Integer] := Module[{len,m},
	len = js // filter[_j] // First // Length;
	If[dir > len, Print["dir > len"],
		m = UnitVector[len, dir];
		js /. j[n__] :> j @@ ({n} + m)
	]
]


(* ::Text:: *)
(*Factoring out matrices for Cext and Cstd*)


coeffarr[dir_,list_] := With[{derdir = der[std,dir]},
	If[Intersection[derdir, list] == {},
		ConstantArray[0, Length /@ {list,list}],
		CoefficientArrays[der[std,dir], list][[2]] // Normal
	]
]


(* ::Text:: *)
(*Exponents*)


btoed[expr__] := Module[{bs},
	bs = expr // Variables // Map@sym2ind // filter[_b] // Map@ind2sym;
	bs = Join[{
		First @ bs -> \[Epsilon]}, 
		Thread[Rest @ bs -> - \[Epsilon] \[Delta]]
	];
	expr /. bs
]


(* ::Text:: *)
(*Rewriting and simplifying differential operators coming from asir*)


niceop[op_,dxs_] := op // btoed // dxtoj[dxs] // jsimp


(* ::Text:: *)
(*Import Macaulay matrix data and replace b1 \[RightArrow] eps and bi \[RightArrow] - eps * del*)


mmimport[mmdata_String] := mmdata // RightComposition[
	Import,
	StringReplace[{"(1)*<<" -> "j[",">>"     -> "]"}],
	ToExpression,
	btoed,
	Factor
]


(* ::Section:: *)
(*Matrix utils*)


(* ::Text:: *)
(*Row reduction and nullspace.*)


deleteZeroRows = Apply[List] /* DeleteCases[x_ /; SameQ[x["Density"], 0.]] /* SparseArray;
rowReduce = RowReduce /* SparseArray /* deleteZeroRows;
getNullSpace = NullSpace /* SparseArray;


(* ::Text:: *)
(*Independent rows.*)


indrows[mat_List] := mat // RightComposition[
	Transpose,
	rowReduce,
	Transpose,
	Apply @ List,
	Position[x_ /; Length[x["NonzeroValues"]] === 1],
	Flatten,
	Rest
]


(* ::Text:: *)
(*Expand a tensor in `var` around `0`, convert it to `j`-notation in `{var, dvar}`*)
(*To restore the original tensor in `j`-notation: `Apply[Dot]`*)


mseries[{var_, dvar_}, max_] := RightComposition[
    Normal,
    Series[#, {var, 0, max}]&,
    Normal,
    dxtoj[{var, dvar}],
    {#, # // filter[_j] // SortBy[#, Last /* Multiply[-1]]&}&,
    {
        # // Apply[CoefficientArrays] // Last,
        #[[2]]
    }&,
    Identity
];

mseries1[var_, max_] := RightComposition[
    Series[#, {var, 0, max}]&,
    Normal,
    dxtoj[{var}],
    {#, # // filter[_j]}&,
    {
        #[[2]],
        # // RightComposition[
            Apply[CoefficientArrays],
            Last,
            (*Bring `j`'s index to the left*)
            Transpose[#, # // Dimensions // Length // Range // List // Cycles]&,
            Identity
        ]
    }&,
    Identity
];


(* ::Text:: *)
(*Conjugate matrix*)


conj[x_] := Inverse[x] . # . x&;


(* ::Text:: *)
(*Gauge transformation*)


gauge[x_, var_] := conj[x] /* Add[-Inverse[x] . D[x, var]];


(* ::Text:: *)
(*Residue*)


res[var_] := mseries1[var, -1] /* (Part[#, 2, 1]&);


(* ::Text:: *)
(*Jordan form*)


jord := (# // JordanDecomposition // Last // Factor // MatrixForm)&


(* ::Text:: *)
(*Integrability check*)


integ[A_,B_][x_,y_] := D[A, y] - D[B, x] + A . B - B . A // RightComposition[
	FullSimplify,
	Flatten,
	ContainsOnly[#,{0}]&
]


(* ::Section:: *)
(*asir script*)


(* ::Text:: *)
(*Mathematica list to asir list*)


asirlist[mmalist_] /; Length@Dimensions@mmalist > 1 := mmalist // RightComposition[
	ToString,
	StringReplace[{ 
		"{"  ->  "[" , 
		"}"  ->  "]" 
	}],
	StringReplace[{
		", [" ->  ",\n\t["  ,
		"[["  ->  "[\n\t["  ,
		"]]"  ->  "]\n];"
	}],
	StringDelete[" "],
	Identity
]

asirlist[mmalist_] /; Length@Dimensions@mmalist == 1 := mmalist // RightComposition[
	ToString,
	StringReplace[{ 
		"{"  ->  "[" , 
		"}"  ->  "];" 
	}],
	StringDelete[" "],
	Identity
]


(* ::Text:: *)
(*Generate asir script for computing D-module basis elements given an R-vector representing an integral basis*)


asirscript[A_, beta_, sigma_, rvec_] /; Length@Dimensions@rvec > 1 := Module[{xs,dxs,strA,strAp,strbeta,strsigma,strrvec},

	{xs, dxs} = 
		simplexvars[A,sigma] // Map[asirlist] // StringDelete[#, ";"]&;
		
	{strA, strAp, strbeta, strsigma, strrvec} = 
		{A, A // Ap, beta, sigma, rvec} // Map[asirlist];
	
	StringJoin[
		"A = " <> strA <> "\n",
		"Ap = " <> strAp <> "\n",
		"Beta = " <> strbeta <> "\n",
		"Sigma = " <> strsigma <> "\n",
		"Rvec = " <> strrvec <> "\n",
		"NumDen = mt_gkz.rvec_red2(Rvec, A, Ap, Beta, Sigma);\n",
		"Num = map(mt_gkz.to_partial3, NumDen[0], " <> xs <> ");\n",
		"Num = map(dp_dtop, Num, " <> dxs <> ");\n",
		"Den = NumDen[1];\n",
		"NumDen = [];\n",
		"for (I=0; I<length(Num); I++) NumDen = append(NumDen, [Num[I]/Den[I]]);\n\n",
		"mkfile(\"d_module_basis.txt\", NumDen);"
	]
]


(* ::Section:: *)
(*Maple*)
(*from https://github.com/magv/alibrary *)

(* Convert the items into a string, and write it into a given file object. *)
WrString[f_, items__] := {items} // Flatten // Map[BinaryWrite[f, # // ToString]&]

(* Convert the items into a string and write it into the file. *)
MkFile[filename_, items__] := Module[{fd},
  (* The BinaryFormat is needed for the BinaryWrite in WrString. *)
  fd = OpenWrite[MkString[filename], BinaryFormat->True];
  If[fd === $Failed, Error["MkFile: failed to open ", filename, " for writing"]];
  WrString[fd, {items}];
  Close[fd];
]

(* Read a Maple file created by 'save(var, "filename")'. Strip
 * the var name, only return the content.
 *)
MapleGet[filename_String] := Module[{text},
    text = ReadString[filename];
    If[text === $Failed,
        Error["! Failed to read from ", filename];
    ];
    FromMaple[text]
]

(* Convert a string in the Maple format into a Mathematica expression. *)
FromMaple[text_String] := text // RightComposition[
    (* Drop the final '\' on a line. *)
    StringReplace[RegularExpression["(?m)\\\\$"] -> ""],
    (* Drop the whitespace. *)
    StringReplace[RegularExpression["\\s"] -> ""],
    (* Drop the enclosing 'var := ' and ';'. *)
    StringReplace[RegularExpression["^[\\w]+:=|[;:]$"] -> ""],
    (* Transform simple indices like 'zeta[2]' into function
     * calls like 'zeta(2)'. Note that if index arguments contain
     * '[]', this regex will fail.
     *)
    StringReplace[RegularExpression["(\\w)\\[([^]]+)\\]"] -> "$1($2)"],
    (* Every other occurrence of '[' and ']' are lists. *)
    StringReplace[{"[" -> "{", "]" -> "}"}],
    (* Maple's array rules: '(1,2)=x' -> '{1,2}->x' *)
    StringReplace[
        "(" ~~ x:(DigitCharacter | ",").. ~~ ")=" ~~ y:Except[","].. ~~ z:("," | "}")
        :>
        "{" ~~ x ~~ "}" ~~ "->" ~~ y ~~ z
    ],
    (* Maple's array dimensions *)
    StringReplace[
        "1.." ~~ x:DigitCharacter.. ~~ ",1.." ~~ y:DigitCharacter..
        :>
        "{" ~~ x ~~ "," ~~ y ~~ "}"
    ],
    StringReplace[
        "array(" ~~ x__ ~~ ")"
        :>
        "array[" ~~ x ~~ "]"
    ],
    ToExpression[#, TraditionalForm, Hold] &,
    ReplaceAll[O->OO],
    ReplaceAll[FromMaple$Map],
    ReleaseHold,
    Identity
];

FromMaple$Map = {
    HoldPattern[psi[x_]] :> PolyGamma[x],
    HoldPattern[psi[n_, x_]] :> PolyGamma[n, x],
    (* Note that HyperInt uses the original Goncharov notation
     * for Li and MZV, unlike HPL/HypExp, which use the reverse
     * one.
     *)
    (*HoldPattern[zeta[n__]] :> MZV[Reverse[{n}]],*)
    HoldPattern[zeta[n__]] :> Mzv @@ Reverse[{n}],
    HoldPattern[polylog[n_, x_]] :> PolyLog[n, x],
    HoldPattern[Complex[yy_]] :> Complex[0, yy],
    HoldPattern[Complex[xx_, yy_]] :> Complex[xx, yy],
    HoldPattern[array[dims_, rules_]] :> SparseArray[rules, dims],
    Nothing
};

(* Save an expression in a format suitable for Maple\[CloseCurlyQuote]s `read()`
 * command.
 *
 * The name of the variable is set automatically to the
 * basename of the file, so `MaplePut[..., "x.mma"]` would
 * result in a variable `x` being defined after `read("x.mma")`
 * is executed.
 *)
MaplePut[expression_, filename_String] :=
    MaplePut[expression, filename, FileBaseName[filename]]

MaplePut[expression_, filename_String, varname_String] := Module[{fd},
    fd = OpenWrite[filename, BinaryFormat -> True];
    If[fd === $Failed, Error["MaplePut: failed to open ", filename, " for writing"]];
    WrString[fd, varname, " := ", expression // ToMaple, ":\n"];
    Close[fd];
]

(* Convert a Mathematica expression into a a string with an
 * equivalent Maple expression. *)
ToMaple[expression_] :=
    ToString[expression /. ToMaple$Map, InputForm] //
        StringReplace[{" " -> "", "[" -> "(", "]" -> ")", "{" -> "[", "}" -> "]"}]

ToMaple$Map = {
    HoldPattern[Log[x_]] :> ln[x],
    HoldPattern[PolyGamma[x_]] :> psi[x],
    HoldPattern[PolyGamma[n_, x_]] :> psi[n, x],
    HoldPattern[PolyLog[n_, x_]] :> polylog[n, x],
    (* There's no Nielsen polylog on the Maple side; we'll convert
     * it into 'Hpl' from HyperInt.
     *)
    HoldPattern[PolyLog[n_, p_, x_]] :> Hpl[x, Join[Table[0, n], Table[1, p]]],
    (* Convert 'Zeta', 'HPL', 'MZV' and 'Mzv' into HyperInt equivalents.
     *)
    (*HoldPattern[Zeta[n_]] :> zeta[n],*)
    HoldPattern[Zeta[n_]] :> Hpl[1, {n}],
    HoldPattern[HPL[w_List, x_]] :> Hpl[x, w],
    (* Note that HyperInt uses the original Goncharov notation
     * for Li and MZV, unlike HPL/HypExp, which use the reverse
     * one.
     *)
    (* HoldPattern[MZV[n_List]] :> zeta @@ Reverse[n] *)
    HoldPattern[MZV[{w__}]] :> MzvToHpl[Mzv[w]],
    HoldPattern[z_Mzv] :> MzvToHpl[z]
};


