# multigrid.jl : define multigrid cycle operations

"""
    MultigridMethod(A, sz, cycle_type)
    MultigridMethod(A, sz, cycle_type; kwargs...)
    MultigridMethod(f, sz, cycle_type)
    MultigridMethod(f, sz, cycle_type; kwargs...)

Geometric Multigrid method of type `cycle_type` for matrix `A`, that results from a discretization of a PDE on [0,1]^d using an `sz`-point mesh. A Galerkin approach is used to compose the coarse matrices, unless a function `f` is provided for direct discretization.

Inputs
======
A          : SparseMatrixCSC, sparse matrix with discretized PDE
f          : Function, function used for direct-discretization of the PDE
sz         : NTuple, PDE grid size, e.g., `(n,m)`
cycle_type : Multigrid cycle type, can be `V(ν₁,ν₂)` (V-cycle), `W(ν₁,ν₂)` (W-cycle) or `F(ν₀,ν₁,ν₂)` (FMG)

Options
=======
* max_iter : maximum number of iterations. Default value is `20` For FMG using the `F()`-cycle, this is set to 1.
* R_op     : restriction operator type, can be `Injection()` or `FullWeighting()` (default)
* P_op     : interpolation operator type, can be `Injection()`, `FullWeighting()` (default), or `Cubic()`
* ngrids   : total number of grids to use, default is `min.(⌊log₂(sz)⌋)`
* smoother : smoother, can be `GaussSeidel()` of `Jacobi()`
"""
MultigridMethod(A::Union{AbstractMatrix,Function}, sz::NTuple, cycle_type::MultigridCycle; max_iter::Int=20, R_op::TransferKind=FullWeighting(), P_op::TransferKind=FullWeighting(), ngrids=factor_twos.(sz), smoother::Smoother=GaussSeidel()) = MultigridIterable(coarsen(A, sz, R_op, P_op, ngrids), max_iter, cycle_type, smoother, Vector{Float64}(undef, 0))

"""
    V_cycle(A, sz)
    V_cycle(A, sz; kwargs...)

Geometric semicoarsened Multigrid V(2,1)-cycle for matrix `A`, that results from a discretization of a PDE on [0,1]^d using an `sz`-point mesh.

Inputs
======
A          : SparseMatrixCSC, sparse matrix with discretized PDE
sz         : NTuple, PDE grid size, e.g., `(n,m)`

For other options, see `MultigridMethod`.
"""
V_cycle(A::Union{AbstractMatrix,Function}, sz::NTuple; kwargs...) = MultigridMethod(A, sz, V(); kwargs...)

"""
    W_cycle(A, sz)
    W_cycle(A, sz; kwargs...)

Geometric semicoarsened Multigrid W(2,1)-cycle for matrix `A`, that results from a discretization of a PDE on [0,1]^d using an `sz`-point mesh.

Inputs
======
A          : SparseMatrixCSC, sparse matrix with discretized PDE
sz         : NTuple, PDE grid size, e.g., `(n,m)`

For other options, see `MultigridMethod`.
"""
W_cycle(A::Union{AbstractMatrix,Function}, sz::NTuple; kwargs...) = MultigridMethod(A, sz, W(); kwargs...)

"""
    F_cycle(A, sz)
    F_cycle(A, sz; kwargs...)

Geometric semicoarsened Full Multigrid for matrix `A`, that results from a discretization of a PDE on [0,1]^d using an `sz`-point mesh.

Inputs
======
A          : SparseMatrixCSC, sparse matrix with discretized PDE
sz         : NTuple, PDE grid size, e.g., `(n,m)`

For other options, see `MultigridMethod`.
"""
F_cycle(A::Union{AbstractMatrix,Function}, sz::NTuple; kwargs...) = MultigridMethod(A, sz, F(); max_iter=1, kwargs...)

# semicoarsening cycle
filter_at(R, s) = Base.Iterators.filter(i->sum(i.I)==s, R)

filter_fit(I1, Rall, R) = Base.Iterators.filter(i->i[2] ∈ Rall, enumerate(R))

child_iter(Rall, I1, I) = filter_fit(I1, Rall, Base.Iterators.reverse(filter_at(CartesianIndices(UnitRange.(Tuple(I-I1), Tuple(I))), sum(Tuple(I))-1)))

parent_iter(Rall, I1, I) = filter_fit(I1, Rall, filter_at(CartesianIndices(UnitRange.(Tuple(I), Tuple(I+I1))), sum(Tuple(I))+1))

grids_at_level(R::CartesianIndices{d}, s) where d = filter_at(R, s+d-1)

function μ_cycle!(grids::Array{G} where {G<:Grid}, μ::Int, ν₁::Int, ν₂::Int, grid_ptr::Int, smoother::Smoother)
    R = CartesianIndices(size(grids))
    I1, Iend = first(R), last(R)
    for I in grids_at_level(R, grid_ptr)
        smooth!(grids[I], ν₁, smoother)
    end
    if grid_ptr == sum(Tuple(Iend-I1)) + 1
        grids[Iend].x .= grids[Iend].A\grids[Iend].b # exact solve
    else
        for I in grids_at_level(R, grid_ptr+1)
            R_child = child_iter(R, I1, I)
            grids[I].b .= mean(map(i->grids[last(i)].R[first(i)]*residu(grids[last(i)]), R_child))
            fill!(grids[I].x, zero(eltype(grids[I].x)))
        end
        for i in 1:μ
            μ_cycle!(grids, μ, ν₁, ν₂, grid_ptr+1, smoother)
        end
        for I in grids_at_level(R, grid_ptr)
            R_parent = parent_iter(R, I1, I)
            # matrix-dependent prolongation
            λ = map(i->grids[I].A * high_freq_mode(first(i), grids[I].sz), R_parent)
            λ² = broadcast(i->broadcast(j->j^2, i), λ)
            ω = map(i->λ²[i]./sum(λ²), 1:length(λ)) # weight factors from [Naik, Van Rosendale]
            ip = map(i->grids[last(i)].P[first(i)]*grids[last(i)].x, R_parent)
            coarse_grid_correction!(grids, I, sum(map(i->ω[i].*ip[i], 1:length(ω))))
            smooth!(grids[I], ν₂, smoother)
        end
    end
end

high_freq_mode(dir,sz) = view(map(i->-iseven(i[dir])+isodd(i[dir]), CartesianIndices(sz.-1)), :)

function F_cycle!(grids::Array{G} where {G<:Grid}, ν₀::Int, ν₁::Int, ν₂::Int, grid_ptr::Int, smoother::Smoother)
    R = CartesianIndices(size(grids))
    I1, Iend = first(R), last(R)
    if grid_ptr == sum(Tuple(Iend-I1)) + 1
        fill!(grids[grid_ptr].x, zero(eltype(grids[grid_ptr].x)))
    else
        for I in grids_at_level(R, grid_ptr+1)
            R_child = child_iter(R, I1, I)
            grids[I].b .= mean(map(i->grids[last(i)].R[first(i)]*grids[last(i)].b, R_child))
        end
        F_cycle!(grids, ν₀, ν₁, ν₂, grid_ptr+1, smoother)
        for I in grids_at_level(R, grid_ptr)
            R_parent = parent_iter(R, I1, I)
            # matrix-dependent prolongation
            λ = map(i->grids[I].A * high_freq_mode(first(i), grids[I].sz), R_parent)
            λ² = broadcast(i->broadcast(j->j^2, i), λ)
            ω = map(i->λ²[i]./sum(λ²),1:length(λ)) # weight factors from [Naik, Van Rosendale]
            ip = map(i->P̃(first(i), Cubic(), grids[last(i)].sz...) * grids[last(i)].x, R_parent)
            grids[I].x .= sum(map(i->ω[i].*ip[i], 1:length(ω)))
        end
    end
    for i in 1:ν₀
        μ_cycle!(grids, 1, ν₁, ν₂, grid_ptr, smoother)
    end
end


#=
Generate steady states and the stability state

Reference
---------
1. Source code from https://www.nature.com/articles/s41598-017-15889-2
=#

steadies1, jac_ms1, stab_modes1 = fixedpoint_gen(sp.ODE!, sp.u0, sp.p, [-10.:1.0:10.,-10.:1.0:10.])

steadies2, jac_ms2, stab_modes2 = fixedpoint_gen(sp.ODE!, sp.p, [-10.:1.0:10.,-10.:1.0:10.])

@test length(steadies2) == sum(length.([sp.ODE_features["StableStates"]
                                       ,sp.ODE_features["UnstableStates"]]))

@test sum(getfield.(stab_modes2, :stable)) == length(sp.ODE_features["StableStates"])

@test sum(getfield.(stab_modes2, :unstable)) == length(sp.ODE_features["UnstableStates"])


using ModelingToolkit
using Graphs
using DiffEqBase
using Test
using UnPack

# Define some variables
@parameters t L g
@variables x(t) y(t) w(t) z(t) T(t) xˍt(t) yˍt(t) xˍˍt(t) yˍˍt(t)
D = Differential(t)

eqs2 = [D(D(x)) ~ T*x,
        D(D(y)) ~ T*y - g,
        0 ~ x^2 + y^2 - L^2]
pendulum2 = ODESystem(eqs2, t, [x, y, T], [L, g], name=:pendulum)
lowered_sys = ModelingToolkit.ode_order_lowering(pendulum2)

lowered_eqs = [D(xˍt) ~ T*x,
               D(yˍt) ~ T*y - g,
               D(x) ~ xˍt,
               D(y) ~ yˍt,
               0 ~ x^2 + y^2 - L^2,]
@test ODESystem(lowered_eqs, t, [xˍt, yˍt, x, y, T], [L, g], name=:pendulum) == lowered_sys
@test isequal(equations(lowered_sys), lowered_eqs)

# Simple pendulum in cartesian coordinates
eqs = [D(x) ~ w,
       D(y) ~ z,
       D(w) ~ T*x,
       D(z) ~ T*y - g,
       0 ~ x^2 + y^2 - L^2]
pendulum = ODESystem(eqs, t, [x, y, w, z, T], [L, g], name=:pendulum)

state = TearingState(pendulum)
@unpack graph, var_to_diff = state.structure
@test StructuralTransformations.maximal_matching(graph, eq->true, v->var_to_diff[v] === nothing) == map(x -> x == 0 ? StructuralTransformations.unassigned : x, [1, 2, 3, 4, 0, 0, 0, 0, 0])

using ModelingToolkit
@parameters t L g
@variables x(t) y(t) w(t) z(t) T(t) xˍt(t) yˍt(t)
D = Differential(t)
idx1_pendulum = [D(x) ~ w,
                 D(y) ~ z,
                 #0 ~ x^2 + y^2 - L^2,
                 D(w) ~ T*x,
                 D(z) ~ T*y - g,
                 # intermediate 1: 0 ~ 2x*D(x) + 2y*D(y) - 0,
                 # intermediate 2(a): 0 ~ 2x*w + 2y*z - 0, (substitute D(x) and D(y))
                 #0 ~ 2x*w + 2y*z,
                 # D(D(x)) ~ D(w) and substitute the rhs
                 D(xˍt) ~ T*x,
                 # D(D(y)) ~ D(z) and substitute the rhs
                 D(yˍt) ~ T*y - g,
                 # 2x*D(D(x)) + 2*D(x)*D(x) + 2y*D(D(y)) + 2*D(y)*D(y) and
                 # substitute the rhs
                 0 ~ 2x*(T*x) + 2*xˍt*xˍt + 2y*(T*y - g) + 2*yˍt*yˍt]
@named idx1_pendulum = ODESystem(idx1_pendulum, t, [x, y, w, z, xˍt, yˍt, T], [L, g])
first_order_idx1_pendulum = ode_order_lowering(idx1_pendulum)

using OrdinaryDiffEq
using LinearAlgebra
prob = ODEProblem(ODEFunction(first_order_idx1_pendulum),
        #  [x, y, w, z, xˍt, yˍt, T]
           [1, 0, 0, 0,   0,   0, 0.0],# 0, 0, 0, 0],
           (0, 10.0),
           [1, 9.8],
           mass_matrix=calculate_massmatrix(first_order_idx1_pendulum))
sol = solve(prob, Rodas5());
#plot(sol, vars=(1, 2))

new_sys = dae_index_lowering(ModelingToolkit.ode_order_lowering(pendulum2))

prob_auto = ODEProblem(new_sys,
                       [D(x)=>0,
                        D(y)=>0,
                        x=>1,
                        y=>0,
                        T=>0.0],
                       (0, 100.0),
                       [1, 9.8])
sol = solve(prob_auto, Rodas5());
#plot(sol, vars=(x, y))

# Define some variables
@parameters t L g
@variables x(t) y(t) T(t)
D = Differential(t)

eqs2 = [D(D(x)) ~ T*x,
        D(D(y)) ~ T*y - g,
        0 ~ x^2 + y^2 - L^2]
pendulum2 = ODESystem(eqs2, t, [x, y, T], [L, g], name=:pendulum)

# Turn into a first order differential equation system
first_order_sys = ModelingToolkit.ode_order_lowering(pendulum2)

# Perform index reduction to get an Index 1 DAE
new_sys = dae_index_lowering(first_order_sys)

u0 = [
  D(x)    => 0.0,
  D(y)    => 0.0,
  x       => 1.0,
  y       => 0.0,
  T       => 0.0
]

p = [
    L => 1.0,
    g => 9.8
]

prob_auto = ODEProblem(new_sys,u0,(0.0,10.0),p)
sol = solve(prob_auto, Rodas5());
#plot(sol, vars=(D(x), y))

let pss_pendulum2 = partial_state_selection(pendulum2)
    # This currently selects `T` rather than `x` at top level. Needs tearing priorities to fix.
    @test_broken length(equations(pss_pendulum2)) == 3
    @test_broken length(equations(ModelingToolkit.ode_order_lowering(pss_pendulum2))) == 4
end

eqs = [D(x) ~ w,
       D(y) ~ z,
       D(w) ~ T*x,
       D(z) ~ T*y - g,
       0 ~ x^2 + y^2 - L^2]
pendulum = ODESystem(eqs, t, [x, y, w, z, T], [L, g], name=:pendulum)

let pss_pendulum = partial_state_selection(pendulum)
    # This currently selects `T` rather than `x` at top level. Needs tearing priorities to fix.
    @test_broken length(equations(pss_pendulum)) == 3
end


@test modeorthonorm("Hermite",100,1e-10) #test orthonormality
@test modeorthonorm("Hermite",100,1e-10,3.) #test scaling with frequency
@test modeorthonorm("Hermite",370,1e-8) #test maximum cutoff


#=
Created November 2, 2013

Author: Spencer Lyon

This Julia script implements routines described in 'Numerically Stable
and Accurate Stochastic Simulation Approaches for Solving Dynamic
Economic Models' by Kenneth L. Judd, Lilia Maliar and Serguei Maliar,
(2011), Quantitative Economics 2/2, 173-210 (henceforth, JMM, 2011).'

As such, it is adapted from the file 'Num_Stab_Approx.m'.
=#
using MathProgBase  # load in linprog function

function normalize_data(X, Y, intercept=true)

    T, n = size(X)
    X1 = (X[:, 2:n] .- ones(T) * mean(X[:, 2:n], 1)) ./ (ones(T) * std(X[:, 2:n], 1))
    Y1 = (Y .- ones(T) * mean(Y)) ./ (ones(T) * std(Y))
    return X1, Y1
end

function de_normalize(X, Y, beta)

    # Allocate memory for new coefs
    T, n = size(X)
    B = Array(Float64, size(beta, 1) + 1, size(beta, 2))

    # Infer de-normalized slope coefficients
    B[2:end, :] = (1.0 ./ std(X[:, 2:n], 1)') * std(Y) .* beta

    # Infer intercept from others.
    B[1, :] .= mean(Y) .- mean(X[:, 2:n], 1) * B[2:end, :]

    return B
end

function OLS(X, Y, normalize=true)
    # Normal OLS.
    # Verified on 11-2-13
    T, n = size(X)
    if normalize
        X1, Y1 = normalize_data(X, Y)
        B1 = X1 \ Y1
        B = de_normalize(X, Y, B1)
    else
        B = X \ Y
    end

    return B
end

function LS_SVD(X, Y, normalize=true)
    # OLS using singular value decomposition
    # Verified on 11-2-13
    if normalize
        X1, Y1 = normalize_data(X, Y)
        U, S, V = svd(X1, thin=true)
        S_inv = diagm(0 => 1.0 ./ S)
        B1 = V*S_inv * U' * Y1
        B = de_normalize(X, Y, B1)
    else
        U, S, V = svd(X, thin=true)
        S_inv = diagm(1.0 ./ S)
        B = V * S_inv * U' * Y
    end
    return B
end

function LAD_PP(X, Y, normalize=true)
    T, n1 = size(X)
    N = size(Y, 2)

    if normalize
        n1 -= 1
        X1, Y1 = normalize_data(X, Y)
    else
        X1 = X
        Y1 = Y
    end

    #lower and upper bound
    LB = [zeros(n1)-100; zeros(2*T)]
    UB = [zeros(n1)+100; fill(Inf, 2T)]

    f = [zeros(n1); ones(2*T)]
    Aeq =  [X1 Matrix(I, T, T) -Matrix(I, T, T)]
    B1 = zeros(size(X1, 2), N)

    for j = 1:N
       beq = Y1[:, j]
       sol = linprog(f, Aeq, '=', beq, LB, UB)
       B1[:, j] = sol.sol[1:n1]
    end

    B = normalize ? de_normalize(X, Y, B1) : B1

    return B
end

function LAD_DP(X, Y, normalize=true)
    T, n1 = size(X)
    N = size(Y, 2)

    if normalize
        n1 -= 1
        X1, Y1 = normalize_data(X, Y)
    else
        X1 = X
        Y1 = Y
    end

    LB = - ones(T)
    UB = ones(T)
    Aeq = X1'
    beq = zeros(n1)
    B1 = zeros(size(X1, 2), N)

    for j=1:N
        f = - Y1[:, j]
        sol = linprog(f, Aeq, '=', beq, LB, UB)
        B1[:, j] .= -sol.attrs[:lambda][1:n1]
    end

    B = normalize ? de_normalize(X, Y, B1) : B1

    return B
end

function RLS_T(X, Y, penalty=-7)

    T, n = size(X)
    n1 = n - 1
    X1, Y1 = normalize_data(X, Y)
    # B1 = inv(X1' * X1 + T / n1 * eye(n1) * 10.0^ penalty) * X1' * Y1
    B1 = inv(X1' * X1 + T / n1 * I * 10.0^ penalty) * X1' * Y1
    B = de_normalize(X, Y, B1)
    return B
end

function RLS_TSVD(X, Y, penalty=7)
    T, n = size(X)
    n1 = n - 1
    X1, Y1 = normalize_data(X, Y)
    U, S, V = svd(X1; thin=true)
    r = sum((maximum(S)./ S) .<= 10.0^penalty)
    Sr_inv = zeros(Float64, n1, n1)
    Sr_inv[1:r, 1:r] = diagm(1./ S[1:r])
    B1 = V*Sr_inv*U'*Y1
    B = de_normalize(X, Y, B1)
    return B
end

function RLAD_PP(X, Y, penalty=7)
    # TODO: There is a bug here. linprog returns wrong answer, even when
    #       MATLAB gets it right (lame)

    T, n1 = size(X)
    N = size(Y, 2)

    n1 -= 1
    X1, Y1 = normalize_data(X, Y)

    LB = 0.0  # lower bound is 0
    UB =  Inf  # no upper bound

    f = [10.0^penalty*ones(n1*2)*T/n1; ones(2*T)]
    # Aeq =  [X1 -X1 eye(T, T) -eye(T,T)]
    Aeq =  [X1 -X1 I -I]

    B1 = zeros(size(X1, 2), N)

    for j=1:N
        beq = Y1[:, j]
        sol = linprog(f, Aeq, '=', beq, LB, UB)
        xlp = sol.sol
        B1[:, j] = xlp[1:n1] - xlp[n1+1:2*n1]
    end

    B = de_normalize(X, Y, B1)

    return B

end


using SavitzkyGolay

t = LinRange(-4, 4, 500)
y1 = collect(0:20)
y2 = exp.(-t.^2) .+ 0.05 .* (1.0 .+ randn(length(t)))

# Plain call
y1_ = savitzky_golay(y1, 11, 2)
y2_ = savitzky_golay(y2, 21, 4)

# With constructor
f1 = SGolay(11, 2)
y1c_ = f1(y1)

f2 = SGolay(21, 4)
y2c_ = f2(y2)


#!/usr/bin/env julia

using ArgParse
using Printf

# smallest cube:
#
# local_1d_3d_map = Dict(1 => (-1, -1,  1),
#                        2 => ( 1, -1,  1),
#                        3 => ( 1, -1, -1),
#                        4 => (-1, -1, -1),
#                        5 => (-1,  1, -1),
#                        6 => ( 1,  1, -1),
#                        7 => ( 1,  1,  1),
#                        8 => (-1,  1,  1))

basic_curve = [-1 +1 +1 -1 -1 +1 +1 -1;
               -1 -1 -1 -1 +1 +1 +1 +1;
               +1 +1 -1 -1 -1 -1 +1 +1]

rot_eye = [1 0 0; 0 1 0; 0 0 1]
rot_x_c = [1 0 0; 0 0 1; 0 -1 0]
rot_x_a = [1 0 0; 0 0 -1; 0 1 0]
rot_y_c = [0 0 -1; 0 1 0; 1 0 0]
rot_y_a = [0 0 1; 0 1 0; -1 0 0]
rot_z_c = [0 1 0; -1 0 0; 0 0 1]
rot_z_a = [0 -1 0; 1 0 0; 0 0 1]

function hilbert3d(n_order::Int)
    # =========
    # variables
    # =========
    n_nodes = 8^n_order
    node_chain = zeros(Int, (3, n_nodes))

    if n_order == 1
        node_chain = basic_curve
    else
        size_node  = 8^(n_order - 1)
        for node in 1:8
            coor_shift = basic_curve[:, node] .* 2^(n_order-1)
            node_coors = hilbert3d(n_order - 1)

            if node == 1
                rot_mat = rot_z_c * rot_y_a
            elseif node == 2
                rot_mat = rot_z_a * rot_x_c
            elseif node == 3
                rot_mat = rot_eye
            elseif node == 4
                rot_mat = rot_z_c * rot_x_c
            elseif node == 5
                rot_mat = rot_z_a * rot_x_a
            elseif node == 6
                rot_mat = rot_eye
            elseif node == 7
                rot_mat = rot_z_c * rot_x_a
            elseif node == 8
                rot_mat = rot_z_a * rot_y_a
            end
            node_chain[:, (node - 1) * size_node + 1 : node * size_node] = rot_mat * node_coors .+ coor_shift
        end
    end
    return node_chain
end

function parse_commandline()
    s = ArgParseSettings()

    @add_arg_table s begin
        "n_order"
        help     = "Order of the curve."
        required = true
        arg_type = Int
    end

    return parse_args(s)
end

if abspath(PROGRAM_FILE) == @__FILE__
    args = parse_commandline()
    curve = hilbert3d(args["n_order"])

    of_name = @sprintf("HILBERT_3D_CURVE_ORDER_%02d.dat", args["n_order"])
    of = open(of_name, "w")
    for i in 1:size(curve)[2]
        @printf(of, " %5d   %8d %8d %8d \n", i, curve[1, i], curve[2, i], curve[3, i])
    end
    close(of)
end


using HypothesisTests, Test
using HypothesisTests: default_tail

@testset "Mann-Whitney" begin
@testset "Basic exact test" begin
	@test abs(pvalue(ExactMannWhitneyUTest([1:10;], [2.1:2:21;])) - 0.0232) <= 1e-4
	@test abs(pvalue(ExactMannWhitneyUTest([2.1:2:21;], [1:10;])) - 0.0232) <= 1e-4
	@test abs(pvalue(ExactMannWhitneyUTest([1.5:10:100;], [2.1:2:21;])) - 0.0068) <= 1e-4
	@test abs(pvalue(ExactMannWhitneyUTest([2.1:2:21;], [1.5:10:100;])) - 0.0068) <= 1e-4
	@test default_tail(ExactMannWhitneyUTest([1:10;], [2.1:2:21;])) == :both
	show(IOBuffer(), ExactMannWhitneyUTest([1:10;], [2.1:2:21;]))
end

@testset "Exact with ties" begin
	@test abs(pvalue(ExactMannWhitneyUTest([1:10;], [1:10;])) - 1) <= 1e-4
	@test abs(pvalue(ExactMannWhitneyUTest([1:10;], [2:11;])) - 0.5096) <= 1e-4
	@test abs(pvalue(ExactMannWhitneyUTest([2:11;], [1:10;])) - 0.5096) <= 1e-4
	@test abs(pvalue(ExactMannWhitneyUTest([1:10;], [1:5; ones(5)])) - 0.0057) <= 1e-4
	@test abs(pvalue(ExactMannWhitneyUTest([1:5; ones(5)], [1:10;])) - 0.0057) <= 1e-4
	show(IOBuffer(), ExactMannWhitneyUTest([1:10;], [1:10;]))
end

@testset "Exact with ties and unequal lengths" begin
	@test abs(pvalue(ExactMannWhitneyUTest([1:10;], [2:2:24;])) - 0.0118) <= 1e-4
	@test abs(pvalue(ExactMannWhitneyUTest([2:2:24;], [1:10;])) - 0.0118) <= 1e-4
	show(IOBuffer(), ExactMannWhitneyUTest([1:10;], [2:2:24;]))
end

@testset "Approximate test" begin
	@test abs(pvalue(ApproximateMannWhitneyUTest([1:10;], [2.1:2:21;])) - 0.0257) <= 1e-4
	@test abs(pvalue(ApproximateMannWhitneyUTest([2.1:2:21;], [1:10;])) - 0.0257) <= 1e-4
	@test abs(pvalue(ApproximateMannWhitneyUTest([1.5:10:100;], [2.1:2:21;])) - 0.0091) <= 1e-4
	@test abs(pvalue(ApproximateMannWhitneyUTest([2.1:2:21;], [1.5:10:100;])) - 0.0091) <= 1e-4
	@test default_tail(ApproximateMannWhitneyUTest([1:10;], [2.1:2:21;])) == :both
	show(IOBuffer(), ApproximateMannWhitneyUTest([1:10;], [2.1:2:21;]))
end

@testset "Approximate with ties" begin
	@test abs(pvalue(ApproximateMannWhitneyUTest([1:10;], [1:10;])) - 1) <= 1e-4
	@test abs(pvalue(ApproximateMannWhitneyUTest([1:10;], [2:11;])) - 0.4948) <= 1e-4
	@test abs(pvalue(ApproximateMannWhitneyUTest([2:11;], [1:10;])) - 0.4948) <= 1e-4
	@test abs(pvalue(ApproximateMannWhitneyUTest([1:10;], [1:5; ones(5)])) - 0.0076) <= 1e-4
	@test abs(pvalue(ApproximateMannWhitneyUTest([1:5; ones(5)], [1:10;])) - 0.0076) <= 1e-4
	show(IOBuffer(), ApproximateMannWhitneyUTest([1:10;], [1:10;]))
end

@testset "Tests for automatic selection" begin
	@test abs(pvalue(MannWhitneyUTest([1:10;], [2.1:2:21;])) - 0.0232) <= 1e-4
	@test abs(pvalue(MannWhitneyUTest([1:10;], [2:11;])) - 0.4948) <= 1e-4
	show(IOBuffer(), MannWhitneyUTest([1:10;], [2.1:2:21;]))
end

@testset "Issue #39" begin
	@test abs(pvalue(ExactMannWhitneyUTest(Float32[1:10;], Float32[2:11;])) - 0.5096) <= 1e-4
end
end


# Quadratic programming in MATLAB-like style
#
#   minimize x^2 + xy + y^2 + yz + z^2
#
#   s.t.    x + 2 y + 3 z >= 4
#           x +   y       >= 1
#

using Xpress, Compat.Test
using Compat.SparseArrays
using Compat.LinearAlgebra

@testset "QP 2" begin
    model = xpress_model(
    name = "qp_02",
    f = [0., 0., 0.],
    H = [2. 1. 0.; 1. 2. 1.; 0. 1. 2.],
    A = -[1. 2. 3.; 1. 1. 0.],
    b = -[4., 1.])

    @test Xpress.getq_upper(model) == triu(sparse([2. 1. 0.; 1. 2. 1.; 0. 1. 2.]))
    @test Xpress.getq(model) == sparse([2. .5 0.; .5 2. .5; 0. .5 2.])
    @test Xpress.get_obj(model) == [0,0,0]
    @test Xpress.get_constrmatrix(model) == sparse([-1 -2 -3; -1 -1 0])
    @test Xpress.get_rhs(model) == [-4,-1]
    @test Xpress.get_lb(model) == [-1e20, -1e20, -1e20]
    @test Xpress.get_ub(model) == [1e20, 1e20, 1e20]
    @test Xpress.get_sense(model) == [:(<=),:(<=)]

    optimize(model)

    ans = get_optiminfo(model)
    @test ans.status_lp == :optimal

    @test isapprox(get_solution(model), [0.5714, .4285, .8571]; atol = 1e-3)

    @test isapprox(get_objval(model), 1.857; atol = 1e-3)
end


using DelimitedFiles

function solve(frequencyStart,input)
    changes = readdlm(input)
    # part 1
    finalFrequency = frequencyStart + sum(changes)
    println("final frequency: $finalFrequency")
    # part 2
    frequency = frequencyStart
    reached = Set()
    iteration = 0
    num_changes = length(changes)
    while ! in(frequency,reached)
        push!(reached,frequency)
        frequency+=changes[iteration%num_changes+1]
        iteration+=1
    end
    println("first duplicated frequency: $frequency at iteration $iteration")
end

solve(0,"input.txt")


import DistributedArrays

# utilities

function tensorcontractmatrix!(dest::AbstractArray{T,N}, src::AbstractArray{T,N},
							   mtx::StridedMatrix{T}, n::Int;
							   transpose::Bool=false, method::Symbol=:BLAS) where {T <: Number, N}
	#@info "TTM: dest=$(size(dest)) src=$(size(src)) mtx=$(size(mtx)) n=$n transpose=$transpose method=$method"
	# @show ntuple(i -> i<n ? i : (i+1), N-1)
	# @show (n,)
	# @show transpose ? (1,) : (2,), transpose ? (2,) : (1,)
	# @show ntuple(i -> i<n ? i : (i==n ? N : i-1), N)
	# @info "before"
	# @show size(src)
	# @show typeof(src)
	# @show src
	# @show size(mtx)
	# @show typeof(mtx)
	# @show mtx
	# @show size(dest)
	# @show typeof(dest)
	# @show dest
	# n1 = ntuple(i -> i<n ? i : (i+1), N-1)
	# n2 = transpose ? (1,) : (2,)
	# n3 = transpose ? (2,) : (1,)
	# n4 = ntuple(i -> i<n ? i : (i==n ? N : i-1), N)
	# @show n1
	# @show n2
	# @show n3
	# @show n4
	TensorOperations.contract!(1, src, :N, mtx, :N, 0, dest,
							   ntuple(i -> i<n ? i : (i+1), N-1),
							   (n,),
							   transpose ? (1,) : (2,),
							   transpose ? (2,) : (1,),
							   ntuple(i -> i<n ? i : (i==n ? N : i-1), N)
								 )
	# @info "after"
	# @show size(src)
	# @show typeof(src)
	# @show src
	# @show size(mtx)
	# @show typeof(mtx)
	# @show mtx
	# @show size(dest)
	# @show typeof(dest)
	# @show dest
end

tensorcontractmatrix(tnsr::AbstractArray{T,N}, mtx::StridedMatrix{T}, n::Int; transpose::Bool=false, method=nothing) where {T <: Number, N} =
	tensorcontractmatrix!(Array{T, N}(undef, ntuple(i -> i != n ? size(tnsr, i) : size(mtx, transpose ? 1 : 2), N)), tnsr, mtx, n, transpose=transpose, method=method)

"""
Contract N-mode tensor and M matrices.

  * `dest` array to hold the result
  * `src`  source tensor to contract
  * `matrices` matrices to contract
  * `modes` corresponding modes of matrices to contract
  * `transpose` if true, matrices are contracted along their columns
"""
function tensorcontractmatrices!(dest::Array{T,N}, src::Array{T,N}, matrices::Any,
								 modes::Any = 1:length(matrices);
								 transpose::Bool=false, method::Symbol=:BLAS) where {T <: Number, N}
	length(matrices) == length(modes) ||
		throw(ArgumentError("The number of matrices doesn't match the length of mode sequence"))
	for i in 1:length(matrices)-1
		src = tensorcontractmatrix(src, matrices[i], modes[i],
								   transpose=transpose, method=method)
	end
	tensorcontractmatrix!(dest, src, matrices[end], modes[end],
						  transpose=transpose, method=method)
end

"""
Contract N-mode tensor and M matrices.

  * `tensor` tensor to contract
  * `matrices` matrices to contract
  * `modes` corresponding modes of matrices to contract
  * `transpose` if true, matrices are contracted along their columns
"""
function tensorcontractmatrices(tensor::AbstractArray, matrices::Any,
								modes::Any = 1:length(matrices);
								transpose::Bool=false, method::Symbol=:BLAS)
	length(matrices) == length(modes) ||
		throw(ArgumentError("The number of matrices doesn't match the length of mode sequence"))
	res = tensor
	for i in eachindex(matrices)
		res = tensorcontractmatrix(res, matrices[i], modes[i],
								   transpose=transpose, method=method)
	end
	return res
end

"""
Generates random factor matrices for Tucker/CANDECOMP etc decompositions.

  * `orig_dims` original tensor dimensions
  * `core_dims` core tensor dimensions

Returns:
  * a vector of `N` (orig[n], core[n])-sized matrices
"""
_random_factors(orig_dims::NTuple{N, Int}, core_dims::NTuple{N, Int}) where {N} =
	Matrix{Float64}[randn(o_dim, c_dim) for (o_dim, c_dim) in zip(orig_dims, core_dims)]

"""
Generates random factor matrices for Tucker/CANDECOMP decompositions if core tensor is `r^N` hypercube.

Returns:
  * a vector of `N` (orig[n], r)-sized matrices
"""
_random_factors(dims::NTuple{N, Int}, r::Integer) where {N} =
	_random_factors(dims, ntuple(_ -> r, N))

"""
	khatrirao!(dest::AbstractMatrix{T}, mtxs::NTuple{N, <:AbstractMatrix{T}})

In-place Khatri-Rao matrices product (column-wise Kronecker product) calculation.
"""
@generated function khatrirao!(dest::AbstractMatrix{T},
							   mtxs::NTuple{N, <:AbstractMatrix{T}}) where {N, T}
	(N === 1) && return quote
		size(dest) == size(mtxs[1]) ||
			throw(DimensionMismatch("Output and single input matrix have different sizes ($(size(dest)) and $(size(mtxs[1])))"))
		return copyto!(dest, mtxs[1])
	end
	# generate the code for looping over the matrices 2:N
	_innerloop = Base.Cartesian.lreplace(:(desti[offsj_k + 1] = destij_k), :k, N)
	for k in N:-1:2
		_innerloop = Base.Cartesian.lreplace(quote
			for j_k in axes(mtxs[k], 1)
				destij_k = destij_{k-1}*coli_k[j_k]
				offsj_k = offsj_{k-1}*size(mtxs[k], 1) + j_k - 1
				$_innerloop
			end
		end, :k, k)
	end
	# main code
	quote
	# dimensions check
	ncols = size(dest, 2)
	for i in 1:length(mtxs)
		(size(mtxs[i], 2) == ncols) ||
			throw(DimensionMismatch("Output matrix and input matrix #$i have different number of columns ($ncols and $(size(mtxs[i], 2)))"))
	end
	nrows = prod(@ntuple($N, i -> size(mtxs[i], 1)))
	size(dest, 1) == nrows ||
		throw(DimensionMismatch("Output matrix rows and the expected number of rows do not match ($(size(dest, 1)) and $nrows)"))
	# multiplication
	@inbounds for i in axes(dest, 2)
		@nexprs($N, k -> (coli_k = view(mtxs[k], :, i)))
		desti = view(dest, :, i)
		for j_1 in axes(mtxs[1], 1)
			destij_1 = coli_1[j_1]
			offsj_1 = j_1 - 1
			$_innerloop
		end
	end
	return dest
	end
end

khatrirao!(dest::AbstractMatrix, mtxs::AbstractVector) =
	khatrirao!(dest, tuple(mtxs...))
khatrirao!(dest::AbstractMatrix, mtxs...) = khatrirao!(dest, tuple(mtxs...))

"""
	khatrirao(mtxs::NTuple{N, <:AbstractMatrix{T}})

Calculates Khatri-Rao product of a sequence of matrices (column-wise Kronecker product).
"""
function khatrirao(mtxs::NTuple{N, <:AbstractMatrix{T}}) where {N, T}
	(N === 1) && return copy(first(mtxs))
	ncols = size(first(mtxs), 2)
	for i in 2:length(mtxs)
		(size(mtxs[i], 2) == ncols) ||
			throw(DimensionMismatch("Input matrices have different number of columns ($ncols and $(size(mtxs[i], 2)))"))
	end
	return khatrirao!(Matrix{T}(undef, prod(ntuple(i -> size(mtxs[i], 1), N)), ncols), mtxs)
end

khatrirao(mtxs::AbstractVector) = khatrirao(tuple(mtxs...))
khatrirao(mtxs...) = khatrirao(tuple(mtxs...))

"""
Unfolds the tensor into matrix, such that the specified
group of modes becomes matrix rows and the other one becomes columns.

  * `row_modes` vector of modes to be unfolded as rows
  * `col_modes` vector of modes to be unfolded as columns
"""
function _unfold(tnsr::StridedArray, row_modes::Vector{Int}, col_modes::Vector{Int})
	length(row_modes) + length(col_modes) == ndims(tnsr) ||
		throw(ArgumentError("column and row modes should be disjoint subsets of 1:$(ndims(tnsr))"))

	dims = size(tnsr)
	return reshape(permutedims(tnsr, [row_modes; col_modes]), prod(dims[row_modes]), prod(dims[col_modes]))
end
function _unfold(tnsr::DistributedArrays.DArray{T,N,Array{T,N}}, row_modes::Vector{Int}, col_modes::Vector{Int}) where {T <: Number, N}
	length(row_modes) + length(col_modes) == ndims(tnsr) ||
		throw(ArgumentError("column and row modes should be disjoint subsets of 1:$(ndims(tnsr))"))
	dims = size(tnsr)
	if row_modes[1] == 1
		row_modes_new = row_modes
		col_modes_new = col_modes
	else
		row_modes_new = [1]
		col_modes_new = [i for i = 2:N]
		a = [i for i = 1:N]
		k = a .!= row_modes[1]
		v = [row_modes[1], a[k]...]
		dims = dims[v]
	end
	s = reshape(permutedims(Core.eval(TensorDecompositions, Symbol("dX$(row_modes[1])")), [row_modes_new; col_modes_new]), prod(dims[row_modes_new]), prod(dims[col_modes_new]))
	return s
end

"""
Unfolds the tensor into matrix such that the specified mode becomes matrix row.
"""
_row_unfold(tnsr::AbstractArray, mode::Integer) =
	_unfold(tnsr, [mode], [1:mode-1; mode+1:ndims(tnsr)])

"""
Unfolds the tensor into matrix such that the specified mode becomes matrix column.
"""
_col_unfold(tnsr::AbstractArray, mode::Integer) =
	_unfold(tnsr, [1:mode-1; mode+1:ndims(tnsr)], [mode])

function _iter_status(converged::Bool, niters::Integer, maxiter::Integer)
	converged ? @info("Algorithm converged after $(niters) iterations.") :
				@warn("Maximum number $(maxiter) of iterations exceeded.")
end

_check_sign(v::StridedVector) = sign(v[findmax(abs.(v))[2]]) * v

"""
Checks the validity of the core tensor dimensions.
"""
function _check_tensor(tnsr::AbstractArray{T, N}, core_dims::NTuple{N, Int}) where {T<:Real,N}
	ndims(tnsr) > 2 || throw(ArgumentError("This method does not support scalars, vectors, or matrices input."))
	for i in 1:N
		0 < core_dims[i] <= size(tnsr, i) ||
			throw(ArgumentError("core_dims[$i]=$(core_dims[i]) given, 1 <= core_dims[$i] <= size(tensor, $i) = $(size(tnsr, i)) expected."))
	end
	#isreal(T) || throw(ArgumentError("This package currently only supports real-number-valued tensors."))
	return N
end

"""
Checks the validity of the core tensor dimensions, where core tensor is `r^N` hypercube.
"""
_check_tensor(tensor::AbstractArray{<:Number, N}, r::Integer) where N =
	_check_tensor(tensor, ntuple(_ -> r, N))

function distribute(X::AbstractArray{T,N}) where {T <: Number, N}
	sz = size(X)
	@info "Dimension 1: Distributing array with size $(sz)"
	dX = DistributedArrays.distribute(X)
	global dX1 = dX
	for i = 2:N
		a = [i for i = 1:N]
		k = a .!= i
		v = [i, a[k]...]
		vb = "dX$i"
		@info "Dimension $i: Distributing array with size $(sz[v])"
		@eval(global $(Symbol(vb)) = DistributedArrays.distribute(permutedims($X, $v)))
	end
	return dX
end


@testset "projection scales" begin
    # inspired by Statutory Valuation of Life Insurance Liabilities, vol 5 chapter 21
    @testset "basic values" begin
        imp = MortalityTables.table("Projection Scale G2 – Male, ANB")
        @test imp[0] == 0.01
        @test imp[105] == 0.000
        @test imp[65] == 0.015
    end

    @testset "improved mortality" begin
        imp = MortalityTables.table("Projection Scale G2 – Male, ANB")
        mort = MortalityTables.table("2012 IAM Period Table – Male, ANB").ultimate

        @test mort[66] * prod(1 .- imp[65:65]) ≈ 8.41978 / 1000
        @test isapprox(mort[85] * prod(1 .- imp[65:84]), 44.60138 / 1000, atol=1e-5)

    end
end


"""
Constrained-optimisation-based control allocator.
"""
struct ConstrainedAllocator <: StaticAllocator
    B
    u_min
    u_max
    u
    function ConstrainedAllocator(B, u_min, u_max; p=Inf)
        @assert size(u_min) == size(u_max)
        dim_input = length(u_min)
        u = Convex.Variable(dim_input)
        new(B, u_min, u_max, u)
    end
end

"""
# Variables
ν: virtual input
# Notes
ν = B*u where u: control input
"""
function (allocator::ConstrainedAllocator)(ν, Λ=Diagonal(ones(size(ν)));
                                           p=Inf, silent_solver=true)
    @unpack u_min, u_max, u, B = allocator
    prob = minimize(
                    norm(u, p)
                    # sum(u)
                    + 1e5*norm(ν - B*Λ*u, 1)  # exact penalty method
                   )
    prob.constraints += [
                         u .>= u_min;  # min
                         u .<= u_max;  # max
                         # ν == B*Λ*u  # equality; replaced by exact penalty method
                        ]
    Convex.solve!(prob,
                  # SCS.Optimizer();
                  Mosek.Optimizer();  # best
                  silent_solver=silent_solver,
                  verbose=false)
    u.value[:]  # dim: n×1 -> n
end


using Soss, MeasureTheory

hmm = @model begin
    ε ~ Exponential() #  transition
    σ ~ Exponential() # Observation noise
    x ~ Chain(Normal()) do xj
        Normal(xj, ε)
    end

    y ~ For(x) do xj
        Normal(xj, σ)
    end
end

using Random

_rng = Random.GLOBAL_RNG

# sourceRand(hmm)

σ = rand(_rng, Exponential())
ε = rand(_rng, Exponential())
x = rand(_rng, Chain(Normal()) do xj
            #= REPL[2]:5 =#
            Normal(xj, ε)
        end)
y = rand(_rng, For(x) do xj
            #= REPL[2]:8 =#
            Normal(xj, σ)
        end)

rand(hmm())


# Dantzig-Wolfe Reformulation and Column Generation
# Edward J. Xu
# 2019.5.28
########################################################################################################################

function SSCFLRand(m::Int64,n::Int64, maxXY::Int64, minF::Int64, maxF::Int64, minD::Int64, maxD::Int64, seed)
    # m:  number of facilities
    # n: number of customers
    Random.seed!(seed)
    # Cost of opening facitlities
    f = rand(minF:maxF, m)
    # Coordinates for facilities (chosen in the box [0,maxXY]x[0,maxXY])
    fac = rand(0:maxXY, m, 2)
    # Coordinates for customers (chosen in the box [0,maxXY]x[0,maxXY])
    cust = rand(0:maxXY, n, 2)
    d = rand(minD:maxD, n)
    totDemand = sum(d)
    capPerFacility = ceil(1.3 * totDemand / m)
    q = capPerFacility*ones(m)

    c = zeros(m, n)
    # compute distances
    for i = 1:m
        for j = 1:n
            dist = round(norm(fac[i,1:2] - cust[j,1:2]))  # norm function is in package LinearAlgebra
            c[i,j] = dist
        end
    end
    return (c, f, q, d)
end

function solveSSCFL(c, f, q, d)
    (m,n) = size(c)
    timeStart = time()
    mod = Model(solver = GurobiSolver())
    #mod = Model(solver = CplexSolver())
    @variable(mod, x[1:m, 1:n], Bin )
    @variable(mod, y[1:m], Bin )
    @objective(mod, Min, sum(c[i,j] * x[i,j] for i=1:m for j=1:n) + sum(f[i]*y[i] for i=1:m))
    @constraint(mod, [j = 1:n], sum(x[i,j] for i=1:m) == 1)
    #@constraint(mod, [i=1:m,j=1:n], x[i,j] <= y[i])
    @constraint(mod, [i = 1:m], sum(d[j]*x[i,j] for j=1:n) <= q[i]y[i])
    # writeLP(mod, "SSCFLP.lp", genericnames = false)
    status = solve(mod, relaxation = true)
    println("#### Objective value, relaxation: ", getobjectivevalue(mod))
    println("#### y = ", getvalue(y))
    println("#### Elapsed time is $(time() - timeStart) seconds.\n",
            "################################################################################") ########################
    status = solve(mod)
    println("#### Objective value: ", getobjectivevalue(mod))
    #println("x = ", getvalue(x))
    #println("y = ", getvalue(y))
    println("#### Number of facilities opened: ", sum(getvalue(y)))
    println("#### Elapsed time is $(time() - timeStart) seconds.\n",
            "#### End #######################################################################") ########################
end

function getData(whi::Float64)
    if whi == 1  # Convexify􏰁 constraints into one sub-problem
        vec_c = - [292  453  359  219  268  736  291  443  403  498  400  967]
        mat_a = [
            Matrix{Float64}(I, 5, 5) zeros(5, 1) Matrix{Float64}(I, 5, 5) zeros(5, 1);
            14 20 6 16 10 -43 zeros(1, 6);
            zeros(1, 6) 14 20 6 16 10 -43
            ]
        vec_b = hcat([ones(5,1); 0; 0])
        matIndexSub = hcat(collect(1:12))'
        numSub = 1
        numXInSub = 12
        m_mat_h = 5
        vecRowMatD = [2]
    elseif whi == 2  # Convexify􏰁 constraints into two sub-problems
        vec_c = - [292  453  359  219  268  736  291  443  403  498  400  967]
        mat_a = [
            Matrix{Float64}(I, 5, 5) zeros(5, 1) Matrix{Float64}(I, 5, 5) zeros(5, 1);
            14 20 6 16 10 -43 zeros(1, 6);
            zeros(1, 6) 14 20 6 16 10 -43
            ]
        vec_b = hcat([ones(5,1); 0; 0])
        matIndexSub = Array{Int64}(undef, 2, 6)
        matIndexSub[1, :] = collect(1:6)
        matIndexSub[2, :] = collect(7:12)
        numSub = 2
        numXInSub = 6
        m_mat_h = 5
        vecRowMatD = [1, 1]
    elseif whi == 3  # Convexify􏰁 constraints into one sub-problems, restrain the y
        vec_c = - [292  453  359  219  268  736  291  443  403  498  400  967]
        mat_a = [
            Matrix{Float64}(I, 6, 6) Matrix{Float64}(I, 6, 6);
            14 20 6 16 10 -43 zeros(1, 6);
            zeros(1, 6) 14 20 6 16 10 -43
            ]
        vec_b = hcat([ones(5,1); 2; 0; 0])
        matIndexSub = Array{Int64}(undef, 2, 6)
        matIndexSub[1, :] = collect(1:6)
        matIndexSub[2, :] = collect(7:12)
        numSub = 1
        numXInSub = 12
        m_mat_h = 6
        vecRowMatD = [2]
    elseif whi == 4  # Convexify􏰁 constraints into two sub-problems, restrain the y
        vec_c = - [292  453  359  219  268  736  291  443  403  498  400  967]
        mat_a = [
            Matrix{Float64}(I, 6, 6) Matrix{Float64}(I, 6, 6);
            14 20 6 16 10 -43 zeros(1, 6);
            zeros(1, 6) 14 20 6 16 10 -43
            ]
        vec_b = hcat([ones(5,1); 2; 0; 0])
        matIndexSub = Array{Int64}(undef, 2, 6)
        matIndexSub[1, :] = collect(1:6)
        matIndexSub[2, :] = collect(7:12)
        numSub = 2
        numXInSub = 6
        m_mat_h = 6
        vecRowMatD = [1, 1]
    elseif whi >= 5
        if round(whi) == 5
            (c, f, q, d) = SSCFLRand(2, 5, 500, 500, 1000, 1, 20, 3)
        elseif round(whi) == 6
            (c, f, q, d) = SSCFLRand(10, 25, 500, 500, 1000, 1, 20, 3)
        elseif round(whi) == 7
            (c, f, q, d) = SSCFLRand(20, 50, 500, 500, 1000, 1, 20, 3)
        end
        vec_c = - reshape(hcat(c, f)', 1, :)
        (m, n) = size(c)
        mat_a_atom = zeros(m, m + n, n + 1)
        for i = 1:m
            mat_a_atom[i, :, :] = [
                Matrix{Float64}(I, n, n) zeros(n, 1);
                zeros(m, n + 1)
                ]
            mat_a_atom[i, n + i, :] = vcat(d, - q[i])
        end
        mat_a = zeros(n + m, m * (n + 1))
        for i = 1:m
            mat_a[:, ((i - 1) * (n + 1) + 1):(i * (n + 1))] = mat_a_atom[i, :, :]
        end
        vec_b = hcat(zeros(n + m, 1))
        vec_b[1:n, 1] = ones(n, 1)
        if ((whi == 5.1) | (whi == 6.1) | (whi == 7.1))
            matIndexSub = hcat(collect(1:(m * (n + 1))))'
            numSub = 1
            numXInSub = m * (n + 1)
            m_mat_h = n
            vecRowMatD = Array{Int64}(undef, 1)
            vecRowMatD = [m]
        elseif ((whi == 5.2) | (whi == 6.2) | (whi == 7.2))
            matIndexSub = Array{Int64}(undef, m, n + 1)
            for i = 1:m
                matIndexSub[i, :] = collect(((i - 1) * (n + 1) + 1):(i * (n + 1)))
            end
            numSub = m
            numXInSub = n + 1
            m_mat_h = n
            vecRowMatD = Array{Int64}(undef, 1)
            vecRowMatD = ones(1, m)
            println("#### 0,  Solve SSCFL Model using Gurobi#########################################") ####################
            solveSSCFL(c, f, q, d)
            println("#### End #######################################################################") ####################
        else
            throw("Wrong \"whi\" input!")
        end
    else
        throw("Wrong \"whi\" input!")
    end
    return (vec_c, mat_a, vec_b, matIndexSub, numSub, numXInSub, m_mat_h, vecRowMatD)
end


#=
# 2D Poisson, Dirichlet bc
# Uses Matlab imported as a custom gen target
=#
if !@isdefined(Femshop)
    include("../Femshop.jl");
    using .Femshop
end
init_femshop("poisson2dcustomnew");

# Try making an optional log
useLog("poisson2dcustomnewlog")

# Generate for the target in customtarget.jl
generateFor("target_matlab_cg.jl")

# Set up the configuration (order doesn't matter)
domain(2)
functionSpace(order=2)

# Specify the problem
mesh(QUADMESH, elsperdim=30)

u = variable("u")

testSymbol("v")

boundary(u, 1, DIRICHLET, "sin(3*pi*x)")

# Write the weak form
coefficient("f", "-2*pi*pi*sin(pi*x)*sin(pi*y)")
weakForm(u, "-dot(grad(u),grad(v)) - f*v")

solve(u);

finalize_femshop()


using LinearAlgebra
include("Pauli.jl")

function BlochRotation(axis,phase)

#Axis is the axis of rotation, should be a 1x3 array. It is not normalized,
#therefore being axis=angle*axis
#phase is the general phase

p=[Pauli(1), Pauli(2), Pauli(3)]
a=axis

exp(-phase*im)exp(-(p[1]*axis[1]+p[2]*a[2]+p[3]*a[3])/2*im)

end


using Grid
using PyPlot

f(x) = 2 .* cos(6x) .+ sin(14x) .+ 2.5
c_grid = 0:.2:1

Af = CoordInterpGrid(c_grid, f(c_grid), BCnil, InterpLinear)

f_grid = linspace(0, 1, 150)

fig = figure()
ax = fig[:add_subplot](111)

ax[:plot](f_grid, f(f_grid), "b-", lw=2, alpha=0.8, label="true function")
ax[:plot](f_grid, Af[f_grid], "g-", lw=2, alpha=0.8,
          label="linear approximation")

ax[:vlines]([c_grid], [c_grid] * 0, f(c_grid), linestyle="dashed", alpha=0.5)
ax[:legend](loc="upper center")


export Frobenius
"""
[`Frobenius` matrix](http://en.wikipedia.org/wiki/Frobenius_matrix)

Frobenius matrices or Gaussian elimination matrices of the form

```
[ 1 0 ...     0 ]
[ 0 1 ...     0 ]
[ .........     ]
[ ... 1 ...     ]
[ ... c1 1 ...  ]
[ ... c2 0 1 ...]
[ ............. ]
[ ... ck ...   1]
```

i.e. an identity matrix with nonzero subdiagonal elements along a single column.

```
 In this implementation, the subdiagonal of the nonzero column is stored as a
 dense vector, so that the size can be inferred automatically as j+k where j is
 the index of the column and k is the number of subdiagonal elements.
```

```julia
julia> using SpecialMatrices

julia> F=Frobenius(3, [1.0,2.0,3.0]) #Specify subdiagonals of column 3
6x6 Frobenius{Float64}:
 1.0  0.0  0.0  0.0  0.0  0.0
 0.0  1.0  0.0  0.0  0.0  0.0
 0.0  0.0  1.0  0.0  0.0  0.0
 0.0  0.0  1.0  1.0  0.0  0.0
 0.0  0.0  2.0  0.0  1.0  0.0
 0.0  0.0  3.0  0.0  0.0  1.0

julia> inv(F) #Special form of inverse
6x6 Frobenius{Float64}:
 1.0  0.0   0.0  0.0  0.0  0.0
 0.0  1.0   0.0  0.0  0.0  0.0
 0.0  0.0   1.0  0.0  0.0  0.0
 0.0  0.0  -1.0  1.0  0.0  0.0
 0.0  0.0  -2.0  0.0  1.0  0.0
 0.0  0.0  -3.0  0.0  0.0  1.0

julia> F*F #Special form preserved if the same column has the subdiagonals
6x6 Frobenius{Float64}:
 1.0  0.0  0.0  0.0  0.0  0.0
 0.0  1.0  0.0  0.0  0.0  0.0
 0.0  0.0  1.0  0.0  0.0  0.0
 0.0  0.0  2.0  1.0  0.0  0.0
 0.0  0.0  4.0  0.0  1.0  0.0
 0.0  0.0  6.0  0.0  0.0  1.0

julia> F*Frobenius(2, [5.0,4.0,3.0,2.0]) #Promotes to Matrix
6x6 Array{Float64,2}:
 1.0   0.0  0.0  0.0  0.0  0.0
 0.0   1.0  0.0  0.0  0.0  0.0
 0.0   5.0  1.0  0.0  0.0  0.0
 0.0   9.0  1.0  1.0  0.0  0.0
 0.0  13.0  2.0  0.0  1.0  0.0
 0.0  17.0  3.0  0.0  0.0  1.0

julia> F*[10.0,20,30,40,50,60.0]
6-element Array{Float64,1}:
  10.0
  20.0
  30.0
  70.0
 110.0
 150.0
```
"""
struct Frobenius{T} <: AbstractArray{T, 2}
    colidx :: Int
    c :: Vector{T}
end

#Basic property computations
size(F::Frobenius, r::Int) = (r==1 || r==2) ? F.colidx + length(F.c) :
    throw(ArgumentError("Frobenius matrix is of rank 2"))

function size(F::Frobenius)
    n = F.colidx + length(F.c)
    n, n
end

#XXX Inefficient but works
getindex(F::Frobenius, i, j) = getindex(Matrix(F), i, j)
isassigned(F::Frobenius, i::Integer, j::Integer) = isassigned(Matrix(F), i, j)

function Matrix(F::Frobenius{T}) where T
    M = Matrix{T}(I, size(F))
    M[F.colidx+1:end, F.colidx] = F.c
    M
end

#Linear algebra stuff
function mul!(F::Frobenius{T}, b::Vector{T}) where T
    (n = size(F, 2)) == length(b) || throw(DimensionMismatch("$n $(length(b))"))
    for i=F.colidx+1:n
        b[i] += F.c[i-F.colidx] * b[F.colidx]
    end
    b
end
*(F::Frobenius{T}, b::Vector{T}) where T = mul!(F, copy(b))

function *(F::Frobenius{T}, G::Frobenius{T}) where T
    (n = size(F, 2)) == size(G, 2) || throw(DimensionMismatch(""))
    if F.colidx == G.colidx #Answer is still expressible as a Frobenius
        return Frobenius(F.colidx, F.c + G.c)
    else
        M = Matrix(F)
        M[G.colidx+1:end, G.colidx] = F.colidx < G.colidx ? G.c :
            Frobenius(F.colidx-G.colidx, F.c) * G.c
        return M
    end
end

inv(F::Frobenius) = Frobenius(F.colidx, -F.c)


@doc doc"""
XY model with energy $E = -\sum_{ij} J_{ij} \cos(\theta_i - \theta_j)$,
where $\theta_i = 2\pi \sigma_i$ and $\sigma_i \in [0, 1)$.
"""
mutable struct XY <: Model
    lat :: Lattice
    spins :: Matrix{Float64}
    rng :: Random.MersenneTwister

    function XY(lat::Lattice, rng::Random.AbstractRNG)
        spins = rand(rng, 1, numsites(lat))
        return new(lat, spins, rng)
    end
end

XY(lat::Lattice) = XY(lat, Random.seed!(Random.MersenneTwister(0)))
XY(lat::Lattice, seed) = XY(lat, Random.seed!(Random.MersenneTwister(0), seed...))

@doc doc"""
   XY(param)

Generates `XY` using `param["Lattice"]`,  and `param["Seed"]` (if defined).
Each spin $\sigma_i$ will be initialized randomly and independently.
"""
function XY(param::Parameter)
    lat = generatelattice(param)
    if "Seed" in keys(param)
        return XY(lat, param["Seed"])
    else
        return XY(lat)
    end
end

include("update.jl")
include("estimator.jl")
include("postproc.jl")


#############################################################################
# union.jl
# Handles unions of combinatorial sets.
#############################################################################

import Base: union

export union, get_elements

mutable struct UnionAtom <: CombiSet
  head::Symbol
  id_hash::UInt64
  children::Tuple
  elements::ValOrNothing
  baseset::AbstractArray
  cardinality::Int

  function UnionAtom(elements::ValOrNothing, children::Tuple, baseset::AbstractArray)
    this = new(:union, 0, children, elements, baseset, length(baseset))
    this.id_hash = object_id(this)
    return this
  end
end

function union(set::AllCombiSet...)
  sets = collect(set)
  baseset = Set([])
  n = length(sets)
  if n < 2
    error("Cannot compute the union of one set.")
  else
    if isa(sets, Array{Number})
      elements = []
      for i = 1:n
        elements = vcat(elements, sets[i])
      end
      elements = collect(Set(elements))
      return elements
    else
      for i = 1:n
        if isa(sets[i], CombiSet)
          baseset = union(baseset, Set(sets[i].baseset))
        else
          seti = Set(sets[i])
          baseset = union(baseset, seti)
        end
      end
      baseset = collect(baseset)
      newset = UnionAtom(nothing, set, baseset)
      return newset
    end
  end
end

function get_elements(x::UnionAtom)
  elements = []
  n = length(x.children)
  for i = 1:n
    if isa(x.children[i], CombiSet)
      elements = union(elements, Set(get_elements(x.children[i])))
    else
      seti = Set(x.children[i])
      elements = union(elements, seti)
    end
  end
  x.elements = collect(elements)
  return x.elements
end


module RobustnessEvaluation

using Flux
using Statistics
using NeuralVerification
using Plots
using Measures
using Random
using ProgressMeter
using NeuralVerification: compute_output
using NeuralVerification: init_vars, BoundedMixedIntegerLP, encode_network!, _ẑᵢ₊₁, TOL
using GLPK
using JuMP
using LazySets
using StatsBase
import JuMP.MOI.OPTIMAL, JuMP.MOI.INFEASIBLE, JuMP.MOI.INFEASIBLE_OR_UNBOUNDED

struct RobustnessSummary
    samples
    errors_nominal
    mae_nominal
    errors_adversarial_true
    mae_adversarial_true
    errors_adversarial_pred
    mae_adversarial_pred
    saliency_maps
    ig_maps
end

export nominal_errors
include("nominal.jl")

export adversarial_errors
include("adversarial.jl")

export saliency_map, integrated_gradients, get_mask
include("explanations.jl")

export plot_summary, plot_nominal_errors, plot_adversarial_errors, plot_adversarial_perturbations, plot_explanations
include("plotting.jl")

export robustness_summary

function robustness_summary(network, X, Y; N_samples = 100, coeffs = [-0.74, -0.44],
    δ = 0.02, solver = GLPK.Optimizer, use_gurobi = false, PGD = false,
    step_factor = 0.1, iterations = 600, baseline = zeros(128))
    """ Runs all functions to get robustness summary for a given network and dataset
        Args:
            - network: neural network to analyze (flux model)
            - X: training examples (each column is an example)
            - Y: labels (each column is a label)
            ---------------
            - N_samples: number of samples to calculate on (actual indices will be chosen randomly)
            - coeffs: coefficients to get output (linear combination of network outputs)
            - δ: radius for hyperrectangle to check error within
            - solver: solver for mixed integer programs (ignored if PGD true)
            - use_gurobi: whether we are using gurobi (still need to put in solver)
            - PGD: whether or not to use PGD (if not, solve exactly with MIP Verify)
            - step_factor: step factor for PGD (ignored if PGD false)
            - iterations: number of steps for PGD (ignored if PGD false)
            - baseline: IG baseline (should be same size as input)
    """
    N = size(Y, 2)
    nnv_network = NeuralVerification.network(network)

    # Get samples to work with
    samples = randperm(N)[1:N_samples]
    X_samples = X[:, samples]
    Y_samples = Y[:, samples]

    # Nominal
    println("Calculating nominal errors...")
    errors_nominal = nominal_errors(nnv_network, X_samples, Y_samples, coeffs = coeffs)
    mae_nominal = mean(errors_nominal)

    # Adversarial
    println("Calculating adversarial errors...")
    errors_adversarial_true, errors_adversarial_pred = adversarial_errors(nnv_network,
        X_samples, Y_samples, coeffs = coeffs, δ = δ, solver = solver,
        use_gurobi = use_gurobi, PGD = PGD, step_factor = step_factor,
        iterations = iterations)
    mae_adversarial_true = mean(errors_adversarial_true)
    mae_adversarial_pred = mean(errors_adversarial_pred)

    # Explanations
    println("Getting explanations...")
    saliency_maps = zeros(size(X_samples))
    ig_maps = zeros(size(X_samples))
    for i = 1:N_samples
        saliency_maps[:, i] = saliency_map(network, X[:, i], coeffs = coeffs)
        ig_maps[:, i] = integrated_gradients(network, X[:, i], baseline, coeffs = coeffs)
    end

    return RobustnessSummary(samples, errors_nominal, mae_nominal, errors_adversarial_true,
        mae_adversarial_true, errors_adversarial_pred, mae_adversarial_pred,
        saliency_maps, ig_maps)
end

end # module


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

using Markdown
using InteractiveUtils

# ╔═╡ 413bc522-7d27-11eb-0c68-e9b6228728c3
begin
	using LinearAlgebra;
	using Random;
	using Plots;
end

# ╔═╡ 046e441e-7d39-11eb-128d-9965981464f1
begin
using PyCall
np = pyimport("numpy")
end

# ╔═╡ 919e2ef6-7d24-11eb-3ebd-8f770adc87c9
md"""
# Volatility is rough (2014)

[https://arxiv.org/pdf/1410.3394.pdf](https://arxiv.org/pdf/1410.3394.pdf)

Log-prices $Y_{t}$ often modelled as a semi-martingale

$dY_{t} = \mu_{t} dt + \sigma dW_{t}$

+ Dupire Local Volatility - $\sigma(Y_{t}, t)$ is a deterministic function of state $Y_{t}$ and time
+ Stochastic Volatility - Hull and White, Heston, SABR
+ Neither model the implied volatility surface well

Volatility as a *long memory* process

This makes Fractional Brownian Motion possibly better suited to capture this, given long-memory given the Hurst Exponent h < 1/2

$\mathbb{E}[|W_{t + \Delta}^{H} - W_{t}^{H}|^{q}] = \mathbb{E}[|Z|^{q}] \Delta^{q H}$

$\psi(\tau) := |\frac{\partial}{\partial k} \sigma_{BS}(k , \tau)|_{k = 0} , k = log(S/K)$

"""

# ╔═╡ 222d672a-7d25-11eb-00ab-3f2572366f63
begin

function cov_at(s, t, H)
	(abs(s) ^ (2 * H) + abs(t) ^ (2 * H) - abs(s - t) ^ (2 * H)) / 2
end

N = 100
H = (1:9) ./ 10 #hurst parameter h < 1/2 negatively corrleated;  h > 1/2 positively correlated
dt = 0.1
rng = MersenneTwister(1234);
p = plot()
for h = H
	cov_matrix = [cov_at(i, j, h) for i = (1:N) .* dt for  j = (1:N) .* dt]
	cov_matrix = reshape(cov_matrix, N, N)
	C = cholesky(cov_matrix)
	Ws = randn(rng, Float32, N)
	Bs = C.L * Ws
	print(Bs)
	plot!(p, (1:N) .* dt, Bs, label="H=$h");
end
p
end

# ╔═╡ 2f512d8c-7d2b-11eb-2f62-4fa0a8e0e102
md"""
ON  VIX  FUTURES  IN  THE  ROUGH  BERGOMI  MODEL

https://arxiv.org/pdf/1701.04260.pdf

Functional central limit theorems for rough volatility

https://arxiv.org/abs/1711.03078

+ https://github.com/search?p=2&q=rough+volatility&type=Repositories

Deep Learning Volatility

ON  VIX  FUTURES  IN  THE  ROUGH  BERGOMI  MODEL

The quadratic rough heston model and the joint s&p 500/vixsmile calibration problem

Turbo-Charged Rough Bergomi

Rough Volatility with Python: https://tpq.io/p/rough_volatility_with_python.html

Perfect hedging in rough Heston models

https://arxiv.org/pdf/1703.05049.pdf

Modelling Volatility with FractionalBrownian Motion and the RFSV Model, Bachelor Thesis

http://www.peterhinz.net/pdf/Bachelor_Thesis.pdf

Turbocharging Monte Carlo pricingfor the rough Bergomi model

https://arxiv.org/pdf/1708.02563.pdf

Christian Bayer, Peter Friz, Jim Gatheral, Pricing Under Rough Volatility (2015) https://papers.ssrn.com/sol3/papers.cfm?abstract_id=2554754

Portfolio Optimization under Fast Mean-Reverting and Rough Fractional Stochastic Environment https://www.tandfonline.com/doi/abs/10.1080/1350486X.2019.1584532

Martin Forde, Hongzhong Zhang, Asymptotics for rough stochasticvolatility models (2016) https://arxiv.org/abs/1610.08878

Rough Volatility Modelling (MSc Thesis), https://www.duo.uio.no/bitstream/handle/10852/64079/ThesisXiaoyanZhang.pdf?sequence=1&isAllowed=y

https://sites.google.com/site/roughvol/home/risks-1

Rough stochastic volatility and applications of the rough Bergomi model (MSc Thesis) https://www.imperial.ac.uk/media/imperial-college/faculty-of-natural-sciences/department-of-mathematics/math-finance/TEMEL_TURKER_01421814.pdf

"""

# ╔═╡ Cell order:
# ╠═919e2ef6-7d24-11eb-3ebd-8f770adc87c9
# ╠═413bc522-7d27-11eb-0c68-e9b6228728c3
# ╠═222d672a-7d25-11eb-00ab-3f2572366f63
# ╠═046e441e-7d39-11eb-128d-9965981464f1
# ╠═2f512d8c-7d2b-11eb-2f62-4fa0a8e0e102


############################################################################
#                        TEST FUNCIONES DE ACTIVACIÓN
# Comprobamos que satisfacen alguno de sus requisitos teóricos
# Así como que al evaluarlas no de ningún error
############################################################################
using Test

######################### TEST  #########################
@testset "Activations functions" begin

    id(x)=x
    # Propiedas asintóticas funciones clásicas
    @test @ThresholdFunction(id,0)(-1) ≈ -1
    @test @ThresholdFunction(id,0)(1) ≈ 1
    @test CosineSquasher(-100) ≈ 0
    @test CosineSquasher(10) ≈ 1
    @test @IndicatorFunction(0)(-0.1) == 0
    @test @IndicatorFunction(0)(0.1) == 1
    @test RampFunction(-1)== 0
    @test RampFunction(2) == 1
    @test ReLU(-1)== 0
    @test ReLU(2) >0
    @test Sigmoid(9999999) ≈ 1
    @test Sigmoid(-9999999) ≈ 0
    @test HardTanh(9999999) ≈ 1
    @test HardTanh(-9999999) ≈ -1
    @test @LReLU(0.01)(10) == 10
    @test @LReLU(-0.01)(-1) == 0.01

    # Medibles en la práctica equivale a que tome reales
    valores_compacto = [ (rand()-0.5)*1000 for _ in 1:10]

    funciones_activacion = [
        CosineSquasher,
        RampFunction,
        ReLU,
        Sigmoid,
        HardTanh
    ]
    # Comprobamos funciones de activación que no son macros
    for f in funciones_activacion
        for x in valores_compacto
            @test typeof(f(x)) <: Real
        end
    end
    # Las macro me sale error al meternas en la lista :(
    for x in valores_compacto
        @test typeof(@ThresholdFunction(id,0)(x)) <: Real
    end
    for x in valores_compacto
        @test typeof(@IndicatorFunction(0)(x)) <: Real
    end
    for x in valores_compacto
        @test typeof(@LReLU(0.01)(x)) <: Real
    end
end


struct Dataset{T<:Real}

    X::AbstractMatrix{T}
    y::AbstractVector{T}
    n::Int
    nfeatures::Int
    weighted::Bool
    weights::Union{AbstractVector{T}, Nothing}
    varMap::Array{String, 1}

end

"""
    Dataset(X::AbstractMatrix{T}, y::AbstractVector{T};
            weights::Union{AbstractVector{T}, Nothing}=nothing,
            varMap::Union{Array{String, 1}, Nothing}=nothing)

Construct a dataset to pass between internal functions.
"""
function Dataset(
        X::AbstractMatrix{T},
        y::AbstractVector{T};
        weights::Union{AbstractVector{T}, Nothing}=nothing,
        varMap::Union{Array{String, 1}, Nothing}=nothing
       ) where {T<:Real}

    n = size(X, 2)
    nfeatures = size(X, 1)
    weighted = weights !== nothing
    if varMap == nothing
        varMap = ["x$(i)" for i=1:nfeatures]
    end

    return Dataset{T}(X, y, n, nfeatures, weighted, weights, varMap)

end


using GeometricFlux
using GraphSignals
using Graphs
using SparseArrays
using Plots
using GraphPlot
using Clustering
using Cairo, Compose

clusters = [1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 1, 1, 1, 1, 2, 2, 1, 1, 2, 1, 2, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2]

int2col_str(x::Int) = x==1 ? "lightblue" : "red"

g = smallgraph(:karate)
fg = FeaturedGraph(g)
vectors = node2vec(fg; walks_per_node=10, len=80, p=1.0, q=1.0)
R = kmeans(vectors, 2)

learned_clusters = copy(assignments(R))
# ensure that the cluster containing node 1 is cluster 1
if assignments(R)[1] != 1
    learned_clusters = [i == 1 ? 2 : 1 for i in assignments(R)]
end

output_plot_name = "karateclub.pdf"
draw(
    PDF(output_plot_name, 16cm, 16cm),
    gplot(g,
        nodelabel=map(string, 1:34),
        nodefillc=[int2col_str(learned_clusters[i]) for i in 1:34],
        nodestrokec=["white" for _ in 1:34]
    )
)

incorrect = sum(learned_clusters .!= clusters)
println(incorrect, " incorrect cluster labelings")
println("Drawn graph to ", output_plot_name)


using ITensors,
      Test

import ITensors: In, Out, Neither

@testset "QN Index" begin

  @testset "hasqns function" begin
    i = Index(4,"i")
    @test hasqns(i) == false
    j = Index(QN(0)=>1,QN(1)=>1)
    @test hasqns(j) == true
  end

  @testset "Array of QN Constructor" begin
    i = Index([QN(0)=>1,QN(1)=>2],"i")
    @test hasqns(i)
    @test dim(i) == 3
    @test hastags(i,"i")
  end

  @testset "Vararg Constructor" begin
    i = Index(QN(0)=>1,QN(1)=>2;tags="i")
    @test hasqns(i)
    @test dim(i) == 3
    @test hastags(i,"i")
    @test dir(i) == Out

    j = Index(QN(0)=>1,QN(1)=>2;tags="j",dir=In)
    @test hasqns(j)
    @test dim(j) == 3
    @test hastags(j,"j")
    @test dir(j) == In
  end

  @testset "flux and qn" begin
    i = dag(Index([QN(0)=>2, QN(1)=>2], "i"))

    @test flux(i => 1) == QN(0)
    @test flux(i => 2) == QN(0)
    @test flux(i => 3) == QN(-1)
    @test flux(i => 4) == QN(-1)
    @test flux(i => Block(1)) == QN(0)
    @test flux(i => Block(2)) == QN(-1)

    @test qn(i => 1) == QN(0)
    @test qn(i => 2) == QN(0)
    @test qn(i => 3) == QN(1)
    @test qn(i => 4) == QN(1)
    @test qn(i => Block(1)) == QN(0)
    @test qn(i => Block(2)) == QN(1)

  end
end

nothing


using Distributed
addprocs(8)
@everywhere include("reproducibility.jl")
include("reproducibility.jl")

##
@everywhere using LightGraphs
@everywhere using LinearAlgebra
@everywhere using SparseArrays
@everywhere using Expokit
using LambertW
using Base.Iterators
using NPZ

##

@everywhere function hamiltonian_and_normed_eigen(g)
    h = -Matrix{Float64}(laplacian_matrix(g))
    eigdec = eigen(h)
    # eigenvalues are nonnegative
    h ./= abs(minimum(eigdec.values))
    eigdec.values ./= abs(minimum(eigdec.values))
    # eigenvalues are in [-1,0]
    eigdec.values .+= 1
    h += I
    # eigenvalues are in [0,1], spectral gap next to 1

    λ1, λ2, λn = eigdec.values[[end, end-1, 1]]
    a = -(λ2 + λn)/2
    eigdec.values .+= a
    h += a*I
    eigdec.values ./= 1 + a
    h ./= 1 + a
    sparse(h), eigdec
end

@everywhere function true_probability(h, eigdec)
    n = size(h, 1)
    a = eigdec.vectors[1,1:end-1]
    eigs = eigdec.values[1:end-1]

    gamma = sum(abs2.(a) ./ (1 .- eigs)) / sum(abs2.(a))
    opt_time = pi*sqrt(n)/2
    #opt_time = pi/2/abs(a[1,end]) * sqrt(sum(abs2.(a) ./ (1 .- eigs).^2))

    init_state = fill(1/sqrt(n), n)
    hamiltonian = gamma*h
    hamiltonian[1,1] += 1.
    f_state = expmv(opt_time, -1im*hamiltonian, init_state)
    abs2(f_state[1])
end

@everywhere function create_data(n::Int, p::Float64)
    g = erdos_renyi(n, p)
    cc = connected_components(g)
    g = induced_subgraph(g, cc[findmax(length.(cc))[2]])[1]
    @assert is_connected(g)
    @assert nv(g) > n/2

    h, eigdec = hamiltonian_and_normed_eigen(g)
    c = eigdec.values[end-1]

    (1-c)/(1+c), true_probability(h, eigdec)
end

p0s = .5:.5:2
@everywhere p(p0, n) = p0*log(n)/n
repeating = 50

data_n = 100:100:2000

for p0 = p0s
    data = zeros(length(data_n), repeating, 2)
    for (ind_n, n) = enumerate(data_n)
        println("n $n, p0 $p0")
        @time data_tmp = pmap(_ -> create_data(n, p(p0, n)), 1:repeating)
        for (m, d_tmp) = enumerate(data_tmp)
            data[ind_n,m,:] .= d_tmp
        end
    end
    npzwrite("hiding_search/sparse_ER_probs_and_bounds_$p0.npz", data)
end
rmprocs()


using ConditionalMean
using Base.Test

# use all(isapprox.(x,y)) elementwise isapprox of arrays

A=zeros(5,5)
A[:]=1:25
# all finite values

# write your own tests here
@time condmean(A[:],isfinite)[1] == mean(A)
@test condmean(A[:],isfinite)[1] == mean(A)
@time condmean(A[:],isfinite)[1] == mean(A)

f(x)=2x
c(x)=x>5
@test condmean(f,A,c)[2:5] == condmean(f,A,c,1)[2:5]
@test isnan(condmean(f,A,c)[1])
@test condmean(f,A,c,2) == mean(2*A,2)+5 # this fails for

# use a missing value
mv=-999.0
B=zeros(5,5)
B[:]=A[:]
B[3,:]=mv
@test condmean(B,mv,1) == mean(A,1)
B[:]=A[:]
B[2,:]=mv
@test condmean(B,mv,1) == mean(A,1)+0.25

# test using NaN and mapping the function through the mean
B[:]=A[:]
B[3,:]=NaN
# easy to predict for a linear operator
@test all(isapprox.(condmean(x->2*x,B,isfinite,1), 2*[3.0 8.0 13.0 18.0 23.0]))

# test the offsetting version
@show condmean(A,isfinite,2)
@show mean(A,2)
@test condmean(A,isfinite,2) == mean(A,2)

# test precision issues; first no missing values
B=Array{Float64}(A)
B[:]=A[:]+1e6
Bs=Array{Float32}(B)

# first single precision instance requires compilation
@time residual32 = (condmean(Bs[:],NaN)[1] - 1e6 - 13.0)
@time residual32 = (condmean(Bs[:],NaN)[1] - 1e6 - 13.0)
@time residual32 = (condmean(Bs[:],NaN)[1] - 1e6 - 13.0)
@show residual32

@time residual64 = (condmean(B[:],NaN)[1] - 1e6 - 13.0)
@time residual64 = (condmean(B[:],NaN)[1] - 1e6 - 13.0)
@time residual64 = (condmean(B[:],NaN)[1] - 1e6 - 13.0)
@show residual64

if     !(residual64 ≈ Float64(0.0))
  warn("ConditionalMean: condmean imprecise for double precision.")
elseif !(residual32 ≈ Float32(0.0))
  warn("ConditionalMean: condmean imprecise for single precision.")
end
# require to be precise at double precision
@test residual64 ≈ Float64(0.0)

# test precision with missing value
B[3,:]=Float32(mv)
@test condmean(B[:],mv)[1] - 1e6 ≈ 13.0
@test all(isapprox.(condmean(B,mv,1) - 1e6, [3.0 8.0 13.0 18.0 23.0]))
# with NaN
B[3,:]=Float32(NaN)
@test all(isapprox.(condmean(B,isfinite,1) - 1e6, [3.0 8.0 13.0 18.0 23.0]))


zero_like(x) = x - x

type Beam
    """Complex beam parameter."""
    q

    """Wavelength."""
    λ
end

beam(waist_radius, offset, λ) = Beam(-offset + 1im * rayleigh_range(waist_radius, λ), λ)

rayleigh_range(b::Beam) = imag(b.q)
rayleigh_range(waist_radius, λ) = π * waist_radius^2 / λ

waist_offset(b::Beam) = -real(b.q)
waist_radius(b::Beam) = sqrt(rayleigh_range(b) * b.λ / π)

radius(b::Beam) = sqrt(-b.λ / (π * imag(1 / b.q)))
curvature(b::Beam) = 1 / real(1 / b.q)

abstract type Element end

struct Lens <: Element
    focal_length
end

struct FreeSpace <: Element
    distance
end

function expand_system(positions_and_elems, up_to_pos=nothing)
    last_pos = zero_like(up_to_pos == nothing ? positions_and_elems[1][1] : up_to_pos)
    system = Element[]
    for (p, e) in sort(positions_and_elems, by=a -> a[1])
        if up_to_pos != nothing && p > up_to_pos
            break
        end
        delta_pos = p - last_pos
        if delta_pos != zero_like(p)
            push!(system, FreeSpace(delta_pos))
        end
        last_pos = p

        push!(system, e)
    end
    if up_to_pos != nothing && last_pos < up_to_pos
        push!(system, FreeSpace(up_to_pos - last_pos))
    end
    system
end

function apply_abcd(beam::Beam, matrix)
    Beam((matrix[1, 1] * beam.q + matrix[1, 2]) /
        (matrix[2, 1] * beam.q + matrix[2, 2]), beam.λ)
end

function apply(beam::Beam, elem::Lens)
    apply_abcd(beam, [1 zero_like(elem.focal_length); (-1 / elem.focal_length) 1])
end

function apply(beam::Beam, elem::FreeSpace)
    apply_abcd(beam, [1 elem.distance; zero_like(1 / elem.distance) 1])
end

function propagate{E<:Element}(elements::Array{E, 1}, beam::Beam)
    reduce(apply, beam, elements)
end

export Beam, beam, Element, Lens, FreeSpace, expand_system, propagate


#Written By Johnathan Bizzano
exppow(x) = MathConstants.e ^ x
softmax(s, x) = x / s

"multinomial logistic regression and is often used as the last activation function of a neural network to
normalize the output of a network to a probability distribution over predicted output classes - Wikipedia"
function Softmax()
    return LayerGenerator(
        function(pb, sett, input_shape)
                LayerDesign(Sig(:SoftMaxLayer), input_shape, input_shape, (0), nothing,
                    function (ins, outs, cons)
                            return function ()
                                        broadcast!(exppow, ins, ins)
                                        broadcast!(softmax, outs, sum(ins), ins)
                                   end
                    end)
        end)
end

export SoftmaxLayer
SoftMaxLayer() = SoftMax()


#
# Squared distance function, considering Periodic Boundary Condition
#

function r2pbc(p1, p2, bdim)
    r2 = 0.
    for i = 1:2
        Δx = abs(p1[i] - p2[i]) % bdim
        if Δx > bdim/2
            Δx -= bdim
        end
        r2 += Δx^2
    end
    return r2
end


### A Pluto.jl notebook ###
# v0.12.18

using Markdown
using InteractiveUtils

# ╔═╡ 92081fb8-5d90-11eb-2078-ddbf87421051
begin
	using ReinforcementLearning
	using Flux
	using Statistics
	using Plots
	using Distributions
end

# ╔═╡ 3254385c-5d90-11eb-2c99-07a3a40ae467
md"""
Again, we'll describe the car rental problem with a distributional model.
"""

# ╔═╡ a1a5176c-5d90-11eb-2ad6-2f108e1531e2
begin
	const PoissonUpperBound = 10
	const MaxCars= 20
	const MaxMoves = 5
	const CostPerMove = 2
	const CarRentalCartesianIndices = CartesianIndices((0:MaxCars,0:MaxCars))
	const CarRentalLinearIndices = LinearIndices((0:MaxCars,0:MaxCars))
	const Actions = -MaxMoves:MaxMoves
	const RequestDist_1 = Poisson(3)
	const RequestDist_2 = Poisson(4)
	const ReturnDist_1 = Poisson(3)
	const ReturnDist_2 = Poisson(2)

	decode_state(s::Int) = Tuple(CarRentalCartesianIndices[s])
	encode_state(s1::Int, s2::Int) = CarRentalLinearIndices[CartesianIndex(s1+1, s2+1)]
	decode_action(a::Int) = a - MaxMoves - 1
	encode_action(a::Int) = a + MaxMoves + 1

	function merge_prob(dist)
		merged = Dict()
		for (s′, r, p) in dist
			if haskey(merged, (s′, r))
				merged[(s′, r)] += p
			else
				merged[(s′, r)] = p
			end
		end
		[(r, false, s′) => p for ((s′, r), p) in merged]
	end

	function nextstep(state::Int, action::Int)
		(s1, s2), a = decode_state(state), decode_action(action)
		move = a > 0 ? min(a, s1) : max(a, -s2)
		reward = -CostPerMove*abs(move)
		s1′, s2′ = min(s1 - move, MaxCars), min(s2 + move, MaxCars)
		merge_prob(
			(
				encode_state(
					min(max(s1′-req_1, 0)+ret_1, MaxCars),
					min(max(s2′-req_2, 0)+ret_2, MaxCars)
					),
				reward + (min(s1′, req_1) + min(s2′, req_2)) * 10,
				(
					pdf(RequestDist_1, req_1) *
					pdf(RequestDist_2, req_2) *
					pdf(ReturnDist_1, ret_1) *
					pdf(ReturnDist_2, ret_2)
				)
			)
			for req_1 in 0:PoissonUpperBound,
				req_2 in 0:PoissonUpperBound,
				ret_1 in 0:PoissonUpperBound,
				ret_2 in 0:PoissonUpperBound
		)
	end

	struct CarRentalEnvModel <: AbstractEnvironmentModel
		dist
	end

	function CarRentalEnvModel()
		CarRentalEnvModel(
			Dict(
				(s,a) => nextstep(s,a)
				for s in 1:(MaxCars+1)^2 for a in 1:length(Actions)
			)
		)
	end

	(m::CarRentalEnvModel)(s, a) = m.dist[(s, a)]

	RLBase.state_space(m::CarRentalEnvModel) = Base.OneTo((MaxCars+1)^2)
	RLBase.action_space(m::CarRentalEnvModel) = Base.OneTo(length(Actions))
end

# ╔═╡ 20e3c8d6-5d92-11eb-1591-9911e62805bd
model = CarRentalEnvModel()

# ╔═╡ 0c099578-5d92-11eb-2a8b-2f4889026185
V = TabularVApproximator(n_state=length(state_space(model)), opt=Descent(1.0))

# ╔═╡ dda27350-5d94-11eb-1a82-c302513da719
p = TabularPolicy(;table=Dict(s=>1 for s in state_space(model)), n_action=length(action_space(model)))

# ╔═╡ 021f9b90-5d95-11eb-303a-ed2fb3de3d08
policy_iteration!(;V=V, π=p, model=model, γ=0.9, max_iter=300)

# ╔═╡ 8bbe08e6-5d98-11eb-0c02-4712c4659830
heatmap(0:MaxCars, 0:MaxCars, reshape([decode_action(p(x)) for x in state_space(model)], 1+MaxCars,1+MaxCars))

# ╔═╡ 91d90bf6-5d98-11eb-23a5-4dfc2839a707
heatmap(0:MaxCars, 0:MaxCars, reshape(V.table, 1+MaxCars,1+MaxCars))

# ╔═╡ Cell order:
# ╟─3254385c-5d90-11eb-2c99-07a3a40ae467
# ╠═92081fb8-5d90-11eb-2078-ddbf87421051
# ╠═a1a5176c-5d90-11eb-2ad6-2f108e1531e2
# ╠═20e3c8d6-5d92-11eb-1591-9911e62805bd
# ╠═0c099578-5d92-11eb-2a8b-2f4889026185
# ╠═dda27350-5d94-11eb-1a82-c302513da719
# ╠═021f9b90-5d95-11eb-303a-ed2fb3de3d08
# ╠═8bbe08e6-5d98-11eb-0c02-4712c4659830
# ╠═91d90bf6-5d98-11eb-23a5-4dfc2839a707


### A Pluto.jl notebook ###
# v0.15.1

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ a7ef92d7-41c1-4f7d-a0f9-b8037237af4f
begin
	using Pkg
	Pkg.activate(mktempdir())
	Pkg.add(url="https://github.com/mkitti/GaussianBeamLens.git")
	Pkg.add("GR")
	Pkg.add("PlutoUI")
end

# ╔═╡ bca271a4-ee26-48ab-bef8-e6c28ae76351
begin
	using GaussianBeamLens
	using GR
	using PlutoUI
end

# ╔═╡ 10e09fe7-cc23-4487-b1dc-b0a545bf2489
begin
	#cm = GR.COLORMAP_INFERNO
	colormaps = [
			string(getfield(GR,name)) => replace(string(name),"COLORMAP_" => "")
			for name in names(GR, all=true)
			if startswith(string(name),"COLORMAP_")
				]
	println(colormaps)
	md"""
	Colormap:
		$(@bind cm Select(colormaps; default = "45"))
	"""
end

# ╔═╡ b1e13be1-d608-43a1-a7d5-29a6c3f98229
md"""
x0:
$(@bind x0 Slider(-240:10:240; default=0, show_value=true))

w0:
$(@bind w0 Slider(2:1:50; default=25, show_value=true))

f:
$(@bind f Slider(0:10:2000; default=1000, show_value=true))

Lens\_z:
$(@bind Lens_z Slider(0:10:2000; default=1000, show_value=true))
"""

# ╔═╡ 4b7b50c6-1304-4c81-b2ac-92ecf7a3d4d1
begin
xlabel("z")
ylabel("x")
GR.heatmap(GaussianBeamLens.z,GaussianBeamLens.x,GaussianBeamLensPropagate(x0, w0, f, Lens_z)'; colormap = parse(Int,cm))
end

# ╔═╡ Cell order:
# ╠═a7ef92d7-41c1-4f7d-a0f9-b8037237af4f
# ╠═bca271a4-ee26-48ab-bef8-e6c28ae76351
# ╠═10e09fe7-cc23-4487-b1dc-b0a545bf2489
# ╠═4b7b50c6-1304-4c81-b2ac-92ecf7a3d4d1
# ╠═b1e13be1-d608-43a1-a7d5-29a6c3f98229


include("../01_perturbation_model/perturbation_model.jl")

function sqrt_epsilon_fun(ω, z, t, z1, z2, zmax, dN, z0, d_z, t_r, t_d, tau_r, tau_d)
    N_model = N_model_fun(z, t, z1, z2, zmax, dN, z0, d_z, t_r, t_d, tau_r, tau_d)
    ωₚ = 2pi*Ne_to_fpe(N_model)
    if ωₚ > ω
        return NaN
    end
    return sqrt(1 - ωₚ^2/ω^2)
end

function sqrt_eps_integral(z_step, ω, t, z1, z2, zmax, dN, z0, d_z, t_r, t_d, tau_r, tau_d)
    z = z1
    S = sqrt_epsilon_fun(ω, z, t, z1, z2, zmax, dN, z0, d_z, t_r, t_d, tau_r, tau_d)/2
    while true
        z = z + z_step
        s = sqrt_epsilon_fun(ω, z, t, z1, z2, zmax, dN, z0, d_z, t_r, t_d, tau_r, tau_d)
        if isnan(s)
            S = S - sqrt_epsilon_fun(ω, z-z_step, t, z1, z2, zmax, dN, z0, d_z, t_r, t_d, tau_r, tau_d)/2
            break
        end
        S = S + s
    end
    return S*z_step
end

# sqrt_eps_integral(1000., 45e5, 5., z1, z2, zmax, dN, z0, d_z, t_r, t_d, tau_r, tau_d)

function phase_variations(ω, t1, t2, z_step, z1, z2, zmax, dN, z0, d_z, t_r, t_d, tau_r, tau_d)
    c = 299792458.
    phase2 = 2ω/c * sqrt_eps_integral(z_step, ω, t2, z1, z2, zmax, dN, z0, d_z, t_r, t_d, tau_r, tau_d)
    phase1 = 2ω/c * sqrt_eps_integral(z_step, ω, t1, z1, z2, zmax, dN, z0, d_z, t_r, t_d, tau_r, tau_d)
    return phase2 - phase1
end

# f_axe = 4e6:1000.:5e6
# ω_axe = 2pi*f_axe
# t_axe = 0.0:0.1:5.
# z_step = 10.
# DPhi = zeros(length(t_axe),length(f_axe))

# for i=2:length(t_axe)
#     println(i)
#     for j=1:length(f_axe)
#         # println(i, " ",j)
#         DPhi[i,j] = phase_variations(ω_axe[j], t_axe[i-1], t_axe[i], z_step, z1, z2, zmax, dN, z0, d_z, t_r, t_d, tau_r, tau_d)
#     end
# end

# using Plots

# i=12; plot(f_axe/1e6, DPhi[i,:])
# title(string("t=",t_axe[i]))
# ylims((0,15))


using LowRankModels
import StatsBase: sample, Weights

# tests OvALoss

srand(1);
m,n,k = 200,50,2;
kfit = k+1
K = 4; # number of categories
d = n*K;
# matrix to encode
X_real, Y_real = randn(m,k), randn(k,d);
XY = X_real*Y_real;
# subtract the mean so we can compare the truth with the fit;
# the loss function is invariant under shifts
losses = fill(OvALoss(K),n)
yidxs = get_yidxs(losses)
for i=1:m
	for j=1:n
		mef = mean(XY[i,yidxs[j]])
		XY[i,yidxs[j]] = XY[i,yidxs[j]] - mef
	end
end

A = zeros(Int, (m, n))
for i=1:m
	for j=1:n
		wv = Weights(Float64[exp(-XY[i, K*(j-1) + l]) for l in 1:K])
		l = sample(wv)
		A[i,j] = l
	end
end

# and the model
losses = fill(OvALoss(K),n)
rx, ry = QuadReg(), QuadReg();
glrm = GLRM(A,losses,rx,ry,kfit, scale=false, offset=false, X=randn(kfit,m), Y=randn(kfit,d));

# fit w/o initialization
@time X,Y,ch = fit!(glrm);
XYh = X'*Y;
@show ch.objective
println("After fitting, parameters differ from true parameters by $(vecnorm(XY - XYh)/sqrt(prod(size(XY)))) in RMSE")
A_imputed = impute(glrm)
println("After fitting, $(sum(A_imputed .!= A) / prod(size(A))*100)\% of imputed entries are wrong")
println("(Picking randomly, $((K-1)/K*100)\% of entries would be wrong.)\n")

# initialize
init_svd!(glrm)
XYh = glrm.X' * glrm.Y
println("After initialization with the svd, parameters differ from true parameters by $(vecnorm(XY - XYh)/sqrt(prod(size(XY)))) in RMSE")
A_imputed = impute(glrm)
println("After initialization with the svd, $(sum(A_imputed .!= A) / prod(size(A))*100)\% of imputed entries are wrong")
println("(Picking randomly, $((K-1)/K*100)\% of entries would be wrong.)\n")

# fit w/ initialization
@time X,Y,ch = fit!(glrm);
XYh = X'*Y;
@show ch.objective
println("After fitting, parameters differ from true parameters by $(vecnorm(XY - XYh)/sqrt(prod(size(XY)))) in RMSE")
A_imputed = impute(glrm)
println("After fitting, $(sum(A_imputed .!= A) / prod(size(A))*100)\% of imputed entries are wrong")
println("(Picking randomly, $((K-1)/K*100)\% of entries would be wrong.)\n")


module SimpleGF2
using LinearAlgebra

import Base: show, +, -, *, /, xor, &, |, abs
import Base: convert, promote_rule, isless, real, rand

export GF2

struct GF2 <: Number
  val::UInt8
  function GF2(x::T) where T<:Integer
    new(x&1)
  end
end

convert(::Type{GF2}, x::T) where T<:Integer = GF2(x)
convert(::Type{T}, x::GF2) where T<:Integer = T(x.val)

+(x::GF2,y::GF2) = GF2(x.val+y.val)    # addition
-(x::GF2,y::GF2) = GF2(x.val+y.val)    # subtraction
-(x::GF2) = x                          # unary minus
*(x::GF2,y::GF2) = GF2(x.val & y.val)  # multiplication
function /(x::GF2,y::GF2)              # division
  if y.val == 0x00
    error("Division by zero in GF(2)")
  end
  return x
end

# Bitwise operators
xor(x::GF2,y::GF2) = GF2(xor(x.val,y.val))   # XOR
(&)(x::GF2,y::GF2) = GF2(x.val&y.val)   # AND
(|)(x::GF2,y::GF2) = GF2(x.val|y.val)   # OR

# isless defined giving all relations
isless(x::GF2,y::GF2) = isless(x.val,y.val)
isless(x::GF2,y::T) where T<:Real = isless(x.val,y)
isless(x::T,y::GF2) where T<:Real = isless(x,y.val)

abs(x::GF2) = x
real(x::GF2) = x.val

# random values and matrices
rand(::Type{GF2}) = GF2(rand(Int))
#rand(::Type{GF2},dims::Integer...) = map(GF2,rand(Int,dims...))
rand(::Type{GF2},dims::Integer...) = GF2.(rand(Int,dims...))

promote_rule(::Type{GF2}, ::Type{Int} ) = GF2

function show(io::IO, x::GF2)
  if x.val == 0x00
    print(io,"GF2(0)")
  else
    print(io,"GF2(1)")
  end
end

include("solving.jl")
# include("GF2inv.jl") -- now moved to "solving"

end # end of module


#!/usr/bin/env julia

import GR

srand(0)
xd = -2 + 4 * rand(100)
yd = -2 + 4 * rand(100)
zd = [xd[i] * exp(-xd[i]^2 - yd[i]^2) for i = 1:100]

GR.setviewport(0.1, 0.95, 0.1, 0.95)
GR.setwindow(-2, 2, -2, 2)
GR.setspace(-0.5, 0.5, 0, 90)
GR.setmarkersize(1)
GR.setmarkertype(GR.MARKERTYPE_SOLID_CIRCLE)
GR.setcharheight(0.024)
GR.settextalign(2, 0)
GR.settextfontprec(3, 0)

x, y, z = GR.gridit(xd, yd, zd, 200, 200)
h = linspace(-0.5, 0.5, 20)
GR.surface(x, y, z, 5)
GR.contour(x, y, h, z, 0)
GR.polymarker(xd, yd)
GR.axes(0.25, 0.25, -2, -2, 2, 2, 0.01)

GR.updatews()


struct DeepEnsemble
        models::Array
        DeepEnsemble(generator, N::Int) = new([generator() for _=1:N])
end

Flux.@functor DeepEnsemble

Flux.trainable(m::DeepEnsemble) = (Flux.trainable(model) for model in m.models)

# Get the mean and variance estimate from each network individually
function individual_forward(m::DeepEnsemble, x)
        os = [m(x) for m in m.models]
        ndim = Int(size(os[1],1) / 2)
        μs = [o[1:ndim, :] for o in os]
        σ²s = [softplus.(o[ndim+1:end, :]) .+ 1f-3 for o in os]
        μs, σ²s
end

# Compute the prediction and uncertainty of the ensemble
function (m::DeepEnsemble)(x)
        μs, σ²s = individual_forward(m, x)
        t1 = [σ² .+ μ .^2 for (μ, σ²) in zip(μs, σ²s)]
        mean(μs), mean(t1) .- mean(μs).^2
end

# The equation for the gaussian logpdf
de_gaussian_logpdf(μ, σ², y) = -log.(σ²) ./ 2f0 .- (y .- μ).^2 ./ (2f0 .*  σ²)

# Logpdf of some some x and y pair according to the full ensemble
Distributions.logpdf(m::DeepEnsemble, x, y) = de_gaussian_logpdf(m(x)..., y)

# Gets the mean negative log pdf for each network
function training_loss(m::DeepEnsemble, x, y, weights = ones(Float32, size(y)...))
        μs, σ²s = individual_forward(m, x)
        mean([-mean(weights .* de_gaussian_logpdf(μ, σ², y)) for (μ, σ²) in zip(μs, σ²s)])
end

struct DeepClassificationEnsemble
        models::Array
        DeepClassificationEnsemble(generator, N::Int) = new([generator() for _=1:N])
end

Flux.@functor DeepClassificationEnsemble

Flux.trainable(m::DeepClassificationEnsemble) = (Flux.trainable(model) for model in m.models)

# Get the mean and variance estimate from each network individually
function individual_forward(m::DeepClassificationEnsemble, x)
        ps = [softmax(m(x)) for m in m.models]
end

# Compute the prediction and uncertainty of the ensemble
function (m::DeepClassificationEnsemble)(x)
        ps = individual_forward(m, x)
        mean(ps)
end

# Logpdf of some some x and y pair according to the full ensemble
function Distributions.logpdf(m::DeepClassificationEnsemble, x, y)
        log.(sum(m(x) .* y, dims=1) .+ 1f-10)
end

# Gets the mean negative log pdf for each network
function training_loss(m::DeepClassificationEnsemble, x, y, weights = ones(Float32, size(y)...))
        ps = individual_forward(m, x)
        mean([Flux.Losses.crossentropy(p, y) for p in ps])
end


using LinearMaps
using Magnus

using PyPlot

function particle_in_box(M = 20,L = 1)
    dx = L/(M+1)
    x = linspace(dx,L-dx,M)
    T = SymTridiagonal(2ones(M), -ones(M-1))/dx^2
    ei = eigfact(T)
    T,ei,x
end

function state(ei, c::Vector, t::Real = 0)
    c /= sqrt(sum(abs2,c))
    sel = 1:length(c)
    c = c.*exp.(-im*ei[:values][sel]*t)
    psi = ei[:vectors][:,sel]*c
    E = dot(ei[:values][sel], abs2.(c))
    psi,E
end

function test_propagation(name, psi0, tmax, period, N, propagator)
    @time psi = integrate(psi0,
                          tmax*period, N,
                          propagator;
                          save_intermediate = true,
                          verbose = true)[:V]
    t = linspace(0,tmax,N+1)
    n = Vector{Float64}(length(t))
    p = Vector{Complex128}(length(t))
    for i = eachindex(t)
        n[i] = norm(psi[:,i])
        p[i] = dot(psi0,psi[:,i])
        psi[:,i] -= p[i]*psi0
    end

    figure(name,figsize=(8,10))
    clf()
    subplot(311)
    pcolormesh(t, x, abs2.(psi), rasterized=true, cmap = plt[:cm][:get_cmap]("viridis"))
    margins(0,0)
    gca()[:set_xticklabels]([])
    title(L"|\Psi-\Psi_0\langle\Psi_0|\Psi\rangle|^2")
    ylabel(L"x")
    subplot(312)
    plot(t,field(t))
    margins(0,0.1)
    gca()[:set_xticklabels]([])
    ylabel(L"F(t)")
    subplot(313)
    plot(t,n, "-")
    margins(0,0.1)
    xlabel(L"t/T")
    ylabel(L"|\Psi|^2")
    a2 = gca()[:twinx]()
    a2[:plot](t,abs.(p), "--")
    margins(0,0.1)
    ylabel(L"|\langle\Psi_0|\Psi\rangle|")
    tight_layout()
end

M = 50
T,ei,x = particle_in_box(M)
psi0,E0 = state(ei, [1])

H₀ = LinearMap(sparse(T))
F₀ = 1e4
period = 1.0/abs(ei[:values][1]-ei[:values][2])
ω = 2π/4
mx = Diagonal(collect(x-mean(x)))
D = LinearMap(sparse(mx))
tmax = 50
field = t -> F₀*exp.(-(t-tmax/2).^2/2dt^2).*sin.(ω*t)
F = t -> field(t*period)
Hᵢ = t -> F(t)*D
H = t -> H₀ + Hᵢ(t)

midpoint = MidpointPropagator(H, -im, LanczosExponentiator(30,psi0))
cfet = CFET4BfCPropagator(H₀, F, D, -im, LanczosExponentiator(30,psi0))

ndt = 300
dt = 6

test_propagation("midpoint", psi0, tmax, period, tmax*300, midpoint)
test_propagation("CFET", psi0, tmax, period, tmax*30, cfet)


using AbstractIterativeSolvers, LinearAlgebra
using Test

@testset "AbstractIterativeSolvers.jl: GMRES" begin
    n = 20
    A = randn(n,n)/(2*sqrt(n)) + I
    b = randn(n,1)
    out = GMRES(x -> A*x, b, (x,y) -> (y'*x)[1], 1e-10, n+1, x -> x)
    sol = sum([out[2][i]*out[1][i] for i=1:length(out[2]) ])
    @test norm(A*sol-b) < 1e-10
end


include("euler/euler.jl")
using .Calculus: sum_arithmetic_series

function compute(a::Int, b::Int, n::Int)::Int
    n -= 1
    ab = a * b
    return sum_arithmetic_series(a, a, n ÷ a) + sum_arithmetic_series(b, b, n ÷ b) - sum_arithmetic_series(ab, ab, n ÷ ab)
end


SKD = ScikitLearn.Skcore.pyimport("sklearn.datasets")

rmse(y, ŷ) = sqrt(sum(abs2, ŷ .- y) / length(y))

function simple_regression(n=100, p=3; sigma=0.1, σ=sigma, seed=6168901234)
    r = StableRNG(seed)
    X = hcat(randn(r, n, p), ones(n))
    θ = randn(r, p + 1)
    y = X * θ + σ * randn(r, n)
    ŷ = X * (X \ y)
    return X, y, rmse(y, ŷ)
end

function test_regression(m, X, y, ls)
    f, _, _ = MB.fit(m, 1, X, y)
    ŷ = MB.predict(m, f, X)
    return rmse(y, ŷ) / ls, f
end

function simple_regression_2(n=100, p=3;
                             sigma=0.1, σ=sigma, seed=6168901234)
    r  = StableRNG(seed)
    X  = hcat(randn(r, n, p), ones(n))
    θ  = randn(r, p + 1)
    y1 = X * θ + σ * randn(r, n)
    y2 = X * θ + σ * randn(r, n)
    ŷ  = X * (vcat(X, X) \ vcat(y1, y2))
    return X, hcat(y1, y2), (rmse(y1, ŷ)+rmse(y2, ŷ))/2
end

function test_regression_2(m, X, y, ls)
    f, _, _ = MB.fit(m, 1, X, y)
    ŷ = MB.predict(m, f, X)
    return sqrt(sum(abs2, vec(MB.matrix(y) .- MB.matrix(ŷ)))) / ls, f
end

function simple_binaryclf(n=100, p=3; sigma=0.1, seed=616866614)
    X, y = SKD.make_blobs(n_samples=n, n_features=p, centers=2, random_state=seed)
    return MB.table(X), MB.categorical(y)
end

function test_clf(m, X, y)
    f, _, _ = MB.fit(m, 1, X, y)
    if typeof(m) <: MB.Deterministic
        ŷ = MB.predict(m, f, X)
    else
        ŷ = MB.predict_mode(m, f, X)
    end
    return MB.accuracy(ŷ, y), f
end


export OptimizationProblem, primalDualSolve
######
# Simple Primal-Dual IP for the problem
#
# min_x f(x)
#   Cx -h <= 0.0

immutable OptimizationProblem
  d :: Int64
  m :: Int64
  f_g_h :: Function
  #R^d |----> R, returns a tuple with function value, gradient, hessian
  C :: Matrix{Float64} #C \in R^(d,m)
  h :: Vector{Float64}
end

# BV page 610
# The newton system is
# [\nabla^2 f_0, C' ]
# [-diag(\lambda) *C, -diag(C*x -h ) ]

function primalDualSolve(p :: OptimizationProblem, x :: Vector{Float64}, lambda :: Vector{Float64}, mu :: Float64, e_feas :: Float64, e :: Float64, max_iters :: Int64, alpha :: Float64, beta :: Float64)
  assert(mu > 1)
  assert(0.0 < beta < 1.0)
  assert(p.d == length(x))
  assert(p.m == length(lambda))
  assert(all(p.C*x .<= p.h))
  assert(all(lambda .> 0.0))
  assert(e > 0.0 && e_feas > 0.0)
  eta_hat(x,lambda) = -dot(lambda,p.C*x-p.h)

  function compute_r_t(f,g,x,lambda,t)
    [g + p.C'*lambda;
    -lambda.*(p.C*x - p.h) - (1.0/t)*ones(p.m)]
  end

  iter = 1
  for iter = 1:max_iters
    # println("ITER :", iter, ". Printing values of x: ", x)
    e_hat = eta_hat(x,lambda)
    t = (p.m*mu)/e_hat
    (f,g,h) = p.f_g_h(x)
    assert(length(g) == p.d)
    assert(size(h) == (p.d,p.d))

    #check termination conditions... Here?
    r_old = norm(compute_r_t(f,g,x,lambda,t))
    if (iter == max_iters)
      println(norm(g + p.C'*lambda))
      println(e_hat)
      println()
    end
    if (norm(g + p.C'*lambda) < e_feas && e_hat < e)
      break;
    end
    #form newton system...
    newton_system = [h p.C';
      -diagm(lambda)*p.C -diagm(p.C*x - p.h)]
    #target
    r_dual = g + p.C'*lambda
    r_cent = -lambda.*(p.C*x - p.h) - (1.0/t)*ones(p.m)
    #solve
    search_direction = -(newton_system\[r_dual; r_cent])
    sd_x = search_direction[1:p.d]
    # println(" Printing values of sd_x:  ", sd_x)
    sd_lambda = search_direction[p.d+1:end]
    #linesearch
    neg_inds = find(sd_lambda .< 0.0) #bullshit
    s_max = 1.0
    if length(neg_inds) > 0
      s_max = min(1.0, minimum(-lambda[neg_inds]./sd_lambda[neg_inds]))
    end
    #hardcoded in BV
    s = 0.995*s_max
    #backtracking for feasibility
    loopingIndex2 = 1
    while(!all(p.C*(x + s*sd_x) .<= p.h ))
      s = beta*s
      if loopingIndex2 ==1000
	  println("########################################################")
	  println("########################################################")
	  println("### WARINING ENTERING INFINITE LOOP, EXIT WITH ERROR ###")
    println("########################################################")
	  println("########################################################")
	  println("########################################################")
	  error()
      end
      loopingIndex2 = loopingIndex2 + 1
    end

    #backtracking for some other thing
    loopingIndex = 1
    while(true)
      if mod(loopingIndex,100) == 0
        println("### loopin' ",loopingIndex," !########")
      end
      loopingIndex = loopingIndex + 1
      x_plus = x + s*sd_x
      lambda_plus = lambda + s*sd_lambda
      f_t,g_t,h_t = p.f_g_h(x_plus) #only need f_g
      #compute norm of r_t
      r_t = compute_r_t(f_t,g_t,x_plus,lambda_plus,t)
      if (vecnorm(r_t) <= (1-alpha*s)*r_old)
        break;
      end
      s = beta*s
    end
    x = x + s*sd_x
    lambda = lambda + s*sd_lambda
  end
  if iter == max_iters
    warn("Hit max iters in interior point method!")
  end
  return x,lambda
end


using LinearAlgebra, SparseArrays

# State ∈ [0, up, down, updown] = [00, 10, 01, 11]
# 2x2 lattice -> 4 (exponent) -> 4^4 states = 256

function state_from_integer(int, sites=4, substates_per_site=2)
    int > (2substates_per_site)^sites-1 && return
    out = BitArray(undef, (substates_per_site, sites))
    out.chunks[1] = int
    out
end

function create(state, site, substate)
    # No state case
    state == 0 && return 0, 0.0
    # create(|1⟩) -> no state
    state[substate, site] && return 0, 0.0

    # fermionic sign
    lindex = 2(site-1) + substate
    _sign = iseven(sum(state[1:lindex-1])) ? +1.0 : -1.0

    # create(|0⟩) -> |1⟩
    s = copy(state)
    s[substate, site] = true
    return _sign, s
end

function annihilate(state, site, substate)
    # No state
    state == 0 && return 0, 0.0
    # annihilate(|0⟩) -> no state
    !state[substate, site] && return 0, 0.0

    # fermionic sign
    lindex = 2(site-1) + substate
    _sign = iseven(sum(state[1:lindex-1])) ? +1.0 : -1.0

    # annihilate(|1⟩) -> |0⟩
    s = copy(state)
    s[substate, site] = false
    return _sign, s
end

function HamiltonMatrix(model::HubbardModelAttractive)
    lattice = model.l
    t = model.t
    U = model.U
    mu = model.mu

    H = zeros(Float64, 4^lattice.sites, 4^lattice.sites)

    # -t ∑_ijσ c_iσ^† c_jσ
    # +U ∑_i (n_i↑ - 1/2)(n_i↓ - 1/2)
    # -μ ∑_i n_i
    for i in 1:4^lattice.sites
        lstate = state_from_integer(i-1, lattice.sites)
        for j in 1:4^lattice.sites
            rstate = state_from_integer(j-1, lattice.sites)

            E = 0
            # hopping (hermitian conjugate implied/included by lattice generation)
            for substate in [1, 2]
                for source in 1:lattice.sites
                    for target in lattice.neighs[:, source]
                        _sign1, state = annihilate(rstate, source, substate)
                        _sign2, state = create(state, target, substate)
                        if state != 0 && lstate == state
                            E -= _sign1 * _sign2 * t
                        end
                    end
                end
            end

            # # U, μ terms
            for p in 1:4
                up_occ = rstate[1, p]
                down_occ = rstate[2, p]
                if lstate == rstate
                    E += U * (up_occ - 0.5) * (down_occ - 0.5)
                    E -= mu * (up_occ + down_occ)
                end
            end

            H[i, j] = E
        end
    end
    H
end

# Greens function

function Greens(site1, site2, substate1, substate2)
    s -> begin
        _sign1, state = create(s, site1, substate1)
        state == 0 && return typeof(s)[], Float64[]
        _sign2, state = annihilate(state, site2, substate2)
        state == 0 && return typeof(s)[], Float64[]
        return [state], [_sign1 * _sign2]
    end
end

# According to p188 QMCM this is equivalent
# can be used to check if Greens/ED is correct
function Greens_permuted(site1, site2, substate1, substate2)
    s -> begin
        _sign1, state = annihilate(s, site2, substate2)
        _sign2, state = create(state, site1, substate1)
        delta = ((site1 == site2) && (substate1 == substate2)) ? 1.0 : 0.0

        if state == 0 && delta == 0.0
            # off-diagonal
            return typeof(s)[], Float64[]
        elseif state == 0
            # only delta function triggers (|0⟩ state)
            return [s], [delta]
        else
            # both trigger (|1⟩ state)
            return [state], [delta - _sign1 * _sign2]
        end
    end
end

function expectation_value(
        observable::Function,
        H;
        T=1.0, beta = 1.0 / T,
        N_sites = 4,
        N_substates = 2
    )

    vals, vecs = eigen(H)
    Z = 0.0
    O = 0.0
    for i in eachindex(vals)
        # exp(βEᵢ)
        temp = exp(-beta * vals[i])
        Z += temp

        # ⟨ψᵢ|Ô|ψᵢ⟩
        right_coefficients = zeros(eltype(vecs), size(vecs, 1))
        for j in 1:size(vecs, 1)
            state = state_from_integer(j-1, N_sites, N_substates)
            states, values = observable(state)
            for (s, v) in zip(states, values)
                # Assuming no (s, v) pair if state destroyed
                k = s.chunks[1]+1
                right_coefficients[k] += v * vecs[j, i]
            end
        end
        O += temp * dot(vecs[:, i], right_coefficients)
    end
    O / Z
end

function calculate_Greens_matrix(H, lattice; beta=1.0, N_substates=2)
    G = Matrix{Float64}(
        undef,
        lattice.sites*N_substates,
        lattice.sites*N_substates
    )
    for substate1 in 1:N_substates, substate2 in 1:N_substates
        for site1 in 1:lattice.sites, site2 in 1:lattice.sites
            G[
                lattice.sites * (substate1-1) + site1,
                lattice.sites * (substate2-1) + site2
            ] = expectation_value(
                Greens(site1, site2, substate1, substate2),
                H,
                beta = beta,
                N_sites = lattice.sites,
                N_substates=N_substates
            )
        end
    end
    G
end

# utility

function state2string(state)
    sub, L = size(state)
    spin = ["↑", "↓"]
    chunks = String[]
    for p in 1:L
        str = ""
        for s in 1:2
            str *= state[s, p] ? spin[s] : "⋅"
        end
        push!(chunks, str)
    end
    join(chunks, " ")
end


# ------------------------------------------------------------------
# Licensed under the ISC License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    imfilter_fft(img, kern)

Perform filtering on `img` with kernel `kern` using the FFT algorithm.
"""
imfilter_fft(img::AbstractArray{T,N}, kern::AbstractArray{K,N}) where {T<:Real,K<:Real,N} =
  imfilter(img, centered(kern), Inner(), Algorithm.FFT())

"""
    convdist(img, patch, weights)

Compute distance between all patches of `img` and a single `patch`
using `weights` for each pixel in the `patch`.
"""
function convdist(img::AbstractArray, patch::AbstractArray, weights::AbstractArray)
  wpatch = weights.*patch

  A² = imfilter_fft(img.^2, weights)
  AB = imfilter_fft(img, wpatch)
  B² = sum(wpatch .* patch)

  abs.(A² .- 2AB .+ B²)
end

"""
    selectpatch(img, patchsize, center)

Helper function to extract a patch of size
`patchsize` from `img` centered at Cartesian
index `center`.
"""
function selectpatch(img, patchsize, center)
  start  = CartesianIndex(@. center.I - (patchsize-1)÷2)
  finish = CartesianIndex(@. center.I + (patchsize  )÷2)
  view(img, start:finish)
end


"""
    StatsModelComparisons

This module implements several model comparison methods such as PSIS, WAIC and DIC.

Included functions
------------------
psisloo
    Pareto smoothed importance sampling leave-one-out log predictive densities.

psislw
    Pareto smoothed importance sampling.

gpdfitnew
    Estimate the paramaters for the Generalized Pareto Distribution (GPD).

gpinv
    Inverse Generalised Pareto distribution function.

waic
    ...

dic
    ...

deviance
    ...

aic
    Akaike Information Criterion

bic
    Bayesian Information Criterion

References
----------
Aki Vehtari, Andrew Gelman and Jonah Gabry (2015). Efficient implementation
of leave-one-out cross-validation and WAIC for evaluating fitted Bayesian
models. arXiv preprint arXiv:1507.04544.

Aki Vehtari and Andrew Gelman (2015). Pareto smoothed importance sampling.
arXiv preprint arXiv:1507.02646.
"""
module StatsModelComparisons

using RecipesBase
using StatsFuns
using Statistics

include("psisloo.jl")
include("psislw.jl")
include("gpdfitnew.jl")
include("gpinv.jl")
include("waic.jl")
include("pk_utilities.jl")
include("dic.jl")
include("legacy.jl")

export deviance, dic, psisloo, gpdfitnew, gpinv
export var2, waic, aic, bic
export psislw
export pk_qualify
end


using Test
using Nord, ColorTypes, FixedPointNumbers

@test Nord.nord0 == RGB{N0f8}(46/255, 52/255, 64/255)


# output = wx
immutable Linear{T,P<:Params} <: Learnable
    nin::Int
    nout::Int
    input::SumNode{T,1}
    output::OutputNode{T,1}
    params::P
end

function Linear{T}(::Type{T}, nin::Int, nout::Int,
                    θ::AbstractVector = zeros(T, nout*nin),
                    ∇::AbstractVector = zeros(T, nout*nin))
    input = InputNode(T, nin)
    output = OutputNode(T, nout)
    params = Params(θ, ∇, ((nout,nin),))
    w = params.views[1]
    initialize_weights!(w)
    Linear(nin, nout, input, output, params)
end
Linear(nin::Int, nout::Int, args...) = Linear(Float64, nin, nout, args...)

function Base.show(io::IO, t::Linear)
    print(io, "Linear{$(t.nin)-->$(t.nout)}")
end

function transform!(t::Linear)
    w = t.params.views[1]
    x = t.input.val
    y = t.output.val
    for o=1:t.nout
        y[o] = sum(w[o,i] * x[i] for i=1:t.nin)
    end
    y
end

function grad!{T}(t::Linear{T})
    w = t.params.views[1]
    ∇w = t.params.∇_views[1]
    x = t.input.val
    ∇x = t.input.∇
    ∇y = t.output.∇

    # ∇x, ∇w
    for i=1:t.nin
        ∇xᵢ = zero(T)
        for o=1:t.nout
            ∇xᵢ += w[o,i] * ∇y[o]
            ∇w[o,i] = x[i] * ∇y[o]
        end
        ∇x[i] = ∇xᵢ
    end
    return
end