(* Define 0^0 to be 1 *)
Clear[power];
power[x_,y_] := Piecewise[{{1,x==0 && y==0}},x^y];

(* Define x/0 to be 0 *)
Clear[divide];
divide[x_,y_] := Piecewise[{{0,y==0}},x/y];

(* Fast computation of x^y mod q *)
Clear[repeatedSquares];
repeatedSquares = Compile[{{x, _Integer}, {y, _Integer}, {q, _Integer}}, Module[{square, product,X,Y, i},
If[y==0, Return[1];];
If[y==1, Return[Mod[x,q]];];

X =Mod[x,q];
Y ={0};
If[CoprimeQ[x,q],Y=IntegerDigits[Mod[y,EulerPhi[q]],2];,Y=IntegerDigits[y,2];];
square = X;
product = 1;
For[i= Length[Y],i≥ 1 && product ≠ 0,i--,
If[Y[[i]] == 1,product = Mod[product*square,q];];
square = Mod[square^2, q];
];
Return[product];
];
-1
];

(* Check whether or not g is a multiplicative generator mod p^r *)
Clear[isGenerator];
isGenerator[g_,p_,r_]:=Module[{factorization, n, q,i,isgenerator},
If[g==0,Return[False];];
n = (p-1)*p^(r-1);
q =p^r;
factorization = FactorInteger[n];
isgenerator=True;
Do[
If[repeatedSquares[g,n/factorization[[i]][[1]],q] == 1, isgenerator=False;];,
{i,1,Length[factorization]}
];
Return[isgenerator];
];

(* Find a generator of the multiplicative group mod p^r *)
Clear[findGenerator];
findGenerator[p_,r_]:= Module[{g},
g = 2;

While[!isGenerator[g, p, 1], g++];
If[r==1, Return[g];];
If[isGenerator[g, p, 2], Return[g];, Return[g+p];];
];

(* Find generators of the multiplicative groups modulo the prime power divisors of q *)
Clear[findGeneratorList];
findGeneratorList[q_]:=Module[{G,factorization,i,p,r},
G={};
factorization = FactorInteger[q];
i=1;
While[i≤ Length[factorization],
p=factorization[[i]][[1]];
r=factorization[[i]][[2]];
AppendTo[G,Piecewise[{{findGenerator[p,r],p≠2},{{-1,5},p==2&&r>2},{-1,p==2&&r==2},{1,p==2&&r==1}}]];
i++;
];
Return[G,Module];
];

(* Solve nA + mB = 1 over integers; returns {A,B} *)
Clear[sumTo1];
sumTo1= Compile[{{n, _Integer}, {m, _Integer}}, Module[{result1, result2, N, M},
Assert[CoprimeQ[n,m]];
N = n;
M = m;
result1 = {1,0};
result2 = {0,1};

While[N≠1 && M ≠1,
If[N < M, result2 = result2 - Floor[M/N]*result1; M = Mod[M, N];,
result1 = result1 - Floor[N/M]*result2; N = Mod[N, M];];
];
If[N==1, Return[result1];, Return[result2];];
];
{1,0}
];

(* List the powers of g mod q *)
Clear[computeTable];
computeTable[g_, q_] := Module[{table, h},
h=g;
table = {1};

While[h≠1, AppendTo[table,h]; h = Mod[h*g,q];];
Return[table, Module];
];

(* Solve g^x = c mod p^a for x quickly. Table1 is computeTable[g^(p^(a-1)), p^a]
and Table2 is computeTable[g^((p-1)*p^(a-1)), p^a]. These are stored outside of the function to avoid redundant computations.
Runs faster than built-in function for large p^a or with repetitive usage. *)
Clear[discreteLog];
discreteLog[p_,a_,g_,c_,Table1_,Table2_]:= Module[{modulus,order, g′,c′,Y, x,l1,l2,j},
modulus =p^a;
Which[
p==2,
Assert[a>2];
order=2^(a-2);
l2=0;
c′=c*(2-Mod[c,4]);
Do[
l2=Piecewise[{{l2,repeatedSquares[Mod[c′*repeatedSquares[5,Mod[-l2,order], modulus],modulus],2^(a-j-3),modulus]==1}},l2+Power[2,j]];,
{j,0,a-3}];
Return[l2,Module];,

p≠2,
order=(p-1)*modulus/p;
l1 = Position[Table1, repeatedSquares[c, order/(p-1), modulus]][[1]][[1]]- 1;
If[a==1,x=l1,
g′=repeatedSquares[g,p-1,modulus];
c′=repeatedSquares[c,p-1,modulus];
l2 = 0;
Do[
l2=l2+Power[p,j]*(Position[Table2,repeatedSquares[Mod[c′*repeatedSquares[g′,Mod[-l2,order],modulus],modulus],p^(a-j-2),modulus]][[1]][[1]]-1);,
{j,0,a-2}];
Y=sumTo1[order/(p-1),p-1];
x=Mod[l1*Y[[1]]*order/(p-1) + l2*Y[[2]]*(p-1),order];
Return[x,Module];
];
];
];

(* Calculate DirichletCharacter(q, n, c) quickly. Here chiInfo = {q, M, G, Table1List, Table2List},
M is a choice of place to send each generator mod q, G is a list of generators mod q, and n is the index of DirichletCharacter that corresponds with this choice of M and G. *)
Clear[χ];
χ[chiInfo_, c_] := Module[{q, M, G, Table1List, Table2List,product,factorization,i, p, a, modulus, order, g, v},
q = chiInfo[[1]]; M= chiInfo[[2]]; G = chiInfo[[3]]; Table1List=chiInfo[[4]]; Table2List=chiInfo[[5]];
If[!CoprimeQ[q,c],Return[0,Module];];
factorization = FactorInteger[q];
product=1;
Do[
p=factorization[[i]][[1]];
a=factorization[[i]][[2]];
modulus=p^a;
Which[
p==2&&a==2,
product=product*Power[-1,M[[i]]*(Mod[c,4]-1)/2];,

p==2&&a>2,
order=2^(a-2);
v=discreteLog[p,a,0,c,Table1List[[i]],Table2List[[i]]];
product=product*Power[-1,M[[i]][[1]]*(Mod[c,4]-1)/2]* Exp[2I π v* M[[i]][[2]] / order];,

p≠2,
order=(p-1)*modulus/p;
g = G[[i]];
v=discreteLog[p,a,g,c,Table1List[[i]],Table2List[[i]]];
product = product*Exp[2I π v* M[[i]] / order];
];,
{i,1,Length[factorization]}];
Return[product];
];

(* Calculate a list of Table1's for each prime power in q *)
Clear[generateTable1List];
generateTable1List[q_, G_] := Module[{Table1List, factorization,i, p, a, modulus, order, g},
Table1List={};
factorization = FactorInteger[q];
Do[
p=factorization[[i]][[1]];
a=factorization[[i]][[2]];
modulus=p^a;
If[p==2,AppendTo[Table1List,{}];,
g=G[[i]];
order=(p-1)*modulus/p;
AppendTo[Table1List,computeTable[repeatedSquares[g, order/(p-1), modulus], modulus]];
];,
{i,1,Length[factorization]}];
Return[Table1List,Module];
];

(* Calculate a list of Table2's for each prime power in q *)
Clear[generateTable2List];
generateTable2List[q_, G_] := Module[{Table2List, factorization,i, p, a, modulus, order, g},
Table2List={};
factorization = FactorInteger[q];
Do[
p=factorization[[i]][[1]];
a=factorization[[i]][[2]];
modulus=p^a;
If[p==2 || a==1,AppendTo[Table2List,{}];,
g=G[[i]];
order=(p-1)*modulus/p;
AppendTo[Table2List,computeTable[repeatedSquares[g,order/p, modulus], modulus]];
];,
{i,1,Length[factorization]}];
Return[Table2List,Module];
];

(* Calculate the i'th derivative of g_K at z directly with a sum *)
Clear[simpleG];
simpleG[K_, z_] := Sum[Exp[k z], {k, 0, K-1}];

simpleG[K_, i_, z_] := Sum[power[k,i] Exp[k z], {k, 0, K-1}];

(* Approximate the l'th derivative of 1/(1+Exp[z]) using the first 2n-l-1 terms of the Taylor expansion*)
Clear[y];
y[l_,z_,n_]:= (-1)^l l!/z^(l+1)-Piecewise[{{1/2,l==0}},0]+Sum[BernoulliB[2i]/(2i) z^(2i-l-1)/(2i-l-1)!,{i,Ceiling[(l+1)/2],n}];

y[l_,z_,n_,Bernoullis_] := (-1)^l l!/z^(l+1)-Piecewise[{{1/2,l==0}},0]+Sum[Bernoullis[[2i]]Product[2i-l +L, {L,0,l-1}]z^(2i-l-1),{i,Ceiling[(l+1)/2],n}];

(* Precompute a table of values of the above function *)
Clear[YDerivatives,NYDerivatives];

YDerivatives[m_,z_,n_, Bernoullis_] := Table[y[l,z,n, Bernoullis],{l,0,m}];

NYDerivatives[K_,m_,z_,n_,precision_, Bernoullis_] := Table[N[divide[y[l,z,n,Bernoullis],power[K,l]],precision],{l,0,m}]

(* Calculate the i'th derivative of g_K at z using y[l,z,n] *)
Clear[g,Ng];

g[K_,i_,z_,n_] :=Piecewise[{{0, K==0}},Exp[K z] Sum[Binomial[i,l]y[l,z,n]K^(i-l),{l,0,i}]-y[i,z,n]];

g[K_,i_,z_,n_,YDerivatives_]:=Piecewise[{{0, K==0}},Exp[K z] Sum[Binomial[i,l]YDerivatives[[l+1]]K^(i-l),{l,0,i}]-YDerivatives[[i+1]]];

(* These two actually are scaled down by (2K)^i to control the size of numbers that are being rounded early *)
Ng[K_,i_,z_,n_,precision_] :=Piecewise[{{0,K==0}},Exp[K z] Sum[N[Binomial[i,l]/2^i, precision]N[y[l,z,n]/K^l, precision],{l,0,i}]-N[y[i,z,n]/(2 K)^i,precision]];

Ng[K_,i_,z_,n_,precision_,NYDerivatives_] :=
Piecewise[{{0,K==0}},Exp[K z] Sum[N[Binomial[i,l]/2^i, precision]NYDerivatives[[l+1]],{l,0,i}]-NYDerivatives[[i+1]]/2^i];

(* Calculate g_K(z, χ, v) by its definition *)
sumChiG[K_,j_, z_,chiInfo_,v_]:= Sum[χ[chiInfo,v+k]power[k,j]Exp[k*z], {k, 0, K-1}];

(* Calculate g_K(z, χ, v) using simpleG[K, z] and Lemma 1.3 in "An alternative to Riemann-Siegel type formulas" *)
Clear[simpleChiG];
simpleChiG[K_, z_, chiInfo_, v_] := Module[{p, a, b},
p=FactorInteger[chiInfo[[1]]][[1]][[1]];
a=FactorInteger[chiInfo[[1]]][[1]][[2]];
b = Ceiling[a/2];
Return[Sum[χ[chiInfo, v+d] Exp[z d] simpleG[Ceiling[(K-d)/p^b], p^b z + I W[v, d, chiInfo]], {d, 0, p^b - 1}], Module];
];

(* Calculate the i'th derivative of g_K(x, χ, v) wrt x at x = z using the product rule and Lemma 1.3 in "An alternative to Riemann-Siegel type formulas" with simpleG *)
simpleChiG[K_, j_, z_, chiInfo_, v_] :=Module[{p, a, b},
p=FactorInteger[chiInfo[[1]]][[1]][[1]];
a=FactorInteger[chiInfo[[1]]][[1]][[2]];
b = Ceiling[a/2];
Return[Sum[Binomial[j,i]χ[chiInfo, v+d]power[d,j-i]Exp[z d] p^(i b)simpleG[Ceiling[(K-d)/p^b],i,p^b z + I W[v, d, chiInfo]], {d, 0, p^b - 1}, {i, 0, j}], Module];
];


(* Calculate the j'th derivative of g_K(x, χ, v) wrt x at x = z using the product rule and Lemma 1.3 in "An alternative ... formulas" with given Gderivatives *)
Clear[chiG, NchiG];

chiG[K_, j_, z_, chiInfo_, v_, Gderivatives_] := Module[{p, a, b},
p=FactorInteger[chiInfo[[1]]][[1]][[1]];
a=FactorInteger[chiInfo[[1]]][[1]][[2]];
b = Ceiling[a/2];
Return[Sum[Binomial[j,i]χ[chiInfo, v+d]power[d,j-i]Exp[z d] p^(i b)Gderivatives[[d+1]][[i+1]], {d, 0, p^b - 1}, {i, 0, j}], Module];
];

(* This one actually computes chiG/v^j to control the size of numbers that are being rounded early, and assumes that the values of Gderivatives have been scaled down by (2H_d)^i *)
NchiG[K_, j_, z_, chiInfo_, v_, Gderivatives_, precision_] := Module[{p, a, b},
p=FactorInteger[chiInfo[[1]]][[1]][[1]];
a=FactorInteger[chiInfo[[1]]][[1]][[2]];
b = Ceiling[a/2];
Return[Sum[Binomial[j,i]χ[chiInfo, v+d]power[d,j-i] Exp[z d] p^(i b)N[power[2Ceiling[(K-d)/p^b],i]/v^j, precision]Gderivatives[[d+1]][[i+1]], {d, 0, p^b - 1}, {i, 0, j}], Module];
];

(* Calculate f_s(z) *)
Clear[simpleF];
simpleF[s_,z_]:=Exp[s z]/(1+z)^s;


(* Calculate j'th derivatives of f_s at 0 for j from 0 to i using the recurrence relation specified in "An alternative ... formulas" *)
Clear[FDerivatives];
FDerivatives[s_, i_] := Module[{currentRow, derivatives, j, l},
If[i<0, Return[{}, Module];];
currentRow = {1};
currentRow = Join[currentRow, ConstantArray[0, i]];
derivatives = {1};
For[j=1, j≤i, j++,
For[l=i-j, l≥0, l--,
currentRow[[l+2]]=(l+1)currentRow[[l+2]]-s Sum[(-1)^α currentRow[[l+1-α]], {α,1,l}];
];
currentRow = Delete[currentRow,1];
AppendTo[derivatives, currentRow[[1]]];
];
Return[derivatives, Module];
];

(* Calculate W_d *)
Clear[W];
W[v_,d_,chiInfo_] :=Module[{L, w,p,a,b},
p=FactorInteger[chiInfo[[1]]][[1]][[1]];
a=FactorInteger[chiInfo[[1]]][[1]][[2]];
b = Ceiling[a/2];
If[!CoprimeQ[v+d, p],Return[0, Module];, Return[-Log[χ[chiInfo, 1+p^b]]I PowerMod[v+d, -1, p^a], Module];
];
];

(* Mods imaginary part by 2Pi so that it lands between -Pi and Pi *)
Clear[roundOff];
roundOff[z_] := Re[z]+I Mod[Im[z],2 Pi, -Pi];

(* This is the block that (v_r)^(-s) B_r(s, χ, m) estimates, computed directly *)
Clear[blockSum];
blockSum[v_,K_,s_,chiInfo_] := Sum[χ[chiInfo,n]/n^s,{n,v,v+K-1}]

(* Calculate B_r(s, χ, m) directly using simpleF and simpleChiG *)
Clear[simpleB];
simpleB[v_,K_,s_,chiInfo_,m_] := Sum[((D[simpleF[s,z],{z,j}]/.z->0)/j!)*(simpleChiG[K,j,-s/v,chiInfo,v]/v^j), {j, 0, m}];

(* Calculate B_r(s, χ, m) using the y function in equation (27) from "An alternative ... formulas" to calculate the i'th derivative of g_K(z) when z is bounded away from 0 and simpleG when z is not *)
Clear[B,NB];

B[v_,K_,s_,chiInfo_,m_,Fderivatives_,n_] := Module[{p,a, b,Zs,Hs,Gderivatives, Bernoullis, YDerivativeTable},
p=FactorInteger[chiInfo[[1]]][[1]][[1]];
a=FactorInteger[chiInfo[[1]]][[1]][[2]];
b=Ceiling[a/2];
Zs = Table[roundOff[-s p^b/v + I W[v,d,chiInfo]], {d, 0, p^b-1}];
Hs = Table[Ceiling[(K-d)/p^b], {d, 0, p^b-1}];
Bernoullis = Table[BernoulliB[i],{i,1,2n}];
YDerivativeTable = Table[Piecewise[{{YDerivatives[m,Zs[[d+1]],n, Bernoullis],Abs[Zs[[d+1]]]>divide[m+1,Hs[[d+1]]]}},{}],{d,0,p^b-1}];
Gderivatives =Table[Piecewise[{{g[Hs[[d+1]],i,Zs[[d+1]],n,YDerivativeTable[[d+1]]],Abs[Zs[[d+1]]]>divide[m+1,Hs[[d+1]]]}},simpleG[Hs[[d+1]],i,Zs[[d+1]]]], {d,0,p^b-1},{i,0,m}];

Return[Sum[(Fderivatives[[j+1]]/j!)(chiG[K, j, -s/v, chiInfo, v, Gderivatives]/v^j),{j,0,m}],Module];
];

NB[v_,K_,s_,chiInfo_,m_,Fderivatives_,n_, precision_] := Module[{p,a, b,Zs,Hs,Gderivatives, Bernoullis, NYDerivativeTable},
p=FactorInteger[chiInfo[[1]]][[1]][[1]];
a=FactorInteger[chiInfo[[1]]][[1]][[2]];
b=Ceiling[a/2];
Zs = Table[roundOff[-s p^b/v + I W[v,d,chiInfo]], {d, 0, p^b-1}];
Hs = Table[Ceiling[(K-d)/p^b], {d, 0, p^b-1}];
Bernoullis = Table[Piecewise[{{0,Mod[i,2]==1}},N[2Zeta[i]/ (2Pi I)^(2n+2),precision]],{i,1,2n}];
NYDerivativeTable = Table[Piecewise[{{NYDerivatives[Hs[[d+1]],m,Zs[[d+1]],n, precision,Bernoullis], Abs[Zs[[d+1]]]>divide[m+1,Hs[[d+1]]]}},{}],{d,0,p^b-1}];
Gderivatives =Table[Piecewise[{{Ng[Hs[[d+1]],i,Zs[[d+1]],n,precision,NYDerivativeTable[[d+1]]],Abs[Zs[[d+1]]]>divide[m+1,Hs[[d+1]]]}},N[Simplify[divide[simpleG[Hs[[d+1]],i,Zs[[d+1]]],power[2 Hs[[d+1]],i]]],precision]], {d,0,p^b-1},{i,0,m}];

Return[Sum[(Fderivatives[[j+1]]/j!)(NchiG[K, j, -s/v, chiInfo, v, Gderivatives, precision]),{j,0,m}],Module];
];


(* Calculate L(χ,s) using simpleB, B, and NB, respectively *)
Clear[simpleL,L,NL];

(* Calculate set of v_r given u = u_0, v = v_0, and M *)
Clear[calcV];
calcV[u_, v_, M_] := Module[{V,tempV},
V={};
tempV=v;
While[tempV<M,
AppendTo[V,tempV];
tempV=Ceiling[tempV/u]+tempV;
];
AppendTo[V,M];
Return[V,Module];
];


simpleL[s_,chiInfo_,u_,v_,M_,m_,precision_] := Module[{V, sum},
Assert[Length[FactorInteger[chiInfo[[1]]]]==1];
V = calcV[u, v, M];
sum = Sum[χ[chiInfo,N]/N^s, {N, 1, v-1}];
sum = sum + Sum[N[simpleB[V[[r]],V[[r+1]]-V[[r]],s,chiInfo,m], precision]/V[[r]]^s, {r, 1, Length[V]-1}];
Return[sum,Module];
];

L[s_,chiInfo_,u_,v_,M_,m_,n_, precision_] := Module[{V, sum, Fderivatives},
Assert[Length[FactorInteger[chiInfo[[1]]]]==1];
V = calcV[u, v, M];
sum = Sum[χ[chiInfo,N]/N^s, {N, 1, v-1}];
Fderivatives =FDerivatives[s,m];
sum = sum + Sum[N[B[V[[r]],V[[r+1]]-V[[r]],s,chiInfo,m, Fderivatives, n], precision]/V[[r]]^s, {r, 1, Length[V]-1}];
Return[sum,Module];
];

NL[s_,chiInfo_,u_,v_,M_,m_,n_, precision_] := Module[{V, sum, Fderivatives, Bs, B2s},
Assert[Length[FactorInteger[chiInfo[[1]]]]==1];
V = calcV[u,v,M];
sum = Sum[χ[chiInfo,N]/N^s, {N, 1, v-1}];
Fderivatives =FDerivatives[s,m];
Bs = Table[NB[V[[r]],V[[r+1]]-V[[r]],s,chiInfo,m, Fderivatives, n, precision]/V[[r]]^s, {r, 1, Length[V]-1}];
sum = sum + Sum[Bs[[r]], {r, 1, Length[V]-1}];
Return[sum,Module];
];

(* Calculate bounds on error terms *)
Clear[ℬ,ϵ,𝒯,𝒜,TotalErrorBound];

ℬ[M_, s_, u_, V_] := Sum[V[[r]]^(-Re[s]) * Min[simpleG[V[[r+1]]-V[[r]],-Re[s]/V[[r]]],Abs[Csc[Re[s]/(2V[[r]])]]],{r,1,Length[V]-1}];

ϵ[m_,s_,u_] := Piecewise[{{(7/2 Exp[78(m+1)/100]/(m+1)^((m+1)/2))(Abs[s]^((m+1)/2)/u^(m+1)), m≤Abs[s]/4}}, 2^m Exp[194Abs[s]/1000]/u^m];

𝒯[M_, m_, s_, u_, V_] := ℬ[M, s, u, V] ϵ[m,s,u];

𝒜[M_, s_] := (Abs[s]+3)/(Re[s] M^(Re[s]));

TotalErrorBound[M_, m_, s_, u_, V_] := 𝒯[M, m, s, u, V]+𝒜[M, s];

(* Initialize example *)
Clear[p,a,b,q,s,G,chiInfo,u,v,M,V,m,n,precision];
p=3;
a=4;
b=Ceiling[a/2];
q=p^a;
s=1/2 + 10I;
G=findGeneratorList[q];
chiInfo={q,{1},G, generateTable1List[q,G],generateTable2List[q,G]};
u=2Max[Ceiling[Sqrt[Abs[s]+3]],6,Re[s]];
v=u*p^b;
M=1000;
V=calcV[u,v,M];
m=15;
n=15;
precision = 50;

N[TotalErrorBound[M,m,s,u,V]]

(* Compute example *)
Clear[temp1,temp2];
temp1=NL[s,chiInfo,u,v,M,m,n,precision] // AbsoluteTiming
temp2=N[blockSum[1,M-1,s,chiInfo],precision] // AbsoluteTiming
Abs[temp1 - temp2]