module Intermediate

#
#   This module provides some intial calculations for things like numbers and
#   ages of children, number of carers and so on, for benefit units and households.
#   It is intended to simplify subsequent modelling of benefits, etc..
#   and to reduce the number of repeated calculations. It also holds code for some age and
#   child count calculations that are also used stand-alone elsewhere.
#
#   FIXME some of the names should be globally changed.
#

using Base: Bool, String
using Dates: TimeType, Date, now, Year

using ScottishTaxBenefitModel
using .Definitions

using .ModelHousehold:
    BenefitUnit,
    Household,
    Person,
    between_ages,
    count,
    default_bu_allocation,
    empl_status_in,
    ge_age,
    get_benefit_units,
    get_head,
    has_carer_member,
    has_children,
    has_disabled_member,
    has_income,
    is_lone_parent,
    is_severe_disability,
    is_single,
    le_age,
    num_adults,
    num_carers,
    num_people,
    pers_is_carer,
    pers_is_disabled,
    search

using .STBParameters:
    AgeLimits,
    HoursLimits,
    ChildLimits,
    reached_state_pension_age

using .Utils:
    has_non_z,
    haskeys,
    mult,
    to_md_table

using .STBIncomes

export
    MTIntermediate,
    apply_2_child_policy,
    born_before,
    has_limited_capactity_for_work_activity,
    has_limited_capactity_for_work,
    is_working_hours,
    make_intermediate,
    make_recipient,
    num_born_before,
    to_string,
    working_disabled,
    HHIntermed

"""
examples:

2 children born before start_date, 1 after
 => allowable = 2
3  children born before start_date, 1 after
 => allowable = 3
1  children born before start_date, 2 after
 => allowable = 2
0  children born before start_date, 2 after
 => allowable = 2
 @return number of children allowed
"""
function apply_2_child_policy(
    bu      :: BenefitUnit,
    limits  :: ChildLimits
    ;
    model_run_date :: TimeType = now() ) :: Integer
    before_children = 0
    after_children = 0
    for pid in bu.children
        ch = bu.people[pid]
        if born_before( ch.age, limits.policy_start, model_run_date )
            before_children += 1
        else
            after_children += 1
        end
    end
    allowable = before_children + min( max(limits.max_children-before_children,0), after_children )
end

function born_before( age :: Integer,
    start_date     :: TimeType = Date( 2017, 4, 6 ), # 6th April 2017
    model_run_date :: TimeType = now() )
    bdate = model_run_date - Year(age)
    return bdate < start_date
end

"""
Very loose implementation of CPAG 2020/1 ch 45
"""
function has_limited_capactity_for_work( pers :: Person ) :: Bool
    # based on cpag 2020/1 ch45 pp 1100-
    l =  pers.employment_status in
        [   Retired,
            Permanently_sick_or_disabled,
            Temporarily_sick_or_injured,
            Other_Inactive] &&  # FIXME - retired/temp sick/looking after family ??
        pers_is_disabled( pers )
    l = l || (pers.esa_type != no_jsa)
    return l
end

"""
Very loose implementation of CPAG 2020/1 ch 45. Until we have
something better, maybe regression-based.
`Work Activity` here means attending interviews, assessments and the like.
"""
function has_limited_capactity_for_work_activity( pers :: Person ) :: Bool
    l = false
    if has_limited_capactity_for_work( pers )
        has_high_esa = false
        if has_non_z( pers.income, employment_and_support_allowance )
            esa = pers.income[employment_and_support_allowance]
            if( pers.age < 25 )
                if esa > 80 # crudely including support component
                    has_high_esa = true
                end
            else
                if esa > 100 # crudely including support component
                    has_high_esa = true
                end
            end
        end # crude ESA check
        if is_severe_disability( pers ) || has_high_esa
            l = true
        end
    end
    return l
end

function num_born_before(
    bu             :: BenefitUnit,
    start_date     :: TimeType = Date( 2017, 4, 6 ), # 6th April 2017
    model_run_date :: TimeType = now()) :: Integer
    nb = 0
    for pid in bu.children
        ch = bu.people[pid]
        if born_before( ch.age, start_date, model_run_date )
            nb += 1
        end
    end
    return nb
end

function is_working_hours( pers :: Person, hours... ) :: Bool
    if length(hours) == 1
        return (pers.usual_hours_worked >= hours[1])
    elseif length(hours) == 2
        return (pers.usual_hours_worked >= hours[1]) &&
               (pers.usual_hours_worked <= hours[2])
    end
end

function is_working_lt_hours( pers :: Person, hours :: Real ) :: Bool
    return pers.usual_hours_worked < hours
end

"""
See CPAG ch 61 p 1426 and appendix 5
"""
function working_disabled( pers::Person, hrs :: HoursLimits ) :: Bool
    if pers.usual_hours_worked >= hrs.lower || pers.employment_status in [Full_time_Employee, Full_time_Self_Employed]
        if pers.registered_blind || pers.registered_partially_sighted || pers.registered_deaf
            return true
        end
        for (dis, t ) in pers.disabilities
            return true
        end
        if haskeys( pers.income,
            [
                Incapacity_Benefit,
                Severe_Disability_Allowance,
                Employment_and_Support_Allowance ])
            return true
        end
    end
    return false
end

#
# FIXME some names here need clarified
#
mutable struct MTIntermediate

    benefit_unit_number :: Int
    num_people :: Int
    age_youngest_adult :: Int
    age_oldest_adult :: Int
    age_youngest_child :: Int
    age_oldest_child :: Int
    num_adults :: Int
    someone_pension_age  :: Bool
    someone_pension_age_2016 :: Bool
    all_pension_age :: Bool
    someone_working_ft  :: Bool
    #
    someone_working_ft_and_25_plus :: Bool

    num_not_working :: Int
    num_working_ft :: Int
    num_working_pt :: Int
    num_working_24_plus :: Int
    num_working_16_or_less :: Int
    total_hours_worked :: Int
    someone_is_carer :: Bool
    num_carers :: Int

    is_sparent  :: Bool
    is_sing  :: Bool # FIXME RENAME: is_single
    is_disabled :: Bool

    num_disabled_adults :: Int
    num_disabled_children :: Int
    num_severely_disabled_adults :: Int
    num_severely_disabled_children :: Int

    num_job_seekers :: Int

    num_children :: Int
    num_allowed_children :: Int
    num_children_born_before :: Int
    ge_16_u_pension_age  :: Bool
    limited_capacity_for_work  :: Bool
    has_children  :: Bool
    economically_active :: Bool
    working_disabled :: Bool
    num_benefit_units :: Int
end

struct HHIntermed
    hhint :: MTIntermediate
    buint ::  Vector{MTIntermediate}
end

function to_string( it :: MTIntermediate )::String
    s = to_md_table( it, depth=2 )
    return s
end

function to_string( hh :: HHIntermed ) :: String
    s = "## Household\n"
    s *= to_string( hh.hhint )
    for i in eachindex( hh.buint )
        s *= "### Bu $i"
        s *= to_string( hh.buint[i])
    end
    return s
end

"""
Assign CB, etc. to 1st adult female in a benefit unit, everything else to head of bu
"""
function make_recipient( bu :: BenefitUnit, which :: Incomes ) :: BigInt
    ## FIXME assert check here for non bu level benefits e.g. PIPs
    #
    # FIXME maybe also check relationships to children?
    if which in [CHILD_TAX_CREDIT, CHILD_BENEFIT, SCOTTISH_CHILD_PAYMENT ]
        # first_adult_female
        for pid in bu.adults
            if bu.people[pid].sex == Female
                return pid
            end
        end
    end
    ## FIXME expand this to check for if head is retired for WTC/UC
    # see the thing in .UniversalCredit.jl
    return get_head( bu ).pid
end

function aggregate!( sum :: MTIntermediate, add :: MTIntermediate )
    # benefit_unit_number :: Int
    sum.num_people += add.num_people
    sum.age_youngest_adult = min( sum.age_youngest_adult, add.age_youngest_adult )
    sum.age_oldest_adult = max( sum.age_oldest_adult, add.age_oldest_adult )
    sum.age_youngest_child = min( sum.age_youngest_child, add.age_youngest_child )
    sum.age_oldest_child = max( sum.age_oldest_child, add.age_oldest_child )
    sum.num_adults += add.num_adults
    sum.someone_pension_age = add.someone_pension_age || add.someone_pension_age
    sum.someone_pension_age_2016 = sum.someone_pension_age_2016 || add.someone_pension_age_2016
    sum.all_pension_age = sum.all_pension_age && add.all_pension_age
    sum.someone_working_ft  =  sum.someone_working_ft  || add.someone_working_ft # someone FIXME
    sum.someone_working_ft_and_25_plus =  sum.someone_working_ft_and_25_plus  || add.someone_working_ft_and_25_plus # someone FIXME

    sum.num_working_pt += add.num_working_pt
    sum.num_working_ft += add.num_working_ft
    sum.num_not_working += add.num_not_working

    sum.num_working_24_plus += add.num_working_24_plus
    sum.num_working_16_or_less += add.num_working_16_or_less
    sum.total_hours_worked += add.total_hours_worked
    sum.someone_is_carer = sum.someone_is_carer || add.someone_is_carer # rename someone FIXME
    sum.num_carers += add.num_carers

    sum.is_sparent = sum.is_sparent || add.is_sparent
    sum.is_sing = false # since we're adding a 2nd bu
    sum.is_disabled = sum.is_disabled || add.is_disabled

    sum.num_disabled_adults += add.num_disabled_adults
    sum.num_disabled_children += add.num_disabled_children
    sum.num_severely_disabled_adults += add.num_severely_disabled_adults
    sum.num_severely_disabled_children += add.num_severely_disabled_children

    sum.num_job_seekers += add.num_job_seekers

    sum.num_children += add.num_children
    sum.num_allowed_children += add.num_allowed_children
    sum.num_children_born_before += add.num_children_born_before
    sum.ge_16_u_pension_age = sum.ge_16_u_pension_age || add.ge_16_u_pension_age # SOMEONE: FIXME RENAME
    sum.limited_capacity_for_work = sum.limited_capacity_for_work || add.limited_capacity_for_work
    sum.has_children = sum.has_children || add.has_children
    sum.economically_active = sum.economically_active || add.economically_active
    sum.working_disabled = sum.working_disabled || add.working_disabled
end

function make_intermediate(
    buno :: Int,
    bu   :: BenefitUnit,
    hrs  :: HoursLimits,
    age_limits :: AgeLimits,
    child_limits :: ChildLimits,
    num_benefit_units :: Int ) :: MTIntermediate
    # {RT} where RT
    age_youngest_adult :: Int = 9999
    age_oldest_adult :: Int = -9999
    age_youngest_child :: Int = 9999
    age_oldest_child :: Int = -9999

    is_working_disabled :: Bool = false
    num_adlts :: Int = num_adults( bu )
    # check both & die as we need to think about counts again
    # if we're going back in time to when these were unequal
    #
    num_pens_age :: Int = 0
    ge_16_u_pension_age  :: Bool = false
    someone_pension_age_2016 :: Bool = false
    limited_capacity_for_work = false
    num_working_16_or_less  = 0
    # :: Int = count( bu, is_working_lt_hours, hrs.lower )

    for pid in bu.adults
        pers = bu.people[pid]
        if has_limited_capactity_for_work( pers )
            limited_capacity_for_work = true
        end
        if pers.usual_hours_worked < hrs.lower
            num_working_16_or_less += 1
        end
        if reached_state_pension_age(
            age_limits,
            pers.age,
            pers.sex )
            num_pens_age += 1
        else
            ge_16_u_pension_age = true
        end
        if reached_state_pension_age(
            age_limits,
            pers.age,
            pers.sex,
            age_limits.savings_credit_to_new_state_pension )
            someone_pension_age_2016 = true
        end

    end
    someone_pension_age  :: Bool = num_pens_age > 0
    all_pension_age :: Bool = num_adlts == num_pens_age
    num_working_ft :: Int = count( bu, is_working_hours, hrs.higher )
    num_working_pt :: Int = count( bu, is_working_hours, hrs.lower, hrs.higher-0.000001 ) # eps(FloatType)
    num_not_working :: Int = count( bu, empl_status_in,
       Unemployed,
       Retired,
       Student,
       Looking_after_family_or_home,
       Permanently_sick_or_disabled,
       Temporarily_sick_or_injured,
       Other_Inactive
        )

    num_job_seekers = 0
    num_working_24_plus :: Int = count( bu, is_working_hours, hrs.med )
    someone_working_ft  :: Bool = num_working_ft > 0
    someone_working_ft_and_25_plus :: Bool = false

    total_hours_worked :: Int = 0
    num_carrs :: Int = num_carers( bu )
    someone_is_carer :: Bool = has_carer_member( bu )
    is_sparent  :: Bool = is_lone_parent( bu )
    is_sing  :: Bool = is_single( bu )

    has_children  :: Bool = ModelHousehold.has_children( bu )
    economically_active = search( bu, empl_status_in,
        Full_time_Employee,
        Part_time_Employee,
        Full_time_Self_Employed,
        Part_time_Self_Employed,
        Unemployed,
        Temporarily_sick_or_injured )

    # can't think of a simple way of doing the rest with searches..

    is_disabled = has_disabled_member( bu )
    num_disabled_adults = 0
    num_severely_disabled_adults = 0
    for pid in bu.adults
        pers = bu.people[pid]

        if( pers.age >= 25 ) && ( pers.usual_hours_worked >= hrs.higher )
            someone_working_ft_and_25_plus = true
        end

        if (pers.employment_status == Unemployed) || ( pers.jsa_type != no_jsa )
            num_job_seekers += 1
        end

        if pers.age > age_oldest_adult
            age_oldest_adult = pers.age
        end
        if pers.age < age_youngest_adult
            age_youngest_adult = pers.age
        end
        total_hours_worked += round(pers.usual_hours_worked)
        if working_disabled( pers, hrs )
            is_working_disabled = true
        end
        if pers_is_disabled( pers )
            num_disabled_adults += 1
            if is_severe_disability( pers )
                num_severely_disabled_adults += 1
            end
        end
    end
    @assert 120 >= age_oldest_adult >= age_youngest_adult >= 16
    num_children = size( bu.children )[1] # count( bu, le_age, 16 )
    num_children_born_before = num_born_before( bu ) # fixme parameterise
    num_disabled_children = 0
    num_severely_disabled_children :: Int = 0
    for pid in bu.children
        pers = bu.people[pid]
        if pers.age > age_oldest_child
            age_oldest_child = pers.age
        end
        if pers.age < age_youngest_child
            age_youngest_child = pers.age
        end
        if pers_is_disabled( pers )
            num_disabled_children += 1
            if is_severe_disability( pers )
                num_severely_disabled_children += 1
            end
        end
    end

    ## fixme parameterise this
    num_allowed_children :: Int = apply_2_child_policy( bu, child_limits )
    @assert (!has_children)||(19 >= age_oldest_child >= age_youngest_child >= 0)

    return MTIntermediate(
        buno,
        num_people( bu ),
        age_youngest_adult,
        age_oldest_adult,
        age_youngest_child,
        age_oldest_child,
        num_adlts,
        someone_pension_age,
        someone_pension_age_2016,
        all_pension_age,
        someone_working_ft,

        someone_working_ft_and_25_plus,

        num_not_working,
        num_working_ft,
        num_working_pt,
        num_working_24_plus,
        num_working_16_or_less,
        total_hours_worked,
        someone_is_carer,
        num_carrs,
        is_sparent,
        is_sing,
        is_disabled,
        num_disabled_adults,
        num_disabled_children,
        num_severely_disabled_adults,
        num_severely_disabled_children,

        num_job_seekers,

        num_children,
        num_allowed_children,
        num_children_born_before,
        ge_16_u_pension_age,
        limited_capacity_for_work,
        has_children,
        economically_active,
        is_working_disabled,
        num_benefit_units
    )
end

function make_intermediate(
    hh   :: Household,
    hrs  :: HoursLimits,
    age_limits :: AgeLimits,
    child_limits :: ChildLimits,
    allocator :: Function=default_bu_allocation ) :: HHIntermed

    bus = get_benefit_units( hh, allocator )
    n = size( bus )[1]
    buint = Vector{MTIntermediate}(undef,n)
    for buno in 1:n
        buint[buno] = make_intermediate(
            buno,
            bus[buno],
            hrs,
            age_limits,
            child_limits,
            n )
    end
    hhint = deepcopy( buint[1] )
    for buno in 2:n
        aggregate!( hhint, buint[buno])
    end
    return HHIntermed( hhint, buint )
end

end # module Intermediate


# Importing Gurobi in this way avoids a warning with Requires
using .Gurobi: Gurobi

function run_scenario_gurobi(; solver_kwargs::Union{Dict,Nothing}=nothing, kwargs...)
    solver_kwargs = something(solver_kwargs, Dict("Method" => 2, "Crossover" => 0))
    try
        global env = Gurobi.Env()
    catch e
        println("Error encountered starting Gurobi.")
        throw(e)
    end
    try
        global m = run_scenario(;
            optimizer_factory=env, solver_kwargs=solver_kwargs, kwargs...
        )
        Gurobi.finalize(JuMP.backend(m))
    finally
        Gurobi.finalize(env)
        println("Connection to Gurobi closed successfully!")
    end
    # Return `nothing` to prevent `m` from the `try` block from being returned
    return nothing
end

function new_model(env::Gurobi.Env)
    return JuMP.direct_model(Gurobi.Optimizer(env))
end


using GPUInspector
using Test
using CUDA
using InteractiveUtils: subtypes
using LinearAlgebra

if !GPUInspector.functional()
    error("Can't run testsuite since CUDA/GPU not present or not functional!")
end
if Threads.nthreads() == 1
    error("Can't run testsuite in single-threaded mode. Start Julia with multiple threads.")
end
if Threads.nthreads() < length(CUDA.devices()) + 1
    # we should have at least one thread per gpu + one monitoring thread
    @warn(
        "You should run the tests with at least $(length(CUDA.devices()) + 1) Julia threads.",
        Threads.nthreads(),
        length(CUDA.devices())
    )
end

@testset "GPUInspector.jl" begin
    @testset "UnitPrefixedBytes" begin
        # general stuff
        b = B(40_000_000)
        m = MiB(38.14697265625)
        m̃ = MB(40.0)
        @test typeof(b) == B
        @test typeof(m) == MiB
        @test typeof(m̃) == MB
        @test string(b) == "4.0e7 B"
        @test string(m) == "~38.15 MiB"
        @test string(m̃) == "40.0 MB"
        @test value(b) == 40_000_000
        @test value(m) ≈ 38.14697265625
        @test value(m̃) ≈ 40.0
        @test simplify(b) ≈ m
        @test simplify(m) ≈ m
        @test simplify(m̃) ≈ m
        @test simplify(b; base=10) ≈ m̃
        @test simplify(m; base=10) ≈ m̃
        @test simplify(m̃; base=10) ≈ m̃
        @test change_base(b) == b
        @test change_base(m) == m̃
        @test change_base(m̃) == m
        @test convert(MB, b) == m̃
        @test convert(MiB, b) == m
        @test convert(MB, m) == m̃
        @test convert(MiB, m̃) == m
        @test bytes(b) == 4.0e7
        @test bytes(m) == 4.0e7
        @test bytes(m̃) == 4.0e7
        @test bytes(4.0e7) ≈ m

        # GiB <-> GB conversion because it's particularly important
        @test convert(GiB, GB(300)) ≈ GiB(279.39677238464355)
        @test change_base(GB(300)) ≈ GiB(279.39677238464355)
        @test convert(GB, GiB(300)) ≈ GB(322.12254720000004)
        @test change_base(GiB(300)) ≈ GB(322.12254720000004)

        # ≈, ==, ===
        @test b == m
        @test b == m̃
        @test m == m̃
        @test b ≈ m
        @test b ≈ m̃
        @test m ≈ m̃
        @test b === b
        @test b !== m
        @test b !== m̃
        @test m !== m̃

        # basic arithmetics
        types = subtypes(UnitPrefixedBytes)
        for T in types
            @test T(1.2) + T(2.3) ≈ T(3.5)
            @test T(3.5) - T(2.3) ≈ T(1.2)
            @test 2 * T(1.234) ≈ T(2.468)
            @test T(10) / 2 == T(5)
        end
        for T in types, S in types
            @test bytes(T(1.2) + S(2.3)) ≈ bytes(T(1.2)) + bytes(S(2.3))
            @test abs(bytes(T(1.2) - S(2.3))) ≈ abs(bytes(T(1.2)) - bytes(S(2.3)))
        end
        @test B(40_000_000) + MB(3) - 2 * KiB(2) ≈ MB(42.995904)
    end

    @testset "toggle_tensorcoremath" begin
        @test isnothing(toggle_tensorcoremath(true; verbose=false))
        @test CUDA.math_mode() == CUDA.FAST_MATH
        @test isnothing(toggle_tensorcoremath(false; verbose=false))
        @test CUDA.math_mode() == CUDA.DEFAULT_MATH
        # test toggle
        @test isnothing(toggle_tensorcoremath(; verbose=false))
        @test CUDA.math_mode() == CUDA.FAST_MATH
        @test isnothing(toggle_tensorcoremath(; verbose=false))
        @test CUDA.math_mode() == CUDA.DEFAULT_MATH
    end

    @testset "gpu stresstest" begin
        @testset "different kinds" begin
            for dev in (device(), devices()) # single- and multi-gpu
                @test isnothing(stresstest(dev; duration=2, verbose=false))
                @test isnothing(stresstest(dev; enforced_duration=2, verbose=false))
                @test isnothing(stresstest(dev; approx_duration=2, verbose=false))
                @test isnothing(stresstest(dev; niter=100, verbose=false))
                @test isnothing(stresstest(dev; niter=100, verbose=false))
                @test isnothing(stresstest(dev; mem=0.2, verbose=false))
            end
        end

        @testset "keyword options" begin
            dev = device()
            @test isnothing(stresstest(dev; duration=2, verbose=false))
            @test isnothing(stresstest(dev; duration=2, size=3000, verbose=false))
            @test isnothing(stresstest(dev; duration=2, dtype=Float16, verbose=false))
            @test isnothing(stresstest(dev; duration=2, clearmem=true, verbose=false))
            @test isnothing(stresstest(dev; duration=2, batch_duration=1, verbose=false))
            # TODO: kwargs: threads, parallel
        end

        @testset "monitoring" begin
            @testset "automatic (monitoring=true)" begin
                for dev in (device(), devices()) # single- and multi-gpu
                    @test typeof(
                        stresstest(dev; duration=2, verbose=false, monitoring=true)
                    ) == MonitoringResults
                end
            end
            @testset "manual" begin
                for devs in ([device()], devices()) # single- and multi-gpu
                    @test isnothing(monitoring_start(; freq=1, devices=devs, verbose=false))
                    @test isnothing(
                        stresstest(devs; duration=2, verbose=false, monitoring=false)
                    )
                    @test typeof(monitoring_stop(; verbose=false)) == MonitoringResults
                end
                # thread kwarg
                # @test isnothing(monitoring_start(; freq=1, devices=[device()], verbose=false, thread=2))
                # sleep(0.5)
                # TODO: How to find the threadid() of a Task?
                # @test isnothing(monitoring_stop(; verbose=false))
            end
        end

        @testset "monitoring results" begin
            @testset "MonitoringResults" begin
                r = stresstest(device(); duration=2, verbose=false, monitoring=true)
                @test typeof(r) == MonitoringResults
                @test typeof(r.times) == Vector{Float64}
                @test typeof(r.devices) == Vector{Tuple{String,Base.UUID}}
                quant = first(keys(r.results))
                @test r.results[quant] == getproperty(r, quant)
            end
            @testset "save / load" begin
                d = Dict{Symbol,Vector{Vector{Float64}}}()
                ndevs = length(CUDA.devices())
                d[:asd] = [rand(ndevs) for _ in 1:5]
                d[:qwe] = [rand(ndevs) for _ in 1:5]
                d[:jkl] = [rand(ndevs) for _ in 1:5]
                devices = Tuple{String,Base.UUID}[
                    (GPUInspector._device2string(dev), uuid(dev)) for
                    dev in collect(CUDA.devices())
                ]
                r = MonitoringResults(rand(5), devices, d)
                cd(mktempdir()) do
                    save_monitoring_results("tmp.h5", r)
                    r2 = load_monitoring_results("tmp.h5")
                    @test r == r2
                end
            end
            @testset "plot" begin
                r = load_monitoring_results(joinpath(@__DIR__, "test.h5"))
                @test isnothing(plot_monitoring_results(r))
                @test isnothing(plot_monitoring_results(r, (:compute, :mem)))
            end

            @testset "savefig" begin
                using CairoMakie
                r = load_monitoring_results(joinpath(@__DIR__, "test.h5"))
                @test isnothing(savefig_monitoring_results(r))
                @test isnothing(
                    savefig_monitoring_results(r, (:compute, :mem))
                )
                @test isnothing(savefig_monitoring_results(r; ext=:png))
                @test isnothing(savefig_monitoring_results(r; ext=:pdf))
                rm.(filter(endswith(".pdf"), readdir())) # cleanup
            end
        end
    end

    @testset "gpuinfo / gpus" begin
        @test isnothing(gpus())
        @test isnothing(gpuinfo())
        @test isnothing(gpuinfo(0))
        @test isnothing(gpuinfo(device()))
    end

    @testset "p2p_bandwidth" begin
        @testset "unidirectional" begin
            # p2p_bandwidth
            @test typeof(p2p_bandwidth(; verbose=false)) == Float64
            @test 0 ≤ p2p_bandwidth(; verbose=false)
            # options
            @test typeof(p2p_bandwidth(MB(100); verbose=false)) == Float64
            @test typeof(
                p2p_bandwidth(; src=CuDevice(0), dst=CuDevice(1), verbose=false)
            ) == Float64
            @test typeof(p2p_bandwidth(; dtype=Float16, verbose=false)) == Float64
            @test typeof(p2p_bandwidth(; nbench=10, verbose=false)) == Float64
            @test typeof(p2p_bandwidth(; hist=true, verbose=true)) == Float64
            # p2p_bandwidth_all
            @test typeof(p2p_bandwidth_all(; verbose=false)) ==
                Matrix{Union{Nothing,Float64}}
            Mp2p = p2p_bandwidth_all(; verbose=false)
            @test all(isnothing, diag(Mp2p))
            @test all(
                !isnothing(Mp2p[i, j]) for i in axes(Mp2p, 1), j in axes(Mp2p, 2) if i != j
            )
        end
        @testset "bidirectional" begin
            # p2p_bandwidth_bidirectional
            @test typeof(p2p_bandwidth_bidirectional(; verbose=false)) == Float64
            @test 0 ≤ p2p_bandwidth_bidirectional(; verbose=false)
            # options
            @test typeof(p2p_bandwidth_bidirectional(MB(100); verbose=false)) == Float64
            @test typeof(p2p_bandwidth_bidirectional(; dtype=Float16, verbose=false)) ==
                Float64
            @test typeof(p2p_bandwidth_bidirectional(; nbench=10, verbose=false)) ==
                Float64
            @test typeof(p2p_bandwidth_bidirectional(; hist=true, verbose=true)) ==
                Float64
            # p2p_bandwidth_bidirectional_all
            @test typeof(p2p_bandwidth_bidirectional_all(; verbose=false)) ==
                Matrix{Union{Nothing,Float64}}
            Mp2p = p2p_bandwidth_bidirectional_all(; verbose=false)
            @test all(isnothing, diag(Mp2p))
            @test all(
                !isnothing(Mp2p[i, j]) for i in axes(Mp2p, 1), j in axes(Mp2p, 2) if i != j
            )
        end
    end

    @testset "host2device_bandwidth" begin
        @test isnothing(host2device_bandwidth())
        @test isnothing(host2device_bandwidth(MB(100)))
        @test isnothing(host2device_bandwidth(; dtype=Float16))
    end

    @testset "memory_bandwidth" begin
        @test typeof(memory_bandwidth()) == Float64
        @test typeof(memory_bandwidth(MiB(10))) == Float64
        @test typeof(memory_bandwidth(; dtype=Float32)) == Float64

        @test typeof(memory_bandwidth_saxpy()) == Float64
        @test typeof(memory_bandwidth_saxpy(; size=2^20*2)) == Float64
        @test typeof(memory_bandwidth_saxpy(; dtype=Float32)) == Float64
    end

    @testset "peakflops_gpu (CUDA cores)" begin
        @test typeof(peakflops_gpu(; verbose=false, tensorcores=false)) == Float64
        @test typeof(peakflops_gpu(; dtype=Float32, verbose=false, tensorcores=false)) == Float64
        @test typeof(peakflops_gpu(; dtype=Float64, verbose=false, tensorcores=false)) == Float64
    end

    @testset "peakflops_gpu (Tensor cores)" begin
        @test typeof(peakflops_gpu(; verbose=false, tensorcores=true)) == Float64
        @test typeof(peakflops_gpu(; dtype=Float16, verbose=false, tensorcores=true)) == Float64
    end

    @testset "peakflops_gpu_matmul / scaling" begin
        @test typeof(peakflops_gpu_matmul(; verbose=false)) == Float64
        @test typeof(peakflops_gpu_matmul(; size=1024, dtype=Float64, verbose=false)) == Float64
        @test typeof(peakflops_gpu_matmul(; nmatmuls=2, nbench=2, verbose=false)) == Float64
        @test typeof(peakflops_gpu_matmul_scaling(; verbose=false)) == Tuple{Vector{Int64}, Vector{Float64}}
    end
end


facts("Population") do

  context("FloatVectorPopulation constructor") do
    p1 = FloatVectorPopulation(10, 2)
    @fact typeof(p1) => FloatVectorPopulation

    @fact size(p1.individuals, 1) => 10
    @fact size(p1.individuals, 2) => 2

    @fact size(p1.fitness, 1) => 10
    @fact size(p1.fitness, 2) => 2 # This is a bug since we cannot now the size of the fitness vectors before we have evaluated one...

    @fact size(p1.top, 1) => 10
    @fact size(p1.top, 2) => 2

    @fact size(p1.top_fitness, 1) => 10
    @fact size(p1.top_fitness, 2) => 2 # This is a bug since we cannot now the size of the fitness vectors before we have evaluated one...
  end

end


export prepare, getedge

using PyCall

# https://github.com/dhoegh/Hawk.jl/blob/master/src/Hawk.jl

filename = abspath(joinpath(dirname(@__FILE__), "hed.py"))

@pyimport imp
(path, name) = dirname(filename), basename(filename)
(name, ext) = rsplit(name, '.')

(file, filename, data) = imp.find_module(name, [path])
conv = imp.load_module(name, file, filename, data)

struct EdgeModel
    predictor
end

function prepare(model_path=Pkg.dir("HED","src","HED_pretrained_bsds.npy"); gpu=false)::EdgeModel
    println("model_path ", model_path)
    predictor =conv[:prepare](model_path)
    EdgeModel(predictor)
end

function getedge(m::EdgeModel, image_path, output=nothing)
    conv[:getedge](m.predictor, image_path, output)
end


struct InstrumentData <: JAXTAMData
    mission    :: Mission
    instrument :: Symbol
    obsid      :: String
    events     :: DataFrames.DataFrame
    gtis       :: DataFrames.DataFrame
    start      :: Number
    stop       :: Number
    src_ctrate :: Union{Missing, Number}
    bkg_ctrate :: Union{Missing, Number}
    header     :: DataFrames.DataFrame
end

"""
    _read_fits_hdu(fits_file::FITS, hdu_id::String; cols="auto")

Reads the HDU `hdu_id` from the loaded FITS file `fits_file`, returns the HDU data

Cannot read `BitArray` type columns due to `FITSIO` limitations
"""
function _read_fits_hdu(fits_file::FITS, hdu_id::String; cols="auto")
    # Columns read individually, instead of using the FITSIO function to read
    # all the header, as if header contains non-supported column then that function
    # just fails. This avoids that problem and throws warning for non-supported columns
    if cols == "auto"
        fits_cols_events = Array{String,1}

        fits_cols_events = FITSIO.colnames(fits_file[hdu_id])
    else
        fits_cols_events = cols
    end

    fits_hdu_data = DataFrame()

    for col in fits_cols_events
        try
            fits_col_data = read(fits_file[hdu_id], col)

            if ndims(fits_col_data) == 1
                fits_hdu_data[Symbol(col)] = fits_col_data
            end
        catch error
            @warn "$col not supported by FITSIO, skipped - $error"
        end
    end

    return fits_hdu_data
end

"""
    _read_fits_event(fits_path::String)

Reads the standard columns for timing analysis ("TIME", "PI", "GTI") from a FITS file,
returns `InstrumentData` type filled with the relevant data
"""
function _read_fits_event(mission::Mission, fits_path::String)
    @info "Loading $fits_path"
    fits_file   = FITS(fits_path)

    fits_header = read_header(fits_file["EVENTS"])

    instrument_name = Symbol(fits_header["INSTRUME"])
    fits_telescope  = lowercase(fits_header["TELESCOP"])
    @assert fits_telescope == _mission_name(mission) "Mission names do not match"
    fits_telescope  = _mission_symbol_to_type(Symbol(lowercase(fits_telescope)))

    fits_events_df = _read_fits_hdu(fits_file, "EVENTS"; cols=String["TIME", "PI"])

    fits_gtis_df = _read_fits_hdu(fits_file, "GTI")

    fits_obsid = fits_header["OBS_ID"]
    fits_start = fits_header["TSTART"]
    fits_stop  = fits_header["TSTOP"]

    fits_header_df = DataFrame()

    for (i, key) in enumerate(keys(fits_header))
        key = Symbol(replace(key, '-', '_')) # Colnames with `-` don't behave well with DataFrames/Feather
        if typeof(fits_header[i]) == Nothing || fits_header[i] == ""
            fits_header_df[Symbol(key)] =  "empty" # Have to write something, or Feahter.jl errors saving ""
        else
            fits_header_df[Symbol(key)] = fits_header[i]
        end
    end

    # Hacky fix to not create DF with col type of missing
    src_rt = Array{Union{Float64,Missing},1}(undef, 1)
    bkg_rt = Array{Union{Float64,Missing},1}(undef, 1)

    total_gti_time    = sum(fits_gtis_df[:STOP] .- fits_gtis_df[:START])
    total_event_count = size(fits_events_df, 1)

    src_rt[1] = total_event_count/total_gti_time
    bkg_rt[1] = missing # TODO: Add in support for background count rates

    fits_header_df[:SRC_RT] = src_rt
    fits_header_df[:BKG_RT] = bkg_rt
    fits_header_df[:TELESCOP] = lowercase.(fits_header_df[:TELESCOP])

    close(fits_file)
    @info "Finished FITS loading"

    return InstrumentData(fits_telescope, instrument_name, fits_obsid,
        fits_events_df, fits_gtis_df, fits_start, fits_stop,
        src_rt[1], bkg_rt[1],
        fits_header_df)
end

"""
    read_cl_fits(mission::Symbol, obs_row::DataFrames.DataFrame)

Reads in FITS data for an observation, returns a `Dict{Symbol,InstrumentData}`, with the
symbol as the instrument name. So `instrument_data[:XTI]` works for NICER, and either
`instrument_data[:FPMA]` or `instrument_data[:FPMB]` work for NuSTAR
"""
function read_cl_fits(mission::Mission, obs_row::DataFrames.DataFrameRow{DataFrames.DataFrame}, file_path::String)
    if isfile(file_path)
        @info "Found: $file_path"
    elseif isfile(file_path[1:end-3]) # Check for files ending in .evt.gz too
        file_path = file_path[1:end-1]
        @info "Found: $file_path"
    else
        throw(SystemError("attempted to open:\n\t$file_path\n\t$(file_path[1:end-3])\n", 2, nothing))
    end

    instrument_data = _read_fits_event(mission, file_path)
    return Dict(instrument_data.instrument => instrument_data)
end

"""
    _save_cl_feather(feather_dir::String, instrument::Union{String,Symbol},
        fits_events_df::DataFrame, fits_gtis_df::DataFrame, fits_meta_df::DataFrame)

Due to Feather file restrictions, cannot save all the event and GTI data in one, so
they are split up into three files: `events`, `gtis`, and `meta`. The `meta` file contains
just the mission name, obsid, and observation start and stop times
"""
function _save_cl_feather(mission::Mission, obs_row::DataFrames.DataFrameRow{DataFrames.DataFrame}, instrument::Union{String,Symbol},
        feather_dir::String, fits_events_df::DataFrames.DataFrame, fits_gtis_df::DataFrames.DataFrame, fits_meta_df::DataFrames.DataFrame;
        log=true)
    @info "Saving '$instrument' cl files to $feather_dir"
    mkpath(feather_dir)

    path_events = joinpath(feather_dir, "$(instrument)_events.feather")
    path_gtis   = joinpath(feather_dir, "$(instrument)_gtis.feather")
    path_meta   = joinpath(feather_dir, "$(instrument)_meta.feather")

    Feather.write(path_events, fits_events_df)
    Feather.write(path_gtis, fits_gtis_df)
    Feather.write(path_meta, fits_meta_df)

    if log
        _log_add(mission, obs_row,
            Dict("data" =>
                Dict(:feather_cl =>
                    Dict(instrument =>
                        Dict(
                            :path_events => path_events,
                            :path_gtis   => path_gtis,
                            :path_meta   => path_meta,
                        )
                    )
                )
            )
        )
    end

    @info "Saved $instrument feather file"

    return path_events, path_gtis, path_meta
end

function _read_cl_feather(path_events::String, path_gtis::String, path_meta::String)
    data_events = Feather.read(path_events)
    data_gtis   = Feather.read(path_gtis)
    data_meta   = Feather.read(path_meta)

    meta_missn = _mission_symbol_to_type(Symbol(data_meta[1, :TELESCOP]))
    meta_inst  = data_meta[1, :INSTRUME]
    meta_obsid = data_meta[1, :OBS_ID]
    meta_start = data_meta[1, :TSTART]
    meta_stop  = data_meta[1, :TSTOP]
    meta_srcrt = data_meta[1, :SRC_RT]
    meta_bkgrt = data_meta[1, :BKG_RT]

    return InstrumentData(meta_missn, meta_inst, meta_obsid,
        data_events, data_gtis, meta_start, meta_stop,
        meta_srcrt, meta_bkgrt, data_meta)
end

"""
    read_cl(mission::Symbol, obs_row::DataFrames.DataFrameRow{DataFrames.DataFrame}; overwrite=false)

Attempts to read saved (feather) data, if none is found then the `read_cl_fits` function is ran
and the data is saved with `_save_cl_feather` for future use
"""
function read_cl(mission::Mission, obs_row::DataFrames.DataFrameRow{DataFrames.DataFrame}; overwrite=false)
    obsid       = obs_row[:obsid]
    instruments = _mission_instruments(mission)

    cl_files_feather = _log_query(mission, obs_row, "data", :feather_cl)
    cl_files_raw     = _obs_files_cl(mission, obs_row)

    total_src_ctrate = 0.0
    instrument_data = Dict{Symbol,JAXTAM.InstrumentData}()
    for instrument in instruments
        if ismissing(cl_files_feather) || !haskey(cl_files_feather, instrument) || overwrite
            file_path = cl_files_raw[instrument]

            @info "Missing feather_cl for $instrument"
            @info "Processing $instrument cl fits"

            current_instrument = read_cl_fits(mission, obs_row, file_path)[instrument]

            if size(current_instrument.events, 1) == 0
                @warn "No events found in $instrument observation $obsid"
                throw(JAXTAMError("Unable to construct InstrumentData from empty DataFrame. Observation likely has no events", :read_cl))
            else
                cl_path = abspath(_obs_path_local(mission, obs_row; kind=:jaxtam), "JAXTAM", _log_query_path(; category=:data, kind=:feather_cl))
                path_events, path_gtis, path_meta = _save_cl_feather(mission, obs_row, instrument, cl_path,
                    current_instrument.events, current_instrument.gtis, current_instrument.header)
            end

            total_src_ctrate += current_instrument.header[1, :SRC_RT]
            instrument_data[instrument] = current_instrument
        else
            @info "Loading feather cl files for '$instrument'"
            feather_paths = cl_files_feather[instrument]
            instrument_data[instrument] = _read_cl_feather(feather_paths[:path_events], feather_paths[:path_gtis], feather_paths[:path_meta])
        end
    end

    if ismissing(_log_query(mission, obs_row, "meta", :countrates, :raw))
        total_src_ctrate = total_src_ctrate/length(instruments)
        _log_add(mission, obs_row, Dict("meta" => Dict(:countrates => Dict(:raw => total_src_ctrate))))
    end

    return instrument_data
end

function read_cl(mission::Mission, master_df::DataFrames.DataFrame, obsid::String; overwrite=false)
    obs_row = master_query(master_df, :obsid, obsid)

    return read_cl(mission, obs_row; overwrite=overwrite)
end

function read_cl(mission::Mission, obsid::String; overwrite=false)
    return read_cl(mission, master(mission), obsid; overwrite=overwrite)
end


######### Freeze factor node ##########
function freezeFactor!(gbp::ContinuousModel; factor = 0::Int64)
    errorNodeIndex(factor, gbp.graph.Nfactor; name = "factor")

    factorLocal = gbp.graph.dynamic[factor]
    errorFactorLocal(factorLocal; state = "frozen")

    whereIs = 0
    @inbounds for i = 1:length(gbp.graph.iterateFactor)
        if gbp.graph.iterateFactor[i] == factorLocal
            whereIs = i
            break
        end
    end
    errorWhereIs(whereIs; name = "factor node", state = "frozen")

    deleteat!(gbp.graph.iterateFactor, whereIs)
end

######### Defreeze factor node ##########
function defreezeFactor!(gbp::ContinuousModel; factor = 0::Int64)
    errorNodeIndex(factor, gbp.graph.Nfactor; name = "factor")

    factorLocal = gbp.graph.dynamic[factor]
    errorFactorLocal(factorLocal; state = "defrozen")

    whereIs = 0
    @inbounds for i = 1:length(gbp.graph.iterateFactor)
        if gbp.graph.iterateFactor[i] > factorLocal
            whereIs = i
            break
        end
    end

    if whereIs != 0
        errorInside(whereIs, gbp.graph.iterateFactor, factorLocal; name = "factor node", state = "defrozen")
        insert!(gbp.graph.iterateFactor, whereIs, factorLocal)
    else
        errorLast(gbp.graph.iterateFactor, factorLocal; name = " factor node", state = "defrozen")
        push!(gbp.graph.iterateFactor, factorLocal)
    end
end

######### Freeze variable node ##########
function freezeVariable!(gbp::ContinuousModel; variable = 0::Int64)
    errorNodeIndex(variable, gbp.graph.Nvariable; name = "variable")

    whereIs = 0
    @inbounds for i = 1:length(gbp.graph.iterateVariable)
        if gbp.graph.iterateVariable[i] == variable
            whereIs = i
            break
        end
    end
    errorWhereIs(whereIs; name = "variable node", state = "frozen")

    deleteat!(gbp.graph.iterateVariable, whereIs)
end

########## Defreeze variable node ##########
function defreezeVariable!(gbp::ContinuousModel; variable = 0::Int64)
    errorNodeIndex(variable, gbp.graph.Nvariable; name = "variable")

    whereIs = 0
    @inbounds for i = 1:length(gbp.graph.iterateVariable)
        if gbp.graph.iterateVariable[i] > variable
            whereIs = i
            break
        end
    end

    if whereIs != 0
        errorInside(whereIs, gbp.graph.iterateVariable, variable; name = "variable node", state = "defrozen")
        insert!(gbp.graph.iterateVariable, whereIs, variable)
    else
        errorLast(gbp.graph.iterateVariable, variable; name = "variable node", state = "defrozen")
        push!(gbp.graph.iterateVariable, variable)
    end
end

######### Freeze Egdge: From variable to factor node ##########
function freezeVariableFactor!(gbp::ContinuousModel; variable = 0::Int64, factor = 0::Int64)
    errorNodeIndex(factor, gbp.graph.Nfactor; name = "factor")
    errorNodeIndex(variable, gbp.graph.Nvariable; name = "variable")
    errorFactorLocal(gbp.graph.dynamic[factor]; state = "frozen")
    errorEdge(gbp.system.jacobian, factor, variable)

    whereIs = 0
    @inbounds for (k, i) in enumerate(gbp.graph.colptr[variable])
        if gbp.inference.fromFactor[i] == factor
            whereIs = k
        end
    end
    errorWhereIs(whereIs; name = "edge", state = "frozen")

    deleteat!(gbp.graph.colptr[variable], whereIs)
    gbp.graph.toFactor[factor, variable] = 0
end

######### Defreeze Egdge: From variable to factor node ##########
function defreezeVariableFactor!(gbp::ContinuousModel; variable = 0::Int64, factor = 0::Int64)
    errorNodeIndex(factor, gbp.graph.Nfactor; name = "factor")
    errorNodeIndex(variable, gbp.graph.Nvariable; name = "variable")
    errorFactorLocal(gbp.graph.dynamic[factor]; state = "frozen")
    errorEdge(gbp.system.jacobian, factor, variable)

    @inbounds for i = 1:gbp.graph.Nlink
        if gbp.inference.fromFactor[i] == factor && gbp.inference.toVariable[i] == variable
            for j in gbp.graph.colptr[variable]
                if j == i
                    error("The edge is already defrozen.")
                end
            end
            push!(gbp.graph.colptr[variable], i)
            break
        end
    end

    @inbounds for i = 1:gbp.graph.Nlink
        if gbp.inference.fromVariable[i] == variable && gbp.inference.toFactor[i] == factor
            gbp.graph.toFactor[factor, variable] = i
            break
        end
    end
    sort!(gbp.graph.colptr[variable])
end

######### Freeze Egdge: From factor to variable node ##########
function freezeFactorVariable!(gbp::ContinuousModel; factor = 0::Int64, variable = 0::Int64)
    errorNodeIndex(factor, gbp.graph.Nfactor; name = "factor")
    errorNodeIndex(variable, gbp.graph.Nvariable; name = "variable")

    factorLocal = gbp.graph.dynamic[factor]
    errorFactorLocal(factorLocal; state = "frozen")
    errorEdge(gbp.system.jacobian, factor, variable)

    whereIs = 0
    @inbounds for (k, i) in enumerate(gbp.graph.rowptr[factorLocal])
        if gbp.inference.fromVariable[i] == variable
            whereIs = k
        end
    end
    errorWhereIs(whereIs; name = "edge", state = "frozen")

    deleteat!(gbp.graph.rowptr[factorLocal], whereIs)
    gbp.graph.toVariable[variable, factor] = 0
end

######### Defreeze Egdge: From factor to variable node ##########
function defreezeFactorVariable!(gbp::ContinuousModel; factor = 0::Int64, variable = 0::Int64)
    errorNodeIndex(factor, gbp.graph.Nfactor; name = "factor")
    errorNodeIndex(variable, gbp.graph.Nvariable; name = "variable")

    factorLocal = gbp.graph.dynamic[factor]
    errorFactorLocal(factorLocal; state = "frozen")
    errorEdge(gbp.system.jacobian, factor, variable)

    @inbounds for i = 1:gbp.graph.Nlink
        if gbp.inference.fromVariable[i] == variable && gbp.inference.toFactor[i] == factor
            for j in gbp.graph.rowptr[factorLocal]
                if j == i
                    error("The edge is already defrozen.")
                end
            end
            push!(gbp.graph.rowptr[factorLocal], i)
            break
        end
    end

    @inbounds for i = 1:gbp.graph.Nlink
        if gbp.inference.fromFactor[i] == factor && gbp.inference.toVariable[i] == variable
            gbp.graph.toVariable[variable, factor] = i
            break
        end
    end
    sort!(gbp.graph.rowptr[factorLocal])
end

######### Hide factor node ##########
function hideFactor!(gbp::ContinuousModel; factor = 0::Int64)
    errorNodeIndex(factor, gbp.graph.Nfactor; name = "factor")

    factorLocal = gbp.graph.dynamic[factor]
    if factorLocal != 0
        whereIs = 0
        @inbounds for i = 1:length(gbp.graph.iterateFactor)
            if gbp.graph.iterateFactor[i] == factorLocal
                whereIs = i
                break
            end
        end
        errorWhereIs(whereIs; name = "factor node", state = "hidden")
        deleteat!(gbp.graph.iterateFactor, whereIs)

        deleteLinksRow = gbp.graph.rowptr[factorLocal]
        variables = gbp.inference.fromVariable[deleteLinksRow]
        deleteLinksCol = Int[]
        @inbounds for i in variables
            links = gbp.graph.colptr[i]
            for (k, j) in enumerate(links)
                if gbp.inference.fromFactor[j] == factor
                    push!(deleteLinksCol, j)
                    deleteat!(gbp.graph.colptr[i], k)
                    deleteat!(gbp.graph.colptrMarginal[i], k)
                end
            end
        end

        @inbounds for i in deleteLinksRow
            gbp.inference.meanVariableFactor[i] = 0.0
            gbp.inference.varianceVariableFactor[i] = 0.0
        end
        @inbounds for i in deleteLinksCol
            gbp.inference.meanFactorVariable[i] = 0.0
            gbp.inference.varianceFactorVariable[i] = 0.0
        end
    else
        idx = gbp.system.jacobianTranspose.colptr[factor]
        variable = gbp.system.jacobianTranspose.rowval[idx]

        gbp.graph.meanDirect[variable] -= gbp.system.observation[factor] * gbp.system.jacobianTranspose[variable, factor] / gbp.system.variance[factor]
        gbp.graph.weightDirect[variable] -= gbp.system.jacobianTranspose[variable, factor]^2 / gbp.system.variance[factor]
        if gbp.graph.weightDirect[variable] == 0.0
            gbp.graph.weightDirect[variable] = 1 /  gbp.graph.virtualVariance
            gbp.graph.meanDirect[variable] =  gbp.graph.virtualMean /  gbp.graph.virtualVariance
        end
    end
    gbp.system.observation[factor] = 0.0
    @inbounds for i in gbp.system.jacobianTranspose.colptr[factor]:(gbp.system.jacobianTranspose.colptr[factor + 1] - 1)
        col = gbp.system.jacobianTranspose.rowval[i]
        gbp.system.jacobianTranspose.nzval[i] = 0.0
        gbp.system.jacobian[factor, col] = 0.0
        gbp.graph.toFactor[factor, col] = 0
    end
end

######### Add factor nodes ##########
function addFactors!(gbp::ContinuousModel; mean = 0.0, variance = 0.0, jacobian = [])
    ### Update system model
    append!(gbp.system.observation, mean)
    append!(gbp.system.variance, variance)
    newFactors = sparse(jacobian)
    newFactorsTranspose = copy(transpose(newFactors))
    gbp.system.jacobian = [gbp.system.jacobian; newFactors]
    gbp.system.jacobianTranspose = copy(transpose(gbp.system.jacobian))

    ### Add singly connected factor nodes and update dynamic vector
    Nfactor = length(mean)
    Nlink = 0; Nindirect = 0
    append!(gbp.graph.dynamic, fill(0, Nfactor))
    @inbounds for i = 1:Nfactor
        NvariableInRow = newFactorsTranspose.colptr[i + 1] - newFactorsTranspose.colptr[i]
        if NvariableInRow == 1
            variable = newFactorsTranspose.rowval[newFactorsTranspose.colptr[i]]
            if gbp.graph.weightDirect[variable] == 1 / gbp.graph.virtualVariance &&  gbp.graph.meanDirect[variable] == gbp.graph.virtualMean /  gbp.graph.virtualVariance
                gbp.graph.meanDirect[variable] = 0.0
                gbp.graph.weightDirect[variable] = 0.0
            end
            gbp.graph.meanDirect[variable] += mean[i] * newFactors[i, variable] / variance[i]
            gbp.graph.weightDirect[variable] += newFactors[i, variable]^2 / variance[i]
        else
            Nlink += NvariableInRow
            Nindirect += 1
            gbp.graph.dynamic[i + gbp.graph.Nfactor] = Nindirect + gbp.graph.Nindirect
        end
    end

    ### Update vectors related with indirect factor nodes
    toFactorLocal = fill(0, Nlink); temp = fill(0.0, Nlink)
    append!(gbp.graph.rowptr, [Int[] for i = 1:Nindirect])
    append!(gbp.graph.coefficient, temp)
    append!(gbp.graph.meanIndirect, temp)
    append!(gbp.graph.varianceIndirect, temp)
    append!(gbp.inference.toFactor, toFactorLocal)
    append!(gbp.inference.fromVariable, toFactorLocal)
    append!(gbp.inference.meanVariableFactor, temp)
    append!(gbp.inference.varianceVariableFactor, temp)
    append!(gbp.graph.alphaOld, temp)
    append!(gbp.graph.alphaNew, fill(1.0, Nlink))

    varianceInitial = fill(0.0, gbp.graph.Nvariable)
    meanInitial = fill(0.0, gbp.graph.Nvariable)
    idxr = gbp.graph.Nindirect + 1
    idxi = gbp.graph.Nlink + 1
    prev = 1
    @inbounds for col = 1:Nfactor
        if (newFactorsTranspose.colptr[col + 1] - newFactorsTranspose.colptr[col]) != 1
            for i = newFactorsTranspose.colptr[col]:(newFactorsTranspose.colptr[col + 1] - 1)
                row = newFactorsTranspose.rowval[i]

                gbp.graph.coefficient[idxi] = newFactorsTranspose[row, col]
                gbp.graph.meanIndirect[idxi] = mean[col]
                gbp.graph.varianceIndirect[idxi] = variance[col]
                gbp.inference.toFactor[idxi] = col + gbp.graph.Nfactor
                gbp.inference.fromVariable[idxi] = row

                toFactorLocal[idxi - gbp.graph.Nlink] = col
                push!(gbp.graph.rowptr[idxr], idxi)

                if meanInitial[row] == 0 && varianceInitial[row] == 0
                    Mcol = gbp.graph.meanDirect[row]; Wcol = gbp.graph.weightDirect[row]
                    for j in gbp.graph.colptrMarginal[row]
                        Mcol += gbp.inference.meanFactorVariable[j] / gbp.inference.varianceFactorVariable[j]
                        Wcol += 1 / gbp.inference.varianceFactorVariable[j]
                    end
                    varianceInitial[row] = 1 / Wcol
                    meanInitial[row] = Mcol * varianceInitial[row]
                end

                idxi += 1
                rown = newFactorsTranspose.rowval[newFactorsTranspose.colptr[col + 1] - 1]
                if CartesianIndex(rown, col) == CartesianIndex(row, col)
                    idxr += 1
                end
            end
        end
    end

    @inbounds for i = gbp.graph.Nlink + 1:gbp.graph.Nlink + Nlink
        vari = gbp.inference.fromVariable[i]
        gbp.inference.meanVariableFactor[i] = meanInitial[vari]
        gbp.inference.varianceVariableFactor[i] = varianceInitial[vari]
    end

    start = gbp.graph.Nlink + 1
    last = length(gbp.inference.toFactor)
    links = collect(start:last)
    sendToFactor = sparse(toFactorLocal, gbp.inference.fromVariable[start:last], links, Nfactor, gbp.graph.Nvariable)
    gbp.graph.toFactor = [gbp.graph.toFactor; sendToFactor]

    ### Update vectors related with variable nodes
    append!(gbp.inference.toVariable, toFactorLocal)
    append!(gbp.inference.fromFactor, toFactorLocal)
    gbp.graph.toVariable = [gbp.graph.toVariable transpose(sendToFactor)]
    idxi = 1; prev = 1
    @inbounds for col = 1:gbp.graph.Nvariable
        for i = gbp.system.jacobian.colptr[col]:(gbp.system.jacobian.colptr[col + 1] - 1)
            row = gbp.system.jacobian.rowval[i]
            if gbp.system.jacobianTranspose.colptr[row + 1] - gbp.system.jacobianTranspose.colptr[row] != 1
                if prev == col
                    gbp.graph.colptr[col] = Int64[]
                    gbp.graph.colptrMarginal[col] = Int64[]
                end
                prev = col + 1

                if gbp.graph.toFactor[row, col] != 0
                    push!(gbp.graph.colptr[col], idxi)
                    push!(gbp.graph.colptrMarginal[col], idxi)
                end
                gbp.inference.toVariable[idxi] = col
                gbp.inference.fromFactor[idxi] = row

                if gbp.graph.toVariable[col, row] != 0
                    gbp.graph.toVariable[col, row] = idxi
                end
                idxi += 1
            end
        end
    end
    append!(gbp.inference.meanFactorVariable, temp)
    append!(gbp.inference.varianceFactorVariable, temp)
    append!(gbp.graph.iterateFactor, collect(gbp.graph.Nindirect + 1:gbp.graph.Nindirect + Nindirect))

    ### Update factor graf statistics
    gbp.graph.Nfactor += Nfactor
    gbp.graph.Nindirect += Nindirect
    gbp.graph.Nlink += Nlink
end

######### Error node ##########
@inline function errorNodeIndex(node, limit; name = "")
    if node < 0 || node > limit
        error("The $name node does not exist.")
    end
end

######### Error where is ##########
@inline function errorWhereIs(whereIs; name = "", state = "")
    if whereIs == 0
        error("The $name does not exist or it is already $state.")
    end
end

######### Error local factor ##########
@inline function errorFactorLocal(factorLocal; state = "")
    if factorLocal == 0
        error("The singly connected factor node cannot be $state.")
    end
end

######### Error inside ##########
@inline function errorInside(whereIs, iterate, node; name = "", state = "")
    if whereIs != 1 && iterate[whereIs - 1] == node || iterate[1] == node
        error("The $name is already $state.")
    end
end

######### Error last ##########
@inline function errorLast(iterate, node; name = "", state = "")
    if iterate[end] == node
        error("The $name is already $state.")
    end
end

######### Error egde ##########
@inline function errorEdge(jacobian, factor, variable)
    if jacobian[factor, variable] == 0.0
        error("The edge does not exist.")
    end
end


using CSV
using DataFrames
using Colors
using PyPlot

const rexeltype = r"N0f8|N0f16|Float32"
const idxeltype = Dict("N0f8"=>1, "N0f16"=>2, "Float32"=>3)
const eltypeidx = Dict(v=>k for (k, v) in idxeltype)
const rexsize = r"_(\d+)$"
const special_replacements = ["components" => "label_components", "dblcone" => "distance_transform", "spiral" => "flood"]

randomtag(fn) = (occursin(r"rgb", fn), idxeltype[match(rexeltype, fn).match], parse(Int, only(match(rexsize, fn).captures)))
randomlabel(tag) = (tag[1] ? "RGB" : "Gray") * "{" * eltypeidx[tag[2]] * "}, $(tag[3])"

function imagelt(image1, image2)
    image1 ∈ ("2d", "3d") && image2 ∈ ("2d", "3d") && return image1 < image2

    getcolor(str) = match(r"(Gray|RGB)", str).captures[1]::AbstractString
    geteltype(str) = idxeltype[match(r"(N0f8|N0f16|Float32)", str).captures[1]]
    getsz(str) = parse(Int, match(r", (\d)+", str).captures[1])

    col1, col2 = getcolor(image1), getcolor(image2)
    col1 != col2 && return col1 < col2
    et1, et2 = geteltype(image1), geteltype(image2)
    et1 != et2 && return et1 < et2
    return getsz(image1) < getsz(image2)
end

hexhash(c) = "#"*hex(c)
const language_colors = (cols = hexhash.(distinguishable_colors(4, [colorant"white", colorant"black"]; dropseed=true));
                         Dict("python-skimage"=>cols[1], "Matlab"=>cols[2], "OpenCV"=>cols[3], "Julia"=>cols[4]))

function taskplots(jobtag::Function, jobname::Function, filelabel::Pair...; colordict=language_colors, tasknames=nothing, pervoxel::Bool=true)
    # Read all the data
    benchdata = Dict{String,GroupedDataFrame{DataFrame}}()
    for (file, label) in filelabel
        benchdata[label] = groupby(DataFrame(CSV.File(file; delim=',')), "Benchmark")
    end
    # Ensure all inputs have the same tasks
    ks = keys(first(benchdata).second)
    knames = Set(only(Tuple(k)) for k in ks)
    for (_, df) in benchdata
        dfknames = Set(only(Tuple(k)) for k in keys(df))
        # kdf ⊆ ks && ks ⊆ kdf || error("disjoint tasks")
        dfknames == knames || error("disjoint tasks")
    end
    # Determine the plotting order
    if tasknames === nothing
        tasknames = sort(collect(knames))
    end
    # Assign colors to each suite
    if colordict === nothing
        colordict = Dict(zip(keys(benchdata), hexhash.(distinguishable_colors(length(benchdata), [colorant"white"]; dropseed=true))))
    end
    ncols = floor(Int, sqrt(length(tasknames)))
    nrows = ceil(Int, length(tasknames)/ncols)
    fig, axes = plt.subplots(nrows=nrows, ncols=ncols, gridspec_kw=Dict("bottom"=>0.2))
    for i = length(tasknames)+1:length(axes)
        axes[i].set_axis_off()  # close()
    end
    legend_labels, hlines = String[], []
    for (axidx, ax, k) in zip(CartesianIndices(axes), axes, tasknames)
        xlbls = nothing
        for (label, data) in benchdata
            data = sort(data[(Benchmark=k,)], "File"; by=jobtag)
            if xlbls === nothing
                xlbls = jobname.(jobtag.(data."File"))
            end
            y = data."Time(s)"
            if pervoxel
                nvox = map(data."File") do fn
                    jobtag(fn)[end]^2
                end
                y ./= nvox
            end
            hline = ax.plot(y; color=colordict[label], marker="x")
            if length(legend_labels) < length(filelabel)
                push!(legend_labels, label)
                push!(hlines, first(hline))
            end
        end
        if axidx[1] == size(axes, 1) || LinearIndices(axes)[axidx] == length(tasknames)
            ax.set_xticks(0:length(xlbls)-1)
            ax.set_xticklabels(xlbls; rotation=90)
            ax.set_xlabel("Jobs")
        else
            ax.set_xticks([])
        end
        ax.set_yscale("log")
        if axidx[2] == 1
            ax.set_ylabel("Time (s)")
        end
        ax.set_title(k)
    end
    fig.legend(hlines, legend_labels)
    return fig
end

# taskplots(randomtag, randomlabel, "julia/julia_generics.csv" => "Julia",
#                                   "python-skimage/python_generics.csv" => "python-skimage",
#                                   "matlab/matlab_generics.csv" => "Matlab",
#                                   "opencv/opencv_generics.csv" => "OpenCV",
# )

function specialplots(filelabel::Pair...; colordict=language_colors)
    # Read all the data
    benchdata = Dict{String,DataFrame}()
    for (file, label) in filelabel
        benchdata[label] = DataFrame(CSV.File(file; delim=','))
    end
    # Ensure all inputs have the same tasks
    tasknames = sort(first(benchdata).second."Benchmark")
    all(df -> Set(df."Benchmark") == Set(tasknames), values(benchdata))
    # Assign colors to each suite
    if colordict === nothing
        colordict = Dict(zip(keys(benchdata), hexhash.(distinguishable_colors(length(benchdata), [colorant"white"]; dropseed=true))))
    end
    fig, ax = plt.subplots()
    x = 1:length(tasknames)
    w = 1/(2*length(benchdata))
    for (i, (label, df)) in enumerate(benchdata)
        p = Int.(indexin(df."Benchmark", tasknames))
        dfs = df[invperm(p),:]
        display(dfs)
        ax.bar(x .+ (i-1)*w .- w/2, dfs."Time(s)", w; label=label, color=colordict[label])
    end
    ax.set_xticks(x)
    repnames = map(tasknames) do tn
        foldl((str, rep) -> replace(str, rep), special_replacements; init=tn)
    end
    ax.set_xticklabels(repnames; rotation=90)
    ax.set_yscale("log")
    ax.set_ylabel("Time (s)")
    fig.legend()
    fig.tight_layout()
    return fig
end

# specialplots("julia/julia_special.csv" => "Julia",
#              "python-skimage/python_special.csv" => "python-skimage",
#              "matlab/matlab_special.csv" => "Matlab",
#              #= "opencv/opencv_special.csv" => "OpenCV", =#
# )

function combine_all(prs...; szs=(r"64_64", r"2048_2048"), eltypes=(r"gray_N0f8", r"gray_Float32", r"rgb_N0f8", r"rgb_Float32"))
    df = DataFrame("Framework"=>String[], "Operation"=>String[], "Image"=>String[], "Time"=>Float64[])
    for (file, framework) in prs
        for row in CSV.Rows(file)
            if length(row) == 3  # one of the generics
                any(occursin(row.File), szs) || continue
                any(occursin(row.File), eltypes) || continue
                tag = randomtag(row.File)
                t = getproperty(row, Symbol("Time(s)"))
                push!(df, (framework, row.Benchmark, randomlabel(tag), parse(Float64, t)))
            else  # one of the special
                imagetag = occursin("3d", row.Benchmark) ? "3d" : "2d"
                optag = foldl((str, rep) -> replace(str, rep), special_replacements; init=row.Benchmark[1:end-2])
                t = getproperty(row, Symbol("Time(s)"))
                push!(df, (framework, optag, imagetag, parse(Float64, t)))
            end
        end
    end
    return df
end

# df = combine_all("julia/julia_generics.csv" => "Julia",
#                  "python-skimage/python_generics.csv" => "python-skimage",
#                  "matlab/matlab_generics.csv" => "Matlab",
#                  "opencv/opencv_generics.csv" => "OpenCV",
#                  "julia/julia_special.csv" => "Julia",
#                  "python-skimage/python_special.csv" => "python-skimage",
#                  "matlab/matlab_special.csv" => "Matlab")

function ratio_df(df; operation_order=["complement", "mean", "histeq", "blur", "gradient", "label_components", "flood", "distance_transform"])
    jdata = filter(:Framework=>isequal("Julia"), df)
    jgroups = groupby(jdata, :Operation)
    rdf = DataFrame("Framework"=>String[], "Operation"=>String[], "Image"=>String[], "Ratio"=>Float64[])
    for f in ("python-skimage", "Matlab", "OpenCV")
        fdata = filter(:Framework=>isequal(f), df)
        fgroups = groupby(fdata, :Operation)
        for op in operation_order
            k = (Operation=op,)
            jref = sort(jgroups[k], :Image; lt=imagelt)
            try
                fval = sort(fgroups[k], :Image; lt=imagelt)
                for (jrow, frow) in zip(eachrow(jref), eachrow(fval))
                    @assert jrow.Image == frow.Image
                    push!(rdf, [f, op, frow.Image, frow.Time / jrow.Time])
                end
            catch
                for jrow in eachrow(jref)
                    push!(rdf, [f, op, jrow.Image, NaN])
                end
            end
        end
    end
    return rdf
end

function plot_ratio(rdf; colordict=language_colors)
    frameworks = unique(rdf.Framework)
    operations = unique(rdf.Operation)
    n = Int(nrow(rdf)/length(frameworks))  # deliberately errors if not an even divisor
    x = 1:n+length(operations)-1
    xlabels = String[]
    opgaps = Int[]
    gdf = groupby(rdf, :Framework)
    fig, axouter = plt.subplots(figsize=(7.5,4), subplot_kw=Dict("position" => [0.1, 0.05, 0.9, 0.95]))
    insetgap = 0.4
    ax = axouter.inset_axes([0, insetgap, 1, 1-insetgap])
    for k in keys(gdf)
        g = gdf[k]
        col = colordict[k.Framework]
        oldop = ""
        y = Float64[]
        for (i, (op, t)) in enumerate(zip(g.Operation, g.Ratio))
            if op != oldop && i != 1
                push!(y, NaN)
                if length(xlabels) < n && i != 1
                    # push!(xlabels, "")
                    push!(opgaps, length(y))
                end
            end
            oldop = op
            push!(y, t)
            if length(xlabels) < n
                push!(xlabels, g.Image[i])
            end
        end

        ax.plot(x, y, color=col, label=k.Framework, marker="x")
    end
    ax.plot(x, fill(1, length(x)), "k--", label=nothing)
    ax.set_xticks(setdiff(x, opgaps))
    ax.set_xticklabels(xlabels; rotation=90, fontsize=6)
    ax.set_yscale(:log)
    ax.set_ylim((1/3000, 3000))
    ax.set_yticks([1/1000, 1/100, 1/10, 1, 10, 100, 1000])
    ax.set_ylabel("Time relative to Julia")
    ax.legend(loc="lower left", fontsize=8)
    pushfirst!(opgaps, 0)
    push!(opgaps, length(x)+1)
    xc = (opgaps[1:end-1] .+ opgaps[2:end])./2
    axouter.set_xlim(ax.get_xlim())
    axouter.set_xticks(xc)
    axouter.set_xticklabels(operations; rotation=90, fontsize=9)
    for obj in axouter.get_xticklines()
        obj.set_visible(false)
    end
    axouter.get_yaxis().set_visible(false)
    axouter.set_frame_on(false)
    # for spkey in axouter.spines
    #     axouter.spines[spkey].set_visible(false)
    # end
    # for sp in ("left", "bottom", "right", "top")
    #     axouter.spines[sp].set_visible(false)
    # end
    fig.tight_layout()
    return fig
end


module Altro

import TrajectoryOptimization
import RobotDynamics
using StaticArrays
using BenchmarkTools
using Interpolations
using SolverLogging
using Crayons

using SparseArrays
using LinearAlgebra
using Logging
using Statistics
using TimerOutputs
using ForwardDiff
using FiniteDiff
import Octavian

const TO = TrajectoryOptimization
const RD = RobotDynamics

using TrajectoryOptimization:
    num_constraints, get_trajectory

import TrajectoryOptimization: rollout!, get_constraints, get_model, get_objective,
    evaluate_constraints!, constraint_jacobians!
import RobotDynamics: discrete_dynamics, dynamics, dynamics!, evaluate, evaluate!

using TrajectoryOptimization:
    Problem,
    ConstraintList,
    AbstractObjective, Objective, #QuadraticObjective,
    # SampledTrajectory,
    # DynamicsExpansion, # TODO: Move to ALTRO
    # ALConstraintSet,
    # DynamicsConstraint,
    states, controls,
    Equality, Inequality, SecondOrderCone,
    cost

using RobotDynamics:
    AbstractModel, DiscreteDynamics, DiscreteLieDynamics,
    QuadratureRule, Implicit, Explicit,
    FunctionSignature, InPlace, StaticReturn,
    DiffMethod, ForwardAD, FiniteDifference, UserDefined,
    AbstractKnotPoint, KnotPoint, StaticKnotPoint,
    state_dim, control_dim, output_dim, dims,
    state, control, SampledTrajectory

# types
export
    ALTROSolverOld,
    ALTROSolver,
    # iLQRSolverOld,
    # AugmentedLagrangianSolver,
    SolverStats,
    SolverOptions

export
    solve!,
    benchmark_solve!,
    iterations,
    set_options!,
    max_violation,
    status

# modules
export
    Problems

const ColonSlice = Base.Slice{Base.OneTo{Int}}
const SparseView{T,I} = SubArray{T, 2, SparseMatrixCSC{T, I}, Tuple{UnitRange{I}, UnitRange{I}}, false}
const VectorView{T,I} = SubArray{T, 1, Vector{T}, Tuple{UnitRange{I}}, true}

# Select the matix multiplication kernel
const USE_OCTAVIAN = parse(Bool, get(ENV, "ALTRO_USE_OCTAVIAN", "true"))
@static if USE_OCTAVIAN
    const matmul! = Octavian.matmul!
else
    const matmul! = mul!
end

# Include the QDLDL wrapper
include("qdldl.jl")
using .Cqdldl

# High level structs
include("utils.jl")
include("infeasible_model.jl")
include("solvers.jl")
include("solver_opts.jl")

# iLQR Solver
include("ilqr/cost_expansion.jl")
include("ilqr/dynamics_expansion.jl")
include("ilqr/ilqr_solver.jl")
include("ilqr/backwardpass.jl")
include("ilqr/forwardpass.jl")
include("ilqr/ilqr_solve.jl")

# Augmented Lagrangian Solver
include("augmented_lagrangian/alcon.jl")
include("augmented_lagrangian/alconset.jl")
include("augmented_lagrangian/al_objective.jl")
include("augmented_lagrangian/al_solver.jl")
include("augmented_lagrangian/al_solve.jl")
include("direct/sparseblocks.jl")

# Projected Newton Solver
include("direct/pncon.jl")
include("direct/pnconset.jl")
include("direct/pn_solver.jl")
include("direct/pn_solve.jl")

# ALTRO Solver
include("altro/altro_solver.jl")
include("altro/altro_solve.jl")

# Example problems submodule
include("problems.jl")

end # module


"""
This struct holds the market price data for each product.
    energy_price: Dictionary of zonal energy prices with hourly resolution for each scenario.
    reserveup_price: Dictionary of zonal reserve up prices with hourly resolution for each scenario.
    reservedown_price: Dictionary of zonal reserve down prices with hourly resolution for each scenario.
    capacity_price: Dictionary of annual capacity prices for each scenario.
    rec_price: Dictionary of annual REC prices for each scenario.
"""

mutable struct MarketPrices
    energy_price::Union{Nothing, Dict{String, AxisArrays.AxisArray{Float64, 3}}}
    reserveup_price::Union{Nothing, Dict{String, AxisArrays.AxisArray{Float64, 3}}}
    reservedown_price::Union{Nothing, Dict{String, AxisArrays.AxisArray{Float64, 3}}}
    capacity_price::Union{Nothing, Dict{String, AxisArrays.AxisArray{Float64, 1}}}
    rec_price::Union{Nothing, Dict{String, AxisArrays.AxisArray{Float64, 1}}}
end

function MarketPrices()
    return MarketPrices(nothing, nothing, nothing, nothing, nothing)
end

get_energy_price(prices::MarketPrices) = prices.energy_price
get_reserveup_price(prices::MarketPrices) = prices.reserveup_price
get_reservedown_price(prices::MarketPrices) = prices.reservedown_price
get_capacity_price(prices::MarketPrices) = prices.capacity_price
get_rec_price(prices::MarketPrices) = prices.rec_price

get_prices(prices::MarketPrices, prod::T) where T<: Product = nothing
get_prices(prices::MarketPrices, prod::Energy) = get_energy_price(prices)
get_prices(prices::MarketPrices, prod::OperatingReserve{ReserveUpEMIS}) = get_reserveup_price(prices)
get_prices(prices::MarketPrices, prod::OperatingReserve{ReserveDownEMIS}) = get_reservedown_price(prices)
get_prices(prices::MarketPrices, prod::Capacity) = get_capacity_price(prices)
get_prices(prices::MarketPrices, prod::REC) = get_rec_price(prices)

function set_energy_price!(prices::MarketPrices, scenario_name::String, energy_price::AxisArrays.AxisArray{Float64, 3})
    if !isnothing(prices.energy_price)
        prices.energy_price[scenario_name] = energy_price
    else
        prices.energy_price = Dict(scenario_name => energy_price)
    end
    return
end

function set_reserveup_price!(prices::MarketPrices, scenario_name::String, reserveup_price::AxisArrays.AxisArray{Float64, 3})
    if !isnothing(prices.reserveup_price)
        prices.reserveup_price[scenario_name] = reserveup_price
    else
        prices.reserveup_price = Dict(scenario_name => reserveup_price)
    end
    return
end

function set_reservedown_price!(prices::MarketPrices, scenario_name::String, reservedown_price::AxisArrays.AxisArray{Float64, 3})
    if !isnothing(prices.reservedown_price)
        prices.reservedown_price[scenario_name] = reservedown_price
    else
        prices.reservedown_price = Dict(scenario_name => reservedown_price)
    end
    return
end

function set_capacity_price!(prices::MarketPrices, scenario_name::String, capacity_price::AxisArrays.AxisArray{Float64, 1})
    if !isnothing(prices.capacity_price)
        prices.capacity_price[scenario_name] = capacity_price
    else
        prices.capacity_price = Dict(scenario_name => capacity_price)
    end
    return
end

function set_rec_price!(prices::MarketPrices, scenario_name::String, rec_price::AxisArrays.AxisArray{Float64, 1})
    if !isnothing(prices.rec_price)
        prices.rec_price[scenario_name] = rec_price
    else
        prices.rec_price = Dict(scenario_name => rec_price)
    end
    return
end


module OpenBLASHighCoreCount

using LinearAlgebra.BLAS
using OpenBLASHighCoreCount_jll

__init__() = BLAS.lbt_forward(OpenBLASHighCoreCount_jll.libopenblas_path)

end # module


# Avoid variables, constraints and objective being added when loading
@enum State Loading Modeling

mutable struct Data
    variables::Set{SetVariableRef}
    constraints::Dict{ConstraintIndex, SetConstraint}
    transformed_constraints::Dict{ConstraintIndex, JuMP.ConstraintRef}
    names::Dict{ConstraintIndex, String}
    last_index::Int
    objective_sense::MOI.OptimizationSense
    objective::Union{Nothing, AbstractScalarFunction}
    objective_variable::Union{Nothing, JuMP.VariableRef}
    perspective_polyvar::Union{Nothing, SpaceVariable}
    space::Space
    state::State
    spaces::Union{Nothing, Spaces}
end

function JuMP.copy_extension_data(::SetProg.Data, ::JuMP.Model, ::JuMP.Model)
end

function lift_space_variables(data::Data,
                              x::Vector{SpaceVariable})
    @assert data.perspective_polyvar > x[1]
    SpaceVariable[data.perspective_polyvar; x]
end

function set_space(cur::Space, space::Space)
    if space == Undecided || cur == Undecided || cur == space
        return space
    else
        error("Incompatible constraints/objective, some require to do the modeling in the primal space and some in the dual space.")
    end
end
function set_space(d::Data, model::JuMP.Model)
    space = Undecided
    for (index, constraint) in d.constraints
        space = set_space(space, constraint)
    end
    if d.objective !== nothing
        space = set_space(space, d.objective, model)
    end
    if space == Undecided
        space = PrimalSpace
    end
    d.space = space
end

# Remove spaces creating in earlier `optimize!`
function clear_spaces(d::Data)
    for vref in d.variables
        clear_spaces(vref)
    end
    for (index, constraint) in d.constraints
        clear_spaces(constraint)
    end
end
# No space index stored in constants
function clear_spaces(::Union{Polyhedra.Rep, Sets.AbstractSet}) end

# No perspective variable
function Sets.perspective_variable(::Polyhedra.Rep) end

synchronize_perspective(::Nothing, ::Nothing) = nothing
synchronize_perspective(::Nothing, z::SpaceVariable) = z
synchronize_perspective(z::SpaceVariable, ::Nothing) = z
function synchronize_perspective(z1::SpaceVariable, z2::SpaceVariable)
    if z1 !== z2
        throw(ArgumentError("Perspective variables do not match"))
    end
    return z1
end
function synchronize_perspective(d::Data)
    z = nothing
    for (index, constraint) in d.constraints
        z = synchronize_perspective(z, Sets.perspective_variable(constraint))
    end
    if z === nothing
        DynamicPolynomials.@polyvar z
    end
    d.perspective_polyvar = z
end

function create_spaces(set::Polyhedra.Rep, spaces::Spaces)
    return new_space(spaces, Polyhedra.fulldim(set))
end
function create_spaces(set::Union{Sets.AbstractSet, Sets.Projection}, spaces::Spaces)
    polyvars = Sets.space_variables(set)
    if polyvars === nothing
        return new_space(spaces, Sets.dimension(set))
    else
        return new_space(spaces, polyvars)
    end
end

# Create spaces and identify objects lying the the same space
function create_spaces(d::Data)
    spaces = Spaces()
    for vref in d.variables
        create_spaces(vref, spaces)
    end
    for (index, constraint) in d.constraints
        create_spaces(constraint, spaces)
    end
    d.spaces = spaces
end

function data(model::JuMP.Model)
    if !haskey(model.ext, :SetProg)
        model.ext[:SetProg] = Data(Set{SetVariableRef}(),
                                   Dict{ConstraintIndex, SetConstraint}(),
                                   Dict{ConstraintIndex, JuMP.ConstraintRef}(),
                                   Dict{ConstraintIndex, String}(), 0,
                                   MOI.FEASIBILITY_SENSE, nothing, nothing,
                                   nothing, Undecided, Modeling, nothing)
        model.optimize_hook = optimize_hook
    end
    return model.ext[:SetProg]
end


"""
Describes the maximum discharge for a dam output, which can be constant (`Constant*DamOutput`) or depend on the
hydraulic head with a conveyance factor (`Conveyance*DamOutput`); the latter make the assumption that the reservoir level
is the hydraulic head. This discharge can be conditional or not, which is indicated by the `condition` attribute
(whose type must be `DamOutputCondition`).

These releases are considered as degrees of freedom of the operator, including for flood management: the operator
can decide on the exact discharge through each dam output (if the conditions are met). These releases are not
supposed to fulfil any dam purpose (like drinking water), which are either constant or valued (see `Purpose` for these).

Each output must have a unique identifier, whose type is `Symbol`. It is always available with the `getId()` function.
For implementers, a default implementation uses the field `id`. Likewise, the condition is available through
`getCondition()`, and a default implementation retrieves the value from the field `condition`.
"""
abstract DamOutput

getId(dout::DamOutput) = dout.id
getCondition(dout::DamOutput) = dout.condition
hasCondition(dout::DamOutput) = ! isempty(getCondition(dout))

id(dout::DamOutput) = getId(dout)
condition(dout::DamOutput) = getCondition(dout)
isConditional(dout::DamOutput) = hasCondition(dout)

"""
Describes the maximum discharge for a dam output using a conveyance factor, be it an open-channel spillway
(`ConveyanceSpillway`) or a closed conduit (`ConveyancePipe`).

A default implementation of `getConveyanceFactor()` takes its value from the `conveyance` field.
"""
abstract ConveyanceDamOutput <: DamOutput

hasConveyanceFactor(dout::DamOutput) = false
hasConveyanceFactor(dout::ConveyanceDamOutput) = true
getConveyanceFactor(dout::ConveyanceDamOutput) = dout.conveyance
conveyance(dout::DamOutput) = getConveyanceFactor(dout)

"""
Condition at which a dam output is enabled. For example, a sufficient reservoir level is needed for the spillways to be
reachable.
"""
abstract DamOutputCondition

"""
Describes a dam output condition: the dam release is always possible.
"""
immutable NoCondition <: DamOutputCondition; end

isempty(c::DamOutputCondition) = false
isempty(c::NoCondition) = true

"""
Describes a constant dam output, independent of all parameters. Its discharge is given in 10^6 m^3/s.
"""
immutable ConstantDamOutput <: DamOutput
  id::Symbol
  discharge::Float64
  condition::DamOutputCondition

  ConstantDamOutput(discharge::Discharge, condition::DamOutputCondition=NoCondition()) =
    new(:constant, _from_unitful(discharge), condition)
  ConstantDamOutput(id::Symbol, discharge::Discharge, condition::DamOutputCondition=NoCondition()) =
    new(id, _from_unitful(discharge), condition)
end

getDischarge(dout::DamOutput) = error("Discharge has no meaning for output " * string(getId(dout)) * ".")
getDischarge(dout::ConstantDamOutput) = dout.discharge
getDischarge(dout::DamOutput, p::Period) = getDischarge(dout) * Dates.days(p) * 86400
isConstant(dout::DamOutput) = false
isConstant(dout::ConstantDamOutput) = true

discharge(dout::DamOutput) = getDischarge(dout)
discharge(dout::DamOutput, p::Period) = getDischarge(dout, p)

"""
Describes the penstock driving a hydropower unit.

This object is automatically created when creating a `Reservoir` data structure, and thus does not need to be used
outside internal code.
"""
immutable HydropowerDamOutput <: DamOutput
  hp::HydropowerUnit
  condition::DamOutputCondition

  HydropowerDamOutput(hp::HydropowerUnit, cond::DamOutputCondition=NoCondition()) = new(hp, cond)
end

getId(hp::HydropowerDamOutput) = getDamOutputId(hp.hp)
getDischarge(hp::HydropowerDamOutput) = getMaximumDischarge(hp.hp)

isHydropower(dout::DamOutput) = false
isHydropower(hp::HydropowerDamOutput) = true
getHydropowerUnit(hp::HydropowerDamOutput) = hp.hp

hydropower(hp::HydropowerDamOutput) = getHydropowerUnit(hp)

"""
Describes the maximum flow of an open-channel spillway using a conveyance factor. This discharge is limited by the
available hydraulic head.

This kind of dam output obeys the following law:
    Q = conveyance (head)^(3/2)

TODO: How to deal with units? Q is 10^6 m^3/period!
"""
immutable ConveyanceSpillway <: ConveyanceDamOutput
  id::Symbol
  conveyance::Float64
  condition::DamOutputCondition

  ConveyanceSpillway(conveyance::Float64, condition::DamOutputCondition=NoCondition()) =
    new(:spillway, conveyance, condition)
  ConveyanceSpillway(id::Symbol, conveyance::Float64, condition::DamOutputCondition=NoCondition()) =
    new(id, conveyance, condition)
end

"""
Describes the maximum flow of a closed conduit using a conveyance factor. This discharge is limited by the available
hydraulic head.

This kind of dam output obeys the following law:
    Q = conveyance (head)^(1/2)

TODO: Is this law the right one?
TODO: How to deal with units? Q is 10^6 m^3/period, while the coefficient will be given for m^3/s!
"""
immutable ConveyancePipe <: ConveyanceDamOutput
  id::Symbol
  conveyance::Float64
  condition::DamOutputCondition

  ConveyancePipe(conveyance::Float64, condition::DamOutputCondition=NoCondition()) =
    new(:penstock, conveyance, condition)
  ConveyancePipe(id::Symbol, conveyance::Float64, condition::DamOutputCondition=NoCondition()) =
    new(id, conveyance, condition)
end

"""
Describes a dam output condition: the release is possible only if this condition is met, i.e. if the reservoir level is
sufficiently high. This type represents physical constraints on the actual reservoir level (not an approximation for
hydraulic head).
"""
immutable MinimumReservoirLevelCondition <: DamOutputCondition
  min::Float64
  MinimumReservoirLevelCondition(min::Height) = new(_from_unitful(min))
end

getMinimumLevel(c::MinimumReservoirLevelCondition) = c.min

getMinimumLevel(c::DamOutputCondition) = error("Minimum reservoir level has no meaning for the condition.")
minLevel(c::DamOutputCondition) = getMinimumLevel(c)

# Equivalent functions for volumes: see output_fwd.jl (they need the declarations of other stuff).

"""
Builds an array of dam outputs for the most common cases, i.e. one set of bottom outlets and one spillway.
"""
function ConstantDamOutputs(;bottomOutlets::Discharge=0.m^3/s,
                             spillway::Discharge=0.m^3/s, spillway_minLevel::Height=0.m)
  list = DamOutput[] # All data structures expect a generic array of dam outputs, be they constant or not.
  if bottomOutlets != 0.m^3/s
    push!(list, ConstantDamOutput(:bottomOutlets, bottomOutlets))
  end
  if spillway != 0.m^3/s
    if spillway_minLevel != 0.m
      push!(list, ConstantDamOutput(:spillway, spillway, MinimumReservoirLevelCondition(spillway_minLevel)))
    else
      push!(list, ConstantDamOutput(:spillway, spillway))
    end
  end
  return list
end


# module for representing agent
# decide control order to robot

mutable struct Agent
    speed
    yaw_rate
    time_interval
    estimator
    prev_spd
    prev_yr

    # init
    function Agent(speed::Float64, yaw_rate::Float64;
                   time_interval::Float64=0.1,
                   estimator=nothing)
        self = new()
        self.speed = speed
        self.yaw_rate = yaw_rate
        self.time_interval = time_interval
        self.estimator = estimator
        self.prev_spd = 0.0
        self.prev_yr = 0.0
        return self
    end
end

function draw_decision!(self::Agent, observation)
    if self.estimator !== nothing
        motion_update(self.estimator, self.prev_spd, self.prev_yr, self.time_interval)
        self.prev_spd, self.prev_yr = self.speed, self.yaw_rate
        observation_update(self.estimator, observation)
        draw!(self.estimator)
    end
    return self.speed, self.yaw_rate
end


POMDPs.solve(solver::DPWSolver, mdp::Union{POMDP,MDP}) = DPWPlanner(solver, mdp)

"""
Delete existing decision tree.
"""
function clear_tree!(p::DPWPlanner)
    p.tree = nothing
end

"""
Construct an MCTSDPW tree and choose the best action.
"""
POMDPs.action(p::DPWPlanner, s) = first(action_info(p, s))

"""
Construct an MCTSDPW tree and choose the best action. Also output some information.
"""
function POMDPModelTools.action_info(p::DPWPlanner, s; tree_in_info=false)
    local a::actiontype(p.mdp)
    info = Dict{Symbol, Any}()
    try
        if isterminal(p.mdp, s)
            error("""
                  MCTS cannot handle terminal states. action was called with
                  s = $s
                  """)
        end

        S = statetype(p.mdp)
        A = actiontype(p.mdp)
        if p.solver.keep_tree
            if p.tree == nothing
                tree = DPWTree{S,A}(p.solver.n_iterations)
                p.tree = tree
            else
                tree = p.tree
            end
            if haskey(tree.s_lookup, s)
                snode = tree.s_lookup[s]
            else
                snode = insert_state_node!(tree, s, true)
            end
        else
            tree = DPWTree{S,A}(p.solver.n_iterations)
            p.tree = tree
            snode = insert_state_node!(tree, s, p.solver.check_repeat_state)
        end

        nquery = 0
        start_us = CPUtime_us()
        for i = 1:p.solver.n_iterations
            nquery += 1
            simulate(p, snode, p.solver.depth) # (not 100% sure we need to make a copy of the state here)
            if CPUtime_us() - start_us >= p.solver.max_time * 1e6
                break
            end
        end
        info[:search_time_us] = CPUtime_us() - start_us
        info[:tree_queries] = nquery
        if p.solver.tree_in_info || tree_in_info
            info[:tree] = tree
        end

        best_Q = -Inf
        sanode = 0
        for child in tree.children[snode]
            if tree.q[child] > best_Q
                best_Q = tree.q[child]
                sanode = child
            end
        end
        # XXX some publications say to choose action that has been visited the most
        a = tree.a_labels[sanode] # choose action with highest approximate value
    catch ex
        a = convert(actiontype(p.mdp), default_action(p.solver.default_action, p.mdp, s, ex))
        info[:exception] = ex
    end

    return a, info
end

"""
Return the reward for one iteration of MCTSDPW.
"""
function simulate(dpw::DPWPlanner, snode::Int, d::Int)
    S = statetype(dpw.mdp)
    A = actiontype(dpw.mdp)
    sol = dpw.solver
    tree = dpw.tree
    s = tree.s_labels[snode]
    if d == 0 || isterminal(dpw.mdp, s)
        return 0.0
    end

    # action progressive widening
    if dpw.solver.enable_action_pw
        if length(tree.children[snode]) <= sol.k_action*tree.total_n[snode]^sol.alpha_action # criterion for new action generation
            a = next_action(dpw.next_action, dpw.mdp, s, DPWStateNode(tree, snode)) # action generation step
            if !sol.check_repeat_action || !haskey(tree.a_lookup, (snode, a))
                n0 = init_N(sol.init_N, dpw.mdp, s, a)
                insert_action_node!(tree, snode, a, n0,
                                    init_Q(sol.init_Q, dpw.mdp, s, a),
                                    sol.check_repeat_action
                                   )
                tree.total_n[snode] += n0
            end
        end
    elseif isempty(tree.children[snode])
        for a in actions(dpw.mdp, s)
            n0 = init_N(sol.init_N, dpw.mdp, s, a)
            insert_action_node!(tree, snode, a, n0,
                                init_Q(sol.init_Q, dpw.mdp, s, a),
                                false)
            tree.total_n[snode] += n0
        end
    end

    best_UCB = -Inf
    sanode = 0
    ltn = log(tree.total_n[snode])
    for child in tree.children[snode]
        n = tree.n[child]
        q = tree.q[child]
        c = sol.exploration_constant # for clarity
        if (ltn <= 0 && n == 0) || c == 0.0
            UCB = q
        else
            UCB = q + c*sqrt(ltn/n)
        end
        @assert !isnan(UCB) "UCB was NaN (q=$q, c=$c, ltn=$ltn, n=$n)"
        @assert !isequal(UCB, -Inf)
        if UCB > best_UCB
            best_UCB = UCB
            sanode = child
        end
    end

    a = tree.a_labels[sanode]

    # state progressive widening
    new_node = false
    if tree.n_a_children[sanode] <= sol.k_state*tree.n[sanode]^sol.alpha_state
        sp, r = generate_sr(dpw.mdp, s, a, dpw.rng)

        if sol.check_repeat_state && haskey(tree.s_lookup, sp)
            spnode = tree.s_lookup[sp]
        else
            spnode = insert_state_node!(tree, sp, sol.keep_tree || sol.check_repeat_state)
            new_node = true
        end

        push!(tree.transitions[sanode], (spnode, r))

        if !sol.check_repeat_state
            tree.n_a_children[sanode] += 1
        elseif !((sanode,spnode) in tree.unique_transitions)
            push!(tree.unique_transitions, (sanode,spnode))
            tree.n_a_children[sanode] += 1
        end
    else
        spnode, r = rand(dpw.rng, tree.transitions[sanode])
        sp = tree.s_labels[spnode]
    end

    # THE NEXT TWO LINES WERE ADDED TO ALLOW FOR TERMINAL ACTIONS
    if isterminal(dpw.mdp, sp, a)
        q = r
    elseif new_node
        q = r + discount(dpw.mdp)*estimate_value(dpw.solved_estimate, dpw.mdp, sp, d-1)
    else
        q = r + discount(dpw.mdp)*simulate(dpw, spnode, d-1)
    end

    tree.n[sanode] += 1
    tree.total_n[snode] += 1

    tree.q[sanode] += (q - tree.q[sanode])/tree.n[sanode]

    return q
end


mutable struct AllMask
    featureMask
    valueMask
end

function getMask(mask::Dict)
    #получение маски из прочитанного словаря (из аттрибутов) в новый словарь с featureMask и valueMask
    newMask = Dict{String, Any}()
    for (key, val) in mask
        if !isa(val,Dict)
            newMask[key] = AllMask(val,val)
            newMask["Not"*key] = AllMask(val,0)
        else
            for (key2, val2) in val
                if val2 != 0 && !isequal(key2,"Mask")
                    newMask[key*"."*key2] = AllMask(val["Mask"],val2)
                end
            end
        end
    end
    return(newMask)
end

function parseType(atr)
    #преобразование аттрибутов в Dict маски типов
    mask = Dict{String, Any}()
    for str in atr
        if match(r"Type",str[1]) !== nothing
            sbstr = split(str[1],".")
            if length(sbstr[2:end]) > 1
                if !haskey(mask, sbstr[2])
                    mask[sbstr[2]] = Dict{String, Any}()
                end
                mask[sbstr[2]][sbstr[3]] = str[2]
            else
                mask[sbstr[2]] = str[2]
            end
        end
    end
    return(mask)
end

#получение типа int по набору признаков
function detIntType(mask::Dict,types::String)
    type = mask[types].featureMask & mask[types].valueMask
    typeint = Int(type)
end

#возвращает по типу имена входящих в него признаков
function detStringType(mask::Dict,type::Int)
    fields = mask.keys[mask.slots .!= 0]
    r = [checkType(type,mask[f]) for f in fields]
    fields[r]
end

#устанавливает в types биты, соответствующие bitName
function setbit(mask::Dict,types::Vector{T},bitName::Vector{String}) where T
    m = [mask[btN] for btN in bitName if haskey(mask, btN)]
    L = length(types)
    newTypes = Array{Int,1}()
    for i = 1:L
        type = types[i]
        for m1 in m
            type = type-type&m1.featureMask + m1.featureMask & m1.valueMask
        end
        push!(newTypes,type)
    end
    return newTypes
end

#построение дерева классов на основе класса и подкласса, с сохранением данных индексации попавших в категорию элементов
function buildQRStree(QPoint,Class, SubClass)

    unCl = unique(Class)
    QPoint = QPoint[1:length(Class)] #обрезаем по кол-ву классифицированных комплексов

    unSubCl = unique(SubClass)

    treeClasses = Dict{String,Any}()
    treeClasses["nodes"] = []
    treeClasses["name"] = "allClasses"
    treeClasses["size"] = length(Class)

    for i = 1:length(unCl)
        isCl = Class.==unCl[i]
        node =  Dict{String,Any}()
        node["nodes"] = []
        node["name"] = unCl[i]
        node["size"] = sum(isCl)
        node["QPoint"] = QPoint[isCl]
        for j = 1:length(unSubCl) #теперь добавляем подклассы
            isSubCl= (SubClass.== unSubCl[j]) .& isCl

            subnode =  Dict{String,Any}()
            subnode["nodes"] = []
            subnode["name"] = unSubCl[j]
            subnode["size"] = sum(isSubCl)
            subnode["QPoint"] = QPoint[isSubCl]
            push!(node["nodes"],subnode)
        end
        push!(treeClasses["nodes"],node)
    end

    result = JSON.json(treeClasses)
    return result
end


function rebuild_conv!(dst, src)
    shape = dst |> size
    filter_x, filter_y = shape[1:2] .+ 1
    for (i, j, k, m) in Iterators.product([1:s for s in shape]...)
        dst[filter_x - i, filter_y - j, k, m] = src[m, k, j, i]
    end
end

function _load_entry!(entry::Chain, params)
    rebuild_conv!(entry[1].weight, params["conv1.weight"])
    copyto!(entry[2].γ, params["bn1.weight"])
    copyto!(entry[2].β, params["bn1.bias"])
    copyto!(entry[2].μ, params["bn1.running_mean"])
    copyto!(entry[2].σ², params["bn1.running_var"])
end

function _load_basic_layer!(basic::Chain, params, layer_id::Int)
    for repeat in 1:length(basic)
        skipcon = basic[repeat].block
        layer = skipcon.layers
        lk = "layer$layer_id.$(repeat - 1)."

        rebuild_conv!(layer[1].weight, params[lk * "conv1.weight"])
        copyto!(layer[2].γ, params[lk * "bn1.weight"])
        copyto!(layer[2].β, params[lk * "bn1.bias"])
        copyto!(layer[2].μ, params[lk * "bn1.running_mean"])
        copyto!(layer[2].σ², params[lk * "bn1.running_var"])

        rebuild_conv!(layer[3].weight, params[lk * "conv2.weight"])
        copyto!(layer[4].γ, params[lk * "bn2.weight"])
        copyto!(layer[4].β, params[lk * "bn2.bias"])
        copyto!(layer[4].μ, params[lk * "bn2.running_mean"])
        copyto!(layer[4].σ², params[lk * "bn2.running_var"])

        if skipcon.connection isa Shortcut
            connection = skipcon.connection.s
            lk = "layer$layer_id.0.downsample."
            rebuild_conv!(connection[1].weight, params[lk * "0.weight"])
            copyto!(connection[2].γ, params[lk * "1.weight"])
            copyto!(connection[2].β, params[lk * "1.bias"])
            copyto!(connection[2].μ, params[lk * "1.running_mean"])
            copyto!(connection[2].σ², params[lk * "1.running_var"])
        end
    end
end

function _load_layers!(layers::Chain, params, model_size::Int)
    for i in 1:length(layers)
        if model_size < 50
            _load_basic_layer!(layers[i], params, i)
        else
            throw("Not implemented")
        end
    end
end

function from_pretrained(
    model_size::Int; cache_dir::Union{String, Nothing} = nothing, kwargs...,
)
    # url_base = "https://download.pytorch.org/models/"
    url_base = "https://dl.fbaipublicfiles.com/semiweaksupervision/model_files/"
    url_map = Dict(
        18 => "semi_supervised_resnet18-d92f0530.pth",
        50 => "semi_supervised_resnet50-08389792.pth",
        # 18 => "resnet18-5c106cde.pth",
        # 34 => "resnet34-333f7ec4.pth",
        # 101 => "resnet101-5d3b4d8f.pth",
        # 152 => "resnet152-b121ed2d.pth",
    )

    if !(model_size in keys(url_map))
        error(
            "Invalid model model_size: $model_size. " *
            "Supported pretrained models: $(keys(url_map)) "
        )
    end

    if cache_dir ≡ nothing
        cache_dir = joinpath(homedir(), ".cache", "ResNet.jl")
        !isdir(cache_dir) && mkdir(cache_dir)
        @info "Using default cache dir $cache_dir"
    end

    params_file = url_map[model_size]
    params_path = joinpath(cache_dir, params_file)
    if !isfile(params_path)
        download_url = url_base * params_file
        @info(
            "Downloading ResNet$model_size params:\n" *
            "\t- from URL: $download_url \n" *
            "\t- to directory: $params_path"
        )
        download(download_url, params_path)
        @info "Finished downloading params."
    end

    params = Pickle.Torch.THload(params_path)
    model = ResNetModel(model_size; kwargs...)

    _load_entry!(model.entry, params)
    _load_layers!(model.layers, params, model.size)
    if model.head ≢ nothing
        copyto!(model.head[end].weight, params["fc.weight"])
        copyto!(model.head[end].bias, params["fc.bias"])
    end

    model
end