#* Functions that allow to compare different power flow solutions
#*------------------------------------------------------------------------------

#* NOTE: Utility functions used are defined in PlotUtils.jl

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

#=
Scatters branch loads contained in y_ndd versus branch loads in x_ndd. Possible plot settings can be seen in PlotUtils.jl.
=#
function scatter_branchloads(
        x_ndd::Dict{String,<:Any},
        y_ndd::Dict{String,<:Any},
        settings::Dict{String,<:Any}; # plot settings
        figpath = "" # where to save the figure
    )

    ### Set plot settings
    settings = _recursive_merge(
        _default_settings(:scatter_branchloads), settings
    )
    
    ### Get branch loadings from both network data dictionaries
    x_bl_dict = Dict(
        i => br[settings["x_pf_type"]] for (i, br) in x_ndd["branch"] 
        if br["br_status"] == 1
    )
    y_bl_dict = Dict(
        i => br[settings["y_pf_type"]] for (i, br) in y_ndd["branch"]
        if br["br_status"] == 1
    )

    # @assert length(x_bl_dict) == length(y_bl_dict) # same number of branches?
    L = length(y_bl_dict)
    # println(L)
    x_bl, y_bl = zeros(L), zeros(L) # arrays for branch loads
    ### Write branch loads with correct ordering into arrays    
    for (n, key) in enumerate(keys(y_bl_dict))
        x_bl[n] = x_bl_dict[key]
        y_bl[n] = y_bl_dict[key]
    end

    ### Plot straight lines with slope 1 as reference
    x = [i for i in 0:0.01:1]
    plt.plot(x, x, color="k", alpha=.7) # plot line with slope 1

    ### Scatter branch loads versus each other
    diff = y_bl - x_bl # deviation from ratio 1 for coloring
    diff[diff .== 0.0] .= -10 # map exact zero values to -10 for distinct color
    # mcolors = pyimport("matplotlib.colors")
    # offset = mcolors.TwoSlopeNorm(
    #     vcenter=0, vmin=minimum(diff), vmax=maximum(diff)
    # ) # match diff to values between 0 and 1
    cmap = plt.get_cmap(settings["cmap"])
    cmap.set_under("tab:green")
    sc = plt.scatter(
        x_bl, y_bl,       
        s = settings["size"],
        c = diff,
        # c = offset(diff),
        cmap = cmap,
        vmin = -1.,
        vmax = 1.,
        edgecolors = settings["ec"],
        linewidths = settings["lw"],
        alpha = settings["alpha"]
    )

    ### Add colorbar
    sm = plt.cm.ScalarMappable(cmap=cmap, norm=plt.Normalize(-1., 1.))
    # sm = plt.cm.ScalarMappable(cmap=cmap, offset)
    cbar = plt.colorbar(sm)
    cbar.ax.plot([-1, 1], [0.0, 0.0], "tab:green")
    # cbar.add_lines(CS2)
    cbar.ax.set_ylabel(
        L"Difference $y-x$", rotation=-90, va="bottom"
    )

    ### Plot labels
    plt.xlabel(settings["xlabel"])
    plt.ylabel(settings["ylabel"])

    plt.savefig(figpath, bbox_inches="tight") # save figure
    plt.close("all") # close figure
    return nothing
end

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

#=
Plots a power grid similar to plot_pg in PlotPG.jl and colors the edges according to the difference in branch loads in ndd and ref_ndd. Possible plot settings can be seen in PlotUtils.jl.
=#
function plot_pg_pf_diff(
        ndd::Dict{String,<:Any}, # NDD with (new) power flow
        ref_ndd::Dict{String,<:Any}, # NDD with reference power flow
        settings = Dict{String,Any}(); # plot settings
        figpath = "" # 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_pf_diff), settings)
    _plot_pg_pf_diff!(ax, ndd, ref_ndd, settings, figpath)

    return nothing
end

function _plot_pg_pf_diff!(
        ax::PyObject, # axes to draw power grid onto
        ndd::Dict{String,<:Any}, # NDD with (new) power flow
        ref_ndd::Dict{String,<:Any}, # NDD with reference power flow
        settings = Dict{String,Any}(), # plot settings
        figpath = "" # where to save the figure
    )
    
    ### Draw power grid graph
    nx = pyimport("networkx")
    G::PyObject = nx.Graph() # empty graph

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

    ### Plot branches into graph
    G, br_markers, br_labels, br_cbar = _draw_br_pf_diff!(
        G, ndd, ref_ndd, 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

    plt.savefig(figpath, bbox_inches="tight") # save figure
    plt.close("all") # close figure
    return ax, G
end

function _draw_br_pf_diff!(
        G::PyObject, # power grid graph
        ndd::Dict{String,<:Any}, # NDD with (new) power flow
        ref_ndd::Dict{String,<:Any}, # NDD with reference power flow
        settings::Dict{String,<:Any} # dictionary containing plot settings
    )
        
    pos = Dict(
        b["index"] => (b["bus_lon"], b["bus_lat"]) 
        for b in collect(values(ndd["bus"]))
    ) # geographic bus locations
    branches = collect(values(ndd["branch"])) # branch dictionaries

    ### Get edges contained in the NDD and the loading difference
    edges = Array{Tuple{Int64,Int64},1}() # array for edges
    bl_diff = Array{Float64,1}() # array for branch load differences
    for (i, br) in ndd["branch"]
        if br["br_status"] == 1
            push!(edges, (br["f_bus"], br["t_bus"]))
            pf_type = settings["Branches"]["pf_type"]
            push!(bl_diff, br[pf_type] - ref_ndd["branch"][i][pf_type])
        end
    end
    
    ### Draw edges and color them according to bl_diff
    mcolors = pyimport("matplotlib.colors")
    offset = mcolors.TwoSlopeNorm(
        vcenter=0, vmin=minimum(bl_diff), vmax=maximum(bl_diff)
    )
    vmin, vmax = minimum(offset(bl_diff)), maximum(offset(bl_diff))
    cmap = plt.get_cmap(settings["Branches"]["cmap"])
    nx = pyimport("networkx")
    drawnedges = nx.draw_networkx_edges(
        G, pos, 
        width = settings["Branches"]["br_lw"],
        edgelist = edges,
        edge_color = offset(bl_diff),
        edge_vmin = vmin,
        edge_vmax = vmax,
        edge_cmap = cmap
    )

    ### Add colorbar
    sm = plt.cm.ScalarMappable(cmap=cmap, offset)
    cbar = plt.colorbar(sm)

    cbar.ax.set_ylabel(
        settings["Branches"]["cbar_label"], rotation=-90, va="bottom"
    )

    return G, [], [], cbar
end