(* Copyright (C) 2011, 2012 by Itai Seggev 
 *
 * This work is licensed under the Creative Commons Attribution-NonCommercial-ShareAlike 3.0 Unported License. 
 * To view a copy of this license, visit http://creativecommons.org/licenses/by-nc-sa/3.0/ or send a letter to 
 * Creative Commons, 444 Castro Street, Suite 900, Mountain View, California, 94041, USA.
 *
 * THIS SOFTWARE IS PROVIDED ``AS IS'' AND WITHOUT ANY EXPRESS OR
 * IMPLIED WARRANTIES, INCLUDING, WITHOUT LIMITATION, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE.
 *
 * Version 1.0.0, 2/2/12:  First release
 * Version 1.1.0, 8/16/14: Update for M10 compatability
 * Version 1.1.1, 8/20/14: Improved PlotTheme support
 *)


BeginPackage["BifurcationPlots`"]

Unprotect[BifurcationPlot3D];
ClearAll[BifurcationPlot3D];


BifurcationPlot3D::usage = "\<BifurcationPlot3D[\!\(\*
StyleBox[\"f\", \"TI\"]\),{\!\(\*
StyleBox[\"c\", \"TI\"]\),\!\(\*SubscriptBox[
StyleBox[\"c\", \"TI\"], 
StyleBox[\"min\", \"TI\"]]\),\!\(\*SubscriptBox[
StyleBox[\"c\", \"TI\"], 
StyleBox[\"max\", \"TI\"]]\)},{\!\(\*
StyleBox[\"y\", \"TI\"]\),\!\(\*SubscriptBox[
StyleBox[\"y\", \"TI\"], 
StyleBox[\"min\", \"TI\"]]\),\!\(\*SubscriptBox[
StyleBox[\"y\", \"TI\"], 
StyleBox[\"max\", \"TI\"]]\)}] generates the interactive 3D bifurcation diagram for the system \!\(\*
StyleBox[\"y\", \"TI\"]\)'==\!\(\*
StyleBox[\"f\", \"TI\"]\)(\!\(\*
StyleBox[\"y\", \"TI\"]\),\!\(\*
StyleBox[\"c\", \"TI\"]\)).  Options EquilibriumCurveStyle, EquilbriumPlaneStyle, RateFunctionCurveStyle, and RateFunctionPlaneStyle \
accept graphics directives and set the named style.  Option EquilibriumCurve controls how the intersections of the rate function with \
the equilibrium plane are computed with default Automatic; other values are None (do not compute), \"FastNumeric\" (use a fast \
numerical algorithm), \"AccurateNumeric\" (use a more accurate but potentially quiet slow numeric algorithm).  It also accepts all \
the options of Plot3D which it uses to control the overall plot.\>";


Options[BifurcationPlot3D] = Join[
	Options[Plot3D],
	{
		EquilibriumCurves -> Automatic, 
		EquilibriumCurveStyle -> Directive[Thickness[0.015],Black],
		EquilibriumPlaneStyle -> Directive[Black, Opacity[0.5]],
		RateFunctionCurveStyle -> Directive[Thickness[0.015],Red],
		RateFunctionPlaneStyle -> Directive[Red, Opacity[0.5]] 
	}
];
SetOptions[BifurcationPlot3D, {AxesLabel->Automatic, BoxRatios->Automatic}];


BifurcationPlot3D::nsmet = "Cannot solve for equilibrium curve using Solve; drawing using fast numerical approximation.";
BifurcationPlot3D::slmsg = "Possible error during Solve.  Consider setting option EquilbriumCurves.";
BifurcationPlot3D::bdshp = StringJoin[
	"Option value `1` should be Automatic, None, \"FastNumeric\", \"AccurateNumeric\", ",
	"or a graph or list of graphs in the `2``3`-plane."
];
BifurcationPlot3D::bdgrs = "Option value `1` contains invalid expression `2` which is not a graph in the `3``4`-plane";
BifurcationPlot3D::bdgrp = "Option value `1` contains invalid expressions `2` which are not graphs in the `3``4`-plane";
BifurcationPlot3D::bdrte = "Rate function `1` is not an expression in plotting variables `2` and `3`";
BifurcationPlot3D::bdstl = "Style specification `1` for option `2` should be a graphics primitive or directive.";
BifurcationPlot3D::bdvar = "Expression `1` in argument `2` is not a valid variable";
BifurcationPlot3D::nknwn = "An unknown error occured while plotting the rate function; double check input";
BifurcationPlot3D::nolst = "Rate function `1` cannot be a list";

Begin["`Private`"]


BifurcationPlot3D[
	expr: Except[_List],  
	{c_?validVariableQ, cMin_?NumericQ, cMax_?NumericQ}, 
	{y_?validVariableQ, yMin_?NumericQ, yMax_?NumericQ},
	opts : OptionsPattern[]
] := Block[{regionFunc, zeroCurves, plot3d, zMin, zMax, axesLabels, ecStyle, epStyle, rfcStyle, rfpStyle, localStyles, epPlot},
	
	error = Catch[
		axesLabels = Replace[
			OptionValue[BifurcationPlot3D, {opts}, AxesLabel],
			Automatic -> (Style[#,Large]& /@ {c,y, RawBoxes["f"]})
		];
		plot3d = Check[
			Plot3D[expr, {c, cMin, cMax}, {y, yMin, yMax},
				(* This comes first becuase intercepted above *)
				AxesLabel -> axesLabels,
				Evaluate[Sequence @@ FilterRules[{opts}, Options[Plot3D]]],
				(* Set default BoxRatios, might have been overridden above. *)
				BoxRatios->Automatic,
				Evaluate @ If[$VersionNumber > 9, {PlotTheme->"Classic", Lighting->"Classic"}, {}]
			],
			If[(Head[plot3d] =!= Graphics3D && Not @ MatchQ[plot3d,Legended[_Graphics3D,___]]) || MatchQ[plot3d,Graphics3D[{},___] | Legended[Graphics3D[{},___],___]],
				 Throw[True] (* a fatal error occured*)
			]
		];
		(* check if Plot3D silently failed an unevaluated Plot3D*)
		If[ (Head[plot3d] =!= Graphics3D && Not @ MatchQ[plot3d,Legended[_Graphics3D,___]]) || MatchQ[plot3d,Graphics3D[{},___] | Legended[Graphics3D[{},___],___]], 
			If[ Internal`DependsOnQ[expr, c] && Internal`DependsOnQ[expr, y],
				Message[BifurcationPlot3D::nknwn],
				Message[BifurcationPlot3D::bdrte, expr, c, y]
			];
			Throw[True] (* an error occured*)
		];
	
		zeroCurves = processZeroCurves[expr, c, y, OptionValue[BifurcationPlot3D, {opts}, EquilibriumCurves]];
		
		localStyles = {EquilibriumCurveStyle, EquilibriumPlaneStyle, RateFunctionCurveStyle, RateFunctionPlaneStyle};
		{ecStyle, epStyle, rfcStyle, rfpStyle} = OptionValue[BifurcationPlot3D, {opts}, localStyles];
		MapIndexed[
			If[ ! MatchQ[#, _?Charting`generalDirective | {__?Charting`generalDirective}],
				Message[BifurcationPlot3D::bdstl, #,  localStyles[[ First @ #2 ]] ];
				Throw[True] (* an error occured*)
			]&,
			{ecStyle, epStyle, rfcStyle, rfpStyle}
		]		
	];
	

	(* Finish computation, assuming no errors occured.  We're assuming here that if the Plot3D suceeded, so will the ContourPlot *)
	CompoundExpression[
		regionFunc = (Im[expr] == 0. && Slot[3] < expr &) /. {c->Slot[1], y->Slot[2]};
		(* Need quiet since certain plot themes make AbsoluteOptions noisy *)
		{zMin,zMax} = Quiet[Last[PlotRange /. AbsoluteOptions[plot3d]]];
		epPlot = Switch[zeroCurves,
			_Missing,
	 		With[{z=Unique[], zi=zMin/10, za = zMax/10 },
	 			ContourPlot3D[{expr-z==0, z==0},{c, cMin, cMax}, {y, yMin, yMax}, {z,zi,za},
					ContourStyle->{Opacity[0], epStyle}, 
					Mesh->None, 
					BoundaryStyle->{{1,2} -> ecStyle},
					PerformanceGoal->If[zeroCurves === Missing["RequestedAccurateNumeric"], "Quality", "Speed"]
				]
	 		],
			
			{_, {}},
			Plot3D[0,{c, cMin, cMax}, {y, yMin, yMax},
				PlotStyle->epStyle, 
				Mesh->None, 
				PerformanceGoal->"Speed"
			],
			
			_,
			Sequence @@ {
			Plot3D[0,{c, cMin, cMax}, {y, yMin, yMax}, 
				PlotStyle->epStyle, 
				Mesh->None, 
				PerformanceGoal->"Speed"
			],	
	 		ParametricPlot3D[Last @ zeroCurves, {c, cMin, cMax}, PlotStyle->{ecStyle}]
			}
		],
		
		(* Second half of CompoundExpression*)
		With[
			{
				yZC = First @ zeroCurves, z = Unique[], d = Unique[],
				rg = regionFunc, plot3d = plot3d, epPlot = epPlot, zMin = zMin, zMax=zMax, 
				ecStyle = ecStyle, rfcStyle = rfcStyle, rfpStyle = rfpStyle
			},
			Manipulate[
				Show[
					plot3d,
					epPlot,
					(* Plot the vertical plane at c = current value, below the rate function*)
					ContourPlot3D[c==d,{c, cMin, cMax}, {y, yMin, yMax}, {z,zMin,zMax},
						ContourStyle->rfpStyle, 
						RegionFunction->rg,
						Mesh->None
					],
					(* The first plots intersection of the rate function with the c-plane, the second the equilbrium curve*)
					ParametricPlot3D[Evaluate[{{d, y, expr /. c->d}, Sequence @@ yZC}], {y,yMin,yMax}, 
						PlotStyle->Join[{rfcStyle}, ConstantArray[ecStyle, Length @ yZC]]
					]
				], 
				{{d, Divide[cMin+cMax,2], Style[c,"TraditionalForm",18, FontFamily->"Times"]}, cMin, cMax, Appearance->"Labeled"}
			]
		] 
	]/; error =!= True
] 


BifurcationPlot3D[expr: _List, _, _, OptionsPattern[]] := $Failed /; Message[BifurcationPlot3D::nolst, expr]  


BifurcationPlot3D[_, arg: {_?validVariableQ, bad: Except[_?NumericQ], _}, _, OptionsPattern[]] := $Failed /; Message[BifurcationPlot3D::plln, bad, arg] 
BifurcationPlot3D[_, arg: {_?validVariableQ, _?NumericQ, bad: Except[_?NumericQ]},  _, OptionsPattern[]] := $Failed /; Message[BifurcationPlot3D::plln, bad, arg] 
BifurcationPlot3D[_, _, arg: {_?validVariableQ, bad: Except[_?NumericQ],  _}, OptionsPattern[]] := $Failed /; Message[BifurcationPlot3D::plln, bad, arg] 
BifurcationPlot3D[_, _, arg: {_?validVariableQ, _?NumericQ, bad: Except[_?NumericQ]}, OptionsPattern[]] := $Failed /; Message[BifurcationPlot3D::plln, bad, arg] 


BifurcationPlot3D[_, arg: {bad: Except[_?validVariableQ], _, _}, _, OptionsPattern[]] := $Failed /; Message[ BifurcationPlot3D::bdvar, bad, arg] 
BifurcationPlot3D[_, {_?validVariableQ,_,_}, arg: {bad: Except[_?validVariableQ], _, _}, OptionsPattern[]] := $Failed /; Message[ BifurcationPlot3D::bdvar, bad, arg] 


BifurcationPlot3D[_, bad: Except[{_,_,_}], _, OptionsPattern[]] := $Failed /; Message[BifurcationPlot3D::pllim, bad]
BifurcationPlot3D[_, {_,_,_}, bad: Except[{_,_,_}], OptionsPattern[]] := $Failed /; Message[BifurcationPlot3D::pllim, bad]


BifurcationPlot3D[] := $Failed /; Message[BifurcationPlot3D::argm, BifurcationPlot3D, 0, 3]
BifurcationPlot3D[_] := $Failed /; Message[BifurcationPlot3D::argmu, BifurcationPlot3D, 3]
BifurcationPlot3D[_,_] := $Failed /; Message[BifurcationPlot3D::argm, BifurcationPlot3D, 2, 3]
expr: BifurcationPlot3D[_,_, _, rest__] := $Failed /; If[Not @ MatchQ[{rest}, {__?OptionQ}],
	Message[
		BifurcationPlot3D::nonopt, 
		First @ Cases[{rest}, x_ /; Not @ OptionQ[x], 1, 1], (* extract first non-option*) 
		3, (* position where we no longer want regular arguments*)
		HoldForm[expr] (* hold expression *)
	]
]


Clear[orderedPairQ, processZeroCurves, validVariableQ];


orderedPairQ[{Except[_List], Except[_List]}] := True
orderedPairQ[_] := False  


processZeroCurves[expr_, c_, y_, eqCurve: None] = {{},{}};


processZeroCurves[expr_, c_, y_, eqCurve: "FastNumeric" | "AccurateNumeric"] := Missing["Requested" <> eqCurve]


processZeroCurves[expr_, c_, y_, eqCurve: Automatic | All] := Module[{zeroCurves},
	Check[
		zeroCurves = TimeConstrained[Solve[expr == 0, y,Method->"Legacy"],15];
		If[ MatchQ[zeroCurves,	{} | _Solve | $Failed],  
			Message[BifurcationPlot3D::nsmet]; 
			Return[Missing["NotAvailable"], Module]
		],
		Message[BifurcationPlot3D::slmsg]
	];
	{{}, Thread[{c, y /. zeroCurves, 0}]}	
]


processZeroCurves[expr_, c_, y_, eqCurve_] := Module[{zeroCurves, yDepCurves, cDepCurves, droppedCurves},
	Which[ 
		orderedPairQ[eqCurve],
		zeroCurves = {eqCurve},
		
		MatchQ[eqCurve, {__?orderedPairQ}],
		zeroCurves = eqCurve,
		
		True,
		Message[BifurcationPlot3D::bdshp, eqCurve, c, y];
		Throw[True] (* an error occured *)
	];
	(* Need the NumericQ[g] to account for constants!*)
	yDepCurves = Cases[zeroCurves, {g_, y} /; (Internal`DependsOnQ[g, y] && Not @ Internal`DependsOnQ[g, c]) || NumericQ[g]];
	cDepCurves = Cases[zeroCurves, {c, g_} /; (Internal`DependsOnQ[g, c] && Not @ Internal`DependsOnQ[g, y]) || NumericQ[g]];
	droppedCurves = Complement[zeroCurves, yDepCurves, cDepCurves];
	If[ droppedCurves =!= {},
		If[Length[droppedCurves] === 1,
			Message[BifurcationPlot3D::bdgrs, eqCurve, First @ droppedCurves, c, y],
			Message[BifurcationPlot3D::bdgrp, eqCurve, Row[droppedCurves, ","], c, y]
		];
		Throw[True] (* an error occured *),
		
		(* else *)
		{Thread[{yDepCurves, 0}], Thread[{cDepCurves, 0}]}
	] 
]


validVariableQ[e_] :=  Quiet @ Check[
	Plot[e, {e,0,1}];
	Solve[{e==1}, {e}];
	True,
	
	False
]


End[]


Protect[BifurcationPlot3D]


EndPackage[]
