text:
using CSV
using LinearAlgebra
using LabelledGraphs

function _energy(config::Dict, couplings::Dict, cedges::Dict, n::Int)
    eng = zeros(1,n)
    for (i, j) ∈ keys(cedges)
        for (k, l) ∈ values(cedges[i, j])
            for m ∈ 1:length(config[k])
                s = config[k][m]
                r = config[l][m]
                J = couplings[k, l]
                if k == l
                    eng[m] += dot(s, J)
                else
                    eng[m] += dot(s, J, r)
                end
            end
       end
    end
    eng
end

function _energy(ig::LabelledGraph, config::Array)
    s = size(config, 1)
    eng = zeros(s)
    for i ∈ 1:s
        eng[i] = energy(config[i, :], ig)
    end
    eng
end

@testset "Ising graph cannot be created" begin

    @testset "if input instance contains duplicate edges" begin
        @test_throws ArgumentError ising_graph(
        Dict(
                (1, 1) => 2.0,
                (1, 2) => 0.5,
                (2, 1) => -1.0
            )
        )
    end
end

for (instance, source) ∈ (
    ("$(@__DIR__)/instances/example.txt", "file"),
    (
        Dict(
            (1, 1) => 0.1,
            (2, 2) => 0.5,
            (1, 4) => -2.0,
            (4, 2) => 1.0,
            (1, 2) => -0.3
        ),
         "array"
    )
)
@testset "Ising graph created from $(source)" begin
    expected_num_vertices = 3
    expected_biases = [0.1, 0.5, 0.0]
    expected_couplings = Dict(
        LabelledEdge(1, 2) => -0.3,
        LabelledEdge(1, 4) => -2.0,
        LabelledEdge(2, 4) => 1.0
    )
    expected_J_matrix = [
        [0 -0.3 -2.0];
        [0 0 1.0];
        [0 0 0];
    ]

    ig = ising_graph(instance)

    @testset "contains the same number vertices as original instance" begin
        @test nv(ig) == expected_num_vertices
    end

    @testset "has collection of edges comprising all interactions from instance" begin
        # This test uses the fact that edges iterates in the lex ordering.
        @test collect(edges(ig)) == [LabelledEdge(e...) for e in [(1, 2), (1, 4), (2, 4)]]
    end

    @testset "stores biases as property of vertices" begin
        @test biases(ig) == expected_biases
    end

    @testset "stores couplings both as property of edges and its own property" begin
        @test couplings(ig) == expected_J_matrix
    end

    @testset "has default rank stored for each active vertex" begin
        @test get_prop(ig, :rank) == Dict(1 => 2, 2 => 2, 4 => 2)
    end
end
end

@testset "Ising graph created with additional parameters" begin
    expected_biases = [-0.1, -0.5, 0.0]
    expected_couplings = Dict(
        Edge(1, 2) => 0.3,
        Edge(1, 3) => 2.0,
        Edge(2, 3) => -1.0
    )
    expected_J_matrix = [
        [0 0.3 2.0 ];
        [0 0 -1.0];
        [0 0 0];
    ]

    ig = ising_graph(
        "$(@__DIR__)/instances/example.txt",
        sgn=-1,
        rank_override=Dict(1 => 3, 4 => 4)
    )

    @testset "has rank overriden by rank_override dict" begin
        # TODO: update default value of 2 once original implementation
        # is also updated.
        @test get_prop(ig, :rank) == Dict(1 => 3, 2 => 2, 4 => 4)
    end

    @testset "has coefficients multiplied by given sign" begin
        @test biases(ig) == expected_biases
        @test couplings(ig) == expected_J_matrix
    end
end

@testset "Ising" begin
    L = 4
    N = L^2
    instance = "$(@__DIR__)/instances/$(N)_001.txt"

    ig = ising_graph(instance)

    @test nv(ig) == N

    for i ∈ 1:N
        @test has_vertex(ig, i)
    end

    A = adjacency_matrix(ig)

    B = zeros(Int, N, N)
    for i ∈ 1:N
        nbrs = SpinGlassNetworks.unique_neighbors(ig, i)
        for j ∈ nbrs
            B[i, j] = 1
        end
    end

    @test B + B' == A

    @testset "Naive brute force for +/-1" begin
        k = 2^N

        sp = brute_force(ig, num_states=k)

        s = 5

        @test sp.energies ≈ energy.(sp.states, Ref(ig))

        β = rand(Float64)
        ρ = gibbs_tensor(ig, β)

        @test size(ρ) == Tuple(fill(2, N))

        r = exp.(-β .* sp.energies)
        R = r ./ sum(r)

        @test sum(R) ≈ 1
        @test sum(ρ) ≈ 1

        @test [ ρ[idx.(σ)...] for σ ∈ sp.states ] ≈ R
    end

    @testset "Naive brute force for general spins" begin
        L = 4
        instance = "$(@__DIR__)/instances/$(L)_001.txt"

        ig = ising_graph(instance)

        set_prop!(ig, :rank, [3,2,5,4])
        rank = get_prop(ig, :rank)

        all = prod(rank)
        sp = brute_force(ig, num_states=all)

        β = rand(Float64)
        ρ = exp.(-β .* sp.energies)

        ϱ = ρ ./ sum(ρ)
        ϱ̃ = gibbs_tensor(ig, β)

        @test [ ϱ̃[idx.(σ)...] for σ ∈ sp.states ] ≈ ϱ
    end

    @testset "Reading from Dict" begin
        instance_dict = Dict()
        ising = CSV.File(instance, types=[Int, Int, Float64], header=0, comment = "#")

        for (i, j, v) ∈ ising
            push!(instance_dict, (i, j) => v)
        end

        ig = ising_graph(instance)
        ig_dict = ising_graph(instance_dict)

        @test nv(ig_dict) == nv(ig)

        @test collect(edges(ig)) == collect(edges(ig_dict))
    end

    @testset "Ground state energy for pathological instance " begin
        m = 3
        n = 4
        t = 3

        β = 1

        instance = "$(@__DIR__)/instances/pathological/test_$(m)_$(n)_$(t).txt"

        ising = CSV.File(instance, types=[Int, Int, Float64], header=0, comment = "#")
        ig = ising_graph(instance)

        conf = [
            -1 1 1 -1 -1 -1 1 1 1 -1 1 1 -1 1 -1 1;
            -1 1 1 -1 -1 -1 1 1 1 -1 1 1 -1 1 -1 -1;
            -1 1 1 -1 -1 1 1 1 1 -1 1 1 -1 1 -1 1;
            -1 1 1 -1 -1 1 1 1 1 -1 1 1 -1 1 -1 -1]

        eng = _energy(ig, conf)

        couplings = Dict()
        for (i, j, v) ∈ ising
            push!(couplings, (i, j) => v)
        end

        cedges = Dict()
        push!(cedges, (1, 2) => [(1, 4), (1, 5), (1, 6)])
        push!(cedges, (1, 5) => [(1, 13)])

        push!(cedges, (2, 3) => [(4, 7), (5, 7), (6, 8), (6, 9)])
        push!(cedges, (2, 6) => [(6, 16), (6, 18), (5, 16)])

        push!(cedges, (5, 6) => [(13, 16), (13, 18)])

        push!(cedges, (6, 10) => [(18, 28)])
        push!(cedges, (10, 11) => [(28, 31), (28, 32), (28, 33), (29, 31), (29, 32), (29, 33), (30, 31), (30, 32), (30, 33)])

        push!(cedges, (2, 2) => [(4, 5), (4, 6), (5, 6), (6, 6)])
        push!(cedges, (3, 3) => [(7, 8), (7, 9)])
        push!(cedges, (6, 6) => [(16, 18), (16, 16)])
        push!(cedges, (10, 10) => [(28, 29), (28, 30), (29, 30)])

        config = Dict()
        push!(config, 1 => [-1, -1, -1, -1])
        push!(config, 2 => [0, 0, 0, 0])
        push!(config, 3 => [0, 0, 0, 0])
        push!(config, 4 => [1, 1, 1, 1])
        push!(config, 5 => [1, 1, 1, 1])
        push!(config, 6 => [-1, -1, -1, -1])
        push!(config, 7 => [-1, -1, -1, -1])
        push!(config, 8 => [-1, -1, 1, 1])
        push!(config, 9 => [1, 1, 1, 1])
        push!(config, 10 => [0, 0, 0, 0])
        push!(config, 11 => [0, 0, 0, 0])
        push!(config, 12 => [0, 0, 0, 0])
        push!(config, 13 => [1, 1, 1, 1])
        push!(config, 14 => [0, 0, 0, 0])
        push!(config, 15 => [0, 0, 0, 0])
        push!(config, 16 => [1, 1, 1, 1])
        push!(config, 17 => [0, 0, 0, 0])
        push!(config, 18 => [-1, -1, -1, -1])
        push!(config, 19 => [0, 0, 0, 0])
        push!(config, 20 => [0, 0, 0, 0])
        push!(config, 21 => [0, 0, 0, 0])
        push!(config, 22 => [0, 0, 0, 0])
        push!(config, 23 => [0, 0, 0, 0])
        push!(config, 24 => [0, 0, 0, 0])
        push!(config, 25 => [0, 0, 0, 0])
        push!(config, 26 => [0, 0, 0, 0])
        push!(config, 27 => [0, 0, 0, 0])
        push!(config, 28 => [1, 1, 1, 1])
        push!(config, 29 => [1, 1, 1, 1])
        push!(config, 30 => [-1, -1, -1, -1])
        push!(config, 31 => [1, 1, 1, 1])
        push!(config, 32 => [-1, -1, -1, -1])
        push!(config, 33 => [1,-1, 1, -1])
        push!(config, 34 => [0, 0, 0, 0])
        push!(config, 35 => [0, 0, 0, 0])
        push!(config, 36 => [0, 0, 0, 0])

        num_config = length(config[1])
        exact_energy = _energy(config, couplings, cedges, num_config)

        low_energies = [-16.4, -16.4, -16.4, -16.4, -16.1, -16.1, -16.1, -16.1, -15.9, -15.9, -15.9, -15.9, -15.9, -15.9, -15.6, -15.6, -15.6, -15.6, -15.6, -15.6, -15.4, -15.4]

        for i ∈ 1:num_config
            @test exact_energy[i] == low_energies[i] == eng[i]
        end

    end
end

@testset "Pruning" begin

    @testset "No vertices of degree zero" begin
        instance = Dict(
            (1, 1) => 0.1,
            (2, 2) => 0.5,
            (1, 4) => -2.0,
            (4, 2) => 1.0,
            (1, 2) => -0.3,
        )

        ig = ising_graph(instance)
        ng = prune(ig)
        @test nv(ig) == nv(ng)
    end

    @testset "All vertices of degree zero with no local fields" begin
        instance = Dict(
            (1, 1) => 0.0,
            (2, 2) => 0.0,
        )

        ig = ising_graph(instance)
        ng = prune(ig)
        @test nv(ng) == 0
    end

    @testset "All vertices of degree zero, but with local fields" begin
        instance = Dict(
            (1, 1) => 0.1,
            (2, 2) => 0.5,
        )

        ig = ising_graph(instance)
        ng = prune(ig)
        @test nv(ng) == 2
    end

    @testset "Some vertices of degree zero, but nonzero field" begin
        instance = Dict(
                (1, 1) => 0.1,
                (2, 2) => 0.5,
                (1, 4) => -2.0,
                (4, 2) => 1.0,
                (1, 2) => -0.3,
                (5, 5) => 0.1
            )

        ig = ising_graph(instance)
        ng = prune(ig)

        @test nv(ng) == nv(ig)
        @test vertices(ng) == collect(1:nv(ng))
    end

    @testset "Some vertices of degree zero and zero field" begin
        instance = Dict(
                (1, 1) => 0.1,
                (2, 2) => 0.5,
                (1, 4) => -2.0,
                (4, 2) => 1.0,
                (1, 2) => -0.3,
                (5, 5) => 0.0
            )

        ig = ising_graph(instance)
        ng = prune(ig)

        @test nv(ng) == nv(ig) - 1
        @test vertices(ng) == collect(1:nv(ng))
    end
end


text:
module RadiationWrapper

using JavaCall

const AbsoluteDate = @jimport org.orekit.time.AbsoluteDate
const AbstractForceModel = @jimport org.orekit.forces.AbstractForceModel
const AbstractRadiationForceModel = @jimport org.orekit.forces.radiation.AbstractRadiationForceModel
const Class = @jimport java.lang.Class
const ECOM2 = @jimport org.orekit.forces.radiation.ECOM2
const ExtendedPVCoordinatesProvider = @jimport org.orekit.utils.ExtendedPVCoordinatesProvider
const Field = @jimport org.hipparchus.Field
const FieldAbsoluteDate = @jimport org.orekit.time.FieldAbsoluteDate
const FieldRotation = @jimport org.hipparchus.geometry.euclidean.threed.FieldRotation
const FieldSpacecraftState = @jimport org.orekit.propagation.FieldSpacecraftState
const FieldTimeDerivativesEquations = @jimport org.orekit.propagation.numerical.FieldTimeDerivativesEquations
const FieldVector3D = @jimport org.hipparchus.geometry.euclidean.threed.FieldVector3D
const ForceModel = @jimport org.orekit.forces.ForceModel
const Frame = @jimport org.orekit.frames.Frame
const IsotropicRadiationCNES95Convention = @jimport org.orekit.forces.radiation.IsotropicRadiationCNES95Convention
const IsotropicRadiationClassicalConvention = @jimport org.orekit.forces.radiation.IsotropicRadiationClassicalConvention
const IsotropicRadiationSingleCoefficient = @jimport org.orekit.forces.radiation.IsotropicRadiationSingleCoefficient
const JString = @jimport java.lang.String
const Object = @jimport java.lang.Object
const ParameterDriver = @jimport org.orekit.utils.ParameterDriver
const RadiationSensitive = @jimport org.orekit.forces.radiation.RadiationSensitive
const RealFieldElement = @jimport org.hipparchus.RealFieldElement
const Rotation = @jimport org.hipparchus.geometry.euclidean.threed.Rotation
const SolarRadiationPressure = @jimport org.orekit.forces.radiation.SolarRadiationPressure
const SpacecraftState = @jimport org.orekit.propagation.SpacecraftState
const Stream = @jimport java.util.stream.Stream
const TimeDerivativesEquations = @jimport org.orekit.propagation.numerical.TimeDerivativesEquations
const Vector3D = @jimport org.hipparchus.geometry.euclidean.threed.Vector3D

export AbstractRadiationForceModel
export ECOM2
export IsotropicRadiationCNES95Convention
export IsotropicRadiationClassicalConvention
export IsotropicRadiationSingleCoefficient
export RadiationSensitive
export SolarRadiationPressure
export acceleration
export add_contribution
export depends_on_position_only
export equals
export get_class
export get_events_detectors
export get_field_events_detectors
export get_lighting_ratio
export get_parameter_driver
export get_parameters
export get_parameters_drivers
export get_radiation_parameters_drivers
export hash_code
export init
export is_supported
export notify
export notify_all
export radiation_pressure_acceleration
export to_string
export wait

include("abstract_radiation_force_model.jl")
include("eco_m2.jl")
include("isotropic_radiation_classical_convention.jl")
include("isotropic_radiation_cne_s95_convention.jl")
include("isotropic_radiation_single_coefficient.jl")
include("radiation_sensitive.jl")
include("solar_radiation_pressure.jl")

end

ext: jl


text:
@doc raw"""
Penning trap with uniform magnetic field in (x,y,z) coordinates.
Based on Yanyan Shi, Yajuan Sun, Yulei Wang, Jian Liu, Study of adaptive symplectic methods for
      simulating charged particle dynamics, Journal of Computational Dynamics 6, 429-448, 2019.

The covariant components of the vector potential are given by
```math
A (x,y,z) = B_0 \, ( 0, x, 0)^T ,
```
resulting in the magnetic field with covariant components
```math
B (x,y,z) = B_0 \, ( 0, 0, 1)^T ,
```
and the electrostatic potential given by
```math
\varphi (x,y,z) = E_0 \, ( x^2 / 2 + y^2 / 2 - z^2) ,
```
resulting in the electric field with covariant components
```math
E (x,y,z) = E_0 \, ( x, y, - 2 z)^T .
```

Parameters:
* `B₀`: B-field strength
* `E₀`: E-field strength
"""
module PenningTrapUniform

    using RecipesBase

    import ..ElectromagneticFields
    import ..ElectromagneticFields: CartesianEquilibrium, code
    import ..AnalyticCartesianField: X, Y, Z

    export  PenningTrapUniformEquilibrium

    const DEFAULT_B₀ = 100.0
    const DEFAULT_E₀ = 10.0

    struct PenningTrapUniformEquilibrium{T <: Number} <: CartesianEquilibrium
        name::String
        B₀::T
        E₀::T

        function PenningTrapUniformEquilibrium{T}(B₀::T, E₀::T) where T <: Number
            new("PenningTrapUniformEquilibrium", B₀, E₀)
        end
    end

    PenningTrapUniformEquilibrium(B₀::T=DEFAULT_B₀, E₀::T=DEFAULT_E₀) where T <: Number = PenningTrapUniformEquilibrium{T}(B₀, E₀)

    function init(B₀=DEFAULT_B₀, E₀=DEFAULT_E₀)
        PenningTrapUniformEquilibrium(B₀, E₀)
    end

    macro code(B₀=DEFAULT_B₀, E₀=DEFAULT_E₀)
        code(init(B₀, E₀); escape=true)
    end

    function Base.show(io::IO, equ::PenningTrapUniformEquilibrium)
        print(io, "Penning trap with uniform magnetic field in (x,y,z) coordinates with\n")
        print(io, "  B₀ = ", equ.B₀, "\n")
        print(io, "  E₀ = ", equ.E₀)
    end

    ElectromagneticFields.A₁(x::AbstractVector, equ::PenningTrapUniformEquilibrium) = zero(eltype(x))
    ElectromagneticFields.A₂(x::AbstractVector, equ::PenningTrapUniformEquilibrium) = equ.B₀ * X(x,equ)
    ElectromagneticFields.A₃(x::AbstractVector, equ::PenningTrapUniformEquilibrium) = zero(eltype(x))
    ElectromagneticFields.φ(x::AbstractVector, equ::PenningTrapUniformEquilibrium) = - equ.E₀ * (X(x,equ)^2 / 2 + Y(x,equ)^2 / 2 - Z(x,equ)^2)

    ElectromagneticFields.get_functions(::PenningTrapUniformEquilibrium) = (X=X, Y=Y, Z=Z)

end


text: Dict{String, Any}("c" => Dict{String, Any}("a" => Dict{String, Any}("b" => Dict{String, Any}("value" => "1", "type" => "integer"))), "arr" => Any[Dict{String, Any}("T" => Dict{String, Any}("a" => Dict{String, Any}("b" => Dict{String, Any}("value" => "1", "type" => "integer"))), "t" => Dict{String, Any}("a" => Dict{String, Any}("b" => Dict{String, Any}("value" => "1", "type" => "integer")))), Dict{String, Any}("T" => Dict{String, Any}("a" => Dict{String, Any}("b" => Dict{String, Any}("value" => "2", "type" => "integer"))), "t" => Dict{String, Any}("a" => Dict{String, Any}("b" => Dict{String, Any}("value" => "2", "type" => "integer"))))], "e" => Dict{String, Any}("a" => Dict{String, Any}("b" => Dict{String, Any}("value" => "1", "type" => "integer"))), "inline" => Dict{String, Any}("a" => Dict{String, Any}("b" => Dict{String, Any}("value" => "42", "type" => "integer"))), "many" => Dict{String, Any}("dots" => Dict{String, Any}("here" => Dict{String, Any}("dot" => Dict{String, Any}("dot" => Dict{String, Any}("dot" => Dict{String, Any}("a" => Dict{String, Any}("b" => Dict{String, Any}("c" => Dict{String, Any}("value" => "1", "type" => "integer"), "d" => Dict{String, Any}("value" => "2", "type" => "integer"))))))))), "b" => Dict{String, Any}("a" => Dict{String, Any}("b" => Dict{String, Any}("value" => "1", "type" => "integer"))), "tbl" => Dict{String, Any}("x" => Dict{String, Any}("a" => Dict{String, Any}("b" => Dict{String, Any}("c" => Dict{String, Any}("d" => Dict{String, Any}("e" => Dict{String, Any}("value" => "1", "type" => "integer")))))), "a" => Dict{String, Any}("b" => Dict{String, Any}("c" => Dict{String, Any}("d" => Dict{String, Any}("e" => Dict{String, Any}("value" => "1", "type" => "integer")))))), "a" => Dict{String, Any}("a" => Dict{String, Any}("b" => Dict{String, Any}("value" => "1", "type" => "integer"))), "d" => Dict{String, Any}("a" => Dict{String, Any}("b" => Dict{String, Any}("value" => "1", "type" => "integer"))))


text:
function stepOcean_Flow!(
    ocn  :: Ocean;
    cfgs...
)

    adv_scheme = cfgs[:adv_scheme]
    do_convadjust = cfgs[:do_convadjust]
    Δt         = cfgs[:Δt]

    if adv_scheme == :static
        return
    end

    reset!(ocn.wksp)

    #old_intT = sum(ocn.ASUM.filter_T * ocn.ASUM.T_Δvol_T * view(ocn.Ts, :))
    # Determine the temperature / salinity of FLDO layer
    @loop_hor ocn i j let

        ocn.ΔT[i, j] = mixFLDO!(
            qs   = ocn.cols.Ts[i, j],
            zs   = ocn.cols.zs[i, j],
            hs   = ocn.cols.hs[i, j],
            q_ML = ocn.T_ML[i, j],
            FLDO = ocn.FLDO[i, j],
            FLDO_ratio_top = ocn.FLDO_ratio_top[i, j],
            FLDO_ratio_bot = ocn.FLDO_ratio_bot[i, j],
        )

        ocn.ΔS[i, j] = mixFLDO!(
            qs   = ocn.cols.Ss[i, j],
            zs   = ocn.cols.zs[i, j],
            hs   = ocn.cols.hs[i, j],
            q_ML = ocn.S_ML[i, j],
            FLDO = ocn.FLDO[i, j],
            FLDO_ratio_top = ocn.FLDO_ratio_top[i, j],
            FLDO_ratio_bot = ocn.FLDO_ratio_bot[i, j],
        )

    end

    # Pseudo code
    # 1. assign velocity field
    # 2. calculate temperature & salinity flux
    # 3. calculate temperature & salinity flux divergence
    # Gov eqn adv + diff: ∂T/∂t = - 1 / (ρ H1) ( ∇⋅(M1 T1) - (∇⋅M1) Tmid )

#    println("##### QUICK NO SPEEDUP")
#=
    calDiffAdv_QUICK!(
        ocn,
        qs          = ocn.Ts,
        FLUX_bot      = ocn.TFLUX_bot,
        dΔqdt       = ocn.dΔTdt,
        FLUX_CONV   = ocn.TFLUX_CONV,
        FLUX_CONV_h = ocn.TFLUX_CONV_h,
        FLUX_DEN_x  = ocn.TFLUX_DEN_x,
        FLUX_DEN_y  = ocn.TFLUX_DEN_y,
        FLUX_DEN_z  = ocn.TFLUX_DEN_z,
        Dh = ocn.Dh_T,
        Dv = ocn.Dv_T,
    )

    calDiffAdv_QUICK!(
        ocn,
        qs          = ocn.Ss,
        FLUX_bot      = ocn.SFLUX_bot,
        dΔqdt       = ocn.dΔSdt,
        FLUX_CONV   = ocn.SFLUX_CONV,
        FLUX_CONV_h = ocn.SFLUX_CONV_h,
        FLUX_DEN_x  = ocn.SFLUX_DEN_x,
        FLUX_DEN_y  = ocn.SFLUX_DEN_y,
        FLUX_DEN_z  = ocn.SFLUX_DEN_z,
        Dh = ocn.Dh_S,
        Dv = ocn.Dv_S,
    )
=#
    #println("##### QUICK SPEEDUP")
    #println("before calDiffAdv T_ML[60,70] = ", ocn.T_ML[60,70])

    calDiffAdv_QUICK_SpeedUp!(
        ocn,
        qs          = ocn.Ts,
        dΔqdt       = ocn.dΔTdt,
        FLUX_CONV   = ocn.TFLUX_CONV,
        FLUX_CONV_h = ocn.TFLUX_CONV_h,
        FLUX_DEN_x  = ocn.TFLUX_DEN_x,
        FLUX_DEN_y  = ocn.TFLUX_DEN_y,
        FLUX_DEN_z  = ocn.TFLUX_DEN_z,
        Dh = ocn.Dh_T,
        Dv = ocn.Dv_T,
        Δt = Δt,
    )

    #println("after calDiffAdv T_ML[60,70] = ", ocn.T_ML[60,70])
    #println("after calDiffAdv TFLUX_CONV[1, 60,70] = ", ocn.TFLUX_CONV[1, 60,70])

    calDiffAdv_QUICK_SpeedUp!(
        ocn,
        qs          = ocn.Ss,
        dΔqdt       = ocn.dΔSdt,
        FLUX_CONV   = ocn.SFLUX_CONV,
        FLUX_CONV_h = ocn.SFLUX_CONV_h,
        FLUX_DEN_x  = ocn.SFLUX_DEN_x,
        FLUX_DEN_y  = ocn.SFLUX_DEN_y,
        FLUX_DEN_z  = ocn.SFLUX_DEN_z,
        Dh = ocn.Dh_S,
        Dv = ocn.Dv_S,
        Δt = Δt,
    )

    @loop_hor ocn i j let

        Nz = ocn.Nz[i, j]
        zs   = ocn.cols.zs[i, j]
        hs   = ocn.cols.hs[i, j]
        h_ML = ocn.h_ML[i, j]
        FLDO = ocn.FLDO[i, j]

        for k = 1:ocn.Nz[i, j]
            ocn.Ts[k, i, j] += Δt * ocn.TFLUX_CONV[k, i, j]
            ocn.Ss[k, i, j] += Δt * ocn.SFLUX_CONV[k, i, j]
        end

        # Adjust ΔT, ΔS
        ocn.ΔT[i, j] += ocn.dΔTdt[i, j] * Δt
        ocn.ΔS[i, j] += ocn.dΔSdt[i, j] * Δt

        ocn.T_ML[i, j] = unmixFLDOKeepDiff!(;
            qs   = ocn.cols.Ts[i, j],
            zs   = zs,
            hs   = hs,
            h_ML = h_ML,
            FLDO = FLDO,
            Nz   = Nz,
            Δq   = ocn.ΔT[i, j],
        )

        ocn.S_ML[i, j] = unmixFLDOKeepDiff!(;
            qs   = ocn.cols.Ss[i, j],
            zs   = zs,
            hs   = hs,
            h_ML = h_ML,
            FLDO = FLDO,
            Nz   = Nz,
            Δq   = ocn.ΔS[i, j],
        )

        OC_updateB!(ocn, i, j)

        if do_convadjust
            OC_doConvectiveAdjustment!(ocn, i, j)
        end

    end

    #new_intT = sum(ocn.ASUM.filter_T * ocn.ASUM.T_Δvol_T * view(ocn.Ts, :))
    #println("Total integrated T: ", new_intT, "; change=", (new_intT - old_intT)/old_intT * 100, " % ")
end

max_stars_repo_name: meteorologytoday/IOM

max_issues_repo_name: meteorologytoday/IOM

max_forks_repo_name: meteorologytoday/IOM


text:
module TrieLab

include("staticdict.jl")
include("keys.jl")
include("trie.jl")

end


text:
export BoundingVolumeHierarchy

"""
    AccelerationStructure

Base Type for all Acceleration Structures. These can be used to speed up
rendering by a significant amount. The main design behind an acceleration
structure should be such that it can be used as a simple replacement for
a vector of [`Object`](@ref)s.
"""
abstract type AccelerationStructure end

"""
    BVHNode

A node in the graph created from the scene list using [`BoundingVolumeHierarchy`](@ref).

### Fields:

* `x_min`       - Minimum x-value for the bounding box
* `x_max`       - Maximum x-value for the bounding box
* `y_min`       - Minimum y-value for the bounding box
* `y_max`       - Maximum y-value for the bounding box
* `z_min`       - Minimum z-value for the bounding box
* `z_max`       - Maximum z-value for the bounding box
* `index_start` - Starting triangle in the scene list
* `index_end`   - Last triangle in the scene list
* `left_child`  - Can be `nothing` or another [`BVHNode`](@ref)
* `right_child` - Can be `nothing` or another [`BVHNode`](@ref)
"""
struct BVHNode{T}
    x_min::T
    x_max::T
    y_min::T
    y_max::T
    z_min::T
    z_max::T
    index_start::Int
    index_end::Int
    left_child::Union{Nothing, BVHNode}
    right_child::Union{Nothing, BVHNode}
end

"""
    BoundingVolumeHierarchy

An [`AccelerationStructure`](@ref) which constructs bounding boxes around
groups of triangles to speed up intersection checking. A detailed description
of ths technique is present [here](https://www.scratchapixel.com/lessons/advanced-rendering/introduction-acceleration-structure/bounding-volume-hierarchy-BVH-part1).

### Fields:

* `scene_list` - The scene list passed into the BVH constructor but in
                 sorted order
* `root_node`  - Root [`BVHNode`](@ref)

### Constructors:

* `BoundingVolumeHierarchy(scene::Vector)`

"""
struct BoundingVolumeHierarchy{T, S} <: AccelerationStructure
    scene_list::Vector{T}
    root_node::BVHNode{S}
end

const BVH = BoundingVolumeHierarchy

function BoundingVolumeHierarchy(scene::Vector)
    x_min, x_max = extrema(hcat([[s.v1.x[], s.v2.x[], s.v3.x[]] for s in scene]...))
    y_min, y_max = extrema(hcat([[s.v1.y[], s.v2.y[], s.v3.y[]] for s in scene]...))
    z_min, z_max = extrema(hcat([[s.v1.z[], s.v2.z[], s.v3.z[]] for s in scene]...))

    longest_direction = getindex([:x, :y, :z],
                                 argmax([x_max - x_min, y_max - y_min, z_max - z_min]))

    centroids = map(t -> (t.v1 + t.v2 + t.v3) / 3, scene)

    centroid_dict = IdDict{eltype(scene), Vec3}()
    for (s, c) in zip(scene, centroids)
        centroid_dict[s] = c
    end

    scene = sort(scene, by = x -> getproperty(centroid_dict[x], longest_direction)[])

    search_space = map(x -> getproperty(x, longest_direction)[], centroids)
    split_value = median(search_space)
    split_index = searchsortedfirst(search_space, split_value)

    if split_index - 1 == length(scene) || split_index == 1
        bvhnode = BVHNode(x_min, x_max, y_min, y_max, z_min, z_max, 1, length(scene),
                          nothing, nothing)
    else
        left_child, sc = BVHNode(scene[1:split_index - 1], 1, centroid_dict)
        scene[1:split_index - 1] .= sc
        right_child, sc = BVHNode(scene[split_index:end], split_index, centroid_dict)
        scene[split_index:end] .= sc
        bvhnode = BVHNode(x_min, x_max, y_min, y_max, z_min, z_max, 1, length(scene),
                          left_child, right_child)
    end
    return BoundingVolumeHierarchy(scene, bvhnode)
end

function BVHNode(scene, index, centroid_dict)
    length(scene) == 0 && return (nothing, scene)

    x_min, x_max = extrema(hcat([[s.v1.x[], s.v2.x[], s.v3.x[]] for s in scene]...))
    y_min, y_max = extrema(hcat([[s.v1.y[], s.v2.y[], s.v3.y[]] for s in scene]...))
    z_min, z_max = extrema(hcat([[s.v1.z[], s.v2.z[], s.v3.z[]] for s in scene]...))

    if length(scene) <= 100
        return BVHNode(x_min, x_max, y_min, y_max, z_min, z_max, index, index + length(scene) -1,
                       nothing, nothing), scene
    end

    longest_direction = getindex([:x, :y, :z],
                                 argmax([x_max - x_min, y_max - y_min, z_max - z_min]))

    scene = sort(scene, by = x -> getproperty(centroid_dict[x], longest_direction)[])

    centroids = [centroid_dict[s] for s in scene]

    search_space = map(x -> getproperty(x, longest_direction)[], centroids)
    split_value = median(search_space)
    split_index = searchsortedfirst(search_space, split_value)

    if split_index - 1 == length(scene) || split_index == 1
        return BVHNode(x_min, x_max, y_min, y_max, z_min, z_max, index, index + length(scene) - 1,
                       nothing, nothing), scene
    end

    left_child, sc = BVHNode(scene[1:split_index - 1], index, centroid_dict)
    scene[1:split_index - 1] .= sc

    right_child, sc = BVHNode(scene[split_index:end], index + split_index - 1, centroid_dict)
    scene[split_index:end] .= sc

    bvhnode = BVHNode(x_min, x_max, y_min, y_max, z_min, z_max, index, index + length(scene) - 1,
                      left_child, right_child)

    return bvhnode, scene
end

bvh_depth(bvh::BoundingVolumeHierarchy) = bvh_depth(bvh.root_node)

bvh_depth(bvh::BVHNode) = 1 + max(bvh_depth(bvh.left_child), bvh_depth(bvh.right_child))

bvh_depth(::Nothing) = 0

isleafnode(bvh::BVHNode) = isnothing(bvh.left_child) && isnothing(bvh.right_child)

function intersection_check(bvh::BVHNode, origin, direction)
    function check_intersection_ray(dir_x, dir_y, dir_z,
                                    ori_x, ori_y, ori_z)
        a, b = dir_x < 0 ? (bvh.x_max, bvh.x_min) : (bvh.x_min, bvh.x_max)
        tmin = (a - ori_x) / dir_x
        tmax = (b - ori_x) / dir_x
        a, b = dir_y < 0 ? (bvh.y_max, bvh.y_min) : (bvh.y_min, bvh.y_max)
        tymin = (a - ori_y) / dir_y
        tymax = (b - ori_y) / dir_y
        (tmin > tymax || tymin > tmax) && return false
        tmin = max(tmin, tymin)
        tmax = min(tmax, tymax)
        a, b = dir_z < 0 ? (bvh.z_max, bvh.z_min) : (bvh.z_min, bvh.z_max)
        tzmin = (a - ori_z) / dir_z
        tzmax = (b - ori_z) / dir_z
        (tmin > tzmax || tzmin > tmax) && return false
        return true
    end
    return broadcast(check_intersection_ray, direction.x, direction.y,
                     direction.z, origin.x, origin.y, origin.z)
end

intersect(bvh::BoundingVolumeHierarchy, origin, direction) =
    intersect(bvh.root_node, bvh, origin, direction)

function intersect(bvh::BVHNode, bvh_structure::BoundingVolumeHierarchy,
                   origin, direction)
    hit = intersection_check(bvh, origin, direction)
    dir_rays = extract(hit, direction)
    ori_rays = extract(hit, origin)
    if isleafnode(bvh)
        t_values = IdDict{Triangle, Vector}(
                       map(x -> (x, intersect(x, ori_rays, dir_rays)),
                           bvh_structure.scene_list[bvh.index_start:bvh.index_end]))
    else
        t_values = intersect(bvh.left_child, bvh_structure, ori_rays, dir_rays)
        # We do not need to check the ones which have already intersected but let's
        # keep the code simple for now
        # if !isnothing(bvh.right_child)
        merge!(t_values, intersect(bvh.right_child, bvh_structure, ori_rays, dir_rays))
        # end
    end
    for k in keys(t_values)
        t_values[k] = place(t_values[k], hit, Inf)
    end
    return t_values
end


text:
using Plots
gr(titlefontsize=12,size=(500,300),transpose=true,colorbar=false)

using FourierGPE
# Units
# this example works in oscillator units

## Initialize simulation

L = (40.0,40.0)
N = (256,256)
sim = Sim(L,N)
@unpack_Sim sim

## simulation parameters
μ = 15.0

# potential
import FourierGPE.V
V(x,y,t) = 0.5*(x^2 + y^2)

# TF state
ψtf(x,y,μ,g) = sqrt(μ/g)*sqrt(max(1.0-V(x,y,0.0)/μ,0.0)+im*0.0)

# make initial state
x,y = X
ψi = ψtf.(x,y',μ,g)
ϕi = kspace(ψi,sim)

@pack_Sim! sim

## evolve
sol = runsim(sim)

## pull out the ground state
ϕg = sol[end]
ψg = xspace(ϕg,sim)
showpsi(x,y,ψg)

R(w) = sqrt(2*μ/w^2)
R(1)
plot(x, abs2.(ψtf.(x,0.,μ,g)))
plot!(x, abs2.(ψg[:,128]))

## free expansion
V(x,y,t) = 0.0

# sim parameters
γ = 0.0
tf = 10.0
t = LinRange(ti,tf,Nt)
ϕi = ϕg

sime = Sim(L,N)
@pack! sime = tf,t,γ,ϕi

## evolve
sole = runsim(sime)

ϕf = sole[end]
ψf = xspace(ϕf,sime)
showpsi(x,y,ψf)

anim = @animate for i=1:Nt-6
    ψ = xspace(sole[i],sime)
    showpsi(x,y,ψ)
end

gif(anim,"./examples/expand.gif",fps=30)


text:
################################################################################
#                             EXPECTATION METHODS
################################################################################
## Determine if parameter(s) use a distrubution set
# IndependentParameter and UniDistributionSet
function _has_distribution_set(pref::InfiniteOpt.IndependentParameterRef,
                               set::InfiniteOpt.UniDistributionSet)::Bool
    return true
end

# DependentParameter and MultiDistributionSet
function _has_distribution_set(pref::InfiniteOpt.DependentParameterRef,
                               set::InfiniteOpt.MultiDistributionSet)::Bool
    return true
end

# DependentParameter and CollectionSet
function _has_distribution_set(pref::InfiniteOpt.DependentParameterRef,
                               set::InfiniteOpt.CollectionSet)::Bool
    sets = InfiniteOpt.collection_sets(set)
    return sets[InfiniteOpt._param_index(pref)] isa InfiniteOpt.UniDistributionSet
end

# Fallback
function _has_distribution_set(pref, set)::Bool
    return false
end

# Define coefficient function for expectations
_expect_coeffs(x::Array) = ones(size(x)[end]) ./ size(x)[end]

"""
    expect(expr::JuMP.AbstractJuMPScalar,
           prefs::Union{GeneralVariableRef, AbstractArray{GeneralVariableRef};
           [min_num_supports::Int = DefaultNumSupports])::GeneralVariableRef

Creates a measure that represents the expected value of an expression based on
`prefs`. If `prefs` are not random parameters then this will be equivalent
to the following call:
```julia
1/total_num_supports * support_sum(expr, prefs)
```
Note that min_num_supports should be 0 if a single dependent parameter is given.
Also, note that it is preferred to call [`@expect`](@ref) when `expr` is not
just a single variable reference.

**Example**
```julia-repl
julia> @infinite_parameter(model, x in Normal())
x

julia> @infinite_variable(model, f(x))
f(x)

julia> meas = expect(f, min_num_supports = 2)
𝔼{x}[f(x)]

julia> expand(meas)
0.5 f(0.6791074260357777) + 0.5 f(0.8284134829000359)
```
"""
function expect(expr::JuMP.AbstractJuMPScalar,
    prefs::Union{InfiniteOpt.GeneralVariableRef, AbstractArray{InfiniteOpt.GeneralVariableRef}};
    min_num_supports::Int = InfiniteOpt.DefaultNumSupports
    )::InfiniteOpt.GeneralVariableRef
    # check the inputs
    InfiniteOpt._check_params(prefs)
    dpref = InfiniteOpt.dispatch_variable_ref(first(prefs))
    # update properly for single dependent parameters
    if prefs isa InfiniteOpt.GeneralVariableRef && dpref isa InfiniteOpt.DependentParameterRef
        if !(min_num_supports in [0, InfiniteOpt.DefaultNumSupports])
            @warn("Cannot specify a nonzero `min_num_supports` for individual " *
                  "dependent parameters.")
        end
        min_num_supports = 0
    end
    # prepare the label
    set = InfiniteOpt._parameter_set(dpref)
    if _has_distribution_set(dpref, set)
        label = InfiniteOpt.WeightedSample
        is_expect = true
    else
        label = InfiniteOpt.All # TODO maybe do something more rigorous
        is_expect = false
    end
    ordered_prefs = InfiniteOpt._make_ordered_vector(prefs)
    length(ordered_prefs) == 1 ? bounds = NaN : bounds = map(e -> NaN, ordered_prefs)
    # make the data
    data = InfiniteOpt.FunctionalDiscreteMeasureData(ordered_prefs, _expect_coeffs,
                                                     min_num_supports, label,
                                                     InfiniteOpt.default_weight,
                                                     bounds, bounds, is_expect)
    # make the measure
    return InfiniteOpt.measure(expr, data, name = "expect")
end

"""
    @expect(expr::JuMP.AbstractJuMPScalar,
            prefs::Union{GeneralVariableRef, AbstractArray{GeneralVariableRef};
            [min_num_supports::Int = DefaultNumSupports])::GeneralVariableRef

An efficient wrapper for [`expect`](@ref). Please see its doc string more
information.
"""
macro expect(expr, prefs, args...)
    _error(str...) = InfiniteOpt._macro_error(:integral, (expr, prefs, args...),
                                              str...)
    extra, kw_args, requestedcontainer = InfiniteOpt._extract_kw_args(args)
    if length(extra) > 0
        _error("Unexpected positional arguments." *
               "Must be of form @expect(expr, prefs, min_num_supports = some_integer).")
    end
    if !isempty(filter(kw -> kw.args[1] != :min_num_supports, kw_args))
        _error("Unexpected keyword arugments.")
    end
    expression = :( JuMP.@expression(InfiniteOpt._Model, $expr) )
    mref = :( expect($expression, $prefs; ($(kw_args...))) )
    return esc(mref)
end

"""
    𝔼(expr::JuMP.AbstractJuMPScalar,
      prefs::Union{GeneralVariableRef, AbstractArray{GeneralVariableRef}};
      [min_num_supports::Int = DefaultNumSupports]
      )::GeneralVariableRef)

A convenient wrapper for [`expect`](@ref). The unicode symbol `𝔼` is produced by
`\\bbE`.
"""
function 𝔼(expr::JuMP.AbstractJuMPScalar,
    prefs::Union{InfiniteOpt.GeneralVariableRef, AbstractArray{InfiniteOpt.GeneralVariableRef}};
    min_num_supports::Int = InfiniteOpt.DefaultNumSupports
    )::InfiniteOpt.GeneralVariableRef
    return expect(expr, prefs, min_num_supports = min_num_supports)
end

"""
    @𝔼(expr::JuMP.AbstractJuMPScalar,
       prefs::Union{GeneralVariableRef, AbstractArray{GeneralVariableRef};
       [min_num_supports::Int = DefaultNumSupports])::GeneralVariableRef

A convenient wrapper for [`@expect`](@ref). The unicode symbol `𝔼` is produced by
`\\bbE`.
"""
macro 𝔼(expr, prefs, args...)
    return esc(:( @expect($expr, $prefs, $(args...)) ))
end


text:
struct Delta <: Distances.PreMetric
end

@inline function Distances._evaluate(::Delta, a::AbstractVector, b::AbstractVector)
    @boundscheck if length(a) != length(b)
        throw(DimensionMismatch(
            "first array has length $(length(a)) which does not match the length of the " *
            "second, $(length(b)).",
        ))
    end
    return a == b
end

Distances.result_type(::Delta, Ta::Type, Tb::Type) = promote_type(Ta, Tb)

@inline (dist::Delta)(a::AbstractArray, b::AbstractArray) = Distances._evaluate(dist, a, b)
@inline (dist::Delta)(a::Number, b::Number) = a == b


text:
#
# This file is part of the DiscreteEvents.jl Julia package, MIT license
#
# Paul Bayer, 2020
#
# This is a Julia package for discrete event simulation
#

# -------------------------------------------------------------------
# methods for RTClocks
# -------------------------------------------------------------------
tau(rtc::RTClock) = rtc.time

"""
    resetClock!(rtc::RTClock)

Reset a real time clock. Set its time to zero and delete all scheduled and
sampling events.
"""
resetClock!(rtc::RTClock) = put!(rtc.cmd, Reset(false))

# -------------------------------------------------------------------
# RTClock state machine operations
# -------------------------------------------------------------------

# reset an rt clock
function step!(RTC::RTClock, ::Union{Idle, Busy}, ::Reset)
    RTC.t0 = time_ns()*1e-9
    RTC.clock.evcount = 0
    RTC.clock.scount = 0
end

# register an event to the rt clock
step!(RTC::RTClock, ::Union{Idle, Busy}, σ::Register) = _register!(RTC.clock, σ.x)

# fallback transition
step!(RTC::RTClock, q::ClockState, σ::ClockEvent) = error("transition q=$q, σ=$σ not implemented")

# -------------------------------------------------------------------
# event loop for an RTClock.
# -------------------------------------------------------------------
function _RTClock(rtc::RTClock)
    sf = Array{Base.StackTraces.StackFrame,1}[]
    exc = nothing
    rtc.clock.state = Idle()
    rtc.task = current_task()

    while true
        if isready(rtc.cmd)
            σ = take!(rtc.cmd)
            if σ isa Stop
                break
            elseif σ isa Diag
                put!(rtc.back, Response((exc, sf)))
            elseif handle_exceptions()
                try
                    step!(rtc, rtc.clock.state, σ)
                catch exc
                    sf = stacktrace(catch_backtrace())
                    @warn "clock $(rtc.id), thread $(rtc.thread) exception: $exc"
                end
            else
                step!(rtc, rtc.clock.state, σ)
            end
        end
        rtc.time = time_ns()*1e-9 - rtc.t0
        while !isempty(rtc.clock.sc.samples) && rtc.clock.tn ≤ rtc.time
            _tick!(rtc.clock)
            rtc.clock.tn += rtc.clock.Δt
        end
        while !isempty(rtc.clock.sc.events) && _nextevtime(rtc.clock) ≤ rtc.time
            _event!(rtc.clock)
        end
        wait(rtc.Timer)
    end
    close(rtc.Timer)
end

# ---------------------------------------------------------
# starting and destroying real time clocks
# ---------------------------------------------------------
"""
    createRTClock(T::Float64, id::Int, thrd::Int=nthreads(); ch_size::Int=256)

Create, start and return a real time Clock.

The clock takes the current system time and starts to count in seconds with
the given period `T`. Events or sampling functions can then be scheduled to it.

# Arguments
- `T::Float64`:           period (clock resolution) in seconds, T ≥ 0.001
- `id::Int`:              clock identification number other than 0:nthreads()
- `thrd::Int=nthreads()`: thread, the clock task should run in
- `ch_size::Int=256`:     clock communication channel size
"""
function createRTClock(T::Float64, id::Int, thrd::Int=nthreads(); ch_size::Int=256)
    T ≥ 0.001 || throw(ArgumentError("RTClock cannot have a period of $T < 0.001 seconds"))
    id ∉ 0:nthreads() || throw(ArgumentError("RTClock id $id forbidden!"))
    rtc = RTClock(
            Timer(T, interval=T), Clock(),
            Channel{ClockEvent}(ch_size), Channel{ClockEvent}(ch_size),
            id, thrd, 0.0, time_ns()*1e-9, T, current_task())
    rtc.clock.id = id
    onthread(thrd, wait=false) do
        _RTClock(rtc)
    end
    return rtc
end

"""
    stopRTClock(rtc::RTClock)

Stop a real time clock.
"""
stopRTClock(rtc::RTClock) = put!(rtc.cmd, Stop())


text:
struct EmptyCone <: MOI.AbstractSymmetricMatrixSetTriangle end
MOI.side_dimension(::EmptyCone) = 0

"""
    struct PositiveSemidefinite2x2ConeTriangle <: MOI.AbstractSymmetricMatrixSetTriangle end

Cone of positive semidefinite matrices of 2 rows and 2 columns.
"""
struct PositiveSemidefinite2x2ConeTriangle <: MOI.AbstractSymmetricMatrixSetTriangle end
MOI.side_dimension(::PositiveSemidefinite2x2ConeTriangle) = 2

# isbits types, nothing to copy
function Base.copy(set::Union{EmptyCone, PositiveSemidefinite2x2ConeTriangle})
    return set
end


text:
import Base

"""
"""
get_order(x::T) where T<:Real = parse(Int, match(r"e(.*)", Printf.@sprintf("%e", x))[1])

"""
    matrix(mat; )

```jldoctest
julia> matrix(1:3)
3×1 Array{Int64,2}:
 1
 2
 3
```
"""
function matrix(lst::Vector{T}) where T<:AbstractMatrix
    !isvector(first(lst)) && error("Error in concatenating the list")

    # If the input is a list of *COLUMN* matrices,
    findmin(size(first(lst)))[2] == 1 && (lst = LinearAlgebra.adjoint.(lst))

    return _matrix(lst)
end

matrix(lst) = lst
matrix(lst::Vector{T}) where T<:Union{Vector,Real} = _matrix(lst)
matrix(lst::Matrix{T}) where T<:Union{Vector,Real} = isvector(lst) && _matrix(lst)

matrix(T::UnionAll, N; value=0) = T(fill(value, N, N))
matrix(T::LinearAlgebra.UniformScaling, N) = T(N)*1.0

_matrix(lst) = convert(Matrix, cat(lst...; dims=2)')

"""
    vectorize(x::Any)
This method returns a list of vectors.

```jldoctest
julia> mat =  [[1,4]  [2,5]  [3,6]]
2×3 Array{Int64,2}:
 1  2  3
 4  5  6

julia> vectorize(mat)
2-element Array{Array{Int64,1},1}:
 [1, 2, 3]
 [4, 5, 6]
```
"""
vectorize(x::Vararg{Any}) = vectorize(Tuple.(hcat.(x...)))
vectorize(x::Matrix) = collect.(collect(eachrow(x)))
vectorize(x::Vector) = x
vectorize(x::Any) = [x;]
vectorize(x::Poly) = x

"These methods extend `Base.sign` to `Waterfall` types."
Base.sign(data::T) where T <: Data = sign.(get_value(data))
Base.sign(cascade::Cascade{Data}) = sign.(collect_data(cascade))

"These methods extend `Base.length` to `Waterfall` types."
Base.length(x::Data) = length(x.order)
Base.length(x::T) where T <: Geometry = x.nsample
Base.length(x::T) where T <: Cascade = length(x.start)
Base.length(x::T) where T <: Plot = length(x.cascade)
Base.length(x::T) where T <: Shape = x.position
Base.length(x::Handle) = 1
Base.length(x::Annotation) = 1
Base.length(x::Coloring) = 1
Base.length(x::Blending) = 1
# Base.length(x::T) where T <: Axis = length(x.ticks)

"These methods extend `Base.value` to `Waterfall` Types"
Base.values(x::Data) = _values(x)
Base.values(x::Cascade) = _values(x)
Base.values(x::Axis) = _values(x)
Base.values(x::Plot) = _values(x)
Base.values(x::Coloring) = _values(x)
Base.values(x::Blending) = _values(x)
Base.values(x::Handle) = _values(x)
Base.values(x::Annotation) = _values(x)
_values(x::T) where T <: Any = Tuple([getproperty(x, f) for f in fieldnames(T)])

"These methods extend `Base.copy` to `Waterfall` and `Luxor` Types"
Base.copy(x::Symbol) = x
Base.copy(x::Axis) = _copy(x)
Base.copy(x::Data) = _copy(x)
Base.copy(x::Vector{Data}) = _copy.(x)

Base.copy(x::Luxor.Point) = _copy(x)
Base.copy(x::Tuple{Luxor.Point,Luxor.Point}) = _copy.(x)

Base.copy(x::Luxor.RGB) = _copy(x)
Base.copy(x::Vector{Luxor.RGB}) = _copy.(x)
Base.copy(x::Coloring) = _copy(x)
Base.copy(x::Blending) = _copy(x)

Base.copy(x::T) where T<:Shape = T(copy.(_values(x))...)
Base.copy(x::Cascade) = Cascade(copy.(_values(x))...)
Base.copy(x::Plot) = Plot(copy.(_values(x))...)

_copy(x::T) where T <: Any = T(values(x)...)

"""
    Base.minimum(pos::Vector{Luxor.Point}; kwargs...)
    Base.minimum(pos::Vector{Tuple}; kwargs...)
These methods extend `Base.minimum` to find the minimum `x` or `y` value of a list of points
or list of pairs of points, with `dims` indicating whether to return the `x` (`dims=1`) or
`y` (`dims=2`) element.

# Keyword Arguments
- `dims`
"""
Base.maximum(lst::Vector{T}; dims) where T<:Luxor.Point = maximum(getindex.(lst,dims))
Base.maximum(lst::Vector{T}; kwargs...) where T<:Tuple = maximum(vcat(collect.(lst)...); kwargs...)

"""
    Base.maximum(pos::Vector{Luxor.Point}; kwargs...)
    Base.maximum(pos::Vector{Tuple}; kwargs...)
These methods extend `Base.maximum` to find the maximum `x` or `y` value of a list of points
or list of pairs of points, with `dims` indicating whether to return the `x` (`dims=1`) or
`y` (`dims=2`) element.

# Keyword Arguments
- `dims`
"""
Base.minimum(lst::Vector{T}; dims) where T<:Luxor.Point = minimum(getindex.(lst,dims))
Base.minimum(lst::Vector{T}; kwargs...) where T<:Tuple = minimum(vcat(collect.(lst)...); kwargs...)

"""
This function finds the midpoint between the minimum and maximum values of the list of
input points or pairs of points.

# Keyword Arguments
- `dims`
"""
mid(pos; kwargs...) = Statistics.mean([minimum(pos; kwargs...), maximum(pos; kwargs...)])

"This function returns all permutations of the elements in the input vectors or vector of vectors"
permute(x::Vector{Vector{T}}) where T<: Any = permute(x...)
permute(x::Vararg{Any}) = vec(collect(collect.(Iterators.product(x...))))

"This function returns all of the ways to combine the elements in `lst`. At the moment,
due to memory size, this will truncate to a maximum length of `maxlen=6`."
function combinate(lst::AbstractArray; maxlen=6)
    lst = collect(lst)[1:max(length(lst),maxlen)]
    x = fill(collect(lst), length(lst))
    x = permute(x)
    return sort(x[vec(length.(unique.(x)).==length.(x))])
end

"""
    Base.convert(::Type{SparseArrays.SparseMatrixCSC}, order::AbstractVector{Int})
This method converts an `Nx1` list of indices into a sparce matrix ``S`` that,
if multiplied by a matrix `A`:
- ``S \\cdot A`` will reorder the **rows** of ``A`` to `order`
- ``A \\cdot S`` will reorder the **columns** of ``A`` to `order`

```jldoctest
julia> order = [1,3,2];

julia> sparse_matrix = convert(SparseArrays.SparseMatrixCSC, order)
3×3 SparseArrays.SparseMatrixCSC{Int64,Int64} with 3 stored entries:
  [1, 1]  =  1
  [3, 2]  =  1
  [2, 3]  =  1

julia> Matrix(sparse_matrix)
  3×3 Array{Int64,2}:
   1  0  0
   0  0  1
   0  1  0

julia> sparse_matrix * collect(1:length(order)) == order
   true
```

    convert(::Type{DataFrames.DataFrame}, cascade::Cascade)
This method converts `cascade` into a "sample-by-step"-dimension DataFrame with the step
labels as property property names.

    Base.convert(::Type{Matrix}, lst::AbstractVector; dims=2)
"""
function Base.convert(::Type{SparseArrays.SparseMatrixCSC}, order::T) where T <: AbstractVector{Int}
    N = length(order)
    return SparseArrays.sparse(1:N, order, fill(1,N))
end

function Base.convert(::Type{DataFrames.DataFrame}, x::Cascade)
    data = collect_data(x)
    sublabel = getindex.(match.(r"(\S*)", get_sublabel.(data)),1)
    label = get_label.(data)
    [label[ii] = "$(label[ii]) ($(sublabel[ii]))" for ii in [1,length(data)]]

    value = get_value(data)

    return DataFrames.DataFrame(value', label; makeunique=true)
end

"""
    isinteger(x::String)
This method returns true if `x` can be parsed as an integer.
"""
Base.isinteger(x::String) = getindex(match(r"(\d*)", x),1)==x

"""
    tryinteger(x::String)
This method parses `x` as an `Int64` if possible, and otherwise returns `x`.

    tryinteger(lst::Vector{String})
This methods `lst` as `Vector{Int64}` if this is possible for **all** array elements.
"""
tryinteger(x::String) = isinteger(x) ? parse(Int64,x) : x
tryinteger(lst::AbstractArray) = all(isinteger.(lst)) ? parse.(Int64,lst) : lst

"""
    swapat!(lst::AbstractArray, a, b)
This function swaps the values of `lst` at elements `a` and `b`.
"""
function swapat!(lst::AbstractArray, a::Int, b::Int)
    tmp = lst[a]
    lst[a] = lst[b]
    lst[b] = tmp
    return lst
end


text:
module Kaleidoscope

import LLVM

include("lexer.jl")
include("ast.jl")
include("scope.jl")
include("codegen.jl")
include("main.jl")
include("utilities.jl")

end # module


text:
using Test
using Printf: @printf
using LinearAlgebra: tr, dot
using StaticArrays: SVector
using QuantumLattices.Essentials.Frameworks
using QuantumLattices.Essentials: update, reset!
using QuantumLattices.Essentials.Spatials: Point, AbstractPID, PID, Bond, Bonds, Lattice, acrossbonds, zerothbonds, decompose
using QuantumLattices.Essentials.DegreesOfFreedom: SimpleIID, SimpleInternal, IIDSpace, Coupling, Subscript, Subscripts
using QuantumLattices.Essentials.DegreesOfFreedom: Term, Hilbert, Index, Table, OID, OIDToTuple, @couplings
using QuantumLattices.Essentials.QuantumOperators: ID, Operator, Operators, id, idtype, Identity
using QuantumLattices.Prerequisites: Float
using QuantumLattices.Interfaces:  expand!, expand, add!
using QuantumLattices.Prerequisites.Traits: contentnames, reparameter

import QuantumLattices.Essentials.Frameworks: Parameters
import QuantumLattices.Essentials: prepare!, run!, update!
import QuantumLattices.Essentials.DegreesOfFreedom: ishermitian, couplingcenters
import QuantumLattices.Prerequisites.VectorSpaces: shape

struct FID{N<:Union{Int, Symbol}} <: SimpleIID
    nambu::N
end
@inline Base.adjoint(sl::FID{Int}) = FID(3-sl.nambu)
function Base.angle(id::OID{<:Index{<:AbstractPID, FID{Int}}}, vectors::AbstractVector{<:AbstractVector{Float}}, values::AbstractVector{Float})
    phase=  (length(vectors) == 1) ? 2pi*dot(decompose(id.icoord, vectors[1]), values) :
            (length(vectors) == 2) ? 2pi*dot(decompose(id.icoord, vectors[1], vectors[2]), values) :
            (length(vectors) == 3) ? 2pi*dot(decompose(id.icoord, vectors[1], vectors[2], vectors[3]), values) :
            error("angle error: not supported number of input basis vectors.")
    (id.index.iid.nambu == 1) ? phase : -phase
end

struct FFock <: SimpleInternal{FID{Int}}
    nnambu::Int
end
@inline shape(f::FFock) = (1:f.nnambu,)
@inline FID(i::CartesianIndex, vs::FFock) = FID(i.I...)
@inline CartesianIndex(did::FID{Int}, vs::FFock) = CartesianIndex(did.nambu)
@inline shape(iidspace::IIDSpace{FID{Symbol}, FFock}) = (1:iidspace.internal.nnambu,)
@inline shape(iidspace::IIDSpace{FID{Int}, FFock}) = (iidspace.iid.nambu:iidspace.iid.nambu,)
@inline shape(iidspace::IIDSpace{FID{Symbol}, FFock}) = (1:iidspace.internal.nnambu,)
@inline shape(iidspace::IIDSpace{FID{Int}, FFock}) = (iidspace.iid.nambu:iidspace.iid.nambu,)

const FCoupling{V, I<:ID{FID}, C<:Subscripts} = Coupling{V, I, C}
@inline couplingcenters(::(Coupling{V, <:ID{FID}} where V), ::Bond, ::Val) = (1, 2)
@inline FCoupling(value, nambus::Tuple{Vararg{Int}}) = Coupling(value, ID(FID, nambus), Subscripts((nambu=Subscript(nambus),)))
@inline FCoupling(value, nambus::Subscript) = Coupling(value, ID(FID, convert(Tuple, nambus)), Subscripts((nambu=nambus,)))

@inline ishermitian(::Type{<:Term{:Mu}}) = true
@inline ishermitian(::Type{<:Term{:Hp}}) = false

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

@testset "Boundary" begin
    op = Operator(4.5,
        OID(Index(PID(1), FID(2)), SVector(0.5, 0.5), SVector(0.0, 0.0)),
        OID(Index(PID(2), FID(1)), SVector(1.5, 1.5), SVector(1.0, 1.0))
        )
    bound = Boundary{(:θ₁, :θ₂)}([0.1, 0.2], [[1.0, 0.0], [0.0, 1.0]])
    M = reparameter(typeof(op), :value, Complex{Float64})
    @test valtype(typeof(bound), typeof(op)) == M
    @test keys(bound) == keys(typeof(bound)) == (:θ₁, :θ₂)
    @test Parameters(bound) == (θ₁=0.1, θ₂=0.2)

    another = Boundary{(:θ₁, :θ₂)}([0.0, 0.0], [[2.0, 0.0], [0.0, 2.0]])
    @test merge!(deepcopy(bound), another) == another

    @test bound(op) ≈ replace(op, 4.5*exp(2im*pi*0.3))
    @test bound(op, origin=[0.05, 0.15]) ≈ replace(op, 4.5*exp(2im*pi*0.1))
    update!(bound, θ₁=0.3)
    @test bound(op) ≈ replace(op, 4.5*exp(2im*pi*0.5))
    @test bound(op, origin=[0.1, 0.1]) ≈ replace(op, 4.5*exp(2im*pi*0.3))

    ops = Operators{M}(op)
    @test valtype(typeof(bound), typeof(ops)) == typeof(ops)
    @test bound(ops) ≈ Operators(replace(op, 4.5*exp(2im*pi*0.5)))
    @test map!(bound, ops) ≈ ops ≈ Operators(replace(op, 4.5*exp(2im*pi*0.5)))

    @test valtype(typeof(plain), typeof(op)) == typeof(op)
    @test valtype(typeof(plain), typeof(ops)) == typeof(ops)
    @test plain(op) == op
    @test plain(ops) == ops
    @test update!(plain) == plain
end

@testset "Formulation" begin
    A(t, μ, Δ; k) = [
          2t*cos(k[1])+2t*cos(k[2])+μ   2im*Δ*sin(k[1])+2Δ*sin(k[2]);
        -2im*Δ*sin(k[1])+2Δ*sin(k[2])   -2t*cos(k[1])-2t*cos(k[2])-μ
    ]
    f = Formulation(A, (t=1.0, μ=0.0, Δ=0.1))
    @test Parameters(f) == (t=1.0, μ=0.0, Δ=0.1)
    @test f(k=[0.0, 0.0]) ≈ [4 0; 0 -4]

    update!(f; μ=0.3)
    @test f(k=[pi/2, pi/2]) ≈ [0.3 0.2+0.2im; 0.2-0.2im -0.3]
end

@testset "Generator" begin
    @test contentnames(CompositeGenerator) == (:operators, :table)
    @test contentnames(Generator) == (:operators, :terms, :bonds, :hilbert, :half, :table)
    @test contentnames(Image) == (:operators, :transformation, :table, :sourceid)

    lattice = Lattice(:Tuanzi, [Point(PID(1), [0.0]), Point(PID(2), [0.5])], vectors=[[1.0]], neighbors=1)
    bonds = Bonds(lattice)
    hilbert = Hilbert{FFock}(pid->FFock(2), lattice.pids)
    table = Table(hilbert, OIDToTuple(:scope, :site))
    boundary = Boundary{(:θ,)}([0.1], lattice.vectors)

    t = Term{:Hp}(:t, 2.0, 1, couplings=@couplings(FCoupling(1.0, (2, 1))))
    μ = Term{:Mu}(:μ, 1.0, 0, couplings=@couplings(FCoupling(1.0, (2, 1))), modulate=true)
    tops₁ = expand(t, filter(acrossbonds, bonds, Val(:exclude)), hilbert, half=true, table=table)
    tops₂ = boundary(expand(one(t), filter(acrossbonds, bonds, Val(:include)), hilbert, half=true, table=table))
    μops = expand(one(μ), filter(zerothbonds, bonds, Val(:include)), hilbert, half=true, table=table)
    i = Identity()

    optp = Operator{Complex{Float}, ID{OID{Index{PID, FID{Int}}, SVector{1, Float}}, 2}}
    entry = Entry(tops₁, (μ=μops,), (t=tops₂, μ=Operators{optp}()), (t=2.0, μ=1.0), boundary)
    @test entry == Entry((t, μ), bonds, hilbert; half=true, table=table, boundary=boundary)
    @test isequal(entry, i(entry))
    @test Parameters(entry) == (t=2.0, μ=1.0, θ=0.1)
    @test entry|>valtype == entry|>typeof|>valtype == Operators{optp, idtype(optp)}
    @test entry|>eltype == entry|>typeof|>eltype == optp
    @test expand(entry) == expand!(Operators{optp}(), entry) ≈ tops₁+tops₂*2.0+μops

    another = Entry(empty(tops₁), (μ=empty(μops),), (t=empty(tops₂), μ=Operators{optp}()), (t=2.0, μ=1.0), boundary)
    @test empty(entry) == empty!(deepcopy(entry)) == another
    @test merge!(deepcopy(another), entry) == merge(another, entry) == entry

    nb = update!(deepcopy(boundary); θ=0.5)
    another = Entry(tops₁, (μ=μops,), (t=nb(tops₂, origin=[0.1]), μ=Operators{optp}()), (t=2.0, μ=1.5), nb)
    dest = deepcopy(entry)
    update!(dest; μ=1.5, θ=0.5)
    @test dest == another
    @test expand(dest) ≈ tops₁+another.boundary(tops₂, origin=[0.1])*2.0+μops*1.5
    dest = deepcopy(entry)
    update!(dest, i, another)
    @test dest == another
    @test expand(dest) ≈ tops₁+another.boundary(tops₂, origin=[0.1])*2.0+μops*1.5

    @test reset!(deepcopy(another), (t, μ), bonds, hilbert, half=true, table=table, boundary=boundary) == entry
    @test reset!(deepcopy(another), i, entry) == entry

    cgen = Generator((t, μ), bonds, hilbert; half=true, table=table, boundary=boundary)
    @test cgen == deepcopy(cgen) && isequal(cgen, deepcopy(cgen))
    @test cgen|>eltype == cgen|>typeof|>eltype == optp
    @test cgen|>valtype == cgen|>typeof|>valtype == Operators{optp, idtype(optp)}
    @test collect(cgen) == collect(expand(cgen))
    @test Parameters(cgen) == (t=2.0, μ=1.0, θ=0.1)
    @test expand!(Operators{optp}(), cgen) == expand(cgen) ≈ tops₁+tops₂*2.0+μops
    @test expand(cgen, :t) ≈ tops₁+tops₂*2.0
    @test expand(cgen, :μ) ≈ μops
    @test expand(cgen, 1)+expand(cgen, 2)+expand(cgen, 3)+expand(cgen, 4) ≈ expand(cgen)
    @test expand(cgen, :μ, 1)+expand(cgen, :μ, 2) ≈ μops
    @test expand(cgen, :t, 3) ≈ tops₁
    @test expand(cgen, :t, 4) ≈ tops₂*2.0
    @test empty!(deepcopy(cgen)) == Generator((t, μ), empty(bonds), empty(hilbert), half=true, table=empty(table), boundary=boundary) == empty(cgen)
    @test reset!(empty(cgen), lattice) == cgen
    @test update!(cgen, μ=1.5)|>expand ≈ tops₁+tops₂*2.0+μops*1.5

    sgen = i(cgen; table=table)
    @test sgen == Image(cgen.operators, i, table, objectid(cgen))
    @test Parameters(sgen) == (t=2.0, μ=1.5, θ=0.1)
    @test expand!(Operators{optp}(), sgen) == expand(sgen) ≈ tops₁+tops₂*2.0+μops*1.5
    @test empty!(deepcopy(sgen)) == Image(empty(cgen.operators), i, empty(table), objectid(cgen)) == empty(sgen)
    @test update!(sgen, μ=3.5)|>expand ≈ tops₁+tops₂*2.0+μops*3.5
    @test update!(sgen, cgen)|>expand ≈ tops₁+tops₂*2.0+μops*1.5
    @test reset!(empty(sgen), cgen; table=table) == sgen
end

mutable struct VCA <: Engine
    t::Float
    U::Float
    dim::Int
end
function update!(vca::VCA; kwargs...)
    vca.t = get(kwargs, :t, vca.t)
    vca.U = get(kwargs, :U, vca.U)
    return vca
end
@inline Parameters(vca::VCA) = Parameters{(:t, :U)}(vca.t, vca.U)

mutable struct GF <: Action
    count::Int
end
@inline prepare!(gf::GF, vca::VCA) = Matrix{Float}(undef, vca.dim, vca.dim)
@inline run!(alg::Algorithm{VCA}, assign::Assignment{GF}) = (assign.action.count += 1; assign.data[:, :] .= alg.engine.t+alg.engine.U)

mutable struct DOS <: Action
    mu::Float
end
@inline update!(eb::DOS; kwargs...) = (eb.mu = get(kwargs, :mu, eb.mu); eb)
@inline prepare!(dos::DOS, vca::VCA) = 0.0
@inline function run!(alg::Algorithm{VCA}, assign::Assignment{DOS})
    rundependences!(alg, assign)
    gf = alg.assignments[first(assign.dependences)]
    assign.data = tr(gf.data)
end
@inline dosmap(params::Parameters) = Parameters{(:t, :U, :mu)}(params.t, params.U, -params.U/2)

@testset "Framework" begin
    vca = VCA(1.0, 8.0, 4)
    @test vca == deepcopy(vca)
    @test isequal(vca, deepcopy(vca))
    @test repr(vca) == "VCA"
    @test string(vca) == "VCA"

    gf = GF(0)
    @test gf == deepcopy(gf)
    @test isequal(gf, deepcopy(gf))
    @test update!(gf) == gf

    vca = Algorithm(:test, vca)
    @test vca == deepcopy(vca)
    @test isequal(vca, deepcopy(vca))
    @test Parameters(vca) == vca.parameters
    @test repr(vca, ≠(:U)) == "test(VCA)_1.0"
    @test repr(vca) == "test(VCA)_1.0_8.0"

    add!(vca, :GF, GF(0))
    rex = r"Action DOS\(DOS\)\: time consumed [0-9]*\.[0-9]*(e[+-][0-9]*)*s."
    @test_logs (:info, rex) vca(:DOS, DOS(-3.5), parameters=(U=7.0,), map=dosmap, dependences=(:GF,))

    dos = vca.assignments[:DOS]
    @test dos == deepcopy(dos)
    @test isequal(dos, deepcopy(dos))
    @test Parameters(dos) == dos.parameters
    @test valtype(dos) == valtype(typeof(dos)) == Float
    @test dos.data == 32.0
    @test dos.action.mu == -3.5
    @test nameof(vca, dos) == "test(VCA)_1.0_7.0_DOS"

    update!(dos, U=6.0)
    vca(:DOS, info=false)
    @test dos.parameters == (t=1.0, U=6.0)
    @test dos.data == 28.0
    @test dos.action.mu == -3.0
end


text:
include("./neural_net.jl")
using BSON: @save, @load
using Dates
using Flux.Tracker
using Plots
using Random
export BATCH_SIZE, save_model, save_data, load_model, load_data

const BATCH_SIZE = 512

const MODEL_PATH = "./model/"
!isdir(MODEL_PATH) && mkpath(MODEL_PATH)
const DATA_PATH = "./data/"
!isdir(DATA_PATH) &&  mkpath(DATA_PATH)
const LOG_PATH = "./logs/"
!isdir(LOG_PATH) &&  mkpath(LOG_PATH)
const EVALUA_PATH = "./trained_model/"
!isdir(EVALUA_PATH) &&  mkpath(EVALUA_PATH)

use_gpu = false
const UPDATE_WINNING_RATE = 0.6
const PNG_PATH = "./logs/pictures/"

function save_model(nn, trained::Bool=false)
    "checking for path to save"
    p = trained ? EVALUA_PATH : MODEL_PATH
    path = p * Dates.format(now(), "mm-dd-HH-MM") * ".model.bson"
    @save path nn
end

function load_model()
    "load agent model"
    path = MODEL_PATH * maximum(readdir(MODEL_PATH))
    @load path nn
    return nn
end

function save_data(onebatch::AbstractArray)
    "saving a batch of data"

    path = DATA_PATH * Dates.format(now(), "mm-dd-HH:MM") * ".data.bson"
    @save path onebatch
    nothing
end

function load_data(num::Int=1)
    "load the latest data"
    files = readdir(DATA_PATH)
    num = min(num, length(files))
    filenames = sort(files, rev=true)[1:num]
    data = asyncmap(1:num, ntasks=10) do i
           @load DATA_PATH*filenames[i] onebatch
           return onebatch
    end
    return data
end

function load_data()
    path = DATA_PATH*rand(readdir(DATA_PATH))
    @load path onebatch
    return onebatch
end

function save_log(filename::String, object)
    path = LOG_PATH*filename*".log"

    open(path, "a+") do io
        println(io, Dates.format(now(), "mm-dd-HH:MM"))
        println(io, object)
    end
end

function checkQuality(winning_rate::Float64, curr_nn)
    path = LOG_PATH * "evaluation.log"
    open(path, "a+") do io
        println(io, "the winning rate is ", winning_rate)
        println("the winning rate is ", winning_rate)
        if winning_rate ≥ UPDATE_WINNING_RATE
            println(io, "update new agent! saving new trained model!")
            println("update new agent! saving new trained model!")
            @async save_model(curr_nn)
        else
            println(io, "fail to defeat the latest agent! not save model")
            println("fail to defeat the latest agent! not save model")
        end
    end
end

function get_animPlot()
    path = PNG_PATH * Dates.format(now(), "mm-dd-HH:MM") * "/"
    !isdir(path) && mkpath(path)
    versionid = length(readdir(PNG_PATH))-1
    anim, plt = Animation(path, String[]), plot([[], []], title="Trainning Loss: Version$versionid\n\n", xlabel="epoch", label=["Each Epoch Loss", "Moving Average Loss"])
    plt.attr[:size] = (1200, 600)
    plt.attr[:dpi] = 170
    anim, plt
end

function save_retrain(anim, plt, losses)
    path = anim.dir
    gif(anim, path*"loss.gif")
    png(plt, path*"losses.png")
    @save path*"losses.data" losses
    nothing
end


text:
module Interfaces

export ⊕,⊗,⋅
export add!,sub!,mul!,div!
export id,rank,dimension,kind,sequence,dims,index,inds
export update!,reset!,register!,prepare!,run!,expand,expand!,decompose,decompose!,regularize,regularize!
export permute,vector,matrix

"Direct sum."
function ⊕ end

"Direct product."
function ⊗ end

"Dot product."
function ⋅ end

"Inplace addition."
function add! end

"Inplace subtraction."
function sub! end

"Inplace multiplication."
function mul! end

"Inplace division."
function div! end

"Inplace Update."
function update! end

"Inplace reset."
function reset! end

"Inplace register."
function register! end

"Inplace prepare."
function prepare! end

"Inplace run."
function run! end

"Id."
function id end

"Rank."
function rank end

"Dimension."
function dimension end

"Kind."
function kind end

"Sequence."
function sequence end

"Dimensions."
function dims end

"Index."
function index end

"Indices."
function inds end

"Get the expansion."
function expand end

"In place expansion."
function expand! end

"Decompose."
function decompose end

"In place decomposition."
function decompose! end

"Regularize."
function regularize end

"In place regularization."
function regularize! end

"Get the permutation."
function permute end

"Vector representation."
function vector end

"Matrix representation."
function matrix end

end # module


text:
"""
    DescriptorSystems.dss(sys::AbstractStateSpace)

Convert `sys` to a descriptor statespace system from [DescriptorSystems.jl](https://andreasvarga.github.io/DescriptorSystems.jl/dev/index.html)
"""
function DescriptorSystems.dss(sys::AbstractStateSpace)
    A,B,C,D = ssdata(sys)
    DescriptorSystems.dss(A, B, C, D; Ts = isdiscrete(sys) ? sys.Ts : 0)
end

function ControlSystems.ss(G::DescriptorSystems.DescriptorStateSpace)
    try
        G,r = DescriptorSystems.gss2ss(G)
        if r < size(G.A, 1)
            G = DescriptorSystems.dss2ss(G, fast=false)
        end
    catch
        G,_ = DescriptorSystems.dss2ss(G, simple_infeigs = false, fast=false)
    end
    G.Ts >= 0 ? ss(G.A, G.B, G.C, G.D) : ss(G.A, G.B, G.C, G.D, G.Ts)
end

"""
    n, ω = hinfnorm2(sys::LTISystem; kwargs...)

A numerically robust version of `hinfnorm` using DescriptorSystems.jl

For keyword arguments, see the docstring of `DescriptorSystems.ghinfnorm`, reproduced below
$(@doc(DescriptorSystems.ghinfnorm))
"""
function hinfnorm2(sys::LTISystem; kwargs...)
    DescriptorSystems.ghinfnorm(dss(ss(sys)); kwargs...)
end

"""
    n = h2norm(sys::LTISystem; kwargs...)

A numerically robust version of `norm` using DescriptorSystems.jl

For keyword arguments, see the docstring of `DescriptorSystems.gh2norm`, reproduced below
$(@doc(DescriptorSystems.gh2norm))
"""
function h2norm(sys::LTISystem; kwargs...)
    DescriptorSystems.gh2norm(dss(ss(sys)); kwargs...)
end

"""
    n, hsv = hankelnorm(sys::LTISystem; kwargs...)

Compute the hankelnorm and the hankel singular values

For keyword arguments, see the docstring of `DescriptorSystems.ghanorm`, reproduced below
$(@doc(DescriptorSystems.ghanorm))
"""
function hankelnorm(sys::LTISystem; kwargs...)
    DescriptorSystems.ghanorm(dss(ss(sys)); kwargs...)
end

"""
    nugap(sys0::LTISystem, sys1::LTISystem; kwargs...)

Compute the ν-gap metric between two systems. See also [`ncfmargin`](@ref).

For keyword arguments, see the docstring of `DescriptorSystems.gnugap`, reproduced below
$(@doc(DescriptorSystems.gnugap))
"""
function nugap(sys0::LTISystem, sys1::LTISystem; kwargs...)
    DescriptorSystems.gnugap(dss(ss(sys0)), dss(ss(sys1)); kwargs...)
end

const νgap = nugap

"""
    baltrunc2(sys::LTISystem; residual=false, n=missing, kwargs...)

Compute the a balanced truncation of order `n` and the hankel singular values

For keyword arguments, see the docstring of `DescriptorSystems.gbalmr`, reproduced below
$(@doc(DescriptorSystems.gbalmr))
"""
function baltrunc2(sys::LTISystem; residual=false, n=missing, kwargs...)
    sysr, hs = DescriptorSystems.gbalmr(dss(sys); matchdc=residual, ord=n, kwargs...)
    ss(sysr), hs
end

##

function Base.:\(G1::AbstractStateSpace, G2::AbstractStateSpace)
    g1,g2 = dss(G1), dss(G2)
    try
        ss(g1\g2)
    catch
        error("Failed to solve G1\\G2, the product might not be proper. To keep the product on descriptor form, use dss(G1)\\dss(G2)")
    end
end


text:
Base.@kwdef mutable struct JStruct
    JRhoW = nothing
    JWTh = nothing
    JWRho = nothing
    JThW = nothing
    JRhoWB = nothing
    JWThB = nothing
    JWRhoB = nothing
    JThWB = nothing
    JWW = nothing
end

function JacSchur!(J,U,CG,Param)
nz=Param.Grid.nz;
dz=Param.Grid.dz;
nCol=size(U,1);
nJ=nCol*nz;
nz=Param.Grid.nz;
RhoPos=Param.RhoPos;
uPos=Param.uPos;
vPos=Param.vPos;
wPos=Param.wPos;
ThPos=Param.ThPos;

Pres = Param.Cache1
D = Param.Cache2
Dp = Param.Cache2
Dm = Param.Cache3
dPdTh = Param.Cache4

Rho = view(U,:,:,RhoPos)
Th = view(U,:,:,ThPos)
if strcmp(Param.Thermo,"Energy")
  W=[zeros(size(U,1),1)
    ,U[:,:,wPos]];
  WC=0.5*(W[:,1:nz]+W[:,2:nz+1]);
  KE=0.5*(U[:,:,uPos].*U[:,:,uPos]+U[:,:,vPos].*U[:,:,vPos]+WC.*WC);
  Pres=Pressure(U[:,:,ThPos],U[:,:,RhoPos],KE,Param);
  PresCol=permute(Pres,[2 1]);
  PhiCol=Param.Grav*repmat(Param.Grid.zP,1,size(U,1));
else
  @views Pres .= Pressure(U[:,:,ThPos],U[:,:,ThPos],U[:,:,ThPos],Param);
end

@views D[:,1:nz-1] .= -0.5 .* (Rho[:,1:nz-1] .+ Rho[:,2:nz]) ./ dz;
@views J.JRhoW[1,:] .= reshape(PermutedDimsArray(D,(2,1)),nJ,1)
@views J.JRhoW[2,1:nJ-1] .= -reshape(PermutedDimsArray(D,(2,1)),nJ,1)[1:nJ-1]
if strcmp(Param.Thermo,"Energy")
  dPdTh .= (Param.Rd/Param.Cvd)./RhoCol;
else
  dPdTh .= dPresdTh(Th,Param);
end
@views Dp[:,1:nz-1] .= dPdTh[:,2:nz] ./  (0.5 .* (Rho[:,1:nz-1] .+ Rho[:,2:nz])) ./ dz;
@views Dm[:,1:nz-1] .= dPdTh[:,1:nz-1] ./ (0.5 .* (Rho[:,1:nz-1]+Rho[:,2:nz])) ./ dz;
@views J.JWTh[1,2:nJ] .= reshape(-PermutedDimsArray(Dp,(2,1)),nJ,1)[1:nJ-1]
@views J.JWTh[2,:] .= reshape(PermutedDimsArray(Dm,(2,1)),nJ,1)

if strcmp(Param.Thermo,"Energy")
# Dp[:,1:nz-1]=(Param.Rd/Param.Cvd)./
#   (0.5*(RhoCol[1:nz-1,:]+RhoCol[2:nz,:])).*
#   (PhiCol[2:nz,:])/dz;
# Dm[1:nz-1,:]=(Param.Rd/Param.Cvd)./
#   (0.5*(RhoCol[1:nz-1,:]+RhoCol[2:nz,:])).*
#   (PhiCol[1:nz-1,:])/dz;
# D[1:nz-1,:]=-0.5*Param.Grav./
#   (0.5*(RhoCol[1:nz-1,:]+RhoCol[2:nz,:]));
# J.JWRho=spdiags([[reshape(-Dm+D,nJ,1);0] [0;reshape(Dp+D,nJ,1)]],[0 1]
#   ,nJ,nJ);
else
  @views D[:,1:nz-1] .= 0.5 .* (Pres[:,2:nz] .- Pres[:,1:nz-1])./ dz ./
    (0.5 .* (Rho[:,1:nz-1] .+ Rho[:,2:nz])).^2;
  @views J.JWRho[1,2:nJ] .= reshape(PermutedDimsArray(D,(2,1)),nJ,1)[1:nJ-1]
  @views J.JWRho[2,1:nJ] .= reshape(PermutedDimsArray(D,(2,1)),nJ,1)
end

if strcmp(Param.Thermo,"Energy")
# D[1:nz-1,:]=-0.5*(ThCol[1:nz-1,:]+ThCol[2:nz,:] +
#     PresCol[1:nz-1,:]+PresCol[2:nz,:])/dz;
else
  @views D[:,1:nz-1] .= -0.5 .* (Th[:,1:nz-1] .+ Th[:,2:nz]) ./ dz;
end
@views J.JThW[1,:] .= reshape(PermutedDimsArray(D,(2,1)),nJ,1)
@views J.JThW[2,1:nJ-1] .= -reshape(PermutedDimsArray(D,(2,1)),nJ,1)[1:nJ-1]
if Param.Damping
  K=permute(DampingKoeff(CG,Param),[2 1]);
  J.JWW .= reshape(K,nJ,1)
end
end

function JacSchur(U,CG,Param)
nz=Param.Grid.nz;
dz=Param.Grid.dz;
nCol=size(U,1);
nJ=nCol*nz;
nz=Param.Grid.nz;
RhoPos=Param.RhoPos;
uPos=Param.uPos;
vPos=Param.vPos;
wPos=Param.wPos;
ThPos=Param.ThPos;
# W=[zeros(size(U,1),1)
#    ,U[:,:,wPos]];

RhoCol=permute(U[:,:,RhoPos],[2 1]);
ThCol=permute(U[:,:,ThPos],[2 1]);
if strcmp(Param.Thermo,"Energy")
  W=[zeros(size(U,1),1)
    ,U[:,:,wPos]];
  WC=0.5*(W[:,1:nz]+W[:,2:nz+1]);
  KE=0.5*(U[:,:,uPos].*U[:,:,uPos]+U[:,:,vPos].*U[:,:,vPos]+WC.*WC);
  Pres=Pressure(U[:,:,ThPos],U[:,:,RhoPos],KE,Param);
  PresCol=permute(Pres,[2 1]);
  PhiCol=Param.Grav*repmat(Param.Grid.zP,1,size(U,1));
else
  Pres=Pressure(U[:,:,ThPos],U[:,:,ThPos],U[:,:,ThPos],Param);
  PresCol=permute(Pres,[2 1]);
end

J = JStruct(;)
if strcmp(Param.ModelType,"Curl")
  D=zeros(nz,nCol);
  D[1:nz-1,:]=-0.5*(RhoCol[1:nz-1,:]+RhoCol[2:nz,:])/dz;
  J.JRhoW=spdiags([reshape(-D,nJ,1) reshape(D,nJ,1)],[-1 0]
    ,nJ,nJ);
# J.JRhoW=spdiagm(nJ, nJ,  0 => reshape(D,nJ), -1 => reshape(-D,nJ)[1:end-1])
  if strcmp(Param.Thermo,"Energy")
    dPdTh=(Param.Rd/Param.Cvd)./RhoCol;
  else
    dPdTh=dPresdTh(ThCol,Param);
  end
  Dp=zeros(nz,nCol);
  Dm=zeros(nz,nCol);
  Dp[1:nz-1,:]=dPdTh[2:nz,:] ./(0.5*(RhoCol[1:nz-1,:]+RhoCol[2:nz,:])) /dz;
  Dm[1:nz-1,:]=dPdTh[1:nz-1,:] ./(0.5*(RhoCol[1:nz-1,:]+RhoCol[2:nz,:])) /dz;
  J.JWTh=spdiags([[reshape(Dm,nJ,1);0] [0;reshape(-Dp,nJ,1)]],[0 1]
    ,nJ,nJ);

  if strcmp(Param.Thermo,"Energy")
    Dp=zeros(nz,nCol);
    Dm=zeros(nz,nCol);
    Dp[1:nz-1,:]=(Param.Rd/Param.Cvd)./
      (0.5*(RhoCol[1:nz-1,:]+RhoCol[2:nz,:])).*
      (PhiCol[2:nz,:])/dz;
    Dm[1:nz-1,:]=(Param.Rd/Param.Cvd)./
      (0.5*(RhoCol[1:nz-1,:]+RhoCol[2:nz,:])).*
      (PhiCol[1:nz-1,:])/dz;
    D[1:nz-1,:]=-0.5*Param.Grav./
      (0.5*(RhoCol[1:nz-1,:]+RhoCol[2:nz,:]));
    J.JWRho=spdiags([[reshape(-Dm+D,nJ,1);0] [0;reshape(Dp+D,nJ,1)]],[0 1]
      ,nJ,nJ);
  else
    D[1:nz-1,:]=0.5*(PresCol[2:nz,:]-PresCol[1:nz-1,:])/Param.Grid.dz ./
      (0.5*(RhoCol[1:nz-1,:]+RhoCol[2:nz,:])).^2;
    J.JWRho=spdiags([[reshape(D,nJ,1);0] [0;reshape(D,nJ,1)]],[0 1]
      ,nJ,nJ);
  end

  if strcmp(Param.Thermo,"Energy")
    D[1:nz-1,:]=-0.5*(ThCol[1:nz-1,:]+ThCol[2:nz,:] +
        PresCol[1:nz-1,:]+PresCol[2:nz,:])/dz;
  else
    D[1:nz-1,:]=-0.5*(ThCol[1:nz-1,:]+ThCol[2:nz,:])/dz;
  end
  J.JThW=spdiags([reshape(-D,nJ,1) reshape(D,nJ,1)],[-1 0]
    ,nJ,nJ);
  if Param.Damping
    K=permute(DampingKoeff(CG,Param),[2 1]);
    J.JWW=spdiags(reshape(K,nJ,1),0,nJ,nJ);
  end
else
  # Conservative
  D=zeros(nz,nCol);
  D[1:nz-1,:]=-1/dz;
  J.JRhoW=spdiags([reshape(-D,nJ,1) reshape(D,nJ,1)],[-1 0]
    ,nJ,nJ);

  dPdTh=dPresdTh(ThCol,Param);
  Dp=zeros(nz,nCol);
  Dm=zeros(nz,nCol);
  Dp[1:nz-1,:]=dPdTh[2:nz,:]/dz;
  Dm[1:nz-1,:]=dPdTh[1:nz-1,:]/dz;
  J.JWTh=spdiags([[reshape(Dm,nJ,1);0] [0;reshape(-Dp,nJ,1)]],[0 1]
    ,nJ,nJ);
  D[1:nz-1,:]=-0.5*Param.Grav;
  J.JWRho=spdiags([[reshape(D,nJ,1);0] [0;reshape(D,nJ,1)]],[0 1]
    ,nJ,nJ);
  Temp=ThCol./RhoCol;
  D[1:nz-1,:]=-0.5*(Temp[1:nz-1,:]+Temp[2:nz,:])/dz;
  J.JThW=spdiags([reshape(-D,nJ,1) reshape(D,nJ,1)],[-1 0]
    ,nJ,nJ);
end
return J
end


text:
using Test
using CUDA
using ChowLiuTrees

@testset "Pairwise marginals and mutual information for binary data" begin
    # Binary
    x_bit = Bool.([ 0  0  0  0
                1  0  1  1
                0  1  0  1
                1  1  0  0
                1  0  1  0
                1  0  1  0
                0  1  1  1
                0  0  1  1
                1  0  0  0
                0  1  1  1])

    mi_bit = pairwise_MI(x_bit, pseudocount=0.0)

    @test mi_bit[1,1] ≈ 0.6931471805599453
    @test mi_bit[1,2] ≈ 0.08630462173553435
    @test mi_bit[1,3] ≈ 0.0
    @test mi_bit[2,2] ≈ 0.6730116670092565
    @test mi_bit[2,3] ≈ 0.013844293808390695
    @test mi_bit[3,3] ≈ 0.6730116670092565

    data_types = if CUDA.functional()
        [CuMatrix{Bool}, Matrix{Bool}, Matrix{Int}, ] # CuArray{Int}
    else
        [Matrix{Bool}, Matrix{Int}]
    end

    for pseudocount in [0.0, 1.0, 100.0]
        for weights in [nothing, ones(10),
                       [1.0, 0.5, 1.0, 2.0, 0.7, 1.0, 0.5, 1.0, 1.3, 1.0]]
            for Float in [Float64, Float32]
                mar1 = pairwise_marginal(x_bit,
                        weights=weights, pseudocount=pseudocount, Float=Float)
                mi1 = pairwise_MI(x_bit,
                        weights=weights, pseudocount=pseudocount, Float=Float)
                for T in data_types
                    x2 = T(x_bit)
                    if T == Matrix{Int}
                        x2 .+= 1
                    end
                    mar2 = pairwise_marginal(x2;
                            weights=weights, pseudocount=pseudocount, Float=Float)
                    mi2 = pairwise_MI(x2;
                            weights=weights, pseudocount=pseudocount, Float=Float)

                    if T == CuMatrix{Bool}
                        mar2 = Array(mar2)
                        mi2 = Array(mi2)
                    end
                    if T == Matrix{Int}
                        @test all(mar2[:, :, 1, 1] .== mar1[:, :, 1])
                        @test all(mar2[:, :, 1, 2] .== mar1[:, :, 2])
                        @test all(mar2[:, :, 2, 1] .== mar1[:, :, 3])
                        @test all(mar2[:, :, 2, 2] .== mar1[:, :, 4])
                    else
                        @test all(mar1 .== mar2)
                    end
                    @test all(isapprox(mi1, mi2; atol=1e-5))
                end
            end
        end
    end
end

@testset "Pairwise marginals and mutual information for categorical data" begin
    # Binary
    x = [ 2  1  3  3
        1  3  3  2
        1  2  2  3
        3  2  2  2
        3  3  2  2
        1  3  3  2
        2  3  2  3
        1  2  3  3
        3  1  2  3
        1  2  3  1]

    m = pairwise_marginal(x, pseudocount=0.0)
    mi = pairwise_MI(x, pseudocount=0.0)

    @test mi[1,1] ≈ 1.0296530140645737
    @test mi[1,2] ≈ 0.2502012117690939
    @test mi[1,3] ≈ 0.30431653267886216
    @test mi[2,2] ≈ 1.054920167986144
    @test mi[2,3] ≈ 0.0
    @test mi[2,4] ≈ 0.3025260261455486

    # TODO: compare with GPU
end


text:
#=
Data aggregation of any kind
=#
function nomissing(da::AbstractArray{Union{T,Missing},N}) where {T,N}
    any(ismissing, da) && error("array contains missing values")
    return Array{T,N}(da)
end
nomissing(a::AbstractArray, args...) = a
function nomissing(da::Array{Union{T,Missing},N}, value) where {T,N}
    return replace(da, missing => T(value))
end
nomissing(da::AbDimArray) = DimensionalData.rebuild(da, nomissing(da.data))
nomissing(da::ClimArray) = ClimArray(nomissing(da.data), da.dims, da.refdims, da.name, da.attrib)

#########################################################################
# Aggregation of data, dropagg missings, dimensions, etc.
#########################################################################
# TODO: Define dropagg with weights as well (simple to do, same as timeagg but
# without the complexity of temporal averaging)

export dropagg, nomissing, collapse, drop

"""
    dropagg(f, A, dims)
Apply aggregating function `f` (e.g. `sum`) on array `A` across dimension(s) `dims` and drop
the corresponding dimension(s) from the result (Julia inherently keeps singleton dimensions).

If `A` is one dimensional, `dropagg` will return the single number of applying `f(A)`.
"""
function dropagg(f, A, dims)
    length(size(A)) == 1 && return f(A)
    return dropdims(f(A; dims = dims); dims = dims)
end

"""
    collapse(f, A, dim)
Reduce `A` towards dimension `dim` using the collapsing function `f` (e.g. `mean`).
This means that `f` is applied across all other dimensions of `A`, each of which are
subsequently dropped, leaving only the collapsed result of `A` vs. the remaining dimension.
"""
function collapse(f, A, dim)
    di = dimindex(A, dim)
    dimsA = dims(A)
    touse = setdiff(eachindex(dimsA), di)
    return dropagg(f, A; dims = ntuple(i -> dimsA[touse[i]], length(touse)))
end

dimindex(A, i::Int) = i
# this method is necessary because of "reshaping" happening
# in DimensionalData.jl...
function dropagg(f, A::AbDimArray, dims)
    length(size(A)) == 1 && return f(A)
    r = dropdims(f(A; dims = dims); dims = dims)
    DimensionalData.rebuild(r, Array(r.data))
end

dimindex(A::AbDimArray, Dim) = DimensionalData.dimnum(A, Dim)

#########################################################################
# Other dimensions
#########################################################################
export otherdims
export otheridxs

"""
    otheridxs(A::ClimArray, Dim)
Return an iterator of indices, that when used can access all indices of `A` *except* those
belonging to dimension(s) `Dim`.

For example, if `A` has dims `(Lon, Lat, Time)` you can get all timeseries of `A`:
```julia
for i in otheridxs(A, Time())
    x = A[i...] # this is a timeseries (Vector) for each lon × lat combination
end
```
or all time+latitude slices with
```julia
for i in otheridxs(A, (Time(), Lat()))
    x = A[i...] # matrix of time × latitude slice for each longitude
end
```
(notice that splatting `i...` is necessary because `otheridxs` generates tuples)
"""
function otheridxs(A, D)
    z = otherdims(A, D)
    az = DimensionalData.basetypeof.(z)
    iters = [(Dim(i) for i in 1:size(A, Dim)) for Dim in az]
    return Iterators.product(iters...)
end

# This is the version @rafaqz suggested, but has worse performance...
# otheridxs(A, D) = map(identity, DimensionalData.dimwise_generators(otherdims(A, D)))

#########################################################################
# Dimensionwise
#########################################################################
#=
"""
    dimwise(f, A, B)
Apply bivariate function `f` (e.g. multiplication `*`) across the (first) matching dimension
of `A` and `B`, where `B` has only one dimension.
For normal Julia arrays match is done on the first matching dimension size.

For example, if `A` has dimensions (Lon, Lat, Time) and `B` has dimension (Lat,)
(and of course both Lat are actually the same thing),
then the result would be an array `C` with same structure as `A` but with values
```julia
@. C[:, i, :] = f(A[:, i, :], B[i])
```
for every `i` across the matching dimension.

Useful when wanting to weight a field `A` based on its latitude or time.
"""
function dimwise(f, A::AbDimArray, B::AbDimArray)
    @assert length(size(B)) == 1 "3rd argument of dimwise must be a vector"
    m = find_matching_dim(A, B)
    # Array(data) necessary because reshapes somewhow fail to play well with the code...
    C = dimwise(f, Array(data(A)), Array(data(B)), m)
    return DimensionalData.basetypeof(A)(C, dims(A))
end

dimwise(f, A::AbstractArray, B, m::Int = find_matching_dim(A, B)) =
dimwise!(copy(A), f, A, B, m)

function dimwise!(C, f, A::AbstractArray{T, N}, B::AbstractVector, m::Int) where {T, N}
    shape = ntuple(i -> (i==m) ? length(B) : 1, N)
    b = reshape(B, shape)
    C .= f.(A, b)
    return C
end

function find_matching_dim(A::AbDimArray, B::AbDimArray)::Int
    dimB = dims(B)[1]
    m = findfirst(isequal(dimB), dims(A))
    isnothing(m) && error("A and B have no matching dimensions.")
    return m
end

function find_matching_dim(A::AbstractArray, B::AbstractVector)::Int
    m = findfirst(isequal(length(B)), size(A))
    isnothing(m) && error("A and B have no matching dimensions.")
    return m
end

=#
export dimwise


text:
import PowerSimulations:
    open_store,
    HdfSimulationStore,
    HDF_FILENAME,
    SimulationStoreParams,
    SimulationStoreProblemParams,
    SimulationStoreProblemRequirements,
    CacheFlushRules,
    KiB,
    MiB,
    GiB,
    STORE_CONTAINER_VARIABLES,
    CachePriority,
    initialize_problem_storage!,
    add_rule!,
    write_result!,
    read_result,
    has_dirty,
    get_cache_hit_percentage

function _initialize!(store, sim, variables, stage_defs, cache_rules)
    stages = OrderedDict{Symbol, SimulationStoreProblemParams}()
    stage_reqs = Dict{Symbol, SimulationStoreProblemRequirements}()
    num_param_containers = 0
    for stage in keys(stage_defs)
        execution_count = stage_defs[stage]["execution_count"]
        horizon = stage_defs[stage]["horizon"]
        num_rows = execution_count * sim["num_steps"]

        stage_params = SimulationStoreProblemParams(
            execution_count,
            horizon,
            stage_defs[stage]["interval"],
            stage_defs[stage]["resolution"],
            stage_defs[stage]["end_of_interval_step"],
            stage_defs[stage]["base_power"],
            stage_defs[stage]["system_uuid"],
        )
        reqs = SimulationStoreProblemRequirements()

        for (name, array) in stage_defs[stage]["variables"]
            reqs.variables[name] = Dict(
                "columns" => stage_defs[stage]["names"],
                "dims" => (horizon, length(stage_defs[stage]["names"]), num_rows),
            )
            keep_in_cache = variables[name]["keep_in_cache"]
            cache_priority = variables[name]["cache_priority"]
            add_rule!(
                cache_rules,
                stage,
                STORE_CONTAINER_VARIABLES,
                name,
                keep_in_cache,
                cache_priority,
            )
        end

        stages[stage] = stage_params
        stage_reqs[stage] = reqs
        num_param_containers += length(reqs.variables)
    end

    params = SimulationStoreParams(
        sim["initial_time"],
        sim["step_resolution"],
        sim["num_steps"],
        stages,
    )
    initialize_problem_storage!(store, params, stage_reqs, cache_rules)
end

function _run_sim_test(path, sim, variables, stage_defs, cache_rules, seed)
    rng = MersenneTwister(seed)
    type = STORE_CONTAINER_VARIABLES
    open_store(HdfSimulationStore, path, "w") do store
        sim_time = sim["initial_time"]
        _initialize!(store, sim, variables, stage_defs, cache_rules)
        for step in 1:sim["num_steps"]
            for stage in keys(stage_defs)
                stage_time = sim_time
                for i in 1:stage_defs[stage]["execution_count"]
                    for name in keys(variables)
                        data = rand(rng, size(stage_defs[stage]["variables"][name])...)
                        write_result!(store, stage, type, name, stage_time, data)
                        columns = stage_defs[stage]["names"]
                        _verify_data(data, store, stage, type, name, stage_time, columns)
                    end

                    stage_time += stage_defs[stage]["resolution"]
                end
            end

            sim_time += sim["step_resolution"]
        end

        for param_cache in values(store.cache.data)
            @test get_cache_hit_percentage(param_cache) == 100.0
        end

        flush(store)
        @test !has_dirty(store.cache)
    end
end

function _verify_read_results(path, sim, variables, stage_defs, seed)
    rng = MersenneTwister(seed)
    type = STORE_CONTAINER_VARIABLES
    open_store(HdfSimulationStore, path, "r") do store
        sim_time = sim["initial_time"]
        for step in 1:sim["num_steps"]
            for stage in keys(stage_defs)
                stage_time = sim_time
                for i in 1:stage_defs[stage]["execution_count"]
                    for name in keys(variables)
                        data = rand(rng, size(stage_defs[stage]["variables"][name])...)
                        columns = stage_defs[stage]["names"]
                        _verify_data(data, store, stage, type, name, stage_time, columns)
                    end

                    stage_time += stage_defs[stage]["resolution"]
                end
            end

            sim_time += sim["step_resolution"]
        end

        for param_cache in values(store.cache.data)
            @test get_cache_hit_percentage(param_cache) == 0.0
        end
    end
end

function _verify_data(expected, store, stage, type, name, time, columns)
    expected_df = DataFrames.DataFrame(expected, columns)
    df = read_result(DataFrames.DataFrame, store, stage, type, name, time)
    @test expected_df == df

    # TODO read_result with JuMP.Containers.DenseAxisArray
end

@testset "Test SimulationStore 2-d arrays" begin
    sim = Dict(
        "initial_time" => Dates.DateTime("2020-01-01T00:00:00"),
        "step_resolution" => Dates.Hour(24),
        "num_steps" => 2,
    )
    variables = Dict(
        :P__ThermalStandard =>
            Dict("cache_priority" => CachePriority.HIGH, "keep_in_cache" => true),
        :P__InterruptibleLoad =>
            Dict("cache_priority" => CachePriority.LOW, "keep_in_cache" => false),
    )
    stage_defs = OrderedDict(
        :ED => Dict(
            "execution_count" => 24,
            "horizon" => 12,
            "names" => [:dev1, :dev2, :dev3, :dev4, :dev5],
            "variables" => Dict(x => ones(12, 5) for x in keys(variables)),
            "interval" => Dates.Hour(1),
            "resolution" => Dates.Hour(1),
            "end_of_interval_step" => 1,
            "base_power" => 100.0,
            "system_uuid" => Base.UUID("4076af6c-e467-56ae-b986-b466b2749572"),
        ),
        :UC => Dict(
            "execution_count" => 1,
            "horizon" => 24,
            "names" => [:dev1, :dev2, :dev3],
            "variables" => Dict(x => ones(24, 3) for x in keys(variables)),
            "interval" => Dates.Hour(1),
            "resolution" => Dates.Hour(24),
            "end_of_interval_step" => 12,
            "base_power" => 100.0,
            "system_uuid" => Base.UUID("4076af6c-e467-56ae-b986-b466b2749572"),
        ),
    )
    cache_rules = CacheFlushRules(max_size = 1 * MiB, min_flush_size = 4 * KiB)

    path = mktempdir()
    # Use this seed to produce the same randomly generated arrays for write and verify.
    seed = 1234
    _run_sim_test(path, sim, variables, stage_defs, cache_rules, seed)
    _verify_read_results(path, sim, variables, stage_defs, seed)
end

@testset "Test ParamResultCache" begin
    key = PSI.ParamCacheKey((:ED, :variables, :var))
    cache = PSI.ParamResultCache(key, PSI.CacheFlushRule())
    @test !PSI.has_clean(cache)
    @test !PSI.is_dirty(cache, Dates.now())

    timestamp1 = Dates.DateTime("2020-01-01T00:00:00")
    timestamp2 = Dates.DateTime("2020-01-01T01:00:00")
    timestamp3 = Dates.DateTime("2020-01-01T02:00:00")
    PSI.add_result!(cache, timestamp1, ones(2), false)
    @test PSI.is_dirty(cache, timestamp1)
    PSI.add_result!(cache, timestamp2, ones(2), false)
    @test PSI.is_dirty(cache, timestamp2)

    @test length(cache.data) == 2
    @test length(cache.dirty_timestamps) == 2

    popfirst!(cache.dirty_timestamps)
    @test !PSI.is_dirty(cache, timestamp1)
    @test PSI.has_clean(cache)
    PSI.add_result!(cache, timestamp3, ones(2), true)
    @test length(cache.data) == 2

    empty!(cache)
    @test isempty(cache.data)
    @test isempty(cache.dirty_timestamps)

    PSI.add_result!(cache, timestamp1, ones(2), false)
    PSI.add_result!(cache, timestamp2, ones(2), false)
    PSI.discard_results!(cache, [timestamp1, timestamp2])
    @test isempty(cache.data)
end

# TODO: test optimizer stats
# TODO: unit tests of individual functions, size checks
# TODO: 3-d arrays
# TODO: profiling of memory performance and GC


text:
using Test
import UUIDs
import Pluto: PlutoRunner, Notebook, WorkspaceManager, Cell, ServerSession, ClientSession, update_run!

@testset "Macro analysis" begin
    🍭 = ServerSession()
    🍭.options.evaluation.workspace_use_distributed = false

    fakeclient = ClientSession(:fake, nothing)
    🍭.connected_clients[fakeclient.id] = fakeclient

    @testset "Base macro call" begin
        notebook = Notebook([
            Cell("@enum Fruit 🍎 🍐"),
            Cell("my_fruit = 🍎"),
            Cell("jam(fruit::Fruit) = cook(fruit)"),
        ])
        cell(idx) = notebook.cells[idx]

        update_run!(🍭, notebook, notebook.cells)

        @test cell(1) |> noerror
        @test [:🍎, :🍐] ⊆ notebook.topology.nodes[cell(1)].definitions
        @test :Fruit ∈ notebook.topology.nodes[cell(1)].funcdefs_without_signatures
        @test Symbol("@enum") ∈ notebook.topology.nodes[cell(1)].references

        @test cell(2) |> noerror
        @test :🍎 ∈ notebook.topology.nodes[cell(2)].references

        @test cell(3) |> noerror
        @test :Fruit ∈ notebook.topology.nodes[cell(3)].references
    end

    @testset "User defined macro 1" begin
        notebook = Notebook([
            Cell("""macro my_assign(sym, val)
              :(\$(esc(sym)) = \$(val))
            end"""),
            Cell("@my_assign x 1+1"),
        ])
        cell(idx) = notebook.cells[idx]

        update_run!(🍭, notebook, notebook.cells)

        @test :x ∈ notebook.topology.nodes[cell(2)].definitions
        @test Symbol("@my_assign") ∈ notebook.topology.nodes[cell(2)].references

        update_run!(🍭, notebook, notebook.cells)

        # Works on second time because of old workspace
        @test :x ∈ notebook.topology.nodes[cell(2)].definitions
        @test Symbol("@my_assign") ∈ notebook.topology.nodes[cell(2)].references
    end

    @testset "User defined macro 2" begin
        notebook = Notebook([
            Cell("@my_identity(f(123))"),
            Cell(""),
            Cell(""),
        ])
        cell(idx) = notebook.cells[idx]

        update_run!(🍭, notebook, notebook.cells)

        setcode(cell(2), """macro my_identity(expr)
            esc(expr)
        end""")
        update_run!(🍭, notebook, cell(2))

        setcode(cell(3), "f(x) = x")
        update_run!(🍭, notebook, cell(3))

        @test cell(1) |> noerror
        @test cell(2) |> noerror
        @test cell(3) |> noerror
        @test cell(1).output.body == "123"

        update_run!(🍭, notebook, cell(1))
        @test cell(1) |> noerror
        @test cell(2) |> noerror
        @test cell(3) |> noerror
    end

    @testset "User defined macro 3" begin
        notebook = Notebook([
            Cell("""
            macro mymap()
                quote
                    [1, 2, 3] .|> sqrt
                end
            end
            """),
            Cell("@mymap")
        ])
        cell(idx) = notebook.cells[idx]

        update_run!(🍭, notebook, notebook.cells)

        @test cell(1) |> noerror
        @test cell(2) |> noerror

        update_run!(🍭, notebook, cell(1))

        @test cell(2) |> noerror
    end

    @testset "User defined macro 4" begin
        notebook = Notebook([
            Cell("""macro my_assign(ex)
                esc(ex)
            end"""),
            Cell("@macroexpand @my_assign 1+1"),
        ])
        cell(idx) = notebook.cells[idx]

        update_run!(🍭, notebook, notebook.cells)

        @test Symbol("@my_assign") ∈ notebook.topology.nodes[cell(2)].references
    end

    @testset "User defined macro 5" begin
        notebook = Notebook([
            Cell("""macro dynamic_values(ex)
                [:a, :b, :c]
            end"""),
            Cell("myarray = @dynamic_values()"),
        ])
        references(idx) = notebook.topology.nodes[notebook.cells[idx]].references

        update_run!(🍭, notebook, notebook.cells)

        @test :a ∉ references(2)
        @test :b ∉ references(2)
        @test :c ∉ references(2)
    end

    @testset "User defined macro 6" begin
        notebook = Notebook([
            Cell("""macro my_macro()
                esc(:(y + x))
            end"""),
            Cell("""function my_function()
                @my_macro()
            end"""),
            Cell("my_function()"),
            Cell("x = 1"),
            Cell("y = 2"),
        ])
        cell(idx) = notebook.cells[idx]

        update_run!(🍭, notebook, notebook.cells)

        @test [Symbol("@my_macro"), :x, :y] ⊆ notebook.topology.nodes[cell(2)].references
        @test cell(3).output.body == "3"
    end

    @testset "Function docs" begin
        notebook = Notebook([
            Cell("""
                "my function doc"
                f(x) = 2x
            """),
            Cell("f"),
        ])
        cell(idx) = notebook.cells[idx]

        update_run!(🍭, notebook, notebook.cells)

        @test :f ∈ notebook.topology.nodes[cell(1)].funcdefs_without_signatures
        @test :f ∈ notebook.topology.nodes[cell(2)].references
    end

    @testset "Expr sanitization" begin
        struct A; end
        f(x) = x
        unserializable_expr = Expr(:call, f, A(), A[A(), A(), A()], PlutoRunner, PlutoRunner.sanitize_expr)

        get_expr_types(other) = typeof(other)
        get_expr_types(ex::Expr) = get_expr_types.(ex.args)

        flatten(x, acc=[]) = push!(acc, x)
        function flatten(arr::AbstractVector, acc=[]) foreach(x -> flatten(x, acc), arr); acc end

        sanitized_expr = PlutoRunner.sanitize_expr(unserializable_expr)
        types = sanitized_expr |> get_expr_types |> flatten |> Set

        # Checks that no fancy type is part of the serialized expression
        @test Set([Symbol, QuoteNode]) == types

        @test Meta.isexpr(sanitized_expr.args[3], :vect, 3)
        @test sanitized_expr.args[2] == :A
        @test sanitized_expr.args[1] == :(Main.f)
    end

    @testset "Package macro 1" begin
        notebook = Notebook([
            Cell("using Dates"),
            Cell("df = dateformat\"Y-m-d\""),
        ])
        cell(idx) = notebook.cells[idx]

        update_run!(🍭, notebook, cell(2))

        @test cell(2).errored == true
        @test occursinerror("UndefVarError: @dateformat_str", cell(2)) == true

        update_run!(🍭, notebook, notebook.cells)

        @test cell(1) |> noerror
        @test cell(2) |> noerror

        notebook = Notebook([
            Cell("using Dates"),
            Cell("df = dateformat\"Y-m-d\""),
        ])
        update_run!(🍭, notebook, notebook.cells)

        @test cell(1) |> noerror
        @test cell(2) |> noerror
    end

    @testset "Package macro 2" begin
        🍭.options.evaluation.workspace_use_distributed = true

        notebook = Notebook([
            Cell("z = x^2 + y"),
            Cell("@variables x y"),
            Cell("""
            begin
                import Pkg
                Pkg.activate(mktempdir())
                Pkg.add(Pkg.PackageSpec(name="Symbolics", version="1"))
                using Symbolics
            end
            """),
        ])
        cell(idx) = notebook.cells[idx]

        update_run!(🍭, notebook, notebook.cells)

        @test cell(1) |> noerror
        @test cell(2) |> noerror
        @test cell(3) |> noerror

        update_run!(🍭, notebook, cell(2))

        @test cell(1) |> noerror
        @test cell(2) |> noerror

        setcode(cell(2), "@variables 🐰 y")
        update_run!(🍭, notebook, cell(2))

        @test cell(1).errored
        @test cell(2) |> noerror

        setcode(cell(1), "z = 🐰^2 + y")
        update_run!(🍭, notebook, cell(1))

        @test cell(1) |> noerror
        @test cell(2) |> noerror

        WorkspaceManager.unmake_workspace((🍭, notebook))

        🍭.options.evaluation.workspace_use_distributed = false
    end

    @testset "Previous workspace for unknowns" begin
        notebook = Notebook([
            Cell("""macro my_identity(expr)
              expr
            end"""),
            Cell("(@__MODULE__, (@my_identity 1 + 1))"),
            Cell("@__MODULE__"),
        ])
        cell(idx) = notebook.cells[idx]

        update_run!(🍭, notebook, cell(1))
        update_run!(🍭, notebook, notebook.cells[2:end])

        @test cell(1) |> noerror
        @test cell(2) |> noerror
        @test cell(3) |> noerror

        module_from_cell2 = cell(2).output.body[:elements][1][2][1]
        module_from_cell3 = cell(3).output.body

        @test module_from_cell2 == module_from_cell3
    end

    @testset "Definitions" begin
        notebook = Notebook([
            Cell("""macro my_assign(sym, val)
                :(\$(esc(sym)) = \$(val))
            end"""),
            Cell("c = :hello"),
            Cell("@my_assign b c"),
            Cell("b"),
        ])
        cell(idx) = notebook.cells[idx]

        update_run!(🍭, notebook, notebook.cells)
        update_run!(🍭, notebook, notebook.cells)

        @test ":hello" == cell(3).output.body
        @test ":hello" == cell(4).output.body
        @test :b ∈ notebook.topology.nodes[cell(3)].definitions
        @test [:c, Symbol("@my_assign")] ⊆ notebook.topology.nodes[cell(3)].references
        @test cell(3).cell_dependencies.contains_user_defined_macrocalls == true

        setcode(notebook.cells[2], "c = :world")
        update_run!(🍭, notebook, cell(2))

        @test ":world" == cell(3).output.body
        @test ":world" == cell(4).output.body
    end

    @testset "Macros using import" begin
        notebook = Notebook(Cell.([
            """
            @option "option_a" struct OptionA
                option_value::option_type
            end
            """,
            "option_type = String",
            "import Configurations: @option",
        ]))
        cell(idx) = notebook.cells[idx]

        update_run!(🍭, notebook, notebook.cells)

        @test :option_type ∈ notebook.topology.nodes[cell(1)].references
        @test cell(1) |> noerror
    end
end