using .Networks

using Dates
using DataStructures
using JuMP
using Printf
using Parameters

@with_kw mutable struct RSOAssessmentModel <: AbstractModelContainer
    model::Model = Model()

    # Uncertainty model
    #bus,ts
    uncertain_load::SortedDict{Tuple{String, DateTime},VariableRef} =
        SortedDict{Tuple{String, DateTime},VariableRef}()
    #limitableUnit,ts
    uncertain_prod::SortedDict{Tuple{String, DateTime},VariableRef} =
        SortedDict{Tuple{String, DateTime},VariableRef}()

    # Market model
    #unit, ts
    p_injected::SortedDict{Tuple{String, DateTime},VariableRef} =
        SortedDict{Tuple{String, DateTime},VariableRef}()
    #pilotable, ts
    b_in::SortedDict{Tuple{String, DateTime},VariableRef} =
        SortedDict{Tuple{String, DateTime},VariableRef}()
    b_marg::SortedDict{Tuple{String, DateTime},VariableRef} =
        SortedDict{Tuple{String, DateTime},VariableRef}()
    b_out::SortedDict{Tuple{String, DateTime},VariableRef} =
        SortedDict{Tuple{String, DateTime},VariableRef}()

    #ts : rho (subproblem variable)
    overflow::SortedDict{DateTime,VariableRef} =
        SortedDict{DateTime,VariableRef}()

    # duals : branch,ts (subproblem)
    pos_flow_limit_duals::SortedDict{Tuple{String,DateTime},VariableRef} =
        SortedDict{Tuple{String,DateTime},VariableRef}()
    pos_flow_limit_indicators::SortedDict{Tuple{String,DateTime},VariableRef} =
        SortedDict{Tuple{String,DateTime},VariableRef}()
    neg_flow_limit_duals::SortedDict{Tuple{String,DateTime},VariableRef} =
        SortedDict{Tuple{String,DateTime},VariableRef}()
    neg_flow_limit_indicators::SortedDict{Tuple{String,DateTime},VariableRef} =
        SortedDict{Tuple{String,DateTime},VariableRef}()

    #sum_bus ptdf[branch,bus] * (prod_bus - conso_bus)
    signed_flow::SortedDict{Tuple{String,DateTime},AffExpr} =
        SortedDict{Tuple{String,DateTime},AffExpr}()

    #objective
    full_obj::AffExpr = AffExpr(0.)
end

function add_load_uncertainties_vars!(model_container::RSOAssessmentModel,
                                    network, TS, assessment_uncertainties)
    for bus in Networks.get_buses(network)
        bus_id = Networks.get_id(bus)
        lb_l = get_assessment_uncertainties_lb(assessment_uncertainties, bus_id)
        ub_l = get_assessment_uncertainties_ub(assessment_uncertainties, bus_id)
        for ts in TS
            name =  @sprintf("uncertain_load[%s,%s]", bus_id, ts);
            model_container.uncertain_load[bus_id,ts] = @variable(model_container.model, base_name = name,
                                                                lower_bound = lb_l, upper_bound = ub_l)
        end
    end

    return model_container.uncertain_load
end
function add_prod_uncertainties_vars!(model_container::RSOAssessmentModel,
                                    network, TS, assessment_uncertainties, tso_actions)
    for limitable_gen in Networks.get_generators_of_type(network, Networks.LIMITABLE)
        gen_id = Networks.get_id(limitable_gen)
        lb_l = get_assessment_uncertainties_lb(assessment_uncertainties, gen_id)
        for ts in TS
            # this allows to set p_inj[gen_id] == uncertain_prod without needing to express min(uncertain, limit) in the model
            ub_l = min(get_assessment_uncertainties_ub(assessment_uncertainties, gen_id),
                    get_limitation(tso_actions, gen_id, ts),
                    Networks.get_p_max(limitable_gen))
            name =  @sprintf("uncertain_prod[%s,%s]", gen_id, ts);
            model_container.uncertain_prod[gen_id, ts] = @variable(model_container.model, base_name=name,
                                                                lower_bound = lb_l, upper_bound = ub_l)
        end
    end

    return model_container.uncertain_prod
end
function add_uncertainties_vars!(model_container::RSOAssessmentModel,
                                network, TS, assessment_uncertainties, tso_actions)
    add_load_uncertainties_vars!(model_container,
                                network, TS, assessment_uncertainties)
    add_prod_uncertainties_vars!(model_container,
                                network, TS, assessment_uncertainties, tso_actions)
    return model_container
end

function add_limitable_prod_vars!(model_container::RSOAssessmentModel, network, TS)
    for limitable_gen in Networks.get_generators_of_type(network, Networks.LIMITABLE)
        gen_id = Networks.get_id(limitable_gen)
        p_max = Networks.get_p_max(limitable_gen)
        for ts in TS
            name =  @sprintf("p_injected[%s,%s]", gen_id, ts)
            model_container.p_injected[gen_id,ts] =
                @variable(model_container.model, base_name=name, lower_bound=0., upper_bound=p_max)
        end
    end

    return model_container.p_injected
end

function add_use_limitables_constraints!(model_container::RSOAssessmentModel, network, TS, tso_actions)
    for limitable_gen in Networks.get_generators_of_type(network, Networks.LIMITABLE)
        gen_id = Networks.get_id(limitable_gen)
        for ts in TS
            limit_l = get_limitation(tso_actions, gen_id, ts)
            @assert(upper_bound(model_container.uncertain_prod[gen_id, ts]) <= limit_l)
            @constraint(model_container.model,
                        model_container.p_injected[gen_id,ts] == model_container.uncertain_prod[gen_id, ts])
        end
    end
    return model_container
end

function add_pilotable_prod_vars!(model_container_p::RSOAssessmentModel, network, TS, tso_actions)
    p_injected_l = model_container_p.p_injected
    b_in_l = model_container_p.b_in
    b_marg_l = model_container_p.b_marg
    b_out_l = model_container_p.b_out

    for pilotable_gen in Networks.get_generators_of_type(network, Networks.PILOTABLE)
        gen_id = Networks.get_id(pilotable_gen)
        for ts in TS
            pmin_l, pmax_l = safeget_imposition(tso_actions, gen_id, ts)
            name_l =  @sprintf("b_in[%s,%s]", gen_id, ts);
            b_in_l[gen_id,ts] = @variable(model_container_p.model, base_name = name_l, binary = true)
            name_l =  @sprintf("b_marginal[%s,%s]", gen_id, ts);
            b_marg_l[gen_id,ts] = @variable(model_container_p.model, base_name = name_l, binary = true)
            name_l =  @sprintf("b_out[%s,%s]", gen_id, ts);
            b_out_l[gen_id,ts] = @variable(model_container_p.model, base_name = name_l, binary = true)
            name_l =  @sprintf("p_injected[%s,%s]", gen_id, ts);
            p_injected_l[gen_id,ts] = @variable(model_container_p.model, base_name = name_l,
                                                lower_bound=0., upper_bound = pmax_l);

            @constraint(model_container_p.model,
                        p_injected_l[gen_id,ts] >= ( pmax_l * b_in_l[gen_id,ts] +  pmin_l * b_marg_l[gen_id,ts] ));
            @constraint(model_container_p.model,
                        p_injected_l[gen_id,ts] <= (1 - b_out_l[gen_id,ts]) * pmax_l);
            @constraint(model_container_p.model,
                        b_out_l[gen_id,ts] + b_marg_l[gen_id,ts] + b_in_l[gen_id,ts] == 1.);
        end
    end

    return model_container_p
end

function add_market_vars!(model_container::RSOAssessmentModel, network, TS, tso_actions)
    add_limitable_prod_vars!(model_container, network, TS)
    add_pilotable_prod_vars!(model_container, network, TS, tso_actions)
end

function add_cheapest_prod_constraints!(model_container::RSOAssessmentModel, network, TS)
    pilotables = Networks.get_generators_of_type(network, Networks.PILOTABLE)
    for (index, gen_1) in enumerate(pilotables)
        cost_1 = Networks.get_prop_cost(gen_1)
        gen_id_1 = Networks.get_id(gen_1)
        for gen_2 in pilotables[index:end]
            cost_2 = Networks.get_prop_cost(gen_2)
            if cost_1 < cost_2
                gen_id_2 = Networks.get_id(gen_2)
                for ts in TS
                    @constraint(model_container.model, model_container.b_in[gen_id_1,ts] >= model_container.b_in[gen_id_2,ts] );
                end
            end
        end
    end

    return model_container
end

function add_eod_constraint!(model_container_p::RSOAssessmentModel, network, TS)

    for ts in TS

        supply_l = AffExpr(0)
        for gen in Networks.get_generators(network)
            gen_id = Networks.get_id(gen)
            supply_l += model_container_p.p_injected[gen_id, ts];
        end

        demand_l = AffExpr(0.)
        for bus in Networks.get_buses(network)
            bus_id = Networks.get_id(bus)
            demand_l += model_container_p.uncertain_load[bus_id, ts];
        end

        @constraint(model_container_p.model, supply_l == demand_l);
    end

    return model_container_p
end

function add_overflow_var!(model_container::RSOAssessmentModel, TS, upper_bound)
    for ts in TS
        name = @sprintf("overflow[%s]", ts)
        model_container.overflow[ts] =
            @variable(model_container.model, base_name=name, lower_bound=0., upper_bound=upper_bound)
    end
    return model_container.overflow
end

function create_objective!(model_container::RSOAssessmentModel, network, TS)::AffExpr
    model_container.full_obj = AffExpr(0.)
    add_coeffxsum_cost!(model_container.full_obj, model_container.overflow, 1.)
    return model_container.full_obj
end

function add_uncertainties_constraints!(model_container_l::RSOAssessmentModel, network, TS, assessment_uncertainties, configs)
    #FIXME constraints on uncertainties go here !
    #e.g | u[bus1] - u[bus2] | < 5
    #e.g ( 1 - u[bus1] / ub[limitable_1] ) <  ( 1 - u[limitable_1] / ub[limitable_1] )
end

function create_flow_limit_duals_and_indicators(model_container::RSOAssessmentModel, network, TS)
    for branch in Networks.get_branches(network)
        branch_id = Networks.get_id(branch)
        for ts in TS
            name=@sprintf("pos_flow_duals[%s,%s]", branch_id, ts)
            model_container.pos_flow_limit_duals[branch_id, ts] =
                @variable(model_container.model, base_name=name, lower_bound=0., upper_bound=1.)
            name=@sprintf("pos_flow_indicators[%s,%s]", branch_id, ts)
            model_container.pos_flow_limit_indicators[branch_id, ts] =
                @variable(model_container.model, base_name=name, binary=true)
            name=@sprintf("neg_flow_duals[%s,%s]", branch_id, ts)
            model_container.neg_flow_limit_duals[branch_id, ts] =
                @variable(model_container.model, base_name=name, lower_bound=0., upper_bound=1.)
            name=@sprintf("neg_flow_indicators[%s,%s]", branch_id, ts)
            model_container.neg_flow_limit_indicators[branch_id, ts] =
                @variable(model_container.model, base_name=name, binary=true)
        end
    end

    return model_container
end

function create_flow_expressions(model_container::RSOAssessmentModel, network, TS)
    for branch in Networks.get_branches(network)
        branch_id = Networks.get_id(branch)

        for ts in TS

            branch_flow_expr = AffExpr(0.)
            for bus in Networks.get_buses(network)
                bus_id = Networks.get_id(bus)
                ptdf = Networks.safeget_ptdf_elt(network, branch_id, bus_id)

                # + injections
                for gen in Networks.get_generators(bus)
                    gen_id = Networks.get_id(gen)
                    branch_flow_expr += ptdf * model_container.p_injected[gen_id,ts]
                end

                # - loads
                branch_flow_expr -= ptdf * model_container.uncertain_load[bus_id,ts]
            end

            model_container.signed_flow[branch_id, ts] = branch_flow_expr
        end
    end
end

function add_flow_constraints(model_container::RSOAssessmentModel, network, TS)
    model_l = model_container.model
    for branch in Networks.get_branches(network)
        branch_id = Networks.get_id(branch)
        flow_limit_l = Networks.get_limit(branch)

        for ts in TS
            overflow_l = model_container.overflow[ts]
            signed_flow_l = model_container.signed_flow[branch_id, ts]

            name = @sprintf("c_pos_flow_limit[%s,%s]",branch_id,ts)
            @constraint(model_l, signed_flow_l <= flow_limit_l + overflow_l , base_name=name)
            name = @sprintf("c_neg_flow_limit[%s,%s]",branch_id,ts)
            @constraint(model_l, -signed_flow_l <= flow_limit_l + overflow_l , base_name=name)

        end
    end
end

function add_kkt_stationarity_constraints(model_container::RSOAssessmentModel, network, TS)
    for ts in TS
        duals_sum_l = AffExpr(0.)
        for branch_id in Networks.get_id.(Networks.get_branches(network))
            add_to_expression!(duals_sum_l,
                model_container.pos_flow_limit_duals[branch_id,ts] + model_container.neg_flow_limit_duals[branch_id,ts])
        end
        @constraint(model_container.model, duals_sum_l == 1.)
    end

    return model_container
end

function add_kkt_complementarity_constraints(model_container, network, TS)
    for branch in Networks.get_branches(network)
        branch_id = Networks.get_id(branch)
        flow_limit_l = Networks.get_limit(branch)

        for ts in TS
            # pos_flow_limit_duals
            kkt_var = model_container.pos_flow_limit_duals[branch_id, ts]
            cstr_expr = - model_container.signed_flow[branch_id, ts] + flow_limit_l + model_container.overflow[ts]
            b_indicator = model_container.pos_flow_limit_indicators[branch_id, ts]
            ub_cstr = compute_ub(cstr_expr)
            ub_kkt = upper_bound(kkt_var)
            formulate_complementarity_constraints!(model_container.model,
                                                    kkt_var, cstr_expr, b_indicator, ub_kkt, ub_cstr)
            # neg_flow_limit_duals
            kkt_var = model_container.neg_flow_limit_duals[branch_id, ts]
            cstr_expr = model_container.signed_flow[branch_id, ts] + flow_limit_l + model_container.overflow[ts]
            b_indicator = model_container.neg_flow_limit_indicators[branch_id, ts]
            ub_cstr = compute_ub(cstr_expr)
            ub_kkt = upper_bound(kkt_var)
            formulate_complementarity_constraints!(model_container.model,
                                                    kkt_var, cstr_expr, b_indicator, ub_kkt, ub_cstr)
        end
    end

    return model_container
end

function formulate_overflow_subproblem(model_container::RSOAssessmentModel, network, TS)
    create_flow_limit_duals_and_indicators(model_container, network, TS)
    create_flow_expressions(model_container, network, TS)
    add_flow_constraints(model_container, network, TS)
    add_kkt_stationarity_constraints(model_container, network, TS)
    add_kkt_complementarity_constraints(model_container, network, TS)

    return model_container
end

function formulate_rso_assessment(network, TS, assessment_uncertainties, tso_actions, configs)
    model_container_l = RSOAssessmentModel()

    add_uncertainties_vars!(model_container_l, network, TS, assessment_uncertainties, tso_actions)
    add_market_vars!(model_container_l, network, TS, tso_actions)
    add_overflow_var!(model_container_l, TS, configs.BIG_M)

    add_use_limitables_constraints!(model_container_l, network, TS, tso_actions)
    add_cheapest_prod_constraints!(model_container_l, network, TS)
    add_eod_constraint!(model_container_l, network, TS)
    add_uncertainties_constraints!(model_container_l, network, TS, assessment_uncertainties, configs)

    formulate_overflow_subproblem(model_container_l, network, TS)

    create_objective!(model_container_l, network, TS)
    @objective(get_model(model_container_l), Max, model_container_l.full_obj)

    return model_container_l
end

function is_validated(model_container::RSOAssessmentModel)
    return value(model_container.full_obj) < 1e-09
end

function has_positive_slack(model_container::RSOAssessmentModel)::Bool
    return false # ( has_positive_value(model_container.lower.lol_model.p_loss_of_load)
           #  || has_positive_value(model_container.lower.lol_model.p_cut_prod) )
end

function get_assessment_uncertainties_lb(assessment_uncertainties, bus_or_limitable::String)
    return assessment_uncertainties[bus_or_limitable][1]
end
function get_assessment_uncertainties_ub(assessment_uncertainties, bus_or_limitable::String)
    return assessment_uncertainties[bus_or_limitable][2]
end


SubWorker{T,A,PT} = RemoteChannel{Channel{Vector{SubProblem{T,A,PT}}}}
ScenarioProblemChannel{S} = RemoteChannel{Channel{ScenarioProblems{S}}}
Work = RemoteChannel{Channel{Int}}
Progress{T <: AbstractFloat} = Tuple{Int,Int,SubproblemSolution{T}}
ProgressQueue{T <: AbstractFloat} = RemoteChannel{Channel{Progress{T}}}

function initialize_subproblems!(ph::AbstractProgressiveHedging,
                                 subworkers::Vector{SubWorker{T,A,PT}},
                                 scenarioproblems::DistributedScenarioProblems,
                                 penaltyterm::AbstractPenaltyterm) where {T <: AbstractFloat,
                                                                          A <: AbstractVector,
                                                                          PT <: AbstractPenaltyterm}
    # Create subproblems on worker processes
    @sync begin
        for w in workers()
            subworkers[w-1] = RemoteChannel(() -> Channel{Vector{SubProblem{T,A,PT}}}(1), w)
            prev = map(2:(w-1)) do p
                scenarioproblems.scenario_distribution[p-1]
            end
            start_id = isempty(prev) ? 0 : sum(prev)
            @async remotecall_fetch(initialize_subworker!,
                                    w,
                                    subworkers[w-1],
                                    scenarioproblems[w-1],
                                    penaltyterm,
                                    start_id)
        end
    end
    return nothing
end

function update_dual_gap!(ph::AbstractProgressiveHedging,
                          subworkers::Vector{<:SubWorker{T}}) where T <: AbstractFloat
    # Update δ₂
    partial_δs = Vector{Float64}(undef, nworkers())
    @sync begin
        for (i,w) in enumerate(workers())
            @async partial_δs[i] = remotecall_fetch(
                w,
                subworkers[w-1],
                ph.ξ) do sw, ξ
                    subproblems = fetch(sw)
                    return mapreduce(+, subproblems, init = zero(T)) do subproblem
                        π = subproblem.probability
                        x = subproblem.x
                        π * norm(x - ph.ξ, 2) ^ 2
                    end
                end
        end
    end
    ph.data.δ₂ = sum(partial_δs)
    return nothing
end

function initialize_subworker!(subworker::SubWorker{T,A,PT},
                               scenarioproblems::ScenarioProblemChannel,
                               penaltyterm::AbstractPenaltyterm,
                               start_id::Integer) where {T <: AbstractFloat,
                                                         A <: AbstractArray,
                                                         PT <: AbstractPenaltyterm}
    sp = fetch(scenarioproblems)
    subproblems = Vector{SubProblem{T,A,PT}}(undef, num_subproblems(sp))
    for i = 1:num_subproblems(sp)
        subproblems[i] = SubProblem(
            subproblem(sp, i),
            start_id + i,
            T(probability(sp, i)),
            copy(penaltyterm))
    end
    put!(subworker, subproblems)
    return nothing
end

function restore_subproblems!(subworkers::Vector{<:SubWorker})
    @sync begin
        for w in workers()
            @async remotecall_fetch(w, subworkers[w-1]) do sw
                for subproblem in fetch(sw)
                    restore_subproblem!(subproblem)
                end
            end
        end
    end
    return nothing
end

function resolve_subproblems!(subworker::SubWorker{T,A,PT},
                              ξ::AbstractVector,
                              r::AbstractFloat) where {T <: AbstractFloat,
                                                       A <: AbstractArray,
                                                       PT <: AbstractPenaltyterm}
    subproblems::Vector{SubProblem{T,A,PT}} = fetch(subworker)
    Qs = Vector{SubproblemSolution{T}}(undef, length(subproblems))
    # Reformulate and solve sub problems
    for (i,subproblem) in enumerate(subproblems)
        reformulate_subproblem!(subproblem, ξ, r)
        Qs[i] = subproblem(ξ)
    end
    # Return current objective value
    return sum(Qs)
end

function collect_primals(subworker::SubWorker{T,A,PT}, n::Integer) where {T <: AbstractFloat,
                                                                          A <: AbstractArray,
                                                                          PT <: AbstractPenaltyterm}
    subproblems::Vector{SubProblem{T,A,PT}} = fetch(subworker)
    return mapreduce(+, subproblems, init = zeros(T, n)) do subproblem
        π = subproblem.probability
        x = subproblem.x
        π * x
    end
end

function calculate_objective_value(subworkers::Vector{<:SubWorker{T}}) where T <: AbstractFloat
    partial_objectives = Vector{Float64}(undef, nworkers())
    @sync begin
        for (i, w) in enumerate(workers())
            @async partial_objectives[i] = remotecall_fetch(w, subworkers[w-1]) do sw
                return mapreduce(+, fetch(sw), init = zero(T)) do subproblem
                    _objective_value(subproblem)
                end
            end
        end
    end
    return sum(partial_objectives)
end

function work_on_subproblems!(subworker::SubWorker{T,A,PT},
                              work::Work,
                              finalize::Work,
                              progress::ProgressQueue{T},
                              x̄::RemoteRunningAverage{A},
                              δ::RemoteRunningAverage{T},
                              iterates::RemoteIterates{A},
                              r::IteratedValue{T}) where {T <: AbstractFloat,
                                                          A <: AbstractArray,
                                                          PT <: AbstractPenaltyterm}
    subproblems::Vector{SubProblem{T,A,PT}} = fetch(subworker)
    if isempty(subproblems)
       # Workers has nothing do to, return.
       return nothing
    end
    x̄ = fetch(x̄)
    δ = fetch(δ)
    quit = false
    while true
        t::Int = try
            if isready(finalize)
                quit = true
                take!(finalize)
            else
                wait(work)
                take!(work)
            end
        catch err
            if err isa InvalidStateException
                # Master closed the work/finalize channel. Worker finished
                return nothing
            end
        end
        t == -1 && continue
        ξ::A = fetch(iterates, t)
        if t > 1
            update_subproblems!(subproblems, ξ, fetch(r,t-1))
        end
        for (i,subproblem) in enumerate(subproblems)
            !quit && subtract!(δ, i)
            !quit && subtract!(x̄, i)
            x = subproblem.x
            π = subproblem.probability
            !quit && add!(δ, i, norm(x - ξ, 2) ^ 2, π)
            reformulate_subproblem!(subproblem, ξ, fetch(r, t))
            Q::SubproblemSolution{T} = subproblem(ξ)
            !quit && add!(x̄, i, π)
            !quit && put!(progress, (t, subproblem.id, Q))
        end
        if quit
            # Worker finished
            return nothing
        end
    end
end


type LabeledArray{T,N,A<:AbstractArray,L<:AbstractArray} <: AbstractArray{T,N}
    data::A
    label::L
    colors::Vector{RGB}

    # Enforce that the label has Integer type
    LabeledArray{Ti<:Integer}(data::AbstractArray{T,N}, label::AbstractArray{Ti}, colors::Vector{RGB}) =
        new(data, label, colors)
end
LabeledArray{T,N}(data::AbstractArray{T,N}, label::AbstractArray, colors::Vector{RGB}) =
    LabeledArray{T,N,typeof(data),typeof(label)}(data, label, colors)

size(A::LabeledArray) = size(A.data)
size(A::LabeledArray, i::Integer) = size(A.data, i)
eltype{T}(A::LabeledArray{T}) = T
ndims{T,N}(A::LabeledArray{T,N}) = N

for N = 1:4
    @eval begin
        # All but the last of these are inside the @eval loop simply to avoid ambiguity warnings.
        # These first two are additional ones needed to avoid ambiguity warnings.
        _uint32color_gray!{T}(buf::Array{Uint32}, A::LabeledArray{T,$N}, scalei::ScaleSigned) = error("Cannot use ScaleSigned with a labeled array")
        _uint32color_gray!{T,L<:LabeledArray}(buf::Array{Uint32}, A::SubArray{T,$N,L}, scalei::ScaleSigned) = error("Cannot use ScaleSigned with a labeled array")

        function _uint32color_gray!{T}(buf::Array{Uint32}, A::LabeledArray{T,$N}, scalei::ScaleInfo = scaleinfo(Uint8, A))
            if size(buf) != size(A)
                error("Size mismatch")
            end
            dat = A.data
            label = A.label
            col = A.colors
            for i = 1:length(dat)
                gr = scale(scalei, dat[i])
                lbl = label[i]
                if lbl == 0
                    buf[i] = rgb24(gr,gr,gr)
                else
                    buf[i] = convert(RGB24, gr*col[lbl])
                end
            end
            buf
        end

        # For SubArrays, we can't efficiently use linear indexing, and in any event
        # we want to broadcast label where necessary
        function _uint32color_gray!{T,A<:LabeledArray}(buf::Array{Uint32}, S::SubArray{T,$N,A}, scalei::ScaleInfo = scaleinfo(Uint8, A))
            if size(buf) != size(S)
                error("Size mismatch")
            end
            indexes = S.indexes
            dat = slice(S.parent.data, indexes)
            plabel = S.parent.label
            newindexes = RangeIndex[size(plabel,i)==1 ? (isa(indexes[i], Int) ? 1 : (1:1)) : indexes[i]  for i = 1:ndims(plabel)]
            label = slice(plabel, newindexes...)
            col = S.parent.colors
            _uint32color_labeled(buf, dat, label, col, scalei) # type of label can't be inferred, use function boundary
        end

        function _uint32color_labeled{T}(buf, dat::AbstractArray{T,$N}, label, col, scalei)
            k = 0
            @inbounds @nloops $N i buf begin
                val = @nref $N dat i
                gr = scale(scalei, val)
                lbl = @nref $N label i
                if lbl == 0
                    buf[k+=1] = rgb24(gr,gr,gr)
                else
                    buf[k+=1] = convert(RGB24, gr*col[lbl])
                end
            end
            buf
        end
    end
end


# All the types supported by mshadow.
typealias DType Union{Float32, Float64, Float16, UInt8, Int32}
@enum TypeFlag kFloat32 kFloat64 kFloat16 kUint8 kInt32
typealias DEFAULT_DTYPE Float32

function toTypeFlag{T <: DType}(:: Type{T})
  if T == Float32
    return kFloat32
  elseif T == Float64
    return kFloat64
  elseif T == Float16
    return kFloat16
  elseif T == UInt8
    return kUint8
  elseif T == Int32
    return kInt32
  else
    throw(ArgumentError("Can't convert $T to DType."))
  end
end

function fromTypeFlag(T :: TypeFlag)
  if T == kFloat32
    return Float32
  elseif T == kFloat64
    return Float64
  elseif T == kFloat16
    return Float16
  elseif T == kUint8
    return UInt8
  elseif T == kInt32
    return Int32
  else
    throw(ArgumentError("Can't convert DType $T."))
  end
end

# create a NDArray handle of specific shape
function _ndarray_alloc{N}(shape :: NTuple{N, Int}, ctx :: Context, delay_alloc :: Bool)
  h_ref  = Ref{MX_handle}(0)
  shape  = flipdim(MX_uint[shape...],1)
  @mxcall(:MXNDArrayCreate, (Ptr{MX_uint}, MX_uint, Cint, Cint, Cint, Ref{MX_handle}),
      shape, length(shape), ctx.device_type, ctx.device_id, delay_alloc, h_ref)
  handle = MX_NDArrayHandle(h_ref[])
  return handle
end

# create a NDArray handle of specific shape type
function _ndarray_alloc{T <: DType,N}(:: Type{T}, shape :: NTuple{N, Int}, ctx :: Context, delay_alloc :: Bool)
  h_ref  = Ref{MX_handle}(0)
  shape  = flipdim(MX_uint[shape...],1)
  dtype  = toTypeFlag(T)
  @mxcall(:MXNDArrayCreateEx, (Ptr{MX_uint}, MX_uint, Cint, Cint, Cint, Cint, Ref{MX_handle}),
      shape, length(shape), ctx.device_type, ctx.device_id, delay_alloc, dtype, h_ref)
  handle = MX_NDArrayHandle(h_ref[])
  return handle
end

# create a handle to an empty NDArray, this handle can be used to hold
# results returned by libmx API calls
function _ndarray_alloc()
  h_ref = Ref{MX_handle}(0)
  @mxcall(:MXNDArrayCreateNone, (Ref{MX_handle},), h_ref)
  return MX_NDArrayHandle(h_ref[])
end

################################################################################
# NDArray Type
################################################################################
"""
    NDArray

Wrapper of the `NDArray` type in `libmxnet`. This is the basic building block
of tensor-based computation.

!!! note
      since C/C++ use row-major ordering for arrays while Julia follows a
      column-major ordering. To keep things consistent, we keep the underlying data
      in their original layout, but use *language-native* convention when we talk
      about shapes. For example, a mini-batch of 100 MNIST images is a tensor of
      C/C++/Python shape (100,1,28,28), while in Julia, the same piece of memory
      have shape (28,28,1,100).
"""
type NDArray
  handle   :: MX_NDArrayHandle
  writable :: Bool

  function NDArray(handle, writable=true)
    new(handle, writable)
  end
end

function Base.show(io :: IO, arr :: NDArray)
  print(io, "mx.NDArray{$(eltype(arr))}$(size(arr))")
end

function NDArray{T<:Real}(data :: Array{T})
  copy(data, cpu())
end

function Base.unsafe_convert(::Type{MX_handle}, obj::NDArray)
  Base.unsafe_convert(MX_handle, obj.handle)
end
Base.convert(t::Type{MX_handle}, obj::NDArray) = Base.unsafe_convert(t, obj)
Base.cconvert(t::Type{MX_handle}, obj::NDArray) = Base.unsafe_convert(t, obj)

################################################################################
# NDArray functions exported to the users
################################################################################
"""
    context(arr :: NDArray)

Get the context that this `NDArray` lives on.
"""
function context(arr :: NDArray)
  ref_typeid = Ref{Cint}(0)
  ref_devid  = Ref{Cint}(0)
  @mxcall(:MXNDArrayGetContext, (MX_handle, Ref{Cint}, Ref{Cint}),
          arr, ref_typeid, ref_devid)
  return Context(ref_typeid[], ref_devid[])
end

"""
    empty(DType, shape :: Tuple, ctx :: Context)
    empty(DType, shape :: Tuple)
    empty(DType, dim1, dim2, ...)

Allocate memory for an uninitialized `NDArray` with a specified type.
"""
function empty{N,T<:DType}(::Type{T}, shape :: NTuple{N, Int})
  empty(T, shape, cpu())
end
function empty{N,T<:DType}(:: Type{T}, shape :: NTuple{N, Int}, ctx :: Context)
  NDArray(_ndarray_alloc(T, shape, ctx, false))
end
function empty{T<:DType}(:: Type{T}, shape :: Int...)
  empty(T, shape)
end

"""
    empty(shape :: Tuple, ctx :: Context)
    empty(shape :: Tuple)
    empty(dim1, dim2, ...)

Allocate memory for an uninitialized `NDArray` with specific shape of type Float32.
"""
function empty{N}(shape :: NTuple{N, Int})
  empty(shape, cpu())
end
function empty{N}(shape :: NTuple{N, Int}, ctx :: Context)
  NDArray(_ndarray_alloc(shape, ctx, false))
end
function empty(shape :: Int...)
  empty(shape)
end

"""
    zeros(DType, shape :: Tuple, ctx :: Context)
    zeros(DType, shape :: Tuple)
    zeros(DType, dim1, dim2, ...)

Create zero-ed `NDArray` with specific shape and type
"""
function zeros{N,T<:DType}(:: Type{T}, shape :: NTuple{N, Int})
  zeros(T, shape, cpu())
end
function zeros{N,T<:DType}(:: Type{T}, shape :: NTuple{N, Int}, ctx :: Context)
  arr = empty(T, shape, ctx)
  arr[:] = zero(T)
  return arr
end
function zeros{T<:DType}(:: Type{T}, shape :: Int...)
  zeros(T, shape)
end

"""
    zeros(shape :: Tuple, ctx :: Context)
    zeros(shape :: Tuple)
    zeros(dim1, dim2, ...)

Create zero-ed `NDArray` with specific shape.
"""
function zeros{N}(shape :: NTuple{N, Int})
  zeros(shape, cpu())
end
function zeros{N}(shape :: NTuple{N, Int}, ctx :: Context)
  arr = empty(shape, ctx)
  arr[:] = 0
  return arr
end
function zeros(shape :: Int...)
  zeros(shape)
end

"""
    ones(DType, shape :: Tuple, ctx :: Context)
    ones(DType, shape :: Tuple)
    ones(DType, dim1, dim2, ...)

Create an `NDArray` with specific shape & type, and initialize with 1.
"""
function ones{N,T<:DType}(:: Type{T}, shape :: NTuple{N, Int})
  ones(T, shape, cpu())
end
function ones{N,T<:DType}(:: Type{T}, shape :: NTuple{N, Int}, ctx :: Context)
  arr = empty(T, shape, ctx)
  arr[:] = one(T)
  return arr
end
function ones{T<:DType}(:: Type{T}, shape :: Int...)
  ones(T, shape)
end

"""
    ones(shape :: Tuple, ctx :: Context)
    ones(shape :: Tuple)
    ones(dim1, dim2, ...)

Create an `NDArray` with specific shape and initialize with 1.
"""
function ones{N}(shape :: NTuple{N, Int})
  ones(shape, cpu())
end
function ones{N}(shape :: NTuple{N, Int}, ctx :: Context)
  arr = empty(shape, ctx)
  arr[:] = 1
  return arr
end
function ones(shape :: Int...)
  ones(shape)
end

import Base: size, length, ndims, eltype

"""
    size(arr :: NDArray)
    size(arr :: NDArray, dim :: Int)

Get the shape of an `NDArray`. The shape is in Julia's column-major convention. See
also the notes on NDArray shapes [`NDArray`](@ref).
"""
function size(arr :: NDArray)
  ref_ndim  = Ref{MX_uint}(0)
  ref_shape = Ref{Ptr{MX_uint}}(0)
  @mxcall(:MXNDArrayGetShape, (MX_handle, Ref{MX_uint}, Ref{Ptr{MX_uint}}),
          arr, ref_ndim, ref_shape)
  tuple(map(Int, flipdim(unsafe_wrap(Array, ref_shape[], ref_ndim[]),1))...)
end
function size(arr :: NDArray, dim :: Int)
  size(arr)[dim]
end

"""
    length(arr :: NDArray)

Get the number of elements in an `NDArray`.
"""
function length(arr :: NDArray)
  prod(size(arr))
end

"""
    ndims(arr :: NDArray)

Get the number of dimensions of an `NDArray`. Is equivalent to `length(size(arr))`.
"""
function ndims(arr :: NDArray)
  length(size(arr))
end

"""
    eltype(arr :: NDArray)

Get the element type of an `NDArray`.
"""
function eltype{T <: Union{NDArray, MX_NDArrayHandle}}(arr :: T)
  dtype_ref = Ref{Cint}(0)
  @mxcall(:MXNDArrayGetDType, (MX_handle, Ptr{Cint}), arr, dtype_ref)

  if dtype_ref[] == -1 # arr->is_none()
    warn("Eltype of $arr is not defined")
    Base.show_backtrace(STDOUT,backtrace())
    println()
    return Float32
  else
    return fromTypeFlag(TypeFlag(dtype_ref[]))
  end
end

import Base: slice
"""
    slice(arr :: NDArray, start:stop)

Create a view into a sub-slice of an `NDArray`. Note only slicing at the slowest
changing dimension is supported. In Julia's column-major perspective, this is the last
dimension. For example, given an `NDArray` of shape (2,3,4), `slice(array, 2:3)` will create
a `NDArray` of shape (2,3,2), sharing the data with the original array. This operation is
used in data parallelization to split mini-batch into sub-batches for different devices.
"""
function slice(arr :: NDArray, ::Colon)
  arr
end
function slice(arr :: NDArray, slice::UnitRange{Int})
  dim1 = size(arr)[end]
  @assert(1 <= slice.start <= slice.stop <= dim1)
  if slice.start == 1 && slice.stop == dim1
    return arr
  end

  hdr_ref = Ref{MX_handle}(0)
  # note Julia is 1-based, inclusive-inclusive indexing, while C++ is
  # 0-based, inclusive-exclusive indexing. So 1:3 in Julia should
  # translates into 0:3 in C++.
  @mxcall(:MXNDArraySlice, (MX_handle, MX_uint, MX_uint, Ref{MX_handle}),
          arr, slice.start-1, slice.stop, hdr_ref)
  return NDArray(MX_NDArrayHandle(hdr_ref[]), arr.writable)
end

import Base: setindex!

"""
    setindex!(arr :: NDArray, val, idx)

Assign values to an `NDArray`. Elementwise assignment is not implemented, only the following
scenarios are supported

* `arr[:] = val`: whole array assignment, `val` could be a scalar or an array (Julia `Array`
  or `NDArray`) of the same shape.
* `arr[start:stop] = val`: assignment to a *slice*, `val` could be a scalar or an array of
  the same shape to the slice. See also [`slice`](@ref).
"""
function setindex!(arr :: NDArray, val :: Real, ::Colon)
  @assert(arr.writable)
  _set_value(convert(eltype(arr), val), arr)
  return arr
end
function setindex!{T<:Real}(arr :: NDArray, val :: Array{T}, ::Colon)
  copy!(arr, val)
end
function setindex!(arr :: NDArray, val :: NDArray, ::Colon)
  copy!(arr, val)
end
function setindex!{T<:Real}(arr :: NDArray, val :: Union{T,Array{T},NDArray}, idx::UnitRange{Int})
  setindex!(slice(arr, idx), val, Colon())
end

import Base: getindex
"""
    getindex(arr :: NDArray, idx)

Shortcut for [`slice`](@ref). A typical use is to write

```julia
  arr[:] += 5
```

which translates into

```julia
  arr[:] = arr[:] + 5
```

which furthur translates into

```julia
  setindex!(getindex(arr, Colon()), 5, Colon())
```

!!! note
    The behavior is quite different from indexing into Julia's `Array`. For example, `arr[2:5]`
    create a **copy** of the sub-array for Julia `Array`, while for `NDArray`, this is
    a *slice* that shares the memory.
"""
function getindex(arr :: NDArray, ::Colon)
  return arr
end

"""
Shortcut for [`slice`](@ref). **NOTE** the behavior for Julia's built-in index slicing is to create a
copy of the sub-array, while here we simply call `slice`, which shares the underlying memory.
"""
function getindex(arr :: NDArray, idx::UnitRange{Int})
  slice(arr, idx)
end

import Base: copy!, copy, convert
"""
    copy!(dst :: Union{NDArray, Array}, src :: Union{NDArray, Array})

Copy contents of `src` into `dst`.
"""
function copy!(dst :: NDArray, src :: NDArray)
  @assert(dst.writable)
  if dst.handle == src.handle
    warn("Copying an NDArray to itself")
    return
  end

  _copyto(src, dst)
  return dst
end

function copy!{T<:DType}(dst :: Array{T}, src :: NDArray)
  @assert T == eltype(src)
  @assert size(dst) == size(src)
  @mxcall(:MXNDArraySyncCopyToCPU, (MX_handle, Ptr{Void}, Csize_t),
          src, pointer(dst), length(dst))
  return dst
end
function copy!{T<:Real}(dst :: Array{T}, src :: NDArray)
  copy!(dst, copy(src))
end

function copy!{T<:Real}(dst :: NDArray, src :: Array{T})
  @assert dst.writable
  @assert size(dst) == size(src)
  src = convert(Array{eltype(dst)}, src) # this might involve copying
  @mxcall(:MXNDArraySyncCopyFromCPU, (MX_handle, Ptr{Void}, Csize_t),
          dst.handle, pointer(src), length(src))
  return dst
end

function copy_ignore_shape!{T<:Real}(dst :: NDArray, src :: Array{T})
  @assert dst.writable
  @assert length(dst) == length(src)
  src = convert(Array{eltype(dst)}, src) # this might involve copying
  @mxcall(:MXNDArraySyncCopyFromCPU, (MX_handle, Ptr{Void}, Csize_t),
          dst.handle, pointer(src), length(src))
  return dst
end

"""
    copy(arr :: NDArray)
    copy(arr :: NDArray, ctx :: Context)
    copy(arr :: Array, ctx :: Context)

Create a copy of an array. When no `Context` is given, create a Julia `Array`.
Otherwise, create an `NDArray` on the specified context.
"""
# Create copy: NDArray -> Julia Array
function copy(arr :: NDArray)
  j_arr = Array{eltype(arr)}(size(arr))
  copy!(j_arr, arr)
end

# Create copy: NDArray -> NDArray in a given context
function copy(arr :: NDArray, ctx :: Context)
  dst = NDArray(_ndarray_alloc(eltype(arr), size(arr), ctx, true))
  copy!(dst, arr)
end

# Create copy: Julia Array -> NDArray in a given context
function copy{T<:DType}(arr :: Array{T}, ctx :: Context)
  dst = empty(T, size(arr), ctx)
  copy!(dst, arr)
end

"""
    convert(::Type{Array{T}}, arr :: NDArray)

Convert an `NDArray` into a Julia `Array` of specific type. Data will be copied.
"""
function convert{T<:Real}(t::Type{Array{T}}, arr :: NDArray)
  convert(t, copy(arr))
end

"""
    @inplace

Julia does not support re-definiton of `+=` operator (like `__iadd__` in python),
When one write `a += b`, it gets translated to `a = a+b`. `a+b` will allocate new
memory for the results, and the newly allocated `NDArray` object is then assigned
back to a, while the original contents in a is discarded. This is very inefficient
when we want to do inplace update.

This macro is a simple utility to implement this behavior. Write

```julia
  @mx.inplace a += b
```

will translate into

```julia
  mx.add_to!(a, b)
```

which will do inplace adding of the contents of `b` into `a`.
"""
macro inplace(stmt)
  if stmt.head == :+= || stmt.head == :.+=
    Expr(:call, :add_to!, esc(stmt.args[1]), esc(stmt.args[2]))
  elseif stmt.head == :-= || stmt.head == :.-=
    Expr(:call, :sub_from!, esc(stmt.args[1]), esc(stmt.args[2]))
  elseif stmt.head == :.*=
    Expr(:call, :mul_to!, esc(stmt.args[1]), esc(stmt.args[2]))
  elseif stmt.head == :./=
    Expr(:call, :div_from!, esc(stmt.args[1]), esc(stmt.args[2]))
  else
    error("unsupported inplace translation for $stmt")
  end
end

"""
    add_to!(dst :: NDArray, args :: Union{Real, NDArray}...)

Add a bunch of arguments into `dst`. Inplace updating.
"""
function add_to!(dst :: NDArray, args :: Union{Real, NDArray}...)
  @assert dst.writable
  for arg in args
    if isa(arg, Real)
      _plus_scalar(dst, convert(eltype(dst), arg), dst)
    else
      _plus(dst, arg, dst)
    end
  end
  return dst
end

import Base: +, .+

"""
    +(args...)
    .+(args...)

Summation. Multiple arguments of either scalar or `NDArray` could be
added together. Note at least the first or second argument needs to be an `NDArray` to
avoid ambiguity of built-in summation.
"""
function +(arg0 :: NDArray, args :: Union{Real, NDArray}...)
  ret = copy(arg0, context(arg0))
  add_to!(ret, args...)
end
function .+(arg0 :: NDArray, args :: Union{Real, NDArray}...)
  +(arg0, args...)
end
function +(arg0 :: Real, arg1 :: NDArray, args :: Union{Real, NDArray}...)
  +(arg1, arg0, args...)
end
function .+(arg0 :: Real, arg1 :: NDArray, args :: Union{Real, NDArray}...)
  .+(arg1, arg0, args...)
end

"""
    sub_from!(dst :: NDArray, args :: Union{Real, NDArray}...)

Subtract a bunch of arguments from `dst`. Inplace updating.
"""
function sub_from!(dst :: NDArray, arg :: Union{Real, NDArray})
  @assert dst.writable
  if isa(arg, Real)
    _minus_scalar(dst, convert(eltype(dst), arg), dst)
  else
    _minus(dst, arg, dst)
  end
end

import Base: -, .-

"""
    -(arg0, arg1)
    -(arg0)
    .-(arg0, arg1)

Subtraction `arg0 - arg1`, of scalar types or `NDArray`. Or create
the negative of `arg0`.
"""
function -(arg0 :: NDArray, arg1 :: Union{Real, NDArray})
  ret = copy(arg0, context(arg0))
  sub_from!(ret, arg1)
end
function .-(arg0 :: NDArray, arg1 :: Union{Real, NDArray})
  -(arg0, arg1)
end
function -(arg0 :: Real, arg1 :: NDArray)
  ret = -arg1
  add_to!(ret, arg0)
  return ret
end
function .-(arg0 :: Real, arg1 :: NDArray)
  -(arg0, arg1)
end

function -(arg0 :: NDArray)
  _mul_scalar(arg0, -one(eltype(arg0)))
end

"""
    mul_to!(dst :: NDArray, arg :: Union{Real, NDArray})

Elementwise multiplication into `dst` of either a scalar or an `NDArray` of the same shape.
Inplace updating.
"""
function mul_to!(dst :: NDArray, arg :: Union{Real, NDArray})
  @assert dst.writable
  if isa(arg, Real)
    _mul_scalar(dst, convert(eltype(dst), arg), dst)
  else
    _mul(dst, arg, dst)
  end
  return dst
end

import Base: .*, *

"""
    .*(arg0, arg1)

Elementwise multiplication of `arg0` and `arg`, could be either scalar or `NDArray`.
"""
function .*(arg0 :: NDArray, arg :: Union{Real, NDArray})
  ret = copy(arg0, context(arg0))
  mul_to!(ret, arg)
end
function .*(arg0 :: Real, arg :: NDArray)
  .*(arg, arg0)
end

"""
    *(arg0, arg1)

Currently only multiplication a scalar with an `NDArray` is implemented. Matrix multiplication
is to be added soon.
"""
function *(arg0 :: NDArray, arg :: Real)
  ret = copy(arg0, context(arg0))
  mul_to!(ret, arg)
end
function *(arg0 :: Real, arg :: NDArray)
  *(arg, arg0)
end

"""
    div_from!(dst :: NDArray, arg :: Union{Real, NDArray})

Elementwise divide a scalar or an `NDArray` of the same shape from `dst`. Inplace updating.
"""
function div_from!(dst :: NDArray, arg :: Union{Real, NDArray})
  @assert dst.writable
  if isa(arg, Real)
    _div_scalar(dst, convert(eltype(dst), arg), dst)
  else
    _div(dst, arg, dst)
  end
end

import Base: ./, /
"""
    ./(arg0 :: NDArray, arg :: Union{Real, NDArray})

Elementwise dividing an `NDArray` by a scalar or another `NDArray` of the same shape.
"""
function ./(arg0 :: NDArray, arg :: Union{Real, NDArray})
  ret = copy(arg0, context(arg0))
  div_from!(ret, arg)
end

"""
    /(arg0 :: NDArray, arg :: Real)

Divide an `NDArray` by a scalar. Matrix division (solving linear systems) is not implemented yet.
"""
function /(arg0 :: NDArray, arg :: Real)
  ./(arg0, arg)
end

"""
Manipulating as Julia Arrays
----------------------------

    @nd_as_jl(captures..., statement)

A convenient macro that allows to operate `NDArray` as Julia Arrays. For example,

```julia
  x = mx.zeros(3,4)
  y = mx.ones(3,4)
  z = mx.zeros((3,4), mx.gpu())

  @mx.nd_as_jl ro=(x,y) rw=z begin
    # now x, y, z are just ordinary Julia Arrays
    z[:,1] = y[:,2]
    z[:,2] = 5
  end
```

Under the hood, the macro convert all the declared captures from `NDArray` into Julia
Arrays, by using `try_get_shared`. And automatically commit the modifications back into
the `NDArray` that is declared as `rw`. This is useful for fast prototyping and when
implement non-critical computations, such as `AbstractEvalMetric`.

!!! note
* Multiple `rw` and / or `ro` capture declaration could be made.
* The macro does **not** check to make sure that `ro` captures are not modified. If the
  original `NDArray` lives in CPU memory, then it is very likely the corresponding
  Julia Array shares data with the `NDArray`, so modifying the Julia Array will also
  modify the underlying `NDArray`.
* More importantly, since the `NDArray` is
  asynchronized, we will wait for *writing* for `rw` variables but wait only for *reading*
  in `ro` variables. If we write into those `ro` variables, **and** if the memory is
  shared, racing condition might happen, and the behavior is undefined.
* When an `NDArray` is declared to be captured as `rw`, its contents is always sync
  back in the end.
* The execution results of the expanded macro is always `nothing`.
* The statements are wrapped in a `let`, thus locally introduced new variables will not be
  available after the statements. So you will need to declare the variables before calling the
  macro if needed.
"""
macro nd_as_jl(m_args...)
  @assert(length(m_args) > 0)
  stmts = m_args[end]
  @assert(isa(stmts, Expr) && stmts.head == :block,
          "The last argument should be a statement block (begin-end); but get $stmts")
  stmts = esc(stmts)

  dclrs  = m_args[1:end-1]
  nd_ro  = []
  nd_rw  = []
  nd_all = []
  for declr in dclrs
    @assert(isa(declr, Expr) && declr.head == :(=) && length(declr.args)==2 && declr.args[1] ∈ (:ro,:rw),
            "Invalid declaration, should be rw=(x,y) or ro=z; but get $declr")

    declr_vars = declr.args[2]
    if isa(declr_vars, Symbol)
      declr_vars = (declr_vars,)
    elseif isa(declr_vars, Expr)
      @assert(declr_vars.head ∈ (:tuple, :vect),
              "Capture declaration should be a variable or a tuple of variables; but got $declr_vars")
      declr_vars = declr_vars.args
    else
      @assert(false, "Capture declaration should be a variable or a tuple of variables; but got $declr_vars")
    end
    for declr_var in declr_vars
      @assert(isa(declr_var, Symbol),
              "Captured ndarrays in ro/rw declaration should be variables, but get $(declr_var)")
    end
    append!(nd_all, [declr_vars...])
    if declr.args[1] == :ro
      append!(nd_ro, [declr_vars...])
    else
      append!(nd_rw, [declr_vars...])
    end
  end

  nd_ro    = map(esc, nd_ro)
  nd_rw    = map(esc, nd_rw)
  nd_all   = map(esc, nd_all)
  rw_origs = [gensym() for _ in nd_rw]

  save_statements  = Expr(:block, [:($v_orig = $v) for (v_orig, v) in zip(rw_origs, nd_rw)]...)
  wait_statements  = Expr(:block, [:(_wait_to_read($v)) for v in nd_ro]...,
                                  [:(_wait_to_write($v)) for v in nd_rw]...)
  clear_statements = Expr(:block, [:($v_orig = nothing) for v_orig in rw_origs]...)
  let_assignments  = [:($v = try_get_shared($v)) for v in nd_all]
  sync_statements  = map(rw_origs, nd_rw) do v_orig, v
    quote
      if !is_shared($v, $v_orig)
        # copy data back if not or no longer sharing data
        copy!($v_orig, $v)
      end
    end
  end
  sync_statements  = Expr(:block, sync_statements...)

  let_statement = Expr(:let, quote
    $stmts
    $sync_statements
  end, let_assignments...)
  m_body = quote
    $wait_statements
    $save_statements
    $let_statement
    $clear_statements
    nothing # the final results is always nothing
  end

  m_body
end

# NOTE: internal use only. Accessing pointers on a different device (e.g. accessing GPU
# pointers from CPU) leads to undefined behavior.
import Base.pointer
function pointer(arr :: NDArray)
  pdata = Ref{Ptr{MX_float}}(0)
  @mxcall(:MXNDArrayGetData, (MX_handle, Ref{Ptr{MX_float}}), arr, pdata)
  return convert(Ptr{eltype(arr)}, pdata[])
end
function _wait_to_read(arr :: NDArray)
  @mxcall(:MXNDArrayWaitToRead, (MX_handle,), arr)
end
function _wait_to_write(arr :: NDArray)
  @mxcall(:MXNDArrayWaitToWrite, (MX_handle,), arr)
end

"""
    try_get_shared(arr)

Try to create a Julia array by sharing the data with the underlying `NDArray`.

# Arguments:
* `arr::NDArray`: the array to be shared.

!!! note
    The returned array does not guarantee to share data with the underlying `NDArray`.
    In particular, data sharing is possible only when the `NDArray` lives on CPU.
"""
function try_get_shared(arr :: NDArray)
  if context(arr).device_type == CPU
    # try to do data sharing
    return unsafe_wrap(Array, pointer(arr), size(arr))
  else
    # impossible to share, just copying
    return copy(arr)
  end
end

"""
    is_shared(j_arr, arr)

Test whether `j_arr` is sharing data with `arr`.

# Arguments:
* Array j_arr: the Julia Array.
* NDArray arr: the `NDArray`.
"""
function is_shared(j_arr :: Array, arr :: NDArray)
  false
end
function is_shared{T<:DType}(j_arr :: Array{T}, arr :: NDArray)
  if length(j_arr) != length(arr)
    return false
  end
  if context(arr).device_type != CPU
    return false
  end
  return pointer(j_arr) == pointer(arr)
end

"""
    load(filename, ::Type{NDArray})

Load NDArrays from binary file.

# Arguments:
* `filename::String`: the path of the file to load. It could be S3 or HDFS address.

Returns either `Dict{Symbol, NDArray}` or `Vector{NDArray}`.

`filename` can point to `s3` or `hdfs` resources if the `libmxnet` is built with the
corresponding components enabled. Examples:
* `s3://my-bucket/path/my-s3-ndarray`
* `hdfs://my-bucket/path/my-hdfs-ndarray`
* `/path-to/my-local-ndarray`
"""
function load(filename::AbstractString, ::Type{NDArray})
  out_size      = Ref{MX_uint}(0)
  out_hdrs      = Ref{Ptr{MX_handle}}(0)
  out_name_size = Ref{MX_uint}(0)
  out_names     = Ref{char_pp}(0)
  @mxcall(:MXNDArrayLoad, (char_p, Ref{MX_uint}, Ref{Ptr{MX_handle}}, Ref{MX_uint}, Ref{char_pp}),
          filename, out_size, out_hdrs, out_name_size, out_names)
  out_name_size = out_name_size[]
  out_size      = out_size[]
  if out_name_size == 0
    return [NDArray(MX_NDArrayHandle(hdr)) for hdr in unsafe_wrap(Array, out_hdrs[], out_size)]
  else
    @assert out_size == out_name_size
    return Dict([(Symbol(unsafe_wrap(String, k)), NDArray(MX_NDArrayHandle(hdr))) for (k,hdr) in
                 zip(unsafe_wrap(Array, out_names[], out_size), unsafe_wrap(Array, out_hdrs[], out_size))])
  end
end

"""
    save(filename :: AbstractString, data)

Save NDarrays to binary file. Filename could be S3 or HDFS address, if `libmxnet` is built
with corresponding support (see `load`).

* `filename::String`: path to the binary file to write to.
* `data`: data to save to file. Data can be a`NDArray`, a `Vector{NDArray}`, or a `Dict{Base.Symbol, NDArray}`.
"""
function save(filename::String, data::NDArray)
  save(filename, [data])
end
function save(filename::String, data::Vector{NDArray})
  @mxcall(:MXNDArraySave, (char_p, MX_uint, Ptr{MX_handle}, char_pp),
          filename, length(data), MX_handle[data...], char_pp(0))
end
function save(filename::String, data::Dict{Base.Symbol,NDArray})
  names  = [k for k in keys(data)]
  arrays = MX_handle[data[k] for k in names]
  names  = String[string(k) for k in names]

  @mxcall(:MXNDArraySave, (char_p, MX_uint, Ptr{MX_handle}, char_pp),
          filename, length(names), arrays, names)
end

################################################################################
# NDArray functions dynamically imported from libmxnet
################################################################################
function _invoke_mxfunction(func_handle::MX_handle, use_vars, scalars, mut_vars; kwargs...)
  names = String[string(entry[1]) for entry in kwargs]
  args = String[string(entry[2]) for entry in kwargs]
  @mxcall(:MXFuncInvokeEx,
          (MX_handle, Ptr{MX_handle}, Ptr{MX_float}, Ptr{MX_handle}, Cint, char_pp, char_pp),
          func_handle, use_vars, scalars, mut_vars, length(names), names, args)
end

@enum(LIBMX_FUNC_TYPE_MASK,
  NDARRAY_ARG_BEFORE_SCALAR = 1,
  ACCEPT_EMPTY_MUTATE_TARGET = (1 << 2)
)

# Import corresponding math functions from base so the automatically defined libmxnet
# functions can overload them
import Base: sqrt

"""
The libxmnet APIs are automatically imported from `libmxnet.so`. The functions listed
here operate on `NDArray` objects. The arguments to the functions are typically ordered
as

```julia
  func_name(arg_in1, arg_in2, ..., scalar1, scalar2, ..., arg_out1, arg_out2, ...)
```

unless `NDARRAY_ARG_BEFORE_SCALAR` is not set. In this case, the scalars are put before the input arguments:

```julia
  func_name(scalar1, scalar2, ..., arg_in1, arg_in2, ..., arg_out1, arg_out2, ...)
```

If `ACCEPT_EMPTY_MUTATE_TARGET` is set. An overloaded function without the output arguments will also be defined:

```julia
  func_name(arg_in1, arg_in2, ..., scalar1, scalar2, ...)
```

Upon calling, the output arguments will be automatically initialized with empty NDArrays.

Those functions always return the output arguments. If there is only one output (the typical situation), that
object (`NDArray`) is returned. Otherwise, a tuple containing all the outputs will be returned.
"""

function _get_ndarray_functions()
  n = Ref{MX_uint}(0)
  handles = Ref{Ptr{MX_handle}}(0)

  @mxcall(:MXListFunctions, (Ref{MX_uint}, Ref{Ptr{MX_handle}}), n, handles)

  funcs = unsafe_wrap(Array, handles[], n[])
  return funcs
end

const _function_cache = Dict{Symbol, MX_handle}()
function _get_function(name :: Symbol)
  if !haskey(_function_cache, name)
    handle = Ref{MX_handle}(0)

    @mxcall(:MXGetFunction, (Cstring, Ref{MX_handle}), name, handle)
    _function_cache[name] = handle[]
    return handle[]
  else
    return _function_cache[name]
  end
end

function _get_function_description(handle :: MX_handle)
  # get function information (human readable)
  ref_name = Ref{char_p}(0)
  ref_desc = Ref{char_p}(0)
  ref_narg = Ref{MX_uint}(0)

  ref_arg_names = Ref{char_pp}(0)
  ref_arg_types = Ref{char_pp}(0)
  ref_arg_descs = Ref{char_pp}(0)

  ref_ret_type  = Ref{char_p}(0)

  @mxcall(:MXFuncGetInfo,
         (MX_handle, Ref{char_p}, Ref{char_p}, Ref{MX_uint}, Ref{char_pp},
          Ref{char_pp}, Ref{char_pp}, Ref{char_p}),
          handle, ref_name, ref_desc, ref_narg, ref_arg_names,
          ref_arg_types, ref_arg_descs, ref_ret_type)

  name = Symbol(unsafe_wrap(String, ref_name[]))
  signature = _format_signature(Int(ref_narg[]), ref_arg_names)
  desc = "    " * string(name) * "(" * signature * ")\n\n"
  desc *= unsafe_wrap(String, ref_desc[]) * "\n\n"
  desc *= "# Arguments\n"
  desc *= _format_docstring(Int(ref_narg[]), ref_arg_names, ref_arg_types, ref_arg_descs)
  return name, desc
end

function _get_function_expressions(handle :: MX_handle, name)
  # get function specification
  ref_n_use_vars = Ref{MX_uint}(0)
  ref_n_scalars  = Ref{MX_uint}(0)
  ref_n_mut_vars = Ref{MX_uint}(0)
  ref_type_mask  = Ref{Cint}(0)
  @mxcall(:MXFuncDescribe,
         (MX_handle, Ref{MX_uint}, Ref{MX_uint}, Ref{MX_uint}, Ref{Cint}),
          handle, ref_n_use_vars, ref_n_scalars, ref_n_mut_vars, ref_type_mask)

  n_used_vars   = ref_n_use_vars[]
  n_scalars     = ref_n_scalars[]
  n_mutate_vars = ref_n_mut_vars[]
  type_mask     = ref_type_mask[]
  accept_empty_mutate = (type_mask & convert(Cint,ACCEPT_EMPTY_MUTATE_TARGET)) != 0
  arg_before_scalar   = (type_mask & convert(Cint,NDARRAY_ARG_BEFORE_SCALAR)) != 0

  # general ndarray function
  if arg_before_scalar
    args = vcat([Expr(:(::), Symbol("in$i"), NDArray) for i=1:n_used_vars],
                [Expr(:(::), Symbol("sca$i"), Real) for i=1:n_scalars],
                [Expr(:(::), Symbol("out$i"), NDArray) for i=1:n_mutate_vars])
  else
    args = vcat([Expr(:(::), Symbol("sca$i"), Real) for i=1:n_scalars],
                [Expr(:(::), Symbol("in$i"), NDArray) for i=1:n_used_vars],
                [Expr(:(::), Symbol("out$i"), NDArray) for i=1:n_mutate_vars])
  end

  _use_vars = Expr(:ref, :MX_handle, [Symbol("in$i") for i=1:n_used_vars]...)
  _scalars  = Expr(:ref, :MX_float, [Symbol("sca$i") for i=1:n_scalars]...)
  _mut_vars = Expr(:ref, :MX_handle, [Symbol("out$i") for i=1:n_mutate_vars]...)

  # XXX: hacky way of solving the problem that the arguments of `dot` should be swapped
  # See https://github.com/dmlc/MXNet.jl/issues/55
  if name == :dot
    _use_vars.args[2:end] = flipdim(_use_vars.args[2:end], 1)
  end

  # XXX: hacky way of solving the semantic difference of the axes parameter in Julia
  # and in libmxnet.
  # See https://github.com/dmlc/MXNet.jl/pull/123
  if name == :transpose
    transform = quote
      kwargs = Any[key != :axes ? (key, arg) : (key, reverse(map(i->length(arg)-i, arg))) for (key, arg) in kwargs]
    end
  else
    transform = :()
  end

  stmt_call = quote
    local handle = _get_function($(QuoteNode(name)))
    _invoke_mxfunction(handle, $_use_vars, $_scalars, $_mut_vars; kwargs...)
  end
  if n_mutate_vars == 1
    stmt_ret = :(return out1)
  else
    stmt_ret = Expr(:return, Expr(:tuple, [Symbol("out$i") for i=1:n_mutate_vars]...))
  end

  func_def = quote
    function $name($(args...); kwargs...)
      $transform
      $stmt_call
      $stmt_ret
    end
  end

  if accept_empty_mutate
    args0      = args[1:n_used_vars+n_scalars]
    _mut_vars0 = [:(NDArray(_ndarray_alloc())) for i=1:n_mutate_vars]

    func_def0 = quote
      function $name($(args0...); kwargs...)
        $name($(args0...), $(_mut_vars0...); kwargs...)
      end
    end
    return func_def, func_def0
  else
    return func_def, :()
  end
end

macro _import_ndarray_functions()
  funcs = _get_ndarray_functions()
  func_exprs = Expr[]

  for i = 1:length(funcs)
    handle = funcs[i]

    name, desc = _get_function_description(handle)
    func_def, func_def0 = _get_function_expressions(handle, name)

    expr = quote
      $(isdefined(Base, name) ? :(import Base.$name) : :())
      @doc $desc ->
      $func_def
      $func_def0
    end

    push!(func_exprs, expr)
  end

  esc(quote
    $(func_exprs...)
  end)
end

@_import_ndarray_functions()


###
### Bitparallel
###

"""
A generator for efficient bitwise sequence operations.
"""
function compile_bitpar(funcname::Symbol;
                        arguments::Tuple  =  (),
                        init_code::Expr   = :(),
                        head_code::Expr   = :(),
                        body_code::Expr   = :(),
                        tail_code::Expr   = :(),
                        return_code::Expr = :())
    functioncode = :(function $(funcname)() end)
    for arg in arguments
        push!(functioncode.args[1].args, arg)
    end
    functioncode.args[2] = quote
        $(init_code)
        ind = bitindex(seq, 1)
        stop = bitindex(seq, lastindex(seq) + 1)
        data = seq.data
        @inbounds begin
            if !iszero(offset(ind)) & (ind < stop)
                # align the bit index to the beginning of a block boundary
                o = offset(ind)
                chunk = (data[index(ind)] >> o) & bitmask(stop - ind)
                $(head_code)
                ind += 64 - o
            end

            lastind = index(stop - bits_per_symbol(seq))
            lastind -= !iszero(offset(stop))
            for i in index(ind):lastind
                chunk = data[i]
                $(body_code)
                ind += 64
            end

            if ind < stop
                chunk = data[index(ind)] & bitmask(offset(stop))
                $(tail_code)
            end
        end
        $(return_code)
    end
    return functioncode
end

function get_arg_name(arg)
    if typeof(arg) === Expr
        if arg.head === :(::)
            return first(arg.args)
        end
    end
    return arg
end

function check_arguments(args::Tuple)
    # Check all arguments are symbols or expressions
    for arg in args
        if !(isa(arg, Symbol) || isa(arg, Expr))
            error("Argument ", arg, " is not an expression or symbol")
        end
    end

    a1 = first(args)
    if isa(a1, Symbol)
        @assert a1 == :seqa
    elseif isa(a1, Expr)
        @assert a1.head === :(::)
        @assert first(a1.args) === :seqa
    end

    a2 = args[2]
    if isa(a2, Symbol)
        @assert a2 == :seqb
    elseif isa(a2, Expr)
        @assert a2.head === :(::)
        @assert first(a2.args) === :seqb
    end
end

function add_arguments!(exp, args)
    check_arguments(args)
    @assert exp.head === :function
    if exp.args[1].head === :call
        for arg in args
            push!(exp.args[1].args, arg)
        end
    elseif exp.args[1].head === :where
        @assert exp.args[1].args[1].head === :call
        for arg in args
            push!(exp.args[1].args[1].args, arg)
        end
    else
        error("Expression in function is not a :call or :where!")
    end
end

function compile_2seq_bitpar(funcname::Symbol;
                             arguments::Tuple  =  (),
                             parameters::Tuple =  (),
                             init_code::Expr   = :(),
                             head_code::Expr   = :(),
                             body_code::Expr   = :(),
                             tail_code::Expr   = :(),
                             return_code::Expr = :())

    # TODO: Check the parameters provided.

    if isempty(parameters)
        functioncode = :(function $(funcname)() end)
    else
        functioncode = :(function $(funcname)() where {$(parameters...)} end)
    end

    add_arguments!(functioncode, arguments)

    argument_names = map(get_arg_name, arguments)
    argument_names = tuple(argument_names[2], argument_names[1], argument_names[3:end]...)

    functioncode.args[2] = quote
        if length(seqa) > length(seqb)
            return $funcname($(argument_names...))
        end
        @assert length(seqa) ≤ length(seqb)

        nexta = bitindex(seqa, 1)
        stopa = bitindex(seqa, lastindex(seqa) + 1)
        nextb = bitindex(seqb, 1)
        stopb = bitindex(seqb, lastindex(seqb) + 1)
        adata = encoded_data(seqa)
        bdata = encoded_data(seqb)

        $(init_code)

        # The first thing we need to sort out is to correctly align the head of
        # sequence / subsequence `a`s data is aligned such that the offset of
        # `nexta` is essentially reduced to 0.
        # With sequence / subsequence `a` aligned, from there, we only need to
        # worry about the alignment of sequence / subsequence `b` with respect
        # to `a`.
        if nexta < stopa && offset(nexta) != 0
            # Here we shift the first data chunks to the right so as the first
            # nucleotide of the seq/subseq is the first nibble / pair of bits.
            x = adata[index(nexta)] >> offset(nexta)
            y = bdata[index(nextb)] >> offset(nextb)
            # Here it was assumed that there is something to go and get from
            # the next chunk of `b`, yet that may not be true.
            # We know that if this is not true of `b`, then it is certainly not
            # true of `a`.
            # We check if the end of the sequence is not contained in the same
            # integer like so: `64 - offset(nextb) < stopb - nextb`.
            #
            # This edge case was found and accounted for by Ben J. Ward @BenJWard.
            # Ask this maintainer for more information.
            if offset(nextb) > offset(nexta) && 64 - offset(nextb) < stopb - nextb
                y |= bdata[index(nextb) + 1] << (64 - offset(nextb))
            end
            # Here we need to check something, we need to check if the
            # integer of `a` we are currently aligning contains the end of
            # seq/subseq `a`. Because if so it's something we need to take into
            # account of when we mask x and y.
            #
            # In other words if `64 - offset(nexta) > stopa - nexta, we know
            # seq or subseq a's data ends before the end of this data chunk,
            # and so the mask used needs to be defined to account for this:
            # `mask(stopa - nexta)`, otherwise the mask simply needs to be
            # `mask(64 - offset(nexta))`.
            #
            # This edge case was found and accounted for by Ben Ward @Ward9250.
            # Ask this maintainer for more information.
            offs = ifelse(64 - offset(nexta) > stopa - nexta, stopa - nexta, 64 - offset(nexta))
            m = bitmask(offs)

            x &= m
            y &= m

            $(head_code)

            # Here we move our current position markers by k, meaning they move
            # to either, A). The next integer, or B). The end of the sequence if
            # it is in the current integer.
            nexta += offs
            nextb += offs
        end

        if offset(nextb) == 0  # data are aligned with each other
            while stopa - nexta ≥ 64 # Iterate through body of data
                x = adata[index(nexta)]
                y = bdata[index(nextb)]

                $(body_code)

                nexta += 64
                nextb += 64
            end

            if nexta < stopa
                x = adata[index(nexta)]
                y = bdata[index(nextb)]

                offs = stopa - nexta
                m = bitmask(offs)
                x &= m
                y &= m

                $(tail_code)
            end
        elseif nexta < stopa # Data are unaligned
            y = bdata[index(nextb)]
            nextb += 64
            # Note that here, updating `nextb` by 64, increases the chunk index,
            # but the `offset(nextb)` will remain the same.
            while stopa - nexta ≥ 64 # processing body of data
                x = adata[index(nexta)]
                z = bdata[index(nextb)]
                y = y >> offset(nextb) | z << (64 - offset(nextb))

                $(body_code)

                y = z
                nexta += 64
                nextb += 64
            end

            if nexta < stopa # processing tail of data
                x = adata[index(nexta)]
                y = y >> offset(nextb)
                if 64 - offset(nextb) < stopa - nexta
                    y |= bdata[index(nextb)] << (64 - offset(nextb))
                end
                offs = stopa - nexta
                m = bitmask(offs)
                x &= m
                y &= m

                $(tail_code)
            end
        end
        $(return_code)
    end
    return functioncode
end


module Parse

using RBNF
using RBNF: Token

using ..Types

struct QASMLang end

second((a, b)) = b
second(vec::V) where {V<:AbstractArray} = vec[2]

# roses are red
# violets are blue
# pirates are good
RBNF.crate(::Type{Symbol}) = gensym(:qasm)
RBNF.crate(::Type{VersionNumber}) = VersionNumber("0.0.0")

Base.convert(::Type{VersionNumber}, t::Token) = VersionNumber(t.str)
Base.convert(::Type{String}, t::Token) = t.str
Base.convert(::Type{Int}, t::Token{:int}) = Base.parse(Int, t.str)
Base.convert(::Type{Float64}, t::Token{:float64}) = Base.parse(Float64, t.str)
Base.convert(::Type{Symbol}, t::Token{:id}) = Symbol(t.str)
Base.convert(::Type{Symbol}, t::Token{:reserved}) = Symbol(t.str)
Base.convert(::Type{String}, t::Token{:str}) = String(t.str[2:end-1])

RBNF.@parser QASMLang begin
    # define ignorances
    ignore{space, comment}

    @grammar
    # define grammars
    mainprogram::MainProgram := ["OPENQASM", version = float64, ';', prog = program]
    program = statement{*}
    statement = (regdecl | gate | opaque | qop | ifstmt | barrier | inc)
    # stmts
    ifstmt::IfStmt := [:if, '(', left = id, :(==), right = int, ')', body = qop]
    opaque::Opaque := [:opaque, name = id, ['(', [cargs = idlist].?, ')'].?, qargs = idlist, ';']
    barrier::Barrier := [:barrier, qargs = bitlist, ';']
    regdecl::RegDecl := [type = :qreg | :creg, name = id, '[', size = int, ']', ';']
    inc::Include := [:include, file = str, ';']
    # gate
    gate::Gate := [decl = gatedecl, [body = goplist].?, '}']
    gatedecl::GateDecl := [:gate, name = id, ['(', [cargs = idlist].?, ')'].?, qargs = idlist, '{']

    goplist = (uop | barrier){*}

    # qop
    qop = (uop | measure | reset)
    reset::Reset := [:reset, qarg = bit, ';']
    measure::Measure := [:measure, qarg = bit, :(->), carg = bit, ';']

    uop = (inst | ugate | csemantic_gate)
    inst::Instruction := [name = id, ['(', [cargs = explist].?, ')'].?, qargs = bitlist, ';']
    ugate::UGate := [:U, '(', z1 = exp, ',', y = exp, ',', z2 = exp, ')', qarg = bit, ';']
    csemantic_gate::CXGate := [:CX, ctrl = bit, ',', qarg = bit, ';']

    idlist = @direct_recur begin
        init = [id]
        prefix = [recur..., (',', id) % second]
    end

    bit::Bit := [name = id, ['[', address = int, ']'].?]
    bitlist = @direct_recur begin
        init = [bit]
        prefix = [recur..., (',', bit) % second]
    end

    explist = @direct_recur begin
        init = [exp]
        prefix = [recur..., (',', exp) % second]
    end

    con = (float64 | int | :pi | id | call)
    num = (['(', exp, ')'] % second) | neg | con
    add = ( :+ | :- )
    mul = ( :* | :/ )
    exp = @direct_recur begin
        init = term
        prefix = (recur, add, term)
    end
    term = @direct_recur begin
        init = num
        prefix = (recur, mul, term)
    end
    # term = (add | sub | num)
    neg::Neg := [:-, val = num]
    call::Call := [name=fn, "(", args = exp, ")"]
    fn = (:sin | :cos | :tan | :exp | :ln | :sqrt)
    # binop = (:+ | :- | :* | :/)

    # define tokens
    @token
    id := r"\G[a-z]{1}[A-Za-z0-9_]*"
    float64 := r"\G([0-9]+\.[0-9]*|[0-9]*\.[0.9]+)([eE][-+]?[0-9]+)?"
    int := r"\G([1-9]+[0-9]*|0)"
    space := r"\G\s+"
    comment := r"\G//.*"
    str := @quote ("\"", "\\\"", "\"")
end

end # Parse


function save_mesh(::Val{:FLUENT}, io::IO, m::AirfoilMesh)
    NI, NJ = size(m.grid)
    NV = NI * NJ - m.ncon
    NF = NI * (NJ - 1) + NJ * (NI - 1) - max(m.ncon - 1, 0)
    NC = (NI - 1) * (NJ - 1)

    NLI = Matrix(LinearIndices((NI, NJ)))
    NCI = Dict{Int,CartesianIndex{2}}()
    # try do remove duplicated
    if m.ncon > 0
        if m.topo == :O
            @assert m.ncon == NJ
            NLI[:, NJ] .= NLI[:, 1]
        elseif m.topo == :C
            NLI[NI:-1:(NI - m.ncon + 1), 1] .= NLI[1:(m.ncon), 1]
        end
        old_indicies = sort(collect(Set(NLI)))
        @assert length(old_indicies) == NV
        o2n = Dict{Int,Int}()
        for i in eachindex(old_indicies)
            o2n[old_indicies[i]] = i
        end
        for i in eachindex(NLI)
            NLI[i] = o2n[NLI[i]]
        end
    end
    for j in 1:NJ
        for i in 1:NI
            li = NLI[i, j]
            if !haskey(NCI, li)
                NCI[li] = CartesianIndex(i, j)
            end
        end
    end

    CLI = LinearIndices((NI - 1, NJ - 1))
    f_int = Vector{NTuple{4,Int}}()
    f_wall = Vector{NTuple{4,Int}}()
    f_far = Vector{NTuple{4,Int}}()
    f_min = Vector{NTuple{4,Int}}()
    f_max = Vector{NTuple{4,Int}}()
    # INTERIOR
    for j in 1:(NJ - 1)
        for i in 2:(NI - 1)
            p0 = NLI[i, j]
            p1 = NLI[i, j + 1]
            c0 = CLI[i - 1, j]
            c1 = CLI[i, j]
            push!(f_int, (p0, p1, c0, c1))
        end
    end
    for j in 2:(NJ - 1)
        for i in 1:(NI - 1)
            p0 = NLI[i, j]
            p1 = NLI[i + 1, j]
            c0 = CLI[i, j]
            c1 = CLI[i, j - 1]
            push!(f_int, (p0, p1, c0, c1))
        end
    end
    # J = 1
    if m.topo == :C
        for i in 1:(m.ncon - 1)
            p0 = NLI[i, 1]
            p1 = NLI[i + 1, 1]
            c0 = CLI[i, 1]
            c1 = CLI[NI - i, 1]
            push!(f_int, (p0, p1, c0, c1))
        end
        for i in (m.ncon):(NI - m.ncon)
            p0 = NLI[i, 1]
            p1 = NLI[i + 1, 1]
            c0 = CLI[i, 1]
            push!(f_wall, (p0, p1, c0, 0))
        end
    else
        for i in 1:(NI - 1)
            p0 = NLI[i, 1]
            p1 = NLI[i + 1, 1]
            c0 = CLI[i, 1]
            push!(f_wall, (p0, p1, c0, 0))
        end
    end
    # J = NJ
    for i in 1:(NI - 1)
        p0 = NLI[i + 1, NJ] # reversed at J=NJ
        p1 = NLI[i, NJ]
        c0 = CLI[i, NJ - 1]
        c1 = 0
        push!(f_far, (p0, p1, c0, c1))
    end
    # I = 1
    for j in 1:(NJ - 1)
        p0 = NLI[1, j + 1] # reversed
        p1 = NLI[1, j]
        c0 = CLI[1, j]
        if m.topo == :O
            c1 = CLI[NI - 1, j]
            push!(f_int, (p0, p1, c0, c1))
        else
            push!(f_min, (p0, p1, c0, 0))
        end
    end
    # I = NI
    if m.topo != :O
        for j in 1:(NJ - 1)
            p0 = NLI[NI, j]
            p1 = NLI[NI, j + 1]
            c0 = CLI[NI - 1, j]
            push!(f_max, (p0, p1, c0, 0))
        end
    end
    NFINT = length(f_int)
    I_INT = (1, NFINT)
    NFWAL = length(f_wall)
    I_WAL = (I_INT[2] + 1, I_INT[2] + NFWAL)
    NFFAR = length(f_far)
    I_FAR = (I_WAL[2] + 1, I_WAL[2] + NFFAR)
    NFMIN = length(f_min)
    I_MIN = (I_FAR[2] + 1, I_FAR[2] + NFMIN)
    NFAMX = length(f_max)
    I_MAX = (I_MIN[2] + 1, I_MIN[2] + NFAMX)
    @assert I_MAX[2] == NF
    # informations
    _fluent_comment(io, "Exported from AirfoilMeshers.jl")
    _fluent_comment(io, "Dim = 2")
    println(io, "(2 2)")
    _fluent_comment(io, "# of nodes: $(NV)")
    println(io, "(10 (0 1 $(string(NV, base=16)) 0 2))")
    _fluent_comment(io, "# of cells: $(NC)")
    println(io, "(12 (0 1 $(string(NC, base=16)) 0))")
    _fluent_comment(io, "# of faces: $(NF)")
    println(io, "(13 (0 1 $(string(NF, base=16)) 0))")
    # Zone 1: Nodes
    izone = 1
    _fluent_comment(io, "Zone $izone: nodes [$NV]")
    println(io, "(10 ($izone 1 $(string(NV, base=16)) 1 2)(")
    for li in 1:NV
        p = m.grid[CartesianIndex(NCI[li])]
        @printf io "%24.16E %24.16E\n" p[1] p[2]
    end
    println(io, "))")
    # Zone 2: Cells
    izone += 1
    _fluent_comment(io, "Zone $izone: cells [$NC]")
    println(io, "(12 ($izone 1 $(string(NC, base=16)) 1 3))")
    println(io, "(45 ($izone fluid flowfield)())")
    # Zone 3: Interior
    izone += 1
    _fluent_comment(io, "Zone $izone: interior faces $(I_INT[1]):$(I_INT[2])")
    println(
        io, "(13 ($izone $(string(I_INT[1], base=16)) $(string(I_INT[2], base=16)) 2 2)("
    )
    for f in f_int
        println(io, join(string.(f, base=16), ' '))
    end
    println(io, "))")
    println(io, "(45 ($izone interior interior-flow)())")
    # Zone 4: Wall
    izone += 1
    _fluent_comment(io, "Zone $izone: wall faces $(I_WAL[1]):$(I_WAL[2])")
    println(
        io, "(13 ($izone $(string(I_WAL[1], base=16)) $(string(I_WAL[2], base=16)) 3 2)("
    )
    for f in f_wall
        println(io, join(string.(f, base=16), ' '))
    end
    println(io, "))")
    println(io, "(45 ($izone wall wall)())")
    # Zone 5: Far
    izone += 1
    _fluent_comment(io, "Zone $izone: far faces $(I_FAR[1]):$(I_FAR[2])")
    println(
        io, "(13 ($izone $(string(I_FAR[1], base=16)) $(string(I_FAR[2], base=16)) a 2)("
    )
    for f in f_far
        println(io, join(string.(f, base=16), ' '))
    end
    println(io, "))")
    println(io, "(45 ($izone velocity-inlet far)())")
    if NFMIN > 0
        # Zone 6: IMin
        izone += 1
        _fluent_comment(io, "Zone $izone: imin faces $(I_MIN[1]):$(I_MIN[2])")
        println(
            io,
            "(13 ($izone $(string(I_MIN[1], base=16)) $(string(I_MIN[2], base=16)) 5 2)(",
        )
        for f in f_min
            println(io, join(string.(f, base=16), ' '))
        end
        println(io, "))")
        println(io, "(45 ($izone pressure-outlet imin)())")
        # Zone 7: IMax
        izone += 1
        _fluent_comment(io, "Zone $izone: imax faces $(I_MAX[1]):$(I_MAX[2])")
        println(
            io,
            "(13 ($izone $(string(I_MAX[1], base=16)) $(string(I_MAX[2], base=16)) 5 2)(",
        )
        for f in f_max
            println(io, join(string.(f, base=16), ' '))
        end
        println(io, "))")
        println(io, "(45 ($izone pressure-outlet imax)())")
    end

    return nothing
end

function _fluent_comment(io::IO, msg::AbstractString, tag::Int=0)
    return println(io, "($tag \"$msg\")")
end


# Contains functions related to distribution networks but not specific to a particular model

## Lookup functions, to build the constraint selection logic

"Return whether the `f_bus` of branch `i` is the reference bus."
function is_frb_branch(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    return haskey(_PM.ref(pm, nw, :frb_branch), i)
end

"Return whether the `f_bus` of ne_branch `i` is the reference bus."
function is_frb_ne_branch(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    return haskey(_PM.ref(pm, nw, :frb_ne_branch), i)
end

"Return whether branch `i` is an OLTC."
function is_oltc_branch(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    return haskey(_PM.ref(pm, nw, :oltc_branch), i)
end

"Return whether ne_branch `i` is an OLTC."
function is_oltc_ne_branch(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    return haskey(_PM.ref(pm, nw, :oltc_ne_branch), i)
end

## Variables

function variable_oltc_branch_transform(pm::_PM.AbstractWModels; kwargs...)
    variable_oltc_branch_transform_magnitude_sqr_inv(pm; kwargs...)
end

"variable: `0 <= ttmi[l]` for `l` in `oltc_branch`es"
function variable_oltc_branch_transform_magnitude_sqr_inv(pm::_PM.AbstractPowerModel; nw::Int=_PM.nw_id_default, bounded::Bool=true, report::Bool=true)
    ttmi = _PM.var(pm, nw)[:ttmi] = JuMP.@variable(pm.model,
        [i in _PM.ids(pm, nw, :oltc_branch)], base_name="$(nw)_ttmi",
        lower_bound = 0.0,
        start = 1.0 / _PM.ref(pm,nw,:oltc_branch,i,"tap")^2
    )

    if bounded
        for (i, br) in _PM.ref(pm, nw, :oltc_branch)
            JuMP.set_lower_bound(ttmi[i], 1.0 / br["tm_max"]^2 )
            JuMP.set_upper_bound(ttmi[i], 1.0 / br["tm_min"]^2 )
        end
    end

    report && _PM.sol_component_value(pm, nw, :branch, :ttmi, _PM.ids(pm, nw, :oltc_branch), ttmi)
end

function variable_oltc_ne_branch_transform(pm::_PM.AbstractWModels; kwargs...)
    variable_oltc_ne_branch_transform_magnitude_sqr_inv(pm; kwargs...)
end

"variable: `0 <= ttmi_ne[l]` for `l` in `oltc_ne_branch`es"
function variable_oltc_ne_branch_transform_magnitude_sqr_inv(pm::_PM.AbstractPowerModel; nw::Int=_PM.nw_id_default, bounded::Bool=true, report::Bool=true)
    ttmi_ne = _PM.var(pm, nw)[:ttmi_ne] = JuMP.@variable(pm.model,
        [i in _PM.ids(pm, nw, :oltc_ne_branch)], base_name="$(nw)_ttmi_ne",
        lower_bound = 0.0,
        start = 1.0 / _PM.ref(pm,nw,:oltc_ne_branch,i,"tap")^2
    )

    if bounded
        for (i, br) in _PM.ref(pm, nw, :oltc_ne_branch)
            JuMP.set_lower_bound(ttmi_ne[i], 1.0 / br["tm_max"]^2 )
            JuMP.set_upper_bound(ttmi_ne[i], 1.0 / br["tm_min"]^2 )
        end
    end

    report && _PM.sol_component_value(pm, nw, :ne_branch, :ttmi, _PM.ids(pm, nw, :oltc_ne_branch), ttmi_ne)
end

## Constraint templates that group several other constraint templates, provided for convenience

function constraint_dist_branch_tnep(pm::_PM.AbstractBFModel, i::Int; nw::Int=_PM.nw_id_default)
    if isempty(ne_branch_ids(pm, i; nw = nw))
        if is_frb_branch(pm, i; nw = nw)
            if is_oltc_branch(pm, i; nw = nw)
                constraint_power_losses_oltc(pm, i; nw = nw)
                constraint_voltage_magnitude_difference_oltc(pm, i; nw = nw)
            else
                constraint_power_losses_frb(pm, i; nw = nw)
                constraint_voltage_magnitude_difference_frb(pm, i; nw = nw)
            end
        else
            _PM.constraint_power_losses(pm, i; nw = nw)
            _PM.constraint_voltage_magnitude_difference(pm, i; nw = nw)
        end
        _PM.constraint_voltage_angle_difference(pm, i; nw = nw)
        _PM.constraint_thermal_limit_from(pm, i; nw = nw)
        _PM.constraint_thermal_limit_to(pm, i; nw = nw)
    else
        expression_branch_indicator(pm, i; nw = nw)
        # constraint_branch_complementarity(pm, i; nw = nw) is best added once per year; if added here, redundant constraints would be generated
        if is_frb_branch(pm, i; nw = nw)
            if is_oltc_branch(pm, i; nw = nw)
                constraint_power_losses_oltc_on_off(pm, i; nw = nw)
                constraint_voltage_magnitude_difference_oltc_on_off(pm, i; nw = nw)
            else
                constraint_power_losses_frb_on_off(pm, i; nw = nw)
                constraint_voltage_magnitude_difference_frb_on_off(pm, i; nw = nw)
            end
        else
            constraint_power_losses_on_off(pm, i; nw = nw)
            constraint_voltage_magnitude_difference_on_off(pm, i; nw = nw)
        end
        _PM.constraint_voltage_angle_difference_on_off(pm, i; nw = nw)
        _PM.constraint_thermal_limit_from_on_off(pm, i; nw = nw)
        _PM.constraint_thermal_limit_to_on_off(pm, i; nw = nw)
    end
end

function constraint_dist_ne_branch_tnep(pm::_PM.AbstractBFModel, i::Int; nw::Int=_PM.nw_id_default)
    if ne_branch_replace(pm, i, nw = nw)
        if is_frb_ne_branch(pm, i, nw = nw)
            if is_oltc_ne_branch(pm, i, nw = nw)
                constraint_ne_power_losses_oltc(pm, i, nw = nw)
                constraint_ne_voltage_magnitude_difference_oltc(pm, i, nw = nw)
            else
                constraint_ne_power_losses_frb(pm, i, nw = nw)
                constraint_ne_voltage_magnitude_difference_frb(pm, i, nw = nw)
            end
        else
            constraint_ne_power_losses(pm, i, nw = nw)
            constraint_ne_voltage_magnitude_difference(pm, i, nw = nw)
        end
        _PM.constraint_ne_thermal_limit_from(pm, i, nw = nw)
        _PM.constraint_ne_thermal_limit_to(pm, i, nw = nw)
    else
        if is_frb_ne_branch(pm, i, nw = nw)
            if is_oltc_ne_branch(pm, i, nw = nw)
                Memento.error(_LOGGER, "addition of a candidate OLTC in parallel to an existing OLTC is not supported")
            else
                constraint_ne_power_losses_frb_parallel(pm, i, nw = nw)
                constraint_ne_voltage_magnitude_difference_frb_parallel(pm, i, nw = nw)
            end
        else
            constraint_ne_power_losses_parallel(pm, i, nw = nw)
            constraint_ne_voltage_magnitude_difference_parallel(pm, i, nw = nw)
        end
        constraint_ne_thermal_limit_from_parallel(pm, i, nw = nw)
        constraint_ne_thermal_limit_to_parallel(pm, i, nw = nw)
    end
    _PM.constraint_ne_voltage_angle_difference(pm, i, nw = nw)
end

## Constraint templates

"Defines voltage drop over a a branch whose `f_bus` is the reference bus"
function constraint_voltage_magnitude_difference_frb(pm::_PM.AbstractBFModel, i::Int; nw::Int=_PM.nw_id_default)
    branch = _PM.ref(pm, nw, :branch, i)
    f_bus = branch["f_bus"]
    t_bus = branch["t_bus"]
    f_idx = (i, f_bus, t_bus)
    t_idx = (i, t_bus, f_bus)

    r = branch["br_r"]
    x = branch["br_x"]
    g_sh_fr = branch["g_fr"]
    b_sh_fr = branch["b_fr"]
    tm = branch["tap"]

    constraint_voltage_magnitude_difference_frb(pm, nw, i, f_bus, t_bus, f_idx, t_idx, r, x, g_sh_fr, b_sh_fr, tm)
end

"Defines voltage drop over a transformer branch that has an OLTC"
function constraint_voltage_magnitude_difference_oltc(pm::_PM.AbstractBFModel, i::Int; nw::Int=_PM.nw_id_default)
    branch = _PM.ref(pm, nw, :branch, i)
    f_bus = branch["f_bus"]
    t_bus = branch["t_bus"]
    f_idx = (i, f_bus, t_bus)
    t_idx = (i, t_bus, f_bus)

    r = branch["br_r"]
    x = branch["br_x"]
    g_sh_fr = branch["g_fr"]
    b_sh_fr = branch["b_fr"]

    constraint_voltage_magnitude_difference_oltc(pm, nw, i, f_bus, t_bus, f_idx, t_idx, r, x, g_sh_fr, b_sh_fr)
end

"Defines branch flow model power flow equations for a branch whose `f_bus` is the reference bus"
function constraint_power_losses_frb(pm::_PM.AbstractBFModel, i::Int; nw::Int=_PM.nw_id_default)
    branch = _PM.ref(pm, nw, :branch, i)
    f_bus = branch["f_bus"]
    t_bus = branch["t_bus"]
    f_idx = (i, f_bus, t_bus)
    t_idx = (i, t_bus, f_bus)

    r = branch["br_r"]
    x = branch["br_x"]
    tm = branch["tap"]
    g_sh_fr = branch["g_fr"]
    g_sh_to = branch["g_to"]
    b_sh_fr = branch["b_fr"]
    b_sh_to = branch["b_to"]

    constraint_power_losses_frb(pm, nw, i, f_bus, t_bus, f_idx, t_idx, r, x, g_sh_fr, g_sh_to, b_sh_fr, b_sh_to, tm)
end

"Defines branch flow model power flow equations for a transformer branch that has an OLTC"
function constraint_power_losses_oltc(pm::_PM.AbstractBFModel, i::Int; nw::Int=_PM.nw_id_default)
    branch = _PM.ref(pm, nw, :branch, i)
    f_bus = branch["f_bus"]
    t_bus = branch["t_bus"]
    f_idx = (i, f_bus, t_bus)
    t_idx = (i, t_bus, f_bus)

    r = branch["br_r"]
    x = branch["br_x"]
    g_sh_fr = branch["g_fr"]
    g_sh_to = branch["g_to"]
    b_sh_fr = branch["b_fr"]
    b_sh_to = branch["b_to"]

    constraint_power_losses_oltc(pm, nw, i, f_bus, t_bus, f_idx, t_idx, r, x, g_sh_fr, g_sh_to, b_sh_fr, b_sh_to)
end

""
function constraint_ne_power_losses_frb(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    branch = _PM.ref(pm, nw, :ne_branch, i)
    f_bus = branch["f_bus"]
    t_bus = branch["t_bus"]
    f_idx = (i, f_bus, t_bus)
    t_idx = (i, t_bus, f_bus)

    r = branch["br_r"]
    x = branch["br_x"]
    tm = branch["tap"]
    g_sh_fr = branch["g_fr"]
    g_sh_to = branch["g_to"]
    b_sh_fr = branch["b_fr"]
    b_sh_to = branch["b_to"]

    vad_min = _PM.ref(pm, nw, :off_angmin)
    vad_max = _PM.ref(pm, nw, :off_angmax)

    constraint_ne_power_losses_frb(pm, nw, i, f_bus, t_bus, f_idx, t_idx, r, x, g_sh_fr, g_sh_to, b_sh_fr, b_sh_to, tm, vad_min, vad_max)
end

""
function constraint_ne_power_losses_oltc(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    branch = _PM.ref(pm, nw, :ne_branch, i)
    f_bus = branch["f_bus"]
    t_bus = branch["t_bus"]
    f_idx = (i, f_bus, t_bus)
    t_idx = (i, t_bus, f_bus)

    r = branch["br_r"]
    x = branch["br_x"]
    g_sh_fr = branch["g_fr"]
    g_sh_to = branch["g_to"]
    b_sh_fr = branch["b_fr"]
    b_sh_to = branch["b_to"]

    vad_min = _PM.ref(pm, nw, :off_angmin)
    vad_max = _PM.ref(pm, nw, :off_angmax)

    constraint_ne_power_losses_oltc(pm, nw, i, f_bus, t_bus, f_idx, t_idx, r, x, g_sh_fr, g_sh_to, b_sh_fr, b_sh_to, vad_min, vad_max)
end

"""
Defines voltage drop over a branch, linking from and to side voltage magnitude
"""
function constraint_ne_voltage_magnitude_difference_frb(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    branch = _PM.ref(pm, nw, :ne_branch, i)
    f_bus = branch["f_bus"]
    t_bus = branch["t_bus"]
    f_idx = (i, f_bus, t_bus)
    t_idx = (i, t_bus, f_bus)

    r = branch["br_r"]
    x = branch["br_x"]
    g_sh_fr = branch["g_fr"]
    b_sh_fr = branch["b_fr"]
    tm = branch["tap"]

    constraint_ne_voltage_magnitude_difference_frb(pm, nw, i, f_bus, t_bus, f_idx, t_idx, r, x, g_sh_fr, b_sh_fr, tm)
end

"""
Defines voltage drop over a branch, linking from and to side voltage magnitude
"""
function constraint_ne_voltage_magnitude_difference_oltc(pm::_PM.AbstractPowerModel, i::Int; nw::Int=_PM.nw_id_default)
    branch = _PM.ref(pm, nw, :ne_branch, i)
    f_bus = branch["f_bus"]
    t_bus = branch["t_bus"]
    f_idx = (i, f_bus, t_bus)
    t_idx = (i, t_bus, f_bus)

    r = branch["br_r"]
    x = branch["br_x"]
    g_sh_fr = branch["g_fr"]
    b_sh_fr = branch["b_fr"]
    tm_min = branch["tm_min"]
    tm_max = branch["tm_max"]

    constraint_ne_voltage_magnitude_difference_oltc(pm, nw, i, f_bus, t_bus, f_idx, t_idx, r, x, g_sh_fr, b_sh_fr, tm_min, tm_max)
end

## Constraint implementations not limited to a specific model type

""
function constraint_ne_voltage_magnitude_difference_frb(pm::_PM.AbstractBFAModel, n::Int, i, f_bus, t_bus, f_idx, t_idx, r, x, g_sh_fr, b_sh_fr, tm)
    branch = _PM.ref(pm, n, :ne_branch, i)
    fr_bus = _PM.ref(pm, n, :bus, f_bus)
    to_bus = _PM.ref(pm, n, :bus, t_bus)
    M_hi   =  1.0^2/tm^2 - to_bus["vmin"]^2
    M_lo   = -1.0^2/tm^2 + to_bus["vmax"]^2
    p_fr = _PM.var(pm, n, :p_ne, f_idx)
    q_fr = _PM.var(pm, n, :q_ne, f_idx)
    w_to = _PM.var(pm, n, :w, t_bus)
    z    = _PM.var(pm, n, :branch_ne, i)
    # w_fr is assumed equal to 1.0

    JuMP.@constraint(pm.model, (1.0/tm^2) - w_to <= 2*(r*p_fr + x*q_fr) + M_hi*(1-z) )
    JuMP.@constraint(pm.model, (1.0/tm^2) - w_to >= 2*(r*p_fr + x*q_fr) - M_lo*(1-z) )
end

""
function constraint_ne_voltage_magnitude_difference_oltc(pm::_PM.AbstractBFAModel, n::Int, i, f_bus, t_bus, f_idx, t_idx, r, x, g_sh_fr, b_sh_fr, tm_min, tm_max)
    branch = _PM.ref(pm, n, :ne_branch, i)
    fr_bus = _PM.ref(pm, n, :bus, f_bus)
    to_bus = _PM.ref(pm, n, :bus, t_bus)
    M_hi   =  1.0^2/tm_min^2 - to_bus["vmin"]^2
    M_lo   = -1.0^2/tm_max^2 + to_bus["vmax"]^2
    p_fr = _PM.var(pm, n, :p_ne, f_idx)
    q_fr = _PM.var(pm, n, :q_ne, f_idx)
    ttmi = _PM.var(pm, n, :ttmi_ne, i)
    w_to = _PM.var(pm, n, :w, t_bus)
    z    = _PM.var(pm, n, :branch_ne, i)
    # w_fr is assumed equal to 1.0 to preserve the linearity of the model

    JuMP.@constraint(pm.model, 1.0*ttmi - w_to <= 2*(r*p_fr + x*q_fr) + M_hi*(1-z) )
    JuMP.@constraint(pm.model, 1.0*ttmi - w_to >= 2*(r*p_fr + x*q_fr) - M_lo*(1-z) )
end


function initialize_pss!(
    device_states,
    device::PSY.DynamicGenerator{M, S, A, TG, PSY.PSSFixed},
) where {M <: PSY.Machine, S <: PSY.Shaft, A <: PSY.AVR, TG <: PSY.TurbineGov}

end

#Currently not working properly.
#=
function initialize_pss!(
    device_states,
    device::PSY.DynamicGenerator{M, S, A, TG, PSY.PSSSimple},
) where {M <: PSY.Machine, S <: PSY.Shaft, A <: PSY.AVR, TG <: PSY.TurbineGov}

    ω0 = PSY.get_ext(device)[CONTROL_REFS][ω_ref_index]
    static_gen = PSY.get_static_injector(device)
    P0 = PSY.get_active_power(static_gen)
    τe = get_inner_vars(device)[τe_var]

    #Get parameters
    pss = PSY.get_pss(device)
    K_ω = PSY.get_K_ω(pss)
    K_p = PSY.get_K_p(pss)

    #It assumes that the system is balanced so:
    get_inner_vars(device)[V_pss_var] = K_ω * (ω0 - ω0) + K_p * (ω0 * τe - P0)

end
=#