text:
using ModelingToolkitStandardLibrary.Electrical, ModelingToolkit, OrdinaryDiffEq, Test
using ModelingToolkitStandardLibrary.Electrical: Set, Reset, _and, _or, _not

@parameters t

@named set1   = Set()
@named reset1 = Reset()
@named set2   = Set()
@named reset2 = Reset()
@named out = DigitalPin()

@testset "Not gate" begin
    @named set = Set()
    @named reset = Reset()
    @named not = Not()
    sources = [set, reset]
    for source in sources
        not_eqs = [
            connect(source.d, not.x)
            connect(out, not.y)
                ]
        @named not_model = ODESystem(not_eqs, t, systems = [out, not, source])
        sys = alias_elimination(not_model)
        u0 = [
            not.y.val => 0.0
        ]
        prob = ODEProblem(sys, u0, (0, 1.5))
        sol  = solve(prob, Rosenbrock23())

        @test .!(sol[not.x.val] .> 0.5) == Bool.(sol[not.y.val])
    end
end

@testset "And, Nand gate" begin
    @named and   = And()
    @named nand  = Nand()
    for one in [set1, reset1], two in [set2, reset2]
        and_eqs = [
            connect(one.d, and.x1)
            connect(two.d, and.x2)
            connect(out, and.y)
        ]
        @named and_model = ODESystem(and_eqs, t, systems=[and, one, two, out])
        sys = alias_elimination(and_model)

        u0 = []
        prob = ODEProblem(sys, u0, (0, 1.5))
        sol = solve(prob, Rosenbrock23())
        @test sol.retcode == :Success
        @test sol[and.y.val] == _and.(sol[one.d.val], sol[two.d.val])

        nand_eqs = [
            connect(one.d, nand.x1)
            connect(two.d, nand.x2)
            connect(out, nand.y)
            ]
        @named nand_model = ODESystem(nand_eqs, t, systems=[nand, one, two, out])
        sys = alias_elimination(nand_model)

        u0 = []
        prob = ODEProblem(sys, u0, (0, 1.5))
        nsol = solve(prob, Rosenbrock23())
        @test nsol.retcode == :Success
        @test nsol[nand.y.val] == _not.(sol[and.y.val])
    end
end

@testset "Or gate" begin
    @named or = Or()
    @named nor = Nor()
    for one in [set1, reset1], two in [set2, reset2]
        or_eqs = [
            connect(one.d, or.x1)
            connect(two.d, or.x2)
            connect(out, or.y)
        ]
        @named or_model = ODESystem(or_eqs, t, systems = [or, one, two, out])
        sys = alias_elimination(or_model)

        u0 = []
        prob = ODEProblem(sys, u0, (0, 1.5))
        sol = solve(prob, Rosenbrock23())
        @test sol.retcode == :Success
        @test sol[or.y.val] == _or.(sol[one.d.val], sol[two.d.val])

        nor_eqs = [
            connect(one.d, nor.x1)
            connect(two.d, nor.x2)
            connect(out, nor.y)
        ]
        @named nor_model = ODESystem(nor_eqs, t, systems=[nor, one, two, out])
        sys = alias_elimination(nor_model)

        u0 = []
        prob = ODEProblem(sys, u0, (0, 1.5))
        nsol = solve(prob, Rosenbrock23())
        @test nsol.retcode == :Success
        @test nsol[nor.y.val] == _not.(sol[or.y.val])
    end
end

@testset "Xor gate" begin
    @named xor  = Xor()
    @named xnor = Xnor()
    for one in [set1, reset1], two in [set2, reset2]
        xor_eqs = [
            connect(one.d, xor.x1)
            connect(two.d, xor.x2)
            connect(out, xor.y)
        ]
        @named xor_model = ODESystem(xor_eqs, t, systems = [xor, one, two, out])
        sys = alias_elimination(xor_model)

        u0 = []
        prob = ODEProblem(sys, u0, (0, 1.5))
        sol = solve(prob, Rosenbrock23())
        @test sol.retcode == :Success
        @test sol[xor.y.val] == _or.(_and.(sol[one.d.val], _not.(sol[two.d.val])),
                                          _and.(sol[two.d.val], _not.(sol[one.d.val])))

        xnor_eqs = [
            connect(one.d, xnor.x1)
            connect(two.d, xnor.x2)
            connect(out, xnor.y)
        ]
        @named xnor_model = ODESystem(xnor_eqs, t, systems=[xnor, one, two, out])
        sys = alias_elimination(xnor_model)

        u0 = []
        prob = ODEProblem(sys, u0, (0, 1.5))
        nsol = solve(prob, Rosenbrock23())
        @test nsol.retcode == :Success
        @test nsol[xnor.y.val] == _not.(sol[xor.y.val])
    end
end

@testset "Half Adders" begin
    @named ha = HalfAdder()
    @named out1 = DigitalPin()
    @named out2 = DigitalPin()
    for one in [set1, reset1], two in [set2, reset2]
        ha_eqs = [
            connect(one.d, ha.x1)
            connect(two.d, ha.x2)
            connect(out1, ha.y0)
            connect(out2, ha.y1)
        ]
        @named ha_model = ODESystem(ha_eqs, t, systems = [ha, one, two, out1, out2])
        sys = alias_elimination(ha_model)

        u0 = []

        prob = ODEProblem(sys, u0, (0, 1.5))
        sol = solve(prob, Rosenbrock23())

        @test sol.retcode == :Success
        @test sol[ha.y0.val] == sol[ha.sum] == _xor.(sol[one.d.val], sol[two.d.val])
    end
end

@testset "Full Adder" begin
    @named set3   = Set()
    @named reset3 = Reset()
    @named out1 = DigitalPin()
    @named out2 = DigitalPin()
    for one in [set1, reset1], two in [set2, reset2], three in [set3, reset3]
        @named fa = FullAdder()
        fa_eqs = [
            connect(one.d, fa.x1)
            connect(two.d, fa.x2)
            connect(three.d, fa.x3)
            connect(out1, fa.y0)
            connect(out2, fa.y1)
        ]
        @named fa_model = ODESystem(fa_eqs, t,
                            systems=[fa, one, two, three, out1, out2])
        sys = structural_simplify(fa_model)

        u0 = []

        prob = ODEProblem(sys, u0, (0, 1.5))
        sol = solve(prob, Rosenbrock23())

        @test sol.retcode == :Success
        @test sol[fa.y0.val] == sol[fa.sum] == _xor(sol[one.d.val], sol[two.d.val], sol[three.d.val])
    end
end

@testset "Multiplexers" begin
    @named mux2x1 = MUX(N=2)
    @named mux4x1 = MUX(N=4)
    @named demux1x2 = DEMUX(N=2)
    out = map(0:7) do i
        DigitalPin(; name=Symbol(:out, i))
    end
    input = map(0:7) do i
        rand([Set(name=Symbol(:s, i)), Reset(name=Symbol(:r, i))])
    end
    @named select0 = Set()
    @named select1 = Reset()

    @info "Building a 2:1 MUX and 1:2 DEMUX..."
    mux_eqs = [
        connect(input[1].d, mux2x1.d0)
        connect(input[2].d, mux2x1.d1)
        connect(select0.d, mux2x1.s0)
        connect(mux2x1.y, demux1x2.d)
        connect(select1.d, demux1x2.s0)
        connect(out[1], demux1x2.y0)
        connect(out[2], demux1x2.y1)
    ]
    @named mux_model = ODESystem(mux_eqs, t,
                                systems=[mux2x1, demux1x2,
                                        out[2], out[1], select0, select1,
                                        input[2], input[1]])
    sys = alias_elimination(mux_model)

    u0 = []
    prob = ODEProblem(sys, u0, (0, 1.5))
    sol = solve(prob, Rosenbrock23())

    @test sol.retcode == :Success
    @test sol[mux2x1.y.val] == sol[mux2x1.d1.val]
    @test sol[mux2x1.y.val] == sol[demux1x2.y0.val]

    @info "Building a 4:1 MUX..."
    @named seta = Set()
    @named setb = Set()
    @named reseta = Reset()
    @named resetb = Reset()
    select0, select1 = [reseta, seta], [resetb, setb]
    for idx1 in 1:2, idx0 in 1:2
        mux_eqs = [
            connect(input[1].d, mux4x1.d0)
            connect(input[2].d, mux4x1.d1)
            connect(input[3].d, mux4x1.d2)
            connect(input[4].d, mux4x1.d3)
            connect(select0[idx0].d, mux4x1.s0)
            connect(select1[idx1].d, mux4x1.s1)
            connect(out[1], mux4x1.y)
        ]
        @named mux_model = ODESystem(mux_eqs, t,
                                    systems=[mux4x1, out[1], select0[idx0], select1[idx1],
                                            input[1], input[2], input[3], input[4]])
        sys = alias_elimination(mux_model)

        u0 = []
        prob = ODEProblem(sys, u0, (0, 1.5))
        sol = solve(prob, Rosenbrock23())

        @test sol.retcode == :Success
        (idx1 == 1 && idx0 == 1) && @test sol[mux4x1.y.val] == sol[mux4x1.d0.val]
        (idx1 == 1 && idx0 == 2) && @test sol[mux4x1.y.val] == sol[mux4x1.d1.val]
        (idx1 == 2 && idx0 == 1) && @test sol[mux4x1.y.val] == sol[mux4x1.d2.val]
        (idx1 == 2 && idx0 == 2) && @test sol[mux4x1.y.val] == sol[mux4x1.d3.val]
    end
end

@testset "Demultiplexers" begin
    @named demux1x2 = DEMUX(N=2)
    @named demux1x4 = DEMUX(N=4)
    @named demux1x8 = DEMUX(N=8)
    @named input = Set()
    out = map(1:4) do i
        DigitalPin(; name=Symbol(:out, i))
    end

    @named seta = Set()
    @named setb = Set()
    @named reseta = Reset()
    @named resetb = Reset()
    select0, select1 = [reseta, seta], [resetb, setb]
    for idx1 in 1:2, idx0 in 1:2

        @info "Building 1:4 DEMUX..."
        eqs = [
            connect(input.d, demux1x4.d)
            connect(select0[idx0].d, demux1x4.s0)
            connect(select1[idx1].d, demux1x4.s1)
            connect(out[1], demux1x4.y0)
            connect(out[2], demux1x4.y1)
            connect(out[3], demux1x4.y2)
            connect(out[4], demux1x4.y3)
        ]
        @named demux_model = ODESystem(eqs, t,
                                    systems=[demux1x4, select0[idx0], select1[idx1], input,
                                    out[1], out[2], out[3], out[4]])
        sys = alias_elimination(demux_model)

        u0 = []
        prob = ODEProblem(sys, u0, (0, 1.5))
        sol = solve(prob, Rosenbrock23())

        @test sol.retcode == :Success
        (idx1 == 1 && idx0 == 1) && @test sol[demux1x4.d.val] == sol[demux1x4.y0.val]
        (idx1 == 1 && idx0 == 2) && @test sol[demux1x4.d.val] == sol[demux1x4.y1.val]
        (idx1 == 2 && idx0 == 1) && @test sol[demux1x4.d.val] == sol[demux1x4.y2.val]
        (idx1 == 2 && idx0 == 2) && @test sol[demux1x4.d.val] == sol[demux1x4.y3.val]
    end
end

@testset "Encoder and Decoder" begin
    @named enc4x2 = Encoder(N=4)
    @named dec2x4 = Decoder(n=2)
    out = map(1:4) do i
        DigitalPin(; name=Symbol(:out, i))
    end

    input = map(0:7) do i
        rand([Set(name=Symbol(:s, i)), Reset(name=Symbol(:r, i))])
    end

    @info "Building a 4:2 Encoder..."
    enc_eqs = [
        connect(input[1].d, enc4x2.d0)
        connect(input[2].d, enc4x2.d1)
        connect(input[3].d, enc4x2.d2)
        connect(input[4].d, enc4x2.d3)
        connect(out[1], enc4x2.y0)
        connect(out[2], enc4x2.y1)
    ]
    @named enc_model = ODESystem(enc_eqs, t,
                                 systems=[enc4x2, out[1], out[2],
                                          input[1], input[2], input[3], input[4]])
    sys = alias_elimination(enc_model)

    u0 = []
    prob = ODEProblem(sys, u0, (0, 1.5))
    sol = solve(prob, Rosenbrock23())

    @test sol.retcode == :Success
    @test sol[enc4x2.y0.val] == _or.(sol[input[4].d.val], sol[input[2].d.val])
    @test sol[enc4x2.y1.val] == _or.(sol[input[4].d.val], sol[input[3].d.val])

    @info "Building a 2:4 Decoder..."
    dec_eqs = [
        connect(input[1].d, dec2x4.d0)
        connect(input[2].d, dec2x4.d1)
        connect(out[1], dec2x4.y0)
        connect(out[2], dec2x4.y1)
        connect(out[3], dec2x4.y2)
        connect(out[4], dec2x4.y3)
    ]
    @named dec_model = ODESystem(dec_eqs, t,
        systems=[dec2x4, out[1], out[2],
                 out[3], out[4],
                 input[1], input[2]])
    sys = alias_elimination(dec_model)

    u0 = []
    prob = ODEProblem(sys, u0, (0, 1.5))
    sol = solve(prob, Rosenbrock23())

    @test sol.retcode == :Success
    @test sol[dec2x4.y0.val] == _and.(_not.(sol[input[2].d.val]),
                                                _not.(sol[input[1].d.val]))
    @test sol[dec2x4.y1.val] == _and.(_not.(sol[input[2].d.val]),
                                                (sol[input[1].d.val]))
    @test sol[dec2x4.y2.val] == _and.((sol[input[2].d.val]),
                                                _not.(sol[input[1].d.val]))
    @test sol[dec2x4.y3.val] == _and.((sol[input[2].d.val]), (sol[input[1].d.val]))
end

@testset "Sources and DigitalPin" begin
    @named out = DigitalPin()
    @named and = And()
    @named setₐ = Set()
    @named resetₐ = Reset()
    @named setᵦ = Set()
    @named resetᵦ = Reset()

    for α in [setₐ, resetₐ], β in [setᵦ, resetᵦ]
        eqs = [connect(α.d, and.x1)
            connect(β.d, and.x2)
            connect(out, and.y)
            ]

        @named pul = ODESystem(eqs, t, systems=[α, β, and, out])
        sys = alias_elimination(pul)
        # sys = structural_simplify(pul)
        u0 = []
        prob = ODEProblem(sys, u0, (0, 1.5))
        sol = solve(prob, Rosenbrock23())

        @test sol.retcode == :Success
        @test ModelingToolkitStandardLibrary._and.(sol[and.x1.val], sol[and.x2.val]) == sol[and.y.val]
    end

    # Test Pulse
    @named pulse = Pulse()
    @named pulseD = PulseDiff()
    @named not = Not()
    eqs = [connect(pulseD.d, not.x)
           connect(out, not.y)
          ]
    @named pul = ODESystem(eqs, t, systems=[pulseD, not, out])
    sys = alias_elimination(pul)
    # sys = structural_simplify(pul)

    u0 = []
    prob = ODEProblem(sys, u0, (0, 1.5))
    # sol = solve(prob, Rosenbrock23())
end


text:
# Runge-Kutta 3rd/4th order coefficients
if RKo == 3     # version 2
    const RKa = Numtype.([1/4,0.,3/4])
    const RKb = Numtype.([1/3,2/3])
elseif RKo == 4
    const RKa = Numtype.([1/6,1/3,1/3,1/6])
    const RKb = Numtype.([.5,.5,1.])
end

# for the ghost point copy/tangential boundary conditions
const one_minus_α = Numtype(1-lbc)

# for the laplace operator
const minus_4 = Numtype(-4.)

# for the interpolation functions
const one_half = Numtype(0.5)
const one_twelve = Numtype(1/12)
const one_quart = Numtype(0.25)

# will be used for the Bernoulli potential
const g = Numtype(gravity)

# for the bottom friction (include grid spacing as gradient operators are dimensionless)
const c_D = Numtype(-Δ*drag)
const r_B = Numtype(-Δ/(τdrag*24*3600))  # [1/s]

# frequency [1/s] of interface relaxation (for non-dimensional gradients, γ contains the grid spacing Δ)
const γ = Numtype(Δ/(t_relax*3600*24))

# for biharmonic diffusion
const cSmag = Numtype(-c_smag)
const νB = Numtype(-ν_const/30000)   # linear scaling based on 540m^s/s at Δ=30km

# for semi-Lagrangian advection / interpolation
const zeero = Numtype(0.)
const oone = Numtype(1.)
const r_SST = Numtype(dtadvint/(τSST*3600*24))    # [1]

# for analysis the old operators with boundary conditions are kept. Constants for these
const one_over_Δ = Numtype(1/Δ)
const α = Numtype(lbc)
const minus_α = Numtype(-lbc)
const one_minus_α_half = Numtype(1-0.5*lbc)
const α_over_Δ = Numtype(lbc/Δ)
const one_quarter = Numtype(0.25)
const minus_3_minus_α = Numtype(-3-lbc)

ext: jl

max_stars_repo_name: eschnett/juls

max_issues_repo_name: eschnett/juls

max_forks_repo_name: eschnett/juls


text:
import CSV
import TimeSeries

@testset "Test Timeseries formats" begin
    formats = [
        (IS.TimeseriesFormatYMDPeriodAsColumn,
         joinpath(FORECASTS_DIR, "YMDPeriodAsColumn.csv"),
         nothing,
        ),
        (IS.TimeseriesFormatYMDPeriodAsHeader,
         joinpath(FORECASTS_DIR, "YMDPeriodAsHeader.csv"),
         "fake",
        ),
        (IS.TimeseriesFormatComponentsAsColumnsNoTime,
         joinpath(FORECASTS_DIR, "ComponentsAsColumnsNoTime.csv"),
         nothing),
        # TODO: add a file that has a column name with a DateTime.
        # TODO: add a file that more than one unique timestamp so that we can fully test
        # IS.get_step_time().
    ]

    for (format, filename, component_name) in formats
        file = CSV.File(filename)
        @test format == IS.get_timeseries_format(file)

        data = IS.read_time_series(filename, component_name)
        @test data isa TimeSeries.TimeArray
    end
end


text:
using DynamicEnergyBudgets, Unitful, Test

reftemp = 310.0u"K"
arrtemp = 2000.0u"K"
lowerbound = 280.0u"K"
arrlower = 20000.0u"K"
upperbound = 315.0u"K"
arrupper = 70000.0u"K"

tc = TempCorrLowerUpper(reftemp, arrtemp, lowerbound, arrlower, upperbound, arrupper)
@test tempcorr(100.0u"°C", tc) ≈ 0.0 atol=1e-10
@test tempcorr(-100.0u"°C", tc) ≈ 0.0 atol=1e-10
@test tempcorr(reftemp, tc) == 1.0

tc = TempCorrLower(reftemp, arrtemp, lowerbound, arrlower)
@test tempcorr(100.0u"°C", tc) > 1.0
@test tempcorr(-100.0u"°C", tc) ≈ 0.0 atol=1e-10
@test tempcorr(reftemp, tc) == 1.0

tc = TempCorr(reftemp, arrtemp)
@test tempcorr(100.0u"°C", tc) > 1.0
@test tempcorr(-100.0u"°C", tc) > 0.0
@test tempcorr(-100.0u"°C", tc) < 1.0
@test tempcorr(reftemp, tc) == 1.0


text:
# examples/blink.jl
#   This is an example of an Application that uses a package containing binary
#   dependencies, in this case Blink.jl.
#   When building such an Application, you must provide its dependencies
#   via `-L` and `-R`.
# Build this file with the following command:
# $ julia ~/.julia/v0.6/ApplicationBuilder/build_app.jl \
#      -R ~/.julia/v0.6/Blink/deps/Julia.app \
#      -R ~/.julia/v0.6/Blink/src/AtomShell/main.js \
#      -R ~/.julia/v0.6/Blink/src/content/main.html \
#      -R ~/.julia/v0.6/Blink/res \
#      -L ~/.julia/v0.6/HttpParser/deps/usr/lib/libhttp_parser.dylib \
#      -L ~/.julia/v0.6/MbedTLS/deps/usr/lib/libmbedcrypto.2.7.1.dylib \
#      examples/blink.jl "HelloBlink"

using Blink

using ApplicationBuilderAppUtils

# In order to distribute this Application, we've copied all its dependencies
# into the .app bundle via the -R and -L build flags.
# Now, here, we have to update the Packages that use those dependencies to find
# them in the right place. These changes take place at compile time. We override
# the paths to all be relative paths, so that the app bundle can be moved.
if get(ENV, "COMPILING_APPLE_BUNDLE", "false") == "true"
    println("Overriding Blink dependency paths.")
    println("Overriding Blink dependency paths.")
    Core.eval(Blink.AtomShell, :(_electron = "Julia.app/Contents/MacOS/Julia"))
    Core.eval(Blink.AtomShell, :(mainjs = "main.js"))
    Core.eval(Blink, :(buzz = "main.html"))
    Core.eval(Blink, :(resources = Dict("spinner.css" => "res/spinner.css",
                             "blink.js" => "res/blink.js",
                             "blink.css" => "res/blink.css",
                             "reset.css" => "res/reset.css")))
    # Clear out Blink.__inits__, since it will attempt to evaluate hardcoded paths.
    # (We've defined all the variables manually, above: `resources` and `port`.)
    Core.eval(Blink, :(empty!(__inits__)))

    Core.eval(Blink.Mux.HTTP.MbedTLS, :(const libmbedcrypto = basename(libmbedcrypto)))

#    WebSockets = Blink.WebSockets
#    Core.eval(WebSockets, :(using HttpServer))  # needed to cause @require lines to execute & compile
#    Core.eval(WebSockets,
#        :(include(joinpath(Pkg.dir("WebSockets"),"src/HttpServer.jl"))))  # Manually load this from the @requires line.

    println("Done changing dependencies.")
end

# Simple HTML example: "Hello" with an input field to change name being greeted.
html() = """
    Hello, World! -- Love, Blink.

Type a new name and press enter: <input id="nameentry"></input>

 """

function helloFromBlink()
    # Set Blink port randomly before anything else.
    Blink.port[] = get(ENV, "BLINK_PORT", rand(2_000:10_000))

    # Create Blink window and load HTML.
    win = Blink.Window(Blink.shell(), Dict(:width=>850)); sleep(5.0)
    Blink.body!(win, html(); fade=false) ; sleep(1)
    Blink.tools(win)
    sleep(2)  # wait for js to initialize

    # Example javascript interaction.
    Blink.@js_ win console.log("HELLO!")
    cwd_str = "pwd(): $(pwd())"
    Blink.@js_ win console.log($cwd_str)
    Blink.@js_ win setName("Nathan")  # Calling a js function.

    # Set up the js to call a julia callback that calls back to js! haha :)
    # (This could be all done within js, but it's fun to see the full round trip.)
    Blink.handlers(win)["changeNameJulia"] = (n) -> (println("msg from js: $n"); Blink.@js_ win setName($n))
    # Set the input field to call the above julia callback when it's changed.
    Blink.@js_ win nameEntry.onchange = (e) -> (Blink.msg("changeNameJulia", nameEntry.value); console.log("sent msg to julia!"); e.returnValue=false)

    while Blink.active(win)
        sleep(1)
    end
end

Base.@ccallable function julia_main(args::Vector{String})::Cint
    ApplicationBuilderAppUtils.cd_to_bundle_resources()

    # Apparently starting Electron too quickly means the OS doesn't get a
    # chance to find the name of the application...
    sleep(2)

    helloFromBlink()
    return 0
end


text:
const SKIP_PM_VALIDATION = false

const SYSTEM_KWARGS = Set((
    :branch_name_formatter,
    :bus_name_formatter,
    :config_path,
    :frequency,
    :gen_name_formatter,
    :generator_mapping,
    :internal,
    :load_name_formatter,
    :load_zone_formatter,
    :runchecks,
    :shunt_name_formatter,
    :time_series_directory,
    :time_series_in_memory,
    :time_series_read_only,
    :timeseries_metadata_file,
    :unit_system,
    :pm_data_corrections,
    :import_all,
))

# This will be used in the future to handle serialization changes.
const DATA_FORMAT_VERSION = "1.0.1"

"""
System

A power system defined by fields for base_power, components, and time series.

```julia
System(base_power)
System(base_power, buses, components...)
System(base_power, buses, generators, loads, branches, storage, services; kwargs...)
System(base_power, buses, generators, loads; kwargs...)
System(file; kwargs...)
System(; buses, generators, loads, branches, storage, base_power, services, kwargs...)
System(; kwargs...)
```

# Arguments
- `base_power::Float64`: the base power value for the system
- `buses::Vector{Bus}`: an array of buses
- `components...`: Each element must be an iterable containing subtypes of Component.

# Keyword arguments
- `ext::Dict`: Contains user-defined parameters. Should only contain standard types.
- `runchecks::Bool`: Run available checks on input fields and when add_component! is called.
  Throws InvalidRange if an error is found.
- `time_series_in_memory::Bool=false`: Store time series data in memory instead of HDF5.
- `config_path::String`: specify path to validation config file
"""
struct System <: IS.InfrastructureSystemsType
    data::IS.SystemData
    frequency::Float64 # [Hz]
    bus_numbers::Set{Int}
    runchecks::Base.RefValue{Bool}
    units_settings::SystemUnitsSettings
    internal::IS.InfrastructureSystemsInternal

    function System(
        data,
        units_settings::SystemUnitsSettings,
        internal::IS.InfrastructureSystemsInternal;
        kwargs...,
    )
        # Note to devs: if you add parameters to kwargs then consider whether they need
        # special handling in the deserialization function in this file.
        # See deserialize for System.

        # Implement a strict check here to make sure that SYSTEM_KWARGS can be used
        # elsewhere.
        unsupported = setdiff(keys(kwargs), SYSTEM_KWARGS)
        !isempty(unsupported) && error("Unsupported kwargs = $unsupported")
        if !isnothing(get(kwargs, :unit_system, nothing))
            @warn(
                "unit_system kwarg ignored. The value in SystemUnitsSetting takes precedence"
            )
        end
        bus_numbers = Set{Int}()
        frequency = get(kwargs, :frequency, DEFAULT_SYSTEM_FREQUENCY)
        runchecks = Base.RefValue{Bool}(get(kwargs, :runchecks, true))
        sys = new(data, frequency, bus_numbers, runchecks, units_settings, internal)
        return sys
    end
end

function System(data, base_power, internal; kwargs...)
    unit_system_ = get(kwargs, :unit_system, "SYSTEM_BASE")
    unit_system = UNIT_SYSTEM_MAPPING[unit_system_]
    units_settings = SystemUnitsSettings(base_power, unit_system)
    return System(data, units_settings, internal; kwargs...)
end

"""Construct an empty `System`. Useful for building a System while parsing raw data."""
function System(base_power; kwargs...)
    return System(_create_system_data_from_kwargs(; kwargs...), base_power; kwargs...)
end

"""Construct a `System` from `InfrastructureSystems.SystemData`"""
function System(data, base_power; internal = IS.InfrastructureSystemsInternal(), kwargs...)
    return System(data, base_power, internal; kwargs...)
end

"""
System constructor when components are constructed externally.
"""
function System(base_power::Float64, buses::Vector{Bus}, components...; kwargs...)
    data = _create_system_data_from_kwargs(; kwargs...)
    sys = System(data, base_power; kwargs...)

    for bus in buses
        add_component!(sys, bus)
    end

    for component in Iterators.flatten(components)
        add_component!(sys, component)
    end

    if get(kwargs, :runchecks, true)
        check(sys)
    end

    return sys
end

"""Constructs a non-functional System for demo purposes."""
function System(
    ::Nothing;
    buses = [
        Bus(;
            number = 0,
            name = "init",
            bustype = BusTypes.REF,
            angle = 0.0,
            magnitude = 0.0,
            voltage_limits = (min = 0.0, max = 0.0),
            base_voltage = nothing,
            area = nothing,
            load_zone = nothing,
            ext = Dict{String, Any}(),
        ),
    ],
    generators = [ThermalStandard(nothing), RenewableFix(nothing)],
    loads = [PowerLoad(nothing)],
    branches = nothing,
    storage = nothing,
    base_power = 100.0,
    services = nothing,
    kwargs...,
)
    _services = isnothing(services) ? [] : services
    return System(base_power, buses, generators, loads, _services; kwargs...)
end

"""Constructs a System from a file path ending with .m, .RAW, or .json"""
function System(file_path::AbstractString; kwargs...)
    ext = splitext(file_path)[2]
    if lowercase(ext) in [".m", ".raw"]
        pm_kwargs = Dict(k => v for (k, v) in kwargs if !in(k, SYSTEM_KWARGS))
        return System(PowerModelsData(file_path; pm_kwargs...); kwargs...)
    elseif lowercase(ext) == ".json"
        unsupported = setdiff(keys(kwargs), SYSTEM_KWARGS)
        !isempty(unsupported) && error("Unsupported kwargs = $unsupported")
        runchecks = get(kwargs, :runchecks, true)
        # File paths in the JSON are relative. Temporarily change to this directory in order
        # to find all dependent files.
        orig_dir = pwd()
        new_dir = dirname(file_path)
        if isempty(new_dir)
            new_dir = "."
        end
        cd(new_dir)
        try
            time_series_read_only = get(kwargs, :time_series_read_only, false)
            time_series_directory = get(kwargs, :time_series_directory, nothing)
            sys = deserialize(
                System,
                basename(file_path);
                time_series_read_only = time_series_read_only,
                runchecks = runchecks,
                time_series_directory = time_series_directory,
            )
            runchecks && check(sys)
            return sys
        finally
            cd(orig_dir)
        end
    else
        throw(DataFormatError("$file_path is not a supported file type"))
    end
end

"""
Parse static and dynamic data directly from PSS/e text files. Automatically generates
all the relationships between the available dynamic injection models and the static counterpart

Each dictionary indexed by id contains a vector with 5 of its components:
* Machine
* Shaft
* AVR
* TurbineGov
* PSS

Files must be parsed from a .raw file (PTI data format) and a .dyr file.

## Examples:
```julia
raw_file = "Example.raw"
dyr_file = "Example.dyr"
sys = System(raw_file, dyr_file)
```

"""
function System(sys_file::AbstractString, dyr_file::AbstractString; kwargs...)
    ext = splitext(sys_file)[2]
    if lowercase(ext) in [".raw"]
        pm_kwargs = Dict(k => v for (k, v) in kwargs if !in(k, SYSTEM_KWARGS))
        sys = System(PowerModelsData(sys_file; pm_kwargs...); kwargs...)
    else
        throw(DataFormatError("$sys_file is not a .raw file type"))
    end
    bus_dict_gen = _parse_dyr_components(dyr_file)
    add_dyn_injectors!(sys, bus_dict_gen)
    return sys
end

"""
Serializes a system to a JSON string.

# Arguments
- `sys::System`: system
- `filename::AbstractString`: filename to write

# Keyword arguments
- `force::Bool = false`: whether to overwrite existing files
- `check::Bool = false`: whether to run system validation checks

Refer to [`check_component`](@ref) for exceptions thrown if `check = true`.
"""
function IS.to_json(sys::System, filename::AbstractString; force = false, runchecks = false)
    if runchecks
        check(sys)
        check_components(sys)
    end

    IS.prepare_for_serialization!(sys.data, filename; force = force)
    data = to_json(sys)
    open(filename, "w") do io
        write(io, data)
    end
    @info "Serialized System to $filename"
end

function Base.deepcopy(sys::System)
    # We store time series in an HDF5 file that would not be copied as part of deepcopy.
    # The HDF5 file could have data buffered in memory, so we would have to close it, make
    # a copy, and attach it to a new system.
    # A simpler solution is to serialize to a tmp dir and deserialize.
    path = mktempdir()
    filename = joinpath(path, "sys.json")
    to_json(sys, filename)
    return System(filename)
end

"""
Return the internal of the system
"""
IS.get_internal(sys::System) = sys.internal

"""
Return a user-modifiable dictionary to store extra information.
"""
get_ext(sys::System) = IS.get_ext(sys.internal)

"""
Return the system's base power.
"""
get_base_power(sys::System) = sys.units_settings.base_value

"""
Return the system's frequency.
"""
get_frequency(sys::System) = sys.frequency

"""
Clear any value stored in ext.
"""
clear_ext!(sys::System) = IS.clear_ext!(sys.internal)

"""
Return true if checks are enabled on the system.
"""
get_runchecks(sys::System) = sys.runchecks[]

"""
Enable or disable system checks.
Applies to component addition as well as overall system consistency.
"""
function set_runchecks!(sys::System, value::Bool)
    sys.runchecks[] = value
    @info "Set runchecks to $value"
end

function set_units_setting!(
    component::Component,
    settings::Union{SystemUnitsSettings, Nothing},
)
    set_units_info!(get_internal(component), settings)
    return
end

"""
Sets the units base for the getter functions on the devices. It modifies the behavior of all getter functions
"""
function set_units_base_system!(system::System, settings::String)
    system.units_settings.unit_system = UNIT_SYSTEM_MAPPING[uppercase(settings)]
    @info "Unit System changed to $(UNIT_SYSTEM_MAPPING[uppercase(settings)])"
    return
end

function set_units_base_system!(system::System, settings::UnitSystem)
    system.units_settings.unit_system = settings
    @info "Unit System changed to $settings"
    return
end

function get_units_base(system::System)
    return string(system.units_settings.unit_system)
end

function get_units_setting(component::T) where {T <: Component}
    return get_units_info(get_internal(component))
end

function has_units_setting(component::T) where {T <: Component}
    return !isnothing(get_units_setting(component))
end

"""
Add a component to the system.

Throws ArgumentError if the component's name is already stored for its concrete type.
Throws ArgumentError if any Component-specific rule is violated.
Throws InvalidRange if any of the component's field values are outside of defined valid
range.

# Examples
```julia
sys = System(100.0)

# Add a single component.
add_component!(sys, bus)

# Add many at once.
buses = [bus1, bus2, bus3]
generators = [gen1, gen2, gen3]
foreach(x -> add_component!(sys, x), Iterators.flatten((buses, generators)))
```
"""
function add_component!(
    sys::System,
    component::T;
    skip_validation = false,
    kwargs...,
) where {T <: Component}
    set_units_setting!(component, sys.units_settings)
    @assert has_units_setting(component)

    check_attached_buses(sys, component)
    check_component_addition(sys, component; kwargs...)

    deserialization_in_progress = _is_deserialization_in_progress(sys)
    if !deserialization_in_progress
        # Services are attached to devices at deserialization time.
        check_for_services_on_addition(sys, component)
    end

    skip_validation = _validate_or_skip!(sys, component, skip_validation)
    _kwargs = Dict(k => v for (k, v) in kwargs if k !== :static_injector)
    IS.add_component!(
        sys.data,
        component;
        deserialization_in_progress = deserialization_in_progress,
        skip_validation = skip_validation,
        _kwargs...,
    )

    if !deserialization_in_progress
        # Whatever this may change should have been validated above in
        # check_component_addition, so this should not fail.
        # Doesn't run at deserialization time because the changes made by this function
        # occurred when the original addition ran and do not apply to that scenario.
        handle_component_addition!(sys, component; kwargs...)
        # Special condition required to populate the bus numbers in the system after
    elseif component isa Bus
        handle_component_addition!(sys, component; kwargs...)
    end

    return
end

"""
Add a dynamic injector to the system.

Throws ArgumentError if the name does not match the static_injector name.
Throws ArgumentError if the static_injector is not attached to the system.

All rules for the generic add_component! method also apply.
"""
function add_component!(
    sys::System,
    dyn_injector::DynamicInjection,
    static_injector::StaticInjection;
    kwargs...,
)
    add_component!(sys, dyn_injector; static_injector = static_injector, kwargs...)
end

function _add_service!(
    sys::System,
    service::Service,
    contributing_devices;
    skip_validation = false,
    kwargs...,
)
    skip_validation = _validate_or_skip!(sys, service, skip_validation)
    for device in contributing_devices
        device_type = typeof(device)
        if !(device_type <: Device)
            throw(ArgumentError("contributing_devices must be of type Device"))
        end
        throw_if_not_attached(device, sys)
    end

    set_units_setting!(service, sys.units_settings)
    # Since this isn't atomic, order is important. Add to system before adding to devices.
    IS.add_component!(sys.data, service; skip_validation = skip_validation, kwargs...)

    for device in contributing_devices
        add_service_internal!(device, service)
    end
end

"""
Similar to [`add_component!`](@ref) but for services.

# Arguments
- `sys::System`: system
- `service::Service`: service to add
- `contributing_devices`: Must be an iterable of type Device
"""
function add_service!(sys::System, service::Service, contributing_devices; kwargs...)
    _add_service!(sys, service, contributing_devices; kwargs...)
end

"""
Similar to [`add_component!`](@ref) but for services.

# Arguments
- `sys::System`: system
- `service::Service`: service to add
- `contributing_device::Device`: Valid Device
"""
function add_service!(sys::System, service::Service, contributing_device::Device; kwargs...)
    _add_service!(sys, service, [contributing_device]; kwargs...)
end

"""
Similar to [`add_service!`](@ref) but for Service and Device already stored in the system.
Performs validation checks on the device and the system

# Arguments
- `device::Device`: Device
- `service::Service`: Service
- `sys::System`: system
"""
function add_service!(device::Device, service::Service, sys::System)
    throw_if_not_attached(service, sys)
    throw_if_not_attached(device, sys)
    add_service_internal!(device, service)
end

"""
Similar to [`add_component!`](@ref) but for StaticReserveGroup.

# Arguments
- `sys::System`: system
- `service::StaticReserveGroup`: service to add
"""
function add_service!(
    sys::System,
    service::StaticReserveGroup;
    skip_validation = false,
    kwargs...,
)
    skip_validation = _validate_or_skip!(sys, service, skip_validation)

    for _service in get_contributing_services(service)
        throw_if_not_attached(_service, sys)
    end

    set_units_setting!(service, sys.units_settings)
    IS.add_component!(sys.data, service; skip_validation = skip_validation, kwargs...)
end

"""Set StaticReserveGroup contributing_services with check"""
function set_contributing_services!(
    sys::System,
    service::StaticReserveGroup,
    val::Vector{<:Service},
)
    for _service in val
        throw_if_not_attached(_service, sys)
    end
    service.contributing_services = val
end

"""
Similar to [`add_component!`](@ref) but for StaticReserveGroup.

# Arguments
- `sys::System`: system
- `service::StaticReserveGroup`: service to add
- `contributing_services`: contributing services to the group
"""
function add_service!(
    sys::System,
    service::StaticReserveGroup,
    contributing_services::Vector{<:Service};
    skip_validation = false,
    kwargs...,
)
    skip_validation = _validate_or_skip!(sys, service, skip_validation)
    set_contributing_services!(sys, service, contributing_services)

    set_units_setting!(service, sys.units_settings)
    IS.add_component!(sys.data, service; skip_validation = skip_validation, kwargs...)
end

"""
Add time series data from a metadata file or metadata descriptors.

# Arguments
- `sys::System`: system
- `metadata_file::AbstractString`: metadata file for timeseries
  that includes an array of IS.TimeSeriesFileMetadata instances or a vector.
- `resolution::DateTime.Period=nothing`: skip time series that don't match this resolution.
"""
function add_time_series!(sys::System, metadata_file::AbstractString; resolution = nothing)
    return IS.add_time_series_from_file_metadata!(
        sys.data,
        Component,
        metadata_file;
        resolution = resolution,
    )
end

"""
Add time series data from a metadata file or metadata descriptors.

# Arguments
- `sys::System`: system
- `timeseries_metadata::Vector{IS.TimeSeriesFileMetadata}`: metadata for timeseries
- `resolution::DateTime.Period=nothing`: skip time series that don't match this resolution.
"""
function add_time_series!(
    sys::System,
    file_metadata::Vector{IS.TimeSeriesFileMetadata};
    resolution = nothing,
)
    return IS.add_time_series_from_file_metadata!(
        sys.data,
        Component,
        file_metadata;
        resolution = resolution,
    )
end

function IS.add_time_series_from_file_metadata_internal!(
    data::IS.SystemData,
    ::Type{<:Component},
    cache::IS.TimeSeriesParsingCache,
    file_metadata::IS.TimeSeriesFileMetadata,
)
    IS.set_component!(file_metadata, data, PowerSystems)
    component = file_metadata.component
    if isnothing(component)
        return
    end

    ts = IS.make_time_series!(cache, file_metadata)
    if component isa AggregationTopology && file_metadata.scaling_factor_multiplier in
       ["get_max_active_power", "get_max_reactive_power"]
        uuids = Set{UUIDs.UUID}()
        for bus in _get_buses(data, component)
            push!(uuids, IS.get_uuid(bus))
        end
        for _component in (
            load for load in IS.get_components(ElectricLoad, data) if
            IS.get_uuid(get_bus(load)) in uuids
        )
            IS.add_time_series!(data, _component, ts; skip_if_present = true)
        end
        file_metadata.scaling_factor_multiplier =
            replace(file_metadata.scaling_factor_multiplier, "max" => "peak")
        area_ts = IS.make_time_series!(cache, file_metadata)
        IS.add_time_series!(data, component, area_ts; skip_if_present = true)
    else
        IS.add_time_series!(data, component, ts)
    end
end

"""
Iterates over all components.

# Examples
```julia
for component in iterate_components(sys)
    @show component
end
```

See also: [`get_components`](@ref)
"""
function iterate_components(sys::System)
    return IS.iterate_components(sys.data)
end

"""
Remove all components from the system.
"""
function clear_components!(sys::System)
    return IS.clear_components!(sys.data)
end

"""
Remove all components of type T from the system.

Throws ArgumentError if the type is not stored.
"""
function remove_components!(::Type{T}, sys::System) where {T <: Component}
    for component in IS.remove_components!(T, sys.data)
        handle_component_removal!(sys, component)
    end
end

function clear_units!(component::Component)
    get_internal(component).units_info = nothing
    return
end

"""
Remove a component from the system by its value.

Throws ArgumentError if the component is not stored.
"""
function remove_component!(sys::System, component::T) where {T <: Component}
    check_component_removal(sys, component)
    IS.remove_component!(sys.data, component)
    handle_component_removal!(sys, component)
    clear_units!(component)
end

"""
Throws ArgumentError if a PowerSystems rule blocks removal from the system.
"""
function check_component_removal(sys::System, service::T) where {T <: Service}
    if T == StaticReserveGroup
        return
    end
    groupservices = get_components(StaticReserveGroup, sys)
    for groupservice in groupservices
        if service ∈ get_contributing_services(groupservice)
            throw(
                ArgumentError(
                    "service $(get_name(service)) cannot be removed with an attached StaticReserveGroup",
                ),
            )
            return
        end
    end
end

"""
Remove a component from the system by its name.

Throws ArgumentError if the component is not stored.
"""
function remove_component!(
    ::Type{T},
    sys::System,
    name::AbstractString,
) where {T <: Component}
    component = IS.remove_component!(T, sys.data, name)
    handle_component_removal!(sys, component)
end

"""
Check to see if the component of type T with name exists.
"""
function has_component(::Type{T}, sys::System, name::AbstractString) where {T <: Component}
    return IS.has_component(T, sys.data.components, name)
end

"""
Get the component of type T with name. Returns nothing if no component matches. If T is an abstract
type then the names of components across all subtypes of T must be unique.

See [`get_components_by_name`](@ref) for abstract types with non-unique names across subtypes.

Throws ArgumentError if T is not a concrete type and there is more than one component with
    requested name
"""
function get_component(::Type{T}, sys::System, name::AbstractString) where {T <: Component}
    return IS.get_component(T, sys.data, name)
end

"""
Returns an iterator of components. T can be concrete or abstract.
Call collect on the result if an array is desired.

# Examples
```julia
iter = PowerSystems.get_components(ThermalStandard, sys)
iter = PowerSystems.get_components(Generator, sys)
iter = PowerSystems.get_components(Generator, sys, x->(PowerSystems.get_available(x)))
generators = collect(PowerSystems.get_components(Generator, sys))
```

See also: [`iterate_components`](@ref)
"""
function get_components(::Type{T}, sys::System) where {T <: Component}
    return IS.get_components(T, sys.data, nothing)
end

# These two methods are defined independently instead of  filter_func::Union{Function, Nothing} = nothing
# because of a documenter error
# that has no relation with the code https://github.com/JuliaDocs/Documenter.jl/issues/1296
#
function get_components(
    ::Type{T},
    sys::System,
    filter_func::Function,
) where {T <: Component}
    return IS.get_components(T, sys.data, filter_func)
end

# These are helper functions for debugging problems.
# Searches components linearly, and so is slow compared to the other get_component functions
get_component(sys::System, uuid::Base.UUID) = IS.get_component(sys.data, uuid)
get_component(sys::System, uuid::String) = IS.get_component(sys.data, Base.UUID(uuid))

function _get_components_by_name(abstract_types, data::IS.SystemData, name::AbstractString)
    _components = []
    for subtype in abstract_types
        component = IS.get_component(subtype, data, name)
        if !isnothing(component)
            push!(_components, component)
        end
    end

    return _components
end

"""
Get the components of abstract type T with name. Note that PowerSystems enforces unique
names on each concrete type but not across concrete types.

See [`get_component`](@ref) if the concrete type is known.

Throws ArgumentError if T is not an abstract type.
"""
function get_components_by_name(
    ::Type{T},
    sys::System,
    name::AbstractString,
) where {T <: Component}
    return IS.get_components_by_name(T, sys.data, name)
end

"""
Gets components availability. Requires type T to have the method get_available implemented.
"""

function get_available_components(::Type{T}, sys::System) where {T <: Component}
    return get_components(T, sys, x -> get_available(x))
end

"""
Return true if the component is attached to the system.
"""
function is_attached(component::T, sys::System) where {T <: Component}
    return is_attached(T, get_name(component), sys)
end

function is_attached(::Type{T}, name, sys::System) where {T <: Component}
    return !isnothing(get_component(T, sys, name))
end

"""
Throws ArgumentError if the component is not attached to the system.
"""
function throw_if_not_attached(component::Component, sys::System)
    if !is_attached(component, sys)
        throw(ArgumentError("$(summary(component)) is not attached to the system"))
    end
end

"""
Return a vector of devices contributing to the service.
"""
function get_contributing_devices(sys::System, service::T) where {T <: Service}
    throw_if_not_attached(service, sys)
    return [x for x in get_components(Device, sys) if has_service(x, service)]
end

struct ServiceContributingDevices
    service::Service
    contributing_devices::Vector{Device}
end

const ServiceContributingDevicesKey = NamedTuple{(:type, :name), Tuple{DataType, String}}
const ServiceContributingDevicesMapping =
    Dict{ServiceContributingDevicesKey, ServiceContributingDevices}

"""
Returns a ServiceContributingDevices object.
"""
function _get_contributing_devices(sys::System, service::T) where {T <: Service}
    uuid = IS.get_uuid(service)
    devices = ServiceContributingDevices(service, Vector{Device}())
    for device in get_components(Device, sys)
        for _service in get_services(device)
            if IS.get_uuid(_service) == uuid
                push!(devices.contributing_devices, device)
                break
            end
        end
    end
    return devices
end

"""
Return an instance of ServiceContributingDevicesMapping.
"""
function get_contributing_device_mapping(sys::System)
    services = ServiceContributingDevicesMapping()
    for service in get_components(Service, sys)
        key = ServiceContributingDevicesKey((typeof(service), get_name(service)))
        services[key] = _get_contributing_devices(sys, service)
    end

    return services
end

"""
Return a vector of components with buses in the AggregationTopology.
"""
function get_components_in_aggregation_topology(
    ::Type{T},
    sys::System,
    aggregator::AggregationTopology,
) where {T <: StaticInjection}
    buses = Set{String}((get_name(x) for x in get_buses(sys, aggregator)))
    components = Vector{T}()
    for component in get_components(T, sys)
        bus = get_bus(component)
        bus_name = get_name(bus)
        if bus_name in buses
            push!(components, component)
        end
    end

    return components
end

"""
Return a mapping of AggregationTopology name to vector of buses within it.
"""
function get_aggregation_topology_mapping(
    ::Type{T},
    sys::System,
) where {T <: AggregationTopology}
    mapping = Dict{String, Vector{Bus}}()
    accessor_func = get_aggregation_topology_accessor(T)
    for bus in get_components(Bus, sys)
        aggregator = accessor_func(bus)
        name = get_name(aggregator)
        buses = get(mapping, name, nothing)
        if isnothing(buses)
            mapping[name] = Vector{Bus}([bus])
        else
            push!(buses, bus)
        end
    end

    return mapping
end

"""
Return a vector of buses contained within the AggregationTopology.
"""
function get_buses(sys::System, aggregator::AggregationTopology)
    return _get_buses(sys.data, aggregator)
end

function _get_buses(data::IS.SystemData, aggregator::T) where {T <: AggregationTopology}
    accessor_func = get_aggregation_topology_accessor(T)
    buses = Vector{Bus}()
    for bus in IS.get_components(Bus, data)
        _aggregator = accessor_func(bus)
        if IS.get_uuid(_aggregator) == IS.get_uuid(aggregator)
            push!(buses, bus)
        end
    end

    return buses
end

"""
Add time series data to a component.

Throws ArgumentError if the component is not stored in the system.

"""
function add_time_series!(sys::System, component::Component, time_series::TimeSeriesData)
    return IS.add_time_series!(sys.data, component, time_series)
end

"""
Add the same time series data to multiple components.

This is significantly more efficent than calling `add_time_series!` for each component
individually with the same data because in this case, only one time series array is stored.

Throws ArgumentError if a component is not stored in the system.
"""
function add_time_series!(sys::System, components, time_series::TimeSeriesData)
    return IS.add_time_series!(sys.data, components, time_series)
end

"""
Efficiently add all time series data in one component to another by copying the underlying
references.

# Arguments
- `dst::Component`: Destination component
- `src::Component`: Source component
- `name_mapping::Dict = nothing`: Optionally map src names to different dst names.
  If provided and src has a time series with a name not present in name_mapping, that
  time series will not copied. If name_mapping is nothing then all time series will be
  copied with src's names.
- `scaling_factor_multiplier_mapping::Dict = nothing`: Optionally map src
  scaling_factor_multipliers to dst scaling_factor_multipliers. Same behaviors as
  name_mapping.
"""
function copy_time_series!(
    dst::Component,
    src::Component;
    name_mapping::Union{Nothing, Dict{String, String}} = nothing,
    scaling_factor_multiplier_mapping::Union{Nothing, Dict{String, String}} = nothing,
)
    IS.copy_time_series!(
        dst,
        src,
        name_mapping = name_mapping,
        scaling_factor_multiplier_mapping = scaling_factor_multiplier_mapping,
    )
end

#=
# TODO 1.0: do we need this functionality? not currently present in IS
"""
Return a vector of time series data from a metadata file.

# Arguments
- `data::SystemData`: system
- `metadata_file::AbstractString`: path to metadata file
- `resolution::{Nothing, Dates.Period}`: skip data that doesn't match this resolution

See InfrastructureSystems.TimeSeriesFileMetadata for description of what the file
should contain.
"""
function make_time_series(sys::System, metadata_file::AbstractString; resolution = nothing)
    return IS.make_time_series(
        sys.data,
        metadata_file,
        PowerSystems;
        resolution = resolution,
    )
end

"""
Return a vector of time series data from a vector of TimeSeriesFileMetadata values.

# Arguments
- `data::SystemData`: system
- `timeseries_metadata::Vector{TimeSeriesFileMetadata}`: metadata values
- `resolution::{Nothing, Dates.Period}`: skip data that doesn't match this resolution
"""
function make_time_series(
    sys::System,
    metadata::Vector{IS.TimeSeriesFileMetadata};
    resolution = nothing,
)
    return IS.make_time_series!(sys.data, metadata; resolution = resolution)
end
=#

"""
Return the initial times for all forecasts.
"""
get_forecast_initial_times(sys::System) = IS.get_forecast_initial_times(sys.data)

"""
Return the total period covered by all forecasts.
"""
get_forecast_total_period(sys::System) = IS.get_forecast_total_period(sys.data)

"""
Return the window count for all forecasts.
"""
get_forecast_window_count(sys::System) = IS.get_forecast_window_count(sys.data)

"""
Return the horizon for all forecasts.
"""
get_forecast_horizon(sys::System) = IS.get_forecast_horizon(sys.data)

"""
Return the initial_timestamp for all forecasts.
"""
get_forecast_initial_timestamp(sys::System) = IS.get_forecast_initial_timestamp(sys.data)

"""
Return the interval for all forecasts.
"""
get_forecast_interval(sys::System) = IS.get_forecast_interval(sys.data)

"""
Return the resolution for all time series.
"""
get_time_series_resolution(sys::System) = IS.get_time_series_resolution(sys.data)

"""
Return an iterator of time series in order of initial time.

Note that passing a filter function can be much slower than the other filtering parameters
because it reads time series data from media.

Call `collect` on the result to get an array.

# Arguments
- `data::SystemData`: system
- `filter_func = nothing`: Only return time series for which this returns true.
- `type = nothing`: Only return time series with this type.
- `name = nothing`: Only return time series matching this value.

# Examples
```julia
for time_series in get_time_series_multiple(sys)
    @show time_series
end

ts = collect(get_time_series_multiple(sys; type = SingleTimeSeries))
```
"""
function IS.get_time_series_multiple(
    sys::System,
    filter_func = nothing;
    type = nothing,
    name = nothing,
)
    return get_time_series_multiple(sys.data, filter_func; type = type, name = name)
end

"""
Remove all time series data from the system.
"""
function clear_time_series!(sys::System)
    return IS.clear_time_series!(sys.data)
end

"""
Remove the time series data for a component.
"""
function remove_time_series!(
    sys::System,
    ::Type{T},
    component::Component,
    name::String,
) where {T <: TimeSeriesData}
    return IS.remove_time_series!(sys.data, T, component, name)
end

"""
Transform all instances of SingleTimeSeries to DeterministicSingleTimeSeries.
"""
function transform_single_time_series!(sys::System, horizon::Int, interval::Dates.Period)
    IS.transform_single_time_series!(
        sys.data,
        IS.DeterministicSingleTimeSeries,
        horizon,
        interval,
    )
end

"""
Sanitize component values.
"""
sanitize_component!(component::Component, sys::System) = true

"""
Validate the component fields using only those fields. Refer to
[`validate_component_with_system`](@ref) to use other System data for the
validation.

Return true if the instance is valid.
"""
validate_component(component::Component) = true

"""
Validate a component against System data. Return true if the instance is valid.

Refer to [`validate_component`](@ref) if the validation logic only requires data contained
within the instance.
"""
validate_component_with_system(component::Component, sys::System) = true

Base.@deprecate validate_struct(sys::System, component::Component) validate_component_with_system(
    component,
    sys,
) false

# Keeps the code working with IS.
IS.validate_struct(component::Component) = validate_component(component)

"""
Check system consistency and validity.
"""
function check(sys::System)
    buses = get_components(Bus, sys)
    slack_bus_check(buses)
    buscheck(buses)
    critical_components_check(sys)
    adequacy_check(sys)
    return
end

"""
Check the values of all components. See [`check_component`](@ref) for exceptions thrown.
"""
function check_components(sys::System)
    for component in iterate_components(sys)
        check_component(sys, component)
    end
end

"""
Check the values of a component.

Throws InvalidRange if any of the component's field values are outside of defined valid
range.

Throws InvalidValue if the custom validate method for the type fails its check.
"""
function check_component(sys::System, component::Component)
    if !validate_component_with_system(component, sys)
        throw(IS.InvalidValue("Invalid value for $component"))
    end
    IS.check_component(sys.data, component)
end

function IS.serialize(sys::T) where {T <: System}
    data = Dict{String, Any}()
    data["data_format_version"] = DATA_FORMAT_VERSION
    for field in fieldnames(T)
        # Exclude bus_numbers because they will get rebuilt during deserialization.
        if field != :bus_numbers
            data[string(field)] = serialize(getfield(sys, field))
        end
    end

    return data
end

function IS.deserialize(
    ::Type{System},
    filename::AbstractString;
    time_series_read_only = false,
    time_series_directory = nothing,
    kwargs...,
)
    raw = open(filename) do io
        JSON3.read(io, Dict)
    end

    # Read any field that is defined in System but optional for the constructors and not
    # already handled here.
    handled = ("data", "units_settings", "bus_numbers", "internal", "data_format_version")
    parsed_kwargs = Dict{Symbol, Any}()
    for field in setdiff(keys(raw), handled)
        parsed_kwargs[Symbol(field)] = raw[field]
    end

    # The user can override the serialized runchecks value by passing a kwarg here.
    if haskey(kwargs, :runchecks)
        parsed_kwargs[:runchecks] = kwargs[:runchecks]
    end

    units = IS.deserialize(SystemUnitsSettings, raw["units_settings"])
    data = IS.deserialize(
        IS.SystemData,
        raw["data"];
        time_series_read_only = time_series_read_only,
        time_series_directory = time_series_directory,
    )
    internal = IS.deserialize(InfrastructureSystemsInternal, raw["internal"])
    sys = System(data, units; internal = internal, kwargs...)

    if raw["data_format_version"] != DATA_FORMAT_VERSION
        pre_deserialize_conversion!(raw, sys)
    end

    ext = get_ext(sys)
    ext["deserialization_in_progress"] = true
    deserialize_components!(sys, raw["data"]["components"])
    pop!(ext, "deserialization_in_progress")
    isempty(ext) && clear_ext!(sys)

    if !get_runchecks(sys)
        @warn "The System was deserialized with checks disabled, and so was not validated."
    end

    if raw["data_format_version"] != DATA_FORMAT_VERSION
        post_deserialize_conversion!(sys, raw)
    end

    return sys
end

function deserialize_components!(sys::System, raw)
    # Convert the array of components into type-specific arrays to allow addition by type.
    data = Dict{Any, Vector{Dict}}()
    for component in raw
        type = IS.get_type_from_serialization_data(component)
        components = get(data, type, nothing)
        if components === nothing
            components = Vector{Dict}()
            data[type] = components
        end
        push!(components, component)
    end

    # Maintain a lookup of UUID to component because some component types encode
    # composed types as UUIDs instead of actual types.
    component_cache = Dict{Base.UUID, Component}()

    # Add each type to this as we parse.
    parsed_types = Set()

    function is_matching_type(type, types)
        return any(x -> type <: x, types)
    end

    function deserialize_and_add!(;
        skip_types = nothing,
        include_types = nothing,
        post_add_func = nothing,
    )
        for (type, components) in data
            type in parsed_types && continue
            if !isnothing(skip_types) && is_matching_type(type, skip_types)
                continue
            end
            if !isnothing(include_types) && !is_matching_type(type, include_types)
                continue
            end
            for component in components
                comp = deserialize(type, component, component_cache)
                add_component!(sys, comp)
                component_cache[IS.get_uuid(comp)] = comp
                if !isnothing(post_add_func)
                    post_add_func(comp)
                end
            end
            push!(parsed_types, type)
        end
    end

    # Run in order based on type composition.
    deserialize_and_add!(; include_types = [Area, LoadZone])
    deserialize_and_add!(; include_types = [AGC])
    deserialize_and_add!(; include_types = [Bus])
    deserialize_and_add!(;
        include_types = [Arc, Service],
        skip_types = [StaticReserveGroup],
    )
    deserialize_and_add!(; include_types = [Branch], skip_types = [DynamicBranch])
    deserialize_and_add!(; include_types = [DynamicBranch])
    # Static injection devices can contain dynamic injection devices.
    deserialize_and_add!(; include_types = [StaticReserveGroup, DynamicInjection])
    deserialize_and_add!()
end

"""
Return bus with name.
"""
function get_bus(sys::System, name::String)
    return get_component(Bus, sys, name)
end

"""
Return bus with bus_number.
"""
function get_bus(sys::System, bus_number::Int)
    for bus in get_components(Bus, sys)
        if bus.number == bus_number
            return bus
        end
    end

    return nothing
end

"""
Return all buses values with bus_numbers.
"""
function get_buses(sys::System, bus_numbers::Set{Int})
    buses = Vector{Bus}()
    for bus in get_components(Bus, sys)
        if bus.number in bus_numbers
            push!(buses, bus)
        end
    end

    return buses
end

function _is_deserialization_in_progress(sys::System)
    ext = get_ext(sys)
    return get(ext, "deserialization_in_progress", false)
end

check_for_services_on_addition(sys::System, component::Component) = nothing

function check_for_services_on_addition(sys::System, component::Device)
    if supports_services(component) && length(get_services(component)) > 0
        throw(ArgumentError("type Device cannot be added with services"))
    end
    return
end

# Needed because get_services returns the services of the underlying struct
function check_for_services_on_addition(sys::System, component::RegulationDevice)
    for d in get_services(component)
        isa(d, AGC) && throw(ArgumentError("type Device cannot be added with services"))
    end
    return
end

"""
Throws ArgumentError if any bus attached to the component is invalid.
"""
check_attached_buses(sys::System, component::Component) = nothing

function check_attached_buses(sys::System, component::StaticInjection)
    throw_if_not_attached(get_bus(component), sys)
end

function check_attached_buses(sys::System, component::Branch)
    throw_if_not_attached(get_from_bus(component), sys)
    throw_if_not_attached(get_to_bus(component), sys)
end

function check_attached_buses(sys::System, component::DynamicBranch)
    check_attached_buses(sys, get_branch(component))
end

function check_attached_buses(sys::System, component::Arc)
    throw_if_not_attached(get_from(component), sys)
    throw_if_not_attached(get_to(component), sys)
end

"""
Throws ArgumentError if a PowerSystems rule blocks addition to the system.

This method is tied with handle_component_addition!. If the methods are re-implemented for
a subtype then whatever is added in handle_component_addition! must be checked here.
"""
check_component_addition(sys::System, component::Component; kwargs...) = nothing

"""
Throws ArgumentError if a PowerSystems rule blocks removal from the system.
"""
check_component_removal(sys::System, component::Component) = nothing

function check_component_removal(sys::System, static_injector::StaticInjection)
    if get_dynamic_injector(static_injector) !== nothing
        name = get_name(static_injector)
        throw(ArgumentError("$name cannot be removed with an attached dynamic injector"))
    end
end

"""
Refer to docstring for check_component_addition!
"""
handle_component_addition!(sys::System, component::Component; kwargs...) = nothing

handle_component_removal!(sys::System, component::Component) = nothing

function check_component_addition(sys::System, branch::Branch; kwargs...)
    arc = get_arc(branch)
    throw_if_not_attached(get_from(arc), sys)
    throw_if_not_attached(get_to(arc), sys)
end

function check_component_addition(sys::System, dyn_branch::DynamicBranch; kwargs...)
    if !_is_deserialization_in_progress(sys)
        throw_if_not_attached(dyn_branch.branch, sys)
    end
    arc = get_arc(dyn_branch)
    throw_if_not_attached(get_from(arc), sys)
    throw_if_not_attached(get_to(arc), sys)
end

function check_component_addition(sys::System, dyn_injector::DynamicInjection; kwargs...)
    if _is_deserialization_in_progress(sys)
        # Ordering of component addition makes these checks impossible.
        return
    end

    static_injector = get(kwargs, :static_injector, nothing)
    if static_injector === nothing
        throw(
            ArgumentError("static_injector must be passed when adding a DynamicInjection"),
        )
    end

    if get_name(dyn_injector) != get_name(static_injector)
        throw(
            ArgumentError(
                "static_injector must have the same name as the DynamicInjection",
            ),
        )
    end

    throw_if_not_attached(static_injector, sys)
end

function check_component_addition(sys::System, bus::Bus; kwargs...)
    number = get_number(bus)
    if number in sys.bus_numbers
        throw(ArgumentError("bus number $number is already stored in the system"))
    end

    area = get_area(bus)
    if !isnothing(area)
        throw_if_not_attached(area, sys)
    end

    load_zone = get_load_zone(bus)
    if !isnothing(load_zone)
        throw_if_not_attached(load_zone, sys)
    end
end

function handle_component_addition!(sys::System, bus::Bus; kwargs...)
    number = get_number(bus)
    @assert !(number in sys.bus_numbers) "bus number $number is already stored"
    push!(sys.bus_numbers, number)
end

function handle_component_addition!(sys::System, component::RegulationDevice; kwargs...)
    copy_time_series!(component, component.device)
    if !isnothing(get_component(typeof(component.device), sys, get_name(component.device)))
        # This will not be true during deserialization, and so won't run then.
        remove_component!(sys, component.device)
        # The line above removed the component setting so needs to be added back
        set_units_setting!(component.device, component.internal.units_info)
    end
    return
end

function handle_component_addition!(sys::System, component::Branch; kwargs...)
    handle_component_addition_common!(sys, component)
end

function handle_component_addition!(sys::System, component::DynamicBranch; kwargs...)
    handle_component_addition_common!(sys, component)
    remove_component!(sys, component.branch)
end

function handle_component_addition!(sys::System, dyn_injector::DynamicInjection; kwargs...)
    static_injector = kwargs[:static_injector]
    static_base_power = get_base_power(static_injector)
    set_base_power!(dyn_injector, static_base_power)
    set_dynamic_injector!(static_injector, dyn_injector)
end

function handle_component_addition_common!(sys::System, component::Branch)
    # If this arc is already attached to the system, assign it to the branch.
    # Else, add it to the system.
    arc = get_arc(component)
    _arc = get_component(Arc, sys, get_name(arc))
    if isnothing(_arc)
        add_component!(sys, arc)
    else
        set_arc!(component, _arc)
    end
end

"""
Throws ArgumentError if the bus number is not stored in the system.
"""
function handle_component_removal!(sys::System, bus::Bus)
    number = get_number(bus)
    @assert number in sys.bus_numbers "bus number $number is not stored"
    pop!(sys.bus_numbers, number)
end

function handle_component_removal!(sys::System, device::Device)
    # This may have to be refactored if handle_component_removal! needs to be implemented
    # for a subtype.
    clear_services!(device)
end

function handle_component_removal!(sys::System, service::Service)
    for device in get_components(Device, sys)
        _remove_service!(device, service)
    end
end

function handle_component_removal!(sys::System, value::T) where {T <: AggregationTopology}
    for device in get_components(Bus, sys)
        if get_aggregation_topology_accessor(T)(device) == value
            _remove_aggregration_topology!(device, value)
        end
    end
end

function handle_component_removal!(sys::System, dyn_injector::DynamicInjection)
    injectors = get_components_by_name(StaticInjection, sys, get_name(dyn_injector))
    found = false
    for static_injector in injectors
        _dyn_injector = get_dynamic_injector(static_injector)
        _dyn_injector === nothing && continue
        if _dyn_injector === dyn_injector
            @assert !found
            set_dynamic_injector!(static_injector, nothing)
            found = true
        end
    end

    @assert found
end

"""
Return a sorted vector of bus numbers in the system.
"""
function get_bus_numbers(sys::System)
    return sort(collect(sys.bus_numbers))
end

function IS.compare_values(x::System, y::System)
    match = true

    if x.units_settings.unit_system != x.units_settings.unit_system
        @debug "unit system does not match" x.units_settings.unit_system y.units_settings.unit_system
        match = false
    end

    if get_base_power(x) != get_base_power(y)
        @debug "base_power does not match" get_base_power(x) get_base_power(y)
        match = false
    end

    if !IS.compare_values(x.data, y.data)
        @debug "SystemData values do not match"
        match = false
    end

    return match
end

function IS.compare_values(x::T, y::T) where {T <: Union{StaticInjection, DynamicInjection}}
    # Must implement this method because a device of one of these subtypes might have a
    # reference to its counterpart, and vice versa, and so infinite recursion will occur
    # in the default function.
    match = true
    for (fieldname, fieldtype) in zip(fieldnames(T), fieldtypes(T))
        val1 = getfield(x, fieldname)
        val2 = getfield(y, fieldname)
        if val1 isa StaticInjection || val2 isa DynamicInjection
            uuid1 = IS.get_uuid(val1)
            uuid2 = IS.get_uuid(val2)
            if uuid1 != uuid2
                @debug "values do not match" T fieldname uuid1 uuid2
                match = false
                break
            end
        elseif !isempty(fieldnames(typeof(val1)))
            if !IS.compare_values(val1, val2)
                @debug "values do not match" T fieldname val1 val2
                match = false
                break
            end
        elseif val1 isa AbstractArray
            if !IS.compare_values(val1, val2)
                match = false
            end
        else
            if val1 != val2
                @debug "values do not match" T fieldname val1 val2
                match = false
                break
            end
        end
    end

    return match
end

function _create_system_data_from_kwargs(; kwargs...)
    validation_descriptor_file = nothing
    time_series_in_memory = get(kwargs, :time_series_in_memory, false)
    time_series_directory = get(kwargs, :time_series_directory, nothing)
    validation_descriptor_file =
        get(kwargs, :config_path, POWER_SYSTEM_STRUCT_DESCRIPTOR_FILE)

    return IS.SystemData(;
        validation_descriptor_file = validation_descriptor_file,
        time_series_in_memory = time_series_in_memory,
        time_series_directory = time_series_directory,
    )
end

"""
Converts a Line component to a MonitoredLine component and replaces the original in the
system
"""
function convert_component!(
    linetype::Type{MonitoredLine},
    line::Line,
    sys::System;
    kwargs...,
)
    new_line = linetype(
        line.name,
        line.available,
        line.active_power_flow,
        line.reactive_power_flow,
        line.arc,
        line.r,
        line.x,
        line.b,
        (from_to = line.rate, to_from = line.rate),
        line.rate,
        line.angle_limits,
        line.services,
        line.ext,
        InfrastructureSystems.TimeSeriesContainer(),
        line.internal,
    )
    IS.assign_new_uuid!(line)
    add_component!(sys, new_line)
    copy_time_series!(new_line, line)
    remove_component!(sys, line)
end

"""
Converts a MonitoredLine component to a Line component and replaces the original in the
system
"""
function convert_component!(
    linetype::Type{Line},
    line::MonitoredLine,
    sys::System;
    kwargs...,
)
    force = get(kwargs, :force, false)
    if force
        @warn("Possible data loss converting from $(typeof(line)) to $linetype")
    else
        error(
            "Possible data loss converting from $(typeof(line)) to $linetype, add `force = true` to convert anyway.",
        )
    end

    new_line = linetype(
        line.name,
        line.available,
        line.active_power_flow,
        line.reactive_power_flow,
        line.arc,
        line.r,
        line.x,
        line.b,
        line.rate,
        line.angle_limits,
        line.services,
        line.ext,
        InfrastructureSystems.TimeSeriesContainer(),
        line.internal,
    )
    IS.assign_new_uuid!(line)
    add_component!(sys, new_line)
    copy_time_series!(new_line, line)
    remove_component!(sys, line)
end

function _validate_or_skip!(sys, component, skip_validation)
    if skip_validation && get_runchecks(sys)
        @warn(
            "skip_validation is deprecated; construct System with runchecks = true or call set_runchecks!. Disabling System.runchecks"
        )
        set_runchecks!(sys, false)
    end

    # Always skip if system checks are disabled.
    if !skip_validation && !get_runchecks(sys)
        skip_validation = true
    end

    if !skip_validation
        sanitize_component!(component, sys)
        if !validate_component_with_system(component, sys)
            throw(IS.InvalidValue("Invalid value for $component"))
        end
    end

    return skip_validation
end


text:
__precompile__()

module ImageTracking

using Images
using ImageFiltering
using Interpolations
using StaticArrays

include("core.jl")
include("optical_flow.jl")
include("haar.jl")

export

	# main functions
  optical_flow,

	# other functions
	haar_coordinates,
	haar_features,

	# other functions
	polynomial_expansion,

	# optical flow algorithms
	LK,
	Farneback

end


text:
"""
Created on Mon Apr 13 2020

@author: Yoann Pradat

Running NMF algorithms on a matrix that can be factorized exactly.
"""

include("../../../src/VNMF.jl")
using .VNMF
using Distributions
using Random
include("../util/plot_cost.jl")

gr() # set the backend to GR

println("current working directory: " * pwd())

#### # 1. SIMULATE DATA
#### # ################################################################################################################

#### simulate matrix for testing NMF
F = 10
N = 25
K = 5
rng = MersenneTwister(123)

#### truncated-normal factors with normal noise
dist_W = truncated(Normal(1, 1), 0, Inf)
dist_H = truncated(Normal(1, 1), 0, Inf)

W = rand(rng, dist_W, F, K)
H = rand(rng, dist_H, K, N)

V = W * H

fo = "../../plot/simulation/nmf_exact"
mkpath(fo)

#### global params
global_params = NMFParams(
    init          = :random,
    dist          = Uniform(0,1),
    rank          = K,
    max_iter      = 20_000,
    stopping_crit = :none,
    verbose       = true,
)

#### # 2. FEVOTTE-IDIER NMF ALGORITHMS WITH KNOWN RANK
#### # ################################################################################################################

list_res = Array{NMFResults{Float64}, 1}()
list_lab = Array{String, 1}()
list_ls  = Array{Symbol, 1}()
list_col = Array{String, 1}()

local_params = FIParams(
    scale_W_iter = true,
    scale_W_last = true,
    α_H          = 0,
    α_W          = 0,
)

####
#### 2.1 β divergences, FI algorithms
####

for (β, col) in zip([2, 1, 0], ["black", "red", "green"])
    local_params.β = β

    run_mm = true
    run_me = true
    run_h  = true

    if β == 2
        run_h = false
    elseif β == 1
        run_h  = false
        run_me = false
    end

    if run_mm
        local_params.alg = :mm

        push!(list_res, nmf_FI(V, global_params, local_params))
        push!(list_lab, raw"$\mathrm{MM} \quad \beta =" * "$(β)" * raw"$")
        push!(list_ls, :dash)
        push!(list_col, col)
    end

    if run_me
        local_params.alg = :me

        push!(list_res, nmf_FI(V, global_params, local_params))
        push!(list_lab, raw"$\mathrm{ME} \quad \beta =" * "$(β)" * raw"$")
        push!(list_ls, :solid)
        push!(list_col, col)
    end

    if run_h
        local_params.alg = :h

        push!(list_res, nmf_FI(V, global_params, local_params))
        push!(list_lab, raw"$\mathrm{H} \quad \beta =" * "$(β)" * raw"$")
        push!(list_ls, :dot)
        push!(list_col, col)
    end

end

####
#### 2.2 PLOT
####

fn = joinpath(fo, "nmf_FI_exact.pdf")
plot_cost(list_res, list_lab, list_ls, list_col, fn)

#### # 3. MU NMF ALGORITHMS WITH KNOWN RANK
#### # ################################################################################################################

list_res = Array{NMFResults{Float64}, 1}()
list_lab = Array{String, 1}()
list_ls  = Array{Symbol, 1}()
list_col = Array{String, 1}()

####
#### 2.1 β divergences
####

for (β, col) in zip([2, 1, 0], ["black", "red", "green"])
    local_params = MUParams(
        β       = β,
        scale_W = true,
        div     = :β,
        alg     = :mu
    )

    push!(list_res, nmf_MU(V, global_params, local_params))
    push!(list_lab, raw"$\mathrm{MU} \quad \beta =" * "$(β)" * raw"$")
    push!(list_ls, :solid)
    push!(list_col, col)
end

####
#### 2.2 α divergences
####

for (α, col) in zip([2, 1, 0], ["black", "red", "green"])
    local_params = MUParams(
        α            = α,
        scale_W_iter = true,
        scale_W_last = true,
        div          = :α,
        alg          = :mu
    )

    push!(list_res, nmf_MU(V, global_params, local_params))
    push!(list_lab, raw"$\mathrm{MU} \quad \alpha =" * "$(α)" * raw"$")
    push!(list_ls, :dash)
    push!(list_col, col)
end

####
#### 2.3 PLOT
####

fn = joinpath(fo, "nmf_MU_exact.pdf")
plot_cost(list_res, list_lab, list_ls, list_col, fn)


text:
import Base: iterate, eltype, length, rand, repeat, summary, show
using Base.Iterators: Cycle
using Base: @propagate_inbounds, tail
using Random: randperm
using Knet: atype
include("../utils.jl")

mutable struct VariantData{T}
    x::Any
    y::Any
    length::Any
    imax::Any
    indices::Any
    shuffle::Any
    xsize::Any
    ysize::Any
    xtype::Any
    ytype::Any
    add_random_mirroring::Bool
end

function variant_minibatch(
    x,
    y;
    shuffle = false,
    xsize = size(x),
    ysize = size(y),
    xtype = (eltype(x) <: AbstractFloat ? atype() : (typeof(x).name.wrapper){eltype(x)}),
    ytype = (eltype(y) <: AbstractFloat ? atype() : (typeof(y).name.wrapper){eltype(y)}),
    add_random_mirroring = false,
)
    nx = size(x)[end]
    if nx != size(y)[end]
        throw(DimensionMismatch())
    end

    imax = nx
    VariantData{Tuple{xtype,ytype}}(
        x,
        y,
        nx,
        imax,
        1:nx,
        shuffle,
        xsize,
        ysize,
        xtype,
        ytype,
        add_random_mirroring,
    )
end

@propagate_inbounds function iterate(d::VariantData, i = 0)
    if i >= d.imax
        return nothing
    end
    if d.shuffle && i == 0
        d.indices = randperm(d.length)
    end
    nexti = min(i + 1, d.length)
    id = d.indices[nexti]

    if d.add_random_mirroring && rand(1:2) == 1
        xbatch = mirror_raw_batch_item(d.x[id], false, d.xtype)
        if d.y == nothing
            return (xbatch, nexti)
        else
            ybatch = mirror_raw_batch_item(d.y[id], true, d.ytype)
            return ((xbatch, ybatch), nexti)
        end
    else

        if !(isa(d.x[id], d.xtype))
            xbatch = convert(d.xtype, d.x[id])
        else
            xbatch = d.x[id]
        end

        if d.y == nothing
            return (xbatch, nexti)
        else

            if !(isa(d.y[id], d.ytype))
                ybatch = convert(d.ytype, d.y[id])
            else
                ybatch = d.y[id]
            end

            return ((xbatch, ybatch), nexti)
        end
    end
end

eltype(::Type{VariantData{T}}) where {T} = T

function length(d::VariantData)
    n = d.length
    floor(Int, n)
end

function rand(d::VariantData)
    i = rand(0:(d.length-1))
    return iterate(d, i)[1]
end

# Give length info in summary:
summary(d::VariantData) = "$(length(d))-element $(typeof(d))"
show(io::IO, d::VariantData) = print(IOContext(io, :compact => true), summary(d))
show(io::IO, ::MIME"text/plain", d::VariantData) = show(io, d)


text:
module LearningStrategiesTest

import Base.Iterators: repeated

using LearningStrategies, Statistics, Test
import LearningStrategies: update!

mutable struct Model n::Int end
struct NewStrat <: LearningStrategy end
update!(m::Model, s::NewStrat, item) = (m.n += 1)

#-----------------------------------------------------------------------# Verbose
@info("Verbose")
learn!(nothing, Verbose(MaxIter(5)), 1:10)
learn!(nothing, Verbose(TimeLimit(.5)))
learn!(Model(5), Verbose(Converged(m -> m.n == 5)))

#-----------------------------------------------------------------------# Type stability
@info("Type Stability and Printing")
@testset "Type Stability and Printing" begin
    strat_list = [
        MaxIter(20),
        TimeLimit(2),
        Converged(m -> m),
        ConvergedTo(m -> m, ones(2)),
        IterFunction((m, i) -> i),
        Breaker((m, i) -> true),
        Tracer(Float64, (m, i) -> mean(m))
    ]
    println(strategy(strat_list...))
    m = ones(2)
    i = 5
    for s in strat_list
        println("  > ", s)
        @inferred setup!(s, m, nothing)
        @inferred cleanup!(s, m)
        @inferred hook(s, m, i)
        @inferred finished(s, m, nothing, i)
        @inferred update!(m, s, i)
    end
end

#-----------------------------------------------------------------------# "Real" Tests
@testset "MaxIter" begin
    for j in 1:10
        m = Model(0)
        learn!(m, strategy(MaxIter(j), NewStrat()), 1:100)
        @test m.n == j
    end
end

@testset "TimeLimit" begin
    t1 = time()
    learn!(nothing, TimeLimit(.5))
    elapsed = time() - t1
    @test .5 <= elapsed < 1
end

@testset "Converged" begin
    learn!(nothing, Converged(x -> 1))
end

@testset "ConvergedTo" begin
    model = ones(5)
    s = strategy(ConvergedTo(m -> m, ones(5)))
    @inferred learn!(model, s)
end

@testset "IterFunction" begin
    s = IterFunction(1, (m,i) -> println("Print twice:  > $i"))
    learn!(nothing, strategy(s, MaxIter(2)))
end

@testset "Tracer" begin
    t = Tracer(Int, (mod,i) -> i)
    learn!(nothing, strategy(MaxIter(100), t))
    @test collect(t) == collect(1:100)

    # collect will copy
    A = collect(t)
    A[1] = 42
    @test collect(t) == collect(1:100)
end

@testset "Breaker" begin
    m = Model(0)
    learn!(m, strategy(NewStrat(), Breaker((m,i) -> m.n==10)))
    @test m.n == 10
end

#-------------------------------------------------------------------# Linear Regression Example
struct LinRegModel
    β::Vector
end
struct LinRegSolver <: LearningStrategy end
update!(m::LinRegModel, s::LinRegSolver, item) = (m.β[:] = item[1] \ item[2])

@testset "LinRegModel" begin
    n, p = 1000, 50
    x = randn(n, p)
    y = x * range(-1, stop=1, length=p) + randn(n)

    model = LinRegModel(zeros(p))
    s = strategy(MaxIter(1), LinRegSolver())
    data = repeated((x, y))
    learn!(model, s, data)
    @test model.β == x \ y
end

#-------------------------------------------------------------------# Test Counter

mutable struct Counter <: LearningStrategy
    n::Int
    Counter() = new(0)
end

function update!(m, s::Counter, i, item)
    s.n += 1
    @test s.n == i
end

@testset "update!(m, s, i, item)" begin
  s = strategy(Verbose(MaxIter(5)), Counter())
    learn!(nothing, s)

    learn!(nothing, Counter(), repeated(42, 5))
end

end  # module LearningStrategiesTest


text:
## Demo: Sonar Equations
# Demonstrate calculation of the sonar equations, with focus on signal excess and probability of detection.
using Plots; gr()
using SpecialFunctions
using ColorSchemes

include("../SonarEquations.jl")

## Single Propagation
B = 1. # bandwidth
t = 1. # integration time
SL = 52. # source level
NL = 5. # noise level
p_fal = 1e-4 # probability of false alarm

c = 1500. # sound speed
f = 100. # frequency
λ = c/f # wavelength
k = 2π/λ # wavenumber

r₀ = 0. # source range
z₀ = 100. # source depth

R(r, z) = sqrt.((r - r₀)^2 + (z - z₀)^2) # distance from source
p(r, z) = R(r, z)\exp(im*k*R(r, z)) # complex pressure
TL(r, z) = -10log10(abs(p(r, z))) # transmission loss
DI = 20 # directivity index
# DT(r, z) = 5log10(d(r, z)*B/t)
DT(r, z) = SonarEquations.detection_threshold(d(r, z), B, t)
# SE(r, z) = SL - TL(r, z) - NL + DI - DT(r, z)
SE(r, z) = SonarEquations.signal_excess_passive(SL, TL(r, z), NL, DI, DT(r, z))

# d(r, z) = B*t*((SL - TL(r, z))/(B*NL))^2
d(r, z) = SonarEquations.detection_index_gaussian(SL, TL(r, z), NL, B, t)
# p_dtc(r, z) = erfc(erfcinv(2p_fal) - sqrt(d(r, z)/2))/2
# POD(r, z) = 100p_dtc(r, z)
POD(r, z) = 100SonarEquations.probability_of_detection_gaussian(d(r, z), p_fal)

r = range(1., 1e3, length = 100)
z = range(1., 400., length = 50)

pTL = heatmap(r, z, TL,
	seriescolor = cgrad(:jet, rev = true),
	title = "TL (dB)",
	yaxis = ("Depth (m)", :flip))

pd = heatmap(r, z, d,
	seriescolor = :jet,
	title = "d",
	yaxis = :flip)

pSE = heatmap(r, z, SE,
	seriescolor = :jet,
	title = "SE (dB)",
	xaxis = "Range (m)",
	yaxis = ("Depth (m)", :flip))

pPOD = heatmap(r, z, POD,
	seriescolor = :jet,
	title = "POD (%)",
	xaxis = "Range (m)",
	yaxis = :flip)

l = @layout [a b; c d]

pt = plot(pTL, pd, pSE, pPOD,
	layout = l)
display(pt)
savefig(pt, "img/SonarEqs_SimplePropagation.png")

## Simple Lloyd's Mirror
include("../LloydsMirror.jl")

SL = 50.
NL = 5.
B = 1.
t = 1.
DI = 20.
p_fal = 1e-2

c = 1500.0
f = 1e2
λ = c/f
r_src = 0.0
z_src = 2λ
r = range(0, 3e2, length = 201)
z = range(0, 10λ, length = 101)
zTemp = 5λ

TL(r, z) = LloydsMirror.lloydsmirror_singlereflection.(c, f, r_src, r, z_src, z)
d(r, z) = SonarEquations.detection_index_gaussian(SL, TL(r, z), NL, B, t)
DT(r, z) = SonarEquations.detection_threshold(d(r, z), B, t)
SE(r, z) = SonarEquations.signal_excess_passive(SL, TL(r, z), NL, DI, DT(r, z))
POD(r, z) = 100SonarEquations.probability_of_detection_gaussian(d(r, z), p_fal)

pTL = heatmap(r, z, TL,
	seriescolor = cgrad(:jet, rev = true),
	title = "TL (dB)",
	yaxis = ("Depth (m)", :flip))

pd = heatmap(r, z, d,
	seriescolor = :jet,
	title = "d",
	yaxis = :flip)

pSE = heatmap(r, z, SE,
	seriescolor = :jet,
	title = "SE (dB)",
	xaxis = "Range (m)",
	yaxis = ("Depth (m)", :flip))

pPOD = heatmap(r, z, POD,
	seriescolor = :jet,
	title = "POD (%)",
	xaxis = "Range (m)",
	yaxis = :flip)

l = @layout [a b; c d]

pt = plot(pTL, pd, pSE, pPOD,
	layout = l)
display(pt)
savefig(pt, "img/SonarEqs_LloydsMirror.png")

heatmap(r, z, DT,
	yaxis = :flip)

## Multipath Propagation