#* Functions for calculating the power flow in a power grid described by a NDD
#*------------------------------------------------------------------------------

#=
Adds or updates the power flow solution in a NDD using a power flow solution obtained with PowerModels.jl and returns the NDD. If model=:ac, the AC power flow equations are used to calculate the flows in the branches. If model=:dc, the DC power flow equations are used. 
=#
function update_pf_data!(
        network_data::Dict{String,<:Any}, 
        pf_result::Dict{String,<:Any};
        model = :ac # how to calculate power flows
    )
    
    ### Add solution (voltage angles and magnitudes) to NDD
    update_data!(network_data, pf_result["solution"])
    
    ### Calculate power flows with PowerModels.jl depending on chosen model
    if model == :dc
        flows = calc_branch_flow_dc(network_data) # DC-PF
        update_data!(network_data, flows) # add power flows to NDD
        calc_branchloads!(network_data, model=:dc)
    elseif model == :ac
        flows = calc_branch_flow_ac(network_data) # AC-PF
        update_data!(network_data, flows) # add power flows to NDD
        calc_branchloads!(network_data, model=:ac)
    else
        throw(ArgumentError("Unknown model $model."))
    end

    return network_data
end

#*------------------------------------------------------------------------------

#=
Reads the power grid data saved in a file and calculates the AC (model=:ac) or DC (model=:dc) power flow solution using PowerModels.jl. The solution is added to a network data dictionary (NDD) of PowerModels.jl that is then returned.
=#
function calc_init_op(file::String; model=:ac)
    network_data = parse_file(file, import_all=false) # read data file
    
    ### Calculate power flow for chosen model
    if model == :dc
        calc_dc_pf!(network_data) # calculate DC-PF solution and add it to NDD
    elseif model == :ac
        calc_ac_pf!(network_data) # calculate AC-PF solution and add it to NDD
    else
        throw(ArgumentError("Unknown model $model."))
    end

    return network_data # return NDD
end

#=
Calculates the AC power flow solution for the given NDD and updates it in the NDD.
CAUTION: Overwrites any previous AC solution contained in the NDD!
=#
function calc_ac_pf!(network_data::Dict{String,<:Any})
    ### Calculate AC-PF solution using PowerModels.jl and Ipopt
    pf_result = run_ac_pf(
        network_data, 
        optimizer_with_attributes(Ipopt.Optimizer, "print_level" => 0)
    )
    update_pf_data!(network_data, pf_result, model=:ac) # update PF in NDD
    return network_data
end

#=
Calculates the DC power flow solution for the given NDD and updates it in the NDD.
CAUTION: Overwrites any previous DC solution contained in the NDD!
=#
function calc_dc_pf!(network_data::Dict{String,<:Any})
    ### Calculate AC-PF solution using PowerModels.jl and Ipopt
    pf_result = run_dc_pf(
        network_data, 
        optimizer_with_attributes(Ipopt.Optimizer, "print_level" => 0)
    )
    update_pf_data!(network_data, pf_result, model=:dc) # update PF in NDD
    return network_data
end

#=
Calculates the loading (flow/capacity) of all branches in the NDD based on the current power flow solution and adds it to the NDD. If model=:dc, only the loading due to the active power flow is calculated. For model=:ac, also the loading due to reactive power and apparent power is calculated.
=#
function calc_branchloads!(network_data::Dict{String,<:Any}; model=:ac)
    branchloads = Dict{String,Any}(
        "baseMVA"  => 100.0,
        "branch" => Dict{String,Any}(), 
        "per_unit" => true
    )
    L = length(network_data["branch"]) # number of branches
    sizehint!(branchloads["branch"], L)

    ### Go through branches and calculate power flows
    for (i, branch) in network_data["branch"]
        if model == :dc # DC-PF
            ### Active power flow
            mw_fr = sqrt(branch["pf"]^2)
            mw_to = sqrt(branch["pt"]^2)
            ### Add active power flow to branchloads
            branchloads["branch"][i] = Dict(
                "MW-loading" => max(mw_fr, mw_to) / branch["rate_a"]
            )
        elseif model == :ac # AC-PF
            ### Active power flow
            mw_fr = sqrt(branch["pf"]^2)
            mw_to = sqrt(branch["pt"]^2)
            ### Reactive power flow
            mvar_fr = sqrt(branch["qf"]^2)
            mvar_to = sqrt(branch["qt"]^2)
            ### Apparent power flow
            mva_fr = sqrt(branch["pf"]^2 + branch["qf"]^2)
            mva_to = sqrt(branch["pt"]^2 + branch["qt"]^2)
            ### Add all flows to branchloads
            branchloads["branch"][i] = Dict(
                "MW-loading" => max(mw_fr, mw_to) / branch["rate_a"],
                "Mvar-loading" => max(mvar_fr, mvar_to) / branch["rate_a"],
                "MVA-loading" => max(mva_fr, mva_to) / branch["rate_a"]
            )
        else
            throw(ArgumentError("Unknown model $model."))
        end
    end

    update_data!(network_data, branchloads) # add branchloads to NDD

    return branchloads
end

#*------------------------------------------------------------------------------