text:
using Caesar, DelimitedFiles, JLD
using IncrementalInference
using Gadfly, Cairo, Fontconfig
using KernelDensityEstimatePlotting

include(joinpath(@__DIR__,"slamUtils.jl"))
include(joinpath(@__DIR__,"plotSASUtils.jl"))
include(joinpath(@__DIR__,"expDataUtils.jl"))

function main(expID::String, rangegap::Int, wstart::Int, wend::Int, trialID::Int)

    datawindow = collect(wstart:wend);
    allpaths = getPaths(expID,"range", currloc = "lj", trialID = trialID);
    scriptHeader = joinpath(allpaths[3],expID*"_rgap$(rangegap)_trial$(trialID)_f$(wstart)t$(wend)_")

    igt = loadGT(datawindow,expID,allpaths);
    mfin = loadRanges(datawindow,allpaths)
    posData = loadNav(datawindow,allpaths)

    dposData = deepcopy(posData)
    cumulativeDrift!(dposData,[0.0;0],[0.2,0.2])

    window = 1:datawindow[end]-datawindow[1];
    fg = initfg();
    beacon = :l1
    addVariable!(fg, beacon, Point2 )

    for i in window
        sym = Symbol("x$i")
        addVariable!(fg, sym, Point2)
    end

    rtkCov = Matrix(Diagonal([0.1;0.1].^2));
    dpσ = Matrix(Diagonal([0.2;0.2].^2))

    for i in window
        sym = Symbol("x$i")
        nextsym = Symbol("x$(i+1)")

        if i == window[1]
            pp = PriorPoint2(MvNormal(posData[i,:], rtkCov))
            addFactor!(fg, [sym;], pp, autoinit=true)

            dx = dposData[i+1,1] - dposData[i,1];
            dy = dposData[i+1,2] - dposData[i,2];
            p2p2 = Point2Point2(MvNormal([dx;dy],dpσ))
            addFactor!(fg, [sym;nextsym], p2p2, autoinit=true)
        elseif  i == window[end]
            pp = PriorPoint2(MvNormal(posData[i,:], rtkCov))
            addFactor!(fg, [sym;], pp, autoinit=false)
        else
            dx = dposData[i+1,1] - dposData[i,1];
            dy = dposData[i+1,2] - dposData[i,2];
            p2p2 = Point2Point2(MvNormal([dx;dy],dpσ))
            addFactor!(fg, [sym;nextsym], p2p2, autoinit=true)
        end
    end

    rangewindow = window[1]:rangegap:window[end]
    for i in rangewindow
        sym = Symbol("x$i")
        ppR = Point2Point2Range(AliasingScalarSampler(mfin[:,1,i],exp.(mfin[:,2,i]),SNRfloor=0.8))
        addFactor!(fg, [sym;beacon], ppR, autoinit=false)
    end

    writeGraphPdf(fg,viewerapp="", engine="neato", filepath=scriptHeader*"fg.pdf")

    getSolverParams(fg).drawtree = false
    getSolverParams(fg).showtree = false
    getSolverParams(fg).limititers=500
    getSolverParams(fg).maxincidence = 400

    tree, smt, hist = solveTree!(fg)
    tree, smt, hist = solveTree!(fg)
    drawTree(tree, filepath=scriptHeader*"bt.pdf")

    plotSASDefault(fg,expID, posData,igt,dposData,datadir=allpaths[1],savedir=scriptHeader*"SASdefault.pdf")

    #RangeOnly PLOTTING ------------
    plk= [];
    push!(plk,plotBeaconGT(igt));
    push!(plk,plotBeaconMax(fg));
    push!(plk,plotBeaconMean(fg));

    for var in rangewindow
        mysym = Symbol("x$var")
        push!(plk, plotPoint(getVal(fg,mysym), colorIn = colorant"orange"))
    end

    # counter = 1;
    for var in rangewindow       #Plot only for range factors
        sym = Symbol("x$var")
        for mysym in ls(fg,sym)
            if occursin(r"l1",string(mysym))
                # if counter < 4       #Plot one or more approxConv
                    L1ac = approxConv(fg,mysym,:l1)
                    # K1 = plotKDEContour(kde!(L1ac),xlbl="X (m)", ylbl="Y (m)",levels=3,layers=true)
                    # push!(plk,K1...)
                    push!(plk,layer(x=L1ac[1,:],y=L1ac[2,:],Geom.histogram2d(xbincount=400, ybincount=400)))
                    # push!(pkde,plotKDE(kde!(L1ac),layers=true)...)
                    # onetime = false
                    # push!(plk,Gadfly.Theme(key_position = :none));
                    # push!(plk, Guide.xlabel("X (m)"), Guide.ylabel("Y (m)"))
                    # counter +=1;
                # end

            end
        end
    end

    plotKDEMeans!(plk,fg);
    push!(plk,plotPath(posData));
    push!(plk,plotPath(dposData,colorIn=colorant"blue"));
    plotBeaconContours!(plk,fg);

    if expID == "dock"
        push!(plk, Coord.cartesian(xmin=-40, xmax=140, ymin=-140, ymax=30,fixed=true))
    elseif expID == "drift"
        push!(plk, Coord.cartesian(xmin=20, xmax=200, ymin=-220, ymax=0,fixed=true))
    end
    savefile = scriptHeader*"showRanges.pdf"
    Gadfly.plot(plk...) |> PDF(savefile)

    plk= [];
    push!(plk,plotBeaconGT(igt));
    # push!(plk,plotBeaconMax(fg));
    # push!(plk,plotBeaconMean(fg));
    L1pd = predictbelief(fg, :l1,ls(fg, :l1))
    L1ac = L1pd[1];

    L1 = getKDEMax(kde!(L1ac))
    push!(plk,layer(x=[L1[1];],y=[L1[2];], label=String["Beacon Max";],Geom.point,Geom.label(hide_overlaps=false), order=2, Theme(default_color=colorant"darkorange",highlight_width = 0pt)));

    L1 = getKDEMean(kde!(L1ac))
    push!(plk,layer(x=[L1[1];],y=[L1[2];], label=String["Beacon Mean";],Geom.point,Geom.label(hide_overlaps=false), order=2, Theme(default_color=colorant"orangered",highlight_width = 0pt)))

    for var in rangewindow
        mysym = Symbol("x$var")
        push!(plk, plotPoint(getVal(fg,mysym), colorIn = colorant"orange"))
    end

    push!(plk,layer(x=L1ac[1,:],y=L1ac[2,:],Geom.histogram2d(xbincount=300, ybincount=300)))

    K1 = plotKDEContour(kde!(L1ac),xlbl="X (m)", ylbl="Y (m)",levels=5,layers=true);
    push!(plk,K1...)
    push!(plk,Gadfly.Theme(key_position = :none));

    plotKDEMeans!(plk,fg);
    push!(plk,plotPath(posData));
    push!(plk,plotPath(dposData,colorIn=colorant"blue"));
    # plotBeaconContours!(plk,fg);
    if expID == "dock"
        push!(plk, Coord.cartesian(xmin=-40, xmax=140, ymin=-140, ymax=30,fixed=true))
    elseif expID == "drift"
        push!(plk, Coord.cartesian(xmin=20, xmax=200, ymin=-220, ymax=0,fixed=true))
    end
    savefile = scriptHeader*"predBel.pdf"
    Gadfly.plot(plk...) |> PDF(savefile)

    if expID == "dock"
        l1fit = getKDEMean(getBelief(fg,:l1))
        meanerror = sqrt((igt[1]-l1fit[1])^2+(igt[2]-l1fit[2])^2)

        l1max = getKDEMax(getBelief(fg,:l1))
        maxerror = sqrt((igt[1]-l1max[1])^2+(igt[2]-l1max[2])^2)

        mykde = getBelief(fg,:l1)
        mynorm = kde!(rand(fit(MvNormal,getVal(fg,:l1)),200))
        kld = min(abs(KernelDensityEstimate.kld(mynorm,mykde)),abs(KernelDensityEstimate.kld(mykde,mynorm)))
    elseif expID == "drift"
        l1fit = getKDEMean(getBelief(fg,:l1))
        meanerror = sqrt((mean(igt[:,1])-l1fit[1])^2+(mean(igt[:,2])-l1fit[2])^2)

        l1max = getKDEMax(getBelief(fg,:l1))
        maxerror = sqrt((mean(igt[:,1])-l1max[1])^2+(mean(igt[:,2])-l1max[2])^2)

        mykde = getBelief(fg,:l1)
        mynorm = kde!(rand(fit(MvNormal,getVal(fg,:l1)),200))
        kld = min(abs(KernelDensityEstimate.kld(mynorm,mykde)),abs(KernelDensityEstimate.kld(mykde,mynorm)))
    end
    writedlm(scriptHeader*"stats.txt", [meanerror maxerror kld], ",")
end


text:
using LinearAlgebra
using Makie
using Combinatorics
import Quaternions
H = Quaternions.Quaternion

const TOLERANCE = 1e-7

"""
σ(q)

Sends a point on a 3-sphere to a point in the plane x₄=0 with the given point.
This is the stereographic projection of a 3-sphere. S³ ↦ R³
"""
function σ(q)
    Point3f0(q.s, q.v1, q.v2) ./ (1 - q.v3)
end

"""
π(q)

Sends a point on a unit 3-sphere to a point on a unit 2-sphere with the given
point. S³ ↦ S²
"""
function π(q)
    x₁ = q.s
    x₂ = q.v1
    x₃ = q.v2
    x₄ = q.v3
    Point3f0(2(x₁*x₃ + x₂*x₄), 2(x₂*x₃ - x₁*x₄), x₁^2+x₂^2-x₃^2-x₄^2)
end

"""
S¹action(α, q)

Performs a group action corresponding to moving along the circumference of a
circle with the given angle and the point.
"""
function S¹action(α, q)
    z₁ = Complex(q.s, q.v1)
    z₂ = Complex(q.v2, q.v3)
    αz₁, αz₂ = exp(im * α) * z₁, exp(im * α) * z₂
    H(real(αz₁), imag(αz₁), real(αz₂), imag(αz₂))
end

"""
rotate3D(point, q)

Rotates a point in the 3D space with the given point and the unit quaternion.
"""
function rotate3D(point, q)
    p = H(0, point[1], point[2], point[3])
    R = conj(q) * p * q
    Point3f0(R.v1, R.v2, R.v3)
end

"""
get_center(A, B, C)

Finds the center point of the fiber circle under stereographic projection
with the given 3 points on the circle circumference.
"""
function get_center(A, B, C)
    a = norm(B - C)
    b = norm(A - C)
    c = norm(A - B)
    numerator = a^2 * (b^2 + c^2 - a^2) * A +
                b^2 * (a^2 + c^2 - b^2) * B +
                c^2 * (a^2 + b^2 - c^2) * C
    denominator = a^2 * (b^2 + c^2 - a^2) +
                  b^2 * (a^2 + c^2 - b^2) +
                  c^2 * (a^2 + b^2 - c^2)
    numerator / denominator
end

"""
get_fiber(point, segments, samples; r=0.025)

Calculates a torus of revolution for building a surface in a specific way with
the given point in the base space, the number of segments, the number of
samples and the radius of the smaller circle in the torus of revolution.
"""
function get_fiber(point; segments = 30, r = 0.025)
    # Find 3 points on the circle
    A = σ(S¹action(pi / 6, point))
    B = σ(S¹action(pi / 4, point))
    C = σ(S¹action(pi / 3, point))
    # The fiber circle center
    Q = get_center(A, B, C)
    # The bigger radius
    R = norm(Q - A)
    # Get the normal to the plane containing the points
    n = cross(A - Q, B - Q)
    n = n / norm(n)
    # The initial normal to the circle
    i = [0.0, 0.0, 1.0]
    # The axis of rotation
    u = cross(n, i)
    u = u / norm(u)
    # The angle of rotation
    angle = acos(dot(n, i)) / 2.0
    q = H(cos(angle), sin(angle)*u[1], sin(angle)*u[2], sin(angle)*u[3])
    # Construct a torus of revolution grid
    fiber = Array{Point3f0}(undef, segments, segments)
    for i in 1:segments
        for j in 1:segments
            longitude = j * 2pi / (segments - 1)
            latitude = i * 2pi / (segments - 1)
            x₁ = Q[1] + (R + r * cos(longitude)) * cos(latitude) / R
            x₂ = Q[2] + (R + r * cos(longitude)) * sin(latitude) / R
            x₃ = Q[3] + r * sin(longitude) / R
            fiber[i, j] = rotate3D(Point3f0(x₁, x₂, x₃), q)
        end
    end
    fiber
end

"""
compressed_σ(q)

Sends a point on a 3-sphere to a point in the compressed plane x₄=0 with the
given point. This is the stereographic projection of a 3-sphere. S³ ↦ R³
"""
function compressed_σ(q)
    sigmoid(x) = exp(x) / (1 + exp(x))
    2 .* (sigmoid.(σ(q)) .- 0.5)
end

"""
get_rotation_axis(i, n)

Calculates a unit quaternion corresponding to the axis of a 3D rotation with
the given initial and end unit vectors.
"""
function get_rotation_axis(i, n)
    u = normalize(cross(n, i))
    # The angle of rotation
    dotproduct = dot(n, i)
    if isapprox(dotproduct, 1, atol=TOLERANCE)
        θ = acos(1)
    elseif isapprox(dotproduct, -1, atol=TOLERANCE)
        θ = acos(-1)
    else
        θ = acos(dotproduct)
    end
    H(cos(θ/2), sin(θ/2)*u[1], sin(θ/2)*u[2], sin(θ/2)*u[3])
end

"""
get_cylinder(A, B; segments = 30, r = 0.025)

Calculates a cylinder with the given 2 points, the number of segments and the
radius. The points correspond to the beginning and the end of the cylinder.
"""
function get_cylinder(A, B; segments = 30, r = 0.025)
    cylinder = Array{Point3f0}(undef, segments, segments)
    i = [0.0, 0.0, 1.0]
    n = normalize(B - A)
    q = get_rotation_axis(i, n)
    for i in 1:segments
        for j in 1:segments
            α = (j-1) * 2pi / (segments - 1)
            x₁ = r * sin(α)
            x₂ = r * cos(α)
            x₃ = (i-1) * norm(B - A) / (segments - 1)
            cylinder[i, j] = (rotate3D(Point3f0(x₁, x₂, x₃), q) + A)
        end
    end
    cylinder
end

"""
get_torus(A, B, C; segments = 30, r = 0.025)

Calculates a section of a 2-torus with the given 3 points, the number of
segments and the radius of the smaller circle. The points correspond to the
beginning, middle and the end of the arc of the bigger circle.
"""
function get_torus(A, B, C; segments = 30, r = 0.025)
    torus = Array{Point3f0}(undef, segments, segments)
    # The fiber circle center
    Q = get_center(A, B, C)
    # The bigger radius
    R = norm(Q - A)
    # Find the arc sweep angle
    QA = normalize(A - Q)
    QC = normalize(C - Q)
    dotproduct = dot(QA, QC)
    if isapprox(dotproduct, 1, atol=TOLERANCE)
        α = acos(1)
    elseif isapprox(dotproduct, -1, atol=TOLERANCE)
        α = acos(-1)
    else
        α = acos(dotproduct)
    end
    # The normal to the initial circle
    i = [0.0, 0.0, 1.0]
    # Get the normal to the rotated circle
    n = normalize(normalize(cross(A - Q, B - Q)) +
                  normalize(cross(A - Q, C - Q)) +
                  normalize(cross(B - Q, C - Q)))
    # The first rotation versor
    q = get_rotation_axis(i, n)
    # The vector pointing to the positive direction of the x axis on the
    # plane of the rotated circle
    i = rotate3D([1.0, 0.0, 0.0], q)
    # The second rotation versor
    q′ = get_rotation_axis(i, QA)
    # Construct a torus of revolution
    for i in 1:segments
        for j in 1:segments
            longitude = (j-1) * 2pi / (segments - 1)
            latitude = (i-1) * α / (segments - 1)
            x₁ = (R + r * cos(longitude)) * cos(latitude)
            x₂ = (R + r * cos(longitude)) * sin(latitude)
            x₃ = r * sin(longitude)
            torus[i, j] = rotate3D(Point3f0(x₁, x₂, x₃), q * q′) + Q
        end
    end
    torus
end

"""
get_manifold(points, segments, radius)

Calculates a grid of points in order to build a surface with the given 3 points
that are on a circle, the number of segments and the radius of torus/cylinder.
The points should be quaternions.
"""
function get_manifold(points, segments, radius)
    # Use stereographic projection to send the points to R³
    A, B, C = map((x) -> compressed_σ(x), points)
    # If the points are collinear, then build a cylinder
    if isapprox(norm(cross(A - B, C - B)), 0, atol = TOLERANCE)
        manifold = get_cylinder(A, C, segments = segments, r = radius)
        # manifold = get_torus(A, B, C, segments = segments, r = radius)
    # And if they are not collinear, then build a torus of revolution
    else
        # manifold = get_torus(A, B, C, segments = segments, r = radius)
        manifold = get_cylinder(A, C, segments = segments, r = radius)
    end
    manifold
end

"""
get_vertices()

Returns the versors corresponding to the 24 vertices of a 24-cell inscribed
inside of a 3-sphere.
"""
function get_vertices()
    set = []
    set1 = [H(getindex.(Ref([-1/2, 1/2]), 1 .+ digits(i-1; base=2, pad=4))...)
            for i=1:2^4]
    set2 = [H(i...) for i in permutations([1, 0, 0, 0]) |> collect]
    set3 = [H(i...) for i in permutations([-1, 0, 0, 0]) |> collect]
    push!(set, set1..., set2..., set3...)
    unique(set)
end

# The scene object that contains other visual objects
scene = Scene(backgroundcolor = :navyblue, show_axis=false, resolution = (360, 360))
# Use a slider for rotating the base space in an interactive way
sθ, θ = textslider(0:0.01:2pi, "θ", start = pi/4)
sϕ, ϕ = textslider(0:0.01:2pi, "ϕ", start = 0)
sψ, ψ = textslider(0:0.01:2pi, "ψ", start = 0)
# Calculate a unit quaternion as the rotation axis
u = @lift([cos($ϕ) * cos($ψ), cos($ϕ) * sin($ψ), sin($ϕ)])
q = @lift(H(cos($θ), sin($θ) * $u[1], sin($θ) * u[2], sin($θ) * $u[3]))
vertices = get_vertices()
segments = 30
radius = 0.025

edges = []
for i in 1:length(vertices)
    for j in i:length(vertices)
        g = vertices[i]
        h = vertices[j]
        distance = norm(g - h)
        if isapprox(distance, 1)
            m = normalize((g + h) / 2)
            push!(edges, [g, m, h])
        end
    end
end
faces = []
for i in 1:length(vertices)
    for j in i:length(vertices)
        for k in j:length(vertices)
            a = vertices[i]
            b = vertices[j]
            c = vertices[k]
            ab = norm(a - b)
            ac = norm(a - c)
            bc = norm(b - c)
            if isapprox(ab, 1, atol = TOLERANCE) &&
               isapprox(ac, 1, atol = TOLERANCE) &&
               isapprox(bc, 1, atol = TOLERANCE)
                push!(faces, [a, b, c])
            end
        end
    end
end
gold = [0.8314, 0.6863, 0.2157]
saturated = [0.9098, 0.7255, 0.1373]
rosegold = [0.718, 0.431, 0.475]
platinum = [0.898, 0.894, 0.886]
silver = [0.752, 0.752, 0.752]
manifolds = []
for i in 1:length(edges)
    # First use 3-sphere rotations to rotate the points
    rotated = @lift(map((x) -> $q * x, edges[i]))
    manifold = @lift(get_manifold($rotated, segments, radius))
    x = @lift([m[1] for m in $manifold])
    y = @lift([m[2] for m in $manifold])
    z = @lift([m[3] for m in $manifold])
    if i > length(edges) / 2
        color = fill(RGBAf0(gold..., 0.9), segments, segments)
    else
        color = fill(RGBAf0(saturated..., 0.9), segments, segments)
    end
    surface!(scene, x, y, z, color = color, shading = true)
end

for vertex in vertices
    rotated = @lift(compressed_σ($q * vertex))
    x = @lift([$rotated[1]])
    y = @lift([$rotated[2]])
    z = @lift([$rotated[3]])
    color = RGBAf0(rosegold..., 0.9)
    meshscatter!(scene, x, y, z, markersize = 2radius, color = color, shading = true)
end

for i in 1:length(faces)
    a, b, c = faces[i]
    rotateda = @lift(compressed_σ($q * a))
    rotatedb = @lift(compressed_σ($q * b))
    rotatedc = @lift(compressed_σ($q * c))
    l = @lift([($rotateda[1], $rotateda[2], $rotateda[3]),
               ($rotatedb[1], $rotatedb[2], $rotatedb[3]),
               ($rotatedc[1], $rotatedc[2], $rotatedc[3])])
    indices = [1, 2, 3,   2, 3, 1,   3, 1, 2]
    color = [RGBAf0(gold..., rand()), RGBAf0(saturated..., rand()), RGBAf0(rosegold..., rand())]
    mesh!(scene, l, color = color)
end

# update eye position
eye_position, lookat, upvector = Vec3f0(2, 2, 2), Vec3f0(0), Vec3f0(0.0, 0.0, 0.0001)
update_cam!(scene, eye_position, lookat)
scene.center = false # prevent scene from recentering on display

#fullscene = hbox(scene, vbox(sθ, sϕ, sψ), parent = Scene(resolution = (500, 500)))

record(scene, "24-cell.gif") do io
    frames = 90
    for i in 1:frames
        recordframe!(io) # record a new frame
        # animate scene
        ϕ[] = cos(i/frames*pi/2)*2pi
        ψ[] = sin(i/frames*pi/2)*2pi
    end
end

ext: jl

max_issues_repo_name: iamazadi/Porta

max_forks_repo_name: iamazadi/Porta


text:
using HTTP

HTTP.WebSockets.open("ws://127.0.0.1:8081") do ws
    write(ws, "Hello")
    x = readavailable(ws)
    @show x
    println(String(x))
end;

max_stars_repo_name: terasakisatoshi/juliaExer

max_issues_repo_name: terasakisatoshi/juliaExer

max_forks_repo_name: terasakisatoshi/juliaExer


text:
@otestset "Visualization" begin
    # test that this all at least runs
    surf1 = AcceleratedParametricSurface(TestData.beziersurface(), 15)
    surf2 = AcceleratedParametricSurface(TestData.upsidedownbeziersurface(), 15)
    m = csgintersection(leaf(surf1, translation(-0.5, -0.5, 0.0)), csgintersection(leaf(Cylinder(0.3, 5.0)), leaf(surf1, RigidBodyTransform{Float64}(0.0, Float64(π), 0.0, 0.5, -0.5, 0.0))))()
    @test_nowarn Vis.draw(m)
    m = csgintersection(leaf(surf1, translation(-0.5, -0.5, 0.0)), csgintersection(leaf(Cylinder(0.3, 5.0)), leaf(surf2, translation(-0.5, -0.5, 0.0))))()
    @test_nowarn Vis.draw(m)
    m = csgunion(leaf(Cylinder(0.5, 3.0)), (leaf(Sphere(1.0))))()
    @test_nowarn Vis.draw(m)
    m = csgintersection(leaf(Cylinder(0.5, 3.0)), (leaf(Sphere(1.0))))()
    @test_nowarn Vis.draw(m)
    m = csgdifference(leaf(Cylinder(0.5, 3.0)), leaf(Sphere(1.0)))()
    @test_nowarn Vis.draw(m)
end # testset Visualization


text:
# ---
# title: 790. Domino and Tromino Tiling
# id: problem790
# author: Indigo
# date: 2021-01-29
# difficulty: Medium
# categories: Dynamic Programming
# link: <https://leetcode.com/problems/domino-and-tromino-tiling/description/>
# hidden: true
# ---
#
# We have two types of tiles: a 2x1 domino shape, and an "L" tromino shape.
# These shapes may be rotated.
#
#
#
#     XX  <- domino
#
#     XX  <- "L" tromino
#     X
#
#
# Given N, how many ways are there to tile a 2 x N board? **Return your answer
# modulo 10^9 + 7**.
#
# (In a tiling, every square must be covered by a tile. Two tilings are
# different if and only if there are two 4-directionally adjacent cells on the
# board such that exactly one of the tilings has both squares occupied by a
# tile.)
#
#
#
#     **Example:**
#     Input: 3
#     Output: 5
#     Explanation:
#     The five different ways are listed below, different letters indicates different tiles:
#     XYZ XXZ XYY XXY XYY
#     XYZ YYZ XZZ XYY XXY
#
# **Note:**
#
#   * N  will be in range `[1, 1000]`.
#
#
#
#
## @lc code=start
using LeetCode

function num_tilings(N::Int)
    f0, f1= 1, 1
    g0, g1 = 0, 1
    for i in 1:N-1
        f0, f1, g0, g1 = f1, f0 + f1 + 2 * g0, g1, g1 + f1
    end
    f1
end
## @lc code=end


text:
using Test
using DFTK
import DFTK: mpi_sum

function run_scf_and_compare(T, basis, ref_evals, ref_etot;
                             n_ignored=0, test_tol=1e-6, scf_tol=1e-6, test_etot=true, kwargs...)
    n_kpt = length(ref_evals)
    n_bands = length(ref_evals[1])
    kpt_done = zeros(Bool, n_kpt)

    scfres = self_consistent_field(basis; tol=scf_tol, n_bands=n_bands, kwargs...)
    for (ik, ik_global) in enumerate(basis.krange_thisproc)
        @test eltype(scfres.eigenvalues[ik]) == T
        @test eltype(scfres.ψ[ik]) == Complex{T}
        # println(ik_global, "  ", abs.(ref_evals[ik] - scfres.eigenvalues[ik][1:n_bands]))
    end
    for (ik, ik_global) in enumerate(basis.krange_thisproc)
        # Ignore last few bands, because these eigenvalues are hardest to converge
        # and typically a bit random and unstable in the LOBPCG
        diff = abs.(ref_evals[ik_global] - scfres.eigenvalues[ik][1:n_bands])
        @test maximum(diff[1:n_bands - n_ignored]) < test_tol

        kpt_done[ik_global] = true
    end
    @test mpi_sum(sum(kpt_done), basis.comm_kpts) == n_kpt

    @test eltype(scfres.energies.total) == T
    test_etot && (@test scfres.energies.total ≈ ref_etot atol=test_tol)

    scfres
end


text:
mutable struct ClampBoundOptim{V} <: Flux.Optimise.AbstractOptimiser
  os::Vector{Any}
  lb::V
  ub::V
end
ClampBoundOptim(lb,ub,o...) = ClampBoundOptim{typeof(lb)}(Any[o...],lb,ub)

Flux.@forward ClampBoundOptim.os Base.getindex, Base.first, Base.last, Base.lastindex, Base.push!, Base.setindex!
Flux.@forward ClampBoundOptim.os Base.iterate

Base.getindex(c::ClampBoundOptim, i::AbstractArray) = ClampBoundOptim(c.lb,c.ub,c.os[i]...)

function Flux.Optimise.apply!(o::ClampBoundOptim{V}, x, Δ) where V
  for opt in o.os
    Δ = Flux.Optimise.apply!(opt, x, Δ)
  end
  return Δ
end

function Flux.Optimise.update!(opt::ClampBoundOptim{V}, x, x̄) where V
  x̄r = Flux.Optimise.ArrayInterface.restructure(x, x̄) # address some cases where Zygote's
                                          # output are not mutable, see #1510
  x .-= Flux.Optimise.apply!(opt, x, x̄r)
  x .= map(i -> clamp(x[i], opt.lb[i], opt.ub[i]), 1:length(x))
end

load_model(chain::Flux.Chain, T::DataType=Float32) = (Flux.destructure(chain)..., get_bounds(chain,T)...)
load_model(chain::FastChain, T::DataType=Float32) = (initial_params(chain), chain, get_bounds(chain,T)...)
load_model(model, T::DataType=Float32) = (initial_params(model), model, get_bounds(model,T)...)

function loss_collocation(p)
    cost = zero(first(p))
    for i in 1:size(du,2)
      _du = dudt2(@view(u[:,i]),p)
      dui = @view du[:,i]
      cost += sum(abs2,dui .- _du)
    end
    sqrt(cost)
end

function optimize(chain, loss, cb, opt, train_dl, epochs=1, T::DataType=Float32, AD=GalacticOptim.AutoZygote())
  pp, model, lb, ub = load_model(chain, T)

  println("--------------- optimize ---------------")
  println("# training samples:         $(length(train_dl))")
  println("# parameters:               $(length(pp))")
  println("typeof(p):                  $(typeof(pp))")
  println("# epochs:                   $(epochs)")
  println("# lb:                       $(length(lb))")
  println("# ub:                       $(length(ub))")
  println("typeof(lb):                 $(typeof(lb))")

  pp = T.(pp)
  lb = T.(lb)
  ub = T.(ub)

  # mycb = LTC.MyCallback(T, cb, epochs, length(train_dl))
  train_dlnc = epochs > 1 ? ncycle(train_dl, epochs) : train_dl

  # bboptimize(θ->loss(θ,model,first(train_dlnc)...)[1]; SearchRange = [(lb[i],ub[i]) for i in 1:length(lb)], NumDimensions = length(lb))

  f = (θ,p,x,y) -> loss(θ,model,x,y)
  optfun = GalacticOptim.OptimizationFunction(f, AD)
  #optfunc = GalacticOptim.instantiate_function(optfun, pp, AD, nothing)
  optprob = GalacticOptim.OptimizationProblem(optfun, pp, lb=lb, ub=ub,
                                grad = true, hess = false, #sparse = true,
                                parallel=ModelingToolkit.MultithreadedForm()
                                )

  sol = GalacticOptim.solve(optprob, opt, train_dlnc; cb)

  # optfun = GalacticOptim.OptimizationFunction(f, AD)
  # optfunc = GalacticOptim.instantiate_function(optfun, sol.u, AD, nothing)
  # optprob = GalacticOptim.OptimizationProblem(optfunc, sol.u, lb=lb, ub=ub,
  #                               grad = true, hess = true, #sparse = true,
  #                               #parallel=ModelingToolkit.MultithreadedForm()
  #                               )

  # sol = GalacticOptim.solve(optprob, BFGS(initial_stepnorm=0.01), train_dlnc, cb = cb)

  for i in 1:length(pp)
    pp[i] == sol[i] && println(i)
  end
  println(pp)
  println(sol)
  sol
end


text:
using CompScienceMeshes
using BEAST

Γ = readmesh(joinpath(@__DIR__,"sphere2.in"))
X = brezzidouglasmarini(Γ)

κ = 1.0
t = Maxwell3D.singlelayer(wavenumber=κ)
E = Maxwell3D.planewave(direction=ẑ, polarization=x̂, wavenumber=κ)
e = (n × E) × n

@hilbertspace j
@hilbertspace k
efie = @discretise t[k,j]==e[k]  j∈X k∈X
u = gmres(efie)

include("utils/postproc.jl")
include("utils/plotresults.jl")


text:
using ModelingToolkit, StaticArrays, LinearAlgebra
using DiffEqBase, SparseArrays
using Test
using NonlinearSolve
using ModelingToolkit: value

canonequal(a, b) = isequal(simplify(a), simplify(b))

# Define some variables
@parameters t σ ρ β
@variables x y z

function test_nlsys_inference(name, sys, vs, ps)
    @testset "NonlinearSystem construction: $name" begin
        @test Set(states(sys)) == Set(value.(vs))
        @test Set(parameters(sys)) == Set(value.(ps))
    end
end

# Define a nonlinear system
eqs = [0 ~ σ * (y - x),
    0 ~ x * (ρ - z) - y,
    0 ~ x * y - β * z]
@named ns = NonlinearSystem(eqs, [x, y, z], [σ, ρ, β], defaults = Dict(x => 2))
@test eval(toexpr(ns)) == ns
test_nlsys_inference("standard", ns, (x, y, z), (σ, ρ, β))
@test begin
    f = eval(generate_function(ns, [x, y, z], [σ, ρ, β])[2])
    du = [0.0, 0.0, 0.0]
    f(du, [1, 2, 3], [1, 2, 3])
    du ≈ [1, -3, -7]
end

# Now nonlinear system with only variables
@variables x y z
@parameters σ ρ β

# Define a nonlinear system
eqs = [0 ~ σ * (y - x),
    y ~ x * (ρ - z),
    β * z ~ x * y]
@named ns = NonlinearSystem(eqs, [x, y, z], [σ, ρ, β])
jac = calculate_jacobian(ns)
@testset "nlsys jacobian" begin
    @test canonequal(jac[1, 1], σ * -1)
    @test canonequal(jac[1, 2], σ)
    @test canonequal(jac[1, 3], 0)
    @test canonequal(jac[2, 1], ρ - z)
    @test canonequal(jac[2, 2], -1)
    @test canonequal(jac[2, 3], x * -1)
    @test canonequal(jac[3, 1], y)
    @test canonequal(jac[3, 2], x)
    @test canonequal(jac[3, 3], -1 * β)
end
nlsys_func = generate_function(ns, [x, y, z], [σ, ρ, β])
jac_func = generate_jacobian(ns)
f = @eval eval(nlsys_func)

# Intermediate calculations
a = y - x
# Define a nonlinear system
eqs = [0 ~ σ * a,
    0 ~ x * (ρ - z) - y,
    0 ~ x * y - β * z]
@named ns = NonlinearSystem(eqs, [x, y, z], [σ, ρ, β])
nlsys_func = generate_function(ns, [x, y, z], [σ, ρ, β])
nf = NonlinearFunction(ns)
jac = calculate_jacobian(ns)

@test ModelingToolkit.jacobian_sparsity(ns).colptr == sparse(jac).colptr
@test ModelingToolkit.jacobian_sparsity(ns).rowval == sparse(jac).rowval

jac = generate_jacobian(ns)

sH = calculate_hessian(ns)
@test getfield.(ModelingToolkit.hessian_sparsity(ns), :colptr) ==
      getfield.(sparse.(sH), :colptr)
@test getfield.(ModelingToolkit.hessian_sparsity(ns), :rowval) ==
      getfield.(sparse.(sH), :rowval)

prob = NonlinearProblem(ns, ones(3), ones(3))
sol = solve(prob, NewtonRaphson())
@test sol.u[1] ≈ sol.u[2]

@test_throws ArgumentError NonlinearProblem(ns, ones(4), ones(3))

@variables u F s a
eqs1 = [
    0 ~ σ * (y - x) + F,
    0 ~ x * (ρ - z) - u,
    0 ~ x * y - β * z,
    0 ~ x + y - z - u,
]

lorenz = name -> NonlinearSystem(eqs1, [x, y, z, u, F], [σ, ρ, β], name = name)
lorenz1 = lorenz(:lorenz1)
@test_throws ArgumentError NonlinearProblem(lorenz1, zeros(5))
lorenz2 = lorenz(:lorenz2)
@named connected = NonlinearSystem([s ~ a + lorenz1.x
                                    lorenz2.y ~ s
                                    lorenz1.F ~ lorenz2.u
                                    lorenz2.F ~ lorenz1.u], [s, a], [],
                                   systems = [lorenz1, lorenz2])
@test_nowarn alias_elimination(connected)

# system promotion
using OrdinaryDiffEq
@variables t
D = Differential(t)
@named subsys = convert_system(ODESystem, lorenz1, t)
@named sys = ODESystem([D(subsys.x) ~ subsys.x + subsys.x], t, systems = [subsys])
sys = structural_simplify(sys)
u0 = [subsys.x => 1, subsys.z => 2.0]
prob = ODEProblem(sys, u0, (0, 1.0), [subsys.σ => 1, subsys.ρ => 2, subsys.β => 3])
sol = solve(prob, Rodas5())
@test sol[subsys.x] + sol[subsys.y] - sol[subsys.z] ≈ sol[subsys.u]
@test_throws ArgumentError convert_system(ODESystem, sys, t)

@parameters t σ ρ β
@variables x y z

# Define a nonlinear system
eqs = [0 ~ σ * (y - x),
    0 ~ x * (ρ - z) - y,
    0 ~ x * y - β * z]
@named ns = NonlinearSystem(eqs, [x, y, z], [σ, ρ, β])
np = NonlinearProblem(ns, [0, 0, 0], [1, 2, 3], jac = true, sparse = true)
@test calculate_jacobian(ns, sparse = true) isa SparseMatrixCSC

# issue #819
@testset "Combined system name collisions" begin
    function makesys(name)
        @parameters a
        @variables x f

        NonlinearSystem([0 ~ -a * x + f], [x, f], [a]; name)
    end

    function issue819()
        sys1 = makesys(:sys1)
        sys2 = makesys(:sys1)
        @test_throws ArgumentError NonlinearSystem([sys2.f ~ sys1.x, sys1.f ~ 0], [], [],
                                                   systems = [sys1, sys2], name = :foo)
    end
    issue819()
end

# issue #1115
@testset "Extending a NonlinearSystem with no iv" begin
    @parameters a b
    @variables x y
    eqs1 = [
        0 ~ a * x,
    ]
    eqs2 = [
        0 ~ b * y,
    ]

    @named sys1 = NonlinearSystem(eqs1, [x], [a])
    @named sys2 = NonlinearSystem(eqs2, [y], [b])
    @named sys3 = extend(sys1, sys2)

    @test isequal(union(Set(parameters(sys1)), Set(parameters(sys2))),
                  Set(parameters(sys3)))
    @test isequal(union(Set(states(sys1)), Set(states(sys2))), Set(states(sys3)))
    @test isequal(union(Set(equations(sys1)), Set(equations(sys2))), Set(equations(sys3)))
end

# observed variable handling
@variables t x(t) RHS(t)
@parameters τ
@named fol = NonlinearSystem([0 ~ (1 - x) / τ], [x], [τ]; observed = [RHS ~ (1 - x) / τ])
@test isequal(RHS, @nonamespace fol.RHS)
RHS2 = RHS
@unpack RHS = fol
@test isequal(RHS, RHS2)

# issue #1358
@variables t
@variables v1(t) v2(t) i1(t) i2(t)
eq = [v1 ~ sin(2pi * t)
      v1 - v2 ~ i1
      v2 ~ i2
      i1 ~ i2]
@named sys = ODESystem(eq)
@test length(equations(structural_simplify(sys))) == 0

#1504
let
    @variables u[1:4]

    eqs = [u[1] ~ 1,
        u[2] ~ 1,
        u[3] ~ 1,
        u[4] ~ 1]

    sys = NonlinearSystem(eqs, collect(u[1:4]), Num[], defaults = Dict([]), name = :test)
    prob = NonlinearProblem(sys, ones(length(sys.states)))

    sol = NonlinearSolve.solve(prob, NewtonRaphson())

    @test sol[u] ≈ ones(4)
end


text:
include("linearsystem.jl")
using BenchmarkTools

function tbt(δ)
    u22 = 1.0 - 1.0/δ
    b2 = 2.0 - (1.0 + δ)/δ
    x2 = b2/u22
    e2 = x2-1

    b1 = 1 + δ
    x1 = (b1-x2)/δ
    e1 = x1 - 1
    return e1, e2
end

function tbtp(δ)
    u22 = 1.0 - δ
    b2 =  (1.0 + δ) - 2.0*δ
    x2 = b2/u22
    e2 = x2-1

    b1 = 2
    x1 = (b1-x2)
    e1 = x1 - 1
    return e1, e2
end

δ = eps(1.0)/4

e1,e2 = tbt(δ)
println((e1,e2))
e1,e2 = tbtp(δ)
println((e1,e2))

# Test LU factorization
#A = [1.0 2 3; 4 5 6; 7 8 9]
#A = rand(1000,1000)
#B = deepcopy(A)
#@time L, U, p = lu!(A)
#L*U - B[p, :]

# Test Cholesky factorization
A = rand(20,20)
C = A'*A
D = deepcopy(C)
@time L = chol(C)
#@benchmark L = chol3!(D)

#@benchmark L2 = cholesky(C)

# test gaussian elimination with partial pivot
n = 10
A = rand(n,n)
x = rand(n)
b = A*x
@time xt = gepp!(A,b)
err = norm(xt-x)
println(err)

A = [2.0 -1.0 0; -1.0 2.0 -1.0; 0 -1.0 2.0]
x = rand(3)
b = A*x
x0 = rand(3)
@time x1, it = gaussseidel(A, b, x0)
err = norm(x1-x)
println(err)
println(x1)
println(it)


text:
using JLD2

ENV["GKSwstype"] = "nul"
using Plots
using Printf
using Statistics
using CUDA

using Oceananigans
using Oceananigans.Advection
using Oceananigans.AbstractOperations
using Oceananigans.OutputWriters
using Oceananigans.Grids
using Oceananigans.Fields
using Oceananigans.Forcings

include("Bickley.jl")

using .Bickley

make_name(Nh, advection) = "oceananigans_unstable_bickley_jet_Nh$(Nh)_$(typeof(advection).name.wrapper)"

function run(; Nh = 128,
               output_time_interval = 2,
               stop_time = 200,
               arch = CPU(),
               ν = 0,
               advection = WENO5())

    name = make_name(Nh, advection)

    grid = RegularCartesianGrid(size=(Nh, Nh, 1),
                                x = (-2π, 2π), y=(-2π, 2π), z=(0, 1),
                                topology = (Periodic, Periodic, Bounded))

    model = IncompressibleModel(architecture = arch,
                                 timestepper = :RungeKutta3,
                                   advection = advection,
                                        grid = grid,
                                     tracers = :c,
                                     closure = IsotropicDiffusivity(ν=ν, κ=ν),
                                    buoyancy = nothing)

    # ** Initial conditions **
    #
    # u, v: Large-scale jet + vortical perturbations
    #    c: Sinusoid

    # Parameters
    ϵ = 0.1 # perturbation magnitude
    ℓ = 0.5 # Gaussian width
    k = 0.5 # Sinusoidal wavenumber

    # Total initial conditions
    uᵢ(x, y, z) = Bickley.U(y, grid.Ly) + ϵ * Bickley.ũ(x, y, ℓ, k)
    vᵢ(x, y, z) = ϵ * Bickley.ṽ(x, y, ℓ, k)
    cᵢ(x, y, z) = Bickley.C(y, grid.Ly)

    set!(model, u=uᵢ, v=vᵢ, c=cᵢ)

    # Subtract off mean velocity for comparison with GeophysicalFlows
    DeviceArray = arch isa CPU ? Array : CuArray
    U = DeviceArray(zeros(1, 1, 1))
    mean!(U, interior(model.velocities.u))
    model.velocities.u.data.parent .-= U

    progress(sim) = @info(@sprintf("Iter: %d, time: %.1f, Δt: %.3f, max|u|: %.2f",
                                   sim.model.clock.iteration, sim.model.clock.time,
                                   sim.Δt.Δt, maximum(abs, u.data.parent)))

    wizard = TimeStepWizard(cfl=0.1, Δt=1e-4, max_change=1.1, max_Δt=10.0)

    simulation = Simulation(model, Δt=wizard, stop_time=stop_time,
                            iteration_interval=10, progress=progress)

    # Output: primitive fields + computations
    u, v, w, c = primitives = merge(model.velocities, model.tracers)

    ζ   = ComputedField(∂x(v) - ∂y(u))

    save_grid = (file, model) -> file["serialized/grid"] = model.grid

    simulation.output_writers[:fields] =
        JLD2OutputWriter(model, merge(model.velocities, model.tracers, (ζ=ζ,)),
                                schedule = TimeInterval(output_time_interval),
                                init = save_grid,
                                prefix = name * "_fields",
                                field_slicer = nothing,
                                force = true)

    @info "Running a simulation of an unstable Bickley jet with $(Nh)² degrees of freedom..."

    start_time = time_ns()

    run!(simulation)

    run_time = (time_ns() - start_time) * 1e-9

    @info "The simulation with Nh = $Nh ran for $run_time seconds"
    @show cost = run_time / (model.clock.iteration * Nh^2)

    return name
end

using Oceananigans.Grids: AbstractGrid
using Oceananigans.Fields: offset_data
import Oceananigans.Fields: Field

Field(loc::Tuple, grid::AbstractGrid, raw_data::Array) = Field(loc, CPU(), grid, nothing, offset_data(raw_data, grid, loc))
field_from_file(file, loc, name, iter) = Field(loc, file["serialized/grid"], file["timeseries/$name/$iter"])

function visualize(name, contours=false)
    @info "Making a fun movie about an unstable Bickley jet..."

    file = jldopen(name * "_fields.jld2")

    iterations = parse.(Int, keys(file["timeseries/t"]))
    grid = file["serialized/grid"]

    xu, yu, zu = nodes((Face, Cell, Cell), grid)
    xζ, yζ, zζ = nodes((Face, Face, Cell), grid)
    xc, yc, zc = nodes((Cell, Cell, Cell), grid)

    anim = @animate for (i, iteration) in enumerate(iterations)

        @info "    Plotting frame $i from iteration $iteration..."

        t = file["timeseries/t/$iteration"]

        ζ = field_from_file(file, (Face, Face, Cell), :ζ, iteration)
        c = field_from_file(file, (Cell, Cell, Cell), :c, iteration)

        ζi = interior(ζ)[:, :, 1]
        ci = interior(c)[:, :, 1]

        kwargs = Dict(
                      :aspectratio => 1,
                      :linewidth => 0,
                      :colorbar => :none,
                      :ticks => nothing,
                      :clims => (-1, 1),
                      :xlims => (-grid.Lx/2, grid.Lx/2),
                      :ylims => (-grid.Ly/2, grid.Ly/2)
                     )

        contours && (kwargs[:levels] = range(-1, 1, length=31))
        plotter = contours ? contourf : heatmap

        ζ_plot = plotter(xζ, yζ, clamp.(ζi, -1, 1)'; color = :balance, kwargs...)
        c_plot = plotter(xc, yc, clamp.(ci, -1, 1)'; color = :thermal, kwargs...)

        ζ_title = @sprintf("ζ at t = %.1f", t)
        c_title = @sprintf("c at t = %.1f", t)

        plot(ζ_plot, c_plot, title = [ζ_title c_title], size = (4000, 2000))
    end

    gif(anim, name * ".gif", fps = 8)

    return nothing
end

for Nh in (32, 64, 128, 256, 512)
    name = run(Nh=Nh, arch=GPU(), advection=WENO5())
    visualize(name)
end


text:
struct MountainCarParams{T} <:Any where {T<:Real}
    ucoeff::T # action coefficient for acceleration
    g::T # gravity coeff
    h::T # cosine frequency parameter

    MountainCarParams() = new{Float64}(0.001, 0.0025, 3.)
    MountainCarParams(T::Type) = new{T}(T(0.001), T(0.0025), T(3.))
	MountainCarParams(T::Type, ucoeff, g, h) = new{T}(T(ucoeff), T(g), T(h))
end

function random_mountaincar_params()
	ubase = 0.001
	u = rand(rng, Uniform(0.8*ubase, 1.2*ubase))  # coefficient for action force
	# fix these to keep the problem around the same difficulty
	g = 0.0025   	# "gravity" parameter
	h = 3.  		# mountain frequency parameter  (fixed because it changes the state width. Need to recompute xlims)

	params = MountainCarParams(T, u, g, h)
	return params
end

function mountaincar_finitetime(;randomize=false, maxT=5000, Atype=:Discrete, droptime=true, stochastic_start=false)
	if randomize
		params = random_mountaincar_params()
	else
		params = MountainCarParams()
	end
	return create_finitetime_mountaincar(params; maxT=maxT, Atype=Atype, droptime=droptime, stochastic_start=stochastic_start)
end

function create_finitetime_mountaincar(params::MountainCarParams; maxT=5000, Atype=:Discrete, droptime=true, stochastic_start=false)
	X = zeros((2,2))
    X[1,:] .= [-1.2, 0.5]       # x range
	X[2,:] .= [-0.07, 0.07]     # xDot range
    S = ([0. maxT],				# time range
		X)

	if Atype==:Discrete
		A = 1:3
	else
		A = [-1.0 1.0]
    end
    dt = 1.0
	if droptime
		X = S[2]
		function get_outcome1(s,a,params,dt,maxT,stochastic_start)
			t, x, r, γ = mountaincar_step!(s,a, params, dt, maxT, stochastic_start)
			s = (t,x)
			return s,x,r,γ
		end
		p = (s,a)->get_outcome1(s,a,params,dt,maxT,stochastic_start)
		function mcd0_obs()
			t,x = mc_sample_initial(stochastic_start)
			return (t,x), x
		end
		d0 = mcd0_obs
	else
		X = S
		function get_outcome2(s,a,params,dt,maxT,stochastic_start)
			t, x, r, γ = mountaincar_step!(s,a, params, dt, maxT,stochastic_start)
			s = (t,x)
			return s,s,r,γ
		end
		p = (s,a)->get_outcome2(s,a,params,dt,maxT,stochastic_start)
		function mcd0()
			s = mc_sample_initial(stochastic_start)
			return s,s
		end
		d0 = mcd0
    end

	meta = Dict{Symbol,Any}()
    meta[:minreward] = -1.0
    meta[:maxreward] = -1.0
    meta[:minreturn] = -ceil(maxT)
    meta[:maxreturn] = -80
    meta[:stochastic] = false
    meta[:minhorizon] = 80
    meta[:maxhorizon] = ceil(Int, maxT)
	meta[:discounted] = false
	meta[:episodes] = 100
	meta[:threshold] = -150  # maybe this should depend on environment params

	render = (state,clearplot=false)->mountaincarplot(state, params)

	m = SequentialProblem(S,X,A,p,d0,meta, render)
	return m
end

function mountaincar_sim!(state, u, params)
	x, xDot = state
	xDot = xDot + params.ucoeff * u - params.g * cos(params.h * x)
	xDot = clamp(xDot, -0.07, 0.07)
	x += xDot

	if x < -1.2
		x = -1.2
		xDot = 0.
	end

	state[1:2] .= x, xDot
end

function mcget_torque(action::Int)
    if action <= 0 || action > 3
        error("Action needs to be an integer in [1, 3]")
    end
    u = 0.0
    u = (Float64(action) - 2.0)
    return u
end

function mcget_torque(action::Float64)
    u = clamp(action, -1., 1.)
    return u
end

function mc_sample_initial(stochastic_start)
	t = 0.0

	x = zeros(2)
	x[1] = -0.5
	x[2] = 0.

	if stochastic_start
		x[1] = rand(Uniform(-0.6, -0.4))
	end

    return t, x
end

function mountaincar_step!(state, action, params, dt, maxT, stochastic_start)
    u = mcget_torque(action)
    t, x = state
	mountaincar_sim!(x, u, params)
	t += 1

    reward = -1.0
    γ = 1.0
	done = mc_terminal(t, x, maxT)
	if done
        γ = 0.0
        t, x = mc_sample_initial(stochastic_start)
	end

	return t, x, reward, γ
end

function mc_terminal(t, state, maxT)::Bool
	goalcond = state[1] ≥ 0.5
	timecond = t ≥ maxT
	done = goalcond | timecond

	return done
end

@userplot MountainCarPlot
@recipe function f(ap::MountainCarPlot)
	state, params = ap.args
	t, (x, xvel) = state
	h = params.h
	y = sin(h * x)

	xpts = range(-1.2, stop=0.5, length=50)
	ypts = map(x->sin(h*x), xpts)

	legend := false
	xlims := (-1.2, 0.5)
	ylims := (min(ypts...), max(ypts...))
	grid := false
	ticks := nothing
	foreground_color := :white
	aspect_ratio := 1.

	# mountain
	@series begin
		seriestype := :line
		linecolor := :black
		linewidth := 2

		xpts, ypts
	end

	# car
	@series begin
		seriestype := :shape
		linecolor := nothing
		seriescolor := :blue
		aspect_ratio := 1.

		circle_shape(x, y, 0.05)
	end
end


text:
# Configuration for nearly neutral simuation
function dfe_mixd( x::Int64 )
  dfe_mixed( x::Int64; adv_probability=0.01,
    alpha_disadv=dfe_disadv_alpha, beta_disadv=dfe_disadv_beta,
    alpha_adv=dfe_adv_alpha, beta_adv=dfe_adv_beta )
end

const nn_simtype = 1    # nearly neutral infinite alleles
type_str = "M2"
const mu_list_flag = false
@everywhere const N_list     = [25,50,100,200,400,800,1600,3200,6400]     # popsize
const N_mu_list= [0.5,1.0,2.0,4.0]   # Mutation rate as a multiple of 1.0/N
const L = 1           # number of loci
const ngens = 400000      # Generations after burn-in
const fix_minimum = 0.45  # Fraction of population needed to "fix"
const burn_in= 2.0    # generations of burn_in as a multiple of 1/mu
dfe = dfe_mixd
dfe_str = "dfe_mixd advprob=0.01 alpha_adv=1.0 alpha_disadv=0.2 beta_adv=0.01 beta_disadv=0.5"
dfe_adv_prob=0.01
dfe_adv_alpha=1.0
dfe_disadv_alpha=0.2
dfe_adv_beta=0.01
dfe_disadv_beta=0.5

const use_poplist=false   # if true, save a list of populations as a check of correctness

max_stars_repo_name: um-tech-evolution/nearly-neutral

max_issues_repo_name: um-tech-evolution/nearly-neutral

max_forks_repo_name: um-tech-evolution/nearly-neutral


text:
@test Nemo.accuracy_bits(GibbsTypePriors.Cnk(6, 5, 0.5, 200)) ≤ 200
@test Nemo.accuracy_bits(GibbsTypePriors.Cnk(6, 5, 0.5, 2000)) ≥ 200
@test Nemo.accuracy_bits(GibbsTypePriors.Cnk_rec(600, 5, 0.5, 200)) ≤ 200
@test Nemo.accuracy_bits(GibbsTypePriors.Cnk_rec(600, 5, 0.5, 2000)) ≥ 200

max_stars_repo_name: konkam/GibbsTypePriors

max_issues_repo_name: konkam/GibbsTypePriors

max_forks_repo_name: konkam/GibbsTypePriors


text:
export Jacobian

"""
`Jacobian(A::AbstractOperator,x)`

Shorthand constructor:

`jacobian(A::AbstractOperator,x)`

Returns the jacobian of `A` evaluated at `x` (which in the case of a `LinearOperator` is `A` itself).

```julia
julia> Jacobian(DFT(10),randn(10))
ℱ  ℝ^10 -> ^ℂ10

julia> Jacobian(Sigmoid((10,)),randn(10))
J(σ)  ℝ^10 -> ℝ^10

```

"""
struct Jacobian{T <: NonLinearOperator, X<:AbstractArray} <: LinearOperator
	A::T
	x::X
end

#Jacobian of LinearOperator
Jacobian(L::T,x::X) where {T<:LinearOperator,X<:AbstractArray} = L
#Jacobian of Scale
Jacobian(S::T2,x::AbstractArray) where {T,L,T2<:Scale{T,L}} = Scale(S.coeff,Jacobian(S.A,x))
##Jacobian of DCAT
function Jacobian(H::DCAT{N,L,P1,P2},b) where {N,L,P1,P2}
  x = b.x
	A = ()
	c = 0
    for (k, idx) in enumerate(H.idxD)
		if length(idx) == 1
            A = (A...,jacobian(H.A[k],x[idx]))
        else
            xx = ([x[i] for i in idx]...,)
            A = (A...,jacobian(H.A[k],xx))
        end
	end
	DCAT(A,H.idxD,H.idxC)
end
#Jacobian of HCAT
function Jacobian(H::HCAT{N,L,P,C},b::ArrayPartition) where {N,L,P,C,D}
  x = b.x
	A = ()
    for (k, idx) in enumerate(H.idxs)
		if length(idx) == 1
            A = (A...,jacobian(H.A[k],x[idx]))
        else
            xx = ArrayPartition([x[i] for i in idx]...,)
            A = (A...,jacobian(H.A[k],xx))
        end
	end
	HCAT(A,H.idxs,H.buf)
end
#Jacobian of VCAT
function Jacobian(V::VCAT{N,L,P,C},x::AbstractArray) where {N,L,P,C,D}
    JJ = ([Jacobian(a,x) for a in V.A]...,)
    VCAT(JJ, V.idxs, V.buf)
end
#Jacobian of Compose
function Jacobian(L::Compose, x::X) where {X<:AbstractArray}
	Compose(Jacobian.(L.A,(x,L.buf...)),L.buf)
end

function Jacobian(L::Compose, x::X) where {N,X<:NTuple{N,AbstractArray}}
	Compose(Jacobian.(L.A,(x,L.buf...)),L.buf)
end
#Jacobian of Reshape
Jacobian(R::Reshape{N,L},x::AbstractArray) where {N,L} = Reshape(Jacobian(R.A,x),R.dim_out)
#Jacobian of Sum
Jacobian(S::Sum{K,C,D},x::D) where {K,C,D} =
Sum(([Jacobian(a,x) for a in S.A]...,),S.bufC,S.bufD)
#Jacobian of Transpose
Jacobian(T::Transpose{A}, x::AbstractArray) where {A <: AbstractOperator} = T
#Jacobian of BroadCast
Jacobian(B::A, x::AbstractArray) where {A <: BroadCast} = BroadCast(Jacobian(B.A,x),B.dim_out)
#Jacobian of AffineAdd
Jacobian(B::A, x) where {A <: AffineAdd} = Jacobian(B.A,x)

# Properties

fun_name(L::Jacobian)  = "J("*fun_name(L.A)*")"
size(L::Jacobian) = size(L.A,1), size(L.A,2)

domainType(L::Jacobian) = domainType(L.A)
codomainType(L::Jacobian) = codomainType(L.A)


text:
using Plots

x = -3:0.1:3
f(x)=x^2
y=f.(x)

plotlyjs()
#plot(x,y,label="line")
#scatter!(x,y,label="points")

globaltemperatures = [14.4,14.5,14.8,15.2,15.5,15.8]
numpirates = [45000,20000,15000,5000,400,17]

plot(numpirates,globaltemperatures,legend=false)
scatter!(numpirates,globaltemperatures,legend=false)
xflip!()

xlabel!("Number of pirates")
ylabel!("Global Temperature (C)")
title!("Influence of pirate population on global warming")

max_stars_repo_name: luca-mm/Bachelor-Thesis

max_issues_repo_name: luca-mm/Bachelor-Thesis

max_forks_repo_name: luca-mm/Bachelor-Thesis


text:
function check_generator_bounds(env::AdmmEnv, xbar)
    generators = env.data.generators
    gen_start = env.model.gen_mod.gen_start

    ngen = length(generators)

    if env.use_gpu
        pgmin = CuArray{Float64}(undef, ngen)
        pgmax = CuArray{Float64}(undef, ngen)
        qgmin = CuArray{Float64}(undef, ngen)
        qgmax = CuArray{Float64}(undef, ngen)
    else
        pgmin = Array{Float64}(undef, ngen)
        pgmax = Array{Float64}(undef, ngen)
        qgmin = Array{Float64}(undef, ngen)
        qgmax = Array{Float64}(undef, ngen)
    end

    Pmin = Float64[generators[g].Pmin for g in 1:ngen]
    Pmax = Float64[generators[g].Pmax for g in 1:ngen]
    Qmin = Float64[generators[g].Qmin for g in 1:ngen]
    Qmax = Float64[generators[g].Qmax for g in 1:ngen]
    copyto!(pgmin, Pmin)
    copyto!(pgmax, Pmax)
    copyto!(qgmin, Qmin)
    copyto!(qgmax, Qmax)

    max_viol_real = 0.0
    max_viol_reactive = 0.0

    for g=1:ngen
        pidx = gen_start + 2*(g-1)
        qidx = gen_start + 2*(g-1) + 1

        real_err = 0.0
        reactive_err = 0.0
        CUDA.@allowscalar real_err = max(max(0.0, xbar[pidx] - pgmax[g]), max(0.0, pgmin[g] - xbar[pidx]))
        CUDA.@allowscalar reactive_err = max(max(0.0, xbar[qidx] - qgmax[g]), max(0.0, qgmin[g] - xbar[qidx]))

        max_viol_real = (max_viol_real < real_err) ? real_err : max_viol_real
        max_viol_reactive = (max_viol_reactive < reactive_err) ? reactive_err : max_viol_reactive
    end

    return max_viol_real, max_viol_reactive
end

function check_voltage_bounds(env::AdmmEnv, xbar)

    buses = env.data.buses
    nbus = length(buses)
    bus_start = env.model.bus_start

    max_viol = 0.0

    for b=1:nbus
        bidx = bus_start + 2*(b-1)
        err = 0.0
        CUDA.@allowscalar err = max(max(0.0, xbar[bidx] - buses[b].Vmax^2), max(0.0, buses[b].Vmin^2 - xbar[bidx]))
        max_viol = (max_viol < err) ? err : max_viol
    end

    return max_viol
end

function check_power_balance_violation(env::AdmmEnv, xbar)
    data = env.data
    model = env.model
    gen_start, line_start, bus_start, YshR, YshI = model.gen_mod.gen_start, model.line_start, model.bus_start, model.YshR, model.YshI

    baseMVA = data.baseMVA
    nbus = length(data.buses)

    max_viol_real = 0.0
    max_viol_reactive = 0.0
    for b=1:nbus
        real_err = 0.0
        reactive_err = 0.0
        for g in data.BusGenerators[b]
            CUDA.@allowscalar real_err += xbar[gen_start + 2*(g-1)]
            CUDA.@allowscalar reactive_err += xbar[gen_start + 2*(g-1)+1]
        end

        CUDA.@allowscalar real_err -= (env.model.Pd[b] / baseMVA)
        CUDA.@allowscalar reactive_err -= (env.model.Qd[b] / baseMVA)

        #real_err -= (data.buses[b].Pd / baseMVA)
        #reactive_err -= (data.buses[b].Qd / baseMVA)

        for l in data.FromLines[b]
            CUDA.@allowscalar real_err -= xbar[line_start + 4*(l-1)]
            CUDA.@allowscalar reactive_err -= xbar[line_start + 4*(l-1) + 1]
        end

        for l in data.ToLines[b]
            CUDA.@allowscalar real_err -= xbar[line_start + 4*(l-1) + 2]
            CUDA.@allowscalar reactive_err -= xbar[line_start + 4*(l-1) + 3]
        end

        CUDA.@allowscalar real_err -= YshR[b] * xbar[bus_start + 2*(b-1)]
        CUDA.@allowscalar reactive_err += YshI[b] * xbar[bus_start + 2*(b-1)]

        max_viol_real = (max_viol_real < abs(real_err)) ? abs(real_err) : max_viol_real
        max_viol_reactive = (max_viol_reactive < abs(reactive_err)) ? abs(reactive_err) : max_viol_reactive
    end

    return max_viol_real, max_viol_reactive
end

function get_branch_bus_index(data::OPFData; use_gpu=false)
    lines = data.lines
    BusIdx = data.BusIdx
    nline = length(lines)

    brBusIdx = [ x for l=1:nline for x in (BusIdx[lines[l].from], BusIdx[lines[l].to]) ]

    if use_gpu
        cu_brBusIdx = CuArray{Int}(undef, 2*nline)
        copyto!(cu_brBusIdx, brBusIdx)
        return cu_brBusIdx
    else
        return brBusIdx
    end
end

function init_solution!(env::AdmmEnv, sol::SolutionTwoLevel, ybus::Ybus, rho_pq, rho_va)
    data, model = env.data, env.model
    gen_start, line_start, bus_start = model.gen_mod.gen_start, model.line_start, model.bus_start

    lines = data.lines
    buses = data.buses
    BusIdx = data.BusIdx
    ngen = length(data.generators)
    nline = length(data.lines)
    nbus = length(data.buses)

    YffR = ybus.YffR; YffI = ybus.YffI
    YttR = ybus.YttR; YttI = ybus.YttI
    YftR = ybus.YftR; YftI = ybus.YftI
    YtfR = ybus.YtfR; YtfI = ybus.YtfI

    # Init
    fill!(sol.x_curr, 0.0)
    fill!(sol.xbar_curr, 0.0)
    fill!(sol.z_outer, 0.0)
    fill!(sol.z_curr, 0.0)
    fill!(sol.z_prev, 0.0)
    fill!(sol.l_curr, 0.0)
    fill!(sol.lz, 0.0)
    fill!(sol.rho, 0.0)
    fill!(sol.rp, 0.0)
    fill!(sol.rp_old, 0.0)
    fill!(sol.rd, 0.0)
    fill!(sol.Ax_plus_By, 0.0)

    u_curr = view(sol.x_curr, 1:model.nvar_u)
    v_curr = view(sol.x_curr, model.nvar_u+1:model.nvar)
    rho_u = view(sol.rho, 1:model.nvar_u)
    rho_v = view(sol.rho, model.nvar_u+1:model.nvar)

    rho_u .= rho_pq
    rho_v .= rho_pq
    rho_v[bus_start:end] .= rho_va

    for g=1:ngen
        pg_idx = gen_start + 2*(g-1)
        CUDA.@allowscalar sol.xbar_curr[pg_idx] = 0.5*(data.generators[g].Pmin + data.generators[g].Pmax)
        CUDA.@allowscalar sol.xbar_curr[pg_idx+1] = 0.5*(data.generators[g].Qmin + data.generators[g].Qmax)
    end

    fill!(sol.wRIij, 0.0)
    for l=1:nline
        fr_idx = BusIdx[lines[l].from]
        to_idx = BusIdx[lines[l].to]

        wij0 = (buses[fr_idx].Vmax^2 + buses[fr_idx].Vmin^2) / 2
        wji0 = (buses[to_idx].Vmax^2 + buses[to_idx].Vmin^2) / 2
        wR0 = sqrt(wij0 * wji0)

        u_pij_idx = line_start + 8*(l-1)
        v_pij_idx = line_start + 4*(l-1)
        CUDA.@allowscalar v_curr[v_pij_idx] = u_curr[u_pij_idx] = YffR[l] * wij0 + YftR[l] * wR0
        CUDA.@allowscalar v_curr[v_pij_idx+1] = u_curr[u_pij_idx+1] = -YffI[l] * wij0 - YftI[l] * wR0
        CUDA.@allowscalar v_curr[v_pij_idx+2] = u_curr[u_pij_idx+2] = YttR[l] * wji0 + YtfR[l] * wR0
        CUDA.@allowscalar v_curr[v_pij_idx+3] = u_curr[u_pij_idx+3] = -YttI[l] * wji0 - YtfI[l] * wR0

        rho_u[u_pij_idx+4:u_pij_idx+7] .= rho_va

        CUDA.@allowscalar sol.wRIij[2*(l-1)+1] = wR0
        CUDA.@allowscalar sol.wRIij[2*l] = 0.0
    end

    for b=1:nbus
        CUDA.@allowscalar sol.xbar_curr[bus_start + 2*(b-1)] = (buses[b].Vmax^2 + buses[b].Vmin^2) / 2
        CUDA.@allowscalar sol.xbar_curr[bus_start + 2*(b-1)+1] = 0.0
    end

    return
end

function update_xbar(env::AdmmEnv, u, v, xbar, zu, zv, lu, lv, rho_u, rho_v)
    data = env.data
    ngen = length(data.generators)
    nline = length(data.lines)
    nbus = length(data.buses)

    model = env.model
    gen_start = model.gen_mod.gen_start
    line_start = model.line_start
    bus_start = model.bus_start
    FrStart = model.FrStart
    ToStart = model.ToStart
    FrIdx = model.FrIdx
    ToIdx = model.ToIdx

    gen_end = gen_start + 2*ngen - 1
    xbar[gen_start:gen_end] .= (lu[gen_start:gen_end] .+ rho_u[gen_start:gen_end].*(u[gen_start:gen_end] .+ zu[gen_start:gen_end]) .+
                                lv[gen_start:gen_end] .+ rho_v[gen_start:gen_end].*(v[gen_start:gen_end] .+ zv[gen_start:gen_end])) ./
                               (rho_u[gen_start:gen_end] .+ rho_v[gen_start:gen_end])

    ul_cur = line_start
    vl_cur = line_start
    for j=1:nline
        xbar[vl_cur:vl_cur+3] .= (lu[ul_cur:ul_cur+3] .+ rho_u[ul_cur:ul_cur+3].*(u[ul_cur:ul_cur+3] .+ zu[ul_cur:ul_cur+3]) .+
                                  lv[vl_cur:vl_cur+3] .+ rho_v[vl_cur:vl_cur+3].*(v[vl_cur:vl_cur+3] .+ zv[vl_cur:vl_cur+3])) ./
                                 (rho_u[ul_cur:ul_cur+3] .+ rho_v[vl_cur:vl_cur+3])
        ul_cur += 8
        vl_cur += 4
    end

    for b=1:nbus
        wi_sum = 0.0
        ti_sum = 0.0
        rho_wi_sum = 0.0
        rho_ti_sum = 0.0

        for j=FrStart[b]:FrStart[b+1]-1
            u_pij_idx = line_start + 8*(FrIdx[j]-1)
            wi_sum += lu[u_pij_idx+4] + rho_u[u_pij_idx+4]*(u[u_pij_idx+4] + zu[u_pij_idx+4])
            ti_sum += lu[u_pij_idx+6] + rho_u[u_pij_idx+6]*(u[u_pij_idx+6] + zu[u_pij_idx+6])
            rho_wi_sum += rho_u[u_pij_idx+4]
            rho_ti_sum += rho_u[u_pij_idx+6]
        end
        for j=ToStart[b]:ToStart[b+1]-1
            u_pij_idx = line_start + 8*(ToIdx[j]-1)
            wi_sum += lu[u_pij_idx+5] + rho_u[u_pij_idx+5]*(u[u_pij_idx+5] + zu[u_pij_idx+5])
            ti_sum += lu[u_pij_idx+7] + rho_u[u_pij_idx+7]*(u[u_pij_idx+7] + zu[u_pij_idx+7])
            rho_wi_sum += rho_u[u_pij_idx+5]
            rho_ti_sum += rho_u[u_pij_idx+7]
        end

        bus_cur = bus_start + 2*(b-1)
        wi_sum += lv[bus_cur] + rho_v[bus_cur]*(v[bus_cur] + zv[bus_cur])
        rho_wi_sum += rho_v[bus_cur]
        ti_sum += lv[bus_cur+1] + rho_v[bus_cur+1]*(v[bus_cur+1] + zv[bus_cur+1])
        rho_ti_sum += rho_v[bus_cur+1]
        xbar[bus_cur] = wi_sum / rho_wi_sum
        xbar[bus_cur+1] = ti_sum / rho_ti_sum
    end
end

function update_xbar_generator_kernel(n::Int, gen_start::Int, u, v, xbar, zu, zv, lu, lv, rho_u, rho_v)
    tx = threadIdx().x + (blockDim().x * (blockIdx().x - 1))

    if tx <= n
        i = gen_start + 2*(tx - 1)
        @inbounds begin
            xbar[i] = (lu[i] + rho_u[i]*(u[i] + zu[i]) + lv[i] + rho_v[i]*(v[i] + zv[i])) / (rho_u[i] + rho_v[i])
            xbar[i+1] = (lu[i+1] + rho_u[i+1]*(u[i+1] + zu[i+1]) + lv[i+1] + rho_v[i+1]*(v[i+1] + zv[i+1])) / (rho_u[i+1] + rho_v[i+1])
        end
    end

    return
end

function update_xbar_branch_kernel(n::Int, line_start::Int, u, v, xbar, zu, zv, lu, lv, rho_u, rho_v)
    tx = threadIdx().x + (blockDim().x * (blockIdx().x - 1))

    if tx <= n
        ul_cur = line_start + 8*(tx-1)
        vl_cur = line_start + 4*(tx-1)
        @inbounds begin
            xbar[vl_cur] = (lu[ul_cur] + rho_u[ul_cur]*(u[ul_cur] + zu[ul_cur]) +
                            lv[vl_cur] + rho_v[vl_cur]*(v[vl_cur] + zv[vl_cur])) /
                           (rho_u[ul_cur] + rho_v[vl_cur])
            xbar[vl_cur+1] = (lu[ul_cur+1] + rho_u[ul_cur+1]*(u[ul_cur+1] + zu[ul_cur+1]) +
                              lv[vl_cur+1] + rho_v[vl_cur+1]*(v[vl_cur+1] + zv[vl_cur+1])) /
                             (rho_u[ul_cur+1] + rho_v[vl_cur+1])
            xbar[vl_cur+2] = (lu[ul_cur+2] + rho_u[ul_cur+2]*(u[ul_cur+2] + zu[ul_cur+2]) +
                             lv[vl_cur+2] + rho_v[vl_cur+2]*(v[vl_cur+2] + zv[vl_cur+2])) /
                            (rho_u[ul_cur+2] + rho_v[vl_cur+2])
            xbar[vl_cur+3] = (lu[ul_cur+3] + rho_u[ul_cur+3]*(u[ul_cur+3] + zu[ul_cur+3]) +
                              lv[vl_cur+3] + rho_v[vl_cur+3]*(v[vl_cur+3] + zv[vl_cur+3])) /
                             (rho_u[ul_cur+3] + rho_v[vl_cur+3])
        end
    end

    return
end

function update_xbar_bus_kernel(n::Int, line_start::Int, bus_start::Int, FrStart, FrIdx, ToStart, ToIdx, u, v, xbar, zu, zv, lu, lv, rho_u, rho_v)
    b = threadIdx().x + (blockDim().x * (blockIdx().x - 1))

    if b <= n
        wi_sum = 0.0
        ti_sum = 0.0
        rho_wi_sum = 0.0
        rho_ti_sum = 0.0

        @inbounds begin
            for j=FrStart[b]:FrStart[b+1]-1
                u_pij_idx = line_start + 8*(FrIdx[j]-1)
                wi_sum += lu[u_pij_idx+4] + rho_u[u_pij_idx+4]*(u[u_pij_idx+4] + zu[u_pij_idx+4])
                ti_sum += lu[u_pij_idx+6] + rho_u[u_pij_idx+6]*(u[u_pij_idx+6] + zu[u_pij_idx+6])
                rho_wi_sum += rho_u[u_pij_idx+4]
                rho_ti_sum += rho_u[u_pij_idx+6]
            end
            for j=ToStart[b]:ToStart[b+1]-1
                u_pij_idx = line_start + 8*(ToIdx[j]-1)
                wi_sum += lu[u_pij_idx+5] + rho_u[u_pij_idx+5]*(u[u_pij_idx+5] + zu[u_pij_idx+5])
                ti_sum += lu[u_pij_idx+7] + rho_u[u_pij_idx+7]*(u[u_pij_idx+7] + zu[u_pij_idx+7])
                rho_wi_sum += rho_u[u_pij_idx+5]
                rho_ti_sum += rho_u[u_pij_idx+7]
            end

            bus_cur = bus_start + 2*(b-1)
            wi_sum += lv[bus_cur] + rho_v[bus_cur]*(v[bus_cur] + zv[bus_cur])
            rho_wi_sum += rho_v[bus_cur]
            ti_sum += lv[bus_cur+1] + rho_v[bus_cur+1]*(v[bus_cur+1] + zv[bus_cur+1])
            rho_ti_sum += rho_v[bus_cur+1]
            xbar[bus_cur] = wi_sum / rho_wi_sum
            xbar[bus_cur+1] = ti_sum / rho_ti_sum
        end
    end

    return
end

# TODO: This function is not used.
function update_lu(data::OPFData, gen_start, line_start, bus_start, u, xbar, zu, l, rho)
    lines = data.lines
    BusIdx = data.BusIdx
    ngen = length(data.generators)
    nline = length(data.lines)

    gen_end = gen_start + 2*ngen - 1
    l[gen_start:gen_end] .+= rho[gen_start:gen_end] .* (u[gen_start:gen_end] .- xbar[gen_start:gen_end] .+ zu[gen_start:gen_end])

    ul_cur = line_start
    xl_cur = line_start
    for j=1:nline
        fr_idx = bus_start + 2*(BusIdx[lines[j].from]-1)
        to_idx = bus_start + 2*(BusIdx[lines[j].to]-1)

        l[ul_cur:ul_cur+3] .+= rho[ul_cur:ul_cur+3] .* (u[ul_cur:ul_cur+3] .- xbar[xl_cur:xl_cur+3] .+ zu[ul_cur:ul_cur+3])
        l[ul_cur+4] += rho[ul_cur+4] * (u[ul_cur+4] - xbar[fr_idx] + zu[ul_cur+4])
        l[ul_cur+5] += rho[ul_cur+5] * (u[ul_cur+5] - xbar[to_idx] + zu[ul_cur+5])
        l[ul_cur+6] += rho[ul_cur+6] * (u[ul_cur+6] - xbar[fr_idx+1] + zu[ul_cur+6])
        l[ul_cur+7] += rho[ul_cur+7] * (u[ul_cur+7] - xbar[to_idx+1] + zu[ul_cur+7])
        ul_cur += 8
        xl_cur += 4
    end
end

function update_zu(env::AdmmEnv, u, xbar, z, l, rho, lz, beta)
    data = env.data
    lines = data.lines
    BusIdx = data.BusIdx
    ngen = length(data.generators)
    nline = length(data.lines)

    model = env.model
    gen_start, line_start, bus_start = model.gen_mod.gen_start, model.line_start, model.bus_start

    gen_end = gen_start + 2*ngen - 1
    z[gen_start:gen_end] .= (-(lz[gen_start:gen_end] .+ l[gen_start:gen_end] .+ rho[gen_start:gen_end].*(u[gen_start:gen_end] .- xbar[gen_start:gen_end]))) ./ (beta .+ rho[gen_start:gen_end])

    ul_cur = line_start
    xl_cur = line_start
    for j=1:nline
        fr_idx = bus_start + 2*(BusIdx[lines[j].from]-1)
        to_idx = bus_start + 2*(BusIdx[lines[j].to]-1)

        z[ul_cur:ul_cur+3] .= (-(lz[ul_cur:ul_cur+3] .+ l[ul_cur:ul_cur+3] .+ rho[ul_cur:ul_cur+3].*(u[ul_cur:ul_cur+3] .- xbar[xl_cur:xl_cur+3]))) ./ (beta .+ rho[ul_cur:ul_cur+3])
        z[ul_cur+4] = (-(lz[ul_cur+4] + l[ul_cur+4] + rho[ul_cur+4]*(u[ul_cur+4] - xbar[fr_idx]))) / (beta + rho[ul_cur+4])
        z[ul_cur+5] = (-(lz[ul_cur+5] + l[ul_cur+5] + rho[ul_cur+5]*(u[ul_cur+5] - xbar[to_idx]))) / (beta + rho[ul_cur+5])
        z[ul_cur+6] = (-(lz[ul_cur+6] + l[ul_cur+6] + rho[ul_cur+6]*(u[ul_cur+6] - xbar[fr_idx+1]))) / (beta + rho[ul_cur+6])
        z[ul_cur+7] = (-(lz[ul_cur+7] + l[ul_cur+7] + rho[ul_cur+7]*(u[ul_cur+7] - xbar[to_idx+1]))) / (beta + rho[ul_cur+7])
        ul_cur += 8
        xl_cur += 4
    end
end

function update_zu_generator_kernel(n::Int, gen_start::Int, u, xbar, z, l, rho, lz, beta)
    tx = threadIdx().x + (blockDim().x * (blockIdx().x - 1))

    if tx <= n
        i = gen_start  + 2*(tx - 1)
        @inbounds begin
            z[i] = (-(lz[i] + l[i] + rho[i]*(u[i] - xbar[i]))) / (beta + rho[i])
            z[i+1] = (-(lz[i+1] + l[i+1] + rho[i+1]*(u[i+1] - xbar[i+1]))) / (beta + rho[i+1])
        end
    end

    return
end

function update_zu_branch_kernel(n::Int, line_start::Int, bus_start::Int, brBusIdx, u, xbar, z, l, rho, lz, beta)
    tx = threadIdx().x + (blockDim().x * (blockIdx().x - 1))

    if tx <= n
        ul_cur = line_start + 8*(tx - 1)
        xl_cur = line_start + 4*(tx - 1)

        @inbounds begin
            fr_idx = bus_start + 2*(brBusIdx[2*(tx-1)+1] - 1)
            to_idx = bus_start + 2*(brBusIdx[2*tx] - 1)

            z[ul_cur] = (-(lz[ul_cur] + l[ul_cur] + rho[ul_cur]*(u[ul_cur] - xbar[xl_cur]))) / (beta + rho[ul_cur])
            z[ul_cur+1] = (-(lz[ul_cur+1] + l[ul_cur+1] + rho[ul_cur+1]*(u[ul_cur+1] - xbar[xl_cur+1]))) / (beta + rho[ul_cur+1])
            z[ul_cur+2] = (-(lz[ul_cur+2] + l[ul_cur+2] + rho[ul_cur+2]*(u[ul_cur+2] - xbar[xl_cur+2]))) / (beta + rho[ul_cur+2])
            z[ul_cur+3] = (-(lz[ul_cur+3] + l[ul_cur+3] + rho[ul_cur+3]*(u[ul_cur+3] - xbar[xl_cur+3]))) / (beta + rho[ul_cur+3])
            z[ul_cur+4] = (-(lz[ul_cur+4] + l[ul_cur+4] + rho[ul_cur+4]*(u[ul_cur+4] - xbar[fr_idx]))) / (beta + rho[ul_cur+4])
            z[ul_cur+5] = (-(lz[ul_cur+5] + l[ul_cur+5] + rho[ul_cur+5]*(u[ul_cur+5] - xbar[to_idx]))) / (beta + rho[ul_cur+5])
            z[ul_cur+6] = (-(lz[ul_cur+6] + l[ul_cur+6] + rho[ul_cur+6]*(u[ul_cur+6] - xbar[fr_idx+1]))) / (beta + rho[ul_cur+6])
            z[ul_cur+7] = (-(lz[ul_cur+7] + l[ul_cur+7] + rho[ul_cur+7]*(u[ul_cur+7] - xbar[to_idx+1]))) / (beta + rho[ul_cur+7])
        end
    end

    return
end

function update_zv_kernel(n::Int, v, xbar, z, l, rho, lz, beta)
    tx = threadIdx().x + (blockDim().x * (blockIdx().x - 1))

    if tx <= n
        @inbounds begin
            z[tx] = (-(lz[tx] + l[tx] + rho[tx]*(v[tx] - xbar[tx]))) / (beta + rho[tx])
        end
    end

    return
end

function update_l_kernel(n::Int, l, z, lz, beta)
    tx = threadIdx().x + (blockDim().x * (blockIdx().x - 1))

    if tx <= n
        @inbounds begin
            l[tx] = -(lz[tx] + beta*z[tx])
        end
    end

    return
end

function compute_primal_residual_u(env::AdmmEnv, rp_u, u, xbar, z)
    data = env.data
    lines = data.lines
    BusIdx = data.BusIdx
    ngen = length(data.generators)
    nline = length(data.lines)

    model = env.model
    gen_start, line_start, bus_start = model.gen_mod.gen_start, model.line_start, model.bus_start

    gen_end = gen_start + 2*ngen - 1
    rp_u[gen_start:gen_end] .= u[gen_start:gen_end] .- xbar[gen_start:gen_end] .+ z[gen_start:gen_end]

    ul_cur = line_start
    xl_cur = line_start
    for j=1:nline
        fr_idx = bus_start + 2*(BusIdx[lines[j].from]-1)
        to_idx = bus_start + 2*(BusIdx[lines[j].to]-1)

        rp_u[ul_cur:ul_cur+3] .= u[ul_cur:ul_cur+3] .- xbar[xl_cur:xl_cur+3] .+ z[ul_cur:ul_cur+3]
        rp_u[ul_cur+4] = u[ul_cur+4] - xbar[fr_idx] + z[ul_cur+4]
        rp_u[ul_cur+5] = u[ul_cur+5] - xbar[to_idx] + z[ul_cur+5]
        rp_u[ul_cur+6] = u[ul_cur+6] - xbar[fr_idx+1] + z[ul_cur+6]
        rp_u[ul_cur+7] = u[ul_cur+7] - xbar[to_idx+1] + z[ul_cur+7]

        ul_cur += 8
        xl_cur += 4
    end
end

function compute_primal_residual_u_generator_kernel(n::Int, gen_start::Int, rp, u, xbar, z)
    tx = threadIdx().x + (blockDim().x * (blockIdx().x - 1))

    if tx <= n
        i = gen_start + 2*(tx - 1)
        @inbounds begin
            rp[i] = u[i] - xbar[i] + z[i]
            rp[i+1] = u[i+1] - xbar[i+1] + z[i+1]
        end
    end

    return
end

function compute_primal_residual_u_branch_kernel(n::Int, line_start::Int, bus_start::Int, brBusIdx, rp, u, xbar, z)
    tx = threadIdx().x + (blockDim().x * (blockIdx().x - 1))

    if tx <= n
        ul_cur = line_start + 8*(tx - 1)
        xl_cur = line_start + 4*(tx - 1)
        fr_idx = bus_start + 2*(brBusIdx[2*(tx-1)+1]-1)
        to_idx = bus_start + 2*(brBusIdx[2*tx]-1)

        rp[ul_cur] = u[ul_cur] - xbar[xl_cur] + z[ul_cur]
        rp[ul_cur+1] = u[ul_cur+1] - xbar[xl_cur+1] + z[ul_cur+1]
        rp[ul_cur+2] = u[ul_cur+2] - xbar[xl_cur+2] + z[ul_cur+2]
        rp[ul_cur+3] = u[ul_cur+3] - xbar[xl_cur+3] + z[ul_cur+3]
        rp[ul_cur+4] = u[ul_cur+4] - xbar[fr_idx] + z[ul_cur+4]
        rp[ul_cur+5] = u[ul_cur+5] - xbar[to_idx] + z[ul_cur+5]
        rp[ul_cur+6] = u[ul_cur+6] - xbar[fr_idx+1] + z[ul_cur+6]
        rp[ul_cur+7] = u[ul_cur+7] - xbar[to_idx+1] + z[ul_cur+7]
    end

    return
end

function compute_primal_residual_v_kernel(n::Int, rp, v, xbar, z)
    tx = threadIdx().x + (blockDim().x * (blockIdx().x - 1))

    if tx <= n
        rp[tx] = v[tx] - xbar[tx] + z[tx]
    end

    return
end

function vector_difference(n::Int, c, a, b)
    tx = threadIdx().x + (blockDim().x * (blockIdx().x - 1))

    if tx <= n
        c[tx] = a[tx] - b[tx]
    end

    return
end

function update_lz_kernel(n::Int, max_limit::Float64, z, lz, beta)
    tx = threadIdx().x + (blockDim().x * (blockIdx().x - 1))

    if tx <= n
        lz[tx] += max(-max_limit, min(max_limit, beta*z[tx]))
    end

    return
end

# TODO: Not used
function update_rho(rho, rp, rp_old, theta, gamma)
    for i=1:length(rho)
        if abs(rp[i]) > theta*abs(rp_old[i])
            rho[i] = min(gamma*rho[i], 1e24)
        end
    end
end

function admm_solve!(env::AdmmEnv, sol::SolutionTwoLevel; outer_iterlim=10, inner_iterlim=800, scale=1e-4)
    data, par, mod = env.data, env.params, env.model

    # -------------------------------------------------------------------
    # Variables are of two types: u and v
    #   u contains variables for generators and branches and
    #   v contains variables for buses.
    #
    # Variable layout:
    #   u: | (pg,qg)_g | (pij,qij,pji,qji,w_i_ij,w_j_ij,a_i_ij, a_i_ji)_ij |
    #   v: | (pg,qg)_g | (pij,qij,pji,qji)_ij | (w_i,theta_i)_i |
    #   xbar: same as v
    # -------------------------------------------------------------------

    x_curr = sol.x_curr
    xbar_curr = sol.xbar_curr
    z_outer = sol.z_outer
    z_curr = sol.z_curr
    z_prev = sol.z_prev
    l_curr = sol.l_curr
    lz = sol.lz
    rho = sol.rho
    rp = sol.rp
    rd = sol.rd
    rp_old = sol.rp_old
    Ax_plus_By = sol.Ax_plus_By
    wRIij = sol.wRIij

    u_curr = view(x_curr, 1:mod.nvar_u)
    v_curr = view(x_curr, mod.nvar_u+1:mod.nvar)
    zu_curr = view(z_curr, 1:mod.nvar_u)
    zv_curr = view(z_curr, mod.nvar_u+1:mod.nvar)
    lu_curr = view(l_curr, 1:mod.nvar_u)
    lv_curr = view(l_curr, mod.nvar_u+1:mod.nvar)
    lz_u = view(lz, 1:mod.nvar_u)
    lz_v = view(lz, mod.nvar_u+1:mod.nvar)
    rho_u = view(rho, 1:mod.nvar_u)
    rho_v = view(rho, mod.nvar_u+1:mod.nvar)
    rp_u = view(rp, 1:mod.nvar_u)
    rp_v = view(rp, mod.nvar_u+1:mod.nvar)

    nblk_gen = div(mod.gen_mod.ngen, 32, RoundUp)
    nblk_br = mod.nline
    nblk_bus = div(mod.nbus, 32, RoundUp)

    beta = 1e3
    gamma = 6.0 # TODO: not used
    c = 6.0
    theta = 0.8
    sqrt_d = sqrt(mod.nvar_u + mod.nvar_v)
    OUTER_TOL = sqrt_d*(par.outer_eps)

    outer = 0
    inner = 0

    time_gen = time_br = time_bus = 0
    shift_lines = 0
    shmem_size = sizeof(Float64)*(14*mod.n+3*mod.n^2) + sizeof(Int)*(4*mod.n)

    mismatch = Inf
    z_prev_norm = z_curr_norm = Inf

    # Return status
    status = INITIAL

    while outer < outer_iterlim
        outer += 1

        if !env.use_gpu
            z_outer .= z_curr
            z_prev_norm = norm(z_outer)
        else
            @cuda threads=64 blocks=(div(mod.nvar-1, 64)+1) copy_data_kernel(mod.nvar, z_outer, z_curr)
            z_prev_norm = CUDA.norm(z_curr)
            CUDA.synchronize()
        end

        inner = 0
        while inner < inner_iterlim
            inner += 1

            if !env.use_gpu
                z_prev .= z_curr
                rp_old .= rp

                tcpu = generator_kernel_two_level(mod.gen_mod, data.baseMVA, u_curr, xbar_curr, zu_curr, lu_curr, rho_u)
                time_gen += tcpu.time

                #scale = min(scale, (2*1e4) / maximum(abs.(rho_u)))
                if env.use_polar
                    tcpu = @timed auglag_it, tron_it = polar_kernel_two_level_cpu(mod.n, mod.nline, mod.line_start, mod.bus_start, scale,
                                                                                u_curr, xbar_curr, zu_curr, lu_curr, rho_u,
                                                                                shift_lines, env.membuf, mod.YffR, mod.YffI, mod.YftR, mod.YftI,
                                                                                mod.YttR, mod.YttI, mod.YtfR, mod.YtfI, mod.FrBound, mod.ToBound, mod.brBusIdx)
                else
                    tcpu = @timed auglag_it, tron_it = auglag_kernel_cpu(mod.n, mod.nline, inner, par.max_auglag, mod.line_start, par.mu_max,
                                                                        u_curr, v_curr, l_curr, rho,
                                                                        wRIij, env.membuf, mod.YffR, mod.YffI, mod.YftR, mod.YftI,
                                                                        mod.YttR, mod.YttI, mod.YtfR, mod.YtfI, mod.FrBound, mod.ToBound)
                end
                time_br += tcpu.time

                if !env.allow_infeas
                    tcpu = @timed bus_kernel_two_level_cpu(data.baseMVA, mod.nbus, mod.gen_mod.gen_start, mod.line_start, mod.bus_start,
                                                        mod.FrStart, mod.FrIdx, mod.ToStart, mod.ToIdx, mod.GenStart, mod.GenIdx,
                                                        mod.Pd, mod.Qd, v_curr, xbar_curr, zv_curr, lv_curr, rho_v, mod.YshR, mod.YshI)
                else
                    tcpu = @timed bus_kernel_two_level_cpu(data.baseMVA, mod.nbus, mod.gen_mod.gen_start, mod.line_start, mod.bus_start,
                                                        mod.FrStart, mod.FrIdx, mod.ToStart, mod.ToIdx, mod.GenStart, mod.GenIdx,
                                                        mod.Pd, mod.Qd, v_curr, xbar_curr, zv_curr, lv_curr, rho_v, mod.YshR, mod.YshI,
                                                        sol.s_curr, par.rho_sigma)

                end
                time_bus += tcpu.time

                update_xbar(env, u_curr, v_curr, xbar_curr, zu_curr, zv_curr, lu_curr, lv_curr, rho_u, rho_v)

                update_zu(env, u_curr, xbar_curr, zu_curr, lu_curr, rho_u, lz_u, beta)
                zv_curr .= (-(lz_v .+ lv_curr .+ rho_v.*(v_curr .- xbar_curr))) ./ (beta .+ rho_v)

                l_curr .= -(lz .+ beta.*z_curr)

                compute_primal_residual_u(env, rp_u, u_curr, xbar_curr, zu_curr)
                rp_v .= v_curr .- xbar_curr .+ zv_curr

                #=
                if inner > 1
                    update_rho(rho, rp, rp_old, theta, gamma)
                end
                =#

                rd .= z_curr .- z_prev
                Ax_plus_By .= rp .- z_curr

                primres = norm(rp)
                dualres = norm(rd)
                z_curr_norm = norm(z_curr)
                mismatch = norm(Ax_plus_By)
                eps_pri = sqrt_d / (2500*outer)

                if par.verbose > 0
                    if inner == 1 || (inner % 50) == 0
                        @printf("%8s  %8s  %10s  %10s  %10s  %10s  %10s  %10s\n",
                                "Outer", "Inner", "PrimRes", "EpsPrimRes", "DualRes", "||z||", "||Ax+By||", "Beta")
                    end

                    @printf("%8d  %8d  %10.3e  %10.3e  %10.3e  %10.3e  %10.3e  %10.3e\n",
                            outer, inner, primres, eps_pri, dualres, z_curr_norm, mismatch, beta)
                end

                if primres <= eps_pri || dualres <= par.DUAL_TOL
                    break
                end
            else
                @cuda threads=64 blocks=(div(mod.nvar-1, 64)+1) copy_data_kernel(mod.nvar, z_prev, z_curr)
                @cuda threads=64 blocks=(div(mod.nvar-1, 64)+1) copy_data_kernel(mod.nvar, rp_old, rp)
                CUDA.synchronize()

                tgpu = generator_kernel_two_level(mod.gen_mod, data.baseMVA, u_curr, xbar_curr, zu_curr, lu_curr, rho_u)
                # MPI routines to be implemented:
                #  - Broadcast v_curr and l_curr to GPUs.
                #  - Collect u_curr.
                #  - div(nblk_br / number of GPUs, RoundUp)

                time_gen += tgpu.time
                if env.use_polar
                    tgpu = CUDA.@timed @cuda threads=32 blocks=nblk_br shmem=shmem_size polar_kernel_two_level(mod.n, mod.nline, mod.line_start, mod.bus_start, scale,
                                                              u_curr, xbar_curr, zu_curr, lu_curr, rho_u,
                                                              shift_lines, env.membuf, mod.YffR, mod.YffI, mod.YftR, mod.YftI,
                                                              mod.YttR, mod.YttI, mod.YtfR, mod.YtfI, mod.FrBound, mod.ToBound, mod.brBusIdx)
                else
                    tgpu = CUDA.@timed @cuda threads=32 blocks=nblk_br shmem=shmem_size auglag_kernel(mod.n, inner, par.max_auglag, mod.line_start, scale, par.mu_max,
                                                                                u_curr, v_curr, l_curr, rho,
                                                                                wRIij, env.membuf, mod.YffR, mod.YffI, mod.YftR, mod.YftI,
                                                                                mod.YttR, mod.YttI, mod.YtfR, mod.YtfI, mod.FrBound, mod.ToBound)
                end
                time_br += tgpu.time

                if !env.allow_infeas
                    tgpu = CUDA.@timed @cuda threads=32 blocks=nblk_bus bus_kernel_two_level(data.baseMVA, mod.nbus, mod.gen_mod.gen_start, mod.line_start, mod.bus_start,
                                                                        mod.FrStart, mod.FrIdx, mod.ToStart, mod.ToIdx, mod.GenStart, mod.GenIdx,
                                                                        mod.Pd, mod.Qd, v_curr, xbar_curr, zv_curr, lv_curr,
                                                                        rho_v, mod.YshR, mod.YshI)
                else
                    tgpu = CUDA.@timed @cuda threads=32 blocks=nblk_bus bus_kernel_two_level(data.baseMVA, mod.nbus, mod.gen_mod.gen_start, mod.line_start, mod.bus_start,
                                                                        mod.FrStart, mod.FrIdx, mod.ToStart, mod.ToIdx, mod.GenStart, mod.GenIdx,
                                                                        mod.Pd, mod.Qd, v_curr, xbar_curr, zv_curr, lv_curr,
                                                                        rho_v, mod.YshR, mod.YshI,
                                                                        sol.s_curr, par.rho_sigma)

                end
                time_bus += tgpu.time

                # Update xbar.
                @cuda threads=64 blocks=(div(mod.gen_mod.ngen-1, 64)+1) update_xbar_generator_kernel(mod.gen_mod.ngen, mod.gen_mod.gen_start, u_curr, v_curr,
                                               xbar_curr, zu_curr, zv_curr, lu_curr, lv_curr, rho_u, rho_v)
                @cuda threads=64 blocks=(div(mod.nline-1, 64)+1) update_xbar_branch_kernel(mod.nline, mod.line_start, u_curr, v_curr,
                                               xbar_curr, zu_curr, zv_curr, lu_curr, lv_curr, rho_u, rho_v)
                @cuda threads=64 blocks=(div(mod.nbus-1, 64)+1) update_xbar_bus_kernel(mod.nbus, mod.line_start, mod.bus_start, mod.FrStart, mod.FrIdx, mod.ToStart, mod.ToIdx, u_curr, v_curr, xbar_curr, zu_curr, zv_curr, lu_curr, lv_curr, rho_u, rho_v)
                CUDA.synchronize()

                # Update z.
                @cuda threads=64 blocks=(div(mod.gen_mod.ngen-1, 64)+1) update_zu_generator_kernel(mod.gen_mod.ngen, mod.gen_mod.gen_start, u_curr,
                                               xbar_curr, zu_curr, lu_curr, rho_u, lz_u, beta)
                @cuda threads=64 blocks=(div(mod.nline-1, 64)+1) update_zu_branch_kernel(mod.nline, mod.line_start, mod.bus_start, mod.brBusIdx, u_curr, xbar_curr, zu_curr, lu_curr, rho_u, lz_u, beta)
                @cuda threads=64 blocks=(div(mod.nvar_v-1, 64)+1) update_zv_kernel(mod.nvar_v, v_curr, xbar_curr, zv_curr,
                                               lv_curr, rho_v, lz_v, beta)
                CUDA.synchronize()

                # Update multiiplier and residuals.
                @cuda threads=64 blocks=(div(mod.nvar-1, 64)+1) update_l_kernel(mod.nvar, l_curr, z_curr, lz, beta)
                @cuda threads=64 blocks=(div(mod.gen_mod.ngen-1, 64)+1) compute_primal_residual_u_generator_kernel(mod.gen_mod.ngen, mod.gen_mod.gen_start, rp_u, u_curr, xbar_curr, zu_curr)
                @cuda threads=64 blocks=(div(mod.nline-1, 64)+1) compute_primal_residual_u_branch_kernel(mod.nline, mod.line_start, mod.bus_start, mod.brBusIdx, rp_u, u_curr, xbar_curr, zu_curr)
                @cuda threads=64 blocks=(div(mod.nvar_v-1, 64)+1) compute_primal_residual_v_kernel(mod.nvar_v, rp_v, v_curr, xbar_curr, zv_curr)
                @cuda threads=64 blocks=(div(mod.nvar-1, 64)+1) vector_difference(mod.nvar, rd, z_curr, z_prev)
                CUDA.synchronize()

                CUDA.@sync @cuda threads=64 blocks=(div(mod.nvar-1, 64)+1) vector_difference(mod.nvar, Ax_plus_By, rp, z_curr)
                mismatch = CUDA.norm(Ax_plus_By)

                primres = CUDA.norm(rp)
                dualres = CUDA.norm(rd)
                z_curr_norm = CUDA.norm(z_curr)
                eps_pri = sqrt_d / (2500*outer)

                if par.verbose > 0
                    if inner == 1 || (inner % 50) == 0
                        @printf("%8s  %8s  %10s  %10s  %10s  %10s  %10s  %10s  %10s\n",
                                "Outer", "Inner", "PrimRes", "EpsPrimRes", "DualRes", "||z||", "||Ax+By||", "OuterTol", "Beta")
                    end

                    @printf("%8d  %8d  %10.3e  %10.3e  %10.3e  %10.3e  %10.3e  %10.3e  %10.3e\n",
                            outer, inner, primres, eps_pri, dualres, z_curr_norm, mismatch, OUTER_TOL, beta)
                end

                if primres <= eps_pri || dualres <= par.DUAL_TOL
                    break
                end
            end
        end

        if mismatch <= OUTER_TOL
            status = HAS_CONVERGED
            break
        end

        if !env.use_gpu
            lz .+= max.(-par.MAX_MULTIPLIER, min.(par.MAX_MULTIPLIER, beta .* z_curr))
        else
            CUDA.@sync @cuda threads=64 blocks=(div(mod.nvar-1, 64)+1) update_lz_kernel(mod.nvar, par.MAX_MULTIPLIER, z_curr, lz, beta)
        end

        if z_curr_norm > theta*z_prev_norm
            beta = min(c*beta, 1e24)
        end
    end

    # Move x onto host
    hxbar_curr = Vector(xbar_curr)
    sol.objval = sum(data.generators[g].coeff[data.generators[g].n-2]*(data.baseMVA*hxbar_curr[mod.gen_mod.gen_start+2*(g-1)])^2 +
                data.generators[g].coeff[data.generators[g].n-1]*(data.baseMVA*hxbar_curr[mod.gen_mod.gen_start+2*(g-1)]) +
                data.generators[g].coeff[data.generators[g].n]
                for g in 1:mod.gen_mod.ngen)

    if outer >= outer_iterlim
        sol.status = MAXIMUM_ITERATIONS
    else
        sol.status = status
    end

    if par.verbose > 0
        # Test feasibility of global variable xbar:
        pg_err, qg_err = check_generator_bounds(env, xbar_curr)
        vm_err = check_voltage_bounds(env, xbar_curr)
        real_err, reactive_err = check_power_balance_violation(env, xbar_curr)
        @printf(" ** Violations of global variable xbar\n")
        @printf("Real power generator bounds     = %.6e\n", pg_err)
        @printf("Reactive power generator bounds = %.6e\n", qg_err)
        @printf("Voltage bounds                  = %.6e\n", vm_err)
        @printf("Real power balance              = %.6e\n", real_err)
        @printf("Reaactive power balance         = %.6e\n", reactive_err)

        rateA_nviols, rateA_maxviol, rateC_nviols, rateC_maxviol = check_linelimit_violation(data, u_curr)
        @printf(" ** Line limit violations\n")
        @printf("RateA number of violations = %d (%d)\n", rateA_nviols, mod.nline)
        @printf("RateA maximum violation    = %.2f\n", rateA_maxviol)
        @printf("RateC number of violations = %d (%d)\n", rateC_nviols, mod.nline)
        @printf("RateC maximum violation    = %.2f\n", rateC_maxviol)

        @printf(" ** Time\n")
        @printf("Generator = %.2f\n", time_gen)
        @printf("Branch    = %.2f\n", time_br)
        @printf("Bus       = %.2f\n", time_bus)
        @printf("Total     = %.2f\n", time_gen + time_br + time_bus)
        @printf("Objective value = %.6e\n", sol.objval)
    end

    return
end

function admm_rect_gpu_two_level(
    case::String;
    outer_iterlim=10, inner_iterlim=800, rho_pq=400.0, rho_va=40000.0, scale=1e-4,
    use_gpu=false, use_polar=true, allow_infeas=false, rho_sigma=1e8,
    gpu_no=0, verbose=1, outer_eps=2e-4
)
    env = AdmmEnv(
        case, use_gpu, rho_pq, rho_va; use_polar=use_polar, use_twolevel=true,
        allow_infeas=allow_infeas, rho_sigma=rho_sigma,
        gpu_no=gpu_no, verbose=verbose,
    )
    admm_restart!(env, outer_iterlim=outer_iterlim, inner_iterlim=inner_iterlim, scale=scale)
    return env
end