using Plots, LinearAlgebra, ForwardDiff




f(x, y) = x^2 + y^2

xs = range(-5, stop=5, length=50)
ys = range(-5, stop=5, length=50)
zs = [f(x,y) for y in ys, x in xs]

surface(xs, ys, zs)

using SymPy
@vars x y real=true
zs = [x^2 - y^2 for y in ys, x in xs]
plot(xs, ys, N.(zs))

f(x, y) = x^2 + y^2
zs = [f(x,y) for y in ys, x in xs]
contour(xs, ys, zs)

f(x, y) = x^2 + y^2
zs = [f(x,y) for y in ys, x in xs]

using Interact

@manipulate for azim in -90:90, elev in 0:90; do
  plot(xs, ys, zs, camera=(azim, elev))
end

xs_ys(vs) = (A=hcat(vs...); Tuple([A[i,:] for i in eachindex(vs[1])]))
xs_ys(v,vs...) = xs_ys([v, vs...])
xs_ys(r::Function, a, b, n=100) = xs_ys(r.(range(a, stop=b, length=n)))

function arrow!(plt::Plots.Plot, p, v; kwargs...)
  if length(p) == 2
     quiver!(plt, xs_ys([p])..., quiver=Tuple(xs_ys([v])); kwargs...)
  elseif length(p) == 3
    # 3d quiver needs support
    # https://github.com/JuliaPlots/Plots.jl/issues/319#issue-159652535
    # headless arrow instead
    plot!(plt, xs_ys(p, p+v)...; kwargs...)
	end
end
arrow!(p,v;kwargs...) = arrow!(Plots.current(), p, v; kwargs...)

g(x,y) = x^2 + 2y^2
xs = ys = range(-2, stop=2, length=50)
zs = [g(x,y) for y in ys, x in xs]

contour(xs, ys, zs, levels=[1,2,3,4,5])
arrow!([ 1,1], 0.25*[ 1,-1])  ## start at [1,1] in direction [1,-1]
arrow!([-1,1], 0.25*[-1, 1])

## Spherical
X(r,theta,phi) = r * sin(theta) * sin(phi)
Y(r,theta,phi) = r * sin(theta) * cos(phi)
Z(r,theta,phi) = r * cos(theta)

thetas = range(0, stop=pi,   length=250)
phis   = range(0, stop=pi/2, length=250)
xs = [X(1, theta, phi) for theta in thetas, phi in phis]
ys = [Y(1, theta, phi) for theta in thetas, phi in phis]
zs = [Z(1, theta, phi) for theta in thetas, phi in phis]

surface(xs, ys, zs)

H(theta, phi) = [X(1, theta, phi), Y(1, theta, phi), Z(1, theta, phi)]
out = [ [H(theta, phi)[i] for theta in thetas, phi in phis] for i in 1:3]

surface(out...)

p = [1,1,1]
u = [1,0,1]
v = [0,1,1]
@assert norm(u × v) > 0   # \times<tab> for ×
plane(a,b) = p + a*u + b*v

as = bs = range(-1, stop=1, length=50)

xs = [plane(a, b)[1] for b in bs, a in as]
ys = [plane(a, b)[2] for b in bs, a in as]
zs = [plane(a, b)[3] for b in bs, a in as]

surface(xs, ys, zs)
arrow!(p, u, color="red")
arrow!(p, v, color="yellow")

f(x,y) = x^2 - 2x*y + y^3

f(x) = x[1]^2 + 2x[1] * x[2] + x[2]^3

f(x,y) = x^2 - 2x*y + y^3
f(x) = f(x...)

using ForwardDiff
import ForwardDiff: gradient
gradient(f, [1,2])

dx(x,y) = 2x - 2y
dy(x,y) = -2x + 3y^2
(dx(1,2), dy(1,2))

@vars x y real=true
f(x,y)  = x^2 - 2x*y + y^3
ex = f(x, y)
partial_x = diff(ex, x, 1)
partial_y = diff(ex, y)    # 1 is default degree
[partial_x, partial_y]     # The gradient, or
[diff(ex, s) for s in (x, y)]  # finds gradient in one swoop

f(x,y)  = x^2 - 2x*y + y^3
ex = f(x, y)

hess = [diff(diff(ex, a), b) for b in (x,y), a in (x,y)]

sympy.hessian(ex, (x,y))

Grad(ex::SymPy.Sym, vars=free_symbols(ex)) = [diff(ex, v) for v in vars]

f(x,y) = x^2 + y^2
grad_f(x,y) = [2*x, 2*y]
xs = ys = range(-3, stop=3, length=50)
zs = [f(x,y) for y in ys, x in xs]
contour(xs, ys, zs, levels = [1,2,4,9], aspect_ratio=:equal)

## add gradient
arrow!([sqrt(3)/2,1/2], grad_f(sqrt(3)/2,1/2))

## tangent to contour by implicit differentiation solves 2x + 2y y' = 0
## or y' = -x/y (which is in direction [y, -x])
arrow!([sqrt(3)/2,1/2], [1/2, -sqrt(3)/2])

F = SymFunction("F")
@vars x y real=true

ex = x^2 + y^2
tmp = diff(subs(ex, y => F(x)), x)  # x^2 + F(x)^2
ex1 = solve(tmp, diff(F(x),x))    # solve for dF/dx in 2x + 2F(x) dF/dx, gives [-x/F(x)]
tl = subs(ex1[1], F(x) => y)                # now it is [-x/y]

[diff(ex,s) for s in [x,y]] ⋅ [1, tl]      # 0 as orthogonal.

f1(x,y) = y^2 - y^4 - x^2

xs = ys = range(-1, stop=1, length=50)
zs = [f1(x,y) for y in ys, x in xs]
surface(xs, ys, zs)

g(x,y) =  2 + 2x + 2y - x^2 - y^2
h(x,y) = (x>0) & (y>0) & (y < 9-x) ? g(x,y) : NaN

xs = ys = range(0, stop=9, length=250)
zs = [h(x,y) for y in ys, x in xs]

surface(xs, ys, zs)

@vars x y real=true
grad = [diff(g(x,y), x), diff(g(x,y),y)]
solve(grad)

subs.(grad, x => 1, y => 1)

hess = sympy.hessian(g(x,y), (x,y))
det(hess)

g(1,1)

h(y) = g(0,y)
out = solve(diff(h(y), y))
append!(out, [0, 9])
h.(out)

h(x) = g(x,0)
out = solve(diff(h(x), x))
append!(out, [0, 9])
h.(out)

@vars t

h(t) = g(t, 9-t)
out = solve(diff(h(t), t))
append!(out, [0, 9])
h.(out)

g(x,y) = x^2 + y^2
∇g(x,y) = [2x, 2y]

f(x,y) = 5 + x + y^2
∇f(x,y) = [1, 2*y]

uvec(u) = u/norm(u)

xs = ys = range(-3, stop=3, length=50) ## should be square!
ts = [0, pi/6, pi/3, pi/2]
zcs = [g(x,y) for y in ys, x in xs]

p = contour(xs, ys, zcs,  levels=[1,4,9])

## parameterize g(x,y) == 1
rad(t) = [cos(t), sin(t)]

for t in ts
    arrow!(p, rad(t), uvec(∇g(rad(t)...)), color=:blue)
    arrow!(p, rad(t), uvec(∇f(rad(t)...)), color=:red )
end

p

using SymPy
@vars x y real=true

f(x,y) = 5 + x + y^2
g(x,y) = x^2 + y^2
gradf = Grad(f(x,y), [x,y])    # add [x,y] to ensure order
gradg = Grad(g(x,y), [x,y])
[gradf gradg]   # gradf in first column, gradg in second

@vars lambda real=true
out = solve(gradf - lambda * gradg)

ex = subs(g(x,y), x => 1//2) # or out[1][x] to fish out the 1/2
solve(ex - 1, y)           # two answers

exs = gradf - lambda * gradg

push!(exs, g(x,y) - 1)
out = solve(exs)

[f(o[x], o[y]) for o in out]

@vars a b c real=true

Vol(a,b,c) = a*b*c
SA(a,b,c) = 2*(a*b + a*c + b*c)

gradf = Grad(Vol(a,b,c), [a,b,c])
gradg = Grad(SA(a,b,c),  [a,b,c])

[gradf  gradg]   # gradf in first column,...

@vars lambda  real=true

exs = gradf - lambda * gradg
push!(exs, SA(a,b,c) - 64)
out = solve(exs, [lambda, a, b, c])
out[2]

@vars a b c real=true
@vars lambda  real=true

Vol(a,b,c) = a*b*c
SA(a,b,c) = 2*(a*b + a*c + b*c)

gradf = Grad(Vol(a,b,c), [a,b,c])
gradg = Grad(SA(a,b,c),  [a,b,c])

exs = gradf - lambda * gradg
push!(exs, Vol(a,b,c) - 64)

solve(exs)

@vars radius height real=true
@vars lambda real=true

Vol(r,h) = PI*r^2* h
SA(r,h) = 2*PI*r^2 + 2*PI*r*h
gradf = Grad(Vol(radius, height), [radius, height])
gradg = Grad(SA(radius, height),  [radius, height])


exs = gradf - lambda * gradg
push!(exs, Vol(radius,height) - 355)

solve(exs)[1]

solve(gradf - lambda * gradg)[3]  # third one has answer