text:
using Test
using Dates

using ScottishTaxBenefitModel

using .ModelHousehold:
    Household,
    Person,
    People_Dict,
    default_bu_allocation,
    get_benefit_units,
    get_head,
    get_spouse,
    is_single,
    pers_is_carer,
    pers_is_disabled,
    search

using .IncomeTaxCalculations:
    calc_income_tax!

using .Definitions

using .LegacyMeansTestedBenefits:
    LMTResults,
    calc_allowances,
    calc_incomes,
    calc_legacy_means_tested_benefits!,
    calc_NDDs,
    calc_premia,
    calculateHB_CTR!,
    calcWTC_CTC!,
    is_working_hours,
    make_lmt_benefit_applicability,
    tariff_income,
    working_disabled

using .LocalLevelCalculations:
    apply_rent_restrictions,
    calc_council_tax

using .STBIncomes

using .Intermediate:
    MTIntermediate,
    apply_2_child_policy,
    make_intermediate

using .NonMeansTestedBenefits:
    calc_pre_tax_non_means_tested!,
    calc_post_tax_non_means_tested!

using .LocalLevelCalculations:
    calc_council_tax

using .STBParameters:
    ChildLimits,
    HoursLimits,
    IncomeRules,
    LegacyMeansTestedBenefitSystem

using .Results:
    BenefitUnitResult,
    LMTResults,
    LMTCanApplyFor,
    init_household_result,
    init_benefit_unit_result,
    to_string
using .Utils:
    eq_nearest_p,
    to_md_table
using .ExampleHelpers

## FIXME don't need both
lmt = LegacyMeansTestedBenefitSystem{Float64}()
sys = get_system( scotland=true )

@testset "2 child policy" begin
    sph = get_example(single_parent_hh)
    sparent = get_benefit_units(sph)[1]
    println( "keys of initial children  $(sparent.children)" )

    @test num_children( sparent ) == 2
    @test apply_2_child_policy( sparent, sys.child_limits ) == 2

    np = add_child!( sph, 10, Female )
    sparent = get_benefit_units(sph)[1]
    println( "keys of children after 10yo added $(sparent.children) new pid = $np" )
    @test num_children( sparent ) == 3
    @test apply_2_child_policy( sparent, sys.child_limits ) == 3

    np = add_child!( sph, 1, Female )
    sparent = get_benefit_units(sph)[1]
    @test num_children( sparent ) == 4
    @test apply_2_child_policy( sparent, sys.child_limits ) == 3
    child_limits = ChildLimits(Date( 2017, 4, 6 ),5)
    @test apply_2_child_policy( sparent, child_limits ) == 4

    child_limits = ChildLimits(Date( 2010, 4, 6 ), 2 )
    @test apply_2_child_policy( sparent, child_limits ) == 2
end

@testset "CPAG income and capital chapters 20, 21, 22, 23" begin

    income = [110.0,145.0,325,755.0,1_000.0]

    #  basic check of tariff incomes; cpag ch 21 ?? page
    caps = Dict(2000=>0,6000=>0,6000.01=>1, 6253=>2,8008=>9)
    for (c,t) in caps
        ci = tariff_income(c,
            lmt.income_rules.capital_min,
            lmt.income_rules.capital_tariff)
        println("ci=$ci t=$t c=$c")
        @test (ci == t)
    end

    ntests = size(income)[1]

    for (hht,hh) in get_all_examples()
        println( "on hhld '$hht'")
        bus = get_benefit_units( hh )
        bu = bus[1]
        intermed = make_intermediate(
            1,
            bu,
            lmt.hours_limits,
            sys.age_limits,
            sys.child_limits,
            1 )

        @test (size(bus)[1] == 1) || (hht == mbu )
        spouse = nothing
        head = get_head(bu)
        hdwork = is_working_hours( head, lmt.hours_limits.lower )
        head.usual_hours_worked = 5
        head.employment_status = Unemployed
        if hht in [cpl_w_2_children_hh childless_couple_hh]
            spouse = get_spouse(bu)
            spouse.employment_status = Full_time_Employee
            spouse.usual_hours_worked = 45
            working = search( bu, is_working_hours, lmt.hours_limits.lower )
            println( "working $working ")
            @test working
            spouse.usual_hours_worked = 5
            spouse.employment_status = Unemployed
            spwork = is_working_hours( spouse, lmt.hours_limits.lower )
            working = search( bu, is_working_hours, lmt.hours_limits.lower )
            println( "working $working spwork $spwork hdwork $hdwork")
            @test ! working
            @test ! is_single( bu )
        end
        head.usual_hours_worked = 0
        head.employment_status = Unemployed
        for i in 1:ntests
            if spouse !== nothing
                spouse.income[wages] = income[i]
            else
                head.income[wages] = income[i]
            end
            for ben in [esa hb is jsa pc]
                bur = init_benefit_unit_result( bu )
                println("on $ben")
                inc = calc_incomes(
                    ben,
                    bu,
                    bur,
                    intermed,
                    lmt.income_rules,
                    lmt.hours_limits )
                @test inc.tariff_income ≈ 0.0
                # cpl_w_2_children_hh single_parent_hh single_hh childless_couple_hh

                if hht == single_parent_hh
                    if ben == hb
                        @test inc.disregard == 25.0
                    else
                        @test inc.disregard == 20.0
                    end
                elseif hht == cpl_w_2_children_hh
                    if ben == hb
                        @test inc.disregard == 10.0
                    elseif ben == esa
                        @test inc.disregard == 20.0
                    else
                        @test inc.disregard == 10.0
                    end

                end
                ## TODO: finish income tests ..
                # println( "inc res for wage $(income[i])ben $ben = \n $inc")
            end # bens loop
        end # incomes loop
    end # households loop

    # Evan and Mia example p 433
    e_and_m = get_benefit_units( get_example( cpl_w_2_children_hh ) )[1]
    evan = get_head( e_and_m )
    mia = get_spouse( e_and_m )
    intermed = make_intermediate(
        1,
        e_and_m,
        lmt.hours_limits,
        sys.age_limits,
        sys.child_limits,
        1 )
    empty!(mia.income)
    mia.income[wages] = 136.0
    mia.usual_hours_worked = 17
    empty!(evan.income)
    ep = evan.pid
    emr = init_benefit_unit_result(  e_and_m )
    emr.pers[ep].income[NON_CONTRIB_EMPLOYMENT_AND_SUPPORT_ALLOWANCE] = 1.0
    inc = calc_incomes(
        hb,
        e_and_m,
        emr,
        intermed,
        lmt.income_rules,
        lmt.hours_limits )
    @test inc.net_earnings ≈ 98.90
    @test inc.total_income ≈ 99.90
    @test inc.disregard ≈ 37.10

end # test set

@testset "Applicability Tests" begin
    cpl = get_benefit_units(get_example(cpl_w_2_children_hh))[1]
    sparent = get_benefit_units(get_example(single_parent_hh))[1]
    intermed = make_intermediate(
        1,
        cpl,
        sys.lmt.hours_limits,
        sys.age_limits,
        sys.child_limits,
        1 )
    eligs_cpl = make_lmt_benefit_applicability(
        sys.lmt, intermed, sys.lmt.hours_limits )
    head = get_head( cpl )
    println( "head=$head")
    spouse = get_spouse( cpl )
    println( "sp=$spouse" )
    println( "couple $eligs_cpl" )
    @test ! eligs_cpl.esa
    # @test ! eligs_cpl.hb
    @test ! eligs_cpl.is
    @test ! eligs_cpl.jsa
    @test ! eligs_cpl.pc
    # @test ! eligs_cpl.ndds
    @test eligs_cpl.wtc
    @test eligs_cpl.ctr
    #2 not_working couple - jsa
    unemploy!( head )
    unemploy!( spouse )
    intermed = make_intermediate(
        1,
        cpl,
        sys.lmt.hours_limits,
        sys.age_limits,
        sys.child_limits,
        1 )
    eligs_cpl = make_lmt_benefit_applicability( sys.lmt, intermed, sys.lmt.hours_limits )
    println( "not_working: sp=$spouse" )
    println( "not_working couple $eligs_cpl" )
    @test ! eligs_cpl.esa
    # @test ! eligs_cpl.hb
    @test ! eligs_cpl.is
    @test eligs_cpl.jsa
    @test ! eligs_cpl.pc
    # @test ! eligs_cpl.ndds
    @test ! eligs_cpl.wtc
    @test eligs_cpl.ctr

    disable_slightly!( spouse )
    intermed = make_intermediate(
        1,
        cpl,
        sys.lmt.hours_limits,
        sys.age_limits,
        sys.child_limits,
        1 )
    eligs_cpl = make_lmt_benefit_applicability( sys.lmt, intermed, sys.lmt.hours_limits )
    println( "not_working: sp=$spouse" )
    println( "not_working couple $eligs_cpl" )
    @test eligs_cpl.esa
    # @test ! eligs_cpl.hb
    @test ! eligs_cpl.is
    @test eligs_cpl.ctr
    enable!( spouse )
    unemploy!( spouse )
    carer!( spouse )
    carer!( head )
    println( "head.employment_status=$(head.employment_status) spouse.employment_status=$(spouse.employment_status)")
    intermed = make_intermediate(
        1,
        cpl,
        sys.lmt.hours_limits,
        sys.age_limits,
        sys.child_limits,
        1 )
    eligs_cpl = make_lmt_benefit_applicability( sys.lmt, intermed, sys.lmt.hours_limits )
    println( eligs_cpl )
    @test ! eligs_cpl.esa
    @test eligs_cpl.is
    @test eligs_cpl.ctr
    @test ! eligs_cpl.wtc

    intermed = make_intermediate(
        1,
        sparent,
        sys.lmt.hours_limits,
        sys.age_limits,
        sys.child_limits,
        1 )
    eligs_sp = make_lmt_benefit_applicability( sys.lmt, intermed, sys.lmt.hours_limits )
    println( "single parent $eligs_sp" )
    @test ! eligs_sp.is
    @test ! eligs_sp.jsa
    @test ! eligs_sp.pc
    # @test ! eligs_sp.ndds
    @test eligs_sp.wtc
    @test eligs_sp.ctr
    head = get_head( sparent )
    unemploy!( head )
    intermed = make_intermediate(
        1,
        sparent,
        sys.lmt.hours_limits,
        sys.age_limits,
        sys.child_limits,
        1 )
    eligs_sp = make_lmt_benefit_applicability( sys.lmt, intermed, sys.lmt.hours_limits )
    println( "single parent $eligs_sp" )
    @test ! eligs_sp.is
    @test eligs_sp.jsa
    @test ! eligs_sp.pc
    # @test ! eligs_sp.ndds
    @test ! eligs_sp.wtc
    @test eligs_sp.ctr
    carer!( head )
    intermed = make_intermediate(
        1,
        sparent,
        sys.lmt.hours_limits,
        sys.age_limits,
        sys.child_limits,
        1 )
    eligs_sp = make_lmt_benefit_applicability( sys.lmt, intermed, sys.lmt.hours_limits )
    println( "single parent $eligs_sp" )
    @test eligs_sp.is
    @test ! eligs_sp.jsa
    @test ! eligs_sp.pc
    # @test ! eligs_sp.ndds
    @test ! eligs_sp.wtc
    @test eligs_sp.ctr
    head.age = 70
    intermed = make_intermediate(
        1,
        sparent,
        sys.lmt.hours_limits,
        sys.age_limits,
        sys.child_limits,
        1 )
    eligs_sp = make_lmt_benefit_applicability( sys.lmt, intermed, sys.lmt.hours_limits )
    println( "single parent $eligs_sp" )
    @test ! eligs_sp.is
    @test ! eligs_sp.jsa
    @test eligs_sp.pc
    # @test ! eligs_sp.ndds
    @test ! eligs_sp.wtc
    @test eligs_sp.ctr
    retire!( head )
    intermed = make_intermediate(
        1,
        sparent,
        sys.lmt.hours_limits,
        sys.age_limits,
        sys.child_limits,
        1 )
    eligs_sp = make_lmt_benefit_applicability(
        sys.lmt,
        intermed,
        sys.lmt.hours_limits )
    @test ! eligs_sp.is
    @test ! eligs_sp.jsa
    @test eligs_sp.pc
    # @test ! eligs_sp.ndds
    @test ! eligs_sp.wtc  # this is right - could be on both ctr and wtc
    @test eligs_sp.ctr
    employ!( head )
    intermed = make_intermediate(
        1,
        sparent,
        sys.lmt.hours_limits,
        sys.age_limits,
        sys.child_limits,
        1 )
    eligs_sp = make_lmt_benefit_applicability( sys.lmt, intermed, sys.lmt.hours_limits )
    @test ! eligs_sp.is
    @test ! eligs_sp.jsa
    @test eligs_sp.pc
    # @test ! eligs_sp.ndds
    @test eligs_sp.wtc  # this is right - could be on both pc and wtc
    @test eligs_sp.ctr

end

@testset "Intermediates" begin
    sys = get_system( scotland=true )
    cpl = get_benefit_units(get_example(cpl_w_2_children_hh))[1]
    spouse = get_spouse( cpl )
    head = get_head( cpl )
    println( "sp=$spouse" )
    intermed = make_intermediate(
        1,
        cpl,
        sys.lmt.hours_limits,
        sys.age_limits,
        sys.child_limits,
        1 )
    println( intermed )
    @test ! intermed.someone_pension_age
    @test ! intermed.all_pension_age
    @test intermed.someone_working_ft
    @test intermed.num_working_pt == 0
    @test intermed.num_working_24_plus == 2
    @test intermed.total_hours_worked > 40
    @test ! intermed.someone_is_carer
    @test ! intermed.is_sparent
    @test ! intermed.is_sing
    @test ! intermed.is_disabled
    @test intermed.num_children > 1
    @test intermed.ge_16_u_pension_age
    @test ! intermed.limited_capacity_for_work
    @test intermed.has_children
    @test intermed.economically_active
    @test intermed.num_working_ft == 2
    @test intermed.num_not_working == 0
    @test intermed.num_working_pt == 0

    unemploy!( head )
    unemploy!( spouse )
    intermed = make_intermediate(
        1,
        cpl,
        sys.lmt.hours_limits,
        sys.age_limits,
        sys.child_limits,
        1 )
    @test ! intermed.someone_pension_age
    @test ! intermed.all_pension_age
    @test ! intermed.someone_working_ft
    @test intermed.num_working_pt == 0
    @test intermed.num_working_24_plus == 0
    @test intermed.total_hours_worked == 0
    @test ! intermed.someone_is_carer
    @test ! intermed.is_sparent
    @test ! intermed.is_sing
    @test ! intermed.is_disabled
    @test intermed.num_children > 1
    @test intermed.ge_16_u_pension_age
    @test ! intermed.limited_capacity_for_work
    @test intermed.has_children
    @test intermed.economically_active # not_working is active
    @test intermed.num_working_ft == 0

    disable_slightly!( spouse )
    intermed = make_intermediate(
        1,
        cpl,
        sys.lmt.hours_limits,
        sys.age_limits,
        sys.child_limits,
        1 )
    @test ! intermed.someone_working_ft
    @test intermed.num_working_pt == 0
    @test intermed.num_working_24_plus == 0
    @test intermed.total_hours_worked == 0
    @test ! intermed.someone_is_carer
    @test ! intermed.is_sparent
    @test ! intermed.is_sing
    @test intermed.is_disabled
    @test intermed.num_children == 2
    @test intermed.ge_16_u_pension_age
    @test intermed.limited_capacity_for_work
    @test intermed.has_children
    @test intermed.economically_active # not_working is active
    @test intermed.num_working_ft == 0
    @test intermed.num_not_working == 2 # ! 1 not_working/1 inactive?
    @test intermed.num_working_pt == 0

    carer!( head )
    intermed = make_intermediate(
        1,
        cpl,
        sys.lmt.hours_limits,
        sys.age_limits,
        sys.child_limits,
        1 )
    @test ! intermed.someone_working_ft
    @test intermed.num_not_working == 2 # ! 1 not_working/1 inactive?
    @test intermed.num_working_pt == 0
    @test intermed.num_working_ft == 0
    @test intermed.num_working_24_plus == 0
    @test intermed.total_hours_worked == 0
    @test intermed.someone_is_carer
    @test ! intermed.is_sparent
    @test ! intermed.is_sing
    @test intermed.is_disabled
    @test intermed.num_children == 2
    @test intermed.ge_16_u_pension_age
    @test intermed.limited_capacity_for_work
    @test intermed.has_children
    @test ! intermed.economically_active # not_working is active

    sparent = get_benefit_units(get_example(single_parent_hh))[1]
    intermed = make_intermediate(
        1,
        sparent,
        sys.lmt.hours_limits,
        sys.age_limits,
        sys.child_limits,
        1 )
    @test intermed.someone_working_ft
    @test intermed.num_working_pt == 0
    @test intermed.num_working_24_plus == 1
    @test intermed.total_hours_worked >= 40
    @test ! intermed.someone_is_carer
    @test intermed.is_sparent
    @test ! intermed.is_sing
    @test ! intermed.is_disabled
    @test intermed.num_children > 0
    @test intermed.ge_16_u_pension_age
    @test ! intermed.limited_capacity_for_work
    @test intermed.has_children
    @test intermed.economically_active # not_working is active
    @test intermed.num_working_ft == 1
    @test intermed.num_not_working == 0 # ! 1 not_working/1 inactive?
    @test intermed.num_working_pt == 0
    @test ! intermed.working_disabled

end

@testset "Allowances" begin
    sys = get_system( scotland=true )
    # TODO
end

@testset "NDDS" begin
    sys = get_system( scotland=true )
    spers = get_benefit_units(get_example(single_hh))[1]
    head = get_head( spers ) # the only person, obvs..
    head.age = 30
    empty!(head.income)
    bur = init_benefit_unit_result(  spers )
    wage = [600.0,400,200,0]
    ndds = [100.65, 91.70, 35.85, 15.60]
    nt = size(wage)[1]
    employ!( head )
    println( "FT Employed")
    println( head )
    for i in 1:nt
        bur.pers[head.pid].income[WAGES] = wage[i]
        intermed = make_intermediate(
            1,
            spers,
            sys.lmt.hours_limits,
            sys.age_limits,
            sys.child_limits,
            1 )
        incomes = calc_incomes(
            hb,
            spers,
            bur,
            intermed,
            sys.lmt.income_rules,
            sys.lmt.hours_limits )
        println( "incomes $(wage[i])" )
        println(  to_md_table( incomes ))
        ndd = calc_NDDs( spers, bur, intermed, incomes, sys.lmt.hb )
        @test ndd ≈ ndds[i]
    end # loop round various incomes
    println( "Unemployed")
    unemploy!( head )
    for i in 1:nt
        head.income[wages] = wage[i]
        intermed = make_intermediate(
            1,
            spers,
            sys.lmt.hours_limits,
            sys.age_limits,
            sys.child_limits,
            1 )
        incomes = calc_incomes(
            hb,
            spers,
            bur,
            intermed,
            sys.lmt.income_rules,
            sys.lmt.hours_limits )
        ndd = calc_NDDs( spers, bur, intermed, incomes, sys.lmt.hb )
        @test ndd ≈ ndds[nt]
    end # loop round various incomes, unemployed
    disable_slightly!( head )
    head.income[attendance_allowance] = 100.0
    bur = init_benefit_unit_result(  spers )
    println( "Disabled")
    for i in 1:nt
        # head.income[wages] = wage[i]
        bur.pers[head.pid].income[WAGES] = wage[i]
        intermed = make_intermediate(
            1,
            spers,
            sys.lmt.hours_limits,
            sys.age_limits,
            sys.child_limits,
            1 )
        incomes = calc_incomes(
            hb,
            spers,
            bur,
            intermed,
            sys.lmt.income_rules,
            sys.lmt.hours_limits )
        ndd = calc_NDDs( spers, bur, intermed, incomes, sys.lmt.hb )
        @test ndd ≈ 0.0
    end # loop ro
    empty!(head.income)
    enable!( head )
    blind!( head )
    bur = init_benefit_unit_result(  spers )
    println( "Blind")
    for i in 1:nt
        # head.income[wages] = wage[i]
        bur.pers[head.pid].income[WAGES] = wage[i]
        intermed = make_intermediate(
            1,
            spers,
            sys.lmt.hours_limits,
            sys.age_limits,
            sys.child_limits,
            1 )
        incomes = calc_incomes(
            hb,
            spers,
            bur,
            intermed,
            sys.lmt.income_rules,
            sys.lmt.hours_limits )
        ndd = calc_NDDs( spers, bur, intermed, incomes, sys.lmt.hb )
        @test ndd ≈ 0.0
    end # loop ro
    # couples
    cpl = get_benefit_units(get_example( cpl_w_2_children_hh ))[1]
    spouse = get_spouse( cpl )
    head = get_head( cpl )
    employ!( spouse )
    empty!( head.income )
    unemploy!( head )
    bur = init_benefit_unit_result(  spers )
    bur = init_benefit_unit_result(  cpl )
    println( "FT Employed")
    println( spouse )
    for i in 1:nt
        # spouse.income[wages] = wage[i]
        bur.pers[spouse.pid].income[WAGES] = wage[i]
        intermed = make_intermediate(
            1,
            cpl,
            sys.lmt.hours_limits,
            sys.age_limits,
            sys.child_limits,
            1 )
        incomes = calc_incomes(
            hb,
            cpl,
            bur,
            intermed,
            sys.lmt.income_rules,
            sys.lmt.hours_limits )
        ndd = calc_NDDs( cpl, bur, intermed, incomes, sys.lmt.hb )
        @test ndd ≈ ndds[i]
    end # loop round various incomes
    employ!( head )
    for i in 1:nt
        # spouse.income[wages] = wage[i] / 3.0
        # head.income[wages] = wage[i] * 2 / 3.0
        bur.pers[spouse.pid].income[WAGES] = wage[i] / 3.0
        bur.pers[head.pid].income[WAGES] = wage[i]* 2 / 3.0
        intermed = make_intermediate(
            1,
            cpl,
            sys.lmt.hours_limits,
            sys.age_limits,
            sys.child_limits,
            1 )
        incomes = calc_incomes(
            hb,
            cpl,
            bur,
            intermed,
            sys.lmt.income_rules,
            sys.lmt.hours_limits )
        ndd = calc_NDDs( cpl, bur, intermed, incomes, sys.lmt.hb )
        @test ndd ≈ ndds[i]
    end # loop round various incomes

    # no NDDs for CTB (I think)
    bur = init_benefit_unit_result(  spers )
    head = get_head( spers )
    employ!( head )
    bur = init_benefit_unit_result(  spers )
    println( "FT Employed - CTB")
    println( head )
    for i in 1:nt
        # head.income[wages] = wage[i]
        bur.pers[head.pid].income[WAGES] = wage[i]
        intermed = make_intermediate(
            1,
            spers,
            sys.lmt.hours_limits,
            sys.age_limits,
            sys.child_limits,
            1 )
        incomes = calc_incomes(
            hb,
            spers,
            bur,
            intermed,
            sys.lmt.income_rules,
            sys.lmt.hours_limits )
        ndd = calc_NDDs( spers, bur, intermed, incomes, sys.lmt.ctr )
        @test ndd ≈ 0.0
    end # loop round various incomes
end

@testset "Allowances" begin
    sys = get_system( scotland=true )
    sing = get_example(single_hh)
    singbu = get_benefit_units(sing)[1]
    head = get_head( singbu )
    head.age = 25
    empty!( head.income )
    intermed = make_intermediate(
            1,
            singbu,
            sys.lmt.hours_limits,
            sys.age_limits,
            sys.child_limits,
            1 )

    for ben in [hb,ctr,is,jsa,esa]
        allow = calc_allowances(
            ben,
            intermed,
            sys.lmt.allowances,
            sys.age_limits,
            0.0 # oo h costs
        )
        @test allow ==  sys.lmt.allowances.age_25_and_over
    end
    head.age = 17
    intermed = make_intermediate(
            1,
            singbu,
            sys.lmt.hours_limits,
            sys.age_limits,
            sys.child_limits,
            1 )
    for ben in [hb,ctr,is,jsa,esa]
        allow = calc_allowances(
            ben,
            intermed,
            sys.lmt.allowances,
            sys.age_limits,
            0.0 # oo h costs
        )
        println("on ben $ben")
        if ben != esa
            @test allow == sys.lmt.allowances.age_18_24 # you get the 18 except in odd cases; cpag p336
        else
            @test allow == sys.lmt.allowances.age_25_and_over
        end
    end
    #
    # 17 yo single parent - should be same as single
    #

    spar = get_example(single_parent_hh)
    sparbu = get_benefit_units(spar)[1]
    head = get_head(sparbu)

    intermed = make_intermediate(
            1,
            sparbu,
            sys.lmt.hours_limits,
            sys.age_limits,
            sys.child_limits,
            1 )
    head.age = 17
    intermed = make_intermediate(
            1,
            sparbu,
            sys.lmt.hours_limits,
            sys.age_limits,
            sys.child_limits,
            1 )
    @test intermed.num_children == 2
    @test intermed.num_allowed_children == 2
    for ben in [hb,ctr,is,jsa,esa]
        allow = calc_allowances(
            ben,
            intermed,
            sys.lmt.allowances,
            sys.age_limits,
            0.0
        )
        println( "on $ben allow=$allow" )
        if ben in [hb,ctr]
            @test allow == sys.lmt.allowances.age_18_24 + 2*sys.lmt.allowances.child
        elseif ben in [esa]
            @test allow == sys.lmt.allowances.age_25_and_over
        else
            @test allow ==  sys.lmt.allowances.age_18_24 # you get the 18 except in odd cases; cpag p336
        end
    end
    #
    # funny age combinations
    # a) 2 < 18
    cplhh = get_example(cpl_w_2_children_hh)
    cpl = get_benefit_units(cplhh)[1]
    spouse = get_spouse( cpl )
    head = get_head( cpl )
    spouse.age = 17
    head.age = 17
    println( "sp=$spouse" )
    intermed = make_intermediate(
        1,
        cpl,
        sys.lmt.hours_limits,
        sys.age_limits,
        sys.child_limits,
        1 )
    println( intermed )
    for ben in [hb,ctr,is,jsa,esa]
        allow = calc_allowances(
            ben,
            intermed,
            sys.lmt.allowances,
            sys.age_limits,
            0.0  )
        if ben in [hb,ctr]
            @test allow ≈ 87.50 + 2*sys.lmt.allowances.child
        else
            @test allow ≈ 87.50 # p336 weird stuff about ESA stages
        end

    end
    spouse.age = 17
    head.age = 18
    intermed = make_intermediate(
        1,
        cpl,
        sys.lmt.hours_limits,
        sys.age_limits,
        sys.child_limits,
        1 )
    println( intermed )
    for ben in [hb,ctr,is,jsa,esa]
        allow = calc_allowances(
            ben,
            intermed,
            sys.lmt.allowances,
            sys.age_limits,
            0.0 )
        if ben in [hb,ctr]
            @test allow ≈ 114.85 + 2*sys.lmt.allowances.child
        else
            @test allow ≈ 114.85 # p336 weird stuff about ESA stages
        end
    end
    spouse.age = 18
    head.age = 18
    intermed = make_intermediate(
        1,
        cpl,
        sys.lmt.hours_limits,
        sys.age_limits,
        sys.child_limits,
        1 )
    println( intermed )
    for ben in [hb,ctr,is,jsa,esa]
        allow = calc_allowances(
            ben,
            intermed,
            sys.lmt.allowances,
            sys.age_limits,
            0.0 )
        if ben in [hb,ctr]
            @test allow ≈ 114.85 + 2*sys.lmt.allowances.child
        else
            @test allow ≈ 114.85 # p336 weird stuff about ESA stages
        end
    end
    spouse.age = 60
    head.age = 60
    intermed = make_intermediate(
        1,
        cpl,
        sys.lmt.hours_limits,
        sys.age_limits,
        sys.child_limits,
        1 )
    for ben in [hb,ctr,is,jsa,esa]
        allow = calc_allowances(
            ben,
            intermed,
            sys.lmt.allowances,
            sys.age_limits,
            0.0 )
        if ben in [hb,ctr]
            @test allow ≈ 114.85 + 2*sys.lmt.allowances.child
        else
            @test allow ≈ 114.85 # p336 weird stuff about ESA stages
        end
    end
end

@testset "Premia" begin
    sys = get_system( scotland=true )
    cplhh = get_example(cpl_w_2_children_hh)
    cpl = get_benefit_units(cplhh)[1]
    bures = init_benefit_unit_result(cpl)
    spouse = get_spouse( cpl )
    head = get_head( cpl )
    spouse.age = 60
    head.age = 60
    intermed = make_intermediate(
        1,
        cpl,
        sys.lmt.hours_limits,
        sys.age_limits,
        sys.child_limits,
        1 )
    for ben in [hb,ctr,is,jsa,esa]
        premia, premset = calc_premia(
            ben,
            cpl,
            bures,
            intermed,
            sys.lmt.premia,
            sys.nmt_bens,
            sys.age_limits )
        @test premia ≈ 0.0
        @test length(premset)==0
    end
    disable_slightly!( head )
    bures.pers[head.pid].income[PERSONAL_INDEPENDENCE_PAYMENT_DAILY_LIVING] =
        sys.nmt_bens.pip.dl_standard
    intermed = make_intermediate(
        1,
        cpl,
        sys.lmt.hours_limits,
        sys.age_limits,
        sys.child_limits,
        1 )
    @test pers_is_disabled( head )
    @test intermed.num_disabled_adults == 1
    @test intermed.working_disabled
    for ben in [hb,ctr,is,jsa,esa]
        premia, premset = calc_premia(
            ben,
            cpl,
            bures,
            intermed,
            sys.lmt.premia,
            sys.nmt_bens,
            sys.age_limits )
        println( "ben=$ben premia=$premia" )
        if ben == esa
            @test premia ≈ 0.0
            @test length(premset)==0
        else
            @test premia == sys.lmt.premia.disability_single
            @test (disability_single ∈ premset) && length(premset)==1
        end
    end
    disable_seriously!( head )
    intermed = make_intermediate(
        1,
        cpl,
        sys.lmt.hours_limits,
        sys.age_limits,
        sys.child_limits,
        1 )
    @test pers_is_disabled( head )
    @test intermed.num_disabled_adults == 1
    @test intermed.working_disabled
    for ben in [hb,ctr,is,jsa,esa]
        premia, premset = calc_premia(
            ben,
            cpl,
            bures,
            intermed,
            sys.lmt.premia,
            sys.nmt_bens,
            sys.age_limits )
        println( "ben=$ben premia=$premia" )
        if ben == esa
            @test premia ≈ 0.0 # sys.lmt.premia.enhanced_disability_single
            @test length(premset)==0 ## FIXME check this ESA allowed enhanced_disability_single
        else
            @test premia == sys.lmt.premia.disability_single # +sys.lmt.premia.enhanced_disability_single - no since 2 people but 1 disabled
            @test (disability_single ∈ premset) && length(premset)==1
        end
    end
end

@testset "HB/CTB" begin
    # CPAG 19/20 p190
    sys = get_system( scotland=true )
    joplings = get_example(childless_couple_hh)
    jbu = get_benefit_units(joplings)[1]
    spouse = get_spouse( jbu )
    head = get_head( jbu )
    empty!( spouse.income )
    empty!( head.income )
    employ!( spouse )
    unemploy!( head )
    joplings.gross_rent = 120.00 # eligible rent
    joplings.bedrooms = 1
    spouse.income[wages] = 201.75
    spouse.usual_hours_worked = 21
    head.income[wages] = 73.10 # FIXME needs to be jobseekers_allowance] = 73.10
    hhres = init_household_result( joplings )
    intermed = make_intermediate(
        joplings,
        sys.hours_limits,
        sys.age_limits,
        sys.child_limits
    )
    hhres.housing = apply_rent_restrictions( joplings, intermed.hhint, sys.hr )
    calculateHB_CTR!(
        hhres,
        hb,
        joplings,
        intermed,
        sys.lmt,
        sys.age_limits,
        sys.nmt_bens )
    println( "Jopling: council $(joplings.council)")
    println( "MTBens for Joblings:\n$(hhres.bus[1].legacy_mtbens)\n" )
    println( "housing $(hhres.housing)")
    println( "Incomes $(hhres.bus[1].legacy_mtbens.hb_incomes)\n")
    @test hhres.bus[1].pers[head.pid].income[HOUSING_BENEFIT] ≈ 22.50

    mr_h = get_example(single_hh)
    mrhbu = get_benefit_units(mr_h)[1]
    head = get_head( mrhbu )
    println( "mr h:: initial: head.dla_self_care_type $(head.dla_self_care_type)" )
    head.dla_self_care_type = missing_lmh # turn off the dla in the spreadsheet
    retire!( head )
    empty!( head.income )
    head.age = 67
    head.income[private_pensions]=60.0
    head.income[state_pension]=127.75
    mr_h.gross_rent = 88.50
    hhres = init_household_result( mr_h )
    intermed = make_intermediate(
        mr_h,
        sys.hours_limits,
        sys.age_limits,
        sys.child_limits
        )
    println( "mr h:: after changes: head.dla_self_care_type $(head.dla_self_care_type)" )
    calc_pre_tax_non_means_tested!(
        hhres,
        mr_h,
        sys.nmt_bens,
        sys.hours_limits,
        sys.age_limits )

    hhres.bus[1].pers[head.pid].income[STATE_PENSION] = 127.75 # override calculated pension to march cpag example

    hhres.housing = apply_rent_restrictions( mr_h, intermed.hhint, sys.hr )
    calculateHB_CTR!(
        hhres,
        hb,
        mr_h,
        intermed,
        sys.lmt,
        sys.age_limits,
        sys.nmt_bens )
    println("Head's incomes:")
    println( inctostr( hhres.bus[1].pers[head.pid].income))
    println( "MTBens for Mr. H:\n$(hhres.bus[1].legacy_mtbens)\n" )
    println( "Incomes $(hhres.bus[1].legacy_mtbens.hb_incomes)\n")
    println( "premia $(hhres.bus[1].legacy_mtbens.premia)\n")
    @test length( hhres.bus[1].legacy_mtbens.premia ) == 0
    @test hhres.bus[1].legacy_mtbens.hb_allowances == 181.00
    println( "hhres.bus[1].pers[head.pid].income[HOUSING_BENEFIT]=$(hhres.bus[1].pers[head.pid].income[HOUSING_BENEFIT])")
    @test to_nearest_p( hhres.bus[1].pers[head.pid].income[HOUSING_BENEFIT], 84.11 )

    ## FIXME MORE TESTS NEEDED HERE: passporting NDDS
    # PENSION CREDIT
end

@testset "NDDs" begin
    sys = get_system( scotland=true )
    bu3 = get_example(mbu)
    bus = get_benefit_units( bu3 )
    @assert size(bus)[1] == 3
    head = get_head( bu3 )
    @assert head.age == 26
    bu2p = get_head( bus[2])
    employ!( bu2p )
    # cpag 19/0 ch 10 p194
    incomes = [500,380,300.0,250.0,175.0]
    ndds = [100.65,91.70,80.55,49.20,35.85,15.60]
    ntests = size( incomes)[1]
    for i in 1:ntests
        bu2p.income[wages] = incomes[i]
        println( bu2p.income )
        println( "bu2p.usual_hours_worked $(bu2p.usual_hours_worked)" )
        println( "hours_limits $(sys.lmt.hours_limits) ")
        bures = init_benefit_unit_result( Float64, bus[2])
        intermed = make_intermediate(
            bu3,
            sys.hours_limits,
            sys.age_limits,
            sys.child_limits
            )
        lmt_incomes = calc_incomes(
            hb,
            bus[2],
            bures,
            intermed.buint[2],
            sys.lmt.income_rules,
            sys.lmt.hours_limits )
        println( "lmt_incomes $(lmt_incomes)" )
        ndd = calc_NDDs(
            bus[2],
            bures,
            intermed.buint[2],
            lmt_incomes,
            sys.lmt.hb )
        @test ndd ≈ ndds[i]
    end
    unemploy!( bu2p )
    intermed = make_intermediate(
        bu3,
        sys.hours_limits,
        sys.age_limits,
        sys.child_limits
        )
    bures = init_benefit_unit_result( Float64, bus[2] )
    lmt_incomes = calc_incomes(
        hb,
        bus[2],
        bures,
        intermed.buint[2],
        sys.lmt.income_rules,
        sys.lmt.hours_limits )
    ndd = calc_NDDs(
        bus[2],
        bures,
        intermed.buint[2],
        lmt_incomes,
        sys.lmt.hb )
    @test ndd ≈ 15.60
end

@testset "Passporting" begin
    #TODO
end

@testset "CTC/WTC" begin
    # TODO
end

@testset "Full Legacy Benefits" begin
#TODO
    for (key,hhc) in get_all_examples()
        println( "on $key")
        head = get_head( hhc )
        intermed = make_intermediate(
            hhc,
            sys.hours_limits,
            sys.age_limits,
            sys.child_limits
            )
        hhres = init_household_result( hhc )
        calc_legacy_means_tested_benefits!(
            hhres,
            hhc,
            intermed,
            sys.lmt,
            sys.age_limits,
            sys.hours_limits,
            sys.nmt_bens,
            sys.hr )

        # head is unemployed
        unemploy!( head )
        intermed = make_intermediate(
            hhc,
            sys.hours_limits,
            sys.age_limits,
            sys.child_limits )
        hhres = init_household_result( hhc )
        calc_legacy_means_tested_benefits!(
            hhres,
            hhc,
            intermed,
            sys.lmt,
            sys.age_limits,
            sys.hours_limits,
            sys.nmt_bens,
            sys.hr )

        # head is seriously disabled
        disable_seriously!( head )
        intermed = make_intermediate(
            hhc,
            sys.hours_limits,
            sys.age_limits,
            sys.child_limits )
        hhres = init_household_result( hhc )
        calc_legacy_means_tested_benefits!(
            hhres,
            hhc,
            intermed,
            sys.lmt,
            sys.age_limits,
            sys.hours_limits,
            sys.nmt_bens,
            sys.hr )
        # head is an informal carer
        enable!( head )
        carer!( head )
        intermed = make_intermediate(
            hhc,
            sys.hours_limits,
            sys.age_limits,
            sys.child_limits )
        hhres = init_household_result( hhc )
        calc_legacy_means_tested_benefits!(
            hhres,
            hhc,
            intermed,
            sys.lmt,
            sys.age_limits,
            sys.hours_limits,
            sys.nmt_bens,
            sys.hr )

    end
end

@testset "PC/SC" begin
    # cpag19/20 EXAMPLES on p274
    sys = get_system( scotland=true )
    bhh= get_example(single_hh)
    barbara = get_head(bhh)
    retire!( barbara )
    disable_seriously!( barbara )
    barbara.age = 68
    bpid = barbara.pid
    barbara.pip_daily_living_type = enhanced_pip
    hhres = init_household_result( bhh )
    bres = hhres.bus[1].pers[bpid]
    bres.income .= 0.0
    bres.income[STATE_PENSION] = 129.20
    bres.income[ATTENDANCE_ALLOWANCE] = sys.nmt_bens.attendance_allowance.higher

    intermed = make_intermediate(
        bhh,
        sys.hours_limits,
        sys.age_limits,
        sys.child_limits )
    calc_legacy_means_tested_benefits!(
        hhres,
        bhh,
        intermed,
        sys.lmt,
        sys.age_limits,
        sys.hours_limits,
        sys.nmt_bens,
        sys.hr )
    lmt = hhres.bus[1].legacy_mtbens

    println( to_string( hhres ))

    @test lmt.can_apply_for.pc
    @test ! lmt.can_apply_for.sc # needs to be > pension age in 2016 - test make age 80 or
    @test ! lmt.can_apply_for.ctc
    @test ! lmt.can_apply_for.esa

    @test lmt.mig ≈ 233.10
    @test lmt.pc_premia ≈ 65.85
    @test hhres.bus[1].pers[bpid].income[PENSION_CREDIT] ≈ 103.90
    @test lmt.hb_passported
    @test lmt.ctr_passported

    m_and_j= get_example(childless_couple_hh)
    maria = get_head( m_and_j )
    maria.age = 66
    retire!( maria )
    maria.dla_self_care_type = high

    geoff = get_spouse( m_and_j )
    retire!( geoff )
    geoff.attendance_allowance_type = high
    geoff.age = 67
    geoff.assets[A_Stocks_Shares_Bonds_etc] = 11_500.0

    newpid = add_non_dependent!( m_and_j, 24, Female )
    bus = get_benefit_units( m_and_j )
    @test size( bus )[1] == 2
    daught = get_head(bus[2])
    hhres = init_household_result( m_and_j )

    mres = hhres.bus[1].pers[maria.pid]
    gres = hhres.bus[1].pers[geoff.pid]
    dres = hhres.bus[2].pers[daught.pid]

    mres.income .= 0
    mres.income[STATE_PENSION] = 77.45

    gres.income .= 0
    gres.income[STATE_PENSION] = 129.20
    gres.income[PRIVATE_PENSIONS] = 90

    dres.income .= 0
    dres.income[NON_CONTRIB_JOBSEEKERS_ALLOWANCE] = 100.0 # anything ..
    intermed = make_intermediate(
        m_and_j,
        sys.hours_limits,
        sys.age_limits,
        sys.child_limits )

    calc_legacy_means_tested_benefits!(
        hhres,
        m_and_j,
        intermed,
        sys.lmt,
        sys.age_limits,
        sys.hours_limits,
        sys.nmt_bens,
        sys.hr )
    lmt = hhres.bus[1].legacy_mtbens
    @test lmt.mig ≈ 255.25
    @test lmt.pc_premia ≈ 0
    @test hhres.bus[1].pers[maria.pid].income[PENSION_CREDIT] ≈ 0
    @test ! lmt.hb_passported
    @test ! lmt.ctr_passported
    @test lmt.pc_incomes.tariff_income ≈ 3
    @test lmt.pc_incomes.total_income ≈ 299.65

    # SAVINGS_CREDIT

    # from CPAG 2011/12 updated by me.

    t_and_j= get_example(childless_couple_hh)
    june = get_head( t_and_j )
    retire!( june )
    june.age = 75

    terry = get_spouse( t_and_j )
    retire!( terry )
    terry.age = 75 # should still qualify on June's age
    intermed = make_intermediate(
        t_and_j,
        sys.hours_limits,
        sys.age_limits,
        sys.child_limits )
    @test intermed.buint[1].someone_pension_age_2016

    hhres = init_household_result( t_and_j )
    jres = hhres.bus[1].pers[june.pid]
    jres.income[PRIVATE_PENSIONS] = 30.0
    tres = hhres.bus[1].pers[terry.pid]
    tres.income .= 0
    tres.income[STATE_PENSION] = 193.35
    calc_legacy_means_tested_benefits!(
        hhres,
        t_and_j,
        intermed,
        sys.lmt,
        sys.age_limits,
        sys.hours_limits,
        sys.nmt_bens,
        sys.hr )
    lmt = hhres.bus[1].legacy_mtbens
    @test lmt.can_apply_for.pc
    @test lmt.can_apply_for.sc
    @test lmt.mig ≈ 255.25
    @test lmt.pc_premia ≈ 0
    @test hhres.bus[1].pers[june.pid].income[SAVINGS_CREDIT] ≈ 0 # below cpl threshold 244.12 15.35
    @test lmt.hb_passported
    @test lmt.ctr_passported
    @test lmt.sc_incomes.total_income ≈ 0 # not calculated 193.35+30
    jres.income[PRIVATE_PENSIONS] = 60
    calc_legacy_means_tested_benefits!(
        hhres,
        t_and_j,
        intermed,
        sys.lmt,
        sys.age_limits,
        sys.hours_limits,
        sys.nmt_bens,
        sys.hr )
    lmt = hhres.bus[1].legacy_mtbens
    @test eq_nearest_p(hhres.bus[1].pers[june.pid].income[SAVINGS_CREDIT], 14.21 )
    jres.income[PRIVATE_PENSIONS] = 100
    calc_legacy_means_tested_benefits!(
        hhres,
        t_and_j,
        intermed,
        sys.lmt,
        sys.age_limits,
        sys.hours_limits,
        sys.nmt_bens,
        sys.hr )
    lmt = hhres.bus[1].legacy_mtbens
    @test eq_nearest_p(hhres.bus[1].pers[june.pid].income[SAVINGS_CREDIT], 0.11 )
end

@testset "CTC/WTC/Childcare CPAG Ch 61" begin
    relevant_period = 366
    cpag_wpy = 366/365.25
    sys = get_system( scotland=true )
    mt_ben_sys = sys.lmt

    tracyh = get_example(single_parent_hh)
    @test num_children( tracyh ) == 2
    for pid in child_pids( tracyh )
        tracyh.people[pid].hours_of_childcare = 40
        tracyh.people[pid].cost_of_childcare = 100.0
        tracyh.people[pid].age = 4
    end
    tracy = get_head( tracyh )
    tracy.usual_hours_worked = 20

    intermed = make_intermediate(
        tracyh,
        sys.hours_limits,
        sys.age_limits,
        sys.child_limits )
    bus = get_benefit_units( tracyh )
    bures = init_benefit_unit_result( Float64, bus[1])
    bures.pers[tracy.pid].income .= 0.0
    bures.pers[tracy.pid].income[WAGES] = 20*8

    #
    # needed since we use IT calculated incomes
    #
    calc_income_tax!(
        bures,
        tracy,
        nothing,
    sys.it )

    #
    # benefit unit level - this doesn't include ct/hb
    # which are done once for the household
    #
    calc_legacy_means_tested_benefits!(
            bures,
            bus[1],
            intermed.buint[1],
            sys.lmt,
            sys.age_limits,
            sys.hours_limits,
            sys.nmt_bens,
            tracyh )

    #
    # calculatint these from scratch
    # in since I can't work out the CPAG
    # DWP period stuff, which has variously 366 and 365
    # day years in their example, while I have 365.25 day
    # years
    #
    targetwtc = sys.lmt.working_tax_credit.basic +
        sys.lmt.working_tax_credit.lone_parent +
        bures.legacy_mtbens.cost_of_childcare
    targetwtc -= sys.lmt.working_tax_credit.taper*max( 0.0,
        bures.pers[tracy.pid].income[WAGES] - sys.lmt.working_tax_credit.threshold )

    @test bures.legacy_mtbens.cost_of_childcare ≈ 7280.0/52
    println( bures.legacy_mtbens.can_apply_for )
    println( to_md_table(intermed.buint[1]))
    println( to_md_table(bures))
    println( inctostr( bures.pers[tracy.pid].income ))
    targetctc = 2*(sys.lmt.child_tax_credit.child)
    @test bures.pers[tracy.pid].income[CHILD_TAX_CREDIT] ≈ targetctc
    targetwtc = sys.lmt.working_tax_credit.basic +
        sys.lmt.working_tax_credit.lone_parent +
        bures.legacy_mtbens.cost_of_childcare
    targetwtc -= sys.lmt.working_tax_credit.taper*max( 0.0,
        bures.pers[tracy.pid].income[WAGES] - sys.lmt.working_tax_credit.threshold )
    @test bures.pers[tracy.pid].income[WORKING_TAX_CREDIT] ≈ targetwtc

    bures.pers[tracy.pid].income .= 0.0
    bures.pers[tracy.pid].income[WAGES] = 20*10

    #
    # needed since we use IT calculated incomes
    #
    calc_income_tax!(
        bures,
        tracy,
        nothing,
    sys.it )

    #
    # benefit unit level - this doesn't include ct/hb
    # which are done once for the household
    #
    calc_legacy_means_tested_benefits!(
            bures,
            bus[1],
            intermed.buint[1],
            sys.lmt,
            sys.age_limits,
            sys.hours_limits,
            sys.nmt_bens,
            tracyh )

    #
    # calculatint these from scratch
    # in since I can't work out the CPAG
    # DWP period stuff, which has variously 366 and 365
    # day years in their example, while I have 365.25 day
    # years
    #
    targetwtc = sys.lmt.working_tax_credit.basic +
        sys.lmt.working_tax_credit.lone_parent +
        bures.legacy_mtbens.cost_of_childcare
    targetwtc -= sys.lmt.working_tax_credit.taper*max( 0.0,
        bures.pers[tracy.pid].income[WAGES] - sys.lmt.working_tax_credit.threshold )

    @test bures.legacy_mtbens.cost_of_childcare ≈ 7280.0/52
    println( bures.legacy_mtbens.can_apply_for )
    println( to_md_table(intermed.buint[1]))
    println( to_md_table(bures))
    println( inctostr( bures.pers[tracy.pid].income ))
    targetctc = 2*(sys.lmt.child_tax_credit.child)
    @test bures.pers[tracy.pid].income[CHILD_TAX_CREDIT] ≈ targetctc
    targetwtc = sys.lmt.working_tax_credit.basic +
        sys.lmt.working_tax_credit.lone_parent +
        bures.legacy_mtbens.cost_of_childcare
    targetwtc -= sys.lmt.working_tax_credit.taper*max( 0.0,
        bures.pers[tracy.pid].income[WAGES] - sys.lmt.working_tax_credit.threshold )
    @test bures.pers[tracy.pid].income[WORKING_TAX_CREDIT] ≈ targetwtc

    unemploy!( tracy )
    intermed = make_intermediate(
        tracyh,
        sys.hours_limits,
        sys.age_limits,
        sys.child_limits )
    calc_income_tax!(
        bures,
        tracy,
        nothing,
        sys.it )
    bures = init_benefit_unit_result( Float64, bus[1])
    bures.pers[tracy.pid].income .= 0.0
    calc_legacy_means_tested_benefits!(
        bures,
        bus[1],
        intermed.buint[1],
        sys.lmt,
        sys.age_limits,
        sys.hours_limits,
        sys.nmt_bens,
        tracyh )
    @test bures.pers[tracy.pid].income[CHILD_TAX_CREDIT] ≈ targetctc
    @test bures.pers[tracy.pid].income[WORKING_TAX_CREDIT] ≈ 0
    println( inctostr( bures.pers[tracy.pid].income ))
    println( ModelHousehold.to_string( tracyh ))

    # same at hhld level
    hhres = init_household_result( tracyh )
    hhres.bus[1].pers[tracy.pid].income .= 0.0
    calc_income_tax!(
        hhres.bus[1],
        tracy,
        nothing,
        sys.it )

    hhres.bus[1].pers[tracy.pid].income[LOCAL_TAXES] =
        calc_council_tax( tracyh, intermed.hhint, sys.loctax.ct )

    calc_legacy_means_tested_benefits!(
        hhres,
        tracyh,
        intermed,
        sys.lmt,
        sys.age_limits,
        sys.hours_limits,
        sys.nmt_bens,
        sys.hr
    )
    println( to_string(hhres))
    println( inctostr( hhres.bus[1].pers[tracy.pid].income ))
    @test hhres.bus[1].pers[tracy.pid].income[CHILD_TAX_CREDIT] ≈ targetctc
    @test hhres.bus[1].pers[tracy.pid].income[WORKING_TAX_CREDIT] ≈ 0

end

max_stars_repo_name: grahamstark/ScottishTaxBenefitModel

max_issues_repo_name: grahamstark/ScottishTaxBenefitModel

max_forks_repo_name: grahamstark/ScottishTaxBenefitModel


text:
### A Pluto.jl notebook ###
# v0.14.8

using Markdown
using InteractiveUtils

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

# ╔═╡ 6e3ef066-d115-11eb-2338-013a707dfe8a
begin
	using PlutoUI
	using DataFrames
	using Statistics
	using RollingFunctions
	using HTTP
	using JSON
	using Chain
	using Tables
	using PerceptualColourMaps
	using Gadfly
	using Match
	Gadfly.set_default_plot_size(16cm, 14cm)
end

# ╔═╡ 0da5a4ba-6b23-4f23-8165-1bd6183abbc1
html"<button onclick='present()'>Present</button>"

# ╔═╡ ba990146-058e-440e-b172-4d35b2b63985
md"""
## Looking Back

### Updates from last time
- Switched codebase to PyTorch + Huggingface
  - No more stability issues
  - Smaller memory footprint (can now use more machines)
  - Comes with helper functions for linguistic analysis
- Showed comparison between PPO with/without GPT on Breakout.
  - Both versions were comparable.

### Goals set last time
- How fast is new code?
  - How long to run a standard amount of steps for Atari?
- What is the published PPO benchmark on each Atari game?
- How is performance on other Atari games?
  - How do we compare to the baseline across Atari library?
- How does performance compare when using different numbers of embeddings for perception -> GPT input.
- How does model with pretrained GPT compare to model with randomly initialized GPT?
- Linguistic analysis
  - How similar are perception module's outputs to GPT's vocabulary of token embeddings?
  - For `n_embeddings > 1`, is the sequence of embeddings grammatical according to GPT?
    - (Does the sequence of embeddings have high probability under GPT's language model?)
"""

# ╔═╡ 731054d8-e5df-45d0-bd4f-91083884f0c8
md"""## What is the published PPO benchmark on each Atari game?
We have added a purple horizontal line to our graphs indicating the 10-million-step performance reported in the [PPO paper](https://arxiv.org/pdf/1707.06347.pdf).
"""

# ╔═╡ a15380dc-fb70-42ed-8025-77ba49dfb073
md"""
## How fast is the new code?
**Assessment:** The new code (PPO) runs at about the same speed as the old code (A2C).

The FPS for PPO is much lower (probably due to the high number of gradient updates performed per frame) but the PPO algorithm is more sample efficient.
"""

# ╔═╡ 5084f38e-b69a-4ee5-90c3-f23d1e6abad4
md"""
Here we compare the time to reach 10 million frames (the number of frames used in the PPO paper). As the graph demonstrates, A2C reaches 10 million frames in much less time.
"""

# ╔═╡ 76f3a217-3303-4952-ab6e-d546c8c05ffe
md"""
## Next Steps
- Lingustic Analysis
  - We need to re-run to record the necessary data. Any particular runs we should prioritize?
  - Correlation between particular state attributes and perception embeddings?
- Generalization-Focused Tasks ([BabyAI](https://github.com/mila-iqia/babyai), [ProcGen](https://openai.com/blog/procgen-benchmark/))
  - Finish implementation of few-shot learning algorithm
  - Implement learning-from-tokens algorithm
- Comparison with randomized parameters:
  - Search for good hyperparameters with randomized parameters
  - Compare with non-randomized on 5 Atari games
"""

# ╔═╡ 087ee9a9-e5e6-4607-a547-824d2fea470e
md"""
However, if we extend both lines, we see that the actual performance curves are nearly identcal, measured against time.

"""

# ╔═╡ 07b1a576-f7a1-404d-9043-2827da1f4d07
md"""
Furthermore, PPO learns much more quickly than A2C when performance is measured against time steps.
"""

# ╔═╡ 0be25163-f9c1-41bc-82cf-9f63ff8e609a
md"## Compare GPT with randomly initialized GPT"

# ╔═╡ 654ad80a-a802-49d5-8373-b0b4056fd8f1
md"""
The following graph compares the performance of our GPT architecture with and without randomized (but still frozen) GPT parameters. The randomized parameter architecture is still running, so results are currently inconclusive.
"""

# ╔═╡ ab745b36-164c-4534-be1f-a703a2010f3e
md"""
## Graph comparing GPT architectures by number of embeddings
The current architecture passes each observation through a series of CNNs and then reshapes the results into $n$ input embeddings to be fed into GPT, where $n$ is a hyperparameter.

The following graphs depict hyperparameter searches on the medium (355M parameter) GPT architecture with lines colored by $n$.

**Assessment:** There is not a clear linear correlation between the number of embeddings and performance, but it is with noting that 8-embedding architectures are among the top performers for all games except Seaquest.

Note that the first graph only runs to 5 million frames which explains the poor relative performance of the GPT architectures.
"""

# ╔═╡ 9b367184-ecfa-40b1-b439-d6f517d022eb
md"""
## Other Atari games
As the preceeding graphs demonstrate, hyperparameter search is underway on four Atari games besides Breakout. Once this search is complete, we will select a single, best-performing set of parameters to run on the complete Atari suite.
"""

# ╔═╡ 3838787a-f138-45ec-9a58-7879aeab9b99
md"""
## Linguistic Analysis
We are currently rerunning the best BeamRider run, saving
- the input observation
- the perception output
- the action logits
at set intervals.

We chose BeamRider because we significantly outperformed the baseline on this task. When the run is complete we will run linguistic analysis on the saved data.
"""

# ╔═╡ 8798b46c-1b0b-4d34-93ef-1d5062f8a632
md"""
## Architectures
This report compares the performance of a standard PPO baseline against a "GPT-in-the-middle" variant. The architectures are as follows:

#### PPO baseline

- Convolution with
  - output-size $32$
  - kernel-shape $8\times 8$
  - stride $4\times 4$
- ReLU
- Convolution with
  - output-size $64$
  - kernel-shape $4\times 4$
  - stride $2\times 2$
- ReLU
- Convolution with
  - output-size $64$
  - kernel-shape $3\times 3$
  - stride $1\times 1$
- ReLU
- Linear layer with output-size $512$
- Parallel value head and action head

#### GPT
(differences in **bold**)

- Convolution with
  - output-size $32$
  - kernel-shape $8\times 8$
  - stride $4\times 4$
- ReLU
- Convolution with
  - output-size $64$
  - kernel-shape $4\times 4$
  - stride $2\times 2$
- ReLU
- Convolution with
  - output-size $64$
  - kernel-shape $3\times 3$
  - stride $1\times 1$
- ReLU
- **Linear layer with output-size = $n \times e$ where $n$ is the number of embeddings (a hyperparameter) and $e$ is the GPT embedding size (see table below).**
- **GPT torso**
- **one `hidden_size` layer**
- one layer projecting to action and value

| Architecture | parameters | embedding size ($e$) |
|--------------|------------|----------------|
| small        | 117M       | 768            |
| medium       | 345M       | 1024           |
| large        | 774M       | 1280           |
| xl           | 1558M      | 1600           |
"""

# ╔═╡ 33909132-16c2-4eb3-97d3-607010a82e52
md"""
## Next Steps
- How to choose best hyperparameters (one set across all games for each `num_embeddings`).
- Graph comparing the speed of different GPT sizes (e.g. run 1 million samples for each size and see how long it takes).
- "Time Steps" -> "Steps"
- Display interesting hyperparameters for the runs below each graph (e.g. how many embs for the first few graphs in this notebook).
"""

# ╔═╡ 5027f577-6604-40e0-a40c-17110a1bc3fa
md"""
## TODO:
"""

# ╔═╡ 623ed8e1-7f9e-4092-bf39-b9a3d90a7d23
# Colormap Gallery: https://colorcet.com/gallery.html
colors(n) = cmap("I1"; N=n);

# ╔═╡ 16ce1ff4-3a0e-48e5-ae84-30afb618804c
begin
	num_embeddings = [1,2,4,6,8]
	color_scale = Scale.color_discrete_hue(colors, levels=reverse(num_embeddings))
end;

# ╔═╡ 5d3bfc57-6a5a-41e5-8f63-6902e4958936
begin
	struct HTMLDocument
		embedded
	end
	function Base.show(io::IO, mime::MIME"text/html", doc::HTMLDocument)
		println(io, "<html>")
		show(io, mime, doc.embedded)
		println(io, "</html>")
	end
end;

# ╔═╡ 9af550be-4ef4-4938-a710-41f96392b05a

# ╔═╡ 41bd06a2-e7bd-46c6-9249-6e69223b0e11
begin
	HASRUA_ENDPOINT_URL = "http://rldl12.eecs.umich.edu:1200/v1/graphql"
	HASURA_ADMIN_SECRET = "hNuxFSTJMk28GMiFDpZDbviKtelFQamcb20UJiYZSIZ4LYsyLCCwQFYA6Y8HRSUo"
end;

# ╔═╡ 94185c13-b6d2-4337-b5ce-336f5e128032
function gql_query(query:: String; variables:: Dict = nothing)
	r = HTTP.request(
		"POST",
		HASRUA_ENDPOINT_URL;
		verbose=3,
		body= JSON.json(Dict("query" => query, "variables" => variables)) ,
		headers=Dict("x-hasura-admin-secret" => HASURA_ADMIN_SECRET),
	)
	@chain r begin
		_.body
		String
		JSON.parse(_, null=missing)
		_["data"]
	end
end;

# ╔═╡ ce12c840-ece3-48bd-8696-40c1f9802695
for (x,y) in [1=>2]
	x
end

# ╔═╡ c5b77c36-436f-4682-8d52-945ae25fb47c
gadfly_theme = :default;

# ╔═╡ 0f54629c-e045-47bf-90a5-36e44c05b8f0
function sweep_runs(sweep_ids::AbstractVector{Int}, max_step::Int)
	query = """
		query getSweepRuns(\$ids: [Int!], \$max_step: Int!) {
		  logs_less_than_step(args: {max_step: \$max_step}, where: {run: {sweep_id: {_in: \$ids}}}) {
			log
			run_id
			run {
			  metadata
			  sweep_id
			}
		  }
		}
  	"""
	rows = @chain gql_query(query; variables=Dict("ids" => sweep_ids, "max_step" => max_step)) begin
		_["logs_less_than_step"]
		map(d -> Dict(
				"run_id" => d["run_id"],
				"sweep_id" => d["run"]["sweep_id"],
				d["log"]...,
				d["run"]["metadata"]["parameters"]...
				), _)
		map(d -> Dict(
				d...,
				[k => v for (k1, v1, k2, v2) in [
							(
								"hours", get(d, "time-delta", 0) / 3600,
								"time-delta", get(d, "hours", 0) * 3600,
							),
							(
								"env", get(d, "env_name", nothing),
								"env_name", get(d, "env", nothing),
							)
						]
						for (k, v) in [
								(k1, get(d, k1, v1)),
								(k2, get(d, k2, v2)),
								]]...,
				[name => get(d, name, false) for name in [
							"randomize_parameters"
						]]...,
				[name => get(d, name, nothing) for name in [
							"action_hidden_size",
							"gpt",
							"time",
							"gae",
							"gradient_clip",
							"nonlinearity",
							"normalize_observation",
							"normalize_torso_output",
							"optimizer",
							"num_embeddings",
							"save_interval",
							"save_path",
							"config",
							# This is where the compatibility with old gpt stuff starts:
							"action loss",
							"alpha",
							"clip_param",
							"cuda",
							"linguistic_analysis_path",
							"entropy",
							"eps",
							"eval_interval",
							"gae_lambda",
							"gpt_size",
							"linear_lr_decay",
							"lr",
							"max_grad_norm",
							"num_env_steps",
							"num_mini_batch",
							"num_processes",
							"ppo_epoch",
							"recurrent_policy",
							"use_proper_time_limits",
							"value loss",
							"value_coef",
							"async_envs",
							"charts_path",
							"episode length",
							"epsilon",
							"jit",
							"lambda_",
							"learning_rate",
							"logger",
							"loss",
							"num_envs",
							"resnet",
							"rnn",
							"training_steps",
						]]...
				), _)
		collect
	end
	vcat(DataFrame.(rows)...)

end;

# ╔═╡ 03f1dfc7-970e-4920-9df9-79dd9f048e65
function a2c_vs_ppo(a2c_steps, ppo_steps)
	a2c = sweep_runs([85], a2c_steps * 4)
	time_min = minimum(a2c[!, :time])
	seconds =  (a2c[!, :time] .- time_min) ./ 1000000
	hours =  seconds ./ 3600
	a2c[!, :hours] = hours
	a2c[!, :step] /= 4
	insertcols!(a2c, :algorithm=>"A2C")
	ppo = sweep_runs([672], ppo_steps)
	insertcols!(ppo, :algorithm=>"PPO")
	vcat(a2c, ppo)
end;

# ╔═╡ 2ab4c326-8ae0-40a3-9dc9-e5fd1cf17dfb
begin

	Gadfly.with_theme(gadfly_theme) do
		plot(
			a2c_vs_ppo(10000000, 10000000),
			x=:hours, y="episode return",
			yintercept=[274.8],
			group=:run_id,
			color="algorithm",
			Guide.xlabel("Hours"),
			Guide.ylabel("Episode Return"),
			Geom.line,
			Geom.hline(style=:dot,color="#b88fff"),
			Scale.color_discrete(colors),
			Guide.colorkey(title="Algorithm"),
			Guide.title("Breakout"),
			alpha=[0.5]
		) |> HTMLDocument
	end
end

# ╔═╡ d3352f83-5228-489a-9db6-b944fb2c1f72
begin

	Gadfly.with_theme(gadfly_theme) do
		plot(
			a2c_vs_ppo(80000000, 80000000),
			x=:hours, y="episode return",
			group=:run_id,
			color="algorithm",
			yintercept=[274.8],
			Geom.line,
			Geom.hline(style=:dot, color="#b88fff"),
			Guide.xlabel("Hours"),
			Guide.ylabel("Episode Return"),
			Scale.color_discrete(colors),
			Guide.colorkey(title="Algorithm"),
			Guide.title("Breakout"),
			alpha=[0.5]
		) |> HTMLDocument
	end
end

# ╔═╡ f05ae4d2-9b03-4013-a4d2-cb5ff3e4b61e
begin

	Gadfly.with_theme(gadfly_theme) do
		plot(
			a2c_vs_ppo(10000000, 10000000),
			x=:step, y="episode return",
			yintercept=[274.8],
			group=:run_id,
			color="algorithm",
			Guide.xlabel("Time Steps"),
			Guide.ylabel("Episode Return"),
			Geom.line,
			Geom.hline(style=:dot,color="#b88fff"),
			Scale.color_discrete(colors),
			Guide.colorkey(title="Algorithm"),
			Guide.title("Breakout"),
			alpha=[0.5]
		) |> HTMLDocument
	end
end

# ╔═╡ 67944dbb-0ebb-44d9-b6f3-79e8d5610f61
begin
	Gadfly.with_theme(gadfly_theme) do
		plot(
			sweep_runs([672, 687], 50000000),
			x=:step, y="episode return",
			yintercept=[274.8],
			group=:run_id,
			color=:randomize_parameters,
			Guide.xlabel("Step"),
			Guide.ylabel("Episode Return"),
			Geom.line,
			Geom.hline(style=:dot,color="#b88fff"),
			Scale.color_discrete(colors),
			Guide.colorkey(title="Random Parameters"),
			Guide.title("Breakout"),
			alpha=[0.5]
		) |> HTMLDocument
	end
end

# ╔═╡ f4368e0e-dd79-4a80-bae0-84a09de65e97
begin
	Gadfly.with_theme(gadfly_theme) do
		plot(
			sweep_runs([646], 100000000),
			x=:step, y="episode return",
			yintercept=[274.8],
			group=:run_id,
			color="num_embeddings",
			Guide.xlabel("Time Step"),
			Guide.ylabel("Episode Return"),
			Geom.line,
			Geom.hline(style=:dot,color="purple"),
			color_scale,
			Guide.colorkey(title="#embeddings"),
			Guide.title("Breakout"),
			alpha=[0.5]
		) |> HTMLDocument
	end
end

# ╔═╡ 9000d2d6-353d-499f-ae30-0338dc85bfe9
begin
	Gadfly.with_theme(gadfly_theme) do
		plot(
			sweep_runs([677], 50000000),
			x=:step, y="episode return",
			yintercept=[1204.5],
			group=:run_id,
			color="num_embeddings",
			Guide.xlabel("Time Step"),
			Guide.ylabel("Episode Return"),
			Geom.line,
			Geom.hline(style=:dot,color="#b88fff"),
			color_scale,
			Guide.colorkey(title="#embeddings"),
			Guide.title("Seaquest"),
			alpha=[0.5]
		) |> HTMLDocument
	end
end

# ╔═╡ 0a801028-935d-463a-b09c-01fca6eca913
begin
	Gadfly.with_theme(gadfly_theme) do
		plot(
			sweep_runs([676], 50000000),
			x=:step, y="episode return",
			yintercept=[14293.3],
			group=:run_id,
			color="num_embeddings",
			Guide.xlabel("Time Step"),
			Guide.ylabel("Episode Return"),
			Geom.line,
			Geom.hline(style=:dot,color="#b88fff"),
			color_scale,
			Guide.colorkey(title="#embeddings"),
			Guide.title("Qbert"),
			alpha=[0.5]
		) |> HTMLDocument
	end
end

# ╔═╡ 8420320c-222d-4144-9669-5f4f2a3b8b11
begin
	Gadfly.with_theme(gadfly_theme) do
	plot(
        sweep_runs([675], 50000000),
		x=:step, y="episode return",
		yintercept=[20.7],
		group=:run_id,
		color="num_embeddings",
		Guide.xlabel("Time Step"),
		Guide.ylabel("Episode Return"),
		Geom.line,
		Geom.hline(style=:dot,color="#b88fff"),
		color_scale,
		Guide.colorkey(title="#embeddings"),
		Guide.title("Pong"),
		alpha=[0.5]
	) |> HTMLDocument
	end
end

# ╔═╡ 3fd9c768-3dcc-4aa5-8671-c83ab39bd887
begin
	Gadfly.with_theme(gadfly_theme) do
		plot(
			sweep_runs([674], 50000000),
			x=:step, y="episode return",
			yintercept=[1590.0],
			group=:run_id,
			color="num_embeddings",
			Guide.xlabel("Time Step"),
			Guide.ylabel("Episode Return"),
			Geom.line,
			Geom.hline(style=:dot,color="#b88fff"),
			color_scale,
			Guide.colorkey(title="#embeddings"),
			Guide.title("BeamRider"),
			alpha=[0.5]
		) |> HTMLDocument
	end
end

# ╔═╡ bbed8889-d401-4167-9370-9927aabbc83b
# @bind window_size Select(string.(collapse_runs ? (100:100:500) : (1:10)), default=collapse_runs ? "200" : "5")
window_size = "11";

# ╔═╡ Cell order:
# ╟─0da5a4ba-6b23-4f23-8165-1bd6183abbc1
# ╟─ba990146-058e-440e-b172-4d35b2b63985
# ╠═731054d8-e5df-45d0-bd4f-91083884f0c8
# ╟─a15380dc-fb70-42ed-8025-77ba49dfb073
# ╟─5084f38e-b69a-4ee5-90c3-f23d1e6abad4
# ╟─76f3a217-3303-4952-ab6e-d546c8c05ffe
# ╟─2ab4c326-8ae0-40a3-9dc9-e5fd1cf17dfb
# ╟─087ee9a9-e5e6-4607-a547-824d2fea470e
# ╟─d3352f83-5228-489a-9db6-b944fb2c1f72
# ╟─07b1a576-f7a1-404d-9043-2827da1f4d07
# ╟─f05ae4d2-9b03-4013-a4d2-cb5ff3e4b61e
# ╟─03f1dfc7-970e-4920-9df9-79dd9f048e65
# ╟─0be25163-f9c1-41bc-82cf-9f63ff8e609a
# ╟─654ad80a-a802-49d5-8373-b0b4056fd8f1
# ╠═67944dbb-0ebb-44d9-b6f3-79e8d5610f61
# ╟─ab745b36-164c-4534-be1f-a703a2010f3e
# ╠═f4368e0e-dd79-4a80-bae0-84a09de65e97
# ╠═9000d2d6-353d-499f-ae30-0338dc85bfe9
# ╠═0a801028-935d-463a-b09c-01fca6eca913
# ╠═8420320c-222d-4144-9669-5f4f2a3b8b11
# ╠═3fd9c768-3dcc-4aa5-8671-c83ab39bd887
# ╟─9b367184-ecfa-40b1-b439-d6f517d022eb
# ╟─3838787a-f138-45ec-9a58-7879aeab9b99
# ╟─8798b46c-1b0b-4d34-93ef-1d5062f8a632
# ╠═33909132-16c2-4eb3-97d3-607010a82e52
# ╠═5027f577-6604-40e0-a40c-17110a1bc3fa
# ╠═16ce1ff4-3a0e-48e5-ae84-30afb618804c
# ╠═623ed8e1-7f9e-4092-bf39-b9a3d90a7d23
# ╠═5d3bfc57-6a5a-41e5-8f63-6902e4958936
# ╠═6e3ef066-d115-11eb-2338-013a707dfe8a
# ╟─9af550be-4ef4-4938-a710-41f96392b05a
# ╠═41bd06a2-e7bd-46c6-9249-6e69223b0e11
# ╠═94185c13-b6d2-4337-b5ce-336f5e128032
# ╠═ce12c840-ece3-48bd-8696-40c1f9802695
# ╠═c5b77c36-436f-4682-8d52-945ae25fb47c
# ╠═0f54629c-e045-47bf-90a5-36e44c05b8f0
# ╟─bbed8889-d401-4167-9370-9927aabbc83b

ext: jl

max_stars_repo_name: GPT-RL/generalization

max_issues_repo_name: GPT-RL/generalization

max_forks_repo_name: GPT-RL/generalization


text:
import StatsFuns: log2π

@node NormalMeanPrecision Stochastic [ out, (μ, aliases = [ mean ]), (τ, aliases = [ invcov, precision ]) ]

conjugate_type(::Type{ <: NormalMeanPrecision }, ::Type{ Val{ :out } }) = NormalMeanPrecision
conjugate_type(::Type{ <: NormalMeanPrecision }, ::Type{ Val{ :μ } })   = NormalMeanPrecision
conjugate_type(::Type{ <: NormalMeanPrecision }, ::Type{ Val{ :τ } })   = Gamma

@average_energy NormalMeanPrecision (q_out::Any, q_μ::Any, q_τ::Any) = begin
    μ_mean, μ_var     = mean(q_μ), var(q_μ)
    out_mean, out_var = mean(q_out), var(q_out)
    return 0.5 * (log2π - logmean(q_τ) + mean(q_τ) * (μ_var + out_var + abs2(μ_mean - out_mean)))
end


text:
function DiffEqBase.__solve(prob::DiffEqBase.DEProblem,args...;default_set=false,kwargs...)
  if length(args)>0 && !(typeof(args[1]) <: DiffEqBase.DEAlgorithm)
    error("Inappropiate solve command. The arguments do not make sense. Likely, you gave an algorithm which does not actually exist (or does not <:DiffEqBase.DEAlgorithm)")
  else
    DiffEqBase.__solve(prob::DiffEqBase.DEProblem,nothing,args...;default_set=false,kwargs...)
  end
end

function DiffEqBase.__solve(prob::DiffEqBase.DEProblem,alg::Union{Nothing,DiffEqBase.DEAlgorithm},
               args...;default_set=false,kwargs...)
  if default_set == true
    error("The chosen algorithm, $alg, does not exist. Please verify that the appropriate solver package has been installed.")
  end
  alg,extra_kwargs = default_algorithm(prob;kwargs...)
  DiffEqBase.__solve(prob,alg,args...;default_set=true,kwargs...,extra_kwargs...)
end

function DiffEqBase.__solve(prob::DiffEqJump.JumpProblem,alg::Nothing,
               args...;default_set=false,kwargs...)
  if default_set == true
    error("The chosen algorithm, $alg, does not exist. Please verify that the appropriate solver package has been installed.")
  end
  alg,extra_kwargs = default_algorithm(prob.prob;kwargs...)
  DiffEqBase.__solve(prob,alg,args...;default_set=true,kwargs...,extra_kwargs...)
end

export default_algorithm, DefaultAlgorithm


text:
module Scale

using Color
using Compat
using Compose
using DataArrays
using DataStructures
using Gadfly
using Showoff

import Gadfly: element_aesthetics, isconcrete, concrete_length,
               nonzero_length
import Distributions: Distribution

include("color.jl")

# Return true if var is categorical.
function iscategorical(scales::Dict{Symbol, Gadfly.ScaleElement}, var::Symbol)
    return haskey(scales, var) && isa(scales[var], DiscreteScale)
end

# Apply some scales to data in the given order.
#
# Args:
#   scales: An iterable object of ScaleElements.
#   aess: Aesthetics (of the same length as datas) to update with scaled data.
#   datas: Zero or more data objects. (Yes, I know "datas" is not a real word.)
#
# Returns:
#   nothing
#
function apply_scales(scales,
                      aess::Vector{Gadfly.Aesthetics},
                      datas::Gadfly.Data...)
    for scale in scales
        apply_scale(scale, aess, datas...)
    end

    for (aes, data) in zip(aess, datas)
        aes.titles = data.titles
    end
end

# Apply some scales to data in the given order.
#
# Args:
#   scales: An iterable object of ScaleElements.
#   datas: Zero or more data objects.
#
# Returns:
#   A vector of Aesthetics of the same length as datas containing scaled data.
#
function apply_scales(scales, datas::Gadfly.Data...)
    aess = Gadfly.Aesthetics[Gadfly.Aesthetics() for _ in datas]
    apply_scales(scales, aess, datas...)
    aess
end

# Transformations on continuous scales
immutable ContinuousScaleTransform
    f::Function     # transform function
    finv::Function  # f's inverse

    # A function taking one or more values and returning an array of
    # strings.
    label::Function
end

function identity_formatter(xs::AbstractArray, format=:auto)
    return showoff(xs, format)
end

const identity_transform =
    ContinuousScaleTransform(identity, identity, identity_formatter)

function log10_formatter(xs::AbstractArray, format=:plain)
    [@sprintf("10%s", x) for x in showoff(xs, format)]
end

const log10_transform =
    ContinuousScaleTransform(log10, x -> 10^x, log10_formatter)

function log2_formatter(xs::AbstractArray, format=:plain)
    [@sprintf("2%s", x) for x in showoff(xs, format)]
end

const log2_transform =
    ContinuousScaleTransform(log2, x -> 2^x, log2_formatter)

function ln_formatter(xs::AbstractArray, format=:plain)
    [@sprintf("e%s", x) for x in showoff(xs, format)]
end

const ln_transform =
    ContinuousScaleTransform(log, exp, ln_formatter)

function asinh_formatter(xs::AbstractArray, format=:plain)
    [@sprintf("asinh(%s)", x) for x in showoff(xs, format)]
end

const asinh_transform =
    ContinuousScaleTransform(asinh, sinh, asinh_formatter)

function sqrt_formatter(xs::AbstractArray, format=:plain)
    [@sprintf("%s2", x) for x in showoff(xs, format)]
end

const sqrt_transform = ContinuousScaleTransform(sqrt, x -> x^2, sqrt_formatter)

# Continuous scale maps data on a continuous scale simple by calling
# `convert(Float64, ...)`.
immutable ContinuousScale <: Gadfly.ScaleElement
    vars::Vector{Symbol}
    trans::ContinuousScaleTransform
    minvalue
    maxvalue
    minticks
    maxticks
    labels::Union(Nothing, Function)
    format
    scalable

    function ContinuousScale(vars::Vector{Symbol},
                             trans::ContinuousScaleTransform;
                             labels=nothing,
                             minvalue=nothing, maxvalue=nothing,
                             minticks=2, maxticks=10,
                             format=nothing,
                             scalable=true)
        if minvalue != nothing && maxvalue != nothing && minvalue > maxvalue
            error("Cannot construct a ContinuousScale with minvalue > maxvalue")
        end
        new(vars, trans, minvalue, maxvalue, minticks, maxticks, labels,
            format, scalable)
    end
end

function make_labeler(scale::ContinuousScale)
    if scale.labels != nothing
        function f(xs)
            return [scale.labels(x) for x in xs]
        end
    elseif scale.format == nothing
        scale.trans.label
    else
        function f(xs)
            return scale.trans.label(xs, scale.format)
        end
    end
end

const x_vars = [:x, :xmin, :xmax, :xintercept, :xviewmin, :xviewmax]
const y_vars = [:y, :ymin, :ymax, :yintercept, :middle,
                :upper_fence, :lower_fence, :upper_hinge, :lower_hinge,
                :yviewmin, :yviewmax]

function continuous_scale_partial(vars::Vector{Symbol},
                                  trans::ContinuousScaleTransform)
    function f(; minvalue=nothing, maxvalue=nothing, labels=nothing, format=nothing, minticks=2,
                 maxticks=10, scalable=true)
        ContinuousScale(vars, trans, minvalue=minvalue, maxvalue=maxvalue,
                        labels=labels, format=format, minticks=minticks,
                        maxticks=maxticks, scalable=scalable)
    end
end

# Commonly used scales.
const x_continuous = continuous_scale_partial(x_vars, identity_transform)
const y_continuous = continuous_scale_partial(y_vars, identity_transform)
const x_log10      = continuous_scale_partial(x_vars, log10_transform)
const y_log10      = continuous_scale_partial(y_vars, log10_transform)
const x_log2       = continuous_scale_partial(x_vars, log2_transform)
const y_log2       = continuous_scale_partial(y_vars, log2_transform)
const x_log        = continuous_scale_partial(x_vars, ln_transform)
const y_log        = continuous_scale_partial(y_vars, ln_transform)
const x_asinh      = continuous_scale_partial(x_vars, asinh_transform)
const y_asinh      = continuous_scale_partial(y_vars, asinh_transform)
const x_sqrt       = continuous_scale_partial(x_vars, sqrt_transform)
const y_sqrt       = continuous_scale_partial(y_vars, sqrt_transform)

const size_continuous = continuous_scale_partial([:size], identity_transform)

function element_aesthetics(scale::ContinuousScale)
    return scale.vars
end

# Apply a continuous scale.
#
# Args:
#   scale: A continuos scale.
#   datas: Zero or more data objects.
#   aess: Aesthetics (of the same length as datas) to update with scaled data.
#
# Return:
#   nothing
#
function apply_scale(scale::ContinuousScale,
                     aess::Vector{Gadfly.Aesthetics}, datas::Gadfly.Data...)
    for (aes, data) in zip(aess, datas)
        for var in scale.vars
            vals = getfield(data, var)
            if vals === nothing
                continue
            end

            # special case for function arrays bound to :y
            # pass the function values through and wait for the scale to
            # be reapplied by Stat.func
            if var == :y && eltype(vals) == Function
                aes.y = vals
                continue
            end

            # special case for Distribution values bound to :x or :y. wait for
            # scale to be re-applied by Stat.qq
            if in(var, [:x, :y]) && typeof(vals) <: Distribution
                setfield!(aes, var, vals)
                continue
            end

            T = Any
            for (i, d) in enumerate(vals)
                if isconcrete(d)
                    T = typeof(scale.trans.f(d))
                    break
                end
            end

            hasna = false
            if isa(vals, AbstractDataArray)
                for d in vals
                    if isna(d)
                        hasna = true
                        break
                    end
                end
            end

            ds = hasna ? DataArray(T, length(vals)) : Array(T, length(vals))
            apply_scale_typed!(ds, vals, scale)

            if var == :xviewmin || var == :xviewmax ||
               var == :yviewmin || var == :yviewmax
                setfield!(aes, var, ds[1])
            else
                setfield!(aes, var, ds)
            end

            if var in x_vars
                label_var = :x_label
            elseif var in y_vars
                label_var = :y_label
            else
                label_var = symbol(@sprintf("%s_label", string(var)))
            end

            if in(label_var, Set(fieldnames(aes)))
                setfield!(aes, label_var, make_labeler(scale))
            end
        end

        if scale.minvalue != nothing
            if scale.vars === x_vars
                aes.xviewmin = scale.trans.f(scale.minvalue)
            elseif scale.vars === y_vars
                aes.yviewmin = scale.trans.f(scale.minvalue)
            end
        end

        if scale.maxvalue != nothing
            if scale.vars === x_vars
                aes.xviewmax = scale.trans.f(scale.maxvalue)
            elseif scale.vars === y_vars
                aes.yviewmax = scale.trans.f(scale.maxvalue)
            end
        end
    end
end

function apply_scale_typed!(ds, field, scale::ContinuousScale)
    for (i, d) in enumerate(field)
        if isconcrete(d)
            ds[i] = scale.trans.f(d)
            i += 1
        else
            ds[i] = d
        end
    end
end

# Reorder the levels of a pooled data array
function reorder_levels(da::PooledDataArray, order::AbstractVector)
    level_values = levels(da)
    if length(order) != length(level_values)
        error("Discrete scale order is not of the same length as the data's levels.")
    end
    permute!(level_values, order)
    return PooledDataArray(da, level_values)
end

function discretize_make_pda(values::Vector, levels=nothing)
    if levels == nothing
        return PooledDataArray(values)
    else
        return PooledDataArray(convert(Vector{eltype(levels)}, values), levels)
    end
end

function discretize_make_pda(values::DataArray, levels=nothing)
    if levels == nothing
        return PooledDataArray(values)
    else
        return PooledDataArray(convert(DataArray{eltype(levels)}, values), levels)
    end
end

function discretize_make_pda(values::Range, levels=nothing)
    if levels == nothing
        return PooledDataArray(collect(values))
    else
        return PooledDataArray(collect(values), levels)
    end
end

function discretize_make_pda(values::PooledDataArray, levels=nothing)
    if levels == nothing
        return values
    else
        return PooledDataArray(values, convert(Vector{eltype(values)}, levels))
    end
end

function discretize(values, levels=nothing, order=nothing,
                    preserve_order=true)
    if levels == nothing
        if preserve_order
            levels = OrderedSet()
            for value in values
                push!(levels, value)
            end
            da = discretize_make_pda(values, collect(eltype(values), levels))
        else
            da = discretize_make_pda(values)
        end
    else
        da = discretize_make_pda(values, levels)
    end

    if order != nothing
        return reorder_levels(da, order)
    else
        return da
    end
end

immutable DiscreteScaleTransform
    f::Function
end

immutable DiscreteScale <: Gadfly.ScaleElement
    vars::Vector{Symbol}

    # Labels are either a function that takes an array of values and returns
    # an array of string labels, a vector of string labels of the same length
    # as the number of unique values in the discrete data, or nothing to use
    # the default labels.
    labels::Union(Nothing, Function)

    # If non-nothing, give values for the scale. Order will be respected and
    # anything in the data that's not represented in values will be set to NA.
    levels::Union(Nothing, AbstractVector)

    # If non-nothing, a permutation of the pool of values.
    order::Union(Nothing, AbstractVector)

    function DiscreteScale(vals::Vector{Symbol};
                           labels=nothing, levels=nothing, order=nothing)
        new(vals, labels, levels, order)
    end
end

const discrete = DiscreteScale

element_aesthetics(scale::DiscreteScale) = scale.vars

function x_discrete(; labels=nothing, levels=nothing, order=nothing)
    return DiscreteScale(x_vars, labels=labels, levels=levels, order=order)
end

function y_discrete(; labels=nothing, levels=nothing, order=nothing)
    return DiscreteScale(y_vars, labels=labels, levels=levels, order=order)
end

function group_discrete(; labels=nothing, levels=nothing, order=nothing)
    return DiscreteScale([:group], labels=labels, levels=levels, order=order)
end

function apply_scale(scale::DiscreteScale, aess::Vector{Gadfly.Aesthetics},
                     datas::Gadfly.Data...)
    for (aes, data) in zip(aess, datas)
        for var in scale.vars
            label_var = symbol(@sprintf("%s_label", string(var)))

            if getfield(data, var) === nothing
                continue
            end

            disc_data = discretize(getfield(data, var), scale.levels, scale.order)

            setfield!(aes, var, PooledDataArray(round(Int64, disc_data.refs)))

            # The leveler for discrete scales is a closure over the discretized data.
            if scale.labels === nothing
                function default_labeler(xs)
                    lvls = levels(disc_data)
                    vals = Any[1 <= x <= length(lvls) ? lvls[x] : "" for x in xs]
                    if all([isa(val, FloatingPoint) for val in vals])
                        return showoff(vals)
                    else
                        return [string(val) for val in vals]
                    end
                end

                labeler = default_labeler
            else
                function explicit_labeler(xs)
                    lvls = levels(disc_data)
                    return [string(scale.labels(lvls[x])) for x in xs]
                end

                labeler = explicit_labeler
            end

            if in(label_var, Set(fieldnames(aes)))
                setfield!(aes, label_var, labeler)
            end
        end
    end
end

immutable NoneColorScale <: Gadfly.ScaleElement
end

const color_none = NoneColorScale

function element_aesthetics(scale::NoneColorScale)
    [:color]
end

function apply_scale(scale::NoneColorScale,
                     aess::Vector{Gadfly.Aesthetics}, datas::Gadfly.Data...)
    for aes in aess
        aes.color = nothing
    end
end

immutable DiscreteColorScale <: Gadfly.ScaleElement
    f::Function # A function f(n) that produces a vector of n colors.

    # If non-nothing, give values for the scale. Order will be respected and
    # anything in the data that's not represented in values will be set to NA.
    levels::Union(Nothing, AbstractVector)

    # If non-nothing, a permutation of the pool of values.
    order::Union(Nothing, AbstractVector)

    # If true, order levels as they appear in the data
    preserve_order::Bool

    function DiscreteColorScale(f::Function; levels=nothing, order=nothing,
                                preserve_order=true)
        new(f, levels, order, preserve_order)
    end
end

function element_aesthetics(scale::DiscreteColorScale)
    [:color]
end

# Common discrete color scales
function color_discrete_hue(; levels=nothing, order=nothing,
                            preserve_order=true)
    DiscreteColorScale(
        h -> convert(Vector{ColorValue},
             distinguishable_colors(h, ColorValue[LCHab(70, 60, 240)],
                                    transform=c -> deuteranopic(c, 0.5),
                                    lchoices=Float64[65, 70, 75, 80],
                                    cchoices=Float64[0, 50, 60, 70],
                                    hchoices=linspace(0, 330, 24))),
        levels=levels, order=order, preserve_order=preserve_order)
end

@deprecate discrete_color_hue(; levels=nothing, order=nothing, preserve_order=true) color_discrete_hue(; levels=levels, order=order, preserve_order=preserve_order)

const color_discrete = color_discrete_hue

@deprecate discrete_color(; levels=nothing, order=nothing, preserve_order=true) color_discrete(; levels=levels, order=order, preserve_order=preserve_order)

function color_discrete_manual(colors...; levels=nothing, order=nothing)
    cs = ColorValue[color(c) for c in colors]
    function f(n)
        distinguishable_colors(n, cs)
    end
    DiscreteColorScale(f, levels=levels, order=order)
end

@deprecate discrete_color_manual(colors...; levels=nothing, order=nothing) color_discrete_manual(colors...; levels=levels, order=order)

function apply_scale(scale::DiscreteColorScale,
                     aess::Vector{Gadfly.Aesthetics}, datas::Gadfly.Data...)
    levelset = OrderedSet()
    for (aes, data) in zip(aess, datas)
        if data.color === nothing
            continue
        end
        for d in data.color
            if !isna(d)
                push!(levelset, d)
            end
        end
    end

    if scale.levels == nothing
        scale_levels = [levelset...]
        if !scale.preserve_order
            sort!(scale_levels)
        end
    else
        scale_levels = scale.levels
    end
    if scale.order != nothing
        permute!(scale_levels, scale.order)
    end
    colors = convert(Vector{RGB{Float32}}, scale.f(length(scale_levels)))

    color_map = @compat Dict([color => string(label)
                              for (label, color) in zip(scale_levels, colors)])
    function labeler(xs)
        [color_map[x] for x in xs]
    end

    for (aes, data) in zip(aess, datas)
        if data.color === nothing
            continue
        end
        ds = discretize(data.color, scale_levels)
        colorvals = Array(RGB{Float32}, nonzero_length(ds.refs))
        i = 1
        for k in ds.refs
            if k != 0
                colorvals[i] = colors[k]
                i += 1
            end
        end

        colored_ds = PooledDataArray(colorvals, colors)
        aes.color = colored_ds

        aes.color_label = labeler
        aes.color_key_colors = OrderedDict()
        for (i, c) in enumerate(colors)
            aes.color_key_colors[c] = i
        end
    end
end

immutable ContinuousColorScale <: Gadfly.ScaleElement
    # A function of the form f(p) where 0 <= p <= 1, that returns a color.
    f::Function

    minvalue
    maxvalue

    function ContinuousColorScale(f::Function; minvalue=nothing, maxvalue=nothing)
        new(f, minvalue, maxvalue)
    end
end

element_aesthetics(::ContinuousColorScale) = [:color]

function color_continuous_gradient(;minvalue=nothing, maxvalue=nothing)

    # TODO: this should be made more general purpose. I.e. define some
    # more color scales.
    function lch_diverge2(l0=30, l1=100, c=40, h0=260, h1=10, hmid=20, power=1.5)
        lspan = l1 - l0
        hspan1 = hmid - h0
        hspan0 = h1 - hmid
        function f(r)
            r2 = 2r - 1
            return LCHab(min(80, l1 - lspan * abs(r2)^power), max(10, c * abs(r2)),
                         (1-r)*h0 + r * h1)
        end
    end

    ContinuousColorScale(
        lch_diverge2(),
        minvalue=minvalue, maxvalue=maxvalue)
end

@deprecate continuous_color_gradient(;minvalue=nothing, maxvalue=nothing) color_continuous_gradient(;minvalue=minvalue, maxvalue=maxvalue)

const color_continuous = color_continuous_gradient

@deprecate continuous_color(;minvalue=nothing, maxvalue=nothing) color_continuous(;minvalue=nothing, maxvalue=nothing)

function apply_scale(scale::ContinuousColorScale,
                     aess::Vector{Gadfly.Aesthetics}, datas::Gadfly.Data...)
    cmin = Inf
    cmax = -Inf
    for data in datas
        if data.color === nothing
            continue
        end

        for c in data.color
            if c === NA
                continue
            end

            c = convert(Float64, c)
            if c < cmin
                cmin = c
            end

            if c > cmax
                cmax = c
            end
        end
    end

    if cmin == Inf || cmax == -Inf
        return nothing
    end

    if scale.minvalue != nothing
        cmin = scale.minvalue
    end

    if scale.maxvalue  != nothing
        cmax = scale.maxvalue
    end

    cmin, cmax = promote(cmin, cmax)

    ticks, viewmin, viewmax = Gadfly.optimize_ticks(cmin, cmax)
    if ticks[1] == 0 && cmin >= 1
        ticks[1] = 1
    end

    cmin = ticks[1]
    cmax = ticks[end]
    cspan = cmax != cmin ? cmax - cmin : 1.0

    for (aes, data) in zip(aess, datas)
        if data.color === nothing
            continue
        end

        aes.color = DataArray(RGB{Float32}, length(data.color))
        apply_scale_typed!(aes.color, data.color, scale, cmin, cspan)

        color_key_colors = Dict{ColorValue, Float64}()
        color_key_labels = Dict{ColorValue, String}()

        tick_labels = identity_formatter(ticks)
        for (i, j, label) in zip(ticks, ticks[2:end], tick_labels)
            r = (i - cmin) / cspan
            c = scale.f(r)
            color_key_colors[c] = r
            color_key_labels[c] = label
        end
        c = scale.f((ticks[end] - cmin) / cspan)
        color_key_colors[c] = (ticks[end] - cmin) / cspan
        color_key_labels[c] = tick_labels[end]

        function labeler(xs)
            [get(color_key_labels, x, "") for x in xs]
        end

        aes.color_function = scale.f
        aes.color_label = labeler
        aes.color_key_colors = color_key_colors
        aes.color_key_continuous = true
    end
end

function apply_scale_typed!(ds, field, scale::ContinuousColorScale,
                            cmin::Float64, cspan::Float64)
    for (i, d) in enumerate(field)
        if isconcrete(d)
            ds[i] = convert(RGB{Float32}, scale.f((convert(Float64, d) - cmin) / cspan))
        else
            ds[i] = NA
        end
    end
end

# Label scale is always discrete, hence we call it 'label' rather
# 'label_discrete'.
immutable LabelScale <: Gadfly.ScaleElement
end

function apply_scale(scale::LabelScale,
                     aess::Vector{Gadfly.Aesthetics}, datas::Gadfly.Data...)
    for (aes, data) in zip(aess, datas)
        if data.label === nothing
            continue
        end

        aes.label = discretize(data.label)
    end
end

element_aesthetics(::LabelScale) = [:label]

const label = LabelScale

# Scale applied to grouping aesthetics.
immutable GroupingScale <: Gadfly.ScaleElement
    var::Symbol
end

function xgroup(; labels=nothing, levels=nothing, order=nothing)
    return DiscreteScale([:xgroup], labels=labels, levels=levels, order=order)
end

function ygroup(; labels=nothing, levels=nothing, order=nothing)
    return DiscreteScale([:ygroup], labels=labels, levels=levels, order=order)
end

# Catchall scale for when no transformation of the data is necessary
immutable IdentityScale <: Gadfly.ScaleElement
    var::Symbol
end

function element_aesthetics(scale::IdentityScale)
    return [scale.var]
end

function apply_scale(scale::IdentityScale,
                     aess::Vector{Gadfly.Aesthetics}, datas::Gadfly.Data...)
    for (aes, data) in zip(aess, datas)
        if getfield(data, scale.var) === nothing
            continue
        end

        setfield!(aes, scale.var, getfield(data, scale.var))
    end
end

function z_func()
    return IdentityScale(:z)
end

function y_func()
    return IdentityScale(:y)
end

function x_distribution()
    return IdentityScale(:x)
end

function y_distribution()
    return IdentityScale(:y)
end

end # module Scale


text:
# ------------------------------------------------------------------- #
# Constant Volatility Exogenous Process
# ------------------------------------------------------------------- #
"""
This is the type that holds the parameters for a constant volatility
exogenous process from the BCFL paper. Namely the process follows:

```
\log(\hat{z}_{t+1}) = (1 - 2γ) \log(\hat{z}_{t}) + (σz1/2) ϵ1 - (σz2/2) ϵ2
\log(g_{t+1}) = (σz1/2) ϵ1 + (σz2/2) ϵ2
```

The number of epsilons we use for quadrature in each i is given by nϵi
"""
immutable ExogCV <: AbstractExog
    γ::Float64
    σz1::Float64
    σz2::Float64
    nϵ1::Int
    nϵ2::Int
end

# Small self-explanatory helper functions
init(::ExogCV) = [0.0, 0.0]
n_exog(::ExogCV) = 2
_unpack(e::ExogCV) = (e.γ, e.σz1, e.σz2, e.nϵ1, e.nϵ2)

const _cv_exog_ordering_docstr_extra = """
Because lzhp is a time t+1 state variable, but lgp is not, the ordering of
variables here is significant. Specifically, at points in the code it is assumed
that lzhp comes first and lgp comes last. This is also true in the stochastic
volatiltiy version of the model where there there is an additional state `vp`
placed between the other two.
"""

"""
Takes one step for exogenous process, given current levels and innovations

$_cv_exog_ordering_docstr_extra
"""
function exog_step(exog::ExogCV, lzh, ϵ1, ϵ2)
    γ, σz1, σz2, nϵ1, nϵ2 = _unpack(exog)
    lzhp = (1-2γ)*lzh .+ 0.5*(σz1.*ϵ1 .- σz2*ϵ2)
    lgp = 0.0*lzh .+ 0.5*σz1.*ϵ1 .+ 0.5*σz2*ϵ2

    lzhp, lgp
end

# API method on above
exog_step(exog::ExogCV, current::Vector, ϵ::Vector) =
    exog_step(exog, current[1], ϵ[1], ϵ[2])

# ------------------------------------------------------------------- #
# Constant Volatility Model
# ------------------------------------------------------------------- #
type ModelCV <: AbstractModel
    agent1::Agent
    agent2::Agent
    exog::ExogCV
    m_ls::Tuple{LinSpace{Float64}, LinSpace{Float64}}
    grid::Matrix{Float64}
    grid_transpose::Matrix{Float64}
    lλ☆p_pol::Matrix{Float64}
    ϵ::Matrix{Float64}
    Π::Vector{Float64}
    lzhp::Matrix{Float64}
    lgp::Matrix{Float64}
    a1_itp::ScaledInterpolation
    b2_itp::ScaledInterpolation
    other::OtherItems

    """
    Given the paramters of the model as well as information on the coarseness
    of the grids, do initial setup work to define the model. This involves
    the following steps:

    - Construct instances of `Agent` to hold preference paramters for
      both agents
    - Construct an object to hold the paramters of the exogenous process
    - Define grids over state variables (lλ☆, lzh)
    - Define an initial guess for the policy of lλ☆p
    - Build quadrature weights and nodes to approximate the expectations
    - Precompute solutions the allocation in terms of lλ☆ and lzh

    Note that all parameters match the notation from the paper and are settable
    one at a time using keyword argument syntax (`ModelCV(β2=0.99)`) changes
    only β2 relative to the default values. The only notational exception is
    that the standard deviation of the shock to `zi` is `σzi` here and
    `sqrt(vi)` in the paper.

    Other arguments to this constructor are:

    - `lλ☆bds`: bound on grid for lλ☆. The grid is assumed to be symmetric,
    extending from `-lλ☆bds` to `lλ☆bds`.
    - `nX` for various `X`: the number of grid points for variable `X`
    - `σz1` and `σz2`: standard deviation of shocks to `lz1` and `lz2` (see above)
    """
    function ModelCV(;ρ1::Real=-1.0, α1::Real=-9.0, β1::Real=0.98,
                      ω1::Real=0.1, σ1::Real=0.0,
                      ρ2::Real=-1.0, α2::Real=-9.0, β2::Real=0.98,
                      ω2::Real=0.1, σ2::Real=0.0,
                      γ::Real=0.10, σz1::Real=0.015, σz2::Real=0.015,
                      lλ☆bds=4.6, nλ☆::Int=451, nlzh::Int=15,
                      nϵ1::Int=5, nϵ2::Int=5)

        # pack up agents and exog
        agent1 = Agent(ρ1, α1, β1, ω1, σ1)
        agent2 = Agent(ρ2, α2, β2, ω2, σ2)
        exog = ExogCV(γ, σz1, σz2, nϵ1, nϵ2)

        # Get linspaces that define grids  for lλ☆, zh
        lλ☆_grid = linspace(-lλ☆bds, lλ☆bds, nλ☆)

        # Get exogenous grid bounds from simulation of 5,000,000 periods
        lzh_bds, _ = get_exog_bounds(exog)
        lzh_grid = linspace(lzh_bds[1], lzh_bds[2], nlzh)
        m_ls = (lλ☆_grid, lzh_grid)

        # Create a grid
        grid = gridmake(collect(lλ☆_grid), collect(lzh_grid))
        grid_transpose = grid'
        lλ☆_pol = repeat(grid_transpose[1, :], inner=[nϵ1*nϵ2, 1])

        # build exog grid for next period.  Both lgp, ξp are Neps × Ngrid
        ϵ, Π = qnwnorm([nϵ1, nϵ2], [0.0, 0.0], eye(2))
        lzhp, lgp = exog_step(exog, grid[:, 2]', ϵ[:, 1], ϵ[:, 2])

        # incompletely initialize the ModelCV so we can pass it to routines that
        # construct the `OtherItems` for us
        m = new(agent1, agent2, exog, m_ls, grid, grid_transpose, lλ☆_pol,
                ϵ, Π, lzhp, lgp)

        # Solve for b2 on the grid and get the coefficients for its approximant
        a1_itp, b2_itp = solve_for_a1b2(m, (-lλ☆bds, lλ☆bds), lzh_bds)
        m.a1_itp = a1_itp
        m.b2_itp = b2_itp

        # # solve for rest of allocation and fill in other items field of ModelCV
        a1 = Float64[a1_itp[grid_transpose[:, i]...] for i=1:nλ☆*nlzh]
        b2 = Float64[b2_itp[grid_transpose[:, i]...] for i=1:nλ☆*nlzh]

        a2 = exp(grid[:, 2]) - a1
        b1 = 1.0 ./ exp(grid[:, 2]) - b2

        c1 = _agg_c(m.agent1, a1, b1)
        c2 = _agg_c(m.agent2, b2, a2)

        m.other = OtherItems(a1, b2, c1, c2)

        m
    end

end

"""
Extracts time t+1 levels of lzh and lgp on the grid in state i at time t for all
time t+1 states

$_cv_exog_ordering_docstr_extra
"""
function get_exog_tomorrow(m::ModelCV, i::Int)
    lzhp_i = m.lzhp[:, i]
    lgp_i = m.lgp[:, i]
    lzhp_i, lgp_i
end

"""
Extracts time t+1 levels of lzh and lgp on the grid in state j at time t+1 and i
at time t

$_cv_exog_ordering_docstr_extra
"""
function get_exog_tomorrow(m::ModelCV, i::Int, j::Int)
    lzhp_ij = m.lzhp[j, i]
    lgp_ij = m.lgp[j, i]
    lzhp_ij, lgp_ij
end

jld_prefix(::ModelCV) = "curr_soln_cv"
jld_prefix(::Type{ModelCV}) = "curr_soln_cv"


text:
using LinearAlgebra: tril!, triu!

function batched_tril!(x::A, d) where {T, N, A <: AbstractArray{T, N}}
  if N < 2
    error("MethodError: no method matching tril!(::Array{Float64,1})")
  elseif N == 2
    return tril!(x, d)
  else
    s = size(x)
    m = (s[1], s[2])
    ms = s[1] * s[2]
    len = Int(length(x) // ms)
    Wt = Core.apply_type(A.name.wrapper, T, 2)
    Threads.@threads for i = 1:len
      tril!(Base.unsafe_wrap(Wt, Base.pointer(x, (ms * (i - 1) + 1)), m), d)
    end
    return x
  end
end

function batched_triu!(x::A, d) where {T, N, A <: AbstractArray{T, N}}
  if N < 2
    error("MethodError: no method matching triu!(::Array{Float64,1})")
  elseif N == 2
    return triu!(x, d)
  else
    s = size(x)
    m = (s[1], s[2])
    ms = s[1] * s[2]
    len = Int(length(x) // ms)
    Wt = Core.apply_type(A.name.wrapper, T, 2)
    Threads.@threads for i = 1:len
      triu!(Base.unsafe_wrap(Wt, Base.pointer(x, (ms * (i - 1) + 1)), m), d)
    end
    return x
  end
end


text:
using CSV, DataFrames
@everywhere  begin
    using Jedi
    using LinearAlgebra
    using Distributions
    using SharedArrays
    emat = Matrix{Float64}(I, 4, 4)
    f = fermi_fitness(f0=100/20000)
end

results = SharedArray{Float64, 2}(4, 26)
@everywhere function run(rho)
    Gamma_arr = zeros(Float64, 26)
    for (i,l) in enumerate(5:30)
        for k in 1:100
            pop = driver_trailer(N=10000, L=30, l=l)
            initiate_rand!(pop, 20)
            mut_arr = rand(Poisson(1), 20000000)
            rho_arr = rand(Poisson(rho), 20000000)
            for j in 1:20000000
                for m in 1:mut_arr[j]
                    mutation!(pop)
                end
                for m in 1:rho_arr[j]
                    driver_mutation!(pop)
                end
                sample_gen!(pop, f, emat; remove=true)
            end
            Gamma_arr[i] += sum(get_energy(pop, emat) .* pop.freqs) / pop.N / l
        end
        println("Length $l done. Gamma= $(Gamma_arr[i]/100)")
    end
    return Gamma_arr
end

rho = [0, 0.00001, 0.0001, 0.0005]
@sync @distributed for r in 1:4
    results[r, :] = run(rho[r]) ./ 100
end

df_results = vcat(results'...)
df_l = vcat((collect(5:30) * ones(4)')...)
df_rho = vcat((rho * ones(26)')'...)

df = DataFrame(rho=df_rho, l=df_l, gamma=df_results)
CSV.write("script1_results.csv", df)