#* Functions for plotting a power grid described by a network data dictionary
#*------------------------------------------------------------------------------

#= 
Function for recursively merging two dictionaries (used for updating the dictionary containing plot settings).
=#
_recursive_merge(x::AbstractDict,y::AbstractDict) = merge(_recursive_merge,x,y)
_recursive_merge(x,y) = y

#=
Returns the minimal and maximal values of bus longitude and latitude coordinates contained a dictionary. An offset can be used to arbitrarily increase the area.
=#
function _get_pg_area(pos::Dict{Int64,Tuple{Float64,Float64}}, offset=0.)
    locs = collect(values(pos))
    lon_min = minimum(p -> p[1], locs) - offset
    lon_max = maximum(p -> p[1], locs) + offset
    lat_min = minimum(p -> p[2], locs) - offset
    lat_max = maximum(p -> p[2], locs) + offset
    return lon_min, lon_max, lat_min, lat_max
end

function _get_pg_area(network_data::Dict{String,<:Any}, offset=0.)
    pos = Dict(
        b["index"] => (b["bus_lon"], b["bus_lat"]) 
        for b in collect(values(network_data["bus"]))
    ) # geographic bus locations

    return _get_pg_area(pos, offset)
end

#=
Returns a dictionary of default plot settings for various plotting functions.
=#
function _default_settings(func::Symbol)
    if func == :plot_pg
        defaults = Dict(
            ### Settings for _draw_buses!
            "Buses" => Dict(
                "Generator" => Dict(
                    "marker" => "s",
                    "size" => 30,
                    "color" => "darkorange",
                    "alpha" => 1,
                    "label" => "Generators",
                    "show" => true
                ),
                "Load and generator" => Dict(
                    "marker" => "o",
                    "size" => 30,
                    "color" => "darkorange",
                    "alpha" => 1,
                    "label" => "nolabel",
                    "show" => true
                ),
                "Slack" => Dict(
                    "marker" => "s",
                    "size" => 70,
                    "color" => "red",
                    "alpha" => 0.6,
                    "label" => "Slack",
                    "show" => true
                ),
                "Load" => Dict(
                    "marker" => "o",
                    "size" => 15,
                    "color" => "limegreen",
                    "alpha" => 0.7,
                    "label" => "Loads",
                    "show" => true
                ),
                "Empty bus" => Dict(
                    "marker" => "o",
                    "size" => 2,
                    "color" => "black",
                    "alpha" => 0.3,
                    "label" => "Empty buses",
                    "show" => true
                )
            ),
            ### Settings for _draw_branches!
            "Branches" => Dict(
                "br_status" => "active", # plot only "active" or "all" branches
                "br_coloring" => "equal", # how to color branches (can be set to "equal", "voltage", "MW-loading", "Mvar-loading" or "MVA-loading")
                "br_color" => "k", # default color for "equal" coloring
                "br_lw" => 2,
                "br_alpha" => 1
            ),
            "draw_ticks" => [true, true, true, true], # whether to show ticks
            "draw_legend" => true, # whether to draw a legend
            "xlabel" => L"Longitude in $°$",
            "ylabel" => L"Latitude in $°$",
        )
    end

    return defaults
end

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

function plot_pg_map(
        network_data::Dict{String,<:Any},
        settings = Dict{String,Any}(); # dictionary containing plot settings
        figpath::String # where to save the figure
    )

    ### Python imports
    cartopy = pyimport("cartopy")
    cticker = pyimport("cartopy.mpl.ticker")
    
    ### Setup figure and plot geographic map
    fig::Figure = plt.figure()
    w::Float64, h::Float64 = plt.figaspect(2/3)::Vector{Float64}
    fig.set_size_inches(1.5w, 1.5h)
    ax = fig.add_subplot(projection=cartopy.crs.PlateCarree())
    ax.set_aspect("equal")

    ### Set area to plot
    xmin, xmax, ymin, ymax = _get_pg_area(network_data, 0.4)
    ax.set_extent([xmin, xmax, ymin, ymax])

    ### Get state borders
    states_provinces = cartopy.feature.NaturalEarthFeature(
        category="cultural",
        name="admin_1_states_provinces_lines",
        scale="50m",
        facecolor="none"
    )
    ### Add wanted features to plot
    ax.add_feature(cartopy.feature.LAND)
    ax.add_feature(cartopy.feature.OCEAN)
    ax.add_feature(cartopy.feature.COASTLINE)
    ax.add_feature(cartopy.feature.BORDERS)
    ax.add_feature(states_provinces, edgecolor="gray")

    ### Show longitude and latitude values
    gl = ax.gridlines(crs=cartopy.crs.PlateCarree(), draw_labels=true)
    gl.xlabels_top = false
    gl.ylabels_right = false
    gl.xlines = false
    gl.ylines = false
    gl.xformatter = cartopy.mpl.gridliner.LONGITUDE_FORMATTER
    gl.yformatter = cartopy.mpl.gridliner.LATITUDE_FORMATTER

    ### Set plot settings and plot power grid
    settings = _recursive_merge(_default_settings(:plot_pg), settings)
    _plot_pg!(ax, network_data, settings, figpath)

    return nothing
end

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

#=
Plots the power grid described by the network data dictionary (NDD). Possible plot settings are shown in _default_settings.
=#
function plot_pg(
        network_data::Dict{String,<:Any},
        settings = Dict{String,Any}(); # dictionary containing plot settings
        figpath::String # where to save the figure
    )

    ### Setup figure
    figure::Figure, ax::PyObject = plt.subplots()::Tuple{Figure,PyObject}
    w::Float64, h::Float64 = plt.figaspect(2/3)::Vector{Float64}
    figure.set_size_inches(1.5w, 1.5h)
    ax.set_aspect("equal")

    ### Set plot settings and plot power grid
    settings = _recursive_merge(_default_settings(:plot_pg), settings)
    _plot_pg!(ax, network_data, settings, figpath)
    
    return nothing
end

#=
Plots the power grid described by the network data dictionary (NDD) onto an already existing axes. Possible plot settings are shown in _default_settings.
=#
function _plot_pg!(
        ax::PyObject, # axes to draw power grid onto
        network_data::Dict{String,<:Any},
        settings::Dict{String,<:Any}, # dictionary containing plot settings
        figpath::String # where to save the figure
    )

    ### Draw power grid graph
    nx = pyimport("networkx")
    G::PyObject = nx.Graph() # empty graph
    _draw_pg!(ax, G, network_data, settings) # draw power grid onto axes
   
    plt.savefig(figpath, bbox_inches="tight")

    plt.close("all") # close figure
    return nothing
end

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

#=
Draws a graph for the power grid described by the NDD with options according to the dictionary "settings" (see _default_settings for possible options).
=#
function _draw_pg!(
        ax::PyObject, # axes to draw power grid onto
        G::PyObject, # power grid graph
        network_data::Dict{String,<:Any}, 
        settings::Dict{String,<:Any} # dictionary containing plot settings
    )

    ### Plot buses into graph
    G, bus_markers, bus_labels = _draw_buses!(G, network_data, settings)

    ### Plot branches into graph
    G, br_markers, br_labels, cbar = _draw_branches!(G, network_data, settings)

    ### Check for a predefined area to show
    if haskey(settings, "area")
        area = settings["area"]
        plt.xlim(area[1], area[2])
        plt.ylim(area[3], area[4])
    end

    ### Axes settings
    ax.tick_params(
        left = settings["draw_ticks"][1],
        bottom = settings["draw_ticks"][2], 
        labelleft = settings["draw_ticks"][3], 
        labelbottom = settings["draw_ticks"][4]
    )
    plt.xlabel(settings["xlabel"])
    plt.ylabel(settings["ylabel"], rotation=90)

    ### Draw legend, if wanted
    if settings["draw_legend"] == true
        all_markers = vcat(bus_markers, br_markers)
        all_labels = vcat(bus_labels, br_labels)
        plt.legend(all_markers, all_labels)
    end

    return ax, G
end

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

#=
Draws buses contained in the NDD as nodes into graph G. The nodes are displayed according to the settings dictionary (see _default_settings).
=#
function _draw_buses!(
        G::PyObject, # power grid graph
        network_data::Dict{String,<:Any},
        settings::Dict{String,<:Any} # dictionary containing plot settings 
    )

    ### Python imports
    nx = pyimport("networkx")
    mlines = pyimport("matplotlib.lines")
    
    bustypes = get_bustypes(network_data) # types of all buses
    pos = Dict(
        b["index"] => (b["bus_lon"], b["bus_lat"]) 
        for b in collect(values(network_data["bus"]))
    ) # geographic bus locations
    bus_markers = [
        mlines.Line2D([], [], color=b["color"], marker=b["marker"], ls="None")
        for b in collect(values(settings["Buses"]))
        if b["label"] != "nolabel"
    ] # markers for legend
    bus_labels = [
        b["label"] for b in collect(values(settings["Buses"])) 
        if b["label"] != "nolabel"
    ] # labels for legend

    ### Draw different buses as nodes
    for (type, buses) in bustypes
        bus_settings = settings["Buses"][type]
        if bus_settings["show"] == true
            nx.draw_networkx_nodes(
                G, pos,
                nodelist = buses,
                node_shape = bus_settings["marker"],
                node_size = bus_settings["size"],
                node_color = bus_settings["color"],
                alpha = bus_settings["alpha"]
            )
        end
    end

    return G, bus_markers, bus_labels
end

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

#=
Draws branches contained in the NDD as edges into graph G. The branches are displayed according to the settings dictionary (see _default_settings).
=#
function _draw_branches!(
        G::PyObject, # power grid graph
        network_data::Dict{String,<:Any},
        settings::Dict{String,<:Any} # dictionary containing plot settings 
    )

    br_settings = settings["Branches"]
    br_coloring = br_settings["br_coloring"]

    ### Draw branches according to coloring mode
    if br_coloring == "equal"
        G, br_markers, br_labels, cbar = _draw_br_equal!(
            G, network_data, br_settings
        )
    elseif br_coloring == "voltage"
        G, br_markers, br_labels, cbar = _draw_br_voltage!(
            G, network_data, br_settings
        )
    elseif br_coloring in ["MW-loading","Mvar-loading","MVA-loading"]
        G, br_markers, br_labels, cbar = _draw_br_branchloads!(
            G, network_data, br_settings
        )
    else
        throw(ArgumentError("Unknown branch coloring $br_coloring."))
    end

    return G, br_markers, br_labels, cbar
end

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

#=
Draws branches contained in the NDD with coloring mode "equal". All branches are displayed using the same color.
=#
function _draw_br_equal!(
        G::PyObject, # power grid graph
        network_data::Dict{String,<:Any},
        br_settings::Dict{String,<:Any} # dictionary containing plot settings 
    )
    
    ### Python imports
    nx = pyimport("networkx")
    
    pos = Dict(
        b["index"] => (b["bus_lon"], b["bus_lat"]) 
        for b in collect(values(network_data["bus"]))
    ) # geographic bus locations
    branches = collect(values(network_data["branch"])) # branch dictionaries
    
    ### Get edges contained in the NDD
    if br_settings["br_status"] == "active" # only plot active branches
        edges = [(b["f_bus"],b["t_bus"]) for b in branches if b["br_status"]==1]
    elseif br_settings["br_status"] == "all" # plot all branches
        edges = [(b["f_bus"],b["t_bus"]) for b in branches]
    else
        br_status = br_settings["br_status"]
        throw(ArgumentError("Unknown branch status $br_status."))
    end

    ### Draw edges
    drawn_edges = nx.draw_networkx_edges(
        G, pos, 
        edgelist = edges,
        width = br_settings["br_lw"],
        edge_color = br_settings["br_color"],
        alpha = br_settings["br_alpha"]
    )

    return G, [], [], nothing    
end

#=
Draws branches contained in the NDD with coloring mode "MW-loading", "Mvar-loading" or "MVA-loading". The branches are colored depending on their loading (flow/capacity).
=#
function _draw_br_branchloads!(
        G::PyObject, # power grid graph
        network_data::Dict{String,<:Any},
        br_settings::Dict{String,<:Any} # dictionary containing plot settings
    )

    ### Python imports
    nx = pyimport("networkx")

    pos = Dict(
        b["index"] => (b["bus_lon"], b["bus_lat"]) 
        for b in collect(values(network_data["bus"]))
    ) # geographic bus locations
    branches = collect(values(network_data["branch"])) # branch dictionaries
    br_coloring = br_settings["br_coloring"] # what kind of loading to use
    br_status = br_settings["br_status"]

    ### Get edges contained in the NDD and their loadings
    if br_status == "active" # only plot active branches
        edges = [(b["f_bus"],b["t_bus"]) for b in branches if b["br_status"]==1]
        branchloads = [b[br_coloring] for b in branches if b["br_status"]==1]
    elseif br_status == "all" # plot all branches
        edges = [(b["f_bus"],b["t_bus"]) for b in branches]
        branchloads = [b[br_coloring] for b in branches]
    else
        throw(ArgumentError("Unknown branch status $br_status."))
    end

    ### Draw edges
    cmap = plt.cm.inferno_r
    vmin, vmax = 0., 1.
    drawnedges = nx.draw_networkx_edges(
        G, pos, 
        edgelist = edges,
        width = br_settings["br_lw"],
        edge_color = branchloads,
        edge_cmap = cmap,
        edge_vmin = vmin,
        edge_vmax = vmax,
        alpha = br_settings["br_alpha"]
    )

    ### Add colorbar
    sm = plt.cm.ScalarMappable(cmap=cmap, norm=plt.Normalize(vmin, vmax))
    cbar = plt.colorbar(sm)
    cbar.ax.set_ylabel(
        "Line $br_coloring " * L"$F_{ij}/C_{ij}$", rotation=-90, va="bottom"
    )

    return G, [], [], cbar
end

#=
Draws branches contained in the NDD with coloring mode "voltage". Transmission lines are colored according to their voltage levels.
=#
function _draw_br_voltage!(
        G::PyObject, # power grid graph
        network_data::Dict{String,<:Any},
        br_settings::Dict{String,<:Any} # dictionary containing plot settings
    )
    
    ### Python imports
    nx = pyimport("networkx")
    mlines = pyimport("matplotlib.lines")

    pos = Dict(
        b["index"] => (b["bus_lon"], b["bus_lat"]) 
        for b in collect(values(network_data["bus"]))
    ) # geographic bus locations
    branches = collect(values(network_data["branch"])) # branch dictionaries
    br_markers = Array{PyObject,1}() # markers for legend
    br_labels = Array{String,1}() # labels for legend
    
    ### Get edges contained in the NDD and their voltage levels
    if br_settings["br_status"] == "active" # only plot active branches
        edges = [
            (b["f_bus"],b["t_bus"]) for b in branches 
            if b["br_status"] == 1
        ]
        voltages = [
            string(b["tl_voltage"]) for b in branches 
            if b["br_status"] == 1
        ]
    elseif br_settings["br_status"] == "all" # plot all branches
        edges = [(b["f_bus"],b["t_bus"]) for b in branches]
        voltages = [string(b["tl_voltage"]) for b in branches]
    else
        br_status = br_settings["br_status"]
        throw(ArgumentError("Unknown branch status $br_status."))
    end

    ### Assign colors to voltage levels and add markers and labels for legend
    voltages[voltages .== "0.0"] .= "k" # transformers
    mcolors = pyimport("matplotlib.colors")
    tableau = [
        key for key in keys(mcolors.TABLEAU_COLORS) 
        if key ∉ ["tab:orange", "tab:green"] # orange and green used for buses
    ]
    for (i, v) in enumerate(sort(unique(filter(v -> v != "k", voltages))))
        voltages[voltages .== v] .= tableau[i]
        push!(br_markers, mlines.Line2D([], [], color=tableau[i], ls="-"))
        push!(br_labels, string(v) * " kV")
    end

    ### Draw edges
    drawnedges = nx.draw_networkx_edges(
        G, pos, 
        edgelist = edges,
        width = br_settings["br_lw"],
        edge_color = voltages,
        alpha = br_settings["br_alpha"]
    )
    
    return G, br_markers, br_labels, nothing
end