using Test
using Printf: @sprintf
using LinearAlgebra: dot
using StaticArrays: SVector
using QuantumLattices.Essentials.Terms
using QuantumLattices.Essentials.Spatials: Point, PID, Bond, Bonds, Lattice, decompose, acrossbonds, zerothbonds
using QuantumLattices.Essentials.DegreesOfFreedom: IDFConfig, Table, IID, Index, Internal, FilteredAttributes, OID, Operator, Operators, coordabsent, coordpresent, Boundary
using QuantumLattices.Interfaces: rank, update!, expand!, id, kind, reset!
using QuantumLattices.Prerequisites: Float, decimaltostr
using QuantumLattices.Prerequisites.CompositeStructures: NamedContainer
using QuantumLattices.Mathematics.AlgebraOverFields: ID, SimpleID, idtype
import QuantumLattices.Interfaces: dimension, expand
import QuantumLattices.Essentials.DegreesOfFreedom: isHermitian, otype
import QuantumLattices.Essentials.Terms: couplingcenter, couplingcenters, abbr

struct TID <: IID nambu::Int end
Base.adjoint(sl::TID) = TID(3-sl.nambu)

struct TIndex{S} <: Index{PID{S}, TID}
    scope::S
    site::Int
    nambu::Int
end
Base.fieldnames(::Type{<:TIndex}) = (:scope, :site, :nambu)
Base.union(::Type{P}, ::Type{I}) where {P<:PID, I<:TID} = TIndex{fieldtype(P, :scope)}

struct TFock <: Internal{TID} end
dimension(f::TFock) = 2
Base.getindex(f::TFock, i::Int) = TID(i)

struct TCID <: SimpleID
    center::Int
    nambu::Int
end
Base.fieldnames(::Type{<:TCID}) = (:center, :nambu)

struct TCoupling{V<:Number, I<:ID{TCID}} <: Coupling{V, I}
    value::V
    id::I
end
Base.repr(tc::TCoupling) = @sprintf "%s ph(%s)@(%s)" decimaltostr(tc.value) join(tc.id.nambus, ", ") join(tc.id.centers, ", ")
function expand(tc::TCoupling, pids::NTuple{R, PID}, focks::NTuple{R, TFock}, kind::Union{Val{K}, Nothing}=nothing) where {R, K}
    nambus = tc.id.nambus
    pids = NTuple{rank(tc), eltype(pids)}(pids[tc.id.centers[i]] for i = 1:rank(tc))
    return ((tc.value, NTuple{rank(tc), TIndex{fieldtype(pids|>eltype, :scope)}}(TIndex(pids[i].scope, pids[i].site, nambus[i]) for i = 1:rank(tc))),)
end

struct TOperator{V<:Number, I<:ID{OID}} <: Operator{V, I}
    value::V
    id::I
end

abbr(::Type{<:Term{ST, :TermMu} where ST}) = :tmu
isHermitian(::Type{<:Term{ST, :TermMu} where ST}) = true
otype(T::Type{<:Term{ST, :TermMu} where ST}, I::Type{<:OID}) = TOperator{T|>valtype, ID{I, T|>rank}}

abbr(::Type{<:Term{ST, :TermHopping} where ST}) = :thp
isHermitian(::Type{<:Term{ST, :TermHopping} where ST}) = false
otype(T::Type{<:Term{ST, :TermHopping} where ST}, I::Type{<:OID}) = TOperator{T|>valtype, ID{I, T|>rank}}

@testset "Subscript" begin
    sub = @subscript (x1, 4, 4, x2) with x1 < x2
    @test sub == deepcopy(sub)
    @test isequal(sub, deepcopy(sub))
    @test rank(sub) == rank(typeof(sub)) == 2
    @test dimension(sub) == dimension(typeof(sub)) == 4
    @test string(sub) == "(x1, 4, 4, x2) with $(sub.identifier)"
    @test sub(Val('M'), 1, 2) == (1, 4, 4, 2)
    @test sub(Val('C'), 1, 2) == true
    @test sub(Val('C'), 2, 1) == false

    sub = @subscript (x1, x2, x1, x2)
    @test sub(Val('M'), 1, 2) == (1, 2, 1, 2)
    @test sub(Val('C'), 1, 2) == true

    sub = Subscript{4}()
    @test rank(sub) == 1 && dimension(sub) == 4
    @test string(sub) == "(*, *, *, *)"
    @test sub(Val('M'), 2) == (2, 2, 2, 2)
    @test sub(Val('C'), 2) == true

    sub = Subscript((1, 2, 2, 1))
    @test rank(sub) == 0 && dimension(sub) == 4
    @test string(sub) == "(1, 2, 2, 1)"
    @test sub(Val('M')) == (1, 2, 2, 1)
    @test sub(Val('C')) == true
end

@testset "Subscripts" begin
    sub1 = @subscript (x1, 2) with x1 < 2
    sub2 = @subscript (y1, y2, y1, y2) with y1 < y2
    subs = Subscripts(sub1, sub2)
    @test rank(subs) == rank(typeof(subs)) == 3
    @test rank(subs, 1) == rank(typeof(subs), 1) == 1
    @test rank(subs, 2) == rank(typeof(subs), 2) == 2
    @test dimension(subs) == dimension(typeof(subs)) == 6
    @test dimension(subs, 1) == dimension(typeof(subs), 1) == 2
    @test dimension(subs, 2) == dimension(typeof(subs), 2) == 4
    @test subs(Val('M'), (1, 2, 3)) == (1, 2, 2, 3, 2, 3)
    @test subs(Val('C'), (1, 2, 3)) == true
    @test subs(Val('C'), (2, 2, 3)) == false
    @test subs(Val('C'), (1, 3, 2)) == false
    @test subs(Val('C'), (2, 3, 2)) == false

    sub3 = Subscript((6, 6))
    @test subs == sub1 * sub2
    @test subs * sub3 == Subscripts(sub1, sub2, sub3)
    @test sub3 * subs == Subscripts(sub3, sub1, sub2)
    @test subs * Subscripts(sub3) == Subscripts(sub1, sub2, sub3)
    @test expand(subs * sub3, (3, 3, 2, 3, 2, 3, 8, 9))|>collect == [(1, 2, 1, 2, 1, 2, 6, 6), (1, 2, 1, 3, 1, 3, 6, 6), (1, 2, 2, 3, 2, 3, 6, 6)]
end

@testset "centerrelated" begin
    @test couplingcenter(Coupling, 2, 2, Val(1)) == 1
    @test couplingcenter(Coupling, 1, 2, Val(1)) == 1
    @test couplingcenters(Coupling, ('*', '*'), Val(1)) == (1, 1)
    @test couplingcenters(Coupling, (1, 2, 3), Val(3)) == (1, 2, 3)
    @test couplingcenters("(1, 2, 3, 4)") == (1, 2, 3, 4)
end

@testset "couplings" begin
    @test @couplings(TCoupling(1.0, ID(TCID(1, 1), TCID(2, 1)))) == Couplings(TCoupling(1.0, ID(TCID(1, 1), TCID(2, 1))))
    @test @couplings(TCoupling(1.0, ID(TCID(1, 1), TCID(2, 1)))+TCoupling(1.0, ID(TCID(1, 2), TCID(2, 2)))) == TCoupling(1.0, ID(TCID(1, 1), TCID(2, 1)))+TCoupling(1.0, ID(TCID(1, 2), TCID(2, 2)))
end

@testset "TermFunction" begin
    ta = TermAmplitude()
    @test ta() == 1

    ta = TermAmplitude(x->x+3.0)
    @test ta(1.0) == 4.0

    tcs = TCoupling(1.0, ID(TCID(1, 1), TCID(2, 1))) + TCoupling(2.0, ID(TCID(1, 2), TCID(2, 2)))
    termcouplings = TermCouplings(tcs)
    @test termcouplings == deepcopy(TermCouplings(tcs))
    @test isequal(termcouplings, deepcopy(TermCouplings(tcs)))
    @test termcouplings() == tcs

    tcs1 = TCoupling(1.0, ID(TCID(1, 1), TCID(2, 1))) + TCoupling(1.0, ID(TCID(1, 2), TCID(2, 2)))
    tcs2 = TCoupling(1.0, ID(TCID(1, 2), TCID(2, 1))) + TCoupling(1.0, ID(TCID(2, 1), TCID(1, 2)))
    termcouplings = TermCouplings(i->(tcs1, tcs2)[(i-1)%2+1])
    @test termcouplings(1) == tcs1
    @test termcouplings(2) == tcs2

    termmodulate = TermModulate(:t)
    @test termmodulate(t=1) == 1
    @test termmodulate(mu=1) == nothing
    @test termmodulate() == nothing

    termmodulate = TermModulate(:t, t->t*2.0)
    @test termmodulate(2) == 4
end

@testset "Term" begin
    point = Point(PID(1, 1), (0.0, 0.0), (0.0, 0.0))
    config = IDFConfig{TFock}(pid->TFock(), [PID(1, 1)])
    term = Term{'F', :TermMu, 2}(:mu, 1.5, 0, couplings=TCoupling(1.0, ID(TCID(1, 2), TCID(1, 1))), amplitude=(bond->3), modulate=false)
    @test term|>statistics == term|>typeof|>statistics == 'F'
    @test term|>kind == term|>typeof|>kind == :TermMu
    @test term|>id == term|>typeof|>id == :mu
    @test term|>valtype == term|>typeof|>valtype == Float
    @test term|>rank == term|>typeof|>rank == 2
    @test term|>abbr == term|>typeof|>abbr == :tmu
    @test term|>isHermitian == term|>typeof|>isHermitian == true
    @test term == deepcopy(term)
    @test isequal(term, deepcopy(term))
    @test string(term) == "TermMu{2F}(id=mu, value=1.5, bondkind=0, factor=1.0)"
    @test repr(term, point, config) == "tmu: 4.5 ph(2, 1)@(1, 1)"
    @test +term == term
    @test -term == term*(-1) == replace(term, factor=-term.factor)
    @test 2*term == term*2 == replace(term, factor=term.factor*2)
    @test term/2 == term*(1/2) == replace(term, factor=term.factor/2)

    p1 = Point(PID(1, 1), (0.0, 0.0), (0.0, 0.0))
    p2 = Point(PID(1, 2), (0.0, 0.0), (0.0, 0.0))
    config = IDFConfig{TFock}(pid->TFock(), [PID(1, 1), PID(1, 2)])
    tcs1 = Couplings(TCoupling(1.0, ID(TCID(1, 2), TCID(1, 2))))
    tcs2 = Couplings(TCoupling(1.0, ID(TCID(1, 1), TCID(1, 1))))
    term = Term{'F', :TermMu, 2}(:mu, 1.5, 0, couplings=bond->(tcs1, tcs2)[bond.pid.site%2+1], amplitude=bond->3, modulate=true)
    @test repr(term, p1, config) == "tmu: 4.5 ph(1, 1)@(1, 1)"
    @test repr(term, p2, config) == "tmu: 4.5 ph(2, 2)@(1, 1)"
    @test one(term) == replace(term, value=1.0)
    @test zero(term) == replace(term, value=0.0)
    @test term.modulate(mu=4.0) == 4.0
    @test term.modulate(t=1.0) == nothing
    @test update!(term, mu=4.25) == replace(term, value=4.25)
    @test term.value == 4.25
end

@testset "expand" begin
    point = Point(PID(1, 1), (0.0, 0.0), (0.0, 0.0))
    config = IDFConfig{TFock}(pid->TFock(), [PID(1, 1)])
    table = Table(config, filter(attr->(attr ≠ :nambu), FilteredAttributes(TIndex)))
    term = Term{'F', :TermMu, 2}(:mu, 1.5, 0, couplings=TCoupling(1.0, ID(TCID(1, 2), TCID(1, 1))), amplitude=bond->3.0, modulate=true)
    operators = Operators(TOperator(+2.25, (TIndex(1, 1, 2), TIndex(1, 1, 1)), seqs=(1, 1)))
    @test expand(term, point, config, table, true, coordabsent) == operators
    @test expand(term, point, config, table, false, coordabsent) == operators*2

    bond = Bond(1, Point(PID('b', 2), (1.5, 1.5), (1.0, 1.0)), Point(PID('a', 1), (0.5, 0.5), (0.0, 0.0)))
    config = IDFConfig{TFock}(pid->TFock(), [PID('a', 1), PID('b', 2)])
    table = Table(config, filter(attr->(attr ≠ :nambu), FilteredAttributes(TIndex)))
    term = Term{'F', :TermHopping, 2}(:t, 1.5, 1, couplings=TCoupling(1.0, ID(TCID(1, 2), TCID(2, 1))), amplitude=bond->3.0, modulate=true)
    operators = Operators(TOperator(4.5, (TIndex('a', 1, 2), TIndex('b', 2, 1)), rcoords=(SVector(0.5, 0.5), SVector(1.5, 1.5)), icoords=(SVector(0.0, 0.0), SVector(1.0, 1.0)), seqs=(1, 2)))
    @test expand(term, bond, config, table, true) == operators
    @test expand(term, bond, config, table, false) == operators+operators'

    lattice = Lattice("Tuanzi", [Point(PID(1, 1), (0.0, 0.0), (0.0, 0.0))], vectors=[[1.0, 0.0]], neighbors=1)
    bonds = Bonds(lattice)
    config = IDFConfig{TFock}(pid->TFock(), lattice.pids)
    table = Table(config, filter(attr->(attr ≠ :nambu), FilteredAttributes(TIndex)))
    term = Term{'F', :TermMu, 2}(:mu, 1.5, 0, couplings=TCoupling(1.0, ID(TCID(1, 2), TCID(1, 1))), amplitude=bond->3.0, modulate=true)
    operators = Operators(TOperator(+2.25, (TIndex(1, 1, 2), TIndex(1, 1, 1)), seqs=(1, 1)))
    @test expand(term, bonds, config, table, true, coordabsent) == operators
    @test expand(term, bonds, config, table, false, coordabsent) == operators*2
end

@testset "Parameters" begin
    ps1 = Parameters{(:t1, :t2, :U)}(1.0im, 1.0, 2.0)
    ps2 = Parameters{(:t1, :U)}(1.0im, 2.0)
    ps3 = Parameters{(:t1, :U)}(1.0im, 2.1)
    @test match(ps1, ps2) == true
    @test match(ps1, ps3) == false
end

@testset "Generator" begin
    lattice = Lattice("Tuanzi", [Point(PID(1, 1), (0.0, 0.0), (0.0, 0.0)), Point(PID(1, 2), (0.5, 0.0), (0.0, 0.0))], vectors=[[1.0, 0.0]], neighbors=1)
    bonds = Bonds(lattice)
    config = IDFConfig{TFock}(pid->TFock(), lattice.pids)
    table = Table(config, filter(attr->(attr ≠ :nambu), FilteredAttributes(TIndex)))
    boundary = Boundary()
    t = Term{'F', :TermHopping, 2}(:t, 2.0, 1, couplings=TCoupling(1.0, ID(TCID(1, 2), TCID(2, 1))))
    μ = Term{'F', :TermMu, 2}(:μ, 1.0, 0, couplings=TCoupling(1.0, ID(TCID(1, 2), TCID(1, 1))), modulate=true)
    tops1 = expand(t, filter(acrossbonds, bonds, Val(:exclude)), config, table, true, coordpresent)
    tops2 = expand(one(t), filter(acrossbonds, bonds, Val(:include)), config, table, true, coordpresent)
    μops = expand(one(μ), filter(zerothbonds, bonds, Val(:include)), config, table, true, coordpresent)

    optp = TOperator{Float, ID{OID{TIndex{Int}, SVector{2, Float}, SVector{2, Float}, Int}, 2}}
    genops = GenOperators(tops1, NamedContainer{(:μ,)}((μops,)), NamedContainer{(:t, :μ)}((tops2, Operators{optp|>idtype, optp}())))
    @test genops == deepcopy(genops) && isequal(genops, deepcopy(genops))
    @test genops == GenOperators((t, μ), bonds, config, table, true, coordpresent)
    @test genops|>eltype == genops|>typeof|>eltype == optp
    @test genops|>idtype == genops|>typeof|>idtype == optp|>idtype
    @test expand!(Operators{idtype(optp), optp}(), genops, boundary, t=2.0, μ=1.5) == tops1+tops2*2.0+μops*1.5
    @test empty!(deepcopy(genops)) == empty(genops) == GenOperators(empty(μops), NamedContainer{(:μ,)}((empty(μops),)), NamedContainer{(:t, :μ)}((empty(μops), empty(μops))))
    @test reset!(deepcopy(genops), (t, μ), bonds, config, table, true, coordpresent) == genops

    gen = Generator((t, μ), bonds, config, table, true, boundary)
    @test gen == deepcopy(gen) && isequal(gen, deepcopy(gen))
    @test Parameters(gen) == Parameters{(:t, :μ)}(2.0, 1.0)
    @test expand!(Operators{idtype(optp), optp}(), gen) == expand(gen) == tops1+tops2*2.0+μops
    @test expand(gen, :t) == tops1+tops2*2.0
    @test expand(gen, :μ) == μops
    @test expand(gen, 1)+expand(gen, 2)+expand(gen, 3)+expand(gen, 4) == expand(gen)
    @test expand(gen, :μ, 1)+expand(gen, :μ, 2) == μops
    @test expand(gen, :t, 3) == tops1
    @test expand(gen, :t, 4) == tops2*2.0
    @test empty!(deepcopy(gen)) == Generator{coordpresent}((t, μ), empty(bonds), empty(config), empty(table), true, boundary) == empty(gen)
    @test reset!(empty(gen), lattice) == gen
    @test update!(gen, μ=1.5)|>expand == tops1+tops2*2.0+μops*1.5
end


# This hack will go away once CLBLAS doesn't
# rely on the clunky cl.CLArray type in OpenCL anymore. Sadly, we need cl.Buffer for that, which  overwrites the default
# constructor registering a finalizer, which we don't want here, so we need to
# get the real constructor.
@generated function clbuffer(ptr::OwnedPtr{T}, len::Integer) where T
    Expr(:new, cl.Buffer{T}, true, :(pointer(ptr)), :(len), :(ptr.mapped), Ptr{T}(C_NULL))
end

function clbuffer(A::CLArray)
    ptr = pointer(A)
    clbuffer(ptr, length(A))
end

import CLFFT

# figure out a gc safe way to store plans.
# weak refs won't work, since the caching should keep them alive.
# But at the end, we need to free all of these, otherwise CLFFT will crash
# at closing time.
# An atexit hook here, which will empty the dictionary seems to introduce racing
# conditions.
#const plan_dict = Dict()
import Base: *, plan_ifft!, plan_fft!, plan_fft, plan_ifft, size, plan_bfft, plan_bfft!

immutable CLFFTPlan{Direction, Inplace, T, N} <: Base.FFTW.FFTWPlan{T, Direction, Inplace}
    plan::CLFFT.Plan{T}
    function (::Type{CLFFTPlan{Direction, Inplace}})(A::CLArray{T, N}) where {T, N, Direction, Inplace}
        ctx = global_context(A)
        p = CLFFT.Plan(T, ctx, size(A))
        CLFFT.set_layout!(p, :interleaved, :interleaved)
        if Inplace
            CLFFT.set_result!(p, :inplace)
        else
            CLFFT.set_result!(p, :outofplace)
        end
        CLFFT.set_scaling_factor!(p, Direction, 1f0)
        CLFFT.bake!(p, global_queue(A))
        new{Direction, Inplace, T, N}(p)
    end
end
size(x::CLFFTPlan) = (CLFFT.lengths(x.plan)...,)

# ignore flags, but have them to make it base compatible.
# TODO can we actually implement the flags?
function plan_fft(A::CLArray; flags = nothing, timelimit = Inf)
    CLFFTPlan{:forward, false}(A)
end
function plan_fft!(A::CLArray; flags = nothing, timelimit = Inf)
    CLFFTPlan{:forward, true}(A)
end
function plan_bfft(A::CLArray, region; flags = nothing, timelimit = Inf)
    CLFFTPlan{:backward, false}(A)
end
function plan_bfft!(A::CLArray, region; flags = nothing, timelimit = Inf)
    CLFFTPlan{:backward, true}(A)
end

const _queue_ref = Vector{cl.CmdQueue}(1)
function *(plan::CLFFTPlan{Direction, true, T, N}, A::CLArray{T, N}) where {T, N, Direction}
    _queue_ref[] = global_queue(A)
    CLFFT.enqueue_transform(plan.plan, Direction, _queue_ref, clbuffer(A), nothing)
    A
end
function *(plan::CLFFTPlan{Direction, false, T, N}, A::CLArray{T, N}) where {T, N, Direction}
    _queue_ref[] = global_queue(A)
    y = typeof(A)(size(plan))
    CLFFT.enqueue_transform(plan.plan, Direction, _queue_ref, clbuffer(A), clbuffer(y))
    y
end

# Enable BLAS support via GPUArrays
import CLBLAS
import GPUArrays: blas_module, blasbuffer
blas_module(::CLArray) = CLBLAS
blasbuffer(A::CLArray) = cl.CLArray(clbuffer(A), global_queue(A), size(A))


## INSTRUCTIONS:
## RUN IN LINE WITH SIMULATION MAIN FLOW CODE
## REQUIRES "data_file"

### UPDATE FUNCTION FOR DATA FILE ARRAYS

function UpdateArray(data_dictionary,string,alterations)
  # access array
  #println(data_dictionary)
  #println(alterations)
  data_array = data_dictionary[string];
  # set modifications for all conditions here
  for i in 1:Int(size(alterations)[1])
    updates = alterations[Int(i)];
    index = updates[1];
    value = updates[2];
    data_array[Int(index)]=value;
  data_dictionary[string] = data_array;
  end
  return data_array
end

##

#tgfb forward rate

forward_master = 0.5 #0.5

tgfb_add_forward_rate = forward_master #4.0
vegf_add_forward_rate = forward_master*10
vivit_add_forward_rate = forward_master*10

##
#### set values for multiple variables

# # # TGFB SIGNALING RATES
tgfb12_to_ap1_sp1_gain = 1.0 #1.5
tgfb3_to_smad_gain_rate = 2.0 *1.0 # TGFB3_R_phosphorylation_Smad: Smad =(TGFB3_R)=> pSmad # 2.0
tgfb12_to_smad_rate = 0.25
MAPK_activation_smad_rate = 0.2
# # # BETA-CATEININ RATES
bcat_deg_rate = 0.3 # default 0.1
# bcat_deg_gsk3_rate = 0.05 # default 0.1

bcat_trans_rate = 0.5-0.1
# # # BETA-CATEININ COMPLEXING
ecad_complex_gain = 12.0 *2 # 6 or 8
ecad_complex_reverse = 1.0 *2
ecad_bcat_degradation_rate = 0.3 + 0.3
lef1_complex_gain = 4.0
tcf4_complex_gain = 4.0 + 4
apc_complex_forward_rate=11.0 # 11.0
apc_bcat_deg_rate = 0.1
tgfb3_pik3_rate = 0.5  #5.0
gsk3_p_rate = 0.5 # 3.0
#ecad_bais_factor = 1.0 # multiply ecadherin-bcat only
ecad_trans_rate = 1.0+0.7+0.5# 1.5
ecad_deg_rate= 1.3 -0.9 -0.2 # 0.1

ecad_transc_rate = 1.0+0.4 # 1.5
ecad_mrna_deg_rate= 0.1 + 0.1 # 0.1

pi3k_translation_rate = 0.5 # 0.5
gsk3_translation_rate = 0.5 # 0.5

apc_complex_trans_rate = 0.6
lef1_trans_rate = 1.0 # 1.0
tcf4_trans_rate = 1.0 # 0.5

snail_slug_trans_rate = 1.0 # increases lef1 mrna too much

yreg_deg_rate = 1.5 # 1.5
# for ox/ko
yreg_ox_deg_rate = 0.15
yreg_ko_deg_rate_protein = 15

yreg_translation_rate = 1

### SELECTED DATA FILE INDEXES

# INDEX VALUES CHANGE WITH UPDATES TO MODEL.NET

# initial conditions
tgfb12_initial_cond_index = 88 # TGFB12
R_initial_cond_index = 87
vegf_initial_cond_index = 93 # VEGF

yreg_translation_rate_index = 152 # translation_mRNA_YREG1: mRNA_YREG1 = YREG1
yreg_protein_deg_rate_index = 48 # YREG1_degradation: YREG1 =([])=> []
# TGFB signaling
tgfb3_binding_rate_index = 3 # TGFB3_binding: TGFB3+R =([])=> TGFB3_R
R_deg_rate_index = 57 # degradation_R: R =([])=> []
TGFB12_R_activation_RAS_rate_index = 16 # TGFB12_R_activation_RAS: RAS =(TGFB12_R)=> RAS_GTP
RAS_activation_RAF_rate_index  = 17 # RAS_GTP_phosphorylation_RAF: RAF =(RAS_GTP)=> RAF_P
RAF_activation_MAPK_rate_index = 18 # RAF_P_phosphorylation_MAPK: MAPK =(RAF_P)=> pMAPK
MAPK_activation_AP1_SP1_rate_index = 20 # MAPK_phosphorylation_AP1_SP1: AP1_SP1 =(pMAPK)=> AP1_SP1_P
MAPK_activation_smad_rate_index = 19 # MAPK_phosphorylation_Smad: Smad =(pMAPK)=> pSmad
smad_deg_rate_index = 45 # Smad_degradation: Smad =([])=> []
psmad_deg_rate_index = 66 # pSmad_degradation: pSmad =([])=> []

# Bcatenin complex formation
bcat_ecad_forward_index = 21 # Binding_Ecadherin_Bcatenin: Ecadherin+Bcatenin =([])=> Ecadherin_Bcatenin
bcat_ecad_reverse_index = 22 # Binding_Ecadherin_Bcatenin: Ecadherin_Bcatenin =([])=> Ecadherin+Bcatenin (reverse)
#bcat_apc_reverse_index = 23
bcat_tcf4_forward_index = 25 # TCF4_Bcatenin_complex_formation: TCF4+Bcatenin =([])=> Active_TCF4
bcat_tcf4_reverse_index = 26 # TCF4_Bcatenin_complex_formation: Active_TCF4 =([])=> TCF4+Bcatenin (reverse)
bcat_lef1_forward_index = 27 # LEF1_Bcatenin_complex_formation: LEF1+Bcatenin =([])=> Active_LEF1
bcat_lef1_reverse_index = 28 # LEF1_Bcatenin_complex_formation: Active_LEF1 =([])=> LEF1+Bcatenin (reverse)
bcat_deg_rate_index = 43 # Bcatenin_degradation: Bcatenin =([])=> []
bcat_trans_rate_index = 149 # translation_mRNA_Bcatenin: mRNA_Bcatenin = Bcatenin

ecad_bcat_degradation_rate_index = 58 # Ecadherin_Bcatenin_degradation: Ecadherin_Bcatenin =([])=> []

# VEGF signaing
PKC_activation_RAF_rate_index = 11 # PKC_activation_RAF: RAF =(Active_PKC)=> RAF_P

snail_slug_trans_rate_index = 143 # induction_gene_SNAIL_SLUG: gene_SNAIL_SLUG = mRNA_SNAIL_SLUG
lef1_trans_rate_index = 153 # translation_mRNA_LEF1: mRNA_LEF1 = LEF1

apc_complex_forward_rate_index = 23 # Binding_APC_AXIN_forw: APC_AXIN+Bcatenin =([])=> APC_AXIN_Bcatenin
tgfb3_pik3_rate_index = 14 # TGFB3_R_phosphorylation_PI3K: PI3K =(TGFB3_R)=> PI3K_P
apc_bcat_deg_rate_index = 77 # degradation_APC_AXIN_Bcatenin: APC_AXIN_Bcatenin =([])=> []
gsk3_p_rate_index = 15 # PI3K_phosphorylation_GSK3: GSK3 =(PI3K_P)=> GSK3_P
pi3k_translation_rate_index = 169 # translation_mRNA_PI3K: mRNA_PI3K = PI3K
gsk3_translation_rate_index = 168	# translation_mRNA_GSK3: mRNA_GSK3 = GSK3
apc_complex_trans_rate_index = 167 # translation_mRNA_APC_AXIN: mRNA_APC_AXIN = APC_AXIN

#ecad bcat balance
ecad_trans_rate_index = 148 # translation_mRNA_Ecadherin: mRNA_Ecadherin = Ecadherin
ecad_deg_rate_index = 44 # Ecadherin_degrdation: Ecadherin =([])=> []
ecad_mrna_deg_rate_index = 102 # degradation_mRNA_Ecadherin: mRNA_Ecadherin = []
ecad_transc_rate_index = 125 # induction_gene_Ecadherin: gene_Ecadherin = mRNA_Ecadherin
# TGFB dosing
tgfb12_forward_rate_index = 91 # addition_TGFb12: [] =([])=> TGFB12
tgfb12_reverse_rate_index = 92 # addition_TGFb12: TGFB12 =([])=> [] (reverse)
tgfb12_binding_rate_index = 1 # TGFB12_binding: TGFB12+R =([])=> TGFB12_R

# VEGF dosing
vegf_forward_rate_index = 93 # addition_VEGF: [] =([])=> VEGF
vegf_reverse_rate_index = 94 # addition_VEGF: VEGF =([])=> [] (reverse)
vegf_binding_rate_index = 5 # VEGF_binding: VEGF+VEGFR =([])=> VEGF_VEGFR

# TGFB3_Ab dosing
TGFB3_Ab_foward_rate_index = 86 # addition_TGFB3_Ab: [] =([])=> TGFB3_Ab (reverse)
TGFB3_Ab_reverse_rate_index = 85 # addition_TGFB3_Ab: TGFB3_Ab =([])=> []

# U0126 dosing
U0126_foward_rate_index = 98 # addition_U0126: [] =([])=> U0126 (reverse)
U0126_reverse_rate_index = 97 # addition_U0126: U0126 =([])=> []

# TGFB3 dosing
TGFB3_foward_rate_index = 95 # addition_TGFB3: [] =([])=> TGFB3
TGFB3_reverse_rate_index = 96 # addition_TGFB3: TGFB3 =([])=> [] (reverse)
tgfb12_to_smad_rate_index = 12 # TGFB12_R_phosphorylation_Smad: Smad =(TGFB12_R)=> pSmad
tgfb3_to_smad_gain_rate_index = 13 # TGFB3_R_phosphorylation_Smad: Smad =(TGFB3_R)=> pSmad

# vivit

vivit_forward_rate_index  = 100 # addition_vivit: [] =([])=> vivit (reverse)
vivit_reverse_rate_index = 99 # addition_vivit: vivit =([])=> []

# Snail slug transfeciton
ss_virus_forward_rate_index = 88 # addition_ss_virus: [] =([])=> ss_virus (reverse)
ss_virus_reverse_rate_index = 87 # addition_ss_virus: ss_virus =([])=> []

# Ecadherin transfection
ecad_virus_forward_rate_index = 90 # addition_ecad_virus: [] =([])=> ecad_virus (reverse)
ecad_virus_reverse_rate_index = 89 # addition_ecad_virus: ecad_virus =([])=> []
# DN smad
# DN SMAD temoprary
smad_induction_rate_index = 124 # induction_gene_Smad: gene_Smad = mRNA_Smad
# redundant psmad_deg_rate_index = 66 # pSmad_degradation: pSmad =([])=> []
# reverse convention based on product term
#smad_virus_forward_rate_index = 146 # addition_smad_virus: [] =([])=> smad_virus (reverse)
#smad_virus_reverse_rate_index = 122 # addition_smad_virus: smad_virus =([])=> []

# DN 1
lef1_induction_rate_index = 130 # induction_gene_LEF1: gene_LEF1 = mRNA_LEF1
# reverse convention based on product term
#LEF1_virus_forward_rate_index = 18 # addition_lef1_virus: [] =([])=> lef1_virus (reverse)
#LEF1_virus_reverse_rate_index = 111 # addition_lef1_virus: lef1_virus =([])=> []

lef1_deg_rate_index = 50 # LEF1_degradation: LEF1 =([])=> []
Active_LEF1_deg_rate_index = 51 # Active_LEF1_degradation: Active_LEF1 =([])=> []

#yreg_deg_rate_index = 122 # degradation_mRNA_YREG1: mRNA_YREG1 = []
yreg_deg_rate_index = 106 # degradation_mRNA_YREG1: mRNA_YREG1 = []

tcf4_trans_rate_index = 150 # translation_mRNA_TCF4: mRNA_TCF4 = TCF4

# MAPK inhibitor
# ?
raf_mapk_phos_rate_index = 18 # RAF_P_phosphorylation_MAPK: MAPK =(RAF_P)=> pMAPK

#### THESE SETTINGS ARE FOR STEADY STATE AND ARE DEFAULT ####

# initialize arrays to hold updates for each array in data_file
initial_cond_update_array = [];
rate_constant_update_array = [];
control_parameter_update_array = [];

# initial conditions
push!(initial_cond_update_array,[tgfb12_initial_cond_index,0.0]) # TGFB12 initial conditions
push!(initial_cond_update_array,[R_initial_cond_index,6.0]) # R initial condition
push!(initial_cond_update_array,[vegf_initial_cond_index,0.0]) # VEGF

# rates
# TGFB dosing
push!(rate_constant_update_array,[tgfb12_forward_rate_index,0.0]) # addition_TGFb12: [] =([])=> TGFB12
push!(rate_constant_update_array,[tgfb12_reverse_rate_index,0.0]) # addition_TGFb12: TGFB12 =([])=> [] (reverse)
push!(rate_constant_update_array,[tgfb12_binding_rate_index,3.0]) # TGFB12_binding: TGFB12+R =([])=> TGFB12_R

# VEGF dosing
push!(rate_constant_update_array,[vegf_forward_rate_index,0.0]) # addition_VEGF: [] =([])=> VEGF
push!(rate_constant_update_array,[vegf_reverse_rate_index,0.0]) # addition_VEGF: VEGF =([])=> [] (reverse)
push!(rate_constant_update_array,[vegf_binding_rate_index,3.0]) # VEGF_binding: VEGF+VEGFR =([])=> VEGF_VEGFR

# ecad deg and trans
push!(rate_constant_update_array,[ecad_deg_rate_index,ecad_deg_rate]) #default 0.1
push!(rate_constant_update_array,[ecad_trans_rate_index ,ecad_trans_rate]) #default 1.0

## Inhibitor dosing

# SNAIL_SLUG dosing
# push!(rate_constant_update_array,[snail_slug_forward_rate_index,0.0]) # addition_VEGF: [] =([])=> VEGF
# push!(rate_constant_update_array,[snail_slug_reverse_rate_index,0.0]) # addition_VEGF: VEGF =([])=> [] (reverse)

# TGFB3_Ab dosing
push!(rate_constant_update_array,[TGFB3_Ab_foward_rate_index,0.0])
push!(rate_constant_update_array,[TGFB3_Ab_reverse_rate_index,0.0])

# U0126 dosing
push!(rate_constant_update_array,[U0126_foward_rate_index ,0.0])
push!(rate_constant_update_array,[U0126_reverse_rate_index,0.0])

# vivit dosing
push!(rate_constant_update_array,[vivit_forward_rate_index,0.0])
push!(rate_constant_update_array,[vivit_reverse_rate_index,0.0])

## Virus dosing

push!(rate_constant_update_array,[ss_virus_forward_rate_index,0.0])
push!(rate_constant_update_array,[ss_virus_reverse_rate_index,0.0])

push!(rate_constant_update_array,[ecad_virus_forward_rate_index,0.0])
push!(rate_constant_update_array,[ecad_virus_reverse_rate_index,0.0])

#push!(rate_constant_update_array,[smad_virus_forward_rate_index,0.0])
#push!(rate_constant_update_array,[smad_virus_reverse_rate_index,0.0])

#push!(rate_constant_update_array,[LEF1_virus_forward_rate_index,0.0])
#push!(rate_constant_update_array,[LEF1_virus_reverse_rate_index,0.0])

# TGFB3 dosing

push!(rate_constant_update_array,[TGFB3_foward_rate_index,0.0])
push!(rate_constant_update_array,[TGFB3_reverse_rate_index,0.0])

##
# TGFB signaling
push!(rate_constant_update_array,[tgfb3_binding_rate_index,3.5]) # TGFB3_binding: TGFB3+R =([])=> TGFB3_R
push!(rate_constant_update_array,[tgfb3_to_smad_gain_rate_index,tgfb3_to_smad_gain_rate]) # TGFB3_binding: TGFB3+R =([])=> TGFB3_R

push!(rate_constant_update_array,[tgfb12_to_smad_rate_index,tgfb12_to_smad_rate]) # TGFB3_binding: TGFB3+R =([])=> TGFB3_R

push!(rate_constant_update_array,[R_deg_rate_index,0.01]) # R_degradation: R =([])=> []
push!(rate_constant_update_array,[TGFB12_R_activation_RAS_rate_index,1.0]) #TGFB12_R_activation_RAS: RAS =(TGFB12_R)=> RAS_GTP
push!(rate_constant_update_array,[RAS_activation_RAF_rate_index,tgfb12_to_ap1_sp1_gain]) # RAS_GTP_phosphorylation_RAF: RAF =(RAS_GTP)=> RAF_P
push!(rate_constant_update_array,[RAF_activation_MAPK_rate_index,tgfb12_to_ap1_sp1_gain]) # RAF_P_phosphorylation_MAPK: MAPK =(RAF_P)=> pMAPK
push!(rate_constant_update_array,[MAPK_activation_AP1_SP1_rate_index,tgfb12_to_ap1_sp1_gain]) # MAPK_phosphorylation_AP1_SP1: AP1_SP1 =(pMAPK)=> AP1_SP1_P
push!(rate_constant_update_array,[MAPK_activation_smad_rate_index,MAPK_activation_smad_rate]) # MAPK_phosphorylation_AP1_SP1: AP1_SP1 =(pMAPK)=> AP1_SP1_P

# VEGF signaing
push!(rate_constant_update_array,[PKC_activation_RAF_rate_index,0.0]) # PKC_activation_RAF: RAF =(Active_PKC)=> RAF_P

# Bcatenin complex formation
push!(rate_constant_update_array,[bcat_ecad_forward_index,ecad_complex_gain]) # 19	Binding_Ecadherin_Bcatenin: Ecadherin+Bcatenin =([])=> Ecadherin_Bcatenin
push!(rate_constant_update_array,[bcat_ecad_reverse_index,ecad_complex_reverse]) # 20	Binding_Ecadherin_Bcatenin: Ecadherin_Bcatenin =([])=> Ecadherin+Bcatenin (reverse)
push!(rate_constant_update_array,[bcat_tcf4_forward_index,tcf4_complex_gain]) # 21	TCF4_Bcatenin_complex_formation: TCF4+Bcatenin =([])=> Active_TCF4
push!(rate_constant_update_array,[bcat_tcf4_reverse_index,1.0]) # 22	TCF4_Bcatenin_complex_formation: Active_TCF4 =([])=> TCF4+Bcatenin (reverse)
push!(rate_constant_update_array,[bcat_lef1_forward_index,lef1_complex_gain]) # 23	LEF1_Bcatenin_complex_formation: LEF1+Bcatenin =([])=> Active_LEF1
push!(rate_constant_update_array,[bcat_lef1_reverse_index,1.0]) # 24	LEF1_Bcatenin_complex_formation: Active_LEF1 =([])=> LEF1+Bcatenin (reverse)

push!(rate_constant_update_array,[bcat_deg_rate_index,bcat_deg_rate]) # Bcatenin_degradation: Bcatenin =([])=> []
#push!(rate_constant_update_array,[bcat_deg_gsk3_mediated_index,bcat_deg_gsk3_rate])
#push!(rate_constant_update_array,[pmapk_inact_gsk3_index,pmapk_inact_gsk3_rate])

push!(rate_constant_update_array,[yreg_deg_rate_index,yreg_deg_rate])
push!(rate_constant_update_array,[bcat_trans_rate_index,bcat_trans_rate])
push!(rate_constant_update_array,[ecad_bcat_degradation_rate_index,ecad_bcat_degradation_rate])

# LEF1
push!(rate_constant_update_array,[lef1_trans_rate_index,lef1_trans_rate])

#
push!(rate_constant_update_array,[apc_complex_forward_rate_index,apc_complex_forward_rate])
push!(rate_constant_update_array,[tgfb3_pik3_rate_index,tgfb3_pik3_rate])
push!(rate_constant_update_array,[gsk3_p_rate_index,gsk3_p_rate])
push!(rate_constant_update_array,[apc_bcat_deg_rate_index,apc_bcat_deg_rate])
push!(rate_constant_update_array,[pi3k_translation_rate_index,pi3k_translation_rate])
push!(rate_constant_update_array,[gsk3_translation_rate_index,gsk3_translation_rate])

push!(rate_constant_update_array,[apc_complex_trans_rate_index,apc_complex_trans_rate])

push!(rate_constant_update_array,[tcf4_trans_rate_index ,tcf4_trans_rate])

push!(rate_constant_update_array,[yreg_translation_rate_index ,yreg_translation_rate])

push!(rate_constant_update_array,[ecad_transc_rate_index,ecad_transc_rate]) # addition_TGFb12: [] =([])=> TGFB12
push!(rate_constant_update_array,[ecad_mrna_deg_rate_index,ecad_mrna_deg_rate]) # addition_TGFb12: [] =([])=> TGFB12

# control parameter array
#num_regs = 15; # number of inducing or repressing intereactions (could have multiple for species)
# do forced update for now
#for i in 1:num_regs
  #push!(control_parameter_update_array[i,1],GAIN_MASTER) # Gain for transcriptional activity
#end
#push!(control_parameter_update_array,[7,GAIN_TCF4_TGFB3])
#push!(control_parameter_update_array,[6,GAIN_BCAT_TGFB3])
#push!(control_parameter_update_array,[5,GAIN_ACTIVE_LEF1_ECAD])

# push modifications to RAM
#println(data_dictionary["INITIAL_CONDITION_ARRAY"])
#println(data_dictionary)
#UpdateArray(data_dictionary,"INITIAL_CONDITION_ARRAY",initial_cond_update_array)
#UpdateArray(data_dictionary,"RATE_CONSTANT_ARRAY",rate_constant_update_array)
#UpdateArray(data_dictionary,"CONTROL_PARAMETER_ARRAY",control_parameter_update_array)
###  I can't figure out how to get data_dictionary to be passed into this streamed file,
### It would seem that the stream opens a new variable space and then merges it with the main variable space
### Hence, these will need to be called post-stream

#


const NInts{N} = Tuple{Vararg{Union{Int,StaticInt},N}}
struct PtrArray{T,N,S<:NInts{N},X<:NInts{N}} <: DenseArray{T,N}
  ptr::Ptr{T}
  size::S
  strides::X
end
ispointerstruct(::Type{<:PtrArray}) = true
# `(... + 7) & -8` is eliminated whenever `sizeof(T) == 8`
@inline offset(A::PtrArray{T}) where {T} = (sizeof(T) * length(A) + 7) & -8

@inline Base.size(A::PtrArray) = map(Int, getfield(A,:size))
@inline Base.length(A::PtrArray) = prod(size(A))
# @inline Base.axes(A::PtrArray) = map(CloseOpen, getfield(A,:size))
@inline Base.axes(A::PtrArray) = map(SafeCloseOpen, getfield(A,:size))
@inline Base.strides(A::PtrArray) = map(Int, getfield(A,:strides))
@inline Base.unsafe_convert(::Type{Ptr{T}}, A::PtrArray{T}) where {T} = getfield(A,:ptr)
@inline Base.pointer(A::PtrArray) = getfield(A, :ptr)

@inline sxcumprod(c, sz::Tuple{}) = ()
@inline sxcumprod(c, sz::NInts) = (c, sxcumprod(c * getfield(sz,1), Base.tail(sz))...)
@inline function ptrarray(ptr::Ptr, sz::NInts{N}) where {N}
  sx = sxcumprod(One(), sz)
  PtrArray(ptr, sz, sx), ptr + last(sz)*last(sx)
end
@inline ptrarray(ptr::Ptr, sz::Vararg{Integer,N}) where {N} = ptrarray(ptr, sz)

const PtrVector{T,L} = PtrArray{T,1,Tuple{L},Tuple{StaticInt{1}}}
const PtrMatrix{T,M,N} = PtrArray{T,2,Tuple{M,N},Tuple{StaticInt{1},M}}
const DPtrVector{T} = PtrVector{T,Int}
const DPtrMatrix{T} = PtrMatrix{T,Int,Int}

@inline filtercolon(p::Ptr, s::Tuple{}, x::Tuple{}, i::Tuple{}) = p, (), ()
@inline function filtercolon(p::Ptr{T}, s::Tuple{S}, x::Tuple{X}, i::Tuple{AbstractUnitRange}) where {T,S,X}
  r = getfield(i,1)
  x1 = getfield(x,1)
  p + sizeof(T)*(x1*first(r)), (length(r),), (x1,)
end
@inline filtercolon(p::Ptr, s::Tuple{S}, x::Tuple{X}, i::Tuple{Colon}) where {S,X} = p, (getfield(s,1),), (getfield(x,1),)
@inline filtercolon(p::Ptr{T}, s::Tuple{S}, x::Tuple{X}, i::Tuple{Integer}) where {S,T,X} = p + sizeof(T)*(getfield(x,1)*getfield(i,1)), (), ()

@inline function filtercolon(p::Ptr{T}, s::Tuple{S1,S2,Vararg}, x::Tuple{X1,X2,Vararg}, i::Tuple{I1,I2,Vararg}) where {T,S1,S2,X1,X2,I1<:AbstractUnitRange,I2}
  p2, s2, x2 = filtercolon(p, Base.tail(s), Base.tail(x), Base.tail(i))
  r = getfield(i,1); x1 = getfield(x,1)
  p2 + sizeof(T)*(x1*first(r)), (length(r), s2...), (x1, x2...)
end
@inline function filtercolon(p::Ptr, s::Tuple{S1,S2,Vararg}, x::Tuple{X1,X2,Vararg}, i::Tuple{Colon,I1,Vararg}) where {S1,S2,X1,X2,I1}
  p2, s2, x2 = filtercolon(p, Base.tail(s), Base.tail(x), Base.tail(i))
  p2, (getfield(s,1), s2...), (getfield(x,1), x2...)
end
@inline function filtercolon(p::Ptr{T}, s::Tuple{S1,S2,Vararg}, x::Tuple{X1,X2,Vararg}, i::Tuple{Integer,I1,Vararg}) where {T,S1,S2,X1,X2,I1}
  p2, s2, x2 = filtercolon(p, Base.tail(s), Base.tail(x), Base.tail(i))
  p2 + sizeof(T)*(getfield(x,1)*getfield(i,1)), s2, x2
end

@inline _getindex(p::Ptr, s::Tuple{}, x::Tuple{}) = unsafe_load(p)
@inline _getindex(p::Ptr, s::Tuple{S,Vararg}, x::Tuple{X,Vararg}) where {S,X} = PtrArray(p, s, x)
Base.@propagate_inbounds function Base.getindex(A::PtrArray{T,N}, i::Vararg{Union{Integer,Colon},N}) where {T,N}
  Base.@boundscheck checkbounds(A, i...)
  p, s, x = filtercolon(pointer(A), A.size, A.strides, i)
  _getindex(p, s, x)
end
Base.@propagate_inbounds function Base.getindex(A::PtrArray{T,N}, i::Integer) where {T,N}
  Base.@boundscheck checkbounds(A, i)
  _getindex(pointer(A) + i*sizeof(T), (), ())
end
Base.@propagate_inbounds function Base.getindex(A::PtrArray{T,1}, i::Integer) where {T}
  Base.@boundscheck checkbounds(A, i)
  _getindex(pointer(A) + i*sizeof(T), (), ())
end
Base.@propagate_inbounds function Base.getindex(A::PtrArray{T,1}, r::AbstractUnitRange) where {T}
  Base.@boundscheck checkbounds(A, r)
  p, s, x = filtercolon(pointer(A), A.size, A.strides, (r,))
  _getindex(p, s, x)
end

@inline _setindex!(p::Ptr{T}, v, ::Tuple{}, ::Tuple{}) where {T} = unsafe_store!(p, convert(T, v))
@inline _setindex!(p::Ptr{T}, v::AbstractArray{<:Any,0}, ::Tuple{}, ::Tuple{}) where {T} = unsafe_store!(p, convert(T, v[]))
@inline function _setindex!(p::Ptr, v::AbstractArray{<:Any,N}, s::NInts{N}, x::NInts{N}) where {N}
  d = PtrArray(p, s, x)
  for i in eachindex(d, v)
    @inbounds d[i] = v[i]
  end
end

Base.@propagate_inbounds function Base.setindex!(A::PtrArray{T,N}, v, i::Vararg{Union{Integer,Colon},N}) where {T,N}
  Base.@boundscheck checkbounds(A, i...)
  p, s, x = filtercolon(pointer(A), A.size, A.strides, i)
  _setindex!(p, v, s, x)
end
Base.@propagate_inbounds function Base.setindex!(A::PtrArray{T,N}, v, i::Integer) where {T,N}
  Base.@boundscheck checkbounds(A, i)
  _setindex!(pointer(A) + sizeof(T)*i, v, (), ())
end
Base.@propagate_inbounds function Base.setindex!(A::PtrArray{T,1}, v, i::Integer) where {T}
  Base.@boundscheck checkbounds(A, i)
  _setindex!(pointer(A) + sizeof(T)*i, v, (), ())
end

Base.@propagate_inbounds Base.getindex(A::PtrArray{<:Any,N}, I::CartesianIndices{N}) where {N} = A[Tuple(I)...]
Base.@propagate_inbounds Base.setindex!(A::PtrArray{<:Any,N}, v, I::CartesianIndices{N}) where {N} = A[Tuple(I)...] = v

function Base.getindex(a::PtrVector{T}, i::PtrVector{I}) where {T,I<:Integer}
  b = Vector{T}(undef, length(a))
  @inbounds for j ∈ eachindex(i)
    b[j+1] = a[i[j]]
  end
  return b
end

_indstyle(::StaticInt{1}) = Base.IndexLinear()
_indstyle(_) = Base.IndexCartesian()
Base.IndexStyle(A::PtrVector) = _indstyle(getfield(A.strides,1))

Base.LinearIndices(x::PtrVector) = CloseOpen(length(x))
function Base.Array(A::PtrArray{T}) where {T}
  B = Array{T}(undef, size(A))
  @inbounds for I ∈ eachindex(A)
    B[I + one(I)] = A[I]
  end
  return B
end
Base.Array{<:Any,N}(A::PtrArray{<:Any,N}) where {N} = Array(A)


"""
    fit!(model, obs::AbstractArray)

Fit a model to data with least squares regression, using `curve_fit` from
LsqFit.jl. The passed in model should be initialised with sensible defaults,
these will be used as the initial parameters for the optimization.

Any (nested) `Real` fields on the struct are flattened to a parameter vector using
[Flatten.jl](http://github.com/rafaqz/Flatten.jl). Fields can be marked to ignore
using the `@flattenable` macro from [FieldMetadata.jl](http://github.com/rafaqz/FieldMetadata.jl).

## Arguments

- `model`: Any constructed [`RateModel`](@ref) or a `Tuple` of `RateModel`.
- `obs`: A `Vector` of `(val, rate)` tuples where `val` is the value of the
  x-axis variable (such as temperature), and `rate` is the growth rate observed.

## Returns

An updated `Model` containing the fitted parameter values.
"""
function fit!(model::Model, obs::AbstractArray)
    fit = curve_fit(first.(obs), last.(obs), collect(model[:val])) do xs, vals
        model[:val] = vals
        GrowthMaps.conditionalrate.(Ref(stripparams(model)), xs)
    end
    model[:val] = fit.param
    return model
end

"""
    manualfit!(model::Model, data::NamedTuple; obs=[], kw...) =

Returns the fitted model.

# Arguments

- `obs`: A `Vector` of `(val, rate)` tuples where `val` is the value of the
  x-axis variable (such as temperature), and `rate` is the growth rate observed.
- `data`: A `NamedTuple` of `AbstractVector`. The `NamedTuple` key/s must match
  the key required by the `Layer`/s.
- `obs`: Optional observations to scatter-plot over the curve. A `Vector` of `(val, rate)`
  tuples where `val` is the value of the x-axis variable (such as temperature), and `rate`
  is the growth rate observed.
- `kwargs`: passed to `plot`

# Example

```julia
p = 3e-01
ΔH_A = 3e4cal/mol
ΔH_L = -1e5cal/mol
ΔH_H = 3e5cal/mol
Thalf_L = 2e2K
Thalf_H = 3e2K
T_ref = K(25.0°C)
growthmodel = SchoolfieldIntrinsicGrowth(p, ΔH_A, ΔH_L, Thalf_L, ΔH_H, Thalf_H, T_ref)
model = Model(Layer(:surface_temp, K, growthmodel))
obs = []

tempdata=(surface_temp=(270.0:0.1:310.0)K,)
manualfit!(model, tempdata; obs=obs)

To use the interface in a desktop app, use Blink.jl:

```julia; eval=false
using Blink
w = Blink.Window()
body!(w, interface)
```
"""
function manualfit!(
    model::AbstractModel, data::NamedTuple{<:Any,<:Tuple{Vararg{<:AbstractVector}}};
    observations=[],
    throttle=0.1,
    kwargs...
)
    InteractModel(model; throttle=throttle, submodel=RateModel) do updated_model
        @show params(updated_model)
        ModelParameters.setparent!(model, updated_model)
        manualfit(stripparams(updated_model), (observations, data); kwargs...)
    end
end

function manualfit(model, (observations, data); kwargs...)
    predictions = combinelayers(model, data)
    p = plot(first(data), predictions; label="predicted", legend=false, kwargs...)
    scatter!(p, observations; label="observed")
    return p
end

"""
    mapfit!(model::Model, modelkwargs;
            occurrence=[],
            precomputed=nothing,
            throttle=0.1,
            window=(Band(1),),
            kwargs...
    )

Fit a model to the map.

# Arguments

- `occurence`: a `Vector` of occurence locations, as `(lon, lat)` tuples.
- `modelkwargs`: are passed to the `mapgrowth` with the model.
- `mapkwargs`: are passed to the `plot` function the plots the `GeoArray`
- `throttle`: the response time of Interact.jl sliders.
- `window`: selects a window of the output to plot. By default this is `(Band(1),)`,
  which just removes the `Band` dimension from a `GeoArray`, if it exists.
- `kwargs`: passed to `Plots.scatter!`

# Example

```julia
model = Model(wiltstress, coldstress, heatstress)
interface = mapfit!(model, modelkwargs;
    occurrence=occurrence,
    precomputed=precomputed,
    throttle=0.2,
    window=(Band(1),),
    markershape=:cross,
    markercolor=:lightblue,
    markeropacity=0.4
)
display(interface)
```

To use the interface in a desktop app, use Blink.jl:

```julia; eval=false
using Blink
w = Blink.Window()
body!(w, interface)
```
"""
function mapfit!(model::AbstractModel, modelkwargs;
    mapkwargs=(),
    occurrence=[],
    precomputed=nothing,
    throttle=0.1,
    scatterkwargs...
)
    title = "GrowthMaps: mapfit interface"
    InteractModel(model; throttle=throttle, title=title, submodel=RateModel) do updated_model
        ModelParameters.setparent!(model, updated_model)
        mapfit(
            stripparams(updated_model), (modelkwargs, occurrence, precomputed);
            scatterkwargs...
        )
    end
end

function mapfit(model, (modelkwargs, occurrence, precomputed);
    window=(Band(1),),
    levels=10,
    mapkwargs=(),
    markercolor=:white,
    markersize=2.0,
    clims=(0.0, 0.25),
    scatterkwargs...
)
    output = mapgrowth(model; modelkwargs...)
    output = isnothing(precomputed) ? output : output .+ precomputed
    windowed = output[window...]
    p = plot(windowed; legend=:none, levels=levels, clims=clims, mapkwargs...)
    scatter(; t...) = scatter!(
        p, occurrence;
        t..., markercolor=markercolor, markersize=markersize, scatterkwargs...
    )
    if hasdim(windowed, Ti())
        for t in 1:length(dims(windowed, Ti()))
            scatter(; subplot=t)
        end
    else
        scatter()
    end
    return p
end


import Adapt
import CUDAapi
import Flux
import GPUArrays
import NNlib
import ProgressMeter
import ValueHistories
import Zygote


macro ArbFloat(x)
    convert(ArbFloat, string(:($x)))
end
macro ArbFloat(p,x)
    convert(ArbFloat{:($p)}, string(:($x)))
end

# interconvert Arb with Arf

function convert{P}(::Type{ArbFloat{P}}, x::ArfFloat{P})
    z = initializer(ArbFloat{P})
    ccall(@libarb(arb_set_arf), Void, (Ptr{ArbFloat{P}}, Ptr{ArfFloat{P}}), &z, &x)
    z
end

function convert{P}(::Type{ArfFloat{P}}, x::ArbFloat{P})
    z = initializer(ArfFloat{P})
    z.mid_exp  = x.mid_exp
    z.mid_size = x.mid_size
    z.mid_d1   = x.mid_d1
    z.mid_d2   = x.mid_d2
    z
end

#interconvert ArbFloat{P} with ArbFloat{Q}

function convert{P,Q}(::Type{ArbFloat{Q}}, a::ArbFloat{P})
    if (Q < P)
        a = round(a, Q, 2)
    end

    z = initializer(ArbFloat{Q})
    z.mid_exp  = a.mid_exp
    z.mid_size = a.mid_size
    z.mid_d1   = a.mid_d1
    z.mid_d2   = a.mid_d2
    z.rad_exp  = a.rad_exp
    z.rad_man  = a.rad_man

    z
end

#

function convert{P}(::Type{ArbFloat{P}}, x::UInt)
    z = initializer(ArbFloat{P})
    ccall(@libarb(arb_set_ui), Void, (Ptr{ArbFloat{P}}, UInt), &z, x)
    z
end
if sizeof(Int)==sizeof(Int64)
   convert{P}(::Type{ArbFloat{P}}, x::UInt32) = convert(ArbFloat{P}, convert(UInt64,x))
else
   convert{P}(::Type{ArbFloat{P}}, x::UInt64) = convert(ArbFloat{P}, convert(UInt32,x))
end
convert{P}(::Type{ArbFloat{P}}, x::UInt16) = convert(ArbFloat{P}, convert(UInt,x))

function convert{P}(::Type{ArbFloat{P}}, x::Int)
    z = initializer(ArbFloat{P})
    ccall(@libarb(arb_set_si), Void, (Ptr{ArbFloat{P}}, Int), &z, x)
    z
end
if sizeof(Int)==sizeof(Int64)
   convert{P}(::Type{ArbFloat{P}}, x::Int32) = convert(ArbFloat{P}, convert(Int64,x))
else
   convert{P}(::Type{ArbFloat{P}}, x::Int64) = convert(ArbFloat{P}, convert(Int32,x))
end
convert{P}(::Type{ArbFloat{P}}, x::Int16) = convert(ArbFloat{P}, convert(Int,x))

function convert{P}(::Type{ArbFloat{P}}, x::Float64)
    z = initializer(ArbFloat{P})
    fp=copy(x)
    ccall(@libarb(arb_set_d), Void, (Ptr{ArbFloat{P}}, Float64), &z, fp)
    z
end
convert{P}(::Type{ArbFloat{P}}, x::Float32) = convert(ArbFloat{P}, convert(Float64,x))
convert{P}(::Type{ArbFloat{P}}, x::Float16) = convert(ArbFloat{P}, convert(Float64,x))

function convert{P}(::Type{ArbFloat{P}}, x::String)
    s = String(x)
    z = initializer(ArbFloat{P})
    ccall(@libarb(arb_set_str), Void, (Ptr{ArbFloat}, Ptr{UInt8}, Int), &z, s, P)
    z
end

convert(::Type{BigInt}, x::String) = parse(BigInt,x)
convert(::Type{BigFloat}, x::String) = parse(BigFloat,x)

function convert{P}(::Type{ArbFloat{P}}, x::BigFloat)
     x = round(x,P,2)
     s = string(x)
     ArbFloat{P}(s)
end

function convert{P}(::Type{BigFloat}, x::ArbFloat{P})
     s = string(midpoint(x))
     round(parse(BigFloat, s),P,2)
end

function convert{I<:Integer,P}(::Type{Rational{I}}, x::ArbFloat{P})
    bf = convert(BigFloat, x)
    convert(Rational{I}, bf)
end

convert{P}(::Type{ArbFloat{P}}, x::BigInt)   = convert(ArbFloat{P}, convert(BigFloat,x))
convert{P}(::Type{ArbFloat{P}}, x::Rational) = convert(ArbFloat{P}, convert(BigFloat,x))
convert{P,S}(::Type{ArbFloat{P}}, x::Irrational{S}) =
    convert(ArbFloat{P}, convert(BigFloat,x))

#= returns 256.0 for convert(big(1.5))
function convert{P}(::Type{ArbFloat{P}}, x::BigFloat)
    z = initializer(ArbFloat{P})
    ccall(@libarb(arb_set_round_fmpz), Void, (Ptr{ArbFloat}, Ptr{BigFloat}, Int), &z, &x, P)
    z
end
=#

convert{P}(::Type{ArbFloat{P}}, y::ArbFloat{P}) = y

function convert{P}(::Type{Float64}, x::ArbFloat{P})
    s = smartarbstring(x)
    parse(Float64,s)
end
convert{P}(::Type{Float32}, x::ArbFloat{P}) = convert(Float32,convert(Float64,x))

function convert{I<:Integer,P}(::Type{I}, x::ArbFloat{P})
    s = smartarbstring(x)
    parse(I,split(s,".")[1])
end

for T in (:Int128, :Int64, :Int32, :Int16, :Float64, :Float32, :Float16,
          :(Rational{Int64}), :(Rational{Int32}), :(Rational{Int16}),
          :String)
  @eval convert(::Type{ArbFloat}, x::$T) = convert(ArbFloat{precision(ArbFloat)}, x)
end

# Promotion
for T in (:Int128, :Int64, :Int32, :Int16, :Float64, :Float32, :Float16,
          :(Rational{Int64}), :(Rational{Int32}), :(Rational{Int16}),
          :String)
  @eval promote_rule{P}(::Type{ArbFloat{P}}, ::Type{$T}) = ArbFloat{P}
end

promote_rule{P}(::Type{ArbFloat{P}}, ::Type{BigFloat}) = BigFloat
promote_rule{P}(::Type{ArbFloat{P}}, ::Type{BigInt}) = ArbFloat{P}
promote_rule{P}(::Type{ArbFloat{P}}, ::Type{Rational{BigInt}}) = Rational{BigInt}

promote_rule{P,Q}(::Type{ArbFloat{P}}, ::Type{ArbFloat{Q}}) =
    ifelse(P>Q, ArbFloat{P}, ArbFloat{Q})


###################################################
## network/iterative/greedy.jl
## Iterative state with fixed subse
###################################################

"""
    `GreedyEdges`, Iterative state optimizing on static trip list (subset of data)
    - each trip has a fixed max number of paths
    - Default number of paths per trip is infinite
    - Set of independent edges is updated at every iteration
"""
mutable struct GreedyEdges <: IterativeState
    # inherited attributes
    data::NetworkData
    timings::NetworkTimings
    trips::Vector{NetworkTrip}
    paths::Vector{Vector{Dict{Edge,Float64}}}  # for each trip, a vector of paths
    pathDiff::Float64

    "max number of paths per trip" #temporarily cannot be more than 1
    pathsPerTrip::Int64
    "independent edges"
    independent::Vector{Edge}
    "dependent edges"
    dependent::Vector{Edge}
    "original non-simplified paths"
    origPaths::Vector{Vector{Dict{Edge,Float64}}}
    "dependency matrix"
    dependencies::SparseMatrixCSC{Float64,Int}
    "edge map"
    edgeMap::Dict{Edge,Int}
    "number of edges to remove"
    numEdges::Int
    "number of iterations before removal"
    numIter::Int
    "minimum number of independent edges"
    minIndep::Int
    "state variable, remembers whether or not to update"
    state::Int
    "number of independent edges each dependent edge depends on"
    numDeps::Int
end

"""
    Creates initial fixed number of paths per trip state
    takes in NetworkData object, initial timings (as link times or full timings),
    and max number of paths per trip as integer
"""
function GreedyEdges(data::NetworkData, startSolution::NetworkTimings; pathsPerTrip::Int = typemax(Int), maxTrip::Int=1000, numEdges::Int = 10, numIter::Int = 3, minIndep::Int = 100, numDeps::Int = 3)
    if pathsPerTrip < 1
        error("Must have at least one path per trip")
    end
    # randomly select the trips
    Random.seed!(1991)
    trips = shuffle(data.trips)[1:min(maxTrip,length(data.trips))]
    # One path per trip: the initial shortest path
    paths = [Dict{Edge,Float64}[getFullPathEdges(t, startSolution)] for t in trips]
    println(typeof(paths))
    origPaths = [Dict{Edge,Float64}[getFullPathEdges(t, startSolution)] for t in trips]
    println(typeof(origPaths))
    independent = collect(edges(data.network.graph))
    dependent = Edge[]
    dependencies, edgeMap = findNetworkDependence(data.network, independent, dependent, numDeps = numDeps)
    return GreedyEdges(data,startSolution,trips,paths,Inf,pathsPerTrip,independent, dependent,origPaths,dependencies,edgeMap,numEdges, numIter, minIndep, 0, numDeps)
end

GreedyEdges(data::NetworkData, initTimes::AbstractArray{Float64, 2}; args...) =
    GreedyEdges(data, NetworkTimings(data.network, initTimes); args...)

"""
    Update paths of GreedyEdges object given new times
"""
function updateState!(s::GreedyEdges, times::AbstractArray{Float64, 2})
    # update state variable
    s.state += 1
    # update the timings and compute shortest paths
    newTimes = evaluateTimes(s.data.network, s.dependencies, times, s.independent, s.edgeMap)
    s.timings = NetworkTimings(s.data.network, newTimes)
    # update independent set and dependencies if necessary
    if s.state == s.numIter
        s.independent, s.dependent = updateIndependentEdges(s.paths, s.independent, s.dependent, s.numEdges, s.minIndep)
        s.dependencies, s.edgeMap = findNetworkDependence(s.data.network, s.independent, s.dependent, numDeps = s.numDeps)
        s.state = 0
        origPaths = [Dict{Edge,Float64}[getFullPathEdges(t, s.timings)] for t in s.trips]
        s.pathDiff = sum([length(symdiff(collect(keys(origPaths[d][1])), collect(keys(s.origPaths[d][1])))) for d=eachindex(s.trips)])/(2 * length(s.trips))
        s.origPaths = origPaths
    else
        s.pathDiff = 0.
        # update paths normally
        for (d,t) in enumerate(s.trips)
            sp = getFullPathEdges(t, s.timings)
            s.pathDiff += length(symdiff(collect(keys(sp)), collect(keys(s.origPaths[d][1]))))/2
            if s.pathsPerTrip == 1
                s.origPaths[d][1] = sp
            else
                index = findfirst(isequal(sp), s.origPaths[d])
                if index != nothing   # if so, put it in first position
                    s.origPaths[d][index], s.origPaths[d][1] = s.origPaths[d][1], s.origPaths[d][index] #swap
                elseif length(s.origPaths[d]) < s.pathsPerTrip # if not, and if enough room to add, add it in first position
                    pushfirst!(s.origPaths[d], sp)
                else        # replace least useful path
                    worstIndex = argmax([pathEdgesTime(s.timings, collect(keys(p))) for p in s.paths[d]])
                    s.origPaths[d][worstIndex] = s.origPaths[d][1]
                    s.origPaths[d][1] = sp
                end
            end
        end
        s.pathDiff /= length(s.trips)
    end
    s.paths = [[simplifyPath(p, s.independent, s.dependencies, s.edgeMap) for p in s.origPaths[d]] for d = eachindex(s.trips)]
end


using Plots
path = dirname(@__FILE__)

# Initialize the plotting backend
gr()

# Initialize model object
m = AnSchorfheide(testing = true)
m <= Setting(:saveroot, tempdir())
m <= Setting(:date_forecast_start, quartertodate("2015-Q4"))
m <= Setting(:date_conditional_end, quartertodate("2015-Q4"))
m <= Setting(:use_population_forecast, true)
m <= Setting(:forecast_horizons, 12)
m <= Setting(:impulse_response_horizons, 8)

# Run full-distribution forecast
estroot = normpath(joinpath(dirname(@__FILE__), "..", "reference"))
overrides = forecast_input_file_overrides(m)
overrides[:mode] = joinpath(estroot, "optimize.h5")
overrides[:full] = joinpath(estroot, "metropolis_hastings.h5")
output_vars = add_requisite_output_vars([:histobs, :forecastobs, :shockdecobs, :irfobs])

@everywhere using DSGE
m <= Setting(:forecast_block_size, 5)
forecast_one(m, :full, :none, output_vars, verbose = :none)
compute_meansbands(m, :full, :none, output_vars; verbose = :none)
println("The following warning is expected test behavior:")

# Plot history and forecast
plot_history_and_forecast(m, :obs_nominalrate, :obs, :full, :none,
                          bdd_and_unbdd = true,
                          start_date = DSGE.quartertodate("2007-Q1"),
                          verbose = :none)

# Plot forecast comparison
plot_forecast_comparison(m, m, :obs_nominalrate, :obs, :full, :none,
                         bdd_and_unbdd = true,
                         start_date = DSGE.quartertodate("2007-Q1"),
                         verbose = :none)

# Plot shock decomposition
plot_shock_decomposition(m, :obs_nominalrate, :obs, :full, :none, verbose = :none)

# Plot IRF
plot_impulse_response(m, :rm_sh, collect(keys(m.observables)), :obs, :full, :none,
                      verbose = :none)
plot_impulse_response(m, :rm_sh, collect(keys(m.observables)), :obs, :full, :none,
                      flip = true, verbose = :none)

# Plot scenario, zeroing out measurement error
for para in [:e_y, :e_π, :e_R]
    m[para].value = 0
end
alt = Scenario(:altscen, "Test Alternative Scenario", [:obs_gdp, :obs_cpi], [:g_sh, :rm_sh], "REF")
forecast_scenario(m, alt, verbose = :none)
scenario_means_bands(m, alt, verbose = :none)
plot_scenario(m, :obs_nominalrate, :obs, alt, untrans = true, verbose = :none)
plot_scenario(m, :obs_nominalrate, :obs, alt, verbose = :none)
plot_scenario(m, :obs_nominalrate, :obs, alt, fourquarter = true, verbose = :none)
@testset "Test scenario exception" begin
    @test_throws ErrorException plot_scenario(m, :obs_nominalrate, :obs, alt,
                                              untrans = true, fourquarter = true)
end

# Hair plot
realized = load_data(m, verbose = :none)
hist_mb = read_mb(m, :full, :none, :histobs)
fcast_mb = read_mb(m, :full, :none, :bddforecastobs)
hair_plot(:obs_nominalrate, realized, [hist_mb], [fcast_mb];
          plotroot = saveroot(m), verbose = :none)


using JSON
using Random
using CUDA
using Statistics
using Dates
using DataStructures

include("environments/collect_points.jl")
include("environments/dummy_app.jl")
include("brains/continuous_time_rnn.jl")
include("optimizers/cma_es.jl")
include("optimizers/randomizer.jl")
include("tools/episode_runner.jl")
include("tools/write_results.jl")

struct TrainingCfg
    number_generations::Int
    number_validation_runs::Int
    number_rounds::Int
    maximum_env_seed::Int
    environment::OrderedDict
    brain::OrderedDict
    optimizer::OrderedDict
    experiment_id::Int

    function TrainingCfg(configuration::OrderedDict)
        new(
            configuration["number_generations"],
            configuration["number_validation_runs"],
            configuration["number_rounds"],
            configuration["maximum_env_seed"],
            configuration["environment"],
            configuration["brain"],
            configuration["optimizer"],
            get(configuration, "experiment_id", -1),
        )
    end
end

function main()

    configuration_file = "configurations/CMA_ES_Deap_CTRNN_Dense_Dummy_App.json"

    # Load configuration file
    configuration = JSON.parsefile(configuration_file, dicttype = OrderedDict)

    config = TrainingCfg(configuration)

    # Get environment type from configuration
    # TODO: Choose environment type from configuration
    environment_type = DummyApp

    # Get brain type from configuration
    # TODO: Choose brain type from configuration
    brain_type = ContinuousTimeRNN

    # Get optimizer type from configuration
    if configuration["optimizer"]["type"] == "Randomizer"
        optimizer_type = OptimizerRandom
    elseif configuration["optimizer"]["type"] == "CMA-ES-Deap"
        optimizer_type = OptimizerCmaEs
    else
        error("No valid optimizer type")
    end

    number_individuals = config.optimizer["population_size"]

    # Initialize environments
    environments = environment_type(config.environment, number_individuals)

    number_inputs = get_number_inputs(environments)
    number_outputs = get_number_outputs(environments)

    # Initialize brains
    brains = brain_type(config.brain, number_inputs, number_outputs, number_individuals)

    individual_size = get_individual_size(brains)

    # Initialize optimizer
    optimizer = optimizer_type(individual_size, config.optimizer)

    best_genome_overall = nothing
    best_reward_overall = typemin(Int32)

    required_shared_memory = get_memory_requirements(brains) + get_memory_requirements(environments)

    # Get start time of training and date
    start_time_training = now()

    log = OrderedDict()

    # Run evolutionary training for given number of generations
    for generation = 1:config.number_generations

        start_time_generation = time()

        # Environment seed for this generation (excludes validation environment seeds)
        env_seed = Random.rand(config.number_validation_runs:config.maximum_env_seed)

        # Ask optimizer for new population
        individuals = ask(optimizer)

        # Training runs for current generation
        rewards_training = training_runs(
            individuals,
            number_individuals,
            brains,
            environments,
            env_seed = env_seed,
            number_rounds = config.number_rounds,
            threads = brains.number_neurons,
            blocks = number_individuals,
            shared_memory = required_shared_memory,
        )

        # Tell optimizer new rewards
        tell(optimizer, rewards_training)

        best_genome_current_generation = individuals[(findmax(rewards_training))[2], :]

        # Validation runs for best individual in current generation
        rewards_validation = validation_runs(
            best_genome_current_generation,
            individual_size,
            config.number_validation_runs,
            brains,
            environments,
            threads = brains.number_neurons,
            blocks = config.number_validation_runs,
            shared_memory = required_shared_memory,
        )

        best_reward_current_generation = mean(rewards_validation)

        # Better individual found in current generation?
        if best_reward_current_generation > best_reward_overall
            best_genome_overall = best_genome_current_generation
            best_reward_overall = best_reward_current_generation
        end

        # Logging and printing
        elapsed_time_current_generation = round(time() - start_time_generation, digits = 3)
        log_line = OrderedDict()
        log_line["gen"] = generation
        log_line["min"] = minimum(rewards_training)
        log_line["mean"] = mean(rewards_training)
        log_line["max"] = maximum(rewards_training)
        log_line["best"] = best_reward_overall
        log_line["elapsed_time"] = elapsed_time_current_generation
        log[generation] = log_line
        println(
            "Generation: ", generation,
            "   Min: ", minimum(rewards_training),
            "   Mean: ", mean(rewards_training),
            "   Max: ", maximum(rewards_training),
            "   Best: ", best_reward_overall,
            "   Elapsed time: ", elapsed_time_current_generation, "s"
        )
    end

    result_directory =
        "Simulation_results/" *
        Dates.format(floor(start_time_training, Second), DateFormat("yyyy-mm-dd_HH-MM-SS"))

    mkdir(result_directory)
    write_results_to_textfile(
        result_directory * "/Log.txt",
        configuration,
        log,
        number_inputs,
        number_outputs,
        number_individuals,
        individual_size,
        now() - start_time_training,
    )
end

main()


"""
    couple_models(aero::LiftingLine, stru::GEBT)

Create a coupled model using a lifting line aerodynamic model and a
geometrically exact beam theory model.  This model introduces additional
parameters corresponding to the freestream air density ``\\rho_\\infty``,
followed by the external loads ``F_{x,i}, F_{y,i}, F_{z,i}, M_{x,i}, M_{y,i},
M_{z,i}`` or displacements ``u_{x,i}, u_{y,i}, u_{z,i}, \\theta_{x,i},
\\theta_{y,i}, \\theta_{z,i}`` for each node, followed by the constant
distributed loads ``f_{x,i}, f_{y,i}, f_{z,i}, m_{x,i}, m_{y,i}, m_{z,i}``
applied on each beam element (excluding aerodynamic loads), followed by the
body frame linear and angular velocities ``u, v, w, p, q, r``.

**NOTE: When using this model, the local frame for each beam element should be
oriented with the x-axis along the beam's axis, the y-axis forward, and the
z-axis normal to the surface**
"""
couple_models(aero::LiftingLine, stru::GEBT) = (aero, stru)

# --- traits --- #

function number_of_additional_parameters(aero::LiftingLine, stru::GEBT)
    return 1 + 6*length(stru.icol_point) + 6*length(stru.icol_elem) + 6
end

function coupling_inplaceness(::Type{LiftingLine{NA,TA}}, ::Type{<:GEBT}) where {NA,TA}
    return InPlace()
end

function coupling_rate_jacobian_type(::Type{LiftingLine{NA,TA}}, ::Type{<:GEBT}) where {NA,TA}
    model_types = TA.parameters
    if all(isempty.(coupling_rate_jacobian_type.(model_types, Ref(LiftingLineSection))))
        return Empty()
    elseif all(iszero.(coupling_rate_jacobian_type.(model_types, Ref(LiftingLineSection))))
        return Zeros()
    elseif all(isidentity.(coupling_rate_jacobian_type.(model_types, Ref(LiftingLineSection))))
        return Identity()
    elseif all(isinvariant.(coupling_rate_jacobian_type.(model_types, Ref(LiftingLineSection))))
        return Invariant()
    elseif all(isconstant.(coupling_rate_jacobian_type.(model_types, Ref(LiftingLineSection))))
        return Constant()
    elseif all(islinear.(coupling_rate_jacobian_type.(model_types, Ref(LiftingLineSection))))
        return Linear()
    else
        return Nonlinear()
    end
end

function coupling_state_jacobian_type(::Type{LiftingLine{NA,TA}}, ::Type{<:GEBT}) where {NA,TA}
    model_types = TA.parameters
    if all(isempty.(coupling_state_jacobian_type.(model_types, Ref(LiftingLineSection))))
        return Empty()
    elseif all(iszero.(coupling_state_jacobian_type.(model_types, Ref(LiftingLineSection))))
        return Zeros()
    elseif all(isidentity.(coupling_state_jacobian_type.(model_types, Ref(LiftingLineSection))))
        return Identity()
    elseif all(isinvariant.(coupling_state_jacobian_type.(model_types, Ref(LiftingLineSection))))
        return Invariant()
    elseif all(isconstant.(coupling_state_jacobian_type.(model_types, Ref(LiftingLineSection))))
        return Constant()
    elseif all(islinear.(coupling_state_jacobian_type.(model_types, Ref(LiftingLineSection))))
        return Linear()
    else
        return Nonlinear()
    end
end

coupling_parameter_jacobian_type(::Type{<:LiftingLine}, ::Type{<:GEBT}) = Nonlinear()

function coupling_time_gradient_type(::Type{LiftingLine{NA,TA}}, ::Type{<:GEBT}) where {NA,TA}
    model_types = TA.parameters
    if all(isempty.(coupling_time_gradient_type.(model_types, Ref(LiftingLineSection))))
        return Empty()
    elseif all(iszero.(coupling_time_gradient_type.(model_types, Ref(LiftingLineSection))))
        return Zeros()
    elseif all(isinvariant.(coupling_time_gradient_type.(model_types, Ref(LiftingLineSection))))
        return Invariant()
    elseif all(isconstant.(coupling_time_gradient_type.(model_types, Ref(LiftingLineSection))))
        return Constant()
    elseif all(islinear.(coupling_time_gradient_type.(model_types, Ref(LiftingLineSection))))
        return Linear()
    else
        return Nonlinear()
    end
end

# --- methods --- #

function get_coupling_inputs!(y, aero::LiftingLine{NA,TA}, stru::GEBT, dx, x, p, t) where {NA,TA}

    # extract number of state variables, inputs, and parameters
    nua = number_of_states(aero) # number of aerodynamic states
    nus = number_of_states(stru) # number of structural states
    nya = number_of_inputs(aero) # number of aerodynamic inputs
    nys = number_of_inputs(stru) # number of structural inputs
    npa = number_of_parameters(aero) # number of aerodynamic parameters
    nps = number_of_parameters(stru) # number of structural parameters
    npadd = number_of_additional_parameters(aero, stru) # number of additional parameters

    # get indices for state variables, inputs, and parameters
    iua = 1:nua # indices of aerodynamic states
    iya = 1:nya # indices of aerodynamic inputs
    ipa = 1:npa # indices of aerodynamic parameters
    ius = nua + 1 : nua + nus # indices of structural states
    iys = nya + 1 : nya + nys # indices of structural inputs
    ips = npa + 1 : npa + nps # indices of structural parameters
    ipadd = npa + nps + 1 : npa + nps + npadd # indices of additional parameters
    iuas = state_indices(aero.models) # indices of aerodynamic states for each section
    iyas = input_indices(aero.models) # indices of aerodynamic inputs for each section
    ipas = parameter_indices(aero.models) # indices of aerodynamic inputs for each section

    # separate state variables, inputs, and parameters
    dua = view(dx, iua) # aerodynamic rates
    ua = view(x, iua) # aerodynamic states
    ya = view(y, iya) # aerodynamic inputs
    pa = view(p, ipa) # aerodynamic parameters
    dus = view(dx, ius) # structural rates
    us = view(x, ius) # structural states
    ys = view(y, iys) # structural inputs
    ps = view(p, ips) # structural parameters
    padd = view(p, ipadd) # additional parameters for coupled model
    duas = view.(Ref(dua), iuas) # aerodynamic rates for each section
    uas = view.(Ref(ua), iuas) # aerodynamic states for each section
    yas = view.(Ref(ya), iyas) # aerodynamic inputs for each section
    pas = view.(Ref(pa), ipas) # aerodynamic parameters for each section

    # number of points and elements
    npoint = length(stru.icol_point)
    nelem = length(stru.icol_elem)

    # global parameters
    ρ = padd[1]
    ur = padd[1 + 6*npoint + 6*nelem + 1]
    vr = padd[1 + 6*npoint + 6*nelem + 2]
    wr = padd[1 + 6*npoint + 6*nelem + 3]
    pr = padd[1 + 6*npoint + 6*nelem + 4]
    qr = padd[1 + 6*npoint + 6*nelem + 5]
    rr = padd[1 + 6*npoint + 6*nelem + 6]

    # body linear and angular velocity
    V = SVector(ur, vr, wr)
    Ω = SVector(pr, qr, rr)

    # construct assembly from structural parameters
    assembly = gxbeam_assembly(ps, npoint, nelem, stru.start, stru.stop)

    # set prescribed point loads/displacements
    for ip = 1:npoint
        yoff = 6*(ip-1)
        poff = 1 + 6*(ip-1)
        ys[yoff+1] = padd[poff+1]
        ys[yoff+2] = padd[poff+2]
        ys[yoff+3] = padd[poff+3]
        ys[yoff+4] = padd[poff+4]
        ys[yoff+5] = padd[poff+5]
        ys[yoff+6] = padd[poff+6]
    end

    # set aerodynamic inputs and distributed loads
    for i = 1:NA
        # models for this section
        section_aero = aero.models[i]
        section_stru = LiftingLineSection()
        section_models = (section_aero, section_stru)

        # model dimensions for this section
        Nuai = number_of_states(section_aero)
        Nyai = number_of_inputs(section_aero)
        Npai = number_of_parameters(section_aero)
        Nusi = number_of_states(section_stru)
        Nysi = number_of_inputs(section_stru)
        Npsi = number_of_parameters(section_stru)

        # section properties
        icol = stru.icol_elem[i]
        element = assembly.elements[i]
        u_elem = SVector(us[icol], us[icol+1], us[icol+2]) # linear displacement
        θ_elem = SVector(us[icol+3], us[icol+4], us[icol+5]) # angular displacement
        F_elem = SVector(us[icol+6], us[icol+7], us[icol+8]) .* stru.force_scaling # internal force
        M_elem = SVector(us[icol+9], us[icol+10], us[icol+11]) .* stru.force_scaling # internal moment
        P_elem = SVector(us[icol+12], us[icol+13], us[icol+14]) .* stru.mass_scaling # linear momentum
        H_elem = SVector(us[icol+15], us[icol+16], us[icol+17]) .* stru.mass_scaling # angular momentum

        du_elem = SVector(dus[icol], dus[icol+1], dus[icol+2]) # linear displacement
        dθ_elem = SVector(dus[icol+3], dus[icol+4], dus[icol+5]) # angular displacement
        dF_elem = SVector(dus[icol+6], dus[icol+7], dus[icol+8]) .* stru.force_scaling # internal force
        dM_elem = SVector(dus[icol+9], dus[icol+10], dus[icol+11]) .* stru.force_scaling # internal moment
        dP_elem = SVector(dus[icol+12], dus[icol+13], dus[icol+14]) .* stru.mass_scaling # linear momentum
        dH_elem = SVector(dus[icol+15], dus[icol+16], dus[icol+17]) .* stru.mass_scaling # angular momentum

        # convert rotation parameter to Wiener-Milenkovic parameters
        scaling = GXBeam.rotation_parameter_scaling(θ_elem)
        θ_elem *= scaling # angular displacement (Wiener-Milenkovic parameters)
        dθ_elem *= scaling # angular displacement rate (Wiener-Milenkovic parameters)

        # local to body transformation
        CtCab = GXBeam.get_C(θ_elem)'*element.Cab # local to body transformation

        # local structural to aerodynamic transformation
        R = @SMatrix [0 -1 0; 1 0 0; 0 0 1] # local structural to aerodynamic transformation

        # local freestream linear and angular velocities
        vi = -R*(CtCab'*V + GXBeam.element_linear_velocity(element, P_elem, H_elem)) # local linear freestream velocity
        ωi = R*GXBeam.element_angular_velocity(element, P_elem, H_elem) # local angular freestream velocity

        # local freestream linear and angular accelerations
        dvi = -R*GXBeam.element_linear_velocity(element, dP_elem, dH_elem)
        dωi = R*GXBeam.element_angular_velocity(element, dP_elem, dH_elem)

        # NOTE: We assume that the local beam frame y-axis is oriented towards
        # the leading edge and the z-axis is oriented up
        # NOTE: We assume the local aerodynamic frame is oriented with the x-axis
        # in the chordwise direction and the z-axis in the (airfoil) normal
        # direction.

        # section state variable rates
        duai = SVector{Nuai}(duas[i]) # aerodynamic state variables
        dusi = vcat(dvi, dωi) # structural state variables
        dui = vcat(duai, dusi)

        # section state variables
        uai = SVector{Nuai}(uas[i]) # aerodynamic state variables
        usi = vcat(vi, ωi) # structural state variables
        ui = vcat(uai, usi)

        # section parameters
        pai = SVector{Npai}(pas[i]) # aerodynamic parameters
        psi = SVector(ρ) # structural parameters
        pi = vcat(pai, psi)

        # section inputs
        yi = get_coupling_inputs(section_models, dui, ui, pi, t)

        # separate inputs
        yai = view(yi, 1:Nyai)
        ysi = view(yi, Nyai + 1 : Nyai + Nysi)

        # save local inputs
        y[iya[iyas[i]]] .= yai # aerodynamic inputs

        # section aerodynamic loads (in body frame)
        fi = CtCab*R'*SVector(ysi[1], ysi[2], ysi[3])
        mi = CtCab*R'*SVector(ysi[4], ysi[5], ysi[6])

        # add constant distributed loads (in body frame)
        poff = 1 + 6*npoint + 6*(i-1)
        fi += SVector(padd[poff+1], padd[poff+2], padd[poff+3])
        mi += SVector(padd[poff+4], padd[poff+5], padd[poff+6])

        # save distributed loads for this element (in body frame)
        yoff = 6*npoint + 6*(i-1)
        y[iys[yoff+1 : yoff+3]] = fi
        y[iys[yoff+4 : yoff+6]] = mi
    end

    # save body frame linear/angular velocities
    yoff = 6*npoint + 6*nelem
    y[iys[yoff+1]] = ur
    y[iys[yoff+2]] = vr
    y[iys[yoff+3]] = wr
    y[iys[yoff+4]] = pr
    y[iys[yoff+5]] = qr
    y[iys[yoff+6]] = rr

    return y
end

# --- performance overloads --- #

# TODO

# --- convenience methods --- #

function set_additional_parameters!(padd, aero::LiftingLine, stru::GEBT;
    rho, point_conditions, element_loads, u, v, w, p, q, r) where {NA,TA}

    np = length(stru.icol_point)
    ne = length(stru.icol_elem)

    padd[1] = rho

    for ip = 1:np
        padd[1+6*(ip-1)+1] = point_conditions[1,ip]
        padd[1+6*(ip-1)+2] = point_conditions[2,ip]
        padd[1+6*(ip-1)+3] = point_conditions[3,ip]
        padd[1+6*(ip-1)+4] = point_conditions[4,ip]
        padd[1+6*(ip-1)+5] = point_conditions[5,ip]
        padd[1+6*(ip-1)+6] = point_conditions[6,ip]
    end

    for ie = 1:ne
        padd[1+6*np+6*(ie-1)+1] = element_loads[1,ie]
        padd[1+6*np+6*(ie-1)+2] = element_loads[2,ie]
        padd[1+6*np+6*(ie-1)+3] = element_loads[3,ie]
        padd[1+6*np+6*(ie-1)+4] = element_loads[4,ie]
        padd[1+6*np+6*(ie-1)+5] = element_loads[5,ie]
        padd[1+6*np+6*(ie-1)+6] = element_loads[6,ie]
    end

    padd[1+6*np+6*ne+1] = u
    padd[1+6*np+6*ne+2] = v
    padd[1+6*np+6*ne+3] = w
    padd[1+6*np+6*ne+4] = p
    padd[1+6*np+6*ne+5] = q
    padd[1+6*np+6*ne+6] = r

    return padd
end

function separate_additional_parameters(aero::LiftingLine, stru::GEBT, padd)

    np = length(stru.icol_point)
    ne = length(stru.icol_elem)

    rho = padd[1]

    point_conditions = reshape(view(padd, 1 + 1 : 1 + 6*np), 6, np)

    element_loads = reshape(view(padd, 1 + 6*np + 1 : 1 + 6*np + 6*ne), 6, ne)

    u = padd[1 + 6*np + 6*ne + 1]
    v = padd[1 + 6*np + 6*ne + 2]
    w = padd[1 + 6*np + 6*ne + 3]
    p = padd[1 + 6*np + 6*ne + 4]
    q = padd[1 + 6*np + 6*ne + 5]
    r = padd[1 + 6*np + 6*ne + 6]

    return (rho = rho, point_conditions = point_conditions, element_loads = element_loads,
        u = u, v = v, w = w, p = p, q = q, r = r)
end


const mutated = nothing
# TODO: define a singleton on Python side

_gen_doc(stem, jl) = """
    __$(stem)__(self, other)

Python's $(_pymodelref("__$(stem)__")).
Default implementation is roughly: `broadcast($jl, self, other)`.
"""

_gen_rdoc(stem) = """
    __r$(stem)__(self, other)

Python's $(_pymodelref("__r$(stem)__")).
Default implementation: `PyBase.__$(stem)__(other, self)`.
"""

_gen_idoc(stem, jl) = """
    __i$(stem)__(self, other)

Python's $(_pymodelref("__i$(stem)__")).
If `self` is mutated, `PyBase.mutated` must be returned.
Default implementation:

```julia
function __i$(stem)__(self, other)
    broadcast!($jl, self, self, other)
    return PyBase.mutated
end
```
"""

for (stem, jl) in [
        (:add, +),
        (:sub, -),
        (:mul, *),
        (:truediv, /),
        (:floordiv, ÷),
        (:mod, Base.mod),
        (:pow, ^),
        (:lshift, <<),
        (:rshift, >>),
        (:and, &),
        (:xor, Base.xor),
        (:or, |),
        ]
    py = Symbol(:__, stem, :__)
    @eval $py(a, b) = _wrap(broadcast($jl, a, b))
    @eval $(Symbol(:__r, stem, :__))(a, b) = $py(b, a)
    @eval function $(Symbol(:__i, stem, :__))(a, b)
        broadcast!($jl, a, a, b)
        return mutated
    end
    doc = _gen_doc(stem, jl)
    rdoc = _gen_rdoc(stem)
    idoc = _gen_idoc(stem, jl)
    @eval @doc $doc $py
    @eval @doc $rdoc $(Symbol(:__r, stem, :__))
    @eval @doc $idoc $(Symbol(:__i, stem, :__))
end

__pow__(a, b, modulo) = _wrap(broadcast(powermod, a, b, modulo))

__divmod__(a, b) = (_wrap(broadcast(÷, a, b)),
                    _wrap(broadcast(Base.mod, a, b)))
__rdivmod__(a, b) = __divmod__(b, a)

__matmul__(a, b) = _wrap(a * b)
__rmatmul__(a, b) = __matmul__(b, a)

function __imatmul__(y, a, b)
    if applicable(mul!, y, a, b)
        mul!(y, a, b)
        return mutated
    else
        return _wrap(a * b)
    end
end

for (stem, jl) in [
        (:neg, -),
        (:pos, +),
        (:abs, abs),
        (:invert, ~),
        (:complex, ComplexF64),
        (:int, Int128),
        (:float, Float64),
        # TODO: Use BigFloat/BigInt if mpmath is installed?
        (:trunc, trunc),
        (:floor, floor),
        (:ceil, ceil),
        ]
    py = Symbol(:__, stem, :__)
    @eval $py(a) = _wrap($jl(a))
    doc = """
        __$(stem)__(self)

    Python's $(_pymodelref("__$(stem)__")).
    Default implementation is roughly: `$jl(self)`.
    """
    @eval @doc $doc $py
end

"""
    __round__(self, [ndigits = nothing])

Python's $(_pymodelref("__round__")).
Default implementation is roughly:

```julia
__round__(self) = round(self)
__round__(self, ndigits) = round(self, digits=ndigits)
```
"""
__round__(self) = _wrap(round(self))
__round__(self, ::Nothing) = __round__(self)
__round__(self, ndigits) = _wrap(round(self, digits=ndigits))


####
#### Local minimization with NLopt. Loaded on demand when `NLopt` is.
####

import NLopt

import MultistartOptimization: local_minimization

export NLoptLocalMethod

const NLopt_ret_success = Set([:SUCCESS, :STOPVAL_REACHED, :FTOL_REACHED, :XTOL_REACHED,
                               :MAXEVAL_REACHED, :MAXTIME_REACHED])

Base.@kwdef struct NLoptLocalMethod{S}
    algorithm::NLopt.Algorithm
    xtol_abs::Float64 = 1e-8
    xtol_rel::Float64 = 1e-8
    maxeval::Int = 0
    maxtime::Float64 = 0.0
    "Return values which are considered as “success”."
    ret_success::S = NLopt_ret_success
end

"""
$(SIGNATURES)

A wrapper for algorithms supported by `NLopt`. Used to construct the corresponding
optimization problem. All return values in `ret_success` are considered valid (`ret` is also
kept in the result), all negative return values are considered invalid.

See the NLopt documentation for the options. Defaults are changed slightly.
"""
function NLoptLocalMethod(algorithm::NLopt.Algorithm; options...)
    NLoptLocalMethod(; algorithm = algorithm, options...)
end

"""
$(SIGNATURES)

Solve `minimization_problem` using `local_method`, starting from `x`. Return a
`LocationValue`.
"""
function local_minimization(local_method::NLoptLocalMethod,
                            minimization_problem::MinimizationProblem, x)
    @unpack algorithm, xtol_abs, xtol_rel, maxeval, maxtime, ret_success = local_method
    @unpack objective, lower_bounds, upper_bounds = minimization_problem
    opt = NLopt.Opt(algorithm, length(x))
    opt.lower_bounds = lower_bounds
    opt.upper_bounds = upper_bounds

    # if a method objective(x,grad) exists, use it. otherwise assume objective is not differentiable (which was the previous behavior)
    opt.min_objective = applicable(objective,x,x) ? objective : nlopt_nondifferentiable_wrapper(objective)
    opt.xtol_abs = xtol_abs
    opt.xtol_rel = xtol_rel
    opt.maxeval = maxeval
    opt.maxtime = maxtime
    optf, optx, ret = NLopt.optimize(opt, x)
    ret ∈ ret_success ? (value = optf, location = optx, ret = ret) : nothing
end

function nlopt_nondifferentiable_wrapper(fn)
    function f̃(x,grad)              # wrapper for NLopt
        @argcheck isempty(grad)     # ensure no derivatives are asked for
        return fn(x)
    end
    return f̃
end


# This file includes functions that build subnetwork models
# with extensions from PowerModels.jl

function decompose(
    data::Dict{String, Any},
    N_gs::Vector{Vector{Int64}},
    modeltype::Type{T},
    build_function::Function;
    extra_ref_extensions=[]
    )::NetDecModel where T <: PM.AbstractPowerModel

    return decompose(data, [Set(i) for i in N_gs], modeltype, build_function, extra_ref_extensions = extra_ref_extensions)
end

function decompose(
    data::Dict{String, Any},
    N_gs::Vector{Set{Int64}},
    modeltype::Type{T},
    build_function::Function;
    extra_ref_extensions=[]
    )::NetDecModel where T <: PM.AbstractPowerModel

    models = T[]
    shared_vars_dict = Dict{Int64, Dict}()
    denet_refs = [ref_add_cut_bus!, ref_add_cut_branch!]
    if T <: PM.AbstractWRMModel
        push!(denet_refs, ref_add_global_bus!)
    end
    for i in eachindex(N_gs)
        N_g = N_gs[i]
        sub_data = generate_subnet_data(data, N_g)
        push!(models, PM.instantiate_model(sub_data, modeltype, build_function, ref_extensions = push!(denet_refs, extra_ref_extensions...)))
        shared_vars_dict[i] = collect_split_vars(models[i])
    end

    return NetDecModel(N_gs, models, shared_vars_dict)
end

# modified from original build_opf
# get rid of ref buses and apply power balance to only nodes in partition
function build_opf_mod(pm::PM.AbstractPowerModel)
    # variable_bus_voltage(pm, bounded=false)
    # vr = var(pm, :vr)
    # vi = var(pm, :vi)
    # for (i, bus) in setdiff(ref(pm, :bus), ref(pm, :cut_bus))
    #     JuMP.set_lower_bound(vr[i], -bus["vmax"])
    #     JuMP.set_upper_bound(vr[i],  bus["vmax"])
    #     JuMP.set_lower_bound(vi[i], -bus["vmax"])
    #     JuMP.set_upper_bound(vi[i],  bus["vmax"])
    #     constraint_voltage_magnitude_bounds(pm, i)
    # end
    PM.variable_bus_voltage(pm)
    PM.variable_gen_power(pm)
    PM.variable_branch_power(pm)
    PM.variable_dcline_power(pm)

    objective_min_fuel_and_flow_cost_mod(pm)

    PM.constraint_model_voltage(pm)

    for i in setdiff(PM.ids(pm, :bus), PM.ids(pm, :cut_bus))
        PM.constraint_power_balance(pm, i)
    end

    for i in PM.ids(pm, :branch)
        PM.constraint_ohms_yt_from(pm, i)
        PM.constraint_ohms_yt_to(pm, i)

        PM.constraint_voltage_angle_difference(pm, i)

        PM.constraint_thermal_limit_from(pm, i)
        PM.constraint_thermal_limit_to(pm, i)
    end

    for i in PM.ids(pm, :dcline)
        PM.constraint_dcline_power_losses(pm, i)
    end
end

function build_opf_mod(pm::PM.AbstractWRMModel)
    function variable_bus_voltage(pm::PM.AbstractWRMModel; nw::Int=PM.nw_id_default, bounded::Bool=true, report::Bool=true)
        wr_min, wr_max, wi_min, wi_max = PM.ref_calc_voltage_product_bounds(PM.ref(pm, nw, :all_buspairs))
        bus_ids = PM.ids(pm, nw, :all_bus)

        w_index = 1:length(bus_ids)
        lookup_w_index = Dict((bi,i) for (i,bi) in enumerate(bus_ids))

        WR_start = zeros(length(bus_ids), length(bus_ids)) + I

        WR = PM.var(pm, nw)[:WR] = JuMP.@variable(pm.model,
            [i=1:length(bus_ids), j=1:length(bus_ids)], Symmetric, base_name="$(nw)_WR", start=WR_start[i,j]
        )
        if report
            PM.sol(pm, nw)[:WR] = WR
        end

        WI = PM.var(pm, nw)[:WI] = JuMP.@variable(pm.model,
            [1:length(bus_ids), 1:length(bus_ids)], base_name="$(nw)_WI", start=0.0
        )
        if report
            PM.sol(pm, nw)[:WI] = WI
        end

        # bounds on diagonal
        for (i, bus) in PM.ref(pm, nw, :all_bus)
            w_idx = lookup_w_index[i]
            wr_ii = WR[w_idx,w_idx]
            wi_ii = WR[w_idx,w_idx]

            if bounded
                JuMP.set_lower_bound(wr_ii, (bus["vmin"])^2)
                JuMP.set_upper_bound(wr_ii, (bus["vmax"])^2)
            else
                JuMP.set_lower_bound(wr_ii, 0)
            end
        end

        # bounds on off-diagonal
        for (i,j) in PM.ids(pm, nw, :all_buspairs)
            wi_idx = lookup_w_index[i]
            wj_idx = lookup_w_index[j]

            if bounded
                JuMP.set_upper_bound(WR[wi_idx, wj_idx], wr_max[(i,j)])
                JuMP.set_lower_bound(WR[wi_idx, wj_idx], wr_min[(i,j)])

                JuMP.set_upper_bound(WI[wi_idx, wj_idx], wi_max[(i,j)])
                JuMP.set_lower_bound(WI[wi_idx, wj_idx], wi_min[(i,j)])
            end
        end

        PM.var(pm, nw)[:w] = Dict{Int,Any}()
        for (i, bus) in PM.ref(pm, nw, :all_bus)
            w_idx = lookup_w_index[i]
            PM.var(pm, nw, :w)[i] = WR[w_idx,w_idx]
        end
        report && PM.sol_component_value(pm, nw, :all_bus, :w, PM.ids(pm, nw, :all_bus), PM.var(pm, nw)[:w])

        PM.var(pm, nw)[:wr] = Dict{Tuple{Int,Int},Any}()
        PM.var(pm, nw)[:wi] = Dict{Tuple{Int,Int},Any}()
        for (i,j) in PM.ids(pm, nw, :all_buspairs)
            w_fr_index = lookup_w_index[i]
            w_to_index = lookup_w_index[j]

            PM.var(pm, nw, :wr)[(i,j)] = WR[w_fr_index, w_to_index]
            PM.var(pm, nw, :wi)[(i,j)] = WI[w_fr_index, w_to_index]
        end
    end
    function constraint_voltage_angle_difference_all(pm::PM.AbstractPowerModel, i::Int; nw::Int=PM.nw_id_default)
        branch = PM.ref(pm, nw, :all_branch, i)
        f_bus = branch["f_bus"]
        t_bus = branch["t_bus"]
        f_idx = (i, f_bus, t_bus)
        pair = (f_bus, t_bus)
        buspair = PM.ref(pm, nw, :all_buspairs, pair)

        if buspair["branch"] == i
            PM.constraint_voltage_angle_difference(pm, nw, f_idx, buspair["angmin"], buspair["angmax"])
        end
    end
    variable_bus_voltage(pm) # re-define W matrix
    PM.variable_gen_power(pm)
    PM.variable_branch_power(pm)
    PM.variable_dcline_power(pm)

    objective_min_fuel_and_flow_cost_mod(pm) # allow for no generator case

    PM.constraint_model_voltage(pm)

    for i in setdiff(PM.ids(pm, :bus), PM.ids(pm, :cut_bus))
        PM.constraint_power_balance(pm, i)
    end

    for i in PM.ids(pm, :branch)
        PM.constraint_ohms_yt_from(pm, i)
        PM.constraint_ohms_yt_to(pm, i)

        # PM.constraint_voltage_angle_difference(pm, i)

        PM.constraint_thermal_limit_from(pm, i)
        PM.constraint_thermal_limit_to(pm, i)
    end

    for i in PM.ids(pm, :all_branch)
        constraint_voltage_angle_difference_all(pm, i)
    end

    for i in PM.ids(pm, :dcline)
        PM.constraint_dcline_power_losses(pm, i)
    end
end

#=
function PM.variable_bus_voltage_on_off(pm::PM.AbstractACRModel; kwargs...)
    variable_bus_voltage(pm; kwargs...)
end

function PM.constraint_model_voltage_on_off(pm::PM.AbstractACRModel; kwargs...)
end

function PM.constraint_ohms_yt_from_on_off(pm::AbstractACRModel, n::Int, f_bus, t_bus, f_idx, t_idx, g, b, g_fr, b_fr, tr, ti, tm)
    p_fr = var(pm, n, :p, f_idx)
    q_fr = var(pm, n, :q, f_idx)
    wrr_ff = var(pm, n, :wrr)[f_bus, f_bus]
    wii_ff = var(pm, n, :wii)[f_bus, f_bus]
    wrr_ft = var(pm, n, :wrr)[f_bus, t_bus]
    wii_ft = var(pm, n, :wii)[f_bus, t_bus]
    wri_ft = var(pm, n, :wri)[f_bus, t_bus]
    wri_tf = var(pm, n, :wri)[t_bus, f_bus]
    z = var(pm, n, :z_branch, i)

    JuMP.@constraint(pm.model, p_fr == z * ((g+g_fr)/tm^2*(wrr_ff + wii_ff) + (-g*tr+b*ti)/tm^2*(wrr_ft + wii_ft) + (-b*tr-g*ti)/tm^2*(wri_tf - wri_ft)) )
    JuMP.@constraint(pm.model, q_fr == z * (-(b+b_fr)/tm^2*(wrr_ff + wii_ff) - (-b*tr-g*ti)/tm^2*(wrr_ft + wii_ft) + (-g*tr+b*ti)/tm^2*(wri_tf - wri_ft)) )
end
=#

# modified from original build_ots
# get rid of ref buses and apply power balance to only nodes in partition
function build_ots_mod(pm::PM.AbstractPowerModel)
    PM.variable_branch_indicator(pm)
    PM.variable_bus_voltage_on_off(pm)
    PM.variable_gen_power(pm)
    PM.variable_branch_power(pm)
    PM.variable_dcline_power(pm)

    objective_min_fuel_and_flow_cost_mod(pm)

    PM.constraint_model_voltage_on_off(pm)

    for i in setdiff(PM.ids(pm, :bus), PM.ids(pm, :cut_bus))
        PM.constraint_power_balance(pm, i)
    end

    for i in PM.ids(pm, :branch)
        PM.constraint_ohms_yt_from_on_off(pm, i)
        PM.constraint_ohms_yt_to_on_off(pm, i)

        PM.constraint_voltage_angle_difference_on_off(pm, i)

        PM.constraint_thermal_limit_from_on_off(pm, i)
        PM.constraint_thermal_limit_to_on_off(pm, i)
    end

    for i in PM.ids(pm, :dcline)
        PM.constraint_dcline_power_losses(pm, i)
    end
end

# modified from original build_opf_bf
# get rid of ref buses and apply power balance to only nodes in partition
function build_opf_bf_mod(pm::PM.AbstractPowerModel)
    PM.variable_bus_voltage(pm)
    PM.variable_gen_power(pm)
    PM.variable_branch_power(pm)
    PM.variable_branch_current(pm)
    PM.variable_dcline_power(pm)

    objective_min_fuel_and_flow_cost_mod(pm)

    PM.constraint_model_current(pm)

    for i in setdiff(PM.ids(pm, :bus), PM.ids(pm, :cut_bus))
        PM.constraint_power_balance(pm, i)
    end

    for i in PM.ids(pm, :branch)
        PM.constraint_power_losses(pm, i)
        PM.constraint_voltage_magnitude_difference(pm, i)

        PM.constraint_voltage_angle_difference(pm, i)

        PM.constraint_thermal_limit_from(pm, i)
        PM.constraint_thermal_limit_to(pm, i)
    end

    for i in PM.ids(pm, :dcline)
        PM.constraint_dcline_power_losses(pm, i)
    end
end

# modified from original objective_min_fuel_and_flow_cost to allow for
# the case with no generators
function objective_min_fuel_and_flow_cost_mod(pm::PM.AbstractPowerModel; kwargs...)
    model = PM.check_cost_models(pm)
    if model == 1
        return PM.objective_min_fuel_and_flow_cost_pwl(pm; kwargs...)
    elseif model == 2
        return PM.objective_min_fuel_and_flow_cost_polynomial(pm; kwargs...)
    elseif model === nothing
        return JuMP.@objective(pm.model, Min, 0)
    else
        Memento.error(_LOGGER, "Only cost models of types 1 and 2 are supported at this time, given cost model type of $(model)")
    end
end

function variable_w_matrix(pm::PM.ACRPowerModel; nw::Int=PM.nw_id_default)
    wrr = PM.var(pm, nw)[:wrr] = JuMP.@variable(pm.model,
        [i in PM.ids(pm, nw, :bus), j in PM.ids(pm, nw, :bus)], base_name="$(nw)_wrr",
        start = PM.comp_start_value(PM.ref(pm, nw, :bus, i), "vr_start", 1.0) * PM.comp_start_value(PM.ref(pm, nw, :bus, j), "vr_start", 1.0)
        )
    wri = PM.var(pm, nw)[:wri] = JuMP.@variable(pm.model,
        [i in PM.ids(pm, nw, :bus), j in PM.ids(pm, nw, :bus)], base_name="$(nw)_wri",
        start = PM.comp_start_value(PM.ref(pm, nw, :bus, i), "vr_start", 1.0) * PM.comp_start_value(PM.ref(pm, nw, :bus, j), "vi_start")
        )
    wii = PM.var(pm, nw)[:wii] = JuMP.@variable(pm.model,
        [i in PM.ids(pm, nw, :bus), j in PM.ids(pm, nw, :bus)], base_name="$(nw)_wii",
        start = PM.comp_start_value(PM.ref(pm, nw, :bus, i), "vi_start") * PM.comp_start_value(PM.ref(pm, nw, :bus, j), "vi_start")
        )
end

#=
function build_acopf_with_free_lines(pm::PM.AbstractPowerModel)
    # this is for adding w variables to the model
    build_opf_mod(pm)
    PM.variable_bus_voltage_magnitude_sqr(pm)
    PM.variable_buspair_voltage_product(pm)

    w  = PM.var(pm,  :w)
    wr = PM.var(pm, :wr)
    wi = PM.var(pm, :wi)
    vr = PM.var(pm, :vr)
    vi = PM.var(pm, :vi)

    for (i, bus) in PM.ref(pm, :bus)
        JuMP.@constraint(pm.model, w[i] == vr[i]^2 + vi[i]^2)
    end

    for (_, branch) in PM.ref(pm, :branch)
        fbus = branch["f_bus"]
        tbus = branch["t_bus"]
        JuMP.@constraint(pm.model, wr[(fbus, tbus)] == vr[fbus] * vr[tbus] + vi[fbus] * vi[tbus])
        JuMP.@constraint(pm.model, wi[(fbus, tbus)] == vi[fbus] * vr[tbus] - vr[fbus] * vi[tbus])
    end
end
=#

function build_acopf_with_free_lines(pm::PM.AbstractPowerModel)
    build_opf_mod(pm)
    variable_w_matrix(pm)
    wrr = PM.var(pm, :wrr)
    wri = PM.var(pm, :wri)
    wii = PM.var(pm, :wii)
    vr = PM.var(pm, :vr)
    vi = PM.var(pm, :vi)
    for (i, _) in PM.ref(pm, :bus), (j, _) in PM.ref(pm, :bus)
        JuMP.@constraint(pm.model, wrr[i,j] == vr[i] * vr[j])
        JuMP.@constraint(pm.model, wri[i,j] == vr[i] * vi[j])
        JuMP.@constraint(pm.model, wii[i,j] == vi[i] * vi[j])
    end
end

function build_acopf_with_free_lines(pm::PM.SDPWRMPowerModel)
    build_opf_mod(pm)
end

function build_socwr_with_free_lines(pm::PM.AbstractWRModel)
    build_opf_mod(pm)
end

function build_socbf_with_free_lines(pm::PM.AbstractBFModel)
    build_opf_bf_mod(pm)
end

function build_acots_with_free_lines(pm::PM.AbstractACPModel)
    build_ots_mod(pm)
end

# function build_acots_with_free_lines(pm::PM.AbstractACRModel)
#     variable_w_matrix(pm)
#     build_ots_mod(pm)
# end

function collect_split_vars(pm::PM.ACPPowerModel)
    p = PM.var(pm, :p)
    q = PM.var(pm, :q)
    z = PM.var(pm, :z_branch)
    shared_vars_dict = Dict{String, Dict{Tuple, JuMP.VariableRef}}()
    shared_vars_dict["z"] = Dict{Tuple{Int64, Int64, Int64}, JuMP.VariableRef}()
    shared_vars_dict["p"] = Dict{Tuple{Int64, Int64, Int64}, JuMP.VariableRef}()
    shared_vars_dict["q"] = Dict{Tuple{Int64, Int64, Int64}, JuMP.VariableRef}()
    cut_arcs_from = PM.ref(pm, :cut_arcs_from)
    for (l,i,j) in cut_arcs_from
        shared_vars_dict["z"][(l,i,j)] = z[l]
        shared_vars_dict["p"][(l,i,j)] = p[(l,i,j)]
        shared_vars_dict["p"][(l,j,i)] = p[(l,j,i)]
        shared_vars_dict["q"][(l,i,j)] = q[(l,i,j)]
        shared_vars_dict["q"][(l,j,i)] = q[(l,j,i)]
    end
    return shared_vars_dict
end

function collect_split_vars(pm::PM.ACRPowerModel)
    wrr = PM.var(pm, :wrr)
    wri = PM.var(pm, :wri)
    wii = PM.var(pm, :wii)
    p  = PM.var(pm,  :p)
    q  = PM.var(pm,  :q)

    shared_vars_dict = Dict{String, Dict{Tuple, JuMP.VariableRef}}()
    shared_vars_dict["wrr"] = Dict{Tuple{Int64, Int64}, JuMP.VariableRef}()
    shared_vars_dict["wri"] = Dict{Tuple{Int64, Int64}, JuMP.VariableRef}()
    shared_vars_dict["wir"] = Dict{Tuple{Int64, Int64}, JuMP.VariableRef}()
    shared_vars_dict["wii"] = Dict{Tuple{Int64, Int64}, JuMP.VariableRef}()
    shared_vars_dict["p"] = Dict{Tuple{Int64, Int64, Int64}, JuMP.VariableRef}()
    shared_vars_dict["q"] = Dict{Tuple{Int64, Int64, Int64}, JuMP.VariableRef}()

    cut_arcs_from = PM.ref(pm, :cut_arcs_from)
    for (l,i,j) in cut_arcs_from
        if !((i,j) in keys(shared_vars_dict["wrr"]))
            shared_vars_dict["wrr"][(i,j)] = wrr[i,j]
            shared_vars_dict["wri"][(i,j)] = wri[i,j]
            shared_vars_dict["wir"][(i,j)] = wri[j,i]
            shared_vars_dict["wii"][(i,j)] = wii[i,j]
        end
        shared_vars_dict["p"][(l,i,j)] = p[(l,i,j)]
        shared_vars_dict["p"][(l,j,i)] = p[(l,j,i)]
        shared_vars_dict["q"][(l,i,j)] = q[(l,i,j)]
        shared_vars_dict["q"][(l,j,i)] = q[(l,j,i)]
    end
    return shared_vars_dict
end

function collect_split_vars(pm::PM.AbstractWRModel)
    wr = PM.var(pm, :wr)
    wi = PM.var(pm, :wi)
    p  = PM.var(pm,  :p)
    q  = PM.var(pm,  :q)

    shared_vars_dict = Dict{String, Dict{Tuple, JuMP.VariableRef}}()
    # shared_vars_dict["wr"] = Dict{Tuple{Int64, Int64}, JuMP.VariableRef}()
    # shared_vars_dict["wi"] = Dict{Tuple{Int64, Int64}, JuMP.VariableRef}()
    shared_vars_dict["p"] = Dict{Tuple{Int64, Int64, Int64}, JuMP.VariableRef}()
    shared_vars_dict["q"] = Dict{Tuple{Int64, Int64, Int64}, JuMP.VariableRef}()

    cut_arcs_from = PM.ref(pm, :cut_arcs_from)
    for (l,i,j) in cut_arcs_from
        # if !((i,j) in keys(shared_vars_dict["wr"]))
        #     shared_vars_dict["wr"][(i,j)] = wr[(i,j)]
        #     shared_vars_dict["wi"][(i,j)] = wi[(i,j)]
        # end
        shared_vars_dict["p"][(l,i,j)] = p[(l,i,j)]
        shared_vars_dict["p"][(l,j,i)] = p[(l,j,i)]
        shared_vars_dict["q"][(l,i,j)] = q[(l,i,j)]
        shared_vars_dict["q"][(l,j,i)] = q[(l,j,i)]
    end
    return shared_vars_dict
end

function collect_split_vars(pm::PM.AbstractBFModel)
    p  = PM.var(pm,  :p)
    q  = PM.var(pm,  :q)
    ccm = PM.var(pm, :ccm)

    shared_vars_dict = Dict{String, Dict}()
    shared_vars_dict["ccm"] = Dict{Int64, JuMP.VariableRef}()
    shared_vars_dict["p"] = Dict{Tuple{Int64, Int64, Int64}, JuMP.VariableRef}()
    shared_vars_dict["q"] = Dict{Tuple{Int64, Int64, Int64}, JuMP.VariableRef}()

    cut_arcs_from = PM.ref(pm, :cut_arcs_from)
    for (l,i,j) in cut_arcs_from
        shared_vars_dict["ccm"][l] = ccm[l]
        shared_vars_dict["p"][(l,i,j)] = p[(l,i,j)]
        shared_vars_dict["p"][(l,j,i)] = p[(l,j,i)]
        shared_vars_dict["q"][(l,i,j)] = q[(l,i,j)]
        shared_vars_dict["q"][(l,j,i)] = q[(l,j,i)]
    end
    return shared_vars_dict
end

function collect_split_vars(pm::PM.AbstractWRMModel)
    p = PM.var(pm, :p)
    q = PM.var(pm, :q)
    WR = PM.var(pm, :WR)
    WI = PM.var(pm, :WI)
    bus_ids = PM.ids(pm, :all_bus) # all buses from the full network
    lookup_w_index = Dict((bi,i) for (i,bi) in enumerate(bus_ids))

    shared_vars_dict = Dict{String, Dict{Tuple, JuMP.VariableRef}}()
    shared_vars_dict["p"] = Dict{Tuple{Int64, Int64, Int64}, JuMP.VariableRef}()
    shared_vars_dict["q"] = Dict{Tuple{Int64, Int64, Int64}, JuMP.VariableRef}()
    shared_vars_dict["WR"] = Dict{Tuple{Int64, Int64}, JuMP.VariableRef}()
    shared_vars_dict["WI"] = Dict{Tuple{Int64, Int64}, JuMP.VariableRef}()
    cut_arcs_from = PM.ref(pm, :cut_arcs_from)
    for (l,i,j) in cut_arcs_from # penalizing cut lines
        shared_vars_dict["p"][(l,i,j)] = p[(l,i,j)]
        shared_vars_dict["p"][(l,j,i)] = p[(l,j,i)]
        shared_vars_dict["q"][(l,i,j)] = q[(l,i,j)]
        shared_vars_dict["q"][(l,j,i)] = q[(l,j,i)]
    end
    bus = PM.ids(pm, :bus)
    for i in bus_ids, j in bus_ids
        shared_vars_dict["WR"][(i,j)] = WR[lookup_w_index[i], lookup_w_index[j]]
        shared_vars_dict["WI"][(i,j)] = WI[lookup_w_index[i], lookup_w_index[j]]
        if i != j
            shared_vars_dict["WR"][(j,i)] = WR[lookup_w_index[j], lookup_w_index[i]]
            shared_vars_dict["WI"][(j,i)] = WI[lookup_w_index[j], lookup_w_index[i]]
        end
    end
    return shared_vars_dict
end

######################################################
###             EXPERIMENTAL FUNCTIONS             ###
######################################################

function constraint_model_voltage_mod(pm::PM.AbstractPowerModel)
    PM.constraint_model_voltage(pm)
end

# function constraint_model_voltage_mod(pm::PM.AbstractWRMModel)
#     PM._check_missing_keys(PM.var(pm), [:WR,:WI], typeof(pm))

#     WR = PM.var(pm)[:WR]
#     WI = PM.var(pm)[:WI]

#     bus_ids = PM.ids(pm, :bus)
#     lookup_w_index = Dict((bi,i) for (i,bi) in enumerate(bus_ids))

#     local_buses = setdiff(PM.ids(pm, :bus), PM.ids(pm, :cut_bus))
#     local_w_idx = [lookup_w_index[i] for i in local_buses]

#     WR_local = WR[local_w_idx, local_w_idx]
#     WI_local = WI[local_w_idx, local_w_idx]
#     JuMP.@constraint(pm.model, [WR_local WI_local; -WI_local WR_local] in JuMP.PSDCone())
# end