text:
using FixedPointDecimals
using FixedPointDecimals: FD, value
using Test
using Printf
using Base.Checked: checked_mul

include("utils.jl")

const SFD2 = FixedDecimal{Int16, 2}
const SFD4 = FixedDecimal{Int16, 4}
const FD1 = FixedDecimal{Int, 1}
const FD2 = FixedDecimal{Int, 2}
const FD3 = FixedDecimal{Int, 3}
const FD4 = FixedDecimal{Int, 4}
const WFD2 = FixedDecimal{Int128, 2}
const WFD4 = FixedDecimal{Int128, 4}

const CONTAINER_TYPES = Base.BitInteger_types  # Integer concrete subtypes which are bits

# these arrays should be kept sorted manually
const keyvalues = Dict(
    FD2 => [typemin(FD2),  # near minimum range
            FD2(-289.64),  # randomly generated
            FD2(-1),
            FD2(-0.01),    # near zero
            FD2(0),
            FD2(0.01),
            FD2(0.14),     # fraction-like
            FD2(0.33),
            FD2(0.5),
            FD2(1),        # one
            FD2(592.57),   # randomly generated
            typemax(FD2)], # near maximum range
    WFD4 => [typemin(WFD4),
             reinterpret(WFD4, -157030247204331916472131926508185768261),
             reinterpret(WFD4, -64628160301714851880492570261792833470),
             reinterpret(WFD4, -11679287782747983139362515984380939763),
             WFD4(-1),
             WFD4(-0.0001),
             WFD4(0),
             WFD4(0.0001),
             WFD4(0.01),
             WFD4(1),
             reinterpret(WFD4, 164435910993133062409572187012743929911),
             typemax(WFD4)])

# Floating point values written as integer strings. Useful for testing behaviours of
# trunc, floor, and ceil.
const INTS = Dict(
    v => replace(@sprintf("%.200f", v), "." => "")
    for v in [
        1.22,
        1.23,
        1.51,
        2.2,
        2.3,
    ]
)
const smaller_than_decimal = [1.22, 1.23, 2.3]
const bigger_than_decimal = [1.51, 2.2]

# numbers that may cause overflow
islarge(x) = x == typemin(x) || abs(x) > 1000

# numbers that can never cause overflow
issmall(x) = -1 < x ≤ 1

function parse_int(::Type{FD{T, f}}, val::AbstractString; ceil::Bool=false) where {T, f}
    reinterpret(FD{T, f}, parse(T, val[1:(f + 1)]) + T(ceil))
end

@testset "FixedPointDecimals" begin

# Basic tests for the methods created above
@testset "alt" begin
    @test trunc_alt(FD2, 0.0) == FD2(0)
    @test floor_alt(FD2, 0.0) == FD2(0)
    @test ceil_alt(FD2,  0.0) == FD2(0)

    @test trunc_alt(FD2, 2.149) == FD2(2.14)
    @test floor_alt(FD2, 2.149) == FD2(2.14)
    @test ceil_alt(FD2,  2.149) == FD2(2.15)

    @test trunc_alt(FD2, -2.149) == FD2(-2.14)
    @test floor_alt(FD2, -2.149) == FD2(-2.15)
    @test ceil_alt(FD2,  -2.149) == FD2(-2.14)

    @test trunc_alt(FD2, nextfloat(0.0)) == FD2(0)
    @test floor_alt(FD2, nextfloat(0.0)) == FD2(0)
    @test ceil_alt(FD2,  nextfloat(0.0)) == FD2(0.01)

    @test trunc_alt(FD2, prevfloat(0.0)) == FD2(0)
    @test floor_alt(FD2, prevfloat(0.0)) == FD2(-0.01)
    @test ceil_alt(FD2,  prevfloat(0.0)) == FD2(0)
end

@testset "max_exp10" begin
    @test FixedPointDecimals.max_exp10(Int8) == 2
    @test FixedPointDecimals.max_exp10(Int64) == 18
    @test FixedPointDecimals.max_exp10(Int128) == 38
    @test FixedPointDecimals.max_exp10(UInt8) == 2
    @test FixedPointDecimals.max_exp10(UInt64) == 19
    @test FixedPointDecimals.max_exp10(UInt128) == 38
    @test FixedPointDecimals.max_exp10(BigInt) == -1

    for T in CONTAINER_TYPES
        x = FixedPointDecimals.max_exp10(T)
        @test T(10)^x == widen(T(10))^x
    end

    @testset "custom integer types" begin
        @eval begin
            primitive type Int24 <: Integer 24 end
            Base.typemax(::Type{Int24}) = 2^24
            Base.widen(::Type{Int24}) = Int32
        end

        @test FixedPointDecimals.max_exp10(Int24) == 7

        # Note: we're just pretending that this is unbounded
        @eval primitive type IntUnbounded <: Integer 256 end
        @test_throws MethodError FixedPointDecimals.max_exp10(IntUnbounded)
    end
end

# ensure that the coefficient multiplied by the highest and lowest representable values of
# the container type do not result in overflow.
@testset "coefficient" begin
    @testset "overflow $T" for T in CONTAINER_TYPES
        f = FixedPointDecimals.max_exp10(T)
        powt = FixedPointDecimals.coefficient(FD{T, f})
        @test powt % 10 == 0
        @test checked_mul(widen(powt), typemax(T)) == widemul(powt, typemax(T))
        @test checked_mul(widen(powt), typemin(T)) == widemul(powt, typemin(T))
    end
end

@testset "constructor" begin
    @testset "invalid $T" for T in CONTAINER_TYPES
        f = FixedPointDecimals.max_exp10(T) + 1
        @test_throws ArgumentError reinterpret(FD{T,f}, 0)
        @test_throws ArgumentError reinterpret(FD{T,-1}, 0)
    end
end

@testset "conversion" begin
    @testset for x in keyvalues[FD2]
        @testset for T in [Rational{Int128}, WFD2, WFD4]
            @test convert(FD2, convert(T, x)) == x
            @test T(x) == convert(T, x)
        end
        if 0 ≤ abs(x) < 2
            @testset for T in [SFD2, SFD4, FD4]
                @test convert(FD2, convert(T, x)) == x
            end
        end
    end

    @testset "to float" begin
        # Convert the rational 5//7 into a FixedDecimal with as much precision as we can
        # without using BigInt.
        T = Int128
        f = FixedPointDecimals.max_exp10(T)
        powt = FixedPointDecimals.coefficient(FD{T,f})
        val = T(trunc(BigInt, widemul(5//7, powt)))

        fd = reinterpret(FD{T,f}, val)
        @test convert(Float64, fd) != convert(BigFloat, fd)
        @test convert(Float64, fd) == T(val) / T(powt)
        @test convert(BigFloat, fd) == BigInt(val) / BigInt(powt)
    end

    @testset "to rational" begin
        fd = reinterpret(FD2, 25)
        @test convert(Rational, fd) == 1//4
    end

    @testset "invalid" begin
        @test_throws InexactError convert(FD2, FD4(0.0001))
        @test_throws InexactError convert(FD4, typemax(FD2))
        @test_throws InexactError convert(SFD2, typemax(FD2))
        @test_throws InexactError convert(FD2, 1//3)
        @test_throws InexactError convert(FD{Int8,1}, 1//4)
    end

    @testset "limits of $T" for T in CONTAINER_TYPES
        max_exp = FixedPointDecimals.max_exp10(T)
        f = max_exp
        powt = widen(FixedPointDecimals.coefficient(FD{T,f}))

        # Smallest positive integer which is out-of-bounds for the FD
        x = max_exp - f + 1
        oob = T(10)^(x > 0 ? x : 0)

        # ideally we would just use `typemax(T)` but due to precision issues with
        # floating-point its possible the closest float will exceed `typemax(T)`.
        # Note: we should be doing `trunc(T, ...)` but truncating a BigFloat can be
        # problematic (https://github.com/JuliaLang/julia/issues/21914)
        max_int = trunc(BigInt, prevfloat(typemax(T) / powt) * powt)
        min_int = trunc(BigInt, nextfloat(typemin(T) / powt) * powt)

        @test max_int <= typemax(T)
        @test value(convert(FD{T,f}, max_int / powt)) == max_int
        @test min_int >= typemin(T)
        @test value(convert(FD{T,f}, min_int / powt)) == min_int

        @test convert(FD{T,f}, typemax(T) // powt) == reinterpret(FD{T,f}, typemax(T))
        @test convert(FD{T,f}, typemin(T) // powt) == reinterpret(FD{T,f}, typemin(T))

        @test_throws InexactError convert(FD{T,f}, oob)

        # Converting to a floating-point
        fd = reinterpret(FD{T,f}, typemax(T))
        @test convert(Float32, fd)  == Float32(typemax(T) / powt)
        @test convert(Float64, fd)  == Float64(typemax(T) / powt)
        @test convert(BigFloat, fd) == BigInt(typemax(T)) / powt

        fd = reinterpret(FD{T,f}, typemin(T))
        @test convert(Float32, fd)  == Float32(typemin(T) / powt)
        @test convert(Float64, fd)  == Float64(typemin(T) / powt)
        @test convert(BigFloat, fd) == BigInt(typemin(T)) / powt

        # Converting to a rational
        fd = reinterpret(FD{T,f}, typemax(T))
        @test convert(Rational, fd)  == typemax(T) // powt

        fd = reinterpret(FD{T,f}, typemin(T))
        @test convert(Rational, fd)  == typemin(T) // powt

        # The following tests require that the number of decimal places allow for
        # `-10 < x < 10` where x is a FD{T,f}. Needed to test `convert(::FD, ::Integer)`.
        max_int = typemax(T) ÷ powt * powt
        min_int = typemin(T) ÷ powt * powt

        @test convert(FD{T,f}, max_int ÷ powt) == reinterpret(FD{T,f}, max_int)
        @test convert(FD{T,f}, min_int ÷ powt) == reinterpret(FD{T,f}, min_int)

        @test_throws InexactError convert(FD{T,f}, max_int ÷ powt + oob)
        @test_throws InexactError convert(FD{T,f}, min_int ÷ powt - oob)  # Overflows with Unsigned
    end

    @testset "limits from $U to $T" for T in CONTAINER_TYPES, U in CONTAINER_TYPES
        f = FixedPointDecimals.max_exp10(T)
        g = FixedPointDecimals.max_exp10(U)
        powt = div(
            FixedPointDecimals.coefficient(FD{T, f}),
            FixedPointDecimals.coefficient(FD{U, g}),
        )

        val = typemax(U)
        expected = widemul(typemax(U), powt)

        # Mixed usage of signed and unsigned types makes testing with typemin hard.
        if f >= g && expected <= typemax(T)
            @test convert(FD{T,f}, reinterpret(FD{U,g}, val)) == reinterpret(FD{T,f}, expected)
        else
            @test_throws InexactError convert(FD{T,f}, reinterpret(FD{U,g}, val))
        end
    end
end

@testset "promotion" begin
    @test 1//10 + FD2(0.1) === 1//5
    @test 0.1 + FD2(0.1) === 0.2
    @test 1 + FD2(0.1) === FD2(1.1)
    @test FD2(0.1) + FD4(0.0001) === FD4(0.1001)
    @test WFD2(0.1) + FD4(0.0001) === WFD4(0.1001)
end

@testset "float" begin
    @test float(-one(SFD2)) === -1.0f0
    @test float(zero(SFD2)) === 0.0f0
    @test float(one(SFD2)) === 1.0f0
    @test float(-one(FD2)) === -1.0
    @test float(zero(FD2)) === 0.0
    @test float(one(FD2)) === 1.0
end

@testset "comparison" begin
    @testset for T in [FD2, WFD4]
        @testset for (i, x) in enumerate(keyvalues[T])
            @test x == x
            @testset for y in keyvalues[T][i+1:end]
                @test x ≠ y
                @test x < y
                @test x ≤ y
                @test y ≠ x
                @test y > x
                @test y ≥ x
            end
        end
    end
end

@testset "traits" begin
    @testset "zero, one" begin
        @test FD2(0) == zero(FD2)
        @test FD2(42.42) + FD2(0) == FD2(42.42)
        @test FD2(1) == one(FD2)
        @test FD2(42.42) * FD2(1) == FD2(42.42)
    end

    @testset "eps, floatmin, floatmax" begin
        @test floatmin(FD2) == eps(FD2) == FD2(0.01)
        @test eps(FD2(1.11)) == FD2(0.01)
        for x in keyvalues[FD2]
            if x ≠ typemax(FD2)
                @test x + eps(x) > x
            end
            if x ≠ typemin(FD2)
                @test x - eps(x) < x
                if x ≠ 0
                    @test floatmin(FD2) ≤ abs(x) ≤ floatmax(FD2)
                end
            end
        end
    end
end

@testset "addition" begin
    @test FD2(0) + FD2(0) == FD2(0)
    @test FD2(1.11) + FD2(2.22) == FD2(3.33)
    @test FD2(0.01) + FD2(0.01) == FD2(0.02)
    @test FD2(0.01) + FD2(-0.01) == FD2(0)

    # overflow
    @test typemax(FD2) + eps(FD2) == typemin(FD2)
end

@testset "subtraction" begin
    for x in keyvalues[FD2]
        @test x - x == 0
        for y in keyvalues[FD2]
            @test x + y - y == x
            @test y + x - y == x
        end
    end
end

@testset "multiply" begin
    @testset "with integer, $T" for T in [FD2, WFD4]
        for x in keyvalues[T]
            @test 1 * x == x * 1 == x
            @test one(x) * x == x * one(x) == x
            @test (-1) * x == x * (-1) == -x
            @test 2 * x == x + x == (one(x) + one(x)) * x
        end
    end

    @testset "binary" begin
        @test FD2(0.33) * FD2(1.00) == FD2(0.33)
        @test FD2(0.33) * FD2(3.00) == FD2(0.99)
        @test FD2(0.33) * FD2(0.50) == FD2(0.16)
        @test FD2(0.33) * FD2(0.33) == FD2(0.11)
        @test FD2(0.67) * FD2(0.67) == FD2(0.45)
    end

    @testset "key values $T" for T in [FD2, WFD4]
        totest = [(x, y) for x in keyvalues[T] for y in keyvalues[T] if
                  issmall(x) || issmall(y) || (!islarge(x) && !islarge(y))]
        @testset for (x, y) in totest
            # test the multiplication result is correctly rounded
            @test x * y == round(typeof(x), Base.widemul(x, y))
        end
        @test prod(keyvalues[T]) == 0
    end

    @testset "without promotion" begin
        @test_throws InexactError FD{Int8,1}(20)
        @test 20 * FD{Int8,1}(0.1) == FD{Int8,1}(2.0)
        @test FD{Int8,1}(0.1) * 20 == FD{Int8,1}(2.0)
    end

    @testset "limits of $T" for T in CONTAINER_TYPES
        f = FixedPointDecimals.max_exp10(T)
        scalar = convert(FD{T,f}, 1 // 10)  # 0.1

        # Since multiply will round the result we'll make sure our value does not
        # always rounds down.
        max_int = typemax(T) - (typemax(T) % 10)
        min_int = typemin(T) - (typemin(T) % 10)

        @test reinterpret(FD{T,f}, max_int) * scalar ==
            reinterpret(FD{T,f}, div(max_int, 10))
        @test reinterpret(FD{T,f}, min_int) * scalar ==
            reinterpret(FD{T,f}, div(min_int, 10))
    end
end

@testset "division" begin
    @testset "division by 1" begin
        @testset for x in keyvalues[FD2]
            @test x / one(x) == x

            # signed integers using two's complement have one additional negative value
            if x < 0 && x == typemin(x)
                @test_throws InexactError x / -one(x)
            else
                @test x / -one(x) == -x
            end
        end
    end

    @testset "division by 2" begin
        # even targets
        for x in FD2[-0.02, 0, 0.02, 1.00]
            for y in [2x-eps(x), 2x, 2x+eps(x)]
                @test y / 2 == y / 2one(y) == x == y * FD2(0.5)
            end
        end

        # odd targets
        for x in FD2[-0.01, 0.01, 1.01]
            y = 2x
            @test y / 2 == y / 2one(x) == x == y * FD2(0.5)
        end

        # big numbers
        for T in [SFD2, SFD4, FD2, FD4, WFD2, WFD4]
            @test typemin(T) / 2 * 2 == typemin(T)
            @test (typemax(T) / 2 - eps(T)) * 2 == typemax(T) - eps(T)
        end
    end

    @testset "division by 3" begin
        @test FD2(10) / 3 == FD2(3.33)
        @test FD2(20) / 3 == FD2(6.67)
        @test FD2(-1.50) / 3 == FD2(-0.50)
        @test FD2(-20) / 3 == FD2(-6.67)

        # should work with big numbers
        @test typemin(FD2) / 3 < 0
        @test typemax(FD2) / 3 > 0
    end

    @testset "reciprocal $x" for x in filter(!iszero, keyvalues[FD2])
        # convert each keyvalue to rational to take an exact reciprocal
        # and check that it rounds to the in-type reciprocal.
        r = 1/Rational{BigInt}(x)
        @test round(FD2, r) == 1/x
    end

    @testset "divide $x by 0" for x in keyvalues[FD2]
        @test_throws DivideError x/FD2(0)
    end

    @testset "rounding" begin
        # RoundNearest: 1.27 / 2 == 0.635 rounds up to 0.64
        @test FD2(1.27) / FD2(2) == FD2(0.64)
        @test FD2(-1.27) / FD2(2) == FD2(-0.64)
        @test FD2(1.27) / 2 == FD2(0.64)
        @test FD2(-1.27) / 2 == FD2(-0.64)
        @test 127 / FD2(200) == FD2(0.64)
        @test -127 / FD2(200) == FD2(-0.64)

        # RoundNearest: 1.29 / 2 == 0.645 rounds down to 0.64
        @test FD2(1.29) / FD2(2) == FD2(0.64)
        @test FD2(-1.29) / FD2(2) == FD2(-0.64)
        @test FD2(1.29) / 2 == FD2(0.64)
        @test FD2(-1.29) / 2 == FD2(-0.64)
        @test 129 / FD2(200) == FD2(0.64)
        @test -129 / FD2(200) == FD2(-0.64)

        # Use of Float or BigFloat internally should not change the calculated result
        @test round(Int, 109 / 200 * 100) == 55
        @test round(Int, BigInt(109) / 200 * 100) == 54  # Correct

        x = FD2(1.09)
        y = FD2(200)
        for T in [FD2, Int8, Int128, BigInt]
            @test x / T(2) == FD2(0.54)
            @test T(109) / y == FD2(0.54)
        end
    end

    @testset "without promotion" begin
        @test_throws InexactError FD{Int8,1}(20)
        @test Int8(20) / FD{Int8,1}(2) == FD{Int8,1}(10.0)
        @test FD{Int8,1}(2) / Int8(20) == FD{Int8,1}(0.1)
    end

    @testset "limits" begin
        @test_throws InexactError Int8(1) / FD{Int8,2}(0.4)
        @test_throws InexactError FD{Int8,2}(1) / FD{Int8,2}(0.4)
    end

    @testset "limits of $T" for T in CONTAINER_TYPES
        max_exp = FixedPointDecimals.max_exp10(T)
        f = max_exp
        scalar = convert(FD{T,f}, 1 // 10)  # 0.1

        # Should be outside of the bounds of a FD{T,f}
        x = T(10)
        @test_throws InexactError FD{T,f}(x)

        # Since multiply will round the result we'll make sure our value always
        # rounds down.
        max_int = typemax(T) - (typemax(T) % 10)
        min_int = typemin(T) - (typemin(T) % 10)
        max_fd = reinterpret(FD{T,f}, max_int)
        min_fd = reinterpret(FD{T,f}, min_int)

        @test (max_fd * scalar) / scalar == max_fd
        @test (min_fd * scalar) / scalar == min_fd
        @test max_fd / x == reinterpret(FD{T,f}, div(max_int, x))
        @test min_fd / x == reinterpret(FD{T,f}, div(min_int, x))
    end
end

@testset "truncating div" begin
    @testset "div by 1" begin
        @testset for x in keyvalues[FD2]
            @test x ÷ one(x) === trunc(x)

            # signed integers using two's complement have one additional negative value
            if x < 0 && trunc(x) === typemin(x)
                @test_throws InexactError x ÷ -one(x)
            else
                @test x ÷ -one(x) === -trunc(x)
            end
        end
    end
    @testset "div by 2" begin
        @testset for x in keyvalues[FD2]
            @test x ÷ 2one(x) === x ÷ 2 === FD2(x.i ÷ FD2(2).i)
        end
    end
    @testset "return types" begin
        @test div(2one(FD2), 3) isa FD2
        @test one(FD2) ÷ one(FD2) isa FD2

        # Promotion to bigger type
        @test one(FD4) ÷ one(FD2) isa FD4
        @test one(FD2) ÷ one(FD4) isa FD4

        @test one(FD{Int32, 2}) ÷ one(FD{Int64, 6}) isa FD{Int64, 6}
    end

    @testset "div with rounding modes" begin
        if VERSION >= v"1.4.0-"
            @testset for x in keyvalues[FD2]
                # TODO: Test RoundFromZero -- https://github.com/JuliaLang/julia/issues/34519
                for R in (RoundToZero, RoundUp, RoundDown, RoundNearest, RoundNearestTiesAway)
                    @test div(x, 2one(x), R) === div(x, 2, R) === FD2(div(x.i, FD2(2).i, R))
                end
            end
        end
        @testset for x in keyvalues[FD2], f in (fld, cld, fld1, div)
            @test f(x, 2one(x)) === f(x, 2) === FD2(f(x.i, FD2(2).i))
        end
    end
end

@testset "abs, sign" begin
    @testset for T in [FD2, WFD4]
        for x in keyvalues[T]
            @test sign(x)^2 ∈ [0, 1]
            @test abs(abs(x)) == abs(x)
            @test abs(x) * sign(x) == x
            @test abs(x) == abs(-x)
            if x ≠ typemin(x)
                @test abs(x) ≥ 0
                @test sign(x) == -sign(-x)
            end
            @test (abs(x) == 0) === (x == 0)
        end
    end
end

@testset "isinteger" begin
    # Note: Test cannot be used unless we can construct `FD{Int8,6}`
    # @testset "overflow" begin
    #     # Note: After overflow `Int8(10)^6 == 64`
    #     @test !isinteger(reinterpret(FD{Int8,6}, 64))  # 0.000064
    # end

    @testset "limits of $T" for T in CONTAINER_TYPES
        f = FixedPointDecimals.max_exp10(T)

        max_fd = typemax(FD{T,f})
        min_fd = typemin(FD{T,f})

        @test !isinteger(max_fd)
        @test isinteger(trunc(max_fd))

        @test isinteger(min_fd) == (min_fd == zero(min_fd))
        @test isinteger(trunc(min_fd))
    end
end

@testset "round" begin
    @testset "to Int" begin
        @test round(Int, FD2(-0.51)) === -1
        @test round(Int, FD2(-0.50)) === 0
        @test round(Int, FD2(-0.49)) === 0
        @test round(Int, FD2(0.50)) === 0
        @test round(Int, FD2(0.51)) === 1
        @test round(Int, FD2(1.50)) === 2
    end

    @testset "rounding invariant $x" for x in filter(!islarge, keyvalues[FD2])
        @test isinteger(round(x))
        @test x - FD2(1//2) ≤ round(x) ≤ x + FD2(1//2)
        if x - FD2(1//2) == round(x) || x + FD2(1//2) == round(x)
            @test iseven(convert(Int, 100round(x)))
        end
        @testset "to Int" for T in [Int64, Int32]
            @test round(T, x) == round(x)
        end
        # to FD1
        @test x - FD2(1//20) ≤ round(FD1, x) ≤ x + FD2(1//20)
        # to FD2
        @test x == round(FD2, x)
    end

    @testset "limits of $T" for T in CONTAINER_TYPES
        f = FixedPointDecimals.max_exp10(T)
        powt = FixedPointDecimals.coefficient(FD{T,f})

        # Ideally we would just use `typemax(T)` but due to precision issues with
        # floating-point its possible the closest float will exceed `typemax(T)`.
        # Additionally, when the division results in a `BigFloat` we need to first truncate
        # to a `BigInt` before we can truncate the type we want.
        max_int = T(trunc(BigInt, prevfloat(typemax(T) / powt) * powt))
        min_int = T(trunc(BigInt, nextfloat(typemin(T) / powt) * powt))

        @test round(FD{T,f}, max_int / powt) == reinterpret(FD{T,f}, max_int)
        @test round(FD{T,f}, min_int / powt) == reinterpret(FD{T,f}, min_int)

        @test round(FD{T,f}, typemax(T) // powt) == reinterpret(FD{T,f}, typemax(T))
        @test round(FD{T,f}, typemin(T) // powt) == reinterpret(FD{T,f}, typemin(T))

        # Note: rounding away from zero will result in an exception.
        max_int = typemax(T)
        min_int = typemin(T)

        max_dec = max_int / powt
        min_dec = min_int / powt

        if round(T, max_dec) == trunc(T, max_dec)
            @test round(reinterpret(FD{T,f}, max_int)) == FD{T,f}(round(T, max_dec))
        else
            @test_throws InexactError round(reinterpret(FD{T,f}, max_int))
        end

        if round(T, min_dec) == trunc(T, min_dec)
            @test round(reinterpret(FD{T,f}, min_int)) == FD{T,f}(round(T, min_dec))
        else
            @test_throws InexactError round(reinterpret(FD{T,f}, min_int))
        end
    end
end

@testset "trunc" begin
    @test trunc(Int, FD2(0.99)) === 0
    @test trunc(Int, FD2(-0.99)) === 0
    @test trunc(Int, FD2(1)) === 1
    @test trunc(Int, FD2(-1)) === -1
    @test trunc(typemax(FD2)) ≤ typemax(FD2)
    @test trunc(Int, typemax(FD2)) ≤ typemax(FD2)
    @test trunc(typemin(FD2)) ≥ typemin(FD2)
    @test trunc(Int, typemin(FD2)) ≥ typemin(FD2)
    @test trunc(eps(FD2)) == 0
    @test trunc(-eps(FD2)) == 0

    @testset "truncate invariant" for x in keyvalues[FD2]
        @test isinteger(trunc(x))
        if x ≠ typemin(FD2)
            @test abs(x) - 1 < abs(trunc(x)) ≤ abs(x)
        else
            @test abs(trunc(x)) > 0
        end

        # to FD1
        @test isinteger(Base.widemul(10, FD2(trunc(FD1, x))))
        @test abs(FD2(trunc(FD1, x))) ≥ 0
    end

    @testset "truncate precision" begin
        for x in smaller_than_decimal
            @test trunc(FD2, x) ≠ trunc(FD3, x)
            @test trunc(FD2, x) == FD2(x - 0.01)
            @test trunc(FD3, x) == FD3(x - 0.001)

            for f in 0:18
                @test trunc(FD{Int64, f}, x) == parse_int(FD{Int64, f}, INTS[x])
            end
            for f in 0:200
                @test trunc(FD{BigInt, f}, x) == parse_int(FD{BigInt, f}, INTS[x])
            end
        end

        for x in bigger_than_decimal
            exactval = FD3(x)
            for f in 3:14
                @test trunc(FD{Int64, f}, x) == exactval
            end
            for f in 0:18
                @test trunc(FD{Int64, f}, x) == parse_int(FD{Int64, f}, INTS[x])
            end
        end
    end

    @testset "limits of $T" for T in CONTAINER_TYPES
        f = FixedPointDecimals.max_exp10(T)
        powt = FixedPointDecimals.coefficient(FD{T,f})

        # When converting from typemax to a floating-point it is possible that due to
        # precision issues that the closest possible float will exceed the typemax.
        max_float = prevfloat(convert(AbstractFloat, typemax(FD{T,f})))
        min_float = nextfloat(convert(AbstractFloat, typemin(FD{T,f})))

        @test trunc(FD{T,f}, max_float) == trunc_alt(FD{T,f}, max_float)
        @test trunc(FD{T,f}, min_float) == trunc_alt(FD{T,f}, min_float)

        @test trunc(reinterpret(FD{T,f}, typemax(T))) == FD{T,f}(div(typemax(T), powt))
        @test trunc(reinterpret(FD{T,f}, typemin(T))) == FD{T,f}(div(typemin(T), powt))
    end
end

# eps that works for integers too
epsi(::Type{T}) where T <: Integer = one(T)::T
epsi(::Type{T}) where T = eps(T)

@testset "floor, ceil" begin
    @testset for x in filter(!islarge, keyvalues[FD2])
        @test floor(x) ≤ x < floor(x) + 1
        @test ceil(x) - 1 < x ≤ ceil(x)
        @test isinteger(floor(x))
        @test isinteger(ceil(x))

        @testset for T in [Int32, Int64, FD1, FD2, FD4, WFD2, WFD4]
            @test floor(T, x) ≤ x < floor(T, x) + epsi(T)
            @test ceil(T, x) - epsi(T) < x ≤ ceil(T, x)
        end
    end

    @testset "floor, ceil precision" begin
        for x in smaller_than_decimal
            @test floor(FD2, x) != floor(FD3, x)
            @test floor(FD2, x) == FD2(x - 0.01)
            @test floor(FD3, x) == FD3(x - 0.001)

            for f in 0:18
                @test floor(FD{Int64, f}, x) == parse_int(FD{Int64, f}, INTS[x])
            end

            @test ceil(FD3, x) == ceil(FD4, x) == FD4(x)
        end

        for x in bigger_than_decimal
            @test ceil(FD2, x) ≠ ceil(FD3, x)
            @test ceil(FD2, x) == FD2(x + 0.01)
            @test ceil(FD3, x) == FD3(x + 0.001)

            for f in 0:18
                @test ceil(FD{Int64, f}, x) == parse_int(FD{Int64, f}, INTS[x], ceil=true)
            end

            @test floor(FD3, x) == floor(FD4, x) == FD4(x)
        end
    end

    @testset "limits of $T" for T in CONTAINER_TYPES
        f = FixedPointDecimals.max_exp10(T)
        powt = FixedPointDecimals.coefficient(FD{T,f})

        # When converting from typemax to a floating-point it is possible that due to
        # precision issues that the closest possible float will exceed the typemax.
        max_float = prevfloat(convert(AbstractFloat, typemax(FD{T,f})))
        min_float = nextfloat(convert(AbstractFloat, typemin(FD{T,f})))

        @test floor(FD{T,f}, max_float) == floor_alt(FD{T,f}, max_float)
        @test floor(FD{T,f}, min_float) == floor_alt(FD{T,f}, min_float)

        @test ceil(FD{T,f}, max_float) == ceil_alt(FD{T,f}, max_float)
        @test ceil(FD{T,f}, min_float) == ceil_alt(FD{T,f}, min_float)

        # Note: rounding away from zero will result in an exception.
        max_int = typemax(T)
        min_int = typemin(T)

        max_dec = max_int / powt
        min_dec = min_int / powt

        @test floor(reinterpret(FD{T,f}, max_int)) == FD{T,f}(floor(T, max_dec))
        if floor(T, min_dec) == trunc(T, min_dec)
            @test floor(reinterpret(FD{T,f}, min_int)) == FD{T,f}(floor(T, min_dec))
        else
            @test_throws InexactError floor(reinterpret(FD{T,f}, min_int))
        end

        if ceil(T, max_dec) == trunc(T, max_dec)
            @test ceil(reinterpret(FD{T,f}, max_int)) == FD{T,f}(ceil(T, max_dec))
        else
            @test_throws InexactError ceil(reinterpret(FD{T,f}, max_int))
        end
        @test ceil(reinterpret(FD{T,f}, min_int)) == FD{T,f}(ceil(T, min_dec))
    end
end

@testset "type stability" begin
    # Test that basic operations are type stable for all the basic integer types.
    fs = [0, 1, 2, 7, 16, 38]  # To save time, don't test all possible combinations.
    @testset for T in (CONTAINER_TYPES..., BigInt,)
        maxF = FixedPointDecimals.max_exp10(T)
        frange = filter(f->f<=maxF, fs)
        # Unary operations
        @testset for f in frange
            @test @inferred(zero(FD{T,f}(1))) === FD{T,f}(0)
            @test @inferred(one(FD{T,f}(1))) === FD{T,f}(1)
            @test @inferred(ceil(FD{T,f}(1))) === FD{T,f}(1)
            @test @inferred(round(FD{T,f}(1))) === FD{T,f}(1)
            @test @inferred(abs(FD{T,f}(1))) === FD{T,f}(1)
            @test @inferred(FD{T,f}(1)^2) === FD{T,f}(1)
            @test @inferred(typemax(FD{T,f})) isa FD{T,f}
        end
        # Binary operations
        @testset for (f1,f2) in Iterators.product(frange, frange)
            fmax = max(f1,f2)
            @test @inferred(FD{T,f1}(1) + FD{T,f2}(0)) === FD{T,fmax}(1)
            @test @inferred(FD{T,f1}(1) - FD{T,f2}(0)) === FD{T,fmax}(1)
            @test @inferred(FD{T,f1}(1) * FD{T,f2}(1)) === FD{T,fmax}(1)
            @test @inferred(FD{T,f1}(1) / FD{T,f2}(1)) === FD{T,fmax}(1)
            @test @inferred(FD{T,f1}(1) ÷ FD{T,f2}(1)) === FD{T,fmax}(1)
            @test @inferred(max(FD{T,f1}(1), FD{T,f2}(0))) === FD{T,fmax}(1)
            @test @inferred(min(FD{T,f1}(1), FD{T,f2}(0))) === FD{T,fmax}(0)
        end
    end
end

@testset "print" begin
    @test string(FD2(1.00)) == "1.00"
    @test string(FD2(1.23)) == "1.23"
    @test string(FD2(42.40)) == "42.40"
    @test string(FD2(-42.40)) == "-42.40"
    @test string(FD2(-0.01)) == "-0.01"
    @test string(FD2(0)) == "0.00"
    @test string(FixedDecimal{Int,0}(123.4)) == "123"

    # Displaying a decimal could be incorrect when using a decimal place precision which is
    # close to or at the limit for our storage type.
    @testset "limits of $T" for T in CONTAINER_TYPES
        f = FixedPointDecimals.max_exp10(T)

        function fmt(val, f)
            str = string(val)
            neg = ""
            if str[1] == '-'
                neg = "-"
                str = str[2:end]
            end
            return string(neg, str[1], ".", rpad(str[2:end], f, '0'))
        end

        @test string(reinterpret(FD{T,f}, typemax(T))) == fmt(typemax(T), f)
        @test string(reinterpret(FD{T,f}, typemin(T))) == fmt(typemin(T), f)
    end
end

@testset "show" begin
    @testset "compact" begin
        @test sprint(show, FD2(1.00), context=:compact=>true) == "1.0"
        @test sprint(show, FD2(1.23), context=:compact=>true) == "1.23"
        @test sprint(show, FD2(42.40), context=:compact=>true) == "42.4"
        @test sprint(show, FD2(-42.40), context=:compact=>true) == "-42.4"
        @test sprint(show, FD2(-0.01), context=:compact=>true) == "-0.01"
        @test sprint(show, FD2(0), context=:compact=>true) == "0.0"

        @test repr(typemin(FixedDecimal{Int64, 2})) ==
              "FixedDecimal{Int64,2}(-92233720368547758.08)"
        @test repr(typemax(FixedDecimal{Int64, 2})) ==
              "FixedDecimal{Int64,2}(92233720368547758.07)"
        @test repr(typemin(FixedDecimal{Int32, 2})) ==
              "FixedDecimal{Int32,2}(-21474836.48)"
        @test repr(typemax(FixedDecimal{Int32, 2})) ==
              "FixedDecimal{Int32,2}(21474836.47)"
    end
end

@testset "string" begin
    for x in keyvalues[FD2]
        if 0 ≤ abs(x) < 1000
            @test Core.eval(@__MODULE__, Meta.parse(string(x))) == x
        end
    end
end

@testset "parse_round" begin
    @test FixedPointDecimals.parse_round(Int, "44", RoundNearest) == 0
    @test FixedPointDecimals.parse_round(Int, "45", RoundNearest) == 0
    @test FixedPointDecimals.parse_round(Int, "46", RoundNearest) == 1
    @test FixedPointDecimals.parse_round(Int, "54", RoundNearest) == 0
    @test FixedPointDecimals.parse_round(Int, "55", RoundNearest) == 1
    @test FixedPointDecimals.parse_round(Int, "56", RoundNearest) == 1

    # Handle a number of digits that exceeds the storage capacity of Int128
    @test FixedPointDecimals.parse_round(Int8, "9"^40, RoundNearest) == 1
end

@testset "parse" begin
    # Note: the underscore used in the reinterpreted integer is used to indicate the decimal
    # place.
    @testset "decimal position" begin
        @test parse(FD2, "123")   == reinterpret(FD2, 123_00)
        @test parse(FD2, "0.123") == reinterpret(FD2, 0_12)
        @test parse(FD2, ".123")  == reinterpret(FD2, 0_12)
        @test parse(FD2, "1.23")  == reinterpret(FD2, 1_23)
        @test parse(FD2, "12.3")  == reinterpret(FD2, 12_30)
        @test parse(FD2, "123.")  == reinterpret(FD2, 123_00)
        @test parse(FD2, "123.0") == reinterpret(FD2, 123_00)

        @test parse(FD2, "-123")   == reinterpret(FD2, -123_00)
        @test parse(FD2, "-0.123") == reinterpret(FD2, -0_12)
        @test parse(FD2, "-.123")  == reinterpret(FD2, -0_12)
        @test parse(FD2, "-1.23")  == reinterpret(FD2, -1_23)
        @test parse(FD2, "-12.3")  == reinterpret(FD2, -12_30)
        @test parse(FD2, "-123.")  == reinterpret(FD2, -123_00)
        @test parse(FD2, "-123.0") == reinterpret(FD2, -123_00)
    end

    @testset "scientific notation" begin
        @test parse(FD4, "12e0")   == reinterpret(FD4, 00012_0000)
        @test parse(FD4, "12e3")   == reinterpret(FD4, 12000_0000)
        @test parse(FD4, "12e-3")  == reinterpret(FD4, 00000_0120)
        @test parse(FD4, "1.2e0")  == reinterpret(FD4, 00001_2000)
        @test parse(FD4, "1.2e3")  == reinterpret(FD4, 01200_0000)
        @test parse(FD4, "1.2e-3") == reinterpret(FD4, 00000_0012)
        @test parse(FD4, "1.2e-4") == reinterpret(FD4, 00000_0001)

        @test parse(FD4, "-12e0")   == reinterpret(FD4, -00012_0000)
        @test parse(FD4, "-12e3")   == reinterpret(FD4, -12000_0000)
        @test parse(FD4, "-12e-3")  == reinterpret(FD4, -00000_0120)
        @test parse(FD4, "-1.2e0")  == reinterpret(FD4, -00001_2000)
        @test parse(FD4, "-1.2e3")  == reinterpret(FD4, -01200_0000)
        @test parse(FD4, "-1.2e-3") == reinterpret(FD4, -00000_0012)

        @test parse(FD2, "999e-1") == reinterpret(FD2, 99_90)
        @test parse(FD2, "999e-2") == reinterpret(FD2, 09_99)
        @test parse(FD2, "999e-3") == reinterpret(FD2, 01_00)
        @test parse(FD2, "999e-4") == reinterpret(FD2, 00_10)
        @test parse(FD2, "999e-5") == reinterpret(FD2, 00_01)
        @test parse(FD2, "999e-6") == reinterpret(FD2, 00_00)

        @test parse(FD2, "-999e-1") == reinterpret(FD2, -99_90)
        @test parse(FD2, "-999e-2") == reinterpret(FD2, -09_99)
        @test parse(FD2, "-999e-3") == reinterpret(FD2, -01_00)
        @test parse(FD2, "-999e-4") == reinterpret(FD2, -00_10)
        @test parse(FD2, "-999e-5") == reinterpret(FD2, -00_01)
        @test parse(FD2, "-999e-6") == reinterpret(FD2, -00_00)

        @test parse(FD4, "9"^96 * "e-100") == reinterpret(FD4, 0_001)
    end

    @testset "round to nearest" begin
        @test parse(FD2, "0.444") == reinterpret(FD2, 0_44)
        @test parse(FD2, "0.445") == reinterpret(FD2, 0_44)
        @test parse(FD2, "0.446") == reinterpret(FD2, 0_45)
        @test parse(FD2, "0.454") == reinterpret(FD2, 0_45)
        @test parse(FD2, "0.455") == reinterpret(FD2, 0_46)
        @test parse(FD2, "0.456") == reinterpret(FD2, 0_46)

        @test parse(FD2, "-0.444") == reinterpret(FD2, -0_44)
        @test parse(FD2, "-0.445") == reinterpret(FD2, -0_44)
        @test parse(FD2, "-0.446") == reinterpret(FD2, -0_45)
        @test parse(FD2, "-0.454") == reinterpret(FD2, -0_45)
        @test parse(FD2, "-0.455") == reinterpret(FD2, -0_46)
        @test parse(FD2, "-0.456") == reinterpret(FD2, -0_46)

        @test parse(FD2, "0.009")  == reinterpret(FD2,  0_01)
        @test parse(FD2, "-0.009") == reinterpret(FD2, -0_01)

        @test parse(FD4, "1.5e-4") == reinterpret(FD4, 0_0002)
    end

    @testset "round to zero" begin
        @test parse(FD2, "0.444", RoundToZero) == reinterpret(FD2, 0_44)
        @test parse(FD2, "0.445", RoundToZero) == reinterpret(FD2, 0_44)
        @test parse(FD2, "0.446", RoundToZero) == reinterpret(FD2, 0_44)
        @test parse(FD2, "0.454", RoundToZero) == reinterpret(FD2, 0_45)
        @test parse(FD2, "0.455", RoundToZero) == reinterpret(FD2, 0_45)
        @test parse(FD2, "0.456", RoundToZero) == reinterpret(FD2, 0_45)

        @test parse(FD2, "-0.444", RoundToZero) == reinterpret(FD2, -0_44)
        @test parse(FD2, "-0.445", RoundToZero) == reinterpret(FD2, -0_44)
        @test parse(FD2, "-0.446", RoundToZero) == reinterpret(FD2, -0_44)
        @test parse(FD2, "-0.454", RoundToZero) == reinterpret(FD2, -0_45)
        @test parse(FD2, "-0.455", RoundToZero) == reinterpret(FD2, -0_45)
        @test parse(FD2, "-0.456", RoundToZero) == reinterpret(FD2, -0_45)

        @test parse(FD2, "0.009", RoundToZero)  == reinterpret(FD2, 0_00)
        @test parse(FD2, "-0.009", RoundToZero) == reinterpret(FD2, 0_00)

        @test parse(FD4, "1.5e-4", RoundToZero) == reinterpret(FD4, 0_0001)
    end

    @testset "round throws" begin
        @test parse(FD2, "0.44", RoundThrows)  == reinterpret(FD2, 0_44)
        @test parse(FD2, "0.440", RoundThrows) == reinterpret(FD2, 0_44)

        @test_throws InexactError parse(FD2, "0.444", RoundThrows)
        @test_throws InexactError parse(FD2, "0.445", RoundThrows)
        @test_throws InexactError parse(FD2, "0.446", RoundThrows)
        @test_throws InexactError parse(FD2, "0.454", RoundThrows)
        @test_throws InexactError parse(FD2, "0.455", RoundThrows)
        @test_throws InexactError parse(FD2, "0.456", RoundThrows)

        @test_throws InexactError parse(FD2, "-0.444", RoundThrows)
        @test_throws InexactError parse(FD2, "-0.445", RoundThrows)
        @test_throws InexactError parse(FD2, "-0.446", RoundThrows)
        @test_throws InexactError parse(FD2, "-0.454", RoundThrows)
        @test_throws InexactError parse(FD2, "-0.455", RoundThrows)
        @test_throws InexactError parse(FD2, "-0.456", RoundThrows)

        @test_throws InexactError parse(FD2, "0.009", RoundThrows)
        @test_throws InexactError parse(FD2, "-0.009", RoundThrows)

        @test_throws InexactError parse(FD4, "1.5e-4", RoundThrows)
    end

    @testset "invalid" begin
        @test_throws OverflowError parse(FD4, "1.2e100")
        @test_throws ArgumentError parse(FD4, "foo")
        @test_throws ArgumentError parse(FD4, "1.2.3")
        @test_throws ArgumentError parse(FD4, "1.2", RoundUp)
    end

    @testset "hashing" begin
        fd1 = FixedDecimal{Int, 4}(2.5)
        fd2 = FixedDecimal{Int, 5}(2.5)
        fd3 = FixedDecimal{Int, 4}(3.5)

        @test hash(fd1) == hash(fd2)
        @test hash(fd1) != hash(fd3)
        @test hash(fd1) != hash(fd1.i)
        @test hash(FD2(1//10)) == hash(1//10)
        @test hash(FD2(1//10)) ≠ hash(0.1)
    end
end

end  # global testset


text:
"""
    KiteConnect

Julia module to interface with Zerodha's KiteConnect API
"""
module KiteConnect

export Bar

using HTTP
using JSON3
using SHA

include("quote.jl")
include("historical_data.jl")
include("charges.jl")

const API_ENDPOINT = "https://api.kite.trade"
ACCESS_TOKEN = ""

function get_http_headers(access_token::String)
    return [
        "X-Kite-Version" => "3",
        "Authorization" => "token " * ENV["KITE_API_KEY"] * ":$access_token",
    ]
end

function http_get(url_fragment::String, access_token::String)
    url = "$API_ENDPOINT/$url_fragment"
    r = HTTP.request("GET", url, get_http_headers(access_token))
    return r.body |> String |> JSON3.read
end

"""
  `gen_access_token(request_token::String)`

Generate the access token by passing in your request token
"""
function gen_access_token(request_token::String)
    checksum = (ENV["KITE_API_KEY"] * request_token * ENV["KITE_API_SECRET"]) |> sha256 |> bytes2hex
    url = "$API_ENDPOINT/session/token"

    headers = [
        "X-Kite-Version" => "3",
        "Content-Type"   => "application/x-www-form-urlencoded",
    ]
    body = "api_key=" * ENV["KITE_API_KEY"] * "&request_token=$request_token&checksum=$checksum"

    res = HTTP.request("POST", url, headers, body)
    r = res.body |> String |> JSON3.read
    return r["data"]["access_token"]
end

end # module

ext: jl


text:
using JLD2
using FileIO
import Reproduce

function save_setup(parsed;
    save_dir_key="save_dir",
    def_save_file="results.jld2",
    filter_keys=["verbose",
                 "working",
                 "exp_loc",
                 "visualize",
                 "progress",
                 "synopsis"])

    savefile = def_save_file
    Reproduce.create_info!(
        parsed,
        parsed[save_dir_key];
        filter_keys=filter_keys)
    savepath = Reproduce.get_save_dir(parsed)
    joinpath(savepath, def_save_file)
end

function save_results(savefile, results)
    JLD2.@save savefile results
end

function check_save_file_loadable(savefile)
    try
        JLD2.@load savefile results
    catch
        return false
    end
    return true
end

function experiment_wrapper(exp_func::Function, parsed, logger_init_info, working)
    save_file = save_setup(parsed)
    if isfile(save_file) && Curiosity.check_save_file_loadable(save_file) && !working
        println("Experiment Already Done. Skipping.")
        return
    end
    logger = Logger(parsed, logger_init_info, save_file)

    ret = exp_func(parsed, logger)

    if working
        ret, logger
    else
        logger_save(logger)
    end
end

max_stars_repo_name: MatthewMcLeod/curiosity

max_issues_repo_name: MatthewMcLeod/curiosity

max_forks_repo_name: MatthewMcLeod/curiosity


text:
module Test_iterate_jkp

include("../src/iterate_jkp.jl")

using Test
using Random
using SparseArrays

function test(F,top_k)
    n = size(F,1)
    sF = sparse(F)
    Fp,Fi = sF.colptr,sF.rowval
    F = Matrix(F)

    @testset "js" begin
        @test [j for (j,kps) in iterate_jkp(Fp,Fi,Val(top_k))] == (1:n)
    end

    @testset "ks" begin
        for (j,kps) in iterate_jkp(Fp,Fi,Val(top_k))
            ks = sort!([k for (k,ps) in kps])
            ref = findall(!iszero, F[j,1:j-1])
            top_k && filter!(k->all(F[k+1:j-1,k] .== 0), ref)
            @test ks == ref
        end
    end

    @testset "ps" begin
        for (j,kps) in iterate_jkp(Fp,Fi,Val(top_k))
            for (k,ps) in kps
                @test Fi[ps] == (j-1) .+ findall(!iszero, F[j:n,k])
            end
        end
    end
end

function bit_vector(nbits,i)
    b = Vector{Bool}(undef,nbits)
    for k = 1:nbits
        b[k] = i%2
        i = i >> 1
    end
    return b
end

@testset "iterate_jkp" begin
    @testset for top_k = (false,true)
        # Test all matrices with n <= 3
        @testset for n = 1:3
            @testset for i = 1:2^(n^2)
                F = reshape(bit_vector(n^2,i-1),(n,n))
                test(F,top_k)
            end
        end

        # Test a random sample of matrices with n > 3
        @testset "n in 4:100" begin
            @testset for seed = 1:100
                Random.seed!(seed)
                n = rand(4:100)
                F = sprand(n,n,rand())
                test(F,top_k)
            end
        end
    end
end

end # module

max_stars_repo_name: ettersi/SparseFactorizations

max_issues_repo_name: ettersi/SparseFactorizations

max_forks_repo_name: ettersi/SparseFactorizations


text:
using DifferentialEquations, DiffEqFlux, DiffEqSensitivity, Flux
using Optim
using StatsBase
using Distributions
using BSON
using Random

Random.seed!(1)

include("meshgrid.jl") # meshgrid functions
include("example_dynamics.jl") # definitions of differential equations used for synthetic datasets

# ********************************************************* Model ****************************************************************** #

# Hyperparameters
T = 1001        # time stamps rounded to the nearest 1ms
L = 3           # number of latents assumed by the model
L_true = 3      # true number of latents
D = 150         # number of observations
N = 343         # number of trials
dt = 0.001      # 1-ms bin
κ = 1.          # scale parameter for prior initial value distribution covariance
tspan = (0.0f0,40.0f0) # 40 seconds long
trange = range(tspan[1],tspan[2],length=T)
u0 = Float32[1.; 1.; 1.]

f = FastChain(FastDense(L, 17, swish), FastDense(17,23, swish), FastDense(23,17, swish), FastDense(17,L)) # Feedforward Neural Network
_nn_ode = NeuralODE(f, tspan, AutoTsit5(TRBDF2(autodiff=false)), saveat = trange, reltol=1e-6, abstol=1e-6) # Neural ODE
_logλ_true = FastDense(L_true, D) # mapping from latent space (true) to observations
_logλ = FastDense(L, D) # mapping from latent space (model) to observations
p2 = initial_params(_logλ) # initialize parameters for the mapping

nn_ode = (u, p) -> _nn_ode(u, p[1:length(_nn_ode.p)])
logλ = (u, p) -> _logλ(u, p[length(_nn_ode.p)+1:length(_nn_ode.p)+length(p2)])

# ***************************************************** Generate data ************************************************************** #

# Generate loading matrix and bias for this dataset
C = (rand(D, L_true) .+ 2) .* sign.(randn(D, L_true)) # loading matrix -- low mean firing rates
#C = (rand(D, L_true) .+ 8) .* sign.(randn(D, L_true)) # loading matrix -- high mean firing rates
params_λ_true = [collect(Iterators.flatten(C)); zeros(D)]

# Generate initial values for training dataset
x,y,z = ([range(-0.5,0.5,length=7);],[range(-0.5,0.5,length=7);],[range(-0.5,0.5,length=7);])
X,Y,Z = meshgrid(x,y,z)
u0s = [X[:] Y[:] Z[:]]'

# Generate latent trajectories for training dataset
prob_true = ODEProblem(spiral,u0,tspan)
prob_func(prob,i,repeat) = remake(prob, u0=u0s[:,i])
ens_prob = EnsembleProblem(prob_true, prob_func=prob_func)
z_true_mat = Array(solve(ens_prob, Tsit5(), saveat=trange, trajectories=N)) # latent trajectories [latents x timebins x trials]

# Generate training dataset spikes from discretized intensity
spike_times_disc = zeros(D, T, N)
for i = 1:N; spike_times_disc[:,:,i] = (rand.(Poisson.(dt*exp.(_logλ_true(z_true_mat[:,:,i], params_λ_true))))) .> 0; end
spike_times_disc = spike_times_disc .> 0
spikes = permutedims(spike_times_disc, (1, 3, 2)) # spike trains [neurons x trials x timebins]

# Generate initial values for test dataset
u0s_test = rand(L_true, N) .- 0.5

# Generate latent trajectories for test dataset
prob_func_test(prob,i,repeat) = remake(prob, u0=u0s_test[:,i])
ens_prob_test = EnsembleProblem(prob_true, prob_func=prob_func_test)
z_true_test = Array(solve(ens_prob_test, Tsit5(), saveat=trange, trajectories=N)) # latent trajectories [latents x timebins x trials]

# Generate test dataset spikes from discretized intensity
spike_times_test = zeros(D, T, N)
for i = 1:N; spike_times_test[:,:,i] = (rand.(Poisson.(dt*exp.(_logλ_true(z_true_test[:,:,i], params_λ_true))))) .> 0; end
spike_times_test = spike_times_test .> 0
spikes_test = permutedims(spike_times_test, (1, 3, 2)) # spike trains [neurons x trials x timebins]

# ********************************************************** Loss ****************************************************************** #

# Initialize model parameters before training
θ = Float32[Flux.glorot_uniform(length(_nn_ode.p)); randn(D*(L+1)); randn(L*N); -10 .* ones(L*N)]

function loss_nn_ode(p) # negative of the Evidence Lower BOund (ELBO) in Eq. 13
    u0_m = reshape(p[end-L*N-L*N+1:end-L*N], L, :)
    u0_s = clamp.(reshape(p[end-L*N+1:end], L, :), -1e8, 0)
    u0s = u0_m .+ exp.(u0_s) .* randn(size(u0_s))
    kld = 0.5 .* (N .* L .* log(κ) .- sum(2 .* u0_s) - N .* L .+ sum(exp.(2 .* u0_s)) ./ κ .+ sum(abs2, u0_m) ./ κ)
    z_hat = Array(nn_ode(u0s, p))
    λ_hat = exp.(logλ(reshape(z_hat, L, :), p))
    λ_hat = reshape(λ_hat, D, size(z_hat, 2), size(z_hat, 3))
    Nlogλ = spikes[1:size(λ_hat, 1), 1:size(λ_hat, 2), 1:size(λ_hat, 3)] .* log.(dt.*λ_hat .+ sqrt(eps(Float32)))
    loss = ( sum(dt.*λ_hat .- Nlogλ) .+ kld ) ./ N
    return loss
end

cb = function (p, l)
    println(l)
    return false
end

# Burn-in loss
cb(θ,loss_nn_ode(θ))

# ******************************************************** Training **************************************************************** #

"""
Training PLNDE by "iteratively growing fit". Optimizing both the generative and variational parameters.

References:
[1] Rackauckas et al. (2020). Universal differential equations for scientific machine learning. arXiv.
"""

_nn_ode = NeuralODE(f, (0.0, 4.0), AutoTsit5(TRBDF2(autodiff=false)), saveat = trange[trange .< 4.], reltol=1e-6, abstol=1e-6)
res1 = DiffEqFlux.sciml_train(loss_nn_ode, θ, ADAM(0.01), cb = cb, maxiters = 300)

_nn_ode = NeuralODE(f, (0.0, 8.0), AutoTsit5(TRBDF2(autodiff=false)), saveat = trange[trange .< 8.], reltol=1e-6, abstol=1e-6)
res2 = DiffEqFlux.sciml_train(loss_nn_ode, res1.minimizer, ADAM(0.005), cb = cb, maxiters = 400)

_nn_ode = NeuralODE(f, (0.0, 12.0), AutoTsit5(TRBDF2(autodiff=false)), saveat = trange[trange .< 12.], reltol=1e-6, abstol=1e-6)
res3 = DiffEqFlux.sciml_train(loss_nn_ode, res2.minimizer, ADAM(0.005), cb = cb, maxiters = 400)

_nn_ode = NeuralODE(f, (0.0, 16.0), AutoTsit5(TRBDF2(autodiff=false)), saveat = trange[trange .< 16.], reltol=1e-6, abstol=1e-6)
res4 = DiffEqFlux.sciml_train(loss_nn_ode, res3.minimizer, ADAM(0.005), cb = cb, maxiters = 400)

_nn_ode = NeuralODE(f, (0.0, 20.0), AutoTsit5(TRBDF2(autodiff=false)), saveat = trange[trange .< 20.], reltol=1e-6, abstol=1e-6)
res5 = DiffEqFlux.sciml_train(loss_nn_ode, res4.minimizer, ADAM(0.001), cb = cb, maxiters = 400)

_nn_ode = NeuralODE(f, (0.0, 24.0), AutoTsit5(TRBDF2(autodiff=false)), saveat = trange[trange .< 24.], reltol=1e-6, abstol=1e-6)
res6 = DiffEqFlux.sciml_train(loss_nn_ode, res5.minimizer, ADAM(0.001), cb = cb, maxiters = 400)

_nn_ode = NeuralODE(f, (0.0, 28.0), AutoTsit5(TRBDF2(autodiff=false)), saveat = trange[trange .< 28.], reltol=1e-6, abstol=1e-6)
res7 = DiffEqFlux.sciml_train(loss_nn_ode, res6.minimizer, ADAM(0.001), cb = cb, maxiters = 400)

_nn_ode = NeuralODE(f, (0.0, 32.0), AutoTsit5(TRBDF2(autodiff=false)), saveat = trange[trange .< 32.], reltol=1e-6, abstol=1e-6)
res8 = DiffEqFlux.sciml_train(loss_nn_ode, res7.minimizer, ADAM(0.001), cb = cb, maxiters = 600)

_nn_ode = NeuralODE(f, (0.0, 36.0), AutoTsit5(TRBDF2(autodiff=false)), saveat = trange[trange .< 36.], reltol=1e-6, abstol=1e-6)
res9 = DiffEqFlux.sciml_train(loss_nn_ode, res8.minimizer, ADAM(0.001), cb = cb, maxiters = 800)

_nn_ode = NeuralODE(f, tspan, AutoTsit5(TRBDF2(autodiff=false)), saveat = trange, reltol=1e-6, abstol=1e-6)
res10 = DiffEqFlux.sciml_train(loss_nn_ode, res9.minimizer, ADAM(0.001), cb = cb, maxiters = 100)

θ_opt = res10.minimizer

BSON.@save "spiral_fit_train_PLNDE.bson" θ_opt spikes z_true_mat params_λ_true

# ******************************************************** Testing ***************************************************************** #

θ_iv = θ_opt[end-L*N-L*N+1:end]

# Redefine the loss such that the generative parameters are frozen and only the variational parameters are optimized
function loss_nn_ode(p)
    u0_m = reshape(p[end-L*N-L*N+1:end-L*N], L, :)
    u0_s = clamp.(reshape(p[end-L*N+1:end], L, :), -1e8, 0)
    u0s1 = u0_m .+ exp.(u0_s) .* randn(size(u0_s))
    kld = 0.5 .* (N .* L .* log(κ) .- sum(2 .* u0_s) - N .* L .+ sum(exp.(2 .* u0_s)) ./ κ .+ sum(abs2, u0_m) ./ κ)
    z_hat = Array(nn_ode(u0s1, θ_opt))
    λ_hat = exp.(logλ(reshape(z_hat, L, :), θ_opt))
    λ_hat = reshape(λ_hat, D, size(z_hat, 2), size(z_hat, 3))
    Nlogλ = spikes_test[1:size(λ_hat, 1), 1:size(λ_hat, 2), 1:size(λ_hat, 3)] .* log.(dt.*λ_hat .+ sqrt(eps(Float32)))
    loss = ( sum(dt.*λ_hat .- Nlogλ) .+ kld ) ./ N
    return loss
end

# Display the ODE with the initial parameter values - burn in loss
cb(θ_iv,loss_nn_ode(θ_iv))

res11 = DiffEqFlux.sciml_train(loss_nn_ode, θ_iv, ADAM(0.01), cb = cb, maxiters = 100)
res12 = DiffEqFlux.sciml_train(loss_nn_ode, res11.minimizer, ADAM(0.001), cb = cb, maxiters = 450)

θ_opt_test = θ_opt
θ_opt_test[end-L*N-L*N+1:end] = res12.minimizer

BSON.@save "spiral_fit_test_PLNDE.bson" θ_opt_test spikes_test z_true_test params_λ_true

max_stars_repo_name: timkimd/plnde

max_issues_repo_name: timkimd/plnde

max_forks_repo_name: timkimd/plnde


text:
using DrWatson
@quickactivate :ExtendedConv

using PyPlot
using Seaborn
using Random
using Flux
using JOLI
using Flux, Flux.Data.MNIST

set_style("whitegrid")
rc("font", family="serif", size=10)
sfmt=matplotlib.ticker.ScalarFormatter(useMathText=true)
sfmt.set_powerlimits((0, 0))
plt.ioff()

sim_name = "schematics"
save_path = plotsdir(sim_name)

nx = 8
ny = 8
nc = 1
batchsize = 1

# Conv Op
w = collect(range(-0.9999f0, 1f0, length=9))
W = reshape(w, 3, 3, 1, 1)
bias = zeros(Float32, 1)

G = Conv(W, bias, identity; stride=1, pad=1)
Gadj = ConvTranspose(W, bias, identity; stride=1, pad=1)

N = nx * ny * nc * batchsize
K = joLinearFunctionFwd_T(
    N, N,
    x -> vec(G(reshape(x, nx, ny, nc, batchsize))),
    y -> vec(Gadj(reshape(y, nx, ny, nc, batchsize))),
    Float32,Float32, name="Conv"
)

# Form explicit convolution matrix
K=K*jo_eye(Float32,nx*ny)

K_vis = deepcopy(K)
K_vis[K_vis.==0f0] .= NaN

fig, ax = subplots(1, 1, figsize=(4, 4))
im = ax.imshow(
    W[end:-1:1, end:-1:1, 1, 1], cmap="tab20",
    vmin=-1, vmax=1)
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
wsave(joinpath(save_path, "3x3-kernel.png"), fig)
close(fig)

# Toeplitz matrix
fig, ax = subplots(1, 1, figsize=(4, 4))
ax.imshow(
    K_vis, cmap="tab20",
    vmin=-1, vmax=1)
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
wsave(joinpath(save_path, "toeplitz.png"), fig)
close(fig)

# Matricize multiple images
Random.seed!(13)
x1 = 1f0*ones(Float32, nx, ny) + 7f-1*randn(Float32, nx, ny)
x2 = 2f0*ones(Float32, nx, ny) + 7f-1*randn(Float32, nx, ny)
x3 = 3f0*ones(Float32, nx, ny) + 7f-1*randn(Float32, nx, ny)

X = hcat(vec(x1), vec(x2), vec(x3))

fig, ax = subplots(1, 1, figsize=(4, 4))
im = ax.imshow(
    x1, cmap="Set3",
    vmin=0, vmax=4)
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
wsave(joinpath(save_path, "image-1.png"), fig)
close(fig)

fig, ax = subplots(1, 1, figsize=(4, 4))
im = ax.imshow(
    x2, cmap="Set3",
    vmin=0, vmax=4)
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
wsave(joinpath(save_path, "image-2.png"), fig)
close(fig)

fig, ax = subplots(1, 1, figsize=(4, 4))
im = ax.imshow(
    x3, cmap="Set3",
    vmin=0, vmax=4)
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
wsave(joinpath(save_path, "image-3.png"), fig)
close(fig)

fig, ax = subplots(1, 1, figsize=(4, 4))
im = ax.imshow(
    X, cmap="Set3", aspect=1,
    vmin=0, vmax=4)
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
wsave(joinpath(save_path, "X.png"), fig)
close(fig)

# MNIST images
train_imgs = MNIST.images()
X = Array{Float32}(undef, size(train_imgs[1])..., 1, 3)
Y = Array{Float32}(undef, size(train_imgs[1])..., 1, 3)

x1 = Float32.(train_imgs[1])
x2 = Float32.(train_imgs[2])
x3 = Float32.(train_imgs[3])

X = hcat(vec(x1), vec(x2), vec(x3))

fig, ax = subplots(1, 1, figsize=(4, 4))
im = ax.imshow(
    x1, interpolation="nearest", cmap="gray",
    vmin=0, vmax=1)
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
wsave(joinpath(save_path, "mnist-1.png"), fig)
close(fig)

fig, ax = subplots(1, 1, figsize=(4, 4))
im = ax.imshow(
    x2, interpolation="nearest", cmap="gray",
    vmin=0, vmax=1)
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
wsave(joinpath(save_path, "mnist-2.png"), fig)
close(fig)

fig, ax = subplots(1, 1, figsize=(4, 4))
im = ax.imshow(
    x3, interpolation="nearest", cmap="gray",
    vmin=0, vmax=1)
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
wsave(joinpath(save_path, "mnist-3.png"), fig)
close(fig)

fig, ax = subplots(1, 1, figsize=(4, 4))
im = ax.imshow(
    X, aspect=.1, interpolation="nearest", cmap="gray",
    vmin=0, vmax=1)
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
wsave(joinpath(save_path, "X-mnist.png"), fig)
close(fig)

# MNIST Conv
nx, ny = size(x1)
nc = 1
batchsize = 1

# Conv Op
Random.seed!(1)
w = randn(Float32, 3, 3)/sqrt(9f0)
W = reshape(w, 3, 3, 1, 1)
bias = zeros(Float32, 1)

G = Conv(W, bias, identity; stride=1, pad=1)
Gadj = ConvTranspose(W, bias, identity; stride=1, pad=1)

N = nx * ny * nc * batchsize
K = joLinearFunctionFwd_T(
    N, N,
    x -> vec(G(reshape(x, nx, ny, nc, batchsize))),
    y -> vec(Gadj(reshape(y, nx, ny, nc, batchsize))),
    Float32,Float32, name="Conv"
)

# Form explicit convolution matrix
K=K*jo_eye(Float32,nx*ny)
Y = K*X

fig, ax = subplots(1, 1, figsize=(4, 4))
im = ax.imshow(
    reshape(Y[:, 1], nx, ny), interpolation="nearest", cmap="gray",
    vmin=0, vmax=.2)
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
wsave(joinpath(save_path, "y-1.png"), fig)
close(fig)

fig, ax = subplots(1, 1, figsize=(4, 4))
im = ax.imshow(
    reshape(Y[:, 2], nx, ny), interpolation="nearest", cmap="gray",
    vmin=0, vmax=.2)
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
wsave(joinpath(save_path, "y-2.png"), fig)
close(fig)

fig, ax = subplots(1, 1, figsize=(4, 4))
im = ax.imshow(
    reshape(Y[:, 3], nx, ny), interpolation="nearest", cmap="gray",
    vmin=0, vmax=.2)
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
wsave(joinpath(save_path, "y-3.png"), fig)
close(fig)

fig, ax = subplots(1, 1, figsize=(4, 4))
im = ax.imshow(
    Y, aspect=.1, interpolation="nearest", cmap="gray",
    vmin=0, vmax=.2)
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
wsave(joinpath(save_path, "Y.png"), fig)
close(fig)


text:
using Oceananigans.AbstractOperations: UnaryOperation, Derivative, BinaryOperation, MultiaryOperation
using Oceananigans.AbstractOperations: KernelFunctionOperation
using Oceananigans.Operators: ℑxyᶜᶠᵃ, ℑxyᶠᶜᵃ
using Oceananigans.Fields: PressureField, compute_at!
using Oceananigans.BuoyancyModels: BuoyancyField

function compute_derivative(model, ∂)
    T, S = model.tracers
    S.data.parent .= π
    @compute ∂S = ComputedField(∂(S))
    result = Array(interior(∂S))
    return all(result .≈ zero(eltype(model.grid)))
end

function compute_unary(unary, model)
    set!(model; S=π)
    T, S = model.tracers
    @compute uS = ComputedField(unary(S), data=model.pressures.pHY′.data)
    result = Array(interior(uS))
    return all(result .≈ unary(eltype(model.grid)(π)))
end

function compute_plus(model)
    set!(model; S=π, T=42)
    T, S = model.tracers
    @compute ST = ComputedField(S + T, data=model.pressures.pHY′.data)
    result = Array(interior(ST))
    return all(result .≈ eltype(model.grid)(π+42))
end

function compute_many_plus(model)
    set!(model; u=2, S=π, T=42)
    T, S = model.tracers
    u, v, w = model.velocities
    @compute uTS = ComputedField(@at((Center, Center, Center), u + T + S))
    result = Array(interior(uTS))
    return all(result .≈ eltype(model.grid)(2+π+42))
end

function compute_minus(model)
    set!(model; S=π, T=42)
    T, S = model.tracers
    @compute ST = ComputedField(S - T, data=model.pressures.pHY′.data)
    result = Array(interior(ST))
    return all(result .≈ eltype(model.grid)(π-42))
end

function compute_times(model)
    set!(model; S=π, T=42)
    T, S = model.tracers
    @compute ST = ComputedField(S * T, data=model.pressures.pHY′.data)
    result = Array(interior(ST))
    return all(result .≈ eltype(model.grid)(π*42))
end

function compute_kinetic_energy(model)
    u, v, w = model.velocities
    set!(u, 1)
    set!(v, 2)
    set!(w, 3)

    kinetic_energy_operation = @at (Center, Center, Center) (u^2 + v^2 + w^2) / 2
    @compute kinetic_energy = ComputedField(kinetic_energy_operation, data=model.pressures.pHY′.data)
    result = Array(interior(kinetic_energy))

    return all(result .≈ 7)
end

function horizontal_average_of_plus(model)
    Ny, Nz = model.grid.Ny, model.grid.Nz

    S₀(x, y, z) = sin(π * z)
    T₀(x, y, z) = 42 * z
    set!(model; S=S₀, T=T₀)
    T, S = model.tracers

    @compute ST = AveragedField(S + T, dims=(1, 2))

    zC = znodes(Center, model.grid)
    correct_profile = @. sin(π * zC) + 42 * zC

    result = Array(interior(ST))[:]

    return all(result .≈ correct_profile)
end

function zonal_average_of_plus(model)
    Ny, Nz = model.grid.Ny, model.grid.Nz

    S₀(x, y, z) = sin(π*z) * sin(π*y)
    T₀(x, y, z) = 42*z + y^2
    set!(model; S=S₀, T=T₀)
    T, S = model.tracers

    @compute ST = AveragedField(S + T, dims=1)

    yC = ynodes(Center, model.grid, reshape=true)
    zC = znodes(Center, model.grid, reshape=true)
    correct_slice = @. sin(π * zC) * sin(π * yC) + 42*zC + yC^2

    result = Array(interior(ST))

    return all(result[1, :, :] .≈ correct_slice[1, :, :])
end

function volume_average_of_times(model)
    Ny, Nz = model.grid.Ny, model.grid.Nz

    S₀(x, y, z) = 1 + sin(2π*x)
    T₀(x, y, z) = y
    set!(model; S=S₀, T=T₀)
    T, S = model.tracers

    @compute ST = AveragedField(S * T, dims=(1, 2, 3))

    result = Array(interior(ST))

    return all(result[1, 1, 1] .≈ 0.5)
end

function horizontal_average_of_minus(model)
    Ny, Nz = model.grid.Ny, model.grid.Nz

    S₀(x, y, z) = sin(π*z)
    T₀(x, y, z) = 42*z
    set!(model; S=S₀, T=T₀)
    T, S = model.tracers

    @compute ST = AveragedField(S - T, dims=(1, 2))

    zC = znodes(Center, model.grid)
    correct_profile = @. sin(π * zC) - 42 * zC

    result = Array(interior(ST))

    return all(result[1, 1, 1:Nz] .≈ correct_profile)
end

function horizontal_average_of_times(model)
    Ny, Nz = model.grid.Ny, model.grid.Nz

    S₀(x, y, z) = sin(π*z)
    T₀(x, y, z) = 42*z
    set!(model; S=S₀, T=T₀)
    T, S = model.tracers

    @compute ST = AveragedField(S * T, dims=(1, 2))

    zC = znodes(Center, model.grid)
    correct_profile = @. sin(π * zC) * 42 * zC

    result = Array(interior(ST))

    return all(result[1, 1, 1:Nz] .≈ correct_profile)
end

function multiplication_and_derivative_ccf(model)
    Ny, Nz = model.grid.Ny, model.grid.Nz

    w₀(x, y, z) = sin(π * z)
    T₀(x, y, z) = 42 * z
    set!(model; enforce_incompressibility=false, w=w₀, T=T₀)

    w = model.velocities.w
    T = model.tracers.T

    @compute wT = AveragedField(w * ∂z(T), dims=(1, 2))

    zF = znodes(Face, model.grid)
    correct_profile = @. 42 * sin(π * zF)

    result = Array(interior(wT))

    # Omit both halos and boundary points
    return all(result[1, 1, 2:Nz] .≈ correct_profile[2:Nz])
end

const C = Center
const F = Face

function multiplication_and_derivative_ccc(model)
    Ny, Nz = model.grid.Ny, model.grid.Nz

    w₀(x, y, z) = sin(π*z)
    T₀(x, y, z) = 42*z
    set!(model; enforce_incompressibility=false, w=w₀, T=T₀)

    w = model.velocities.w
    T = model.tracers.T

    wT_ccc = @at (C, C, C) w * ∂z(T)
    @compute wT_ccc_avg = AveragedField(wT_ccc, dims=(1, 2))

    zF = znodes(Face, model.grid)
    sinusoid = sin.(π * zF)
    interped_sin = [(sinusoid[k] + sinusoid[k+1]) / 2 for k in 1:model.grid.Nz]
    correct_profile = interped_sin .* 42

    result = Array(interior(wT_ccc_avg))

    # Omit boundary-adjacent points from comparison
    return all(result[1, 1, 2:Nz-1] .≈ correct_profile[2:Nz-1])
end

function computation_including_boundaries(arch)
    topo = (Periodic, Bounded, Bounded)
    grid = RegularRectilinearGrid(topology=topo, size=(13, 17, 19), extent=(1, 1, 1))
    model = NonhydrostaticModel(architecture=arch, grid=grid)

    u, v, w = model.velocities
    @. u.data = 1 + rand()
    @. v.data = 2 + rand()
    @. w.data = 3 + rand()

    op = @at (Center, Face, Face) u * v * w
    @compute uvw = ComputedField(op)

    return all(interior(uvw) .!= 0)
end

function operations_with_computed_field(model)
    u, v, w = model.velocities
    uv = ComputedField(u * v)
    @compute uvw = ComputedField(uv * w)
    return true
end

function operations_with_averaged_field(model)
    u, v, w = model.velocities
    UV = AveragedField(u * v, dims=(1, 2))
    wUV = ComputedField(w * UV)
    compute!(wUV)
    return true
end

function pressure_field(model)
    p = PressureField(model)
    u, v, w = model.velocities
    @compute up = ComputedField(u * p)
    return true
end

function computations_with_buoyancy_field(arch, buoyancy)
    grid = RegularRectilinearGrid(size=(1, 1, 1), extent=(1, 1, 1))
    tracers = buoyancy isa BuoyancyTracer ? :b : (:T, :S)
    model = NonhydrostaticModel(architecture=arch, grid=grid,
                                tracers=tracers, buoyancy=buoyancy)

    b = BuoyancyField(model)
    u, v, w = model.velocities

    compute!(b)

    ub = ComputedField(b * u)
    vb = ComputedField(b * v)
    wb = ComputedField(b * w)

    compute!(ub)
    compute!(vb)
    compute!(wb)

    return true # test that it doesn't error
end

function computations_with_averaged_fields(model)
    u, v, w, T, S = fields(model)

    set!(model, enforce_incompressibility = false, u = (x, y, z) -> z, v = 2, w = 3)

    # Two ways to compute turbulent kinetic energy
    U = AveragedField(u, dims=(1, 2))
    V = AveragedField(v, dims=(1, 2))

    tke_op = @at (Center, Center, Center) ((u - U)^2  + (v - V)^2 + w^2) / 2
    tke = ComputedField(tke_op)
    compute!(tke)

    return all(interior(tke)[2:3, 2:3, 2:3] .== 9/2)
end

function computations_with_averaged_field_derivative(model)

    set!(model, enforce_incompressibility = false, u = (x, y, z) -> z, v = 2, w = 3)

    u, v, w, T, S = fields(model)

    # Two ways to compute turbulent kinetic energy
    U = AveragedField(u, dims=(1, 2))
    V = AveragedField(v, dims=(1, 2))

    # This tests a vertical derivative of an AveragedField
    shear_production_op = @at (Center, Center, Center) u * w * ∂z(U)
    shear = ComputedField(shear_production_op)
    compute!(shear)

    set!(model, T = (x, y, z) -> 3 * z)

    return all(interior(shear)[2:3, 2:3, 2:3] .== interior(T)[2:3, 2:3, 2:3])
end

function computations_with_computed_fields(model)
    u, v, w, T, S = fields(model)

    set!(model, enforce_incompressibility = false, u = (x, y, z) -> z, v = 2, w = 3)

    # Two ways to compute turbulent kinetic energy
    U = AveragedField(u, dims=(1, 2))
    V = AveragedField(v, dims=(1, 2))

    u′ = ComputedField(u - U)
    v′ = ComputedField(v - V)

    tke_op = @at (Center, Center, Center) (u′^2  + v′^2 + w^2) / 2
    tke = ComputedField(tke_op)
    compute!(tke)

    return all(interior(tke)[2:3, 2:3, 2:3] .== 9/2)
end

for arch in archs
    @testset "ComputedFields [$(typeof(arch))]" begin
        @info "  Testing ComputedFields [$(typeof(arch))]..."

        grid = RegularRectilinearGrid(size=(4, 4, 4), extent=(1, 1, 1),
                                      topology=(Periodic, Periodic, Bounded))

        buoyancy = SeawaterBuoyancy(gravitational_acceleration = 1,
                                             equation_of_state = LinearEquationOfState(α=1, β=1))

        model = NonhydrostaticModel(architecture = arch,
                                            grid = grid,
                                        buoyancy = buoyancy)

        @testset "Derivative computations [$(typeof(arch))]" begin
            @info "      Testing compute! derivatives..."
            @test compute_derivative(model, ∂x)
            @test compute_derivative(model, ∂y)
            @test compute_derivative(model, ∂z)
        end

        @testset "Unary computations [$(typeof(arch))]" begin
            @info "      Testing compute! unary operations..."
            for unary in (sqrt, sin, cos, exp, tanh)
                @test compute_unary(unary, model)
            end
        end

        @testset "Binary computations [$(typeof(arch))]" begin
            @info "      Testing compute! binary operations..."
            @test compute_plus(model)
            @test compute_minus(model)
            @test compute_times(model)

            # Basic compilation test for nested BinaryOperations...
            u, v, w = model.velocities
            @test try compute!(ComputedField(u + v - w)); true; catch; false; end
        end

        @testset "Multiary computations [$(typeof(arch))]" begin
            @info "      Testing compute! multiary operations..."
            @test compute_many_plus(model)

            @info "      Testing compute! kinetic energy..."
            @test compute_kinetic_energy(model)
        end

        @testset "Computations with KernelComputedField [$(typeof(arch))]" begin
            @test begin
                @inline trivial_kernel_function(i, j, k, grid) = 1
                op = KernelFunctionOperation{Center, Center, Center}(trivial_kernel_function, grid)
                f = ComputedField(op, arch)
                compute!(f)
                f isa ComputedField && f.operand === op
            end

            @test begin
                @inline trivial_parameterized_kernel_function(i, j, k, grid, μ) = μ
                op = KernelFunctionOperation{Center, Center, Center}(trivial_parameterized_kernel_function, grid, parameters=0.1)
                f = ComputedField(op, arch)
                compute!(f)
                f isa ComputedField && f.operand === op
            end

            @test begin
                u, v, w = model.velocities
                ζ_op = KernelFunctionOperation{Face, Face, Center}(ζ₃ᶠᶠᵃ, grid, computed_dependencies=(u, v))
                ζ = ComputedField(ζ_op) # identical to `VerticalVorticityField`
                compute!(ζ)
                ζ isa ComputedField && ζ.operand.kernel_function === ζ₃ᶠᶠᵃ
            end
        end

        @testset "Operations with ComputedField and PressureField [$(typeof(arch))]" begin
            @info "      Testing operations with ComputedField..."
            @test operations_with_computed_field(model)

            @info "      Testing PressureField..."
            @test pressure_field(model)
        end

        @testset "Horizontal averages of operations [$(typeof(arch))]" begin
            @info "      Testing horizontal averges..."
            @test horizontal_average_of_plus(model)
            @test horizontal_average_of_minus(model)
            @test horizontal_average_of_times(model)

            @test multiplication_and_derivative_ccf(model)
            @test multiplication_and_derivative_ccc(model)
        end

        @testset "Zonal averages of operations [$(typeof(arch))]" begin
            @info "      Testing zonal averges..."
            @test zonal_average_of_plus(model)
        end

        @testset "Volume averages of operations [$(typeof(arch))]" begin
            @info "      Testing volume averges..."
            @test volume_average_of_times(model)
        end

        @testset "ComputedField boundary conditions [$(typeof(arch))]" begin
            @info "      Testing boundary conditions for ComputedField..."

            set!(model; S=π, T=42)
            T, S = model.tracers

            @compute ST = ComputedField(S + T, data=model.pressures.pHY′.data)

            Nx, Ny, Nz = size(model.grid)

            @test all(ST.data[0, 1:Ny, 1:Nz]  .== ST.data[Nx+1, 1:Ny, 1:Nz])
            @test all(ST.data[1:Nx, 0, 1:Nz]  .== ST.data[1:Nx, Ny+1, 1:Nz])
            @test all(ST.data[1:Nx, 1:Ny, 0]  .== ST.data[1:Nx, 1:Ny, 1])
            @test all(ST.data[1:Nx, 1:Ny, Nz] .== ST.data[1:Nx, 1:Ny, Nz+1])

            @compute ST_face = ComputedField(@at (Center, Center, Face) S * T)

            @test all(ST_face.data[1:Nx, 1:Ny, 0] .== 0)
            @test all(ST_face.data[1:Nx, 1:Ny, Nz+2] .== 0)
        end

        @testset "Operations with AveragedField [$(typeof(arch))]" begin
            @info "      Testing operations with AveragedField..."

            T, S = model.tracers
            TS = AveragedField(T * S, dims=(1, 2))
            @test operations_with_averaged_field(model)
        end

        @testset "Compute! on faces along bounded dimensions" begin
            @info "      Testing compute! on faces along bounded dimensions..."
            @test computation_including_boundaries(arch)
        end

        EquationsOfState = (LinearEquationOfState, SeawaterPolynomials.RoquetEquationOfState,
                            SeawaterPolynomials.TEOS10EquationOfState)

        buoyancies = (BuoyancyTracer(), SeawaterBuoyancy(),
                      (SeawaterBuoyancy(equation_of_state=eos()) for eos in EquationsOfState)...)

        for buoyancy in buoyancies
            @testset "Computations with BuoyancyFields [$(typeof(arch)), $(typeof(buoyancy).name.wrapper)]" begin
                @info "      Testing computations with BuoyancyField " *
                      "[$(typeof(arch)), $(typeof(buoyancy).name.wrapper)]..."

                @test computations_with_buoyancy_field(arch, buoyancy)
            end
        end

        @testset "Computations with AveragedFields [$(typeof(arch))]" begin
            @info "      Testing computations with AveragedField [$(typeof(arch))]..."

            @test computations_with_averaged_field_derivative(model)

            u, v, w = model.velocities

            set!(model, enforce_incompressibility = false, u = (x, y, z) -> z, v = 2, w = 3)

            # Two ways to compute turbulent kinetic energy
            U = AveragedField(u, dims=(1, 2))
            V = AveragedField(v, dims=(1, 2))

            # Build up compilation tests incrementally...
            u_prime              = u - U
            u_prime_ccc          = @at (Center, Center, Center) u - U
            u_prime_squared      = (u - U)^2
            u_prime_squared_ccc  = @at (Center, Center, Center) (u - U)^2
            horizontal_twice_tke = (u - U)^2 + (v - V)^2
            horizontal_tke       = ((u - U)^2 + (v - V)^2) / 2
            horizontal_tke_ccc   = @at (Center, Center, Center) ((u - U)^2 + (v - V)^2) / 2
            twice_tke            = (u - U)^2  + (v - V)^2 + w^2
            tke                  = ((u - U)^2  + (v - V)^2 + w^2) / 2
            tke_ccc              = @at (Center, Center, Center) ((u - U)^2  + (v - V)^2 + w^2) / 2

            @test try compute!(ComputedField(u_prime             )); true; catch; false; end
            @test try compute!(ComputedField(u_prime_ccc         )); true; catch; false; end
            @test try compute!(ComputedField(u_prime_squared     )); true; catch; false; end
            @test try compute!(ComputedField(u_prime_squared_ccc )); true; catch; false; end
            @test try compute!(ComputedField(horizontal_twice_tke)); true; catch; false; end
            @test try compute!(ComputedField(horizontal_tke      )); true; catch; false; end
            @test try compute!(ComputedField(twice_tke           )); true; catch; false; end

            @test try compute!(ComputedField(horizontal_tke_ccc  )); true; catch; false; end
            @test try compute!(ComputedField(tke                 )); true; catch; false; end
            @test try compute!(ComputedField(tke_ccc             )); true; catch; false; end

            computed_tke = ComputedField(tke_ccc)
            @test (compute!(computed_tke); all(interior(computed_tke)[2:3, 2:3, 2:3] .== 9/2))
        end

        @testset "Computations with ComputedFields [$(typeof(arch))]" begin
            @info "      Testing computations with ComputedField [$(typeof(arch))]..."
            @test computations_with_computed_fields(model)
        end

        @testset "Conditional computation of ComputedField and BuoyancyField [$(typeof(arch))]" begin
            @info "      Testing conditional computation of ComputedField and BuoyancyField " *
                  "[$(typeof(arch))]..."

            set!(model, u=2, v=0, w=0, T=3, S=0)
            u, v, w, T, S = fields(model)

            uT = ComputedField(u * T)

            α = model.buoyancy.model.equation_of_state.α
            g = model.buoyancy.model.gravitational_acceleration
            b = BuoyancyField(model)

            compute_at!(uT, 1.0)
            compute_at!(b, 1.0)
            @test all(interior(uT) .== 6)
            @test all(interior(b) .== g * α * 3)

            set!(model, u=2, T=4)
            compute_at!(uT, 1.0)
            compute_at!(b, 1.0)
            @test all(interior(uT) .== 6)
            @test all(interior(b) .== g * α * 3)

            compute_at!(uT, 2.0)
            compute_at!(b, 2.0)
            @test all(interior(uT) .== 8)
            @test all(interior(b) .== g * α * 4)
        end
    end
end


text:
using Catalyst
using PWS
using Test
using StaticArrays
using Distributions

sn = @reaction_network begin
    κ, ∅ --> S
    λ, S --> ∅
end κ λ

κ = 0.25
λ = 0.005

dist = PWS.distribution(sn, [κ, λ])

traj = PWS.Trajectory([[50.0], [51.0], [50.0]], [1.0, 2.0, 3.0], [1, 2])

p_wait(s, dt) = exp(- (κ + s * λ) * dt)

wait1 = log(p_wait(50, 1.0))
wait2 = log(p_wait(51, 1.0))
wait3 = log(p_wait(50, 1.0))
reac1 = log(κ)
reac2 = log(51 * λ)

@test PWS.logpdf(dist, traj) ≈ sum((wait1, wait2, wait3, reac1, reac2))
@test PWS.trajectory_energy(dist, traj) ≈ sum((wait1, wait2, wait3, reac1, reac2))
@test PWS.trajectory_energy(dist, traj, tspan=(1.0,2.0)) ≈ sum((wait2, reac2))

cumulative = PWS.cumulative_logpdf(dist, traj, 0:0.5:3)
@test cumulative ≈ [
    0.0,
    0.5wait1,
    wait1 + reac1,
    wait1 + reac1 + 0.5wait2,
    wait1 + reac1 + wait2 + reac2,
    wait1 + reac1 + wait2 + reac2 + 0.5wait3,
    wait1 + reac1 + wait2 + reac2 + wait3,
]

subset = PWS.cumulative_logpdf(dist, traj, 2:0.33:3)
@test subset ≈ [
    0.0,
    0.33wait3,
    0.66wait3,
    0.99wait3
]

traj2 = PWS.Trajectory([[50.0]], [1.0])
@test PWS.trajectory_energy(dist, traj2, tspan=(0.0,0.5)) == -0.25


text:
module NeuralDataScience

include("helper_file.jl")

f(2,3)

export f

end


text:
__precompile__()

module Gencode

using SymPy

#export syminit, gencode, compilecode
export syminit, gencode, compilecode, runcode, checkcompilers, setcompilers, tring2cmd

include("syminit.jl");
include("varsassign.jl");
include("sympyassign.jl");
include("gencodebou.jl");
include("gencodeelem.jl");
include("gencodeelem2.jl");
include("gencodeelem3.jl");
include("gencodeface.jl");
include("gencodeface2.jl");
include("nocodeelem.jl");
include("nocodeelem2.jl");
include("nocodeelem3.jl");
include("nocodeface.jl");
include("nocodeface2.jl");
include("string2cmd.jl");
include("genlib.jl");
include("checkcompilers.jl");
include("setcompilers.jl");
include("compilecode.jl");
include("runcode.jl");

function gencode(app)

print("generate code...\n");
if !isdir("app")
    mkdir("app");
else
    if isfile("app/opuApp.a")
        rm("app/opuApp.a")
    end
    if isfile("app/cpuApp.a")
        rm("app/cpuApp.a")
    end
    if isfile("app/gpuApp.a")
        rm("app/gpuApp.a")
    end
end

xdg, udg, udg1, udg2, wdg, wdg1, wdg2, odg, odg1, odg2, uhg, nlg, tau, uinf, param, time = syminit(app);

ncu = app.ncu;
u = udg[1:ncu];
u1 = udg1[1:ncu];
u2 = udg2[1:ncu];
if app.nc>app.ncu
    q = udg[ncu+1:end];
    q1 = udg1[ncu+1:end];
    q2 = udg2[ncu+1:end];
else
    q = [];
    q1 = [];
    q2 = [];
end

if isdefined(Main, Symbol("flux"))
    #f = Main.flux(xdg, udg, odg, wdg, uinf, param, time);
    f = Main.flux(u, q, wdg, odg, xdg, time, param, uinf);
    if length(f)==1
        f = reshape([f],1,1);
    end
    f = f[:];
    gencodeelem("Flux", f, xdg, udg, odg, wdg, uinf, param, time);
else
    error("app.Flux is empty");
end
if isdefined(Main, Symbol("source"))
    #f = Main.source(xdg, udg, odg, wdg, uinf, param, time);
    f = Main.source(u, q, wdg, odg, xdg, time, param, uinf);
    if length(f)==1
        f = reshape([f],1,1);
    end
    f = f[:];
    gencodeelem("Source", f, xdg, udg, odg, wdg, uinf, param, time);
else
    nocodeelem("Source");
end
if isdefined(Main, Symbol("mass"))
    #f = Main.mass(xdg, udg, odg, wdg, uinf, param, time);
    f = Main.mass(u, q, wdg, odg, xdg, time, param, uinf);
    if length(f)==1
        f = reshape([f],1,1);
    end
    f = f[:];
    gencodeelem("Tdfunc", f, xdg, udg, odg, wdg, uinf, param, time);
else
    if app.model=="ModelW" || app.model == "modelW" || app.tdep==1
        error("pde.mass is not defined");
    else
        nocodeelem("Tdfunc");
    end
end
if isdefined(Main, Symbol("avifield"))
    #f = Main.avfield(xdg, udg, odg, wdg, uinf, param, time);
    f = Main.avfield(u, q, wdg, odg, xdg, time, param, uinf);
    if length(f)==1
        f = reshape([f],1,1);
    end
    f = f[:];
    gencodeelem2("Avfield", f, xdg, udg, odg, wdg, uinf, param, time);
else
    nocodeelem2("Avfield");
end
if isdefined(Main, Symbol("output"))
    #f = Main.output(xdg, udg, odg, wdg, uinf, param, time);
    f = Main.output(u, q, wdg, odg, xdg, time, param, uinf);
    if length(f)==1
        f = reshape([f],1,1);
    end
    f = f[:];
    gencodeelem2("Output", f, xdg, udg, odg, wdg, uinf, param, time);
else
    nocodeelem2("Output");
end
if isdefined(Main, Symbol("fbou"))
    #f = Main.fbou(xdg, udg, odg, wdg, uhg, nlg, tau, uinf, param, time);
    f = Main.fbou(u, q, wdg, odg, xdg, time, param, uinf, uhg, nlg, tau);
    if length(f)==1
        f = reshape([f],1,1);
    end
    f = f[:];
    f = reshape(f,app.ncu,Int(length(f)/app.ncu));
    gencodeface("Fbou", f, xdg, udg, odg, wdg, uhg, nlg, tau, uinf, param, time);
else
    error("app.Fbou is empty");
end
if isdefined(Main, Symbol("ubou"))
    #f = Main.ubou(xdg, udg, odg, wdg, uhg, nlg, tau, uinf, param, time);
    f = Main.ubou(u, q, wdg, odg, xdg, time, param, uinf, uhg, nlg, tau);
    if length(f)==1
        f = reshape([f],1,1);
    end
    f = f[:];
    f = reshape(f,app.ncu,Int(length(f)/app.ncu));
    gencodeface("Ubou", f, xdg, udg, odg, wdg, uhg, nlg, tau, uinf, param, time);
else
    nocodeface("Ubou");
end
if isdefined(Main, Symbol("Fhat"))
    #f = Main.fhat(xdg, udg1, udg2, odg1, odg2, wdg1, wdg2, uhg, nlg, tau, uinf, param, time);
    f = Main.fhat(u1, q1, wdg1, odg1, xdg, time, param, uinf, uhg, nlg, tau, u2, q2, wdg2, odg2);
    if length(f)==1
        f = reshape([f],1,1);
    end
    f = f[:];
    gencodeface2("Fhat", f, xdg, udg1, udg2, odg1, odg2, wdg1, wdg2, uhg, nlg, tau, uinf, param, time);
else
    nocodeface2("Fhat");
end
if isdefined(Main, Symbol("uhat"))
    #f = Main.uhat(xdg, udg1, udg2, odg1, odg2, wdg1, wdg2, uhg, nlg, tau, uinf, param, time);
    f = Main.uhat(u1, q1, wdg1, odg1, xdg, time, param, uinf, uhg, nlg, tau, u2, q2, wdg2, odg2);
    if length(f)==1
        f = reshape([f],1,1);
    end
    f = f[:];
    gencodeface2("Uhat", f, xdg, udg1, udg2, odg1, odg2, wdg1, wdg2, uhg, nlg, tau, uinf, param, time);
else
    nocodeface2("Uhat");
end
if isdefined(Main, Symbol("stab"))
    #f = Main.stab(xdg, udg1, udg2, odg1, odg2, wdg1, wdg2, uhg, nlg, tau, uinf, param, time);
    f = Main.stab(u1, q1, wdg1, odg1, xdg, time, param, uinf, uhg, nlg, tau, u2, q2, wdg2, odg2);
    if length(f)==1
        f = reshape([f],1,1);
    end
    f = f[:];
    gencodeface2("Stab", f, xdg, udg1, udg2, odg1, odg2, wdg1, wdg2, uhg, nlg, tau, uinf, param, time);
else
    nocodeface2("Stab");
end
if isdefined(Main, Symbol("initu"))
    udg = Main.initu(xdg, param, uinf);
    if length(udg)==1
        udg = reshape([udg],1,1);
    end
    udg = udg[:];
    gencodeelem3("Initu", udg, xdg, uinf, param);
else
    error("initu is not defined");
end
if isdefined(Main, Symbol("initw"))
    wdg = Main.initw(xdg, param, uinf);
    if length(wdg)==1
        wdg = reshape([wdg],1,1);
    end
    wdg = wdg[:];
    gencodeelem3("Initwdg", wdg, xdg, uinf, param);
else
    nocodeelem3("Initwdg");
end
if isdefined(Main, Symbol("initv"))
    odg = Main.initv(xdg, param, uinf);
    if length(odg)==1
        odg = reshape([odg],1,1);
    end
    odg = odg[:];
    gencodeelem3("Initodg", odg, xdg, uinf, param);
else
    nocodeelem3("Initodg");
end
if isdefined(Main, Symbol("initq"))
    qdg = Main.initq(xdg, param, uinf);
    if length(qdg)==1
        qdg = reshape([qdg],1,1);
    end
    qdg = qdg[:];
    gencodeelem3("Initq", qdg, xdg, uinf, param);

    udg = Main.initu(xdg, param, uinf);
    if length(udg)==1
        udg = reshape([udg],1,1);
    end
    udg = udg[:];

    udg = [udg; qdg];
    gencodeelem3("Initudg", udg, xdg, uinf, param);
else
    if app.model == "ModelW" || app.model == "modelW"
        error("initq is not defined");
    else
        nocodeelem3("Initq");
        nocodeelem3("Initudg");
    end
end

# if isdefined(Main, Symbol("initq"))
#     qdg = Main.initq(xdg, uinf, param);
#     if length(qdg)==1
#         qdg = reshape([qdg],1,1);
#     end
#     qdg = qdg[:];
#     gencodeelem3("Initq", qdg, xdg, uinf, param);
# else
#     nocodeelem3("Initq");
# end
# if isdefined(Main, Symbol("inituq"))
#     udg = Main.inituq(xdg, uinf, param);
#     if length(udg)==1
#         udg = reshape([udg],1,1);
#     end
#     udg = udg[:];
#     gencodeelem3("Initudg", udg, xdg, uinf, param);
# else
#     if app.model == "ModelW" || app.model == "modelW"
#         error("inituq is not defined");
#     else
#         nocodeelem3("Initudg");
#     end
# end

end

end

max_stars_repo_name: glwagner/Exasim

max_issues_repo_name: glwagner/Exasim

max_forks_repo_name: glwagner/Exasim


text:
using AutoHOOT
using Test

const ad = AutoHOOT.autodiff
const T = AutoHOOT.backend

T.set_backend("numpy")

@testset "test einsum" begin
    x2 = ad.Variable(name = "x2", shape = [3, 2])
    x3 = ad.Variable(name = "x3", shape = [2, 3])
    matmul = ad.einsum("ik,kj->ij", x2, x3)
    y = ad.sum(matmul)

    grad_x2, grad_x3 = ad.gradients(y, [x2, x3])

    executor = ad.Executor([y, grad_x2, grad_x3])
    x2_val = T.tensor([[1, 2], [3, 4], [5, 6]])  # 3x2
    x3_val = T.tensor([[7, 8, 9], [10, 11, 12]])  # 2x3

    y_val, grad_x2_val, grad_x3_val =
        executor.run(feed_dict = Dict(x2 => x2_val, x3 => x3_val))

    expected_grad_sum = T.ones_like(T.dot(x2_val, x3_val))
    expected_yval = T.sum(T.dot(x2_val, x3_val))
    expected_grad_x2_val = T.dot(expected_grad_sum, T.transpose(x3_val))
    expected_grad_x3_val = T.dot(T.transpose(x2_val), expected_grad_sum)

    @test T.array_equal(y_val, expected_yval)
    @test T.array_equal(grad_x2_val, expected_grad_x2_val)
    @test T.array_equal(grad_x3_val, expected_grad_x3_val)
end

@testset "test vjps" begin
    x = ad.Variable(name = "x", shape = [2])
    A = ad.Variable(name = "A", shape = [3, 2])
    v = ad.Variable(name = "v", shape = [3])
    y = ad.einsum("ab, b->a", A, x)

    transposed_vjp_x, = ad.transposed_vjps(y, [x], v)

    executor = ad.Executor([y, transposed_vjp_x])
    x_val = T.tensor([1.0, 2.0])  # 1x3
    A_val = T.tensor([[1.0, 2.0], [3.0, 4.0], [5, 6]])
    v_val = T.tensor([1.0, 2.0, 3.0])

    y_val, transposed_vjp_x_val =
        executor.run(feed_dict = Dict(x => x_val, A => A_val, v => v_val))

    expected_yval = T.einsum("ab, b->a", A_val, x_val)
    expected_transposed_vjp_x_val = T.einsum("b, ba->a", v_val, A_val)

    @test T.array_equal(y_val, expected_yval)
    @test T.array_equal(transposed_vjp_x_val, expected_transposed_vjp_x_val)
end

@testset "test inner product hvp" begin
    x = ad.Variable(name = "x", shape = [3, 1])
    v = ad.Variable(name = "v", shape = [3, 1])
    y = ad.sum(ad.einsum("ab,bc->ac", ad.transpose(x), x))

    grad_x, = ad.gradients(y, [x])
    Hv, = ad.hvp(output_node = y, node_list = [x], vector_list = [v])

    executor = ad.Executor([y, grad_x, Hv])
    x_val = T.tensor([[1.0], [2.0], [3]])  # 3x1
    v_val = T.tensor([[1.0], [2.0], [3]])  # 3x1
    y_val, grad_x_val, Hv_val = executor.run(feed_dict = Dict(x => x_val, v => v_val))

    expected_yval = T.sum(T.dot(T.transpose(x_val), x_val))
    expected_grad_x_val = 2 * x_val
    expected_hv_val = 2 * v_val

    @test T.array_equal(y_val, expected_yval)
    @test T.array_equal(grad_x_val, expected_grad_x_val)
    @test T.array_equal(Hv_val, expected_hv_val)
end


text:
include("euler/euler.jl")
using .Calculus: get_fibonacci

compute(n::Int)::Int = get_fibonacci(n + 2)

max_stars_repo_name: Dynortice/Project-Euler

max_issues_repo_name: Dynortice/Project-Euler

max_forks_repo_name: Dynortice/Project-Euler


text:
# First slider
step1 = data_type_slider(options = [:raster, :network])

val1 = dynamically_get_val(slider1)

sel2 = Observable{String}("")
step2 = Observable{Any}()
on(val1) do val1
    sel2[] = "Pairwise"
    if val1 == "Raster"
        slider2 = modelling_mode_slider(options = [:pairwise,
                                                :advanced, :onetoall, :alltoone:])
        val2 = dyanamically_get_val(slider2)
        on(val2) do x
            sel2[] = x
        end
    else
        slider2 = modelling_mode_slider(options = [:pairwise, :advanced])
        val2 = dynamically_get_val(slider2)
    end

    if val2 == "Pairwise"
        ui = dynamically_generate_pairwise_ui()
    elseif val2 == "Advanced"
        ui = dynamically_generate_advanced_ui()
    else
        ui = dynamically_generate_pairwise_ui()
    end
        step2[] = ui
    end
end

# output_ui = generate_output_ui() # With the logging box and the pretty printed resistance matrix
val1["value"][] = "Raster"

push_to_page(vbox(step1,
                  step1,
                  ui))
                  #output_ui))


text:
using JuMP, Plots, LinearAlgebra, SparseArrays, Ipopt

#A JuMP model is initialized
mod = Model(optimizer_with_attributes(Ipopt.Optimizer ,  "max_iter" => 100,
                                        "mumps_mem_percent" => 500))
#The parameters of the 2-D body
λ = 45.0                            # Conductivity
c = 460.0                           # Specific heat capacitivity
ρ = 7800.0                          # Material density
L = 0.1                             #Length and width of body are assumed equivalent
α = λ / (c*ρ)                       # Diffusivity
N = 11                              # Number of Discretization points N_x = N_y
h = L / (N - 1)                     #Δx = Δy = h

#System , input and output matrices: θ' = Aθ(t) + Bu(t)
A = spdiagm(0 => -4*ones(N^2))
for i in 1:N^2
    if i%N == 1
        A[i+1,i] = 2
    elseif i%N == 0
        A[i-1,i] = 2
    else
        A[i-1,i] = 1
        A[i+1,i] = 1
    end
    if (i <= N)
        A[i+N,i] = 2
    elseif N^2 - i < N
        A[i- N,i] = 2
    else
        A[i-N,i] = 1
        A[i+N,i] = 1
    end
end
A *= α / (h^2);

B = spzeros(3,N^2);
B[1,1] = 1;
B[2,6] = 1;
B[3,11] = 1;
B *= (2/ (c * ρ * h));
#y(t) = Cθ(t)
C = spzeros(N^2,1)
C[111] = 1.0;
C[115] = 1.0;
C[end] = 1.0;
#Integration from t0 = 0 to tf
tf = 40000.0;                   #Final time
dt = 1.6;                       #Sampling period Δt
steps = round(Int,tf / dt);     #number of time steps in simulation.

u_min = 0.0                     #lower input bound
u_max = 15000.0                 #upper input bound
@variable(mod,θ[1:N^2,1:steps]);                                      # state vector
@variable(mod,u_min <= u[i = 1:3,k = 1:steps - 1] <= u_max);          # Bounded Input u

θinit = 273.0                           #initial temperature of the body
θref = 500.0                            #Reference temperature

@constraint(mod, θ[:,1] .== θinit)      #Initial value added as constraints
#System dynamics
for k in 1:steps - 1,i in 1:N^2
    @constraint(mod,θ[i,k + 1] .==  θ[i,k] + dt * ((A[:,i]' * θ[:,k]) + B[:,i]' * u[:,k]))
end
#Q and R weighing matrices
Q = 10000 * C * C';
R = 0.0001 * I;
#Cost function
@NLobjective(mod, Min, 0.5 * dt * sum(sum(Q[i,i] *(θref - θ[i,k])^2 for i in 1:N^2)
+ sum(R[i,i] * u[i,k]^2 for i in 1:3) for k in 1:steps - 1))
#Optimization process
optimize!(mod)

#Plotting results
states_θstart1 = zeros(steps)
states_θstart2 = zeros(steps)
states_θstart3 = zeros(steps)
states_θend3 = zeros(steps)
states_θend2 = zeros(steps)
states_θend1 = zeros(steps)
input_u1 = zeros(steps - 1)
input_u2 = zeros(steps - 1)
input_u3 = zeros(steps - 1)
for i = 1 : steps-1
    states_θstart1[i] = JuMP.value(θ[1,i])
    states_θstart2[i] = JuMP.value(θ[6,i])
    states_θstart3[i] = JuMP.value(θ[11,i])
    states_θend3[i] = JuMP.value(θ[115,i])
    states_θend2[i] = JuMP.value(θ[85,i])
    states_θend1[i] = JuMP.value(θ[45,i])
    input_u1[i] = JuMP.value(u[1,i])
    input_u2[i] = JuMP.value(u[2,i])
    input_u3[i] = JuMP.value(u[3,i])
end

states_θend3[end] = JuMP.value(θ[end,end])
states_θend2[end] = JuMP.value(θ[end-5,end])
states_θend1[end] = JuMP.value(θ[end-10,end])
states_θstart1[end] = JuMP.value(θ[1,end])
states_θstart2[end] = JuMP.value(θ[6,end])
states_θstart3[end] = JuMP.value(θ[11,end])
#Plotting temperatures at the input nodes
p2 = plot(states_θstart1, label = "θ_in1",legend = false)
p2 = plot!(states_θstart2,label = "θ_in2",legend = false)
p2 = plot!(states_θstart3,label = "θ_in3",legend = false)
savefig("lq_optimal_2D.png")
#Plotting temperatures at the output
p1 = plot(states_θend3,label = "output-3",legend = false)
p1 = plot!(states_θend2,label = "output-2",legend = false)
p1 = plot!(states_θend1,label = "output-1",legend = false)
savefig("lq_optimal_2D-output-bounded.png")
#Plotting the input signals
p3 = plot(input_u1,label = "in_1",legend = false)
p3 = plot!(input_u2,label = "in_2",legend = false)
p3 = plot!(input_u3,label = "in_3",legend = false)
savefig("lq_optimal_2D_inputbounded.png")


text:
module DRMSMIP

using JuMP, Ipopt, GLPK
using DualDecomposition, BundleMethod

const DD = DualDecomposition
const BM = BundleMethod

include("AmbiguitySet.jl")
include("ScenarioTree.jl")

include("BlockModel.jl")
include("LagrangeDual.jl")

end


text:
module DFMisc
  using DFControl
  using DFWannier
# package code goes here
  include("types.jl")
  include("io.jl")
  include("construct.jl")
  include("calcs.jl")
end # module


text:
module MatrixFactorizations

println("You have decided to learn Applications of Matrix Factorizations using Julia. Please refer to the document Julia_MatFact.pdf")

end


text:
export LocalCSRMatrix, numRows, numCols, getRowView

struct LocalCSRMatrix{Data, IndexType <: Integer}
    graph::LocalCSRGraph{IndexType, IndexType}
    values::Union{Vector{Data}, SubArray{Data, 1, Vector{Data}, Tuple{UnitRange{IndexType}}, true}}

    numCols::IndexType
end

"""
    LocalCSRMatrix{Data, IndexType}()

Creates an empty LocalCSRMatrix
"""
function LocalCSRMatrix{Data, IndexType}() where {Data, IndexType}
    LocalCSRMatrix(LocalCSRGraph{IndexType, IndexType}(), Data[], IndexType(0))
end

"""
    LocalCSRMatrix(nRows::Integer, nCols::Integer, vals::AbstractArray{Data, 1}, rows::AbstractArray{IndexType, 1}, cols::AbstractArray{IndexType, 1}}

Creates the specified LocalCSRMatrix
"""
function LocalCSRMatrix(nRows::Integer, nCols::Integer,
        vals::AbstractArray{Data, 1}, rows::AbstractArray{IndexType, 1},
        cols::AbstractArray{IndexType, 1}) where {Data, IndexType}
    if length(rows) != nRows + 1
        throw(InvalidArgumentError("length(rows) = $(length(rows)) != nRows+1 "
                * "= $(nRows + 1)"))
    end
    LocalCSRMatrix(LocalCSRGraph(cols, rows), vals, IndexType(nCols))
end

"""
    LocalCSRMatrix(numCols::IndexType, values::AbstractArray{Data, 1}, localGraph::LocalCSRGraph{IndexType, IndexType}) where {IndexType, Data <: Number}

Creates the specified LocalCSRMatrix
"""
function LocalCSRMatrix(numCols::IndexType, values::AbstractArray{Data, 1},
        localGraph::LocalCSRGraph{IndexType, IndexType}) where {IndexType, Data <: Number}
    if numCols < 0
        throw(InvalidArgumentError("Cannot have a negative number of rows"))
    end

    LocalCSRMatrix(localGraph, values, numCols)
end

"""
    numRows(::LocalCSRMatrix{Data, IndexType})::IndexType

Gets the number of rows in the matrix
"""
numRows(matrix::LocalCSRMatrix) = numRows(matrix.graph)

"""
    numCols(::LocalCSRMatrix{Data, IndexType})::IndexType

Gets the number of columns in the matrix
"""
numCols(matrix::LocalCSRMatrix) = matrix.numCols

"""
    getRowView((matrix::LocalCSRMatrix{Data, IndexType}, row::Integer)::SparseRowView{Data, IndexType}

Gets a view of the requested row
"""
function getRowView(matrix::LocalCSRMatrix{Data, IndexType},
        row::Integer)::SparseRowView{Data, IndexType} where {Data, IndexType}
    start = matrix.graph.rowMap[row]
    count = matrix.graph.rowMap[row+1] - start

    if count == 0
        SparseRowView(Data[], IndexType[])
    else
        SparseRowView(matrix.values, matrix.graph.entries, count, start)
    end
end