(* :Title: Fields *)

(* :Author: Steven Wilkinson, Northern Kentucky University *)

(* :Summary:  	VectorFieldPlot[] and VectorFieldListPlot 
						plots the two dimensional vector field of a
						vector valued function of two variables.
						DirectionFieldPlot[] and DirectionFieldListPlot
						plots the line elements whose slopes are
						given by a scalar-valued function of two variables.
*)

(* :Context: Graphics`Fields` *)

(* :Package Version: 2.0 *)

(* :Keywords: vector fields, direction fields
*)

BeginPackage[
	"Graphics`Fields`","Utilities`FilterOptions`"
];

DirectionFieldPlot::usage =
	"DirectionFieldPlot[exp,{x,xmin,xmax},{y,ymin,ymax}] generates a
plot of the line elements whose slopes are given by exp, a
function of x and y, for values of x between xmin and xmax,
and values of y between ymin and ymax.";

DirectionFieldListPlot::usage =
	"DirectionFieldListPlot[func,pts] generates a plot of the
line elements at each point of pts with slope given by evaluating
pure function func at that point.";
	
VectorFieldPlot::usage =
	"VectorFieldPlot[vectorExp,{x,xmin,xmax},{y,ymin,ymax}]
generates a plot of the two dimensional vectors given by
vectorExp, a vector-valued expression in x and y, for values
of x between xmin and xmax, and values of y between ymin and ymax.";

VectorFieldListPlot::usage =
	"VectorFieldListPlot[vectorFunc,pts] generates a plot of
the two dimensional vectors based at pts given by vectorFunc,
a vector-valued pure function, evaluated on pts.";

PlotPoints::usage =
	"PlotPoints is an option for plotting functions that
specifies how many sample points to use.";

PlotStyle::usage =
	"PlotStyle is an option for various plotting functions
that specifies the style of lines or points to be plotted.";

ArrowheadStyle::usage =
	"ArrowheadStyle is an option for VectorFieldPlot and
VectorFieldListPlot. ArrowheadStyle->None draws no head.
ArrowheadStyle->Filled gives a filled arrowhead.
ArrowheadStyle->Unfilled gives an unfilled arrowhead.";

TooLongArrowheadStyle::usage =
	"TooLongArrowheadStyle is an option for VectorFieldPlot and
VectorFieldListPlot that specifies the arrowhead style for
vectors that are longer than the maximum vector length given
by the option MaxVectorLength. TooLongArrowheadStyle->Filled
gives a filled arrowhead. TooLongArrowheadStyle->Unfilled
gives an unfilled arrowhead.";

ArrowheadAngle::usage =
	"ArrowheadAngle is an option for VectorFieldPlot and
VectorFieldListPlot that specifies the angle of the
arrowhead with the vector stem.";

TooLongArrowheadAngle::usage =
	"TooLongArrowheadAngle is an option for VectorFieldPlot and
VectorFieldListPlot that specifies the angle of the
arrowhead with the vector stem for vectors that are longer
than the maximum vector length given by the option MaxVectorLength.";

ArrowheadLength::usage =
	"ArrowheadLength is an option for VectorFieldPlot and
VectorFieldListPlot that specifies the length of the
arrowhead as a fraction of the length of the vector stem.";

TooLongArrowheadLength::usage =
	"TooLongArrowheadLength is an option for VectorFieldPlot and
VectorFieldListPlot that specifies the length of the
arrowhead as a fraction of the length of the vector stem for
vectors that are longer than the maximum vector length given
by the option MaxVectorLength.";

MaxVectorLength::usage =
	"MaxVectorLength is an option for VectorFieldPlot and
VectorFieldListPlot that specifies the longest vector
length, if any, that will be plotted.";

ScaleFactor::usage =
	"ScaleFactor is an option for VectorFieldPlot and
VectorFieldListPlot. If it is set to
a pure function then all the vectors' lengths will be scaled to
the length given by the function applied to their original lengths.
If it is set to None no scaling occurs, If is set to a number then
the vectors' lengths is scaled by that number. ScaleFactor->UnitLength
scales all the vectors to unit length. ScaleFactor->UnitLengthScaled
scales all the vectors to the minimum distance between adjacent points.
ScaleFactor->Scaled scales the vectors' lengths by a number that
makes the longest vector have length equal to the minimum distance
between adjacent points.";

ScaleLength::usage =
	"ScaleLength is an option for DirectionFieldPlot and
DirectionFieldListPlot that scales the length of each line element
to its value times the minimum distance between adjacent points.";

TooLongPlotStyle::usage =
	"TooLongPlotStyle is an option for VectorFieldPlot and
VectorFieldListPlot that specifies the plot style for vectors
that are longer than the maximum vector length given by the
option MaxVectorLength.";

LengthScaled::usage =
	"LengthScaled can be used in the option ScaleFactor to
rescale all vectors so there are no overlapping vectors.";

LogLengthScaled::usage =
	"LogLengthScaled can be used in the option ScaleFactor to
rescale all vectors to length log(1+length), further scaled
so there are no overlapping vectors.";

UnitLength::usage =
	"UnitLength can be used in the option ScaleFactor to rescale all
nonzero vectors to unit lengths.";

UnitLengthScaled::usage =
	"UnitLengthScaled can be used in the option ScaleFactor to rescale all
nonzero vectors to the length of the minimum distance between adjacent
points.";

Unfilled::usage =
	"Unfilled can be used in the option ArrowheadStyle to obtain unfilled
arrowheads.";

Filled ::usage =
	"Filled can be used in the option ArrowheadStyle to obtain filled
arrowheads.";

DirectionFieldPlot::arg =
	"DirectionFieldPlot must have at least 3 arguments, where the
second and third arguments are both of the form {x, xmin, xmax}.";

DirectionFieldListPlot::argfunc =
	"The first argument, `1` must be a pure function or the head of
a function (no arguments).";

DirectionFieldListPlot::arg =
	"DirectionFieldListPlot must have at least 2 arguments.";

FieldListPlot::argpts =
	"The second argument, `1`, must be a list of pairs of real numbers.";
	
VectorFieldPlot::arg =
	"VectorFieldPlot must have at least 3 arguments, where the
second and third arguments are both of the form {x, xmin, xmax}.";

VectorFieldPlot::vfptr =
	"`1` does not evaluate to a pair of real numbers at
`2` = `4` and `3` = `5`.";

VectorFieldListPlot::argfunc =
	"The first argument, `1` must be a pure vector-valued function
or the head of a vector-valued function (no arguments).";

VectorFieldListPlot::arg =
	"VectorFieldListPlot must have at least 2 arguments.";
	
Options[DirectionFieldPlot] =
	{PlotPoints->20,
	PlotStyle->Automatic,
	ScaleLength->.6};
	
Options[DirectionFieldListPlot] =
	{PlotStyle->Automatic,
	ScaleLength->.6};

Options[VectorFieldPlot] =
	{ArrowheadStyle->Unfilled,
	ArrowheadAngle->Pi/6,
	TooLongArrowheadAngle->Pi/6,
	ArrowheadLength->0.3,
	TooLongArrowheadLength->0.3,
	MaxVectorLength->None,
	PlotPoints->10,
	PlotStyle->Automatic,
	TooLongPlotStyle->Automatic,
	ScaleFactor->LengthScaled};
	
Options[VectorFieldListPlot] =
	{ArrowheadStyle->Unfilled,
	ArrowheadAngle->Pi/6,
	TooLongArrowheadAngle->Pi/6,
	ArrowheadLength->0.3,
	TooLongArrowheadLength->0.3,
	MaxVectorLength->None,
	PlotStyle->Automatic,
	TooLongPlotStyle->Automatic,
	ScaleFactor->LengthScaled};

Begin["`Private`"];
	
magnitude[v_List] = Sqrt[v.v];

linePlot[pt_,slope_,lngth_,aspRat_] :=
	Line[
		{
			pt - lngth/Sqrt[1+slope*slope] * {aspRat,slope},
			pt + lngth/Sqrt[1+slope*slope] * {aspRat,slope}
		}
	] /; Head[slope] === Real;

linePlot[pt_,slope_,lngth_,aspRat_] :=
	Line[
		{pt - lngth * {0.,1.}, pt + lngth * {0.,1.} }
	] /; slope === ComplexInfinity;

linePlot[pt_,slope_,lngth_,aspRat_] :=
	Point[pt];
	
vectorPlot[pt_,vector_,other__] :=
	{Point[pt]} /; vector==={0.0,0.0};

vectorPlot[pt_,vector_,arrowInfo__] :=
	Join[
		{Line[{pt,pt+vector}]},
		arrowhead[pt+vector,vector,arrowInfo]
	];
	
arrowhead[tip_,vector_,matrix1_,matrix2_,style_] :=
	Flatten[
		{style[
			{tip + matrix1.vector,
			tip,
			tip + matrix2.vector}
		]}
	];


findMinDist[minDist_?NumberQ,pts_] := {minDist,minDist};

findMinDist[minDist_,pts_] := minDist /; VectorQ[minDist,NumberQ];

findMinDist[minDist_,pts_] := 1 /; Length[pts] == 1;

findMinDist[minDist_,pts_] :=
	Min[
		Apply[
			Plus,
			Flatten[
				Map[
					subtract[pts[[#]],Take[pts,#-Length[pts]]]&,
					Range[Length[pts]-1]
				],1
			],1
		]
	];
	
point1[{x_,y_,z_}] := Point[y];

generatePoints[xRange_List,yRange_List,opts___] :=
	Module[{dx,dy,plotPts,xmin,xmax,ymin,ymax},
		plotPts = PlotPoints /. {opts};
		If[NumberQ[plotPts],
			plotPts = {plotPts,plotPts}
		];
		plotPts = plotPts - {1,1};
		If[plotPts[[1]] < 1 || plotPts[[2]] < 1,
			Message[VectorField::vfpp]
		];
		{xmin,xmax} = Sort[Take[xRange,-2]];
		{ymin,ymax} = Sort[Take[yRange,-2]];
		{dx,dy} = N[{xmax-xmin,ymax-ymin}/plotPts];
		{
			Flatten[
				Table[
					N[{xmin + m*dx, ymin + n*dy}],
					{m,0,plotPts[[1]]},{n,0,plotPts[[2]]}
				],1
			],
			opts,
			MinimumDistance->{dx,dy},
			PlotRange->{{xmin-dx,xmax+dx},{ymin-dy,ymax+dy}}
		}
	];
		
scaleVector[vector_,vectorLength_,newVectorLength_] :=
			vector /; vectorLength == 0.;
			
scaleVector[vector_,vectorLength_,newVectorLength_] :=
			newVectorLength/vectorLength vector;

subtract[x_List,lst_?MatrixQ] :=
	Map[Abs[x-#]&,lst] /; Length[x] == Dimensions[lst][[2]];
	
arrowheadInformation[{{xmin_,xmax_},{ymin_,ymax_}},opts___] :=
	Module[{aspect,arrowhdAngle,arrowhdLength,m,arrowhdRotation,
					style},
		aspect = N[AspectRatio /. {opts} /. Options[Graphics]];
		arrowhdAngle =
			N[ArrowheadAngle /. {opts} /. Options[VectorFieldListPlot]];
		arrowhdLength =
			N[ArrowheadLength /. {opts} /. Options[VectorFieldListPlot]];
		If[NumberQ[aspect],
			m = {{1.,0.},{0., N[(xmax-xmin)/(ymax-ymin) aspect]}},
			m = {{1.,0.},{0.,1.}}
		]; 
	(* corrects arrowhead distortion due to aspect ratio *)
	
		style =
			(ArrowheadStyle /. {opts} /. Options[VectorFieldPlot]) /.
					{Filled->Function[{Line[#],Polygon[#]}],
					Unfilled->Line};
		If[style=!=Line && style=!=Function[{Line[#],Polygon[#]}],
			style=point1
		];
		arrowhdRotation = N[{{Cos[Pi - arrowhdAngle],-Sin[Pi - arrowhdAngle]},
					{Sin[Pi - arrowhdAngle],Cos[Pi - arrowhdAngle]}}];
		{ 
			arrowhdLength*Inverse[m].arrowhdRotation.m,
			arrowhdLength*Inverse[m].Transpose[arrowhdRotation].m,
			style
		}
	];

DirectionFieldPlot[
	exp_,xRange:{x_,x0_,x1_},yRange:{y_,y0_,y1_},opts___
] :=
	DirectionFieldListPlot[
		Function[{x,y},exp],
		Apply[
			Sequence,
			generatePoints[
				xRange,
				yRange,
				opts,
				Apply[Sequence,Options[DirectionFieldPlot]]
			]
		]
	];

DirectionFieldPlot[arg1___] :=
	Message[DirectionFieldPlot::arg];
	
	
DirectionFieldListPlot[func_,points_,opts___] :=
	Module[{a,aspect,dir,minDist,pltRange,pts=N[points]},
		aspect = 
			N[AspectRatio /. {opts} /. Options[Graphics]];
		pltRange = 
			PlotRange /. {opts};
		If[!MatrixQ[pltRange,NumberQ],
			pltRange = 
				Map[
					{Min[#],Max[#]}&,
					Transpose[pts]
				]
		];
		If[NumberQ[aspect],
			a = 
				aspect*Apply[Divide,Apply[Subtract,pltRange,{1}]],
			a = 1.
		];
		minDist = 
			Min[
				{1/a,1}*
				findMinDist[
					(MinimumDistance /. {opts}),pts
				]
			];
		lngth =
			.5*minDist*ScaleLength /. {opts} /.
					Options[DirectionFieldListPlot];
		dir =
			Switch[
				DirectionType /. {opts} /. 
					Options[DirectionFieldListPlot],
				Angle,
					Composition[Tan,func],
				DirectionVector,
					Composition[#[[2]]/#[[1]]&,func],
				PerpendicularVector,
					Composition[-#[[1]]/#[[2]]&,func],
				_,
					func
			];
			dir = Composition[Times[a,#]&,dir];
			Show[
			Graphics[
				Join[
					Flatten[
						{PlotStyle /. {opts} /. Options[DirectionFieldListPlot] /. Automatic->{}}
					],
					Map[
						linePlot[#,Apply[dir,#],lngth,a]&,
						pts
					]
				]
			], 
			FilterOptions[
				Graphics,
				Sequence @@ Join[{opts},Options[DirectionFieldListPlot]]
			], Axes->True
		]
	] /; MatrixQ[points,NumberQ] && 
				(Head[func]===Function || Head[func]===Symbol);
				
DirectionFieldListPlot[func_,other___] :=
	Message[DirectionFieldListPlot::argfunc,func] /;
		Head[func]=!=Function && Head[func]=!=Symbol;

DirectionFieldListPlot[func_,pts_,opts___] :=
	Message[FieldListPlot::argpts,pts] /;
		!MatrixQ[N[pts],NumberQ] || Last[Dimensions[pts]]=!=2;
	
DirectionFieldListPlot[func___] :=
	Message[DirectionFieldListPlot::arg];

VectorFieldPlot[
	exp_,xRange:{x_,x0_,x1_},yRange:{y_,y0_,y1_},opts___
] :=
	VectorFieldListPlot[
		Function[{x,y},exp],
		Apply[
			Sequence,
			generatePoints[
				xRange,
				yRange,
				opts,
				Apply[Sequence,Options[VectorFieldPlot]]
			]
		]
	] /; Head[exp] === List && Length[exp] == 2;

VectorFieldPlot[arg1___] :=
	Message[VectorFieldPlot::arg];
		
VectorFieldListPlot[func_,points_List,opts___] :=
	Module[{arrowheadInfo,pts=N[points],minDistance,vectors,
					lengths,ptsIndex,badPts,goodPts,tooLongPts,pltRange,
					tooLongArrowheadInfo,maxLength,scaleFunc,scaledLengths},
		vectors = Apply[func,pts,1];
		lengths = Map[magnitude,vectors];
		ptsIndex = Range[Length[pts]];
	(* 
		Divide points into groups of bad points, points with vector
		length too long, and others.
	*)
		badPts = 
			Select[
					ptsIndex,
					!(VectorQ[vectors[[#]],NumberQ] && Length[vectors[[#]] ]===2)&
			];
		Map[
			Message[VectorFieldPlot::vfptr,func,x,y,pts[[#,1]],pts[[#,2]] ]&,
			badPts
		];
		vectors = ReplacePart[vectors,{0.,0.},Map[List,badPts] ];
		lengths = ReplacePart[lengths,0.,Map[List,badPts] ];
		(* 
			Return error message if input vector field is bad and
			replace with the zero vector.
		*)
		
		goodPts = Complement[ptsIndex,badPts];
		maxLength = MaxVectorLength /. {opts} /. Options[VectorFieldListPlot];
		If[maxLength===None,
		
			maxLength = Max[lengths];
			tooLongPts = {},
			
			tooLongPts =
			Select[
				goodPts,
				lengths[[#]] > maxLength&
			];
			vectors = 
				MapAt[maxLength*#/magnitude[#]&,vectors,Map[List,tooLongPts]];
			lengths = ReplacePart[lengths,.9*maxLength,Map[List,tooLongPts]]
		];
		goodPts = Complement[goodPts,tooLongPts];
	
		aspect = 
			N[AspectRatio /. {opts} /. Options[Graphics]];
		pltRange = 
			PlotRange /. {opts};
		If[!MatrixQ[pltRange,NumberQ],
			pltRange = 
				Map[
					{Min[#],Max[#]}&,
					Transpose[Join[pts,pts+vectors]]
				]
		];
		If[NumberQ[aspect],
			a = aspect*Apply[Divide,Apply[Subtract,pltRange,{1}]],
			a = 1.
		];

		minDistance = 
			Min[
				{1,1/a}*
				findMinDist[
					(MinimumDistance /. {opts}),pts
				]
			];
	(* Scale the vectors according to the option ScaleFunction. *)		
		scaleFunc =
			ScaleFactor /. {opts} /. Options[VectorFieldListPlot] /.
			{LengthScaled -> Function[#*minDistance/maxLength],
			LogLengthScaled -> Function[Log[#+1]*minDistance/Log[1+maxLength]],
			UnitLengthScaled -> Function[minDistance],
			UnitLength -> Function[1],
			Automatic -> Identity,
			None -> Identity};
		If[
			NumberQ[scaleFunc],
			scaleFunc = Function[Evaluate[#*scaleFunc]],
		];

		scaledLengths = Map[scaleFunc,lengths];
		vectors = MapThread[scaleVector,{vectors,lengths,scaledLengths}];
	
	(* Determine arrowhead information.*)
		arrowheadInfo = 
			Apply[Sequence,arrowheadInformation[pltRange,opts] ];

		tooLongArrowheadInfo =
			Apply[
				Sequence,
				arrowheadInformation[
					pltRange,
					Apply[
						Sequence,
						Join[
							Map[
								ToExpression[StringDrop[ToString[#],7]]&,
								Select[
									{opts},
									StringMatchQ[StringTake[ToString[#],7],"TooLong"]&
								]
							],{opts}
						]
					]
				]
			];
	(* Plot the vector field. *)
		Show[
			Graphics[
				Join[
					Flatten[
						{PlotStyle} /. {opts} /. Options[VectorFieldListPlot] /.
						Automatic->{}
					],
					Map[
						vectorPlot[pts[[#]],vectors[[#]],arrowheadInfo]&,
						goodPts
					],
					TooLongPlotStyle /. {opts} /. Options[VectorFieldListPlot] /. Automatic->{},
					Map[
						vectorPlot[pts[[#]],vectors[[#]],tooLongArrowheadInfo]&,
						tooLongPts
					]
				]
			],
			FilterOptions[
				Graphics,
				Sequence @@ Join[{opts},Options[VectorFieldListPlot]]
			], Axes->True
		]
	] /; MatrixQ[points,NumberQ] && 
				(Head[func]===Function || Head[func]===Symbol);
					
	VectorFieldListPlot[func_,other___] :=
		Message[VectorFieldListPlot::argfunc,func] /;
			Head[func]=!=Function && Head[func]=!=Symbol;

	VectorFieldListPlot[func_,pts_,opts___] :=
		Message[FieldListPlot::argpts,pts] /;
			!MatrixQ[N[pts],NumberQ] || Last[Dimensions[pts]]=!=2;
	
	VectorFieldListPlot[func___] :=
		Message[VectorFieldListPlot::arg];

 End[]; (* "`Private`" *)

EndPackage[];  (* "Graphics`Fields`" *)
