"""
```julia
    ConstantStep()
```

Use a constant step size of ``( √L̅ - √̅μ̅  ) / ( √L̅ + √̅μ̅  )``, where ``L`` is the maximum eigenvalue
of the Hessian and ``μ`` is the minimum eigenvalue of the Hessian.

See also: [`VariableStep`](@ref)
"""
mutable struct ConstantStep <: AbstractStep
    β

    function ConstantStep()
        return new( 1.0 )
    end
end

function init!( step::ConstantStep, L, μ )
    step.β = ( sqrt( L ) - sqrt( μ ) ) / ( sqrt( L ) + sqrt( μ ) )
end

struct ConstantStepIterator{T}
    β::T
end

iterate( cs::ConstantStep ) = (cs.β, ConstantStepIterator{typeof( cs.β )}( cs.β ))
iterate( cs::ConstantStep, step::ConstantStepIterator ) = (cs.β, step)

"""
```julia
    VariableStep()
```

Use a variable step size  computed by finding the roots of a simple polynomial.

See Algorithm II.1 in [1] for more information.

[1] S. Richter, C. N. Jones, and M. Morari, ‘Computational Complexity Certification for Real-Time MPC With Input
    Constraints Based on the Fast Gradient Method’, IEEE Transactions on Automatic Control, vol. 57, no. 6,
    pp. 1391–1403, 2012, doi: 10.1109/TAC.2011.2176389.

See also: [`ConstantStep`](@ref)
"""
mutable struct VariableStep <: AbstractStep
    L
    μ
    func

    function VariableStep()
        return new( 0.0, 0.0, nothing )
    end
end

function init!( step::VariableStep, L, μ )
    step.L = L
    step.μ = μ

    # This let block creates a local copy of the variables for the closure, which
    # speeds up the computations (https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-captured)
    step.func = let μ = μ, L = L
        (α, αₙ) -> ( 1 - αₙ ) * α^2 + μ/L * αₙ - αₙ^2
    end
end

mutable struct VariableStepIterator{T} <: AbstractStep
    α::T
    αₙ::T
end

function iterate( vs::VariableStep )
    α = sqrt( vs.μ / vs.L )

    step = VariableStepIterator{typeof( α )}( α, α )

    return iterate( vs, step )
end

function iterate( vs::VariableStep, step::VariableStepIterator )
    step.α = step.αₙ

    # Search for the zero crossing inside (0,1)
    step.αₙ = find_zero( αₙ -> vs.func( step.α, αₙ ), (0, 1) )

    β = ( step.α * ( 1 - step.α ) ) / ( step.α^2 + step.αₙ )

    return (β, step)
end


using Test
using Forecast

@testset "sma" begin

    @test_throws AssertionError sma(1:5,10)

    x = sma(1:5,3)
    @test length(x) == 3
    @test x == [2.0, 3.0, 4.0]

    x = sma(1:5,3,true)
    @test length(x) == 5
    @test length(collect(skipmissing(x))) == 3
    @test findlast(!ismissing,x) == 4
    @test collect(skipmissing(x)) == [2.0, 3.0, 4.0]

    x = sma(1:5,3,false)
    @test length(x) == 5
    @test length(collect(skipmissing(x))) == 3
    @test findlast(!ismissing,x) == 3
    @test collect(skipmissing(x)) == [2.0, 3.0, 4.0]

    r = rand(100)
    x = sma(r,10)
    y = collect(skipmissing(sma(r,10,false)))
    z = collect(skipmissing(sma(r,10,true)))
    @test x == y == z

end


"""
solve for the parameters in the similarity function in the rep-FAM model
"""
function solve_rep_param(; d=[1.0, 10.0], p=[.01, .5])
  alpha = log(log(1-p[1]) / log(1-p[2])) / log(abs(d[1] / d[2]))
  phi = - abs(d[1]) ^ alpha / log(1-p[1])
  return (alpha=alpha, phi=phi)
end

#= Test
using Distributions
using RCall
include("repFAM.jl")
d = [7, 10]
p = [.01, .99]
alpha, phi = solve_rep_param(d=d, p=p)
@assert abs(1-exp(-d[2]^alpha/phi) - p[2]) < 1E-6

R"d = 0:32; plot(d, 1-exp(-d^$alpha/$phi), type='o', pch=20); abline(v=$d[2], h=$p[2], lty=2)"
=#

"""
Simulate random bit vectors

For a bit vector of length K, estimate the probability of any two random bit vectors having
at most k ∈ 1:K repeats.
"""
function simulate_prob_differing_bits(bit_vec_length::Int; nsims=100)
  base_bitvec = fill(false, bit_vec_length)
  sim_vec() = [rand() > .5 for j in 1:bit_vec_length]
  num_diffs() = sum(sim_vec() .!= base_bitvec)
  est_prob_num_diffs(k_repeats::Int) = mean(num_diffs() <= k_repeats for i in 1:nsims)

  return [est_prob_num_diffs(k) for k in 1:bit_vec_length]
end

"""
Find the tolerable number of repeated elements in a bit vec of length `bit_vec_length`,
based on a maximum probability of repeats being `thresh_prob`.
"""
function tolerable_num_diffs(bit_vec_length::Int; thresh_prob::Float64=.01, nsims::Int=100)
  sims = simulate_prob_differing_bits(bit_vec_length, nsims=nsims)
  return minimum(findall(sims .> thresh_prob))
end

"""
default parameters psi=(alpha, phi) for repFAM similarity function
"""
function default_psi(bit_vec_length::Int; thresh_probs=(.01, .5), nsims::Int=100)
  tol_num_diff_bits = tolerable_num_diffs(bit_vec_length, thresh_prob=thresh_probs[1], nsims=nsims)
  return solve_rep_param(d=[tol_num_diff_bits, tol_num_diff_bits+1], p=thresh_probs)
end

"""
return default similarity function for repFAM
"""
function gen_similarity_fn(bit_vec_length::Int, thresh_probs=(.01, .5), nsims::Int=100)
  psi = default_psi(bit_vec_length, thresh_probs=thresh_probs, nsims=nsims)
  return (z1::Vector{Bool}, z2::Vector{Bool}) -> exp(-sum(abs.(z1 - z2)) ^ psi.alpha / psi.phi)
end

"""
return default similarity function for repFAM, given probability of repeats
"""
function gen_similarity_fn(; repeats=(1,3), thresh_probs=(.01, .8))
  psi = solve_rep_param(d=repeats, p=thresh_probs)
  return (z1::Vector{Bool}, z2::Vector{Bool}) -> exp(-sum(abs.(z1 - z2)) ^ psi.alpha / psi.phi)
end

"""
Generate a function of threshold, which returns similarity of one if distance
between two vectors < threshold, and zero otherwise.
"""
function sim_fn_abs(threshold::Int)::Function
  return (z1::Vector{Bool}, z2::Vector{Bool}) -> sum(abs.(z1 - z2)) <= threshold ? 1.0 : 0.0
end

function sim_fn_exp_decay(z1::Vector{Bool}, z2::Vector{Bool}, phi::Float64)
  return exp(-sum(abs.(z1 - z2)) / phi)
end

function penalty_exp_decay(z1::Vector{Bool}, z2::Vector{Bool}, phi::Float64)
  return 1 - sim_fn_exp_decay(z1, z2, phi)
end

function log_penalty_exp_decay(z1::Vector{Bool}, z2::Vector{Bool}, phi::Float64)
  return MCMC.log1m(sim_fn_exp_decay(z1, z2, phi))
end

function sim_fn_exp_decay_generator(phi::Float64)::Function
  return (z1::Vector{Bool}, z2::Vector{Bool}) -> begin
    sim_fn_exp_decay(z1, z2, phi)
  end
end

function log_repulsive_fn_generator(phi::Float64)
  return (z1::Vector{Bool}, z2::Vector{Bool}) -> let
    d = sum(abs.(z1 - z2))
    phi * MCMC.log1m(exp(-d))
  end
end

#= Test
using RCall, Distributions
include("repFAM.jl")
@time x = simulate_prob_differing_bits(25, nsims=1000)
R"plot(1:length($x), $x, xlab='k = number of diffs', ylab='probability of at most k differences', type='o')"
R"abline(h=.01)"

tolerable_num_diffs(25, thresh_prob=.01, nsims=1000)
=#

#=
using RCall, Distributions

J = 32
sim_fn = gen_similarity_fn()
sim_fn = sim_fn_abs(2)
z1 = [rand() > .5 ? 1 : 0 for i in 1:J]
z2 = z1 .+ 0

p = []
for j in 1:J
  z2[j] = z2[j] == 1 ? 0 : 1
  append!(p, sim_fn(z1, z2))
end

R"plot(1:$J, $(1 .- p), type='o', xlab='differences', ylab='prior prob')"
=#

include("update_Z_repFAM.jl")


#Modifying impute function from https://github.com/invenia/Impute.jl/blob/master/src/imputors/interp.jl

"""
Uses linear interpolation to fill for missing values.
# Example:
```
x = rand(1:10.0, 10)
x[5] = missing
linear_interpolation(x)
```
"""
function linear_interpolation(timeseries)
    if  count(i->(ismissing(i)), timeseries) > length(timeseries) *1/2
        return zero(1:length(timeseries))
    end
    data = copy(timeseries)
    i = findfirst(!ismissing, data) + 1
    while i < length(data)
        if ismissing(data[i])
            prev_idx = i - 1
            next_idx = findnext(!ismissing, data, i + 1)
            if next_idx !== nothing
                gap_sz = (next_idx - prev_idx) - 1

                diff = data[next_idx] - data[prev_idx]
                incr = diff / (gap_sz + 1)
                val = data[prev_idx] + incr
                for j in i:(next_idx - 1)
                    data[j] = val
                    val += incr
                end
                i = next_idx
            else
                break
            end
        end
        i += 1
    end
    if count(i->(ismissing(i)), data)>0
        data[1:findfirst(!ismissing, data)] .= data[findfirst(!ismissing, data)]
        data = reverse(data)
        data[1:findfirst(!ismissing, data)] .= data[findfirst(!ismissing, data)]
        data = reverse(data)
    end
    return data
end


# PROX_SEPARABLE   Evaluate the prox operator of a fully separable function.
#
# Arguments:
#
#  v is the point at which to evaluate the operator.
#  fp is a subgradient oracle for the function.
#  lambda (optional) is the proximal parameter; defaults to 1.
#  l (optional) is a lower bound for x; defaults to -Inf.
#  u (optional) is an upper bound for x; defaults to Inf.
#  x0 (optional) is a value at which to warm start the algorithm.
#  tol (optional) is a stopping tolerance.
#  MAX_ITER (optional) is the maximum number of iterations.
#
# Examples:
#
#  v = randn(n,1);
#  x = prox_separable(v, 1, @(w) sign(w));
#  [x iter] = prox_separable(v, 1, @(w) sign(w));
#
# This function can be called in vectorized form if fp is vectorized,
# i.e., if fp works elementwise, then v, l, u, and x0 can be vectors.
function prox_separable(
    v,
    lambda::Real = NaN,
    fp::Function,
    l,
    u,
    x0,
    tol::Real = NaN,
    maxiter::Integer = 0
)
    n = length(v)

    if isnan(lambda)
        lambda = 1.0
    end
    rho = 1 / lambda

    if any(isnan(l)) || isempty(l)
        l = -inf(n, 1)
    end
    if any(isnan(u)) || isempty(u)
        u = inf(n, 1)
    end
    if any(isnan(x0)) || isempty(x0)
        x0 = zeros(n, 1)
    end
    if isnan(tol)
        tol = 1e-8
    end
    if maxiter == 0
        maxiter = 500
    end

    iter = 0
    x = max(l, min(x0, u))

    while any(u .- l .> tol) && iter < MAX_ITER
        g = fp(x) .+ rho .* (x .- v)

        idx = g .> 0
        l[idx] = max(l[idx], x[idx] .- g[idx] ./ rho)
        u[idx] = x[idx]

        idx = !idx
        u[idx] = min(u[idx], x[idx] .- g[idx] ./ rho)
        l[idx] = x[idx]

        x = (l .+ u) ./ 2
        iter += 1
    end

    if any(u .- l .> tol)
        @printf(
            STDERR,
            "Warning: %d entries did not converge; max interval size = %f.",
            sum(u .- l .> tol),
            max(u .- l)
        )
    end

    return x, iter
end


"""
	random_initial_guess(lx::Vector,ux::Vector,large_val::Float64=100.0)

Generate random initial guess
"""
function random_initial_guess(lx::Vector,ux::Vector,large_val::Float64=100.0)
	nx = length(lx)
	rand_arr = rand(nx)
	x0 = []
	for i = 1:nx
		if lx[i] == -Inf
			lx_i = -large_val
		else
			lx_i = lx[i]
		end
		if ux[i] == Inf
			ux_i = large_val
		else
			ux_i = ux[i]
		end
		push!(x0, lx_i + rand_arr[i]*(ux_i-lx_i))
	end
	return x0
end


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.3'
#       jupytext_version: 0.8.6
#   kernelspec:
#     display_name: Julia 1.0.3
#     language: julia
#     name: julia-1.0
# ---

# # Load Packages for Julia

using DataFrames
using Distributions
using Plots

# # Calculate Discrete Erlang Probabilities
#
# The following function is to calculate the discrete truncated Erlang probability, given k and γ:
#
# $$\begin{equation} p_i = \frac{1}{C(n^{E})} \Bigl(\sum_{j=0}^{k-1} \frac{e^{-(i-1)\gamma} \times ((i-1)\gamma)^{j}} {j!} -\sum_{j=0}^{k-1} \frac{e^{-i\gamma} \times (i\gamma)^{j}} {j!}\Bigr),\quad\text{for i=1,…,nE}. \end{equation} $$
#
# where
#
# $$\begin{equation} n^{E} = argmin_n\Bigl(C(n) = 1 - \sum_{j=0}^{k-1} \frac{e^{-n\gamma} \times (n\gamma)^{j}} {j!} > 0.99 \Bigr) \end{equation}$$
#
# **N.B. The formula of pi here is slightly different from what is shown in the original paper because the latter (which is likely to be wrong) would lead to negative probabilities.**

#' @param k     The shape parameter of the Erlang distribution.
#' @param gamma The rate parameter of the Erlang distribution.
#' @return      A vector containing all p_i values, for i = 1 : n.
function compute_erlang_discrete_prob(k::Int64, gamma::Float64)
    n_bin = 0
    factorials = zeros(Int64, k + 1)
    factorials[1] = 1  ## factorials[1] = 0!
    for i = 1 : k
        factorials[i + 1] = factorials[i] * i  ## factorial[i + 1] = i!
    end

    one_sub_cummulative_probs = Float64[]
    cummulative_prob = 0
    while cummulative_prob <= 0.99
        n_bin = n_bin + 1
        push!(one_sub_cummulative_probs, 0)

        for j = 0 : (k - 1)
            one_sub_cummulative_probs[end] =
                one_sub_cummulative_probs[end] +
                (
                    exp( -n_bin * gamma )
                    * ( (n_bin * gamma) ^ j )
                    / factorials[j + 1]    ## factorials[j + 1] = j!
                )
        end
        cummulative_prob = 1 - one_sub_cummulative_probs[end]
    end
    one_sub_cummulative_probs = pushfirst!(one_sub_cummulative_probs, 1)

    density_prob = one_sub_cummulative_probs[1 : end - 1] - one_sub_cummulative_probs[2 : end]
    density_prob = density_prob / cummulative_prob

    return density_prob
end

# The implementation above calculates discrete probabilities pi’s base on the cummulative density function of the Erlang distribution:
#
# $$\begin{equation} p_i = CDF_{Erlang}(x = i) - CDF_{Erlang}(x = i-1) \end{equation}$$
#
# Meanwhile, the estimates of pi’s in the original paper seems to be based on the probability density function:
#
# $$\begin{equation} p_i = PDF_{Erlang}(x = i) \end{equation}$$
#
# While the two methods give slightly different estimates, they do not lead to any visible differences in the results of the subsequent simulations. This implementation uses the CDF function since it leads to faster runs.

# # Simulate the SEIR Dynamics

# The next function is to simulate the SEIR (susceptible, exposed, infectious, recovered) dynamics of an epidemic, assuming that transmission is frequency-dependent, i.e.
#
# $$\begin{equation} \beta = \beta_0 \frac{I(t)}{N} \end{equation}$$
#
# where N is the population size, I(t) is the number of infectious people at time t, and β0 is the base transmission rate.
#
# This model does not consider births and deads (i.e. N is constant).
#
# The rates at which individuals move through the E and the I classes are assumed to follow Erlang distributions of given shapes (kE, kI) and rates (γE, γI).

function seir_simulation( initial_state::Dict, parameters::Dict, max_time::Int64 )
    S_1 = get(initial_state, "S", -1)
    E_1 = get(initial_state, "E", -1)
    I_1 = get(initial_state, "I", -1)
    R_1 = get(initial_state, "R", -1)
    (S_1 != -1) || error("An initial value for S is required.")
    (E_1 != -1) || error("An initial value for E is required.")
    (I_1 != -1) || error("An initial value for I is required.")
    (R_1 != -1) || error("An initial value for R is required.")

    k_E     = get(parameters, "k_E", -1)
    gamma_E = get(parameters, "gamma_E", -1.0)
    k_I     = get(parameters, "k_I", -1)
    gamma_I = get(parameters, "gamma_I", -1.0)
    beta    = get(parameters, "beta", -1.0)
    (k_E != -1)       || error("Parameter k_E must be specified.")
    (gamma_E != -1.0) || error("Parameter gamma_E must be specified.")
    (k_I != -1)       || error("Parameter k_I must be specified.")
    (gamma_I != -1.0) || error("Parameter gamma_I must be specified.")
    (beta != -1.0)    || error("Parameter beta must be specified.")

    population_size = S_1 + E_1 + I_1 + R_1
    sim_data = DataFrame( S = Int64[], E = Int64[], I = Int64[], R = Int64[] )
    push!( sim_data, (S_1, E_1, I_1, R_1) )

    ## Initialise a matrix to store the states of the exposed sub-blocks over time.
    exposed_block_adm_rates = compute_erlang_discrete_prob(k_E, gamma_E)
    n_exposed_blocks = length(exposed_block_adm_rates)
    exposed_blocks = zeros(Int64, max_time, n_exposed_blocks)
    exposed_blocks[1, n_exposed_blocks] = sim_data[1, :E]

    ## Initialise a matrix to store the states of the infectious sub-blocks over time.
    infectious_block_adm_rates = compute_erlang_discrete_prob(k_I, gamma_I)
    n_infectious_blocks = length(infectious_block_adm_rates)
    infectious_blocks = zeros(Int64, max_time, n_infectious_blocks)
    infectious_blocks[1, n_infectious_blocks] = sim_data[1, :I]

    ## Run the simulation from time t = 2 to t = max_time
    for time = 2 : max_time
        transmission_rate = beta * sim_data[(time - 1), :I] / population_size
        exposure_prob = 1 - exp(-transmission_rate)
        distribution = Binomial(sim_data[(time - 1), :S], exposure_prob)
        new_exposed = rand(distribution)
        new_infectious = exposed_blocks[time - 1, 1]
        new_recovered  = infectious_blocks[time - 1, 1]

        if new_exposed > 0
            distribution =  Multinomial(new_exposed, exposed_block_adm_rates)
            exposed_blocks[time, :] = rand(distribution)
        end
        exposed_blocks[time, (1 : end - 1)] =
        exposed_blocks[time, (1 : end - 1)] + exposed_blocks[(time - 1), (2 : end)]

        if new_infectious > 0
            distribution =  Multinomial(new_infectious, infectious_block_adm_rates)
            infectious_blocks[time, :] = rand(distribution)
        end
        infectious_blocks[time, (1 : end - 1)] =
        infectious_blocks[time, (1 : end - 1)] + infectious_blocks[(time - 1), (2 : end)]

        push!( sim_data, ( sim_data[time - 1, :S] - new_exposed,
                           sum(exposed_blocks[time, :]),
                           sum(infectious_blocks[time, :]),
                           sim_data[time - 1, :R] + new_recovered )
        )
    end

    sim_data[:time] = collect(1 : max_time)
    return sim_data
end

# To run a simulation, simply call the seirsimulation(…) method above.
#
# Below is an example simulation where kE=5, γE=1, kI=10, γI=1, and β0=0.25 (R0=β0kIγI=2.5). The population size is N=10,000. The simmulation starts with 1 exposed case and everyone else belongs to the susceptible class. These settings are the same the the simulation 11 of the original paper.
#
# **N.B. Since this is a stochastic model, there is chance for the outbreak not to occur even with a high R0.**

# +
using Random
function main()
    sim = seir_simulation( Dict("S" => 9999, "E" => 1, "I" => 0, "R" => 0),
                       Dict("k_E" =>  5, "gamma_E" => 1.0,
                            "k_I" => 10, "gamma_I" => 1.0,
                            "beta" => 0.25),
                       300 )
    print(sim[end - 9 : end, :])

    gr(fmt = :png)
    plot( sim[:, :time],
           [sim[:, :S], sim[:, :E], sim[:, :I], sim[:, :R]],
           label = ["Susceptible", "Exposed", "Infectious", "Recovered"],
           xlabel = "Time", ylabel = "Number of Individuals",
           linealpha = 0.5, linewidth = 3 )
    Random.seed!(12345)
    test_sim = seir_simulation( Dict("S" => 9999, "E" => 1, "I" => 0, "R" => 0),
                                Dict("k_E" =>  5, "gamma_E" => 1.0,
                                     "k_I" => 10, "gamma_I" => 1.0,
                                     "beta" => 0.25),
                                100 )
    test_result = convert(Matrix, test_sim[end-2 : end, :])

    correct_result = [ 6416  1017  1298  1269   98 ;
                       6210  1045  1379  1366   99 ;
                       6010  1086  1442  1462  100 ]

    println("\n-------------------")
    if correct_result == test_result
        println("    Test PASSED")
    else
        println("    Test FAILED")
    end
    println("-------------------\n")
end


using Revise
using AdFem
using PyPlot
using LinearAlgebra
using Statistics
using MAT
SMALL_SIZE = 20
MEDIUM_SIZE = 20
BIGGER_SIZE = 20

matplotlib.rc("font", size=SMALL_SIZE)          # controls default text sizes
matplotlib.rc("axes", titlesize=SMALL_SIZE)     # fontsize of the axes title
matplotlib.rc("axes", labelsize=MEDIUM_SIZE)    # fontsize of the x and y labels
matplotlib.rc("xtick", labelsize=SMALL_SIZE)    # fontsize of the tick labels
matplotlib.rc("ytick", labelsize=SMALL_SIZE)    # fontsize of the tick labels
matplotlib.rc("legend", fontsize=SMALL_SIZE)    # legend fontsize
matplotlib.rc("figure", titlesize=BIGGER_SIZE)  # fontsize of the figure title

function f(x, y)
    0.0001*(10/(1+x^2) + x * y + 10*y^2)
end

Δt = 0.01
mmesh = Mesh(30, 30, 1/30)

E = eval_f_on_gauss_pts(f, mmesh)

E2 = matread("fenics/bwd2-61.mat")["nu"]
E3 = matread("fenics/bwd3-61.mat")["nu"]
E4 = matread("fenics/bwd4-61.mat")["nu"]

loss2 = matread("fenics/bwd2-61.mat")["loss"]
loss3 = matread("fenics/bwd3-61.mat")["loss"]
loss4 = matread("fenics/bwd4-61.mat")["loss"]

close("all")
figure(figsize=(30, 18))
subplot(331)
title("Exact")
visualize_scalar_on_gauss_points(E, mmesh)
gca().set_rasterized(true)
subplot(332)
title("DNN")
visualize_scalar_on_gauss_points(E2, mmesh)
xlabel("")
subplot(333)
title("Error")
visualize_scalar_on_gauss_points(abs.(E-E2), mmesh)
xlabel("")
subplot(334)
semilogy(loss2, label="DNN")
semilogy(loss3, label="Discretization (Gauss Points)")
semilogy(loss4, label="Discretization (Piecewise Constant)")
legend(loc="upper left", bbox_to_anchor=(0.1, -0.3))
xlabel("Iterations")
ylabel("Loss")
grid("on")
grid(b=true, which="minor")
plt.minorticks_on()
subplot(335)
title("Discretization (Gauss Points)")
visualize_scalar_on_gauss_points(E3, mmesh)
xlabel("")
subplot(336)
title("Error")
visualize_scalar_on_gauss_points(abs.(E-E3), mmesh)
xlabel("")
subplot(338)
title("Discretization (Piecewise Constant)")
visualize_scalar_on_gauss_points(E4, mmesh)
subplot(339)
title("Error")
visualize_scalar_on_gauss_points(abs.(E-E4), mmesh)
savefig("Burgers.png")
savefig("Burgers.pdf")


using Einsum
using BenchmarkTools

srand(1234)
benchmarks = BenchmarkGroup()

### CP DECOMPOSITION BENCHMARK ###
"""
Composes a 5th-order tensor from 5 factors. This is a
CP decomposition, PARAFAC, or CANDECOMP decomposition of
tensors.
"""
function cpd(dim,r)
    A = randn(dim,r)
    B = randn(dim,r)
    C = randn(dim,r)
    D = randn(dim,r)
    E = randn(dim,r)
    @einsum X[a,b,c,d,e] := A[r,a]*B[r,b]*C[r,c]*D[r,d]*E[r,e]
    return X
end
benchmarks["cpd"] = @benchmarkable cpd(1000,10)

### COMMON FATE BENCHMARK ###
"""
Common fate bechmark proposed by @faroit
"""
function commonfate(dim,r)
    A = randn(r,dim,dim,dim)
    H = randn(r,dim)
    C = randn(r,dim)
    @einsum P[a,b,f,t,c] := A[j,a,b,f]*H[j,t]*C[j,c]
    return P
end
benchmarks["commonfate"] = @benchmarkable commonfate(30,10)


#___  ____       ___ _____   _   _       _ _
#|  \/  | |     |_  /  __ \ | | | |     (_| |
#| .  . | |       | | /  \/ | | | |_ __  _| |_ ___
#| |\/| | |       | | |     | | | | '_ \| | __/ _ \
#| |  | | |___/\__/ | \__/\ | |_| | | | | | || (_) |
#_______\_____\____/ \____/  _____|_|___|_____\_____ _____ _____
#| ___ \        (_)         | | \ \ / / / __  |  _  / __  |  _  |
#| |_/ _ __ ___  _  ___  ___| |_ \ V /  `' / /| |/' `' / /| |/' |
#|  __| '__/ _ \| |/ _ \/ __| __|/   \    / / |  /| | / / |  /| |
#| |  | | | (_) | |  __| (__| |_/ /^\ \ ./ /__\ |_/ ./ /__\ |_/ /
#\_|  |_|  \___/| |\___|\___|\__\/   \/ \_____/\___/\_____/\___/
#              _/ |
#             |__/
#
# This code is part of the proposal of the team "MLJC UniTo" - University of Turin
# for "ProjectX 2020" Climate Change for AI.
# The code is licensed under MIT 3.0
# Please read readme or comments for credits and further information.

# Compiler: Julia 1.5.3

# Short description of this file:This script simulates the Fisher-KPP equation
#and fits a neural PDE to the data with the growth (aka reaction) term replaced
# by a feed-forward neural network and the diffusion term with a CNN.
#The resulting equation is solved with a RNN.

#Warning: Reset Julia Kernel (pre-compilation required) to avoid errors in
#redefinitions of constants.

cd(@__DIR__)
using Pkg#; Pkg.activate("."); Pkg.instantiate()
Pkg.add("BSON")
Pkg.add("DifferentialEquations")

using PyPlot, Printf
using LinearAlgebra
using Flux, DiffEqFlux, Optim, DiffEqSensitivity
using BSON: @save, @load
using Flux: @epochs
using DifferentialEquations

#parameter
D = 0.01; #diffusion
r = 1.0; #reaction rate

#domain
X = 1.0; T = 5.0;
dx = 0.04; dt = T/10;
x = collect(0:dx:X);
t = collect(0:dt:T);
tTrain  = collect(0:dt:T/2);
tValid  = collect(T/2:dt:T);
Nx = Int64(X/dx+1);
Nt = Int64(T/dt+1);

#initial conditions
Amp = 1.0;
Delta = 0.2
#IC-1
rho0 = Amp*(tanh.((x .- (0.5 - Delta/2))/(Delta/10)) - tanh.((x .- (0.5 + Delta/2))/(Delta/10)))/2
#IC-2
#rho0 = Amp*(1 .- tanh.((x .- 0.2)/(Delta/6)))/2.

save_folder = "data"

if isdir(save_folder)
    rm(save_folder, recursive=true)
end
mkdir(save_folder)

close("all")
figure()
plot(x, rho0)
title("Initial Condition")
gcf()

########################
# Generate training data
########################
reaction(u) = r * u .* (1 .- u)
lap = diagm(0 => -2.0 * ones(Nx), 1=> ones(Nx-1), -1 => ones(Nx-1)) ./ dx^2
#Periodic BC
lap[1,end] = 1.0/dx^2
lap[end,1] = 1.0/dx^2
#Neumann BC
#lap[1,2] = 2.0/dx^2
#lap[end,end-1] = 2.0/dx^2

function rc_ode(rho, p, t)
    #finite difference
    D * lap * rho + reaction.(rho)
end

prob = ODEProblem(rc_ode, rho0, (0.0, T), saveat=dt)
sol = DifferentialEquations.solve(prob, RK4()) #Tsit5()

ode_data = Array(sol);

figure(figsize=(8,3))

subplot(121)
pcolor(x,t,ode_data')
xlabel("x"); ylabel("t");
colorbar()

subplot(122)
for i in 1:2:Nt
    plot(x, ode_data[:,i], label="t=$(sol.t[i])")
end
xlabel("x"); ylabel(L"$\rho$")
legend(frameon=false, fontsize=7, bbox_to_anchor=(1, 1), loc="upper left", ncol=1)
tight_layout()
savefig(@sprintf("%s/training_data.pdf", save_folder))
gcf()

########################
# Define the neural PDE
########################
n_weights = 10

#for the reaction term
rx_nn = Chain(Dense(1, n_weights, tanh),
                Dense(n_weights, 2*n_weights, tanh),
                Dense(2*n_weights, n_weights, tanh),
                Dense(n_weights, 1),
                x -> x[1])

#conv with bias with initial values as 1/dx^2
w_err = 0.0
init_w = reshape([1.1 -2.5 1.0], (3, 1, 1, 1))
diff_cnn_ = Conv(init_w, [0.], pad=(0,0,0,0))

#initialize D0 close to D/dx^2
D0 = [6.5]

p1,re1 = Flux.destructure(rx_nn)
p2,re2 = Flux.destructure(diff_cnn_)
p = [p1;p2;D0]
full_restructure(p) = re1(p[1:length(p1)]), re2(p[(length(p1)+1):end-1]), p[end]

function nn_ode(u,p,t)
    rx_nn = re1(p[1:length(p1)])

    u_cnn_1   = [p[end-4] * u[end] + p[end-3] * u[1] + p[end-2] * u[2]]
    u_cnn     = [p[end-4] * u[i-1] + p[end-3] * u[i] + p[end-2] * u[i+1] for i in 2:Nx-1]
    u_cnn_end = [p[end-4] * u[end-1] + p[end-3] * u[end] + p[end-2] * u[1]]

    # Equivalent using Flux, but slower!
    #CNN term with periodic BC
    #diff_cnn_ = Conv(reshape(p[(end-4):(end-2)],(3,1,1,1)), [0.0], pad=(0,0,0,0))
    #u_cnn = reshape(diff_cnn_(reshape(u, (Nx, 1, 1, 1))), (Nx-2,))
    #u_cnn_1 = reshape(diff_cnn_(reshape(vcat(u[end:end], u[1:1], u[2:2]), (3, 1, 1, 1))), (1,))
    #u_cnn_end = reshape(diff_cnn_(reshape(vcat(u[end-1:end-1], u[end:end], u[1:1]), (3, 1, 1, 1))), (1,))

    [rx_nn([u[i]])[1] for i in 1:Nx] + p[end] * vcat(u_cnn_1, u_cnn, u_cnn_end)
end

########################
# Soving the neural PDE and setting up loss function
########################
prob_nn = ODEProblem(nn_ode, rho0, (0.0, T), p)
sol_nn = concrete_solve(prob_nn,Tsit5(), rho0, p)

function predict_rd(θ)
  # No ReverseDiff if using Flux
  Array(concrete_solve(prob_nn,Tsit5(),rho0,θ,saveat=dt,sensealg=InterpolatingAdjoint(autojacvec=ReverseDiffVJP())))
end

#match data and force the weights of the CNN to add up to zero
function loss_rd(p)
    pred = predict_rd(p)
    sum(abs2, ode_data .- pred) + 10^2 * abs(sum(p[end-4 : end-2])), pred
end

########################
# Training
########################
#Optimizer
opt = ADAM(0.001)

global count = 0
global save_count = 0
save_freq = 50

train_arr = Float64[]
diff_arr = Float64[]
w1_arr = Float64[]
w2_arr = Float64[]
w3_arr = Float64[]

#callback function to observe training
cb = function (p,l,pred)
    rx_nn, diff_cnn_, D0 = full_restructure(p)
    push!(train_arr, l)
    push!(diff_arr, p[end])

    weight = diff_cnn_.weight[:]
    push!(w1_arr, weight[1])
    push!(w2_arr, weight[2])
    push!(w3_arr, weight[3])

    println(@sprintf("Loss: %0.4f\tD0: %0.4f Weights:(%0.4f,\t %0.4f, \t%0.4f) \t Sum: %0.4f"
            ,l, D0[1], weight[1], weight[2], weight[3], sum(weight)))

    global count

    if count==0
        fig = figure(figsize=(8,2.5));
        ttl = fig.suptitle(@sprintf("Epoch = %d", count), y=1.05)
        global ttl
        subplot(131)
        pcolormesh(x,t,ode_data')
        xlabel(L"$x$"); ylabel(L"$t$"); title("Data")
        colorbar()

        subplot(132)
        img = pcolormesh(x,t,pred')
        global img
        xlabel(L"$x$"); ylabel(L"$t$"); title("Prediction")
        colorbar(); clim([0, 1]);

        ax = subplot(133); global ax
        u = collect(0:0.01:1)
        rx_line = plot(u, rx_nn.([[elem] for elem in u]), label="NN")[1];
        global rx_line
        plot(u, reaction.(u), label="True")
        title("Reaction Term")
        legend(loc="upper right", frameon=false, fontsize=8);
        ylim([0, r*0.25+0.2])

        subplots_adjust(top=0.8)
        tight_layout()
    end

    if count>0
        println("updating figure")
        img.set_array(pred[1:end-1, 1:end-1][:])
        ttl.set_text(@sprintf("Epoch = %d", count))

        u = collect(0:0.01:1)
        rx_pred = rx_nn.([[elem] for elem in u])
        rx_line.set_ydata(rx_pred)
        u = collect(0:0.01:1)

        min_lim = min(minimum(rx_pred), minimum(reaction.(u)))-0.1
        max_lim = max(maximum(rx_pred), maximum(reaction.(u)))+0.1

        ax.set_ylim([min_lim, max_lim])
    end

    global save_count
    if count%save_freq == 0
        println("saved figure")
        savefig(@sprintf("%s/pred_%05d.png", save_folder, save_count), dpi=200, bbox_inches="tight")
        save_count += 1
    end

    display(gcf())
    count += 1

    false
end

#train
res1 = DiffEqFlux.sciml_train(loss_rd, p, ADAM(0.01), cb=cb, maxiters = 10)
res2 = DiffEqFlux.sciml_train(loss_rd, res1.minimizer, ADAM(0.01), cb=cb, maxiters = 50)
res3 = DiffEqFlux.sciml_train(loss_rd, res2.minimizer, BFGS(), cb=cb, maxiters = 100)

pstar = res3.minimizer

### Recurrent Neural Network

master_rnn = Flux.RNN(26, 1, (x -> x))

function predict_rnn(x)
  out = master_rnn.(x)[end]
  Flux.reset!(master_rnn)
  out
end

function loss_rnn(p)
    #pred = predict_rd(p)
    pred_rnn = predict_rnn(p)
    sum(abs2, ode_data .- pred_rnn), pred_rnn
end

ps = Flux.params(master_rnn)

train_data = [nn_ode(predict_rd(p)[:,k], p, tTrain) for k in 1:length(tTrain)]
test_data = [nn_ode(predict_rd(p)[:,k], p, tTrain) for k in length(tTrain):length(tTrain)+ length(tValid) -1]

train_labels = [predict_rd(p)[:,k] for k in 1:length(tTrain)]
test_labels = [predict_rd(p)[:,k] for k in length(tTrain):length(tTrain) + length(tValid) - 1]

loss(x, y) = abs(sum((predict_rnn(x) .- y)))

n_epochs = 50

@epochs n_epochs Flux.train!(loss, ps, zip(train_data, train_labels), opt)

using BenchmarkTools
@btime pred_test = [predict_rnn(test_data[k]) for k in 1:length(test_data)]

train_lab = reshape(train_labels[1], : , 1)
for k in 2:6
    train_lab = cat(train_lab, reshape(train_labels[k], : , 1), dims= (2))

end

fig = figure(figsize=(8,2.5));
subplot(131)
pcolormesh(x,tTrain,train_lab')
xlabel(L"$x$"); ylabel(L"$t$"); title("Data")
colorbar()

pred_test_fig = reshape(pred_test[1], : , 1)
for k in 2:6
    pred_test_fig = cat(pred_test_fig, reshape(pred_test[k], : , 1), dims= (2))

end

subplot(132)
img = pcolormesh(x,tValid,pred_test_fig')
global img
xlabel(L"$x$"); ylabel(L"$t$"); title("Prediction")
colorbar(); clim([0, 1]);

display(gcf())

## Save trained model
"""
function pipeline()
    u = u0
    for i in t
        uderiv = nn_ode(u, p, t)
        u = predict_rnn(zip(u, uderiv))
    end
end
"""


abstract type AbstractPotential end

struct LennardJones <: AbstractPotential
    C12::Float64
    C6::Float64
end

function LennardJones(ϵ::Unitful.Energy, σ::Unitful.Length)
    s = auconvert(σ)
    e = auconvert(ϵ)
    A = 4*e*s^12
    B = 4*e*s^6
    return LennardJones(austrip(A),austrip(B))
end

function LennardJones(ϵ::Unitful.Temperature, σ)
    return LennardJones(ϵ*u"k", σ)
end

#(p::LennardJones)(r) = p.C12*r^-12 - p.C6*r^-6

#(p::LennardJones)(r::AbstractArray) = sum(p.C12.*r.^-12 .- p.C6.*r.^-6)

function (p::LennardJones)(r)
    @avx r6 = r.^-6
    return sum(muladd.(p.C12, r6, -p.C6) .* r6)
end


using AxisKeys
using DataFrames, CSV
using StanSample
using Statistics, Test

df = CSV.read(joinpath(@__DIR__, "..", "..", "data", "chimpanzees.csv"), DataFrame);

# Define the Stan language model

stan10_4 = "
data{
    int N;
    int N_actors;
    int pulled_left[N];
    int prosoc_left[N];
    int condition[N];
    int actor[N];
}
parameters{
    vector[N_actors] a;
    real bp;
    real bpC;
}
model{
    vector[N] p;
    bpC ~ normal( 0 , 10 );
    bp ~ normal( 0 , 10 );
    a ~ normal( 0 , 10 );
    for ( i in 1:504 ) {
        p[i] = a[actor[i]] + (bp + bpC * condition[i]) * prosoc_left[i];
        p[i] = inv_logit(p[i]);
    }
    pulled_left ~ binomial( 1 , p );
}
";

data = (N = size(df, 1), N_actors = length(unique(df.actor)),
    actor = df.actor, pulled_left = df.pulled_left,
    prosoc_left = df.prosoc_left, condition = df.condition);

# Sample using cmdstan

m10_4s = SampleModel("m10.4s", stan10_4)
rc10_4s = stan_sample(m10_4s; data);

# Result rethinking

rethinking = "
      mean   sd  5.5% 94.5% n_eff Rhat
bp    0.84 0.26  0.43  1.26  2271    1
bpC  -0.13 0.29 -0.59  0.34  2949    1

a[1] -0.74 0.27 -1.16 -0.31  3310    1
a[2] 10.88 5.20  4.57 20.73  1634    1
a[3] -1.05 0.28 -1.52 -0.59  4206    1
a[4] -1.05 0.28 -1.50 -0.60  4133    1
a[5] -0.75 0.27 -1.18 -0.32  4049    1
a[6]  0.22 0.27 -0.22  0.65  3877    1
a[7]  1.81 0.39  1.22  2.48  3807    1
";

# Update sections

@testset "KeyedArray manipulations" begin
    if success(rc10_4s)
        ka = read_samples(m10_4s, :keyedarray)
        kb = ka(Symbol("a.1"))

        # Other manipulations

        @test Tables.istable(ka) == true
        # All draws of :a_1
        @test size(vcat(kb...)) == (4000,)
        # All of parameters @testset "Basic HelpModel" begin
        kar = reshape(ka.data, 4000, 9);
        @test size(kar) == (4000, 9)
        # Axes ranges
        @test axes(ka) == (Base.OneTo(1000), Base.OneTo(4), Base.OneTo(9))
        # Axes keys
        @test axiskeys(ka) == (
            1:1000, 1:4,
            [Symbol("a.1"), Symbol("a.2"), Symbol("a.3"), Symbol("a.4"),
                Symbol("a.5"), Symbol("a.6"), Symbol("a.7"), :bp, :bpC]
        )
        # A single axis
        @test axiskeys(ka, :param) == vcat([Symbol("a.$i") for i in 1:7], :bp, :bpC)

        # Test combining vector param 'a'
        ma = matrix(ka, "a");
        rma = reshape(ma.data, 4000, size(ma, 3))
        @test mean(rma, dims=1) ≈ [-0.7 10.9 -1 -1 -0.7 0.2 1.8] atol=0.7
    end
end


# testing discretization procedures
# example is based on Walter Gautschi "Orthogonal Polynomials: Computation and Approximation"
# examples 2.36

using PolyChaos, Test, StaticArrays
import LinearAlgebra: norm

nodes = [40, 80, 160, 320]
tol = 1e-14

@testset "Stieltjes procedure" begin
    for n in nodes
        myfile = open("fej$n.txt")
        βref = parse.(Float64,readlines(myfile))
        βcom = stieltjes(n,fejer(n)...)[2]
        nodes, weights = fejer(n)
        @test stieltjes(n,SVector(nodes...),SVector(weights...)) == stieltjes(n,nodes,weights) == stieltjes(n,nodes,SVector(weights...)) == stieltjes(n,SVector(nodes...),weights)
        @test isapprox(norm(βref-βcom,Inf),0.;atol=tol)
        close(myfile)
    end
end

@testset "Lanczos procedure" begin
    for n in nodes
        myfile = open("fej$n.txt")
        βref = parse.(Float64,readlines(myfile))
        nodes, weights = fejer(n)
        βcom = lanczos(n,nodes,weights)[2]
        @test lanczos(n,SVector(nodes...),SVector(weights...)) == lanczos(n,nodes,weights) == lanczos(n,nodes,SVector(weights...)) == lanczos(n,SVector(nodes...),weights)
        @test isapprox(norm(βref-βcom,Inf),0.;atol=tol)
        close(myfile)
    end
end


function dist_matrix(G::SimpleDigraph)
    VV = vlist(G)
    n = length(VV)
    A = zeros(Int, n, n)
    for u = 1:n
        d = dist(G, VV[u])
        for v = 1:n
            A[u, v] = d[VV[v]]
        end
    end
    return A
end

function diam(G::SimpleDigraph)
    A = dist_matrix(G)
    if minimum(A) < 0
        return -1
    end
    return maximum(A)
end


# This file includes and example file for 2D integral

using FractalTools
using GeometryBasics
using Makie

# Generate data
f(x, y) = [
    x^2 + y^2 + 1,
    x^2 - y^2 + 1,
]
ngon = Triangle(
    Point(BigFloat(0.), BigFloat(0.)),
    Point(BigFloat(1.), BigFloat(0.)),
    Point(BigFloat(0.5), BigFloat(1.)))
npts = 100
pts = getdata(f, ngon, npts)
interp = integrate(pts, HInterp2D(0.001 * ones(2,2)))


using ANema
using ATools
using DataFrames
using Test

dir = "../data/"
conf = "n3-w-20mm-phot"
path = joinpath(dir, conf)

@testset "ANema.jl" begin
    ndf = readdf(path)
    @test nrow(ndf) == 180079
    @test length(names(ndf)) ==  46

    ndfq =selectinterval(ndf, "q1", "q2", 1700.0, 2100.0);
    @test nrow(ndfq) == 42875

    ndfz =selectinterval(ndfq, "zstd1", "zstd2", 1.0, 16.0)
    @test nrow(ndfz) == 42806

    fz, pz = fit_profile(ndfz.zstd1, ndfz.r1, "σz", "r", "pol2")
    fpars = fz.fitpar
    fstds = fz.fitstd
    @test isapprox(fpars[1], 372.96, atol=0.1)

    r1z = fz.g.(ndfz.zstd1)
    fg,p = fitg1(r1z - ndfz.r1, "r1 - r1z", 200, -5.0, 5.0, xgmin=-1.5, xgmax=1.5)
    @test isapprox(fg.std[1], 0.4, atol=0.1)

end


### Economic dispatch without renewable energy sources
function EconomicDispatch(time,
						  nodes::Nodes,
			   		      powerplants::PowerPlants,
			   			  )

	G = powerplants.unit
	T = time

	# Initialise JuMP model: Economic dispatch
	ED_mod = JuMP.Model(with_optimizer(Gurobi.Optimizer))

	# Variables
	@variable(ED_mod, P[T, G] >= 0)

	# Constraints
	# Market clearing/power balance
	@constraint(ED_mod, MarketClearing[t = T],
	    sum(P[t, g] for g in G) == nodes.systemload[t]);

	# Upper generation limit
	@constraint(ED_mod, GenerationLimitUp[g = G, t = T],
	    P[t, g] <= powerplants.pmax[g]);

	# Lower generation limit
	@constraint(ED_mod, GenerationLimitDown[g = G, t = T],
	    P[t, g] >= powerplants.pmin[g]);

	# Ramp-up limit
	@constraint(ED_mod, RampUp[g = G, t = T; t > T[1]],
		P[t, g] - P[t-1, g] <= powerplants.rup[g])

	# Ramp-down limit
	@constraint(ED_mod, RampDown[g = G, t = T; t > T[1]],
		-(P[t, g] - P[t-1, g]) <= powerplants.rdn[g])

	# Objective function
	@objective(ED_mod, Min,
	    sum(powerplants.mc[g][1] * P[t, g] for g in G, t in T));

	# Initiate optimisation process
	JuMP.optimize!(ED_mod)

	# Export results
	P_opt = JuMP.value.(P)
	price = JuMP.dual.(MarketClearing)

	return(ED_mod,
	       P_opt,
		   price
	       )
end

### Economic dispatch with renewable energy sources
function EconomicDispatch(time,
						  nodes::Nodes,
			   		      powerplants::PowerPlants,
						  renewables::Renewables,
			   			  )

	T = time
	G = powerplants.unit
	R = renewables.unit

	# Initialise JuMP model: Economic dispatch
	ED_mod = JuMP.Model(with_optimizer(Gurobi.Optimizer))

	# Variables
	@variable(ED_mod, P[T, G] >= 0)
	@variable(ED_mod, P_R[T, R] >= 0)

	# Constraints
	# Market clearing/power balance
	@constraint(ED_mod, MarketClearing[t = T],
	    sum(P[t, g] for g in G) +
			sum(P_R[t, r] for r in R) == nodes.systemload[t]);

	# Upper generation limit
	@constraint(ED_mod, GenerationLimitUp[g = G, t = T],
	    P[t, g] <= powerplants.pmax[g]);

	# Lower generation limit
	@constraint(ED_mod, GenerationLimitDown[g = G, t = T],
	    P[t, g] >= powerplants.pmin[g]);

	# Upper generation limit for renewables
	@constraint(ED_mod, ResGenerationLimitUp[r = R, t = T],
	    P_R[t, r] <= renewables.generation[r][t]);

	# Ramp-up limit
	@constraint(ED_mod, RampUp[g = G, t = T; t > T[1]],
		P[t, g] - P[t-1, g] <= powerplants.rup[g])

	# Ramp-down limit
	@constraint(ED_mod, RampDown[g = G, t = T; t > T[1]],
		-(P[t, g] - P[t-1, g]) <= powerplants.rdn[g])

	# Objective function
	@objective(ED_mod, Min,
	    sum(powerplants.mc[g][1] * P[t, g] for g in G, t in T) +
			sum(renewables.mc[r] * P_R[t, r] for r in R, t in T));

	# Initiate optimisation process
	JuMP.optimize!(ED_mod)

	# Export results
	P_opt = JuMP.value.(P)
	P_R_opt = JuMP.value.(P_R)
	price = JuMP.dual.(MarketClearing)

	# Correct RES dispatch
	# Total summed RES generation from dispatch for each t
	ResTotalGeneration = zeros(length(T), 1)

	for t in T
		ResTotalGeneration[t-(T[1]-1)] = sum(P_R_opt.data[t-(T[1]-1), :])
	end

	# Maximum output possible
	ResMaxGeneration = zeros(length(T), 1)
	for t in T
		ResMaxGeneration[t-(T[1]-1)] = sum(renewables.generation[r][t] for r in R)
	end

	# Redistribution among RES
	ResUnitGeneration = zeros(length(T), length(R))
	for t in T
		for r in R
			ResUnitGeneration[t-(T[1]-1), r] = (renewables.generation[r][t] / ResMaxGeneration[t-(T[1]-1)]) *
			ResTotalGeneration[t-(T[1]-1)]
		end
	end

	# Distributed, corrected P_R_opt
	P_R_opt_dist = JuMP.Containers.DenseAxisArray(ResUnitGeneration, T, R)

	return(ED_mod,
	       P_opt,
		   P_R_opt,
		   P_R_opt_dist,
		   price
	       )
end


using Test
using GRASP.Symmetries

@testset "Symmetries" begin

#
# Parity
# ------------------------------------------------------------------------------------------
@test isdefined(Symmetries, :even)
@test isdefined(Symmetries, :odd)

@test Parity('+') == Symmetries.even
@test Parity('-') == Symmetries.odd
@test_throws ArgumentError Parity(' ')

@test Parity(1) == Symmetries.even
@test Parity(-1) == Symmetries.odd
@test_throws ArgumentError Parity(0)
@test_throws ArgumentError Parity(-2)
@test_throws ArgumentError Parity(42)

@test string(Symmetries.even) == "+"
@test string(Symmetries.odd)  == "-"

@test parse(Parity, "+")   == Symmetries.even
@test parse(Parity, "-")   == Symmetries.odd
@test parse(Parity, " + ") == Symmetries.even
@test parse(Parity, "\t+") == Symmetries.even

@test_throws Meta.ParseError parse(Parity, "")
@test_throws Meta.ParseError parse(Parity, " ")
@test_throws Meta.ParseError parse(Parity, "++")
@test_throws Meta.ParseError parse(Parity, "α")

#
# Angular momentum
# ------------------------------------------------------------------------------------------
@test_throws ArgumentError AngularMomentum(-2)
@test_throws ArgumentError AngularMomentum(-1//2)

@test string(AngularMomentum(0))     == "0"
@test string(AngularMomentum(1//2))  == "1/2"
@test string(AngularMomentum(2//2))  == "1"
@test string(AngularMomentum(3//2))  == "3/2"
@test string(AngularMomentum(5))     == "5"
@test string(AngularMomentum(21//2)) == "21/2"
@test string(AngularMomentum(40//4)) == "10"

@test parse(AngularMomentum, "")     == AngularMomentum(0)
@test parse(AngularMomentum, "  ")   == AngularMomentum(0)
@test parse(AngularMomentum, "0")    == AngularMomentum(0)
@test parse(AngularMomentum, "1")    == AngularMomentum(1)
@test parse(AngularMomentum, "22")   == AngularMomentum(22)
@test parse(AngularMomentum, "1/2")  == AngularMomentum(1//2)
@test parse(AngularMomentum, "11/2") == AngularMomentum(11//2)
@test parse(AngularMomentum, "10/2") == AngularMomentum(10//2)

@test convert(Rational, AngularMomentum(0)) == 0//1
@test convert(Rational, AngularMomentum(0)) == 0
@test convert(Rational, AngularMomentum(22)) == 22
@test convert(Rational, AngularMomentum(22)) == 22//1
@test convert(Rational, AngularMomentum(1//2)) == 1//2

@test AngularMomentum(0) + AngularMomentum(1//2) == AngularMomentum(1//2)
@test AngularMomentum(1//2) + AngularMomentum(0) == AngularMomentum(1//2)
@test AngularMomentum(1//2) + AngularMomentum(1//2) == AngularMomentum(1)
@test AngularMomentum(1) + AngularMomentum(1//2) == AngularMomentum(3//2)
@test AngularMomentum(9//2) + AngularMomentum(10) == AngularMomentum(29//2)

@test absdiff(AngularMomentum(0), AngularMomentum(0)) == AngularMomentum(0)
@test absdiff(AngularMomentum(0), AngularMomentum(5)) == AngularMomentum(5)
@test absdiff(AngularMomentum(5), AngularMomentum(0)) == AngularMomentum(5)
@test absdiff(AngularMomentum(1), AngularMomentum(1)) == AngularMomentum(0)
@test absdiff(AngularMomentum(3//2), AngularMomentum(3//2)) == AngularMomentum(0)
@test absdiff(AngularMomentum(1//2), AngularMomentum(5)) == AngularMomentum(9//2)

@test AngularMomentum(0) < AngularMomentum(1//2)
@test AngularMomentum(0) <= AngularMomentum(0)
@test AngularMomentum(3//2) <= AngularMomentum(3//2)
@test AngularMomentum(3//2) < AngularMomentum(2)
@test AngularMomentum(3//2) <= AngularMomentum(2)
@test !(AngularMomentum(3//2) < AngularMomentum(3//2))
@test !(AngularMomentum(4//2) <= AngularMomentum(3//2))
@test AngularMomentum(4//2) > AngularMomentum(3//2)

#
# Angular symmetries
# ------------------------------------------------------------------------------------------
@test string(AngularSymmetry(AngularMomentum(0), Symmetries.even))   == "0+"
@test string(AngularSymmetry(AngularMomentum(1//2), Symmetries.odd)) == "1/2-"

@test string(AngularSymmetry(0, Symmetries.even))   == "0+"
@test string(AngularSymmetry(1//2, Symmetries.odd)) == "1/2-"

@test string(AngularSymmetry(AngularMomentum(0), '-'))    == "0-"
@test string(AngularSymmetry(AngularMomentum(1//2), '+')) == "1/2+"

@test string(AngularSymmetry(0, true))     == "0+"
@test string(AngularSymmetry(1//2, false)) == "1/2-"
@test string(AngularSymmetry(1, '-'))      == "1-"
@test string(AngularSymmetry(3//2, '+'))   == "3/2+"

end # @testset "Symmetries"


"""
    randomBinaryGenerator(rng::Random.AbstractRNG, genotypeSize::Integer)::BinaryGenotype

Generates a binary individual with a uniform probability for each gen to take *true* or *false*.

`Selfprovided Arguments` are already provided by the library, so only `User Arguments` must be provided.

# Self-provided Arguments
- `rng::Random.AbstractRNG`: random number generator for random number consistency along an experiment.

# User Arguments
- `genotypeSize::Integer`: number of genes.
"""
function randomBinaryGenerator(rng::Random.AbstractRNG, genotypeSize::Integer)::BinaryGenotype

    # Memory allocation
    genotype = Array{Bool}(undef, genotypeSize)

    # For each gen
    for i = eachindex(genotype)
        if rand(rng, UInt8) > 127
            genotype[i] = true
        else
            genotype[i] = false
        end
    end

    return BinaryGenotype(genotype)
end # function


### Implementation of a nodal polynomial basis for elements in 1 and 2 dimensions
#
# Polynomials are written in coordinates on reference elements.
###

# 1D nodal basis

using Polynomials

# Chebyshev nodes
nodes1(p) = reverse([cos((2i-1)/(2p)*pi) for i = 1:p])
nodes(p::DG1D) = nodes1(nbf(p))
#nodes(p) = [-1,-1/3,1/3,1]

phi1(nodes::Vector) = (x = Poly([0.0,1.0]); Poly{Float64}[Poly([1.0])*prod([k == i ? 1 : (x-nodes[i])/(nodes[k]-nodes[i]) for i=1:length(nodes)]) for k=1:length(nodes)])
phi(p::DG1D) = phi1(nodes(p))
dphi(p::DG1D) = polyder(phi(p))

# 2D nodal basis

#
# Consider the folling sketch of a reference element. We label
# reference coordinates as ξ and η (not x and y as those refer to real
# coordinates in the mesh)
#
#
#               η
#               ▲
#               |
#             1 -      v3
#               |      |\
#               |      |  \
#               |      |    \
#               |      |      \
#               |      |        \
#               |      |          \
#             0 -      |____________\
#               |      v1           v2
#               |
#               |
#               +------|-------------|----▶ ξ
#                      0             1
#
#
#
# Note the choice of counterclockwise orientation when numbering
# The vertices.
#

## p = 0 ##
#
# One node and (1/3,1/3) with constant basis functions
#
###########

p0ϕ1(ξ,η) = 1.
∇p0ϕ1{T}(ξ::T,η::T) = Vector2{T}(0.,0.)
const P0 = [p0ϕ1]
const DP0 = [∇p0ϕ1]
const N0 = Vector2{Float64}[Vector2{Float64}(1//3,1//3)]

## p = 1 ##
#
# Three nodes at the vertices of the reference element with linear
# basis functions
#
###########

p1ϕ1(ξ,η) = 1 - ξ - η
p1ϕ2(ξ,η) = ξ
p1ϕ3(ξ,η) = η

∇p1ϕ1{T}(ξ::T,η::T) = Vector2{T}(-1,-1)
∇p1ϕ2{T}(ξ::T,η::T) = Vector2{T}(1,0)
∇p1ϕ3{T}(ξ::T,η::T) = Vector2{T}(0,1)

const P1 = [p1ϕ1,p1ϕ2,p1ϕ3]
const DP1 = [∇p1ϕ1,∇p1ϕ2,∇p1ϕ3]
const N1 = Vector2{Float64}[Vector2{Float64}(0,0),Vector2{Float64}(1,0),Vector2{Float64}(0,1)]

## p = 2 ##
#
# Six nodes at the vertices + midpoints of the reference element with
# quadratic functions
#
###########

p2ϕ1(ξ,η) = 2.0*(0.5 - ξ - η)*(1.0 - ξ - η)
p2ϕ2(ξ,η) = 2.0ξ*(ξ - 0.5)
p2ϕ3(ξ,η) = 2.0η*(η - 0.5)
p2ϕ4(ξ,η) = 4.0ξ*(1. - ξ - η)
p2ϕ5(ξ,η) = 4.0η*ξ
p2ϕ6(ξ,η) = 4.0η*(1. - ξ - η)

# TODO: Make Vector2?
∇p2ϕ1{T}(ξ::T,η::T) = Vector2{T}(4ξ+4η-3,4ξ+4η-3)
∇p2ϕ2{T}(ξ::T,η::T) = Vector2{T}(4ξ - 1,0)
∇p2ϕ3{T}(ξ::T,η::T) = Vector2{T}(0,4η - 1)
∇p2ϕ4{T}(ξ::T,η::T) = Vector2{T}(-8ξ-4η+4,-4ξ)
∇p2ϕ5{T}(ξ::T,η::T) = Vector2{T}(4η,4ξ)
∇p2ϕ6{T}(ξ::T,η::T) = Vector2{T}(-4η,-8η-4ξ+4)

const P2 = [p2ϕ1,p2ϕ2,p2ϕ3,p2ϕ4,p2ϕ5,p2ϕ6]
const DP2 = [∇p2ϕ1,∇p2ϕ2,∇p2ϕ3,∇p2ϕ4,∇p2ϕ5,∇p2ϕ6]
const N2 = Vector2{Float64}[Vector2{Float64}(0.,0.),Vector2{Float64}(1.,0.),Vector2{Float64}(0.,1.),
            Vector2{Float64}(0.5,0),Vector2{Float64}(0.5,0.5),Vector2{Float64}(0,0.5)]

### End of P=2 Basis functions

for f in [P0,P1,P2,DP0,DP1,DP2]
    @eval $(f.env.name)(v::Vertex2) = $(f.env.name)(v.coords.e1,v.coords.e2)
    @eval $(f.env.name)(v::Vector2{Float64}) = $(f.env.name)(v.e1,v.e2)
end

getPhi(p) = (@assert 0 <= p <= 2; p==0 ? P0 : p == 1 ? P1 : P2)
getDPhi(p) = (@assert 0 <= p <= 2; p==0 ? DP0 : p == 1 ? DP1 : DP2)
nodes2(p) = (@assert 0 <= p <= 2; p==0 ? N0 : p == 1 ? N1 : N2)
phi(p::Galerkin2D) = getPhi(porder(p))
dphi(p::Galerkin2D) = getDPhi(porder(p))
nodes(p::Galerkin2D) = nodes2(porder(p))

# Evaluating coefficients on a basis
function evaluate_ref{N,T}(coeffs::Coeffs{N,T},basis)
    order(coeffs) == order(basis) || error("Basis and vector must agree")
    sum([a*f for (f,a) in zip(coeffs,basis)])
end

order(x::Array) = length(x)

function evaluate_ref(coeffs,basis,p)
    order(coeffs) == order(basis) || error("Basis and vector must agree (got $basis and $coeffs)")
    result = coeffs[1]*fapply(basis[1],p)
    for i = 2:length(basis)
        result += coeffs[i]*fapply(basis[i],p)
    end
    result
end

# Below are the hottest functions in the system. We unroll them manually
# For maximum performce. This is usually a bad idea, but because they are
# so important, we really have no other choice.

macro unroll_eval(call, coefff, basis)
    coefff = eval(coefff)
    funcs = eval(basis)
    adds = Expr(:block)
    for i = 2:length(funcs)
        addend = coefff(i,:($(Base.function_name(funcs[i]))(p)))
        push!(adds.args,:(@inbounds result += $addend))
    end
    Expr(:function,esc(call),quote
        @inbounds result = $(coefff(1,:($(Base.function_name(funcs[1]))(p))))
        $adds
        result
    end)
end

evalrefcoefff = (k,e)->:(coeffs[$k]*$e)
@unroll_eval evaluate_ref2dp0(coeffs,p) evalrefcoefff P0
@unroll_eval evaluate_ref2dp1(coeffs,p) evalrefcoefff P1
@unroll_eval evaluate_ref2dp2(coeffs,p) evalrefcoefff P2

cidevalrefcoefff = (k,e)->:(coeffs[1,cid,$k]*$e)
@unroll_eval evaluate_ref2dp0(coeffs,cid,p) cidevalrefcoefff P0
@unroll_eval evaluate_ref2dp1(coeffs,cid,p) cidevalrefcoefff P1
@unroll_eval evaluate_ref2dp2(coeffs,cid,p) cidevalrefcoefff P2

function evaluate_ref2d{T}(p,coeffs::Array{T},a,b)
    po = porder(p)
    if po == 0
        return evaluate_ref2dp0(coeffs,a,b)
    elseif po == 1
        return evaluate_ref2dp1(coeffs,a,b)
    elseif po == 2
        return evaluate_ref2dp2(coeffs,a,b)
    else
        error("Not specialized for this case")
    end
end
function evaluate_ref2d{T}(p,coeffs::Array{T},a)
    po = porder(p)
    if po == 0
        return evaluate_ref2dp0(coeffs,a)
    elseif po == 1
        return evaluate_ref2dp1(coeffs,a)
    elseif po == 2
        return evaluate_ref2dp2(coeffs,a)
    else
        error("Not specialized for this case")
    end
end
evaluate_ref2d(p,coeffs::Array{Coeffs},point) = evaluate_ref2d(p,coeffs.coeffs,point)

macro unroll_map(call, basis)
    funcs = eval(basis)
    cats = Expr(:vcat)
    for i = 1:length(funcs)
        push!(cats.args,:($(Base.function_name(funcs[i]))(p)))
    end
    Expr(:function,esc(call),quote
        $cats
    end)
end

@unroll_map evaluate_basis2dp0(p) P0
@unroll_map evaluate_basis2dp1(p) P1
@unroll_map evaluate_basis2dp2(p) P2

function evaluate_basis2d(p,point)
    po = porder(p)
    if po == 0
        return evaluate_basis2dp0(point)
    elseif po == 1
        return evaluate_basis2dp1(point)
    elseif po == 2
        return evaluate_basis2dp2(point)
    else
        error("Not specialized for this case")
    end
end

evalbasiscoefff = (k,e)->:(vRHS[$k] += mul*$e)
@unroll_eval evaluate_basis2dp0!(vRHS,mul,p) evalbasiscoefff P0
@unroll_eval evaluate_basis2dp1!(vRHS,mul,p) evalbasiscoefff P1
@unroll_eval evaluate_basis2dp2!(vRHS,mul,p) evalbasiscoefff P2

function evaluate_basis2d!(vRHS,p,mul,point)
    po = porder(p)
    if po == 0
        return evaluate_basis2dp0!(vRHS,mul,point)
    elseif po == 1
        return evaluate_basis2dp1!(vRHS,mul,point)
    elseif po == 2
        return evaluate_basis2dp2!(vRHS,mul,point)
    else
        error("Not specialized for this case")
    end
end

const evaluate = evaluate_ref


using PtFEM, Test

data = Dict(
  # Frame(nels, nn, ndim, nst, nip, finite_element(nod, nodof))
  :struc_el => Frame(12, 10, 3, 1, 1, Line(2, 3)),
  :properties => [
    1.0 1.0e4 1.0e4 1.0 1.0 1.0 1.0e8;
     ],
  :gamma => [
    0.0, 0.0, 0.0, 0.0,
    0.0, 0.0, 0.0, 0.0,
    0.0, 0.0, 0.0, 0.0],
  :x_coords => [
    0.0, 1.0, 1.0, 2.0,
    2.0, 2.0, 3.0, 3.0,
    3.0, 3.0],
  :y_coords => [
    0.0, 0.0, 1.0, 0.0,
    1.0, 2.0, 0.0, 1.0,
    2.0, 3.0],
  :z_coords => [
    0.0, 0.0, 0.0, 0.0,
    0.0, 0.0, 0.0, 0.0,
    0.0, 0.0],
  :g_num => [
    1 2 2 3 4 5 4 5 6 7 8  9;
    2 3 4 5 5 6 7 8 9 8 9 10],
  :support => [
    (1, [0 0 0 0 0 0]),
    (3, [0 0 0 0 0 0]),
    (6, [0 0 0 0 0 0]),
    (10, [0 0 0 0 0 0])
    ],
  :loaded_nodes => [
    (7, [0.0 0.0 1.0 0.0 0.0 0.0])
    ],
  :limit => 200,
  :tol => 0.00001,
  :incs => 5,
  :dload => [0.5; 0.5; 0.5; 0.05; 0.05]
)

@time m = p45(data)

@test m[2] ≈ [0.0,0.0,0.0,0.0007322670153097881,-0.0007520032370654202,
  -0.00046008115273005434]


using ADCME
using PyCall
using LinearAlgebra
using PyPlot
using Random
# Random.seed!(233)

function spatial_fem_stiffness(hmat,m,n,h)
    spatial_fem_stiffness_ = load_op_and_grad("./build/libSpatialFemStiffness","spatial_fem_stiffness", multiple=true)
    hmat,m,n,h = convert_to_tensor([hmat,m,n,h], [Float64,Int32,Int32,Float64])
    spatial_fem_stiffness_(hmat,m,n,h)
end

# TODO: specify your input parameters
m = n = 20
h = 0.1
D = diagm(0=>[1,1,0.5])
Ds = zeros(4*m*n, 3, 3)
for i = 1:m
    for j = 1:n
        Ds[4*((j-1)*m+i-1)+1,:,:] = D
        Ds[4*((j-1)*m+i-1)+2,:,:] = D
        Ds[4*((j-1)*m+i-1)+3,:,:] = D
        Ds[4*((j-1)*m+i-1)+4,:,:] = D
    end
end
u = spatial_fem_stiffness(Ds,m,n,h)
sess = Session(); init(sess)
run(sess, u);

# uncomment it for testing gradients
# error()

# TODO: change your test parameter to `m`
#       in the case of `multiple=true`, you also need to specify which component you are testings
# gradient check -- v
function scalar_function(m_)
    return sum(spatial_fem_stiffness(m_,m,n,h)[3]^4)
end

# TODO: change `m_` and `v_` to appropriate values
m_ = constant(rand(size(Ds)...))
v_ = rand(4*m*n,3,3)
y_ = scalar_function(m_)
dy_ = gradients(y_, m_)
ms_ = Array{Any}(undef, 5)
ys_ = Array{Any}(undef, 5)
s_ = Array{Any}(undef, 5)
w_ = Array{Any}(undef, 5)
gs_ =  @. 1e-4 / 10^(1:5)

for i = 1:5
    g_ = gs_[i]
    ms_[i] = m_ + g_*v_
    ys_[i] = scalar_function(ms_[i])
    s_[i] = ys_[i] - y_
    w_[i] = s_[i] - g_*sum(v_.*dy_)
end

sess = Session(); init(sess)
sval_ = run(sess, s_)
wval_ = run(sess, w_)
close("all")
loglog(gs_, abs.(sval_), "*-", label="finite difference")
loglog(gs_, abs.(wval_), "+-", label="automatic differentiation")
loglog(gs_, gs_.^2 * 0.5*abs(wval_[1])/gs_[1]^2, "--",label="\$\\mathcal{O}(\\gamma^2)\$")
loglog(gs_, gs_ * 0.5*abs(sval_[1])/gs_[1], "--",label="\$\\mathcal{O}(\\gamma)\$")

plt.gca().invert_xaxis()
legend()
xlabel("\$\\gamma\$")
ylabel("Error")


"""
    module PermPlain

Permuation types and operations.

Types representing permutations: `PList`, `PCycles`, `PMatrix`, `PDict`

Functions for converting among reprentations: `permlist`, `permcycles`, `permmatrix`, `permsparse`.

Functions: `isperm`, `permsgn`, `canoncycles`, `cycletype`, `cyclelengths`, `numcycles`, `permlength`,
      `permorder`, `permpower`, `permpower2`, `permcommute`, `permcompose`, `permcompose!`,
     `permkron`, `ltpermlist`, `support`, `supportsize`, `greatestmoved`, `leastmoved`, `preimage`, `same`, `fixed`
"""
module PermPlain

using DataStructures: counter
import LinearAlgebra
using SparseArrays: sparse
using Random
import Base: isperm, ==

include("collect.jl")

export permlist, permcycles, permsparse, permmatrix, canoncycles, permpower, permpower2,
    permsgn, cycletype, support, supportsize, greatestmoved, leastmoved, fixed
export cyclelengths, numcycles, permlength, permorder, getdata, cyc_pow_perm, permcompose, permcompose!,
    permcommute, permkron, ltpermlist, preimage
export PCycles, PList, PMatrix, PDict

randperm(::Type{T}, n::Integer) where T = collect(T, randperm(n))

eye(::Type{T}, n::Integer) where T =  Matrix{T}(LinearAlgebra.I, n, n)
speye(::Type{T}, n) where T = sparse(LinearAlgebra.I * one(T), n, n)

PCyclesVec{T} = AbstractVector{<:AbstractVector{T}} where {T <: Real}

"""
    PCycles{T}

Cycle representation of a permutation.
"""
mutable struct PCycles{T}
    data::PCyclesVec{T}
    len::Int
end

"""
    permlength(p)

The number of elements in the permutation.
"""
permlength(c::PCycles) = c.len
getdata(c::PCycles) = c.data
==(c1::PCycles, c2::PCycles) = c1.data == c2.data && c1.len == c2.len

"""
    PList{T} = AbstractVector{T} where {T <: Real}

List representation of a permtuation.
"""
PList{T} = AbstractVector{T} where {T <: Real}

permlength(p::PList) = length(p)

"""
    PMatrix{T} = AbstractMatrix{T} where {T <: Real}

Permutation matrix.
"""
PMatrix{T} = AbstractMatrix{T} where {T <: Real}
permlength(m::PMatrix) = size(m, 1)

PDictData{T} = Dict{T, T} where T

"""
    PDict{T}

Sparse representation of a permtuation as a `Dict`.
"""
mutable struct PDict{T}
    data::PDictData{T}
    len::Int
end

permlength(c::PDict) = c.len
getdata(c::PDict) = c.data
==(c1::PDict, c2::PDict) = c1.data == c2.data && c1.len == c2.len

# Compute cyclic decomposition.
# Builds a cycle list in the canonical order.
# See pari for possible efficiency improvement.
"""
    permcycles(p::AbstractVector{T}) where {T <: Real}

Return a cycle representation of the permuation `p`.

### Example
```jldoctest
julia> p = [10, 8, 5, 1, 9, 4, 6, 3, 2, 7];

julia> permcycles(p)
2-element Array{Array{Int64,1},1}:
 [1, 10, 7, 6, 4]
 [2, 8, 3, 5, 9]
```
"""
function permcycles(p::PList{T}) where T
    n = length(p)
    visited = falses(n)
    cycles = Vector{Vector{T}}(undef, 0)
@inbounds for k in convert(T,1):convert(T,n)
        if ! visited[k]
            knext = k
            cycle = Array{T}(undef, 0)
            while ! visited[knext]
                push!(cycle,knext)
                visited[knext] = true
                knext = p[knext]
            end
            length(cycle) > 1 && push!(cycles,cycle)
        end
    end
    return PCycles(cycles, n)
end

"""
    permcycles(m::AbstractMatrix)

Return a cycle representation of the permutation `m`.

### Example
```jldoctest
julia> p = [6, 1, 5, 3, 4, 2];

julia> m = permmatrix(p)
6×6 Array{Int64,2}:
 0  0  0  0  0  1
 1  0  0  0  0  0
 0  0  0  0  1  0
 0  0  1  0  0  0
 0  0  0  1  0  0
 0  1  0  0  0  0

julia> permcycles(m)
2-element Array{Array{Int64,1},1}:
 [1, 6, 2]
 [3, 5, 4]
```
"""
permcycles(m::PMatrix) = permcycles(permlist(m))

"""
    permcycles(sp::Dict{T, T})

Return a cycle representation of the permutation `sp`.

### Example
```jldoctest
julia> p = [6, 1, 5, 3, 4, 2];

julia> sp = permsparse(p)
(Dict(4=>3,2=>1,3=>5,5=>4,6=>2,1=>6), 6)

julia> p = [10, 8, 5, 1, 9, 4, 6, 3, 2, 7];

julia> sp, maxn = permsparse(p)
(Dict(7=>6,9=>2,4=>1,10=>7,2=>8,3=>5,8=>3,5=>9,6=>4,1=>10…), 10)

julia> permcycles(sp)
2-element Array{Array{Int64,1},1}:
 [7, 6, 4, 1, 10]
 [9, 2, 8, 3, 5]
```
"""
permcycles(sp::PDict) = sparsetocycles(sp)

function sparsetocycles(dsp::PDict{T}) where T
    cycs = Array{Vector{T}}(undef, 0)
    permlength(dsp) == 0 && return PCycles(cycs, permlength(dsp))
    sp = getdata(dsp)
    ks = collect(keys(sp))
    n = length(ks)
    seen = Dict{T, Bool}()
    for k in ks seen[k] = false end
@inbounds   k = ks[1]
    nseen = 0
@inbounds while nseen <= n
        foundunseen = false
        for i in ks  # could be more efficient starting search after last found
            if seen[i] == false
                k = i
                foundunseen = true
                break
            end
        end
        foundunseen == false && break
        kcyclestart = k
        cyc = Array{T}(undef, 0)
        while true
            push!(cyc,k)
            if seen[k] == true
                error("Algorithm error: double setting seen k=$k, nseen=$nseen, kcyclestart=$kcyclestart")
            end
            seen[k] = true
            k = sp[k]
            nseen = nseen + 1
            if k == kcyclestart
                break
            end
        end
        push!(cycs,cyc)
    end
    return PCycles(cycs, permlength(dsp))
end

"""
    permlist(m::PMatrix)

Return a list representation of the permutation `m`.

### Example
```jldoctest
julia> m = [0 0 0 0 0 1; 1 0 0 0 0 0; 0 0 0 0 1 0; 0 0 1 0 0 0; 0 0 0 1 0 0; 0 1 0 0 0 0]
6×6 Array{Int64,2}:
 0  0  0  0  0  1
 1  0  0  0  0  0
 0  0  0  0  1  0
 0  0  1  0  0  0
 0  0  0  1  0  0
 0  1  0  0  0  0

julia> print(permlist(m))
[6, 1, 5, 3, 4, 2]
```
"""
permlist(m::PMatrix) = mattoperm(m)
mattoperm(m::PMatrix{T}) where T = mattoperm!(Vector{T}(undef, size(m, 1)), m)

function mattoperm!(p, m::PMatrix{T}) where T
    n = size(m)[1]
    maxk = zero(T)
@inbounds  for i in 1:n
        for j in 1:n
            if m[j,i] != 1
                continue
            end
            p[j] = i
        end
    end
    p
end

"""
    permlist(cycs::PCycles, pmax::Real = 0)

Return the list representation of the permuatation `cycs`, where `pmax` is
the length of the permutation.

### Examples
```jldoctest
julia> pcyc = [[4, 3, 5], [2, 1, 6]]
2-element Array{Array{Int64,1},1}:
 [4, 3, 5]
 [2, 1, 6]

julia> print(permlist(pcyc))
[6, 1, 5, 3, 4, 2]
```
"""
permlist(cycs::PCycles) = _permlist(getdata(cycs), permlength(cycs))

_permlist(cycsvec::PCyclesVec, pmax::Real = 0) = cycstoperm(cycsvec, pmax)

function cycstoperm(cycsvec::PCyclesVec{T}, pmax::Integer = 0) where T
    isempty(cycsvec) && return collect(one(T):convert(T,pmax))
    cmaxes = [maximum(c) for c in cycsvec]
    cmax = maximum(cmaxes)  # must be a faster way
#    perm = [one(T): (pmax > cmax ? convert(T,pmax) : convert(T,cmax))]
    perm = collect(one(T): (pmax > cmax ? convert(T,pmax) : convert(T,cmax)))
@inbounds  for c in cycsvec
        for i in 2:length(c)
            perm[c[i-1]] = c[i]
        end
        perm[c[end]] = c[1]
    end
    perm[cmax+1:pmax] = collect(cmax+1:pmax)
    return perm
end

"""
    permlist(sp::PDict)

Return the list representation of the permuatation `cycs`.

### Examples
```jldoctest
julia> sp = Dict(4=>3,2=>1,3=>5,5=>4,6=>2,1=>6);

julia> print(permlist(sp))
[6, 1, 5, 3, 4, 2]
```
"""
permlist(sp::PDict) = sparsetolist(sp)

function sparsetolist(sp::PDict{T}) where T
    p = collect(one(T):convert(T, permlength(sp)))
    data = getdata(sp)
@inbounds for (i,v) in data
        p[i] = v
    end
    return p
end

## permsparse ##
"""
    permsparse(m::PMatrix)

Return a sparse (`PDict`) representation of the permutation `m`.
"""
permsparse(m::PMatrix) = mattosparse(m)

function mattosparse(m::PMatrix{T}) where T
    p = Dict{T, T}()
    return PDict{T}(mattoperm!(p, m), permlength(m))
end

permsparse(p::PList) = listtosparse(p)

function listtosparse(p::PList{T}) where T
    data = Dict{T,T}()
    maxk = zero(T)
    length(p) == 0 && return PDict{T}(data, maxk)
@inbounds for i in p
        pv = p[i]
        if pv != i
            data[i] = pv
            pv > maxk ? maxk = pv : nothing
        end
    end
    return PDict(data, permlength(p))
end

permsparse(cycs::PCycles) = cycstosparse(cycs)

function cycstosparse(cycs::PCycles{T}) where T
    data = Dict{T,T}()
    maxk = zero(T)
@inbounds for c in getdata(cycs)
        pv = c[1]
        data[c[end]] = pv
        c[end] > maxk ? maxk = c[end] : nothing
        for i in 1:length(c)-1
            pv = c[i]
            data[pv] = c[i+1]
            pv > maxk ? maxk = pv : nothing
        end
    end
    return PDict(data, permlength(cycs))
end

## permmatrix ##

"""
    permmatrix(p, sparse::Bool = false)

Convert `p` to a matrix representation.
"""
permmatrix(p::PList, sparse::Bool = false) = permtomat(p, sparse)
# Convert PLIST to PMAT
function permtomat(p::PList{T}, sparse::Bool = false) where T <: Real
    n::T = length(p)
    A = sparse ? speye(T, n) : eye(T, n)
    return A[p,:]
end

function permmatrix(sp::PDict{T}) where {T}
    n = convert(T, permlength(sp))
    ot = one(T)
    z = zero(T)
    m = eye(Int, n)
    data = getdata(sp)
@inbounds for (i, v) in data
        m[i,i] = z
        m[i,v] = ot
    end
    return m
end

# Convert cyclic decomposition to canonical form
# used by gap, Mma, and Arndt thesis.
# Note that Arndt uses a different order internally to store the cycles as a single list.
"""
    canoncycles(cycs::PCycles)

Convert cyclic decomposition `cycs` to a canonical form.
"""
function canoncycles(cycs::PCycles{T}) where T
    ocycsdata = Array{Array{T,1}}(undef, 0)
    for cyc in getdata(cycs)
        push!(ocycsdata, circshift(cyc, -argmin(cyc) + 1))
    end
    sort!(ocycsdata, by=(x)->x[1])
    return PCycles(ocycsdata, permlength(cycs))
end

## cyclelengths. Find cyclic decomposition, but only save cycle lengths ##

"""
    cyclelengths(p)

Return a `Vector` of the lengths of the cycles in the permutation `p`.

`cyclelengths(p)` is equivalent to, but more efficient than `lenth(permcycles(p))`.
"""
cyclelengths(c::PCycles) = [length(x) for x in getdata(c)]
cyclelengths(m::PMatrix) = cyclelengths(permlist(m))
function cyclelengths(p::PList{T}) where T
    n = length(p)
    visited = falses(n)
    lengths = Array{Int}(undef, 0)
@inbounds for k in one(T):convert(T,n)
        if ! visited[k]
            knext = k
            len = 0
            while ! visited[knext]
                len += 1
                visited[knext] = true
                knext = p[knext]
            end
            len > 1 && push!(lengths,len)
        end
    end
    return lengths
end

# Gives cyclelengths in canonical order.
function cyclelengths(spd::PDict{T}) where T
    sp = getdata(spd)
    cyclens = Array{Int}(undef, 0)
    isempty(sp) && return cyclens
    ks = sort(collect(keys(sp)))
    n = length(ks)
    seen = Dict{T, Bool}()
    for k in ks seen[k] = false end
    k = ks[1]
    nseen = 0
    while nseen <= n
        didsee = false
        for i in ks
            if seen[i] == false
                k = i
                didsee = true
                break
            end
        end
        didsee == false && break
        k1 = k
        nincyc = 0
        while true
            nincyc += 1
            if seen[k] == true
                error("Algorithm error: double setting seen k=$k, nseen=$nseen, k1=$k1")
            end
            seen[k] = true
            k = sp[k]
            nseen = nseen + 1
            if k == k1
                break
            end
        end
        push!(cyclens,nincyc)  # inefficient
    end
    return cyclens
end

"""
    numcycles(p)

Return the number of cycles in the permutation `p`.

`numcycles(p)` is in general more efficient than
`length(cyclelengths(p))`.
"""
numcycles(p) = length(cyclelengths(p))
numcycles(c::PCycles) = length(getdata(c))
numcycles(m::PMatrix) = numcycles(permlist(m))
function numcycles(p::PList{T}) where T
    n = length(p)
    visited = falses(n)
    Ncycles = 0
@inbounds for k in one(T):convert(T, n)
        if ! visited[k]
            knext = k
            len = 0
            while ! visited[knext]
                len += 1
                visited[knext] = true
                knext = p[knext]
            end
            len > 1 ? Ncycles += 1 : nothing
        end
    end
    return Ncycles
end

# Wikipedia says some authors also require no fixed points.
# Here, we allow fixed points
"""
    iscyclic(p)

Return `true` if permutation `p` is cyclic.
"""
iscyclic(p) = numcycles(p) == 1

## cycle type, sign ##

"""
    cycletype(p)

The cycle type of permutation `p`.
"""
cycletype(p) = counter(cyclelengths(p))

permsgn_from_lengths(lens) = (-1)^(length(lens)+sum(lens))

"""
    permsgn(p)

The sign of permutation `p`.
"""
permsgn(p) = permsgn_from_lengths(cyclelengths(p))

function permorder_from_lengths(clengths)
    result = 1
    for c in clengths
        result = lcm(result, c)
    end
    return result
end

"""
    permorder(p)

Return the order of permutation `p`.
"""
permorder(p) = permorder_from_lengths(cyclelengths(p))

# distance between two PLISTs
# could use a macro for this and ==.
# is there a penalty for using swap macro on p and q instead of two branches ?

"""
    permdistance(p::PList, q::PList)

Compute the distance between permutations `p` and `q`.
"""
function permdistance(p::PList{T}, q::PList{T}) where T
    lp = length(p)
    lq = length(q)
    count = 0
    if lp < lq
        for i in 1:lp
            p[i] != q[i] ? count += 1 : nothing
        end
        for i in lp+1:lq
            q[i] != i ? count += 1 : nothing
        end
    else  # could factor code with refs, prbly not worth the trouble
        for i in 1:lq
            p[i] != q[i] ? count += 1 : nothing
        end
        for i in lq+1:lp
            p[i] != i ? count += 1 : nothing
        end
    end
    return count
end

# composition (multiplication) of two PLISTs

"""
    permcompose(p::PList, q::PList)

Compute the composition of permutations `p` and `q`.
"""
function permcompose(p::PList{T}, q::PList) where {T}
    lq = length(q)
    lp = length(p)
    lq == lp && return p[q]  # Much faster to use the Base method
    lr = lq < lp ? lp : lq
    r = Array{T}(undef, lr)
    @inbounds  if lq <= lp
        r[1:lq] = p[q]
        r[lq+1:lp] = p[lq+1:lp]
    else
      @inbounds  for k in 1:lq
            i = q[k]
            r[k] = (i < lp ? p[i] : i)
        end
    end
    return r
end

function permcompose(q::PDict{T}, p::PDict{V}) where {T, V}
    dout = Dict{T,T}()
    z = zero(T)
    maxk = z
    seen = Dict{T,Bool}()
    for k in keys(q) seen[k] = false end  # wasteful!
    for (k,pofk) in p
        qofpofk = get(q,pofk,z)
        seen[pofk] = true
        k == qofpofk ? continue : nothing  # ignore 1-cycles
        dout[k] = (qofpofk == z ? pofk : qofpofk)
          qofpofk > maxk ? maxk = qofpofk : nothing
    end
    for j in keys(q)
        seen[j] ? nothing : dout[j] = q[j]
    end
    return PDict(dout, max(permlength(q), permlength(p)))
end

permcompose(p::PMatrix, q::PMatrix) = p * q

## permapply ##

# Vector ?
function permapply(q::PDict{T}, a::AbstractArray) where T
    aout = copy(a)
    len = length(aout)
    for (k,v) in getdata(q)
        if k <= len && v <= len
            aout[k] = a[v]
        end
    end
    return aout
end

function permapply(q::PList, a::AbstractArray)
    aout = copy(a)
    lenq = length(q)
    lena = length(a)
    len =  lenq < lena ? lenq : lena
    for k in 1:len
        v = q[k]
        v <= len  && (aout[k] = a[v])
    end
    return aout
end

permapply(q::Union{PDict{T}, PList{T}}, a::String) where T = String(permapply(q, a.data))

# power of PLIST. output is PLIST
# The method permpower below is usually preferred because it does less allocation
function permpower2(p::PList{T}, n::Integer) where T
    n == 0 && return [one(T):convert(T,length(p))]
    n == 1 && return copy(p) # for consistency, don't return ref
    n < 0  && return permpower2(invperm(p),-n)
    q = permpower2(p, Int(floor(n/2)))
    q = q[q]
    return iseven(n) ? q : p[q]
end

function permpower!(p::PList{T},
                    pret::PList{T},
                    ptmp::PList{T},
                    n::Integer) where T
    onep = one(T)
    lenp = convert(T,length(p))
    n == 0 && (for i in onep:lenp pret[i] = i end; return )
    n < 0  && (permpower!(invperm(p), pret, ptmp, -n); return )
    n == 1 && (copy!(pret,p); return)
    permpower!(p, ptmp, pret,Int(floor(n/2)))
    if iseven(n)
        for i in onep:lenp pret[i] = ptmp[ptmp[i]] end
    else
        for i in onep:lenp pret[i] = p[ptmp[ptmp[i]]] end
    end
end

# This does less allocation (in general) than permpower2, and is faster in benchmarks
function permpower(p::PList{T}, n::Integer) where T
    n == 0 && return [one(T):convert(T,length(p))]
    n == 1 && return copy(p) # for consistency, don't return ref
    pret = similar(p)
    ptmp = similar(p)
    permpower!(p,pret,ptmp,n)
    return pret
end

# Compute power of permutation. Both input and output are PCYC
# Translated from pari perm.c
# Careful of degeneracy, and empty array may be returned.
function permpower(pcyc::PCycles{T}, exp::Integer) where {T}
    cyc = getdata(pcyc)
    r = 1
    for j in 1:length(cyc)
        r += gcd(length(cyc[j])-1,exp)
    end
    c = Vector{Vector{T}}(undef, 0)
    for j in 1:length(cyc)
        v = cyc[j]
        n = length(v)
        e = mod(exp,n)
        g = gcd(n,e)
        m = div(n,g)
        if m == 1 continue end
        for i in 1:g
            p = Array{Int}(undef, 0)
            l = i
            for k in 1:m
                push!(p,v[l])
                l += e
                if l > n  l -= n  end
            end
            push!(c,p)
        end
    end
    return PCycles(c, permlength(pcyc))
end

permpower(q::PDict{T}, exp::Integer) where T = cycstosparse(permpower(sparsetocycles(q),exp))

# power of PCYC. output is PLIST
# see pari perm.c
function cyc_pow_perm(pcyc::PCycles{T}, exp::Integer) where {T}
    n = 0
    cyc = getdata(pcyc)
    # cmaxes = [maximum(c) for c in cyc]
    # n = maximum(cmaxes)
    n = permlength(pcyc)
# This routine assumes cycles of length 1 are included. Still don't understand it.
#    for j = 1:length(cyc)
#        n += length(cyc[j])+1
#    end
    p = collect(T, 1:n) # wasteful
    for j = 1:length(cyc)
        v = cyc[j]
        n = length(v)
        e = mod(exp,n)
        l = e
        for k in 1:n
            nv = v[l+1]
            ind = v[k]
            if (ind>length(p))
                error("ind is $ind length p is $(length(p))")
            end
            p[ind] = nv
            l += 1
            l == n ? l = 0 : nothing
        end
    end
    return p
end

"""
    permcommute(p, q)

Return `true` if permutations `p` and `q` commuted.
"""
function permcommute(p::PList, q::PList)
    length(q) < length(p) ? (p,q) = (q,p) : nothing
    for i in length(p)
        q[p[i]] == p[q[i]] || return false
    end
    for i in length(p)+1:length(q)
        q[i] == i || return false
    end
    return true
end

permcommute(p::PMatrix, q::PMatrix) = p * q == q * p

"""
    isperm(m::AbstractMatrix)

Return `true` if `m` is a permutation matrix.
"""
function isperm(m::PMatrix{T}) where T
    sx,sy = size(m)
    sx == sy || return false
    z = zero(T)
    o = one(T)
    seen = falses(sx)
    for i in 1:sx
        for j in 1:sx
            val = m[j,i]
            if val == o
                seen[j] == true && return false
                seen[j] = true
            elseif val != z
                return false
            end
        end
    end
    true
end

# is cyclic decomposition a permutation
function isperm(pcycs::PCycles)
    cycs = getdata(pcycs)
    isempty(cycs) && return true
    seen = counter(eltype(cycs[1])) # inefficient
    for c in cycs
        for i in c
            push!(seen,i) == 1 || return false
        end
    end
    return true
end

# is sparse representation a permutation
function isperm(dsp::PDict)
    sp = getdata(dsp)
    sort(collect(keys(sp))) == sort(collect(values(sp)))  # inefficient
end

# TODO: Could be isone ? No,..
function isid(p::PList)
    for i in 1:length(p)
        i == p[i] || return false
    end
    return true
end

function permlistisequal(p::PList{T}, q::PList) where T
    lp = length(p)
    lq = length(q)
    s = one(T)
    if lp < lq
        for i in s:lp
            p[i] == q[i] || return false
        end
        for i in lp+s:lq
            q[i] == i || return false
        end
    else  # could factor code with refs, prbly not worth the trouble
        for i in s:lq
            p[i] == q[i] || return false
        end
        for i in lq+s:lp
            p[i] == i || return false
        end
    end
    return true
end

# is p < q with lexicographical ordering ?
for (f,ret) in ((:ltpermlist, :false), (:lepermlist, :true))
    @eval begin
function ($f)(p::PList{T}, q::PList{V}) where {T, V}
    o = one(T)
    lp = length(p)
    lq = length(q)
    if ( lp < lq )
        for k in o:lp
            if q[k] != p[k]
                return p[k] < q[k] ? true : false
            end
        end
        for k in lp+o:lq
            q[k] != k && return false
        end
    else
        for k in o:lq
            if q[k] != p[k]
                return p[k] < q[k] ? true : false
            end
        end
        for k in lq+o:lp
            p[k] != k && return false
        end
    end
    return $ret
end
end
end

"""
    preimage(p, k::Int)

The preimage of `k` under permutation `p`.
"""
function preimage(p::PList, k::Int)
    k > length(p) && return k
    for i in 1:length(p)
        if p[i] == k
            return i
        end
    end
    return k  #  make preimage consistent with image
#    error("Can't find inverse image of $k.")
end

function preimage(p::PDict, i::Int)
    for (k,v) in getdata(p)
        v == i && return k
    end
    return i  #  make preimage consistent with image
#    error("Can't find inverse image of $k.")
end

# List of points mapped to same point by p and q
#function same(pin::PermList, qin::PermList)
"""
    same(p::PList, q::PList)

A list of points mapped to same point by p and q
"""
function same(p::PList, q::PList)
    lp = length(p)
    lq = length(q)
    d = Array{eltype(p)}(undef, 0)
    if lp < lq
        for i in 1:lp
            p[i] == q[i] ? push!(d,p[i]) : nothing
        end
        for i in lp+1:lq
            q[i] == i ? push!(d,q[i]) : nothing
        end
    else  # could factor code with refs, prbly not worth the trouble
        for i in 1:lq
            p[i] == q[i] ? push!(d,p[i]) : nothing
        end
        for i in lq+1:lp
            p[i] != i ? push!(d,p[i]) : nothing
        end
    end
    return d
end

# The return type depends on value of input. How to get around this ?
# There is no integer Inf.
# agrees with gap (except definition, use of Inf is different)
"""
    leastmoved(p::PList)

The least element in permutation `p` that is moved.
"""
function leastmoved(p::PList)
    lp = length(p)
    lm = lp+1
    for i in 1:lp
        k = p[i]
        k == i ? nothing :
           k < lm ? lm = k : nothing
    end
    return lm > lp ? Inf : lm
end

# agrees with gap
"""
    greatestmoved(p::PList)

The greatest element in permutation `p` that is moved.
"""
function greatestmoved(p::PList)
    lp = length(p)
    gm = 0
    for i in 1:lp
        k = p[i]
        k == i ? nothing :
           k > gm ? gm = k : nothing
    end
    return gm
end

"""
    supportsize(p::PList)

The number of elements in the support of permutation `p`.
"""
function supportsize(p::PList)
    lp = length(p)
    count = 0
    for i in 1:lp
        k = p[i]
        k != i ? count += 1 : nothing
    end
    count
end

"""
    support(p::PList)

The support of permutation `p`. That is, the elements
that are moved under the permutation.
"""
function support(p::PList)
    lp = length(p)
    mov = Array{eltype(p)}(undef, 0)
    for i in 1:lp
        k = p[i]
        k != i ? push!(mov,i) : nothing
    end
    return mov
end

"""
    fixed(p::PList)

The elements that are not moved under permutation `p`.

`fixed` is the complement of `support`.
"""
function fixed(p::PList)
    lp = length(p)
    fixedel = Array{eltype(p)}(undef, 0)
    for i in 1:lp
        k = p[i]
        k == i ? push!(fixedel,i) : nothing
    end
    return fixedel
end

## Output ##

function cycleprint(io::IO, v::Vector)
    len = length(v)
    print(io,"(")
    for i in 1:len
        print(io,v[i])
        i != len ? print(" ") : nothing
    end
    print(io,")")
end

function permarrprint(io::IO, v::Array)
    if isempty(v)
        print(io, "()")
    else
        pd = ndigits(maximum(v))
        len = length(v)
        print(io,"( ")
        for k in v
            print(io,rpad(k,pd,' ')," ")
        end
        print(io,")")
    end
end

cycleprint(v::Array) = cycleprint(STDOUT,v)
permarrprint(v::Array) = permarrprint(STDOUT,v)

# Copied from Base string.jl
function pprint_to_string(printfunc::Function, xs...)
    s = IOBuffer(Array{Uint8}(isa(xs[1],String) ? endof(xs[1]) : 0), true, true)
    truncate(s,0)
    for x in xs
        printfunc(s, x)
    end
    takebuf_string(s)
end

permarrstring(xs...) = pprint_to_string(permarrprint, xs...)
cyclestring(xs...) = pprint_to_string(cycleprint, x...)

include("matrixops.jl")
include("pivot.jl")

end # module Permplain


# This is for the PriorityQueue
using DataStructures

struct Leaf
    weight::Int64
    key::Char
end

struct Branch
    right::Union{Leaf, Branch}
    left::Union{Leaf, Branch}
    weight::Int64
end

const Node = Union{Leaf, Branch}
isbranch(branch::Branch) = true
isbranch(other::T) where {T} = false

function codebook_recurse!(leaf::Leaf, code::String,
                          dict::Dict{Char,String})
    dict[leaf.key] = code
end

function codebook_recurse!(branch::Branch, code::String,
                          dict::Dict{Char,String})
    codebook_recurse!(branch.left, string(code, "1"), dict)
    codebook_recurse!(branch.right, string(code, "0"), dict)
end

# This will depth-first search through the tree
# to create bitstrings for each character.
# Note: Any depth-first search method will work
# This outputs encoding Dict to be used for encoding
function create_codebook(n::Node)
    codebook = Dict{Char,String}()
    codebook_recurse!(n, "", codebook)
    return codebook
end

# This outputs huffman tree to generate dictionary for encoding
function create_tree(phrase::String)

    # creating weights
    weights = PriorityQueue()
    for i in phrase
        temp_string = string(i)
        if (haskey(weights, temp_string))
            weights[temp_string] += 1
        else
            weights[temp_string] = 1
        end
    end

    # Creating all nodes to iterate through
    nodes = PriorityQueue{Node, Int64}()
    while(length(weights) > 0)
        weight = peek(weights)[2]
        key = dequeue!(weights)[1]
        temp_node = Leaf(weight, key)
        enqueue!(nodes, temp_node, weight)
    end

    while(length(nodes) > 1)
        node1 = dequeue!(nodes)
        node2 = dequeue!(nodes)
        temp_node = Branch(node1, node2, node1.weight + node2.weight)
        enqueue!(nodes, temp_node, temp_node.weight)
    end

    huffman_tree = dequeue!(nodes)
    return huffman_tree

end

function encode(codebook::Dict{Char, String}, phrase::String)
    final_bitstring = ""
    for i in phrase
        final_bitstring = final_bitstring * codebook[i]
    end

    return final_bitstring
end

function decode(huffman_tree::Node, bitstring::String)
    current = huffman_tree
    final_string = ""
    for i in bitstring
        if (i == '1')
            current = current.left
        else
            current = current.right
        end
        if (!isbranch(current))
            final_string = final_string * string(current.key)
            current = huffman_tree
        end
    end

    return final_string
end

function two_pass_huffman(phrase::String)
    huffman_tree = create_tree(phrase)
    codebook = create_codebook(huffman_tree)
    println(codebook)
    bitstring = encode(codebook, phrase)
    final_string = decode(huffman_tree, bitstring)
    println(bitstring)
    println(final_string)
end

two_pass_huffman("bibbity bobbity")


include("src/TALib.jl")

using TALib
using TALib: TA_MAType_SMA

using DataFrames

print("Initialize: ")
retCode = Initialize()
println(retCode)

for f in [GetVersionString, GetVersionMajor, GetVersionMinor, GetVersionBuild, GetVersionDate, GetVersionTime]
    print(string(f) * ": ")
    println(f())
end

s_xml = FunctionDescriptionXML()
#println(s_xml)

println("Input")
inReal = [0, pi/2, pi, 3pi/2, 0, pi/2, pi, 3pi/2]
println(inReal)

println("COS")
outReal = COS(inReal)
println(outReal)

println("ACOS")
outReal = ACOS(outReal)
println(outReal)

println("SIN")
outReal = SIN(inReal)
println(outReal)

println("ASIN")
outReal = ASIN(outReal)
println(outReal)

println("TAN")
outReal = TAN(inReal)
println(outReal)

println("ATAN")
outReal = ATAN(outReal)
println(outReal)

filename = "test/ford_2012.csv"
println("Read '$filename'")
dfOHLCV = readtable(filename)
dfOHLCV[:Date] = Date(dfOHLCV[:Date])
println(dfOHLCV)
opn = Array(dfOHLCV[:Open])
hig = Array(dfOHLCV[:High])
low = Array(dfOHLCV[:Low])
cls = Array(dfOHLCV[:Close])
price = Array(dfOHLCV[:Close])
vol = Array(dfOHLCV[:Volume])
println(price)

println("MA")
#indic = MA(price)
indic = MA(price, time_period=30, ma_type=TA_MAType_SMA)
println(indic)

println(MA(dfOHLCV, price=:Close))

using PyPlot
#plot(dfOHLCV[:Date], dfOHLCV[:Close])
plot(dfOHLCV[:Date], dfOHLCV[:Close], dfOHLCV[:Date], indic)
#=
plot(dfOHLCV[:Date], dfOHLCV[:Open],
     dfOHLCV[:Date], dfOHLCV[:High],
     dfOHLCV[:Date], dfOHLCV[:Low],
     dfOHLCV[:Date], dfOHLCV[:Close],
)
=#

println("BBANDS")
println(price)
#indic = BBANDS(price)
#indic = BBANDS(price, time_period=30, deviations_up=2.0, deviations_down=2.0, ma_type=TA_MAType_SMA)
#println(indic)
indic = BBANDS(dfOHLCV, price=:Close)
println(indic)
#=
plot(dfOHLCV[:Date], dfOHLCV[:Close],
    dfOHLCV[:Date], indic[:UpperBand],
    dfOHLCV[:Date], indic[:MiddleBand],
    dfOHLCV[:Date], indic[:LowerBand],
)
=#

pause(10)

#angles = readdlm("test/angles.csv")
#angles = reshape(angles, length(angles))  # convert Array{Float64,2} to Array{Float64,1}

print("Shutdown: ")
retCode = Shutdown()
println(retCode)


f0 = Complex(0.0, 0.0)
f1 = Complex(1.0, 0.0)
f1i = Complex(0.0, 1.0)
f11 = Complex(1.0, 1.0)
f2h = Complex(2.0, 0.5)

d0 = Complex(Double64(0), Double64(0))
d1 = Complex(Double64(1), Double64(0))
d1i = Complex(Double64(0), Double64(1))
d11 = Complex(Double64(1), Double64(1))
d2h = Complex(Double64(2), Double64(0.5))

@testset "Complex" begin
  @test isapprox(log(f11), Complex{Float64}(log(d11)))
  @test isapprox(exp(f11), Complex{Float64}(exp(d11)))

  @test isapprox(sin(f11), Complex{Float64}(sin(d11)))
  @test isapprox(cos(f11), Complex{Float64}(cos(d11)))
  @test isapprox(tan(f11), Complex{Float64}(tan(d11)))
  @test isapprox(csc(f11), Complex{Float64}(csc(d11)))
  @test isapprox(sec(f11), Complex{Float64}(sec(d11)))
  @test isapprox(cot(f11), Complex{Float64}(cot(d11)))

  @test isapprox(sinh(f11), Complex{Float64}(sinh(d11)))
  @test isapprox(cosh(f11), Complex{Float64}(cosh(d11)))
  @test isapprox(tanh(f11), Complex{Float64}(tanh(d11)))
  @test isapprox(csch(f11), Complex{Float64}(csch(d11)))
  @test isapprox(sech(f11), Complex{Float64}(sech(d11)))
  @test isapprox(coth(f11), Complex{Float64}(coth(d11)))

  @test isapprox(asin(f11), Complex{Float64}(asin(d11)))
  @test isapprox(acos(f11), Complex{Float64}(acos(d11)))
  @test isapprox(atan(f2h), Complex{Float64}(atan(d2h)))
  @test isapprox(acsc(f11), Complex{Float64}(acsc(d11)))
  @test isapprox(asec(f11), Complex{Float64}(asec(d11)))
  @test isapprox(acot(f2h), Complex{Float64}(acot(d2h)))

  @test isapprox(asinh(f11), Complex{Float64}(asinh(d11)))
  @test isapprox(acosh(f11), Complex{Float64}(acosh(d11)))
  @test isapprox(atanh(f11), Complex{Float64}(atanh(d11)))
  @test isapprox(acsch(f11), Complex{Float64}(acsch(d11)))
  @test isapprox(asech(f11), Complex{Float64}(asech(d11)))
  @test isapprox(acoth(f11), Complex{Float64}(acoth(d11)))
end

@testset "Complex Accuracy" begin
  x = ComplexDF64(sqrt(df64"2"), cbrt(df64"3"))
  y = Complex{BigFloat}(x)

  @test abs(sin(y) - sin(x)) < 1.0e-31
  @test abs(cos(y) - cos(x)) < 1.0e-31
  @test abs(tan(y) - tan(x)) < 1.0e-31
  @test abs(csc(y) - csc(x)) < 1.0e-31
  @test abs(sec(y) - sec(x)) < 1.0e-31
  @test abs(tan(y) - tan(x)) < 1.0e-31

  @test abs(sinh(y) - sinh(x)) < 1.0e-31
  @test abs(cosh(y) - cosh(x)) < 1.0e-31
  @test abs(tanh(y) - tanh(x)) < 1.0e-31
  @test abs(csch(y) - csch(x)) < 1.0e-31
  @test abs(sech(y) - sech(x)) < 1.0e-31
  @test abs(tanh(y) - tanh(x)) < 1.0e-31

  @test abs(asin(y) - asin(x)) < 1.0e-31
  @test abs(acos(y) - acos(x)) < 1.0e-31
  @test abs(atan(y) - atan(x)) < 1.0e-31
  @test abs(acsc(y) - acsc(x)) < 1.0e-31
  @test abs(asec(y) - asec(x)) < 1.0e-31
  @test abs(atan(y) - atan(x)) < 1.0e-31

  @test abs(asinh(y) - asinh(x)) < 1.0e-31
  @test abs(acosh(y) - acosh(x)) < 1.0e-31
  @test abs(atanh(y) - atanh(x)) < 1.0e-31
  @test abs(acsch(y) - acsch(x)) < 1.0e-31
  @test abs(asech(y) - asech(x)) < 1.0e-31
  @test abs(atanh(y) - atanh(x)) < 1.0e-31

end

@testset "Complex exp(huge)" begin
    x = Double64(-1.168869267851442e79) + Double64(4.140780141319954e87)*im
    @test ComplexDF64(exp(x)) == exp(ComplexDF64(x))
end

@testset "Complex f(Inf)" begin
   @test isapprox(exp(ComplexDF64(Inf)), exp(ComplexF64(Inf)))
   @test isapprox(exp(ComplexDF64(-Inf)), exp(ComplexF64(-Inf)))
   @test isapprox(log(ComplexDF64(Inf)), log(ComplexF64(Inf)))
   @test isapprox(log(ComplexDF64(-Inf)), log(ComplexF64(-Inf)))

   @test isapprox(sinh(ComplexDF64(Inf)), sinh(ComplexF64(Inf)))
   @test isapprox(sinh(ComplexDF64(-Inf)), sinh(ComplexF64(-Inf)))
   @test isapprox(cosh(ComplexDF64(Inf)), cosh(ComplexF64(Inf)))
   @test isapprox(cosh(ComplexDF64(-Inf)), cosh(ComplexF64(-Inf)))
   @test isapprox(tanh(ComplexDF64(Inf)), tanh(ComplexF64(Inf)))
   @test isapprox(tanh(ComplexDF64(-Inf)), tanh(ComplexF64(-Inf)))
   @test isapprox(csch(ComplexDF64(Inf)), csch(ComplexF64(Inf)))
   @test isapprox(csch(ComplexDF64(-Inf)), csch(ComplexF64(-Inf)))
   @test isapprox(sech(ComplexDF64(Inf)), sech(ComplexF64(Inf)))
   @test isapprox(sech(ComplexDF64(-Inf)), sech(ComplexF64(-Inf)))
   @test isapprox(coth(ComplexDF64(Inf)), coth(ComplexF64(Inf)))
   @test isapprox(coth(ComplexDF64(-Inf)), coth(ComplexF64(-Inf)))

   @test isapprox(asin(ComplexDF64(Inf)), asin(ComplexF64(Inf)))
   @test isapprox(asin(ComplexDF64(-Inf)), asin(ComplexF64(-Inf)))
   @test isapprox(acos(ComplexDF64(Inf)), acos(ComplexF64(Inf)))
   @test isapprox(acos(ComplexDF64(-Inf)), acos(ComplexF64(-Inf)))
   @test isapprox(atan(ComplexDF64(Inf)), atan(ComplexF64(Inf)))
   @test isapprox(atan(ComplexDF64(-Inf)), atan(ComplexF64(-Inf)))
   @test isapprox(acsc(ComplexDF64(Inf)), acsc(ComplexF64(Inf)))
   @test isapprox(acsc(ComplexDF64(-Inf)), acsc(ComplexF64(-Inf)))
   @test isapprox(asec(ComplexDF64(Inf)), asec(ComplexF64(Inf)))
   @test isapprox(asec(ComplexDF64(-Inf)), asec(ComplexF64(-Inf)))
   @test isapprox(acot(ComplexDF64(Inf)), acot(ComplexF64(Inf)))
   @test isapprox(acot(ComplexDF64(-Inf)), acot(ComplexF64(-Inf)))

   @test isapprox(asinh(ComplexDF64(Inf)), asinh(ComplexF64(Inf)))
   @test isapprox(asinh(ComplexDF64(-Inf)), asinh(ComplexF64(-Inf)))
   @test isapprox(acosh(ComplexDF64(Inf)), acosh(ComplexF64(Inf)))
   @test isapprox(acosh(ComplexDF64(-Inf)), acosh(ComplexF64(-Inf)))
   @test isapprox(atanh(ComplexDF64(Inf)), atanh(ComplexF64(Inf)))
   @test isapprox(atanh(ComplexDF64(-Inf)), atanh(ComplexF64(-Inf)))
   @test isapprox(acsch(ComplexDF64(Inf)), acsch(ComplexF64(Inf)))
   @test isapprox(acsch(ComplexDF64(-Inf)), acsch(ComplexF64(-Inf)))
  # @test isapprox(asech(ComplexDF64(Inf)), asech(ComplexF64(Inf)))     needs fix in Julia v1
  # @test isapprox(asech(ComplexDF64(-Inf)), asech(ComplexF64(-Inf)))
   @test isapprox(acoth(ComplexDF64(Inf)), acoth(ComplexF64(Inf)))
   @test isapprox(acoth(ComplexDF64(-Inf)), acoth(ComplexF64(-Inf)))
end


@testset "Sorts" begin
    x=[3,5,1,4,2]
    BubbleSort!(x)
    @test x == [1,2,3,4,5]
    x=[3,5,1,4,2]
    InsertionSort!(x)
    @test x == [1,2,3,4,5]
    x=[3,5,1,4,2]
    SelectionSort!(x)
    @test x == [1,2,3,4,5]
end