using Pkg
Pkg.add("ForwardDiff")
using LinearAlgebra
using AutoGrad
using ForwardDiff

struct LDder{v,w}
    val::v
    der::Array{w,1}
end

# Identifies independent variable for differentiation
struct var{v}
    val::v
end

# output the number of directions in the LDder
ndir(A::LDder) = length(A.der)
ndir(A::Array) = length(A[1].der)

# Converts var inputs in an expression to LDderivatives
# If no direction matrix is specified, the identity is used
function derinit(ex::Expr,M)
    args = ex.args
    
    for i in 2:length(args)
        x = args[i]
        
        if isa(x, Expr)
            args[i] = derinit(x,M)
        else
            args[i] = @eval derinit($x,$M)
        end
    end
    eval(ex)
end
function derinit(x::var, M)
    val = x.val 
    lx = length(val)
    X = Array{LDder}(undef,lx)
    
    if M == 0
        M = diagm(0 => ones(lx))
    end
    
    for i = 1:lx
        X[i] = LDder(val[i],M[i,:])
    end
    X
end
derinit(p,M) = p

# Converts array of LDder structures to separate arrays of values and derivatives
function LD2mat(X::Array)
    lx = length(X)
    nd = ndir(X)
    val = zeros(lx)
    der = zeros(lx,nd)
    for i = 1:lx
        val[i] = X[i].val
        der[i,:] = X[i].der
    end
    val, der
end
function LD2mat(x::LDder)
    val = x.val;
    der = x.der;
    val, der
end

# Calculates Jacobian from LDder structures
function jacobian(X)
    v,d = LD2mat(X)
    d
end
function jacobian(X::LDder,M)
    v,d = LD2mat(X)
    J = d'*inv(M)
end
function jacobian(X::Array,M)
    v,d = LD2mat(X)
    J = d*inv(M)
end

# Macro for calculating LDder for existing functions
macro nsdiff(ex, M = 0)
    derinit(ex,M)
end

# Macros for isolating values, LD-derivatives, and Jacobians from LD-derivative objects
macro val(ex)
    v,d = @eval LD2mat($ex)
    v
end
macro der(ex)
    v,d = @eval LD2mat($ex)
    d
end
macro jacobian(ex)
    @eval jacobian($ex)
end
macro jacobian(ex,M)
    @eval jacobian($ex,M)
end

Base.:-(A::LDder) = LDder(-A.val,-A.der)
Base.:-(A::LDder,B::LDder) = LDder(A.val-B.val,A.der-B.der)
Base.:-(A::LDder,b) = LDder(A.val-b,A.der)
Base.:-(b,A::LDder) = LDder(b-A.val,-A.der)

Base.:+(A::LDder,B::LDder) = LDder(A.val+B.val,A.der+B.der)
Base.:+(A::LDder,b) = LDder(A.val+b,A.der)
Base.:+(b,A::LDder) = LDder(A.val+b,A.der)

Base.:*(A::LDder,B::LDder) = LDder(A.val*B.val,A.der*B.val+B.der*A.val)
Base.:*(A::LDder,b) = LDder(b*A.val,b*A.der)
Base.:*(b,A::LDder) = LDder(b*A.val,b*A.der)

Base.inv(A::LDder) = LDder(1/A.val, -1/(A.val^2)*A.der)
Base.:/(A::LDder,B::LDder) = LDder(A.val/B.val, (B.val*A.der-A.val*B.der)/(B.val^2))
Base.:/(A::LDder,b) = A*inv(b)
Base.:/(b,A::LDder) = LDder(b/A.val, -b/(A.val^2)*A.der)
Base.:\(A::LDder,B::LDder) = LDder(B.val/A.val, (A.val*B.der-B.val*A.der)/(A.val^2))
Base.:\(A::LDder,b) = LDder(b/A.val, -b/(A.val^2)*A.der)
Base.:\(b,A::LDder) = A*inv(b)

Base.:^(A::LDder,b) = LDder(A.val^b,b*A.val^(b-1)*A.der)
Base.:^(a, B::LDder) = LDder(a^B.val,log(a)*a^B.val*B.der)

Base.sin(A::LDder) = LDder(sin(A.val),cos(A.val)*A.der)
Base.cos(A::LDder) = LDder(cos(A.val),-sin(A.val)*A.der)
Base.tan(A::LDder) = LDder(tan(A.val),A.der/(cos(A.val)^2))

function fsign(x)
    i = 1;
    lx = length(x)
    while (x[i] == 0) & (i<lx)
        i += 1
    end
    sign(x[i])
end
lmin(x,y) = fsign(x-y) <= 0 ? x : y
lmax(x,y) = fsign(x-y) >= 0 ? x : y
function lmid(x,y,z)
    xy = fsign(x-y)
    yz = fsign(y-z)
    zx = fsign(z-x)
    if (xy>=0 && zx>=0) | (xy<=0 && zx<=0)
        f = x
    elseif (xy>=0 && yz>=0) | (xy<=0 && yz<=0)
        f = y
    else
        f = z
    end
    f
end
double(X::LDder) = [X.val,X.der...]
Base.abs(A::LDder) = LDder(abs(A.val),fsign([A.val, A.der...])*A.der)
function Base.min(A::LDder, B::LDder)
    z = lmin([A.val, A.der...],[B.val, B.der...])
    LDder(min(A.val,B.val),z[2:end])
end 
function Base.min(A::LDder, b)
    z = lmin([A.val, A.der...],[b, zeros(ndir(A))...])
    LDder(min(A.val,b),z[2:end])
end 
Base.min(b, A::LDder) = min(A,b)
function Base.max(A::LDder, B::LDder)
    z = lmax([A.val, A.der...],[B.val, B.der...])
    LDder(max(A.val,B.val),z[2:end])
end 
function Base.max(A::LDder, b)
    z = lmax([A.val, A.der...],[b, zeros(ndir(A))...])
    LDder(max(A.val,b),z[2:end])
end 
Base.max(b, A::LDder) = max(A,b)
function mid(x,y,z)
    if (x>=y && x<=z) | (x<=y && x>=z)
        f = x
    elseif (y>=x && y<=z) | (y<=x && y>=z)
        f = y
    else
        f = z
    end
end
function mid(X::LDder,y,z)
    nd = ndir(X)
    d = zeros(nd)
    Y = LDder(y,d)
    Z = LDder(z,d)
    mid(X,Y,Z)
end
mid(x,Y::LDder,z) = mid(Y,x,z)
mid(x, y, Z::LDder) = mid(Z,x,y)
function mid(X::LDder,Y::LDder,z)
    nd = ndir(X)
    d = zeros(nd)
    Z = LDder(z,d)
    mid(X,Y,Z)
end
mid(X::LDder, y, Z::LDder) = mid(X,Z,y)
mid(x, Y::LDder, Z::LDder) = mid(Y,Z,x)
function mid(X::LDder, Y::LDder, Z::LDder)
    x = double(X)
    y = double(Y)
    z = double(Z)
    f = lmid(x,y,z)
    LDder(mid(X.val,Y.val,Z.val),f[2:end])
end

function hserve(x)
    a = x[1]
    v = x[2]
    h = x[3]
    f = v^2*cos(a)^2/32*(tan(a)+(tan(a)^2+64*h/(v^2*cos(a)^2))^0.5)
end

y = var([0.35,44,9])
@nsdiff a = hserve(y)
@jacobian a

M = [-1 0 0; 0 -1 0; 0 0 -1]
@nsdiff a = hserve(y) M
@jacobian a M

sScalar(x) = x[1]*x[1]+2x[1]/x[2]
x = var([1;1])
f = @nsdiff2 sScalar(x)
@hessian f

nsVec1(x) = [abs(x[1]-x[2]);max(0,min(x[1],x[2]))]

x = [0,0]
ForwardDiff.jacobian(nsVec1,x)

y = var([0,0])

a = @nsdiff nsVec1(y)

@jacobian a

M = [-1 0; 0 -1]
b = @nsdiff nsVec1(y) M
@jacobian b M

function nsVec2(x)
    x1 = x[1]
    x2 = x[2]
    f1 = mid(x1+1,0,x1-1)
    f2 = mid(x2+1,x2-x1,x2-1)
    [f1,f2]
end

M = [0 -1; -1 0]
x = var([-1,-1])
@nsdiff c = nsVec2(x)
@jacobian c