PowerFlowWithJuMP
import SparseArrays, PGLib; RawDict = PGLib.pglib("pglib_opf_case118_ieee.m");
import JuMP, Gurobi; GRB_ENV = Gurobi.Env();
import LinearAlgebra.dot as dot
function rectify_dir!(bft)
    for r in 1:size(bft, 1)
        if bft[r, 1] > bft[r, 2]
            bft[r, 1], bft[r, 2] = bft[r, 2], bft[r, 1]
        elseif bft[r, 1] == bft[r, 2]
            error("self-branch")
        end
    end
end; function assert_is_connected(bft)
    B = size(bft, 1)
    pb = Vector(1:B) # primal bs
    sn = [1] # subnet
    for ite in 1:B # only du a full iteration as a conservative choice
        progress = false
        for (i, b) in enumerate(pb) # take out branch b
            if bft[b, 1] in sn
                on = bft[b, 2] # the other node
                on ∉ sn && push!(sn, on)
            elseif bft[b, 2] in sn
                on = bft[b, 1] # the other node
                on ∉ sn && push!(sn, on)
            else
                continue
            end
            popat!(pb, i); progress = true; break # fathom branch b
        end
        if progress == false
            error("ite = $ite. All the rest branches are not connected to the subnet being investigated. Check the current subnet")
        end
        1:maximum(bft) ⊆ sn && return  # The graph is proved to be connected
    end
    error("here shouldn't be reached")
end; function bft_2_A(bft)
    B, N = size(bft, 1), maximum(bft)
    return SparseArrays.sparse([Vector(1:B); Vector(1:B)], vec(bft), [-ones(Int, B); ones(Int, B)], B, N)
end; function adjnode(n)
    nv = Int[]
    for ci in findall(x -> x == n, bft)
        b, c = ci.I # row(= branch), col
        on = bft[b, 3 - c]
        on ∉ nv && push!(nv, on)
    end
    sort(nv)
end; function get_b(n, m) # 🟠 use this after `rectify_dir!(bft)` had been done
    n > m && ((n, m) = (m, n))
    for r in 1:size(bft, 1)
        (bft[r, 1] == n && bft[r, 2] == m) && return r
    end
    error("no such branch exist")
end;
function bpf(b, Dvec = Dref, pvec = pg) # dispatch function of power flow on branch b
    b == 9 && return -pvec[findfirst(x -> x == 10, Gnode)::Int] # b8[9] = -3
    b == 7 && return bpf(9) # b8[7] = -1
    b == 94 && return bpf(93) # b8[94] = -1
    b == 127 && return -bpf(126) # b8[127] = -1
    b == 134 && return -pvec[findfirst(x -> x == 87, Gnode)::Int] # b8[134] = -3
    b == 176 && return -pvec[findfirst(x -> x == 111, Gnode)::Int] # b8[176] = -3
    b == 184 && return Dvec[findfirst(x -> x == 117, Dnode)::Int] # b8[184] = -2
    dup_bvec = [67, 76, 99, 86, 124, 139, 142]
    b ∈ dup_bvec && return false # b8[dup_bvec] .= 0; b8[dup_bvec .- 1] .= 2
    f = findfirst(x -> x == b, F2B)
    return pf[f] # the JuMP variable
end
B, N = 186, 118; # fixed for this network
G = 54; # ✅ in case118, there is no such case that 2 generators connecting to the same bus
D = 99; # demands are data
bft = let # load a rectified bft
    bft = [RawDict["branch"]["$r"][c == 1 ? "f_bus" : "t_bus"] for r in 1:length(RawDict["branch"]), c in 1:2]; # raw bft
    assert_is_connected(bft);
    rectify_dir!(bft);
    bft
end; A = bft_2_A(bft);
Dnode = [RawDict["load"]["$d"]["load_bus"] for d in 1:D]; # function n2d(n) return findfirst(x -> x == n, Dnode)::Int end;
Dref  = [RawDict["load"]["$d"]["pd"]       for d in 1:D]; # demand vector (std reference)
Gnode = [ RawDict["gen"]["$g"]["gen_bus"]  for g in 1:G]; # function n2g(n) return findfirst(x -> x == n, Gnode)::Int end;
Cgen = let
    Cgen = zeros(G)
    for g in 1:G
        v = RawDict["gen"]["$g"]["cost"]
        isempty(v) || (Cgen[g] = v[1] / 1000)
    end
    Cgen
end;
n8 = let # elements being
    # [1]: a bus with load only, and whose degree is larger than 2
    # [2]: a bus with gen only, and whose degree is larger than 2
    # [3]: a bus with both power injection and load
    # [-1]: a bus with load only, and whose degree is 1, we won't add KCL constr here
    # [-3]: a bus with gen only, and whose degree is 1, we won't add KCL constr here
    # [-2]: a bare bus whose degree is 2, we won't add KCL constr here
    # [0]: a bare joint bus
    n8 = zeros(Int8, N)
    n8[Dnode] .+= 1
    n8[Gnode] .+= 2
    n8[117] = -1
    n8[[10, 87, 111]] .= -3
    n8[[9, 63, 81]] .= -2
    n8
end; b8 = let # Note this procedure is specific to case118, therefore executing once is sufficient
    # elements being
    # [2]: this line represents a 2-parallel branch, works as a normal line
    # [1]: a normal line
    # [0]: this line was banned (pf ≡ 0) due to parallel redundancy, we won't allocate pf variable here
    # [-1]: power flow of this line hinges on one of its adjacent line, we won't allocate pf variable here
    # [-2]: power flow of this line hinges on the corresponding pure load, we won't allocate pf variable here
    # [-3]: power flow of this line hinges on the corresponding pure gen, we won't allocate pf variable here
    b8 = ones(Int8, size(bft, 1))
    b8[9] = -3
    b8[7] = -1
    b8[94] = -1
    b8[127] = -1
    b8[134] = -3
    b8[176] = -3
    b8[184] = -2
    dup_bvec = [67, 76, 99, 86, 124, 139, 142]
    b8[dup_bvec] .= 0; b8[dup_bvec .- 1] .= 2
    b8
end; F2B = findall(s -> s ∈ 1:2, b8); F = length(F2B);
model = JuMP.Model(() -> Gurobi.Optimizer(GRB_ENV)); JuMP.@variable(model, pg[g = 1:G] >= 0);
JuMP.@variable(model, pf[f = 1:F]);
let # KCL constrs
    for n in findall(s -> s == 0, n8)
        bv, sgv = SparseArrays.findnz(view(A, :, n))
        JuMP.@constraint(model, dot(sgv, bpf.(bv)) == 0)
    end
    for n in findall(s -> s == 1, n8)
        bv, sgv = SparseArrays.findnz(view(A, :, n))
        JuMP.@constraint(model, dot(sgv, bpf.(bv)) == Dref[findfirst(x -> x == n, Dnode)::Int])
    end
    for n in findall(s -> s == 3, n8)
        bv, sgv = SparseArrays.findnz(view(A, :, n))
        JuMP.@constraint(model, dot(sgv, bpf.(bv)) + pg[findfirst(x -> x == n, Gnode)::Int] == Dref[findfirst(x -> x == n, Dnode)::Int])
    end
    for n in findall(s -> s == 2, n8)
        bv, sgv = SparseArrays.findnz(view(A, :, n))
        JuMP.@constraint(model, dot(sgv, bpf.(bv)) + pg[findfirst(x -> x == n, Gnode)::Int] == 0)
    end
end
JuMP.@objective(model, Min, dot(Cgen, pg))
JuMP.optimize!(model); JuMP.assert_is_solved_and_feasible(model; allow_local = false);
pg_val = JuMP.value.(pg) # output power of all generators
bpf_val = @. JuMP.value(bpf(1:B)) # power flow on all branches