import DataFrames.skipmissing

sim_param = setup_sim_param_model()

include("misc_functions.jl")

##### To load the DR25 KOI catalog and apply the necessary cuts:

Q1Q17_DR25 = CSV.read("q1_q17_dr25_koi.tab_selectcols_new.csv", header=19, allowmissing=:all)

#If using our own cuts applied to the stellar and koi catalogs:

#=
Q1Q17_DR25_stellar = CSV.read("q1_q17_dr25_stellar_koi.tab_selectcols.csv", header=30, allowmissing=:all)
Q1Q17_DR25_stellar_all = CSV.read("q1_q17_dr25_stellar_koi.tab_all.csv", header=1, allowmissing=:all)

N_Kepler_targets = sum((Q1Q17_DR25_stellar_all[:teff] .> 4000.) .& (Q1Q17_DR25_stellar_all[:teff] .< 7000.) .& (Q1Q17_DR25_stellar_all[:logg] .> 4.))
println("Total number of Kepler targets satisfying our cuts: ", N_Kepler_targets)

table_confirmed = Q1Q17_DR25[(Q1Q17_DR25[:koi_disposition] .== "CONFIRMED") .| (Q1Q17_DR25[:koi_disposition] .== "CANDIDATE"), :] #Table containing only the confirmed and candidate objects
table_stellar = Q1Q17_DR25_stellar

table_confirmed = table_confirmed[(table_confirmed[:koi_period] .> get_real(sim_param,"min_period")) .& (table_confirmed[:koi_period] .< get_real(sim_param,"max_period")), :] #to make additional cuts in period P to be comparable to our simulated sample
table_confirmed = table_confirmed[(table_confirmed[:koi_prad] .> get_real(sim_param,"min_radius")/ExoplanetsSysSim.earth_radius) .& (table_confirmed[:koi_prad] .< get_real(sim_param,"max_radius")/ExoplanetsSysSim.earth_radius) .& (.~ismissing.(table_confirmed[:koi_prad])), :] #to make additional cuts in planetary radii to be comparable to our simulated sample

#To make cuts based on stellar properties of T_eff and logg:
teff_confirmed = zeros(Int64, size(table_confirmed,1)) #list to be filled with the T_eff (K) for the objects
logg_confirmed = zeros(Float64, size(table_confirmed,1)) #list to be filled with the logg(cgs) for the objects
cdpp5_confirmed = zeros(Float64, size(table_confirmed,1)) #list to be filled with the RMS CDPP 5h values for the objects
cdpp6_confirmed = zeros(Float64, size(table_confirmed,1)) #list to be filled with the RMS CDPP 6h values for the objects
for (i,KepID) in enumerate(table_confirmed[:kepid])
    teff_confirmed[i] = table_stellar[:teff][table_stellar[:kepid] .== KepID, :][1]
    logg_confirmed[i] = table_stellar[:logg][table_stellar[:kepid] .== KepID, :][1]
    cdpp5_confirmed[i] = table_stellar[:rrmscdpp05p0][table_stellar[:kepid] .== KepID, :][1]
    cdpp6_confirmed[i] = table_stellar[:rrmscdpp06p0][table_stellar[:kepid] .== KepID, :][1]
end

cdpp_cut = 250.
println("Fraction of CONFIRMED and CANDIDATE planets after CDPP cut: ", size(table_confirmed[(teff_confirmed .> 4000.) .& (teff_confirmed .< 7000.) .& (logg_confirmed .> 4.) .& (cdpp5_confirmed .< cdpp_cut), :], 1), "/", size(table_confirmed[(teff_confirmed .> 4000.) .& (teff_confirmed .< 7000.) .& (logg_confirmed .> 4.), :], 1))

table_confirmed = table_confirmed[(teff_confirmed .> 4000.) .& (teff_confirmed .< 7000.) .& (logg_confirmed .> 4.) .& (cdpp5_confirmed .< cdpp_cut), :]
=#

#If using a stellar table with cuts already made and just matching kepids to get a koi catalog:

stellar_catalog = setup_star_table_christiansen(sim_param)
N_Kepler_targets = get_int(sim_param,"num_kepler_targets")

table_confirmed = Q1Q17_DR25[(Q1Q17_DR25[:koi_disposition] .== "CONFIRMED") .| (Q1Q17_DR25[:koi_disposition] .== "CANDIDATE"), :] #Table containing only the confirmed and candidate objects

in_stellar_catalog = [] #will be filled with booleans indicating whether each koi in 'table_confirmed' is found in the 'stellar_catalog' or not
for i in 1:length(table_confirmed[:kepid])
    if any(x->x==table_confirmed[i,:kepid], stellar_catalog[:kepid])
        push!(in_stellar_catalog, true)
    else
        push!(in_stellar_catalog, false)
    end
end
in_stellar_catalog_indices = findall(in_stellar_catalog)

table_confirmed = table_confirmed[in_stellar_catalog_indices, :]

for (i,kepid) in enumerate(table_confirmed[:kepid]) #to replace the stellar and planetary radii in 'table_confirmed' with the more reliable values as derived from the stellar properties in 'stellar_catalog'
    stellar_radii_new = stellar_catalog[stellar_catalog[:kepid] .== kepid, :radius][1]
    table_confirmed[i, :koi_srad] = stellar_radii_new
    table_confirmed[i, :koi_prad] = (1 ./ExoplanetsSysSim.earth_radius)*stellar_radii_new*sqrt(table_confirmed[i, :koi_depth]/(1e6))
end

table_confirmed = table_confirmed[(table_confirmed[:koi_period] .> get_real(sim_param,"min_period")) .& (table_confirmed[:koi_period] .< get_real(sim_param,"max_period")), :] #to make additional cuts in period P to be comparable to our simulated sample
table_confirmed = table_confirmed[(table_confirmed[:koi_prad] .> get_real(sim_param,"min_radius")/ExoplanetsSysSim.earth_radius) .& (table_confirmed[:koi_prad] .< get_real(sim_param,"max_radius")/ExoplanetsSysSim.earth_radius) .& (.~ismissing.(table_confirmed[:koi_prad])), :] #to make additional cuts in planetary radii to be comparable to our simulated sample

#=
#If we want to write the stellar radii to a file:
CSV.write("stellar_radii_q1_q17_dr25_gaia_fgk.txt", stellar_catalog[[:kepid, :radius]])
=#

#=
#If we want to write the kep_oi's of the remaining koi catalog to a file:
f = open("kepoi_names.txt", "w")
println(f, table_confirmed[:kepoi_name])
close(f)
=#

#####To compute arrays of the observables (multiplicities, periods, period ratios, transit durations, transit depths, period-normalized transit duration ratios (xi), and transit depth ratios) from the remaining sample of planets:

KOI_systems = [x[1:6] for x in table_confirmed[:kepoi_name]]
checked_bools = zeros(size(table_confirmed,1)) #0's denote KOI that were not checked yet; 1's denote already checked KOI

M_confirmed = Int64[] #list to be filled with the planet multiplicities of the systems
R_confirmed = Float64[] #list to be filled with period ratios of adjacent planet pairs
xi_confirmed = Float64[] #list to be filled with the period-normalized transit duration ratios of adjacent planet pairs
D_ratio_confirmed = Float64[] #list to be filled with the transit depth ratios of adjacent planet pairs
P_confirmed = table_confirmed[:koi_period] #array of the periods (days)
P_confirmed = collect(skipmissing(P_confirmed))
t_D_confirmed = table_confirmed[:koi_duration] #array of the transit durations (hrs)
t_D_confirmed = collect(skipmissing(t_D_confirmed))
D_confirmed = table_confirmed[:koi_depth]/(1e6) #array of the transit depths (fraction)
D_confirmed = D_confirmed[ismissing.(D_confirmed) .== false] #to get rid of NA values

D_above_confirmed = Float64[] #list to be filled with the transit depths of planets above the photoevaporation boundary in Carrera et al 2018
D_below_confirmed = Float64[] #list to be filled with the transit depths of planets below the boundary
D_ratio_above_confirmed = Float64[] #list to be filled with the transit depth ratios of adjacent planet pairs, both above the boundary
D_ratio_below_confirmed = Float64[] #list to be filled with the transit depth ratios of adjacent planet pairs, both below the boundary
D_ratio_across_confirmed = Float64[] #list to be filled with the transit depth ratios of adjacent planet pairs, across the boundary

for i in 1:length(KOI_systems)
    if checked_bools[i] == 0 #if the KOI has not been checked (included while looking at another planet in the same system)
        system_i = (1:length(KOI_systems))[KOI_systems .== KOI_systems[i]]
        checked_bools[system_i] .= 1

        #To get the periods and transit durations in this system:
        system_P = table_confirmed[:koi_period][system_i] #periods of all the planets in this system
        system_t_D = table_confirmed[:koi_duration][system_i] #transit durations of all the planets in this system
        system_D = table_confirmed[:koi_depth][system_i] #transit depths (in ppm) of all the planets in this system
        system_radii = table_confirmed[:koi_prad][system_i] #radii of all the planets in this system
        system_sort_i = sortperm(system_P) #indices that would sort the periods of the planets in this system
        system_P = system_P[system_sort_i] #periods of all the planets in this system, sorted
        system_t_D = system_t_D[system_sort_i] #transit durations of all the planets in this system, sorted by period
        system_D = system_D[system_sort_i] #transit depths of all the planets in this system, sorted by period
        system_radii = system_radii[system_sort_i] #radii of all the planets in this system, sorted by period

        #To count the total number of planets in this system:
        push!(M_confirmed, length(system_P))

        #To compute the period ratios, period-normalized transit duration ratios, and transit depth ratios in this system:
        system_R = system_P[2:end]./system_P[1:end-1] #period ratios of all the adjacent planet pairs in this system
        system_D_ratio = system_D[2:end]./system_D[1:end-1] #transit depth ratios of all the adjacent planet pairs in this system
        system_xi = (system_t_D[1:end-1]./system_t_D[2:end]).*(system_P[2:end]./system_P[1:end-1]).^(1//3) #period-normalized transit duration ratios of all the adjacent planet pairs in this system

        append!(R_confirmed, system_R)
        append!(D_ratio_confirmed, system_D_ratio)
        append!(xi_confirmed, system_xi)

        #To separate the planets in the system as above and below the boundary:
        system_above_bools = [photoevap_boundary_Carrera2018(system_radii[x], system_P[x]) for x in 1:length(system_P)]
        #if length(system_above_bools) > 1 println(system_above_bools) end

        #To record the transit depths of the planets above and below the boundary:
        for (j,D) in enumerate(system_D)
            if system_above_bools[j] == 1
                append!(D_above_confirmed, D)
            elseif system_above_bools[j] == 0
                append!(D_below_confirmed, D)
            end
        end

        #To record the transit depth ratios of the planets above, below, and across the boundary:
        for (j,D_ratio) in enumerate(system_D_ratio)
            if system_above_bools[j] + system_above_bools[j+1] == 2 #both planets are above the boundary
                append!(D_ratio_above_confirmed, D_ratio)
            elseif system_above_bools[j] + system_above_bools[j+1] == 1 #one planet is above, the other is below the boundary
                append!(D_ratio_across_confirmed, D_ratio)
            elseif system_above_bools[j] + system_above_bools[j+1] == 0 #both planets are below the boundary
                append!(D_ratio_below_confirmed, D_ratio)
            end
        end
    end
end


@generated function saturated_add(x::I, y::I) where {I <: IntegerTypesHW}
    typ = "i$(8sizeof(I))"
    s = I <: Signed ? 's' : 'u'
    f = "@llvm.$(s)add.sat.$typ"
    decl = "declare $typ $f($typ, $typ)"
    instrs = """
        %res = call $typ $f($typ %0, $typ %1)
        ret $typ %res
    """
    llvmcall_expr(decl, instrs, I, :(Tuple{$I,$I}), typ, [typ,typ], [:x,:y])
end
@generated function saturated_add(x::Vec{W,I}, y::Vec{W,I}) where {W,I}

end

@eval @inline function assume(b::Bool)
    $(llvmcall_expr("declare void @llvm.assume(i1)", "%b = trunc i8 %0 to i1\ncall void @llvm.assume(i1 %b)\nret void", :Cvoid, :(Tuple{Bool}), "void", ["i8"], [:b]))
end

@eval @inline function expect(b::Bool)
    $(llvmcall_expr("declare i1 @llvm.expect.i1(i1, i1)", """
    %b = trunc i8 %0 to i1
    %actual = call i1 @llvm.expect.i1(i1 %b, i1 true)
    %byte = zext i1 %actual to i8
    ret i8 %byte""", :Bool, :(Tuple{Bool}), "i8", ["i8"], [:b]))
end
@generated function expect(i::I, ::Val{N}) where {I <: Integer, N}
    ityp = 'i' * string(8sizeof(I))
    llvmcall_expr("declare i1 @llvm.expect.$ityp($ityp, i1)", """
    %actual = call $ityp @llvm.expect.$ityp($ityp %0, $ityp $N)
    ret $ityp %actual""", I, :(Tuple{$I}), ityp, [ityp], [:i])
end

# for (op,f) ∈ [("abs",:abs)]
# end
if Base.libllvm_version ≥ v"12"
    for (op,f,S) ∈ [("smax",:max,:Signed),("smin",:min,:Signed),("umax",:max,:Unsigned),("umin",:min,:Unsigned)]
        @eval @generated Base.$f(v1::Vec{W,T}, v2::Vec{W,T}) where {W, T <: $S} = llvmcall_expr($op, W, T, (W, W), (T, T))
    end
else
    @inline Base.max(v1::Vec{W,<:Integer}, v2::Vec{W,<:Integer}) where {W} = ifelse(v1 > v2, v1, v2)
    @inline Base.min(v1::Vec{W,<:Integer}, v2::Vec{W,<:Integer}) where {W} = ifelse(v1 < v2, v1, v2)
end
# for T ∈ [Float32, Float64]
#     W = 2
#     while W * sizeof(T) ≤ REGISTER_SIZE
# floating point
for (op,f) ∈ [("sqrt",:sqrt),("fabs",:abs),("floor",:floor),("ceil",:ceil),("trunc",:trunc),("nearbyint",:round)
              ]
    # @eval @generated Base.$f(v1::Vec{W,T}) where {W, T <: Union{Float32,Float64}} = llvmcall_expr($op, W, T, (W,), (T,), "nsz arcp contract afn reassoc")
    @eval @generated Base.$f(v1::Vec{W,T}) where {W, T <: Union{Float32,Float64}} = llvmcall_expr($op, W, T, (W,), (T,), "fast")
end

@generated function Base.round(::Type{Int64}, v1::Vec{W,T}) where {W, T <: Union{Float32,Float64}}
    llvmcall_expr("lrint", W, Int64, (W,), (T,), "nsz arcp contract afn reassoc")
end
@generated function Base.round(::Type{Int32}, v1::Vec{W,T}) where {W, T <: Union{Float32,Float64}}
    llvmcall_expr("lrint", W, Int32, (W,), (T,), "nsz arcp contract afn reassoc")
end
@inline Base.trunc(::Type{I}, v::AbstractSIMD{W,T}) where {W, I<:IntegerTypesHW, T <: NativeTypes} = convert(I, v)

# """
#    setbits(x::Unsigned, y::Unsigned, mask::Unsigned)

# If you have AVX512, setbits of vector-arguments will select bits according to mask `m`, selecting from `y` if 0 and from `x` if `1`.
# For scalar arguments, or vector arguments without AVX512, `setbits` requires the additional restrictions on `y` that all bits for
# which `m` is 1, `y` must be 0.
# That is for scalar arguments or vector arguments without AVX512, it requires the restriction that
# ((y ⊻ m) & m) == m
# """
# @inline setbits(x, y, m) = (x & m) | y

"""
   bitselect(m::Unsigned, x::Unsigned, y::Unsigned)

If you have AVX512, setbits of vector-arguments will select bits according to mask `m`, selecting from `x` if 0 and from `y` if `1`.
For scalar arguments, or vector arguments without AVX512, `setbits` requires the additional restrictions on `y` that all bits for
which `m` is 1, `y` must be 0.
That is for scalar arguments or vector arguments without AVX512, it requires the restriction that
((y ⊻ m) & m) == m
"""
@inline bitselect(m, x, y) = ((~m) & x) | (m & y)
if AVX512F
    # AVX512 lets us use 1 instruction instead of 2 dependent instructions to set bits
    @generated function vpternlog(m::Vec{W,UInt64}, x::Vec{W,UInt64}, y::Vec{W,UInt64}, ::Val{L}) where {W, L}
        @assert W ∈ (2,4,8)
        bits = 64W
        decl64 = "declare <$W x i64> @llvm.x86.avx512.mask.pternlog.q.$(bits)(<$W x i64>, <$W x i64>, <$W x i64>, i32, i8)"
        instr64 = """
            %res = call <$W x i64> @llvm.x86.avx512.mask.pternlog.q.$(bits)(<$W x i64> %0, <$W x i64> %1, <$W x i64> %2, i32 $L, i8 -1)
            ret <$W x i64> %res
        """
        arg_syms = [:(data(m)), :(data(x)), :(data(y))]
        llvmcall_expr(decl64, instr64, :(_Vec{$W,UInt64}), :(Tuple{_Vec{$W,UInt64},_Vec{$W,UInt64},_Vec{$W,UInt64}}), "<$W x i64>", ["<$W x i64>", "<$W x i64>", "<$W x i64>"], arg_syms)
    end
    @generated function vpternlog(m::Vec{W,UInt32}, x::Vec{W,UInt32}, y::Vec{W,UInt32}, ::Val{L}) where {W, L}
        if W ∉ (4,8,16)
            return Expr(:block, Expr(:meta, :inline), :(((~m) & x) | (m & y)))
        end
        bits = 32W
        decl32 = "declare <$W x i32> @llvm.x86.avx512.mask.pternlog.d.$(bits)(<$W x i32>, <$W x i32>, <$W x i32>, i32, i16)"
        instr32 = """
            %res = call <$W x i32> @llvm.x86.avx512.mask.pternlog.d.$(bits)(<$W x i32> %0, <$W x i32> %1, <$W x i32> %2, i32 $L, i16 -1)
            ret <$W x i32> %res
        """
        arg_syms = [:(data(m)), :(data(x)), :(data(y))]
        llvmcall_expr(decl32, instr32, :(_Vec{$W,UInt32}), :(Tuple{_Vec{$W,UInt32},_Vec{$W,UInt32},_Vec{$W,UInt32}}), "<$W x i32>", ["<$W x i32>", "<$W x i32>", "<$W x i32>"], arg_syms)
    end
    # @eval @generated function setbits(x::Vec{W,T}, y::Vec{W,T}, m::Vec{W,T}) where {W,T <: Union{UInt32,UInt64}}
    #     ex = if W*sizeof(T) ∈ (16,32,64)
    #         :(vpternlog(x, y, m, Val{216}()))
    #     else
    #         :((x & m) | y)
    #     end
    #     Expr(:block, Expr(:meta, :inline), ex)
    # end
    @generated function bitselect(m::Vec{W,T}, x::Vec{W,T}, y::Vec{W,T}) where {W,T <: Union{UInt32,UInt64}}
        ex = if W*sizeof(T) ∈ (16,32,64)
            :(vpternlog(m, x, y, Val{172}()))
        else
            :(((~m) & x) | (m & y))
        end
        Expr(:block, Expr(:meta, :inline), ex)
    end
    @inline Base.copysign(v1::Vec{W,Float64}, v2::Vec{W,Float64}) where {W} = reinterpret(Float64, bitselect(Vec{W,UInt64}(0x8000000000000000), reinterpret(UInt64, v1), reinterpret(UInt64, v2)))
    @inline Base.copysign(v1::Vec{W,Float32}, v2::Vec{W,Float32}) where {W} = reinterpret(Float32, bitselect(Vec{W,UInt32}(0x80000000), reinterpret(UInt32, v1), reinterpret(UInt32, v2)))
end

for (op,f) ∈ [("minnum",:min),("maxnum",:max),("copysign",:copysign),
              ]
    @eval @generated function Base.$f(v1::Vec{W,T}, v2::Vec{W,T}) where {W, T <: Union{Float32,Float64}}
        llvmcall_expr($op, W, T, (W for _ in 1:2), (T for _ in 1:2), "nsz arcp contract afn reassoc")
    end
end
@inline _signbit(v::Vec{W, I}) where {W, I<:Signed} = v & Vec{W,I}(typemin(I))
@inline Base.copysign(v1::Vec{W,I}, v2::Vec{W,I}) where {W, I <: Signed} = ifelse(_signbit(v1) == _signbit(v2), v1, -v1)

@inline Base.copysign(x::Float32, v::Vec{W}) where {W} = copysign(vbroadcast(Val{W}(), x), v)
@inline Base.copysign(x::Float64, v::Vec{W}) where {W} = copysign(vbroadcast(Val{W}(), x), v)
@inline Base.copysign(x::Float32, v::VecUnroll{N,W}) where {N,W} = copysign(vbroadcast(Val{W}(), x), v)
@inline Base.copysign(x::Float64, v::VecUnroll{N,W}) where {N,W} = copysign(vbroadcast(Val{W}(), x), v)
@inline Base.copysign(v::Vec, u::VecUnroll) = VecUnroll(fmap(copysign, v, u.data))
@inline Base.copysign(v::Vec{W,T}, x::NativeTypes) where {W,T} = copysign(v, Vec{W,T}(x))
@inline Base.copysign(v1::Vec{W,T}, v2::Vec{W}) where {W,T} = copysign(v1, convert(Vec{W,T}, v2))

# ternary
for (op,f) ∈ [("fma",:fma),("fmuladd",:muladd)]
    @eval @generated function Base.$f(v1::Vec{W,T}, v2::Vec{W,T}, v3::Vec{W,T}) where {W, T <: Union{Float32,Float64}}
        llvmcall_expr($op, W, T, (W for _ in 1:3), (T for _ in 1:3), $(f === :fma ? nothing : "nsz arcp contract afn reassoc"))
    end
end
# floating vector, integer scalar
# @generated function Base.:(^)(v1::Vec{W,T}, v2::Int32) where {W, T <: Union{Float32,Float64}}
#     llvmcall_expr("powi", W, T, (W, 1), (T, Int32), "nsz arcp contract afn reassoc")
# end
for (op,f) ∈ [
    ("experimental.vector.reduce.v2.fadd",:vsum),
    ("experimental.vector.reduce.v2.fmul",:vprod)
]
    @eval @generated function $f(v1::T, v2::Vec{W,T}) where {W, T <: Union{Float32,Float64}}
        llvmcall_expr($op, -1, T, (1, W), (T, T), "nsz arcp contract afn reassoc")
    end
end
vsum(s::T, v::Vec{W,T}) where {W,T} = Base.FastMath.add_fast(s, vsum(v))
vprod(s::T, v::Vec{W,T}) where {W,T} = Base.FastMath.mul_fast(s, vprod(v))
for (op,f) ∈ [
    ("experimental.vector.reduce.fmax",:vmaximum),
    ("experimental.vector.reduce.fmin",:vminimum)
]
    @eval @generated function $f(v1::Vec{W,T}) where {W, T <: Union{Float32,Float64}}
        llvmcall_expr($op, -1, T, (W,), (T,), "nsz arcp contract afn reassoc")
    end
end
for (op,f,S) ∈ [
    ("experimental.vector.reduce.add",:vsum,:Integer),
    ("experimental.vector.reduce.mul",:vprod,:Integer),
    ("experimental.vector.reduce.and",:vall,:Integer),
    ("experimental.vector.reduce.or",:vany,:Integer),
    ("experimental.vector.reduce.xor",:vxorreduce,:Integer),
    ("experimental.vector.reduce.smax",:vmaximum,:Signed),
    ("experimental.vector.reduce.smin",:vminimum,:Signed),
    ("experimental.vector.reduce.umax",:vmaximum,:Unsigned),
    ("experimental.vector.reduce.umin",:vminimum,:Unsigned)
]
    @eval @generated function $f(v1::Vec{W,T}) where {W, T <: $S}
        llvmcall_expr($op, -1, T, (W,), (T,))
    end
end

#         W += W
#     end
# end
@inline vsum(v::Vec{W,T}) where {W,T<:Union{Float32,Float64}} = vsum(-zero(T), v)
@inline vprod(v::Vec{W,T}) where {W,T<:Union{Float32,Float64}} = vprod(one(T), v)
@inline vsum(x, v::Vec{W,T}) where {W,T<:Union{Float32,Float64}} = vsum(convert(T, x), v)
@inline vprod(x, v::Vec{W,T}) where {W,T<:Union{Float32,Float64}} = vprod(convert(T, x), v)

for (f,f_to,op,reduce,twoarg) ∈ [
    (:reduced_add,:reduce_to_add,:+,:vsum,true),(:reduced_prod,:reduce_to_prod,:*,:vprod,true),
    (:reduced_max,:reduce_to_max,:max,:vmaximum,false),(:reduced_min,:reduce_to_min,:min,:vminimum,false)
]
    @eval begin
        @inline $f_to(x::NativeTypes, y::NativeTypes) = x
        @inline $f_to(x::AbstractSIMD, y::AbstractSIMD) = x
        @inline $f_to(x::AbstractSIMD, y::NativeTypes) = $reduce(x)
        @inline $f(x::NativeTypes, y::NativeTypes) = $op(x,y)
        @inline $f(x::AbstractSIMD, y::AbstractSIMD) = $op(x,y)
    end
    if twoarg
        # @eval @inline $f(y::T, x::AbstractSIMD{W,T}) where {W,T} = $reduce(y, x)
        @eval @inline $f(x::AbstractSIMD, y::NativeTypes) = $reduce(y, x)
        # @eval @inline $f(x::AbstractSIMD, y::NativeTypes) = ((y2,x2,r) = @show (y, x, $reduce(y, x)); r)
    else
        # @eval @inline $f(y::T, x::AbstractSIMD{W,T}) where {W,T} = $op(y, $reduce(x))
        @eval @inline $f(x::AbstractSIMD, y::NativeTypes) = $op(y, $reduce(x))
    end
end

@inline roundint(x::Float32) = round(Int32, x)
if AVX512DQ
    @inline roundint(x::Float64) = round(Int, x)
    @inline roundint(v::Vec{W,Float32}) where {W} = round(Int32, v)
    @inline roundint(v::Vec{W,Float64}) where {W} = round(Int64, v)
else
    @inline roundint(x::Float64) = round(Int32, x)
    @inline roundint(v::Vec{W}) where {W} = round(Int32, v)
end
# binary

function count_zeros_func(W, I, op, tf = 1)
    typ = "i$(8sizeof(I))"
    vtyp = "<$W x $typ>"
    instr = "@llvm.$op.v$(W)$(typ)"
    decl = "declare $vtyp $instr($vtyp, i1)"
    instrs = "%res = call $vtyp $instr($vtyp %0, i1 $tf)\nret $vtyp %res"
    llvmcall_expr(decl, instrs, _Vec{W,I}, Tuple{_Vec{W,I}}, vtyp, (vtyp,), (:(data(v)),))
end
# @generated Base.abs(v::Vec{W,I}) where {W, I <: Integer} = count_zeros_func(W, I, "abs", 0)
@generated Base.leading_zeros(v::Vec{W,I}) where {W, I <: Integer} = count_zeros_func(W, I, "ctlz")
@generated Base.trailing_zeros(v::Vec{W,I}) where {W, I <: Integer} = count_zeros_func(W, I, "cttz")

for (op,f) ∈ [("ctpop", :count_ones)]
    @eval @generated Base.$f(v1::Vec{W,T}) where {W,T} = llvmcall_expr($op, W, T, (W,), (T,))
end

for (op,f) ∈ [("fshl",:funnel_shift_left),("fshr",:funnel_shift_right)
              ]
    @eval @generated function $f(v1::Vec{W,T}, v2::Vec{W,T}, v3::Vec{W,T}) where {W,T}
        llvmcall_expr($op, W, T, (W for _ in 1:3), (T for _ in 1:3))
    end
end
funnel_shift_left(a, b, c) = (a << c) | (b >>> (8sizeof(b) - c))
funnel_shift_right(a, b, c) = (a >>> c) | (b << (8sizeof(b) - c))
@inline rotate_left(a, b) = funnel_shift_left(a, a, b)
@inline rotate_right(a, b) = funnel_shift_right(a, a, b)

# for T ∈ [UInt8,UInt16,UInt32,UInt64]
#     bytes = sizeof(T)
#     W = 2
#     while W * bytes ≤ REGISTER_SIZE

#         for (op,f) ∈ [("ctpop", :count_ones)]
#             @eval @inline Base.$f(v1::Vec{$W,$T}) = $(llvmcall_expr(op, W, T, (W,), (T,)))
#         end

#         for (op,f) ∈ [("fshl",:funnel_shift_left),("fshr",:funnel_shift_right)
#                       ]
#             @eval @inline function $f(v1::Vec{$W,$T}, v2::Vec{$W,$T})
#                 $(llvmcall_expr(op, W, T, (W for _ in 1:3), (T for _ in 1:3)))
#             end
#         end

#         W += W
#     end
# end

@inline vfmadd(a, b, c) = muladd(a, b, c)
@inline vfnmadd(a, b, c) = muladd(-a, b, c)
@inline vfmsub(a, b, c) = muladd(a, b, -c)
@inline vfnmsub(a, b, c) = -muladd(a, b, c)

@inline vfmadd231(a, b, c) = vfmadd(a, b, c)
@inline vfnmadd231(a, b, c) = vfnmadd(a, b, c)
@inline vfmsub231(a, b, c) = vfmsub(a, b, c)
@inline vfnmsub231(a, b, c) = vfnmsub(a, b, c)

if FMA
    @eval begin
        @generated function vfmadd231(a::Vec{W,T}, b::Vec{W,T}, c::Vec{W,T}) where {W, T <: Union{Float32,Float64}}
            if !(ispow2(W) && (W * sizeof(T) ≤ REGISTER_SIZE) && (W ≥ (T === Float32 ? 4 : 2)))
                return Expr(:block, Expr(:meta, :inline), :(vfmadd(a, b, c)))
            end
            typ = LLVM_TYPES[T]
            suffix = T == Float32 ? "ps" : "pd"
            vfmadd_str = """%res = call <$W x $(typ)> asm "vfmadd231$(suffix) \$3, \$2, \$1", "=v,0,v,v"(<$W x $(typ)> %2, <$W x $(typ)> %1, <$W x $(typ)> %0)
                ret <$W x $(typ)> %res"""
            quote
                $(Expr(:meta, :inline))
                Vec(llvmcall($vfmadd_str, _Vec{$W,$T}, Tuple{_Vec{$W,$T},_Vec{$W,$T},_Vec{$W,$T}}, data(a), data(b), data(c)))
            end
        end
        @generated function vfnmadd231(a::Vec{W,T}, b::Vec{W,T}, c::Vec{W,T}) where {W, T <: Union{Float32,Float64}}
            if !(ispow2(W) && (W * sizeof(T) ≤ REGISTER_SIZE) && (W ≥ (T === Float32 ? 4 : 2)))
                return Expr(:block, Expr(:meta, :inline), :(vfnmadd(a, b, c)))
            end
            typ = LLVM_TYPES[T]
            suffix = T == Float32 ? "ps" : "pd"
            vfnmadd_str = """%res = call <$W x $(typ)> asm "vfnmadd231$(suffix) \$3, \$2, \$1", "=v,0,v,v"(<$W x $(typ)> %2, <$W x $(typ)> %1, <$W x $(typ)> %0)
                ret <$W x $(typ)> %res"""
            quote
                $(Expr(:meta, :inline))
                Vec(llvmcall($vfnmadd_str, _Vec{$W,$T}, Tuple{_Vec{$W,$T},_Vec{$W,$T},_Vec{$W,$T}}, data(a), data(b), data(c)))
            end
        end
        @generated function vfmsub231(a::Vec{W,T}, b::Vec{W,T}, c::Vec{W,T}) where {W, T <: Union{Float32,Float64}}
            if !(ispow2(W) && (W * sizeof(T) ≤ REGISTER_SIZE) && (W ≥ (T === Float32 ? 4 : 2)))
                return Expr(:block, Expr(:meta, :inline), :(vfmsub(a, b, c)))
            end
            typ = LLVM_TYPES[T]
            suffix = T == Float32 ? "ps" : "pd"
            vfmsub_str = """%res = call <$W x $(typ)> asm "vfmsub231$(suffix) \$3, \$2, \$1", "=v,0,v,v"(<$W x $(typ)> %2, <$W x $(typ)> %1, <$W x $(typ)> %0)
                ret <$W x $(typ)> %res"""
            quote
                $(Expr(:meta, :inline))
                Vec(llvmcall($vfmsub_str, _Vec{$W,$T}, Tuple{_Vec{$W,$T},_Vec{$W,$T},_Vec{$W,$T}}, data(a), data(b), data(c)))
            end
        end
        @generated function vfnmsub231(a::Vec{W,T}, b::Vec{W,T}, c::Vec{W,T}) where {W, T <: Union{Float32,Float64}}
            if !(ispow2(W) && (W * sizeof(T) ≤ REGISTER_SIZE) && (W ≥ (T === Float32 ? 4 : 2)))
                return Expr(:block, Expr(:meta, :inline), :(vfnmsub(a, b, c)))
            end
            typ = LLVM_TYPES[T]
            suffix = T == Float32 ? "ps" : "pd"
            vfnmsub_str = """%res = call <$W x $(typ)> asm "vfnmsub231$(suffix) \$3, \$2, \$1", "=v,0,v,v"(<$W x $(typ)> %2, <$W x $(typ)> %1, <$W x $(typ)> %0)
                ret <$W x $(typ)> %res"""
            quote
                $(Expr(:meta, :inline))
                Vec(llvmcall($vfnmsub_str, _Vec{$W,$T}, Tuple{_Vec{$W,$T},_Vec{$W,$T},_Vec{$W,$T}}, data(a), data(b), data(c)))
            end
        end
    end
    if AVX512BW
        @eval begin
            @generated function ifelse(::typeof(vfmadd231), m::Mask{W,U}, a::Vec{W,T}, b::Vec{W,T}, c::Vec{W,T}) where {W,U<:Unsigned,T<:Union{Float32,Float64}}
                if !((W ≥ 8) && ispow2(W) && (W * sizeof(T) ≤ REGISTER_SIZE))
                    return Expr(:block, Expr(:meta, :inline), :(ifelse(vfmadd, m, a, b, c)))
                end
                typ = LLVM_TYPES[T]
                suffix = T == Float32 ? "ps" : "pd"
                vfmaddmask_str = """%res = call <$W x $(typ)> asm "vfmadd231$(suffix) \$3, \$2, \$1 {\$4}", "=v,0,v,v,^Yk"(<$W x $(typ)> %2, <$W x $(typ)> %1, <$W x $(typ)> %0, i$W %3)
                                ret <$W x $(typ)> %res"""
                quote
                    $(Expr(:meta,:inline))
                    Vec(llvmcall($vfmaddmask_str, _Vec{$W,$T}, Tuple{_Vec{$W,$T},_Vec{$W,$T},_Vec{$W,$T},$U}, data(a), data(b), data(c), data(m)))
                end
            end
            @generated function ifelse(::typeof(vfnmadd231), m::Mask{W,U}, a::Vec{W,T}, b::Vec{W,T}, c::Vec{W,T}) where {W,U<:Unsigned,T<:Union{Float32,Float64}}
                if !((W ≥ 8) && ispow2(W) && (W * sizeof(T) ≤ REGISTER_SIZE))
                    return Expr(:block, Expr(:meta, :inline), :(ifelse(vfmmadd, m, a, b, c)))
                end
                typ = LLVM_TYPES[T]
                suffix = T == Float32 ? "ps" : "pd"
                vfnmaddmask_str = """%res = call <$W x $(typ)> asm "vfnmadd231$(suffix) \$3, \$2, \$1 {\$4}", "=v,0,v,v,^Yk"(<$W x $(typ)> %2, <$W x $(typ)> %1, <$W x $(typ)> %0, i$W %3)
                            ret <$W x $(typ)> %res"""
                quote
                    $(Expr(:meta,:inline))
                    Vec(llvmcall($vfnmaddmask_str, _Vec{$W,$T}, Tuple{_Vec{$W,$T},_Vec{$W,$T},_Vec{$W,$T},$U}, data(a), data(b), data(c), data(m)))
                end
            end
            @generated function ifelse(::typeof(vfmsub231), m::Mask{W,U}, a::Vec{W,T}, b::Vec{W,T}, c::Vec{W,T}) where {W,U<:Unsigned,T<:Union{Float32,Float64}}
                if !((W ≥ 8) && ispow2(W) && (W * sizeof(T) ≤ REGISTER_SIZE))
                    return Expr(:block, Expr(:meta, :inline), :(ifelse(vfmsub, m, a, b, c)))
                end
                typ = LLVM_TYPES[T]
                suffix = T == Float32 ? "ps" : "pd"
                vfmsubmask_str = """%res = call <$W x $(typ)> asm "vfmsub231$(suffix) \$3, \$2, \$1 {\$4}", "=v,0,v,v,^Yk"(<$W x $(typ)> %2, <$W x $(typ)> %1, <$W x $(typ)> %0, i$W %3)
                            ret <$W x $(typ)> %res"""
                quote
                    $(Expr(:meta,:inline))
                    Vec(llvmcall($vfmsubmask_str, _Vec{$W,$T}, Tuple{_Vec{$W,$T},_Vec{$W,$T},_Vec{$W,$T},$U}, data(a), data(b), data(c), data(m)))
                end
            end
            @generated function ifelse(::typeof(vfnmsub231), m::Mask{W,U}, a::Vec{W,T}, b::Vec{W,T}, c::Vec{W,T}) where {W,U<:Unsigned,T<:Union{Float32,Float64}}
                if !((W ≥ 8) && ispow2(W) && (W * sizeof(T) ≤ REGISTER_SIZE))
                    return Expr(:block, Expr(:meta, :inline), :(ifelse(vfnmsub, m, a, b, c)))
                end
                typ = LLVM_TYPES[T]
                suffix = T == Float32 ? "ps" : "pd"
                vfnmsubmask_str = """%res = call <$W x $(typ)> asm "vfnmsub231$(suffix) \$3, \$2, \$1 {\$4}", "=v,0,v,v,^Yk"(<$W x $(typ)> %2, <$W x $(typ)> %1, <$W x $(typ)> %0, i$W %3)
                            ret <$W x $(typ)> %res"""
                quote
                    $(Expr(:meta,:inline))
                    Vec(llvmcall($vfnmsubmask_str, _Vec{$W,$T}, Tuple{_Vec{$W,$T},_Vec{$W,$T},_Vec{$W,$T},$U}, data(a), data(b), data(c), data(m)))
                end
            end
        end
    end
end
    # for T ∈ [Float32,Float64]
    #     W = 16 ÷ sizeof(T)
    #     local suffix = T == Float32 ? "ps" : "pd"
    #     typ = LLVM_TYPES[T]
    #     while W <= VectorizationBase.REGISTER_SIZE ÷ sizeof(T)
    #         vfmadd_str = """%res = call <$W x $(typ)> asm "vfmadd231$(suffix) \$3, \$2, \$1", "=v,0,v,v"(<$W x $(typ)> %2, <$W x $(typ)> %1, <$W x $(typ)> %0)
    #             ret <$W x $(typ)> %res"""
    #         vfnmadd_str = """%res = call <$W x $(typ)> asm "vfnmadd231$(suffix) \$3, \$2, \$1", "=v,0,v,v"(<$W x $(typ)> %2, <$W x $(typ)> %1, <$W x $(typ)> %0)
    #             ret <$W x $(typ)> %res"""
    #         vfmsub_str = """%res = call <$W x $(typ)> asm "vfmsub231$(suffix) \$3, \$2, \$1", "=v,0,v,v"(<$W x $(typ)> %2, <$W x $(typ)> %1, <$W x $(typ)> %0)
    #             ret <$W x $(typ)> %res"""
    #         vfnmsub_str = """%res = call <$W x $(typ)> asm "vfnmsub231$(suffix) \$3, \$2, \$1", "=v,0,v,v"(<$W x $(typ)> %2, <$W x $(typ)> %1, <$W x $(typ)> %0)
    #             ret <$W x $(typ)> %res"""
    #         @eval begin
    #             @inline function vfmadd231(a::Vec{$W,$T}, b::Vec{$W,$T}, c::Vec{$W,$T})
    #                 Vec(llvmcall($vfmadd_str, _Vec{$W,$T}, Tuple{_Vec{$W,$T},_Vec{$W,$T},_Vec{$W,$T}}, data(a), data(b), data(c)))
    #             end
    #             @inline function vfnmadd231(a::Vec{$W,$T}, b::Vec{$W,$T}, c::Vec{$W,$T})
    #                 Vec(llvmcall($vfnmadd_str, Vec{$W,$T}, Tuple{_Vec{$W,$T},_Vec{$W,$T},_Vec{$W,$T}}, data(a), data(b), data(c)))
    #             end
    #             @inline function vfmsub231(a::Vec{$W,$T}, b::Vec{$W,$T}, c::Vec{$W,$T})
    #                 Vec(llvmcall($vfmsub_str, Vec{$W,$T}, Tuple{_Vec{$W,$T},_Vec{$W,$T},_Vec{$W,$T}}, data(a), data(b), data(c)))
    #             end
    #             @inline function vfnmsub231(a::Vec{$W,$T}, b::Vec{$W,$T}, c::Vec{$W,$T})
    #                 Vec(llvmcall($vfnmsub_str, Vec{$W,$T}, Tuple{_Vec{$W,$T},_Vec{$W,$T},_Vec{$W,$T}}, data(a), data(b), data(c)))
    #             end
    #         end
    #         if VectorizationBase.AVX512BW && W ≥ 8
    #             vfmaddmask_str = """%res = call <$W x $(typ)> asm "vfmadd231$(suffix) \$3, \$2, \$1 {\$4}", "=v,0,v,v,^Yk"(<$W x $(typ)> %2, <$W x $(typ)> %1, <$W x $(typ)> %0, i$W %3)
    #                 ret <$W x $(typ)> %res"""
    #             vfnmaddmask_str = """%res = call <$W x $(typ)> asm "vfnmadd231$(suffix) \$3, \$2, \$1 {\$4}", "=v,0,v,v,^Yk"(<$W x $(typ)> %2, <$W x $(typ)> %1, <$W x $(typ)> %0, i$W %3)
    #                 ret <$W x $(typ)> %res"""
    #             vfmsubmask_str = """%res = call <$W x $(typ)> asm "vfmsub231$(suffix) \$3, \$2, \$1 {\$4}", "=v,0,v,v,^Yk"(<$W x $(typ)> %2, <$W x $(typ)> %1, <$W x $(typ)> %0, i$W %3)
    #                 ret <$W x $(typ)> %res"""
    #             vfnmsubmask_str = """%res = call <$W x $(typ)> asm "vfnmsub231$(suffix) \$3, \$2, \$1 {\$4}", "=v,0,v,v,^Yk"(<$W x $(typ)> %2, <$W x $(typ)> %1, <$W x $(typ)> %0, i$W %3)
    #                 ret <$W x $(typ)> %res"""
    #             U = VectorizationBase.mask_type(W)
    #             @eval begin
    #                 @inline function ifelse(::typeof(vfmadd231), m::Mask{$W,$U}, a::Vec{$W,$T}, b::Vec{$W,$T}, c::Vec{$W,$T})
    #                     Vec(llvmcall($vfmaddmask_str, Vec{$W,$T}, Tuple{_Vec{$W,$T},_Vec{$W,$T},_Vec{$W,$T},$U}, data(a), data(b), data(c), data(m)))
    #                 end
    #                 @inline function ifelse(::typeof(vfnmadd231), m::Mask{$W,$U}, a::Vec{$W,$T}, b::Vec{$W,$T}, c::Vec{$W,$T})
    #                     Vec(llvmcall($vfnmaddmask_str, Vec{$W,$T}, Tuple{_Vec{$W,$T},_Vec{$W,$T},_Vec{$W,$T},$U}, data(a), data(b), data(c), data(m)))
    #                 end
    #                 @inline function ifelse(::typeof(vfmsub231), m::Mask{$W,$U}, a::Vec{$W,$T}, b::Vec{$W,$T}, c::Vec{$W,$T})
    #                     Vec(llvmcall($vfmsubmask_str, Vec{$W,$T}, Tuple{_Vec{$W,$T},_Vec{$W,$T},_Vec{$W,$T},$U}, data(a), data(b), data(c), data(m)))
    #                 end
    #                 @inline function ifelse(::typeof(vfnmsub231), m::Mask{$W,$U}, a::Vec{$W,$T}, b::Vec{$W,$T}, c::Vec{$W,$T})
    #                     Vec(llvmcall($vfnmsubmask_str, Vec{$W,$T}, Tuple{_Vec{$W,$T},_Vec{$W,$T},_Vec{$W,$T},$U}, data(a), data(b), data(c), data(m)))
    #                 end
    #             end
    #         end
    #         W += W
    #     end
# end

@inline ifelse(f::F, m::Mask, a::Vararg{Any,K}) where {F<:Function,K} = ifelse(m, f(a...), a[K])

"""
Fast approximate reciprocal.

Guaranteed accurate to at least 2^-14 ≈ 6.103515625e-5.

Useful for special funcion implementations.
"""
@inline inv_approx(x) = Base.FastMath.inv_fast(x)
@inline inv_approx(v::VecUnroll) = VecUnroll(fmap(inv_approx, v.data))
@generated function inv_approx(v::Vec{W,T}) where {W, T <: Union{Float32, Float64}}
    ((Sys.ARCH === :x86_64) || (Sys.ARCH === :i686)) || return Expr(:block, Expr(:meta, :inline), :(inv(v)))
    bits = 8sizeof(T) * W
    if (AVX512F && (bits === 512)) || (AVX512VL && (bits ∈ (128, 256)))
        typ = T === Float64 ? "double" : "float"
        vtyp = "<$W x $(typ)>"
        dors = T === Float64 ? "d" : "s"
        f = "@llvm.x86.avx512.rcp14.p$(dors).$(bits)"
        decl = "declare $(vtyp) $f($(vtyp), $(vtyp), i$(min(8,W))) nounwind readnone"
        instrs = "%res = call $(vtyp) $f($vtyp %0, $vtyp zeroinitializer, i$(min(8,W)) -1)\nret $(vtyp) %res"
        return llvmcall_expr(decl, instrs, :(_Vec{$W,$T}), :(Tuple{_Vec{$W,$T}}), vtyp, [vtyp], [:(data(v))])
    end
    if (AVX && (W == 8)) && (T === Float32)
        decl = "declare <8 x float> @llvm.x86.avx.rcp.ps.256(<8 x float>) nounwind readnone"
        instrs = "%res = call <8 x float> @llvm.x86.avx.rcp.ps.256(<8 x float> %0)\nret <8 x float> %res"
        return llvmcall_expr(decl, instrs, :(_Vec{8,Float32}), :(Tuple{_Vec{8,Float32}}), "<8 x float>", ["<8 x float>"], [:(data(v))])
    elseif W == 4
        decl = "declare <4 x float> @llvm.x86.sse.rcp.ps(<4 x float>) nounwind readnone"
        instrs = "%res = call <4 x float> @llvm.x86.sse.rcp.ps(<4 x float> %0)\nret <4 x float> %res"
        if T === Float32
            return llvmcall_expr(decl, instrs, :(_Vec{4,Float32}), :(Tuple{_Vec{4,Float32}}), "<4 x float>", ["<4 x float>"], [:(data(v))])
        else#if T === Float64
            argexpr = [:(data(convert(Float32, v)))]
            call = llvmcall_expr(decl, instrs, :(_Vec{4,Float32}), :(Tuple{_Vec{4,Float32}}), "<4 x float>", ["<4 x float>"], argexpr, true)
            return Expr(:block, Expr(:meta, :inline), :(convert(Float64, $call)))
        end
    end
    Expr(:block, Expr(:meta, :inline), :(inv(v)))
end


@testset "1220.count-vowels-permutation.jl" begin
    @test count_vowel_permutation(1) == 5
    @test count_vowel_permutation(2) == 10
    @test count_vowel_permutation(5) == 68
    @test count_vowel_permutation(10^4) == 76428576
end


using Test
@isdefined(TestHelpers) || include("../TestHelpers.jl")

TestHelpers.@timed_testset "layers/" begin
    include("layers/conv2d.jl")
    include("layers/flatten.jl")
    include("layers/linear.jl")
    include("layers/masked_relu.jl")
    include("layers/pool.jl")
    include("layers/relu.jl")
end


function make_operation(@nospecialize(op), args)
    if op === (*)
        args = filter(!_isone, args)
        if isempty(args)
            return 1
        end
    elseif op === (+)
        args = filter(!_iszero, args)
        if isempty(args)
            return 0
        end
    end
    return op(args...)
end

function detime_dvs(op)
    if !istree(op)
        op
    elseif operation(op) isa Sym
        Sym{Real}(nameof(operation(op)))
    else
        similarterm(op, operation(op),detime_dvs.(arguments(op)))
    end
end

function retime_dvs(op::Sym,dvs,iv)
    Sym{FnType{Tuple{symtype(iv)}, Real}}(nameof(op))(iv)
end

function retime_dvs(op, dvs, iv)
    istree(op) ?
        similarterm(op, operation(op), retime_dvs.(arguments(op),(dvs,),(iv,))) :
        op
end

modified_states!(mstates, e::Equation, statelist=nothing) = get_variables!(mstates, e.lhs, statelist)

macro showarr(x)
    n = string(x)
    quote
        y = $(esc(x))
        println($n, " = ", summary(y))
        Base.print_array(stdout, y)
        println()
        y
    end
end

@deprecate substitute_expr!(expr,s) substitute(expr,s)

function states_to_sym(states::Set)
    function _states_to_sym(O)
        if O isa Equation
            Expr(:(=), _states_to_sym(O.lhs), _states_to_sym(O.rhs))
        elseif istree(O)
            op = operation(O)
            args = arguments(O)
            if op isa Sym
                O in states && return tosymbol(O)
                # dependent variables
                return build_expr(:call, Any[nameof(op); _states_to_sym.(args)])
            else
                canonical, O = canonicalexpr(O)
                return canonical ? O : build_expr(:call, Any[op; _states_to_sym.(args)])
            end
        elseif O isa Num
            return _states_to_sym(value(O))
        else
            return toexpr(O)
        end
    end
end
states_to_sym(states) = states_to_sym(Set(states))

function todict(d)
    eltype(d) <: Pair || throw(ArgumentError("The variable-value mapping must be a Dict."))
    d isa Dict ? d : Dict(d)
end

_merge(d1, d2) = merge(todict(d1), todict(d2))


export write_extxyz
export read_extxyz

import ExtXYZ

"""
    write_extxyz(file, atoms, R, cell)

Uses ExtXYZ.jl to write a .extxyz file

R can be either a single configuration (Matrix) or many (Vector{<:Matrix}).
"""
write_extxyz(file, atoms, R::Matrix, cell) =
    ExtXYZ.write_frame(file, to_extxyz_dict(atoms, R, cell))

write_extxyz(file, atoms, R::Vector{<:Matrix}, cell) =
    ExtXYZ.write_frames(file, to_extxyz_dict.(Ref(atoms), R, Ref(cell)))

"""
    read_extxyz(file)

Uses ExtXYZ.jl to read a .extxyz file
"""
function read_extxyz(file)

    dict = ExtXYZ.read_frame(file)
    atoms, R, cell = from_extxyz_dict(dict)
    positions = [R]

    for dict in ExtXYZ.read_frames(file, Iterators.countfrom(2))
        _, R, _ = from_extxyz_dict(dict)
        push!(positions, R)
    end

    return atoms, positions, cell
end

function to_extxyz_dict(atoms::Atoms, R::Matrix, cell::PeriodicCell)

    dict = Dict{String, Any}()
    dict["N_atoms"] = length(atoms)

    dict["pbc"] = cell.periodicity

    dict["cell"] = au_to_ang.(permutedims(cell.vectors, (2,1)))

    dict["arrays"] = Dict{String,Any}()
    dict["arrays"]["species"] = String.(atoms.types) |> Array
    dict["arrays"]["pos"] = au_to_ang.(R)

    dict["info"] = Dict{String,Any}()

    return dict
end

function from_extxyz_dict(dict::Dict{String,Any})

    atoms = Atoms(Symbol.(dict["arrays"]["species"]))
    R = ang_to_au.(dict["arrays"]["pos"])
    c = permutedims(dict["cell"], (2,1))
    cell = PeriodicCell{eltype(c)}(ang_to_au.(c), dict["pbc"])

    return atoms, R, cell
end


module StringDistances

using Distances
import Distances: evaluate, result_type

isnormalized(dist::SemiMetric) = false

include("utils.jl")
include("edit.jl")
include("qgram.jl")
include("modifier.jl")

const StringDistance = Union{Jaro, Levenshtein, DamerauLevenshtein, RatcliffObershelp, QGramDistance, Winkler, Partial, TokenSort, TokenSet, TokenMax}
function result_type(dist::StringDistance, s1, s2)
    typeof(evaluate(dist, "", ""))
end

"""
    compare(s1, s2, dist)

return a similarity score between 0 and 1 for the strings `s1` and
`s2` based on the distance `dist`.

### Examples
```julia-repl
julia> compare("martha", "marhta", Levenshtein())
0.6666666666666667
```
"""
compare(s1, s2, dist::StringDistance; min_score = 0.0) = 1 - evaluate(normalize(dist), s1, s2, 1 - min_score)

const metrics = [Cosine, DamerauLevenshtein, Jaccard, Jaro, Levenshtein, Overlap, Partial, QGram,
    RatcliffObershelp, SorensenDice, StringDistance,  TokenMax,
    TokenSet, TokenSort, Winkler]

for M in metrics
    @eval @inline (dist::$M)(a::AbstractString, b::AbstractString) = evaluate(dist, a, b)
end

include("find.jl")

##############################################################################
##
## Export
##
##############################################################################

export
StringDistance,
Levenshtein,
DamerauLevenshtein,
Jaro,
RatcliffObershelp,
QGram,
Cosine,
Jaccard,
SorensenDice,
Overlap,
Winkler,
Partial,
TokenSort,
TokenSet,
TokenMax,
evaluate,
compare,
result_type,
qgrams,
normalize
end

##############################################################################
##
## Some things about Strings

# length: number of characters
# ncodeunits: Return the number of code units in a string (aking to index of vector).
# Not all such indices  are valid – they may not be the start of a character,.
# sizeof:  Size, in bytes, of the string str. Equal to the number of code units in str
# multiplied by the size, in bytes, of one code unit in str.

# lastindex: Return the last index of a collection
# nextinds(s, i):  return the index of the start of the character whose encoding starts after index i
# nextind(s, 0, N): return the index of the Nth character of s (or, if there are
# less than N characters, return ncodeunits(str) + (N - length(s))

##############################################################################


#### EDMF model kernels

using CLIMAParameters.Planet: e_int_v0, grav, day, R_d, R_v, molmass_ratio
using Printf
using ClimateMachine.Atmos: nodal_update_auxiliary_state!, Advect

using ClimateMachine.BalanceLaws

using ClimateMachine.MPIStateArrays: MPIStateArray
using ClimateMachine.DGMethods: LocalGeometry, DGModel

import ClimateMachine.Atmos: atmos_source!

import ClimateMachine.BalanceLaws:
    vars_state,
    prognostic_vars,
    flux,
    source,
    eq_tends,
    update_auxiliary_state!,
    init_state_prognostic!,
    flux_first_order!,
    flux_second_order!,
    compute_gradient_argument!,
    compute_gradient_flux!

import ClimateMachine.TurbulenceConvection:
    init_aux_turbconv!,
    turbconv_nodal_update_auxiliary_state!,
    turbconv_boundary_state!,
    turbconv_normal_boundary_flux_second_order!

using ClimateMachine.Thermodynamics: air_pressure, air_density

include(joinpath("helper_funcs", "nondimensional_exchange_functions.jl"))
include(joinpath("helper_funcs", "lamb_smooth_minimum.jl"))
include(joinpath("helper_funcs", "utility_funcs.jl"))
include(joinpath("helper_funcs", "subdomain_statistics.jl"))
include(joinpath("helper_funcs", "diagnose_environment.jl"))
include(joinpath("helper_funcs", "subdomain_thermo_states.jl"))
include(joinpath("helper_funcs", "save_subdomain_temperature.jl"))
include(joinpath("closures", "entr_detr.jl"))
include(joinpath("closures", "pressure.jl"))
include(joinpath("closures", "mixing_length.jl"))
include(joinpath("closures", "turbulence_functions.jl"))
include(joinpath("closures", "surface_functions.jl"))

function vars_state(m::NTuple{N, Updraft}, st::Auxiliary, FT) where {N}
    return Tuple{ntuple(i -> vars_state(m[i], st, FT), N)...}
end

function vars_state(::Updraft, ::Auxiliary, FT)
    @vars(
        buoyancy::FT,
        a::FT,
        E_dyn::FT,
        Δ_dyn::FT,
        E_trb::FT,
        T::FT,
        θ_liq::FT,
        q_tot::FT,
        w::FT,
    )
end

function vars_state(::Environment, ::Auxiliary, FT)
    @vars(T::FT, cld_frac::FT, buoyancy::FT)
end

function vars_state(m::EDMF, st::Auxiliary, FT)
    @vars(
        environment::vars_state(m.environment, st, FT),
        updraft::vars_state(m.updraft, st, FT)
    )
end

function vars_state(::Updraft, ::Prognostic, FT)
    @vars(ρa::FT, ρaw::FT, ρaθ_liq::FT, ρaq_tot::FT,)
end

function vars_state(::Environment, ::Prognostic, FT)
    @vars(ρatke::FT, ρaθ_liq_cv::FT, ρaq_tot_cv::FT, ρaθ_liq_q_tot_cv::FT,)
end

function vars_state(m::NTuple{N, Updraft}, st::Prognostic, FT) where {N}
    return Tuple{ntuple(i -> vars_state(m[i], st, FT), N)...}
end

function vars_state(m::EDMF, st::Prognostic, FT)
    @vars(
        environment::vars_state(m.environment, st, FT),
        updraft::vars_state(m.updraft, st, FT)
    )
end

function vars_state(::Updraft, ::Gradient, FT)
    @vars(w::FT,)
end

function vars_state(::Environment, ::Gradient, FT)
    @vars(
        θ_liq::FT,
        q_tot::FT,
        w::FT,
        tke::FT,
        θ_liq_cv::FT,
        q_tot_cv::FT,
        θ_liq_q_tot_cv::FT,
        θv::FT,
        e::FT,
    )
end

function vars_state(m::NTuple{N, Updraft}, st::Gradient, FT) where {N}
    return Tuple{ntuple(i -> vars_state(m[i], st, FT), N)...}
end

function vars_state(m::EDMF, st::Gradient, FT)
    @vars(
        environment::vars_state(m.environment, st, FT),
        updraft::vars_state(m.updraft, st, FT)
    )
end

function vars_state(m::NTuple{N, Updraft}, st::GradientFlux, FT) where {N}
    return Tuple{ntuple(i -> vars_state(m[i], st, FT), N)...}
end

function vars_state(::Updraft, st::GradientFlux, FT)
    @vars(∇w::SVector{3, FT},)
end

function vars_state(::Environment, ::GradientFlux, FT)
    @vars(
        ∇θ_liq::SVector{3, FT},
        ∇q_tot::SVector{3, FT},
        ∇w::SVector{3, FT},
        ∇tke::SVector{3, FT},
        ∇θ_liq_cv::SVector{3, FT},
        ∇q_tot_cv::SVector{3, FT},
        ∇θ_liq_q_tot_cv::SVector{3, FT},
        ∇θv::SVector{3, FT},
        ∇e::SVector{3, FT},
    )

end

function vars_state(m::EDMF, st::GradientFlux, FT)
    @vars(
        S²::FT, # should be conditionally grabbed from atmos.turbulence
        environment::vars_state(m.environment, st, FT),
        updraft::vars_state(m.updraft, st, FT)
    )
end

abstract type EDMFPrognosticVariable <: PrognosticVariable end

abstract type EnvironmentPrognosticVariable <: EDMFPrognosticVariable end
struct en_ρatke <: EnvironmentPrognosticVariable end
struct en_ρaθ_liq_cv <: EnvironmentPrognosticVariable end
struct en_ρaq_tot_cv <: EnvironmentPrognosticVariable end
struct en_ρaθ_liq_q_tot_cv <: EnvironmentPrognosticVariable end

abstract type UpdraftPrognosticVariable{i} <: EDMFPrognosticVariable end
struct up_ρa{i} <: UpdraftPrognosticVariable{i} end
struct up_ρaw{i} <: UpdraftPrognosticVariable{i} end
struct up_ρaθ_liq{i} <: UpdraftPrognosticVariable{i} end
struct up_ρaq_tot{i} <: UpdraftPrognosticVariable{i} end

prognostic_vars(m::EDMF) =
    (prognostic_vars(m.environment)..., prognostic_vars(m.updraft)...)
prognostic_vars(m::Environment) =
    (en_ρatke(), en_ρaθ_liq_cv(), en_ρaq_tot_cv(), en_ρaθ_liq_q_tot_cv())

function prognostic_vars(m::NTuple{N, Updraft}) where {N}
    t_ρa = ntuple(i -> up_ρa{i}(), N)
    t_ρaw = ntuple(i -> up_ρaw{i}(), N)
    t_ρaθ_liq = ntuple(i -> up_ρaθ_liq{i}(), N)
    t_ρaq_tot = ntuple(i -> up_ρaq_tot{i}(), N)
    t = (t_ρa..., t_ρaw..., t_ρaθ_liq..., t_ρaq_tot...)
    return t
end

# Dycore tendencies
eq_tends(
    pv::PV,
    m::EDMF,
    ::Flux{SecondOrder},
) where {PV <: Union{Momentum, Energy, TotalMoisture}} = ()
# (SGSFlux{PV}(),) # to add SGSFlux back to grid-mean

# Turbconv tendencies
eq_tends(
    pv::PV,
    m::AtmosModel,
    tt::Flux{O},
) where {O, PV <: EDMFPrognosticVariable} = eq_tends(pv, m.turbconv, tt)

eq_tends(pv::PV, m::EDMF, ::Flux{O}) where {O, PV <: EDMFPrognosticVariable} =
    ()

eq_tends(
    pv::PV,
    m::EDMF,
    ::Flux{FirstOrder},
) where {PV <: EDMFPrognosticVariable} = (Advect{PV}(),)

struct SGSFlux{PV <: Union{Momentum, Energy, TotalMoisture}} <:
       TendencyDef{Flux{SecondOrder}, PV} end

"""
    init_aux_turbconv!(
        turbconv::EDMF{FT},
        m::AtmosModel{FT},
        aux::Vars,
        geom::LocalGeometry,
    ) where {FT}

Initialize EDMF auxiliary variables.
"""
function init_aux_turbconv!(
    turbconv::EDMF{FT},
    m::AtmosModel{FT},
    aux::Vars,
    geom::LocalGeometry,
) where {FT}
    N_up = n_updrafts(turbconv)

    # Aliases:
    en_aux = aux.turbconv.environment
    up_aux = aux.turbconv.updraft

    en_aux.cld_frac = FT(0)
    en_aux.buoyancy = FT(0)

    @unroll_map(N_up) do i
        up_aux[i].buoyancy = FT(0)
        up_aux[i].θ_liq = FT(0)
        up_aux[i].q_tot = FT(0)
        up_aux[i].w = FT(0)
    end
end;

function turbconv_nodal_update_auxiliary_state!(
    turbconv::EDMF{FT},
    m::AtmosModel{FT},
    state::Vars,
    aux::Vars,
    t::Real,
) where {FT}
    N_up = n_updrafts(turbconv)
    save_subdomain_temperature!(m, state, aux)

    en_aux = aux.turbconv.environment
    up_aux = aux.turbconv.updraft
    gm = state
    en = state.turbconv.environment
    up = state.turbconv.updraft

    # Recover thermo states
    ts = recover_thermo_state_all(m, state, aux)

    # Get environment variables
    env = environment_vars(state, aux, N_up)

    # Compute buoyancies of subdomains
    ρ_inv = 1 / gm.ρ
    _grav::FT = grav(m.param_set)

    z = altitude(m, aux)

    ρ_en = air_density(ts.en)
    en_aux.buoyancy = -_grav * (ρ_en - aux.ref_state.ρ) * ρ_inv

    @unroll_map(N_up) do i
        ρ_i = air_density(ts.up[i])
        up_aux[i].buoyancy = -_grav * (ρ_i - aux.ref_state.ρ) * ρ_inv
        up_aux[i].a = up[i].ρa * ρ_inv
        up_aux[i].θ_liq = up[i].ρaθ_liq / up[i].ρa
        up_aux[i].q_tot = up[i].ρaq_tot / up[i].ρa
        up_aux[i].w = up[i].ρaw / up[i].ρa
    end
    b_gm = grid_mean_b(state, aux, N_up)

    # remove the gm_b from all subdomains
    @unroll_map(N_up) do i
        up_aux[i].buoyancy -= b_gm
    end
    en_aux.buoyancy -= b_gm

    EΔ_up = ntuple(N_up) do i
        entr_detr(m, m.turbconv.entr_detr, state, aux, t, ts, env, i)
    end

    E_dyn, Δ_dyn, E_trb = ntuple(i -> map(x -> x[i], EΔ_up), 3)

    @unroll_map(N_up) do i
        up_aux[i].E_dyn = E_dyn[i]
        up_aux[i].Δ_dyn = Δ_dyn[i]
        up_aux[i].E_trb = E_trb[i]
    end

end;

function compute_gradient_argument!(
    turbconv::EDMF{FT},
    m::AtmosModel{FT},
    transform::Vars,
    state::Vars,
    aux::Vars,
    t::Real,
) where {FT}
    N_up = n_updrafts(turbconv)
    z = altitude(m, aux)

    # Aliases:
    up_tf = transform.turbconv.updraft
    en_tf = transform.turbconv.environment
    gm = state
    up = state.turbconv.updraft
    en = state.turbconv.environment

    # Recover thermo states
    ts = recover_thermo_state_all(m, state, aux)

    # Get environment variables
    env = environment_vars(state, aux, N_up)

    @unroll_map(N_up) do i
        up_tf[i].w = up[i].ρaw / up[i].ρa
    end
    _grav::FT = grav(m.param_set)

    ρ_inv = 1 / gm.ρ
    θ_liq_en = liquid_ice_pottemp(ts.en)
    q_tot_en = total_specific_humidity(ts.en)

    # populate gradient arguments
    en_tf.θ_liq = θ_liq_en
    en_tf.q_tot = q_tot_en
    en_tf.w = env.w

    en_tf.tke = enforce_positivity(en.ρatke) / (env.a * gm.ρ)
    en_tf.θ_liq_cv = enforce_positivity(en.ρaθ_liq_cv) / (env.a * gm.ρ)
    en_tf.q_tot_cv = enforce_positivity(en.ρaq_tot_cv) / (env.a * gm.ρ)
    en_tf.θ_liq_q_tot_cv = en.ρaθ_liq_q_tot_cv / (env.a * gm.ρ)

    # TODO: is this supposed to be grabbed from grid mean?
    en_tf.θv = virtual_pottemp(ts.gm)
    e_kin = FT(1 // 2) * ((gm.ρu[1] * ρ_inv)^2 + (gm.ρu[2] * ρ_inv)^2 + env.w^2)
    en_tf.e = total_energy(e_kin, _grav * z, ts.en)
end;

function compute_gradient_flux!(
    turbconv::EDMF{FT},
    m::AtmosModel{FT},
    diffusive::Vars,
    ∇transform::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
) where {FT}
    N_up = n_updrafts(turbconv)

    # Aliases:
    gm = state
    up_dif = diffusive.turbconv.updraft
    tc_dif = diffusive.turbconv
    up_∇tf = ∇transform.turbconv.updraft
    en_dif = diffusive.turbconv.environment
    en_∇tf = ∇transform.turbconv.environment

    @unroll_map(N_up) do i
        up_dif[i].∇w = up_∇tf[i].w
    end

    ρ_inv = 1 / gm.ρ
    # first moment grid mean coming from environment gradients only
    en_dif.∇θ_liq = en_∇tf.θ_liq
    en_dif.∇q_tot = en_∇tf.q_tot
    en_dif.∇w = en_∇tf.w
    # second moment env cov
    en_dif.∇tke = en_∇tf.tke
    en_dif.∇θ_liq_cv = en_∇tf.θ_liq_cv
    en_dif.∇q_tot_cv = en_∇tf.q_tot_cv
    en_dif.∇θ_liq_q_tot_cv = en_∇tf.θ_liq_q_tot_cv

    en_dif.∇θv = en_∇tf.θv
    en_dif.∇e = en_∇tf.e

    tc_dif.S² = ∇transform.u[3, 1]^2 + ∇transform.u[3, 2]^2 + en_dif.∇w[3]^2
end;

struct TurbconvSource <: AbstractSource end

function atmos_source!(
    ::TurbconvSource,
    m::AtmosModel{FT},
    source::Vars,
    state::Vars,
    diffusive::Vars,
    aux::Vars,
    t::Real,
    direction,
) where {FT}
    turbconv = m.turbconv
    N_up = n_updrafts(turbconv)

    # Aliases:
    gm = state
    en = state.turbconv.environment
    up = state.turbconv.updraft
    en_src = source.turbconv.environment
    up_src = source.turbconv.updraft
    en_dif = diffusive.turbconv.environment
    up_aux = aux.turbconv.updraft

    # Recover thermo states
    ts = recover_thermo_state_all(m, state, aux)

    # Get environment variables
    env = environment_vars(state, aux, N_up)

    EΔ_up = ntuple(N_up) do i
        entr_detr(m, m.turbconv.entr_detr, state, aux, t, ts, env, i)
    end
    E_dyn, Δ_dyn, E_trb = ntuple(i -> map(x -> x[i], EΔ_up), 3)

    # get environment values
    _grav::FT = grav(m.param_set)
    ρ_inv = 1 / gm.ρ
    θ_liq_en = liquid_ice_pottemp(ts.en)
    q_tot_en = total_specific_humidity(ts.en)
    tke_en = enforce_positivity(en.ρatke) * ρ_inv / env.a
    θ_liq = liquid_ice_pottemp(ts.gm)
    a_min = turbconv.subdomains.a_min
    a_max = turbconv.subdomains.a_max

    ρa_up = vuntuple(N_up) do i
        gm.ρ * enforce_unit_bounds(up[i].ρa * ρ_inv, a_min, a_max)
    end
    ρq_tot = m.moisture isa DryModel ? FT(0) : gm.moisture.ρq_tot

    @unroll_map(N_up) do i

        ρa_up_i = ρa_up[i]
        w_up_i = up[i].ρaw / ρa_up_i
        ρa_up_i_inv = FT(1) / ρa_up_i

        # first moment sources - for now we compute these as aux variable
        dpdz = perturbation_pressure(
            m,
            m.turbconv.pressure,
            state,
            diffusive,
            aux,
            t,
            env,
            i,
        )

        # entrainment and detrainment
        up_src[i].ρa += E_dyn[i] - Δ_dyn[i]
        up_src[i].ρaw +=
            ((E_dyn[i] + E_trb[i]) * env.w - (Δ_dyn[i] + E_trb[i]) * w_up_i)
        up_src[i].ρaθ_liq += (
            (E_dyn[i] + E_trb[i]) * θ_liq_en -
            (Δ_dyn[i] + E_trb[i]) * up[i].ρaθ_liq * ρa_up_i_inv
        )
        up_src[i].ρaq_tot += (
            (E_dyn[i] + E_trb[i]) * q_tot_en -
            (Δ_dyn[i] + E_trb[i]) * up[i].ρaq_tot * ρa_up_i_inv
        )

        # add buoyancy and perturbation pressure in subdomain w equation
        up_src[i].ρaw += up[i].ρa * (up_aux[i].buoyancy - dpdz)
        # microphysics sources should be applied here

        # environment second moments:
        en_src.ρatke += (
            Δ_dyn[i] * (w_up_i - env.w) * (w_up_i - env.w) * FT(0.5) +
            E_trb[i] * (env.w - gm.ρu[3] * ρ_inv) * (env.w - w_up_i) -
            (E_dyn[i] + E_trb[i]) * tke_en
        )

        en_src.ρaθ_liq_cv += (
            Δ_dyn[i] *
            (up[i].ρaθ_liq * ρa_up_i_inv - θ_liq_en) *
            (up[i].ρaθ_liq * ρa_up_i_inv - θ_liq_en) +
            E_trb[i] *
            (θ_liq_en - θ_liq) *
            (θ_liq_en - up[i].ρaθ_liq * ρa_up_i_inv) +
            E_trb[i] *
            (θ_liq_en - θ_liq) *
            (θ_liq_en - up[i].ρaθ_liq * ρa_up_i_inv) -
            (E_dyn[i] + E_trb[i]) * en.ρaθ_liq_cv
        )

        en_src.ρaq_tot_cv += (
            Δ_dyn[i] *
            (up[i].ρaq_tot * ρa_up_i_inv - q_tot_en) *
            (up[i].ρaq_tot * ρa_up_i_inv - q_tot_en) +
            E_trb[i] *
            (q_tot_en - ρq_tot * ρ_inv) *
            (q_tot_en - up[i].ρaq_tot * ρa_up_i_inv) +
            E_trb[i] *
            (q_tot_en - ρq_tot * ρ_inv) *
            (q_tot_en - up[i].ρaq_tot * ρa_up_i_inv) -
            (E_dyn[i] + E_trb[i]) * en.ρaq_tot_cv
        )

        en_src.ρaθ_liq_q_tot_cv += (
            Δ_dyn[i] *
            (up[i].ρaθ_liq * ρa_up_i_inv - θ_liq_en) *
            (up[i].ρaq_tot * ρa_up_i_inv - q_tot_en) +
            E_trb[i] *
            (θ_liq_en - θ_liq) *
            (q_tot_en - up[i].ρaq_tot * ρa_up_i_inv) +
            E_trb[i] *
            (q_tot_en - ρq_tot * ρ_inv) *
            (θ_liq_en - up[i].ρaθ_liq * ρa_up_i_inv) -
            (E_dyn[i] + E_trb[i]) * en.ρaθ_liq_q_tot_cv
        )

        # pressure tke source from the i'th updraft
        en_src.ρatke += ρa_up_i * (w_up_i - env.w) * dpdz
    end
    l_mix, ∂b∂z_env, Pr_t = mixing_length(
        m,
        m.turbconv.mix_len,
        state,
        diffusive,
        aux,
        t,
        Δ_dyn,
        E_trb,
        ts,
        env,
    )

    K_m = m.turbconv.mix_len.c_m * l_mix * sqrt(tke_en)
    K_h = K_m / Pr_t
    Shear² = diffusive.turbconv.S²
    ρa₀ = gm.ρ * env.a
    Diss₀ = m.turbconv.mix_len.c_d * sqrt(tke_en) / l_mix

    # production from mean gradient and Dissipation
    en_src.ρatke += ρa₀ * K_m * Shear² # tke Shear source
    en_src.ρatke += -ρa₀ * K_h * ∂b∂z_env   # tke Bouyancy source
    en_src.ρatke += -ρa₀ * Diss₀ * tke_en  # tke Dissipation

    en_src.ρaθ_liq_cv +=
        ρa₀ *
        (K_h * en_dif.∇θ_liq[3] * en_dif.∇θ_liq[3] - Diss₀ * en.ρaθ_liq_cv)
    en_src.ρaq_tot_cv +=
        ρa₀ *
        (K_h * en_dif.∇q_tot[3] * en_dif.∇q_tot[3] - Diss₀ * en.ρaq_tot_cv)
    en_src.ρaθ_liq_q_tot_cv +=
        ρa₀ * (
            K_h * en_dif.∇θ_liq[3] * en_dif.∇q_tot[3] -
            Diss₀ * en.ρaθ_liq_q_tot_cv
        )
    # covariance microphysics sources should be applied here
end;

function compute_ρa_up(m, state, aux)
    # Aliases:
    turbconv = m.turbconv
    gm = state
    up = state.turbconv.updraft
    N_up = n_updrafts(turbconv)
    a_min = turbconv.subdomains.a_min
    a_max = turbconv.subdomains.a_max
    # in future GCM implementations we need to think about grid mean advection
    ρa_up = vuntuple(N_up) do i
        gm.ρ * enforce_unit_bounds(up[i].ρa / gm.ρ, a_min, a_max)
    end
    return ρa_up
end

function flux(::Advect{up_ρa{i}}, m, state, aux, t, ts, direction) where {i}
    up = state.turbconv.updraft
    ẑ = vertical_unit_vector(m, aux)
    return up[i].ρaw * ẑ
end
function flux(::Advect{up_ρaw{i}}, m, state, aux, t, ts, direction) where {i}
    up = state.turbconv.updraft
    ẑ = vertical_unit_vector(m, aux)
    ρa_up = compute_ρa_up(m, state, aux)
    return up[i].ρaw * up[i].ρaw / ρa_up[i] * ẑ
end
function flux(
    ::Advect{up_ρaθ_liq{i}},
    m,
    state,
    aux,
    t,
    ts,
    direction,
) where {i}
    up = state.turbconv.updraft
    ẑ = vertical_unit_vector(m, aux)
    ρa_up = compute_ρa_up(m, state, aux)
    return up[i].ρaw / ρa_up[i] * up[i].ρaθ_liq * ẑ
end
function flux(
    ::Advect{up_ρaq_tot{i}},
    m,
    state,
    aux,
    t,
    ts,
    direction,
) where {i}
    up = state.turbconv.updraft
    ẑ = vertical_unit_vector(m, aux)
    ρa_up = compute_ρa_up(m, state, aux)
    return up[i].ρaw / ρa_up[i] * up[i].ρaq_tot * ẑ
end

function flux(::Advect{en_ρatke}, m, state, aux, t, ts, direction)
    en = state.turbconv.environment
    env = environment_vars(state, aux, n_updrafts(m.turbconv))
    ẑ = vertical_unit_vector(m, aux)
    return en.ρatke * env.w * ẑ
end
function flux(::Advect{en_ρaθ_liq_cv}, m, state, aux, t, ts, direction)
    en = state.turbconv.environment
    env = environment_vars(state, aux, n_updrafts(m.turbconv))
    ẑ = vertical_unit_vector(m, aux)
    return en.ρaθ_liq_cv * env.w * ẑ
end
function flux(::Advect{en_ρaq_tot_cv}, m, state, aux, t, ts, direction)
    en = state.turbconv.environment
    env = environment_vars(state, aux, n_updrafts(m.turbconv))
    ẑ = vertical_unit_vector(m, aux)
    return en.ρaq_tot_cv * env.w * ẑ
end
function flux(::Advect{en_ρaθ_liq_q_tot_cv}, m, state, aux, t, ts, direction)
    en = state.turbconv.environment
    env = environment_vars(state, aux, n_updrafts(m.turbconv))
    ẑ = vertical_unit_vector(m, aux)
    return en.ρaθ_liq_q_tot_cv * env.w * ẑ
end

# # in the EDMF first order (advective) fluxes exist only in the grid mean (if <w> is nonzero) and the uprdafts
function flux_first_order!(
    turbconv::EDMF{FT},
    m::AtmosModel{FT},
    flux::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
    ts,
    direction,
) where {FT}
    # Aliases:
    up_flx = flux.turbconv.updraft
    en_flx = flux.turbconv.environment
    N_up = n_updrafts(turbconv)
    # in future GCM implementations we need to think about grid mean advection
    tend = Flux{FirstOrder}()
    args = (m, state, aux, t, ts, direction)

    @unroll_map(N_up) do i
        up_flx[i].ρa = Σfluxes(eq_tends(up_ρa{i}(), m, tend), args...)
        up_flx[i].ρaw = Σfluxes(eq_tends(up_ρaw{i}(), m, tend), args...)
        up_flx[i].ρaθ_liq = Σfluxes(eq_tends(up_ρaθ_liq{i}(), m, tend), args...)
        up_flx[i].ρaq_tot = Σfluxes(eq_tends(up_ρaq_tot{i}(), m, tend), args...)
    end
    en_flx.ρatke = Σfluxes(eq_tends(en_ρatke(), m, tend), args...)
    en_flx.ρaθ_liq_cv = Σfluxes(eq_tends(en_ρaθ_liq_cv(), m, tend), args...)
    en_flx.ρaq_tot_cv = Σfluxes(eq_tends(en_ρaq_tot_cv(), m, tend), args...)
    en_flx.ρaθ_liq_q_tot_cv =
        Σfluxes(eq_tends(en_ρaθ_liq_q_tot_cv(), m, tend), args...)
end;

# in the EDMF second order (diffusive) fluxes
# exist only in the grid mean and the environment
function flux_second_order!(
    turbconv::EDMF{FT},
    m::AtmosModel{FT},
    flux::Grad,
    state::Vars,
    diffusive::Vars,
    aux::Vars,
    t::Real,
) where {FT}
    N_up = n_updrafts(turbconv)

    # Aliases:
    gm = state
    up = state.turbconv.updraft
    en = state.turbconv.environment
    gm_flx = flux
    en_flx = flux.turbconv.environment
    en_dif = diffusive.turbconv.environment

    # Recover thermo states
    ts = recover_thermo_state_all(m, state, aux)

    # Get environment variables
    env = environment_vars(state, aux, N_up)

    ρ_inv = FT(1) / gm.ρ
    _grav::FT = grav(m.param_set)
    z = altitude(m, aux)
    a_min = turbconv.subdomains.a_min
    a_max = turbconv.subdomains.a_max

    EΔ_up = ntuple(N_up) do i
        entr_detr(m, m.turbconv.entr_detr, state, aux, t, ts, env, i)
    end

    E_dyn, Δ_dyn, E_trb = ntuple(i -> map(x -> x[i], EΔ_up), 3)

    l_mix, _, Pr_t = mixing_length(
        m,
        turbconv.mix_len,
        state,
        diffusive,
        aux,
        t,
        Δ_dyn,
        E_trb,
        ts,
        env,
    )
    tke_en = enforce_positivity(en.ρatke) / env.a * ρ_inv
    K_m = m.turbconv.mix_len.c_m * l_mix * sqrt(tke_en)
    K_h = K_m / Pr_t

    #TotalFlux(ϕ) = Eddy_Diffusivity(ϕ) + MassFlux(ϕ)
    e_int = internal_energy(m, state, aux)

    e_kin = vuntuple(N_up) do i
        FT(1 // 2) * (
            (gm.ρu[1] * ρ_inv)^2 +
            (gm.ρu[2] * ρ_inv)^2 +
            (up[i].ρaw / up[i].ρa)^2
        )
    end
    e_tot_up = ntuple(i -> total_energy(e_kin[i], _grav * z, ts.up[i]), N_up)
    ρa_up = vuntuple(N_up) do i
        gm.ρ * enforce_unit_bounds(up[i].ρa * ρ_inv, a_min, a_max)
    end

    massflux_e = sum(
        vuntuple(N_up) do i
            up[i].ρa *
            ρ_inv *
            (gm.ρe * ρ_inv - e_tot_up[i]) *
            (gm.ρu[3] * ρ_inv - up[i].ρaw / ρa_up[i])
        end,
    )
    ρq_tot = m.moisture isa DryModel ? FT(0) : gm.moisture.ρq_tot
    massflux_q_tot = sum(
        vuntuple(N_up) do i
            up[i].ρa *
            ρ_inv *
            (ρq_tot * ρ_inv - up[i].ρaq_tot / up[i].ρa) *
            (gm.ρu[3] * ρ_inv - up[i].ρaw / ρa_up[i])
        end,
    )

    massflux_w = sum(
        vuntuple(N_up) do i
            up[i].ρa *
            ρ_inv *
            (gm.ρu[3] * ρ_inv - up[i].ρaw / up[i].ρa) *
            (gm.ρu[3] * ρ_inv - up[i].ρaw / ρa_up[i])
        end,
    )

    # update grid mean flux_second_order
    ρe_sgs_flux = -gm.ρ * env.a * K_h * en_dif.∇e[3] + massflux_e
    ρq_tot_sgs_flux = -gm.ρ * env.a * K_h * en_dif.∇q_tot[3] + massflux_q_tot
    ρu_sgs_flux = -gm.ρ * env.a * K_m * en_dif.∇w[3] + massflux_w

    # for now the coupling to the dycore is commented out

    # gm_flx.ρe              += SVector{3,FT}(0,0,ρe_sgs_flux)
    # gm_flx.moisture.ρq_tot += SVector{3,FT}(0,0,ρq_tot_sgs_flux)
    # gm_flx.ρu              += SMatrix{3, 3, FT, 9}(
    #     0, 0, 0,
    #     0, 0, 0,
    #     0, 0, ρu_sgs_flux,
    # )

    ẑ = vertical_unit_vector(m, aux)
    # env second moment flux_second_order
    en_flx.ρatke = -gm.ρ * env.a * K_m * en_dif.∇tke[3] * ẑ
    en_flx.ρaθ_liq_cv = -gm.ρ * env.a * K_h * en_dif.∇θ_liq_cv[3] * ẑ
    en_flx.ρaq_tot_cv = -gm.ρ * env.a * K_h * en_dif.∇q_tot_cv[3] * ẑ
    en_flx.ρaθ_liq_q_tot_cv =
        -gm.ρ * env.a * K_h * en_dif.∇θ_liq_q_tot_cv[3] * ẑ
end;

# First order boundary conditions
function turbconv_boundary_state!(
    nf,
    bc::EDMFBottomBC,
    m::AtmosModel{FT},
    state⁺::Vars,
    aux⁺::Vars,
    n⁻,
    state⁻::Vars,
    aux⁻::Vars,
    t,
    state_int::Vars,
    aux_int::Vars,
) where {FT}

    turbconv = m.turbconv
    N_up = n_updrafts(turbconv)
    up = state⁺.turbconv.updraft
    en = state⁺.turbconv.environment
    gm = state⁺
    gm_a = aux⁺

    zLL = altitude(m, aux_int)
    a_up_surf,
    θ_liq_up_surf,
    q_tot_up_surf,
    θ_liq_cv,
    q_tot_cv,
    θ_liq_q_tot_cv,
    tke = subdomain_surface_values(turbconv.surface, turbconv, m, gm, gm_a, zLL)

    @unroll_map(N_up) do i
        up[i].ρaw = FT(0)
        up[i].ρa = a_up_surf[i] * gm.ρ
        up[i].ρaθ_liq = up[i].ρa * θ_liq_up_surf[i]
        up[i].ρaq_tot = up[i].ρa * q_tot_up_surf[i]
    end
    a_en = environment_area(gm, gm_a, N_up)
    en.ρatke = gm.ρ * a_en * tke
    en.ρaθ_liq_cv = gm.ρ * a_en * θ_liq_cv
    en.ρaq_tot_cv = gm.ρ * a_en * q_tot_cv
    en.ρaθ_liq_q_tot_cv = gm.ρ * a_en * θ_liq_q_tot_cv
end;
function turbconv_boundary_state!(
    nf,
    bc::EDMFTopBC,
    m::AtmosModel{FT},
    state⁺::Vars,
    aux⁺::Vars,
    n⁻,
    state⁻::Vars,
    aux⁻::Vars,
    t,
    state_int::Vars,
    aux_int::Vars,
) where {FT}
    nothing
end;

# The boundary conditions for second-order unknowns
function turbconv_normal_boundary_flux_second_order!(
    nf,
    bc::EDMFBottomBC,
    m::AtmosModel{FT},
    fluxᵀn::Vars,
    n⁻,
    state⁻::Vars,
    diff⁻::Vars,
    hyperdiff⁻::Vars,
    aux⁻::Vars,
    state⁺::Vars,
    diff⁺::Vars,
    hyperdiff⁺::Vars,
    aux⁺::Vars,
    t,
    _...,
) where {FT}
    nothing
end;
function turbconv_normal_boundary_flux_second_order!(
    nf,
    bc::EDMFTopBC,
    m::AtmosModel{FT},
    fluxᵀn::Vars,
    n⁻,
    state⁻::Vars,
    diff⁻::Vars,
    hyperdiff⁻::Vars,
    aux⁻::Vars,
    state⁺::Vars,
    diff⁺::Vars,
    hyperdiff⁺::Vars,
    aux⁺::Vars,
    t,
    _...,
) where {FT}
    turbconv = m.turbconv
    N_up = n_updrafts(turbconv)
    up_flx = fluxᵀn.turbconv.updraft
    en_flx = fluxᵀn.turbconv.environment
    # @unroll_map(N_up) do i
    #     up_flx[i].ρaw = -n⁻ * FT(0)
    #     up_flx[i].ρa = -n⁻ * FT(0)
    #     up_flx[i].ρaθ_liq = -n⁻ * FT(0)
    #     up_flx[i].ρaq_tot = -n⁻ * FT(0)
    # end
    # en_flx.∇tke = -n⁻ * FT(0)
    # en_flx.∇e_int_cv = -n⁻ * FT(0)
    # en_flx.∇q_tot_cv = -n⁻ * FT(0)
    # en_flx.∇e_int_q_tot_cv = -n⁻ * FT(0)

end;


#fin. gen. submodules of Z^n and F_p^n (and possibly more)
#
#import Base.show, Base.inv, Nemo.inv, Base.solve, Hecke.solve,
#       Hecke.lift, Hecke.rational_reconstruction, Hecke.elementary_divisors,
#       Hecke.rank, Hecke.det

export det_mc, id, isid, isupper_triangular, norm2, hadamard_bound2,
       hnf, hnf!, echelon_with_trafo

const p = next_prime(2^20)

add_verbose_scope(:HNF)

add_assert_scope(:HNF)
set_assert_level(:HNF, 0)

function show(io::IO, M::ModuleCtxNmod)
  println("Sparse module over $(M.R) of (current) rank $(rows(M.basis)) and $(rows(M.gens))")
end

function show(io::IO, M::ModuleCtx_fmpz)
  println("Sparse module over FlintZZ of (current) rank $(rows(M.bas_gens)) and further $(rows(M.rel_gens))")
  if isdefined(M, :basis_idx)
    println("current index: $(M.basis_idx)")
  end
  if isdefined(M, :essential_elementary_divisors)
    println("current structure: $(M.essential_elementary_divisors)")
  end
end

function add_gen!(M::ModuleCtxNmod, g::SRow{nmod})
  if M.basis.r == M.basis.c
    return false
  end
  h = reduce(M.basis, g)
  if !iszero(h)
    i = 1
    while i<= rows(M.basis) && M.basis.rows[i].pos[1] < h.pos[1]
      i += 1
    end
    @hassert :HNF 2  i > rows(M.basis) || M.basis[i].pos[1] > h.pos[1]
    insert!(M.basis.rows, i, h)
    M.basis.r += 1
    M.basis.nnz += length(h)
    M.basis.c = max(M.basis.c, h.pos[end])
    push!(M.gens, g)
    return true
  end
  return false
end

function add_gen!(M::ModuleCtx_fmpz, g::SRow{fmpz})
  gp = SRow(g, M.Mp.R)
  M.new = true
  if add_gen!(M.Mp, gp)
    push!(M.bas_gens, g)
    return true
  else
    push!(M.rel_gens, g)
  end
  return false
end

function check_index(M::ModuleCtx_fmpz)
  if rows(M.Mp.basis) < cols(M.Mp.basis)
    return fmpz(0)
  end

  if !M.new
    return M.basis_idx
  end

  if isdefined(M, :trafo)  #once we have trafo, we need to keep it!
    module_trafo_assure(M)
    return M.basis_idx
  end

  M.new = false

  if isdefined(M, :basis)
    C = copy(M.basis)
  else
    d = abs(det(M.bas_gens))
    C = M.max_indep
    C.c = M.bas_gens.c
    for ii = M.bas_gens
      h = reduce(C, ii, 2*d) #to avoid problems with diag being 1...1 d
      @hassert :HNF 2  !iszero(h)
      i = 1
      while i<= rows(C) && C.rows[i].pos[1] < h.pos[1]
        i += 1
      end
      @hassert :HNF 2  i > rows(C) || C[i].pos[1] > h.pos[1]
      insert!(C.rows, i, h)
      C.r += 1
      C.nnz += length(h)
      C.c = max(C.c, h.pos[end])
    end
    M.max_indep = copy(C)
  end

  for i=length(M.rel_reps_p)+1:length(M.rel_gens)
    reduce(C, M.rel_gens[i])
    push!(M.rel_reps_p, solve_ut(M.Mp.basis, SRow(M.rel_gens[i], M.Mp.R)))
  end

#=
  for l=1:5
    mis = find(i->C[i,i] != 1, 1:rows(C))
    if length(mis) == 0
      break
    end
#    println("mis: $mis")
    for i = mis
      if C[i,i] == 1
#        println("miracle for $i")
        continue
      end
      r = find(x->i in M.rel_reps_p[x].pos, 1:length(M.rel_reps_p))
#      println("found $(length(r)) rows")
      if length(r) == 0
        break
      end
      g = M.rel_gens[rand(r)]
      for j=1:min(5, div(length(r), 2))
        g += rand(-10:10)*M.rel_gens[rand(r)]
      end
      reduce(C, g)
      if C[i,i] == 1
#        println("bingo for i=$i")
      end
    end
  end
=#
  M.basis = C
  M.basis_idx = prod([C[i,i] for i=1:rows(C)])

  return M.basis_idx
end

function elementary_divisors(M::ModuleCtx_fmpz)
  if !isdefined(M, :basis)
    i = check_index(M)
    if i == 0
      return fmpz[]
    end
  end
  C = M.basis
  f = find(i -> C[i,i] != 1, 1:rows(C))
  if length(f) == 0
    M.essential_elementary_divisors = fmpz[]
    return M.essential_elementary_divisors
  end
  e = minimum(f)
  m = fmpz_mat(sub(C, e:rows(C), e:cols(C)))
  s = snf(m)
  M.essential_elementary_divisors = [s[i,i] for i=1:rows(s)]
  return M.essential_elementary_divisors
end

function missing_pivot(M::ModuleCtx_fmpz)
  C = M.Mp.basis
  return setdiff(IntSet(1:cols(C)), [x.pos[1] for x=C])
end

function non_trivial_pivot(M::ModuleCtx_fmpz)
  h = check_index(M)
  if h == 0
    return missing_pivot(M)
  end
  C = M.basis
  @hassert :HNF 2  C.r == C.c
  return setdiff(IntSet(1:cols(C)), find(i->C[i].values[1] == 1, 1:C.c))
end

function rank(M::ModuleCtx_fmpz)
  return M.bas_gens.r
end

function rank(M::ModuleCtxNmod)
  return M.basis.r
end

function module_trafo_assure(M::ModuleCtx_fmpz)
  if !M.new && isdefined(M, :trafo)
    return
  end
  z = vcat(M.bas_gens, M.rel_gens)
  h, t = hnf_kannan_bachem(z, Val{true})
  M.trafo = t
  M.basis = h
  M.basis_idx = det(h) # h is upp_triangular, hence det is trivial
  M.new = false
  nothing
end


include(joinpath(@__DIR__, "..", "src/dynamics", "hopper_2D", "visuals.jl"))
T = Float64
vis = Visualizer()
open(vis)

# simulation model
include(joinpath(module_dir(), "src/dynamics/hopper_2D/model.jl"))

model_sim = Hopper2D(Dimensions(nq, nu, nw, nc, nb),
			   1.2 * mb, ml, 1.2 * Jb, Jl,
			   μ_world, μ_joint, g,
			   BaseMethods(), DynamicsMethods(), ContactMethods(),
			   ResidualMethods(), ResidualMethods(),
			   SparseStructure(spzeros(0, 0), spzeros(0, 0)),
			   SVector{4}(zeros(4)),
			   environment_2D_flat())

expr_base = generate_base_expressions(model, M_analytical=true)
instantiate_base!(model_sim.base, expr_base)
expr_dyn = generate_dynamics_expressions(model_sim)
instantiate_dynamics!(model_sim.dyn, expr_dyn)
expr_res, rz_sp, rθ_sp = generate_residual_expressions(model_sim)
instantiate_residual!(model_sim.res, expr_res)
model_sim.spa.rz_sp = rz_sp
model_sim.spa.rθ_sp = rθ_sp

# get hopper model
model = get_model("hopper_2D")
nq = model.dim.q
nu = model.dim.u
nc = model.dim.c
nb = model.dim.b
nd = nq + nc + nb
nr = nq + nu + nc + nb + nd

# get trajectory
# ref_traj = get_trajectory("hopper_2D", "gait_in_place", load_type=:joint_traj)
ref_traj = get_trajectory("hopper_2D", "gait_forward", load_type=:joint_traj)
ref_traj_copy = deepcopy(ref_traj)

# time
H = ref_traj.H
h = ref_traj.h
N_sample = 5
H_mpc = 10
h_sim = h / N_sample
H_sim = 5000

# barrier parameter
κ_mpc = 1.0e-4

obj = TrackingObjective(H_mpc, model.dim,
    q = [Diagonal(1.0e-1 * [1,3,1,3])   for t = 1:H_mpc],
    u = [Diagonal(1.0e-0 * [1e-2, 1e0]) for t = 1:H_mpc],
    γ = [Diagonal(1.0e-100 * ones(model.dim.c)) for t = 1:H_mpc],
    b = [Diagonal(1.0e-100 * ones(model.dim.b)) for t = 1:H_mpc])

p = linearized_mpc_policy(ref_traj, model, obj,
    H_mpc = H_mpc,
    N_sample = N_sample,
    κ_mpc = κ_mpc,
    n_opts = NewtonOptions(
        r_tol = 3e-4,
        max_iter = 5),
    mpc_opts = LinearizedMPCOptions(
        altitude_update = true,
        altitude_impact_threshold = 0.5,
        altitude_verbose = true,
        )
    )

q1_ref = copy(ref_traj.q[2])
q0_ref = copy(ref_traj.q[1])
q1_sim = SVector{model.dim.q}(q1_ref)
q0_sim = SVector{model.dim.q}(copy(q1_sim - (q1_ref - q0_ref) / N_sample))
@assert norm((q1_sim - q0_sim) / h_sim - (q1_ref - q0_ref) / h) < 1.0e-8

sim = ContactImplicitMPC.simulator(model_sim, q0_sim, q1_sim, h_sim, H_sim,
    p = p,
	uL = [-0.25, -0.75], uU = [0.25, 0.75],
    ip_opts = ContactImplicitMPC.InteriorPointOptions(
        r_tol = 1.0e-8,
        κ_init = 1.0e-8,
        κ_tol = 2.0e-8),
    sim_opts = ContactImplicitMPC.SimulatorOptions(warmstart = true))

@time status = ContactImplicitMPC.simulate!(sim)

(sim.traj.q[end][1] - sim.traj.q[1][1]) / (h_sim * H_sim)
anim = visualize_robot!(vis, sim.model, sim.traj, sample=20)

plot(hcat(sim.traj.u...)[1:2, 1:10:end]')


abstract type Unit end

mutable struct Goblin <: Unit
    hp::Int64
    ap::Int64
end

mutable struct Elf <: Unit
    hp::Int64
    ap::Int64
end

mutable struct Empty <: Unit end
mutable struct Wall <: Unit end

mutable struct Node{T}
    loc::Tuple{Int64,Int64}
    data::T
    neighbors::Dict{String,Node{T}}
end

mutable struct Graph{T}
    data::Dict{Tuple{Int64,Int64},Node{T}}
    width::Int64
    height::Int64
end

function make_grid()
    f = open("inputs/15_test.txt", "r")
    lines = readlines(f)
    ysize = length(lines)
    xsize = maximum(length.(lines))
    d::Dict{Tuple{Int64,Int64},Node{Unit}} = Dict()
    graph = Graph(d,xsize,ysize)
    for row in 1:ysize
        for col in 1:xsize
            neighbors::Dict{String,Node{Unit}} = Dict()
            newnode = Node((col,row),Empty(),neighbors)
            graph.data[(col,row)] = newnode
            if row > 1
                # connect above
                nodeabove = graph.data[(col,row-1)]
                nodeabove.neighbors["south"] = newnode
                newnode.neighbors["north"] = nodeabove
            end
            if col > 1
                # connect left
                nodeleft = graph.data[(col-1,row)]
                nodeleft.neighbors["east"] = newnode
                newnode.neighbors["west"] = nodeleft
            end
        end
    end
    for y in 1:length(lines)
        for x in 1:length(lines[y])
            c = lines[y][x]
            if c == '#'
                # remove edges
                node = graph.data[(x,y)]
                try
                    n = node.neighbors["north"]
                    delete!(n.neighbors, "south")
                    delete!(node.neighbors, "north")
                catch
                end
                try
                    n = node.neighbors["south"]
                    delete!(n.neighbors, "north")
                    delete!(node.neighbors, "south")
                catch
                end
                try
                    n = node.neighbors["west"]
                    delete!(n.neighbors, "east")
                    delete!(node.neighbors, "west")
                catch
                end
                try
                    n = node.neighbors["east"]
                    delete!(n.neighbors, "west")
                    delete!(node.neighbors, "east")
                catch
                end
                delete!(graph.data, (x,y))
            elseif c == 'G'
                goblin = Goblin(200,3)
                graph.data[(x,y)].data = goblin
            elseif c == 'E'
                elf = Elf(200,3)
                graph.data[(x,y)].data = elf
            end
        end
    end

    return graph
end

function path(currentloc::Tuple{Int64,Int64},visited::Set{Tuple{Int64,Int64}},target::Tuple{Int64,Int64},grid::Graph{Unit})::Array{Tuple{Int64,Int64}}
    options = []
    try
        north = grid.data[currentloc].neighbors["north"]
        if north.loc == target
            return [north.loc]
        elseif !(in(north.loc, visited))
            v2 = copy(visited)
            push!(v2,currentloc)
            push!(vcat([currentloc], path(north.loc, v2, target, grid)), options)
        end
    catch
    end
    try
        west = grid.data[currentloc].neighbors["west"]
        if west.loc == target
            return [west.loc]
        elseif !(in(west.loc, visited))
            v2 = copy(visited)
            push!(v2,currentloc)
            push!(vcat([currentloc], path(west.loc, v2, target, grid)), options)
        end
    catch
    end
    try
        east = grid.data[currentloc].neighbors["east"]
        if east.loc == target
            return [east.loc]
        elseif !(in(east.loc, visited))
            v2 = copy(visited)
            push!(v2,currentloc)
            push!(vcat([currentloc], path(east.loc, v2, target, grid)), options)
        end
    catch
    end
    try
        south = grid.data[currentloc].neighbors["south"]
        if south.loc == target
            return [south.loc]
        elseif !(in(south.loc, visited))
            v2 = copy(visited)
            push!(v2,currentloc)
            push!(vcat([currentloc], path(south.loc, v2, target, grid)), options)
        end
    catch
    end
    shortest = sort(map(x->(length(x),x), options), by=x->x[1])[1]
    return shortest
end

function _empty(x)
    if x == (0,0)
        return true
    else
        return false
    end
end

function update_grid(direction::String,coords::Tuple{Int64,Int64},paths::Array{Tuple{Int64,Int64},2},grid::Graph, target::Tuple{Int64,Int64})
    try
        node = grid.data[coords].neighbors[direction]
        if paths[node.loc...] == (0,0) && (node.loc == target || typeof(grid.data[node.loc...].data) == Empty)
            paths[node.loc...] = coords
            return true
        end
        return false
    catch
        return false
    end
end

function _path(target::Tuple{Int64,Int64},start::Tuple{Int64,Int64},grid::Graph{Unit})
    paths::Array{Tuple{Int64,Int64},2} = map(x->(0,0),zeros(grid.height,grid.width))
    paths[start...] = (-1,-1)
    active = [start]
    while length(active) > 0
        n = active[1]
        deleteat!(active,1)
        for direction in ["west", "north", "south", "east"]
            result = update_grid(direction, n, paths, grid, target)
            if result
                push!(active, grid.data[n].neighbors[direction].loc)
            end
        end
    end
    if paths[target...] == (0,0)
        return []
    end
    retval = []
    push!(retval, target)
    while retval[end] != (-1,-1)
        push!(retval, paths[retval[end]...])
    end
    return retval[2:end-1]
end

function nearest_target(x::Tuple{Int64,Int64},grid::Graph)
    target = Elf
    if typeof(grid.data[x].data) == Elf
        target = Goblin
    end
    println(target)
    targets = []
    for pair in grid.data
        if typeof(pair[2].data) == target
            push!(targets, pair[1])
        end
    end
    paths = [_path(x,y,grid) for y in targets]
    distances = [length(x) for x in paths]
    if length(distances) == 0
        return x
    end
    attack_positions = []
    if minimum(distances) == 1
        # skip move
        return x
    else
        # find closest attack position
        println("closing in")
        for target in targets
            push!(attack_positions, (target[1]+1,target[2]))
            push!(attack_positions, (target[1],target[2]+1))
            push!(attack_positions, (target[1],target[2]-1))
            push!(attack_positions, (target[1]-1,target[2]))
        end
        attack_positions = filter(y->length(_path(x,y,grid)) > 0, attack_positions)
        println(attack_positions)
    end
    closest_distance = minimum(map(y->length(_path(x,y,grid)), attack_positions))
    println(closest_distance)
    positions = filter(y->length(_path(x,y,grid))==closest_distance, attack_positions)
    println(positions)
    if length(positions) == 0
        return x
    end
    position_to_go_to = positions[end]
    return position_to_go_to
end

function move(x::Tuple{Int64,Int64}, grid::Graph)
    println("moving " * string(x))
    nt = nearest_target(x,grid)
    draw_grid(grid)
    if x != nt
        println("Actually moving")
        _p = _path(x,nt,grid)
        println(_p)
        p = _p[1]
        u = grid.data[x].data
        grid.data[p].data = u
        grid.data[x].data = Empty()
    end
    println("...")
    draw_grid(grid)
end

function _get_units(grid::Graph)
    units = []
    for row in 1:grid.height
        for col in 1:grid.width
            try
                n = grid.data[(row,col)]
                if typeof(n.data) == Elf || typeof(n.data) == Goblin
                    push!(units, (row,col))
                end
            catch
            end
        end
    end
    return units
end

function move_units(grid::Graph)
    types_left = Set()
    for unit in _get_units(grid)
        push!(types_left, typeof(grid.data[unit].data))
    end
    if length(collect(types_left)) > 1
        for unit in _get_units(grid)
            move(unit, grid)
        end
    end
end

function draw_grid(grid::Graph)
    retval = ""
    for row in 1:grid.height
        for col in 1:grid.width
            try
                if typeof(grid.data[(row,col)].data) == Wall
                    retval *= '#'
                end
                if typeof(grid.data[(row,col)].data) == Empty
                    retval *= '.'
                end
                if typeof(grid.data[(row,col)].data) == Goblin
                    retval *= 'G'
                end
                if typeof(grid.data[(row,col)].data) == Elf
                    retval *= 'E'
                end
            catch
                retval *= '#'
            end
        end
        retval *= '\n'
    end
    print(retval)
end

grid = make_grid();
move_units(grid)
draw_grid(grid)


#=
Simulation:
- Julia version:
- Author: inaba
- Date: 2021-11-20
=#
module Simulation
using Dates: format, now
using LightGraphs
using Random
using Statistics: mean

mutable struct Agent
    id::Int
    is_cooperator::Bool
    is_punisher::Bool
    next_is_cooperator::Bool
    next_is_punisher::Bool
    payoff::Float64

    Agent(id::Int) = new(id, rand([true, false]), rand([true, false]), false, false, 0.0)
end

struct Model
    agents::Vector{Agent}
    graph::SimpleGraph{Int64}
    hop::Int  # neighbor parameter
    neighbours::Vector{Vector{Int64}}  # neighbors of every nodes

    g::Float64  # game interaction locality parameter
    p::Float64  # punishment locality parameter
    a::Float64  # adaptation (learning) locality parameter

    b::Float64  # benefit of a token
    c::Float64  # game cost
    f::Float64  # fine
    s::Float64  # sanction for punishment

    μ::Float64  # mutation rate

    Model(agents::Vector{Agent}, graph::SimpleGraph; g::Float64, p::Float64, a::Float64) = new(
        agents, graph, 2,
        [[_n for _n in neighborhood(graph, i, 2) if _n != i] for i in vertices(graph)],
        g, p, a,
        2.0, 1.0, 6.0, 3.0, 0.01
    )
end

cooperator_rate(model::Model)::Float64 = mean([agent.is_cooperator for agent in model.agents])
punisher_rate(model::Model)::Float64 = mean([agent.is_punisher for agent in model.agents])

function calc_payoffs!(model::Model)
    n_size = length(model.agents)
    global_cooperator_rate = cooperator_rate(model)
    global_punisher_rate = punisher_rate(model)

    for agent in model.agents
        local_cooperator_rate = mean([model.agents[neighbor].is_cooperator for neighbor in model.neighbours[agent.id]])
        local_punisher_rate = mean([model.agents[neighbor].is_punisher for neighbor in model.neighbours[agent.id]])

        _global_cooperator_rate = (global_cooperator_rate * n_size - Int(agent.is_cooperator)) / (n_size - 1)
        _global_punisher_rate = (global_punisher_rate * n_size - Int(agent.is_punisher)) / (n_size - 1)

        # ゲームの中で協力者が占める割合
        _cooperator_rate = local_cooperator_rate * (1. - model.g) + _global_cooperator_rate * model.g

        # benefit (ゲームの中で協力者が占める割合 (自分を含む))
        _b = (_cooperator_rate * model.hop * 2 + Int(agent.is_cooperator)) / (model.hop * 2 + 1)
        # game cost
        _c = agent.is_cooperator ? 1 : 0
        # sanction (罰を課すコスト)
        _s = agent.is_punisher ? (1 - local_cooperator_rate) * (1 - model.p) + (1 - _global_cooperator_rate) * model.p : 0.0
        # Fine (罰金)
        _f = agent.is_cooperator ? 0.0 : local_punisher_rate * (1 - model.p) + _global_punisher_rate * model.p

        agent.payoff = _b * model.b - _c * model.c - _s * model.s - _f * model.f
    end
end

function set_next_strategies!(model::Model)
    for agent in model.agents
        if rand() < model.μ
            agent.next_is_cooperator = rand([true, false])
            agent.next_is_punisher = rand([true, false])
        else
            y = if rand() < 0.5 # a
                rand([_a.id for _a in model.agents if _a.id != agent.id])
            else
                rand(model.neighbours[agent.id])
            end

            if model.agents[y].payoff > agent.payoff
                agent.next_is_cooperator = model.agents[y].is_cooperator
                agent.next_is_punisher = model.agents[y].is_punisher
            else
                agent.next_is_cooperator = agent.is_cooperator
                agent.next_is_punisher = agent.is_punisher
            end
        end
    end
end

function update_strategies!(model::Model)
    for agent in model.agents
        agent.is_cooperator = agent.next_is_cooperator
        agent.is_punisher = agent.next_is_punisher
    end
end

function run(;
    n_size::Int = 100,
    generation::Int = 100,
    trial::Int = 100,
    gs::Vector = [0.0, 0.25, 0.5, 0.75, 1.0],
    ps::Vector = [0.0, 0.25, 0.5, 0.75, 1.0],
    as::Vector = [0.0, 0.25, 0.5, 0.75, 1.0],
    file_name = "out/jl$(format(now(), "yyyymmdd_HHMMSS")).csv")

    for g = gs, p = ps, a = as
        cooperator_rate_vec = []
        punisher_rate_vec = []
        for t in 1:trial
            agents = [Agent(id) for id in 1:n_size]
            graph = cycle_graph(n_size)
            model = Model(agents, graph; g, p, a)

            for step = 1:generation
                calc_payoffs!(model)
                set_next_strategies!(model)
                update_strategies!(model)
            end

            push!(cooperator_rate_vec, cooperator_rate(model))
            push!(punisher_rate_vec, punisher_rate(model))
        end

        csv = "$(g),$(p),$(a),$(mean(cooperator_rate_vec)),$(mean(punisher_rate_vec))"
        println(csv)
        open(file_name, "a") do out
            println(out, csv)
        end
    end
end
end

if abspath(PROGRAM_FILE) == @__FILE__
    using .Simulation
    Simulation.run()
    # Simulation.run(gs = [0.0], ps = [1.0], as = [0.0])
end


# Optimization interface

"""
    optimize(f[, F], indiv, algo[, opts])
    optimize(f[, F], constr, algo[, opts])
    optimize(f[, F], constr, indiv, algo[, opts])
    optimize(f[, F], constr, algo, poplt[, opts])

Perform optimization of the function `f` using the alorithm `algo` with the population, composed of
the initial population `poplt`, or individuals similar to the original individual `indiv`,
or generated from the constraints `constr`, with the options `opts`.
- For multi-objective optimization, the objective value `F` *must* be provided.
"""
optimize(f, individual, method::M,
         opts::Options = Options(;default_options(method)...)) where {M<:AbstractOptimizer} =
    optimize(f, NoConstraints(), individual, method, opts)
optimize(f, F::AbstractVector, individual, method::M,
         opts::Options = Options(;default_options(method)...)) where {M<:AbstractOptimizer} =
    optimize(f, F, NoConstraints(), individual, method, opts)
optimize(f, bounds::ConstraintBounds, method::M,
         opts::Options = Options(;default_options(method)...)) where {M<:AbstractOptimizer} =
    optimize(f, BoxConstraints(bounds), method, opts)
optimize(f, F, bounds::ConstraintBounds, method::M,
         opts::Options = Options(;default_options(method)...)) where {M<:AbstractOptimizer} =
    optimize(f, F, BoxConstraints(bounds), method, opts)
function optimize(f, constraints::C, method::M,
                  opts::Options = Options(;default_options(method)...)
                  ) where {M<:AbstractOptimizer, C<:AbstractConstraints}
    population = initial_population(method, bounds(constraints))
    optimize(f, constraints, method, population, opts)
end
function optimize(f, F, constraints::C, method::M,
                  opts::Options = Options(;default_options(method)...)
                  ) where {M<:AbstractOptimizer, C<:AbstractConstraints}
    population = initial_population(method, bounds(constraints))
    optimize(f, F, constraints, method, population, opts)
end
function optimize(f, constraints::C, individual, method::M,
                  opts::Options = Options(;default_options(method)...)
                  ) where {M<:AbstractOptimizer, C<:AbstractConstraints}
    population = initial_population(method, individual)
    optimize(f, constraints, method, population, opts)
end
function optimize(f, F, constraints::C, individual, method::M,
                  opts::Options = Options(;default_options(method)...)
                  ) where {M<:AbstractOptimizer, C<:AbstractConstraints}
    population = initial_population(method, individual)
    optimize(f, F, constraints, method, population, opts)
end
function optimize(f, constraints::C, method::M, population,
                  opts::Options = Options(;default_options(method)...)
                 ) where {M<:AbstractOptimizer, C<:AbstractConstraints}
    @assert length(population) > 0 "Population is empty"
    objfun = EvolutionaryObjective(f, first(population); eval=opts.parallelization)
    optimize(objfun, constraints, method, population, opts)
end
function optimize(f, F, constraints::C, method::M, population,
                  opts::Options = Options(;default_options(method)...)
                 ) where {M<:AbstractOptimizer, C<:AbstractConstraints}
    @assert length(population) > 0 "Population is empty"
    objfun = EvolutionaryObjective(f, first(population), F; eval=opts.parallelization)
    optimize(objfun, constraints, method, population, opts)
end

function optimize(objfun::D, constraints::C, method::M, population::AbstractArray,
                  options::Options = Options(;default_options(method)...),
                  state = initial_state(method, options, objfun, population)
                 ) where {D<:AbstractObjective, C<:AbstractConstraints, M<:AbstractOptimizer}
    # setup trace
    tr = OptimizationTrace{typeof(value(objfun)), typeof(method)}()
    tracing = options.store_trace || options.show_trace || options.callback !== nothing

    # prepare iteration counter (used to make "initial state" trace entry)
    iteration = 0
    t0 = time()
    stopped, stopped_by_callback = false, false
    converged, counter_tol = false, 0 # tolerance convergence
    is_moo = ismultiobjective(objfun)

    options.show_trace && print_header(method)
    trace!(tr, iteration, objfun, state, population, method, options, time()-t0)

    _time = time()
    while !converged && !stopped && iteration < options.iterations
        iteration += 1

        # perform state update
        update_state!(objfun, constraints, state, population, method, options, iteration) && break

        # evaluate convergence
        converged = assess_convergence(objfun, state, method, options)

        # update the function value
        !is_moo && value!(objfun, minimizer(state))

        # check convergence persistence
        counter_tol = converged ? counter_tol+1 : 0
        converged = converged && (counter_tol > options.successive_f_tol)

        # custom state-based termination
        converged = converged || terminate(state)

        if tracing
            # update trace; callbacks can stop routine early by returning true
            stopped_by_callback = trace!(tr, iteration, objfun, state, population, method, options, time()-t0)
        end

        _time = time()
        stopped_by_time_limit = _time-t0 > options.time_limit

        if stopped_by_callback || stopped_by_time_limit
            stopped = true
        end
    end

    after_while!(objfun, state, method, options)

    return EvolutionaryOptimizationResults(
        method,
        minimizer(state),
        is_moo ? NaN : value(objfun),
        iteration,
        iteration == options.iterations,
        converged,
        options.abstol,
        options.reltol,
        abschange(objfun, state),
        relchange(objfun, state),
        tr,
        f_calls(objfun),
        options.time_limit,
        _time-t0,
    )
end

after_while!(objfun, state, method, options) = nothing