text:
#### EDMF model kernels

using CLIMAParameters.Planet: e_int_v0, grav, day, R_d, R_v, molmass_ratio
using Printf
using ClimateMachine.Atmos: nodal_update_auxiliary_state!

using ClimateMachine.BalanceLaws: number_states

using ClimateMachine.MPIStateArrays: MPIStateArray
using ClimateMachine.DGMethods: LocalGeometry, DGModel

import ClimateMachine.BalanceLaws:
    vars_state,
    update_auxiliary_state!,
    init_state_prognostic!,
    flux_first_order!,
    flux_second_order!,
    compute_gradient_argument!,
    compute_gradient_flux!

import ClimateMachine.TurbulenceConvection:
    init_aux_turbconv!,
    turbconv_nodal_update_auxiliary_state!,
    turbconv_boundary_state!,
    turbconv_normal_boundary_flux_second_order!

using ClimateMachine.Thermodynamics: air_pressure, air_density

include(joinpath("helper_funcs", "nondimensional_exchange_functions.jl"))
include(joinpath("helper_funcs", "lamb_smooth_minimum.jl"))
include(joinpath("helper_funcs", "utility_funcs.jl"))
include(joinpath("helper_funcs", "subdomain_statistics.jl"))
include(joinpath("helper_funcs", "diagnose_environment.jl"))
include(joinpath("helper_funcs", "subdomain_thermo_states.jl"))
include(joinpath("helper_funcs", "save_subdomain_temperature.jl"))
include(joinpath("closures", "entr_detr.jl"))
include(joinpath("closures", "pressure.jl"))
include(joinpath("closures", "mixing_length.jl"))
include(joinpath("closures", "turbulence_functions.jl"))
include(joinpath("closures", "surface_functions.jl"))

function vars_state(m::NTuple{N, Updraft}, st::Auxiliary, FT) where {N}
    return Tuple{ntuple(i -> vars_state(m[i], st, FT), N)...}
end

function vars_state(::Updraft, ::Auxiliary, FT)
    @vars(
        buoyancy::FT,
        a::FT,
        ε_dyn::FT,
        δ_dyn::FT,
        ε_trb::FT,
        T::FT,
        θ_liq::FT,
        q_tot::FT,
        w::FT,
    )
end

function vars_state(::Environment, ::Auxiliary, FT)
    @vars(T::FT, cld_frac::FT, buoyancy::FT)
end

function vars_state(m::EDMF, st::Auxiliary, FT)
    @vars(
        environment::vars_state(m.environment, st, FT),
        updraft::vars_state(m.updraft, st, FT)
    )
end

function vars_state(::Updraft, ::Prognostic, FT)
    @vars(ρa::FT, ρaw::FT, ρaθ_liq::FT, ρaq_tot::FT,)
end

function vars_state(::Environment, ::Prognostic, FT)
    @vars(ρatke::FT, ρaθ_liq_cv::FT, ρaq_tot_cv::FT, ρaθ_liq_q_tot_cv::FT,)
end

function vars_state(m::NTuple{N, Updraft}, st::Prognostic, FT) where {N}
    return Tuple{ntuple(i -> vars_state(m[i], st, FT), N)...}
end

function vars_state(m::EDMF, st::Prognostic, FT)
    @vars(
        environment::vars_state(m.environment, st, FT),
        updraft::vars_state(m.updraft, st, FT)
    )
end

function vars_state(::Updraft, ::Gradient, FT)
    @vars(w::FT,)
end

function vars_state(::Environment, ::Gradient, FT)
    @vars(
        θ_liq::FT,
        q_tot::FT,
        w::FT,
        tke::FT,
        θ_liq_cv::FT,
        q_tot_cv::FT,
        θ_liq_q_tot_cv::FT,
        θv::FT,
        e::FT,
    )
end

function vars_state(m::NTuple{N, Updraft}, st::Gradient, FT) where {N}
    return Tuple{ntuple(i -> vars_state(m[i], st, FT), N)...}
end

function vars_state(m::EDMF, st::Gradient, FT)
    @vars(
        environment::vars_state(m.environment, st, FT),
        updraft::vars_state(m.updraft, st, FT)
    )
end

function vars_state(m::NTuple{N, Updraft}, st::GradientFlux, FT) where {N}
    return Tuple{ntuple(i -> vars_state(m[i], st, FT), N)...}
end

function vars_state(::Updraft, st::GradientFlux, FT)
    @vars(∇w::SVector{3, FT},)
end

function vars_state(::Environment, ::GradientFlux, FT)
    @vars(
        ∇θ_liq::SVector{3, FT},
        ∇q_tot::SVector{3, FT},
        ∇w::SVector{3, FT},
        ∇tke::SVector{3, FT},
        ∇θ_liq_cv::SVector{3, FT},
        ∇q_tot_cv::SVector{3, FT},
        ∇θ_liq_q_tot_cv::SVector{3, FT},
        ∇θv::SVector{3, FT},
        ∇e::SVector{3, FT},
    )

end

function vars_state(m::EDMF, st::GradientFlux, FT)
    @vars(
        S²::FT, # should be conditionally grabbed from atmos.turbulence
        environment::vars_state(m.environment, st, FT),
        updraft::vars_state(m.updraft, st, FT)
    )
end

"""
    init_aux_turbconv!(
        turbconv::EDMF{FT},
        m::AtmosModel{FT},
        aux::Vars,
        geom::LocalGeometry,
    ) where {FT}

Initialize EDMF auxiliary variables.
"""
function init_aux_turbconv!(
    turbconv::EDMF{FT},
    m::AtmosModel{FT},
    aux::Vars,
    geom::LocalGeometry,
) where {FT}
    N_up = n_updrafts(turbconv)

    # Aliases:
    en_aux = aux.turbconv.environment
    up_aux = aux.turbconv.updraft

    en_aux.cld_frac = FT(0)
    en_aux.buoyancy = FT(0)

    @unroll_map(N_up) do i
        up_aux[i].buoyancy = FT(0)
        up_aux[i].θ_liq = FT(0)
        up_aux[i].q_tot = FT(0)
        up_aux[i].w = FT(0)
    end
end;

function turbconv_nodal_update_auxiliary_state!(
    turbconv::EDMF{FT},
    m::AtmosModel{FT},
    state::Vars,
    aux::Vars,
    t::Real,
) where {FT}
    N_up = n_updrafts(turbconv)
    save_subdomain_temperature!(m, state, aux)

    en_aux = aux.turbconv.environment
    up_aux = aux.turbconv.updraft
    gm = state
    en = state.turbconv.environment
    up = state.turbconv.updraft

    # Recover thermo states
    ts = recover_thermo_state_all(m, state, aux)

    # Get environment variables
    env = environment_vars(state, aux, N_up)

    # Compute buoyancies of subdomains
    ρ_inv = 1 / gm.ρ
    _grav::FT = grav(m.param_set)

    z = altitude(m, aux)

    ρ_en = air_density(ts.en)
    en_aux.buoyancy = -_grav * (ρ_en - aux.ref_state.ρ) * ρ_inv

    @unroll_map(N_up) do i
        ρ_i = air_density(ts.up[i])
        up_aux[i].buoyancy = -_grav * (ρ_i - aux.ref_state.ρ) * ρ_inv
        up_aux[i].a = up[i].ρa * ρ_inv
        up_aux[i].θ_liq = up[i].ρaθ_liq / up[i].ρa
        up_aux[i].q_tot = up[i].ρaq_tot / up[i].ρa
        up_aux[i].w = up[i].ρaw / up[i].ρa
    end
    b_gm = grid_mean_b(state, aux, N_up)

    # remove the gm_b from all subdomains
    @unroll_map(N_up) do i
        up_aux[i].buoyancy -= b_gm
    end
    en_aux.buoyancy -= b_gm

    εδ_up = ntuple(N_up) do i
        entr_detr(m, m.turbconv.entr_detr, state, aux, t, ts, env, i)
    end

    ε_dyn, δ_dyn, ε_trb = ntuple(i -> map(x -> x[i], εδ_up), 3)

    @unroll_map(N_up) do i
        up_aux[i].ε_dyn = ε_dyn[i]
        up_aux[i].δ_dyn = δ_dyn[i]
        up_aux[i].ε_trb = ε_trb[i]
    end

end;

function compute_gradient_argument!(
    turbconv::EDMF{FT},
    m::AtmosModel{FT},
    transform::Vars,
    state::Vars,
    aux::Vars,
    t::Real,
) where {FT}
    N_up = n_updrafts(turbconv)
    z = altitude(m, aux)

    # Aliases:
    up_tf = transform.turbconv.updraft
    en_tf = transform.turbconv.environment
    gm = state
    up = state.turbconv.updraft
    en = state.turbconv.environment

    # Recover thermo states
    ts = recover_thermo_state_all(m, state, aux)

    # Get environment variables
    env = environment_vars(state, aux, N_up)

    @unroll_map(N_up) do i
        up_tf[i].w = up[i].ρaw / up[i].ρa
    end
    _grav::FT = grav(m.param_set)

    ρ_inv = 1 / gm.ρ
    θ_liq_en = liquid_ice_pottemp(ts.en)
    q_tot_en = total_specific_humidity(ts.en)

    # populate gradient arguments
    en_tf.θ_liq = θ_liq_en
    en_tf.q_tot = q_tot_en
    en_tf.w = env.w

    en_tf.tke = enforce_positivity(en.ρatke) / (env.a * gm.ρ)
    en_tf.θ_liq_cv = enforce_positivity(en.ρaθ_liq_cv) / (env.a * gm.ρ)
    en_tf.q_tot_cv = enforce_positivity(en.ρaq_tot_cv) / (env.a * gm.ρ)
    en_tf.θ_liq_q_tot_cv = en.ρaθ_liq_q_tot_cv / (env.a * gm.ρ)

    # TODO: is this supposed to be grabbed from grid mean?
    en_tf.θv = virtual_pottemp(ts.gm)
    e_kin = FT(1 // 2) * ((gm.ρu[1] * ρ_inv)^2 + (gm.ρu[2] * ρ_inv)^2 + env.w^2)
    en_tf.e = total_energy(e_kin, _grav * z, ts.en)
end;

function compute_gradient_flux!(
    turbconv::EDMF{FT},
    m::AtmosModel{FT},
    diffusive::Vars,
    ∇transform::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
) where {FT}
    N_up = n_updrafts(turbconv)

    # Aliases:
    gm = state
    up_dif = diffusive.turbconv.updraft
    tc_dif = diffusive.turbconv
    up_∇tf = ∇transform.turbconv.updraft
    en_dif = diffusive.turbconv.environment
    en_∇tf = ∇transform.turbconv.environment

    @unroll_map(N_up) do i
        up_dif[i].∇w = up_∇tf[i].w
    end

    ρ_inv = 1 / gm.ρ
    # first moment grid mean coming from environment gradients only
    en_dif.∇θ_liq = en_∇tf.θ_liq
    en_dif.∇q_tot = en_∇tf.q_tot
    en_dif.∇w = en_∇tf.w
    # second moment env cov
    en_dif.∇tke = en_∇tf.tke
    en_dif.∇θ_liq_cv = en_∇tf.θ_liq_cv
    en_dif.∇q_tot_cv = en_∇tf.q_tot_cv
    en_dif.∇θ_liq_q_tot_cv = en_∇tf.θ_liq_q_tot_cv

    en_dif.∇θv = en_∇tf.θv
    en_dif.∇e = en_∇tf.e

    tc_dif.S² = ∇transform.u[3, 1]^2 + ∇transform.u[3, 2]^2 + en_dif.∇w[3]^2
end;

function turbconv_source!(
    m::AtmosModel{FT},
    source::Vars,
    state::Vars,
    diffusive::Vars,
    aux::Vars,
    t::Real,
    direction,
) where {FT}
    turbconv = m.turbconv
    N_up = n_updrafts(turbconv)

    # Aliases:
    gm = state
    en = state.turbconv.environment
    up = state.turbconv.updraft
    en_src = source.turbconv.environment
    up_src = source.turbconv.updraft
    en_dif = diffusive.turbconv.environment
    up_aux = aux.turbconv.updraft

    # Recover thermo states
    ts = recover_thermo_state_all(m, state, aux)

    # Get environment variables
    env = environment_vars(state, aux, N_up)

    εδ_up = ntuple(N_up) do i
        entr_detr(m, m.turbconv.entr_detr, state, aux, t, ts, env, i)
    end
    ε_dyn, δ_dyn, ε_trb = ntuple(i -> map(x -> x[i], εδ_up), 3)

    # get environment values
    _grav::FT = grav(m.param_set)
    ρ_inv = 1 / gm.ρ
    θ_liq_en = liquid_ice_pottemp(ts.en)
    q_tot_en = total_specific_humidity(ts.en)
    tke_en = enforce_positivity(en.ρatke) * ρ_inv / env.a
    θ_liq = liquid_ice_pottemp(ts.gm)
    a_min = turbconv.subdomains.a_min
    a_max = turbconv.subdomains.a_max

    ρa_up = vuntuple(N_up) do i
        gm.ρ * enforce_unit_bounds(up[i].ρa * ρ_inv, a_min, a_max)
    end

    @unroll_map(N_up) do i

        ρa_up_i = ρa_up[i]
        w_up_i = up[i].ρaw / ρa_up_i
        ρa_up_i_inv = FT(1) / ρa_up_i

        # first moment sources - for now we compute these as aux variable
        dpdz = perturbation_pressure(
            m,
            m.turbconv.pressure,
            state,
            diffusive,
            aux,
            t,
            env,
            i,
        )

        # entrainment and detrainment
        up_src[i].ρa += up[i].ρaw * (ε_dyn[i] - δ_dyn[i])
        up_src[i].ρaw +=
            up[i].ρaw *
            ((ε_dyn[i] + ε_trb[i]) * env.w - (δ_dyn[i] + ε_trb[i]) * w_up_i)
        up_src[i].ρaθ_liq +=
            up[i].ρaw * (
                (ε_dyn[i] + ε_trb[i]) * θ_liq_en -
                (δ_dyn[i] + ε_trb[i]) * up[i].ρaθ_liq * ρa_up_i_inv
            )
        up_src[i].ρaq_tot +=
            up[i].ρaw * (
                (ε_dyn[i] + ε_trb[i]) * q_tot_en -
                (δ_dyn[i] + ε_trb[i]) * up[i].ρaq_tot * ρa_up_i_inv
            )

        # add buoyancy and perturbation pressure in subdomain w equation
        up_src[i].ρaw += up[i].ρa * (up_aux[i].buoyancy - dpdz)
        # microphysics sources should be applied here

        # environment second moments:
        en_src.ρatke += (
            up[i].ρaw *
            δ_dyn[i] *
            (w_up_i - env.w) *
            (w_up_i - env.w) *
            FT(0.5) +
            up[i].ρaw *
            ε_trb[i] *
            (env.w - gm.ρu[3] * ρ_inv) *
            (env.w - w_up_i) - up[i].ρaw * (ε_dyn[i] + ε_trb[i]) * tke_en
        )

        en_src.ρaθ_liq_cv += (
            up[i].ρaw *
            δ_dyn[i] *
            (up[i].ρaθ_liq * ρa_up_i_inv - θ_liq_en) *
            (up[i].ρaθ_liq * ρa_up_i_inv - θ_liq_en) +
            up[i].ρaw *
            ε_trb[i] *
            (θ_liq_en - θ_liq) *
            (θ_liq_en - up[i].ρaθ_liq * ρa_up_i_inv) +
            up[i].ρaw *
            ε_trb[i] *
            (θ_liq_en - θ_liq) *
            (θ_liq_en - up[i].ρaθ_liq * ρa_up_i_inv) -
            up[i].ρaw * (ε_dyn[i] + ε_trb[i]) * en.ρaθ_liq_cv
        )

        en_src.ρaq_tot_cv += (
            up[i].ρaw *
            δ_dyn[i] *
            (up[i].ρaq_tot * ρa_up_i_inv - q_tot_en) *
            (up[i].ρaq_tot * ρa_up_i_inv - q_tot_en) +
            up[i].ρaw *
            ε_trb[i] *
            (q_tot_en - gm.moisture.ρq_tot * ρ_inv) *
            (q_tot_en - up[i].ρaq_tot * ρa_up_i_inv) +
            up[i].ρaw *
            ε_trb[i] *
            (q_tot_en - gm.moisture.ρq_tot * ρ_inv) *
            (q_tot_en - up[i].ρaq_tot * ρa_up_i_inv) -
            up[i].ρaw * (ε_dyn[i] + ε_trb[i]) * en.ρaq_tot_cv
        )

        en_src.ρaθ_liq_q_tot_cv += (
            up[i].ρaw *
            δ_dyn[i] *
            (up[i].ρaθ_liq * ρa_up_i_inv - θ_liq_en) *
            (up[i].ρaq_tot * ρa_up_i_inv - q_tot_en) +
            up[i].ρaw *
            ε_trb[i] *
            (θ_liq_en - θ_liq) *
            (q_tot_en - up[i].ρaq_tot * ρa_up_i_inv) +
            up[i].ρaw *
            ε_trb[i] *
            (q_tot_en - gm.moisture.ρq_tot * ρ_inv) *
            (θ_liq_en - up[i].ρaθ_liq * ρa_up_i_inv) -
            up[i].ρaw * (ε_dyn[i] + ε_trb[i]) * en.ρaθ_liq_q_tot_cv
        )

        # pressure tke source from the i'th updraft
        en_src.ρatke += ρa_up_i * (w_up_i - env.w) * dpdz
    end
    l_mix = mixing_length(
        m,
        m.turbconv.mix_len,
        state,
        diffusive,
        aux,
        t,
        δ_dyn,
        ε_trb,
        ts,
        env,
    )
    K_eddy = m.turbconv.mix_len.c_m * l_mix * sqrt(tke_en)
    Shear² = diffusive.turbconv.S²
    ρa₀ = gm.ρ * env.a
    Diss₀ = m.turbconv.mix_len.c_d * sqrt(tke_en) / l_mix

    # production from mean gradient and Dissipation
    en_src.ρatke += ρa₀ * (K_eddy * Shear² - Diss₀ * tke_en)
    en_src.ρaθ_liq_cv +=
        ρa₀ *
        (K_eddy * en_dif.∇θ_liq[3] * en_dif.∇θ_liq[3] - Diss₀ * en.ρaθ_liq_cv)
    en_src.ρaq_tot_cv +=
        ρa₀ *
        (K_eddy * en_dif.∇q_tot[3] * en_dif.∇q_tot[3] - Diss₀ * en.ρaq_tot_cv)
    en_src.ρaθ_liq_q_tot_cv +=
        ρa₀ * (
            K_eddy * en_dif.∇θ_liq[3] * en_dif.∇q_tot[3] -
            Diss₀ * en.ρaθ_liq_q_tot_cv
        )
    # covariance microphysics sources should be applied here
end;

# # in the EDMF first order (advective) fluxes exist only in the grid mean (if <w> is nonzero) and the uprdafts
function flux_first_order!(
    turbconv::EDMF{FT},
    m::AtmosModel{FT},
    flux::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
) where {FT}
    # Aliases:
    gm = state
    up = state.turbconv.updraft
    up_flx = flux.turbconv.updraft
    N_up = n_updrafts(turbconv)

    ρ_inv = 1 / gm.ρ
    ẑ = vertical_unit_vector(m, aux)
    a_min = turbconv.subdomains.a_min
    a_max = turbconv.subdomains.a_max
    # in future GCM implementations we need to think about grid mean advection

    ρa_up = vuntuple(N_up) do i
        gm.ρ * enforce_unit_bounds(up[i].ρa * ρ_inv, a_min, a_max)
    end

    @unroll_map(N_up) do i
        ρa_i = ρa_up[i]
        up_flx[i].ρa = up[i].ρaw * ẑ
        w_up_i = up[i].ρaw / ρa_i
        up_flx[i].ρaw = up[i].ρaw * w_up_i * ẑ
        up_flx[i].ρaθ_liq = w_up_i * up[i].ρaθ_liq * ẑ
        up_flx[i].ρaq_tot = w_up_i * up[i].ρaq_tot * ẑ
    end
end;

# in the EDMF second order (diffusive) fluxes
# exist only in the grid mean and the environment
function flux_second_order!(
    turbconv::EDMF{FT},
    m::AtmosModel{FT},
    flux::Grad,
    state::Vars,
    diffusive::Vars,
    aux::Vars,
    t::Real,
) where {FT}
    N_up = n_updrafts(turbconv)

    # Aliases:
    gm = state
    up = state.turbconv.updraft
    en = state.turbconv.environment
    gm_flx = flux
    en_flx = flux.turbconv.environment
    en_dif = diffusive.turbconv.environment

    # Recover thermo states
    ts = recover_thermo_state_all(m, state, aux)

    # Get environment variables
    env = environment_vars(state, aux, N_up)

    ρ_inv = FT(1) / gm.ρ
    _grav::FT = grav(m.param_set)
    z = altitude(m, aux)
    a_min = turbconv.subdomains.a_min
    a_max = turbconv.subdomains.a_max

    εδ_up = ntuple(N_up) do i
        entr_detr(m, m.turbconv.entr_detr, state, aux, t, ts, env, i)
    end

    ε_dyn, δ_dyn, ε_trb = ntuple(i -> map(x -> x[i], εδ_up), 3)

    l_mix = mixing_length(
        m,
        turbconv.mix_len,
        state,
        diffusive,
        aux,
        t,
        δ_dyn,
        ε_trb,
        ts,
        env,
    )
    tke_en = enforce_positivity(en.ρatke) / env.a * ρ_inv
    K_eddy = m.turbconv.mix_len.c_m * l_mix * sqrt(tke_en)

    #TotalFlux(ϕ) = Eddy_Diffusivity(ϕ) + MassFlux(ϕ)
    e_int = internal_energy(m, state, aux)

    e_kin = vuntuple(N_up) do i
        FT(1 // 2) * (
            (gm.ρu[1] * ρ_inv)^2 +
            (gm.ρu[2] * ρ_inv)^2 +
            (up[i].ρaw / up[i].ρa)^2
        )
    end
    e_tot_up = ntuple(i -> total_energy(e_kin[i], _grav * z, ts.up[i]), N_up)
    ρa_up = vuntuple(N_up) do i
        gm.ρ * enforce_unit_bounds(up[i].ρa * ρ_inv, a_min, a_max)
    end

    massflux_e = sum(
        vuntuple(N_up) do i
            up[i].ρa *
            ρ_inv *
            (gm.ρe * ρ_inv - e_tot_up[i]) *
            (gm.ρu[3] * ρ_inv - up[i].ρaw / ρa_up[i])
        end,
    )

    massflux_q_tot = sum(
        vuntuple(N_up) do i
            up[i].ρa *
            ρ_inv *
            (gm.moisture.ρq_tot * ρ_inv - up[i].ρaq_tot / up[i].ρa) *
            (gm.ρu[3] * ρ_inv - up[i].ρaw / ρa_up[i])
        end,
    )

    massflux_w = sum(
        vuntuple(N_up) do i
            up[i].ρa *
            ρ_inv *
            (gm.ρu[3] * ρ_inv - up[i].ρaw / up[i].ρa) *
            (gm.ρu[3] * ρ_inv - up[i].ρaw / ρa_up[i])
        end,
    )

    # update grid mean flux_second_order
    ρe_sgs_flux = -gm.ρ * env.a * K_eddy * en_dif.∇e[3] + massflux_e
    ρq_tot_sgs_flux = -gm.ρ * env.a * K_eddy * en_dif.∇q_tot[3] + massflux_q_tot
    ρu_sgs_flux = -gm.ρ * env.a * K_eddy * en_dif.∇w[3] + massflux_w

    # for now the coupling to the dycore is commented out

    # gm_flx.ρe              += SVector{3,FT}(0,0,ρe_sgs_flux)
    # gm_flx.moisture.ρq_tot += SVector{3,FT}(0,0,ρq_tot_sgs_flux)
    # gm_flx.ρu              += SMatrix{3, 3, FT, 9}(
    #     0, 0, 0,
    #     0, 0, 0,
    #     0, 0, ρu_sgs_flux,
    # )

    ẑ = vertical_unit_vector(m, aux)
    # env second moment flux_second_order
    en_flx.ρatke = -gm.ρ * env.a * K_eddy * en_dif.∇tke[3] * ẑ
    en_flx.ρaθ_liq_cv = -gm.ρ * env.a * K_eddy * en_dif.∇θ_liq_cv[3] * ẑ
    en_flx.ρaq_tot_cv = -gm.ρ * env.a * K_eddy * en_dif.∇q_tot_cv[3] * ẑ
    en_flx.ρaθ_liq_q_tot_cv =
        -gm.ρ * env.a * K_eddy * en_dif.∇θ_liq_q_tot_cv[3] * ẑ
end;

# First order boundary conditions
function turbconv_boundary_state!(
    nf,
    bc::EDMFBCs,
    m::AtmosModel{FT},
    state⁺::Vars,
    aux⁺::Vars,
    n⁻,
    state⁻::Vars,
    aux⁻::Vars,
    bctype,
    t,
    state_int::Vars,
    aux_int::Vars,
) where {FT}

    turbconv = m.turbconv
    N_up = n_updrafts(turbconv)
    up = state⁺.turbconv.updraft
    en = state⁺.turbconv.environment
    gm = state⁺
    gm_a = aux⁺
    if bctype == 1 # bottom
        zLL = altitude(m, aux_int)
        a_up_surf,
        θ_liq_up_surf,
        q_tot_up_surf,
        θ_liq_cv,
        q_tot_cv,
        θ_liq_q_tot_cv,
        tke = subdomain_surface_values(
            turbconv.surface,
            turbconv,
            m,
            gm,
            gm_a,
            zLL,
        )

        @unroll_map(N_up) do i
            up[i].ρaw = FT(0)
            up[i].ρa = a_up_surf[i] * gm.ρ
            up[i].ρaθ_liq = up[i].ρa * θ_liq_up_surf[i]
            up[i].ρaq_tot = up[i].ρa * q_tot_up_surf[i]
        end
        a_en = environment_area(gm, gm_a, N_up)
        en.ρatke = gm.ρ * a_en * tke
        en.ρaθ_liq_cv = gm.ρ * a_en * θ_liq_cv
        en.ρaq_tot_cv = gm.ρ * a_en * q_tot_cv
        en.ρaθ_liq_q_tot_cv = gm.ρ * a_en * θ_liq_q_tot_cv
    end
end;

# The boundary conditions for second-order unknowns
function turbconv_normal_boundary_flux_second_order!(
    nf,
    bc::EDMFBCs,
    m::AtmosModel{FT},
    fluxᵀn::Vars,
    n⁻,
    state⁻::Vars,
    diff⁻::Vars,
    hyperdiff⁻::Vars,
    aux⁻::Vars,
    state⁺::Vars,
    diff⁺::Vars,
    hyperdiff⁺::Vars,
    aux⁺::Vars,
    bctype,
    t,
    _...,
) where {FT}
    turbconv = m.turbconv
    N = n_updrafts(turbconv)
    up_flx = fluxᵀn.turbconv.updraft
    en_flx = fluxᵀn.turbconv.environment
    if bctype == 2 # top
        @unroll_map(N_up) do i
            up_flx[i].ρaw = -n⁻ * FT(0)
            up_flx[i].ρa = -n⁻ * FT(0)
            up_flx[i].ρaθ_liq = -n⁻ * FT(0)
            up_flx[i].ρaq_tot = -n⁻ * FT(0)
        end
        en_flx.∇tke = -n⁻ * FT(0)
        en_flx.∇e_int_cv = -n⁻ * FT(0)
        en_flx.∇q_tot_cv = -n⁻ * FT(0)
        en_flx.∇e_int_q_tot_cv = -n⁻ * FT(0)
    end
end;

max_stars_repo_name: mwarusz/CLIMA

max_issues_repo_name: mwarusz/CLIMA

max_forks_repo_name: mwarusz/CLIMA


text:
const DEFAULT_BRACKETING_s = 1E-2
const DEFAULT_BRACKETING_k = 2.0

function bracket_minimum(f, x=0.0; s=DEFAULT_BRACKETING_s, k=DEFAULT_BRACKETING_k, nmax=DEFAULT_BRACKETING_nmax)
    a, ya = x, f(x)
    b, yb = a + s, f(a + s)
    if yb > ya
        a, b = b, a
        ya, yb = yb, ya
        s = -s
    end
    for _ in 1:nmax
        c, yc = b + s, f(b + s)
        # println("a=$a, b=$b, c=$c, f(a)=$ya, f(b)=$yb, f(c)=$yc")
        if yc > yb
            return a < c ? (a,c) : (c,a)
        end
        a, ya, b, yb = b, yb, c, yc
        s *= k
    end
    error("Unable to bracket f starting at x = $x.")
end

ext: jl


text:
"""
KitML.jl: The lightweight module of neural differential equations in Kinetic.jl

Copyright (c) 2021 Tianbai Xiao & Steffen Schotthöfer <tianbaixiao@gmail.com>
"""

module KitML

using KitBase
using KitBase.JLD2
using KitBase.PyCall
using KitBase.Reexport
@reexport using Solaris
using Solaris.DiffEqFlux

include("Equation/equation.jl")
include("Closure/closure.jl")
include("Solver/solver.jl")

end


text:
#
# Get atomic weigth from atom representation
#

function getaw(atom; awdict::Dict = atomic_weight)
    return awdict[lowercase(atom)]
end


text:
import DataFrames
import PredictMD
import Test

x = Union{Missing, String}["foo", "bar", "foo", missing, "bar"]

Test.@test(
    length(PredictMD.get_unique_values(x; skip_missings = true)) == 2
    )
Test.@test(
    length(PredictMD.get_unique_values(x; skip_missings = false)) == 3
    )
Test.@test(
    length(PredictMD.get_unique_values_skip_missings(x)) == 2
    )
Test.@test(
    length(PredictMD.get_unique_values_include_missings(x)) == 3
    )

Test.@test(
    PredictMD.get_number_of_unique_values(x; skip_missings = true) == 2
    )
Test.@test(
    PredictMD.get_number_of_unique_values(x; skip_missings = false) == 3
    )
Test.@test(
    PredictMD.get_number_of_unique_values_skip_missings(x) == 2
    )
Test.@test(
    PredictMD.get_number_of_unique_values_include_missings(x) == 3
    )

df = DataFrames.DataFrame()
df[:a] = [1,2,3,4,5,6,7]
df[:b] = [1,1,1,1,1,1,1]
df_constant_cols = PredictMD.find_constant_columns(df)
Test.@test length(df_constant_cols) == 1
Test.@test df_constant_cols == [:b]
Test.@test all(df_constant_cols .== [:b])


text:
# ### m13.2t.jl

using Pkg, DrWatson

@quickactivate "StatisticalRethinkingTuring"
using Turing
using StatisticalRethinking
Turing.setprogress!(false)

df = CSV.read(sr_datadir("UCBadmit.csv"), DataFrame);

dept_map = Dict(key => idx for (idx, key) in enumerate(unique(df.dept)))
df.male = ifelse.(df.gender .== "male", 1, 0)
df.dept_id = [dept_map[de] for de in df.dept]

@model ppl13_2(applications, dept_id, male, admit) = begin
    sigma_dept ~ truncated(Cauchy(0, 2), 0, Inf)
    bm ~ Normal(0, 1)
    a ~ Normal(0, 10)
    a_dept ~ filldist(Normal(a, sigma_dept), 6)

    logit_p = a_dept[dept_id] + bm*male

    admit .~ BinomialLogit.(applications, logit_p)
end

m13_2t = ppl13_2(df.applications, df.dept_id, df.male, df.admit)
nchains = 4; sampler = NUTS(0.65); nsamples=2000
nchains = 1 # Fails frequently with nchains=4
chns13_2t = mapreduce(c -> sample(m13_2t, sampler, nsamples), chainscat, 1:nchains)

m_13_2s_results = """
Inference for Stan model: 359c2483e3bdbf74fd0484be27c2909b.
    3 chains, each with iter=4500; warmup=500; thin=1;
    post-warmup draws per chain=4000, total post-warmup draws=12000.

                   mean se_mean   sd     2.5%      25%      50%      75%    97.5%
    a_dept[1]      0.67    0.00 0.10     0.48     0.61     0.67     0.74     0.87
    a_dept[2]      0.63    0.00 0.12     0.40     0.55     0.63     0.71     0.85
    a_dept[3]     -0.58    0.00 0.08    -0.73    -0.63    -0.58    -0.53    -0.44
    a_dept[4]     -0.62    0.00 0.09    -0.79    -0.67    -0.62    -0.56    -0.45
    a_dept[5]     -1.06    0.00 0.10    -1.26    -1.13    -1.06    -0.99    -0.86
    a_dept[6]     -2.61    0.00 0.16    -2.92    -2.71    -2.60    -2.50    -2.31
    a             -0.59    0.01 0.66    -1.94    -0.97    -0.59    -0.21     0.73
    bm            -0.09    0.00 0.08    -0.25    -0.15    -0.09    -0.04     0.06
    sigma_dept     1.49    0.01 0.62     0.78     1.09     1.35     1.71     3.01
    lp__       -2602.27    0.04 2.27 -2607.74 -2603.48 -2601.90 -2600.61 -2598.95
               n_eff Rhat
    a_dept[1]   8191    1
    a_dept[2]   8282    1
    a_dept[3]  10974    1
    a_dept[4]  10034    1
    a_dept[5]  10783    1
    a_dept[6]  10344    1
    a           5839    1
    bm          7137    1
    sigma_dept  4626    1
    lp__        4200    1
"""

# End of m13.2.jl


text:
mutable struct Simulation
    system::PSY.System
    reset::Bool
    problem::DiffEqBase.DAEProblem
    perturbations::Vector{<:Perturbation}
    x0_init::Vector{Float64}
    initialized::Bool
    tstops::Vector{Float64}
    callbacks::DiffEqBase.CallbackSet
    solution::Union{Nothing, DiffEqBase.DAESolution}
    simulation_folder::String
    ext::Dict{String, Any}
end

function Simulation(
    simulation_folder::String,
    system::PSY.System,
    tspan::NTuple{2, Float64},
    perturbations::Vector{<:Perturbation} = Vector{Perturbation}();
    kwargs...,
)
    check_folder(simulation_folder)
    simulation_system = deepcopy(system)
    PSY.set_units_base_system!(simulation_system, "device_base")
    check_kwargs(kwargs, SIMULATION_ACCEPTED_KWARGS, "Simulation")
    initialized = false
    DAE_vector = _index_dynamic_system!(simulation_system)
    var_count = get_variable_count(simulation_system)

    flat_start = zeros(var_count)
    bus_count = length(PSY.get_components(PSY.Bus, simulation_system))
    flat_start[1:bus_count] .= 1.0
    x0_init = get(kwargs, :initial_guess, flat_start)

    initialize_simulation = get(kwargs, :initialize_simulation, true)
    if initialize_simulation
        @info("Initializing Simulation States")
        _add_aux_arrays!(simulation_system, Real)
        initialized = calculate_initial_conditions!(simulation_system, x0_init)
    end

    dx0 = zeros(var_count)
    callback_set, tstops = _build_perturbations(simulation_system, perturbations)
    _add_aux_arrays!(simulation_system, Float64)
    prob = DiffEqBase.DAEProblem(
        system!,
        dx0,
        x0_init,
        tspan,
        simulation_system,
        differential_vars = DAE_vector;
        kwargs...,
    )

    if get(kwargs, :system_to_file, false)
        PSY.to_json(
            simulation_system,
            joinpath(simulation_folder, "initialized_system.json"),
        )
    end
    return Simulation(
        simulation_system,
        false,
        prob,
        perturbations,
        x0_init,
        initialized,
        tstops,
        callback_set,
        nothing,
        simulation_folder,
        Dict{String, Any}(),
    )
end

function Simulation(
    simulation_folder::String,
    system::PSY.System,
    tspan::NTuple{2, Float64},
    perturbation::Perturbation;
    initialize_simulation::Bool = true,
    kwargs...,
)
    return Simulation(
        simulation_folder,
        system,
        tspan,
        [perturbation];
        initialize_simulation = initialize_simulation,
        kwargs...,
    )
end

function _add_aux_arrays!(system::PSY.System, T)
    bus_count = get_bus_count(system)
    aux_arrays = Dict(
        1 => collect(zeros(T, bus_count)),                       #I_injections_r
        2 => collect(zeros(T, bus_count)),                       #I_injections_i
        3 => collect(zeros(T, get_n_injection_states(system))),  #injection_ode
        4 => collect(zeros(T, get_n_branches_states(system))),   #branches_ode
        5 => collect(zeros(Complex{T}, bus_count)),              #I_bus
        6 => collect(zeros(T, 2 * bus_count)),                   #I_balance
    )
    system.internal.ext[AUX_ARRAYS] = aux_arrays
    return
end

function _build_perturbations(system::PSY.System, perturbations::Vector{<:Perturbation})
    isempty(perturbations) && return DiffEqBase.CallbackSet(), [0.0]
    perturbations_count = length(perturbations)
    callback_vector = Vector{DiffEqBase.DiscreteCallback}(undef, perturbations_count)
    tstops = Vector{Float64}(undef, perturbations_count)
    for (ix, pert) in enumerate(perturbations)
        condition = (x, t, integrator) -> t in [pert.time]
        affect = get_affect(system, pert)
        callback_vector[ix] = DiffEqBase.DiscreteCallback(condition, affect)
        tstops[ix] = pert.time
    end
    callback_tuple = Tuple(cb for cb in callback_vector)
    callback_set = DiffEqBase.CallbackSet((), callback_tuple)
    return callback_set, tstops
end

function _index_local_states!(
    component_state_index::Vector{Int64},
    local_states::Vector{Symbol},
    component::PSY.DynamicComponent,
)
    for (ix, s) in enumerate(component.states)
        component_state_index[ix] = findfirst(x -> x == s, local_states)
    end
    return
end

function _attach_ports!(component::PSY.DynamicComponent)
    component.ext[PORTS] = Ports(component)
    return
end

function _attach_inner_vars!(
    device::PSY.DynamicGenerator,
    ::Type{T} = Real,
) where {T <: Real}
    device.ext[INNER_VARS] = zeros(T, 8)
    return
end

function _attach_inner_vars!(
    device::PSY.DynamicInverter,
    ::Type{T} = Real,
) where {T <: Real}
    device.ext[INNER_VARS] = zeros(T, 14)
    return
end

function _attach_control_refs!(device::PSY.DynamicInjection)
    device.ext[CONTROL_REFS] = [
        PSY.get_V_ref(device),
        PSY.get_ω_ref(device),
        PSY.get_P_ref(device),
        PSY.get_reactive_power(PSY.get_static_injector(device)),
    ]
    return
end

function _index_port_mapping!(
    index_component_inputs::Vector{Int64},
    local_states::Vector{Symbol},
    component::PSY.DynamicComponent,
)
    _attach_ports!(component)
    for i in component.ext[PORTS].state
        tmp = [(ix, var) for (ix, var) in enumerate(local_states) if var == i]
        isempty(tmp) && continue
        push!(index_component_inputs, tmp[1][1])
    end

    return
end

function _get_Ybus(sys::PSY.System)
    n_buses = length(PSY.get_components(PSY.Bus, sys))
    dyn_lines = PSY.get_components(PSY.DynamicBranch, sys)
    if !isempty(PSY.get_components(PSY.ACBranch, sys))
        Ybus_ = PSY.Ybus(sys)
        Ybus = Ybus_[:, :]
        lookup = Ybus_.lookup[1]
        if !isempty(dyn_lines)
            for br in dyn_lines
                ybus!(Ybus, br, lookup, -1.0)
            end
        end
    else
        Ybus = SparseMatrixCSC{Complex{Float64}, Int64}(zeros(n_buses, n_buses))
        lookup = Dict{Int.Int}()
    end
    return Ybus, lookup
end

function _make_device_index!(device::PSY.DynamicInjection, sys_basepower::Float64)
    states = PSY.get_states(device)
    device_state_mapping = Dict{Type{<:PSY.DynamicComponent}, Vector{Int64}}()
    input_port_mapping = Dict{Type{<:PSY.DynamicComponent}, Vector{Int64}}()
    _attach_inner_vars!(device)
    _attach_control_refs!(device)

    for c in PSY.get_dynamic_components(device)
        device_state_mapping[typeof(c)] = Vector{Int64}(undef, length(c.states))
        input_port_mapping[typeof(c)] = Vector{Int64}()
        _index_local_states!(device_state_mapping[typeof(c)], states, c)
        _index_port_mapping!(input_port_mapping[typeof(c)], states, c)
        device.ext[LOCAL_STATE_MAPPING] = device_state_mapping
        device.ext[INPUT_PORT_MAPPING] = input_port_mapping
    end

    return
end

function _add_states_to_global!(
    global_state_index::Dict{String, Dict{Symbol, Int64}},
    state_space_ix::Vector{Int64},
    device::PSY.Device,
)
    device_state_ix = Dict{Symbol, Int}()
    for s in PSY.get_states(device)
        state_space_ix[1] += 1
        device_state_ix[s] = state_space_ix[1]
    end
    global_state_index[PSY.get_name(device)] = device_state_ix
    return
end

function _index_dynamic_system!(sys::PSY.System)
    n_buses = length(PSY.get_components(PSY.Bus, sys))
    DAE_vector = collect(falses(n_buses * 2))
    global_state_index = Dict{String, Dict{Symbol, Int64}}()
    state_space_ix = [n_buses * 2]
    current_buses_ix = collect(1:n_buses)
    static_bus_var_count = 2 * length(current_buses_ix)
    voltage_buses_ix = Vector{Int}()
    total_states = 0
    first_dyn_branch_point = -1
    branches_n_states = 0
    global_vars = Dict{Symbol, Number}(
        :ω_sys => 1.0,
        :ω_sys_index => -1, #To define 0 if infinite source, bus_number otherwise,
    )
    total_shunts = Dict{Int, Float64}()
    found_ref_bus = false
    sys_basepower = PSY.get_base_power(sys)

    Ybus, lookup = _get_Ybus(sys)
    dyn_branches = PSY.get_components(PSY.DynamicBranch, sys)

    if !(isempty(dyn_branches))
        first_dyn_branch_point = state_space_ix[1] + 1
        for br in dyn_branches
            arc = PSY.get_arc(br)
            n_states = PSY.get_n_states(br)
            from_bus_number = PSY.get_number(arc.from)
            to_bus_number = PSY.get_number(arc.to)
            bus_ix_from = lookup[from_bus_number]
            bus_ix_to = lookup[to_bus_number]
            merge!(
                +,
                total_shunts,
                Dict(
                    bus_ix_from => 1 / PSY.get_b(br).from,
                    bus_ix_to => 1 / PSY.get_b(br).to,
                ),
            )
            push!(voltage_buses_ix, bus_ix_from, bus_ix_to)
            DAE_vector[bus_ix_from] = DAE_vector[bus_ix_from + n_buses] = true
            DAE_vector[bus_ix_to] = DAE_vector[bus_ix_to + n_buses] = true
            DAE_vector = push!(DAE_vector, collect(trues(n_states))...)
            total_states += n_states
            _add_states_to_global!(global_state_index, state_space_ix, br)
        end

        for (ix, val) in enumerate(DAE_vector[1:n_buses])
            if val
                #This V_ix should be V_number.
                global_state_index["V_$(ix)"] = Dict(:R => ix, :I => ix + n_buses)
                total_states += 2
                static_bus_var_count -= 2
                push!(voltage_buses_ix, ix)
                @assert static_bus_var_count >= 0
            end
        end
        branches_n_states = state_space_ix[1] - n_buses * 2
    else
        @debug("System doesn't contain Dynamic Branches")
    end
    unique!(voltage_buses_ix)
    sources = PSY.get_components(PSY.Source, sys)
    for s in sources
        btype = PSY.get_bustype(PSY.get_bus(s))
        if (btype == PSY.BusTypes.REF) && found_ref_bus
            throw(IS.ConflictingInputsError("The system can't have more than one source or generator in the REF Bus"))
        end
        btype != PSY.BusTypes.REF && continue
        global_vars[:ω_sys_index] = 0 #To define 0 if infinite source, bus_number otherwise,
        found_ref_bus = true
    end
    dynamic_injection = PSY.get_components(PSY.DynamicInjection, sys)
    isempty(dynamic_injection) &&
        error("System doesn't contain any DynamicInjection devices")
    for d in dynamic_injection
        if !(:states in fieldnames(typeof(d)))
            continue
        end
        device_bus = PSY.get_bus(d)
        btype = PSY.get_bustype(device_bus)
        if (btype == PSY.BusTypes.REF) && found_ref_bus
            throw(IS.ConflictingInputsError("The system can't have more than one source or generator in the REF Bus"))
        end
        _make_device_index!(d, sys_basepower)
        device_n_states = PSY.get_n_states(d)
        DAE_vector = push!(DAE_vector, collect(trues(device_n_states))...)
        total_states += device_n_states
        _add_states_to_global!(global_state_index, state_space_ix, d)
        btype != PSY.BusTypes.REF && continue
        global_vars[:ω_sys_index] = global_state_index[PSY.get_name(d)][:ω] #To define 0 if infinite source, bus_number otherwise,
        found_ref_bus = true
    end
    injection_n_states = state_space_ix[1] - branches_n_states - n_buses * 2
    @assert total_states == state_space_ix[1] - static_bus_var_count
    @debug total_states
    setdiff!(current_buses_ix, voltage_buses_ix)
    sys_ext = Dict{String, Any}()
    counts = Base.ImmutableDict(
        :total_states => total_states,
        :injection_n_states => injection_n_states,
        :branches_n_states => branches_n_states,
        :first_dyn_injection_pointer => 2 * n_buses + branches_n_states + 1,
        :first_dyn_branch_point => first_dyn_branch_point,
        :total_variables => total_states + static_bus_var_count,
        :bus_count => n_buses,
    )

    sys_ext[LITS_COUNTS] = counts
    sys_ext[GLOBAL_INDEX] = global_state_index
    sys_ext[VOLTAGE_BUSES_IX] = voltage_buses_ix
    sys_ext[CURRENT_BUSES_IX] = current_buses_ix
    sys_ext[YBUS] = Ybus
    sys_ext[TOTAL_SHUNTS] = total_shunts
    sys_ext[GLOBAL_VARS] = global_vars
    sys_ext[DYN_LINES] = !isempty(PSY.get_components(PSY.DynamicBranch, sys))
    sys_ext[LOOKUP] = lookup
    @assert sys_ext[GLOBAL_VARS][:ω_sys_index] != -1
    sys.internal.ext = sys_ext
    return DAE_vector
end

get_injection_pointer(sys::PSY.System) =
    PSY.get_ext(sys)[LITS_COUNTS][:first_dyn_injection_pointer]
get_branches_pointer(sys::PSY.System) =
    PSY.get_ext(sys)[LITS_COUNTS][:first_dyn_branch_point]
get_n_injection_states(sys::PSY.System) = PSY.get_ext(sys)[LITS_COUNTS][:injection_n_states]
get_n_branches_states(sys::PSY.System) = PSY.get_ext(sys)[LITS_COUNTS][:branches_n_states]
get_system_state_count(sys::PSY.System) = PSY.get_ext(sys)[LITS_COUNTS][:total_states]
get_variable_count(sys::PSY.System) = PSY.get_ext(sys)[LITS_COUNTS][:total_variables]
get_device_index(sys::PSY.System, device::D) where {D <: PSY.DynamicInjection} =
    PSY.get_ext(sys)[GLOBAL_INDEX][device.name]
get_inner_vars(device::PSY.DynamicInjection) = device.ext[INNER_VARS]
get_ω_sys(sys::PSY.System) = PSY.get_ext(sys)[GLOBAL_VARS][:ω_sys]
get_current_bus_ix(sys::PSY.System) = PSY.get_ext(sys)[CURRENT_BUSES_IX]
get_voltage_bus_ix(sys::PSY.System) = PSY.get_ext(sys)[VOLTAGE_BUSES_IX]
get_total_shunts(sys::PSY.System) = PSY.get_ext(sys)[TOTAL_SHUNTS]
get_bus_count(sys::PSY.System) = PSY.get_ext(sys)[LITS_COUNTS][:bus_count]
get_lookup(sys::PSY.System) = PSY.get_ext(sys)[LOOKUP]

function _get_internal_mapping(
    device::PSY.DynamicInjection,
    key::AbstractString,
    ty::Type{T},
) where {T <: PSY.DynamicComponent}
    device_index = PSY.get_ext(device)[key]
    val = get(device_index, ty, nothing)
    @assert !isnothing(val)
    return val
end

function get_local_state_ix(
    device::PSY.DynamicInjection,
    ty::Type{T},
) where {T <: PSY.DynamicComponent}
    return _get_internal_mapping(device, LOCAL_STATE_MAPPING, ty)
end

function get_input_port_ix(
    device::PSY.DynamicInjection,
    ty::Type{T},
) where {T <: PSY.DynamicComponent}
    return _get_internal_mapping(device, INPUT_PORT_MAPPING, ty)
end

function run_simulation!(sim::Simulation, solver; kwargs...)
    if sim.reset
        @error("Reset the simulation")
    end

    sim.solution = DiffEqBase.solve(
        sim.problem,
        solver;
        callback = sim.callbacks,
        tstops = sim.tstops,
        kwargs...,
    )
    return
end

function _change_vector_type(sys::PSY.System)
    for d in PSY.get_components(PSY.DynamicInjection, sys)
        _attach_inner_vars!(d, Real)
    end
end

function _determine_stability(vals::Vector{Complex{Float64}})
    stable = true
    for real_eig in real(vals)
        real_eig > 0.0 && return false
    end
    return true
end

function small_signal_analysis(sim::Simulation; kwargs...)
    if sim.reset
        @error("Reset the simulation")
    end
    _change_vector_type(sim.system)
    _add_aux_arrays!(sim.system, Real)
    var_count = LITS.get_variable_count(sim.system)
    dx0 = zeros(var_count) #Define a vector of zeros for the derivative
    bus_count = length(PSY.get_components(PSY.Bus, sim.system))
    sysf! = (out, x) -> LITS.system!(
        out,            #output of the function
        dx0,            #derivatives equal to zero
        x,              #states
        sim.system,     #Parameters
        0.0,            #time equals to zero.
    )
    out = zeros(var_count) #Define a vector of zeros for the output
    x_eval = get(kwargs, :operating_point, sim.x0_init)
    jacobian = ForwardDiff.jacobian(sysf!, out, x_eval)
    n_buses = length(PSY.get_components(PSY.Bus, sim.system))
    diff_states = collect(trues(var_count))
    diff_states[1:(2 * n_buses)] .= false
    for b_ix in get_voltage_bus_ix(sim.system)
        diff_states[b_ix] = true
        diff_states[b_ix + n_buses] = true
    end
    alg_states = .!diff_states
    fx = @view jacobian[diff_states, diff_states]
    gy = jacobian[alg_states, alg_states]
    fy = @view jacobian[diff_states, alg_states]
    gx = @view jacobian[alg_states, diff_states]
    # TODO: Make operation using BLAS!
    reduced_jacobian = fx - fy * inv(gy) * gx
    vals, vect = LinearAlgebra.eigen(reduced_jacobian)
    sources = collect(PSY.get_components(PSY.Source, sim.system))
    if isempty(sources)
        @warn("No Infinite Bus found. Confirm stability directly checking eigenvalues.\nIf all eigenvalues are on the left-half plane and only one eigenvalue is zero, the system is small signal stable.")
        info_evals = "Eigenvalues are:\n"
        for i in vals
            info_evals = info_evals * string(i) * "\n"
        end
        @info(info_evals)
    end
    return SmallSignalOutput(
        reduced_jacobian,
        vals,
        vect,
        _determine_stability(vals),
        x_eval,
    )
end


text:
using MultivariateStats
using LinearAlgebra
using Test
import Statistics: mean, cov, var
import Random
import SparseArrays
import StatsBase

Random.seed!(34568)

## PCA with zero mean

X = randn(5, 10)
Y = randn(3, 10)

W = qr(randn(5, 5)).Q[:, 1:3]
σ² = 0.1
M = PPCA(Float64[], W, σ²)

@test indim(M) == 5
@test outdim(M) == 3
@test mean(M) == zeros(5)
@test loadings(M) == W
@test var(M) == σ²

T = inv(W'*W .+ σ²*Matrix(I, 3, 3))*W'
@test transform(M, X[:,1]) ≈ T * X[:,1]
@test transform(M, X) ≈ T * X

R = W*inv(W'W)*(W'W .+ σ²*Matrix(I, 3, 3))
@test reconstruct(M, Y[:,1]) ≈ R * Y[:,1]
@test reconstruct(M, Y) ≈ R * Y

## PCA with non-zero mean

mval = rand(5)
M = PPCA(mval, W, σ²)

@test indim(M) == 5
@test outdim(M) == 3
@test mean(M) == mval
@test loadings(M) == W
@test var(M) == σ²

@test transform(M, X[:,1]) ≈ T * (X[:,1] .- mval)
@test transform(M, X) ≈ T * (X .- mval)

@test reconstruct(M, Y[:,1]) ≈ R * Y[:,1] .+ mval
@test reconstruct(M, Y) ≈ R * Y .+ mval

## prepare training data

d = 5
n = 1000

R = collect(qr(randn(d, d)).Q)
@test R'R ≈ Matrix(I, 5, 5)
rmul!(R, Diagonal(sqrt.([0.5, 0.3, 0.1, 0.05, 0.05])))

X = R'randn(5, n) .+ randn(5)
mval = vec(mean(X, dims=2))
Z = X .- mval

M0 = fit(PCA, X; mean=mval, maxoutdim = 4)

## ppcaml (default)

M = fit(PPCA, X)
P = projection(M)
W = loadings(M)

@test indim(M) == 5
@test outdim(M) == 4
@test mean(M) == mval
@test P'P ≈ Matrix(I, 4, 4)
@test reconstruct(M, transform(M, X)) ≈ reconstruct(M0, transform(M0, X))

M = fit(PPCA, X; mean=mval)
@test loadings(M) ≈ W

M = fit(PPCA, Z; mean=0)
@test loadings(M) ≈ W

M = fit(PPCA, X; maxoutdim=3)
P = projection(M)
W = loadings(M)

@test indim(M) == 5
@test outdim(M) == 3
@test P'P ≈ Matrix(I, 3, 3)

# ppcaem

M = fit(PPCA, X; method=:em)
P = projection(M)
W = loadings(M)

@test indim(M) == 5
@test outdim(M) == 4
@test mean(M) == mval
@test P'P ≈ Matrix(I, 4, 4)
@test all(isapprox.(reconstruct(M, transform(M, X)), reconstruct(M0, transform(M0, X)), atol=1e-3))

M = fit(PPCA, X; method=:em, mean=mval)
@test loadings(M) ≈ W

M = fit(PPCA, Z; method=:em, mean=0)
@test loadings(M) ≈ W

M = fit(PPCA, X; method=:em, maxoutdim=3)
P = projection(M)

@test indim(M) == 5
@test outdim(M) == 3
@test P'P ≈ Matrix(I, 3, 3)

@test_throws StatsBase.ConvergenceException fit(PPCA, X; method=:em, maxiter=1)

# bayespca
M0 = fit(PCA, X; mean=mval, maxoutdim = 3)

M = fit(PPCA, X; method=:bayes)
P = projection(M)
W = loadings(M)

@test indim(M) == 5
@test outdim(M) == 3
@test mean(M) == mval
@test P'P ≈ Matrix(I, 3, 3)
@test reconstruct(M, transform(M, X)) ≈ reconstruct(M0, transform(M0, X))

M = fit(PPCA, X; method=:bayes, mean=mval)
@test loadings(M) ≈ W

M = fit(PPCA, Z; method=:bayes, mean=0)
@test loadings(M) ≈ W

M = fit(PPCA, X; method=:em, maxoutdim=2)
P = projection(M)

@test indim(M) == 5
@test outdim(M) == 2
@test P'P ≈ Matrix(I, 2, 2)

@test_throws StatsBase.ConvergenceException fit(PPCA, X; method=:bayes, maxiter=1)

# test that fit works with Float32 values
X2 = convert(Array{Float32,2}, X)
# Float32 input, default pratio
M = fit(PPCA, X2; maxoutdim=3)
M = fit(PPCA, X2; maxoutdim=3, method=:em)
M = fit(PPCA, X2; maxoutdim=3, method=:bayes)

# views
M = fit(PPCA, view(X2, :, 1:100), maxoutdim=3)
M = fit(PPCA, view(X2, :, 1:100), maxoutdim=3, method=:em)
M = fit(PPCA, view(X2, :, 1:100), maxoutdim=3, method=:bayes)
# sparse
@test_throws AssertionError fit(PCA, SparseArrays.sprandn(100d, n, 0.6))


text:
function compute_solution2_fast(con_file::String, inl_file::String, raw_file::String, rop_file::String, time_limit::Int, scoring_method::Int, network_model::String; output_dir::String="", scenario_id::String="none")
    time_data_start = time()
    goc_data = parse_goc_files(con_file, inl_file, raw_file, rop_file, scenario_id=scenario_id)
    network = build_pm_model(goc_data)
    load_time = time() - time_data_start

    ###### Prepare Solution 2 ######

    time_contingencies_start = time()

    gen_cont_total = length(network["gen_contingencies"])
    branch_cont_total = length(network["branch_contingencies"])
    cont_total = gen_cont_total + branch_cont_total

    cont_order = contingency_order(network)

    #for cont in cont_order
    #    println(cont.label)
    #end

    workers = Distributed.workers()

    process_data = []

    cont_per_proc = cont_total/length(workers)

    for p in 1:length(workers)
        cont_start = trunc(Int, ceil(1+(p-1)*cont_per_proc))
        cont_end = min(cont_total, trunc(Int,ceil(p*cont_per_proc)))
        pd = (
            pid = p,
            processes = length(workers),
            con_file = con_file,
            inl_file = inl_file,
            raw_file = raw_file,
            rop_file = rop_file,
            scenario_id = scenario_id,
            output_dir = output_dir,
            cont_range = cont_start:cont_end,
        )
        #println(pd)
        push!(process_data, pd)
    end

    for (i,pd) in enumerate(process_data)
        info(LOGGER, "worker task $(pd.pid): $(length(pd.cont_range)) / $(pd.cont_range)")
    end

    #for (i,pd) in enumerate(process_data)
    #    println(pd.pid)
    #    for cont in cont_order[pd.cont_range]
    #        println("$(cont.label)")
    #    end
    #end

    solution2_files = pmap(solution2_solver_fast, process_data)

    sort!(solution2_files)

    #println("pmap result: $(solution2_files)")

    time_contingencies = time() - time_contingencies_start
    info(LOGGER, "contingency eval time: $(time_contingencies)")

    info(LOGGER, "combine $(length(solution2_files)) solution2 files")
    combine_files(solution2_files, "solution2.txt"; output_dir=output_dir)
    remove_files(solution2_files)

    println("")

    data = [
        "----",
        "scenario id",
        "bus",
        "branch",
        "gen_cont",
        "branch_cont",
        "runtime (sec.)",
    ]
    println(join(data, ", "))

    data = [
        "DATA_SSS",
        goc_data.scenario,
        length(network["bus"]),
        length(network["branch"]),
        length(network["gen_contingencies"]),
        length(network["branch_contingencies"]),
        time_contingencies,
    ]
    println(join(data, ", "))

    write_file_paths(goc_data.files; output_dir=output_dir)

    #println("")
    #write_evaluation_summary(goc_data, network, objective_lb=-Inf, load_time=load_time, contingency_time=time_contingencies, output_dir=output_dir)
end

@everywhere function solution2_solver_fast(process_data)
    #println(process_data)
    time_data_start = time()
    PowerModels.silence()
    goc_data = parse_goc_files(
        process_data.con_file, process_data.inl_file, process_data.raw_file,
        process_data.rop_file, scenario_id=process_data.scenario_id)
    network = build_pm_model(goc_data)

    sol = read_solution1(network, output_dir=process_data.output_dir)
    PowerModels.update_data!(network, sol)
    time_data = time() - time_data_start

    for (i,bus) in network["bus"]
        if haskey(bus, "evhi")
            bus["vmax"] = bus["evhi"]
        end
        if haskey(bus, "evlo")
            bus["vmin"] = bus["evlo"]
        end
    end

    contingencies = contingency_order(network)[process_data.cont_range]

    pad_size = trunc(Int, ceil(log(10,process_data.processes)))
    padded_pid = lpad(string(process_data.pid), pad_size, "0")
    solution_filename = "solution2-$(padded_pid).txt"

    if length(process_data.output_dir) > 0
        solution_path = joinpath(process_data.output_dir, solution_filename)
    else
        solution_path = solution_filename
    end
    if isfile(solution_path)
        warn(LOGGER, "removing existing solution2 file $(solution_path)")
        rm(solution_path)
    end
    open(solution_path, "w") do sol_file
        # creates an empty file in the case of workers without contingencies
    end

    for cont in contingencies
        if cont.type == "gen"
            info(LOGGER, "working on: $(cont.label)")
            time_start = time()
            sol_cont = deepcopy(sol)
            debug(LOGGER, "contingency copy time: $(time() - time_start)")

            sol_cont["label"] = cont.label
            sol_cont["feasible"] = true
            sol_cont["cont_type"] = "gen"
            sol_cont["cont_comp_id"] = cont.idx

            sol_cont["delta"] = 0.0
            sol_cont["gen"]["$(cont.idx)"]["pg"] = 0.0
            sol_cont["gen"]["$(cont.idx)"]["qg"] = 0.0

            correct_contingency_solution!(network, sol_cont)
            open(solution_path, "a") do sol_file
                sol2 = write_solution2_contingency(sol_file, network, sol_cont)
            end
        elseif cont.type == "branch"
            info(LOGGER, "working on: $(cont.label)")
            time_start = time()
            sol_cont = deepcopy(sol)
            debug(LOGGER, "contingency copy time: $(time() - time_start)")
            sol_cont = deepcopy(sol)

            sol_cont["label"] = cont.label
            sol_cont["feasible"] = true
            sol_cont["cont_type"] = "branch"
            sol_cont["cont_comp_id"] = cont.idx

            sol_cont["delta"] = 0.0

            correct_contingency_solution!(network, sol_cont)
            open(solution_path, "a") do sol_file
                sol2 = write_solution2_contingency(sol_file, network, sol_cont)
            end
        else
            @assert("contingency type $(cont.type) not known")
        end
    end

    return solution_path
end


text:
using pigpiod_if2_jll
include("./constants.jl")
#  SPI
#  spi_open	Opens a SPI device
#  int spi_open(int pi, unsigned spi_channel, unsigned baud, unsigned spi_flags)
#  This function returns a handle for the SPI device on the channel.
#  Data will be transferred at baud bits per second.
#  The flags may be used to modify the default behaviour of 4-wire operation, mode 0, active low chip select.
function spi_open(p::Int32, channel::Int, baud::Int, flags::Int)
    return ccall((:spi_open, libpigpiod_if2), Int32, (Int32, UInt32, UInt32, UInt32), p, channel, baud, flags)
end

#  spi_close	Closes a SPI device
#  int spi_close(int pi, unsigned handle)
#  This functions closes the SPI device identified by the handle.
function spi_close(p::Int32, handle::Int32)
    return ccall((:spi_close, libpigpiod_if2), Int32, (Int32, UInt32), p, handle)
end

#  spi_read	Reads bytes from a SPI device
#  int spi_read(int pi, unsigned handle, char *buf, unsigned count)
#  This function reads count bytes of data from the SPI device associated with the handle.
function spi_read(p::Int32, handle::Int32, buf::Vector{Int}, count::Int)
    return ccall((:spi_read, libpigpiod_if2), Int32, (Int32, UInt32, Ref{UInt8}, UInt32), p, handle, buf, count)
end

#  spi_write	Writes bytes to a SPI device
#  int spi_write(int pi, unsigned handle, char *buf, unsigned count)
#  This function writes count bytes of data from buf to the SPI device associated with the handle.
function spi_write(p::Int32, handle::Int32, buf::Vector{Int}, count::Int)
    return ccall((:spi_write, libpigpiod_if2), Int32, (Int32, UInt32, Ref{UInt8}, UInt32), p, handle, buf, count)
end

#  spi_xfer	Transfers bytes with a SPI device
#  int spi_xfer(int pi, unsigned handle, char *txBuf, char *rxBuf, unsigned count)
#  This function transfers count bytes of data from txBuf to the SPI device associated with the handle. Simultaneously count bytes of data are read from the device and placed in rxBuf.
function spi_xfer(p::Int32, handle::Int32, txBuf::Vector{Int}, rxBuf::Vector{Int}, count::Int)
    return ccall((:spi_xfer, libpigpiod_if2), Int32, (Int32, UInt32, Ref{UInt8}, Ref{UInt8}, UInt32), p, handle, txBuf, rxBuf, count)
end

#  SPI BIT BANG
#  bb_spi_open	Opens GPIO for bit banging SPI
#  int bb_spi_open(int pi, unsigned CS, unsigned MISO, unsigned MOSI, unsigned SCLK, unsigned baud, unsigned spi_flags)
#  This function selects a set of GPIO for bit banging SPI at a specified baud rate.
function bb_spi_open(p::Int32, cs::Int, miso::Int, mosi::Int, sclk::Int, baud::Int, flags::Int)
    return ccall((:bb_spi_open, libpigpiod_if2), Int32, (Int32, UInt32, UInt32, UInt32, UInt32, UInt32, UInt32), p, cs, miso, mosi, sclk, baud, flags)
end

#  bb_spi_close	Closes GPIO for bit banging SPI
#  int bb_spi_close(int pi, unsigned CS)
#  This function stops bit banging SPI on a set of GPIO opened with bbSPIOpen.
function bb_spi_close(p::Int32, cs::Int)
    return ccall((:bb_spi_close, libpigpiod_if2), Int32, (Int32, UInt32), p, cs)
end

#  bb_spi_xfer	Transfers bytes with bit banging SPI
#  int bb_spi_xfer(int pi, unsigned CS, char *txBuf, char *rxBuf, unsigned count)
#  This function executes a bit banged SPI transfer.
function bb_spi_xfer(p::Int32, cs::Int, txBuf::Vector{Int}, rxBuf::Vector{Int}, count::Int)
    return ccall((:bb_spi_xfer, libpigpiod_if2), Int32, (Int32, UInt32, Ref{UInt8}, Ref{UInt8}, UInt32), p, cs, txBuf, rxBuf, count)
end


text:
module Incentive

include("./../../constraints/geometric/constants.jl")
using JuMP
using .Constants

export incentive

function incentive(m)

    @expression(m, midpoint_of_support_x[i=1:length(boxes)],
        (m[:max_flbx_of_supporting_i][i] + m[:min_flbx_of_supporting_i][i])/2
    )

    @expression(m, midpoint_of_support_y[i=1:length(boxes)],
        (m[:max_flby_of_supporting_i][i] + m[:min_flby_of_supporting_i][i])/2
    )

    #Want to minimise the absolute value distance away from the midpoints that are suggested.
    @variable(m, incentive_x[i=1:length(boxes)])
    @constraint(m, incentive_x_min[i=1:length(boxes)],
        m[:incentive_x][i] >=
        m[:two_or_more_below][i] *
        (m[:midpoint_of_support_x][i] - (m[:flbx_front_left_bot_x][i] + m[:box_x_axis_cover][i]/2))
    )

    @constraint(m, incentive_x_max[i=1:length(boxes)-1],
        m[:incentive_x][i] >=
        m[:two_or_more_below][i] *
        (m[:flbx_front_left_bot_x][i] + m[:box_x_axis_cover][i]/2) - m[:midpoint_of_support_x][i]
    )

    @variable(m, incentive_y[i=1:length(boxes)])
    @constraint(m, incentive_y_min[i=1:length(boxes)],
        m[:incentive_y][i] >=
        m[:two_or_more_below][i] *
        m[:midpoint_of_support_y][i] - (m[:flby_front_left_bot_y][i] + m[:box_y_axis_cover][i]/2)
    )

    @constraint(m, incentive_y_max[i=1:length(boxes)],
        m[:incentive_y][i] >=
        m[:two_or_more_below][i] *
        (m[:flby_front_left_bot_y][i] + m[:box_y_axis_cover][i]/2) - m[:midpoint_of_support_y][i]
    )

    @expression(m, ds_incentive_obj,
        sum(m[:incentive_x]) +
        sum(m[:incentive_y])
    )

    return m
end

end


text:
md"""
# ЛЕКЦИЯ 6
- Пользовательские типы
- Обобщенные функции
(продолжение темы)

Рассмотрим задачу

## Задача 1
ДАНО: Робот - в некоторой клетке поля
РЕЗУЛЬТАТ: Робот в исходном положении и все клетки поля промаркированы в шахматном порядке так,
что в клетке с Роботом стоит маркер.

Следуя технологии сверху-вниз, и воспользовавшись некоторым ранее уже написанным кодом,
можем записать решение этой задачи в следующем виде.
"""

function chess_mark!(r)
    back_path = BackPath(r)
    #УТВ: Робот - в Ю-З углу

    chess_mark_from_sw!(r, iseven(num_steps(back_path)))
    #Поле промаркировано в шахматном порядке

    movements!(r,Sud)
    movements!(r,West)
    #УТВ: Робот снова находится в Ю-З углу
    back!(r, back_path)
    #УТВ: Робот - в исходном положении
end

num_steps(back_path::BackPath) = sum(back_path.path)

#-------------------------------------------------
md"""
Остаётся только реализовать функцию chess_mark_from_sw!,
которая должна будет перремещать робота змейкой по горизонтальным рядам, и расставлять маркеры в
шахматном порядке.

Для её реализации можно, например, ввести глобальнцю переменную flag::Bool, с тем чтобы при каждом
перемещении робота на 1 шаг ее значение менялось на противоположное, а маркер ставился бы только при
условии, что flag == true

Хотя решение с использованием глобальной переменной следует признать не очень хорошим, и такие решения
следует избегать, но начнем всё же с этого варианта (плохого!).
"""

flag = true

function chess_mark_from_sw!(r, marker_sw::Bool)
    # Робот - в юго-западном углу

    global flag = marker_sw

    side = Ost
    mark_line!(r,side)
    while !isborder(r,Nord)
        move!(r,Nord); flag = !flag
        side = inverse(side)
        mark_line!(r, side)
    end
end

function mark_line!(r, side::HorizonSide)
    global flag
    if flag
        putmarker!(r)
    end
    putmarkers!(r, side)
end

function putmarkers!(r, side::HorizonSide)
    global flag
    while !isborder(r,side)
        move!(r,side); flag = !flag
        if flag
            putmarker!(r)
        end
    end
end

md"""
Вариант с использованием глобальной переменной плох тем, что в результате мы получили функции, которые между
собой связаны через эту глобальную переменную в одно целое, их не возможно отлаживать и использовать
независимо друг от друга.

Кроме того, нам пришлось изменить код ранее написанных функций, которые составляли нашу собственную библиотеку
полезных функций.

Но от всех этих неприятностей можно легко избавиться, если специально для решения данной задачи сначала
спроектировать НОВОГО РОБОТА (новый пользовательский тип данных).
"""

#--------------------
mutable struct ChessRobot <: AbstractRobot
    robot::Robot
    flag::Bool
end

get(r::ChessRobot) = r.robot

function move!(r::ChessRobot, side::HorizonSide)
    move!(get(r),side)
    r.flag = !r.flag
end

function putmarker!(r::ChessRobot)
    if r.flag
        putmarker!(get(r))
    end
end

#----------------------------------------------

md"""
Тогда следующая функция mark_from_sw!(r::Any) будет обобщенной, т.е. довольно универсалной,
с помощью которой можно будет и замаркировать все клетки поля, и клетки поля в шахматном порядке,
и, может быть, в каком-либо еще порядке.

Конкретный результат будет зависеть от фактического типа ее аргумента.
Т.е. если аргументом будет робот типа Robot, то маркироваться будет все поле подряд,
если аргументом будет робот типа ChessRobot, то - в шахматном порядке.
"""

function mark_from_sw!(r) # эта функция ОБОБЩЕННАЯ
    # Робот - в юго-западном углу
    side = Ost
    mark_line!(r,side)
    while !isborder(r,Nord)
        move!(r,Nord); flag = !flag
        side = inverse(side)
        mark_line!(r, side)
    end
end

function mark_line!(r, side::HorizonSide)
    putmarker!(r)
    putmarkers!(r, side)
end

function putmarkers!(r, side::HorizonSide)
    while !isborder(r,side)
        move!(r,side)
        putmarker!(r)
    end
end

#-------------------------------------
md"""

Вот так теперь полученным обобщенным кодом можно воспользоваться для решения задачи 1:
"""

function chess_mark!(r::Robot)
    back_path = BackPath(r)
    #УТВ: Робот - в Ю-З углу

    r = ChessRobot(r, iseven(num_steps(back_path)))
    mark_from_sw!(r)
    #Поле промаркировано в шахматном порядке

    movements!(r,Sud)
    movements!(r,West)
    #УТВ: Робот снова находится в Ю-З углу
    back!(r, back_path)
    #УТВ: Робот - в исходном положении
end

num_steps(back_path::BackPath) = sum(back_path.path)

md"""
Тот же самый обобщённый код можно будет использовать и в случае, если понадобится ставить маркеры
каким-то ещё другим способом. Для этого понадобится только спроектировать робота специального типа
(специальный тип данных).

## Задача 2
ДАНО: Робот - в некоторой клетке поля
РЕЗУЛЬТАТ: Робот в исходном положении и все клетки поля промаркированы в шахматном порядке так,
что в клетке с Роботом стоит маркер.
"""

mutable struct NChessRobot <: AbstractCoordRobot
    robot::Robot
    N::Int
    coord::Coord
    NChessRobot(r,N::Int) = new(r,N,Coord())
end

get(r::NChessRobot) = r.robot

function putmarker!(r::ChessRobot)
    x,y = get_coord(r.coord)
    x = x % (2*r.N)
    y = y  % (2*r.N)
    if x in (0:r.N-1) && y in (0:r.N-1) || x in (r.N:2*r.N-1) && y in (r.N:2*r.N-1)
        putmarker!(get(r))
    end
end

md"""
## ЗАМЕЧАНИЕ
Фактически тип аргумента r разработанных нами обобщенных функций: mark_from_sw!, mark_line!, putmarkers!
есть тип Any.

Но это слишком широкий тип, и мы могли бы сузить допустимы тип, описав его как объединение
конкретного типа Robot и абстрактного типа AbstrsctRobot
"""
function mark_from_sw!(r::Union{Robot,AbstractRobot}) # эта функция ОБОБЩЕННАЯ
    # Робот - в юго-западном углу
    side = Ost
    mark_line!(r,side)
    while !isborder(r,Nord)
        move!(r,Nord); flag = !flag
        side = inverse(side)
        mark_line!(r, side)
    end
end

function mark_line!(r::Union{Robot,AbstractRobot}, side::HorizonSide)
    putmarker!(r)
    putmarkers!(r, side)
end

function putmarkers!(r::Union{Robot,AbstractRobot}, side::HorizonSide)
    while !isborder(r,side)
        move!(r,side)
        putmarker!(r)
    end
end

max_stars_repo_name: Igetsuga/AlgorithmicsHomework

max_issues_repo_name: Igetsuga/AlgorithmicsHomework

max_forks_repo_name: Igetsuga/AlgorithmicsHomework


text:
using MatrixAlgebra
using LinearAlgebra
using Test

@testset "MatrixAlgebra" begin
    @testset "tridiagonalize" begin include("tridiagonalize.jl") end
    @testset "zrst" begin include("zrst.jl") end
    @testset "determinant" begin include("determinant.jl") end
    @testset "rational approximation" begin include("rational.jl") end
end


text:
"""
    struct CartesianSpace <: EuclideanSpace{ℝ}

A real euclidean space `ℝ^d`, which is therefore self-dual. `CartesianSpace` has no
additonal structure and is completely characterised by its dimension `d`. This is the
vector space that is implicitly assumed in most of matrix algebra.
"""
struct CartesianSpace <: EuclideanSpace{ℝ}
    d::Int
end
CartesianSpace(V::CartesianSpace) = V

Base.getindex(::RealNumbers) = CartesianSpace
Base.getindex(::RealNumbers, d::Int) = CartesianSpace(d)
Base.:^(::RealNumbers, d::Int) = CartesianSpace(d)

# Corresponding methods:
#------------------------
dim(V::CartesianSpace) = V.d
Base.axes(V::CartesianSpace) = Base.OneTo(dim(V))

Base.oneunit(::Type{CartesianSpace}) = CartesianSpace(1)
⊕(V1::CartesianSpace, V2::CartesianSpace) = CartesianSpace(V1.d+V2.d)
fuse(V1::CartesianSpace, V2::CartesianSpace) = CartesianSpace(V1.d*V2.d)
flip(V::CartesianSpace) = V

Base.min(V1::CartesianSpace, V2::CartesianSpace) = CartesianSpace(min(V1.d, V2.d))
Base.max(V1::CartesianSpace, V2::CartesianSpace) = CartesianSpace(max(V1.d, V2.d))

Base.show(io::IO, V::CartesianSpace) = print(io, "ℝ^$(V.d)")
Base.show(io::IO, ::Type{CartesianSpace}) = print(io, "CartesianSpace")


text:
# Examples of abstract types:
abstract type Number end
abstract type Real     <: Number end
abstract type AbstractFloat <: Real end
abstract type Integer  <: Real end
abstract type Signed   <: Integer end
abstract type Unsigned <: Integer end

# Examples of primitive types:
primitive type Char 32 end
primitive type Bool <: Integer 8 end

primitive type Int8    <: Signed   8 end
primitive type Int16   <: Signed   16 end
primitive type Int32   <: Signed   32 end
primitive type Int64   <: Signed   64 end

primitive type UInt8   <: Unsigned 8 end
primitive type UInt16  <: Unsigned 16 end
primitive type UInt32  <: Unsigned 32 end
primitive type UInt64  <: Unsigned 64 end

primitive type Float16 <: AbstractFloat 16 end
primitive type Float32 <: AbstractFloat 32 end
primitive type Float64 <: AbstractFloat 64 end

max_stars_repo_name: elucian/julia

max_issues_repo_name: elucian/julia

max_forks_repo_name: elucian/julia


text:
module RNNAutoencoder

import Flux:onehotbatch
using Flux
using Flux:reset!
using Random
using Printf

struct SequentialData
    n_char
    n_seq
    charlist
    id_sos
    id_eos
    id_pad
    strs
    sequence_data

    function SequentialData(strs::Vector{String}, charlist::Vector{String})
        n_seq = maximum(map(length, strs)) + 2
        n_char = length(charlist)
        id_sos = n_char - 2
        id_eos = n_char - 1
        id_pad = n_char

        strs_edited = map(x -> ["<SOS>", string.(collect(x))..., "<EOS>"], strs)
        data = map(x->append!(x, collect(Iterators.repeated("<PAD>", n_seq-length(x)))), strs_edited)
        data_char = reduce(hcat, data)
        sequence_data = [onehotbatch(data_char[i, :], charlist) for i in 1:n_seq]

        return new(n_char, n_seq, charlist, id_sos, id_eos, id_pad, strs, sequence_data)
    end
end

function get_shuffled_traindata(sequence_data, batch_size)
    n_data = size(sequence_data[1])[2]
    vindices = collect(Iterators.partition(randperm(n_data), batch_size))
    return [([x[:, indices] for x in sequence_data],) for indices in vindices]
end

onehot_to_indices(onehot) = reduce(vcat, mapslices(x->argmax(x), onehot, dims=1))

struct Encoder
    embed
    gru
    n_char
    function Encoder(features, n_char)
        embed = Flux.glorot_uniform(features, n_char) |> param
        gru = GRU(features, features)
        return new(embed, gru, n_char)
    end
end
Flux.@treelike Encoder # params()の呼び出しを可能にする

function (e::Encoder)(x)
    reset!(e.gru)
    vindices = map(onehot_to_indices, x)
    x = map(x1 -> e.embed[:, x1], vindices)
    hidden = e.gru.(x)[end]
    return hidden
end

struct Decoder
    embed
    gru
    dense
    n_char
    n_seq
    function Decoder(features, n_char, n_seq)
        embed = Flux.glorot_uniform(features, n_char) |> param
        gru = GRU(features, features)
        dense = Dense(features, n_char, relu)
        return new(embed, gru, dense, n_char, n_seq)
    end
end
Flux.@treelike Decoder

function (d::Decoder)(h)
    reset!(d.gru)
    _, batch_size = size(h)
    d.gru.state = h
    sos = d.embed[:, [d.n_char-1 for _ in 1:batch_size]]
    x = d.gru(sos)
    x = d.dense(x)
    categorical = softmax(x)
    for i in 1:d.n_seq-2 # exclude <SOS>
        indices = onehot_to_indices(x)
        x2 = d.embed[:, indices]
        x = d.gru(x2) |> d.dense
        categorical = hcat(categorical, softmax(x))
    end
    categorical
end

struct Seq2Seq
    encoder::Encoder
    decoder::Decoder
    function Seq2Seq(features, n_char, n_seq)
        encoder = Encoder(features, n_char)
        decoder = Decoder(features, n_char, n_seq)
        return new(encoder, decoder)
    end
end
Flux.@treelike Seq2Seq

loss_total = 0
function loss_func(m::Seq2Seq)
    weight = ones(m.decoder.n_char)
    weight[end] = 0 # end is <PAD> index
    return function loss(seq_x)
        hidden_state = m.encoder(reverse(seq_x))
        predict = m.decoder(hidden_state)
        t = reduce(hcat, map(x->x, seq_x[2:end])) # 1 is <SOS> so skip
        loss_batch = Flux.crossentropy(predict, t, weight=weight)
        global loss_total += loss_batch.data
        return loss_batch
    end
end

end # module

max_stars_repo_name: Soyukke/FluxModels

max_issues_repo_name: Soyukke/FluxModels

max_forks_repo_name: Soyukke/FluxModels


text:
# Autogenerated wrapper script for JpegTurbo_jll for x86_64-apple-darwin
export cjpeg, djpeg, jpegtran, libjpeg, libturbojpeg

JLLWrappers.@generate_wrapper_header("JpegTurbo")
JLLWrappers.@declare_executable_product(cjpeg)
JLLWrappers.@declare_executable_product(djpeg)
JLLWrappers.@declare_executable_product(jpegtran)
JLLWrappers.@declare_library_product(libjpeg, "@rpath/libjpeg.62.dylib")
JLLWrappers.@declare_library_product(libturbojpeg, "@rpath/libturbojpeg.0.dylib")
function __init__()
    JLLWrappers.@generate_init_header()
    JLLWrappers.@init_executable_product(
        cjpeg,
        "bin/cjpeg",
    )

    JLLWrappers.@init_executable_product(
        djpeg,
        "bin/djpeg",
    )

    JLLWrappers.@init_executable_product(
        jpegtran,
        "bin/jpegtran",
    )

    JLLWrappers.@init_library_product(
        libjpeg,
        "lib/libjpeg.62.3.0.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libturbojpeg,
        "lib/libturbojpeg.0.2.0.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
"""
  DomainNorm(xy,cm,u;normID="Inf")

Computes either the L-1, L-2, or L-inf norm on a given domain. Expects quadrilateral elements, and order 1 or 2 polynomial basis.

ASSUMPTIONS:
- Uses strings for L-norm specification ... change this to symbols

INPUT:
  xy:     [scalar array]
            Node List for node positions
  cm:     [Matrix]
            Connectivity Matrix
  u:      [Array]
            Scalar field to put in integrand
  normID: [String]
            indicates L-1,2,Inf norm
OUTPUT:
  norm:  [scalar array]
            (1,2,Inf)-Norm of field u
"""
function DomainNorm(xy,cm,u;normID="Inf")
  norm  = 0.0
  Nel   = length(cm)
  ngpt  = length(cm[1].NodeList)
  order = (ngpt==4 ? 1 : 2)

  # generate basis values at gauss points
  w,s,t = GaussQuadPoints2D(order+1)

  # L1 or L2 norm
  if (normID=="1" || normID=="2")
    normIDnum = str2float(normID)

    for el=1:Nel
      xNodes = [xy[cm[el].NodeList[i]].x for i=1:ngpt]
      yNodes = [xy[cm[el].NodeList[i]].y for i=1:ngpt]

      for gpt=1:ngpt
        phi,_,_,jac = derivShape2D(s[gpt],t[gpt],xNodes,yNodes,order)

	      ucoefs = [u[cm[el].NodeList[i]] for i=1:ngpt]
	      ugpt = shapeEval(ucoefs,phi)

	      norm += w[gpt]*abs(ugpt)^normIDnum*jac
      end
    end

    # take sqrt of norm value for L2 norm
    if normID=="2"
      norm = sqrt(norm)
    end

  # infinity norm (aka max norm)
  elseif normID=="Inf"
    norm = maximum(abs,u)
  else
    error("invalid normID given. Choices are \"1\", \"2\", or \"Inf\".")
  end

  return norm
end

"""
  hCalc(mesh)

Computes the spacing parameter h for a given mesh. Works on quadrilateral meshes

ASSUMPTIONS:
- Most effective when all elements are roughly the same size. This
    algorithm will not measure coarseness of stretched meshes accurately.

INPUT:
  mesh:  [AbstractMesh]
            Mesh that includes nodes and connectivity

OUTPUT:
  h:     [AbstractFloat]
            Maximum edge distance of mesh, measure for mesh refinement
"""
function hCalc(mesh)
  Nel = length(mesh.cm)
  h   = 0.0

  for el=1:Nel
    for i=1:4
      j = mod(i,4)+1
      ci = mesh.cm[el].NodeList[i]
      cj = mesh.cm[el].NodeList[j]

      xdiff = mesh.xy[ci].x - mesh.xy[cj].x
      ydiff = mesh.xy[ci].y - mesh.xy[cj].y

      htemp = sqrt(xdiff^2 + ydiff^2)

      if htemp > h
        h = htemp
      end

    end
  end

  return h
end


text:
using Test
using MLDatasets
using MLDatasets: SupervisedDataset, UnsupervisedDataset, AbstractDataset
using MLDatasets: Graph
using FileIO
using DataDeps
using DataFrames, CSV, Tables
using HDF5
using JLD2
using ColorTypes
using FixedPointNumbers
using MLDatasets
using DataDeps

ENV["DATADEPS_ALWAYS_ACCEPT"] = true

include("test_utils.jl")

dataset_tests = [
    "datasets/graphs.jl",
    "datasets/misc.jl",
    "datasets/vision/emnist.jl",
    "datasets/vision/fashion_mnist.jl",
    "datasets/vision/mnist.jl",
    "datasets/text.jl",
]

no_ci_dataset_tests = [
    "datasets/graphs_no_ci.jl",
    "datasets/vision/cifar10.jl",
    "datasets/vision/cifar100.jl",
    "datasets/vision/svhn2.jl",
    ]

@assert isempty(intersect(dataset_tests, no_ci_dataset_tests))

deprecated_interface = [
    # "datasets/vision/deprecated_fashion_mnist.jl",
    # "datasets/vision/deprecated_mnist.jl",
    # "datasets/vision/deprecated_cifar10.jl",
    # "datasets/vision/deprecated_cifar100.jl",
    # "datasets/vision/deprecated_svhn2.jl", # NOT OK
    # "datasets/graphs/deprecated_citeseer.jl",
    # "datasets/graphs/deprecated_cora.jl",
    # "datasets/graphs/deprecated_pubmed.jl",
    # "datasets/graphs/deprecated_tudataset.jl",
    # "datasets/graphs/deprecated_polblogs.jl",
    # "datasets/graphs/deprecated_karateclub.jl",
    # "datasets/text_deprecated.jl",
    # "datasets/misc_deprecated.jl",
]

container_tests = [
    "containers/filedataset.jl",
    "containers/tabledataset.jl",
    "containers/hdf5dataset.jl",
    "containers/jld2dataset.jl",
    "containers/cacheddataset.jl",
]

@testset "Datasets" begin
    @testset "$(split(t,"/")[end])" for t in dataset_tests
        include(t)
    end

    if !parse(Bool, get(ENV, "CI", "false"))
        @info "Testing larger datasets"
        @testset "$(split(t,"/")[end])" for t in no_ci_dataset_tests
            include(t)
        end
    else
        @info "CI detected: skipping tests on large datasets"
    end

    if !isempty(deprecated_interface)
        @info "Testing deprecated dataset interface"
        @testset "$(split(t,"/")[end])" for t in deprecated_interface
            include(t)
        end
    end
end

@testset "Containers" begin
    for t in container_tests
        include(t)
    end
end

nothing