#
#----------------------------------------------------------------------------#
export MeshNet
#--------------------------------------#
struct MeshNet
    nprt::Integer
    ndim::Integer

    pou::AbstractPOU
    bsp::AbstractBanachSpace
    C  ::AbstractArray
end
#--------------------------------------#
function MeshNet(pou::AbstractPOU
                ,bsp::AbstractBanachSpace)

    nprt = pou.nprt
    ndim = bsp.ndim

    init = Flux.glorot_uniform
    C = softmax(init(ndim,nprt);dims=1)

    return MeshNet(nprt,ndim,pou,bsp,C)
end
#--------------------------------------#
function MeshNet(nprt,ndim
                ,pouType::DataType=nnPOU
                ,bspType::DataType=Cosines)

    pou = pouType(nprt)
    bsp = bspType(nprt,ndim)

    return MeshNet(pou,bsp)
end
#--------------------------------------#
Flux.@functor MeshNet
DiffEqFlux.initial_params(net::MeshNet) = Flux.destructure(net)[1]
#--------------------------------------#
function (net::MeshNet)(x)  # [1,N]

    @unpack nprt,ndim = net
    @unpack pou,bsp,C = net

    N = size(x)[end]

    # partition of unity
    ϕ = pou(x)              # [nprt,N]
    ϕ = reshape(ϕ,1,nprt,N) # [1,nprt,N]

    # function space
    ψ = bsp(x)              # [ndim,nprt,N]

    u = @. ϕ * ψ * C        # [ndim,nprt,N]
    u = sum(u,dims=(1,2))   # [1,1,N]
    u = reshape(u,1,N)      # [1,N]

    return u
end
#----------------------------------------------------------------------------#
export MeshNet2D
#--------------------------------------#
struct MeshNet2D
    nprt::Integer
    ndim::Integer

    pou::AbstractPOU
    bsp::AbstractBanachSpace
    C  ::AbstractArray
end
#--------------------------------------#
function MeshNet2D(pou::AbstractPOU
                ,bsp::AbstractBanachSpace)

    nprt = pou.nprt
    ndim = bsp.ndim

    init = Flux.glorot_uniform
    C = softmax(init(ndim,nprt);dims=1)

    return MeshNet2D(nprt,ndim,pou,bsp,C)
end
#--------------------------------------#
function MeshNet2D(nprt,ndim
                ,pouType::DataType=nnPOU
                ,bspType::DataType=Cosines)

    pou = pouType(nprt)
    bsp = bspType(nprt,ndim)

    return MeshNet2D(pou,bsp)
end
#--------------------------------------#
Flux.@functor MeshNet2D
DiffEqFlux.initial_params(net::MeshNet2D) = Flux.destructure(net)[1]
#--------------------------------------#
function (net::MeshNet2D)(x,y)  # [1,N]

    @assert size(x) == size(y)

    @unpack nprt,ndim = net
    @unpack pou,bsp,C = net

    N = size(x)[end]

    # partition of unity
    ϕ = pou(x,y)             # [nprt,N]
    ϕ = reshape(ϕ,1,nprt,N) # [1,nprt,N]

    # function space
    ψ = bsp(x,y)            # [ndim,nprt,N]

    u = @. ϕ * ψ * C        # [ndim,nprt,N]
    u = sum(u,dims=(1,2))   # [1,1,N]
    u = reshape(u,1,N)      # [1,N]

    return u
end
#--------------------------------------#
function (net::MeshNet2D)(xy) # [2,N]
    x = xy[1,:]
    y = xy[2,:]
    return net(x,y)
end
#--------------------------------------#
export leastSq!
#--------------------------------------#
function leastSq!(net::MeshNet,x,u)

    @unpack nprt,ndim = net
    @unpack pou,bsp,C = net

    N = size(x)[end]

    # form g
    ϕ = pou(x)                  # [nprt,N]
    ϕ = reshape(ϕ,1,nprt,N)     # [1,nprt,N]
    ψ = bsp(x)                  # [ndim,nprt,N]
    g  = @. ϕ * ψ               # [ndim,nprt,N]

    # least square solve
    g  = reshape(g,ndim*nprt,N) # size [N,ndim*nprt]
    c  = lsq(u',g')
    C .= reshape(c,ndim,nprt)   # size [ndim,nprt]

    return
end
#--------------------------------------#
function leastSq!(net::MeshNet2D,x,y,u)

    @unpack nprt,ndim = net
    @unpack pou,bsp,C = net

    N = size(x)[end]

    # form g
    ϕ = pou(x,y)                # [nprt,N]
    ϕ = reshape(ϕ,1,nprt,N)     # [1,nprt,N]
    ψ = bsp(x,y)                # [ndim,nprt,N]
    g  = @. ϕ * ψ               # [ndim,nprt,N]

    # least square solve
    g  = reshape(g,ndim*nprt,N) # size [N,ndim*nprt]
    c  = lsq(u',g')
    C .= reshape(c,ndim,nprt)   # size [ndim,nprt]

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