#=============================================================================
# Plotting functions
#=============================================================================
# Marc A. Murison
# Astronomical Applications Dept.
# U.S. Naval Observatory
# 3450 Massachusetts Ave., NW
# Washignton, DC 20392
# murison@aa.usno.navy.mil
# http://aa.usno.navy.mil/murison/
#=============================================================================

plotting := module()

	 description "a package of handy plotting functions";
	 option package;
	 export	imageplot, psplot, jpgplot, gifplot, vrmlplot, ODEplot, curveplot2D, 
	 		implicitplot2D, scalarplot, coordplot2D, curveplot3D, curvatureplot3D, 
	 		tubecurve, tubecurveplot, surfaceplot, implicitplot3D, contourplot3D, 
	 		coordplot3D, ZernikeWavefrontPlot, myarrow, surfarrows;


#---------------------------------------------------------------
# imageplot( dev, plot, filename, filedir, plot options )
#---------------------------------------------------------------
imageplot := proc( dev::name, p::{PLOT,PLOT3D}, fname::string, fdir::string )
    local f, pltoptions, ftype;
    if dev<>'ps' and dev<>'gif' and dev<>'jpeg' and dev<>'vrml' then
        ERROR("dev must be one of: gif, jpg, ps, vrml");
    end if;
    if dev='vrml' then
        plottools[vrml]( p, cat(fdir,fname,".wrl") );
    else
        if dev='ps' then
            pltoptions := "portrait,noborder,color=cmyk,width=8in,height=8in";
            ftype := ".ps";
        elif dev='gif' then
            pltoptions := "height=600,width=600,colors=256";
            ftype := ".gif";
        elif dev='jpeg' then
            pltoptions := "height=600,width=600";
            ftype := ".jpg";
        else
            ERROR("dev must be one of: gif, jpg, ps, vrml");
        end if;
        if nargs > 4 then
            if nargs > 5 then
                ERROR("Too many arguments!");
            end if;
            if not type(args[5],string) then
                ERROR("Plot options must be in the form of a string!");
            end if;
            pltoptions := cat(pltoptions,",",args[5]);
        end if;
        plotsetup( dev,
                   plotoutput=cat(fdir,fname,ftype),
                   plotoptions=pltoptions);
        print(p);
        plotsetup(default);
    end if;
end proc;

#---------------------------------------------------------------
# psplot( plot, filename, filedir, plot options )
#---------------------------------------------------------------
psplot := proc( p::{PLOT,PLOT3D}, fname::string, fdir::string )
    imageplot(ps,args);
end proc;

#---------------------------------------------------------------
# jpgplot( plot, filename, filedir, plot options )
#---------------------------------------------------------------
jpgplot := proc( p::{PLOT,PLOT3D}, fname::string, fdir::string )
    imageplot(jpeg,args);
end proc;

#---------------------------------------------------------------
# gifplot( plot, filename, filedir, plot options )
#---------------------------------------------------------------
gifplot := proc( p::{PLOT,PLOT3D}, fname::string, fdir::string )
    imageplot(gif,args);
end proc;

#---------------------------------------------------------------
# vrmlplot( plot, filename, filedir, plot options )
#---------------------------------------------------------------
vrmlplot := proc( p::{PLOT,PLOT3D}, fname::string, fdir::string )
    imageplot('vrml',args);
end proc;



#=============================================================================
# ODE Plotting
#=============================================================================

#---------------------------------------------------------------
# ODE solution curves
#---------------------------------------------------------------
ODEplot := proc( ODEs::{`=`,list,set}, depvars::{list,set},
                            ivar_range::{name=range}, ICs::list )

    DEtools[DEplot]( ODEs, depvars, ivar_range, [eval(ICs)], axes=normal,
                     linecolor=navy, thickness=1, axes=normal,
                     titlefont=[TIMES,ITALIC,18], axesfont=[HELVETICA,12],
                     labelfont=[TIMES,ROMAN,14], scaling=unconstrained,
                     method=dverk78,
                     args[5..nargs] );

end proc;



#=============================================================================
# 2D Plotting
#=============================================================================

#---------------------------------------------------------------
# curveplot2D( curve, t=t1..t2, plot_options )
#
# where curve = vector( [ x(t), y(t) ] )
# and the options include, besides the usual plot options, the
# keywords radial, evolute, involute, pedal, and inversion. 
# If "pedal" is specified, the reference point must be 
# listed as curveplot2D(..., refpoint=[...]), where the list 
# contains the x, y, (and z) components.  If "inversion"
# is listed, then in addition to the reference point there must
# be specified the inversion radius: curveplot2D(..., radius=r),
# where r is the radius.
#---------------------------------------------------------------
curveplot2D := proc( curve::{vector,list}, 
                               t_limits::name=range(numeric) )

    local p, q, curvelist, plotargs, plotlist, cplot, var, k, rp, r;

    #pick out the plot arguments and any derived curves to be plotted
    curvelist := [];
    plotargs  := [args[3..nargs]];
    for p in args[3..nargs] do
        if p=`radial` then 
            curvelist := [op(curvelist),p];
            plotargs  := remove( has, plotargs, `radial` );
        elif p=`evolute` then 
            curvelist := [op(curvelist),p];
            plotargs  := remove( has, plotargs, `evolute` );
        elif p=`involute` then 
            curvelist := [op(curvelist),p];
            plotargs  := remove( has, plotargs, `involute` );
        elif p=`pedal` then
            curvelist := [op(curvelist),p];
            plotargs  := remove( has, plotargs, `pedal` );
            for q in args[3..nargs] do
                if type(q,name=list) and lhs(q)=`refpoint` then
                    rp       := vec( rhs(q) );
                    plotargs := remove( has, plotargs, `refpoint` );
                    break;
                end if;
            od;
        elif p=`inversion` then
            curvelist := [op(curvelist),p];
            plotargs  := remove( has, plotargs, `inversion` );
            for q in args[3..nargs] do
                if type(q,name=list) and lhs(q)=`refpoint` then
                    rp       := vec( rhs(q) );
                    plotargs := remove( has, plotargs, `refpoint` );
                elif type(q,name=algebraic) and lhs(q)=`radius` then
                    r        := vec( rhs(q) );
                    plotargs := remove( has, plotargs, `radius` );
                end if;
            od;
        end if;
    od;

    #define a convenient plotting function
    cplot := proc( curve::{vector,list}, t_limits::name=range(numeric) )
        plot( [op(convert(curve,list)), t_limits],
              axes=frame, numpoints=100, args[3..nargs] );
    end;

    #always plot at least the curve
    plotlist := [ cplot( curve, t_limits, thickness=2, 
                         color=black, op(plotargs) ) ];

    #generate the derived curve plots
    var := lhs(t_limits);
    for k from 1 to size(curvelist) do
        if curvelist[k]=`radial` then
            plotlist := [ op(plotlist), 
                          cplot( radial(curve,var), t_limits, 
                                 thickness=1, color=sienna, op(plotargs) ) ];
        elif curvelist[k]=`evolute` then
            plotlist := [ op(plotlist), 
                          cplot( evolute(curve,var), t_limits, 
                                 thickness=1, color=red, op(plotargs) ) ];
        elif curvelist[k]=`involute` then
            plotlist := [ op(plotlist), 
                          cplot( involute(curve,var), t_limits, 
                                 thickness=1, color=blue, op(plotargs) ) ];
        elif curvelist[k]=`pedal` then
            plotlist := [ op(plotlist), 
                          cplot( pedal(curve,var,rp), t_limits, 
                                 thickness=1, color=magenta, op(plotargs) ) ];
        elif curvelist[k]=`inversion` then
            plotlist := [ op(plotlist), 
                          cplot( inversion(curve,var,rp,r), t_limits, 
                                 thickness=1, color=aquamarine, op(plotargs) ) ];
        end if;
    od;

    #finally, put all the curves on one plot
    display( plotlist, scaling=constrained, op(plotargs) );

end proc;


#---------------------------------------------------------------
# implicitplot2D( eqn, x=x1..x2, y=y1..y2, plot_options )
#---------------------------------------------------------------
implicitplot2D := proc( eqn, x_limits::(name=range), y_limits::(name=range) )

    implicitplot( eqn, x_limits, y_limits, axes=frame, grid=[50,50], color=blue, args[4..nargs] );

end proc;



#---------------------------------------------------------------
# scalarplot( curve, t=t1..t2, plot_options )
#---------------------------------------------------------------
scalarplot := proc( curve::algebraic, t_limits::(name=range) )

    plot( curve, t_limits, numpoints=100, args[3..nargs] );

end proc;



#------------------------------------------------------------
# Plot n constant curves in a given coordinate system
#
# xyeqs = [x(u,v),y(u,v)]
#------------------------------------------------------------
coordplot2D := proc( xyeqs::list, urange::{name=range},
                                vrange::{name=range}, n::integer )

    local uplots, vplots, umin, umax, vmin, vmax, i,
          U, V, plotargs, x, y;

    umin := op([2,1],urange);
    umax := op([2,2],urange);
    vmin := op([2,1],vrange);
    vmax := op([2,2],vrange);
    plotargs := args[5..nargs];

    for i from 1 to n do
        U := umin + (i-1)/(n-1)*(umax-umin);
        uplots[i] := plot( [op(subs(lhs(urange)=U,xyeqs)), vrange], color=blue, plotargs );
    od;
    for i from 1 to n do
        V := vmin + (i-1)/(n-1)*(vmax-vmin);
        vplots[i] := plot( [op(subs(lhs(vrange)=V,xyeqs)), urange], color=red, plotargs );
    od;

    display( [ op(convert(uplots,list)), op(convert(vplots,list)) ], 
             scaling=constrained, axes=box, labels=["x","y"], plotargs );

end proc;



#=============================================================================
# 3D Plotting
#=============================================================================
#------------------------------------------------------------
# Plot a 3D space curve
# curveplot3D( curve, t=t1..t2, plot3d_options )
# where curve = vector([x(t),y(t),z(t)])
#------------------------------------------------------------
curveplot3D := proc( curve::{vector,list}, t_limits::(name=range) )

    plots[spacecurve]( [op(convert(curve,list))], t_limits, scaling=unconstrained, 
            projection=0.8, numpoints=100, args[3..nargs] );

end proc;



#------------------------------------------------------------
# Plot the curvature of a 2D curve as a space curve, with the 
# curve itself plotted on the xy plane.
#
# curvatureplot3D( curve, curvature, t=t1..t2, title, plot3d_options )
# where curve     = [x(t),y(t)]
#       curvature = an expression in t
#       title     = plot title as a string
#------------------------------------------------------------
curvatureplot3D := proc( curve::{vector,list}, curvature::algebraic, 
                                    t_limits::(name=range), plottitle::string )
    local f, c1, c2;
    c1 := curveplot3D( [op(convert(curve,list)), curvature], t_limits, 
                       labels=["x","y","k"], title=plottitle,
                       thickness=3, args[5..nargs] ):
    c2 := curveplot2D( curve, t_limits, color=black, linestyle=1 ):
    f  := plottools[transform]( (x,y)->[x,y,0] );
    display( c1, f(c2), projection=0.8, labels=["x","y","k"], title=plottitle, args[5..nargs] );
end proc;



#---------------------------------------------------------
# tubecurve( curve, t, radius, theta )
#
# Create a "tube" surrounding a space curve, including
# torsion or "twisting" effects.
#---------------------------------------------------------
tubecurve := proc( curve::{vector,list}, t::name, r::algebraic, theta::algebraic )
    local normvec, binorm, tube;
    normvec := UnitNormal(curve,t);
    binorm  := Binormal(curve,t);
    if has(curve,sin(t)) or has(curve,cos(t)) then
        normvec := map( simplify, evalm(normvec), trig );
        binorm  := map( simplify, evalm(binorm), trig );
    end if;
    tube := evalm( curve +
            r*( cos(theta)*normvec + sin(theta)*binorm ) );
end proc;



#---------------------------------------------------------
# Plot a tube curve
#
# tubecurveplot( curve, radius, t=t1..t2,
#                theta=theta1..theta2,
#                plot3d_options )
#
# Use this plotting function, instead of the Maple
# tubeplot(), to explicitely show torsion effects.
#---------------------------------------------------------
tubecurveplot := proc( curve::{vector,list}, radius::algebraic, 
                                  t_limits::(name=range), theta_limits::(name=range) )

    local tube, t, theta, r;

    t     := lhs(t_limits);
    theta := lhs(theta_limits);

    #temporarily replace numerical "radius" with name "r"
    tube  := tubecurve( curve, t, r, theta );
    tube  := map( radsimp, tube );
    tube  := map( collect, tube,
                    [ r, cos(theta), sin(theta) ] );
    tube  := subs( r=radius, evalm(tube) );

    plot3d( tube, t_limits, theta_limits, 
            grid=[100,10], args[5..nargs] );
end proc;



#------------------------------------------------------------
# Plot a 3D surface
#
# surfaceplot( [x(u,v),y(u,v),z(u,v)], u=u1..u2, v=v1..v2,
#              plot3d_options )
#------------------------------------------------------------
surfaceplot := proc( eqns::{list,set}, u_limits::(name=range), 
                                v_limits::(name=range) )

    plot3d( eqns, u_limits, v_limits, grid=[20,20], args[4..nargs] );

end proc;



#------------------------------------------------------------
# Plot a 3D surface implicitly
#
# implicitplot3D( eqn, x=x1..x2, y=y1..y2, z=z1..z2,
#                 plot3d_options )
#------------------------------------------------------------
implicitplot3D := proc( eqn, x_limits::(name=range), 
                                   y_limits::(name=range), z_limits::(name=range) )

    plots[implicitplot3d]( eqn, x_limits, y_limits, z_limits,
        grid=[20,20,20], args[5..nargs] );

end proc;



#------------------------------------------------------------
# Plot a 3D contour plot
#
# contourplot3D( eqn, x=x1..x2, y=y1..y2, plot3d_options )
#
# additional plot options are filled=true and contours=n
#------------------------------------------------------------
contourplot3D := proc( eqn, x_limits::(name=range), y_limits::(name=range) )

    plots[contourplot3d]( eqn, x_limits, y_limits,
        filled=true, coloring=[yellow,red], 
        contours=10, args[4..nargs] );

end proc;



#------------------------------------------------------------
# Plot three constant surfaces in a given coordinate system
#
# xyz = [x(u,v,w),y(u,v,w),z(u,v,w)]
#------------------------------------------------------------
coordplot3D := proc( xyz::list, uvaleqn::{name=numeric},
                                vvaleqn::{name=numeric}, wvaleqn::{name=numeric},
                                urange::{name=range}, vrange::{name=range},
                                wrange::{name=range} )
    local p, q, r, g, X, Y, Z, u, v, w, uval, vval, wval;
    uval := rhs(uvaleqn);
    vval := rhs(vvaleqn);
    wval := rhs(wvaleqn);
    g := 30;
    p := surfaceplot( subs(lhs(urange)=uval,xyz), vrange, wrange,
                      grid=[g,g], shading=zhue, style=patch, args[8..nargs] );
    q := surfaceplot( subs(lhs(vrange)=vval,xyz), urange, wrange,
                      grid=[g,g], shading=zhue, style=patch, args[8..nargs] ):
    r := surfaceplot( subs(lhs(wrange)=wval,xyz), urange, vrange,
                      grid=[g,g], shading=zhue, style=patch, args[8..nargs] ):
    display( [p,q,r], scaling=constrained, labels=["X","Y","Z"], args[8..nargs] );
end proc;



#------------------------------------------------------------
# Plot a wavefront with specified Zernike terms subtracted
#
# PlotZernikeWavefront( opd, Zcoeffs, sublist, units, pert,
#                       f, R, plot3d_options )
#
#------------------------------------------------------------
ZernikeWavefrontPlot := proc( opd::algebraic, epsilon::{name,list}, c::table,
                                         sublist::list, units::string, pert::{float,list},
                                         focal_length::float, beam_radius::float )

    local G, F, i, p, n, m, k, scale, plotargs, tlist:

    #
    # Build up the terms to subtract from the wavefront.
    #
    p := 0:
    for i from 1 to nops(sublist) do
        n := sublist[i][1]:
        m := sublist[i][2]:
        p := p + ZernikeTerm( n, m, c );
    od:

    #
    # Subtract the desired terms and simplify.
    #
    if type(epsilon,list) then
        tlist := [op(epsilon),rho];
    else
        tlist := [epsilon,rho];
    end if;
    G := collect( opd-p, tlist, factor ):
    G := combine( G, trig ):
    G := collect( G, tlist, factor ):
    debug_print( procname, "wavefront:", 1, G ):

    #
    # Make an optimized procedure out of what's left...
    #
    if units="cm" then
        scale := 1:
    elif units="mm" then
        scale := 10;
    elif units="microns" or units="micron" or units="um" then
        scale := 10000:
    elif units="nm" then
        scale := 10000000:
    elif units="pm" then
        scale := 10000000000:
    elif units="fm" then
        scale := 10000000000000:
    else
        ERROR(" Unknown units!"):
    fi:
    if type(epsilon,list) then
        tlist := [rho,phi,op(epsilon),f,R];
    else
        tlist := [rho,phi,epsilon,f,R];
    end if;
    F := optimize( makeproc(-scale*G,parameters=tlist), tryhard ):

    #
    # ...and create a color plot from it.
    #
    if type(epsilon,list) then
        tlist := [rho,phi,op(pert),focal_length,beam_radius];
    else
        tlist := [rho,phi,pert,focal_length,beam_radius];
    end if;
    plotargs := [[rho*cos(phi), rho*sin(phi), F(op(tlist))],
                 phi=0..2*Pi, rho=0..beam_radius,
                 grid=[35,20], orientation=[110,65], 
                 light=[60,90,2,2,2], ambientlight=[0.4,0.4,0.4],
                 labels=["radius (cm)","","OPD "],
                 title="Wavefront Error (".units.")"];
    if nargs > 8 then
        for i from 9 to nargs do
            plotargs := [op(plotargs),args[i]];
        od;
    end if;
    plot3d( op(plotargs) );

end proc;




#---------------------------------------------------------------------------
# bug fixes in the Waterloo Maple VRML package
#---------------------------------------------------------------------------
#`vrml/font` := proc(n) 
#    local t; 
#    t := table([('TIMES')      = ('SERIF'), 
#                ('COURIER')    = ('TYPEWRITER'), 
#                ('HELVETICA')  = ('SANS'),      #Note: 'SANS' is okay with Cosmo but causes WorldView to bomb
#                ('OBLIQUE')    = ('ITALIC'),    #Waterloo forgot to include this
#                ('ROMAN')      = ('NONE'), 
#                ('BOLD')       = ('BOLD'), 
#                ('ITALIC')     = ('ITALIC'), 
#                ('BOLDITALIC') = ('BOLD') 
#               ]); 
#    t[n];
#end proc;




#---------------------------------------------------------------------------
# Plots an arrow as a line with a pyramidal head.
# usage: myarrow(base, vect, options)
# "base" is point for base of arrow, "vect" is tip - base
# These can be lists or vectors.
# "options" are the usual plot3d options.
# Colour is red by default.
# More options could be added to control the proportions of the arrow.
#
# From Robert Israel
#---------------------------------------------------------------------------
myarrow := proc(base::{list,vector},vect::{list,vector})

    local evl,b1,v1,stem,r,hbase,hpoint,v2,v3,head,opts;
    evl := t -> convert(evalm(t), list):
    b1 := evalf(evl(base));
    v1 := evalf(evl(vect));
    if nops(b1) <> 3 or nops(v1) <> 3 then
        ERROR("points need three dimensions")
    end if;
    if not type(b1,[realcons,realcons,realcons]) and
        type(v1,[realcons,realcons,realcons]) then
        ERROR("points must evaluate to real constants")
    end if;
    hbase := evl(b1+.8*v1);
    hpoint := evl(b1+v1);
    opts := args[3 .. nargs];
    if indets({opts}) intersect {color,colour} = {} then
        opts := opts,colour = red
    end if;
    stem := plots[spacecurve]([b1,hbase],thickness = 3,opts);
    if v1 = [0,0,0] then return stem; end if;
    r := sqrt(linalg[norm](v1,2));
    v2 := evl(linalg[crossprod]([1,0,0],v1));
    if linalg[dotprod](v2,v2) < .01*linalg[dotprod](v1,v1) then
        v2 := evl(linalg[crossprod]([0,1,0],v1))
    end if;
    v2 := evl(.05*r*v2/linalg[norm](v2,2));
    v3 := linalg[crossprod](v1,v2);
    v3 := evl(.05*r*v3/linalg[norm](v3,2));
    head := plots[polygonplot3d]([hpoint,evl(hbase+v2),
        evl(hbase+v3), evl(hbase-v2),
        evl(hbase-v3), evl(hbase+v2)]
        ,style = patch,opts);
    display({stem,head})
end proc;

surfarrows := proc(fld,surf,r1,r2)
## Plots a 3d surface (parametric or cartesian) with arrows from
## a vector field at points on the surface
## Usage:
##   surfarrows(F(s,t), [x(s,t),y(s,t),z(s,t)], s=a1 .. b1, t=a2 .. b2, options)
##     for vectors F(s,t) at points [x(s,t), y(s,t), z(s,t)]
##     where F(s,t) is a list expression
##   surfarrows(F(x,y), z(x,y), x=a1 .. b1, y=a2 .. b2, options)
##     for vectors F(x,y) at points [x, y, z(x,y)]
##     where F(x,y) is a list expression
##   surfarrows(F, [x(s,t),y(s,t),z(s,t)], s=a1 .. b1, t=a2 .. b2, options)
##     for vectors F(x(s,t),y(s,t),z(s,t)) at [x(s,t),y(s,t),z(s,t)]
##     where F is a list-valued function of three variables
##   surfarrows(F, z(s,t), x=a1 .. b1, y=a2 .. b2, options)
##     for vectors F(x,y,z(s,t)] at [x,y,z(x,y)]
##     where F is a list-valued function of three variables
##   "options" are the usual 3d plot options, plus
##         arrowgrid= [n1, n2]
##     for arrows in an n1 by n2 grid, equally spaced in terms of the parameters
##     of the surface (default is arrowgrid=[8,8])
## Caution: inputs are not checked for validity

local evl,opts,gspec,p1,p2,s2,f2,v1,v2,a1,a2,b1,b2,n1,n2,h1,h2,i,j;
    evl := t -> convert(evalm(t), list):
    opts := args[5 .. nargs];
    v1 := op(1,r1);
    v2 := op(1,r2);
    a1 := op(1,op(2,r1));
    b1 := op(2,op(2,r1));
    a2 := op(1,op(2,r2));
    b2 := op(2,op(2,r2));
    gspec := indets({opts},identical('arrowgrid') = list);
    if gspec = {} then gspec := arrowgrid = [8,8]
    else opts := op({opts} minus gspec); gspec := gspec[1]
    end if;
    p1 := plot3d(surf,r1,r2,opts);
    opts :=
        op({opts} minus indets({opts},identical('grid') = list))
        ;
    n1 := op(2,gspec)[1]-1;
    n2 := op(2,gspec)[2]-1;
    h1 := (b1-a1)/n1;
    h2 := (b2-a2)/n2;
    if type(surf,list) then s2 := surf
    elif type(surf,vector) then s2 := evl(surf)
    else s2 := [v1,v2,surf]
    end if;
    if type(fld,procedure) then f2:= fld(op(s2))
    else f2:= evl(fld)
    end if;
    p2 := seq(seq(
     myarrow(op(subs(v1 = a1+h1*i,v2 = a2+j*h2,[s2,f2])),opts),
     i = 0 .. n1),j = 0 .. n2);
    display({p2,p1},opts)
end proc;

## Examples:
#
#> surfarrows((x,y,z) -> [-y, x, 0], x^2 + y^2, x=-1 .. 1, y=-1 .. 1);
#
#> torus:= [(2+cos(u))*cos(v), (2+cos(u))*sin(v), sin(u)]:
#> with(linalg):
#> unit:= V -> evl(V/norm(V,2)):
#> tnormal:= unit(crossprod(diff(torus,v),diff(torus,u))):
#> surfarrows( tnormal/2, torus, u=0 .. 2*Pi, v=0 .. 2*Pi, scaling=constrained);


end module:

#==============================================================================
savelibname := LIB_PATH:
savelib(plotting):
savelibname := 'savelibname':
#==============================================================================