#* Functions for handling the Network Data Dictionary (NDD) of PowerModels.jl
#*------------------------------------------------------------------------------

#=
Builds a network data dictionary conform with PowerModels.jl from several files containing certain input data:
model - whether to calculate AC- or DC-PF (:ac or :dc)
pgfile - .RAW or .m file that contains the detailed grid data (e.g. generator and load information etc.)
busfile - .csv file that contains geographic locations of all buses
branchfile - optional .csv file that contains transmission line lengths
nddfile - .jld2 file in which the final network data dictionary will be saved
=#
function build_network_data(
        model = :ac; 
        pgfile::String,  busfile::String, nddfile::String, branchfile = ""
    )

    ### Read grid data and calculate initial operation point
    network_data = calc_init_op(pgfile, model=model)
    
    ### Add additional data to dictionary
    add_locs!(network_data, busfile) # add geographic bus locations
    if isempty(branchfile) == false # check whether to add line lengths
        add_tl_lengths!(network_data, branchfile) 
    end
    add_tl_voltages!(network_data) # add transmission line voltages
    
    save(nddfile, "network_data", network_data) # save network data

    return nothing
end

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

#=
Reads the geographic bus locations (longitude and latitude coordinates in degrees) from a .csv file and adds them to the network data dictionary (NDD). Returns the resulting NDD, in which each bus has a "bus_lat" and "bus_lon" entry. The columns containg the longitude and latitude coordinates in the .csv file should be called "SubLongitude" and "SubLatitude", respectively.
=#
function add_locs!(network_data::Dict{String,<:Any}, csvfile::String)
    pos = Dict{String,Any}(
        "baseMVA"  => 100.0,
        "bus" => Dict{String,Any}(), 
        "per_unit" => true
    )

    BusData::DataFrame = CSV.File(csvfile) |> DataFrame!
    N = size(BusData, 1) # number of buses
    sizehint!(pos["bus"], N)

    for bus in eachrow(BusData)
        pos["bus"][string(bus[:Number])] = Dict(
            "bus_lon" => bus[:SubLongitude], 
            "bus_lat" => bus[:SubLatitude]
        )
    end
    update_data!(network_data, pos) # add positions to network_data
    
    return network_data
end

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

#=
Reads transmission line lengths from a .csv file and adds them to the network data dictionary.
=#
function add_tl_lengths!(network_data::Dict{String,<:Any}, csvfile::String)
    len = Dict{String,Any}(
        "baseMVA"  => 100.0,
        "branch" => Dict{String,Any}(), 
        "per_unit" => true
    )

    BranchData::DataFrame = CSV.File(csvfile) |> DataFrame!
    N = size(BranchData, 1) # number of branches
    sizehint!(len["branch"], N)

    for branch in eachrow(BranchData)
        from, to = branch[:BusNumFrom], branch[:BusNumTo]
        indices = findall(
            b -> ((b["f_bus"], b["t_bus"]) == (from, to) 
            || (b["f_bus"], b["t_bus"]) == (to, from)), network_data["branch"]
        )
        for index in indices
            len["branch"][index] = Dict("length" => branch[:GICLineDistanceKm])
        end
    end
    update_data!(network_data, len) # add positions to network_data
    
    return network_data
end

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

#=
Determines voltage levels of transmission lines and adds them to the network data dictionary.
=#
function add_tl_voltages!(network_data::Dict{String,<:Any})
    voltages = Dict{String,Any}(
        "baseMVA"  => 100.0,
        "branch" => Dict{String,Any}(), 
        "per_unit" => true
    )
    L = length(network_data["branch"])
    sizehint!(voltages["branch"], L)

    for (i, branch) in network_data["branch"]
        if branch["transformer"] == true
            ### Transformers have a transmission line voltage of 0 kV
            voltages["branch"][i] = Dict(
                "tl_voltage" => 0. 
            )
        else
            ### Transmission lines get base voltage of from bus
            from = branch["f_bus"]
            voltages["branch"][i] = Dict(
                "tl_voltage" => network_data["bus"][string(from)]["base_kv"]
            )
        end
    end
    update_data!(network_data, voltages)

    return network_data
end

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

#=
Returns the bus types of all buses in the NDD in form an ordered dictionary. The different bus types are "Generator", "Load and generator", "Load" and "Empty bus". The ordering is fixed so that the dictionary can be used to plot the buses in a specific order.
=#
#? How to call "empty" buses? Load and generator buses in a substation are connected to these empty buses? Is there only one empty bus per substation?
function get_bustypes(network_data::Dict{String,<:Any})
    ### Active loads
    load = unique(
        [l["load_bus"] for l in collect(values(network_data["load"]))
        if l["status"] == 1]
    )

    ### Active generators
    gen = unique(
        [g["gen_bus"] for g in collect(values(network_data["gen"]))
        if g["gen_status"] == 1]
    )
    
    ### Buses with both load and generation
    load_and_gen = [i for i in load if i in gen]
    
    ### Slack bus
    slack = [
        b["index"] for b in collect(values(network_data["bus"])) 
        if b["bus_type"] == 3
    ]

    ### Remove buses with load and generation from pure loads and generators
    filter!(i -> i ∉ load_and_gen, load)
    filter!(i -> i ∉ load_and_gen, gen)
    ### Remove slack bus from other arrays
    filter!(i -> i ∉ slack, load)
    filter!(i -> i ∉ slack, gen)
    filter!(i -> i ∉ slack, load_and_gen)

    
    ### Empty buses with neither load nor generation
    empty = unique(
        [b["index"] for b in collect(values(network_data["bus"]))
        if b["index"] ∉ vcat(load, gen, load_and_gen, slack)]
    )

    ### Check whether the right number of bus types was obtained
    sum = (
        length(load) + length(gen) + length(load_and_gen) + length(empty) + length(slack)
    )
    N::Int64 = length(network_data["bus"]) # number of buses
    @assert sum == N "$sum bus types were obtained instead of $N"

    return OrderedDict(
        "Generator" => gen,
        "Load and generator" => load_and_gen,
        "Load" => load,
        "Slack" => slack,
        "Empty bus" => empty       
    )
end

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

#=
Returns the (string) indices of transmission lines in the network data dictionary that would be considered as underground lines depending on the used maximum length of underground lines and minimum MW-load served. If mode = :sum, the sum of MW-loads of both ends of a line has to exceed min_MW_load, and if mode = :single, one end alone has to exceed min_MW_load.
=#
function get_underground_tl(
        network_data::Dict{String,<:Any}, 
        max_length = 12.875, # in km
        min_MW_load = 2.; # in per-unit (assumed 100 MW base)
        mode = :sum # :sum or :single
    )

    MW_loads = get_MW_loads(network_data) # dictionary of MW loads

    underground_tl = []
    if mode == :single # one end alone has to exceed min_MW_load
        filter!(load -> last(load) >= min_MW_load, MW_loads)
        large_MW_loads = collect(keys(MW_loads))
        ### Go through branches and identify underground transmission lines
        for (i, branch) in network_data["branch"]
            if branch["transformer"] == false && branch["length"] < max_length
                from, to = branch["f_bus"], branch["t_bus"]
                if from in large_MW_loads || to in large_MW_loads
                    push!(underground_tl, i)
                end
            end
        end
    elseif mode == :sum # both ends together have to exceed min_MW_load
        ### Go through branches and identify underground transmission lines
        for (i, branch) in network_data["branch"]
            if branch["transformer"] == false && branch["length"] < max_length
                from, to = branch["f_bus"], branch["t_bus"]
                sum_MW_load = 0.
                if haskey(MW_loads, from)
                    sum_MW_load += MW_loads[from]
                end
                if haskey(MW_loads, to)
                    sum_MW_load += MW_loads[to]
                end
                if sum_MW_load >= min_MW_load
                    push!(underground_tl, i)
                end
            end
        end
    else
        throw(ArgumentError("Unknown mode $mode."))
    end

    return underground_tl
end

#=
Returns a dictionary with load bus indices (Int64) as keys and their respective MW load as values.
=#
function get_MW_loads(network_data::Dict{String,<:Any})
    load_dict = Dict{Int64,Float64}()
    for (i, load) in network_data["load"]
        sub_dict = Dict(load["load_bus"] => load["pd"])
        merge!(+, load_dict, sub_dict)
    end
    return load_dict
end

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

#=
Returns an array of tuples containing branch indices of branches with a minimum loading of min_loading and their actual loading (considering the power flow type pf_type). Choosing min_loading = 1. yields overloaded branches.
=#
function get_branches(
        network_data::Dict{String,<:Any}; 
        min_loading = 0.0, # minimum loading of branches
        pf_type = "MVA-loading" # or "MW-loading", "Mvar-loading"
    )

    branches = [
        (i, br[pf_type]) for (i, br) in network_data["branch"]
        if br["br_status"] == 1 && br[pf_type] >= min_loading
    ]

    return branches
end

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

#=
Deactivates branches that have been overloaded due to a power flow redistribution (flow/capacity > 1.). The branches are identified by their string indices in branch_ids. The status of these branches is set to 0 in network_data. Parallel branches are not automatically deactivated. Afterwards PowerModels function simplify_network! is run in order to remove dangling buses etc. that might occur due to the failures.
=#
function disable_branches!(
        network_data::Dict{String,<:Any}, 
        branch_ids::Array
    )
    ### Go through branches to deactivate
    for id in branch_ids
        network_data["branch"][id]["br_status"] = 0 # deactivate branch
    end

    simplify_network!(network_data) # remove dangling buses etc.

    return nothing
end

#=
Function that deactivates (overhead) transmission lines identified by their indices (string) in tl_ids that have been destroyed by a hurricane. For any given line index all parallel lines also become deactivated. The status of these lines is set to 0 in network_data. Afterwards PowerModels function simplify_network! is run in order to remove dangling buses etc. that might occur due to the failures.
=#
function destroy_tl!(
        network_data::Dict{String,<:Any}, 
        tl_ids::Array # string ids of destroyed transmission lines
    )

    ### Go through transmission lines to destroy
    for id in tl_ids
        ## Deactivate id and all parallel lines
        from = network_data["branch"][id]["f_bus"]
        to = network_data["branch"][id]["t_bus"]
        ## Find all parallel branches
        ids = findall(
            b -> (b["f_bus"], b["t_bus"]) == (from, to), 
            network_data["branch"]
        )
        for idx in ids
            network_data["branch"][idx]["br_status"] = 0
        end
    end

    simplify_network!(network_data) # remove dangling buses etc.

    return nothing
end

#=
Deactivates all components (generator, loads, etc.) in a connected component.
=#
function _deactivate_cc!(
        network_data::Dict{String,<:Any}, 
        active_cc_gens::Array{Int64,1}, 
        active_cc_loads::Array{Int64,1}
    )

    ### Deactivate all active generators in connected component
    for g in active_cc_gens
        network_data["gen"]["$g"]["gen_status"] = 0
    end

    ### Deactivate all active loads in connected component
    for l in active_cc_loads
        network_data["load"]["$l"]["status"] = 0
    end

    ### Deactivate rest of connected component (branches, dangling buses etc.)
    simplify_network!(network_data)
    
    return nothing
end