text:
using PSCOPF

using .PSCOPFFixtures

using Test
using JuMP
using Dates
using DataStructures

@testset verbose=true "test_balance_market" begin

    #=
    ECH = 10h
    TS: [11h]
    S: [S1,S2]
                        bus 1                   bus 2
                        |                      |
    (limitable) wind_1_1|       "1_2"          |
    Pmin=0, Pmax=100    |                      |
    Csta=0, Cprop=1     |                      |
      S1: 50            |----------------------|
      S2: 40            |         35           |
                        |                      |
                        |                      |
    (pilotable) prod_1_1|                      |(pilotable) prod_2_1
    Pmin=10, Pmax=100   |                      | Pmin=10, Pmax=100
    Csta=45k, Cprop=10  |                      | Csta=80k, Cprop=15
    ON                  |                      | ON
                        |                      |
           load(bus_1)  |                      |load(bus_2)
      S1: 100           |                      | S1: 40
      S2: 100           |                      | S2: 40
    =#
    function create_instance(ech, next_ech, TS;
                            limit_1::Union{Missing,Float64}=missing,
                            impos_1_s1::Union{Missing,Tuple{Float64,Float64}}=missing, impos_1_s2::Union{Missing,Tuple{Float64,Float64}}=missing,
                            impos_2_s1::Union{Missing,Tuple{Float64,Float64}}=missing, impos_2_s2::Union{Missing,Tuple{Float64,Float64}}=missing,
                            )
        network = PSCOPFFixtures.network_2buses()
        # Limitables
        PSCOPF.Networks.add_new_generator_to_bus!(network, "bus_1", "wind_1_1", PSCOPF.Networks.LIMITABLE,
                                                0., 100.,
                                                0., 1.,
                                                Dates.Second(0), Dates.Second(0))
        # Pilotables
        PSCOPF.Networks.add_new_generator_to_bus!(network, "bus_1", "prod_1_1", PSCOPF.Networks.PILOTABLE,
                                                10., 100.,
                                                45000., 10.,
                                                Dates.Second(0), Dates.Second(0))
        PSCOPF.Networks.add_new_generator_to_bus!(network, "bus_2", "prod_2_1", PSCOPF.Networks.PILOTABLE,
                                                10., 100.,
                                                80000., 15.,
                                                Dates.Second(0), Dates.Second(0))
        # initial generators state
        generators_init_state = SortedDict(
                        "prod_1_1" => PSCOPF.ON,
                        "prod_2_1" => PSCOPF.ON
                    )
        # Uncertainties
        uncertainties = PSCOPF.Uncertainties()
        PSCOPF.add_uncertainty!(uncertainties, ech, "wind_1_1", TS[1], "S1", 50.)
        PSCOPF.add_uncertainty!(uncertainties, ech, "wind_1_1", TS[1], "S2", 40.)
        PSCOPF.add_uncertainty!(uncertainties, ech, "bus_1", TS[1], "S1", 100.)
        PSCOPF.add_uncertainty!(uncertainties, ech, "bus_1", TS[1], "S2", 100.)
        PSCOPF.add_uncertainty!(uncertainties, ech, "bus_2", TS[1], "S1", 40.)
        PSCOPF.add_uncertainty!(uncertainties, ech, "bus_2", TS[1], "S2", 40.)

        context = PSCOPF.PSCOPFContext(network, TS, PSCOPF.PSCOPF_MODE_1,
                                        generators_init_state,
                                        uncertainties, nothing,
                                        "debug_balance")

        tso_actions_l = PSCOPF.get_tso_actions(context)
        if !ismissing(limit_1)
            PSCOPF.set_limitation_value!(tso_actions_l, "wind_1_1", TS[1], limit_1)
        end
        if !ismissing(impos_1_s1)
            PSCOPF.set_imposition_value!(tso_actions_l, "prod_1_1", TS[1], "S1", impos_1_s1[1], impos_1_s1[2])
        end
        if !ismissing(impos_1_s2)
            PSCOPF.set_imposition_value!(tso_actions_l, "prod_1_1", TS[1], "S2", impos_1_s2[1], impos_1_s2[2])
        end
        if !ismissing(impos_2_s1)
            PSCOPF.set_imposition_value!(tso_actions_l, "prod_2_1", TS[1], "S1", impos_2_s1[1], impos_2_s1[2])
        end
        if !ismissing(impos_2_s2)
            PSCOPF.set_imposition_value!(tso_actions_l, "prod_2_1", TS[1], "S2", impos_2_s2[1], impos_2_s2[2])
        end

        market = PSCOPF.BalanceMarket()
        firmness = PSCOPF.compute_firmness(market, ech, next_ech, TS, context)
        result = PSCOPF.run(market, ech, firmness,
                    PSCOPF.get_target_timepoints(context),
                    context)
        PSCOPF.update_market_schedule!(context, ech, result, firmness, market)

        return context, result
    end

    ech = DateTime("2015-01-01T10:00:00")
    TS = [DateTime("2015-01-01T11:00:00")]

    #=
    No limit, No impositions

    ECH = 10h
    TS: [11h]
    S: [S1,S2]
                        bus 1                   bus 2
                        |                      |
    (limitable) wind_1_1|       "1_2"          |
    Pmin=0, Pmax=100    |                      |
    Csta=0, Cprop=1     |                      |
    S1: 50              |----------------------|
    S2: 40              |         35           |
                        |                      |
                        |                      |
    (pilotable) prod_1_1|                      |(pilotable) prod_2_1
    Pmin=10, Pmax=100   |                      | Pmin=10, Pmax=100
    Csta=45k, Cprop=10  |                      | Csta=80k, Cprop=15
    ON                  |                      | ON
                        |                      |
        load(bus_1)     |                      |load(bus_2)
    S1: 100             |                      | S1: 40
    S2: 100             |                      | S2: 40

    wind_1_1 will produce at maximum allowed ;
    S1:50 , S2:40
    prod_1_1 is cheaper => will be used :
    S1:90 , S2:100
    prod_2_1 : shutdown
    =#
    @testset "no_limitations_no_impositions" begin
        context, result = create_instance(ech, ech+Minute(1), TS)
        m_schedule_l = PSCOPF.get_market_schedule(context)

        # Solution is optimal
        @test PSCOPF.get_status(result) == PSCOPF.pscopf_OPTIMAL

        # Limitable respects limitations
        @test 50. ≈ PSCOPF.get_prod_value(m_schedule_l, "wind_1_1", TS[1], "S1")
        @test 40. ≈ PSCOPF.get_prod_value(m_schedule_l, "wind_1_1", TS[1], "S2")
        @test value(result.limitable_model.p_global_capping[TS[1], "S1"]) < 1e-09
        @test value(result.limitable_model.p_global_capping[TS[1], "S2"]) < 1e-09

        @test 90. ≈ PSCOPF.get_prod_value(m_schedule_l, "prod_1_1", TS[1], "S1")
        @test 100. ≈ PSCOPF.get_prod_value(m_schedule_l, "prod_1_1", TS[1], "S2")
        @test PSCOPF.get_prod_value(m_schedule_l, "prod_2_1", TS[1], "S1") < 1e-09
        @test PSCOPF.get_prod_value(m_schedule_l, "prod_2_1", TS[1], "S2") < 1e-09
    end

    @testset "balance_market_respects_limitations" begin

        #=
        Limit wind_1_1 to 45MW => only limits S1

        ECH = 10h
        TS: [11h]
        S: [S1,S2]
                            bus 1                   bus 2
                            |                      |
        (limitable) wind_1_1|       "1_2"          |
        Pmin=0, Pmax=100    |                      |
        Csta=0, Cprop=1     |                      |
        S1: 50              |----------------------|
        S2: 40              |         35           |
                            |                      |
                            |                      |
        (pilotable) prod_1_1|                      |(pilotable) prod_2_1
        Pmin=10, Pmax=100   |                      | Pmin=10, Pmax=100
        Csta=45k, Cprop=10  |                      | Csta=80k, Cprop=15
        ON                  |                      | ON
                            |                      |
            load(bus_1)     |                      |load(bus_2)
        S1: 100             |                      | S1: 40
        S2: 100             |                      | S2: 40

        wind_1_1 will produce at maximum allowed ;
        S1:45 , S2:40
        prod_1_1 is cheaper => will be used :
        S1:95 , S2:100
        prod_2_1 : shutdown
        =#
        @testset "limit_affects_one_scenario" begin
            context, result = create_instance(ech, ech+Minute(1), TS, limit_1=45.)
            m_schedule_l = PSCOPF.get_market_schedule(context)

            # Solution is optimal
            @test PSCOPF.get_status(result) == PSCOPF.pscopf_OPTIMAL

            # Limitable respects limitations
            @test 45. ≈ PSCOPF.get_prod_value(m_schedule_l, "wind_1_1", TS[1], "S1")
            @test 40. ≈ PSCOPF.get_prod_value(m_schedule_l, "wind_1_1", TS[1], "S2")
            @test 5. ≈ value(result.limitable_model.p_global_capping[TS[1], "S1"])
            @test value(result.limitable_model.p_global_capping[TS[1], "S2"]) < 1e-09

            @test 95. ≈ PSCOPF.get_prod_value(m_schedule_l, "prod_1_1", TS[1], "S1")
            @test 100. ≈ PSCOPF.get_prod_value(m_schedule_l, "prod_1_1", TS[1], "S2")
            @test PSCOPF.get_prod_value(m_schedule_l, "prod_2_1", TS[1], "S1") < 1e-09
            @test PSCOPF.get_prod_value(m_schedule_l, "prod_2_1", TS[1], "S2") < 1e-09
        end

        #=
        Limit wind_1_1 to 35MW => limits S1 and S2

        ECH = 10h
        TS: [11h]
        S: [S1,S2]
                            bus 1                   bus 2
                            |                      |
        (limitable) wind_1_1|       "1_2"          |
        Pmin=0, Pmax=100    |                      |
        Csta=0, Cprop=1     |                      |
        S1: 50              |----------------------|
        S2: 40              |         35           |
                            |                      |
                            |                      |
        (pilotable) prod_1_1|                      |(pilotable) prod_2_1
        Pmin=10, Pmax=100   |                      | Pmin=10, Pmax=100
        Csta=45k, Cprop=10  |                      | Csta=80k, Cprop=15
        ON                  |                      | ON
                            |                      |
            load(bus_1)     |                      |load(bus_2)
        S1: 100             |                      | S1: 40
        S2: 100             |                      | S2: 40

        wind_1_1 will produce at maximum allowed ;
        S1:35 , S2:35
        prod_1_1 is cheaper => want to use it at 100MW but prod_2_1 has a minimum capacity of 10MW :
        S1:95 , S2:95
        prod_2_1 :
        S1:10 , S2:10
        =#
        @testset "limit_affects_both_scenarios" begin
            context, result = create_instance(ech, ech+Minute(1), TS, limit_1=35.)
            m_schedule_l = PSCOPF.get_market_schedule(context)

            # Solution is optimal
            @test PSCOPF.get_status(result) == PSCOPF.pscopf_OPTIMAL

            # Limitable respects limitations
            @test 35. ≈ PSCOPF.get_prod_value(m_schedule_l, "wind_1_1", TS[1], "S1")
            @test 35. ≈ PSCOPF.get_prod_value(m_schedule_l, "wind_1_1", TS[1], "S2")
            @test 15. ≈ value(result.limitable_model.p_global_capping[TS[1], "S1"]) #35MW /50
            @test 5. ≈ value(result.limitable_model.p_global_capping[TS[1], "S2"]) #35MW /40

            @test 95. ≈ PSCOPF.get_prod_value(m_schedule_l, "prod_1_1", TS[1], "S1")
            @test 95. ≈ PSCOPF.get_prod_value(m_schedule_l, "prod_1_1", TS[1], "S2")
            @test 10. ≈ PSCOPF.get_prod_value(m_schedule_l, "prod_2_1", TS[1], "S1")
            @test 10. ≈ PSCOPF.get_prod_value(m_schedule_l, "prod_2_1", TS[1], "S2")
        end

    end

    @testset "balance_market_respects_impositions" begin

        #=
        impose on prod_1_1 to be in:
        S1 : [10, 85]
        S2 : [10, 95]
        The market can no longer choose to shutdown prod_1_1

        ECH = 10h
        TS: [11h]
        S: [S1,S2]
                            bus 1                   bus 2
                            |                      |
        (limitable) wind_1_1|       "1_2"          |
        Pmin=0, Pmax=100    |                      |
        Csta=0, Cprop=1     |                      |
        S1: 50              |----------------------|
        S2: 40              |         35           |
                            |                      |
                            |                      |
        (pilotable) prod_1_1|                      |(pilotable) prod_2_1
        Pmin=10, Pmax=100   |                      | Pmin=10, Pmax=100
        Csta=45k, Cprop=10  |                      | Csta=80k, Cprop=15
        ON                  |                      | ON
                            |                      |
            load(bus_1)     |                      |load(bus_2)
        S1: 100             |                      | S1: 40
        S2: 100             |                      | S2: 40

        wind_1_1 will produce at maximum allowed ;
            S1:50 , S2:40
        prod_1_1 is cheaper => will be used as allowed by impositions:
            S1:80 (not 85 cause need to use prod_2_1 and prod_2_1 has a min capacity of 10),
            S2:90 (not 95 cause need to use prod_2_1 and prod_2_1 has a min capacity of 10),
        prod_2_1 at p_min
            S1:10 , S2:10
        =#
        @testset "imposition_on_prod_1_1" begin
            impos_1_s1_l = (10., 85.)
            impos_1_s2_l = (10., 95.)

            context, result = create_instance(ech, ech+Minute(1), TS,
                                            impos_1_s1=impos_1_s1_l,
                                            impos_1_s2=impos_1_s2_l,
                                            )
            m_schedule_l = PSCOPF.get_market_schedule(context)

            # Solution is optimal
            @test PSCOPF.get_status(result) == PSCOPF.pscopf_OPTIMAL

            # Pilotable respects impositions
            @test 80. ≈ value(result.pilotable_model.p_injected["prod_1_1",TS[1], "S1"])
            @test 80. ≈ PSCOPF.get_prod_value(m_schedule_l, "prod_1_1", TS[1], "S1")
            @test PSCOPF.in_bounds(PSCOPF.get_prod_value(m_schedule_l, "prod_1_1", TS[1], "S1"), impos_1_s1_l[1], impos_1_s1_l[2])
            @test 90. ≈ value(result.pilotable_model.p_injected["prod_1_1",TS[1], "S2"])
            @test 90. ≈ PSCOPF.get_prod_value(m_schedule_l, "prod_1_1", TS[1], "S2")
            @test PSCOPF.in_bounds(PSCOPF.get_prod_value(m_schedule_l, "prod_1_1", TS[1], "S2"), impos_1_s2_l[1], impos_1_s2_l[2])

            @test 10. ≈ PSCOPF.get_prod_value(m_schedule_l, "prod_2_1", TS[1], "S1")
            @test 10. ≈ PSCOPF.get_prod_value(m_schedule_l, "prod_2_1", TS[1], "S2")

            # Limitable is used at maximum allowed
            @test 50. ≈ PSCOPF.get_prod_value(m_schedule_l, "wind_1_1", TS[1], "S1")
            @test 40. ≈ PSCOPF.get_prod_value(m_schedule_l, "wind_1_1", TS[1], "S2")
            @test value(result.limitable_model.p_global_capping[TS[1], "S1"]) < 1e-09
            @test value(result.limitable_model.p_global_capping[TS[1], "S2"]) < 1e-09
        end

        #=
        impose on prod_1_1 to be in:
        S1 : [0, 75]
        S2 : [0, 80]
        The market can choose to shutdown prod_1_1 or not
        Here it will use it cause it's cheaper than prod_2_1

        ECH = 10h
        TS: [11h]
        S: [S1,S2]
                            bus 1                   bus 2
                            |                      |
        (limitable) wind_1_1|       "1_2"          |
        Pmin=0, Pmax=100    |                      |
        Csta=0, Cprop=1     |                      |
        S1: 50              |----------------------|
        S2: 40              |         35           |
                            |                      |
                            |                      |
        (pilotable) prod_1_1|                      |(pilotable) prod_2_1
        Pmin=10, Pmax=100   |                      | Pmin=10, Pmax=100
        Csta=45k, Cprop=10  |                      | Csta=80k, Cprop=15
        ON                  |                      | ON
                            |                      |
            load(bus_1)     |                      |load(bus_2)
        S1: 100             |                      | S1: 40
        S2: 100             |                      | S2: 40

        wind_1_1 will produce at maximum allowed ;
            S1:50 , S2:40
        prod_1_1 is cheaper => will be used as allowed by impositions:
            S1:75
            S2:80
        prod_2_1 at p_min
            S1:15 , S2:20
        =#
        @testset "different_impositions_by_scenario" begin
            impos_1_s1_l=(0., 75.)
            impos_1_s2_l=(0., 80.)

            context, result = create_instance(ech, ech+Minute(1), TS,
                                            impos_1_s1=impos_1_s1_l,
                                            impos_1_s2=impos_1_s2_l,
                                            )
            m_schedule_l = PSCOPF.get_market_schedule(context)

            # Solution is optimal
            @test PSCOPF.get_status(result) == PSCOPF.pscopf_OPTIMAL

            # Pilotable respects impositions
            @test 75. ≈ value(result.pilotable_model.p_injected["prod_1_1",TS[1], "S1"])
            @test 75. ≈ PSCOPF.get_prod_value(m_schedule_l, "prod_1_1", TS[1], "S1")
            @test PSCOPF.in_bounds(PSCOPF.get_prod_value(m_schedule_l, "prod_1_1", TS[1], "S1"), impos_1_s1_l[1], impos_1_s1_l[2])
            @test 80. ≈ value(result.pilotable_model.p_injected["prod_1_1",TS[1], "S2"])
            @test 80. ≈ PSCOPF.get_prod_value(m_schedule_l, "prod_1_1", TS[1], "S2")
            @test PSCOPF.in_bounds(PSCOPF.get_prod_value(m_schedule_l, "prod_1_1", TS[1], "S2"), impos_1_s2_l[1], impos_1_s2_l[2])

            @test 15. ≈ PSCOPF.get_prod_value(m_schedule_l, "prod_2_1", TS[1], "S1")
            @test 20. ≈ PSCOPF.get_prod_value(m_schedule_l, "prod_2_1", TS[1], "S2")

            # Limitable is used at maximum allowed
            @test 50. ≈ PSCOPF.get_prod_value(m_schedule_l, "wind_1_1", TS[1], "S1")
            @test 40. ≈ PSCOPF.get_prod_value(m_schedule_l, "wind_1_1", TS[1], "S2")
            @test value(result.limitable_model.p_global_capping[TS[1], "S1"]) < 1e-09
            @test value(result.limitable_model.p_global_capping[TS[1], "S2"]) < 1e-09
        end

        #=
        impose on prod_2_1 to be in:
        S1 : [0, 75]
        S2 : [0, 80]
        The market can choose to shutdown prod_2_1 or not
        Here it will shut it down cause it's cheaper to use prod_1_1

        ECH = 10h
        TS: [11h]
        S: [S1,S2]
                            bus 1                   bus 2
                            |                      |
        (limitable) wind_1_1|       "1_2"          |
        Pmin=0, Pmax=100    |                      |
        Csta=0, Cprop=1     |                      |
        S1: 50              |----------------------|
        S2: 40              |         35           |
                            |                      |
                            |                      |
        (pilotable) prod_1_1|                      |(pilotable) prod_2_1
        Pmin=10, Pmax=100   |                      | Pmin=10, Pmax=100
        Csta=45k, Cprop=10  |                      | Csta=80k, Cprop=15
        ON                  |                      | ON
                            |                      |
            load(bus_1)     |                      |load(bus_2)
        S1: 100             |                      | S1: 40
        S2: 100             |                      | S2: 40

        wind_1_1 will produce at maximum allowed ;
            S1:50 , S2:40
        prod_1_1 is cheaper => will be used as allowed by impositions:
            S1:90
            S2:100
        prod_2_1 at p_min
            S1:0 , S2:0
        =#
        @testset "imposition_levels_may_leave_commitment_decision_to_market" begin
            impos_2_s1_l=(0., 75.)
            impos_2_s2_l=(0., 80.)

            context, result = create_instance(ech, ech+Minute(1), TS,
                                            impos_2_s1=impos_2_s1_l,
                                            impos_2_s2=impos_2_s2_l,
                                            )
            m_schedule_l = PSCOPF.get_market_schedule(context)

            # Solution is optimal
            @test PSCOPF.get_status(result) == PSCOPF.pscopf_OPTIMAL

            # Pilotable respects impositions
            @test 90. ≈ PSCOPF.get_prod_value(m_schedule_l, "prod_1_1", TS[1], "S1")
            @test 100. ≈ PSCOPF.get_prod_value(m_schedule_l, "prod_1_1", TS[1], "S2")

            @test value(result.pilotable_model.p_injected["prod_2_1",TS[1], "S1"]) < 1e-09
            @test PSCOPF.get_prod_value(m_schedule_l, "prod_2_1", TS[1], "S1") < 1e-09
            @test PSCOPF.in_bounds(PSCOPF.get_prod_value(m_schedule_l, "prod_2_1", TS[1], "S1"), impos_2_s1_l[1], impos_2_s1_l[2])
            @test value(result.pilotable_model.p_injected["prod_2_1",TS[1], "S2"])  < 1e-09
            @test PSCOPF.get_prod_value(m_schedule_l, "prod_2_1", TS[1], "S2")  < 1e-09
            @test PSCOPF.in_bounds(PSCOPF.get_prod_value(m_schedule_l, "prod_2_1", TS[1], "S2"), impos_2_s2_l[1], impos_2_s2_l[2])

            # Limitable is used at maximum allowed
            @test 50. ≈ PSCOPF.get_prod_value(m_schedule_l, "wind_1_1", TS[1], "S1")
            @test 40. ≈ PSCOPF.get_prod_value(m_schedule_l, "wind_1_1", TS[1], "S2")
            @test value(result.limitable_model.p_global_capping[TS[1], "S1"]) < 1e-09
            @test value(result.limitable_model.p_global_capping[TS[1], "S2"]) < 1e-09
        end

    end

end

max_stars_repo_name: klorel/pscopf

max_issues_repo_name: klorel/pscopf

max_forks_repo_name: klorel/pscopf


text:
module Wordle

export play

import StatsBase: denserank
import StaticArrays: MVector, MMatrix
import Printf: @printf
import LinearAlgebra: dot

include("core_types/utils.jl")
include("core_types/letter_counts.jl")
include("core_types/letter_flags.jl")
include("core_types/candidate.jl")
include("core_types/result_code.jl")
include("core_types/result.jl")
include("core_types/trial.jl")
include("core_types/constraints.jl")

include("pretty_print/trial.jl")
include("pretty_print/constraints.jl")

include("sourcing/candidate_pool.jl")

include("filtering/update_constraints.jl")
include("filtering/satisfies_constraints.jl")
include("filtering/filter_candidates.jl")

include("scoring/scoring.jl")

include("selection/selection.jl")

include("evaluation.jl")

include("play.jl")

end # module Wordle

ext: jl


text:
mutable struct FieldHandler
    fields::Vector{Field}
    cellset::Set{Int}
    bc_values::Vector{BCValues} # for boundary conditions
end
"""
    FieldHandler(fields::Vector{Field}, cellset)

Construct a `FieldHandler` based on an array of `Field`s and assigns it a set of cells.
"""
function FieldHandler(fields::Vector{Field}, cellset)
    # TODO for now, only accept isoparamtric mapping
    bc_values = [BCValues(field.interpolation, field.interpolation) for field in fields]
    FieldHandler(fields, cellset, bc_values)
end

struct CellVector{T}
    values::Vector{T}
    offset::Vector{Int}
    length::Vector{Int}
end

function Base.getindex(elvec::CellVector, el::Int)
    offset = elvec.offset[el]
    return elvec.values[offset:offset + elvec.length[el]-1]
 end

struct MixedDofHandler{dim,C,T} <: JuAFEM.AbstractDofHandler
    fieldhandlers::Vector{FieldHandler}
    cell_dofs::CellVector{Int}
    cell_nodes::CellVector{Int}
    cell_coords::CellVector{Vec{dim,T}}
    closed::ScalarWrapper{Bool}
    grid::Grid{dim,C,T}
    ndofs::ScalarWrapper{Int}
end

function MixedDofHandler(grid::Grid{dim,C,T}) where {dim,C,T}
    MixedDofHandler{dim,C,T}(FieldHandler[], CellVector(Int[],Int[],Int[]), CellVector(Int[],Int[],Int[]), CellVector(Vec{dim,T}[],Int[],Int[]), JuAFEM.ScalarWrapper(false), grid, JuAFEM.ScalarWrapper(-1))
end

getfieldnames(fh::FieldHandler) = [field.name for field in fh.fields]
getfielddims(fh::FieldHandler) = [field.dim for field in fh.fields]
getfieldinterpolations(fh::FieldHandler) = [field.interpolation for field in fh.fields]
ndofs_per_cell(dh::MixedDofHandler, cell::Int=1) = dh.cell_dofs.length[cell]
nnodes_per_cell(dh::MixedDofHandler, cell::Int=1) = dh.cell_nodes.length[cell]

function celldofs!(global_dofs::Vector{Int}, dh::MixedDofHandler, i::Int)
    @assert isclosed(dh)
    @assert length(global_dofs) == ndofs_per_cell(dh, i)
    unsafe_copyto!(global_dofs, 1, dh.cell_dofs.values, dh.cell_dofs.offset[i], length(global_dofs))
    return global_dofs
end

function celldofs(dh::MixedDofHandler, i::Int)
    @assert isclosed(dh)
    return dh.cell_dofs[i]
end

function cellcoords!(global_coords::Vector{Vec{dim,T}}, dh::MixedDofHandler, i::Int) where {dim,T}
    @assert isclosed(dh)
    @assert length(global_coords) == nnodes_per_cell(dh, i)
    unsafe_copyto!(global_coords, 1, dh.cell_coords.values, dh.cell_coords.offset[i], length(global_coords))
    return global_coords
end

function cellnodes!(global_nodes::Vector{Int}, dh::MixedDofHandler, i::Int) where {dim,T}
    @assert isclosed(dh)
    @assert length(global_nodes) == nnodes_per_cell(dh, i)
    unsafe_copyto!(global_nodes, 1, dh.cell_nodes.values, dh.cell_nodes.offset[i], length(global_nodes))
    return global_nodes
end

"""
    getfieldnames(dh::MixedDofHandler)

Returns the union of all the fields. Can be used as an iterable over all the fields in the problem.
"""
function getfieldnames(dh::MixedDofHandler)
    fieldnames = Vector{Symbol}()
    for fh in dh.fieldhandlers
        for name in getfieldnames(fh)
            push!(fieldnames, name)
        end
    end
    return unique!(fieldnames)
end

"""
    getfielddim(dh::MixedDofHandler, name::Symbol)

Returns the dimension of a specific field, given by name. Note that it will return the dimension of the first field found among the `FieldHandler`s.
"""
function getfielddim(dh::MixedDofHandler, name::Symbol)

    for fh in dh.fieldhandlers
        field_pos = findfirst(i->i == name, getfieldnames(fh))
        if field_pos > 0
            return fh.fields[field_pos].dim
        end
    end
    error("did not find field $name")
end

"""
    nfields(dh::MixedDofHandler)

Returns the number of unique fields defined.
"""
nfields(dh::MixedDofHandler) = length(getfieldnames(dh))

function Base.push!(dh::MixedDofHandler, fh::FieldHandler)
    @assert !isclosed(dh)
    push!(dh.fieldhandlers, fh)
    return dh
end

function Base.push!(dh::MixedDofHandler, name::Symbol, dim::Int)
    celltype = getcelltype(dh.grid)
    isconcretetype(celltype) || error("If you have more than one celltype in Grid, you must use push!(dh::MixedDofHandler, fh::FieldHandler)")
    push!(dh, name, dim, default_interpolation(celltype))
end

function Base.push!(dh::MixedDofHandler, name::Symbol, dim::Int, ip::Interpolation)
    @assert !isclosed(dh)

    if length(dh.fieldhandlers) == 0
        cellset = Set(1:getncells(dh.grid))
        push!(dh.fieldhandlers, FieldHandler(Field[],cellset,BCValues[]))
    elseif length(dh.fieldhandlers) > 1
        error("If you have more than one FieldHandler, you must specify field")
    end
    fh = first(dh.fieldhandlers)

    field = Field(name,ip,dim)
    bc_value = BCValues(field.interpolation, field.interpolation)

    push!(fh.fields, field)
    push!(fh.bc_values, bc_value)

    return dh
end

"""
    close!(dh::MixedDofHandler)

Closes the dofhandler and creates degrees of freedom for each cell.
Dofs are created in the following order: Go through each FieldHandler in the order they were added. For each field in the FieldHandler, create dofs for the cell. This means that dofs on a particular cell will be numbered according to the fields; first dofs for field 1, then field 2, etc.
"""
function close!(dh::MixedDofHandler{dim}) where {dim}

    @assert !JuAFEM.isclosed(dh)
    field_names = JuAFEM.getfieldnames(dh)  # all the fields in the problem
    numfields =  length(field_names)

    # Create dicts that stores created dofs
    # Each key should uniquely identify the given type
    vertexdicts = [Dict{Int, Array{Int}}() for _ in 1:numfields]
    edgedicts = [Dict{Tuple{Int,Int}, Array{Int}}() for _ in 1:numfields]
    facedicts = [Dict{Tuple{Int,Int}, Array{Int}}() for _ in 1:numfields]
    celldicts = [Dict{Int, Array{Int}}() for _ in 1:numfields]

    # Set initial values
    nextdof = 1  # next free dof to distribute

    append!(dh.cell_dofs.offset, zeros(getncells(dh.grid)))
    append!(dh.cell_dofs.length, zeros(getncells(dh.grid)))

    @debug "\n\nCreating dofs\n"
    for fh in dh.fieldhandlers
        # sort the cellset since we want to loop through the cells in a fixed order
        cellnumbers = sort(collect(fh.cellset))
        nextdof = _close!(
            dh,
            cellnumbers,
            JuAFEM.getfieldnames(fh),
            JuAFEM.getfielddims(fh),
            JuAFEM.getfieldinterpolations(fh),
            nextdof,
            vertexdicts,
            edgedicts,
            facedicts,
            celldicts)
    end
    dh.ndofs[] = maximum(dh.cell_dofs.values)
    dh.closed[] = true

    #Create cell_nodes and cell_coords (similar to cell_dofs)
    push!(dh.cell_nodes.offset, 1)
    push!(dh.cell_coords.offset, 1)
    for cell in dh.grid.cells
        for nodeid in cell.nodes
            push!(dh.cell_nodes.values, nodeid)
            push!(dh.cell_coords.values, dh.grid.nodes[nodeid].x)
        end
        push!(dh.cell_nodes.offset, length(dh.cell_nodes.values)+1)
        push!(dh.cell_coords.offset, length(dh.cell_coords.values)+1)
        push!(dh.cell_nodes.length, length(cell.nodes))
        push!(dh.cell_coords.length, length(cell.nodes))
    end

    return dh
end

function _close!(dh::MixedDofHandler{dim}, cellnumbers, field_names, field_dims, field_interpolations, nextdof, vertexdicts, edgedicts, facedicts, celldicts) where {dim}

    ip_infos = JuAFEM.InterpolationInfo[]
    for interpolation in field_interpolations
        ip_info = JuAFEM.InterpolationInfo(interpolation)
        # these are not implemented yet (or have not been tested)
        @assert(ip_info.nvertexdofs <= 1)
        @assert(ip_info.nedgedofs <= 1)
        @assert(ip_info.nfacedofs <= 1)
        @assert(ip_info.ncelldofs <= 1)  # not tested but probably works
        push!(ip_infos, ip_info)
    end

    # loop over all the cells, and distribute dofs for all the fields
    for ci in cellnumbers
        dh.cell_dofs.offset[ci] = length(dh.cell_dofs.values)+1

        cell = dh.grid.cells[ci]
        cell_dofs = Int[]  # list of global dofs for each cell
        @debug "Creating dofs for cell #$ci"

        for fi in 1:length(field_names)
            @debug "\tfield: $(field_names[fi])"
            ip_info = ip_infos[fi]

            if ip_info.nvertexdofs > 0
                nextdof = add_vertex_dofs(cell_dofs, cell, vertexdicts[fi], field_dims[fi], ip_info.nvertexdofs, nextdof)
            end

            if ip_info.nfacedofs > 0
                nextdof = add_face_dofs(cell_dofs, cell, facedicts[fi], field_dims[fi], ip_info.nfacedofs, nextdof)
            end

            if ip_info.nedgedofs > 0
                nextdof = add_edge_dofs(cell_dofs, cell, edgedicts[fi], field_dims[fi], ip_info.nedgedofs, nextdof)
            end

            if ip_info.ncelldofs > 0
                nextdof = add_cell_dofs(cell_dofs, ci, celldicts[fi], field_dims[fi], ip_info.ncelldofs, nextdof)
            end

        end
        # after done creating dofs for the cell, push them to the global list
        push!(dh.cell_dofs.values, cell_dofs...)
        dh.cell_dofs.length[ci] = length(cell_dofs)

        @debug "Dofs for cell #$ci:\n\t$cell_dofs"
    end # cell loop
    return nextdof
end

"""
Returns the next global dof number and an array of dofs.
If dofs have already been created for the object (vertex, face) then simply return those, otherwise create new dofs.
"""
function get_or_create_dofs!(nextdof, field_dim; dict, key)

    token = Base.ht_keyindex2!(dict, key)
    if token > 0  # vertex, face etc. visited before
        # reuse stored dofs (TODO unless field is discontinuous)
        @debug "\t\tkey: $key dofs: $(dict[key])  (reused dofs)"
        return nextdof, dict[key]

    else  # create new dofs
        dofs = collect(nextdof : nextdof + field_dim-1)
        @debug "\t\tkey: $key dofs: $dofs"
        Base._setindex!(dict, dofs, key, -token) #
        nextdof += field_dim
        return nextdof, dofs
    end
end

function add_vertex_dofs(cell_dofs, cell, vertexdict, field_dim, nvertexdofs, nextdof)
    for vertex in JuAFEM.vertices(cell)
        @debug "\tvertex #$vertex"
        nextdof, dofs = get_or_create_dofs!(nextdof, field_dim, dict=vertexdict, key=vertex)
        push!(cell_dofs, dofs...)
    end
    return nextdof
end

function add_face_dofs(cell_dofs, cell, facedict, field_dim, nfacedofs, nextdof)
    @assert nfacedofs == 1 "Currently only supports interpolations with nfacedofs = 1"

    for face in JuAFEM.faces(cell)
        sface = JuAFEM.sortface(face)
        @debug "\tface #$sface"
        nextdof, dofs = get_or_create_dofs!(nextdof, field_dim, dict=facedict, key=sface)
        push!(cell_dofs, dofs...)
    end
    return nextdof
end

function add_edge_dofs(cell_dofs, cell, edgedict, field_dim, nedgedofs, nextdof)
    @assert nedgedofs == 1 "Currently only supports interpolations with nedgedofs = 1"
    for edge in JuAFEM.edges(cell)
        sedge, dir = JuAFEM.sortedge(edge)
        @debug "\tedge #$sedge"
        nextdof, dofs = get_or_create_dofs!(nextdof, field_dim, dict=edgedict, key=sedge)
        push!(cell_dofs, dofs...)
    end
    return nextdof
end

function add_cell_dofs(cell_dofs, cell, celldict, field_dim, ncelldofs, nextdof)
    for celldof in 1:ncelldofs
        @debug "\tcell #$cell"
        nextdof, dofs = get_or_create_dofs!(nextdof, field_dim, dict=celldict, key=cell)
        push!(cell_dofs, dofs...)
    end
    return nextdof
end

# TODO if not too slow it can replace the "Grid-version"
function _create_sparsity_pattern(dh::MixedDofHandler, sym::Bool)
    ncells = getncells(dh.grid)
    N::Int = 0
    for element_id = 1:ncells  # TODO check for correctness
        n = ndofs_per_cell(dh, element_id)
        N += sym ? div(n*(n+1), 2) : n^2
    end
    N += ndofs(dh) # always add the diagonal elements
    I = Int[]; resize!(I, N)
    J = Int[]; resize!(J, N)

    cnt = 0
    for element_id in 1:ncells
        n = ndofs_per_cell(dh, element_id)
        global_dofs = zeros(Int, n)
        celldofs!(global_dofs, dh, element_id)
        @inbounds for j in 1:n, i in 1:n
            dofi = global_dofs[i]
            dofj = global_dofs[j]
            sym && (dofi > dofj && continue)
            cnt += 1
            if cnt > length(J)
                resize!(I, trunc(Int, length(I) * 1.5))
                resize!(J, trunc(Int, length(J) * 1.5))
            end
            I[cnt] = dofi
            J[cnt] = dofj
        end
    end
    @inbounds for d in 1:ndofs(dh)
        cnt += 1
        if cnt > length(J)
            resize!(I, trunc(Int, length(I) + ndofs(dh)))
            resize!(J, trunc(Int, length(J) + ndofs(dh)))
        end
        I[cnt] = d
        J[cnt] = d
    end
    resize!(I, cnt)
    resize!(J, cnt)
    V = zeros(length(I))
    K = sparse(I, J, V)
    return K

end

@inline create_sparsity_pattern(dh::MixedDofHandler) = _create_sparsity_pattern(dh, false)

function find_field(fh::FieldHandler, field_name::Symbol)
    j = findfirst(i->i == field_name, getfieldnames(fh))
    j == 0 && error("did not find field $field_name")
    return j
end

# Calculate the offset to the first local dof of a field
function field_offset(fh::FieldHandler, field_name::Symbol)
    offset = 0
    for i in 1:find_field(fh, field_name)-1
        offset += getnbasefunctions(getfieldinterpolations(fh)[i])::Int * getfielddims(fh)[i]
    end
    return offset
end

find_field(dh::MixedDofHandler, field_name::Symbol) = find_field(first(dh.fieldhandlers), field_name)
field_offset(dh::MixedDofHandler, field_name::Symbol) = field_offset(first(dh.fieldhandlers), field_name)
getfieldinterpolation(dh::MixedDofHandler, field_idx::Int) = dh.fieldhandlers[1].fields[field_idx].interpolation
getfielddim(dh::MixedDofHandler, field_idx::Int) = dh.fieldhandlers[1].fields[field_idx].dim
getbcvalue(dh::MixedDofHandler, field_idx::Int) =  dh.fieldhandlers[1].bc_values[field_idx]


text:
x = 0.1(0:66) .- 3.3
f = t -> t^5 - 13*t^3 + 36*t
y = f.(x)
yy = [y y .+ 10x.^2]

plot(x, f)
plot(x, y)
plot(y)
plot(x, yy)
plot(yy)
plot(x .+ [-1 1], yy)

grid(false)

plot(-2π:0.01:2π, sin)
xlim(-2π, 2π)
ylim((-sqrt(2), sqrt(2)), true)
xticks(π/8, 4)
xticklabels(x -> Base.Printf.@sprintf("%0.1f\\pi", float(x/π)))

# Does not pass CI for Windows!!
# panzoom(0.4, 0)
# xlim(); ylim()
# zoom(1.5)

plot(x, y, ratio = 16//9)
oplot(x, x -> x^3 + x^2 + x)

x = LinRange(-2, 2, 40)
y = 2 .* x .+ 4
stair(x, y)
stair(x, x -> x^3 + x^2 + x)
stair(y)
stair(y, where="pre")
stair(y, where="mid")
stair(y, where="post")

x = LinRange(-2, 2, 40)
y = 0.2 .* x .+ 0.4
scatter(x, y)
scatter(x, x -> 0.2 * x + 0.4)
scatter(y)
x = LinRange(0, 1, 11)
y = LinRange(0, 1, 11)
s = LinRange(50, 400, 11)
c = LinRange(0, 255, 11)
scatter(x, y, s, c)

x = LinRange(-2, 2, 40)
y = 0.2 .* x .+ 0.4
stem(x, y)
stem(x, x -> x^3 + x^2 + x + 6)
stem(y)

x = LinRange(0, 30, 1000)
y = cos.(x) .* x
z = sin.(x) .* x
plot3(x, y, z)

angles = LinRange(0, 2pi, 40)
radii = LinRange(0, 2, 40)
polar(angles, radii)
polar(angles, r -> cos(r) ^ 2)

x = 2 .* rand(100) .- 1
y = 2 .* rand(100) .- 1
z = 2 .* rand(100) .- 1
c = 999 .* rand(100) .+ 1
# Plot the points
scatter3(x, y, z)
# Plot the points with colors
scatter3(x, y, z, c)

population = Dict("Africa" => 1216,
                 "America" => 1002,
                 "Asia" => 4436,
                 "Europe" => 739,
                 "Oceania" => 38)
barplot(keys(population), values(population))
barplot(keys(population), values(population), baseline=1000)
barplot(keys(population), values(population), baseline=500, ylog=true)
barplot(keys(population), values(population), horizontal=true)

x = 2 .* randn(100) .- 1
histogram(x)
histogram(x, ylog=true)
histogram(x, nbins=19)
histogram(x, horizontal=true)

x = 360 .* rand(100)
subplot(1,2,1)
histogram(x)
subplot(1,2,2)
polarhistogram(x, alpha=0.5)
subplot(1,2,1)
histogram(x, nbins=19)
subplot(1,2,2)
polarhistogram(x, nbins=19, alpha=0.5)

Figure()
x = LinRange(0, 1, 100)
plot(x, x.^2)
hold(true)
plot(x, x.^4, label="power 4")
plot(x, x.^8)
legend(location=11)
legend("square", location=2)
hold(false)

# Create example point data
x = 8 .* rand(100) .- 4
y = 8 .* rand(100) .- 4
f = (x,y) -> sin(x) + cos(y)
# Other
contour(x, y, z)
contour(x, y, z, levels=10)
contour(x, y, z, majorlevels=3)
contour(x, y, z, colorbar=false)
# Filled
contourf(x, y, z)
contourf(x, y, z, levels=10)
contourf(x, y, z, majorlevels=3)
contourf(x, y, z, colorbar=false)
# Tricontour
tricont(x, y, z)
tricont(x, y, z, levels=10)
tricont(x, y, z, colorbar=false)
# Surface
# surface(x, y, z) # gr3
trisurf(x, y, z)
# Create example grid data
x = LinRange(-2, 2, 40)
y = LinRange(0, pi, 20)
z = sin.(x') .+ cos.(y) # instead of # x, y = meshgrid(x, y)
# Draw the contour plot using a callable
contour(x, y, f)
contourf(x, y, z)
# surface(x, y, z) # gr3
# surface(x, y, z, accelerate=false) # compile gr3
wireframe(x, y, z)
heatmap(z)
polarheatmap(z)
imshow(z/4 .+ 0.5)

x = randn(100000)
y = randn(100000)
hexbin(x, y)

# s = LinRange(-4, 4, 50)
# v = cos.(s) .+ cos.(s)' .+ cos.(reshape(s,1,1,:))
# isosurface(v, 0.5, tilt=120, color=(0.6, 1.0, 0.85),
#     cameradistance=3.0, twist=12)

# Create point data
x = randn(100_000)
y = randn(100_000)
shade(x, y)
# Create line data with NaN as polyline separator
x = [randn(10000); NaN; randn(10000) .+ 5 ]
y = [randn(10000); NaN; randn(10000) .+ 5]
shade(x, y, xform=3)

# # Create example data
# x = LinRange(-1, 1, 40)
# y = LinRange(-1, 1, 40)
# z = LinRange(-1, 1, 40)
# v = 1 .- (x.^2 .+ y'.^2 .+ reshape(z,1,1,:).^2).^0.5 - 0.25 .* rand(40, 40, 40)
# # Draw the 3d volume data
# volume(v)
# # Draw the 3d volume data using an emission model
# volume(v, algorithm=2)


text:
module CUSOLVER

using ..APIUtils

using ..CUDA
using ..CUDA: CUstream, cuComplex, cuDoubleComplex, libraryPropertyType, cudaDataType
using ..CUDA: libcusolver, libcusolvermg, @allowscalar, assertscalar, unsafe_free!, @retry_reclaim, @context!

using ..CUBLAS: cublasFillMode_t, cublasOperation_t, cublasSideMode_t, cublasDiagType_t
using ..CUSPARSE: cusparseMatDescr_t

using CEnum

# core library
include("libcusolver_common.jl")
include("error.jl")
include("libcusolver.jl")

# low-level wrappers
include("base.jl")
include("sparse.jl")
include("dense.jl")
include("multigpu.jl")

# high-level integrations
include("linalg.jl")

# cache for created, but unused handles
const idle_dense_handles = HandleCache{CuContext,cusolverDnHandle_t}()
const idle_sparse_handles = HandleCache{CuContext,cusolverSpHandle_t}()

function dense_handle()
    cuda = CUDA.active_state()

    # every task maintains library state per device
    LibraryState = @NamedTuple{handle::cusolverDnHandle_t, stream::CuStream}
    states = get!(task_local_storage(), :CUSOLVER_dense) do
        Dict{CuContext,LibraryState}()
    end::Dict{CuContext,LibraryState}

    # get library state
    @noinline function new_state(cuda)
        new_handle = pop!(idle_dense_handles, cuda.context) do
            cusolverDnCreate()
        end

        finalizer(current_task()) do task
            push!(idle_dense_handles, cuda.context, new_handle) do
                @context! skip_destroyed=true cuda.context cusolverDnDestroy(new_handle)
            end
        end

        cusolverDnSetStream(new_handle, cuda.stream)
        (; handle=new_handle, cuda.stream)
    end
    state = get!(states, cuda.context) do
        new_state(cuda)
    end

    # update stream
    @noinline function update_stream(cuda, state)
        cusolverDnSetStream(state.handle, cuda.stream)
        (; state.handle, cuda.stream)
    end
    if state.stream != cuda.stream
        states[cuda.context] = state = update_stream(cuda, state)
    end

    return state.handle
end

function sparse_handle()
    cuda = CUDA.active_state()

    # every task maintains library state per device
    LibraryState = @NamedTuple{handle::cusolverSpHandle_t, stream::CuStream}
    states = get!(task_local_storage(), :CUSOLVER_sparse) do
        Dict{CuContext,LibraryState}()
    end::Dict{CuContext,LibraryState}

    # get or create handle
    @noinline function new_state(cuda)
        new_handle = pop!(idle_sparse_handles, cuda.context) do
            cusolverSpCreate()
        end

        finalizer(current_task()) do task
            push!(idle_sparse_handles, cuda.context, new_handle) do
                @context! skip_destroyed=true cuda.context cusolverSpDestroy(new_handle)
            end
        end

        cusolverSpSetStream(new_handle, cuda.stream)
        (; handle=new_handle, cuda.stream)
    end
    state = get!(states, cuda.context) do
        new_state(cuda)
    end

    # update stream
    @noinline function update_stream(cuda, state)
        cusolverSpSetStream(state.handle, cuda.stream)
        (; state.handle, cuda.stream)
    end
    if state.stream != cuda.stream
        states[cuda.context] = state = update_stream(cuda, state)
    end

    return state.handle
end

function mg_handle()
    cuda = CUDA.active_state()

    # every task maintains library state per set of devices
    LibraryState = @NamedTuple{handle::cusolverMgHandle_t}
    states = get!(task_local_storage(), :CUSOLVERmg) do
        Dict{UInt,LibraryState}()
    end::Dict{UInt,LibraryState}

    # derive a key from the active and selected devices
    key = hash(cuda.context)
    for dev in devices()
        # we hash the device context to support device resets
        key = hash(context(dev), key)
    end

    # get library state
    @noinline function new_state(cuda)
        # we can't reuse cusolverMg handles because they can only be assigned devices once
        new_handle = cusolverMgCreate()

        finalizer(current_task()) do task
            @context! skip_destroyed=true cuda.context cusolverMgDestroy(new_handle)
        end

        devs = convert.(Cint, devices())
        cusolverMgDeviceSelect(new_handle, length(devs), devs)

        (; handle=new_handle)
    end
    state = get!(states, key) do
        new_state(cuda)
    end

    return state.handle
end

function devices!(devs::Vector{CuDevice})
    task_local_storage(:CUSOLVERmg_devices, sort(devs; by=deviceid))
    return
end

devices() = get!(task_local_storage(), :CUSOLVERmg_devices) do
    # by default, select only the first device
    [first(CUDA.devices())]
    # TODO: select all devices
    #sort(collect(CUDA.devices()); by=deviceid)
end::Vector{CuDevice}

ndevices() = length(devices())

end


text:
## SMA performance: buffer vector version with recursive mean instead of direct summation
using TimeSeries, MarketTechnicals
using Reactive

nma = 50
ohlc_BA = TimeSeries.readtimearray("../../data/OHLC_BA.csv")
sma_ta = sma(ohlc_BA["Close"], nma)

s_ohlc = Reactive.Input(ohlc_BA.values[nma,:])

function smarecbuffer!(prev, val_ohlc)
    buffer, prev_sma = prev

    # update SMA value recursively
    new_sma = prev_sma + (val_ohlc[4] - buffer[1])/float(nma)

    # update buffer
    push!(buffer, val_ohlc[4]) # add new close price
    shift!(buffer) # remove the earliest element

    # new_sma becomes prev_sma in the next call
    return (buffer, new_sma)
end
s_sma = lift(s -> s[2], foldl(smarecbuffer!, (ohlc_BA["Close"].values[1:nma], mean(ohlc_BA["Close"].values[1:nma])), s_ohlc))

vsig = [s_sma.value]
function vsigupd!(vsig::Vector{Float64}, nma::Int64, nmax::Int64)
    for i in (nma + 1):nmax
        push!(s_ohlc, ohlc_BA.values[i,:])
        push!(vsig, s_sma.value)
    end
end

# verify sma values
vsigupd!(vsig, nma, length(ohlc_BA))
println("check if SMA values are correct:")
println(mean(abs(vsig - sma_ta.values)) < 1e-7)
println("")

println("SMA foldl-buffer version with recursive calculation:")
println("reactive signal updates timing")
expr_run = parse("@time vsigupd!(vsig, nma, length(ohlc_BA))")
for i = 1:5
    eval(expr_run)
end

### results
# elapsed time: 0.266140198 seconds (23986336 bytes allocated, 16.14% gc time)
### longer and more complex than mean(buffer): no point for this recursive version

Profile.clear()  # in case we have any previous profiling data
@profile vsigupd!(vsig, nma, length(ohlc_BA))
Profile.print()


text:
const NGSIM_TIMESTEP = 0.1 # [sec]
const SMOOTHING_WIDTH_POS = 0.5 # [s]

include(Pkg.dir("NGSIM", "src", "trajectory_smoothing.jl"))

# from "Estimating Acceleration and Lane-Changing
#       Dynamics Based on NGSIM Trajectory Data"
function symmetric_exponential_moving_average(
    arr :: Vector{Float64},
    T   :: Float64; # smoothing width [s]
    dt  :: Float64 = 0.1 # sampling period [s]
    )

    Δ = T / dt

    N = length(arr)
    retval = Array(Float64, N)

    for i = 1 : N

        Z = 0.0
        x = 0.0

        D = min(round(Int, 3Δ), i-1)
        if i+D > N
            D = N-i
        end

        for k in (i-D):(i+D)
            e = exp(-abs(i-k)/Δ)
            Z += e
            x += arr[k] * e
        end

        retval[i] = x / Z
    end

    retval
end

###############

type NGSIMTrajdata
    df         :: DataFrame
    car2start  :: Dict{Int, Int}         # maps carindex to starting index in the df
    frame2cars :: Dict{Int, Vector{Int}} # maps frame to list of carids in the scene
    roadway    :: Roadway

    function NGSIMTrajdata(input_path::AbstractString, roadway::Roadway)

        @assert(isfile(input_path))

        df = readtable(input_path, separator=' ', header = false)
        col_names = [:id, :frame, :n_frames_in_dataset, :epoch, :local_x, :local_y, :global_x, :global_y, :length, :width, :class, :speed, :acc, :lane, :carind_front, :carind_rear, :dist_headway, :time_headway]
        for (i,name) in enumerate(col_names)
            rename!(df, symbol(@sprintf("x%d", i)), name)
        end

        df[:global_heading] = fill(NaN, nrow(df))

        car2start = Dict{Int, Int}()
        frame2cars = Dict{Int, Vector{Int}}()

        for (dfind, carid) in enumerate(df[:id])
            if !haskey(car2start, carid)
                car2start[carid] = dfind
            end

            frame = convert(Int, df[dfind, :frame])
            if !haskey(frame2cars, frame)
                frame2cars[frame] = [carid]
            else
                frame2cars[frame] = push!(frame2cars[frame], carid)
            end
        end

        new(df, car2start, frame2cars, roadway)
    end
end

nframes(trajdata::NGSIMTrajdata) = maximum(keys(trajdata.frame2cars))
carsinframe(trajdata::NGSIMTrajdata, frame::Int) = get(trajdata.frame2cars, frame, Int[]) # NOTE: memory allocation!
carid_set(trajdata::NGSIMTrajdata) = Set(keys(trajdata.car2start)) # NOTE: memory allocation!
nth_carid(trajdata::NGSIMTrajdata, frame::Int, n::Int) = trajdata.frame2cars[frame][n]
first_carid(trajdata::NGSIMTrajdata, frame::Int) = nth_carid(trajdata, frame, 1)
iscarinframe(trajdata::NGSIMTrajdata, carid::Int, frame::Int) = in(carid, carsinframe(trajdata, frame))

function car_df_index(trajdata::NGSIMTrajdata, carid::Int, frame::Int)
    #=
    given frame and carid, find index of car in trajdata
    Returns 0 if it does not exist
    =#

    df = trajdata.df

    lo = trajdata.car2start[carid]
    framestart = df[lo, :frame]

    retval = 0

    if framestart == frame
        retval = lo
    elseif frame ≥ framestart
        retval = frame - framestart + lo
        n_frames = df[lo, :n_frames_in_dataset]
        if retval > lo + n_frames
            retval = 0
        end
    end

    retval
end
function get_frame_range(trajdata::NGSIMTrajdata, carid::Int)
    lo = trajdata.car2start[carid]
    framestart = trajdata.df[lo, :frame]

    n_frames = trajdata.df[lo, :n_frames_in_dataset]
    frameend = framestart + n_frames - 1

    framestart:frameend
end

function pull_vehicle_headings!(trajdata::NGSIMTrajdata;
    v_cutoff::Float64 = 2.5, # speeds below this will use a linearly interpolated heading [fps]
    smoothing_width::Float64 = 0.5, # [s]
    )

    df = trajdata.df

    for carid in carid_set(trajdata)

        frames = collect(get_frame_range(trajdata, carid))
        states = get_state_list_global(trajdata, carid, frames)

        arr_x = map(p->p.x, states)-states[1].x
        arr_y = map(p->p.y, states)-states[1].y

        arr_v = map(i->hypot(states[i]-states[i-1]), 2:length(states)) / 0.1
        slow_indeces = find(v->v < v_cutoff, arr_v)

        slow_segments = Tuple{Int,Int}[]
        i = 0
        while i < length(slow_indeces)
            i += 1
            i_lo = i
            while i < length(slow_indeces) && slow_indeces[i+1] == slow_indeces[i]+1
                i += 1
            end
            i_hi = i
            push!(slow_segments, (slow_indeces[i_lo], slow_indeces[i_hi]))
        end
        slow_segments

        arr_heading = map(i->atan2(convert(VecE2, states[i]-states[i-1])), 2:length(states))
        unshift!(arr_heading, arr_heading[1])

        arr_x_smoothed = symmetric_exponential_moving_average(arr_x, smoothing_width)
        arr_y_smoothed = symmetric_exponential_moving_average(arr_y, smoothing_width)
        arr_dx_smoothed = arr_x_smoothed[2:end] - arr_x_smoothed[1:end-1]
        arr_dy_smoothed = arr_y_smoothed[2:end] - arr_y_smoothed[1:end-1]
        arr_heading2 = map(i->atan2(arr_y_smoothed[i]-arr_y_smoothed[i-1], arr_x_smoothed[i]-arr_x_smoothed[i-1]), 2:length(states))
        unshift!(arr_heading2, arr_heading2[1])

        arr_heading3 = deepcopy(arr_heading2)
        for (lo,hi) in slow_segments
            heading_lo = arr_heading2[max(lo-1, 1)]
            heading_hi = arr_heading2[min(hi+1, length(arr_heading2))]
            for f in lo : hi
                arr_heading3[f] = heading_lo + (heading_hi - heading_lo)*(f-lo)/(hi-lo+1)
            end
        end

        for (i,frame) in enumerate(frames)
            df_ind = car_df_index(trajdata, carid, frame)
            df[df_ind, :global_heading] = arr_heading3[i]
        end
    end

    @assert(findfirst(v->isnan(v), df[:global_heading]) == 0)

    trajdata
end

type FilterTrajectoryResult
    carid::Int
    x_arr::Vector{Float64}
    y_arr::Vector{Float64}
    θ_arr::Vector{Float64}
    v_arr::Vector{Float64}

    function FilterTrajectoryResult(trajdata::NGSIMTrajdata, carid::Int)
        dfstart = trajdata.car2start[carid]
        N = trajdata.df[dfstart, :n_frames_in_dataset]

        # these are our observations
        x_arr = fill(NaN, N)
        y_arr = fill(NaN, N)
        θ_arr = fill(NaN, N)
        v_arr = fill(NaN, N)

        for i in 1 : N
            x_arr[i] = trajdata.df[dfstart + i - 1, :global_x]
            y_arr[i] = trajdata.df[dfstart + i - 1, :global_y]
        end

        # choose an initial belief
        θ_arr[1] = atan2(y_arr[5] - y_arr[1], x_arr[5] - x_arr[1])
        v_arr[1] = trajdata.df[dfstart, :speed] #hypot(ftr.y_arr[lookahead] - y₀, ftr.x_arr[lookahead] - x₀)/ν.Δt
        if v_arr[1] < 1.0 # small speed
            # estimate with greater lookahead
            θ_arr[1] = atan2(y_arr[end] - y_arr[1], x_arr[end] - x_arr[1])
        end

        new(carid, x_arr, y_arr, θ_arr, v_arr)
    end
end
Base.length(ftr::FilterTrajectoryResult) = length(ftr.x_arr)
function filter_trajectory!(ftr::FilterTrajectoryResult, ν::VehicleSystem = VehicleSystem())

    μ = [ftr.x_arr[1], ftr.y_arr[1], ftr.θ_arr[1], ftr.v_arr[1]]

    σ = 1e-1
    Σ = diagm([σ*0.01, σ*0.01, σ*0.1, σ])

    # assume control is centered
    u = [0.0, 0.0]
    z = [NaN, NaN]

    for i in 2 : length(ftr)

        # pull observation
        z[1] = ftr.x_arr[i]
        z[2] = ftr.y_arr[i]

        # apply extended Kalman filter
        μ, Σ = EKF(ν, μ, Σ, u, z)

        # store results
        ftr.x_arr[i] = μ[1]
        ftr.y_arr[i] = μ[2]
        ftr.θ_arr[i] = μ[3]
        ftr.v_arr[i] = μ[4]
    end

    ftr
end
function Base.copy!(trajdata::NGSIMTrajdata, ftr::FilterTrajectoryResult)

    dfstart = trajdata.car2start[ftr.carid]
    N = trajdata.df[dfstart, :n_frames_in_dataset]

    # copy results back to trajdata
    for i in 1 : N
        trajdata.df[dfstart + i - 1, :global_x] = ftr.x_arr[i]
        trajdata.df[dfstart + i - 1, :global_y] = ftr.y_arr[i]
        # trajdata.df[dfstart + i - 1, :speed]   = ftr.v_arr[i]
        if i > 1
            trajdata.df[dfstart + i - 1, :speed]   = hypot(ftr.x_arr[i] - ftr.x_arr[i-1], ftr.y_arr[i] - ftr.y_arr[i-1]) / NGSIM_TIMESTEP
        else
            trajdata.df[dfstart + i - 1, :speed]   = hypot(ftr.x_arr[i+1] - ftr.x_arr[i], ftr.y_arr[i+1] - ftr.y_arr[i]) / NGSIM_TIMESTEP
        end
        trajdata.df[dfstart + i - 1, :global_heading] = ftr.θ_arr[i]
    end

    trajdata
end

function filter_trajectory!(trajdata::NGSIMTrajdata, carid::Int)
    #=
    Filters the given vehicle's trajectory using an Extended Kalman Filter
    =#

    ftr = FilterTrajectoryResult(trajdata, carid)

    # run pre-smoothing
    ftr.x_arr = symmetric_exponential_moving_average(ftr.x_arr, SMOOTHING_WIDTH_POS)
    ftr.y_arr = symmetric_exponential_moving_average(ftr.y_arr, SMOOTHING_WIDTH_POS)

    filter_trajectory!(ftr)

    copy!(trajdata, ftr)
    trajdata
end
function load_ngsim_trajdata(filepath::AbstractString, roadway::Roadway)

    print("loading from file: "); tic()
    tdraw = NGSIMTrajdata(filepath, roadway)
    toc()

    if splitext(filepath)[2] == ".txt" # txt is original
        print("filtering:         "); tic()
        for carid in carid_set(tdraw)
            filter_trajectory!(tdraw, carid)
        end
        toc()
    end

    tdraw
end

function Base.convert(::Type{Trajdata}, tdraw::NGSIMTrajdata)

    df = tdraw.df

    vehdefs = Dict{Int, VehicleDef}()
    states = Array(TrajdataState, nrow(df))
    frames = Array(TrajdataFrame, nframes(tdraw))

    for (id, dfind) in tdraw.car2start
        vehdefs[id] = VehicleDef(id, df[dfind, :class], df[dfind, :length]*METERS_PER_FOOT, df[dfind, :width]*METERS_PER_FOOT)
    end

    state_ind = 0
    for frame in 1 : nframes(tdraw)

        frame_lo = state_ind+1

        for id in carsinframe(tdraw, frame)
            dfind = car_df_index(tdraw, id, frame)

            posG = VecSE2(df[dfind, :global_x]*METERS_PER_FOOT, df[dfind, :global_y]*METERS_PER_FOOT, df[dfind, :global_heading])
            speed = df[dfind, :speed]*METERS_PER_FOOT

            states[state_ind += 1] = TrajdataState(id, VehicleState(posG, tdraw.roadway, speed))
        end

        frame_hi = state_ind
        t = NGSIM_TIMESTEP * (frame-1)
        frames[frame] = TrajdataFrame(frame_lo, frame_hi, t)
    end

    Trajdata(tdraw.roadway, vehdefs, states, frames)
end

function convert_raw_ngsim_to_trajdatas()
    for filename in ("i101_trajectories-0750am-0805am.txt",
                     "i101_trajectories-0805am-0820am.txt",
                     "i101_trajectories-0820am-0835am.txt",
                     "i80_trajectories-0400-0415.txt",
                     "i80_trajectories-0500-0515.txt",
                     "i80_trajectories-0515-0530.txt")

        println("converting ", filename); tic()

        filepath = Pkg.dir("NGSIM", "data", filename)
        roadway = startswith(filename, "i101") ? ROADWAY_101 : ROADWAY_80
        tdraw = NGSIM.load_ngsim_trajdata(filepath, roadway)
        td = convert(Trajdata, tdraw)

        outpath = Pkg.dir("NGSIM", "data", "trajdata_"*filename)
        open(io->write(io, td), outpath, "w")

        toc()
    end
end

const TRAJDATA_PATHS = [
                        Pkg.dir("NGSIM", "data", "trajdata_i101_trajectories-0750am-0805am.txt"),
                        Pkg.dir("NGSIM", "data", "trajdata_i101_trajectories-0805am-0820am.txt"),
                        Pkg.dir("NGSIM", "data", "trajdata_i101_trajectories-0820am-0835am.txt"),
                        Pkg.dir("NGSIM", "data", "trajdata_i80_trajectories-0400-0415.txt"),
                        Pkg.dir("NGSIM", "data", "trajdata_i80_trajectories-0500-0515.txt"),
                        Pkg.dir("NGSIM", "data", "trajdata_i80_trajectories-0515-0530.txt"),
                       ]

function load_trajdata(filepath::AbstractString)
    td = open(io->read(io, Trajdata), filepath, "r")
    td.roadway = startswith(splitdir(filepath)[2], "trajdata_i101") ? ROADWAY_101 : ROADWAY_80
    td
end
load_trajdata(i::Int) = load_trajdata(TRAJDATA_PATHS[i])


text:
module Nonconvex

using Reexport
@reexport using NonconvexCore

include("load.jl")

end


text:
#!/usr/bin/env julia

using PyCall
using SQLite
using MySQL

os = pyimport("os")

Retriever.check_for_updates()

# datasets to test
test_datasets = ["bird-size", "Iris"]

# Service host names
# Use service names on travis as host names else localhost
# Set in Docker compose
pgdb = mysqldb = "localhost"
os_password = ""
if haskey(ENV, "TRAVIS_OR_DOCKER") == true
    pgdb = "pgdb"
    mysqldb = "mysqldb"
    os_password = "Password12!"
end

sqlite_opts = Dict("engine" => "sqlite",
                    "file" => "dbfile",
                    "table_name" => "{db}.{table}")
postgres_opts =  Dict("engine" =>  "postgres",
                        "user" =>  "postgres",
                        "host" =>  pgdb,
                        "password" => os_password,
                        "port" =>  5432,
                        "database" =>  "testdb",
                        "database_name" =>  "testschema2",
                        "table_name" => "{db}.{table}")
csv_opts = Dict("engine" =>  "csv",
                "table_name" => "{db}_{table}.csv")
mysql_opt = Dict("engine" =>  "mysql",
                "user" => "travis",
                "password"=> os_password,
                "host"=> mysqldb,
                "port"=>3306,
                "database_name"=>"testdb",
                "table_name"=>"{db}.{table}")
json_opt = Dict("engine" =>  "json",
                "table_name" => "{db}_{table}.json")
xml_opt = Dict("engine" =>  "xml",
                "table_name" => "{db}_{table}.xml")

function setup()
    # result
end

function teardown()
end

function reset_reload_scripts()
  # Test reset and reload_scripts
  dataset = test_datasets[0]
  Retriever.reset_retriever(dataset)
  Retriever.reload_scripts()
  @test dataset in rdataretriever::datasets()["offline"] == FALSE
  Retriever.get_updates()
  Retriever.reload_scripts()
  @test dataset in rdataretriever::datasets()["offline"] == TRUE
end

function empty_files(path, ext)
    for (root, dirs, files) in walkdir(path)
        for file in files
            relative_path = joinpath(root, file)
            if endswith(relative_path, ext)
                println(relative_path)
                @test filesize(relative_path) > 10
            end
        end
    end
end

function install_csv_engine(data_arg)
    try
        mktempdir() do dir_tmp
            cd(dir_tmp) do
                    # Install dataset into Json database
                    Retriever.install_csv(data_arg,
                        table_name = csv_opts["table_name"])
                    empty_files(dir_tmp, ".csv")
            end
        end
      return true
    catch
        return false
    end
end

my_tempdir = tempdir()
@test isdir(my_tempdir) == true

function install_json_engine(data_arg)
    try
        mktempdir() do dir_tmp
            cd(dir_tmp) do
                    # Install dataset into Json database
                    Retriever.install_json(data_arg,
                        table_name = json_opt["table_name"])
                    empty_files(dir_tmp, ".json")
            end
        end
      return true
    catch
        return false
    end
end

function install_xml_engine(data_arg)
    try
        mktempdir() do dir_tmp
            cd(dir_tmp) do
                    # Install dataset into Json database
                    Retriever.install_xml(data_arg,
                        table_name = xml_opt["table_name"])
                    empty_files(dir_tmp, ".xml")
            end
        end
      return true
    catch
        return false
    end
end

function install_mysql_engine(data_arg)
    try
      # Drop database
      conn = MySQL.connect(mysql_opt["host"], mysql_opt["user"],
          mysql_opt["password"], db = mysql_opt["database_name"])
      db = mysql_opt["database_name"]
      command = "DROP TABLE IF EXISTS $db"
      MySQL.Stmt(conn, command)
      # dframe = mysql_execute(con, command)
      MySQL.disconnect(conn)
      # Install dataset into mysql database
      Retriever.install_mysql(data_arg,
          user = mysql_opt["user"],
          password = mysql_opt["password"],
          host = mysql_opt["host"],
          port = mysql_opt["port"],
          database_name = mysql_opt["database_name"],
          table_name = mysql_opt["table_name"])
      return true
    catch
        return false
    end
end

function install_postgres_engine(data_arg::String)
    try
      # Use python to drop table.
      usr = postgres_opts["user"]
      prt = postgres_opts["port"]
      cmd = "psql -U $usr -d testdb -h $pgdb -p $prt -w -c"
      drop_sql = "\"DROP SCHEMA IF EXISTS testschema CASCADE\""
      query_stm = "$cmd $drop_sql"
      os.system(query_stm)

      # Install dataset into mysql database
      Retriever.install_postgres(data_arg,
          user = postgres_opts["user"],
          password = postgres_opts["password"],
          host = postgres_opts["host"],
          port = postgres_opts["port"],
          database_name = postgres_opts["database_name"],
          table_name = postgres_opts["table_name"])
      return true
    catch
        return false
    end
end

function install_sqlite_engine(data_arg)
    try
        mktempdir() do dir_tmp
            cd(dir_tmp) do
                # Install dataset into SQLite database
                Retriever.install_sqlite(data_arg,
                    file = sqlite_opts["file"],
                    table_name = sqlite_opts["table_name"])
                return true
            end
        end
    catch
        return false
    end
end

@testset "Regression" begin
    @test true
    work_dir = pwd()
    for datset_n in test_datasets

        # Data DB test
        @test true == install_mysql_engine(datset_n)

        # Postgres is currently unstable, December 2018
        @test true == install_postgres_engine(datset_n)
        @test true == install_sqlite_engine(datset_n)

        # File engines use a temporary directory for tests
        @test true == install_csv_engine(datset_n)
        @test true == install_json_engine(datset_n)
        @test true == install_xml_engine(datset_n)
    end

end # @testset Regression


text:
arg_dict = Dict{String, Any}(
"var4" => 1,
"smb_name" => "a string"
)


text:
################################################################################
#
# Generate synthetic data for the BG Compass model
#
# Author: Mathias Loubout and Gabbrio Rizzuti
# Date: September 2020
################################################################################
### Module loading

using JUDI, JLD2

### Load true model
base_path = dirname(pathof(JUDI))*"/../examples/twri/"
@load string(base_path*"data//BGCompass_tti.jld") n d o m

# Model

model_true = Model(n, d, o, m; nb=80)

### Acquisition geometrye
dt = 4f0
T = 2500f0 # total recording time [ms]

# Source wavelet
freq_peak = 0.01f0
wavelet = ricker_wavelet(T, dt, freq_peak)

# Sources
nsrc = 51
x_src = convertToCell(range(0f0, stop = (n[1]-1)*d[1], length = nsrc))
y_src = convertToCell(range(0f0, stop = 0f0, length = nsrc))
z_src = convertToCell(range(0f0, stop = 0f0, length = nsrc))

# Receivers
nrcv = n[1]
x_rcv = range(0f0, stop = (n[1]-1)*d[1], length = nrcv)
y_rcv = 0f0
z_rcv = range(12.5f0, stop = 12.5f0, length = nrcv)

# Geometry structures
src_geom = Geometry(x_src, y_src, z_src; dt = dt, t = T)
rcv_geom = Geometry(x_rcv, y_rcv, z_rcv; dt = dt, t = T, nsrc = nsrc)

# Source function
fsrc = judiVector(src_geom, wavelet)

# Set up info structure for linear operators
ntComp = get_computational_nt(src_geom, rcv_geom, model_true)
info = Info(prod(n), nsrc, ntComp)

# Setup operators
F = judiModeling(info,model_true, src_geom,rcv_geom)
dat = F*fsrc

### Saving data
base_path = dirname(pathof(JUDI))*"/../examples/twri/"
@save string(base_path*"data/BGCompass_data_acou.jld") model_true fsrc dat

max_stars_repo_name: nogueirapeterson/JUDI

max_issues_repo_name: nogueirapeterson/JUDI

max_forks_repo_name: nogueirapeterson/JUDI


text:
include("records.jl")

function nonmem(model, population, fixeffs, obstimes; records = CURRENT_RECORDS, name = missing, rep = 1)
    result = []
    for _ in 1:rep
        simulated = simobs(model, population, fixeffs, obstimes = obstimes)
        push!(result, simulated)
    end
    run = Run(result; records = records, name = name)
    record!(run; records = records)
    return run
end


text:
using TensorNetworkCodes
using TensorNetworkCodes.TNDistance
using Test

@testset "tn_distance" begin
    code = TensorNetworkCode(steane_code())
    code = contract(code, code, [[1, 1], [3, 3]])
    @test tn_distance(code) == 2
end

@testset "tn_distance-contract-same-code" begin
    code = TensorNetworkCode(steane_code())
    code = contract(code, code, [[1, 1], [3, 3]])
    code = contract(code, code, [[1, 1], [7, 7]])
    brute_distance = find_distance_logicals(code)[1]
    tensor_distance = tn_distance(code)
    @test tensor_distance == brute_distance == 2
end

@testset "tn_operator_weights" begin
    code = TensorNetworkCode(five_qubit_code())
    code = contract(code, code, [[1, 1], [3, 3]])
    expected = OperatorWeights([1, 0, 0, 0, 9, 0, 6], [1, 0, 9, 24, 99, 72, 51])
    actual = tn_operator_weights(code)
    @test actual.stabilizer_weights == expected.stabilizer_weights
    @test actual.all_operator_weights == expected.all_operator_weights
    @test actual.distance == expected.distance
    plot_operator_weights(actual)  # no error thrown
end