module ReducedNetwork
using DifferentialEquations, LightGraphs, SimpleWeightedGraphs, Statistics,
    Plots, GraphPlot, Compose, Cairo, Fontconfig, Printf # for plotting
using SpikingNeuralNets.GraphUtils: mesh_graph

# Effective transfer function
x(S1, S2, Inet, I, Inoise, t, p) = p.JNex*S1 + p.JNin*S2 + Inet + p.Iext + I*(t >= 0.0) + Inoise
function H(x1, x2, p)
    v = p.a*x1 - p.fA(x2) - p.b
    return v / (1 - exp(p.d * v))
end

"""
    Inetwork(S::AbstractVector{T}, g::AbstractGraph)::Vector{T} where {T <: Number}

For each vertex in `g` calculate a weighted average of S from all inneighbors of `v`.
"""
function Inetwork(S::AbstractVector{T}, g::AbstractGraph)::Vector{T} where {T <: Number}
    I = zeros(T, size(S))
    @inbounds for v in vertices(g)
        in = inneighbors(g, v)
        if !iszero(length(in))
            I[v] = sum(weights(g)[in,v] .* S[in])
        end
    end
    return I
end

"""
    ReducedModel(du, u, p, t, W)

Calculate the change mean gating variables `du` given the current gating
variables `u`, named parameters `p`, time `t`, and noise currents `W`.
"""
function ReducedModel(du, u, p, t, W)
    Sa, Sb = eachcol(u)
    Inoisea, Inoiseb = eachcol(W)

    Ineta = Inetwork(p.Jex*Sa + p.Jin*Sb, p.g)          # network inputs
    Inetb = Inetwork(p.Jex*Sb + p.Jin*Sa, p.g)
    xa = x.(Sa, Sb, Ineta, p.Ia, Inoisea, t, Ref(p))    # intermediate variables
    xb = x.(Sb, Sa, Inetb, p.Ib, Inoiseb, t, Ref(p))
    du[:,1] .= .-(Sa/p.τs) .+ (1 .- Sa) .* p.γ .* H.(xa, xb, Ref(p))
    du[:,2] .= .-(Sb/p.τs) .+ (1 .- Sb) .* p.γ .* H.(xb, xa, Ref(p))
end

"""
    run(g::AbstractGraph; c=fill(0, nv(g)), μ0=fill(30.0, nv(g)), S0=zeros(nv(g), 2),
        Inoise0=zeros(nv(g), 2), tstart=-0.5, tstop=1.0, reltol=1e-15,
        abstol=1e-15, dt=0.1e-3, Jex=1.0, Jin=0.0, Jext=0.2243e-3, JNex=0.1561,
        JNin=-0.0264, JAex=9.9026e-4, JAin=-6.5177e-5, τs=100e-3, τAMPA=2e-3,
        γ=0.641, Iext=0.35, σnoise=0.007, display=true, color1=:red, color2=:blue,
        kwargs...)

Run a reduced 2-variable model of neural populations during a random dot motion task.
Return a tuple (g, t, Sa, Sb, Inoisea, Inoiseb, Ineta, Inetb, xa, xb, ra, rb).

`g`: a graph of neural populations, where each vertex represents a population of
     neurons, and each edge represents a connection between neural populations.
     Weighted graphs (e.g. `SimpleWeightedDiGraph`) a.
`c`: the coherence of external input to each node.
`μ0`: the firing rate (in Hz) of external input to each node,
`S0`: the initial average NMDA gating variables for each node,
`Inoise0`: are the initial noise currents for each node,
`tstart`: the starting time of the simulation (negative times having no stimulus)
`tstop`: is the stopping time of the simulation
`dt`: is the timestep of the simulation.
`Jex`: the weight of recurrent input from excitatory sub-populations
`Jin`: the weight of recurrent input from inhibitory sub-populations
`Jext`: the weight of external inputs
`JNex`: the excitatory (> 0) NMDA weight
`JNin`: the inhibitory (< 0) NMDA weight
`JAex`: the excitatory (> 0) AMPA weight
`JAin`: the excitatory (> 0) AMPA weight
`τs`: the time constant of NMDA channels
`τAMPA`: the time constant of AMPA channels
`γ`: a factor applied to firing rate to update gating variables
`Iext`: baseline external input current (nV)
`σnoise`: standard deviation of noise current (nV).
`display`: if true, draw a plot of gating variables and firing rates over time
`color1`: the line color of the correct direction in the output plot
`color2`: the line color of the incorrect direction in the output plot
Additional keyword arguments are passed to DifferentialEquations.solve.
"""
function run(g::AbstractGraph; c=fill(0, nv(g)), μ0=fill(30.0, nv(g)), S0=zeros(nv(g), 2),
             Inoise0=zeros(nv(g), 2), tstart=-0.5, tstop=1.0, reltol=1e-15,
             abstol=1e-15, dt=0.1e-3, Jex=1.0, Jin=0.0, Jext=0.2243e-3, JNex=0.1561,
             JNin=-0.0264, JAex=9.9026e-4, JAin=-6.5177e-5, τs=100e-3, τAMPA=2e-3,
             γ=0.641, Iext=0.35, σnoise=0.007, display=true, color1=:red, color2=:blue,
             kwargs...)
    # set up parameters
    Ia = @. Jext*μ0*(1 + c/100.0)
    Ib = @. Jext*μ0*(1 - c/100.0)
    a = 239400JAex + 270
    b = 97000JAex + 108
    d = 30JAex - 0.154
    fA(x) = JAin*(276x - 106) * (x >= 0.4)
    params = (g=g, a=a, b=b, c=c, d=d, fA=fA, μ0=μ0, Ia=Ia, Ib=Ib, Jex=Jex, Jin=Jin,
              Jext=Jext, JNex=JNex, JNin=JNin, JAex=JAex, JAin=JAin,
              τs=τs, τAMPA=τAMPA, γ=γ, Iext=Iext)

    # noise process: τAMPA dI(t)/dt = -I(t) + W(t)*σnoise*√τAMPA
    W = OrnsteinUhlenbeckProcess(τAMPA, 0.0, σnoise*√τAMPA, tstart, Inoise0)
    prob = RODEProblem(ReducedModel, S0, (tstart, tstop), params, noise=W)
    sol = solve(prob, reltol=reltol, abstol=abstol, dt=dt, kwargs...)

    if display
        # Reconstruct the variables
        Sa, Sb = sol[:,1,:], sol[:,2,:]
        Inoisea, Inoiseb = sol.W[:,1,1:length(sol.t)], sol.W[:,2,1:length(sol.t)]
        Ineta = mapslices((s)->Inetwork(s, g), Jex*Sa + Jin*Sb; dims=1)
        Inetb = mapslices((s)->Inetwork(s, g), Jex*Sb + Jin*Sa; dims=1)

        # reconstruct the firing rates
        xa, xb = similar(Sa), similar(Sb)
        @views for v in vertices(g)
            xa[v,:] .= x.(Sa[v,:], Sb[v,:], Ineta[v,:], Ia[v], Inoisea[v,:], sol.t, Ref(params))
            xb[v,:] .= x.(Sb[v,:], Sa[v,:], Inetb[v,:], Ib[v], Inoiseb[v,:], sol.t, Ref(params))
        end
        ra = H.(xa, xb, Ref(params))
        rb = H.(xb, xa, Ref(params))

        Base.display(plot(plot(sol.t, [Sa' Sb'], title="Gating Variables"),
                     plot(sol.t, [ra' rb'], title="Firing Rates (Hz)"),
                     leg=false, linecolor=[fill(:red, 1, nv(g)) fill(:blue, 1, nv(g))], layout=(2, 1)))
        return (g, sol.t, Sa, Sb, Inoisea, Inoiseb, Ineta, Inetb, xa, xb, ra, rb)
    else
        return g, sol
    end
end

"""
    run(g::AbstractGraph, inputs; c=0.0, μ0=30.0, kwargs...)

Run a reduced 2-variable model of neural populations during a random dot motion task,
where `inputs` contains all of the vertices of `g` receiving external input, and
`c` and `μ0` are the coherence and strength of external input. Return a tuple
(g, t, Sa, Sb, Inoisea, Inoiseb, Ineta, Inetb, xa, xb, ra, rb).
"""
function run(g, inputs; c=0.0, μ0=30.0, kwargs...)
    C = fill(c, nv(g))
    μ = zeros(nv(g))
    μ[inputs] .= μ0
    return run(g; c=C, μ0=μ, kwargs...)
end

"""
    run(dims::Tuple{Vararg{<:Integer, N}}; inputs=ifelse(N==1, 1, 1:first(dims)), weight=0.015, c=0, μ0=30.0, kwargs...) where N

Run a reduced 2-variable model of a mesh_graph of neural populations during a
random dot motion task, where `dims` is the size of the mesh graph, `inputs`
contains all of the vertices of `g` receiving external input, and `c` and `μ0`
are the coherence strength of external input. Return a tuple
(g, t, Sa, Sb, Inoisea, Inoiseb, Ineta, Inetb, xa, xb, ra, rb).
"""
function run(dims::Tuple{Vararg{<:Integer, N}}; inputs=ifelse(N==1, 1, 1:first(dims)), weight=0.015, c=0, μ0=30.0, kwargs...) where N
    return run(mesh_graph(dims; self=false, weight=weight), inputs; kwargs...)
end

"""
    run(dims::Vararg{<:Integer, N}; kwargs...) where N

Run a reduced 2-variable model of a mesh_graph of neural populations during a
random dot motion task, where `dims` is the size of the mesh graph, `inputs`
contains all of the vertices of `g` receiving external input, and `c` and `μ0`
are the coherence strength of external input. Return a tuple
(g, t, Sa, Sb, Inoisea, Inoiseb, Ineta, Inetb, xa, xb, ra, rb).
"""
function run(dims::Vararg{<:Integer, N}; kwargs...) where N
    return run(dims; kwargs...)
end

"""
    RBplot(g::AbstractGraph, Sa::Vector{<:Real}, Sb::Vector{<:Real};
           color1=colorant"red", color2=colorant"yellow",
           locs_x=nothing, locs_y=nothing, kwargs...)

Draw a plot of the graph g, where each vertex `v` is colored by the values
`Sa[v]` and `Sb[v]` on a scale from `color1` to `color2`.
"""
function RBplot(g::AbstractGraph, Sa::Vector{<:Real}, Sb::Vector{<:Real};
                color1=colorant"red", color2=colorant"yellow",
                locs_x=nothing, locs_y=nothing, kwargs...)
    nodefillc = @. weighted_color_mean((Sa-Sb+1)/2, color1, color2)
    if locs_x != nothing && locs_y != nothing
        return gplot(DiGraph(weights(g)), locs_x, locs_y; nodefillc=nodefillc, kwargs...)
    end
    return gplot(DiGraph(weights(g)); nodefillc=nodefillc, kwargs...)
end

"""
    RBgif(g::AbstractGraph, Sa::Matrix{<:Real}, Sb::Matrix{<:Real};
          layout::Function=spring_layout,fps:Integer=100, skip::Integer=1,
          width::Integer=800, height::Integer=600, kwargs...)

Draw a gif of the graph g, where each vertex `v` is colored over time by the
values `Sa[v, t]` and `Sb[v, t]` on a scale from `color1` to `color2`.
"""
function RBgif(g::AbstractGraph, Sa::Matrix{<:Real}, Sb::Matrix{<:Real};
               layout::Function=spring_layout,
               filename=nothing, fps::Integer=100, skip::Integer=1,
               width::Integer=800, height::Integer=600, kwargs...)
    # use the Animation object manually since frame() doesn't work on Compose objects
    anim = Animation()
    locs_x, locs_y = layout(g)
    for t in 1:skip:size(Sa, 2)
        fname = @sprintf("%06d.png", length(anim.frames)+1)
        draw(PNG(joinpath(anim.dir, fname), width, height),
             RBplot(g, Sa[:,t], Sb[:,t]; locs_x=locs_x, locs_y=locs_y, kwargs...))
        push!(anim.frames, fname)
    end
    if isnothing(filename)
        return gif(anim; fps=fps)
    end
    return gif(anim, filename; fps=fps)
end

function phaseplot(g::AbstractGraph, inputs; iterations=1, threaded=true,
                   Jex=-1.0:0.01:1.0, Jin=-1.0:0.01:1.0, kwargs...)
    diff = Array{Float64}(undef, length(Jex), length(Jin), nv(g))

    function difference(Jex, Jin)
        g, sol = run(g, inputs; Jex=Jex, Jin=Jin, display=false, kwargs...)
        return sol[:,1,end] .- sol[:,2,end]
    end

    if threaded
        @inbounds @Threads.threads for i in eachindex(Jex)
            for j in eachindex(Jin)
                D = [difference(Jex[i], Jin[j]) for itr in 1:iterations]
                diff[i, j, :] .= mean(D)
            end
        end
    else
        @inbounds for i in eachindex(Jex), j in eachindex(Jin)
            diff[i, j, :] .= mean(difference(Jex[i], Jin[j]) for itr in 1:iterations)
        end
    end
    return diff
end

end


using Sampling

function bench_categorical_sampler(s, pv::Vector{Float64}, ns::Int)
    # warm up
    for i = 1:10; rand(s); end
    # measure
    et = @elapsed for i = 1:ns; rand(s); end
    mps = (ns * 1.0e-6) / et
    @printf("%-13s with K = %4d =>  %9.4f MPS\n", string(typeof(s)), length(pv), mps)
end

for K in 2.^(1:12)
    pv = fill(1/K, K)
    bench_categorical_sampler(AliasTable(pv), pv, 10^6)
end


### A Pluto.jl notebook ###
# v0.17.3

using Markdown
using InteractiveUtils

# ╔═╡ 8d375e12-5f3c-11ec-0551-0d5da3263f13
using DelimitedFiles

# ╔═╡ da0f707b-5c2a-4ab5-89c2-c48c04ef9404
d = readdlm("inputs/day_02.txt")

# ╔═╡ da7e5513-c0be-4bc6-9126-80ecfefd73db
directions = Dict(
	"forward" => [1, 0],
	"down" => [0, 1],
	"up" => [0, -1]
)

# ╔═╡ 127bd9b0-6b9b-4203-a325-3ddfae44dec2
instructions = map(command -> directions[command], d[:, 1])

# ╔═╡ 55369b8a-fd99-462f-a05a-6775a11652c3
steps = d[:, 2]

# ╔═╡ 4d5841f6-4f6f-44c3-af4a-084ce4e59864
*(sum(instructions .* steps)...)

# ╔═╡ 265dca33-2dd8-4adc-8612-4e905036ef5f
md"""
# Part II

In addition to horizontal position and depth, you'll also need to track a third value, aim, which also starts at 0. The commands also mean something entirely different than you first thought:

- down X increases your aim by X units.
- up X decreases your aim by X units.
- forward X does two things: It increases your horizontal position by X units. It increases your depth by your aim multiplied by X.

Again note that since you're on a submarine, down and up do the opposite of what you might expect: "down" means aiming in the positive direction.

"""

# ╔═╡ 8dca4413-94d8-4778-bda7-2143da76160b
function move_submarine(instructions, steps)
	aim = 0
	location = [0, 0]

	for (inst, step) in zip(instructions, steps)
		if inst == "up"
			aim -= step
		elseif inst == "down"
			aim += step
		else
			move = [step, step * aim]
			location += move
		end
	end
	return location
end

# ╔═╡ b9c2d28c-dc6c-49f9-9a6e-7ef2ba2dfec9
*(move_submarine(d[:, 1], steps)...)

# ╔═╡ 1b2e827f-31fc-4ba7-9d64-86ec46ce4ee8

# ╔═╡ 00000000-0000-0000-0000-000000000001
PLUTO_PROJECT_TOML_CONTENTS = """
[deps]
DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab"
"""

# ╔═╡ 00000000-0000-0000-0000-000000000002
PLUTO_MANIFEST_TOML_CONTENTS = """
# This file is machine-generated - editing it directly is not advised

julia_version = "1.7.0"
manifest_format = "2.0"

[[deps.DelimitedFiles]]
deps = ["Mmap"]
uuid = "8bb1440f-4735-579b-a4ab-409b98df4dab"

[[deps.Mmap]]
uuid = "a63ad114-7e13-5084-954f-fe012c677804"
"""

# ╔═╡ Cell order:
# ╠═8d375e12-5f3c-11ec-0551-0d5da3263f13
# ╠═da0f707b-5c2a-4ab5-89c2-c48c04ef9404
# ╠═da7e5513-c0be-4bc6-9126-80ecfefd73db
# ╠═127bd9b0-6b9b-4203-a325-3ddfae44dec2
# ╠═55369b8a-fd99-462f-a05a-6775a11652c3
# ╠═4d5841f6-4f6f-44c3-af4a-084ce4e59864
# ╠═265dca33-2dd8-4adc-8612-4e905036ef5f
# ╠═8dca4413-94d8-4778-bda7-2143da76160b
# ╠═b9c2d28c-dc6c-49f9-9a6e-7ef2ba2dfec9
# ╠═1b2e827f-31fc-4ba7-9d64-86ec46ce4ee8
# ╟─00000000-0000-0000-0000-000000000001
# ╟─00000000-0000-0000-0000-000000000002


struct BradleyTerry
	rating::Dict
	playergames::Dict
	default_rating::Float64
end

function BT(; rating = Dict{Int64, Float64}(), playergames = Dict(), default_rating = 1.0)
    #Set up a new Elo struct, either using key words or defaults if not provided
    return BradleyTerry(rating, playergames, default_rating)
end

#Add games to a BradleyTerry object
function add_games!(bt::BradleyTerry, games::DataFrame; dummy_games::Bool = true)
    #Duplicate whole rating dataframe with reversed results
    games = dupe_for_rating(games)

    #Split up by player (P1) and create Dictionary of games belonging to each player (won and lost)
    for player in groupby(games, 1)

		#Add all games for that player
		p = Tuple.(eachrow(player))
		i = player[1,1]
        bt.playergames[i] = vcat(get(bt.playergames, i, []), p)

		sort!(bt.playergames[i], by = x -> x[4])

		#See whether we've got our dummy games in right place
		if dummy_games
			bt.playergames[i] = add_dummy_games(i, bt.playergames[i])
		end
    end
end

function add_dummy_games(i, games)
	if games[1][2] == 0
		#if dummy games exist as the earliest games, do nothing
		return games
	else
		#if dummy games are not the first games, then add them
		firstday = games[1][4]
		return vcat([(i, 0, 1.0, firstday, 1), (i, 0, 0.0, firstday, 1)], games)
		#if we have dummy games but they are not the first ones then Delete existing dummy games and create new ones
		#ADD THIS. DOES NOT MATTER IF ALWAYS ADDING FUTURE GAMES BUT WILL MATTER LATER
	end
end

#Iterate the ratings for all players
function iterate!(bt::BradleyTerry, iterations::Int64; exclude_non_ford::Bool = false)

	#Check and get rid of any players not Ford-connected
	if exclude_non_ford
		non_ford = check_ford(bt)
		connected_players = filter(x -> ∉(first(x), non_ford) , bt.playergames)
		#println(length(non_ford), " players excluded due to not being connected. ", length(connected_players), " players being rated.")
	else
		connected_players = bt.playergames
	end

	#println("Initial loglikelihood ", loglikelihood(bt, connected_players))

	#Iterate over players selected
	for i = 1:iterations
		for (player, _) in connected_players
			bt.rating[player] = update_rating(bt, player)
		end
		#println(loglikelihood(bt, connected_players))
	end
end

#Iterate the rating of selection of players
function iterate!(bt::BradleyTerry, iterations::Int64, players)
	for i = 1:iterations
		for player in players
			bt.rating[player] = update_rating(bt, player)
		end
		println(loglikelihood(bt, players))
	end
end

#Update function allowing for period in which there may be duplicate games
function update_rating(bt::BradleyTerry, player::Int64)
	p = sum(w for (_, _, w, _) in bt.playergames[player])
	opponents = unique(j for (_, j, _, _) in bt.playergames[player])
	q = sum(zermelement(bt, player, opponent) for opponent in opponents)
	return p / q
end

function loglikelihood(bt::BradleyTerry; players = bt.playergames)
	tally = 0.0
	for (player, games) in players
		wonplayergames = filter(y -> y[3] == 1.0, games)
		tally += sum(llelement(bt, player, opponent) for (_, opponent, _, _) in wonplayergames)
		tally
	end
	return tally
end

function llelement(bt::BradleyTerry, i::Int64, j::Int64)
	j == 0 && return 1.0 * (log(get(bt.rating, i, bt.default_rating)) - log(get(bt.rating, j, bt.default_rating) + 1))
	player_won_games = length(filter(y -> y[2] == j && y[3] == 1.0, get(bt.playergames, i, [])))
	return player_won_games * (log(get(bt.rating, i, bt.default_rating)) - log(get(bt.rating, i, bt.default_rating) + get(bt.rating, j, bt.default_rating)))
end

#Calculate the element of the Zermelo algorithm given in Hunter paper
function zermelement(bt::BradleyTerry, i::Int64, j::Int64)
	i == j && return 0.0 #Don't include self
	#j == 0 && return 2 / (bt.rating[i] + 1.0) #Give quick answer for dummy games where all players won and lost to player 0 with rating 1.
	Nij = length(filter(y -> y[2] == j, get(bt.playergames, i, [])))  #Number of times players have played one another
	Nij == 0 && return 0.0 #If no matches, immediately return zero
	return Nij / (get(bt.rating, i, bt.default_rating) + get(bt.rating, j, bt.default_rating))
end

function predict(bt::BradleyTerry, P1::Int64, P2::Int64)
	#Generate standard Bradley Terry probability for P1 winning a game
	return get(bt.rating, P1, bt.default_rating) / (get(bt.rating, P1, bt.default_rating) + get(bt.rating, P2, bt.default_rating))
end

function check_ford(bt::BradleyTerry)
	#Check for players who have either never won or never lost - the Ford criterion
	unconnected = []
	for (player, games) in bt.playergames
		if length(filter(y -> y[3] == 0.0, games)) == 0
			unconnected = vcat(unconnected, player)
		end
		if length(filter(y -> y[3] == 1.0, games)) == 0
			unconnected = vcat(unconnected, player)
		end
	end
	return unconnected
end

function bt_predictable_games(bt, games)
    unratedplayers = check_ford(bt)
    println(length((unratedplayers)))
    Player1history = map(x -> x ∉ unratedplayers, games.P1)
    Player2history = map(x -> x ∉ unratedplayers, games.P2)
    return Player1history .* Player2history
end


"""
    simulate_random_snparray(s::String, n::Integer, p::Integer;
        [mafs::Vector{Float64}], [min_ma::Integer])

Creates a random SnpArray in the current directory without missing value,
where each SNP has ⫺5 (default) minor alleles.

Note: if supplied minor allele frequency is extremely small, it could take a
long time for the simulation to generate samples where at least `min_ma`
(defaults to 5) are present.

# Arguments:
- `s`: name of SnpArray that will be created in the current directory. To not
    create file, use `undef`.
- `n`: number of samples
- `p`: number of SNPs

# Optional Arguments:
- `mafs`: vector of desired minor allele freuqencies (uniform(0,0.5) by default)
- `min_ma`: the minimum number of minor alleles that must be present for each
    SNP (defaults to 5)
"""
function simulate_random_snparray(s::Union{String, UndefInitializer}, n::Int64,
    p::Int64; mafs::Vector{Float64}=zeros(Float64, p), min_ma::Int = 5)

    @assert all(0.0 .<= mafs .<= 0.5) "Minor allele frequencies not in (0, 0.5)"

    if mafs != zeros(Float64, p)
        return _random_snparray(s, n, p, mafs, min_ma=min_ma)
    end

    #first simulate a random {0, 1, 2} matrix with each SNP drawn from Binomial(2, r[i])
    A1 = BitArray(undef, n, p)
    A2 = BitArray(undef, n, p)
    for j in 1:p
        minor_alleles = 0
        maf = 0
        while minor_alleles <= min_ma
            maf = 0.5rand()
            for i in 1:n
                A1[i, j] = rand(Bernoulli(maf))
                A2[i, j] = rand(Bernoulli(maf))
            end
            minor_alleles = sum(view(A1, :, j)) + sum(view(A2, :, j))
        end
        mafs[j] = maf
    end

    #fill the SnpArray with the corresponding x_tmp entry
    return _make_snparray(s, A1, A2)
end

"""
This function requires a vector of minor alleles frequencies to perform simulation.
It will create a random SnpArray in the current directory without missing value,
where each SNP has at least 5 minor alleles.
"""
function _random_snparray(s::Union{String, UndefInitializer}, n::Int64,
        p::Int64, mafs::Vector{Float64}; min_ma::Int = 5)
    all(0.0 .<= mafs .<= 0.5) || throw(ArgumentError("vector of minor allele frequencies must be in (0, 0.5)"))
    any(mafs .<= 0.0005) && @warn("Provided minor allele frequencies contain entries smaller than 0.0005, simulation may take long if sample size is small and min_ma = $min_ma is large")

    #first simulate a random {0, 1, 2} matrix with each SNP drawn from Binomial(2, r[i])
    A1 = BitArray(undef, n, p)
    A2 = BitArray(undef, n, p)
    for j in 1:p
        minor_alleles = 0
        maf = mafs[j]
        while minor_alleles <= min_ma
            for i in 1:n
                A1[i, j] = rand(Bernoulli(maf))
                A2[i, j] = rand(Bernoulli(maf))
            end
            minor_alleles = sum(view(A1, :, j)) + sum(view(A2, :, j))
        end
    end

    #fill the SnpArray with the corresponding x_tmp entry
    return _make_snparray(s, A1, A2)
end

"""
Make a SnpArray from 2 BitArrays.
"""
function _make_snparray(s::Union{String, UndefInitializer}, A1::BitArray, A2::BitArray)
    n, p = size(A1)
    x = SnpArray(s, n, p)
    for i in 1:(n*p)
        c = A1[i] + A2[i]
        if c == 0
            x[i] = 0x00
        elseif c == 1
            x[i] = 0x02
        elseif c == 2
            x[i] = 0x03
        else
            throw(MissingException("matrix shouldn't have missing values!"))
        end
    end
    return x
end

"""
    simulate_correlated_snparray(s, n, p; block_length, hap, prob)

Simulates a SnpArray with correlation. SNPs are divided into blocks where each
adjacent SNP is the same with probability prob. There are no correlation between blocks.

# Arguments:
- `n`: number of samples
- `p`: number of SNPs
- `s`: name of SnpArray that will be created (memory mapped) in the current directory. To not memory map, use `undef`.

# Optional arguments:
- `block_length`: length of each LD block
- `hap`: number of haplotypes to simulate for each block
- `prob`: with probability `prob` an adjacent SNP would be the same.
"""
function simulate_correlated_snparray( s::Union{String, UndefInitializer},
    n::Int64, p::Int64; block_length::Int64=20, hap::Int=20, prob::Float64=0.75)

    @assert mod(p, block_length) == 0 "block_length ($block_length) is not divible by p ($p)"
    @assert 0 < prob < 1 "transition probably should be between 0 and 1, got $prob"

    x = SnpArray(s, n, p)
    haplotypes = zeros(hap, block_length)
    snps = zeros(block_length)
    blocks = Int(p / block_length)

    @inbounds for b in 1:blocks

        #create pool of haplotypes for each block
        _sample_haptotypes!(haplotypes, prob)

        for i in 1:n
            #sample 2 haplotypes with replacement from the pool of haplotypes
            row1 = rand(1:hap)
            row2 = rand(1:hap)
            for j in 1:block_length
                snps[j] = haplotypes[row1, j] + haplotypes[row2, j]
            end

            #copy haplotypes into x
            _copy_blocks!(x, i, snps, b, block_length)
        end
    end

    return x
end

function _sample_haptotypes!(haplotypes::Matrix, prob::Float64)
    n, p = size(haplotypes)
    fill!(haplotypes, 0)

    @inbounds for i in 1:n
        cur_row_sum = 0
        while cur_row_sum == 0
            curr = rand(0:1)
            haplotypes[i, 1] = curr
            cur_row_sum += curr
            for j in 2:p
                stay = rand(Bernoulli(prob)) #stay = 1 means retain the current value
                curr = (stay == 1 ? curr : 1 - curr)
                haplotypes[i, j] = curr
                cur_row_sum += curr
            end
        end
    end
end

function _copy_blocks!(x::SnpArray, row, snps, cur_block, block_length)
    #copy sampled snps into SnpArray
    @inbounds for k in 1:length(snps)
        c = snps[k]
        col = (cur_block - 1) * block_length + k
        if c == 0
            x[row, col] = 0x00
        elseif c == 1
            x[row, col] = 0x02
        elseif c == 2
            x[row, col] = 0x03
        else
            throw(ArgumentError("SNP values should be 0, 1, or 2 but was $c"))
        end
    end
end

"""
    simulate_random_response(x, k, d, l; kwargs...)

This function simulates a random response (trait) vector `y`. When the
distribution `d` is from Poisson, Gamma, or Negative Binomial, we simulate
`β ∼ N(0, 0.3)` to roughly ensure the mean of response `y` doesn't become too
large. For other distributions, we choose `β ∼ N(0, 1)`.

# Arguments
- `x`: Design matrix
- `k`: the true number of predictors.
- `d`: The distribution of the simulated trait (note `typeof(d) = UnionAll` but `typeof(d())` is an actual distribution: e.g. Normal)
- `l`: The link function. Input `canonicallink(d())` if you want to use the canonical link of `d`.

# Optional arguments
- `r`: The number of success until stopping in negative binomial regression, defaults to 10
- `α`: Shape parameter of the gamma distribution, defaults to 1
- `Zu`: Effect of non-genetic covariates. `Zu` should have dimension `n × 1`.
"""
function simulate_random_response(x::AbstractMatrix, k::Int,
    d::UnionAll, l::Link; r = 10, α = 1, Zu::AbstractVector = zeros(size(x, 1)))
    n, p = size(x)
    if (typeof(d) <: NegativeBinomial) || (typeof(d) <: Gamma)
        l == LogLink() || throw(ArgumentError("Distribution $d must use LogLink!"))
    end

    #simulate a random model β from a normal distribution
    true_b = zeros(p)
    if d == Poisson || d == Gamma || d == NegativeBinomial
        true_b[1:k] = rand(Normal(0, 0.3), k)
    else
        true_b[1:k] = randn(k)
    end
    shuffle!(true_b)
    correct_position = findall(x -> x != 0, true_b)

    #simulate phenotypes (e.g. vector y)
    if d == Normal || d == Poisson || d == Bernoulli
        prob = linkinv.(l, x * true_b + Zu)
        clamp!(prob, -20, 20)
        y = [rand(d(i)) for i in prob]
    elseif d == NegativeBinomial
        μ = linkinv.(l, x * true_b + Zu)
        clamp!(μ, -20, 20)
        prob = 1 ./ (1 .+ μ ./ r)
        y = [rand(d(r, i)) for i in prob] #number of failtures before r success occurs
    elseif d == Gamma
        μ = linkinv.(l, x * true_b + Zu)
        β = 1 ./ μ # here β is the rate parameter for gamma distribution
        y = [rand(d(α, i)) for i in β] # α is the shape parameter for gamma
    end
    y = Float64.(y)

    return y, true_b, correct_position
end
"""
    simulate_random_response(x, k, traits)

Simulates a response matrix `Y` where each row is an independent multivariate
Gaussian with length `trait`. There are `k` non-zero `β` over all traits. Each
trait shares `overlap` causal SNPs. The covariance matrix `Σ` is positive definite
and symmetric.

# Arguments
- `x`: Design matrix of dimension `n × p`. Each row is a sample.
- `k`: the total true number of causal SNPs (predictors)
- `traits`: Number of traits

# Optional arguments
- `Zu`: Effect of non-genetic covariates. `Zu` should have dimension `n × traits`.
- `overlap`: Number of causal SNPs shared by all traits. Shared SNPs does not have the same effect size.

# Outputs
- `Y`: Response matrix where each row is sampled from a multivariate normal with mean `μ[i] = X[i, :] * true_b` and variance `Σ`
- `Σ`: the symmetric, positive definite covariance matrix used
- `true_b`: A sparse matrix containing true beta values.
- `correct_position`: Non-zero indices of `true_b`
"""
function simulate_random_response(x::AbstractMatrix, k::Int, traits::Int;
    Zu::AbstractMatrix = zeros(size(x, 1), traits), overlap::Int = 0
    )
    n, p = size(x)
    traits * overlap ≤ k || throw(ArgumentError("traits * overlap cannot exceed k!"))

    #simulate a random model β. Each trait can have different number of causal SNPs
    true_b = zeros(p, traits)
    if overlap == 0
        causal_snps = sample(1:(traits * p), k, replace=false)
        true_b[causal_snps] = randn(k)
    else
        shared_snps = sample(1:p, overlap, replace=false)
        weight_vector = aweights(1 / (traits * (p - overlap)) * ones(traits * p))
        for i in 1:traits
            weight_vector[i*shared_snps] .= 0.0 # avoid sampling from shared snps
        end
        @assert sum(weight_vector) ≈ 1.0
        # simulate β for shared predictors
        for i in 1:traits
            true_b[shared_snps, i] = randn(overlap)
        end
        # simulate β for none shared predictors
        nonshared_snps = sample(1:(traits * p), weight_vector, k - traits * overlap, replace=false)
        true_b[nonshared_snps] = randn(k - traits * overlap)
    end
    correct_position = findall(x -> x != 0, true_b)

    # simulate random covariance matrix
    Σ = random_covariance_matrix(traits)

    # simulate multivariate normal phenotype for each sample
    μ = x * true_b + Zu

    # simulate response
    Y = zeros(n, traits)
    for i in 1:n
        μi = @view(μ[i, :])
        Y[i, :] = rand(MvNormal(μi, Σ))
    end

    return Y, Σ, true_b, correct_position
end

"""
    random_covariance_matrix(n::Int, [κ=10])

Generates a `n × n` positive definite, symmetric matrix with condition number less than `κ`
https://discourse.julialang.org/t/generate-a-positive-definite-matrix/48582

Both issymmetric(random_covariance_matrix(n)) and isposdef(random_covariance_matrix(n))
will always return true, and cond(random_covariance_matrix(n)) ≤ κ
"""
function random_covariance_matrix(n::Int, κ=10)
    d = Uniform(1, sqrt(κ))
    Q, _ = qr(randn(n, n))
    σ = rand(d, n)
    D = Diagonal(σ)
    A = Q * D * Q'
    return A' * A
end

"""
    adhoc_add_correlation!(x::SnpArray, ρ::Float64, pos::Int64, location::Vector{Int})

Makes 1 SNP (a column of `x`) correlate with SNPs in `location` with correlation coefficient roughly `ρ`.

# Arguments:
- `x`: the snparray
- `ρ`: correlation coefficient
- `pos`: the position of the target SNP that everything would be correlated to
- `location`: All the SNPs that shall be correlated with the SNP at position `pos` with correlation `ρ`.
"""
function adhoc_add_correlation!(x::SnpArray, ρ::Float64, pos::Int64, location::Vector{Int})
    @assert 0 <= ρ <= 1 "correlation coefficient must be in (0, 1) but was $ρ"

    for loc in location
        for i in 1:size(x, 1)
            rand(Bernoulli(ρ)) && (x[i, loc] = x[i, pos]) #make 2nd column the same as 1st ~ρ% of the time
        end
        corr = cor(x[:, loc], x[:, pos])
    end
end

"""
    make_bim_fam_files(x::SnpArray, y, name::String)

Creates .bim and .bed files from a SnpArray.

# Arguments:
- `x`: A SnpArray (i.e. `.bed` file on the disk) for which you wish to create corresponding `.bim` and `.fam` files.
- `name`: string that should match the `.bed` file (Do not include `.bim` or `.fam` extensions in `name`).
- `y`: Trait vector that will go in to the 6th column of `.fam` file.
"""
function make_bim_fam_files(x::SnpArray, y::AbstractVecOrMat, name::String)
    ly = size(y, 1)
    n, p = size(x)
    @assert n == ly "dimension mismatch: phenotype data has length $ly but SnpArray has $n samples"

    #create .bim file structure: https://www.cog-genomics.org/plink2/formats#bim
    open(name * ".bim", "w") do f
        for i in 1:p
            write(f, "1\tsnp$i\t0\t1\t1\t2\n")
        end
    end

    #create .fam file structure: https://www.cog-genomics.org/plink2/formats#fam
    traits = size(y, 2)
    open(name * ".fam", "w") do f
        for i in 1:n
            write(f, "$i\t1\t0\t0\t1")
            for j in 1:traits
                write(f, "\t$(y[i, j])")
            end
            write(f, "\n")
        end
    end
end


using Statistics
using Plots
using BenchmarkTools
using Profile
using FFTW
using Statistics
using Optim
using Measures
using Images, FileIO

#using Profile
using LinearAlgebra
using SparseArrays

# put the cwd on the path so Julia can find the module
push!(LOAD_PATH, pwd()*"/main")
using DHC_2DUtils

@time fhash = fink_filter_hash(1, 8, nx=64, pc=1, wd=1)

function realspace_filter(Nx, f_i, f_v)

    zarr = zeros(ComplexF64, Nx, Nx)
    for i = 1:length(f_i)
        zarr[f_i[i]] = f_v[i] # filter*image in Fourier domain
    end
    filt = ifft(zarr)  # real space, complex
    return filt
end

Nx = 256
bar = rand(2500,100)
fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)
#M = S2_iso_matrix(fhash, sparse=false)
#Ms = sparse(M)
Ms = fhash["S2_iso_mat"]
M = Array(Ms)
@benchmark M*bar
@benchmark Ms*bar

# M*bar test
# Nx         full   sparse
# 128       0.8 ms     4 μs
# 128x100   45  ms   500 μs
# 256       1.9 ms     7 μs
# 256x100   90  ms   800 μs

function myrealifft(fbar)
    (Nx, Ny)  = size(fbar)
    fbars = circshift(fbar[end:-1:1,end:-1:1],(1,1))
    myarg = fbar[1:Ny÷2+1,:]+conj(fbars[1:Ny÷2+1,:])
    irfft(myarg,Ny).* 0.5
end

function wst_S1S2_derivfast(image::Array{Float64,2}, filter_hash::Dict)
    #Works now.
    #Possible Bug: You assumed zero freq in wrong place and N-k is the wrong transformation.
    # Use 2 threads for FFT
    FFTW.set_num_threads(2)

    # array sizes
    (Nx, Ny)  = size(image)
    (Nf, )    = size(filter_hash["filt_index"])

    # allocate output array
    dS1dp  = zeros(Float64, Nx, Nx, Nf)
    dS2dp  = zeros(Float64, Nx, Nx, Nf, Nf)

    # allocate image arrays for internal use
    im_rdc_0_1 = Array{ComplexF64, 3}(undef, Nx, Ny, Nf)  # real domain, complex
    im_rd_0_1  = Array{Float64, 3}(undef, Nx, Ny, Nf)  # real domain, complex

    ## 1st Order
    im_fd_0 = fft(image)  # total power=1.0

    # unpack filter_hash
    f_ind   = filter_hash["filt_index"]  # (J, L) array of filters represented as index value pairs
    f_val   = filter_hash["filt_value"]

    f_ind_rev = [[CartesianIndex(mod1(Nx+2 - ci[1],Nx), mod1(Nx+2 - ci[2],Nx)) for ci in f_Nf] for f_Nf in f_ind]

    # tmp arrays for all tmp[f_i] = arr[f_i] .* f_v opns
    zarr1     = zeros(ComplexF64, Nx, Nx)
    fz_fψ1    = zeros(ComplexF64, Nx, Nx)
    fterm_a   = zeros(ComplexF64, Nx, Nx)
    fterm_ct1 = zeros(ComplexF64, Nx, Nx)
    fterm_ct2 = zeros(ComplexF64, Nx, Nx)
    #rterm_bt2 = zeros(ComplexF64, Nx, Nx)
    # make a FFTW "plan" for an array of the given size and type
    P = plan_ifft(im_fd_0)  # P is an operator, P*im is ifft(im)

    # Loop over filters
    for f1 = 1:Nf
        f_i1 = f_ind[f1]  # CartesianIndex list for filter1
        f_v1 = f_val[f1]  # Values for f_i1
        f_i1rev = f_ind_rev[f1]

        zarr1[f_i1] = f_v1 .* im_fd_0[f_i1] #F[z]
        zarr1_rd = P*zarr1 #for frzi
        fz_fψ1[f_i1] = f_v1 .* zarr1[f_i1]
        fz_fψ1_rd = P*fz_fψ1
        dS1dp[:,:,f1] = 2 .* real.(fz_fψ1_rd) #real.(conv(I_λ, ψ_λ))
        #CHECK; that this equals derivS1fast code

        #dS2 loop prep
        #ψ_λ1  = realspace_filter(Nx, f_i1, f_v1) #Slow
        #fcψ_λ1 = fft(conj.(ψ_λ1)) #Slow
        I_λ1 = sqrt.(abs2.(zarr1_rd))
        fI_λ1 = fft(I_λ1)
        rterm_bt1 = zarr1_rd./I_λ1 #Z_λ1/I_λ1_bar.
        rterm_bt2 = conj.(zarr1_rd)./I_λ1
        for f2 = 1:Nf
            f_i2 = f_ind[f2]  # CartesianIndex list for filter2
            f_v2 = f_val[f2]  # Values for f_i2

            fterm_a[f_i2] = fI_λ1[f_i2] .* (f_v2).^2 #fterm_a = F[I_λ1].F[ψ_λ]^2
            rterm_a = P*fterm_a                #Finv[fterm_a]

            fterm_bt1 = fft(rterm_a .* rterm_bt1) #F[Finv[fterm_a].*Z_λ1/I_λ1_bar] for T1
            fterm_bt2 = fft(rterm_a .* rterm_bt2) #F[Finv[fterm_a].*Z_λ1_bar/I_λ1_bar] for T2
            fterm_ct1[f_i1] = fterm_bt1[f_i1] .* f_v1    #fterm_b*F[ψ_λ]
            #fterm_ct2 = fterm_bt2 .* fcψ_λ1             #Slow
            #println(size(fterm_ct2[f_i1rev]),size(fterm_bt2[f_i1rev]),size(f_v1),size(conj.(f_v1)))
            fterm_ct2[f_i1rev] = fterm_bt2[f_i1rev] .* f_v1 #f_v1 is real
            #fterm_ct2slow = fterm_bt2 .* fcψ_λ1
            dS2dp[:, :, f1, f2] = real.(P*(fterm_ct1 + fterm_ct2))
            #println("Term2",fterm_ct2)

            #Reset f2 specific tmp arrays to 0
            fterm_a[f_i2] .= 0
            fterm_ct1[f_i1] .=0
            fterm_ct2[f_i1rev] .=0

        end
        # reset all reused variables to 0
        zarr1[f_i1] .= 0
        fz_fψ1[f_i1] .= 0

    end
    return dS1dp, dS2dp
end

function derivtestS1S2(Nx)
    eps = 1e-4
    fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1)
    im = rand(Float64, Nx, Nx).*0.1
    im[6,6]=1.0
    im1 = copy(im)
    im0 = copy(im)
    im1[2,3] += eps/2
    im0[2,3] -= eps/2
    dS1dp, dS2dp = wst_S1S2_derivfast(im, fhash)
    println(size(dS2dp))
    der0=DHC_compute(im0,fhash,doS2=true)
    der1=DHC_compute(im1,fhash,doS2=true)
    dS = (der1-der0) ./ eps

    Nf = length(fhash["filt_index"])
    i0 = 3+Nf
    blarg = dS2dp[2,3,:,:]
    diff = dS[i0:end]-reshape(blarg,Nf*Nf)
    println(dS[i0:end])
    println("and")
    println(blarg)
    println("stdev: ",std(diff))
    println()
    println(diff)

    return
end

derivtestS1S2(8)

println(1)

function derivtest2(Nx)
    eps = 1e-4
    fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1)
    im = rand(Float64, Nx, Nx).*0.1
    im[6,6]=1.0
    im1 = copy(im)
    im0 = copy(im)
    im1[2,3] += eps/2
    im0[2,3] -= eps/2
    dS20dp = wst_S20_deriv(im, fhash)

    der0=DHC_compute(im0,fhash,doS2=false,doS20=true)
    der1=DHC_compute(im1,fhash,doS2=false,doS20=true)
    dS = (der1-der0) ./ eps

    Nf = length(fhash["filt_index"])
    i0 = 3+Nf
    blarg = dS20dp[2,3,:,:]
    diff = dS[i0:end]-reshape(blarg,Nf*Nf)
    println(dS[i0:end])
    println("and")
    println(blarg)
    println("stdev: ",std(diff))
    println()
    println(diff)

    return
end

derivtest2(8)

Nx=128
fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1)
im = zeros(Float64, Nx, Nx)
im[6,6]=1.0
@benchmark blarg = wst_S1_deriv(im, fhash)

# S1 deriv time, Jan 30 version (old) compared to Feb 14 version from NM
# Nx    old      Feb 14
#   32     17 ms   0.5 ms
#   64     34 ms   3.5 ms
#  128    115 ms    20 ms
#  256    520 ms    92 ms
#  512   2500 ms   720 ms   540 ms with 2 threads
# 1024   9500 ms  3300 ms  2500 ms with 2

Nx = 128
fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1)
im = zeros(Float64, Nx, Nx)
im[6,6]=1.0
@benchmark blarg,blafb2 = wst_S1S2_derivfast(im, fhash)

# S20 vs S2 deriv time, Feb 15
# Nx     S20
#   8      1 ms    27 ms
#  16      7      104
#  32     50      300
#  64    400 ms   1.1 sec
# 128    3.3 sec  5.0
# 256   17.2 sec
# 512   ---

Nx = 32
fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1)
im = zeros(Float64, Nx, Nx)
im[6,6]=1.0
@benchmark blarg = wst_S20_deriv(im, fhash)

Profile.clear()
@profile blarg = wst_S20_deriv(im, fhash)
Juno.profiler()

# S20 deriv time, Jan 30
# Nx     Jan 30  Feb 14
#   8     28 ms   1 ms
#  16    112      7
#  32    320     50
#  64   1000    400
# 128   5 sec     3.3 sec
# 256   ---      17.2 sec
# 512   ---

print(1)

# wst_S1_deriv agrees with brute force at 1e-10 level.
function derivtest(Nx)
    eps = 1e-4
    fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1)
    im = rand(Float64, Nx, Nx).*0.1
    im[6,6]=1.0
    im1 = copy(im)
    im0 = copy(im)
    im1[2,3] += eps/2
    im0[2,3] -= eps/2
    blarg = wst_S1_deriv(im, fhash)

    der0=DHC_compute(im0,fhash,doS2=false)
    der1=DHC_compute(im1,fhash,doS2=false)
    dS = (der1-der0) ./ eps

    diff = dS[3:end]-blarg[2,3,:]
    println(dS[3:end])
    println("and")
    println(blarg[2,3,:])
    println("stdev: ",std(diff))

    #plot(dS[3:end])
    #plot!(blarg[2,3,:])
    #plot(diff)
    return
end

derivtest(128)

Nx=32
im = rand(Nx,Nx)
@time fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1)
@benchmark Sarr = DHC_compute(im, fhash, doS2=false)  # 0.36 ms
@benchmark Sarr = DHC_compute(im, fhash, doS2=true)   # 6.4 ms
@benchmark Sarr = DHC_compute(im, fhash, doS2=false, doS12=true)  # 0.7 ms
@benchmark Sarr = DHC_compute(im, fhash, doS2=false, doS20=true)  # 2.1 ms

Nx=64
im = rand(Nx,Nx)
@time fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1)
@benchmark Sarr = DHC_compute(im, fhash, doS2=false)  # 1.3 ms
@benchmark Sarr = DHC_compute(im, fhash, doS2=true)   # 19 ms
@benchmark Sarr = DHC_compute(im, fhash, doS2=false, doS12=true)  # 2.8 ms
@benchmark Sarr = DHC_compute(im, fhash, doS2=false, doS20=true)  # 6.6 ms

Nx=256
im = rand(Nx,Nx)
@time fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1)
@benchmark Sarr = DHC_compute(im, fhash, doS2=false)  # 25 ms
@benchmark Sarr = DHC_compute(im, fhash, doS2=true)   # 375 ms
@benchmark Sarr = DHC_compute(im, fhash, doS2=false, doS12=true)  # 55 ms
@benchmark Sarr = DHC_compute(im, fhash, doS2=false, doS20=true)  # 130 ms

Nx = 64
fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1)
im = rand(Nx,Nx)
S_targ = DHC_compute(im, fhash, doS2=false)

im = rand(Nx,Nx)

function wst_synth(im_init, fixmask)
    # fixmask -  0=float in fit   1=fixed

    function wst_synth_chisq(vec_in)

        thisim = copy(im_init)
        thisim[indfloat] = vec_in

        Sarr  = DHC_compute(thisim, fhash, doS2=false)
        diff  = (Sarr .- S_targ)[3:end]

        # should have some kind of weight here
        chisq = diff'*diff
        return chisq

    end

    function wst_synth_dchisq(storage, vec_in)

        thisim = copy(im_init)
        thisim[indfloat] = vec_in
        dSdp   = wst_S1_deriv(thisim, fhash)
        S1arr  = DHC_compute(thisim, fhash, doS2=false)
        diff  = (S1arr - S_targ)[3:end]

        # dSdp matrix * S1-S_targ is dchisq
        dchisq_im = (reshape(dSdp, Nx*Nx, Nf) * diff).*2
        dchisq = reshape(dchisq_im, Nx, Nx)[indfloat]

        storage .= dchisq
    end

    (Nx, Ny)  = size(im_init)
    if Nx != Ny error("Input image must be square") end
    (Nf, )    = size(fhash["filt_index"])

    # index list of pixels to float in fit
    indfloat = findall(fixmask .== 0)

    # initial values for floating pixels
    vec_init = im_init[indfloat]
    println(length(vec_init), " pixels floating")

    eps = zeros(size(vec_init))
    eps[1] = 1e-4
    chisq1 = wst_synth_chisq(vec_init+eps./2)
    chisq0 = wst_synth_chisq(vec_init-eps./2)
    brute  = (chisq1-chisq0)/1e-4

    clever = zeros(size(vec_init))
    _bar = wst_synth_dchisq(clever, vec_init)
    println("Brute:  ",brute)
    println("Clever: ",clever[1])

    # call optimizer
    res = optimize(wst_synth_chisq, wst_synth_dchisq, vec_init, BFGS())

    # copy results into pixels of output image
    im_synth = copy(im_init)
    im_synth[indfloat] = Optim.minimizer(res)
    println(res)

    return im_synth
end

Nx    = 128
fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1)
im    = rand(Nx,Nx)
fixmask = im .> 0.5
S_targ = DHC_compute(im, fhash, doS2=false)

init = copy(im)
init[findall(fixmask .==0)] .= 0

foo = wst_synth(init, fixmask)

#Optim.minimizer(optimize(FOM, im, LBFGS()))

# using dchisq function
# size   t(BFGS) t(LBFGS) [sec]
# 16x16    1
# 32x32    4        9
# 64x64    20      52
# 128x128         153 (fitting 50% of pix)

# without dchisq function
# size   t(BFGS) t(LBFGS) [sec]
# 16x16   27
# 32x32  355      110
# 64x64          1645
# 128x128        est 28,000  (8 hrs)

function readdust()

    RGBA_img = load(pwd()*"/scratch_DF/t115_clean_small.png")
    img = reinterpret(UInt8,rawview(channelview(RGBA_img)))

    return img[1,:,:]
end

function wst_synthS20(im_init, fixmask, S_targ, S20sig; iso=false)
    # fixmask -  0=float in fit   1=fixed

    function wst_synth_chisq(vec_in)

        thisim = copy(im_init)
        thisim[indfloat] = vec_in

        M20 = fhash["S2_iso_mat"]

        S20 = DHC_compute(thisim, fhash, doS2=false, doS20=true, norm=false, iso=iso)
        i0 = 3+(iso ? N1iso : Nf)
        diff  = ((S20[i0:end] - S_targ)./S20sig)

        # should have some kind of weight here
        chisq = diff'*diff
        println(chisq)
        return chisq

    end

    function wst_synth_dchisq(storage, vec_in)

        thisim = copy(im_init)
        thisim[indfloat] = vec_in

        dS20dp = reshape(wst_S20_deriv(thisim, fhash), Nx*Nx, Nf*Nf)
        if iso
            M20 = fhash["S2_iso_mat"]
            dS20dp = dS20dp * M20'
        end
        i0 = 3+(iso ? N1iso : Nf)

        S20arr = (DHC_compute(thisim, fhash, doS2=false, doS20=true, norm=false, iso=iso))[i0:end]

        # put both factors of S20sig in this array to weight
        diff   = (S20arr - S_targ)./(S20sig.^2)
        #println("size of diff", size(diff))
        #println("size of dS20dp", size(reshape(dS20dp, Nx*Nx, Nf*Nf)))
        # dSdp matrix * S1-S_targ is dchisq
        #dchisq_im = (reshape(dS20dp, Nx*Nx, Nf*Nf) * diff).*2
        dchisq_im = (dS20dp * diff).*2
        dchisq = reshape(dchisq_im, Nx, Nx)[indfloat]

        storage .= dchisq
    end

    (Nx, Ny)  = size(im_init)
    if Nx != Ny error("Input image must be square") end
    (N1iso, Nf)    = size(fhash["S1_iso_mat"])

    # index list of pixels to float in fit
    indfloat = findall(fixmask .== 0)

    # initial values for floating pixels
    vec_init = im_init[indfloat]
    println(length(vec_init), " pixels floating")

    eps = zeros(size(vec_init))
    eps[1] = 1e-4
    chisq1 = wst_synth_chisq(vec_init+eps./2)
    chisq0 = wst_synth_chisq(vec_init-eps./2)
    brute  = (chisq1-chisq0)/1e-4

    clever = zeros(size(vec_init))
    _bar = wst_synth_dchisq(clever, vec_init)
    println("Brute:  ",brute)
    println("Clever: ",clever[1])

    # call optimizer
    #res = optimize(wst_synth_chisq, wst_synth_dchisq, copy(vec_init), BFGS())
    res = optimize(wst_synth_chisq, wst_synth_dchisq, copy(vec_init),
        ConjugateGradient(), Optim.Options(iterations=1000))

    # copy results into pixels of output image
    im_synth = copy(im_init)
    im_synth[indfloat] = Optim.minimizer(res)
    println(res)

    return im_synth
end

function S20_weights(im, fhash, Nsam=10; iso=iso)

    (Nx, Ny)  = size(im)
    if Nx != Ny error("Input image must be square") end
    (N1iso, Nf)    = size(fhash["S1_iso_mat"])

    # fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1)

    S20   = DHC_compute(im, fhash, doS2=false, doS20=true, norm=false, iso=iso)
    Ns    = length(S20)
    S20arr = zeros(Float64, Ns, Nsam)
    for j=1:Nsam
        noise = rand(Nx,Nx)
        S20arr[:,j] = DHC_compute(im+noise, fhash, doS2=false, doS20=true, norm=false, iso=iso)
    end

    wt = zeros(Float64, Ns)
    for i=1:Ns
        wt[i] = std(S20arr[i,:])
    end

    return wt
end

# read dust map
dust = Float64.(readdust())
dust = dust[1:256,1:256]

Nx     = 16
doiso  = false
fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)
(N1iso, Nf)    = size(fhash["S1_iso_mat"])
i0 = 3+(doiso ? N1iso : Nf)
img    = imresize(dust,(Nx,Nx))
#fixmask = rand(Nx,Nx) .< 0.1
fixmask = falses((Nx, Nx))

S_targ = DHC_compute(img, fhash, doS2=false, doS20=true, norm=false, iso=doiso)
S_targ = S_targ[i0:end]

init = copy(img)
floatind = findall(fixmask .==0)
init[floatind] .+= rand(length(floatind)).*50 .-25

S20sig = S20_weights(img, fhash, 100, iso=doiso)
S20sig = S20sig[i0:end]
foo = wst_synthS20(init, fixmask, S_targ, S20sig, iso=doiso)
S_foo = DHC_compute(foo, fhash, doS2=false, doS20=true, norm=false, iso=doiso)

plot_synth_QA(im, init, foo, fhash)

function plot_synth_QA(ImTrue, ImInit, ImSynth, fhash; fname="test2.png")

    # -------- define plot1 to append plot to a list
    function plot1(ps, image; clim=nothing, bin=1.0, fsz=16, label=nothing)
        # ps    - list of plots to append to
        # image - image to heatmap
        # clim  - tuple of color limits (min, max)
        # bin   - bin down by this factor, centered on nx/2+1
        # fsz   - font size
        marg   = 1mm
        nx, ny = size(image)
        nxb    = nx/round(Integer, 2*bin)

        # -------- center on nx/2+1
        i0 = max(1,round(Integer, (nx/2+2)-nxb-1))
        i1 = min(nx,round(Integer, (nx/2)+nxb+1))
        lims = [i0,i1]
        subim = image[i0:i1,i0:i1]
        push!(ps, heatmap(image, aspect_ratio=:equal, clim=clim,
            xlims=lims, ylims=lims, size=(400,400),
            legend=false, xtickfontsize=fsz, ytickfontsize=fsz,#tick_direction=:out,
            rightmargin=marg, leftmargin=marg, topmargin=marg, bottommargin=marg))
        if label != nothing
            annotate!(lims'*[.96,.04],lims'*[.09,.91],text(label,:left,:white,32))
        end
        return
    end

    # -------- initialize array of plots
    ps   = []
    clim  = (0,200)
    clim2 = (0,200).-100

    # -------- 6 panel QA plot
    plot1(ps, ImTrue, clim=clim, label="True")
    plot1(ps, ImSynth, clim=clim, label="Synth")
    plot1(ps, ImInit, clim=clim, label="Init")
    plot1(ps, ImInit-ImTrue, clim=clim2, label="Init-True")
    plot1(ps, ImInit-ImSynth, clim=clim2, label="Init-Synth")
    plot1(ps, ImSynth-ImTrue, clim=clim2, label="Synth-True")

    myplot = plot(ps..., layout=(3,2), size=(1400,2000))
    savefig(myplot, fname)
end

plot_synth_QA(im, init, foo, fhash)

# using dchisq function with S20
# size   t(BFGS) t(LBFGS) [sec]
# 8x8      30
# 16x16    90
# 32x32   189                           74
# 64x64  1637                          531      642 iso
# 128x128                              2 hrs


abstract type ObjectiveFunc end
abstract type Objective1D  <: ObjectiveFunc end
abstract type Objective2D  <: ObjectiveFunc end

Base.string(o::ObjectiveFunc) = o.name

#Test functions
mutable struct DistributionFunc <: Objective1D
    name::String
    d::Distribution
    σ::Float64
    xlim::Tuple{Float64,Float64}
    g_max::Float64
    g_min::Float64
    x_max::Float64

    function DistributionFunc(name::String, d::Distribution, σ::Float64; xmin=-1.0, xmax=1.0, n=1000)
        o = new(name, d, σ, (xmin,xmax))
        xs = linspace(xmin, xmax, n)
        o.g_max, i_max = findmax(truth.(o, xs))
        o.x_max = xs[i_max]
        o.g_min = minimum(truth.(o, xs))
        o
    end
end
function NarrowBump(d=Normal(-0.25,0.05); σ=0.3)
    DistributionFunc("NarrowBump", d, σ)
end
function WideBump(d=Normal(0.15,0.25); σ=0.3)
    DistributionFunc("WideBump", d, σ)
end
function TwoBumps(d=MixtureModel([Normal(-0.5,0.12), Normal(0.1,0.3)],
                                  [0.35,0.65]); σ=0.3)
    DistributionFunc("TwoBumps", d, σ)
end
function ThreeBumps(d=MixtureModel([Normal(-0.5,0.15),Normal(0.0,0.3),Normal(0.5,0.07)],
                                   [0.27,0.57,0.16]); σ=0.3)
    DistributionFunc("ThreeBumps", d, σ)
end
truth(o::DistributionFunc, x) = pdf(o.d, x)
function (o::DistributionFunc)(x, rng::AbstractRNG=Base.GLOBAL_RNG)
    truth(o,x) + o.σ*randn(rng)
end

mutable struct SincFunc <: Objective1D
    name::String
    a::Float64
    b::Float64
    c::Float64
    σ::Float64
    xlim::Tuple{Float64,Float64}
    g_max::Float64
    x_max::Float64
    g_min::Float64

    function SincFunc(name::String, a::Float64, b::Float64, c::Float64, σ::Float64; xmin=-1.0, xmax=1.0, n=1000)
        o = new(name, a, b, c, σ, (xmin,xmax))
        xs = linspace(xmin, xmax, n)
        o.g_max, i_max = findmax(truth.(o, xs))
        o.x_max = xs[i_max]
        o.g_min = minimum(truth.(o, xs))
        o
    end
end
function Sinc(; a=20.0, b=0.3, c=0.5, σ=0.3)
    SincFunc("Sinc", a,b,c,σ)
end
sinc_truth(a::Float64, b::Float64, c::Float64, x::Float64) = sinc(a * (x-b) + c)
truth(o::SincFunc, x) = sinc_truth(o.a, o.b, o.c, x)
function (o::SincFunc)(x, rng::AbstractRNG=Base.GLOBAL_RNG)
    truth(o,x) + o.σ*randn(rng)
end
function circle(x,y,r)
    c = Shape(Plots.partialcircle(0,2π,20,r))
    translate!(c, x, y)
    c
end
@recipe function plot(o::Objective1D; xmin=-1.0, xmax=1.0, n=100)
    xlim := o.xlim
    ylim := (o.g_min - 1.0, o.g_max + 1.0)
    xlabel := "x"
    ylabel := "f(x)"
    @series begin
        seriestype := :path
        label := "true mean"
        xs = linspace(xmin, xmax, n)
        xs, truth.(o, xs)
    end
end
@recipe function plot(o::Objective1D, rng::AbstractRNG; xmin=-1.0, xmax=1.0, n=100)
    xlim := o.xlim
    ylim := (o.g_min - 1.0, o.g_max + 1.0)
    xlabel := "x"
    ylabel := "f(x)"
    @series begin
        seriestype := :path
        label := "noisy sample"
        xs = linspace(xmin, xmax, n)
        xs, o.(xs, rng)
    end
end

Base.maximum(o::Objective1D) = o.g_max
Base.minimum(o::Objective1D) = o.g_min

#####
# 2D
mutable struct DistributionFunc2D <: Objective2D
    name::String
    d::Distribution
    σ::Float64
    xlim::Tuple{Float64,Float64}
    ylim::Tuple{Float64,Float64}
    g_max::Float64
    g_min::Float64
    x_max::Vector{Float64}

    function DistributionFunc2D(name::String, d::Distribution, σ::Float64;
                                xlim=(-2.0,2.0), ylim=(-2.0,2.0), n=1000)
        o = new(name, d, σ, xlim, ylim)
        Xs = [[x,y] for x in linspace(xlim..., n), y in linspace(ylim..., n)]
        o.g_max, i_max = findmax(truth.(o,Xs))
        o.x_max = Xs[i_max]
        o.g_min = minimum(truth.(o,Xs))
        o
    end
end
function NarrowBump2D(d=MvNormal([0.4,0.4],[0.2,0.2]); σ=2.0)
    DistributionFunc2D("NarrowBump2D", d, σ)
end
function WideBump2D(d=MvNormal([0.25,0.25],[0.2,0.7]); σ=0.5)
    DistributionFunc2D("WideBump2D", d, σ)
end
truth(o::DistributionFunc2D, x) = pdf(o.d, x)
function (o::DistributionFunc2D)(x, rng::AbstractRNG=Base.GLOBAL_RNG)
    truth(o,x) + o.σ*randn(rng)
end

mutable struct Rosenbrock <: Objective2D
    name::String
    σ::Float64
    xlim::Tuple{Float64,Float64}
    ylim::Tuple{Float64,Float64}
    g_max::Float64
    g_min::Float64
    x_max::Vector{Float64}

    function Rosenbrock(σ::Float64=10.0; xlim=(-2.0,2.0), ylim=(-2.0,2.0), n=1000)
        o = new("Rosenbrock", σ, xlim, ylim)
        Xs = [[x,y] for x in linspace(xlim..., n), y in linspace(ylim..., n)]
        o.g_max, i_max = findmax(truth.(o,Xs))
        o.x_max = Xs[i_max]
        o.g_min = minimum(truth.(o,Xs))
        o
    end
end
truth(o::Rosenbrock, x) = -rosenbrock(x)
function (o::Rosenbrock)(x, rng::AbstractRNG=Base.GLOBAL_RNG)
    truth(o,x) + o.σ*randn(rng)
end

@recipe function plot(o::Objective2D; n=100)
    xlim := o.xlim
    ylim := o.ylim
    zlim := (o.g_min - 1.0, o.g_max + 1.0)
    xlabel := "x"
    ylabel := "y"
    zlabel := "f(x,y)"
    @series begin
        seriestype := :heatmap
        label := "true mean"
        xs = linspace(o.xlim..., n)
        ys = linspace(o.ylim..., n)
        xs, ys, (x,y)->truth(o, [x,y])
    end
end
@recipe function plot(o::Objective2D, rng::AbstractRNG; n=100)
    xlim := o.xlim
    ylim := o.ylim
    zlim := (o.g_min - 1.0, o.g_max + 1.0)
    xlabel := "x"
    ylabel := "y"
    zlabel := "f(x,y)"
    @series begin
        seriestype := :heatmap
        label := "noisy sample"
        xs = linspace(o.xlim..., n)
        ys = linspace(o.ylim..., n)
        xs, ys, (x,y)->o([x,y],rng)
    end
end

Base.maximum(o::Objective2D) = o.g_max
Base.minimum(o::Objective2D) = o.g_min


using Gadfly, DataFrames

xs = 0:0.1:20

df = DataFrame(
    x=xs,
    y=cos(xs),
    ymin=cos(xs) .- 0.5,
    ymax=cos(xs) .+ 0.5,
)

plot(df, x=:x, y=:y, ymin=:ymin, ymax=:ymax, Geom.line, Geom.ribbon)


function Simple_AON_Louvain(H::hypergraph;startclusters="singletons",gamma=1.0,
	clusterpenalty=0,kmax=maximum(keys(H.E)),maxits::Int64=100,
	randflag::Bool = false,verbose=true,Z0 = collect(1:length(H.D)))
	"""
    Simplest version of AON Louvain, the doesn't require setting an initial
    intensity function Ω. Intensity function is generated from one of a few
	different choices:
		singletons: learn omega from all nodes being in a singleton cluster
		cliquelouvain: use clique expansion + graph louvain warm start with parameter gamma
		starlouvain: use star expansion + graph louvain warm start with parameter gamma
	"""

	n = length(H.D)
	if startclusters == "singletons"
		Z0 = collect(1:n)
	elseif startclusters == "cliquelouvain"
		Z0 = CliqueExpansionModularity(H,gamma)
	elseif startclusters == "starlouvain"
		Z0 = StarExpansionModularity(H,gamma)
	else
		error("There is no starting clustering called $startclusters. Choose from: [\"singletons\", \"cliquelouvain\",\"starlouvain\"]")
	end

	d = Float64.(H.D)
	n = length(d)
	kmax = maximum(keys(H.E))
	He2n, weights = hypergraph2incidence(H)
    e2n = incidence2elist(He2n);
    n2e = incidence2elist(SparseArrays.sparse(He2n'))
    m = length(e2n)
    elen = zeros(Int64,m)
    for e = 1:m
        elen[e] = length(e2n[e])
    end
	β, γ,omega = learn_omega_aon(e2n,weights,Z0,kmax,d,n)
	Zset = AON_Louvain(n2e,e2n,weights,d,elen,β,γ,kmax,randflag,maxits,verbose,Z0,clusterpenalty)
	Z = Zset[:,end]
	return Z
end

function AON_Louvain(H::hypergraph,β::Vector{Float64},γ::Vector{Float64};maxits::Int64=100,
	verbose::Bool=true,randflag::Bool = false,clusterpenalty=0.0,Z0 = collect(1:length(H.D)))
	"""
	Calling AON_Louvain without explicit intensity function Ω. Just requires
	cut and volume penalties β and γ. Assumes hyperedges are of size >= 2.
	"""
	if length(H.E[1]) > 0
		println("This code assumes a hypergraph where hyperedges have at least two nodes.")
	end
	if γ[1] != 0 || β[1] != 0
		println("Code assumes there are no hyperedges of size 1. Setting β[1] = γ[1] = 0.")
		β[1] = 0
		γ[1] = 0
	end
	d = Float64.(H.D)
	n = length(d)
	kmax = maximum(keys(H.E))
	@assert(kmax == length(β))
	He2n, weights = hypergraph2incidence(H)
    e2n = incidence2elist(He2n);
    n2e = incidence2elist(SparseArrays.sparse(He2n'))
    m = length(e2n)
    elen = zeros(Int64,m)
    for e = 1:m
        elen[e] = length(e2n[e])
    end
	Zset = AON_Louvain(n2e,e2n,weights,d,elen,β,γ,kmax,randflag,maxits,verbose,Z0,clusterpenalty)
	Z = Zset[:,end];
end

function alternate_hypergraph_storage(H::hypergraph)
	"""
	It will be quicker in some cases to store the hypergraph in an alternate format.
	"""
	d = Float64.(H.D)
	n = length(d)
	kmax = maximum(keys(H.E))
	He2n, weights = HyperModularity.hypergraph2incidence(H)
	e2n = incidence2elist(He2n);
	n2e = incidence2elist(SparseArrays.sparse(He2n'))
	m = length(e2n)
	elen = zeros(Int64,m)
    for e = 1:m
        elen[e] = length(e2n[e])
    end
	return d,n,kmax,e2n,n2e,weights,elen
end

function AON_Louvain(H::hypergraph,Ω::IntensityFunction;α,clusterpenalty=0,kmax=maximum(keys(H.E)),maxits::Int64=100,bigInt::Bool=true,verbose=true,scan_order="random",Z0 = collect(1:length(H.D)))
	"""
	Call AON_Louvain with a function call similar to the symmetric HMLL function.
	"""
	randflag = !(scan_order == "lexical")
    β, γ, e2n, n2e,w,d,elen = AON_Inputs(H,Ω.ω,α,kmax)
    Zset = AON_Louvain(n2e,e2n,w,d,elen,β,γ,kmax,randflag,maxits,verbose,Z0,clusterpenalty)
    Z = Zset[:,end]
    return Z
end

function AON_Louvain(node2edges::Vector{Vector{Int64}},
    edge2nodes::Vector{Vector{Int64}},w::Vector{Float64},
    d::Vector{Float64},elen::Vector{Int64},
    β::Vector{Float64},γ::Vector{Float64},
    kmax::Int64,randflag::Bool=false,maxits::Int64=100,verbose::Bool=true,
    Zwarm::Vector{Int64}=Vector{Int64}(),clusterpenalty = 0)
    """
    Supernode version of Hypergraph Louvain for planted partition model
        (all-or-nothing cut penalties).

        * node2edges[i] = list of edge labels that node i is in
        * edge2nodes[e] = list of node labels that edge e contains
        * w[e] = weight of hyperedge e
        * d[i] = degree or weight of node i
        * elen[e] = number of nodes in hyperedge e
                     note that this may be > length(edge2nodes[e])
                     if some of the nodes are supernodes
        * β = scaling values of for hyperedge cuts
        * γ = scaling values for volumes
        * kmax = maximum hyperedge size
        * Zwarm = warm start clustering (optional)
        * randflag = whether or not to permute node order
        * maxits = Maximum number of greedy passes over the node set
        * clusterpenalty = include term clusterpenalty*log(k) in objective
            so if clusterpenalty > 0, there is incentive to form fewer clusters
    """

    n = length(d)
    # Step 1: greedy moves until no more improvement
    Z, improved = ANHL_Step(node2edges,edge2nodes,w,d,elen,β,γ,kmax,
                            randflag,maxits,verbose,Zwarm,clusterpenalty)

    # Store all clusterings found
    if improved
        Zs = Z
        Z_old = copy(Z)
    else
        Zs = Z
    end

    # As long as something is still improving each time you call step 1, keep going
    while improved

        # Step 2: Collapse the clustering into supernodes
        @assert(minimum(Z_old) > 0)
        e2n = deepcopy(edge2nodes)
        Change_id_H!(e2n,Z_old)
        # Merge hyperedges that are the same and remove size 1 hyperedges
        e2n, wSuper, elenSuper =  Condense_H(e2n,w,elen)
        He2n = elist2incidence(e2n,maximum(Z_old))
        n2e = incidence2elist(He2n,true)
        dSuper = condense_d(d,Z_old)
        ########
        # Step 1: Go back to greedy local moves, this time on the reduced hypergraph
        Zsuper, improved = ANHL_Step(n2e,e2n,wSuper,dSuper,elenSuper,β,γ,
                                kmax,randflag,maxits,verbose,Vector{Int64}(),clusterpenalty)
        N = length(Zsuper)    # N = number of supernodes = number clusters from last round

        @assert(minimum(Zsuper)>0)
        @assert(length(Zsuper)==length(dSuper))
        # println("$N, $improved")
        # Undo the procedure that merged nodes into super nodes, so that you get a clustering vector of length n
        # Do this only if the last call to Step 1 led to at least one greedy move.
        if improved
            # Extract what that new clustering is
            Z_new = zeros(Int64,n)

            # For each supernode, place all original node IDs that make it up
            # into a cluster of the supernodes label
            for i = 1:N

                # Get the cluster that supernode i is in
                SuperI_Cluster = Zsuper[i]

                # Get individual node ID that are in supernode i.
                SuperI_nodes = findall(x->x==i,Z_old)
                Z_new[SuperI_nodes] .= SuperI_Cluster
            end
            @assert(minimum(Z_new)>0)
            Zs = [Zs Z_new]
            Z_old = copy(Z_new)
        end
    end

    return Zs
end

function ANHL_Step(node2edges::Vector{Vector{Int64}},edge2nodes::Vector{Vector{Int64}},
    w::Vector{Float64},d::Vector{Float64},elen::Vector{Int64},
    β::Vector{Float64},γ::Vector{Float64},kmax::Int64,
    randflag::Bool=false,maxits::Int64=100,verbose::Bool=true,Zwarm::Vector{Int64}=Vector{Int64}(),clusterpenalty=0)
    """
    Basic step Louvain algorithm: iterate through nodes and greedily move
    nodes to adjacent clusters. Does not form supernodes and does not recurse.

    Should work also for degenerate hyperedges, though some bugs are possible

    Features:
    * has weight for size of hyperedge rather than computing the weights
    * allows you to warm start with a clustering
    * allows you to randomize node order
    """
    if verbose println("One step of all-or-nothing HyperLouvain") end
    m = length(edge2nodes)
    n = length(node2edges)

    # Randomize node visit order
    if randflag
        node_order = Random.randperm(n)
        # NOTE: may be faster to re-arrange
        #   all the data structures so that we march through
        #   Z in order.
    else
        node_order = 1:n
    end

    if verbose println("") end

    # Warm start clustering
    if length(Zwarm) > 0
        Z = renumber(Zwarm)
    else
        # Default is to put all nodes in their own cluster
        Z = collect(1:n)
    end

	K = maximum(Z)   			# number of clusters
    ClusVol = zeros(K)			# volume of clusters
	ClusSize = zeros(K)			# size of each cluster
	for i = 1:n
		ClusVol[Z[i]] += d[i]
		ClusSize[Z[i]] += 1
	end

    # Store cut penalties for each edge
    cutpenalty = zeros(m)
    for e = 1:m
        k = elen[e]      				# size of the edge
        cutpenalty[e] = β[k]*w[e] 	# penalty for cutting it
    end

    # Store node neighbors of each node
    Neighbs = NeighborList(node2edges, edge2nodes)

    # Store all the edges for the node, minus that node id
    # This makes computing cut changes faster
    edgelists = Vector{Vector{Vector{Int64}}}()
    for i = 1:n
        iedges = Vector{Vector{Int64}}()    # each node has a list of hyperedges it's in, i.e. a list of lists of nodes.
        for e in node2edges[i]
            edge = edge2nodes[e]
            push!(iedges, setdiff(edge,i))  # don't save i itself
        end
        push!(edgelists,iedges)
    end

    improving = true
    iter = 0
    changemade = false

    # Initialize vols and cuts to track efficient local changes
    r = kmax
    toler = 1e-8

    mainstart = time()
    while improving && iter < maxits

        iter += 1
        if mod(iter,1) == 0
            if verbose println("Louvain Iteration $iter") end
        end
        improving = false

        # visit each node in turn
        for i = node_order

            # Cluster index for node i
            Ci_ind = Z[i]

            # Get the size of the cluster
			clus_size = ClusSize[Ci_ind]

            # Get the indices of i's neighbors--these define clusters we might move to
            Ni = Neighbs[i]

            # Get the neighboring clusters of i
            NC = unique(Z[Ni])

            # The default is to not move the node
            BestC = Z[i]
            BestImprove = 0

            # Set of edges that i is in
            Cv = node2edges[i]

            # Set of hyperedges that node i is in, not including i itself
            Cv_list = edgelists[i]

            # Volume of the set currently
            vS = ClusVol[Ci_ind]
            dv = d[i]

            for j = 1:length(NC)

				# Check how much it would improve to move i to to cluster j
                Cj_ind = NC[j]

                if Cj_ind == Ci_ind
                    continue
                end

                # Change in volume
                vJ = ClusVol[Cj_ind]
                Δvol = 0
                for k = 1:kmax
					# Better if this is smaller
                    Δvol += γ[k]*((vS-dv)^k + (vJ+dv)^k - vS^k - vJ^k)
                end

                # Change in cut
                Δcut = 0
                for eid = 1:length(Cv)
                    # change in cut for edge e when v moves from
                    # cluster Ci to cluster Cj
                    e = Cv[eid]
                    edge_noi = Cv_list[eid]

                    if elen[e] > 1
						we = cutpenalty[e]
                        Δcut += move_cut(i,Z,edge_noi,Ci_ind,Cj_ind,we)
                    end
                end

                # Change in objective due to change in cluster number
                Δclus = 0
                if clusterpenalty > 0 && clus_size == 1
                    Δclus = clusterpenalty*(log(K-1)-log(K))
                end
                change = Δcut + Δvol + Δclus # want this to be negative

                # Check if this is currently the best possible greedy move to make
                # The tolerance helps with making some behavior more stable:
                #   you only move if there's a numerically nonzero reason to move.
                if change < BestImprove - toler
                    BestImprove = change
                    BestC = Cj_ind
                    improving = true
                end
            end

            # Move i to the best new cluster, only if it strictly improves modularity
            if BestImprove < -toler

                # update clustering
                ci_old = Z[i]
                Z[i] = BestC

                # Remove i from its old cluster...
                ClusVol[ci_old] -= dv
				ClusSize[ci_old] -= 1

                # ...and add it to its new cluster
                ClusVol[BestC] += dv
				ClusSize[BestC] += 1
                changemade = true
                improving = true # we have a reason to keep iterating!

                if clus_size == 1
                    # we moved a singleton to another cluster
                    K -= 1
                end

            end
        end
    end
    mainloop = time()-mainstart
    if ~changemade
        improved = false
        if verbose println("No nodes moved clusters") end
    else
        improved = true
    end
    if verbose println("Main loop took $mainloop seconds") end
	Z = renumber(Z)
    return Z, improved

end

"""
Compute the change from moving node v from Ci to Cj
"""
function move_cut(v::Int64,Z::Vector{Int64},erest::Vector{Int64},Ci_ind::Int64,
    Cj_ind::Int64,we::Float64)

    # p = Z[erest]            # set of clusters in e, not counting e
    p1 = Z[erest[1]]
    na = notsame(erest,Z,p1)

    if na
        # if the edge is cut regardless of whether v is moved, return 0
        return 0
    else
        # p1 = p[1]   # this the the cluster all other nodes are in
        if Ci_ind == p1
            # then moving v cuts the edge
            return we
        elseif Cj_ind == p1
            # then moving v to Cj_ind "uncuts" the edge
            return -we
        else
            # moving doesn't matter
            return 0
        end
    end
end

"""
e is the set of nodes in the hyperedge. Z stores their clusters
"""
function notsame(e::Vector{Int64},Z::Vector{Int64},p1::Int64)
    for i = 2:length(e)
        if Z[e[i]] != p1
            return true
        end
    end
    return false
end

function Condense_H(edge2nodes::Vector{Vector{Int64}},w::Vector{Float64})
    """
    Condense Hypergraph.

        A hypergraph stored as an edgelist may contain some edges
        that are identical. This function finds them and condenses
        them instead into a weighted hypergraph edgelist
    """
    Hdict = Dict{Vector{Int64},Float64}()   # maps from set of nodes to a weight

    for e = 1:length(edge2nodes)
        edge = edge2nodes[e]
        sort!(edge)
        we = w[e]
        wt = get(Hdict,edge,0)
        Hdict[edge] = wt+we
    end

    newlist = Vector{Vector{Int64}}()
    for edge in keys(Hdict)
        push!(newlist,edge)
        push!(wnew,Hdict[edge])
    end

    return newlist, wnew
end

function Condense_H(edge2nodes::Vector{Vector{Int64}},w::Vector{Float64},label::Vector{Int64})
    """
    Condense Hypergraph.

        A hypergraph stored as an edgelist may contain some edges
        that are identical. This function finds them and condenses
        them instead into a weighted hypergraph edgelist.

        In this version, each hyperedge is additionally associated with a
        label, and we only merge nodes if they have the same node set
        AND the same label.

        The edge label is mainly meant to store original hyperedge size.
        Since some of the nodes may be supernodes, we need to keep track of
        the original hyperedge size.

        Also, get rid of size 1 hyperedges in the process
    """
    Hdict = Dict{Vector{Int64},Float64}()   # maps from set of nodes to a weight

    @assert(length(label)==length(edge2nodes))

    for e = 1:length(edge2nodes)
        edge = edge2nodes[e]
        lab = label[e]
        sort!(edge)
        ky = [lab; edge]    # the key is the label + edge
        we = w[e]
        wt = get(Hdict,ky,0)
        Hdict[ky] = wt+we
    end

    newlist = Vector{Vector{Int64}}()
    wnew = Vector{Float64}()
    labelsnew = Vector{Int64}() # indices of edges you keep
    for ky in keys(Hdict)
        edge = ky[2:end]
        lab = ky[1]
        if length(edge) > 1
            push!(newlist,edge)
            push!(wnew,Hdict[ky])
            push!(labelsnew,lab)
        end
    end

    return newlist, wnew, labelsnew
end

function Change_id_H!(edge2nodes::Vector{Vector{Int64}},Z::Vector{Int64},uniqueflag::Bool=true)
"""
Given a clustering vector Z, change node IDs in the hypergraph
so that node i becomes node Z[i]. This will potentially lead
to many degenerate hyperedges where the same node shows up multiple times.

This is useful as a super-node step in clustering heuristics like louvain.

If uniqueflag == true, don't store degenerate hyperedges. Instead, just save
one copy of the vector
"""

    for e = 1:length(edge2nodes)
        for i = 1:length(edge2nodes[e])
            edge2nodes[e][i] = Z[edge2nodes[e][i]]
        end
        if uniqueflag
            unique!(edge2nodes[e])
        end
    end

    return edge2nodes
end

function condense_d(d::Vector{Float64},Z::Vector{Int64})
    """
    Condense a degree vector to a weights vector by summing up volumes of
    clusters.
    """

    n = maximum(Z)
    dnew = zeros(n)
    for i = 1:n
        S = findall(x->x==i,Z)
        dnew[i] = sum(d[S])
    end
    return dnew
end

function Elist_to_Hypergraph(elist::Vector{Vector{Int64}}, maxsize::Int64=25)
    """
    Converts a binary hypergraph incidence matrix to type "hypergraph".
    Assumes edges are properly sorted already
    """
    E = Dict{Integer, Dict}()
    for edge in elist
        # sort!(edge)
        if length(edge) > maxsize; continue; end
        sz = length(edge)
        if !haskey(E, sz)
            E[sz] = Dict{}()
        end
        E[sz][edge] = 1
    end

    D = zeros(Int64, n)
    for (sz, edges) in E
        for (e, _) in edges
            D[e] .+= 1
        end
    end
    N = 1:n
    return hypergraph(N, E, D)
end

function Hmat_to_Hypergraph(H::SparseArrays.SparseMatrixCSC, maxsize::Int64=25)
    """
    Converts a binary hypergraph incidence matrix to type "hypergraph"
    """
    E = Dict{Integer, Dict}()
    elist = incidence2elist(H)
    for edge in elist
        sort!(edge)
        if length(edge) > maxsize; continue; end
        sz = length(edge)
        if !haskey(E, sz)
            E[sz] = Dict{}()
        end
        E[sz][edge] = 1
    end

    D = zeros(Int64, n)
    for (sz, edges) in E
        for (e, _) in edges
            D[e] .+= 1
        end
    end
    N = 1:n
    return hypergraph(N, E, D)
end

function AON_Inputs(H,ω,α,kmax)
    """
    Given a hyperedge H, ω function, and α which parameterizes ω,
    return all of the data structures you need to run the
    all-or-nothing louvain algorithm
    """
    β = zeros(kmax)
    γ = zeros(kmax)
    for k = 2:kmax
        β[k] = log(ω([1,k],α))-log(ω([0,k],α))
        γ[k] = ω([1,k],α)-ω([0,k],α)
    end

    He2n, edge_weights = hypergraph2incidence(H)
    e2n = incidence2elist(He2n);
    n2e = incidence2elist(SparseArrays.sparse(He2n'))
    m = length(e2n)
    edge_len = zeros(Int64,m)
    for e = 1:m
        edge_len[e] = length(e2n[e])
    end
    # deg = vec(sum(He2n,dims = 1)) # only works if hypergraph is unweighted
	deg = 1.0*H.D # coercion to float for compatibility with AON_louvain
    # @show H.D-deg
    # @assert(deg == H.D)

    return β, γ, e2n, n2e, edge_weights,deg,edge_len

end

function learn_omega_aon(e2n,weights,Z,kmax,d,n)
    """"
    Given a fixed clustering, this computes the maximum likelihood estimate
    for the hyperedge intensity parameters omega specifically for the
    all-or-nothing modularity objective.

	These esimates are obtained by directly deriving the AON modularity
	with respect to omega_{1,k} and omega_{0,k}, probabilities for turning size-k
	tuples into hypergraphs, if they are completely contained inside a cluster
	or not.

    Assumes kmin is 2.
    e2n is the edge-to-node list
    d is the degree vector
    n is number of nodes
    Z is the clustering
    """

    L = maximum(Z)
    m = length(e2n)
    ClusVol = zeros(L)			# volume of clusters
    for i = 1:n
        ClusVol[Z[i]] += d[i]
    end

	# top row is weight of edges; bottom row is weight of cut edges
    EdgesAndCuts = zeros(2,kmax)

	# initialize to small positive number as a heuristic, to avoid taking log of 0
    EdgesAndCuts[1,:] = 0.01*ones(1,kmax)

    # Compute the cut penalty for each hyperedge size
    for j = 1:m
        edge = e2n[j]
		we = weights[j]
        k = length(edge)
        EdgesAndCuts[1,k] += we     # just counting number of hyperedges
        iscut = notsame(edge,Z,Z[edge[1]])
        if iscut
            EdgesAndCuts[2,k] += we
        end
    end

    # top row is in probability, bottom is out probability
    omega = zeros(2,kmax)
    volV = sum(ClusVol)

    for k = 2:kmax
        # get the volume piece
        volsums = 0
        for l = 1:L
            volsums += ClusVol[l]^k
        end
        volVk = volV^k

        omega[1,k] = (EdgesAndCuts[1,k] - EdgesAndCuts[2,k])/volsums
        omega[2,k] = EdgesAndCuts[2,k]/(volVk - volsums)
    end

    β = zeros(kmax)
    γ = zeros(kmax)
    for k = 2:kmax
        β[k] = log(omega[1,k])-log(omega[2,k])
		# if β[k] == Inf
		# 	# In cases where β[k] = Inf,
		# 	# we replace with a high but finite value instead
		# 	β[k] = 1e3
		# end
        γ[k] = omega[1,k]-omega[2,k]
    end
    return β, γ,omega
end

function learn_omega_aon(H::hypergraph,Z::Vector{Int64})
	"""
	Same function as above, but takes input type Hypergraph.
	"""
	d = H.D
	n = length(d)
	kmax = maximum(keys(H.E))
	e2n,weights = hyperedge_formatting(H)
	return learn_omega_aon(e2n,weights,Z,kmax,d,n)
end

function Kaminski_default(H::hypergraph)
	kmax = maximum(keys(H.E))
	EdgeList,weights = hyperedge_formatting(H)
	m = length(EdgeList)
	Mvec = zeros(kmax)
	for i = 1:m
		edge = EdgeList[i]
		we = weights[i]
		k = length(edge)
		Mvec[k] += we
	end
	volH = sum(H.D)
	beta = ones(kmax)
	beta[1] = 0
	gamma = zeros(kmax)
	for k = 2:kmax
		gamma[k] = Mvec[k]/(volH)^k
	end
	return beta, gamma
end

function omega_to_betagamma(omega)
	"""
	Go from AON intensity parameters to parameters omega and gamma.
	Page 7 in text.
	"""
	kmax = size(omega,2)
	β = zeros(kmax)
	γ = zeros(kmax)
	for k = 1:kmax
		β[k] = log(omega[1,k]/omega[2,k])
		γ[k] = (omega[1,k]-omega[2,k])/β[k]
	end
	return β, γ
end

function betagamma_to_omega(β, γ)
	"""
	Go from parameters omega and gamma to AON intensity parameters.
	Page 7 in text.
	"""
	kmax = length(β)
	omega = zeros(2,kmax)
	for k = 1:kmax
		omega[2,k] = β[k]*γ[k]/(exp(β[k])-1)
		omega[1,k,] = exp(β[k])*omega[2,k]
	end

	return omega
end

function allsame(v::Vector{Float64})
	"""
	Checks whether all entries of a vector are the same.
	"""
    v1 = v[1]
    for j = 2:length(v)
        if v[j] != v1
            return false
        end
    end
    return true
end

function allsame(v::Vector{Int64})
	"""
	Checks whether all entries of a vector are the same.
	"""
    v1 = v[1]
    for j = 2:length(v)
        if v[j] != v1
            return false
        end
    end
    return true
end

function modularity_aon(H::hypergraph,Z::Vector{Int64},omega::Array{Float64,2};likelihood=false)
	"""
	Computes the all or nothing modularity function.
		Equation (4), function Q in manuscript.
	"""
	β, γ = omega_to_betagamma(omega)
	n = length(Z)
	kmax = maximum(keys(H.E))
	kmin = maximum(keys(H.E))
	if kmin == 1
		error("Function requires hyperedges to have at least size 2.")
	end

	# m_k = Mvec[k] = weighted number of hyperedges of size k
	mvec = zeros(kmax)
	for k = kmin:kmax
        El = H.E[k]
        for edge in keys(El)
            mvec[k] = El[edge]
        end
    end

	L = maximum(Z)
    ClusVol = zeros(L)			# volume of clusters
    for i = 1:n
        ClusVol[Z[i]] += H.D[i]
    end

	obj = 0
	for k = kmin:kmax
		# Compute cut penalty of the objective
		El = H.E[k]
		for edge in keys(El)
			if ~allsame(Z[edge])
				obj -= El[edge]*β[k]
			end
		end

		# Volume penalty of the objective
		for l = 1:L
			obj -= β[k]*γ[k]*ClusVol[l]^k
		end
	end

	# Additional term that is constant with respect to clustervec
	for k = kmin:kmax
		obj += β[k]*mvec[k]
		El = H.E[k]
		for edge in keys(El)
			obj += El[edge]*log(omega[2,k])
			p = partitionize(Z[edge])
			bR = Combinatorics.multinomial(p...)
			obj -= prod(H.D[edge])*omega[2,k]*bR
		end
	end

	if likelihood == true
		# extra part for log-likelihood computation
		loglikhood = obj
		for k = kmin:kmax
			El = H.E[k]
			for edge in keys(El)
				aR = El[edge]
				loglikhood += aR*log(prod(H.D[edge]))
				p = partitionize(Z[edge])
				bR = Combinatorics.multinomial(p...)
				loglikhood -=  aR*log(bR) - log(factorial(aR))
			end
		end
		return obj, loglikhood
	else
		return obj
	end
end


include("include.jl")
include("read_lp.jl")

####################################
### PLOT A TOY PROBLEM
####################################
using JuMP, Ipopt, PyPlot
m = Model()

A = [[1.0 1.0 0.0]; [1.0 0.0 -1.0]]
b, c, lbounds, ubounds = [0.0 0.0]', [0.0 0.0 0.0]', [0.0 0.0 0.0]', [1e30 1e30 1e30]'

ls = Dict(1e-3 => "-", 1e-6 => "-.", 1e-9 => "--", 0.0 => ":")
delta_set = [1e-3, 1e-6, 1e-9, 0.0]

ylims = [1e-2, 1e10]

subplot(121)
for delta = delta_set
  max_iter = 100
  hist = run_IPOPT(A, b, c, lbounds, ubounds, max_iter, delta)
  dual_hist = compute_dual_hist(hist)

  semilogy(1:length(dual_hist), dual_hist, color="black", linestyle=ls[delta], label="δ = $delta", basey=10)
end
title("IPOPT")
ax = gca()
ax[:set_ylim](ylims)
xlabel("iteration")
ylabel("maximum dual variable")

legend()

function run_one_phase(jump_model)
    nlp_raw = MathProgNLPModel(jump_model)
    nlp = Class_CUTEst(nlp_raw)
    timer = class_advanced_timer()
    start_advanced_timer(timer)
    start_advanced_timer(timer, "INIT")
    intial_it = init(nlp, my_par, timer)
    pause_advanced_timer(timer, "INIT")
    @assert(is_feasible(intial_it, my_par.ls.comp_feas))
    iter, status, hist, t, err = one_phase_IPM(intial_it, my_par, timer);
    pause_advanced_timer(timer)

    return hist
end

subplot(122)
for delta = delta_set
  jump_model = read_lp_into_JuMP(A, b, c, lbounds, ubounds, delta, true)
  hist = run_one_phase(jump_model)
  dual_hist = compute_dual_hist(hist)

  semilogy(1:length(dual_hist), dual_hist, color="black", linestyle=ls[delta], label="δ = $delta", basey=10)
end
title("Well behaved IPM")
xlabel("iteration")
ax = gca()
ax[:set_ylim](ylims)

#ax[:legend]()[:draggable]()


using LazyArrays, Test
using LazyArrays.SetOperations
import LazyArrays.SetOperations: ∪, ∩, unioned, intersected, setdiffed

@testset "Set operations" begin
    @test SetOperations.Style(Int) == SetOperations.VectorStyle()
    @test SetOperations.Style(Set) == SetOperations.SetStyle()
    @test 1 ∈ unioned(1,2.0)
    @test 2 ∈ unioned(1,2.0)
    @test 3 ∉ unioned(1,2.0)
    @test issubset([1,2],  unioned(1,2.0))

    @test 1 ∪ 2.0 == [1.0,2.0]
    @test [1,2] ∪ Set([2,4]) ==  Set([2,4]) ∪ [1,2] == SetOperations.union(Set([2,4]) , [1,2]) ==
        SetOperations.union([1,2], Set([2,4])) == Set([1,2,4])
    @test Base.union([1,2], Set([2,4])) == Base.:∪([1,2], Set([2,4])) ==
                [1,2,4]

    @test 1 ∈ intersected([1,2], Set([1.0]))
    @test 2 ∉ intersected([1,2], Set([1.0]))
    @test 1 ∉ setdiffed([1,2], Set([1.0]))
    @test 2 ∈ setdiffed([1,2], Set([1.0]))

    @test 1 ∈ intersected([1,2], Set([1.0]), 1f0)
end


# from McCollum et al,
# "The sorting direct method for stochastic simulation of biochemical systems with varying reaction execution behavior"
# Comp. Bio. and Chem., 30, pg. 39-49 (2006).

using DiffEqBase, DiffEqJump, DiffEqBiological
using Base.Test

using Plots; plotlyjs()
doplot = false
using BenchmarkTools
dobenchmark = true

dotestmean   = false
doprintmeans = false

# SSAs to test
#SSAalgs = (SortingDirect(),Direct()) #, DirectFW(), FRM(), FRMFW())
SSAalgs = (Direct(), SortingDirect(),)

# Nsims        = 80000
# tf           = 100000.0
Nsims        = 100
tf           = 500.0
u0           = [1,0,0,0]
expected_avg = 5.926553750000000e+02
reltol       = .01

# average number of proteins in a simulation
function runSSAs(jump_prob)
    Psamp = zeros(Int, Nsims)
    for i in 1:Nsims
        sol = solve(jump_prob, SSAStepper())
        Psamp[i] = sol[3,end]
    end
    mean(Psamp)
end

# MODEL SETUP
N = 10
genenetwork = "@reaction_network gtype begin\n"
for i in 1:N
    genenetwork *= "\t 10.0, G$(2*i-1) --> G$(2*i-1) + M$(2*i-1)\n"
    genenetwork *= "\t 10.0, M$(2*i-1) --> M$(2*i-1) + P$(2*i-1)\n"
    genenetwork *= "\t 1.0,  M$(2*i-1) --> 0\n"
    genenetwork *= "\t 1.0,  P$(2*i-1) --> 0\n"

    genenetwork *= "\t 5.0, G$(2*i) --> G$(2*i) + M$(2*i)\n"
    genenetwork *= "\t 5.0, M$(2*i) --> M$(2*i) + P$(2*i)\n"
    genenetwork *= "\t 1.0,  M$(2*i) --> 0\n"
    genenetwork *= "\t 1.0,  P$(2*i) --> 0\n"

    genenetwork *= "\t 0.0001, G$(2*i) + P$(2*i-1) --> G$(2*i)_ind \n"
    genenetwork *= "\t 100., G$(2*i)_ind --> G$(2*i)_ind + M$(2*i)\n"
end
genenetwork *= "end"
rs = eval( parse(genenetwork) )
u0 = zeros(Int, length(rs.syms))
for i = 1:(2*N)
    u0[findfirst(rs.syms, Symbol("G$(i)"))] = 1
end

# TESTING:
prob = DiscreteProblem(u0, (0.0, tf))

# plotting one full trajectory
if doplot
    plothand = plot(reuse=false)
    for alg in SSAalgs
        jump_prob = JumpProblem(prob, alg, rs, save_positions=(false,false))
        sol = solve(jump_prob, SSAStepper(), saveat=(tf/1000.))
        plot!(plothand, sol.t, sol[[3,6],:]', lab=string(alg))
    end
    display(plothand)
end

# # test the means
# if dotestmean
#     means = zeros(Float64,length(SSAalgs))
#     for (i,alg) in enumerate(SSAalgs)
#         jump_prob = JumpProblem(prob, alg, majumps, save_positions=(false,false))
#         means[i]  = runSSAs(jump_prob)
#         relerr = abs(means[i] - expected_avg) / expected_avg
#         if doprintmeans
#             println("Mean from method: ", typeof(alg), " is = ", means[i], ", rel err = ", relerr)
#         end

#         if dobenchmark
#             @btime (runSSAs($jump_prob);)
#         end

#         @test abs(means[i] - expected_avg) < reltol*expected_avg
#     end
# end

# # benchmark performance
if dobenchmark
    # exact methods
    for alg in SSAalgs
        println("Solving with method: ", typeof(alg), ", using SSAStepper")
        jump_prob = JumpProblem(prob, alg, rs, save_positions=(false,false))
        @btime solve($jump_prob, SSAStepper(), saveat=(tf/1000.))
        #@btime (runSSAs($jump_prob);)
    end
    println()
end


# ==============================================================================
# Parseur lisant une instance de probleme de partitionnement d'ensembles (SPA) bi-objectif

function loadInstance2SPA(fname::String)
    f = open(fname)
    nbcontraintes, nbvar = parse.(Int, split(readline(f))) # nombre de contraintes , nombre de variables
    L = zeros(Int, nbcontraintes, nbvar)    # matrice des contraintes
    c1 = zeros(Int, nbvar)                  # vecteur des couts
    c2 = zeros(Int, nbvar)                  # deuxième vecteur des couts
    nb = zeros(Int, nbvar)
    for i in 1:nbvar
        flag = 1
        for valeur in split(readline(f))
            if flag == 1
                c1[i] = parse(Int, valeur)
                flag +=1
            elseif flag == 2
                c2[i] = parse(Int, valeur)
                flag +=1
            elseif flag == 3
                nb[i] = parse(Int, valeur)
                flag +=1
            else
                j = parse(Int, valeur)
                L[j,i] = 1
            end
        end
    end
    close(f)
    return nbvar, nbcontraintes, L, c1, c2
end


# TODO: quasi-exact line search via interpolation, GPs
# IDEA: g .*= -g'g / (g'A*g) # optimal stepsize for quadratic problems

################################################################################
# inputs: searchcondition function, i.e. Armijo, Wolfe
# stepupdate is step size updating policy
function linesearch!(xn::AbstractVector, x::AbstractVector, α::Real,
                     direction::AbstractVector, searchcondition!, maxbacktrack::Int, decrease::Real = 3.)
    for i in 1:maxbacktrack
        @. xn = x + α * direction # update
        if searchcondition!(α, xn)
            return true
        end
        α /= decrease
    end
    return false
end
# linesearch for scalar parameters
function linesearch(x::Real, α::Real, direction::Real, searchcondition,
                    maxbacktrack::Int, decrease::Real = 3.)
    for i in 1:maxbacktrack
        xn = x + α * direction # update
        if searchcondition(α, xn)
            return xn
        end
        α /= decrease
    end
    return x
end
########################### linesearch types ###################################
struct DecreasingStep{T, D, A<:Real, S<:Real} <: Update{T}
    direction::D # αξία και κατεύθυνση
    x::T # storage for value
    α::A # stepsize
    c::S # coefficient in decreasing rule
    decrease::S # μείωση factor
    # increase::S # αύξηση factor
    maxbacktrack::Int
    function DecreasingStep(direction::Direction, x, α::Real = 1., c::Real = 1.0,
                            decrease::Real = 3.0, maxbacktrack::Int = 16) # increase::Real = 2,
        T, D, A, S = typeof(x), typeof(direction), typeof(α), promote_type(typeof(c), typeof(decrease)) #, typeof(increase))
        new{T, D, A, S}(direction, copy(x), α, c, decrease, maxbacktrack) # makes sure that we use a copy of x
    end
end

objective(D::DecreasingStep) = objective(D.direction)
function update!(D::DecreasingStep, x::AbstractVector)
    val, dir = valdir(D.direction, x)
    descent = DescentCondition(objective(D), val, x, D.c)
    success = linesearch!(D.x, x, D.α, dir, descent, D.maxbacktrack, D.decrease)
    return success ? copyto!(x, D.x) : x
end
# update for scalar parameter
function update!(D::DecreasingStep, x::Real)
    val, dir = valdir(D.direction, x)
    descent = DescentCondition(objective(D), val, x, D.c)
    xn = linesearch(x, D.α, dir, descent, D.maxbacktrack, D.decrease)
    success = xn != x
    return success ? xn : x
end

# combines descent condition and local minimum w.r.t. step size α condition
mutable struct DescentCondition{T, F, X}
    f::F # objective function
    value::T # value f(x)
    new_value::T # value at most recently proposed point f(xn)
    xn::X
    c::T # descent factor
end
function DescentCondition(objective::Function, value, x, c::Real = 1.)
    DescentCondition(objective, value, typeof(value)(Inf), zero(x), c)
end
function (D::DescentCondition)(α, xn)
    decreasing = D.new_value ≤ D.c * D.value # was the last proposal decreasing
    new_value = D.f(xn)
    alpha_minimum = D.new_value ≤ new_value # was the last proposal a stationary point w.r.t. α?
    if decreasing && alpha_minimum
        if xn isa Real
            xn = D.xn
        else
            @. xn = D.xn # assign last proposal to current one
        end
        return true
    else
        if xn isa Real
            D.xn = xn
        else
            @. D.xn = xn
        end
        D.new_value = new_value
        return false
    end
end

# weakest rule, no good theoretical guarantees, but works in practice and is fast
function descent_condition(objective::Function, value, c::Real = 1.)
    return descent(α, xn) = (objective(xn) ≤ c * value)
end

################################################################################
# TODO: what if this instead had a direction field?
struct ArmijoStep{T, D, G, A, S} <: Update{T}
    direction::D # αξία και κατεύθυνση
    x::T # storage for value
    gradient::G # need gradient for Armijo rule
    α::A # last α, has to be 0D array, or armijo is mutable (went with mutable)
    c::S # coefficient in Armijo rule
    decrease::S # μείωση factor
    maxbacktrack::Int
    function ArmijoStep(dir::D, x::T, grad::G = x->FD.gradient(objective(dir), x),
            α::S = 1., c = 0., decrease = 3., maxbacktrack::Int = 16) where {T, D, G, S}
            new{T, D, G, S, S}(dir, copy(x), grad, α, S(c), S(decrease), maxbacktrack)
    end
end
# function ArmijoStep(x, obj, valdir,
#         grad = x->FD.gradient(obj, x), α::S = 1., c = 0., decrease = 3.,
#         increase = 2., maxbacktrack::Int = 16) where {S}
#     # α = fill(α) # creating 0D array
#     x = fill(zero(eltype(x)), size(x))
#     ArmijoStep(copy(x), obj, valdir, grad, α,
#                 S(c), S(decrease), S(increase), maxbacktrack)
# end
objective(A::ArmijoStep) = objective(A.direction)
function update!(A::ArmijoStep, x)
    val, dir = valdir(A.direction, x)
    grad = A.gradient(x)
    armijo = armijo_condition(objective(A), val, dir, grad, A.c)
    success = linesearch!(A.x, x, A.α, dir, armijo, A.maxbacktrack, A.decrease)
    return success ? copyto!(x, A.x) : x
end

# returns a function which calculates the armijo condition
# if direction = gradient, pass it to the function to avoid duplicate calculation
# value = objective(x)
function armijo_condition(objective::Function, value, direction::T, ∇::T,
                                                    c::Real = 1e-4) where {T}
    direction_gradient = dot(direction, ∇) # just to make sure its eager
    return armijo(α, xn) = (objective(xn) ≤ value + α * c * direction_gradient)
end

################################################################################
# could also abstract away step update:
# backtrack(α, decrease) = α * decrease
# backtrack(decrease) = α -> backtrackstep(α, decrease)
# integral control?
# AKA InexactLinesearch, could have separate type for exact
# struct LineSearch{T, D<:Direction, S, C} <: Update{T}
#     direction::D
#     x::T # intermediate storage for value
#     # stepupdate::S
#     searchcondition::C
#     decrease::S
#     maxbacktrack::Int
#     function LineSearch(dir::D, x::T, α::S = 1., c = 1., decrease = 3.,
#                     maxbacktrack::Int = 16) where {T, D<:Direction, S}
#         # makes sure that we use a copy of x
#         new{T, D, S, C}(dir, copy(x), S(decrease), maxbacktrack)
#     end
# end
# function update!(L::LineSearch, x)
#     val, dir = valdir(L.direction, x)
#
#     d = L.direction(x)
#     # cond = L.searchcondition(x, t)
#     descent = L.search_condition(objective(D), val, dir, D.c)
#     success = linesearch!(L.x, x, dir, L.searchcondition, L.maxbacktrack, L.decrease) #L.stepupdate)
#     return success ? copy!(x, D.x) : x
# end
# needs to compute function which returns searchcondition for each x
# function armijo(d::Direction)
#     value, direction = A.valdir(x, t...)
#     grad = A.gradient(x)
#     cond = armijo_condition(A.objective, value, direction, grad, A.c)
#     success = linesearch!(A.x, x, A.α, direction, armijo, A.maxbacktrack, A.decrease)
#     return success ? copy!(x, A.x) : x
# end

################################################################################
# strongest condition, theoretical guarantees
function wolfe_condition(x::T, objective::Function, direction::T, ∇::T,
                            c₁::Real = 1e-4, c₂::Real = .99) where {T}
    fx = objective(x)
    direction_gradient = dot(direction, ∇) # just to make sure its eager
    function wolfe(α, xn, ∇n = gradient(f, xn))
        armijo = (objective(xn) ≤ fx + α * c₁ * direction_gradient)
        curv = abs(dot(direction, ∇n)) ≤ abs(c₂*direction_gradient) # weak rule: -d'∇(x + α*d) ≤ -c*d'∇(x)
        return armijo && curv
    end
    return wolfe
end

# TODO: Could rename Damped, so type reads e.g. Damped{Newton}
# 0 < c₁ < c₂ < 1
# mutable struct WolfeStep{T, D<:Direction{T}} <: Update{T}
#     d::D
#     c₁::T
#     c₂::T
#     α::T
#     decrease::T # μείωση factor
#     increase::T # αύξηση factor
#     maxbacktrack::Int
# end
# function WolfeStep(d::D, c::T = .99, α = 1., decrease = 7., increase = 2.,
#                             maxbacktrack::Int = 16) where {T, D<:Direction}
#     ArmijoStep{T, D}(d, c, T(α), T(decrease), T(increase), maxbacktrack)
# end

############################ Step Size Selectors ###############################
# generalizes search for an α such that (x - α * direction) leads to minimizer
# struct ConstantStep{T, D<:Direction} <: Direction{T}
#     d::D
#     α::T
# end
# # decide if I want to represent the multiplication lazily
# # direction(C::ConstantStep) = C.α * direction(C.d) # introduces temporary
# function direction!(C::ConstantStep, x)
#     direction!(C.d, x) # update direction vector
#     direction(C.d) .*= C.α # multiply direction with step-size
# end
# function update!(C::ConstantStep, x::AbstractVector)
#     direction!(C.d, x) # update direction vector
#     x .+= C.α .* direction(C.d)
# end

################################################################################
# struct TimeDependentStep{T, D<:Direction} <: Direction{T}
#     d::D
#     α::T
# end
# function update!(C::TimeDependentStep, x::AbstractVector, t::Int)
#     direction!(C.d, x)
#     C.d .*= C.α(t)
#     x .+= direction(C.d)
# end

################################################################################
# two most important rules to consider when selecting step sizes
# function armijo_rule(x, f::Function, d, c, α, ∇ = gradient(f, x))
#     f(x + α*d) ≤ f(x) + α * c * (d'∇) # RHS could be computed independently of α
# end

# direction_gradient is dot product of the two
# function armijo_rule(xn, value::Real, f, direction_gradient, c::Real, α::Real)
#     f(xn) ≤ value + α * c * direction_gradient
# end
# implements the strong wolfe
# function curvature_rule(x, f, d, c, α)
#     abs(d'gradient(f, x + α*d)) ≤ abs(c*d'gradient(f, x)) # weak rule: -d'∇(x + α*d) ≤ -c*d'∇(x)
# end
# function curvature_rule(xn, f, dg, c, α)
#     abs(d'gradient(f, xn)) ≤ abs(c*dg) # weak rule: -d'∇(x + α*d) ≤ -c*d'∇(x)
# end

################################## WIP #########################################
# the minimization is expensive as we have to reevaluate the gradient of the objective
# better with interpolation?
# step size calculated with approximate minimization via secant method
# struct SecantStep{T, D<:Direction{T}} <: Update{T}
#     direction::D
# end
#
# objective(D::SecantStep) = objective(D.direction)
# function update!(D::SecantStep, x)
#     val, dir = valdir(D.direction, x)
#     descent = descent_condition(objective(D), val, dir, D.c)
#     success = linesearch!(D.x, x, D.α, dir, descent, D.maxbacktrack, D.decrease)
#     # α, success = linesearch!(D.x, x, D.α, dir, descent, D.maxbacktrack, D.decrease)
#     # D.α = α * D.increase
#     function value(α)
#         @. xn = x + α * direction # update
#         value(D, xn)
#     end
#     function value_direction(α)
#         @. xn = x + α * direction # update
#         val, dir = value_direction(D, xn)
#         val, dot(dir, direction)
#     end
#     return success ? copy!(x, D.x) : x
# end


module Farey

export farey

farey(x) = sort(unique(n//d for d=1:x for n=0:d))

end # module


# supertype for all logit model results

using StatsBase

abstract type LogitModel <: RegressionModel end

StatsBase.coefnames(r::LogitModel) = r.coefnames
StatsBase.coef(r::LogitModel) = r.coefs
StatsBase.islinear(r::LogitModel) = false
StatsBase.loglikelihood(r::LogitModel) = r.final_ll
StatsBase.vcov(r::LogitModel) = r.vcov
StatsBase.stderror(r::LogitModel) = sqrt.(diag(StatsBase.vcov(r)))