From 005ce11313ca2ff927f7bc29da74743e02d19278 Mon Sep 17 00:00:00 2001
From: Thomas Breloff <tom@breloff.com>
Date: Thu, 28 Apr 2016 22:03:52 -0400
Subject: [PATCH] InputWrapper and pyplot surface zcolor

---
 src/Plots.jl           |  2 ++
 src/args.jl            |  4 +++-
 src/backends/pyplot.jl | 26 +++++++++++++++++++++++---
 src/types.jl           |  8 ++++++++
 4 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/src/Plots.jl b/src/Plots.jl
index 1367ca67..f92ee7c7 100644
--- a/src/Plots.jl
+++ b/src/Plots.jl
@@ -20,6 +20,8 @@ export
     AVec,
     AMat,
     KW,
+
+    wrap,
     set_theme,
     add_theme,
 
diff --git a/src/args.jl b/src/args.jl
index 6d279a42..2efcd01b 100644
--- a/src/args.jl
+++ b/src/args.jl
@@ -754,19 +754,21 @@ end
 
 # 1-row matrices will give an element
 # multi-row matrices will give a column
+# InputWrapper just gives the contents
 # anything else is returned as-is
 # getArgValue(v::Tuple, idx::Int) = v[mod1(idx, length(v))]
 function getArgValue(v::AMat, idx::Int)
     c = mod1(idx, size(v,2))
     size(v,1) == 1 ? v[1,c] : v[:,c]
 end
+getArgValue(wrapper::InputWrapper, idx) = wrapper.obj
 getArgValue(v, idx) = v
 
 
 # given an argument key (k), we want to extract the argument value for this index.
 # if nothing is set (or container is empty), return the default.
 function setDictValue(d_in::KW, d_out::KW, k::Symbol, idx::Int, defaults::KW)
-    if haskey(d_in, k) && !(typeof(d_in[k]) <: @compat(Union{AbstractArray, Tuple}) && isempty(d_in[k]))
+    if haskey(d_in, k) && !(typeof(d_in[k]) <: Union{AbstractArray, Tuple} && isempty(d_in[k]))
         d_out[k] = getArgValue(d_in[k], idx)
     else
         d_out[k] = defaults[k]
diff --git a/src/backends/pyplot.jl b/src/backends/pyplot.jl
index 29a118a6..dad7693d 100644
--- a/src/backends/pyplot.jl
+++ b/src/backends/pyplot.jl
@@ -11,7 +11,8 @@ function _initialize_backend(::PyPlotBackend)
         const pypatches = PyPlot.pywrap(PyPlot.pyimport("matplotlib.patches"))
         const pyfont = PyPlot.pywrap(PyPlot.pyimport("matplotlib.font_manager"))
         const pyticker = PyPlot.pywrap(PyPlot.pyimport("matplotlib.ticker"))
-        # const pycolorbar = PyPlot.pywrap(PyPlot.pyimport("matplotlib.colorbar"))
+        const pycmap = PyPlot.pywrap(PyPlot.pyimport("matplotlib.cm"))
+        const pynp = PyPlot.pywrap(PyPlot.pyimport("numpy"))
     end
 
     if !isa(Base.Multimedia.displays[end], Base.REPL.REPLDisplay)
@@ -51,6 +52,15 @@ getPyPlotColorMap(v::AVec, α=nothing) = getPyPlotColorMap(ColorGradient(v), α)
 # anything else just gets a bluesred gradient
 getPyPlotColorMap(c, α=nothing) = getPyPlotColorMap(default_gradient(), α)
 
+function getPyPlotCustomShading(c, z, α=nothing)
+    cmap = getPyPlotColorMap(c, α)
+    # sm = pycmap.pymember("ScalarMappable")(cmap = cmap)
+    # sm[:set_array](z)
+    # sm
+    ls = pycolors.pymember("LightSource")(270,45)
+    ls[:shade](z, cmap, vert_exag=0.1, blend_mode="soft")
+end
+
 # get the style (solid, dashed, etc)
 function getPyPlotLineStyle(linetype::Symbol, linestyle::Symbol)
     linetype == :none && return " "
@@ -508,8 +518,13 @@ function _add_series(pkg::PyPlotBackend, plt::Plot, d::KW)
                 z = z'
             end
             if lt == :surface
-                extrakw[:cmap] = pyfillcolormap(d)
-                needs_colorbar = true
+                if d[:marker_z] != nothing
+                    extrakw[:facecolors] = getPyPlotCustomShading(d[:fillcolor], d[:marker_z], d[:fillalpha])
+                    extrakw[:shade] = false
+                else
+                    extrakw[:cmap] = pyfillcolormap(d)
+                    needs_colorbar = true
+                end
             end
             handle = ax[lt == :surface ? :plot_surface : :plot_wireframe](x, y, z;
                 label = d[:label],
@@ -536,6 +551,11 @@ function _add_series(pkg::PyPlotBackend, plt::Plot, d::KW)
                 end
             end
 
+            # no colorbar if we are creating a surface LightSource
+            if haskey(extrakw, :facecolors)
+                needs_colorbar = false
+            end
+
         elseif typeof(z) <: AbstractVector
             # tri-surface plot (http://matplotlib.org/mpl_toolkits/mplot3d/tutorial.html#tri-surface-plots)
             handle = ax[:plot_trisurf](x, y, z;
diff --git a/src/types.jl b/src/types.jl
index 61f34e04..fac40a3e 100644
--- a/src/types.jl
+++ b/src/types.jl
@@ -8,6 +8,14 @@ abstract AbstractBackend
 abstract AbstractPlot{T<:AbstractBackend}
 
 typealias KW Dict{Symbol,Any}
+
+immutable InputWrapper{T}
+    obj::T
+end
+
+wrap{T}(obj::T) = InputWrapper{T}(obj)
+Base.isempty(wrapper::InputWrapper) = false
+
 # -----------------------------------------------------------
 # Plot
 # -----------------------------------------------------------