using Random: randstring

@testset "G" begin
    @testset "gcpool" begin
        try
            data = [randstring() for i in 1:10]
            pcpool("pcpool_test", data)
            cvals = gcpool("pcpool_test")
            @test data == cvals
            @test gcpool("norbert") === nothing
        finally
            kclear()
        end
    end
    @testset "gdpool" begin
        try
            exp = collect(1.0:10.0)
            pdpool("array", exp)
            act = gdpool("array")
            @test exp == act
            act = gdpool("array", start = 8)
            @test [8.0, 9.0, 10.0] == act
            act = gdpool("array", room = 8)
            @test collect(1.0:8.0) == act
            @test gdpool("norbert") === nothing
        finally
            kclear()
        end
    end
    @testset "georec" begin
        try
            furnsh(path(CORE, :pck))
            radii = bodvrd("EARTH", "RADII", 3)
            flat = (radii[1] - radii[3]) / radii[1]
            lon = deg2rad(118.0)
            lat = deg2rad(32.0)
            alt = 0.0
            output = georec(lon, lat, alt, radii[1], flat)
            expected = [-2541.74621567, 4780.329376, 3360.4312092]
            @test expected ≈ output
            @test_throws SpiceError georec(lon, lat, alt, -1.0, flat)
            @test_throws SpiceError georec(lon, lat, alt, radii[1], 2.0)
        finally
            kclear()
        end
    end
    @testset "getelm" begin
        try
            tle = ["1 18123U 87 53  A 87324.61041692 -.00000023  00000-0 -75103-5 0 00675",
                   "2 18123  98.8296 152.0074 0014950 168.7820 191.3688 14.12912554 21686"]
            furnsh(path(CORE, :lsk))
            epoch, elems = getelm(1950, tle)
            expected_elems = [-6.969196665949579e-13, 0.0, -7.510300000000001e-06,
                              1.724901918428988, 2.653029617396028, 0.001495,
                              2.9458016181010693, 3.3400156455905243, 0.06164994027515544,
                              -382310404.79526937]
            expected_epoch = -382310404.79526937
            @test expected_elems ≈ elems
            @test epoch ≈ expected_epoch
        finally
            kclear()
        end
    end
    @testset "getfat" begin
    arch, outtype = getfat(path(CORE, :lsk))
        @test arch == "KPL"
        @test outtype == "LSK"
    end
    @testset "getfov" begin
        try
            kernel = tempname()
            open(kernel, "w") do kernel_file
                write(kernel_file, "\\begindata\n")
                write(kernel_file, "INS-999004_FOV_SHAPE            = 'POLYGON'\n")
                write(kernel_file, "INS-999004_FOV_FRAME            = 'SC999_INST004'\n")
                write(kernel_file, "INS-999004_BORESIGHT            = (  0.0,  1.0,  0.0 )\n")
                write(kernel_file, "INS-999004_FOV_BOUNDARY_CORNERS = (  0.0,  0.8,  0.5,\n")
                write(kernel_file, "                                     0.4,  0.8, -0.2,\n")
                write(kernel_file, "                                    -0.4,  0.8, -0.2,\n")
                write(kernel_file, "\\begintext\n")
            end
            furnsh(kernel)
            shape, frame, bsight, bounds = getfov(-999004, 4, 32, 32)
            @test shape == "POLYGON"
            @test frame == "SC999_INST004"
            @test bsight ≈ [0.0, 1.0, 0.0]
            @test length(bounds) == 3
    expected = [[0.0, 0.8, 0.5], [0.4, 0.8, -0.2], [-0.4, 0.8, -0.2]]
            @test expected ≈ bounds
        finally
            kclear()
        end
    end
    @testset "gfbail" begin
        @test !Bool(SPICE.gfbail())
    end
    @testset "gfclrh" begin
        SPICE.gfclrh()
        @test !Bool(SPICE.gfbail())
    end
    @testset "gfdist" begin
        try
            furnsh(path(CORE, :lsk), path(CORE, :pck), path(CORE, :spk))
            et0 = str2et("2007 JAN 01 00:00:00 TDB")
            et1 = str2et("2007 APR 01 00:00:00 TDB")
            cnfine = SpiceDoubleCell(2)
            wninsd!(cnfine, et0, et1)
            result = gfdist("moon", "none", "earth", ">", 400000, 0.0, spd(), 1000, cnfine)
            count = wncard(result)
            @test count == 4
            results = String[]
            for i in 1:count
                left, right = wnfetd(result, i)
                timstr_left = timout(left, "YYYY-MON-DD HR:MN:SC.###### (TDB) ::TDB ::RND", 41)
                timstr_right = timout(right, "YYYY-MON-DD HR:MN:SC.###### (TDB) ::TDB ::RND", 41)
                push!(results, timstr_left)
                push!(results, timstr_right)
            end
            expected = ["2007-JAN-08 00:11:07.661897 (TDB)", "2007-JAN-13 06:37:47.937762 (TDB)",
                        "2007-FEB-04 07:02:35.320555 (TDB)", "2007-FEB-10 09:31:01.829206 (TDB)",
                        "2007-MAR-03 00:20:25.228066 (TDB)", "2007-MAR-10 14:04:38.482902 (TDB)",
                        "2007-MAR-29 22:53:58.186230 (TDB)", "2007-APR-01 00:00:00.000000 (TDB)"]
            @test results == expected
        finally
            kclear()
        end
    end
    @testset "gfevnt" begin
        try
            furnsh(path(CORE, :lsk), path(CORE, :pck), path(CORE, :spk))
            et_start = str2et("2001 JAN 01 00:00:00.000")
            et_end   = str2et("2001 DEC 31 00:00:00.000")
            cnfine   = SpiceDoubleCell(2)
            wninsd!(cnfine, et_start, et_end)
            result   = SpiceDoubleCell(1000)
            qpnams   = ["TARGET", "OBSERVER", "ABCORR"]
            qcpars   = ["MOON  ", "EARTH   ", "LT+S  "]
            # Set the step size to 1/1000 day and convert to seconds
            gfsstp(0.001 * spd())
            qdpars = zeros(3)
            qipars = zeros(Int, 3)
            qlpars = falses(3)
            result = gfevnt(gfstep, gfrefn, "DISTANCE", qpnams, qcpars,
                            qdpars, qipars, qlpars, "LOCMAX", 0, 1.e-6, 0,
                            true, gfrepi, gfrepu, gfrepf, 10000, cnfine)
            @test length(result) == 26
            s_timout = "YYYY-MON-DD HR:MN:SC.###### (TDB) ::TDB ::RND"
            @test startswith(timout(result[1], s_timout), "2001-JAN-24 19:22:01.41871")
            @test startswith(timout(result[2], s_timout), "2001-JAN-24 19:22:01.41871")
            @test startswith(timout(result[3], s_timout), "2001-FEB-20 21:52:07.90087")
            @test startswith(timout(result[4], s_timout), "2001-FEB-20 21:52:07.90087")
        finally
            gfsstp(0.5)
            kclear()
        end
    end
    @testset "gffove" begin
        try
            furnsh(path(CORE, :lsk),
                   path(CORE, :pck),
                   path(CORE, :spk),
                   path(CASSINI, :ck),
                   path(CASSINI, :fk),
                   path(CASSINI, :ik),
                   path(CASSINI, :pck),
                   path(CASSINI, :sclk),
                   path(CASSINI, :tour_spk),
                   path(CASSINI, :sat_spk))
            et_start = str2et("2013-FEB-25 10:00:00.000")
            et_end   = str2et("2013-FEB-25 11:45:00.000")
            cnfine   = SpiceDoubleCell(2)
            wninsd!(cnfine, et_start, et_end)
            result = SpiceDoubleCell(1000)
            gfsstp(1.0)
            gffove!("CASSINI_ISS_NAC", "ELLIPSOID", [0.0, 0.0, 0.0], "ENCELADUS",
                    "IAU_ENCELADUS", "LT+S", "CASSINI", 1.e-6, gfstep, gfrefn, true, gfrepi,
                    gfrepu, gfrepf, cnfine, result)
            # Verify the expected results
            @test length(result) == 2
            s_timout = "YYYY-MON-DD HR:MN:SC UTC ::RND"
            @test timout(result[1], s_timout) == "2013-FEB-25 10:42:33 UTC"
            @test timout(result[2], s_timout) == "2013-FEB-25 11:45:00 UTC"
        finally
            gfsstp(0.5)
            kclear()
        end
    end
    @testset "gfilum" begin
        try
            furnsh(path(CORE, :lsk), path(CORE, :pck), path(CORE, :spk),
                   path(EXTRA, :mars_spk))
            pos = [3376.17890941875839416753, -325.55203839445334779157, -121.47422900638389364758]
            start_et = str2et("1971 OCT 02 00:00:00 UTC")
            end_et = str2et("1971 NOV 30 12:00:00 UTC")
            cnfine = SpiceDoubleCell(2000)
            wninsd!(cnfine, start_et, end_et)
            wnsolr = gfilum("ELLIPSOID", "INCIDENCE", "MARS", "SUN",
                            "IAU_MARS", "CN+S", "PHOBOS", pos,
                            "<", deg2rad(60.0), 0.0, 21600.0,
                            1000, cnfine)
            result = gfilum("Ellipsoid", "EMISSION", "MARS", "SUN",
                            "IAU_MARS", "CN+S", "PHOBOS", pos,
                            "<", deg2rad(20.0), 0.0, 900.0,
                            1000, wnsolr)
            @test wncard(result) > 0
            start_epoch = timout(result[1],  "YYYY MON DD HR:MN:SC.###### UTC")
            end_epoch   = timout(result[end], "YYYY MON DD HR:MN:SC.###### UTC")
            @test startswith(start_epoch, "1971 OCT 02")
            @test startswith(end_epoch, "1971 NOV 29")
        finally
            kclear()
        end
    end
    @testset "gfocce" begin
        try
            furnsh(path(CORE, :lsk), path(CORE, :pck), path(CORE, :spk))
            et0 = str2et("2001 DEC 01 00:00:00 TDB")
            et1 = str2et("2002 JAN 01 00:00:00 TDB")
            cnfine = SpiceDoubleCell(2)
            wninsd!(cnfine, et0, et1)
            result = SpiceDoubleCell(1000)
            gfsstp(20.0)
            gfocce!("Any", "moon", "ellipsoid", "iau_moon", "sun",
                    "ellipsoid", "iau_sun", "lt", "earth", 1.e-6,
                    gfstep, gfrefn, true, gfrepi, gfrepu, gfrepf,
                    cnfine, result)
            @test wncard(result) == 1
        finally
            gfsstp(0.5)
            kclear()
        end
    end
    @testset "gfoclt" begin
        try
            furnsh(path(CORE, :lsk), path(CORE, :pck), path(CORE, :spk))
            et0 = str2et("2001 DEC 01 00:00:00 TDB")
            et1 = str2et("2002 JAN 01 00:00:00 TDB")
            cnfine = SpiceDoubleCell(2)
            wninsd!(cnfine, et0, et1)
            result = gfoclt("ANY", "MOON", "ELLIPSOID", "IAU_MOON", "SUN",
                            "ELLIPSOID", "IAU_SUN", "LT", "EARTH", 15 * 60.0, cnfine)
            count = wncard(result)
            @test count == 1
            start, stop = wnfetd(result, 1)
            start_time = timout(start, "YYYY-MON-DD HR:MN:SC.###### (TDB) ::TDB ::RND", 41)
            end_time = timout(stop, "YYYY-MON-DD HR:MN:SC.###### (TDB) ::TDB ::RND", 41)
            @test start_time == "2001-DEC-14 20:10:14.203347 (TDB)"
            @test end_time == "2001-DEC-14 21:35:50.328804 (TDB)"
        finally
            kclear()
        end
    end
    @testset "gfpa" begin
        relate = ["=", "<", ">", "LOCMIN", "ABSMIN", "LOCMAX", "ABSMAX"]
        expected = Dict("=" => [
                                "2006-DEC-02 13:31:34.425",
                                "2006-DEC-02 13:31:34.425",
                                "2006-DEC-07 14:07:55.480",
                                "2006-DEC-07 14:07:55.480",
                                "2007-JAN-01 00:00:00.007",
                                "2007-JAN-01 00:00:00.007",
                                "2007-JAN-06 08:16:25.522",
                                "2007-JAN-06 08:16:25.522",
                                "2007-JAN-30 11:41:32.568",
                                "2007-JAN-30 11:41:32.568",
                               ],
                        "<" => [
                                "2006-DEC-02 13:31:34.425",
                                "2006-DEC-07 14:07:55.480",
                                "2007-JAN-01 00:00:00.007",
                                "2007-JAN-06 08:16:25.522",
                                "2007-JAN-30 11:41:32.568",
                                "2007-JAN-31 00:00:00.000",
                               ],
                        ">" => [
                                "2006-DEC-01 00:00:00.000",
                                "2006-DEC-02 13:31:34.425",
                                "2006-DEC-07 14:07:55.480",
                                "2007-JAN-01 00:00:00.007",
                                "2007-JAN-06 08:16:25.522",
                                "2007-JAN-30 11:41:32.568",
                               ],
                        "LOCMIN" => [
                                     "2006-DEC-05 00:16:50.327",
                                     "2006-DEC-05 00:16:50.327",
                                     "2007-JAN-03 14:18:31.987",
                                     "2007-JAN-03 14:18:31.987",
                                    ],
                        "ABSMIN" => [
                                     "2007-JAN-03 14:18:31.987",
                                     "2007-JAN-03 14:18:31.987",
                                    ],
                        "LOCMAX" => [
                                     "2006-DEC-20 14:09:10.402",
                                     "2006-DEC-20 14:09:10.402",
                                     "2007-JAN-19 04:27:54.610",
                                     "2007-JAN-19 04:27:54.610",
                                    ],
                        "ABSMAX" => [
                                     "2007-JAN-19 04:27:54.610",
                                     "2007-JAN-19 04:27:54.610",
                                    ])
        try
            furnsh(path(CORE, :lsk), path(CORE, :pck), path(CORE, :spk))
            et0 = str2et("2006 DEC 01")
            et1 = str2et("2007 JAN 31")
            cnfine = SpiceDoubleCell(2)
            wninsd!(cnfine, et0, et1)
            result = SpiceDoubleCell(2000)
            for relation in relate
                result = gfpa("Moon", "Sun", "LT+S", "Earth", relation, 0.57598845,
                              0.0, spd(), 5000, cnfine)
                count = wncard(result)
                @test count > 0
                if count > 0
                    temp_results = String[]
                    for i in 1:count
                        left, right = wnfetd(result, i)
                        timstr_left = timout(left, "YYYY-MON-DD HR:MN:SC.###")
                        timstr_right = timout(right, "YYYY-MON-DD HR:MN:SC.###")
                        push!(temp_results, timstr_left)
                        push!(temp_results, timstr_right)
                    end
                    @test temp_results == expected[relation]
                end
            end
        finally
            kclear()
        end
    end
    @testset "gfposc" begin
        try
            furnsh(path(CORE, :lsk), path(CORE, :pck), path(CORE, :spk))
            et0 = str2et("2007 JAN 01")
            et1 = str2et("2008 JAN 01")
            cnfine = SpiceDoubleCell(2)
            wninsd!(cnfine, et0, et1)
            result = gfposc("sun", "iau_earth", "none", "earth", "latitudinal", "latitude",
                            "absmax", 0.0, 0.0, 90.0 * spd(), 1000, cnfine)
            count = wncard(result)
            @test count == 1
            start, stop = wnfetd(result, 1)
            start_time = timout(start, "YYYY-MON-DD HR:MN:SC.###### (TDB) ::TDB ::RND", 41)
            end_time = timout(stop, "YYYY-MON-DD HR:MN:SC.###### (TDB) ::TDB ::RND", 41)
            @test start_time == end_time
            @test start_time == "2007-JUN-21 17:54:13.201561 (TDB)"
        finally
            kclear()
        end
    end
    @testset "gfrefn" begin
        s1 = [true, false]
        s2 = [true, false]
        for i in 1:2
            for j in 1:2
                scale = 10.0 * i + j
                t1 = 5.0 * scale
                t2 = 7.0 * scale
                t  = gfrefn(t1, t2, s1[i], s2[j])
                @test t ≈ scale * 6.0
            end
        end
        for i in 1:2
            for j in 1:2
                scale = 10.0 * i + j
                t1 = 15.0 * scale
                t2 = 7.0 * scale
                t  = gfrefn(t1, t2, s1[i], s2[j])
                @test t ≈ scale * 11.0
            end
        end
        for i in 1:2
            for j in 1:2
                scale = 10.0 * i + j
                t1 = -scale
                t2 = -scale
                t  = gfrefn(t1, t2, s1[i], s2[j])
                @test t ≈ -scale
            end
        end
    end
    @testset "gfrepi" begin
        window = SpiceDoubleCell(4)
        wninsd!(window, 0., 100.)
        gfrepi(window, "x", "y")
        # BEGMSS or ENDMSS empty, too long, or containing non-printing characters
        @test_throws SpiceError gfrepi(window, "", "y")
        @test_throws SpiceError gfrepi(window, "x", "")
        @test_throws SpiceError gfrepi(window, "x"^1000, "y")
        @test_throws SpiceError gfrepi(window, "x", "y"^1000)
        @test_throws SpiceError gfrepi(window, "y\n", "y")
        @test_throws SpiceError gfrepi(window, "x", "y\n")
        gfrepf()
    end
    @testset "gfrepu" begin
        window = SpiceDoubleCell(4)
        wninsd!(window, 0., 100.)
        gfrepi(window, "x", "y")
        gfrepu(0., 100., 50.)
        gfrepu(0., 100., 100.)
        @test_throws SpiceError gfrepu(100., 0., 100.)
        @test_throws SpiceError gfrepu(0., 100., -1.)
        @test_throws SpiceError gfrepu(0., 100., 1011.)
        gfrepu(0., 100., 100.)
        gfrepf()
    end
    @testset "gfrfov" begin
        try
            furnsh(path(CORE, :lsk),
                   path(CORE, :pck),
                   path(CORE, :spk),
                   path(CASSINI, :ck),
                   path(CASSINI, :fk),
                   path(CASSINI, :ik),
                   path(CASSINI, :pck),
                   path(CASSINI, :sclk),
                   path(CASSINI, :tour_spk),
                   path(CASSINI, :sat_spk))
            inst = "CASSINI_ISS_WAC"
            et_start1 = str2et("2013-FEB-25 07:20:00.000")
            et_end1 = str2et("2013-FEB-25 11:45:00.000")
            et_start2 = str2et("2013-FEB-25 11:55:00.000")
            et_end2 = str2et("2013-FEB-26 14:25:00.000")
            cnfine = SpiceDoubleCell(4)
            wninsd!(cnfine, et_start1, et_end1)
            wninsd!(cnfine, et_start2, et_end2)
            et_nom = str2et("2013-FEB-25 11:50:00.000")
            raydir, lt = spkpos("Enceladus", et_nom, "J2000", "NONE", "Cassini")
            result = gfrfov(inst, raydir, "J2000", "NONE", "Cassini", 10.0, cnfine)
            @test length(result) == 4
            s_timout = "YYYY-MON-DD HR:MN:SC UTC ::RND"
            @test timout(result[1], s_timout) == "2013-FEB-25 11:26:46 UTC"
            @test timout(result[2], s_timout) == "2013-FEB-25 11:45:00 UTC"
            @test timout(result[3], s_timout) == "2013-FEB-25 11:55:00 UTC"
            @test timout(result[4], s_timout) == "2013-FEB-25 12:05:33 UTC"
        finally
            kclear()
        end
    end
    @testset "gfrr" begin
        relate = ["=", "<", ">", "LOCMIN", "ABSMIN", "LOCMAX", "ABSMAX"]
        expected = Dict("=" => ["2007-JAN-02 00:35:19.583",
                                "2007-JAN-02 00:35:19.583",
                                "2007-JAN-19 22:04:54.905",
                                "2007-JAN-19 22:04:54.905",
                                "2007-FEB-01 23:30:13.439",
                                "2007-FEB-01 23:30:13.439",
                                "2007-FEB-17 11:10:46.547",
                                "2007-FEB-17 11:10:46.547",
                                "2007-MAR-04 15:50:19.940",
                                "2007-MAR-04 15:50:19.940",
                                "2007-MAR-18 09:59:05.966",
                                "2007-MAR-18 09:59:05.966"],
                        "<" => ["2007-JAN-02 00:35:19.583",
                                "2007-JAN-19 22:04:54.905",
                                "2007-FEB-01 23:30:13.439",
                                "2007-FEB-17 11:10:46.547",
                                "2007-MAR-04 15:50:19.940",
                                "2007-MAR-18 09:59:05.966"],
                        ">" => ["2007-JAN-01 00:00:00.000",
                                "2007-JAN-02 00:35:19.583",
                                "2007-JAN-19 22:04:54.905",
                                "2007-FEB-01 23:30:13.439",
                                "2007-FEB-17 11:10:46.547",
                                "2007-MAR-04 15:50:19.940",
                                "2007-MAR-18 09:59:05.966",
                                "2007-APR-01 00:00:00.000"],
                        "LOCMIN" => ["2007-JAN-11 07:03:59.001",
                                     "2007-JAN-11 07:03:59.001",
                                     "2007-FEB-10 06:26:15.451",
                                     "2007-FEB-10 06:26:15.451",
                                     "2007-MAR-12 03:28:36.414",
                                     "2007-MAR-12 03:28:36.414"],
                        "ABSMIN" => ["2007-JAN-11 07:03:59.001",
                                     "2007-JAN-11 07:03:59.001"],
                        "LOCMAX" => ["2007-JAN-26 02:27:33.772",
                                     "2007-JAN-26 02:27:33.772",
                                     "2007-FEB-24 09:35:07.822",
                                     "2007-FEB-24 09:35:07.822",
                                     "2007-MAR-25 17:26:56.158",
                                     "2007-MAR-25 17:26:56.158"],
                        "ABSMAX" => ["2007-MAR-25 17:26:56.158",
                                     "2007-MAR-25 17:26:56.158"])
        try
            furnsh(path(CORE, :lsk), path(CORE, :pck), path(CORE, :spk))
            et0 = str2et("2007 JAN 01")
            et1 = str2et("2007 APR 01")
            cnfine = SpiceDoubleCell(2)
            wninsd!(cnfine, et0, et1)
            for relation in relate
                result = gfrr("MOON", "NONE", "SUN", relation, 0.3365, 0.0, spd(), 2000, cnfine)
                count = wncard(result)
                if count > 0
                    results = String[]
                    for i in 1:count
                        left, right = wnfetd(result, i)
                        timstr_left = timout(left, "YYYY-MON-DD HR:MN:SC.###", 41)
                        timstr_right = timout(right, "YYYY-MON-DD HR:MN:SC.###", 41)
                        push!(results, timstr_left)
                        push!(results, timstr_right)
                    end
                    @test results == expected[relation]
                end
            end
        finally
            kclear()
        end
    end
    @testset "gfsep" begin
        try
            furnsh(path(CORE, :lsk), path(CORE, :pck), path(CORE, :spk))
            expected = ["2007-JAN-03 14:20:24.628017 (TDB)",
                        "2007-FEB-02 06:16:24.111794 (TDB)",
                        "2007-MAR-03 23:22:42.005064 (TDB)",
                        "2007-APR-02 16:49:16.145506 (TDB)",
                        "2007-MAY-02 09:41:43.840096 (TDB)",
                        "2007-JUN-01 01:03:44.537483 (TDB)",
                        "2007-JUN-30 14:15:26.586223 (TDB)",
                        "2007-JUL-30 01:14:49.010797 (TDB)",
                        "2007-AUG-28 10:39:01.398087 (TDB)",
                        "2007-SEP-26 19:25:51.519413 (TDB)",
                        "2007-OCT-26 04:30:56.635336 (TDB)",
                        "2007-NOV-24 14:31:04.341632 (TDB)",
                        "2007-DEC-24 01:40:12.245932 (TDB)"]
            et0 = str2et("2007 JAN 01")
            et1 = str2et("2008 JAN 01")
            cnfine = SpiceDoubleCell(2)
            wninsd!(cnfine, et0, et1)
            result = gfsep("MOON", "SPHERE", "NULL", "SUN", "SPHERE", "NULL", "NONE", "EARTH",
                           "LOCMAX", 0.0, 0.0, 6.0 * spd(), 1000, cnfine)
            count = wncard(result)
            @test count == 13
            results = []
        for i in 1:count
                start, stop = wnfetd(result, i)
                @test start == stop
                push!(results, timout(start, "YYYY-MON-DD HR:MN:SC.###### (TDB) ::TDB ::RND", 41))
            end
            @test results == expected
        finally
            kclear()
        end
    end
    @testset "gfsntc" begin
        kclear()
        kernel = tempname()
        open(kernel, "w") do kernel_file
            write(kernel_file, "\\begindata\n")
            write(kernel_file, "FRAME_SEM                     =  10100000\n")
            write(kernel_file, "FRAME_10100000_NAME           = 'SEM'\n")
            write(kernel_file, "FRAME_10100000_CLASS          =  5\n")
            write(kernel_file, "FRAME_10100000_CLASS_ID       =  10100000\n")
            write(kernel_file, "FRAME_10100000_CENTER         =  10\n")
            write(kernel_file, "FRAME_10100000_RELATIVE       = 'J2000'\n")
            write(kernel_file, "FRAME_10100000_DEF_STYLE      = 'PARAMETERIZED'\n")
            write(kernel_file, "FRAME_10100000_FAMILY         = 'TWO-VECTOR'\n")
            write(kernel_file, "FRAME_10100000_PRI_AXIS       = 'X'\n")
            write(kernel_file, "FRAME_10100000_PRI_VECTOR_DEF = 'OBSERVER_TARGET_POSITION'\n")
            write(kernel_file, "FRAME_10100000_PRI_OBSERVER   = 'SUN'\n")
            write(kernel_file, "FRAME_10100000_PRI_TARGET     = 'EARTH'\n")
            write(kernel_file, "FRAME_10100000_PRI_ABCORR     = 'NONE'\n")
            write(kernel_file, "FRAME_10100000_SEC_AXIS       = 'Y'\n")
            write(kernel_file, "FRAME_10100000_SEC_VECTOR_DEF = 'OBSERVER_TARGET_VELOCITY'\n")
            write(kernel_file, "FRAME_10100000_SEC_OBSERVER   = 'SUN'\n")
            write(kernel_file, "FRAME_10100000_SEC_TARGET     = 'EARTH'\n")
            write(kernel_file, "FRAME_10100000_SEC_ABCORR     = 'NONE'\n")
            write(kernel_file, "FRAME_10100000_SEC_FRAME      = 'J2000'\n")
        end
        try
            furnsh(path(CORE, :lsk), path(CORE, :pck), path(CORE, :spk))
            furnsh(kernel)
            et0 = str2et("2007 JAN 01")
            et1 = str2et("2008 JAN 01")
            cnfine = SpiceDoubleCell(2)
            wninsd!(cnfine, et0, et1)
            result = gfsntc("EARTH", "IAU_EARTH", "Ellipsoid", "NONE", "SUN", "SEM", [1.0, 0.0, 0.0],
                            "LATITUDINAL", "LATITUDE", "=", 0.0, 0.0, 90.0 * spd(), 1000, cnfine)
            count = wncard(result)
            @test count > 0
            beg, stop = wnfetd(result, 1)
            begstr = timout(beg, "YYYY-MON-DD HR:MN:SC.###### (TDB) ::TDB ::RND", 80)
            endstr = timout(stop, "YYYY-MON-DD HR:MN:SC.###### (TDB) ::TDB ::RND", 80)
            @test begstr == "2007-MAR-21 00:01:25.527303 (TDB)"
            @test endstr == "2007-MAR-21 00:01:25.527303 (TDB)"
            beg, stop = wnfetd(result, 2)
            begstr = timout(beg, "YYYY-MON-DD HR:MN:SC.###### (TDB) ::TDB ::RND", 80)
            endstr = timout(stop, "YYYY-MON-DD HR:MN:SC.###### (TDB) ::TDB ::RND", 80)
            @test startswith(begstr, "2007-SEP-23 09:46:39.60698")
            @test startswith(endstr, "2007-SEP-23 09:46:39.60698")
        finally
            kclear()
        end
    end
    @testset "gfstep/gfsstp" begin
        gfsstp(0.5)
        @test gfstep(0.5) == 0.5
    end
    @testset "gfsubc" begin
        try
            furnsh(path(CORE, :lsk), path(CORE, :pck), path(CORE, :spk))
            et0 = str2et("2007 JAN 01")
            et1 = str2et("2008 JAN 01")
            cnfine = SpiceDoubleCell(2)
            wninsd!(cnfine, et0, et1)
            result = gfsubc("earth", "iau_earth", "Near point: ellipsoid", "none", "sun",
                            "geodetic", "latitude", ">", deg2rad(16.0), 0.0, spd() * 90.0, 1000,
                            cnfine)
            count = wncard(result)
            @test count > 0
            start, stop = wnfetd(result, 1)
            start_time = timout(start, "YYYY-MON-DD HR:MN:SC.###### (TDB) ::TDB ::RND", 41)
            end_time = timout(stop, "YYYY-MON-DD HR:MN:SC.###### (TDB) ::TDB ::RND", 41)
            @test start_time == "2007-MAY-04 17:08:56.724320 (TDB)"
            @test end_time == "2007-AUG-09 01:51:29.307830 (TDB)"
        finally
            kclear()
        end
    end
    @testset "gftfov" begin
        try
            furnsh(path(CORE, :lsk), path(CORE, :pck), path(CORE, :spk),
                   path(CASSINI, :ck),
                   path(CASSINI, :fk),
                   path(CASSINI, :ik),
                   path(CASSINI, :pck),
                   path(CASSINI, :sclk),
                   path(CASSINI, :tour_spk),
                   path(CASSINI, :sat_spk))
            et_start1 = str2et("2013-FEB-25 07:20:00.000")
            et_end1 = str2et("2013-FEB-25 11:45:00.000")
            et_start2 = str2et("2013-FEB-25 11:55:00.000")
            et_end2 = str2et("2013-FEB-26 14:25:00.000")
            cnfine = SpiceDoubleCell(4)
            wninsd!(cnfine, et_start1, et_end1)
            wninsd!(cnfine, et_start2, et_end2)
            result = gftfov("CASSINI_ISS_NAC", "ENCELADUS", "ELLIPSOID", "IAU_ENCELADUS", "LT", "CASSINI", 10.0, 5000, cnfine)
            @test card(result) == 4
            s_timout = "YYYY-MON-DD HR:MN:SC UTC ::RND"
            @test timout(result[1], s_timout) == "2013-FEB-25 10:42:33 UTC"
            @test timout(result[2], s_timout) == "2013-FEB-25 11:45:00 UTC"
            @test timout(result[3], s_timout) == "2013-FEB-25 11:55:00 UTC"
            @test timout(result[4], s_timout) == "2013-FEB-25 12:04:30 UTC"
        finally
            kclear()
        end
    end
    @testset "gfudb" begin
        try
            furnsh(path(CORE, :lsk), path(CORE, :pck), path(CORE, :spk))
            et_start = str2et("Jan 1 2001")
            et_end = str2et("Jan 1 2002")
            result = SpiceDoubleCell(40000)
            cnfine = SpiceDoubleCell(2)
            wninsd!(cnfine, et_start, et_end)
            step = 5.0 * spd()

            function gfq(_, et)
                state, lt = spkez(301, et, "IAU_EARTH", "NONE", 399)
                state[3] >= 0.0 && state[6] > 0.0
            end

            gfudb!(_ -> 0.0, gfq, step, cnfine, result)
            @test length(result) > 20 # true value is 28
        finally
            kclear()
        end
    end
    @testset "gfudb2" begin
        try
            furnsh(path(CORE, :lsk), path(CORE, :pck), path(CORE, :spk))
            et_start = str2et("Jan 1 2001")
            et_end = str2et("Jan 1 2002")
            result = SpiceDoubleCell(40000)
            cnfine = SpiceDoubleCell(2)
            wninsd!(cnfine, et_start, et_end)
            step = 60.0 * 60.0

            function gfq(et)
                pos, _ = spkezp(301, et, "IAU_EARTH", "NONE", 399)
                pos[3]
            end

            function gfb(udfuns, et)
                value = udfuns(et)
                -1000.0 <= value <= 1000.0
            end

            gfudb!(gfq, gfb, step, cnfine, result)
            @test length(result) > 50  # true value is 56
        finally
            kclear()
        end
    end
    @testset "gfuds" begin
        try
            furnsh(path(CORE, :lsk), path(CORE, :pck), path(CORE, :spk))
            relations = ["=", "<", ">", "LOCMIN", "ABSMIN", "LOCMAX", "ABSMAX"]
            et_start = str2et("Jan 1 2007")
            et_end = str2et("Apr 1 2007")
            step = spd()
            adjust = 0.0
            refval  = 0.3365

            function gfq(et)
                state, _ = spkez(301, et, "J2000", "NONE", 10)
                dvnorm(state)
            end

            function gfdecrx(udfuns, et)
                h = 1.0
                dx = uddf(udfuns, et, h)
                dx < 0.0
            end

            @testset for r in relations
                result = SpiceDoubleCell(40000)
                cnfine = SpiceDoubleCell(2)
                wninsd!(cnfine, et_start, et_end)
                result = gfuds!(gfq, gfdecrx, r, refval, adjust, step, 20000, cnfine, result)
                @test length(result) > 0
            end
        finally
            kclear()
        end
    end
    @testset "gipool" begin
        try
            data = collect(1:10)
            pipool("pipool_array", data)
            ivals = gipool("pipool_array")
            @test data == ivals
        finally
            kclear()
        end
    end
    @testset "gnpool" begin
        try
            furnsh(path(CORE, :pck))
            var = "BODY599*"
            index = 1
            room = 10
            expected = ["BODY599_POLE_DEC", "BODY599_LONG_AXIS", "BODY599_PM", "BODY599_RADII",
                        "BODY599_POLE_RA", "BODY599_NUT_PREC_PM", "BODY599_NUT_PREC_DEC",
                        "BODY599_NUT_PREC_RA"]
            kervar = gnpool(var, index, room)
            @test Set(expected) == Set(kervar)
            @test gnpool("BLOB", index, room) === nothing
        finally
            kclear()
        end
    end
end


# MIT License
#
# Copyright (c) 2018 Martin Biel
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

@reexport module ProgressiveHedging

# Standard library
using LinearAlgebra
using SparseArrays
using Distributed
using Printf

# External libraries
using Parameters
using JuMP
using MathOptInterface
using StochasticPrograms
using StochasticPrograms: AcceptableTermination
using StochasticPrograms: UnspecifiedInstantiation, ScenarioDecompositionStructure, AbstractScenarioProblems, ScenarioProblems, DistributedScenarioProblems
using StochasticPrograms: AbstractExecution, ExecutionParameter, Serial, Synchronous, Asynchronous
using StochasticPrograms: AbstractStructuredOptimizer, set_master_optimizer!, set_subproblem_optimizer!
using StochasticPrograms: DecisionMap, get_decisions, all_decisions, remove_decision!, set_decision!, set_stage!, SingleDecisionSet, DecisionCoefficientChange, KnownDecision, NoSpecifiedConstraint
using StochasticPrograms: add_subscript
using StochasticPrograms: AbstractPenaltyTerm, Quadratic, InfNorm, ManhattanNorm, initialize_penaltyterm!, update_penaltyterm!, remove_penalty!
using ProgressMeter

import Base: show, put!, wait, isready, take!, fetch
import StochasticPrograms: supports_structure, num_iterations, default_structure, load_structure!, restore_structure!, optimize!, optimizer_name, master_optimizer, subproblem_optimizer, num_subproblems

const MOI = MathOptInterface
const MOIU = MOI.Utilities

export
    AbstractProgressiveHedgingAttribute,
    ProgressiveHedgingAlgorithm,
    PrimalTolerance,
    DualTolerance,
    Penalizer,
    PenalizationParameter,
    RawPenalizationParameter,
    set_penalization_attribute,
    set_penalization_attributes,
    PenaltyTerm,
    Fixed,
    Adaptive

# Include files
include("types/types.jl")
include("penalties/penalization.jl")
include("execution/execution.jl")
include("solver.jl")
include("MOI_wrapper.jl")

end # module


abstract type TTM_Potential end
include("constants.jl")

@inline function compute_M_site_crd!(ttm::TTM_Potential, coords::AbstractMatrix{Float64})
    """
    Takes the 3xN coordinates you get coming in and stores the coordinates with M site in the TTM_Potential struct.
    """
    num_M_sites::Int = 0
    for i in 1:Int(size(coords, 2) / 3) * 4
        if i % 4 == 0
            @inbounds ttm.M_site_coords[i] = convert(SVector{3, Float64}, ttm.constants.γ_1 * coords[:,i-3 - num_M_sites] + ttm.constants.γ_2 * (coords[:,i-2 - num_M_sites] + coords[:, i-1 - num_M_sites]))
            num_M_sites += 1
        else
            @inbounds ttm.M_site_coords[i] = convert(SVector{3, Float64}, coords[:, (i - num_M_sites)])
        end
    end
end

@inline function triu_linear_index(i::Int, j::Int, n::Int)
    return Int(n*(n-1)/2 - (n-i+1)*(n-i)/2 + j - i)
end

function distance_vectors(coords::Vector{SVector{3, Float64}})
    # make this take the storage as most time will be spent allocating
    n::Int = length(coords)
    distances = [@SVector zeros(3) for _ in 1:(n*(n-1)/2)]
    for i in 1:n-1
        for j in i+1:n
            distances[triu_linear_index(i,j,n)] = coords[i]- coords[j]
        end
    end
    return distances
end

function get_dipoles(ttm::TTM_Potential)
    dipoles = zeros(3, ttm.num_waters)
    for i in 1:length(ttm.M_site_coords)
        if i % 4 != 1 # static dipoles from Partridge-Schwenke
            dipoles[:, (i-1)÷4+1] += ttm.elec_data.q[i] * ttm.M_site_coords[i]
        else # add in the induced dipoles
            dipoles[:, i÷4+1] += ttm.elec_data.dipoles[i]   # O
            dipoles[:, i÷4+1] += ttm.elec_data.dipoles[i+1] # H
            dipoles[:, i÷4+1] += ttm.elec_data.dipoles[i+2] # H
            dipoles[:, i÷4+1] += ttm.elec_data.dipoles[i+3] # M
        end
    end
    return dipoles / chargecon() * debye()
end

function get_dipole_magnitudes(ttm::TTM_Potential)
    return norm.(eachcol(get_dipoles(ttm)))
end

function PS_energies_and_gradients!(ttm::TTM_Potential, coords::AbstractMatrix{Float64}, grads_q::Vector{MArray{Tuple{3, 3, 3}, Float64, 3, 27}}, grads::Union{Matrix{Float64}, Nothing}=nothing, ttm3f::Bool=false)
    E_int::Float64 = 0.0
    q3 = @MVector zeros(3)
    dq3 = @MArray zeros(3,3,3)
    for n in 1:ttm.num_waters

        if (grads !== nothing)
            @inbounds E_int += pot_nasa(convert(SMatrix{3, 3, Float64}, @view(coords[:, (3*n-2):3*n])), @view(grads[:, (3*n-2):3*n]))
            @inbounds dms_nasa!(convert(SMatrix{3, 3, Float64}, @view(coords[:, (3*n-2):3*n])), q3, dq3, ttm3f)
        else
            @inbounds E_int += pot_nasa(convert(SMatrix{3, 3, Float64}, @view(coords[:, (3*n-2):3*n])), nothing)
            @inbounds dms_nasa!(convert(SMatrix{3, 3, Float64}, @view(coords[:, (3*n-2):3*n])), q3, nothing, ttm3f)
        end

        # TTM_Potential2.1-F assignment
        tmp::Float64 = 0.5 * ttm.constants.γ_M / (1.0 - ttm.constants.γ_M)
        @inbounds ttm.elec_data.q[4*(n-1) + 1] = 0.0                                  # O
        @inbounds ttm.elec_data.q[4*(n-1) + 2] = q3[2] + tmp * (q3[2] + q3[3])        # H1
        @inbounds ttm.elec_data.q[4*(n-1) + 3] = q3[3] + tmp * (q3[2] + q3[3])        # H2
        @inbounds ttm.elec_data.q[4*(n-1) + 4] = q3[1] / (1.0 - ttm.constants.γ_M) # M

        if (grads === nothing)
            continue
        end
        # charge gradients
        @inbounds grads_q[n][:, 1, :] = dq3[:, 1, :] +   tmp*(dq3[:, 1, :] + dq3[:, 2, :])
        @inbounds grads_q[n][:, 2, :] = dq3[:, 2, :] +   tmp*(dq3[:, 2, :] + dq3[:, 1, :])
        @inbounds grads_q[n][:, 3, :] = dq3[:, 3, :] - 2*tmp*(dq3[:, 1, :] + dq3[:, 2, :])
    end

    return E_int
end

function add_electrostatic_and_dms_derivatives!(ttm::TTM_Potential, grads::AbstractMatrix{Float64}, grads_E::Vector{MVector{3, Float64}}, grads_q::Vector{MArray{Tuple{3, 3, 3}, Float64, 3, 27}})
    # add electrostatic derivatives
    for n in 1:ttm.num_waters
        # add O H H gradients from PS and electrostatics (which has M site)
        for k in 1:3
            @inbounds @views grads[:, 3 * (n-1) + k] += grads_E[4 * (n-1) + k]
        end

        # redistribute the M-site derivatives
        @inbounds @views grads[:, 3 * (n-1) + 1] += (1.0 - ttm.constants.γ_M) * grads_E[4*(n-1) + 4] # O
        @inbounds @views grads[:, 3 * (n-1) + 2] +=  0.5 * ttm.constants.γ_M  * grads_E[4*(n-1) + 4] # H
        @inbounds @views grads[:, 3 * (n-1) + 3] +=  0.5 * ttm.constants.γ_M  * grads_E[4*(n-1) + 4] # H
    end

    # derivatives from the adjustable charges of the NASA's PES
    for n in 1:ttm.num_waters
        i_O  = 3 * (n-1) + 1
        i_H1 = 3 * (n-1) + 2
        i_H2 = 3 * (n-1) + 3

        @inbounds grads[:, i_H1] += (grads_q[n][:, 1, 1]*ttm.elec_data.ϕ[4*(n-1) + 2]   # phi(h1)
                         + grads_q[n][:, 2, 1]*ttm.elec_data.ϕ[4*(n-1) + 3]   # phi(h2)
                         + grads_q[n][:, 3, 1]*ttm.elec_data.ϕ[4*(n-1) + 4])  # phi(M)

        @inbounds grads[:, i_H2] += (grads_q[n][:, 1, 2]*ttm.elec_data.ϕ[4*(n-1) + 2]   # phi(h1)
                         + grads_q[n][:, 2, 2]*ttm.elec_data.ϕ[4*(n-1) + 3]   # phi(h2)
                         + grads_q[n][:, 3, 2]*ttm.elec_data.ϕ[4*(n-1) + 4])  # phi(M)

        @inbounds grads[:, i_O]  += (grads_q[n][:, 1, 3]*ttm.elec_data.ϕ[4*(n-1) + 2]   # phi(h1)
                         + grads_q[n][:, 2, 3]*ttm.elec_data.ϕ[4*(n-1) + 3]   # phi(h2)
                         + grads_q[n][:, 3, 3]*ttm.elec_data.ϕ[4*(n-1) + 4])  # phi(M)
    end
end


struct QuickMDP{ID,S,A,D<:NamedTuple} <: MDP{S,A}
    data::D
end

"""
    QuickMDP(gen::Function, [id]; kwargs...)

Construct a generative MDP model with the function `gen` and keyword arguments.

`gen` should take three arguments: a state, an action, and a random number generator. It should return a `NamedTuple` with keys `sp` for the next state and `r` for the reward.

Keywords can be static objects or functions. See the QuickPOMDPs.jl documentation for more information.
"""
QuickMDP(gen::Function, id=uuid4(); kwargs...) = QuickMDP(id; gen=gen, kwargs...)

"""
    QuickMDP([id]; kwargs...)

Construct an MDP model with keyword arguments. Keywords can be static objects or functions. See the QuickPOMDPs.jl documentation for more information.
"""
function QuickMDP(id=uuid4(); kwargs...)
    kwd = Dict{Symbol, Any}(kwargs)

    for (k, v) in pairs(kwd)
        kwd[k] = preprocess(v)
    end
    quick_defaults!(kwd)

    S = infer_statetype(kwd)
    A = infer_actiontype(kwd)
    d = namedtuple(keys(kwd)...)(values(kwd)...)
    qm = QuickMDP{id, S, A, typeof(d)}(d)
    return qm
end

id(::QuickMDP{ID}) where ID = ID

struct QuickPOMDP{ID,S,A,O,D<:NamedTuple} <: POMDP{S,A,O}
    data::D
end

"""
    QuickPOMDP(gen::Function, [id]; kwargs...)

Construct a generative POMDP model with the function `gen` and keyword arguments.

`gen` should take three arguments: a state, an action, and a random number generator. It should return a `NamedTuple` with keys `sp` for the next state, `o` for the observation, and `r` for the reward.

Keywords can be static objects or functions. See the QuickPOMDPs.jl documentation for more information.
"""
QuickPOMDP(gen::Function, id=uuid4(); kwargs...) = QuickPOMDP(id; gen=gen, kwargs...)

"""
    QuickPOMDP([id]; kwargs...)

Construct an POMDP model with keyword arguments. Keywords can be static objects or functions. See the QuickPOMDPs.jl documentation for more information.
"""
function QuickPOMDP(id=uuid4(); kwargs...)
    kwd = Dict{Symbol, Any}(kwargs)

    for (k, v) in pairs(kwd)
        kwd[k] = preprocess(v)
    end
    quick_defaults!(kwd)
    quick_warnings(kwd)

    S = infer_statetype(kwd)
    A = infer_actiontype(kwd)
    O = infer_obstype(kwd)
    d = namedtuple(keys(kwd)...)(values(kwd)...)
    qm = QuickPOMDP{id, S, A, O, typeof(d)}(d)
    return qm
end

id(::QuickPOMDP{ID}) where ID = ID

const QuickModel = Union{QuickMDP, QuickPOMDP}

"Function that is called on each keyword argument before anything else is done. This was designed as a hook to allow other packages to handle PyObjects."
preprocess(x) = x

function quick_defaults!(kwd::Dict)
    kwd[:discount] = get(kwd, :discount, 1.0)
    kwd[:isterminal] = get(kwd, :isterminal, false)

    if !haskey(kwd, :stateindex)
        if haskey(kwd, :states)
            states = _call(Val(:states), kwd[:states], ())
            if hasmethod(length, typeof((states,))) && length(states) < Inf
                kwd[:stateindex] = Dict(s=>i for (i,s) in enumerate(states))
            end
        end
    end

    if !haskey(kwd, :actionindex)
        if haskey(kwd, :actions)
            actions = _call(Val(:actions), kwd[:actions], ())
            if hasmethod(length, typeof((actions,))) && length(actions) < Inf
                kwd[:actionindex] = Dict(s=>i for (i,s) in enumerate(actions))
            end
        end
    end

    if !haskey(kwd, :obsindex)
        if haskey(kwd, :observations)
            observations = _call(Val(:observations), kwd[:observations], ())
            if hasmethod(length, typeof((observations,))) && length(observations) < Inf
                kwd[:obsindex] = Dict(s=>i for (i,s) in enumerate(observations))
            end
        end
    end
end

function quick_warnings(kwd)
    if haskey(kwd, :initialstate)
        isd = _call(Val(:initialstate), kwd[:initialstate], ())
        try rand(MersenneTwister(0), isd)
        catch ex
            if ex isa MethodError || ex isa ArgumentError
                @warn("Unable to call rand(rng, $isd). Is the `initialstate` that you supplied a distribution?")
            else
                rethrow(ex)
            end
        end
    end

    if haskey(kwd, :reward) && !(kwd[:reward] isa Function)
        @warn("`reward` must be a function; got $(kwd[:reward])")
    end
end

function infer_statetype(kwd)
    if haskey(kwd, :statetype)
        st = _call(Val(:statetype), kwd[:statetype], (), NamedTuple())
    elseif haskey(kwd, :states)
        st = eltype(_call(Val(:states), kwd[:states], (), NamedTuple()))
    elseif haskey(kwd, :initialstate)
        st = typeof(rand(MersenneTwister(0), _call(Val(:initialstate), kwd[:initialstate], (), NamedTuple())))
    else
        st = Any
    end
    if st == Any
        @warn("Unable to infer state type for a Quick(PO)MDP; using Any. This may have significant performance consequences. Use the statetype keyword argument to specify a concrete state type.")
    end
    return st
end

function infer_actiontype(kwd)
    if haskey(kwd, :actiontype)
        at = _call(Val(:actiontype), kwd[:actiontype], (), NamedTuple())
    elseif haskey(kwd, :actions)
        at = eltype(_call(Val(:actions), kwd[:actions], (), NamedTuple()))
    else
        at = Any
    end
    if at == Any
        @warn("Unable to infer action type for a Quick(PO)MDP; using Any. This may have significant performance consequences. Use the actiontype keyword argument to specify a concrete action type.")
    end
    return at
end

function infer_obstype(kwd)
    if haskey(kwd, :obstype)
        ot = _call(Val(:obstype), kwd[:obstype], (), NamedTuple())
    elseif haskey(kwd, :observations)
        ot = eltype(_call(Val(:observations), kwd[:observations], (), NamedTuple()))
    elseif haskey(kwd, :initialobs) && haskey(kwd, :initialstate)
        s0 = rand(MersenneTwister(0), _call(Val(:initialstate), kwd[:initialstate], (), NamedTuple()))
        ot = typeof(rand(MersenneTwister(0), _call(Val(:initialobs), kwd[:initialobs], (s0,), NamedTuple())))
    else
        ot = Any
    end
    if ot == Any
        @warn("Unable to infer observation type for a QuickPOMDP; using Any. This may have significant performance consequences. Use the obstype keyword argument to specify a concrete observation type.")
    end
    return ot
end

function _call(namev::Val{name}, m::QuickModel, args, kwargs=NamedTuple()) where name
    _call(namev,
          get(m.data, name) do
              throw(MissingQuickArgument(m, name))
          end,
          args,
          kwargs)
end

_call(::Val, f::Function, args, kwargs=NamedTuple()) = f(args...; kwargs...)
_call(v::Val, object, args, kwargs=NamedTuple()) = object
_call(v::Val, d::Dict, args, kwargs=NamedTuple()) = d[args...]

macro forward_to_data(f)
    @assert f.head == :. "@forward_to_data must be used with a module-qualified function expression, e.g. @forward_to_data POMDPs.discount"
    quote
        $f(m::QuickModel, args...; kwargs...) = _call(Val($(f.args[2])), m, args, kwargs)
    end
end

function POMDPs.transition(m::QuickModel, s, a)
    if haskey(m.data, :transition)
        return m.data.transition(s, a)
    else
        throw(MissingQuickArgument(m, :transition, types=[Function], also=[:gen]))
    end
end

function POMDPs.observation(m::QuickPOMDP, args...)
    if haskey(m.data, :observation)
        obs = m.data[:observation]
        if static_hasmethod(obs, typeof(args))
            return obs(args...)
        elseif length(args) == 3 && static_hasmethod(obs, typeof(args[2:3]))
            return obs(args[2:3]...)
        else
            return obs(args...)
        end
        return m.data.observation(args...)
    else
        throw(MissingQuickArgument(m, :observation, types=[Function], also=[:gen]))
    end
end

function POMDPs.reward(m::QuickModel, args...)
    if haskey(m.data, :reward)
        r = m.data[:reward]
        if static_hasmethod(r, typeof(args)) # static_hasmethod could cause issues, but I think it is worth doing in this single spot
            return r(args...)
        elseif m isa POMDP && length(args) == 4
            if static_hasmethod(r, typeof(args[1:3])) # (s, a, sp, o) -> (s, a, sp)
                return r(args[1:3]...)
            elseif static_hasmethod(r, typeof(args[1:2])) # (s, a, sp, o) -> (s, a)
                return r(args[1:2]...)
            end
        elseif length(args) == 3 && static_hasmethod(r, typeof(args[1:2])) # (s, a, sp) -> (s, a)
            return r(args[1:2]...)
        else
            return r(args...)
        end
    else
        throw(MissingQuickArgument(m, :reward))
    end
end

@forward_to_data POMDPs.initialstate
@forward_to_data POMDPs.initialobs

function POMDPs.gen(m::QuickModel, s, a, rng)
    if haskey(m.data, :gen)
        return m.data.gen(s, a, rng)
    else
        return NamedTuple()
    end
end

@forward_to_data POMDPs.states
@forward_to_data POMDPs.actions
@forward_to_data POMDPs.observations
@forward_to_data POMDPs.discount

@forward_to_data POMDPs.stateindex
@forward_to_data POMDPs.actionindex
@forward_to_data POMDPs.obsindex
@forward_to_data POMDPs.isterminal

function POMDPModelTools.obs_weight(m::QuickPOMDP, args...)
    if haskey(m.data, :obs_weight)
        return _call(Val(:obs_weight), m, args)
    elseif haskey(m.data, :observation)
        return pdf(_call(Val(:observation), m, args[1:end-1]), args[end])
    else
        throw(MissingQuickArgument(m, :obs_weight, types[Function], also=[:observation]))
    end
end
@forward_to_data POMDPModelTools.render

function POMDPModelTools.StateActionReward(m::Union{QuickPOMDP,QuickMDP})
    if hasmethod(m.data[:reward], Tuple{statetype(m), actiontype(m)})
        return FunctionSAR(m)
    else
        return LazyCachedSAR(m)
    end
end


const RTR = ReverseDiff.TrackedReal
const RTV = ReverseDiff.TrackedVector
const RTM = ReverseDiff.TrackedMatrix
const RTA = ReverseDiff.TrackedArray
using ReverseDiff: SpecialInstruction
import NaNMath
using ForwardDiff: Dual
import SpecialFunctions: logbeta
import Distributions: Gamma

Base.:*(A::Adjoint{<:Real, <:RTV{<:Real}}, B::AbstractVector{<:Real}) = dot(A, B)
Base.:*(A::Adjoint{<:Real, <:RTV{<:Real}}, B::RTV{<:Real}) = dot(A, B)
Base.:*(A::AbstractVector{<:Real}, B::Adjoint{<:Real, <:RTV{<:Real}}) = dot(A, B)
Base.:*(A::RTV{<:Real}, B::Adjoint{<:Real, <:RTV{<:Real}}) = dot(A, B)

Gamma(α::RTR, θ::Real; check_args=true) = pgamma(α, θ, check_args = check_args)
Gamma(α::Real, θ::RTR; check_args=true) = pgamma(α, θ, check_args = check_args)
Gamma(α::RTR, θ::RTR; check_args=true) = pgamma(α, θ, check_args = check_args)
pgamma(α, θ; check_args=true) = Gamma(promote(α, θ)..., check_args = check_args)
Gamma(α::T; check_args=true) where {T <: RTR} = Gamma(α, one(T), check_args = check_args)
function Gamma(α::T, θ::T; check_args=true) where {T <: RTR}
    check_args && Distributions.@check_args(Gamma, α > zero(α) && θ > zero(θ))
    return Gamma{T}(α, θ)
end

# Work around to stop TrackedReal of Inf and -Inf from producing NaN in the derivative
function Base.minimum(d::LocationScale{T}) where {T <: RTR}
    if isfinite(minimum(d.ρ))
        return d.μ + d.σ * minimum(d.ρ)
    else
        return convert(T, ReverseDiff.@skip(minimum)(d.ρ))
    end
end
function Base.maximum(d::LocationScale{T}) where {T <: RTR}
    if isfinite(minimum(d.ρ))
        return d.μ + d.σ * maximum(d.ρ)
    else
        return convert(T, ReverseDiff.@skip(maximum)(d.ρ))
    end
end

for T in (:RTV, :RTM)
    @eval begin
        function Distributions.logpdf(d::MvNormal{<:Any, <:PDMats.ScalMat}, x::$T)
            logpdf(TuringScalMvNormal(d.μ, d.Σ.value), x)
        end
        function Distributions.logpdf(d::MvNormal{<:Any, <:PDMats.PDiagMat}, x::$T)
            logpdf(TuringDiagMvNormal(d.μ, d.Σ.diag), x)
        end
        function Distributions.logpdf(d::MvNormal{<:Any, <:PDMats.PDMat}, x::$T)
            logpdf(TuringDenseMvNormal(d.μ, d.Σ.chol), x)
        end

        function Distributions.logpdf(d::MvLogNormal{<:Any, <:PDMats.ScalMat}, x::$T)
            logpdf(TuringMvLogNormal(TuringScalMvNormal(d.normal.μ, d.normal.Σ.value)), x)
        end
        function Distributions.logpdf(d::MvLogNormal{<:Any, <:PDMats.PDiagMat}, x::$T)
            logpdf(TuringMvLogNormal(TuringDiagMvNormal(d.normal.μ, d.normal.Σ.diag)), x)
        end
        function Distributions.logpdf(d::MvLogNormal{<:Any, <:PDMats.PDMat}, x::$T)
            logpdf(TuringMvLogNormal(TuringDenseMvNormal(d.normal.μ, d.normal.Σ.chol)), x)
        end
    end
end

# zero mean, dense covariance
MvNormal(A::RTM) = TuringMvNormal(A)

# zero mean, diagonal covariance
MvNormal(σ::RTV) = TuringMvNormal(σ)

# dense mean, dense covariance
MvNormal(m::AbstractVector{<:Real}, A::RTM{<:Real}) = TuringMvNormal(m, A)
MvNormal(m::RTV{<:Real}, A::Matrix{<:Real}) = TuringMvNormal(m, A)
MvNormal(m::RTV{<:Real}, A::RTM{<:Real}) = TuringMvNormal(m, A)

# dense mean, diagonal covariance
function MvNormal(
    m::RTV{<:Real},
    D::Diagonal{<:RTR, <:RTV{<:Real}},
)
    return TuringMvNormal(m, D)
end
function MvNormal(
    m::AbstractVector{<:Real},
    D::Diagonal{<:RTR, <:RTV{<:Real}},
)
    return TuringMvNormal(m, D)
end
function MvNormal(
    m::RTV{<:Real},
    D::Diagonal{T, <:AbstractVector{T}} where {T<:Real},
)
    return TuringMvNormal(m, D)
end

# dense mean, diagonal covariance
MvNormal(m::RTV{<:Real}, σ::RTV{<:Real}) = TuringMvNormal(m, σ)
MvNormal(m::RTV{<:Real}, σ::AbstractVector{<:Real}) = TuringMvNormal(m, σ)
MvNormal(m::RTV{<:Real}, σ::Vector{<:Real}) = TuringMvNormal(m, σ)
MvNormal(m::AbstractVector{<:Real}, σ::RTV{<:Real}) = TuringMvNormal(m, σ)

# dense mean, constant variance
MvNormal(m::RTV{<:Real}, σ::RTR) = TuringMvNormal(m, σ)
MvNormal(m::RTV{<:Real}, σ::Real) = TuringMvNormal(m, σ)
MvNormal(m::AbstractVector{<:Real}, σ::RTR) = TuringMvNormal(m, σ)

# dense mean, constant variance
function MvNormal(m::RTV{<:Real}, A::UniformScaling{<:RTR})
    return TuringMvNormal(m, A)
end
function MvNormal(m::AbstractVector{<:Real}, A::UniformScaling{<:RTR})
    return TuringMvNormal(m, A)
end
function MvNormal(m::RTV{<:Real}, A::UniformScaling{<:Real})
    return TuringMvNormal(m, A)
end

# zero mean,, constant variance
MvNormal(d::Int, σ::RTR) = TuringMvNormal(d, σ)

# zero mean, dense covariance
MvLogNormal(A::RTM) = TuringMvLogNormal(TuringMvNormal(A))

# zero mean, diagonal covariance
MvLogNormal(σ::RTV) = TuringMvLogNormal(TuringMvNormal(σ))

# dense mean, dense covariance
MvLogNormal(m::RTV{<:Real}, A::RTM{<:Real}) = TuringMvLogNormal(TuringMvNormal(m, A))
MvLogNormal(m::RTV{<:Real}, A::Matrix{<:Real}) = TuringMvLogNormal(TuringMvNormal(m, A))
MvLogNormal(m::AbstractVector{<:Real}, A::RTM{<:Real}) = TuringMvLogNormal(TuringMvNormal(m, A))

# dense mean, diagonal covariance
function MvLogNormal(
    m::RTV{<:Real},
    D::Diagonal{<:RTR, <:RTV{<:Real}},
)
    return TuringMvLogNormal(TuringMvNormal(m, D))
end
function MvLogNormal(
    m::AbstractVector{<:Real},
    D::Diagonal{<:RTR, <:RTV{<:Real}},
)
    return TuringMvLogNormal(TuringMvNormal(m, D))
end
function MvLogNormal(
    m::RTV{<:Real},
    D::Diagonal{T, <:AbstractVector{T}} where {T<:Real},
)
    return TuringMvLogNormal(TuringMvNormal(m, D))
end

# dense mean, diagonal covariance
MvLogNormal(m::RTV{<:Real}, σ::RTV{<:Real}) = TuringMvLogNormal(TuringMvNormal(m, σ))
MvLogNormal(m::RTV{<:Real}, σ::AbstractVector{<:Real}) = TuringMvLogNormal(TuringMvNormal(m, σ))
MvLogNormal(m::RTV{<:Real}, σ::Vector{<:Real}) = TuringMvLogNormal(TuringMvNormal(m, σ))
MvLogNormal(m::AbstractVector{<:Real}, σ::RTV{<:Real}) = TuringMvLogNormal(TuringMvNormal(m, σ))

# dense mean, constant variance
function MvLogNormal(m::RTV{<:Real}, σ::RTR)
    return TuringMvLogNormal(TuringMvNormal(m, σ))
end
function MvLogNormal(m::RTV{<:Real}, σ::Real)
    return TuringMvLogNormal(TuringMvNormal(m, σ))
end
function MvLogNormal(m::AbstractVector{<:Real}, σ::RTR)
    return TuringMvLogNormal(TuringMvNormal(m, σ))
end

# dense mean, constant variance
function MvLogNormal(m::RTV{<:Real}, A::UniformScaling{<:RTR})
    return TuringMvLogNormal(TuringMvNormal(m, A))
end
function MvLogNormal(m::AbstractVector{<:Real}, A::UniformScaling{<:RTR})
    return TuringMvLogNormal(TuringMvNormal(m, A))
end
function MvLogNormal(m::RTV{<:Real}, A::UniformScaling{<:Real})
    return TuringMvLogNormal(TuringMvNormal(m, A))
end

# zero mean,, constant variance
MvLogNormal(d::Int, σ::RTR) = TuringMvLogNormal(TuringMvNormal(d, σ))

function LinearAlgebra.cholesky(A::RTM; check=true)
    factors, info = turing_chol(A, check)
    return Cholesky{eltype(factors), typeof(factors)}(factors, 'U', info)
end

function turing_chol(x::ReverseDiff.TrackedArray{V,D}, check) where {V,D}
    tp = ReverseDiff.tape(x)
    x_value = ReverseDiff.value(x)
    check_value = ReverseDiff.value(check)
    C, back = pullback(_turing_chol, x_value, check_value)
    out = ReverseDiff.track(C.factors, D, tp)
    ReverseDiff.record!(tp, SpecialInstruction, turing_chol, (x, check), out, (back, issuccess(C)))
    return out, C.info
end

@noinline function ReverseDiff.special_reverse_exec!(instruction::SpecialInstruction{typeof(turing_chol)})
    output = instruction.output
    instruction.cache[2] || throw(PosDefException(C.info))
    input = instruction.input
    input_deriv = ReverseDiff.deriv(input[1])
    P = instruction.cache[1]
    input_deriv .+= P((factors = ReverseDiff.deriv(output),))[1]
    ReverseDiff.unseed!(output)
    return nothing
end

@noinline function ReverseDiff.special_forward_exec!(instruction::SpecialInstruction{typeof(turing_chol)})
    output, input = instruction.output, instruction.input
    C = cholesky(ReverseDiff.value(input[1]), check = ReverseDiff.value(input[2]))
    ReverseDiff.value!(output, C.factors)
    return nothing
end

Distributions.Dirichlet(alpha::RTV) = TuringDirichlet(alpha)
Distributions.Dirichlet(d::Integer, alpha::RTR) = TuringDirichlet(d, alpha)

function Distributions.logpdf(d::MatrixBeta, X::AbstractArray{<:RTM{<:Real}})
    return mapvcat(x -> logpdf(d, x), X)
end

Distributions.Wishart(df::RTR, S::Matrix{<:Real}) = TuringWishart(df, S)
Distributions.Wishart(df::RTR, S::AbstractMatrix{<:Real}) = TuringWishart(df, S)
Distributions.Wishart(df::Real, S::RTM) = TuringWishart(df, S)
Distributions.Wishart(df::RTR, S::RTM) = TuringWishart(df, S)

Distributions.InverseWishart(df::RTR, S::Matrix{<:Real}) = TuringInverseWishart(df, S)
Distributions.InverseWishart(df::RTR, S::AbstractMatrix{<:Real}) = TuringInverseWishart(df, S)
Distributions.InverseWishart(df::Real, S::RTM) = TuringInverseWishart(df, S)
Distributions.InverseWishart(df::RTR, S::RTM) = TuringInverseWishart(df, S)

function Distributions.logpdf(d::Wishart, X::RTM)
    return logpdf(TuringWishart(d), X)
end
function Distributions.logpdf(d::Wishart, X::AbstractArray{<:RTM})
    return logpdf(TuringWishart(d), X)
end

function Distributions.logpdf(d::InverseWishart, X::RTM)
    return logpdf(TuringInverseWishart(d), X)
end
function Distributions.logpdf(d::InverseWishart, X::AbstractArray{<:RTM})
    return logpdf(TuringInverseWishart(d), X)
end

# Modified from Tracker.jl

Base.vcat(xs::RTM...) = _vcat(xs...)
Base.vcat(xs::RTV...) = _vcat(xs...)
function _vcat(xs::Union{RTV{<:Any, D}, RTM{<:Any, D}}...) where {D}
    tp = ReverseDiff.tape(xs...)
    xs_value = ReverseDiff.value.(xs)
    out_value = vcat(xs_value...)
    function back(Δ)
        start = 0
        Δs = [begin
          x = map(_ -> :, size(xsi))
          i = isempty(x) ? x : Base.tail(x)
          d = Δ[start+1:start+size(xsi,1), i...]
          start += size(xsi, 1)
          d
        end for xsi in xs]
        return (Δs...,)
    end
    out = ReverseDiff.track(out_value, D, tp)
    ReverseDiff.record!(tp, SpecialInstruction, vcat, xs, out, (back,))
    return out
end

@noinline function ReverseDiff.special_reverse_exec!(instruction::SpecialInstruction{typeof(vcat)})
    output = instruction.output
    input = instruction.input
    input_derivs = ReverseDiff.deriv.(input)
    P = instruction.cache[1]
    jtvs = P(ReverseDiff.deriv(output))
    for i in 1:length(input_derivs)
        input_derivs[i] .+= jtvs[i]
    end
    ReverseDiff.unseed!(output)
    return nothing
end

@noinline function ReverseDiff.special_forward_exec!(instruction::SpecialInstruction{typeof(vcat)})
    output, input = instruction.output, instruction.input
    out_value = vcat(ReverseDiff.value.(input)...)
    ReverseDiff.value!(output, out_value)
    return nothing
end

Base.hcat(xs::RTM...) = _hcat(xs...)
Base.hcat(xs::RTV...) = _hcat(xs...)
function _hcat(xs::Union{RTV{<:Any, D}, RTM{<:Any, D}}...) where {D}
    tp = ReverseDiff.tape(xs...)
    xs_value = ReverseDiff.value.(xs)
    out_value = hcat(xs_value...)
    function back(Δ)
        start = 0
        Δs = [begin
          d = if ndims(xsi) == 1
            Δ[:, start+1]
          else
            i = map(_ -> :, size(xsi)) |> Base.tail |> Base.tail
            Δ[:, start+1:start+size(xsi,2), i...]
          end
          start += size(xsi, 2)
          d
        end for xsi in xs]
        return (Δs...,)
    end
    out = ReverseDiff.track(out_value, D, tp)
    ReverseDiff.record!(tp, SpecialInstruction, hcat, xs, out, (back,))
    return out
end

@noinline function ReverseDiff.special_reverse_exec!(instruction::SpecialInstruction{typeof(hcat)})
    output = instruction.output
    input = instruction.input
    input_derivs = ReverseDiff.deriv.(input)
    P = instruction.cache[1]
    jtvs = P(ReverseDiff.deriv(output))
    for i in 1:length(input_derivs)
        input_derivs[i] .+= jtvs[i]
    end
    ReverseDiff.unseed!(output)
    return nothing
end

@noinline function ReverseDiff.special_forward_exec!(instruction::SpecialInstruction{typeof(hcat)})
    output, input = instruction.output, instruction.input
    out_value = hcat(ReverseDiff.value.(input)...)
    ReverseDiff.value!(output, out_value)
    return nothing
end

Base.cat(Xs::RTA...; dims) = _cat(dims, Xs...)
Base.cat(Xs::RTV...; dims) = _cat(dims, Xs...)
function _cat(dims, Xs::Union{RTV{<:Any, D}, RTM{<:Any, D}}...) where {D}
    tp = ReverseDiff.tape(dims, Xs...)
    Xs_value = ReverseDiff.value.(Xs)
    out_value = cat(Xs_value...; dims = dims)
    function back(Δ)
        start = ntuple(i -> 0, Val(ndims(Δ)))
        Δs = [begin
          dim_xs = 1:ndims(xs)
          till_xs = ntuple((i -> i in dims ? (i in dim_xs ? size(xs,i) : 1) : 0), Val(ndims(Δ)))
          xs_in_Δ = ntuple(i -> till_xs[i] > 0 ? (start[i]+1:start[i]+till_xs[i]) : Colon(), Val(ndims(Δ)))
          d = reshape(Δ[xs_in_Δ...],size(xs))
          start = start .+ till_xs
          d
        end for xs in Xs]
        return (Δs...,)
    end
    out = ReverseDiff.track(out_value, D, tp)
    ReverseDiff.record!(tp, SpecialInstruction, cat, (dims, Xs...), out, (back,))
    return out
end

@noinline function ReverseDiff.special_reverse_exec!(instruction::SpecialInstruction{typeof(cat)})
    output = instruction.output
    input = instruction.input
    input_derivs = ReverseDiff.deriv.(Base.tail(input))
    P = instruction.cache[1]
    jtvs = P(ReverseDiff.deriv(output))
    for i in 1:length(jtvs)
        input_derivs[i] .+= jtvs[i]
    end
    ReverseDiff.unseed!(output)
    return nothing
end

@noinline function ReverseDiff.special_forward_exec!(instruction::SpecialInstruction{typeof(cat)})
    output, input = instruction.output, instruction.input
    dims = ReverseDiff.value(input[1])
    Xs = ReverseDiff.value.(Base.tail(input))
    out_value = cat(Xs..., dims = dims)
    ReverseDiff.value!(output, out_value)
    return nothing
end

###########

# Broadcasting

using ReverseDiff: ForwardOptimize
using Base.Broadcast: Broadcasted
import Base.Broadcast: materialize
const RDBroadcasted{F, T} = Broadcasted{<:Any, <:Any, F, T}

_materialize(f, args) = broadcast(ForwardOptimize(f), args...)

for (M, f, arity) in ReverseDiff.DiffRules.diffrules()
    if arity == 1
        @eval @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{RTA}}) = _materialize(bc.f, bc.args)
    elseif arity == 2
        @eval begin
            @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{RTA, RTA}}) = _materialize(bc.f, bc.args)
            @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{RTA, RTR}}) = _materialize(bc.f, bc.args)
            @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{RTR, RTA}}) = _materialize(bc.f, bc.args)
        end
        for A in ReverseDiff.ARRAY_TYPES
            @eval begin
                @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{$A, RTA}}) = _materialize(bc.f, bc.args)
                @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{RTA, $A}}) = _materialize(bc.f, bc.args)
                @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{$A, RTR}}) = _materialize(bc.f, bc.args)
                @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{RTR, $A}}) = _materialize(bc.f, bc.args)
            end
        end
        for R in ReverseDiff.REAL_TYPES
            @eval begin
                @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{$R, RTA}}) = _materialize(bc.f, bc.args)
                @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{RTA, $R}}) = _materialize(bc.f, bc.args)
            end
        end
    end
end


module VTUFileHandler

using Base64, CodecZlib
using XMLParser

include(joinpath(".","VTUFileHandler","defs.jl"))

"""
    VTUHeader(::Type{T},input::Vector{UInt8}) where {T<:Union{UInt32,UInt64}}

Computes the VTUHeader based on the headertype and a Base64 decoded input data array.

# Constructor
- `::Type{T}`: headertype, either UInt32 or UInt64
- `input::Vector{UInt8}`: input data

# Fields
- `num_blocks::T` : number of blocks
- `blocksize::T` : size of blocks
- `last_blocksize::T` : size of last block (can be different)
- `compressed_blocksizes::T` : size of compressed blocks
"""
struct VTUHeader{T<:Union{UInt32,UInt64}}
	num_blocks::T
	blocksize::T
	last_blocksize::T
	compressed_blocksizes::Vector{T}
	function VTUHeader(::Type{T}) where {T<:Union{UInt32,UInt64}}
		return new{T}(zero(T),zero(T),zero(T),T[])
	end
	function VTUHeader(::Type{T},input::Vector{UInt8}) where {T<:Union{UInt32,UInt64}}
		dat = reinterpret(T,input)
		num_blocks = dat[1]
		blocksize = dat[2]
		last_blocksize = dat[3]
		compressed_blocksizes=Vector{T}(undef,num_blocks)
		for (j,i) = enumerate(4:length(dat))
			compressed_blocksizes[j] = dat[i]
		end
		return new{T}(num_blocks,blocksize,last_blocksize,compressed_blocksizes)
	end
	function VTUHeader(a::T,b::T,c::T,d::Vector{T}) where {T<:Union{UInt32,UInt64}}
		return new{T}(a,b,c,d)
	end
end

include(joinpath(".","VTUFileHandler","vtuheader_utils.jl"))
include(joinpath(".","VTUFileHandler","io.jl"))

"""
    VTUDataField(x::Vector{T}) where {T} = new{T}(x)

Container for VTU field data.

# Constructor
- `x::Vector{T}`: Vector with VTU field data

# Fields
- `dat::Vector{T}` : VTU field data
"""
struct VTUDataField{T}
	dat::Vector{T}
	VTUDataField(x::Vector{T}) where {T} = new{T}(x)
	VTUDataField(::Type{T}) where {T} = new{T}()
	VTUDataField(x::Base.ReinterpretArray{T, A, B, Vector{C}, D}) where {T,A,B,C,D} = new{T}(x)
end

"""
    VTUData(dataarrays::Vector{XMLElement},appendeddata::Vector{XMLElement},headertype::Union{Type{UInt32},Type{UInt64}},offsets::Vector{Int},compressed_dat::Bool)

Container for VTU data.

# Constructor
- `dataarrays::Vector{XMLElement}`: Vector with all XML elements with tag `DataArray` of VTU file
- `appendeddata::Vector{XMLElement}`: Vector with all XML elements with tag `AppendedData`
- `headertype::Union{Type{UInt32},Type{UInt64}}` : Type of VTU header
- `offsets::Vector{Int}` : Offset of each field data in the compressed appended data
- `compressed_dat::Bool` : True if data is compressed

# Fields
- `names::Vector{String}` : Name of each data field
- `header::Vector{VTUHeader}` : Vector with [`VTUHeader`](@ref)s.
- `data::Vector{VTUDataField}` : Vector with [`VTUDataField`](@ref)s
- `interp_data::Vector{VTUDataField{Float64}}` : Vector with [`VTUDataField`](@ref)s
- `idat::Vector{Int}` : Vector indexing the `data`-fields onto which the math operators should be applied

For a [`VTUDataField`](@ref) to appear in `data`, the appropriate keyword has to be added via [`add_uncompress_keywords`](@ref) or [`add_interpolation_keywords`](@ref).
For a [`VTUDataField`](@ref) to appear in `interp_data` the appropriate keyword has to be added via [`add_interpolation_keywords`](@ref).
For more information, see [`VTUKeyWords`](@ref).
"""
struct VTUData
	names::Vector{String}
	header::Vector{VTUHeader}
	data::Vector{VTUDataField}
	interp_data::Vector{VTUDataField{Float64}}
	idat::Vector{Int}
	VTUData(names,header,data,interp_data,idat) = new(names,header,data,interp_data,idat)
	function VTUData(dataarrays::Vector{XMLElement},appendeddata::Vector{XMLElement},headertype::Union{Type{UInt32},Type{UInt64}},offsets::Vector{Int},compressed_dat::Bool)
		names = Vector{String}()
		data = Vector{VTUDataField}()
		header = Vector{VTUHeader}()
		for (i,el) in enumerate(dataarrays)
			_type = replace(getAttribute(el,"type"),"\""=>"")
			type = eval(Meta.parse(_type))
			_format = getAttribute(el,"format")
			_name = getAttribute(el,"Name")
			if findfirst(x->x==_name,vtukeywords.interpolation_keywords) != nothing	|| findfirst(x->x==_name,vtukeywords.uncompress_keywords) != nothing
				if _format == "appended" || _format == "\"appended\""
					vtuheader = readappendeddata!(data,i,appendeddata,offsets,type,headertype,compressed_dat)
				else
					vtuheader = readdataarray!(data,el,type,headertype,compressed_dat)
				end
				push!(names,_name)
				push!(header,vtuheader)
			end
		end
		idat = findall(map(y->findfirst(x->x==y,vtukeywords.interpolation_keywords)!=nothing,names))
		interp_data = data[idat]
		return new(names,header,data,interp_data,idat)
	end
end

include(joinpath(".","VTUFileHandler","vtudata_utils.jl"))
include(joinpath(".","VTUFileHandler","vtudata_math.jl"))

function deletefieldata!(xmlroot)
	dataarrays = getElements(xmlroot,"DataArray")
	fielddata = getElements(xmlroot,"FieldData")[1]
	_offset = parse(Int,replace(getAttribute(dataarrays[length(fielddata.content)+1],"offset"),"\""=>""))
	empty!(fielddata.content)
	els = getElements(xmlroot,"DataArray")
	for el in els
		setAttribute(el,"offset",string(parse(Int,replace(getAttribute(el,"offset"),"\""=>""))-_offset))
	end
	appendeddata = getElements(xmlroot,"AppendedData")
	appendeddata[1].content[1] = String(deleteat!(collect(appendeddata[1].content[1]),2:_offset+1))
end

"""
    VTUFile(name::String)

Loads a VTU file.
Don't forget to set the proper fieldnames via `set_uncompress_keywords` and `set_interpolation_keywords`.

# Constructor
- `name::String`: path to vtu file

# Fields
- `name::String`: path to vtu file; destination for file writing
- `xmlroot::XMLElement`: VTU file in XML represantation
- `dataarrays::Vector{XMLElement}`: Vector with all XML elements with tag `DataArray` of VTU file
- `headertype::Union{Type{UInt32},Type{UInt64}}`: type of header
- `offsets::Vector{Int}`: Offset of each field data in the compressed appended data
- `data::VTUData`: Conatainer with  [`VTUData`](@ref)
- `compressed_dat::Bool` : True if data is compressed

# Example
```julia
set_uncompress_keywords("temperature","points")
set_interpolation_keywords("temperature")
vtufile = VTUFile("./path-to-vtu/example.vtu");
```
"""
mutable struct VTUFile
	name::String
	xmlroot::XMLElement
	dataarrays::Vector{XMLElement}
	appendeddata::Vector{XMLElement}
	headertype::Union{Type{UInt32},Type{UInt64}}
	offsets::Vector{Int}
	data::VTUData
	compressed_dat::Bool
	VTUFile(name,xmlroot,dataarrays,appendeddata,headertype,offsets,data) = new(name,xmlroot,dataarrays,appendeddata,headertype,offsets,data,true)
	VTUFile(name,xmlroot,dataarrays,appendeddata,headertype,offsets,data,compr_dat) = new(name,xmlroot,dataarrays,appendeddata,headertype,offsets,data,compr_dat)
	function VTUFile(name::String)
		state = IOState(name);
		xmlroot = readXMLElement(state);
		if !isempty(getElements(xmlroot,"FieldData"))
			deletefieldata!(xmlroot)
		end
		dataarrays = getElements(xmlroot,"DataArray")
		appendeddata = getElements(xmlroot,"AppendedData")
		vtkfile = getElements(xmlroot,"VTKFile")
		@assert length(vtkfile) == 1
		compressed_dat = false
		if hasAttributekey(vtkfile[1],"compressor")
			compressed_dat = true
		end
		#headertype = eval(Symbol(getAttribute(first(vtkfile),"header_type")))
		attr = replace(getAttribute(first(vtkfile),"header_type"),"\""=>"")
		headertype = eval(Meta.parse(attr))
		offsets = Vector{Int}()
		for el in dataarrays
			if hasAttributekey(el,"offset")
				_offset = replace(getAttribute(el,"offset"),"\""=>"")
				offset = parse(Int,_offset)
				push!(offsets,offset)
			else
				push!(offsets,0)
			end
		end
		retval = new(name,xmlroot,dataarrays,appendeddata,headertype,offsets,VTUData(dataarrays,appendeddata,headertype,offsets,compressed_dat),compressed_dat)
		return retval
	end
end

include(joinpath(".","VTUFileHandler","utils.jl"))
include(joinpath(".","VTUFileHandler","vtufile_math.jl"))
include(joinpath(".","VTUFileHandler","globaltolocal.jl"))

export VTUHeader, VTUFile, set_uncompress_keywords, add_uncompress_keywords, set_interpolation_keywords, add_interpolation_keywords, add!, minus!, pow!, div!, min!, max!

end #module VTUHandler


include("../imports.jl")

@testset ExtendedTestSet "ImagePreprocessing" begin
    step = ImagePreprocessing((0, 0, 0), (.5, .5, .5))
    image = rand(RGB, 100, 100)
    x = FastAI.run(step, Training(), image)
    @test size(x) == (100, 100, 3)
    @test eltype(x) == Float32

    image2 = rand(RGB, 100, 100)
    buf = copy(x)
    @test_nowarn FastAI.run!(buf, step, Training(), image2)
    @test !(buf ≈ x)
end


using GraphPlot
using Plots
pyplot()
using PyPlot
pygui(false)
include("loadbuild.jl")
include("graphing.jl")
include("beliefprop.jl")
include("Psis.jl")

EDGEFILE = "data/save.feather"
BIASFILE = "data/bias.csv"
plotdir = "results/plots/"
figext = "eps"

problem_types = ["bias", "fake"]
epsilon = 0.44
MAXITER = 2

Plots.PyPlot.rc("lines", lw=10)
# Plots.PyPlot.rc("font", family="monospace")
# Plots.PyPlot.rc("font", weight="bold")
Plots.PyPlot.rc("font", size=22)

if EDGEFILE == "data/save.feather"
  by_domain_flag = true
elseif EDGEFILE == "data/save2.feather"
  by_domain_flag = false
end

save = loadSave(EDGEFILE)

# size(unique(save[:sdom]))

#println(size(unique(vcat(save[:rsrc],save[:rdest]))))

bias = loadBias(BIASFILE)

cats = buildWeb(save, by_domain_flag) #only mutual links for "a"

edgeunion = countunion(cats["a"], cats["img"])
@time for x in ["link", "script"]
    countunion!(edgeunion, cats[x])
end

sets = [("union", 1)] #[("union",10)], [("mut", 1000), ("mut", 50), ("mut", 2), ("mut", 1), ("a", 1000), ("a", 500),("a", 100), ("img", 200), ("img", 50)]

#s2p = Dict{String, Array}()

roc_x = Vector{Float64}()
roc_y = Vector{Float64}()

k = 3 #number of folds

for (sn, set) in enumerate(sets)
   edgetype = set[1]
    if edgetype == "union"
      edgeset = edgeunion
    else
      edgeset = cats[set[1]]
    end

    # if ~by_domain_flag
    #   domainmap = Dict{String, String}()
    #   for j in edgeset
    #     domainmap[j.first] = split(j.first, '/')[1]
    #     for i in j.second
    #       domainmap[i.first] = split(i.first, '/')[1]
    #     end
    #   end
    # end

    G, s2i, i2s = constructGraph(edgeset, minsamp=set[2], giant=false, by_domain=by_domain_flag)
    #G2, s2i2, i2s2 = constructGraph(edgeset, minsamp=set[2], giant=true, by_domain=by_domain_flag)
    mindeg = minimum([length([x for x in all_neighbors(G, v) if x != v]) for v in 1:nv(G)])
    println("Minimum Degree: $mindeg")
    maxdeg = maximum(degree(G))
    println("Maximum Degree: $maxdeg")
    # lx, ly = spring_layout(G)
    #lx2, ly2 = spring_layout(G2)

    for key in problem_types
        p = Plots.plot()
        plot!(p, [0, 1], [0, 1], line=:dash, label="")
        println(key*string(set))
        try mkdir(plotdir*string(set[1])) catch x end
        try mkdir(plotdir*string(set[1])*"/"*string(set[2])) catch x end
        mydir = plotdir*string(set[1])*"/"*string(set[2])*"/"

        # plotGraph(bias, G, s2i, i2s, lx, ly, color=key,
        #             path=mydir*key*"train.png")
        # ϕ = makeBeliefs(bias, G, s2i, i2s, key)
        # pr, b, lodds = beliefprop(G, ϕ, Psis(0.44), 1);
        # plotGraph(bias, G, s2i, i2s, lx, ly, color=b[:,1],
                    # path=mydir*key*"prop.png")

        # plotGraph(bias, G2, s2i2, i2s2, lx2, ly2, color=key,
        #             path=mydir*key*"train_main.png")
        # ϕ = makeBeliefs(bias, G2, s2i2, i2s2, key)
        # pr, b, lodds = beliefprop(G2, ϕ, Psis(0.44), 1);
        # plotGraph(bias, G2, s2i2, i2s2, lx2, ly2, color=b[:,1],
        #             path=mydir*key*"prop_main.png")

        folds = makeFolds(G, k) # makeFolds(G, k, i2s, domainmap)
        aucs = []
        for f in 1:k
          println("Fold " *string(f))
            ϕ = makeBeliefs(bias, G, s2i, i2s, key, folds=folds, fold=f)
            pr, b, lodds = beliefprop(G, ϕ, Psis(epsilon), MAXITER);
            # for v in 1:nv(G)
            #     if folds[v] == f
            #         site = i2s[v]
            #         if ~haskey(s2p, site)
            #             s2p[site] = ones(1, size(sets, 1)+2)./2.0
            #             s2p[site][size(sets, 1)+1] = reality(bias, site, "bias")
            #             s2p[site][size(sets, 1)+2] = folds[v]
            #         end
            #         s2p[site][sn] = b[v]
            #     end
            # end

            roc_x, roc_y, acc, lvls, cms = getROC(bias, G, s2i, i2s, key, b, folds=folds, fold=f)
            auc = AUC(roc_x, roc_y)
            #println("Max "*string(maximum(acc))*" at "*string(lvls[find(x->x==maximum(acc), acc)[1]]))
            println("AUC : "*string(auc))
            best_cutoff = collect(keys(acc))[indmax(collect(values(acc)))]
            println("best cutoff threshold : "*string(best_cutoff))
            println("confusion matrix (best threshold) : ")
            show(cms[best_cutoff])
            println("")
            println("confusion matrix (0.5 threshold) : ")
            show(cms[0.5])
            println("")
            #println(typeof(roc_x))
            #println(typeof(roc_y))
            #println(roc_x)
            #println(roc_y)
            plot!(p, roc_x, roc_y,label="Fold $f (AUC=$(@sprintf("%.3f", auc)))")
            push!(aucs, auc)
        end
        avg_auc = sum(aucs)/k
        title!("Receiver Operating Characteristic (Average AUC=$(@sprintf("%.3f", avg_auc)))")
        println("Avg AUC : "*string(avg_auc))
        figpath = mydir*key*"_ROC_curve.$figext"
        println("Saving figure to: $figpath")
        savefig(figpath)

        open(mydir*"score.txt", "w") do f
            write(f, string(sum(aucs)/k))
        end
    end
end

# for site in keys(s2p)
#     s2p[site][10] = reality(bias, site, "bias")
# end

# s2p

# s2p["cnn.com"]
# s2p["bluebirdbanter.com"]

# using JSON
# open("metadata.json", "w") do f
#     write(f, JSON.json(s2p))
# end


using MCPhylo
using Test

@testset "insert_node!" begin
    tree = MCPhylo.parsing_newick_string("(A,B,(C,D,E)F)G;")
    insert_tree = MCPhylo.parsing_newick_string("(A,B,((C,D,E)no_name)F)G;")
    insert_tree_newick = newick(insert_tree)
    insert_tree2 = MCPhylo.parsing_newick_string("(A,B,(((C,D)no_name,E)no_name)F)G;")
    insert_tree_newick2 = newick(insert_tree2)
    MCPhylo.number_nodes!(tree)

    children = [find_by_name(tree, "C"), find_by_name(tree, "D"),
                find_by_name(tree, "E")]
    mother = find_by_name(tree, "F")
    inserted_node = insert_node!(mother, children)
    @test newick(tree) == insert_tree_newick

    pop!(children)
    insert_node!(inserted_node, children)
    @test newick(tree) == insert_tree_newick2

    new_node = Node("G")
    append!(children, new_node)
    @test_throws AssertionError insert_node!(mother, children)
 end

@testset "delete_node!" begin
    tree = MCPhylo.parsing_newick_string("(A,B,(C,D)E)F;")
    MCPhylo.number_nodes!(tree)
    remove_tree = MCPhylo.parsing_newick_string("(A,B,C,D)F;")
    remove_tree_newick = newick(remove_tree)
    remove_tree2 = MCPhylo.parsing_newick_string("(B,C,D)F;")
    remove_tree2_newick = newick(remove_tree2)

    delete_node!(find_by_name(tree, "E"))
    @test newick(tree) == remove_tree_newick
    delete_node!(find_by_name(tree, "A"))
    @test newick(tree) == remove_tree2_newick
    @test_throws ArgumentError delete_node!(find_by_name(tree, "F"))
end

@testset "find_lca" begin
    tree = MCPhylo.parsing_newick_string("(A,B,(C,(D,E)F)G)H;")
    MCPhylo.set_binary!(tree)

    @testset "find_lca by name" begin
        @test find_lca(tree, ["A","C"]) == tree
        @test find_lca(tree, ["D", "E"]) == find_by_name(tree, "F")
        @test find_lca(tree, ["D", "F"]) == find_by_name(tree, "F")
        @test find_lca(tree, ["C", "D"]) == find_by_name(tree, "G")
        @test find_lca(tree, ["A","B","C"]) == tree
        @test find_lca(tree, ["D", "E", "F", "G"]) == find_by_name(tree, "G")
        @test find_lca(tree, ["D", "E", "C"]) == find_by_name(tree, "G")
    end # find_lca by name

    @testset "find_lca by nodes" begin
        nodes = [find_by_name(tree, "A"), find_by_name(tree, "B"),
                 find_by_name(tree, "C")]
        @test find_lca(tree, nodes) == tree

        nodes = [find_by_name(tree, "D"), find_by_name(tree, "E"),
                 find_by_name(tree, "F"), find_by_name(tree, "G")]
        @test find_lca(tree, nodes) == find_by_name(tree, "G")
    end # find_lca by nodes
end

@testset "check_leafsets" begin
    tree1 = MCPhylo.parsing_newick_string("(((A,B),C),(D,E))")
    tree2 = MCPhylo.parsing_newick_string("((A,C),(B,D,E))")
    tree3 = MCPhylo.parsing_newick_string("(((G,C),A),D,F)")
    tree4 = MCPhylo.parsing_newick_string("(((B,C),A),D,E)")
    tree5 = MCPhylo.parsing_newick_string("(((B,C),A),D,F)")

    trees = [tree1, tree2, tree3, tree4, tree5]
    MCPhylo.number_nodes!.(trees)
    MCPhylo.set_binary!.(trees)

    @test_throws ArgumentError MCPhylo.check_leafsets(trees)
end


abstract type AbstractContact{T} <: AbstractObject{T} end

"""
    mutable struct Contact{T} <: AbstractContact{T}

T: Type of precision.
"""
mutable struct Contact{T} <: AbstractContact{T}
    potential::T
    material::NamedTuple
    id::Int
    name::String
    geometry::AbstractGeometry{T}
    geometry_positive::Vector{AbstractGeometry{T}}
    geometry_negative::Vector{AbstractGeometry{T}}
end

function Contact{T}(dict::Union{Dict{String,Any}, Dict{Any, Any}}, inputunit_dict::Dict{String,Unitful.Units})::Contact{T} where {T <: SSDFloat}
    haskey(dict, "channel") ? channel = dict["channel"] : channel = -1
    haskey(dict, "material") ? material = material_properties[materials[dict["material"]]] : material = material_properties[materials["HPGe"]]
    haskey(dict,"name") ? name = dict["name"] : name = ""
    geometry =  Geometry(T, dict["geometry"], inputunit_dict )
    geometry_positive, geometry_negative = get_decomposed_volumes(geometry)
    return Contact{T}( dict["potential"], material, channel, name, geometry, geometry_positive, geometry_negative )
end

function println(io::IO, d::Contact{T}) where {T <: SSDFloat}
    println("\t________"*"Contact $(d.id)"*"________\n")
    println("\t---General Properties---")
    println("\t-Potential: \t\t $(d.potential)")
    println("\t-Contact Material: \t $(d.material.name)")
    println()

end

function show(io::IO, d::Contact{T}) where {T <: SSDFloat} println(d) end
function print(io::IO, d::Contact{T}) where {T <: SSDFloat} println(d) end
function display(io::IO, d::Contact{T} ) where {T <: SSDFloat} println(d) end
function show(io::IO,::MIME"text/plain", d::Contact) where {T <: SSDFloat}
    show(io, d)
end


__precompile__()

module PiecewiseLinearOpt

import JuMP

export PWLFunction, UnivariatePWLFunction, BivariatePWLFunction, piecewiselinear

include("types.jl")
include("jump.jl")

end # module


function asap2(style = Rectangular(2,1))
    arch = TopLevel{2}("asap2")
    arch.metadata["style"] = style

    processor = build_processor_tile(style)
    for r in 0:11, c in 0:12
        add_child(arch, processor, Address(r,c))
    end
    # Add the two processors on the bottom that are not memory processors.
    # Do not include "viterbi", "fft", or "motion estimation" for now.
    add_child(arch, processor, Address(12, 4))
    add_child(arch, processor, Address(12, 9))

    # Add the memory processors.
    memory_processor = build_processor_tile(style, include_memory = true)
    for r in 12, c in (2, 3, 5, 6, 7, 8)
        add_child(arch, memory_processor, Address(r, c))
    end

    # Add the memories. Use left corner as address.
    memory = build_memory(2)
    for r in 13, c in (2, 5, 7)
        add_child(arch, memory, Address(r, c))
    end

    input_handler = build_input_handler(12)
    add_child(arch, input_handler, Address(0,-1))

    output_handler = build_output_handler(12)
    add_child(arch, output_handler, Address(0,13))

    connect_processors(arch, style)
    connect_memories(arch, style)
    connect_io_asap2(arch)

    return arch
end

function connect_io_asap2(arch)

    # Build metadata dictionary for capacity and cost
    metadata = Dict(
        "cost"          => 1.0,
        "capacity"      => 1,
        "link_class"    => "link_class"
    )

    # Connect input manually
    for i in 0:11
        src_address = Address(0,-1)
        dst_address = Address(i,0)
        src_name = getname(arch, src_address)
        dst_name = getname(arch, dst_address)

        src_port = Path{Port}(src_name, "out[$i]")
        dst_port = Path{Port}(dst_name, "west_in[0]")
        add_link(arch, src_port, dst_port, metadata = metadata)
    end

    # Connect output manually
    for i in 0:11
        src_address = Address(i,12)
        dst_address = Address(0,13)
        src_name = getname(arch, src_address)
        dst_name = getname(arch, dst_address)

        src_port = Path{Port}(src_name, "east_out[0]")
        dst_port = Path{Port}(dst_name, "in[$i]")
        add_link(arch, src_port, dst_port, metadata = metadata)
    end
end


module TestTargetedLearning

using FactCheck
# FactCheck.setstyle(:compact)

is_a(T) = x -> isa(x, T)

using TargetedLearning

include("LReg.jl")
include("automaticics.jl")
include("Qmodels.jl")
include("Parameters.jl")
include("TMLEs.jl")
include("CTMLEs.jl")

FactCheck.exitstatus()
end


using Test

using FlightMechanicsUtils

using FlightMechanicsSimulator

# Create different dynamic systems
x_sixdofaeroeuler = SixDOFAeroEuler([
    502.0 * FT2M,
    0.2392628,
    0.0005061803,
    1.366289,
    0.05000808,
    0.2340769,
    -0.01499617,
    0.2933811,
    0.06084932,
    0.0 * FT2M,
    0.0 * FT2M,
    0.0 * FT2M,
    64.12363,
])

x_sixdofbodyeuler = SixDOFBodyEuler([
    wind2body(502.0 * FT2M, 0, 0,  0.2392628, 0.0005061803)...,
    1.366289,
    0.05000808,
    0.2340769,
    -0.01499617,
    0.2933811,
    0.06084932,
    0.0 * FT2M,
    0.0 * FT2M,
    0.0 * FT2M,
    64.12363,
])

DSS_ARR = [x_sixdofaeroeuler, x_sixdofbodyeuler]

# Test the interface for each system.
@testset "Interface $(typeof(dss))" for dss in DSS_ARR

    # Check contruction of a DSState type from any other DSState type (including itself)
    @testset "Constructor from $(typeof(dss2))" for dss2 in DSS_ARR
        @test hasmethod(typeof(dss), (typeof(dss2),))
        # TODO: check that DSState are equivalent
    end

    # Contructor for trimmer
    @testset "Constructor for Trimmer" begin
        @test hasmethod(
            typeof(dss),
            (
                FlightMechanicsSimulator.TrimConditions,
                FlightMechanicsSimulator.TrimSolution,
                FlightMechanicsSimulator.Aircraft,
            ),
        )
    end

    methods_scalar_output = [
        get_height,
        get_tas,
        get_α,
        get_β,
        get_engine_power,
    ]

    # Check that methods that return a scalar exist and have right return type
    @testset "$im" for im in methods_scalar_output
        @test hasmethod(im, (typeof(dss),))
        @test isa(im(dss), Number)
    end

    methods_vec3_output = [
        get_earth_position,
        get_tasαβ,
        get_euler_angles,
        get_body_velocity,
        get_horizon_velocity,
        get_ang_vel_body,
        get_euler_angles_rates,
    ]

    # Check that methods that return a len 3 vector exist and have right return type
    @testset "$im" for im in methods_vec3_output
        @test hasmethod(im, (typeof(dss),))
        rv = im(dss)
        @test isa(rv, AbstractArray)
        @test length(rv) == 3
        @test eltype(rv) <: Number
    end

    # Check get_n_states implementation
    @testset "$im" for im in [get_n_states]
        @test hasmethod(im, (typeof(dss),))
        rv = im(dss)
        @test isa(rv, Integer)
        @test length(dss.x) == rv
    end

    # Check get_x implementation
    @testset "$im" for im in [get_x]
        @test hasmethod(im, (typeof(dss),))
        rv = im(dss)
        @test isa(rv, AbstractArray)
        @test length(rv) == get_n_states(dss)
        @test eltype(rv) <: Number
    end

    # Check get_x_names implementation
    @testset "$im" for im in [get_x_names]
        @test hasmethod(im, (typeof(dss),))
        rv = im(dss)
        @test isa(rv, AbstractArray)
        @test length(rv) == get_n_states(dss)
        @test eltype(rv) == Symbol
    end

    # Check state_eqs implementation
    @testset "state_eqs" begin
        @test hasmethod(
            state_eqs,
            (
                typeof(dss),
                Number,
                Number,
                AbstractMatrix,
                AbstractVector,
                AbstractVector,
                AbstractVector,
                Number,
            ),
        )

        time = 0.0
        mass = 10500.0
        inertia = [555.5 0.0 1234.5; 0.0 5674.0 0.0; 987.2 0.0  7895.3]
        forces = [1.0, 2.0, 3.0]
        moments = [4.0, 5.0, 6.0]
        hx = [12.0, 20.0, 30.0]
        pdot = 12.5

        dssd = state_eqs(dss, time, mass, inertia, forces, moments, hx, pdot)

        @test isa(dssd, DSStateDot)
    end

    # Check DSStateDot for each subtype: constructor and mandatory methods.
    @testset "DSStateDot{$(typeof(dss)), $(get_n_states(dss)), $(eltype(dss))}" begin
        # Test constructor
        dssd = DSStateDot(dss, get_x(dss))

        @testset "$im" for im in [get_ds_state]
            @test hasmethod(im, (typeof(dssd),))
            rv = im(dssd)
            @test isa(rv, typeof(dss))
        end

        @testset "$im" for im in [get_xdot]
            @test hasmethod(im, (typeof(dssd),))
            rv = im(dssd)
            @test isa(rv, AbstractArray)
            @test length(rv) == get_n_states(dss)
            @test eltype(rv) <: Number
        end

        @testset "$im" for im in [get_n_states]
            @test hasmethod(im, (typeof(dss),))
            rv = im(dssd)
            @test isa(rv, Integer)
            @test length(dss.x) == rv
        end

        @testset "$im" for im in [get_x]
            @test hasmethod(im, (typeof(dss),))
            rv = im(dssd)
            @test isa(rv, AbstractArray)
            @test length(rv) == get_n_states(dss)
            @test eltype(rv) <: Number
        end

        @testset "$im" for im in [get_x_names]
            @test hasmethod(im, (typeof(dss),))
            rv = im(dssd)
            @test isa(rv, AbstractArray)
            @test length(rv) == get_n_states(dss)
            @test eltype(rv) == Symbol
        end

        methods_scalar_output = [
            get_height,
            get_tas,
            get_α,
            get_β,
            get_engine_power,
            get_tas_dot,
            get_α_dot,
            get_β_dot,
            get_engine_power_dot,
        ]

        @testset "$im" for im in methods_scalar_output
            @test hasmethod(im, (typeof(dssd),))
            @test isa(im(dssd), Number)
        end

        methods_vec3_output = [
            get_earth_position,
            get_tasαβ,
            get_euler_angles,
            get_body_velocity,
            get_horizon_velocity,
            get_ang_vel_body,
            get_euler_angles_rates,
            get_tasαβ_dot,
            get_accel_body,
            get_ang_vel_body,
        ]

        @testset "$im" for im in methods_vec3_output
            @test hasmethod(im, (typeof(dssd),))
            rv = im(dssd)
            @test isa(rv, AbstractArray)
            @test length(rv) == 3
            @test eltype(rv) <: Number
        end

        @testset "$im" for im in [get_trimmer_cost]
            @test hasmethod(im, (typeof(dssd),))
            rv = im(dssd)
            @test length(rv) == 6
        end
    end
end