(*  DiscreteToPDE -- transition from a discrete lattice-based transport model
 *    to a mean-field PDE formulation.
 *  Copyright (C) 2015  Christoph Koutschan
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version, see http://www.gnu.org/licenses/
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *)

(* We introduce the symbol Partial to denote partial derivatives. *)
Format[Partial[v__][expr_]] := Subscript["\[PartialD]", StringJoin @@ (ToString /@ {v})][expr]

(* Convert between our Partial-syntax and Mathematica expression involving Derivative. *)
D2P[expr_, fus_, vs_] :=
  expr //. {
    (f_ /; MemberQ[fus, f])[Sequence @@ vs] -> f,
    Derivative[v__][f_ /; MemberQ[fus, f]][Sequence @@ vs] :>
      (Partial @@ Flatten[MapThread[Table[#1, {#2}]&, {vs, {v}}]])[f]};
P2D[expr_, fus_, vs_] :=
  expr /. f_ /; MemberQ[fus, f] :> f @@ vs //. Partial[v__][a_ /; FreeQ[a, Partial]] :> D[a, v];

(* Expand nested expressions involving our Partial-notation (apply product rule, ect.). *)
ExpandDerivatives[expr_, fus_] :=
  Expand[Expand[expr] //. {
    Partial[][a_] :> a,
    Partial[v___, x_][a_^b_Integer] :> Partial[v][b*a^(b - 1)*Partial[x][a]],
    Partial[v___, x_][a_ * b_] :> Partial[v][a * Partial[x][b] + Partial[x][a] * b],
    Partial[v__][a_ + b_] :> Partial[v][a] + Partial[v][b],
    Partial[v__][a_ /; FreeQ[a, Alternatives @@ fus]] -> 0,
    Partial[v1___][Partial[v2___][a_]] :> (Partial @@ Sort[Join[{v1}, {v2}]])[a]}];

(* DiscreteToPDE using formal Taylor expansions *)
DiscreteToPDE::usage = "DiscreteToPDE[tprobs, cv] derives to a given discrete model the " <>
"corresponding mean-field PDEs. The transition probabilities tprobs must be given in the " <>
"form T[b, {k,l}] == expr, where expr is the probability that an individual of species b " <>
"moves from (x_i,y_j) to (x_{i+k},y_{j+l}). The second argument cv is a list of continuous " <>
"variables to be used in the PDE, e.g., {x, y}. With the option Order -> 2 the order of " <>
"the Taylor expansions can be specified.";
Options[DiscreteToPDE] = {Order -> 2};
DiscreteToPDE[tprobs_List, cv1_, opts:((_Rule|_RuleDelayed)...)] :=
Module[{fus, cv = Flatten[{cv1}],
  d, f, f0, r, tr, sh, eqns = {}},
  If[Not[MatchQ[tprobs, {(T[_, _List] == _)..}]], Throw["Bad input"]];
  r = Order /. {opts} /. Options[DiscreteToPDE];
  fus = Union[#[[1,1]]& /@ tprobs];
  d = Length[cv];
  Do[
    tr = Cases[tprobs, T[fus[[i]], _] == _];
    f0 = fus[[i]] @@ Table[0, {d}];
    sh = Function[{expr, w}, expr /. (f_ /; MemberQ[fus, f])[v__] :> (f @@ ({v} - w))];
    AppendTo[eqns, -Total[Last /@ tr] * f0 + Total[sh[#[[2]] * f0, #[[1,2]]]& /@ tr]];
  , {i, Length[fus]}];

  (* Taylor expansion *)
  eqns = (eqns / h) //. (f_ /; MemberQ[fus, f])[v__] /; {v}.{v} =!= 0 :>
    With[{i = Position[{v}, Except[0], Heads -> False][[1, 1]]},
      Sum[(h*{v}[[i]])^k/k! * Nest[Partial[cv[[i]]], f @@ ReplacePart[{v}, i -> 0], k], {k, 0, r}]];

  (* Omit variables *)
  eqns = eqns /. (f_ /; MemberQ[fus, f])[(0)..] -> f;

  (* Expand derivatives *)
  eqns = D2P[P2D[eqns, fus, cv] //. a_ /; Depth[a] === 6 :> PolynomialMod[Expand[a], h^(r+1)], fus, cv];
  (* Alternatives: *)
  (* eqns = eqns //. a_ /; Depth[a] === 6 :> PolynomialMod[ExpandDerivatives[a, fus], h^(r+1)];*)
  (* eqns = ExpandDerivatives[eqns, fus];*)

  (* Discard all terms with h^i where i >= r. *)
  eqns = PolynomialMod[Expand[eqns], h^r];

  Return[eqns];
];


(* Partial integration procedure *)
(* The algorithm is described in the paper, here we use the same notation. *)
PartialIntegrate::usage = "PartialIntegrate[expr, {f, g, ...}, {x, y, ...}] performs partial integration on expr, " <>
"which involves the unspecified functions f, g, ... in the variables x, y, ... . " <>
"The order in which the functions and the variables are given may influence the result. " <>
"The option Depth -> d specifies that the output has integration depth at most d. " <>
"The option Print -> True can be used for verbose mode.";
Clear[PartialIntegrate];
Options[PartialIntegrate] = {Depth -> Infinity, Print -> False};

PartialIntegrate[expr_, opts:((_Rule|_RuleDelayed)...)] := PartialIntegrate[expr,
  Variables[Cases[expr, Partial[__][a_ /; FreeQ[a, Partial]] :> a, {0, Infinity}]], opts];

PartialIntegrate[expr_, fus:Except[_Rule|_RuleDelayed], opts:((_Rule|_RuleDelayed)...)] :=
  PartialIntegrate[expr, fus, Union[Cases[expr, Partial[a__] :> a, Infinity, Heads -> True]], opts];

PartialIntegrate[expr1_, fus_, x:Except[_Rule|_RuleDelayed], opts:((_Rule|_RuleDelayed)...)] :=
 Module[{expr = expr1, vs, res, dp, pr},
  {dp, pr} = {Depth, Print} /. {opts} /. Options[PartialIntegrate];
  ccnt = 1; (* Counter for integration constants *)

  (* All variables with respect to which a partial derivative occurs. *)
  vs = Union[Flatten[{x, Cases[expr, Partial[a__] :> a, Infinity, Heads -> True]}]];

  (* Rewrite our Partial[x,y,...][f] notation into Mathematica syntax. *)
  expr = expr /. f_ /; MemberQ[fus, f] :> f @@ vs //. Partial[v__][f_ /; FreeQ[f, Partial]] :> D[f, v];

  (* Call the main procedure *)
  res = PartialIntegrate[expr, fus, Flatten[{x}], vs, dp, pr];

  (* Convert back to our Partial notation. *)
  res = res //. (Derivative[a__][f_] @@ vs) :>
    (Partial @@ Flatten[MapThread[Table[#1, {#2}] &, {vs, {a}}]])[f @@ vs];

  (* Pull out constant factors from partial derivatives. *)
  res = res //. Partial[v__][a1_ * a2_ /; FreeQ[a1, Alternatives @@ vs]] :> a1 * Partial[v][a2];
  res = res //. Partial[v__][a_Plus] :>
    With[{cont = Times @@ (#1^#2& @@@ Select[FactorList[a], FreeQ[#, Alternatives @@ vs]&])},
      cont * "Partial"[v][# / cont & /@ a]] /. "Partial" -> Partial;

  (* Omit the dependence on the variables. *)
  res = res /. ((f_ /; MemberQ[fus, f]) @@ vs) -> f;

  Return[res];
];


(* The main procedure for partial integration *)
PartialIntegrate[expr1_, fus_List, xs_List, vs_List, dp_, pr:(True|False)] :=
Module[{f, ff, ff1, g, m, n, p, x, int = 0, expr = expr1, rem = 0, indent, print},
  indent = StringJoin @@ Table["   ", {Count[Stack[], Module]}];
  print = If[pr, Print[indent, ##]&, Null];
  If[expr === 0 || dp === 0 || xs === {}, Return[Factor[expr]]];
  print["Called with ", expr, ", xs = ", xs, ", dp = ", dp];
  x = First[xs];
  p = Position[vs, x][[1, 1]];
  n = Max[Cases[expr, (Derivative[a__][Alternatives @@ fus] @@ vs) :> {a}[[p]], Infinity]];
  print["n = ", n];
  If[n < 0, Return[PartialIntegrate[expr, fus, Rest[xs], vs, dp, pr]]];
  Do[
    print["i = ", i];
    Do[
      f = fus[[j]];
      print["function ", f];
      ff = D[f @@ vs, {x, i}];
      ff1 = D[f @@ vs, {x, i - 1}];
      If[(m = Exponent[expr, ff]) > 1,
        g = Drop[CoefficientList[expr, ff], 2].ff^Range[2, m];
        rem += g;
        expr -= g;
        print["m > 1: rem = ", rem, ", expr = ", expr];
      ];
      m = 1;
      While[m > 0,
        g = Coefficient[expr, ff];
        m = Exponent[g, ff1];
        int += Expand[1/(m + 1)*ff1*g];
        expr = Expand[expr - (ff*g + ff1*D[g, x])/(m + 1)];
        print["int = ", int, ", expr = ", expr];
      ];
    , {j, Length[fus]}];
  , {i, n, 1, -1}];
  Return[If[int =!= 0, Partial[x][PartialIntegrate[int, fus, xs, vs, dp - 1, pr] + 0 * xC[ccnt++]], 0] + 
    PartialIntegrate[rem + expr, fus, Rest[xs], vs, dp, pr]];
];


(* The following procedure is related to the paper:
 *   Fernando, Landman, Simpson:
 *   Nonlinear diffusion and exclusion processes with contact interactions,
 *   Physical Review E 81, 011903 (2010)
 * (in the following, all formula numbers refer to that paper).
 * Implemented is the general model with combined contact-forming or
 * contact-breaking and contact-maintaining interactions, as given in (13).
 * We omit the factor P here as it is irrelevant for the computations.
 * Input:
 *   nn is the set of nearest neighbors, e.g. {{0, 1}, {0, -1}, {1, 0}, {-1, 0}}.
 *   in is the set of interacting neighbors.
 * Output:
 *   set of transition probabilities, in the form required by DiscreteToPDE. *)
DiscreteModel[nn_List, in_List] :=
Module[{v, cmn, cfn, cbn, P, tprobs = {}},
  (* This is the P defined in (5). *)
  P = Function[set, 1 - Times @@ (1 - c @@ # & /@ set)];
  Do[
    v = nn[[i]];
    (* contact-maintaining neighbors *)
    cmn = Intersection[in, v + # & /@ in];
    (* contact-forming neighbors *)
    cfn = Complement[v + # & /@ in, Append[in, {0, 0}]];
    (* contact-breaking neighbors *)
    cbn = Complement[in, v + # & /@ Append[in, {0, 0}]];
    (* Formula (13) *)
    AppendTo[tprobs, T[c, v] == 1 / Length[nn] * (1 - c @@ v) *
      (u * P[cmn] + w * (1 - P[cmn])) *
      (s * (1 - P[cfn]) * (1 - P[cbn]) +
        p * P[cfn] * (1 - P[cbn]) + q * (1 - P[cfn]) * P[cbn] +
        r * P[cfn] * P[cbn])];
  , {i, Length[nn]}];
  Return[tprobs];
];
