using Distributions
using PyPlot
using RDatasets
using MAT

include("./cluster_eval.jl")

srand(0) #seed

@everywhere function kpp_init(X::Array{Float64,2}, K::Int64)

    n, d = size(X)
    mu = zeros(K,d)
    dist = Inf * ones(n)
    Lambda = 1

    idx = round(Int64, rand(Uniform(1,n)))
    mu[1,:] = X[idx,:]
    for i = 2:K
        D = X - repmat(mu[i-1,:]',n,1)
        dist = min(dist, sum(D.*D,2))
        idx_cond = rand(Uniform(0,1)) .< cumsum(dist/sum(dist))
        idx_mu = find(idx_cond)[1]
        mu[i,:] = X[idx_mu,:]
        Lambda = maximum(dist)
    end
    Lambda
end

@everywhere function dpmeans_fit(X::Array{Float64, 2})

    #init params
    K = 1
    K_init = 9
    d = size(X,2)
    z = mod(randperm(size(X,1)),K)+1
    mu = rand(MvNormal(zeros(d), 10*eye(d)))
    sigma = 1
    nk = zeros(K)
    pik = ones(K)/K

    mu = mean(X,1)
    Lambda = kpp_init(X, K_init)
    println("lambda: ", Lambda)

    max_iter = 100
    obj = zeros(max_iter)
    em_time = zeros(max_iter)

    obj_tol = 1e-3
    n, d = size(X)

    for iter = 1:max_iter
        dist = zeros(n,K)

        #assignment step
        for kk = 1:K
            Xm = X - repmat(mu[kk,:]',n,1)
            dist[:,kk] = sum(Xm.*Xm,2)
        end

        #update labels
        dmin = minimum(dist,2)
        z = mapslices(indmin, dist, 2)
        #find(dist->dist==0, dist)
        idx = find(dmin->dmin > Lambda, dmin)
        if (size(idx)[1] > 0)
            K = K + 1
            z[idx] = K
            mu = cat(1, mu, mean(X[idx,:],1))
            Xm = X - repmat(mu[K,:]',n,1)
            dist = cat(2, dist, sum(Xm.*Xm,2))
        end

        #update step
        nk = zeros(K)
        for kk = 1:K
            nk[kk] = counts(z, kk)[1]
            idx2 = find(z->z==kk, z)
            if (length(idx2)>0)
                mu[kk,:] = mean(X[idx2,:],1)
            end
        end
        pik = nk / sum(nk)

        #compute objective
        for kk = 1:K
            idx = find(z->z==kk, z)
            obj[iter] = obj[iter] + sum(dist[idx,kk])
        end
        obj[iter] = obj[iter] + Lambda * K

        #check convergence
        if (iter > 1 && abs(obj[iter] - obj[iter-1]) < obj_tol * obj[iter])
            println("converged in ", iter, " iterations.")
            break
        end
        em_time[iter] = 0

    end

    z, obj, em_time
end

#generate data (X is NxD)
load_data = "synthetic"  #"synthetic", "iris", "wine", "soybean", "pima", "car"

if load_data == "synthetic"
    println("synthetic dataset")
    file = matread("./data/clusters.mat")
    X = transpose(file["data"])
    z_true = round(Int64, file["labels"])
elseif load_data == "iris"
    println("iris dataset")
    iris = dataset("datasets","iris")
    X = convert(Array, iris[:,1:4])
    species = convert(Array, iris[:,5])
    z_true = zeros(length(species))
    cnt = 1
    for sp in unique(species)
        idx = find(species->species == sp, species)
        z_true[idx] = cnt
        cnt = cnt + 1
    end
    z_true = round(Int64, z_true)
elseif load_data == "wine"
    println("wine dataset")
    file = matread("./data/wine.mat")
    X = 1.0 .* file["wine"]
    z_true = round(Int64, file["label"])
    z_true = squeeze(z_true, 2)
elseif load_data == "soybean"
    println("soybean dataset")
    file = matread("./data/soybean.mat")
    X = file["data"]
    z_true = round(Int64, file["label"])
    z_true = squeeze(z_true, 2)
elseif load_data == "pima"
    println("pima dataset")
    file = matread("./data/pima.mat")
    X = file["data"]
    z_true = round(Int64, file["label"])
    z_true = squeeze(z_true, 2)
elseif load_data == "car"
    println("car dataset")
    file = matread("./data/car.mat")
    X = file["data"]
    z_true = round(Int64, file["label"])
    z_true = squeeze(z_true, 2)
else
    println("choose a dataset")
    quit()
end

#DP means
tic()
labels, obj, em_time = dpmeans_fit(X)
ftime = toc()

#evaluate clustering
labels_ord = zeros(length(labels))
cnt = 1
for kk in unique(labels)
    cidx = find(labels->labels == kk, labels)
    labels_ord[cidx] = cnt
    cnt = cnt + 1
end
labels_ord = round(Int64, labels_ord)
println("K true: ", length(unique(z_true)))
println("K dpmeans: ", length(unique(labels_ord)))
cluster_eval(labels_ord, z_true)

#generate plots
figure()
scatter(X[:,1], X[:,2], c=z_true)
title("ground truth clustering")

figure()
scatter(X[:,1], X[:,2], c=labels_ord)
title("dp-means clustering")


using Pkg
Pkg.activate(".")

using Roots
using Dierckx
using PyPlot

function find_c(β)
    -log(1-β) / β
end
β = 0.01:0.005:0.99
plot(find_c.(β), β, linewidth=6)

function find_β(c)
    function f(β)
        β + exp(-c*β) - 1.0
    end
    find_zero(f, (0.00001, 1))
end

c = 1.005:0.005:4.65
plot(c, find_β.(c))

spl_c = Spline1D(c, find_β.(c))
spl_β = Spline1D(find_c.(β), β)
plot(c, spl_c.(c) - spl_β.(c))


"""
cgs_efield!(incident_electric_field, fft_tmp, G, vector_potential_3d, dielectric, vector_potential_1d, AToE, Ig, pfft, pifft, x, p, r, rt, u, v, q, uq; tol = 1e-15, maxit = 18)\n
Inplace version for solving the VIE. the result will be in x.\n
"""
function cgs_efield!(
    incident_electric_field,
    fft_tmp,
    G,
    vector_potential_3d,
    dielectric,
    vector_potential_1d,
    AToE,
    Ig,
    pfft,
    pifft,
    x,
    p,
    r,
    rt,
    u,
    v,
    q,
    uq;
    tol = 1e-15,
    maxit = 18
)
    T = ComplexF32
    n2b = norm(incident_electric_field)
    n2b = n2b*tol
    #copyto!(dst,src)
    copyto!(x, incident_electric_field)
    _total_minus_scattered_electric_field!(p, x, vector_potential_1d, vector_potential_3d, G, dielectric, Ig, AToE, fft_tmp, pfft, pifft, -one(T))
    # p = b - Ax
    axpy!(one(T), incident_electric_field, p)
    # r= b - Ax
    copyto!(r, p)
    # u = b - Ax
    copyto!(u, p)
    # v = A(b - Ax)
    _total_minus_scattered_electric_field!(v, p, vector_potential_1d, vector_potential_3d, G, dielectric, Ig, AToE, fft_tmp, pfft, pifft, one(T))
    copyto!(rt, r)
    ρ = dot(rt, r)
    for _ in 1:maxit
        σ = dot(rt, v)
        α = ρ/σ
        # q = u-α*v
        caxpy!(q, -α, v, u)
        # uq = u + q
        caxpy!(uq, one(T), u, q)
        # x = x + α*uq
        axpy!(α, uq, x)

        _total_minus_scattered_electric_field_for_residual!(r, uq, vector_potential_1d, vector_potential_3d, G, dielectric, Ig, AToE, fft_tmp, pfft, pifft, -α)

        if norm(r) <= n2b
          break #return xx,xy,xz
        end

        ρold = ρ
        ρ = dot(rt, r)
        β = ρ/ρold
        # u = r + β*q
        caxpy!(u, β, q, r)
        #p = u + β*(q+β*p)
        axpby!(one(T), q, β, p)
        axpby!(one(T), u, β, p)
        _total_minus_scattered_electric_field!(v, p, vector_potential_1d, vector_potential_3d, G, dielectric, Ig, AToE, fft_tmp, pfft, pifft, one(T))
    end

    return nothing
end


function roll_non_degenerate_tet()
    v1 = randn(SVector{3,Float64})
    v2 = randn(SVector{3,Float64})
    v3 = randn(SVector{3,Float64})
    v4 = randn(SVector{3,Float64})
    if volume(v1, v2, v3, v4) < 0.25
        return roll_non_degenerate_tet()
    else
        return v1, v2, v3, v4
    end
end

function make_4_sided()
    v1, v2, v3, v4 = roll_non_degenerate_tet()
    (0.0 < volume(v1, v2, v3, v4)) || error("bad tet")
    A = asMatOnePad(SVector{4,SVector{3,Float64}}(v1, v2, v3, v4))
    inv_A = inv(A)
    n̂ = normalize(randn(SVector{3,Float64}))
    plane = SMatrix{1,4,Float64,4}(n̂[1], n̂[2], n̂[3], randn())
    c = clip_plane_tet(plane, A)
    if length(c) == 4
        return c
    else
        return make_4_sided()
    end
end

dist_from_plane(plane::SMatrix{1,4,Float64,4}, v::SVector{3,Float64}) = dot(v, unPad(plane)) + plane[4]
project_into_plane(plane::SMatrix{1,4,Float64,4}, v::SVector{3,Float64}) = v - dist_from_plane(plane, v) * unPad(plane)


# Operations on 2D and 3D vectors, and colors

const Vec3{T<:AbstractFloat} = SVector{3, T}

# Adapted from @Christ_Foster's: https://discourse.julialang.org/t/ray-tracing-in-a-week-end-julia-vs-simd-optimized-c/72958/40
# Christ added this feature directly in StaticArrays.jl on 2021-12-17, so
# this will soon be no longer needed.
@inline @inbounds function Base.getproperty(v::SVector{3}, name::Symbol)
    name === :x || name === :r ? getfield(v, :data)[1] :
    name === :y || name === :g ? getfield(v, :data)[2] :
    name === :z || name === :b ? getfield(v, :data)[3] : getfield(v, name)
end

@inline @inbounds function Base.getproperty(v::SVector{2}, name::Symbol)
    name === :x || name === :r ? getfield(v, :data)[1] :
    name === :y || name === :g ? getfield(v, :data)[2] : getfield(v, name)
end

@inline squared_length(v::Vec3) = v ⋅ v
@inline near_zero(v::Vec3) = squared_length(v) < 1e-5
@inline rgb(v::Vec3) = RGB(v...)
@inline rgb_gamma2(v::Vec3) = RGB(sqrt.(v)...)


module TMP

using Gridap
using ChainRulesCore
import ChainRulesCore: rrule

domain = (0,1,0,1)
cells = (10,10)
model = CartesianDiscreteModel(domain,cells)

order = 2
Ω = Triangulation(model)
dΩ = Measure(Ω,2*order)

# Manufactured solution
u(x) = x[1]+x[2]

# Weak form functions
a(p,u,v) = ∫( (p*p)*∇(v)⋅∇(u) )dΩ
l(v) = ∫( -v*Δ(u) )dΩ

# Space for the design dofs
reffe_S = ReferenceFE(lagrangian,Float64,0)
S = TestFESpace(model,reffe_S)
R = TrialFESpace(S)

# Space for the material params
reffe_Q = ReferenceFE(lagrangian,Float64,1)
Q = TestFESpace(model,reffe_Q)
P = TrialFESpace(Q)

# Space for the primal solution
reffe_V = ReferenceFE(lagrangian,Float64,order)
V = TestFESpace(model,reffe_V,dirichlet_tags="boundary")
U = TrialFESpace(V,u)

"""
    f(uvec) -> Number

Objective function
"""
function f(uvec)
  function α(uh)
    eh = u-uh
    eh*eh
  end
  uh = FEFunction(U,uvec)
  sum(∫(α(uh))dΩ)
end

"""
    uvec = g(pvec)

Direct problem from material parameters
"""
function g(pvec)
  ph = FEFunction(P,pvec)
  op = AffineFEOperator((u,v)->a(ph,u,v),l,U,V)
  uh = solve(op)
  uvec = get_free_values(uh)
  uvec
end

"""
    pvec = h(rvec)

Material parameters from design dofs `rvec`.
"""
function h(rvec)
  rh = FEFunction(R,rvec)
  op = AffineFEOperator((p,q)->h1(p,q),q->h1(rh,q),P,Q)
  ph = solve(op)
  pvec = get_free_values(ph)
  pvec
end
h1(p,q) = ∫( ∇(p)⋅∇(q) + p*q )dΩ

"""
   F(rvec) -> Number

Objective function from design dofs
"""
function F(rvec)
  pvec = h(rvec) # Smoothing
  uvec = g(pvec) # Forward PDE
  f(uvec) # eval objective
end

# reverse rules

function rrule(::typeof(f),uvec)
  function f_pullback(dFdF)
    NO_FIELDS, dFdF*df_du(uvec)
  end
  f(uvec), f_pullback
end

function df_du(uvec)
  function dα(uh,du)
    eh = uh-u
    de = du
    ∫( 2*eh*de )dΩ
  end
  uh = FEFunction(U,uvec)
  assemble_vector(du->dα(uh,du),V)
end

function rrule(::typeof(g),pvec)
  uvec = g(pvec)
  function g_pullback(dFdu)
    NO_FIELDS, dg_dp(dFdu,uvec,pvec)
  end
  uvec, g_pullback
end

function dg_dp(dFdu,uvec,pvec)

  # Adjoint problem
  ph = FEFunction(P,pvec)
  A = assemble_matrix((λ,v)->a(ph,v,λ),U,V)
  λvec = A\dFdu

  # Total derivative wrt pvec
  da(p,u,v,dp) = ∫( (2*p*dp)*∇(v)⋅∇(u) )dΩ
  uh = FEFunction(U,uvec)
  λh = FEFunction(V,λvec)
  -1*assemble_vector((dp)->da(ph,uh,λh,dp),Q)
end

function rrule(::typeof(h),rvec)
  pvec = h(rvec)
  function h_pullback(dFdp)
    NO_FIELDS, dh_dr(dFdp)
  end
  pvec, h_pullback
end

function dh_dr(dFdp)
  A = assemble_matrix((p,q)->h1(q,p),P,Q)
  ϕvec = A\dFdp

  ϕh = FEFunction(S,ϕvec)
  assemble_vector((dr)->h1(dr,ϕh),S)
end

rvec = rand(num_free_dofs(R))
@show F(rvec)

pvec, h_pullback = rrule(h,rvec)
uvec, g_pullback = rrule(g,pvec)
obj, f_pullback = rrule(f,uvec)

dFdF = 1
_, dFdu = f_pullback(dFdF)
_, dFdp = g_pullback(dFdu)
_, dFdr = h_pullback(dFdp)

display(dFdr)

@show num_free_dofs(U)
@show num_free_dofs(P)
@show num_free_dofs(R)

using Zygote
dFdr2, = Zygote.gradient(F,rvec)

display(dFdr2)

end # module


#Parameters used to estimate canopy transpiration
mutable struct TranspirationParameter{T<:Float64}
    ρₐ::T #density of dry air (Kg m⁻³)
    cₚₐ::T #Specific heat capacity of dry air (J kg⁻¹ K⁻¹)
    γ::T #Psychrometric constant (Pa K⁻¹)
    λ::T #Latent heat of water vaporization (J kg⁻¹)
end
TranspirationParameter(;ρₐ::T=1.204,cₚₐ::T=1004.0,γ::T=66.1,λ::T=2.454e6) where {T<:Float64} = TranspirationParameter(ρₐ,cₚₐ,γ,λ)

#derivative of saturation Vapor Pressure (Pa) with respect to temperature at T degree celsius
#Derived from Equation 21 in Alduchov and Eskridge (1996)
dSVPₜ(T::AbstractFloat) = 610.94*17.625*243.04*exp(17.625*T/(243.04+T))/(243.04+T)^2

function CalcE_eq(ϕₙₐ::T,Tₐ::T;par::TranspirationParameter=TranspirationParameter()) where {T<:Float64}
    #Equation for calcualting the equilibrium evaporation rate E_eq (Kg s⁻¹ m⁻²)
    #ϕₙₐ net radiation absorbed by the canopy (J s⁻¹ m⁻²)
    #Tₐ air temperature (°C)
    s = dSVPₜ(Tₐ)
    E_eq = s*ϕₙₐ/((s+par.γ)*par.λ)
    return E_eq
end

function CalcE_imp(g_C::T,VPD::T;par::TranspirationParameter=TranspirationParameter()) where {T<:Float64}
    #Equation for calcualting the imposed evaporation rate E_imp (Kg s⁻¹ m⁻²)
    #g_C canopy conductance (water based conductance) (m³ s⁻¹ m⁻²)
    #VPD vapour-pressure deficit (Pa)
    E_imp = par.ρₐ*par.cₚₐ*g_C*VPD/(par.γ*par.λ)
    return E_imp
end

function CalcE(ϕₙₐ::T,Tₐ::T,g_C::T,VPD::T,Ω::T;par::TranspirationParameter=TranspirationParameter()) where {T<:Float64}
    #Equaiton for estimating the canopy tranpiraiton E (Kg s⁻¹ m⁻²)
    #ϕₙₐ net radiation absorbed by the canopy (J s⁻¹ m⁻²)
    #Tₐ air temperature (°C)
    #g_C canopy conductance (water based conductance) (m³ s⁻¹ m⁻²)
    #VPD vapour-pressure deficit (Pa)
    #Ω a dimentionless number in the range 0.0-1.0
    E_eq = CalcE_eq(ϕₙₐ,Tₐ;par=par)
    E_imp = CalcE_imp(g_C,VPD;par=par)
    E = Ω*E_eq+(1- Ω)*E_imp
    return E
end

function Estg_C(E::T,ϕₙₐ::T,Tₐ::T,VPD::T,Ω::T;par::TranspirationParameter=TranspirationParameter()) where {T<:Float64}
    #Equaiton for estimating the canopy conductance g_C (m³ s⁻¹ m⁻²)
    #E canopy tranpiraiton (Kg s⁻¹ m⁻²)
    #ϕₙₐ net radiation absorbed by the canopy (J s⁻¹ m⁻²)
    #Tₐ air temperature (°C)
    #VPD vapour-pressure deficit (Pa)
    #Ω a dimentionless number in the range 0.0-1.0
    E_eq = CalcE_eq(ϕₙₐ,Tₐ;par)
    g_C = (E-Ω*E_eq)/(1-Ω)*(par.γ*par.λ)/(par.ρₐ*par.cₚₐ*VPD)
    return g_C
end

function PAR2Rad(PAR::T;conv_coeff::T=10^6/2.3) where {T<:Float64}
    #Funciton for converting From PAR (mol s⁻¹ m⁻²) to global
    #radiation (Rad, J s⁻¹ m⁻²)
    Rad = PAR*conv_coeff
    return Rad
end

function gₛ2g_C(gₛ::T,PAR::T,LAI::T;cons::CCPH.Constants=CCPH.Constants(),treepar::CCPH.TreePar=CCPH.TreePar(),
    I_gₛ::T=50.0e-6) where {T<:Float64}
    #Estimates canopy conductance g_C (m³ s⁻¹ m⁻²)
    #gₛ stomatal conductanceg (mol CO₂ s⁻¹ m⁻²)
    #PAR at canopy top (mol s⁻¹ m⁻²)
    #LAI Canopy leaf area (m² leaf area m⁻² ground area)
    #I_gₛ Irradiace level at which gₛ is half of canopy top gₛ (mol s⁻¹ m⁻²)

    g_C = cons.r*cons.M_H2O/cons.ρ_H2O*gₛ*log((I_gₛ+PAR)/(I_gₛ+PAR*exp(-treepar.k*LAI)))/treepar.k

    return g_C
end

function Calctranpiration(VPD_z::T,Sₑ::T;E_cm_ref::T=1.812,s_VPD_z::T=3.121,
    s_Sₑ::T=18.342) where {T<:Float64}
    #Tor-Ngern et al. 2017 model for estimating canopy transpiration E_c (mm/d)
    #Standard paramater values (E_cm_ref,s_VPD_z,s_Sₑ) are taken from the
    #original paper for the Rosinedal scots pine stand
    #--Input--
    #VPD_z vapor pressure deficit (Pa)
    #Sₑ effective saturation (-)

    E_cm = E_cm_ref*(1-exp(-s_VPD_z*VPD_z*10^-3))
    E_c = E_cm*(1-exp(-s_Sₑ*Sₑ))

    return E_c
end

function Est_gₛ(LAI::T,daylight::T;cons::CCPH.Constants=CCPH.Constants(),
    env::CCPH.EnvironmentStruct=CCPH.EnvironmentStruct()) where {T<:Float64}
    # Estimate stomatal conductance gₛ (mol C m⁻² leaf area s⁻¹) based on
    # Weather data and leaf area index, using the canopy transpiration from
    # Tor-Ngern et al. 2017
    #--Input--
    # LAI leaf area index (m² leaf area m⁻² ground area)
    # daylight lenght of a single day (s)
    VPD = env.VPD
    θₛ = env.θₛ
    P = env.P
    VPD_z = VPD.*daylight/(24*3600)
    Sₑ = CCPH.CalcSₑ.(θₛ)
    E_c = Calctranpiration.(VPD_z,Sₑ) #mm/day
    E_c *= 10^-3*cons.ρ_H2O/cons.M_H2O/daylight #mol m⁻² s⁻¹
    gₛ = E_c*P/(VPD*cons.r*LAI)
    return gₛ
end


"""
    $(FUNCTIONNAME)(FT=Float64)

Creates a `LeafDistribution` following a planophile (mostly horizontal) distribution using a Beta distribution

Standard values for θₗ (26.76) and s (18.51) from Bonan https://doi.org/10.1017/9781107339217.003, Table 2.1
"""
function planophile_leaves(FT=Float64)
    θ = FT(26.76)
    s = FT(18.51)
    α, β =  βparameters(deg2rad(θ), deg2rad(s));
    return LeafDistribution(Beta(FT(α), FT(β)), FT(2/π));
end

function planophile_leaves2(FT=Float64)
    return LeafDistribution(Beta(FT(1), FT(2.768)), FT(2/π));
end

"""
    $(FUNCTIONNAME)(FT=Float64)

Creates a `LeafDistribution` following a uniform distribution (all angles equally likely)
"""
function uniform_leaves(FT=Float64)
    return LeafDistribution(Uniform(FT(0), FT(1)),FT(2/π));
end

"""
    $(FUNCTIONNAME)(FT=Float64)

Creates a `LeafDistribution` following a plagiophile (mostly 45degrees) distribution using a Beta distribution

Standard values for θₗ (45.0) and s (16.27) from Bonan https://doi.org/10.1017/9781107339217.003, Table 2.1
"""
function plagiophile_leaves(FT=Float64)
    θ = FT(45)
    s = FT(16.27)
    α, β =  βparameters(deg2rad(θ), deg2rad(s));
    return LeafDistribution(Beta(FT(α), FT(β)), FT(2/π));
end

"""
    $(FUNCTIONNAME)(FT=Float64)

Creates a `LeafDistribution` following a erectophile (mostly vertical) distribution using a Beta distribution

Standard values for θₗ (63.24) and s (18.5) from Bonan https://doi.org/10.1017/9781107339217.003, Table 2.1
"""
function erectophile_leaves(FT=Float64)
    θ = FT(63.24)
    s = FT(18.5)
    α, β =  βparameters(deg2rad(θ), deg2rad(s));
    return LeafDistribution(Beta(FT(α), FT(β)), FT(2/π));
end

"""
    $(FUNCTIONNAME)(FT=Float64)

Creates a `LeafDistribution` following a spherical (distributed as facets on a sphere) distribution using a Beta distribution

Standard values for θₗ (57.3) and s (21.55) from Bonan https://doi.org/10.1017/9781107339217.003, Table 2.1
"""
function spherical_leaves(FT=Float64)
    θ = FT(57.3)
    s = FT(21.55)
    α, β =  βparameters(deg2rad(θ), deg2rad(s));
    return LeafDistribution(Beta(FT(α), FT(β)), FT(2/π));
end

function flat_leaves(FT=Float64)
    θ = FT(57.3)
    s = FT(21.55)
    α, β =  βparameters(deg2rad(θ), deg2rad(s));
    return LeafDistribution(Beta(FT(α), FT(β)), FT(2/π));
end


# sample model as in BayesExamples1 and 2, but using
# a specialized package and more sophisticated sampling
using Distributions, Turing, StatsPlots
function main()
# sample is from exponential, prior is lognormal
y = rand(Exponential(3.0), 300)
# the model: prior and likelihood
@model function ExpModel(y)
    θ ~ LogNormal(1.,1.)
    y ~ Exponential(θ)
end
# get the chain
chain = sample(ExpModel(y), NUTS(200, 0.65), 10000)
end
## run it
chain = main()
# examine and plot chain
display(chain)
plot(chain[end-1000:end])


using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:4] <= 1)

                     @variable(m, -12.503038214935698 <= q <= 11.138318577023822)

                     add_NL_constraint(m, :(gelu(-0.2825986234777189 + -0.06477303738755547*gelu(0.91111659897616 + 0.32194568310871174*$(x[1]) + -0.04079743189253193*$(x[2]) + -0.7834407613653229*$(x[3]) + -0.5105294854228863*$(x[4])) + 0.23435836886254213*gelu(-0.27698098400696525 + 0.765733032903245*$(x[1]) + 0.6374561618281618*$(x[2]) + -0.006430667300674209*$(x[3]) + -0.6595935184751247*$(x[4])) + 0.7516401137686608*gelu(0.7789283911146407 + -0.8853253179157212*$(x[1]) + -0.6613182381252614*$(x[2]) + 0.5007525575333553*$(x[3]) + 0.7960733427996214*$(x[4]))) + gelu(0.12692347233027856 + 0.060984281016326136*gelu(0.91111659897616 + 0.32194568310871174*$(x[1]) + -0.04079743189253193*$(x[2]) + -0.7834407613653229*$(x[3]) + -0.5105294854228863*$(x[4])) + 0.4767325000240725*gelu(-0.27698098400696525 + 0.765733032903245*$(x[1]) + 0.6374561618281618*$(x[2]) + -0.006430667300674209*$(x[3]) + -0.6595935184751247*$(x[4])) + -0.9515049437955176*gelu(0.7789283911146407 + -0.8853253179157212*$(x[1]) + -0.6613182381252614*$(x[2]) + 0.5007525575333553*$(x[3]) + 0.7960733427996214*$(x[4]))) + gelu(-0.022625917828582143 + -0.9740734403326825*gelu(0.91111659897616 + 0.32194568310871174*$(x[1]) + -0.04079743189253193*$(x[2]) + -0.7834407613653229*$(x[3]) + -0.5105294854228863*$(x[4])) + -0.9306328050206538*gelu(-0.27698098400696525 + 0.765733032903245*$(x[1]) + 0.6374561618281618*$(x[2]) + -0.006430667300674209*$(x[3]) + -0.6595935184751247*$(x[4])) + 0.61848968844355*gelu(0.7789283911146407 + -0.8853253179157212*$(x[1]) + -0.6613182381252614*$(x[2]) + 0.5007525575333553*$(x[3]) + 0.7960733427996214*$(x[4]))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m


using Manopt, ManifoldsBase, Manifolds, Test
using Random
@testset "Particle Swarm" begin
    # Test the particle swarm algorithm
    A = [1.0 3.0 4.0; 3.0 -2.0 -6.0; 4.0 -6.0 5.0]
    @testset "Eucliedean Particle Swarm" begin
        M = Euclidean(3)
        F(::Euclidean, x) = (x' * A * x) / (x' * x)
        x_start = [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]
        x_start2 = deepcopy(x_start)

        Random.seed!(35)
        o = particle_swarm(M, F; x0=x_start, return_options=true)
        g = get_solver_result(o)

        Random.seed!(35)
        g2 = particle_swarm(M, F; x0=x_start2, return_options=false)
        @test isequal(g, g2)

        # the cost of g and the p[i]'s are not greater after one step
        j = argmin([F(M, y) for y in x_start])
        g0 = deepcopy(x_start[j])
        @test F(M, g) <= F(M, g0)
        for i in 1:3
            @test F(M, o.p[i]) <= F(M, x_start[i])
            # the cost of g is not greater than the cost of any p[i]
            @test F(M, g) <= F(M, o.p[i])
        end
    end
    @testset "Spherical Particle Swarm" begin
        Random.seed!(42)
        M = Sphere(2)
        F(::Sphere, x) = transpose(x) * A * x
        x_start = [random_point(M) for i in 1:3]
        v_start = [random_tangent(M, y) for y in x_start]
        p = CostProblem(M, F)
        o = ParticleSwarmOptions(x_start, v_start)
        initialize_solver!(p, o)
        step_solver!(p, o, 1)
        for i in 1:3
            # check that the new particle locations are on the manifold
            @test is_point(M, o.x[i], true)
            # check that the new velocities are tangent vectors of the original particle locations
            @test is_vector(M, o.x[i], o.velocity[i], true; atol=5 * 10^(-16))
        end
    end
end


function sampleFromDiscrete(probs,Pw,nonzeroLength)
# draw a sample from unnormalized discrete distribution, probs
# Pw is the total of discrete distributions
    temp = rand()*Pw;
    total = 0.0;
    for i = 1:nonzeroLength
        total = total + probs[i];
        if temp<total
            return i;
            break
        end
    end
end

function generateAlias(p, k)
# creates a 2xnumTopics alias table
    # p is the unnormalized probability distribution of size 1xk
    # k is the number of outcomes
    # returns a generated alias table of size kx2

    # initialization...
    AT = zeros(k,2); # generated alias table of size kx2
    # temporary array to store low and high probabilities
    L_probInd = Int64[]; L_prob = Float64[];H_probInd = Int64[]; H_prob = Float64[];
    Lind = 1; Hind = 1; # current index of AT and L or H
    kInv = 1/k; # average probability of the distribution
    sum_p = sum(p);

    # generation alias table...
    for i = 1:k
        p[i] = p[i] ./ sum_p; # normalizing the previously unnormalized discrete distribution
        if p[i]<=kInv
            push!(L_probInd,i);
            push!(L_prob,p[i]);
            Lind+= 1;
        else
            push!(H_probInd,i);
            push!(H_prob,p[i]);
            Hind+= 1;
        end
        AT[i,1]=i; AT[i,2]=i*kInv; # initial formation of alias table
    end
    Lind-= 1; Hind-= 1; # this points to the last added element
    while Lind>=1 && Hind>=1
        ATind = L_probInd[Lind]; # extract last index from L_probInd
        AT[ATind,1] = H_probInd[Hind]; # extract new index from H_probInd and map it to AT
        AT[ATind,2] = (ATind-1)*kInv+L_prob[Lind]; # place low probability value in AT
        Ph = H_prob[Hind]-(kInv-L_prob[Lind]); # cut alias probability from high value probabilities
        Lind-= 1; Hind-= 1; # points to next entry we can extract
        if Ph>kInv
            Hind+= 1;
            H_probInd[Hind] = AT[ATind,1]; H_prob[Hind] = Ph;
        else
            Lind+= 1;
            L_probInd[Lind] = AT[ATind,1];L_prob[Lind] = Ph;
        end
    end
    return AT;
end

function sampleAlias(AT,k)
# generates a sample from alias table
    bin = rand(1:k);
    if AT[bin,2]>rand()
        s = bin;
    else
        s = convert(Int64,AT[bin,1]); # to make sure we are getting integer samples from alias table
    end
    return s;
end

function aliasMethodOnWord(word,Alpha,wordTopicCounts,topicCounts,Beta,BetaSum,numTopics,numCachedSamples,cachedSamples,cached_qw)
    qw = Array{Float64}(numTopics); # unnormalized discrete distribution of dense bucket
    for topic = 1:numTopics
        qw[topic] = Alpha[topic] .* (wordTopicCounts[word,topic] + Beta[word]) ./ (topicCounts[topic] + BetaSum);
        cached_qw[word,topic] = qw[topic];
    end
    aliasTable = generateAlias(qw,numTopics);
    for sample = 1:numCachedSamples
        cachedSamples[word,sample] = sampleAlias(aliasTable,numTopics);
    end

    return cachedSamples,cached_qw;
end


###################################
## Demo script for Bayesian Dimensionality Reduction

using PyPlot, PyCall
@pyimport sklearn.datasets as datasets

push!(LOAD_PATH,".")
import DimensionalityReduction

function load_facedata(skip::Int)
    face = datasets.fetch_olivetti_faces()
    Y_raw = face["images"]
    N, S_raw, _ = size(Y_raw)

    L = round(Int, S_raw / skip)
    Y_tmp = Y_raw[:,1:skip:end, 1:skip:end]
    Y = convert(Array{Float64, 2}, reshape(Y_tmp, N, size(Y_tmp,2)*size(Y_tmp,3))')
    D = size(Y, 1)

    return Y, D, L
end

function visualize(Y::Array{Float64,2}, L::Int)
    D, N = size(Y)
    base = round(Int, sqrt(N))
    v = round(Int, (L*ceil(N / base)))
    h = L * base
    pic = zeros(v, h)

    for n in 1 : N
        i = round(Int, (L*ceil(n / base)))
        idx1 = i - L + 1 : i
        idx2 = L*mod(n-1, base)+1 : L*(mod(n-1, base) + 1)
        pic[idx1,idx2] = reshape(Y[:,n], L, L)
    end
    imshow(pic, cmap=ColorMap("gray"))
end

function visualize(Y::Array{Float64,2}, L::Int, mask::BitArray{2})
    # for missing
    D, N = size(Y)
    base = round(Int, sqrt(N))
    v = round(Int, (L*ceil(N / base)))
    h = L * base
    pic = zeros(v, h, 3)

    Y_3dim = zeros(D, N, 3)
    for i in 1 : 3
        if i == 2
            Y_tmp = deepcopy(Y)
            Y_tmp[mask] = 1
            Y_3dim[:,:,i] = Y_tmp
        else
            Y_tmp = deepcopy(Y)
            Y_tmp[mask] = 0
            Y_3dim[:,:,i] = Y_tmp
        end
    end

    for n in 1 : N
        i = round(Int, (L*ceil(n / base)))
        idx1 = i - L + 1 : i
        idx2 = L*mod(n-1, base)+1 : L*(mod(n-1, base) + 1)
        for i in 1 : 3
            pic[idx1,idx2,i] = reshape(Y_3dim[:,n,i], L, L)
        end
    end
    imshow(pic, cmap=ColorMap("gray"))
end

"""
Run a demo script of missing data interpolation for face dataset.
"""
function test_face_missing()
    # load data
    skip = 2
    Y, D, L = load_facedata(skip)

    # mask
    missing_rate = 0.50
    mask = rand(size(Y)) .< missing_rate
    Y_obs = deepcopy(Y)
    Y_obs[mask] = NaN

    # known parames
    M = 16
    sigma2_y = 0.001
    Sigma_W = zeros(M,M,D)
    Sigma_mu = 1.0 * eye(D)
    for d in 1 : D
        Sigma_W[:,:,d] = 0.1 * eye(M)
    end
    prior = DimensionalityReduction.DRModel(D, M, sigma2_y, zeros(M, D), Sigma_W, zeros(D), Sigma_mu)

    # learn & generate
    max_iter = 100
    posterior, X_est = DimensionalityReduction.VI(deepcopy(Y_obs), prior, max_iter)
    Y_est = posterior.m_W'*X_est + repmat(posterior.m_mu, 1, size(X_est, 2))
    Y_itp = deepcopy(Y_obs)
    Y_itp[mask] = Y_est[mask]

    #visualize
    N_show = 4^2

    figure("Observation")
    clf()
    visualize(Y_obs[:,1:N_show], L, mask[:,1:N_show])
    title("Observation")

    #figure("Estimation")
    #clf()
    #visualize(Y_est[:,1:N_show], L)
    #title("Estimation")

    figure("Interpolation")
    clf()
    visualize(Y_itp[:,1:N_show], L)
    title("Interpolation")

    figure("Truth")
    clf()
    visualize(Y[:,1:N_show], L)
    title("Truth")
    show()
end

"""
Run a dimensionality reduction demo using Iris dataset.
"""
function test_iris()
    ##################
    # load data
    iris = datasets.load_iris()
    Y_obs = iris["data"]'
    label_list = [iris["target_names"][elem+1] for elem in iris["target"]]
    D, N = size(Y_obs)

    ##################
    # 2D compression

    # model
    M = 2
    sigma2_y = 0.001
    Sigma_W = zeros(M,M,D)
    Sigma_mu = 1.0 * eye(D)
    for d in 1 : D
        Sigma_W[:,:,d] = 0.1 * eye(M)
    end
    prior = DimensionalityReduction.DRModel(D, M, sigma2_y, zeros(M, D), Sigma_W, zeros(D), Sigma_mu)

    # learn & generate
    max_iter = 100
    posterior, X_est = DimensionalityReduction.VI(deepcopy(Y_obs), prior, max_iter)

    # visualize
    figure("2D plot")
    clf()
    scatter(X_est[1,1:50], X_est[2,1:50], color="r")
    scatter(X_est[1,51:100], X_est[2,51:100], color="g")
    scatter(X_est[1,101:end], X_est[2,101:end], color="b")
    xlabel("\$x_1\$", fontsize=20)
    ylabel("\$x_2\$", fontsize=20)
    legend([label_list[1], label_list[51], label_list[101]], fontsize=16)

    ##################
    # 3D compression

    # model
    M = 3
    sigma2_y = 0.001
    Sigma_W = zeros(M,M,D)
    Sigma_mu = 1.0 * eye(D)
    for d in 1 : D
        Sigma_W[:,:,d] = 0.1 * eye(M)
    end
    prior = DimensionalityReduction.DRModel(D, M, sigma2_y, zeros(M, D), Sigma_W, zeros(D), Sigma_mu)

    # learn & generate
    max_iter = 100
    posterior, X_est = DimensionalityReduction.VI(deepcopy(Y_obs), prior, max_iter)

    # visualize
    figure("3D plot")
    clf()
    scatter3D(X_est[1,1:50], X_est[2,1:50], X_est[3,1:50], c="r")
    scatter3D(X_est[1,51:100], X_est[2,51:100], X_est[3,51:100], c="g")
    scatter3D(X_est[1,101:end], X_est[2,101:end], X_est[3,101:end], c="b")
    legend([label_list[1], label_list[51], label_list[101]], fontsize=16)
    xlabel("\$x_1\$", fontsize=20)
    ylabel("\$x_2\$", fontsize=20)
    zlabel("\$x_3\$", fontsize=20)
    show()
end

#test_face_missing()
test_iris()


"""
    multinomial_hyper(α::AbstractArray{Float32,1})

[Dirichlet Distribution](https://en.wikipedia.org/wiki/Dirichlet_distribution)
"""
struct multinomial_hyper <: distribution_hyper_params
    α::AbstractArray{Float32,1}
end

mutable struct multinomial_sufficient_statistics <: sufficient_statistics
    N::Float32
    points_sum::AbstractArray{Float32,1}
end

function calc_posterior(prior:: multinomial_hyper, suff_statistics::multinomial_sufficient_statistics)
    if suff_statistics.N == 0
        return prior
    end
    return multinomial_hyper(prior.α + suff_statistics.points_sum)
end

function sample_distribution(hyperparams::multinomial_hyper)
    return multinomial_dist(log.(rand(Dirichlet(Float64.(hyperparams.α)))))
end

function create_sufficient_statistics(hyper::multinomial_hyper,posterior::multinomial_hyper,points::AbstractArray{Float32,2}, pts_to_group = 0)
    pts = copy(points)
    points_sum = sum(pts, dims = 2)[:]
    S = pts * pts'
    return multinomial_sufficient_statistics(size(points,2),points_sum)
end

function log_marginal_likelihood(hyper::multinomial_hyper, posterior_hyper::multinomial_hyper, suff_stats::multinomial_sufficient_statistics)
    D = length(suff_stats.points_sum)
    logpi = log(pi)
    val = logabsgamma(sum(hyper.α))[1] -logabsgamma(sum(posterior_hyper.α))[1] + sum((x-> logabsgamma(x)[1]).(posterior_hyper.α) - (x-> logabsgamma(x)[1]).(hyper.α))
    return val
end

function aggregate_suff_stats(suff_l::multinomial_sufficient_statistics, suff_r::multinomial_sufficient_statistics)
    return multinomial_sufficient_statistics(suff_l.N+suff_r.N, suff_l.points_sum + suff_r.points_sum)
end


module Bessels

export besselj0
export besselj1

export bessely0
export bessely1

export besseli
export besselix
export besseli0
export besseli0x
export besseli1
export besseli1x

export besselk
export besselkx
export besselk0
export besselk0x
export besselk1
export besselk1x

include("besseli.jl")
include("besselj.jl")
include("besselk.jl")
include("bessely.jl")
include("constants.jl")

include("Float128/besseli.jl")
include("Float128/besselj.jl")
include("Float128/besselk.jl")
include("Float128/bessely.jl")
include("Float128/constants.jl")

include("math_constants.jl")
include("U_polynomials.jl")
include("recurrence.jl")
include("misc.jl")
#include("hankel.jl")

end


using AsyPlots, LaTeXStrings, SpecialFunctions

x = range(0, stop=1, length = 400)

p = plot(collect(x),x.^44 .* (1 .- x).^11 / beta(44,11),xticks=Ticks([1]),yticks=Ticks([1])) + Plot(Path([4/5 -0.04; 4/5 0.04])) + Plot(Label(L"\frac{\alpha}{\alpha+\beta}",(4/5,-0.2));width=128)


"""
### Nuclear attraction term
"""
function Aterm(i::Int64,r::Int64,u::Int64,l1::Int64,l2::Int64,ax::Float64,bx::Float64,cx::Float64,gamma::Float64)
  term1 = (-1)^(i+u)*binomial_prefactor(i,l1,l2,ax,bx)
  term2 = factorial(i)*cx^(i-2r-2u)
  term3 = (1/4/gamma)^(r+u)/factorial(r)/factorial(u)/factorial(i-2r-2u)
  return term1*term2*term3
end

function Aarray(l1::Int64,l2::Int64,a::Float64,b::Float64,c::Float64,g::Float64)
  Imax = l1+l2+1
  A = zeros(Float64,Imax)
  for i in 0:(Imax-1)
    for r in 0:div(i,2)
      for u in 0:div(i-2r,2)
        I = i-2r-u+1
        A[I] += Aterm(i,r,u,l1,l2,a,b,c,g)
      end
    end
  end
  return A
end

function nuclear_attraction(aexpn::Float64,ax::Float64,ay::Float64,az::Float64,
              aI::Int64,aJ::Int64,aK::Int64,
              bexpn::Float64,bx::Float64,by::Float64,bz::Float64,
              bI::Int64,bJ::Int64,bK::Int64,
              cx::Float64,cy::Float64,cz::Float64)
  #print("na($aexpn,$ax,$ay,$az,$aI,$aJ,$aK,$bexpn,$bx,$by,$bz,$bI,$bJ,$bK,$cx,$cy,$cz)=")
  px,py,pz = gaussian_product_center(aexpn,ax,ay,az,bexpn,bx,by,bz)
  gamma = aexpn+bexpn
  rab2 = dist2(ax-bx,ay-by,az-bz)
  rcp2 = dist2(cx-px,cy-py,cz-pz)
  Ax = Aarray(aI,bI,px-ax,px-bx,px-cx,gamma)
  Ay = Aarray(aJ,bJ,py-ay,py-by,py-cy,gamma)
  Az = Aarray(aK,bK,pz-az,pz-bz,pz-cz,gamma)
  total = 0
  for I in 0:(aI+bI)
    for J in 0:(aJ+bJ)
      for K in 0:(aK+bK)
        total += Ax[I+1]*Ay[J+1]*Az[K+1]*Fgamma(I+J+K,rcp2*gamma)
      end
    end
  end
  val=-2pi*exp(-aexpn*bexpn*rab2/gamma)*total/gamma
  #println(val)
  #println((Ax,Ay,Az,rcp2*gamma,Fgamma(0,rcp2*gamma)))
  return val
end

function nuclear_attraction(a::PGBF,b::PGBF,cx::Float64,cy::Float64,cz::Float64)
  return a.norm*b.norm*nuclear_attraction(a.expn,a.x,a.y,a.z,a.I,a.J,a.K,
                      b.expn,b.x,b.y,b.z,b.I,b.J,b.K,cx,cy,cz)
end
nuclear_attraction(a::PGBF,b::PGBF,c::Atom) = c.atno*nuclear_attraction(a,b,c.x,c.y,c.z)
nuclear_attraction(a::PGBF,b::PGBF,m::Molecule) = sum([nuclear_attraction(a,b,c) for c in m.atomlist])

function Fgamma(m::Int64,x::Float64,SMALL::Float64=1e-12)
  #println("Fgamma($m,$x)")
  x = max(x,SMALL) # Evidently needs underflow protection
  return 0.5*x^(-m-0.5)*gammainc(m+0.5,x)
end

function gammainc(a::Float64,x::Float64)
  # This is the series version of gamma from pyquante. For reasons I don't get, it
  # doesn't work around a=1. This works alright, but is only a stopgap solution
  # until Julia gets an incomplete gamma function programmed
  if abs(a-1) < 1e-3
    println("Warning: gammainc_series is known to have problems for a ~ 1")
  end
  if x < (a+1.0)
    #Use the series representation
    gam,gln = gser(a,x)
  else
    #Use continued fractions
    gamc,gln = gcf(a,x)
    gam = 1-gamc
  end
  return exp(gln)*gam
end

function gser(a::Float64,x::Float64,ITMAX::Int64=100,EPS::Float64=3e-9)
  # Series representation of Gamma. NumRec sect 6.1.
  gln=lgamma(a)
  if x == 0
    return 0,gln
  end
  ap = a
  delt = s = 1/a
  for i in 1:ITMAX
    ap += 1
    delt *= (x/ap)
    s += delt
    if abs(delt) < abs(s)*EPS
      break
    end
  end
  return s*exp(-x+a*log(x)-gln),gln
end

function gcf(a::Float64,x::Float64,ITMAX::Int64=200,EPS::Float64=3e-9,FPMIN::Float64=1e-30)
  #Continued fraction representation of Gamma. NumRec sect 6.1"
  gln=lgamma(a)
  b=x+1.-a
  c=1./FPMIN
  d=1./b
  h=d
  for i in 1:ITMAX
    an=-i*(i-a)
    b=b+2.
    d=an*d+b
    if abs(d) < FPMIN
      d=FPMIN
    end
    c=b+an/c
    if abs(c) < FPMIN
      c=FPMIN
    end
    d=1./d
    delt=d*c
    h=h*delt
    if abs(delt-1.) < EPS
      break
    end
  end
  gammcf = exp(-x+a*log(x)-gln)*h
  return gammcf,gln
end

# Need a nested scope to squeeze this into the contract function
function nuclear_attraction(a::CGBF,b::CGBF,cx::Float64,cy::Float64,cz::Float64)
  na(a,b) = nuclear_attraction(a,b,cx,cy,cz)
  contract(na,a,b)
end
function nuclear_attraction(a::CGBF,b::CGBF,c::Atom)
  na(a,b) = nuclear_attraction(a,b,c)
  contract(na,a,b)
end
function nuclear_attraction(a::CGBF,b::CGBF,m::Molecule)
  na(a,b) = nuclear_attraction(a,b,m)
  contract(na,a,b)
end

function test_a_terms()
  @assert Aterm(0,0,0,0,0,0.,0.,0.,0.) == 1.0
  @assert Aarray(0,0,0.,0.,0.,1.) == [1.0]
  @assert Aarray(0,1,1.,1.,1.,1.) == [1.0, -1.0]
  @assert Aarray(1,1,1.,1.,1.,1.) == [1.5, -2.5, 1.0]
  @assert Aterm(0,0,0,0,0,0.,0.,0.,1.) == 1.0
  @assert Aterm(0,0,0,0,1,1.,1.,1.,1.) == 1.0
  @assert Aterm(1,0,0,0,1,1.,1.,1.,1.) == -1.0
  @assert Aterm(0,0,0,1,1,1.,1.,1.,1.) == 1.0
  @assert Aterm(1,0,0,1,1,1.,1.,1.,1.) == -2.0
  @assert Aterm(2,0,0,1,1,1.,1.,1.,1.) == 1.0
  @assert Aterm(2,0,1,1,1,1.,1.,1.,1.) == -0.5
  @assert Aterm(2,1,0,1,1,1.,1.,1.,1.) == 0.5
end

function test_gamma()
  # gammainc test functions. Test values taken from Mathematica
  # println("a=0.5 test")
  @assert maximum([gammainc(0.5,float(x)) for x in 0:10]
      -Any[0, 1.49365, 1.69181, 1.7471, 1.76416, 1.76968,
        1.77151, 1.77213, 1.77234, 1.77241, 1.77244]) < 1e-5

  # println("a=1.5 test")
  @assert maximum([gammainc(1.5,float(x)) for x in 0:10]
      -Any[0, 1.49365, 1.69181, 1.7471, 1.76416, 1.76968,
        1.77151, 1.77213, 1.77234, 1.77241, 1.77244]) < 1e-5
  # println("a=2.5 test")
  @assert maximum([gammainc(2.5,float(x)) for x in 0:10]
      -Any[0, 0.200538, 0.59898, 0.922271, 1.12165, 1.22933,
        1.2831, 1.30859, 1.32024, 1.32542, 1.32768]) < 1e-5
end

function test_na()
  s = pgbf(1.0)
  c = cgbf(0.0,0.0,0.0)
  push!(c,1,1)
  @assert isapprox(amplitude(c,0,0,0),0.71270547)
  @assert isapprox(nuclear_attraction(s,s,0.,0.,0.),-1.59576912)
  @assert isapprox(nuclear_attraction(c,c,0.,0.,0.),-1.595769)
end

function test_fgamma()
  for (x,res) in Any[(0.,1),
          (30.,0.161802),
          (60.,0.1144114),
          (90.,0.0934165),
          (120.,0.08090108),
          (300.,0.051166336)]
    #@assert isapprox(res,Fgamma(0,x))
    @printf("res, Fgamma(0,x) %18.10f %18.10f\n", x, Fgamma(0,x))
  end
  @printf "test_fgamma is passed\n"
end

#todo make into a test
function test_one()
  s1 = pgbf(1)
  s2 = pgbf(1,0,1,0)
  x=y=0.
  println("S: $(overlap(s1,s2))")
  println("T: $(kinetic(s1,s2))")
  for z in linspace(0,1,5)
    println("V: $z $(nuclear_attraction(s1,s2,x,y,z))")
  end
end

function test_na2()
  li,h = lih.atomlist
  bfs = build_basis(lih)
  s1,s2,x,y,z,h1s = bfs.bfs
  @assert isapprox(nuclear_attraction(s1,s1,lih),-8.307532656)
end

nuclear_repulsion(a::Atom,b::Atom)= a.atno*b.atno/sqrt(dist2(a.x-b.x,a.y-b.y,a.z-b.z))
function nuclear_repulsion(mol::Molecule)
  nr = 0
  for (i,j) in pairs(nat(mol),"subdiag")
    nr += nuclear_repulsion(mol.atomlist[i],mol.atomlist[j])
  end
  return nr
end

function test_geo_basis()

  @printf("\nCalling test_geo_basis\n")

  E1     = nuclear_repulsion(h2)
  E1_ref = 0.7223600367
  @printf("H2: nuclear_repulsion %18.10f, err = %18.10f\n", E1, abs(E1_ref-E1))

  @assert nel(h2) == 2
  @assert nel(h2o) == 10
  @assert length(sto3g)==10

  bfs = build_basis(h2)
  @assert length(bfs.bfs)==2

  l,r = bfs.bfs
  @assert isapprox(overlap(l,l),1)
  @assert isapprox(overlap(r,r),1)

  #@assert isapprox(overlap(l,r),0.66473625)
  ovl1 = overlap(l,r)
  ovl1_ref = abs(ovl1 - 0.66473625)
  @printf("overlap(l,r), err %18.10f, %18.10f\n", ovl1, ovl1_ref)

  @assert isapprox(kinetic(l,l),0.76003188)
  @assert isapprox(kinetic(r,r),0.76003188)
  @assert isapprox(kinetic(l,r),0.24141861181119084)
  @assert isapprox(coulomb(l,l,l,l), 0.7746059439196398)
  @assert isapprox(coulomb(r,r,r,r), 0.7746059439196398)
  @assert isapprox(coulomb(l,l,r,r), 0.5727937653511646)
  @assert isapprox(coulomb(l,l,l,r), 0.4488373301593464)
  @assert isapprox(coulomb(l,r,l,r), 0.3025451156654606)
  bfs = build_basis(h2o)

  s1,s2,px,py,pz,hl,hr = bfs.bfs
  @assert isapprox(coulomb(s1,s2,hl,hr),0.03855344493645537)
  @assert isapprox(coulomb(s1,pz,hl,hr),-0.0027720110485359053)
  @assert isapprox(coulomb(s1,hl,pz,hr),-0.010049491284827426)
  @assert coulomb(s1,py,hl,hr)==0
  @assert coulomb(s1,hl,py,hr)==0
end


using KrylovKit, ArnoldiMethod, LinearMaps, LinearAlgebra, RecursiveArrayTools

struct PrecPartialSchur{Ts, Tu, Tsm1, Teigen}
	S::Ts
	U::Tu
	Sm1::Tsm1
	eigenvalues::Teigen
end

function LinearAlgebra.ldiv!(out, Pl::PrecPartialSchur, rhs::AbstractArray)
	########################################################
	# U * Sm1 * Ut + (I - U Ut)
	# U(Sm1 - I)Ut + I
	########################################################
	# y = (Pl.U * Pl.Sm1 * transpose(Pl.U) .+ (I - Pl.U * transpose(Pl.U))) * rhs
	y = transpose(Pl.U) * rhs
	y .= Pl.Sm1 * y .- y
	out .= rhs .+ Pl.U * y
end

function LinearAlgebra.ldiv!(Pl::PrecPartialSchur, rhs::AbstractArray)
	########################################################
	# U * Sm1 * Ut + (I - U Ut)
	# U(Sm1 - I)Ut + I
	########################################################
	# y = (Pl.U * Pl.Sm1 * transpose(Pl.U) .+ (I - Pl.U * transpose(Pl.U))) * rhs
	y = transpose(Pl.U) * rhs
	y .= Pl.Sm1 * y .- y
	rhs .= rhs .+ Pl.U * y
end

function LinearAlgebra.:\(Pl::PrecPartialSchur, rhs::AbstractArray)
	out = similar(rhs)
	ldiv!(out, Pl, rhs)
	return out
end

"""
	PrecPartialSchurKrylovKit(J, x0, nev, which = :LM; krylovdim = max(2nev, 20), verbosity = 0)

Builds a preconditioner based on deflation of `nev` eigenvalues chosen according to `which`. A partial Schur decomposition is computed (Matrix-Free), using the package `KrylovKit.jl`, from which a projection is built. The options are similar to the ones of `EigKrylovKit()`.
"""
function PrecPartialSchurKrylovKit(J, x0, nev, which = :LM; krylovdim = max(2nev, 20), verbosity = 0)
	H, V, vals, info = KrylovKit.schursolve(J, x0, nev, which, Arnoldi(krylovdim = krylovdim, verbosity = verbosity))
	Q, S = qr(H)
	U = VectorOfArray(V) * Q
	return PrecPartialSchur(S, U, inv(S), vals)
end

"""
	PrecPartialSchurArnoldiMethod(J, N, nev, which = LM(); tol = 1e-9, kwargs...)

Builds a preconditioner based on deflation of `nev` eigenvalues chosen according to `which`. A partial Schur decomposition is computed (Matrix-Free), using the package `ArnoldiMethod.jl`, from which a projection is built. See the package `ArnoldiMethod.jl` for how to pass the proper options.
"""
function PrecPartialSchurArnoldiMethod(J, N, nev, which = LM(); tol = 1e-9, kwargs...)
	if J isa AbstractMatrix
		decomp, history = ArnoldiMethod.partialschur(J; nev = nev, tol = tol, which = which, kwargs...)
	else
		Jmap = LinearMap{Float64}(J, N, N ; ismutating = false)
		decomp, history = ArnoldiMethod.partialschur(Jmap; nev = nev, tol = tol, which = which, kwargs...)
	end
	return PrecPartialSchur(decomp.R, decomp.Q, inv(decomp.R), decomp.eigenvalues)
end

PrecPartialSchurArnoldiMethod(J::AbstractArray, nev, which = LM(); tol = 1e-9, kwargs...) = PrecPartialSchurArnoldiMethod(J, size(J)[1], nev, which ; tol = tol, kwargs...)


import Base.==

"""
Determine if `p, q` is a handle for `braid0`.
"""
function dehornoyishandle(braid0::BraidElement, p::Int, q::Int)::Bool
    braid = simplifybraid(braid0)
    if (braid.crossings[p][2] + braid.crossings[q][2] == 0) && (braid.crossings[p][1] == braid.crossings[q][1])
        v = braid.crossings[(p+1):(q-1)]
        j = braid.crossings[p][1]
        jp1 = filter(x -> x[1] == (j+1), v)
        jp1pos = filter(x -> x[2] > 0, jp1)
        jp1neg = filter(x -> x[2] < 0, jp1)
        if length(jp1pos) > 0 && length(jp1neg) > 0
            return false
        end
        for k in v
            if !(k[1] < (j-1) || k[1] > j) # Original
            #if !(k[1] < (j-1) || k[1] > (j+1)) # Ambiguity in the paper
                return false
            end
        end
        true
    else
        false
    end
end

"""
Find the first handle in `braid0` if one exists.
"""
function dehornoyfirsthandle(braid0::BraidElement)::Union{Tuple{Int,Int},Void}
    braid = simplifybraid(braid0)
    #for q = length(braid.crossings):-1:2
    for q = 2:length(braid.crossings)
    #for p = 1:(length(braid.crossings) - 1)
        #for q = (p+1):length(braid.crossings)
        #for p = 1:(q-1)
        for p = (q-1):-1:1
            if dehornoyishandle(braid0, p, q)
                return (p, q)
            end
        end
    end
    return nothing
end

"""
Reduce according to the alphabetic homomorphism of Dehornoy.
"""
function dehornoyalphareduce(braid0::BraidElement, handlep::Int, handleq::Int)::BraidElement
    nh = []
    j, e = braid0.crossings[handlep]
    for s in braid0.crossings[handlep:handleq]
        idx, ex = s
        if (idx != j) && (idx != (j+1))
            push!(nh, s)
        elseif idx == (j+1)
            push!(nh, (idx, -1*e))
            push!(nh, (j, ex))
            push!(nh, (idx, e))
        end
    end
    nw = cat(1, braid0.crossings[1:handlep-1], nh, braid0.crossings[handleq+1:end])
    b = BraidElement(braid0.group, nw)
    simplifybraid(b)
end

"""
Reduce according to the algorithm of Dehornoy.
"""
function dehornoyreduce(braid0::BraidElement)::BraidElement
    braid = unprettybraid(simplifybraid(braid0))
    while true
        handle = dehornoyfirsthandle(braid)
        if handle == nothing
            return prettybraid(simplifybraid(braid))
        else
            p, q = handle
            braid = dehornoyalphareduce(braid, p, q)
        end
    end
end

"""
Is a braid empty?
"""
function isempty(b::BraidElement)::Bool
    length(b.crossings) == 0
end

"""
Are two braids equal (in terms of ambient isotopy)?
"""
function ==(b1::BraidElement, b2::BraidElement)
    isempty(dehornoyreduce(b1 * inv(b2)))
end


###########################
# Function Cell Interface #
###########################
"""
	pca([;kwargs])

Constructs an untrained cell that when piped data inside, calculates the PCA
projection matrix of the input data. We assume that `d` is the dimesionality
of the input dataset, `n` the number of samples

# Keyword arguments (same as in `MultivariateStats`)
  * `method` is the PCA method used `:auto`, `:cov`, `:svd` (default `:auto` e.g. use `:cov` when the number of dimensions
  smaller than the number of samples, `svd` otherwise)
  * `maxoutdim` is the number of output dimensions, default `<number of variables>`
  * `pratio` is the ratio of variaces preserved in the principal subspace (default `0.99`)
  * `mean` is the mean vector, can be `nothing`, `0` or precomputed mean vector (default `nothing`)

Read the `MultivariateStats.jl` documentation for more information.
"""
pca(;kwargs...) = FunctionCell(pca, (), ModelProperties(), kwtitle("PCA", kwargs);kwargs...)

############################
# DataCell/Array Interface #
############################
"""
	pca(x, [;kwargs])

Trains a function cell that when piped data into will compute the principal components of
the observations based on the projection matrix extracted from `x`.
"""
# Training
pca(x::T where T<:CellData; kwargs...) = pca(getx!(x);kwargs...)
pca(x::T where T<:AbstractVector; kwargs...) = pca(mat(x, LearnBase.ObsDim.Constant{2}());kwargs...)
pca(x::T where T<:AbstractMatrix; kwargs...)  = begin

	ppadata = fit(MultivariateStats.PCA, getobs(x); kwargs...)

	# Build model properties
	modelprops = ModelProperties(size(ppadata.proj,1), size(ppadata.proj,2))

	# Returned trained cell
	FunctionCell(pca, Model(ppadata, modelprops), kwtitle("PCA",kwargs))
end

# Execution
pca(x::T where T<:CellData, model::Model{<:MultivariateStats.PCA}) =
	datacell(pca(getx!(x), model), gety(x))
pca(x::T where T<:AbstractVector, model::Model{<:MultivariateStats.PCA}) =
	pca(mat(x, LearnBase.ObsDim.Constant{2}()), model)
pca(x::T where T<:AbstractMatrix, model::Model{<:MultivariateStats.PCA}) =
	MultivariateStats.transform(model.data, getobs(x))

"""
	pcar(M::Cell), where `M=pca(x)`

Trains a cell based on a previously trained PCA function cell. When piped data into,
the returned output will be the reconstruction of the observations in the original
space. Obviously, the input data (e.g. data to be reconstructed) is expected to be
principal components extracted with `M`.

Read the `MultivariateStats.jl` documentation for more information.

# Examples
```
julia> P= rand(5,1000) |> pca(maxoutdim=4)
PCA: maxoutdim=4, 5 -> 4, trained

julia> PR = pcar(P)
PCA: reconstruct, 4 -> 5, trained

julia> ([1.0 2 3 4 5]' |> P ) |> pcar(P)
5×1 Array{Float64,2}:
 0.305402
 1.16768
 2.25037
 5.23595
 1.77198
```
"""
# Training for the reconstruction tranform (fixed cell)
pcar(x::T where T<:CellFunT{<:Model{<:MultivariateStats.PCA}}) = begin

	# Build model properties
	modelprops = ModelProperties(getx(x).properties.odim, getx(x).properties.idim)

	FunctionCell(pcar, Model(getx(x).data, modelprops),"PCA: reconstruct")
end

# Execution
pcar(x::T where T<:CellData, model::Model{<:MultivariateStats.PCA}) = datacell(pcar(getx!(x), model), gety(x))
pcar(x::T where T<:AbstractVector, model::Model{<:MultivariateStats.PCA}) = pcar(mat(x, LearnBase.ObsDim.Constant{2}()), model)
pcar(x::T where T<:AbstractMatrix, model::Model{<:MultivariateStats.PCA})::Matrix{Float64} = MultivariateStats.reconstruct(model.data, getobs(x))


using Luxor, Test

function polyconvextest()

    Drawing(100, 100, :png)
    origin()

    # polygons are convex
    for i in 3:16
        @test ispolyconvex(ngon(O, 100, i, vertices=true)) == true
    end

    # 4+ pointed stars aren't convex
    # 3 pointed star is effectively a triangle
    for i in 4:16
        @test ispolyconvex(star(O, 100, i, 0.5, 0, vertices=true)) == false
    end

    pts = [Point(x, 0) for x in 0:100]
    @test ispolyconvex(pts) == true

    @test finish() == true
end

function polyclockwisetest()

    Drawing(100, 100, :png)
    origin()

    # polygons are clockwise
    for i in 3:16
        @test ispolyclockwise(ngon(O, 100, i, vertices=true)) == true
    end
    # stars are too
    for i in 3:16
        @test ispolyclockwise(star(O, 100, i, 0.5, 0, vertices=true)) == true
    end

    pts = [Point(x, 0) for x in 0:100]
    @test ispolyclockwise(pts) == false

    pts = [Point(10cos(x), 10sin(x)) for x in 0:100]
    @test ispolyclockwise(pts) == true

    @test finish() == true
end

polyconvextest()
polyclockwisetest()

println("...finished convex and clockwise tests")


import LinearAlgebra

################################################################################
#                          Negative log-likelihood                             #
################################################################################

"""
Negative log-likelihood. ŷ should be a vector of normalized log probabilities.
"""
function nllloss(ŷ::AbstractMatrix, c::AbstractVector{<:Real})
    # loss = @tullio (+) r = ŷ[c[k], k]
    loss = LinearAlgebra.tr(NNlib.gather(ŷ', c))
    return -loss / length(c)
end

function ∇nllloss(Δ::Real, ŷ::AbstractMatrix, c::AbstractVector{<:Real})
    dŷ = zero(ŷ)
    val = -Δ / length(c)
    @tullio dŷ[c[k], k] = val
end

function ChainRulesCore.rrule(::typeof(nllloss), ŷ::AbstractMatrix, c::AbstractVector{<:Real})
    y = nllloss(ŷ, c)
    function nllloss_pullback(Δ)
        return NoTangent(), ∇nllloss(unthunk(Δ), ŷ, c), NoTangent()
    end
    return y, nllloss_pullback
end

################################################################################
#                               Cross Entropy                                  #
################################################################################

crossentropyloss(x::AbstractMatrix, c::Union{AbstractMatrix{<:Real}, AbstractVector{<:Real}}) =
    nllloss(logsoftmax(x), c)

################################################################################
#                                 Mean Squared Loss                            #
################################################################################

mseloss(inp::AbstractArray, target::AbstractArray) = mean((inp .- target) .^ 2)


# Test for solving div-grad system both complex and real and single and multiple rhs
using MUMPSjInv
using Test

include("getDivGrad.jl")
A = getDivGrad(32, 32, 16)
n = size(A, 1)

# REAL: test for single rhs
println("Test for real SPD matrix: one rhs")
rhs = randn(n)
x = solveMUMPS(A, rhs, 1)
err  =  norm(A * x - rhs) / norm(rhs)
@test err < 1e-14

# REAL: test for multiple rhs
println("Test for real SPD matrix: multiple rhs")
nrhs = 10
rhs = randn(n, nrhs)
x = solveMUMPS(A, rhs, 1)
err = zeros(nrhs)
for i= 1:nrhs
    err[i] =  norm(A * x[:, i] - rhs[:, i]) / norm(rhs[:, i])
end
@test eltype(x) == Float64
@test maximum(err) < 1e-14

# COMPLEX: test for multiple rhs
println("Test for complex symmetric matrix: one rhs")
r = rand(n)
A = A + im * spdiagm(0 => r)
rhs = randn(n) + im * randn(n)
x = solveMUMPS(A, rhs,1)
err=  norm(A * x - rhs) / norm(rhs)
@test eltype(x) == ComplexF64
@test err < 1e-14

# COMPLEX : test for multiple rhs
println("Test for complex symmetric matrix: multiple rhs")
nrhs = 10
rhs = randn(n, nrhs) + im * randn(n, nrhs)
x = solveMUMPS(A, rhs, 2)
err = zeros(nrhs)
for i = 1:nrhs
    err[i] = norm(A * x[:, i] - rhs[:, i]) / norm(rhs[:, i])
end
@test eltype(x) == ComplexF64
@test maximum(err) < 1e-14


using LinearAlgebra: Eigen
import LinearAlgebra: eigen, \, det, logdet, inv
import Base: +

"""
    eigen(K::AbstractKroneckerProduct)

Wrapper around `eigen` from the `LinearAlgebra` package. Performs Eigenvalue decompositon
on the matrices of a `AbstractKroneckerProduct` instances and returns an
`Eigen` type. The functions `size`, `\\`, `inv`, `det`,
and `logdet` are overloaded to efficiently work with this type.
"""
function eigen(K::AbstractKroneckerProduct)
    squarecheck(K)
    A, B = getmatrices(K)
    A_λ, A_Γ = eigen(A)
    B_λ, B_Γ = eigen(B)
    return Eigen(kron(A_λ, B_λ), kronecker(A_Γ, B_Γ))
end

+(E::Eigen, B::UniformScaling) = Eigen(E.values .+ B.λ, E.vectors)
+(A::UniformScaling, E::Eigen) = E + A

"""
    function collect(E::Eigen{<:Number, <:Number, <:AbstractKroneckerProduct})

Collects eigenvalue decomposition of a `AbstractKroneckerProduct` type into a
matrix.
"""
function collect(E::Eigen{<:Number, <:Number, <:AbstractKroneckerProduct})
    λ, Γ = E
    return Γ * Diagonal(λ) * Γ'
end

function \(E::Eigen{<:Number, <:Number, <:AbstractKroneckerProduct}, v::AbstractVector{<:Number})
    λ, Γ = E
    return Γ * (Diagonal(λ) \ (Γ' * v))
end

"""
    det(K::Eigen)

Compute the determinant of the eigenvalue decomp of aKronecker product.
"""
det(E::Eigen{<:Number, <:Number, <:AbstractKroneckerProduct}) = prod(E.values)

"""
    logdet(K::Eigen)

Compute the logarithm of the determinant of the eigenvalue decomp of aKronecker
product.
"""
logdet(E::Eigen{<:Number, <:Number, <:AbstractKroneckerProduct}) = sum(log, E.values)

"""
    inv(K::Eigen)

Compute the inverse of the eigenvalue decomp of aKronecker product. Returns
another type of `Eigen`.
"""
inv(E::Eigen{<:Number, <:Number, <:AbstractKroneckerProduct}) = Eigen(inv.(E.values), E.vectors)


using Mass2

srand(hash("mockpluses.jl"))
pg = TupacLikeTwoExponentialPulseGenerator(Int, Normal(1000,0.5))
getcleanpulse(pg,1000)
gettriggeredpulse!(pg,10)

function compute_whitenoise_filter(pulse, selection_good)
	filter = mean(pulse[selection_good])
	filter -= mean(filter) # 0 average
	normalization = (maximum(filter)-minimum(filter))./dot(filter, filter)
	filter*normalization
end
function filter1lag(pulse, filter)
	out = Array(Float64, length(pulse))
	for i=1:length(pulse)
		out[i] = dot(pulse[i], filter)
	end
	out
end
cleanpulses = [getcleanpulse(pg,1000) for i = 1:10000]
filter = compute_whitenoise_filter(cleanpulses, trues(length(cleanpulses)))
cleanpulses_filt_values = filter1lag(cleanpulses, filter)
@test abs(mean(cleanpulses_filt_values)-1000)<1
@test abs(std(cleanpulses_filt_values)-.655)<.1


# установка пакетов, выполняется один раз!
using Pkg
Pkg.add("CSV")
Pkg.add("RDatasets")
Pkg.add("Econometrics")
Pkg.add("DataFrames")
# то, что выше, делаем один раз :)

# матрички :)
a = [5 7 8 -1] # строка
b = [5 7 8 -1]' # столбец
b = [5, 7, 8, -1] # столбец
b = [5; 7; 8; -1] # столбец

y = [3 4 5 -2]' # столбец зависимой переменной
X = [1 3; 1 7; 1 5; 1 5] # матрица регрессоров (предикторов)
X'y
X'X
beta_hat = (X'X)^(-1) * X'y

y .* y # поточечное (поэлеметное умножение)

f(x) = 2x^3 + 3x - 7
f(4)

[2x + 5y for x = 2:8, y in [5, -3, 9]]

using RDatasets, Econometrics # аналог import numpy...

psid = dataset("Ecdat", "PSID") # Panel Study of Income Dynamics
psid

psid.Age
psid."Age"
psid[!, "Age"]
psid[2:5, ["Age", "Educatn"]]

model_0 = fit(EconometricModel,
    @formula(log(1 + Earnings) ~ Age + Educatn), psid)


using Knet, KnetModules

#=
First specify the model
(Param) is a type used to register
our trainable weights to global context.
=#
type Lreg <: KnetModule
    w::Param
    b::Param
end

# Now, let's write a constructor
function Lreg(input::Int)
    w = Param(0.1randn(1, input))
    b = Param(zeros(1,1))
    return Lreg(w, b)
end

# Last step is the declare the forward execution logic,
# by adding a method to KnetModules.forward function
import KnetModules.forward

# Here, ctx is a dependency where parameters are stored
forward(ctx, l::Lreg, x) =
    val(ctx, l.w) * x .+ val(ctx, l.b)

# Now, let's load the data and train our model
include(Knet.dir("data","housing.jl"))
x,y = housing()
model = Lreg(size(x,1))

# We built the entire forward pass logic, now the backward
loss(ypred, ygold) = mean(abs2, ygold .- ypred)
gradfn = grad(model, loss)

# Now, gradfn is a function (x, y) -> grads
lr = 0.1
for i = 1:10
    grads = gradfn(x, y)
    #=
    - params(model) returns the parameters in model
    - aval(p) returns the default active context
    - getgrad(p, g) returns grad of p in returned grads g

    params will return any parameter that lives in
    - Fields of type Param
    - Param fields of associative fields
    - Param fields of Array or Tuple
    - Param fields of other fields that inherit KnetModule
    This process is recursively continued
    =#
    for p in params(model)
        setval!(p, aval(p) - lr * getgrad(p, grads))
    end
    # @run evaluates the model
    ypred = @run model(x)
    println("Loss at iter $i ->", loss(ypred, y))
end

# The above training loop is very low-level, but it
# introduces you important accessor functions available

# let's now create the model
info("Resetting model")
model = switch_clean_ctx!(Lreg(size(x,1)))
# you don't have to call switch_clean_ctx!, but
# if you don't do it, old model's weights will continue to be stored
gradfn = grad(model, loss)
# optimizers are similar to Knet.optimizers, works with KnetModule objects
optims = optimizers(model, Sgd; lr=0.1)
for i = 1:10
    grads = gradfn(x, y)
    update!(model, grads, optims)
    ypred = @run model(x)
    println("Loss at iter $i -> ", loss(ypred, y))
end

info("Saving the final model")
save_module("lreg.jld", model)


export Disk,𝔻

# T is Real or Complex
# TT is (Real,Real) or Complex
"""
    Disk(c,r)

represents the disk centred at `c` with radius `r`.
"""
struct Disk{T} <: Domain{T}
    center::T
    radius::Float64
end

Disk(c,r) = Disk{typeof(c)}(c,r)
Disk(c::Real,r) = Disk(complex(c),r)
Disk(c::Integer,r) = Disk(float(c),r)
Disk(c::Complex{<:Integer},r) = Disk(float(c),r)

Disk() = Disk(Vec(0.,0.),1.)
Disk(::AnyDomain) = Disk(NaN,(NaN,NaN))

const 𝔻 = Disk()

isambiguous(d::Disk) = isnan(d.radius) && all(isnan,d.center)

polar(r,θ) = Vec(r*cos(θ),r*sin(θ))
ipolar(x,y) = Vec(sqrt(abs2(x)+abs2(y)), atan(y,x))
polar(rθ::Vec) = polar(rθ...)
ipolar(xy::Vec) = ipolar(xy...)

fromcanonical(D::Disk{T},xy::Vec) where {T<:Vec} = Vec((D.radius.*xy .+ D.center)...)
tocanonical(D::Disk{T},x,y) where {T<:Vec} = Vec((xy .- D.center)./D.radius)

checkpoints(d::Disk) = [fromcanonical(d,Vec(.1,.2243)),fromcanonical(d,Vec(-.212423,-.3))]

# function points(d::Disk,n,m,k)
#     ptsx=0.5*(1-gaussjacobi(n,1.,0.)[1])
#     ptst=points(PeriodicSegment(),m)
#
#     Float64[fromcanonical(d,x,t)[k] for x in ptsx, t in ptst]
# end

boundary(d::Disk) = Circle(d.center,d.radius)


@testset "Creation of an identity multiple" begin
    for n in [1, 2, 1000]
        In = IdentityMultiple(1.0I, n) # same as IdentityMultiple(UniformScaling(1.0), n))
        @test size(In) == (n, n)
        @test In[1, 1] == 1.0 && In[1, 2] == 0.0
    end
end

@testset "Operations between identity multiples" begin
    I2 = IdentityMultiple(UniformScaling(1.0), 2)
    I10 = IdentityMultiple(UniformScaling(1.0), 10)

    @test (I2 + I2).M == UniformScaling(2.0)
    @test (10.0 * I2).M == UniformScaling(10.0)
    @test (I2 * I2).M == I2.M

    @test_throws AssertionError I2 + I10
    @test_throws AssertionError I2 * I10
end

@testset "Create a continuous system with one matrix being a multiple of the identity" begin
    # See #40
    A, B = rand(4, 4), rand(4, 2)
    X = rand(Hyperrectangle, dim=4)
    U = rand(Ball2, dim=2)
    I4 = I(4)
    s = ConstrainedLinearControlContinuousSystem(A, I4, X, B*U);

    @test statedim(s) == 4
    @test stateset(s) == X
end

@testset "Matrix operations for identity multiple" begin
    I2 = I(2)
    @test getindex(I2, 1) == 1.
    @test getindex(I2, 2) == 0.
    @test getindex(I2, 3) == 0.
    @test getindex(I2, 4) == 1.
    x = I2 * rand(2)
    @test x isa Vector && length(x) == 2
    A = I2 * rand(2, 2)
    @test A isa Matrix && size(A) == (2, 2)
end


using CompScienceMeshes

export meshwaveguidepost

"""
    meshwaveguidepost(width, height, length, post_radius, delta)

Creates a mesh for a waveguide of width (along the x-axis) `width`, height (along
    the y-axis) `height` and length (along the z-axis) `length` by parsing a .geo script
    incorporating these parameters into the GMSH mesher.

A cylindrical post of radius `post_radius` is positioned vertically (i.e.
    parallel to y-axis) midway along the waveguide.

The target edge size is `delta`.
"""

function meshwaveguidepost(width, height, length, post_radius, delta)
    s =
"""

lc = $delta;

//--------------WAVEGUIDE POINTS--------------

Point(1) = {0.0, 0.0, 0.0, lc};
Point(2) = {0.0, $height, 0.0, lc};
Point(3) = {$width, $height, 0.0, lc};
Point(4) = {$width, 0.0, 0.0, lc};

Point(5) = {0.0, 0.0, $length, lc};
Point(6) = {0.0, $height, $length, lc};
Point(7) = {$width, $height, $length, lc};
Point(8) = {$width, 0.0, $length, lc};

//--------------POST POINTS-------------------

Point(9) = {($width)/2,0,($length)/2,lc};

Point(10) = {($width)/2,0,(($length)/2)-(($post_radius)/2),lc};
Point(11) = {(($width)/2)-(($post_radius)/2),0,($length)/2,lc};
Point(12) = {($width)/2,0,(($length)/2)+(($post_radius)/2),lc};
Point(13) = {(($width)/2)+(($post_radius)/2),0,($length)/2,lc};

Point(14) = {($width)/2,$height,($length)/2,lc};

Point(15) = {($width)/2,$height,(($length)/2)-(($post_radius)/2),lc};
Point(16) = {(($width)/2)-(($post_radius)/2),$height,($length)/2,lc};
Point(17) = {($width)/2,$height,(($length)/2)+(($post_radius)/2),lc};
Point(18) = {(($width)/2)+(($post_radius)/2),$height,($length)/2,lc};

//--------------WAVEGUIDE-POST CONTINUOUS SURFACE DEFINITION---------

//Part 1 - waveguide boundaries

Line(1)={1,2};
Line(2)={2,3};
Line(3)={3,4};
Line(4)={4,1};

Line(5)={5,6};
Line(6)={6,7};
Line(7)={7,8};
Line(8)={8,5};

Line(9)={1,5};
Line(10)={2,6};
Line(11)={3,7};
Line(12)={4,8};

//Part 2 - discontinuity boundaries

Circle(13) = {10,9,11};
Circle(14) = {11,9,12};
Circle(15) = {12,9,13};
Circle(16) = {13,9,10};

Circle(17) = {15,14,16};
Circle(18) = {16,14,17};
Circle(19) = {17,14,18};
Circle(20) = {18,14,15};

Line(21) = {10,15};
Line(22) = {11,16};
Line(23) = {12,17};
Line(24) = {13,18};

//Part 3 - creating skeleton

Line Loop(1)={-1,9,5,-10};

Line Loop(2)={-2,10,6,-11};
Line Loop(3)={17,18,19,20};

Line Loop(4)={-3,11,7,-12};

Line Loop(5)={-4,12,8,-9};
Line Loop(6)={-16,-15,-14,-13};

Line Loop(7) = {13,22,-17,-21};
Line Loop(8) = {14,23,-18,-22};
Line Loop(9) = {15,24,-19,-23};
Line Loop(10) = {16,21,-20,-24};

//Part 4 - creating surfaces

Plane Surface(1)={1};
Plane Surface(2)={3,2};
Plane Surface(3)={4};
Plane Surface(4)={6,5};

Ruled Surface(5) = {7};
Ruled Surface(6) = {8};
Ruled Surface(7) = {9};
Ruled Surface(8) = {10};

"""

    fn = tempname()
    io = open(fn, "w")
    try
        print(io, s)
    finally
        close(io)
    end

    # feed the file to gmsh
    fno = tempname()
    run(`gmsh $fn -2 -format msh -o $fno`)
    fdo = open(fno,"r")
    # include(Pkg.dir("CompScienceMeshes","src","gmsh.jl"))
    m = CompScienceMeshes.read_gmsh_mesh(fdo)
    close(fdo)
    rm(fno)
    rm(fn)

    return m

end

# Γ = meshwaveguidepost(22,22,50,3.81,0.8)
# RT = raviartthomas(Γ)
#
# include(Pkg.dir("CompScienceMeshes","examples","plotlyjs_patches.jl"))
# A = rand(numcells(Γ))
# p = patch(Γ, A)
# PlotlyJS.plot([p])


using CommutativeRings: len, GFImpl

@testset "iterate" begin
    @test length(collect(ZZ/7)) == 7
    @test length(collect(GF(3,5))) == 3^5
    @test length(collect(Monic((ZZ/2)[:x], 5))) == 2^5
    @test eltype(Monic((ZZ/2)[:x], 5)) == (ZZ/2)[:x]
end

@testset "ofindex $T" for (T, r) in ((Int, 0), (UInt, 0), (ZZ{Int8}, 0), (ZZ/77, 77), (GF(2,5),2^5), (QQ{Int}, 0))
    @test iszero(ofindex(0, T))
    @test isone(ofindex(1, T))
    @test ofindex(10, T) isa T
    @test len(T) == r
end

@testset "ofindex $T degree $n" for (T, n) in ((GF(2)[:z], 3),)
    @test ofindex(0, T, n) == monom(T, n)
    @test ofindex(1, T, n) == monom(T, n) + 1
    @test ofindex(10, T, n) isa T
end

@testset "ofindex Quotient" begin
    P = (ZZ/5)[:x]
    x = monom(P)
    Q = P / ( x^3 + x + 1)
    n = length(Q)
    @test length(Q) == 125
    @test sort(ofindex.(0:n-1, Ref(Q))) == sort(collect(Q))
end

@testset "factors" begin
    @test length(factors(12)) == 6
    @test factors(12) |> collect |> sort == [1, 2, 3, 4, 6, 12]
end


function get_nonlinearcoeff(res, mode, source)
    getω(source)/c * res.material.nonlinearindex * mode.corefraction(source) / mode.effectivearea(source)
end

function get_ωmesh(tmesh)
    dt_mesh = tmesh[2]-tmesh[1]
    num_pts = length(tmesh)
    ω = collect(2π*(-num_pts/2:num_pts/2-1)/(num_pts*dt_mesh) .|> THz2π)
    ω
end


# module model_fit

using DifferentialEquations, Flux, Optim, DiffEqFlux, DiffEqSensitivity, Plots
using BlackBoxOptim
using BenchmarkTools
using RecursiveArrayTools
# -------------------------------------------------------------------------

function generate_data(sol)
    t = collect(range(0,stop=10,length=101))

    solution = sum(sol(t),dims=1)
    # randomized = VectorOfArray([(sol(t[i]) + .02randn(size(sol,1))) for i in 1:length(t)])
    # dataset = convert(Array,randomized)
    randomized = solution .+ .01randn(size(solution))
    dataset = randomized
    display(scatter!(t,dataset'))
    return dataset
end

function generate_data(sol,t)

    solution = sum(sol(t),dims=1)
    # randomized = VectorOfArray([(sol(t[i]) + .02randn(size(sol,1))) for i in 1:length(t)])
    # dataset = convert(Array,randomized)
    randomized = solution .+ .01randn(size(solution))
    dataset = randomized
    display(scatter!(t,dataset'))
    return dataset
end

# function run_model(model,p,data,densities,constants,t) # version of run function with multiple models
#
#     # Σ_sol_stack = Array{Any}(undef, 1, size(data,2))
#     Σ_sol_stack = zeros(1, size(data,2))
#     # display(1:size(data,1))
#     for i in 1:size(data,1)
#         # run model with given densities
#         densities_i = densities[i,:]
#         f = (du,u,p,t) -> model(du,u,10 .^ p,t,densities_i,constants)
#         tmp_prob = ODEProblem(f,u₀,tspan,p)
#         # tmp_sol = solve(tmp_prob,Rosenbrock23(),saveat=t)
#         tmp_sol = solve(tmp_prob,Rodas4(),saveat=t)
#         # stack temp solution
#         # stack Σ of solution across n species
#         Σ_sol = sum(Array(tmp_sol),dims=1)
#         Σ_sol_stack = vcat(Σ_sol_stack,Σ_sol)
#     end
#     return Σ_sol_stack[2:end,:]
# end

function run_model(model,p,data,densities,constants,t) # version of run function with multiple models

    # Σ_sol_stack = Array{Any}(undef, 1, size(data,2))
    # Σ_sol_stack = zeros(1, size(data,2))

    f = (du,u,p,t) -> model(du,u,10 .^ p,t,densities,constants)
    tmp_prob = ODEProblem(f,u₀,tspan,p)
    # tmp_sol = solve(tmp_prob,Rosenbrock23(),saveat=t)
    # tmp_sol = solve(tmp_prob,Rodas4(),saveat=t)
    tmp_sol = solve(tmp_prob,Tsit5(),saveat=t)
    # stack temp solution
    # stack Σ of solution across n species
    Σ_sol = sum(Array(tmp_sol),dims=1)

end

function loss(p,data,model,dens,cons,t)
    Σ_sol = run_model(model,p,data,dens,cons,t)
    # display(Σ_sol)
    # sum(abs2, (Σ_sol - data) .* t), Σ_sol
    sum(abs2, (Σ_sol - data))#, Σ_sol
end

function callback(p,l,tmp_sol,data,model,dens,cons,t)
    @show l

    Σ_sol = run_model(model,p,data,dens,cons,t)

    fig = plot(t,Σ_sol')
    scatter!(fig,t,data',legend=:bottomright)
    display(fig)

    false
end

function scan_param_space(p,data,model,dens,cons,t,lb,ub)
        #create rand mesh
        n = 10000
        ps = rand(n,size(p,1)) .* 9 .- 7
        errors = zeros(n)
        for i in 1:n
            # error, a = loss(ps[i,:],data,model,dens,cons,t)
            error = loss(ps[i,:],data,model,dens,cons,t)
            errors[i] = error
            # @show error
        end
        min_ind = argmin(errors)
        min_error = minimum(errors)
        @show findmin(errors)
        p_best = ps[min_ind,:]
        @show p_best
        if size(p,1) == 2
            fig = scatter(ps[:,1],ps[:,2],log10.(errors),
            zcolor = log10.(errors),camera=[0,90])
            scatter!([p_best[1]],[p_best[2]],[log10(min_error)],markercolor=:green,markersize=5)
            display(fig)
        end
        return p_best
end

function optimization(model,data,dens,cons,options,t)
# fit function - given model, data, densities, options get best parameters

    loss_in = (p) -> loss(p,data,model,dens,cons,t)
    callback_in = (p,l,tmp_sol) -> callback(p,l,tmp_sol,data,model,dens,cons,t)

    # @unpack lb,ub = options
    lb,ub = options
    # f = (du,u,p,t) -> model(du,u,p,t,dens,cons) # optimize in log space
    # prob = ODEProblem(f,u₀,tspan,p)

    # start = scan_param_space(pinit,data,model,dens,cons,t,lb,ub) # !! change behavior so that it doesnt need pinit
    # use ADAM first to find ballparck solution
    # res = DiffEqFlux.sciml_train(loss_in,pinit,ADAM(0.001),cb=callback_in,maxiters=300)
    # res = DiffEqFlux.sciml_train(loss_in,pinit,BFGS(initial_stepnorm = 0.0001),cb=callback_in)
    start = [-4.,-1.]
    res = DiffEqFlux.sciml_train(loss_in,start,BFGS(),cb=callback_in)
    # res = DiffEqFlux.sciml_train(loss_in,start,NelderMead(),cb=callback_in)
    @show res
    rates = res.minimizer
    # rates = start

    f = (du,u,p,t) -> model(du,u,p,t,dens,cons) # optimize in log space
    prob = ODEProblem(f,u₀,tspan,10 .^ rates)
    sol = solve(prob,Rodas4())
    h = plot!(sol)
    display(h)
    display([model,rates,start,dens,cons])
    return rates
end

# end


drop(gen::Function, n::Integer) = (for _ in 1:n gen() end; gen)
take(gen::Function, n::Integer) = collect(gen() for _ in 1:n)

function pgen(n::Number)
    x = 0
    return () -> (x += 1) ^ n
end

function genfilter(g1::Function, g2::Function)
    local r1
    local r2 = g2()
    return () -> begin
        r1 = g1()
        while r2 <  r1 r2 = g2() end
        while r1 == r2 r1 = g1() end
        return r1
    end
end

@show take(drop(genfilter(pgen(2), pgen(3)), 20), 10)


## Generate a funky signal
fs  = 5e6
dt  = 1/fs
t   = collect(0:dt:0.00007)
f   = 1e5
tau = 0.00001
y   = sin.(2*pi*f*t) .* exp.(-t/tau)
y[10] = 2;
y = vcat(y, reverse(y))
t = collect(0:dt:dt*(length(y)-1))
# add some noise
zg = 0.01*randn(length(y))
y .= y .+ zg
thr = 0.05 #absolute threshold

## Basic usage:
# 1. base function
@time tc,yc, idx2save = AdaptiveSampling.anchor_point(t,y,thr)
@info "compression_ratio" compression_ratio = length(t)/length(tc)
# 2. split the input array in chunks, compress each of them and concatenate -> faster
@time tc,yc, idx2save = AdaptiveSampling.anchor_point(t,y,thr,true)
@info "compression_ratio" compression_ratio = length(t)/length(tc)
# 3. no split of the input array, same as 1.
@time tc,yc, idx2save = AdaptiveSampling.anchor_point(t,y,thr,false)
@info "compression_ratio" compression_ratio = length(t)/length(tc)

## Generating a gif
using Plots
# The anchor_point_gif function was commented because it's just for animation purposes
# and it would burden the package with a heavy dependency: Plots
# don't use this for real compression tasks! it's slow because of plotting
tc,yc, idx2save, anim = AdaptiveSampling.anchor_point_gif(t,y,thr)

gif(anim, "anchor_point.gif", fps = 20)


using FixedSizeArrays
using PlotlyJS
include("../src/CHull2D.jl")

#
# LineSegment Intersection
#

# Simple example
p1 = Point(0.0, 0.0)
p2 = Point(0.5, 0.5)
p3 = Point(0.0, 0.5)
p4 = Point(0.5, 0.0)
l1 = CHull2D.LineSegment(p1, p2)
l2 = CHull2D.LineSegment(p3, p4)
println("Should be true", ":", intersect(l1, l2))

# TODO: Should be able to handle Inf instead of 1e12
p5 = Point(0.0, 0.25)
p6 = Point(1e12, 0.25)
l3 = CHull2D.LineSegment(p5, p6)
println("Should be true", ":", intersect(l1, l3))

#
# Convex Hull with Quadrant
#
p = [Point(randn(), randn()) for i=1:25]
ch = CHull2D.convexhull(p)
q = CHull2D.Quadrant(Point(0.0, 0.0), Inf, Inf)
chiq = intersect(ch, q)

t1 = scatter(;x=[foo[1] for foo in chiq], y=[foo[2] for foo in chiq])
t2 = scatter(;x=[foo[1] for foo in ch.extremepoints], y=[foo[2] for foo in ch.extremepoints])
Plot([t1, t2])


module TestGramian
using CovarianceFunctions
using Test
using LinearAlgebra
using ToeplitzMatrices
using CovarianceFunctions
using CovarianceFunctions: PeriodicInput, LazyMatrixSum
using BlockFactorizations

@testset "Gramian properties" begin
    n = 8
    x = randn(Float64, n)
    k = CovarianceFunctions.EQ()
    G = gramian(k, x)

    @test issymmetric(G)
    M = Matrix(G)
    @test M ≈ AbstractMatrix(G)
    @test G' ≈ M'
    @test transpose(G) ≈ transpose(M)
    @test M ≈ k.(x, x')
    @test issymmetric(G) && issymmetric(M)
    @test ishermitian(G) && ishermitian(M)
    @test isposdef(G) && isposdef(M) # if M is very low-rank, this can fail -> TODO: ispsd()

    G = gramian(k, x, x .+ randn(size(x)))
    M = Matrix(G)
    @test !issymmetric(G) && !issymmetric(M)
    @test !isposdef(G) && !isposdef(M)

    G = gramian(k, x, copy(x))
    @test issymmetric(G) # testing that value check also works
    @test isposdef(G)

    # testing standard Gramian
    G = gramian(x)
    @test G ≈ dot.(x, x')

    x = randn(Float64, n)
    k = CovarianceFunctions.EQ()
    G = gramian(k, x)
    # type promotion
    @test k(x[1], x[2]) isa typeof(G[1, 2]) # test promotion in inner constructor
    @test k(x[1], x[2]) isa eltype(G) # test promotion in inner constructor
    x = tuple.(x, x)
    G = gramian(k, x)
    @test k(x[1], x[2]) isa typeof(G[1,1])  # with tuples (or arrays)
    @test k(x[1], x[2]) isa eltype(G)

    # adding diagonal to Gramian is done lazily
    D = 1e-6I(n)
    @test D + G isa LazyMatrixSum
    @test G + D isa LazyMatrixSum

    # rectangular multiply
    x = randn(n)
    y = randn(2n)
    G = gramian(k, x, y)
    @test size(G) == (n, 2n)
    a = randn(2n)
    b = G*a
    @test length(b) == n
    @test b ≈ Matrix(G) * a

    y = randn(2n)
    G = gramian(k, x, y)
    @test size(G) == (n, 2n)
    p = 3
    A = randn(2n, p)
    B = G*A
    @test size(B) == (n, p)
    @test B ≈ Matrix(G) * A

    # testing scalar indexing
    G = gramian(k, x, y)
    for i in 1:n
        for j in 1:n
            @test G[i, j] ≈ k(x[i], y[j])
        end
    end

    # testing non-scalar indexing
    i, j = 2:n-1, 3:2n-4
    Gij = G[i, j]
    @test Gij isa AbstractMatrix
    @test Gij ≈ gramian(k, x[i], y[j])

    i, j = 2:n-1, 2
    Gij = G[i, j]
    @test Gij isa AbstractVector
    @test Gij ≈ vec(gramian(k, x[i], [y[j]]))

    Gij = G[j, i]
    @test Gij isa AbstractVector
    @test Gij ≈ vec(gramian(k, [x[j]], y[i]))
end

@testset "Gramian factorization" begin
    rtol = 1e-6
    k = CovarianceFunctions.EQ()
    n = 64
    x = randn(n)
    G = gramian(k, x)
    F = factorize(G)
    @test F isa CholeskyPivoted
    @test rank(F) < n ÷ 2 # this should be very low rank
    @test isapprox(Matrix(F), G, rtol = rtol)

    # pivoted cholesky
    G = gramian(k, x)
    F = cholesky(G, Val(true), check = false, tol = 1e-6)
    @test F isa CholeskyPivoted
    @test isapprox(Matrix(F), G, rtol = rtol)

    # regular cholesky
    k = CovarianceFunctions.Exponential() # does not yield low rank matrix
    G = gramian(k, x)
    F = cholesky(G)
    @test F isa Cholesky
    @test isapprox(Matrix(F), G, rtol = rtol)

    # testing gramian of matrix-valued anonymous kernel
    d = 3
    A = randn(d, d)
    k = (x, y)->A
    G = gramian(k, x)
    GA = Matrix(G)
    a = randn(d*n)
    @test size(G) == (d*n, d*n)
    @test G isa BlockFactorization
    @test G*a ≈ GA*a

    # adding diagonal by default maintains laziness
    @test I(d*n) + G isa LazyMatrixSum
    @test G + I(d*n) isa LazyMatrixSum

    # same thing for block diagonal matrices
    D = BlockFactorization(Diagonal([randn(d, d) for _ in 1:n]))
    @test D + G isa LazyMatrixSum
    @test G + D isa LazyMatrixSum
end

@testset "toeplitz structure" begin
    n = 32
    x = range(-1, 1, n)
    k = CovarianceFunctions.EQ()
    G = gramian(k, x)
    @test G isa SymmetricToeplitz
    @test size(G) == (n, n)
    @test G*x ≈ Matrix(G)*x

    @test Matrix(G) ≈ Matrix(Gramian(k, x))
    G = gramian(k, x, PeriodicInput()) # periodic boundary conditions
    @test G isa Circulant
    @test size(G) == (n, n)
    @test G*x ≈ Matrix(G)*x

    k = (x, y)->CovarianceFunctions.EQ()(x, y)
    G = gramian(k, x)
    @test G isa Gramian # if we can't infer the input trait of k, falls back to lazy representation
    @test Matrix(G) ≈ Matrix(Gramian(k, x))

    # testing that we can hook anonymous kernel into the system
    CovarianceFunctions.input_trait(::typeof(k)) = CovarianceFunctions.IsotropicInput()
    G = gramian(k, x)
    @test G isa SymmetricToeplitz
    @test Matrix(G) ≈ Matrix(Gramian(k, x))

    # testing non-symmetric Toeplitz systems
    y = x .+ randn() # shifted version of x with same step length
    G = gramian(k, x, y)
    @test G isa Toeplitz
    @test Matrix(G) ≈ Matrix(Gramian(k, x, y))

    y = range(-1, 1, n÷2) # if y has different step length, defaults to gramian
    G = gramian(k, x, y)
    @test G isa Gramian
    @test Matrix(G) ≈ Matrix(Gramian(k, x, y))

end # test solves and multiplications

end # TestGramian


#Exchangers and transporters on the inner mitochondirla membrane
using Parameters
include("common.jl")

# NHE (mitochondrial Na-H exchanger) parameters
@with_kw struct NHEParams
	# Intracelllular proton
	H_I = 1E-4
    # NHE forward rate constant
    K1_P = 0.0252
    # NHE backward rate constant
    K1_M = 0.0429
    # NHE forward rate constant
    K4_P = 0.16
    # NHE backward rate constant
    K4_M = 0.0939
    # Na dissociation constant
    KM_NA = 24.0
    # H+ dissociation constant
    KM_H = 1.584E-4
	# Proton inhibitory constant
	H_KI = 3.02E-6
    # Hill coefficient for proton binding
    NI = 3
    # NHE concentration
    ET = 0.00785
    # Dependence factor of mitochondrial proton
    DENOM = 1 + (H_KI / H_I)^ NI
	VMAX = ET / DENOM
end

# Mitochondrial Na-H exchanger (NHE). scalar version
function vnhe(na_m, na_i, h_m, h_i, K1_P, K1_M, K4_P, K4_M, KM_H, KM_NA, VMAX)
	β₁⁺ = K1_P * _mm(na_m * _mm(KM_H, h_m), KM_NA)
	β₁⁻ = K1_M * _mm(na_i * _mm(KM_H, h_i), KM_NA)
	β₂⁺ = K4_P * _mm(h_i * _mm(KM_NA, na_i), KM_H)
	β₂⁻ = K4_M * _mm(h_m * _mm(KM_NA, na_m), KM_H)
    vNHE = VMAX * (β₁⁺ * β₂⁺ - β₁⁻ * β₂⁻) / (β₁⁺ + β₂⁺ + β₁⁻ + β₂⁻)
end

function vnhe(na_m, na_i, h_m, h_i, pNHE::NHEParams)
	@unpack K1_P, K1_M, K4_P, K4_M, KM_H, KM_NA, VMAX = pNHE
	vNHE = vnhe(na_m, na_i, h_m, h_i, K1_P, K1_M, K4_P, K4_M, KM_H, KM_NA, VMAX)
end

# PIC (phosphate carrier)
@with_kw struct PICParams
    # PIC concentration (mM)
    ET = 4.9
    # Cytosolic Pi binding constant
    KM_PI_IN = 11.06
    # Mitochondrial Pi binding constant
    KM_PI_MT = 11.06
    # Cytosolic OH binding constant
    KM_OH_IN = 4.08E-5
    # Mitochondrial OH binding constant
    KM_OH_MT = 4.08E-5
    # Max forward rate of PIC (mM/ms)
    VF = 90 / 60E3
    VF_MAX = VF * ET
    # Max backward rate of PIC (mM/ms)
    VB = 90 / 60E3
    VB_MAX = VB * ET
end

# Phosphate carrier (PiC)
# Manual correction from dissociation
function _vpic(H₂PO₄⁻_m, H₂PO₄⁻_i, h_m, h_i, KM_PI_IN, KM_PI_MT, KM_OH_IN, KM_OH_MT, VF_MAX, VB_MAX)
	f_pi = H₂PO₄⁻_i / KM_PI_IN
    f_pm = H₂PO₄⁻_m / KM_PI_MT
    f_ohi = _oh(h_i) / KM_OH_IN
    f_ohm = _oh(h_m) / KM_OH_MT
	ϕf = f_pi * f_ohm
	ϕb = f_pm * f_ohi
    den = 1 + f_pi + f_pm + f_ohi + f_ohm + ϕf + ϕb
    vPIC = (VF_MAX * ϕf - VB_MAX * ϕb) / den
end

# Auto correction from dissociation
function vpic(pi_m, pi_i, h_m, h_i, KM_PI_IN, KM_PI_MT, KM_OH_IN, KM_OH_MT, VF_MAX, VB_MAX)
	H₂PO₄⁻_m = _H2PO4(pi_m, h_m)
	H₂PO₄⁻_i = _H2PO4(pi_i, h_i)
	vPIC = _vpic(H₂PO₄⁻_m, H₂PO₄⁻_i, h_m, h_i, KM_PI_IN, KM_PI_MT, KM_OH_IN, KM_OH_MT, VF_MAX, VB_MAX)
end

function vpic(pi_m, pi_i, h_m, h_i, pPIC::PICParams)
	@unpack KM_PI_IN, KM_PI_MT, KM_OH_IN, KM_OH_MT, VF_MAX, VB_MAX = pPIC
	vPIC = vpic(pi_m, pi_i, h_m, h_i, KM_PI_IN, KM_PI_MT, KM_OH_IN, KM_OH_MT, VF_MAX, VB_MAX)
end

# Mitochondrial NCE parameters
@with_kw struct MNCEParams
    VMAX = 4.665E-5
    B = 0.5
    KM_NA = 9.4
    KM_CA = 3.75E-4
    N = 3
end

# Mitochondrial Ca uniporter
@with_kw struct MCUParams
	VMAX = 1.2295E-4
    DPSI_OFFSET = 91
    KACT = 3.8E-4
    KTRANS = 0.019
    L = 110
    N = 2.8
end

# Mitochondrial Calcium uniporter (MCU) flux (mM/ms)
# scalar version
function vuni(ca_i, dpsi, DPSI_OFFSET, KTRANS, KACT, N, L, VMAX)
	v2frt = 2 * F_RT * (dpsi - DPSI_OFFSET)
    f_ca = ca_i / KTRANS
    f_ca_p1 = f_ca + one(f_ca)
	f_act =  L * (_mm(KACT, ca_i) ^ N)
	vUni = VMAX * f_ca * f_ca_p1^3 * v2frt / ((f_ca_p1^4 + f_act) * (-expm1(-v2frt)))
end

function vuni(ca_i, dpsi, pMCU::MCUParams)
	@unpack DPSI_OFFSET, KTRANS, KACT, N, L, VMAX = pMCU
	vUni = vuni(ca_i, dpsi, DPSI_OFFSET, KTRANS, KACT, N, L, VMAX)
end

# Mitochondrial NCE (mNCE) flux (mM/ms)
function vnaca(ca_i, ca_m, na_i, dpsi, KM_NA, KM_CA, N, B, VMAX, DPSI_OFFSET)
	f_na = _mm(na_i, KM_NA) ^ N
    f_ca = _mm(ca_m, KM_CA)
	ϕ_ca = ca_m / ca_i
	vNaCa = VMAX * _ra(B * (dpsi - DPSI_OFFSET)) * ϕ_ca * f_na * f_ca
end

function vnaca(ca_i, ca_m, na_i, dpsi, pNCE::MNCEParams, pMCU::MCUParams)
	@unpack DPSI_OFFSET = pMCU
	@unpack KM_NA, KM_CA, N, B, VMAX = pNCE
	vNaCa = vnaca(ca_i, ca_m, na_i, dpsi, KM_NA, KM_CA, N, B, VMAX, DPSI_OFFSET)
end


abstract type Sense end
struct Max<:Sense end
struct Min<:Sense end

mutable struct NonDomPoints{T}
	xe::Vector{Vector{Int}}
	yn::Vector{Tuple{Int,Int}}
	λ::Tuple{Int,Int}
	nadirs::Vector{Tuple{Int,Int}}
	worst_nadir::Tuple{Int,Int}
	function NonDomPoints(t::Type{T}, xe, yn) where T<:Sense
		@assert issorted(yn, by = y-> y[1])
		λ = round(Int, yn[1][2] - yn[end][2]), round(Int, yn[end][1] - yn[1][1])
		λ = λ[1]÷gcd(λ[1], λ[2]), λ[2]÷gcd(λ[1], λ[2])
		res = new{T}([round.(Int, x) for x in xe], [round.(Int, x) for x in yn], λ, [])
		update_ln!(res)
		res
	end
end

function update_ln!(p::NonDomPoints{Max})
	p.nadirs = [(p.yn[i][1], p.yn[i+1][2]) for i = 1:length(p.yn)-1]
	if isempty(p.nadirs)
		p.worst_nadir = p.yn[1].+1
		return p
	end
	p.worst_nadir = p.nadirs[1].+1
	for n in p.nadirs
		if sum((n.+1).*p.λ) < sum(p.worst_nadir.*p.λ)
			p.worst_nadir = n.+1
		end
	end
	p
end

function update_ln!(p::NonDomPoints{Min})
	p.nadirs = [(p.yn[i+1][1], p.yn[i][2]) for i = 1:length(p.yn)-1]
	if isempty(p.nadirs)
		p.worst_nadir = p.yn[1].-1
		return p
	end
	p.worst_nadir = p.nadirs[1].-1
	for n in p.nadirs
		if sum((n.-1).*p.λ) > sum(p.worst_nadir.*p.λ)
			p.worst_nadir = n.-1
		end
	end
	p
end

nadirs(p::NonDomPoints) = p.nadirs
worst_nadir(p::NonDomPoints) = p.worst_nadir
isfathomable(z, p::NonDomPoints{Min}) = z > sum(worst_nadir(p).*p.λ)
isfathomable(z, p::NonDomPoints{Max}) = z < sum(worst_nadir(p).*p.λ)
local_nadir(p::NonDomPoints{Min}) = (p.yn[end][1], p.yn[1][2])
local_nadir(p::NonDomPoints{Max}) = (p.yn[1][1], p.yn[end][2])

function isweaklydominated(z1, z2, p::NonDomPoints{Min})
	for (v1, v2) in p.yn
		if v1 <= z1 && v2 <= z2
			return true
		end
	end
	false
end

function isweaklydominated(z1, z2, p::NonDomPoints{Max})
	for (v1, v2) in p.yn
		if v1 >= z1 && v2 >= z2
			return true
		end
	end
	false
end

Base.length(p::NonDomPoints) = length(p.yn)

function Base.push!(p::NonDomPoints{S}, x, y) where S<:Sense
	#@assert !isweaklydominated(y..., p) "y : $y \n p : $p"
	if isweaklydominated(y..., p)
		return p
	end
	ind = searchsortedfirst(p.yn, y, by = y -> first(y))
	insert!(p.xe, ind, x)
	insert!(p.yn, ind, y)
	inds_delete = find(x->dominates(S, y, x), p.yn)

	deleteat!(p.xe, inds_delete)
	deleteat!(p.yn, inds_delete)
	update_ln!(p)
end


# Based on https://www.tandfonline.com/doi/abs/10.1080/01966324.1987.10737221
# TODO: make a package out of this function. Kombinator?

function _ossb_exact_convex_combination(instance::CombinatorialInstance{T}, t::Dict{T, Float64}) where T
  # Input: a combinatorial problem, a (continuous) solution to the OSSB problem,
  # *scaled so that v=1*. This ensures that the point is within the polytope.

  # The solution to return: a set of solutions (each of them being a set of arms),
  # their weights in the convex combination.
  # This representation assumes that all vertices of the polytope are actual solutions.
  points = Vector{Int}[]
  weights = Float64[]

  # Convert the current point t as a vector (this is required to perform linear algebra).
  arms = all_arm_indices(instance) # Assimilate an arm and its index.
  idx_arm = Dict(i => arms[i] for i in 1:length(arms))
  d = instance.n_arms
  x = zeros(Float64, d)
  for (arm, weight) in t
    x[idx_arm[arm]] = weight
  end

  # Counters for the algorithm.
  point = t # What remains to decompose.
  δ = 1.0 # Weight still to assign.

  # Start the algorithm.

  # Find the currently tight constraints to create the matrix G.
  _, tight_matrix = _ossb_find_tight_constraints(instance, x, idx_arm)

  while rank(tight_matrix) < d
    # Find a new vertex by decomposing (what's left of) x.
    y_bar = x
    while rank(tight_matrix) < d
      # Find a solution dir such that G dir = 0 with dir ≠ 0.
      if length(tight_matrix) == 0 # Easier to compute than rank(G) == 0
        # No constraint to satisfy. Take an arbitrary "solution". This is a pathologically
        # underdetermined system.
        # Normalise by dividing by the norm.
        dir = ones(d) ./ sqrt(d)
      else
        # This system is underdetermined, so that the usual syntax A\b is not guaranteed to work
        # (Julia will return a minimum-norm solution in the underdetermined case, which is very likely zero).
        dir = vec(nullspace(tight_matrix)[:, 1])

        # Also normalise this vector, in case the previous computation gave strange results.
        n = norm(dir, 2)
        @assert n > 1.0e-6 # Just in case...
        dir ./= n
      end

      # Compute a point y = y_bar - α d that touches a face of the polyhedron (at least one more
      # constraint is tight there).
      α_val, y_bar = _ossb_touch_facet(copy(instance), y_bar, dir, idx_arm)

      # Update tight_matrix for this vertex-to-be (or already a vertex).
      _, tight_matrix = _ossb_find_tight_constraints(instance, y_bar, idx_arm)
    end

    # Retrieve that point from x to continue the decomposition.
    dir = x - y_bar
    # @show dir
    γ_val, x = _ossb_touch_facet(copy(instance), x, dir, idx_arm)
    δ_scale = γ_val / (1 + γ_val)

    # This point y is a vertex that plays a role in the convex decomposition.
    push!(points, round.(Int, y_bar))
    push!(weights, δ_scale * δ)

    δ *= 1 - δ_scale

    # Update tight_matrix for the rest of the point to decompose.
    _, tight_matrix = _ossb_find_tight_constraints(instance, x, idx_arm)
  end

  # The last iterate x must be an extreme point, as the rank of the tight constraints is d.
  push!(points, round.(Int, x))
  push!(weights, δ)

  return points, weights
end

function _ossb_find_tight_constraints(instance::CombinatorialInstance{T}, x::Vector{Float64}, idx_arm::Dict{T, Int}) where T
  # Assumption: point within the polytope; the combinatorial instance has an LP formulation.

  # Get the LP formulation, without objective function (_).
  m, _, vars = get_lp_formulation(instance, Dict(arm => 0.0 for arm in all_arm_indices(instance)))

  # vars is a dictionary of arms to variables, but we will need the reverse operation.
  var_to_arm = Dict(values(vars) .=> keys(vars))

  # Iterate over the constraints.
  tight_single = JuMP.ConstraintRef[]
  tight_linear = JuMP.ConstraintRef[]
  for (F, S) in list_of_constraint_types(m)
    # Ignore the case where this is an integrality constraint (it does not add bounds). Simple bounds are required!
    if S == MOI.Integer
      continue
    end

    if S == MOI.ZeroOne
      @assert F == JuMP.VariableRef # TODO: is MOI.SingleVariable possible here? Other functions should not be allowed.
      for c in all_constraints(m, F, S)
        lhs = value(jump_function(constraint_object(c)), vr -> x[var_to_arm[vr]])
        if abs(lhs) <= 1.0e-6 || abs(lhs - 1.0) <= 1.0e-6
          push!(tight_single, c)
        end
      end
    else
      for c in all_constraints(m, F, S)
        # If the constraint is tight at t, keep it.
        lhs = value(c, vr -> x[var_to_arm[vr]])
        rhs = MOI.constant(moi_set(constraint_object(c)))
        if abs(lhs - rhs) <= 1.0e-6
          push!(tight_linear, c)
        end
      end
    end
  end

  # Retrieve the coefficients for each constraint.
  var_to_idx = Dict(vars[arm] => idx for (arm, idx) in idx_arm) # Must be rebuild for each new JuMP model.

  tight_to_coefficients = Dict{JuMP.ConstraintRef, Vector{Float64}}()
  sizehint!(tight_to_coefficients, length(tight_single) + length(tight_linear))

  nvars = MOI.get(m, MOI.NumberOfVariables())

  for c in tight_single
    tight_to_coefficients[c] = spzeros(nvars)
    var = jump_function(constraint_object(c))
    tight_to_coefficients[c][var_to_idx[var]] = 1.0
  end
  for c in tight_linear
    tight_to_coefficients[c] = spzeros(nvars)
    for (var, coeff) in jump_function(constraint_object(c)).terms
      tight_to_coefficients[c][var_to_idx[var]] = coeff
    end
  end

  # The returned set may very well be empty, if the solution hit the bound on v.
  tight_matrix = (length(tight_to_coefficients) > 0) ? vcat(collect(values(tight_to_coefficients))'...) : zeros(0, 0)
  return tight_to_coefficients, tight_matrix
end

function _ossb_touch_facet(instance::CombinatorialInstance{T}, x::Vector{Float64}, dir::Vector{Float64}, idx_arm::Dict{T, Int}) where T
  # Maximise α such that ȳ = y - 2 α dir ∈ polytope(instance), where y = x - dir.
  # Return this optimum α and ȳ = y - 2 α⋆ dir.
  d = length(idx_arm)

  # First, get the standard formulation, without objective function (_).
  msub, _, varssub = get_lp_formulation(instance, Dict(arm => 0.0 for arm in all_arm_indices(instance)))
  var_to_idx = Dict(varssub[arm] => idx for (arm, idx) in idx_arm)

  # Then, build a similar model for this specific problem.
  # TODO: could also do it by inspection of all constraints.
  m = Model(instance.solver.solver) # TODO: only for LP-based solvers.
  # TODO: write a dedicated solver for this (highly simplified) problem? It could work by simply shortening an interval, iterating through the constraints.
  set_silent(m) # TODO: maybe allow to configure this?
  @variable(m, α)
  @objective(m, Max, α)

  # Iterate through the constraints and build the new ones.
  y_bar = x .+ α .* dir

  for (F, S) in list_of_constraint_types(msub)
    # Ignore the case where this is an integrality constraint (it does not add bounds).
    if S == MOI.Integer
      continue
    end

    if S == MOI.ZeroOne
      @assert F == JuMP.VariableRef # TODO: is MOI.SingleVariable possible here? Other functions should not be allowed.
      for c in all_constraints(msub, F, S)
        var = jump_function(constraint_object(c)) # A simple variable reference.
        lhs = y_bar[var_to_idx[var]]
        @constraint(m, 0 <= lhs <= 1)
      end
    else
      for c in all_constraints(msub, F, S)
        lhs = AffExpr() + jump_function(constraint_object(c)).constant # The constant should be zero (only present in MOI set).
        for (var, coeff) in jump_function(constraint_object(c)).terms
          lhs += coeff * y_bar[var_to_idx[var]]
        end
        @constraint(m, lhs ∈ moi_set(constraint_object(c)))
      end
    end
  end

  # Solve the new problem. (It should be really easy, as there is only one variable.)
  optimize!(m)
  α_val = value(α)
  y_bar_val = value.(y_bar)

  return α_val, y_bar_val
end


using Convex, SCS
set_default_solver(SCSSolver(verbose=0, max_iters=1000000))

println("small test to compile and see if we throw warnings when warmstarting what can't be warmstarted.")
n = 10
y = rand(n)
x = Variable(n)

lambda = [100]
problem = minimize(sumsquares(y - x) + lambda[1] * sumsquares(x - 10))
@time solve!(problem, warmstart = true)

println("\nnow try first solving and then warmstarting")
n = 1000
y = rand(n)
x = Variable(n)

lambda = [100]
problem = minimize(sumsquares(y - x) + lambda[1] * sumsquares(x - 10))
@time solve!(problem)
@show problem.optval
lambda[1] = 105
println("this run should be faster than the last and have slightly higher optimal value")
@time solve!(problem, warmstart=true)
@show problem.optval

println("\nnow try with scalar variables")
lambda = 100
problem = minimize(sumsquares(y - x) + lambda * sumsquares(x - 10))
@time solve!(problem)
@show problem.optval
lambda = 105
println("this run should be faster than the last and have slightly higher optimal value")
@time solve!(problem, warmstart=true)
@show problem.optval


function DCVoltageSource(; name, V = 1.0)
    @named oneport = OnePort()
    @unpack v = oneport
    ps = @parameters V=V
    eqs = [
        V ~ v
    ]
    extend(ODESystem(eqs, t, [], ps; name), oneport)
end

# V = Vm sin(ωt+ϕ)
function ACVoltageSource(; name, Vm = 1.0, ω = 2π, ϕ = 0)
    @named oneport = OnePort()
    @unpack v, i = oneport
    ps = @parameters Vm=Vm ω=ω ϕ=ϕ
    eqs = [
        v ~ Vm * sin(ω * t + ϕ)
    ]
    extend(ODESystem(eqs, t, [], ps; name), oneport)
end


export Bins, count!, hellinger, data_discrepancy, hellinger_stat_simulations

struct Bins
    Z::Array{Function} # lenght(z) = n List of membership functions
    q::Array{Float64} # lenght(q) = n empirical probabilities associated with each interval
    count::Array{Int64} # simulation counts in each bin
    Bins(Z,q,count) = length(Z) == length(q) ==length(count) ? new(Z,q,count) : error("length mismatch")
end

function Bins(data; min = -Inf, max = +Inf) # censoring cutoffs of continuous measurement and number of measurements outside cutoff
    lows = sum(data .< min)
    highs = sum(data .>= max)
    datapoints = sort(filter(x -> min <= x < max, data)) # get things in order
    midpoints = [(datapoints[ii] + datapoints[ii+1])/2 for ii in 1:length(datapoints)-1]
    boundaries = cat(-Inf, min, midpoints, max, Inf; dims =1)
    Z = [x -> boundaries[ii] <= x <= boundaries[ii+1] for ii in 1:length(boundaries)-1]
    q = cat(lows, ones(length(datapoints)), highs; dims=1)
    q = q/sum(q)
    count = zeros(Int64,length(q))
    Bins(Z,q,count)
end

function count!(x,b::Bins)
    ind = findfirst(z -> b.Z[z](x) ,eachindex(b.Z))
    b.count[ind] += 1
end

function hellinger(b::Bins)
    total = sum(b.count)
    sum( (sqrt(q) - sqrt(c/total))^2 for (c,q) in zip(b.count,b.q))
end

function observed_discrepancy(reference_data, simulated_data; min = -Inf, max = +Inf)
    bins = Bins(reference_data)
    for x in simulated_data
        count!(x,bins)
    end
    hellinger(bins)
end

function data_discrepancy(bins::Bins, θ_vec, ab_profiles... ; samples = 10^4)
    # the data are expected to be minus-log-mutant-frequency objects.
    # The cutoffs are the minimum and maximum data observed.
    b = bins
    b.count .= 0
    for θ in θ_vec
        sampler_array = map(x->WrightSampler(x...),
            hcat([θ .*  profile for profile in ab_profiles]...))
        for _ in 1:samples
            x = -log(escape_freq(sample!.(sampler_array)))
            count!(x, b)
        end
    end
    hellinger(b)
end

function hellinger_stat_simulations(θ_vec, ab_profiles... ;
    data_size=16, min=0.1, max=56/3, # parameters of the data
    samples = 10^3, # parameters of the simulation
    trials =10^3) # * length of θ_vec = generated samples
    datastream = []
    for θ in θ_vec
        sampler_array = map(x->WrightSampler(x...),
                hcat(θ .* [ ab_profiles...]...))
        for ii in 1:trials*data_size
        push!(datastream, -log(escape_freq(sample!.(sampler_array))))
        end
    end
    out =[]
    while length(datastream) >= data_size
        data = [];
        for jj in 1:data_size
            ind = rand(1:length(datastream))
            push!(data, splice!(datastream, ind))
        end
        b = Bins(data; min=min,max=max)
        push!(out, data_discrepancy(b, θ_vec, ab_profiles... ; samples = samples))
    end
    return out
end