#
#----------------------------------------------------------------------------#
abstract type AbstractPOU end # partition of unity
#----------------------------------------------------------------------------#
export nnPOU
#--------------------------------------#
struct nnPOU <: AbstractPOU
    nprt::Integer
    NN::Chain
end
#--------------------------------------#
function nnPOU(nprt::Integer;
               depth=2,
               act=relu,
               init=Flux.glorot_uniform)

    """
     relu activation function for sharp (piecewise linear)
     partition functions necessary for representing piecewise
     continuous functions
    """
    
    NN = Chain(Dense(init(nprt,1),init(nprt),act)
              ,Chain([Dense(init(nprt,nprt),init(nprt),act)
                                           for i in 1:(depth-1)]...)
              ,softmax)

    return nnPOU(nprt,NN)
end
#--------------------------------------#
Flux.@functor nnPOU
DiffEqFlux.initial_params(pou::nnPOU) = Flux.destructure(pou)[1]
#--------------------------------------#
(pou::nnPOU)(x) = pou.NN(x)
#----------------------------------------------------------------------------#
export nnPOU2D
#--------------------------------------#
struct nnPOU2D <: AbstractPOU
    nprt::Integer
    sdf
    tol::Number
    NN::Chain
end
#--------------------------------------#
function nnPOU2D(nprt::Integer,sdf
                ;tol=1f-8
                ,depth=4
                ,act=relu
                ,init=Flux.glorot_uniform)

    """
     integrate signed distance function knowledge in neural network
    """

    NN = Chain(Dense(init(nprt,2),init(nprt),act)
              ,Chain([Dense(init(nprt,nprt),init(nprt),act)
                                           for i in 1:(depth-1)]...)
              ,softmax)

    return nnPOU2D(nprt,sdf,tol,NN)
end
#--------------------------------------#
Flux.@functor nnPOU2D
DiffEqFlux.initial_params(pou::nnPOU2D) = Flux.destructure(pou)[1]
#--------------------------------------#
function (pou::nnPOU2D)(x,y)
    @unpack sdf, tol, NN = pou

    ϕ = NN(vcat(x,y))

    Ω = Zygote.ignore() do
        Ω = sdf(x,y) .<= 0f0 + tol
        return Ω
    end

    ϕ = ϕ .* Ω # clip partitions at domain boundary

    return ϕ
end
#----------------------------------------------------------------------------#
export femPOU
#--------------------------------------#
"""
 FEM partition of unity with arbitrary shape function
"""
struct femPOU <: AbstractPOU
    nprt::Integer
    shape::Function
    xe::AbstractVector # interior points
    x0::Number         # end points
    x1::Number
end

#--------------------------------------#
function femPOU(nprt::Integer
               ,shape = x -> x # linear FEM hat function
               ;x0::Number=-1f0
               ,x1::Number= 1f0)

    """
     shape: shape function defined on [0,1]
     s.t shape(0.0) = 0, shape(1.0) = 1
    """

    xg = Array(range(x0,x1,length=nprt))

    return femPOU(nprt,shape,xg[2:end-1],x0,x1)
end
#--------------------------------------#
Flux.@functor femPOU
DiffEqFlux.initial_params(pou::femPOU) = Flux.destructure(pou)[1]
#--------------------------------------#
function (pou::femPOU)(x)
    @unpack shape,nprt,xe,x0,x1 = pou

    xg = addEndPts(xe,x0,x1)

    # allow partitions to collapse onto themselves
#   xg = sort(xg)

    n  = length(xg)
    h  = diff(xg)
    hi = 1.0 ./ h 

    RR,b2 = Zygote.ignore() do
        RR = restrictUpper(n) + restrictLower(n)
        RR = RR'
        b2 = zero(xg); b2[1] = 1
        return RR,b2
    end

    hh = -RR*hi
    W2 = Tdg(hi,hh,hi)

    ϕ = Chain(x -> relu.(x .- xg)
             ,u -> (W2*u .+ b2)
             ,@. shape)

    u = ϕ(x)

    return u
end
#----------------------------------------------------------------------------#
export femPOU2D
#--------------------------------------#
"""
 2D FEM partition of unity with arbitrary shape function
"""
struct femPOU2D <: AbstractPOU
    nprt::Integer
    shape::Function
    sdf  ::Function # range [-1,0]
end
#--------------------------------------#
Flux.@functor femPOU2D
DiffEqFlux.initial_params(pou::femPOU2D) = Flux.destructure(pou)[1]
#--------------------------------------#
function (pou::femPOU2D)(x,y)

    @unpack nprt, shape, sdf = pou

    N = size(x)[end]

    d  = sdf(x,y) # range [-1,0]
    dd = 1 .+ d   # range [0 ,1]
    s = shape(dd)

    u = vcat(s, 1 .- s)

    return u
end
#----------------------------------------------------------------------------#
export indicatorPOU
#--------------------------------------#
"""
 Indicator function over grid
"""
struct indicatorPOU <: AbstractPOU
    nprt::Integer
    xe::AbstractVector # interior points
    x0::Number         # end points
    x1::Number
end

function indicatorPOU(nprt
                     ;x0::Number=-1f0
                     ,x1::Number= 1f0)


    xg = Array(range(x0,x1,length=nprt+1))

    return indicatorPOU(nprt,xg[2:end-1],x0,x1)
end
#--------------------------------------#
Flux.@functor indicatorPOU
DiffEqFlux.initial_params(pou::indicatorPOU) = Flux.destructure(pou)[1]
#--------------------------------------#
function (pou::indicatorPOU)(x) # [1,N]

    @unpack nprt,xe,x0,x1 = pou

    N = size(x)[end]
    y = zeros(nprt,N)

    xg = addEndPts(xe,x0,x1) # [nprt+1]

    Ru,Rl = Zygote.ignore() do
        Ru = restrictUpper(nprt+1)
        Rl = restrictLower(nprt+1)
        return Ru,Rl
    end

    xm = Rl*xg # [nprt]
    xM = Ru*xg

    """
    julia> xp
    1×1000 adjoint(::Vector{Float32}) with eltype Float32:
     -1.0  -0.997998  -0.995996  …  0.993994  0.995996  0.997998  1.0
    julia> sign.(xp)
    1×1000 Matrix{Float32}:
     -1.0  -1.0  -1.0  -1.0  -1.0  -1.0  …  1.0  1.0  1.0  1.0  1.0  1.0  1.0

    #########

    julia> fun(x)=sum(x .< 1);ps=Flux.params(xp);
    julia> gs = gradient(() -> fun(xp),ps);
    julia> gs[xp]
    gs[xp]
    > nothing

    #########

    julia> fun(x)=sum(sign.(x));ps=Flux.params(xp);
    julia> gs = gradient(() -> fun(xp),ps);
    julia> gs[xp]
    > 1×1000 Matrix{Float32}:
     0.0  0.0  0.0  0.0  0.0  0.0  0.0  …  0.0  0.0  0.0  0.0  0.0  0.0  0.0
    """
    y = xm .<= x .<= xM # [nprt,N]

    return y
end
#--------------------------------------#
