Data.jl

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

#=
Reads the geographic bus locations (longitude and latitude coordinates in degrees) from a CSV file and adds them to the 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

#=
Returns a dictionary containing geographic bus locations that are included in the NDD. The dicitonary structure is bus-index => (lon, lat). If getlims=true, the minmimum and maximum longitude and latitude coordinates are returned as well. Can be useful for plotting the grid.
=#
function get_locs(network_data::Dict{String,<:Any}; getlims=false)
    pos = Dict(
        parse(Int64, i) => (b["bus_lon"], b["bus_lat"]) 
        for (i, b) in network_data["bus"]
    )

    ### Check whether to return lon- and lat-limits or not
    if getlims == true
        N = length(network_data["bus"]) # number of buses
        lats = [
            network_data["bus"]["$i"]["bus_lat"] 
            for i in keys(network_data["bus"])
        ] # latitudes
        lons = [
            network_data["bus"]["$i"]["bus_lon"] 
            for i in keys(network_data["bus"])
        ] # longitudes
        xlims = (minimum(lons), maximum(lons))
        ylims = (minimum(lats), maximum(lats))
        return pos, xlims, ylims
    else
        return pos
    end
end

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

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