import .StatsFuns
using .StatsFuns: xlogx, xlogy, logistic, logit, log1psq, log1pexp,
    logsumexp, logaddexp, logsubexp
using Base.Broadcast: broadcasted

@adjoint function xlogx(x::Real)
    result, dx = ∇xlogx(x)
    back(δ) = (dx * δ,)
    return result, back
end
@adjoint function broadcasted(::typeof(xlogx), x::Numeric)
    result, dx = ∇xlogx(x)
    back(δ) = (nothing, unbroadcast(x, δ .* dx))
    return result, back
end
function ∇xlogx(x::Numeric)
    logx = log.(x)
    xlogx = x .* logx
    result = @. ifelse(iszero(x), zero(xlogx), xlogx)
    dx = @. one(x) + logx
    return result, dx
end

@adjoint function logistic(x::Real)
    y = logistic(x)
    return y, Δ->(Δ * y * (1 - y),)
end

@adjoint logit(x::Real) = logit(x), Δ->(Δ / (x * (1 - x)),)

@adjoint log1psq(x::Real) = log1psq(x), Δ->(Δ * 2x / (1 + abs2(x)),)

@adjoint function log1pexp(x::Real)
    dx = ∂log1pexp(x)
    return log1pexp(x), δ -> (δ * dx,)
end
@adjoint function broadcasted(::typeof(log1pexp), x::Numeric)
    dx = ∂log1pexp.(x)
    return log1pexp.(x), δ -> (nothing, unbroadcast(x, δ .* dx))
end
∂log1pexp(x::Real)    = x < 18.0 ? logistic(x) : x < 33.3 ? one(x) - exp(-x) : oftype(exp(x), 1)
∂log1pexp(x::Float32) = x < 9f0  ? logistic(x) : x < 16f0 ? one(x) - exp(-x) : oftype(exp(x), 1)

@adjoint function logsumexp(X::AbstractArray{<:Real}; dims=:)
    lse = logsumexp(X; dims=dims)
    return lse, Δ -> (Δ .* exp.(X .- lse),)
end

@adjoint function xlogy(x::Real, y::Real)
    result, dx, dy = ∇xlogy(x, y)
    back(δ) = (δ * dx, δ * dy)
    return result, back
end
@adjoint function broadcasted(::typeof(xlogy), x::Numeric, y::Numeric)
    result, dx, dy = ∇xlogy(x, y)
    back(δ) = (nothing, unbroadcast(x, δ .* dx), unbroadcast(y, δ .* dy))
    return result, back
end
function ∇xlogy(x::Numeric, y::Numeric)
    dx = logy = log.(y)
    dy = x ./ y
    xlogy = x .* logy
    result = @. ifelse(iszero(x) & !isnan(y), zero(xlogy), xlogy)
    return result, dx, dy
end

@adjoint function logaddexp(x::Real, y::Real)
    result, dx, dy = ∇logaddexp(x, y)
    back(δ) = (δ * dx, δ * dy)
    return result, back
end
@adjoint function broadcasted(::typeof(logaddexp), x::Numeric, y::Numeric)
    result, dx, dy = ∇logaddexp(x, y)
    back(δ) = (nothing, unbroadcast(x, δ .* dx), unbroadcast(y, δ .* dy))
    return result, back
end
function ∇logaddexp(x::Numeric, y::Numeric)
    result = logaddexp.(x, y)
    t = @. exp(-abs(x - y))
    dx, dy = select(x .≥ y, inv.(one.(t) .+ t), t ./ (one.(t) .+ t))
    return result, dx, dy
end

@adjoint function logsubexp(x::Real, y::Real)
    result, dx, dy = ∇logsubexp(x, y)
    back(δ) = (δ * dx, δ * dy)
    return result, back
end
@adjoint function broadcasted(::typeof(logsubexp), x::Numeric, y::Numeric)
    result, dx, dy = ∇logsubexp(x, y)
    back(δ) = (nothing, unbroadcast(x, δ .* dx), unbroadcast(y, δ .* dy))
    return result, back
end
function ∇logsubexp(x::Numeric, y::Numeric)
    result = logsubexp.(x, y)
    t = @. -inv(expm1(-abs(x - y)))
    dx, dy = select(x .≥ y, t, one.(t) .- t)
    return result, dx, dy
end

"""
	select(cond, x, y)

The call `a, b = select(cond, x, y)` constructs two arrays `a, b`, where
`a[i], b[i] = x[i], y[i]` if `cond[i]` is `true`, and `a[i], b[i] = y[i], x[i]`
 if `cond[i]` is `false`.
"""
select(cond, x, y) = ifelse.(cond, x, y), ifelse.(cond, y, x)


@doc raw"""
    guess_density(basis)

Build a superposition of atomic densities (SAD) guess density.

We take for the guess density a gaussian centered around the atom, of
length specified by `atom_decay_length`, normalized to get the right number of electrons
```math
\hat{ρ}(G) = Z \exp\left(-(2π \text{length} |G|)^2\right)
```
"""
guess_density(basis::PlaneWaveBasis) = guess_density(basis, basis.model.atoms)
@timing function guess_density(basis::PlaneWaveBasis{T}, atoms) where {T}
    model = basis.model
    ρ = zeros(complex(T), basis.fft_size)
    # If no atoms, start with a zero initial guess
    isempty(atoms) && return from_fourier(basis, ρ)
    # fill ρ with the (unnormalized) Fourier transform, ie ∫ e^{-iGx} ρ(x) dx
    for (spec, positions) in atoms
        n_el_val = n_elec_valence(spec)
        decay_length::T = atom_decay_length(spec)
        for (iG, G) in enumerate(G_vectors(basis))
            Gsq = sum(abs2, model.recip_lattice * G)
            form_factor::T = n_el_val * exp(-Gsq * decay_length^2)
            for r in positions
                ρ[iG] += form_factor * cis(-2T(π) * dot(G, r))
            end
        end
    end

    # projection in the normalized plane wave basis
    from_fourier(basis, ρ / sqrt(model.unit_cell_volume))
end

@doc raw"""
Get the lengthscale of the valence density for an atom with `n_elec_core` core
and `n_elec_valence` valence electrons.
```
"""
function atom_decay_length(n_elec_core, n_elec_valence)
    # Adapted from ABINIT/src/32_util/m_atomdata.F90,
    # from which also the data has been taken.

    n_elec_valence = round(Int, n_elec_valence)
    if n_elec_valence == 0
        return 0.0
    end

    data = if n_elec_core < 0.5
        # Bare ions: Adjusted on 1H and 2He only
        [0.6, 0.4, 0.3, 0.25, 0.2]
    elseif n_elec_core < 2.5
        # 1s2 core: Adjusted on 3Li, 6C, 7N, and 8O
        [1.8, 1.4, 1.0, 0.7, 0.6, 0.5, 0.4, 0.35, 0.3]
    elseif n_elec_core < 10.5
        # Ne core (1s2 2s2 2p6): Adjusted on 11na, 13al, 14si and 17cl
        [2.0, 1.6, 1.25, 1.1, 1.0, 0.9, 0.8, 0.7 , 0.7, 0.7, 0.6]
    elseif n_elec_core < 12.5
        # Mg core (1s2 2s2 2p6 3s2): Adjusted on 19k, and on n_elec_core==10
        [1.9, 1.5, 1.15, 1.0, 0.9, 0.8, 0.7, 0.6 , 0.6, 0.6, 0.5]
    elseif n_elec_core < 18.5
        # Ar core (Ne + 3s2 3p6): Adjusted on 20ca, 25mn and 30zn
        [2.0, 1.8, 1.5, 1.2, 1.0, 0.9, 0.85, 0.8, 0.75, 0.7, 0.65, 0.65, 0.6]
    elseif n_elec_core < 28.5
        # Full 3rd shell core (Ar + 3d10): Adjusted on 31ga, 34se and 38sr
        [1.5, 1.25, 1.15, 1.05, 1.00, 0.95, 0.95, 0.9, 0.9, 0.85, 0.85, 0.80,
         0.8 , 0.75, 0.7]
    elseif n_elec_core < 36.5
        # Krypton core (Ar + 3d10 4s2 4p6): Adjusted on 39y, 42mo and 48cd
        [2.0, 2.00, 1.60, 1.40, 1.25, 1.10, 1.00, 0.95, 0.90, 0.85, 0.80, 0.75, 0.7]
    else
        # For the remaining elements, consider a function of n_elec_valence only
        [2.0 , 2.00, 1.55, 1.25, 1.15, 1.10, 1.05, 1.0 , 0.95, 0.9, 0.85, 0.85, 0.8]
    end
    data[min(n_elec_valence, length(data))]
end
atom_decay_length(sp::Element) = atom_decay_length(n_elec_core(sp), n_elec_valence(sp))


module ConvDiffMIPDECO

using jInv.Mesh
using jInv.ForwardShare
using jInv.Utils
using jInv.LinearSolvers
using jInv.InverseSolve
using KrylovMethods
using LinearAlgebra
using SparseArrays
using Printf
using DSP

function getBICGSTB(;PC=:jac,maxIter=1000,out=0,tol=1e-10)
	bicg = (A,b; M=identity,tol=1e-10,maxIter=500,out=1)->
	bicgstb(A,b,M1=identity,tol=tol,maxIter=maxIter,out=out,tolRho=1e-60)
	return getIterativeSolver(bicg,PC=PC,maxIter=maxIter,out=out,tol=tol)
end

# files are organized in the src (containing functions) and test (containing some unit tests)
import jInv.ForwardShare.ForwardProbType

export ConvDiffParam, getConvDiffParam

"""
type ConvDiffParam <: ForwardProbType

description of stationary convection diffusion forward problem

(- sig*Laplacian + v dot GRAD) u = f
with boundary conditions du/dn = gn on Omega1, u=gd on Omega2

Construct an instance using getConvDiffParam(M,v,kwargs...)

Fields:
	M      - describes computational mesh
	v      - velocities at cell-centers
	sig    - scalar, >0
	bc     - vector describing boundary conditions
	Fields - stores fields, i.e., u
	Ainv   - factorization of PDE

"""
mutable struct ConvDiffParam <: ForwardProbType
	M :: RegularMesh
	A :: SparseMatrixCSC{Float64}
	P :: AbstractArray{Float64}
	sig::Float64
	bc::Array{Float64}
	Fields::Array{Float64,1}
	Ainv::AbstractSolver
end

"""
function getConvDiffParam(M,v)

constructs and returns ConvDiffParam

The PDE is discretized and factorized and stored in field Ainv.

Inputs:
	M            - mesh

Keyword arguments:
	v            - velocity, (vector or function)
	gd::Function - Dirichlet boundary conditions
	gn::Function - Neuman boundary condition
	P            - measurement operator
	sig          - viscosity
	bc           - description of boundary conditions
	Fields       - storing PDE solutions
	Ainv         - description of linear solver

"""
function getConvDiffParam(M::RegularMesh,v;
	gd::Function=X->zeros(size(X,1)), gn::Function=X->zeros(size(X,1)), P=Diagonal(ones(M.nc)),sig::Number=0.01,bc=(:dir,:neu,:neu,:neu),Fields=zeros(0),Ainv=getJuliaSolver())

	# get boundary conditions
	iddir, idneu, iddirn, idneun, idint = getBoundaryIndices(M,bc)

	A, Adir, Aneu = getConvDiffMatrix(M,sig,v,iddir,idneu,iddirn,idneun,idint)

	gdir, gneu = getBoundaryFuncVal(gd,gn,M,bc)

	bc = -2*Adir*gdir + Aneu*gneu

	return ConvDiffParam(M,A,P,sig,bc,Fields,Ainv)
end

include("getBoundaryCondition.jl")
include("getGhostIndices.jl")
include("getConvDiffMatrix.jl")
include("getBoundaryIndices.jl")
include("getBoundaryFuncVal.jl")
include("getDiffOps.jl")
include("getData.jl")
include("utils.jl")
include("ConvDiffFEMParam.jl")
include("FEM.jl")
include("getDataFEM.jl")
include("getConvDiffFEMConstraintsAMPL.jl")
include("mipdecoHeuristic.jl")
include("rounding.jl")
include("dilation.jl")
include("regularizers.jl")
end


lines = readlines("input.txt")
octomap = parse.(Int, reduce(hcat, collect.(lines)))
idx = CartesianIndices(octomap)
minval = typemin(Int)

o = CartesianIndex(1,1)
V = CartesianIndices((-1:1, -1:1))
padded = fill(minval, size(octomap) .+ 2)
padded[idx .+ o] = octomap[idx]

Σ = 0
for step = 1:100
    padded .+= 1
    flashmap = fill(false, size(padded))
    while (yxs = findall(padded .>= 10)) |> length > 0
        flashmap[yxs] .= true
        padded[yxs] .= minval # Prevent any further flashing for this step
        [padded[yx .+ V] .+= 1 for yx = yxs]
    end
    padded[findall(flashmap)] .= 0
    global Σ += sum(flashmap)
end

println(Σ)


using Interpolations, Test, LinearAlgebra, Random

@testset "ScalingNoInterpTests" begin
    xs = -pi:2pi/10:pi
    f1(x) = sin(x)
    f2(x) = cos(x)
    f3(x) = sin(x) .* cos(x)
    f(x,y) = y == 1 ? f1(x) : (y == 2 ? f2(x) : (y == 3 ? f3(x) : error("invalid value for y (must be 1, 2 or 3, you used $y)")))
    ys = 1:3

    A = hcat(map(f1, xs), map(f2, xs), map(f3, xs))

    itp = interpolate(A, (BSpline(Quadratic(Periodic(OnGrid()))), NoInterp()))
    sitp = scale(itp, xs, ys)

    for (ix,x0) in enumerate(xs[1:end-1]), y0 in ys
        x,y = x0, y0
        @test ≈(sitp(x,y),f(x,y),atol=0.05)
    end

    @test length(Interpolations.gradient(sitp, pi/3, 2)) == 1

    # check for case where initial/middle indices are NoInterp but later ones are <:BSpline
    isdefined(Random, :seed!) ? Random.seed!(1234) : srand(1234) # `srand` was renamed to `seed!`
    z0 = rand(10,10)
    za = copy(z0)
    zb = copy(z0')

    itpa = interpolate(za, (BSpline(Linear()), NoInterp()))
    itpb = interpolate(zb, (NoInterp(), BSpline(Linear())))

    rng = range(1.0, stop=19.0, length=10)
    sitpa = scale(itpa, rng, 1:10)
    sitpb = scale(itpb, 1:10, rng)
    @test Interpolations.gradient(sitpa, 3.0, 3) ==  Interpolations.gradient(sitpb, 3, 3.0)

end


using LinearAlgebra: norm

#=
  Solvers for GLM
=#
abstract type AbstractSolver end

# Weights functions
@inline function W(distrib::AbstractDistribution, link::AbstractLink, mu::Array{T, 1}, eta::Array{T, 1}) where {T <: AbstractFloat}
  return ( (deta_dmu(link, mu, eta).^2) .* variance(distrib, mu)).^(-0.5)
end
function W(distrib::AbstractDistribution, link::AbstractLink, mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}) where {T <: AbstractFloat}
  nBlocks::Int64 = length(mu)
  return [( (deta_dmu(link, mu[i], eta[i]).^2) .* variance(distrib, mu[i])).^(-0.5) for i in 1:nBlocks]
end
function W(::Block1DParallel, distrib::AbstractDistribution, link::AbstractLink, mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}) where {T <: AbstractFloat}
  nBlocks::Int64 = length(mu)
  ret::Array{Array{T, 1}, 1} = Array{Array{T, 1}, 1}(undef, nBlocks)
  @threads for i in 1:nBlocks
    ret[i] = ((deta_dmu(link, mu[i], eta[i]).^2) .* variance(distrib, mu[i])).^(-0.5)
  end
  return ret
end

#=
  GLM Solver uses (X'WX)^(-1) * X'Wy
=#
struct VanillaSolver <: AbstractSolver end

# Weights for the VanillaSolver
@inline function W(::VanillaSolver, distrib::AbstractDistribution, link::AbstractLink, mu::Array{T, 1}, eta::Array{T, 1}) where {T <: AbstractFloat}
  return ((deta_dmu(link, mu, eta).^2) .* variance(distrib, mu)).^(-1)
end
function W(::VanillaSolver, distrib::AbstractDistribution, link::AbstractLink, mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}) where {T <: AbstractFloat}
  nBlocks::Int64 = length(mu)
  return [((deta_dmu(link, mu[i], eta[i]).^2) .* variance(distrib, mu[i])).^(-1) for i in 1:nBlocks]
end
function W(::Block1DParallel, ::VanillaSolver, distrib::AbstractDistribution, link::AbstractLink, mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}) where {T <: AbstractFloat}
  nBlocks::Int64 = length(mu)
  ret::Array{Array{T, 1}, 1} = Array{Array{T, 1}, 1}(undef, nBlocks)
  @threads for i in 1:nBlocks
    ret[i] = ((deta_dmu(link, mu[i], eta[i]).^2) .* variance(distrib, mu[i])).^(-1)
  end
  return ret
end

#=
  Solver Using QR Decomposition
=#
struct QRSolver <: AbstractSolver end

# Weights for the QRSolver
@inline function W(::QRSolver, distrib::AbstractDistribution, link::AbstractLink, mu::Array{T, 1}, eta::Array{T, 1}) where {T <: AbstractFloat}
  return ((deta_dmu(link, mu, eta).^2) .* variance(distrib, mu)).^(-0.5)
end
function W(::QRSolver, distrib::AbstractDistribution, link::AbstractLink, mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}) where {T <: AbstractFloat}
  nBlocks::Int64 = length(mu)
  return [((deta_dmu(link, mu[i], eta[i]).^2) .* variance(distrib, mu[i])).^(-0.5) for i in 1:nBlocks]
end
function W(::Block1DParallel, ::QRSolver, distrib::AbstractDistribution, link::AbstractLink, mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}) where {T <: AbstractFloat}
  nBlocks::Int64 = length(mu)
  ret::Array{Array{T, 1}, 1} = Array{Array{T, 1}, 1}(undef, nBlocks)
  @threads for i in 1:nBlocks
    ret[i] = ((deta_dmu(link, mu[i], eta[i]).^2) .* variance(distrib, mu[i])).^(-0.5)
  end
  return ret
end

#=
  TODO:
  1. Unify the QRSolve and Vanilla solve interface as in
     the D implementation.
=#
function solve(::Type{<: VanillaSolver}, X::Array{T, 2}, y::Array{T, 1}) where {T <: AbstractFloat}
  ret::Array{T, 1} = Array{T, 1}(undef, 0)
  try
    ret = (X'*X)/(X'*y)
  catch
    ret = inv(X'*X)*(X'*y)
  end
  return ret
  #return inv(X'*X)*(X'*y)
end

import LinearAlgebra.LAPACK.geqrf!
import LinearAlgebra.LAPACK.orgqr!
import LinearAlgebra.BLAS.gemv
import LinearAlgebra.LAPACK.trtrs!

# import LinearAlgebra.UpperTriangular
# import LinearAlgebra.LAPACK.trtri!
# import LinearAlgebra.BLAS.trsm!

#=
  The QR Solver
=#
function getUpperR(QR::Array{T, 2}) where {T <: AbstractFloat}
  _, p = size(QR)
  R = zeros(T, (p, p))
  for i in 1:p
    for j in 1:i
      R[j, i] = QR[j, i]
    end
  end
  return R
end

function solve(::Type{<: QRSolver}, Q::Array{T, 2}, y::Array{T, 1}) where {T <: AbstractFloat}
  Q = copy(Q)
  _, p = size(Q)
  tau = zeros(T, p)
  geqrf!(Q, tau)
  R = getUpperR(Q)
  orgqr!(Q, tau, p)
  z = gemv('T', 1.0, Q, y)
  # z = Q' * y
  trtrs!('U', 'N', 'N', R, z)
  return z
end

#============================  Matrix Inverses ============================#
#= Matrix Inverse By LU Decomposition =#
import LinearAlgebra.LAPACK.getrf!
import LinearAlgebra.LAPACK.getri!

#=
  Inverse Types
=#
abstract type AbstractInverse end

#= GETRI Matrix Inverse =#
struct GETRIInverse <: AbstractInverse end
#= POTRI Matrix Inverse =#
struct POTRIInverse <: AbstractInverse end
#= SYTRF Matrix Inverse =#
struct SYTRFInverse <: AbstractInverse end
#= GESVD Matrix Inverse =#
struct GESVDInverse <: AbstractInverse end

#=
  Inverse of a square matrix using the GETRI function
  uses LU decomposition. The geqrf() does the QR decomposition
  and the getri() function carries out an inverse from the output
  of geqrf().
=#
function inv(::GETRIInverse, A::Array{T, 2}) where {T <: AbstractFloat}
  A = copy(A)
  p, _ = size(A)
  A, ipiv, info = getrf!(A)
  @assert(info == 0, "getrf! error in argument " * string(info))
  getri!(A, ipiv)
  return A
end

#= Matrix Inverse By Cholesky Decomposition For Positive Definite Matrices =#
import LinearAlgebra.LAPACK.potrf!
import LinearAlgebra.LAPACK.potri!

#=
  # Examples:
  x = rand(100, 10);
  x = x' * x
  sum(abs.(Base.inv(x) .- inv(POTRIInverse(), x)))
=#
function inv(::POTRIInverse, A::Array{T, 2}) where {T <: AbstractFloat}
  p, _ = size(A)
  A = copy(A)
  A, info = potrf!('U', A)
  @assert(info == 0, "potrf error in argument " * string(info))
  potri!('U', A)
  for i in 1:p
    for j in 1:i
      A[i, j] = A[j, i]
    end
  end
  return A
end

#= Matrix Inverse By LU Decomposition For Symmetric Indefinite Matrices =#
import LinearAlgebra.LAPACK.sytri!# (uplo, A, ipiv)
import LinearAlgebra.LAPACK.sytrf!# (uplo, A) -> (A, ipiv, info)

#=
  # Examples:
  x = rand(100, 10);
  x = x' * x
  sum(abs.(Base.inv(x) .- inv(SYTRFInverse(), x)))
=#
function inv(::SYTRFInverse, A::Array{T, 2}) where {T <: AbstractFloat}
  p, _ = size(A)
  A = copy(A)
  A, ipiv, info = sytrf!('U', A)
  sytri!('U', A, ipiv)
  for i in 1:p
    for j in 1:i
      A[i, j] = A[j, i]
    end
  end
  return A
end

#= Generalized Matrix Inverse By SVD  =#
import LinearAlgebra.LAPACK.gesvd! # (jobu, jobvt, A) -> (U, S, VT)

#= Matrix Vector Sweep Function =#
function multSweep(xw::Array{T, 2}, w::Array{T, 1})::Array{T, 2} where {T <: AbstractFloat}
  n, p = size(xw)
  for j in 1:p
    for i in 1:n
      xw[i, j] = xw[i, j] * w[i]
    end
  end
  return xw
end

#=
  # Examples:
  x = rand(100, 10);
  x = x' * x
  sum(abs.(Base.inv(x) .- inv(GESVDInverse(), x)))
=#
function inv(::GESVDInverse, A::Array{T, 2}) where {T <: AbstractFloat}
  p, _ = size(A)
  A = copy(A)
  U, S, Vt = gesvd!('A', 'A', A)
  S = [el > 1E-9 ? 1/el : el for el in S]
  return multSweep(Vt, S)' * U'
end

#============================== Linear Equation Solvers ==============================#
#= For Ax = b =#
#==============================  GESV Solver ==============================#

function calcXWX(xwx::Array{T, 2}, x::Array{T, 2}, w::Array{T, 1}, z::Array{T, 1}) where {T <: AbstractFloat}
  # n, p = size(x)
  xw = copy(x)
  xw = multSweep(xw, w)
  xwx = xw' * x
  xwz = xw' * z
  return xwx, xwz
end

function calcXWX(xwx::Array{T, 2}, x::Array{Array{T, 2}, 1}, w::Array{Array{T, 1}, 1}, z::Array{Array{T, 1}, 1}) where {T <: AbstractFloat}
  p::Int64 = size(x[1])[2]
  nBlocks::Int64 = length(x)
  xwx = zeros(T, (p, p))
  xwz = zeros(T, (p, 1))
  for i in 1:nBlocks
    xw = copy(x[i])
    xw = multSweep(xw, w[i])
    xwx += xw' * x[i]
    xwz += xw' * z[i]
  end
  return xwx, xwz
end

function calcXWX(::Block1DParallel, xwx::Array{T, 2}, x::Array{Array{T, 2}, 1}, w::Array{Array{T, 1}, 1}, z::Array{Array{T, 1}, 1}) where {T <: AbstractFloat}
  p::Int64 = size(x[1])[2]
  nBlocks::Int64 = length(x)
  xwx = [zeros(T, (p, p)) for i in 1:nthreads()]
  xwz = [zeros(T, (p, 1)) for i in 1:nthreads()]
  @threads for i in 1:nBlocks
    xw = copy(x[i])
    xw = multSweep(xw, w[i])
    xwx[threadid()] += xw' * x[i]
    xwz[threadid()] += xw' * z[i]
  end
  for i in 2:nthreads()
    xwx[1] .+= xwx[i]
    xwz[1] .+= xwz[i]
  end
  return xwx[1], xwz[1]
end

function calcXW(xwx::Array{T, 2}, x::Array{T, 2}, w::Array{T, 1}, z::Array{T, 1}) where {T <: AbstractFloat}
  # n, p = size(x)
  xw = copy(x)
  xw = multSweep(xw, w)
  return xw, z .* w
end

struct GESVSolver <: AbstractSolver end
# Weights for the GESVSolver
@inline function W(::GESVSolver, distrib::AbstractDistribution, link::AbstractLink, mu::Array{T, 1}, eta::Array{T, 1}) where {T <: AbstractFloat}
  return ((deta_dmu(link, mu, eta).^2) .* variance(distrib, mu)).^(-1)
end
function W(::GESVSolver, distrib::AbstractDistribution, link::AbstractLink, mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}) where {T <: AbstractFloat}
  nBlocks::Int64 = length(mu)
  return [((deta_dmu(link, mu[i], eta[i]).^2) .* variance(distrib, mu[i])).^(-1) for i in 1:nBlocks]
end
function W(matrixType::Block1DParallel, ::GESVSolver, distrib::AbstractDistribution, link::AbstractLink, mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}) where {T <: AbstractFloat}
  nBlocks::Int64 = length(mu)
  ret::Array{Array{T, 1}, 1} = Array{Array{T, 1}, 1}(undef, nBlocks)
  @threads for i in 1:nBlocks
    ret[i] = ((deta_dmu(link, mu[i], eta[i]).^2) .* variance(distrib, mu[i])).^(-1)
  end
  return ret
end

# Solver
import LinearAlgebra.LAPACK.gesv!
function solve!(::GESVSolver, R::Array{T, 2},
    xwx::Array{T, 2}, xw::Array{T, 2},
    x::Array{T, 2}, z::Array{T, 1}, w::Array{T, 1},
    coef::Array{T, 1}) where {T <: AbstractFloat}
  xwx, xwz = calcXWX(xwx, x, w, z)
  coef, _, ipiv = gesv!(copy(xwx), xwz)
  return (R, xwx, xw, x, z, w, coef)
end
function solve!(::GESVSolver, R::Array{T, 2},
  xwx::Array{T, 2}, xw::Array{Array{T, 2},1},
  x::Array{Array{T, 2}, 1}, z::Array{Array{T, 1}, 1},
  w::Array{Array{T, 1}, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}
  xwx, xwz = calcXWX(xwx, x, w, z)
  coef, _, ipiv = gesv!(copy(xwx), xwz)
  return (R, xwx, xw, x, z, w, coef)
end
function solve!(matrixType::Block1DParallel, ::GESVSolver, R::Array{T, 2},
  xwx::Array{T, 2}, xw::Array{Array{T, 2},1},
  x::Array{Array{T, 2}, 1}, z::Array{Array{T, 1}, 1},
  w::Array{Array{T, 1}, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}
  xwx, xwz = calcXWX(matrixType, xwx, x, w, z)
  coef, _, ipiv = gesv!(copy(xwx), xwz)
  return (R, xwx, xw, x, z, w, coef)
end
# Default Covariance Function
function cov(::AbstractSolver, invType::AbstractInverse, R::Array{T, 2}, xwx::Array{T, 2}, xw::Array{T, 2})::Array{T, 2} where {T <: AbstractFloat}
  return inv(invType, xwx)
end

function cov(::AbstractSolver, invType::AbstractInverse, R::Array{T, 2}, xwx::Array{T, 2})::Array{T, 2} where {T <: AbstractFloat}
  return inv(invType, xwx)
end

#==============================  POSV Solver ==============================#
#= Solver For Positive Definite Matrices =#
# Weights for posv!(uplo, A, B) -> (A, B)
import LinearAlgebra.LAPACK.posv! # (uplo, A, B) -> (A, B)

struct POSVSolver <: AbstractSolver end
# Weights for the POSVSolver
@inline function W(::POSVSolver, distrib::AbstractDistribution, link::AbstractLink, mu::Array{T, 1}, eta::Array{T, 1}) where {T <: AbstractFloat}
  return ((deta_dmu(link, mu, eta).^2) .* variance(distrib, mu)).^(-1)
end
function W(::POSVSolver, distrib::AbstractDistribution, link::AbstractLink, mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}) where {T <: AbstractFloat}
  nBlocks::Int64 = length(mu)
  return [((deta_dmu(link, mu[i], eta[i]).^2) .* variance(distrib, mu[i])).^(-1) for i in 1:nBlocks]
end
function W(matrixType::Block1DParallel, ::POSVSolver, distrib::AbstractDistribution, link::AbstractLink, mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}) where {T <: AbstractFloat}
  nBlocks::Int64 = length(mu)
  ret::Array{Array{T, 1}, 1} = Array{Array{T, 1}, 1}(undef, nBlocks)
  @threads for i in 1:nBlocks
    ret[i] = ((deta_dmu(link, mu[i], eta[i]).^2) .* variance(distrib, mu[i])).^(-1)
  end
  return ret
end
# Solver
function solve!(::POSVSolver, R::Array{T, 2},
  xwx::Array{T, 2}, xw::Array{T, 2},
  x::Array{T, 2}, z::Array{T, 1}, w::Array{T, 1},
  coef::Array{T, 1}) where {T <: AbstractFloat}
  xwx, xwz = calcXWX(xwx, x, w, z)
  _, coef = posv!('U', copy(xwx), xwz)
  return (R, xwx, xw, x, z, w, coef)
end
function solve!(::POSVSolver, R::Array{T, 2},
  xwx::Array{T, 2}, xw::Array{Array{T, 2},1},
  x::Array{Array{T, 2}, 1}, z::Array{Array{T, 1}, 1},
  w::Array{Array{T, 1}, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}
  xwx, xwz = calcXWX(xwx, x, w, z)
  _, coef = posv!('U', copy(xwx), xwz)
  return (R, xwx, xw, x, z, w, coef)
end
function solve!(matrixType::Block1DParallel, ::POSVSolver, R::Array{T, 2},
  xwx::Array{T, 2}, xw::Array{Array{T, 2},1},
  x::Array{Array{T, 2}, 1}, z::Array{Array{T, 1}, 1},
  w::Array{Array{T, 1}, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}
  xwx, xwz = calcXWX(matrixType, xwx, x, w, z)
  _, coef = posv!('U', copy(xwx), xwz)
  return (R, xwx, xw, x, z, w, coef)
end
#==============================  SYSV Solver ==============================#
#= Solver For Symmetric Indefinite Matrices =#
import LinearAlgebra.LAPACK.sysv! # (uplo, A, B) -> (B, A, ipiv)

struct SYSVSolver <: AbstractSolver end
# Weights for the SYSVSolver
@inline function W(::SYSVSolver, distrib::AbstractDistribution, link::AbstractLink, mu::Array{T, 1}, eta::Array{T, 1}) where {T <: AbstractFloat}
  return ((deta_dmu(link, mu, eta).^2) .* variance(distrib, mu)).^(-1)
end
function W(::SYSVSolver, distrib::AbstractDistribution, link::AbstractLink, mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}) where {T <: AbstractFloat}
  nBlocks::Int64 = length(mu)
  return [((deta_dmu(link, mu[i], eta[i]).^2) .* variance(distrib, mu[i])).^(-1) for i in 1:nBlocks]
end
function W(matrixType::Block1DParallel, ::SYSVSolver, distrib::AbstractDistribution, link::AbstractLink, mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}) where {T <: AbstractFloat}
  nBlocks::Int64 = length(mu)
  ret::Array{Array{T, 1}, 1} = Array{Array{T, 1}, 1}(undef, nBlocks)
  @threads for i in 1:nBlocks
    ret[i] = ((deta_dmu(link, mu[i], eta[i]).^2) .* variance(distrib, mu[i])).^(-1)
  end
  return ret
end
# Solver
function solve!(::SYSVSolver, R::Array{T, 2},
  xwx::Array{T, 2}, xw::Array{T, 2},
  x::Array{T, 2}, z::Array{T, 1}, w::Array{T, 1},
  coef::Array{T, 1}) where {T <: AbstractFloat}
  xwx, xwz = calcXWX(xwx, x, w, z)
  coef, _, _ = sysv!('U', copy(xwx), xwz)
  return (R, xwx, xw, x, z, w, coef)
end
function solve!(::SYSVSolver, R::Array{T, 2},
  xwx::Array{T, 2}, xw::Array{Array{T, 2},1},
  x::Array{Array{T, 2}, 1}, z::Array{Array{T, 1}, 1},
  w::Array{Array{T, 1}, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}
  xwx, xwz = calcXWX(xwx, x, w, z)
  coef, _, _ = sysv!('U', copy(xwx), xwz)
  return (R, xwx, xw, x, z, w, coef)
end
function solve!(matrixType::Block1DParallel, ::SYSVSolver, R::Array{T, 2},
  xwx::Array{T, 2}, xw::Array{Array{T, 2},1},
  x::Array{Array{T, 2}, 1}, z::Array{Array{T, 1}, 1},
  w::Array{Array{T, 1}, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}
  xwx, xwz = calcXWX(matrixType, xwx, x, w, z)
  coef, _, _ = sysv!('U', copy(xwx), xwz)
  return (R, xwx, xw, x, z, w, coef)
end
#============================== Least Squares Solvers ==============================#
#==============================  GELS Solver ==============================#
#= Solver For Symmetric Indefinite Matrices =#
import LinearAlgebra.LAPACK.gels! # (trans, A, B) -> (F, B, ssr)

struct GELSSolver <: AbstractSolver end
# Weights for the GELSSolver
@inline function W(::GELSSolver, distrib::AbstractDistribution, link::AbstractLink, mu::Array{T, 1}, eta::Array{T, 1}) where {T <: AbstractFloat}
  return ((deta_dmu(link, mu, eta).^2) .* variance(distrib, mu)).^(-0.5)
end
function W(::GELSSolver, distrib::AbstractDistribution, link::AbstractLink, mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}) where {T <: AbstractFloat}
  nBlocks::Int64 = length(mu)
  return [((deta_dmu(link, mu[i], eta[i]).^2) .* variance(distrib, mu[i])).^(-0.5) for i in 1:nBlocks]
end
function W(matrixType::Block1DParallel, ::GELSSolver, distrib::AbstractDistribution, link::AbstractLink, mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}) where {T <: AbstractFloat}
  nBlocks::Int64 = length(mu)
  ret::Array{Array{T, 1}, 1} = Array{Array{T, 1}, 1}(undef, nBlocks)
  @threads for i in 1:nBlocks
    ret[i] = ((deta_dmu(link, mu[i], eta[i]).^2) .* variance(distrib, mu[i])).^(-0.5)
  end
  return ret
end

# Solver
function solve!(::GELSSolver, R::Array{T, 2},
  xwx::Array{T, 2}, xw::Array{T, 2},
  x::Array{T, 2}, z::Array{T, 1}, w::Array{T, 1},
  coef::Array{T, 1}) where {T <: AbstractFloat}
  _, p = size(x)
  xw, zw = calcXW(xwx, x, w, z)
  F, coef, ssr = gels!('N', xw, zw)
  for j in 1:p
    for i in 1:j
      if i <= j
        R[i, j] = xw[i, j]
      end
    end
  end
  return (R, xwx, xw, x, z, w, coef)
end
function cov(::GELSSolver, invType::AbstractInverse, R::Array{T, 2}, xwx::Array{T, 2}, xw::Array{T, 2})::Array{T, 2} where {T <: AbstractFloat}
  xwx = R' * R
  return inv(invType, xwx)
end
#==============================  GELSY Solver ==============================#
import LinearAlgebra.LAPACK.gelsy! # (A, B, rcond) -> (B, rnk)
struct GELSYSolver <: AbstractSolver end
# Weights for the GELSYSolver
@inline function W(::GELSYSolver, distrib::AbstractDistribution, link::AbstractLink, mu::Array{T, 1}, eta::Array{T, 1}) where {T <: AbstractFloat}
  return ((deta_dmu(link, mu, eta).^2) .* variance(distrib, mu)).^(-0.5)
end
function W(::GELSYSolver, distrib::AbstractDistribution, link::AbstractLink, mu::Array{T, 1}, eta::Array{T, 1}) where {T <: AbstractFloat}
  nBlocks::Int64 = length(mu)
  return [((deta_dmu(link, mu[i], eta[i]).^2) .* variance(distrib, mu[i])).^(-0.5) for i in 1:nBlocks]
end
function W(matrixType::Block1DParallel, ::GELSYSolver, distrib::AbstractDistribution, link::AbstractLink, mu::Array{T, 1}, eta::Array{T, 1}) where {T <: AbstractFloat}
  nBlocks::Int64 = length(mu)
  ret::Array{Array{T, 1}, 1} = Array{Array{T, 1}, 1}(undef, nBlocks)
  @threads for i in 1:nBlocks
    ret[i] = ((deta_dmu(link, mu[i], eta[i]).^2) .* variance(distrib, mu[i])).^(-0.5)
  end
  return ret
end

# Solver
function solve!(::GELSYSolver, R::Array{T, 2},
  xwx::Array{T, 2}, xw::Array{T, 2},
  x::Array{T, 2}, z::Array{T, 1}, w::Array{T, 1},
  coef::Array{T, 1}) where {T <: AbstractFloat}
  _, p = size(x); rcond = T(0)
  xw, zw = calcXW(xwx, x, w, z)
  coef, rnk = gelsy!(copy(xw), zw, rcond)
  return (R, xwx, xw, x, z, w, coef)
end
function cov(::GELSYSolver, invType::AbstractInverse, R::Array{T, 2}, xwx::Array{T, 2}, xw::Array{T, 2})::Array{T, 2} where {T <: AbstractFloat}
  xwx = xw' * xw
  return inv(invType, xwx)
end
#==============================  GELSD Solver ==============================#
#= Solver For Symmetric Indefinite Matrices =#
import LinearAlgebra.LAPACK.gelsd! # (A, B, rcond) -> (B, rnk)

struct GELSDSolver <: AbstractSolver end
# Weights for the GELSDSolver
@inline function W(::GELSDSolver, distrib::AbstractDistribution, link::AbstractLink, mu::Array{T, 1}, eta::Array{T, 1}) where {T <: AbstractFloat}
  return ((deta_dmu(link, mu, eta).^2) .* variance(distrib, mu)).^(-0.5)
end
function W(::GELSDSolver, distrib::AbstractDistribution, link::AbstractLink, mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}) where {T <: AbstractFloat}
  nBlocks::Int64 = length(mu)
  return [((deta_dmu(link, mu[i], eta[i]).^2) .* variance(distrib, mu[i])).^(-0.5) for i in 1:nBlocks]
end
# Solver
function solve!(::GELSDSolver, R::Array{T, 2},
  xwx::Array{T, 2}, xw::Array{T, 2},
  x::Array{T, 2}, z::Array{T, 1}, w::Array{T, 1},
  coef::Array{T, 1}) where {T <: AbstractFloat}
  _, p = size(x)
  rcond::T = T(0)
  xw, zw = calcXW(xwx, x, w, z)
  coef, _ = gelsd!(copy(xw), zw, rcond)
  return (R, xwx, xw, x, z, w, coef)
end
function cov(::GELSDSolver, invType::AbstractInverse, R::Array{T, 2}, xwx::Array{T, 2}, xw::Array{T, 2})::Array{T, 2} where {T <: AbstractFloat}
  xwx = xw' * xw
  return inv(invType, xwx)
end

#==============================  GRADIENT DESCENT SOLVERS ==============================#
#=======================================================================================#

function XWX(x::Array{T, 2}, w::Array{T, 1}) where {T <: AbstractFloat}
  xw = copy(x)
  xw = multSweep(xw, w)
  xwx = xw' * x
  return xwx
end

function XWX(x::Array{Array{T, 2}, 1}, w::Array{Array{T, 1}, 1}) where {T <: AbstractFloat}
  p::Int64 = size(x[1])[2]
  nBlocks::Int64 = length(x)
  xwx = zeros(T, (p, p))
  for i in 1:nBlocks
    xw = copy(x[i])
    xw = multSweep(xw, w[i])
    xwx += xw' * x[i]
  end
  return xwx
end

function XWX(::Block1DParallel, x::Array{Array{T, 2}, 1}, w::Array{Array{T, 1}, 1}) where {T <: AbstractFloat}
  p::Int64 = size(x[1])[2]
  nBlocks::Int64 = length(x)
  xwx = [zeros(T, (p, p)) for i in 1:nthreads()]
  @threads for i in 1:nBlocks
    xw = copy(x[i])
    xw = multSweep(xw, w[i])
    xwx[threadid()] += xw' * x[i]
  end
  for i in 2:nthreads()
    xwx[1] .+= xwx[i]
  end
  return xwx[1]
end

abstract type AbstractGradientDescentSolver end

# Weights for the GradientDescentSolver
@inline function W(::AbstractGradientDescentSolver,
            distrib::AbstractDistribution, link::AbstractLink,
            mu::Array{T, 1}, eta::Array{T, 1}) where {T <: AbstractFloat}
  return ((deta_dmu(link, mu, eta).^2) .* variance(distrib, mu)).^(-1)
end
function W(::AbstractGradientDescentSolver, distrib::AbstractDistribution,
            link::AbstractLink, mu::Array{Array{T, 1}, 1},
            eta::Array{Array{T, 1}, 1}) where {T <: AbstractFloat}
  nBlocks::Int64 = length(mu)
  return [((deta_dmu(link, mu[i], eta[i]).^2) .* variance(distrib, mu[i])).^(-1) for i in 1:nBlocks]
end
function W(matrixType::Block1DParallel, ::AbstractGradientDescentSolver,
            distrib::AbstractDistribution, link::AbstractLink,
            mu::Array{Array{T, 1}, 1},
            eta::Array{Array{T, 1}, 1}) where {T <: AbstractFloat}
  nBlocks::Int64 = length(mu)
  ret::Array{Array{T, 1}, 1} = Array{Array{T, 1}, 1}(undef, nBlocks)
  @threads for i in 1:nBlocks
    ret[i] = ((deta_dmu(link, mu[i], eta[i]).^2) .* variance(distrib, mu[i])).^(-1)
  end
  return ret
end

# GD Covariance Function
function cov(::AbstractGradientDescentSolver, invType::AbstractInverse, xwx::Array{T, 2})::Array{T, 2} where {T <: AbstractFloat}
  return inv(invType, xwx)
end

#==============================  GRADIENTS ==============================#

#=
  Internal gradient function
=#
function _gradient(distrib::AbstractDistribution, link::AbstractLink,
            y::Array{T, 1}, x::Array{T, 2}, mu::Array{T, 1},
            eta::Array{T, 1}) where {T <: AbstractFloat}
  n, p = size(x)
  grad = zeros(T, p)
  numer::Array{T, 1} = y .- mu
  X2::T = sum((numer.^2) ./ variance(distrib, mu))
  tmp::Array{T, 1} = numer ./(deta_dmu(link, mu, eta) .* variance(distrib, mu))
  for j in 1:p
    grad[j] += sum(tmp .* x[:, j])
  end
  return (X2 = X2, grad = grad)
end

#=
  Gradient function for regular data
=#
function gradient(::RegularData, distrib::AbstractDistribution, link::AbstractLink,
            y::Array{T, 1}, x::Array{T, 2}, mu::Array{T, 1},
            eta::Array{T, 1}) where {T <: AbstractFloat}
  n, p = size(x)
  tmp = _gradient(distrib, link, y, x, mu, eta)
  df = n - p
  @assert(df > 0, "Number of items not greater than number of parameters")
  phi::T = tmp.X2/df
  grad = tmp.grad ./phi
  return grad
end

#=
  Gradient function for blocked data
=#
function gradient(::Block1D, distrib::AbstractDistribution, link::AbstractLink,
  y::Array{Array{T, 1}, 1}, x::Array{Array{T, 2}, 1},
  mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}) where {T <: AbstractFloat}

  nBlocks::Int64 = length(y); n::Int64 = 0
  p::Int64 = size(x[1])[2]; grad::Array{T, 1} = zeros(T, p)
  X2::T = T(0)

  for i in 1:nBlocks
    n += length(y[i])
    tmp = _gradient(distrib, link, y[i], x[i], mu[i], eta[i])
    # println("size(grad): ", size(grad))
    # println("size(tmp.grad): ", size(tmp.grad))
    grad .+= tmp.grad
    X2 += tmp.X2
  end

  df = n - p
  phi::T = X2/df
  @assert(df > 0, "Number of items not greater than number of parameters")
  grad ./= phi
  return grad
end

function gradient(::Block1DParallel, distrib::AbstractDistribution, link::AbstractLink,
  y::Array{Array{T, 1}, 1}, x::Array{Array{T, 2}, 1},
  mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}) where {T <: AbstractFloat}

  nBlocks::Int64 = length(y); p::Int64 = size(x[1])[2]; n::Int64 = 0;
  # grad::Array{T, 1} = zeros(T, p); X2::T = T(0)

  nStore = [Int64(0) for i in 1:nthreads()]
  X2Store = [T(0) for i in 1:nthreads()]
  gradStore = [Base.zeros(T, p) for i in 1:nthreads()]

  distribStore = [copy(distrib) for i in 1:nthreads()]
  linkStore = [copy(link) for i in 1:nthreads()]

  @threads for i in 1:nBlocks
    tmp = _gradient(distribStore[threadid()], linkStore[threadid()], y[i], x[i], mu[i], eta[i])
    gradStore[threadid()] .+= tmp.grad
    X2Store[threadid()] += tmp.X2
    nStore[threadid()] += length(y[i])
  end

  n = Int64(0); grad = zeros(T, p); X2 = T(0)
  for i in 1:nthreads()
    # nStore[1] += nStore[i]
    # gradStore[1] .+= gradStore[i]
    # X2Store[1] += X2Store[i]
    n += nStore[i]
    grad .+= gradStore[i]
    X2 += X2Store[i]
  end

  # df = nStore[1] - p
  # phi::T = X2Store[1]/df
  df = n - p
  phi::T = X2/df
  @assert(df > 0, "Number of items not greater than number of parameters")

  # return gradStore[1]./phi
  grad ./= phi
  return grad
end

#==============================  Simple Gradient Descent Solver ==============================#

mutable struct GradientDescentSolver{T} <: AbstractGradientDescentSolver
  learningRate::T
  function GradientDescentSolver(learningRate::T = T(1E-6)) where {T <: AbstractFloat}
    return new{T}(learningRate)
  end
end
#= Copy constructor =#
function copy(solver::GradientDescentSolver)
  return GradientDescentSolver(solver.learningRate)
end

function solve!(dataType::RegularData, solver::GradientDescentSolver, distrib::AbstractDistribution,
  link::AbstractLink, y::Array{T, 1}, x::Array{T, 2}, mu::Array{T, 1},
  eta::Array{T, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}

  grad = gradient(dataType, distrib, link, y, x, mu, eta);
  coef .+= solver.learningRate .* grad;
  return solver, coef
end

function solve!(dataType::Block1D, solver::GradientDescentSolver, distrib::AbstractDistribution,
  link::AbstractLink, y::Array{Array{T, 1}, 1}, x::Array{Array{T, 2}, 1},
  mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}

  grad = gradient(dataType, distrib, link, y, x, mu, eta);
  coef .+= solver.learningRate .* grad;
  return solver, coef
end

function solve!(dataType::Block1DParallel, solver::GradientDescentSolver,
  distrib::AbstractDistribution, link::AbstractLink,
  y::Array{Array{T, 1}, 1}, x::Array{Array{T, 2}, 1},
  mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}

  grad = gradient(dataType, distrib, link, y, x, mu, eta);
  coef .+= solver.learningRate .* grad;
  return solver, coef
end

#==============================  Gradient Descent Momentum Solver ==============================#

mutable struct MomentumSolver{T} <: AbstractGradientDescentSolver
  learningRate::T
  momentum::T
  delta::Array{T, 1}
  function MomentumSolver(learningRate::T, momentum::T, p::Int64) where {T <: AbstractFloat}
    return new{T}(learningRate, momentum, zeros(T, p))
  end
  function MomentumSolver(learningRate::T, momentum::T, delta::Array{T, 1}) where {T <: AbstractFloat}
    return new{T}(learningRate, momentum, delta)
  end
end
#= Copy constructor =#
function copy(solver::MomentumSolver)
  return MomentumSolver(solver.learningRate, solver.momentum, copy(solver.delta))
end

function solve!(dataType::RegularData, solver::MomentumSolver, distrib::AbstractDistribution,
  link::AbstractLink, y::Array{T, 1}, x::Array{T, 2}, mu::Array{T, 1},
  eta::Array{T, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}

  grad = gradient(dataType, distrib, link, y, x, mu, eta)
  solver.delta .= (solver.momentum .* solver.delta) .+ solver.learningRate .* grad
  coef .+= solver.delta
  return solver, coef
end

function solve!(dataType::Block1D, solver::MomentumSolver, distrib::AbstractDistribution,
  link::AbstractLink, y::Array{Array{T, 1}, 1}, x::Array{Array{T, 2}, 1},
  mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}

  grad = gradient(dataType, distrib, link, y, x, mu, eta)
  solver.delta .= (solver.momentum .* solver.delta) .+ solver.learningRate .* grad
  coef .+= solver.delta
  return solver, coef
end

function solve!(dataType::Block1DParallel, solver::MomentumSolver,
  distrib::AbstractDistribution, link::AbstractLink,
  y::Array{Array{T, 1}, 1}, x::Array{Array{T, 2}, 1},
  mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}

  grad = gradient(dataType, distrib, link, y, x, mu, eta)
  solver.delta .= (solver.momentum .* solver.delta) .+ solver.learningRate .* grad
  coef .+= solver.delta
  return solver, coef
end

#==============================  Gradient Descent Nesterov Solver ==============================#
#= Default Modify/Unmodify coefficients for Nesterov =#
function NesterovModifier(obj::AbstractGradientDescentSolver, coef::Array{T, 1}) where {T <: AbstractFloat}
  return coef
end
function NesterovUnModifier(obj::AbstractGradientDescentSolver, coef::Array{T, 1}) where {T <: AbstractFloat}
  return coef
end

mutable struct NesterovSolver{T} <: AbstractGradientDescentSolver
  learningRate::T
  momentum::T
  delta::Array{T, 1}
  function NesterovSolver(learningRate::T, momentum::T, p::Int64) where {T <: AbstractFloat}
    return new{T}(learningRate, momentum, zeros(T, p))
  end
  function NesterovSolver(learningRate::T, momentum::T, delta::Array{T, 1}) where {T <: AbstractFloat}
    return new{T}(learningRate, momentum, delta)
  end
end
#= Copy constructor =#
function copy(solver::NesterovSolver)
  return NesterovSolver(solver.learningRate, solver.momentum, copy(solver.delta))
end

function solve!(dataType::RegularData, solver::NesterovSolver, distrib::AbstractDistribution,
  link::AbstractLink, y::Array{T, 1}, x::Array{T, 2}, mu::Array{T, 1},
  eta::Array{T, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}

  grad = gradient(dataType, distrib, link, y, x, mu, eta)
  solver.delta .= (solver.momentum .* solver.delta) .+ solver.learningRate .* grad
  coef .+= solver.delta
  return solver, coef
end

function solve!(dataType::Block1D, solver::NesterovSolver, distrib::AbstractDistribution,
  link::AbstractLink, y::Array{Array{T, 1}, 1}, x::Array{Array{T, 2}, 1},
  mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}

  grad = gradient(dataType, distrib, link, y, x, mu, eta)
  solver.delta .= (solver.momentum .* solver.delta) .+ solver.learningRate .* grad
  coef .+= solver.delta
  return solver, coef
end

function solve!(dataType::Block1DParallel, solver::NesterovSolver,
  distrib::AbstractDistribution, link::AbstractLink,
  y::Array{Array{T, 1}, 1}, x::Array{Array{T, 2}, 1},
  mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}

  grad = gradient(dataType, distrib, link, y, x, mu, eta)
  solver.delta .= (solver.momentum .* solver.delta) .+ solver.learningRate .* grad
  coef .+= solver.delta
  return solver, coef
end

function NesterovModifier(obj::NesterovSolver, coef::Array{T, 1}) where {T <: AbstractFloat}
  coef .-= obj.momentum .* obj.delta
  return coef
end
function NesterovUnModifier(obj::NesterovSolver, coef::Array{T, 1}) where {T <: AbstractFloat}
  coef .+= obj.momentum .* obj.delta
  return coef
end

#==============================  Gradient Descent Adagrad Solver ==============================#
#= Modify/Unmodify coefficients for Adagrad =#
mutable struct AdagradSolver{T} <: AbstractGradientDescentSolver
  learningRate::T
  G::Array{T, 1}
  epsilon::T
  function AdagradSolver(learningRate::T, p::Int64, epsilon::T) where {T <: AbstractFloat}
    return new{T}(learningRate, zeros(T, p), epsilon)
  end
  function AdagradSolver(learningRate::T, G::Array{T, 1}, epsilon::T) where {T <: AbstractFloat}
    return new{T}(learningRate, G, epsilon)
  end
end
#= Copy constructor =#
function copy(solver::AdagradSolver)
  return AdagradSolver(solver.learningRate, copy(solver.G), solver.epsilon)
end

function solve!(dataType::RegularData, solver::AdagradSolver, distrib::AbstractDistribution,
  link::AbstractLink, y::Array{T, 1}, x::Array{T, 2}, mu::Array{T, 1},
  eta::Array{T, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}

  grad = gradient(dataType, distrib, link, y, x, mu, eta)
  solver.G += coef.^2
  coef .+= (solver.learningRate .* grad)./((solver.G .+ solver.epsilon).^0.5)
  return solver, coef
end

function solve!(dataType::Block1D, solver::AdagradSolver, distrib::AbstractDistribution,
  link::AbstractLink, y::Array{Array{T, 1}, 1}, x::Array{Array{T, 2}, 1},
  mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}

  grad = gradient(dataType, distrib, link, y, x, mu, eta)
  solver.G += coef.^2
  coef .+= (solver.learningRate .* grad)./((solver.G .+ solver.epsilon).^0.5)
  return solver, coef
end

function solve!(dataType::Block1DParallel, solver::AdagradSolver,
  distrib::AbstractDistribution, link::AbstractLink,
  y::Array{Array{T, 1}, 1}, x::Array{Array{T, 2}, 1},
  mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}

  grad = gradient(dataType, distrib, link, y, x, mu, eta)
  solver.G += coef.^2
  coef .+= (solver.learningRate .* grad)./((solver.G .+ solver.epsilon).^0.5)
  return solver, coef
end

#==============================  Gradient Descent Adadelta Solver ==============================#
#= Modify/Unmodify coefficients for Adagrad =#
mutable struct AdadeltaSolver{T} <: AbstractGradientDescentSolver
  momentum::T
  G::Array{T, 1}
  B::Array{T, 1}
  epsilon::T
  function AdadeltaSolver(momentum::T, p::Int64, epsilon::T) where {T <: AbstractFloat}
    return new{T}(momentum, zeros(T, p), zeros(T, p), epsilon)
  end
  function AdadeltaSolver(momentum::T, G::Array{T, 1}, B::Array{T, 1}, epsilon::T) where {T <: AbstractFloat}
    return new{T}(momentum, G, B, epsilon)
  end
end
#= Copy constructor =#
function copy(solver::AdadeltaSolver)
  return AdadeltaSolver(solver.momentum, copy(solver.G), copy(solver.B), solver.epsilon)
end

macro AdadeltaUpdate()
  expr = quote
    grad = gradient(dataType, distrib, link, y, x, mu, eta)
    momentumM1 = (1 - solver.momentum)
    solver.G = (solver.momentum .* solver.G) .+ (momentumM1 .* (grad.^2))
    diff = (grad .* ((solver.B .+ solver.epsilon).^0.5)) ./((solver.G .+ solver.epsilon).^0.5)
    coef .+= diff
    solver.B .+= (solver.momentum .* solver.B) .+ (momentumM1 .* (diff.^2))
    return solver, coef
  end
  return esc(expr)
end

function solve!(dataType::RegularData, solver::AdadeltaSolver, distrib::AbstractDistribution,
  link::AbstractLink, y::Array{T, 1}, x::Array{T, 2}, mu::Array{T, 1},
  eta::Array{T, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}

  @AdadeltaUpdate
end

function solve!(dataType::Block1D, solver::AdadeltaSolver, distrib::AbstractDistribution,
  link::AbstractLink, y::Array{Array{T, 1}, 1}, x::Array{Array{T, 2}, 1},
  mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}

  @AdadeltaUpdate
end

function solve!(dataType::Block1DParallel, solver::AdadeltaSolver,
  distrib::AbstractDistribution, link::AbstractLink,
  y::Array{Array{T, 1}, 1}, x::Array{Array{T, 2}, 1},
  mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}

  @AdadeltaUpdate
end

#==============================  Gradient Descent RMSprop Solver ==============================#
#= Modify/Unmodify coefficients for Adagrad =#
mutable struct RMSpropSolver{T} <: AbstractGradientDescentSolver
  learningRate::T
  G::Array{T, 1}
  epsilon::T
  function RMSpropSolver(learningRate::T, p::Int64, epsilon::T) where {T <: AbstractFloat}
    return new{T}(learningRate, zeros(T, p), epsilon)
  end
  function RMSpropSolver(learningRate::T, G::Array{T, 1}, epsilon::T) where {T <: AbstractFloat}
    return new{T}(learningRate, G, epsilon)
  end
end
#= Copy constructor =#
function copy(solver::RMSpropSolver)
  return RMSpropSolver(solver.learningRate, copy(solver.G), solver.epsilon)
end

macro RMSpropUpdate()
  expr = quote
    grad = gradient(dataType, distrib, link, y, x, mu, eta)
    solver.G = (0.9 .* solver.G) .+ (0.1 .* (grad.^2))
    diff = (solver.learningRate .* grad) ./((solver.G .+ solver.epsilon).^0.5)
    coef .+= diff
    return solver, coef
  end
  return esc(expr)
end

function solve!(dataType::RegularData, solver::RMSpropSolver, distrib::AbstractDistribution,
  link::AbstractLink, y::Array{T, 1}, x::Array{T, 2}, mu::Array{T, 1},
  eta::Array{T, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}

  @RMSpropUpdate
end

function solve!(dataType::Block1D, solver::RMSpropSolver, distrib::AbstractDistribution,
  link::AbstractLink, y::Array{Array{T, 1}, 1}, x::Array{Array{T, 2}, 1},
  mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}

  @RMSpropUpdate
end

function solve!(dataType::Block1DParallel, solver::RMSpropSolver,
  distrib::AbstractDistribution, link::AbstractLink,
  y::Array{Array{T, 1}, 1}, x::Array{Array{T, 2}, 1},
  mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}

  @RMSpropUpdate
end

#==============================  Gradient Descent Adam Solver ==============================#
function iteration(obj::AbstractGradientDescentSolver, iter::Int64)
  return obj
end
#= Adam =#
mutable struct AdamSolver{T} <: AbstractGradientDescentSolver
  learningRate::T
  b1::T
  b2::T
  m::Array{T, 1}
  v::Array{T, 1}
  epsilon::T
  iter::Int64
  function AdamSolver(learningRate::T, b1::T, b2::T, p::Int64, epsilon::T) where {T <: AbstractFloat}
    return new{T}(learningRate, b1, b2, zeros(T, p), zeros(T, p), epsilon, Int64(1))
  end
  function AdamSolver(learningRate::T, b1::T, b2::T, m::Array{T, 1}, v::Array{T, 1}, epsilon::T, iter::Int64) where {T <: AbstractFloat}
    return new{T}(learningRate, b1, b2, m, v, epsilon, iter)
  end
end
#= Copy constructor =#
function copy(solver::AdamSolver)
  return AdamSolver(solver.learningRate, solver.b1, solver.b2, copy(solver.m), copy(solver.v), solver.epsilon, solver.iter)
end
function iteration(obj::AdamSolver, iter::Int64)
  obj.iter = iter
  return obj
end

macro AdamUpdate()
  expr = quote
    grad = gradient(dataType, distrib, link, y, x, mu, eta)
    solver.m = (solver.b1 .* solver.m) .+ ((1 - solver.b1) .* grad)
    solver.v = (solver.b2 .* solver.v) .+ ((1 - solver.b2) .* (grad.^2))
    mp = solver.m ./(1 - (solver.b1^solver.iter))
    vp = solver.v ./(1 - (solver.b2^solver.iter))
    diff = (solver.learningRate .* mp)./(solver.epsilon .+ (vp.^0.5))
    coef .+= diff
    return solver, coef
  end
  return esc(expr)
end

function solve!(dataType::RegularData, solver::AdamSolver, distrib::AbstractDistribution,
  link::AbstractLink, y::Array{T, 1}, x::Array{T, 2}, mu::Array{T, 1},
  eta::Array{T, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}

  @AdamUpdate
end

function solve!(dataType::Block1D, solver::AdamSolver, distrib::AbstractDistribution,
  link::AbstractLink, y::Array{Array{T, 1}, 1}, x::Array{Array{T, 2}, 1},
  mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}

  @AdamUpdate
end

function solve!(dataType::Block1DParallel, solver::AdamSolver,
  distrib::AbstractDistribution, link::AbstractLink,
  y::Array{Array{T, 1}, 1}, x::Array{Array{T, 2}, 1},
  mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}

  @AdamUpdate
end

#==============================  Gradient Descent AdaMax Solver ==============================#
#= AdaMax =#
mutable struct AdaMaxSolver{T} <: AbstractGradientDescentSolver
  learningRate::T
  b1::T
  b2::T
  m::Array{T, 1}
  v::T
  iter::Int64
  function AdaMaxSolver(learningRate::T, b1::T, b2::T, p::Int64) where {T <: AbstractFloat}
    return new{T}(learningRate, b1, b2, zeros(T, p), T(0), Int64(1))
  end
  function AdaMaxSolver(learningRate::T, b1::T, b2::T, m::Array{T, 1}, v::T, iter::Int64) where {T <: AbstractFloat}
    return new{T}(learningRate, b1, b2, m, v, iter)
  end
end
#= Copy constructor =#
function copy(solver::AdaMaxSolver)
  return AdaMaxSolver(solver.learningRate, solver.b1, solver.b2, copy(solver.m), solver.v, solver.iter)
end
function iteration(obj::AdaMaxSolver, iter::Int64)
  obj.iter = iter
  return obj
end

macro AdaMaxUpdate()
  expr = quote
    grad = gradient(dataType, distrib, link, y, x, mu, eta)

    absG = norm(grad)
    solver.v = (solver.b2 * solver.v) + ((1 - solver.b2) * absG)
    u = max(solver.b2 * solver.v, absG)
    solver.m = (solver.b1 .* solver.m) .+ ((1 - solver.b1) .* grad)
    mp = solver.m ./(1 - (solver.b1^solver.iter))
    diff = (solver.learningRate .* mp)./u

    coef .+= diff
    return solver, coef
  end
  return esc(expr)
end

function solve!(dataType::RegularData, solver::AdaMaxSolver, distrib::AbstractDistribution,
  link::AbstractLink, y::Array{T, 1}, x::Array{T, 2}, mu::Array{T, 1},
  eta::Array{T, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}

  @AdaMaxUpdate
end

function solve!(dataType::Block1D, solver::AdaMaxSolver, distrib::AbstractDistribution,
  link::AbstractLink, y::Array{Array{T, 1}, 1}, x::Array{Array{T, 2}, 1},
  mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}

  @AdaMaxUpdate
end

function solve!(dataType::Block1DParallel, solver::AdaMaxSolver,
  distrib::AbstractDistribution, link::AbstractLink,
  y::Array{Array{T, 1}, 1}, x::Array{Array{T, 2}, 1},
  mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}

  @AdaMaxUpdate
end

#==============================  Gradient Descent NAdam Solver ==============================#
#= NAdam =#
mutable struct NAdamSolver{T} <: AbstractGradientDescentSolver
  learningRate::T
  b1::T
  b2::T
  m::Array{T, 1}
  v::Array{T, 1}
  epsilon::T
  iter::Int64
  function NAdamSolver(learningRate::T, b1::T, b2::T, p::Int64, epsilon::T) where {T <: AbstractFloat}
    return new{T}(learningRate, b1, b2, zeros(T, p), zeros(T, p), epsilon, Int64(1))
  end
  function NAdamSolver(learningRate::T, b1::T, b2::T, m::Array{T, 1}, v::Array{T, 1}, epsilon::T, iter::Int64) where {T <: AbstractFloat}
    return new{T}(learningRate, b1, b2, m, v, epsilon, iter)
  end
end
#= Copy constructor =#
function copy(solver::NAdamSolver)
  return NAdamSolver(solver.learningRate, solver.b1, solver.b2, copy(solver.m), copy(solver.v), solver.epsilon, solver.iter)
end
function iteration(obj::NAdamSolver, iter::Int64)
  obj.iter = iter
  return obj
end

macro NAdamUpdate()
  expr = quote
    grad = gradient(dataType, distrib, link, y, x, mu, eta)

    solver.m .= (solver.b1 .* solver.m) .+ ((1 - solver.b1) .* grad)
    mp = solver.m ./(1 - (solver.b1^solver.iter))
    solver.v .= (solver.b2 .* solver.v) .+ ((1 - solver.b2) .* (grad.^2))
    vp = solver.v ./(1 - (solver.b2.^solver.iter))
    tmp1 = solver.learningRate ./((vp.^0.5) .+ solver.epsilon);
    tmp2 = (solver.b1 .* mp) .+ ((1 - solver.b1) .* grad) ./(1 - (solver.b1^solver.iter))

    coef .+= (tmp1 .* tmp2)

    return solver, coef
  end
  return esc(expr)
end

function solve!(dataType::RegularData, solver::NAdamSolver, distrib::AbstractDistribution,
  link::AbstractLink, y::Array{T, 1}, x::Array{T, 2}, mu::Array{T, 1},
  eta::Array{T, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}

  @NAdamUpdate
end

function solve!(dataType::Block1D, solver::NAdamSolver, distrib::AbstractDistribution,
  link::AbstractLink, y::Array{Array{T, 1}, 1}, x::Array{Array{T, 2}, 1},
  mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}

  @NAdamUpdate
end

function solve!(dataType::Block1DParallel, solver::NAdamSolver,
  distrib::AbstractDistribution, link::AbstractLink,
  y::Array{Array{T, 1}, 1}, x::Array{Array{T, 2}, 1},
  mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}

  @NAdamUpdate
end

#==============================  Gradient Descent AMSGrad Solver ==============================#
#= AMSGrad =#
mutable struct AMSGradSolver{T} <: AbstractGradientDescentSolver
  learningRate::T
  b1::T
  b2::T
  m::Array{T, 1}
  v::Array{T, 1}
  vp::Array{T, 1}
  epsilon::T
  iter::Int64
  function AMSGradSolver(learningRate::T, b1::T, b2::T, p::Int64, epsilon::T) where {T <: AbstractFloat}
    return new{T}(learningRate, b1, b2, zeros(T, p), zeros(T, p), zeros(T, p), epsilon, Int64(1))
  end
  function AMSGradSolver(learningRate::T, b1::T, b2::T, m::Array{T, 1}, v::Array{T, 1}, vp::Array{T, 1}, epsilon::T, iter::Int64) where {T <: AbstractFloat}
    return new{T}(learningRate, b1, b2, m, v, vp, epsilon, iter)
  end
end
#= Copy constructor =#
function copy(solver::AMSGradSolver)
  return AMSGradSolver(solver.learningRate, solver.b1, solver.b2, copy(solver.m), copy(solver.v), copy(solver.vp), solver.epsilon, solver.iter)
end
function iteration(obj::AMSGradSolver, iter::Int64)
  obj.iter = iter
  return obj
end

macro AMSGradUpdate()
  expr = quote
    grad = gradient(dataType, distrib, link, y, x, mu, eta)

    solver.m .= (solver.b1 .* solver.m) .+ ((1 - solver.b1) .* grad)
    solver.v .= (solver.b2 .* solver.v) .+ ((1 - solver.b2) .* (grad.^2))
    solver.vp .= max.(solver.vp, solver.v)
    diff = (solver.learningRate .* solver.m)./((solver.vp .^0.5) .+ solver.epsilon)

    coef .+= diff

    return solver, coef
  end
  return esc(expr)
end

function solve!(dataType::RegularData, solver::AMSGradSolver, distrib::AbstractDistribution,
  link::AbstractLink, y::Array{T, 1}, x::Array{T, 2}, mu::Array{T, 1},
  eta::Array{T, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}

  @AMSGradUpdate
end

function solve!(dataType::Block1D, solver::AMSGradSolver, distrib::AbstractDistribution,
  link::AbstractLink, y::Array{Array{T, 1}, 1}, x::Array{Array{T, 2}, 1},
  mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}

  @AMSGradUpdate
end

function solve!(dataType::Block1DParallel, solver::AMSGradSolver,
  distrib::AbstractDistribution, link::AbstractLink,
  y::Array{Array{T, 1}, 1}, x::Array{Array{T, 2}, 1},
  mu::Array{Array{T, 1}, 1}, eta::Array{Array{T, 1}, 1}, coef::Array{T, 1}) where {T <: AbstractFloat}

  @AMSGradUpdate
end

#==============================  GRADIENT DESCENT INITIALIZERS ==============================#

# init type return.
InitDouble{T} = Tuple{Array{T, 1}, Array{T, 1}} # where {T <: AbstractFloat}

# Initializer for gradient descent solvers
function init!(::AbstractGradientDescentSolver,
               ::AbstractDistribution, y::Array{T},
               wts::Array{T, 1}) #= ::InitDouble{T} =# where {T <: AbstractFloat}
  y = y[:, 1]
  return y, wts
end

function init!(::AbstractGradientDescentSolver,
               ::AbstractDistribution, y::Array{Array{T, 2}, 1},
               wts::Array{Array{T, 1}, 1}) #= ::InitDouble{Array{T, 1}}=# where {T <: AbstractFloat}
  nBlocks::Int64 = length(y)
  if length(size(y[1])) == 2
    y = [y[i][:, 1] for i in 1:nBlocks]
  end
  return y, wts;
end

function init!(::AbstractGradientDescentSolver,
               ::Block1DParallel, ::AbstractDistribution,
               y::Array{Array{T, 2}, 1},
               wts::Array{Array{T, 1}, 1}) #=::InitDouble{Array{T, 1}} =# where {T <: AbstractFloat}
  nBlocks::Int64 = length(y)
  isMat::Bool = length(size(y[1])) == 2
  yNew = Array{Array{T, 1}, 1}(undef, nBlocks);
  @threads for i in 1:nBlocks
    yNew[i] = isMat ? y[i][:, 1] : y[i]
  end
  return yNew, wts;
end

#===================== BinomialDistribution Initializers ====================#
function init!(::AbstractGradientDescentSolver,
      ::BinomialDistribution, y::Array{T},
      wts::Array{T, 1}) #= ::InitDouble{T}=# where {T <: AbstractFloat}
  if size(y)[2] == 1
    y = y[:, 1]
  elseif size(y)[2] == 2
    n = size(y)[1]
    events = y[:, 1]
    N = events .+ y[:, 2]
    y = [if N[i] != 0; events[i]/N[i]; else T(0) end for i in 1:n]
    if length(wts) != 0
      wts .*= N
    else
      wts = N
    end
  else
    error("There was a problem")
  end
  return y, wts
end

function init!(::AbstractGradientDescentSolver,
               ::BinomialDistribution, y::Array{Array{T, 2}, 1},
               wts::Array{Array{T, 1}, 1}) #= ::InitDouble{Array{T, 1}} =# where {T <: AbstractFloat}
  nBlocks::Int64 = length(y)
  if size(y[1])[2] == 1
    y = [y[i][:, 1] for i in 1:nBlocks]
  elseif size(y[1])[2] == 2
    events = [y[i][:, 1] for i in 1:nBlocks]
    N = [events[i] .+ y[i][:, 2] for i in 1:nBlocks]
    yNew = Array{Array{T, 1}, 1}(undef, nBlocks)
    for i in 1:nBlocks
      n = size(y[i])[1]
      tmp = zeros(T, n)
      for j in 1:n
        tmp[j] = if N[i][j] != 0; events[i][j]/N[i][j]; else T(0) end
      end
      yNew[i] = tmp
    end
    y = yNew
    if length(wts) != 0
      [wts[i] .*= N[i] for i in 1:nBlocks]
    else
      wts = N
    end
  else
    error("There was a problem with the dimensions of y")
  end
  return y, wts
end

function init!(::AbstractGradientDescentSolver,
          ::Block1DParallel, ::BinomialDistribution,
          y::Array{Array{T, 2}, 1},
          wts::Array{Array{T, 1}, 1}) #= ::InitDouble{Array{T, 1}} =# where {T <: AbstractFloat}
  nBlocks::Int64 = length(y)
  yNew = Array{Array{T, 1}, 1}(undef, nBlocks)
  events = Array{Array{T, 1}, 1}(undef, nBlocks)
  N = Array{Array{T, 1}, 1}(undef, nBlocks)
  if size(y[1])[2] == 1
    @threads for i in 1:nBlocks
      yNew[i] = y[i][:, 1]
    end
    y = yNew
  elseif size(y[1])[2] == 2
    @threads for i in 1:nBlocks
      events[i] = y[i][:, 1]
    end
    @threads for i in 1:nBlocks
      N[i] = events[i] .+ y[i][:, 2]
    end
    @threads for i in 1:nBlocks
      n = size(y[i])[1]
      tmp = zeros(T, n)
      for j in 1:n
        tmp[j] = if N[i][j] != 0; events[i][j]/N[i][j]; else T(0) end
      end
      yNew[i] = tmp
    end
    y = yNew
    if length(wts) != 0
      @threads for i in 1:nBlocks
        wts[i] = wts[i] .* N[i]
      end
    else
      wts = N
    end
  else
    error("There was a problem with the dimensions of y")
  end
  return y, wts
end


using YaoBlocks.AD
using SparseArrays, LuxurySparse
using YaoArrayRegister
using Test

@testset "outer prod" begin
    D = 8
    T = ComplexF64
    for (l, r) in [(randn(T, D), randn(T, D)), (randn(T, D, 3), randn(T, D, 3))]
        a = outerprod(l, r)
        A = Matrix(a)
        B = randn(T, D, D)
        V = randn(T, D)
        @test A * B ≈ a * B
        @test A * A ≈ a * a
        @test B * A ≈ B * a
        @test V' * A ≈ V' * a
        @test A * V ≈ a * V
        @test size(a) == size(A)
        @test size(a, 2) == size(A, 2)
        @test A[4] == a[4]
        @test A[4, 4] == a[4, 4]
        @test A' == a'
        @test transpose(A) == transpose(a)
        @test conj(A) == conj(a)
        @test conj!(A) == conj!(a)
        outerprod(ArrayReg(l), ArrayReg(r)) == a
    end
end

@testset "projection" begin
    T = ComplexF64
    D = 8
    for y in [pmrand(T, D), sprand(T, D, D, 0.5), Diagonal(randn(T, D))]
        @test projection(y, y) == y
        l, r = randn(T, D), randn(T, D)
        op = outerprod(l, r)
        @test projection(y, op) == projection(y, Matrix(op))
    end
end


using QMatrices
using LinearAlgebra: LinearAlgebra

@testset "Basic algebraic props" begin
    for m in (X, Y, Z, H)
        @test isapprox(m^2, LinearAlgebra.I) # TODO: (slow) norm is unnecessary
    end
end

@testset "control" begin
    toffoli = [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]

    @test control(X, 2) == toffoli
end

@testset "Bell state" begin
    k = zeroket(2)
    @test CX * (H ⊗ I2) * k == [0.7071067811865475, 0.0, 0.0, 0.7071067811865475]
end


struct Cost{T}
    #TODO: types for methods
    evaluate
    gradient_state
    gradient_action
    hessian_state_state
    hessian_action_action
    hessian_action_state
    evaluate_cache::Vector{T}
    gradient_state_cache::Vector{T}
    gradient_action_cache::Vector{T}
    hessian_state_state_cache::Matrix{T}
    hessian_action_action_cache::Matrix{T}
    hessian_action_state_cache::Matrix{T}
end

function Cost(f::Function, num_state::Int, num_action::Int; num_parameter::Int=0)
    #TODO: option to load/save methods
    @variables x[1:num_state], u[1:num_action], w[1:num_parameter]

    evaluate = f(x, u, w)
    gradient_state = Symbolics.gradient(evaluate, x)
    gradient_action = Symbolics.gradient(evaluate, u)
    hessian_state_state = Symbolics.jacobian(gradient_state, x)
    hessian_action_action = Symbolics.jacobian(gradient_action, u)
    hessian_action_state = Symbolics.jacobian(gradient_action, x)

    evaluate_func = eval(Symbolics.build_function([evaluate], x, u, w)[2])
    gradient_state_func = eval(Symbolics.build_function(gradient_state, x, u, w)[2])
    gradient_action_func = eval(Symbolics.build_function(gradient_action, x, u, w)[2])
    hessian_state_state_func = eval(Symbolics.build_function(hessian_state_state, x, u, w)[2])
    hessian_action_action_func = eval(Symbolics.build_function(hessian_action_action, x, u, w)[2])
    hessian_action_state_func = eval(Symbolics.build_function(hessian_action_state, x, u, w)[2])

    return Cost(evaluate_func,
        gradient_state_func, gradient_action_func,
        hessian_state_state_func, hessian_action_action_func, hessian_action_state_func,
        zeros(1),
        zeros(num_state), zeros(num_action),
        zeros(num_state, num_state), zeros(num_action, num_action), zeros(num_action, num_state))
end

Objective{T} = Vector{Cost{T}} where T

function cost(costs::Vector{Cost{T}}, states, actions, parameters) where T
    J = 0.0
    for (t, cost) in enumerate(costs)
        cost.evaluate(cost.evaluate_cache, states[t], actions[t], parameters[t])
        J += cost.evaluate_cache[1]
    end
    return J
end

function cost_gradient!(gradient_states, gradient_actions, costs::Vector{Cost{T}}, states, actions, parameters) where T
    H = length(costs)
    for (t, cost) in enumerate(costs)
        cost.gradient_state(cost.gradient_state_cache, states[t], actions[t], parameters[t])
        @views gradient_states[t] .= cost.gradient_state_cache
        fill!(cost.gradient_state_cache, 0.0) # TODO: confirm this is necessary
        t == H && continue
        cost.gradient_action(cost.gradient_action_cache, states[t], actions[t], parameters[t])
        @views gradient_actions[t] .= cost.gradient_action_cache
        fill!(cost.gradient_action_cache, 0.0) # TODO: confirm this is necessary
    end
end

function cost_hessian!(hessian_state_state, hessian_action_action, hessian_action_state, costs::Vector{Cost{T}}, states, actions, parameters) where T
    H = length(costs)
    for (t, cost) in enumerate(costs)
        cost.hessian_state_state(cost.hessian_state_state_cache, states[t], actions[t], parameters[t])
        @views hessian_state_state[t] .+= cost.hessian_state_state_cache
        fill!(cost.hessian_state_state_cache, 0.0) # TODO: confirm this is necessary
        t == H && continue
        cost.hessian_action_action(cost.hessian_action_action_cache, states[t], actions[t], parameters[t])
        cost.hessian_action_state(cost.hessian_action_state_cache, states[t], actions[t], parameters[t])
        @views hessian_action_action[t] .+= cost.hessian_action_action_cache
        @views hessian_action_state[t]  .+= cost.hessian_action_state_cache
        fill!(cost.hessian_action_action_cache, 0.0) # TODO: confirm this is necessary
        fill!(cost.hessian_action_state_cache, 0.0) # TODO: confirm this is necessary
    end
end


module IntrinsicMetrics

using NearestNeighbors
using StaticArrays
using Distances
using LightGraphs, SimpleWeightedGraphs
using RecipesBase

include("intrinsic_metric.jl")
include("plotting.jl")

export IntrinsicMetric,
    ambientmetric, adjgraph, points, npoints, ambientdim, points3d,
    evaluate, pairwise!, pairwise, colwise!, colwise,
    Euclidean, Minkowski, Cityblock, Chebyshev

end


@testset "kl_divergence" begin
    p = MvNormal(zeros(2), ones(2))
    q = MvNormal(zeros(2), ones(2))
    @test all(kl_divergence(p, q) .≈ 0.0)

    cp = ConditionalMvNormal(identity)

    kld = kl_divergence(cp, q, zeros(2,10))
    @test size(kld) == (1,10)
    @test all(kld .≈ 0.0)
end


# LuGreModels.jl

# The Lu-Gre state equation (bristle velocity)
function heuristic_state_equation(
    mdl::LuGreModel,
    t::T,
    nrm::T,
    pos::T,
    vel::T,
    z::Array{T}
) where T <: Number
    z0 = z[1]
    Fc = mdl.coulomb_coefficient * nrm
    Fs = mdl.static_coefficient * nrm
    g = Fc + (Fs - Fc) * exp(-(vel / mdl.stribeck_velocity)^2)
    return [vel - mdl.bristle_stiffness * abs(vel) * z0 / g]
end

# The Lu-Gre friction force equation.
function heuristic_force_equation(
    mdl::LuGreModel,
    t::T,
    nrm::T,
    pos::T,
    vel::T,
    z::Array{T},
    dzdt::Array{T}
) where T <: Number
    return mdl.bristle_stiffness * z[1] + mdl.bristle_damping * dzdt[1] +
        mdl.viscous_damping * vel
end

# ------------------------------------------------------------------------------
# Required routines to support friction model fitting.

function model_from_array(mdl::LuGreModel, x::Array{T}) where T <: Number
    LuGreModel(x[1], x[2], x[3], x[4], x[5], x[6])
end

function model_to_array(x::LuGreModel)
    [
        x.static_coefficient,
        x.coulomb_coefficient,
        x.stribeck_velocity,
        x.bristle_stiffness,
        x.bristle_damping,
        x.viscous_damping
    ]
end


using LinearAlgebra
### Solver Options ###
opts = TrajectoryOptimization.SolverOptions()
opts.verbose = true
######################

### Set up model, objective, solver ###
# Model
dt = 0.1
traj_folder = joinpath(dirname(pathof(TrajectoryOptimization)),"..")
urdf_folder = joinpath(traj_folder, "dynamics/urdf")
urdf_cartpole = joinpath(urdf_folder, "cartpole.urdf")

model_urdf = Model(urdf_cartpole,[1.;0.])
n = model.n
m = model.m
# Objective
Q = 0.01*Diagonal(I,n)
Qf = 1000.0*Diagonal(I,n)
R = 0.01*Diagonal(I,m)

x0 = [0.;pi;0.;0.]
xf = [0.;0.;0.;0.]

tf = 5.0

obj_uncon = LQRObjective(Q, R, Qf, tf, x0, xf)

# -Constraints
u_min = -20
u_max = 1000
x_min = [-1000; -1000; -1000; -1000]
x_max = [1000; 1000; 1000; 1000]

obj_con = TrajectoryOptimization.ConstrainedObjective(obj_uncon, u_min=u_min, u_max=u_max, x_min=x_min, x_max=x_max) # constrained objective

# Solver (foh & zoh)
# solver_foh = Solver(model, obj_con, integration=:rk3_foh, dt=dt, opts=opts)
solver_zoh = Solver(model, obj_uncon, integration=:rk3, dt=dt, opts=opts)

# -Initial control and state trajectories
U = ones(solver_zoh.model.m,solver_zoh.N)
X_interp = line_trajectory(solver_zoh)
X_rollout = TrajectoryOptimization.rollout(solver,U)
# X_interp = ones(solver_zoh.model.n,solver.N)
#######################################

### Solve ###
# @time sol_foh, = TrajectoryOptimization.solve(solver_foh,X_interp,U)
solver_zoh.opts.verbose = false
@time sol_zoh, = TrajectoryOptimization.solve(solver_zoh,U)
solve_dircol(solver_zoh,X_rollout,U,method=:hermite_simpson)
#############

# ### Results ###
# if opts.verbose
#     # println("Final state (foh): $(sol_foh.X[:,end])")
#     println("Final state (zoh): $(sol_zoh.X[:,end])")
#
#     # println("Termination index\n foh: $(sol_foh.termination_index)\n zoh: $(sol_foh.termination_index)")
#
#     # println("Final cost (foh): $(sol_foh.cost[sol_foh.termination_index])")
#     println("Final cost (zoh): $(sol_zoh.cost[sol_zoh.termination_index])")
#
#     # plot((sol_foh.cost[1:sol_foh.termination_index]))
#     plot!((sol_zoh.cost[1:sol_zoh.termination_index]))
#
#     # plot(sol_foh.U')
#     plot!(sol_zoh.U')
#
#     # plot(sol_foh.X[1:2,:]')
#     plot!(sol_zoh.X[1:2,:]')
# end
###############


module TimeseriesSurrogates

using Random
using Distributions
using Distances # Will be used by the LombScargle method
using StatsBase
using AbstractFFTs
using DSP
using Interpolations
using Wavelets
using Requires

include("api.jl")

include("utils/testsystems.jl")

# The different surrogate routines
include("methods/randomshuffle.jl")
include("methods/large_shuffle.jl")
include("methods/randomfourier.jl")
include("methods/aaft.jl")
include("methods/iaaft.jl")
include("methods/truncated_fourier.jl")
include("methods/partial_randomization.jl")
include("methods/wavelet_based.jl")
include("methods/pseudoperiodic.jl")
include("methods/pseudoperiodic_twin.jl")
include("methods/multidimensional.jl")
include("methods/ar.jl")
include("methods/trend_based.jl")

# Methods for irregular time series
include("methods/lombscargle.jl")

# Visualization routine for time series + surrogate + periodogram/acf/histogram
using Requires
function __init__()
    @require UncertainData="dcd9ba68-c27b-5cea-ae21-829cd07325bf" begin
        include("utils/uncertaindatasets.jl")
    end
    @require Makie="ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a" begin
        include("plotting/surrogate_plot.jl")
    end
end

export surroplot_path
function surroplot_path()
    joinpath(@__DIR__, "plotting", "surrogate_plot.jl")
end

end # module


# TODO test jldoctests
function genGaussWGrid(aG, mG, AG;
                  xmin=nothing, xmax=nothing, nPts=nothing)
    """
    Calculate sum of Gaussians given amplitudes, means, and covariance matrices.

    Evaluates nGaussian Gaussian functions on a grid, returns their sum.
    The grid models space with nPts discretization points along nDims dimensions.
    LaTeX pseudocode:

       sum_i aG_i * exp( - (x-mG_i)' * AG_i * (x-mG_i)).

    Arguments nGauss and nDims are not declared explicitely, but calculated from
    means vector mG, see below.

    # Arguments
    - `aG::Array{Float64}`: amplitudes of Gaussians, vector of length nGauss
    - `mG::Array{Float64}`: coordinates of Gaussians,
                            matrix of size nDims x nGauss
    - `AG::Array{Float64}`: covariance matrices of Gaussians,
                            tensor of size nDims x nDims x nGauss

    # Keyword Arguments
    - `xmin[=-ones(nDims)]`: lower bound of the discretization box
    - `xmax[= ones(nDims)]`: upper bound of the discretization box
    - `nPts[= 67 .* ones(Int64, nDims)]`: number of discretization points
                                          along each dimension

    # Examples for GenGauss (different handle for this function, cf. below)
    ```jldoctest
    julia> # Generate 3 Gaussians
    julia> aG = ones(3);
    julia> mG = zeros(2,3);
    julia> mG[1,:] = [0.0, 0.7, -0.3];
    julia> mG[2,:] = [0.0, -0.6, -0.06];
    julia> AG = zeros(2,2,3);
    julia> AG[:,:,1] = [100 0.05; 0 100];
    julia> AG[:,:,2] = [100 0.15; 0 100];
    julia> AG[:,:,3] = [100 -0.3; 0.4 100];
    julia>
    julia> g = FPRS.genGauss(aG, mG, AG);
    julia> g[32:35,33:36]
    3×3 Array{Float64,2}:
     0.912746  0.832567  0.632048
     1.00009   0.912317  0.692624
     0.912275  0.832193  0.631775
    julia>
    julia> # Double the resolution
    julia> # Make grid bounds larger if necessary
    julia> xmin = [-1.0 -2.0];
    julia> xmax = [2.0   1.0];
    julia> g = FPRS.genGauss(aG, mG, AG, nPts = [127,127], xmin=xmin, xmax=xmax);
    julia> g[71:73, 59:61]
    3×3 Array{Float64,2}:
     0.862873  0.892834  0.824811
     0.955636  0.988733  0.913326
     0.944926  0.977569  0.902937
    julia>
    julia> # At which x1 coordinate was this Gaussian? Calculate with
    julia> x1 = range(xmin[1], stop=xmax[1], length=127);
    julia> collect(x1[73:75])
    3-element Array{Float64,1}:
     0.7142857142857143
     0.7380952380952381
     0.7619047619047619
    ```
    """
    # Convenience variables
    nDims  = size(mG,1)
    nGauss = size(mG,2)

    # Check input size
    size(aG) == (nGauss,) || throw(DimensionMismatch("size of aG, $(size(aG)), does not equal (size(mG,2),), $((size(mG,2),))"))
    size(AG) == (nDims, nDims, nGauss) || throw(DimensionMismatch("size of AG, $(size(AG)), does not equal (size(mG,1),size(mG,1),size(mG,2)), $((size(mG,1),size(mG,1),(size(mG,2))))"))

    # instantiate keyword args
    xmin == nothing ? xmin = - ones(Float64, nDims) : nothing
    xmax == nothing ? xmax = + ones(Float64, nDims) : nothing
    nPts == nothing ? nPts = 67 .* ones(Int64, nDims) : nothing

    # One-dimensional case
    if nDims == 1
        x1 = range(xmin[1], stop=xmax[1], length=nPts[1]);
        xGrid = (x1); # To return underlying grid in the end
        g = zeros(nPts[1]);

        for iG = 1:nGauss
            g = g + aG[iG] .*
                exp.( - (x1 .- mG[iG]) .^2 * AG[iG] .^2)
        end
    end

    # Two-dimensional case
    if nDims == 2
        x1 = range(xmin[1], stop=xmax[1], length=nPts[1]);
        x2 = range(xmin[2], stop=xmax[2], length=nPts[2]);
        xGrid = (x1, x2); # To return underlying grid in the end
        g = zeros(nPts[1], nPts[2]);

        for iG = 1:nGauss
            for ix2 = 1:nPts[2], ix1 = 1:nPts[1]
                x = [x1[ix1]; x2[ix2]]
                g[ix1,ix2] = g[ix1,ix2] +
                    aG[iG] * exp.(- transpose(x .- mG[:,iG]) *
                                  AG[:,:,iG] * (x .- mG[:,iG]))
            end
        end
    end

    if nDims == 3
        x1 = range(xmin[1], stop=xmax[1], length=nPts[1]);
        x2 = range(xmin[2], stop=xmax[2], length=nPts[2]);
        x3 = range(xmin[3], stop=xmax[3], length=nPts[3]);
        xGrid = (x1, x2, x3); # To return underlying grid in the end
        g = zeros(nPts[1], nPts[2], nPts[3]);

        for iG = 1:nGauss
            for ix3 = 1:nPts[3], ix2 = 1:nPts[2], ix1 = 1:nPts[1]
                x = [x1[ix1]; x2[ix2]; x3[ix3]]
                g[ix1,ix2,ix3] = g[ix1,ix2,ix3] +
                    aG[iG] * exp.(- transpose(x .- mG[:,iG]) *
                                  AG[:,:,iG] * (x .- mG[:,iG]))
            end
        end
    end

    return g, xGrid
end

function genGauss(aG, mG, AG;
                  xmin=nothing, xmax=nothing, nPts=nothing)
    """
    Return only the density of genGaussWGrid (do not return the grid)
    """
    res, ~ = genGaussWGrid(aG, mG, AG; xmin=xmin, xmax=xmax, nPts=nPts)
    return res
end

function genRandGaussWGrid(; nGauss=30,
                          nPts = [67; 67],
                          aGmin = 0.2,
                          aGmax = 2.0,
                          mGrad = 0.6,
                          AGmin = 100,
                          AGmax = 800,
                          randSeed = nothing,
                          suppShape = "ball")
    """
    Evaluate sum of Gaussians with random amplitudes, means, and covariance matrices.

    Generates random amplitudes, means, and covariance matrices, and calls
    genGauss function. Dimension is determined from the size of nPts arguments.

    # Keyword Arguments
    - `nGauss::Int64`: number of Gaussians
    - `nPts::Array{Int64}`: number of discr. pts along each dimension
    - `aGmin::Float64`
    - `aGmax::Float64` :
          amplitudes are uniformely distributed between these values.
    - `mGrad::Float64` :
          means (each coordinate) are  uniformely distributed between
          -mGrad and mGrad, if suppShape = "box"
    - `mGrad::Float64` :
          means are  uniformely distributed in the nDim-dimensional
          ball with radius mGrad, if suppShape = "ball"
    - `AGmin::Float64`
    - `AGmax::Float64` :
          eigenvalues of covariance matrices are uniformely distributed
          between these values. AGmin small --> slowly decaying gaussians
          (try AGmin = 10, AGmax = 15). AGmin large --> fast decaying gaussians
          (try AGmin = 500, AGmax = 1000).
    - `randSeed::Int64` : set random seed.
    - `suppShape::String` : "ball" or "box", see mGrad variable above.

    # Examples
    ```jldoctest
    julia> # Default 2D density with 30 Gaussians
    julia> g = genRandGauss();
    julia> # With specific random seed:
    julia> g = genRandGauss(randSeed = 0);
    julia> # The same with better resolution:
    julia> g = genRandGauss(randSeed = 0, nPts=[512,512]);
    julia> # One-dimensional, 10 Gaussians, smaller support
    julia> FPRS.genRandGauss(nGauss = 10, nPts=[512], AGmin=30, AGmax=45, mGrad=0.5, randSeed=0)
    julia> # 50 Gaussians, three-dimensional, with means in a box and faster decay
    julia> g = FPRS.genRandGauss(nGauss = 50, nPts=[80,80,80], AGmin=100, AGmax=1000, mGrad=0.5, randSeed=0, suppShape="box")
    julia> # E.g., plot after integrating with >>> heatmap(sum(g,dims=3)[:,:,1])
    ```
    """

    if suppShape != "box" && suppShape != "ball"
         throw(ArgumentError("suppShape, $(suppShape), must be either 'box' or 'ball'"))
    end

    # Convenience variables
    nDims = length(nPts)

    if randSeed == nothing
        seedVal = Random.seed!()
    else
        seedVal = Int(randSeed)
        Random.seed!(seedVal)
    end

    aG = (aGmax-aGmin) * rand(nGauss) .+ aGmin

    if nDims == 1
        mG = zeros(1,nGauss);
        mG[1,:] = 2.0 * mGrad * rand(nGauss) .- mGrad
        AG = zeros(1,1,nGauss);
        AG[1,1,:] = (AGmax-AGmin) * rand(nGauss) .+ AGmin
    end

    if nDims == 2
        # Generate means of Gaussians
        if suppShape == "box"
            mG = zeros(2,nGauss)
            mG[1,:] = 2.0 * mGrad * rand(nGauss) .- mGrad
            mG[2,:] = 2.0 * mGrad * rand(nGauss) .- mGrad
        elseif suppShape == "ball"
            mG = mGrad * randBall(2,nGauss)
        end

        # Generate variance matrices of Gaussians
        AG = zeros(2,2,nGauss);
        for iG = 1:nGauss
            # Prepare random rotation matrix (SO(2))
            theta = 2. * pi * rand();
            Omx = [cos(theta) -sin(theta); sin(theta) cos(theta)]
            Dmx = Diagonal((AGmax-AGmin) * rand(2) .+ AGmin)
            Amx = transpose(Omx) * Dmx * Omx
            AG[:,:,iG] = Amx[:,:]
        end
    end

    if nDims == 3
        # Generate means of Gaussians
        if suppShape == "box"
            mG = zeros(3,nGauss)
            mG[1,:] = 2.0 * mGrad * rand(nGauss) .- mGrad
            mG[2,:] = 2.0 * mGrad * rand(nGauss) .- mGrad
            mG[3,:] = 2.0 * mGrad * rand(nGauss) .- mGrad
        elseif suppShape == "ball"
            mG = mGrad * randBall(3,nGauss)
        end

        # Generate variance matrices of Gaussians
        AG = zeros(3,3,nGauss);
        for iG = 1:nGauss
            # Prepare random rotation matrix (SO(2))
            Omx = randRotationMatrix()
            Dmx = Diagonal((AGmax-AGmin) * rand(3) .+ AGmin)
            Amx = transpose(Omx) * Dmx * Omx
            AG[:,:,iG] = Amx[:,:]
        end
    end

    return genGaussWGrid(aG, mG, AG; nPts=nPts)
end

function genRandGauss(; nGauss=30,
                      nPts = [67; 67],
                      aGmin = 0.2,
                      aGmax = 2.0,
                      mGrad = 0.6,
                      AGmin = 100,
                      AGmax = 800,
                      randSeed = nothing,
                      suppShape = "ball")
    res, ~ = genRandGaussWGrid(; nGauss=nGauss,
                              nPts = nPts,
                              aGmin = aGmin,
                              aGmax = aGmax,
                              mGrad = mGrad,
                              AGmin = AGmin,
                              AGmax = AGmax,
                              randSeed = randSeed,
                              suppShape = suppShape)
    return res
end


const N_GRADS_2D_EXPONENT = 7
const N_GRADS_2D = 1 << N_GRADS_2D_EXPONENT

const NORMALIZER_2D = 0.05481866495625118

struct Grad2
    x::Float32
    y::Float32
end

const grad2 = Grad2[
    Grad2( 0.38268343236509,   0.923879532511287),
    Grad2( 0.923879532511287,  0.38268343236509),
    Grad2( 0.923879532511287, -0.38268343236509),
    Grad2( 0.38268343236509,  -0.923879532511287),
    Grad2(-0.38268343236509,  -0.923879532511287),
    Grad2(-0.923879532511287, -0.38268343236509),
    Grad2(-0.923879532511287,  0.38268343236509),
    Grad2(-0.38268343236509,   0.923879532511287),
    #-------------------------------------------#
    Grad2( 0.130526192220052,  0.99144486137381),
    Grad2( 0.608761429008721,  0.793353340291235),
    Grad2( 0.793353340291235,  0.608761429008721),
    Grad2( 0.99144486137381,   0.130526192220051),
    Grad2( 0.99144486137381,  -0.130526192220051),
    Grad2( 0.793353340291235, -0.60876142900872),
    Grad2( 0.608761429008721, -0.793353340291235),
    Grad2( 0.130526192220052, -0.99144486137381),
    Grad2(-0.130526192220052, -0.99144486137381),
    Grad2(-0.608761429008721, -0.793353340291235),
    Grad2(-0.793353340291235, -0.608761429008721),
    Grad2(-0.99144486137381,  -0.130526192220052),
    Grad2(-0.99144486137381,   0.130526192220051),
    Grad2(-0.793353340291235,  0.608761429008721),
    Grad2(-0.608761429008721,  0.793353340291235),
    Grad2(-0.130526192220052,  0.99144486137381),
]

const GRADIENTS_2D = Vector{Grad2}(undef, N_GRADS_2D)

function grad(a, seed, xsvp, ysvp, dx, dy)
    hash = xor(seed, xsvp, ysvp) * HASH_MULTIPLIER
    g = GRADIENTS_2D[(xor(hash, hash >> (64 - N_GRADS_2D_EXPONENT)) & (N_GRADS_2D - 1)) + 1]
    Float32(p4(a) * (g.x * dx + g.y * dy) / NORMALIZER_2D)
end

function fill_2d()
    j = 0
    len = length(grad2)
    for i = 1:N_GRADS_2D
        GRADIENTS_2D[i] = grad2[j += 1]
        j == len && (j = 0)
    end
end

fill_2d()


#=

  Zoo, buses and kids problem in Julia ConstraintSolver.jl

  https://dmcommunity.org/challenge/challenge-july-2018/
  """
  300 kids need to travel to the London zoo. The school may rent 40 seats and
  30 seats buses for 500 and 400 £. How many buses of each to minimize cost?
  """

  Also, see Alex Fleisher
  "Optimization : simply do more with less, zoo, buses and kids (Part2, python, java, C++…)"
  ""
  https://alexfleischer-84755.medium.com/optimization-simply-do-more-with-less-zoo-buses-and-kids-part2-python-java-c-cc04558e49b5

  This model was created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/

=#

using ConstraintSolver, JuMP
using Cbc, GLPK, Ipopt
const CS = ConstraintSolver
include("constraints_utils.jl")

function zoo(problem,print_solutions=true,all_solutions=true)

    cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
    glpk_optimizer = optimizer_with_attributes(GLPK.Optimizer)
    ipopt_optimizer = optimizer_with_attributes(Ipopt.Optimizer)

    model = Model(optimizer_with_attributes(CS.Optimizer,   # "all_solutions"=> all_solutions,
                                                            "all_optimal_solutions"=>all_solutions,
                                                            "logging"=>[],

                                                            "traverse_strategy"=>:BFS,
                                                            # "traverse_strategy"=>:DFS,
                                                            # "traverse_strategy"=>:DBFS,

                                                            # "branch_split"=>:Smallest,
                                                            # "branch_split"=>:Biggest,
                                                            "branch_split"=>:InHalf,

                                                            # https://wikunia.github.io/ConstraintSolver.jl/stable/options/#branch_strategy-(:Auto)
                                                            "branch_strategy" => :IMPS, # default
                                                            # "branch_strategy" => :ABS, # Activity Based Search
                                                            # "activity.decay" => 0.999, # default 0.999
                                                            # "activity.max_probes" => 10, # default, 10
                                                            # "activity.max_confidence_deviation" => 20, # default 20

                                                            # "simplify"=>false,
                                                            # "simplify"=>true, # default

                                                            "time_limit"=>6,

                                                            # "backtrack" => false, # default true
                                                            # "backtrack_sorting" => false, # default true

                                                            # "lp_optimizer" => cbc_optimizer,
                                                            # "lp_optimizer" => glpk_optimizer,
                                                            # "lp_optimizer" => ipopt_optimizer,
                                        ))
    num_kids  = problem[:num_kids]
    num_seats = problem[:num_seats]
    price     = problem[:price]
    n         = problem[:n]

    @variable(model, 1 <= x[1:n] <= 100, Int)
    @variable(model, 0 <= z <= sum(price.*num_seats), Int)

    scalar_product(model, x,num_seats,:(>=), num_kids)
    @constraint(model, z == sum(x.*price))

    @objective(model,Min,z)

    # Solve the problem
    optimize!(model)

    status = JuMP.termination_status(model)
    # println("status:$status")
    if status == MOI.OPTIMAL
        num_sols = MOI.get(model, MOI.ResultCount())
        println("num_sols:$num_sols\n")
        if print_solutions
            for sol in 1:num_sols
                println("solution #$sol")
                x_val = convert.(Integer,JuMP.value.(x; result=sol))
                z_val = convert.(Integer,JuMP.value.(z; result=sol))
                println("x:$x_val")
                println("z:$z_val")
                println()

            end
        end
    else
        println("status:$status")
    end

    return status
end

zoo_problem = Dict(
   :num_kids  => 300,
   :num_seats => [30,40],
   :price     => [400,500],
   :n         => 2,

)

@time zoo(zoo_problem)


const pubchem_features = [:acceptor, :anion, :cation, :donor, :hydrophobe, :rings, :aromaticrings]

"""
    feat = atoms_to_feature(atoms, dirs=nothing)

Combines Gaussian distributions specified by `atoms` to create a feature repressented by a
single Gaussian. The feature is centered at the center of mass of the distributions, and the
width `σ` of the feature is the average of the distance of individual features from the center
added to their widths.

Geometric constraints for the feature can be specified by `dirs`. The returned feature has no
direction by default.
"""
function atoms_to_feature(atoms::AbstractVector{<:AbstractIsotropicGaussian}, dirs=nothing)
    if length(atoms)==1
        μ = atoms[1].μ
        σ = atoms[1].σ
        ϕ = atoms[1].ϕ
    else
        μ = center_of_mass(atoms)
        # σ = mean([norm(μ-a.μ) + a.σ for (i,a) in enumerate(atoms)])
        σ = sum([norm(μ.-a.μ) for (i,a) in enumerate(atoms)])/length(atoms)
        ϕ = sum([a.ϕ for a in atoms])/length(atoms)
    end
    if isnothing(dirs)
        dirs = typeof(μ)[]
    end

    return IsotropicGaussian(μ, σ, ϕ, dirs)
end
atoms_to_feature(gmm::AbstractSingleGMM, dirs=nothing) = atoms_to_feature(gmm.gaussians, dirs)
atoms_to_feature(mol::UndirectedGraph, nodeset, dirs=nothing; kwargs...) = atoms_to_feature(MolGMM(nodesubgraph(mol,nodeset); kwargs...), dirs)
atoms_to_feature(submol::SubgraphView, dirs=nothing; kwargs...) = atoms_to_feature(submol.graph, nodeset(submol), dirs; kwargs...)

#TODO: fit single Gaussian to points sampled from sum of the features?
function combine_features(feats::AbstractVector{<:AbstractIsotropicGaussian}, weights=[f.ϕ for f in feats])
    ϕ = sum(weights)
    μ = center_of_mass(feats, weights)
    σ = sum([f.σ * weights[i]/ϕ for (i,f) in enumerate(feats)])
    return IsotropicGaussian(μ, σ, ϕ)
end


function part1(input)
    joltages = sort(parse.(Int, readlines(input)))
    num_one_diffs = count(==(1), diff(vcat(0, joltages)))
    num_three_diffs = count(==(3), diff(vcat(0, joltages))) + 1
    return num_one_diffs * num_three_diffs
end

function part2(input)
    joltages = sort(parse.(Int, readlines(input)))
    num_combinations = Dict{Int, Int}()
    num_combinations[0] = 1
    for joltage in joltages
        num_combinations[joltage] = sum(get(num_combinations, joltage - n, 0)
                                        for n in 1:3)
    end
    return num_combinations[last(joltages)]
end


# test deterministic vs. stochastic
using Distributed;
addprocs(20);
@everywhere include("loadMod.jl");
@everywhere const GUROBI_ENV = Gurobi.Env();

caseList = [13,33,123];
NN = 1000;
Δt = 0.25;
N = 5;
iterMax = 20;
pathListDRaw = load("pathHist_1000.jld");
TList = [24,36,48,72,96];

for ci in 1:length(caseList)
    pathDictA = pathListDRaw["pathDict"][ci];
    detOut = Dict();
    stochOut = Dict();
    detCostDict = Dict();
    for T in TList
        τ = Int64(1/6*T);
        for j in procs()
            remotecall_fetch(readInData,j,ci,caseList,T,τ);
            #remotecall_fetch(readInData_old,j,T,ωSet0,10000,0);
        end

        # select a preset pathDict
        pathDict = pathDictA[T];
        solDet,costDet = exeDet(T, Δt, fData, bData, dData, pDistr, NN, pathDict);
        listDet = [costDet[i] for i in 1:NN];
        meanDet = mean(listDet);
        sigmaDet = std(listDet);
        println(round(meanDet,digits = 2)," ",
                round(meanDet - 1.96*sigmaDet,digits = 2)," ",
                round(meanDet + 1.96*sigmaDet,digits = 2));
        detOut[T] = [costDet,listDet,meanDet,sigmaDet];

        # train the stochastic programming strategy
        cutDictPG = preGen(T, Δt, N, iterMax);
        # cutDict,LBHist,UBHist,UBuHist,UBlHist,timeHist = solveMain(τ, T, Δt, N, true, false, 10, 10, cutDictPG);
        cutDict,LBHist,UBHist,UBuHist,UBlHist,timeHist = solveMain(T, Δt, N, false, false,
            max(Int64(round(500/N)),20), max(Int64(round(500/N)),20),cutDictPG);
        for j in procs()
            remotecall_fetch(cutIni,j,cutDict);
        end

        solSDDP, LBSDDP, costSDDP = exeForward(T, Δt, NN, false, pathDict);
        listSDDP = [costSDDP[i] for i in 1:NN];
        meanSDDP = mean(listSDDP);
        sigmaSDDP = std(listSDDP);
        println(round(meanSDDP,digits = 2)," ",
                round(meanSDDP - 1.96*sigmaSDDP,digits = 2)," ",
                round(meanSDDP + 1.96*sigmaSDDP,digits = 2));
        stochOut[T] = [LBSDDP, costSDDP, listSDDP, meanSDDP, sigmaSDDP, LBHist];

        detSol,detObj = detBuild(Δt, T, fData, bData, dData);
        detCostDict[T] = detObj;

        #println("Output: ",round(meanDet,1)," & ",round(meanSDDP,1)," & ", round(meanDet - meanSDDP,1)," & ", round((meanDet - meanSDDP)/meanDet*100,1));
        #save("detResults_old.jld", "detOut", detOut, "stochOut", stochOut, "nomOut", detCostDict);
        save("detResults_$(ci).jld", "detOut", detOut, "stochOut", stochOut, "nomOut", detCostDict);
    end
end

# for T in TList
#     println(T," & ",round(detOut[T][3],2)," & ",round(stochOut[T][4],2)," & ",round((detOut[T][3] - stochOut[T][4])/detOut[T][3]*100,2)," & ");
# end

# pathListData = pmap(i -> simuPath(τ,T,pDistr), 1:NN);
# pathDict = Dict();
# for i in 1:length(pathListData)
#     pathDict[i] = pathListData[i];
# end
# solDet,costDet = exeDet(τ, T, Δt, fData, bData, dData, pDistr, NN, pathDict);
# listDet = [costDet[i] for i in 1:NN];
# meanDet = mean(listDet);
# sigmaDet = std(listDet);
# println(round(meanDet,2)," ",round(meanDet - 1.96*sigmaDet,2)," ",round(meanDet + 1.96*sigmaDet,2));
#
# cutDict,LBHist,UBHist,UBuHist,UBlHist = solveMain(τ, T, Δt, fData, pDistr, bData, dData, N, 20, 20);
# solSDDP, LBSDDP, costSDDP = exeForward(τ, T, Δt, fData, bData, dData, pDistr, NN, cutDict, pathDict);
# listSDDP = [costSDDP[i] for i in 1:NN];
# meanSDDP = mean(listSDDP);
# sigmaSDDP = std(listSDDP);
# println(round(meanSDDP,2)," ",round(meanSDDP - 1.96*sigmaSDDP,2)," ",round(meanSDDP + 1.96*sigmaSDDP,2));
#
# outputData = [solDet,costDet,meanDet,sigmaDet,
#     cutDict,LBHist,UBHist,UBuHist,UBlHist,solSDDP,LBSDDP,costSDDP,meanSDDP,sigmaSDDP];
# save("detOut.jld","data",outputData);


#=**************************************************************************

This script imports and runs an ozone NOx formaldehyde mechanism, producing 3
text files to be used in training physics constrained neural networks.

Written in Julia Version 1.2.0, tested up to 1.4.2

Written by Obin Sturm (posturm@ucdavis.edu)

The inverse problem tries to solve for S according to

  A*S = ∆C

 where S is the integrated reaction rates over a given timestep, and ∆C is the
 change in concentration.

The output files are "C.txt", "S.txt" and "delC.txt" for use in machine
learning research and inverse problems.  This was first used in a paper
introducing a framework for conserving mass in machine learning applications
for air quality and climate modeling: https://doi.org/10.5194/gmd-13-4435-2020

 C and delC correspond to the following species:
         1) O3
         2) NO
         3) NO2
         4) HCHO
         5) HO2
         6) HO2H
         7) HO.
         8) O
         9) HNO3
        10) CO
        11) H2

 S corresponds to the following reactions:
         1) NO2 + HV = NO + O
         2) O + O2 = O3
         3) O3 + NO = NO2 + O2
         4) HCHO + HV = 2 HO2. + CO
         5) HCHO + HV = H2 + CO
         6) HCHO + HO. = HO2. + CO + H2O
         7) HO2. + NO = HO. + NO2
         8) HO. + NO2 = HNO3
         9) HO2H + HV = 2 HO.
        10) HO2H + HO. = H2O + HO2
The stoichiometry matrix A, which can also be interpreted as the directed,
weighted incidence matrix of the species-reaction network, is

A =           [0  1 -1  0  0  0  0  0  0  0;
               1  0 -1  0  0  0 -1  0  0  0;
              -1  0  1  0  0  0  1 -1  0  0;
               0  0  0 -1 -1 -1  0  0  0  0;
               0  0  0  2  0  1 -1  0  0  1;
               0  0  0  0  0  0  0  0 -1 -1;
               0  0  0  0  0 -1  1 -1  2 -1;
               1 -1  0  0  0  0  0  0  0  0;
               0  0  0  0  0  0  0  1  0  0;
               0  0  0  1  1  1  0  0  0  0;
               0  0  0  0  1  0  0  0  0  0]
****************************************************************************=#

mechanism = include("assign3_driver.jl") # Include photochemical mechanism module
using DelimitedFiles

mechanism = include("assign3_driver.jl")
C,S,J = mechanism.gas(0.0,24.0*60.0,6.0,5000) # Run 5000 full days, sampling every 6 minutes
# delC = diff([C ; zeros(1,11)],dims = 1)

  open("S.txt", "w") do io
                         writedlm(io, S,',')
                   end
  open("C.txt", "w") do io
                          writedlm(io, C,',')
                   end
  open("J.txt", "w") do io
                         writedlm(io, J,',')
                   end