BeginPackage["Schofield`"]

$Assumptions = {θc > 0, θc > 1, gC[_] ∈ Reals, B > 0}

β[D_:2] := Piecewise[
  {
    {1/8, D == 2},
    {0.326419, D == 3},
    {1/2, D == 4},
    {β, True}
  }
]

δ[D_:2] := Piecewise[
  {
    {15, D == 2},
    {4.78984, D == 3},
    {3, D == 4},
    {δ, True}
  }
]

α[D_:2] := Piecewise[
  {
    {0, D == 2},
    {0.11008, D == 3},
    {0, D == 4},
    {α, True}
  }
]

Δ[D_:2] := β[D] δ[D]

os := 2^(1/12) Exp[-1/8] Glaisher^(3/2)

t[θ_] := ((θ/1)^2 - 1)
h[n_][θ_] := (1 - (θ/θc)^2) Sum[gC[i] θ^(2*i+1), {i, 0, n}]
η[g_][θ_] := t[θ] / (g[θ] / I)^(1 / Δ[2])

fLow[B_, θc_][θ_] := (θc Exp[1/(B θc)] ExpIntegralEi[-1/(B θc)] + (θ - θc) Exp[-1/(B (θ - θc))] ExpIntegralEi[1/(B (θ - θc))]) / π

RFLow[B_, θc_][θ_] := fLow[B, θc][θ] + fLow[B, θc][-θ]
RFHigh[ξ0_][ξ_] := (ξ^2+ξ0^2)^(5/6) - (ξ0^2)^(5/6)

RFReg[n_][θ_] := AL fLow[B, θc][-θ] + AH RFHigh[θ0][θ] + Sum[A[i] θ^(2 i), {i, 1, n}]
RF[n_][θ_] := RFReg[n][θ] + AL fLow[B, θc][θ]
dfLow[AL_, B_][m_] := Piecewise[{{AL m! Gamma[m - 1] B^(m - 1) / π, m > 1}, {AL θc Exp[1/(B θc)] ExpIntegralEi[-1/(B θc)] / π, m == 0}, {0, True}}]
dRFc[n_][m_] := dfLow[AL, B][m] + D[RFReg[n][θ], {θ, m}] /. θ -> θc

RFC[n_][θ_] := RF[n][θ] + AL I Sign[Im[θ]] ((θ-θc)Exp[-1/(B(θ-θc))]-(-θ-θc)Exp[-1/(B(-θ-θc))])

invDerivativeList[n_][f_][x_] := Module[
    {xp, dfs, fp, Pns},
  dfs = Rest[NestList[D[#, xp] &, f[xp], n]] /. xp -> x;
  Pns = FoldList[{Pm, m} \[Function]
        fp'[xp] D[Pm, xp] - (2 m - 1) fp''[xp] Pm, 1, Range[n - 1]] /. 
    Derivative[m_][fp][xp] :> dfs[[m]];
  MapIndexed[{Pn, i} \[Function] Pn/dfs[[1]]^(2 i[[1]] - 1), Pns]
  ]

dGdξLowList[n_, h_][m_] := Module[
  { ds, dF, df },
  ds = invDerivativeList[m+1][Function[θ, h[θ] / t[θ]^Δ[2]]][θc];
  dF = NestList[Function[f, D[f, θ]], RFReg[n][θ], m] + Table[dfLow[AL, B][k], {k, 0, m}] /. θ -> θc;
  df = NestList[D[#, \[Theta]] &, 
      fp[\[Theta]]/t[\[Theta]]^2 - 1/(8 \[Pi]) Log[t[\[Theta]]^2], 
      m] /. 
   Map[Derivative[#][fp][\[Theta]] -> dF[[# + 1]] &, 
    Range[0, m]] /. θ -> θc;
  Table[Sum[df[[k+1]] BellY[j, k, ds[[;; j - k + 1]]], {k, 0, j}]/(j!), {j, 0, m}]
]

dGdξList[n_, h_][m_, θp_] := Module[
  { ds, dF, df },
  ds = invDerivativeList[m+1][Function[θ, h[θ] / RealAbs[t[θ]]^Δ[2]]][θp];
  dF = NestList[Function[f, D[f, θ]], RF[n][θ], m] /. θ -> θp;
  df = NestList[D[#, \[Theta]] &, 
      fp[\[Theta]]/t[\[Theta]]^2 - 1/(8 \[Pi]) Log[t[\[Theta]]^2], 
      m] /. 
   Map[Derivative[#][fp][\[Theta]] -> dF[[# + 1]] &, 
    Range[0, m]] /. θ -> θp;
  Table[Sum[df[[k+1]] BellY[j, k, ds[[;; j - k + 1]]], {k, 0, j}]/(j!), {j, 0, m}]
]

dΦdηList[n_, h_][m_, θp_] := Module[
  { ds, dF, df },
  ds = invDerivativeList[m+1][Function[θ, t[θ] / h[θ]^(1 / Δ[2])]][θp];
  dF = NestList[Function[f, D[f, θ]], RF[n][θ], m] /. θ -> θp;
  df = NestList[D[#, \[Theta]] &, 
      hh[θ]^(-2 / Δ[]) (fp[θ] - t[θ]^2 Log[hh[θ]^2] / (8 π Δ[])), 
      m] /. 
   Map[Derivative[#][fp][\[Theta]] -> dF[[# + 1]] &, 
    Range[0, m]] /. hh -> h /. θ -> θp;
  Table[Sum[df[[k+1]] BellY[j, k, ds[[;; j - k + 1]]], {k, 0, j}]/(j!), {j, 0, m}]
]

ruleB[g_] := B -> (2 * os / π) * (- g'[θc] / t[θc]^Δ[2])
ξYL[g_] := g[I θ0]/(-t[I θ0])^Δ[2]/I
AYL[g_] := AH / ((g[I θ0]/ I)^(2 / Δ[2]) * (-η[g]'[I θ0] / (2 θ0 I))^(5/6))
ruleθ0[g_] := ξYL[g] - Around[0.18930, 0.00005]
ruleAL[g_] := AL -> Exp[Δ[2] t[θc]^(Δ[2] - 1) t'[θc] / (2 os / π g'[θc]) - t[θc]^Δ[2] g''[θc] / (4 os / π g'[θc]^2)] t[θc]^(1/8) os / (2 π) * g'[θc]
ruleAH[g_] := AYL[g] + Around[1.37, 0.02]

eqLowRHSReg[n_][m_] := dRFc[n][m]

eqLowLHS[h_][m_] :=D[
  t[θ]^2 (Gl[h[θ] t[θ]^-Δ[2]] + Log[t[θ]^2]/(8 π)),
  {θ, m} ] /. θ -> θc

eqLow[n_, h_][m_] := (eqLowRHSReg[n][m] - eqLowLHS[h][m]) / m!

eqHighRHS[F_][m_] := D[F[θ], {θ, m} ] /. θ -> 0

eqHighLHS[h_][m_] := D[(-t[θ])^2 (Gh[h[θ] (-t[θ])^-Δ[2]] + Log[(-t[θ])^2]/(8 π)), {θ, m} ] /. θ -> 0

eqHigh[n_, h_][m_] := (eqHighRHS[RF[n]][m] - eqHighLHS[h][m]) / m!

eqMid[F_, h_][m_] := D[
  F[θ] - t[θ]^2 Log[h[θ]^2]/(8 Δ[2]π) - h[θ]^((2-α[2])/Δ[2]) Φ[η]
    /. η -> t[θ] / h[θ]^(1 / Δ[2]),
  {θ, m} ] / m! /. θ -> 1

Φs = {
  -1.19773338379799339,
  -0.31881012489061,
  Around[0.110886196683, 2.0 10^-12],
  Around[0.01642689465,  1.0 10^-11],
  Around[-2.639978 10^-4, 1.0 10^-10],
  Around[-5.140526 10^-4, 1.0 10^-10],
  Around[2.08865 10^-4,   1.0 10^-9],
  Around[-4.4819 10^-5,   1.0 10^-9],
  Around[3.194 10^-7,     1.0 10^-9],
  4.31 10^-6,
  -1.99 10^-6
}

Gls = {
  0,
  -os,
  −1.000960328725262189480934955172097320572505951770117 Sqrt[2]/((2 )^(-7/8) (2^(3/16)/os)^2)/2/(12 \[Pi]),
  Around[ 0.038863932, 3.0 10^(-9)],
  Around[−0.068362119, 2.0 10^(-9)],
  Around[ 0.18388371,  1.0 10^(-8)],
  Around[-0.659170,    1.0 10^(-6)],
  Around[ 2.937665,    3.0 10^(-6)],
  Around[-15.61,       1.0 10^(-2)],
  Around[ 96.76,       1.0 10^(-2)],
  Around[-6.79 10^2,   1.0],
  Around[5.34 10^3,    10.],
  Around[-4.66 10^4,   0.01 10^4],
  4.46 10^5,
  -4.66 10^6
}

Ghs = {
  0,
  0,
  -1.000815260440212647119476363047210236937534925597789 Sqrt[2]/((2 )^(-7/8) (2^(3/16)/os)^2)/2,
  0,
  Around[  8.333711750, 5.0 10^(-9)],
  0,
  Around[-95.16897,     3.0 10^(-5)],
  0,
  Around[1457.62,       3.0 10^(-2)],
  0,
  Around[-2.5891 10^4,  2.0],
  0,
  5.02 10^5,
  0,
  -1.04 10^7
}

dRule[sym_][f_, i_] := Derivative[i[[1]] - 1][sym][0] -> f (i[[1]] - 1)!

ΦRules = MapIndexed[dRule[Φ], Φs];
GlRules = MapIndexed[dRule[Gl], Gls];
GhRules = MapIndexed[dRule[Gh], Ghs];

rules[g_] := Join[ΦRules, GlRules, GhRules, {ruleAL[g], ruleB[g], gC[0] -> 1}]

eqAround[n_, g_][m_, p_, q_] := Flatten[Join[{ruleAH[g], ruleθ0[g]}, eqLow[n, g][#] & /@ Range[0, m],eqMid[RF[n], g][#] & /@ Range[0, p], eqHigh[n, g] /@ Range[2, q, 2]]] //. rules[g]

formResiduals[data_, functions_, δ_:10^(-15)] := If[Head[#1]===Around,
  (#2-#1["Value"]) / Max[#1["Uncertainty"], δ],
  (#2-#1) / δ] & @@@ Thread[{data, functions}]

resLow[n_, g_, δ_][m_] := formResiduals[Gls[[;;m+1]], dGdξLowList[n, g][m], δ]

resHigh[n_, g_, δ_][m_] := Rest[formResiduals[Ghs[[;;m+1]], dGdξList[n, g][m, 0], δ][[;;;;2]]]

resMid[n_, g_, δ_][m_] := formResiduals[Φs[[;;m+1]], dΦdηList[n, g][m, 1], δ]

res[F_, g_, δ_][m_] := Join[resLow[F, g, δ][m], resHigh[F, g, δ][m]] //. rules[g]
chiSquared[F_, g_, δ_][m_] := Total[res[F, g, δ][m]^2]
resAll[F_, g_, δ_][m_] := Join[resLow[F, g, δ][m], resHigh[F, g, δ][m], resMid[F, g, δ][m](*, {ruleθ0[g] / 0.00005 /. Around[x_, _] :> x}*)] //. rules[g]

newSol[eqs_, oldSol_, newVars_, δ_:0, γ_:0, opts___] := FindRoot[
  eqs,
  Join[
    {#1, #2 + γ * RandomVariate[NormalDistribution[]]} & @@@ (oldSol /. Rule -> List),
    Thread[{newVars, δ * RandomVariate[NormalDistribution[], Length[newVars]]}]
  ],
  MaxIterations -> 10000,
  opts
]

levenbergMarquardtStep[M_, λ_, g_] := Quiet[
  - LinearSolve[M + λ DiagonalMatrix[Max[10^-5, #] & /@Diagonal[M]], g]
  , LinearSolve::luc
]

levenbergMarquardtAcceleration[M_, λ_, v_, J_, rold_, rf_, h_, β_] := Quiet[
  - Re[LinearSolve[M + λ DiagonalMatrix[Max[10^-5, #] & /@Diagonal[M]], Transpose[J].(2/h((rf[β[[All,2]] + h v] - rold) / h - J . v))] / 2]
  , LinearSolve::luc
]

levenbergMarquardt[r_, β0_, λ0_ : 1, ν_ : 2, ε_ : 10^-15] := 
 Module[
  {
    n = Length[β0],
    rf, Jf
   },
  PrintTemporary["Compiling the Jacobian..."];
  Quiet[
    rf = Compile[{{x, _Real, 1}}, 
      Evaluate[r /. Thread[Rule[First /@ β0, Part[x, #] & /@ Range[n]]]]
    ];
    Jf = Compile[{{x, _Real, 1}}, 
      Evaluate[D[r, {First /@ β0}] /. Thread[Rule[First /@ β0, Part[x, #] & /@ Range[n]]]]
    ];
  , Part::partd];
  levenbergMarquardtHelper[1, rf, Jf, β0, λ0, ν, ε]
 ]


levenbergMarquardtHelper[Δ_, rf_, Jf_, β0_, λ0_ : 1, ν_ : 1.5, ε_ : 10^-15] := 
 Module[
  {
    n = Length[β0],
    β = β0,
    λ = λ0,
    newβ = β0,
    x, oldJ, oldr, newr, M, g, δ, oldC, newC, vlast,
    h = 0.1, α = 0.1, a, v, μ = 5
   },
  oldr = rf[Append[β[[All, 2]], Δ]];
  oldJ = Jf[Append[β[[All, 2]], Δ]];
  oldC = Re[Total[oldr^2]];
  newC = oldC;
  g = Re[Transpose[oldJ] . oldr];
  M = Re[Transpose[oldJ] . oldJ];
  v = levenbergMarquardtStep[M, λ, g];
  a = levenbergMarquardtAcceleration[M, λ, v, oldJ, oldr, rf, h, β, Δ];
  δ = v + a;
  vlast = v;
  PrintTemporary["Current cost value: ", Dynamic[oldC]];
  While[Norm[δ] > ε,
    newβ[[All, 2]] += δ;
    newr = rf[Append[newβ[[All, 2]], Δ]];
    newC = Re[Total[newr^2]];
    While[((1 - v . vlast / Norm[v] / Norm[vlast])newC > oldC) || (2 Norm[a] / Norm[v] > α),
      λ *= ν;
      v = levenbergMarquardtStep[M, λ, g];
      a = levenbergMarquardtAcceleration[M, λ, v, oldJ, oldr, rf, h, β, Δ];
      δ = v + a;
      newβ = β;
      newβ[[All, 2]] += δ;
      newr = rf[Append[newβ[[All, 2]], Δ]];
      newC = Re[Total[newr^2]];
    ];
    oldC = newC;
    oldr = newr;
    β = newβ;
    oldJ = Jf[Append[β[[All, 2]], Δ]];
    g = Re[Transpose[oldJ] . oldr];
    M = Re[Transpose[oldJ] . oldJ];
    vlast = v;
    λ /= μ;
    v = levenbergMarquardtStep[M, λ, g];
    a = levenbergMarquardtAcceleration[M, λ, v, oldJ, oldr, rf, h, β, Δ];
    δ = v + a;
  ];
  {newC, β}
 ]

annealFit[r_, β0_, γ_ : 0, λ0_ : 1, ν_ : 2, μ_ : 3, ε_: 10^-14, h_ : 0.1, α_ : 0.1] :=
Module[ {
    rf, Jf,
    n = Length[β0],
    β = β0,
    λ = λ0,
    newβ = β0,
    x, oldJ, oldr, newr, M, g, δβ, oldC, newC, vlast,
    a, v
    },
  Quiet[
    rf = Compile[{{x, _Real, 1}}, 
      Evaluate[r /. Thread[Rule[First /@ β0, Part[x, #] & /@ Range[n]]]]
    ];
    Jf = Compile[{{x, _Real, 1}}, 
      Evaluate[D[r, {First /@ β0}] /. Thread[Rule[First /@ β0, Part[x, #] & /@ Range[n]]]]
    ];
  , Part::partd];
  If[γ > 0,
  β[[All, 2]] += RandomVariate[NormalDistribution[0, γ], n]];
  newβ = β;
  λ = λ0;
  oldr = rf[β[[All, 2]]];
  oldJ = Jf[β[[All, 2]]];
  oldC = Re[Total[oldr^2]];
  newC = oldC;
  g = Re[Transpose[oldJ] . oldr];
  M = Re[Transpose[oldJ] . oldJ];
  v = levenbergMarquardtStep[M, λ, g];
  a = levenbergMarquardtAcceleration[M, λ, v, oldJ, oldr, rf, h, β];
  δβ = v + a;
  vlast = v;
  While[Norm[δβ] > ε,
    newβ[[All, 2]] += δβ;
    newr = rf[newβ[[All, 2]]];
    newC = Re[Total[newr^2]];
    While[((1 - v . vlast / Norm[v] / Norm[vlast])newC > oldC) || (2 Norm[a] / Norm[v] > α),
      λ *= ν;
      v = levenbergMarquardtStep[M, λ, g];
      a = levenbergMarquardtAcceleration[M, λ, v, oldJ, oldr, rf, h, β];
      δβ = v + a;
      newβ = β;
      newβ[[All, 2]] += δβ;
      newr = rf[newβ[[All, 2]]];
      newC = Re[Total[newr^2]];
    ];
    oldC = newC;
    oldr = newr;
    β = newβ;
    oldJ = Jf[β[[All, 2]]];
    g = Re[Transpose[oldJ] . oldr];
    M = Re[Transpose[oldJ] . oldJ];
    vlast = v;
    λ /= μ;
    v = levenbergMarquardtStep[M, λ, g];
    a = levenbergMarquardtAcceleration[M, λ, v, oldJ, oldr, rf, h, β];
    δβ = v + a;
  ];
  {oldC, β}
]

fitCov[res_, s_] := 
 With[{J = D[res, {s[[All, 1]]}]}, 
  Symmetrize[Inverse[SetPrecision[Transpose[J] . J /. s, 60]]]]

EndPackage[]