text:
import Base.Broadcast
Base.Broadcast.dottable(x::Function) = true
RuntimeGeneratedFunctions.init(@__MODULE__)
"""
Algorithm for solving Physics-Informed Neural Networks problems.

Arguments:
* `chain`: a Flux.jl chain with a d-dimensional input and a 1-dimensional output,
* `strategy`: determines which training strategy will be used,
* `init_params`: the initial parameter of the neural network,
* `phi`: a trial solution,
* `derivative`: method that calculates the derivative.

"""
abstract type AbstractPINN{isinplace} <: SciMLBase.SciMLProblem end

struct PhysicsInformedNN{isinplace,C,T,P,PH,DER,PE,AL,K} <: AbstractPINN{isinplace}
  chain::C
  strategy::T
  init_params::P
  phi::PH
  derivative::DER
  param_estim::PE
  additional_loss::AL
  kwargs::K

 @add_kwonly function PhysicsInformedNN{iip}(chain,
                                             strategy;
                                             init_params = nothing,
                                             phi = nothing,
                                             derivative = nothing,
                                             param_estim=false,
                                             additional_loss=nothing,
                                             kwargs...) where iip
        if init_params == nothing
            if chain isa AbstractArray
                initθ = DiffEqFlux.initial_params.(chain)
            else
                initθ = DiffEqFlux.initial_params(chain)
            end

        else
            initθ = init_params
        end
        type_initθ = if (typeof(chain) <: AbstractVector) Base.promote_typeof.(initθ)[1] else  Base.promote_typeof(initθ) end
        parameterless_type_θ = DiffEqBase.parameterless_type(type_initθ)

        if phi == nothing
            if chain isa AbstractArray
                _phi = get_phi.(chain,parameterless_type_θ)
            else
                _phi = get_phi(chain,parameterless_type_θ)
            end
        else
            _phi = phi
        end

        if derivative == nothing
            _derivative = get_numeric_derivative()
        else
            _derivative = derivative
        end
        new{iip,typeof(chain),typeof(strategy),typeof(initθ),typeof(_phi),typeof(_derivative),typeof(param_estim),typeof(additional_loss),typeof(kwargs)}(chain,strategy,initθ,_phi,_derivative,param_estim,additional_loss,kwargs)
    end
end
PhysicsInformedNN(chain,strategy,args...;kwargs...) = PhysicsInformedNN{true}(chain,strategy,args...;kwargs...)

SciMLBase.isinplace(prob::PhysicsInformedNN{iip}) where iip = iip

abstract type TrainingStrategies  end

"""
* `dx`: the discretization of the grid.
"""
struct GridTraining <: TrainingStrategies
    dx
end

"""
* `points`: number of points in random select training set,
* `bcs_points`: number of points in random select training set for boundry conditions (by default, it equals `points`).
"""
struct StochasticTraining <:TrainingStrategies
    points:: Int64
    bcs_points:: Int64
end
function StochasticTraining(points;bcs_points = points)
    StochasticTraining(points, bcs_points)
end
"""
* `points`:  the number of quasi-random points in a sample,
* `bcs_points`: the number of quasi-random points in a sample for boundry conditions (by default, it equals `points`),
* `sampling_alg`: the quasi-Monte Carlo sampling algorithm,
* `resampling`: if it's false - the full training set is generated in advance before training,
   and at each iteration, one subset is randomly selected out of the batch.
   if it's true - the training set isn't generated beforehand, and one set of quasi-random
   points is generated directly at each iteration in runtime. In this case `minibatch` has no effect,
* `minibatch`: the number of subsets, if resampling == false.

For more information look: QuasiMonteCarlo.jl https://github.com/SciML/QuasiMonteCarlo.jl
"""
struct QuasiRandomTraining <:TrainingStrategies
    points:: Int64
    bcs_points:: Int64
    sampling_alg::QuasiMonteCarlo.SamplingAlgorithm
    resampling:: Bool
    minibatch:: Int64
end
function QuasiRandomTraining(points;bcs_points = points, sampling_alg = LatinHypercubeSample(),resampling =true, minibatch=0)
    QuasiRandomTraining(points,bcs_points,sampling_alg,resampling,minibatch)
end
"""
* `quadrature_alg`: quadrature algorithm,
* `reltol`: relative tolerance,
* `abstol`: absolute tolerance,
* `maxiters`: the maximum number of iterations in quadrature algorithm,
* `batch`: the preferred number of points to batch.

For more information look: Quadrature.jl https://github.com/SciML/Quadrature.jl
"""
struct QuadratureTraining <: TrainingStrategies
    quadrature_alg::DiffEqBase.AbstractQuadratureAlgorithm
    reltol::Float64
    abstol::Float64
    maxiters::Int64
    batch::Int64
end

function QuadratureTraining(;quadrature_alg=CubatureJLh(),reltol= 1e-6,abstol= 1e-3,maxiters=1e3,batch=100)
    QuadratureTraining(quadrature_alg,reltol,abstol,maxiters,batch)
end

"""
Create dictionary: variable => unique number for variable

# Example 1

Dict{Symbol,Int64} with 3 entries:
  :y => 2
  :t => 3
  :x => 1

# Example 2

 Dict{Symbol,Int64} with 2 entries:
  :u1 => 1
  :u2 => 2
"""
get_dict_vars(vars) = Dict( [Symbol(v) .=> i for (i,v) in enumerate(vars)])

# Wrapper for _transform_expression
function transform_expression(ex,dict_indvars,dict_depvars,chain,eltypeθ,strategy)
    if ex isa Expr
        ex = _transform_expression(ex,dict_indvars,dict_depvars,chain,eltypeθ,strategy)
    end
    return ex
end

function get_ε(dim, der_num,eltypeθ)
    epsilon = cbrt(eps(eltypeθ))
    ε = zeros(eltypeθ, dim)
    ε[der_num] = epsilon
    ε
end

θ = gensym("θ")

"""
Transform the derivative expression to inner representation

# Examples

1. First compute the derivative of function 'u(x,y)' with respect to x.

Take expressions in the form: `derivative(u(x,y), x)` to `derivative(phi, u, [x, y], εs, order, θ)`,
where
 phi - trial solution
 u - function
 x,y - coordinates of point
 εs - epsilon mask
 order - order of derivative
 θ - weight in neural network
"""
function _transform_expression(ex,dict_indvars,dict_depvars,chain,eltypeθ,strategy)
    _args = ex.args
    for (i,e) in enumerate(_args)
        if !(e isa Expr)
            if e in keys(dict_depvars)
                depvar = _args[1]
                num_depvar = dict_depvars[depvar]
                indvars = _args[2:end]
                cord = :cord
                ex.args = if !(typeof(chain) <: AbstractVector)
                    [:u, cord, :($θ), :phi]
                else
                    [:u, cord, Symbol(:($θ),num_depvar), Symbol(:phi,num_depvar)]
                end
                break
            elseif e isa ModelingToolkit.Differential
                derivative_variables = Symbol[]
                order = 0
                while (_args[1] isa ModelingToolkit.Differential)
                    order += 1
                    push!(derivative_variables, toexpr(_args[1].x))
                    _args = _args[2].args
                end
                depvar = _args[1]
                num_depvar = dict_depvars[depvar]
                indvars = _args[2:end]
                cord = :cord
                dim_l = length(indvars)
                εs = [get_ε(dim_l,d,eltypeθ) for d in 1:dim_l]
                undv = [dict_indvars[d_p] for d_p  in derivative_variables]
                εs_dnv = [εs[d] for d in undv]
                ex.args = if !(typeof(chain) <: AbstractVector)
                    [:derivative, :phi, :u, cord, εs_dnv, order, :($θ)]
                else
                    [:derivative, Symbol(:phi,num_depvar), :u, cord, εs_dnv, order, Symbol(:($θ),num_depvar)]
                end
                break
            end
        else
            ex.args[i] = _transform_expression(ex.args[i],dict_indvars,dict_depvars,chain,eltypeθ,strategy)
        end
    end
    return ex
end

"""
Parse ModelingToolkit equation form to the inner representation.

Example:

1)  1-D ODE: Dt(u(t)) ~ t +1

    Take expressions in the form: 'Equation(derivative(u(t), t), t + 1)' to 'derivative(phi, u_d, [t], [[ε]], 1, θ) - (t + 1)'

2)  2-D PDE: Dxx(u(x,y)) + Dyy(u(x,y)) ~ -sin(pi*x)*sin(pi*y)

    Take expressions in the form:
     Equation(derivative(derivative(u(x, y), x), x) + derivative(derivative(u(x, y), y), y), -(sin(πx)) * sin(πy))
    to
     (derivative(phi,u, [x, y], [[ε,0],[ε,0]], 2, θ) + derivative(phi, u, [x, y], [[0,ε],[0,ε]], 2, θ)) - -(sin(πx)) * sin(πy)

3)  System of PDEs: [Dx(u1(x,y)) + 4*Dy(u2(x,y)) ~ 0,
                    Dx(u2(x,y)) + 9*Dy(u1(x,y)) ~ 0]

    Take expressions in the form:
    2-element Array{Equation,1}:
        Equation(derivative(u1(x, y), x) + 4 * derivative(u2(x, y), y), ModelingToolkit.Constant(0))
        Equation(derivative(u2(x, y), x) + 9 * derivative(u1(x, y), y), ModelingToolkit.Constant(0))
    to
      [(derivative(phi1, u1, [x, y], [[ε,0]], 1, θ1) + 4 * derivative(phi2, u, [x, y], [[0,ε]], 1, θ2)) - 0,
       (derivative(phi2, u2, [x, y], [[ε,0]], 1, θ2) + 9 * derivative(phi1, u, [x, y], [[0,ε]], 1, θ1)) - 0]
"""

function build_symbolic_equation(eq,_indvars,_depvars,chain,eltypeθ,strategy)
    depvars,indvars,dict_indvars,dict_depvars = get_vars(_indvars, _depvars)
    parse_equation(eq,dict_indvars,dict_depvars,chain,eltypeθ,strategy)
end

function parse_equation(eq,dict_indvars,dict_depvars,chain,eltypeθ,strategy)
    eq_lhs = isequal(expand_derivatives(eq.lhs), 0) ? eq.lhs : expand_derivatives(eq.lhs)
    eq_rhs = isequal(expand_derivatives(eq.rhs), 0) ? eq.rhs : expand_derivatives(eq.rhs)

    left_expr = transform_expression(toexpr(eq_lhs),dict_indvars,dict_depvars,chain,eltypeθ,strategy)
    right_expr = transform_expression(toexpr(eq_rhs),dict_indvars,dict_depvars,chain,eltypeθ,strategy)
    left_expr = Broadcast.__dot__(left_expr)
    right_expr = Broadcast.__dot__(right_expr)
    loss_func = :($left_expr .- $right_expr)
end

"""
Build a loss function for a PDE or a boundary condition

# Examples: System of PDEs:

Take expressions in the form:

[Dx(u1(x,y)) + 4*Dy(u2(x,y)) ~ 0,
 Dx(u2(x,y)) + 9*Dy(u1(x,y)) ~ 0]

to

:((cord, θ, phi, derivative, u)->begin
          #= ... =#
          #= ... =#
          begin
              (θ1, θ2) = (θ[1:33], θ"[34:66])
              (phi1, phi2) = (phi[1], phi[2])
              let (x, y) = (cord[1], cord[2])
                  [(+)(derivative(phi1, u, [x, y], [[ε, 0.0]], 1, θ1), (*)(4, derivative(phi2, u, [x, y], [[0.0, ε]], 1, θ2))) - 0,
                   (+)(derivative(phi2, u, [x, y], [[ε, 0.0]], 1, θ2), (*)(9, derivative(phi1, u, [x, y], [[0.0, ε]], 1, θ1))) - 0]
              end
          end
      end)
"""
function build_symbolic_loss_function(eqs,_indvars,_depvars,
                                      phi, derivative,chain,initθ,strategy;
                                      bc_indvars=nothing,
                                      eq_params = SciMLBase.NullParameters(),
                                      param_estim = false,
                                      default_p=nothing)
    # dictionaries: variable -> unique number
    depvars,indvars,dict_indvars,dict_depvars = get_vars(_indvars, _depvars)
    bc_indvars = bc_indvars==nothing ? indvars : bc_indvars
    return build_symbolic_loss_function(eqs,indvars,depvars,
                                        dict_indvars,dict_depvars,
                                        phi,derivative,chain,initθ,strategy;
                                        bc_indvars = bc_indvars,
                                        eq_params = eq_params,
                                        param_estim = param_estim,
                                        default_p=default_p)
end

function get_indvars_ex(bc_indvars)
    i_=1
    indvars_ex = map(bc_indvars) do u
        if u isa Symbol
             ex = :($:cord[[$i_],:])
             i_+=1
             ex
        else
           :(fill($u,size($:cord[[1],:])))
        end
    end
    indvars_ex
end

function build_symbolic_loss_function(eqs,indvars,depvars,
                                      dict_indvars,dict_depvars,
                                      phi,derivative,chain,initθ,strategy;
                                      eq_params = SciMLBase.NullParameters(),
                                      param_estim = param_estim,
                                      default_p=default_p,
                                      bc_indvars = indvars)
    if chain isa AbstractArray
        eltypeθ = eltype(initθ[1])
    else
        eltypeθ = eltype(initθ)
    end

    loss_function = parse_equation(eqs,dict_indvars,dict_depvars,chain,eltypeθ,strategy)
    vars = :(cord, $θ, phi, derivative,u,p)
    ex = Expr(:block)
    if typeof(chain) <: AbstractVector
        θ_nums = Symbol[]
        phi_nums = Symbol[]
        for v in depvars
            num = dict_depvars[v]
            push!(θ_nums,:($(Symbol(:($θ),num))))
            push!(phi_nums,:($(Symbol(:phi,num))))
        end

        expr_θ = Expr[]
        expr_phi = Expr[]

        acum =  [0;accumulate(+, length.(initθ))]
        sep = [acum[i]+1 : acum[i+1] for i in 1:length(acum)-1]

        for i in eachindex(depvars)
            push!(expr_θ, :($θ[$(sep[i])]))
            push!(expr_phi, :(phi[$i]))
        end

        vars_θ = Expr(:(=), build_expr(:tuple, θ_nums), build_expr(:tuple, expr_θ))
        push!(ex.args,  vars_θ)

        vars_phi = Expr(:(=), build_expr(:tuple, phi_nums), build_expr(:tuple, expr_phi))
        push!(ex.args,  vars_phi)
    end
    #Add an expression for parameter symbols
    if param_estim == true && eq_params != SciMLBase.NullParameters()
        param_len = length(eq_params)
        last_indx =  [0;accumulate(+, length.(initθ))][end]
        params_symbols = Symbol[]
        expr_params = Expr[]
        for (i , eq_param) in enumerate(eq_params)
            push!(expr_params, :($θ[$(i+last_indx:i+last_indx)]))
            push!(params_symbols, Symbol(:($eq_param)))
        end
        params_eq = Expr(:(=), build_expr(:tuple, params_symbols), build_expr(:tuple, expr_params))
        push!(ex.args,  params_eq)
    end

    if eq_params != SciMLBase.NullParameters() && param_estim == false
        params_symbols = Symbol[]
        expr_params = Expr[]
        for (i , eq_param) in enumerate(eq_params)
            push!(expr_params, :(ArrayInterface.allowed_getindex(p,$i:$i)))
            push!(params_symbols, Symbol(:($eq_param)))
        end
        params_eq = Expr(:(=), build_expr(:tuple, params_symbols), build_expr(:tuple, expr_params))
        push!(ex.args,  params_eq)
    end

    if strategy isa QuadratureTraining

        indvars_ex = get_indvars_ex(bc_indvars)

        left_arg_pairs, right_arg_pairs = indvars,indvars_ex
        vcat_expr =  :(cord = vcat($(indvars...)))
        vcat_expr_loss_functions = Expr(:block,vcat_expr,loss_function) #TODO rename
        vars_eq = Expr(:(=), build_expr(:tuple, left_arg_pairs), build_expr(:tuple, right_arg_pairs))

    else
        indvars_ex = [:($:cord[[$i],:]) for (i, u) ∈ enumerate(indvars)]
        left_arg_pairs, right_arg_pairs = indvars,indvars_ex
        vars_eq = Expr(:(=), build_expr(:tuple, left_arg_pairs), build_expr(:tuple, right_arg_pairs))
        vcat_expr_loss_functions = loss_function #TODO rename
    end

    let_ex = Expr(:let, vars_eq, vcat_expr_loss_functions)
    push!(ex.args,  let_ex)

    expr_loss_function = :(($vars) -> begin $ex end)
end

function build_loss_function(eqs,_indvars,_depvars,phi,derivative,
                             chain,initθ,strategy;
                             bc_indvars=nothing,
                             eq_params=SciMLBase.NullParameters(),
                             param_estim=false,
                             default_p=nothing)
    # dictionaries: variable -> unique number
    depvars,indvars,dict_indvars,dict_depvars = get_vars(_indvars, _depvars)
    bc_indvars = bc_indvars==nothing ? indvars : bc_indvars
    return build_loss_function(eqs,indvars,depvars,
                               dict_indvars,dict_depvars,
                               phi,derivative,chain,initθ,strategy;
                               bc_indvars = bc_indvars,
                               eq_params=eq_params,
                               param_estim=param_estim,
                               default_p=default_p)
end

function build_loss_function(eqs,indvars,depvars,
                             dict_indvars,dict_depvars,
                             phi,derivative,chain,initθ,strategy;
                             bc_indvars = indvars,
                             eq_params=SciMLBase.NullParameters(),
                             param_estim=false,
                             default_p=nothing)
     expr_loss_function = build_symbolic_loss_function(eqs,indvars,depvars,
                                                       dict_indvars,dict_depvars,
                                                       phi,derivative,chain,initθ,strategy;
                                                       bc_indvars = bc_indvars,eq_params = eq_params,
                                                       param_estim=param_estim,default_p=default_p)
    u = get_u()
    _loss_function = @RuntimeGeneratedFunction(expr_loss_function)
    loss_function = (cord, θ) -> begin
        _loss_function(cord, θ, phi, derivative, u, default_p)
    end
    return loss_function
end

function get_vars(indvars_, depvars_)
    depvars = [nameof(value(d)) for d in depvars_]
    indvars = [nameof(value(i)) for i in indvars_]
    dict_indvars = get_dict_vars(indvars)
    dict_depvars = get_dict_vars(depvars)
    return depvars,indvars,dict_indvars,dict_depvars
end

function get_variables(eqs,_indvars::Array,_depvars::Array)
    depvars,indvars,dict_indvars,dict_depvars = get_vars(_indvars, _depvars)
    return get_variables(eqs,dict_indvars,dict_depvars)
end

function get_variables(eqs,dict_indvars,dict_depvars)
    bc_args = get_argument(eqs,dict_indvars,dict_depvars)
    return map(barg -> filter(x -> x isa Symbol, barg), bc_args)
end

function get_number(eqs,dict_indvars,dict_depvars)
    bc_args = get_argument(eqs,dict_indvars,dict_depvars)
    return map(barg -> filter(x -> x isa Number, barg), bc_args)
end

function find_thing_in_expr(ex::Expr, thing; ans = Expr[])
    for e in ex.args
        if e isa Expr
            if thing in e.args
                push!(ans,e)
            end
            find_thing_in_expr(e,thing; ans=ans)
        end
    end
    return collect(Set(ans))
end

# Get arguments from boundary condition functions
function get_argument(eqs,_indvars::Array,_depvars::Array)
    depvars,indvars,dict_indvars,dict_depvars = get_vars(_indvars, _depvars)
    get_argument(eqs,dict_indvars,dict_depvars)
end
function get_argument(eqs,dict_indvars,dict_depvars)
    exprs = toexpr.(eqs)
    vars = map(exprs) do expr
        _vars =  map(depvar -> find_thing_in_expr(expr,  depvar), collect(keys(dict_depvars)))
        f_vars = filter(x -> !isempty(x), _vars)
        map(x -> first(x), f_vars)
    end
    args_ = map(vars) do _vars
        map(var -> var.args[2:end] , _vars)
    end
    return first.(args_) #TODO for all arguments
end

function generate_training_sets(domains,dx,eqs,bcs,eltypeθ,_indvars::Array,_depvars::Array)
    depvars,indvars,dict_indvars,dict_depvars = get_vars(_indvars, _depvars)
    return generate_training_sets(domains,dx,eqs,bcs,eltypeθ,dict_indvars,dict_depvars)
end
# Generate training set in the domain and on the boundary
function generate_training_sets(domains,dx,eqs,bcs,eltypeθ,dict_indvars::Dict,dict_depvars::Dict)
    if dx isa Array
        dxs = dx
    else
        dxs = fill(dx,length(domains))
    end

    spans = [infimum(d.domain):dx:supremum(d.domain) for (d,dx) in zip(domains,dxs)]
    dict_var_span = Dict([Symbol(d.variables) => infimum(d.domain):dx:supremum(d.domain) for (d,dx) in zip(domains,dxs)])

    bound_args = get_argument(bcs,dict_indvars,dict_depvars)
    bound_vars = get_variables(bcs,dict_indvars,dict_depvars)

    dif = [eltypeθ[] for i=1:size(domains)[1]]
    for _args in bound_args
        for (i,x) in enumerate(_args)
            if x isa Number
                push!(dif[i],x)
            end
        end
    end
    cord_train_set = collect.(spans)
    bc_data = map(zip(dif,cord_train_set)) do (d,c)
        setdiff(c, d)
    end

    dict_var_span_ = Dict([Symbol(d.variables) => bc for (d,bc) in zip(domains,bc_data)])

    bcs_train_sets = map(bound_args) do bt
        span = map(b -> get(dict_var_span, b, b), bt)
        _set = adapt(eltypeθ,hcat(vec(map(points -> collect(points), Iterators.product(span...)))...))
    end

    pde_vars = get_variables(eqs,dict_indvars,dict_depvars)
    pde_args = get_argument(eqs,dict_indvars,dict_depvars)

    pde_train_set = adapt(eltypeθ, hcat(vec(map(points -> collect(points), Iterators.product(bc_data...)))...))

    pde_train_sets = map(pde_args) do bt
        span = map(b -> get(dict_var_span_, b, b), bt)
        _set = adapt(eltypeθ,hcat(vec(map(points -> collect(points), Iterators.product(span...)))...))
    end
    [pde_train_sets,bcs_train_sets]
end

function get_bounds(domains,eqs,bcs,eltypeθ,_indvars::Array,_depvars::Array,strategy)
    depvars,indvars,dict_indvars,dict_depvars = get_vars(_indvars, _depvars)
    return get_bounds(domains,eqs,bcs,eltypeθ,dict_indvars,dict_depvars,strategy)
end

function get_bounds(domains,eqs,bcs,eltypeθ,_indvars::Array,_depvars::Array,strategy::QuadratureTraining)
    depvars,indvars,dict_indvars,dict_depvars = get_vars(_indvars, _depvars)
    return get_bounds(domains,eqs,bcs,eltypeθ,dict_indvars,dict_depvars,strategy)
end

function get_bounds(domains,eqs,bcs,eltypeθ,dict_indvars,dict_depvars,strategy::QuadratureTraining)
    dict_lower_bound = Dict([Symbol(d.variables) => infimum(d.domain) for d in domains])
    dict_upper_bound = Dict([Symbol(d.variables) => supremum(d.domain) for d in domains])

    pde_args = get_argument(eqs,dict_indvars,dict_depvars)

    pde_lower_bounds= map(pde_args) do pd
        span = map(p -> get(dict_lower_bound, p, p), pd)
        map(s -> adapt(eltypeθ,s) + cbrt(eps(eltypeθ)), span)
    end
    pde_upper_bounds= map(pde_args) do pd
        span = map(p -> get(dict_upper_bound, p, p), pd)
        map(s -> adapt(eltypeθ,s) - cbrt(eps(eltypeθ)), span)
    end
    pde_bounds= [pde_lower_bounds,pde_upper_bounds]

    bound_vars = get_variables(bcs,dict_indvars,dict_depvars)

    bcs_lower_bounds = map(bound_vars) do bt
        map(b -> dict_lower_bound[b], bt)
    end
    bcs_upper_bounds = map(bound_vars) do bt
        map(b -> dict_upper_bound[b], bt)
    end
    bcs_bounds= [bcs_lower_bounds,bcs_upper_bounds]

    [pde_bounds, bcs_bounds]
end

function get_bounds(domains,eqs,bcs,eltypeθ,dict_indvars,dict_depvars,strategy)
    dx = 1 / strategy.points
    dict_span = Dict([Symbol(d.variables) => [infimum(d.domain)+dx, supremum(d.domain)-dx] for d in domains])
    # pde_bounds = [[infimum(d.domain),supremum(d.domain)] for d in domains]
    pde_args = get_argument(eqs,dict_indvars,dict_depvars)

    pde_bounds= map(pde_args) do pd
        span = map(p -> get(dict_span, p, p), pd)
        map(s -> adapt(eltypeθ,s), span)
    end

    bound_args = get_argument(bcs,dict_indvars,dict_depvars)
    dict_span = Dict([Symbol(d.variables) => [infimum(d.domain), supremum(d.domain)] for d in domains])

    bcs_bounds= map(bound_args) do bt
        span = map(b -> get(dict_span, b, b), bt)
        map(s -> adapt(eltypeθ,s), span)
    end
    [pde_bounds,bcs_bounds]
end

function get_phi(chain,parameterless_type_θ)
    # The phi trial solution
    if chain isa FastChain
        phi = (x,θ) -> chain(adapt(parameterless_type_θ,x),θ)
    else
        _,re  = Flux.destructure(chain)
        phi = (x,θ) -> re(θ)(adapt(parameterless_type_θ,x))
    end
    phi
end

function get_u()
    u = (cord, θ, phi)-> phi(cord, θ)
end

Base.Broadcast.broadcasted(::typeof(get_u()), cord, θ, phi) = get_u()(cord, θ, phi)

# the method to calculate the derivative
function get_numeric_derivative()
    derivative =
        (phi,u,x,εs,order,θ) ->
        begin
            _epsilon = one(eltype(θ)) / (2*cbrt(eps(eltype(θ))))
            ε = εs[order]
            ε = adapt(DiffEqBase.parameterless_type(θ),ε)
            x = adapt(DiffEqBase.parameterless_type(θ),x)
            if order > 1
                return (derivative(phi,u,x .+ ε,εs,order-1,θ)
                      .- derivative(phi,u,x .- ε,εs,order-1,θ)) .* _epsilon
            else
                return (u(x .+ ε,θ,phi) .- u(x .- ε,θ,phi)) .* _epsilon
            end
        end
end

Base.Broadcast.broadcasted(::typeof(get_numeric_derivative()), phi,u,x,εs,order,θ) = get_numeric_derivative()(phi,u,x,εs,order,θ)

function get_loss_function(loss_function, train_set, eltypeθ,parameterless_type_θ, strategy::GridTraining;τ=nothing)
    loss = (θ) -> mean(abs2,loss_function(train_set, θ))
end

@nograd function generate_random_points(points, bound, eltypeθ)
    function f(b)
      if b isa Number
           fill(eltypeθ(b),(1,points))
       else
           lb, ub =  b[1], b[2]
           lb .+ (ub .- lb) .* rand(eltypeθ,1,points)
       end
    end
    vcat(f.(bound)...)
end

function get_loss_function(loss_function, bound, eltypeθ, parameterless_type_θ, strategy::StochasticTraining;τ=nothing)
    points = strategy.points

    loss = (θ) -> begin
        sets = generate_random_points(points, bound,eltypeθ)
        sets_ = adapt(parameterless_type_θ,sets)
        mean(abs2,loss_function(sets_, θ))
    end
    return loss
end

@nograd function generate_quasi_random_points(points, bound, eltypeθ, sampling_alg)
    function f(b)
      if b isa Number
           fill(eltypeθ(b),(1,points))
       else
           lb, ub =  eltypeθ[b[1]], [b[2]]
           QuasiMonteCarlo.sample(points,lb,ub,sampling_alg)
       end
    end
    vcat(f.(bound)...)
end

function generate_quasi_random_points_batch(points, bound, eltypeθ, sampling_alg,minibatch)
    map(bound) do b
        if !(b isa Number)
            lb, ub =  [b[1]], [b[2]]
            set_ = QuasiMonteCarlo.generate_design_matrices(points,lb,ub,sampling_alg,minibatch)
            set = map(s -> adapt(eltypeθ,s), set_)
        else
            set = fill(eltypeθ(b),(1,points))
        end
    end
end

function get_loss_function(loss_function, bound, eltypeθ,parameterless_type_θ,strategy::QuasiRandomTraining;τ=nothing)
    sampling_alg = strategy.sampling_alg
    points = strategy.points
    resampling = strategy.resampling
    minibatch = strategy.minibatch

    point_batch = nothing
    point_batch = if resampling == false
        generate_quasi_random_points_batch(points, bound,eltypeθ,sampling_alg,minibatch)
    end
    loss =
        if resampling == true
            θ -> begin
                sets = generate_quasi_random_points(points, bound, eltypeθ, sampling_alg)
                sets_ = adapt(parameterless_type_θ,sets)
                mean(abs2,loss_function(sets_, θ))
            end
        else
            θ -> begin
                sets =  [point_batch[i] isa Array{eltypeθ,2} ?
                         point_batch[i] : point_batch[i][rand(1:minibatch)]
                                            for i in 1:length(point_batch)] #TODO
                sets_ = vcat(sets...)
                sets__ = adapt(parameterless_type_θ,sets_)
                mean(abs2,loss_function(sets__, θ))
            end
        end
    return loss
end

function get_loss_function(loss_function, lb,ub ,eltypeθ, parameterless_type_θ,strategy::QuadratureTraining;τ=nothing)

    if length(lb) == 0
        loss = (θ) -> mean(abs2,loss_function(rand(eltypeθ,1,10), θ))
        return loss
    end
    area = eltypeθ(prod(abs.(ub .-lb)))
    f_ = (lb,ub,loss_,θ) -> begin
        # last_x = 1
        function _loss(x,θ)
            # last_x = x
            # mean(abs2,loss_(x,θ), dims=2)
            # size_x = fill(size(x)[2],(1,1))
            x = adapt(parameterless_type_θ,x)
            sum(abs2,loss_(x,θ), dims=2) #./ size_x
        end
        prob = QuadratureProblem(_loss,lb,ub,θ,batch = strategy.batch,nout=1)
        solve(prob,
              strategy.quadrature_alg,
              reltol = strategy.reltol,
              abstol = strategy.abstol,
              maxiters = strategy.maxiters)[1]
    end
    loss = (θ) -> 1/area* f_(lb,ub,loss_function,θ)
    return loss
end

function SciMLBase.symbolic_discretize(pde_system::PDESystem, discretization::PhysicsInformedNN)
    eqs = pde_system.eqs
    bcs = pde_system.bcs

    domains = pde_system.domain
    eq_params = pde_system.ps
    defaults = pde_system.defaults
    default_p = eq_params == SciMLBase.NullParameters() ? nothing : [defaults[ep] for ep in eq_params]

    param_estim = discretization.param_estim
    additional_loss = discretization.additional_loss

    # dimensionality of equation
    dim = length(domains)
    depvars,indvars,dict_indvars,dict_depvars = get_vars(pde_system.indvars, pde_system.depvars)

    chain = discretization.chain
    initθ = discretization.init_params
    flat_initθ = if (typeof(chain) <: AbstractVector) reduce(vcat,initθ) else initθ end
    eltypeθ = eltype(flat_initθ)
    parameterless_type_θ =  DiffEqBase.parameterless_type(flat_initθ)
    phi = discretization.phi
    derivative = discretization.derivative
    strategy = discretization.strategy
    if !(eqs isa Array)
        eqs = [eqs]
    end
    symbolic_pde_loss_functions = [build_symbolic_loss_function(eq,indvars,depvars,
                                                                dict_indvars,dict_depvars,
                                                                phi, derivative,chain,initθ,strategy,
                                                                eq_params=eq_params,
                                                                param_estim=param_estim,
                                                                default_p=default_p)
                                                                for eq in eqs]

    bc_indvars = if strategy isa QuadratureTraining
         get_argument(bcs,dict_indvars,dict_depvars)
    else
         get_variables(bcs,dict_indvars,dict_depvars)
    end
    symbolic_bc_loss_functions = [build_symbolic_loss_function(bc,indvars,depvars,
                                                               dict_indvars,dict_depvars,
                                                               phi, derivative,chain,initθ,strategy;
                                                               eq_params=eq_params,
                                                               param_estim=param_estim,
                                                               default_p=default_p,
                                                               bc_indvars = bc_indvar)
                                                               for (bc,bc_indvar) in zip(bcs,bc_indvars)]
    symbolic_pde_loss_functions,symbolic_bc_loss_functions
end

# Convert a PDE problem into an OptimizationProblem
function SciMLBase.discretize(pde_system::PDESystem, discretization::PhysicsInformedNN)
    eqs = pde_system.eqs
    bcs = pde_system.bcs

    domains = pde_system.domain
    eq_params = pde_system.ps
    defaults = pde_system.defaults
    default_p = eq_params == SciMLBase.NullParameters() ? nothing : [defaults[ep] for ep in eq_params]

    param_estim = discretization.param_estim
    additional_loss = discretization.additional_loss

    # dimensionality of equation
    dim = length(domains)
    depvars,indvars,dict_indvars,dict_depvars = get_vars(pde_system.indvars,pde_system.depvars)

    chain = discretization.chain
    initθ = discretization.init_params
    flat_initθ = if (typeof(chain) <: AbstractVector) reduce(vcat,initθ) else  initθ end
    eltypeθ = eltype(flat_initθ)
    parameterless_type_θ =  DiffEqBase.parameterless_type(flat_initθ)

    flat_initθ = if param_estim == false flat_initθ else vcat(flat_initθ, adapt(typeof(flat_initθ),default_p)) end
    phi = discretization.phi
    derivative = discretization.derivative
    strategy = discretization.strategy
    if !(eqs isa Array)
        eqs = [eqs]
    end
    _pde_loss_functions = [build_loss_function(eq,indvars,depvars,
                                               dict_indvars,dict_depvars,
                                               phi, derivative,chain,initθ,strategy;
                                               eq_params=eq_params,
                                               param_estim=param_estim,
                                               default_p=default_p) for eq in eqs]
    bc_indvars = if strategy isa QuadratureTraining
         get_argument(bcs,dict_indvars,dict_depvars)
    else
         get_variables(bcs,dict_indvars,dict_depvars)
    end

    _bc_loss_functions = [build_loss_function(bc,indvars,depvars,
                                              dict_indvars,dict_depvars,
                                              phi,derivative,chain,initθ,strategy;
                                              eq_params=eq_params,
                                              param_estim=param_estim,
                                              default_p=default_p,
                                              bc_indvars = bc_indvar) for (bc,bc_indvar) in zip(bcs,bc_indvars)]

    pde_loss_functions, bc_loss_functions =
    if strategy isa GridTraining
        dx = strategy.dx

        train_sets = generate_training_sets(domains,dx,eqs,bcs,eltypeθ,
                                            dict_indvars,dict_depvars)

        # the points in the domain and on the boundary
        pde_train_sets, bcs_train_sets = train_sets

        pde_train_sets = adapt.(parameterless_type_θ,pde_train_sets)
        bcs_train_sets =  adapt.(parameterless_type_θ,bcs_train_sets)
        pde_loss_functions = [get_loss_function(_loss,_set,eltypeθ,parameterless_type_θ,strategy)
                                                for (_loss,_set) in zip(_pde_loss_functions,pde_train_sets)]

        bc_loss_functions =  [get_loss_function(_loss,_set,eltypeθ,parameterless_type_θ,strategy)
                                                for (_loss,_set) in zip(_bc_loss_functions, bcs_train_sets)]
        (pde_loss_functions, bc_loss_functions)
    elseif strategy isa StochasticTraining
          bounds = get_bounds(domains,eqs,bcs,eltypeθ,dict_indvars,dict_depvars,strategy)
          pde_bounds, bcs_bounds = bounds

          pde_loss_functions = [get_loss_function(_loss,bound,eltypeθ,parameterless_type_θ,strategy)
                                                  for (_loss,bound) in zip(_pde_loss_functions, pde_bounds)]

          strategy_ = StochasticTraining(strategy.bcs_points)
          bc_loss_functions = [get_loss_function(_loss,bound,eltypeθ,parameterless_type_θ,strategy_)
                                                 for (_loss,bound) in zip(_bc_loss_functions, bcs_bounds)]
          (pde_loss_functions, bc_loss_functions)
    elseif strategy isa QuasiRandomTraining
         bounds = get_bounds(domains,eqs,bcs,eltypeθ,dict_indvars,dict_depvars,strategy)
         pde_bounds, bcs_bounds = bounds

         pde_loss_functions = [get_loss_function(_loss,bound,eltypeθ,parameterless_type_θ,strategy)
                                                 for (_loss,bound) in zip(_pde_loss_functions, pde_bounds)]

         strategy_ = QuasiRandomTraining(strategy.bcs_points;
                                         sampling_alg = strategy.sampling_alg,
                                         resampling = strategy.resampling,
                                         minibatch = strategy.minibatch)
         bc_loss_functions = [get_loss_function(_loss,bound,eltypeθ,parameterless_type_θ,strategy_)
                                                for (_loss,bound) in zip(_bc_loss_functions, bcs_bounds)]
         (pde_loss_functions, bc_loss_functions)
    elseif strategy isa QuadratureTraining
        bounds = get_bounds(domains,eqs,bcs,eltypeθ,dict_indvars,dict_depvars,strategy)
        pde_bounds, bcs_bounds = bounds

        lbs,ubs = pde_bounds
        pde_loss_functions = [get_loss_function(_loss,lb,ub,eltypeθ,parameterless_type_θ,strategy)
                                                for (_loss,lb,ub) in zip(_pde_loss_functions, lbs,ubs )]
        lbs,ubs = bcs_bounds
        bc_loss_functions = [get_loss_function(_loss,lb,ub,eltypeθ,parameterless_type_θ,strategy)
                                               for (_loss,lb,ub) in zip(_bc_loss_functions, lbs,ubs)]

        (pde_loss_functions, bc_loss_functions)
    end

    pde_loss_function = θ -> sum(map(l->l(θ) ,pde_loss_functions))
    bcs_loss_function = θ -> sum(map(l->l(θ) ,bc_loss_functions))
    loss_function = θ -> pde_loss_function(θ) + bcs_loss_function(θ)

    function loss_function_(θ,p)
        if additional_loss isa Nothing
            return loss_function(θ)
        else
            function _additional_loss(phi,θ)
                θ_ = θ[1:end - length(default_p)]
                p = θ[(end - length(default_p) + 1):end]
                return additional_loss(phi,θ_,p)
            end
            return loss_function(θ) + _additional_loss(phi,θ)
        end
    end

    f = OptimizationFunction(loss_function_, GalacticOptim.AutoZygote())
    GalacticOptim.OptimizationProblem(f, flat_initθ)
end


text:
#################################################################################
# Comments
#
# - Ideally the net_injection variables would be bounded.  This can be done using an adhoc data model extention
#
#################################################################################
# Model Definitions

""
function build_nip_model(data::Dict{String, Any},
                         model_constructor;
                         multinetwork=true, kwargs...)
    return PM.build_model(data, model_constructor, post_nip; multinetwork=multinetwork, kwargs...)
end

""
function post_nip(pm::PM.AbstractPowerModel)
    for (n, network) in PM.nws(pm)
        @assert !PM.ismulticonductor(pm, nw=n)
        PM.variable_voltage(pm, nw=n)
        variable_net_injection(pm, nw=n)
        PM.variable_branch_flow(pm, nw=n, bounded=false)
        PM.variable_dcline_flow(pm, nw=n)

        PM.constraint_model_voltage(pm, nw=n)

        for i in PM.ids(pm, :ref_buses, nw=n)
            PM.constraint_theta_ref(pm, i, nw=n)
        end

        for i in PM.ids(pm, :bus, nw=n)
            constraint_power_balance_ni(pm, i, nw=n)
        end

        for i in PM.ids(pm, :branch, nw=n)
            PM.constraint_ohms_yt_from(pm, i, nw=n)
            PM.constraint_ohms_yt_to(pm, i, nw=n)

            PM.constraint_voltage_angle_difference(pm, i, nw=n)

            #PM.constraint_thermal_limit_from(pm, i, nw=n)
            #PM.constraint_thermal_limit_to(pm, i, nw=n)
        end

        for i in PM.ids(pm, :dcline)
            PM.constraint_dcline(pm, i, nw=n)
        end
    end

    return

end

""
function build_nip_expr_model(data::Dict{String, Any}, model_constructor; multinetwork=true, kwargs...)
    return PM.build_model(data, model_constructor, post_nip_expr; multinetwork=multinetwork, kwargs...)
end

""
function post_nip_expr(pm::PM.AbstractPowerModel)
    for (n, network) in PM.nws(pm)
        @assert !PM.ismulticonductor(pm, nw=n)
        PM.variable_voltage(pm, nw=n)
        PM.variable_branch_flow(pm, nw=n; bounded = false)
        PM.variable_dcline_flow(pm, nw=n)

        PM.constraint_model_voltage(pm, nw=n)

        for i in PM.ids(pm, :ref_buses, nw=n)
            PM.constraint_theta_ref(pm, i, nw=n)
        end

        for i in PM.ids(pm, :bus, nw=n)
            constraint_power_balance_ni_expr(pm, i, nw=n)
        end

        for i in PM.ids(pm, :branch, nw=n)
            PM.constraint_ohms_yt_from(pm, i, nw=n)
            PM.constraint_ohms_yt_to(pm, i, nw=n)

            PM.constraint_voltage_angle_difference(pm, i, nw=n)

            #PM.constraint_thermal_limit_from(pm, i, nw=n)
            #PM.constraint_thermal_limit_to(pm, i, nw=n)
        end

        for i in PM.ids(pm, :dcline)
            PM.constraint_dcline(pm, i, nw=n)
        end
    end

    return

end

#################################################################################
# Model Extention Functions

"generates variables for both `active` and `reactive` net injection"
function variable_net_injection(pm::PM.AbstractPowerModel; kwargs...)
    variable_active_net_injection(pm; kwargs...)
    variable_reactive_net_injection(pm; kwargs...)

    return

end

""
function variable_active_net_injection(pm::PM.AbstractPowerModel; nw::Int=pm.cnw, cnd::Int=pm.ccnd)
    PM.var(pm, nw, cnd)[:pni] = JuMP.@variable(pm.model,
        [i in PM.ids(pm, nw, :bus)], base_name="$(nw)_$(cnd)_pni",
        start = 0.0
    )

    return

end

""
function variable_reactive_net_injection(pm::PM.AbstractPowerModel; nw::Int=pm.cnw, cnd::Int=pm.ccnd)
    PM.var(pm, nw, cnd)[:qni] = JuMP.@variable(pm.model,
        [i in PM.ids(pm, nw, :bus)], base_name="$(nw)_$(cnd)_qni",
        start = 0.0
    )

    return
end

""
function constraint_power_balance_ni(pm::PM.AbstractPowerModel, i::Int; nw::Int=pm.cnw, cnd::Int=pm.ccnd)
    if !haskey(PM.con(pm, nw, cnd), :power_balance_p)
        PM.con(pm, nw, cnd)[:power_balance_p] = Dict{Int, JuMP.ConstraintRef}()
    end
    if !haskey(PM.con(pm, nw, cnd), :power_balance_q)
        PM.con(pm, nw, cnd)[:power_balance_q] = Dict{Int, JuMP.ConstraintRef}()
    end

    bus = PM.ref(pm, nw, :bus, i)
    bus_arcs = PM.ref(pm, nw, :bus_arcs, i)
    bus_arcs_dc = PM.ref(pm, nw, :bus_arcs_dc, i)

    constraint_power_balance_ni(pm, nw, cnd, i, bus_arcs, bus_arcs_dc)

    return

end

""
function constraint_power_balance_ni(pm::PM.AbstractPowerModel,
                           n::Int, c::Int, i::Int,
                           bus_arcs, bus_arcs_dc)
    p = PM.var(pm, n, c, :p)
    q = PM.var(pm, n, c, :q)
    pni = PM.var(pm, n, c, :pni, i)
    qni = PM.var(pm, n, c, :qni, i)
    p_dc = PM.var(pm, n, c, :p_dc)
    q_dc = PM.var(pm, n, c, :q_dc)

    PM.con(pm, n, c, :power_balance_p)[i] = JuMP.@constraint(pm.model, sum(p[a] for a in bus_arcs) + sum(p_dc[a_dc] for a_dc in bus_arcs_dc) == pni)
    PM.con(pm, n, c, :power_balance_q)[i] = JuMP.@constraint(pm.model, sum(q[a] for a in bus_arcs) + sum(q_dc[a_dc] for a_dc in bus_arcs_dc) == qni)

    return

end

""
function constraint_power_balance_ni_expr(pm::PM.AbstractPowerModel,
                                i::Int; nw::Int=pm.cnw, cnd::Int=pm.ccnd)
    if !haskey(PM.con(pm, nw, cnd), :power_balance_p)
        PM.con(pm, nw, cnd)[:power_balance_p] = Dict{Int, JuMP.ConstraintRef}()
    end
    if !haskey(PM.con(pm, nw, cnd), :power_balance_q)
        PM.con(pm, nw, cnd)[:power_balance_q] = Dict{Int, JuMP.ConstraintRef}()
    end

    bus = PM.ref(pm, nw, :bus, i)
    bus_arcs = PM.ref(pm, nw, :bus_arcs, i)
    bus_arcs_dc = PM.ref(pm, nw, :bus_arcs_dc, i)

    pni_expr = PM.ref(pm, nw, :bus, i, "pni")
    qni_expr = PM.ref(pm, nw, :bus, i, "qni")

    constraint_power_balance_ni_expr(pm, nw, cnd, i, bus_arcs, bus_arcs_dc, pni_expr, qni_expr)

    return

end

""
function constraint_power_balance_ni_expr(pm::PM.AbstractPowerModel,
                                n::Int, c::Int, i::Int,
                                bus_arcs, bus_arcs_dc, pni_expr, qni_expr)
    p = PM.var(pm, n, c, :p)
    q = PM.var(pm, n, c, :q)
    p_dc = PM.var(pm, n, c, :p_dc)
    q_dc = PM.var(pm, n, c, :q_dc)

    PM.con(pm, n, c, :power_balance_p)[i] = JuMP.@constraint(pm.model, sum(p[a] for a in bus_arcs) + sum(p_dc[a_dc] for a_dc in bus_arcs_dc) == pni_expr)
    PM.con(pm, n, c, :power_balance_q)[i] = JuMP.@constraint(pm.model, sum(q[a] for a in bus_arcs) + sum(q_dc[a_dc] for a_dc in bus_arcs_dc) == qni_expr)

    return

end

"active power only models ignore reactive power variables"
function variable_reactive_net_injection(pm::PM.AbstractActivePowerModel; kwargs...)
    return
end

"active power only models ignore reactive power flows"
function constraint_power_balance_ni(pm::PM.AbstractActivePowerModel,
                                    n::Int, c::Int, i::Int,
                                    bus_arcs, bus_arcs_dc)
    p = PM.var(pm, n, c, :p)
    pni = PM.var(pm, n, c, :pni, i)
    p_dc = PM.var(pm, n, c, :p_dc)

    PM.con(pm, n, c, :power_balance_p)[i] = JuMP.@constraint(pm.model, sum(p[a] for a in bus_arcs) + sum(p_dc[a_dc] for a_dc in bus_arcs_dc) == pni)

    return

end

""
function constraint_power_balance_ni_expr(pm::PM.AbstractActivePowerModel,
                                        n::Int, c::Int, i::Int,
                                        bus_arcs, bus_arcs_dc, pni_expr, qni_expr)
    p = PM.var(pm, n, c, :p)
    p_dc = PM.var(pm, n, c, :p_dc)

    PM.con(pm, n, c, :power_balance_p)[i] = JuMP.@constraint(pm.model, sum(p[a] for a in bus_arcs) + sum(p_dc[a_dc] for a_dc in bus_arcs_dc) == pni_expr)

    return

end

""
function powermodels_network!(canonical_model::CanonicalModel,
                              system_formulation::Type{S},
                              sys::PSY.System) where {S<:PM.AbstractPowerModel}

    time_steps = model_time_steps(canonical_model)
    pm_data, PM_map = pass_to_pm(sys, time_steps[end])
    buses = PSY.get_components(PSY.Bus, sys)

    _remove_undef!(canonical_model.expressions[:nodal_balance_active])
    _remove_undef!(canonical_model.expressions[:nodal_balance_reactive])

    for t in time_steps, bus in buses
        pm_data["nw"]["$(t)"]["bus"]["$(bus.number)"]["pni"] = canonical_model.expressions[:nodal_balance_active][bus.number, t]
        pm_data["nw"]["$(t)"]["bus"]["$(bus.number)"]["qni"] = canonical_model.expressions[:nodal_balance_reactive][bus.number, t]
    end

    canonical_model.pm_model = build_nip_expr_model(pm_data, system_formulation, jump_model=canonical_model.JuMPmodel);
    canonical_model.pm_model.ext[:PMmap] = PM_map

    return

end

""
function powermodels_network!(canonical_model::CanonicalModel,
                              system_formulation::Type{S},
                              sys::PSY.System) where {S<:PM.AbstractActivePowerModel}

    time_steps = model_time_steps(canonical_model)
    pm_data, PM_map = pass_to_pm(sys, time_steps[end])
    buses = PSY.get_components(PSY.Bus, sys)

    _remove_undef!(canonical_model.expressions[:nodal_balance_active])

    for t in time_steps, bus in buses
        pm_data["nw"]["$(t)"]["bus"]["$(PSY.get_number(bus))"]["pni"] = canonical_model.expressions[:nodal_balance_active][PSY.get_number(bus), t]
        #pm_data["nw"]["$(t)"]["bus"]["$(bus.number)"]["qni"] = 0.0
    end

    canonical_model.pm_model = build_nip_expr_model(pm_data, system_formulation, jump_model=canonical_model.JuMPmodel);
    canonical_model.pm_model.ext[:PMmap] = PM_map

    return

end

#### PM accessor functions ########

function PMvarmap(system_formulation::Type{S}) where {S<:PM.AbstractDCPModel}
    pm_var_map = Dict{Type,Dict{Symbol, Union{Symbol,NamedTuple}}}()

    pm_var_map[PSY.Bus] = Dict(:va => :theta)
    pm_var_map[PSY.ACBranch] = Dict(:p => (from_to = :Fp, to_from = nothing))
    pm_var_map[PSY.DCBranch] = Dict(:p_dc => (from_to = :Fp, to_from = nothing))

    return pm_var_map
end

function PMvarmap(system_formulation::Type{S}) where {S<:PM.AbstractActivePowerModel}
    pm_var_map = Dict{Type,Dict{Symbol, Union{Symbol,NamedTuple}}}()

    pm_var_map[PSY.Bus] = Dict(:va => :theta)
    pm_var_map[PSY.ACBranch] = Dict(:p => (from_to = :FpFT, to_from = :FpTF))
    pm_var_map[PSY.DCBranch] = Dict(:p_dc => (from_to = :FpFT, to_from = :FpTF))

    return pm_var_map
end

function PMvarmap(system_formulation::Type{S}) where {S<:PM.AbstractPowerModel}
    pm_var_map = Dict{Type,Dict{Symbol, Union{Symbol,NamedTuple}}}()

    pm_var_map[PSY.Bus] = Dict(:va => :theta,
                                :vm => :Vm)
    pm_var_map[PSY.ACBranch] = Dict(:p => (from_to = :FpFT, to_from = :FpTF),
                                    :q => (from_to = :FqFT, to_from = :FqTF))
    pm_var_map[PSY.DCBranch] = Dict(:p_dc => (from_to = :FpFT, to_from = :FpTF),
                                    :q_dc => (from_to = :FqFT, to_from = :FqTF))

    return pm_var_map
end

function add_pm_var_refs!(canonical_model::CanonicalModel,
                          system_formulation::Type{S},
                          sys::PSY.System) where {S<:PM.AbstractPowerModel}

    time_steps = model_time_steps(canonical_model)
    bus_dict = canonical_model.pm_model.ext[:PMmap].bus
    ACbranch_dict = canonical_model.pm_model.ext[:PMmap].arcs
    ACbranch_types = typeof.(values(ACbranch_dict))
    DCbranch_dict = canonical_model.pm_model.ext[:PMmap].arcs_dc
    DCbranch_types = typeof.(values(DCbranch_dict))

    pm_var_names = keys(canonical_model.pm_model.var[:nw][1][:cnd][1])

    pm_var_map = PMvarmap(system_formulation)

    for (pm_v, ps_v) in pm_var_map[PSY.Bus]
        if pm_v in pm_var_names
            canonical_model.variables[ps_v] = PSI._container_spec(canonical_model.JuMPmodel,
                                                        (PSY.get_name(b) for b in values(bus_dict)),
                                                        time_steps)
            for t in time_steps, (pm_bus, bus) in bus_dict
                name = PSY.get_name(bus)
                canonical_model.variables[ps_v][name, t] = PM.var(canonical_model.pm_model, t, 1, pm_v)[1] #pm_vars[pm_v][pm_bus]
            end
        end
    end

    add_pm_var_refs!(canonical_model, PSY.ACBranch, ACbranch_types, ACbranch_dict, pm_var_map, pm_var_names, time_steps)
    add_pm_var_refs!(canonical_model, PSY.DCBranch, DCbranch_types, DCbranch_dict, pm_var_map, pm_var_names, time_steps)

end

function add_pm_var_refs!(canonical_model::CanonicalModel,
                          d_class::Type,
                          device_types::Vector,
                          pm_map::Dict,
                          pm_var_map::Dict,
                          pm_var_names::Base.KeySet,
                          time_steps::UnitRange{Int64})

    for d_type in Set(device_types)
        devices = [d for d in pm_map if typeof(d[2]) == d_type]
        for (pm_v, ps_v) in pm_var_map[d_class]
            if pm_v in pm_var_names
                for dir in fieldnames(typeof(ps_v))
                    isnothing(getfield(ps_v,dir)) && continue
                    var_name = Symbol("$(getfield(ps_v,dir))_$(d_type)")
                    canonical_model.variables[var_name] = PSI._container_spec(canonical_model.JuMPmodel,
                                                                        (PSY.get_name(d[2]) for d in devices),
                                                                        time_steps)
                    for t in time_steps, (pm_d, d) in devices
                        canonical_model.variables[var_name][PSY.get_name(d), t] = PM.var(canonical_model.pm_model, t, 1, pm_v, getfield(pm_d, dir))
                    end
                end
            end
        end
    end
end

ext: jl


text:
using WeakValueDicts
using Test

mutable struct V
    v
end

struct ImmutableStruct
    v
end

@testset "WeakValueDict.jl" begin
    @testset "Basic Dict operations" begin
        foo = V("foo")
        bar = V("bar")
        wvd = WeakValueDict(
            "foo" => foo,
            "bar" => bar,
        )
        @test typeof(wvd) === WeakValueDict{String, V}
        @test length(wvd) == 2
        @test !isempty(wvd)

        @test copy(wvd) !== wvd
        @test copy(wvd) == wvd

        @test wvd["foo"] === foo
        @test get(wvd, "foo", nothing) === foo
        @test get!(wvd, "foo", V(nothing)) === foo

        @test pop!(wvd, "foo") === foo
        @test !haskey(wvd, "foo")
        @test_throws KeyError pop!(wvd, "foo")
        @test_throws KeyError wvd["foo"]
        @test delete!(wvd, "foo") === wvd

        @test get!(wvd, "foo", foo) === foo
        @test wvd["foo"] === foo
        @test delete!(wvd, "foo") === wvd
        @test !haskey(wvd, "foo")

        @test delete!(wvd, "bar") == empty(wvd)
        @test isempty(wvd)

        wvd["bar"] = V("bar")
        @test empty!(wvd) == empty(wvd)
        @test isempty(wvd)

        @test WeakValueDict{String, V}() == empty(WeakValueDict(), String, V)
    end

    @testset "Finalization" begin
        wvd = WeakValueDict{Int, V}()
        for i in 1:1000
            wvd[i] = V(i)
        end
        Base.GC.gc(true)
        @test length(wvd) < 1000
    end

    @testset "Constructors" begin
        wvd = WeakValueDict{String, String}()
        @test typeof(wvd) === WeakValueDict{String, String}

        wvd = WeakValueDict()
        @test typeof(wvd) === WeakValueDict{Any, Any}

        wvd = WeakValueDict(Tuple{}())
        @test typeof(wvd) === WeakValueDict{Any, Any}

        ps = Pair{String, V}["foo" => V(1), "bar" => V(2)]
        wvd = WeakValueDict(ps...)
        @test typeof(wvd) === WeakValueDict{String, V}

        ps = Pair{String}["foo" => "foo", "bar" => V(123)]
        wvd = WeakValueDict(ps...)
        @test typeof(wvd) === WeakValueDict{String, Any}

        wvd = WeakValueDict("foo" => "bar", 123 => "eggs")
        @test typeof(wvd) === WeakValueDict{Any, String}

        wvd = WeakValueDict("foo" => "bar", V(1) => V(2))
        @test typeof(wvd) === WeakValueDict{Any, Any}

        wvd = WeakValueDict(Dict("foo" => "bar", "spam" => "eggs"))
        @test typeof(wvd) === WeakValueDict{String, String}

        wvd = WeakValueDict{String, String}(Dict("foo" => "bar"))
        @test typeof(wvd) === WeakValueDict{String, String}

        wvd = WeakValueDict{String, String}("foo" => "bar")
        @test typeof(wvd) === WeakValueDict{String, String}

        wvd = WeakValueDict{String, String}("foo" => "bar", "spam" => "eggs")
        @test typeof(wvd) === WeakValueDict{String, String}

        wvd = WeakValueDict(Dict("foo" => "bar"))
        @test typeof(wvd) === WeakValueDict{String, String}

        bar = V("bar")
        wvd = WeakValueDict([("foo", bar), ("spam", bar)])
        @test typeof(wvd) === WeakValueDict{String, V}
        @test wvd["foo"] === bar

        @test_throws ArgumentError WeakValueDict("foo")
        @test_throws ArgumentError WeakValueDict(123)
        @test_throws MethodError WeakValueDict("foo", "bar", "spam", "eggs")
    end

    @testset "map! and filter!" begin
        foo = V("foo")
        bar = V("bar")
        wvd = WeakValueDict(
            "foo" => foo,
            "bar" => bar,
        )

        map!(values(wvd)) do value
            return V(uppercase(value.v))
        end
        @test wvd["foo"].v == "FOO"
        @test wvd["bar"].v == "BAR"

        filter!(wvd) do p
            return p.first == "foo"
        end
        @test wvd["foo"].v == "FOO"
        @test !haskey(wvd, "bar")
        @test length(wvd) == 1
    end

    @testset "iterate" begin
        foo = V("foo")
        bar = V("bar")
        wvd = WeakValueDict(
            "foo" => foo,
            "bar" => bar,
        )

        for (k, v) in wvd
            @test v.v == k
        end
    end

    @testset "finalizer deferal" begin
        wvd = WeakValueDict()
        for i in 1:1000
            wvd[i] = V(i)
        end
        lock(wvd) do
            @test islocked(wvd)
            @test islocked(wvd.lock)

            l = length(wvd)
            @test l > 0

            # This GC should invoke the finalizer for some of the values in the
            # WeakValueDict, but since its locked, it will "defer" the finalizer
            # until later (at which point the WVD should be unlocked).
            Base.GC.gc(true)

            # We shouldn't delete any keys while locked because none of the
            # finalizers should have executed in the "normal" (unlocked) path.
            @test length(wvd) == l
        end

        # This should invoke the finalizers for all the entries in the dict.
        # The finalizers *should* be the finalizers that were scheduled later
        # and the GC'd values should have their entries in the dict cleared.
        Base.GC.gc(true)
        @test length(wvd) < 1000
    end
end


text:
# This file is based on a part of Julia LinearAlgebra/test/bidiag.jl. License is MIT: https://julialang.org/license
using Test, LazyBandedMatrices, SparseArrays, Random, FillArrays
import LinearAlgebra
import LazyBandedMatrices: Bidiagonal, SymTridiagonal, Tridiagonal
import LinearAlgebra: mul!

@testset "Bidiagonal" begin
    n = 10 #Size of test matrix
    Random.seed!(1)

    @testset for relty in (Int, Float32, Float64, BigFloat), elty in (relty, Complex{relty})
        if relty <: AbstractFloat
            dv = convert(Vector{elty}, randn(n))
            ev = convert(Vector{elty}, randn(n-1))
            if (elty <: Complex)
                dv += im*convert(Vector{elty}, randn(n))
                ev += im*convert(Vector{elty}, randn(n-1))
            end
        elseif relty <: Integer
            dv = convert(Vector{elty}, rand(1:10, n))
            ev = convert(Vector{elty}, rand(1:10, n-1))
            if (elty <: Complex)
                dv += im*convert(Vector{elty}, rand(1:10, n))
                ev += im*convert(Vector{elty}, rand(1:10, n-1))
            end
        end
        dv0 = zeros(elty, 0)
        ev0 = zeros(elty, 0)

        @testset "Constructors" begin
            for (x, y) in ((dv0, ev0), (dv, ev), (GenericArray(dv), GenericArray(ev)))
                # from vectors
                ubd = Bidiagonal(x, y, :U)
                lbd = Bidiagonal(x, y, :L)
                @test ubd != lbd || x === dv0
                @test ubd.dv === x
                @test lbd.ev === y
                @test_throws ArgumentError Bidiagonal(x, y, :R)
                x == dv0 || @test_throws DimensionMismatch Bidiagonal(x, x, :U)
                @test_throws MethodError Bidiagonal(x, y)
                # from matrix
                @test Bidiagonal(ubd, :U) == Bidiagonal(Matrix(ubd), :U) == ubd
                @test Bidiagonal(lbd, :L) == Bidiagonal(Matrix(lbd), :L) == lbd
            end
            @test eltype(Bidiagonal{elty}([1,2,3,4], [1.0f0,2.0f0,3.0f0], :U)) == elty
            @test eltype(Bidiagonal([1,2,3,4], [1.0f0,2.0f0,3.0f0], :U)) == Float32 # promotion test
            @test isa(Bidiagonal{elty,Vector{elty},Vector{elty}}(GenericArray(dv), ev, :U), Bidiagonal{elty,Vector{elty},Vector{elty}})
            @test Bidiagonal(dv, GenericArray(ev), :U) == Bidiagonal(dv, ev, :U)
            @test Bidiagonal(GenericArray(dv), ev, :U) == Bidiagonal(dv, ev, :U)
            BI = Bidiagonal([1,2,3,4], [1,2,3], :U)
            @test Bidiagonal(BI) === BI
            @test isa(Bidiagonal{elty}(BI), Bidiagonal{elty})
        end

        @testset "getindex, setindex!, size, and similar" begin
            ubd = Bidiagonal(dv, ev, :U)
            lbd = Bidiagonal(dv, ev, :L)
            # bidiagonal getindex / upper & lower
            @test_throws BoundsError ubd[n + 1, 1]
            @test_throws BoundsError ubd[1, n + 1]
            @test ubd[2, 2] == dv[2]
            # bidiagonal getindex / upper
            @test ubd[2, 3] == ev[2]
            @test iszero(ubd[3, 2])
            # bidiagonal getindex / lower
            @test lbd[3, 2] == ev[2]
            @test iszero(lbd[2, 3])
            # bidiagonal setindex! / upper
            cubd = copy(ubd)
            @test_throws ArgumentError ubd[2, 1] = 1
            @test_throws ArgumentError ubd[3, 1] = 1
            @test (cubd[2, 1] = 0; cubd == ubd)
            @test ((cubd[1, 2] = 10) == 10; cubd[1, 2] == 10)
            # bidiagonal setindex! / lower
            clbd = copy(lbd)
            @test_throws ArgumentError lbd[1, 2] = 1
            @test_throws ArgumentError lbd[1, 3] = 1
            @test (clbd[1, 2] = 0; clbd == lbd)
            @test ((clbd[2, 1] = 10) == 10; clbd[2, 1] == 10)
            # bidiagonal setindex! / upper & lower
            @test_throws BoundsError ubd[n + 1, 1] = 1
            @test_throws BoundsError ubd[1, n + 1] = 1
            @test ((cubd[2, 2] = 10) == 10; cubd[2, 2] == 10)
            # bidiagonal size
            @test_throws ArgumentError size(ubd, 0)
            @test size(ubd, 1) == size(ubd, 2) == n
            @test size(ubd, 3) == 1
            # bidiagonal similar
            @test isa(similar(ubd), Bidiagonal{elty})
            @test similar(ubd).uplo == ubd.uplo
            @test isa(similar(ubd, Int), Bidiagonal{Int})
            @test similar(ubd, Int).uplo == ubd.uplo
            # @test isa(similar(ubd, (3, 2)), SparseMatrixCSC)
            # @test isa(similar(ubd, Int, (3, 2)), SparseMatrixCSC{Int})

            # setindex! when off diagonal is zero bug
            Bu = Bidiagonal(rand(elty, 10), zeros(elty, 9), 'U')
            Bl = Bidiagonal(rand(elty, 10), zeros(elty, 9), 'L')
            @test_throws ArgumentError Bu[5, 4] = 1
            @test_throws ArgumentError Bl[4, 5] = 1
        end

        # @testset "show" begin
        #     BD = Bidiagonal(dv, ev, :U)
        #     dstring = sprint(Base.print_matrix,BD.dv')
        #     estring = sprint(Base.print_matrix,BD.ev')
        #     @test sprint(show,BD) == "$(summary(BD)):\n diag:$dstring\n super:$estring"
        #     BD = Bidiagonal(dv,ev,:L)
        #     @test sprint(show,BD) == "$(summary(BD)):\n diag:$dstring\n sub:$estring"
        # end

        @testset for uplo in (:U, :L)
            T = Bidiagonal(dv, ev, uplo)
            Tfull = Array(T)
            @testset "Constructor and basic properties" begin
                @test size(T, 1) == size(T, 2) == n
                @test size(T) == (n, n)
                @test Array(T) == diagm(0 => dv, (uplo == :U ? 1 : -1) => ev)
                @test Bidiagonal(Array(T), uplo) == T
                @test big.(T) == T
                @test Array(abs.(T)) == abs.(diagm(0 => dv, (uplo == :U ? 1 : -1) => ev))
                @test Array(real(T)) == real(diagm(0 => dv, (uplo == :U ? 1 : -1) => ev))
                @test Array(imag(T)) == imag(diagm(0 => dv, (uplo == :U ? 1 : -1) => ev))
            end

            @testset for func in (conj, transpose, adjoint)
                @test func(func(T)) == T
            end

            @testset "triu and tril" begin
                zerosdv = zeros(elty, length(dv))
                zerosev = zeros(elty, length(ev))
                bidiagcopy(dv, ev, uplo) = Bidiagonal(copy(dv), copy(ev), uplo)

                @test istril(Bidiagonal(dv,ev,:L))
                @test istril(Bidiagonal(dv,ev,:L), 1)
                @test !istril(Bidiagonal(dv,ev,:L), -1)
                @test istril(Bidiagonal(zerosdv,ev,:L), -1)
                @test !istril(Bidiagonal(zerosdv,ev,:L), -2)
                @test istril(Bidiagonal(zerosdv,zerosev,:L), -2)
                @test !istril(Bidiagonal(dv,ev,:U))
                @test istril(Bidiagonal(dv,ev,:U), 1)
                @test !istril(Bidiagonal(dv,ev,:U), -1)
                @test !istril(Bidiagonal(zerosdv,ev,:U), -1)
                @test istril(Bidiagonal(zerosdv,zerosev,:U), -1)
                @test tril!(bidiagcopy(dv,ev,:U),-1) == Bidiagonal(zerosdv,zerosev,:U)
                @test tril!(bidiagcopy(dv,ev,:L),-1) == Bidiagonal(zerosdv,ev,:L)
                @test tril!(bidiagcopy(dv,ev,:U),-2) == Bidiagonal(zerosdv,zerosev,:U)
                @test tril!(bidiagcopy(dv,ev,:L),-2) == Bidiagonal(zerosdv,zerosev,:L)
                @test tril!(bidiagcopy(dv,ev,:U),1)  == Bidiagonal(dv,ev,:U)
                @test tril!(bidiagcopy(dv,ev,:L),1)  == Bidiagonal(dv,ev,:L)
                @test tril!(bidiagcopy(dv,ev,:U))    == Bidiagonal(dv,zerosev,:U)
                @test tril!(bidiagcopy(dv,ev,:L))    == Bidiagonal(dv,ev,:L)
                @test_throws ArgumentError tril!(bidiagcopy(dv, ev, :U), -n - 2)
                @test_throws ArgumentError tril!(bidiagcopy(dv, ev, :U), n)

                @test istriu(Bidiagonal(dv,ev,:U))
                @test istriu(Bidiagonal(dv,ev,:U), -1)
                @test !istriu(Bidiagonal(dv,ev,:U), 1)
                @test istriu(Bidiagonal(zerosdv,ev,:U), 1)
                @test !istriu(Bidiagonal(zerosdv,ev,:U), 2)
                @test istriu(Bidiagonal(zerosdv,zerosev,:U), 2)
                @test !istriu(Bidiagonal(dv,ev,:L))
                @test istriu(Bidiagonal(dv,ev,:L), -1)
                @test !istriu(Bidiagonal(dv,ev,:L), 1)
                @test !istriu(Bidiagonal(zerosdv,ev,:L), 1)
                @test istriu(Bidiagonal(zerosdv,zerosev,:L), 1)
                @test triu!(bidiagcopy(dv,ev,:L),1)  == Bidiagonal(zerosdv,zerosev,:L)
                @test triu!(bidiagcopy(dv,ev,:U),1)  == Bidiagonal(zerosdv,ev,:U)
                @test triu!(bidiagcopy(dv,ev,:U),2)  == Bidiagonal(zerosdv,zerosev,:U)
                @test triu!(bidiagcopy(dv,ev,:L),2)  == Bidiagonal(zerosdv,zerosev,:L)
                @test triu!(bidiagcopy(dv,ev,:U),-1) == Bidiagonal(dv,ev,:U)
                @test triu!(bidiagcopy(dv,ev,:L),-1) == Bidiagonal(dv,ev,:L)
                @test triu!(bidiagcopy(dv,ev,:L))    == Bidiagonal(dv,zerosev,:L)
                @test triu!(bidiagcopy(dv,ev,:U))    == Bidiagonal(dv,ev,:U)
                @test_throws ArgumentError triu!(bidiagcopy(dv, ev, :U), -n)
                @test_throws ArgumentError triu!(bidiagcopy(dv, ev, :U), n + 2)
                @test !isdiag(Bidiagonal(dv,ev,:U))
                @test !isdiag(Bidiagonal(dv,ev,:L))
                @test isdiag(Bidiagonal(dv,zerosev,:U))
                @test isdiag(Bidiagonal(dv,zerosev,:L))
            end

            @testset "iszero and isone" begin
                for uplo in (:U, :L)
                    BDzero = Bidiagonal(zeros(elty, 10), zeros(elty, 9), uplo)
                    BDone = Bidiagonal(ones(elty, 10), zeros(elty, 9), uplo)
                    BDmix = Bidiagonal(zeros(elty, 10), zeros(elty, 9), uplo)
                    BDmix[end,end] = one(elty)

                    @test iszero(BDzero)
                    @test !isone(BDzero)
                    @test !iszero(BDone)
                    @test isone(BDone)
                    @test !iszero(BDmix)
                    @test !isone(BDmix)
                end
            end

            @testset "diag" begin
                @test (@inferred diag(T))::typeof(dv) == dv
                @test (@inferred diag(T, uplo == :U ? 1 : -1))::typeof(dv) == ev
                @test (@inferred diag(T,2))::typeof(dv) == zeros(elty, n-2)
                @test_throws ArgumentError diag(T, -n - 1)
                @test_throws ArgumentError diag(T,  n + 1)
                # test diag with another wrapped vector type
                gdv, gev = GenericArray(dv), GenericArray(ev)
                G = Bidiagonal(gdv, gev, uplo)
                @test (@inferred diag(G))::typeof(gdv) == gdv
                @test (@inferred diag(G, uplo == :U ? 1 : -1))::typeof(gdv) == gev
                @test (@inferred diag(G,2))::typeof(gdv) == GenericArray(zeros(elty, n-2))
            end

            @testset "Binary operations" begin
                @test -T == Bidiagonal(-T.dv,-T.ev,T.uplo)
                @test convert(elty,-1.0) * T == Bidiagonal(-T.dv,-T.ev,T.uplo)
                @test T / convert(elty,-1.0) == Bidiagonal(-T.dv,-T.ev,T.uplo)
                @test T * convert(elty,-1.0) == Bidiagonal(-T.dv,-T.ev,T.uplo)
                @testset for uplo2 in (:U, :L)
                    dv = convert(Vector{elty}, relty <: AbstractFloat ? randn(n) : rand(1:10, n))
                    ev = convert(Vector{elty}, relty <: AbstractFloat ? randn(n-1) : rand(1:10, n-1))
                    T2 = Bidiagonal(dv, ev, uplo2)
                    Tfull2 = Array(T2)
                    for op in (+, -, *)
                        @test Array(op(T, T2)) ≈ op(Tfull, Tfull2)
                    end
                end
                # test pass-through of mul! for SymTridiagonal*Bidiagonal
                TriSym = SymTridiagonal(T.dv, T.ev)
                @test Array(TriSym*T) ≈ Array(TriSym)*Array(T)
                # test pass-through of mul! for AbstractTriangular*Bidiagonal
                Tri = UpperTriangular(diagm(1 => T.ev))
                Dia = Diagonal(T.dv)
                @test Array(Tri*T) ≈ Array(Tri)*Array(T)
                # test mul! itself for these types
                for AA in (Tri, Dia)
                    for f in (identity, transpose, adjoint)
                        C = rand(elty, n, n)
                        D = copy(C) + 2.0 * Array(f(AA) * T)
                        mul!(C, f(AA), T, 2.0, 1.0) ≈ D
                    end
                end
                # test mul! for BiTrySym * adjoint/transpose AbstractMat
                for f in (identity, transpose, adjoint)
                    C = relty == Int ? rand(float(elty), n, n) : rand(elty, n, n)
                    B = rand(elty, n, n)
                    D = copy(C) + 2.0 * Array(T*f(B))
                    mul!(C, T, f(B), 2.0, 1.0) ≈ D
                end

                # Issue #31870
                # Bi/Tri/Sym times Diagonal
                Diag = Diagonal(rand(elty, 10))
                BidiagU = Bidiagonal(rand(elty, 10), rand(elty, 9), 'U')
                BidiagL = Bidiagonal(rand(elty, 10), rand(elty, 9), 'L')
                Tridiag = Tridiagonal(rand(elty, 9), rand(elty, 10), rand(elty, 9))
                SymTri = SymTridiagonal(rand(elty, 10), rand(elty, 9))

                mats = [Diag, BidiagU, BidiagL, Tridiag, SymTri]
                for a in mats
                    for b in mats
                        @test a*b ≈ Matrix(a)*Matrix(b)
                    end
                end

                @test typeof(BidiagU*Diag) <: LinearAlgebra.Bidiagonal
                @test typeof(BidiagL*Diag) <: LinearAlgebra.Bidiagonal
                @test typeof(Tridiag*Diag) <: LinearAlgebra.Tridiagonal
                @test typeof(SymTri*Diag)  <: LinearAlgebra.Tridiagonal

                @test typeof(BidiagU*Diag) <: LinearAlgebra.Bidiagonal
                @test typeof(Diag*BidiagL) <: LinearAlgebra.Bidiagonal
                @test typeof(Diag*Tridiag) <: LinearAlgebra.Tridiagonal
                @test typeof(Diag*SymTri)  <: LinearAlgebra.Tridiagonal
            end

            # using BandedMatrices.jl backup, which is broken
            @test_skip inv(T)*Tfull ≈ Matrix(I, n, n)
            @test_broken factorize(T) === T
        end
        BD = Bidiagonal(dv, ev, :U)
        @test Matrix{ComplexF64}(BD) == BD
    end

    # Issue 10742 and similar
    let A = Bidiagonal([1,2,3], [0,0], :U)
        @test istril(A)
        @test isdiag(A)
    end

    # test construct from range
    @test Bidiagonal(1:3, 1:2, :U) == [1 1 0; 0 2 2; 0 0 3]

    @testset "promote_rule" begin
        A = Bidiagonal(fill(1f0,10),fill(1f0,9),:U)
        B = rand(Float64,10,10)
        C = Tridiagonal(rand(Float64,9),rand(Float64,10),rand(Float64,9))
        @test promote_rule(Matrix{Float64}, Bidiagonal{Float64}) == Matrix{Float64}
        @test promote(B,A) == (B, convert(Matrix{Float64}, A))
        @test promote(B,A) isa Tuple{Matrix{Float64}, Matrix{Float64}}
        @test promote(C,A) == (C,Tridiagonal(zeros(Float64,9),convert(Vector{Float64},A.dv),convert(Vector{Float64},A.ev)))
        @test promote(C,A) isa Tuple{Tridiagonal, Tridiagonal}
    end

    @testset "pathological promotion (#24707)" begin
        @test promote_type(Matrix{Int}, Bidiagonal{Tuple{S}} where S<:Integer) <: Matrix
        @test promote_type(Matrix{Tuple{T}} where T<:Integer, Bidiagonal{Tuple{S}} where S<:Integer) <: Matrix
        @test promote_type(Matrix{Tuple{T}} where T<:Integer, Bidiagonal{Int}) <: Matrix
        @test promote_type(Tridiagonal{Int}, Bidiagonal{Tuple{S}} where S<:Integer) <: Tridiagonal
        @test promote_type(Tridiagonal{Tuple{T}} where T<:Integer, Bidiagonal{Tuple{S}} where S<:Integer) <: Tridiagonal
        @test promote_type(Tridiagonal{Tuple{T}} where T<:Integer, Bidiagonal{Int}) <: Tridiagonal
    end

    @testset "solve with matrix elements" begin
        A = triu(tril(randn(9, 9), 3), -3)
        b = randn(9)
        Alb = Bidiagonal(Any[tril(A[1:3,1:3]), tril(A[4:6,4:6]), tril(A[7:9,7:9])],
                        Any[triu(A[4:6,1:3]), triu(A[7:9,4:6])], 'L')
        Aub = Bidiagonal(Any[triu(A[1:3,1:3]), triu(A[4:6,4:6]), triu(A[7:9,7:9])],
                        Any[tril(A[1:3,4:6]), tril(A[4:6,7:9])], 'U')
        bb = Any[b[1:3], b[4:6], b[7:9]]
        @test vcat((Alb\bb)...) ≈ LowerTriangular(A)\b
        @test vcat((Aub\bb)...) ≈ UpperTriangular(A)\b
    end

    @testset "sum, mapreduce" begin
        Bu = Bidiagonal([1,2,3], [1,2], :U)
        Budense = Matrix(Bu)
        Bl = Bidiagonal([1,2,3], [1,2], :L)
        Bldense = Matrix(Bl)
        @test sum(Bu) == 9
        @test sum(Bl) == 9
        @test_throws ArgumentError sum(Bu, dims=0)
        @test sum(Bu, dims=1) == sum(Budense, dims=1)
        @test sum(Bu, dims=2) == sum(Budense, dims=2)
        @test sum(Bu, dims=3) == sum(Budense, dims=3)
        @test typeof(sum(Bu, dims=1)) == typeof(sum(Budense, dims=1))
        @test mapreduce(one, min, Bu, dims=1) == mapreduce(one, min, Budense, dims=1)
        @test mapreduce(one, min, Bu, dims=2) == mapreduce(one, min, Budense, dims=2)
        @test mapreduce(one, min, Bu, dims=3) == mapreduce(one, min, Budense, dims=3)
        @test typeof(mapreduce(one, min, Bu, dims=1)) == typeof(mapreduce(one, min, Budense, dims=1))
        @test mapreduce(zero, max, Bu, dims=1) == mapreduce(zero, max, Budense, dims=1)
        @test mapreduce(zero, max, Bu, dims=2) == mapreduce(zero, max, Budense, dims=2)
        @test mapreduce(zero, max, Bu, dims=3) == mapreduce(zero, max, Budense, dims=3)
        @test typeof(mapreduce(zero, max, Bu, dims=1)) == typeof(mapreduce(zero, max, Budense, dims=1))
        @test_throws ArgumentError sum(Bl, dims=0)
        @test sum(Bl, dims=1) == sum(Bldense, dims=1)
        @test sum(Bl, dims=2) == sum(Bldense, dims=2)
        @test sum(Bl, dims=3) == sum(Bldense, dims=3)
        @test typeof(sum(Bl, dims=1)) == typeof(sum(Bldense, dims=1))
        @test mapreduce(one, min, Bl, dims=1) == mapreduce(one, min, Bldense, dims=1)
        @test mapreduce(one, min, Bl, dims=2) == mapreduce(one, min, Bldense, dims=2)
        @test mapreduce(one, min, Bl, dims=3) == mapreduce(one, min, Bldense, dims=3)
        @test typeof(mapreduce(one, min, Bl, dims=1)) == typeof(mapreduce(one, min, Bldense, dims=1))
        @test mapreduce(zero, max, Bl, dims=1) == mapreduce(zero, max, Bldense, dims=1)
        @test mapreduce(zero, max, Bl, dims=2) == mapreduce(zero, max, Bldense, dims=2)
        @test mapreduce(zero, max, Bl, dims=3) == mapreduce(zero, max, Bldense, dims=3)
        @test typeof(mapreduce(zero, max, Bl, dims=1)) == typeof(mapreduce(zero, max, Bldense, dims=1))

        Bu = Bidiagonal([2], Int[], :U)
        Budense = Matrix(Bu)
        Bl = Bidiagonal([2], Int[], :L)
        Bldense = Matrix(Bl)
        @test sum(Bu) == 2
        @test sum(Bl) == 2
        @test_throws ArgumentError sum(Bu, dims=0)
        @test sum(Bu, dims=1) == sum(Budense, dims=1)
        @test sum(Bu, dims=2) == sum(Budense, dims=2)
        @test sum(Bu, dims=3) == sum(Budense, dims=3)
        @test typeof(sum(Bu, dims=1)) == typeof(sum(Budense, dims=1))
    end

    @testset "empty sub-diagonal" begin
        # `mul!` must use non-specialized method when sub-diagonal is empty
        A = [1 2 3 4]'
        @test A * Tridiagonal(ones(1, 1)) == A
    end

    @testset "generalized dot" begin
        for elty in (Float64, ComplexF64)
            dv = randn(elty, 5)
            ev = randn(elty, 4)
            x = randn(elty, 5)
            y = randn(elty, 5)
            for uplo in (:U, :L)
                B = Bidiagonal(dv, ev, uplo)
                @test dot(x, B, y) ≈ dot(B'x, y) ≈ dot(x, Matrix(B), y)
            end
        end
    end

    @testset "multiplication of bidiagonal and triangular matrix" begin
        n = 5
        for eltyB in (Int, ComplexF64)
            if eltyB == Int
                BU = Bidiagonal(rand(1:7, n), rand(1:7, n - 1), :U)
                BL = Bidiagonal(rand(1:7, n), rand(1:7, n - 1), :L)
            else
                BU = Bidiagonal(randn(eltyB, n), randn(eltyB, n - 1), :U)
                BL = Bidiagonal(randn(eltyB, n), randn(eltyB, n - 1), :L)
            end
            for eltyT in (Int, ComplexF64)
                for TriT in (LowerTriangular, UnitLowerTriangular, UpperTriangular, UnitUpperTriangular)
                    if eltyT == Int
                        T = TriT(rand(1:7, n, n))
                    else
                        T = TriT(randn(eltyT, n, n))
                    end
                    for B in (BU, BL)
                        MB = Matrix(B)
                        MT = Matrix(T)
                        for transB in (identity, adjoint, transpose), transT in (identity, adjoint, transpose)
                            @test transB(B) * transT(T) ≈ transB(MB) * transT(MT)
                            @test transT(T) * transB(B) ≈ transT(MT) * transB(MB)
                        end
                    end
                end
            end
        end
    end

    struct MyNotANumberType
        n::Float64
    end
    Base.zero(n::MyNotANumberType)      = MyNotANumberType(zero(Float64))
    Base.zero(T::Type{MyNotANumberType}) = MyNotANumberType(zero(Float64))
    Base.copy(n::MyNotANumberType)      = MyNotANumberType(copy(n.n))
    Base.transpose(n::MyNotANumberType) = n

    @testset "transpose for a non-numeric eltype" begin
        @test !(MyNotANumberType(1.0) isa Number)
        a = [MyNotANumberType(1.0), MyNotANumberType(2.0), MyNotANumberType(3.0)]
        b = [MyNotANumberType(5.0), MyNotANumberType(6.0)]
        B = Bidiagonal(a, b, :U)
        tB = transpose(B)
        @test tB == Bidiagonal(a, b, :L)
        @test transpose(copy(tB)) == B
    end

    @testset "empty bidiagonal matrices" begin
        dv0 = zeros(0)
        ev0 = zeros(0)
        zm = zeros(0, 0)
        ubd = Bidiagonal(dv0, ev0, :U)
        lbd = Bidiagonal(dv0, ev0, :L)
        @test size(ubd) == (0, 0)
        @test_throws BoundsError getindex(ubd, 1, 1)
        @test_throws BoundsError setindex!(ubd, 0.0, 1, 1)
        @test similar(ubd) == ubd
        @test similar(lbd, Int) == zeros(Int, 0, 0)
        @test ubd == zm
        @test lbd == zm
        @test ubd == lbd
        @test ubd * ubd == ubd
        @test lbd + lbd == lbd
        @test lbd' == ubd
        @test ubd' == lbd
        @test triu(ubd, 1) == ubd
        @test triu(lbd, 1) == ubd
        @test tril(ubd, -1) == ubd
        @test tril(lbd, -1) == ubd
        @test_throws ArgumentError triu(ubd)
        @test_throws ArgumentError tril(ubd)
        @test sum(ubd) == 0.0
        @test reduce(+, ubd) == 0.0
        @test reduce(+, ubd, dims=1) == zeros(1, 0)
        @test reduce(+, ubd, dims=2) == zeros(0, 1)
        @test hcat(ubd, ubd) == zm
        @test vcat(ubd, lbd) == zm
        @test hcat(lbd, ones(0, 3)) == ones(0, 3)
        @test fill!(copy(ubd), 1.0) == ubd
        @test map(abs, ubd) == zm
        @test lbd .+ 1 == zm
        @test lbd + ubd isa Bidiagonal
        @test_broken lbd .+ ubd isa Bidiagonal
        @test ubd * 5 == ubd
        @test ubd .* 3 == ubd
    end

    @testset "mixed" begin
        B = Bidiagonal(1:5, Ones{Int}(4), :U)
        @test B.dv ≡ 1:5
        @test B.ev ≡ Ones{Int}(4)
    end
end # module TestBidiagonal


text:
function raindrops(number)
    has_factor(f) = number % f == 0
    rain = (has_factor(3) ? "Pling" : "") * (has_factor(5) ? "Plang" : "") * (has_factor(7) ? "Plong" : "")
    length(rain) == 0 ? string(number) : rain
end

max_stars_repo_name: vladflore/exercism

max_issues_repo_name: vladflore/exercism

max_forks_repo_name: vladflore/exercism


text:
# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

using .JuMP
using .Ipopt

function _lsif_coeffs(H, h, dre::LSIF, optlib::Type{JuMPLib})
  # retrieve parameters
  λ, b = dre.λ, length(h)

  model = Model(optimizer_with_attributes(Ipopt.Optimizer, "print_level" => 0, "sb" => "yes"))
  @variable(model, α[1:b])
  @objective(model, Min, (1/2) * dot(α, H*α - 2h) + λ * sum(α))
  @constraint(model, α .≥ 0)

  # solve the problem
  optimize!(model)

  # density ratio
  value.(α)
end