(* :Title: CompartmentalModeling.wl *)

(* :Author: Robert B. Nachbar *)

(* :Summary:
   Functions for constructing, visualizing, and simulating compartmental models.
   Application areas supported: epidemiology, population ecology, pharmacokinetics, 
   chemical kinetics.
 *)

(* :Copyright: (c) 2020 by Robert B. Nachbar *)

(* :Package Version: 1.0 *)

(* :Mathematica Version: 12.0 *)

(* :History:
   1.0 14-Mar-2020
*)

(* ::Package:: *)

BeginPackage["CompartmentalModeling`"]
ClearAll["`*"]

Transition::usage = "Transition[c1, r, c2] specifies a Transition from compartment \
c1 to compartment c2 with with transmission rate r. In the notebook FrontEnd, a \
Transition is rendered as c1 \!\(\*OverscriptBox[\(\[RightArrow]\), \(r\)]\) c2. The same \
form can be used as input where the arrow is \\[RightArrow].
Transition[\[EmptySet], r, c2] specifies a source (birth) to compartment c2 with transmission \
rate r, and \[EmptySet] is the \\[EmptySet] character. In the notebook FrontEnd, a source \
Transition is rendered as \!\(\*OverscriptBox[\(\[RightTeeArrow]\), \(r\)]\) c2. The same \
form can be used as input where the arrow is \\[RightTeeArrow].
Transition[c1, r, \[EmptySet]] specifies a sink (death) from compartment c1 with transmission \
rate r, and \[EmptySet] is the \\[EmptySet] character. In the notebook FrontEnd, a sink \
Transition is rendered as c1 \!\(\*OverscriptBox[\(\[RightArrowBar]\), \(r\)]\). The same \
form can be used as input where the arrow is \\[RightArrowBar]."; 

CompartmentalModelGraph::usage = "CompartmentalModelGraph[model] gives a \
Graph of model comprised of Transitions after converting to a Petri net. \
Vertices representing compartments are colored blue, and those representing \
Transitions are colored red."; 

ResolveCompartmentalModel::usage = "ResolveCompartmentalModel[model] returns a \
list of the independent components of the input compartmental model that is \
expressed as a list of transitions, which are considered coupled when they have \
compartments and/or symbolic rates in common."; 

KineticCompartmentalModel::usage = "KineticCompartmentalModel[model, t] returns a \
list {v, c, S, r, e} of the variable vector (v), component transition vector (c), \
stoichiotmetric matrix (S), rate vector (r), and vector of ordinary differential \
equations (e) for the set of transitions in model; the argument t is the time \
variable. The length of the component transition vector c may be longer than the \
input model because equilibria are decomposed into their forward and reverse reactions."; 

(* messages *)

KineticCompartmentalModel::ifcomp = "Warning: the compartment `1` appears to be \
ill-formed."; 

KineticCompartmentalModel::xtrm = "Warning: unexpected product of compartments `1`."; 

KineticCompartmentalModel::ntzr = "Warning: compartment `1` has a net contribution \
of zero (`2`)."; 

KineticCompartmentalModel::subsys = "Warning: the input consists of `1` independent \
subsystems."; 

Begin["`Private`"]
ClearAll["`*"]

(* Transition *)

Transition /: MakeBoxes[Transition[Global`\[EmptySet], rate_, c2_], 
		form:StandardForm|TraditionalForm] := 
	RowBox[{OverscriptBox["\[RightTeeArrow]", MakeBoxes[rate, form]], 
		MakeBoxes[c2, form]}]
Transition /: MakeBoxes[Transition[c1_, rate_, Global`\[EmptySet]], 
		form:StandardForm|TraditionalForm] := 
	RowBox[{MakeBoxes[c1, form], 
		OverscriptBox["\[RightArrowBar]", MakeBoxes[rate, form]]}]
Transition /: MakeBoxes[Transition[c1_, rate_, c2_], 
		form:StandardForm|TraditionalForm] := 
	RowBox[{MakeBoxes[c1, form], 
		OverscriptBox["\[RightArrow]", MakeBoxes[rate, form]], 
		MakeBoxes[c2, form]}]

MakeExpression[RowBox[{OverscriptBox["\[RightTeeArrow]", r_], c2_}], 
		StandardForm]:=
	MakeExpression[RowBox[{"Transition", "[", 
		RowBox[{Global`\[EmptySet], ",", r, ",", c2}], 
		"]"}], StandardForm]

MakeExpression[RowBox[{c1_, OverscriptBox["\[RightArrowBar]", r_]}], 
		StandardForm]:=
	MakeExpression[RowBox[{"Transition", "[", 
		RowBox[{c1, ",", r, ",", Global`\[EmptySet]}], 
		"]"}], StandardForm]

MakeExpression[RowBox[{c1_, OverscriptBox["\[RightArrow]" | "\[Rule]", r_], c2_}], 
		StandardForm] := 
	MakeExpression[RowBox[{"Transition", "[", 
		RowBox[{c1, ",", r, ",", c2}], 
		"]"}], StandardForm]

MakeExpression[RowBox[{c1_, 
		UnderoverscriptBox["\[Equilibrium]" | "\[LeftRightArrow]", r2_, r1_], c2_}], 
		StandardForm] := 
	MakeExpression[RowBox[{"{", RowBox[{
		RowBox[{"Transition", "[", 
			RowBox[{c1, ",", r1, ",", c2}], 
			"]"}], 
		",", 
		RowBox[{"Transition", "[", 
			RowBox[{c2, ",", r2, ",", c1}], 
			"]"}]
		}], "}"}], StandardForm]

(* CompartmentalModelGraph *)

$C=.;
$R=.;
vsf[style_][{x_, y_}, $C[s_], {w_, h_}] := 
	Inset[Framed[Style[s, Black, Sequence @@ style], 
			FrameStyle -> Blue, 
			Background -> LightBlue], 
		{x, y}]
vsf[style_][{x_, y_}, $R[_, Global`\[EmptySet]][r_], {w_, h_}] := 
	Inset[Framed[Style[r, Black, Sequence @@ style], 
			RoundingRadius -> 25, 
			FrameStyle -> Red, 
			Background -> LightRed], 
		{x, y}]
vsf[style_][{x_, y_}, $R[_][r_], {w_, h_}] := 
	Inset[Framed[Style[r, Black, Sequence @@ style], 
			FrameStyle -> Red, 
			Background -> LightRed], 
		{x, y}]

Options[CompartmentalModelGraph] = Join[
	{}, 
	Options[Graph]
	]; 

CompartmentalModelGraph[model_, options:OptionsPattern[]] := 
	Module[{transitions, edges, vertices, i = 0, baseStyle}, 
		baseStyle = {OptionValue[BaseStyle]} /. Automatic -> {} // Flatten; 

		transitions = Union @ model; 
		
		edges = transitions /. 
			{
			Transition[a_, r_, b_] :> 
				(Transition[a //. {Plus -> List, n_Integer c_ :> Table[c, n]}, 
					r, 
					b //. {Plus -> List, n_Integer c_ :> Table[c, n]}] /. l_List :> Flatten[l])
			} /. 
			{
			Transition[Global`\[EmptySet], r_, b_] :> 
				Thread @ DirectedEdge[$R[++i, Global`\[EmptySet] (* needed for esf for source *)][r], Thread @ $C[b]], 
			Transition[a_, r_, Global`\[EmptySet]] :> 
				Thread @ DirectedEdge[Thread @ $C[a], $R[++i][r]],
			Transition[a_, r_, b_] :> 
				Thread /@ {DirectedEdge[Thread @ $C[a], $R[++i][r]], 
					DirectedEdge[$R[i][r], Thread @ $C[b]]}
			} // Flatten (* // Echo*); 
		
		vertices = Union @@ List @@@ edges(* // Echo*); 
		
(* TODO: check that root vertex is present *)

		Graph[vertices, edges, 
			VertexShapeFunction -> vsf[baseStyle], 
			PerformanceGoal -> "Quality", 
			FilterRules[{options}, Options[Graph]] /. 
				("RootVertex" -> root_) :> ("RootVertex" -> FirstCase[vertices, $C[root] | $R[__][root]])
			]
		]

(* ResolveCompartmentalModel *)

Options[ResolveCompartmentalModel] = {}; 

ResolveCompartmentalModel[OptionsPattern[]] /; Message[ResolveCompartmentalModel::argx, ResolveCompartmentalModel, 0, 1] := Null

ResolveCompartmentalModel[r:(_Transition), options:OptionsPattern[]] := 
	ResolveCompartmentalModel[{r}, options]

ResolveCompartmentalModel[model:{(_Transition)...}, options:OptionsPattern[]] := 
	Module[{transitions, transitionEdges, edges, graph, componentEdges}, 
		transitions = Flatten @ model; 
		edges = transitions /. 
			{
			Transition[a_, r_, b_] :> 
				(Transition[a //. {Plus -> List, n_Integer c_ :> Table[c, n]}, 
					r, 
					b //. {Plus -> List, n_Integer c_ :> Table[c, n]}] /. l_List :> Flatten[l])
			} /. 
			{
			Transition[c1_, r_, c2_] :> {Thread @ DirectedEdge[c1, r], Thread @ DirectedEdge[r, c2]}
			}; 
		transitionEdges = Thread[{transitions, Flatten /@ edges}](* // Echo*); 
		edges = Flatten @ edges // DeleteCases[#, DirectedEdge[Global`\[EmptySet], _] | DirectedEdge[_, Global`\[EmptySet]]]&; 
		graph = Graph[edges]; 
		componentEdges = EdgeList /@ WeaklyConnectedGraphComponents[graph](* // Echo*); 
		Cases[transitionEdges, 
			{trns_, edg_} /; (MatchQ[Intersection[edg, #], {__}]) :> trns]& /@ 
			componentEdges
		]

(* KineticCompartmentalModel *)

compartmentQ[s_] := MatchQ[s, 
	_Symbol | 
	_String | 
	Subscript[__] | 
	Superscript[__] | 
	Subsuperscript[__] | 
	CenterDot[__] | 
	SuperStar[_] | 
	SuperDagger[_] | 
	SuperPlus[_] | 
	SuperMinus[_] | 
	_String[__?AtomQ] | 
	System`C[_Integer] | 
	(h_Symbol[__?AtomQ] /; Names["System`"<>ToString[h]]==={})
	]

Options[KineticCompartmentalModel] = {
	}; 

KineticCompartmentalModel[_, OptionsPattern[]] /; 
	Message[KineticCompartmentalModel::argr, KineticCompartmentalModel, 2] := Null

KineticCompartmentalModel[trans:(_Transition), t_, options:OptionsPattern[]] := 
	KineticCompartmentalModel[{trans}, t, options]

KineticCompartmentalModel[transitionList:{(_Transition)...}, t_Symbol, options:OptionsPattern[]] := 
	With[{attemptEvaluation=iKineticCompartmentalModel[transitionList, t, options]}, 
		attemptEvaluation /; MatchQ[attemptEvaluation, Except[$Failed]]
		]

iKineticCompartmentalModel[transitionList:{(_Transition)...}, t_Symbol, 
		options:OptionsPattern[KineticCompartmentalModel]] := 
	Module[{transitions, reactants, rates, products, 
			stoichiometry, crossTerms, crossTermStoichiometry, 
			inputOrders, defaultOrders, orders, orderRules, 
			terms, lhs, rhs, optlist, 
			ok, 
			compartments, ignore,
			vars, eqns, flux}, 
		optlist = {options} // Flatten;

		ok = True; 
		Scan[If[FreeQ[Keys @ Options[KineticCompartmentalModel], Keys @ #], 
				Message[KineticCompartmentalModel::optx, #, 
					HoldForm[KineticCompartmentalModel[transitionList, t, options]]]; 
				ok=False]&, 
			optlist];
		If[!ok, Return[$Failed]]; 

		With[{n=Length[ResolveCompartmentalModel[transitionList]]},
			If[n > 1, Message[KineticCompartmentalModel::subsys, n]]
			]; 

		transitions = Union @ Flatten @ transitionList; 

		compartments = Join @@ Cases[transitions, Transition[c1_, _, c2_] :> {c1, c2}] // 
			Variables // DeleteCases[#, Global`\[EmptySet]]&;
		Scan[If[!compartmentQ[#], Message[KineticCompartmentalModel::ifcomp, #]]&, compartments];

		{reactants, rates, products} = 
			Transpose @ Replace[List @@@ transitions, {} -> {{}, {}, {}}]; 

		vars = Through[compartments[t]]; 

		stoichiometry = Outer[Coefficient, products - reactants, compartments]; 
		crossTerms = Join @@ Table[compartments[[i]] compartments[[j]], 
			{i, Length[compartments]}, {j, i}]; 
		crossTermStoichiometry = Outer[Coefficient, products - reactants, 
			crossTerms]; 

		Scan[If[!MatchQ[#, {_, {(0)...}}], Message[KineticCompartmentalModel::xtrm, First[#]]]&, 
			{crossTerms, Transpose[crossTermStoichiometry]} // Thread]; 
		defaultOrders = (Thread[compartments -> #1]&) /@ 
			Outer[Coefficient, reactants, compartments](* // Echo[#, "defaultOrders"]&*); 
		orders = compartments /. defaultOrders(* // Echo[#, "orders"]&*); 
		orderRules = (Thread[compartments -> #1]&) /@ 
			(vars^#1&) /@ orders(* // Echo[#, "orderRules"]&*); 
		terms = rates (Times @@ compartments /. orderRules); 

		flux = terms; 

		lhs = D[vars, t];
		rhs = flux . stoichiometry; 
		eqns = lhs == rhs // Thread; 

		Scan[If[MatchQ[#, Derivative[1][_][t] == 0], 
			Message[KineticCompartmentalModel::ntzr, #[[1, 0, 1]]]]&, eqns]; 

		{vars, transitions, stoichiometry, flux, eqns}
		]

End[]; 

EndPackage[]; 
