include("../../MaximinOPF/src/MaximinOPF.jl")
using PowerModels
using JuMP
using SCS
using Ipopt
using ProxSDP
PowerModels.silence()

function SolveFP(pm_data,pm_form,pm_optimizer; x_vals=Dict{Int64,Float64}() )
  pm = MaximinOPF.PF_FeasModel(pm_data, pm_form, x_vals)
  JuMP.set_optimizer(pm.model,pm_optimizer)
  JuMP.optimize!(pm.model)
  status=JuMP.termination_status(pm.model)
  if status != OPTIMAL
    println("FLAGGING: Solve status=",status)
  end
  return pm
end

testcase = Dict(
	"file" => "data/case9.m",
	"PMOption" => SparseSDPWRMPowerModel,
 	"name" => "case9K3ACR",
 	"attack_budget" => 0,
 	"inactive_indices" => [1,4,7],
 	"protected_indices" => []
	)

pm_data = PowerModels.parse_file(testcase["file"])
pm_data["attacker_budget"] = testcase["attack_budget"] ###Adding another key and entry
pm_data["inactive_branches"] = testcase["inactive_indices"] ###Adding another key and entry
pm_data["protected_branches"] = testcase["protected_indices"] ###Adding another key and entry

sdp_solver=with_optimizer(ProxSDP.Optimizer, log_verbose=false, tol_primal=1e-6, tol_dual=1e-6 )
#sdp_solver=with_optimizer(Mosek.Optimizer,MSK_IPAR_LOG=0)
#conic_solver=with_optimizer(Mosek.Optimizer,MSK_IPAR_LOG=0)
conic_solver=with_optimizer(SCS.Optimizer,verbose=0)
ip_solver=with_optimizer(Ipopt.Optimizer,print_level=0)

# nonconvex AC forms
nonconvex_ac=[ACPPowerModel, ACRPowerModel, ACTPowerModel]
for pm_form in nonconvex_ac
    println("Formulating and solving the form ",pm_form)
    pm = MaximinOPF.PF_FeasModel(pm_data, pm_form)
    JuMP.set_optimizer(pm.model,ip_solver)
    JuMP.optimize!(pm.model)
    println("\tOptimal value using powerform ", pm_form, " is: ",JuMP.objective_value(pm.model), "with status ",JuMP.termination_status(pm.model))
end

# linear approximations
linear_approx=[DCPPowerModel, DCMPPowerModel, NFAPowerModel]
for pm_form in linear_approx
    println("Formulating and solving the form ",pm_form)
    pm = MaximinOPF.PF_FeasModel(pm_data, pm_form)
    JuMP.set_optimizer(pm.model,conic_solver)
    JuMP.optimize!(pm.model)
    println("\tOptimal value using powerform ", pm_form, " is: ",JuMP.objective_value(pm.model), "with status ",JuMP.termination_status(pm.model))
end

# quadratic approximations
quadratic_approx=[DCPLLPowerModel, LPACCPowerModel]
for pm_form in quadratic_approx
    println("Formulating and solving the form ",pm_form)
    pm = MaximinOPF.PF_FeasModel(pm_data, pm_form)
    JuMP.set_optimizer(pm.model,ip_solver)
    JuMP.optimize!(pm.model)
    println("\tOptimal value using powerform ", pm_form, " is: ",JuMP.objective_value(pm.model), "with status ",JuMP.termination_status(pm.model))
end
# quadratic relaxations
quadratic_relax=[SOCWRPowerModel, SOCBFPowerModel, QCRMPowerModel, QCLSPowerModel]
for pm_form in quadratic_relax
    println("Formulating and solving the form ",pm_form)
    pm = MaximinOPF.PF_FeasModel(pm_data, pm_form)
    JuMP.set_optimizer(pm.model,ip_solver)
    JuMP.optimize!(pm.model)
    println("\tOptimal value using powerform ", pm_form, " is: ",JuMP.objective_value(pm.model), "with status ",JuMP.termination_status(pm.model))
end
quad_conic_relax=[SOCWRConicPowerModel, SOCBFConicPowerModel]
for pm_form in quad_conic_relax
    println("Formulating and solving the form ",pm_form)
    pm = MaximinOPF.PF_FeasModel(pm_data, pm_form)
    JuMP.set_optimizer(pm.model,conic_solver)
    JuMP.optimize!(pm.model)
    println("\tOptimal value using powerform ", pm_form, " is: ",JuMP.objective_value(pm.model), "with status ",JuMP.termination_status(pm.model))
end
# sdp relaxations
sdp_relax=[SDPWRMPowerModel, SparseSDPWRMPowerModel]
for pm_form in sdp_relax
    println("Formulating and solving the form ",pm_form)
    pm = MaximinOPF.PF_FeasModel(pm_data, pm_form)
    JuMP.set_optimizer(pm.model,sdp_solver)
    JuMP.optimize!(pm.model)
    println("\tOptimal value using powerform ", pm_form, " is: ",JuMP.objective_value(pm.model), "with status ",JuMP.termination_status(pm.model))
end


using ABCDistances
using Base.Test

include("abc.jl")


#=
Continuation of data, to e.g. fill missing values
=#
using SignalDecomposition
export sinusoidal_continuation

_numbertype(T) = _numbertype(eltype(T))
_numbertype(::Type{Union{T, Missing}}) where T = T
_numbertype(::Type{T}) where {T <: Real} = T

# TODO: perhaps it is worth considering a version that makes true time in seconds
# instead of months, which in the future would allow support of diurnal cycles
# TODO: YES! Make true time be seconds!

"""
    sinusoidal_continuation(T::ClimArray, fs = [1, 2]; Tmin = -Inf, Tmax = Inf)
Fill-in the missing values of spatiotemporal field `T`, by fitting sinusoidals
to the non-missing values, and then using the fitted sinusoidals for the missing values.

Frequencies are given per year (frequency 2 means 1/2 of a year).

`Tmin, Tmax` limits are used to clamp the result into this range (no clamping in the
default case).
"""
function sinusoidal_continuation(T, frequencies = [1.0, 2.0]; Tmin = -Inf, Tmax = Inf)
    E = _numbertype(T)
    lpv = Sinusoidal(E.(frequencies ./ DAYS_IN_ORBIT))
    fullT = copy(T)
    # TODO: this must be extended to a general "true time" function
    truetime = realtime_days(dims(T, Time).val, E)
    for i in otheridxs(T, Time)
        x = T[i...]
        any(ismissing, x) || continue # this timeseries needs no correction
        mi = findall(!ismissing, x)
        x0 = Array{Float32}(x[mi])
        t0 = truetime[mi]
        sea, res = decompose(t0, x0, lpv)
        c = cosines(truetime, lpv)
        imi = findall(ismissing, x)
        x[imi] .= c[imi]
        fullT[i...] .= x
    end
    if !isinf(Tmin) && !isinf(Tmin)
        fullT = clamp.(fullT, Tmin, Tmax)
    end
    return ClimArray(nomissing(fullT.data), T.dims; name = T.name, attrib = T.attrib)
end

function cosines(t, lpv::Sinusoidal)
    s = fill(lpv.A[1], length(t))
    for i in 2:length(lpv.A)
        @. s += lpv.A[i] * cos(2π * lpv.fs[i-1] * t + lpv.φ[i])
    end
    return s
end


"""Given a matrix of energies `E` (in Hartrees), where each column
corresponds to a different angular symmetry, a vector of matrices `V`
containing the basis vectors for each angular symmetry, generate a
field-free Hamiltonian on diagonal form and transform all operators
`O` to the field-free basis. Optionally, `cutoff` (in eV) limits the
spectrum. Also returned is the index of the lowest, i.e., groundstate
energy."""
function hamiltonian{T<:AbstractFloat,
                     M<:AbstractMatrix}(E::Vector{Vector{T}},
                                        V::Vector{M},
                                        O::Vector{Operator},
                                        cutoff::Real = 0,
                                        verbose::Bool = false)
    n = length(E)
    sel = [(1:length(E[j])) for j = 1:n]

    # Find range of elements with energies below the cutoff energy.
    if cutoff != 0
        cutoff /= 27.211
        for j = 1:n
            i = findlast(e -> e <= cutoff, E[j])
            verbose && i == sel[j][end] && warn("Largest energy $(E[j][end]) smaller than cutoff $(cutoff).")
            i == 0 && error("No energy smaller than cutoff in partial wave $j.")
            sel[j] = 1:i
        end
    end

    # Find block coordinates
    coords = Vector{UnitRange{Int64}}(n)
    coords[1] = sel[1]
    for j = 2:n
        coords[j] = sel[j] + coords[j-1][end]
    end

    # Reduce the space
    E = vcat([E[j][sel[j]] for j = 1:n]...)
    V = Matrix{T}[V[j][:,sel[j]] for j in 1:n]

    # Field-free Hamiltonian in its eigenbasis
    H = sparse(Diagonal(E))

    # Transform all operators O into the eigenbasis
    O = map(O) do o
        ot = spzeros(eltype(first(o.o).m), size(H)...)
        for oo in o.o
            o12 = sparse(V[oo.k1]'*oo.m*V[oo.k2])
            ot[coords[oo.k1],coords[oo.k2]] += o12
            ot[coords[oo.k2],coords[oo.k1]] += o12'
        end
        ot
    end

    # Find groundstate
    gst = indmin(E)

    H,O,gst
end

function hamiltonian{T<:AbstractFloat,
                     M<:AbstractMatrix}(E::Vector{Vector{T}},
                                        V::Vector{M},
                                        O::Operator,
                                        cutoff::Real = 0,
                                        verbose::Bool = false)
    H,O,gst = hamiltonian(E, V, [O], cutoff, verbose)
    H,O[1],gst
end

export hamiltonian


using SIIPExamples
pkgpath = dirname(dirname(pathof(SIIPExamples)))
include(
    joinpath(pkgpath, "test", "3_PowerSimulations_examples", "01_operations_problems.jl"),
);

using Ipopt
solver = optimizer_with_attributes(Ipopt.Optimizer)

ed_template = template_economic_dispatch(network = StandardPTDFModel)

ed_template.devices[:Hydro] = DeviceModel(HydroEnergyReservoir, HydroDispatchRunOfRiver)

PTDF_matrix = PTDF(sys)

problem = OperationsProblem(
    EconomicDispatchProblem,
    ed_template,
    sys,
    horizon = 4,
    optimizer = solver,
    balance_slack_variables = true,
    constraint_duals = [:CopperPlateBalance, :network_flow],
    PTDF = PTDF_matrix,
)

res = solve!(problem);

λ = convert(Array, res.dual_values[:CopperPlateBalance])
μ = convert(Array, res.dual_values[:network_flow])

buses = get_components(Bus, sys)
congestion_lmp = Dict()
for bus in buses
    congestion_lmp[get_name(bus)] = μ * PTDF_matrix[:, get_number(bus)]
end
congestion_lmp = DataFrame(congestion_lmp)

LMP = λ .- congestion_lmp

# This file was generated using Literate.jl, https://github.com/fredrikekre/Literate.jl


using Oceananigans.Architectures
using Oceananigans.Architectures: device
import Oceananigans.Architectures: architecture
using CUDA, CUDA.CUSPARSE
using KernelAbstractions: @kernel, @index

using LinearAlgebra, SparseArrays, IncompleteLU
using SparseArrays: nnz

import LinearAlgebra.ldiv!

"""
`JacobiPreconditioner`
    stores only the diagonal `d` of D⁻¹ where `D = diag(A)`
    is applied to `r` with a vector multiplication `d .* r`

`ILUFactorization`
    stores two sparse lower and upper trianguilar matrices `L` and `U` such that `LU ≈ A`
    is applied to `r` with `forward_substitution!(L, r)` followed by `backward_substitution(U, r)`
    constructed with `ilu(A, τ = drop_tolerance)`

`SparseInversePreconditioner`
    stores a sparse matrix `M` such that `M ≈ A⁻¹`
    is applied to `r` with a matrix multiplication `M * r`
    constructed with
    `simplified_inverse_preconditioner(A)`
        -> assumes that the sparsity of `M` is the same as the sparsity of `A`
    `sparse_approximate_preconditioner(A, ε = tolerance, nzrel = relative_maximum_number_of_elements)`
        -> starts constructing the sparse inverse of A from identity matrix until, either a tolerance (ε) is met or nnz(M) = nzrel * nnz(A)

The suggested preconditioners are

on the `CPU`
`ilu()` (superior to everything always and in every situation!)

on the `GPU`
`sparse_inverse_preconditioner()` (if `Δt` is constant)
`simplified_inverse_preconditioner()` or `JacobiPreconditioner` (if Δt is variable)

`ilu()` cannot be used on the GPU because preconditioning the solver with a direct LU (or Choleski) type
of preconditioner would require too much computation for the `ldiv!(P, r)` step completely hindering the performances

"""
validate_settings(T, arch, settings)                                 = settings
validate_settings(::Val{:Default}, arch, settings)                   = arch isa CPU ? (τ = 0.001,) : (ε = 0.1, nzrel = 1.0)
validate_settings(::Val{:SparseInverse}, arch, settings::Nothing)    = (ε = 0.1, nzrel = 1.0)
validate_settings(::Val{:ILUFactorization}, arch, settings::Nothing) = (τ = 0.001, )

validate_settings(::Val{:ILUFactorization}, arch, settings) = haskey(settings, :τ) ?
                                                                     settings :
                                                                     throw(ArgumentError("τ has to be specified for ILUFactorization"))
validate_settings(::Val{:SparseInverse}, arch, settings)    = haskey(settings, :ε) && haskey(settings, :nzrel) ?
                                                                     settings :
                                                                     throw(ArgumentError("both ε and nzrel have to be specified for SparseInverse"))

function build_preconditioner(::Val{:Default}, matrix, settings)
    default_method = architecture(matrix) isa CPU ? :ILUFactorization : :SparseInverse
    return build_preconditioner(Val(default_method), matrix, settings)
end

build_preconditioner(::Val{:None},              A, settings) = Identity()
build_preconditioner(::Val{:Jacobi},            A, settings) = JacobiPreconditioner(A)
build_preconditioner(::Val{:SparseInverse},     A, settings) = sparse_inverse_preconditioner(A, ε = settings.ε, nzrel = settings.nzrel)
build_preconditioner(::Val{:SimplifiedInverse}, A, settings) = simplified_inverse_preconditioner(A)
build_preconditioner(::Val{:ILUFactorization},  A, settings) = ilu(A, τ = settings.τ)

@inline architecture(::CuSparseMatrixCSC) = GPU()
@inline architecture(::SparseMatrixCSC)   = CPU()

struct JacobiPreconditioner{V}
    invdiag::V
end

function JacobiPreconditioner(A::AbstractMatrix)
    M    = size(A, 1)
    arch = architecture(A)
    dev  = device(arch)

    invdiag = arch_array(arch, zeros(eltype(A), M))
    col, row, val = unpack_constructors(arch, constructors(arch, A))

    loop! = _get_inv_diag!(dev, 256, M)
    event = loop!(invdiag, col, row, val; dependencies=Event(dev))
    wait(dev, event)

    return JacobiPreconditioner(invdiag)
end

@kernel function _multiply_in_place!(u, invdiag, v)
    i = @index(Global, Linear)
    @inbounds u[i] = invdiag[i] * v[i]
end

function  LinearAlgebra.ldiv!(u, precon::JacobiPreconditioner, v)
    invdiag = precon.invdiag
    arch = architecture(invdiag)
    dev  = device(arch)

    M = length(invdiag)
    loop! = _multiply_in_place!(dev, 256, M)
    event = loop!(u, invdiag, v; dependencies=Event(dev))
    wait(dev, event)
end

function  LinearAlgebra.ldiv!(precon::JacobiPreconditioner, v)
    ldiv!(v,precon,v)
end

abstract type AbstractInversePreconditioner{M} end

function  LinearAlgebra.ldiv!(u, precon::AbstractInversePreconditioner, v)
    mul!(u, matrix(precon), v)
end

function  LinearAlgebra.ldiv!(precon::AbstractInversePreconditioner, v)
    mul!(v, matrix(precon), v)
end

struct SparseInversePreconditioner{M} <: AbstractInversePreconditioner{M}
    Minv :: M
end

@inline matrix(p::SparseInversePreconditioner)  = p.Minv

function simplified_inverse_preconditioner(A::AbstractMatrix)

    arch                  = architecture(A)
    constr                = deepcopy(constructors(arch, A))
    colptr, rowval, nzval = copy_unpack_constructors(arch, constr)
    dev                   = device(arch)

    M = size(A, 1)

    diag = arch_array(arch, zeros(eltype(nzval), M))

    loop! = _get_diag!(dev, 256, M)
    event = loop!(diag, colptr, rowval, nzval; dependencies=Event(dev))
    wait(dev, event)

    loop! = _initialize_simplified_inverse_preconditioner!(dev, 256, M)
    event = loop!(nzval, colptr, rowval, diag; dependencies=Event(dev))
    wait(dev, event)

    constr_new = (colptr, rowval, nzval)

    Minv = arch_sparse_matrix(arch, constructors(arch, M, constr_new))

    return SparseInversePreconditioner(Minv)
end

@kernel function _initialize_simplified_inverse_preconditioner!(nzval, colptr, rowval, diag)
    col = @index(Global, Linear)

    for idx = colptr[col] : colptr[col+1] - 1
        if rowval[idx] == col
            nzval[idx] = diag[col]
        else
            nzval[idx] = - nzval[idx] * diag[rowval[idx]]
        end
    end
end

function sparse_inverse_preconditioner(A::AbstractMatrix; ε, nzrel)

   # let's choose an initial sparsity => diagonal
   A_cpu    = arch_sparse_matrix(CPU(), A)
   Minv_cpu = spai_preconditioner(A_cpu, ε = ε, nzrel = nzrel)

   Minv = arch_sparse_matrix(architecture(A), Minv_cpu)
   return SparseInversePreconditioner(Minv)
end


using Flux, HDF5, ProgressMeter
using MLDataUtils, BSON
using Statistics, Random
using Base.Iterators: partition
#regression neural net
function regressionNN(dimIn, dimOut, neuronas)
    Chain(
        Dense(dimIn, neuronas, Flux.relu),
        Dense(neuronas, neuronas, Flux.relu),
        Dense(neuronas, neuronas, Flux.relu),
        Dense(neuronas, neuronas, Flux.relu),
        Dense(neuronas, neuronas, Flux.relu),
        #Dense(neuronas, neuronas, Flux.relu),
        Dense(neuronas, neuronas, Flux.relu),
        Dense(neuronas, dimOut)) # |> gpu
end
function machine(xdata, ydata, xval, yval, seed;
        neurons=40, nepochs=100, bs=512, best_val = 5.0,
        last_improvement = 0,patience=5, gpu = false)
    #saving best model
    dict_models = Dict()
    dict_models["model_bm$(seed)"] = NaN
    #optimizer and loss function
    Random.seed!(seed)
    model = regressionNN(size(xdata)[1], size(ydata)[1], neurons)
    # if training on gpu's is an option
    model = gpu == true ? model |> gpu : model
    xdata = gpu == true ? xdata |> gpu : xdata
    ydata = gpu == true ? ydata |> gpu : ydata

    opt = ADAM()
    loss(x,y) = Flux.mse(model(x), y)
    #nepochs = 100
    lossVal = zeros(nepochs)
    #mini-batch
    #batch_size 512 is ok too, similar performance at test accuracy
    mb_idxs = partition(1:size(xdata)[2], bs)
    #xs, ys = shuffleobs((Float32.(xdatacoef), ycdata))
    p = Progress(nepochs, dt=1, barglyphs=BarGlyphs("[=> ]"), color = :yellow)
    for epoch_indx in 1:nepochs
        xs, ys = shuffleobs((Float32.(xdata), Float32.(ydata)))
        train_set = [(xs[:,p], ys[:,p]) for p in mb_idxs]
        Flux.train!(loss, params(model), train_set, opt)

        validation_loss = loss(xval, yval)
        lossVal[epoch_indx] = validation_loss
        # If this is the best val_loss we've seen so far, save the model out
        if validation_loss <= best_val
            #@info(" -> New best val_loss! Saving model out to model_bm$(seed).bson")
            last_improvement = epoch_indx
            dict_models["model_bm$(seed)"] = model # model_bestmodel
            best_val = validation_loss
        end
        if epoch_indx - last_improvement >= patience
            #@info(" -> Early-exiting iteration $(r) and epoch $(epoch_indx): no more patience")
            break
        end
        next!(p; showvalues = [(:epoch,epoch_indx), (:Model_seed, seed)])
        #ProgressMeter.ijulia_behavior(:append)
        ProgressMeter.ijulia_behavior(:clear)
    end
    model = dict_models["model_bm$(seed)"]
    #BSON.@save "models_bm$(seed).bson" dict_models lossVal
    model, lossVal
end

function ensembleMachines(xtrain, ytrain, xval, yval;
        nmodels = 10, nepoch = 100, neurons=40, bs = 512, best_val = 5.0,
        last_improvement = 0, patience=10, gpu = false)

    losses = zeros(nepoch, nmodels)
    machines = []
    dict_models = Dict()
    #pro = Progress(nmodels, dt=0.5, barglyphs=BarGlyphs("[=> ]"), color = :black)
    for seed in 1:nmodels
        modelo1, lossVal = machine(xtrain, ytrain, xval, yval,seed;
            neurons= neurons, nepochs= nepoch, bs=bs, best_val = best_val,
            last_improvement = last_improvement, patience=patience, gpu = gpu)

        losses[:,seed] = lossVal
        dict_models["model_bm$(seed)"] = modelo1
        #next!(pro; showvalues = [(:Model, seed)])
        #ProgressMeter.ijulia_behavior(:append)
        #ProgressMeter.ijulia_behavior(:clear)
    end
    #BSON.@save "models_bm$(nmodels).bson" dict_models losses
    dict_models, losses
end

function splitdata(xdata, ydata)
    #μnoi = mean(xdata)
    #stdnoi = std(xdata)
    #μref = mean(ydata_1)
    #stdref = std(ydata_1)
    #xdata = (xdata .- μnoi)./stdnoi
    #ydata = (ydata .- μref)./stdref

    xtrain = xdata[:,1:32000]
    ytrain = ydata[:,1:32000]
    xval = xdata[:,32001:36000]
    yval = ydata[:,32001:36000]
    xtest = xdata[:,36001:end]
    ytest = ydata[:,36001:end]
    xtrain, ytrain, xval, yval, xtest, ytest #, [μnoi, stdnoi, μref, stdref]
end


# Quantify the error of the 6th order approximation to transverse Mercator
# projections by the invertability of points

using Geodesy
using PyPlot

z = 31
h = true
t1 = UTMfromLLA(z,h,wgs84)
t2 = LLAfromUTM(z,h,wgs84)

lats = Float64[0:1:80;]
lons = Float64[0:1:80;]
errors = Float64[distance(LLA(lats[i],lons[j]+3.0,0.0), transform(t2, transform(t1, LLA(lats[i],lons[j]+3.0,0.0)))) for i=1:length(lats), j in 1:length(lons)]
errors = convert(Matrix{Float64}, errors) # Compehension is broken (code_warntype is correct for distance but not for the array...)

figure()
imshow(log10(errors))

error_lon = [maximum(errors[:,i]) for i in 1:length(lons)]

figure()
plot(lons, error_lon)


importall DynamicDiscreteModels
#need to `importall` rather than `using` because we will extend `coef!()` and `dim()`
#in real applications this would typically be handled in a front-end package ToyModel.jl

import ForwardDiff

type ToyModel <: DynamicDiscreteModel

	#DynamicDiscreteModel fields
	m::Array{Float64,4}			  	#the transition matrix given as m[x,y,x',y']
	mu::Array{Float64,2}  			#initial distribution (dx,dy)

	#DynamicDiscreteModel's technical fields
	mjac::Array{Float64,5}			#jacobian
	rho::Array{Float64,1}
	phi::Array{Float64,1}
	psi::Array{Float64,1}
	rhojac::Array{Float64,1}
	phijac::Array{Float64,2}
	psijac::Array{Float64,2}
end

dx=2
dy=3

toymodel()=ToyModel(
	Array(Float64,dx,dy,dx,dy), 			#m
	fill(1/6,2,3), 							#mu
	Array(Float64,dx,dy,dx,dy,2),			#mjac
	Array(Float64,1), 						#rho
	Array(Float64,dx),						#phi
	Array(Float64,dx),						#psi
	Array(Float64,2), 						#rhojac
	Array(Float64,dx,2),					#phijac
	Array(Float64,dx,2)						#psijac
	)

function coef!(model::ToyModel,theta::Tuple)
	p1,p2=theta[1],theta[2]
	a=[p1 1-p1;1-p1 p1]
	p3=(1-p2)/2
	b=[p2 p3 p3; p3 p3 p2]
	hmm2ddm!(model,a,b)
end

theta2eta(theta::Tuple)=[log(theta[1]/(1-theta[1])),log(theta[2]/(1-theta[2]))]
eta2theta(eta::Array)=(exp(eta[1])/(1+exp(eta[1])),exp(eta[2])/(1+exp(eta[2])))
coef!(model::ToyModel,eta::Array)=coef!(model,eta2theta(eta))

dim(model::ToyModel)=2

function eta2ab(eta)
	theta=eta2theta(eta)
	a=[theta[1],1-theta[1],1-theta[1],theta[1]]
	p2=theta[2]
	p3=(1-p2)/2
	b=[p2, p3, p3, p3, p3, p2]
	(a,b)
end

#quick implementation of derivatives with numerical gradient
#explicit expressions should be easy to implement
function coef_jac!(model::ToyModel,eta)
	a,b=eta2ab(eta)
	fa(eta)=eta2ab(eta)[1]
	fb(eta)=eta2ab(eta)[2]
	ajac=reshape(ForwardDiff.jacobian(fa)(eta),2,2,2)
	bjac=reshape(ForwardDiff.jacobian(fb)(eta),2,3,2)
	hmm2ddm!(model::DynamicDiscreteModel,reshape(a,2,2),reshape(b,2,3),ajac,bjac)
end

# Plotting function used for the README plot

function lkcontour()
	theta0=(.65, .5)
	model=toymodel()
	coef!(model,theta0)
	data=rand(model,100,100)
	thetahat=eta2theta(mle(model,data))
	thetahat2=eta2theta(em(model,data))
	thetahat3=eta2theta(mle(model,data[1:10]))

	xx=linspace(.1,.9,50)
	yy=linspace(.1,.9,50)
	function ff(x,y,n)
		coef!(model,(x,y))
		loglikelihood(model,data[1:n])
	end
	gg1=[ff(x,y,10)::Float64 for x=xx,y=yy]
	gg2=[ff(x,y,100)::Float64 for x=xx,y=yy]

	qlevels=[0,.25,.5,.75,0.8,0.9,0.95,0.975,0.99,1]
	levels=quantile(vec(gg),qlevels)

	contour(xx,yy,gg,fill=true,levels=levels,color=ColorGradient(:heat,[0,0.98,1]))
	scatter!([theta0[1]],[theta0[2]],c=:green,leg=false)
	scatter!([thetahat[1]],[thetahat[2]],c=:blue,leg=false)
	contour(xx,yy,gg,fill=true,levels=levels,color=ColorGradient(:heat,[0,0.98,1]))
	scatter!([theta0[1]],[theta0[2]],c=:green,leg=false)
	scatter!([thetahat3[1]],[thetahat3[2]],c=:blue,leg=false)
end


using HoloProcessing
using Test
using FFTW
using BenchmarkTools
using LinearAlgebra: norm

@testset "HoloProcessing.jl" begin
    img = [
        16.0 2 3 13
        5 11 10 8
        9 7 6 12
        4 14 15 1
    ]

    @test reconst(img, plan_rfft(img), 1) |> norm ==
          reconst(img, plan_rfft(img), 1; nthreads = true) |> norm ≈
          1.1215103508786834
    @test reconst(img, plan_fft(img), 1) |> norm ==
          reconst(img, plan_fft(img), 1; nthreads = true) |> norm ≈
          1.1375929179890423
    @test reconst(img, plan_fft(img), 1; shift = true) |> norm ==
          reconst(img, plan_fft(img), 1; shift = true, nthreads = true) |> norm ≈
          1.137592917989042
    @test normalize(img) == normalize!(img)
    @test img |> color .|> Float64 == img
    # path = "/home/qling/Documents/Denoising/hologram"
    # holo = load_holo(path, "head.bmp")
    # FFTW.set_num_threads(Sys.CPU_THREADS)
    # P1 = plan_rfft(holo)
    # @benchmark sdm($holo, $(size(holo) .÷ 2), ($50, $100), $P1, $600)

    # @benchmark rse($holo, $(size(holo) .÷ 2), ($50, $100), $P1, $600)

    # P2 = plan_rfft(similar(holo, size(holo) .÷ 2))

    # @benchmark ldr($holo, $2, ($50, $100), $P2, $600)
end


module HearThereTests

using HearThere
using DataFrame
using Base.Test

function generate_rawdata(n::Integer)
    generate_random_rawdata(n, 100, 0, 100)
end

function generate_rawdata(n::Integer, μ::Real, σ::Real)
    generate_random_rawdata(n, fill(μ, 9), fill(σ, 9))
end

function generate_rawdata(n::Integer, μ::Array{Real}, σ::Array{Real})
    DataFrame(
        timestamp=linspace(0, 10, n),
        accel_x=convert(Array{Int32}, map(round, randn(n) * σ[1] + μ[1])),
        accel_y=convert(Array{Int32}, map(round, randn(n) * σ[2] + μ[2])),
        accel_z=convert(Array{Int32}, map(round, randn(n) * σ[3] + μ[3])),
        gyro_x=convert(Array{Int32}, map(round, randn(n) * σ[4] + μ[4])),
        gyro_y=convert(Array{Int32}, map(round, randn(n) * σ[5] + μ[5])),
        gyro_z=convert(Array{Int32}, map(round, randn(n) * σ[6] + μ[6])),
        mag_x=convert(Array{Int32}, map(round, randn(n) * σ[7] + μ[7])),
        mag_y=convert(Array{Int32}, map(round, randn(n) * σ[8] + μ[8])),
        mag_z=convert(Array{Int32}, map(round, randn(n) * σ[9] + μ[9])))
end

function test_gyro_bias()
    n = 1000
    df = generate_rawdata(n, 2, 100)
end

test_gyro_bias()

end


using Zygote
using LinearAlgebra
using GalacticOptim,NLopt,Optim
using Plots
using GLOQ

using Random
Random.seed!(14);

# System parameters for a simple two level open quantum system
N_states = 2; # number of states
freqs = [4.1] # transition frequency in GHz
omegas = 2.0*pi.*freqs # change to angular frequency
gamma1   = [1.0/(45.0*GLOQ.GLOQ_MICRO_SEC)] # Reciprocal of relaxation time - T1 (in units of ns)
gamma2   = [1.0/(24.0*GLOQ.GLOQ_MICRO_SEC)] # Reciprocal of dephasing time - T2 (in units of ns)
omr_ramsey = 2.0*pi*(4.1 - 5.0e-4) # drive frequency for the Ramsey experiment
omr_echo = 2.0*pi*4.1  # drive frequency for the Echo experiment
omr_t1 = 2.0*pi*4.1  # drive frequency for the T1 experiment
TC = 2.5*17.0 # total control time

# Initial state
initial_state = 0
state_u0 = [0.0;0.0]
state_v0 = [0.0;0.0]
state_u0[initial_state+1] = 1.0

# Forward solve to generate synthetic data for the Ramsey experiment
# Duration of the Ramsey experiment, largest dark time
T_Ramsey = 10.0*GLOQ.GLOQ_MICRO_SEC # convert micro-sec to nano-sec
# total number of dark time samples
N_dark_ramsey = 201
dt_ramsey = T_Ramsey/N_dark_ramsey
t_dark_ramsey = collect(range(0.0, T_Ramsey, length=N_dark_ramsey))
# Forward solve
rho_synthetic_ramsey_u,rho_synthetic_ramsey_v = GLOQ.RamseyForwardSolve(
				 state_u0,state_v0, # initial values, u for the real part, v for the imaginary part
			     omegas,omr_ramsey, # transition frequencies, drive frequency
				 gamma1,gamma2, # decay and dephasing
				 initial_state, # initial state
				 TC,t_dark_ramsey,N_states) # control time, dark time, total number of states
population_ramsey_synthetic = GLOQ.get_population(rho_synthetic_ramsey_u)

# Echo experiment
# Duration of the Echo experiment, largest dark time
T_Echo = 30.0*GLOQ.GLOQ_MICRO_SEC # convert micro-sec to nano-sec
# total number of dark time samples
N_dark_echo = 601
# Step size of the Echo experiment
dt_echo = T_Echo/N_dark_echo
t_dark_echo = collect(range(0.0, T_Echo, length=N_dark_echo))
rho_synthetic_echo_u,rho_synthetic_echo_v = GLOQ.EchoForwardSolve(
				 state_u0,state_v0, # initial values, u for the real part, v for the imaginary part
			     omegas,omr_echo, # transition frequencies, drive frequency
				 gamma1,gamma2, # decay and dephasing
				 initial_state, # initial state
				 TC,t_dark_echo,N_states) # control time, dark time, total number of states
population_echo_synthetic = GLOQ.get_population(rho_synthetic_echo_u)

# T1
# Duration of the T1-decay experiment, largest dark time
T_t1 = 40.0*GLOQ.GLOQ_MICRO_SEC # convert micro-sec to nano-sec
# total number of dark time samples
N_dark_t1 = 401
# Step size of the T1 experiment
dt_t1 = T_t1/N_dark_t1
t_dark_t1 = collect(range(0.0, T_t1, length=N_dark_t1))
rho_synthetic_t1_u,rho_synthetic_t1_v = GLOQ.T1ForwardSolve(
				 state_u0,state_v0, # initial values, u for the real part, v for the imaginary part
			     omegas,omr_t1, # transition frequencies, drive frequency
				 gamma1,gamma2, # decay and dephasing
				 initial_state, # initial state
				 TC,t_dark_t1,N_states) # control time, dark time, total number of states
population_t1_synthetic = GLOQ.get_population(rho_synthetic_t1_u)

# Plot the synthetic data
#fig=plot(t_dark_times./GLOQ.GLOQ_MICRO_SEC,population_synthetic)
#display(fig)

# Define the loss function for the GalacticOptim
# p: phyiscal parameters:
#	 p[1] = transition frequency in GHz
#    p[2] = gamma1
#    p[3] = gamma2
# dummy_parameter: needed by GalacticOptim, one can just put [] here
function loss(p,dummy_parameter)
	# Ramsey
	_rho_ramsey_u,_rho_ramsey_v = GLOQ.RamseyForwardSolve(state_u0,state_v0,
				     (2*pi).*[p[1]],omr_ramsey,
					 [p[2]],[p[3]],#gamma1,gamma2,
					 initial_state, # initial state
					 TC,t_dark_ramsey,N_states)
	_population_ramsey = GLOQ.get_population(_rho_ramsey_u)
	# Echo
    _rho_echo_u,_rho_echo_v = GLOQ.EchoForwardSolve(state_u0,state_v0,
				     (2*pi).*[p[1]],omr_echo,
					 [p[2]],[p[3]],#gamma1,gamma2,
					 initial_state, # initial state
					 TC,t_dark_echo,N_states)
	_population_echo = GLOQ.get_population(_rho_echo_u)
	# T1
    _rho_t1_u,_rho_t1_v = GLOQ.T1ForwardSolve(state_u0,state_v0,
				     (2*pi).*[p[1]],omr_t1,
					 [p[2]],[p[3]],#gamma1,gamma2,
					 initial_state, # initial state
					 TC,t_dark_t1,N_states)
	_population_t1 = GLOQ.get_population(_rho_t1_u)

	_loss = sum(abs2,_population_ramsey-population_ramsey_synthetic)*dt_ramsey+
			sum(abs2,_population_echo-population_echo_synthetic)*dt_echo+
			sum(abs2,_population_t1-population_t1_synthetic)*dt_t1

	return _loss
end

plot_callback = function(p,other_args)
	# Ramsey
	rho_ramsey_u,rho_ramsey_v = GLOQ.RamseyForwardSolve(state_u0,state_v0,
					 (2*pi).*[p[1]],omr_ramsey,
					 [p[2]],[p[3]],#gamma1,gamma2,
					 initial_state, # initial state
					 TC,t_dark_ramsey,N_states)
	population_ramsey = GLOQ.get_population(rho_ramsey_u)
	# Echo
    rho_echo_u,rho_echo_v = GLOQ.EchoForwardSolve(state_u0,state_v0,
				     (2*pi).*[p[1]],omr_echo,
					 [p[2]],[p[3]],#gamma1,gamma2,
					 initial_state, # initial state
					 TC,t_dark_echo,N_states)
	population_echo = GLOQ.get_population(rho_echo_u)

	# T1
    rho_t1_u,rho_t1_v = GLOQ.T1ForwardSolve(state_u0,state_v0,
				     (2*pi).*[p[1]],omr_t1,
					 [p[2]],[p[3]],#gamma1,gamma2,
					 initial_state, # initial state
					 TC,t_dark_t1,N_states)
	population_t1 = GLOQ.get_population(rho_t1_u)

	# Plot
	# Ramsey
	fig_ramsey = plot(t_dark_ramsey./GLOQ.GLOQ_MICRO_SEC,population_ramsey_synthetic,label=["Syn-0" "Syn-1"],
			          line = (:dash,0.0), marker = ([:hex :hex], 5, 0.5),legend=:outerright,
					  title="Ramsey");
	plot!(fig_ramsey,t_dark_ramsey./GLOQ.GLOQ_MICRO_SEC,population_ramsey,label=["Opt-0" "Opt-1"],legend=:outerright);
	# Echo
	fig_echo = plot(t_dark_echo./GLOQ.GLOQ_MICRO_SEC,population_echo_synthetic,label=["Syn-0" "Syn-1"],
		  		     line = (:dash,0.0), marker = ([:hex :hex], 5, 0.5),legend=:outerright,
					 title="Echo");
	plot!(fig_echo,t_dark_echo./GLOQ.GLOQ_MICRO_SEC,population_echo,label=["Opt-0" "Opt-1"],legend=:outerright);
	# T1
	fig_t1=plot(t_dark_t1./GLOQ.GLOQ_MICRO_SEC,population_t1_synthetic,label=["Syn-0" "Syn-1"],
			    line = (:dash,0.0), marker = ([:hex :hex], 5, 0.5),legend=:outerright,
				title="T1");
	plot!(fig_t1,t_dark_t1./GLOQ.GLOQ_MICRO_SEC,population_t1,label=["Opt-0" "Opt-1"],legend=:outerright);

	display( plot(fig_ramsey,fig_echo,fig_t1,layout=grid(3,1),size=[1000,1500],
				  legendfontsize=15,xtickfontsize=15,ytickfontsize=15,titlefontsize=18) )
	return false
end

p_true = [freqs;gamma1;gamma2] # values to generate synthetic data
# initial guess for the optimization
p_initial = [freqs.-0.5e-4;0.75.*gamma1;1.25.*gamma2]
# bounds for the optimization
lower_bound = (0.5).*p_true
upper_bound = (1.5).*p_true

# construct optimization object, use Zygote auto-differentiation to compute the gradient
loss_gradient = GalacticOptim.OptimizationFunction(loss, GalacticOptim.AutoZygote())
opt_prob = GalacticOptim.OptimizationProblem(loss_gradient, p_initial,
										 lb = lower_bound, ub = upper_bound)

println("Optim Fminbox(LBFGS) Optimization starts")
@time sol = GalacticOptim.solve(opt_prob ,Fminbox(LBFGS()),
								cb = plot_callback,
								outer_iterations = 20,
								iterations = 10,
								show_trace=true,
								f_tol = 1e-5,
								outer_f_tol = 1e-5)
println("Optim Fminbox(LBFGS) Optimization done")

# present the solutions
println("\nOptimized results: ",sol.u,
        "\nLoss: ",sol.minimum,
		"\nError: ",sol.u-p_true)


# This file illustrates the simulation of feedthrough dynamical system in a unity feedback.

using Causal
using Plots

# Construct the model
x0 = ones(1)

@defmodel model begin
    @nodes begin
        gen = FunctionGenerator(readout=sin)
        adder = Adder(signs=(+,-))
        ds = ContinuousLinearSystem(state=x0)
        writer = Writer()
    end
    @branches begin
        gen[1]      =>      adder[1]
        adder       =>      ds
        ds[1]       =>      adder[2]
        ds          =>      writer
    end
end

# Simulate the model
simulate!(model, 0., 0.01, 10.)

# Read simulation data
t, ys = read(getnode(model, :writer).component)

# Compoute simulation error
r = getnode(model,:gen).component.readout.(t)
xa = (x0[1] + 2 / 13) * exp.(-3 / 2 * t) + 3 / 13 * sin.(t) - 2 / 13 * cos.(t)
ya = (xa + r) / 2
er = ys - ya

# Plot results.
p1 = plot(t, ys, label=:simulation)
    plot!(t, ya, label=:analytic)
p2 = plot(t, er, label=:error)
display(plot(p1, p2, layout=(2, 1)))


using GeoData, Test, ArchGDAL
using GeoData: reproject, convertmode

@testset "reproject" begin
    cea = ProjString("+proj=cea +lon_0=0 +lat_ts=30 +x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs +ellps=WGS84 +towgs84=0,0,0")
    wktcea = convert(WellKnownText, cea)
    projcea = convert(ProjString, cea)
    wkt4326 = convert(WellKnownText, EPSG(4326))
    proj4326 = convert(ProjString, EPSG(4326))

    @test reproject(proj4326, cea, Lat(), -30.0) ≈ -3.658789324855012e6
    @test reproject(wktcea, EPSG(4326), Lat(), -3.658789324855012e6) ≈ -30.0
    @test reproject(projcea, wkt4326, Lat(), -3.658789324855012e6) .≈ -30.0
    @test reproject(cea, proj4326, Lat(), [-3.658789324855012e6]) ≈ [-30.0]

    @test reproject(proj4326, cea, Lon(), 180.0) ≈ 1.7367530445161372e7
    @test reproject(cea, EPSG(4326), Lon(), 1.7367530445161372e7) ≈ 180.0
    @test reproject(cea, wkt4326, Lon(), 1.7367530445161372e7) .≈ 180.0
    @test reproject(cea, proj4326, Lon(), [1.7367530445161372e7]) ≈ [180.0]
end

@testset "convertmode" begin
    projcea = ProjString("+proj=cea")
    proj4326 = convert(ProjString, EPSG(4326))

    lonstart, lonend = 0.5, 179.5
    cealonstart, cealonend = reproject(proj4326, projcea, Lon(), [lonstart, lonend])
    cealonrange = LinRange(cealonstart, cealonend, 180)
    lon = Lon(cealonrange; mode=Projected(Ordered(), Regular(step(cealonrange)),
              Intervals(Center()), projcea, proj4326))
    convertedlon = convertmode(Converted, lon)
    @test all(isapprox.(bounds(convertedlon), (0.0, 180.0); atol=1e-10))
    @test val(convertedlon) ≈ 0.5:179.5

    projectedlon = convertmode(Projected, convertedlon)
    @test val(projectedlon) ≈ val(lon)
    @test all(isapprox.(bounds(projectedlon), bounds(lon); atol=1e-10))

    # Lat in not linear
    latbounds = -90.0, 90.0
    ceabounds = reproject(proj4326, projcea, Lat(), latbounds)
    # Vals are the bounding range without the end bound
    cealatrange = LinRange(ceabounds..., 181)[1:180]
    lat = Lat(cealatrange; mode=Projected(Ordered(), Regular(step(cealatrange)),
              Intervals(Start()), projcea, proj4326))
    convertedlat = convertmode(Converted, lat)
    @test first(convertedlat) ≈ -90.0 atol=1e-5
    @test all(isapprox.(bounds(convertedlat), (-90.0, 90.0); atol=1e-5))
    @test val(lat) ≈ reproject(proj4326, projcea, Lat(), val(convertedlat))

    projectedlat = convertmode(Projected, convertedlat)
    @test val(projectedlat) ≈ val(lat)
    @test all(bounds(projectedlat) .≈ bounds(lat))

    A = DimArray(zeros(length(lon), length(lat)), (lon, lat))
    Aconv = convertmode(Converted, A)

    @test dims(Aconv) == (convertedlon, convertedlat)
end


export getData

function getData(m,pFor::EikonalInvParam,doClear::Bool=false)

    # extract pointers
    Mesh   	= pFor.Mesh
    Q     	= pFor.Sources
    P     	= pFor.Receivers
	n_nodes = pFor.Mesh.n+1;
    nrec  	= size(P,2)
    nsrc  	= size(Q,2)

    # allocate space for data and fields
    D  = zeros(nrec,nsrc)
    pEik = Array{EikonalParam}(nsrc);

	if pFor.useFilesForFields
		tfilename = getFieldsFileName();
		tfile = matopen(tfilename, "w");
	end

	pMem = getEikonalTempMemory(n_nodes);
	ntup = tuple(n_nodes...);
	T = zeros(Float64,ntup) # V is a temporary array of Float64.
    T1_temp = zeros(Float32,ntup);
	ordering_temp = zeros(FactoredEikonalFastMarching.EikIdxType,prod(n_nodes));
	OP_temp 	  = zeros(Int8,ntup)
	for k=1:nsrc
		src_k_loc = Q.rowval[k];
		if Mesh.dim==2
			src = zeros(Int64,2);
			cs2loc(src,src_k_loc,n_nodes);
		else
			src = zeros(Int64,3);
			cs2loc3D(src,src_k_loc,n_nodes);
		end
		m = reshape(m,ntup);
		pEik[k] = getEikonalParam(Mesh,m,src,pFor.HO);
		pEik[k].T1 = T; # Here, we set T (of Float64) for the calculation, but reuse the memory.
		pEik[k].ordering = ordering_temp;
		pEik[k].OP		 = OP_temp;
		solveFastMarchingUpwindGrad(pEik[k],pMem);
		T1_temp[:] = T;
		if Mesh.dim==2
			selfMultiplyWithAnalyticSolution2D(n_nodes,Mesh.h,src,T);
		else
			selfMultiplyWithAnalyticSolution3D(n_nodes,Mesh.h,src,T)
		end
		D[:,k] = P'*T[:];
		if doClear
			FactoredEikonalFastMarching.clear!(pEik[k]);
		elseif pFor.useFilesForFields
			write(tfile,string("T1_",k),T1_temp);
			write(tfile,string("ordering_",k),pEik[k].ordering);
			write(tfile,string("OP_",k),pEik[k].OP);
			FactoredEikonalFastMarching.clear!(pEik[k]);
		else
			# Only here we actually save the data.
			pEik[k].T1 = copy(T1_temp);
			pEik[k].ordering = copy(ordering_temp);
			pEik[k].OP = copy(OP_temp);
		end
	end
	if pFor.useFilesForFields
		close(tfile);
		gc()
	end
	pFor.eikonalParams = pEik;
    return D,pFor
end


# Tests of intstats.jl

using MLBase
using Test

# classify

ss = rand(8, 50)
for i = 1:size(ss,2)
    ss_i = ss[:,i]
    kmax = argmax(ss_i)
    kmin = argmin(ss_i)
    vmax = ss_i[kmax]
    vmin = ss_i[kmin]

    @test classify(ss_i) == kmax
    @test classify(ss_i, Forward) == kmax
    @test classify(ss_i, Reverse) == kmin

    @test classify(ss_i, 0.8) == (vmax >= 0.8 ? kmax : 0)
    @test classify(ss_i, 0.8, Forward) == (vmax >= 0.8 ? kmax : 0)
    @test classify(ss_i, 0.2, Reverse) == (vmin <= 0.2 ? kmin : 0)

    @test classify_withscore(ss_i) == (kmax, ss_i[kmax])
    @test classify_withscore(ss_i, Forward) == (kmax, ss_i[kmax])
    @test classify_withscore(ss_i, Reverse) == (kmin, ss_i[kmin])
end

rmax = Int[argmax(ss[:,i]) for i = 1:size(ss,2)]
rmin = Int[argmin(ss[:,i]) for i = 1:size(ss,2)]
vmax = ss[LinearIndices(size(ss))[CartesianIndex.(rmax, 1:size(ss,2))]]
vmin = ss[LinearIndices(size(ss))[CartesianIndex.(rmin, 1:size(ss,2))]]

trmax = copy(rmax); trmax[vmax .< 0.8] .= 0
trmin = copy(rmin); trmin[vmin .> 0.2] .= 0

@test classify(ss) == rmax
@test classify(ss, Forward) == rmax
@test classify(ss, Reverse) == rmin

@test classify_withscores(ss) == (rmax, vmax)
@test classify_withscores(ss, Forward) == (rmax, vmax)
@test classify_withscores(ss, Reverse) == (rmin, vmin)

@test classify(ss, 0.8) == trmax
@test classify(ss, 0.8, Forward) == trmax
@test classify(ss, 0.2, Reverse) == trmin

# labelmap & labelencode

xs = ["a", "a", "b", "b", "a", "b", "c", "a"]
labels = [1, 1, 2, 2, 1, 2, 3, 1]
lmap = labelmap(xs)

@test keys(lmap) == ["a", "b", "c"]
@test labelencode(lmap, xs) == labels
@test labeldecode(lmap, labels) == xs

# groupindices

gs = Any[[1,2,5,8],[3,4,6],[7]]

@test groupindices(3, labels) == gs
@test groupindices(lmap, xs) == gs


# Parts of this code were taken / derived from NetworkX. See LICENSE for
# licensing details.

"""
    CompleteGraph(n)

Create an undirected [complete graph](https://en.wikipedia.org/wiki/Complete_graph)
with `n` vertices.
"""
function CompleteGraph(n::Integer)
    g = Graph(n)
    for i = 1:n, j=1:n
        if i < j
            add_edge!(g, Edge(i,j))
        end
    end
    return g
end

@doc_str """
    CompleteBipartiteGraph(n1, n2)

Create an undirected [complete bipartite graph](https://en.wikipedia.org/wiki/Complete_bipartite_graph)
with `n1 + n2` vertices.
"""
function CompleteBipartiteGraph(n1::Integer, n2::Integer)
    g = Graph(n1+n2)
    for i = 1:n1, j=n1+1:n1+n2
        add_edge!(g, Edge(i,j))
    end
    return g
end

"""
    CompleteDiGraph(n)

Create a directed [complete graph](https://en.wikipedia.org/wiki/Complete_graph)
with `n` vertices.
"""
function CompleteDiGraph(n::Integer)
    g = DiGraph(n)
    for i = 1:n, j=1:n
        if i != j
            add_edge!(g, Edge(i,j))
        end
    end
    return g
end

"""
    StarGraph(n)

Create an undirected [star graph](https://en.wikipedia.org/wiki/Star_(graph_theory))
with `n` vertices.
"""
function StarGraph(n::Integer)
    g = Graph(n)
    for i = 2:n
        add_edge!(g, Edge(1,i))
    end
    return g
end

"""
    StarDiGraph(n)

Create a directed [star graph](https://en.wikipedia.org/wiki/Star_(graph_theory))
with `n` vertices.
"""
function StarDiGraph(n::Integer)
    g = DiGraph(n)
    for i = 2:n
        add_edge!(g, Edge(1,i))
    end
    return g
end

"""
    PathGraph(n)

Create an undirected [path graph](https://en.wikipedia.org/wiki/Path_graph)
with `n` vertices.
"""
function PathGraph(n::Integer)
    g = Graph(n)
    for i = 2:n
        add_edge!(g, Edge(i-1, i))
    end
    return g
end

"""
    PathDiGraph(n)

Creates a directed [path graph](https://en.wikipedia.org/wiki/Path_graph)
with `n` vertices.
"""
function PathDiGraph(n::Integer)
    g = DiGraph(n)
    for i = 2:n
        add_edge!(g, Edge(i-1, i))
    end
    return g
end

"""
    CycleGraph(n)

Create an undirected [cycle graph](https://en.wikipedia.org/wiki/Cycle_graph)
with `n` vertices.
"""
function CycleGraph(n::Integer)
    g = Graph(n)
    for i = 1:n-1
        add_edge!(g, Edge(i, i+1))
    end
    add_edge!(g, Edge(n, 1))
    return g
end

"""
    CycleDiGraph(n)

Create a directed [cycle graph](https://en.wikipedia.org/wiki/Cycle_graph)
with `n` vertices.
"""
function CycleDiGraph(n::Integer)
    g = DiGraph(n)
    for i = 1:n-1
        add_edge!(g, Edge(i, i+1))
    end
    add_edge!(g, Edge(n, 1))
    return g
end

"""
    WheelGraph(n)

Create an undirected [wheel graph](https://en.wikipedia.org/wiki/Wheel_graph)
with `n` vertices.
"""
function WheelGraph(n::Integer)
    g = StarGraph(n)
    for i = 3:n
        add_edge!(g, Edge(i-1, i))
    end
    if n != 2
        add_edge!(g, Edge(n, 2))
    end
    return g
end

"""
    WheelDiGraph(n)

Create a directed [wheel graph](https://en.wikipedia.org/wiki/Wheel_graph)
with `n` vertices.
"""
function WheelDiGraph(n::Integer)
    g = StarDiGraph(n)
    for i = 3:n
        add_edge!(g, Edge(i-1, i))
    end
    if n != 2
        add_edge!(g, Edge(n, 2))
    end
    return g
end

@doc_str """
    Grid(dims; periodic=false)

Create a ``|dims|``-dimensional cubic lattice, with length `dims[i]`
in dimension `i`.

### Optional Arguments
- `periodic=false`: If true, the resulting lattice will have periodic boundary
condition in each dimension.
"""
function Grid(dims::AbstractVector; periodic=false)
    if periodic
        g = CycleGraph(dims[1])
        for d in dims[2:end]
            g = cartesian_product(CycleGraph(d), g)
        end
    else
        g = PathGraph(dims[1])
        for d in dims[2:end]
            g = cartesian_product(PathGraph(d), g)
        end
    end
    return g
end

"""
    BinaryTree(k::Integer)

Create a [binary tree](https://en.wikipedia.org/wiki/Binary_tree)
of depth `k`.
"""

function BinaryTree(k::Integer)
    g = Graph(Int(2^k-1))
    for i in 0:k-2
        for j in 2^i:2^(i+1)-1
            add_edge!(g, j, 2j)
            add_edge!(g, j, 2j+1)
        end
    end
    return g
end

"""
    BinaryTree(k::Integer)

Create a double complete binary tree with `k` levels.

### References
- Used as an example for spectral clustering by Guattery and Miller 1998.
"""
function DoubleBinaryTree(k::Integer)
    gl = BinaryTree(k)
    gr = BinaryTree(k)
    g = blkdiag(gl, gr)
    add_edge!(g,1, nv(gl)+1)
    return g
end

"""
    RoachGraph(k)

Create a Roach Graph of size `k`.

### References
- Guattery and Miller 1998
"""
function RoachGraph(k::Integer)
    dipole = CompleteGraph(2)
    nopole = Graph(2)
    antannae = crosspath(k, nopole)
    body = crosspath(k,dipole)
    roach = blkdiag(antannae, body)
    add_edge!(roach, nv(antannae)-1, nv(antannae)+1)
    add_edge!(roach, nv(antannae), nv(antannae)+2)
    return roach
end

"""
    CliqueGraph(k, n)

Create a graph consisting of `n` connected `k`-cliques.
"""
function CliqueGraph(k::Integer, n::Integer)
    g = Graph(k*n)
    for c=1:n
        for i=(c-1)*k+1:c*k-1, j=i+1:c*k
            add_edge!(g, i, j)
        end
    end
    for i=1:n-1
        add_edge!(g, (i-1)*k+1, i*k+1)
    end
    add_edge!(g, 1, (n-1)*k+1)
    return g
end


@testset "Basic Vec3 Functionality" begin
    # Check the constructors
    a = Vec3(1.0)

    @test a.x[] == 1.0 && a.y[] == 1.0 && a.z[] == 1.0

    a = Vec3([1.0, 1.0])

    @test a.x == [1.0, 1.0] && a.y == [1.0, 1.0] && a.z == [1.0, 1.0]

    # Operators +, -, * | Both Vec3
    for op in (:+, :-, :*)
        len = rand(1:25)

        a1 = rand(len)
        b1 = rand(len)
        c1 = rand(len)

        a2 = rand(len)
        b2 = rand(len)
        c2 = rand(len)

        vec3_1 = Vec3(a1, b1, c1)
        vec3_2 = Vec3(a2, b2, c2)

        res = @eval $(op)($vec3_1, $vec3_2)
        res_x = @eval broadcast($(op), $a1, $a2)
        res_y = @eval broadcast($(op), $b1, $b2)
        res_z = @eval broadcast($(op), $c1, $c2)

        @test res.x == res_x && res.y == res_y && res.z == res_z
    end

    # Operators +, -, *, /, % | Only 1 Vec3
    for op in (:+, :-, :*, :/, :%)
        len = rand(1:25)

        a = rand(len)
        b = rand(len)
        c = rand(len)

        num = rand()

        vec3_1 = Vec3(a, b, c)

        res1 = @eval $(op)($vec3_1, $num)
        res1_x = @eval broadcast($(op), $a, $num)
        res1_y = @eval broadcast($(op), $b, $num)
        res1_z = @eval broadcast($(op), $c, $num)

        @test res1.x == res1_x && res1.y == res1_y && res1.z == res1_z

        res2 = @eval $(op)($num, $vec3_1)
        res2_x = @eval broadcast($(op), $a, $num)
        res2_y = @eval broadcast($(op), $b, $num)
        res2_z = @eval broadcast($(op), $c, $num)

        @test res2.x == res2_x && res2.y == res2_y && res2.z == res2_z
    end

    # Other Operators
    a_vec = Vec3(1.0, 2.0, 3.0)
    b_vec = Vec3(4.0, 5.0, 6.0)
    neg_a_vec = - a_vec
    cross_prod = RayTracer.cross(a_vec, b_vec)
    b_clamped = clamp(b_vec, 4.5, 5.5)

    @test neg_a_vec.x[] == -1.0 && neg_a_vec.y[] == -2.0 && neg_a_vec.z[] == -3.0
    @test RayTracer.dot(a_vec, b_vec)[] == 32.0
    @test RayTracer.l2norm(a_vec)[] == 14.0
    @test RayTracer.l2norm(RayTracer.normalize(a_vec))[] == 1.0
    @test maximum(a_vec) == 3.0
    @test minimum(a_vec) == 1.0
    @test size(a_vec) == (1,)
    @test getindex(a_vec, 1) == (x = 1.0, y = 2.0, z = 3.0)
    @test cross_prod.x[] == -3.0 && cross_prod.y[] == 6.0 && cross_prod.z[] == -3.0
    @test b_clamped.x[] == 4.5 && b_clamped.y[] == 5.0 && b_clamped.z[] == 5.5

end


#=
max_percent_diff.jl
Copyright 2000-9-16, Jeff Fessler, University of Michigan
=#

export max_percent_diff

"""
    d = max_percent_diff(s1, s2, [options])

Compute the "maximum percent difference" between two signals: `s1, s2`.

Default is to normalize by `maximum(abs.(s1))`.

options
- `maxboth::Bool` use max of both arguments to normalize; default `false`
- `normalize::Bool` normalize each before comparing; default `false`

"""
function max_percent_diff(s1, s2 ;
        maxboth::Bool = false, normalize::Bool = false)

    # check to make sure that we have comparable signals
    size(s1) != size(s2) &&
        throw(DomainError("Dimension mismatch $(size(s1)) $(size(s2))"))

    any(isnan.(s1)) && @warn("NaN values in s1 input to max_percent_diff")
    any(isnan.(s2)) && @warn("NaN values in s2 input to max_percent_diff")

    if normalize
        normFactor = sum(abs.(s2)) / sum(abs.(s1))
        isinteractive() && (@info "Normalization factor (matrix 2 / matrix 1): $normFactor \n")
        s1 /= sum(abs.(s1))
        s2 /= sum(abs.(s2))
    end

    if maxboth
        denom = max(maximum(abs.(s1)), maximum(abs.(s2)))
        d = denom == 0 ? 0 : maximum(abs.(s1 - s2)) / denom
    else
        denom = maximum(abs.(s1))
        denom = denom == 0 ? denom : maximum(abs.(s2))
        d = denom == 0 ? 0 : maximum(abs.(s1 - s2)) / denom
    end

    return 100 * d
end


# -*- coding: utf-8 -*-
# Load necessary libraries

# +
using Gen
using Statistics
using ProfileView #Make sure this is loaded before PyPlot to avoid stupid gtk version issues
using Profile
using PyPlot

include("DistributionsBacked.jl")
include("IterDeep.jl")
include("AnimatedPyplot.jl")
include("nuts.jl")

include("unrealizedNormals.jl")

const my_inv_gamma = DistributionsBacked{Float64}((alpha, theta) -> Distributions.InverseGamma(alpha, theta), [true, true], true)
# -
# Space for declaring constants.

disable_sample_logging()
INV_GAMMA_PRIOR_CONSTANT = 0.1;

# Canonical Rats data, from from section 6 of Gelfand et al (1990). See http://www.openbugs.net/Examples/Rats.html

ys_raw = ([151, 145, 147, 155, 135, 159, 141, 159, 177, 134,
    160, 143, 154, 171, 163, 160, 142, 156, 157, 152, 154, 139, 146,
    157, 132, 160, 169, 157, 137, 153, 199, 199, 214, 200, 188, 210,
    189, 201, 236, 182, 208, 188, 200, 221, 216, 207, 187, 203, 212,
    203, 205, 190, 191, 211, 185, 207, 216, 205, 180, 200, 246, 249,
    263, 237, 230, 252, 231, 248, 285, 220, 261, 220, 244, 270, 242,
    248, 234, 243, 259, 246, 253, 225, 229, 250, 237, 257, 261, 248,
    219, 244, 283, 293, 312, 272, 280, 298, 275, 297, 350, 260, 313,
    273, 289, 326, 281, 288, 280, 283, 307, 286, 298, 267, 272, 285,
    286, 303, 295, 289, 258, 286, 320, 354, 328, 297, 323, 331, 305,
    338, 376, 296, 352, 314, 325, 358, 312, 324, 316, 317, 336, 321,
    334, 302, 302, 323, 331, 345, 333, 316, 291, 324])
ys = reshape([Float64(y) for y = ys_raw],30,5)
xs = [8.0, 15.0, 22.0, 29.0, 36.0]
ys = ys ./ 100.
(N,T) = size(ys)

# These data are the weights of 30 rats, measured at 5 common age values in days. The model is a simple hierarchical/random effects one; a random intercept and random slope for each rat, plus an error term (presumably representing not measurement error, but just the random variation of individual growth curves).

for i in 1:30
    plot(xs,ys[i,:])
end
xlabel("Age (days)")
ylabel("Weight (g)")
title("Rat growth")

# +
@gen function oneRat(xs, T, mu_alpha, sigmasq_alpha, mu_beta, sigmasq_beta, sigmasq_y)

    alpha ~ normal(mu_alpha, sqrt(sigmasq_alpha))
    beta ~ normal(mu_beta, sqrt(sigmasq_beta))
    y ~ broadcasted_normal([alpha + beta * (xs[t]) for t = 1:T],
                                        sqrt(sigmasq_y))
end

@gen function ratsModel(xs::Vector,N::Int32,)
    T = length(xs)
    xbar = mean(xs) #could be precomputed, but YKWTS about premature optimization...

    mu_alpha ~ normal(0, 10.0)
    mu_beta ~ normal(0, 1.0)
    sigmasq_y ~ my_inv_gamma(INV_GAMMA_PRIOR_CONSTANT, INV_GAMMA_PRIOR_CONSTANT)
    sigmasq_alpha ~ my_inv_gamma(INV_GAMMA_PRIOR_CONSTANT, INV_GAMMA_PRIOR_CONSTANT)
    sigmasq_beta ~ my_inv_gamma(INV_GAMMA_PRIOR_CONSTANT, INV_GAMMA_PRIOR_CONSTANT)

    alpha = Vector{Float64}(undef,N)
    beta = Vector{Float64}(undef,N)
    y = Vector{Vector{Float64}}(undef,N)
    for n in 1:N
        {:data => n} ~ oneRat(xs .- xbar, T, mu_alpha, sigmasq_alpha, mu_beta, sigmasq_beta, sigmasq_y)
    end
end

model_transformations = choicemap()
model_transformations[:sigmasq_y] = transform_log
model_transformations[:sigmasq_alpha] = transform_log
model_transformations[:sigmasq_beta] = transform_log

;
# -

function make_constraints(ys::Array)
    (N,T) = size(ys)
    constraints = Gen.choicemap()
    for n in 1:N
        constraints[:data => n => :y] = ys[n,:]
    end
    constraints
end
;

function mcmc_inference(xs, ys, num_iters, update, model = ratsModel)
    results = Vector{Gen.DynamicDSLTrace{DynamicDSLFunction{Any}}}(undef, num_iters+1)
    N = size(ys)[1]
    observations = make_constraints(ys)
    (trace, _) = generate(model, (xs, N), observations)
    results[1] = trace
    selection = Gen.complement(select(observations))
    for iter=1:num_iters
        trace = update(trace,N,selection)
        results[iter+1] = trace
    end
    results
end
;

# +

@gen function oneDeferredRat(xs, T, mu_alpha, sigmasq_alpha, mu_beta, sigmasq_beta, sigmasq_y)

    alpha ~ unrealizedIndyNormal(mu_alpha, sqrt(sigmasq_alpha))
    beta ~ unrealizedIndyNormal(mu_beta, sqrt(sigmasq_beta))
    for t = 1:T
        @trace(observedNormal(alpha + beta * (xs[t]), sqrt(sigmasq_y)), "y$(t)")
    end
end

@gen function deferredRatsModel(xs::Vector,N::Int32,)
    T = length(xs)
    xbar = mean(xs) #could be precomputed, but YKWTS about premature optimization...

    mu_alpha ~ normal(0, 10.0)
    mu_beta ~ normal(0, 1.0)
    sigmasq_y ~ my_inv_gamma(INV_GAMMA_PRIOR_CONSTANT, INV_GAMMA_PRIOR_CONSTANT)
    sigmasq_alpha ~ my_inv_gamma(INV_GAMMA_PRIOR_CONSTANT, INV_GAMMA_PRIOR_CONSTANT)
    sigmasq_beta ~ my_inv_gamma(INV_GAMMA_PRIOR_CONSTANT, INV_GAMMA_PRIOR_CONSTANT)

    alpha = Vector{Float64}(undef,N)
    beta = Vector{Float64}(undef,N)
    y = Vector{Vector{Float64}}(undef,N)
    for n in 1:N
        {:data => n} ~ oneDeferredRat(xs .- xbar, T, mu_alpha, sigmasq_alpha, mu_beta, sigmasq_beta, sigmasq_y)
    end
end

model_transformations = choicemap()
model_transformations[:sigmasq_y] = transform_log
model_transformations[:sigmasq_alpha] = transform_log
model_transformations[:sigmasq_beta] = transform_log

;

# +
Profile.clear()

function block_mh_d(tr,N,selection)
    (tr, _) = mh(tr, select(:mu_alpha, :mu_beta))
    (tr, _) = mh(tr, select(:sigmasq_alpha, :sigmasq_beta, :sigmasq_y))

    for n in 1:N
        (tr, _) = mh(tr, select(:data => n => :y1,:data => n => :y2,:data => n => :y3,:data => n => :y4,
                                :data => n => :alpha,:data => n => :beta))
    end
    tr
end
;

VIZ_INTERVAL = 1
VIZ_FRAMES = 20
trs3= Profile.@profile mcmc_inference(xs, ys, VIZ_INTERVAL * VIZ_FRAMES, block_mh_d, deferredRatsModel)
tr3 = trs3[end]
print("mu_alpha: $(tr3[:mu_alpha]), mu_beta: $(tr3[:mu_beta]), sigma_y: $(sqrt(tr3[:sigmasq_y]))\n")
;

# +
c = make_constraints(ys)
alphahats = mean(ys,dims=2)
alphahat = mean(alphahats)
betahats = transpose(ys)\xs
betahat = mean(betahats)
c[:mu_alpha] = alphahat
c[:mu_beta] = betahat

c[:sigmasq_beta] = var(betahats) * 2.1601
c[:sigmasq_alpha] = var(alphahats) * 2.1601
c[:sigmasq_y] = 0.000001

tr, w = generate(deferredRatsModel, (ys, N), c) #w == -N/2

s = get_score(tr)

c[:sigmasq_y] = 1000.

c[:sigmasq_beta] = var(betahats) * 2.1603
c[:sigmasq_alpha] = var(alphahats) * 2.1603
c[:sigmasq_y] = 0.000001

tr2, w2 = generate(deferredRatsModel, (ys, N), c) #w == -N/2

s2 = get_score(tr2)

w, s, w2, s2
# -

#
# ...
#
# ...
#
# ...
#
# ...
#
# ...
#
# ...
#
# ...
#
# ...
#
# ...
#
# ...
#
# ...
#
#
#
#
#
#
#
#

function block_mh(tr,N,selection)
    (tr, _) = mh(tr, select(:mu_alpha, :mu_beta))
    (tr, _) = mh(tr, select(:sigmasq_alpha, :sigmasq_beta, :sigmasq_y))

    for n in 1:N
        (tr, _) = mh(tr, select(:data => n => :y,
                                :data => n => :alpha,:data => n => :beta))
    end
    tr
end
;

function simple_hmc(tr,N,selection)
    (tr, _) = hmc(tr, selection)
    tr
end
;

VIZ_INTERVAL = 2
VIZ_FRAMES = 10
trs= mcmc_inference(xs, ys, VIZ_INTERVAL * VIZ_FRAMES, block_mh)
tr = trs[end]
get_choices(tr)
;

print("mu_alpha: $(tr[:mu_alpha]), mu_beta: $(tr[:mu_beta]), sigma_y: $(sqrt(tr[:sigmasq_y]))")

# +
fewcolors = ["r", "b", "g"] #In order to avoid busy graphs, I graph just 3 rats per plot
numPlots = 2 #...but tile 2 plots
numRats = length(fewcolors)
CONF95SIGS = 1.96

function make_axes()
    fig = figure("A few rats",figsize=(5,5))
    my_axes = fig.subplots(nrows=numPlots)
    for ax in my_axes
        ax.set_xlim((5,40))
    end
    (fig,my_axes)
end

function visualize_rats(trace,my_axes)
    xbar = mean(xs)
    xgrid = collect(min(xs...):.5:max(xs...))
    xlims = [min(xs...), max(xs...)]
    centeredgrid = xgrid .- xbar
    linesigmas = sqrt.(trace[:sigmasq_alpha] .+ trace[:sigmasq_beta] .* centeredgrid .^ 2)
    meanline = trace[:mu_alpha] .+ trace[:mu_beta] .* centeredgrid
    upperline = meanline .+ CONF95SIGS .* linesigmas
    lowerline = meanline .- CONF95SIGS .* linesigmas
    for (j,ax) in enumerate(my_axes)
        ax.clear()
        ax.fill_between(xgrid,lowerline,upperline,alpha=.1,color="black")
        ax.fill_between([6.,7],
                        fill(250. + CONF95SIGS * sqrt(trace[:sigmasq_y]),2),
                        fill(250. - CONF95SIGS * sqrt(trace[:sigmasq_y]),2),
            alpha=1,color="black")

        ax.set_ylim((min(1.00,lowerline[1],lowerline[end]),
                max(4.00,upperline[1],upperline[end])))
        for k in 1:numRats
            i = j*numRats + k
            ax.plot(xs,ys[i,:],color=fewcolors[k], alpha=1.)

            ratline = trace[:data => i => :alpha] .+ trace[:data => i => :beta] .* (xlims .- xbar)
            ax.plot(xlims, ratline,color=fewcolors[k], alpha=.3)
        end
        ax.set_ylabel("Weight (g)")
        if j == 1
            ax.set_title("Rat growth")
        else
            ax.set_xlabel("Age (days)")
        end

    end
end

# +

visualize_rats(tr,make_axes()[2])
# -

display_animation("rats_mh", make_axes,
    (ax)->visualize_rats(trs[1],ax),
    (ax,i)->visualize_rats(trs[1+i*VIZ_INTERVAL],ax),
    VIZ_FRAMES)

VIZ_INTERVAL = 2
VIZ_FRAMES = 10
trs2= mcmc_inference(xs, ys, VIZ_INTERVAL * VIZ_FRAMES, simple_hmc)
tr2 = trs2[end]
get_choices(tr2)
;

print("mu_alpha: $(tr2[:mu_alpha]), mu_beta: $(tr2[:mu_beta]), sigma_y: $(sqrt(tr2[:sigmasq_y]))\n")

if false
    display_animation("rats_hmc", make_axes,
        (ax)->visualize_rats(trs[1],ax),
        (ax,i)->visualize_rats(trs[1+i*VIZ_INTERVAL],ax),
        VIZ_FRAMES)
end

# +
Profile.clear()

VIZ_INTERVAL = 1
VIZ_FRAMES = 20
trs3= Profile.@profile mcmc_inference(xs, ys, VIZ_INTERVAL * VIZ_FRAMES, (tr, N, selection)->my_nuts(tr, selection, model_transformations))
tr3 = trs3[end]
print("mu_alpha: $(tr3[:mu_alpha]), mu_beta: $(tr3[:mu_beta]), sigma_y: $(sqrt(tr3[:sigmasq_y]))\n")
;
# -

ProfileView.view()

open("profile.txt", "w") do s
    Profile.print(IOContext(s, :displaysize => (24, 500)))
end

display_animation("rats_nuts", make_axes,
    (ax)->visualize_rats(trs3[1],ax),
    (ax,i)->visualize_rats(trs3[1+i*VIZ_INTERVAL],ax),
    VIZ_FRAMES,200)

# The below is just a sketch of how I think "semi-automatic conjugacy" might work. It's not working code, yet.

l = [(4,6)]
for (a, (b,c)) in enumerate(l)
    @warn "hi" a b c
end


# TODO: add keyword `shrink` to give a consistant result on Base
#       when this is done, then we can propose this change to upstream Base
"""
    fdiff(A::AbstractArray; dims::Int, rev=false, boundary=:periodic)

A one-dimension finite difference operator on array `A`. Unlike `Base.diff`, this function doesn't
shrink the array size.

Take vector as an example, it computes `(A[2]-A[1], A[3]-A[2], ..., A[1]-A[end])`.

# Keywords

- `rev::Bool`
  If `rev==true`, then it computes the backward difference
  `(A[end]-A[1], A[1]-A[2], ..., A[end-1]-A[end])`.
- `boundary`
  By default it computes periodically in the boundary, i.e., `:periodic`.
  In some cases, one can fill zero values with `boundary=:zero`.

# Examples

```jldoctest; setup=:(using ImageBase: fdiff)
julia> A = [2 4 8; 3 9 27; 4 16 64]
3×3 $(Matrix{Int}):
 2   4   8
 3   9  27
 4  16  64

julia> diff(A, dims=2) # this function exists in Base
3×2 $(Matrix{Int}):
  2   4
  6  18
 12  48

julia> fdiff(A, dims=2)
3×3 $(Matrix{Int}):
  2   4   -6
  6  18  -24
 12  48  -60

julia> fdiff(A, dims=2, rev=true) # reverse diff
3×3 $(Matrix{Int}):
  -6   2   4
 -24   6  18
 -60  12  48

julia> fdiff(A, dims=2, boundary=:zero) # fill boundary with zeros
3×3 $(Matrix{Int}):
  2   4  0
  6  18  0
 12  48  0
```

See also [`fdiff!`](@ref) for the in-place version.
"""
fdiff(A::AbstractArray; kwargs...) = fdiff!(similar(A), A; kwargs...)

"""
    fdiff!(dst::AbstractArray, src::AbstractArray; dims::Int, rev=false, boundary=:periodic)

The in-place version of [`ImageBase.fdiff`](@ref)
"""
function fdiff!(dst::AbstractArray, src::AbstractArray;
        dims=_fdiff_default_dims(src),
        rev=false,
        boundary::Symbol=:periodic)
    isnothing(dims) && throw(UndefKeywordError(:dims))
    axes(dst) == axes(src) || throw(ArgumentError("axes of all input arrays should be equal. Instead they are $(axes(dst)) and $(axes(src))."))
    N = ndims(src)
    1 <= dims <= N || throw(ArgumentError("dimension $dims out of range (1:$N)"))

    r = axes(src)
    r0 = ntuple(i -> i == dims ? UnitRange(first(r[i]), last(r[i]) - 1) : UnitRange(r[i]), N)
    r1 = ntuple(i -> i == dims ? UnitRange(first(r[i])+1, last(r[i])) : UnitRange(r[i]), N)

    d0 = ntuple(i -> i == dims ? UnitRange(last(r[i]), last(r[i])) : UnitRange(r[i]), N)
    d1 = ntuple(i -> i == dims ? UnitRange(first(r[i]), first(r[i])) : UnitRange(r[i]), N)

    if rev
        dst[r1...] .= view(src, r1...) .- view(src, r0...)
        if boundary == :periodic
            dst[d1...] .= view(src, d1...) .- view(src, d0...)
        elseif boundary == :zero
            dst[d1...] .= zero(eltype(dst))
        else
            throw(ArgumentError("Wrong boundary condition $boundary"))
        end
    else
        dst[r0...] .= view(src, r1...) .- view(src, r0...)
        if boundary == :periodic
            dst[d0...] .= view(src, d1...) .- view(src, d0...)
        elseif boundary == :zero
            dst[d0...] .= zero(eltype(dst))
        else
            throw(ArgumentError("Wrong boundary condition $boundary"))
        end
    end

    return dst
end

_fdiff_default_dims(A) = nothing
_fdiff_default_dims(A::AbstractVector) = 1


# linear-map/getindex.jl
#
# Provides getindex() capabilities like A[:,j] for LinearMap objects.
#
# Currently this provides only a partial set of the possible ways
# one can use indexing for a matrix, because currently
# LinearMap is not a subtype of an AbstractMatrix.
#
# If LinearMap were a subtype of an AbstractMatrix, then all possible
# indexing will be supported by Base.getindex, albeit very likely
# by quite inefficient iterators.
#
# These capabilities are provided as a user convenience.
# The user must recognize that LinearMap objects are not designed
# for efficient element-wise indexing.
#
# 2018-01-19, Jeff Fessler, University of Michigan

using LinearMaps
using Test

# A[end]
function Base.lastindex(A::LinearMap)
    return prod(size(A))
end

# A[?,end] and A[end,?]
function Base.lastindex(A::LinearMap, d::Integer)
    return size(A, d)
end

# A[i,j]
function Base.getindex(A::LinearMap, i::Integer, j::Integer)
	e = zeros(size(A,2)); e[j] = 1
	tmp = A * e
	return tmp[i]
end

# A[k]
function Base.getindex(A::LinearMap, k::Integer)
	c = CartesianIndices(size(A))[k] # is there a more elegant way?
	return A[c[1], c[2]]
#	return A[c[:]] # fails
#	return A[c...] # fails
#	return A[Tuple(c)...] # works
#	return getindex(A, c[1], c[2]) # works
end

# A[:,j]
# it is crucial to provide this function rather than to inherit from
# Base.getindex(A::AbstractArray, ::Colon, ::Integer)
# because Base.getindex does this by iterating (I think).
function Base.getindex(A::LinearMap, ::Colon, j::Integer)
	e = zeros(size(A,2)); e[j] = 1
	return A * e
end

# A[i,:]
function Base.getindex(A::LinearMap, i::Integer, ::Colon)
	# in Julia: A[i,:] = A'[:,i] for real matrix A else need conjugate
	return eltype(A) <: Complex ? conj.(A'[:,i]) : A'[:,i]
end

# A[:,j:k]
# this one is also important for efficiency
function Base.getindex(A::LinearMap, ::Colon, ur::UnitRange)
	return hcat([A[:,j] for j in ur]...)
end

# A[i:k,:]
Base.getindex(A::LinearMap, ur::UnitRange, ::Colon) = A'[:,ur]'

# A[:,:] = Matrix(A)
Base.getindex(A::LinearMap, ::Colon, ::Colon) = Matrix(A)

# A[i:k,j:l]
# this one is inefficient so could be inherited from Base?
Base.getindex(A::LinearMap, r1::UnitRange, r2::UnitRange) = A[:,r2][r1,:]

# A[???]
# informative error message in case we have overlooked any types
#=
function Base.getindex(A::LinearMap, kw...)
	@show kw
	for arg in kw
		@show typeof(arg)
	end
	throw("unsupported indexing type")
end
=#

# tests for getindex for LinearMaps

function ir_LinearMap_test_getindex(A::LinearMap)
	B = Matrix(A)
	@test all(size(A) .>= (4,4)) # required by tests
	@test B[1] == A[1]
	@test B[7] == A[7]
	@test B[:,5] == A[:,5]
	@test B[3,:] == A[3,:]
	@test B[1,3] == A[1,3]
	@test B[:,1:3] == A[:,1:3]
	@test B[1:3,:] == A[1:3,:]
	@test B[1:3,2:4] == A[1:3,2:4]
	@test B == A[:,:]
	@test B'[3] == A'[3]
	@test B'[:,4] == A'[:,4]
	@test B'[2,:] == A'[2,:]
	@test B[end] == A[end] # lastindex
	@test B[3,end-1] == A[3,end-1]

	# The following do not work because currently LinearMap is not a
	# subtype of AbstractMatrix.  If it were such a subtype, then LinearMap
	# would inherit general Base.getindex abilities
#=
	# todo later, if/when LinearMap is a subtype of AbstractMatrix
	@test B[[1, 3, 4]] == A[[1, 3, 4]]
	@test B[:, [1, 3, 4]] == A[:, [1, 3, 4]]
	@test B[[1, 3, 4], :] == A[[1, 3, 4], :]
	@test B[4:7] == A[4:7]
=#

	true
end

# tests for cumsum()
# note: the adjoint of cumsum is reverse(cumsum(reverse(y)))
function ir_lm_test_getindex_cumsum( ; N::Integer=5)
	A = LinearMap(cumsum, y -> reverse(cumsum(reverse(y))), N)
	ir_LinearMap_test_getindex(A)
	@test Matrix(A)' == Matrix(A')
#	@show A[:,4]
	true
end


using Revise
using Oceananigans
using Oceananigans.Units
using Oceananigans.Models.HydrostaticFreeSurfaceModels: ImplicitFreeSurface
using Statistics
using Printf
using LinearAlgebra, SparseArrays
using Oceananigans.Solvers: constructors, unpack_constructors

function geostrophic_adjustment_simulation(free_surface, topology; arch = Oceananigans.CPU())

    Lh = 100kilometers
    Lz = 400meters

    grid = RectilinearGrid(arch,
        size = (64, 3, 1),
        x = (0, Lh), y = (0, Lh), z = (-Lz, 0),
        topology = topology)

    coriolis = FPlane(f = 1e-4)

    model = HydrostaticFreeSurfaceModel(grid = grid,
        coriolis = coriolis,
        free_surface = free_surface)

    gaussian(x, L) = exp(-x^2 / 2L^2)

    U = 0.1 # geostrophic velocity
    L = grid.Lx / 40 # gaussian width
    x₀ = grid.Lx / 4 # gaussian center

    vᵍ(x, y, z) = -U * (x - x₀) / L * gaussian(x - x₀, L)

    g = model.free_surface.gravitational_acceleration
    η = model.free_surface.η

    η₀ = coriolis.f * U * L / g # geostrohpic free surface amplitude

    ηᵍ(x) = η₀ * gaussian(x - x₀, L)

    ηⁱ(x, y) = 2 * ηᵍ(x)

    set!(model, v = vᵍ)
    set!(η, ηⁱ)

    gravity_wave_speed = sqrt(g * grid.Lz) # hydrostatic (shallow water) gravity wave speed
    wave_propagation_time_scale = model.grid.Δxᶜᵃᵃ / gravity_wave_speed
    simulation = Simulation(model, Δt = 2wave_propagation_time_scale, stop_iteration = 300)

    return simulation
end

function run_and_analyze(simulation)
    η = simulation.model.free_surface.η
    u, v, w = simulation.model.velocities
    Δt = simulation.Δt

    ηx = Field(∂x(η))
    compute!(ηx)

    f = simulation.model.free_surface
    @views u₀ = interior(u)[:, 1, 1]
    @views v₀ = interior(v)[:, 1, 1]
    @views η₀ = interior(η)[:, 1, 1]
    @views ηx₀ = interior(ηx)[:, 1, 1]

    if f isa SplitExplicitFreeSurface
        solver_method = "SplitExplicitFreeSurface"
    else
        solver_method = string(simulation.model.free_surface.solver_method)
    end

    simulation.output_writers[:fields] = JLD2OutputWriter(simulation.model, (η, ηx, u, v, w),
        schedule = TimeInterval(Δt),
        prefix = "solution_$(solver_method)",
        force = true)

    progress_message(sim) = @info @sprintf("[%.2f%%], iteration: %d, time: %.3f, max|w|: %.2e",
        100 * sim.model.clock.time / sim.stop_time, sim.model.clock.iteration,
        sim.model.clock.time, maximum(abs, sim.model.velocities.u))

    simulation.callbacks[:progress] = Callback(progress_message, IterationInterval(10))

    run!(simulation)

    compute!(ηx)

    @views u₁ = interior(u)[:, 1, 1]
    @views v₁ = interior(v)[:, 1, 1]
    @views η₁ = interior(η)[:, 1, 1]
    @views ηx₁ = interior(ηx)[:, 1, 1]

    @show mean(Array(η₀))
    @show mean(Array(η₁))

    Δη = Array(η₁) - Array(η₀)

    return (; η₀, η₁, Δη, ηx₀, ηx₁, u₀, u₁, v₀, v₁)
end

# fft_based_free_surface = ImplicitFreeSurface()
pcg_free_surface = ImplicitFreeSurface(solver_method = :PreconditionedConjugateGradient);
matrix_free_surface = ImplicitFreeSurface(solver_method = :HeptadiagonalIterativeSolver);
splitexplicit_free_surface = SplitExplicitFreeSurface()

topology_types = [(Bounded, Periodic, Bounded), (Periodic, Periodic, Bounded)]
topology_types = [topology_types[1]]

archs = [Oceananigans.CPU(), Oceananigans.GPU()]
archs = [archs[1]]

free_surfaces = [pcg_free_surface, matrix_free_surface, splitexplicit_free_surface];
simulations = [geostrophic_adjustment_simulation(free_surface, topology_type, arch = arch) for free_surface in free_surfaces, topology_type in topology_types, arch in archs];
data = [run_and_analyze(sim) for sim in simulations];
# run_and_analyze(simulations[3])

using GLMakie
using JLD2

file1 = jldopen("solution_PreconditionedConjugateGradient.jld2")
file2 = jldopen("solution_HeptadiagonalIterativeSolver.jld2")
file3 = jldopen("solution_SplitExplicitFreeSurface.jld2")

grid = file1["serialized/grid"]

x = grid.xᶜᵃᵃ[1:grid.Nx]
xf = grid.xᶠᵃᵃ[1:grid.Nx+1]
y = grid.yᵃᶜᵃ[1:grid.Ny]

iterations = parse.(Int, keys(file1["timeseries/t"]))
iterations = iterations[1:200]

iter = Node(0) # Node or Observable depending on Makie version
mid = Int(floor(grid.Ny / 2))
η0 = file1["timeseries/1/0"][:, mid, 1]
η1 = @lift(Array(file1["timeseries/1/"*string($iter)])[:, mid, 1])
η2 = @lift(Array(file2["timeseries/1/"*string($iter)])[:, mid, 1])
η3 = @lift(Array(file3["timeseries/1/"*string($iter)])[:, mid, 1])
u1 = @lift(Array(file1["timeseries/3/"*string($iter)])[:, mid, 1])
u2 = @lift(Array(file2["timeseries/3/"*string($iter)])[:, mid, 1])
u3 = @lift(Array(file3["timeseries/3/"*string($iter)])[:, mid, 1])
fig = Figure(resolution = (1400, 1000))
options = (; ylabelsize = 22,
    xlabelsize = 22, xgridstyle = :dash, ygridstyle = :dash, xtickalign = 1,
    xticksize = 10, ytickalign = 1, yticksize = 10, xlabel = "y [m]", xlims = extrema(x))
ax1 = Axis(fig[1, 1]; options..., ylabel = "η [m]", ylims = (-5e-4, 5e-4))
ax2 = Axis(fig[1, 2]; options..., ylabel = "u [m/s]", ylims = (-5e-5, 5e-5))

ηlines0 = lines!(ax1, x, η0, color = :black)
ηlines1 = lines!(ax1, x, η1, color = :red)
ηlines2 = lines!(ax1, x, η2, color = :blue)
ηlines3 = lines!(ax1, x, η3, color = :orange)
axislegend(ax1,
    [ηlines0, ηlines1, ηlines2, ηlines3],
    ["Initial Condition", "PCG", "Matrix", "Split-Explicit"])
ylims!(ax1, (-5e-4, 5e-3))
u0 = Array(file3["timeseries/3/"*string(0)])[:, mid, 1]
xf = length(u0) == length(xf) ? xf : x
ulines1 = lines!(ax2, xf, u1, color = :red)
ulines2 = lines!(ax2, xf, u2, color = :blue)
ulines3 = lines!(ax2, xf, u3, color = :orange)
axislegend(ax2,
    [ulines1, ulines2, ulines3],
    ["PCG", "Matrix", "Split-Explicit"])
ylims!(ax2, (-2e-4, 2e-4))
GLMakie.record(fig, "free_surface_bounded.mp4", iterations, framerate = 12) do i
    @info "Plotting iteration $i of $(iterations[end])..."
    iter[] = i
end


using Fredholm
using Test
using QuadGK, Random
using LinearAlgebra: norm

include("fredholm.jl")


include("plotting.jl")
include("algConnectivity.jl")
include("growOptimalACgraph.jl")
include("utilities.jl")
fldrs = "empirical_results_2016_04_16_08_57_32__1_"
fname = "final"
num_fldrs = 5
empirical_opt_connectivities = zeros(num_fldrs)
random_connectivities = similar(empirical_opt_connectivities)
opt_connectivity=0
n = 0
edg = 0
for i = 1:num_fldrs
    out_fldr = string(fldrs,i)
    f = open(string(joinpath(out_fldr,fname),".dot"),"r")

    g = SimpleGraphs.SimpleGraph()

    for j in enumerate(eachline(f))
        if j[1] == 1 || eof(f)
            continue
        else
            inew = [parse(Int,x) for x in split(j[2],"--")]
            if length(inew) == 1
                add!(g,inew[1])
            else
                add!(g,inew[1],inew[2])
            end
        end
    end
    if i == 1
        n = NV(g)
        edg = NE(g)
        opt_connectivity = growOptimalA(n,edg)[2]
    end
    close(f)
    empirical_opt_connectivities[i] = AlgConnectivity(g,false)
end

# create random graphs and store connectivites
for i = 1:num_fldrs
    g = simple_NE_sgraph(n,edg)
    random_connectivities[i] = AlgConnectivity(g,false)
end
vals = cat(2,empirical_opt_connectivities,random_connectivities)
plot_c_box(vals,n,opt_connectivity,"","test")


using Parameters, StaticArrays, BenchmarkTools
import TrajectoryOptimization: dynamics, discrete_dynamics, RK3, AbstractModel, KnotPoint,
    jacobian, discrete_jacobian

@with_kw mutable struct BicycleCar{T} <: AbstractModel
    k::T = 0.1
    a::T = 1.  # dist to front axle
    b::T = 1.  # dist to rear axle
    h::T = 1.
    mass::T = 1000.  # mass
    Iz::T = 10.
    Cf::T = 1.
    Cr::T = 1.
    μ::T = 0.1
    Cd0::T = 0.1
    Cd1_Npmps::T = 0.1
    g::T = 9.81
end

Base.size(::BicycleCar) = 8,2

function dynamics(car::BicycleCar, x, u)
    δ_dot  = u[1]
    fx_dot = u[2]
    δ  = x[1]
    fx = x[2]
    r  = x[3]
    Uy = x[4]
    Ux = x[5]
    Δψ = x[6]
    e  = x[7]
    t  = x[8]

    # Drag Force
    Fx_drag = 0.0

    # Spatial derivative
    s_dot = (Ux * cos(Δψ) - Uy * sin(Δψ)) / (1 - car.k*e)

    # Get tire forces
    Fxf, Fxr, Fyf, Fyr = logit_lateral_force_model(car, x)

    # State Derivatives
    r_dot  = (car.a * (Fxf * cos(δ) + Fxf*sin(δ)) - car.b * Fyr) / car.Iz
    Ux_dot =  r * Uy + (Fxf * cos(δ) - Fyf * sin(δ) + Fxr - Fx_drag) / car.mass
    Uy_dot = -r * Ux + (Fyf * cos(δ) + Fxf * sin(δ)) / car.mass
    Δψ_dot = r - car.k * s_dot
    e_dot = Ux * sin(Δψ) + Uy * cos(Δψ)
    t_dot = 1 / s_dot

    return @SVector [δ_dot, fx_dot, r_dot, Uy_dot, Ux_dot, Δψ_dot, e_dot, t_dot]
end

function FWD_force_model(car::BicycleCar, fx)
    l = car.a + car.b
    Fxf = max(fx * car.b / l, fx)
    Fxr = min(fx * car.a / l, 0.0)
    Fzf = car.mass * car.g * car.b / l
    Fzr = car.mass * car.g * car.a / l
    return Fxf, Fxr, Fzf, Fzr
end

function logit_lateral_force_model(car::BicycleCar, x)
    fx = x[2]
    r  = x[3]
    Uy = x[4]
    Ux = x[5]
    Cα = 1.0

    α_f = atan(Uy + car.a * r, Ux)
    α_r = atan(Uy - car.b * r, Ux)

    Fxf, Fxr, Fzf, Fzr = FWD_force_model(car, fx)
    Fyf = lateral_tire_model(car.μ, Cα, α_f, Fxf, Fzf)
    Fyr = lateral_tire_model(car.μ, Cα, α_r, Fxr, Fzr)
    return Fxf, Fxr, Fyf, Fyr
end

function lateral_tire_model(μ, Cα, α, Fx, Fz)
    Fy_max = sqrt((μ*Fz)^2 - Fx^2)
    slope = 2*Cα / Fy_max
    Fy = Fy_max * (1-2*exp(slope*α)) / (1 + exp(slope*α))
    return Fy
end

struct ContingencyCar{T} <: AbstractModel
    cars::Vector{BicycleCar{T}}
    ix::Vector{SVector{8,Int}}
    iu::Vector{SVector{2,Int}}
end

function ContingencyCar(num_cars=2)
    cars = [BicycleCar() for i = 1:num_cars]
    n,m = 8,2
    ix = [SVector{n}((1:n) .+ (i-1)*n) for i = 1:num_cars]
    iu = [SVector{m}((1:m) .+ (i-1)*m) for i = 1:num_cars]
    ContingencyCar(cars, ix, iu)
end

Base.size(::ContingencyCar) = 16,4

function dynamics(car::ContingencyCar, x, u)
    ix, iu = car.ix, car.iu
    cars = car.cars
    x1,u1 = x[ix[1]], u[iu[1]]
    ẋ1 = dynamics(cars[1], x1, u1)
    x2,u2 = x[ix[2]], u[iu[2]]
    ẋ2 = dynamics(cars[2], x1, u1)
    return [ẋ1; ẋ2]
end

car = ContingencyCar()
n,m = size(car)
x = @SVector rand(n)
u = @SVector rand(m)
z = KnotPoint(x,u,0.1)

@btime dynamics($car, $x, $u)

@btime discrete_dynamics(RK3, $car, $x, $u, 0.1)

@btime discrete_jacobian(RK3, $car, $z)


using UnionFind

# edges must be pre-sorted according to weight.
function kruskal{T <: Integer}(nodes :: T, edges :: Array{(T, T)})
    uf = UnionFinder(nodes)
    mst = Array{(T, T)}

    for i in 1:length(edges)
        (u, v) = edges[i]
        if find!(uf, u) != find!(uf, v)
            union!(uf, u, v)
            push!(mst, (u, v))
        end
    end
end


include("utils.jl")

binary = false

binarize_data(x) = round(x) <= 0 ? -1 : 1

binarize_state(x) = x <= 0 ? -1 : 1

##

mutable struct RBM

    visibles::Array{Float32}
    hiddens::Array{Float32}

    weights::Array{Float32}

    RBM(in_size, hidden_size) = new(

        [0 for _ in 1:in_size],
        [0 for _ in 1:hidden_size],

        randn(in_size, hidden_size),

    )

end

(rbm::RBM)(input) =
begin

    rbm.hiddens = tanh.(input * rbm.weights)

    binary ? rbm.hiddens = binarize_state.(rbm.hiddens) : ()

end

(rbm::RBM)() =
begin

    rbm.visibles = tanh.(rbm.hiddens * transpose(rbm.weights))

    binary ? rbm.visibles = binarize_state.(rbm.visibles) : ()

end

##

energy(rbm) = -(rbm.visibles * rbm.weights * rbm.hiddens')

propogate_until_convergence!(rbm, input; k=1_000) =
begin

    prev_visibles = nothing
    prev_hiddens = nothing

    ctr = 0

    while (rbm.visibles != prev_visibles || rbm.hiddens != prev_hiddens) && ctr < k

        prev_visibles = rbm.visibles
        prev_hiddens = rbm.hiddens

        ctr +=1

        rbm(input)
        rbm()

        input = rbm.visibles

    end

input
end

##

alternating_gibbs_grads!(rbm, input; converge=false, k=1) =
begin

    rbm(input)

    pos_hiddens = rbm.hiddens

    if converge

        rbm()

        propogate_until_convergence!(rbm, rbm.visibles)

    else

        for _ in 1:k

            rbm()

            rbm(rbm.visibles)

        end

    end

    grads = (transpose(input) * pos_hiddens) .- (transpose(rbm.visibles) * rbm.hiddens)

grads
end

batch_grads(rbm, batch) = threadpool(sum, args->alternating_gibbs_grads!(args...), [[deepcopy(rbm), input] for input in batch]) ./ length(batch)

update_weights!(rbm, grads, lr) =
begin

    rbm.weights += lr .* grads

end

##


@testset "fma" begin

  @test HILO(fma(Float64(pi), Float64(phi), Double64(gamma))) == (5.660419357216792, 4.1344240638440936e-16)
  @test HILO(fma(Float64(pi), Double64(phi), Float64(gamma))) == (5.660419357216792, 2.477303893634162e-16)
  @test HILO(fma(Double64(pi), Float64(phi), Float64(gamma))) == (5.660419357216793, -2.7164108363986695e-16)
  @test HILO(fma(Double64(pi), Double64(phi), Float64(gamma))) == (5.660419357216793, -4.422960158132908e-16)
  @test HILO(fma(Double64(pi), Float64(phi), Double64(gamma))) == (5.660419357216793, -2.765839987922976e-16)
  @test HILO(fma(Double64(pi), Double64(phi), Double64(gamma))) == (5.660419357216792, 4.4093948873440384e-16)

  @test HILO(muladd(Float64(pi), Float64(phi), Double64(gamma))) == (5.660419357216792, 4.1344240638440936e-16)
  @test HILO(muladd(Float64(pi), Double64(phi), Float64(gamma))) == (5.660419357216792, 2.477303893634162e-16)
  @test HILO(muladd(Double64(pi), Float64(phi), Float64(gamma))) == (5.660419357216793, -2.7164108363986695e-16)
  @test HILO(muladd(Double64(pi), Double64(phi), Float64(gamma))) == (5.660419357216793, -4.422960158132908e-16)
  @test HILO(muladd(Double64(pi), Float64(phi), Double64(gamma))) == (5.660419357216793, -2.765839987922976e-16)
  @test HILO(muladd(Double64(pi), Double64(phi), Double64(gamma))) == (5.660419357216792, 4.4093948873440384e-16)

end


module EL

import ..ED
using ..SimLib
using ..SimLib: Maybe
using LinearAlgebra
using SharedArrays: sdata

export EigenstateLocality, ELDataDescriptor, ELData, load_el

### Descriptor

struct ELDataDescriptor{T} <: ED.EDDerivedDataDescriptor
    operator::Maybe{T}
    operatorname::String
    derivedfrom::ED.EDDataDescriptor
end

ELDataDescriptor(operatorname::String, args...; kwargs...) = ELDataDescriptor(nothing, operatorname, args...; kwargs...)
ELDataDescriptor(operator, operatorname::String, args...; kwargs...) = ELDataDescriptor(operator, operatorname, EDDataDescriptor(args...; kwargs...))

### Data obj

struct ELData{T, N} <: SimLib.AbstractSimpleData
    descriptor::ELDataDescriptor{T}
    data::Array{Float64, N}
end

ED._default_folder(::ELDataDescriptor) = "locality"
ED._filename_addition(eldd::ELDataDescriptor) = "_" * eldd.operatorname

"""
    load_el(operatorname, edd)
    load_el(operatorname, model, diagtype[, location])
"""
load_el(args...; kwargs...) = load(ELDataDescriptor(args...; kwargs...))

## functions

eigenstatelocality(eigen, operator) = eigenstatelocality!(Vector{Float64}(undef, length(eigen.values)-1), eigen, operator)

eigenstatelocality!(res, eigen, operator) = eigenstatelocality!(res, eigen.values, eigen.vectors, operator)

function eigenstatelocality!(res, evals, evecs, operator)
    energies = [eva + real(dot(eve, operator, eve)) for (eva, eve) in zip(evals, eachcol(evecs))]
    order = sortperm(energies)
    for (i, (index, nextindex)) in enumerate(zip(order[1:end-1], order[2:end]))
        res[i] = log(abs(dot(view(evecs, :, index), operator, view(evecs, :, nextindex))) / (energies[nextindex] - energies[index]))
    end
    res
end

### Task

mutable struct ELTask{T} <: ED.EDTask
    operatorname::String
    operator::T
    data
end

EigenstateLocality(operatorname, operator) = ELTask(operatorname, operator, nothing)

function ED.initialize!(task::ELTask, arrayconstructor, spectral_size)
    task.data = arrayconstructor(Float64, spectral_size-1)
end

function ED.compute_task!(task::ELTask, evals, evecs, inds...)
    n = min(size(task.data,1), length(evals)-1)
    eigenstatelocality!(view(task.data, 1:n, inds...), view(evals, 1:n+1), evecs, task.operator)
end

function ED.failed_task!(task::ELTask, inds...)
    task.data[:, inds...] .= NaN64
end

function ED.assemble(task::ELTask, edd)
    ELData(ELDataDescriptor(task.operator, task.operatorname, edd), sdata(task.data))
end

Base.summary(task::ELTask) = string(typeof(task)) * "($(task.operatorname))"

end #module


"""
Abstract continuous filter.
"""
abstract type Filter end

"""
    IdentityFilter()

Identity filter, which does not filter.
"""
struct IdentityFilter <: Filter end

"""
    TopHatFilter(width)

Top hat filter, parameterized by a variable filter width.
"""
struct TopHatFilter <: Filter
    width::Function
end

"""
    ConvolutionalFilter(kernel)

Convolutional filter, parameterized by a filter kernel.
"""
struct ConvolutionalFilter <: Filter
    width::Function
    kernel::Function
end

"""
GaussianFilter(h, σ) -> ConvolutionalFilter
GaussianFilter(σ) -> ConvolutionalFilter

Create Gaussian ConvolutionalFilter with domain width `2h` and variance `σ^2`.
"""
GaussianFilter(h, σ) = ConvolutionalFilter(h, x -> 1 / √(2π * σ^2) * exp(-x^2 / 2σ^2))
GaussianFilter(σ) = GaussianFilter(x -> 10σ, σ)


# In theory, for load-based MMF routing, using any monotone function of the load
# should give the same result. We check that fact. We proceed as follows:
# for each topology, for several objective functions (including the simple
# load), compute a routing; compare all those routings.

# For simplicity, we will focus on load functions that can be solved with
# a standard SOCP solver (like CPLEX, Gurobi, Mosek).

# Theoretical reference:
# https://ieeexplore.ieee.org/document/4738463 : Section 4.1.2.

include("topologies.jl")

using PrettyTables, LightGraphs, DelimitedFiles
using JuMP, CPLEX, Gurobi, Mosek, MosekTools

# opt = CPLEX.Optimizer
opt = optimizer_with_attributes(Gurobi.Optimizer, "NonConvex" => 2)
# opt = Mosek.Optimizer

ftopologies = [t for t in topologies if ! (t.name in ["att", "abilene"])]
ε = 1.0e-8

objs = [Load(), KleinrockLoad(true), FortzThorupLoad()]#, AlphaFairness(0.5)]#, AlphaFairness(1.0)]#, AlphaFairness(1.5), AlphaFairness(2.0)]
table = Matrix{Any}(undef, 2 * length(ftopologies), 2 + length(objs))
csv_cols = ["Instance", "Number of edges", "Number of nodes", "Number of demands",
            "Objective function for MMF", "MMF iterations", "Time to create the master problem [ms]",
            "Computation time (per iteration) [ms]", "Maximum iteration time [ms]", "Total iteration time [ms]",
            "Total time [ms]", "Maximum load of the obtained routing", "MMF status code"]
csv = Matrix{String}(undef, 1 + length(ftopologies) * length(objs), length(csv_cols))
csv[1, :] = csv_cols
csv_row = 2

for (i, t) in enumerate(ftopologies)
    println("=> $(t.name)")
    g, k, dm = topology_to_graphs(t)

    loads = Any[]
    n_paths = Any[]

    for obj in objs
        println("  -> $(obj)")

        if typeof(obj) != AlphaFairness
            mt = ModelType(obj, MinMaxFair(ε))
        else
            mt = ModelType(obj, MaxMinFair(ε))
        end

        rd = RoutingData(g, k, opt, mt, name=t.name, traffic_matrix=dm)
        sol = compute_routing(rd)

        r = sol.routings[sol.n_iter]
        max_load = compute_max_load(r)
        push!(loads, max_load)
        push!(n_paths, length(r.paths))

        global csv_row
        csv[csv_row, :] = String[t.name, string(ne(g)), string(nv(g)), string(ne(k)),
                                 string(obj), string(sol.n_iter), string(sol.time_create_master_model_ms),
                                 string(sol.time_solve_ms), string(maximum(values(sol.time_solve_ms))), string(sol.total_time_solve_ms),
                                 string(sol.total_time_ms), string(max_load), string(sol.result)]
        csv_row += 1
    end

    table[2 * i - 1, :] = vcat([t.name, "Max load"], loads)
    table[2 * i, :] = vcat(["", "Nb. paths"], n_paths)
end

@ptconfclean
@ptconf tf = unicode_rounded
@pt :header = ["", "", objs...] table

# println(csv)
writedlm("$(@__DIR__)/mmf_compare_objectives.csv", csv, ',')


abstract type SemiMartingale{dType, tType, N} end

#Type for holding the increments of a Wiener process:
# conv - mode of convergence: "strong" or "weak"
# nd   - dimension of the Wiener process
# nt   - number of increments in the DataSeries
# ns   - number of sample paths of the Wiener process
# Δt   - time increment of the TimeSeries
# ΔW   - variable storing the increments of the Wiener process over Δt, or the discrete random variable \\hat{I}
# ΔZ   - variable holding the time integral of the Wiener process \\int_{tk}^{tk+1} (W(t)-W(tk))dt, or the discrete random variable \\tilde{I}
# dType- type of the elements of the increments of the Wiener process
# tType- type of the time steps
# N    - the number of dimensions of the arrays holding data in ΔW and ΔZ
struct WienerProcess{dType, tType, N} <: SemiMartingale{dType, tType, N}
    conv::String
    nd::Int
    nt::Int
    ns::Int
    Δt::tType
    ΔW::SStochasticDataSeries{dType,N}
    ΔZ::SStochasticDataSeries{dType,N}

    function WienerProcess{dType, tType, N}(nd, nt, ns, Δt, conv="strong") where {dType, tType, N}

        @assert conv=="strong" || conv=="weak"
        @assert dType <: Number
        @assert tType <: Number
        @assert nd > 0
        @assert nt ≥ 1
        @assert ns > 0
        @assert N ∈ (1,2,3)

        if conv=="strong"
            if N == 1
                chi = randn(nt)
                eta = randn(nt)
            elseif N == 2
                chi = randn(nd, nt)
                eta = randn(nd, nt)
            elseif N == 3
                chi = randn(nd, nt, ns)
                eta = randn(nd, nt, ns)
            end

            dW = chi*√Δt
            dZ = Δt^(3/2)/2 * (chi+eta/√3)
        else
            if N == 1
                chi = rand(dType,nt)
                eta = rand(dType,nt)
                dW  = zeros(dType,nt)
                dZ  = sqrt(Δt)*ones(dType,nt)
            elseif N == 2
                chi = rand(dType, nd, nt)
                eta = rand(dType, nd, nt)
                dW  = zeros(dType, nd, nt)
                dZ  = sqrt(Δt)*ones(dType, nd, nt)
            elseif N == 3
                chi = rand(dType, nd, nt, ns)
                eta = rand(dType, nd, nt, ns)
                dW  = zeros(dType, nd, nt, ns)
                dZ  = sqrt(Δt)*ones(dType, nd, nt, ns)
            end

            # Generating the random variable \hat I
            # P(\\hat I=+-sqrt(3*Δt)) = 1/6
            # P(\\hat I=0) = 2/3
            indx     = findall(x->x<1. / 6., chi)
            dW[indx] .= -sqrt(3*Δt)
            indx     = findall(x->x>5. / 6., chi)
            dW[indx] .= sqrt(3*Δt)

            # Generating the random variable \tilde I
            # P(\tilde I=+-sqrt(Δt)) = 1/2
            indx     = findall(x->x<0.5, eta)
            dZ[indx] .= -sqrt(Δt)
        end

        # Creating ΔW, ΔZ
        # nt-1, because SStochasticDataSeries adds one time step
        # ni=1, because we cosider a single IC for the Wiener process
        ΔW = SStochasticDataSeries{dType,N}(nd, nt-1, ns, 1, dW)
        ΔZ = SStochasticDataSeries{dType,N}(nd, nt-1, ns, 1, dZ)

        new(conv, nd, nt, ns, Δt, ΔW, ΔZ)
    end

    function WienerProcess{dType, tType, N}(nd, nt, ns, Δt, dW, dZ, conv="strong") where {dType, tType, N}

        @assert conv=="strong" || conv=="weak"

        new(conv, nd, nt, ns, Δt, SStochasticDataSeries(dW), SStochasticDataSeries(dZ))
    end

end

function WienerProcess(dType, nd::Int, nt::Int, ns::Int, Δt::tType, conv::String="strong") where {tType}

    if nd==ns==1
        N=1
    elseif ns==1
        N=2
    else
        N=3
    end

    WienerProcess{dType, tType, N}(nd, nt, ns, Δt, conv)
end

function WienerProcess(Δt::tType, dW::Array{dType,1}, dZ::Array{dType,1}, conv::String="strong") where {dType, tType}

    @assert size(dW,1)==size(dZ,1)

    nd = 1
    nt = size(dW,1)
    ns = 1

    return WienerProcess{dType, tType, 1}(nd, nt, ns, Δt, dW, dZ, conv)
end

function WienerProcess(Δt::tType, dW::Array{dType,2}, dZ::Array{dType,2}, conv::String="strong") where {dType, tType}

    @assert size(dW,1)==size(dZ,1)
    @assert size(dW,2)==size(dZ,2)

    nd = size(dW,1)
    nt = size(dW,2)
    ns = 1

    return WienerProcess{dType, tType, 2}(nd, nt, ns, Δt, dW, dZ, conv)
end

function WienerProcess(Δt::tType, dW::Array{dType,3}, dZ::Array{dType,3}, conv::String="strong") where {dType, tType}

    @assert size(dW,1)==size(dZ,1)
    @assert size(dW,2)==size(dZ,2)
    @assert size(dW,3)==size(dZ,3)

    nd = size(dW,1)
    nt = size(dW,2)
    ns = size(dW,3)

    return WienerProcess{dType, tType, 3}(nd, nt, ns, Δt, dW, dZ, conv)
end

# Generates a new series of increments for the Wiener process W
function generate_wienerprocess!(W::WienerProcess{dType, tType, N}) where {dType, tType, N}

    if W.conv=="strong"
        if N == 1
            chi = randn(W.nt)
            eta = randn(W.nt)
        elseif N == 2
            chi = randn(W.nd, W.nt)
            eta = randn(W.nd, W.nt)
        elseif N == 3
            chi = randn(W.nd, W.nt, W.ns)
            eta = randn(W.nd, W.nt, W.ns)
        end

        dW = chi*√W.Δt
        dZ = W.Δt^(3/2)/2 * (chi+eta/√3)
    else
        if N == 1
            chi = rand(dType,W.nt)
            eta = rand(dType,W.nt)
            dW  = zeros(dType,W.nt)
            dZ  = sqrt(W.Δt)*ones(dType,W.nt)
        elseif N == 2
            chi = rand(dType, W.nd, W.nt)
            eta = rand(dType, W.nd, W.nt)
            dW  = zeros(dType, W.nd, W.nt)
            dZ  = sqrt(W.Δt)*ones(dType, W.nd, W.nt)
        elseif N == 3
            chi = rand(dType, W.nd, W.nt, W.ns)
            eta = rand(dType, W.nd, W.nt, W.ns)
            dW  = zeros(dType, W.nd, W.nt, W.ns)
            dZ  = sqrt(W.Δt)*ones(dType, W.nd, W.nt, W.ns)
        end

        # Generating the random variable \\hat I
        # P(\\hat I=+-sqrt(3*Δt)) = 1/6
        # P(\\hat I=0) = 2/3
        indx     = findall(x->x<1. / 6., chi)
        dW[indx] .= -sqrt(3*W.Δt)
        indx     = findall(x->x>5. / 6., chi)
        dW[indx] .= sqrt(3*W.Δt)

        # Generating the random variable \\tilde I
        # P(\\tilde I=+-sqrt(Δt)) = 1/2
        indx     = findall(x->x<0.5, eta)
        dZ[indx] .= -sqrt(W.Δt)
    end

    W.ΔW.d.=dW
    W.ΔZ.d.=dZ

    return nothing
end

Base.ndims(ds::WienerProcess{DT,TT,N}) where {DT,TT,N} = N


using DiffEqFlux, GalacticOptim, OrdinaryDiffEq, RecursiveArrayTools

u0 = Float32[0.; 2.]
du0 = Float32[0.; 0.]
tspan = (0.0f0, 1.0f0)
t = range(tspan[1], tspan[2], length=20)

model = FastChain(FastDense(2, 50, tanh), FastDense(50, 2))
p = initial_params(model)
ff(du,u,p,t) = model(u,p)
prob = SecondOrderODEProblem{false}(ff, du0, u0, tspan, p)

function predict(p)
    Array(solve(prob, Tsit5(), p=p, saveat=t, sensealg = InterpolatingAdjoint(autojacvec=ZygoteVJP())))
end

correct_pos = Float32.(transpose(hcat(collect(0:0.05:1)[2:end], collect(2:-0.05:1)[2:end])))

function loss_n_ode(p)
    pred = predict(p)
    sum(abs2, correct_pos .- pred[1:2, :]), pred
end

data = Iterators.repeated((), 1000)
opt = ADAM(0.01)

l1 = loss_n_ode(p)

cb = function (p,l,pred)
    @show l
    l < 0.01 && Flux.stop()
end

optfunc = GalacticOptim.OptimizationFunction((x, p) -> loss_n_ode(x), GalacticOptim.AutoZygote())
optprob = GalacticOptim.OptimizationProblem(optfunc, p)
res = GalacticOptim.solve(optprob, opt, cb=cb, maxiters = 100)
l2 = loss_n_ode(res.minimizer)
@test l2 < l1

function predict(p)
    Array(solve(prob, Tsit5(), p=p, saveat=t, sensealg = QuadratureAdjoint(autojacvec=ZygoteVJP())))
end

correct_pos = Float32.(transpose(hcat(collect(0:0.05:1)[2:end], collect(2:-0.05:1)[2:end])))

function loss_n_ode(p)
    pred = predict(p)
    sum(abs2, correct_pos .- pred[1:2, :]), pred
end

data = Iterators.repeated((), 1000)
opt = ADAM(0.01)

loss_n_ode(p)

cb = function (p,l,pred)
    @show l
    l < 0.01 && Flux.stop()
end
optfunc = GalacticOptim.OptimizationFunction((x, p) -> loss_n_ode(x), GalacticOptim.AutoZygote())
optprob = GalacticOptim.OptimizationProblem(optfunc, p)
res = GalacticOptim.solve(optprob, opt, cb=cb, maxiters = 100)
l2 = loss_n_ode(res.minimizer)
@test l2 < l1

function predict(p)
    Array(solve(prob, Tsit5(), p=p, saveat=t, sensealg = BacksolveAdjoint(autojacvec=ZygoteVJP())))
end

correct_pos = Float32.(transpose(hcat(collect(0:0.05:1)[2:end], collect(2:-0.05:1)[2:end])))

function loss_n_ode(p)
    pred = predict(p)
    sum(abs2, correct_pos .- pred[1:2, :]), pred
end

data = Iterators.repeated((), 1000)
opt = ADAM(0.01)

loss_n_ode(p)

cb = function (p,l,pred)
    @show l
    l < 0.01 && Flux.stop()
end

optfunc = GalacticOptim.OptimizationFunction((x, p) -> loss_n_ode(x), GalacticOptim.AutoZygote())
optprob = GalacticOptim.OptimizationProblem(optfunc, p)
res = GalacticOptim.solve(optprob, opt, cb=cb, maxiters = 100)
l2 = loss_n_ode(res.minimizer)
@test l2 < l1


using Base.Broadcast: broadcastable, Broadcasted
using Swizzles.ValArrays
using Swizzles.ArrayifiedArrays
using Swizzles

foo(x) = x

@testset "ValArrays" begin
    A0 = ValArray(0)
    A1 = ValArray(1)

    @test A1.+A1 == 2

    @test A1.*A1 == 1

    B = rand(3,3)
    B .*= A0
    @test B == zeros(3,3)

    B = rand(3,3)
    B .= A0
    @test B == zeros(3,3)

    @test lift_vals(1) == A1
    @test typeof(lift_vals(1)) == typeof(A1)

    B = reshape(1:27, 3, 3, 3)
    C = reshape(1:9, 3, 3)

    @test Sum(2, 3).(B .* (A1 .+ C .+ A0) .+ A1) == Sum(2, 3).(B .* (1 .+ C .+ 0) .+ 1)
    @test typeof(Sum(2, 3).(B .* (A1 .+ C .+ A0))) == typeof(Sum(2, 3).(B .* (1 .+ C .+ 0)))

    @test typeof(Delay().(Sum(2, 3).(A0, B .* (A1 .+ C .+ A0) .+ A1))) ==
        typeof(lift_vals(Delay().(Sum(2, 3).(B .* (1 .+ C .+ 0) .+ 1))))

end