# Regular functions for evaluating derivatives and gradients of Chebyshev polynomials

# Regular functions for derivatives

function chebyshev_derivative(weights::AbstractArray{T,N},x::AbstractArray{R,1},pos::S,order::Array{S,1},domain=[ones(T,1,N);-ones(T,1,N)]) where {T<:AbstractFloat,R<:Number,N,S<:Integer}

  poly = Array{Array{R,2},1}(undef,N)
  @inbounds for i = 1:N
    if i === pos
      poly[i] = derivative_of_chebyshev_polynomial(order[i],normalize_node(x[i],domain[:,i]))
    else
      poly[i] = chebyshev_polynomial(order[i],normalize_node(x[i],domain[:,i]))
    end
  end

  derivative = zero(T)
  @inbounds for i in CartesianIndices(weights)
    poly_product = poly[1][i[1]]
    @inbounds for j = 2:N
      poly_product *= poly[j][i[j]]
    end
    derivative += weights[i]*poly_product
  end

  return derivative*(2.0/(domain[1,pos]-domain[2,pos]))

end

function chebyshev_derivative(weights::AbstractArray{T,N},x::AbstractArray{R,1},pos::S,order::S,domain=[ones(T,1,N);-ones(T,1,N)]) where {T<:AbstractFloat,R<:Number,N,S<:Integer}

  poly = Array{Array{R,2},1}(undef,N)
  @inbounds for i = 1:N
    if i === pos
      poly[i] = derivative_of_chebyshev_polynomial(order,normalize_node(x[i],domain[:,i]))
    else
      poly[i] = chebyshev_polynomial(order,normalize_node(x[i],domain[:,i]))
    end
  end

  derivative = zero(T)
  @inbounds for i in CartesianIndices(weights)
    if sum(Tuple(i)) <= order+N
      poly_product = poly[1][i[1]]
      @inbounds for j = 2:N
        poly_product *= poly[j][i[j]]
      end
      derivative += weights[i]*poly_product
    end
  end

  return derivative*(2.0/(domain[1,pos]-domain[2,pos]))

end

function chebyshev_derivative(cheb_poly::ChebPoly,x::AbstractArray{R,1},pos::S) where {T<:AbstractFloat,R<:Number,S<:Integer}

  derivative = chebyshev_derivative(cheb_poly.weights,x,pos,cheb_poly.order,cheb_poly.domain)

  return derivative

end

function chebyshev_derivative(cheb::ChebInterpRoots,x::AbstractArray{R,1},pos::S) where {T<:AbstractFloat,R<:Number,S<:Integer}

  weights    = chebyshev_weights_threaded(cheb.data,cheb.nodes,cheb.order,cheb.domain)
  derivative = chebyshev_derivative(weights,x,pos,cheb.order,cheb.domain)

  return derivative

end

function chebyshev_derivative(cheb::ChebInterpExtrema,x::AbstractArray{R,1},pos::S) where {T<:AbstractFloat,R<:Number,S<:Integer}

  weights    = chebyshev_weights_extrema_threaded(cheb.data,cheb.nodes,cheb.order,cheb.domain)
  derivative = chebyshev_derivative(weights,x,pos,cheb.order,cheb.domain)

  return derivative

end

function chebyshev_derivative(cheb::ChebInterpExtended,x::AbstractArray{R,1},pos::S) where {T<:AbstractFloat,R<:Number,S<:Integer}

  weights    = chebyshev_weights_extended_threaded(cheb.data,cheb.nodes,cheb.order,cheb.domain)
  derivative = chebyshev_derivative(weights,x,pos,cheb.order,cheb.domain)

  return derivative

end

# Regular functions for gradients

function chebyshev_gradient(weights::AbstractArray{T,N},x::AbstractArray{R,1},order::Array{S,1},domain=[ones(T,1,N);-ones(T,1,N)]) where {T<:AbstractFloat,R<:Number,N,S<:Integer}

  gradient = Array{R,2}(undef,1,N)

  @inbounds for i = 1:N
    gradient[i] = chebyshev_derivative(weights,x,i,order,domain)
  end

  return gradient

end

function chebyshev_gradient(weights::AbstractArray{T,N},x::AbstractArray{R,1},order::S,domain=[ones(T,1,N);-ones(T,1,N)]) where {T<:AbstractFloat,R<:Number,N,S<:Integer}

  gradient = Array{R,2}(undef,1,N)

  @inbounds for i = 1:N
    gradient[i] = chebyshev_derivative(weights,x,i,order,domain)
  end

  return gradient

end

function chebyshev_gradient(cheb_poly::ChebPoly,x::AbstractArray{R,1}) where {R<:Number}

  gradient = chebyshev_gradient(cheb_poly.weights,x,cheb_poly.order,cheb_poly.domain)

  return gradient

end

function chebyshev_gradient(cheb::ChebInterpRoots,x::AbstractArray{R,1}) where {R<:Number}

  weights  = chebyshev_weights_threaded(cheb.data,cheb.nodes,cheb.order,cheb.domain)
  gradient = chebyshev_gradient(weights,x,cheb.order,cheb.domain)

  return gradient

end

function chebyshev_gradient(cheb::ChebInterpExtrema,x::AbstractArray{R,1}) where {R<:Number}

  weights  = chebyshev_weights_extrema_threaded(cheb.data,cheb.nodes,cheb.order,cheb.domain)
  gradient = chebyshev_gradient(weights,x,cheb.order,cheb.domain)

  return gradient

end

function chebyshev_gradient(cheb::ChebInterpExtended,x::AbstractArray{R,1}) where {R<:Number}

  weights  = chebyshev_weights_extended_threaded(cheb.data,cheb.nodes,cheb.order,cheb.domain)
  gradient = chebyshev_gradient(weights,x,cheb.order,cheb.domain)

  return gradient

end

function chebyshev_gradient(cheb::ChebInterpVertesi,x::AbstractArray{R,1}) where {R<:Number}

  weights  = chebyshev_weights_vertesi_threaded(cheb.data,cheb.nodes,cheb.order,cheb.domain)
  gradient = chebyshev_gradient(weights,x,cheb.order,cheb.domain)

  return gradient

end

function chebyshev_derivative(weights::AbstractArray{T,N},pos::S,order::Array{S,1},domain=[ones(T,1,N);-ones(T,1,N)]) where {T <: AbstractFloat,N,S <: Integer}

  function chebderiv(x::AbstractArray{R,1}) where {R<:Number}

    return chebyshev_derivative(weights,x,pos,order,domain)

  end

  return chebderiv

end

function chebyshev_derivative(weights::AbstractArray{T,N},pos::S,order::S,domain=[ones(T,1,N);-ones(T,1,N)]) where {T <: AbstractFloat,N,S <: Integer}

  function chebderiv(x::AbstractArray{R,1}) where {R<:Number}

    return chebyshev_derivative(weights,x,pos,order,domain)

  end

  return chebderiv

end

function chebyshev_derivative(cheb_poly::ChebPoly,pos::S) where {S <: Integer}

  function chebderiv(x::AbstractArray{R,1}) where {R<:Number}

    return chebyshev_derivative(cheb_poly,x,pos)

  end

  return chebderiv

end

function chebyshev_derivative(cheb::ChebInterpRoots,pos::S) where {S <: Integer}

  weights = chebyshev_weights_threaded(cheb.data,cheb.nodes,cheb.order,cheb.domain)

  function chebderiv(x::AbstractArray{R,1}) where {R<:Number}

    return chebyshev_derivative(weights,x,pos,cheb.order,cheb.domain)

  end

  return chebderiv

end

function chebyshev_derivative(cheb::ChebInterpExtrema,pos::S) where {S <: Integer}

  weights = chebyshev_weights_extrema_threaded(cheb.data,cheb.nodes,cheb.order,cheb.domain)

  function chebderiv(x::AbstractArray{R,1}) where {R<:Number}

    return chebyshev_derivative(weights,x,pos,cheb.order,cheb.domain)

  end

  return chebderiv

end

function chebyshev_derivative(cheb::ChebInterpExtended,pos::S) where {S <: Integer}

  weights = chebyshev_weights_extended_threaded(cheb.data,cheb.nodes,cheb.order,cheb.domain)

  function chebderiv(x::AbstractArray{R,1}) where {R<:Number}

    return chebyshev_derivative(weights,x,pos,cheb.order,cheb.domain)

  end

  return chebderiv

end

function chebyshev_derivative(cheb::ChebInterpVertesi,pos::S) where {S <: Integer}

  weights = chebyshev_weights_vertesi_threaded(cheb.data,cheb.nodes,cheb.order,cheb.domain)

  function chebderiv(x::AbstractArray{R,1}) where {R<:Number}

    return chebyshev_derivative(weights,x,pos,cheb.order,cheb.domain)

  end

  return chebderiv

end

function chebyshev_gradient(weights::AbstractArray{T,N},order::Array{S,1},domain=[ones(T,1,N);-ones(T,1,N)]) where {T <: AbstractFloat,N,S <: Integer}

  function chebgrad(x::AbstractArray{R,1}) where {R<:Number}

    return chebyshev_gradient(weights,x,pos,order,domain)

  end

  return chebgrad

end

function chebyshev_gradient(weights::AbstractArray{T,N},order::S,domain=[ones(T,1,N);-ones(T,1,N)]) where {T <: AbstractFloat,N,S <: Integer}

  function chebgrad(x::AbstractArray{R,1}) where {R<:Number}

    return chebyshev_gradient(weights,x,pos,order,domain)

  end

  return chebgrad

end

function chebyshev_gradient(cheb_poly::ChebPoly) where {S <: Integer}

  function chebgrad(x::AbstractArray{R,1}) where {R<:Number}

    return chebyshev_gradient(cheb_poly,x)

  end

  return chebgrad

end

function chebyshev_gradient(cheb::ChebInterpRoots) where {S <: Integer}

  weights = chebyshev_weights_threaded(cheb.data,cheb.nodes,cheb.order,cheb.domain)

  function chebgrad(x::AbstractArray{R,1}) where {R<:Number}

    return chebyshev_gradient(weights,x,cheb.order,cheb.domain)

  end

  return chebgrad

end

function chebyshev_gradient(cheb::ChebInterpExtrema) where {S <: Integer}

  weights = chebyshev_weights_extrema_threaded(cheb.data,cheb.nodes,cheb.order,cheb.domain)

  function chebgrad(x::AbstractArray{R,1}) where {R<:Number}

    return chebyshev_gradient(weights,x,cheb.order,cheb.domain)

  end

  return chebgrad

end

function chebyshev_gradient(cheb::ChebInterpExtended) where {S <: Integer}

  weights = chebyshev_weights_extended_threaded(cheb.data,cheb.nodes,cheb.order,cheb.domain)

  function chebgrad(x::AbstractArray{R,1}) where {R<:Number}

    return chebyshev_gradient(weights,x,cheb.order,cheb.domain)

  end

  return chebgrad

end

function chebyshev_gradient(cheb::ChebInterpVertesi) where {S <: Integer}

  weights = chebyshev_weights_vertesi_threaded(cheb.data,cheb.nodes,cheb.order,cheb.domain)

  function chebgrad(x::AbstractArray{R,1}) where {R<:Number}

    return chebyshev_gradient(weights,x,cheb.order,cheb.domain)

  end

  return chebgrad

end


include("gibbssampling.jl")
include("hmcsampling.jl")

function log_gp_prior(gp::_MCGP,f::AbstractVector)
    logpdf(MvNormal(gp.μ₀,gp.K),f)
end

function log_joint_model(model::MCGP{T,L,<:SamplingInference},x) where {T,L}
    fs = unsqueeze(model,x)
    log_likelihood(model.likelihood,get_y(model),fs) + sum(log_gp_prior.(model.f),fs)
end

function unsqueeze(model::MCGP,f)
    n = model.nFeatures
    Tuple(f[((i-1)*n+1):(i*n)] for i in 1:model.nLatent)
end

function grad_log_joint_model(model::MCGP{T,L,<:SamplingInference},x) where {T,L}
    fs = unsqueeze(model,x)
    vcat(grad_log_likelihood(model.likelihood,get_y(model),fs)...) + vcat(grad_log_gp_prior.(model.f,fs)...)
end

log_likelihood(l::Likelihood,y::AbstractVector,f::Tuple{<:AbstractVector{T}}) where {T<:Real} = logpdf.(l,y,first(f))

function grad_log_likelihood(l::Likelihood,y::AbstractVector,f::Tuple{<:AbstractVector{T}}) where {T}
    grad_log_pdf.(l,y,f)
end

function grad_log_gp_prior(gp,f)
    gp.K/(f-gp.μ₀)
end


module Poisson_examples
using FinEtools

function Poisson_FE_example()
    println("""

    Heat conduction example described by Amuthan A. Ramabathiran
    http://www.codeproject.com/Articles/579983/Finite-Element-programming-in-Julia:
    Unit square, with known temperature distribution along the boundary,
    and uniform heat generation rate inside.  Mesh of regular linear TRIANGLES,
    in a grid of 1000 x 1000 edges (2M triangles, 1M degrees of freedom).
    Version: 05/29/2017
    """
    )
    t0 = time()

    A = 1.0 # dimension of the domain (length of the side of the square)
    thermal_conductivity =  [i==j ? one(FFlt) : zero(FFlt) for i=1:2, j=1:2]; # conductivity matrix
    Q = -6.0; # internal heat generation rate
    tempf(x) = (1.0 .+ x[:,1].^2 .+ 2.0 .* x[:,2].^2);#the exact distribution of temperature
    N = 1000;# number of subdivisions along the sides of the square domain

    println("Mesh generation")
    @time fens,fes =T3block(A, A, N, N)

    geom = NodalField(fens.xyz)
    Temp = NodalField(zeros(size(fens.xyz,1),1))

    println("Searching nodes  for BC")
    @time l1 = selectnode(fens; box=[0. 0. 0. A], inflate = 1.0/N/100.0)
    @time l2 = selectnode(fens; box=[A A 0. A], inflate = 1.0/N/100.0)
    @time l3 = selectnode(fens; box=[0. A 0. 0.], inflate = 1.0/N/100.0)
    @time l4 = selectnode(fens; box=[0. A A A], inflate = 1.0/N/100.0)
    List = vcat(l1, l2, l3, l4)
    @time setebc!(Temp, List, true, 1, tempf(geom.values[List,:])[:])
    @time applyebc!(Temp)
    @time numberdofs!(Temp)

    t1 = time()

    material = MatHeatDiff(thermal_conductivity)

    femm = FEMMHeatDiff(IntegData(fes, TriRule(1)), material)

    println("Conductivity")
    @time K = conductivity(femm, geom, Temp)
    println("Nonzero EBC")
    @time F2 = nzebcloadsconductivity(femm, geom, Temp);
    println("Internal heat generation")
    # function getsource!(forceout::FFltVec, XYZ::FFltMat, tangents::FFltMat, fe_label::FInt)
    #   forceout[1] = Q; #heat source
    # end
    # fi = ForceIntensity(FFlt, getsource!);# alternative  specification
    fi = ForceIntensity(FFlt[Q]);
    @time F1 = distribloads(FEMMBase(IntegData(fes, TriRule(1))), geom, Temp, fi, 3);

    println("Factorization")
    @time K = cholesky(K)
    println("Solution of the factorized system")
    @time U = K\(F1+F2)
    scattersysvec!(Temp,U[:])

    println("Total time elapsed = $(time() - t0) [s]")
    println("Solution time elapsed = $(time() - t1) [s]")

    Error= 0.0
    for k=1:size(fens.xyz,1)
        Error = Error + abs.(Temp.values[k,1] - tempf(reshape(fens.xyz[k,:], (1,2)))[1])
    end
    println("Error =$Error")

    # File =  "a.vtk"
    # MeshExportModule.vtkexportmesh (File, fes.conn, [geom.values Temp.values], MeshExportModule.T3; scalars=Temp.values, scalars_name ="Temperature")

    true
end

function Poisson_FE_example_algo()
    A= 1.0
    thermal_conductivity =  [i==j ? one(FFlt) : zero(FFlt) for i=1:2, j=1:2]; # conductivity matrix
    magn = -6.0; #heat source
    truetempf(x)=1.0 .+ x[1].^2 .+ 2.0 .* x[2].^2;
    N=20;

    println("""
    Heat conduction example described by Amuthan A. Ramabathiran
    http://www.codeproject.com/Articles/579983/Finite-Element-programming-in-Julia:
    Unit square, with known temperature distribution along the boundary,
    and uniform heat generation rate inside.  Mesh of regular TRIANGLES,
    in a grid of $N x $N edges.
    This version uses the FinEtools algorithm module.
    """
    )
    t0 = time()

    fens,fes = T3block(A, A, N, N)

    # Define boundary conditions
    l1  = selectnode(fens; box=[0. 0. 0. A], inflate = 1.0/N/100.0)
    l2  = selectnode(fens; box=[A A 0. A], inflate = 1.0/N/100.0)
    l3  = selectnode(fens; box=[0. A 0. 0.], inflate = 1.0/N/100.0)
    l4  = selectnode(fens; box=[0. A A A], inflate = 1.0/N/100.0)

    essential1 = FDataDict("node_list"=>vcat(l1, l2, l3, l4),
    "temperature"=>truetempf);
    material = MatHeatDiff(thermal_conductivity)
    femm = FEMMHeatDiff(IntegData(fes, TriRule(1)), material)
    region1 = FDataDict("femm"=>femm, "Q"=>magn)
    # Make model data
    modeldata= FDataDict("fens"=> fens,
    "regions"=>[region1],
    "essential_bcs"=>[essential1]);

    # Call the solver
    modeldata = FinEtools.AlgoHeatDiffModule.steadystate(modeldata)

    println("Total time elapsed = ",time() - t0,"s")

    geom=modeldata["geom"]
    Temp=modeldata["temp"]
    femm=modeldata["regions"][1]["femm"]

    function errfh(loc,val)
        exact = truetempf(loc)
        return ((exact[1]-val[1])*exact)
    end
    E = integratefieldfunction(femm, geom, Temp, errfh, 0.0, m=3)
    println("Error=$E")

end # Poisson_FE_example_algo

function Poisson_FE_example_csys_1()
    println("""
    Heat conduction example described by Amuthan A. Ramabathiran
    http://www.codeproject.com/Articles/579983/Finite-Element-programming-in-Julia:
    Unit square, with known temperature distribution along the boundary,
    and uniform heat generation rate inside.  Mesh of regular linear TRIANGLES,
    in a grid of 1000 x 1000 edges (2M triangles, 1M degrees of freedom).
    The material response is defined in a local coordinate system.
    Version: 05/29/2017
    """
    )
    t0 = time()

    A = 1.0 # dimension of the domain (length of the side of the square)
    thermal_conductivity =  [i==j ? one(FFlt) : zero(FFlt) for i=1:2, j=1:2]; # conductivity matrix
    Q = -6.0; # internal heat generation rate
    function getsource!(forceout::FFltVec, XYZ::FFltMat, tangents::FFltMat, fe_label::FInt)
        forceout[1] = Q; #heat source
    end
    tempf(x) = (1.0 .+ x[:,1].^2 .+ 2.0 .* x[:,2].^2);#the exact distribution of temperature
    N = 1000;# number of subdivisions along the sides of the square domain
    Rm=[-0.9917568452513019 -0.12813414805267656
    -0.12813414805267656 0.9917568452513019]
    Rm=[-0.8020689950104449 -0.5972313850116512
    -0.5972313850116512 0.8020689950104447]

    println("Mesh generation")
    @time fens,fes =T3block(A, A, N, N)

    geom = NodalField(fens.xyz)
    Temp = NodalField(zeros(size(fens.xyz,1),1))

    println("Searching nodes  for BC")
    @time l1 = selectnode(fens; box=[0. 0. 0. A], inflate = 1.0/N/100.0)
    @time l2 = selectnode(fens; box=[A A 0. A], inflate = 1.0/N/100.0)
    @time l3 = selectnode(fens; box=[0. A 0. 0.], inflate = 1.0/N/100.0)
    @time l4 = selectnode(fens; box=[0. A A A], inflate = 1.0/N/100.0)
    List = vcat(l1, l2, l3, l4)
    @time setebc!(Temp, List, true, 1, tempf(geom.values[List,:])[:])
    @time applyebc!(Temp)
    @time numberdofs!(Temp)

    t1 = time()

    material = MatHeatDiff(thermal_conductivity)

    femm = FEMMHeatDiff(IntegData(fes, TriRule(1)), CSys(Rm), material)

    println("Conductivity")
    @time K = conductivity(femm, geom, Temp)
    println("Nonzero EBC")
    @time F2 = nzebcloadsconductivity(femm, geom, Temp);
    println("Internal heat generation")
    fi = ForceIntensity(FFlt[Q]);
    @time F1 = distribloads(femm, geom, Temp, fi, 3);

    println("Factorization")
    @time K = cholesky(K)
    println("Solution of the factorized system")
    @time U = K\(F1+F2)
    scattersysvec!(Temp,U[:])

    println("Total time elapsed = $(time() - t0) [s]")
    println("Solution time elapsed = $(time() - t1) [s]")

    Error= 0.0
    for k=1:size(fens.xyz,1)
        Error = Error + abs.(Temp.values[k,1] - tempf(reshape(fens.xyz[k,:], (1,2)))[1])
    end
    println("Error =$Error")

    # File =  "a.vtk"
    # MeshExportModule.vtkexportmesh (File, fes.conn, [geom.values Temp.values], MeshExportModule.T3; scalars=Temp.values, scalars_name ="Temperature")

    true

end # Poisson_FE_example_csys_1

function Poisson_FE_Q4_example()
    println("""

    Heat conduction example described by Amuthan A. Ramabathiran
    http://www.codeproject.com/Articles/579983/Finite-Element-programming-in-Julia:
    Unit square, with known temperature distribution along the boundary,
    and uniform heat generation rate inside.  Mesh of regular four-node QUADRILATERALS,
    in a grid of 1000 x 1000 edges (1M quads, 1M degrees of freedom).
    Version: 05/29/2017
    """
    )
    t0 = time()

    A = 1.0
    thermal_conductivity =  [i==j ? one(FFlt) : zero(FFlt) for i=1:2, j=1:2]; # conductivity matrix
    function getsource!(forceout::FFltVec, XYZ::FFltMat, tangents::FFltMat, fe_label::FInt)
        forceout[1] = -6.0; #heat source
    end
    tempf(x) = (1.0 .+ x[:,1].^2 .+ 2.0 .* x[:,2].^2);#the exact distribution of temperature
    N = 1000;

    println("Mesh generation")
    @time fens,fes = Q4block(A, A, N, N)

    geom = NodalField(fens.xyz)
    Temp = NodalField(zeros(size(fens.xyz,1),1))

    println("Searching nodes  for BC")
    @time l1 = selectnode(fens; box=[0. 0. 0. A], inflate = 1.0/N/100.0)
    @time l2 = selectnode(fens; box=[A A 0. A], inflate = 1.0/N/100.0)
    @time l3 = selectnode(fens; box=[0. A 0. 0.], inflate = 1.0/N/100.0)
    @time l4 = selectnode(fens; box=[0. A A A], inflate = 1.0/N/100.0)
    List = vcat(l1, l2, l3, l4);
    setebc!(Temp, List, true, 1, tempf(geom.values[List,:])[:])
    applyebc!(Temp)

    numberdofs!(Temp)

    t1 = time()

    m = MatHeatDiff(thermal_conductivity)
    femm = FEMMHeatDiff(IntegData(fes, GaussRule(2, 2)), m)

    println("Conductivity")
    @time K=conductivity(femm, geom, Temp)
    #Profile.print()

    println("Nonzero EBC")
    @time F2 = nzebcloadsconductivity(femm, geom, Temp);
    println("Internal heat generation")
    fi = ForceIntensity(FFlt, 1, getsource!);
    @time F1 = distribloads(femm, geom, Temp, fi, 3);

    println("Factorization")
    @time K = cholesky(K)
    println("Solution of the factorized system")
    @time U=  K\(F1+F2)
    scattersysvec!(Temp, U[:])

    println("Total time elapsed = $(time() - t0) [s]")
    println("Solution time elapsed = $(time() - t1) [s]")

    # using MeshExportModule

    # File =  "a.vtk"
    # MeshExportModule.vtkexportmesh (File, fes.conn, [geom.values Temp.values], MeshExportModule.Q4; scalars=Temp.values, scalars_name ="Temperature")

    Error= 0.0
    for k=1:size(fens.xyz,1)
        Error = Error + abs.(Temp.values[k,1] - tempf(reshape(fens.xyz[k,:], (1,2)))[1])
    end
    println("Error =$Error")

    true
end # Poisson_FE_Q4_example

function allrun()
    println("#####################################################")
    println("# Poisson_FE_example ")
    Poisson_FE_example()
    println("#####################################################")
    println("# Poisson_FE_example_algo ")
    Poisson_FE_example_algo()
    println("#####################################################")
    println("# Poisson_FE_example_csys_1 ")
    Poisson_FE_example_csys_1()
    println("#####################################################")
    println("# Poisson_FE_Q4_example ")
    Poisson_FE_Q4_example()
end # function allrun

end # module Poisson_examples


using Stheno, Random
using Stheno: @model

######################## Example 1: Vanilla Toeplitz ########################

@model function vanilla()
    f = GP(EQ())
    y = f + GP(Noise(0.1))
    return f, y
end
f, y = vanilla();

# Create regularly spaced data. Dispatch handles literally everything else.
x = range(-5.0, stop=5.0, length=10000);

# Notice that we don't struggle to construct large covariances. All of these matrices
# are either Toeplitz or SymmetricToeplitz. This structure is exploited to accelerate
# all computations. These covariance matrices are constucted in O(N)-time.
Σff = cov(f(x));
Σfy = xcov(f(x), y(x));
Σyy = cov(y(x));

# These operations are also quite a lot faster. They occur in O(N^2)-time.
rng = MersenneTwister(123456);
ŷ = rand(rng, y(x));
@show logpdf(y(x), ŷ);

######################## Example 2: Block Vanilla Toeplitz ########################

# Suppose that we have two separate blocks of regularly spaced data eg. missing data.
xl, xr = range(-5.0, stop=0.0, length=5000), range(0.0, stop=5.0, length=5000);

# If we inspect the underlying blocks of the covariance, we see they are all still Toeplitz.
using Stheno: unbox
f_xl_xr = BlockGP([f(xl), f(xr)]);
y_xl_xr = BlockGP([y(xl), y(xr)]);
Σff = cov(f_xl_xr);
display(typeof.(unbox(unbox(Σff)).blocks));

# THIS CURRENTLY SEEMS A BIT SLOW! COMPARE TO DENSE PROPERLY!
rng = MersenneTwister(123456);
ŷ = rand(rng, y_xl_xr);
@show logpdf(y_xl_xr, ŷ);

######################### Example 3: Multi-Output Toeplitz ##########################

# Suppose now that we modify our model slightly. We construct a two-output process from
# three latent processes.
@model function simple_multi_output()
    f₁, f₂, f₃ = GP(EQ()), GP(EQ()), GP(EQ())
    y₁, y₂ = f₁ + f₂, f₂ + f₃
    return y₁, y₂
end
y₁, y₂ = simple_multi_output();

# Observe both at x. The joint covariance matrix is still Toeplitz!
y_1x_2x = BlockGP([y₁(x), y₂(x)]);
Σ12 = cov(y_1x_2x);
display(typeof.(unbox(unbox(Σ12)).blocks));

# Same again for the missing data.
y_1xl_2xr = BlockGP([y₁(xl), y₂(xr)]);
Σ12 = cov(y_1xl_2xr);
display(typeof.(unbox(unbox(Σ12)).blocks));

######################### Example 4: Temporally mis-aligned data ##########################

# What if each output has the same sampling rate, but the data are misaligned?
x₁ = range(-5.0, stop=5.0, length=5000);
x₂ = x₁ .+ 0.3;

# The joint covariance still comprises Toeplitz matrices.
y_x1_x2 = BlockGP([y₁(x₁), y₂(x₂)]);
Σ12 = cov(y_x1_x2);
display(typeof.(unbox(unbox(Σ12)).blocks));

######################### Example 5: Partial structure ###########################

# Suppose that we have some regularly sampled data, and some structure-less data.
x′ = randn(rng, 100);

# Notice that Σ_x_x′ comprised three dense blocks, and one Toeplitz. This is because
# cov(f(x′)) is dense, xcov(f(x), f(x′)) / xcov(f(x′), f(x)) is dense, whilst cov(f(x))
# is Toeplitz.
f_x_x′ = BlockGP([f(x), f(x′)]);
Σ_x_x′ = cov(f_x_x′);
display(typeof.(unbox(unbox(Σ_x_x′)).blocks))

# Note that any combination of these things will (should) also yield some combination of
# dense and Toeplitz matrices.


beta_a_from_mv(m::Float64, v::Float64) = (m*(1-m)/v-1)*m

beta_b_from_mv(m::Float64, v::Float64) = (m*(1-m)/v-1)*(1-m)

function beta_pars_from_mv(m::Float64, v::Float64)
  c = m*(1-m)/v-1
  return c*m, c*(1-m)
end


using ClassicalOrthogonalPolynomials, ContinuumArrays, BandedMatrices, LazyArrays, ForwardDiff, Test
import LazyArrays: rowsupport, colsupport

@testset "Ultraspherical" begin
    @testset "Transforms" begin
        U = Ultraspherical(1)
        x = axes(U,1)
        Un = U[:,Base.OneTo(5)]
        @test factorize(Un) isa ContinuumArrays.TransformFactorization
        @test (Un \ x) ≈ [0,0.5,0,0,0]
        @test (U * (U \ exp.(x)))[0.1] ≈ exp(0.1)
    end

    @testset "Operators" begin
        @testset "Lowering" begin
            λ = 1
            wC1 = Weighted(Ultraspherical(λ))
            wC2 = Weighted(Ultraspherical(λ+1))
            L = wC1 \ wC2
            @test L isa BandedMatrix
            @test bandwidths(L) == (2,0)
            @test L[1,1] ≈ 2λ*(2λ+1)/(4λ*(λ+1))
            @test L[3,1] ≈ -2/(4λ*(λ+1))
            u = [randn(10); zeros(∞)]
            @test wC1[0.1,:]'*(L*u) ≈ wC2[0.1,:]'*u
        end

        @testset "Weighted Derivative" begin
            T = Chebyshev()
            wC1 = Weighted(Ultraspherical(1))
            wC2 = Weighted(Ultraspherical(2))
            x = axes(wC2,1)
            D = Derivative(x)

            @test wC1 \ (D*wC2) isa BandedMatrix
            @test (wC1 \ (D*wC2))[1:5,1:5] == BandedMatrix(-1 => [-1.5,-4,-7.5,-12])

            u = wC2 * [randn(5); zeros(∞)]
            @test (D*u)[0.1] ≈ ((D*T) * (T\u))[0.1]
        end

        @testset "Interrelationships" begin
            @testset "Chebyshev–Ultrashperical" begin
                T = ChebyshevT()
                U = ChebyshevU()
                C = Ultraspherical(2)
                D = Derivative(axes(T,1))

                @test C\C === pinv(C)*C === Eye(∞)
                D₀ = U\(D*T)
                D₁ = C\(D*U)
                @test D₁ isa BandedMatrix
                @test (D₁*D₀)[1:10,1:10] == diagm(2 => 4:2:18)
                @test D₁*D₀ isa MulMatrix
                @test bandwidths(D₁*D₀) == (-2,2)

                S₁ = (C\U)[1:10,1:10]
                @test S₁ isa BandedMatrix{Float64}
                @test S₁ == diagm(0 => 1 ./ (1:10), 2=> -(1 ./ (3:10)))

                @test (U\C)[1:10,1:10] ≈ inv((C\U)[1:10,1:10])
                @test (T\C)[1:10,1:10] ≈ inv((C\T)[1:10,1:10])
                @test bandwidths(U\C) == bandwidths(T\C) == (0,∞)
                @test colsupport(U\C,5) == colsupport(T\C,5) == 1:5
                @test rowsupport(U\C,5) == rowsupport(T\C,5) == 5:∞
            end
            @testset "Legendre" begin
                @test Ultraspherical(0.5) \ (UltrasphericalWeight(0.0) .* Ultraspherical(0.5)) == Eye(∞)
                @test Legendre() \ (UltrasphericalWeight(0.0) .* Ultraspherical(0.5)) == Eye(∞)
                @test (Legendre() \ Ultraspherical(1.5))[1:10,1:10] ≈ inv(Ultraspherical(1.5) \ Legendre())[1:10,1:10]
            end
        end

        @testset "Conversion" begin
            R = Ultraspherical(3.5) \ Ultraspherical(0.5)
            c = [[1,2,3,4,5]; zeros(∞)]
            @test Ultraspherical(3.5)[0.1,:]' * (R * c) ≈ Ultraspherical(0.5)[0.1,:]' * c
            Ri = Ultraspherical(0.5) \ Ultraspherical(3.5)
            @test Ri[1:10,1:10] ≈ inv(R[1:10,1:10])
            @test Ultraspherical(0.5)[0.1,:]' * (Ri * c) ≈ Ultraspherical(3.5)[0.1,:]' * c
        end
    end

    @testset "test on functions" begin
        T = Chebyshev()
        U = Ultraspherical(1)
        D = Derivative(axes(T,1))
        f = T*Vcat(randn(10), Zeros(∞))
        @test (U*(U\f)).args[1] isa Ultraspherical
        @test (U*(U\f))[0.1] ≈ f[0.1]
        @test (D*f)[0.1] ≈ ForwardDiff.derivative(x -> (ChebyshevT{eltype(x)}()*f.args[2])[x],0.1)
    end

    @testset "Evaluation" begin
        C = Ultraspherical(2)
        @test @inferred(C[0.1,Base.OneTo(0)]) == Float64[]
        @test @inferred(C[0.1,Base.OneTo(1)]) == [1.0]
        @test @inferred(C[0.1,Base.OneTo(2)]) == [1.0,0.4]
        @test @inferred(C[0.1,Base.OneTo(3)]) == [1.0,0.4,-1.88]
    end

    @testset "special syntax" begin
        @test ultrasphericalc.(0:5, 2, 0.3) == Ultraspherical(2)[0.3, 1:6]
    end

    @testset "Ultraspherical vs Chebyshev and Jacobi" begin
        @test Ultraspherical(1) == ChebyshevU()
        @test ChebyshevU() == Ultraspherical(1)
        @test Ultraspherical(0) ≠ ChebyshevT()
        @test ChebyshevT() ≠ Ultraspherical(0)
        @test Ultraspherical(1) ≠ Jacobi(1/2,1/2)
        @test Jacobi(1/2,1/2) ≠ Ultraspherical(1)
        @test Ultraspherical(1/2) == Jacobi(0,0)
        @test Ultraspherical(1/2) == Legendre()
        @test Jacobi(0,0) == Ultraspherical(1/2)
        @test Legendre() == Ultraspherical(1/2)

        @test Ultraspherical(1/2) \ (JacobiWeight(0,0) .* Jacobi(0,0)) isa Diagonal
        @test (JacobiWeight(0,0) .* Jacobi(0,0)) \ Ultraspherical(1/2) isa Diagonal
    end

    @testset "D^2 * mapped" begin
        T = chebyshevt(0..1)
        C = ultraspherical(2,0..1)
        r = axes(T,1)
        D = Derivative(r)
        D₂ = C \ (D^2 * T)
        # r²D₂ = C \ (r.^2 .* (D^2 * T))

        c = [randn(100); zeros(∞)]
        @test C[0.1,:]'*(D₂ * c) ≈ 4*(Derivative(axes(ChebyshevT(),1))^2 * (ChebyshevT() * c))[2*0.1-1]
        @test_broken C[0.1,:]'*(r²D₂ * c) ≈ 0.1^2 * C[0.1,:]'*(D₂ * c)
    end

    @testset "BigFloat" begin
        U = Ultraspherical{BigFloat}(1)
        T = ChebyshevT{BigFloat}()
        x = axes(U,1)
        D = Derivative(x)

        @test Weighted(T) \ (D * Weighted(U)) isa BandedMatrix{BigFloat}

        C³ = Ultraspherical{BigFloat}(3)
        c = [1; 2; 3; zeros(BigFloat,∞)]
        @test C³[big(1)/10,:]'*(C³ \ U) * c ≈ U[big(1)/10,:]'c
    end
end


using StanSample, MCMCChains

# Define the Stan language model

m8_2s = "
data{
    int N;
    vector[N] y;
}
parameters{
    real mu;
    real<lower=0> sigma;
}
model{
    y ~ normal( mu , sigma );
}
";

# Define the Stanmodel.

m_8_2s = SampleModel("m8.2s", m8_2s);

# Input data for cmdstan

m8_2_data = Dict("N" => 2, "y" => [-1, 1]);
m8_2_init = Dict("mu" => 0.0, "sigma" => 1.0);

# Sample using cmdstan

rc = stan_sample(m_8_2s; data=m8_2_data, init=m8_2_init);

# Describe the draws
if success(rc)
  chn = read_samples(m_8_2s; output_format=:mcmcchains)
  describe(chn)
end


include("fourier_utils.jl")
include("fourier_operations.jl")
include("fourier_fields.jl")

# Define Interactions
function (∇::FourierOperator)(ϕ::FourierField)
    return FourierField(∇.op .* ϕ.data, ϕ.metadata)
end
function *(∇::FourierOperator, ϕ::FourierField)
    return FourierField(∇.op .* ϕ.data, ϕ.metadata)
end
# Define Naming
function (∇::FourierOperator{S, T})(ϕ::FourierField) where
    {S, T <: FourierOperatorMetaData}
    op_name = ∇.metadata.name
    ϕ_name = ϕ.metadata.name
    new_name = op_name * "(" * ϕ_name * ")"
    fmd = FourierMetaData(new_name, ϕ.metadata.grid, ϕ.metadata.transform)
    return FourierField(∇.op .* ϕ.data, fmd)
end


module Covering

import Voronoi

abstract type Element end

mutable struct Edge <: Element
   origin::Tuple{Real, Real}
   next::Union{Element, Nothing}
end

mutable struct Arc <: Element
   origin::Tuple{Real, Real}
   next::Union{Element, Nothing}
end

mutable struct Section
   center::Union{Tuple{Real, Real}, Nothing}
   borderHead::Union{Element, Nothing}
end

function Section(points::Array{Tuple{Real, Real}, 1})
    n = size(points)[1]

    edge = head = Edge(points[1], nothing)

    for i in 2:n
        edge.next = Edge(points[i], nothing)
        edge = edge.next
    end

    edge.next = head

    return Section((0, 0), head)
end

function Section(section::Section)
    edge = head = Edge(section.borderHead.origin, nothing)

    copying = section.borderHead.next
    while copying != section.borderHead
        edge.next = Edge(copying.origin, nothing)
        edge = edge.next
        copying = copying.next
    end

    edge.next = head

    return Section((0, 0), head)
end

function segmentArcIntersections(segment::Voronoi.Diagram.HalfEdge, center::Tuple{Real, Real}, r::Real)
    return segmentArcIntersections(segment.origin, segment.next.origin, center, r)
end

function segmentArcIntersections(p::Tuple{Real, Real}, q::Tuple{Real, Real}, center::Tuple{Real, Real}, r::Real)
    function dot(a, b)
        return a[1]*b[1] + a[2]*b[2]
    end

    PQ = Voronoi.Geometry.subVector(q, p)
    CP = Voronoi.Geometry.subVector(p, center)

    a = dot(PQ, PQ)
    b = 2*dot(CP, PQ)
    c = dot(CP, CP) - r^2

    Δ = b^2 - 4a*c # discriminant

    if Δ < 0 # no intersection
        return []
    elseif Δ == 0 # one intersection
        α = (-b - sqrt(Δ))/2a

        if 0 <= α <= 1
            return [(p[1] + α*PQ[1], p[2] + α*PQ[2])]
        else
            return []
        end
    else # two intersections
        α1 = (-b - sqrt(Δ))/2a
        α2 = (-b + sqrt(Δ))/2a

        if (α1 < α2)
            if 0 <= α1 <= α2 <= 1
                return [(p[1] + α1*PQ[1], p[2] + α1*PQ[2]), (p[1] + α2*PQ[1], p[2] + α2*PQ[2])]
            elseif 0 <= α1 <= 1
                return [(p[1] + α1*PQ[1], p[2] + α1*PQ[2])]
            elseif 0 <= α2 <= 1
                return [(p[1] + α2*PQ[1], p[2] + α2*PQ[2])]
            else
                return []
            end
        else
            if 0 <= α2 <= α1 <= 1
                return [(p[1] + α2*PQ[1], p[2] + α2*PQ[2]), (p[1] + α1*PQ[1], p[2] + α1*PQ[2])]
            elseif 0 <= α1 <= 1
                return [(p[1] + α1*PQ[1], p[2] + α1*PQ[2])]
            elseif 0 <= α2 <= 1
                return [(p[1] + α2*PQ[1], p[2] + α2*PQ[2])]
            else
                return []
            end
        end
    end
end

function distanceOutsideCircle(a, center, r)
    d = Voronoi.Geometry.distanceSquared(a, center)
    return d - r^2
end

function intersectedSegment(segment::Voronoi.Diagram.HalfEdge, center::Tuple{Real, Real}, r::Real)
    p = segment.origin
    q = segment.next.origin

    docp = distanceOutsideCircle(p, center, r)
    docq = distanceOutsideCircle(q, center, r)

    if docp <= 0 && docq <= 0
        el = head = Edge(p, nothing)
    else
        intersections = segmentArcIntersections(segment, center, r)

        if docq >= 0 && docp >= 0 && length(intersections) == 1
            el = head = Arc(p, nothing)
        else
            if docp < 0
                inside = true
                el = head = Edge(p, nothing)
            else
                inside = false
                el = head = Arc(p, nothing)
            end

            for i in intersections
                if inside
                    el.next = Arc(i, nothing);
                else
                    el.next = Edge(i, nothing);
                end

                el = el.next
                inside = !inside
            end
        end
    end

    return head, el
end

function voronoiRegionToCoveringSection(region::Voronoi.Diagram.Region, r::Real)
    he = region.borderHead.next
    head, tail = intersectedSegment(he, region.generator, r)

    el = head
    while (he != region.borderHead)
        he = he.next
        tail.next, tail = intersectedSegment(he, region.generator, r)
    end
    tail.next = head

    return Section(region.generator, head)
end

mutable struct Partition
    sections::Array{Covering.Section}
    r::Real
end

function voronoiDiagramToPartition(V::Voronoi.Diagram.DCEL, r::Real)
    n = length(V.regions)
    sections = Array{Section, 1}(undef, n)

    for i in 1:n
        sections[i] = voronoiRegionToCoveringSection(V.regions[i], r)
    end

    return Partition(sections, r)
end

function radianToDegrees(rad)
    return rad*180/π
end

function segmentAngle(p, q)
    num = q[2] - p[2]
    den = q[1] - p[1]
    if num == 0 && den > 0
        angle = 0
    elseif num == 0 && den < 0
        angle = π
    elseif den == 0 && num > 0
        angle = π/2
    elseif den == 0 && num < 0
        angle = 3π/2
    elseif num > 0 && den > 0
        angle = atan(num/den)
    elseif num > 0 && den < 0
        angle = π - atan(-num/den)
    elseif num < 0 && den < 0
        angle = π + atan(num/den)
    else # num < 0 && den > 0
        angle = -atan(-num/den)
    end

    return angle
end

function areaAndGradient(section::Section, r::Real)
    if section.borderHead == nothing
        return 0, 0, (0, 0)
    end

    coveredArea = 0
    el = section.borderHead
    gᵣ = 0
    gₛx = gₛy = 0.0

    firstTime = true # gimme do whiles, Julia, dammit!
    while (el != section.borderHead || firstTime)
        firstTime = false
        if (el isa Arc)
            θ1 = segmentAngle(section.center, el.origin)
            θ2 = segmentAngle(section.center, el.next.origin)
            θ = θ2 - θ1
            if θ < 0
                θ += 2π
            end

            coveredArea += section.center[1]*r*(sin(θ2) - sin(θ1))
            coveredArea += ((r^2)/2)*(θ + sin(θ2)*cos(θ2) - sin(θ1)*cos(θ1))

            gᵣ += -r*θ # (θ/(2π))*2πr
            gₛx += r*(sin(θ1) - sin(θ2))
            gₛy += r*(cos(θ2) - cos(θ1))
        else # isa Edge
            v = el.origin
            w = el.next.origin

            coveredArea += (w[2] - v[2])*(v[1] + w[1])/2
        end

        el = el.next
    end

    return coveredArea, gᵣ, (gₛx, gₛy)
end

function areaAndGradient(P::Partition)
    n = length(P.sections)

    coveredArea = gᵣ = 0
    gₛ = fill((0.0, 0.0), n)

    for i in 1:n
        a, g_r, g_s_i = areaAndGradient(P.sections[i], P.r)
        coveredArea += a
        gᵣ += g_r
        gₛ[i] = g_s_i
    end

    return coveredArea, gᵣ, gₛ
end

end # module


abstract type CNFLayer <: Function end
Flux.trainable(m::CNFLayer) = (m.p,)

"""
Constructs a continuous-time recurrent neural network, also known as a neural
ordinary differential equation (neural ODE), with fast gradient calculation
via adjoints [1] and specialized for density estimation based on continuous
normalizing flows (CNF) [2]. At a high level this corresponds to the following steps:

1. Parametrize the variable of interest x(t) as a function f(z,θ,t) of a base variable z(t) with known density p_z;
2. Use the transformation of variables formula to predict the density p_x as a function of the density p_z and the trace of the Jacobian of f;
3. Choose the parameter θ to minimize a loss function of p_x (usually the negative likelihood of the data);

After these steps one may use the NN model and the learned θ to predict the density p_x for new values of x.

```julia
NeuralODE(model,basedist=nothing,monte_carlo=false,tspan,args...;kwargs...)
```
Arguments:
- `model`: A Chain neural network that defines the ̇x.
- `basedist`: Distribution of the base variable. Set to the unit normal by default.
- `monte_carlo`: Method for calculating the trace of the Jacobian. The default monte_carlo = false calculates the Jacobian and its trace directly. monte_carlo = true uses the stochastic approach presented in [3] to provide an unbiased estimate for the trace.
- `tspan`: The timespan to be solved on.
- `kwargs`: Additional arguments splatted to the ODE solver. See the
  [Common Solver Arguments](https://docs.sciml.ai/dev/basics/common_solver_opts/)
  documentation for more details.
Ref
[1]L. S. Pontryagin, Mathematical Theory of Optimal Processes. CRC Press, 1987.
[2]R. T. Q. Chen, Y. Rubanova, J. Bettencourt, D. Duvenaud. Neural Ordinary Differential Equations. arXiv preprint at arXiv1806.07366, 2019.
[3]W. Grathwohl, R. T. Q. Chen, J. Bettencourt, I. Sutskever, D. Duvenaud. FFJORD: Free-Form Continuous Dynamic For Scalable Reversible Generative Models. arXiv preprint at arXiv1810.01367, 2018.

"""
struct FFJORD{M,P,RE,Distribution,Bool,T,A,K} <: CNFLayer
    model::M
    p::P
    re::RE
    basedist::Distribution
    monte_carlo::Bool
    tspan::T
    args::A
    kwargs::K

    function FFJORD(model,tspan,args...;p = nothing,basedist=nothing,monte_carlo=false,kwargs...)
        _p,re = Flux.destructure(model)
        if p === nothing
            p = _p
        end
        if basedist === nothing
            size_input = size(model[1].W)[2]
            basedist = MvNormal(zeros(size_input), I + zeros(size_input,size_input))
        end
        new{typeof(model),typeof(p),typeof(re),typeof(basedist),
            typeof(monte_carlo),typeof(tspan),typeof(args),typeof(kwargs)}(
            model,p,re,basedist,monte_carlo,tspan,args,kwargs)
    end
end

function jacobian_fn(f, x::AbstractVector)
   y::AbstractVector, back = Zygote.pullback(f, x)
   ȳ(i) = [i == j for j = 1:length(y)]
   vcat([transpose(back(ȳ(i))[1]) for i = 1:length(y)]...)
 end

function ffjord(du,u,p,t,re,monte_carlo,e)
    z = @view u[1:end-1]
    m = re(p)
    if monte_carlo
        _, back = Zygote.pullback(m,z)
        eJ = back(e)[1]
        trace_jac = (eJ.*e)[1]
    else
        J = jacobian_fn(m, z)
        trace_jac = length(z) == 1 ? sum(J) : tr(J)
    end
    du[1:end-1] = m(z)
    du[end] = -trace_jac
end

function (n::FFJORD)(x,p=n.p,monte_carlo=n.monte_carlo)
    e = monte_carlo ? randn(Float32,length(x)) : nothing
    ffjord_ = (du,u,p,t)->ffjord(du,u,p,t,n.re,monte_carlo,e)
    prob = ODEProblem{true}(ffjord_,vcat(x,0f0),n.tspan,p)
    pred = solve(prob,n.args...;n.kwargs...)[:,end]
    pz = n.basedist
    z = pred[1:end-1]
    delta_logp = pred[end]
    logpz = log.(pdf(pz, z))
    logpx = logpz .- delta_logp
    return logpx[1]
end


struct mv_gaussian <: distibution_sample
    μ::AbstractArray{Float64,1}
    Σ::AbstractArray{Float64,2}
    invΣ::AbstractArray{Float64,2}
    logdetΣ::Float64
    # mvn::MvNormal
end

function dinvquad!(r,a,x)
    dcolwise_dot!(r,x, a \ x)
end

function log_likelihood!(r::AbstractArray,x::AbstractArray, distibution_sample::mv_gaussian , group::Int64 = -1)
    z = x .- distibution_sample.μ
    dcolwise_dot!(r,z, distibution_sample.invΣ * z)
    r .= -((length(distibution_sample.Σ) * Float64(log(2π)) + logdet(distibution_sample.Σ))/2) .-r
end


# auxiliary model: HAR-RV
# Corsi, Fulvio. "A simple approximate long-memory model
# of realized volatility." Journal of Financial Econometrics 7,
# no. 2 (2009): 174-196.
function HAR(y)
    ylags = lags(y,10)
    X = [ones(size(y,1)) ylags[:,1]  mean(ylags[:,1:4],dims=2) mean(ylags[:,1:10],dims=2)]
    # drop missings
    y = y[11:end]
    X = X[11:end,:]
    βhat = X\y
    σhat = std(y-X*βhat)
    vcat(βhat,σhat)
end


module DL

using Printf
using LinearAlgebra
import Base.*, Base.+

export (*), (+), ReLU, VariableTensor, MSE, ConstantTensor, backward, zero_grad, SGD, Linear, Tensor

# Tensor
abstract type Tensor end
abstract type LeafTensor <: Tensor end
abstract type FunctionalTensor <: Tensor end

#= 掛け算 t = t1 * t2 =#
mutable struct MulTensor <: FunctionalTensor
    x
    tensors::Array{Tensor, 1}
end
(*)(t3::Tensor, t4::Tensor) = abstract_operator(MulTensor, (x1, x2) -> x1 * x2, [t3, t4])
# t2 <- t1 <- t3 * t4
function grad_fn(t1::MulTensor, grad_up)
    t3, t4 = t1.tensors
    # dt2/dt1 * dt1/dt3 = grad_up * t4.x
    # dt2/dt1 * dt1/dt4 = grad_up * t3.x
    return [t4.x * grad_up, grad_up * t3.x]
end

#= 足し算 =#
mutable struct AddTensor <: FunctionalTensor
    x
    tensors::Array{Tensor, 1}
end
# forward
(+)(t3::Tensor, t4::Tensor) = abstract_operator(AddTensor, (x1, x2) -> x1 .+ x2, [t3, t4])
# backward t2 <- t1 <- t3 + t4
function grad_fn(t1::AddTensor, grad_up)
    t3, t4 = t1.tensors
    # dt2/dt1 * dt1/dt3 = grad_up * I
    # dt2/dt1 * dt1/dt4 = grad_up * I
    return [grad_up, grad_up]
end

#= ReLU function t = ReLU(t1) =#
mutable struct ReLUTensor <: FunctionalTensor
    x
    tensors::Array{Tensor, 1}
end
# forward
function ReLU(t::Tensor)
    return abstract_operator(ReLUTensor, x -> x .* (x .> 0), [t])
end
# backward t2 <- t1 <- ReLU(t3)
function grad_fn(t1::ReLUTensor, grad_up)
    t3 = t1.tensors[1]
    # dt2/dt1 * dt1/dt3 = grad_up * d ReLU(x)/dx
    return [grad_up .* (t3.x' .> 0)]
end

#= Mean Squared Error t = MSE(t1, t2) =#
mutable struct MSETensor <: FunctionalTensor
    x
    tensors::Array{Tensor, 1}
end
# forward
function MSE(t1::Tensor, t2::Tensor)
    mse = (x1, x2) -> sum((x1 - x2).^2) / length(x1)
    return abstract_operator(MSETensor, mse, [t1, t2])
end
# t2 <- t1 <- MSE(t3, t4)
function grad_fn(t1::MSETensor, grad_up)
    t3, t4 = t1.tensors
    # dt2/dt1 * dt1/dt3 = grad_up * d/dx sum((x - y)^2) / length(x)
    N = length(t3.x)
    dt1dt3 = 2*(t3.x - t4.x)' / N
    dt1dt4 = 2*(t4.x - t3.x)' / N
    return [grad_up * dt1dt3, grad_up * dt1dt4]
end

#= t2 <- t1 末端の部分 微分あり =#
mutable struct VariableTensor <: LeafTensor
    x
    grad
    VariableTensor(x) = new(x, zero(x))
end

#= ただの定数 =#
mutable struct ConstantTensor <: LeafTensor
    x
end

# backward 計算
# 最初
function backward(t::Tensor)
	backward(t, 1)
end
# 再帰
function backward(t::Tensor, grad_up)
    # 微分はtが持つTensorsの数だけあるから，リストで返すことにする．
    vgrad = grad_fn(t, grad_up)
	# 微分と対応するTensorのペアをつくる
    pair = zip(t.tensors, vgrad)
    # 掘り下げ
	map(x -> backward(x[1], x[2]), pair)
end
# ConstantTensorとか
function backward(t::LeafTensor, grad_up)
end
# VariableTensorは勾配を蓄積していく
function backward(t::VariableTensor, grad_up)
    t.grad += grad_up'
end

# 抽象化された関数
# 演算子を計算するときにConstant同士の演算の出力はConstantであるようにする
# Array{T, 1} where T <: Tensor という書き方でTensor以下のタイプのTのArrayを網羅できる
function abstract_operator(EnzansiTensor::DataType, f::Function, tensors::Array{T, 1} where T <: Tensor)
    # Tensor.xを展開してfに代入する
    y = f(map(tensor -> tensor.x, tensors)...)
    return EnzansiTensor(y, tensors)
end
function abstract_operator(EnzansiTensor::DataType, f::Function, tensors::Array{ConstantTensor, 1})
    y = f(map(tensor -> tensor.x, tensors)...)
    return ConstantTensor(y)
end

function tensor_map(f::Function, t::Tensor)
    #= 任意の関数をVariableTensorに作用させる
    毎回掘り下げる必要はないと思うので考えなおす余地あり
    =#
    # 掘り下げていく
    map(t -> tensor_map(f, t), t.tensors)
end
function tensor_map(f::Function, t::LeafTensor)
    # 末端だけどVariableTensor以外のtypeはパラメータを変えない
    nothing
end
function tensor_map(f::Function, t::VariableTensor)
    # VariableTensorにだけ作用させる
    f(t)
end

function zero_grad(tensor::Tensor)
    f = t -> t.grad = zero(t.grad)
    tensor_map(f, tensor)
end

function SGD(lr=1e-3)
    f = t::Tensor -> t.x = t.x - lr * t.grad
    optimizer = t::Tensor -> tensor_map(f, t)
    return optimizer
end

function Linear(in_features::Int64, out_features::Int64)
    W = VariableTensor(randn(out_features, in_features))
    b = VariableTensor(zeros(out_features, 1))
    @show b.x
    return x -> W*x + b
end

function printarray(x)
    show(stdout, "text/plain", x)
    @printf "\n"
end

end


# Code derived from metabolicEP (https://github.com/anna-pa-m/Metabolic-EP)

function inplaceinverse!(dest::AbstractArray, source::AbstractArray)
    copyto!(dest, source)
    inv!(cholesky!(Hermitian(dest)))
end

Φ(x) = 0.5*(1.0+erf(x/sqrt(2.0)))
ϕ(x) = exp(-x.^2/2.0)/sqrt(2π)
ϕ(x, μ, σ) = inv(σ*sqrt(2π))*exp(-((x - μ)^2)/(2σ^2))


#=
    Given three points (in their natural sequence) and a contained or nearby fourth abcissa,
    interpolate or extrapolate the fourth ordinate(s) using a polynomial or a rational approach.

        quadratic(pt₁, pt₂, pt₃, x)

        thiele3(pt₁, pt₂, pt₃, x)
=#

function quadratic(pt₁, pt₂, pt₃, x)
    result = quadratic(pt₁[1], pt₂[1], pt₃[1],
                       pt₁[2:end], pt₂[2:end], pt₃[2:end], x)
    # check for degenerate case
    if !allfinite(result)
	result = linear(pt₁, pt₂, pt₃, x)
    end
	return result
end

function quadratic(pt₁x, pt₂x, pt₃x, pt₁y, pt₂y, pt₃y, x)
    t1 = @. pt₂x - pt₁x
    t2 = @. pt₁x - pt₃x
    t3 = @. pt₃x - pt₂x
    t4 = @. t2   * pt₂y
    t5 = @. pt₃x * pt₃x
    t6 = @. pt₂x * pt₂x
    t7 = @. pt₁x * pt₁x

    s = @. -(inv(t1) * inv(t2) * inv(t3))

    a = @. (pt₁y * t3 + pt₃y * t1 + t4) * x
    b = @. t5 * (pt₁y - pt₂y)
    c = @. t6 * (pt₃y - pt₁y)
    d = @. t7 * (pt₂y - pt₃y)
    q = @. t6 * (pt₁y * pt₃x - pt₃y * pt₁x)
    r = @. t4 *  pt₁x * pt₃x
    n = @. (-pt₁y * t5 + pt₃y * t7) * pt₂x
    aa = @. a - b - c - d
    bb = @. r - n - q

    result = @. (s * (aa * x + bb))
    return result # may contain !finite values
end

function thiele3(pt₁, pt₂, pt₃, x)
    result = thiele3(pt₁[1], pt₂[1], pt₃[1],
                     pt₁[2:end], pt₂[2:end], pt₃[2:end], x)
    # check for degenerate case
    if !allfinite(result)
        result = quadratic(pt₁, pt₂, pt₃, x)
    end
    return result
end

function thiele3(pt₁x, pt₂x, pt₃x, pt₁y, pt₂y, pt₃y, x)
    t1 = @. pt₁y - pt₂y
    t2 = @. pt₂y - pt₃y
    t1 = @. inv(t1)
    t2 = @. inv(t2)
    t1 = @. (pt₁x - pt₂x) * t1
    t2 = @. -(pt₂x - pt₃x) * t2 + t1
    t2 = @. inv(t2)
    t2 = @. (pt₁x - pt₃x) * t2 - pt₁y + pt₂y
    t2 = @. inv(t2)
    t1 = @. -(pt₂x - x) * t2 + t1
    t1 = @. inv(t1)

    result = @. (-(pt₁x - x) * t1 + pt₁y)
    if !allfinite(result)
        result = (pt₁x, pt₂x, pt₃x, pt₁y, pt₂y, pt₃y, x)
    end

    return result
end


### routines for approximating functions.

function smoothmin(x::Vector{T})::T where T <: Real
    return smoothextrema(x,convert(T,-999.0))
end

function smoothmax(x::Vector{T})::T where T <: Real
    return smoothextrema(x,convert(T,999.0))
end

function smoothextrema(x::Vector{T}, α::T) where T <: Real
    u = collect( log(x[i]) + x[i]*α for i = 1:length(x) )
    v = collect( x[i]*α for i = 1:length(x) )
    return exp( StatsFuns.logsumexp(u) - StatsFuns.logsumexp(v) )
    #return sum( x[i]*exp(x[i]*α) for i = 1:length(x))/sum( exp(x[i]*α) for i = 1:length(x))
end
# ## test code.
# α = 999.0
# D = 5
# x = randn(D)
# t = rand()
# y = t.*x
# smax_x = smoothextrema(x,α)
# smax_y = smoothextrema(y,α)
# println("smooth max of x is ", smax_x)
# println("max of x is ", maximum(x))
# println("smooth max of y is ", smax_y)
# println("max of y is ", maximum(y))
#
# smin_x = smoothextrema(x,-α)
# smin_y = smoothextrema(y,-α)
# println("smooth min of x is ", smin_x)
# println("min of x is ", minimum(x))
# println("smooth min of y is ", smin_y)
# println("min of y is ", minimum(y))


using OrdinaryDiffEq
using GLMakie

#=
 # 1. simplest case

 ̈x =  -sin(x) - λ ̇x

=#
dampedpendulum(du, u, p, t) = @. -sin(u) - p.λ * du

du0 = [0.]
u0  = [1.0]
tspan = (0.0, 100.0)
p = (; λ = 0.1)

prob = SecondOrderODEProblem(dampedpendulum, du0, u0, tspan, p)
sol = solve(prob, VelocityVerlet())

# other solvers, e.g. DPRKN6

begin
    f = Figure()

    ax = Axis(f[1,1], title = "Position")
    lines!(sol.t, sol[1,:])

    ax = Axis(f[1,2], title = "Velocity")
    lines!(sol.t, sol[2,:])

    Label(f[0,1:2], "Damped pendulum", textsize = 24)

    mkpath("plots")
    save("plots/01_damped_pendulum.png", f)
    f
end


# This file is part of IntegerSequences.
# Copyright Peter Luschny. License is MIT.

(@__DIR__) ∉ LOAD_PATH && push!(LOAD_PATH, (@__DIR__))

module PrimeSwingFactorial

using PrimesIterator, Products, Nemo

export ModulePrimeSwingFactorial
export PSfactorial, Swing

"""
 Cf. P. Luschny, [Swing, divide and conquer the factorial](https://oeis.org/A000142/a000142.pdf), excerpt.

* PSfactorial, Swing
"""
const ModulePrimeSwingFactorial = ""

const SwingOddpart = [1,1,1,3,3,15,5,35,35, 315, 63, 693, 231,
   3003, 429, 6435, 6435, 109395,12155,230945,46189,969969,
   88179,2028117, 676039,16900975,1300075,35102025,5014575,
   145422675,9694845,300540195,300540195]

"""
Computes the odd part of the swinging factorial ``n≀`` (cf. A163590).
"""
function swing_oddpart(n::Int)
    n < 33 && return ZZ(SwingOddpart[n+1])

    sqrtn = isqrt(n)
    factors = Primes(div(n,2) + 1, n)
    P = Primes(sqrtn + 1, div(n, 3))
    s = [p for p in P if isodd(div(n, p))]

    for prime in Primes(3, sqrtn)
        p, q = 1, n
        while true
            q = div(q, prime)
            q == 0 && break
            isodd(q) && (p *= prime)
        end
        p > 1 && push!(s, p)
    end

    return ∏(factors)*∏(s)
end

"""
Computes the swinging factorial (a.k.a. Swing numbers n≀) (cf. A056040).
"""
function Swing(n::Int)
    sh = count_ones(div(n, 2))
    swing_oddpart(n) << sh
end

const FactorialOddPart = [1, 1, 1, 3, 3, 15, 45, 315, 315, 2835, 14175, 155925,
    467775, 6081075, 42567525, 638512875, 638512875, 10854718875, 97692469875,
    1856156927625, 9280784638125, 194896477400625, 2143861251406875,
    49308808782358125, 147926426347074375, 3698160658676859375]

"""
Return the largest odd divisor of ``n!``. Cf. A049606.
"""
function factorial_oddpart(n::Int)
    n < length(FactorialOddPart) && return ZZ(FactorialOddPart[n+1])
    swing_oddpart(n)*(factorial_oddpart(div(n,2))^2)
end

"""
Return the factorial ``n! = 1×2× ... ×n``, which is the order of the symmetric group S_n or the number of permutations of n letters (cf. A000142).
"""
function PSfactorial(n::Int)
    n < 0 && ArgumentError("Argument must be >= 0")
    sh = n - count_ones(n)
    factorial_oddpart(n) << sh
end

#START-TEST-########################################################

using Test, BenchmarkTools

function test()
    @testset "PrimeSwingF" begin
        for n in 0:999
            S = PSfactorial(n)
            B = Base.factorial(BigInt(n))
            #@test S == B
            @test all(S .== B)
        end
    end
end

function demo()
end

"""
10^1  128.016 ns (3        allocations:     48 bytes)
10^2    8.148 μs (78       allocations:   3.31 KiB)
10^3   68.933 μs (469      allocations:  16.53 KiB)
10^4  832.578 μs (3780     allocations: 106.00 KiB)
10^5   17.717 ms (31657    allocations: 816.70 KiB)
10^6  288.445 ms (263237   allocations:   6.53 MiB)
10^7    5.484  s (2209289  allocations:  53.95 MiB)
10^8   86.723  s (18926351 allocations: 446.32 MiB)
10^9 1218.261  s (167.46 M allocations: 3.828 GiB)
"""
function superperf()
    GC.gc()
    println(10)
    @btime PSfactorial(10)
    GC.gc()
    println(100)
    @btime PSfactorial(100)
    GC.gc()
    println(1000)
    @btime PSfactorial(1000)
    GC.gc()
    println(10000)
    @btime PSfactorial(10000)
    GC.gc()
    println(100000)
    @btime PSfactorial(100000)
    GC.gc()
    println(1000000)
    @btime PSfactorial(1000000)
    GC.gc()
    println(10000000)
    @btime PSfactorial(10000000)
    GC.gc()
    println(100000000)
    @btime PSfactorial(100000000)
    GC.gc()
    println(1000000000)
    @time PSfactorial(1000000000)
end

function perf()
    GC.gc()
    @time Base.factorial(BigInt(1000000))
    @time PSfactorial(1000000)
end

function main()
    test()
    demo()
    perf()
end

main()

end # module


n = rand(40:100)
A = matrixdepot("gilbert", n)
@test issym(A)

B = matrixdepot("gilbert", n, 0.2)
C = matrixdepot("gilbert", Int, n, 0.5)
@test nnz(B) <= nnz(C)
println("'gilbert' passed test ...")


# ------------------------------------------------------------------
# Code to simulate logit data and estimate nonparametrically
# Plots estimates of own and cross-price elasticities
# Written by James Brand
# ------------------------------------------------------------------
using Statistics, NPDemand
using RCall, DataFrames
@rlibrary ggplot2

J=2; # number of products
T =5000;
beta = -0.4; # price coefficient in utility function
sdxi = 0.15; # standard deviation of xi

S = 50; # number of simulations
G = 10; # size of grid on which to evaluate price elasticities
elast_own = zeros(S,G);
elast_cross = zeros(S,G);
elast_own_dist = zeros(S,T);
elastTrue = zeros(S,G);

s, pt, zt, xi = simulate_logit(J,T, beta, sdxi);

p_points = range(quantile(pt[:,1],.25),stop = quantile(pt[:,1],.75),length = G);
p_points = convert.(Float64, p_points)

# ------------------------------------------------------
# Set options for estimation and elasticity calculation
# ------------------------------------------------------
bernO = 2*ones(2J,1);        # Order of Bernstein Polynomial
iv=0;                       # Order of IV Polynomial = (bernO + iv)
trueS = false;                    # Evaluate at true market shares or not
own = [1,1];                # [derivative of j, with respect to price of k]

nfolds = 5; # number of cross-validation folds
nlam = 10; # number of regularization parameters to try. Actual values chosen automatically by hierNet
strong = true; # boolean for whether or not to impose strong hierarchy constraint
# Note: "strong = true" takes much longer than "strong = false."
nboot = 1; # number of bootstrapped samples to run hierarchical lasso on

# ------------------------------------------------------
# Simulation
# ------------------------------------------------------
included_symmetric_pct = zeros(2J,2J)
included_pct = zeros(2J,2J)
for si = 1:1:S
    # Simulate demand in two groups -- each product only substitutes to J-1 others
    s, pt, zt = simulate_logit(J,T, beta, sdxi);
    s2, pt2, zt2  = simulate_logit(J,T, beta, sdxi);

    s = [s s2];
    s = s ./ 2;
    pt = [pt pt2];
    zt = [zt zt2];
    df = toDataFrame(s,pt,zt);

    # hierNet() Returns two matrices: one, the "raw" selected model, and another
    #   which imposes symmetry. I.e. if j is a substute for k, then k is
    #   a substitute for j as well (this can drastically increase the # parameters
    #    to estimate when s has many columns)
    included, included_symmetric = hierNet_boot(df; nfolds = nfolds, nlam = 10, strong = false, nboot = 1);

    # Estimate demand nonparametrically
        # If you want to include an additional covariate in all demand
        # functions, add an additional argument "marketvars" after included. If it is an
        # additional product characteristic, marketvars should be T x J
    inv_sigma, designs = inverse_demand(df; included = included_symmetric);

    # Calculate price elasticities
    prices = median(pt).*ones(G,2J);
    prices[:,1] = p_points;
    elast, Jacobians, share_vec = price_elasticity(inv_sigma, df, prices; whichProducts = own,
        included = included_symmetric, trueS = trueS);

    trueelast = beta.*p_points.*(1 .- 2 .* share_vec[:,1])

    elast_own[si,:] = elast;
    elastTrue[si,:] = trueelast;
    included_pct[:,:] += included./S; # summarizing selection patterns. See pure_model_selection.jl
    included_symmetric_pct[:,:] += included_symmetric./S;
end

elast025 = zeros(G,1)
elast50 = zeros(G,1)
elast975 = zeros(G,1)
for i = 1:G
    elast025[i] = quantile(elast_own[:,i], 0.1)
    elast50[i] = quantile(elast_own[:,i], 0.5)
    elast975[i] = quantile(elast_own[:,i], 0.9)
end
elast025 = dropdims(elast025,dims = 2)
elast50 = dropdims(elast50,dims = 2)
elast975 = dropdims(elast975,dims = 2)

df = DataFrame(p = p_points, e025 = elast025, e50 = elast50, e975 = elast975)
ggplot(df, aes(x=:p, y=:e50)) + geom_line() + geom_line(aes(y=:e975), color = "gray", linetype = "dashed") +
    geom_line(aes(y=:e025), color = "gray", linetype = "dashed") +
    xlab("Price") + ylab("Own-Elasticity") + theme_light()


using Plots;
using Random;
using LinearAlgebra;
pyplot()

#===============================================#
#                                               #
#                Plots Settings                 #
#                                               #
#===============================================#

upscale = 2                             # 8x upscaling in resolution
fntsm = Plots.font("sans-serif", 10.0*upscale)
fntlg = Plots.font("sans-serif", 14.0*upscale)
default(titlefont=fntlg, guidefont=fntlg, tickfont=fntsm, legendfont=fntsm)
default(size=(800*upscale,600*upscale)) # Plot canvas size
default(dpi=300)                        # Only for PyPlot - presently broken

#===============================================#
#                                               #
#                  Definations                  #
#                                               #
#===============================================#

CORE_BIAS = 3
SAMPLE_SIZE = 50
FEATURE_DIM = 2
NUM_CLASSES = 2
λ = 0.1             # normalize
α = 0.001           # learning rate

#===============================================#
#                                               #
#                Tools functions                #
#                                               #
#===============================================#

function generate_dataset(sample_size::Int, feature_dim::Int=FEATURE_DIM, num_classes::Int=NUM_CLASSES)
    # random get
    Y = rand(1:num_classes,sample_size)
    # make features
    X = randn(sample_size,feature_dim) .+ (CORE_BIAS .* Y)
    # make label
    Y = Int.([y == class for y in Y, class = 1:num_classes])
    return hcat(ones(SAMPLE_SIZE,1), X) ,Y
end

# Gradient Descent
# π: predict(θ,X)
# ∇: gradient(π(θ,X),X,Y)
# loss: loss(π(θ,X),X,Y)
function gradient_descent(X,Y,π::Function,∇::Function,loss::Function,θ::Array,λ = 0)
    LR_DECAY = 0.1                  # lr_decay
    MOMENTUM = 0.9                  # Momentum
    i = 0

    P = π(θ,X)
    losses = Float64[loss(P,X,Y)]
    v = zeros(size(θ))
    while losses[end] > 0.1
        p = π(θ,X)
        Δ = ∇(p,X,Y) + λ * θ ./ length(Y)        # normalize, if λ = 0, equal to no normalization
        lr_i = α / (1 + LR_DECAY * i)
        v = MOMENTUM * v - lr_i * Δ
        θ += v
        push!(losses,loss(p,X,Y) + λ * sum(θ .* θ) ./ 2length(Y))
        println(losses[end])
    end
    return θ,losses
end

function gradient_descent(X,Y,π::Function,∇::Function,loss::Function,λ = 0)
    θ = ones(size(X,2),size(Y,2))
    return gradient_descent(X,Y,π,∇,loss,θ,λ)
end

function newtown(X,Y,π::Function,loss::Function,θ::Array,λ = 0)
    LR_DECAY = 0                      # lr_decay
    MOMENTUM = 0                      # Momentum
    i = 0

    P = π(θ,X)
    losses = Float64[loss(P,X,Y)]
    v = zeros(size(θ))
    while losses[end] > 0.1
        P = π(θ,X)
        A = Diagonal((P .* (1 .- P))[:,])
        U = X' * (Y .- P)
        H = X' * A * X
        lr_i = α / (1 + LR_DECAY * i)
        v = MOMENTUM * v + lr_i * (H^-1 * U - λ * θ ./ length(Y))
        θ = θ + v
        i = i + 1
        push!(losses,loss(P,X,Y) + λ * sum(θ .* θ) ./ 2length(Y))
        println(losses[end])
    end
    return θ,losses
end

function newtown(X,Y,π::Function,loss::Function,λ = 0)
    θ = ones(size(X,2),size(Y,2))
    return newtown(X,Y,π,loss,θ,λ)
end

function draw(θ,X,Y,losses,name::String)
    # draw the line and points
    x1 = Array(1:1:7)
    x2 = (- θ[1,1] .- θ[2,1] .* x1) ./ θ[3,1]
    colors = [[:red, :green][l+1] for l in Y[:,1]]
    scatter(X[:,2], X[:,3],  markercolor = colors)
    plot!(x1, x2)
    savefig("$name-plot")
    # losses plot
    plot(1:size(losses,1),losses[1:end])
    savefig("$name-loss")
end

#===============================================#
#                                               #
#               Logistic Regression             #
#                                               #
#===============================================#

# generate data
X,Y = generate_dataset(SAMPLE_SIZE,FEATURE_DIM,2)
Y = Y[:,1]

# class model
sigmoid(θ,X) = 1 ./ (1 .+ exp.(- X * θ))

# gradient
# ∇(sigmoid(θ,X),X,Y)
∇(π,X,Y) = X' * (π - Y) ./ size(Y,1)

# loss function
# loss(sigmoid(θ,X),X,Y)
loss(π,X,Y) = - sum(Y .* log.(π) + (1 .- Y) .* log.(1 .- π)) ./ length(Y)

θ, losses = gradient_descent(X,Y,sigmoid,∇,loss)
draw(θ,X,Y,losses,"grad")

θ, losses = newtown(X,Y,sigmoid,loss)
draw(θ,X,Y,losses,"newtown")

#===============================================#
#                                               #
#                    Test  UCI                  #
#                                               #
#===============================================#

using DataFrames, CSV;

function toLabel(labels)
    classes = unique(labels)
    function λ(label)
        for (index, value) in enumerate(classes)
            if label == value
                return index
            end
        end
    end
    return Int.([y == class for y in λ.(labels), class = 1:size(classes,1)]),classes
end

datasets = CSV.read("data/iris-cut.csv");
labels = Array{String}(copy(datasets[:Species]))

X = Matrix{Float64}(datasets[:,1:end-1])
X = hcat(ones(size(X,1),1), X)
Y,classes = toLabel(labels)
Y = Y[:,1]

θ, losses = gradient_descent(X,Y,sigmoid,∇,loss)
draw(θ,X,Y,losses,"iris")

#===============================================#
#                                               #
#                Softmax Regression             #
#                                               #
#===============================================#
using DataFrames, CSV;

function toLabel(labels)
    classes = unique(labels)
    function λ(label)
        for (index, value) in enumerate(classes)
            if label == value
                return index
            end
        end
    end
    return Int.([y == class for y in λ.(labels), class = 1:size(classes,1)]),classes
end

# class model
function softmax(θ,X)
    Y = exp.(X * θ)
    total = sum(Y,dims = 2)
    return Y ./ total
end

# gradient
# ∇(softmax(θ,X),X,Y)
∇(π,X,Y) = X' * (π - Y) ./ length(Y)

# loss function
# loss(softmax(θ,X),X,Y)
loss(π,X,Y) = - sum(Y .* log.(π)) ./ length(Y)

datasets = CSV.read("data/iris.csv");
labels = Array{String}(copy(datasets[:Species]))

# generate data
X = Matrix{Float64}(datasets[:,1:end-1])
# set bias row
X = hcat(ones(size(X,1),1), X)
Y,classes = toLabel(labels)

θ, losses = gradient_descent(X,Y,softmax,∇,loss)
θ, losses = gradient_descent(X,Y,softmax,∇,loss,θ)

# softmax
softmax(θ,X)

# losses plot
plot(1:size(losses,1),losses)


using Flux, Zygote, LinearAlgebra, Statistics
println("NNPDEHAN_tests")
using Test, StochasticDiffEq
println("Starting Soon!")
using NeuralPDE

using Random
Random.seed!(100)

# one-dimensional heat equation
x0 = [11.0f0]  # initial points
tspan = (0.0f0,5.0f0)
dt = 0.5   # time step
time_steps = div(tspan[2]-tspan[1],dt)
d = 1      # number of dimensions
m = 10     # number of trajectories (batch size)
batch_size = 8192
train_steps = 400
lr_boundaries = [100, 200, 300]
lr_values = [1e-1,1e-2,1e-3,1e-4]

g(X) = sum(X.^2)   # terminal condition
f(y,z,v_y,v_z,dv_y,dv_z,t) = sin(v_y) - v_z * π^(d/2)*σ_sampling^2 # function from solved equation
μ_f(X,p,t) = 0.0
σ_f(X,p,t) = 1.0
prob = TerminalPDEProblem(g, f, μ_f, σ_f, x0, tspan)

hls = d + 50 #hidden layer size
#neural network approximating solutions at the desired point
# with batch normalisation
# u0 = Flux.Chain(Dense(d,hls),
#                 BatchNorm(hls,tanh),
#                 Dense(hls,hls,tanh),
#                 BatchNorm(hls,tanh),
#                 Dense(hls,1))
# without batch normalisation
u0 = Flux.Chain(Dense(d,hls,tanh),
                Dense(hls,hls,tanh),
                Dense(hls,1))


struct Forcing{T<:Real}
    Fx::Array{T,2}
    Fy::Array{T,2}
    H::Array{T,2}
    η_ref::Array{T,2}
    Fη::Array{T,2}
    #sst_ref::Array{T,2}
    #sst_γ::Array{T,2}
end

function Forcing{T}(P::Parameter,G::Grid) where {T<:Real}

    @unpack wind_forcing_x,wind_forcing_y,topography = P

    if wind_forcing_x == "channel"
        Fx,_ = ChannelWind(T,P,G)
    elseif wind_forcing_x == "shear"
        Fx,_ = ShearWind(T,P,G)
    elseif wind_forcing_x == "double_gyre"
        Fx,_ = DoubleGyreWind(T,P,G)
    elseif wind_forcing_x == "constant"
        Fx,_ = ConstantWind(T,P,G)
    elseif wind_forcing_x == "none"
        Fx,_ = NoWind(T,P,G)
    end

    if wind_forcing_y == "channel"
        _,Fy = ChannelWind(T,P,G)
    elseif wind_forcing_y == "shear"
        _,Fy = ShearWind(T,P,G)
    elseif wind_forcing_y == "double_gyre"
        _,Fy = DoubleGyreWind(T,P,G)
    elseif wind_forcing_y == "constant"
        _,Fy = ConstantWind(T,P,G)
    elseif wind_forcing_x == "none"
        _,Fy = NoWind(T,P,G)
    end

    if topography == "ridge"
        H,_ = Ridge(T,P,G)
    elseif topography == "ridges"
        H = Ridges(T,P,G)
    elseif topography == "seamount"
        H = Seamount(T,P,G)
    elseif topography == "flat"
        H = FlatBottom(T,P,G)
    end

    η_ref = InterfaceRelaxation(T,P,G)
    Fη = KelvinPump(T,P,G)

    return Forcing{T}(Fx,Fy,H,η_ref,Fη)
end

"""Returns the constant forcing matrices Fx,Fy that vary only meriodionally/zonally
as a cosine with strongest forcing in the middle and vanishing forcing at boundaries."""
function ChannelWind(::Type{T},P::Parameter,G::Grid) where {T<:Real}

    @unpack Δ,x_u,y_u,Lx,Ly = G
    @unpack Fx0,Fy0,H,ρ,scale = P

    # for non-dimensional gradients the wind forcing needs to contain the grid spacing Δ
    xx_u,yy_u = meshgrid(x_u,y_u)
    Fx = (scale*Δ*Fx0/ρ/H)*cos.(π*(yy_u/Ly .- 1/2)).^2
    Fy = (scale*Δ*Fy0/ρ/H)*cos.(π*(xx_u/Lx .- 1/2)).^2

    return T.(Fx),T.(Fy)
end

"""Returns the constant forcing matrices Fx,Fy that vary only meriodionally/zonally
as a hyperbolic tangent with strongest shear in the middle."""
function ShearWind(::Type{T},P::Parameter,G::Grid) where {T<:Real}

    @unpack Δ,x_u,y_u,Lx,Ly = G
    @unpack Fx0,Fy0,H,ρ,scale = P

    # for non-dimensional gradients the wind forcing needs to contain the grid spacing Δ
    xx_u,yy_u = meshgrid(x_u,y_u)
    Fx = (scale*Δ*Fx0/ρ/H)*tanh.(2π*(yy_u/Ly .- 1/2))
    Fy = (scale*Δ*Fy0/ρ/H)*tanh.(2π*(xx_u/Lx .- 1/2))

    return T.(Fx),T.(Fy)
end

"""Returns the constant forcing matrices Fx,Fy that vary only meriodionally/zonally
with a superposition of sin & cos for a double gyre circulation.
See Cooper&Zanna 2015 or Kloewer et al 2018."""
function DoubleGyreWind(::Type{T},P::Parameter,G::Grid) where {T<:Real}

    @unpack Δ,x_u,y_u,Lx,Ly = G
    @unpack Fx0,Fy0,H,ρ,scale = P

    # for non-dimensional gradients the wind forcing needs to contain the grid spacing Δ
    xx_u,yy_u = meshgrid(x_u,y_u)
    Fx = (scale*Δ*Fx0/ρ/H)*(cos.(2π*(yy_u/Ly .- 1/2)) + 2*sin.(π*(yy_u/Ly .- 1/2)))
    Fy = (scale*Δ*Fy0/ρ/H)*(cos.(2π*(xx_u/Lx .- 1/2)) + 2*sin.(π*(xx_u/Lx .- 1/2)))
    return T.(Fx),T.(Fy)
end

"""Returns constant in in space forcing matrices Fx,Fy."""
function ConstantWind(::Type{T},P::Parameter,G::Grid) where {T<:Real}

    @unpack Δ,nux,nuy,nvx,nvy = G
    @unpack Fx0,Fy0,H,ρ,scale = P

    # for non-dimensional gradients the wind forcing needs to contain the grid spacing Δ
    Fx = T.(scale*Δ*Fx0/ρ/H)*ones(T,nux,nuy)
    Fy = T.(scale*Δ*Fy0/ρ/H)*ones(T,nvx,nvy)

    return Fx,Fy
end

"""Returns constant in in space forcing matrices Fx,Fy."""
function NoWind(::Type{T},P::Parameter,G::Grid) where {T<:Real}

    @unpack nux,nuy,nvx,nvy = G

    # for non-dimensional gradients the wind forcing needs to contain the grid spacing Δ
    Fx = zeros(T,nux,nuy)
    Fy = zeros(T,nvx,nvy)

    return Fx,Fy
end

"""Returns a reference state for Newtonian cooling/surface relaxation shaped as a
hyperbolic tangent to force the continuity equation."""
function InterfaceRelaxation(::Type{T},P::Parameter,G::Grid) where {T<:Real}

    @unpack x_T,y_T,Ly,Lx = G
    @unpack η_refh,η_refw = P

    # width of a tangent is defined as the distance between 65% of its minimum value and 65% of the max
    xx_T,yy_T = meshgrid(x_T,y_T)
    η_ref = -(η_refh/2)*tanh.(2π*(Ly/(4*η_refw))*(yy_T/Ly .- 1/2))
    return T.(η_ref)
end

"""Returns a matrix of water depth for the whole domain that contains a
Gaussian seamount in the middle. Water depth, heigth and width of the
seamount are adjusted with the constants H, topofeat_height and topofeat_width."""
function Seamount(::Type{T},P::Parameter,G::Grid) where {T<:Real}

    @unpack x_T_halo,y_T_halo,Lx,Ly = G
    @unpack topo_width,topo_height,H = P

    xx_T,yy_T = meshgrid(x_T_halo,y_T_halo)
    bumpx = exp.(-((xx_T .- Lx/2).^2)/(2*topo_width^2))
    bumpy = exp.(-((yy_T .- Ly/2).^2)/(2*topo_width^2))

    return T.(H .- topo_height*bumpx.*bumpy)
end

"""Returns a matrix of water depth for the whole domain that contains a
meridional Gaussian ridge in the middle. Water depth, heigth and width of the
ridge are adjusted with the constants water_depth, topofeat_height and topofeat_width."""
function Ridge(::Type{T},P::Parameter,G::Grid) where {T<:Real}

    @unpack x_T_halo,y_T_halo,Lx,Ly = G
    @unpack topo_width,topo_height,H = P

    xx_T,yy_T = meshgrid(x_T_halo,y_T_halo)
    bumpx = exp.(-((xx_T .- Lx/2).^2)/(2*topo_width^2))
    bumpy = exp.(-((yy_T .- Ly/2).^2)/(2*topo_width^2))

    Hx = H .- topo_height*bumpx
    Hy = H .- topo_height*bumpy
    return T.(Hx),T.(Hy)
end

"""Same as Ridge() but for n ridges at various x positions."""
function Ridges(::Type{T},P::Parameter,G::Grid) where {T<:Real}

    @unpack x_T_halo,y_T_halo,Lx,Ly = G
    @unpack topo_width,topo_height,H = P
    @unpack topo_ridges_positions = P
    n_ridges = length(topo_ridges_positions)

    xx_T,yy_T = meshgrid(x_T_halo,y_T_halo)

    # loop over bumps in x direction
    R = zero(xx_T) .+ H

    for i in 1:n_ridges
        R .-= topo_height*exp.(-((xx_T .- topo_ridges_positions[i]*Lx).^2)/(2*topo_width^2))
    end

    return T.(R)
end

"""Returns a matrix of constant water depth H."""
function FlatBottom(::Type{T},P::Parameter,G::Grid) where {T<:Real}
    @unpack nx,ny,haloη = G
    @unpack H = P
    return fill(T(H),(nx+2*haloη,ny+2*haloη))
end

"""Returns Kelvin wave pumping forcing of the continuity equation."""
function KelvinPump(::Type{T},P::Parameter,G::Grid) where {T<:Real}

    @unpack x_T,y_T,Lx,Ly = G
    @unpack R,ϕ,Δ = G
    @unpack A,ϕk,wk = P

    xx_T,yy_T = meshgrid(x_T,y_T)

    mϕ = 2π*R/360.          # meters per degree latitude
    y0 = Ly/2 - (ϕ-ϕk)*mϕ   # y[m] for central latitude of pumping

    Fη = A*Δ*exp.(-(yy_T.-y0).^2/(2*wk^2))

    return T.(Fη)
end

"""Time evolution of forcing."""
function Ftime(::Type{T},t::Int,ω::Real) where {T<:Real}
    return convert(T,sin(ω*t))
end


using Interpolations
using QuadGK

function calculate_Vdl_interp(dos_f, Vq_min, Vq_max, C_dl)
    Vq_range = range(Vq_min, Vq_max, step=0.001)
    E_min = dos_f.ranges[1][1]
    E_max = dos_f.ranges[1][end]
    CQ = [compute_cq(E_min, E_max, V, dos_f)*1.602*10 for V in Vq_range]
    Vdl_data = []
    Vappl_data = []
    for i = 1:length(Vq_range)
        V_app_i =  Vq_range[i]*((CQ[i]/C_dl) + 1)
        V_dl_i = V_app_i - Vq_range[i]
        push!(Vdl_data, V_dl_i)
        push!(Vappl_data, V_app_i)
    end

    ## sort data and do V_dl interpolation
    inds = sortperm(Vappl_data)
    v_interp = LinearInterpolation(Vappl_data[inds], Vdl_data[inds])
    return v_interp
end

function QC_integrand(E, eVq, dos_func; kT=.026)
    dos_func(E) * sech((E-eVq)/(2*kT))^2/(4*kT)
end

function compute_cq(E_min, E_max, eVq, dos_func; kT=.026)
    fcn = E->QC_integrand(E, eVq, dos_func; kT=kT)
    quadgk(fcn, E_min, E_max)[1]
end


# Test for gausslegendre().

    tol = 1e-14
    n = 42
    x = gausslegendre(n)
    @test length(x[1]) == n && length(x[2]) == n
    @test abs(x[1][37] - 0.910959724904127) < tol
    @test abs(x[2][37] - 0.030479240699603) < tol

    # Test a larger n (using ASY)
    n = 251
    x = gausslegendre(n)
    w = x[2]; x = x[1]
    @test all(length(x) == n) && all(length(w) == n)
    @test abs(x[37] + 0.896467746955729) < tol
    @test abs(w[37] - 0.005535005742012) < tol

    x = gausslegendre(1013);
    @test norm( x[1][2] - -0.999985167586110, Inf ) < tol
    @test norm( x[1][13] - -0.999218995240887, Inf ) < tol
    @test norm( x[2][2] - 1.681691163200592e-05, Inf ) < tol
    @test norm( x[2][13] - 1.224755309137936e-04, Inf ) < tol

    x = gausslegendre(10013);
    @test norm( x[1][2] - -0.999999848054223, Inf ) < tol
    @test norm( x[1][13] - -0.999991998242661, Inf ) < tol
    @test norm( x[2][2] - 1.722757320118474e-07, Inf ) < tol
    @test norm( x[2][13] - 1.254980540032470e-06, Inf ) < tol


# This file is part of the Julia package ModularForms.jl
#
# Copyright (c) 2018-2019: Lynn Engelberts and Alexandru Ghitza.

"""
    prime_range(n)

Return an array consisting of all primes up to and including `n`.
"""
function prime_range(n::Int)

   primes = fill(true, n)
   if n == 0
      return primes
   end

   primes[1] = false
   for p in 2:n
      primes[p] || continue
      for j in 2:div(n, p)
         primes[p*j] = false
       end
   end

   return findall(primes)
end


using Aqua
using Test
using NCDatasets
using EmpiricalOrthogonalFunctions

testdir = @__DIR__
datadir = joinpath(testdir, "data")

ds = NCDataset(joinpath(datadir,"sst.mnmean.1990-present.nc"))

@testset "EmpiricalOrthogonalFunctions" begin

    sst = ds["sst"][:]
    dsin = copy(sst)

    # @testset "EOF Construction" begin
        eof = EmpiricalOrthogonalFunction(dsin)

        dsin = permutedims(sst,(3,1,2))
        eof = EmpiricalOrthogonalFunction(dsin; timedim=1)

        dsin = permutedims(sst,(1,3,2))
        eof = EmpiricalOrthogonalFunction(dsin; timedim=2)
    # end

    # @testset "EOF Analysis" begin

        # @testset "eigenvals" begin
            @test totalanomalyvar(eof) ≈ 2673.0002f0 atol=1e-3
            @test eigenvalues(eof,n=5)[1] ≈ 1332.6342
            @test variancefraction(eof, n=100)[1] ≈ 0.4985537f0 atol=1e-3
        # end

        # @testset "pcs" begin
            @test pcs(eof,n=1)[1,1] ≈ -10.937908 atol=1e-3
            @test pcs(eof,scaling=:divide)[1,1] ≈ -0.299626 atol=1e-3
            @test pcs(eof,scaling=:multiply, n=10)[1,1] ≈ -399.291
        # end

        # @testset "eofs" begin
            @test eofs(eof,n=5)[1,:] ≈ [0.0279191, -0.00658843, 0.02273, -0.00208263, 0.0347027] atol=1e-3
            @test eofs(eof,n=2,scaling=:divide)[1,:] ≈ [0.000764796, -0.000231014] atol=1e-3
            @test eofs(eof,n=3,scaling=:multiply)[1,:] ≈ [ 1.01919, -0.1879, 0.293734] atol=1e-3
        # end

        # @testset "projection" begin
            @test all((reshape(reconstruct(eof), size(sst)) ≈ sst)...)
            @test projectfield(eof, collect(eof.dataset))[1,1] ≈ 230.678 atol=1e-3
        # end

        # @testset "rotation" begin
            reof = orthorotation(eof,n=4)

            @test eofs(reof)[1,:] ≈ [0.02492, -0.00625398, 0.0254053, 0.00619081] atol=1e-3

        # end
    # end
end

close(ds)

Aqua.test_all(EmpiricalOrthogonalFunctions)

# end tests


function basispursuit(A::AbstractMatrix, b::AbstractVector)
    basispursuit(A, b, ones(eltype(A), size(A, 2)))
end
const bp = basispursuit

function basispursuit(A::AbstractMatrix, b::AbstractVector, w::AbstractVector)
    model = JuMP.Model(optimizer_with_attributes(Clp.Optimizer, "LogLevel" => 0))
    n, m = size(A)
    @variable(model, x⁺[1:m] ≥ 0)
    @variable(model, x⁻[1:m] ≥ 0)
    @objective(model, Min, dot(w, x⁺) + dot(w, x⁻))
    @constraint(model, constraint, A * (x⁺ .- x⁻) .== b)
    JuMP.optimize!(model)
    x = @. JuMP.value(x⁺) - JuMP.value.(x⁻)
    sparse(x)
end

function basispursuit_reweighting(A::AbstractMatrix, b::AbstractVector,
                                    reweighting!; maxiter = 8, min_decrease = 1e-8)
    x = basispursuit(A, b)
    w = ones(length(x))
    for i in 2:maxiter
        reweighting!(w, x) # updates w
        xs = basispursuit(A, b, w)
        if norm(xs-x) < min_decrease # stop if we don't make sufficient progress
            return xs
        end
        x .= xs
    end
    x
end

candes_weight(x, ε) = inv(abs(x) + ε)
function candes_weights!(w, x, ε::Real)
    @. w = candes_weight(x, ε)
    if any(isnan, w) || any(isinf, w)
        throw("weights contain NaN or Inf: $w")
    end
end
candes_function(ε::Real) = (w, x) -> candes_weights!(w, x, ε)
function candes_reweighting(A::AbstractMatrix, b::AbstractVector,
                                            ε = 1e-2; maxiter = 8)
    basispursuit_reweighting(A, b, candes_function(ε), maxiter = maxiter)
end
const bp_candes = candes_reweighting

# using LinearAlgebraExtensions: LowRank
# TODO: stop when w has converged, instead of maxiter
function ard_weights!(w, A, x, ε::Real, iter::Int = 8)
    if any(==(0), w)
        error("weights cannot be zero")
    end
    # nzind = x.nzind # this could speed up the weight computation
    # Ai = @view A[:, nzind]
    # wx = @. abs(x.nzval) / w[nzind]
    # K = Woodbury(ε*I(size(A, 1)), Ai, WX, Ai')
    for i in 1:iter
        wx = @. abs(x) / w
        WX = Diagonal(wx)
        K = Woodbury(ε*I(size(A, 1)), A, WX, A')
        K⁻¹ = inverse(factorize(K))
        dK(a) = sqrt(max(dot(a, K⁻¹, a), 0))
        @. w = dK($eachcol(A))
    end
end

ard_function(A::AbstractMatrix, ε::Real) = (w, x) -> ard_weights!(w, A, x, ε)

function ard_reweighting(A::AbstractMatrix, b::AbstractVector, ε::Real = 1e-2;
                                                            maxiter::Int = 8)
    basispursuit_reweighting(A, b, ard_function(A, ε), maxiter = maxiter)
end
const bp_ard = ard_reweighting

############################# Basis Pursuit Denoising ##########################
# noisy variant, where we allow δ-deviations from data
# minimizes l2 norm constrained linear objective (conic program)
# min_x |x|_1 s.t. |Ax-b|_2 < δ
function basis_pursuit_denoising(A::AbstractMatrix, b::AbstractVector, δ::Real,
                                    w::AbstractVector = ones(eltype(A), size(A, 2)))

    model = bpd_ecos_model()
    n, m = size(A)
    @variable(model, x⁺[1:m] ≥ 0)
    @variable(model, x⁻[1:m] ≥ 0)
    @objective(model, Min, dot(w, x⁺) + dot(w, x⁻))
    @constraint(model, [δ; A * (x⁺ .- x⁻) .- b] in SecondOrderCone())
    # return model
    JuMP.optimize!(model)
    try
        x = @. JuMP.value(x⁺) - JuMP.value(x⁻)
        return sparse(x)
    catch e
        println(e)
        x = fill(NaN, size(A, 2))
        return sparse(x)
    end
end
const bpd = basis_pursuit_denoising

function bpd_reweighting(A::AbstractMatrix, b::AbstractVector, δ::Real,
                                    reweighting!; maxiter = 8, min_decrease = 1e-4)
    x = bpd(A, b, δ)
    w = ones(length(x))
    for i in 2:maxiter
        reweighting!(w, x) # updates w
        xs = bpd(A, b, δ, w)
        if norm(xs-x) < min_decrease # stop if we don't make sufficient progress
            return xs
        end
        x .= xs
    end
    x
end

# δ is l2 noise threshold, ε is reweighting constant
# ε below 1e-3 can lead to convergence problems
function bpd_candes(A::AbstractMatrix, b::AbstractVector, δ::Real, ε::Real = δ; maxiter = 8)
    bpd_reweighting(A, b, δ, candes_function(ε), maxiter = maxiter)
end
function bpd_ard(A::AbstractMatrix, b::AbstractVector, δ::Real, ε::Real = δ^2; maxiter = 8)
    bpd_reweighting(A, b, δ, ard_function(A, ε), maxiter = maxiter)
end

# optimizer setup
function bpd_scs_model()
    JuMP.Model(optimizer_with_attributes(SCS.Optimizer,
                "verbose" => true,
                "eps" => 1e-6,
                "alpha" => 1.,
                "cg_rate" => 2.,
                "max_iters" => 5000))
end
function bpd_ecos_model()
    JuMP.Model(optimizer_with_attributes(ECOS.Optimizer, "verbose" => false))
end

################################ (F)ISTA ########################################
struct FISTA end
# TODO: use homotopy to select appropriate λ

# shrinkage and thresholding operator
shrinkage(x::Real, α::Real) = sign(x) * max(abs(x)-α, zero(x))

# TODO fast iterative shrinkage and hard-thresholding algorithm
# for weighted l1-norm minimization
function fista(A::AbstractMatrix, b::AbstractVector, λ::Real,
    x::AbstractVector = spzeros(size(A, 2)); maxiter::Int = 1024, stepsize::Real = 1e-2)
    w = fill(λ, size(x))
    fista(A, b, w, x, maxiter = maxiter, stepsize = stepsize)
end

function l1(x::AbstractVector, w::AbstractVector)
    length(x) == length(w) || throw(DimensionMismatch("length(x) ≠ length(w)"))
    n = zero(eltype(x))
    @simd for i in eachindex(x)
        @inbounds n += w[i] * abs(x[i])
    end
    return n
end

# TODO: stepsize selection
function ista(A, b, λ::Real, x = spzeros(size(A, 2)); maxiter::Int = 1024, stepsize::Real = 1e-2)
    ista(A, b, fill(λ, size(x)), x, maxiter = maxiter, stepsize = stepsize)
end

function ista(A::AbstractMatrix, b::AbstractVector, w::AbstractVector,
                                        x::AbstractVector = spzeros(size(A, 2));
                                        maxiter::Int = 1024, stepsize::Real = 1e-2)
    x = sparse(x)
    r(x) = b-A*x # residual
    f(x) = sum(abs2, r(x)) + l1(x, w)
    g(x) = A'r(x) # negative gradient
    α = stepsize
    fx = f(x)
    for i in 1:maxiter
        ∇ = g(x)
        @. x = shrinkage(x + 2α*∇, w*α)
        dropzeros!(x) # optimize sparse representation
    end
    return x
end

function fista(A::AbstractMatrix, b::AbstractVector, w::AbstractVector,
                                        x::AbstractVector = spzeros(size(A, 2));
                                        maxiter::Int = 1024, stepsize::Real = 1e-2)
    x = sparse(x)
    r(x) = b-A*x # residual
    f(x) = sum(abs2, r(x)) + l1(x, w)
    g(x) = A'r(x) # negative gradient
    α = stepsize
    tk = 1
    for i in 1:maxiter
        ∇ = g(x)
        tkn = (1 + sqrt(1 + 4tk^2)) / 2
        y = xkn + (tk - 1) / tkn * (xkn - xk)
        tk = tkn
        @. x = shrinkage(x + 2*∇, w*α)
        dropzeros!(x) # optimize sparse representation
    end
    return x
end


@testset ExtendedTestSet "squared_euclidean" begin
    @testset ExtendedTestSet "squared_euclidean" begin
        x = boolean_indicator([0, 1, 0, 0, 0, 1, 1, 1, 1, 1, 0])
        tfm = SquaredEuclidean(x)
        answer = [1, 0, 1, 4, 1, 0, 0, 0, 0, 0, 1]
        @test transform(x, tfm) == answer
    end

    @testset ExtendedTestSet "squared_euclidean" begin
        x = boolean_indicator(
            [
                0 1 1 1 0
                1 1 1 1 1
                1 0 0 0 1
                1 0 0 0 1
                1 0 0 0 1
                1 1 1 1 1
                0 1 1 1 0
            ]
        )
        tfm = SquaredEuclidean(x)
        answer = [
            1.0 0.0 0.0 0.0 1.0
            0.0 0.0 0.0 0.0 0.0
            0.0 1.0 1.0 1.0 0.0
            0.0 4.0 4.0 4.0 0.0
            0.0 1.0 1.0 1.0 0.0
            0.0 0.0 0.0 0.0 0.0
            1.0 0.0 0.0 0.0 1.0
        ]
        @test transform(x, tfm) == answer
    end

    @testset ExtendedTestSet "squared_euclidean" begin
        x = boolean_indicator(
            [
                0 1 1 1 0
                1 1 1 1 1
                1 0 0 0 1
                1 0 0 0 1
                1 0 0 0 1
                1 1 1 1 1
                0 1 1 1 0
            ]
        )
        tfm = SquaredEuclidean(x)
        answer = [
            1.0 0.0 0.0 0.0 1.0
            0.0 0.0 0.0 0.0 0.0
            0.0 1.0 1.0 1.0 0.0
            0.0 4.0 4.0 4.0 0.0
            0.0 1.0 1.0 1.0 0.0
            0.0 0.0 0.0 0.0 0.0
            1.0 0.0 0.0 0.0 1.0
        ]
        @test transform!(x, tfm, Threads.nthreads()) == answer
    end
end


# This file is pat of RLinearAlgebra.jl
# 1. Specifies type
# 2. Implements rsubsolve function
# 3. Exports Type

#suing LinearAlgebra

"""
    LinSysVecColProjFO <: LinSysVecColProjection

A mutable structure that represents the standard column projection method
with full orthogonalization against all previous projections.

See Patel, Vivak, Mohammad Jahangoshahi, and Daniel Adrian Maldonado. "An Implicit
    Representation and Iterative Solution of Randomly Sketched Linear Systems."
    SIAM Journal on Matrix Analysis and Applications (2021) 42:2, 800-831.
    https://doi.org/10.1137/19M1259481

# Fields

- `S::Union{Matrix{Float64}, Nothing}` is a matrix used for orthogonalizing against
    all previous search directions.

Calling `LinSysVecColProjFO()` defaults to `LinSysVecColProjFO(nothing)`.
"""
mutable struct LinSysVecColProjFO <: LinSysVecColProjection
    S::Union{Matrix{Float64}, Nothing}
end
LinSysVecColProjFO() = LinSysVecColProjFO(nothing)

# Common rsubsolve interface for linear systems
function rsubsolve!(
    type::LinSysVecColProjFO,
    x::AbstractVector,
    samp::Tuple{U,V,W} where {U<:AbstractVector,V<:AbstractArray,W<:Real},
    iter::Int64,
)
    # samp[1] is the search direction
    # samp[2] is the coefficient matrix A
    # samp[3] is the residual of normal system in A * samp[1], samp[1]' A' * (A * x - b)

    if iter == 1
        d = length(x)
        type.S = diagm(ones(Float64, d))
    end

    # Compute orthogonal component of A * samp[1] using S (projection matrix)
    q = samp[2]' * (samp[2] * samp[1])
    u = type.S * q

    if dot(u,u) < eps()^1.5 * length(x) #1.5 instead of 2 since normal system
        return nothing
    end

    # Update Iterate
    γ = dot(u, q)
    x .= x - u * (samp[3] / γ)

    # Update Projection Matrix
    type.S .= (I - ( u / γ) * q') * type.S

    return nothing
end


"""
    get_proportion_2d(for_prop_arr::AbstractArray{Float64, 2})::AbstractArray{Float64, 2}
Compute the proportional values for each element in the 2d array.
returns a AbstractArray{Float64, 2}
# Example
```
julia> get_proportion_2d([1.0 2.0 ; 4.0 5.0])
=[0.08333333333333333 0.16666666666666666; 0.3333333333333333 0.4166666666666667]
```
"""
function get_proportion_2d(for_prop_arr::AbstractArray{Float64,2})::AbstractArray{Float64,2}
    return for_prop_arr ./ sum(for_prop_arr)
end

"""
    get_proportion_round_2d(for_prop_arr::AbstractArray{Float64, 2};round_digits::Int64)::AbstractArray{Float64, 2}
Compute the proportional values for each element in the 2d array and round to a specified decimal digits.
If the sum of all the elements is not 1 , then the adjustment is done in the 1st largest element
returns a AbstractArray{Float64, 1}
# Example
```
julia> get_proportion_round_2d([1.1101 2.243 ; 4.9898 5.87],round_digits=5)
=[0.07811 0.15781; 0.35108 0.413]
```
"""
function get_proportion_round_2d(for_prop_arr::AbstractArray{Float64,2};round_digits::Int64 = 2)::AbstractArray{Float64,2}
    temparr=  round.(for_prop_arr ./ sum(for_prop_arr), digits=round_digits, base = 10)
    diff = 1 - sum(temparr) #delta naming
    if diff != 0.0
        max = argmax(temparr)
        temparr[max] =  round(temparr[max] + diff, digits=round_digits, base = 10)
    end
    return temparr
end

"""
    get_proportion_round_add1251(for_prop_arr::AbstractArray{Float64, 2};round_digits::Int64)::AbstractArray{Float64, 2}
Compute the proportional values for each element in the 2d array and round to a specified decimal digits.
If the sum of all the elements is not 1 , then add 1251 to the array and try looking for proportional correctness
after 2 itterations , the adjustment is done in the 1st largest element
returns a AbstractArray{Float64, 2}
# Example
```
julia> get_proportion_round_add1251_2d([1.1101 2.243 ; 4.9898 5.87 ; 34.333 6.78],round_digits=8)
=[0.16559399 0.16574381; 0.16610708 0.16622349; 0.16998779 0.16634384]
```
"""
function get_proportion_round_add1251_2d(for_prop_arr::AbstractArray{Float64,2};round_digits::Int64 = 2)::AbstractArray{Float64,2}
    temparr=  round.(for_prop_arr ./ sum(for_prop_arr), digits=round_digits, base = 10)
    diff = 1 - sum(temparr)
    if diff != 0.0
        check = 1
        while check < 2
            sqarr= (for_prop_arr) .+1251
            temparr=  round.(sqarr ./ sum(sqarr), digits=round_digits, base = 10)
            diff = 1 - sum(temparr)
            if diff != 0.00
                check = check + 1
                if check == 2
                    max = argmax(temparr)
                    temparr[max] = round(temparr[max]+diff,digits=round_digits, base = 10)
                end
            else
                break
            end
        end
    end
    return temparr
end


using Pidoh
using Random

@testset "Flip bit string" begin
    @test flip(true) == false
    @test flip(false) == true
end

@testset "OneMax Fitness Function" begin
    @test fitness(BitArray([true, true, true]), OneMax(3)) == 3
    @test fitness(BitArray([false, false, false, false]), OneMax(4)) == 0
    @test fitness(BitArray([]), OneMax(0)) == 0
end

@testset "ZeroMax Fitness Function" begin
    @test fitness(BitArray([true, true, true]), ZeroMax(3)) == 0
    @test fitness(BitArray([false, false, false, false]), ZeroMax(4)) == 4
    @test fitness(BitArray([]), ZeroMax(0)) == 0
end

jump = Jump(10, 3)
@testset "ZeroMax Fitness Function" begin
    @test fitness(BitArray([0, 0, 0, 0, 0, 0, 0, 0, 0, 0]), jump) == jump.jumpsize
    @test fitness(BitArray([1, 0, 0, 0, 0, 0, 0, 0, 0, 0]), jump) == jump.jumpsize + 1
    @test fitness(BitArray([1, 1, 1, 1, 1, 1, 1, 1, 1, 1]), jump) == jump.jumpsize + 10
    @test fitness(BitArray([1, 1, 1, 1, 1, 1, 1, 1, 1, 0]), jump) == 1
    @test fitness(BitArray([1, 1, 1, 1, 1, 1, 1, 0, 0, 0]), jump) == 10
    @test fitness(BitArray([1, 1, 1, 1, 1, 1, 1, 1, 0, 0]), jump) == 2

end

ins1 = Instance(BitArray([true, true, false]), OneMax(3))

@testset "UniformlyIndependentMutation mutation" begin
    @test mutation(ins1, UniformlyIndependentMutation(0.0)).individual == ins1.individual
    @test mutation(ins1, UniformlyIndependentMutation(1.0)).individual != ins1.individual
    @test fitness(mutation(ins1, UniformlyIndependentMutation(1.0))) == 1
end


export DegreeIndependentSet

struct DegreeIndependentSet end

"""
    independent_set(g, DegreeIndependentSet())

Obtain an [independent set](https://en.wikipedia.org/wiki/Independent_set_(graph_theory)) of
`g` using a greedy heuristic based on the degree of the vertices.

### Implementation Notes
A vertex is said to be valid if it is not in the independent set or adjacent to any vertex
in the independent set.
Initilalise the independent set to an empty set and iteratively choose the vertex that is
adjacent to the fewest valid vertices in the independent set until all vertices are invalid.

### Performance
Runtime: O((|V|+|E|)*log(|V|))
Memory: O(|V|)
"""
function independent_set(
    g::AbstractGraph{T},
    alg::DegreeIndependentSet
    ) where T <: Integer

    nvg = nv(g)
    ind_set = Vector{T}()
    sizehint!(ind_set, nvg)
    deleted = falses(nvg)
    degree_queue = PriorityQueue(enumerate(degree(g)))

    while !isempty(degree_queue)
        v = dequeue!(degree_queue)
        deleted[v] && continue
        deleted[v] = true
        push!(ind_set, v)

        for u in neighbors(g, v)
            deleted[u] && continue
            deleted[u] = true
            @inbounds @simd for w in neighbors(g, u)
                if !deleted[w]
                    degree_queue[w] -= 1
                end
            end
        end
    end
    return ind_set
end


using RandomMatrices
using Random: seed!
using Test

@testset "RandomMatrices" begin
seed!(1)
include("GaussianEnsembles.jl")
include("FormalPowerSeries.jl")
include("Haar.jl")
include("StochasticProcess.jl")

@testset "densities" begin
    include("densities/Semicircle.jl")
    include("densities/TracyWidom.jl")
end
end