function get_generator_data(data::OPFData; use_gpu=false)
    ngen = length(data.generators)

    if use_gpu
        pgmin = CuArray{Float64}(undef, ngen)
        pgmax = CuArray{Float64}(undef, ngen)
        qgmin = CuArray{Float64}(undef, ngen)
        qgmax = CuArray{Float64}(undef, ngen)
        c2 = CuArray{Float64}(undef, ngen)
        c1 = CuArray{Float64}(undef, ngen)
        c0 = 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)
        c2 = Array{Float64}(undef, ngen)
        c1 = Array{Float64}(undef, ngen)
        c0 = Array{Float64}(undef, ngen)
    end

    Pmin = Float64[data.generators[g].Pmin for g in 1:ngen]
    Pmax = Float64[data.generators[g].Pmax for g in 1:ngen]
    Qmin = Float64[data.generators[g].Qmin for g in 1:ngen]
    Qmax = Float64[data.generators[g].Qmax for g in 1:ngen]
    coeff0 = Float64[data.generators[g].coeff[3] for g in 1:ngen]
    coeff1 = Float64[data.generators[g].coeff[2] for g in 1:ngen]
    coeff2 = Float64[data.generators[g].coeff[1] for g in 1:ngen]
    copyto!(pgmin, Pmin)
    copyto!(pgmax, Pmax)
    copyto!(qgmin, Qmin)
    copyto!(qgmax, Qmax)
    copyto!(c0, coeff0)
    copyto!(c1, coeff1)
    copyto!(c2, coeff2)

    return pgmin,pgmax,qgmin,qgmax,c2,c1,c0
end

function get_bus_data(data::OPFData; use_gpu=false)
    nbus = length(data.buses)

    FrIdx = [l for b=1:nbus for l in data.FromLines[b]]
    ToIdx = [l for b=1:nbus for l in data.ToLines[b]]
    GenIdx = [g for b=1:nbus for g in data.BusGenerators[b]]
    FrStart = accumulate(+, vcat([1], [length(data.FromLines[b]) for b=1:nbus]))
    ToStart = accumulate(+, vcat([1], [length(data.ToLines[b]) for b=1:nbus]))
    GenStart = accumulate(+, vcat([1], [length(data.BusGenerators[b]) for b=1:nbus]))

    Pd = Float64[data.buses[i].Pd for i=1:nbus]
    Qd = Float64[data.buses[i].Qd for i=1:nbus]

    if use_gpu
        cuFrIdx = CuArray{Int}(undef, length(FrIdx))
        cuToIdx = CuArray{Int}(undef, length(ToIdx))
        cuGenIdx = CuArray{Int}(undef, length(GenIdx))
        cuFrStart = CuArray{Int}(undef, length(FrStart))
        cuToStart = CuArray{Int}(undef, length(ToStart))
        cuGenStart = CuArray{Int}(undef, length(GenStart))
        cuPd = CuArray{Float64}(undef, nbus)
        cuQd = CuArray{Float64}(undef, nbus)

        copyto!(cuFrIdx, FrIdx)
        copyto!(cuToIdx, ToIdx)
        copyto!(cuGenIdx, GenIdx)
        copyto!(cuFrStart, FrStart)
        copyto!(cuToStart, ToStart)
        copyto!(cuGenStart, GenStart)
        copyto!(cuPd, Pd)
        copyto!(cuQd, Qd)

        return cuFrStart,cuFrIdx,cuToStart,cuToIdx,cuGenStart,cuGenIdx,cuPd,cuQd
    else
        return FrStart,FrIdx,ToStart,ToIdx,GenStart,GenIdx,Pd,Qd
    end
end

function get_branch_data(data::OPFData; use_gpu=false)
    buses = data.buses
    lines = data.lines
    BusIdx = data.BusIdx
    nline = length(data.lines)
    ybus = Ybus{Array{Float64}}(computeAdmitances(data.lines, data.buses, data.baseMVA; VI=Array{Int}, VD=Array{Float64})...)
    frBound = [ x for l=1:nline for x in (buses[BusIdx[lines[l].from]].Vmin^2, buses[BusIdx[lines[l].from]].Vmax^2) ]
    toBound = [ x for l=1:nline for x in (buses[BusIdx[lines[l].to]].Vmin^2, buses[BusIdx[lines[l].to]].Vmax^2) ]

    if use_gpu
        cuYshR = CuArray{Float64}(undef, length(ybus.YshR))
        cuYshI = CuArray{Float64}(undef, length(ybus.YshI))
        cuYffR = CuArray{Float64}(undef, nline)
        cuYffI = CuArray{Float64}(undef, nline)
        cuYftR = CuArray{Float64}(undef, nline)
        cuYftI = CuArray{Float64}(undef, nline)
        cuYttR = CuArray{Float64}(undef, nline)
        cuYttI = CuArray{Float64}(undef, nline)
        cuYtfR = CuArray{Float64}(undef, nline)
        cuYtfI = CuArray{Float64}(undef, nline)
        cuFrBound = CuArray{Float64}(undef, 2*nline)
        cuToBound = CuArray{Float64}(undef, 2*nline)
        copyto!(cuYshR, ybus.YshR)
        copyto!(cuYshI, ybus.YshI)
        copyto!(cuYffR, ybus.YffR)
        copyto!(cuYffI, ybus.YffI)
        copyto!(cuYftR, ybus.YftR)
        copyto!(cuYftI, ybus.YftI)
        copyto!(cuYttR, ybus.YttR)
        copyto!(cuYttI, ybus.YttI)
        copyto!(cuYtfR, ybus.YtfR)
        copyto!(cuYtfI, ybus.YtfI)
        copyto!(cuFrBound, frBound)
        copyto!(cuToBound, toBound)

        return cuYshR, cuYshI, cuYffR, cuYffI, cuYftR, cuYftI,
               cuYttR, cuYttI, cuYtfR, cuYtfI, cuFrBound, cuToBound
    else
        return ybus.YshR, ybus.YshI, ybus.YffR, ybus.YffI, ybus.YftR, ybus.YftI,
               ybus.YttR, ybus.YttI, ybus.YtfR, ybus.YtfI, frBound, toBound
    end
end

function init_solution!(env::AdmmEnv, sol::SolutionOneLevel, ybus::Ybus, rho_pq, rho_va)
    data, model = env.data, env.model

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

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

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

    sol.rho .= rho_pq
    fill!(sol.u_curr, 0.0)
    fill!(sol.v_curr, 0.0)

    for l=1:nline
        wij0 = (buses[BusIdx[lines[l].from]].Vmax^2 + buses[BusIdx[lines[l].from]].Vmin^2) / 2
        wji0 = (buses[BusIdx[lines[l].to]].Vmax^2 + buses[BusIdx[lines[l].to]].Vmin^2) / 2
        wR0 = sqrt(wij0 * wji0)

        pij_idx = model.line_start + 8*(l-1)
        CUDA.@allowscalar sol.u_curr[pij_idx] = YffR[l] * wij0 + YftR[l] * wR0
        CUDA.@allowscalar sol.u_curr[pij_idx+1] = -YffI[l] * wij0 - YftI[l] * wR0
        CUDA.@allowscalar sol.u_curr[pij_idx+2] = YttR[l] * wji0 + YtfR[l] * wR0
        CUDA.@allowscalar sol.u_curr[pij_idx+3] = -YttI[l] * wji0 - YtfI[l] * wR0
        #=
        u_curr[pij_idx+4] = wij0
        u_curr[pij_idx+5] = wji0
        u_curr[pij_idx+6] = 0.0
        u_curr[pij_idx+7] = 0.0
        =#
        # wRIij[2*(l-1)+1] = wR0
        # wRIij[2*l] = 0.0

        CUDA.@allowscalar sol.v_curr[pij_idx+4] = wij0
        CUDA.@allowscalar sol.v_curr[pij_idx+5] = wji0
        CUDA.@allowscalar sol.v_curr[pij_idx+6] = 0.0
        CUDA.@allowscalar sol.v_curr[pij_idx+7] = 0.0

        sol.rho[pij_idx+4:pij_idx+7] .= rho_va
    end

    sol.l_curr .= 0
    return
end

function copy_data_kernel(n::Int, dest::CuDeviceArray{Float64,1}, src::CuDeviceArray{Float64,1})
    tx = threadIdx().x + (blockDim().x * (blockIdx().x - 1))

    if tx <= n
        dest[tx] = src[tx]
    end
    return
end

function update_multiplier_kernel(n::Int, l_curr::CuDeviceArray{Float64,1},
                                  u_curr::CuDeviceArray{Float64,1},
                                  v_curr::CuDeviceArray{Float64,1},
                                  rho::CuDeviceArray{Float64,1})
    tx = threadIdx().x + (blockDim().x * (blockIdx().x - 1))

    if tx <= n
        l_curr[tx] += rho[tx] * (u_curr[tx] - v_curr[tx])
    end
    return
end

function primal_residual_kernel(n::Int, rp::CuDeviceArray{Float64,1},
                                u_curr::CuDeviceArray{Float64,1},
                                v_curr::CuDeviceArray{Float64,1})
    tx = threadIdx().x + (blockDim().x * (blockIdx().x - 1))

    if tx <= n
        rp[tx] = u_curr[tx] - v_curr[tx]
    end

    return
end

function dual_residual_kernel(n::Int, rd::CuDeviceArray{Float64,1},
                              v_prev::CuDeviceArray{Float64,1},
                              v_curr::CuDeviceArray{Float64,1},
                              rho::CuDeviceArray{Float64,1})
    tx = threadIdx().x + (blockDim().x * (blockIdx().x - 1))

    if tx <= n
        rd[tx] = -rho[tx] * (v_curr[tx] - v_prev[tx])
    end

    return
end

function check_linelimit_violation(data::OPFData, u)
    lines = data.lines
    nline = length(data.lines)
    line_start = 2*length(data.generators) + 1

    rateA_nviols = 0
    rateA_maxviol = 0.0
    rateC_nviols = 0
    rateC_maxviol = 0.0

    for l=1:nline
        pij_idx = line_start + 8*(l-1)
        ij_val = 0.0
        ji_val = 0.0
        CUDA.@allowscalar ij_val = u[pij_idx]^2 + u[pij_idx+1]^2
        CUDA.@allowscalar ji_val = u[pij_idx+2]^2 + u[pij_idx+3]^2

        limit = (lines[l].rateA / data.baseMVA)^2
        if limit > 0 && limit < 1e10
            if ij_val > limit || ji_val > limit
                rateA_nviols += 1
                rateA_maxviol = max(rateA_maxviol, max(ij_val - limit, ji_val - limit))
            end
        end

        limit = (lines[l].rateC / data.baseMVA)^2
        if limit > 0 && limit < 1e10
            if ij_val > limit || ji_val > limit
                rateC_nviols += 1
                rateC_maxviol = max(rateC_maxviol, max(ij_val - limit, ji_val - limit))
            end
        end
    end
    rateA_maxviol = sqrt(rateA_maxviol)
    rateC_maxviol = sqrt(rateC_maxviol)

    return rateA_nviols, rateA_maxviol, rateC_nviols, rateC_maxviol
end

"""
    admm_restart!

This function restarts the ADMM with a given `env::AdmmEnv` containing solutions and all the other parameters.
"""
admm_restart!(env::AdmmEnv; options...) = admm_solve!(env, env.solution; options...)

function admm_solve!(env::AdmmEnv, sol::SolutionOneLevel; iterlim=800, scale=1e-4)
    if env.use_gpu
        CUDA.device!(env.gpu_no)
    end

    data, par, mod = env.data, env.params, env.model

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

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

    it = 0
    time_gen = time_br = time_bus = 0.0

    @time begin
    while it < iterlim
        it += 1

        if !env.use_gpu
            sol.u_prev .= sol.u_curr
            sol.v_prev .= sol.v_curr
            sol.l_prev .= sol.l_curr

            tcpu = generator_kernel(mod.gen_mod, data.baseMVA, sol.u_curr, sol.v_curr, sol.l_curr, sol.rho)
            time_gen += tcpu.time

            if env.use_polar
                tcpu = @timed auglag_it, tron_it = polar_kernel_cpu(mod.n, mod.nline, mod.line_start, scale,
                                                                    sol.u_curr, sol.v_curr, sol.l_curr, sol.rho,
                                                                    shift_lines, env.membuf, mod.YffR, mod.YffI, mod.YftR, mod.YftI,
                                                                    mod.YttR, mod.YttI, mod.YtfR, mod.YtfI, mod.FrBound, mod.ToBound)
            else
                tcpu = @timed auglag_it, tron_it = auglag_kernel_cpu(mod.n, mod.nline, it, par.max_auglag, mod.line_start, par.mu_max,
                                                                     sol.u_curr, sol.v_curr, sol.l_curr, sol.rho,
                                                                     shift_lines, 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

            tcpu = @timed bus_kernel_cpu(data.baseMVA, mod.nbus, mod.gen_mod.gen_start, mod.line_start,
                                         mod.FrStart, mod.FrIdx, mod.ToStart, mod.ToIdx, mod.GenStart,
                                         mod.GenIdx, mod.Pd, mod.Qd, sol.u_curr, sol.v_curr, sol.l_curr, sol.rho, mod.YshR, mod.YshI)
            time_bus += tcpu.time

            sol.l_curr .+= sol.rho .* (sol.u_curr .- sol.v_curr)
            sol.rd .= -sol.rho .* (sol.v_curr .- sol.v_prev)
            sol.rp .= sol.u_curr .- sol.v_curr
            #sol.rp_old .= sol.u_prev .- sol.v_prev

            primres = norm(sol.rp)
            dualres = norm(sol.rd)

            eps_pri = sqrt(length(sol.l_curr))*par.ABSTOL + par.RELTOL*max(norm(sol.u_curr), norm(-sol.v_curr))
            eps_dual = sqrt(length(sol.u_curr))*par.ABSTOL + par.RELTOL*norm(sol.l_curr)

            (par.verbose > 0) && @printf("[CPU] %10d  %.6e  %.6e  %.6e  %.6e  %6.2f  %6.2f\n",
                    it, primres, dualres, eps_pri, eps_dual, auglag_it, tron_it)

            if primres <= eps_pri && dualres <= eps_dual
                break
            end
        else
            @cuda threads=64 blocks=(div(mod.nvar-1, 64)+1) copy_data_kernel(mod.nvar, sol.u_prev, sol.u_curr)
            @cuda threads=64 blocks=(div(mod.nvar-1, 64)+1) copy_data_kernel(mod.nvar, sol.v_prev, sol.v_curr)
            @cuda threads=64 blocks=(div(mod.nvar-1, 64)+1) copy_data_kernel(mod.nvar, sol.l_prev, sol.l_curr)
            CUDA.synchronize()

            tgpu = generator_kernel(mod.gen_mod, data.baseMVA, sol.u_curr, sol.v_curr, sol.l_curr, sol.rho)

            time_gen += tgpu.time
            if env.use_polar
                tgpu = CUDA.@timed @cuda threads=32 blocks=nblk_br shmem=shmem_size polar_kernel(mod.n, mod.nline, mod.line_start, scale,
                                                                                sol.u_curr, sol.v_curr, sol.l_curr, sol.rho,
                                                                                shift_lines, env.membuf, mod.YffR, mod.YffI, mod.YftR, mod.YftI,
                                                                                mod.YttR, mod.YttI, mod.YtfR, mod.YtfI, mod.FrBound, mod.ToBound)
            else
                tgpu = CUDA.@timed @cuda threads=32 blocks=nblk_br shmem=shmem_size auglag_kernel(mod.n, it, par.max_auglag, mod.line_start, scale, par.mu_max,
                                                                                sol.u_curr, sol.v_curr, sol.l_curr, sol.rho,
                                                                                shift_lines, 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
            tgpu = CUDA.@timed @cuda threads=32 blocks=nblk_bus bus_kernel(data.baseMVA, mod.nbus, mod.gen_mod.gen_start, mod.line_start,
                                                                           mod.FrStart, mod.FrIdx, mod.ToStart, mod.ToIdx, mod.GenStart,
                                                                           mod.GenIdx, mod.Pd, mod.Qd, sol.u_curr, sol.v_curr, sol.l_curr,
                                                                           sol.rho, mod.YshR, mod.YshI)
            time_bus += tgpu.time

            @cuda threads=64 blocks=(div(mod.nvar-1, 64)+1) update_multiplier_kernel(mod.nvar, sol.l_curr, sol.u_curr, sol.v_curr, sol.rho)
            @cuda threads=64 blocks=(div(mod.nvar-1, 64)+1) primal_residual_kernel(mod.nvar, sol.rp, sol.u_curr, sol.v_curr)
            @cuda threads=64 blocks=(div(mod.nvar-1, 64)+1) dual_residual_kernel(mod.nvar, sol.rd, sol.v_prev, sol.v_curr, sol.rho)
            CUDA.synchronize()

            gpu_primres = CUDA.norm(sol.rp)
            gpu_dualres = CUDA.norm(sol.rd)

            gpu_eps_pri = sqrt(length(sol.l_curr))*par.ABSTOL + par.RELTOL*max(CUDA.norm(sol.u_curr), CUDA.norm(sol.v_curr))
            gpu_eps_dual = sqrt(length(sol.u_curr))*par.ABSTOL + par.RELTOL*CUDA.norm(sol.l_curr)

            (par.verbose > 0) && @printf("[GPU] %10d  %.6e  %.6e  %.6e  %.6e\n", it, gpu_primres, gpu_dualres, gpu_eps_pri, gpu_eps_dual)

            if gpu_primres <= gpu_eps_pri && gpu_dualres <= gpu_eps_dual
                break
            end
        end
    end
    end

    u_curr = zeros(mod.nvar)
    copyto!(u_curr, sol.u_curr)
    objval = sum(data.generators[g].coeff[data.generators[g].n-2]*(data.baseMVA*u_curr[mod.gen_mod.gen_start+2*(g-1)])^2 +
                 data.generators[g].coeff[data.generators[g].n-1]*(data.baseMVA*u_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)::Float64
    sol.objval = objval

    if it < iterlim
        sol.status = HAS_CONVERGED
    else
        sol.status = MAXIMUM_ITERATIONS
    end

    if par.verbose > 0
        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", objval)
    end
    return
end

function admm_rect_gpu(case::String; iterlim=800, rho_pq=400.0, rho_va=40000.0, scale=1e-4,
                       use_gpu=false, use_polar=true, gpu_no=0, verbose=1)
    env = AdmmEnv(case, use_gpu, rho_pq, rho_va; use_polar=use_polar, gpu_no=gpu_no, verbose=verbose,)
    admm_restart!(env, iterlim=iterlim, scale=scale)
    return env
end

# TODO: This needs revised to use AdmmEnv.
function admm_rect_gpu_mpi(
    case::String;
    iterlim=800, rho_pq=400.0, rho_va=40000.0, scale=1e-4, use_gpu=false, use_polar=true, gpu_no=1,
    comm::MPI.Comm=MPI.COMM_WORLD,
)
    data = opf_loaddata(case)

    ngen = length(data.generators)
    nline = length(data.lines)
    nbus = length(data.buses)
    nvar = 2*ngen + 8*nline
    if use_gpu
        @assert MPI.has_cuda()
    end

    # MPI settings
    root = 0
    is_master = MPI.Comm_rank(comm) == root
    n_processes = MPI.Comm_size(comm)
    nlines_local = div(nline, n_processes, RoundUp)
    nlines_padded = n_processes * nlines_local

    nvar_padded = 2*ngen + 8 * nlines_padded

    baseMVA = data.baseMVA
    n = (use_polar == true) ? 4 : 10
    mu_max = 1e8
    rho_max = 1e6
    rho_min_pq = 5.0
    rho_min_w = 5.0
    eps_rp = 1e-4
    eps_rp_min = 1e-5
    rt_inc = 2.0
    rt_dec = 2.0
    eta = 0.99
    Kf = 100
    Kf_mean = 10

    if use_gpu
        CUDA.device!(MPI.Comm_rank(comm) % CUDA.ndevices())
    end

    ybus = Ybus{Array{Float64}}(computeAdmitances(data.lines, data.buses, data.baseMVA; VI=Array{Int}, VD=Array{Float64})...)

    pgmin, pgmax, qgmin, qgmax, c2, c1, c0 = get_generator_data(data)
    YshR, YshI, YffR, YffI, YftR, YftI, YttR, YttI, YtfR, YtfI, FrBound, ToBound = get_branch_data(data)
    FrStart, FrIdx, ToStart, ToIdx, GenStart, GenIdx, Pd, Qd = get_bus_data(data)

    cu_pgmin, cu_pgmax, cu_qgmin, cu_qgmax, cu_c2, cu_c1, cu_c0 = get_generator_data(data; use_gpu=use_gpu)
    cuYshR, cuYshI, cuYffR, cuYffI, cuYftR, cuYftI, cuYttR, cuYttI, cuYtfR, cuYtfI, cuFrBound, cuToBound = get_branch_data(data; use_gpu=use_gpu)
    cu_FrStart, cu_FrIdx, cu_ToStart, cu_ToIdx, cu_GenStart, cu_GenIdx, cu_Pd, cu_Qd = get_bus_data(data; use_gpu=use_gpu)

    gen_start = 1
    line_start = 2*ngen + 1

    # Allocations
    u_curr = zeros(nvar_padded)
    v_curr = zeros(nvar_padded)
    l_curr = zeros(nvar_padded)
    u_prev = zeros(nvar_padded)
    v_prev = zeros(nvar_padded)
    l_prev = zeros(nvar_padded)
    rho = zeros(nvar_padded)
    rd = zeros(nvar_padded)
    rp = zeros(nvar_padded)
    rp_old = zeros(nvar_padded)
    rp_k0 = zeros(nvar_padded)
    param = zeros(31, nlines_padded)
    wRIij = zeros(2*nline)

    init_values(data, ybus, gen_start, line_start,
                rho_pq, rho_va, u_curr, v_curr, l_curr, rho, wRIij)

    if use_gpu
        cu_u_curr = CuArray{Float64}(undef, nvar_padded)
        cu_v_curr = CuArray{Float64}(undef, nvar_padded)
        cu_l_curr = CuArray{Float64}(undef, nvar_padded)
        cu_u_prev = CuArray{Float64}(undef, nvar_padded)
        cu_v_prev = CuArray{Float64}(undef, nvar_padded)
        cu_l_prev = CuArray{Float64}(undef, nvar_padded)
        cu_rho = CuArray{Float64}(undef, nvar_padded)
        cu_rd = CuArray{Float64}(undef, nvar_padded)
        cu_rp = CuArray{Float64}(undef, nvar_padded)
        cu_rp_old = CuArray{Float64}(undef, nvar_padded)
        cu_rp_k0 = CuArray{Float64}(undef, nvar_padded)
        cuParam = CuArray{Float64}(undef, (31, nlines_padded))
        cuWRIij = CuArray{Float64}(undef, 2*nline)

        copyto!(cu_u_curr, u_curr)
        copyto!(cu_v_curr, v_curr)
        copyto!(cu_l_curr, l_curr)
        copyto!(cu_rho, rho)
        copyto!(cuParam, param)
        copyto!(cuWRIij, wRIij)
    end

    # MPI: Global info
    if use_gpu
        u_lines_root = @view cu_u_curr[line_start:end]
        l_lines_root = @view cu_l_curr[line_start:end]
        v_lines_root = @view cu_v_curr[line_start:end]
        rho_lines_root = cu_rho[line_start:end]
    else
        u_lines_root = @view u_curr[line_start:end]
        l_lines_root = @view l_curr[line_start:end]
        v_lines_root = @view v_curr[line_start:end]
        rho_lines_root = @view rho[line_start:end]
    end
    # MPI: Local info
    # We need only to transfer info about lines
    if use_gpu
        u_local = CUDA.zeros(Float64, 8 * nlines_local)
        v_local = CUDA.zeros(Float64, 8 * nlines_local)
        l_local = CUDA.zeros(Float64, 8 * nlines_local)
        rho_local = CUDA.zeros(Float64, 8 * nlines_local)
    else
        u_local = zeros(8 * nlines_local)
        v_local = zeros(8 * nlines_local)
        l_local = zeros(8 * nlines_local)
        rho_local = zeros(8 * nlines_local)
    end

    MPI.Scatter!(u_lines_root, u_local, root, comm)
    MPI.Scatter!(rho_lines_root, rho_local, root, comm)

    max_auglag = 50

    nblk_gen = div(ngen, 32, RoundUp)
    nblk_br = nline
    nblk_br_local = nlines_local
    nblk_bus = div(nbus, 32, RoundUp)

    ABSTOL = 1e-6
    RELTOL = 1e-5

    it = 0
    time_gen = time_br = time_bus = 0
    time_br_scatter = time_br_gather = 0

    h_u_curr = zeros(nvar)
    h_param = zeros(31, nline)
    h_wRIij = zeros(2*nline)

    shift_lines = MPI.Comm_rank(comm) * nlines_local
    # GPU settings
    shmem_size = sizeof(Float64)*(14*n+3*n^2) + sizeof(Int)*(4*n)

    while it < iterlim
        it += 1

        # CPU code
        if !use_gpu

            if is_master
                u_prev .= u_curr
                v_prev .= v_curr
                l_prev .= l_curr
                tcpu = @timed generator_kernel_cpu(baseMVA, ngen, gen_start, u_curr, v_curr, l_curr, rho,
                                                pgmin, pgmax, qgmin, qgmax, c2, c1)
                time_gen += tcpu.time
            end

            # MPI routines to be implemented:
            #  - Broadcast cu_v_curr and cu_l_curr to GPUs.
            #  - Collect cu_u_curr.
            #  - div(nblk_br / number of GPUs, RoundUp)
            # scatter / gather

            tcpu_mpi = @timed begin
                MPI.Scatter!(v_lines_root, v_local, root, comm)
                MPI.Scatter!(l_lines_root, l_local, root, comm)
            end
            time_br_scatter += tcpu_mpi.time

            nlines_actual = min(nlines_local, nline - shift_lines)
            tcpu = @timed auglag_it, tron_it = polar_kernel_cpu(n, nlines_actual, 1, scale,
                                                                u_local, v_local, l_local, rho_local,
                                                                shift_lines, param, YffR, YffI, YftR, YftI,
                                                                YttR, YttI, YtfR, YtfI, FrBound, ToBound)

            time_br += tcpu.time
            tcpu_mpi = @timed MPI.Gather!(u_local, u_lines_root, root, comm)
            time_br_gather += tcpu_mpi.time

            if is_master
                tcpu = @timed bus_kernel_cpu(baseMVA, nbus, gen_start, line_start,
                                            FrStart, FrIdx, ToStart, ToIdx, GenStart,
                                            GenIdx, Pd, Qd, u_curr, v_curr, l_curr, rho, YshR, YshI)
                time_bus += tcpu.time

                l_curr .+= rho .* (u_curr .- v_curr)
                rd .= -rho .* (v_curr .- v_prev)
                rp .= u_curr .- v_curr
                #rp_old .= u_prev .- v_prev

                primres = norm(rp)
                dualres = norm(rd)

                eps_pri = sqrt(length(l_curr))*ABSTOL + RELTOL*max(norm(u_curr), norm(-v_curr))
                eps_dual = sqrt(length(u_curr))*ABSTOL + RELTOL*norm(l_curr)

                @printf("[CPU] %10d  %.6e  %.6e  %.6e  %.6e  %6.2f  %6.2f\n",
                        it, primres, dualres, eps_pri, eps_dual, auglag_it, tron_it)
            end
        # GPU code
        else
            if is_master
                @cuda threads=64 blocks=(div(nvar-1, 64)+1) copy_data_kernel(nvar, cu_u_prev, cu_u_curr)
                @cuda threads=64 blocks=(div(nvar-1, 64)+1) copy_data_kernel(nvar, cu_v_prev, cu_v_curr)
                @cuda threads=64 blocks=(div(nvar-1, 64)+1) copy_data_kernel(nvar, cu_l_prev, cu_l_curr)
                CUDA.synchronize()

                tgpu = CUDA.@timed @cuda threads=32 blocks=nblk_gen generator_kernel(baseMVA, ngen, gen_start,
                                                                                cu_u_curr, cu_v_curr, cu_l_curr, cu_rho,
                                                                                cu_pgmin, cu_pgmax, cu_qgmin, cu_qgmax, cu_c2, cu_c1)

                time_gen += tgpu.time
            end

            #  - Broadcast cu_v_curr and cu_l_curr to GPUs.
            tgpu_mpi = @timed begin
                MPI.Scatter!(v_lines_root, v_local, root, comm)
                MPI.Scatter!(l_lines_root, l_local, root, comm)
            end
            time_br_scatter += tgpu_mpi.time

            #  - div(nblk_br / number of GPUs, RoundUp)
            nblk_br_actual = min(nblk_br_local, nline - shift_lines)
            tgpu = CUDA.@timed @cuda threads=32 blocks=nblk_br_actual shmem=shmem_size polar_kernel(n, nline, 1, scale,
                u_local, v_local, l_local, rho_local,
                shift_lines, cuParam, cuYffR, cuYffI, cuYftR, cuYftI,
                cuYttR, cuYttI, cuYtfR, cuYtfI, cuFrBound, cuToBound
            )
            time_br += tgpu.time

            #  - Collect cu_u_curr.
            tgpu_mpi = @timed MPI.Gather!(u_local, u_lines_root, root, comm)
            time_br_gather += tgpu_mpi.time

            if is_master
                tgpu = CUDA.@timed @cuda threads=32 blocks=nblk_bus bus_kernel(baseMVA, nbus, gen_start, line_start,
                                                                            cu_FrStart, cu_FrIdx, cu_ToStart, cu_ToIdx, cu_GenStart,
                                                                            cu_GenIdx, cu_Pd, cu_Qd, cu_u_curr, cu_v_curr, cu_l_curr,
                                                                            cu_rho, cuYshR, cuYshI)
                time_bus += tgpu.time
                @cuda threads=64 blocks=(div(nvar-1, 64)+1) update_multiplier_kernel(nvar, cu_l_curr, cu_u_curr, cu_v_curr, cu_rho)
                @cuda threads=64 blocks=(div(nvar-1, 64)+1) primal_residual_kernel(nvar, cu_rp, cu_u_curr, cu_v_curr)
                @cuda threads=64 blocks=(div(nvar-1, 64)+1) dual_residual_kernel(nvar, cu_rd, cu_v_prev, cu_v_curr, cu_rho)
                CUDA.synchronize()

                gpu_primres = CUDA.norm(cu_rp)
                gpu_dualres = CUDA.norm(cu_rd)

                gpu_eps_pri = sqrt(length(l_curr))*ABSTOL + RELTOL*max(CUDA.norm(cu_u_curr), CUDA.norm(cu_v_curr))
                gpu_eps_dual = sqrt(length(u_curr))*ABSTOL + RELTOL*CUDA.norm(cu_l_curr)

                @printf("[GPU] %10d  %.6e  %.6e  %.6e  %.6e\n", it, gpu_primres, gpu_dualres, gpu_eps_pri, gpu_eps_dual)
            end
        end
    end

    if use_gpu
        copyto!(u_curr, cu_u_curr)
    end

    rank = MPI.Comm_rank(comm)
    @printf(" ** Time\n")
    @printf("[%d] Generator = %.2f\n", rank, time_gen)
    @printf("[%d] Branch    = %.2f\n", rank, time_br)
    @printf("[%d] Bus       = %.2f\n", rank, time_bus)
    @printf("[%d] G+Br+Bus  = %.2f\n", rank, time_gen + time_br + time_bus)
    @printf("[%d] Scatter   = %.2f\n", rank, time_br_scatter)
    @printf("[%d] Gather    = %.2f\n", rank, time_br_gather)
    @printf("[%d] MPI(S+G)  = %.2f\n", rank, time_br_scatter + time_br_gather)

    if is_master
        objval = sum(data.generators[g].coeff[data.generators[g].n-2]*(baseMVA*u_curr[gen_start+2*(g-1)])^2 +
                    data.generators[g].coeff[data.generators[g].n-1]*(baseMVA*u_curr[gen_start+2*(g-1)]) +
                    data.generators[g].coeff[data.generators[g].n]
                    for g in 1:ngen)
        @printf("Objective value = %.6e\n", objval)
    end

    return
end


#=
Copyright (c) 2015, Intel Corporation

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

    * Redistributions of source code must retain the above copyright notice,
      this list of conditions and the following disclaimer.
    * Redistributions in binary form must reproduce the above copyright
      notice, this list of conditions and the following disclaimer in the
      documentation and/or other materials provided with the distribution.
    * Neither the name of Intel Corporation nor the names of its contributors
      may be used to endorse or promote products derived from this software
      without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
=#

include("../../../src/Sparso.jl")
using Sparso

function abiatic(Hdmat, d, nqbits, T)
  total_time = -time()

  nsets = 2^nqbits
  Phi0at0 = 1/sqrt(2^nqbits)*ones(nsets)

  ode_time = -time()

  h = 0.01
  rtol = 1e-6

  # Compute the constants once
  c30 = 3/8
  c31 = -im*3/32
  c32 = -im*9/32
  c40 = 12/13
  c41 = -im*1932/2197
  c42 = im*7200/2197
  c43 = -im*7296/2197
  c51 = -im*439/216
  c52 = im*8
  c53 = -im*3680/513
  c54 = im*845/4104
  c61 = im*8/27
  c62 = -im*2
  c63 = im*3544/2565
  c64 = -im*1859/4104
  c65 = im*11/40
  cz1 = -im*16/135
  cz3 = -im*6656/12825
  cz4 = -im*28561/56430
  cz5 = im*9/50
  cz6 = -im*2/55
  ce1 = -im*1/360
  ce3 = im*128/4275
  ce4 = im*2197/75240
  ce5 = -im*1/50
  ce6 = -im*2/55

  # Absolute tolerance
  atol = 1e-13
  alpha_val = 0.8
  k = 0
  # Initial time moment
  i = 1
  Tv = zeros(0)
  push!(Tv, 0)
  t = 0
  # Initial condition
  meanenergy = zeros(1)
  variance = zeros(1)
  meanenergy[1] = dot(vec(abs(Phi0at0).^2), d)
  variance[1] = dot(vec(abs(Phi0at0).^2), (d - meanenergy[1]).^2)
  wi = Phi0at0
  # If it is the last iteration, then lastit = 1, otherwise lastit = 0
  lastit = 0
  spmv_time = 0

  while lastit == 0
      # Stretch the step if within 10% of b-t
      if t + 1.1*h > T
          h = T - t
          lastit = 1
      end

      # Compute the step
      s = t/T # real
      w = wi
      spmv_time -= time()
      s1 = -(1 - s)*Hdmat*w + s*d.*w
      spmv_time += time()

      s = (t + 0.25*h)/T
      w = wi + -im*0.25*h*s1
      spmv_time -= time()
      s2 = -(1 - s)*Hdmat*w + s*d.*w
      spmv_time += time()

      s = (t + c30*h)/T
      w = wi + c31*h*s1 + c32*h*s2
      spmv_time -= time()
      s3 = -(1 - s)*Hdmat*w + s*d.*w
      spmv_time += time()

      s = (t + c40*h)/T
      w = wi + c41*h*s1 + c42*h*s2 + c43*h*s3
      spmv_time -= time()
      s4 = -(1 - s)*Hdmat*w + s*d.*w
      spmv_time += time()

      s = (t + h)/T
      w = wi + c51*h*s1 + c52*h*s2 + c53*h*s3 + c54*h*s4
      spmv_time -= time()
      s5 = -(1 - s)*Hdmat*w + s*d.*w
      spmv_time += time()

      s = (t + 0.5*h)/T
      w = wi + c61*h*s1 + c62*h*s2 + c63*h*s3 + c64*h*s4 + c65*h*s5
      spmv_time -= time()
      s6 = -(1 - s)*Hdmat*w + s*d.*w
      spmv_time += time()

      z = wi + h*(cz1*s1 + cz3*s3 + cz4*s4 + cz5*s5 + cz6*s6)
      e = h * norm(ce1*s1 + ce3*s3 + ce4*s4 + ce5*s5 + ce6*s6) # real

      # Target tolerance for this step
      target_tol = rtol*norm(wi) + atol
      if e <= target_tol # In case the tolerance is met
          t = t + h
          h = alpha_val*h*(target_tol/e)^0.2
          i = i + 1
          push!(Tv, t)
          wi = z
          push!(meanenergy, dot(abs(z).^2, d))
          push!(variance, dot(abs(z).^2, (d - meanenergy[end]).^2))
          k = 0
      elseif k == 0 # Tolerance is not met for the first time in this step
          h = alpha_val*h*(target_tol/e)^0.2
          k = k + 1
          lastit = 0
      else # Tolerance is not met more than once in this step
          h = h/2
          lastit = 0
      end
  end

  println(size(Tv))

  ode_time += time()

  lanczos_time = -time()
  spectralgap = zeros(length(Tv))

  nit = min(30, nsets)

  x0 = rand(nsets)

  alpha = zeros(nit)
  beta = zeros(nit + 1)

  for i = 1:length(Tv)
    s = Tv[i]/T

    q0 = zeros(nsets)
    q1 = x0/norm(x0)

    for k = 1:nit
      spmv_time -= time()
      uk = -(1 - s)*Hdmat*q1 + s*d.*q1
      spmv_time += time()

      alpha[k] = dot(q1, uk)
      uk -= beta[k]*q0 + alpha[k]*q1
      beta[k + 1] = norm(uk)
      if beta[k + 1] == 0
        println("Error")
        return
      end

      q0 = q1
      q1 = uk/beta[k + 1]
    end

    TT = SymTridiagonal(alpha, beta[2:nit])

    Dt = sort(eig(TT)[1], 1)
    Dt = [Dt[1]; Dt[end]]

    spectralgap[i] = abs(Dt[1] - Dt[2])^2
  end
  lanczos_time += time()

  optimalenergy = minimum(d)

  total_time += time()

  @printf("%10s%10s%12s%12s%12s\n", "time", "eigengap", "optimal", "meanenergy", "variance");
  for i = length(Tv):length(Tv)
    @printf("%10.2f%10.4f%12.4f%12.4f%12.4f\n", Tv[i], spectralgap[i], optimalenergy, meanenergy[i], variance[i])
  end

  println("total_time = $(total_time) ode_time = $(ode_time) lanczos_time = $(lanczos_time) spmv_time = $(spmv_time)")
end

function print_result(Tv, spectralgap, optimalenergy, meanenergy, variance)
  @printf("%10s%10s%12s%12s%12s\n", "time", "eigengap", "optimal", "meanenergy", "variance");
  for i = length(Tv):length(Tv)
    @printf("%10.2f%10.4f%12.4f%12.4f%12.4f\n", Tv[i], spectralgap[i], optimalenergy, meanenergy[i], variance[i])
  end
end

function abiatic_sa(Hdmat, d, nqbits, T)
  total_time = -time()

  set_matrix_property(Dict(
     :Hdmat => SA_SYMM_VALUED | SA_STRUCTURE_ONLY,
     )
  )

  nsets = 2^nqbits
  Phi0at0 = 1/sqrt(2^nqbits)*ones(nsets)

  ode_time = -time()

  h = 0.01
  rtol = 1e-6

  # Compute the constants once
  c30 = 3/8
  c31 = -im*3/32
  c32 = -im*9/32
  c40 = 12/13
  c41 = -im*1932/2197
  c42 = im*7200/2197
  c43 = -im*7296/2197
  c51 = -im*439/216
  c52 = im*8
  c53 = -im*3680/513
  c54 = im*845/4104
  c61 = im*8/27
  c62 = -im*2
  c63 = im*3544/2565
  c64 = -im*1859/4104
  c65 = im*11/40
  cz1 = -im*16/135
  cz3 = -im*6656/12825
  cz4 = -im*28561/56430
  cz5 = im*9/50
  cz6 = -im*2/55
  ce1 = -im*1/360
  ce3 = im*128/4275
  ce4 = im*2197/75240
  ce5 = -im*1/50
  ce6 = -im*2/55

  # Absolute tolerance
  atol = 1e-13
  alpha_val = 0.8
  k = 0
  # Initial time moment
  i = 1
  Tv = zeros(0)
  push!(Tv, 0)
  t = 0
  # Initial condition
  meanenergy = zeros(1)
  variance = zeros(1)
  meanenergy[1] = dot(vec(abs(Phi0at0).^2), d)
  variance[1] = dot(vec(abs(Phi0at0).^2), (d - meanenergy[1]).^2)
  wi = Array{Complex128}(Phi0at0)
  # If it is the last iteration, then lastit = 1, otherwise lastit = 0
  lastit = 0
  spmv_time = 0

  s1 = Array{Complex128}(length(d))
  s2 = Array{Complex128}(length(d))
  s3 = Array{Complex128}(length(d))
  s4 = Array{Complex128}(length(d))
  s5 = Array{Complex128}(length(d))
  s6 = Array{Complex128}(length(d))

  while lastit == 0
      # Stretch the step if within 10% of b-t
      if t + 1.1*h > T
          h = T - t
          lastit = 1
      end

      # Compute the step
      s = t/T # real
      w = wi
      spmv_time -= time()
      s1 = -(1 - s)*Hdmat*w + s*d.*w
      spmv_time += time()

      s = (t + 0.25*h)/T
      w = wi + -im*0.25*h*s1
      spmv_time -= time()
      s2 = -(1 - s)*Hdmat*w + s*d.*w
      spmv_time += time()

      s = (t + c30*h)/T
      w = wi + c31*h*s1 + c32*h*s2
      spmv_time -= time()
      s3 = -(1 - s)*Hdmat*w + s*d.*w
      spmv_time += time()

      s = (t + c40*h)/T
      w = wi + c41*h*s1 + c42*h*s2 + c43*h*s3
      spmv_time -= time()
      s4 = -(1 - s)*Hdmat*w + s*d.*w
      spmv_time += time()

      s = (t + h)/T
      w = wi + c51*h*s1 + c52*h*s2 + c53*h*s3 + c54*h*s4
      spmv_time -= time()
      s5 = -(1 - s)*Hdmat*w + s*d.*w
      spmv_time += time()

      s = (t + 0.5*h)/T
      w = wi + c61*h*s1 + c62*h*s2 + c63*h*s3 + c64*h*s4 + c65*h*s5
      spmv_time -= time()
      s6 = -(1 - s)*Hdmat*w + s*d.*w
      spmv_time += time()

      z = wi + h*(cz1*s1 + cz3*s3 + cz4*s4 + cz5*s5 + cz6*s6)
      e = h * norm(ce1*s1 + ce3*s3 + ce4*s4 + ce5*s5 + ce6*s6) # real

      # Target tolerance for this step
      target_tol = rtol*norm(wi) + atol
      if e <= target_tol # In case the tolerance is met
          t = t + h
          h = alpha_val*h*(target_tol/e)^0.2
          i = i + 1
          push!(Tv, t)
          wi = z
          push!(meanenergy, dot(abs(z).^2, d))
          push!(variance, dot(abs(z).^2, (d - meanenergy[end]).^2))
          k = 0
      elseif k == 0 # Tolerance is not met for the first time in this step
          h = alpha_val*h*(target_tol/e)^0.2
          k = k + 1
          lastit = 0
      else # Tolerance is not met more than once in this step
          h = h/2
          lastit = 0
      end
  end

  println(size(Tv))

  ode_time += time()

  lanczos_time = -time()
  spectralgap = zeros(length(Tv))

  nit = min(30, nsets)

  x0 = rand(nsets)

  alpha = zeros(nit)
  beta = zeros(nit + 1)

  uk = Array{Float64}(length(d))

  dot_time = 0
  for i = 1:length(Tv)
    s = Tv[i]/T

    q0 = zeros(nsets)
    q1 = x0/norm(x0)

    for k = 1:nit
      spmv_time -= time()
      uk = -(1 - s)*Hdmat*q1 + s*d.*q1
      spmv_time += time()

      alpha[k] = dot(q1, uk)
      uk -= beta[k]*q0 + alpha[k]*q1
      beta[k + 1] = norm(uk)
      if beta[k + 1] == 0
        println("Error")
        return
      end

      q0 = q1
      #q1 = uk/beta[k + 1]
      q1 = Sparso.WAXPB(1/beta[k+1], uk, 0)
    end

    TT = SymTridiagonal(alpha, beta[2:nit])

    Dt = sort(eig(TT)[1], 1)
    Dt = [Dt[1]; Dt[end]]

    spectralgap[i] = abs(Dt[1] - Dt[2])^2
  end
  lanczos_time += time()

  optimalenergy = minimum(d)

  total_time += time()

  print_result(Tv, spectralgap, optimalenergy, meanenergy, variance)

  println("total_time = $(total_time) ode_time = $(ode_time) lanczos_time = $(lanczos_time) spmv_time = $(spmv_time)")
end

nqbits = parse(Int, ARGS[1])
T = parse(Int, ARGS[2])

nsets = 2^nqbits
d = 4*collect(-nsets/2 : nsets/2 - 1)

tunnelstr = -1

hdmat_time = -time()
Hdmat = zeros(nsets, nsets)
for i=1:nsets
  for j=0:nqbits-1
    flip = (i - 1)$(1 << j) # flip jth bit in zero-based indexing
    Hdmat[i,flip+1] = 1
  end
end
Hdmat = SparseMatrixCSC{Float64, Int32}(sparse(Hdmat))

if length(ARGS) == 3
  test = ARGS[3]
else
  test = "julia"
end

if test == "julia"
  println("compiler warmup (ignored): ")
  srand(0)
  println(@time(abiatic(Hdmat, d, nqbits, T)))

  println("\nRUN: ")
  srand(0)
  println(@time(abiatic(Hdmat, d, nqbits, T)))
else
  if test == "call-repl"
    set_options(SA_ENABLE, SA_USE_SPMP, SA_REPLACE_CALLS)
  elseif test == "context"
    set_options(SA_ENABLE, SA_USE_SPMP, SA_CONTEXT, SA_REPLACE_CALLS)
  elseif test == "reorder"
    set_options(SA_ENABLE, SA_USE_SPMP, SA_CONTEXT, SA_REORDER, SA_REPLACE_CALLS)
  elseif test == "verbose"
    set_options(SA_ENABLE, SA_USE_SPMP, SA_VERBOSE, SA_CONTEXT, SA_REORDER, SA_REPLACE_CALLS)
  end

  println("compiler warmup (ignored): ")
  srand(0)
  println(@time(@acc abiatic_sa(Hdmat, d, nqbits, T)))

  println("\nRUN: ")
  srand(0)
  println(@time(@acc abiatic_sa(Hdmat, d, nqbits, T)))
end


include("definitions.jl")
include("updateFunctions.jl")
include("utility.jl")

function sparseSCVB0_lda(documents::Array{SparseDocument,1},
                           numWords::Int64,
                           numDocuments::Int64,
                           numTopics::Int64,
                           numDocumentIterations::Int,
                           burnInPerDoc::Int,
                           minibatchSize::Int,
                           Alpha::Float64,
                           Beta::Float64,
                           dict,
                           tau::Float64,
                           kappa::Float64,
                           tau2::Float64,
                           kappa2::Float64,
                           scale::Float64,
                           sampleSize::Int,
                           time_limit::Float64,
                           miniEpoch::Int,
                           numDocuments_miniEpoch::Int64,
						   sparsity::Float64)

# sparse stochastic collapsed variational Bayes (sparseSCVB0) for LDA
    Alpha = ones(1,numTopics).*Alpha;
    Beta = ones(numWords,1).*Beta;

    # initialization of sparseSCVB0 specific model parameters
    #documentTopicCounts, wordTopicCounts, topicCounts, topicMem, sparseCounts = initialize(documents, numWords,numTopics,numDocuments);
    sparseVBparams=initialize(documents,Alpha, Beta, numTopics,numWords,numDocumentIterations,numDocuments);
    miniBatchesPerCorpus = numDocuments ./ minibatchSize;
    stepSize = scale ./ (tau^kappa); # step size for global expected counts
    #stepSize2 = 0.0; # initializing step size for local expected counts to make it soft local scope

    wordTopicCounts_hat = zeros(numWords,numTopics);
    topicCounts_hat = zeros(1,numTopics);
    iter = 0;

    saved_topics = Array{Array{Float64,2}}(0)
    saved_Totaltopics = Array{Array{Float64,2}}(0)
    saved_time = Float64[]
    saved_iters = Int64[]
    #saved_docTopics = Array{Array{Float64,2}}(0)

    prev_doc = 0; # to track miniEpoch docs
    doc_counter = 0;
    # main sparse stochastic collapsed variational Bayes (sparseSCVB0) loop
    while (sparseVBparams.wallClockTime < time_limit) && (doc_counter<numDocuments)
        prev_doc,documentTopicCounts,topicMem,sparseCounts = initialize_miniEpoch(documents,numTopics,numWords,numDocuments_miniEpoch,prev_doc);
        for epoch = 1:miniEpoch
            for epochDoc = 1:numDocuments_miniEpoch
                iter += 1;
                tic();
                #docInd_mini = mod(iter - 1, numDocuments_miniEpoch) + 1;
                #docInd = doc_counter + docInd_mini;
                docInd = doc_counter + epochDoc;
                docLength = documents[docInd].docLength;
                if docLength == 0
                    continue;
                end
                    # burn-in pass for a document
                sparseVBparams,documentTopicCounts,sparseCounts,topicMem = burnInUpdate(sparseVBparams,documentTopicCounts,documents,burnInPerDoc,docInd,epochDoc,docLength,tau2,kappa2,sparseCounts,topicMem,sampleSize,sparsity);

                    # main loop update for a word
                sparseVBparams,documentTopicCounts,wordTopicCounts_hat,topicCounts_hat,sparseCounts,topicMem = mainLoopUpdate(sparseVBparams,documentTopicCounts,documents,burnInPerDoc,docInd,epochDoc,docLength,tau2,kappa2,sparseCounts,
                                                                topicMem,wordTopicCounts_hat,topicCounts_hat,sampleSize);

                # mini-batch update of expected global counts
                if mod(iter, minibatchSize) == 0
                    # compute effective stepsize to account for minibatches
                    stepSize = 1 - (1-stepSize)^minibatchSize
                    # minibatch update of expected global counts
                    #wordTopicCounts,topicCounts,wordTopicCounts_hat,topicCounts_hat=minibatchUpdate(stepSize,miniBatchesPerCorpus,wordTopicCounts,topicCounts,wordTopicCounts_hat,topicCounts_hat);
                    topicCounts = sparseVBparams.topicCounts;
                    wordTopicCounts = sparseVBparams.wordTopicCounts;
                    for topic = 1:numTopics
                        for word = 1:numWords
                            wordTopicCounts[word,topic] = (1 - stepSize) .*  wordTopicCounts[word,topic] + stepSize .* miniBatchesPerCorpus .* wordTopicCounts_hat[word,topic];
                            wordTopicCounts_hat[word,topic] = wordTopicCounts_hat[word,topic]*0.0;
                        end
                        topicCounts[topic] = (1 - stepSize) .* topicCounts[topic] + stepSize .* miniBatchesPerCorpus .* topicCounts_hat[topic];
                        topicCounts_hat[topic] = topicCounts_hat[topic]*0.0;
                    end
                    stepSize = scale ./ (iter + tau)^kappa;
                end
                sparseVBparams.wallClockTime += toq();

                if iter == 1
                    push!(saved_topics,copy(sparseVBparams.wordTopicCounts))
                    push!(saved_Totaltopics,copy(sparseVBparams.topicCounts))
                    push!(saved_time,sparseVBparams.wallClockTime)
                    push!(saved_iters,iter)
                else
                    if mod(iter,200)==0
                        push!(saved_topics,copy(sparseVBparams.wordTopicCounts))
                        push!(saved_Totaltopics,copy(sparseVBparams.topicCounts))
                        push!(saved_time,sparseVBparams.wallClockTime)
                        push!(saved_iters,iter);
                    end
                end
            end
        end
        doc_counter = doc_counter + numDocuments_miniEpoch;
        #push!(saved_docTopics,copy(documentTopicCounts))
    end
    return sparseVBparams,saved_topics,saved_Totaltopics,saved_time,saved_iters,iter;
end

function initialize(documents,Alpha, Beta, numTopics,numWords,numDocumentIterations,numDocuments);
    # initialize sparseSCVB0 model parameters
    sparseVBparams = sparseSCVBO_params_type(Alpha, Beta, numTopics,numWords,numDocumentIterations,numDocuments);
    totalWordsInCorpus = 0;
    for i = 1:numDocuments
        docLength = documents[i].docLength;
        totalWordsInCorpus = totalWordsInCorpus + docLength;
    end
    # initialize expected global counts
    sparseVBparams.topicCounts.*= totalWordsInCorpus ./ numTopics;
    sparseVBparams.wordTopicCounts.*= totalWordsInCorpus ./ (numWords .* numTopics);

    # initializing alias samples from alias table for each word
    sparseVBparams.cachedSamples,sparseVBparams.cached_qw,sparseVBparams.cached_Qw = aliasMethod(sparseVBparams.numWords,sparseVBparams.Alpha,
                                       sparseVBparams.wordTopicCounts,sparseVBparams.topicCounts,sparseVBparams.Beta,
                                       sparseVBparams.BetaSum,sparseVBparams.numTopics,sparseVBparams.numCachedSamples,
                                       sparseVBparams.cachedSamples,sparseVBparams.cached_qw,sparseVBparams.cached_Qw);
    return sparseVBparams;
end

function initialize_miniEpoch(documents,numTopics,numWords,numDocuments_miniEpoch,prev_doc);
    # initialization of topic storage
    topicMem = Array{TopicArray}(0);
    # initialization for keeping track on non-zero sparse counts
    sparseCounts = Array{SparseCounts}(0);
    # initialize sparseSCVB0 miniEpoch parameters
    documentTopicCounts = zeros(numDocuments_miniEpoch,numTopics);
    j=0;
    for i = (prev_doc+1):(prev_doc+numDocuments_miniEpoch)
        j=j+1;
        docLength = documents[i].docLength;
        tokenLength = length(documents[i].tokenIndex); # storing topic for clumped documents
        Assigned = Array{Int64}(tokenLength); # array to store topic assignment for each token of clumped document
        nonzeroDocIndex = Array{Int64}(0); # array to keep track of non-zero document-topic counts
        documentTopicTmp = rand(1,numTopics);
        sum_documentTopicTmp = sum(documentTopicTmp);
        for topic = 1:numTopics
            documentTopicCounts[j,topic] = (documentTopicTmp[topic]./sum_documentTopicTmp).*docLength;
            push!(nonzeroDocIndex,topic);
        end
        Assigned = rand(1:numTopics,tokenLength);
        push!(topicMem,TopicArray(Assigned));
        push!(sparseCounts,SparseCounts(nonzeroDocIndex));
    end
    prev_doc = prev_doc+numDocuments_miniEpoch;
    return prev_doc,documentTopicCounts,topicMem,sparseCounts;
end

function aliasMethod(numWords,Alpha,wordTopicCounts,topicCounts,Beta,BetaSum,numTopics,numCachedSamples,cachedSamples,cached_qw,cached_Qw)
    qw = Array{Float64}(numTopics); # unnormalized discrete distribution of dense bucket
    for word = 1:numWords
        cached_Qw[word] = 0.0;
        for topic = 1:numTopics
            qw[topic] = Alpha[topic] .* (wordTopicCounts[word,topic] + Beta[word]) ./ (topicCounts[topic] + BetaSum);
            cached_qw[word,topic] = qw[topic];
            cached_Qw[word] += qw[topic];
        end
        aliasTable = generateAlias(qw,numTopics);
        for sample = 1:numCachedSamples
            cachedSamples[word,sample] = sampleAlias(aliasTable,numTopics);
        end
    end
    return cachedSamples,cached_qw,cached_Qw;
end


#include("common.jl")
#include("algorithm_parameters.jl")
#include("dynamical_systems.jl")

"""
`TMState` is a named tuple of the state variables of a temporal memory.

- `α`: active neurons
- `Π`: predictive neurons
- `WN`: winning neurons
- `Πₛ`: predictive dendritic segments (to calculate winning segments)
- `Mₛ`: matching dendritic segments (didn't receive enough input to activate, but enough to learn)
- `ovp_Mₛ`:
"""
mutable struct TMState
  state::NamedTuple{
    (:α, :Π, :WN, :Πₛ, :Mₛ, :ovp_Mₛ),  # Ncell, Ncell, Ncell, Nseg, Nseg, Nseg
    Tuple{CellActivity, CellActivity, CellActivity,
          BitArray{1}, BitArray{1}, Vector{Int}}
  }
  init_Nseg::Int
end
TMState(Nₙ::Int, Nseg)= TMState((
  α=falses(Nₙ), Π=falses(Nₙ), WN=falses(Nₙ),
  Πₛ=falses(Nseg), Mₛ=falses(Nseg),
  ovp_Mₛ=zeros(Nseg)
), Nseg)
Base.getproperty(s::TMState, name::Symbol)= name === :state ?
    getfield(s,:state) :
    name === :init_Nseg ? getfield(s, :init_Nseg) :
      getproperty(getfield(s,:state),name)
update_TMState!(s::TMState; Nseg,α,Π,WN,Πₛ,Mₛ,ovp_Mₛ)=
    s.state= (α=α, Π= Π, WN= WN, Πₛ= padfalse(Πₛ,Nseg),
              Mₛ= padfalse(Mₛ,Nseg), ovp_Mₛ= padfalse(ovp_Mₛ,Nseg))
reset!(s::TMState)= begin
  Nₙ= length(s.α);
  update_TMState!(s, Nseg=s.init_Nseg, α=falses(Nₙ), Π=falses(Nₙ), WN=falses(Nₙ),
    Πₛ=falses(s.init_Nseg), Mₛ=falses(s.init_Nseg), ovp_Mₛ=zeros(s.init_Nseg))
end

"""
    TemporalMemory(params::TMParams; recurrent= true, distal_input_size=0)

`TemporalMemory` learns to predict sequences of input Sparse Distributed Representations (SDRs),
usually generated by a [`SpatialPooler`](@ref).
It learns to represent each input symbol in the temporal context of the symbols that come before it in the sequence,
using the individual neurons of each minicolumn and their distal synapses.

When considering a neuron layer with proximal and distal synapses, the spatial pooler is a way to
activate and learn the proximal synapses, while the temporal memory is a way to activate and learn
the distal synapses.

Parameters:

- `params`: [`TMParams`](@ref) defining the layer size and algorithm constants
- `recurrent`: presence of recurrent distal connections between the layer's neurons
- `distal_input_size`: the number of presynaptic neurons for each layer with incoming distal connections (apart from same layer's neurons)

## High-order predictions and Ambiguity

The [neurons-thousand-synapses paper](https://www.frontiersin.org/articles/10.3389/fncir.2016.00023/full) describes the
Temporal Memory's properties, and especially the ability to
- make "high-order" predictions, based on previous inputs potentially going far back in time
- represent ambiguity in the form of simultaneous predictions

For more information see figures 2,3 in the paper.

# TM activation

Overview of the temporal memory's process:

1. Activate neurons (fire, proximal synapses)
2. Predict neurons (depolarize, distal/apical synapses)
3. Learn distal/apical synapses:
    - adapt existing synapses
    - create new synapses/dendrites

## Activation

An SDR input activates some minicolumns of the neuron layer.
If some neurons in the minicolum were predicted at the previous step, they activate faster than the rest and inhibit them.
If no neuron was predicted, all the neurons fire (minicolumn bursting).

TODO

---
See also: [`TMParams`](@ref) for parameter and symbol description, [`DistalSynapses`](@ref), [`TMState`](@ref)
"""
struct TemporalMemory
  params::TMParams
  distalSynapses::DistalSynapses
  previous::TMState
  recurrent::Bool
end
function TemporalMemory(params::TMParams; recurrent= true, distal_input_size=0)
  @unpack Nₙ, Nc, k = params
  distal_input_size= Int(distal_input_size)
  @assert recurrent || distal_input_size > 0 "At least 1 of recurrent or external input connections are needed"

  # NOTE: it's not possible to have independent distal synapses for each presynaptic input,
  # because we need a single address space for postsynaptic segments. Otherwise,
  # it would be impossible to mix signals from multiple sources.
  N_presynaptic= recurrent ? Nₙ + distal_input_size : distal_input_size
  Nseg_init= 0

  TemporalMemory(params,
      DistalSynapses(N_presynaptic, Nc, k;
        Nseg_init= Nseg_init, params= DistalSynapseParams(params)),
      TMState(Nₙ,Nseg_init),
      recurrent)
end

Nₙ(tm::TemporalMemory)= tm.params.Nₙ
reset!(tm::TemporalMemory)= begin
  reset!(tm.distalSynapses)
  reset!(tm.previous)
end

"""
`step!(tm::TemporalMemory, c, distal_input=[])` evolves the Temporal Memory to the next timestep
given the minicolumn activations and any distal input.
Returns the state of the region's neurons: active, predictive, bursting (minicolumns)

See also: [`tm_activate`](@ref), [`tm_predict`](@ref)
"""
# Given a column activation pattern `c` (SP output), step the TM
function step!(tm::TemporalMemory, c, distal_input=falses(0))
  s= tm.distalSynapses; p= tm.previous

  α, B, WN= tm_activate(tm, c, p.Π)
  # If recurrent, concatenate distal input activity with this layer's activity
  distal_input= tm.recurrent ? [α;distal_input] : distal_input
  Π, Πₛ, Mₛ, ovp_Mₛ= tm_predict(tm, distal_input)

  # Learn
  WS, WS_burst= calculate_WS!(s, p.Πₛ,p.ovp_Mₛ,α,B)
  # Update winner neurons with entries from bursting columns
  WN[NS(s)*WS_burst .> 0].= true
  step!(s, p.WN,WS, α, p.α,p.Mₛ,p.ovp_Mₛ)
  update_TMState!(p, Nseg=Nₛ(s),
                  α=α, Π=Π, WN=WN, Πₛ=Πₛ, Mₛ=Mₛ, ovp_Mₛ=ovp_Mₛ)
  return (
    active= α,
    predictive= Π,
    bursting= B
  )
end

"""
Active and predictive cells given the minicolumn activations and any distal input.

See also: [`tm_activate`](@ref), [`tm_predict`](@ref)
"""
(tm::TemporalMemory)(c, distal_input=falses(0))= begin
  α= tm_activate(tm,c,tm.previous.Π)[1]
  # If recurrent, concatenate distal input activity with this layer's activity
  distal_input= tm.recurrent ? [α;distal_input] : distal_input
  (
    active= α,
    predictive= tm_predict(tm,distal_input)[1]
  )
end

"""
`tm_activate(tm::TemporalMemory, c, Π)` calculates

1. which minicolumns burst,
2. which neurons in the layer are activated
3. which become predictive

for minicolumn activation `c` (size `Nc`) given by the [`SpatialPooler`](@ref)
and the previously predictive neurons `Π` (size `Nₙ`).

# Returns

1. `a`: neuron activation (`Nₙ`)
2. `B`: bursting minicolumns (`Nc`)
3. `WN`: "winning" neurons (`Nₙ`)

See also: [`tm_predict`](@ref), [`DistalSynapses`](@ref)
"""
function tm_activate(tm::TemporalMemory, c, Π)
  @unpack Nc, k = tm.params

  # bursting minicolumns (Nc)
  burst(c,Π)= c .& .!@percolumn(any, Π, k)
  # minicolumns with a predictive neuron (k × Nc)
  predicted(c,Π)= @percolumn(&,Π,c, k)
  # activation for the whole layer given the predicted/bursting neurons
  activate(A_pred, B)= (A_pred .| B')|> vec
  # cache parts of the output because they're used 2 times
  B= burst(c,Π)           # Nc
  A_pred= predicted(c,Π)  # k × Nc
  return activate(A_pred,B), B, A_pred|>vec
end

"""
`tm_predict(tm::TemporalMemory, α)` calculates which neurons will be predictive at the next step
given the currently active neurons + distal input `α`.
`size(a)` must match the presynaptic length of `tm.distalSynapses`.

# Returns

1. `Π`: predictive neurons (`Nₙ`)
2. `Πₛ`: predictive dendritic segments ('Nₛ') (caching)
3. `Mₛ`: matching dendritic segments (`Nₛ`) (learning)
4. `ovp_Mₛ`: subthreshold-matching dendritic segments (`Nₛ`) (learning)
"""
function tm_predict(tm::TemporalMemory, α::CellActivity)
  @unpack θ_stimulus_activate, θ_stimulus_learn = tm.params
  distal= tm.distalSynapses

  # Segment depolarization (prediction)
  Πₛ= Wd(distal)'α .> θ_stimulus_activate
  # Neuron depolarization (prediction)
  Π(Πₛ)= NS(distal)*Πₛ .> 0  # NOTE: params.θ_segment_act instead of 0
  # Sub-threshold segment stimulation sufficient for learning
  ovp_Mₛ= distal.Dd'α
  Mₛ= ovp_Mₛ .> θ_stimulus_learn
  return Π(Πₛ),Πₛ, Mₛ,ovp_Mₛ
end


# Simulate a 2D actomyosin network
# Alex Tam, 12/10/2020

# Simulate network
"Control function for actomyosin network simulations"
function actomyosin_network(parN, parA, parM, par, trial)
    # Specify domain width
    Lxx::Float64 = 2.5; Lxy::Float64 = 0; Lyx::Float64 = 0; Lyy::Float64 = 2.5; # Actual domain widths
    # Pre-allocate
    Force = [[0.0, 0.0, 0.0, 0.0] for idx in 1:parN.nT]; # [pN] Network force
    Curvature = [0.0 for idx in 1:parN.nT]; # Mean network curvature
    Index = [0.0 for idx in 1:parN.nT]; # Mean two-filament index
    Filament_Speed = [0.0 for idx in 1:parN.nT]; # [μm/s] Mean filament node speed
    Motor_Speed = [0.0 for idx in 1:parN.nT]; # [μm/s] Mean motor head speed
    Angle_ROC = [0.0 for idx in 1:parN.nT]; # [rad/s] Mean rate of change of angle
    Motor_Pos = [0.0 for idx in 1:parN.nT]; # [--] Mean myosin motor head position
    Motor_Angle = [0.0 for idx in 1:parN.nT]; # [rad] Mean angle between filaments at motor binding site
    Bins = Vector{Float64}(); Hist = Vector{Int}(); # [μm, --] Pre-allocate paired-distance histogram
    # Generate initial conditions
    state = State{Float64}(Vector{Vector{Vector}}(), Vector{Vector}()); # Initialise empty State struct
    mm = Vector{Myosin_Motor}(); # Pre-allocate empty myosin motors
    af, state = actin_ic(state, parN, parA, Lxx, Lyy); # Initialise actin filaments
    xl = intersection_search(state, af, mm); # Initialise cross-links
    mm, xl, state = myosin_ic(state, mm, parN, xl, Lxx, Lxy, Lyx, Lyy); # Initialise myosin motors
    state_old = state; # Store initial state to compute force
    # Time-stepping
    animation = @animate for i = 1:parN.nT
        # Write final DOF vector to file
        if i == parN.nT
            dof = build_dof(state);
            writedlm("dof-$par-trial-$trial-$i.csv", dof);
        end
        # Compute network characteristics
        Curvature[i] = curvature(af, state, Lxx, Lxy, Lyx, Lyy); # savefig("actomyosin_curvature-$par-trial-$trial.svg"); # Mean filament curvature
        Index[i] = two_filament_index(mm, state, Lxx, Lxy, Lyx, Lyy); # savefig("actomyosin_2f_index-$par-trial-$trial.svg"); # Mean two-filament index
        Bins, Hist = pcf(af, state, Lxx, Lyy); savefig("actomyosin_pcf-$par-trial-$trial.svg"); # Paired distances between filament nodes
        Filament_Speed[i] = filament_speed(parN, af, state, state_old, Lxx, Lxy, Lyx, Lyy); # savefig("actomyosin_filament_speed-$par-trial-$trial.svg"); # Filament node speeds
        Motor_Speed[i] = motor_speed(parN, mm, state, state_old, Lxx, Lxy, Lyx, Lyy); # savefig("actomyosin_motor_speed-$par-trial-$trial.svg"); # Motor head speeds
        Angle_ROC[i] = motor_angle_roc(parN, mm, state, state_old, Lxx, Lxy, Lyx, Lyy); # savefig("actomyosin_angle_roc-$par-trial-$trial.svg"); # Angle rate of change
        Motor_Pos[i] = motor_position(mm, state); # savefig("actomyosin_angle_roc-$par-trial-$trial.svg"); # Myosin motor head positions
        Motor_Angle[i] = motor_angle(mm, state, parN, Lxx, Lxy, Lyx, Lyy); # savefig("actomyosin_angle_roc-$par-trial-$trial.svg"); # Myosin motor angles
        Force[i] = network_force(state, state_old, af, xl, mm, parN, parA, parM, Lxx, Lxy, Lyx, Lyy); # Force
        # Draw network
        draw_network(state, af, xl, mm, parN, parA, Force[i], Lxx, Lxy, Lyx, Lyy);
        if i == 1
            savefig("actomyosin_ic-par-$par-trial-$trial.svg"); # Save image of initial condition
        end
        # Simulate one time step
        if i != parN.nT # Ensure correct looping sequence
            # Turnover and polymerisation
            af, mm, state = actin_turnover(state, af, mm, parN, parA, Lxx, Lyy);
            af = segment_translations(state, af); # Update filament translations
            xl = intersection_search(state, af, mm); # Update cross-links
            mm, xl, state = myosin_turnover(state, xl, mm, parN, parM, Lxx, Lxy, Lyx, Lyy); # Myosin binding/unbinding
            # Compute solution and store data
            state_old = state; # Store current state for energy functional
            new_dof = optimise_network(state_old, af, xl, mm, parN, parA, parM, Lxx, Lxy, Lyx, Lyy); # Compute network solution
            state = build_state(new_dof, af, mm); # Construct State from vector
        end
    end
    # Output .gif and image of the network
    gif(animation, "actomyosin-par-$par-trial-$trial.gif", fps = 10); savefig("actomyosin_end-par-$par-trial-$trial.svg");
    # Plot quantities versus time
    draw_stress(parN, Force, parN.nT, Lxx, Lyy); savefig("actomyosin_stress-par-$par-trial-$trial.svg"); # Stress components
    times, Bulk_Stress, Bulk_Stress_Int = draw_bulk_stress(parN, Force, parN.nT, Lxx, Lyy); savefig("actomyosin_bulk_stress-par-$par-trial-$trial.svg"); # Bulk stress
    @printf("Time-averaged bulk stress is %f pN/μm*s.\n", Bulk_Stress_Int[end]/times[end])
    # Compute time-integrated network statistics
    Curvature_Int, Index_Int, Filament_Speed_Int, Motor_Speed_Int, Angle_ROC_Int, Motor_Pos_Int, Motor_Angle_Int = integrated_statistics(parN, Curvature, Index, Filament_Speed, Motor_Speed, Angle_ROC, Motor_Pos, Motor_Angle)
    # Save mean stress and spatial measures time series data to files
    writedlm("times.csv", times);
    writedlm("stress-par-$par-trial-$trial.csv", Bulk_Stress);
    writedlm("curvature-par-$par-trial-$trial.csv", Curvature);
    writedlm("index-par-$par-trial-$trial.csv", Index);
    writedlm("filament_speed-par-$par-trial-$trial.csv", Filament_Speed);
    writedlm("motor_speed-par-$par-trial-$trial.csv", Motor_Speed);
    writedlm("motor_pos-par-$par-trial-$trial.csv", Motor_Pos);
    writedlm("motor_angle-par-$par-trial-$trial.csv", Motor_Angle);
    return state, af, mm, xl, Bulk_Stress_Int[end]/times[end], Curvature_Int[end]/times[end], Index_Int[end]/times[end], Filament_Speed_Int[end]/times[end], Motor_Speed_Int[end]/times[end], Angle_ROC_Int[end]/times[end], Motor_Pos_Int[end]/times[end], Motor_Angle_Int[end]/times[end], Bins, Hist
end


using MathOptInterface
const MOI = MathOptInterface
const AFF = MOI.ScalarAffineFunction{Cdouble}
const EQ = MOI.EqualTo{Cdouble}
const AFFEQ = MOI.ConstraintIndex{AFF,EQ}

mutable struct Optimizer <: MOI.AbstractOptimizer
    objective_constant::Cdouble
    objective_sign::Int
    blockdims::Vector{CSDP_INT}
    varmap::Vector{Tuple{Int, Int, Int}} # Variable Index vi -> blk, i, j
    num_entries::Dict{Tuple{Int, Int}, Int}
    b::Vector{Cdouble}
    C::blockmatrix
    problem::Union{Nothing, LoadingProblem}
    X::blockmatrix
    y::Union{Nothing, Vector{Cdouble}}
    Z::blockmatrix
    status::CSDP_INT
    pobj::Cdouble
    dobj::Cdouble
    solve_time::Float64
    silent::Bool
    options::Dict{Symbol, Any}
    function Optimizer(; kwargs...)
        optimizer = new(
            zero(Cdouble), 1, CSDP_INT[], Tuple{Int, Int, Int}[],
            Dict{Tuple{Int, Int}, Int}(), Cdouble[],
            blockmatrix(), nothing, blockmatrix(), nothing, blockmatrix(),
            -1, NaN, NaN, NaN, false, Dict{Symbol, Any}())
        for (key, value) in kwargs
            MOI.set(optimizer, MOI.RawOptimizerAttribute(String(key)), value)
        end
        # May need to call `free_loaded_prob` and `free_loading_prob`.
        finalizer(MOI.empty!, optimizer)
        return optimizer
    end
end

varmap(optimizer::Optimizer, vi::MOI.VariableIndex) = optimizer.varmap[vi.value]

function MOI.supports(optimizer::Optimizer, param::MOI.RawOptimizerAttribute)
    return Symbol(param.name) in ALLOWED_OPTIONS
end
function MOI.set(optimizer::Optimizer, param::MOI.RawOptimizerAttribute, value)
    if !(param.name isa String)
        Base.depwarn(
            "passing `$(param.name)` to `MOI.RawOptimizerAttribute` as type " *
            "`$(typeof(param.name))` is deprecated. Use a string instead.",
            Symbol("MOI.set")
        )
    end
    if !MOI.supports(optimizer, param)
        throw(MOI.UnsupportedAttribute(param))
    end
    optimizer.options[Symbol(param.name)] = value
end
function MOI.get(optimizer::Optimizer, param::MOI.RawOptimizerAttribute)
    # TODO: This gives a poor error message if the name of the parameter is invalid.
    if !(param.name isa String)
        Base.depwarn(
            "passing `$(param.name)` to `MOI.RawOptimizerAttribute` as type " *
            "`$(typeof(param.name))` is deprecated. Use a string instead.",
            Symbol("MOI.set")
        )
    end
    return optimizer.options[Symbol(param.name)]
end

MOI.supports(::Optimizer, ::MOI.Silent) = true
function MOI.set(optimizer::Optimizer, ::MOI.Silent, value::Bool)
    optimizer.silent = value
end
MOI.get(optimizer::Optimizer, ::MOI.Silent) = optimizer.silent

MOI.get(::Optimizer, ::MOI.SolverName) = "CSDP"
# See table "Return codes for easy_sdp() and CSDP" in `doc/csdpuser.pdf`.
const RAW_STATUS = [
    "Problem solved to optimality.",
    "Problem is primal infeasible.",
    "Problem is dual infeasible.",
    "Problem solved to near optimality.",
    "Maximum iterations reached.",
    "Stuck at edge of primal feasibility.",
    "Stuck at edge of dual feasibility.",
    "Lack of progress.",
    "X, Z, or O is singular.",
    "NaN or Inf values encountered.",
    "Program stopped by signal (SIXCPU, SIGTERM, etc.)"]

function MOI.get(optimizer::Optimizer, ::MOI.RawStatusString)
    return RAW_STATUS[optimizer.status + 1]
end
function MOI.get(optimizer::Optimizer, ::MOI.SolveTimeSec)
    return optimizer.solve_time
end

function MOI.is_empty(optimizer::Optimizer)
    return iszero(optimizer.objective_constant) &&
        isone(optimizer.objective_sign) &&
        isempty(optimizer.blockdims) &&
        isempty(optimizer.varmap) &&
        isempty(optimizer.num_entries) &&
        isempty(optimizer.b) &&
        iszero(optimizer.C.nblocks) &&
        optimizer.C.blocks == C_NULL &&
        optimizer.problem === nothing
end

function MOI.empty!(optimizer::Optimizer)
    optimizer.objective_constant = zero(Cdouble)
    optimizer.objective_sign = 1
    empty!(optimizer.blockdims)
    empty!(optimizer.varmap)
    empty!(optimizer.num_entries)
    empty!(optimizer.b)
    if optimizer.problem !== nothing
        if optimizer.y !== nothing
            free_loaded_prob(optimizer.problem, optimizer.X, optimizer.y, optimizer.Z)
        end
        free_loading_prob(optimizer.problem)
    end
    optimizer.problem = nothing
    optimizer.C.nblocks = 0
    optimizer.C.blocks = C_NULL
    optimizer.X.nblocks = 0
    optimizer.X.blocks = C_NULL
    optimizer.y = nothing
    optimizer.Z.nblocks = 0
    optimizer.Z.blocks = C_NULL
    optimizer.status = -1
    optimizer.pobj = 0.0
    optimizer.dobj = 0.0
end

function MOI.supports(
    optimizer::Optimizer,
    ::Union{MOI.ObjectiveSense,
            MOI.ObjectiveFunction{AFF}})
    return true
end

MOI.supports_add_constrained_variables(::Optimizer, ::Type{MOI.Reals}) = false

const SupportedSets = Union{MOI.Nonnegatives, MOI.PositiveSemidefiniteConeTriangle}
MOI.supports_add_constrained_variables(::Optimizer, ::Type{<:SupportedSets}) = true
function MOI.supports_constraint(::Optimizer, ::Type{AFF}, ::Type{EQ})
    return true
end

function _new_block(optimizer::Optimizer, set::MOI.Nonnegatives)
    push!(optimizer.blockdims, -MOI.dimension(set))
    blk = length(optimizer.blockdims)
    for i in 1:MOI.dimension(set)
        push!(optimizer.varmap, (blk, i, i))
    end
end

function _new_block(optimizer::Optimizer, set::MOI.PositiveSemidefiniteConeTriangle)
    push!(optimizer.blockdims, set.side_dimension)
    blk = length(optimizer.blockdims)
    for i in 1:set.side_dimension
        for j in 1:i
            push!(optimizer.varmap, (blk, i, j))
        end
    end
end

function _add_constrained_variables(optimizer::Optimizer, set::SupportedSets)
    offset = length(optimizer.varmap)
    _new_block(optimizer, set)
    ci = MOI.ConstraintIndex{MOI.VectorOfVariables, typeof(set)}(offset + 1)
    return [MOI.VariableIndex(i) for i in offset .+ (1:MOI.dimension(set))], ci
end

function _error(start, stop)
    error(start, ". Use `MOI.instantiate(CSDP.Optimizer, with_bridge_type = Float64)` ", stop)
end

function constrain_variables_on_creation(
    dest::MOI.ModelLike,
    src::MOI.ModelLike,
    index_map::MOI.Utilities.IndexMap,
    ::Type{S},
) where {S<:MOI.AbstractVectorSet}
    for ci_src in
        MOI.get(src, MOI.ListOfConstraintIndices{MOI.VectorOfVariables,S}())
        f_src = MOI.get(src, MOI.ConstraintFunction(), ci_src)
        if !allunique(f_src.variables)
            _error("Cannot copy constraint `$(ci_src)` as variables constrained on creation because there are duplicate variables in the function `$(f_src)`",
                   "to bridge this by creating slack variables.")
        elseif any(vi -> haskey(index_map, vi), f_src.variables)
            _error("Cannot copy constraint `$(ci_src)` as variables constrained on creation because some variables of the function `$(f_src)` are in another constraint as well.",
                   "to bridge constraints having the same variables by creating slack variables.")
        else
            set = MOI.get(src, MOI.ConstraintSet(), ci_src)::S
            vis_dest, ci_dest = _add_constrained_variables(dest, set)
            index_map[ci_src] = ci_dest
            for (vi_src, vi_dest) in zip(f_src.variables, vis_dest)
                index_map[vi_src] = vi_dest
            end
        end
    end
end

function count_entry(optimizer::Optimizer, con_idx::Integer, blk::Integer)
    key = (con_idx, blk)
    optimizer.num_entries[key] = get(optimizer.num_entries, key, 0) + 1
end

# Loads objective coefficient α * vi
function load_objective_term!(optimizer::Optimizer, α, vi::MOI.VariableIndex)
    blk, i, j = varmap(optimizer, vi)
    # in SDP format, it is max and in MPB Conic format it is min
    coef = optimizer.objective_sign * α
    if i != j
        coef /= 2
    end
    addentry(optimizer.problem, 0, blk, i, j, coef, true)
end

function MOI.copy_to(dest::Optimizer, src::MOI.ModelLike)
    MOI.empty!(dest)
    index_map = MOI.Utilities.IndexMap()

    # Step 1) Compute the dimensions of what needs to be allocated
    constrain_variables_on_creation(
        dest,
        src,
        index_map,
        MOI.Nonnegatives,
    )
    constrain_variables_on_creation(
        dest,
        src,
        index_map,
        MOI.PositiveSemidefiniteConeTriangle,
    )
    vis_src = MOI.get(src, MOI.ListOfVariableIndices())
    if length(vis_src) < length(index_map.var_map)
        _error("Free variables are not supported by CSDP",
               "to bridge free variables into `x - y` where `x` and `y` are nonnegative.")
    end
    cis_src = MOI.get(src, MOI.ListOfConstraintIndices{AFF,EQ}())
    dest.b = Vector{Cdouble}(undef, length(cis_src))
    funcs = Vector{AFF}(undef, length(cis_src))
    for (k, ci_src) in enumerate(cis_src)
        funcs[k] = MOI.get(src, MOI.CanonicalConstraintFunction(), ci_src)
        set = MOI.get(src, MOI.ConstraintSet(), ci_src)
        if isempty(funcs[k].terms)
            throw(ArgumentError("Empty constraint $cis_src: $(funcs[k])-in-$set. Not supported by CSDP."))
        end
        if !iszero(MOI.constant(funcs[k]))
            throw(MOI.ScalarFunctionConstantNotZero{
                Cdouble, AFF, EQ}(
                    MOI.constant(funcs[k])))
        end
        for t in funcs[k].terms
            if !iszero(t.coefficient)
                blk, _, _ = varmap(dest, index_map[t.variable])
                count_entry(dest, k, blk)
            end
        end
        dest.b[k] = MOI.constant(set)
        index_map[ci_src] = AFFEQ(k)
    end

    # Step 2) Allocate CSDP datastructures
    dummy = isempty(dest.b)
    if dummy
        # See https://github.com/coin-or/Csdp/issues/2
        dest.b = [one(Cdouble)]
        dest.blockdims = [dest.blockdims; CSDP_INT(-1)]
        count_entry(dest, 1, length(dest.blockdims))
    end
    dest.C.nblocks = length(dest.blockdims)
    num_entries = zeros(CSDP_INT, length(dest.b), length(dest.blockdims))
    for (key, value) in dest.num_entries
        num_entries[key...] = value
    end
    dest.problem = allocate_loading_prob(Ref(dest.C), dest.blockdims, length(dest.b), num_entries, 3)
    if dummy
        # See https://github.com/coin-or/Csdp/issues/2
        duplicate = addentry(dest.problem, 1, length(dest.blockdims), 1, 1, 1.0, true)
        @assert !duplicate
    end

    # Step 3) Load data in the datastructures
    for k in eachindex(funcs)
        setconstant(dest.problem, k, dest.b[k])
        for term in funcs[k].terms
            if !iszero(term.coefficient)
                blk, i, j = varmap(dest, index_map[term.variable])
                coef = term.coefficient
                if i != j
                    coef /= 2
                end
                duplicate = addentry(dest.problem, k, blk, i, j, coef, true)
                @assert !duplicate
            end
        end
    end

    # Throw error for variable attributes
    MOI.Utilities.pass_attributes(dest, src, index_map, vis_src)
    # Throw error for constraint attributes
    MOI.Utilities.pass_attributes(dest, src, index_map, cis_src)

    # Pass objective attributes and throw error for other ones
    model_attributes = MOI.get(src, MOI.ListOfModelAttributesSet())
    for attr in model_attributes
        if attr != MOI.ObjectiveSense() && attr != MOI.ObjectiveFunction{AFF}()
            throw(MOI.UnsupportedAttribute(attr))
        end
    end
    # We make sure to set `objective_sign` first before setting the objective
    if MOI.ObjectiveSense() in model_attributes
        sense = MOI.get(src, MOI.ObjectiveSense())
        dest.objective_sign = sense == MOI.MIN_SENSE ? -1 : 1
    end
    if MOI.ObjectiveFunction{AFF}() in model_attributes
        func = MOI.get(src, MOI.ObjectiveFunction{AFF}())
        obj = MOI.Utilities.canonical(func)
        dest.objective_constant = obj.constant
        for term in obj.terms
            if !iszero(term.coefficient)
                load_objective_term!(dest, term.coefficient, index_map[term.variable])
            end
        end
    end
    return index_map
end

function MOI.optimize!(optimizer::Optimizer)
    write_prob(optimizer)

    start_time = time()
    optimizer.y = loaded_initsoln(optimizer.problem, length(optimizer.b), Ref(optimizer.X), Ref(optimizer.Z))

    options = optimizer.options
    if optimizer.silent
        options = copy(options)
        options[:printlevel] = 0
    end

    optimizer.status, optimizer.pobj, optimizer.dobj = loaded_sdp(
        optimizer.problem, optimizer.objective_sign * optimizer.objective_constant,
        Ref(optimizer.X), optimizer.y, Ref(optimizer.Z), options,
    )
    optimizer.solve_time = time() - start_time
    return
end

function MOI.get(m::Optimizer, ::MOI.TerminationStatus)
    status = m.status
    if status == -1
        return MOI.OPTIMIZE_NOT_CALLED
    elseif status == 0
        return MOI.OPTIMAL
    elseif status == 1
        return MOI.INFEASIBLE
    elseif status == 2
        return MOI.DUAL_INFEASIBLE
    elseif status == 3
        return MOI.ALMOST_OPTIMAL
    elseif status == 4
        return MOI.ITERATION_LIMIT
    elseif 5 <= status <= 7
        return MOI.SLOW_PROGRESS
    elseif 8 <= status <= 9
        return MOI.NUMERICAL_ERROR
    else
        error("Internal library error: status=$status")
    end
end

function MOI.get(m::Optimizer, attr::MOI.PrimalStatus)
    if attr.result_index > MOI.get(m, MOI.ResultCount())
        return MOI.NO_SOLUTION
    end
    status = m.status
    if status == 0
        return MOI.FEASIBLE_POINT
    elseif status == 1
        return MOI.INFEASIBLE_POINT
    elseif status == 2
        return MOI.INFEASIBILITY_CERTIFICATE
    elseif status == 3
        return MOI.NEARLY_FEASIBLE_POINT
    elseif 4 <= status <= 9
        return MOI.UNKNOWN_RESULT_STATUS
    else
        error("Internal library error: status=$status")
    end
end

function MOI.get(m::Optimizer, attr::MOI.DualStatus)
    if attr.result_index > MOI.get(m, MOI.ResultCount())
        return MOI.NO_SOLUTION
    end
    status = m.status
    if status == 0
        return MOI.FEASIBLE_POINT
    elseif status == 1
        return MOI.INFEASIBILITY_CERTIFICATE
    elseif status == 2
        return MOI.INFEASIBLE_POINT
    elseif status == 3
        return MOI.NEARLY_FEASIBLE_POINT
    elseif 4 <= status <= 9
        return MOI.UNKNOWN_RESULT_STATUS
    else
        error("Internal library error: status=$status")
    end
end

MOI.get(m::Optimizer, ::MOI.ResultCount) = 1
function MOI.get(m::Optimizer, attr::MOI.ObjectiveValue)
    MOI.check_result_index_bounds(m, attr)
    return m.objective_sign * m.pobj
end
function MOI.get(m::Optimizer, attr::MOI.DualObjectiveValue)
    MOI.check_result_index_bounds(m, attr)
    return m.objective_sign * m.dobj
end
struct PrimalSolutionMatrix <: MOI.AbstractModelAttribute end
MOI.is_set_by_optimize(::PrimalSolutionMatrix) = true
MOI.get(optimizer::Optimizer, ::PrimalSolutionMatrix) = optimizer.X

struct DualSolutionVector <: MOI.AbstractModelAttribute end
MOI.is_set_by_optimize(::DualSolutionVector) = true
MOI.get(optimizer::Optimizer, ::DualSolutionVector) = optimizer.y

struct DualSlackMatrix <: MOI.AbstractModelAttribute end
MOI.is_set_by_optimize(::DualSlackMatrix) = true
MOI.get(optimizer::Optimizer, ::DualSlackMatrix) = optimizer.Z

function block(optimizer::Optimizer, ci::MOI.ConstraintIndex{MOI.VectorOfVariables})
    return optimizer.varmap[ci.value][1]
end
function dimension(optimizer::Optimizer, ci::MOI.ConstraintIndex{MOI.VectorOfVariables})
    blockdim = optimizer.blockdims[block(optimizer, ci)]
    if blockdim < 0
        return -blockdim
    else
        return MOI.dimension(MOI.PositiveSemidefiniteConeTriangle(blockdim))
    end
end
function vectorize_block(M, blk::Integer, s::Type{MOI.Nonnegatives})
    return diag(block(M, blk))
end
function vectorize_block(M::AbstractMatrix{Cdouble}, blk::Integer, s::Type{MOI.PositiveSemidefiniteConeTriangle}) where T
    B = block(M, blk)
    d = LinearAlgebra.checksquare(B)
    n = MOI.dimension(MOI.PositiveSemidefiniteConeTriangle(d))
    v = Vector{Cdouble}(undef, n)
    k = 0
    for j in 1:d
        for i in 1:j
            k += 1
            v[k] = B[i, j]
        end
    end
    @assert k == n
    return v
end

function MOI.get(optimizer::Optimizer, attr::MOI.VariablePrimal, vi::MOI.VariableIndex)
    MOI.check_result_index_bounds(optimizer, attr)
    blk, i, j = varmap(optimizer, vi)
    return block(MOI.get(optimizer, PrimalSolutionMatrix()), blk)[i, j]
end

function MOI.get(optimizer::Optimizer, attr::MOI.ConstraintPrimal,
                 ci::MOI.ConstraintIndex{MOI.VectorOfVariables, S}) where S<:SupportedSets
    MOI.check_result_index_bounds(optimizer, attr)
    return vectorize_block(MOI.get(optimizer, PrimalSolutionMatrix()), block(optimizer, ci), S)
end
function MOI.get(optimizer::Optimizer, attr::MOI.ConstraintPrimal, ci::AFFEQ)
    MOI.check_result_index_bounds(optimizer, attr)
    return optimizer.b[ci.value]
end

function MOI.get(optimizer::Optimizer, attr::MOI.ConstraintDual,
                 ci::MOI.ConstraintIndex{MOI.VectorOfVariables, S}) where S<:SupportedSets
    MOI.check_result_index_bounds(optimizer, attr)
    return vectorize_block(MOI.get(optimizer, DualSlackMatrix()), block(optimizer, ci), S)
end
function MOI.get(optimizer::Optimizer, attr::MOI.ConstraintDual, ci::AFFEQ)
    MOI.check_result_index_bounds(optimizer, attr)
    return -MOI.get(optimizer, DualSolutionVector())[ci.value]
end


"""
This function creates a PowerSystems ThermalStandard unit.
"""
function create_PSY_generator(gen::ThermalGenEMIS{<: BuildPhase}, sys::PSY.System)

    tech = get_tech(gen)
    base_power = get_maxcap(gen)

    buses = PSY.get_components(PSY.Bus, sys)
    project_bus = PSY.Bus[]

    for b in buses
        if PSY.get_number(b) == get_bus(tech)
            push!(project_bus, b)
        end
    end

    PSY_gen =  PSY.ThermalStandard(
        get_name(gen), # name
        true,   # available
        true,   # status
        project_bus[1], # bus
        get_maxcap(gen) / base_power,    # active power
        1.0,                # reactive power
        get_maxcap(gen) / base_power,    # rating
        (min = get_mincap(gen) / base_power, max = get_maxcap(gen) / base_power), # active power limits
        nothing,      # reactive power limits
        (up = get_ramp_limits(tech)[:up] / (base_power * 60), down = get_ramp_limits(tech)[:down] / (base_power * 60)), # ramp limits
        get_operation_cost(tech), # operation cost
        base_power, # base power
        get_time_limits(tech), # up and down time limits
        PSY.PrimeMovers.Symbol(get_type(tech)), # primemover
        PSY.ThermalFuels.Symbol(get_fuel(tech)), # fuel type
    )
    return PSY_gen
end

"""
This function creates a PowerSystems RenewableDispatch unit.
"""
function create_PSY_generator(gen::RenewableGenEMIS{<: BuildPhase}, sys::PSY.System)
    tech = get_tech(gen)
    base_power = get_maxcap(gen)

    if get_type(tech) == "WT"
        primemover = PSY.PrimeMovers.WT
    elseif get_type(tech) == "PVe"
        primemover = PSY.PrimeMovers.PVe
    end

    buses = PSY.get_components(PSY.Bus, sys)
    project_bus = PSY.Bus[]

    for b in buses
        if PSY.get_number(b) == get_bus(tech)
            push!(project_bus, b)
        end
    end

    PSY_gen =  PSY.RenewableDispatch(
        get_name(gen),  # name
        true,           # available
        project_bus[1], # bus
        get_maxcap(gen) / base_power, # active power
        1.0,             # reactive power
        get_maxcap(gen) / base_power,    # rating
        primemover,     # primemover
        nothing,        # reactivepower limits
        1.0,            # power factor
        get_operation_cost(tech),
        base_power, # base power
    )
    return PSY_gen
end

"""
This function creates a PowerSystems GenericBattery unit.
"""
function create_PSY_generator(gen::BatteryEMIS{<: BuildPhase}, sys::PSY.System)
    tech = get_tech(gen)
    base_power = get_maxcap(gen)

    buses = PSY.get_components(PSY.Bus, sys)
    project_bus = PSY.Bus[]

    for b in buses
        if PSY.get_number(b) == get_bus(tech)
            push!(project_bus, b)
        end
    end

    PSY_gen =  PSY.GenericBattery(
        get_name(gen),  # name
        true,           # available
        project_bus[1], # bus
        PSY.PrimeMovers.Symbol(get_type(tech)), # primemover
        get_soc(tech) / (get_maxcap(gen) * base_power), #initial state of charge
        (min = get_storage_capacity(tech)[:min] / (get_maxcap(gen) * base_power), max = get_storage_capacity(tech)[:max] / (get_maxcap(gen) * base_power)), # state of charge limits
        get_maxcap(gen) / base_power, # rating
        get_maxcap(gen) / base_power, # active power
        (min = get_input_active_power_limits(tech)[:min] / base_power, max = get_input_active_power_limits(tech)[:max] / base_power), # input active power limits
        (min = get_output_active_power_limits(tech)[:min] / base_power, max = get_output_active_power_limits(tech)[:max] / base_power), # output active power limits
        get_efficiency(tech), # efficiency
        1.0,             # reactive power
        nothing,      # reactive power limits
        base_power, # base power
    )
    return PSY_gen
end


"""
Abstract supertype for organ parameters.
Extend to add additional components to organ parameters.
"""
abstract type AbstractParams end

"""
Model parameters that vary between organs
"""
@default_kw @flattenable @selectable struct Params{As,Sh,Al,Ma,Tr,Re,Ge,Pr} <: AbstractParams
    # Field               | Default                    | _     | Selectable Types
    assimilation_pars::As | ConstantCAssim()           | _     | Union{Nothing,AbstractAssim}
    scaling_pars::Sh      | nothing                    | _     | Union{Nothing,AbstractScaling}
    allometry_pars::Al    | nothing                    | _     | Union{Nothing,AbstractAllometry}
    maturity_pars::Ma     | nothing                    | _     | Union{Nothing,AbstractMaturity}
    activetrans_pars::Tr  | nothing                    | _     | Union{Nothing,ActiveTranslocation}
    passivetrans_pars::Re | LosslessPassiveTranslocation() | _ | PassiveTranslocation
    germination_pars::Ge  | nothing                    | _     | Union{Nothing,AbstractGermination}
    production_pars::Pr   | nothing                    | _     | Union{Nothing,AbstractProduction}
end

for fn in fieldnames(Params)
    @eval $fn(p::Params) = p.$fn
end

"""
Astract supertype for shared parameters.
Extend to change the components that are shared.
"""
abstract type AbstractSharedParams end

"""
    SharedParams(su_pars, core_pars, resorption_pars, tempcorr_pars, catabolism_pars)

Model parameters shared between organs.

# FieldMetadata macros
- `@default_kw` provides default values and constructors,
- `@selectable` provides the set of types that can be used for the parameter,
  so that they can be selected in a live interface.
"""
@default_kw @selectable struct SharedParams{SU,Co,Fe,Te,Ca} <: AbstractSharedParams
    # Field               | Default                   | Selectable Types
    su_pars::SU           | ParallelComplementarySU() | AbstractSynthesizingUnit
    core_pars::Co         | DEBCore()                 | _
    resorption_pars::Fe   | nothing                   | Union{Nothing,AbstractResorption}
    tempcorr_pars::Te     | nothing                   | Union{Nothing,AbstractTemperatureCorrection}
    catabolism_pars::Ca   | CatabolismCN()            | AbstractCatabolism
end

for fn in fieldnames(SharedParams)
    @eval $fn(p::SharedParams) = p.$fn
end

###########################################################################################
# Variables

"""
Model variables.
Allow storing and accessing variables for use by multiple components.
"""
abstract type AbstractVars end

"""
    PlottableVars()

Plottable model variables. These are vectors witih values for each time-step,
to allow plotting and model introspection.
"""
@udefault_kw @units @plottable struct PlottableVars{F,R,E,T,WP,H,TS} <: AbstractVars
    scaling::F        | [0.0] | _                 | _
    rate::R           | [0.0] | mol*mol^-1*d^-1   | _
    E_ctb::E          | [0.0] | mol/hr            | _
    θE::F             | [0.0] | _                 | _
    temp::T           | [0.0] | K                 | _
    tempcorrection::F | [0.0] | _                 | _
    swp::WP           | [0.0] | kPa               | _
    soilcorrection::F | [0.0] | _                 | _
    height::H         | [0.0] | m                 | _
    tstep::TS         | [1]   | _                 | false
end

"""
    Vars()

Mutable struct to allow storing variables
for use by multiple components.
"""
@udefault_kw @units mutable struct Vars{F,R,E,T,WP,H} <: AbstractVars
    scaling::F        | 0.0   | _
    rate::R           | 0.0   | mol*mol^-1*d^-1
    θE::F             | 0.0   | _
    E_ctb::E          | 0.0   | mol/hr
    temp::T           | 0.0   | K
    tempcorrection::F | 0.0   | _
    swp::WP           | 0.0   | kPa
    soilcorrection::F | 0.0   | _
    height::H         | 0.0   | m
end

tstep(v::PlottableVars) = v.tstep[1]
tstep(v::Vars) = nothing

set_tstep!(v::PlottableVars, val) = v.tstep[1] = val
set_tstep!(v::Vars, val) = nothing

depth(v) = height(v)

build_vars(vars::Vars, tspan) = vars
build_vars(vars::PlottableVars, tspan) = begin
    len = length(tspan)
    len <= length(vars.rate) && return vars

    for fname in fieldnames(typeof(vars))
        varvec = getfield(vars, fname)
        append!(varvec, fill(getfield(vars, fname)[1], len - length(varvec)))
    end
    vars
end

###########################################################################################
# Organs and Organisms

"""
Abstract supertype for organs. Inherit from it if you need to difine
behaviour different to that or [`Organ`](@ref).
"""
abstract type AbstractOrgan{P,S} end

params(o::AbstractOrgan) = o.params
shared(o::AbstractOrgan) = o.shared
vars(o::AbstractOrgan) = o.vars
flux(o::AbstractOrgan) = o.J

#= κ functions allow modular addition of destinations for catabolised flux,
by default maturity and active translocation. If those components are
not used the flux fraction will be allocated to κsoma. =#

κtra(o::AbstractOrgan) = κtra(activetrans_pars(o))
κtra(o::Nothing) = 0.0

κmat(o::AbstractOrgan) = κmat(maturity_pars(o))
κmat(::Nothing) = 0.0

κsoma(o::AbstractOrgan) = oneunit(κtra(o)) - κtra(o) - κmat(o)

# Define `scaling` and `setscaling` etc. methods
for fn in fieldnames(Vars)
    setfn = Symbol.(:set_, fn, :!)
    @eval $fn(o::AbstractOrgan) = $fn(vars(o))
    @eval $setfn(o::AbstractOrgan, x) = $setfn(vars(o), x)
    @eval @inline ($fn)(vars::PlottableVars) = vars.$fn[tstep(vars)]
    @eval @inline ($setfn)(vars::PlottableVars, val) = vars.$fn[tstep(vars)] = val
    @eval @inline ($fn)(vars::Vars) = vars.$fn
    @eval @inline ($setfn)(vars::Vars, val) = vars.$fn = val
end

#= Forward variable and parameter getter/setter methods
so they can be acessesed/set directly from the organ,
without knowning where they are actually stored. =#
tstep(o::AbstractOrgan) = tstep(vars(o))
set_tstep!(o::AbstractOrgan, t) = set_tstep!(vars(o), t)

for fn in fieldnames(DEBCore)
    @eval $fn(p::SharedParams) = $fn(core_pars(p))
    @eval $fn(o::AbstractOrgan{<:Any,<:SharedParams}) = $fn(shared(o))
end

for fn in fieldnames(SharedParams)
    @eval $fn(o::AbstractOrgan{<:Any,<:SharedParams}) = $fn(shared(o))
end

for fn in fieldnames(Params)
    @eval $fn(o::AbstractOrgan{<:Params}) = $fn(params(o))
end

"""
    Organ(params, shared, vars, J)

Basic model components. For a plants, organs might be roots, stem and leaves
"""
struct Organ{P,S,V,F} <: AbstractOrgan{P,S}
    params::P
    shared::S
    vars::V
    J::F
end
"""
    Organ(params, shared, records)

Construct an organ from parameters, shared parameters and
views into records arrays for vaiable and flux matrices.
"""
Organ(params::AbstractParams, shared::AbstractSharedParams, records) = begin
    vars = records.vars
    set_tstep!(vars, 1)
    J = view(records.J, Ti(1))
    Organ(params, shared, vars, J)
end

# Records hold vars and flux, allowing storage of each
# timestep for plotting if PlottableVars are used.
abstract type AbstractRecords end

@plottable struct PlottableRecords{V,F} <: AbstractRecords
    vars::V | true
    J::F    | false
end
PlottableRecords(vars::PlottableVars, J::AbstractArray) =
    PlottableRecords{map(typeof,(vars,J))...}(vars, J)
PlottableRecords(vars::PlottableVars, fluxval::Number, fluxaxes::Tuple, tspan::AbstractRange) = begin
    vars = build_vars(vars, tspan)
    J = build_flux(fluxval, fluxaxes..., tspan)
    PlottableRecords(vars, J)
end

@plottable struct Records{V,F} <: AbstractRecords
    vars::V | false
    J::F    | false
end
Records(vars::Vars, J::AbstractArray) =
    Records{map(typeof,(vars,J))...}(vars, J)
Records(vars::Vars, fluxval::Number, fluxaxes::Tuple) = begin
    J = build_flux(fluxval, fluxaxes...)
    Records{map(typeof, (vars,J))...}(vars, J)
end

build_flux(fluxval, x::Tuple, y::Tuple) = begin
    dims = X(Val(x)), Y(Val(y))
    A = zeros(typeof(fluxval), map(length, dims)...)
    DimensionalArray(A, dims)
end
build_flux(fluxval, x::Tuple, y::Tuple, time::AbstractRange) = begin
    dims = X(Val(x)), Y(Val(y)), Ti(time)
    A = zeros(typeof(fluxval), map(length, dims)...)
    DimensionalArray(A, dims)
end

"""
A an Organism is a model object for modelling the growth
of a single organism.

It can be run as a functor:

```julia
(o::AbstactOrganism)(du, u, p, t) =
...
```

So that it can be passed into DifferentialEquations.jl solvers
if required.

Where `du` is the flux to be updated, `u` is the current state,
`p` are new model paremeters or `nothing`, and `t` is the current
time step. See model.jl for implemntations.

When julia 1.5 is released this will be implemented for any
`AbstactOrganism`, but for now it is only implemented for `Plant`
due to current limitations in Julia.
"""
abstract type AbstractOrganism end

params(o::AbstractOrganism) = o.params
shared(o::AbstractOrganism) = o.shared
records(o::AbstractOrganism) = o.records
organs(o::AbstractOrganism) = o.organs
environment(o::AbstractOrganism) = o.environment
environment_start(o::AbstractOrganism) = o.environment_start
dead(o::AbstractOrganism) = o.dead[]
set_dead!(o::AbstractOrganism, val) = o.dead[] = val

"""
    define_organs(o::AbstractOrganism, t)

Organs are constructed with views of Records and J Arrays at time t
"""
define_organs(o::AbstractOrganism, t) =
    define_organs(params(o), shared(o), records(o), t)
define_organs(params::Tuple, shared, records::Tuple, t) =
    map((p, r) -> Organ(p, shared, r), params, records)

update_organs(organs, t) = update_organs(vars(first(organs)), organs, t)
update_organs(::Vars, organs, t) = organs
update_organs(::PlottableVars, organs, t) = begin
    i = floor(Int, ustrip(t))
    map(organs) do organ
        organ.vars.tstep[1] = i
        J = organ.J
        data = J.data
        @set! data.indices[3] = i
        @set! data.offset1 = (i - 1) * length(data.indices[1]) * length(data.indices[2])
        # Setfield.jl isn't type-stable, DD `rebuild` is better
        @set! organ.J = DimensionalData.rebuild(J, data)
        organ
    end
end

"""
    Plant(params, shared, records, environment, environment_start, dead)
    Plant(states=(:V, :C, :N),
          transformations=(:asi, :gro, :mai, :rej, :res),
          params=(ShootParamsCN(), RootParamsCN()),
          vars=(Vars(), Vars()),
          shared=SharedParams(),
          records=nothing,
          environment=nothing,
          time=0.0hr:1.0hr:8760.0hr,
          environment_start=Ref(1.0hr),
          dead=Ref(false))

Plant model.

`params` are a tuple of Parameters objects, one for each organ.

`shared` is a struct of parameters shared between organs.

`vars`: can be `Vars` or `PlottableVars` or custom struct with additional variables.
  `PlottableVars` will be stored for each timestep for plotting.

`environment` can be `nothing`, or a `MicroclimPoint` or `MicroclimControl`
from Microclimates.jl

`time` determines the timespan over which `PlottableVars` will be constructed.
it isn't used with regular `Vars`.

If 4 state model is used, pass in
```julia
states=(:V, :E, :C, :N),
```

Similarly, if additional components like Maturity or ActiveTranslocation
are used, their state label needs to be passed in:

```julia
transformations=(:asi, :gro, :mai, :mat, :rej, :tra, :res),
```
"""
@flattenable @description mutable struct Plant{P,S,R,O,E,ES,D} <: AbstractOrganism
    params::P             | true  | "Model parameters"
    shared::S             | true  | "Parameters shared between organs"
    records::R            | false | "Plotable variables stored in arrays and sliced for each timestep on demand"
    organs::O             | false | ""
    environment::E        | false | "Environment object, provides environmental variables for each timestep"
    environment_start::ES | false | "Start index of environmental data"
    dead::D               | false | "`Bool` flag: has the plant died"
    Plant(params::P, shared::S, records::R, organs::O, environment::E, environment_start::ES, dead::D) where {P,S,R,O,E,ES,D} = begin
        organs = define_organs(params, shared, records, 0)
        new{P,S,R,typeof(organs),E,ES,D}(params, shared, records, organs, environment, environment_start, dead)
    end
end
Plant(; states=(:V, :C, :N),
        transformations=(:asi, :gro, :mai, :rej, :tra, :res),
        params=(
            Params(assimilation_pars=KooijmanSLAPhotosynthesis()),
            Params(assimilation_pars=NAssim()),
        ),
        vars=(Vars(), Vars()),
        shared=SharedParams(),
        environment=nothing,
        time=0.0hr:1.0hr:8760.0hr,
        records=nothing,
        environment_start=Ref(0.0hr),
        dead=Ref(false)
      ) = begin
    fluxaxes = states, transformations
    fluxval = 1.0mol/hr
    records = build_records(records, vars, fluxval, fluxaxes, time)
    Plant(params, shared, records, nothing, environment, environment_start, dead)
end

build_records(records::AbstractRecords, args...) = records
build_records(records::Nothing, vars::Tuple, fluxval, fluxaxes, tspan) =
    map(vars) do v
        build_records(v, fluxval, fluxaxes, tspan)
    end
build_records(vars::Vars, fluxval::Number, fluxaxes::Tuple, tspan::AbstractRange) =
    Records(vars, fluxval, fluxaxes)
build_records(vars::PlottableVars, fluxval::Number, fluxaxes::Tuple, tspan::AbstractRange) =
    PlottableRecords(vars, fluxval, fluxaxes, tspan)

# Define a SubArray constructor so we can increment the time index
ConstructionBase.constructorof(::Type{A}) where A<:SubArray{T,N,P,I,L} where {T,N,P,I,L} =
    (parent, indices, offset1, stride1) -> begin
        SubArray{eltype(parent),ndims(parent),typeof(parent),typeof(indices),L}(parent, indices, offset1, stride1)
    end