# These functions (variants of sa_numerical_method)
# return an instance of a numerical method to solve
# saturation adjustment, for different combinations
# of thermodynamic variable inputs.

# @print only accepts literal strings, so we must
# branch to print which method is being used.
function print_numerical_method(
    ::Type{sat_adjust_method},
) where {sat_adjust_method}
    if sat_adjust_method <: NewtonsMethod
        @print("    Method=NewtonsMethod")
    elseif sat_adjust_method <: NewtonsMethodAD
        @print("    Method=NewtonsMethodAD")
    elseif sat_adjust_method <: SecantMethod
        @print("    Method=SecantMethod")
    elseif sat_adjust_method <: RegulaFalsiMethod
        @print("    Method=RegulaFalsiMethod")
    else
        error("Unsupported numerical method")
    end
end

#####
##### Thermodynamic variable inputs: ρ, e_int, q_tot
#####
function sa_numerical_method(
    ::Type{NM},
    param_set::APS,
    ρ::FT,
    e_int::FT,
    q_tot::FT,
    phase_type::Type{<:PhaseEquil},
) where {FT, NM <: NewtonsMethod}
    _T_min::FT = T_min(param_set)
    T_init =
        max(_T_min, air_temperature(param_set, e_int, PhasePartition(q_tot))) # Assume all vapor
    return NewtonsMethod(
        T_init,
        T_ -> ∂e_int_∂T(param_set, heavisided(T_), e_int, ρ, q_tot, phase_type),
    )
end

function sa_numerical_method(
    ::Type{NM},
    param_set::APS,
    ρ::FT,
    e_int::FT,
    q_tot::FT,
    phase_type::Type{<:PhaseEquil},
) where {FT, NM <: NewtonsMethodAD}
    _T_min::FT = T_min(param_set)
    T_init =
        max(_T_min, air_temperature(param_set, e_int, PhasePartition(q_tot))) # Assume all vapor
    return NewtonsMethodAD(T_init)
end

function sa_numerical_method(
    ::Type{NM},
    param_set::APS,
    ρ::FT,
    e_int::FT,
    q_tot::FT,
    phase_type::Type{<:PhaseEquil},
) where {FT, NM <: SecantMethod}
    _T_min::FT = T_min(param_set)
    q_pt = PhasePartition(q_tot, FT(0), q_tot) # Assume all ice
    T_2 = air_temperature(param_set, e_int, q_pt)
    T_1 = max(_T_min, air_temperature(param_set, e_int, PhasePartition(q_tot))) # Assume all vapor
    T_2 = bound_upper_temperature(T_1, T_2)
    return SecantMethod(T_1, T_2)
end

function sa_numerical_method(
    ::Type{NM},
    param_set::APS,
    ρ::FT,
    e_int::FT,
    q_tot::FT,
    phase_type::Type{<:PhaseEquil},
) where {FT, NM <: RegulaFalsiMethod}
    _T_min::FT = T_min(param_set)
    q_pt = PhasePartition(q_tot, FT(0), q_tot) # Assume all ice
    T_2 = air_temperature(param_set, e_int, q_pt)
    T_1 = max(_T_min, air_temperature(param_set, e_int, PhasePartition(q_tot))) # Assume all vapor
    T_2 = bound_upper_temperature(T_1, T_2)
    return RegulaFalsiMethod(T_1, T_2)
end

#####
##### Thermodynamic variable inputs: ρ, p, q_tot
#####

function sa_numerical_method_ρpq(
    ::Type{NM},
    param_set::APS,
    ρ::FT,
    p::FT,
    q_tot::FT,
    phase_type::Type{<:PhaseEquil},
) where {FT, NM <: NewtonsMethodAD}
    q_pt = PhasePartition(q_tot)
    T_init = air_temperature_from_ideal_gas_law(param_set, p, ρ, q_pt)
    return NewtonsMethodAD(T_init)
end

function sa_numerical_method_ρpq(
    ::Type{NM},
    param_set::APS,
    ρ::FT,
    p::FT,
    q_tot::FT,
    phase_type::Type{<:PhaseEquil},
) where {FT, NM <: RegulaFalsiMethod}
    q_pt = PhasePartition(q_tot)
    T_1 = air_temperature_from_ideal_gas_law(param_set, p, ρ, q_pt) - 5
    T_2 = air_temperature_from_ideal_gas_law(param_set, p, ρ, q_pt) + 5
    return RegulaFalsiMethod(T_1, T_2)
end

#####
##### Thermodynamic variable inputs: p, e_int, q_tot
#####

function sa_numerical_method_peq(
    ::Type{NM},
    param_set::APS,
    p::FT,
    e_int::FT,
    q_tot::FT,
    phase_type::Type{<:PhaseEquil},
) where {FT, NM <: NewtonsMethodAD}
    _T_min::FT = T_min(param_set)
    T_init =
        max(_T_min, air_temperature(param_set, e_int, PhasePartition(q_tot))) # Assume all vapor
    return NewtonsMethodAD(T_init)
end

function sa_numerical_method_peq(
    ::Type{NM},
    param_set::APS,
    p::FT,
    e_int::FT,
    q_tot::FT,
    phase_type::Type{<:PhaseEquil},
) where {FT, NM <: SecantMethod}
    _T_min::FT = T_min(param_set)
    q_pt = PhasePartition(q_tot, FT(0), q_tot) # Assume all ice
    T_2 = air_temperature(param_set, e_int, q_pt)
    T_1 = max(_T_min, air_temperature(param_set, e_int, PhasePartition(q_tot))) # Assume all vapor
    T_2 = bound_upper_temperature(T_1, T_2)
    return SecantMethod(T_1, T_2)
end


using Base58: base58checkdecode

@testset "Transaction" begin
    raw_tx = hex2bytes("0100000001813f79011acb80925dfe69b3def355fe914bd1d96a3f5f71bf8303c6a989c7d1000000006b483045022100ed81ff192e75a3fd2304004dcadb746fa5e24c5031ccfcf21320b0277457c98f02207a986d955c6e0cb35d446a89d3f56100f4d7f67801c31967743a9c8e10615bed01210349fc4e631e3624a545de3f89f5d8684c7b8138bd94bdd531d2e213bf016b278afeffffff02a135ef01000000001976a914bc3b654dca7e56b04dca18f2566cdaf02e8d9ada88ac99c39800000000001976a9141c4bc762dd5423e332166702cb75f40df79fea1288ac19430600")
    @testset "Parsing" begin
        @testset "Version" begin
            stream = IOBuffer(raw_tx)
            tx = txparse(stream)
            @test tx.version == 1
        end
        @testset "Inputs" begin
            stream = IOBuffer(raw_tx)
            tx = txparse(stream)
            @test length(tx.tx_ins) == 1
            want = hex2bytes("d1c789a9c60383bf715f3f6ad9d14b91fe55f3deb369fe5d9280cb1a01793f81")
            @test tx.tx_ins[1].prev_tx == want
            @test tx.tx_ins[1].prev_index == 0
            want = hex2bytes("6b483045022100ed81ff192e75a3fd2304004dcadb746fa5e24c5031ccfcf21320b0277457c98f02207a986d955c6e0cb35d446a89d3f56100f4d7f67801c31967743a9c8e10615bed01210349fc4e631e3624a545de3f89f5d8684c7b8138bd94bdd531d2e213bf016b278a")
            @test Bitcoin.scriptserialize(tx.tx_ins[1].script_sig) == want
            @test tx.tx_ins[1].sequence == 0xfffffffe
        end
        @testset "Outputs" begin
            stream = IOBuffer(raw_tx)
            tx = txparse(stream)
            @test length(tx.tx_outs) == 2
            want = 32454049
            @test tx.tx_outs[1].amount == want
            want = hex2bytes("1976a914bc3b654dca7e56b04dca18f2566cdaf02e8d9ada88ac")
            @test Bitcoin.scriptserialize(tx.tx_outs[1].script_pubkey) == want
            want = 10011545
            @test tx.tx_outs[2].amount == want
            want = hex2bytes("1976a9141c4bc762dd5423e332166702cb75f40df79fea1288ac")
            @test Bitcoin.scriptserialize(tx.tx_outs[2].script_pubkey) == want
        end
        @testset "Locktime" begin
            stream = IOBuffer(raw_tx)
            tx = txparse(stream)
            @test tx.locktime == 410393
        end
    end
    @testset "Serialize" begin
        stream = IOBuffer(raw_tx)
        tx = txparse(stream)
        @test txserialize(tx) == raw_tx
    end
    @testset  "Input Value" begin
        tx_hash = "d1c789a9c60383bf715f3f6ad9d14b91fe55f3deb369fe5d9280cb1a01793f81"
        index = 0
        want = 42505594
        tx_in = Bitcoin.TxIn(hex2bytes(tx_hash), index)
        @test Bitcoin.txinvalue(tx_in) == want
    end
    @testset "Input PubKey" begin
        tx_hash = "d1c789a9c60383bf715f3f6ad9d14b91fe55f3deb369fe5d9280cb1a01793f81"
        index = 0
        want = hex2bytes("1976a914a802fc56c704ce87c42d7c92eb75e7896bdc41ae88ac")
        tx_in = Bitcoin.TxIn(hex2bytes(tx_hash), index)
        @test Bitcoin.scriptserialize(Bitcoin.txin_scriptpubkey(tx_in)) == want
    end
    @testset "Fee" begin
        raw_tx = hex2bytes("0100000001813f79011acb80925dfe69b3def355fe914bd1d96a3f5f71bf8303c6a989c7d1000000006b483045022100ed81ff192e75a3fd2304004dcadb746fa5e24c5031ccfcf21320b0277457c98f02207a986d955c6e0cb35d446a89d3f56100f4d7f67801c31967743a9c8e10615bed01210349fc4e631e3624a545de3f89f5d8684c7b8138bd94bdd531d2e213bf016b278afeffffff02a135ef01000000001976a914bc3b654dca7e56b04dca18f2566cdaf02e8d9ada88ac99c39800000000001976a9141c4bc762dd5423e332166702cb75f40df79fea1288ac19430600")
        stream = IOBuffer(raw_tx)
        tx = txparse(stream)
        @test Bitcoin.txfee(tx) == 40000
        raw_tx = hex2bytes("010000000456919960ac691763688d3d3bcea9ad6ecaf875df5339e148a1fc61c6ed7a069e010000006a47304402204585bcdef85e6b1c6af5c2669d4830ff86e42dd205c0e089bc2a821657e951c002201024a10366077f87d6bce1f7100ad8cfa8a064b39d4e8fe4ea13a7b71aa8180f012102f0da57e85eec2934a82a585ea337ce2f4998b50ae699dd79f5880e253dafafb7feffffffeb8f51f4038dc17e6313cf831d4f02281c2a468bde0fafd37f1bf882729e7fd3000000006a47304402207899531a52d59a6de200179928ca900254a36b8dff8bb75f5f5d71b1cdc26125022008b422690b8461cb52c3cc30330b23d574351872b7c361e9aae3649071c1a7160121035d5c93d9ac96881f19ba1f686f15f009ded7c62efe85a872e6a19b43c15a2937feffffff567bf40595119d1bb8a3037c356efd56170b64cbcc160fb028fa10704b45d775000000006a47304402204c7c7818424c7f7911da6cddc59655a70af1cb5eaf17c69dadbfc74ffa0b662f02207599e08bc8023693ad4e9527dc42c34210f7a7d1d1ddfc8492b654a11e7620a0012102158b46fbdff65d0172b7989aec8850aa0dae49abfb84c81ae6e5b251a58ace5cfeffffffd63a5e6c16e620f86f375925b21cabaf736c779f88fd04dcad51d26690f7f345010000006a47304402200633ea0d3314bea0d95b3cd8dadb2ef79ea8331ffe1e61f762c0f6daea0fabde022029f23b3e9c30f080446150b23852028751635dcee2be669c2a1686a4b5edf304012103ffd6f4a67e94aba353a00882e563ff2722eb4cff0ad6006e86ee20dfe7520d55feffffff0251430f00000000001976a914ab0c0b2e98b1ab6dbf67d4750b0a56244948a87988ac005a6202000000001976a9143c82d7df364eb6c75be8c80df2b3eda8db57397088ac46430600")
        stream = IOBuffer(raw_tx)
        tx = txparse(stream)
        @test Bitcoin.txfee(tx) == 140500
    end
    @testset "Sig Hash" begin
        raw_tx = hex2bytes("0100000001813f79011acb80925dfe69b3def355fe914bd1d96a3f5f71bf8303c6a989c7d1000000006b483045022100ed81ff192e75a3fd2304004dcadb746fa5e24c5031ccfcf21320b0277457c98f02207a986d955c6e0cb35d446a89d3f56100f4d7f67801c31967743a9c8e10615bed01210349fc4e631e3624a545de3f89f5d8684c7b8138bd94bdd531d2e213bf016b278afeffffff02a135ef01000000001976a914bc3b654dca7e56b04dca18f2566cdaf02e8d9ada88ac99c39800000000001976a9141c4bc762dd5423e332166702cb75f40df79fea1288ac19430600")
        stream = IOBuffer(raw_tx)
        tx = txparse(stream)
        want = parse(BigInt, "27e0c5994dec7824e56dec6b2fcb342eb7cdb0d0957c2fce9882f715e85d81a6", base=16)
        @test Bitcoin.txsighash(tx, 0) == want
    end
    @testset "Verify P2PKH" begin
        tx = txfetch("452c629d67e41baec3ac6f04fe744b4b9617f8f859c63b3002f8684e7a4fee03")
        @test txverify(tx)
        tx = txfetch("5418099cc755cb9dd3ebc6cf1a7888ad53a1a3beb5a025bce89eb1bf7f1650a2", true)
        @test txverify(tx)
    end
    @testset "Sign Input" begin
        private_key = PrivateKey(8675309)
        tx_ins = TxIn[]
        prev_tx = hex2bytes("0025bc3c0fa8b7eb55b9437fdbd016870d18e0df0ace7bc9864efc38414147c8")
        push!(tx_ins, TxIn(prev_tx, 0))
        tx_outs = TxOut[]
        h160 = base58checkdecode(b"mzx5YhAH9kNHtcN481u6WkjeHjYtVeKVh2")[2:end]
        push!(tx_outs, TxOut(Int(0.99 * 100000000), Bitcoin.p2pkh_script(h160)))
        h160 = base58checkdecode(b"mnrVtF8DWjMu839VW3rBfgYaAfKk8983Xf")[2:end]
        push!(tx_outs, TxOut(Int(0.1 * 100000000), Bitcoin.p2pkh_script(h160)))
        tx = Tx(1, tx_ins, tx_outs, 0, true)
        @test txsigninput(tx, 0, private_key)
    end
end


# Load modules
using JuMP
using SpineModel
using Cbc

# Custon constraints required for case study A4
function add_constraints(m::Model)

    # Cyclic storage bounds on the first solve
    @fetch stor_state = m.ext[:variables]
    cons = m.ext[:constraints][:stor_cyclic] = Dict()
    filtered_storage__commodity = filter(stor_c -> stor_c.storage in SpineModel.indices(cyclic), storage__commodity())
    stor_start = [first(stor_state_indices(storage=stor, commodity=c)) for (stor, c) in filtered_storage__commodity]
    stor_end = [last(stor_state_indices(storage=stor, commodity=c)) for (stor, c) in filtered_storage__commodity]
    for (stor, c, t_first) in stor_start
        for (stor, c, t_last) in stor_end
            (stor, c, t_first, t_last) in keys(cons) && continue
            cons[stor, c, t_first, t_last] = @constraint(
                m,
                stor_state[stor, c, t_first]
                ==
                stor_state[stor, c, t_last]
            )
        end
    end

    # THIS ISN'T NECESSARY FOR THE BUILDING MODEL!
    # Input-output conversion with online consumption
    @fetch flow, units_on = m.ext[:variables]
    cons = m.ext[:constraints][:conversion] = Dict()
    for (u, c1, c2) in SpineModel.indices(idle_heat_rate)
        for t in SpineModel.t_lowest_resolution(map(x -> x.t, flow_indices(unit=u, commodity=[c1, c2])))
            cons[u, c1, c2, t] = @constraint(
                m,
                + reduce(
                    +,
                    flow[u_, n, c1_, d, t_] * SpineModel.duration(t_)
                    for (u_, n, c1_, d, t_) in flow_indices(
                        unit=u, commodity=c1, direction=direction(:from_node), t=t_in_t(t_long=t)
                    );
                    init=0
                )
                ==
                + variable_heat_rate(unit=u, commodity1=c1, commodity2=c2, t=t)
                * reduce(
                    +,
                    flow[u_, n, c2_, d, t_] * SpineModel.duration(t_)
                    for (u_, n, c2_, d, t_) in flow_indices(
                        unit=u, commodity=c2, direction=direction(:to_node), t=t_in_t(t_long=t)
                    );
                    init=0
                )
                + idle_heat_rate(unit=u, commodity1=c1, commodity2=c2, t=t)
                #* unit_capacity(unit=u, commodity=c2, direction=direction(:to_node))
                * reduce(
                    +,
                    units_on[u_, t_] * SpineModel.duration(t_)
                    for (u_, t_) in units_on_indices(unit=u, t=t_in_t(t_long=t));
                    init=0
                )
            )
        end
    end

end

function update_constraints(m::Model)

    # Remove cyclic storage bounds from subsequent solves
    cons = pop!(m.ext[:constraints], :stor_cyclic, nothing)
    cons === nothing && return
    delete.(m, values(cons))

    # Update input-output convetrsion equation
    @fetch flow, units_on = m.ext[:variables]
    cons = m.ext[:constraints][:conversion]
    for (u, c1, c2) in SpineModel.indices(idle_heat_rate)
        for t in SpineModel.t_lowest_resolution(map(x -> x.t, flow_indices(unit=u, commodity=[c1, c2])))
            # Update idle heat rates
            set_normalized_coefficient(
                cons[u, c1, c2, t],
                units_on[u,t],
                (
                    - idle_heat_rate(unit=u, commodity1=c1, commodity2=c2, t=t)
                    #* unit_capacity(unit=u, commodity=c2, direction=direction(:to_node))
                    * SpineModel.duration(t)
                )
            )
            # Update variable heat rates
            for (u_, n, c2_, d, t_) in flow_indices(unit=u, commodity=c2, direction=direction(:to_node), t=t_in_t(t_long=t))
                set_normalized_coefficient(
                    cons[u_, c1, c2_, t_],
                    flow[u_, n, c2_, d, t_],
                    - variable_heat_rate(unit=u_, commodity1=c1, commodity2=c2_, t=t) * SpineModel.duration(t_)
                )
            end
        end
    end

end

# Catch an error if no db_url is provided
try
    ARGS[1]
catch e
    if isa(e, BoundsError)
        println("!!! Database url required as a command line argument !!!")
    end
end

# Run the model from the chosen database
m = run_spinemodel(
    ARGS[1];
    with_optimizer=optimizer_with_attributes(Cbc.Optimizer, "logLevel"=>1, "allowableGap"=>0, "ratioGap"=>0),
    add_constraints=m->add_constraints(m),
    update_constraints=m->update_constraints(m),
)


import Base: promote, neg_int, add_int, sub_int, mul_int, afoldl, @_inline_meta,
    power_by_squaring, throw_domerr_powbysq, to_power_type
import Base.Checked: checked_neg, checked_add, checked_sub, checked_mul, checked_abs,
    add_with_overflow, sub_with_overflow, mul_with_overflow, SignedInt,
    throw_overflowerr_binaryop

# convert multi-argument calls into nested two-argument calls
unchecked_add(a, b, c, xs...) = @unchecked (@_inline_meta; afoldl(+, (+)((+)(a, b), c), xs...))
unchecked_sub(a, b, c, xs...) = @unchecked (@_inline_meta; afoldl(-, (-)((-)(a, b), c), xs...))
unchecked_mul(a, b, c, xs...) = @unchecked (@_inline_meta; afoldl(*, (*)((*)(a, b), c), xs...))

checked_add(a, b, c, xs...) = @checked (@_inline_meta; afoldl(+, (+)((+)(a, b), c), xs...))
checked_sub(a, b, c, xs...) = @checked (@_inline_meta; afoldl(-, (-)((-)(a, b), c), xs...))
checked_mul(a, b, c, xs...) = @checked (@_inline_meta; afoldl(*, (*)((*)(a, b), c), xs...))

# passthrough for non-numbers
unchecked_neg(x) = Base.:-(x)
unchecked_add(x, y) = Base.:+(x, y)
unchecked_sub(x, y) = Base.:-(x, y)
unchecked_mul(x, y) = Base.:*(x, y)
unchecked_pow(x, y) = Base.:^(x, y)
unchecked_abs(x) = Base.abs(x)

checked_neg(x) = Base.:-(x)
checked_add(x, y) = Base.:+(x, y)
checked_sub(x, y) = Base.:-(x, y)
checked_mul(x, y) = Base.:*(x, y)
checked_pow(x, y) = Base.:^(x, y)
checked_abs(x) = Base.abs(x)

# promote unmatched number types to same type
unchecked_add(x::Number, y::Number) = unchecked_add(promote(x, y)...)
unchecked_sub(x::Number, y::Number) = unchecked_sub(promote(x, y)...)
unchecked_mul(x::Number, y::Number) = unchecked_mul(promote(x, y)...)
unchecked_mul(x::Number, y::Number) = unchecked_mul(promote(x, y)...)

checked_add(x::Number, y::Number) = checked_add(promote(x, y)...)
checked_sub(x::Number, y::Number) = checked_sub(promote(x, y)...)
checked_mul(x::Number, y::Number) = checked_mul(promote(x, y)...)
checked_pow(x::Number, y::Number) = checked_pow(promote(x, y)...)

# passthrough for same-type numbers that aren't integers
unchecked_add(x::T, y::T) where T <: Number = Base.:+(x, y)
unchecked_sub(x::T, y::T) where T <: Number = Base.:-(x, y)
unchecked_mul(x::T, y::T) where T <: Number = Base.:*(x, y)
unchecked_pow(x::T, y::T) where T <: Number = Base.:^(x, y)

checked_add(x::T, y::T) where T <: Number = Base.:+(x, y)
checked_sub(x::T, y::T) where T <: Number = Base.:-(x, y)
checked_mul(x::T, y::T) where T <: Number = Base.:*(x, y)
checked_pow(x::T, y::T) where T <: Number = Base.:^(x, y)

# core methods
unchecked_neg(x::T) where T <: BitInteger = neg_int(x)
unchecked_add(x::T, y::T) where T <: BitInteger = add_int(x, y)
unchecked_sub(x::T, y::T) where T <: BitInteger = sub_int(x, y)
unchecked_mul(x::T, y::T) where T <: BitInteger = mul_int(x, y)
unchecked_pow(x::T, y::S) where {T <: BitInteger, S <: BitInteger} = power_by_squaring(x, y)
unchecked_abs(x::T) where T <: SignedBitInteger = flipsign(x, x)

checked_neg(x::T) where T <: BitInteger = @checked T(0) - x
function checked_add(x::T, y::T) where T <: BitInteger
    @_inline_meta
    z, b = add_with_overflow(x, y)
    b && throw_overflowerr_binaryop(:+, x, y)
    z
end
function checked_sub(x::T, y::T) where T <: BitInteger
    @_inline_meta
    z, b = sub_with_overflow(x, y)
    b && throw_overflowerr_binaryop(:-, x, y)
    z
end
function checked_mul(x::T, y::T) where T <: BitInteger
    @_inline_meta
    z, b = mul_with_overflow(x, y)
    b && throw_overflowerr_binaryop(:*, x, y)
    z
end
function checked_pow(x::T, y::S) where {T <: BitInteger, S <: BitInteger}
    @_inline_meta
    z, b = pow_with_overflow(x, y)
    b && throw_overflowerr_binaryop(:^, x, y)
    z
end
function checked_abs(x::SignedBitInteger)
    @_inline_meta
    r = @unchecked ifelse(x < 0, -x, x)
    r < 0 && throw(OverflowError(string("checked arithmetic: cannot compute |x| for x = ", x, "::", typeof(x))))
    r
 end

 function pow_with_overflow(x_, p::Integer)
    x = to_power_type(x_)
    if p == 1
        return (copy(x), false)
    elseif p == 0
        return (one(x), false)
    elseif p == 2
        return mul_with_overflow(x, x)
    elseif p < 0
        isone(x) && return (copy(x), false)
        isone(-x) && return (iseven(p) ? one(x) : copy(x), false)
        throw_domerr_powbysq(x, p)
    end
    t = trailing_zeros(p) + 1
    p >>= t
    b = false
    while (t -= 1) > 0
        x, b1 = mul_with_overflow(x, x)
        b |= b1
    end
    y = x
    while p > 0
        t = trailing_zeros(p) + 1
        p >>= t
        while (t -= 1) >= 0
            x, b1 = mul_with_overflow(x, x)
            b |= b1
        end
        y, b1 = mul_with_overflow(y, x)
        b |= b1
    end
    return y, b
end
pow_with_overflow(x::Bool, p::Unsigned) = ((p==0) | x, false)
function pow_with_overflow(x::Bool, p::Integer)
    p < 0 && !x && throw_domerr_powbysq(x, p)
    return (p==0) | x, false
end


#=
Julia implementation of a Region Proposal Network (RPN)
=#

struct rpn
    anchors_per_loc
    anchor_stride
    depth
end


module WaveFD

using Base.Threads, CvxCompress, DSP, Distributed, DistributedArrays, FFTW, LinearAlgebra, NearestNeighbors, Random, SpecialFunctions, StaticArrays, Statistics, WaveFD_jll

import
Base.convert,
Base.copy,
Base.get,
Base.min,
Base.max,
Base.maximum,
Base.show,
Base.size

abstract type Language end
struct LangC <: Language end
struct LangJulia <: Language end
show(io::IO, l::LangC) = write(io, "C")
show(io::IO, l::LangJulia) = write(io, "Julia")

abstract type ImagingCondition end
struct ImagingConditionStandard <: ImagingCondition end
struct ImagingConditionWaveFieldSeparationFWI <: ImagingCondition end
struct ImagingConditionWaveFieldSeparationRTM <: ImagingCondition end

include("stencil.jl")
include("spacetime.jl")
include("compressedio.jl")
include("wavelet.jl")
include("absorb.jl")
include("illumination.jl")

include("prop2DAcoIsoDenQ_DEO2_FDTD.jl")
include("prop2DAcoTTIDenQ_DEO2_FDTD.jl")
include("prop2DAcoVTIDenQ_DEO2_FDTD.jl")

include("prop3DAcoIsoDenQ_DEO2_FDTD.jl")
include("prop3DAcoTTIDenQ_DEO2_FDTD.jl")
include("prop3DAcoVTIDenQ_DEO2_FDTD.jl")

export
dtmod,
Prop2DAcoIsoDenQ_DEO2_FDTD,
Prop2DAcoTTIDenQ_DEO2_FDTD,
Prop2DAcoVTIDenQ_DEO2_FDTD,
Prop3DAcoIsoDenQ_DEO2_FDTD,
Prop3DAcoTTIDenQ_DEO2_FDTD,
Prop3DAcoVTIDenQ_DEO2_FDTD,
fieldfile!,
free,
Ginsu,
linearadjoint,
linearforward,
nonlinearforward,
ntmod,
reportfreq!,
sourceillum!,
sourceillum,
traces,
Wavelet,
WaveletDerivRicker,
WaveletMinPhaseRicker,
WaveletCausalRicker,
WaveletOrmsby,
WaveletMinPhaseOrmsby,
WaveletRicker,
WaveletSine

end


### A Pluto.jl notebook ###
# v0.14.2

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 74b008f6-ed6b-11ea-291f-b3791d6d1b35
begin
	import Pkg
	Pkg.activate(mktempdir())
	Pkg.add([
		Pkg.PackageSpec(name="ImageIO", version="0.5"),
		Pkg.PackageSpec(name="ImageShow", version="0.2"),
		Pkg.PackageSpec(name="FileIO", version="1.6"),
		Pkg.PackageSpec(name="PNGFiles", version="0.3.6"),
		Pkg.PackageSpec(name="Colors", version="0.12"),
		Pkg.PackageSpec(name="ColorVectorSpace", version="0.8"),
		Pkg.PackageSpec(name="PlutoUI", version="0.7"),
		Pkg.PackageSpec(name="HypertextLiteral", version="0.5")
	])

	using Colors, ColorVectorSpace, ImageShow, FileIO
	using PlutoUI
	using HypertextLiteral
end

# ╔═╡ 71a1e08a-6abc-48d5-b177-5184dbdd76a8
filter!(LOAD_PATH) do path
	path != "@v#.#"
end;

# ╔═╡ e91d7926-ec6e-41e7-aba2-9dca333c8aa5
html"""

Section 1.1

 Images as Data and Arrays

Lecture Video

"""

# ╔═╡ d07fcdb0-7afc-4a25-b68a-49fd1e3405e7
PlutoUI.TableOfContents(aside=true)

# ╔═╡ 9b49500c-0164-4556-a17b-7595e35c5ede
md"""
#### Intializing packages

_When running this notebook for the first time, this could take up to 15 minutes. Hang in there!_
"""

# ╔═╡ ca1b507e-6017-11eb-34e6-6b85cd189002
md"""
# Images as examples of data  all around us
Welcome to the Computational Thinking using Julia for Real-World Problems, at MIT in Spring 2021!

The aim of this course is to bring together concepts from computer science and applied math with coding in the modern **Julia language**, and to see how to apply these techniques to study interesting applications (and of course to have fun).

We would be pleased if students who have been interested in computer science now become interested in computational science and those interested in scientific applications learn computer science they may not see elsewhere.
... and for all students, we wish to share the value of
the Julia  language as the best of both worlds.
"""

# ╔═╡ e9ff96d8-6bc1-11eb-0f6a-234b9fae047e
md"""

## Alan's Essay: Are all programming languages the same?

>Superficially, many programming languages are very similar.  "Showoffs" will compare functional programming vs imperative programming.  Others will compare compiled languages vs dynamic languages.  I will avoid such fancy terms in this little essay, preferring to provide this course's pedagogical viewpoint.
>
>Generally speaking beginning programmers should learn to create "arrays" write "for loops", "conditionals", "comparisons", express mathematical formulas, etc. So why Julia at a time when Python seems to be the language of teaching, and Java and C++ so prominent in the corporate world?
>
>As you might imagine, we believe Julia is special.   Oh you will still have the nitty gritty of when to use a bracket and a comma.  You might have strong opinions as to whether arrays should begin with 0 or 1 (joke: some say it's time to compromise and use ½.)  Getting past these irrelevant issues,  you will begin to experience one by one what makes Julia so very special.  For starters, a language that runs fast is more fun.  We can have you try things that would just be so slow in other languages it would be boring.  We also think you will start to notice how natural Julia is, how it feels like the mathematics, and how flexible it can be.
>
>Getting to see the true value of fancy terms like multiple dispatch, strong typing, generic programming, and composable software will take a little longer, but stick with us, and you too will see why Julia is so very special.
"""

# ╔═╡ 9111db10-6bc3-11eb-38e5-cf3f58536914
md"""
## Computer Science and Computational Science Working Together
"""

# ╔═╡ fb8a99ac-6bc1-11eb-0835-3146734a1c99
md"""
Applications of computer science in the real world use **data**, i.e. information that we can **measure** in some way. Data take many different forms, for example:

- Numbers that change over time (**time series**):
  - Stock price each second / minute / day
  - Weekly number of infections
  - Earth's global average temperature

- Video:
  - The view from a window of a self-driving car
  - A hurricane monitoring station
  - Ultrasound e.g. prenatal

- Images:
  - Diseased versus healthy tissue in a medical scan
  - Pictures of your favourite dog
"""

# ╔═╡ b795dcb4-6bc3-11eb-20ec-db2cc4b89bfb
md"""
#### Exercise:
> Think of another two examples in each category. Can you think of other categories of data?
"""

# ╔═╡ 8691e434-6bc4-11eb-07d1-8169158484e6
md"""
Computational science can be summed up by a simplified workflow:
"""

# ╔═╡ 546db74c-6d4e-11eb-2e27-f5bed9dbd9ba
md"""
## data ⟶  input  ⟶ process ⟶ model ⟶ visualize ⟶ output
"""

# ╔═╡ 6385d174-6d4e-11eb-093b-6f6fafb79f84
md"""
$(html"
")
To use any data source, we need to **input** the data of interest, for example by downloading it, reading in the resulting file, and converting it into a form that we can use in the computer. Then we need to **process** it in some way to extract information of interest. We usually want to **visualize** the results, and we may want to **output** them, for example by saving to disc or putting them on a website.

We often want to make a mathematical or computational **model** that can help us to understand and predict the behavior of the system of interest.

> In this course we aim to show how programming, computer science and applied math combine to help us with these goals.
"""

# ╔═╡ 132f6596-6bc6-11eb-29f1-1b2478c929af
md"""
# Data: Images (as an example of data)
Let's start off by looking at **images** and how we can process them.
Our goal is to process the data contained in an image in some way, which we will do by developing and coding certain **algorithms**.

Here is the the Fall 2020 version of this lecture (small variations) by 3-Blue-1-Brown (Grant Sanderson) for your reference.
"""

# ╔═╡ 635a03dd-abd7-49c8-a3d2-e68c7d83cc9b
html"""

"""

# ╔═╡ 9eb6efd2-6018-11eb-2db8-c3ce41d9e337
md"""

If we open an image on our computer or the web and zoom in enough, we will see that it consists of many tiny squares, or **pixels** ("picture elements"). Each pixel is a block of one single colour, and the pixels are arranged in a two-dimensional square grid.

You probably already know that these pixels are stored in a computer numerically
perhaps in some form of RGB (red,green,blue) format.  This is the computer's represenation of the data.

Note that an image is already an **approximation** of the real world -- it is a two-dimensional, discrete representation of a three-dimensional reality.

"""

# ╔═╡ e37e4d40-6018-11eb-3e1d-093266c98507
md"""
# Input and Visualize: loading and viewing an Image (in Julia)
"""

# ╔═╡ e1c9742a-6018-11eb-23ba-d974e57f78f9
md"""
Let's use Julia to load  actual images and play around with them. We can download images from the internet, your own file, or your own webcam.
"""

# ╔═╡ 9b004f70-6bc9-11eb-128c-914eadfc9a0e
md"""
## Downloading an image from the internet or a local file
We can use the `Images.jl` package to load an image file in three steps.
"""

# ╔═╡ 62fa19da-64c6-11eb-0038-5d40a6890cf5
md"""
Step 1: (from internet) we specify the URL (web address) to download from:
$(html"
")
(note that Pluto places results before commands because some people believe
output is more interesting than code.  This takes some getting used to.)
"""

# ╔═╡ 34ee0954-601e-11eb-1912-97dc2937fd52
url = "https://user-images.githubusercontent.com/6933510/107239146-dcc3fd00-6a28-11eb-8c7b-41aaf6618935.png"

# ╔═╡ 9180fbcc-601e-11eb-0c22-c920dc7ee9a9
md"""
Step 2: Now we use the aptly-named `download` function to download the image file to our own computer. (Philip is Prof. Edelman's corgi.)
"""

# ╔═╡ 34ffc3d8-601e-11eb-161c-6f9a07c5fd78
philip_filename = download(url) # download to a local file. The filename is returned

# ╔═╡ abaaa980-601e-11eb-0f71-8ff02269b775
md"""
Step 3:
Using the `Images.jl` package (loaded at the start of this notebook; scroll up and take a look.) we can **load** the file, which automatically converts it into usable data. We'll store the result in a variable. (Remember the code is after the output.)
"""

# ╔═╡ aafe76a6-601e-11eb-1ff5-01885c5238da
philip = load(philip_filename)

# ╔═╡ e86ed944-ee05-11ea-3e0f-d70fc73b789c
md"_Hi there Philip_"

# ╔═╡ c99d2aa8-601e-11eb-3469-497a246db17c
md"""
We see that the Pluto notebook has recognised that we created an object representing an image, and automatically displayed the resulting image of Philip, the cute Welsh Pembroke corgi and co-professor of this course.
Poor Philip will undergo quite a few transformations as we go along!
"""

# ╔═╡ 11dff4ce-6bca-11eb-1056-c1345c796ed4
md"""
- Exercise : change the url.
- Exercise: download a file that is already on your own computer.
"""

# ╔═╡ efef3a32-6bc9-11eb-17e9-dd2171be9c21
md"""
## Capturing an Image from your own camera
"""

# ╔═╡ e94dcc62-6d4e-11eb-3d53-ff9878f0091e
md"""
Even more fun is to use your own webcam. Try pressing the enable button below. Then
press the camera to capture an image. Kind of fun to keep pressing the button as you move your hand etc.
"""

# ╔═╡ cef1a95a-64c6-11eb-15e7-636a3621d727
md"""
## Inspecting your data
"""

# ╔═╡ f26d9326-64c6-11eb-1166-5d82586422ed
md"""
### Image size
"""

# ╔═╡ 6f928b30-602c-11eb-1033-71d442feff93
md"""
The first thing we might want to know is the size of the image:
"""

# ╔═╡ 75c5c85a-602c-11eb-2fb1-f7e7f2c5d04b
philip_size = size(philip)

# ╔═╡ 77f93eb8-602c-11eb-1f38-efa56cc93ca5
md"""
Julia returns a pair of two numbers. Comparing these with the picture of the image, we see that the first number is the height, i.e. the vertical number of pixels, and the second is the width.
"""

# ╔═╡ 96b7d801-c427-4e27-ab1f-e2fd18fc24d0
philip_height = philip_size[1]

# ╔═╡ f08d02af-6e38-4ace-8b11-7af4930b64ea
philip_width = philip_size[2]

# ╔═╡ f9244264-64c6-11eb-23a6-cfa76f8aff6d
md"""
### Locations in an image: Indexing

Now suppose that we want to examine a piece of the image in more detail. We need some way of specifying which piece of the image we want.

Thinking of the image as a grid of pixels, we need a way to tell the computer which pixel or group of pixels we want to refer to.
Since the image is a two-dimensional grid, we can use two integers (whole numbers) to give the coordinates of a single pixel.  Specifying coordinates like this is called **indexing**: think of the index of a book, which tells you *on which page* an idea is discussed.

In Julia we use (square) brackets, `[` and `]` for indexing:
"""

# ╔═╡ bd22d09a-64c7-11eb-146f-67733b8be241
a_pixel = philip[200, 100]

# ╔═╡ 28860d48-64c8-11eb-240f-e1232b3638df
md"""
We see that Julia knows to draw our pixel object for us a block of the relevant color.

When we index into an image like this, the first number indicates the *row* in the image, starting from the top, and the second the *column*, starting from the left. In Julia, the first row and column are numbered starting from 1, not from 0 as in some other programming languages.
"""

# ╔═╡ 4ef99715-4d8d-4f9d-bf0b-8df9907a14cf

# ╔═╡ a510fc33-406e-4fb5-be83-9e4b5578717c
md"""
We can also use variables as indices...
"""

# ╔═╡ 13844ebf-52c4-47e9-bda4-106a02fad9d7
md"""
...and these variables can be controlled by sliders!
"""

# ╔═╡ 08d61afb-c641-4aa9-b995-2552af89f3b8
@bind row_i Slider(1:size(philip)[1], show_value=true)

# ╔═╡ 6511a498-7ac9-445b-9c15-ec02d09783fe
@bind col_i Slider(1:size(philip)[2], show_value=true)

# ╔═╡ 94b77934-713e-11eb-18cf-c5dc5e7afc5b
row_i,col_i

# ╔═╡ ff762861-b186-4eb0-9582-0ce66ca10f60
philip[row_i, col_i]

# ╔═╡ c9ed950c-dcd9-4296-a431-ee0f36d5b557
md"""
### Locations in an image: Range indexing

We saw that we can use the **row number** and **column number** to index a _single pixel_ of our image. Next, we will use a **range of numbers** to index _multiple rows or columns_ at once, returning a subarray:
"""

# ╔═╡ f0796032-8105-4f6d-b5ee-3647b052f2f6
philip[550:650, 1:philip_width]

# ╔═╡ b9be8761-a9c9-49eb-ba1b-527d12097362
md"""
Here, we use `a:b` to mean "_all numbers between `a` and `b`_". For example:

"""

# ╔═╡ d515286b-4ad4-449b-8967-06b9b4c87684
collect(1:10)

# ╔═╡ eef8fbc8-8887-4628-8ba8-114575d6b91f
md"""

You can also use a `:` without start and end to mean "_every index_"
"""

# ╔═╡ 4e6a31d6-1ef8-4a69-b346-ad58cfc4d8a5
philip[550:650, :]

# ╔═╡ e11f0e47-02d9-48a6-9b1a-e313c18db129
md"""
Let's get a single row of pixels:
"""

# ╔═╡ 9e447eab-14b6-45d8-83ab-1f7f1f1c70d2
philip[550, :]

# ╔═╡ c926435c-c648-419c-9951-ac8a1d4f3b92
philip_head = philip[470:800, 140:410]

# ╔═╡ 32e7e51c-dd0d-483d-95cb-e6043f2b2975
md"""
#### Scroll in on Philip's nose!

Use the widgets below (slide left and right sides).
"""

# ╔═╡ 4b64e1f2-d0ca-4e22-a89d-1d9a16bd6788
@bind range_rows RangeSlider(1:size(philip_head)[1])

# ╔═╡ 85919db9-1444-4904-930f-ba572cff9460
@bind range_cols RangeSlider(1:size(philip_head)[2])

# ╔═╡ 2ac47b91-bbc3-49ae-9bf5-4def30ff46f4
nose = philip_head[range_rows, range_cols]

# ╔═╡ 5a0cc342-64c9-11eb-1211-f1b06d652497
md"""
# Process: Modifying an image

Now that we have access to image data, we can start to **process** that data to extract information and/or modify it in some way.

We might want to detect what type of objects are in the image, say to detect whether a patient has a certain disease. To achieve a high-level goal like this, we will need to perform mid-level operations, such as detecting edges that separate different objects based on their color. And, in turn, to carry that out we will need to do low-level operations like comparing colors of neighboring pixels and somehow deciding if they are "different".

"""

# ╔═╡ 4504577c-64c8-11eb-343b-3369b6d10d8b
md"""
## Representing colors

We can  use indexing to *modify* a pixel's color. To do so, we need a way to specify a new color.

Color turns out to be a complicated concept, having to do with the interaction of the physical properties of light with the physiological mechanisms and mental processes by which we detect it!

We will ignore this complexity by using a standard method of representing colours in the computer as an **RGB triple**, i.e. a triple of three numbers $(r, g, b)$, giving the amount of red, of green and of blue in a colour, respectively. These are numbers between 0 (none) and 1 (full). The final colour that we perceive is the result of "adding" the corresponding amount of light of each colour; the details are fascinating, but beyond the scope of this course!
"""

# ╔═╡ 40886d36-64c9-11eb-3c69-4b68673a6dde
md"""
We can create a new color in Julia as follows:
"""

# ╔═╡ 552235ec-64c9-11eb-1f7f-f76da2818cb3
RGB(1.0, 0.0, 0.0)

# ╔═╡ c2907d1a-47b1-4634-8669-a68022706861
begin
	md"""
	A pixel with $(@bind test_r Scrubbable(0:0.1:1; default=0.1)) red, $(@bind test_g Scrubbable(0:0.1:1; default=0.5)) green and $(@bind test_b Scrubbable(0:0.1:1; default=1.0)) blue looks like:
	"""
end

# ╔═╡ ff9eea3f-cab0-4030-8337-f519b94316c5
RGB(test_r, test_g, test_b)

# ╔═╡ f6cc03a0-ee07-11ea-17d8-013991514d42
md"""
#### Exercise 2.5
👉 Write a function `invert` that inverts a color, i.e. sends $(r, g, b)$ to $(1 - r, 1-g, 1-b)$.
"""

# ╔═╡ 63e8d636-ee0b-11ea-173d-bd3327347d55
function invert(color::AbstractRGB)

	return RGB(1 - color.r, 1 - color.g, 1 - color.b)
end

# ╔═╡ 2cc2f84e-ee0d-11ea-373b-e7ad3204bb00
md"Let's invert some colors:"

# ╔═╡ b8f26960-ee0a-11ea-05b9-3f4bc1099050
black = RGB(0.0, 0.0, 0.0)

# ╔═╡ 5de3a22e-ee0b-11ea-230f-35df4ca3c96d
invert(black)

# ╔═╡ 4e21e0c4-ee0b-11ea-3d65-b311ae3f98e9
red = RGB(0.8, 0.1, 0.1)

# ╔═╡ 6dbf67ce-ee0b-11ea-3b71-abc05a64dc43
invert(red)

# ╔═╡ 846b1330-ee0b-11ea-3579-7d90fafd7290
md"Can you invert the picture of Philip?"

# ╔═╡ 943103e2-ee0b-11ea-33aa-75a8a1529931
philip_inverted = invert.(philip)

# ╔═╡ 2ee543b2-64d6-11eb-3c39-c5660141787e
md"""

## Modifying a pixel

Let's start by seeing how to modify an image, e.g. in order to hide sensitive information.

We do this by assigning a new value to the color of a pixel:
"""

# ╔═╡ 53bad296-4c7b-471f-b481-0e9423a9288a
let
	temp = copy(philip_head)
	temp[100, 200] = RGB(1.0, 0.0, 0.0)
	temp
end

# ╔═╡ 81b88cbe-64c9-11eb-3b26-39011efb2089
md"""
Be careful: We are actually *modifying* the original image here, even though if we look at the image it is hard to spot, since a single pixel is so small.
"""

# ╔═╡ ab9af0f6-64c9-11eb-13d3-5dbdb75a69a7
md"""
## Groups of pixels

We probably want to examine and modify several pixels at once.
For example, we can extract a horizontal strip 1 pixel tall:
"""

# ╔═╡ e29b7954-64cb-11eb-2768-47de07766055
philip_head[50, 50:100]

# ╔═╡ 8e7c4866-64cc-11eb-0457-85be566a8966
md"""
Here, Julia is showing the strip as a collection of rectangles in a row.

"""

# ╔═╡ f2ad501a-64cb-11eb-1707-3365d05b300a
md"""
And then modify it:
"""

# ╔═╡ 4f03f651-56ed-4361-b954-e6848ac56089
let
	temp = copy(philip_head)
	temp[50, 50:100] .= RGB(1.0, 0.0, 0.0)
	temp
end

# ╔═╡ 2808339c-64cc-11eb-21d1-c76a9854aa5b
md"""
Similarly we can modify a whole rectangular block of pixels:
"""

# ╔═╡ 1bd53326-d705-4d1a-bf8f-5d7f2a4e696f
let
	temp = copy(philip_head)
	temp[50:100, 50:100] .= RGB(1.0, 0.0, 0.0)
	temp
end

# ╔═╡ a5f8bafe-edf0-11ea-0da3-3330861ae43a
md"""
#### Exercise 1.2

👉 Generate a vector of 100 zeros. Change the center 20 elements to 1.
"""

# ╔═╡ b6b65b94-edf0-11ea-3686-fbff0ff53d08
function create_bar()
	bar = zeros(100)
	bar[41:60] .= 1

	return bar
end

# ╔═╡ 693af19c-64cc-11eb-31f3-57ab2fbae597
md"""
## Reducing the size of an image
"""

# ╔═╡ 6361d102-64cc-11eb-31b7-fb631b632040
md"""
Maybe we would also like to reduce the size of this image, since it's rather large. For example, we could take every 10th row and every 10th column and make a new image from the result:
"""

# ╔═╡ ae542fe4-64cc-11eb-29fc-73b7a66314a9
reduced_image = philip[1:10:end, 1:10:end]

# ╔═╡ c29292b8-64cc-11eb-28db-b52c46e865e6
md"""
Note that the resulting image doesn't look very good, since we seem to have lost too much detail.

#### Exercise

> Think about what we might do to reduce the size of an image without losing so much detail.
"""

# ╔═╡ 7b04331a-6bcb-11eb-34fa-1f5b151e5510
md"""
# Model: Creating synthetic images

Think about your favorite Pixar movie (e.g. Monsters Inc.) Movie frames are images that are generated from complicated mathematical models.  Ray tracing (which may be covered in this class)
is a method for making images feel realistic.
"""

# ╔═╡ 5319c03c-64cc-11eb-0743-a1612476e2d3
md"""
# Output: Saving an image to a file

Finally, we want to be able to save our new creation to a file. To do so, you can **right click** on a displayed image, or you can write it to a file. Fill in a path below:
"""

# ╔═╡ 3db09d92-64cc-11eb-0333-45193c0fd1fe
save("reduced_phil.png", reduced_image)

# ╔═╡ 61606acc-6bcc-11eb-2c80-69ceec9f9702
md"""
# $(html"
")
"""

# ╔═╡ dd183eca-6018-11eb-2a83-2fcaeea62942
md"""
# Computer science: Arrays

An image is a concrete example of a fundamental concept in computer science, namely an **array**.

Just as an image is a rectangular grid, where each grid cell contains a single color,
an array is a rectangular grid for storing data. Data is stored and retrieved using indexing, just as in the image examples: each cell in the grid can store a single "piece of data" of a given type.

## Dimension of an array

An array can be one-dimensional, like the strip of pixels above, two-dimensional, three-dimensional, and so on. The dimension tells us the number of indices that we need to specify a unique location in the grid.
The array object also needs to know the length of the data in each dimension.

## Names for different types of array

One-dimensional arrays are often called **vectors** (or, in some other languages, "lists") and two-dimensional arrays are **matrices**. Higher-dimensional arrays are  **tensors**.

## Arrays as data structures

An array is an example of a **data structure**, i.e. a way of arranging data such that we can access it. A key theme in computer science is that of designing different data structures that represent data in different ways.

Conceptually, we can think of an array as a block of data that has a position or location in space. This can be a useful way to arrange data if, for example, we want to represent the fact that values in nearby locations in array are somehow near to one another.

Images are a good example of this: neighbouring pixels often represent different pieces of the same object, for example the rug or floor, or Philip himself, in the photo. We thus expect neighbouring pixels to be of a similar color. On the other hand, if they are not, this is also useful information, since that may correspond to the edge of an object.

"""

# ╔═╡ 8ddcb286-602a-11eb-3ae0-07d3c77a0f8c
md"""
# Julia: constructing arrays

## Creating vectors and matrices
Julia has strong support for arrays of any dimension.

Vectors, or one-dimensional arrays, are written using square brackets and commas:
"""

# ╔═╡ f4b0aa23-2d76-4d88-b2a4-3807e88d27ce
[1, 20, "hello"]

# ╔═╡ 1b2b2b18-64d4-11eb-2d43-e31cb8bc25d1
[RGB(1, 0, 0), RGB(0, 1, 0), RGB(0, 0, 1)]

# ╔═╡ 2b0e6450-64d4-11eb-182b-ff1bd515b56f
md"""
Matrices, or two-dimensional arrays, also use square brackets, but with spaces and new lines instead of commas:
"""

# ╔═╡ 3b2b041a-64d4-11eb-31dd-47d7321ee909
[RGB(1, 0, 0)  RGB(0, 1, 0)
 RGB(0, 0, 1)  RGB(0.5, 0.5, 0.5)]

# ╔═╡ 0f35603a-64d4-11eb-3baf-4fef06d82daa
md"""

## Array comprehensions

It's clear that if we want to create an array with more than a few elements, it will be *very* tedious to do so by hand like this.
Rather, we want to *automate* the process of creating an array by following some pattern, for example to create a whole palette of colors!

Let's start with all the possible colors interpolating between black, `RGB(0, 0, 0)`, and red, `RGB(1, 0, 0)`.  Since only one of the values is changing, we can represent this as a vector, i.e. a one-dimensional array.

A neat method to do this is an **array comprehension**. Again we use square brackets  to create an array, but now we use a **variable** that varies over a given **range** values:
"""

# ╔═╡ e69b02c6-64d6-11eb-02f1-21c4fb5d1043
[RGB(x, 0, 0) for x in 0:0.1:1]

# ╔═╡ fce76132-64d6-11eb-259d-b130038bbae6
md"""
Here, `0:0.1:1` is a **range**; the first and last numbers are the start and end values, and the middle number is the size of the step.
"""

# ╔═╡ 17a69736-64d7-11eb-2c6c-eb5ebf51b285
md"""
In a similar way we can create two-dimensional matrices, by separating the two variables for each dimension with a comma (`,`):
"""

# ╔═╡ 291b04de-64d7-11eb-1ee0-d998dccb998c
[RGB(i, j, 0) for i in 0:0.1:1, j in 0:0.1:1]

# ╔═╡ 647fddf2-60ee-11eb-124d-5356c7014c3b
md"""
## Joining matrices

We often want to join vectors and matrices together. We can do so using an extension of the array creation syntax:
"""

# ╔═╡ 7d9ad134-60ee-11eb-1b2a-a7d63f3a7a2d
[philip_head  philip_head]

# ╔═╡ 8433b862-60ee-11eb-0cfc-add2b72997dc
[philip_head                   reverse(philip_head, dims=2)
 reverse(philip_head, dims=1)  rot180(philip_head)]

# ╔═╡ 5e52d12e-64d7-11eb-0905-c9038a404e24
md"""
# Pluto: Interactivity using sliders
"""

# ╔═╡ 6aba7e62-64d7-11eb-2c49-7944e9e2b94b
md"""
Suppose we want to see the effect of changing the number of colors in our vector or matrix. We could, of course, do so by manually fiddling with the range.

It would be nice if we could do so using a **user interface**, for example with a **slider**. Fortunately, the Pluto notebook allows us to do so!
"""

# ╔═╡ afc66dac-64d7-11eb-1ad0-7f62c20ffefb
md"""
We can define a slider using
"""

# ╔═╡ b37c9868-64d7-11eb-3033-a7b5d3065f7f
@bind number_reds Slider(1:100, show_value=true)

# ╔═╡ b1dfe122-64dc-11eb-1104-1b8852b2c4c5
md"""
[The `Slider` type is defined in the `PlutoUI.jl` package.]
"""

# ╔═╡ cfc55140-64d7-11eb-0ff6-e59c70d01d67
md"""
This creates a new variable called `number_reds`, whose value is the value shown by the slider. When we move the slider, the value of the variable gets updated. Since Pluto is a **reactive** notebook, other cells which use the value of this variable will *automatically be updated too*!
"""

# ╔═╡ fca72490-64d7-11eb-1464-f5e0582c4d18
md"""
Let's use this to make a slider for our one-dimensional collection of reds:
"""

# ╔═╡ 88933746-6028-11eb-32de-13eb6ff43e29
[RGB(red_value / number_reds, 0, 0) for red_value in 0:number_reds]

# ╔═╡ 1c539b02-64d8-11eb-3505-c9288357d139
md"""
When you move the slider, you should see the number of red color patches change!
"""

# ╔═╡ 10f6e6da-64d8-11eb-366f-11f16e73043b
md"""
What is going on here is that we are creating a vector in which `red_value` takes each value in turn from the range from `0` up to the current value of `number_reds`. If we change `number_reds`, then we create a new vector with that new number of red patches.
"""

# ╔═╡ 82a8314c-64d8-11eb-1acb-e33625381178
md"""
#### Exercise

> Make three sliders with variables `r`, `g` and `b`. Then make a single color patch with the RGB color given by those values.
"""

# ╔═╡ e01fb06a-3be2-4f02-af8a-88d1141bf223
@bind r Slider(0:0.01:1, show_value=true)

# ╔═╡ 19956bda-2828-4931-bc15-e6577e66bbe1
@bind g Slider(0:0.01:1, show_value=true)

# ╔═╡ 3b420ccf-4ecb-4d9d-b052-6c6398c6dc1c
@bind b Slider(0:0.01:1, show_value=true)

# ╔═╡ 40331f9c-8a83-4819-bd6d-cfef1fe91fc2
RGB(r, g, b)

# ╔═╡ 576d5e3a-64d8-11eb-10c9-876be31f7830
md"""
We can do the same to create different size matrices, by creating two sliders, one for reds and one for greens. Try it out!
"""

# ╔═╡ ace86c8a-60ee-11eb-34ef-93c54abc7b1a
md"""
# Summary
"""

# ╔═╡ b08e57e4-60ee-11eb-0e1a-2f49c496668b
md"""
Let's summarize the main ideas from this notebook:
- Images are **arrays** of colors
- We can inspect and modify arrays using **indexing**
- We can create arrays directly or using **array comprehensions**
"""

# ╔═╡ 9025a5b4-6066-11eb-20e8-099e9b8f859e
md"""
----
"""

# ╔═╡ 5da8cbe8-eded-11ea-2e43-c5b7cc71e133
begin
	colored_line(x::Vector{<:Real}) = Gray.(Float64.((hcat(x)')))
	colored_line(x::Any) = nothing
end

# ╔═╡ d862fb16-edf1-11ea-36ec-615d521e6bc0
colored_line(create_bar())

# ╔═╡ e074560a-601b-11eb-340e-47acd64f03b2
hint(text) = Markdown.MD(Markdown.Admonition("hint", "Hint", [text]))

# ╔═╡ e0776548-601b-11eb-2563-57ba2cf1d5d1
almost(text) = Markdown.MD(Markdown.Admonition("warning", "Almost there!", [text]))

# ╔═╡ e083bef6-601b-11eb-2134-e3063d5c4253
still_missing(text=md"Replace `missing` with your answer.") = Markdown.MD(Markdown.Admonition("warning", "Here we go!", [text]))

# ╔═╡ e08ecb84-601b-11eb-0e25-152ed3a262f7
keep_working(text=md"The answer is not quite right.") = Markdown.MD(Markdown.Admonition("danger", "Keep working on it!", [text]))

# ╔═╡ e09036a4-601b-11eb-1a8b-ef70105ab91c
yays = [md"Great!", md"Yay ❤", md"Great! 🎉", md"Well done!", md"Keep it up!", md"Good job!", md"Awesome!", md"You got the right answer!", md"Let's move on to the next section."]

# ╔═╡ e09af1a2-601b-11eb-14c8-57a46546f6ce
correct(text=rand(yays)) = Markdown.MD(Markdown.Admonition("correct", "Got it!", [text]))

# ╔═╡ e0a4fc10-601b-11eb-211d-03570aca2726
not_defined(variable_name) = Markdown.MD(Markdown.Admonition("danger", "Oopsie!", [md"Make sure that you define a variable called **$(Markdown.Code(string(variable_name)))**"]))

# ╔═╡ e3394c8a-edf0-11ea-1bb8-619f7abb6881
if !@isdefined(create_bar)
	not_defined(:create_bar)
else
	let
		result = create_bar()
		if ismissing(result)
			still_missing()
		elseif isnothing(result)
			keep_working(md"Did you forget to write `return`?")
		elseif !(result isa Vector) || length(result) != 100
			keep_working(md"The result should be a `Vector` with 100 elements.")
		elseif result[[1,50,100]] != [0,1,0]
			keep_working()
		else
			correct()
		end
	end
end

# ╔═╡ e0a6031c-601b-11eb-27a5-65140dd92897
bigbreak = html"

";

# ╔═╡ 45815734-ee0a-11ea-2982-595e1fc0e7b1
bigbreak

# ╔═╡ e0b15582-601b-11eb-26d6-bbf708933bc8
function camera_input(;max_size=150, default_url="https://i.imgur.com/SUmi94P.png")
"""

			<button id="stop" title="Stop video">✖</button>
			<button id="pop-out" title="Pop out/pop in">⏏</button>

		<video playsinline autoplay></video>

		<button id="shutter" title="Click to take a picture">📷</button>

		Enable webcam

""" |> HTML
end

# ╔═╡ d6742ea0-1106-4f3c-a5b8-a31a48d33f19
@bind webcam_data1 camera_input()

# ╔═╡ 2a94a2cf-b697-4b0b-afd0-af2e35af2bb1
@bind webcam_data camera_input()

# ╔═╡ e891fce0-601b-11eb-383b-bde5b128822e

function process_raw_camera_data(raw_camera_data)
	# the raw image data is a long byte array, we need to transform it into something
	# more "Julian" - something with more _structure_.

	# The encoding of the raw byte stream is:
	# every 4 bytes is a single pixel
	# every pixel has 4 values: Red, Green, Blue, Alpha
	# (we ignore alpha for this notebook)

	# So to get the red values for each pixel, we take every 4th value, starting at
	# the 1st:
	reds_flat = UInt8.(raw_camera_data["data"][1:4:end])
	greens_flat = UInt8.(raw_camera_data["data"][2:4:end])
	blues_flat = UInt8.(raw_camera_data["data"][3:4:end])

	# but these are still 1-dimensional arrays, nicknamed 'flat' arrays
	# We will 'reshape' this into 2D arrays:

	width = raw_camera_data["width"]
	height = raw_camera_data["height"]

	# shuffle and flip to get it in the right shape
	reds = reshape(reds_flat, (width, height))' / 255.0
	greens = reshape(greens_flat, (width, height))' / 255.0
	blues = reshape(blues_flat, (width, height))' / 255.0

	# we have our 2D array for each color
	# Let's create a single 2D array, where each value contains the R, G and B value of
	# that pixel

	RGB.(reds, greens, blues)
end

# ╔═╡ 1d7375b7-7ea6-4d67-ab73-1c69d6b8b87f
myface1 = process_raw_camera_data(webcam_data1);

# ╔═╡ 6224c74b-8915-4983-abf0-30e6ba04a46d
[
	myface1              myface1[   :    , end:-1:1]
	myface1[end:-1:1, :] myface1[end:-1:1, end:-1:1]
]

# ╔═╡ 3e0ece65-b8a7-4be7-ae44-6d7210c2e15b
myface = process_raw_camera_data(webcam_data);

# ╔═╡ 4ee18bee-13e6-4478-b2ca-ab66100e57ec
[
	myface              myface[   :    , end:-1:1]
	myface[end:-1:1, :] myface[end:-1:1, end:-1:1]
]

# ╔═╡ 3ef77236-1867-4d02-8af2-ff4777fcd6d9
exercise_css = html"""

"""

# ╔═╡ 61b29e7d-5aba-4bc8-870b-c1c43919c236
exercise(x, number="") =
@htl("""
	<ct-exercise class="exercise">

Exercise $(number)

$(x)

	</ct-exercise>
	""")

# ╔═╡ a9fef6c9-e911-4d8c-b141-a4832b40a260
quick_question(x, number, options, correct) = let
	name = join(rand('a':'z',16))
@htl("""
	<ct-exercise class="quick-question">

Quick Question $(number)

$(x)
	<ct-answers>
	$(map(enumerate(options)) do (i, o)
		@htl("<ct-answer><input type=radio name=$(name) id=$(i) >$(o)</ct-answer>")
	end)
	</ct-answers>

	</ct-exercise>
	""")
end

# ╔═╡ edf900be-601b-11eb-0456-3f7cfc5e876b
md"_Lecture 1, Spring 2021, version 0_"

# ╔═╡ Cell order:
# ╟─e91d7926-ec6e-41e7-aba2-9dca333c8aa5
# ╟─d07fcdb0-7afc-4a25-b68a-49fd1e3405e7
# ╟─9b49500c-0164-4556-a17b-7595e35c5ede
# ╠═74b008f6-ed6b-11ea-291f-b3791d6d1b35
# ╟─71a1e08a-6abc-48d5-b177-5184dbdd76a8
# ╟─ca1b507e-6017-11eb-34e6-6b85cd189002
# ╟─e9ff96d8-6bc1-11eb-0f6a-234b9fae047e
# ╟─9111db10-6bc3-11eb-38e5-cf3f58536914
# ╟─fb8a99ac-6bc1-11eb-0835-3146734a1c99
# ╟─b795dcb4-6bc3-11eb-20ec-db2cc4b89bfb
# ╟─8691e434-6bc4-11eb-07d1-8169158484e6
# ╟─546db74c-6d4e-11eb-2e27-f5bed9dbd9ba
# ╟─6385d174-6d4e-11eb-093b-6f6fafb79f84
# ╟─132f6596-6bc6-11eb-29f1-1b2478c929af
# ╟─635a03dd-abd7-49c8-a3d2-e68c7d83cc9b
# ╟─9eb6efd2-6018-11eb-2db8-c3ce41d9e337
# ╟─e37e4d40-6018-11eb-3e1d-093266c98507
# ╟─e1c9742a-6018-11eb-23ba-d974e57f78f9
# ╟─9b004f70-6bc9-11eb-128c-914eadfc9a0e
# ╟─62fa19da-64c6-11eb-0038-5d40a6890cf5
# ╠═34ee0954-601e-11eb-1912-97dc2937fd52
# ╟─9180fbcc-601e-11eb-0c22-c920dc7ee9a9
# ╠═34ffc3d8-601e-11eb-161c-6f9a07c5fd78
# ╟─abaaa980-601e-11eb-0f71-8ff02269b775
# ╠═aafe76a6-601e-11eb-1ff5-01885c5238da
# ╟─e86ed944-ee05-11ea-3e0f-d70fc73b789c
# ╟─c99d2aa8-601e-11eb-3469-497a246db17c
# ╟─11dff4ce-6bca-11eb-1056-c1345c796ed4
# ╟─efef3a32-6bc9-11eb-17e9-dd2171be9c21
# ╟─e94dcc62-6d4e-11eb-3d53-ff9878f0091e
# ╟─d6742ea0-1106-4f3c-a5b8-a31a48d33f19
# ╠═1d7375b7-7ea6-4d67-ab73-1c69d6b8b87f
# ╠═6224c74b-8915-4983-abf0-30e6ba04a46d
# ╟─cef1a95a-64c6-11eb-15e7-636a3621d727
# ╟─f26d9326-64c6-11eb-1166-5d82586422ed
# ╟─6f928b30-602c-11eb-1033-71d442feff93
# ╠═75c5c85a-602c-11eb-2fb1-f7e7f2c5d04b
# ╟─77f93eb8-602c-11eb-1f38-efa56cc93ca5
# ╠═96b7d801-c427-4e27-ab1f-e2fd18fc24d0
# ╠═f08d02af-6e38-4ace-8b11-7af4930b64ea
# ╟─f9244264-64c6-11eb-23a6-cfa76f8aff6d
# ╠═bd22d09a-64c7-11eb-146f-67733b8be241
# ╟─28860d48-64c8-11eb-240f-e1232b3638df
# ╟─4ef99715-4d8d-4f9d-bf0b-8df9907a14cf
# ╟─a510fc33-406e-4fb5-be83-9e4b5578717c
# ╠═94b77934-713e-11eb-18cf-c5dc5e7afc5b
# ╠═ff762861-b186-4eb0-9582-0ce66ca10f60
# ╟─13844ebf-52c4-47e9-bda4-106a02fad9d7
# ╠═08d61afb-c641-4aa9-b995-2552af89f3b8
# ╠═6511a498-7ac9-445b-9c15-ec02d09783fe
# ╟─c9ed950c-dcd9-4296-a431-ee0f36d5b557
# ╠═f0796032-8105-4f6d-b5ee-3647b052f2f6
# ╟─b9be8761-a9c9-49eb-ba1b-527d12097362
# ╠═d515286b-4ad4-449b-8967-06b9b4c87684
# ╟─eef8fbc8-8887-4628-8ba8-114575d6b91f
# ╠═4e6a31d6-1ef8-4a69-b346-ad58cfc4d8a5
# ╟─e11f0e47-02d9-48a6-9b1a-e313c18db129
# ╠═9e447eab-14b6-45d8-83ab-1f7f1f1c70d2
# ╠═c926435c-c648-419c-9951-ac8a1d4f3b92
# ╟─32e7e51c-dd0d-483d-95cb-e6043f2b2975
# ╠═4b64e1f2-d0ca-4e22-a89d-1d9a16bd6788
# ╠═85919db9-1444-4904-930f-ba572cff9460
# ╠═2ac47b91-bbc3-49ae-9bf5-4def30ff46f4
# ╟─5a0cc342-64c9-11eb-1211-f1b06d652497
# ╟─4504577c-64c8-11eb-343b-3369b6d10d8b
# ╟─40886d36-64c9-11eb-3c69-4b68673a6dde
# ╠═552235ec-64c9-11eb-1f7f-f76da2818cb3
# ╟─c2907d1a-47b1-4634-8669-a68022706861
# ╠═ff9eea3f-cab0-4030-8337-f519b94316c5
# ╟─f6cc03a0-ee07-11ea-17d8-013991514d42
# ╠═63e8d636-ee0b-11ea-173d-bd3327347d55
# ╟─2cc2f84e-ee0d-11ea-373b-e7ad3204bb00
# ╟─b8f26960-ee0a-11ea-05b9-3f4bc1099050
# ╠═5de3a22e-ee0b-11ea-230f-35df4ca3c96d
# ╠═4e21e0c4-ee0b-11ea-3d65-b311ae3f98e9
# ╠═6dbf67ce-ee0b-11ea-3b71-abc05a64dc43
# ╟─846b1330-ee0b-11ea-3579-7d90fafd7290
# ╠═943103e2-ee0b-11ea-33aa-75a8a1529931
# ╟─2ee543b2-64d6-11eb-3c39-c5660141787e
# ╠═53bad296-4c7b-471f-b481-0e9423a9288a
# ╟─81b88cbe-64c9-11eb-3b26-39011efb2089
# ╟─ab9af0f6-64c9-11eb-13d3-5dbdb75a69a7
# ╠═e29b7954-64cb-11eb-2768-47de07766055
# ╟─8e7c4866-64cc-11eb-0457-85be566a8966
# ╟─f2ad501a-64cb-11eb-1707-3365d05b300a
# ╠═4f03f651-56ed-4361-b954-e6848ac56089
# ╟─2808339c-64cc-11eb-21d1-c76a9854aa5b
# ╠═1bd53326-d705-4d1a-bf8f-5d7f2a4e696f
# ╟─a5f8bafe-edf0-11ea-0da3-3330861ae43a
# ╠═b6b65b94-edf0-11ea-3686-fbff0ff53d08
# ╟─d862fb16-edf1-11ea-36ec-615d521e6bc0
# ╟─e3394c8a-edf0-11ea-1bb8-619f7abb6881
# ╟─693af19c-64cc-11eb-31f3-57ab2fbae597
# ╟─6361d102-64cc-11eb-31b7-fb631b632040
# ╠═ae542fe4-64cc-11eb-29fc-73b7a66314a9
# ╟─c29292b8-64cc-11eb-28db-b52c46e865e6
# ╟─7b04331a-6bcb-11eb-34fa-1f5b151e5510
# ╟─5319c03c-64cc-11eb-0743-a1612476e2d3
# ╠═3db09d92-64cc-11eb-0333-45193c0fd1fe
# ╟─61606acc-6bcc-11eb-2c80-69ceec9f9702
# ╟─dd183eca-6018-11eb-2a83-2fcaeea62942
# ╟─8ddcb286-602a-11eb-3ae0-07d3c77a0f8c
# ╠═f4b0aa23-2d76-4d88-b2a4-3807e88d27ce
# ╠═1b2b2b18-64d4-11eb-2d43-e31cb8bc25d1
# ╟─2b0e6450-64d4-11eb-182b-ff1bd515b56f
# ╠═3b2b041a-64d4-11eb-31dd-47d7321ee909
# ╟─0f35603a-64d4-11eb-3baf-4fef06d82daa
# ╠═e69b02c6-64d6-11eb-02f1-21c4fb5d1043
# ╟─fce76132-64d6-11eb-259d-b130038bbae6
# ╟─17a69736-64d7-11eb-2c6c-eb5ebf51b285
# ╠═291b04de-64d7-11eb-1ee0-d998dccb998c
# ╟─647fddf2-60ee-11eb-124d-5356c7014c3b
# ╠═7d9ad134-60ee-11eb-1b2a-a7d63f3a7a2d
# ╠═8433b862-60ee-11eb-0cfc-add2b72997dc
# ╟─5e52d12e-64d7-11eb-0905-c9038a404e24
# ╟─6aba7e62-64d7-11eb-2c49-7944e9e2b94b
# ╟─afc66dac-64d7-11eb-1ad0-7f62c20ffefb
# ╠═b37c9868-64d7-11eb-3033-a7b5d3065f7f
# ╟─b1dfe122-64dc-11eb-1104-1b8852b2c4c5
# ╟─cfc55140-64d7-11eb-0ff6-e59c70d01d67
# ╟─fca72490-64d7-11eb-1464-f5e0582c4d18
# ╠═88933746-6028-11eb-32de-13eb6ff43e29
# ╟─1c539b02-64d8-11eb-3505-c9288357d139
# ╟─10f6e6da-64d8-11eb-366f-11f16e73043b
# ╟─82a8314c-64d8-11eb-1acb-e33625381178
# ╠═e01fb06a-3be2-4f02-af8a-88d1141bf223
# ╠═19956bda-2828-4931-bc15-e6577e66bbe1
# ╠═3b420ccf-4ecb-4d9d-b052-6c6398c6dc1c
# ╠═40331f9c-8a83-4819-bd6d-cfef1fe91fc2
# ╟─576d5e3a-64d8-11eb-10c9-876be31f7830
# ╠═2a94a2cf-b697-4b0b-afd0-af2e35af2bb1
# ╠═3e0ece65-b8a7-4be7-ae44-6d7210c2e15b
# ╠═4ee18bee-13e6-4478-b2ca-ab66100e57ec
# ╟─ace86c8a-60ee-11eb-34ef-93c54abc7b1a
# ╟─b08e57e4-60ee-11eb-0e1a-2f49c496668b
# ╟─9025a5b4-6066-11eb-20e8-099e9b8f859e
# ╟─45815734-ee0a-11ea-2982-595e1fc0e7b1
# ╟─5da8cbe8-eded-11ea-2e43-c5b7cc71e133
# ╟─e074560a-601b-11eb-340e-47acd64f03b2
# ╟─e0776548-601b-11eb-2563-57ba2cf1d5d1
# ╟─e083bef6-601b-11eb-2134-e3063d5c4253
# ╟─e08ecb84-601b-11eb-0e25-152ed3a262f7
# ╟─e09036a4-601b-11eb-1a8b-ef70105ab91c
# ╟─e09af1a2-601b-11eb-14c8-57a46546f6ce
# ╟─e0a4fc10-601b-11eb-211d-03570aca2726
# ╠═e0a6031c-601b-11eb-27a5-65140dd92897
# ╟─e0b15582-601b-11eb-26d6-bbf708933bc8
# ╟─e891fce0-601b-11eb-383b-bde5b128822e
# ╟─3ef77236-1867-4d02-8af2-ff4777fcd6d9
# ╟─61b29e7d-5aba-4bc8-870b-c1c43919c236
# ╟─a9fef6c9-e911-4d8c-b141-a4832b40a260
# ╟─edf900be-601b-11eb-0456-3f7cfc5e876b


#
using Definitions
using Uprating
using Dates

mutable struct Person
    hid::BigInt # == sernum
    pid::BigInt # == unique id (year * 100000)+
    pno::Integer # person number in household
    default_benefit_unit::Integer
    age::Integer

    sex::Sex
    ethnic_group::Ethnic_Group
    marital_status::Marital_Status
    highest_qualification::Qualification_Type

    sic::SIC_2007
    occupational_classification::Standard_Occupational_Classification
    public_or_private :: Employment_Sector
    principal_employment_type :: Employment_Type

    socio_economic_grouping::Socio_Economic_Group
    age_completed_full_time_education::Integer
    years_in_full_time_work::Integer
    employment_status::ILO_Employment
    actual_hours_worked::Real
    usual_hours_worked::Real

    income::Incomes_Dict
    assets::Asset_Dict
    # contracted_out_of_serps::Bool

    registered_blind::Bool
    registered_partially_sighted::Bool
    registered_deaf::Bool

    disabilities::Disability_Dict
    health_status::Health_Status
    relationships::Relationship_Dict
    relationship_to_hoh :: Relationship;
    is_informal_carer::Bool
    receives_informal_care_from_non_householder::Bool
    hours_of_care_received::Real
    hours_of_care_given::Real
    #
    # Childcare fields; assigned to children
    #
    hours_of_childcare :: Real
    cost_of_childcare :: Real
    childcare_type :: Child_Care_Type
    employer_provides_child_care :: Bool

    company_car_fuel_type :: Fuel_Type
    company_car_value :: Real
    company_car_contribution :: Real
    fuel_supplied :: Real
end

People_Dict = Dict{BigInt,Person}
Pid_Array = Vector{BigInt}

mutable struct Household
    sequence::Integer # position in current generated dataset
    hid::BigInt
    interview_year::Integer
    interview_month::Integer
    quarter::Integer
    tenure::Tenure_Type
    region::Standard_Region
    ct_band::CT_Band
    council_tax::Real
    water_and_sewerage ::Real
    mortgage_payment::Real
    mortgage_interest::Real
    years_outstanding_on_mortgage::Integer
    mortgage_outstanding::Real
    year_house_bought::Integer
    gross_rent::Real # rentg Gross rent including Housing Benefit  or rent Net amount of last rent payment
    rent_includes_water_and_sewerage::Bool
    other_housing_charges::Real # rent Net amount of last rent payment
    gross_housing_costs::Real
    total_income::Real
    total_wealth::Real
    house_value::Real
    weight::Real
    people::People_Dict

end

function uprate!( pid :: BigInt, year::Integer, quarter::Integer, person :: Person )
    for (t,inc) in person.income
            person.income[t] = uprate( inc, year, quarter, UPRATE_MAPPINGS[t])
    end
    for (a,ass) in person.assets
            person.assets[a] = uprate( ass, year, quarter, UPRATE_MAPPINGS[a])
    end
    person.cost_of_childcare = uprate( person.cost_of_childcare, year, quarter, upr_earnings )
end

function uprate!( hh :: Household )

    hh.water_and_sewerage  = uprate( hh.water_and_sewerage , hh.interview_year, hh.quarter, upr_housing_rents )
    hh.mortgage_payment = uprate( hh.mortgage_payment, hh.interview_year, hh.quarter, upr_housing_oo )
    hh.mortgage_interest = uprate( hh.mortgage_interest, hh.interview_year, hh.quarter, upr_housing_oo )
    hh.mortgage_outstanding = uprate( hh.mortgage_outstanding, hh.interview_year, hh.quarter, upr_housing_oo )
    hh.gross_rent = uprate( hh.gross_rent, hh.interview_year, hh.quarter, upr_housing_rents )
    hh.other_housing_charges = uprate( hh.other_housing_charges, hh.interview_year, hh.quarter, upr_nominal_gdp )
    hh.gross_housing_costs = uprate( hh.gross_housing_costs, hh.interview_year, hh.quarter, upr_nominal_gdp )
    hh.total_income = uprate( hh.total_income, hh.interview_year, hh.quarter, upr_nominal_gdp )
    hh.total_wealth = uprate( hh.total_wealth, hh.interview_year, hh.quarter, upr_nominal_gdp )
    hh.house_value = uprate( hh.house_value, hh.interview_year, hh.quarter, upr_housing_oo )
    for (pid,person) in hh.people
        uprate!( pid, hh.interview_year, hh.quarter, person )
    end

end

function oldest_person( people :: People_Dict ) :: NamedTuple
    oldest = ( age=-999, pid=BigInt(0))
    for person in people
        if person.age > oldest.age
            oldest.age = person.age
            oldest.pid = person.pid
        end
    end
    oldest
end

function equivalence_scale( people :: People_Dict ) :: Dict{Equivalence_Scale_Type,Real}
    np = length(people)
    eqp = Vector{EQ_Person}()
    oldest_pid = oldest_person( people )
    for (pid,person) in people
        eqtype = eq_other_adult
        if pid == oldest_pid.pid
            eqtype = eq_head
        else
            if (person.age < 16) || (( person.age < 18 ) & ( person.employment_status in [Student,Other_Inactive]))
                eqtype = eq_dependent_child # needn't actually be dependent, of course
            elseif person.relationships[ oldest_pid.pid ] in [Spouse,Cohabitee,Civil_Partner]
                eqtype = eq_spouse_of_head
            else
                eqtype = eq_other_adult
            end
        end
        push!( eqp, EQ_Person( person.age, eqtype ))
    end
    get_equivalence_scales( eqp )
end

PeopleArray = Vector{Person}

struct BenefitUnit
    people :: People_Dict
    head :: BigInt
    spouse :: BigInt
    children :: Pid_Array
end

function num_people( bu :: BenefitUnit )::Integer
    1 + ((bu.spouse > 0) ? 1 : 0 ) + size( bu.children )[1]
end

function get_head( bu :: BenefitUnit )::Person
    bu.people[bu.head]
end

function get_spouse( bu :: BenefitUnit )::Union{Nothing,Person}
    if bu.spouse <= 0
        return nothing
    end
    bu.people[bu.spouse]
end

BenefitUnits = Vector{BenefitUnit}
BUAllocation = Vector{PeopleArray}

#
# This creates a array of references to each person in the houshold, broken into
# benefit units using the default FRS/EFS benefit unit number.
#
function default_bu_allocation( hh :: Household ) :: BUAllocation
    bua = BUAllocation()
    nbus = 0
    # how many bus
    for (pid,person) in hh.people
        nbus = max( nbus, person.default_benefit_unit )
    end
    ## create benefit units
    for buno in 1:nbus
        push!( bua, PeopleArray())
    end
    sz = size( bua )[1]
    ## allocate people to them
    for (pid,person) in hh.people
        pp = bua[person.default_benefit_unit]
        push!( pp, person )
    end
    ## sort people in each by pid
    for buno in 1:nbus
        sort!( bua[buno], lt=(left,right)->isless(right.age,left.age))
    end
    bua
end

#
# This creates a array of references to each person in the houshold, broken into
# benefit units using the default FRS/EFS benefit unit number.
#
function allocate_to_bus( bua :: BUAllocation ) :: BenefitUnits
    nbus = size(bua)[1]
    bus = BenefitUnits(undef, nbus)
    for i in 1:nbus
        people = People_Dict()
        head_pid :: BigInt = -1
        spouse_pid :: BigInt = -1
        children = Pid_Array()
        npeople = size( bua[i])[1]
        for p in 1:npeople
            person = bua[i][p]
            people[person.pid] = person
            if p == 1
                head_pid = person.pid
            else
                reltohead = person.relationships[head_pid]
                if reltohead in [Spouse,Cohabitee,Civil_Partner]
                    spouse_pid = person.pid
                    # FIXME we need to remove these checks if
                    # we're using a non-default allocation to bus
                    # @assert person.age >= 16
                else
                    # @assert person.age <= 19
                    push!( children, person.pid )
                end
            end
        end
        new_bu = BenefitUnit( people, head_pid, spouse_pid, children )
        bus[i] = new_bu
    end
    bus
end

function get_benefit_units(
    hh :: Household,
    allocator :: Function=default_bu_allocation ) :: BenefitUnits
    allocate_to_bus( allocator(hh))
end

# simple diagnostic prints for testing allocation

function printpids( pers :: Person )
    println( "pid: $(pers.pid) age $(pers.age) sex $(pers.sex) default_benefit_unit $(pers.default_benefit_unit) ")
end

function printpids( bu::BenefitUnit)
      head = get_head( bu )
      print( "HEAD: ")
      printpids( head )
      spouse = get_spouse( bu )
      if spouse != nothing
            print( "SPOUSE: ")
            printpids( spouse )

      end
      for chno in bu.children
            print( "CHILD_$(chno):")
            child = bu.people[chno]
            printpids( child )
      end
end

function printpids( people :: People_Dict )
    ks = sort(collect(keys(people)))
    for k in ks
        print( "PERSON_$k")
        printpids( people[k] )
    end
end

function printpids( buas::BUAllocation )
    nbus = size( buas )[1]
    for i in 1:nbus
        npeople = size( buas[i] )[1]
        println( "bu[$i]; num_people = $npeople:")
        for p in 1:npeople
            print( "PERS_$(p): ")
            printpids( buas[i][p])
        end
    end
end


module AccretionGeometry

import GeodesicTracer: tracegeodesics, DiscreteCallback, terminate!
import GeodesicRendering: rendergeodesics, prerendergeodesics
import Base: in
import RecursiveArrayTools: ArrayPartition
import GeometryBasics
import LinearAlgebra: ×, ⋅

using GeodesicBase
using StaticArrays

include("geometry.jl")
include("meshes.jl")
include("intersections.jl")
include("discs.jl")

function tracegeodesics(
    m::AbstractMetricParams{T},
    init_positions,
    init_velocities,
    time_domain::Tuple{T,T},
    accretion_geometry;
    callback = nothing,
    kwargs...,
) where {T}
    cbs = add_collision_callback(callback, accretion_geometry)
    tracegeodesics(
        m,
        init_positions,
        init_velocities,
        time_domain;
        callback = cbs,
        kwargs...,
    )
end

function rendergeodesics(
    m::AbstractMetricParams{T},
    init_pos,
    max_time::T,
    accretion_geometry;
    callback = nothing,
    kwargs...,
) where {T}
    cbs = add_collision_callback(callback, accretion_geometry)
    rendergeodesics(m, init_pos, max_time; callback = cbs, kwargs...)
end

function prerendergeodesics(
    m::AbstractMetricParams{T},
    init_pos,
    max_time::T,
    accretion_geometry;
    callback = nothing,
    kwargs...,
) where {T}
    cbs = add_collision_callback(callback, accretion_geometry)
    prerendergeodesics(m, init_pos, max_time; callback = cbs, kwargs...)
end

add_collision_callback(::Nothing, accretion_geometry) =
    build_collision_callback(accretion_geometry)
add_collision_callback(callback::Base.AbstractVecOrTuple, accretion_geometry) =
    (callback..., build_collision_callback(accretion_geometry))
add_collision_callback(callback, accretion_geometry) =
    (callback, build_collision_callback(accretion_geometry))

function build_collision_callback(geometry::AbstractAccretionGeometry{T}) where {T}
    DiscreteCallback(collision_callback(geometry), i -> terminate!(i, :Intersected))
end

collision_callback(m::AbstractAccretionGeometry{T}) where {T} =
    (u, λ, integrator) -> intersects_geometry(m, line_element(u, integrator))

export tracegeodesics, rendergeodesics, prerendergeodesics

end # module


"""
   Delegates loading functions & traits for a generic theoretical spectrograph

Author: Eric Ford and collaborators
Created: August 2020
"""

"""
Module providing types and traits and customized functions for a generic theoretical spectrograph.
"""
module TheoreticalInstrument
using ..RvSpectMLBase
using ..RvSpectMLBase.InstrumentsCommon
#import ..RvSpectML: AbstractInstrument, AbstractInstrument1D, AbstractInstrument2D
using DataFrames # , FITSIO

default_theoretical_instrument_resolution = 137500
default_theoretical_instrument_λ_min = 4700.0
default_theoretical_instrument_λ_max = 6000.0

""" Trait for 1D spectra from a theoretical instrument """
struct TheoreticalInstrument1D <: AbstractInstrument1D
    λ_min::Float64
    λ_max::Float64
    resolution::Float64

    function TheoreticalInstrument1D(λl::Real, λh::Real, res::Real)
        @assert 3000 <= λl < λh
        @assert λl < λh <= 10000
        @assert 10000 <= res <= 1000000
        new(λl,λh,res)
    end
end

function TheoreticalInstrument1D(;λ_min::Real = default_theoretical_instrument_λ_min, λ_max::Real = default_theoretical_instrument_λ_max, resolution::Real = default_theoretical_instrument_resolution )
    TheoreticalInstrument1D(λ_min,λ_max,resolution)
end

λ_min(x::TheoreticalInstrument1D) = x.λ_min
λ_max(x::TheoreticalInstrument1D) = x.λ_max

""" Trait for a 2D spectra from a theoretical instrument """
struct TheoreticalInstrument2D <: AbstractInstrument2D
    λ_min::AbstractArray{Float64,1}
    λ_max::AbstractArray{Float64,1}
    pixels_per_order::Int64

    function TheoreticalInstrument2D(λ_in::AbstractArray{Float64,2}; pixels_per_order::Integer = size(λ_in,1))
        @assert size(λ_in,1) >= 6
        @assert size(λ_in,2) >= 1
        @assert 1 <= pixels_per_order <= 16386

        lambda_min = vec(minimum(λ_in,dims=1))
        lambda_max = vec(maximum(λ_in,dims=1))
        new(lambda_min, lambda_max, pixels_per_order)
    end
end

# Trait for any spectra from TheoreticalInstrument (could improve by using SimpleTraits)
const AnyTheoreticalInstrument = Union{TheoreticalInstrument1D,TheoreticalInstrument2D}
export TheoreticalInstrument, TheoreticalInstrument1D, TheoreticalInstrument2D, AnyTheoreticalInstrument

λ_min(x::TheoreticalInstrument2D) = minimum(x.λ_min)
λ_max(x::TheoreticalInstrument2D) = maximum(x.λ_max)

include("traits.jl")
export min_pixel, max_pixel, min_order, max_order, min_pixel_in_order, max_pixel_in_order
export orders_to_use_default, min_col_default, max_col_default
export metadata_symbols_default, metadata_strings_default
export default_ccf_mask_v_width

# Need to think about what to do for IO, if anything.
#include("io.jl")
#export make_manifest, read_metadata, read_data

include("util.jl")

include("gen_data.jl")
export generate_spectrum, generate_spectra_timeseries

end


struct DGModel{BL,G,NFND,NFD,GNF,AS,DS,D,MD}
  balancelaw::BL
  grid::G
  numfluxnondiff::NFND
  numfluxdiff::NFD
  gradnumflux::GNF
  auxstate::AS
  diffstate::DS
  direction::D
  modeldata::MD
end
function DGModel(balancelaw, grid, numfluxnondiff, numfluxdiff, gradnumflux;
                 auxstate=create_auxstate(balancelaw, grid),
                 diffstate=create_diffstate(balancelaw, grid),
                 direction=EveryDirection(), modeldata=nothing)
  DGModel(balancelaw, grid, numfluxnondiff, numfluxdiff, gradnumflux, auxstate,
          diffstate, direction, modeldata)
end

function (dg::DGModel)(dQdt, Q, ::Nothing, t; increment=false)
  bl = dg.balancelaw
  device = typeof(Q.data) <: Array ? CPU() : CUDA()

  grid = dg.grid
  topology = grid.topology

  dim = dimensionality(grid)
  N = polynomialorder(grid)
  Nq = N + 1
  Nqk = dim == 2 ? 1 : Nq
  Nfp = Nq * Nqk
  nrealelem = length(topology.realelems)

  Qvisc = dg.diffstate
  auxstate = dg.auxstate

  FT = eltype(Q)
  nviscstate = num_diffusive(bl, FT)

  Np = dofs_per_element(grid)

  communicate = !(isstacked(topology) &&
                  typeof(dg.direction) <: VerticalDirection)

  aux_comm = update_aux!(dg, bl, Q, t)
  @assert typeof(aux_comm) == Bool

  ########################
  # Gradient Computation #
  ########################
  if communicate
    MPIStateArrays.start_ghost_exchange!(Q)
    if aux_comm
      MPIStateArrays.start_ghost_exchange!(auxstate)
    end
  end

  if nviscstate > 0

    @launch(device, threads=(Nq, Nq, Nqk), blocks=nrealelem,
            volumeviscterms!(bl, Val(dim), Val(N), dg.direction, Q.data,
                             Qvisc.data, auxstate.data, grid.vgeo, t, grid.D,
                             topology.realelems))

    if communicate
      MPIStateArrays.finish_ghost_recv!(Q)
      if aux_comm
        MPIStateArrays.finish_ghost_recv!(auxstate)
      end
    end

    @launch(device, threads=Nfp, blocks=nrealelem,
            faceviscterms!(bl, Val(dim), Val(N), dg.direction,
                           dg.gradnumflux, Q.data, Qvisc.data, auxstate.data,
                           grid.vgeo, grid.sgeo, t, grid.vmapM, grid.vmapP, grid.elemtobndy,
                           topology.realelems))

    if communicate
      MPIStateArrays.start_ghost_exchange!(Qvisc)
    end

    aux_comm = update_aux_diffusive!(dg, bl, Q, t)
    @assert typeof(aux_comm) == Bool

    if aux_comm
      MPIStateArrays.start_ghost_exchange!(auxstate)
    end
  end

  ###################
  # RHS Computation #
  ###################
  @launch(device, threads=(Nq, Nq, Nqk), blocks=nrealelem,
          volumerhs!(bl, Val(dim), Val(N), dg.direction, dQdt.data,
                     Q.data, Qvisc.data, auxstate.data, grid.vgeo, t,
                     grid.ω, grid.D, topology.realelems, increment))

  if communicate
    if nviscstate > 0
      MPIStateArrays.finish_ghost_recv!(Qvisc)
      if aux_comm
        MPIStateArrays.finish_ghost_recv!(auxstate)
      end
    else
      MPIStateArrays.finish_ghost_recv!(Q)
      if aux_comm
        MPIStateArrays.finish_ghost_recv!(auxstate)
      end
    end
  end

  @launch(device, threads=Nfp, blocks=nrealelem,
          facerhs!(bl, Val(dim), Val(N), dg.direction,
                   dg.numfluxnondiff,
                   dg.numfluxdiff,
                   dQdt.data, Q.data, Qvisc.data,
                   auxstate.data, grid.vgeo, grid.sgeo, t, grid.vmapM, grid.vmapP, grid.elemtobndy,
                   topology.realelems))

  # Just to be safe, we wait on the sends we started.
  if communicate
    MPIStateArrays.finish_ghost_send!(Qvisc)
    MPIStateArrays.finish_ghost_send!(Q)
  end
end

function init_ode_state(dg::DGModel, args...;
                        forcecpu=false,
                        commtag=888)
  device = arraytype(dg.grid) <: Array ? CPU() : CUDA()

  bl = dg.balancelaw
  grid = dg.grid

  state = create_state(bl, grid, commtag)

  topology = grid.topology
  Np = dofs_per_element(grid)

  auxstate = dg.auxstate
  dim = dimensionality(grid)
  N = polynomialorder(grid)
  nrealelem = length(topology.realelems)

  if !forcecpu
    @launch(device, threads=(Np,), blocks=nrealelem,
            initstate!(bl, Val(dim), Val(N), state.data, auxstate.data, grid.vgeo,
                     topology.realelems, args...))
  else
    h_state = similar(state, Array)
    h_auxstate = similar(auxstate, Array)
    h_auxstate .= auxstate
    @launch(CPU(), threads=(Np,), blocks=nrealelem,
      initstate!(bl, Val(dim), Val(N), h_state.data, h_auxstate.data, Array(grid.vgeo),
          topology.realelems, args...))
    state .= h_state
  end

  MPIStateArrays.start_ghost_exchange!(state)
  MPIStateArrays.finish_ghost_exchange!(state)

  return state
end

# fallback
function update_aux!(dg::DGModel, bl::BalanceLaw, Q::MPIStateArray, t::Real)
  return false
end

function update_aux_diffusive!(dg::DGModel, bl::BalanceLaw, Q::MPIStateArray, t::Real)
  return false
end

function indefinite_stack_integral!(dg::DGModel, m::BalanceLaw,
                                    Q::MPIStateArray,
                                    auxstate::MPIStateArray,
                                    t::Real)

  device = typeof(Q.data) <: Array ? CPU() : CUDA()

  grid = dg.grid
  topology = grid.topology

  dim = dimensionality(grid)
  N = polynomialorder(grid)
  Nq = N + 1
  Nqk = dim == 2 ? 1 : Nq

  FT = eltype(Q)

  # do integrals
  nelem = length(topology.elems)
  nvertelem = topology.stacksize
  nhorzelem = div(nelem, nvertelem)

  @launch(device, threads=(Nq, Nqk, 1), blocks=nhorzelem,
          knl_indefinite_stack_integral!(m, Val(dim), Val(N),
                                         Val(nvertelem),
                                         Q.data, auxstate.data,
                                         grid.vgeo, grid.Imat, 1:nhorzelem))
end

function reverse_indefinite_stack_integral!(dg::DGModel,
                                            m::BalanceLaw,
                                            Q::MPIStateArray,
                                            auxstate::MPIStateArray, t::Real)

  device = typeof(auxstate.data) <: Array ? CPU() : CUDA()

  grid = dg.grid
  topology = grid.topology

  dim = dimensionality(grid)
  N = polynomialorder(grid)
  Nq = N + 1
  Nqk = dim == 2 ? 1 : Nq

  FT = eltype(auxstate)

  # do integrals
  nelem = length(topology.elems)
  nvertelem = topology.stacksize
  nhorzelem = div(nelem, nvertelem)

  @launch(device, threads=(Nq, Nqk, 1), blocks=nhorzelem,
          knl_reverse_indefinite_stack_integral!(m, Val(dim), Val(N),
                                                 Val(nvertelem),
                                                 Q.data, auxstate.data,
                                                 1:nhorzelem))
end

function nodal_update_aux!(f!, dg::DGModel, m::BalanceLaw, Q::MPIStateArray,
                           t::Real; diffusive=false)
  device = typeof(Q.data) <: Array ? CPU() : CUDA()

  grid = dg.grid
  topology = grid.topology

  dim = dimensionality(grid)
  N = polynomialorder(grid)
  Nq = N + 1
  nrealelem = length(topology.realelems)

  Np = dofs_per_element(grid)

  ### update aux variables
  if diffusive
    @launch(device, threads=(Np,), blocks=nrealelem,
            knl_nodal_update_aux!(m, Val(dim), Val(N), f!,
                            Q.data, dg.auxstate.data, dg.diffstate.data, t,
                            topology.realelems))
  else
    @launch(device, threads=(Np,), blocks=nrealelem,
            knl_nodal_update_aux!(m, Val(dim), Val(N), f!,
                            Q.data, dg.auxstate.data, t,
                            topology.realelems))
  end
end

"""
    courant(local_courant::Function, dg::DGModel, m::BalanceLaw,
            Q::MPIStateArray, direction=EveryDirection())
Returns the maximum of the evaluation of the function `local_courant`
pointwise throughout the domain.  The function `local_courant` is given an
approximation of the local node distance `Δx`.  The `direction` controls which
reference directions are considered when computing the minimum node distance
`Δx`.
An example `local_courant` function is
    function local_courant(m::AtmosModel, state::Vars, aux::Vars,
                           diffusive::Vars, Δx)
      return Δt * cmax / Δx
    end
where `Δt` is the time step size and `cmax` is the maximum flow speed in the
model.
"""
function courant(local_courant::Function, dg::DGModel, m::BalanceLaw,
                 Q::MPIStateArray, Δt, direction=EveryDirection())
    grid = dg.grid
    topology = grid.topology
    nrealelem = length(topology.realelems)

    if nrealelem > 0
        N = polynomialorder(grid)
        dim = dimensionality(grid)
        Nq = N + 1
        Nqk = dim == 2 ? 1 : Nq
        device = grid.vgeo isa Array ? CPU() : CUDA()
        pointwise_courant = similar(grid.vgeo, Nq^dim, nrealelem)
        @launch(device, threads=(Nq, Nq, Nqk), blocks=nrealelem,
        Grids.knl_min_neighbor_distance!(Val(N), Val(dim), direction,
                                         pointwise_courant, grid.vgeo, topology.realelems))
        @launch(device, threads=(Nq*Nq*Nqk,), blocks=nrealelem,
                knl_local_courant!(m, Val(dim), Val(N), pointwise_courant,
                local_courant, Q.data, dg.auxstate.data,
                dg.diffstate.data, topology.realelems, direction, Δt))
        rank_courant_max = maximum(pointwise_courant)
    else
        rank_courant_max = typemin(eltype(Q))
    end

    MPI.Allreduce(rank_courant_max, max, topology.mpicomm)
end

function copy_stack_field_down!(dg::DGModel, m::BalanceLaw,
                                auxstate::MPIStateArray, fldin, fldout)

  device = typeof(auxstate.data) <: Array ? CPU() : CUDA()

  grid = dg.grid
  topology = grid.topology

  dim = dimensionality(grid)
  N = polynomialorder(grid)
  Nq = N + 1
  Nqk = dim == 2 ? 1 : Nq

  # do integrals
  nelem = length(topology.elems)
  nvertelem = topology.stacksize
  nhorzelem = div(nelem, nvertelem)

  @launch(device, threads=(Nq, Nqk, 1), blocks=nhorzelem,
          knl_copy_stack_field_down!(Val(dim), Val(N), Val(nvertelem),
                                     auxstate.data, 1:nhorzelem, Val(fldin),
                                     Val(fldout)))
end

function MPIStateArrays.MPIStateArray(dg::DGModel, commtag=888)
  bl = dg.balancelaw
  grid = dg.grid

  state = create_state(bl, grid, commtag)

  return state
end


#=
This file is auto-generated. Do not edit.
=#
"""
    mutable struct OneDOneQMachine <: Machine
        R::Float64
        Xd::Float64
        Xq::Float64
        Xd_p::Float64
        Xq_p::Float64
        Td0_p::Float64
        Tq0_p::Float64
        ext::Dict{String, Any}
        states::Vector{Symbol}
        n_states::Int
        internal::InfrastructureSystemsInternal
    end

Parameters of 4-states synchronous machine: Simplified Marconato model
 The derivative of stator fluxes (ψd and ψq) is neglected and ωψd = ψd and
 ωψq = ψq is assumed (i.e. ω=1.0). This is standard when
 transmission network dynamics is neglected.

# Arguments
- `R::Float64`: Resistance after EMF in machine per unit, validation range: `(0, nothing)`
- `Xd::Float64`: Reactance after EMF in d-axis per unit, validation range: `(0, nothing)`
- `Xq::Float64`: Reactance after EMF in q-axis per unit, validation range: `(0, nothing)`
- `Xd_p::Float64`: Transient reactance after EMF in d-axis per unit, validation range: `(0, nothing)`
- `Xq_p::Float64`: Transient reactance after EMF in q-axis per unit, validation range: `(0, nothing)`
- `Td0_p::Float64`: Time constant of transient d-axis voltage, validation range: `(0, nothing)`
- `Tq0_p::Float64`: Time constant of transient q-axis voltage, validation range: `(0, nothing)`
- `ext::Dict{String, Any}`
- `states::Vector{Symbol}`: The states are:
	eq_p: q-axis transient voltage,
	ed_p: d-axis transient voltage
- `n_states::Int`: OneDOneQMachine has 2 states
- `internal::InfrastructureSystemsInternal`: power system internal reference, do not modify
"""
mutable struct OneDOneQMachine <: Machine
    "Resistance after EMF in machine per unit"
    R::Float64
    "Reactance after EMF in d-axis per unit"
    Xd::Float64
    "Reactance after EMF in q-axis per unit"
    Xq::Float64
    "Transient reactance after EMF in d-axis per unit"
    Xd_p::Float64
    "Transient reactance after EMF in q-axis per unit"
    Xq_p::Float64
    "Time constant of transient d-axis voltage"
    Td0_p::Float64
    "Time constant of transient q-axis voltage"
    Tq0_p::Float64
    ext::Dict{String, Any}
    "The states are:
	eq_p: q-axis transient voltage,
	ed_p: d-axis transient voltage"
    states::Vector{Symbol}
    "OneDOneQMachine has 2 states"
    n_states::Int
    "power system internal reference, do not modify"
    internal::InfrastructureSystemsInternal
end

function OneDOneQMachine(R, Xd, Xq, Xd_p, Xq_p, Td0_p, Tq0_p, ext=Dict{String, Any}(), )
    OneDOneQMachine(R, Xd, Xq, Xd_p, Xq_p, Td0_p, Tq0_p, ext, [:eq_p, :ed_p], 2, InfrastructureSystemsInternal(), )
end

function OneDOneQMachine(; R, Xd, Xq, Xd_p, Xq_p, Td0_p, Tq0_p, ext=Dict{String, Any}(), states=[:eq_p, :ed_p], n_states=2, internal=InfrastructureSystemsInternal(), )
    OneDOneQMachine(R, Xd, Xq, Xd_p, Xq_p, Td0_p, Tq0_p, ext, states, n_states, internal, )
end

# Constructor for demo purposes; non-functional.
function OneDOneQMachine(::Nothing)
    OneDOneQMachine(;
        R=0,
        Xd=0,
        Xq=0,
        Xd_p=0,
        Xq_p=0,
        Td0_p=0,
        Tq0_p=0,
        ext=Dict{String, Any}(),
    )
end

"""Get [`OneDOneQMachine`](@ref) `R`."""
get_R(value::OneDOneQMachine) = value.R
"""Get [`OneDOneQMachine`](@ref) `Xd`."""
get_Xd(value::OneDOneQMachine) = value.Xd
"""Get [`OneDOneQMachine`](@ref) `Xq`."""
get_Xq(value::OneDOneQMachine) = value.Xq
"""Get [`OneDOneQMachine`](@ref) `Xd_p`."""
get_Xd_p(value::OneDOneQMachine) = value.Xd_p
"""Get [`OneDOneQMachine`](@ref) `Xq_p`."""
get_Xq_p(value::OneDOneQMachine) = value.Xq_p
"""Get [`OneDOneQMachine`](@ref) `Td0_p`."""
get_Td0_p(value::OneDOneQMachine) = value.Td0_p
"""Get [`OneDOneQMachine`](@ref) `Tq0_p`."""
get_Tq0_p(value::OneDOneQMachine) = value.Tq0_p
"""Get [`OneDOneQMachine`](@ref) `ext`."""
get_ext(value::OneDOneQMachine) = value.ext
"""Get [`OneDOneQMachine`](@ref) `states`."""
get_states(value::OneDOneQMachine) = value.states
"""Get [`OneDOneQMachine`](@ref) `n_states`."""
get_n_states(value::OneDOneQMachine) = value.n_states
"""Get [`OneDOneQMachine`](@ref) `internal`."""
get_internal(value::OneDOneQMachine) = value.internal

"""Set [`OneDOneQMachine`](@ref) `R`."""
set_R!(value::OneDOneQMachine, val) = value.R = val
"""Set [`OneDOneQMachine`](@ref) `Xd`."""
set_Xd!(value::OneDOneQMachine, val) = value.Xd = val
"""Set [`OneDOneQMachine`](@ref) `Xq`."""
set_Xq!(value::OneDOneQMachine, val) = value.Xq = val
"""Set [`OneDOneQMachine`](@ref) `Xd_p`."""
set_Xd_p!(value::OneDOneQMachine, val) = value.Xd_p = val
"""Set [`OneDOneQMachine`](@ref) `Xq_p`."""
set_Xq_p!(value::OneDOneQMachine, val) = value.Xq_p = val
"""Set [`OneDOneQMachine`](@ref) `Td0_p`."""
set_Td0_p!(value::OneDOneQMachine, val) = value.Td0_p = val
"""Set [`OneDOneQMachine`](@ref) `Tq0_p`."""
set_Tq0_p!(value::OneDOneQMachine, val) = value.Tq0_p = val
"""Set [`OneDOneQMachine`](@ref) `ext`."""
set_ext!(value::OneDOneQMachine, val) = value.ext = val


import QuantumControlBase
using QuantumControlBase: getcontrols, getcontrolderivs, discretize_on_midpoints, evalcontrols
using QuantumPropagators: init_storage
using ConcreteStructs

# Krotov workspace (for internal use)
@concrete terse struct KrotovWrk

    # a copy of the objectives
    objectives

    # the adjoint objectives, containing the adjoint generators for the
    # backward propagation
    adjoint_objectives

    # The kwargs from the control problem
    kwargs

    # Tuple of the original controls (probably functions)
    controls

    # storage for controls discretized on intervals of tlist
    pulses0 :: Vector{Vector{Float64}}

    # second pulse storage: pulses0 and pulses1 alternate in storing the guess
    # pulses and optimized pulses in each iteration
    pulses1 ::  Vector{Vector{Float64}}

    # values of ∫gₐ(t)dt for each pulse
    g_a_int :: Vector{Float64}

    # update shapes S(t) for each pulse, discretized on intervals
    update_shapes ::  Vector{Vector{Float64}}

    lambda_vals :: Vector{Float64}

    is_parametrized :: Vector{Bool}

    parametrization

    # map of controls to options
    pulse_options

    # Result object

    result

    #################################
    # scratch objects, per objective:

    # backward-propagated states
    # note: storage for fw-propagated states is in result.states
    bw_states

    # dynamical generator at a particular point in time
    G

    control_derivs :: Vector{Vector{Union{Function, Nothing}}}

    vals_dict

    fw_storage # forward storage array (per objective)

    fw_storage2 # forward storage array (per objective)

    bw_storage # backward storage array (per objective)

    fw_prop_wrk

    bw_prop_wrk

    use_threads :: Bool

end

function KrotovWrk(problem::QuantumControlBase.ControlProblem)
    use_threads = get(problem.kwargs, :use_threads, false)
    objectives = [obj for obj in problem.objectives]
    adjoint_objectives = [adjoint(obj) for obj in problem.objectives]
    controls = getcontrols(objectives)
    control_derivs = [
        getcontrolderivs(obj.generator, controls) for obj in objectives
    ]
    tlist = problem.tlist
    kwargs = Dict(problem.kwargs)
    # TODO: check that kwargs has all the required parameters
    pulse_options = problem.pulse_options
    update_shapes = [
        discretize_on_midpoints(
                pulse_options[control][:update_shape],
                tlist
        )
        for control in controls
    ]
    lambda_vals = [
        convert(Float64, pulse_options[control][:lambda_a])
        for control in controls
    ]
    is_parametrized = [
        haskey(pulse_options[control], :parametrization)
        for control in controls
    ]
    parametrization = [
        get(pulse_options[control], :parametrization, NoParametrization())
        for control in controls
    ]
    if haskey(kwargs, :continue_from)
        @info "Continuing previous optimization"
        result = kwargs[:continue_from]
        if !(result isa KrotovResult)
            # account for continuing from a different optimization method
            result = convert(KrotovResult, result)
        end
        result.iter_stop = get(problem.kwargs, :iter_stop, 5000)
        result.converged = false
        result.start_local_time = now()
        result.message = "in progress"
        pulses0 = [
            discretize_on_midpoints(control, tlist)
            for control in result.optimized_controls
        ]
    else
        result = KrotovResult(problem)
        pulses0 = [
            discretize_on_midpoints(control, tlist) for control in controls
        ]
    end
    pulses1 = [copy(pulse) for pulse in pulses0]
    g_a_int = zeros(length(pulses0))
    bw_states = [similar(obj.initial_state) for obj in objectives]
    zero_vals = IdDict(control => zero(pulses0[i][1]) for (i, control) in enumerate(controls))
    G = [evalcontrols(obj.generator, zero_vals) for obj in objectives]
    vals_dict = [copy(zero_vals) for _ in objectives]
    # TODO: forward_storage only if g_b != 0
    fw_storage = [init_storage(obj.initial_state, tlist) for obj in objectives]
    # TODO: second forward storage only if second order
    fw_storage2 = [init_storage(obj.initial_state, tlist) for obj in objectives]
    bw_storage = [init_storage(obj.initial_state, tlist) for obj in objectives]
    fw_prop_method = [
        Val(
            QuantumControlBase.get_objective_prop_method(
                obj, :fw_prop_method, :prop_method; problem.kwargs...
            )
        )
        for obj in objectives
    ]
    bw_prop_method = [
        Val(
            QuantumControlBase.get_objective_prop_method(
                obj, :bw_prop_method, :prop_method; problem.kwargs...
            )
        )
        for obj in objectives
    ]

    fw_prop_wrk = [
        QuantumControlBase.initobjpropwrk(obj, tlist, fw_prop_method[k];
                                            initial_state=obj.initial_state,
                                            kwargs...)
        for (k, obj) in enumerate(objectives)
    ]
    bw_prop_wrk = [
        QuantumControlBase.initobjpropwrk(obj, tlist, bw_prop_method[k];
                                            initial_state=obj.initial_state,
                                            kwargs...)
        for (k, obj) in enumerate(objectives)
    ]
    return KrotovWrk(
        objectives, adjoint_objectives, kwargs, controls, pulses0, pulses1,
        g_a_int, update_shapes, lambda_vals, is_parametrized,
        parametrization, pulse_options, result, bw_states, G,
        control_derivs, vals_dict, fw_storage, fw_storage2,
        bw_storage, fw_prop_wrk, bw_prop_wrk, use_threads
    )
end