text:
using Test
using Pkg, Pkg.PlatformEngines, Pkg.BinaryPlatforms, Pkg.Artifacts
using BinaryBuilderBase
using BinaryBuilderBase: abi_agnostic, get_concrete_platform, march

@testset "Supported Platforms" begin
    all = supported_platforms()
    opt_out_specific = supported_platforms(exclude=[Linux(:x86_64, libc=:glibc)])
    islin(x) = typeof(x) == Linux
    opt_out_fx = supported_platforms(exclude=islin)

    @test length(all) == length(opt_out_specific)+1
    @test any(opt_out_specific .== [Linux(:i686 , libc=:glibc)])
    @test !any(opt_out_fx .== [Linux(:i686 , libc=:glibc)])

    @test sort([Windows(:x86_64), Linux(:i686, libc=:musl), Linux(:i686, libc=:glibc), MacOS(:x86_64)],
               by = triplet) ==
                   [Linux(:i686, libc=:glibc), Linux(:i686, libc=:musl), MacOS(:x86_64), Windows(:x86_64)]
end

@testset "ExtendedPlatform" begin
    p = ExtendedPlatform(Linux(:x86_64; libc=:musl); march = "avx", cuda = "9.2")
    @test p.p == Linux(:x86_64, libc=:musl)
    @test p.ext == Dict("march" => "avx", "cuda" => "9.2")
    @test BinaryPlatforms.platform_name(p) == "ExtendedPlatform"
    @test BinaryPlatforms.arch(p) == :x86_64
    @test BinaryPlatforms.libc(p) == :musl
    @test BinaryPlatforms.call_abi(p) == nothing
    @test BinaryPlatforms.compiler_abi(p) == CompilerABI()
    @test BinaryPlatforms.triplet(p) == "x86_64-linux-musl-cuda+9.2-march+avx"
    @test abi_agnostic(p) == Linux(:x86_64, libc=:musl)
    @test aatriplet(p) == "x86_64-linux-musl"
    @test replace_cxxstring_abi(p, :cxx03) == ExtendedPlatform(Linux(:x86_64; libc=:musl, compiler_abi=CompilerABI(; cxxstring_abi=:cxx03)); march = "avx", cuda = "9.2")
    @test replace_libgfortran_version(p, v"4") == ExtendedPlatform(Linux(:x86_64; libc=:musl, compiler_abi=CompilerABI(; libgfortran_version=v"4")); march = "avx", cuda = "9.2")
    # The concrete platform has C++11 strinb ABI because march=avx requires GCC v5
    @test get_concrete_platform(p; compilers = [:c], preferred_gcc_version = v"4.8", preferred_llvm_version = v"9") == ExtendedPlatform(Linux(:x86_64, libc=:musl, compiler_abi=CompilerABI(libgfortran_version=v"3", cxxstring_abi=:cxx11)); march="avx", cuda="9.2")
    @test ExtendedPlatform(p; cuda="9.2", foo="bar") == ExtendedPlatform(Linux(:x86_64, libc=:musl); march="avx", cuda="9.2", foo="bar")
    @test ExtendedPlatform(p; foo="bar") == ExtendedPlatform(Linux(:x86_64, libc=:musl); march="avx", cuda="9.2", foo="bar")
    @test_throws ErrorException ExtendedPlatform(p; cuda="10.1")
    # We have to split in this way the test on the representation of the type
    # because it may depend on the order with which the keys are extracted from
    # the dictionary
    @test startswith(repr(p), "ExtendedPlatform(Linux(:x86_64, libc=:musl);")
    @test occursin("march=\"avx\"", repr(p))
    @test occursin("cuda=\"9.2\"", repr(p))
    @test endswith(repr(p), ")")
    # Make sure the round trip works
    @test parse(ExtendedPlatform, triplet(p)) == p
    meta = Dict{String,String}()
    Artifacts.pack_platform!(meta, p)
    @test meta == Dict("arch" => "x86_64","libc" => "musl","march" => "avx","os" => "linux", "cuda" => "9.2")

    p = ExtendedPlatform(Linux(:aarch64, libc=:glibc, compiler_abi=CompilerABI(libgfortran_version=v"5", cxxstring_abi=:cxx03)); march="carmel", cuda="10.1", cuda_capability="52")
    @test p.p == Linux(:aarch64, libc=:glibc, compiler_abi=CompilerABI(libgfortran_version=v"5", cxxstring_abi=:cxx03))
    @test p.ext == Dict("march" => "carmel","cuda_capability" => "52","cuda" => "10.1")
    @test BinaryPlatforms.platform_name(p) == "ExtendedPlatform"
    @test BinaryPlatforms.arch(p) == :aarch64
    @test BinaryPlatforms.libc(p) == :glibc
    @test BinaryPlatforms.call_abi(p) == nothing
    @test BinaryPlatforms.compiler_abi(p) == CompilerABI(libgfortran_version=v"5", cxxstring_abi=:cxx03)
    @test BinaryPlatforms.triplet(p) == "aarch64-linux-gnu-libgfortran5-cxx03-cuda+10.1-cuda_capability+52-march+carmel"
    @test replace_cxxstring_abi(p, :cxx11) == ExtendedPlatform(Linux(:aarch64, libc=:glibc, compiler_abi=CompilerABI(libgfortran_version=v"5", cxxstring_abi=:cxx11)); march="carmel", cuda="10.1", cuda_capability="52")
    @test replace_libgfortran_version(p, v"3") == ExtendedPlatform(Linux(:aarch64, libc=:glibc, compiler_abi=CompilerABI(libgfortran_version=v"3", cxxstring_abi=:cxx03)); march="carmel", cuda="10.1", cuda_capability="52")
    @test abi_agnostic(p) == Linux(:aarch64, libc=:glibc)
    @test aatriplet(p) == "aarch64-linux-gnu"
    @test get_concrete_platform(p; compilers = [:c, :go], preferred_gcc_version = v"8", preferred_llvm_version = v"6") == ExtendedPlatform(Linux(:aarch64, libc=:glibc, compiler_abi=CompilerABI(libgfortran_version=v"5", cxxstring_abi=:cxx03)); march="carmel", cuda_capability="52", cuda="10.1")
    @test ExtendedPlatform(p; cuda="10.1", foo="bar") == ExtendedPlatform(Linux(:aarch64, libc=:glibc, compiler_abi=CompilerABI(libgfortran_version=v"5.0.0", cxxstring_abi=:cxx03)); march="carmel", cuda_capability="52", cuda="10.1", foo="bar")
    @test ExtendedPlatform(p; foo="bar") == ExtendedPlatform(Linux(:aarch64, libc=:glibc, compiler_abi=CompilerABI(libgfortran_version=v"5.0.0", cxxstring_abi=:cxx03)); march="carmel", cuda_capability="52", cuda="10.1", foo="bar")
    @test_throws ErrorException ExtendedPlatform(p; cuda="10.1", cuda_capability="80")
    @test startswith(repr(p), "ExtendedPlatform(Linux(:aarch64, libc=:glibc, compiler_abi=CompilerABI(libgfortran_version=v\"5.0.0\", cxxstring_abi=:cxx03)); ")
    @test occursin("march=\"carmel\"", repr(p))
    @test occursin("cuda=\"10.1\"", repr(p))
    @test occursin("cuda_capability=\"52\"", repr(p))
    @test endswith(repr(p), ")")
    @test parse(ExtendedPlatform, triplet(p)) == p
    meta = Dict{String,String}()
    Artifacts.pack_platform!(meta, p)
    @test meta == Dict("libgfortran_version" => "5.0.0","arch" => "aarch64","libc" => "glibc","os" => "linux","cxxstring_abi" => "cxx03", "march" => "carmel", "cuda" => "10.1", "cuda_capability" => "52")

    p = ExtendedPlatform(Linux(:aarch64, libc=:glibc, compiler_abi=CompilerABI(libgfortran_version=v"5", cxxstring_abi=:cxx11)); march="carmel", cuda="10.1", cuda_capability="52")
    @test p.p == Linux(:aarch64, libc=:glibc, compiler_abi=CompilerABI(libgfortran_version=v"5", cxxstring_abi=:cxx11))
    @test p.ext == Dict("march" => "carmel","cuda_capability" => "52","cuda" => "10.1")
    @test BinaryPlatforms.platform_name(p) == "ExtendedPlatform"
    @test BinaryPlatforms.arch(p) == :aarch64
    @test BinaryPlatforms.libc(p) == :glibc
    @test BinaryPlatforms.call_abi(p) == nothing
    @test BinaryPlatforms.compiler_abi(p) == CompilerABI(libgfortran_version=v"5", cxxstring_abi=:cxx11)
    @test BinaryPlatforms.triplet(p) == "aarch64-linux-gnu-libgfortran5-cxx11-cuda+10.1-cuda_capability+52-march+carmel"
    @test abi_agnostic(p) == Linux(:aarch64, libc=:glibc)
    @test aatriplet(p) == "aarch64-linux-gnu"
    @test replace_cxxstring_abi(p, :cxx11) == p
    @test replace_libgfortran_version(p, v"5") == p
    @test get_concrete_platform(p; compilers = [:c], preferred_gcc_version = v"5", preferred_llvm_version = v"9") == ExtendedPlatform(Linux(:aarch64, libc=:glibc, compiler_abi=CompilerABI(libgfortran_version=v"5", cxxstring_abi=:cxx11)); march="carmel", cuda_capability="52", cuda="10.1")
    @test ExtendedPlatform(p; cuda="10.1", foo="bar") == ExtendedPlatform(Linux(:aarch64, libc=:glibc, compiler_abi=CompilerABI(libgfortran_version=v"5.0.0", cxxstring_abi=:cxx11)); march="carmel", cuda_capability="52", cuda="10.1", foo="bar")
    @test ExtendedPlatform(p; foo="bar") == ExtendedPlatform(Linux(:aarch64, libc=:glibc, compiler_abi=CompilerABI(libgfortran_version=v"5.0.0", cxxstring_abi=:cxx11)); march="carmel", cuda_capability="52", cuda="10.1", foo="bar")
    @test_throws ErrorException ExtendedPlatform(p; cuda="10.1", cuda_capability="80")
    @test startswith(repr(p), "ExtendedPlatform(Linux(:aarch64, libc=:glibc, compiler_abi=CompilerABI(libgfortran_version=v\"5.0.0\", cxxstring_abi=:cxx11)); ")
    @test occursin("march=\"carmel\"", repr(p))
    @test occursin("cuda=\"10.1\"", repr(p))
    @test occursin("cuda_capability=\"52\"", repr(p))
    @test endswith(repr(p), ")")
    @test parse(ExtendedPlatform, triplet(p)) == p
    meta = Dict{String,String}()
    Artifacts.pack_platform!(meta, p)
    @test meta == Dict("libgfortran_version" => "5.0.0","arch" => "aarch64","libc" => "glibc","os" => "linux","cxxstring_abi" => "cxx11", "march" => "carmel", "cuda" => "10.1", "cuda_capability" => "52")

    p = ExtendedPlatform(Linux(:armv7l, libc=:glibc, call_abi=:eabihf, compiler_abi=CompilerABI(; libstdcxx_version=v"3.4.24")); march="armv7l", cuda="11.1")
    @test p.p == Linux(:armv7l, libc=:glibc, call_abi=:eabihf, compiler_abi=CompilerABI(; libstdcxx_version=v"3.4.24"))
    @test p.ext == Dict("march" => "armv7l","cuda" => "11.1")
    @test BinaryPlatforms.platform_name(p) == "ExtendedPlatform"
    @test BinaryPlatforms.arch(p) == :armv7l
    @test BinaryPlatforms.libc(p) == :glibc
    @test BinaryPlatforms.call_abi(p) == :eabihf
    @test BinaryPlatforms.compiler_abi(p) == CompilerABI(libstdcxx_version=v"3.4.24")
    @test occursin(r"^arm(v7l)?-linux-gnueabihf-libstdcxx24-cuda\+11.1-march\+armv7l$", BinaryPlatforms.triplet(p))
    @test abi_agnostic(p) == Linux(:armv7l, libc=:glibc, call_abi=:eabihf)
    @test aatriplet(p) == "arm-linux-gnueabihf"
    @test replace_cxxstring_abi(p, :cxx11) == ExtendedPlatform(Linux(:armv7l, libc=:glibc, call_abi=:eabihf, compiler_abi=CompilerABI(; cxxstring_abi=:cxx11, libstdcxx_version=v"3.4.24")); march="armv7l", cuda="11.1")
    @test replace_libgfortran_version(p, v"4") == ExtendedPlatform(Linux(:armv7l, libc=:glibc, call_abi=:eabihf, compiler_abi=CompilerABI(; libgfortran_version=v"4", libstdcxx_version=v"3.4.24")); march="armv7l", cuda="11.1")
    @test get_concrete_platform(p; compilers = [:c], preferred_gcc_version = v"7", preferred_llvm_version = v"9") == ExtendedPlatform(Linux(:armv7l, libc=:glibc, call_abi=:eabihf, compiler_abi=CompilerABI(libgfortran_version=v"4.0.0", libstdcxx_version=v"3.4.24", cxxstring_abi=:cxx11)); march="armv7l", cuda="11.1")
    @test ExtendedPlatform(p; cuda="11.1", foo="bar") == ExtendedPlatform(Linux(:armv7l, libc=:glibc, call_abi=:eabihf, compiler_abi=CompilerABI(libstdcxx_version=v"3.4.24")); march="armv7l", cuda="11.1", foo="bar")
    @test ExtendedPlatform(p; foo="bar") == ExtendedPlatform(Linux(:armv7l, libc=:glibc, call_abi=:eabihf, compiler_abi=CompilerABI(libstdcxx_version=v"3.4.24")); march="armv7l", cuda="11.1", foo="bar")
    @test_throws ErrorException ExtendedPlatform(p; cuda="11.1", march="neon")
    @test startswith(repr(p), "ExtendedPlatform(Linux(:armv7l, libc=:glibc, call_abi=:eabihf, compiler_abi=CompilerABI(libstdcxx_version=v\"3.4.24\")); ")
    @test occursin("march=\"armv7l\"", repr(p))
    @test occursin("cuda=\"11.1\"", repr(p))
    @test endswith(repr(p), ")")
    @test parse(ExtendedPlatform, triplet(p)) == p
    meta = Dict{String,String}()
    Artifacts.pack_platform!(meta, p)
    @test meta == Dict("arch" => "armv7l","libc" => "glibc","march" => "armv7l", "libstdcxx_version" => "3.4.24","os" => "linux", "cuda" => "11.1")

    # Extended platforms with wrong microarchitecture
    @test_throws ArgumentError ExtendedPlatform(Linux(:x86_64); march="carmel")
    @test_throws ArgumentError ExtendedPlatform(Linux(:i686); march="haswell")
    # Extended platform with invalid keys
    @test_throws ArgumentError ExtendedPlatform(Linux(:x86_64); os="windows")
    @test_throws ArgumentError ExtendedPlatform(Linux(:i686); libc="musl")
    # Extended platform with "+" sign in them
    @test_throws ArgumentError ExtendedPlatform(Linux(:x86_64); invalid="plus+sign")
    @test_throws ArgumentError ExtendedPlatform(Linux(:i686), Dict("this+is" => "invalid"))

    # Parse `"any"` as `AnyPlatform`
    @test tryparse(ExtendedPlatform, "any") == AnyPlatform()
    # AnyPlatform shouldn't be extended
    @test isnothing(tryparse(ExtendedPlatform, "any-march+avx"))
    # This is valid standard platform
    @test tryparse(ExtendedPlatform, "x86_64-unknown-freebsd11.1") == FreeBSD(:x86_64)
    # This string doesn't contain any platform at all
    @test isnothing(tryparse(ExtendedPlatform, "abcde"))
    # What comes before the first key-value pair is not a valid platform
    @test isnothing(tryparse(ExtendedPlatform, "armv7l-linux-musleabihf-this-is+not-valid"))
    # This string doesn't contain valid key-value pairs in the extra part
    @test_throws ArgumentError parse(ExtendedPlatform, "x86_64-linux-gnu-this+is-not-valid")

    @testset "base_platform" begin
        @test base_platform(Linux(:armv7l)) == Linux(:armv7l)
        @test base_platform(ExtendedPlatform(FreeBSD(:x86_64); march="avx")) == FreeBSD(:x86_64)
    end

    @testset "march" begin
        @test isnothing(march(Linux(:x86_64)))
        @test isnothing(march(AnyPlatform()))
        @test isnothing(march(ExtendedPlatform(Windows(:x86_64); cuda="9.2")))
        @test march(ExtendedPlatform(FreeBSD(:x86_64); march="avx512")) == "avx512"
    end

    @testset "Sys utilities" begin
        p = ExtendedPlatform(Linux(:x86_64); cuda="9.2")
        @test Sys.islinux(p)
        @test !Sys.isapple(p)
        @test !Sys.isbsd(p)
        @test !Sys.iswindows(p)
        p = ExtendedPlatform(FreeBSD(:x86_64); cuda="9.2")
        @test !Sys.islinux(p)
        @test !Sys.isapple(p)
        @test Sys.isbsd(p)
        @test !Sys.iswindows(p)
        p = ExtendedPlatform(MacOS(:x86_64); march="avx512")
        @test !Sys.islinux(p)
        @test Sys.isapple(p)
        @test Sys.isbsd(p)
        @test !Sys.iswindows(p)
        p = ExtendedPlatform(Windows(:x86_64); foo="bar")
        @test !Sys.islinux(p)
        @test !Sys.isapple(p)
        @test !Sys.isbsd(p)
        @test Sys.iswindows(p)
    end

    @testset "Platform matching" begin
        # Extending same platform
        @test !platforms_match(ExtendedPlatform(Linux(:i686; libc=:musl); cuda="9.2"), ExtendedPlatform(Linux(:i686; libc=:glibc); cuda="9.2"))
        @test platforms_match(ExtendedPlatform(Linux(:i686; libc=:musl); cuda="9.2"), ExtendedPlatform(Linux(:i686; libc=:musl); cuda="9.2"))
        @test platforms_match(ExtendedPlatform(Linux(:aarch64; libc=:musl); cuda="9.2"), ExtendedPlatform(Linux(:aarch64; libc=:musl); cuda="9.2", march="thunderx2"))
        @test !platforms_match(ExtendedPlatform(Linux(:armv7l; libc=:musl); cuda="9.1"), ExtendedPlatform(Linux(:armv7l; libc=:musl); cuda="9.2", march="neon"))
        # Extending the same platform as the other one
        @test platforms_match(ExtendedPlatform(Linux(:x86_64; compiler_abi=CompilerABI(; libgfortran_version=v"5")); march="avx512"), Linux(:x86_64))
        @test !platforms_match(ExtendedPlatform(Windows(:x86_64); march="avx"), Windows(:i686))
        @test platforms_match(MacOS(:x86_64; compiler_abi=CompilerABI(; cxxstring_abi=:cxx11)), ExtendedPlatform(MacOS(:x86_64); march="avx512"))
        @test !platforms_match(FreeBSD(:x86_64; compiler_abi=CompilerABI(; cxxstring_abi=:cxx03)), ExtendedPlatform(FreeBSD(:x86_64; compiler_abi=CompilerABI(; cxxstring_abi=:cxx11)); cuda="10.1"))
        # Extending a platform different from the other one
        @test !platforms_match(ExtendedPlatform(Linux(:i686); cuda="9.2"), MacOS(:x86_64))
        @test !platforms_match(Linux(:i686), ExtendedPlatform(Windows(:x86_64); cuda="9.2"))
        # Extending different platforms
        @test !platforms_match(ExtendedPlatform(Linux(:i686); cuda="10.1"), ExtendedPlatform(FreeBSD(:x86_64); cuda="11.1"))
    end
end

@testset "AnyPlatform" begin
    # Test some AnyPlatform properties
    @test triplet(AnyPlatform()) == "any"
    @test Pkg.BinaryPlatforms.platform_name(AnyPlatform()) == "AnyPlatform"
    @test abi_agnostic(AnyPlatform()) === AnyPlatform()
    @test arch(AnyPlatform()) == :x86_64
    @test repr(AnyPlatform()) == "AnyPlatform()"

    # In the build environment we want AnyPlatform to look like x86_64-linux-musl
    @test get_concrete_platform(AnyPlatform(); compilers = [:c], preferred_gcc_version = v"7", preferred_llvm_version = v"9") ==
        get_concrete_platform(Linux(:x86_64, libc=:musl); compilers = [:c], preferred_gcc_version = v"7", preferred_llvm_version = v"9")
    @test BinaryBuilderBase.choose_shards(AnyPlatform()) == BinaryBuilderBase.choose_shards(Linux(:x86_64, libc=:musl))
    @test BinaryBuilderBase.aatriplet(AnyPlatform()) == BinaryBuilderBase.aatriplet(Linux(:x86_64, libc=:musl))
end

@testset "Target properties" begin
    for p in [Linux(:i686), Windows(:i686), Linux(:armv7l)]
        @test BinaryBuilderBase.nbits(p) == 32
    end

    for p in [Linux(:x86_64), Windows(:x86_64), Linux(:aarch64),
              Linux(:powerpc64le), MacOS()]
        @test BinaryBuilderBase.nbits(p) == 64
    end

    for p in [Linux(:x86_64), MacOS(), Windows(:i686)]
        @test BinaryBuilderBase.proc_family(p) == :intel
    end
    for p in [Linux(:aarch64; libc=:musl), Linux(:armv7l)]
        @test BinaryBuilderBase.proc_family(p) == :arm
    end
    @test BinaryBuilderBase.proc_family(Linux(:powerpc64le)) == :power

    for p in [Linux(:aarch64), FreeBSD(:x86_64)]
        @test BinaryBuilderBase.dlext(p) == "so"
    end
    @test BinaryBuilderBase.dlext(MacOS()) == "dylib"
    @test BinaryBuilderBase.dlext(Windows(:i686)) == "dll"

    for p in [Linux(:x86_64), FreeBSD(:x86_64), Linux(:powerpc64le), MacOS()]
        @test BinaryBuilderBase.exeext(p) == ""
    end
    @test BinaryBuilderBase.exeext(Windows(:x86_64)) == ".exe"
    @test BinaryBuilderBase.exeext(Windows(:i686)) == ".exe"

    @test BinaryPlatforms.platform_dlext(ExtendedPlatform(Linux(:aarch64); march="thunderx2")) == "so"
    @test BinaryPlatforms.platform_dlext(ExtendedPlatform(FreeBSD(:x86_64); march="x86_64")) == "so"
    @test BinaryPlatforms.platform_dlext(ExtendedPlatform(MacOS(:x86_64); march="avx512")) == "dylib"
    @test BinaryPlatforms.platform_dlext(ExtendedPlatform(Windows(:i686); cuda="10.1")) == "dll"
end


text:
using Optim

export steady_state_trim

mutable struct Trimmer
    ac::Aircraft
    aerostate::AeroState
    state::State
    env::Environment
    ψ_dot::Number
    γ::Number
    controls::Controls
    cost_vector::Array{T, 1} where T<:Number
end

Trimmer(ac, aerostate, state, env, ψ_dot, γ, controls) = Trimmer(
    ac, aerostate, state, env, ψ_dot, γ, controls, fill(NaN, 6)
    )

function Trimmer(ac, controls, env, pos, tas, ψ, γ, ψ_dot, α, β)
    # Impose constrains
    att  = trimmer_constrains_attitude(ψ, ψ_dot, α, β, tas, γ)
    ang_vel = turn_rate_angular_velocity(ψ_dot, att.theta, att.phi)

    # Ensure env is updated
    env = calculate_environment(env, pos)

    # There is no information to set acceleration and angular acceleration at this point.
    accel = [0., 0., 0.]
    ang_accel = [0., 0., 0.]

    state, aerostate = generate_state_aerostate(
        pos,
        att,
        tas,
        α,
        β,
        env,
        ang_vel,
        accel,
        ang_accel,
    )

    return Trimmer(ac, aerostate, state, env, ψ_dot, γ, controls)
end

function update_trimmer!(trimmer, α, β, controls)

    ψ_dot = trimmer.ψ_dot
    tas = get_tas(trimmer.aerostate)
    γ = trimmer.γ
    ψ = get_euler_angles(trimmer.state)[1]

    # Impose constrains
    att  = trimmer_constrains_attitude(ψ, ψ_dot, α, β, tas, γ)
    ang_vel = turn_rate_angular_velocity(ψ_dot, att.theta, att.phi)
    # Calculate new state and aerostate
    state, aerostate = generate_state_aerostate(
        get_position(trimmer.state),
        att,
        tas,
        α,
        β,
        trimmer.env,
        ang_vel,
        get_body_accel(trimmer.state),
        get_body_ang_accel(trimmer.state)
    )
    # Calcualte aircraft
    ac = calculate_aircraft(
        trimmer.ac, controls, aerostate, state, trimmer.env.grav;
        consume_fuel = false
        )

    # Evaluate dynamic system to get x_dot
    ds = convert(SixDOFEulerFixedMass, state)
    x = get_x(ds)
    f = get_state_equation(ds)

    # Evaluate x_dot given x, u, parameters
    x_dot = f(
        x,
        get_mass_props(ac).mass,
        get_mass_props(ac).inertia,
        ac.pfm.forces,
        ac.pfm.moments
        )

        # Transform x, x_dot to State
    state = convert(state, SixDOFEulerFixedMass(x, x_dot))

    # Update trimmer
    trimmer.ac = ac
    trimmer.aerostate = aerostate
    trimmer.state = state
    trimmer.controls = controls
    trimmer.state = state
    trimmer.cost_vector = x_dot[1:6]
end

evaluate_cost_function(trimmer::Trimmer) = sum(trimmer.cost_vector.^2)

"""
    steady_state_trim(
        ac::Aircraft, controls::Controls, env::Environment, tas::Number, pos::Position,
        ψ::Number, γ::Number, ψ_dot::Number, α0::Number, β0::Number;
        show_trace = false, g_tol = 1e-25, max_iters = 5000
        )

Find a steady state flight condition.

Three different conditions can be sought:
- Steady level longitudinal flight: bank angle μ=0 and flight-path angle γ=0.
- Steady level turns: turn_rate≠0 and flight-path angle γ=0.
- Steady longitudinal climbs or descents: γ≠0 and bank angle μ=0.
- Steady turning climbs or descents: γ≠0 and bank angle μ≠0.

# Inputs
ac: aircraft to be trimmed.
controls : controls to be used to trim the aircraft with their initial values.
env: environment variables (atmospheric values, wind and gravity).
tas: trimming true airspeed [m/s].
pos: position of the aircraft.
ψ: heading of the aircraft [rad].
γ: aerodynamic rate of climb of the aircraft [rad].
ψ_dot: heading rate of change [rad/s].
α0, β0: Inital AOA and AOS for trimmer algorithm.
show_trace: show trimming information and algorithm trace. Optional, default=true.

# Returns
ac: trimmed aircraft.
aerostate: trimmed aerostate.
state: trimmed state.
controls: trimmed controls

# Notes

Steady flight is defined as flight where the aircraft's linear and angular velocity vectors
are constant in a body-fixed reference frame (ie. body frame or wind frame).

SixDOFEulerFixedMass is always used for state derivative calculation during trimming.

# References

See section 2.5 in [1] for the definition of steady-state flight condition.
See section 3.4 in [1] for the algorithm description.

- [1] Stevens, B. L., Lewis, F. L., (1992). Aircraft control and simulation:
 dynamics, controls design, and autonomous systems. John Wiley & Sons.
 (page 41, formula 1.4-23)
"""
function steady_state_trim(
    ac::Aircraft, controls::Controls, env::Environment, tas::Number, pos::Position,
    ψ::Number, γ::Number, ψ_dot::Number, α0::Number, β0::Number;
    show_trace = false, g_tol = 1e-25, max_iters = 5000
    )

    # Store FCS configuration to be restored at the end.
    # Allow out of range values during trimming.
    fcs = get_fcs(ac)
    allow_oor_controls = get_allow_out_of_range_inputs(fcs)
    err_on_oor_controls = get_throw_error_on_out_of_range_inputs(fcs)
    set_allow_out_of_range_inputs!(fcs, true)
    set_throw_error_on_out_of_range_inputs!(fcs, false)

    # Initialize trimmer
    trimmer = Trimmer(ac, controls, env, pos, tas, ψ, γ, ψ_dot, α0, β0)
    update_trimmer!(trimmer, α0, β0, controls)
    cost = evaluate_cost_function(trimmer)

    # variables to trim the ac -> x0 = [α, β, controls_to_be_trimmed...]
    x0 = [α0, β0, get_controls_array(controls)...]

    # Wrapper for trim_cost_function with trimmer
    trimming_function(x) = trim_cost_function(x, trimmer)

    # Trim if not already trimmed
    # XXX: It is observed that sometimes optimization method returns a minimum slightly
    # above g_tol, so the optmization loop would be entered again in a retrimming.
    # See tests -> trimmer.jl
    if cost > g_tol * 10
        # Trim
        result = optimize(
            trimming_function,
            x0,
            Optim.Options(
                g_tol = g_tol,
                iterations = max_iters,
                show_trace = show_trace,
                show_every = 100,
                );
            )

        if show_trace
            println(result)
            println(Optim.minimizer(result))
        end
    end

    # Reconfigure fcs to initial configuration
    fcs = get_fcs(trimmer.ac)
    set_allow_out_of_range_inputs!(fcs, allow_oor_controls)
    set_throw_error_on_out_of_range_inputs!(fcs, err_on_oor_controls)

    # Return trimmed variables
    return trimmer.ac, trimmer.aerostate, trimmer.state, trimmer.controls
end

function trimmer_constrains_attitude(ψ, ψ_dot, α, β, tas, γ)
    ϕ = coordinated_turn_bank(ψ_dot, α, β, tas, γ)
    θ = climb_theta(γ, α, β, ϕ)
    return Attitude(ψ, θ, ϕ)
end

"""
    trim_cost_function(x, trimmer::Trimmer)

Function minimized to trim the aircraft. Independent variables are passed in x:
- x[1] angle of attack (rad).
- x[2] angle of sideslip (rad).
- x[3:end]: controls expected by set_controls_trimmer! (ie. stick_long, stick_lat, pedals,
throttle)

"""
function trim_cost_function(x, trimmer::Trimmer)

    # Unpack x with trimming values candidates
    # α and β are given by x
    α, β = x[1:2]
    controls = typeof(trimmer.controls)(x[3:end]...)

    update_trimmer!(trimmer, α, β, controls)

    # Calculate objective function cost
    cost = evaluate_cost_function(trimmer)
    return cost
end

ext: jl


text:
#
println("Dd) Test of the PhotoRecombination module with ASF from an internally generated initial- and final-state multiplet.")
#
setDefaults("print summary: open", "zzz-PhotoRecombination.sum")
setDefaults("method: continuum, Galerkin")           ## setDefaults("method: continuum, Galerkin")  "method: continuum, asymptotic Coulomb"
setDefaults("method: normalization, pure sine")      ## setDefaults("method: normalization, pure Coulomb")    setDefaults("method: normalization, pure sine")

if  false
    grid = Radial.Grid(Radial.Grid(true), rnt = 2.0e-5,h = 5.0e-2, hp = 1.0e-2, NoPoints = 900)

    wa = Atomic.Computation(Atomic.Computation(), name="xx", grid=grid, nuclearModel=Nuclear.Model(12.),
                            initialConfigs=[Configuration("1s^2 2s^2 2p^5"), Configuration("1s^2 2s 2p^6") ],
                            finalConfigs  =[Configuration("1s^2 2s^2 2p^6")],
                            process = Rec(),
                            processSettings=PhotoRecombination.Settings([E1, M1], [JAC.UseCoulomb, JAC.UseBabushkin], [10., 30., 50.], [0.],
                                                                        false, true, true, true, true, Tuple{Int64,Int64}[(1,1)]) )

    wb = perform(wa)

elseif  false
    # REC calculations for the capture into 1s, 2s, 2p, 3s, 3p, 3d and Z=5: Comparison with Eichler (ADNDT, 2000)
    setDefaults("unit: energy", "eV") # then, give photonEnergies in eV
    setDefaults("unit: cross section", "barn")

    recSettings = PhotoRecombination.Settings([E1, M1], [JAC.UseCoulomb, JAC.UseBabushkin], [1., 10., 100., 1000., 10000.],
                                              Float64[], false, false, false, true, LineSelection())

    grid = Radial.Grid(Radial.Grid(false), rnt = 4.0e-6, h = 5.0e-2, hp = 1.0e-2, rbox = 10.0)
    wa = Atomic.Computation(Atomic.Computation(), name="xx", grid=grid, nuclearModel=Nuclear.Model(9.),
                            ## initialConfigs = [Configuration("1s^2")],
                            ## finalConfigs   = [Configuration("1s^2 2s"), Configuration("1s^2 2p"),
                            ##                   Configuration("1s^2 3s"), Configuration("1s^2 3p"), Configuration("1s^2 3d")],
                            initialConfigs = [Configuration("1s")],
                            finalConfigs   = [Configuration("1s^2")],
                            process = Rec(), processSettings=recSettings )

    wb = perform(wa)

elseif  true
    # REC calculations for the capture into 2s at given ion energies and Z=92: Comparison with Fritzsche (PRA, 2005)
    setDefaults("unit: energy", "eV") # then, give photonEnergies in eV
    setDefaults("unit: cross section", "barn")

    recSettings = PhotoRecombination.Settings([E1, M1, E2, M2, E3, M3], [JAC.UseBabushkin], Float64[],
                                              [2.18, 21.8, 218.0], true, false, false, true, LineSelection())

    grid = Radial.Grid(Radial.Grid(false), rnt = 4.0e-6, h = 5.0e-2, hp = 0.3e-2, rbox = 10.0)
    wa = Atomic.Computation(Atomic.Computation(), name="xx", grid=grid, nuclearModel=Nuclear.Model(92.),
                            ## initialConfigs = [Configuration("1s^2")],
                            ## finalConfigs   = [Configuration("1s^2 2s"), Configuration("1s^2 2p"),
                            ##                   Configuration("1s^2 3s"), Configuration("1s^2 3p"), Configuration("1s^2 3d")],
                            initialConfigs = [Configuration("1s^2")],
                            finalConfigs   = [Configuration("1s^2 2s")],
                            process = Rec(), processSettings=recSettings )

    wb = perform(wa)

end
setDefaults("print summary: close", "")


text:
@testset "CompressibilityFactor" begin
    p = Parameters(0.4, 1.4737, 7^3, 1e-5, 393216, Float64)
    s = SimulationSystem(p, 3, Float64)
    potential = PseudoHS(50.0, p.kT)
    initialize!(s)
    move!(1000, s, potential; interval = 500)
    zfac = ZFactor(0.0, 0.0)
    move!(1000, s, potential, zfac; interval = 500)

    # Check consistency
    @test zfac.naverage == 1000.0
    # Compute the Carnahan-Starling value
    true_zcomp = 8.352786794040853
    # Check value, larger means that it needs more steps (which is the case)
    @test zfac.zval ≈ true_zcomp
end

@testset "CompressibilityFactorRDF" begin
    p = Parameters(0.4, 1.4737, 7^3, 1e-5, 393216, Float64)
    s = SimulationSystem(p, 3, Float64)
    potential = PseudoHS(50.0, p.kT)
    initialize!(s)
    move!(1000, s, potential; interval = 500)

    gfunc = PairDistributionFunction(s, 2^9)
    zfac = ZFactor(0.0, 0.0)
    move!(1000, s, potential, gfunc, zfac; interval = 500)

    # Check that the RDF is physically consistent
    rdf_values = @view gfunc.gofr[:, 2]
    half_idx = fld1(length(rdf_values), 2)
    @test isapprox(mean(rdf_values[half_idx:end]), 1.0; atol=1e-1)

    # Check consistency
    @test zfac.naverage == 1000.0
    # Compute the Carnahan-Starling value
    true_zcomp = 8.352786794040853
    # Check value, larger means that it needs more steps (which is the case)
    @test zfac.zval ≈ true_zcomp
end


text:
#########################
# Expression Predicates #
#########################

is_call(x) = Base.Meta.isexpr(x, :call) && (x.args[1] !== GlobalRef(Core, :tuple))

function is_method_definition(x)
    if isa(x, Expr)
        if x.head == :function
            return true
        elseif x.head == :(=) && isa(x.args[1], Expr)
            lhs = x.args[1]
            if lhs.head == :where
                lhs = lhs.args[1]
            end
            return lhs.head == :call
        end
    end
    return false
end

############################
# Expression Match/Replace #
############################

function replace_match!(replace, ismatch, x)
    if ismatch(x)
        return replace(x)
    elseif isa(x, Array)
        for i in eachindex(x)
            x[i] = replace_match!(replace, ismatch, x[i])
        end
    elseif isa(x, Expr)
        replace_match!(replace, ismatch, x.args)
    end
    return x
end

############
# Julia IR #
############

function fix_labels_and_gotos!(code::Vector)
    changes = Dict{Int,Int}()
    for (i, stmnt) in enumerate(code)
        if isa(stmnt, LabelNode)
            code[i] = LabelNode(i)
            changes[stmnt.label] = i
        end
    end
    for (i, stmnt) in enumerate(code)
        if isa(stmnt, GotoNode)
            code[i] = GotoNode(get(changes, stmnt.label, stmnt.label))
        elseif Base.Meta.isexpr(stmnt, :enter)
            stmnt.args[1] = get(changes, stmnt.args[1], stmnt.args[1])
        elseif Base.Meta.isexpr(stmnt, :gotoifnot)
            stmnt.args[2] = get(changes, stmnt.args[2], stmnt.args[2])
        end
    end
    return code
end

function copy_prelude_code(code::Vector)
    prelude_code = Any[]
    for stmnt in code
        if isa(stmnt, Nothing) || isa(stmnt, NewvarNode)
            push!(prelude_code, stmnt)
        else
            break
        end
    end
    return prelude_code
end

#################
# Miscellaneous #
#################

unquote(x) = x
unquote(x::QuoteNode) = x.value
unquote(x::Expr) = x.head == :quote ? first(x.args) : x

function unqualify_name(e::Expr)
    @assert e.head == :(.)
    return unqualify_name(last(e.args))
end

unqualify_name(name::Symbol) = name

# define safe loggers for use in generated functions (where task switches are not allowed)
for level in [:debug, :info, :warn, :error]
    @eval begin
        macro $(Symbol("safe_$level"))(ex...)
            macrocall = :(@placeholder $(ex...))
            # NOTE: `@placeholder` in order to avoid hard-coding @__LINE__ etc
            macrocall.args[1] = Symbol($"@$level")
            quote
                old_logger = global_logger()
                global_logger(Logging.ConsoleLogger(Core.stderr, old_logger.min_level))
                ret = $(esc(macrocall))
                global_logger(old_logger)
                ret
            end
        end
    end
end


text:
using BernsteinExpansions
using Base.Test

@time @testset "Univariate monomial, floating point input" begin
@test univariate(3, 3, 1.0, 2.0) == [1.0, 2.0, 4.0, 8.0]
@test univariate(1, 5, 0.0, 1.0) == [0.0, 0.2, 0.4, 0.6, 0.8, 1.0]
end

@time @testset "Univariate monomial, rational input" begin
@test univariate(3, 3, 1//1, 2//1) == [1//1, 2//1, 4//1, 8//1]
@test univariate(1, 5, 0//1, 1//1) == [0//1, 1//5, 2//5, 3//5, 4//5, 1//1]
end

@time @testset "Univariate monomial, rational input" begin
@test univariate(3, 3, 1//1, 2//1) == [1//1, 2//1, 4//1, 8//1]
@test univariate(1, 5, 0//1, 1//1) == [0//1, 1//5, 2//5, 3//5, 4//5, 1//1]
end

m = multivariate([5,5,5,5,5,5,5,5],[10,10,10,10,10,10,10,10],
                 [1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0],
                 [2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0]);
@time @testset "Multivariate monomial, float input, implicit form" begin

# test loop (default) algorithm
@test generate_tensor_form(m, algorithm="loop") == multivariate_tensor([5,5,5,5,5,5,5,5],
                                                     [10,10,10,10,10,10,10,10],
                                                     [1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0],
                                                     [2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0], algorithm="loop")

# test kron algorithm
@test generate_tensor_form(m, algorithm="kron") == multivariate_tensor([5,5,5,5,5,5,5,5],
                                                                       [10,10,10,10,10,10,10,10],
                                                                       [1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0],
                                                                       [2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0], algorithm="kron")
end


text:
using ArgParse

include("tyckiting_client/client.jl")

global args

function main()
  # Main method for running the client. Reads configuration from command line,
  # and starts the client.
  s = ArgParseSettings(description="Destroy 'em all")
  @add_arg_table s begin
    "--host", "-H"
      default="localhost"
      help="Host to connect to"
    "--port", "-P"
      default=3000
      arg_type=Int
      help="Port to connect to"
    "--verbose", "-v"
      action=:store_true
      help="Verbose output"
    "--name", "-n"
      default="bot"
      help="Bot's name"
    "--ai", "-a"
      default="dummy"
      help="Ai package"
  end

  global args = parse_args(s)
  start_client()
end

function start_client()
  global args
  client = TykitingClient(args["host"], args["port"], args["name"], args["ai"])
  @sync run(client, args["verbose"])
end

main()

max_stars_repo_name: Condimentum/spacetyckiting

max_issues_repo_name: Condimentum/spacetyckiting

max_forks_repo_name: Condimentum/spacetyckiting


text:
# todo: ^, fma, muladd, div, fld, cld, rem, mod, mod1, fld1
Base.:+(x::T) where {T<:InfExtendedReal} = T(+x.val)

Base.:-(x::T) where {T<:InfExtendedReal} = T(-x.val)

Base.:+(x::T, y::T) where {T<:InfExtendedReal} = isinf(x) ? isinf(y) ? T(x.val+y.val) : x : isinf(y) ? y : T(x.val+y.val)

Base.:-(x::T, y::T) where {T<:InfExtendedReal} = isinf(x) ? isinf(y) ? T(x.val-y.val) : x : isinf(y) ? -y : T(x.val-y.val)

Base.:*(x::T, y::T) where {T<:InfExtendedReal} =
  if isinf(x)
    if isinf(y)
      T(x.val * y.val)
    else
      iszero(y) ? throw(DivideError()) : T(Infinite(signbit(x) ⊻ signbit(y)))
    end
  else
    if isinf(y)
      iszero(x) ? throw(DivideError()) : T(Infinite(signbit(x) ⊻ signbit(y)))
    else
      T(x.val * y.val)
    end
  end

@generated Base.:/(x::InfExtendedReal{T}, y::InfExtendedReal{T}) where {T<:Real} = :(convert($(InfExtendedReal(typeof(one(T)/one(T)))),
  if isinf(x)
    if isinf(y)
      throw(DivideError())
    else
      Infinite(signbit(x) ⊻ signbit(y))
    end
  else
    if isinf(y)
      zero(T)
    else
      x.val / y.val
    end
  end))

Base.abs(x::InfExtendedReal{T}) where {T<:Real} = InfExtendedReal{T}(abs(x.val))

@generated Base.:(//)(x::InfExtendedReal{T}, y::InfExtendedReal{S}) where {T<:Real,S<:Real} = :(convert($(InfExtendedReal(typeof(one(T)//one(S)))),
  if isinf(x)
    if isinf(y)
      throw(DivideError())
    else
      Infinite(signbit(x) ⊻ signbit(y))
    end
  else
    if isinf(y)
      zero(T)
    else
      x.val // y.val
    end
  end))
Base.:(//)(x::InfExtendedReal, y::Real) = //(promote(x,y)...)
Base.:(//)(x::Real, y::InfExtendedReal) = //(promote(x,y)...)


text:
# This example is from https://spglib.github.io/spglib/definition.html#transformation-to-a-primitive-cell
@testset "Transformation to a primitive cell" begin
    lattice = [[7.17851431, 0, 0], [0, 3.99943947, 0], [0, 0, 8.57154746]]
    positions = [
        [0.0, 0.84688439, 0.1203133],
        [0.0, 0.65311561, 0.6203133],
        [0.0, 0.34688439, 0.3796867],
        [0.0, 0.15311561, 0.8796867],
        [0.5, 0.34688439, 0.1203133],
        [0.5, 0.15311561, 0.6203133],
        [0.5, 0.84688439, 0.3796867],
        [0.5, 0.65311561, 0.8796867],
    ]
    types = fill(8, length(positions))
    cell = Cell(lattice, positions, types)
    primitive_cell = find_primitive(cell)
    # Write transformation matrix directly
    @test primitive_cell.lattice ==
          cell.lattice * [
              1//2 1//2 0
              -1//2 1//2 0
              0 0 1
          ] ==
          [
              3.589257155 3.589257155 0.0
              -1.999719735 1.999719735 0.0
              0.0 0.0 8.57154746
          ]
    @test primitive_cell.positions ≈ [  # Python results
        0.15311561 0.34688439 0.65311561 0.84688439
        0.84688439 0.65311561 0.34688439 0.15311561
        0.1203133 0.6203133 0.3796867 0.8796867
    ]
    @test primitive_cell.types == [8, 8, 8, 8] ./ 8  # Python results
end

@testset "Rotate the basis vectors rigidly in the above example" begin
    lattice = [
        [5.0759761474456697, 5.0759761474456697, 0],
        [-2.8280307701821314, 2.8280307701821314, 0],
        [0, 0, 8.57154746],
    ]
    positions = [
        [0.0, 0.84688439, 0.1203133],
        [0.0, 0.65311561, 0.6203133],
        [0.0, 0.34688439, 0.3796867],
        [0.0, 0.15311561, 0.8796867],
        [0.5, 0.34688439, 0.1203133],
        [0.5, 0.15311561, 0.6203133],
        [0.5, 0.84688439, 0.3796867],
        [0.5, 0.65311561, 0.8796867],
    ]
    types = fill(8, length(positions))
    cell = Cell(lattice, positions, types)
    primitive_cell = find_primitive(cell)
    # Compare with documented results
    @test cell.lattice * [
        1//2 1//2 0
        -1//2 1//2 0
        0 0 1
    ] ≈ [
        3.95200346 1.12397269 0.0
        1.12397269 3.95200346 0.0
        0.0 0.0 8.57154746
    ]
    # Compare with documented and Python results
    @test primitive_cell.lattice ≈ [
        3.58925715 3.58925715 0.0
        -1.99971973 1.99971973 0.0
        0.0 0.0 8.57154746
    ]
    @test primitive_cell.positions ≈ [  # Python results
        0.15311561 0.34688439 0.65311561 0.84688439
        0.84688439 0.65311561 0.34688439 0.15311561
        0.1203133 0.6203133 0.3796867 0.8796867
    ]
    @test primitive_cell.types == [8, 8, 8, 8] ./ 8  # Python results
    @testset "Obtain the rotated primitive cell basis vectors" begin
        @test standardize_cell(cell, to_primitive = true, no_idealize = true).lattice ≈ [
            3.95200346 1.12397269 0.0
            1.12397269 3.95200346 0.0
            0.0 0.0 8.57154746
        ]
    end
end

# From https://github.com/unkcpz/LibSymspg.jl/blob/53d2f6d/test/test_api.jl#L113-L136
@testset "Cell reduce standardize" begin
    @testset "Test `find_primitive`" begin
        lattice = [
            4.0 0.0 0.0
            0.0 4.0 0.0
            0.0 0.0 4.0
        ]
        positions = [
            0.0 0.5
            0.0 0.5
            0.0 0.5
        ]
        types = [1, 1]
        cell = Cell(lattice, positions, types)
        new_cell = find_primitive(cell, 1e-5)
        @test new_cell.lattice ≈ [
            -2.0 2.0 2.0
            2.0 -2.0 2.0
            2.0 2.0 -2.0
        ]
        @test new_cell.positions ≈ [0.0 0.0 0.0]'
        @test new_cell.types == [1]
    end
    @testset "Test `refine_cell`" begin
        lattice = [
            -2.0 2.0 2.0
            2.0 -2.0 2.0
            2.0 2.0 -2.0
        ]
        positions = [0.0 0.0 0.0]'
        types = [1]
        cell = Cell(lattice, positions, types)
        new_cell = refine_cell(cell, 1e-5)
        @test new_cell.lattice ≈ [
            4.0 0.0 0.0
            0.0 4.0 0.0
            0.0 0.0 4.0
        ]
        @test new_cell.positions ≈ [
            0.0 0.5
            0.0 0.5
            0.0 0.5
        ]
        @test new_cell.types == [1, 1]
    end
end


text:
#  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
#  Copyright (C) 2020-2021, UChicago Argonne, LLC. All rights reserved.
#  Released under the modified BSD license. See COPYING.md for more details.

function parse_name(name::AbstractString)::Vector{String}
    # x
    m = match(r"^[-a-z0-9_]*$", name)
    if m !== nothing
        return [name]
    end

    # x[1,2,3]
    m = match(r"^([-a-z0-9_]*)\[([-a-z0-9_,]*)\]$"i, name)
    if m !== nothing
        return [m[1], split(m[2], ",")...]
    end

    # unknown
    error("Could not parse: $(name)")
end


text:
# hidden dynamics system runner.
#
# This file actually runs the hidden dynamics system problem
# and writes the results to file
#

using RobustDMD, Printf, Random, JLD, StatsBase

# most definitions are in another file

dirname = @__DIR__
include(dirname * "/example_hidden_dynamics_defs.jl")

################################################################
################################################################
################################################################
# Actually run the tests...

# make data

xclean, eigs, t, x = hidden_dynamics_system_data(nt=128)

# solver parameters

function prox_mr(ar)
    # project for numerical stability
    k = length(ar) >> 1;
    for i = 1:k
        ar[i<<1-1] = min(2.0, ar[i<<1-1]);
    end
end

tol = 1e-7;
ptf = 10000000;
itm = 1000;
optsl2 = DMD_PG_Options(itm=itm, tol=tol,
                        ptf=ptf,prox = prox_mr);
optshub = DMD_PG_Options(itm=itm, tol=tol,
                         ptf=ptf,prox = prox_mr);
optstrim = DMD_PG_Options(itm=itm, tol=tol,
                          ptf=ptf,prox = prox_mr);

# test parameters

seed = 8675309
Random.seed!(seed)

p = 0.05 # add spikes at this percentage of data points
k = 4 # rank of fit

nj = 5 # number of different levels of background noise
ntest = 20

nkeep = Integer(floor(0.8*length(x))) # number of sensors to keep when trimming

errs = Array{Float64}(undef,4,ntest,3,nj)
eigs_comp = Array{Complex{Float64}}(undef,k,4,ntest,3,nj)
params = Array{Float64}(undef,7,nj)

@printf "==================================================\n\n"
@printf "Running the hidden dynamics example... \n\n"

for j = 1:nj
    @printf "starting noise level %d of %d\n" j nj
    sigma = 10.0^(-6+j) # noise floor
    mu = 1.0 # size of spikes
    A = 1.0 # height of bump
    w = 10.0 # width of bump
    kappa = 5*sigma
    params[1,j] = sigma
    params[2,j] = mu
    params[3,j] = p
    params[4,j] = kappa
    params[5,j] = A
    params[6,j] = w
    params[7,j] = nkeep

    Threads.@threads for i = 1:ntest

        xdat = make_corrupt_data_sparse_spikes(xclean,sigma,mu,p)

        (err1, err2, err3, err4, eig1, eig2, eig3, eig4) = (
            dofits(xdat,t,eigs,kappa,optsl2,optshub,optstrim,nkeep))

        @printf "test %d %d sparse spikes\n"  j i
        @printf "errs %e %e %e %e\n" err1 err2 err3 err4
        errs[1,i,1,j] = err1
        errs[2,i,1,j] = err2
        errs[3,i,1,j] = err3
        errs[4,i,1,j] = err4

        eigs_comp[:,1,i,1,j] = eig1[:]
        eigs_comp[:,2,i,1,j] = eig2[:]
        eigs_comp[:,3,i,1,j] = eig3[:]
        eigs_comp[:,4,i,1,j] = eig4[:]

        xdat = make_corrupt_data_broken_sensors(xclean,sigma,mu,p)

        (err1, err2, err3, err4, eig1, eig2, eig3, eig4) = (
            dofits(xdat,t,eigs,kappa,optsl2,optshub,optstrim,nkeep))

        @printf "test %d %d broken sensors\n"  j i
        @printf "errs %e %e %e %e\n" err1 err2 err3 err4

        errs[1,i,2,j] = err1
        errs[2,i,2,j] = err2
        errs[3,i,2,j] = err3
        errs[4,i,2,j] = err4

        eigs_comp[:,1,i,2,j] = eig1[:]
        eigs_comp[:,2,i,2,j] = eig2[:]
        eigs_comp[:,3,i,2,j] = eig3[:]
        eigs_comp[:,4,i,2,j] = eig4[:]

        xdat = make_corrupt_data_bump(xclean,x,t,sigma,A,w)

        (err1, err2, err3, err4, eig1, eig2, eig3, eig4) = (
            dofits(xdat,t,eigs,kappa,optsl2,optshub,optstrim,nkeep))

        @printf "test %d %d bump\n"  j i
        @printf "errs %e %e %e %e\n" err1 err2 err3 err4

        errs[1,i,3,j] = err1
        errs[2,i,3,j] = err2
        errs[3,i,3,j] = err3
        errs[4,i,3,j] = err4

        eigs_comp[:,1,i,3,j] = eig1[:]
        eigs_comp[:,2,i,3,j] = eig2[:]
        eigs_comp[:,3,i,3,j] = eig3[:]
        eigs_comp[:,4,i,3,j] = eig4[:]

    end
end

fname = dirname * "/results/hidden_dynamics_out.jld"

@printf "writing results to file %s\n" fname

# save the results to file

mkpath(dirname * "/results")
file = jldopen(fname,"w")
file["errs"] = errs
file["eigs_comp"] = eigs_comp
file["params"] = params
close(file)

@printf "\n ...done\n"

@printf "==================================================\n"


text:
### A Pluto.jl notebook ###
# v0.14.5

using Markdown
using InteractiveUtils

# ╔═╡ 774b5c68-c2e8-11eb-0c25-c124276640d6
begin
	using Revise
	using Pkg
	Pkg.activate(tempname())
	Pkg.develop(url="https://github.com/manuelbb-upb/RadialBasisFunctionModels.jl.git")
	Pkg.add("Flux")
	Pkg.add("MLDatasets")
	using Flux, MLDatasets
	using Random, Statistics
	using RadialBasisFunctionModels
end

# ╔═╡ 8c51b60d-e1b5-439d-ab67-ef7813490684
using Flux.Optimise: update!

# ╔═╡ ec8e09e0-2238-4314-baf4-1209794bb4e4
md"Create some test data: `N_d` is the total number of samples."

# ╔═╡ a974f4de-47b2-4981-9381-717abd8ccda2
begin
	F = x -> [
		sum(x.^2);
		exp(x[1]) + sum(x);
		sin( sum(x)/(1+sum(x)) )
	]

	N_d = 100
	features = [ rand(3) for i = 1 : N_d ]
	labels = F.(features)
	data_indices = eachindex(features)

	# training data
	training_percentage = .8
	num_training = ceil(Int, N_d * training_percentage )
	num_test = N_d - num_training

	train_indices = randperm(N_d)[1:num_training]
	X_train = features[ train_indices ]
	Y_train = labels[ train_indices ]

	test_indices = setdiff( data_indices, train_indices )
	X_test = features[ test_indices ]
	Y_test = labels[ test_indices ]

	n_rbf_centers = ceil(Int, num_training/5)
	center_indices = randperm(num_training)[ 1 : n_rbf_centers ]
	nothing
end

# ╔═╡ 4b0db5f4-77cb-4391-a3bb-aae7d3c1c140
md"Build a model by chaining a dense neural net and an RBF Layer."

# ╔═╡ 1f80e0a3-71a3-483d-acc5-93c032655d4e
begin
	function get_rbf_model( m̃ )
		global X_train, Y_train
		X̃ = m̃.( X_train )
		rbf_centers = X̃[ center_indices ]
		return RBFModel( X̃, Y_train, Multiquadric(), -1; centers = rbf_centers )
	end

	function get_model1()
		n_out = 8
		layer1 = Dense(3, 10)
		layer2 = Dense(10, n_out)

		m̃ = Chain( layer1, layer2 )

		r = get_rbf_model( m̃ )

		return Chain(m̃, r)
	end
end

# ╔═╡ da1f41c7-12dd-4c8f-b267-a74bf99bc73b
begin
	batch_size = 10
	train_loader = Flux.Data.DataLoader((X_train, Y_train); batchsize = batch_size)
end

# ╔═╡ 44b9ab7e-c52e-4faa-b925-12eaf7c734da
Markdown.parse("""
## Minibatch Training

The training data is split into batches of size ``N_b = $(batch_size)``.
Suppose the labels have dimension ``k``.
The loss is calculated over the batch ``Y``:
```math
	L(Y) = \\frac{1}{N_b} \\sum_{y_i ∈ Y} \\frac{1}{k} \\sum_{ℓ=1}^k  (y_{i,ℓ} - \\hat{y}_{i,ℓ})^2
```
""")

# ╔═╡ 5cdfe543-d773-41ec-8840-2c79b767fb27
md"Training without affecting the rbf layer at all."

# ╔═╡ 10dc99e7-028b-4378-a549-c32b49f14a65
Flux.trainable(::RBFModel) = ()

# ╔═╡ 2ad4a6ab-57ff-45d6-9087-73bae6e81aff
begin
	opt = ADAM()
end

# ╔═╡ c8f1ce7d-a17f-4d13-acb5-e03d6002a273
begin
	m = get_model1()
	loss_sample = (x, y) -> Flux.Losses.mse(m(x), y);
	loss_batch = function( X, Y )
		mean( loss_sample.(X,Y) )
	end

	ps = params(m)	# not affecting RBF Layer

	test_loss_before = loss_batch( X_test, Y_test )

	Flux.@epochs 2 for (X_batch, Y_batch) ∈ train_loader
		loss_val, pback = Flux.Zygote.pullback(ps) do
			loss_batch(X_batch, Y_batch)
		end
		@show loss_val
		gs = pback( one(loss_val) )
		update!(opt, ps, gs )
	end
	test_loss_after = loss_batch( X_test, Y_test )

	Markdown.parse("$(test_loss_before) -> \n$(test_loss_after)")
end

# ╔═╡ 80b2d7ed-1cb6-46b1-a3da-17e19b434569
begin
	m̃ = get_model1()[1]

	loss_whole_chain = function(X,Y)
		local m
		rbf = get_rbf_model( m̃ )
		m = Chain( m̃, rbf )
		loss_sample = (x, y) -> Flux.Losses.mse(m(x), y);
		return mean( loss_sample.(X,Y) )
	end

	ps_2 = params(m̃)	# not affecting RBF Layer

	test_loss_before_2 = loss_whole_chain( X_test, Y_test )

	Flux.@epochs 1 for (X_batch, Y_batch) ∈ train_loader
		loss_val, pback = Flux.Zygote.pullback(ps_2) do
			loss_whole_chain(X_batch, Y_batch)
		end
		@show loss_val
		gs = pback( one(loss_val) )
		update!(opt, ps_2, gs )
	end
	test_loss_after_2 = loss_whole_chain( X_test, Y_test )

	Markdown.parse("$(test_loss_before_2) -> \n$(test_loss_after_2)")
end

# ╔═╡ 0123fe9d-352e-453c-8482-08c9230da47f
test_l = function(x)
	M = RBFModel( X_train, Y_train; centers = X_train[center_indices] )
	M(x, 1)
end

# ╔═╡ ed05ca08-9da5-48ad-8795-5ae6f3ee07dd
Flux.Zygote.gradient( test_l, rand(3))

# ╔═╡ 9f8ae0a4-adfb-49a9-8591-b3ba8a2f3409
M = RBFModel( X_train, Y_train; centers = X_train[center_indices] )

# ╔═╡ 589e6360-5d11-494b-9081-5ea30111b052
auto_grad(M, rand(3))

# ╔═╡ Cell order:
# ╠═774b5c68-c2e8-11eb-0c25-c124276640d6
# ╠═8c51b60d-e1b5-439d-ab67-ef7813490684
# ╟─ec8e09e0-2238-4314-baf4-1209794bb4e4
# ╠═a974f4de-47b2-4981-9381-717abd8ccda2
# ╟─4b0db5f4-77cb-4391-a3bb-aae7d3c1c140
# ╠═1f80e0a3-71a3-483d-acc5-93c032655d4e
# ╟─44b9ab7e-c52e-4faa-b925-12eaf7c734da
# ╠═da1f41c7-12dd-4c8f-b267-a74bf99bc73b
# ╟─5cdfe543-d773-41ec-8840-2c79b767fb27
# ╠═10dc99e7-028b-4378-a549-c32b49f14a65
# ╠═2ad4a6ab-57ff-45d6-9087-73bae6e81aff
# ╠═c8f1ce7d-a17f-4d13-acb5-e03d6002a273
# ╠═80b2d7ed-1cb6-46b1-a3da-17e19b434569
# ╠═0123fe9d-352e-453c-8482-08c9230da47f
# ╠═ed05ca08-9da5-48ad-8795-5ae6f3ee07dd
# ╠═9f8ae0a4-adfb-49a9-8591-b3ba8a2f3409
# ╠═589e6360-5d11-494b-9081-5ea30111b052


text:
using GeometricIntegrators.Config
using GeometricIntegrators.Integrators
using GeometricIntegrators.Solvers
using GeometricIntegrators.Tableaus
using GeometricIntegrators.TestProblems.LotkaVolterra
using GeometricIntegrators.Utils
using Test

set_config(:nls_atol, 8eps())
set_config(:nls_rtol, 2eps())
set_config(:nls_solver, QuasiNewtonSolver)
set_config(:jacobian_autodiff, true)

using GeometricIntegrators.TestProblems.LotkaVolterra: Δt, nt

idae = lotka_volterra_2d_idae()
pdae = lotka_volterra_2d_pdae()

int  = IntegratorFIRK(lotka_volterra_2d_ode(), getTableauGLRK(8), Δt)
sol  = integrate(int, nt)
refx = sol.q[:,end]

@testset "$(rpad("VPARK integrators",80))" begin

    dint = Integrator(idae, getTableauSymplecticProjection(:pglrk1ps, getCoefficientsGLRK(1), getCoefficientsGLRK(1)), Δt)
    dsol = integrate(dint, nt)

    # println(rel_err(dsol.q, refx))
    @test rel_err(dsol.q, refx) < 1E-6

    dint = Integrator(idae, getTableauSymplecticProjection(:pglrk2ps, getCoefficientsGLRK(2), getCoefficientsGLRK(2)), Δt)
    dsol = integrate(dint, nt)

    # println(rel_err(dsol.q, refx))
    @test rel_err(dsol.q, refx) < 1E-11

    dint = Integrator(idae, getTableauGLRKpSymplectic(1), Δt)
    dsol = integrate(dint, nt)

    # println(rel_err(dsol.q, refx))
    @test rel_err(dsol.q, refx) < 1E-6

    dint = Integrator(idae, getTableauGLRKpSymplectic(2), Δt)
    dsol = integrate(dint, nt)

    # println(rel_err(dsol.q, refx))
    @test rel_err(dsol.q, refx) < 1E-11

    dint = Integrator(idae, getTableauLobIIIAIIIB2pSymplectic(), Δt)
    dsol = integrate(dint, nt)

    # println(rel_err(dsol.q, refx))
    @test rel_err(dsol.q, refx) < 2E-6

    dint = Integrator(idae, getTableauLobIIIAIIIB3pSymplectic(), Δt)
    dsol = integrate(dint, nt)

    # println(rel_err(dsol.q, refx))
    @test rel_err(dsol.q, refx) < 8E-5

end

@testset "$(rpad("VSPARK integrators",80))" begin
    ### VSPARK Integrators ###

    dint = IntegratorVSPARK(idae, getTableauSPARKGLRK(1), Δt)
    # dsol = integrate(dint, nt)
    # TODO
    # println(rel_err(dsol.q, refx))
    # @test rel_err(dsol.q, refx) < 1E-6

    ### VSPARKprimary Integrators ###

    dint = Integrator(idae, getTableauVSPARKGLRKpMidpoint(1), Δt)
    dsol = integrate(dint, nt)

    # println(rel_err(dsol.q, refx))
    @test rel_err(dsol.q, refx) < 1E-6

    dint = Integrator(idae, getTableauVSPARKGLRKpMidpoint(2), Δt)
    dsol = integrate(dint, nt)

    # println(rel_err(dsol.q, refx))
    @test rel_err(dsol.q, refx) < 1E-11

    dint = Integrator(idae, getTableauVSPARKGLRKpSymplectic(1), Δt)
    dsol = integrate(dint, nt)

    # println(rel_err(dsol.q, refx))
    @test rel_err(dsol.q, refx) < 1E-6

    dint = Integrator(idae, getTableauVSPARKGLRKpSymplectic(2), Δt)
    dsol = integrate(dint, nt)

    # println(rel_err(dsol.q, refx))
    @test rel_err(dsol.q, refx) < 1E-11

    dint = Integrator(idae, getTableauVSPARKGLRKpSymmetric(1), Δt)
    dsol = integrate(dint, nt)

    # println(rel_err(dsol.q, refx))
    @test rel_err(dsol.q, refx) < 1E-6

    dint = Integrator(idae, getTableauVSPARKGLRKpSymmetric(2), Δt)
    dsol = integrate(dint, nt)

    # println(rel_err(dsol.q, refx))
    @test rel_err(dsol.q, refx) < 1E-11

    dint = IntegratorVSPARKprimary(idae, getTableauVSPARKLobIIIAIIIB2pSymmetric(), Δt)
    dsol = integrate(dint, nt)

    # println(rel_err(dsol.q, refx))
    @test rel_err(dsol.q, refx) < 2E-6

    dint = IntegratorVSPARKprimary(idae, getTableauVSPARKLobIIIAIIIB3pSymmetric(), Δt)
    dsol = integrate(dint, nt)

    # println(rel_err(dsol.q, refx))
    @test rel_err(dsol.q, refx) < 8E-5

end

@testset "$(rpad("HPARK integrators",80))" begin

    dint = Integrator(pdae, getTableauHPARK(:hpark_glrk1, getCoefficientsGLRK(1), getCoefficientsGLRK(1)), Δt)
    dsol = integrate(dint, nt)

    # println(rel_err(dsol.q, refx))
    @test rel_err(dsol.q, refx) < 2E-6

    dint = Integrator(pdae, getTableauHPARK(:hpark_glrk2, getCoefficientsGLRK(2), getCoefficientsGLRK(2)), Δt)
    dsol = integrate(dint, nt)

    # println(rel_err(dsol.q, refx))
    @test rel_err(dsol.q, refx) < 8E-7

    dint = Integrator(pdae, getTableauHPARKGLRK(1), Δt)
    dsol = integrate(dint, nt)

    # println(rel_err(dsol.q, refx))
    @test rel_err(dsol.q, refx) < 2E-6

    dint = Integrator(pdae, getTableauHPARKGLRK(2), Δt)
    dsol = integrate(dint, nt)

    # println(rel_err(dsol.q, refx))
    @test rel_err(dsol.q, refx) < 8E-7

    dint = Integrator(pdae, getTableauHPARKLobIIIAIIIB2(), Δt)
    # dsol = integrate(dint, nt)
    # TODO
    # println(rel_err(dsol.q, refx))
    # @test rel_err(dsol.q, refx) < 2E-2

    dint = Integrator(pdae, getTableauHPARKLobIIIAIIIB3(), Δt)
    # dsol = integrate(dint, nt)
    # TODO
    # println(rel_err(dsol.q, refx))
    # @test rel_err(dsol.q, refx) < 8E-2

end

@testset "$(rpad("HSPARK integrators",80))" begin
    ### HSPARK Integrators ###

    dint = IntegratorHSPARK(pdae, getTableauSPARKGLRK(1), Δt)
    dsol = integrate(dint, nt)

    # println(rel_err(dsol.q, refx))
    @test rel_err(dsol.q, refx) < 1E-6

    ### HSPARKprimary Integrators ###

    dint = Integrator(pdae, getTableauHSPARKGLRKpSymmetric(1), Δt)
    dsol = integrate(dint, nt)

    # println(rel_err(dsol.q, refx))
    @test rel_err(dsol.q, refx) < 2E-1

    dint = Integrator(pdae, getTableauHSPARKGLRKpSymmetric(2), Δt)
    dsol = integrate(dint, nt)

    # println(rel_err(dsol.q, refx))
    @test rel_err(dsol.q, refx) < 4E-6

    dint = Integrator(pdae, getTableauHSPARKLobIIIAIIIB2pSymmetric(), Δt)
    dsol = integrate(dint, nt)

    # println(rel_err(dsol.q, refx))
    @test rel_err(dsol.q, refx) < 2E-1

    dint = Integrator(pdae, getTableauHSPARKLobIIIAIIIB3pSymmetric(), Δt)
    dsol = integrate(dint, nt)

    # println(rel_err(dsol.q, refx))
    @test rel_err(dsol.q, refx) < 4E-6
end


text:
export ElectronSite,
       electronSites

struct ElectronSite <: Site
  s::Index
  ElectronSite(I::Index) = new(I)
end
ElectronSite(n::Int) = ElectronSite(Index(4,"Site,Electron,n=$n"))

function electronSites(N::Int;kwargs...)::SiteSet
  sites = SiteSet(N)
  for n=1:N
    set(sites,n,ElectronSite(n))
  end
  return sites
end

function operator(site::ElectronSite,
                  opname::AbstractString)::ITensor
    s  = site.s
    sP = prime(site.s)
    Emp   = s(1)
    EmpP  = sP(1)
    Up    = s(2)
    UpP   = sP(2)
    Dn    = s(3)
    DnP   = sP(3)
    UpDn  = s(4)
    UpDnP = sP(4)

    Op = ITensor(dag(s), s')

    if opname == "Nup"
        Op[Up, UpP] = 1.
        Op[UpDn, UpDnP] = 1.
    elseif opname == "Ndn"
        Op[Dn, DnP] = 1.
        Op[UpDn, UpDnP] = 1.
    elseif opname == "Ntot"
        Op[Up, UpP] = 1.
        Op[Dn, DnP] = 1.
        Op[UpDn, UpDnP] = 2.
    elseif opname == "Cup" || opname == "Aup"
        Op[Up, EmpP] = 1.
        Op[UpDn, DnP] = 1.
    elseif opname == "Cdagup" || opname == "Adagup"
        Op[Emp, UpP] = 1.
        Op[Dn, UpDnP] = 1.
    elseif opname == "Cdn" || opname == "Adn"
        Op[Dn, EmpP] = 1.
        Op[UpDn, UpP] = 1.
    elseif opname == "Cdagdn" || opname == "Adagdn"
        Op[Emp, DnP] = 1.
        Op[Up, UpDnP] = 1.
    elseif opname == "FermiPhase" || opname == "FP"
        Op[Up, UpP] = -1.
        Op[Emp, EmpP] = 1.
        Op[Dn, DnP] = -1.
        Op[UpDn, UpDnP] = 1.
    elseif opname == "Fup"
        Op[Emp, EmpP] = 1.
        Op[Up, UpP] = -1.
        Op[Dn, DnP] = 1.
        Op[UpDn, UpDnP] = -1.
    elseif opname == "Fdn"
        Op[Emp, EmpP] = 1.
        Op[Up, UpP] = 1.
        Op[Dn, DnP] = -1.
        Op[UpDn, UpDnP] = -1.
    elseif opname == "Sᶻ" || opname == "Sz"
        Op[Up, UpP] = 0.5
        Op[Dn, DnP] = -0.5
    elseif opname == "Sˣ" || opname == "Sx"
        Op[Up, DnP] = 1.0
        Op[Dn, UpP] = 1.0
    elseif opname == "S⁺" || opname == "Splus"
        Op[Dn, UpP] = 1.
    elseif opname == "S⁻" || opname == "Sminus"
        Op[Up, DnP] = 1.
    elseif opname == "Emp" || opname == "0"
        pEmp = ITensor(s)
        pEmp[Emp] = 1.
        return pEmp
    elseif opname == "Up" || opname == "↑"
        pU = ITensor(s)
        pU[Up] = 1.
        return pU
    elseif opname == "Dn" || opname == "↓"
        pD = ITensor(s)
        pD[Dn] = 1.
        return pD
    elseif opname == "UpDn" || opname == "↑↓"
        pUD = ITensor(s)
        pUD[UpDn] = 1.
        return pUD
    else
        throw(ArgumentError("Operator name $opname not recognized for ElectronSite"))
    end
    return Op
end


text:
import KernelTools
using Base.Test

#### Testing the utility functions

# Convenience function for testing equality of expressions
# Strip line numbers, extraneous :block wrappers, and :escape expressions

# There's one small bit of type-instability: Expr(:block, :a) or Expr(:escape, :a)
# will get turned into a symbol, not an Expr. By confining the type-instability to
# an "outer" wrapper (stripexpr) and making the core algorithm (stripexpr!) type-stable,
# we hope to make this faster. As a consequence, stripexpr! can only modify the args
# of an expression, and can't elide the container expression itself.
function stripexpr(a)
    ex = stripexpr!(copy(a))
    if isa(ex, Expr) && (ex.head == :block || ex.head == :escape) && length(ex.args) == 1
        return ex.args[1]
    end
    ex
end

stripexpr!(a) = a
function stripexpr!(ex::Expr)
    # Strip line number statements
    keep = ![Base.is_linenumber(ex.args[i]) for i = 1:length(ex.args)]
    ex.args = ex.args[keep]
    # Eliminate :escape and unnecessary :block statements
    nargs = length(ex.args)
    for i = 1:nargs
        if isa(ex.args[i], Expr)
            exa = stripexpr!(ex.args[i]::Expr)
            if exa.head == :escape
                length(exa.args) == 1 || error("complex escape expression")
                ex.args[i] = exa.args[1]
            elseif exa.head == :block && length(exa.args) == 1
                ex.args[i] = exa.args[1]
            end
        end
    end
    ex
end

# hoisting expressions that are constant for the purposes of some loops
KernelTools.ktgensym_counter_reset()
ex = KernelTools.gen_hoisted(:(s += A[i]), [:j, :i], [:(1:size(A,2)),:(1:size(A,1))])
@test stripexpr(ex) == stripexpr(:(for j = 1:size(A,2) for i = 1:size(A,1) s += A[i] end end))
ex = KernelTools.gen_hoisted(:(s += A[j]), [:j, :i], [:(1:size(A,2)),:(1:size(A,1))])
@test stripexpr(ex) == stripexpr(:(for j = 1:size(A,2); _kt_1 = A[j]; for i = 1:size(A,1) s += _kt_1 end; nothing; end))

# Symbols associated with setindex!-like statements
@test KernelTools.assignments(:(a = 5)) == Symbol[]
@test KernelTools.assignments(:(a[3] = 5)) == [:a]
@test KernelTools.assignments(:(a[3] += 5)) == [:a]
@test KernelTools.assignments(:(a[3] -= b[3])) == [:a]
@test KernelTools.assignments(:(a *= b[3])) == Symbol[]
@test KernelTools.assignments(:(A[i,j] /= b[3]; c[7] = 8)) == [:A, :c]

# Collecting all reference expressions
# optionally we ignore all those that occur on the left hand side of an assignment
@test KernelTools.refexprs(:(a = 5)) == Expr[]
@test KernelTools.refexprs(:(a[3] = 5)) == [:(a[3])]
@test KernelTools.refexprs(:(a[3] += 5), skiplhs=true) == Expr[]
@test KernelTools.refexprs(:(a[3] -= b[3])) == [:(a[3]), :(b[3])]
@test KernelTools.refexprs(:(a[3] -= b[3]), skiplhs=true) == [:(b[3])]
@test KernelTools.refexprs(:(a *= b[i])) == [:(b[i])]
@test KernelTools.refexprs(:(A[i,j] /= b[3]; c[7] = 8)) == [:(A[i,j]), :(b[3]), :(c[7])]

# Parsing indexes
@test KernelTools.AffineIndex(7) == KernelTools.AffineIndex(:nothing, 0, 7)
@test KernelTools.AffineIndex(:i) == KernelTools.AffineIndex(:i, 1, 0)
@test KernelTools.AffineIndex(:(2k)) == KernelTools.AffineIndex(:k, 2, 0)
@test KernelTools.AffineIndex(:(j+3)) == KernelTools.AffineIndex(:j, 1, 3)
@test KernelTools.AffineIndex(:(3qq-1)) == KernelTools.AffineIndex(:qq, 3, -1)
@test KernelTools.AffineIndex(:(i + 2 - 1)) == KernelTools.AffineIndex(:i, 1, 1)
@test KernelTools.AffineIndex(:(i + 2 + 3i)) == KernelTools.AffineIndex(:i, 4, 2)
@test_throws ErrorException KernelTools.AffineIndex(:(i + 2 + 3j))

## Bounds computations on temporaries for tiling
    # ti, tj = tile sizes for i, j
    # inner indexes ii take values in 0:(ti-1)
    # A statement like A[i] will get turned into A[offset+ii], and A[i-1] into A[offset+ii-1]
    # compute the offset and size of A needed to be able to execute all indexing statements without bounds errors
looprangedict = [:i => :(1:size(A,1)), :j => :(1:size(A,2))]
# A[i]
stmnts = Any[]
offsetsyms, sizesyms, lastsyms = KernelTools.constructbounds!(stmnts, :A, Vector{Any}[[:i]], [:i], [:ti], looprangedict)
@test stripexpr(stmnts[1]) == :(const $(offsetsyms[1]) = 1 - min(0, 1*(ti-1)+0))  # ultimately, this will evaluate to 1
@test stripexpr(stmnts[2]) == :(const $(sizesyms[1]) = $(offsetsyms[1]) + max(0, 1*(ti-1)+0))  # ultimately, this will evaluate to ti
# A[i-1]
stmnts = Any[]
offsetsyms, sizesyms, lastsyms = KernelTools.constructbounds!(stmnts, :A, Vector{Any}[[:(i-1)]], [:i], [:ti], looprangedict)
@test stripexpr(stmnts[1]) == :(const $(offsetsyms[1]) = 1 - min(-1, 1*(ti-1)+(-1)))  # ultimately, this will evaluate to 2
@test stripexpr(stmnts[2]) == :(const $(sizesyms[1]) = $(offsetsyms[1]) + max(-1, 1*(ti-1)+(-1)))  # ultimately, this will evaluate to ti
# A[i-1] and A[i+1]
stmnts = Any[]
offsetsyms, sizesyms, lastsyms = KernelTools.constructbounds!(stmnts, :A, Vector{Any}[[:(i-1)],[:(i+1)]], [:i], [:ti], looprangedict)
@test stripexpr(stmnts[1]) == :(const $(offsetsyms[1]) = 1 - min(-1, 1*(ti-1)+(-1), 1, 1*(ti-1)+1))  # ultimately, this will evaluate to 2
@test stripexpr(stmnts[2]) == :(const $(sizesyms[1]) = $(offsetsyms[1]) + max(-1, 1*(ti-1)+(-1), 1, 1*(ti-1)+1))  # ultimately, this will evaluate to ti+2
# A[i+1,j] and A[i,j+1]
stmnts = Any[]
offsetsyms, sizesyms, lastsyms = KernelTools.constructbounds!(stmnts, :A, Vector{Any}[[:(i+1),:j],[:i,:(j+1)]], [:i,:j], [:ti,:tj], looprangedict)
@test stripexpr(stmnts[1]) == :(const $(offsetsyms[1]) = 1 - min(1, 1*(ti-1)+1, 0, 1*(ti-1)+0))
@test stripexpr(stmnts[2]) == :(const $(sizesyms[1]) = $(offsetsyms[1]) + max(1, 1*(ti-1)+1, 0, 1*(ti-1)+0))
@test stripexpr(stmnts[4]) == :(const $(offsetsyms[2]) = 1 - min(0, 1*(tj-1)+0, 1, 1*(tj-1)+1))
@test stripexpr(stmnts[5]) == :(const $(sizesyms[2]) = $(offsetsyms[2]) + max(0, 1*(tj-1)+0, 1, 1*(tj-1)+1))
# Tiling just j for A[i,j]
stmnts = Any[]
offsetsyms, sizesyms, lastsyms = KernelTools.constructbounds!(stmnts, :A, Vector{Any}[[:i,:j]], [:j], [:tj], looprangedict)
@test stripexpr(stmnts[3]) == :(const $(lastsyms[1]) = 1*last($(looprangedict[:i])) + 0)
@test stripexpr(stmnts[6]) == :(const $(lastsyms[2]) = 1*last($(looprangedict[:j])) + 0)
# Tiling just j for A[i+1,j]
stmnts = Any[]
offsetsyms, sizesyms, lastsyms = KernelTools.constructbounds!(stmnts, :A, Vector{Any}[[:(i+1),:j]], [:j], [:tj], looprangedict)
@test stripexpr(stmnts[3]) == :(const $(lastsyms[1]) = 1*last($(looprangedict[:i])) + 1)
@test stripexpr(stmnts[6]) == :(const $(lastsyms[2]) = 1*last($(looprangedict[:j])) + 0)

##### Loop ordering in @tile
getlooporder(ex::Expr) = getlooporder!(Array(Vector{Symbol}, 0), ex)
function getlooporder!(order, ex::Expr)
    if ex.head == :for
        loopvars = ex.args[1]
        if isa(loopvars, Expr) && (loopvars::Expr).head == :block
            newvars = Symbol[]
            for exlv in (loopvars::Expr).args
                push!(newvars, getloopvar(exlv))
            end
            push!(order, newvars)
        else
            push!(order, [getloopvar(loopvars)])
        end
        getlooporder!(order, ex.args[2])
    else
        for i = 1:length(ex.args)
            if isa(ex.args[i], Expr)
                getlooporder!(order, ex.args[i])
            end
        end
    end
    order
end
function getloopvar(ex::Expr)
    ex.head == :(=) || error("Not an assignment expression")
    ex.args[1]
end

order = getlooporder(macroexpand(quote
        KernelTools.@tile (i,4) for i = 1:10 end
    end))
@test order == Vector{Symbol}[[:_kt_outer_i], [:_kt_inner_i]]
order = getlooporder(macroexpand(quote
        for j = 1:5
            KernelTools.@tile (i,4) for i = 1:10 end
        end
    end))
@test order == Vector{Symbol}[[:j],[:_kt_outer_i], [:_kt_inner_i]]
order = getlooporder(macroexpand(quote
            KernelTools.@tile (i,4) for j = 1:5, i = 1:10 end
    end))
@test order == Vector{Symbol}[[:_kt_outer_i], [:j, :_kt_inner_i]]
order = getlooporder(macroexpand(quote
            KernelTools.@tile (j,2,i,4) for j = 1:5, i = 1:10 end
    end))
@test order == Vector{Symbol}[[:_kt_outer_j,:_kt_outer_i], [:_kt_inner_j, :_kt_inner_i]]
order = getlooporder(macroexpand(quote
            KernelTools.@tile (i,4,j,2) for j = 1:5, i = 1:10 end
    end))
@test order == Vector{Symbol}[[:_kt_outer_i,:_kt_outer_j], [:_kt_inner_j, :_kt_inner_i]]

##### Functional tests of the macros

# Reductions & hoising
A = rand(3,5)
A1 = sum(A, 2)
A2 = sum(A, 1)
A1c = similar(A1); fill!(A1c, 0)
A2c = similar(A2); fill!(A2c, 0)
KernelTools.@loophoist for j = 1:size(A,2), i = 1:size(A,1)
    A1c[i,1] += A[i,j]
end
KernelTools.@loophoist for j = 1:size(A,2), i = 1:size(A,1)
    A2c[1,j] += A[i,j]
end
@test A1c == A1
@test A2c == A2

# Tiling
A = zeros(Int, 5, 6)
KernelTools.@tile (i,4,j,4) begin
    for j = 1:size(A,2), i = 1:size(A,1)
        A[i,j] = i + (j-1)*size(A,1)
    end
end
@test A == reshape(1:length(A), size(A))

# Tiling over a subset of variables
A = zeros(Int, 5, 6)
KernelTools.@tile (j,4) begin
    for j = 1:size(A,2), i = 1:size(A,1)
        A[i,j] = i + (j-1)*size(A,1)
    end
end
@test A == reshape(1:length(A), size(A))

## Tiling matrix multiplication
M = 8
K = 9
N = 7
A = rand(M,K)
B = rand(K,N)
C = A*B
Cc = similar(C)
# Tiling with just a body expression
fill!(Cc, 0)
KernelTools.@tile (m,4,k,4,n,4) begin
    for n = 1:N, k = 1:K, m = 1:M
        Cc[m,n] += A[m,k] * B[k,n]
    end
end
@test_approx_eq C Cc
fill!(Cc, 0)
KernelTools.@tile (m,4,k,4,n,4) begin
    KernelTools.@loophoist for n = 1:N,  k = 1:K, m = 1:M
        Cc[m,n] += A[m,k] * B[k,n]
    end
end
@test_approx_eq C Cc
# Tiling with a pre-expression
fill!(Cc, 0)
KernelTools.@tile (m,4,k,4,n,4) (Bt[n,k] = B[k,n]) begin
    for n = 1:N,  k = 1:K, m = 1:M
        Cc[m,n] += A[m,k] * Bt[n,k]
    end
end
@test_approx_eq C Cc
fill!(Cc, 0)
KernelTools.@tile (m,4,k,4,n,4) (Bt[n,k] = B[k,n]) begin
    KernelTools.@loophoist for n = 1:N,  k = 1:K, m = 1:M
        Cc[m,n] += A[m,k] * Bt[n,k]
    end
end
@test_approx_eq C Cc
# Don't tile over m
fill!(Cc, 0)
KernelTools.@tile (k,4,n,4) (Bt[n,k] = B[k,n]) begin
    KernelTools.@loophoist for n = 1:N,  k = 1:K, m = 1:M
        Cc[m,n] += A[m,k] * Bt[n,k]
    end
end
@test_approx_eq C Cc

# Image blur
A = rand(7,8)
blur_xy = similar(A)
interior_x, interior_y = 2:size(A,1)-1, 2:size(A,2)-1
KernelTools.@tile (x,4,y,4) (blur_y[x,y] = A[x,y-1]+A[x,y]+A[x,y+1]) begin
    for y = interior_y, x = interior_x
        blur_xy[x,y] = (1.0/9)*(blur_y[x-1,y] + blur_y[x,y] + blur_y[x+1,y])
    end
end
Ab = (A[interior_x, interior_y] + A[interior_x-1, interior_y] + A[interior_x+1, interior_y] +
      A[interior_x, interior_y-1] + A[interior_x-1, interior_y-1] + A[interior_x+1, interior_y-1] +
      A[interior_x, interior_y+1] + A[interior_x-1, interior_y+1] + A[interior_x+1, interior_y+1])/9
@test_approx_eq blur_xy[interior_x, interior_y] Ab


text:
cube_distance(a::Array{Int,1}) = Int(sum(abs.(a))/2)

function puzzle11(path::String=joinpath(@__DIR__, "..", "data/11.txt"))
    instructions = readdlm(path, ',')
    location = zeros(Int, 3)

    hexsteps = Dict(
        "n"  => [0, 1, -1],
        "ne" => [1, 0, -1],
        "se" => [1, -1, 0],
        "s"  => [0, -1, 1],
        "sw" => [-1, 0, 1],
        "nw" => [-1, 1, 0]
    )

    furthest_distance = 0
    for step in instructions
        location += hexsteps[step]
        furthest_distance = max(furthest_distance, cube_distance(location))
    end

    return [cube_distance(location), furthest_distance]

end


text:
# create an index for id6
import Base.ht_keyindex
import Base.ht_keyindex2

function buildindex_iter{T}(val::Vector{T})
  index = zeros(Int64, length(val))

  for (i,v) in enumerate(val)
    dindex = ht_keyindex2(index, v)

  end
  return index
end

function sumby_index{T,S}(index::Dict{T, Vector{Int64}}, val::Vector{S})
  return Dict{T,S}(k => sum(val[index[k]]) for k in keys(index))
end

@time iid4_iter = buildindex_iter(id4); #9
@time sumby_index(iid4_iter, v1); #8
@time sumby_index(iid4_iter, v1);  #8

@time iid6_iter = buildindex_iter(id6);  #259
@time sumby_index(iid6, v1); # 30.93
@time sumby_index(iid6, v3); #26.7


text:
# to run in parallel julia has to start with "julia -p auto" and the "JULIA_NUM_THREADS" env variable on the terminal has to be setted properly to the desired number of threads to run this program
# on powershell $env:JULIA_NUM_THREADS=4 is enough

@everywhere using Statistics
@everywhere using Dates
@everywhere include("FileOperations.jl")
@everywhere include("MathFunctions.jl")

#=
     run(star_id::Int64,dt::Float64=10e-4,refine::Float64=10e-3)

Search for star_id star on ./data and find a possible period on the interval 0.1-20 days with an initial precision of dt=0.001 days (1:26 seconds) and a refinement of 0.01 (864ms)

Returns tuple with period, ephemeris, and the squared magnitude of the Fourier transform on the period.
=#
@everywhere function run(star_id::Int64,dt::Float64=10e-3,refine::Float64=10e-3)
    local start::Float64=time()
    local JD::Array{Float64,1}
    local M::Array{Float64,1}
    local err::Array{Float64,1}
    local per::Float64
    local eph::Float64
    local Q::Float64 # period quality factor

    # import data
    JD,M,err=star_id |> OGLE_file_string |> read_float_table |> permutedims |> x->[x[:,i] for i in 1:3]

    # chech for gibberish
    if JD==[0.0,0.0,0.0]
        return (0.0,0.0,0.0)
    end

    # normalize magnitude
    normalize101!(M)

    # rough search
    per,eph,Q=find_period_ephemeris(JD,M,1.0:dt:30.0)

    # refined search
    per,eph,Q=find_period_ephemeris(JD,M,per-2dt:dt*refine:per+2dt)

    #println(star_id,"\tdone in\t",r5(time()-start))
    return (r5(per),r5(eph),r5(Q))
end

# read stars id from data filenames
const F_numbers=parse.(Int64,(x->x[14:17]).(readdir("./data")))

####### Compiling run function #########
run(2)                                ##
##### does not take this seriously #####

# declare an (empty) Channel for threads to look which star work on
const stars = RemoteChannel(()->Channel{Int}(32));
# declare an (empty) Channel of the type of run(star_id) to put results on
const results = RemoteChannel(()->Channel{Tuple}(32));

# function to find a period for all the stars
@everywhere function calculate_stars(stars,results)
    while true
        star_id = take!(stars)
        io = open("$(myid()).dat", append=true);
        s1=time()
        period,ephemeris,quality=run(star_id)
        write(io,"$star_id\t$(round(time()-s1; digits=6))\t$period\t$ephemeris\t$quality\n")
        close(io)
        if period==0.0
            put!(results,("$(myid()), $star_id no data",1))
        else
            put!(results, ("$(myid()), $star_id done",0))
        end
    end
end

# fill (define) the stars Channel
function make_stars(n::Int64)
    for i in F_numbers[1:n]
        put!(stars,i)
    end
end

n=parse(Int64,ARGS[1])

# creates a task on the general queue
@async make_stars(n)

for p in workers() # start tasks on the workers to process requests in parallel
    remote_do(calculate_stars, p, stars, results)
end

start=time()

while n > 0
    s_=take!(results)
    println(s_[1])
    global n = n - 1
end

# read data from workers output
data=[read_float_table("$i.dat") for i in workers()]

star_numbers=vcat([i[1,:] for i in data]...)
exec_times=vcat([i[2,:] for i in data]...)
periods=vcat([i[3,:] for i in data]...)
ephemerides=vcat([i[4,:] for i in data]...)
qualities=vcat([i[5,:] for i in data]...)

# remove workers output
for i in workers()
    rm("$i.dat")
end

# sort data by star_id
sorted_data=sortslices(hcat(star_numbers,exec_times,periods,ephemerides,qualities);dims=1,by = x -> x[1])

# write sorted data
file=open("results.dat",append=true)

for i in 1:length(sorted_data[:,1])
    write(file,join(sorted_data[i,:]," ")*"\n")
end

# Comparison between CPU and elapsed time
println("CPU:\t",r5(sum(exec_times)))
println("elapsed:\t",r5(time()-start))

max_stars_repo_name: santiagohenao/NUDFT-III

max_issues_repo_name: santiagohenao/NUDFT-III

max_forks_repo_name: santiagohenao/NUDFT-III


text:
include("4.1.0.1 (a sin)^m (b trg)^n.jl")
include("4.1.0.2 (a trg)^m (b tan)^n.jl")
include("4.1.0.3 (a csc)^m (b sec)^n.jl")
include("4.1.1.1 (a+b sin)^n.jl")
include("4.1.1.2 (g cos)^p (a+b sin)^m.jl")
include("4.1.1.3 (g tan)^p (a+b sin)^m.jl")
include("4.1.10 (c+d x)^m (a+b sin)^n.jl")
include("4.1.11 (e x)^m (a+b x^n)^p sin.jl")
include("4.1.12 (e x)^m (a+b sin(c+d x^n))^p.jl")
include("4.1.13 (d+e x)^m sin(a+b x+c x^2)^n.jl")
include("4.1.2.1 (a+b sin)^m (c+d sin)^n.jl")
include("4.1.2.2 (g cos)^p (a+b sin)^m (c+d sin)^n.jl")
include("4.1.2.3 (g sin)^p (a+b sin)^m (c+d sin)^n.jl")
include("4.1.3.1 (a+b sin)^m (c+d sin)^n (A+B sin).jl")
include("4.1.4.1 (a+b sin)^m (A+B sin+C sin^2).jl")
include("4.1.4.2 (a+b sin)^m (c+d sin)^n (A+B sin+C sin^2).jl")
include("4.1.5 trig^m (a cos+b sin)^n.jl")
include("4.1.6 (a+b cos+c sin)^n.jl")
include("4.1.7 (d trig)^m (a+b (c sin)^n)^p.jl")
include("4.1.8 trig^m (a+b cos^p+c sin^q)^n.jl")
include("4.1.9 trig^m (a+b sin^n+c sin^(2 n))^p.jl")

max_stars_repo_path: src/rules/4 Trig functions/4.1 Sine/.jl

max_issues_repo_path: src/rules/4 Trig functions/4.1 Sine/.jl

max_forks_repo_path: src/rules/4 Trig functions/4.1 Sine/.jl


text:
using ViewerGL
GL = ViewerGL

GL.VIEW([
	GL.GLFrame2,
 	GL.GLText("PLaSM.or.LAR?", GL.COLORS[1])
]);


text:
using Quine
using Test

@testset "Quine.jl" begin
    # Write your tests here.
end


text:
@testset "Output" begin

    @set printstyle repl
    @set output test_output.dat
    x = @capture_out output("Test Message")

    # Must print message and NOT create file
    @test x == "Test Message\n" && !isfile("test_output.dat")

    @set printstyle both
    x = @capture_out output("Test Message Two!")
    if isfile("test_output.dat")
        y = read("test_output.dat", String)
        # Must print message AND create file
        @test x === "Test Message Two!\n" && y == "Test Message Two!\n"
        rm("test_output.dat")
    else
        # File was not created properly
        @test false
    end

    @set printstyle file
    x = @capture_out output("Test Message Three!")
    if isfile("test_output.dat")
        y = read("test_output.dat", String)
        # Must print message AND NOT create file
        @test x === "" && y == "Test Message Three!\n"
        rm("test_output.dat")
    else
        # File was not created properly
        @test false
    end

    @set printstyle none
    x = @capture_out output("Test Message Four!")
    # Must neither print or create a file
    if isfile("test_output.dat")
        @test false
        rm("test_output.dat")
    else
        @test x == ""
    end

    # Check invalid printstyle
    @set printstyle invalid_printstyle
    @test_throws FermiException output("This will never print")

    @reset
end

max_stars_repo_name: FermiQC/Fermi

max_issues_repo_name: FermiQC/Fermi

max_forks_repo_name: FermiQC/Fermi


text:
# This file is a part of MGVI.jl, licensed under the MIT License (MIT).

module ModelFFTGP

import Base: *, adjoint
import AbstractFFTs
import FFTW: plan_r2r, DHT
import ForwardDiff
import Random: randn, MersenneTwister
import Distributions: Normal
import ValueShapes: NamedTupleDist
import Zygote
import LinearAlgebra: Diagonal

export model, true_params, starting_point

const _dims = 40
const _k = [i < _dims / 2 ? i : _dims-i for i = 0:_dims-1]

# Define the harmonic transform operator as a matrix-like object
const _ht = plan_r2r(zeros(_dims), DHT)

# Unfortunately neither Zygote nor ForwardDiff support planned Hartley
# transformations. While Zygote does not support AbstractFFTs.ScaledPlan,
# ForwardDiff does not overload the appropriate methods from AbstractFFTs.
function _plan_dual_product(trafo::AbstractFFTs.Plan, u::Vector{ForwardDiff.Dual{T, V, N}}) where {T, V, N}
    # Unpack AoS -> SoA
    vs = ForwardDiff.value.(u)
    ps = mapreduce(ForwardDiff.partials, hcat, u)
    # Actual computation
    val = trafo * vs
    jvp = [trafo*t[:] for t in eachrow(ps)]
    # Pack SoA -> AoS (depending on jvp, might need `eachrow`)
    return map((v, p) -> ForwardDiff.Dual{T}(v, p...), val, zip(jvp...))
end

function *(trafo::AbstractFFTs.Plan, u::Vector{ForwardDiff.Dual{T, V, N}}) where {T, V, N}
    _plan_dual_product(trafo, u)
end

function *(trafo::AbstractFFTs.ScaledPlan, u::Vector{ForwardDiff.Dual{T, V, N}}) where {T, V, N}
    _plan_dual_product(trafo, u)
end

Zygote.@adjoint function *(trafo::AbstractFFTs.ScaledPlan, xs)
    # Zygote has implementation for AbstractFFTs.Plan. ScaledPlan doesn't require norm mb?
    # https://github.com/FluxML/Zygote.jl/blob/2308bc8f30ccd6be913a054f7cc938c12a103512/src/lib/array.jl#L824
    # should be adjoint actually: return trafo * xs, Δ -> (nothing, trafo * Δ)
    return trafo * xs, Δ -> (nothing, trafo * Δ)
end

Zygote.@adjoint function inv(trafo::AbstractFFTs.Plan)
    inv_t = inv(trafo)
    return inv_t, function (Δ)
        return (- inv_t * Δ * inv_t,)
    end
end

function _correlated_field(ξ::Vector)
    loglogslope = 2.3
    P = @. 50 / (_k^loglogslope + 1)
    return inv(_ht) * (P .* ξ)
end

function _mean(ξ::Vector)
    return exp.(_correlated_field(ξ))
end

function model(ξ::Vector)
    return NamedTupleDist(unnamed=Normal.(_mean(ξ), 0.4))
end

# ξ := latent variables
const true_params = randn(MersenneTwister(128), _dims)
const starting_point = randn(MersenneTwister(12), _dims)

end # module


text:
################################
                                            ### Problem Set 0 - ECON 603 ###
                                            ################################

                                                #Sebastian Melo-Martin

# Call some useful packages
using Distributions
using Random
using QuantEcon
using Plots

############
# Question 1
############

function tauchen(N::Integer, ρ::T1, σ::T2, μ=zero(promote_type(T1, T2)), n_std::T3=3) where {T1 <: Real, T2 <: Real, T3 <: Real}
    # Discretized space
    a_bar = n_std * sqrt(σ^2 / (1 - ρ^2))
    y = range(-a_bar, stop=a_bar, length=N)
    d = y[2] - y[1]

    # Transition probabilities
    Π = zeros(promote_type(T1, T2), N, N)
    for row = 1:N
        # Do end points first
        Π[row, 1] = cdf(Normal(),((y[1] - ρ*y[row] + d/2) / σ))
        Π[row, N] = 1 - cdf(Normal(),((y[N] - ρ*y[row] - d/2) / σ))

        # Fill in columns
        for col = 2:N-1
            Π[row, col] = (cdf(Normal(),((y[col] - ρ*y[row] + d/2) / σ)) -
                           cdf(Normal(),((y[col] - ρ*y[row] - d/2) / σ)))
        end
    end

    return Π, y

end

z = tauchen(3,0.2,0.4)
prob_tau = z[1]
ybar_tau = @show Vector(z[2])

y_tau = zeros(1000,1)
pos = zeros(1000,1)
y_tau[1] = ybar_tau[rand(1:3)]

for i=2:1000
    x=rand()
    for j=1:3
        if y_tau[i-1]==ybar_tau[j]
            pos[i] = j
        end
    end
    prob_row = prob_tau[floor(Int, pos[i]),:]
    thr_1 = prob_row[1]
    thr_2 = thr_1 + prob_row[2]
    if x<=thr_1
        y_tau[i]=ybar_tau[1]
    elseif x>thr_1 && x<=thr_2
        y_tau[i] = ybar_tau[2]
    else
        y_tau[i] = ybar_tau[3]
    end
end

t = collect(range(1,stop=1000,length=1000))
plot(t, y_tau)

##############
# Question 2 #
##############

#Getting initial vector and probability matrix
function rouwenhorst(N::Integer, ρ::Real, σ::Real, μ::Real)

    σ_y = σ / sqrt(1-ρ^2)
    p  = (1+ρ)/2
    Θ = [p 1-p; 1-p p]
    ψ = sqrt(N-1) * σ_y
    m = μ / (1 - ρ)

    state_values, p = _rouwenhorst(p, p, m, ψ, N)

    return p, state_values
end

function _rouwenhorst(p::Real, q::Real, m::Real, Δ::Real, n::Integer)
    if n == 2
        return [m-Δ, m+Δ],  [p 1-p; 1-q q]
    else
        _, θ_nm1 = _rouwenhorst(p, q, m, Δ, n-1)
        θN = p    *[θ_nm1 zeros(n-1, 1); zeros(1, n)] +
             (1-p)*[zeros(n-1, 1) θ_nm1; zeros(1, n)] +
             q    *[zeros(1, n); zeros(n-1, 1) θ_nm1] +
             (1-q)*[zeros(1, n); θ_nm1 zeros(n-1, 1)]

        θN[2:end-1, :] ./= 2

        return range(m-Δ, stop=m+Δ, length=n), θN
    end
end

z = rouwenhorst(3,0.2,0.4,0)
prob_rou = z[1]
ybar_rou = @show Vector(z[2])

#Simulating the time series
y_rou = zeros(1000,1)
pos = zeros(1000,1)
y_rou[1] = ybar_rou[rand(1:3)]

for i=2:1000
    x=rand()
    for j=1:3
        if y_rou[i-1]==ybar_rou[j]
            pos[i] = j
        end
    end
    prob_row = prob_rou[floor(Int, pos[i]),:]
    thr_1 = prob_row[1]
    thr_2 = prob_row[1] + prob_row[2]
    if x<=thr_1
        y_rou[i]=ybar_rou[1]
    elseif x>thr_1 && x<=thr_2
        y_rou[i] = ybar_rou[2]
    else
        y_rou[i] = ybar_rou[3]
    end
end

t = collect(range(1,stop=1000,length=1000))
plot(t, y_rou)

##############
# Question 3 #
##############

function normal_ar(ρ, N)
    d = Normal(0, 0.16)
    y = [0.0 for i = 1:N]
    Random.seed!(641993)
    ϵ = rand(d, N)
    for i in 1:(N-1)
        y[i+1] = ρ*y[i] + ϵ[i]
    end
    return y
end

ρ=0.2
N=1000
y_ar = normal_ar(ρ,N)
t = collect(range(1,stop=1000,length=1000))

plot(t, y_ar)

##############
# Question 4 #
##############

# Tauchen
z = tauchen(10,0.2,0.4)
prob_tau = z[1]
ybar_tau = @show Vector(z[2])

y_tau = zeros(1000,1)
pos = zeros(1000,1)
y_tau[1] = ybar_tau[rand(1:10)]

for i=2:1000
    x=rand()
    for j=1:10
        if y_tau[i-1]==ybar_tau[j]
            pos[i] = j
        end
    end
    prob_row = prob_tau[floor(Int, pos[i]),:]
    thr_1 = prob_row[1]
    thr_2 = thr_1 + prob_row[2]
    thr_3 = thr_2 + prob_row[3]
    thr_4 = thr_3 + prob_row[4]
    thr_5 = thr_4 + prob_row[5]
    thr_6 = thr_5 + prob_row[6]
    thr_7 = thr_6 + prob_row[7]
    thr_8 = thr_7 + prob_row[8]
    thr_9 = thr_8 + prob_row[9]
    if x<=thr_1
        y_tau[i]=ybar_tau[1]
    elseif x>thr_1 && x<=thr_2
        y_tau[i] = ybar_tau[2]
    elseif x>thr_2 && x<=thr_3
        y_tau[i] = ybar_tau[3]
    elseif x>thr_3 && x<=thr_4
        y_tau[i] = ybar_tau[4]
    elseif x>thr_4 && x<=thr_5
        y_tau[i] = ybar_tau[5]
    elseif x>thr_5 && x<=thr_6
        y_tau[i] = ybar_tau[6]
    elseif x>thr_6 && x<=thr_7
        y_tau[i] = ybar_tau[7]
    elseif x>thr_7 && x<=thr_8
        y_tau[i] = ybar_tau[8]
    elseif x>thr_8 && x<=thr_9
        y_tau[i] = ybar_tau[9]
    else
        y_tau[i] = ybar_tau[10]
    end
end

t = collect(range(1,stop=1000,length=1000))
plot(t, y_tau)

# Rouwenhorst
z = rouwenhorst(10,0.2,0.4,0)
prob_rou = z[1]
ybar_rou = @show Vector(z[2])

y_rou = zeros(1000,1)
pos = zeros(1000,1)
y_rou[1] = ybar_tau[rand(1:10)]

for i=2:1000
    x=rand()
    for j=1:10
        if y_rou[i-1]==ybar_rou[j]
            pos[i] = j
        end
    end
    prob_row = prob_rou[floor(Int, pos[i]),:]
    thr_1 = prob_row[1]
    thr_2 = thr_1 + prob_row[2]
    thr_3 = thr_2 + prob_row[3]
    thr_4 = thr_3 + prob_row[4]
    thr_5 = thr_4 + prob_row[5]
    thr_6 = thr_5 + prob_row[6]
    thr_7 = thr_6 + prob_row[7]
    thr_8 = thr_7 + prob_row[8]
    thr_9 = thr_8 + prob_row[9]
    if x<=thr_1
        y_rou[i]=ybar_rou[1]
    elseif x>thr_1 && x<=thr_2
        y_rou[i] = ybar_rou[2]
    elseif x>thr_2 && x<=thr_3
        y_rou[i] = ybar_rou[3]
    elseif x>thr_3 && x<=thr_4
        y_rou[i] = ybar_rou[4]
    elseif x>thr_4 && x<=thr_5
        y_rou[i] = ybar_rou[5]
    elseif x>thr_5 && x<=thr_6
        y_rou[i] = ybar_rou[6]
    elseif x>thr_6 && x<=thr_7
        y_rou[i] = ybar_rou[7]
    elseif x>thr_7 && x<=thr_8
        y_rou[i] = ybar_rou[8]
    elseif x>thr_8 && x<=thr_9
        y_rou[i] = ybar_rou[9]
    else
        y_rou[i] = ybar_rou[10]
    end
end

t = collect(range(1,stop=1000,length=1000))
plot(t, y_rou)

##############
# Question 5 #
##############

#= The coding for this question is already included previously in a loop
    =#


text:
######################################################################
# gradients.jl
# generation of gradients
#
# TODO: this whole file needs work
#  While technically correct, it is an eyesore and can probably
#   be greatly simplified.
######################################################################

# Only first two states matter for ergodic score and barrier penalty
# Assumes ad has been initialized with zeros; that is, ad[3:end, ni] = 0.0
function gradients!(ad::MF, bd::MF, em::ErgodicManager, tm::TrajectoryManager, xd::VVF, ud::VVF)
	ck = decompose(em, xd)
	gradients!(ad, bd, em, tm, xd, ud, ck)
end

function gradients!(ad::MF, bd::MF, em::ErgodicManager, tm::TrajectoryManager, xd::VVF, ud::VVF, ck)
	ni =  1
	#ck = decompose(em, xd)
	for n = 0:(tm.N-1)

		# ergodic gradients
		an = compute_ans(em, xd, tm, n, ck)
		for i = 1:length(an)
			ad[i,ni] = an[i]
		end

		# quadratic boundary
		if tm.barrier_cost > 0.0
			xnx = xd[n+1][1]
			xny = xd[n+1][2]
			xmax = x_max(em)
			xmin = x_min(em)
			ymax = y_max(em)
			ymin = y_min(em)
			if (xnx > xmax)
				ad[1,ni] += tm.barrier_cost * 2.0 * (xnx - xmax)
			elseif xnx < xmin
				ad[1,ni] += tm.barrier_cost * 2.0 * (xnx - xmin)
			end
			if xny > ymax
				ad[2,ni] += tm.barrier_cost * 2.0 * (xny - ymax)
			elseif xny < ymin
				ad[2,ni] += tm.barrier_cost * 2.0 * (xny - ymin)
			end
		end

		# control gradients
		bd[:,ni] = tm.h * tm.R * ud[ni]

		ni += 1
	end

	n = tm.N
	an = compute_ans(em, xd, tm, n, ck)
	for i = 1:length(an)
		ad[i,ni] = an[i]
	end
	if tm.barrier_cost > 0.0
		xnx = xd[n+1][1]
		xny = xd[n+1][2]
		xmax = x_max(em)
		xmin = x_min(em)
		ymax = y_max(em)
		ymin = y_min(em)
		if (xnx > xmax)
			ad[1,ni] += tm.barrier_cost * 2.0 * (xnx - xmax)
		elseif xnx < xmin
			ad[1,ni] += tm.barrier_cost * 2.0 * (xnx - xmin)
		end
		if xny > ymax
			ad[2,ni] += tm.barrier_cost * 2.0 * (xny - ymax)
		elseif xny < ymin
			ad[2,ni] += tm.barrier_cost * 2.0 * (xny - ymin)
		end
	end
end

# This version is for multi-agent method
function gradients!(ad::MF, bd::MF, em::ErgodicManager, vtm::Vector{TrajectoryManager}, xd::VVF, ud::VVF)

	# must first create xds, where each agent trajectory is its own vector
	xds, uds = vvf2vvvf(xd, ud, vtm)

	ck = decompose(em, xds)
	n_idx = 1
	m_idx = 1
	N = length(xd) - 1
	num_agents = length(vtm)
	for j = 1:num_agents
		tm = vtm[j]
		adt = zeros(tm.dynamics.n, N+1)
		bdt = zeros(tm.dynamics.m, N)
		n_rows = tm.dynamics.n
		m_rows = tm.dynamics.m
		gradients!(adt, bdt, em, tm, xds[j], uds[j], ck)
		ad[n_idx:(n_idx+n_rows-1), :] = adt / num_agents
		bd[m_idx:(m_idx+m_rows-1), :] = bdt / num_agents
		n_idx += n_rows
		m_idx += m_rows
	end
end

# This version is for decentralized multi agents
function gradients!(ad::MF, bd::MF, em::ErgodicManager, vtm::Vector{TrajectoryManager}, xd::VVF, ud::VVF)

	# must first create xds, where each agent trajectory is its own vector
	xds, uds = vvf2vvvf(xd, ud, vtm)

	ck = decompose(em, xds)
	n_idx = 1
	m_idx = 1
	N = length(xd) - 1
	num_agents = length(vtm)
	for j = 1:num_agents
		tm = vtm[j]
		adt = zeros(tm.dynamics.n, N+1)
		bdt = zeros(tm.dynamics.m, N)
		n_rows = tm.dynamics.n
		m_rows = tm.dynamics.m
		gradients!(adt, bdt, em, tm, xds[j], uds[j], ck)
		ad[n_idx:(n_idx+n_rows-1), :] = adt / num_agents
		bd[m_idx:(m_idx+m_rows-1), :] = bdt / num_agents
		n_idx += n_rows
		m_idx += m_rows
	end
end

# TODO: make this functional
function add_barrier_gradient!(ad::MF, em::ErgodicManager, tm::TrajectoryManager)
	if tm.barrier_cost > 0.0
		xnx = xd[n+1][1]
		xny = xd[n+1][2]
		xmax = x_max(em)
		xmin = x_min(em)
		ymax = y_max(em)
		ymin = y_min(em)
		if (xnx > xmax)
			ad[1,ni] += tm.barrier_cost * 2.0 * (xnx - xmax)
		elseif xnx < xmin
			ad[1,ni] += tm.barrier_cost * 2.0 * (xnx - xmin)
		end
		if xny > ymax
			ad[2,ni] += tm.barrier_cost * 2.0 * (xny - ymax)
		elseif xny < ymin
			ad[2,ni] += tm.barrier_cost * 2.0 * (xny - ymin)
		end
	end
end

# computes a_n, derivative of c_k wrt x_n
# returns tuple containing elements of an
# TODO: I shouldn't need to use tm
function compute_ans(em::ErgodicManagerR2, xd::VVF, tm::TrajectoryManager, n::Int, ck::Matrix{Float64})
	x = xd[n + 1][1]
	y = xd[n + 1][2]

	Lx = em.domain.lengths[1]
	Ly = em.domain.lengths[2]

	an_x = 0.0
	an_y = 0.0

	xm = x_min(em)
	ym = y_min(em)

	for k1 = 0:em.K
		for k2 = 0:em.K
			hk = em.hk[k1+1,k2+1]

			dFk_dxn1 = -k1*pi*sin(k1*pi*(x-xm)/Lx)*cos(k2*pi*(y-ym)/Ly) / (hk*Lx)
			dFk_dxn2 = -k2*pi*cos(k1*pi*(x-xm)/Lx)*sin(k2*pi*(y-ym)/Ly) / (hk*Ly)

			c = em.Lambda[k1+1,k2+1] * (ck[k1+1,k2+1] - em.phik[k1+1,k2+1])
			an_x += c*dFk_dxn1
			an_y += c*dFk_dxn2
		end
	end
	an_x *= 2.0/(tm.N+1)
	an_y *= 2.0/(tm.N+1)
	return an_x, an_y
end

function compute_ans(em::ErgodicManagerR3, xd::VVF, tm::TrajectoryManager, n::Int, ck::Array{Float64,3})
	x = xd[n + 1][1]
	y = xd[n + 1][2]
	z = xd[n + 1][3]

	Lx = em.domain.lengths[1]
	Ly = em.domain.lengths[2]
	Lz = em.domain.lengths[3]

	an_x = 0.0
	an_y = 0.0
	an_z = 0.0

	xm = x_min(em)
	ym = y_min(em)
	zm = z_min(em)

	for k1 = 0:em.K
		for k2 = 0:em.K
			for k3 = 0:em.K
				hk = em.hk[k1+1,k2+1,k3+1]

				c1 = cos(k1*pi*(x-xm)/Lx)
				c2 = cos(k2*pi*(y-ym)/Ly)
				c3 = cos(k3*pi*(z-zm)/Lz)

				dFk_dxn1 = -k1*pi* sin(k1*pi*(x-xm)/Lx)*c2*c3 / (hk*Lx)
				dFk_dxn2 = -k2*pi*c1* sin(k2*pi*(y-ym)/Ly)*c3 / (hk*Ly)
				dFk_dxn3 = -k3*pi*c1*c2* sin(k3*pi*(z-zm)/Lz) / (hk*Lz)

				c = em.Lambda[k1+1,k2+1,k3+1] * (ck[k1+1,k2+1,k3+1] - em.phik[k1+1,k2+1,k3+1])
				an_x += c*dFk_dxn1
				an_y += c*dFk_dxn2
				an_z += c*dFk_dxn3
			end
		end
	end
	an_x *= 2.0/(tm.N+1)
	an_y *= 2.0/(tm.N+1)
	an_z *= 2.0/(tm.N+1)
	return an_x, an_y, an_z
end

function compute_ans(em::ErgodicManagerR2T, xd::VVF, tm::TrajectoryManager, n::Int, ck::Array{Float64,3})
	x = xd[n + 1][1]
	y = xd[n + 1][2]
	t = n*tm.h

	Lx = em.domain.lengths[1]
	Ly = em.domain.lengths[2]
	T = tm.N*tm.h

	an_x = 0.0
	an_y = 0.0

	xm = x_min(em)
	ym = y_min(em)

	for k1 = 0:em.K
		for k2 = 0:em.K
			for k3 = 0:em.K
				hk = em.hk[k1+1,k2+1,k3+1]

				c1 = cos(k1*pi*(x-xm)/Ly)
				c2 = cos(k2*pi*(y-ym)/Ly)
				c3 = cos(k3*pi*t/T)

				dFk_dxn1 = -k1*pi* sin(k1*pi*(x-xm)/Lx)*c2*c3 / (hk*Lx)
				dFk_dxn2 = -k2*pi*c1* sin(k2*pi*(y-ym)/Ly)*c3 / (hk*Ly)

				c = em.Lambda[k1+1,k2+1,k3+1] * (ck[k1+1,k2+1,k3+1] - em.phik[k1+1,k2+1,k3+1])
				an_x += c*dFk_dxn1
				an_y += c*dFk_dxn2
			end
		end
	end
	an_x *= 2.0/(tm.N+1)
	an_y *= 2.0/(tm.N+1)
	return an_x, an_y
end

function compute_ans(em::ErgodicManagerSE2, xd::VVF, tm::TrajectoryManager, n::Int, ck)
	x = xd[n + 1][1]
	y = xd[n + 1][2]
	z = xd[n + 1][3]

	an_x = 0.0
	an_y = 0.0
	an_z = 0.0

	i = float(im)

	# polar coords
	r2 = x*x + y*y	# range squared
	r = sqrt(r2)
	psi = atan2(y, x)

    # TODO: this needs to be scrubbed real well, especially since I changed
    #        from em.M being an integer to a range
	#for m = 0:em.M
    #    for n = 0:em.N
    for (mi,m) in enumerate(em.M)
        for (ni,n) in enumerate(em.N)
			# commonly used values
			inm = i^(n-m)
			expi = exp(i * (m*psi + (n-m)*z) )
			#for p = 0:em.P
            for (pl,p) in enumerate(em.P)

				pr = p*r	# commonly used
				bjmn = besselj(m-n, pr)
				dJ_dpr = 0.5 * (besselj(m-n-1, pr) - besselj(m-n+1, pr))

				# TODO: we can reuse even more of this stuff...
				dF_dx = expi * dJ_dpr * p * (x/r)
				dF_dx += bjmn * expi*m*i * (-y/r2)
				dF_dx = inm * dF_dx

				dF_dy = expi * dJ_dpr * p * (y/r)
				dF_dy += bjmn * expi*m*i * (x/r2)
				dF_dy = inm * dF_dy

				dF_dz = inm * bjmn * expi * i * (n-m)

				#c = em.Lambda[m+1,n+1,p+1] * (ck[m+1,n+1,p+1] - em.phik[m+1,n+1,p+1])
				c = em.Lambda[mi,ni,pl] * (ck[mi,ni,pl] - em.phik[mi,ni,pl])
				c_r = real(c)
				c_i = imag(c)

				# recall that c is a complex number...
				an_x += c_r*real(dF_dx) + c_i*imag(dF_dx)
				an_y += c_r*real(dF_dy) + c_i*imag(dF_dy)
				an_z += c_r*real(dF_dz) + c_i*imag(dF_dz)
			end
		end
	end
	an_x *= 2.0 / (tm.N+1)
	an_y *= 2.0 / (tm.N+1)
	an_z *= 2.0 / (tm.N+1)
	return an_x, an_y, an_z
end