mpow2(a::AbstractArray) = a .^ 2

"""
    svd_back(U, S, V, dU, dS, dV)

adjoint for SVD decomposition.

References:
    https://j-towns.github.io/papers/svd-derivative.pdf
    https://giggleliu.github.io/2019/04/02/einsumbp.html
"""
function svd_back(U::AbstractArray, S::AbstractArray{T}, V, dU, dS, dV; η::Real=1e-40) where T
    all(x -> x isa Nothing, (dU, dS, dV)) && return nothing
    η = T(η)
    NS = length(S)
    S2 = mpow2(S)
    Sinv = @. S/(S2+η)
    F = S2' .- S2
    F ./= (mpow2(F) .+ η)

    res = ZeroAdder()
    if !(dU isa Nothing)
        UdU = U'*dU
        J = F.*(UdU)
        res += (J+J')*Diagonal(S) + Diagonal(1im*imag(diag(UdU)) .* Sinv)
    end
    if !(dV isa Nothing)
        VdV = V'*dV
        K = F.*(VdV)
        res += Diagonal(S) * (K+K')
    end
    if !(dS isa Nothing)
        res += Diagonal(dS)
    end

    res = U*res*V'

    if !(dU isa Nothing) && size(U, 1) != size(U, 2)
        res += (dU - U* (U'*dU)) * Diagonal(Sinv) * V'
    end

    if !(dV isa Nothing) && size(V, 1) != size(V, 2)
        res = res + U * Diagonal(Sinv) * (dV' - (dV'*V)*V')
    end
    res
end


using BrainWave
#using PlotlyJS

sf = 512
dur = 1.2
nSamp = round(Int, sf*dur)
tArr = collect(0:nSamp-1)./sf
dat = zeros(2, nSamp)
dat[1,:] = sin.(2*pi*75 .* tArr)
dat[2,:] = sin.(2*pi*120 .* tArr)

p1, f1 = getSpectrum(dat[1,:], sf, powerOfTwo=false)
p1b, f1b = getSpectrum(dat[1,:], sf, powerOfTwo=true)

p2, f2 = getSpectrum(dat[2,:], sf, powerOfTwo=false)
p2b, f2b = getSpectrum(dat[2,:], sf, powerOfTwo=true)

# s1 = scatter(;x=f1, y=10*log10.(p1))
# s1b = scatter(;x=f1b, y=10*log10.(p1b))

# s2 = scatter(;x=f2, y=10*log10.(p2))
# s2b = scatter(;x=f2b, y=10*log10.(p2b))


using HarwellRutherfordBoeing
using Krylov
using LinearOperators
# using ProfileView

# M = HarwellBoeingMatrix("data/illc1033.rra");
M = HarwellBoeingMatrix("data/illc1850.rra");
A = M.matrix;
(m, n) = size(A);
@printf("System size: %d rows and %d columns\n", m, n);

# Define a linear operator with preallocation.
Ap = zeros(m);
Atq = zeros(n);
op = LinearOperator(m, n, false, false,
                    p -> A_mul_B!(1.0,  A, p, 0.0, Ap),
                    q -> At_mul_B!(1.0, A, q, 0.0, Atq),
                    q -> At_mul_B!(1.0, A, q, 0.0, Atq));
λ = 1.0e-3;
λ > 0.0 && (N = 1./λ * opEye(n))

for nrhs = 1 : size(M.rhs, 2)
  b = M.rhs[:,nrhs];
  (x, stats) = lsqr(op, b, λ=λ, sqd=λ > 0, atol=0.0, btol=0.0, N=N);
  # @profile (x, stats) = lsqr(op, b, λ=λ, sqd=λ > 0, atol=0.0, btol=0.0, N=N);
  @time (x, stats) = lsqr(op, b, λ=λ, sqd=λ > 0, atol=0.0, btol=0.0, N=N);
  show(stats);
  resid = norm(A' * (A * x - b) + λ * x) / norm(b);
  @printf("LSQR: Relative residual: %8.1e\n", resid);
  @printf("LSQR: ‖x‖: %8.1e\n", norm(x));
end

# ProfileView.view()


#
# qubit
#

function ITensors.space(::SiteType"qubit";
                        conserve_qns::Bool = false)
  if conserve_qns
    return [QN() => 2]
  end
  return 2
end

ITensors.state(::SiteType"qubit",
               ::ITensors.StateName"0") = 1

ITensors.state(::SiteType"qubit",
               ::ITensors.StateName"1") = 2

ITensors.op(::OpName"Id",
            ::SiteType"qubit",
            s::Index) =
  itensor([1 0; 0 1], s', dag(s))

ITensors.op(::OpName"X",
            ::SiteType"qubit",
            s::Index) =
  itensor([0.0 1.0; 1.0 0.0], s', dag(s))

ITensors.op(::OpName"iY",
            ::SiteType"qubit",
            s::Index) =
  itensor([0 1; -1 0], s', dag(s))

ITensors.op(::OpName"Y",
            ::SiteType"qubit",
            s::Index) =
  itensor([0 -im; im 0], s', dag(s))

ITensors.op(::OpName"Z",
            ::SiteType"qubit",
            s::Index) =
  itensor([1 0; 0 -1], s', dag(s))

ITensors.op(::OpName"H",
            ::SiteType"qubit",
            s::Index) =
  itensor([1/sqrt(2) 1/sqrt(2);
           1/sqrt(2) -1/sqrt(2)], s', dag(s))

ITensors.op(::OpName"Rx",
            ::SiteType"qubit",
            s::Index; θ::Number) =
  itensor([cos(θ/2)      -im*sin(θ/2.);
           -im*sin(θ/2.) cos(θ/2.)], s', dag(s))

ITensors.op(::OpName"Sw",
            ::SiteType"qubit",
            s1::Index,
            s2::Index) =
  itensor([1 0 0 0;
           0 0 1 0;
           0 1 0 0;
           0 0 0 1],
          s1', s2',
          dag(s1), dag(s2))

ITensors.op(::OpName"rand",
            ::SiteType"qubit",
            s::Index...) =
  randomITensor(prime.(s)...,
                dag.(s)...)

ITensors.op(::OpName"Cx",
            ::SiteType"qubit",
            s1::Index,
            s2::Index) =
  itensor([1 0 0 0;
           0 1 0 0;
           0 0 0 1;
           0 0 1 0],
          s1', s2',
          dag(s1), dag(s2))

ITensors.op(::OpName"T",
            ::SiteType"qubit",
            s1::Index,
            s2::Index,
            s3::Index) =
  itensor([1 0 0 0 0 0 0 0;
           0 1 0 0 0 0 0 0;
           0 0 1 0 0 0 0 0;
           0 0 0 1 0 0 0 0;
           0 0 0 0 1 0 0 0;
           0 0 0 0 0 1 0 0;
           0 0 0 0 0 0 0 1;
           0 0 0 0 0 0 1 0],
          s1', s2', s3',
          dag(s1), dag(s2), dag(s3))

ITensors.op(::OpName"noise",
            ::SiteType"qubit",
            s::Index...;
            krausind = hasqns(s[1]) ? Index([QN() => 2], "kraus") : Index(2, "kraus")) =
  randomITensor(prime.(s)..., dag.(s)..., krausind)


module Day01
export part1, part2

fuel1(m) = m÷3 - 2

function fuel2(m)
    acc = fuel1(m)
    f = acc
    while true
        f = fuel1(f)
        f <= 0 && return acc

        acc += f
    end
end

function main(input, fuel)
    open(input) do f
        acc = 0
        for line in readlines(f)
            m = parse(Int, line)
            acc += fuel(m)
        end
        acc
    end
end

part1(input) = main(input, fuel1)
part2(input) = main(input, fuel2)

end # module


immutable Laplace <: ContinuousUnivariateDistribution
    location::Float64
    scale::Float64
    function Laplace(l::Real, s::Real)
        s > zero(s) || error("scale must be positive")
        new(float64(l), float64(s))
    end
end

Laplace(location::Real) = Laplace(location, 1.0)
Laplace() = Laplace(0.0, 1.0)

typealias Biexponential Laplace

## Support
@distr_support Laplace -Inf Inf

## Properties
mean(d::Laplace) = d.location
median(d::Laplace) = d.location
mode(d::Laplace) = d.location

std(d::Laplace) = sqrt2 * d.scale
var(d::Laplace) = 2.0 * d.scale^2
skewness(d::Laplace) = 0.0
kurtosis(d::Laplace) = 3.0

entropy(d::Laplace) = log(2.0 * d.scale) + 1.0

## Functions
pdf(d::Laplace, x::Float64) = 0.5exp(-abs(x - d.location)/d.scale) / d.scale
logpdf(d::Laplace, x::Float64) = -log(2.0 * d.scale) - abs(x - d.location) / d.scale

cdf(d::Laplace, x::Float64) = x < d.location ? 0.5*exp((x-d.location)/d.scale) : 0.5*(2.0-exp((d.location-x)/d.scale))
ccdf(d::Laplace, x::Float64) = x < d.location ? 0.5-0.5*expm1((x-d.location)/d.scale) : 0.5*exp((d.location-x)/d.scale)
logcdf(d::Laplace, x::Float64) = x < d.location ? loghalf + ((x-d.location)/d.scale) : loghalf + log2mexp((d.location-x)/d.scale)
logccdf(d::Laplace, x::Float64) = x < d.location ? loghalf + log2mexp((x-d.location)/d.scale) : loghalf + ((d.location-x)/d.scale)

quantile(d::Laplace, p::Float64) = p < 0.5 ? d.location + d.scale*log(2.0*p) : d.location - d.scale*log(2.0*(1.0-p))
cquantile(d::Laplace, p::Float64) = p >= 0.5 ? d.location + d.scale*log(2.0*(1.0-p)) : d.location - d.scale*log(2.0*p)
invlogcdf(d::Laplace, lp::Float64) = lp < loghalf ? d.location + d.scale*(logtwo + lp) : d.location - d.scale*(logtwo + log1mexp(lp))
invlogccdf(d::Laplace, lp::Float64) = lp >= loghalf ? d.location + d.scale*(logtwo + log1mexp(lp)) : d.location - d.scale*(logtwo + lp)

function gradlogpdf(d::Laplace, x::Float64)
    d.location != x || error("Gradient is undefined at the location point")
    x > d.location ? - 1.0 / d.scale : 1.0 / d.scale
end

function mgf(d::Laplace, t::Real)
    st = d.scale*t
    exp(t * d.location) / ((1.0-st)*(1.0+st))
end
function cf(d::Laplace, t::Real)
    st = d.scale*t
    exp(im * t * d.location) / ((1.0-st)*(1.0+st))
end

## Sampling
function rand(d::Laplace)
    er = Base.Random.randmtzig_exprnd()
    randbool() ? d.location - d.scale * er : d.location + d.scale * er
end

## Fitting
function fit_mle(::Type{Laplace}, x::Array)
    a = median(x)
    Laplace(a, mad(x, a))
end


# Main functions of adaptative MCL algorithms
# Berkeley 2015 - Maxime Legrand

#Pkg.add("Gadfly")
using Gadfly
include("../plot_tools.jl")
include("../matrix_tools.jl")
include("mcl.jl")

##########################
# r-adaptative algorithm #
##########################

"""r-adaptative MCL algorithm
Input : A, input of MCL
    e, expansion coefficient
    p, number of iterations
    rmin, rmax, resp, such that the r tested are those in rmin:resp:rmax
Output : rspot, Array{Any,2} such that for each i in size(rspot,1) :
    - rspot[i,1] is the r-value used for MCL
    - rspot[i,2] is the number of clusters spotted by MCL(r=rspot[i,1])
    - rspot[i,3] is the output of MCL(r=rspot[i,1])"""
function rad_mcl(A; e=2, p=20, rmin=2, rmax=10, resp=0.1)
    rrange = [rmin:resp:rmax]
    routput = map(r -> mcl(A, e=e, r=r, p=p), rrange)
    rspot = Array(Any,0,3)
    nclust2 = mcl_card(routput[1])
    for i = 1:(length(rrange)-1)
        nclust1, nclust2 = nclust2, mcl_card(routput[i+1])
        if nclust1 < nclust2
            temp = Array(Any, 1, 3)
            temp[1, 1:2] = [rrange[i] nclust1]
            temp[1,3] = routput[i]
            rspot = vcat(rspot, temp)
        end
    end
    temp = Array(Any, 1, 3)
    temp[1, 1:2] = [rrange[end] nclust2]
    temp[1,3] = routput[end]
    rspot = vcat(rspot, temp)
    return rspot
end

##########################
# e-adaptative algorithm #
##########################

"""e-adaptative MCL algorithm
Input : A, input of MCL
    r, inflation coefficient
    p, number of iterations
    emin, emax, eesp, such that the e tested are those in emin:eesp:emax (eesp needs to be an integer)
Output : espot, Array{Any,2} such that for each i in size(espot,1) :
    - espot[i,1] is the e-value used for MCL
    - espot[i,2] is the number of clusters spotted by MCL(e=espot[i,1])
    - espot[i,3] is the output of MCL(e=espot[i,1])"""
function ead_mcl(A; r=3, p=20, emin=2, emax=10, eesp=1)
    erange = [emin:eesp:emax]
    eoutput = map(e -> mcl(A, e=e, r=r, p=p), erange)
    espot = Array(Any, 1, 3)
    nclust2 = mcl_card(eoutput[1])
    espot[1, 1:2] = [erange[1] nclust2]
    espot[1, 3] = eoutput[1]
    for i = 1:(length(erange)-1)
        nclust1, nclust2 = nclust2, mcl_card(eoutput[i+1])
        if nclust1 > nclust2
            temp = Array(Any, 1, 3)
            temp[1, 1:2] = [erange[i+1] nclust2]
            temp[1,3] = eoutput[i+1]
            espot = vcat(espot, temp)
        end
    end
    return espot
end

#################
# Leave one out #
#################

##########################
# Arbitrary assignations #
##########################


"""
    rand_init(x::Array; range = (0, 1))

Generate a random initialization of the same size and type as `x`.
If `x` is an integer array, then `range` specifies a set of integers.
Otherwise, `range` specifies the set of real numbers.
"""
rand_init(x::AbstractArray{<:Integer}; range = (0, 255)) = rand(range[1]:range[2], size(x)...)
rand_init(x::AbstractArray; range = (0, 1)) = (range[2] - range[1]) .* rand(size(x)...) .+ range[1]

"""
    proj_lball!(x, δ; ϵ, ϵnorm)

Project the perturbed sample (`x + δ`) onto an `ϵnorm`-ball of size `ϵ` around `x`.

# Arguments
- `x`: the adversarial (perturbed) sample/batch
- `δ`: the pertubations
- `ϵ`: the size of the L-norm ball
- `ϵnorm`: the type of L-norm used
"""
function proj_lball!(x, δ; ϵ, ϵnorm)
    if isinf(ϵnorm)
        x .= x .+ max.(min.(δ, ϵ), -ϵ)
    else
        map!((xi, δi) -> xi .+ δi * (ϵ / max.(norm(reshape(δi, :), ϵnorm), ϵ)),
             x, eachslice(x; dims = ndims(x)), eachslice(δ; dims = ndims(δ)))
    end

    return x
end


# Latafat, Themelis, Patrinos, "Block-coordinate and incremental aggregated
# proximal gradient methods for nonsmooth nonconvex problems."
# arXiv:1906.10053 (2019).
#
# Latafat. "Distributed proximal algorithms for large-scale structured optimization"
# PhD thesis, KU Leuven, 7 2020.
#
# Mairal, "Incremental majorization-minimization optimization with application to
# large-scale machine learning."
# SIAM Journal on Optimization 25, 2 (2015), 829–855.
#
# Defazio, Domke, "Finito: A faster, permutable incremental gradient method
# for big data problems."
# In International Conference on Machine Learning (2014), pp. 1125-1133.
#

using LinearAlgebra
using ProximalOperators
using ProximalAlgorithms.IterationTools
using Printf
using Base.Iterators
using Random
using StatsBase: sample

export solution

include("Finito_basic.jl")
include("Finito_LFinito.jl")
include("Finito_adaptive.jl")

struct Finito{R<:Real}
    γ::Maybe{Union{Array{R},R}}
    sweeping::Int8
    LFinito::Bool
    adaptive::Bool
    minibatch::Tuple{Bool,Int}
    maxit::Int
    verbose::Bool
    freq::Int
    α::R
    tol::R
    tol_b::R
    function Finito{R}(;
        γ::Maybe{Union{Array{R},R}} = nothing,
        sweeping = 1,
        LFinito::Bool = false,
        adaptive::Bool = false,
        minibatch::Tuple{Bool,Int} = (false, 1),
        maxit::Int = 10000,
        verbose::Bool = false,
        freq::Int = 10000,
        α::R = R(0.999),
        tol::R = R(1e-8),
        tol_b::R = R(1e-9),
    ) where {R}
        @assert γ === nothing || minimum(γ) > 0
        @assert maxit > 0
        @assert tol > 0
        @assert tol_b > 0
        @assert freq > 0
        new(γ, sweeping, LFinito, adaptive, minibatch, maxit, verbose, freq, α, tol, tol_b)
    end
end

function (solver::Finito{R})(
    x0::AbstractArray{C};
    F = nothing,
    g = ProximalOperators.Zero(),
    L = nothing,
    N = N,
) where {R,C<:RealOrComplex{R}}

    stop(state) = false

    disp(it, state) = @printf "%5d | %.3e  \n" it state.hat_γ

    F === nothing && (F = fill(ProximalOperators.Zero(), (N,)))
    # dispatching the iterator
    if solver.LFinito
        iter = FINITO_LFinito_iterable(
            F,
            g,
            x0,
            N,
            L,
            solver.γ,
            solver.sweeping,
            solver.minibatch[2],
            solver.α,
        )
    elseif solver.adaptive
        iter = FINITO_adaptive_iterable(
            F,
            g,
            x0,
            N,
            L,
            solver.tol,
            solver.tol_b,
            solver.sweeping,
            solver.α,
        )
    else
        iter = FINITO_basic_iterable(
            F,
            g,
            x0,
            N,
            L,
            solver.γ,
            solver.sweeping,
            solver.minibatch[2],
            solver.α,
        )
    end

    iter = take(halt(iter, stop), solver.maxit)
    iter = enumerate(iter)

    num_iters, state_final = nothing, nothing
    for (it_, state_) in iter  # unrolling the iterator
        # see https://docs.julialang.org/en/v1/manual/interfaces/index.html
        if solver.verbose && mod(it_, solver.freq) == 0
            disp(it_, state_)
        end
        num_iters, state_final = it_, state_
    end
    if solver.verbose && mod(num_iters, solver.freq) !== 0
        disp(num_iters, state_final)
    end # for the final iteration
    return solution(state_final), num_iters
end

"""
    Finito([γ, sweeping, LFinito, adaptive, minibatch, maxit, verbose, freq, tol, tol_b])

Instantiate the Finito algorithm for solving fully nonconvex optimization problems of the form

    minimize 1/N sum_{i=1}^N f_i(x) + g(x)

where `f_i` are smooth and `g` is possibly nonsmooth, all of which may be nonconvex.

If `solver = Finito(args...)`, then the above problem is solved with

	solver(x0, [F, g, N, L])

where F is an array containing f_i's, x0 is the initial point, and L is an array of
smoothness moduli of f_i's; it is optional in the adaptive mode or if γ is provided.

Optional keyword arguments are:
* `γ`: an array of N stepsizes for each coordinate
* `sweeping::Int` 1 for uniform randomized (default), 2 for cyclic, 3 for shuffled
* `LFinito::Bool` low memory variant of the Finito/MISO algorithm
* `adaptive::Bool` to activate adaptive smoothness parameter computation
* `minibatch::(Bool,Int)` to use batchs of a given size
* `maxit::Integer` (default: `10000`), maximum number of iterations to perform.
* `verbose::Bool` (default: `true`), whether or not to print information during the iterations.
* `freq::Integer` (default: `10000`), frequency of verbosity.
* `α::R` parameter where γ_i = αN/L_i
* `tol::Real` (default: `1e-8`), absolute tolerance for the adaptive case
* `tol_b::R` tolerance for the backtrack (default: `1e-9`)
"""

Finito(::Type{R}; kwargs...) where {R} = Finito{R}(; kwargs...)
Finito(; kwargs...) = Finito(Float64; kwargs...)

"""
If `solver = Finito(args...)`, then

    itr = iterator(solver, x0, [F, g, N, L])

is an iterable object. Note that [maxit, verbose, freq] fields of the solver are ignored here.

The solution at any given state can be obtained using solution(state), e.g.,
for state in Iterators.take(itr, maxit)
    # do something using solution(state)
end

See https://docs.julialang.org/en/v1/manual/interfaces/index.html
and https://docs.julialang.org/en/v1/base/iterators/ for a list of iteration utilities
"""

function iterator(
    solver::Finito{R},
    x0::AbstractArray{C};
    F = nothing,
    g = ProximalOperators.Zero(),
    L = nothing,
    N = N,
) where {R,C<:RealOrComplex{R}}
    F === nothing && (F = fill(ProximalOperators.Zero(), (N,)))
    # dispatching the iterator
    if solver.LFinito
        iter = FINITO_LFinito_iterable(
            F,
            g,
            x0,
            N,
            L,
            solver.γ,
            solver.sweeping,
            solver.minibatch[2],
            solver.α,
        )
    elseif solver.adaptive
        iter = FINITO_adaptive_iterable(
            F,
            g,
            x0,
            N,
            L,
            solver.tol,
            solver.tol_b,
            solver.sweeping,
            solver.α,
        )
    else
        iter = FINITO_basic_iterable(
            F,
            g,
            x0,
            N,
            L,
            solver.γ,
            solver.sweeping,
            solver.minibatch[2],
            solver.α,
        )
    end
    return iter
end


# Approximations
include("./approximations/gradient_descent.jl")
include("./approximations/SR1.jl")
include("./approximations/DFP.jl")
include("./approximations/BB.jl")
include("./approximations/BFGS.jl")
include("./approximations/DBFGS.jl")
include("./approximations/LBFGS.jl")
include("./approximations/newton.jl")
include("./approximations/scalings.jl")

include("update_qn.jl")
include("find_direction.jl")


# dependencies
using LFAToolkit
using LinearAlgebra
using Plots
# setup
finep = 4
coarsep = 1
numbercomponents = 1
dimension = 1
mesh = Mesh1D(1.0)

ctofbasis = TensorH1LagrangePProlongationBasis(coarsep+1, finep+1, numbercomponents, dimension)

# diffusion operators
finediffusion = GalleryOperator("diffusion", finep+1, finep+1, mesh)
coarsediffusion = GalleryOperator("diffusion", coarsep+1, finep+1, mesh)

# Chebyshev smoother
nos = Schwarz(finediffusion)

# p-multigrid preconditioner
multigrid = PMultigrid(finediffusion, coarsediffusion, nos, [ctofbasis])

# full operator symbols
numbersteps = 250
maxeigenvalue = 0
θ_min = -π/2
θ_max = 3π/2
θ_step = 2π/(numbersteps-1)
θ_range = θ_min:θ_step:θ_max

# compute and plot smoothing factor
# setup
ω = [2.0/3.0]
eigenvalues = zeros(numbersteps, finep)

# compute
for i in 1:numbersteps
    θ = [θ_range[i]]
    if abs(θ[1]) >  π/512
        A = computesymbols(multigrid, ω, [1,1], θ)
        #currenteigenvalues = [real(val) for val in eigvals(I-A)]
        currenteigenvalues = [real(val) for val in eigvals(A)]
        eigenvalues[i, :] = currenteigenvalues
    end
end

# plot
xrange = θ_range/π
plot(
    xrange,
    xlabel="θ/π",
    xtickfont=font(12, "Courier"),
    eigenvalues,
    ytickfont=font(12, "Courier"),
    ylabel="λ",
    linewidth=3,
    legend=:none,
    title="Spectrum of NOS Symbol",
    palette=palette(:tab10)
)
ylims!(min(0.0, eigenvalues...) * 1.1, max(eigenvalues...) * 1.1)
savefig("pmg_nos_spectrum_4_1")


TriangulatedComplex{N,K}() where {N, K} = TriangulatedComplex{N,K}(
    CellComplex{N,K}(), Dict{Cell{N}, Vector{SignedSimpleSimplex{N}}}())

function TriangulatedComplex(simplices::AbstractVector{Simplex{N, K}}) where {N, K}
    # cache to make sure the same simplex isn't turned into a cell twice
    cells = Dict{Set{Point{N}}, Cell{N}}()
    simplex_mapping = Dict{Cell{N}, Vector{SignedSimpleSimplex{N}}}()
    for k in 1:K
        for simplex in simplices
            for s in subsimplices(simplex, k)
                key = Set(s.points)
                if !(key in keys(cells))
                    s_cell = Cell(s)
                    cells[key] = s_cell
                    simplex_mapping[s_cell] = [(SimpleSimplex(s), true)]
                    for face in subsimplices(s, k-1)
                        face_cell = cells[Set(face.points)]
                        i = findfirst(p -> !(p in face.points), s.points)
                        pm = sign_of_permutation(face.points, face_cell.points)
                        parent!(face_cell, s_cell, pm * (1 - 2 * mod(i-1, 2)) > 0)
                    end
                end
            end
        end
    end
    complex = CellComplex{N,K}(collect(values(cells)))
    return TriangulatedComplex{N,K}(complex, simplex_mapping)
end

import Base: show
show(io::IO, tcomp::TriangulatedComplex{N,K}) where {N, K} = print(
     io, "TriangulatedComplex{$N,$K}$(tuple(map(length, tcomp.complex.cells)...))")

Mesh(tcomp::TriangulatedComplex{N}, center::Function) where N =
    Mesh(tcomp, dual(tcomp.complex, center))

show(io::IO, m::Mesh{N,K}) where {N, K} = print(io, "Mesh{$N,$K}($(m.primal), $(m.dual))")

export signed_volume
"""
    signed_volume(m::Metric{N}, tcomp::TriangulatedComplex{N}, c::Cell{N}) where N

Find the signed volume of a cell in a TriangulatedComplex.
"""
signed_volume(m::Metric{N}, tcomp::TriangulatedComplex{N}, c::Cell{N}) where N =
    sum([volume(m, Simplex(s)) * (2 * b - 1) for (s, b) in tcomp.simplices[c]])

export volume
"""
    signed_volume(m::Metric{N}, tcomp::TriangulatedComplex{N}, c::Cell{N}) where N

Find the absolute value of the signed volume of a cell in a TriangulatedComplex.
"""
volume(m::Metric{N}, tcomp::TriangulatedComplex{N}, c::Cell{N}) where N =
    abs(signed_volume(m, tcomp, c))

SimpleBarySimplex{N} = Vector{SimpleBarycentric{N}}

"""
    elementary_duals!(simplices::Dict{Cell{N}, Vector{Tuple{SimpleBarySimplex{N}, Bool}}},
        center::Function, c::Cell{N}) where N

Compute the elementary dual simplices of `cell`. `simplices` is a mapping from primal cells
to dual elementary simplices, specified as Barycentrics along with signs, and is used for
memoization. `center` is a function that takes a `Simplex{N, K}` to a `Barycentric{N, K}`.
"""
function elementary_duals!(simplices::Dict{Cell{N}, Vector{Tuple{SimpleBarySimplex{N}, Bool}}},
    center::Function, c::Cell{N}) where N
    if !(c in keys(simplices))
        c_center = SimpleBarycentric(center(Simplex(c)))
        simplices[c] = Tuple{SimpleBarySimplex{N}, Bool}[]
        if isempty(c.parents)
            push!(simplices[c], ([c_center], true))
        end
        for p in keys(c.parents)
            for (ps, sign) in elementary_duals!(simplices, center, p)
                opposite_index = first_setdiff_index(ps[1].simplex.points, c.points)
                new_sign = (ps[1].coords[opposite_index] >= 0) == sign
                new_barycentrics = [c_center, ps...]
                push!(simplices[c], (new_barycentrics, new_sign))
            end
        end
    end
    return simplices[c]
end

export dual
"""
    dual(primal::CellComplex{N, K}, center::Function) where {N, K}

Compute the dual TriangulatedComplex of a simplicial complex. `center` is a function that
takes a `Simplex{N, K}` to a `Barycentric{N, K}`.
"""
function dual(primal::CellComplex{N, K}, center::Function) where {N, K}
    dual_tcomp = TriangulatedComplex{N, K}()
    primal_to_elementary_duals = Dict{Cell{N}, Vector{Tuple{SimpleBarySimplex{N}, Bool}}}()
    primal_to_duals = Dict{Cell{N}, Cell{N}}()
    # iterate from high to low dimension so the cells of
    # the dual are constructed in order of dimension
    for k in reverse(1:K)
        for cell in primal.cells[k]
            elementary_duals = elementary_duals!(primal_to_elementary_duals, center, cell)
            signed_elementary_duals = [(SimpleSimplex(p[1]), p[2])
                for p in elementary_duals]
            points = unique(vcat([p[1].points for p in signed_elementary_duals]...))
            dual_cell = Cell(points, K-k+1)
            push!(dual_tcomp.complex, dual_cell)
            dual_tcomp.simplices[dual_cell] = signed_elementary_duals
            primal_to_duals[cell] = dual_cell
            for parent in keys(cell.parents)
                o = cell.parents[parent]
                dual_child = primal_to_duals[parent]
                # for k > 1, use o but for k == 1 use !o. This guarantees
                # that d² = 0 and that the dual mesh is positively oriented.
                parent!(dual_child, dual_cell, k == 1 ? !o : o)
            end
        end
    end
    return dual_tcomp
end

"""
    dual(mesh::Mesh{N, K}, c::Cell{N}) where {N, K}

Find the dual cell of `c` in `mesh`.
"""
function dual(mesh::Mesh{N, K}, c::Cell{N}) where {N, K}
    primal = mesh.primal.complex
    dual = mesh.dual.complex
    comp1, comp2 = (c in primal.cells[c.K]) ? (primal, dual) : (dual, primal)
    i = findfirst(isequal(c), comp1.cells[c.K])
    return comp2.cells[K-c.K+1][i]
end


using Flux, Flux.Tracker
using Flux.@epochs
using PyPlot
using CSV

pygui(false)

debugging = false
raw_data = CSV.read("logs/mturk_trusted_condition_1_mod.csv")

# set random seed for repeatability
#  srand(12345)
# 4 is good
# 21 is good
# 39 is beter, with 0.3 dropout and 10 hidden nodes
srand_num = 39
srand(srand_num)

rand_order = randperm(size(raw_data,1))
x1 = [float(q) for q in raw_data[:tprob][rand_order]]
x2 = [float(q) for q in raw_data[:N][rand_order]]
y = [float(q) for q in raw_data[:X3_1][rand_order]]

nx = length(x1)
train_pct = 0.90
train_limit = convert(Int,round(nx*train_pct))
train_set = 1:train_limit

valid_pct = 1.00
valid_limit = convert(Int,round(nx*valid_pct))
valid_set = train_limit+1:valid_limit

#  test_set = valid_limit+1:length(x1)

x1_train = x1[train_set]
x2_train = x2[train_set]
y_train = y[train_set]'

x_train = [x1_train x2_train]'
#  x_train = x1_train'

data = zip(x_train,y_train)

x1_valid = x1[valid_set]
x2_valid = x2[valid_set]
y_valid = y[valid_set]'

x_valid = [x1_valid x2_valid]'
#  x_valid = x1_valid'

#  x1_test = x1[test_set]
#  x2_test = x2[test_set]
#  y_test = y[test_set]'
#
#  x_test = [x1_test x2_test]'

num_hidden = 10
dropout_rate = 0.30

net = Chain(
            Dense(size(x_train,1),num_hidden,relu),
            Dropout(dropout_rate),
            Dense(num_hidden,num_hidden,relu),
            Dropout(dropout_rate),
            Dense(num_hidden,1)
           )

display(params(net))

loss(x,y) = Flux.mse(net(x),y)

opt = ADAM(params(net) , 0.001; β1 = 0.9, β2 = 0.999, ϵ = 1e-08, decay = 0)

evalcb() = @show(loss(tx,ty))

# array to store the validation error per batch
res = []
tr = []
tst = []

function check_err(x,y,ary)
    Flux.testmode!(net)
    append!(ary,loss(x,y).data)
    Flux.testmode!(net,false)
end

function valid_cb()
    # validate data
    check_err(x_valid,y_valid,res)
end
function train_cb()
    # validate data
    check_err(x_train,y_train,tr)
end
function test_cb()
    # validate data
    check_err(x_test,y_test,tst)
end

function cb_fun()
    valid_cb()
    train_cb()
    #  test_cb()
end

num_epocs = 50

@epochs num_epocs Flux.train!(loss,data,opt,cb = cb_fun)
#  Flux.train!(loss,data,opt,cb = Flux.throttle(evalcb(),3))
#  Flux.train!(loss,data,opt,cb = ()=>println("test"))

fig,ax = PyPlot.subplots(1,1)
fig[:set_size_inches](6,4)

epoc_mark = collect(0+1/length(y_train):1/length(y_train):num_epocs)
#  println(length(res))
#  println(length(epoc_mark))
#  display(epoc_mark)
#  readline()

ax[:plot](epoc_mark,res,color=:black,label="validation (N=$(length(y_valid)))")
ax[:plot](epoc_mark,tr,color=:blue,label="training (N=$(length(y_train)))")
#  ax[:plot](epoc_mark,tst,color=:orange,label="test (N=$(length(y_test)))")
if debugging
    ax[:set_title]("Num Hidden: $num_hidden, its: $num_epocs, srand: $srand_num, dropout: $dropout_rate")
else
    ax[:set_title]("Loss vs Epoch")
end

ax[:set_xlabel]("epoc")
ax[:set_ylabel]("loss (MSE)")
ax[:ticklabel_format](axis="y",style="sci",scilimits=(0,1),useOffset=false)
ax[:legend]()

#  Flux.testmode!(net)
#  test_err = loss(x_test,y_test)
#  Flux.testmode!(net,false)

#  ax[:axhline](y=test_err.data,color=:red)

output_fldr = "sqmodel_convergence"

savefig("$(output_fldr)/$(num_epocs)_$(num_hidden)_$(srand_num).pdf")
savefig("$(output_fldr)/current_fit.pdf")


using DataDrivenDiffEq
using ModelingToolkit
using LinearAlgebra
using OrdinaryDiffEq
using Test
using DataDrivenDiffEq.Optimise

include("./basis.jl")
include("./koopman.jl")
include("./sindy.jl")
include("./isindy.jl")
include("./utils.jl")
include("./optimise.jl")


# check the generlised Euler (Stieltjes) numbers
#    this is mostly just a check that code is correct
#    most errors are arising because of errors in the functions being used
using Polylogarithms
using BenchmarkTools
using DataFrames, CSV
using PyPlot

S = Symbol("S_n")
data1 = CSV.read("../data/stieltjes_test_data_1.csv")
m = size(data1,1)
S1 = zeros(Float64, m)
S1_ref = zeros(Float64, m)
error1 = zeros(Float64, m)
rel_error1 = zeros(Float64, m)
for i=1:m
    S1[i] = stieltjes( data1[i,:n] )
    S1_ref[i] = data1[i,S]
    error1[i] = Float64( S1[i] - S1_ref[i]  )
    rel_error1[i] = error1[i]/Float64( S1_ref[i] )
end
println("max abs. error      = $(maximum( abs.(error1) ))")
println("max rel. abs. error = $(maximum( abs.(rel_error1) ))")


using StochasticDiffEq, Test, Random
using DiffEqGPU
using CUDA

function weak_error(prob,alg,numtraj, batchsize, f_true,trange;abstol=1,reltol=0,ensemblealg=EnsembleCPUArray())
  sol = @time solve(prob,alg,ensemblealg,
    dt=0.01f0,adaptive=true,abstol=abstol,reltol=reltol,
    trajectories=numtraj,batch_size=batchsize,
    saveat = trange
    )
  computed_exp = (sol.u/numtraj)[1,:]
  true_exp = f_true.(trange)
  return sum((computed_exp-true_exp).^2)/length(trange)
end

function prob_func(prob, i, repeat)
    #(i%50000 == 0) && @show i
    remake(prob,seed=seeds[i])
end

function reduction(u,batch,I)
  u.+sum(batch),false
end

function output_func(sol,i)
  #h1(asinh(sol[end][1])),false
  h1.(asinh.(sol)),false
end

# prob 1
u₀ = [0.0f0]
tspan = (0.0f0,2.0f0)
h1(z) = z^3-6*z^2+8*z

function f1!(du,u,p,t)
 @inbounds begin
     du[1] = 1//2*u[1]+sqrt(u[1]^2 +1)
 end
 nothing
end

function g1!(du,u,p,t)
 @inbounds begin
     du[1] = sqrt(u[1]^2 +1)
 end
 nothing
end

f_true1(t) = t^3-3*t^2+2*t

prob1 = SDEProblem(f1!,g1!,u₀,tspan)
ensemble_prob1 = EnsembleProblem(prob1;
        output_func = output_func,
        prob_func = prob_func,
        reduction = reduction,
        u_init=Vector{eltype(prob1.u0)}([0.0])
        )

# prob 2
u₀ = [0.1f0,0.1f0]
function f2!(du,u,p,t)
  @inbounds begin
    du[1] = 3//2*u[1]
    du[2] = 3//2*u[2]
  end
  nothing
end
function g2!(du,u,p,t)
  @inbounds begin
    du[1] = 1//10*u[1]
    du[2] = 1//10*u[2]
  end
  nothing
end

f_true2(t) = 1//10*exp(3//2*t) #1//100*exp(301//100*t)

h2(z) = z #z^2

prob2 = SDEProblem(f2!,g2!,u₀,tspan)
ensemble_prob2 = EnsembleProblem(prob2;
        output_func = (sol,i) -> (h2.(sol),false),
        prob_func = prob_func,
        reduction = reduction,
        u_init=Vector{eltype(prob2.u0)}([0.0, 0.0])
        )

tsave = 0.0f0:0.05f0:2f0

probs = Vector{EnsembleProblem}(undef, 2)
probs[1] = ensemble_prob1
probs[2] = ensemble_prob2

ftrue = Vector{}(undef, 2)
ftrue[1] = f_true1
ftrue[2] = f_true2

numtraj = Int(1e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

for i in 1:2
  @show i

  err1 = weak_error(probs[i],DRI1NM(),numtraj,Int(1e4),ftrue[i],tsave,abstol=1f0,reltol=1f0, ensemblealg=EnsembleGPUArray())
  @show err1
  # err2 = weak_error(probs[i],DRI1NM(),numtraj,Int(1e4),ftrue[i],tsave,abstol=0.1f0,reltol=0.1f0, ensemblealg=EnsembleGPUArray())
  # @show err2
  err3 = weak_error(probs[i],DRI1NM(),numtraj,Int(1e4),ftrue[i],tsave,abstol=0.01f0,reltol=0.01f0, ensemblealg=EnsembleGPUArray())
  @show err3
  @test err1 > err3
  println("")
end

#
# sol = @time solve(probs[1],DRI1NM(),EnsembleGPUArray(),
#   dt=0.001f0,adaptive=false,abstol=0.1f0,reltol=0.1f0,
#   trajectories=numtraj,batch_size=Int(1e1),
#   saveat = tsave
#   )
# computed_exp = (sol.u/numtraj)[1,:]
# true_exp = f_true1.(tsave)
#
# @show  sum((computed_exp-true_exp).^2)/length(tsave)
#
# using Plots; plot(tsave, true_exp)
# plot!(tsave, computed_exp)


# 3D Grid examples
# ===============
#

# ## Quadrilateral
function quadrilateral(;hx=0.25, hy=0.2, hz=0.1)
    X=collect(0:hx:1)
    Y=collect(0:hy:1)
    Z=collect(0:hz:1)
    simplexgrid(X,Y,Z)
end
# ![](quadrilateral.svg)

function mask_bedges()
    grid    = quadrilateral(hx=0.25, hy=0.25, hz=0.25)

    bedgemask!(grid, [0.0, 0.0, 0.0], [0.0, 0.0, 1.0], 1)
    bedgemask!(grid, [0.0, 0.0, 0.0], [0.0, 1.0, 0.0], 2)
    bedgemask!(grid, [0.0, 1.0, 0.0], [0.0, 1.0, 1.0], 3)
    bedgemask!(grid, [0.0, 0.0, 1.0], [0.0, 1.0, 1.0], 4)
    bedgemask!(grid, [0.0, 1.0, 0.0], [0.0, 0.0, 1.0], 5)

    true
end


module HeatEqCalc

function laplacian4!(v, u, dx)
    v[begin] = (u[end] + u[begin+1] - 2u[begin])/dx^2
    @. @views v[begin+1:end-1] = (u[begin:end-2] + u[begin+2:end] - 2u[begin+1:end-1])/dx^2
    v[end] = (u[end-1] + u[begin] - 2u[1end])/dx^2
    return
end

function heateq4(u0, dx, tmax, N=200)
    t = 0:dx:tmax
    dt = step(t)
    u = similar(u0, length(u0), length(t)+1)
    u[:, 1] = u0
    v = similar(u0)
    N = 100
    for i in 2:length(t)+1
        @. @views u[:, i] = u[:, i-1]
        for _ in 1:N
            @views laplacian4!(v, u[:, i], dx)
            @. @views u[:, i] += v*dt/N
        end
    end
    t, u
end

using StaticArrays

function calc_sol(
        n = 200,
        x = range(-π, π; length=n+1)[1:end-1],
        m = 200
    )
    dx = step(x)
    noise1, noise2 = 0.3randn(m), 0.3randn(m)
    f(x) = SVector{m}((1 .+ noise1)*sin(x) .- (1 .+ noise2)*cos(2x))
    u0 = f.(x)
    tmax = 1.0
    t, u = heateq4(u0, dx, tmax)
    (; t, x, u, m)
end

using DelimitedFiles

function save_sol(sol)
    t, x, u, m = sol
    U = vcat((getindex.(u, k) for k in 1:m)...)
    writedlm("sol_m.txt", m)
    writedlm("sol_t.txt", t)
    writedlm("sol_x.txt", x)
    writedlm("sol_U.txt", U)
end

function load_sol(name)
    m = parse(Int, read(name * "_m.txt", String))
    t = vec(readdlm(name * "_t.txt"))
    x = vec(readdlm(name * "_x.txt"))
    U = readdlm(name * "_U.txt")
    u = [SVector((U[k*length(x) + i, j] for k in 0:m-1)...) for i in 1:length(x), j in 1:length(t)]
    (; t, x, u, m)
end

if isfile(joinpath(pkgdir(@__MODULE__), "src", "trace.jl"))
    include("trace.jl")
end

end # module


# Convenience functions

# The following methods are defined:
# viterbi(a, A, L)              -> z
# viterbi(hmm, observations)    -> z

"""
    viterbi(a, A, L; logl) -> Vector

Find the most likely hidden state sequence, see [Viterbi algorithm](https://en.wikipedia.org/wiki/Viterbi_algorithm).
"""
function viterbi(a::AbstractVector, A::AbstractMatrix, L::AbstractMatrix; logl = false)
    T1 = Matrix{Float64}(undef, size(L))
    T2 = Matrix{Int}(undef, size(L))
    z = Vector{Int}(undef, size(L,1))
    if logl
        viterbilog!(T1, T2, z, a, A, L)
    else
        warn_logl(L)
        viterbi!(T1, T2, z, a, A, L)
    end
    z
end

"""
    viterbi(hmm, observations; logl, robust) -> Vector

**Example**
```julia
using Distributions, HMMBase
hmm = HMM([0.9 0.1; 0.1 0.9], [Normal(0,1), Normal(10,1)])
y = rand(hmm, 1000)
zv = viterbi(hmm, y)
```
"""
function viterbi(hmm::AbstractHMM, observations; robust = false, kwargs...)
    L = likelihoods(hmm, observations; robust = robust, kwargs...)
    viterbi(hmm.a, hmm.A, L; kwargs...)
end


function SeisSmooth2(A;L1=7,L2=7,Nrepeat=1)
    a1 = isodd(L1) ? 0 : 1
    a2 = isodd(L2) ? 0 : 1
    kernel = ones(eltype(A),L1,L2)/(L1*L2);
    for irepeat = 1 : Nrepeat
	A = conv2(A,kernel)
	A = A[floor(Int,L1/2) + 1:end-floor(Int,L1/2) + a1,floor(Int,L2/2) + 1:end-floor(Int,L2/2) + a2]
	end
    return A
end


using Equations
consts=@equs(c=3e8, G=6.67e-11, massSun=2e30)
constsu=@equs(c=U(3e8,:m/:s), G=U(6.67e-11,:N*:m^2/:kg^2), massSun=U(2e30,:kg))

F1=-:G*:M*:m/:r^2
F2=:m*:v^2/:r
orbit=Equation(F2,F1)
v=sqrt(orbit/:m*:r) #orbital velocity

a=-:G*:M/:R^2

#Black holes
E=@equ m*v^2/2=G*M*m/R
R=E*:R/:m/:v^2*2
R_BH=R&@equs(v=c, M=massSun)&consts

R=U(6.67e-11,:m^3/:s^2/:kg)*U(2e30,:kg)/U(3e8,:m/:s)/U(3e8,:m/:s)*2
println(R)


using Unitful
import Unitful: Frequency, Hz, kHz, 𝐓, s

samplerate = 48000Hz

abstract type AbstractOscillator end

function maxstable(samplerate::Frequency)
    return samplerate / 2 #TODO: this is not the correct formula
end

function checkstability(frequency::Frequency, samplerate::Frequency)
    if frequency >= maxstable(samplerate)
        error("Oscillator is unstable: $(frequency) is to high for $(samplerate) sample rate. The maximum stable frequency is $(maxstable(samplerate))")
    end
end

mutable struct SineOscillator{T<:AbstractFloat} <: AbstractOscillator
    previous::T
    sample::T
    cosine::T
end

function SineOscillator(frequency::Frequency{T}, amplitude::Complex{T}) where {T<:AbstractFloat}
    checkstability(frequency, samplerate)

    magnitude = abs(amplitude)
    phase = angle(amplitude)

    step = 2 * pi * frequency / samplerate
    previous = magnitude * sin(phase)
    sample = magnitude * sin(phase + step)
    cosine = cos(step)

    return SineOscillator{typeof(step)}(previous, sample, cosine)
end

function SineOscillator(frequency::Frequency{T}, amplitude::T) where {T<:AbstractFloat}
    return SineOscillator(frequency, complex(amplitude))
end

function oscillate!(osc::SineOscillator{T}) where {T<:AbstractFloat}
    nextsample = 2 * osc.sample * osc.cosine - osc.previous
    osc.previous, osc.sample = osc.sample, nextsample
    return nextsample
end

struct VectorOscillator{T<:AbstractFloat}
    oscs::Vector{SineOscillator{T}}
end

function oscillate!(osc::VectorOscillator{T}) where {T<:AbstractFloat}
    sum(oscillate!.(osc.oscs))
end


module Links
#################### TransformDistribution ####################

import Distributions: Continuous, ContinuousDistribution,
ContinuousUnivariateDistribution, ContinuousMultivariateDistribution,
## Univariate Distributions
Truncated, Cauchy, Gumbel, Laplace, Logistic, NoncentralT, Normal,
NormalCanon, TDist, Beta, BetaPrime, Chi, Chisq, Erlang, Exponential,
FDist, Frechet, Gamma, InverseGamma, InverseGaussian, Kolmogorov,
LogNormal, NoncentralChisq, NoncentralF, Rayleigh, Weibull, KSOneSided,
NoncentralBeta

import StatsFuns: logit

using Reexport
@reexport using Distributions

using BayesianTools
@reexport using BayesianTools

#################### Bounded Distribution ####################
const TransformDistribution{T<:ContinuousUnivariateDistribution} = Union{T,Truncated{T},Improper}

function link(d::TransformDistribution, x::Real)
  a, b = minimum(d), maximum(d)
  lowerbounded, upperbounded = isfinite(a), isfinite(b)
  if lowerbounded && upperbounded
    logit((x - a) / (b - a))
  elseif lowerbounded
    log(x - a)
  elseif upperbounded
    log(b - x)
  else
    x
  end
end

function invlink(d::TransformDistribution, x::Real)
  a, b = minimum(d), maximum(d)
  lowerbounded, upperbounded = isfinite(a), isfinite(b)
  if lowerbounded && upperbounded
    (b - a) * invlogit(x) + a
  elseif lowerbounded
    exp(x) + a
  elseif upperbounded
    b - exp(x)
  else
    x
  end
end

function logjacobian(d::TransformDistribution, x::Real)
    # x - a >= 0
    # b - x >= 0
    a, b = minimum(d), maximum(d)
    lowerbounded, upperbounded = isfinite(a), isfinite(b)
    if lowerbounded && upperbounded
      log((b-a)/((b-x)*(x-a)))
    elseif lowerbounded
      -log(x - a)
    elseif upperbounded
      -log(x - b)
    else
      0.0
    end
end
#################### RealDistribution ####################

const RealDistribution =Union{Cauchy, Gumbel, Laplace, Logistic, NoncentralT, Normal,
                              NormalCanon, TDist}

link(d::RealDistribution, x::Real) = x
invlink(d::RealDistribution, x::Real) = x
logjacobian(d::RealDistribution, x::Real) = zero(eltype(x))
#################### PositiveDistribution ####################

const PositiveDistribution = Union{BetaPrime, Chi, Chisq, Erlang, Exponential, FDist, Frechet,
                                   Gamma, InverseGamma, InverseGaussian, Kolmogorov, LogNormal,
                                   NoncentralChisq, NoncentralF, Rayleigh, Weibull}

link(d::PositiveDistribution, x::Real) = log(x)
invlink(d::PositiveDistribution, x::Real) = exp(x)
logjacobian(d::PositiveDistribution, x::Real) = log(x)
#################### UnitDistribution ####################

const UnitDistribution = Union{Beta, KSOneSided, NoncentralBeta}

link(d::UnitDistribution, x::Real) = logit(x)
invlink(d::UnitDistribution, x::Real) = invlogit(x)
## logjacobian is d link / d x = log (1/(1-x))
logjacobian(d::UnitDistribution, x::Real) = -log(x)-log(1.0-x)

invlogit(x::Real) = 1.0 / (exp(-x) + 1.0)
invlogit(x::AbstractArray) = map(invlogit, x)

export link, invlink, logjacobian

end


using Test

function quick_sort(arr)
    if length(arr) < 2 return arr end
    pivot = arr[1]
    less = [i for i in arr[2:end] if i <= pivot]
    greater = [i for i in arr[2:end] if i > pivot]
    return vcat(quick_sort(less), [pivot], quick_sort(greater))
end

@test quick_sort([3,5,2,1,4]) == [1,2,3,4,5]


using SymbolicUtils

"""
A function that creates a tuple of CartesianIndices of unit length and `N` dimensions, one pointing along each dimension.
"""
function unit_indices(N::Int) #create unit CartesianIndex for each dimension
    out = Vector{CartesianIndex{N}}(undef, N)
    null = zeros(Int64, N)
    for i in 1:N
        unit_i = copy(null)
        unit_i[i] = 1
        out[i] = CartesianIndex(Tuple(unit_i))
    end
    Tuple(out)
end

function cartesian_to_linear(I::CartesianIndex, s)  #Not sure if there is a builtin that does this - convert cartesian index to linear index of an array of size s
    out = I[1]
    for i in 1:length(s)-1
        out += (I[i+1]-1)*prod(s[1:i])
    end
    return out
end

# Counts the Differential operators for given variable x. This is used to determine
# the order of a PDE.
function count_differentials(term, x::Symbolics.Symbolic)
    S = Symbolics
    SU = SymbolicUtils
    if !S.istree(term)
        return 0
    else
        op = SU.operation(term)
        count_children = sum(map(arg -> count_differentials(arg, x), SU.arguments(term)))
        if op isa Differential && op.x === x
            return 1 + count_children
        end
        return count_children
    end
end

# return list of differential orders in the equation
function differential_order(eq, x::Symbolics.Symbolic)
    S = Symbolics
    SU = SymbolicUtils
    orders = Set{Int}()
    if S.istree(eq)
        op = SU.operation(eq)
        if op isa Differential
            push!(orders, count_differentials(eq, x))
        else
            for o in map(ch -> differential_order(ch, x), SU.arguments(eq))
                union!(orders, o)
            end
        end
    end
    return filter(!iszero, orders)
end

add_dims(A::AbstractArray, n::Int; dims::Int = 1) = cat(ndims(A) + n, A, dims = dims)

""
perpindex(A, dim::Integer) = A[setdiff(1:length(A), dim)]

"""
the size of A perpendicular to dim
"""
perpsize(A::AbstractArray{T,N}, dim::Integer) where {T,N} = size(A)[setdiff(1:N, dim)]


function backtracking_linesearch!{T}(d::Union{DifferentiableFunction,
                                              TwiceDifferentiableFunction},
                                     x::Vector{T},
                                     s::Vector,
                                     x_scratch::Vector,
                                     gr_scratch::Vector,
                                     lsr::LineSearchResults,
                                     alpha::Real = 1.0,
                                     mayterminate::Bool = false,
                                     c1::Real = 1e-4,
                                     c2::Real = 0.9,
                                     rho::Real = 0.9,
                                     iterations::Integer = 1_000)

    # Count the total number of iterations
    iteration = 0

    # Track calls to function and gradient
    f_calls = 0
    g_calls = 0

    # Count number of parameters
    n = length(x)

    # Store f(x) in f_x
    f_x = d.fg!(x, gr_scratch)
    f_calls += 1
    g_calls += 1

    # Store angle between search direction and gradient
    gxp = vecdot(gr_scratch, s)

    # Tentatively move a distance of alpha in the direction of s
    @simd for i in 1:n
        @inbounds x_scratch[i] = x[i] + alpha * s[i]
    end

    # Backtrack until we satisfy sufficient decrease condition
    f_x_scratch = d.f(x_scratch)
    f_calls += 1
    while f_x_scratch > f_x + c1 * alpha * gxp
        # Increment the number of steps we've had to perform
        iteration += 1

        # Ensure termination
        if iteration > iterations
            error("Too many iterations in backtracking_linesearch!")
        end

        # Shrink proposed step-size
        alpha *= rho

        # Update proposed position
        @simd for i in 1:n
            @inbounds x_scratch[i] = x[i] + alpha * s[i]
        end

        # Evaluate f(x) at proposed position
        f_x_scratch = d.f(x_scratch)
        f_calls += 1
    end

    return alpha, f_calls, g_calls
end


# productmap.jl

"""
A product map is diagonal and acts on each of the components of x separately:
`y = f(x)` becomes `y_i = f_i(x_i)`.
"""
struct ProductMap{MAPS,S <: Tuple,T <: Tuple} <: AbstractMap{S,T}
    # maps has an indexable and iterable type, for example a tuple of maps
    maps    ::  MAPS
end

function ProductMap(maps...)
    MAPS = typeof(maps)
    S = typeof(map(x->zero(domaintype(x)), maps))
    T = typeof(map(x->zero(codomaintype(x)), maps))

    ProductMap{MAPS,S,T}(maps)
end

elements(pm::ProductMap) = pm.maps

(m::ProductMap)(x) = applymap(m, x)

applymap(pm::ProductMap, x) = map(applymap, elements(pm), x)

tensorproduct(map1::AbstractMap, map2::AbstractMap) = ProductMap(map1, map2)
tensorproduct(map1::ProductMap, map2::AbstractMap) = ProductMap(elements(map1)..., map2)
tensorproduct(map1::AbstractMap, map2::ProductMap) = ProductMap(map1, elements(map2)...)
tensorproduct(map1::ProductMap, map2::ProductMap) = ProductMap(elements(map1)..., elements(map2)...)

for op in (:inv, :left_inverse, :right_inverse, :jacobian)
    @eval $op(pm::ProductMap) = ProductMap(map($op, elements(pm))...)
end


@doc raw"""
    thermometer(x::Float64, n::Int)
    thermometer(x::Vector{Float64}, n::Int)
"""
function thermometer(x::Float64, n::Int)
    y = zeros(Bool, n)
    k = trunc(Int, clamp(x, 0.0, 1.0) * n)
    y[1:k] .= true
    return y
end

function thermometer(x::Vector{Float64}, n::Int)
    m = length(x)
    y = zeros(Bool, m * n)
    for i = 1:m
        j = (i - 1) * n + 1
        k = trunc(Int, clamp(x[i], 0.0, 1.0) * n) - 1
        y[j:j+k] .= true
    end
    return y
end


using Printf

# Total circles area: https://rosettacode.org/wiki/Total_circles_area

xc = [1.6417233788, -1.4944608174, 0.6110294452, 0.3844862411, -0.2495892950, 1.7813504266,
     -0.1985249206, -1.7011985145, -0.4319462812, 0.2178372997, -0.6294854565, 1.7952608455,
      1.4168575317, 1.4637371396, -0.5263668798, -1.2197352481, -0.1389358881, 1.5293954595,
     -0.5258728625, -0.1403562064, 0.8055826339, -0.6311979224, 1.4685857879, -0.6855727502,
      0.0152957411]
yc = [1.6121789534, 1.2077959613, -0.6907087527, 0.2923344616, -0.3832854473, 1.6178237031,
     -0.8343333301, -0.1263820964, 1.4104420482, -0.9499557344, -1.3078893852, 0.6281269104,
      1.0683357171, 0.9463877418, 1.7315156631, 0.9144146579, 0.1092805780, 0.0030278255,
      1.3782633069, 0.2437382535, -0.0482092025, 0.7184578971, -0.8347049536, 1.6465021616,
      0.0638919221]
r  = [0.0848270516, 1.1039549836, 0.9089162485, 0.2375743054, 1.0845181219, 0.8162655711,
      0.0538864941, 0.4776976918, 0.7886291537, 0.0357871187, 0.7653357688, 0.2727652452,
      1.1016025378, 1.1846214562, 1.4428514068, 1.0727263474, 0.7350208828, 1.2472867347,
      1.3495508831, 1.3804956588, 0.3327165165, 0.2491045282, 1.3670667538, 1.0593087096,
      0.9771215985]

# Size of my grid -- higher values => higher accuracy.
function main(xc::Vector{<:Real}, yc::Vector{<:Real}, r::Vector{<:Real}, ngrid::Integer=10000)
    r2 = r .* r
    ncircles = length(xc)

    # Compute the bounding box of the circles.
    xmin = minimum(xc .- r)
    xmax = maximum(xc .+ r)
    ymin = minimum(yc .- r)
    ymax = maximum(yc .+ r)
    # Keep a counter.
    inside = 0
    # For every point in my grid.
    for x in linspace(xmin, xmax, ngrid), y = linspace(ymin, ymax, ngrid)
        inside += any(r2 .> (x - xc) .^ 2 + (y - yc) .^ 2)
    end
    boxarea = (xmax - xmin) * (ymax - ymin)
    return boxarea * inside / ngrid ^ 2
end

println(@time main(xc, yc, r, 1000))


function wdot(wx::AbstractArray{T}, x::AbstractArray{T},
                 wy::AbstractArray{T}, y::AbstractArray{T}) where T
    out = zero(T)
    @inbounds @simd for i in 1:length(x)
        out += wx[i]*x[i] * wy[i]*y[i]
    end
    return out
end

wnorm(w, x) = sqrt(wdot(w, x, w, x))
assess_convergence(f, ftol) = assess_convergence(NaN, NaN, f, NaN, ftol)
function assess_convergence(x,
                            x_previous,
                            f,
                            xtol,
                            ftol)
    x_converged, f_converged = false, false

    if !any(isnan, x_previous) && chebyshev(x, x_previous) <= xtol
        x_converged = true
    end

    if maximum(abs, f) <= ftol
        f_converged = true
    end

    converged = x_converged || f_converged

    return x_converged, f_converged, converged
end

function check_isfinite(x::AbstractArray)
    if any((y)->!isfinite(y),x)
        i = findall((!).(isfinite.(x)))
        throw(IsFiniteException(i))
    end
end


# one EIF neuron with constant input

push!(LOAD_PATH, abspath(@__DIR__,".."))
using LinearAlgebra,Statistics,StatsBase
using Plots,NamedColors ; theme(:dark)
using SparseArrays
using SpikingRNNs; const global S = SpikingRNNs
using ProgressMeter

function onesparsemat(w::Real)
  mat=Matrix{Float64}(undef,1,1) ; mat[1,1]=w
  return sparse(mat)
end

## Time parameters
dt = 1E-4
Ttot = 30.
# EIF neuron parameters
# τ::Float64 # time constant (membrane capacitance)
# g_l::Float64 # conductance leak
# v_expt::Float64 # exp term threshold
# steep_exp::Float64 # steepness of exponential term
# v_threshold::Float64 # spiking threshold
# v_reset::Float64 # reset after spike
# v_leak::Float64 # reversal potential for leak term
# τ_refractory::Float64 # refractory time

τ = 0.5
g_l = 1.0
v_expt=10.0
steep_exp = 1.0
v_th = 20.
v_r = -5.0
τrefr = 0.0
v_leak = 0.0

myinput = 11.0
mynt = S.NTEIF(τ,g_l,v_expt,steep_exp,v_th,v_r,v_leak,τrefr)
myps = S.PSEIF(mynt,1)

## one static input above threshold
in_state = S.PSSimpleInput(S.InputSimpleOffset(myinput))
mypop = S.Population(myps,(S.FakeConnection(),),(in_state,))

## that's it, let's make the network
myntw = S.RecurrentNetwork(dt,(mypop,))

## Recorder
myrec1 = S.RecStateNow(myps,10,dt,Ttot)
myrec2 = S.RecCountSpikes(myps,dt)

## Run!
times = (0:myntw.dt:Ttot)
nt = length(times)
# initial conditions
myps.state_now[1] = v_r

S.reset!.([myrec1,myrec2])

for (k,t) in enumerate(times)
  S.dynamics_step!(t,myntw)
  myrec1(t,k,myntw)
  myrec2(t,k,myntw)
end
r0 = S.get_mean_rates(myrec2,dt,Ttot)[1]
##
plot(myrec1.times,myrec1.state_now[1,:])
##

# now, use Fokker-Planck theory and Richardson method
# it is hardcoded for now... maybe later I will convert it to
# a function
dv = 1E-3
v_lb = v_r - 0.1
psifun(V) = steep_exp * exp((V-v_expt)/steep_exp)
ndv = round(Int64,(v_th - v_lb)/dv)
k_re = round(Int64,(v_r - v_lb)/dv)
ks = collect(0:ndv)
Vks = @. v_lb + ks * dv
js = fill(NaN,ndv+1)
ps = fill(NaN,ndv+1)

ksrev = collect(ndv+1:-1:2)

js[ndv+1]=1.0
ps[ndv+1]=0.0

for k in ksrev
  krond = k==k_re ? 1.0 : 0.0
  js[k-1] = js[k] - krond
  ps[k-1] = js[k-1] * mynt.τ / (myinput-Vks[k]+ psifun(Vks[k]))
end

r0 = inv(dv*sum(ps))

##
histogram(myrec1.state_now[:];nbins=100,normalize=true)
plot!(Vks,r0.*ps;color=:white,linewidth=2)

## can I use this for the f-I curve?

function f_lif_num(h_in;dt=1E-4,Ttot=10.0)
  in_state = S.PSSimpleInput(S.InputSimpleOffset(h_in))
  mypop = S.Population(myps,(S.FakeConnection(),),(in_state,))
  myntw = S.RecurrentNetwork(dt,(mypop,))
  rec = S.RecCountSpikes(myps,dt)
  ## Run!
  times = (0:myntw.dt:Ttot)
  # initial conditions
  myps.state_now[1] = v_r
  S.reset!(rec)
  for (k,t) in enumerate(times)
    S.dynamics_step!(t,myntw)
    rec(t,k,myntw)
  end
  r0 = S.get_mean_rates(rec,dt,Ttot)[1]
  return r0
end

function f_lif_semianalitic(h_in;dv=1E-3)
  v_lb = v_r - 0.1
  psifun(V) = steep_exp * exp((V-v_expt)/steep_exp)
  ndv = round(Int64,(v_th - v_lb)/dv)
  k_re = round(Int64,(v_r - v_lb)/dv)
  ks = collect(0:ndv)
  Vks = @. v_lb + ks * dv
  js = fill(NaN,ndv+1)
  ps = fill(NaN,ndv+1)

  ksrev = collect(ndv+1:-1:2)

  js[ndv+1]=1.0
  ps[ndv+1]=0.0

  for k in ksrev
    krond = k==k_re ? 1.0 : 0.0
    js[k-1] = js[k] - krond
    ps[k-1] = js[k-1] * mynt.τ / (h_in-Vks[k]+ psifun(Vks[k]))
  end

  return inv(dv*sum(ps))
end

##
_ = let Ttot = 50.0
  nin = 50
  inputs = range(10.0,80.0;length=nin)
  y1 = similar(inputs)
  y2 = similar(inputs)
  @showprogress for i in eachindex(inputs)
    y1[i]=f_lif_num(inputs[i];Ttot=Ttot)
    y2[i]=f_lif_semianalitic(inputs[i])
  end
  plot(inputs,[y1 y2];
     label=["simulation" "Fokker-Planck"],
     linewidth=2,marker=:star,leg=:bottomright,
     xlabel="input current",
     ylabel="rate", title="exp IF model")
end

savefig("/tmp/tmp.png")


module NKModel

#%% cell
#=
# Solving a New Keynesian model with Julia

This file is part of a computational appendix that accompanies the paper.

> MATLAB, Python, Julia: What to Choose in Economics?
>
> Coleman, Lyon, Maliar, and Maliar (2017)

For some details regarding the model solved in this file, please see the
[companion notebook that describes the model](http://bookshelf.quantecon.org/submission/59fa1b45145fc3772b0cef82).

In order to run the codes in this file you will need to install and configure a
few Julia packages. We recommend downloading
[JuliaPro](https://juliacomputing.com/products/juliapro.html) and/or following
the instructions on
[quantecon.org](https://lectures.quantecon.org/py/getting_started.html).

Once your Julia installation is up and running, there are a few additional packages
you will need in order to run the code here. We use the
[InstantiateFromURL](https://github.com/QuantEcon/InstantiateFromURL.jl) pacakge to make
it easy to get all required packages set up in your Julia environment.  If you have not
previously installed the `InstantiateFromURL` package you will need to uncomment the line
`pkg"add InstantiateFromURL"` and run the cell below.
=#

#%% cell
using Pkg
# pkg"add InstantiateFromURL"
using InstantiateFromURL: activate_github_path
activate_github_path("sglyon/CLMMJuliaPythonMatlab", path="NewKeynesian/julia", activate=true, force=true)

#%% cell
#=
## Julia Code

The Julia version of our algorithm is implemented as a few functions defined on
a core type named `Model`. This type is itself composed of three different types
that hold the model parameters, steady state, and grids needed to describe the
numerical model. Before we get to the types, we need to bring in some
dependencies:
=#

#%% cell
# for plotting the output
using PlotlyJS

# for constructing Sobol sequences
using Sobol
# for basis matrix of complete monomials and monomial quadrature rules
using BasisMatrices, QuantEcon
using MAT
using Parameters

# standard library components
using Printf  # printing messages
using Statistics  # computing `mean`
using Random: seed!
using LinearAlgebra: diagm, cholesky, dot
using InteractiveUtils: versioninfo
using DelimitedFiles: writedlm

seed!(42);  # set seed for reproducibility

include("model.jl")
include("solution.jl")

#%% cell
#=
### Running the code

Now that we've done all the hard work to define the model, its solution and
simulation, and accuracy checks, let's put things together and run the code!
=#

#%% cell
function ensurefile(url, localpath)
    if !isfile(localpath)
        println("Downloading $url to $localpath")
        download(url, localpath)
    end
end

function main(m=Model(); io::IO=stdout)
    ensurefile("https://github.com/sglyon/CLMMJuliaPythonMatlab/raw/master/NewKeynesian/Sobol_grids.mat", "Sobol_grids.mat")
    ensurefile("https://github.com/sglyon/CLMMJuliaPythonMatlab/raw/master/NewKeynesian/epsi_test_NK.mat", "epsi_test_NK.mat")
    ensurefile("https://github.com/sglyon/CLMMJuliaPythonMatlab/raw/master/NewKeynesian/random_grids.mat", "random_grids.mat")
    coefs, solve_time = solve(m);

    # simulate the model
    ts = time(); sim = Simulation(m, coefs); simulation_time = time() - ts

    # check accuracy
    tr = time(); resids = Residuals(m, coefs, sim); resids_time = time() - tr

    err_by_eq = max_E(resids)
    l1 = mean(resids)
    l∞ = max(resids)
    tot_time = solve_time + simulation_time + resids_time
    round3(x) = round(x, digits=3)
    round2(x) = round(x, digits=2)

    println(io, "Solver time (in seconds): $(solve_time)")
    println(io, "Simulation time (in seconds): $(simulation_time)")
    println(io, "Residuals time (in seconds): $(resids_time)")
    println(io, "total time (in seconds): $(tot_time)")
    println(io, "\nAPPROXIMATION ERRORS (log10):");
    println(io, "\ta) mean error in the model equations: $(round3(l1))");
    println(io, "\tb) max error in the model equations: $(round3(l∞))");
    println(io, "\tc) max error by equation:$(round3.(err_by_eq))");

    solve_time, simulation_time, resids_time, coefs, sim, resids

    t = 1:100
    p1 = plot([scatter(x=t, y=sim.S[t], name="S"), scatter(x=t, y=sim.F[t], name="F")],
              Layout(title="Figure 1a. S and F"))

    p2 = plot([scatter(x=t, y=sim.Y[t], name="Y"), scatter(x=t, y=sim.Yn[t], name="Yn")],
              Layout(title="Figure 1b. Output and natural output"))

    p3 = plot([scatter(x=t, y=sim.C[t], name="C"), scatter(x=t, y=sim.L[t], name="L")],
              Layout(title="Figure 1c. Consumption and labor"))

    p4 = plot([scatter(x=t, y=sim.δ[t], name="δ"), scatter(x=t, y=sim.R[t], name="R"),
               scatter(x=t, y=sim.π[t], name="π")],
              Layout(title="Figure 1d. Distortion, interest rate and inflation"))

    p = [p1 p2; p3 p4]
    p.plot.layout["width"] = 1000
    p.plot.layout["height"] = 600

    p, solve_time, simulation_time, resids_time, coefs, sim, resids, l1, l∞
end

#%% cell
function build_paper_table()
    # call main once to precompile all routines
    main()

    to_csv = []
    push!(to_csv, ["pi_star", "zlb", "degree", "solve_time", "l_1", "l_inf"])

    this_dir = @__DIR__

    open(joinpath(this_dir, "output.log"), "w") do f
        for params in [Dict(:πstar=>1.0, :zlb=>false),
                       Dict(:πstar=>1.0, :zlb=>true),
                       Dict(:πstar=>1 + 0.0598/4, :zlb=>false)]
            for deg in 1:5
                println("Starting $params and deg=$deg")
                m = Model(;grid_kind=:sobol, deg=deg, params...)

                println(f, "Starting $params and deg=$deg")
                solve_time, l1, linf = main(m, io=f)[[2, 8, 9]]
                push!(to_csv, [params[:πstar], params[:zlb], deg, solve_time, l1, linf])
                @show to_csv
                println(f, "\n"^5)
                flush(f)
            end
        end
        versioninfo(f, verbose=true)
    end

    writedlm(
        joinpath(this_dir, "output_osx.csv"),
        permutedims(hcat(to_csv...), (2, 1)), ','
    )
end

#%% cell
# if "table" in ARGS
#     build_paper_table()
# else
#     results = main()
#     results[1]  # show plot
# end

end  # module


#tests for "nuisance" inversions with HiQGA.transD_GP

## make options for a purely stationary properties GP
nmin, nmax = 2, 200
pnorm = 2.
λx, λy = 1, 1
x = 0:(0.005λx):λx
y = 0:(0.005λy):0.5λy
nmin, nmax = 2, 200
fbounds = [-2. 2]
δ = 0.1
sdev_prop = [0.1]
sdev_pos = [0.05, 0.05]
K = transD_GP.GP.Mat32()
λ = [0.05maximum(x), 0.05maximum(y)]
demean = false
sampledc = true
##
xall = zeros(2,length(x)*length(y))
for i in 1:size(xall,2)
    xid, yid = Tuple(CartesianIndices((length(x),length(y)))[i])
    xall[:,i] = [x[xid]; y[yid]]
end
xbounds = zeros(Float64,size(xall, 1), 2)
for dim in 1:size(xall, 1)
    xbounds[dim,:] = [minimum(xall[dim,:]), maximum(xall[dim,:])]
end

opt = transD_GP.OptionsStat(nmin = nmin,
                        nmax = nmax,
                        xbounds = xbounds,
                        fbounds = fbounds,
                        xall = xall,
                        λ = λ,
                        δ = δ,
                        demean = demean,
                        sampledc = sampledc,
                        sdev_prop = sdev_prop,
                        sdev_pos = sdev_pos,
                        pnorm = pnorm,
                        needλ²fromlog = false,
                        updatenonstat = false,
                        quasimultid = false,
                        K = K,
                        timesλ = 3,
                        )

@time m = transD_GP.init(opt)

@testset "Creation tests for OptionsNuisance" begin
    nu_sdev = [0.1, 0.2, 0.3]
    nu_bounds = [0. 4.; -1.5 1.0; -1. 0.]
    optn = transD_GP.OptionsNuisance(opt;
        sdev = nu_sdev,
        bounds = nu_bounds,
        updatenuisances = true,
    )
    @testset "Test parameters set properly" begin
        @test isapprox(optn.sdev, [0.4, 0.5, 0.3], rtol=1e-4) #sdev in the struct is in "absolute units"
        @test all(optn.bounds .== nu_bounds)
        @test isapprox(optn.rotatebounds, nu_bounds .- mean(nu_bounds, dims=2), rtol=1e-4) #no covariance matrix was supplied
        @test optn.nnu == 3
        @test optn.updatenuisances == true
    end
    @testset "Constructor does not modify arguments" begin
        @test all(nu_sdev .== [0.1, 0.2, 0.3])
        @test all(nu_bounds .== [0. 4.; -1.5 1.0; -1. 0.])
    end
end


import Base: iterate, elsize, size, getindex, setindex!, parent

abstract type AbstractDiagonalIterator{T, V}  <: AbstractVector{T} end

"""
    DiagonalIterator{T <: AbstractMatrix}

Iterator over the diagonal of an AbstractMatrix
"""
struct DiagonalIterator{T, V <:AbstractMatrix} <: AbstractDiagonalIterator{T, V}
    M::V
    length::Int
end

"""
    altdiagonal(M::AbstractMatrix)

An iterator over the diagonal of `M`. `altdiagonal` is an alternative
to the iterator `diagonal`.  An iterator is constructed much more quickly
with `altdiagonal` than with `diagonal`.
"""
altdiagonal(M::V) where V <: AbstractMatrix  = DiagonalIterator{eltype(M), V}(M, checksquare(M))

iterate(iter::DiagonalIterator, i = 1) = i > iter.length ? nothing : (iter.M[i,i], i + 1)
size(iter::DiagonalIterator) = (iter.length,)
elsize(::Type{DiagonalIterator{T, V}}) where {T, V} = sizeof(T)
getindex(iter::DiagonalIterator, i) = iter.M[i,i]
setindex!(iter::DiagonalIterator, v, i) = (iter.M[i,i] = v)
parent(iter::DiagonalIterator) = iter.M
Base.dataids(d::DiagonalIterator) = Base.dataids(parent(d))


using Flux
using Flux.Tracker

f(x) = 3x^2 + 2x + 1
df(x) = Tracker.gradient(f, x; nest = true)[1]
df(3)

d2f(x) = Tracker.gradient(df, x; nest = true)[1]
d2f(3)

f(W, b, x) = W*x + b
Tracker.gradient(f, 2, 3, 4)

W = param(2) # 2.0 (tracked)
b = param(3) # 3.0 (tracked)

f(x) = W * x + b

grads = Tracker.gradient(() -> f(4), params(W, b))

(grads[W] , grads[b])

W = rand(2, 5)
b = rand(2)

predict(x) = W*x .+ b

function loss(x, y)
  ŷ = predict(x)
  sum((y .- ŷ).^2)
end

x, y = rand(5), rand(2) # Dummy data
loss(x, y) # ~ 3

W = param(W)
b = param(b)

gs = Tracker.gradient(() -> loss(x, y), params(W, b))

gs[W]

gs[b]

using Flux.Tracker: update!

D = gs[W]
update!(W, -0.1D)

loss(x, y)

using Flux, Flux.Data.MNIST, Statistics
using Flux: onehotbatch, onecold, crossentropy, throttle
using Base.Iterators: repeated

# Classify MNIST digits with a simple multi-layer-perceptron

imgs = MNIST.images()

# Stack images into one large batch
X = hcat(float.(reshape.(imgs, :))...)

labels = MNIST.labels()

# One-hot-encode the labels
Y = onehotbatch(labels, 0:9)

m = Chain(
  Dense(28^2, 32, relu),
  Dense(32, 10),
  softmax)

loss(x, y) = crossentropy(m(x), y)
accuracy(x, y) = mean(onecold(m(x)) .== onecold(y))

dataset = repeated((X, Y), 200)
evalcb = () -> @show(loss(X, Y))
opt = ADAM()

Flux.train!(loss, params(m), dataset, opt, cb = throttle(evalcb, 10))
accuracy(X, Y)

# Test set accuracy
tX = hcat(float.(reshape.(MNIST.images(:test), :))...) |> gpu
tY = onehotbatch(MNIST.labels(:test), 0:9) |> gpu

accuracy(tX, tY)

using Metalhead
using Metalhead: classify
using Images

# Switch to dataset images and classify the known (1000) images
cd(homedir()*"/PacktPub/Datasets/Images")
vgg = VGG19()

img01 = load("elephant.jpg")
classify(vgg, img01)

using Unicode

img02 = load("harry.jpg")
animal = uppercase(classify(vgg, img02))
println("Meow, I am a $animal perhaps ???'")

animal_types = Metalhead.labels(vgg)
animal_types[1:5]

# How sure are we?
#
probs = Metalhead.forward(vgg, img02)
sort(probs)[996:1000]

perm = sortperm(probs)
animal_types[perm[1000]]

animal_types[perm[999]]   # <= This is a breed of DOG. sorry Harry.

animal_types[287:288]    # If not a cat -- why not one of these ?

pwd(); include("KnetJL/mnist-mlp/mlp.jl")

MLP.main("--epochs 15")

using Knet, Images
cd(Knet.dir("examples"))
pwd()

include("mnist-mlp/mlp.jl")
MLP.main("--epochs 15")

xtrn,ytrn,xtst,ytst = MLP.mnist()
map(summary,(xtrn,ytrn,xtst,ytst))

using Images
mnistview(x,i) = Images.colorview(Gray,permutedims(x[:,:,1,i],(2,1)))

hcat([mnistview(xtst,i) for i=1:5]...)

ytst[1:5]'

include("fashion-mnist/fashion-mnist.jl")

FashionMNIST.main("--epochs 15")

include("../data/fashion-mnist.jl")

xtrn,ytrn,xtst,ytst = fmnist()
map(maximum,(xtrn,ytrn,xtst,ytst))

fmnistview(x,i) = Images.colorview(Gray,permutedims(x[:,:,1,i],(2,1)))
hcat([fmnistview(xtst,i) for i=1:10]...)

ytst[1:10]'

xtst


"""
Pinched cylinder with diagphram supports and concentrated force
"""
module pinched_cylinder_examples

using FinEtools
using FinEtoolsDeforLinear
using FinEtoolsFlexStructures.FESetShellT3Module: FESetShellT3
using FinEtoolsFlexStructures.FEMMShellT3FFModule
using FinEtoolsFlexStructures.RotUtilModule: initial_Rfield, update_rotation_field!
using VisualStructures: plot_nodes, plot_midline, render, plot_space_box, plot_midsurface, space_aspectratio, save_to_json

function _execute(n = 2, visualize = true)
    E = 3e6;
    nu = 0.3;
    thickness = 3.0;
    # analytical solution for the vertical deflection under the load
    analyt_sol=-1.82488e-5;
    formul = FEMMShellT3FFModule

    @info "Mesh: $n elements per side"

    # Mesh
    R = 300.0;
    L = 600.0;

    tolerance = R/n/1000
    fens, fes = T3block(90/360*2*pi,L/2,n,n);
    fens.xyz = xyz3(fens)
    for i in 1:count(fens)
        a=fens.xyz[i, 1]; y=fens.xyz[i, 2];
        fens.xyz[i, :] .= (R*sin(a), y, R*cos(a))
    end

    mater = MatDeforElastIso(DeforModelRed3D, E, nu)

    sfes = FESetShellT3()
    accepttodelegate(fes, sfes)
    femm = formul.make(IntegDomain(fes, TriRule(1), thickness), mater)
    femm.drilling_stiffness_scale = 0.1
    associategeometry! = formul.associategeometry!
    stiffness = formul.stiffness

    # Construct the requisite fields, geometry and displacement
    # Initialize configuration variables
    geom0 = NodalField(fens.xyz)
    u0 = NodalField(zeros(size(fens.xyz,1), 3))
    Rfield0 = initial_Rfield(fens)
    dchi = NodalField(zeros(size(fens.xyz,1), 6))

    # Apply EBC's
    # rigid diaphragm
    l1 = selectnode(fens; box = Float64[-Inf Inf 0 0 -Inf Inf], inflate = tolerance)
    for i in [1,3]
        setebc!(dchi, l1, true, i)
    end
    # plane of symmetry perpendicular to Y
    l1 = selectnode(fens; box = Float64[-Inf Inf L/2 L/2 -Inf Inf], inflate = tolerance)
    for i in [2,4,6]
        setebc!(dchi, l1, true, i)
    end
    # plane of symmetry perpendicular to X
    l1 = selectnode(fens; box = Float64[0 0 -Inf Inf -Inf Inf], inflate = tolerance)
    for i in [1,5,6]
        setebc!(dchi, l1, true, i)
    end
    # plane of symmetry perpendicular to Z
    l1 = selectnode(fens; box = Float64[-Inf Inf -Inf Inf 0.0 0.0], inflate = tolerance)
    for i in [3,4,5]
        setebc!(dchi, l1, true, i)
    end
    applyebc!(dchi)
    numberdofs!(dchi);

    # Assemble the system matrix
    associategeometry!(femm, geom0)
    K = stiffness(femm, geom0, u0, Rfield0, dchi);

    # Load
    nlc = selectnode(fens; box = Float64[0 0 L/2 L/2 -Inf Inf], inflate = tolerance)
    # If desired, increase the tolerance, for instance to half of the thickness.
    nl = selectnode(fens; box = Float64[0 0 L/2 L/2 -Inf Inf], inflate = tolerance)
    nln = length(nl)
    loadbdry = FESetP1(reshape(nl, nln, 1))
    lfemm = FEMMBase(IntegDomain(loadbdry, PointRule()))
    fi = ForceIntensity(FFlt[0, 0, -1/4/nln, 0, 0, 0]);
    F = distribloads(lfemm, geom0, dchi, fi, 3);

    # Solve
    U = K\F
    scattersysvec!(dchi, U[:])
    @show n,  dchi.values[nlc, 3][1]/analyt_sol*100

    # Visualization
    if visualize
        scattersysvec!(dchi, (L/8)/maximum(abs.(U)).*U)
        update_rotation_field!(Rfield0, dchi)
        plots = cat(plot_space_box([[0 0 -L/2]; [L/2 L/2 L/2]]),
        #plot_nodes(fens),
            plot_midsurface(fens, fes; x = geom0.values, u = dchi.values[:, 1:3], R = Rfield0.values);
            dims = 1)
        pl = render(plots)
    end
    return dchi.values[nlc, 3][1]/analyt_sol*100
end

function test_convergence()
    @info "Pinched cylinder"
    ns = [4, 8, 16, 24, 32]
    results = Float64[]
    for n in ns
        v = _execute(n, false)
        push!(results, v)
    end
    return ns, results
end

end # module

using .pinched_cylinder_examples
ns, results = pinched_cylinder_examples.test_convergence()

# [ Info: Mesh: 4 elements per side
# (n, (dchi.values[nl, 3] / analyt_sol) * 100) = (4, [41.55321964302738])
# [ Info: Mesh: 8 elements per side
# (n, (dchi.values[nl, 3] / analyt_sol) * 100) = (8, [76.87858207030617])
# [ Info: Mesh: 12 elements per side
# (n, (dchi.values[nl, 3] / analyt_sol) * 100) = (12, [87.73634899399313])
# [ Info: Mesh: 16 elements per side
# (n, (dchi.values[nl, 3] / analyt_sol) * 100) = (16, [92.41804796585181])
# [ Info: Mesh: 24 elements per side
# (n, (dchi.values[nl, 3] / analyt_sol) * 100) = (24, [96.49232304717079])
# [ Info: Mesh: 48 elements per side
# (n, (dchi.values[nl, 3] / analyt_sol) * 100) = (48, [99.6766661386644])
# [ Info: Mesh: 96 elements per side
# (n, (dchi.values[nl, 3] / analyt_sol) * 100) = (96, [100.81626349991888])
# [ Info: Mesh: 128 elements per side
# (n, (dchi.values[nl, 3] / analyt_sol) * 100) = (128, [101.05826709146844])
# [ Info: Mesh: 256 elements per side
# (n, (dchi.values[nl, 3] / analyt_sol) * 100) = (256, [101.43102704824867])
# [ Info: Mesh: 400 elements per side
# (n, (dchi.values[nl, 3] / analyt_sol) * 100) = (400, [101.6018373070541])
# [ Info: Mesh: 600 elements per side
# (n, (dchi.values[nl, 3] / analyt_sol) * 100) = (600, [101.73923429243655])

# These results come from Shin, Lee (2014), Table 3.

# Mesh Subdivision
# 4, 8, 16, 24, 32
Allman = [
59.02
92.41
100.43
100.53
100.29
] .* (1.8541/1.82488)
Cook_flat_stiffened = [
53.71
89.73
99.67
100.19
100.10
] .* (1.8541/1.82488)
Providas_Kattis = [
45.35
85.64
98.24
99.52
99.74
] .* (1.8541/1.82488)

using PGFPlotsX

objects = []

all_results = [("Present", results, "*"), ("Allman", Allman, "x"), ("Providas, Kattis", Providas_Kattis, "triangle"), ("Cook, flat, stiffened", Cook_flat_stiffened, "square")]

for r in  all_results
    @pgf p = PGFPlotsX.Plot(
    {
    color = "black",
    line_width  = 0.7,
    style = "solid",
    mark = "$(r[3])"
    },
    Coordinates([v for v in  zip(ns, r[2])])
    )
    push!(objects, p)
    push!(objects, LegendEntry("$(r[1])"))
end

@pgf ax = Axis(
    {
        xlabel = "Number of Elements / side [ND]",
        ylabel = "Normalized Displacement [ND]",
        # xmin = range[1],
        # xmax = range[2],
        xmode = "linear",
        ymode = "linear",
        yminorgrids = "true",
        grid = "both",
        legend_style = {
            at = Coordinate(0.5, 1.05),
            anchor = "south",
            legend_columns = -1
        },
    },
    objects...
)

display(ax)
pgfsave("pinched_cylinder-convergence.pdf", ax)


export dcat

"""
`dcat(A,B)` diagonally concatenate matrices `A` and `B`.
Returns a new matrix of the form `[A 0; 0 B]` where the `0`s
are blocks of zeros of the approriate shape.

More generally, may be called with any positive number of matrix, vector,
or scalar (numeric) arguments.
"""
function dcat(A::Array{S,2}, B::Array{T,2}) where {S,T}
    ST = promote_type(S, T)

    ra, ca = size(A)
    rb, cb = size(B)

    Z1 = zeros(ST, ra, cb)
    Z2 = zeros(ST, rb, ca)

    return [A Z1; Z2 B]
end

dcat(A) where {S} = _matrixify(A)

function dcat(A, B, C...)
    AB = dcat(A, B)
    return dcat(AB, C...)
end

function _matrixify(v::Array{T,1}) where {T}
    n = length(v)
    return reshape(v, n, 1)
end

_matrixify(x::T) where {T<:Number} = _matrixify([x])

_matrixify(A::Array{T,2}) where {T} = A

function dcat(v::T, A::Matrix{S}) where {S,T}
    return dcat(_matrixify(v), A)
end

function dcat(A::Matrix{S}, v::T) where {S,T}
    return dcat(A, _matrixify(v))
end

function dcat(x::T, y::S) where {S,T}
    return dcat(_matrixify(x), _matrixify(y))
end


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

function twolayer(y, x, w1, w2, b1, b2)
    f = x -> begin
        w1 = reshape(w1, 10, 10)
        w2 = reshape(w2, 10, 10)
        z = w2*tanh.(w1*x+b1)+b2
    end
    y[:] = ForwardDiff.jacobian(f, x)[:]
end
two_layer = load_op("build/libTwoLayer", "two_layer")

w1 = rand(100)
w2 = rand(100)
b1 = rand(10)
b2 = rand(10)
x = rand(10)
J = rand(100)
twolayer(J, x, w1, w2, b1, b2)

y = two_layer(constant(x), constant(w1), constant(b1), constant(w2), constant(b2))
sess = Session(); init(sess)
J0 = run(sess, y)
@show norm(J-J0)


# ===============================
# Written by AAE
# TU Delft, Spring 2014
# simulkade.com
# ===============================

# ================================================================
# Changes:
#    2014-12-30 added 2D radial and 3D cylindrical grids
#    2014-12-31 many many bugs fixed
#    2015-01-10 extended to accept nonuniform grids
# ================================================================

function divergenceTerm(F::FaceValue)
d = F.domain.dimension
if d==1
  RHSdiv = divergenceTerm1D(F)
elseif d==1.5
  RHSdiv = divergenceTermCylindrical1D(F)
elseif d==2
  RHSdiv, RHSdivx, RHSdivy = divergenceTerm2D(F)
elseif d==2.5
  RHSdiv, RHSdivx, RHSdivy = divergenceTermCylindrical2D(F)
elseif d==2.8
  RHSdiv, RHSdivx, RHSdivy = divergenceTermRadial2D(F)
elseif d==3
  RHSdiv, RHSdivx, RHSdivy, RHSdivz = divergenceTerm3D(F)
elseif d==3.2
  RHSdiv, RHSdivx, RHSdivy, RHSdivz = divergenceTermCylindrical3D(F)
end
RHSdiv

end

# =============== Divergence 1D Term ============================
function divergenceTerm1D(F::FaceValue)
# This function calculates the divergence of a field
# using its face

# extract data from the mesh structure
Nx = F.domain.dims[1]
G = [1:Nx+2;]
DX = F.domain.cellsize.x[2:end-1]

# define the vector of cell index
row_index = reshape(G[2:Nx+1],Nx) # main diagonal

# compute the divergence
div_x = (F.xvalue[2:Nx+1]-F.xvalue[1:Nx])./DX

# define the RHS Vector
RHSdiv = zeros(Nx+2)

# assign the values of the RHS vector
RHSdiv[row_index] = reshape(div_x,Nx)

return RHSdiv

end

# =============== Divergence Cylindrical 1D Term ============================
function divergenceTermCylindrical1D(F::FaceValue)
# This function calculates the divergence of a field
# using its face

# extract data from the mesh structure
Nx = F.domain.dims[1]
G = [1:Nx+2;]
DX = F.domain.cellsize.x[2:end-1]
rp = F.domain.cellcenters.x
rf = F.domain.facecenters.x

# define the vector of cell index
row_index = reshape(G[2:Nx+1],Nx) # main diagonal

# reassign the east, west, north, and south flux vectors for the
# code readability
Fe = F.xvalue[2:Nx+1]
Fw = F.xvalue[1:Nx]
re = rf[2:Nx+1]
rw = rf[1:Nx]

# compute the divergence
div_x = (re.*Fe-rw.*Fw)./(DX.*rp)

# define the RHS Vector
RHSdiv = zeros(Nx+2)

# assign the values of the RHS vector
RHSdiv[row_index] = reshape(div_x,Nx)

return RHSdiv

end

# =============== Divergence 2D Term ============================
function divergenceTerm2D(F::FaceValue)
# This function calculates the divergence of a field
# using its face

# extract data from the mesh structure
Nx = F.domain.dims[1]
Ny = F.domain.dims[2]
G=reshape([1:(Nx+2)*(Ny+2);], Nx+2, Ny+2)
DX = F.domain.cellsize.x[2:end-1]
DY = zeros( 1, Ny)
DY[:] = F.domain.cellsize.y[2:end-1]

# define the vector of cell index
row_index = reshape(G[2:Nx+1,2:Ny+1],Nx*Ny) # main diagonal

# reassign the east, west, north, and south flux vectors for the
# code readability
Fe = F.xvalue[2:Nx+1,:]
Fw = F.xvalue[1:Nx,:]
Fn = F.yvalue[:,2:Ny+1]
Fs = F.yvalue[:,1:Ny]

# compute the divergence
div_x = (Fe - Fw)./DX
div_y = (Fn - Fs)./DY

# define the RHS Vector
RHSdiv = zeros((Nx+2)*(Ny+2))
RHSdivx = zeros((Nx+2)*(Ny+2))
RHSdivy = zeros((Nx+2)*(Ny+2))

# assign the values of the RHS vector
RHSdiv[row_index] = reshape(div_x+div_y,Nx*Ny)
RHSdivx[row_index] = reshape(div_x,Nx*Ny)
RHSdivy[row_index] = reshape(div_y,Nx*Ny)

(RHSdiv, RHSdivx, RHSdivy)

end

# =============== Divergence 2D Cylindrical Term ============================
function divergenceTermCylindrical2D(F::FaceValue)
# This function calculates the divergence of a field
# using its face

# extract data from the mesh structure
Nr = F.domain.dims[1]
Nz = F.domain.dims[2]
G=reshape([1:(Nr+2)*(Nz+2);], Nr+2, Nz+2)
dr = F.domain.cellsize.x[2:end-1]
dz= zeros( 1, Nz)
dz[:] = F.domain.cellsize.y[2:end-1]
rp = F.domain.cellcenters.x
rf = F.domain.facecenters.x

# define the vector of cell index
row_index = reshape(G[2:Nr+1,2:Nz+1],Nr*Nz) # main diagonal

# reassign the east, west, north, and south flux vectors for the
# code readability
Fe = F.xvalue[2:Nr+1,:]
Fw = F.xvalue[1:Nr,:]
Fn = F.yvalue[:,2:Nz+1]
Fs = F.yvalue[:,1:Nz]
re = rf[2:Nr+1]
rw = rf[1:Nr]

# compute the divergence
div_x = (re.*Fe - rw.*Fw)./(dr.*rp)
div_y = (Fn - Fs)./dz

# define the RHS Vector
RHSdiv = zeros((Nr+2)*(Nz+2))
RHSdivx = zeros((Nr+2)*(Nz+2))
RHSdivy = zeros((Nr+2)*(Nz+2))

# assign the values of the RHS vector
RHSdiv[row_index] = reshape(div_x+div_y,Nr*Nz)
RHSdivx[row_index] = reshape(div_x,Nr*Nz)
RHSdivy[row_index] = reshape(div_y,Nr*Nz)

(RHSdiv, RHSdivx, RHSdivy)

end

# =============== Divergence 2D Radial Term ============================
function divergenceTermRadial2D(F::FaceValue)
# This function calculates the divergence of a field
# using its face

# extract data from the mesh structure
Nr = F.domain.dims[1]
Ntheta = F.domain.dims[2]
G=reshape([1:(Nr+2)*(Ntheta+2);], Nr+2, Ntheta+2)
dr = F.domain.cellsize.x[2:end-1]
dtheta= zeros( 1, Ntheta)
dtheta[:]= F.domain.cellsize.y[2:end-1]
rp = F.domain.cellcenters.x
rf = F.domain.facecenters.x

# define the vector of cell index
row_index = reshape(G[2:Nr+1,2:Ntheta+1],Nr*Ntheta) # main diagonal

# reassign the east, west, north, and south flux vectors for the
# code readability
Fe = F.xvalue[2:Nr+1,:]
Fw = F.xvalue[1:Nr,:]
Fn = F.yvalue[:,2:Ntheta+1]
Fs = F.yvalue[:,1:Ntheta]
re = rf[2:Nr+1]
rw = rf[1:Nr]

# compute the divergence
div_x = (re.*Fe-rw.*Fw)./(dr.*rp)
div_y = (Fn-Fs)./(dtheta.*rp)

# define the RHS Vector
RHSdiv = zeros((Nr+2)*(Ntheta+2))
RHSdivx = zeros((Nr+2)*(Ntheta+2))
RHSdivy = zeros((Nr+2)*(Ntheta+2))

# assign the values of the RHS vector
RHSdiv[row_index] = reshape(div_x+div_y,Nr*Ntheta)
RHSdivx[row_index] = reshape(div_x,Nr*Ntheta)
RHSdivy[row_index] = reshape(div_y,Nr*Ntheta)

(RHSdiv, RHSdivx, RHSdivy)

end

# =============== Divergence 3D Term ============================
function divergenceTerm3D(F::FaceValue)
# This function calculates the divergence of a field
# using its face

# extract data from the mesh structure
Nx = F.domain.dims[1]
Ny = F.domain.dims[2]
Nz = F.domain.dims[3]
G=reshape([1:(Nx+2)*(Ny+2)*(Nz+2);], Nx+2, Ny+2, Nz+2)
dx = F.domain.cellsize.x[2:end-1]
dy = zeros( 1, Ny)
dy[:] = F.domain.cellsize.y[2:end-1]
dz = zeros( 1,1,Nz)
dz[:] = F.domain.cellsize.z[2:end-1]

# define the vector of cell index
row_index = reshape(G[2:Nx+1,2:Ny+1,2:Nz+1],Nx*Ny*Nz) # main diagonal

# reassign the east, west, north, and south flux vectors for the
# code readability
Fe = F.xvalue[2:Nx+1,:,:]
Fw = F.xvalue[1:Nx,:,:]
Fn = F.yvalue[:,2:Ny+1,:]
Fs = F.yvalue[:,1:Ny,:]
Ff = F.zvalue[:,:,2:Nz+1]
Fb = F.zvalue[:,:,1:Nz]

# compute the divergence
div_x = (Fe - Fw)./dx
div_y = (Fn - Fs)./dy
div_z = (Ff - Fb)./dz

# define the RHS Vector
RHSdiv = zeros((Nx+2)*(Ny+2)*(Nz+2))
RHSdivx = zeros((Nx+2)*(Ny+2)*(Nz+2))
RHSdivy = zeros((Nx+2)*(Ny+2)*(Nz+2))
RHSdivz = zeros((Nx+2)*(Ny+2)*(Nz+2))

# assign the values of the RHS vector
RHSdiv[row_index] = reshape(div_x+div_y+div_z,Nx*Ny*Nz)
RHSdivx[row_index] = reshape(div_x,Nx*Ny*Nz)
RHSdivy[row_index] = reshape(div_y,Nx*Ny*Nz)
RHSdivz[row_index] = reshape(div_z,Nx*Ny*Nz)

(RHSdiv, RHSdivx, RHSdivy, RHSdivz)

end

# =============== Divergence 3D Cylindrical Term ============================
function divergenceTermCylindrical3D(F::FaceValue)
# This function calculates the divergence of a field
# using its face

# extract data from the mesh structure
Nx = F.domain.dims[1]
Ny = F.domain.dims[2]
Nz = F.domain.dims[3]
G=reshape([1:(Nx+2)*(Ny+2)*(Nz+2);], Nx+2, Ny+2, Nz+2)
dx = F.domain.cellsize.x[2:end-1]
dy = zeros( 1, Ny)
dy[:] = F.domain.cellsize.y[2:end-1]
dz = zeros( 1,1,Nz)
dz[:] = F.domain.cellsize.z[2:end-1]
rp = F.domain.cellcenters.x

# define the vector of cell index
row_index = reshape(G[2:Nx+1,2:Ny+1,2:Nz+1],Nx*Ny*Nz) # main diagonal

# reassign the east, west, north, and south flux vectors for the
# code readability
Fe = F.xvalue[2:Nx+1,:,:]
Fw = F.xvalue[1:Nx,:,:]
Fn = F.yvalue[:,2:Ny+1,:]
Fs = F.yvalue[:,1:Ny,:]
Ff = F.zvalue[:,:,2:Nz+1]
Fb = F.zvalue[:,:,1:Nz]

# compute the divergence
div_x = (Fe - Fw)./dx
div_y = (Fn - Fs)./(dy.*rp)
div_z = (Ff - Fb)./dz

# define the RHS Vector
RHSdiv = zeros((Nx+2)*(Ny+2)*(Nz+2))
RHSdivx = zeros((Nx+2)*(Ny+2)*(Nz+2))
RHSdivy = zeros((Nx+2)*(Ny+2)*(Nz+2))
RHSdivz = zeros((Nx+2)*(Ny+2)*(Nz+2))

# assign the values of the RHS vector
RHSdiv[row_index] = reshape(div_x+div_y+div_z,Nx*Ny*Nz)
RHSdivx[row_index] = reshape(div_x,Nx*Ny*Nz)
RHSdivy[row_index] = reshape(div_y,Nx*Ny*Nz)
RHSdivz[row_index] = reshape(div_z,Nx*Ny*Nz)

(RHSdiv, RHSdivx, RHSdivy, RHSdivz)

end

# ===================== Gradient Term ============================
function gradientTerm(phi)
  # calculates the gradient of a variable
  # the output is a face variable
  d=phi.domain.dimension
  if d==1 || d==1.5
    dx = 0.5*(phi.domain.cellsize.x[1:end-1]+phi.domain.cellsize.x[2:end])
    FaceValue(phi.domain,
      (phi.value[2:end]-phi.value[1:end-1])./dx,
      [1.0],
      [1.0])
  elseif d==2 || d==2.5
    dx = 0.5*(phi.domain.cellsize.x[1:end-1]+phi.domain.cellsize.x[2:end])
    Ny = phi.domain.dims[2]
    dy = zeros( 1, Ny+1)
    dy[:] = 0.5*(phi.domain.cellsize.y[1:end-1]+phi.domain.cellsize.y[2:end])
    FaceValue(phi.domain,
      (phi.value[2:end,2:end-1]-phi.value[1:end-1,2:end-1])./dx,
      (phi.value[2:end-1,2:end]-phi.value[2:end-1,1:end-1])./dy,
      [1.0])
  elseif d==2.8
    dx = 0.5*(phi.domain.cellsize.x[1:end-1]+phi.domain.cellsize.x[2:end])
    Ntheta = phi.domain.dims[2]
    dtheta = zeros( 1, Ntheta+1)
    dtheta[:] = 0.5*(phi.domain.cellsize.y[1:end-1]+phi.domain.cellsize.y[2:end])
    rp = phi.domain.cellcenters.x
    FaceValue(phi.domain,
      (phi.value[2:end,2:end-1]-phi.value[1:end-1,2:end-1])./dx,
      (phi.value[2:end-1,2:end]-phi.value[2:end-1,1:end-1])./(dtheta.*rp),
      [1.0])
  elseif d==3
    Ny = phi.domain.dims[2]
    Nz = phi.domain.dims[3]
    dx = 0.5*(phi.domain.cellsize.x[1:end-1]+phi.domain.cellsize.x[2:end])
    dy= zeros( 1, Ny+1)
    dy[:] = 0.5*(phi.domain.cellsize.y[1:end-1]+phi.domain.cellsize.y[2:end])
    dz= zeros( 1, 1, Nz+1)
    dz[:] = 0.5*(phi.domain.cellsize.z[1:end-1]+phi.domain.cellsize.z[2:end])
    FaceValue(phi.domain,
      (phi.value[2:end,2:end-1,2:end-1]-phi.value[1:end-1,2:end-1,2:end-1])./dx,
      (phi.value[2:end-1,2:end,2:end-1]-phi.value[2:end-1,1:end-1,2:end-1])./dy,
      (phi.value[2:end-1,2:end-1,2:end]-phi.value[2:end-1,2:end-1,1:end-1])./dz)
  elseif d==3.2
    Ntheta = phi.domain.dims[2]
    Nz = phi.domain.dims[3]
    dx = 0.5*(phi.domain.cellsize.x[1:end-1]+phi.domain.cellsize.x[2:end])
    dy= zeros( 1, Ntheta+1)
    dy[:] = 0.5*(phi.domain.cellsize.y[1:end-1]+phi.domain.cellsize.y[2:end])
    dz= zeros( 1, 1, Nz+1)
    dz[:] = 0.5*(phi.domain.cellsize.z[1:end-1]+phi.domain.cellsize.z[2:end])
    rp = phi.domain.cellcenters.x
    FaceValue(phi.domain,
      (phi.value[2:end,2:end-1,2:end-1]-phi.value[1:end-1,2:end-1,2:end-1])./dx,
      (phi.value[2:end-1,2:end,2:end-1]-phi.value[2:end-1,1:end-1,2:end-1])./(dy.*rp),
      (phi.value[2:end-1,2:end-1,2:end]-phi.value[2:end-1,2:end-1,1:end-1])./dz)

  end
end

# ===================== Gradient Term ============================
function gradientCellTerm(phi)
  # calculates the gradient of a variable
  # the output is a face variable
  d=phi.domain.dimension
  phi_face = linearMean(phi)
  if d==1 || d==1.5
    dx = phi.domain.cellsize.x[2:end-1]
    CellVector(phi.domain,
      (phi_face.xvalue[2:end] .- phi_face.xvalue[1:end-1])./dx,
      [1.0],
      [1.0])
  elseif d==2 || d==2.5
    dx = phi.domain.cellsize.x[2:end-1]
    Ny = phi.domain.dims[2]
    dy = zeros(1, Ny)
    dy[:] = phi.domain.cellsize.y[2:end-1]
    CellVector(phi.domain,
      (phi_face.xvalue[2:end, :] .- phi_face.xvalue[1:end-1, :])./dx,
      (phi_face.yvalue[:, 2:end] .- phi_face.yvalue[:, 1:end-1])./dy,
      [1.0])
  elseif d==2.8
    dx = phi.domain.cellsize.x[2:end-1]
    Ntheta = phi.domain.dims[2]
    dtheta = zeros(1, Ntheta)
    dtheta[:] = phi.domain.cellsize.y[2:end-1]
    rp = phi.domain.cellcenters.x
    CellVector(phi.domain,
      (phi_face.xvalue[2:end, :] .- phi_face.xvalue[1:end-1, :])./dx,
      (phi_face.yvalue[:, 2:end] .- phi_face.yvalue[:, 1:end-1])./(dtheta.*rp),
      [1.0])
  elseif d==3
    Ny = phi.domain.dims[2]
    Nz = phi.domain.dims[3]
    dx = phi.domain.cellsize.x[2:end-1]
    dy= zeros(1, Ny)
    dy[:] = phi.domain.cellsize.y[2:end-1]
    dz= zeros( 1, 1, Nz)
    dz[:] = phi.domain.cellsize.z[2:end-1]
    CellVector(phi.domain,
      (phi_face.xvalue[2:end, :, :] .- phi_face.xvalue[1:end-1, :, :])./dx,
      (phi_face.yvalue[:, 2:end, :] .- phi_face.yvalue[:, 1:end-1, :])./dy,
      (phi_face.zvalue[:, :, 2:end] .- phi_face.zvalue[:, :, 1:end-1])./dz)
  elseif d==3.2
    Ntheta = phi.domain.dims[2]
    Nz = phi.domain.dims[3]
    dx = phi.domain.cellsize.x[2:end-1]
    dy= zeros(1, Ntheta)
    dy[:] = phi.domain.cellsize.y[2:end-1]
    dz= zeros(1, 1, Nz)
    dz[:] = phi.domain.cellsize.z[2:end-1]
    rp = phi.domain.cellcenters.x
    CellVector(phi.domain,
      (phi_face.xvalue[2:end, :, :] .- phi_face.xvalue[1:end-1, :, :])./dx,
      (phi_face.yvalue[:, 2:end, :] .- phi_face.yvalue[:, 1:end-1, :])./(dy.*rp),
      (phi_face.zvalue[:, :, 2:end] .- phi_face.zvalue[:, :, 1:end-1])./dz)
  end
end


# von Mises-Fisher distribution is useful for directional statistics
#
# The implementation here follows:
#
#   - Wikipedia:
#     http://en.wikipedia.org/wiki/Von_Mises–Fisher_distribution
#
#   - R's movMF package's document:
#     http://cran.r-project.org/web/packages/movMF/vignettes/movMF.pdf
#
#   - Wenzel Jakob's notes:
#     http://www.mitsuba-renderer.org/~wenzel/files/vmf.pdf
#

immutable VonMisesFisher <: ContinuousMultivariateDistribution
    μ::Vector{Float64}
    κ::Float64
    logCκ::Float64

    function VonMisesFisher(μ::Vector{Float64}, κ::Float64; checknorm::Bool=true)
        if checknorm
            isunitvec(μ) || error("μ must be a unit vector")
        end
        κ > 0 || error("κ must be positive.")
        new(μ, κ, vmflck(length(μ), κ))
    end
end

@compat VonMisesFisher{T<:Real}(μ::Vector{T}, κ::Real) = VonMisesFisher(Float64(μ), Float64(κ))

VonMisesFisher(θ::Vector{Float64}) = (κ = vecnorm(θ); VonMisesFisher(scale(θ, 1.0 / κ), κ))
VonMisesFisher{T<:Real}(θ::Vector{T}) = VonMisesFisher(Float64(θ))

show(io::IO, d::VonMisesFisher) = show(io, d, (:μ, :κ))

### Basic properties

length(d::VonMisesFisher) = length(d.μ)

meandir(d::VonMisesFisher) = d.μ
concentration(d::VonMisesFisher) = d.κ

insupport{T<:Real}(d::VonMisesFisher, x::DenseVector{T}) = isunitvec(x)

### Evaluation

function _vmflck(p, κ)
    hp = 0.5 * p
    q = hp - 1.0
    q * log(κ) - hp * log(2π) - log(besseli(q, κ))
end
_vmflck3(κ) = log(κ) - log2π - κ - log1mexp(-2.0 * κ)
vmflck(p, κ) = (p == 3 ? _vmflck3(κ) : _vmflck(p, κ))::Float64

_logpdf{T<:Real}(d::VonMisesFisher, x::DenseVector{T}) = d.logCκ + d.κ * dot(d.μ, x)

### Sampling

sampler(d::VonMisesFisher) = VonMisesFisherSampler(d.μ, d.κ)

_rand!(d::VonMisesFisher, x::DenseVector) = _rand!(sampler(d), x)
_rand!(d::VonMisesFisher, x::DenseMatrix) = _rand!(sampler(d), x)

### Estimation

function fit_mle(::Type{VonMisesFisher}, X::Matrix{Float64})
    r = vec(sum(X, 2))
    n = size(X, 2)
    r_nrm = vecnorm(r)
    μ = scale!(r, 1.0 / r_nrm)
    ρ = r_nrm / n
    κ = _vmf_estkappa(length(μ), ρ)
    VonMisesFisher(μ, κ)
end

@compat fit_mle{T<:Real}(::Type{VonMisesFisher}, X::Matrix{T}) = fit_mle(VonMisesFisher, Float64(X))

function _vmf_estkappa(p::Int, ρ::Float64)
    # Using the fixed-point iteration algorithm in the following paper:
    #
    #   Akihiro Tanabe, Kenji Fukumizu, and Shigeyuki Oba, Takashi Takenouchi, and Shin Ishii
    #   Parameter estimation for von Mises-Fisher distributions.
    #   Computational Statistics, 2007, Vol. 22:145-157.
    #

    const maxiter = 200
    half_p = 0.5 * p

    ρ2 = abs2(ρ)
    κ = ρ * (p - ρ2) / (1 - ρ2)
    i = 0
    while i < maxiter
        i += 1
        κ_prev = κ
        a = (ρ / _vmfA(half_p, κ))
        # println("i = $i, a = $a, abs(a - 1) = $(abs(a - 1))")
        κ *= a
        if abs(a - 1.0) < 1.0e-12
            break
        end
    end
    return κ
end

_vmfA(half_p::Float64, κ::Float64) = besseli(half_p, κ) / besseli(half_p - 1.0, κ)


"""
$(TYPEDEF)

One sample superiority test for Mean

Constructors
------------
* `OneSampleMeanSuperior(mu1::Real, mu0::Real, delta::Real, stdunknown::Bool)`

Arguments
---------
* `mu1`: True mean

* `mu0`: Mean under null hypothesis

* `delta`: Superiority margin

* `stdunknown`: If the population standard deviation is known, default = false,
when the population parameter is known, the power / sample size are calculated based
on a standard normal distribution, otherwise a t distribution will be used

Fields
------
$(FIELDS)
"""
type OneSampleMeanSuperior <: TrialTest
    mu1::Real
    mu0::Real
    delta::Real
    stdunknown::Bool

    # Validator
    function OneSampleMeanSuperior(mu1, mu0, delta, stdunknown = false)

        if !((-Inf < mu1 < Inf) & (-Inf < mu0 < Inf))

            error("Mean value must be in (-Inf, Inf)")

        end # end if

        if !(0 <= delta < Inf)

            error("The superiority margin δ must be in [0, Inf)")

        end # end if
        new(mu1, mu0, delta, stdunknown)

    end # end function

end # end type

# One-sample superiority test for mean
function hypotheses{T <: OneSampleMeanSuperior}(test::T, n::Real, std::Real, alpha::Real, side::String)

    diff = test.mu1 - test.mu0 - test.delta
    se = sqrt(1 / n) * std
    ts = diff / se
    if test.stdunknown == false

        p = cdf(ZDIST, ts - quantile(ZDIST, 1 - alpha)) + cdf(ZDIST, -ts - quantile(ZDIST, 1 - alpha))

    else

        p = cdf(ZDIST, ts - quantile(TDist(n - 1), 1 - alpha)) + cdf(TDist(n - 1), -ts - quantile(TDist(n - 1), 1 - alpha))

    end # end if
    # power_result(test, p)
    return p

end # end function


#Functions for the Harmonic correction in 2ls.
using LinearAlgebra
using QuadGK
using FastClosures

function correction_matrix(tf, env_Δ, env_pulse1, env_pulse2)
	nΔ = length(env_Δ)
	npulse1, npulse2 = length(env_pulse1), length(env_pulse2)
	dims = nΔ + npulse1 + npulse2
    M = zeros(3, dims)
    for i = 1: nΔ
        h1 = env_Δ[i]
        func1 = @closure t-> h1(t/tf)*sin(2*c(t, tf))
        int_func1, err = quadgk(func1, 0.0, tf, rtol=1e-10, atol=1e-14)
        func2 = @closure t-> h1(t/tf)*cos(2*c(t, tf))
        int_func2, err = quadgk(func2, 0.0, tf, rtol=1e-10, atol=1e-14)
        M[2, i], M[3, i] = int_func1, int_func2
    end
    for i = 1: npulse1
        h2 = env_pulse1[i]
        func3 = @closure t-> h2(t/tf)*(1 + cos(2*t))/2
        int_func3, err = quadgk(func3, 0.0, tf, rtol=1e-10, atol=1e-14)
        func4 = @closure t-> -h2(t/tf)*sin(2*t)*cos(2*c(t, tf))/2
        int_func4, err = quadgk(func4, 0.0, tf, rtol=1e-10, atol=1e-14)
        func5 = @closure t-> h2(t/tf)*sin(2*t)*sin(2*c(t, tf))/2
        int_func5, err = quadgk(func5, 0.0, tf, rtol=1e-10, atol=1e-14)

        col = nΔ + i
        M[1, col], M[2, col], M[3, col] = int_func3, int_func4, int_func5
    end
    for i = 1: npulse2
        h3 = env_pulse2[i]
        func6 = @closure t-> h3(t/tf)*sin(2*t)/2.0
        int_func6, err = quadgk(func6, 0.0, tf, rtol=1e-10, atol=1e-14)
        func7 = @closure t-> -h3(t/tf)*(1 - cos(2*t))*cos(2*c(t, tf))/2
        int_func7, err = quadgk(func7, 0.0, tf, rtol=1e-10, atol=1e-14)
        func8 = @closure t-> h3(t/tf)*(1 - cos(2*t))*sin(2*c(t, tf))/2
        int_func8, err = quadgk(func8, 0.0, tf, rtol=1e-10, atol=1e-14)

        col = nΔ + npulse1 + i
        M[1, col], M[2, col], M[3, col] = int_func6, int_func7, int_func8
    end
    return M
end

function coefficients_W(Omega, M)
    ih_vec = Omega[1: 3]
    coeffs = pinv(M)*ih_vec
    return coeffs
end

function g(t, tf, coeffs, env_pulse1, env_pulse2)
	npulse1, npulse2 = length(env_pulse1), length(env_pulse2)
	cost, sint = cos(t), sin(t)
	gt = zero(t)
    for i = 1: npulse1
        h2 = env_pulse1[i]
        gt += coeffs[i]*h2(t/tf)*cost
    end
    for i = 1: npulse2
        h3 = env_pulse2[i]
        col = npulse1 + i
        gt += coeffs[col]*h3(t/tf)*sint
    end
	return gt
end
function Δ(t, tf, coeffs, env_Δ)
	nΔ = length(env_Δ)
	Δt = zero(t)
    for i = 1: nΔ
        h1 = env_Δ[i]
        Δt += coeffs[i]*h1(t/tf)
    end
	return Δt
end
function Wharmonic(t, tf, coeffs, env_Δ, env_pulse1, env_pulse2)
	Δt = Δ(t, tf, coeffs[1], env_Δ)
	gt = g(t, tf, coeffs[2], env_pulse1, env_pulse2)

    W = zeros(4)
    W[1] = gt*cos(t)
    W[2] = -gt*sin(t)
    W[3] = Δt
    return W
end

function Wharmonic_I(t, tf, coeffs, env_Δ, env_pulse1, env_pulse2)
	cost, sint = cos(t), sin(t)
	cosθ, sinθ = cos(2*c(t, tf)), sin(2*c(t, tf))
	Δt = Δ(t, tf, coeffs[1], env_Δ)
	gt = g(t, tf, coeffs[2], env_pulse1, env_pulse2)

	W_I = zeros(4)
	W_I[1] = gt*cost
	W_I[2] = Δt*sinθ - gt*sint*cosθ
	W_I[3] = Δt*cosθ + gt*sint*sinθ
	return W_I
end

function get_coeffs(M, coeffs, order, tf, env_Δ, env_pulse1, env_pulse2)
	nΔ = length(env_Δ)
	A = @closure (t-> - V_I(t, tf)
					  - Wharmonic_I(t, tf, coeffs, env_Δ, env_pulse1, env_pulse2))

    Ndims = 4
    Omegaf = solveOmega(tf, order, A, comm!, Ndims)

    vec = sum(Omegaf, dims=2)
    coeffs = coefficients_W(vec, M)
    return (coeffs[1: nΔ], coeffs[nΔ + 1: end])
end
function get_coeffs(M, order, tf, env_Δ, env_pulse1, env_pulse2)
    nΔ = length(env_Δ)
	npulse1, npulse2 = length(env_pulse1), length(env_pulse2)
    coeffs = (zeros(nΔ),  zeros(npulse1 + npulse2))
    for k = 1: order
		coeffs2 = get_coeffs(M, coeffs, k, tf, env_Δ, env_pulse1, env_pulse2)
        coeffs = coeffs .+ coeffs2
    end
    return coeffs
end

function U(tf, σ, coeffs, env_Δ, env_pulse1, env_pulse2)
	H = @closure (t-> H0(t, tf, true)
					  + Wharmonic(t, tf, coeffs, env_Δ, env_pulse1, env_pulse2))

    sol = solveU(H, tf, σ)
    return sol(tf)
end


"""
(D)_{ij} = ∑_k (x_k^{(i)} - x_k^{(j)})^2
"""

function square_metric(x)
	N = size(x, 2)
	D = zeros(eltype(x), N, N)
	@inbounds for i in 1:N
		for j in i+1:N
			dx = norm(view(x, :, i) .- view(x, :, j))
			D[i,j] = D[j,i] = dx * dx
		end
	end
	D
end

function square_metric_back(x, sym_dD)
	dD = 0.5*(sym_dD + sym_dD')  # Cholesky decomposition only use lower triangle, reconstruct the true dD
	m, n = size(x)
	dx = zeros(eltype(x), m, n)
  @inbounds for i in 1:m
    for j in 1:n
      for k in 1:n
        dx[i,j] += 4*dD[j,k]*(x[i,j]-x[i,k])
      end
    end
  end
  dx
end

function square_metric(x, xo)
	N1 = size(x, 2)
	N2 = size(xo, 2)
	D = zeros(eltype(x), N1, N2)
	@inbounds for i in 1:N1
		for j in 1:N2
			dx = norm(view(x, :, i) .- view(xo, :, j))
			D[i,j] = dx * dx
		end
	end
	D
end