BeginPackage["ManifoldNDSolve`"]

(#::usage=ToString[#]<>
  " is a option value for Method used by ManifoldNDSolve[].") & /@ {MidPointSolve, TrapezSolve, LobattoIIIaSolve}

ManifoldNDSolve::usage=
"ManifoldNDSolve[eqns,var,{t,t1,t2}, opts] solve the differential and algebraic equations in eqns for vars numerical.\
The algebraic equations may be things like the energy conservation of a system of differential equations.\
ManifoldNDSolve[] has no stepsize control and the number of steps can be set by the MaxSteps option or \
by ManifoldNDSolve[eqns,var,{t,t1,t2,steps}, opts]. Possible integration methods are the implicit \
mid point rule, the trapez rule and a 3 stage Lobatto III a method."

Begin["`Private`"]

MakeIncrementMatrix[manifold_List, vars_List] :=
  Module[{n, m, gg, g, mat, zero},
    n = Length[vars];
    m = Length[manifold];
    mat = IdentityMatrix[n];
    gg = g = Outer[D, manifold, vars]; (* make G[] *)
    zero = Table[0, {m}, {m}];  
    mat = Join[mat, g];                (* Append G[] on the bottom of the unit matrix *)    
    g = Join[-2*Transpose[g], zero];    (* make the - 2 Transpose[G[]]  and append the zeros*)
    {gg, Join @@@ Transpose[{mat, g}]} (* make the large iteration matrix *)
   ]

newtonSolve::newtit = "Implicit solution of does not terminate after `1` steps."

(* symmetric one step integrators *)

MidPointSolve[yn_List, y1_List, h_?NumericQ,
	   f_, jac_, errPhi_?NumericQ, maxit_Integer:$IterationLimit] :=
  Module[{yn1 = y1, h2 = h/2, delta, i = 0, arg},
    While[i++ < maxit,
      arg = (yn + yn1)/2;
      delta = LinearSolve[IdentityMatrix[Length[yn]] - h2*jac[arg], 
      yn1 - yn - h*f[arg]];
      yn1 -= delta;
      If[errPhi > Sqrt[Dot[#, #]] &[delta],
        Return[yn1];
       ]
     ];
    Message[newtonSolve::newtit, i];
    yn1 
   ]

TrapezSolve[yn_List, y1_List, h_?NumericQ,
	   f_, jac_, errPhi_?NumericQ, maxit_Integer:$IterationLimit] :=
  Module[{yn1 = y1, h2 = h/2, delta, fn, i = 0},
    fn = f[yn];
    While[i++ < maxit,
      delta = LinearSolve[IdentityMatrix[Length[yn]] - h2*jac[yn1], 
      yn1 - yn - h2*(f[yn1] + fn)];
      yn1 -= delta;
      If[errPhi > Sqrt[Dot[#, #]] &[delta],
	Return[yn1];
       ]
     ];
    Message[newtonSolve::newtit, i];
    yn1 
   ]

LobattoIIIaSolve[yn_List, y1_List, h_?NumericQ,
    f_, jac_, errPhi_?NumericQ, maxit_Integer:$IterationLimit] :=
  Module[{k1, k2, k3, delta, one, fn, i = 0, n, arg2, arg3, jac2, jac3},
    n = Length[yn];
    one = IdentityMatrix[n];
    k1 = f[yn];
    k3 = f[y1];
    k2 = (k1 + k3)/2;
    With[{a21 = 5*h/24, a22 = h/3, a23 = -h/24,
          a31 = h/6, a32 = 2*h/3,  a33 = h/6,
           b1 = h/6,  b2 = 2*h/3,   b3 = h/6},
      While[i++ < maxit,
        arg2 = yn + a21*k1 + a22*k2 + a23*k3;
        arg3 = yn + a31*k1 + a32*k2 + a33*k3;
        jac2 = jac[arg2];
        jac3 = jac[arg3];
        delta = 
          LinearSolve[   
            Join @@@ 
              Join @@ Transpose /@
                {{one - a22*jac2,     - a23*jac2},
                 {     -a32*jac3, one - a33*jac3}},
             Join[k2 - f[arg2], k3 - f[arg3]]
            ];
        {k2, k3} -= Partition[delta, n];
        If[errPhi > Sqrt[Dot[#, #]] &[delta],
	     Return[yn + b1*k1 + b2*k2 + b3*k3]];
       ]; (* While[] *)
      Message[newtonSolve::newtit, i];
      yn + b1*k1 + b2*k2 + b3*k3
     ] (* Whith[] *)
   ]

(* We must wrap a Hold[] around the Part[new, _] to  
   prevent that the part of a symbol is serached *)
		
spMakeVars[y_List, new_Symbol] :=
    Module[{x, i},
	Thread[y -> Table[Hold[new[[x]]] /. x -> i, {i, Length[y]}]]
      ]
  
ManifoldNDSolve::noexpl = "No explicit solution found for equations `1`."
ManifoldNDSolve::newton = "Newton iteration for projection fail after `1` steps with error `2`."
ManifoldNDSolve::noninit = "Initial values `3` are not numbers at `1`==`2`."
ManifoldNDSolve::invalmeth = "The Method must be one of `1` instead of `2`."
ManifoldNDSolve::inviter = "MaxIterations must be an integer instead of `1`."

spMakeFun[eqn : {_Equal ..}, manifold_List, y_List, t_Symbol, compile_:True] :=
  Module[{sol, deqn, f, jac, mfoldj, gt, rules, newsym, fj},
    (* make explicit differential equations *)
    deqn = Select[eqn, ! FreeQ[#, t] &];
    sol = Solve[deqn, D[#, t] & /@ y];
    If[sol == {} || Head[sol] === Solve,
      Message[ManifoldNDSolve::noexpl, eqn];
      Return[$Failed];
     ];
    f = D[y, t] /. sol; (* explicit solution *)  
    (* make the Jacobian *)   
    jac = Outer[D[#1, #2] &, #, y] & /@ f;
    {gt, mfoldj} = MakeIncrementMatrix[manifold, y];
    newsym = Unique[];
    rules = spMakeVars[y, newsym];
    fj = Join[First /@ {f, jac}, {mfoldj, gt, manifold}] /.  rules;
    If[compile, (* at first we add the function arguments for 
                   compile with an evaluated newsym,              
                   than Hold[] is wrapped and the Hold[part[_]] 
                   are removed with the last Replace[] *)	      
      fj = Map[Hold, Map[{{{newsym, _Real, 1}}, #} & , fj, {1}] , {1}] /. 
          Hold[a_Part] :> a;
                (* Now we can start with compile. 
                   To prevent a early evaluation of the 
                   Part[] function in the Apply[] call 
                   an Unevaluated[] is needed *)
      fj /. Hold[l_] :> Apply[Compile, Unevaluated[l]]
    (* Else *),
      ReleaseHold[Map[Function[Evaluate[{newsym}], #]  &, fj, {1}]]
     ]
   ]

(* Fluid : 
   $f - right hand side function, 
   $jac - Jacobian, 
   $mfoldj - increment matrix for projection, 
   $manifold - manifold, 
   $gt - Jacobian of the manifold,
   $errPhi - the error of the implicit equation,
   $maxit  - maximum Newton iterations *)

SymmetricProjection1[h_, {y0_, fy0_, mu0_}, m_, method_] :=
  Module[{y1, mu, y0hat, y1hat, delta, i = 0},
    y1 = y0 + h*fy0; (* Euler step *)
    mu = mu0;
    While[i++ < $maxit,
      y0hat = y0 + Dot[mu, $gt[y0]];
      y1hat = y1 - Dot[mu, $gt[y1]];
      y1hat -= method[y0hat, y1hat, h, $f, $jac, $errPhi, $maxit];
      delta = Join[y1hat, $manifold[y1]];
      delta = LinearSolve[$mfoldj[y1], - delta];
      y1 += Drop[delta, -m];
      mu += Take[delta, -m];
      If[Sqrt[Dot[delta, delta]] < $errPhi,
        Return[{y1, $f[y1], mu}]
       ];
     ] ;(* While *)
    Message[ManifoldNDSolve::newton, i,Sqrt[Dot[delta, delta]]];
    {y1, $f[y1], mu} 
   ]

Options[ManifoldNDSolve] = {Compiled -> True, WorkingPrecision -> 0.0000001, 
    MaxIterations :> $IterationLimit, Method -> LobattoIIIaSolve, 
    MaxSteps -> 400}

$IntegrationMethods = {MidPointSolve, TrapezSolve, LobattoIIIaSolve}

(* Help function to use one set of initial values y0 *)

ManifoldNDSolveAux[y0_List, vars_, h_?NumericQ, t0_?NumericQ, n_Integer, m_Integer,t_Symbol] :=
  Block[{mu, nsol},
    mu = Table[1, {m}];
    (* We include the first derivative of 
       the solution to get a smaller error
       in the interpolation function *)
    nsol = Drop[#, -1] & /@ 
        NestList[SymmetricProjection1[h, #, m, meth] &, {y0, $f[y0], mu}, n];
    nsol = MapIndexed[{First[#2 - 1]*h + t0, Sequence @@ #1}  &, nsol];
    (* For the interpolation we need lists of {{t, y1[t], y1'[t]} ..} *)
    nsol = Interpolation[#][t] & /@ 
             Transpose[
               nsol /. 
                 {tau_?NumericQ, yast_, dyast_} :> 
                    ({tau, #} & /@ Transpose[{yast, dyast}])];
    Thread[vars -> nsol]
   ]

Clear[ManifoldNDSolve]

ManifoldNDSolve[
    deqn : {_Equal ..}, 
    vars_List, 
    {t_Symbol, t0_?NumericQ, te_?NumericQ, n_Integer}, 
    opts___?OptionQ] :=
  Block[{comp, mfoldlst, h, $errPhi, $maxit, 
         meth, $f, $jac, $mfoldj, $gt, $manifold, y0, m},
    (* Get the equations for the manifolds,  no derivative but t depend *)      
    mfoldlst = Select[deqn, FreeQ[#, Derivative] && ! FreeQ[#, t] &];
    $manifold = mfoldlst /. a_ == b_ :> a - b;
    (* Calculate the step size *)
    h = (te - t0)/n;
    comp = TrueQ[Compiled /. {opts} /. Options[ManifoldNDSolve]];
    $errPhi = N[WorkingPrecision /. {opts} /. Options[ManifoldNDSolve]];
    $maxit = MaxIterations /. {opts} /. Options[ManifoldNDSolve]; 
    If[! IntegerQ[$maxit],
      Message[ManifoldNDSolve::inviter, $maxit];
      Return[$Failed];
     ];
    meth = Method /. {opts} /. Options[ManifoldNDSolve];
    If[! MemberQ[$IntegrationMethods, meth],
      Message[ManifoldNDSolve::invalmeth, $IntegrationMethods, meth];
      Return[$Failed];
     ];
    y0 = vars /. t -> t0;
    y0 =  y0 /. NSolve[
                  Select[deqn, FreeQ[#, t] &], y0, 
                  WorkingPrecision -> 0.01*$errPhi
                 ];
     h = N[h, Max[$MachinePrecision, $errPhi]];
     If[! (And @@ (VectorQ[#, NumericQ] & /@ y0)),
       Message[ManifoldNDSolve::noninit, t, t0, y0];
       Return[$Failed];
      ];
    {$f, $jac, $mfoldj, $gt, $manifold} = 
       spMakeFun[deqn, $manifold, vars, t, comp];
     m = Length[mfoldlst];
    (* Map over multiple solutions of the initial conditions *)
    ManifoldNDSolveAux[#, vars, h, t0, n, m,t] & /@ y0
   ]

ManifoldNDSolve[
  deqn : {_Equal ..}, 
  vars_List, 
  {t_Symbol, t0_?NumericQ, te_?NumericQ}, opts___?OptionQ] :=
  Module[{steps},
    steps = MaxSteps /. {opts} /. Options[ManifoldNDSolve];
    ManifoldNDSolve[deqn, vars, {t, t0, te, steps}, opts]
   ]

End[]
EndPackage[]