# Command-Line 2048 Example
module CLI2048

using Termbox
using Compat
if VERSION >= v"0.7-"
    using Random
    using Printf
end

const Grid = Matrix

function Base.show(io::IO, g::Grid)
    n = size(g, 1)
    for row in 1:n
        for col in 1:(n-1)
            print(io, g[row, col], '\t')
        end
        print(io, g[row, n], '\n')
    end
end

newGrid(n::Int = 4, T = Int) = zeros(T, n, n)

function randomGrid(n::Int = 4, numNZ::Integer = 2, NonzeroElements::Vector{T} = [2,4]) where {T}
    randomGrid(newGrid(n, T), numNZ, NonzeroElements)
end
function randomGrid(grid::Grid{T}, numNZ::Integer = 2, NonzeroElements::Vector{T} = [2,4]) where {T}
    nGrid = length(grid)
    for ind in randperm(nGrid)[1:numNZ]
        grid[ind] = rand(NonzeroElements)
    end
    grid
end
zeroElements(g::Grid) = findall([gv == 0 for gv in g])
function randomGrid!(grid::Grid{T}, numNZ::Integer = 2, NonzeroElements::Vector{T} = [2,4]) where {T}
    inds = shuffle(zeroElements(grid))[1:numNZ]
    for ind in inds
        grid[ind] = rand(NonzeroElements)
    end
end

#TODO: Move asserts into some kind of test function
v1 = newGrid(4, Int)
@assert length(v1) == 4*4
@assert sum(v1) == 0
v1Rand = randomGrid(v1)
@assert size(v1Rand) == size(v1)
@assert sum(v1Rand) >= 2*2
v2 = randomGrid()
@assert sum(v2) >= 2*2
v3 = randomGrid(4, 3, [2,])
@assert sum(v3) == 3*2

# Functions to complete a single move in one direction
# NOTE: Output from below can be generated ahead of time and cached.
function compactLeft!(v::AbstractVector{T}, n) where {T}
    row = zeros(T, n)
    p = 1
    for (i, vi) in enumerate(v)
        if vi != 0
            row[p] = vi
            p += 1
        end
    end
    for i in eachindex(v)
        v[i] = row[i]
    end
end
function mergeLeft!(v::AbstractVector{T}, diffRef) where {T}
    n = length(v)
    compactLeft!(v, n)
    for i in 1:(n-1)
        vi = v[i]
        vip = v[i + 1]
        if vip == vi
            diffRef[] += vi
            v[i] *= 2
            v[i+1] = 0
        end
    end
    compactLeft!(v, n)
    v
end

function move!(g::Grid, dir::Val{:left}, scoreDiffRef)
    n = size(g, 1)
    for row in 1:n
        mergeLeft!(view(g, row, :), scoreDiffRef)
    end
end
function move!(g::Grid, dir::Val{:right}, scoreDiffRef)
    n = size(g, 1)
    for row in 1:n
        reverse!(mergeLeft!(reverse!(view(g, row, :)), scoreDiffRef))
    end
end
function move!(g::Grid, dir::Val{:up}, scoreDiffRef)
    n = size(g, 1)
    for col in 1:n
        mergeLeft!(view(g, :, col), scoreDiffRef)
    end
end
function move!(g::Grid, dir::Val{:down}, scoreDiffRef)
    n = size(g, 1)
    for col in 1:n
        reverse!(mergeLeft!(reverse!(view(g, :, col)), scoreDiffRef))
    end
end

# Termbox GUI
function putchar(k, i, c)
    tb_change_cell(k, i, c, 0x0001, 0x0000)
end

#TODO: Refactor interface below into boxes
function tb_print(g::Grid, numWidth::Int = 4)
    tb_clear()
    n = size(g, 1)
    rowCounter = 1
    for row in 1:n
        for i in 1:(n*(numWidth+4)+1)
            putchar(i, rowCounter, '#')
        end
        rowCounter += 1
        colCounter = 1
        for col in 1:n
            putchar(colCounter, rowCounter, '#')
            colCounter += 2 + numWidth + 2
            putchar(colCounter, rowCounter, '#')
        end
        rowCounter += 1
        colCounter = 1
        for col in 1:n
            colCounter += 2
            for (i, c) in enumerate(rpad(sprint(show, g[row, col]), numWidth))
                putchar(colCounter, rowCounter, c)
                colCounter += 1
            end
            colCounter += 2
            putchar(colCounter, rowCounter, '#')
        end
        rowCounter += 1
        colCounter = 1
        for col in 1:n
            putchar(colCounter, rowCounter, '#')
            colCounter += 2 + numWidth + 2
            putchar(colCounter, rowCounter, '#')
        end
        rowCounter += 1
    end
    for i in 1:(n*(numWidth+4)+1)
        putchar(i, rowCounter, '#')
    end
end

function handle_char!(c, g)
    if c == Int('r')
        fill!(g, 0)
        randomGrid!(g)
        tb_clear()
        return true
    end
    if c == Int('q')
        return false
    end
    return true
end
function handle_key!(k, g)
    scorediff = Ref(0)
    if k == TB_KEY_ARROW_LEFT
        move!(g, Val(:left), scorediff)
    elseif k == TB_KEY_ARROW_RIGHT
        move!(g, Val(:right), scorediff)
    elseif k == TB_KEY_ARROW_UP
        move!(g, Val(:up), scorediff)
    elseif k == TB_KEY_ARROW_DOWN
        move!(g, Val(:down), scorediff)
    end
    if isempty(zeroElements(g))
        return -1
    end
    if scorediff[] >= 0
        randomGrid!(g, 1)
    end
    return scorediff[]
end
function main_interactive_loop(boardSize::Int = 4)
    do_loop = true
    score = 0
    # g = newGrid(4)
    # g[2, 1] = 2
    # g[2, 2] = 2

    g = randomGrid(boardSize)

    tb_print(g)
    tb_present()
    while do_loop
        ev = tb_event()
        while tb_poll_event(ev) != 0
            if ev.etype == TB_EVENT_KEY
                if ev.key == TB_KEY_ESC
                    do_loop = false
                    break
                elseif ev.key == 0
                    if !(handle_char!(ev.ch, g))
                        do_loop = false
                        break
                    end
                else
                    scorediff = handle_key!(ev.key, g)
                    if scorediff < 0
                        do_loop = false
                        break
                    end
                    score += scorediff
                end
            end #if
            tb_print(g)
            tb_present()
            tb_peek_event(ev, 0)
        end
    end
    score
end

function main_interactive(boardSize::Int = 4)
    ret = tb_init()
    if ret != 0
        @printf stderr "tb_init() failed with error code %d\n" ret
        return 1
    end
    score = 0
    try
        tb_clear()
        tb_select_input_mode(TB_INPUT_ESC)
        tb_select_output_mode(TB_OUTPUT_NORMAL)

        score += main_interactive_loop(boardSize)
    finally
        tb_shutdown()
    end
    @show score
    return 0
end

function main_wargames_loop(sleepTime::Real = 0.05, boardSize::Int = 4, maxMoves::Integer = 1000)
    score = 0
    moves = 0

    g = randomGrid(boardSize)

    tb_print(g)
    tb_present()
    while (moves < maxMoves)
        scorediff = handle_key!(rand([TB_KEY_ARROW_LEFT, TB_KEY_ARROW_RIGHT, TB_KEY_ARROW_DOWN, TB_KEY_ARROW_UP]), g)
        if scorediff < 0
            break
        end
        score += scorediff
        tb_print(g)
        tb_present()
        sleep(sleepTime)
        moves += 1
    end
    score
end

function main_wargames(sleepTime::Real = 0.05, boardSize::Int = 4)
    ret = tb_init()
    if ret != 0
        @printf stderr "tb_init() failed with error code %d\n" ret
        return 1
    end
    score = 0
    try
        tb_clear()
        tb_select_input_mode(TB_INPUT_ESC)
        tb_select_output_mode(TB_OUTPUT_NORMAL)

        score += main_wargames_loop(sleepTime, boardSize)
    finally
        tb_shutdown()
    end
    @show score
    return 0
end
# v2 = newGrid(4)
# v2[2, 1] = 2
# v2[2, 2] = 2
# println(v2)
# println(move!(v2, Val(:left)))
# println(v2)

end # module

# To use this interactively, run
# julia -L 2048Example.jl -e 'CLI2048.main_interactive()'
#
# To see a game played randomly (and quickly), run
# julia -L 2048Example.jl -e 'CLI2048.main_wargames()'
# The latter command is suitable for iteration on the code itself.


using StatsBase
using Base.Test

X = randn(3, 8)

Z1 = X .- mean(X, 1)
Z2 = X .- mean(X, 2)

w1 = rand(3)
w2 = rand(8)

wv1 = weights(w1)
wv2 = weights(w2)

Z1w = X .- mean(X, wv1, 1)
Z2w = X .- mean(X, wv2, 2)

## reference results

S1 = Z1'Z1
S2 = Z2 * Z2'

Sz1 = X'X
Sz2 = X * X'

S1w = Z1w' * diagm(w1) * Z1w
S2w = Z2w * diagm(w2) * Z2w'

Sz1w = X' * diagm(w1) * X
Sz2w = X * diagm(w2) * X'

## scattermat

if VERSION < v"0.5.0-dev+679"
    @test_approx_eq scattermat(X) S1
    @test_approx_eq scattermat(X; vardim=2) S2

    @test_approx_eq scattermat(X; mean=0) Sz1
    @test_approx_eq scattermat(X; mean=0, vardim=2) Sz2

    @test_approx_eq scattermat(X; mean=mean(X,1)) S1
    @test_approx_eq scattermat(X; mean=mean(X,2), vardim=2) S2

    @test_approx_eq scattermat(X; mean=zeros(1,8)) Sz1
    @test_approx_eq scattermat(X; mean=zeros(3), vardim=2) Sz2

    ## weighted scatter mat

    @test_approx_eq scattermat(X, wv1) S1w
    @test_approx_eq scattermat(X, wv2; vardim=2) S2w

    @test_approx_eq scattermat(X, wv1; mean=0) Sz1w
    @test_approx_eq scattermat(X, wv2; mean=0, vardim=2) Sz2w

    @test_approx_eq scattermat(X, wv1; mean=mean(X, wv1, 1)) S1w
    @test_approx_eq scattermat(X, wv2; mean=mean(X, wv2, 2), vardim=2) S2w

    @test_approx_eq scattermat(X, wv1; mean=zeros(1,8)) Sz1w
    @test_approx_eq scattermat(X, wv2; mean=zeros(3), vardim=2) Sz2w
else
    @test_approx_eq scattermat(X) S1
    @test_approx_eq scattermat(X, 2) S2

    @test_approx_eq StatsBase.scattermatm(X, 0) Sz1
    @test_approx_eq StatsBase.scattermatm(X, 0, 2) Sz2

    @test_approx_eq StatsBase.scattermatm(X, mean(X,1)) S1
    @test_approx_eq StatsBase.scattermatm(X, mean(X,2), 2) S2

    @test_approx_eq StatsBase.scattermatm(X, zeros(1,8)) Sz1
    @test_approx_eq StatsBase.scattermatm(X, zeros(3), 2) Sz2

    ## weighted scatter mat

    @test_approx_eq scattermat(X, wv1) S1w
    @test_approx_eq scattermat(X, wv2, 2) S2w

    @test_approx_eq StatsBase.scattermatm(X, 0, wv1) Sz1w
    @test_approx_eq StatsBase.scattermatm(X, 0, wv2, 2) Sz2w

    @test_approx_eq StatsBase.scattermatm(X, mean(X, wv1, 1), wv1) S1w
    @test_approx_eq StatsBase.scattermatm(X, mean(X, wv2, 2), wv2, 2) S2w

    @test_approx_eq StatsBase.scattermatm(X, zeros(1,8), wv1) Sz1w
    @test_approx_eq StatsBase.scattermatm(X, zeros(3), wv2, 2) Sz2w
end

# weighted covariance

if VERSION < v"0.5.0-dev+679"
    @test_approx_eq cov(X, wv1) S1w ./ sum(wv1)
    @test_approx_eq cov(X, wv2; vardim=2) S2w ./ sum(wv2)

    @test_approx_eq cov(X, wv1; mean=0) Sz1w ./ sum(wv1)
    @test_approx_eq cov(X, wv2; mean=0, vardim=2) Sz2w ./ sum(wv2)

    @test_approx_eq cov(X, wv1; mean=mean(X, wv1, 1)) S1w ./ sum(wv1)
    @test_approx_eq cov(X, wv2; mean=mean(X, wv2, 2), vardim=2) S2w ./ sum(wv2)

    @test_approx_eq cov(X, wv1; mean=zeros(1,8)) Sz1w ./ sum(wv1)
    @test_approx_eq cov(X, wv2; mean=zeros(3), vardim=2) Sz2w ./ sum(wv2)
else
    @test_approx_eq cov(X, wv1) S1w ./ sum(wv1)
    @test_approx_eq cov(X, wv2, 2) S2w ./ sum(wv2)

    @test_approx_eq Base.covm(X, 0, wv1) Sz1w ./ sum(wv1)
    @test_approx_eq Base.covm(X, 0, wv2, 2) Sz2w ./ sum(wv2)

    @test_approx_eq Base.covm(X, mean(X, wv1, 1), wv1) S1w ./ sum(wv1)
    @test_approx_eq Base.covm(X, mean(X, wv2, 2), wv2, 2) S2w ./ sum(wv2)

    @test_approx_eq Base.covm(X, zeros(1,8), wv1) Sz1w ./ sum(wv1)
    @test_approx_eq Base.covm(X, zeros(3), wv2, 2) Sz2w ./ sum(wv2)
end

# mean_and_cov
if VERSION < v"0.5.0-dev+679"
    (m, C) = mean_and_cov(X; vardim=1)
    @test m == mean(X, 1)
    @test C == cov(X; vardim=1)

    (m, C) = mean_and_cov(X; vardim=2)
    @test m == mean(X, 2)
    @test C == cov(X; vardim=2)

    (m, C) = mean_and_cov(X, wv1; vardim=1)
    @test m == mean(X, wv1, 1)
    @test C == cov(X, wv1; vardim=1)

    (m, C) = mean_and_cov(X, wv2; vardim=2)
    @test m == mean(X, wv2, 2)
    @test C == cov(X, wv2; vardim=2)
else
    (m, C) = mean_and_cov(X, 1)
    @test m == mean(X, 1)
    @test C == cov(X, 1)

    (m, C) = mean_and_cov(X, 2)
    @test m == mean(X, 2)
    @test C == cov(X, 2)

    (m, C) = mean_and_cov(X, wv1, 1)
    @test m == mean(X, wv1, 1)
    @test C == cov(X, wv1, 1)

    (m, C) = mean_and_cov(X, wv2, 2)
    @test m == mean(X, wv2, 2)
    @test C == cov(X, wv2, 2)
end


using DiffEqGPU, StochasticDiffEq, Test, DiffEqNoiseProcess
using CUDA
using Random

function f!(du,u,p,t)
 @inbounds begin
     du[1] = p[1]*u[1]
     du[2] = -p[1]*u[2]
 end
 nothing
end

function scalar_noise!(du,u,p,t)
 @inbounds begin
     du[1] = p[2]
     du[2] = -p[2]
 end
 nothing
end

function prob_func(prob, i, repeat)
    Random.seed!(seeds[i])
    W = WienerProcess(0.0,0.0,0.0)
    remake(prob,noise=W)
end

# fix seeds
seed = 100
Random.seed!(seed)
numtraj= 100
seeds = rand(UInt, numtraj)
W = WienerProcess(0.0,0.0,0.0)

CUDA.allowscalar(false)
u0 = [-0.1f0,0.0f0]
tspan = (0.0f0,100.0f0)
p = [1.9f0,0.1f0]

prob = SDEProblem(f!,scalar_noise!,u0,tspan,p, noise=W)
ensembleprob = EnsembleProblem(prob, prob_func = prob_func)

@info "scalar SDE"

#Performance check with nvvp
# CUDAnative.CUDAdrv.@profile
# check either on CPU with EnsembleCPUArray() or on GPU with EnsembleGPUArray()
#@test_nowarn sol = @time solve(ensembleprob,DRI1(),EnsembleCPUArray(),trajectories=numtraj)
sol = @time solve(ensembleprob,DRI1NM(),EnsembleGPUArray(),trajectories=numtraj)


## Load Packages
using Omega
using StatsBase
using Random
using Plots
using Distributions

Random.seed!(123)

"""
Run basic Lotka-Volterra Simulation
"""

function gillespie_(rng, transitions, N, t0=0.0)

    """
    Implementation of the Gillespie algorithm for Lotka-Volterra model

    Args:
    transitions(dict): A dictionary with keys "prey" and "pred", and each element
        is a list containing the row of the transition matrix.
    N(int): Number of desired simulation iterations.
    t0(float): initial time point.

    Returns:
    A tuple of lists of the predator, prey and time trajectories
    """

    initial = Dict("prey" => prey_init(rng), "pred" => pred_init(rng))
    prey_list = [initial["prey"]]
    pred_list = [initial["pred"]]
    times = [t0]
    ecology = initial

    theta = Dict("spawn_prey" => spawn_prey(rng),
                 "prey2pred" => prey2pred(rng),
                 "pred_dies" => pred_dies(rng))
    t = times[1]

    for i = 1:N
        hazards = get_hazards(ecology, theta)
        transition = transitions[sample(collect(keys(hazards)), Weights(collect(values(hazards))))]
        t = t + sum(values(hazards))

        ecology["prey"] += transition[1]
        ecology["pred"] += transition[2]
        # Enforce only positive integers
        ecology["prey"] = max(1, ecology["prey"])
        ecology["pred"] = max(1, ecology["pred"])

        append!(prey_list, ecology["prey"])
        append!(pred_list, ecology["pred"])
        append!(times, t)
    end

    return (prey_list, pred_list, times)
end

function get_hazards(ecology, theta)
    """
    Compute the hazard function given the current states.  Note that in this
        model the function depends only on state, not time, though this is not
        the case in general. "spawn_prey" represents the event of a prey being born,
        "prey2pred" represents a predator consuming a new prey and consequently spawning
        a new predator, "pred_dies" represents the death of a predator.

    args:
        ecology(dict): a dictionary containing species counts.
        theta(dict): A dictionary where keys represent events (reactions) and
            values are the location hyperparameters for rate parameter constants
            corresponding to those events.
    """

    return Dict(
        "spawn_prey" => theta["spawn_prey"] * ecology["prey"],
        "prey2pred" => theta["prey2pred"] * ecology["prey"] * ecology["pred"],
        "pred_dies" => theta["pred_dies"] * ecology["pred"]
        )
end

# Define gillespie parameters
Pre = [[1, 0], [1, 1], [0, 1]]
Post = [[2, 0], [0, 2], [0, 0]]
transition_mat = Post - Pre
transitions = Dict("spawn_prey" => transition_mat[1,],
                    "prey2pred" => transition_mat[2,],
                    "pred_dies" => transition_mat[3,])
t0=0.0
N = 1000

# Initiate starting values randomely
prey_init = normal(1., .3)
pred_init = normal(1., .3)

# Random variables for rates
spawn_prey = normal(1.0, .01)
prey2pred = normal(0.1, .0001)
pred_dies = normal(1.0, .0075)

# Make gillespie into random variable
gillespie = ciid(gillespie_, transitions, N, 0.0)

# Sample from random variables
samples = rand((prey_init, pred_init, spawn_prey, prey2pred, pred_dies, gillespie),
                5, alg = RejectionSample)

plot_vals = hcat(values(samples[2][6][1]), values(samples[2][6][2]))
basic_plot = plot(1:N+1, plot_vals[1:N+1,1:2],
        title = "Basic Simulation",
        label = ["Prey" "Pred"],
        lw = 1.25)
#savefig(basic_plot,"./plots/basic_simulation.png")
__________________________________________________________________________________

"""
Run a conditioned model of LV using Gillespie

The conditional statement reads as:
Condition the simulation as the mean number of prey over cond_time_range is over
the mean_limit
"""

#Define helper functions for our conditional statement
function get_prey_list(g_output)
    return g_output[1]
end

function cond_statement_func(list, N)
    mean = 0
    for n = 1:N
        mean += list[end-n]
    end
    return mean/N
end

#define initial parameters for our conditional statement
cond_time_range = 250
mean_limit = 80

#Turn gillespie into a random variable
gillespie = ciid(gillespie_, transitions, N, 0.0)

#Apply our helper functions to get a conditional statement
prey_list = lift(get_prey_list)(gillespie)
cond_statement = lift(cond_statement_func)(prey_list, cond_time_range)

#run the simulation get the conditioned trajectories
samples_cond = rand((prey_init, pred_init, spawn_prey, prey2pred, pred_dies, gillespie, cond_statement),
                cond_statement > mean_limit, 5, alg = RejectionSample)
samples_cond[2][7]

#plot the conditioned predator and prey trajectories
plot_vals = hcat(values(samples_cond[2][6][1]), values(samples_cond[2][6][2]))
cond_plot = plot(1:N+1, plot_vals[1:N+1,1:2],
        title = "Conditioned Model",
        label = ["Prey" "Pred"],
        lw = 1.25)
#savefig(cond_plot,"./plots/conditioned_model.png")

__________________________________________________________________________
"""
Run an Interventional model of LV using Gillespie

The interventional statement reads as:
Intervene the simulation at n_int by setting the (predator, prey) value to
(prey_int, pred_int)
"""
function gillespie_int_(rng, transitions, N, n_int, prey_int, pred_int, t0=0.0)
    """
    Implementation of the Gillespie algorithm with intervention for Lotka-Volterra model

    Args:
    transitions(dict): A dictionary with keys "prey" and "pred", and each element
        is a list containing the row of the transition matrix.
    N(int): Number of desired simulation iterations.
    n_int: time point of intervention
    prey_int: value prey is intervened to at time of intervention
    pred_int: value pred is intervened to at time of intervention
    t0(float): initial time point.

    Returns:
    A tuple of lists of the intervened predator, prey and time trajectories
    """

    initial = Dict("prey" => prey_init(rng), "pred" => pred_init(rng))
    prey_list = [initial["prey"]]
    pred_list = [initial["pred"]]
    times = [t0]
    ecology = initial

    theta = Dict("spawn_prey" => spawn_prey(rng),
                 "prey2pred" => prey2pred(rng),
                 "pred_dies" => pred_dies(rng))
    t = times[1]

    for i = 1:N
        hazards = get_hazards(ecology, theta)
        transition = transitions[sample(collect(keys(hazards)), Weights(collect(values(hazards))))]
        t = t + sum(values(hazards))

        ecology["prey"] += transition[1]
        ecology["pred"] += transition[2]
        # Enforce only positive integers

        if i == n_int
            if !isnothing(prey_int)
                ecology["prey"] = prey_int
            end
            if !isnothing(pred_int)
                ecology["pred"] = pred_int
            end
        end

        ecology["prey"] = max(1, ecology["prey"])
        ecology["pred"] = max(1, ecology["pred"])

        append!(prey_list, ecology["prey"])
        append!(pred_list, ecology["pred"])
        append!(times, t)
    end

    return (prey_list, pred_list, times)
end

# Define parameters for our interventional statement
t_int = 250
prey_int = 10
pred_int= nothing

# Sample from random variables
gillespie_int = ciid(gillespie_int_, transitions, N, t_int, prey_int, pred_int, 0.0)
samples_int = rand((prey_init, pred_init, spawn_prey, prey2pred, pred_dies, gillespie_int),
                5, alg = RejectionSample)

#plot the trajectories of the intervened model
plot_vals = hcat(values(samples_int[2][6][1]), values(samples_int[2][6][2]))
cull_prey_plot = plot(1:N+1, plot_vals[1:N+1,1:2],
        title = "Action: Cull Prey",
        label = ["Prey" "Pred"],
        lw = 1.25)
#savefig(cull_prey_plot,"./plots/cull_prey_plot.png")

#generate histograms for treatment effect
t_int = 250
prey_int = 10
pred_int= nothing
gillespie_int = ciid(gillespie_int_, transitions, N, t_int, prey_int, pred_int, 0.0)
samples_prey_cull_treat = rand((prey_init, pred_init, spawn_prey, prey2pred, pred_dies, gillespie, gillespie_int),
                5, alg = RejectionSample)

prey_cull_treat_effect_hist = Plots.histogram(samples_prey_cull_treat[2][7][1] - samples_prey_cull_treat[2][6][1], bins = 10,
        title = "Prey Cull Treatment Effect")
#savefig(prey_cull_treat_effect_hist,"./plots/prey_cull_treat_effect_hist.png")
t_int = 250
prey_int = nothing
pred_int= 60
gillespie_int = ciid(gillespie_int_, transitions, N, t_int, prey_int, pred_int, 0.0)
samples_pred_cull_treat = rand((prey_init, pred_init, spawn_prey, prey2pred, pred_dies, gillespie, gillespie_int),
                5, alg = RejectionSample)

pred_inc_treat_effect_hist = Plots.histogram(samples_prey_cull_treat[2][7][1] - samples_prey_cull_treat[2][6][1], bins = 10,
        title = "Pred Inc Treatment Effect")
#savefig(pred_inc_treat_effect_hist,"./plots/pred_inc_treat_effect_hist.png")

__________________________________________________________________________
"""
Run a Counterfactual model of LV using Gillespie

The interventional statement reads as:
Given some simulation what would have happened if at n_int the (predator, prey)
values were (prey_int, pred_int)
"""
function gillespie_count_(prey_inc, pred_inc, n_int)
    """
    Implementation of a Counterfactual Gillespie algorithm for Lotka-Volterra model

    Args:
    n_int: time point of intervention
    prey_int: value prey is intervened to at time of intervention
    pred_int: value pred is intervened to at time of intervention
    t0(float): initial time point.

    Returns:
    A gillespie function with the args defined as constants
    """
    function gillespie_(rng, transitions, N, t0=0.0)
        """
        Implementation of the Gillespie algorithm for Lotka-Volterra model

        Args:
        transitions(dict): A dictionary with keys "prey" and "pred", and each element
            is a list containing the row of the transition matrix.
        N(int): Number of desired simulation iterations.
        t0(float): initial time point.

        Returns:
        A tuple of lists of the intervened predator, prey and time trajectories
        """
        initial = Dict("prey" => prey_init(rng), "pred" => pred_init(rng))
        prey_list = [initial["prey"]]
        pred_list = [initial["pred"]]
        times = [t0]
        ecology = initial

        theta = Dict("spawn_prey" => spawn_prey(rng),
                     "prey2pred" => prey2pred(rng),
                     "pred_dies" => pred_dies(rng))
        t = times[1]
        for i = 1:N
            hazards = get_hazards(ecology, theta)
            transition = transitions[sample(collect(keys(hazards)), Weights(collect(values(hazards))))]
            t = t + sum(values(hazards))

            ecology["prey"] += transition[1]
            ecology["pred"] += transition[2]
            # Enforce only positive integers

            if i == n_int
                ecology["prey"] = prey_inc
                ecology["pred"] = pred_inc
            end

            ecology["prey"] = max(1, ecology["prey"])
            ecology["pred"] = max(1, ecology["pred"])

            append!(prey_list, ecology["prey"])
            append!(pred_list, ecology["pred"])
            append!(times, t)
        end
        return (prey_list, pred_list, times)
    end
end

#Define the baseline model
gillespie = ciid(gillespie_, transitions, N, 0.0)
gillispie_count_func_ = gillespie_count_(0, 20, 250)

#Apply the counterfactual step
gillispie_count = replace(gillespie, gillespie_ => gillispie_count_func_)

#Sample and run the simulation of the base and counterfactual models
samples_count = rand((prey_init, pred_init, spawn_prey, prey2pred, pred_dies, gillespie, gillispie_count),
                5, alg = RejectionSample)

#Plot the trajectories of the baseline model
plot_vals = hcat(values(samples_count[2][6][1]), values(samples_count[2][6][2]))
count_none_plot = plot(1:N+1, plot_vals[1:N+1,1:2],
        title = "No Counterfactual",
        label = ["Prey" "Pred"],
        lw = 1.25)
#Plot the trajectories of the counterfactual applied to the base model
plot_vals = hcat(values(samples_count[2][7][1]), values(samples_count[2][7][2]))
count_ind_pred_plot = plot(1:N+1, plot_vals[1:N+1,1:2],
        title = "Counterfactual: Inc Predators",
        label = ["Prey" "Pred"],
        lw = 1.25)

#savefig(count_ind_pred_plot,"./plots/count_ind_pred_plot.png")


struct RBDFAST <: GSAMethod
    num_harmonics::Int
end
RBDFAST(;num_harmonics = 6) = RBDFAST(num_harmonics)

"""
Code based on the theory presented in:
    Saltelli, A. (2008). Global sensitivity analysis: The primer. Chichester: Wiley, pp. 167-169.
and
    S. Tarantola, D. Gatelli and T. Mara (2006)
    "Random Balance Designs for the Estimation of First Order Global Sensitivity Indices",
    Reliability Engineering and System Safety, 91:6, 717-727
"""

using FFTW, Random, Statistics, StatsBase, Distributions
allsame(x) = all(y -> y == first(x), x)

function gsa(f, method::RBDFAST; num_params, N, rng::AbstractRNG = Random.default_rng(), batch = false, kwargs...)
    # Initalize matrix containing range of values of the parametric variable
    # along each column (factor).
    s0 = range(-π, stop = π, length = N)

    # Compute inputs
    s = [s0[randperm(rng, N)] for i in 1:num_params]
    x = hcat([0.5.+asin.(sin.(s[i]))./pi for i in 1:num_params]...)

    # Compute outputs
    if batch
        Y = f(x')
    else
        Y = [f(@view x[i,:]) for i in axes(x, 1)]
    end
    # Iterate over factors

    sensitivites = zeros(num_params)
    for i in 1:num_params

        s_order = sortperm(s[i])
        # Order Ys by how they would occur if they were
        # monotonically increasing as the
        # parametric variable s (not its permutation) increased.
        y_reordered = @view Y[s_order]

        ft = fft(y_reordered)
        ys = abs2.(ft) ./ N
        V = sum(ys[2:N])
        Vi = 2*sum(ys[2:method.num_harmonics+1])
        Si = Vi/V
        # println(ys)
        # unskew the sensitivies
        # lambda = 2*method.num_harmonics/N
        # sensitivites[i] = Si - (lambda / (1 - lambda)) * (1-Si)
        sensitivites[i] = Si
    end

    return sensitivites
end


# -- PRIVATE FUNCTIONS ------------------------------------------- #
function _iv_objective_function(x, assetSet::Set{PSAbstractAsset}, parameters::PSBinaryLatticeModel, baseAssetPriceValue::Float64, optionPriceValue::Float64)

    volatility::Float64
    timeToExercise::Float64
    riskFreeRate::Float64
    dividendRate::Float64
    numberOfLevels::Int64

    # create a new paramter w/new volatility -
    perturbedVolatility = x[1]
    perturbedLatticeModel = PSBinaryLatticeModel(perturbedVolatility, parameters.timeToExercise, parameters.riskFreeRate,
        parameters.dividendRate; numberOfLevels=parameters.numberOfLevels)

    #= We are calling:
    option_contract_price(contractSet::Set{PSAbstractAsset}, latticeModel::PSBinaryLatticeModel, baseUnderlyingPrice::Float64;
        earlyExercise::Bool = true)::PooksoftBase.PSResult
    =#

    # re-compute the option price -
    result = option_contract_price(assetSet, perturbedLatticeModel, baseAssetPriceValue)
    results_tuple = result.value
    estimatedPriceValue = results_tuple.cost_calculation_result.option_contract_price_array[1];

    # compute the error -
    error_term = (optionPriceValue - estimatedPriceValue)

    # return -
    return error_term*error_term
end
# ---------------------------------------------------------------- #

# -- PUBLIC FUNCTIONS -------------------------------------------- #
function estimate_implied_volatility(contract::PSAbstractAsset, parameters::PSBinaryLatticeModel, baseAssetPriceValue::Float64;
    earlyExercise::Bool = false)

    # setup asset set -
    assetSet = Set{PSAbstractAsset}()
    push!(assetSet, contract)

    # what is the premium?
    optionPriceValue = contract.premimumValue

    # what is the initial volatility -
    initialVolatility = [parameters.volatility]

    # setup the objective function -
    OF(p) = _iv_objective_function(p, assetSet, parameters, baseAssetPriceValue, optionPriceValue)

    # call the optimizer -
    opt_result = Optim.optimize(OF, initialVolatility, BFGS())

    # return the result -
    return Optim.minimizer(opt_result)[1]
end

function compute_weighted_volatility(data::DataFrame, weightKey::Symbol, volatilityKey::Symbol)::PSResult

    # initialize -
    tmp_array = Array{Float64,1}()
    volatility_value = nothing

    # sum the wght key -
    iv_array = data[!,volatilityKey]
    wght_array = data[!,weightKey]

    # main loop -
    for (index, iv_value) in enumerate(iv_array)

        # get weight value -
        wght_value = wght_array[index]

        # compute weighted iv value numerator -
        value = wght_value*iv_array[index]
        push!(tmp_array,value)
    end

    # compute value -
    wght_sum = sum(wght_array)
    volatility_value = sum(tmp_array)*(1/wght_sum)

    # return -
    return PSResult(volatility_value)
end
# ---------------------------------------------------------------- #


using .Groebner: _normal_form_reference, _reducegb_reference, _reducegb_reference!,
                    change_ordering, rootn

@testset "Normal Form" begin

    R, (x, y, z) = PolynomialRing(QQ, ["x", "y", "z"])

    Gs = [
        [x], [x, y], [x, y, z]
    ]

    @test _normal_form_reference(x, Gs[1]) == 0
    @test _normal_form_reference(y, Gs[1]) == y
    @test _normal_form_reference(x + y^2, Gs[3]) == 0
    @test _normal_form_reference(5x^2 - y^3 - 1, Gs[2]) == -1
    @test _normal_form_reference(x + y + 3z, Gs[2]) == 3z

    G = [
        x^2 + y*x + 1,
        y^2 - y
    ]

    @test _normal_form_reference(x + y + z, G) == x + y + z
    @test _normal_form_reference(x^2 + y^2, G) == y - y*x - 1
    @test _normal_form_reference(y^3, G) == _normal_form_reference(y^4, G) == y

    G = [
        x*y + y,
        z + 5,
    ]
    @test _normal_form_reference(x*y, G) == -y
    @test _normal_form_reference(3z - 1, G) == -16

    @test _normal_form_reference(8x*y + 88z, [R(1)]) == 0

    R, (x, y, z) = PolynomialRing(QQ, ["x","y","z"], ordering=:degrevlex)
    G = [
        x^2 + y,
        y^2 + x
    ]
    @test _normal_form_reference(x^2 + y^2, G) == -x - y

end

@testset "Change ordering" begin
    R, (x, y, z) = PolynomialRing(QQ, ["x", "y", "z"])

    fs = [1x + 2y*z + 3z^3, y*z, 5x]
    fs_deg = change_ordering(fs, :degrevlex)
    R_deg = parent(first(fs_deg))

    @test ordering(R_deg) == :degrevlex
    @test collect(exponent_vectors(fs_deg[1])) == [[0, 0, 3], [0, 1, 1], [1, 0, 0]]
    @test collect(exponent_vectors(fs_deg[2])) == [[0, 1, 1]]
    @test collect(exponent_vectors(fs_deg[3])) == [[1, 0, 0]]

    fs_vv = change_ordering(fs, :lex)
    @test fs_vv == fs
end


using Test, Neighborhood, StaticArrays, Random, Distances
using Neighborhood: datatype, getmetric, bruteforcesearch
using Neighborhood.Testing

Random.seed!(54525)
data = [rand(SVector{3}) for i in 1:1000]
query = SVector(0.99, 0.99, 0.99)

# Theiler window and skip predicate related
data[48:52] .= [SVector(0.0, 0.0, i*0.001) for i in 1:5]
nidxs = 48:2:52
queries = [data[i] for i in nidxs]
theiler1 = Theiler(1, nidxs)
theiler2 = Theiler(2, nidxs)

r = 0.1
k = 5

@testset "Utils" begin include("util.jl") end
@testset "Neighborhood.Testing" begin include("Testing.jl") end
@testset "Brute force" begin include("bruteforce.jl") end
include("nearestneighbors.jl")


function gschurtest(A::Matrix{T}, B::Matrix{T}, tol;
                   normal=false, commuting=false) where {T<:Complex}
    n = size(A,1)
    ulp = eps(real(T))
    if real(T) <: BlasFloat
        S = GenericSchur.ggschur!(copy(A), copy(B))
    else
        S = schur(A, B)
    end
    # test 1: S.T is upper triangular
    @test norm(tril(S.S,-1)) / (n * norm(A) * ulp) < tol
    @test norm(tril(S.T,-1)) / (n * norm(A) * ulp) < tol
    # @test all(tril(S.S,-1) .== 0)
    # @test all(tril(S.T,-1) .== 0)
    # test 2: norm(A - S.Q * S.T * S.Z') / (n * norm(A) * ulp) < tol
    @test norm(A - S.Q * S.S * S.Z') / (n * norm(A) * ulp) < tol
    @test norm(B - S.Q * S.T * S.Z') / (n * norm(A) * ulp) < tol
    # test 3: S.Z is unitary: norm(I - S.Z * S.Z') / (n * ulp) < tol
    @test norm(I - S.Z * S.Z') / (n * ulp) < tol
    @test norm(I - S.Q * S.Q') / (n * ulp) < tol
    # test 4: S.values are e.v. of T
#    @test all(csort(S.values) .== csort(diag(S.T)))

    # It is tempting to check eigenvalues against LAPACK eigvals(A)
    # for suitable types, but the comparison is misleading without
    # detailed analysis because many of the test matrices are severly
    # non-normal.

    if normal && commuting
        # verify that Schur diagonalizes normal commuting matrices
        @test norm(triu(S.S,1)) / (n * norm(B) * ulp) < tol
        @test norm(triu(S.T,1)) / (n * norm(A) * ulp) < tol
    end

    VR = eigvecs(S)
    evcheck = norm((A * VR) * diagm(0 => S.β) - (B * VR) * diagm(0 => S.α))
    @test evcheck / (n*norm(A) * ulp) < tol
    # someday maybe left eigenvectors too
end

@testset "generalized sanity $T" for T in [ComplexF64, Complex{BigFloat}]

unfl = floatmin(real(T))
ovfl = one(real(T)) / unfl
ulp = eps(real(T))
ulpinv = one(real(T)) / ulp
rtulp = sqrt(ulp)
rtulpi = one(real(T)) / rtulp

# sizes of matrices to test
ens = [4,32]

    for n in ens
        A = rand(T,n,n)
        B = rand(T,n,n)
        gschurtest(A, B, 20)
    end
end # testset

#end # module


function pcontrib{T<:Integer}(p::T, a::T)
    n = one(T)
    pcon = one(T)
    for i in 1:a
        n *= p
        pcon += n
    end
    return pcon
end

function divisorsum{T<:Integer}(n::T)
    dsum = one(T)
    for (p, a) in factor(n)
        dsum *= pcontrib(p, a)
    end
    dsum -= n
end


"""
    divergence(f)
"""
divergence(f) = tr(gradient(f))

function symmetric_gradient end

"""
    symmetric_gradient(f)
"""
symmetric_gradient(f) = symmetic_part(gradient(f))

"""
    const ε = symmetric_gradient

Alias for the symmetric gradient
"""
const ε = symmetric_gradient

"""
    curl(f)
"""
curl(f) = grad2curl(gradient(f))

"""
    grad2curl(∇f)
"""
function grad2curl(f)
  @abstractmethod
end

function laplacian end

"""
    const Δ = laplacian

Alias for the `laplacian` function
"""
const Δ = laplacian

"""
    laplacian(f)
"""
function laplacian(f)
  g = gradient(f)
  h = gradient(g)
  tr(h)
end

"""
    ∇*f

Equivalent to

    divergence(f)
"""
(*)(::typeof(∇),f) = divergence(f)

"""
    outer(∇,f)

Equivalent to

    gradient(f)
"""
outer(::typeof(∇),f) = gradient(f)

"""
    outer(f,∇)

Equivalent to

    transpose(gradient(f))
"""
outer(f,::typeof(∇)) = transpose(gradient(f))

"""
    cross(∇,f)

Equivalent to

    curl(f)
"""
cross(::typeof(∇),f) = curl(f)

# Helpers

grad2curl(f::Field) = apply_kernel_to_field(bcast(_curl_kernel),f)

grad2curl(f::AbstractArray{<:Field}) = apply_to_field_array(bcast(_curl_kernel),f)

function _curl_kernel(∇u::TensorValue{2})
  ∇u[1,2] - ∇u[2,1]
end

function _curl_kernel(∇u::TensorValue{3})
  c1 = ∇u[2,3] - ∇u[3,2]
  c2 = ∇u[3,1] - ∇u[1,3]
  c3 = ∇u[1,2] - ∇u[2,1]
  VectorValue(c1,c2,c3)
end


##                                                                            ##
##              Problem set #4 for Structural modelling course                ##
##                               Nicolás Urdaneta                             ##
##                                                                            ##

using LinearAlgebra, Random, Statistics, DataFrames, CSV, FreqTables, Distributions, Optim, HTTP, GLM, ForwardDiff

url = "https://raw.githubusercontent.com/OU-PhD-Econometrics/fall-2020/master/ProblemSets/PS4-mixture/nlsw88t.csv"
df = CSV.read(HTTP.get(url).body)
X = [df.age df.white df.collgrad]
Z = hcat(df.elnwage1, df.elnwage2, df.elnwage3, df.elnwage4,df.elnwage5, df.elnwage6, df.elnwage7, df.elnwage8)
y = df.occ_code

## 1. multinomial logit as PS3

function mlogit_Z(theta, X, Z, y)

	alpha = theta[1:end-1]
	gamma = theta[end]

	K = size(X,2)
	J = length(unique(y))
	N = length(y)
	bigY = zeros(N,J)

	# Matrix with a column equal 1 if y equals that alternative
	for j=1:J
		bigY[:,j] = y.==j
	end

	## Matrix of parameters, normalice last alternative to zero
	bigAlpha = [reshape(alpha,K,J-1) zeros(K)]

	T = promote_type(eltype(X),eltype(theta))
	num   = zeros(T,N,J)
	den   = zeros(T,N)

	for j=1:J
		num[:,j] = exp.(X*bigAlpha[:,j] .+ (Z[:,j] .- Z[:,J])*gamma)
		den = den .+ num[:,j]
	end

	Prob = num./repeat(den,1,J)

	ll = -sum(bigY.*log.(Prob))

	return ll

end

## Optimize. Not recommended to use random start values (it could take voer 30 minutes. These are the results from PS3)

## startvals = [2*rand(7*size(X,2)).-1; .1]
startvals = [ .0403744; .2439942; -1.57132; .0433254; .1468556; -2.959103; .1020574; .7473086; -4.12005; .0375628; .6884899; -3.65577; .0204543; -.3584007; -4.376929; .1074636; -.5263738; -6.199197; .1168824; -.2870554; -5.322248; 1.307477]

## define as twice differentiable
td = TwiceDifferentiable(theta-> mlogit_Z(theta,X,Z,y), startvals; autodiff = :forward)

## Run

@time theta_hat_optim = optimize(td, startvals, LBFGS(), Optim.Options(g_tol=1e-5, iterations = 100_000, show_trace = true, show_every=50))

theta_hat_mle = theta_hat_optim.minimizer
H  = Optim.hessian!(td, theta_hat_mle)
theta_hat_mle_se = sqrt.(diag(inv(H)))
println([theta_hat_mle theta_hat_mle_se])

## 3. Mixed logit

# a) Quadrature axample normal
d = Normal(0,1)

nodes, weights = lgwt(7,-4,4)

sum(weights.*pdf.(d,nodes))
sum(weights.*nodes.*pdf.(d,nodes))

## b) More examples

# i)
d = Normal(0,2)
nodes, weights = lgwt(7,-10,10)
sum(weights.*nodes.*nodes.*pdf.(d,nodes))

# ii)
d = Normal(0,2)
nodes, weights = lgwt(10,-10,10)
sum(weights.*nodes.*nodes.*pdf.(d,nodes))

#These are aproximations to the variance, should get close to 4

## c) monte Carlo integration

# i)
b = 10
a = -10
c = Uniform(a,b)
draws = rand(c,1_000_000)
println((b-a)*mean(draws.^2 .*pdf.(d,draws)))
println((b-a)*mean(draws .*pdf.(d,draws)))
println((b-a)*mean(pdf.(d,draws)))

## 4. Question 1 code to optimize the expression using quadrature
	## Caution: do not try to run

function mlogit_Z_Quad(theta, X, Z, y, R)

		alpha = theta[1:end-2]
		gamma = theta[end-1]
		sigma = exp(theta[end]) # makes sure sigma is never 0 (which would make the Normal density undefined)

		K = size(X,2)
		J = length(unique(y))
		N = length(y)
		bigY = zeros(N,J)

		# Matrix with a column equal 1 if y equals that alternative
		for j=1:J
			bigY[:,j] = y.==j
		end

		## Matrix of parameters, normalice last alternative to zero
		bigAlpha = [reshape(alpha,K,J-1) zeros(K)]

		T = promote_type(eltype(X),eltype(theta))

		# Write probability as a function of variable of integration
		function ll_int(T,N,J,X,Z,bigAlpha,gamma,sigma,R)
			c = Uniform(-4*sigma,4*sigma)
			draws = rand(c,R)
			#nodes, weights = lgwt(R,-4*sigma,4*sigma)
			out = zeros(T,N)

			for r=1:R
				num = zeros(T,N,J)
				den = zeros(T,N)
				Prob = zeros(T,N,J)
				for j=1:J
					num[:,j] = exp.(X*bigAlpha[:,j] .+ (Z[:,j] .- Z[:,J])*nodes[r])
					den .+= num[:,j]
				end
				Prob = num./repeat(den,1,J)
				out .+= vec( weights[r]*(prod(Prob.^bigY; dims=2))*pdf(Normal(gamma,sigma), nodes[r]))
				# We use vec because vectors in Julia are "flat" and this operator flattens a one dimensional array.
			end
			return out
		end

		intll = ll_int(T,N,J,X,Z,bigAlpha,gamma,sigma,R)
		ll = -sum(log.(intll))

		return ll

end

## 5. Question 1 code to optimize the expression using Monte Carlo simulation
	## Caution: do not try to run

	function mlogit_Z_MonteC(theta, X, Z, y, R)

			alpha = theta[1:end-2]
			gamma = theta[end-1]
			sigma = exp(theta[end]) # makes sure sigma is never 0 (which would make the Normal density undefined)

			K = size(X,2)
			J = length(unique(y))
			N = length(y)
			bigY = zeros(N,J)

			# Matrix with a column equal 1 if y equals that alternative
			for j=1:J
				bigY[:,j] = y.==j
			end

			## Matrix of parameters, normalice last alternative to zero
			bigAlpha = [reshape(alpha,K,J-1) zeros(K)]

			T = promote_type(eltype(X),eltype(theta))

			# Write probability as a function of variable of integration
			function ll_int(T,N,J,X,Z,bigAlpha,gamma,sigma,R)
				nodes, weights = lgwt(R,-4*sigma,4*sigma)
				out = zeros(T,N)

				for r=1:R
					num = zeros(T,N,J)
					den = zeros(T,N)
					Prob = zeros(T,N,J)
					for j=1:J
						num[:,j] = exp.(X*bigAlpha[:,j] .+ (Z[:,j] .- Z[:,J])*nodes[r])
						den .+= num[:,j]
					end
					Prob = num./repeat(den,1,J)
					out .+= vec( (8*sigma/R)*(prod(Prob.^bigY; dims=2))*pdf(Normal(gamma,sigma), draws[r]))
					# We use vec because vectors in Julia are "flat" and this operator flattens a one dimensional array.
				end
				return out
			end

			intll = ll_int(T,N,J,X,Z,bigAlpha,gamma,sigma,R)
			ll = -sum(log.(intll))

			return ll

	end


@testset "runge_kutta" begin
    function f(t, y, θ)
        return y
    end

    x0 = 1.0
    out = rk4(f, 1.0, 100, x0)
    @test abs(run(sess, out)[end]-exp(1))<1e-5
    out = ode45(f, 1.0, 100, x0)
    @test abs(run(sess, out)[end]-exp(1))<1e-5

    function f(t, y, θ)
        return y
    end

    x0 = [1.0;1.0]
    out = rk4(f, 1.0, 100, x0)
    @test norm(run(sess, out)[end,:] .- exp(1))<1e-5
    out = ode45(f, 1.0, 100, x0)
    @test norm(run(sess, out)[end,:] .- exp(1))<1e-5

    x0 = [1.0 1.0
          1.0 1.0]
    out = rk4(f, 1.0, 100, x0)
    @test norm(run(sess, out)[end,:,:] .- exp(1))<1e-5
    out = ode45(f, 1.0, 100, x0)
    @test norm(run(sess, out)[end,:,:] .- exp(1))<1e-5

end

@testset "alpha scheme" begin
    d0 = zeros(2)
    v0 = [2.;3.]
    a0 = zeros(2)
    Δt = 1e-3
    n = 9001
    Δt = Δt * ones(n)
    M = sparse([1.0 0.0;0.0 1.0])
    C = sparse(zeros(2,2))
    K = sparse([4.0 0.0;0.0 9.0])
    F = zeros(n, 2)

    d, v, a = αscheme(M, C, K, F, d0, v0, a0, Δt)
    d_, v_, a_ = run(sess, [d, v, a])

    tspan = 0:1e-3:(n-1)*1e-3
    @test norm(d_[1:end-1,1] - sin.(2tspan))<1e-3
    @test norm(d_[1:end-1,2] - sin.(3tspan))<1e-3
end


module FEMDriver

using Test
using Gridap
using GridapPardiso

model = CartesianDiscreteModel(
  domain=(0,1,0,1,0,1), partition=(10,10,10))

fespace = FESpace(
  reffe=:Lagrangian, order=1, valuetype=Float64,
  conformity=:H1, model=model, diritags="boundary")

V = TestFESpace(fespace)
U = TrialFESpace(fespace,0.0)

trian = Triangulation(model)
quad = CellQuadrature(trian,degree=2)

t_Ω = AffineFETerm(
  (v,u) -> inner(∇(v),∇(u)),
  (v) -> inner(v, (x) -> x[1]*x[2] ),
  trian, quad)

op = LinearFEOperator(V,U,t_Ω)

ls = PardisoSolver()
solver = LinearFESolver(ls)

uh = solve(solver,op)

x = free_dofs(uh)
A = op.mat
b = op.vec

r = A*x - b

tol = 1e-10
@test maximum(abs.(r)) < tol

end #module


function _RobustDelague_latex()
    return """
        ```math
        \\left\\{ r  \\; \\middle| \\begin{array}{ll}
        s.t.  \\quad (\\mathbb{E} [r] - \\hat{r}) ' \\Sigma^{-1} (\\mathbb{E} [r] - \\hat{r}) \\leq \\gamma_1 \\\\
        \\quad \\quad \\mathbb{E} [ (r - \\hat{r}) ' (r - \\hat{r}) ] \\leq \\gamma_2 \\Sigma \\\\
        \\end{array}
        \\right\\} \\\\
        ```
        """
end

"""
    RobustDelague <: AbstractMeanVariance

Delague's Ambiguity set:

$(_RobustDelague_latex())

Atributes:
- `predicted_mean::Array{Float64,1}` (latex notation ``\\hat{r}``): Predicted mean of returns.
- `predicted_covariance::Array{Float64,2}` (latex notation ``\\Sigma``): Predicted covariance of returns.
- `γ1::Float64`: Mean uncertainty parameter (has to be greater than 0).
- `γ2::Float64`: Covariance uncertainty parameter (has to be greater than 1).
- `utility_coeficients::Array{Float64,1}`: Piece-wise utility coeficients (default 1).
- `utility_intercepts::Array{Float64,1}`: Piece-wise utility intercepts (default 0).
"""
struct RobustDelague <: AbstractPortfolioFormulation # AbstractMeanVariance
    predicted_mean::Array{Float64,1}
    predicted_covariance::Array{Float64,2}
    γ1::Float64
    γ2::Float64
    utility_coeficients::Array{Float64,1}
    utility_intercepts::Array{Float64,1}
    number_of_utility_pieces::Int
    number_of_assets::Int
end

function RobustDelague(;
    predicted_mean::Array{Float64,1},
    predicted_covariance::Array{Float64,2},
    γ1::Float64,
    γ2::Float64,
    utility_coeficients::Array{Float64,1},
    utility_intercepts::Array{Float64,1},
)
    number_of_assets = size(predicted_mean, 1)
    number_of_utility_pieces = size(utility_coeficients, 1)
    if number_of_assets != size(predicted_covariance, 1)
        error("size of predicted mean ($number_of_assets) different to size of predicted covariance ($(size(predicted_covariance, 1)))")
    end
    if number_of_utility_pieces != size(utility_intercepts, 1)
        error("number of utility coeficients ($(number_of_utility_pieces)) different to number of utility intercepts ($(size(utility_intercepts,1)))")
    end

    @assert γ1 >= 0
    @assert γ2 >= 1
    @assert issymmetric(predicted_covariance)
    return RobustDelague(
        predicted_mean, predicted_covariance, γ1, γ2, utility_coeficients,
        utility_intercepts, number_of_utility_pieces, number_of_assets
    )
end

"""
    po_max_utility_return(formulation::RobustDelague)

Maximize expected utility of portfolio return under distribution uncertainty defined by
Delague's ambiguity set ([`RobustDelague`](@ref)).
"""
function po_max_utility_return(formulation::RobustDelague;
    current_wealth::Real = 1.0,
    model::JuMP.Model = base_model(formulation.number_of_assets; current_wealth=current_wealth),
    w = model[:w],
)
    # parameters
    r̄ = formulation.predicted_mean
    Σ = formulation.predicted_covariance
    numA = formulation.number_of_assets
    a = formulation.utility_coeficients
    b = formulation.utility_intercepts
    γ1 = formulation.γ1
    γ2 = formulation.γ2
    K = formulation.number_of_utility_pieces
    # dual variables
    @variable(model, P[i=1:numA, j=1:numA])
    @variable(model, p[i=1:numA])
    @variable(model, s)
    @variable(model, Q[i=1:numA, j=1:numA])
    @variable(model, q[i=1:numA])
    @variable(model, r)
    # constraints: from duality
    @constraint(model, p .== -q / 2 - Q * r̄)
    @SDconstraint(model, [[P p]; [p' s]] >= 0)
    for k in 1:K
        @SDconstraint(
            model, [[Q (q / 2 + a[k] * w / 2)]; [(q' / 2 + a[k] * w' / 2) (r + b[k])]] >= 0
        )
    end

    # objective
    @objective(
        model,
        Min, # Min because this is the dual problem
        γ2 * dot(Σ, Q) - sum(r̄'Q * r̄) + r + dot(Σ, P) - 2 * dot(r̄, p) + γ1 * s
    )

    return model
end


export Camera

struct Camera{T}
    origin::Point3{T}
    lower_left_corner::Point3{T}
    horizontal::Vec3{T}
    vertical::Vec3{T}
end

function Camera{T}() where T
    aspect_ratio = T(16) / T(9)
    viewport_height = T(2)
    viewport_width = aspect_ratio * viewport_height
    focal_length = T(1)

    origin = Point3{T}(0, 0, 0)
    horizontal = Vec3{T}(viewport_width, 0, 0)
    vertical = Vec3{T}(0, viewport_height, 0)
    lower_left_corner = origin - horizontal/T(2) - vertical/T(2) - Vec3{T}(0, 0, focal_length)

    return Camera(origin, lower_left_corner, horizontal, vertical)
end

function get_ray(cam, u, v)
    return Ray(cam.origin, cam.lower_left_corner + u*cam.horizontal + v*cam.vertical - cam.origin)
end


part1(input) = simulate(input, 80)
part2(input) = simulate(input, 256)

function simulate(input, n)
    state = zeros(Int, 9)
    for age in parse.(Int, split(readchomp(input), ","))
        state[age + 1] += 1
    end
    for i = 1:n
        state0 = state[1]
        state[1:8] .= state[2:9]
        state[9] = state0
        state[7] += state0
    end
    return sum(state)
end


@testset "Eigenvector" begin
    g1 = smallgraph(:house)
    g2 = CycleDiGraph(4)

    for g in testgraphs(g1)
        y = @inferred(eigenvector_centrality(g))
        @test round.(y, digits=3) == round.([
            0.3577513877490464, 0.3577513877490464, 0.5298987782873977,
            0.5298987782873977, 0.4271328349194304
        ], digits=3)
    end
    for g in testdigraphs(g2)
        y = @inferred(eigenvector_centrality(g))
        @test round.(y, digits=3) == round.([0.5, 0.5, 0.5, 0.5], digits=3)
    end
end


### A Pluto.jl notebook ###
# v0.12.20

using Markdown
using InteractiveUtils

# ╔═╡ e7c05ce8-58ef-11eb-2da6-bfffd4491883
using Pkg, DrWatson

# ╔═╡ 7fb874b8-58f0-11eb-3e17-91468135856f
begin
	#@quickactivate "StatisticalRethinkingStan"
	using StanSample, StanOptimize
	using StatisticalRethinking
end

# ╔═╡ f7e9345a-58ef-11eb-347d-1b1c98f3663b
md" ## Clip-07-12s.jl"

# ╔═╡ 3c15dcca-58f3-11eb-046c-33c66612701d
md" ##### Entropy."

# ╔═╡ f7e8f63e-58ef-11eb-094b-85d2f930cbfc
H(p) = -sum(p .* log.(p))

# ╔═╡ 4a2cd688-58f3-11eb-071a-7d5191b2acc4
md" ##### Cross Entropy."

# ╔═╡ 313dbcfa-58f3-11eb-0b61-73d83a62cb76
H(p, q) = -sum(p .* log.(q))

# ╔═╡ 514bd3c4-58f3-11eb-27eb-055303c639b1
md" ##### Kullback-Leibler divergence."

# ╔═╡ 37a70450-58f0-11eb-00b1-c526f9615ff4
D(p, q) = sum(p .* log.(p ./ q))

# ╔═╡ 5e898a3a-58f4-11eb-1c3d-ef781cf3e9f0
md" ### Snippet 7.12"

# ╔═╡ 48b01fc0-58f0-11eb-10e1-8969cfc5022f
begin
	p = [0.3, 0.7]
	q = [0.25, 0.75]
	earth = [0.7, 0.3]
	mars = [0.01, 0.99]
end

# ╔═╡ 65e714ae-58f0-11eb-0f86-bb6dad500c12
H(p)

# ╔═╡ 90e27340-58f2-11eb-367e-cb2aae23b616
H([0.01, 0.99])

# ╔═╡ c0ee72a4-6668-11eb-1078-a5f09dfbecff
H([0.7, 0.15, 0.15])

# ╔═╡ 69d6127c-58f0-11eb-29c9-0ba1d12c61c2
D(p, q)

# ╔═╡ b6ffed52-58f0-11eb-1ab2-7902bd8a2b52
begin
	qrange = 0.001:0.01:1.0
	res = []
	for qstep in qrange
		qs = [qstep, 1-qstep]
		append!(res, [D(p, qs)])
	end
	plot(qrange, res, xlab="q[1]", ylab="Divergence q from p", leg=false)
	vline!([0.3])
end

# ╔═╡ edf40570-58f3-11eb-0650-03caff8b1c72
H(p, q)

# ╔═╡ 42415416-58f4-11eb-2cc3-c951f0416798
D(p, q)

# ╔═╡ 49cab272-58f4-11eb-02a8-375cce076851
H(p, q) - H(p)

# ╔═╡ 2c131646-58f6-11eb-12f2-f3b8dfb82dcb
md" ##### Divergence from earth -> mars."

# ╔═╡ 09550666-5994-11eb-3c4c-11461379f117
md"
!!! note
    Reverse arguments?
"

# ╔═╡ 755016a2-58f5-11eb-02c6-235a5d44e9c7
D(mars, earth)

# ╔═╡ 40d5189a-58f6-11eb-08c3-95d784d16c87
md" ##### Divergence from mars -> earth."

# ╔═╡ 6bc4b400-58f5-11eb-1de2-c9c8b44ae7ae
D(earth, mars)

# ╔═╡ 0ffd4f9a-58f0-11eb-1f66-236f6b48ab3f
md" ## End of clip-07-12s.jl"

# ╔═╡ Cell order:
# ╟─f7e9345a-58ef-11eb-347d-1b1c98f3663b
# ╠═e7c05ce8-58ef-11eb-2da6-bfffd4491883
# ╠═7fb874b8-58f0-11eb-3e17-91468135856f
# ╟─3c15dcca-58f3-11eb-046c-33c66612701d
# ╠═f7e8f63e-58ef-11eb-094b-85d2f930cbfc
# ╟─4a2cd688-58f3-11eb-071a-7d5191b2acc4
# ╠═313dbcfa-58f3-11eb-0b61-73d83a62cb76
# ╟─514bd3c4-58f3-11eb-27eb-055303c639b1
# ╠═37a70450-58f0-11eb-00b1-c526f9615ff4
# ╟─5e898a3a-58f4-11eb-1c3d-ef781cf3e9f0
# ╠═48b01fc0-58f0-11eb-10e1-8969cfc5022f
# ╠═65e714ae-58f0-11eb-0f86-bb6dad500c12
# ╠═90e27340-58f2-11eb-367e-cb2aae23b616
# ╠═c0ee72a4-6668-11eb-1078-a5f09dfbecff
# ╠═69d6127c-58f0-11eb-29c9-0ba1d12c61c2
# ╠═b6ffed52-58f0-11eb-1ab2-7902bd8a2b52
# ╠═edf40570-58f3-11eb-0650-03caff8b1c72
# ╠═42415416-58f4-11eb-2cc3-c951f0416798
# ╠═49cab272-58f4-11eb-02a8-375cce076851
# ╟─2c131646-58f6-11eb-12f2-f3b8dfb82dcb
# ╟─09550666-5994-11eb-3c4c-11461379f117
# ╠═755016a2-58f5-11eb-02c6-235a5d44e9c7
# ╟─40d5189a-58f6-11eb-08c3-95d784d16c87
# ╠═6bc4b400-58f5-11eb-1de2-c9c8b44ae7ae
# ╟─0ffd4f9a-58f0-11eb-1f66-236f6b48ab3f


@testset "adjmat" begin
    @testset "undirected" begin
        adj = [0 1 0 1;
               1 0 1 0;
               0 1 0 1;
               1 0 1 0]
        deg = [2 0 0 0;
               0 2 0 0;
               0 0 2 0;
               0 0 0 2]
        lap = [2 -1 0 -1;
               -1 2 -1 0;
               0 -1 2 -1;
               -1 0 -1 2]
        norm_lap = [1. -.5 0. -.5;
                   -.5 1. -.5 0.;
                   0. -.5 1. -.5;
                   -.5 0. -.5 1.]
        scaled_lap = [0 -0.5 0 -0.5;
                      -0.5 0 -0.5 0;
                      0 -0.5 0 -0.5;
                      -0.5 0 -0.5 0]
        rw_lap = [1 -0.5 0 -0.5;
                  -0.5 1 -0.5 0;
                  0 -0.5 1 -0.5;
                  -0.5 0 -0.5 1]

        for T in [Int8, Int16, Int32, Int64, Int128, Float16, Float32, Float64]
            D = GraphLaplacians.degree_matrix(adj, T, dir=:out)
            @test D == T.(deg)
            @test D == GraphLaplacians.degree_matrix(adj, T, dir=:in)
            @test D == GraphLaplacians.degree_matrix(adj, T, dir=:both)
            @test eltype(D) == T

            L = Graphs.laplacian_matrix(adj, T)
            @test L == T.(lap)
            @test eltype(L) == T

            SL = GraphLaplacians.signless_laplacian(adj, T)
            @test SL == T.(adj + deg)
            @test eltype(SL) == T
        end

        for T in [Float16, Float32, Float64]
            NL = GraphLaplacians.normalized_laplacian(adj, T)
            @test NL ≈ T.(norm_lap)
            @test eltype(NL) == T

            SL = GraphLaplacians.scaled_laplacian(adj, T)
            @test SL ≈ T.(scaled_lap)
            @test eltype(SL) == T

            RW = GraphLaplacians.random_walk_laplacian(adj, T)
            @test RW == T.(rw_lap)
            @test eltype(RW) == T
        end
    end

    @testset "directed" begin
        adj = [0 2 0 3;
               0 0 4 0;
               2 0 0 1;
               0 0 0 0]
        degs = Dict(
            :out  => diagm(0=>[2, 2, 4, 4]),
            :in   => diagm(0=>[5, 4, 3, 0]),
            :both => diagm(0=>[7, 6, 7, 4]),
        )
        laps = Dict(
            :out  => degs[:out] - adj,
            :in   => degs[:in] - adj,
            :both => degs[:both] - adj,
        )
        norm_laps = Dict(
            :out  => I - diagm(0=>[1/2, 1/2, 1/4, 1/4])*adj,
            :in   => I - diagm(0=>[1/5, 1/4, 1/3, 0])*adj,
        )
        sig_laps = Dict(
            :out  => degs[:out] + adj,
            :in   => degs[:in] + adj,
            :both => degs[:both] + adj,
        )
        rw_laps = Dict(
            :out  => I - diagm(0=>[1/2, 1/2, 1/4, 1/4]) * adj,
            :in   => I - diagm(0=>[1/5, 1/4, 1/3, 0]) * adj,
            :both => I - diagm(0=>[1/7, 1/6, 1/7, 1/4]) * adj,
        )

        for T in [Int8, Int16, Int32, Int64, Int128, Float16, Float32, Float64]
            for dir in [:out, :in, :both]
                D = GraphLaplacians.degree_matrix(adj, T, dir=dir)
                @test D == T.(degs[dir])
                @test eltype(D) == T
            end
            @test_throws DomainError GraphLaplacians.degree_matrix(adj, dir=:other)

            for dir in [:out, :in, :both]
                L = Graphs.laplacian_matrix(adj, T, dir=dir)
                @test L == T.(laps[dir])
                @test eltype(L) == T
            end

            for dir in [:out, :in, :both]
                SL = GraphLaplacians.signless_laplacian(adj, T, dir=dir)
                @test SL == T.(sig_laps[dir])
                @test eltype(SL) == T
            end
        end

        for T in [Float32, Float64]
            for dir in [:out, :in]
                L = GraphLaplacians.normalized_laplacian(adj, T, dir=dir)
                @test L == T.(norm_laps[dir])
                @test eltype(L) == T
            end

            for dir in [:out, :in, :both]
                RW = GraphLaplacians.random_walk_laplacian(adj, T, dir=dir)
                @test RW == T.(rw_laps[dir])
                @test eltype(RW) == T
            end
        end
    end
end


## permkronecker ##

# Kronecker product for matrices induces a Kronecker product on permutations.
"""
    permkron(p, q)

Compute the Kronecker product of permutations `p` and `q`
induced by their matrix representations.
"""
function permkron(p::PList{T}, q::PList{S}) where {T, S}
    np = length(p); nq = length(q)
    dc = Array{promote_type(T, S)}(undef, np*nq)
@inbounds for i in 1:np
    for k in 1:nq
            dc[nq*(i-1) + k] = nq*(p[i]-1)+q[k]
        end
    end
    dc
end

# This is partly broken
function permkron(pt::PDict{T}, qt::PDict) where {T<:Real}
    dout = Dict{T, T}()
    p = getdata(pt); q = getdata(qt)
    (isempty(p) || isempty(q)) && error("Can't yet do sparse kronecker product with identity permutations")
    np = convert(T,maximum(p)[1])
    nq = convert(T,maximum(q)[1])
    maxk = zero(T)
    #    for (i,pi) in p
    for i in 1:np
        pi = get(p,i,zero(T))
        pi = pi == 0 ? i : pi
        #        for (k,qk) in q
        for k in 1:nq
            qk = get(q,k,zero(T))
            qk = qk == 0 ? k : qk
            ind1 = nq *(pi-1) + qk
            ind2 = nq * (i-1) + k
            ind1 == ind2 ? continue : nothing
            z = get(dout,ind2,zero(T))
            z != zero(T) && continue
            dout[ind2] = ind1
            ind2 > maxk ? maxk = ind2 : nothing
            ind1 > maxk ? maxk = ind1 : nothing  # should be superfluous
        end
    end
    return PDict(dout, permlength(pt) * permlength(qt))
end

# This a bit more efficient than full matrix multiplication. Cuts out one loop.
function permkron(a::PList{T}, b::PMatrix{S}) where {T<:Real,S<:Real}
    (nrowa, ncola) = (length(a),length(a))
    (nrowb, ncolb) = size(b)
    R = zeros(promote_type(T,S), nrowa * nrowb, ncola * ncolb)
    d = invperm(a)
    for j = 1:ncola, l = 1:ncolb
        soff = ncola * (j-1)
        i = d[j]
        roff = ncola * (i-1)
        for k = 1:nrowb
            R[roff+k,soff+l] = b[k,l]
        end
    end
    return R
end

# For testing. If there is a difference, it is small
# function permkron2{T<:Real,S<:Real}(a::AbstractVector{T}, b::AbstractMatrix{S})
#     (nrowa, ncola) = size(a)
#     (nrowb, ncolb) = size(a)
#     R = zeros(promote_type(T,S), nrowa * nrowb, ncola * ncolb)
#     d = a.data
#     for j = 1:ncola, l = 1:ncolb
#         soff = ncola * (j-1)
#         i = d[j]
#         roff = ncola * (i-1)
#         for k = 1:nrowb
#             R[soff+l,roff+k] = b[l,k]
#         end
#     end
#     R
# end


type ConfusionMatrix
  classes::Vector
  matrix::Matrix{Int}
  accuracy::FloatingPoint
  kappa::FloatingPoint
end

function show(io::IO, cm::ConfusionMatrix)
  print(io, "Classes: ")
  show(io, cm.classes)
  print(io, "\nMatrix: ")
  display(cm.matrix)
  print(io, "\nAccuracy: ")
  show(io, cm.accuracy)
  print(io, "\nκ: ")
  show(io, cm.kappa)
end

## measures ##

function confusion_matrix(actual::Vector, predicted::Vector)
    @assert length(actual) == length(predicted)
    N = length(actual)
    _actual = zeros(Int,N)
    _predicted = zeros(Int,N)
    classes = sort(unique([actual; predicted]))
    N = length(classes)
    for i in 1:N
        _actual[actual .== classes[i]] = i
        _predicted[predicted .== classes[i]] = i
    end
    CM = zeros(Int,N,N)
    for i in zip(_actual, _predicted)
        CM[i[1],i[2]] += 1
    end
    accuracy = trace(CM) / sum(CM)
    prob_chance = (sum(CM,1) * sum(CM,2))[1] / sum(CM)^2
    kappa = (accuracy - prob_chance) / (1.0 - prob_chance)
    return ConfusionMatrix(classes, CM, accuracy, kappa)
end

function mean_squared_error(actual, predicted)
    @assert length(actual) == length(predicted)
    return mean((actual - predicted).^2)
end

function R2(actual, predicted)
    @assert length(actual) == length(predicted)
    ss_residual = sum((actual - predicted).^2)
    ss_total = sum((actual .- mean(actual)).^2)
    return 1.0 - ss_residual/ss_total
end

## helper functions ##

function _kfoldCV{T<:FloatingPoint, U<:Real}(regressor::Symbol, labels::Vector{T}, features::Matrix{U}, nfolds::Integer,
                                             n_components::Integer, nferns::Integer)
    N = length(labels)
    ntest = _int(floor(N / nfolds))
    inds = randperm(N)
    R2s = zeros(nfolds)
    for i in 1:nfolds
        test_inds = falses(N)
        test_inds[(i - 1) * ntest + 1 : i * ntest] = true
        train_inds = !test_inds
        test_features = features[inds[test_inds],:]
        test_labels = labels[inds[test_inds]]
        train_features = features[inds[train_inds],:]
        train_labels = labels[inds[train_inds]]
        if regressor == :fern
            model = build_fern(train_labels, train_features, n_components)
            predictions, w = apply_fern(model, test_features)
        elseif regressor == :ensemble
            model = build_ensemble(train_labels, train_features, n_components, nferns)
            predictions = apply_ensemble(model, test_features)
        end

        err = mean_squared_error(test_labels, predictions)
        corr = cor(test_labels, predictions)
        r2 = R2(test_labels, predictions)
        R2s[i] = r2

        println("\nFold ", i)
        println("Mean Squared Error:     ", err)
        println("Correlation Coeff:      ", corr)
        println("Coeff of Determination: ", r2)
    end
    println("\nMean Coeff of Determination: ", mean(R2s))
    return R2s
end

kfoldCV_fern{T<:FloatingPoint, U<:Real}(labels::Vector{T}, features::Matrix{U}, nfolds::Integer, n_components::Integer)      = _kfoldCV(:fern, labels, features, nfolds, n_components, 1)
kfoldCV_fernEnsemble{T<:FloatingPoint, U<:Real}(labels::Vector{T}, features::Matrix{U}, nfolds::Integer, n_components::Integer, n_ferns::Integer)  = _kfoldCV(:ensemble, labels, features, nfolds, n_components, n_ferns)


function wider_normal_model(t :: Trace, data :: Vector{Float64})
    loc = sample(t, :loc, Normal(0.0, 10.0))
    scale = sample(t, :scale, LogNormal())
    for i in 1:length(data)
        observe(t, (:obs, i), Normal(loc, scale), data[i])
    end
end

@testset "prior metropolis proposal 1" begin
    data = randn(100) .+ 4.0
    t = trace()
    wider_normal_model(t, data)
    @info "True loc = 4.0"
    @info "True std = 1.0"

    locs = []
    scales = []
    lls = []

    for i in 1:2500
        t = mh_step(t, wider_normal_model; params = (data,))
        push!(locs, t[:loc].value)
        push!(scales, t[:scale].value)
        push!(lls, loglikelihood(t))
    end

    @info "inferred E[loc] = $(mean(locs[1000:end]))"
    @info "inferred E[scale] = $(mean(scales[1000:end]))"
    @info "approximate p(x) = sum_z p(x|z) = $(mean(lls[1000:end]))"
end

@testset "prior metropolis proposal 2" begin
    data = randn(10) .+ 4.0
    @info "True loc = 4.0"
    @info "True std = 1.0"
    results = mh(wider_normal_model; params=(data,))
    mean_loc = mean(results, :loc)
    std_loc = std(results, :loc)
    @info "Mean loc = $mean_loc"
    @info "Std loc = $std_loc"
end

loc_proposal(old_t :: Trace, new_t :: Trace, data) = propose(new_t, :loc, Normal(old_t[:loc].value, 0.25))
scale_proposal(old_t :: Trace, new_t :: Trace, data) = propose(new_t, :scale, truncated(Normal(old_t[:scale].value, 0.25), 0.0, Inf))

@testset "general metropolis proposal 1" begin
    data = randn(100) .+ 4.0
    t = trace()
    wider_normal_model(t, data)

    locs = []
    scales = []
    lls = []

    for i in 1:2500
        t = mh_step(t, wider_normal_model, loc_proposal; params = (data,))
        t = mh_step(t, wider_normal_model, scale_proposal; params = (data,))
        push!(locs, t[:loc].value)
        push!(scales, t[:scale].value)
        push!(lls, loglikelihood(t))
    end

    @info "inferred E[loc] = $(mean(locs[1000:end]))"
    @info "inferred E[scale] = $(mean(scales[1000:end]))"
    @info "approximate p(x) = sum_z p(x|z) = $(mean(lls[1000:end]))"
end

@testset "generic metropolis proposal 2" begin
    data = randn(100) .+ 4.0
    @info "Using full results struct"
    results = mh(wider_normal_model, [loc_proposal, scale_proposal]; params = (data,), inverse_verbosity = Inf)
    @info "Using bare results struct"
    results = mh(wider_normal_model, [loc_proposal, scale_proposal], (:loc,); params = (data,), inverse_verbosity = Inf)
end

function random_sum_model(t :: Trace, data)
    n = sample(t, :n, Geometric(0.1))
    loc = 0.0
    for i in 1:(n + 1)
        loc += sample(t, (:loc, i), Normal())
    end
    obs = Array{Float64, 1}()
    for j in 1:length(data)
        o = observe(t, (:data, j), Normal(loc, 1.0), data[j])
        push!(obs, o)
    end
    obs
end

function random_n_proposal(old_trace, new_trace, params...)
    old_n = float(old_trace[:n].value)
    if old_n > 0
        propose(new_trace, :n, Poisson(old_n))
    else
        propose(new_trace, :n, Poisson(1.0))
    end
end

function gen_loc_proposal(old_trace, new_trace, ix, params...)
    propose(new_trace, (:loc, ix), Normal(old_trace[(:loc, ix)].value, 0.25))
end

@testset "generic metropolis proposal 3" begin
    t = trace()
    data = random_sum_model(t, fill(nothing, 25))
    @info "True :n = $(t[:n].value)"
    random_sum_model(t, data)

    ns = []
    loc_proposals = [
        (o, n, params...) -> gen_loc_proposal(o, n, i, params...) for i in 1:100
    ]

    for i in 1:5000
        t = mh_step(t, random_sum_model, random_n_proposal; params=(data,))
        for j in 1:(t[:n].value + 1)
            t = mh_step(t, random_sum_model, loc_proposals[j]; params=(data,))
        end
        push!(ns, t[:n].value)
    end

    @info "Posterior E[:n] = $(mean(ns[1000:end]))"
end

const Maybe{T} = Union{T, Nothing}

function eye(d :: Int)
    mat = zeros(Float64, d, d)
    for i in 1:d
        mat[i, i] = 1.0
    end
    mat
end

function linear_model!(t :: Trace, out :: Vector{Float64}, X :: Matrix{Float64}, y :: Maybe{Vector{Float64}})
    size_X = size(X)
    D = size_X[1]
    N = size_X[2]
    β = sample(t, :β, MvNormal(zeros(D), eye(D)))
    scale = sample(t, :scale, LogNormal(0.0, 1.0))
    loc = X' * β
    if y === nothing
        for n in 1:N
            out[n] = sample(t, (:y, n), Normal(loc[n], scale))
        end
    else
        for n in 1:N
            out[n] = observe(t, (:y, n), Normal(loc[n], scale), y[n])
        end
    end
end

beta_proposal(old_t :: Trace, new_t :: Trace, _, _, _) = propose(new_t, :β, MvNormal(old_t[:β].value, 0.5 .* eye(length(old_t[:β].value))))
scale_proposal(old_t :: Trace, new_t :: Trace, _, _, _) = propose(new_t, :scale, truncated(Normal(old_t[:scale].value, 0.25), 0.0, Inf))

@testset "generic metropolis proposal 4" begin
    D = 10
    N = 200
    X = randn(D, N)
    out = Vector{Float64}(undef, N)
    @info "Using model to generate dataset"
    t = trace()
    linear_model!(t, out, X, nothing)
    y = [t[(:y, n)].value for n in 1:N]
    β_true = t[:β].value
    @info "True regression coefficients: $β_true"

    @info "Conducting inference using user-defined Metropolis kernels"
    # smoketest only --- low number of iterations. see examples/basic.jl for actual inference results
    @time results = mh(
        linear_model!,
        [beta_proposal, scale_proposal];
        params = (out, X, y),
        burn = 50,
        thin = 10,
        num_iterations = 500
    )
    β_mean = mean(results, :β)
    β_std = std(results, :β)
    @info "Posterior mean regression coefficients = $β_mean"
    @info "Posterior std regression coefficients = $β_std"
end


# This file replicates the testsuite provided by MINPACK for the hybrj
# function. It consists in 14 functions with their analytical Jacobian.
# Julia replication taken from NLSOLVE.jl package
using LeastSquaresOptim, Printf,  LinearAlgebra, SparseArrays, Test

function rosenbrock()
    name = "rosenbrock"
    function f!(fvec, x)
        fvec[1] = 1 - x[1]
        fvec[2] = 10(x[2]-x[1]^2)
    end
    function g!(fjac, x)
        fjac[1,1] = -1
        fjac[1,2] = 0
        fjac[2,1] = -20x[1]
        fjac[2,2] = 10
    end
    x = [-1.2; 1.]
    return name, f!, g!, x
end

function powell_singular()
    name = "powell_singular"
    function f!(fvec, x)
        fvec[1] = x[1] + 10x[2]
        fvec[2] = sqrt(5)*(x[3] - x[4])
        fvec[3] = (x[2] - 2x[3])^2
        fvec[4] = sqrt(10)*(x[1] - x[4])^2
    end
    function g!(fjac, x)
        fill!(fjac, 0)
        fjac[1,1] = 1
        fjac[1,2] = 10
        fjac[2,3] = sqrt(5)
        fjac[2,4] = -fjac[2,3]
        fjac[3,2] = 2(x[2] - 2x[3])
        fjac[3,3] = -2fjac[3,2]
        fjac[4,1] = 2sqrt(10)*(x[1] - x[4])
        fjac[4,4] = -fjac[4,1]
    end
    x = [3.; -1.; 0.; 1.]
    return name, f!, g!, x
end

function powell_badly_scaled()
    name = "powell_badly_scaled"
    c1 = 1e4
    c2 = 1.0001
    function f!(fvec, x)
        fvec[1] = c1*x[1]*x[2] - 1
        fvec[2] = exp(-x[1]) + exp(-x[2]) - c2
    end
    function g!(fjac, x)
        fjac[1,1] = c1*x[2]
        fjac[1,2] = c1*x[1]
        fjac[2,1] = -exp(-x[1])
        fjac[2,2] = -exp(-x[2])
    end
    x = [0.; 1.]
    return name, f!, g!, x
end

function wood()
    name = "wood"
    c3 = 2e2
    c4 = 2.02e1
    c5 = 1.98e1
    c6 = 1.8e2

    function f!(fvec, x)
        temp1 = x[2] - x[1]^2
        temp2 = x[4] - x[3]^2
        fvec[1] = -c3*x[1]*temp1 - (1 - x[1])
        fvec[2] = c3*temp1 + c4*(x[2] - 1) + c5*(x[4] - 1)
        fvec[3] = -c6*x[3]*temp2 - (1 - x[3])
        fvec[4] = c6*temp2 + c4*(x[4] - 1) + c5*(x[2] - 1)
    end

    function g!(fjac, x)
        fill!(fjac, 0)
        temp1 = x[2] - 3x[1]^2
        temp2 = x[4] - 3x[3]^2
        fjac[1,1] = -c3*temp1 + 1
        fjac[1,2] = -c3*x[1]
        fjac[2,1] = -2*c3*x[1]
        fjac[2,2] = c3 + c4
        fjac[2,4] = c5
        fjac[3,3] = -c6*temp2 + 1
        fjac[3,4] = -c6*x[3]
        fjac[4,2] = c5
        fjac[4,3] = -2*c6*x[3]
        fjac[4,4] = c6 + c4
    end
    x = [-3.; -1; -3; -1]
    return name, f!, g!, x
end

function helical_valley()
    name = "helical_valley"
    tpi = 8*atan(1)
    c7 = 2.5e-1
    c8 = 5e-1

    function f!(fvec, x)
        if x[1] > 0
            temp1 = atan(x[2]/x[1])/tpi
        elseif x[1] < 0
            temp1 = atan(x[2]/x[1])/tpi + c8
        else
            temp1 = c7*sign(x[2])
        end
        temp2 = sqrt(x[1]^2+x[2]^2)
        fvec[1] = 10(x[3] - 10*temp1)
        fvec[2] = 10(temp2 - 1)
        fvec[3] = x[3]
    end

    function g!(fjac, x)
        temp = x[1]^2 + x[2]^2
        temp1 = tpi*temp
        temp2 = sqrt(temp)
        fjac[1,1] = 100*x[2]/temp1
        fjac[1,2] = -100*x[1]/temp1
        fjac[1,3] = 10
        fjac[2,1] = 10*x[1]/temp2
        fjac[2,2] = 10*x[2]/temp2
        fjac[2,3] = 0
        fjac[3,1] = 0
        fjac[3,2] = 0
        fjac[3,3] = 1
    end
    x = [-1.; 0; 0]
    return name, f!, g!, x
end

function watson(n::Integer)
    name = "watson"
    c9 = 2.9e1

    function f!(fvec, x)
        fill!(fvec, 0)
        for i = 1:29
            ti = i/c9
            sum1 = 0.0
            temp = 1.0
            for j = 2:n
                sum1 += (j-1)*temp*x[j]
                temp *= ti
            end
            sum2 = 0.0
            temp = 1.0
            for j = 1:n
                sum2 += temp*x[j]
                temp *= ti
            end
            temp1 = sum1 - sum2^2 - 1
            temp2 = 2*ti*sum2
            temp = 1/ti
            for k = 1:n
                fvec[k] += temp*(k-1-temp2)*temp1
                temp *= ti
            end
        end
        temp = x[2] - x[1]^2 - 1
        fvec[1] += x[1]*(1-2temp)
        fvec[2] += temp
    end

    function g!(fjac, x)
        fill!(fjac, 0)
        for i = 1:29
            ti = i/c9
            sum1 = 0.0
            temp = 1.0
            for j = 2:n
                sum1 += (j-1)*temp*x[j]
                temp *= ti
            end
            sum2 = 0.0
            temp = 1.0

            for j = 1:n
                sum2 += temp*x[j]
                temp *= ti
            end
            temp1 = 2(sum1 - sum2^2 - 1)
            temp2 = 2*sum2
            temp = ti^2
            tk = 1.0

            for k = 1:n
                tj = tk
                for j = k:n
                    fjac[k,j] += tj*((Float64(k-1)/ti - temp2)*(Float64(j-1)/ti - temp2) - temp1)
                    tj *= ti
                end
                tk *= temp
            end
        end
        fjac[1,1] += 6x[1]^2 - 2x[2] + 3
        fjac[1,2] -= 2x[1]
        fjac[2,2] += 1
        for k = 1:n
            for j = k:n
                fjac[j,k] = fjac[k,j]
            end
        end
    end
    x = zeros(n)
    return name, f!, g!, x
end

function chebyquad(n::Integer)
    name = "chebyquad"
    tk = 1/n

    function f!(fvec, x)
        fill!(fvec, 0)
        for j = 1:n
            temp1 = 1.0
            temp2 = 2x[j]-1
            temp = 2temp2
            for i = 1:n
                fvec[i] += temp2
                ti = temp*temp2 - temp1
                temp1 = temp2
                temp2 = ti
            end
        end
        iev = -1.0
        for k = 1:n
            fvec[k] *= tk
            if iev > 0
                fvec[k] += 1/(k^2-1)
            end
            iev = -iev
        end
    end

    function g!(fjac, x)
        for j = 1:n
            temp1 = 1.
            temp2 = 2x[j] - 1
            temp = 2*temp2
            temp3 = 0.0
            temp4 = 2.0
            for k = 1:n
                fjac[k,j] = tk*temp4
                ti = 4*temp2 + temp*temp4 - temp3
                temp3 = temp4
                temp4 = ti
                ti = temp*temp2 - temp1
                temp1 = temp2
                temp2 = ti
            end
        end
    end
    x = collect(1:n)/(n+1)
    return name, f!, g!, x
end

function brown_almost_linear(n::Integer)
    name = "brown_almost_linear"
    function f!(fvec, x)
        sum1 = sum(x) - (n+1)
        for k = 1:(n-1)
            fvec[k] = x[k] + sum1
        end
        fvec[n] = prod(x) - 1
    end

    function g!(fjac, x)
        fill!(fjac, 1)
        fjac[diagind(fjac)] .= 2
        prd = prod(x)
        for j = 1:n
            if x[j] == 0.0
                fjac[n,j] = 1.
                for k = 1:n
                    if k != j
                        fjac[n,j] *= x[k]
                    end
                end
            else
                fjac[n,j] = prd/x[j]
            end
        end
    end
    x = 0.5*ones(n)
    return name, f!, g!, x
end

function discrete_boundary_value(n::Integer)
    name = "discrete_boundary_value"
    h = 1/(n+1)

    function f!(fvec, x)
        for k = 1:n
            temp = (x[k] + k*h + 1)^3
            if k != 1
                temp1 = x[k-1]
            else
                temp1 = 0.0
            end
            if k != n
                temp2 = x[k+1]
            else
                temp2 = 0.0
            end
            fvec[k] = 2x[k] - temp1 - temp2 + temp*h^2/2
        end
    end

    function g!(fjac, x)
        for k = 1:n
            temp = 3*(x[k]+k*h+1)^2
            for j = 1:n
                fjac[k,j] = 0
            end
            fjac[k,k] = 2 + temp*h^2/2
            if k != 1
                fjac[k,k-1] = -1
            end
            if k != n
                fjac[k,k+1] = -1
            end
        end
    end
    x = collect(1:n)*h
    x = x .* (x .- 1)
    return name, f!, g!, x
end

function discrete_integral_equation(n::Integer)
    name = "discrete_integral_equation"
    h = 1/(n+1)

    function f!(fvec, x)
        for k = 1:n
            tk = k*h
            sum1 = 0.0
            for j = 1:k
                tj = j*h
                sum1 += tj*(x[j] + tj + 1)^3
            end
            sum2 = 0.0
            kp1 = k+1
            if n >= kp1
                for j = kp1:n
                    tj = j*h
                    sum2 += (1-tj)*(x[j] + tj + 1)^3
                end
            end
            fvec[k] = x[k] + h*((1-tk)*sum1 + tk*sum2)/2
        end
    end

    function g!(fjac, x)
        for k = 1:n
            tk = k*h
            for j = 1:n
                tj = j*h
                fjac[k,j] = h*min(tj*(1-tk), tk*(1-tj))*3(x[j] + tj + 1)^2/2
            end
            fjac[k,k] += 1
        end
    end

    x = collect(1:n)*h
    x = x .* (x .- 1)

    return name, f!, g!, x
end

function trigonometric(n::Integer)
    name = "trigonometric"
    function f!(fvec, x)
        for j = 1:n
            fvec[j] = cos(x[j])
        end
        sum1 = sum(fvec)
        for k = 1:n
            fvec[k] = n+k-sin(x[k]) - sum1 - k*fvec[k]
        end
    end

    function g!(fjac, x)
        for j = 1:n
            temp = sin(x[j])
            for k = 1:n
                fjac[k,j] = temp
            end
            fjac[j,j] = (j+1)*temp - cos(x[j])
        end
    end
    x = ones(n)/n
    return name, f!, g!, x
end

function variably_dimensioned(n::Integer)
    name = "variably_dimensioned"
    function f!(fvec, x)
        sum1 = 0.0
        for j = 1:n
            sum1 += j*(x[j]-1)
        end
        temp = sum1*(1+2sum1^2)
        for k = 1:n
            fvec[k] = x[k] - 1 + k*temp
        end
    end

    function g!(fjac, x)
        sum1 = 0.0
        for j = 1:n
            sum1 += j*(x[j]-1)
        end
        temp = 1 + 6sum1^2
        for k = 1:n
            for j = k:n
                fjac[k,j] = k*j*temp
                fjac[j,k] = fjac[k,j]
            end
            fjac[k,k] += 1
        end
    end
    x = collect(1:n)/n
    return name, f!, g!, x
end

function broyden_tridiagonal(n::Integer)
    name = "broyden_tridiagonal"
    function f!(fvec, x)
        for k = 1:n
            temp = (3-2x[k])*x[k]
            if k != 1
                temp1 = x[k-1]
            else
                temp1 = 0.0
            end
            if k != n
                temp2 = x[k+1]
            else
                temp2 = 0.0
            end
            fvec[k] = temp - temp1 - 2temp2 + 1
        end
    end

    function g!(fjac, x)
        fill!(fjac, 0)
        for k = 1:n
            fjac[k,k] = 3-4x[k]
            if k != 1
                fjac[k,k-1] = -1
            end
            if k != n
                fjac[k,k+1] = -2
            end
        end
    end
    x = -ones(n)
    return name, f!, g!, x
end

function broyden_banded(n::Integer)
    name = "broyden_banded"
    ml = 5
    mu = 1

    function f!(fvec, x)
        for k = 1:n
            k1 = max(1, k-ml)
            k2 = min(k+mu, n)
            temp = 0.0
            for j = k1:k2
                if j != k
                    temp += x[j]*(1+x[j])
                end
            end
            fvec[k] = x[k]*(2+5x[k]^2) + 1 - temp
        end
    end

    function g!(fjac, x)
        fill!(fjac, 0)
        for k = 1:n
            k1 = max(1, k-ml)
            k2 = min(k+mu, n)
            for j = k1:k2
                if j != k
                    fjac[k,j] = -(1+2x[j])
                end
            end
            fjac[k,k] = 2+15x[k]^2
        end
    end
    x = -ones(n)
    return name, f!, g!, x
end

for matrix in (:dense, :sparse)
    for (solver, solver_abbr) in ((LeastSquaresOptim.QR(), :qr), (LeastSquaresOptim.LSMR(), :lsmr))
        if (matrix == :sparse && typeof(solver) <: LeastSquaresOptim.QR)
            # tests below
            continue
        else
            for (optimizer, optimizer_abbr) in ((LeastSquaresOptim.Dogleg, :dl), (LeastSquaresOptim.LevenbergMarquardt, :lm))
                alltests = [rosenbrock();
                powell_singular(); powell_badly_scaled();
                wood();
                helical_valley();
                watson(6); watson(9);
                chebyquad(5); chebyquad(6); chebyquad(7); chebyquad(9);
                brown_almost_linear(10); brown_almost_linear(30); brown_almost_linear(40);
                discrete_boundary_value(10);
                discrete_integral_equation(1); discrete_integral_equation(10);
                trigonometric(10); variably_dimensioned(10);
                broyden_tridiagonal(10); broyden_banded(10)]
                for (name, f!, g!, x) in alltests
                    fcur = similar(x)
                    J = Array{Float64}(undef, length(x), length(x))
                    if matrix == :sparse
                        J = sparse(J)
                    end
                    nls = LeastSquaresProblem(x = x, y = fcur, f! = f!, J = J, g! = g!)
                    r = optimize!(nls, optimizer(solver))
                    @printf("%-6s %4s %2s %30s %5d %5d   %5d   %10e\n", matrix, solver_abbr, optimizer_abbr, name, r.iterations, r.f_calls, r.g_calls, r.ssr)
                    @test r.ssr <= 1e-3
                end
            end
        end
    end
end

# test sparse factorization
alltests =  [rosenbrock();
powell_singular(); powell_badly_scaled();
wood();
helical_valley();
watson(6); watson(9);
chebyquad(5); chebyquad(6); chebyquad(7); chebyquad(9);
brown_almost_linear(10); brown_almost_linear(30); brown_almost_linear(40);
discrete_boundary_value(10);
discrete_integral_equation(1); discrete_integral_equation(10);
trigonometric(10); variably_dimensioned(10);
broyden_tridiagonal(10); broyden_banded(10)]

for (name, f!, g!, x) in alltests
    fcur = similar(x)
    J = ones(length(x), length(x))
    J = sparse(J)
    fill!(J.nzval, 0)
    g!(J, x)
    # try because g! may change symbolic factorizations (see Julia issue #9906)
    try
        nls = LeastSquaresProblem(x = x, y = fcur, f! = f!, J = J, g! = g!)
        r = optimize!(nls, optimizer(solver))
        @test r.converged
        @test r.ssr <= 1e-3
        @printf("%-6s %4s %2s %30s %5d %5d   %5d   %10e\n", :sparse, :fact, :dl, name, r.iterations, r.f_calls, r.g_calls, r.ssr)
    catch y
        if !isa(y, Exception)
            throw(y)
        end
    end
end

alltests = [rosenbrock();
    powell_singular(); powell_badly_scaled();
    wood();
    helical_valley();
    watson(6);
    chebyquad(5); chebyquad(6); chebyquad(7); chebyquad(9);
    brown_almost_linear(10);# brown_almost_linear(30);
    discrete_boundary_value(10);
    discrete_integral_equation(1); discrete_integral_equation(10);
    trigonometric(10); variably_dimensioned(10);
    broyden_tridiagonal(10); broyden_banded(10)]
# test dense cholesky
for (optimizer, optimizer_abbr) in ((LeastSquaresOptim.Dogleg, :dl), (LeastSquaresOptim.LevenbergMarquardt, :lm))
    for (name, f!, g!, x) in alltests
        fcur = similar(x)
        J = Array{Float64}(undef, length(x), length(x))
        nls = LeastSquaresProblem(x = x, y = fcur, f! = f!, J = J, g! = g!)
        r = optimize!(nls, optimizer(LeastSquaresOptim.Cholesky()))
        @printf("%-6s %4s %2s %30s %5d %5d   %5d   %10e\n", :dense, :chol, optimizer_abbr, name, r.iterations, r.f_calls, r.g_calls, r.ssr)
        @test r.converged
        @test r.ssr <= 1e-3
    end
end

alltests = [rosenbrock();
  powell_singular(); powell_badly_scaled();
  wood();
  helical_valley();
  watson(6); watson(9);
  chebyquad(5); chebyquad(6); chebyquad(7); chebyquad(9);
  brown_almost_linear(10); brown_almost_linear(30); brown_almost_linear(40);
  discrete_boundary_value(10);
  discrete_integral_equation(1); discrete_integral_equation(10);
  trigonometric(10); variably_dimensioned(10);
  broyden_tridiagonal(10); broyden_banded(10)]
#test forwarddiff
for (optimizer, optimizer_abbr) in ((LeastSquaresOptim.Dogleg, :dl), (LeastSquaresOptim.LevenbergMarquardt, :lm))
    for (name, f!, g!, x) in alltests
        nls = LeastSquaresProblem(x = x, f! = f!, output_length = length(x))
        r = optimize!(nls, optimizer())
        @printf("%-10s %4s %2s %30s %5d %5d   %5d   %10e\n", :forwarddiff, :fact, optimizer_abbr, name, r.iterations, r.f_calls, r.g_calls, r.ssr)
        @test r.converged
        @test r.ssr <= 1e-3
    end
end

# check default are correctly
name, f!, g!, x = wood()
fcur = similar(x)
J = ones(length(x), length(x))
nls = LeastSquaresProblem(x = x, y = fcur, f! = f!, J = J, g! = g!)
result = optimize!(nls, show_trace = true)
@test result.optimizer == "Dogleg"
nls = LeastSquaresProblem(x = x, y = fcur, f! = f!, J = sparse(J), g! = g!)
result = optimize!(nls, show_trace = true)
@test result.optimizer == "LevenbergMarquardt"

# bigfloat
function newf(x)
    out = similar(x)
    out[1] = 1 - x[1]
    out[2] = 10(x[2]-x[1]^2)
    return out
end
nls = optimize(newf, [BigFloat(0.0), BigFloat(0.0)], Dogleg())
nls = optimize(newf, [BigFloat(0.0), BigFloat(0.0)], LevenbergMarquardt())


@testset "bicgstb" begin
	@testset "real matrix" begin
		A  = sprandn(100,100,.1) + 10*I
		D  = Array(diag(A))
		Af = x -> A*x
		M  = x -> D.\x
		rhs = randn(100)

		# test flag for early stopping
		xt = bicgstb(A,rhs,tol=1e-6,maxIter=3,out=2,storeInterm=true)
		@test xt[2]==-1

		# test handling of zero rhs
		xt = bicgstb(A,zeros(size(A,2)),tol=1e-6,maxIter=3,out=2)
		@test  all(xt[1].==0)
		@test xt[2]==-9

		x0 = bicgstb(Matrix(A),rhs,tol=1e-6)
		x1 = bicgstb(A,rhs,tol=1e-6,out=1)
		x2 = bicgstb(Af,rhs,tol=1e-6)
		x3 = bicgstb(Af,rhs,tol=1e-6,maxIter=100,x=randn(size(rhs)))
		x5 = bicgstb(Af,rhs,tol=1e-6,maxIter=100,M1=M)

		@test norm(A*x0[1]-rhs)/norm(rhs) < 1e-6
		@test norm(A*x1[1]-rhs)/norm(rhs) < 1e-6
		@test norm(A*x2[1]-rhs)/norm(rhs) < 1e-6
		@test norm(A*x3[1]-rhs)/norm(rhs) < 1e-6
		@test norm(A*x5[1]-rhs)/norm(rhs) < 1e-6
		@test norm(x2[1]-x1[1])/norm(x1[1]) < 1e-12
		@test norm(x3[1]-x1[1])/norm(x1[1]) < 1e-5
		@test norm(x0[1]-x1[1])/norm(x1[1]) < 1e-5
	end
	@testset "complex matrix" begin
		A  = sprandn(100,100,.1) + 10*I + im*(sprandn(100,100,.1) + 10*I )
		D  = Array(diag(A))
		Af = x -> A*x
		M  = x -> D.\x
		rhs = randn(100) + 1im * randn(100)

		x1 = bicgstb(A,rhs,tol=1e-6)
		x2 = bicgstb(Af,rhs,tol=1e-6)
		x3 = bicgstb(Af,rhs,tol=1e-6,maxIter=200,x=randn(size(rhs))+im*randn(size(rhs)))
		x4 = bicgstb(Af,rhs,tol=1e-6,maxIter=200,M1=M)

		@test norm(A*x1[1]-rhs)/norm(rhs) < 1e-6
		@test norm(A*x2[1]-rhs)/norm(rhs) < 1e-6
		@test norm(A*x3[1]-rhs)/norm(rhs) < 1e-6
		@test norm(A*x4[1]-rhs)/norm(rhs) < 1e-6
		@test norm(x2[1]-x1[1])/norm(x1[1]) < 1e-12
		@test norm(x3[1]-x1[1])/norm(x1[1]) < 1e-5

	end
end


#############################################################################
# sumlargest.jl
# The sum of the k largest/smallest elements of an expression
# All expressions and atoms are subtpyes of AbstractExpr.
# Please read expressions.jl first.
#############################################################################

struct SumLargestAtom <: AbstractExpr
    head::Symbol
    id_hash::UInt64
    children::Tuple{AbstractExpr}
    size::Tuple{Int, Int}
    k::Int

    function SumLargestAtom(x::AbstractExpr, k::Int)
        if sign(x) == ComplexSign()
            error("Argument should be real instead it is $(sign(x))")
        else
            if k <= 0
                error("sumlargest and sumsmallest only support positive values of k")
            end
            if k > length(x)
                error("k cannot be larger than the number of entries in x")
            end
            children = (x,)
            return new(:sumlargest, hash((children, k)), children, (1,1), k)
        end
    end
end

function sign(x::SumLargestAtom)
    return sign(x.children[1])
end

function monotonicity(x::SumLargestAtom)
    return (Nondecreasing(), )
end

function curvature(x::SumLargestAtom)
    return ConvexVexity()
end

function evaluate(x::SumLargestAtom)
    return sum(sort(vec(evaluate(x.children[1])), rev=true)[1:x.k])
end

function conic_form!(x::SumLargestAtom, unique_conic_forms::UniqueConicForms)
    if !has_conic_form(unique_conic_forms, x)
        c = x.children[1]
        t = Variable(size(c))
        q = Variable()
        # sum k largest given by the solution to
        # minimize sum(t) + k*q
        # subject to c <= t + q, t >= 0
        objective = conic_form!(sum(t) + x.k*q, unique_conic_forms)
        conic_form!(c <= t + q, unique_conic_forms)
        conic_form!(t >= 0, unique_conic_forms)
        cache_conic_form!(unique_conic_forms, x, objective)
    end
    return get_conic_form(unique_conic_forms, x)
end

sumlargest(x::AbstractExpr, k::Int) = k == 0 ? Constant(0) : SumLargestAtom(x, k)
sumsmallest(x::AbstractExpr, k::Int) = k == 0 ? Constant(0) : -SumLargestAtom(-x, k)


export
    # Fock basis
    FockState,
    NumberState,
    VacuumState,
    SinglePhotonState,

    # pure state
    CoherentState,
    SqueezedState,

    # mixed state
    ThermalState,
    SqueezedThermalState

################
# constructors #
################

##### Fock basis #####

"""
    FockState([T::Type{<:Number}], n::Integer, [rep::Type]; dim::Integer)

Construct a Fock state in `rep` representation.

## Arguments

* `T`: Numerical type. Default as `Float64` if not assigned.
* `n`: Photon number of Fock state.
* `rep`: In which representation, default as `Vector` if not assigned.
* `dim`: Maximum photon number for truncate.
"""
function FockState(T::Type{<:Number}, n::Integer, ::Type{Vector}; dim)
    v = zeros(T, dim); v[n+1] = one(T)

    return v
end

function FockState(T::Type{<:Number}, n::Integer, ::Type{Matrix}; dim)
    ρ = zeros(T, dim, dim); ρ[n+1, n+1] = one(T)

    return ρ
end

FockState(n, rep; dim) = FockState(Float64, n, rep, dim=dim)

FockState(n; dim) = FockState(n, Vector, dim=dim)

const NumberState = FockState

"""
    VacuumState([T::Type{<:Number}], [rep::Type]; dim::Integer)

Construct a vacuum state in `rep` representation.

## Arguments

* `T`: Numerical type. Default as `Float64` if not assigned.
* `rep`:In which representation, default as `Vector` if not assigned.
* `dim`: Maximum photon number for truncate.
"""
VacuumState(T, rep; dim) = FockState(T, 0, rep, dim=dim)

VacuumState(rep; dim) = FockState(Float64, 0, rep, dim=dim)

VacuumState(; dim) = FockState(0, Vector, dim=dim)

"""
    SinglePhotonState([T::Type{<:Number}], [rep::Type]; dim::Integer)

Construct a single photon state in `rep` representation.

## Arguments

* `T`: Numerical type. Default as `Float64` if not assigned.
* `rep`:In which representation, default as `Vector` if not assigned.
* `dim`: Maximum photon number for truncate.
"""
SinglePhotonState(T, rep; dim) = FockState(T, 1, rep, dim=dim)

SinglePhotonState(rep; dim) = FockState(Float64, 1, rep, dim=dim)

SinglePhotonState(; dim) = FockState(1, Vector, dim=dim)

##### pure state #####

### coherent state

"""
    CoherentState([T::Type{<:Complex}], r::Real, θ::Real, [rep::Type]; dim::Integer)

Construct a coherent state in `rep` representation.

* `T`: Numerical type. Default as `ComplexF64` if not assigned.
* `r`: The radius in phase space.
* `θ`: The phase angle in phase space.
* `rep`:In which representation, default as `Vector` if not assigned.
* `dim`: Maximum photon number for truncate.
"""
function CoherentState(T::Type{<:Complex}, r, θ, rep; dim)
    return displace(VacuumState(real(T), rep, dim=dim), r, θ)
end

CoherentState(r, θ, rep; dim) = CoherentState(ComplexF64, r, θ, rep, dim=dim)

CoherentState(r, θ; dim) = CoherentState(r, θ, Vector, dim=dim)

### squeezed state

"""
    SqueezedState([T::Type{<:Complex}], r::Real, θ::Real, [rep::Type]; dim::Integer)

Construct a squeezed state in `rep` representation.

Squeezed state is defined if its electric field strength for some phases has a quantum
uncertainty smaller than that of a coherent state.

## Arguments

* `T`: Numerical type. Default as `ComplexF64` if not assigned.
* `r`: The squeezing level.
* `θ`: The squeezing phase.
* `rep`:In which representation, default as `Vector` if not assigned.
* `dim`: Maximum photon number for truncate.

This constructor will construct ``| \\xi \\rangle = \\hat{S}(\\xi) | 0 \\rangle``
"""
function SqueezedState(T::Type{<:Complex}, r, θ, rep; dim)
    return squeeze(VacuumState(real(T), rep, dim=dim), r, θ)
end

SqueezedState(r, θ, rep; dim) = SqueezedState(ComplexF64, r, θ, rep, dim=dim)

SqueezedState(r, θ; dim) = SqueezedState(r, θ, Vector, dim=dim)

##### mixed state #####

### thermal state

bose_einstein(n::Number, n̄::Real) = n̄^n / (1 + n̄)^(n+1)

bose_einstein(n̄) = n -> bose_einstein(n, n̄)

"""
    ThermalState([T::Type{<:Number}], n̄::Real; dim::Integer)

Construct a thermal state.

Thermal state is a mixed state with photon number distribution described by
Bose-Einstein distribution.

## Arguments

* `T`: Numerical type. Default as `Float64` if not assigned.
* `n̄`: Average photon number at a certain temperature.
* `dim`: Maximum photon number for truncate.
"""
ThermalState(T::Type{<:Number}, n̄; dim) = diagm(bose_einstein(n̄).(T.(0:dim-1)))

ThermalState(n̄; dim) = ThermalState(Float64, n̄, dim=dim)

### squeezed thermal state

"""
    SqueezedThermalState([T::Type{<:Complex}], r::Real, θ::Real, n̄::Real; dim::Integer)

Construct a squeezed thermal state.

Squeezed thermal state is defined if its electric field strength for some phases has a
quantum uncertainty smaller than that of a coherent state.

## Arguments

* `T`: Numeric type. Default as `ComplexF64` if not assigned.
* `r`: Squeezing level.
* `θ`: Squeezing phase.
* `n̄`: Average photon number at a certain temperature.
* `dim`: Maximum photon number for truncate.

This constructor will construct ``\\rho = \\hat{S}(\\xi) \\rho_{th} \\hat{S}(\\xi)^{T}``
"""
function SqueezedThermalState(T::Type{<:Complex}, r, θ, n̄; dim)
    return squeeze(ThermalState(real(T), n̄, dim=dim), r, θ)
end

SqueezedThermalState(r, θ, n̄; dim) = SqueezedThermalState(ComplexF64, r, θ, n̄, dim=dim)

###########
# methods #
###########

# TODO: ρ(v)
# TODO: purity(ρ)
# TODO: fidelity(ρ1, ρ2)
# TODO: fidelity(v1, v2) ???


module StarGraph
    export H0, Hnoise, Hamiltonian
    """
        H0(N; ν=1.)

    Constructs the noiseless Hamiltonian for the star graph with ``N`` nodes
    """
    function H0(N; ν=1.)
        H = spzeros(N,N)
        H[1,2:end] = ν
        H = H + H'
        H = spdiagm(0 => sum(H,1)[:]) - H
        return H
    end

    """
        Hnoise(N, noise_id)

    Returns the noisy part of the Hamiltonian for the 1D lattice with `N` nodes
    for a specific state of the noise fluctuators specified by `noise_id`.
    """
    function Hnoise(N, noise_id)
        if noise_id < 0 || noise_id > 2^(N-1)
            throw(ArgumentError("noise_id must be between 1 and 2^(N-1)"))
        end

        if noise_id ==0
            return spzeros(N,N)
        end

        noise = 2 * digits(noise_id - 1, base=2, pad=N-1) - 1

        H = spzeros(N,N)
        H[1, 2:end] = noise
        H = H + H'
        H = spdiagm(0 => sum(H,1)[:]) - H
        return H
    end

    """
        Hamiltonian(N, ν, ν0=1.)

    Returns an array of Hamiltonians for all possible configurations of the noise

    # Arguments
        * `N::Integer`: the dimension of the Hilbert space
        * `ν::Float`: the strength of the noise
        * `ν0::Float`: The rate of the noiseless Hamiltonian (ν0 = 1.)
    """
    function Hamiltonian(N::Integer, ν::Float64, ν0=1.)
        return [H0(N;ν=ν0) + ν*Hnoise(N, id) for id = 1:2^(N-1)]
    end

    function H0_spatial_search(N; ν=1., target::Int=1)
        H = spzeros(N,N)
        H[1,2:end] = ν
        H = H + H'
        H = spdiagm(0 => sum(H,1)[:]) - H
        H[target,target] = -1.
        return H
    end

    function Hamiltonian_spatial_search(N::Integer, ν::Float64; ν0=.3, target::Int=1)
        return [H0_spatial_search(N;ν=ν0) + ν*Hnoise(N, id) for id = 1:2^(N-1)]
    end
end


immutable InverseGamma <: ContinuousUnivariateDistribution
    invd::Gamma
    β::Float64

    function InverseGamma(α::Real, β::Real)
        (α > zero(α) && β > zero(β)) || error("Both shape and scale must be positive.")
        new(Gamma(α, 1.0 / β), float64(β))
    end

    InverseGamma(α::Real) = InverseGamma(α, 1.0)
    InverseGamma() = InverseGamma(1.0, 1.0)
end

@distr_support InverseGamma 0.0 Inf

#### Parameters

shape(d::InverseGamma) = shape(d.invd)
scale(d::InverseGamma) = d.β
rate(d::InverseGamma) = scale(d.invd)

params(d::InverseGamma) = (shape(d), scale(d))

#### Parameters

mean(d::InverseGamma) = ((α, β) = params(d); α  > 1.0 ? β / (α - 1.0) : Inf)

mode(d::InverseGamma) = scale(d) / (shape(d) + 1.0)

function var(d::InverseGamma)
    (α, β) = params(d)
    α > 2.0 ? β^2 / ((α - 1.0)^2 * (α - 2.0)) : Inf
end

function skewness(d::InverseGamma)
    α = shape(d)
    α > 3.0 ? 4.0 * sqrt(α - 2.0) / (α - 3.0) : NaN
end

function kurtosis(d::InverseGamma)
    α = shape(d)
    α > 4.0 ? (30.0 * α - 66.0) / ((α - 3.0) * (α - 4.0)) : NaN
end

function entropy(d::InverseGamma)
    (α, β) = params(d)
    α + lgamma(α) - (1.0 + α) * digamma(α) + log(β)
end

#### Evaluation

pdf(d::InverseGamma, x::Float64) = exp(logpdf(d, x))

function logpdf(d::InverseGamma, x::Float64)
    (α, β) = params(d)
    α * log(β) - lgamma(α) - (α + 1.0) * log(x) - β / x
end

cdf(d::InverseGamma, x::Float64) = ccdf(d.invd, 1.0 / x)
ccdf(d::InverseGamma, x::Float64) = cdf(d.invd, 1.0 / x)
logcdf(d::InverseGamma, x::Float64) = logccdf(d.invd, 1.0 / x)
logccdf(d::InverseGamma, x::Float64) = logcdf(d.invd, 1.0 / x)

quantile(d::InverseGamma, p::Float64) = 1.0 / cquantile(d.invd, p)
cquantile(d::InverseGamma, p::Float64) = 1.0 / quantile(d.invd, p)
invlogcdf(d::InverseGamma, p::Float64) = 1.0 / invlogccdf(d.invd, p)
invlogccdf(d::InverseGamma, p::Float64) = 1.0 / invlogcdf(d.invd, p)

function mgf(d::InverseGamma, t::Real)
    (a, b) = params(d)
    t == zero(t) ? one(float(t)) : 2.0*(-b*t)^(0.5a) / gamma(a) * besselk(a, sqrt(-4.0*b*t))
end

function cf(d::InverseGamma, t::Real)
    (a, b) = params(d)
    t == zero(t) ? complex(one(float(t))) : 2.0*(-im*b*t)^(0.5a) / gamma(a) * besselk(a, sqrt(-4.0*im*b*t))
end

#### Evaluation

rand(d::InverseGamma) = 1.0 / rand(d.invd)

function _rand!(d::InverseGamma, A::AbstractArray)
    s = sampler(d.invd)
    for i = 1:length(A)
    	v = 1.0 / rand(s)
        @inbounds A[i] = v
    end
    A
end