text:
using CSV
using Statistics
using Random

import Gurobi
const GRBENV = Gurobi.Env()

import MathOptInterface
const MOI = MathOptInterface

const SAF = MOI.ScalarAffineFunction
const SAT = MOI.ScalarAffineTerm

import Linda

const DATA_DIR = joinpath(@__DIR__, "../dat")
const RMPDIR = joinpath(@__DIR__, "rmp")

include("../src/DemandResponse.jl")
const DR = DemandResponse

"""
    generate_resources()

# Arguments
- `num_der`: The number of households
- `T`: The number of time-steps
- `devices_rates`: Ownership rates of each device
- `load_norm`: Normalized reference load profile
- `PV_norm`: Normalized PV output profile
- `price`: Electricity price, in dollar per kW.h.
"""
function generate_resources(
    num_der, T,
    devices_rates::Dict{Symbol, Float64},
    load_norm::Vector,
    PV_norm::Vector,
    temp_ext::Vector,
    price::Vector;
    seed=0
)
    # Sanity checks
    T == length(load_norm) || throw(DimensionMismatch(
        "T=$T but load_norm has length $(length(load_norm))"
    ))
    T == length(PV_norm) || throw(DimensionMismatch(
        "T=$T but load_norm has length $(length(PV_norm))"
    ))
    T == length(price) || throw(DimensionMismatch(
        "T=$T but load_norm has length $(length(price))"
    ))

    Random.seed!(seed)
    resources = Vector{DR.DER.Resource}(undef, num_der)

    # Create a pool of `num_der` houses
    # Each household owns appliances according to `devices_rates`
    for i in 1:num_der

        appliances = DR.DER.Resource[]
        γ = 0.5 + rand()  # household's scaling factor

        # ==================
        #    Fixed load
        # ==================
        # All houses have a fixed load
        if rand() < devices_rates[:fixed]
            l = DR.DER.FixedLoad(
                index=i,
                T=T,
                load= γ .* max.(0.0, (load_norm .+ 0.05 .* randn(T)))
            )
            push!(appliances, l)
        end

        # ==================
        #    Fixed load
        # ==================
        # Dishwasher
        if rand() < devices_rates[:dishwasher]
            dishwasher = DR.DER.ShiftableLoad(
                index = 10000000 + i,  # TODO: handle index number properly
                T=T,
                cycle_begin_min=1,
                cycle_begin_max=T-1,
                cycle_length=2,
                cycle_load=[1.0, 1.0]
            )
            push!(appliances, dishwasher)
        end

        # ======================
        #    Shiftable load
        # ======================
        # Clothes washer
        if rand() < devices_rates[:clothes_washer]
            clothes_washer = DR.DER.ShiftableLoad(
                index = 20000000 + i,  # TODO: handle index number properly
                T=T,
                cycle_begin_min=1,
                cycle_begin_max=T-1,
                cycle_length=2,
                cycle_load=[1.0, 1.0]
            )
            push!(appliances, clothes_washer)

            # Clothes dryer
            # Only household with clothes washer may have a clothes dryer
            if rand() < (
                devices_rates[:clothes_dryer] / devices_rates[:clothes_washer]
            )

                clothes_dryer = DR.DER.ShiftableLoad(
                    index = 30000000 + i,  # TODO: handle index number properly
                    T=T,
                    cycle_begin_min=1,
                    cycle_begin_max=T-2,
                    cycle_length=3,
                    cycle_load=[1.0, 1.0, 1.0]
                )
                push!(appliances, clothes_dryer)
            end
        end
        # ===================
        #    Thermal load
        # ===================
        # Electric heating
        if rand() < devices_rates[:thermal]
            th = DR.DER.ThermalLoad(
                index=i, T=T,
                temp_min=fill(18.0, T), temp_max=fill(22.0, T),
                temp_init=20.0,
                temp_ext=temp_ext + 0.5 .* randn(T),
                pwr_min=zeros(T), pwr_max=10.0 .* ones(T),
                C=3.0, η=1.0, μ=0.2,
                binflag=false
            )
            push!(appliances, th)
        end

        # =====================
        #    PV, EV & ES
        # =====================
        if rand() < devices_rates[:renew]

            # Electric vehicle
            if (T % 24) == 0
                ndays = div(T, 24)

                ev = DR.DER.DeferrableLoad(
                    index=i, T=T,
                    ncycles=ndays,
                    cycle_begin=[16 + 24*(d-1) for d in 1:ndays],
                    cycle_end=[24 + 24*(d-1) for d in 1:ndays],
                    cycle_energy_min=fill(10.0, ndays),
                    cycle_energy_max=fill(10.0, ndays),
                    pwr_min=repeat([zeros(15); 1.1*ones(9)], ndays),
                    pwr_max=repeat([zeros(15); 7.7*ones(9)], ndays),
                    binflag=true
                )
                push!(appliances, ev)

            else
                # Not an integer number of days
                # Do not include EV
            end

            # PV
            ζ = rand(T)
            pv = DR.DER.CurtailableLoad(
                index=i,
                T=T,
                load= - γ .* PV_norm .* ζ,
                binflag=true
            )
            push!(appliances, pv)

            # Battery
            bat = DR.DER.Battery(
                index=i,
                T=T,
                soc_min=zeros(T),
                soc_max=13.5 .* ones(T),
                soc_init=0.0,
                selfdischargerate=1.0,
                charge_pwr_min=0.0, charge_pwr_max=5.0,
                discharge_pwr_min=0.0, discharge_pwr_max=5.0,
                charge_eff=0.95, discharge_eff=0.95
            )
            push!(appliances, bat)
        end

        # Household
        h = DR.DER.House(
            index=i,
            T=T,
            netload_min = zeros(T),
            netload_max = 10.0 .* ones(T),
            price=price,
            appliances=appliances
        )
        resources[i] = h
    end

    return resources
end

function main(T, R, T0=408+7)

    # Load data
    ont_load = CSV.read(joinpath(DATA_DIR, "load2016.csv"))
    ont_price = CSV.read(joinpath(DATA_DIR, "price2016.csv"))
    ont_prod = CSV.read(joinpath(DATA_DIR, "prod2016.csv"))
    ont_temp = CSV.read(joinpath(DATA_DIR, "temperature2016.csv"))

    # Market and Time-Of-Use prices
    p_tou = Vector{Float64}(ont_price.TOU[T0:(T0+T-1)])
    p_mkt = Vector{Float64}(ont_price.HOEP[T0:(T0+T-1)])

    # Normalized load
    load_norm = (ont_load.OntDemand / mean(ont_load.OntDemand))[T0:(T0+T-1)]

    # Normalized PV output
    pv_norm = (ont_prod.SOLAR / mean(ont_prod.SOLAR))[T0:(T0+T-1)]

    # Outside temperature
    temp_ext = ont_temp.Temperature[T0:(T0+T-1)]

    # Min and max total net load
    total_load_min = 0.0 * R * ones(T)
    total_load_max = 6.0 * R * ones(T)

    # Devices' ownership rates
    own_rates = Dict(
        :fixed => 1.0,
        :dishwasher => 0.65,
        :clothes_washer => 0.9,
        :clothes_dryer => 0.75,
        :thermal => 0.6,
        :renew => 0.33
    )

    # Generate resources
    resources = generate_resources(
        R, T, own_rates,
        load_norm, pv_norm, temp_ext, p_mkt,
        seed=5
    )

    # Instantiate oracles
    oracles = [
    DR.build_oracle(
            MOI.Bridges.full_bridge_optimizer(Gurobi.Optimizer(GRBENV, OutputFlag=0, Threads=1), Float64),
            hh
        )
        for hh in resources
    ]

    # Instantiate the RMP
    rmp = Gurobi.Optimizer(GRBENV);
    MOI.set(rmp, MOI.RawParameter("Threads"), 1)
    MOI.set(rmp, MOI.RawParameter("Method"), 2)
    MOI.set(rmp, MOI.RawParameter("OutputFlag"), 0)
    M = 1e3

    # Linking variables & artificial slacks for linking constraints
    y  = MOI.add_variables(rmp, T)
    zp = MOI.add_variables(rmp, T)
    zm = MOI.add_variables(rmp, T)
    for t in 1:T
        MOI.add_constraint(rmp, MOI.SingleVariable(zp[t]), MOI.GreaterThan(0.0))
        MOI.add_constraint(rmp, MOI.SingleVariable(zm[t]), MOI.GreaterThan(0.0))
        MOI.add_constraint(rmp, MOI.SingleVariable(y[t]), MOI.Interval(total_load_min[t], total_load_max[t]))
    end

    var_link = vcat(y, zp, zm)

    # Linking constraints
    con_link = MOI.add_constraints(rmp,
        [
            SAF([SAT(-1.0, y[t]), SAT(1.0, zp[t]), SAT(-1.0, zm[t])], 0.0)
            for t in 1:T
        ],
        [
            MOI.EqualTo(0.0)
            for t in 1:T
        ]
    );

    # Convexity constraints
    con_cvx = MOI.add_constraints(rmp,
        [SAF(SAT{Float64}[], 0.0) for _ in 1:R],
        [MOI.EqualTo(1.0) for _ in 1:R]
    );

    # Objective
    MOI.set(rmp, MOI.ObjectiveSense(), MOI.MIN_SENSE)
    fobj = SAF{Float64}([
        [
            SAT(M, zp[t]) for t in 1:T
        ]; [
            SAT(M, zm[t]) for t in 1:T
        ]], 0.0
    )
    MOI.set(rmp, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(), fobj)

    rhs = zeros(T)
    mp = Linda.Master(rmp, con_cvx, con_link, rhs, var_link, MOI.VariableIndex[])

    # Add initial columns
    for oracle in oracles
        Linda.update!(oracle, false, zeros(T), 0.0)
        Linda.optimize!(oracle)

        cols = Linda.get_columns(oracle)
        col, _ = cols[1]
        Linda.add_column!(mp, col)
    end

    env = Linda.LindaEnv()
    env.verbose.val = 1
    env.num_columns_max.val = div(R, 8)
    env.num_cgiter_max.val = 200

    MOI.set(rmp, MOI.Silent(), true)

    cglog = Dict()

    function cg_callback()
        niter::Int = cglog[:n_cg_iter]
        if niter > 0 && niter %10 == 0
            @info "Saving RMP at iter $niter"
            Gurobi.write_model(rmp.inner, joinpath(RMPDIR, "DER_$(T)_$(R)_$(niter).mps.bz2"))
        end
        return nothing
    end

    # Print MP solver
    println("MP solver: ", MOI.get(mp.rmp, MOI.SolverName()), "\n")

    Random.seed!(0)
    Linda.solve_colgen!(env, mp, oracles, cg_log=cglog, callback=cg_callback)
    Gurobi.write_model(rmp.inner, joinpath(RMPDIR, "DER_$(T)_$(R)_$(cglog[:n_cg_iter]).mps.bz2"))
    return
end

if abspath(PROGRAM_FILE) == @__FILE__
    T = parse(Int, ARGS[1])
    R = parse(Int, ARGS[2])
    @info "" T R
    main(T, R)
end


text:
fehmdir = "../src/contamination-model/";
fehmdirresult = "../tests/data/";
srcdir = "../src/";

ext: jl

max_stars_repo_name: josephbakarji/socio-hydro-abm

max_issues_repo_name: josephbakarji/socio-hydro-abm

max_forks_repo_name: josephbakarji/socio-hydro-abm


text:
using PyPlot
using QuantEcon

# == Model parameters == #
r = 0.05
bet = 1 / (1 + r)
T = 45
c_bar = 2.0
sigma = 0.25
mu = 1.0
q = 1e6

# == Formulate as an LQ problem == #
Q = 1.0
R = zeros(2, 2)
Rf = zeros(2, 2); Rf[1, 1] = q
A = [1.0+r -c_bar+mu;
     0.0 1.0]
B = [-1.0, 0.0]
C = [sigma, 0.0]

# == Compute solutions and simulate == #
lq = LQ(Q, R, A, B, C; bet=bet, capT=T, rf=Rf)
x0 = [0.0, 1.0]
xp, up, wp = compute_sequence(lq, x0)

# == Convert back to assets, consumption and income == #
assets = squeeze(xp[1, :], 1)            # a_t
c = squeeze(up .+ c_bar, 1)              # c_t
income = squeeze(wp[1, 2:end] .+ mu, 1)  # y_t

# == Plot results == #
n_rows = 2
fig, axes = subplots(n_rows, 1, figsize=(12, 10))

subplots_adjust(hspace=0.5)
for i=1:n_rows
    axes[i][:grid]()
    axes[i][:set_xlabel]("Time")
end
bbox = [0.0, 1.02, 1.0, 0.102]

# Make first plot
axes[1][:plot](2:T+1, income, "g-", label="non-financial income", lw=2,
               alpha=0.7)
axes[1][:plot](1:T, c, "k-", label="consumption", lw=2, alpha=0.7)
axes[1][:legend](ncol=2, bbox_to_anchor=bbox, loc=3, mode="expand")

# Make second plot
axes[2][:plot](2:T+1, cumsum(income .- mu), "r-",
               label="cumulative unanticipated income", lw=2, alpha=0.7)
axes[2][:plot](1:T+1, assets, "b-", label="assets", lw=2, alpha=0.7)
axes[2][:plot](1:T, zeros(T), "k-")
axes[2][:legend](ncol=2, bbox_to_anchor=bbox, loc=3, mode="expand")


text:
using RSUtils, Test
using Distributions: Distributions

@testset verbose = true "RSUtils.jl" begin
    include("fixedpoint.jl")
    include("statistics.jl")
    include("predicates.jl")
end


text:
using Bridge
using Test, LinearAlgebra, Random
using Distributions, StaticArrays
using Bridge: Gaussian, PSD

@testset "Number" begin
    let
        μ = rand()
        x = rand()
        Σ = rand()^2

        p = pdf(Normal(μ, √Σ), x)
        @test pdf(Gaussian(μ, Σ), x) ≈ p
        @test pdf(Gaussian(μ, Σ*I), x) ≈ p
        @test pdf(Gaussian([μ], [√Σ]*[√Σ]'), [x]) ≈ p

        @test pdf(Gaussian((@SVector [μ]), @SMatrix [Σ]), @SVector [x]) ≈ p
    end
end

@testset "Array" begin
    for d in 1: 3
        μ = rand(d)
        x = rand(d)
        rΣ = tril(rand(d,d))
        Σ = rΣ*rΣ'
        p = pdf(MvNormal(μ, Σ), x)

        @test pdf(Gaussian(μ, Σ), x) ≈ p
        @test pdf(Gaussian(μ, PSD(rΣ)), x) ≈ p
        @test pdf(Gaussian(SVector{d}(μ), SMatrix{d,d}(Σ)), x) ≈ p
        @test pdf(Gaussian(SVector{d}(μ), PSD(SMatrix{d,d}(rΣ))), x) ≈ p
    end
end
@testset "Diagonal" begin
    for d in 1: 3
        μ = rand(d)
        x = rand(d)
        rΣ = rand()
        Σ = Matrix(I, d, d)*rΣ^2
        p = pdf(MvNormal(μ, Σ), x)

        @test pdf(Gaussian(μ, rΣ^2*I), x) ≈ p
        @test pdf(Gaussian(SVector{d}(μ), SDiagonal(rΣ^2*ones(SVector{d}))), x) ≈ p
        @test pdf(Gaussian(SVector{d}(μ), SMatrix{d,d}(Σ)), x) ≈ p
    end
end


text:
using Test, LinearMaps, LinearAlgebra, BlockArrays

@testset "Non-traditional axes" begin

    A = rand(ComplexF64,2,4)
    B = PseudoBlockMatrix{ComplexF64}(undef, [2,3], [3,4])

    ax1 = axes(B)[1]
    ax2 = axes(B)[2]
    fill!(B,0)
    B[Block(1),Block(2)] .= A

    N = @inferred LinearMap(B)
    @test axes(N) == (ax1,ax2)

    @test eltype(N) == eltype(B)

    u = similar(Array{ComplexF64}, ax2)
    v = PseudoBlockVector{ComplexF64}(undef, [3,5])
    w = PseudoBlockVector{ComplexF64}(undef, [4,3])
    # v = similar(Array{ComplexF64}, blockedrange([3,5]))
    # w = similar(Array{ComplexF64}, blockedrange([4,3]))

    for i in eachindex(u) u[i] = rand(ComplexF64) end
    for i in eachindex(v) v[i] = rand(ComplexF64) end
    for i in eachindex(w) w[i] = rand(ComplexF64) end

    @test B*u == N*u
    @test_throws DimensionMismatch N*v

    # Lu = L*u
    Nu = N*u

    @test axes(Nu)[1] == axes(N)[1] == axes(B)[1]
    @test blocklengths(axes(Nu)[1]) == blocklengths(axes(N)[1]) == blocklengths(axes(B)[1]) == [2,3]

    C = B + 2N
    @test axes(C) === axes(B) === axes(N)
    @test C*u ≈ 3*Nu

    Cu = C*u
    @test axes(C)[1] == ax1
    @test blocklengths(axes(C)[1]) == blocklengths(ax1)
end


text:
const blacklist_file = begin
    config_dir = abspath(homedir(), ".julia", "config")
    !isdir(config_dir) && mkpath(config_dir)
    res = abspath(config_dir,"Fezzik.blacklist")
    res
end

function read_blacklist()
    list = Set{String}()
    !isfile(blacklist_file) && return list
    for mod in eachline(blacklist_file)
        push!(list,mod)
    end
    return list
end

function write_blacklist(list::Set{String})
    open(blacklist_file,"w") do io
        for s in list
            println(io,s)
        end
    end
end

function blacklist(mods::String...)
    list = read_blacklist()
    for mod in mods
        push!(list,mod)
    end
    write_blacklist(list)
    list
end

function whitelist(mods::String...)
    list = blacklist()
    for mod in mods
        (mod in list) && pop!(list,mod)
    end
    write_blacklist(list)
    list
end

whitelist(mods) = whitelist(mods...)

export whitelist,blacklist


text:
using Random

function main()
    counts = Dict{Int64,Int64}()
    open("orig-node-labels-mathoverflow-answers.txt") do f
        for line in eachline(f)
            tags = [parse(Int64, t) for t in split(line, ",")]
            for t in tags
                curr = get(counts, t, 0)
                counts[t] = curr + 1
            end
        end
    end

    open("node-labels-mathoverflow-answers.txt", "w") do g
        open("orig-node-labels-mathoverflow-answers.txt") do f
            for line in eachline(f)
                tags = [parse(Int64, t) for t in split(line, ",")]
                cnts = [counts[t] for t in tags]
                most_popular_tags = tags[findall(cnts .== maximum(cnts))]
                shuffle!(most_popular_tags)
                write(g, "$(most_popular_tags[1])\n")
            end
        end
    end
end

max_stars_repo_name: PhilChodrow/HypergraphModularity

max_issues_repo_name: nveldt/TestPkg

max_forks_repo_name: PhilChodrow/HypergraphModularity


text:
"""
    copy(x::MultiAssayExperiment)

Return a copy of `x`, where all components are identically-same as those in `x`.

# Examples
```jldoctest
julia> using MultiAssayExperiments, DataFrames

julia> x = exampleobject();

julia> x2 = copy(x);

julia> setsampledata!(x2, DataFrame(name=["A", "B"]));

julia> size(sampledata(x))
(5, 2)

julia> size(sampledata(x2)) # Change to reference is only reflected in x2.
(2, 1)

julia> stuff = experiments(x);

julia> delete!(stuff, "bar");

julia> collect(keys(experiments(x2)))
1-element Vector{String}:
 "foo"
```
"""
function Base.copy(x::MultiAssayExperiment)
    output = MultiAssayExperiment()

    output.experiments = x.experiments
    output.sampledata = x.sampledata
    output.samplemap = x.samplemap
    output.metadata = x.metadata

    return output
end

"""
    deepcopy(x::MultiAssayExperiment)

Return a deep copy of `x` and all of its components.

# Examples
```jldoctest
julia> using MultiAssayExperiments, DataFrames

julia> x = exampleobject();

julia> x2 = deepcopy(x);

julia> insertcols!(sampledata(x), 2, "WHEE" => 1:5); # References now point to different objects.

julia> names(sampledata(x2))
2-element Vector{String}:
 "name"
 "disease"
```
"""
function Base.deepcopy(x::MultiAssayExperiment)
    output = MultiAssayExperiment()

    output.experiments = deepcopy(x.experiments)
    output.sampledata = deepcopy(x.sampledata)
    output.samplemap = deepcopy(x.samplemap)
    output.metadata = deepcopy(x.metadata)

    return output
end

function scat(io::IO, names::Vector{String})
    if length(names) < 5
        for n in names
            print(io, " " * n)
        end
    else
        print(io, " " * names[1])
        print(io, " " * names[2])
        print(io, " ...")
        print(io, " " * names[length(names)-1])
        print(io, " " * names[length(names)])
    end
end

"""
    print(io::IO, x::MultiAssayExperiment)

Print a summary of `x`.
"""
function Base.show(io::IO, x::MultiAssayExperiment)
    print(io, string(typeof(x)) * " object\n")

    print(io, "  experiments(" * string(length(experiments(x))) * "):")
    scat(io, collect(keys(experiments(x))))
    print(io, "\n")

    print(io, "  " * "sampledata(" * string(size(sampledata(x))[2]) * "):")
    scat(io, names(sampledata(x)))
    print(io, "\n")

    print(io, "  " * "metadata(" * string(length(metadata(x))) * "):")
    scat(io, collect(keys(metadata(x))))
end


text:
""" Define a schema and ACSet which can be used as the base for most individual-based epidemiological models.
    Alone, this module allows for simulation of Markov models.
"""
module SchemaBase

export TheoryIBM, AbstractIBM, IBM,
    npeople, nstate, statelabel, get_index_state,
    queue_state_update, apply_state_updates,
    render_states,
    initialize_states, reset_states,
    create_state_update, create_attr_update,
    simulation_loop

using Catlab
using Catlab.CategoricalAlgebra
using Catlab.CategoricalAlgebra.FinSets
using Catlab.Present
using Catlab.Theories

""" ACSet definition for a basic individual-based model
    See [Catlab.jl documentation](https://algebraicjulia.github.io/Catlab.jl/stable/generated/sketches/smc/#Presentations)
    for description of the `@present` syntax.
"""
@present TheoryIBM(FreeSchema) begin
    Person::Ob
    State::Ob

    state::Hom(Person, State)
    state_update::Hom(Person, State)

    StateLabel::AttrType
    statelabel::Attr(State, StateLabel)
end

""" An abstract ACSet for a basic Markov individual-based model.
"""
@abstract_acset_type AbstractIBM

"""
A concrete ACSet for a basic Markov individual-based model inheriting from `AbstractIBM`.
"""
@acset_type IBM(TheoryIBM, index = [:state, :state_update], unique_index = [:statelabel]) <: AbstractIBM

"""
    npeople(model::AbstractIBM, states)

Return the number of people in some set of `states` (an element of the State Ob).
If called without the argument `states`, simply return the total population size.
"""
npeople(model::AbstractIBM) = nparts(model, :Person)
npeople(model::AbstractIBM, states) = length(incident(model, states, [:state, :statelabel]))

"""
    nstate(model::AbstractIBM)

Return the size of the finite state space.
"""
nstate(model::AbstractIBM) = nparts(model, :State)

"""
    statelabel(model::AbstractIBM)

Return the labels (names) of the states in the finite state space.
"""
statelabel(model::AbstractIBM) = subpart(model, :statelabel)

"""
    get_index_state(model::AbstractIBM, states)

Return an integer vector giving the persons who are in the states specified in `states`.
If called without the argument `states`, simply return everyone's index.
"""
get_index_state(model::AbstractIBM, states) = incident(model, states, [:state, :statelabel])
get_index_state(model::AbstractIBM) = parts(model, :Person)

"""
    queue_state_update(model::AbstractIBM, persons, state)

For persons specified in `persons`, queue a state update to `state`, which will be applied at the
end of the time step.
"""
function queue_state_update(model::AbstractIBM, persons, state)
    if length(persons) > 0
        state_index = only(incident(model, state, :statelabel))
        state_index > 0 || throw(ArgumentError("state $(state) is not is the set of state labels"))
        set_subpart!(model, persons, :state_update, state_index)
    end
end

"""
    render_states(model::AbstractIBM, steps::Integer)

Return a tuple whose first element is a matrix containing counts of
states (columns) by time step (rows), and whose second element is a _process_
function which can be used in the simulation loop.
"""
function render_states(model::AbstractIBM, steps::Integer)
    out = Array{Int64}(undef, steps, nstate(model))

    output_states(t::Int) = begin
        out[t, :] = [length(incident(model, i, :state)) for i = parts(model, :State)]
    end

    return (out, output_states)
end

"""
    initialize_states(model::AbstractIBM, initial_states, state_labels::Vector{String})

Initialize the categorical states of a model. The argument `initial_states` can either
be provided as a vector of integers, corresponding to the internal storage of the ACSet,
or as a vector of strings. It should be equal in length to the population which is to be
simulated.
"""
function initialize_states(model::AbstractIBM, initial_states::Vector{T}, state_labels::Vector{String}) where {T <: Integer}
    length(unique(initial_states)) <= length(state_labels) || throw(ArgumentError("'initial_states' has more unique values than 'state_labels', please fix"))
    if nparts(model, :State) > 0
        reset_states(model, initial_states)
    else
        add_parts!(model, :State, length(state_labels), statelabel = state_labels)
        people = add_parts!(model, :Person, length(initial_states))
        set_subpart!(model, people, :state, initial_states);
    end
end

function initialize_states(model::AbstractIBM, initial_states::Vector{String}, state_labels::Vector{String})
    length(unique(initial_states)) <= length(state_labels) || throw(ArgumentError("'initial_states' has more unique values than 'state_labels', please fix"))
    if nparts(model, :State) > 0
        reset_states(model, initial_states)
    else
        add_parts!(model, :State, length(state_labels), statelabel = state_labels)
        people = add_parts!(model, :Person, length(initial_states))
        set_subpart!(model, people, :state, indexin(initial_states, state_labels));
    end
end

"""
    initialize_states(model::AbstractIBM, initial_states)

Reset a model's categorical states.
"""
function reset_states(model::AbstractIBM, initial_states::Vector{T}) where {T <: Integer}
    nparts(model, :Person) == length(initial_states) || throw(ArgumentError("'initial_states' must be equal to the number of persons in the model"))
    set_subpart!(model, :state_update, 0)
    set_subpart!(model, :state, initial_states);
end

function reset_states(model::AbstractIBM, initial_states::Vector{String})
    nparts(model, :Person) == length(initial_states) || throw(ArgumentError("'initial_states' must be equal to the number of persons in the model"))
    set_subpart!(model, :state_update, 0)
    set_subpart!(model, :state, indexin(initial_states, subpart(model, :statelabel)));
end

"""
    create_state_update(model::AbstractIBM)

Return a function that is called with no arguments that updates any `Ob` that has an accompanying update `Ob`.
For example, if you provide both `risk` and `risk_update`, then `risk_update` will be used to update `risk` and
then cleared.
"""
function create_state_update(model::AbstractIBM)
    s = acset_schema(model)
    homs = s.homs
    homs = map((x)->String(x), homs)
    homs = split.(homs, "_")

    # all the homs that end in 'update'
    update_ix = findall(homs) do x
        if length(x) < 2
            return false
        else
            return x[end] == "update"
        end
    end

    # membership homs that correspond to the update homs
    state_ix = map(update_ix) do x
        for i = 1:length(homs)
            if i == x
                continue
            else
                if homs[x][1:end-1] == homs[i]
                    return i
                end
            end
        end
    end

    # the obs that the update homs are updating
    state_obs = map(update_ix) do x
        s.codoms[s.homs[x]]
    end

    all(.!isnothing.(state_ix)) || throw(AssertionError("some update homs do not have corresponding membership homs, please check your schema"))
    length(update_ix) == length(state_ix) == length(state_obs) || throw(AssertionError("some update homs do not have corresponding membership homs, please check your schema"))

    state_homs = s.homs[state_ix]
    update_homs = s.homs[update_ix]

    function update()
        # all Obs that need updating
        for ob in 1:length(state_obs)
            # all particular states in that Ob
            for state in parts(model, state_obs[ob])
                people_to_update = incident(model, state, update_homs[ob])
                if length(people_to_update) > 0
                    set_subpart!(model, people_to_update, state_homs[ob], state)
                end
            end
            set_subpart!(model, update_homs[ob], 0)
        end
    end

    return update
end

"""
create_attr_update(model::AbstractIBM)

Return a function that is called with no arguments that updates any `Attr` that has an accompanying update `Attr`.
For example, if you provide both `antibody` and `antibody_titre` as `Attr` objects from `Person` to the apropriate `AttrType`,
this function will update `antibody` to the values in `antibody_titre` on each time step. This requires that when the model
is initialized, each attribute and its update have approximately the same starting values.
"""
function create_attr_update(model::AbstractIBM)
    s = acset_schema(model)

    attr_sym = s.attrs

    attrs = String.(s.attrs)
    attrs = split.(attrs, "_")

    update_ix = findall(attrs) do x
        if length(x) < 2
            return false
        else
            return x[end] == "update"
        end
    end

    # if there are attributes which will be updated
    if length(update_ix) > 0

        attr_ix = map(update_ix) do x
            for i = 1:length(attrs)
                if i == x
                    continue
                else
                    if attrs[x][1:end-1] == attrs[i]
                        return i
                    end
                end
            end
        end

        all(.!isnothing.(attr_ix)) || throw(AssertionError("some update attributes do not have corresponding current attributes, please check your schema"))
        length(update_ix) == length(attr_ix) || throw(AssertionError("some update attributes do not have corresponding current attributes, please check your schema"))

        # check that the update and current attrs are approx equal
        for i = 1:length(attr_ix)
            if !isapprox(subpart(model, attr_sym[update_ix[i]]), subpart(model, attr_sym[attr_ix[i]]))
                throw(AssertionError("all update attributes must be approximately equal to their current attribute values, please check this is the case"))
            end
        end

        function update()
            for i = 1:length(attr_ix)
                set_subpart!(model, attr_sym[attr_ix[i]], subpart(model, attr_sym[update_ix[i]]))
            end
        end

        return update

    else
        # no attributes with updates

        # null update fn
        function update_null()

        end

        return update_null
    end
end

"""
    simulation_loop(model::AbstractIBM, processes::Union{Function, AbstractVector{Function}}, steps::Integer)

A simple predefined simulation loop for basic (no events) individual based models. Processes are called first,
followed by state updates.
"""
function simulation_loop(model::AbstractIBM, processes::Union{Function, AbstractVector{Function}}, steps::Integer)
    if processes isa Function
        processes = [processes]
    end

    apply_attr_updates = create_attr_update(model)
    apply_state_updates = create_state_update(model)

    for t = 1:steps
        for p = processes
            p(t)
        end
        apply_attr_updates()
        apply_state_updates()
    end
end

end


text:
import Dagger: DArray, ArrayDomain, Thunk, chunks, domainchunks

# Domain munging conveniences

struct ChunkDomain
  domain::AbstractVector{Int}
  chunks::Vector{<:AbstractVector{Int}}
end

ArrayDomain(d::ChunkDomain...) = ArrayDomain(map(d -> d.domain, d)...)

mapcart(f, xs) = f.(xs)
mapcart(f, xs, ys) = f.(xs, reshape(xs, 1, :))
domprod(ds...) = mapcart(tuple, ds...)

chunkinds(i::ChunkDomain...) = domprod(map(i -> i.chunks, i)...)

chunkinds(xs::DArray, dim::Integer) =
  map(d -> d.indexes[dim],
      domainchunks(xs)[ntuple(i -> i == dim ? (:) : 1 , ndims(xs))...])

domain(xs::DArray, i::Integer) = ChunkDomain(indices(xs, i), chunkinds(xs, i))

# Constructors

DArray{T,N}(dom::NTuple{N,ChunkDomain}, chunks::AbstractArray{<:Any,N}) where {T,N} =
  DArray(T, ArrayDomain(dom...), ArrayDomain.(chunkinds(dom...)), chunks)

DArray{T,N}(f, dom::NTuple{N,ChunkDomain}) where {T,N} =
  DArray{T,N}(dom, f.(chunkinds(dom...)))

# Views

using Dagger: lookup_parts

function catchunks(chs)
  for i = 1:ndims(chs)
    chs = mapslices(xs -> [cat(i, xs...)], chs, i)
  end
  return chs[1]
end

function chslice(xs::DArray, d::ArrayDomain)
  subchunks, subdomains = lookup_parts(chunks(xs), domainchunks(xs), d)
  chsize = size(subdomains)
  subchunks
  Thunk(subchunks...) do subchunks...
    catchunks(reshape(collect(subchunks), chsize))
  end
end

chslice(xs::DArray, i...) = chslice(xs, ArrayDomain(i))

max_stars_repo_name: JuliaTagBot/Tokamak

max_issues_repo_name: JuliaTagBot/Tokamak

max_forks_repo_name: FluxML/Tokamak


text:
using Gadfly

hstack(plot(y=[1,2,3],shape=[Shape.vline], Theme(discrete_highlight_color=x->x)),
       plot(y=[1,2,3],shape=[Shape.hline], Theme(discrete_highlight_color=x->x)))


text:
# The license for this code is available at https://github.com/avik-pal/FastStyleTransfer.jl/blob/master/LICENSE.md

#-----------------------------Transformer Net--------------------------------------

# Replace Instance Norm with Batch Norm until GPU speed issue is fixed

ConvPad(chs::Pair{<:Int,<:Int}, kernel::Tuple{Int,Int}, stride::Tuple{Int,Int} = (1,1)) =
    Conv(kernel, chs, stride = stride, pad = (kernel[1]÷2, kernel[2]÷2))

TransformerNet() = Chain(ConvPad(3=>32, (3,3)),
                         BatchNorm(32, relu),
                         ConvPad(32=>64, (3,3), (2,2)),
                         BatchNorm(64, relu),
                         ConvPad(64=>128, (3,3), (2,2)),
                         BatchNorm(128, relu),
                         [ResidualBlock(128) for i in 1:5]...,
                         UpsamplingBlock(128=>64, (3,3), (1,1), 2),
                         BatchNorm(64),
                         UpsamplingBlock(64=>32, (3,3), (1,1), 2),
                         BatchNorm(32),
                         ConvPad(32=>3, (9,9), (1,1)))

TransformerNet2() = Chain(ConvPad(3=>32, (3,3)),
                          BatchNorm(32, relu),
                          ConvPad(32=>64, (3,3), (2,2)),
                          BatchNorm(64, relu),
                          ConvPad(64=>128, (3,3), (2,2)),
                          BatchNorm(128, relu),
                          [ResidualBlock(128) for i in 1:5]...,
                          ConvTranspose((3, 3), 128=>64, stride = (2, 2), pad = (1, 1), output_pad = (1, 1)),
                          BatchNorm(64, relu),
                          ConvTranspose((3, 3), 64=>32, stride = (2, 2), pad = (1, 1), output_pad = (1, 1)),
                          BatchNorm(32, relu),
                          ConvTranspose((3, 3), 32=>3, pad = (1, 1)))

TransformerNet3() = Chain(ConvPad(3=>32, (3,3)),
                          InstanceNorm(32, relu),
                          ConvPad(32=>64, (3,3), (2,2)),
                          InstanceNorm(64, relu),
                          ConvPad(64=>128, (3,3), (2,2)),
                          InstanceNorm(128, relu),
                          [ResidualBlock(128) for i in 1:5]...,
                          ConvTranspose((3, 3), 128=>64, stride = (2, 2), pad = (1, 1), output_pad = (1, 1)),
                          InstanceNorm(64, relu),
                          ConvTranspose((3, 3), 64=>32, stride = (2, 2), pad = (1, 1), output_pad = (1, 1)),
                          InstanceNorm(32, relu),
                          ConvTranspose((3, 3), 32=>3, pad = (1, 1)))


text:
include("../examples/bubblesort.jl")
a = rand(1:100,100)
@test issorted(sort!(a;alg=BubbleSort))

include("../examples/enum.jl")
@enum TestEnum TestEnum1 TestEnum2 TestEnum3
@test [TestEnum1.n,TestEnum2.n,TestEnum3.n] == [0,1,2]

include("../examples/lru.jl")
include("../examples/lru_test.jl")

include("../examples/modint.jl")
b = ModInts.ModInt{10}(2)
c = ModInts.ModInt{10}(4)
@test b + c == ModInts.ModInt{10}(6)
@test c - b == ModInts.ModInt{10}(2)

include("../examples/ndgrid.jl")
r = repmat(1:10,1,10)
r1, r2 = ndgrid(1:10, 1:10)
@test r1 == r
@test r2 == r'
r3, r4 = meshgrid(1:10,1:10)
@test r3 == r'
@test r4 == r

include("../examples/quaternion.jl")
q = Quaternions.Quaternion(1,0,0,0)
x = Quaternions.Quaternion(0,1,1,1)
@test q*2.0+2 == Quaternions.Quaternion(4,0,0,0)
@test abs((-q+x*2)/4) == 0.9013878188659973

include("../examples/queens.jl")
@test solve(8, 8, 1) == Array{Int,1}[[1,1]]
@test solve(8, 8, 7) == Array{Int,1}[[1,1],[2,3],[3,5],[4,2],[5,8],[7,4],[8,7]]

# At least make sure code laods
include("../examples/plife.jl")

include("../examples/preduce.jl")

include("../examples/wordcount.jl")

max_stars_repo_name: alanedelman/julia

max_issues_repo_name: Aerlinger/julia

max_forks_repo_name: Aerlinger/julia


text:
function label_propagation{V,T<:Real}(graph::AbstractGraph{V};
                                      weights::Vector{T} = Array(Float64, 0),
                                      proximity::Vector{T} = Array(Float64, 0),
                                      initial = Array(Int, 0),
                                      fixed = Array(Bool, 0),
                                      λ::Real = 1.0)
    !is_directed(graph) || error("graph must be undirected.")
    if !isempty(proximity)
        @graph_requires graph edge_map incidence_list
        if length(proximity) != num_edges(graph)
            error("Invalid proximity vector length")
        #elseif minimum(proximity) < 0 || maximum(proximity) > 1
        #    error("Proximity must between 0 and 1")
        end
        mix_label_propagation!(graph, proximity, _initial(graph, initial, fixed)..., λ)
    else
        if !isempty(weights)
            @graph_requires graph edge_map incidence_list
            if length(weights) != num_edges(graph)
                error("Invalid weight vector length")
            elseif minimum(weights) < 0
                error("Weights must be non-negative")
            end
            weighted_label_propagation!(graph, weights, _initial(graph, initial, fixed)...)
        else
            unweighted_label_propagation!(graph, _initial(graph, initial, fixed)...)
        end
    end
end

function _initial{V}(graph::AbstractGraph{V},
                     initial::Vector{Int} = Array(Int, 0),
                     fixed::Set{V} = Set{V}())
    @graph_requires graph vertex_list vertex_map

    N = num_vertices(graph)
    membership = Array(Int, N)
    not_fixed_nodes = Array(V, 0)

    # Do some initial checks
    if !isempty(fixed) && isempty(initial)
        warn("Ignoring fixed vertices as no initial labeling given")
    end

    if !isempty(initial)
        length(initial) == N || error("Invalid initial labeling vector length")
        for i=1:N
            membership[i] = initial[i] > 0 ? initial[i] : 0
        end

        for v in vertices(graph)
            if in(v, fixed)
                if membership[vertex_index(v, graph)] < 1
                    warn("Fixed nodes cannot be unlabeled, ignoring them")
                    push!(not_fixed_nodes, v)
                end
            else
                push!(not_fixed_nodes, v)
            end
        end

        i = maximum(membership)
        i <= N || error("elements of the initial labeling vector must be between 0 and |V|")
        i > 0 || error("at least one vertex must be labeled in the initial labeling")
    else
        membership = [1:N]
        not_fixed_nodes = collect(vertices(graph))
    end

    return membership, not_fixed_nodes
end

function _initial{V}(graph::AbstractGraph{V},
                     initial::Vector{Int} = Array(Int, 0),
                     fixed::Vector{Bool} = Array(Bool, 0))
    @graph_requires graph vertex_list vertex_map

    N = num_vertices(graph)
    membership = Array(Int, N)
    not_fixed_nodes = Array(V, 0)

    # Do some initial checks
    if !isempty(fixed) && isempty(initial)
        warn("Ignoring fixed vertices as no initial labeling given")
    end

    if !isempty(initial)
        length(initial) == N || error("Invalid initial labeling vector length")
        for i=1:N
            membership[i] = initial[i] > 0 ? initial[i] : 0
        end

        if !isempty(fixed)
            length(fixed) == N || error("Invalid fixed node vector length")
            for v in vertices(graph)
                v_idx = vertex_index(v, graph)
                if fixed[v_idx]
                    if membership[v_idx] < 1
                        warn("Fixed nodes cannot be unlabeled, ignoring them")
                        push!(not_fixed_nodes, v)
                    end
                else
                    push!(not_fixed_nodes, v)
                end
            end
        else
            not_fixed_nodes = collect(vertices(graph))
        end

        i = maximum(membership)
        i <= N || error("elements of the initial labeling vector must be between 0 and |V|")
        i > 0 || error("at least one vertex must be labeled in the initial labeling")
    else
        membership = [1:N]
        not_fixed_nodes = collect(vertices(graph))
    end

    membership, not_fixed_nodes
end

function _initial{V}(graph::AbstractGraph{V},
                     initial::Dict{V,Int} = Dict{V,Int}(),
                     fixed::Vector{Bool} = Array(Bool, 0))
    @graph_requires graph vertex_list vertex_map

    N = num_vertices(graph)
    membership = Array(Int, N)
    not_fixed_nodes = Array(V, 0)

    # Do some initial checks
    if !isempty(fixed) && isempty(initial)
        warn("Ignoring fixed vertices as no initial labeling given")
    end

    if !isempty(initial)
        for v in vertices(graph)
            v_idx = vertex_index(v, graph)
            if haskey(initial, v)
                membership[v_idx] = initial[v]
            else
                membership[v_idx] = 0
            end
        end

        if !isempty(fixed)
            length(fixed) == N || error("Invalid fixed node vector length")
            for v in vertices(graph)
                v_idx = vertex_index(v, graph)
                if fixed[v_idx]
                    if membership[v_idx] < 1
                        warn("Fixed nodes cannot be unlabeled, ignoring them")
                        push!(not_fixed_nodes, v)
                    end
                else
                    push!(not_fixed_nodes, v)
                end
            end
        else
            not_fixed_nodes = collect(vertices(graph))
        end

        i = maximum(membership)
        i <= N || error("elements of the initial labeling vector must be between 0 and |V|")
        i > 0 || error("at least one vertex must be labeled in the initial labeling")
    else
        membership = [1:N]
        not_fixed_nodes = collect(vertices(graph))
    end

    return membership, not_fixed_nodes
end

function _initial{V}(graph::AbstractGraph{V},
                     initial::Dict{V,Int} = Dict{V,Int}(),
                     fixed::Set{V} = Set{V}())
    @graph_requires graph vertex_list vertex_map

    N = num_vertices(graph)
    membership = Array(Int, N)
    not_fixed_nodes = Array(V, 0)

    # Do some initial checks
    if !isempty(fixed) && isempty(initial)
        warn("Ignoring fixed vertices as no initial labeling given")
    end

    if !isempty(initial)
        for v in vertices(graph)
            v_idx = vertex_index(v, graph)
            if haskey(initial, v)
                membership[v_idx] = initial[v]
            else
                membership[v_idx] = 0
            end
        end

        for v in vertices(graph)
            if in(v, fixed)
                if membership[vertex_index(v, graph)] < 1
                    warn("Fixed nodes cannot be unlabeled, ignoring them")
                    push!(not_fixed_nodes, v)
                end
            else
                push!(not_fixed_nodes, v)
            end
        end

        i = maximum(membership)
        i <= N || error("elements of the initial labeling vector must be between 0 and |V|")
        i > 0 || error("at least one vertex must be labeled in the initial labeling")
    else
        membership = [1:N]
        not_fixed_nodes = collect(vertices(graph))
    end

    membership, not_fixed_nodes
end

function weighted_label_propagation!{V,T<:Real}(graph::AbstractGraph{V},
                                                weights::Vector{T},
                                                membership::Vector{Int},
                                                not_fixed_nodes::Vector{V})
    @graph_requires graph edge_map incidence_list vertex_map

    label_counters = zeros(T, num_vertices(graph))
    dominant_labels = Int[]
    nonzero_labels = Int[]

    running = true
    STEP = 0
    while running
        STEP += 1
        running = false

        # Shuffle the node ordering vector
        X = shuffle(not_fixed_nodes)

        # In the prescribed order, loop over the vertices and reassign labels
        for v in X
            v_idx = vertex_index(v, graph)

            # Clear dominant_labels and nonzero_labels
            dominant_labels = Int[]
            nonzero_labels = Int[]

            # recount
            max_count = zero(T)

            for w in out_edges(v, graph)
                k = membership[vertex_index(target(w, graph), graph)]

                # skip if it has no label yet
                k != 0 || continue

                was_zero = label_counters[k] == zero(T)
                label_counters[k] += weights[edge_index(w, graph)]
                if was_zero && label_counters[k] != zero(T)
                    # counter just became nonzero
                    push!(nonzero_labels, k)
                end
                if max_count < label_counters[k]
                    max_count = label_counters[k]
                    resize!(dominant_labels, 1)
                    dominant_labels[1] = k
                elseif max_count == label_counters[k]
                    push!(dominant_labels, k)
                end
            end

            if !isempty(dominant_labels)

                # Select randomly from the dominant labels
                k = dominant_labels[rand(1:length(dominant_labels))]

                # Check if the current label of the node is also dominant
                if label_counters[membership[v_idx]] != max_count
                    # Nope, we need at least one more iteration
                    running = true
                end

                # Update label of the current node
                membership[v_idx] = k
            end
            # Clear the nonzero elements in label_counters
            for i in nonzero_labels
                label_counters[i] = zero(T)
            end
        end
    end
    permute_labels!(membership), STEP
end

function nei_label_propagation{V,T<:Real}(graph::AbstractGraph{V},
                                      common_neighbors::Vector{T};
                                          no_common_neighbors::Vector{T} = Array(Float64, 0),
                                      proximity::Vector{T} = Array(Float64, 0),
                                      initial = Array(Int, 0),
                                      fixed = Array(Bool, 0),
                                      λ::Real = 1.0)
    !is_directed(graph) || error("graph must be undirected.")
    if length(common_neighbors) != num_edges(graph)
        error("Invalid weight vector length")
    end
    if isempty(no_common_neighbors)
        @graph_requires graph edge_map incidence_list
        nei_label_propagation!(graph, common_neighbors, _initial(graph, initial, fixed)..., λ)
    else
        if length(no_common_neighbors) != num_edges(graph)
            error("Invalid no_common_neighbors vector length")
        end
        nei_label_propagation!(graph, common_neighbors, no_common_neighbors, _initial(graph, initial, fixed)..., λ)
    end

end

function nei_label_propagation!{V,T<:Real}(graph::AbstractGraph{V},
                                           common_neighbors::Vector{T},
                                           membership::Vector{Int},
                                           not_fixed_nodes::Vector{V},
                                           λ::T)
    @graph_requires graph edge_map incidence_list vertex_map

    label_counters = zeros(T, num_vertices(graph))
    dominant_labels = Int[]
    nonzero_labels = Int[]

    running = true
    STEP = 0
    while running
        STEP += 1
        running = false

        # Shuffle the node ordering vector
        X = shuffle(not_fixed_nodes)

        # In the prescribed order, loop over the vertices and reassign labels
        for v in X
            v_idx = vertex_index(v, graph)

            # Clear dominant_labels and nonzero_labels
            dominant_labels = Int[]
            nonzero_labels = Int[]

            # recount
            max_count = zero(T)

            for w in out_edges(v, graph)
                k = membership[vertex_index(target(w, graph), graph)]

                # skip if it has no label yet
                k != 0 || continue

                was_zero = label_counters[k] == zero(T)
                label_counters[k] += 1 + λ*common_neighbors[edge_index(w, graph)]
                if was_zero && label_counters[k] != zero(T)
                    # counter just became nonzero
                    push!(nonzero_labels, k)
                end
                if max_count < label_counters[k]
                    max_count = label_counters[k]
                    resize!(dominant_labels, 1)
                    dominant_labels[1] = k
                elseif max_count == label_counters[k]
                    push!(dominant_labels, k)
                end
            end

            if !isempty(dominant_labels)

                # Select randomly from the dominant labels
                k = dominant_labels[rand(1:length(dominant_labels))]

                # Check if the current label of the node is also dominant
                if label_counters[membership[v_idx]] != max_count
                    # Nope, we need at least one more iteration
                    running = true
                end

                # Update label of the current node
                membership[v_idx] = k
            end
            # Clear the nonzero elements in label_counters
            for i in nonzero_labels
                label_counters[i] = zero(T)
            end
        end
    end
    permute_labels!(membership), STEP
end

function nei_label_propagation!{V,T<:Real}(graph::AbstractGraph{V},
                                           common_neighbors::Vector{T},
                                           no_common_neighbors::Vector{T},
                                           membership::Vector{Int},
                                           not_fixed_nodes::Vector{V},
                                           λ::T)
    @graph_requires graph edge_map incidence_list vertex_map

    label_counters = zeros(T, num_vertices(graph))
    sub_label_counters = zeros(T, num_vertices(graph))
    dominant_labels = Int[]
    nonzero_labels = Int[]

    running = true
    STEP = 0
    while running
        STEP += 1
        running = false

        # Shuffle the node ordering vector
        X = shuffle(not_fixed_nodes)

        # In the prescribed order, loop over the vertices and reassign labels
        for v in X
            v_idx = vertex_index(v, graph)

            # Clear dominant_labels and nonzero_labels
            dominant_labels = Int[]
            nonzero_labels = Int[]

            # recount
            max_count = zero(T)

            for w in out_edges(v, graph)
                k = membership[vertex_index(target(w, graph), graph)]

                # skip if it has no label yet
                k != 0 || continue

                was_zero = label_counters[k] == zero(T)
                label_counters[k] += 1 + λ*common_neighbors[edge_index(w, graph)]
                sub_label_counters[k] += no_common_neighbors[edge_index(w, graph)]
                if was_zero && label_counters[k] != zero(T)
                    # counter just became nonzero
                    push!(nonzero_labels, k)
                end
                if max_count < label_counters[k]
                    max_count = label_counters[k]
                    resize!(dominant_labels, 1)
                    dominant_labels[1] = k
                elseif max_count == label_counters[k]
                    push!(dominant_labels, k)
                end
            end

            if !isempty(dominant_labels)

                # Select randomly from the dominant labels
                #k = dominant_labels[rand(1:length(dominant_labels))]
                k = dominant_labels[findmin(sub_label_counters[dominant_labels])[2]]

                # Check if the current label of the node is also dominant
                if label_counters[membership[v_idx]] != max_count
                    # Nope, we need at least one more iteration
                    running = true
                end

                # Update label of the current node
                membership[v_idx] = k
            end
            # Clear the nonzero elements in label_counters
            for i in nonzero_labels
                label_counters[i] = zero(T)
            end
        end
    end
    permute_labels!(membership), STEP
end

function unweighted_label_propagation!{V}(graph::AbstractGraph{V},
                                          membership::Vector{Int},
                                          not_fixed_nodes::Vector{V})
    @graph_requires graph vertex_map adjacency_list

    label_counters = zeros(Int, num_vertices(graph))
    dominant_labels = Int[]
    nonzero_labels = Int[]

    running = true
    STEP = 0
    while running
        STEP += 1
        running = false

        # Shuffle the node ordering vector
        X = shuffle(not_fixed_nodes)

        # In the prescribed order, loop over the vertices and reassign labels
        for v in X
            v_idx = vertex_index(v, graph)

            # Clear dominant_labels and nonzero_labels
            dominant_labels = Int[]
            nonzero_labels = Int[]

            # recount
            max_count = 0

            for w in out_neighbors(v, graph)
                k = membership[vertex_index(w, graph)]

                # skip if it has no label yet
                k != 0 || continue

                label_counters[k] += 1

                # counter just became nonzero
                label_counters[k] != 1 || push!(nonzero_labels, k)

                if max_count < label_counters[k]
                    max_count = label_counters[k]
                    resize!(dominant_labels, 1)
                    dominant_labels[1] = k
                elseif max_count == label_counters[k]
                    push!(dominant_labels, k)
                end
            end

            if length(dominant_labels) > 0

                # Select randomly from the dominant labels
                k = dominant_labels[rand(1:length(dominant_labels))]

                # Check if the current label of the node is also dominant
                if label_counters[membership[v_idx]] != max_count
                    # Nope, we need at least one more iteration
                    running = true
                end

                # Update label of the current node
                membership[v_idx] = k
            end
            # Clear the nonzero elements in label_counters
            for i in nonzero_labels
                label_counters[i] = 0
            end
        end
    end
    permute_labels!(membership), STEP
end

function mix_label_propagation!{V,T<:Real}(graph::AbstractGraph{V},
                                           proximity::Vector{T},
                                           membership::Vector{Int},
                                           not_fixed_nodes::Vector{V},
                                           λ::Real)
    @graph_requires graph edge_map incidence_list vertex_map

    label_counters = zeros(T, num_vertices(graph))
    proximity_sum = zeros(T, num_vertices(graph))
    dominant_labels = Int[]
    similarity = Array(T, 0)
    nonzero_labels = Int[]

    running = true
    STEP = 0
    while running
        STEP += 1
        running = false

        #if STEP <= 0
        #    degreeX = Int64[out_degree(u, graph) for u in not_fixed_nodes]
        #    X = not_fixed_nodes[sortperm(degreeX, rev=true)]
        #else
            # Shuffle the node ordering vector
            X = shuffle(not_fixed_nodes)
        #end

        # In the prescribed order, loop over the vertices and reassign labels
        for v in X
            v_idx = vertex_index(v, graph)

            # Clear dominant_labels and nonzero_labels
            dominant_labels = Int[]
            nonzero_labels = Int[]

            # recount
            max_count = zero(T)

            for w in out_edges(v, graph)
                w_idx = edge_index(w, graph)
                k = membership[vertex_index(target(w, graph), graph)]

                # skip if it has no label yet
                k != 0 || continue

                was_zero = label_counters[k] == zero(T)
                label_counters[k] += 1
                proximity_sum[k] += proximity[w_idx]
                if was_zero && label_counters[k] != zero(T)
                    # counter just became nonzero
                    push!(nonzero_labels, k)
                end
                if max_count < label_counters[k]
                    max_count = label_counters[k]
                    resize!(dominant_labels, 1)
                    resize!(similarity, 1)
                    dominant_labels[1] = k
                    similarity[1] = proximity_sum[k]
                elseif max_count == label_counters[k]
                    push!(dominant_labels, k)
                    push!(similarity, proximity_sum[k])
                end
            end

            if !isempty(dominant_labels)

                max_similarity = maximum(similarity)
                candidate_labels = rand() <= λ ? dominant_labels[similarity.>=max_similarity] : dominant_labels

                # Select randomly from the dominant labels

                k = candidate_labels[rand(1:length(candidate_labels))]

                # Check if the current label of the node is also dominant
                if label_counters[membership[v_idx]] != max_count
                    # Nope, we need at least one more iteration
                    running = true
                end

                # Update label of the current node
                membership[v_idx] = k
            end
            # Clear the nonzero elements in label_counters
            for i in nonzero_labels
                label_counters[i] = zero(T)
                proximity_sum[i] = zero(T)
            end
        end
    end
    permute_labels!(membership), STEP
end

function permute_labels!(membership::Vector{Int})
    N = length(membership)
    if maximum(membership) > N || minimum(membership) < 1
        error("Label must between 1 and |V|")
    end
    label_counters = zeros(Int, N)
    j = 1
    for i=1:length(membership)
        k = membership[i]
        if k >= 1
            if label_counters[k] == 0
                # We have seen this label for the first time
                label_counters[k] = j
                k = j
                j += 1
            else
                k = label_counters[k]
            end
        end
        membership[i] = k
    end
    membership
end


text:
@testset "causeSelection.jl" begin
    data = CSO.dataset("cso_raw")
    causes = ["P", "U"]
    colCause = "Code"

    @testset "causedBy!(df, ColCause, causes)" begin
        data = CSO.dataset("cso_raw")
        transformedDF = causedBy!(data, colCause, causes)
        newCauses = unique(map(x -> ismissing(x) ||  x in causes, transformedDF[:, Symbol("Code")]))

        #Df contains only the selected causes
        @test length(newCauses) == length(causes)

        #Ancient DF has been changed
        @test nrow(data) == nrow(transformedDF)
    end

    @testset "causedBy(df, ColCause, causes)" begin
        data = DataFrame(CSV.File("C:\\Users\\Alexandrine\\cso_raw.csv"))
        newDF = causedBy(data, colCause, causes)
        newCauses = unique(map(x -> ismissing(x) ||  x in causes, newDF[:, Symbol("Code")]))

        #Df contains only the selected causes
        @test length(newCauses) == 1 && newCauses[1] == 1

        #Ancient DF has been changed
        @test nrow(data) != nrow(newDF)
    end

end

max_stars_repo_name: AlexLsn/CSO

max_issues_repo_name: AlexLsn/CSO

max_forks_repo_name: AlexLsn/CSO


text:
module NamedVectors

using Printf: @printf
using ..Traits: efficientoperations

export NamedVector, HomoNamedVector

"""
    NamedVector

Abstract type for all named vectors.
"""
abstract type NamedVector end

"""
    getindex(nv::NamedVector, index::Int)

Get the value by the `[]` syntax.
"""
@inline Base.getindex(nv::NamedVector, index::Int) = getfield(nv, index)

"""
    setindex!(nv::NamedVector, value, index::Int)

Set the value by the `[]` syntax if mutable.
"""
@inline Base.setindex!(nv::NamedVector, value, index::Int) = setfield!(nv, index, value)

"""
    ==(nv1::NamedVector, nv2::NamedVector) -> Bool

Overloaded equivalent operator. Two named vector are equal to each other if and only if their keys as well as their values are equal to each other.
!!! note
    It is not necessary for two named vectors to be of the same concrete type to be equal to each other.
"""
@inline Base.:(==)(nv1::NamedVector, nv2::NamedVector) = (keys(nv1) == keys(nv2)) && (values(nv1) == values(nv2))

"""
    isequal(nv1::NamedVector, nv2::NamedVector) -> Bool

Overloaded equivalent operator. Two named vector are equal to each other if and only if their keys as well as their values are equal to each other.
!!! note
    It is not necessary for two named vectors to be of the same concrete type to be equal to each other.
"""
@inline Base.isequal(nv1::NamedVector, nv2::NamedVector) = isequal(keys(nv1), keys(nv2)) && isequal(values(nv1), values(nv2))

"""
    <(nv1::NamedVector, nv2::NamedVector) -> Bool

Compare two named vectors and judge whether the first is less than the second.
"""
@inline Base.:<(nv1::NamedVector, nv2::NamedVector) = <(efficientoperations, nv1, nv2)

"""
    isless(nv1::NamedVector, nv2::NamedVector) -> Bool

Compare two named vectors and judge whether the first is less than the second.
"""
@inline Base.isless(nv1::NamedVector, nv2::NamedVector) = isless(efficientoperations, nv1, nv2)

"""
    show(io::IO, nv::NamedVector)

Show a concrete `NamedVector`.
"""
Base.show(io::IO, nv::NamedVector) = @printf io "%s(%s)" nameof(typeof(nv)) join(repr.(values(nv)), ", ")

"""
    hash(nv::NamedVector, h::UInt)

Hash a concrete `NamedVector`.
"""
@inline Base.hash(nv::NamedVector, h::UInt) = hash(values(nv), h)

"""
    convert(::Type{Tuple}, nv::NamedVector) -> Tuple
    convert(::Type{NV}, nv::Tuple) where {NV<:NamedVector} -> NV

Convert a named vector to tuple and vice versa.
"""
@generated Base.convert(::Type{Tuple}, nv::NamedVector) = Expr(:tuple, (:(getfield(nv, $i)) for i = 1:fieldcount(nv))...)
function Base.convert(::Type{NV}, nv::Tuple) where NV<:NamedVector
    @assert fieldcount(NV) == length(nv) "convert error: mismatched length between $NV($(fieldcount(NV))) and input tuple($(length(nv)))."
    return NV(nv...)
end

"""
    length(::Type{NV}) where NV<:NamedVector -> Int
    length(nv::NamedVector) -> Int

Get the length of a concrete `NamedVector`.
"""
@inline Base.length(::Type{NV}) where {NV<:NamedVector} = fieldcount(NV)
@inline Base.length(nv::NamedVector) = length(typeof(nv))

"""
    zero(::Type{NV}) where NV<:NamedVector -> NV
    zero(nv::NamedVector) -> typeof(nv)

Get a concrete `NamedVector` with all values being zero.
"""
@inline @generated function Base.zero(::Type{NV}) where NV<:NamedVector
    zeros = (zero(fieldtype(NV, i)) for i = 1:fieldcount(NV))
    return :(NV($(zeros...)))
end
@inline Base.zero(nv::NamedVector) = zero(typeof(nv))

"""
    iterate(nv::NamedVector, state=1)
    iterate(rv::Iterators.Reverse{<:NamedVector}, state=length(rv.itr))

Iterate or reversely iterate over the values of a concrete `NamedVector`.
"""
@inline Base.iterate(nv::NamedVector, state=1) = (state > length(nv)) ? nothing : (getfield(nv, state), state+1)
@inline Base.iterate(rv::Iterators.Reverse{<:NamedVector}, state=length(rv.itr)) = (state < 1) ? nothing : (getfield(rv.itr, state), state-1)

"""
    keys(nv::NamedVector) -> NTuple{length(nv), Symbol}

Iterate over the names.
"""
@inline @generated Base.keys(nv::NamedVector) = fieldnames(nv)

"""
    values(nv::NamedVector) -> NTuple{length(nv)}

Iterate over the values.
"""
@inline Base.values(nv::NamedVector) = convert(Tuple, nv)

"""
    pairs(nv::NamedVector) -> Base.Generator

Iterate over the name-value pairs.
"""
@inline Base.pairs(nv::NamedVector) = Base.Generator(=>, keys(nv), values(nv))

"""
    replace(nv::NamedVector; kwargs...) -> typeof(nv)

Return a copy of a concrete `NamedVector` with some of the field values replaced by the keyword arguments.
"""
@inline Base.replace(nv::NamedVector; kwargs...) = replace(efficientoperations, nv; kwargs...)

"""
    map(f, nvs::NV...) where NV<:NamedVector -> NV

Apply function `f` element-wisely on the input named vectors.
"""
@generated function Base.map(f, nvs::NV...) where NV<:NamedVector
    exprs = Vector{Expr}(undef, fieldcount(NV))
    for i = 1:length(exprs)
        tmp = [:(nvs[$j][$i]) for j = 1:length(nvs)]
        exprs[i] = :(f($(tmp...)))
    end
    return :(($NV)($(exprs...)))
end

"""
    HomoNamedVector{T}

Abstract type for all homogeneous named vectors.
"""
abstract type HomoNamedVector{T} <: NamedVector end

"""
    eltype(::Type{<:HomoNamedVector{T}}) where T
    eltype(nv::HomoNamedVector)

Get the type parameter of a concrete `HomoNamedVector`.
"""
@inline Base.eltype(::Type{<:HomoNamedVector{T}}) where T = T
@inline Base.eltype(nv::HomoNamedVector) = eltype(typeof(nv))

end #module


text:
using Base64
# helper functions for crypto_
export seal

preparepubkey(publickey::String) = base64decode(publickey)
preparepubkey(publickey::Vector{UInt8}) = publickey

"""
    seal(message::Base.SecretBuffer, publickey)

Takes a `SecretBuffer` holding a secret message and a base64 encoded public key from the remote.
All remaining arguments to call `libsodium.crypto_box_seal` are inferred automatically.
After that call, message is `shred!`ed and the result is returned as a base64 encoded string.
"""
function seal(message::Base.SecretBuffer, publickey)
    seekstart(message)
    len = bytesavailable(message)
    ciphertext = Vector{Cuchar}(undef, crypto_box_SEALBYTES + len)
    binpublickey = preparepubkey(publickey)
    crypto_box_seal(ciphertext, message.data, len, binpublickey)
    Base.shred!(message)
    return base64encode(ciphertext)
end

"""
    seal(message::Vector{UInt8}, publickey)

Convenience function which wraps the message in a `SecretBuffer` before calling
`seal(::SecretBuffer, publickey)`. Also wipes and empties the message vector.
"""
seal(message::Vector{UInt8}, publickey) = seal(Base.SecretBuffer!(message), publickey)

"""
    seal(message, publickey)

Fallback that creates a secret by `write(secretbuffer, message)`. Won't cleanup message object.
"""
function seal(message, publickey)
    @warn "You need to cleanup the message yourself"
    secret = Base.SecretBuffer()
    write(secret, message)
    seal(secret, publickey)
end


text:
# ------------------------------------------------------------------------------------------
# # Minimizing functions - how a computer learns
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# In the previous notebooks, we've seen that by changing **parameters** in a function, we
# could find a "best fit" of that function to some data. We use a **loss function** to
# quantify the "goodness" of a set of parameters and we look for those parameters that
# **optimize**, in fact **minimize**, the loss function. If we teach a machine how to
# minimize the loss function on its own, we say that machine is able to **learn** how to
# model our data.
#
# In the last notebook, we threw around terms like **derivative**, **gradient**, and
# **gradient descent** to give you a rough sense of how we minimize a function on a
# computer. In this notebook, we will step through these concepts more carefully, with the
# aim of being able to implement them using Julia.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ## Minimizing a 1D function using calculus
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# Let's draw the 1D loss function $L_1$ from a previous notebook again:
# ------------------------------------------------------------------------------------------

σ(x) = 1 / (1 + exp(-x))
f(x, w) = σ(w * x)

x1 = 2
y1 = 0.8

L1(w) = (y1 - f(x1, w))^2

using Plots; gr()
using Interact

plot(L1, -2, 1.5, xlabel="w", ylabel="L1(w)", leg=false)

# ------------------------------------------------------------------------------------------
# By eye, we can see that the minimum is around $w=0.6$. But how can we get the computer to
# work this out on its own?
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# In the previous notebook, we thought of this function plot as a hill, viewed from the
# side. We could find the minimum by making the hill sticky, and letting a ball roll down
# it. The ball will find and settle in the minimum of the function.  Now let's see how to
# teach a computer to do this.
#
# We need to find the downhill direction along the hill, which is related to its *slope*
# (how steep it is). Calculus provides us with tools to calculate that slope!
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# Namely, the slope of a curve $L_1(w)$ at $w$ is given by its **derivative** $L_1'(w)$;
# geometrically, this is the slope of the **tangent line** to the curve at that point, i.e.
# the straight line which touches the curve at that point.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# Calculus provides us with some rules to calculate an analytical formula for the
# derivative, and we will see later how to apply these rules, albeit indirectly, for machine
# learning.
# To gain understanding, however, we will see here how to get the computer to help us out by
# calculating the derivatives numerically instead!
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ## Approximating derivatives
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# Let's recall that the derivative $L_1'(w)$ of a function is defined as
#
# $$L_1'(w) \simeq \frac{L_1(w+h) - L_1(w)}{h},$$
#
# for a small step size $h$. (Strictly speaking, we must take the limit when $h$ tends to
# $0$ to obtain the exact value of the derivative.)
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 1
#
# Write a function to calculate the derivative of a function at a given point. Note that in
# Julia, we can easily pass functions as arguments to other functions!
# The function should take the function, the point $w$ at which to calculate the derivative,
# and the value of $h$, which should have a default value of 0.001.
#
# *Note*: A function input argument can have a *default value* if we set the input argument
# equal to that default value when we define the function. For example,
#
# ```julia
# f(x, a = 3) = a * x^2
# ```
#
# The function `f` will square the input we give it and multiply by `a`. However,  if we
# choose to call `f(x)` *without* passing it an `a` input, it will assume `a` is `3` and
# return `3*x^2`.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 2
#
# Write an interactive visualization of the tangent line to the graph of $L_1$, so that we
# can visualize the tangent at any point on $L_1$. Include the current value of the
# derivative in the title.
#
# *Hint*: Recall that the straight line through the point $(x_0, y_0)$ with slope $m$ is
# given by
#
# $$ \frac{y - y_0}{x - x_0} = m,$$
#
# so
#
# $$y = y_0 + m*(x - x_0).$$
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 3
#
# What is the value of the derivative (slope of the tangent line) at a minimum? Can this
# happen anywhere else?
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ## Minimization by gradient descent
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# The tangent line at a point is a good approximation of the function near that point.
# Therefore the derivative tells us in which direction, and how fast, a function grows or
# shrinks when we move a small distance away from that point.
#
# As we saw in the previous notebook, we can think of the function as being a hill, and
# having a ball that we want to move down the hill. Gravity will automatically pull the ball
# in the direction *down* the hill; we want to emulate this in the computer!
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 4
#
# When the derivative $L_1'(w)$ is positive, that means that $L_1$ increases from left to
# right at point $w$.
#
# If the derivative $L_1'(w)$ is positive (> 0), in which direction should we move $w$ to
# *decrease* $L_1$?
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 5
#
# If the derivative $L_1'(w)$ is negative (< 0), in which direction should we move $w$ to
# *decrease* $L_1$?
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# We can use this information to tell the computer which way to take a step. This
# constitutes the numerical algorithm called **gradient descent**; it is called this since
# we are descending (moving downwards) along the graph of the function by using information
# about its gradient (slope).
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 6
#
# Implement gradient descent by following this prescription for an **iterative (repetitive)
# algorithm**:
#
# 1. Start at an initial guess $w_0$ for $w$.
#
# 2. At each step, calculate the derivative, $L_1'(w_n)$ at the current value of $w_n$ using
# the function that you created above.
#
# 3. Modify the value of $w$ by a small multiple (for example, $\eta=0.01$) of the value of
# the derivative you just created, via $w_{n+1} = w_n - \eta L_1'(w_n)$.
#
# For this problem, start with $w_0 = -2.0$. Repeat steps 2 and 3 a total of `2000` times.
#
# Package this code into a function called `gradient_descent` that takes two inputs, a
# function and a range for values of $w$, and returns the final value of $w$ and $L_1(w)$.
#
# Using `L1` and `-2:0.01:1.5` as your inputs to `gradient_descent`, for what value of $w$
# is $L_1$ at a minimum?
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 7
#
# Modify your code for gradient descent to return the result once it has found an answer
# within some *tolerance*, rather than taking a set number of steps. The new prescription
# for this algorithm is:
#
# 1. Start at an initial guess $w_0$ for $w$.
#
# 2. At each step, calculate the derivative, $L_1'(w_n)$ at the current value of $w_n$,
# using the function that you created above.
#
# 3. Modify the value of $w$ by a small multiple (for example, $\eta=0.01$) of the value of
# the derivative you just created, via $w_{n+1} = w_n - \eta L_1'(w_n)$.
#
# 4. Check how different $w_{n+1}$ is from $w_n$. If you're satisfied that $L_1(w_{n+1})$ is
# minimized, return $w_{n+1}$ and $L_1(w_{n+1})$. Otherwise, go to step (2) and continue.
#
# Edit `gradient_descent` so that it takes three inputs: a function, a range for values of
# $w$, and a tolerance that tells you how close $w_{n+1}$ must be to $w_n$ before you can
# stop iterating.
#
# Using `L1`, `-2:0.01:1.5`, and `.000001` as your inputs to `gradient_descent`, for what
# value of $w$ is $L_1$ at a minimum?
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 8
#
# Alter the function `gradient_descent` so that it stores the results `(w, L1(w))` at each
# step of the algorithm as an array and returns this array. How many steps does the
# algorithm take for input parameters `(L1, -2:0.01:1.5, .000001)` before terminating? You
# should count your starting $w_0$ as your first step.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 9
#
# Overlay the steps taken in the last exercise with a plot of $L_1(w)$ vs. $w$.
#
# Where does our algorithm take the largest steps? (Where does the ball move fastest down
# the hill?)
#
# A) Between w = -2:-1

# B) Between w = -1:0

# C) Between w = 0:.6
#
# *Hint*: It may be easier to see what's going on if you only plot, for example, every 15th
# step.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# B) The following code generates the desired figure:
#
# ```julia
# wsteps = gradient_descent(C1, -2:0.01:1.5, .000001)
# fig1 = plot(C1, -2, 1.5, xlabel="w", ylabel="C1(w)", leg=false)
# for step in wsteps[1:15:end]
#    scatter!([step[1]], [step[2]])
# end
# display(fig1)
# ```
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ## Functions of 2 variables and their derivatives
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# So far, we have been looking at the minimizing a loss function $L_1(w)$ that depends on a
# single parameter, $w$. Now let's turn to the cost function $L_2(w, b)$ from a previous
# notebook, that is a function of *two* parameters, $w$ and $b$, and try to minimize it.
#
# As we've seen, we get a **surface**, instead of a curve, when we graph $L_2$ as a function
# of both of its parameters.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# **Exercise 10**
#
# Draw a surface plot of $L_2$, given by
#
# $$L_{2}(w, b) = \sum_i(y_i - g(x_i, w, b))^2$$
#
# using the `surface` function from `Plots.jl`. For this plot, use the values of `xs` and
# `ys` from notebook 5:
#
# ```julia
# xs = [2, -3, -1, 1]
# ys = [0.8, 0.3, 0.4, 0.4]
# ```
#
# We can get a nice interactive 3D plot by using the Plotly backend of `Plots.jl` by
# executing the command
#
#     plotly()
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ### Finding the minimum
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# We can just about see, by rotating the graph, that $L_2$ has a single minimum. We want to
# find the values of $w$ and $b$ where this minimum value is reached.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# Following what we did for the function $L_1$ above, we expect that we will need to
# calculate derivatives of $L_2$. Since the function is more complicated, though, the
# derivatives are too!
#
# It turns out that the right concept is that of the
# [**gradient**](https://en.wikipedia.org/wiki/Gradient) of $L_2$, denoted $\nabla L_2(w,
# b)$. This is a **vector** consisting of $2$ numbers if there are $2$ parameters [or, in
# general, $n$ numbers if there are $n$ parameters].
#
# The numbers that form the gradient $\nabla L_2(w, b)$ are called the **partial
# derivatives** of $L_2$ with respect to $w$ and $b$,  written as
#
# $$\frac{\partial L_2}{\partial w} \quad \text{and} \quad \frac{\partial L_2}{\partial
# b}.$$
#
# Although this notation might look complicated, all it means is that we calculate
# derivatives just like we did before, except that we fix the value of the other variable.
# For example, to calculate $\frac{\partial L_2}{\partial w}$, the partial derivative of
# $L_2$ with respect to $w$, we fix the value of $b$ and think of the resulting function as
# a function of the single variable $w$; then we use the formula for derivatives of
# functions of a single variable.
#
# [Note that $\frac{\partial L_2}{\partial w}$ is itself a function of $w$ and $b$; we could
# write $\frac{\partial L_2}{\partial w}(w, b)$.]
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 11
#
# Write functions that will allow you to calculate the partial derivatives of a function $f$
# of two variables.
#
# In particular, declare functions called `partial_w` and `partial_b`. Each should take four
# inputs - a function $f$ of two variables, the first input argument to $f$, the second
# input argument to $f$, and a step size `h` with default value `0.001`. `partial_w` should
# return the partial derivative of $f$ with respect to its first input argument and
# `partial_b` should return the partial derivative of $f$ with respect to its second input
# argument.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 12
#
# Use `partial_b` from the last exercise to find the partial derivative of $L_2$ with
# respect to $w$ at b = 0.3, $\frac{\partial L_2}{\partial w}|_{b = 0.3}$ for `w =
# -2:0.01:1`
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 13
#
# Plot the cross section of the surface of $L_2(w, b)$ at $b = 0.3$. Make this plot
# interactive with `@manipulate` to show that the function `partial_w` gives the slope of
# the tangent to this cross section for any point `w` in the range `-2:0.01:1`.
#
# For what value of $w$ in this range is the slope of the cross section closest to -1?
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ## ***Optional**: Functions with $n$ inputs
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# If a function $f$ takes $n$ input arguments, we can write them as $p_1, \ldots, p_n$,
# where $p_i$ means the "$i$th parameter". In Julia, we can wrap them up into a single
# **vector**. Now we can calculate the partial derivative $\frac{\partial L_2}{\partial
# p_i}$ with respect to the $i$th variable.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 14
#
# For the next exercise, you will need to use the splat command, `...`. You can use this
# command to "open up" a collection and pass all the elements of that collection as inputs
# to a function.
#
# For example, say you have an array, `numbers`,
#
# ```julia
# numbers = [4, 3, 2]
# ```
#
# and you want to use `numbers` to create a $4\times3\times3$ randomly populated array via
# `rand`. `rand(numbers)` will not do what you want. You could index into `numbers` to grab
# the values you want and pass them to `rand`, as in
#
# ```julia
# rand(numbers[1], numbers[2], numbers[3])
# ```
#
# or you could use a splat:
#
# ```julia
# rand(numbers...)
# ```
#
# Use `...` to pass the contents of `inputs`
#
# ```julia
# inputs = [30, 12, "cats"]
# ```
#
# to the function `dreams`
#
# ```julia
# dreams(i, j, perfect_mammal) = "I wish I had $(i + j) $perfect_mammal."
# ```
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 15:
#
# Write a function, `partial`, to calculate the $i$th partial derivative of a function. This
# function should have four inputs
#
# * a function, $f$, for which you want to compute the partial derivative
# * an array, *p*, specifying the values of all input arguments to $f$ at the point where
# you want $\frac{\partial f}{\partial p_i}$ computed
# * the index, $i$, of the variable with respect to which you want to calculate the partial
# derivative of $f$
# * a step size with default value 0.001
#
# Hint: you will need to `copy` and modify `p` within `partial`.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ## Gradient descent in 2 dimensions
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# It turns out that the gradient vector of the function $L_2(w, b)$ gives the direction in
# the plane $(w, b)$ in which the function $L_2$ **increases fastest**.
#
# In other words, if you start at the position $(w_0, b_0)$ and take a small step of length
# $\eta$ to new values $(w', b')$, the value of the function will change to a new value
# $L_2(w', b')$. How do we get to the minimum of $L_2(w, b)$ as fast as possible? We want to
# step in the direction where $L_2$ decreases fastest!
#
# In multivariable calculus courses, it is shown that $L_2$ will *increase* fastest if you
# take a step **in the direction of the gradient $\nabla L_2(w, b)$**! To decrease $L_2$ the
# fastest, we should take a step in the *opposite* direction, $-\nabla L_2(w, b)$.
#
# Let's now generalize the gradient descent algorithm that we wrote previously to work our
# 2-dimensional function.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 16
#
# Extend your 1D implementation of `gradient_descent` so that it will minimize the function
# $L_2$.
#
# Requirements:
#
# * Your new method for `gradient_descent` will take four input arguments: the function $f$
# for which you seek the minimum, the range of values for $f$'s first input argument that
# you will consider, the range of values for $f$'s second input argument that you will
# consider, and a tolerance that will specify the maximum allowable step size, $\sum_i \eta
# \frac{\partial f}{\partial p_i}$
#
# * Use $\eta = .01$. For example, for a function $f(w, b)$, update $w$ such that $w_{n+1} =
# w_n - 0.01 * \frac{\partial f}{\partial w_n}$
#
# * Seed `gradient_descent` with the starting coordinates [-2.0, -2.0], i.e. $w_0 = -2.0$
# and $b_0 = -2.0$.
#
# * Return all steps (their coordinates) taken during gradient descent and the values of the
# loss function at these coordinates.
#
# Once you have done this, execute
#
# ```julia
# gradient_descent(L2, -2:0.02:2, -2:0.02:2, .001)
# ```
#
# How many steps were taken by gradient descent?
#
# Hint: Do not count your starting coordinates `[-2.0, -2.0]` as a step.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 17
#
# Use the `surface` and `scatter!` commands to illustrate the path taken by
# `gradient_descent` from [-2.0, -2.0] to wherever the algorithm terminates.
#
# Where do the scattered points representing the steps of gradient descent appear the most
# dense?
#
# A) Near the starting point at [-2.0, -2.0]

# B) Near the point [-1.8, 0] where $C_2$ appears nearly flat

# C) Near the minimum of $C_2$
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ## What we have learnt
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# To recap, in this notebook, we have seen that the computer can calculate approximate
# derivatives, and that we can use those inside a relatively simple algorithm, gradient
# descent, to minimize functions of 1 and 2 variables!
#
# Now that we've used gradient descent to **learn** the best values of `b` and `w` for a
# cost function, we can imagine how we could do our classification of apples vs. bananas by
# allowing the machine to choose our parameters! To do this, we need to define a loss
# function for apple vs. banana classification, using images that have already been
# categorised as apples or bananas.
# ------------------------------------------------------------------------------------------

max_stars_repo_path: .nbexports/introductory-tutorials/broader-topics-and-ecosystem/intro-to-ml/10. Tools - Minimizing functions - how a computer learns.jl

max_stars_repo_name: grenkoca/JuliaTutorials

max_issues_repo_name: grenkoca/JuliaTutorials

max_forks_repo_name: grenkoca/JuliaTutorials


text:
using Random, Distributions, Plots
include_model("double_integrator")

T = 10
h = 1.0
u_dist = Distributions.MvNormal(zeros(model.m),
	Diagonal(1.0e-1 * ones(model.m)))
u_hist = rand(u_dist, T-1)

x1 = [-1.0; 0.0]
x_hist = [x1]

for t = 1:T-1
	w = zeros(model.d)
	x⁺ = propagate_dynamics(model, x_hist[end], u_hist[t], w, h, t;
			solver = :levenberg_marquardt,
			tol_r = 1.0e-8, tol_d = 1.0e-6)
	push!(x_hist, x⁺)
end

plot(hcat(x_hist...)')
plot(hcat(u_hist...)')

# LQR
A, B = get_dynamics(model)
Q = @SMatrix [1.0 0.0; 0.0 0.0]
R = @SMatrix [1.0]

function objective(x, u)
	T = length(x)
	J = 0.0
	for t = 1:T
		J += x[t]' * Q * x[t]
		t < T & continue
		J += u[t]' * R * u[t]
	end
	return J / T
end

K, P = tvlqr([A for t = 1:T-1], [B for t = 1:T-1],
	[Q for t = 1:T], [R for t = 1:T-1])

K_vec = [vec(K[t]) for t = 1:T-1]

plot(hcat(K_vec...)')
K_inf = copy(K[1])

x1 = [-1.0; 0.0]
x_hist = [x1]
u_hist = []

for t = 1:T-1
	w = zeros(model.d)
	# push!(u_hist, -K[t] * x_hist[end])
	push!(u_hist, -K_inf * x_hist[end])

	x⁺ = propagate_dynamics(model, x_hist[end], u_hist[end], w, h, t;
			solver = :levenberg_marquardt,
			tol_r = 1.0e-8, tol_d = 1.0e-6)

	push!(x_hist, x⁺)
end

plot(hcat(x_hist...)')
plot(hcat(u_hist...)')

J = objective(x_hist, u_hist) # 2.297419

# Gaussian policy
sigmoid(z) = 1.0 / (1.0 + exp(-z))
ds(z) = sigmoid(z) * (1.0 - sigmoid(z))
# plot(range(-10.0, stop = 10.0, length = 100), ds.(range(-10.0, stop = 10.0, length = 100)))
get_μ(θ, x) = θ[1:2]' * x
get_σ(θ, x) = ds(θ[3:4]' * x)

function policy(θ, x, a)
	μ = get_μ(θ, x)
	σ = get_σ(θ, x)

	1.0 / (sqrt(2.0 * π) * σ) * exp(-0.5 * ((a - μ) / (σ))^2.0)
end

function rollout(θ, x1, model, T;
		w = zeros(model.d))

	# initialize trajectories
	x_hist = [x1]
	u_hist = []
	∇logπ_hist = []

	π_dist = []

	for t = 1:T-1

		# policy distribution
		μ = get_μ(θ, x_hist[end])
		σ = get_σ(θ, x_hist[end])

		π_dist = Distributions.Normal(μ, σ)

		# sample control
		push!(u_hist, rand(π_dist, 1)[1])

		# gradient of policy
		logp(z) = log.(policy(z, x_hist[end], u_hist[end]))
		push!(∇logπ_hist, ForwardDiff.gradient(logp, θ))

		# step dynamics
		x⁺ = propagate_dynamics(model, x_hist[end], u_hist[t], w, h, t;
				solver = :levenberg_marquardt,
				tol_r = 1.0e-8, tol_d = 1.0e-6)
		push!(x_hist, x⁺)
	end

	# evalute trajectory cost
	J = objective(x_hist, u_hist)

	return x_hist, u_hist, ∇logπ_hist, J
end

θ = 0.1 * rand(m)
# rollout(θ, x1, model, T;
# 		w = zeros(model.d))

function reinforce(m, x1, model, T;
	max_iter = 1,
	batch_size = 1,
	solver = :sgd,
	J_avg_min = Inf)

	Random.seed!(1)

	# parameters
	θ = 1.0e-5 * rand(m)

	# cost history
	J_hist = [1000.0]
	J_avg = [1000.0]

	if solver == :adam
		# Adam
		α = 0.0001
		β1 = 0.9
		β2 = 0.999
		ϵ = 10.0^(-8.0)
		mom = zeros(m)
		vel = zeros(m)

	else
		# sgd
		α = 1.0e-6
	end

	for i = 1:max_iter
		∇logπ_hist = []
		J = []

		for j = 1:batch_size
			_x_hist, _u_hist, _∇logπ_hist, _J = rollout(θ, copy(x1), model, T)
			push!(∇logπ_hist, _∇logπ_hist)
			push!(J, _J)
		end

		G = (1 / batch_size) .* sum([sum(g) for g in ∇logπ_hist]) .* (mean(J) - J_avg[end])#0.0 * x1' * P[1] * x1)

		if solver == :adam
			mom = β1 .* mom + (1.0 - β1) .* G
			vel = β2 .* vel + (1.0 - β2) .* (G.^2.0)
			m̂ = mom ./ (1.0 - β1^i)
			v̂ = vel ./ (1.0 - β2^i)
			θ .-= α * m̂ ./ (sqrt.(v̂) .+ ϵ)
		else
			θ .-= α * G
		end

		push!(J_hist, mean(J))
		push!(J_avg, J_avg[end] + (mean(J) - J_avg[end]) / length(J_hist))

		if J_avg[end] < J_avg_min
			return θ, J_hist, J_avg
		end

		if i % 1000 == 0
			println("i: $i")
			println("	cost: $(mean(J))")
			println("	cost (run. avg.): $(J_avg[end])")
		end
	end

	return θ, J_hist, J_avg
end

θ, J, J_avg = reinforce(4, x1, model, T,
	max_iter = 1e6,
	batch_size = 10,
	solver = :adam,
	J_avg_min = 0.3)

plot(log.(J[1:100:end]), xlabel = "epoch (100 ep.)", label = "cost")
plot!(log.(J_avg[1:100:end]), xlabel = "epoch (100 ep.)", label = "cost (run. avg.)",
	width = 2.0)

x_hist,  = rollout(θ, x1, model, T)

include(joinpath(pwd(), "models/visualize.jl"))
vis = Visualizer()
render(vis)
visualize!(vis, model, x_hist, Δt = h)


text:
const _SERIALIZED_MODEL_FILENAME = "model.bin"

struct OptimizerAttributes
    name::String
    version::String
    attributes::Any
end

function OptimizerAttributes(model::OperationModel, optimizer::MOI.OptimizerWithAttributes)
    jump_model = get_jump_model(model)
    name = JuMP.solver_name(jump_model)
    # Note that this uses private field access to MOI.OptimizerWithAttributes because there
    # is no public method available.
    # This could break if MOI changes their implementation.
    try
        version = MOI.get(JuMP.backend(jump_model), MOI.SolverVersion())
        return OptimizerAttributes(name, version, optimizer.params)
    catch
        @debug "Solver Version not supported by the solver"
        version = "MOI.SolverVersion not supported"
        return OptimizerAttributes(name, version, optimizer.params)
    end
end

function _get_optimizer_attributes(model::OperationModel)
    return get_optimizer(get_settings(model)).params
end

struct ProblemSerializationWrapper
    template::ProblemTemplate
    sys::Union{Nothing, String}
    settings::Settings
    model_type::DataType
    name::String
    optimizer::OptimizerAttributes
end

function serialize_problem(model::OperationModel; optimizer=nothing)
    # A PowerSystem cannot be serialized in this format because of how it stores
    # time series data. Use its specialized serialization method instead.
    sys_to_file = get_system_to_file(get_settings(model))
    if sys_to_file
        sys = get_system(model)
        sys_filename = joinpath(get_output_dir(model), make_system_filename(sys))
        # Skip serialization if the system is already in the folder
        !ispath(sys_filename) && PSY.to_json(sys, sys_filename)
    else
        sys_filename = nothing
    end
    container = get_optimization_container(model)

    if optimizer === nothing
        optimizer = get_optimizer(get_settings(model))
        @assert optimizer !== nothing "optimizer must be passed if it wasn't saved in Settings"
    end

    obj = ProblemSerializationWrapper(
        model.template,
        sys_filename,
        container.settings_copy,
        typeof(model),
        string(get_name(model)),
        OptimizerAttributes(model, optimizer),
    )
    bin_file_name = joinpath(get_output_dir(model), _SERIALIZED_MODEL_FILENAME)
    Serialization.serialize(bin_file_name, obj)
    @info "Serialized OperationModel to" bin_file_name
end

function deserialize_problem(
    ::Type{T},
    directory::AbstractString;
    kwargs...,
) where {T <: OperationModel}
    filename = joinpath(directory, _SERIALIZED_MODEL_FILENAME)
    if !isfile(filename)
        error("$directory does not contain a serialized model")
    end
    obj = Serialization.deserialize(filename)
    if !(obj isa ProblemSerializationWrapper)
        throw(IS.DataFormatError("deserialized object has incorrect type $(typeof(obj))"))
    end
    sys = get(kwargs, :system, nothing)

    if sys === nothing
        if obj.sys === nothing && !settings[:sys_to_file]
            throw(
                IS.DataFormatError(
                    "Operations Problem System was not serialized and a System has not been specified.",
                ),
            )
        elseif !ispath(obj.sys)
            throw(IS.DataFormatError("PowerSystems.System file $(obj.sys) does not exist"))
        end
        sys = PSY.System(obj.sys)
    end
    settings =
        Settings(sys; restore_from_copy(obj.settings; optimizer=kwargs[:optimizer])...)
    model = obj.model_type(obj.template, sys, settings, kwargs[:jump_model]; name=obj.name)
    jump_model = get_jump_model(model)
    if obj.optimizer.name == JuMP.solver_name(jump_model)
        orig_attrs = obj.optimizer.attributes
        new_attrs = _get_optimizer_attributes(model)
        if length(orig_attrs) != length(new_attrs)
            @warn "Different optimizer attributes are set. Original: $orig_attrs New: $new_attrs"
        else
            for attrs in (orig_attrs, new_attrs)
                sort!(attrs, by=x -> x.first.name)
            end
            for i in 1:length(orig_attrs)
                name = orig_attrs[i].first.name
                orig = orig_attrs[i].second
                new = new_attrs[i].second
                if new != orig
                    @warn "Original solver used $name = $orig. New solver uses $new."
                end
            end
        end
    else
        @warn "Original solver was $(obj.optimizer.name), new solver is $(JuMP.solver_name(jump_model))"
    end

    return model
end