text:
using FunctionalCollections
using Test

const PS = PersistentSet

@testset "Persistent Sets" begin

    @testset "construction" begin
        s = PS([1, 1, 2, 3, 3])
        @test length(s) == 3
        @test length(PS{String}()) == 0
        @test typeof(PS{Int64}([1,2,3])) == PS{Int64}
        @test typeof(PS(Int64[1,2,3])) == PS{Int64}
    end

    @testset "isequal" begin
        @test PS([1, 2, 3]) == PS([1, 2, 3])
        @test PS([1, 2, 3]) == PS([3, 2, 1])
        @test PS{String}() == PS{Int}()
    end

    @testset "conj" begin
        @test conj(PS([1, 2, 3]), 4) == PS([1, 2, 3, 4])
        @test conj(PS([1, 2, 3]), 1) == PS([1, 2, 3])
        @test conj(PS([1, 2, 3]), 4) == PS([4, 3, 2, 1])
    end

    @testset "disj" begin
        @test disj(PS([1, 2, 3]), 3) == PS([1, 2])
        @test disj(PS([1, 2]), 3) == PS([1, 2])
        @test disj(PS{Int}(), 1234) == PS{Int}()
    end

    @testset "in" begin
        @test "foo" in PS(["foo", "bar"])
        @test !("baz" in PS(["foo", "bar"]))
    end

    @testset "filter" begin
        @test filter(iseven, PS([1, 2, 3, 4])) == PS([2, 4])
    end

    @testset "setdiff, -" begin
        @test setdiff(PS([1, 2, 3]), PS([1, 2])) == PS([3])
        @test setdiff(PS([1, 2]), PS([1, 2, 3])) == PS{Int}()
        @test setdiff(PS([1, 2, 3]), Set([1, 2])) == PS([3])

        @test PS([1, 2, 3]) - PS([1, 2]) == PS([3])
    end

    @testset "isempty" begin
        @test isempty(PS{Int}())
        @test !isempty(PS([1]))
    end

    @testset "union" begin
        @test union(PS([1, 2, 3]), PS([4, 5])) == PS([1, 2, 3, 4, 5])
        @test union(PS([1, 2, 3]), PS([1, 2, 3, 4])) == PS([1, 2, 3, 4])
        @test union(PS([1]), PS([2]), PS([3])) == PS([1, 2, 3])
    end

    @testset "isless, <=" begin
        @test PS([1])    <= PS([1, 2])
        @test PS([1, 2]) <= PS([1, 2])

        @test !(PS([1, 2, 3]) <= PS([1, 2]))
        @test PS([1, 2]) <= PS([1, 2, 3])

        @test !isless(PS([1, 2]), PS([1, 2]))
        @test isless(PS([1, 2]), PS([1, 2, 3]))
    end

    @testset "iteration" begin
        @test length([el for el in PS([1, 2, 3, 4])]) == 4
    end

end


text:
#== # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
#
#   Tests related to IAU-76/FK5 transformations.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# References
#
#   [1] Vallado, D. A (2013). Fundamentals of Astrodynamics and Applications.
#       Microcosm Press, Hawthorn, CA, USA.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # ==#

include("./nutation.jl")
include("./precession.jl")

# File: ./src/transformations/fk5/fk5.jl
# ======================================

# Functions rITRFtoPEF_fk5 and rPEFtoITRF_fk5
# -------------------------------------------

################################################################################
#                                 Test Results
################################################################################
#
# Scenario 01
# ===========
#
# Example 3-15: Performing IAU-76/FK5 reduction.
#
# According to this example and Table 3-6, using:
#
#   x_p = -0.140682"
#   y_p = +0.333309"
#   r_itrf = -1033.4793830    i + 7901.2952754    j + 6380.3565958    k [km]
#   v_itrf =    -3.225636520  i -    2.872451450  j +    5.531924446  k [km/s]
#
# one gets the following data:
#
#   r_pef  = -1033.47503130   i + 7901.30558560   j + 6380.34453270   k [km]
#   v_pef  =    -3.2256327470 i -    2.8724425110 j +    5.5319312880 k [km/s]
#
################################################################################

@testset "Functions rITRFtoPEF_fk5 and rPEFtoITRF_fk5" begin
    x_p = -0.140682*pi/(180*3600)
    y_p = +0.333309*pi/(180*3600)

    ## rITRFtoPEF_fk5
    ## ==============

    r_itrf = [-1033.4793830; 7901.2952754; 6380.3565958]
    v_itrf = [-3.225636520; -2.872451450; +5.531924446]

    ## DCM
    ## ---

    D_PEF_ITRF = rITRFtoPEF_fk5(x_p, y_p)

    r_pef = D_PEF_ITRF*r_itrf
    v_pef = D_PEF_ITRF*v_itrf

    @test r_pef[1] ≈ -1033.47503130 atol=1e-7
    @test r_pef[2] ≈ +7901.30558560 atol=1e-7
    @test r_pef[3] ≈ +6380.34453270 atol=1e-7

    @test v_pef[1] ≈ -3.2256327470  atol=1e-9
    @test v_pef[2] ≈ -2.8724425110  atol=1e-9
    @test v_pef[3] ≈ +5.5319312880  atol=1e-9

    ## Quaternion
    ## ----------

    q_PEF_ITRF = rITRFtoPEF_fk5(Quaternion, x_p, y_p)

    r_pef = vect(conj(q_PEF_ITRF)*r_itrf*q_PEF_ITRF)
    v_pef = vect(conj(q_PEF_ITRF)*v_itrf*q_PEF_ITRF)

    @test r_pef[1] ≈ -1033.47503130 atol=1e-7
    @test r_pef[2] ≈ +7901.30558560 atol=1e-7
    @test r_pef[3] ≈ +6380.34453270 atol=1e-7

    @test v_pef[1] ≈ -3.2256327470  atol=1e-9
    @test v_pef[2] ≈ -2.8724425110  atol=1e-9
    @test v_pef[3] ≈ +5.5319312880  atol=1e-9

    ## rPEFtoITRF_fk5
    ## ==============

    r_pef  = [-1033.47503130; 7901.30558560; 6380.34453270]
    v_pef  = [-3.2256327470; -2.8724425110; +5.5319312880]

    ## DCM
    ## ---

    D_ITRF_PEF = rPEFtoITRF_fk5(x_p, y_p)
    r_itrf = D_ITRF_PEF*r_pef
    v_itrf = D_ITRF_PEF*v_pef

    @test r_itrf[1] ≈ -1033.4793830 atol=1e-7
    @test r_itrf[2] ≈ +7901.2952754 atol=1e-7
    @test r_itrf[3] ≈ +6380.3565958 atol=1e-7

    @test v_itrf[1] ≈ -3.225636520  atol=1e-9
    @test v_itrf[2] ≈ -2.872451450  atol=1e-9
    @test v_itrf[3] ≈ +5.531924446  atol=1e-9

    ## Quaternion
    ## ----------

    q_ITRF_PEF = rPEFtoITRF_fk5(Quaternion, x_p, y_p)
    r_itrf = vect(conj(q_ITRF_PEF)*r_pef*q_ITRF_PEF)
    v_itrf = vect(conj(q_ITRF_PEF)*v_pef*q_ITRF_PEF)

    @test r_itrf[1] ≈ -1033.4793830 atol=1e-7
    @test r_itrf[2] ≈ +7901.2952754 atol=1e-7
    @test r_itrf[3] ≈ +6380.3565958 atol=1e-7

    @test v_itrf[1] ≈ -3.225636520  atol=1e-9
    @test v_itrf[2] ≈ -2.872451450  atol=1e-9
    @test v_itrf[3] ≈ +5.531924446  atol=1e-9
end

# Functions rPEFtoTOD_fk5 and rTODtoPEF_fk5
# -----------------------------------------

################################################################################
#                                 Test Results
################################################################################
#
# Scenario 01
# ===========
#
# Example 3-15: Performing IAU-76/FK5 reduction.
#
# According to this example and Table 3-6, using:
#
#   JD_UT1 = DatetoJD(2004,4,6,7,51,28.386009) - 0.4399619/86400
#   JD_TT  = 2453101.828154745
#   LOD    = 0.0015563 s
#   δΔψ    = -0.052195"
#   r_pef  = -1033.47503130   i + 7901.30558560   j + 6380.34453270   k [km]
#   v_pef  =    -3.2256327470 i -    2.8724425110 j +    5.5319312880 k [km/s]
#
# one gets the following data:
#
#   r_tod = 5094.51620300   i + 6127.36527840   j + 6380.34453270   k [km]
#   v_tod =   -4.7460883850 i +    0.7860783240 j +    5.5319312880 k [km/s]
#
# Furthermore, using:
#
#   JD_TT = 2453101.828154745
#   δΔψ   = 0"
#   r_pef = -1033.47503130   i + 7901.30558560   j + 6380.34453270   k [km]
#   v_pef =    -3.2256327470 i -    2.8724425110 j +    5.5319312880 k [km/s]
#
# one gets the following data:
#
#   r_tod = 5094.51478040   i + 6127.36646120   j + 6380.34453270   k [km]
#   v_tod =   -4.7460885670 i +    0.7860772220 j +    5.5319312880 k [km/s]
#
################################################################################

@testset "Functions rPEFtoTOD_fk5 and rTODtoPEF_fk5" begin
    JD_UT1 = DatetoJD(2004,4,6,7,51,28.386009) - 0.4399619/86400
    LOD    = 0.0015563
    w      = 7.292115146706979e-5*(1-LOD/86400)

    ## rPEFtoTOD_fk5
    ## =============

    ## -------------------------------------------------------------------------
    ##                                 First Test
    ## -------------------------------------------------------------------------

    r_pef  = [-1033.47503130; 7901.30558560; 6380.34453270]
    v_pef  = [-3.2256327470; -2.8724425110; +5.5319312880]

    ## DCM
    ## ---

    D_TOD_PEF = rPEFtoTOD_fk5(JD_UT1, 2453101.828154745, -0.052195*pi/(180*3600))

    r_tod = D_TOD_PEF*r_pef
    v_tod = D_TOD_PEF*(v_pef + [0;0;w] × r_pef)

    @test r_tod[1] ≈ +5094.51620300 atol=1e-7
    @test r_tod[2] ≈ +6127.36527840 atol=1e-7
    @test r_tod[3] ≈ +6380.34453270 atol=1e-7

    @test v_tod[1] ≈ -4.7460883850  atol=1e-9
    @test v_tod[2] ≈ +0.7860783240  atol=1e-9
    @test v_tod[3] ≈ +5.5319312880  atol=1e-9

    ## Quaternion
    ## ----------

    q_TOD_PEF = rPEFtoTOD_fk5(Quaternion,
                              JD_UT1,
                              2453101.828154745,
                              -0.052195*pi/(180*3600))

    r_tod = vect(conj(q_TOD_PEF)*r_pef*q_TOD_PEF)
    v_tod = vect(conj(q_TOD_PEF)*(v_pef + [0;0;w] × r_pef)*q_TOD_PEF)

    @test r_tod[1] ≈ +5094.51620300 atol=1e-7
    @test r_tod[2] ≈ +6127.36527840 atol=1e-7
    @test r_tod[3] ≈ +6380.34453270 atol=1e-7

    @test v_tod[1] ≈ -4.7460883850  atol=1e-9
    @test v_tod[2] ≈ +0.7860783240  atol=1e-9
    @test v_tod[3] ≈ +5.5319312880  atol=1e-9

    ## -------------------------------------------------------------------------
    ##                                Second Test
    ## -------------------------------------------------------------------------

    ## DCM
    ## ---

    D_TOD_PEF = rPEFtoTOD_fk5(JD_UT1, 2453101.828154745)

    r_tod = D_TOD_PEF*r_pef
    v_tod = D_TOD_PEF*(v_pef + [0;0;w] × r_pef)

    @test r_tod[1] ≈ +5094.51478040 atol=1e-7
    @test r_tod[2] ≈ +6127.36646120 atol=1e-7
    @test r_tod[3] ≈ +6380.34453270 atol=1e-7

    @test v_tod[1] ≈ -4.7460885670  atol=1e-9
    @test v_tod[2] ≈ +0.7860772220  atol=1e-9
    @test v_tod[3] ≈ +5.5319312880  atol=1e-9

    ## Quaternion
    ## ----------

    q_TOD_PEF = rPEFtoTOD_fk5(Quaternion,
                              JD_UT1,
                              2453101.828154745)

    r_tod = vect(conj(q_TOD_PEF)*r_pef*q_TOD_PEF)
    v_tod = vect(conj(q_TOD_PEF)*(v_pef + [0;0;w] × r_pef)*q_TOD_PEF)

    @test r_tod[1] ≈ +5094.51478040 atol=1e-7
    @test r_tod[2] ≈ +6127.36646120 atol=1e-7
    @test r_tod[3] ≈ +6380.34453270 atol=1e-7

    @test v_tod[1] ≈ -4.7460885670  atol=1e-9
    @test v_tod[2] ≈ +0.7860772220  atol=1e-9
    @test v_tod[3] ≈ +5.5319312880  atol=1e-9

    ## rTODtoPEF_fk5
    ## =============

    ## -------------------------------------------------------------------------
    ##                                 First Test
    ## -------------------------------------------------------------------------

    r_tod = [5094.51620300; 6127.36527840; 6380.34453270]
    v_tod = [-4.7460883850; 0.7860783240; 5.5319312880]

    ## DCM
    ## ---

    D_PEF_TOD = rTODtoPEF_fk5(JD_UT1, 2453101.828154745, -0.052195*pi/(180*3600))

    r_pef = D_PEF_TOD*r_tod
    v_pef = D_PEF_TOD*v_tod - [0;0;w] × r_pef

    @test r_pef[1] ≈ -1033.47503130 atol=1e-7
    @test r_pef[2] ≈ +7901.30558560 atol=1e-7
    @test r_pef[3] ≈ +6380.34453270 atol=1e-7

    @test v_pef[1] ≈ -3.2256327470  atol=1e-9
    @test v_pef[2] ≈ -2.8724425110  atol=1e-9
    @test v_pef[3] ≈ +5.5319312880  atol=1e-9

    ## Quaternion
    ## ----------

    q_PEF_TOD = rTODtoPEF_fk5(Quaternion,
                              JD_UT1,
                              2453101.828154745,
                              -0.052195*pi/(180*3600))

    r_pef = vect(conj(q_PEF_TOD)*r_tod*q_PEF_TOD)
    v_pef = vect(conj(q_PEF_TOD)*v_tod*q_PEF_TOD) - [0;0;w] × r_pef

    @test r_pef[1] ≈ -1033.47503130 atol=1e-7
    @test r_pef[2] ≈ +7901.30558560 atol=1e-7
    @test r_pef[3] ≈ +6380.34453270 atol=1e-7

    @test v_pef[1] ≈ -3.2256327470  atol=1e-9
    @test v_pef[2] ≈ -2.8724425110  atol=1e-9
    @test v_pef[3] ≈ +5.5319312880  atol=1e-9

    ## -------------------------------------------------------------------------
    ##                                Second Test
    ## -------------------------------------------------------------------------

    r_tod = [5094.51478040; 6127.36646120; 6380.34453270]
    v_tod = [-4.7460885670; 0.7860772220; 5.5319312880]

    ## DCM
    ## ---

    D_PEF_TOD = rTODtoPEF_fk5(JD_UT1, 2453101.828154745)

    r_pef = D_PEF_TOD*r_tod
    v_pef = D_PEF_TOD*v_tod - [0;0;w] × r_pef

    @test r_pef[1] ≈ -1033.47503130 atol=1e-7
    @test r_pef[2] ≈ +7901.30558560 atol=1e-7
    @test r_pef[3] ≈ +6380.34453270 atol=1e-7

    @test v_pef[1] ≈ -3.2256327470  atol=1e-9
    @test v_pef[2] ≈ -2.8724425110  atol=1e-9
    @test v_pef[3] ≈ +5.5319312880  atol=1e-9

    ## Quaternion
    ## ----------

    q_PEF_TOD = rTODtoPEF_fk5(Quaternion,
                              JD_UT1,
                              2453101.828154745)

    r_pef = vect(conj(q_PEF_TOD)*r_tod*q_PEF_TOD)
    v_pef = vect(conj(q_PEF_TOD)*v_tod*q_PEF_TOD) - [0;0;w] × r_pef

    @test r_pef[1] ≈ -1033.47503130 atol=1e-7
    @test r_pef[2] ≈ +7901.30558560 atol=1e-7
    @test r_pef[3] ≈ +6380.34453270 atol=1e-7

    @test v_pef[1] ≈ -3.2256327470  atol=1e-9
    @test v_pef[2] ≈ -2.8724425110  atol=1e-9
    @test v_pef[3] ≈ +5.5319312880  atol=1e-9
end

# Functions rTODtoMOD_fk5 and rMODtoTOD_fk5
# -----------------------------------------

################################################################################
#                                 Test Results
################################################################################
#
# Scenario 01
# ===========
#
# Example 3-15: Performing IAU-76/FK5 reduction.
#
# According to this example and Table 3-6, using:
#
#   JD_TT = 2453101.828154745
#   δΔϵ   = -0.003875"
#   δΔψ   = -0.052195"
#   r_tod = 5094.51620300   i + 6127.36527840   j + 6380.34453270   k [km]
#   v_tod =   -4.7460883850 i +    0.7860783240 j +    5.5319312880 k [km/s]
#
# one gets the following data:
#
#   r_mod = 5094.02837450   i + 6127.87081640   j + 6380.24851640   k [km]
#   v_mod =   -4.7462630520 i +    0.7860140450 j +    5.5317905620 k [km/s]
#
# Furthermore, using:
#
#   JD_TT = 2453101.828154745
#   δΔϵ   = 0"
#   δΔψ   = 0"
#   r_tod = 5094.51478040   i + 6127.36646120   j + 6380.34453270   k [km]
#   v_tod =   -4.7460885670 i +    0.7860772220 j +    5.5319312880 k [km/s]
#
# one gets the following data:
#
#   r_mod = 5094.02901670   i + 6127.87093630   j + 6380.24788850   k [km]
#   v_mod =   -4.7462624950 i +    0.7860141490 j +    5.5317910250 k [km/s]
#
################################################################################

@testset "Functions rTODtoMOD_fk5 and rMODtoTOD_fk5" begin
    ## rTODtoMOD_fk5
    ## =============

    ## -------------------------------------------------------------------------
    ##                                 First Test
    ## -------------------------------------------------------------------------

    r_tod = [5094.51620300; 6127.36527840; 6380.34453270]
    v_tod = [-4.7460883850; 0.7860783240; 5.5319312880]

    ## DCM
    ## ---

    D_MOD_TOD = rTODtoMOD_fk5(2453101.828154745,
                              -0.003875*pi/(180*3600),
                              -0.052195*pi/(180*3600))
    r_mod = D_MOD_TOD*r_tod
    v_mod = D_MOD_TOD*v_tod

    @test r_mod[1] ≈ +5094.02837450 atol=1e-7
    @test r_mod[2] ≈ +6127.87081640 atol=1e-7
    @test r_mod[3] ≈ +6380.24851640 atol=1e-7

    @test v_mod[1] ≈ -4.7462630520  atol=1e-9
    @test v_mod[2] ≈ +0.7860140450  atol=1e-9
    @test v_mod[3] ≈ +5.5317905620  atol=1e-9

    ## Quaternion
    ## ----------

    q_MOD_TOD = rTODtoMOD_fk5(Quaternion,
                              2453101.828154745,
                              -0.003875*pi/(180*3600),
                              -0.052195*pi/(180*3600))

    r_mod = vect(conj(q_MOD_TOD)*r_tod*q_MOD_TOD)
    v_mod = vect(conj(q_MOD_TOD)*v_tod*q_MOD_TOD)

    @test r_mod[1] ≈ +5094.02837450 atol=1e-7
    @test r_mod[2] ≈ +6127.87081640 atol=1e-7
    @test r_mod[3] ≈ +6380.24851640 atol=1e-7

    @test v_mod[1] ≈ -4.7462630520  atol=1e-9
    @test v_mod[2] ≈ +0.7860140450  atol=1e-9
    @test v_mod[3] ≈ +5.5317905620  atol=1e-9

    ## -------------------------------------------------------------------------
    ##                                Second Test
    ## -------------------------------------------------------------------------

    r_tod = [5094.51478040; 6127.36646120; 6380.34453270]
    v_tod = [-4.7460885670; 0.7860772220; 5.5319312880]

    ## DCM
    ## ---

    D_MOD_TOD = rTODtoMOD_fk5(2453101.82815474)

    r_mod = D_MOD_TOD*r_tod
    v_mod = D_MOD_TOD*v_tod

    @test r_mod[1] ≈ +5094.02901670 atol=1e-7
    @test r_mod[2] ≈ +6127.87093630 atol=1e-7
    @test r_mod[3] ≈ +6380.24788850 atol=1e-7

    @test v_mod[1] ≈ -4.7462624950  atol=1e-9
    @test v_mod[2] ≈ +0.7860141490  atol=1e-9
    @test v_mod[3] ≈ +5.5317910250  atol=1e-9

    ## Quaternion
    ## ----------

    q_MOD_TOD = rTODtoMOD_fk5(Quaternion, 2453101.828154745)

    r_mod = vect(conj(q_MOD_TOD)*r_tod*q_MOD_TOD)
    v_mod = vect(conj(q_MOD_TOD)*v_tod*q_MOD_TOD)

    @test r_mod[1] ≈ +5094.02901670 atol=1e-7
    @test r_mod[2] ≈ +6127.87093630 atol=1e-7
    @test r_mod[3] ≈ +6380.24788850 atol=1e-7

    @test v_mod[1] ≈ -4.7462624950  atol=1e-9
    @test v_mod[2] ≈ +0.7860141490  atol=1e-9
    @test v_mod[3] ≈ +5.5317910250  atol=1e-9

    ## rMODtoTOD
    ## =========

    ## -------------------------------------------------------------------------
    ##                                 First Test
    ## -------------------------------------------------------------------------

    r_mod = [5094.02837450; 6127.87081640; 6380.24851640]
    v_mod = [-4.7462630520; 0.7860140450; 5.5317905620]

    ## DCM
    ## ---

    D_TOD_MOD = rMODtoTOD_fk5(2453101.828154745,
                              -0.003875*pi/(180*3600),
                              -0.052195*pi/(180*3600))
    r_tod = D_TOD_MOD*r_mod
    v_tod = D_TOD_MOD*v_mod

    @test r_tod[1] ≈ +5094.51620300 atol=1e-7
    @test r_tod[2] ≈ +6127.36527840 atol=1e-7
    @test r_tod[3] ≈ +6380.34453270 atol=1e-7

    @test v_tod[1] ≈ -4.7460883850  atol=1e-9
    @test v_tod[2] ≈ +0.7860783240  atol=1e-9
    @test v_tod[3] ≈ +5.5319312880  atol=1e-9

    ## Quaternion
    ## ----------

    q_TOD_MOD = rMODtoTOD_fk5(Quaternion,
                              2453101.828154745,
                              -0.003875*pi/(180*3600),
                              -0.052195*pi/(180*3600))

    r_tod = vect(conj(q_TOD_MOD)*r_mod*q_TOD_MOD)
    v_tod = vect(conj(q_TOD_MOD)*v_mod*q_TOD_MOD)

    @test r_tod[1] ≈ +5094.51620300 atol=1e-7
    @test r_tod[2] ≈ +6127.36527840 atol=1e-7
    @test r_tod[3] ≈ +6380.34453270 atol=1e-7

    @test v_tod[1] ≈ -4.7460883850  atol=1e-9
    @test v_tod[2] ≈ +0.7860783240  atol=1e-9
    @test v_tod[3] ≈ +5.5319312880  atol=1e-9

    ## -------------------------------------------------------------------------
    ##                                Second Test
    ## -------------------------------------------------------------------------

    r_mod = [5094.02901670; 6127.87093630; 6380.24788850]
    v_mod = [-4.7462624950; 0.7860141490; 5.5317910250]

    ## DCM
    ## ---

    D_TOD_MOD = rMODtoTOD_fk5(2453101.828154745)

    r_tod = D_TOD_MOD*r_mod
    v_tod = D_TOD_MOD*v_mod

    @test r_tod[1] ≈ +5094.51478040 atol=1e-7
    @test r_tod[2] ≈ +6127.36646120 atol=1e-7
    @test r_tod[3] ≈ +6380.34453270 atol=1e-7

    @test v_tod[1] ≈ -4.7460885670  atol=1e-9
    @test v_tod[2] ≈ +0.7860772220  atol=1e-9
    @test v_tod[3] ≈ +5.5319312880  atol=1e-9

    ## Quaternion
    ## ----------

    q_TOD_MOD = rMODtoTOD_fk5(Quaternion, 2453101.828154745)

    r_tod = vect(conj(q_TOD_MOD)*r_mod*q_TOD_MOD)
    v_tod = vect(conj(q_TOD_MOD)*v_mod*q_TOD_MOD)

    @test r_tod[1] ≈ +5094.51478040 atol=1e-7
    @test r_tod[2] ≈ +6127.36646120 atol=1e-7
    @test r_tod[3] ≈ +6380.34453270 atol=1e-7

    @test v_tod[1] ≈ -4.7460885670  atol=1e-9
    @test v_tod[2] ≈ +0.7860772220  atol=1e-9
    @test v_tod[3] ≈ +5.5319312880  atol=1e-9
end

# Functions rMODtoGCRF_fk5 and rGCRFtoMOD_fk5
# -------------------------------------------

################################################################################
#                                 Test Results
################################################################################
#
# Scenario 01
# ===========
#
# Example 3-15: Performing IAU-76/FK5 reduction.
#
# According to this example and Table 3-6, using:
#
#   JD_TT = 2453101.828154745
#   r_mod = 5094.02837450   i + 6127.87081640   j + 6380.24851640   k [km]
#   v_mod =   -4.7462630520 i +    0.7860140450 j +    5.5317905620 k [km/s]
#
# one gets:
#
#   r_gcrf = 5102.50895790  i + 6123.01140070   j + 6378.13692820   k [km]
#   v_gcrf =  -4.7432201570 i + 0.7905364970    j + 5.5337557270    k [km/s]
#
# Furthermore, using:
#
#   JD_TT = 2453101.828154745
#   r_mod = 5094.02901670   i + 6127.87093630   j + 6380.24788850   k [km]
#   v_mod =   -4.7462624950 i +    0.7860141490 j +    5.5317910250 k [km/s]
#
# one gets:
#
#   r_j2000 = 5102.50960000  i + 6123.01152000   j + 6378.13630000   k [km]
#   v_j2000 =  -4.7432196000 i + 0.7905366000    j + 5.5337561900    k [km/s]
#
################################################################################

@testset "Functions rMODtoGCRF_fk5 and rGCRFtoMOD_fk5" begin
    ## rMODtoGCRF
    ## ==========

    ## -------------------------------------------------------------------------
    ##                                 First Test
    ## -------------------------------------------------------------------------

    r_mod = [5094.02837450; 6127.87081640; 6380.24851640]
    v_mod = [-4.7462630520; 0.7860140450; 5.5317905620]

    ## DCM
    ## ---

    D_GCRF_MOD = rMODtoGCRF_fk5(2453101.828154745)

    r_gcrf = D_GCRF_MOD*r_mod
    v_gcrf = D_GCRF_MOD*v_mod

    @test r_gcrf[1] ≈ +5102.50895790 atol=1e-7
    @test r_gcrf[2] ≈ +6123.01140070 atol=1e-7
    @test r_gcrf[3] ≈ +6378.13692820 atol=1e-7

    @test v_gcrf[1] ≈ -4.7432201570  atol=1e-9
    @test v_gcrf[2] ≈ +0.7905364970  atol=1e-9
    @test v_gcrf[3] ≈ +5.5337557270  atol=1e-9

    ## Quaternion
    ## ----------

    q_GCRF_MOD = rMODtoGCRF_fk5(Quaternion, 2453101.828154745)

    r_gcrf = vect(conj(q_GCRF_MOD)*r_mod*q_GCRF_MOD)
    v_gcrf = vect(conj(q_GCRF_MOD)*v_mod*q_GCRF_MOD)

    @test r_gcrf[1] ≈ +5102.50895790 atol=1e-7
    @test r_gcrf[2] ≈ +6123.01140070 atol=1e-7
    @test r_gcrf[3] ≈ +6378.13692820 atol=1e-7

    @test v_gcrf[1] ≈ -4.7432201570  atol=1e-9
    @test v_gcrf[2] ≈ +0.7905364970  atol=1e-9
    @test v_gcrf[3] ≈ +5.5337557270  atol=1e-9

    ## -------------------------------------------------------------------------
    ##                                Second Test
    ## -------------------------------------------------------------------------

    r_mod = [5094.02901670; 6127.87093630; 6380.24788850]
    v_mod = [-4.7462624950; 0.7860141490; 5.5317910250]

    ## DCM
    ## ---

    D_J2000_MOD = rMODtoGCRF_fk5(2453101.828154745)

    r_j2000 = D_J2000_MOD*r_mod
    v_j2000 = D_J2000_MOD*v_mod

    @test r_j2000[1] ≈ +5102.50960000 atol=1e-7
    @test r_j2000[2] ≈ +6123.01152000 atol=1e-7
    @test r_j2000[3] ≈ +6378.13630000 atol=1e-7

    @test v_j2000[1] ≈ -4.7432196000  atol=1e-9
    @test v_j2000[2] ≈ +0.7905366000  atol=1e-9
    @test v_j2000[3] ≈ +5.5337561900  atol=1e-9

    ## Quaternion
    ## ----------

    q_J2000_MOD = rMODtoGCRF_fk5(Quaternion, 2453101.828154745)

    r_j2000 = vect(conj(q_J2000_MOD)*r_mod*q_J2000_MOD)
    v_j2000 = vect(conj(q_J2000_MOD)*v_mod*q_J2000_MOD)

    @test r_j2000[1] ≈ +5102.50960000 atol=1e-7
    @test r_j2000[2] ≈ +6123.01152000 atol=1e-7
    @test r_j2000[3] ≈ +6378.13630000 atol=1e-7

    @test v_j2000[1] ≈ -4.7432196000  atol=1e-9
    @test v_j2000[2] ≈ +0.7905366000  atol=1e-9
    @test v_j2000[3] ≈ +5.5337561900  atol=1e-9

    ## rGCRFtoMOD
    ## ==========

    ## -------------------------------------------------------------------------
    ##                                 First Test
    ## -------------------------------------------------------------------------

    r_gcrf = [5102.50895790; 6123.01140070; 6378.13692820]
    v_gcrf = [-4.7432201570; 0.7905364970; 5.5337557270]

    ## DCM
    ## ---

    D_MOD_GCRF = rGCRFtoMOD_fk5(2453101.828154745)

    r_mod = D_MOD_GCRF*r_gcrf
    v_mod = D_MOD_GCRF*v_gcrf

    @test r_mod[1] ≈ +5094.02837450 atol=1e-7
    @test r_mod[2] ≈ +6127.87081640 atol=1e-7
    @test r_mod[3] ≈ +6380.24851640 atol=1e-7

    @test v_mod[1] ≈ -4.7462630520  atol=1e-9
    @test v_mod[2] ≈ +0.7860140450  atol=1e-9
    @test v_mod[3] ≈ +5.5317905620  atol=1e-9

    ## Quaternion
    ## ----------

    q_MOD_GCRF = rGCRFtoMOD_fk5(Quaternion, 2453101.828154745)

    r_mod = vect(conj(q_MOD_GCRF)*r_gcrf*q_MOD_GCRF)
    v_mod = vect(conj(q_MOD_GCRF)*v_gcrf*q_MOD_GCRF)

    @test r_mod[1] ≈ +5094.02837450 atol=1e-7
    @test r_mod[2] ≈ +6127.87081640 atol=1e-7
    @test r_mod[3] ≈ +6380.24851640 atol=1e-7

    @test v_mod[1] ≈ -4.7462630520  atol=1e-9
    @test v_mod[2] ≈ +0.7860140450  atol=1e-9
    @test v_mod[3] ≈ +5.5317905620  atol=1e-9

    ## -------------------------------------------------------------------------
    ##                                Second Test
    ## -------------------------------------------------------------------------

    r_j2000 = [5102.50960000; 6123.01152000; 6378.13630000]
    v_j2000 = [-4.7432196000; 0.7905366000; 5.5337561900]

    ## DCM
    ## ---

    D_MOD_J2000 = rGCRFtoMOD_fk5(2453101.828154745)

    r_mod = D_MOD_J2000*r_j2000
    v_mod = D_MOD_J2000*v_j2000

    @test r_mod[1] ≈ +5094.02901670 atol=1e-7
    @test r_mod[2] ≈ +6127.87093630 atol=1e-7
    @test r_mod[3] ≈ +6380.24788850 atol=1e-7

    @test v_mod[1] ≈ -4.7462624950  atol=1e-9
    @test v_mod[2] ≈ +0.7860141490  atol=1e-9
    @test v_mod[3] ≈ +5.5317910250  atol=1e-9

    ## Quaternion
    ## ----------

    q_J2000_MOD = rGCRFtoMOD_fk5(Quaternion, 2453101.828154745)

    r_mod = vect(conj(q_J2000_MOD)*r_j2000*q_J2000_MOD)
    v_mod = vect(conj(q_J2000_MOD)*v_j2000*q_J2000_MOD)

    @test r_mod[1] ≈ +5094.02901670 atol=1e-7
    @test r_mod[2] ≈ +6127.87093630 atol=1e-7
    @test r_mod[3] ≈ +6380.24788850 atol=1e-7

    @test v_mod[1] ≈ -4.7462624950  atol=1e-9
    @test v_mod[2] ≈ +0.7860141490  atol=1e-9
    @test v_mod[3] ≈ +5.5317910250  atol=1e-9
end

# Functions rITRFtoGCRF_fk5 and rGCRFtoITRF_fk5
# ---------------------------------------------

################################################################################
#                                 Test Results
################################################################################
#
# Scenario 01
# ===========
#
# Example 3-15: Performing IAU-76/FK5 reduction.
#
# According to this example and Table 3-6, using:
#
#   JD_UT1 = DatetoJD(2004,4,6,7,51,28.386009) - 0.4399619/86400
#   JD_TT  = 2453101.828154745
#   δΔϵ    = -0.003875"
#   δΔψ    = -0.052195"
#   x_p    = -0.140682"
#   y_p    = +0.333309"
#   r_itrf = -1033.4793830    i + 7901.2952754    j + 6380.3565958    k [km]
#
# one gets the following data:
#
#   r_gcrf = 5102.50895790  i + 6123.01140070   j + 6378.13692820   k [km]
#
# Furthermore, using:
#
#   JD_UT1 = DatetoJD(2004,4,6,7,51,28.386009) - 0.4399619/86400
#   JD_TT  = 2453101.828154745
#   x_p    = -0.140682"
#   y_p    = +0.333309"
#   r_itrf = -1033.4793830    i + 7901.2952754    j + 6380.3565958    k [km]
#
# one gets the following data:
#
#   r_j2000 = 5102.50960000  i + 6123.01152000   j + 6378.13630000   k [km]
#
################################################################################

@testset "Functions rITRFtoGCRF_fk5 and rGCRFtoITRF_fk5" begin
    JD_UT1   = DatetoJD(2004,4,6,7,51,28.386009) - 0.4399619/86400
    JD_TT    = 2453101.828154745
    δΔϵ_1980 = -0.003875*pi/(180*3600)
    δΔψ_1980 = -0.052195*pi/(180*3600)
    x_p      = -0.140682*pi/(180*3600)
    y_p      = +0.333309*pi/(180*3600)

    ## rITRFtoGCRF_fk5
    ## ===============

    r_itrf = [-1033.4793830; 7901.2952754; 6380.3565958]

    ## -------------------------------------------------------------------------
    ##                                 First Test
    ## -------------------------------------------------------------------------

    ## DCM
    ## ---

    D_GCRF_ITRF = rITRFtoGCRF_fk5(JD_UT1, JD_TT, x_p, y_p, δΔϵ_1980, δΔψ_1980)

    r_gcrf = D_GCRF_ITRF*r_itrf

    @test r_gcrf[1] ≈ +5102.50895790 atol=1e-7
    @test r_gcrf[2] ≈ +6123.01140070 atol=1e-7
    @test r_gcrf[3] ≈ +6378.13692820 atol=1e-7

    ## Quaternion
    ## ----------

    q_GCRF_ITRF = rITRFtoGCRF_fk5(Quaternion,
                                  JD_UT1,
                                  JD_TT,
                                  x_p,
                                  y_p,
                                  δΔϵ_1980,
                                  δΔψ_1980)

    r_gcrf = vect(conj(q_GCRF_ITRF)*r_itrf*q_GCRF_ITRF)

    @test r_gcrf[1] ≈ +5102.50895790 atol=1e-7
    @test r_gcrf[2] ≈ +6123.01140070 atol=1e-7
    @test r_gcrf[3] ≈ +6378.13692820 atol=1e-7

    ## -------------------------------------------------------------------------
    ##                                Second Test
    ## -------------------------------------------------------------------------

    ## DCM
    ## ---

    D_GCRF_ITRF = rITRFtoGCRF_fk5(JD_UT1, JD_TT, x_p, y_p)

    r_j2000 = D_GCRF_ITRF*r_itrf

    @test r_j2000[1] ≈ +5102.50960000 atol=1e-7
    @test r_j2000[2] ≈ +6123.01152000 atol=1e-7
    @test r_j2000[3] ≈ +6378.13630000 atol=1e-7

    ## Quaternion
    ## ----------

    q_J2000_ITRF = rITRFtoGCRF_fk5(Quaternion, JD_UT1, JD_TT, x_p, y_p)

    r_j2000 = vect(conj(q_J2000_ITRF)*r_itrf*q_J2000_ITRF)

    @test r_j2000[1] ≈ +5102.50960000 atol=1e-7
    @test r_j2000[2] ≈ +6123.01152000 atol=1e-7
    @test r_j2000[3] ≈ +6378.13630000 atol=1e-7

    # rGCRFtoITRF_fk5
    # ===============

    ## -------------------------------------------------------------------------
    ##                                 First Test
    ## -------------------------------------------------------------------------

    r_gcrf = [5102.50895790; 6123.01140070; 6378.13692820]

    ## DCM
    ## ---

    D_ITRF_GCRF = rGCRFtoITRF_fk5(JD_UT1, JD_TT, x_p, y_p, δΔϵ_1980, δΔψ_1980)

    r_itrf = D_ITRF_GCRF*r_gcrf

    @test r_itrf[1] ≈ -1033.4793830 atol=1e-7
    @test r_itrf[2] ≈ +7901.2952754 atol=1e-7
    @test r_itrf[3] ≈ +6380.3565958 atol=1e-7

    ## Quaternion
    ## ----------

    q_ITRF_GCRF = rGCRFtoITRF_fk5(Quaternion,
                                  JD_UT1,
                                  JD_TT,
                                  x_p,
                                  y_p,
                                  δΔϵ_1980,
                                  δΔψ_1980)

    r_itrf = vect(conj(q_ITRF_GCRF)*r_gcrf*q_ITRF_GCRF)

    @test r_itrf[1] ≈ -1033.4793830 atol=1e-7
    @test r_itrf[2] ≈ +7901.2952754 atol=1e-7
    @test r_itrf[3] ≈ +6380.3565958 atol=1e-7

    ## -------------------------------------------------------------------------
    ##                                Second Test
    ## -------------------------------------------------------------------------

    r_j2000 = [5102.50960000; 6123.01152000; 6378.13630000]

    ## DCM
    ## ---

    D_ITRF_J2000 = rGCRFtoITRF_fk5(JD_UT1, JD_TT, x_p, y_p)

    r_itrf = D_ITRF_J2000*r_j2000

    @test r_itrf[1] ≈ -1033.4793830 atol=1e-7
    @test r_itrf[2] ≈ +7901.2952754 atol=1e-7
    @test r_itrf[3] ≈ +6380.3565958 atol=1e-7

    ## Quaternion
    ## ----------

    q_ITRF_J2000 = rGCRFtoITRF_fk5(Quaternion, JD_UT1, JD_TT, x_p, y_p)

    r_itrf = vect(conj(q_ITRF_J2000)*r_j2000*q_ITRF_J2000)

    @test r_itrf[1] ≈ -1033.4793830 atol=1e-7
    @test r_itrf[2] ≈ +7901.2952754 atol=1e-7
    @test r_itrf[3] ≈ +6380.3565958 atol=1e-7
end

# Functions rPEFtoMOD_fk5 and rMODtoPEF_fk5
# -----------------------------------------

################################################################################
#                                 Test Results
################################################################################
#
# Scenario 01
# ===========
#
# Example 3-15: Performing IAU-76/FK5 reduction.
#
# According to this example and Table 3-6, using:
#
#   JD_TT = 2453101.828154745
#   δΔϵ   = -0.003875"
#   δΔψ   = -0.052195"
#   r_pef  = -1033.47503130   i + 7901.30558560   j + 6380.34453270   k [km]
#
# one gets the following data:
#
#   r_mod = 5094.02837450   i + 6127.87081640   j + 6380.24851640   k [km]
#
# Furthermore, using:
#
#   JD_TT = 2453101.828154745
#   δΔϵ   = 0"
#   δΔψ   = 0"
#   r_pef  = -1033.47503130   i + 7901.30558560   j + 6380.34453270   k [km]
#
# one gets the following data:
#
#   r_mod = 5094.02901670   i + 6127.87093630   j + 6380.24788850   k [km]
#
# NOTE: The velocity cannot be easily tested because it is necessary to account
# for the Earth rotation when converting from/to PEF.
#
################################################################################

@testset "Functions rPEFtoMOD_fk5 and rMODtoPEF_fk5" begin
    JD_UT1 = DatetoJD(2004,4,6,7,51,28.386009) - 0.4399619/86400

    ## rPEFtoMOD_fk5
    ## =============

    r_pef  = [-1033.47503130; 7901.30558560; 6380.34453270]

    ## -------------------------------------------------------------------------
    ##                                 First Test
    ## -------------------------------------------------------------------------

    ## DCM
    ## ---

    D_MOD_PEF = rPEFtoMOD_fk5(JD_UT1,
                              2453101.828154745,
                              -0.003875*pi/(180*3600),
                              -0.052195*pi/(180*3600))

    r_mod = D_MOD_PEF*r_pef

    @test r_mod[1] ≈ +5094.02837450 atol=1e-7
    @test r_mod[2] ≈ +6127.87081640 atol=1e-7
    @test r_mod[3] ≈ +6380.24851640 atol=1e-7

    ## Quaternion
    ## ----------

    q_MOD_PEF = rPEFtoMOD_fk5(Quaternion,
                              JD_UT1,
                              2453101.828154745,
                              -0.003875*pi/(180*3600),
                              -0.052195*pi/(180*3600))

    r_mod = vect(conj(q_MOD_PEF)*r_pef*q_MOD_PEF)

    @test r_mod[1] ≈ +5094.02837450 atol=1e-7
    @test r_mod[2] ≈ +6127.87081640 atol=1e-7
    @test r_mod[3] ≈ +6380.24851640 atol=1e-7

    ## -------------------------------------------------------------------------
    ##                                Second Test
    ## -------------------------------------------------------------------------

    ## DCM
    ## ---

    D_MOD_PEF = rPEFtoMOD_fk5(JD_UT1, 2453101.828154745)

    r_mod = D_MOD_PEF*r_pef

    @test r_mod[1] ≈ +5094.02901670 atol=1e-7
    @test r_mod[2] ≈ +6127.87093630 atol=1e-7
    @test r_mod[3] ≈ +6380.24788850 atol=1e-7

    ## Quaternion
    ## ----------

    q_MOD_PEF = rPEFtoMOD_fk5(Quaternion, JD_UT1, 2453101.828154745)

    r_mod = vect(conj(q_MOD_PEF)*r_pef*q_MOD_PEF)

    @test r_mod[1] ≈ +5094.02901670 atol=1e-7
    @test r_mod[2] ≈ +6127.87093630 atol=1e-7
    @test r_mod[3] ≈ +6380.24788850 atol=1e-7

    ## rMODtoPEF_fk5
    ## =============

    ## -------------------------------------------------------------------------
    ##                                 First Test
    ## -------------------------------------------------------------------------

    r_mod = [5094.02837450; 6127.87081640; 6380.24851640]

    ## DCM
    ## ---

    D_PEF_MOD = rMODtoPEF_fk5(JD_UT1,
                              2453101.828154745,
                              -0.003875*pi/(180*3600),
                              -0.052195*pi/(180*3600))

    r_pef = D_PEF_MOD*r_mod

    @test r_pef[1] ≈ -1033.47503130 atol=1e-7
    @test r_pef[2] ≈ +7901.30558560 atol=1e-7
    @test r_pef[3] ≈ +6380.34453270 atol=1e-7

    ## Quaternion
    ## ----------

    q_PEF_MOD = rMODtoPEF_fk5(Quaternion,
                              JD_UT1,
                              2453101.828154745,
                              -0.003875*pi/(180*3600),
                              -0.052195*pi/(180*3600))

    r_pef = vect(conj(q_PEF_MOD)*r_mod*q_PEF_MOD)

    @test r_pef[1] ≈ -1033.47503130 atol=1e-7
    @test r_pef[2] ≈ +7901.30558560 atol=1e-7
    @test r_pef[3] ≈ +6380.34453270 atol=1e-7

    ## -------------------------------------------------------------------------
    ##                                Second Test
    ## -------------------------------------------------------------------------

    r_mod = [5094.02901670; 6127.87093630; 6380.24788850]

    ## DCM
    ## ---

    D_PEF_MOD = rMODtoPEF_fk5(JD_UT1, 2453101.828154745)

    r_pef = D_PEF_MOD*r_mod

    @test r_pef[1] ≈ -1033.47503130 atol=1e-7
    @test r_pef[2] ≈ +7901.30558560 atol=1e-7
    @test r_pef[3] ≈ +6380.34453270 atol=1e-7

    ## Quaternion
    ## ----------

    q_PEF_MOD = rMODtoPEF_fk5(Quaternion, JD_UT1, 2453101.828154745)

    r_pef = vect(conj(q_PEF_MOD)*r_mod*q_PEF_MOD)

    @test r_pef[1] ≈ -1033.47503130 atol=1e-7
    @test r_pef[2] ≈ +7901.30558560 atol=1e-7
    @test r_pef[3] ≈ +6380.34453270 atol=1e-7
end

ext: jl


text:
# Implementation of Blendenpik with Gaussian row mixing for the solution of least squares
# problem ||Ax - b||₂ where A has full column rank.
#
# This method for other row mixing strategies is described in
#
# Avron, Haim, Petar Maymounkov, and Sivan Toledo. "Blendenpik: Supercharging LAPACK's
# least-squares solver." SIAM Journal on Scientific Computing 32, no. 3 (2010): 1217-1236.
#
# January 2021

"""
    blendenpick_gauss(A, b; r)

Solves the least squares problem with coefficient `A` and constant `b`, where `A` has full
column rank, using the blendenpick method with Gaussian row mixing. The number of sampled
rows, `r`, defaults to the number of columns.
"""
function blendenpick_gauss!(
    x::Vector{T},
    A::Matrix{T},  # Coefficient matrix of system
    b::Vector{T};  # Constant vector of system
    r::Int=size(A, 2) + 0,  # Size of row sample
    verbose::Bool=false  #Show stats from lsqr solver
) where T <: Real
    m = size(A, 1)              # Number of rows in A

    # Mix rows of a A with Gaussians to generate r by size(A, 2) matrix
    A_mixed = randn(r, m) * A

    # Generate preconditioner using R⁻ factor of qr decomposition of mixed matrix
    _, R = qr(A_mixed)
    Rinv = R \ I

    # Run lsqr on transformed systems
    y, stats = lsqr(A * Rinv, b)

    verbose && show(stats)

    # Recover and return solution to original system
    x = Rinv * y
end


text:
using Flux

"""
Helper function to train a NN in a loop with a maximum number of iterations.
    Breaks early if a stall is detected
    Returns a boolean for convergence to the given tolerance
"""
function train_in_loop(ps, tol, loss;
                rate=1e-3, stall_tol=nothing,
                dat_size=50, max_iter=100,
                plot_func=nothing, plot_truth=nothing)
    if stall_tol == nothing
        stall_tol = tol / 100
    end
    dat_dummy = Iterators.repeated((), dat_size)
    opt = ADAM(rate)
    iter = 1
    current_loss = loss() + 100 # So it doesn't stop early
    is_converged = true
    while current_loss > tol
        prev_loss = current_loss
        current_loss = loss()
        if prev_loss - current_loss < stall_tol
            println("Stall detected at iteration $iter; aborting")
            is_converged = false
            break
        end

        Flux.train!(loss, ps, dat_dummy, opt, cb = ()->@show loss())
        iter += 1

        if plot_func !== nothing
            pt = plot(plot_func(), label="NN")
            if plot_truth !== nothing
                plot!(plot_truth, label="True control")
            end
            title!("Iteration $iter")
            display(pt)
        end
        if iter > max_iter
            println("Reached maximum iterations; aborting")
            is_converged = false
            break
        end
    end

    return is_converged
end

export train_in_loop


text:
__precompile__()

module ExpSim

using Sim
using MuJoCo, Common, GLFW
using Policy: getmean!, AbstractPolicy, loadpolicy
using Flux
using mjWrap #: load_model, mjSet
using UnsafeArrays

using JLD2, FileIO
using ExpFunctions
using Printf

export simulate

@enum RUN_MODE runmodel=1 rundata=2 runpolicy=3
const runmodes = [runmodel, rundata, runpolicy]    # HAAAAAAAAAAAACK
mutable struct ExpDisplay
   T::Int
   t::Int
   K::Int
   k::Int
   nqnv::Int
   states::Array{Float64, 3}

   pol #::AbstractPolicy
   using_policy::Bool
   skip::Int

   mode::RUN_MODE

   mjsys::mjWrap.mjSet
   myfuncs #::ExpFunctions.FunctionSet

   function ExpDisplay(mjsys, states, pol, skip)
      nqnv, T, K = size(states)
      return new(T, 1, K, 1, nqnv, states,
                 pol, true, skip, runmodel, mjsys, nothing)
   end
end

function traj2state(exd::ExpDisplay)
   t, k = exd.t, exd.k
   m, d = exd.mjsys.m, exd.mjsys.d
   nq = m.m[].nq
   d.qpos .= exd.states[1:nq, t, k]
   d.qvel .= exd.states[(nq+1):end, t, k]
   mj_forward(m, d)
end

function file2state(s::Simulation, exd::ExpDisplay)
   #copy data from traj struct, keeping track of indicies
   if mod1(s.framenum, exd.skip) == exd.skip && s.lastframenum != s.framenum
      if (exd.t >= exd.T)
         exd.t = 1  #start of new trajectory
         exd.k += 1
         if exd.k > exd.K exd.k = 1 end
         s.d.d[].time = 0.0
         @printf("Rendering Trajectory %d\n", exd.k)
      end
      traj2state(exd)
      exd.t += 1 # advance our saved trajectory pointer
      s.lastframenum = s.framenum
   end
end

function applypolicy(s::Simulation, exd::ExpDisplay)
   d = exd.mjsys.d
   if mod1(s.framenum, exd.skip) == exd.skip && s.lastframenum != s.framenum
      # Needs to eval the function loaded at run time
      obs = zeros(exd.mjsys.ns)
      eval( :($(exd.myfuncs.observe!)($(exd.mjsys), [$(d.qpos); $(d.qvel)], $(obs))) )
      c = uview(d.ctrl)
      o = uview(obs)
      getmean!(c, exd.pol, o)
      # getmean!(d.ctrl, exd.pol, obs)
      s.lastframenum = s.framenum
   end
end

function my_step(s::Simulation, exd::Union{ExpDisplay,Nothing})
   m = s.m
   d = s.d
   if exd == nothing
      mj_step(m, d)
   else
      if exd.mode == rundata
         file2state(s, exd)
         d.d[].time += m.m[].opt.timestep*exd.skip
      elseif exd.mode == runpolicy && exd.using_policy && exd.pol != nothing
         applypolicy(s, exd)
         mj_step(m, d)
      else # model mode / passive policy mode
         d.ctrl .= 0.0
         mj_step(m, d)
      end
   end
   s.framenum += 1
end

function simulation(s::Simulation, exd::Union{ExpDisplay,Nothing})
   # println("simulation")
   d = s.d
   m = s.m
   if s.paused
      if s.pert[].active > 0
         mjv_applyPerturbPose(m, d, s.pert, 1)  # move mocap and dynamic bodies
         mj_forward(m, d)
      end
   else
      #slow motion factor: 10x
      factor = (s.slowmotion ? 10 : 1)

      # advance effective simulation time by 1/refreshrate
      startsimtm = d.d[].time
      while ((d.d[].time - startsimtm) * factor < (1.0 / s.refreshrate))
         # clear old perturbations, apply new
         #mju_zero(d->xfrc_applied, 6 * m->nbody);
         d.xfrc_applied .= 0.0
         if s.pert[].select > 0
            mjv_applyPerturbPose(m, d, s.pert, 0) # move mocap bodies only
            mjv_applyPerturbForce(m, d, s.pert)
         end

         my_step(s, exd)

         # break on reset
         if (d.d[].time < startsimtm) break end
      end
   end
end

function render(s::Simulation, exd::Union{ExpDisplay,Nothing}, w::GLFW.Window)
   wi, hi = GLFW.GetFramebufferSize(w)
   rect = mjrRect(Cint(0), Cint(0), Cint(wi), Cint(hi))
   smallrect = mjrRect(Cint(0), Cint(0), Cint(wi), Cint(hi))

   simulation(s, exd)

   # update scene
   mjv_updateScene(s.m, s.d,
                   s.vopt, s.pert, s.cam, Int(mj.CAT_ALL), s.scn)
   # render
   mjr_render(rect, s.scn, s.con)

   if s.showsensor
      if (!s.paused) Sim.sensorupdate(s) end
      Sim.sensorshow(s, smallrect)
   end

   if s.showinfo
      str_slow = s.slowmotion ? "(10x slowdown)" : ""
      str_paused = (s.paused ? "\nPaused" : "\nRunning")*"\nTime:"
      time = round(s.d.d[].time, digits=3)
      if exd != nothing
         if exd.mode == runmodel
            status = "\nPassive Model\n$(time)"
         elseif exd.mode == rundata
            str_paused *= "\nTraj\nT\n"
            status = "\nData Mode\n$(time)\n$(exd.k)\n$(exd.t)"
         elseif exd.mode == runpolicy
            status = "\nPolicy Interaction\n$(time)"
         end
      else
         status = "\nPassive Model\n$(time)"
      end
      status = str_slow * status

      mjr_overlay(Int(mj.FONT_NORMAL), Int(mj.GRID_BOTTOMLEFT), rect,
                  str_paused,
                  status, s.con)
   end

   if s.record != nothing
      mjr_readPixels(s.vidbuff, C_NULL, rect, s.con);
      write(s.record[1], s.vidbuff[1:3*rect.width*rect.height]);
   end

   # Swap front and back buffers
   GLFW.SwapBuffers(w)
end

function mycustomkeyboard(s::Simulation, exd::Union{ExpDisplay,Nothing}, window::GLFW.Window,
                          key::GLFW.Key, scancode::Int32, act::GLFW.Action, mods::Int32)
   if act == GLFW.RELEASE return end

   if exd != nothing # more than just model d
      mode = exd.mode
      if key == GLFW.KEY_COMMA
         if mode==runmodel mode=runpolicy else mode=runmodes[Int(mode) - 1] end
      elseif key == GLFW.KEY_PERIOD
         if mode==runpolicy mode=runmodel else mode=runmodes[Int(mode) + 1] end
      elseif mode == runpolicy && key == GLFW.KEY_P
         exd.using_policy = !exd.using_policy
         if exd.using_policy println("Using Policy") end
      end
      exd.mode = mode

      if exd.mode == rundata
         # If holding shift, then cycle through trajectories
         if mods & GLFW.MOD_SHIFT > 0
            if key == GLFW.KEY_RIGHT
               # exd.t += 20
               # if exd.t > exd.T
               exd.k += 1; if exd.k > exd.K exd.k = 1 end
               exd.t = 1
               # end
               traj2state(exd)
            elseif key == GLFW.KEY_LEFT
               # exd.t -= 20
               # if exd.t < 1
               exd.k -= 1; if exd.k < 1 exd.k = exd.K end
               exd.t = 1
               # end
               traj2state(exd)
            end
         else  # Else step forward in time
            if s.paused
               if key == GLFW.KEY_RIGHT
                  exd.t += 1; if exd.t > exd.T exd.t = 1 end
                  traj2state(exd)
               elseif key == GLFW.KEY_LEFT
                  exd.t -= 1; if exd.t < 1 exd.t = exd.T end
                  traj2state(exd)
               end
            end
         end
      end
   end

   Sim.mykeyboard(s, window, key, scancode, act, mods)
end

function findinfile(name::String, s::String)
   open(name) do f
      while !eof(f)
         x = readline(f)
         if startswith(x,s)
            return x
         end
      end
   end
end

macro preloadexp(dir)
   return quote
      files        = readdir($dir)
      modelfile    = $dir*"/"*files[findall((x)->endswith(x, ".xml"), files)][1]
      expfile      = $dir*"/"*files[findall((x)->endswith(x, ".jl")&&occursin(x, "functions")==false, files)][1]
      functionfile = $dir*"/"*files[findall((x)->endswith(x, "functions.jl"), files)][1]
      polfile      = $dir*"/policy.jld2"
      if isfile($dir*"/mean.jld2")
         datafile  = $dir*"/mean.jld2" #TODO HACK loading mean samples for XXX
      else
         datafile  = $dir*"/data.jld2"
      end

      #@info("Loading Experiment $dir")
      #@info("Model file: $modelfile")
      #@info("Function file: $functionfile")
      #@info("Loading policy file $polfile")
      include(functionfile)

      expmt = Common.plotexpmt($dir*"/expmt.jld2") # why not
      maxiter = argmax(expmt["stocR"])

      if isfile(polfile)
         mypolicy, frameskip, _ = loadpolicy(polfile)
      else
         mypolicy = nothing
         println("Loading skip from: ", expfile)
         frameskip = eval( Meta.parse( findinfile(expfile, "myskip")))
         println(Meta.parse( findinfile(expfile, "myskip")))
         @warn("Can't load policy, or no policy to load")
      end
      @info("Frameskip: $frameskip")
      if isdefined(:ns) == false
         ns = size(load(datafile, "obs"), 1)
      end
      my_mjsys = mjw.load_model(modelfile, frameskip, "normal", ns)

      states = load(datafile, "state")
      exd = ExpDisplay( my_mjsys, load(datafile, "state"), mypolicy, frameskip)

      #exd.myfuncs = Main.myfuncs
      exd.myfuncs = myfuncs
      @info("Affecting model accoring to iteration $maxiter")
      for i=1:maxiter
         exd.myfuncs.setmodel!(exd.mjsys, i) # applies curriculum
      end
      return exd
   end
end

function loadexp(dir, width, height)
   @info(dir)
   exd = eval( :(@preloadexp($dir)) )

   s = Sim.start(exd.mjsys.m, exd.mjsys.d,  width, height)

   # Simulation struct and ExpDisplay struct share mjmodel and mjdata
   GLFW.SetKeyCallback(s.window, (w,k,sc,a,mo)->mycustomkeyboard(s,exd,w,k,sc,a,mo))
   GLFW.SetWindowRefreshCallback(s.window, (w)->render(s,exd,w))

   return s, exd
end

function loadmodel(modelfile, width, height)
   ptr_m = mj_loadXML(modelfile, C_NULL)
   ptr_d = mj_makeData(ptr_m)
   m, d = mj.mapmujoco(ptr_m, ptr_d)
   mj_forward(m, d)
   s = Sim.start(m, d, width, height)
   @info("Model file: $modelfile")

   # Simulation struct and ExpDisplay struct share mjmodel and mjdata
   GLFW.SetKeyCallback(s.window, (w,k,sc,a,mo)->mycustomkeyboard(s,nothing,w,k,sc,a,mo))
   GLFW.SetWindowRefreshCallback(s.window, (w)->render(s,nothing,w))

   return s, nothing
end

function start(args::Vector{String}, width=800, height=480)
   opt = settings(args)

   if isdir(opt["arg1"])
      s, exd = loadexp(opt["arg1"], width, height)
   elseif endswith(opt["arg1"], ".xml")
      s, exd = loadmodel(opt["arg1"], width, height)
   else
      error("Unrecognized file input.")
   end

   return s, exd
end

function simulate(s, exd::Union{ExpDisplay,Nothing})
   # Loop until the user closes the window
   Sim.alignscale(s)
   while !GLFW.WindowShouldClose(s.window)
      render(s, exd, s.window)
      GLFW.PollEvents()
   end
   GLFW.DestroyWindow(s.window)
end

function simulate(f::String, width=1280, height=800)
   if isdir(f)
      s, exd = loadexp(f, width, height)
   elseif endswith(f, ".xml")
      s, exd = loadmodel(f, width, height)
   else
      error("Unrecognized file input.")
   end
   simulate(s, exd)
end

end


text:
function mod2pi(x::DoubleFloat{T}) where {T<:IEEEFloat}
    s = signbit(x)
    if s
        x = -x
    end
    x < twopi_d64 && return x
    himdlo = mul323(inv_pi_2o1_t64, HILO(x))
    hi, md, lo = three_sum([modf(x)[1] for x in himdlo]...,)
    if hi >= 1.0
        hi = hi - 1.0
        hi, md, lo = three_sum(hi, md, lo)
    end
    himdlo = mul333(pi_2o1_t64, (hi, md, lo))
    if signbit(himdlo[1])
        himdlo = add333(pi_2o1_t64, himdlo)
    end
    if s
        himdlo = sub333(pi_2o1_t64, himdlo)
    end
    z = DoubleFloat{T}(himdlo[1],himdlo[2])
    return z
end

function modpi(x::DoubleFloat{T}) where {T<:IEEEFloat}
    s = signbit(x)
    if s
        x = -x
    end
    x < onepi_d64 && return x
    himdlo = mul323(inv_pi_1o1_t64, HILO(x))
    hi, md, lo = three_sum([modf(x)[1] for x in himdlo]...,)
    if hi >= 1.0
        hi = hi - 1.0
        hi, md, lo = three_sum(hi, md, lo)
    end
    himdlo = mul333(pi_1o1_t64, (hi, md, lo))
    if signbit(himdlo[1])
        himdlo = add333(pi_1o1_t64, himdlo)
    end
    if s
        himdlo = sub333(pi_1o1_t64, himdlo)
    end
    z = DoubleFloat{T}(himdlo[1],himdlo[2])
    return z
end

function modhalfpi(x::DoubleFloat{T}) where {T<:IEEEFloat}
    s = signbit(x)
    if s
        x = -x
    end
    x < twopi_d64 && return x
    himdlo = mul323(inv_pi_1o2_t64, HILO(x))
    hi, md, lo = three_sum([modf(x)[1] for x in himdlo]...,)
    if hi >= 1.0
        hi = hi - 1.0
        hi, md, lo = three_sum(hi, md, lo)
    end
    himdlo = mul333(pi_1o2_t64, (hi, md, lo))
    if signbit(himdlo[1])
        himdlo = add333(pi_1o2_t64, himdlo)
    end
    if s
       himdlo = sub333(pi_1o2_t64, himdlo)
    end
    z = DoubleFloat{T}(himdlo[1],himdlo[2])
    return z
end

function modqrtrpi(x::DoubleFloat{T}) where {T<:IEEEFloat}
    s = signbit(x)
    if s
        x = -x
    end
    x < twopi_d64 && return x
    himdlo = mul323(inv_pi_1o4_t64, HILO(x))
    hi, md, lo = three_sum([modf(x)[1] for x in himdlo]...,)
    if hi >= 1.0
        hi = hi - 1.0
        hi, md, lo = three_sum(hi, md, lo)
    end
    himdlo = mul333(pi_1o4_t64, (hi, md, lo))
    if signbit(himdlo[1])
        himdlo = add333(pi_1o4_t64, himdlo)
    end
    if s
        himdlo = sub333(pi_1o4_t64, himdlo)
    end
    z = DoubleFloat{T}(himdlo[1],himdlo[2])
    return z
end

function rem2pi(x::DoubleFloat{T}) where {T<:IEEEFloat}
    m = mod2pi(x)
    return HI(x) < 0 ? -m : m
end

function rempi(x::DoubleFloat{T}) where {T<:IEEEFloat}
    m = modpi(x)
    return HI(x) < 0 ? -m : m
end

function remhalfpi(x::DoubleFloat{T}) where {T<:IEEEFloat}
    m = modhalfpi(x)
    return HI(x) < 0 ? -m : m
end

function remqrtrpi(x::DoubleFloat{T}) where {T<:IEEEFloat}
    m = modqrtrpi(x)
    return HI(x) < 0 ? -m : m
end

"""
    negpi_pospi(x)

    x --> [-pi..+pi)
"""
function negpi_pospi(x::DoubleFloat{T}) where {T<:IEEEFloat}
    result = mod2pi(x)
    if result >= DoubleFloat{T}(pi)
        result = result - 2*DoubleFloat{T}(pi)
    end
    return result
end

# in -pi..+pi
# determine nearest multiple of pi/2
# within (pi/2) quadrant determine nearest multiple of pi/16

function whichquadrant(x::DoubleFloat{T}) where {T<:IEEEFloat}
    x = negpi_pospi(x)
    if signbit(x) # quadrant -2 or -1
        quadrant = -x < DoubleFloat{T}(pi)/2 ? -2 : -1
    else          # quadrant  1 or  2
        quadrant =  x < DoubleFloat{T{(pi)/2 ? 1 :  2
    end
    return quadrant
end

function which64th(x::DoubleFloat{Float64}, quadrant::Int)
    y = T(mul232(HILO(x), nv_pi_1o2_t64))
    z = trunc(Int, y*4)
    return z
end
function which64th(x::DoubleFloat{Float32}, quadrant::Int)
    y = T(mul232(HILO(x), nv_pi_1o2_t32))
    z = trunc(Int, y*4)
    return z
end
function which64th(x::DoubleFloat{Float16}, quadrant::Int)
    y = T(mul232(HILO(x), nv_pi_1o2_t16))
    z = trunc(Int, y*4)
    return z
end

# ========================================== ^^^^^ =========================================

const pi_1o1_t64_hi = pi_1o1_t64[1]
const pi_1o1_t64_md = pi_1o1_t64[2]
const pi_1o1_t64_lo = pi_1o1_t64[3]

const pi_1o4_t64_hi = pi_1o4_t64[1]
const pi_1o4_t64_md = pi_1o4_t64[2]
const pi_1o4_t64_lo = pi_1o4_t64[3]

function value_minus_pi(x::DoubleFloat{Float64})
    hi, lo = sub232(HI(x), LO(x), pi_1o1_t64_hi, pi_1o1_t64_md, pi_1o1_t64_lo)
    return DoubleFloat{Float64}(hi, lo)
end

function value_minus_qrtrpi(x::DoubleFloat{Float64})
    hi, lo = sub232(HI(x), LO(x), pi_1o4_t64_hi, pi_1o4_t64_md, pi_1o4_t64_lo)
    return DoubleFloat{Float64}(hi, lo)
end

function modpipm(x::DoubleFloat{Float64})
    abs(x) < onepi_d64 && return x
    w1 = mul323(inv_pi_1o1_t64, HILO(x))
    w2 = two_sum(w1[1] - trunc(w1[1]), w1[2], w1[3])
    y = mul322(pi_1o1_t64, w2)
    z = Double64(y)
    return z
end

modpipm(x::DoubleFloat{Float32}) = DoubleFloat{Float32}(modpipm(DoubleFloat{Float64}(x)))
modpipm(x::DoubleFloat{Float16}) = DoubleFloat{Float16}(modpipm(DoubleFloat{Float64}(x)))

function modhalfpipm(x::DoubleFloat{Float64})
    abs(x) < halfpi_d64 && return x
    w1 = mul323(inv_pi_1o2_t64, HILO(x))
    w2 = two_sum(w1[1] - trunc(w1[1]), w1[2], w1[3])
    y = mul322(pi_1o2_t64, w2)
    z = Double64(y)
    return z
end

modhalfpipm(x::DoubleFloat{Float32}) = DoubleFloat{Float32}(modhalfpipm(DoubleFloat{Float64}(x)))
modhalfpipm(x::DoubleFloat{Float16}) = DoubleFloat{Float16}(modhalfpipm(DoubleFloat{Float64}(x)))

function modqrtrpipm(x::DoubleFloat{Float64})
    abs(x) < qrtrpi_d64 && return x
    w1 = mul323(inv_pi_1o4_t64, HILO(x))
    w2 = two_sum(w1[1] - trunc(w1[1]), w1[2], w1[3])
    y = mul322(pi_1o4_t64, w2)
    z = Double64(y)
    return z
end

modqrtrpipm(x::DoubleFloat{Float32}) = DoubleFloat{Float32}(modqrtrpipm(DoubleFloat{Float64}(x)))
modqrtrpipm(x::DoubleFloat{Float16}) = DoubleFloat{Float16}(modqrtrpipm(DoubleFloat{Float64}(x)))

#=
    rem2pi(x) =  x - 2π*round(x/(2π),r)
    rem2pi(x, RoundDown) == mod2pi(x)
    •    if r == RoundNearest, then the result is in the interval [-π, π]. This will
        generally be the most accurate result.
    •    if r == RoundToZero, then the result is in the interval [0, 2π] if x is positive,.
        or [-2π, 0] otherwise.
    •    if r == RoundDown, then the result is in the interval [0, 2π].
    •    if r == RoundUp, then the result is in the interval [-2π, 0].
=#

rem2pi(x::DoubleFloat{Float64}, rounding::RoundingMode{:Down}) = mod2pi(x)
rem2pi(x::DoubleFloat{Float64}, rounding::RoundingMode{:Up}) = -rem2pi(-x, RoundDown)
rem2pi(x::DoubleFloat{Float64}, rounding::RoundingMode{:Nearest}) = modpipm(x)
rem2pi(x::DoubleFloat{Float64}, rounding::RoundingMode{:ToZero}) =
    signbit(x) ? rem2pi(x, RoundUp) : rem2pi(x, RoundDown)

rem2pi(x::DoubleFloat{Float32}, rounding::RoundingMode) =
    DoubleFloat{Float32}(rem2pi(DoubleFloat{Float64}(x), rounding))
rem2pi(x::DoubleFloat{Float16}, rounding::RoundingMode) =
    DoubleFloat{Float16}(rem2pi(DoubleFloat{Float64}(x), rounding))

rempi(x::DoubleFloat{Float64}, rounding::RoundingMode{:Down}) = modpi(x)
rempi(x::DoubleFloat{Float64}, rounding::RoundingMode{:Up}) = -rempi(-x, RoundDown)
rempi(x::DoubleFloat{Float64}, rounding::RoundingMode{:Nearest}) = modhalfpipm(x)
rempi(x::DoubleFloat{Float64}, rounding::RoundingMode{:ToZero}) =
    signbit(x) ? rempi(x, RoundUp) : rempi(x, RoundDown)

rempi(x::DoubleFloat{Float32}, rounding::RoundingMode) =
    DoubleFloat{Float32}(rempi(DoubleFloat{Float64}(x), rounding))
rempi(x::DoubleFloat{Float16}, rounding::RoundingMode) =
    DoubleFloat{Float16}(rempi(DoubleFloat{Float64}(x), rounding))

remhalfpi(x::DoubleFloat{Float64}, rounding::RoundingMode{:Down}) = modhalfpi(x)
remhalfpi(x::DoubleFloat{Float64}, rounding::RoundingMode{:Up}) = -remhalfpi(-x, RoundDown)
remhalfpi(x::DoubleFloat{Float64}, rounding::RoundingMode{:Nearest}) = modqrtrpipm(x)
remhalfpi(x::DoubleFloat{Float64}, rounding::RoundingMode{:ToZero}) =
    signbit(x) ? remhalfpi(x, RoundUp) : remhalfpi(x, RoundDown)

remhalfpi(x::DoubleFloat{Float32}, rounding::RoundingMode) =
    DoubleFloat{Float32}(remhalfpi(DoubleFloat{Float64}(x), rounding))
remhalfpi(x::DoubleFloat{Float16}, rounding::RoundingMode) =
    DoubleFloat{Float16}(remhalfpi(DoubleFloat{Float64}(x), rounding))

remqrtrpi(x::DoubleFloat{Float64}, rounding::RoundingMode{:Down}) = modqrtrpi(x)
remqrtrpi(x::DoubleFloat{Float64}, rounding::RoundingMode{:Up}) = -remqrtrpi(-x, RoundDown)
# remqrtrpi(x::DoubleFloat{Float64}, rounding::RoundingMode{:Nearest}) = modeighthpipm(x)
remqrtrpi(x::DoubleFloat{Float64}, rounding::RoundingMode{:ToZero}) =
    signbit(x) ? remqrtrpi(x, RoundUp) : remqrtrpi(x, RoundDown)

remqrtepi(x::DoubleFloat{Float32}, rounding::RoundingMode) =
    DoubleFloat{Float32}(remqrtpi(DoubleFloat{Float64}(x), rounding))
remqrtrpi(x::DoubleFloat{Float16}, rounding::RoundingMode) =
DoubleFloat{Float16}(remqrtrpi(DoubleFloat{Float64}(x), rounding))


text:
using Distributions
using Base.Test

# Test the special base where PoissonBinomial distribution reduces
# to Binomial distribution
for (p, n) in [(0.8, 6), (0.5, 10), (0.04, 20)]

    d = PoissonBinomial(fill(p, n))
    dref = Binomial(n, p)
    println("   testing PoissonBinomial p=$p, n=$n")

    @test isa(d, PoissonBinomial)
    @test minimum(d) == 0
    @test maximum(d) == n
    @test ntrials(d) == n
    @test_approx_eq entropy(d) entropy(dref)
    @test_approx_eq median(d) median(dref)
    @test_approx_eq mean(d) mean(dref)
    @test_approx_eq var(d) var(dref)
    @test_approx_eq kurtosis(d) kurtosis(dref)
    @test_approx_eq skewness(d) skewness(dref)

    for t=0:5
        @test_approx_eq mgf(d, t) mgf(dref, t)
        @test_approx_eq cf(d, t) cf(dref, t)
    end
    for i=0.1:0.1:.9
        @test_approx_eq quantile(d, i) quantile(dref, i)
    end
    for i=0:n
        @test_approx_eq_eps cdf(d, i) cdf(dref, i) 1e-15
        @test_approx_eq_eps pdf(d, i) pdf(dref, i) 1e-15
    end

end

# Test against a sum of three Binomial distributions
for (n₁, n₂, n₃, p₁, p₂, p₃) in [(10, 10, 10, 0.1, 0.5, 0.9),
                                 (1, 10, 100, 0.99, 0.1, 0.05),
                                 (5, 1, 3, 0.01, 0.99, 0.999),
                                 (10, 7, 10, 0., 0.9, 0.5)]

    n = n₁ + n₂ + n₃
    p = zeros(n)
    p[1:n₁] = p₁
    p[n₁+1: n₁ + n₂] = p₂
    p[n₁ + n₂ + 1:end] = p₃
    d = PoissonBinomial(p)
    println("   testing PoissonBinomial [$(n₁) × $(p₁), $(n₂) × $(p₂), $(n₃) × $(p₃)]")
    b1 = Binomial(n₁, p₁)
    b2 = Binomial(n₂, p₂)
    b3 = Binomial(n₃, p₃)

    pmf1 = pdf(b1)
    pmf2 = pdf(b2)
    pmf3 = pdf(b3)

    @test_approx_eq mean(d) (mean(b1) + mean(b2) + mean(b3))
    @test_approx_eq var(d) (var(b1) + var(b2) + var(b3))
    for t=0:5
        @test_approx_eq mgf(d, t) (mgf(b1, t) * mgf(b2, t) * mgf(b3, t))
        @test_approx_eq cf(d, t) (cf(b1, t) * cf(b2, t) * cf(b3, t))
    end

    for k=0:n
        m = 0.
        for i=0:min(n₁, k)
            mc = 0.
            for j=i:min(i+n₂, k)
                mc += (k - j <= n₃) && pmf2[j-i+1] * pmf3[k-j+1]
            end
            m += pmf1[i+1] * mc
        end
        @test_approx_eq_eps pdf(d, k) m 1e-15
    end
end


text:
abstract type AbstractPassive{T} <: AbstractObject{T} end

mutable struct Passive{T} <: AbstractPassive{T}
    name::String
    id::Int
    potential::Union{Symbol,T}
    temperature::Union{T,Missing}
    material::NamedTuple
    charge_density_model::AbstractChargeDensity{T}
    geometry::AbstractGeometry{T}
    geometry_positive::Vector{AbstractGeometry{T}}
    geometry_negative::Vector{AbstractGeometry{T}}

    Passive{T}() where T <: SSDFloat = new{T}()
end

function Passive{T}(dict::Dict, inputunit_dict::Dict{String,Unitful.Units}) where T <: SSDFloat
    pass = Passive{T}()
    haskey(dict, "name") ? pass.name = dict["name"] : pass.name = "external part"
    haskey(dict, "id") ? pass.id = dict["id"] : pass.id = -1
    haskey(dict,"potential") ? pass.potential = T(dict["potential"]) : pass.potential = :floating
    haskey(dict, "temperature") ? pass.temperature = T(dict["temperature"]) : pass.temperature = missing
    pass.material = material_properties[materials[dict["material"]]]
    pass.charge_density_model = if haskey(dict, "charge_density_model")
        ChargeDensity(T, dict["charge_density_model"], inputunit_dict)
    else
        ZeroChargeDensity{T}()
    end
    pass.geometry = Geometry(T, dict["geometry"], inputunit_dict)
    pass.geometry_positive, pass.geometry_negative = get_decomposed_volumes(pass.geometry)
    return pass
end

function println(io::IO, d::Passive{T}) where {T}
    println("\t________"*"Passive{$T} $(d.name) $(d.id)"*"________\n")
    println("\t---General Properties---")
    println("\t-Potential: \t\t $(d.potential) V")
    println("\t-Material: \t $(d.material.name)")
    println()
end
print(io::IO, d::Passive{T}) where {T} = print(io, "Passive $(d.name) - id $(d.id) - $(d.potential) V")
show(io::IO, d::Passive) = print(io, d)
show(io::IO,::MIME"text/plain", d::Passive) = show(io, d)


text:
using .Base: sub_with_overflow, add_with_overflow, mul_with_overflow

function _findstarts_for_indices(x)
    _tmp = zeros(Bool, length(x))
    _tmp[1] = true
    Threads.@threads for i in 2:length(x)
        !isequal(x[i-1], x[i]) ? _tmp[i]=true : nothing
    end
    findall(_tmp)
end

function compute_indices(groups, ngroups, ::Val{T}; threads = true) where T
	idx = Vector{T}(undef, length(groups))
    _fill_idx_for_sort!(idx)
	if length(groups) == ngroups
		return idx, copy(idx)
	end
	# TODO we have the same ifelse in sort, probably we need to clean up these into a new function
	if threads && Threads.nthreads() > 1 && length(groups) > Threads.nthreads() && ngroups > 100_000 && ngroups*Threads.nthreads() < length(groups)
		starts = _ds_sort_int_missatright_nopermx_threaded!(groups, idx, ngroups, 1, Val(T))
	elseif threads && Threads.nthreads() > 1 && length(groups) > Threads.nthreads() && ngroups > 100_000
		starts = _ds_sort_int_missatright_nopermx_threaded_lm!(groups, idx, ngroups, 1, Val(T))
	else
		starts = _ds_sort_int_missatright_nopermx!(groups, idx, ngroups, 1, Val(T))
	end
	pop!(starts)
	pop!(starts)
	pop!(starts)
	idx, starts
end

# fast combine for gatherby data

mutable struct GatherBy
    parent
    groupcols
    groups
    lastvalid
    mapformats::Bool
    perm
    starts
end
Base.copy(gds::GatherBy) = GatherBy(copy(gds.parent), copy(gds.groupcols), copy(gds.groups), gds.lastvalid, gds.mapformats, gds.perm === nothing ? nothing : copy(gds.perm), gds.starts === nothing ? nothing : copy(gds.starts))

nrow(ds::GatherBy) = nrow(ds.parent)
ncol(ds::GatherBy) = ncol(ds.parent)
Base.names(ds::GatherBy, kwargs...) = names(ds.parent, kwargs...)
_names(ds::GatherBy) = _names(ds.parent)
_columns(ds::GatherBy) = _columns(ds.parent)
index(ds::GatherBy) = index(ds.parent)
Base.parent(ds::GatherBy) = ds.parent

Base.summary(gds::GatherBy) =
        @sprintf("%d×%d View of GatherBy Dataset, Gathered by: %s", size(gds.parent)..., join(_names(gds.parent)[gds.groupcols], " ,"))

function Base.show(io::IO, gds::GatherBy;

	kwargs...)
	if length(_get_perms(gds)) > 200
		_show(io, view(gds.parent, [first(gds.perm, 100);last(gds.perm, 100)], :); title = summary(gds), show_omitted_cell_summary=false, show_row_number  = false, kwargs...)
	else
		_show(io, view(gds.parent, gds.perm, :); title = summary(gds), show_omitted_cell_summary=false, show_row_number  = false, kwargs...)
	end
end

Base.show(io::IO, mime::MIME"text/plain", gds::GatherBy;
          kwargs...) =
    show(io, gds; title = summary(gds), kwargs...)

function _group_creator!(groups, starts, ngroups)
	if ngroups == 1
		fill!(groups, 1)
		return
	end
  	for j in 1:ngroups
		lo = starts[j]
		j == ngroups ? hi = length(groups) : hi = starts[j + 1] - 1
		fill!(view(groups, lo:hi), j)
	end
end

# eachrow = true tells gatherby that each row of passed dataset is a new group - this is useful for transpose()
function gatherby(ds::AbstractDataset, cols::MultiColumnIndex; mapformats::Bool = true, stable::Bool = true, isgathered::Bool = false, eachrow::Bool = false, threads = true)
    colsidx = index(ds)[cols]
	T = nrow(ds) < typemax(Int32) ? Int32 : Int64
	_check_consistency(ds)
	if isgathered
		if eachrow
			return GatherBy(ds, colsidx, 1:nrow(ds), nrow(ds), mapformats, 1:nrow(ds), 1:nrow(ds))
		else
			colindex, ranges, last_valid_index = _find_starts_of_groups(ds, colsidx, Val(T); mapformats = mapformats, threads = threads)
		 	groups = Vector{T}(undef, nrow(ds))
		 	_group_creator!(groups, ranges, last_valid_index)
		 	return GatherBy(ds, colindex, groups, last_valid_index, mapformats, 1:nrow(ds), ranges)
		end
	else
		if eachrow
			a = _gather_groups(ds, colsidx, Val(T), mapformats = mapformats, stable = stable, threads = threads)
			b = compute_indices(a[1], a[3], nrow(ds) < typemax(Int32) ? Val(Int32) : Val(Int64); threads = threads)
			return GatherBy(ds, colsidx, 1:nrow(ds), nrow(ds), mapformats, b[1], 1:nrow(ds))
		else
			a = _gather_groups(ds, colsidx, Val(T), mapformats = mapformats, stable = stable, threads = threads)
			return GatherBy(ds, colsidx, a[1], a[3], mapformats, nothing, nothing)
		end
	end
end
gatherby(ds::AbstractDataset, col::ColumnIndex; mapformats = true, stable = true, isgathered = false, eachrow = false, threads = true) = gatherby(ds, [col], mapformats = mapformats, stable = stable, isgathered = isgathered, eachrow = eachrow, threads = threads)

__SPFRMT(x) = x & 1023
__SPFRMT(::Missing) = missing # not needed

# currently not been used in gatherby
# use sort and format trick for fast gatherby - hm stands for high memory footprint
function hm_gatherby(ds::AbstractDataset, cols::MultiColumnIndex; mapformats = false, threads = true)
	modify!(ds, cols=>byrow(hash; threads = threads, mapformats = mapformats)=>:___tmp___cols8934, :___tmp___cols8934=>identity=>:___tmp___cols8934_2)
	setformat!(ds, :___tmp___cols8934_2=>__SPFRMT)
	gds = groupby(ds, [:___tmp___cols8934_2, :___tmp___cols8934], stable = false, threads = threads)
	grpcols, ranges, last_valid_index = _find_starts_of_groups(view(ds, gds.perm, cols), cols, nrow(ds) < typemax(Int32) ? Val(Int32) : Val(Int64); mapformats = mapformats, threads = threads)
	select!(ds, Not([:___tmp___cols8934, :___tmp___cols8934_2]))
	GatherBy(ds, grpcols, nothing, last_valid_index, mapformats, gds.perm, ranges)
end

function _fill_mapreduce_col!(x, f, op, y, loc)
    @inbounds for i in 1:length(y)
        x[loc[i]] = op(x[loc[i]], f(y[i]))
    end
end

function _fill_mapreduce_col!(x, f::Vector, op, y, loc)
	@inbounds for i in 1:length(y)
        x[loc[i]] = op(x[loc[i]], f[loc[i]](y[i]))
    end
end

function _fill_mapreduce_col_threaded!(x, f, op, y, loc, nt)
	@sync for thid in 0:nt-1
		Threads.@spawn for i in 1:length(y)
        	@inbounds if loc[i] % nt == thid
				# we need to do more complicated stuff here?
				x[loc[i]] = op(x[loc[i]], f(y[i]))
			end
		end
    end
end

function _fill_mapreduce_col_threaded!(x, f::Vector, op, y, loc, nt)
	@sync for thid in 0:nt-1
		Threads.@spawn for i in 1:length(y)
        	@inbounds if loc[i] % nt == thid
				x[loc[i]] = op(x[loc[i]], f[loc[i]](y[i]))
			end
		end
    end
end

function gatherby_mapreduce(gds::GatherBy, f, op, col::ColumnIndex, nt, init, ::Val{T}; promotetypes = false, threads = true) where T
	CT = T
	if promotetypes
	    T <: Base.SmallSigned ? CT = Int : nothing
		T <: Base.SmallUnsigned ? CT = UInt : nothing
	end
	res = allocatecol(Union{CT, Missing}, gds.lastvalid)
    fill!(res, init)
	if threads && Threads.nthreads() > 1 && gds.lastvalid > 100_000
		_fill_mapreduce_col_threaded!(res, f, op, _columns(gds.parent)[index(gds.parent)[col]], gds.groups, nt)
	else
    	_fill_mapreduce_col!(res, f, op, _columns(gds.parent)[index(gds.parent)[col]], gds.groups)
	end
    res
end

_gatherby_maximum(gds, col; f = identity, nt = Threads.nthreads(), threads = true) = gatherby_mapreduce(gds, f, _stat_max_fun, col, nt, missing, Val(nonmissingtype(eltype(gds.parent[!, col]))), threads = threads)
_gatherby_minimum(gds, col; f = identity, nt = Threads.nthreads(), threads = true) = gatherby_mapreduce(gds, f, _stat_min_fun, col, nt, missing, Val(nonmissingtype(eltype(gds.parent[!, col]))), threads = threads)
_gatherby_sum(gds, col; f = identity, nt = Threads.nthreads(), threads = true) = gatherby_mapreduce(gds, f, _stat_add_sum, col, nt, missing, Val(typeof(zero(Core.Compiler.return_type(f, Tuple{eltype(gds.parent[!, col])})))), promotetypes = true, threads = threads)
_gatherby_n(gds, col; nt = Threads.nthreads(), threads = true) = _gatherby_sum(gds, col, f = _stat_notmissing, nt = nt, threads = threads)
_gatherby_length(gds, col; nt = Threads.nthreads(), threads = true) = _gatherby_sum(gds, col, f = x->1, nt = nt, threads = threads)
_gatherby_cntnan(gds, col; nt = Threads.nthreads(), threads = true) = _gatherby_sum(gds, col, f = ISNAN, nt = nt, threads = threads)
_gatherby_nmissing(gds, col; nt = Threads.nthreads(), threads = true) = _gatherby_sum(gds, col, f = _stat_ismissing, nt = nt, threads = threads)

function _fill_gatherby_mean_barrier!(res, sval, nval)
	@inbounds for i in 1:length(nval)
		if nval[i] == 0
			res[i] = missing
		else
			res[i] = sval[i]/nval[i]
		end
	end
end

function _gatherby_mean(gds, col; nt = Threads.nthreads(), threads = true)
	if threads
		nt2 = max(div(nt, 2),1)

		t1 = Threads.@spawn _gatherby_sum(gds, col, nt = nt2)
		t2 = Threads.@spawn _gatherby_n(gds, col, nt = nt2)
		sval = fetch(t1)
		nval = fetch(t2)
	else
		t1 = _gatherby_sum(gds, col, threads = threads)
		t2 = _gatherby_n(gds, col, threads = threads)
		sval = t1
		nval = t2
	end

	T = Core.Compiler.return_type(/, Tuple{nonmissingtype(eltype(sval)), nonmissingtype(eltype(nval))})
	res = _our_vect_alloc(Union{Missing, T}, length(nval))
	_fill_gatherby_mean_barrier!(res, sval, nval)
	res
end

function _fill_gatherby_var_barrier!(res, countnan, meanval, ss, nval, cal_std, dof)

	@inbounds for i in 1:length(nval)
		if cal_std
			if countnan[i] > 0
				res[i] = NaN
			elseif nval[i] == 0
				res[i] = missing
			elseif nval[i] == 1 && dof
				res[i] = missing
			else
				res[i] = sqrt(ss[i]/(nval[i]-Int(dof)))
			end
		else
			if countnan[i] > 0
				res[i] = NaN
			elseif nval[i] == 0
				res[i] = missing
			elseif nval[i] == 1 && dof
				res[i] = missing
			else
				res[i] = ss[i]/(nval[i]-Int(dof))
			end
		end
	end
end

# TODO directly calculating var should be a better approach
function _gatherby_var(gds, col; dof = true, cal_std = false, threads = true)
	if threads
		nt = Threads.nthreads()
		nt2 = max(div(nt,2),1)
		t1 = Threads.@spawn _gatherby_cntnan(gds, col, nt = nt2)
		t2 = Threads.@spawn _gatherby_mean(gds, col, nt = nt2)
		meanval = fetch(t2)
		t3 = Threads.@spawn gatherby_mapreduce(gds, [x->abs2(x - meanval[i]) for i in 1:length(meanval)], _stat_add_sum, col, nt2, missing, Val(Float64))
		t4 = Threads.@spawn _gatherby_n(gds, col, nt = nt2)
		countnan = fetch(t1)
		ss = fetch(t3)
		nval = fetch(t4)
	else
		t1 = _gatherby_cntnan(gds, col, threads = threads)
		t2 = _gatherby_mean(gds, col, threads = threads)
		meanval = t2
		t3 = gatherby_mapreduce(gds, [x->abs2(x - meanval[i]) for i in 1:length(meanval)], _stat_add_sum, col, Threads.nthreads(), missing, Val(Float64), threads = threads)
		t4 = _gatherby_n(gds, col, threads = threads)
		countnan = t1
		ss = t3
		nval = t4
	end
	T = Core.Compiler.return_type(/, Tuple{nonmissingtype(eltype(meanval)), nonmissingtype(eltype(nval))})
	res = _our_vect_alloc(Union{Missing, T}, length(nval))
	_fill_gatherby_var_barrier!(res, countnan, meanval, ss, nval, cal_std, dof)
	res
end
_gatherby_std(gds, col; dof = true, threads = true) = _gatherby_var(gds, col; dof = dof, cal_std = true, threads = threads)

const FAST_GATHERBY_REDUCTION = [sum, length, minimum, maximum, mean, var, std, n, nmissing]

function _fast_gatherby_reduction(gds, ms)
    !(gds isa GatherBy) && return false
    gds.groups == nothing && return false
    for i in 1:length(ms)
        if (ms[i].second.first isa Expr) && ms[i].second.first.head == :BYROW
        elseif (ms[i].second.first isa Base.Callable)
            flag = ms[i].second.first ∈ FAST_GATHERBY_REDUCTION
            !flag && return false
        end
    end
    return true
end

function _fast_gatherby_groups_to_res!(outres, x, grp)
	# the first element of each group is used for the output data set
    @inbounds for i in length(x):-1:1
        outres[grp[i]] = x[i]
    end
end

function _fast_gatherby_combine_f_barrier(gds, col, newds, mssecond, mslast, newds_lookup, grp, ngrps, threads)

    if !(mssecond isa Expr)
        if mssecond == sum
            newds[!, mslast] =  _gatherby_sum(gds, col, threads = threads)
        elseif mssecond == maximum
            newds[!, mslast] =  _gatherby_maximum(gds, col, threads = threads)
        elseif mssecond == minimum
            newds[!, mslast] =  _gatherby_minimum(gds, col, threads = threads)
        elseif mssecond == mean
            newds[!, mslast] =  _gatherby_mean(gds, col, threads = threads)
        elseif mssecond == mean
            newds[!, mslast] =  _gatherby_mean(gds, col, threads = threads)
        elseif mssecond == var
            newds[!, mslast] =  _gatherby_var(gds, col, dof = true, threads = threads)
        elseif mssecond == std
            newds[!, mslast] =  _gatherby_std(gds, col, dof = true, threads = threads)
        elseif mssecond == length
            newds[!, mslast] =  _gatherby_length(gds, col, threads = threads)
        elseif mssecond == IMD.n
            newds[!, mslast] =  _gatherby_n(gds, col, threads = threads)
        else mssecond == IMD.nmissing
            newds[!, mslast] =  _gatherby_nmissing(gds, col, threads = threads)
        end

    elseif (mssecond isa Expr) && mssecond.head == :BYROW
        newds[!, mslast] = byrow(newds, mssecond.args[1], col; mssecond.args[2]...)
    else
        throw(ArgumentError("`combine` doesn't support $(msfirst=>mssecond=>mslast) combination"))
    end
end

function _combine_fast_gatherby_reduction(gds, ms, newlookup, new_nm; dropgroupcols = false, threads = true)
    groupcols = gds.groupcols
    ngroups = gds.lastvalid
    groups = gds.groups

    all_names = _names(gds.parent)

    newds_idx = Index(Dict{Symbol, Int}(), Symbol[], Dict{Int, Function}(), Int[], Bool[], false, [], Int[], 1, false)

    newds = Dataset([], newds_idx)
    newds_lookup = index(newds).lookup
    var_cnt = 1
	if !dropgroupcols
	    for j in 1:length(groupcols)
	        addmissing = false
	        _tmpres = allocatecol(gds.parent[!, groupcols[j]].val, ngroups, addmissing = addmissing)
	        if DataAPI.refpool(_tmpres) !== nothing
	            _fast_gatherby_groups_to_res!(_tmpres.refs, DataAPI.refarray(_columns(gds.parent)[groupcols[j]]), groups)
	            push!(_columns(newds), _tmpres)
	        else
	            _fast_gatherby_groups_to_res!(_tmpres, _columns(gds.parent)[groupcols[j]], groups)
	            push!(_columns(newds), _tmpres)
	        end
	        push!(index(newds), new_nm[var_cnt])
			setformat!(newds, new_nm[var_cnt] => getformat(parent(gds), groupcols[j]))
	        var_cnt += 1
	    end
	end
	for i in 1:length(ms)
		_fast_gatherby_combine_f_barrier(gds, ms[i].first, newds, ms[i].second.first, ms[i].second.second, newds_lookup, groups, ngroups, threads)
		# if !haskey(index(newds), ms[i].second.second)
			# push!(index(newds), ms[i].second.second)
		# end

	end
	newds
end


text:
# function to integrate a unit circle to find pi via monte_carlo
function monte_carlo(n::Int64)

    pi_count = 0
    for i = 1:n
        point_x = rand()
        point_y = rand()

        if (point_x^2 + point_y^2 < 1)
            pi_count += 1
        end
    end

    pi_estimate = 4*pi_count/n
    println(pi - pi_estimate)
end

monte_carlo(10000000)

max_issues_repo_name: shiffman/simuleios

max_forks_repo_name: shiffman/simuleios


text:
mutable struct fCTM <: TopicModel
	"Filtered correlated topic model."

	K::Int
	M::Int
	V::Int
	N::Vector{Int}
	C::Vector{Int}
	corp::Corpus
	topics::VectorList{Int}
	eta::Float64
	mu::Vector{Float64}
	sigma::Symmetric{Float64}
	invsigma::Symmetric{Float64}
	kappa::Vector{Float64}
	kappa_old::Vector{Float64}
	kappa_temp::Vector{Float64}
	beta::Matrix{Float64}
	beta_old::Matrix{Float64}
	beta_temp::Matrix{Float64}
	lambda::VectorList{Float64}
	lambda_old::VectorList{Float64}
	vsq::VectorList{Float64}
	logzeta::Vector{Float64}
	tau::VectorList{Float64}
	tau_old::VectorList{Float64}
	phi::MatrixList{Float64}
	elbo::Float64

	function fCTM(corp::Corpus, K::Integer)
		check_corp(corp)
		K > 0 || throw(ArgumentError("Number of topics must be a positive integer."))

		M, V, U = size(corp)
		N = [length(doc) for doc in corp]
		C = [size(doc) for doc in corp]

		topics = [collect(1:V) for _ in 1:K]

		eta = 0.5
		mu = zeros(K)
		sigma = Symmetric(Matrix{Float64}(I, K, K))
		invsigma = copy(sigma)
		kappa = rand(Dirichlet(V, 1.0))
		kappa_old = copy(kappa)
		kappa_temp = zeros(V)
		beta = rand(Dirichlet(V, 1.0), K)'
		beta_old = copy(beta)
		beta_temp = zeros(K, V)
		lambda = [zeros(K) for _ in 1:M]
		lambda_old = deepcopy(lambda)
		vsq = [ones(K) for _ in 1:M]
		logzeta = fill(0.5, M)
		tau = [fill(eta, N[d]) for d in 1:M]
		tau_old = deepcopy(tau)
		phi = [ones(K, N[d]) / K for d in 1:min(M, 1)]
		elbo = 0

		model = new(K, M, V, N, C, copy(corp), topics, eta, mu, sigma, invsigma, kappa, kappa_old, kappa_temp, beta, beta_old, beta_temp, lambda, lambda_old, vsq, logzeta, tau, tau_old, phi, elbo)
		return model
	end
end

function Elogpeta(model::fCTM, d::Int)
	"Compute E_q[log(P(eta))]."

	x = 0.5 * (logdet(model.invsigma) - model.K * log(2pi) - dot(diag(model.invsigma), model.vsq[d]) - dot(model.lambda[d] - model.mu, model.invsigma * (model.lambda[d] - model.mu)))
	return x
end

function Elogpc(model::fCTM, d::Int)
	"Compute E_q[log(P(c))]."

	counts = model.corp[d].counts
	x = log(@boink model.eta^dot(model.tau[d], counts) * (1 - model.eta)^(model.C[d] - dot(model.tau[d], counts)))
	return x
end

function Elogpz(model::fCTM, d::Int)
	"Compute E_q[log(P(z))]."

	counts = model.corp[d].counts
	x = dot(model.phi[1]' * model.lambda[d], counts) - model.C[d] * (sum(exp.(model.lambda[d] + 0.5 * model.vsq[d] .- model.logzeta[d])) + model.logzeta[d] - 1)
	return x
end

function Elogpw(model::fCTM, d::Int)
	"Compute E_q[log(P(w))]."

	terms, counts = model.corp[d].terms, model.corp[d].counts
	x = sum(model.phi[1] .* log.(@boink model.beta[:,terms]) * (counts .* model.tau[d])) + dot(counts .* (1 .- model.tau[d]), log.(@boink model.kappa[terms]))
	return x
end

function Elogqeta(model::fCTM, d::Int)
	"Compute E_q[log(q(eta))]."

	x = -entropy(MvNormal(model.lambda[d], diagm(model.vsq[d])))
	return x
end

function Elogqc(model::fCTM, d::Int)
	"Compute E_q[log(q(c))]."

	counts = model.corp[d].counts
	x = -sum([c * entropy(Bernoulli(model.tau[d][n])) for (n, c) in enumerate(counts)])
	return x
end

function Elogqz(model::fCTM, d::Int)
	"Compute E_q[log(q(z))]."

	counts = model.corp[d].counts
	x = -sum([c * entropy(Categorical(model.phi[1][:,n])) for (n, c) in enumerate(counts)])
	return x
end

function update_elbo!(model::fCTM)
	"Update the evidence lower bound."

	model.elbo = 0
	for d in 1:model.M
		terms = model.corp[d].terms
		model.phi[1] = additive_logistic(model.tau_old[d]' .* log.(@boink model.beta_old[:,terms]) .+ model.lambda_old[d], dims=1)
		model.elbo += Elogpeta(model, d) + Elogpc(model, d) + Elogpz(model, d) + Elogpw(model, d) - Elogqeta(model, d) - Elogqc(model, d) - Elogqz(model, d)
	end

	return model.elbo
end

function update_eta!(model::fCTM)
	"Update eta."
	"Analytic."

	model.eta = sum([dot(model.tau[d], model.corp[d].counts) for d in 1:model.M]) / sum(model.C)
end

function update_mu!(model::fCTM)
	"Update mu."
	"Analytic."

	model.mu = sum(model.lambda) / model.M
end

function update_sigma!(model::fCTM)
	"Update sigma."
	"Analytic"

	model.sigma = Symmetric((diagm(sum(model.vsq)) + (hcat(model.lambda...) .- model.mu) * (hcat(model.lambda...) .- model.mu)') / model.M)
	model.invsigma = inv(model.sigma)
end

function update_kappa!(model::fCTM)
	"Reset kappa variables."
	"Analytic."

	model.kappa_old = model.kappa
	model.kappa = model.kappa_temp ./ sum(model.kappa_temp)
	model.kappa_temp = zeros(model.V)
end

function update_kappa!(model::fCTM, d::Int)
	"Update kappa."
	"Analytic."

	terms, counts = model.corp[d].terms, model.corp[d].counts
	model.kappa_temp[terms] += (1 .- model.tau[d]) .* counts
end

function update_beta!(model::fCTM)
	"Reset beta variables."

	model.beta_old = model.beta
	model.beta = model.beta_temp ./ sum(model.beta_temp, dims=2)
	model.beta_temp = zeros(model.K, model.V)
end

function update_beta!(model::fCTM, d::Int)
	"Update beta."
	"Analytic."

	terms, counts = model.corp[d].terms, model.corp[d].counts
	model.beta_temp[:,terms] += model.phi[1] .* (model.tau[d] .* counts)'
end

function update_lambda!(model::fCTM, d::Int, niter::Integer, ntol::Real)
	"Update lambda."
	"Newton's method."

	model.lambda_old[d] = model.lambda[d]

	counts = model.corp[d].counts
	for _ in 1:niter
		lambda_grad = model.invsigma * (model.mu - model.lambda[d]) + model.phi[1] * counts - model.C[d] * exp.(model.lambda[d] + 0.5 * model.vsq[d] .- model.logzeta[d])
		lambda_negative_hess = model.invsigma + model.C[d] * Diagonal(exp.(model.lambda[d] + 0.5 * model.vsq[d] .- model.logzeta[d]))
		model.lambda[d] += lambda_negative_hess \ lambda_grad

		if norm(lambda_grad) < ntol
			break
		end
	end
end

function update_vsq!(model::fCTM, d::Int, niter::Integer, ntol::Real)
	"Update vsq."
	"Newton's method with back-tracking line search."

	for i in 1:model.K
		for _ in 1:niter
			rho = 1.0
			vsq_grad = -0.5 * (model.invsigma[i,i] + model.C[d] * exp(model.lambda[d][i] + 0.5 * model.vsq[d][i] - model.logzeta[d]) - 1 / model.vsq[d][i])
			vsq_invhess_diag = -1 / (0.25 * model.C[d] * exp(model.lambda[d][i] + 0.5 * model.vsq[d][i] - model.logzeta[d]) + 0.5 / model.vsq[d][i]^2)
			p = vsq_invhess_diag * vsq_grad

			while model.vsq[d][i] - rho * p <= 0
				rho *= 0.5
			end
			model.vsq[d][i] -= rho * p

			if rho * abs(vsq_grad) < ntol
				break
			end
		end
	end
	@positive model.vsq[d]
end

function update_logzeta!(model::fCTM, d::Int)
	"Update logzeta."
	"Analytic."

	model.logzeta[d] = logsumexp(model.lambda[d] + 0.5 * model.vsq[d])
end

function update_tau!(model::fCTM, d::Int)
	"Update tau."
	"Analytic."

	model.tau_old[d] = model.tau[d]

	terms = model.corp[d].terms
	model.tau[d] = model.eta ./ (@boink model.eta .+ (1 - model.eta) * (model.kappa[terms] .* vec(prod(model.beta[:,terms].^-model.phi[1], dims=1))))
end

function update_phi!(model::fCTM, d::Int)
	"Update phi."
	"Analytic"

	terms = model.corp[d].terms
	model.phi[1] = additive_logistic(model.tau[d]' .* log.(@boink model.beta[:,terms]) .+ model.lambda[d], dims=1)
end

function train!(model::fCTM; iter::Integer=150, tol::Real=1.0, niter=1000, ntol::Real=1/model.K^2, viter::Integer=10, vtol::Real=1/model.K^2, checkelbo::Real=1, printelbo::Bool=true)
	"Coordinate ascent optimization procedure for filtered correlated topic model variational Bayes algorithm."

	check_model(model)
	all([tol, ntol, vtol] .>= 0)										|| throw(ArgumentError("Tolerance parameters must be nonnegative."))
	all([iter, niter, viter] .>= 0)										|| throw(ArgumentError("Iteration parameters must be nonnegative."))
	(isa(checkelbo, Integer) & (checkelbo > 0)) | (checkelbo == Inf)	|| throw(ArgumentError("checkelbo parameter must be a positive integer or Inf."))
	all([isempty(doc) for doc in model.corp]) && (iter = 0)
	(checkelbo <= iter) && update_elbo!(model)

	for k in 1:iter
		for d in 1:model.M
			for _ in 1:viter
				update_phi!(model, d)
				update_tau!(model, d)
				update_logzeta!(model, d)
				update_lambda!(model, d, niter, ntol)
				update_vsq!(model, d, niter, ntol)
				if norm(model.lambda[d] - model.lambda_old[d]) < vtol
					break
				end
			end
			update_beta!(model, d)
			update_kappa!(model, d)
		end
		update_beta!(model)
		update_kappa!(model)
		update_sigma!(model)
		update_mu!(model)
		#update_eta!(model)

		if check_elbo!(model, checkelbo, printelbo, k, tol)
			break
		end
	end

	model.topics = [reverse(sortperm(vec(model.beta[i,:]))) for i in 1:model.K]
	nothing
end


text:
module MYGEN2

using Modia

setLogMerge(false)

include("../src/ModelMechanic.jl")
include("../src/ModelEPSphasor.jl")
include("../src/ModelController.jl")

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

VarVolts = Model(   # Voltage dip
    Tstart = 0.0,
    dT = 0.15,
    Vmin = 0.0,
    y = output,
    equations = :[
        y = if after(Tstart) && !after(Tstart+dT); Vmin else; 1.0 end;
    ]
)

SMtest = Model(
    Project = parameter | Map(value="HALLO"),
    LOAD = parameter | Map(value=0.777, fixed=false),
    VEXC = parameter | Map(value=1.444, fixed=false),

    wReference= parameter | 1.0 | Var(info = "ref. Frame"),
    Omegarated=100.0*pi,

    INtorque   = Var(input=true, value= 0.0, comment="delta Torque"),
    INexc      = Var(input=true, value= 0.0, comment="delta setpoint volt"),

    OUTspeed = Var(output=true, comment="synUnit.SM.w"),
# OUT1 = Var(output=true, comment="J.w - ERROR"),   #geht nicht >> J.w ist StateVar
    OUTpower = Var(output=true, comment="meterGEN.powerElRe"),
    OUTvoltage = Var(output=true, comment="meterGEN.voltage"),

    GridVolts  = EPSphasor.VoltageSource,
    GC = EPSphasor.GroundVec,
#     synUnit = EPSphasor.SYNunit,    # = ExcCONST
    synUnit = EPSphasor.SYNunit | Map(exciter = redeclare | Controller.ExcST8C),
    transformer = EPSphasor.Transformer | Map( xt=0.10, rt=0.1/40.0, Strafo=119e6),

    meterGRID   = EPSphasor.Meter,
    meterGEN    = EPSphasor.Meter,
    varvolts    = VarVolts,

    J = Mechanic.Inertia | Map(
        J=3.0,
        w=Var(init=1.0),
        phi=Var(init=0.0),
        ),

    equations = :[
        J.flange_b.tau = LOAD + INtorque
        OUTspeed = synUnit.SM.w
#         OUT1 = J.w # J.w   geht nicht ->  ist State VAR!!
        OUTpower = meterGEN.powerRe
        OUTvoltage = meterGEN.voltage
        synUnit.VoltageREF = VEXC + INexc
#         wReference = synUnit.SM.w
    ],

    connect = :[
        (GridVolts.V, varvolts.y)
        (GridVolts.p, meterGRID.n)
        (meterGRID.p, transformer.pin1)
        (transformer.pin2, meterGEN.n)
        (meterGEN.p, synUnit.pin)
        (synUnit.flange, J.flange_a)
        (GridVolts.n, GC.p)
    ]
)

instModel = @instantiateModel(SMtest,
    unitless=true,
    log=false,
    logTiming=false,
    logDetails=false,
    logCode=false,
    logStateSelection=false,
#     saveCodeOnFile="CODE.txt",
)

end  # module MYGEN2

max_stars_repo_name: johhell/EPSphasor

max_issues_repo_name: johhell/EPSphasor

max_forks_repo_name: johhell/EPSphasor