struct SimpleOG{TI <: Integer, T <: Real,
                TR <: AbstractArray{T}, TQ <: AbstractArray{T}}
    B :: TI
    M :: TI
    α :: T
    β :: T
    R :: TR
    Q :: TQ
end

function SimpleOG{T <: Real}(;B::Integer=10, M::Integer=5, α::T=0.5, β::T=0.9)

    u(c) = c^α
    n = B + M + 1
    m = M + 1

    R = Matrix{T}(n, m)
    Q = zeros(Float64,n,m,n)

    for a in 0:M
        Q[:, a + 1, (a:(a + B)) + 1] = 1 / (B + 1)
        for s in 0:(B + M)
            R[s + 1, a + 1] = a<=s ? u(s - a) : -Inf
        end
    end

    return SimpleOG(B, M, α, β, R, Q)
end


### A Pluto.jl notebook ###
# v0.12.16

using Markdown
using InteractiveUtils

# ╔═╡ e077ab16-22e5-11eb-15ed-5f0654bf7328
using Pkg, DrWatson

# ╔═╡ e75a9a06-22e5-11eb-2e36-8d4f62b830ed
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Statistics, DataFrames, CSV
	using Random, Distributions, Plots
	using HypothesisTests
	Random.seed!(0)
end;

# ╔═╡ ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
md"## Listing6.02"

# ╔═╡ 1aa41898-3afa-11eb-2f49-570853ce6435
begin
	data = CSV.read(datadir("machine1.csv"), DataFrame; header=false)[:,1]
	xBar, n = mean(data), length(data)
	s = std(data)
	alpha = 0.1
	t = quantile(TDist(n-1),1-alpha/2)

	Text("Calculating formula: $((xBar - t*s/sqrt(n), xBar + t*s/sqrt(n)))")
end

# ╔═╡ 88f0e5d6-3b0b-11eb-0af9-b964a1170ff3
Text("Using confint() function: $(confint(OneSampleTTest(xBar,s,n),alpha))")

# ╔═╡ 475ff888-22e6-11eb-2354-09f8f40a8e12
md"## End of listing6.02"

# ╔═╡ Cell order:
# ╟─ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
# ╠═e077ab16-22e5-11eb-15ed-5f0654bf7328
# ╠═e75a9a06-22e5-11eb-2e36-8d4f62b830ed
# ╠═1aa41898-3afa-11eb-2f49-570853ce6435
# ╠═88f0e5d6-3b0b-11eb-0af9-b964a1170ff3
# ╟─475ff888-22e6-11eb-2354-09f8f40a8e12


# An example of how the solution for the wave equation
# using a gaussian initial condition converges

using ScalarWaveEquation
using PyPlot

function get_gaussian(N::Int)
	# returns a normalized Gaussian distribution with (mean, std) = (0.5, 0.1)
	mean = 0.5
	std  = 0.1

	xdata      = range(0.0, stop = 1.0, length = N)
	exponent   = -0.5 * (xdata.-mean) .* (xdata.-mean) / (std*std)
	norm_coeff = std*sqrt(2*pi)

	return exp.(exponent)/norm_coeff
end

function get_gaussian_solution(Nx::Int, Nt::Int, Ttot::T) where {T}
	# returns a FunArr2D instance for the solution to the gaussian
	# initial condition, performed at difference resolutions

	# initial conditions are accepted in the form of FunArr instances
	# the x-range is assumed to be [0, 1]
	init_positions  = get_gaussian(Nx)
	init_velocities = zeros(Nx)
	ufun  = FunArr{Float64, Float64}(init_positions)
	utfun = FunArr{Float64, Float64}(init_velocities)

	# use the method of lines to calculate u(x, t) over [0, 1]x[0, Ttot]
	u2fun, energies = scalarwave_and_energy(ufun, utfun, Nt, Ttot)

	return u2fun
end

# get the solution at different resolutions for the same duration
u2fun1 = get_gaussian_solution(100, 150, 0.2)
u2fun2 = get_gaussian_solution(200, 300, 0.2)
u2fun3 = get_gaussian_solution(300, 450, 0.2)
u2fun4 = get_gaussian_solution(400, 600, 0.2)
u2fun5 = get_gaussian_solution(500, 750, 0.2)
u2fun6 = get_gaussian_solution(600, 900, 0.2)

# get the integrated squared difference between solutions
# for successive resolutions, for a 2000x2000 spacing
area12 = integrate_difference(u2fun1, u2fun2, 2000, 2000)
area23 = integrate_difference(u2fun2, u2fun3, 2000, 2000)
area34 = integrate_difference(u2fun3, u2fun4, 2000, 2000)
area45 = integrate_difference(u2fun4, u2fun5, 2000, 2000)
area56 = integrate_difference(u2fun5, u2fun6, 2000, 2000)

# the differences between increasing resolutions decreases
println("The integrated difference between 100x150 and 200x300 is $area12")
println("The integrated difference between 200x300 and 300x450 is $area23")
println("The integrated difference between 300x450 and 400x600 is $area34")
println("The integrated difference between 400x600 and 500x750 is $area45")
println("The integrated difference between 500x750 and 600x900 is $area56")


#=
  From the WebPPL model https://github.com/probmods/ppaml2016/blob/gh-pages/chapters/5-data.md
  """
  Posterior prediction and model checking

  One important use of posterior predictive distributions is to examine the descriptive
  adequacy of a model. The posterior predictive can be viewed as a set of predictions about
  what data the model expects to see, based on the posterior distribution over parameters.
  If these predictions do not match the data already seen, the model is descriptively inadequate.

  Let's say we ran 2 experiments that we believe are conceptually the same (e.g., asking
  10 people whether or not they would vote for "Hillary Clinton or Donald Trump" and asking
  a separate group of 10 people if they would vote for "Donald Trump or Hillary Clinton").
  Suppose we observed the following data from those 2 experiments: k1=0; k2=10.
  """
=#

using Turing
include("jl_utils.jl")

@model function bda2(k1=0,n1=10, k2=10,n2=10)

    # Sample rate from Uniform distribution
    p ~ Uniform(0,1)

    k1 ~ Binomial(n1, p)
    k2 ~ Binomial(n2, p)

    # sample from binomial with updated p
    prior_p ~ Uniform(0,1)
    posteriorPredictive1 ~ Binomial(n1,p)
    posteriorPredictive2 ~ Binomial(n2,p)

end

# data
k1=0
n1=10
k2=10
n2=10
println("k1:$(k1) n1:$(n1) k2:$(k2) n2:$(n2)")
model = bda2(k1,n1,k2,2)

# chns = sample(model, MH(), 10_000)
# chns = sample(model, PG(15), 10_000)
# chns = sample(model, IS(), 10_000)
chns = sample(model, SMC(), 10_000)
# chns = sample(model, SMC(), MCMCThreads(), 10_000, 4)

display(chns)

show_var_dist_pct(chns, :posteriorPredictive1)
show_var_dist_pct(chns, :posteriorPredictive2)


"""
    CompoundRatingCurve(threshold::Vector{<:Real},rc::Vector{RatingCurve})

Construct an object of type CompoundRatingCurve
"""
struct CompoundRatingCurve
    threshold::Vector{<:Real}
    component::Vector{RatingCurve}

    function CompoundRatingCurve(threshold::Vector{<:Real}, component::Vector{RatingCurve})
        length(threshold) == ( length(component) - 1 ) ||
            error("The number of segments does not match the number of rating curves.")
        new(threshold, component)
    end

end

CompoundRatingCurve(threshold::Vector{<:Int}, component::Vector{RatingCurve}) = CompoundRatingCurve(float.(threshold), component)

Base.Broadcast.broadcastable(obj::CompoundRatingCurve) = Ref(obj)

"""
    function bic(crc::CompoundRatingCurve)

BIC of the compound rating curve model.
"""
function bic(crc::CompoundRatingCurve)

    n₁ = length(crc.component[1].gauging)
    n₂ = length(crc.component[2].gauging)
    n = n₁ + n₂
    p = 6

    @assert n>6

    σ̂² = RatingCurves.var(crc)

    return n₁*log(σ̂²[1]) + n₂*log(σ̂²[2]) + p*log(n)

end

"""
    cint(crc::CompoundRatingCurve; nboot::Int=100, α::Real=.05)

Compound rating curve parameter confidence intervals of level `1-α` obtained by a bootstrap sample of size `nboot`.
"""
function cint(crc::CompoundRatingCurve; nboot::Int=100, α::Real=.05)

    k = Vector{Float64}(undef, nboot)

    a₁ = Vector{Float64}(undef, nboot)
    b₁ = Vector{Float64}(undef, nboot)
    c₁ = Vector{Float64}(undef, nboot)

    a₂ = Vector{Float64}(undef, nboot)
    b₂ = Vector{Float64}(undef, nboot)
    c₂ = Vector{Float64}(undef, nboot)

    for i in 1:nboot

        G₁ = RatingCurves.bootstrap(crc.component[1].gauging)
        G₂ = RatingCurves.bootstrap(crc.component[2].gauging)
        crcᵢ = crcfit([G₁..., G₂...])

        k[i] = crcᵢ.threshold[1]

        a₁[i] = crcᵢ.component[1].a
        b₁[i] = crcᵢ.component[1].b
        c₁[i] = crcᵢ.component[1].c

        a₂[i] = crcᵢ.component[2].a
        b₂[i] = crcᵢ.component[2].b
        c₂[i] = crcᵢ.component[2].c

    end

    M = hcat(
            quantile(k, [α/2, 1-α/2]),
            quantile(a₁, [α/2, 1-α/2]),
            quantile(b₁, [α/2, 1-α/2]),
            quantile(c₁, [α/2, 1-α/2]),
            quantile(a₂, [α/2, 1-α/2]),
            quantile(b₂, [α/2, 1-α/2]),
            quantile(c₂, [α/2, 1-α/2])
            )

    return M

end

"""
    discharge(crc::RatingCurve, h::Real)

Compute the estimated discharge at level `h` with the compound rating curve `crc`.
"""
function discharge(crc::CompoundRatingCurve, h::Real)

    y = logdischarge(crc, h)

    return exp(y)

end

"""
    level(crc::RatingCurve, q::Real)

Compute the level corresponding the the discharge `q` and the compound rating curve `crc`.
"""
function level(crc::CompoundRatingCurve, q::Real)

    @assert q>0

    threshold = vcat(crc.threshold, Inf)
    component = crc.component

    i = 1
    h = level(component[i], q)

    while(h > threshold[i])
        i +=1
        h = level(component[i], q)
    end

    return h

end

"""
    logdischarge(crc::RatingCurve, h::Real)

Compute the estimated log discharge at level `h` with the compound rating curve `crc`.
"""
function logdischarge(crc::CompoundRatingCurve, h::Real)

    threshold = vcat(crc.threshold, Inf)

    ind = findfirst(h .< threshold)

    y = logdischarge(crc.component[ind], h)

    return y

end

"""
    pint(crc::CompoundRatingCurve, level::Real, α::Real=0.05, rtol::Real=.05)

`1-α` confidence interval of the estimated discharge at level `h` with the compound rating curve `crc`.

## Details

For each segment of the compound rating curve, the confidence intervals are estimated separately. See [`pint(::RatingCurve, ::Real, ::Real, ::Real)`](@ref)
for the details for the estimation on each segment.
"""
function pint(crc::CompoundRatingCurve, level::Real, α::Real=0.05, rtol::Real=.05)

    res = pintlog(crc, level, α, rtol)

    return exp.(res)
end

"""
    pintlog(crc::CompoundRatingCurve, level::Real, α::Real=0.05, rtol::Real=.05)

`1-α` confidence interval of the estimated log discharge at level `h` with the compound rating curve `crc`.

### Details

`rtol` represents the relative uncertainty of the dishcarge so that the true discharge is included in the interval `q ± 1.96*rtol` 95% of the time
"""
function pintlog(crc::CompoundRatingCurve, level::Real, α::Real=0.05, rtol::Real=.05)

    threshold = vcat(crc.threshold, Inf)

    ind = findfirst(level .< threshold)

    rc = crc.component[ind]

    return pintlog(rc, level, α, rtol)

end

"""
    Base.show(io::IO, obj::EVA)
Override of the show function for the objects of type RatingCurve.
"""
function Base.show(io::IO, obj::CompoundRatingCurve)

    showCompoundRatingCurve(io, obj)

end

function showCompoundRatingCurve(io::IO, obj::CompoundRatingCurve)

    threshold = [0, obj.threshold..., Inf]
    println(io, "CompoundRatingCurve")
    for i in 1:length(obj.component)
        println(io, "")
        println(io, "   for ",threshold[i]," ≤ h < ",threshold[i+1])
        showRatingCurve(io, obj.component[i],"      ")
    end
end

"""
    sse(rc::RatingCurve)

Compute the sum of the squares of residuals between the compound curve and the gauging in the log space.

### Details
A value of the sum of squares is given for each CompoundRatingCurve component.
"""
function sse(crc::CompoundRatingCurve)

    SSE = Float64[]

    for rc in crc.component
        push!(SSE, sse(rc)[])
    end

    return SSE

end

"""
    var(crc::CompoundRatingCurve)

Estimate the variance of the errors in the log space.
"""
function var(crc::CompoundRatingCurve)

    σ̂² = Float64[]

    for rc in crc.component
        push!(σ̂², RatingCurves.var(rc)[])
    end

    return σ̂²

end


# Computes the logarithmic mean: (aR-aL)/(LOG(aR)-LOG(aL)) = (aR-aL)/LOG(aR/aL)
# Problem: if aL~= aR, then 0/0, but should tend to --> 0.5*(aR+aL)
#
# introduce xi=aR/aL and f=(aR-aL)/(aR+aL) = (xi-1)/(xi+1)
# => xi=(1+f)/(1-f)
# => Log(xi) = log(1+f)-log(1-f), and for small f (f^2<1.0E-02) :
#
#    Log(xi) ~=     (f - 1/2 f^2 + 1/3 f^3 - 1/4 f^4 + 1/5 f^5 - 1/6 f^6 + 1/7 f^7)
#                  +(f + 1/2 f^2 + 1/3 f^3 + 1/4 f^4 + 1/5 f^5 + 1/6 f^6 + 1/7 f^7)
#             = 2*f*(1           + 1/3 f^2           + 1/5 f^4           + 1/7 f^6)
#  (aR-aL)/Log(xi) = (aR+aL)*f/(2*f*(1 + 1/3 f^2 + 1/5 f^4 + 1/7 f^6)) = (aR+aL)/(2 + 2/3 f^2 + 2/5 f^4 + 2/7 f^6)
#  (aR-aL)/Log(xi) = 0.5*(aR+aL)*(105/ (105+35 f^2+ 21 f^4 + 15 f^6)
function ln_mean(value1::Float64, value2::Float64)
  epsilon_f2 = 1.0e-4
  ratio = value2/value1
  # f2 = f^2
  f2=(ratio*(ratio-2.)+1.)/(ratio*(ratio+2.)+1.)
  if (f2<epsilon_f2)
    return (value1+value2)*52.5/(105.0 + f2*(35.0 + f2*(21.0 +f2*15.0)))
  else
    return (value2-value1)/log(ratio)
  end
end


# Example of MLE estimation. The data is really drawn from the
# Logit DGP, so the model is well specified, and the MLE has
# the properties discussed in class. E.g., if you make n very
# large you should see that the estimator is very close to the
# true value of theta used to generate data

using Econometrics
include("LogitDGP.jl")
n = 30 # sample size
θ  = [0, 0.5] # true theta for generating data
(y, x) = LogitDGP(n, θ) # generate the data

# now define things for estimation
model = θ -> logit(θ, y, x)
θstart = zeros(size(x,2)) # start values for estimation

# Perform the estimation - Make sure that you examine
# the MLE estimation programs so that you see how this works
θhat, objvalue, V, converged = mleresults(model, θstart, "estimate logit model");


using DifferentialEquations
using Plots
using PyPlot

# Set the required model parameters for the SIRS model with two levels of transmission - Within and between households
N = 10; # Household size - Change to 10 for final analysis
betaHH = 6; # Within household transmission parameter
betaG = 1; # Population wide transmission
gamma = 1; # Rate of recovery from infection
tau = 1; # Rate of loss of protection
params = [betaHH,gamma,tau,betaG,N]; # Put all the parameters together
time = (0.0, 30.0) # Simulation time - note it defined as a float
dim = dim = 0.5*(N+1)*(N+2); # Number of possible configurations - works for three epidemiological classes
y0 = vec(zeros(1,dim)); # Initial condition vector
y0[end-1] = 0.00000001;
y0[end] = 0.99999999;

function hhTransitions(N,dim)
    # Function to generate transition matrices for household model
    # Input: N is the household size

    # Initialize things
    Qinf = zeros(dim,dim);
    Qrec = zeros(dim,dim);
    Qext = zeros(dim,dim);
    Qwane = zeros(dim,dim);
    dataI = Array{Int64}(zeros(dim,3))
    m = 0;
    I = Array{Int64}(zeros(N+1,N+1))

    # To help remember where to store the variables
    for ss = 0:N
        for ii = 0:(N-ss)
            m = m + 1;
            I[ss+1,ii+1] = m
        end
    end

    # Describe the epidemiological transitions

    # Counter for susceptibles
    for ss = 0:N
        # Counter for infecteds
        for ii = 0:(N-ss)
            # If susceptibles and infecteds are more than 1, then infection within the household can occur
            if (ss > 0 && ii > 0)
                Qinf[I[ss+1,ii+1],I[ss,ii+2]] = ii*ss/(N-1);
            end

            # If infecteds are more than 1, recovery can occur
            if ii > 0
                # Rate of recovery
                Qrec[I[ss+1,ii+1],I[ss+1,ii]] = ii;
            end

            # For external infection - just keep track of susceptibles
            if ss > 0
                # Rate of within household infection
                Qext[I[ss+1,ii+1],I[ss,ii+2]] = ss;
            end

            # Loss of protection hence becoming susceptible again. Possible if N-ss-ii = rr > 0
            if (N-ss-ii) > 0
                # Rate of loss of protection
                Qwane[I[ss+1,ii+1],I[ss+2,ii+1]] = N-ss-ii;
            end

            # Store the relevant indices to help identify the household configurations
            dataI[I[ss+1,ii+1],:] = [ss, ii, N-ss-ii];
        end
    end

    Qinf = Qinf - diagm(vec(sum(Qinf,2)),0);
    Qrec = Qrec - diagm(vec(sum(Qrec,2)),0);
    Qext = Qext - diagm(vec(sum(Qext,2)),0);
    Qwane = Qwane - diagm(vec(sum(Qwane,2)),0);

    # Return
    return Qinf, Qrec, Qext, Qwane, dataI
end

Qinf, Qrec, Qext, Qwane, dataI = hhTransitions(N,dim);

function rateSIRS(dy_dt,y0,params,time)

    # Extract the parameters
    betaHH = params[1];
    gamma = params[2];
    tau = params[3];
    betaG = params[4];
    N = params[5];

    # Generate the transition matrices
    Qinf, Qrec, Qext, Qwane, HHconfig = hhTransitions(N,dim);

    # Combine within and external transitions
    Q = betaHH*Qinf + gamma*Qrec + tau*Qwane + (betaG*((HHconfig[:,2]'*y0)/N)*Qext);

    # Generate the rates of change for ODE solver
    dy_dt .= (y0'*Q)';

    # Alternatively this works too
    #=for i=1:length(y0)
        dy_dt[i] = y0'*Q[:,i]
    end=#

end

# Define the ODE problem
prob = ODEProblem(rateSIRS,y0,time,params);

# Solve
sol = solve(prob);

Iconfig = dataI[:,2];
infProp = zeros(length(sol.t),1);

# Prepare the plots
for i = 1:length(sol.t)
    infProp[i,1] = sol[:,i]'*Iconfig/N;
end

# Total infectious in the population
plot(sol.t,infProp,color="blue",xlabel="Time",ylabel="Proportion infectious",label=["Mean infection"],ylims=[0, 1])

# Household profile at endemic prevalence
# Prepare the plots
yprop = zeros(N+1,length(sol.t))
for j = 1:length(sol.t)
    for i = 1:N+1
        index = find(Iconfig.==i-1);
        yprop[i,j] = sum(sol[index,j])
    end
end
step(-0.5:1:10.5,[yprop[:,length(sol.t)];0])

# Household profile at peak prevalence
# Prepare the plots
x = find(infProp.==maximum(infProp))
step(-0.5:1:10.5,[yprop[:,x];0]);


# See https://en.wikipedia.org/wiki/Matrix_(mathematics)
# Examples from https://docs.julialang.org/en/v1/stdlib/LinearAlgebra/

using LinearAlgebra

expressions = [
    :(inv([1 2 3; 4 1 6; 7 8 1]))
    :(det([1 2 3; 4 1 6; 7 8 1]))
    :(eigvals([-4. -17.; 2. 2.]))
    :(eigvecs([-4. -17.; 2. 2.]))
    :(factorize([1.5 2 -4; 3 -1 -6; -10 2.3 4]))
    :(factorize([1.5 2 -4; 2 -1 -3; -4 -3 5]))
    :(([1 0; 1 -2] \ [32; -4]))
    :(([1 0; 1 -2] * [32.0, 18.0]))
    :(dot([1; 1], [2; 3]))
    :(cross([0; 1; 0], [0; 0; 1]))
    :(UpperHessenberg([1 2 3 4; 5 6 7 8; 9 10 11 12; 13 14 15 16]))
    :(cholesky([4. 12. -16.; 12. 37. -43.; -16. -43. 98.]).U)
    :(cholesky([4. 12. -16.; 12. 37. -43.; -16. -43. 98.]).L)
]

for expr in expressions
    println("[] :julia-answer \"\"\"", escape_string(string(expr)), " = ", eval(expr), "\"\"\".")
end


#!/usr/bin/julia

# Trizen
# Date: 18 August 2016
# https://github.com/trizen

# Count the number of partitions of n.

# See also: https://oeis.org/A000041
#           https://en.wikipedia.org/wiki/Partition_(number_theory)

function count_partitions(x::Int64)
    n = 2
    p = Int64[1]

    while (n <= x+1)
        i = 0
        q = 2
        push!(p, 0)

        while q <= n
            p[n] += (i % 4 > 1 ? -1 : 1) * p[n-q+1]
            i += 1
            j = div(i, 2) + 1
            isodd(i) && (j *= -1)
            q = div(j * (3j - 1), 2) + 1
        end

        n += 1
    end

    p[n-1]
end

println("P(200) = ", count_partitions(200))        # 3972999029388


"""
    Chisq(ν)
The *Chi squared distribution* (typically written χ²) with `ν` degrees of freedom has the
probability density function

```math
f(x; k) = \\frac{x^{k/2 - 1} e^{-x/2}}{2^{k/2} \\Gamma(k/2)}, \\quad x > 0.
```

If `ν` is an integer, then it is the distribution of the sum of squares of `ν` independent standard [`Normal`](@ref) variates.

```julia
Chisq(k)     # Chi-squared distribution with k degrees of freedom

params(d)    # Get the parameters, i.e. (k,)
dof(d)       # Get the degrees of freedom, i.e. k
```

External links

* [Chi-squared distribution on Wikipedia](http://en.wikipedia.org/wiki/Chi-squared_distribution)
"""
struct Chisq{T<:Real} <: ContinuousUnivariateDistribution
    ν::T
    Chisq{T}(ν::T) where {T} = new{T}(ν)
end

function Chisq(ν::T) where {T <: Real}
    @check_args(Chisq, ν > zero(ν))
    return Chisq{T}(ν)
end

Chisq(ν::T, ::NoArgCheck) where {T <: Real} = Chisq{T}(ν)
Chisq(ν::Integer) = Chisq(float(ν))

@distr_support Chisq 0.0 Inf

#### Parameters

dof(d::Chisq) = d.ν
params(d::Chisq) = (d.ν,)
@inline partype(d::Chisq{T}) where {T<:Real} = T

### Conversions
convert(::Type{Chisq{T}}, ν::Real) where {T<:Real} = Chisq(T(ν))
convert(::Type{Chisq{T}}, d::Chisq{S}) where {T <: Real, S <: Real} = Chisq(T(d.ν))

#### Statistics

mean(d::Chisq) = d.ν

var(d::Chisq) = 2d.ν

skewness(d::Chisq) = sqrt(8 / d.ν)

kurtosis(d::Chisq) = 12 / d.ν

mode(d::Chisq{T}) where {T<:Real} = d.ν > 2 ? d.ν - 2 : zero(T)

function median(d::Chisq; approx::Bool=false)
    if approx
        return d.ν * (1 - 2 / (9 * d.ν))^3
    else
        return quantile(d, 1//2)
    end
end

function entropy(d::Chisq)
    hν = d.ν/2
    hν + logtwo + lgamma(hν) + (1 - hν) * digamma(hν)
end

#### Evaluation

@_delegate_statsfuns Chisq chisq ν

mgf(d::Chisq, t::Real) = (1 - 2 * t)^(-d.ν/2)

cf(d::Chisq, t::Real) = (1 - 2 * im * t)^(-d.ν/2)

gradlogpdf(d::Chisq{T}, x::Real) where {T<:Real} =  x > 0 ? (d.ν/2 - 1) / x - 1//2 : zero(T)

#### Sampling

rand(rng::AbstractRNG, d::Chisq) =
    (ν = d.ν; rand(rng, Gamma(ν / 2.0, 2.0one(ν))))

sampler(d::Chisq) = (ν = d.ν; sampler(Gamma(ν / 2.0, 2.0one(ν))))


using LinearAlgebra, DataFrames, Statistics

n = 1000
x = [ones(n) rand(1:0.1:10, n)]
b,sigma = [2.5, -1.2], 1
y = x*b + sigma .* randn(n);

# Write types
abstract type StandardErrorEstimator end
struct Bootstrap <: StandardErrorEstimator
    bootstrap_samples::Int
    subsample_size::Int
end
struct Spherical <: StandardErrorEstimator end

# Estimator code
ols(x, y) = ols(x, y, Spherical())
function ols(x, y, se_estimator::Spherical)
    V_inv = inv(x'x)
    b_hat = V_inv*x'y
    y_hat = x*b_hat
    e = y - y_hat

    n = length(y)
    k = size(x, 2)
    s2  = e'e / (n-k)

    se=sqrt.(s2 .* diag(V_inv))

    return DataFrame(
        b_hat=b_hat,
        se=se,
    )
end

function ols(x, y, se_estimator::Bootstrap)
    n = length(y)
    n_samples = se_estimator.bootstrap_samples
    subsample_size = se_estimator.subsample_size
    b_hats = zeros(n_samples, size(x, 2))

    for k in 1:n_samples
        inds = rand(1:n, subsample_size)
        x_sub = x[inds,:]
        y_sub = y[inds]

        b_hats[k, :] = inv(x_sub'x_sub)*x_sub'y_sub
    end

    b_hat = vec(mean(b_hats, dims=1))
    se_hat = vec(std(b_hats, dims=1))

    return DataFrame(
        b_hat=b_hat,
        se_hat=se_hat,
    )
end

fit1 = ols(x, y, Spherical())
fit2 = ols(x, y, Bootstrap(100_000, 10))

println("Spherical")
display(fit1)

println("\nBootstrapped")
display(fit2)


# It was proposed by Christian Goldbach that every odd composite number can be
# written as the sum of a prime and twice a square.
#
# 9 = 7 + 2×1^2
# 15 = 7 + 2×2^2
# 21 = 3 + 2×3^2
# 25 = 7 + 2×3^2
# 27 = 19 + 2×2^2
# 33 = 31 + 2×1^2
#
# It turns out that the conjecture was false.
#
# What is the smallest odd composite that cannot be written as the sum of a
# prime and twice a square?

using ProjectEulerSolutions

# Make a resonable list of twice squares and primes, then sum combinations of
# them and make a set that you compare against the set of odd composite
# numbers.  Would probably be faster to not calculate all the prime numbers
# and odd composite numbers upfront.
function p046solution(n::Integer=100)::Integer
    primearray = sieve_eratosthenes(n)
    primeset = Set(primearray)
    oddnumset = Set(9:2:n)
    oddcompositeset = setdiff(oddnumset, primeset)

    twicesquares = 2 .* (1:Integer(floor(0.75*sqrt(n)))).^2

    sums = Set()
    for ts in twicesquares
        union!(sums, Set(ts .+ primearray))
    end

    oddcomp = setdiff(oddcompositeset, sums)
    if length(oddcomp) > 0
        return minimum(oddcomp)
    else
        return -1
    end
end

p046 = Problems.Problem(p046solution)

Problems.benchmark(p046, 10_000)


#=
Take the number 192 and multiply it by each of 1, 2, and 3:

192 × 1 = 192
192 × 2 = 384
192 × 3 = 576
By concatenating each product we get the 1 to 9 pandigital, 192384576. We will call 192384576 the concatenated product of 192 and (1,2,3)

The same can be achieved by starting with 9 and multiplying by 1, 2, 3, 4, and 5, giving the pandigital, 918273645, which is the concatenated product of 9 and (1,2,3,4,5).

What is the largest 1 to 9 pandigital 9-digit number that can be formed as the concatenated product of an integer with (1,2, ... , n) where n > 1?
=#
function canbepandigital(n)
  stringresult = ""
  i = 1
  while length(stringresult) < 9
    stringresult *= string(n * i)
    i += 1
  end

  if (length(stringresult)) == 9
    result = parse(Int, stringresult)
    if sort(digits(result)) == collect(1:9)
      return (true,result)
    end
  end
  (false, 0)
end

function calc()
  maxpandigital = 0
  for i in 2:10_000
    (result, pandigital) = canbepandigital(i)
    if result && pandigital > maxpandigital
      maxpandigital = pandigital
    end
  end
  maxpandigital
end
@time println(calc())


using Plots

Rotation(R, θ) = [cos(θ) -sin(θ);sin(θ) cos(θ)] * R

function DragonFractal(;order, θ, points)
    θ₁ = θ[1]; θ₂ = θ[2]
    for num in 1:order
        i = 1
        while i < length(points)
            point = (points[i+1] - points[i]) * sin(abs(θ₁))
            point = Rotation(point, θ₁)  + points[i]
            insert!(points,i+1,point)
            i += 2
            point = (points[i+1] - points[i]) * sin(abs(θ₂))
            point = Rotation(point, θ₂)  + points[i]
            insert!(points,i+1,point)
            i += 2
        end
    end
    return points
end

for i in [1,5,10,15]
    Parameters = Dict(:order => i,
                    :θ => [pi/4, -pi/4],
                    :points => [[0.0; 0.0],[1.0; 0.0],[1.0; 1.0]])

    points = DragonFractal(;Parameters...)
    half = Int(floor(length(points)/2))

    plot(hcat(points...)[1,1:half], hcat(points...)[2,1:half], legend = false, color =:blue, border=:none)
    plot!(hcat(points...)[1,half:end], hcat(points...)[2,half:end], legend = false, color =:red, border=:none)
    savefig("C:\\Users\\Yaghoub\\Documents\\GitHub\\ComputationalPhysics-Fall2021\\ProblemSet1\\Figs\\Q2\\Heighway-dragon-O$i.png")

    Parameters = Dict(:order => i,
                    :θ => [-pi/4, -pi/4],
                    :points => [[0.0; 0.0],[1.0; 0.0],[1.0; 1.0]])
    points = DragonFractal(;Parameters...)
    plot(hcat(points...)[1,:], hcat(points...)[2,:], legend = false, border=:none, color =:black, linewidth=0.4)
    savefig("C:\\Users\\Yaghoub\\Documents\\GitHub\\ComputationalPhysics-Fall2021\\ProblemSet1\\Figs\\Q2\\Lévy-C-curve-O$i.png")
end


using LinearAlgebra
using GLM
using LsqFit
using Distributions
using StatsBase
using Printf
using Plots

#linear fit wrapper

function lfit(ndf::DataFrame)
    lr = lm(@formula(y_mean ~ x_mean), ndf)
    c = coef(lr)
    return x -> c[1] + c[2]*x, predict(lr), c
end

struct RFit
	fitpar::Vector{Number}
	fitstd::Vector{Number}
	ci::Vector{Tuple{Number, Number}}
	g::Function
end

function gpol1(ct)
    function f1(z)
        return ct[1] + ct[2] * z
    end
    return f1
end

function gpol2(ct)
    function f2(z)
        return ct[1] + ct[2] * z + ct[3] * z^2
    end
    return f2
end

function gpol3(ct)
    function f3(z)
        return ct[1] + ct[2] * z + ct[3] * z^2 + ct[4] * z^3
    end
    return f3
end

function polfit(pol, x::Vector{<:Real},y::Vector{<:Real}, p0::Vector{<:Real})
    fq = curve_fit(pol, x, y, p0)
    cfq = coef(fq)
    @info "coef(fq)" cfq
    sfq = stderror(fq)
    @info "std(fq)" sfq
    @info "margin_of_error (90%)" margin_error(fq, 0.1)
    @info " confidence_interval (90%)" confidence_interval(fq, 0.1)
    return cfq, sfq
end

function fit_pol1(x::Vector{<:Real},y::Vector{<:Real}, ci=0.1)
    @. pol(x, p) = p[1] + p[2] * x
    p0 = [1.0, 1.0]
    fq = curve_fit(pol, x, y, p0)

    RFit(coef(fq), stderror(fq), confidence_interval(fq, ci),
         gpol1(coef(fq)))
end

function fit_pol2(x::Vector{<:Real},y::Vector{<:Real}, ci=0.1)
    @. pol(x, p) = p[1] + p[2] * x + p[3] * x^2
    p0 = [1.0, 1.0, 1.0]
    fq = curve_fit(pol, x, y, p0)
    RFit(coef(fq), stderror(fq), confidence_interval(fq, ci),
         gpol2(coef(fq)))
end

function fit_pol3(x::Vector{<:Real},y::Vector{<:Real}, ci=0.1)
    @. pol(x, p) = p[1] + p[2] * x + p[3] * x^2 + p[4] * x^3
    p0 = [1.0, 1.0, 1.0, 1.0]
    fq = curve_fit(pol, x, y, p0)
    RFit(coef(fq), stderror(fq), confidence_interval(fq, ci),
         gpol3(coef(fq)))
end

"""
	fit_gauss(y, xmin, xmax, bins=25)

Fit a normal distribution to data
"""

struct FGauss
	mu::Vector{Number}
	std::Vector{Number}
	C::Vector{Number}
	h::Histogram
	X::Vector{Number}
	Y::Vector{Number}
	g::Vector{Function}
end

function gausg(μ::Real, σ::Real, C::Real)
	function gausx(x)
		return C * pdf(Normal(μ, σ,), x)
	end
	return gausx
end

function gausg2(μ1::Real, σ1::Real, C1::Real, μ2::Real, σ2::Real, C2::Real)
	function gausx(x)
		return C1 * pdf(Normal(μ1, σ1,), x) + C2 * pdf(Normal(μ2, σ2,), x)
	end
	return gausx
end

@. gauss1fm(x, p) = p[1]* pdf(Normal(p[2], p[3]), x)
@. gauss1(x, p) = p[1]* pdf(Normal(p[2], p[3]), x)
@. gauss2(x, p) = p[1]* pdf(Normal(p[2], p[3]), x) +  p[4]* pdf(Normal(p[5], p[6]), x)
@. gausscm(x, p) = p[1] * pdf(Normal(p[2], p[3]), x) + p[4] * pdf(Normal(p[2], p[5]), x)

function gaussfm(mu::Real)
	function gauss(x::Vector{<:Real}, p::Vector{<:Real})
		return @. p[1]* pdf(Normal(mu, p[2]), x)
	end
	return gauss
end

function gauss2fm(mu::Real)
	function gauss2(x::Vector{<:Real}, p::Vector{<:Real})
		return @. p[1]* pdf(Normal(mu, p[2]), x) + p[3]* pdf(Normal(mu, p[4]), x)
	end
	return gauss2
end

function gauss1fm(mu)
    function gauss1(x,p)
        return @. p[1]* pdf(Normal(mu, p[2]), x)
    end
    return gauss1
end

function cfit(ffit::Function, x::Vector{Float64}, y::Vector{Float64},
	                         p0::Vector{Float64}, lb::Vector{Float64}, ub::Vector{Float64})
	fq = curve_fit(ffit, x, y, p0, lower=lb, upper=ub)
    cfq = coef(fq)
    @debug "coef(fq)" cfq
    return cfq
end

"""
    cfit

Take into account the errors on the measured variables. Assuming weights
of 1/sigma^2 as in standard least squares.
"""
function cfit(ffit::Function, x::Vector{<:Real}, y::Vector{<:Real},
	          yerr::Vector{<:Real}, p0::Vector{<:Real},
              lb::Vector{<:Real}, ub::Vector{<:Real})
	fq = curve_fit(ffit, x, y, yerr.^-2, p0, lower=lb, upper=ub)
    return fq
end

function fit_gauss(h::Histogram)
	c = centers(h)
	w = h.weights * 1.0
	@debug "histo"  w c
	mu, sigma = mean_and_std(c, Weights(h.weights); corrected = false)
	@debug "mu, std" mu, sigma

	# fit parameters lb, ub, po are lower, upper bounds and pars
    lb = [0., mu - 100.0*sigma, sigma/100.0]
    ub = [100*sum(w), mu + 100.0*sigma, 100.0*sigma]
    p0_bounds = [sum(w), mu, sigma]
	CC, μ, σ  = cfit(gauss1, c, w, p0_bounds, lb, ub)
	gx = gausg(μ, σ, CC)
	return FGauss([μ], [σ], [CC], h, c, gx.(c), [gx])

end

fit_gauss(x::Vector{Float64}, xmin::Float64, xmax::Float64;
	      bins::Integer=50, norm=false) =fit_gauss(hist1d(x, bins, xmin, xmax, norm))

"""
	fit_gauss_fm(y::Vector{Float64}, xmin::Float64, xmax::Float64, bins=50, fm=0.0)

Fit a gaussian with a fixed mean fm
"""
function fit_gauss_fm(y::Vector{Float64}, xmin::Float64, xmax::Float64;
	                  bins=50, norm=false, fm=0.0)

	# fit the unbinned distribution
	x  =  in_range(y, xmin, xmax)
	σ = std(x)
	@debug "gfit_gauss_fm: σ = $σ"

	# bin distribution
    h = hist1d(x, bins, xmin, xmax, norm)
    c = centers(h)
    w = h.weights *1.0
    @debug "histo w and c"  w c

	# fit parameters lb, ub, po are lower, upper bounds and pars
    lb = [0.0, σ/10.0]
    ub = [1.0e+6*sum(w),  σ * 10.0]
    p0_bounds = [0.0, σ]

	g1 = gaussfm(fm)
	CC, sigma  = cfit(g1, c, w, p0_bounds, lb, ub)

	@debug "CC,  sigma"  CC  sigma
	gx = gausg(fm, sigma, CC)

	return FGauss([fm], [sigma], [CC], h, c, gx.(c), [gx])
end

function plot_fit_gauss(x::Vector{Float64}, xs::String, ys::String,
    bins::Integer, xmin::Float64, xmax::Float64;
    xgmin::Float64, xgmax::Float64, gbins::Integer=50)

    h,p = hist1d(x, xs, bins, xmin, xmax, norm=true)
    fg  = fit_gauss(x, xgmin, xgmax, bins=gbins, norm=true)
    gx  = fg.g[1]
    X   = centers(h)
    Y   = h.weights
    σY  = sqrt.(Y)
    lbl = @sprintf " μ=%5.1f, σ =%5.1f " fg.mu[1] fg.std[1]
    lbl =string("gaussian fit:\n", lbl)
    p = scatter(X, Y, yerror=σY,fmt = :png,
    shape = :circle, color = :black, label="data", legend=true)
    p = plot!(p, X, gx.(X), lw=2, label=lbl, legend=true, fmt = :png)
    xlabel!(xs)
    xlabel!(ys)
    return fg, p
end

"""
	fitg1(x, xs, xmin, xmax, xgmin, xgmax; bins=100)

returns the fit and the plot
"""
function fitg1(x, xs, bins, xmin, xmax;
	           xgmin, xgmax, fbins=100, norm=true, fm=0.0, flex_mean=false)

	h, p = hist1d(x, xs, bins, xmin, xmax, norm=norm, legend=true)
    if flex_mean
        fg = fit_gauss(x, xgmin, xgmax, bins=fbins, norm=norm)
    else
        fg = fit_gauss_fm(x, xgmin, xgmax, bins=fbins, norm=norm, fm=fm)
    end
	gx = fg.g[1]
	X = centers(h)
	lbl = @sprintf "σ =%4.1f " fg.std[1]
    p = plot!(p, X, gx.(X), lw=2, label=lbl, legend=true, fmt = :png)
	xlabel!(xs)
	return fg, p
end

"""
	gfit_gauss2_cmean(y, xmin, xmax, bins, sigmas, cs, cmean=0.0)

Fit a double gaussian (with sigmas -->[sigma1, sigma2] cs -->[c1, c2] )
and a cmean to data.
"""
function gfit_gauss2_cmean(y::Vector{Float64}, xmin::Float64, xmax::Float64,
	                       bins::Integer, sigmas::Vector{Float64}, cs::Vector{Float64},
						   norm=false, cmean=0.0)

    x =  in_range(y, xmin, xmax)
    h = hist1d(x, bins, xmin, xmax, norm)
    c = centers(h)
    w = h.weights
    @debug "histo centers and weights in full region"  w c

    g2 = gauss2fm(cmean)
    # fit parameters lb, ub, po are lower, upper bounds and pars

    lb = [cs[1]/100.0, sigmas[1]/5.0, cs[2]/100.0, sigmas[2]/5.0]
    ub = [cs[1]*100.0, sigmas[1]*5.0, cs[2]*100.0, sigmas[2]*5.0]
    p0_bounds = [cs[1], sigmas[1], cs[2], sigmas[2]]

    @debug "pars" p0 lb ub
    # fit double gaussian
    fq = curve_fit(g2, c, w, p0_bounds, lower=lb, upper=ub)
    C1, sigma1, C2,  sigma2   = coef(fq)
    @debug "C1 sigma1 C2 sigma2" C1 sigma1 C2  sigma2

    #
    gx = gausg2(cmean, sigma1, C1, cmean, sigma2, C2)
    gx1 = gausg(cmean, sigma1, C1)
    gx2 = gausg(cmean, sigma2, C2)
    #return (sigma1 = sigma1, sigma2 = sigma2, C1 = C1, C2=C2,
    #        h = h, xg = c, yg = gx.(c), gx = gx, gx1=gx1, gx2=gx2)

	return FGauss([cmean, cmean], [sigma1, sigma2], [C1, C2],
			       h, c, gx.(c), [gx, gx1, gx2])
    end

"""
	fit_2gauss_cmean(data, gp, g1p, g2p, cm)

Fit two gaussian with common mean
"""
function fit_2gauss_cmean(data, gp, g1p, g2p, cm, norm=false)
    gf1 = fit_gauss_fm(data, g1p.xmin,g1p.xmax, bins=g1p.nbin, norm=norm, fm=cm)
    @debug gf1
    gf2 = fit_gauss_fm(data, g2p.xmin,g2p.xmax,bins=g2p.nbin, norm= norm, fm=cm)
    @debug gf2
    gf = gfit_gauss2_cmean(data, gp.xmin,gp.xmax,gp.nbin,
	                       [gf1.std[1], gf2.std[1]], [gf1.C[1], gf2.C[1]], norm)
    @debug gf
    return gf
end

# fit_gauss2(x::Vector{Float64},
#            xmin::Vector{Float64},
# 		   xmax::Vector{Float64},
# 		   bins::Vector{Int64}) = gfit_gauss2(x, xmin, xmax, bins)

"""
    fitg2(x, xs, xmin, xmax, xg1min, xg1max, xg2min, xg2max, xgmin, xgmax; bins=100)
Fits 2 gaussians with common mean (0 by default) to vector x.
"""
function fitg2(x, xs, bins, xmin, xmax;
	           xg1min, xg1max, xg2min, xg2max, xgmin, xgmax, cm=0.0,
         	   g1bins=100, g2bins=100, gbins=100, norm=true)

	hp, p = hist1d(x, xs, bins, xmin, xmax, norm=norm, legend=true)

    g1p = (xmin = xg1min, xmax = xg1max, nbin=g1bins)
    g2p = (xmin= xg2min, xmax=  xg2max, nbin=g2bins)
    gp  = (xmin= xgmin, xmax=  xgmax, nbin=gbins)

    fg  = fit_2gauss_cmean(x, gp, g1p, g2p, cm, norm)
	gx  = fg.g[1]
	gx1 = fg.g[2]
	gx2 = fg.g[3]

	lbl = @sprintf "σt =%4.1f mm, σ =%4.1f mm" fg.std[1] fg.std[2]
	st  = @sprintf "σt =%4.1f mm " fg.std[1]
	sf  = @sprintf "σ  =%4.1f mm" fg.std[2]
    p = plot!(p,fg.X, fg.Y, label=lbl, lw=2, fmt = :png)
    p = plot!(p,fg.X, gx1.(fg.X), label=st, lw=1, fmt = :png)
    p = plot!(p,fg.X, gx2.(fg.X), label=sf, lw=1, fmt = :png)
    #xlabel!(xs)
    return fg, p
end

"""
	fit_profile(x1, x2, tx1, ty1, fit="pol1", bins=25)
    fit_profile(df1, c1, c2, tx1, ty1, fit="pol1", bins=25)
Create and fit a profile with pol1 or poli2 functions.
Return fit parameters, fit function and plot
"""
fit_profile(df1::DataFrame, c1::String, c2::String,
            tx1::String, ty1::String, fit="pol1", bins=25;
            ybin_width::Real=0.1, ymin::Real=352.4, ymax::Real=392.4,
            min_proportion::Real=0.0) =
	fit_profile(df1[!,c1], df1[!,c2], tx1, ty1, fit, bins,
                ybin_width=ybin_width, ymin=ymin, ymax=ymax, min_proportion=min_proportion)

function fit_profile(x1::Vector{<:Real}, x2::Vector{<:Real},
	                 tx1::String, ty1::String, fit="pol1", bins=25;
                     ybin_width::Real=0.1, ymin::Real=minimum(x2),
                     ymax::Real=maximum(x2), min_proportion::Real=0.0)

    pdf1, _ = p1df(x1,x2, bins, ybin_width=ybin_width, ymin=ymin, ymax=ymax, min_proportion=min_proportion)

    if fit == "pol1"
        fr = fit_pol1(pdf1.x_mean, pdf1.y_mean)
    elseif fit == "pol2"
        fr = fit_pol2(pdf1.x_mean, pdf1.y_mean)
	elseif fit == "pol3"
		fr = fit_pol3(pdf1.x_mean, pdf1.y_mean)
	else
		println("option not implemented")
		return nothing
    end

    p1 = scatter(pdf1.x_mean,pdf1.y_mean, yerror=pdf1.y_std,fmt = :png,
	          shape = :circle, color = :black, legend=false)
    p1 = plot!(p1, pdf1.x_mean, fr.g.(pdf1.x_mean), fmt = :png)
    xlabel!(tx1)
    ylabel!(ty1)

   return fr, p1
end


function log_utility(;β = 0.9,
                      ψ = 0.69,
                      Π = 0.5 * ones(2, 2),
                      G = [0.1, 0.2],
                      Θ = ones(2),
                      transfers = false)
    # Derivatives of utility function
    U(c,n) = log(c) + ψ * log(1 - n)
    Uc(c,n) = 1 ./ c
    Ucc(c,n) = -c.^(-2.0)
    Un(c,n) = -ψ ./ (1.0 - n)
    Unn(c,n) = -ψ ./ (1.0 - n).^2.0
    n_less_than_one = true
    return Model(β, Π, G, Θ, transfers,
                 U, Uc, Ucc, Un, Unn, n_less_than_one)
end


# Patrick Keast, MODERATE-DEGREE TETRAHEDRAL QUADRATURE FORMULAS
# http://mech.fsv.cvut.cz/oofem/resources/doc/oofemrefman/gaussintegrationrule_8C_source.html
function _get_gauss_tetdata(n::Int)
    if n == 1
        a = 1. / 4.
        w = 1. / 6.
        xw = [a a a w]
    elseif n == 2
        a = ( 5. + 3. * √(5.) ) / 20.
        b = ( 5. - √(5.) ) / 20.
        w = 1. / 24.
        xw = [a b b w
              b a b w
              b b a w
              b b b w]
    elseif n == 3
        a1 = 1. / 4.
        a2 = 1. / 2.
        b2 = 1. / 6.
        w1 = -2. / 15.
        w2 = 3. / 40.
        xw = [a1 a1 a1 w1
              a2 b2 b2 w2
              b2 a2 b2 w2
              b2 b2 a2 w2
              b2 b2 b2 w2]
    elseif n == 4
        a1 = 1. / 4.;
        w1 = -74. / 5625.;

        a2 = 5. / 70.;
        b2 = 11. / 14.;
        w2 = 343. / 45000.;

        a3 = ( 1. + √(5. / 14.) ) / 4.;
        b3 = ( 1. - √(5. / 14.) ) / 4.;
        w3 = 28. / 1125.;

        xw = [a1 a1 a1 w1
              b2 a2 a2 w2
              a2 b2 a2 w2
              a2 a2 b2 w2
              a2 a2 a2 w2
              a3 a3 b3 w3
              a3 b3 a3 w3
              a3 b3 b3 w3
              b3 a3 a3 w3
              b3 a3 b3 w3
              b3 b3 a3 w3]
    else
        throw(ArgumentError("unsupported order for tetraheder gauss-legendre integration"))
    end
    return xw
end


function _source_f(l::AbstractVector, r::AbstractVector, ϵ::Real)

    L = dot(r, l)

    non_normal = r - L * l

    norm_non_normal = sqrt(norm(non_normal) ^ 2 + ϵ ^ 2)
    nr = sqrt(norm(r) ^ 2 + ϵ ^ 2)

    (l - non_normal * (L / norm_non_normal ^ 2)) / nr

end

"""
Source line influence coefficient
"""
function source_line(colpt::AbstractVector, r1::AbstractVector, r2::AbstractVector, ϵ::Real)

    l = r2 - r1
    nl = norm(l) + 1e-10
    @. l /= nl

    (_source_f(l, colpt - r2, ϵ) - _source_f(l, colpt - r1, ϵ)) / (4 * pi)

end

function _gradPhi(r⃗::AbstractVector, ξ⃗::AbstractVector, η⃗::AbstractVector, ϵ::Real)

    rxi = dot(r⃗, ξ⃗)
    reta = dot(r⃗, η⃗)

    r = sqrt(norm(r⃗) ^ 2 + ϵ ^ 2)

    den = (r ^ 2 - rxi ^ 2 + ϵ ^ 2) * r

    (
        (rxi * η⃗ + reta * ξ⃗) * den - (den * r⃗ / r ^ 2 + 2 * (r⃗ - rxi * ξ⃗) * r) * rxi * reta
    ) / den ^ 2

end

"""
Doublet line influence coefficient
"""
function doublet_line(colpt::AbstractVector, r1::AbstractVector, r2::AbstractVector, η⃗::AbstractVector, ϵ::Real)

    l = r2 - r1
    nl = norm(l) + 1e-10
    @. l /= nl

    (_gradPhi(colpt - r2, l, η⃗, ϵ) - _gradPhi(colpt - r1, l, η⃗, ϵ)) / (4 * pi)

end

"""
Point doublet influence coefficient
"""
function point_doublet_infl(colpt::AbstractVector, pt::AbstractVector, ξ::AbstractVector, ϵ::Real)

    r = colpt .- pt
    nr = sqrt(norm(r) ^ 2 + ϵ ^ 2)

    z = ξ ⋅ r
    nprojr = r .- z .* ξ

    (3 * z) .* nprojr ./ (nr ^ 5 * 4 * π) .- ξ .* ((nr ^ 2 - 3 * z ^ 2) / (4 * π * nr ^ 5))

end

"""
Point source influence coefficient
"""
function point_source_infl(colpt::AbstractVector, pt::AbstractVector, ϵ::Real)

    r = colpt .- pt

    r ./ (4 * π * sqrt(norm(r) + ϵ ^ 2) ^ 3)

end


#
# 1D bound state energies for a QM particle in an infinite well and potential V(x)
#

## Definition of the orthonormal basis states
@everywhere function φ(i) return  function (x) sin(π*i*x) end end

## Coupling strength
@everywhere λ = 1.0*π^2/4.0

## Definition of the centered harmonic potential V(x)
@everywhere function V(x) return λ*(x-0.5)^2 end

## Multiply a potential to calculate <φ1|V|φ2>
@everywhere function Vφ(i) return  function (x) sin(π*i*x)*V(x) end end

## Definition of the inner product <φ1|φ2>
@everywhere function ∫(φ1::Function,φ2::Function)
    ψ(x)=φ1(x)*φ2(x)::Float64
    return 2*quadgk(ψ,0,1,abstol=1e-8)[1]::Float64
end

##  Infinite well energies
@everywhere function ɛ(i::Int,j::Int)
    return  i == j ? i^2 : 0
end

# The numerical Hamiltonian
function hamiltonian(n::Int)
  h = SharedArray(Float64,n^2)
  @sync for j=1:n
    @parallel for i=j:n
      h[i+(j-1)*n] = ɛ(i,j) + ∫(φ(i),Vφ(j))::Float64
     end
  end;
  H = reshape(h,n,n)
  return H+transpose(H-triu(H))
end

# The exact Hamiltonian
function hamiltonian_exact(n::Int)
  return Float64[ i==j? i^2 + (1/12)*λ*(1-6/(i*π)^2) : 0.25*(1+(-1)^(i+j))*(1/(i-j)^2-1/(i+j)^2) for i=1:n,j=1:n ]
end

## Checking orthonomality
function orthonormal(n::Int)
    v = SharedArray(Float64,n^2)
    @sync for j=1:n
        @parallel for i=j:n
            v[i+(j-1)*n] = ∫(φ(i),φ(j))::Float64
        end
    end;
    m = reshape(v,n,n)
    return m+transpose(m-triu(m))
end

function main(n::Int)
   @time println("norm of difference = ", norm(eigvals(hamiltonian(n))-eigvals(hamiltonian_exact(n))))
end

main(400)


## linear.jl : implementation of linear covariance function

## Linear ##
struct Linear{T} <: IsotropicCovarianceStructure{T}
    λ::T
    σ::T
    p::T

    function Linear{T}(λ::T, σ::T, p::T) where T
        λ > 0 || throw(DomainError(λ, "correlation length λ of linear covariance cannot be negative or zero"))
        σ > 0 || throw(DomainError(σ, "marginal standard deviation σ of linear covariance cannot be negative or zero"))
        p >= 1 || throw(DomainError(p, "in p-norm, p must be greater than or equal to 1"))
        isinf(p) && throw(DomainError(p, "in p-norm, p cannot be infinity"))

        new{T}(λ, σ, p)
    end
end

"""
    Linear(λ, [σ = 1], [p = 2])

Linear covariance structure with correlation length `λ`, (optional) marginal standard deviation `σ` and (optional) `p`-norm, defined as

``C(x, y) = \\begin{cases} σ \\left(1 - \\displaystyle\\frac{ρ}{λ}\\right) & \\text{if }ρ ≤ λ\\\\ 0 & \\text{if }ρ>λ\\end{cases}``

with ``ρ = ||x - y||_p``.

# Examples
```jldoctest
julia> Linear(0.1)
linear (λ=0.1, σ=1.0, p=2.0)

julia> Linear(1.0, σ=2)
linear (λ=1.0, σ=2.0, p=2.0)

```
See also: [`Exponential`](@ref), [`Spherical`](@ref), [`Whittle`](@ref), [`Gaussian`](@ref), [`SquaredExponential`](@ref), [`Matern`](@ref)
"""
Linear(λ::Real; σ::Real=1.0, p::Real=2) = Linear{promote_type(typeof(λ),typeof(σ),typeof(p))}(promote(λ, σ, p)...)

# evaluate exponential covariance
apply(l::Linear, x::Real) = max(zero(x), 1 - x / l.λ)

# short name
shortname(::Linear) = "linear"


### A Pluto.jl notebook ###
# v0.14.0

using Markdown
using InteractiveUtils

# ╔═╡ 15b8afb9-ace9-4b5b-894e-dce9dda5bf5f
begin
	using Pkg

	try
		using UnicodePlots

	catch ArgumentError
		Pkg.add("UnicodePlots")
		using UnicodePlots
	end

	using Test
	using PlutoUI
end

# ╔═╡ 45d72214-919c-11eb-2e69-23c0786baacb
md"""
### Functional Programming

Src: [Julia 1.0 Cookbook -  Bogumił Kamiński , Przemysław Szufel - 2018](https://www.packtpub.com/product/julia-1-0-programming-cookbook/9781788998369)
"""

# ╔═╡ 6d80a12a-9649-4438-9a12-36003a57bf86
function deriv(f::Function)::Function
	h = √eps()
	df(x) = (f(x + h) - f(x)) / h
	df
end

# ╔═╡ 5e5c6b37-f3ee-42b1-94cc-686d54171833
md"""
Let us test the `deriv` function. We will use the `UnicodePlots.jl` package to plot a $2×x×x + 5×x - 4$ function and its derivative.
"""

# ╔═╡ b3faae61-dcdc-473d-9a06-cc890ae74bd7
begin
	f(x) = 2x*x + 5x - 4
	x = -5.:3.;
end

# ╔═╡ 75442ce9-cf92-4800-b1aa-bfdd4229224c
begin
	plot = lineplot(x, f.(x), width=45, height=15, canvas=DotCanvas, name="f(x)");
	plot = lineplot!(plot, x, deriv(f).(x), name="f'(x)")
end

# ╔═╡ a2db212a-02f8-493c-a1f2-43c0e63c31fa
md"""
Next, we can construct a function that solves any quadratic equation in the form

$$ax^2 + bx +c = 0$$
"""

# ╔═╡ 415677ce-d5d6-4808-8789-f5fa66d7f7df
function q_solve(f)
	c = f(0.)
	f1 = deriv(f)
	b = f1(0.)
	a = f(1,) - b - c
	d = √(b * b - 4 * a * c)
	((-b - d) / 2a, -(-b + d) / 2a)
end

# ╔═╡ bd8d4613-c666-41e3-b877-e9ed6882fe04
@test q_solve(x -> (x - 1) * (x + 7)) == (-7., -1.)

# ╔═╡ b3ee6944-32c7-4dca-ab51-9f1dc0accd91
@test_throws DomainError q_solve(x -> x * x + 1)  ## no solution in the real domain

# ╔═╡ 6486f3b7-9a29-41d1-8d80-514a9212a174
with_terminal() do
	println(q_solve(x -> x * x + 1 + 0im))
	## (-7.450580707946e-9 - 1.0000000074505im, 7.450580707946e-9 - 1.0000000074505im)
end

# ╔═╡ 75a80d9a-953e-41c0-9569-2f1c08ca16cd
md"""
Our quadratic solver function takes any quadratic function as its argument. Such functions can be presented as $f(x) = a×x^2 + b×x + c$, and we calculate $f(0)$ to get the $c$ value.
Next, we calculate `deriv(f)(0)`, which is $(2×a×x + b)(0)$, and hence obtain a numerical approximation of the value of b.
Finally, we calculate a, and now we can calculate root values of the `f` function.

Since we did not specify types, our function works not only for real numbers but also for complex numbers.
This can be observed when comparing the calls `q_solve(x -> x*x + 1)` and
`q_solve(x -> x*x + 1 + 0im)`.
The first one fails because `x × x + 1` does not have a solution in the real domain. On the other hand - the second one - using complex numbers, works.
"""

# ╔═╡ Cell order:
# ╟─45d72214-919c-11eb-2e69-23c0786baacb
# ╠═15b8afb9-ace9-4b5b-894e-dce9dda5bf5f
# ╠═6d80a12a-9649-4438-9a12-36003a57bf86
# ╟─5e5c6b37-f3ee-42b1-94cc-686d54171833
# ╠═b3faae61-dcdc-473d-9a06-cc890ae74bd7
# ╠═75442ce9-cf92-4800-b1aa-bfdd4229224c
# ╟─a2db212a-02f8-493c-a1f2-43c0e63c31fa
# ╠═415677ce-d5d6-4808-8789-f5fa66d7f7df
# ╠═bd8d4613-c666-41e3-b877-e9ed6882fe04
# ╠═b3ee6944-32c7-4dca-ab51-9f1dc0accd91
# ╠═6486f3b7-9a29-41d1-8d80-514a9212a174
# ╟─75a80d9a-953e-41c0-9569-2f1c08ca16cd


# The number, 197, is called a circular prime because all rotations of the
# digits: 197, 971, and 719, are themselves prime.
#
# There are thirteen such primes below 100: 2, 3, 5, 7, 11, 13, 17, 31, 37, 71,
# 73, 79, and 97.
#
# How many circular primes are there below one million?

using ProjectEulerSolutions

# Helper function to generate an array of circular shifts given an input
# integer.
function circularnumbers(n::Integer)::Array{Integer}
    digs = digits(n)
    numdigits = length(digs)
    multipliers = 10 .^ (0:numdigits-1)

    circular_nums = Integer[]
    for i in 1:numdigits-1
        push!(circular_nums, sum(multipliers .* circshift(digs, i)))
    end
    return circular_nums
end

# Use Sieve of Eratosthenes to get primes, store in set, then cycle through
# them creating circular shifts and checking whether all shifts are prime.  If
# so, perform some set magic to move those primes into the verified set and
# remove from the potential circular primes set.
function p035solution(n::Integer=100)::Integer
    primes = Set{Integer}(sieve_eratosthenes(n))
    verified = Set{Integer}()

    for prime in primes
        if ndigits(prime) == 1
            push!(verified, prime)
        else
            if mapreduce(x-> x % 2, *, digits(prime)) == 0
                # If any digit is even, ignore number
                continue
            end

            circnums = circularnumbers(prime)
            if mapreduce(prime -> in(prime, primes), &, circnums)
                push!(verified, prime)
                union!(verified, circnums)
                delete!(primes, prime)
                setdiff!(primes, circnums)
            end
        end
    end

    return length(verified)
end

p035 = Problems.Problem(p035solution)

Problems.benchmark(p035, 1_000_000)


# Solves the Poisson equation using the Mixed finite element method
using Revise
using AdFem
using DelimitedFiles
using SparseArrays
using PyPlot

n = 50
mmesh = Mesh(n, n, 1/n, degree = BDM1)

function ufunc(x, y)
    x * (1-x) * y * (1-y)
end

function ffunc(x, y)
    -2x*(1-x) -2y*(1-y)
end

function gfunc(x, y)
    (1-2x)*y*(1-y)
end

A = compute_fem_bdm_mass_matrix1(mmesh)
B = compute_fem_bdm_div_matrix1(mmesh)
C = [A -B'
    -B spzeros(mmesh.nelem, mmesh.nelem)]

gD = (x1, y1, x2, y2)->!( (x1<1e-3 && y1>1e-3 && y1<1-1e-3) ||  (x2<1e-3 && y2>1e-3 && y2<1-1e-3))
gN = (x1, y1, x2, y2)->x1<1e-5 && x2<1e-5

D_bdedge = bcedge(gD, mmesh)
N_bdedge = bcedge(gN, mmesh)

t1 = eval_f_on_boundary_edge(ufunc, D_bdedge, mmesh)
g = compute_fem_traction_term1(t1, D_bdedge, mmesh)
t2 = eval_f_on_gauss_pts(ffunc, mmesh)
f = compute_fvm_source_term(t2, mmesh)

gN = eval_f_on_boundary_edge(gfunc, N_bdedge, mmesh)
dof, val = impose_bdm_traction_boundary_condition1(gN, N_bdedge, mmesh)
rhs = [-g; f]
C, rhs = impose_Dirichlet_boundary_conditions(C, rhs, dof, val)

sol = C\rhs
u = sol[mmesh.ndof+1:end]
close("all")
figure(figsize=(15, 5))
subplot(131)
title("Reference")
xy = fvm_nodes(mmesh)
x, y = xy[:,1], xy[:,2]
uf = ufunc.(x, y)
visualize_scalar_on_fvm_points(uf, mmesh)
subplot(132)
title("Numerical")
visualize_scalar_on_fvm_points(u, mmesh)
subplot(133)
title("Absolute Error")
visualize_scalar_on_fvm_points( abs.(u - uf) , mmesh)
savefig("mixed_poisson_neumann.png")


using Distributions

function sliceSampleStep(logf, x0, y0, w=1.0)
    newy = y0 + log(rand())

    # step outward
    L = x0 - w * rand()
    R = L + w
    yL = logf(L)
    yR = logf(R)
    while yL > newy || yR > newy
        if rand(Bool)
            L -= R - L
            yL = logf(L)
        else
            R += R - L
            yR = logf(L)
        end
    end

    #Step inward
    x = rand(Uniform(L, R))
    y = logf(x)
    while y < newy
        if x < x0
            L = x
            yL = logf(L)
        else
            R = x
            yR = logf(R)
        end
        x = rand(Uniform(L, R))
        y = logf(x)
    end

    return (x, y)
end

function sliceSample(logf, n, x0=0.0, y0=0.0, w=1.0)
    (x,y) = (x0,logf(x0))
    xs = zeros(n)
    for i in 1:n
        (x,y) = sliceSampleStep(logf, x, y, w)
        xs[i] = x
    end

    return xs
end

function test(n)
    sliceSample(x -> logpdf(Normal(),x), n)
end

x = test(10000)
scatter(x[2:10000],x[1:9999],α=0.1)


# MIT license
# Copyright (c) Microsoft Corporation. All rights reserved.
# See LICENSE in the project root for full license information.

"""
Module to enclose [Zernike polynomial](https://en.wikipedia.org/wiki/Zernike_polynomials) specific functionality.
"""
module Zernike

"""
    OSAtoNM(j::Int) -> Tuple{Int, Int}

Convert OSA zernike index `j` to `(N,M)` form according to formula `J = N * (N + 2) + M`.
"""
function OSAtoNM(j::Int)::Tuple{Int, Int}
    n = Int(ceil((-3 + sqrt(9 + 8j)) / 2))
    m = 2j - n * (n + 2)
    return (Int(n), Int(m))
end

"""
    NolltoNM(j::Int) -> Tuple{Int, Int}

Convert Noll zernike index `j` to `(N,M)` form.
"""
function NolltoNM(j::Int)
    n = Int(ceil((-3 + sqrt(1 + 8j)) / 2))
    jr = j - Int(n * (n + 1) / 2)
    if mod(n, 4) ∈ (0, 1)
        m1 = jr
        m2 = -(jr - 1)
        if iseven(n - m1)
            m = m1
        else
            m = m2
        end
    else # mod(n,4) ∈ (2,3)
        m1 = jr - 1
        m2 = -(jr)
        if iseven(n - m1)
            m = m1
        else
            m = m2
        end
    end
    return (Int(n), Int(m))
end

"""
    normalisation(::Type{T}, N::Int, M::Int) -> T

Normalisation coefficient for Zernike polynomial term ``Z_{n}^{m}``.
"""
@inline function normalisation(::Type{T}, N::Int, M::Int) where {T <: Real}
    return T(sqrt((2 * (N + 1)) / (1 + (M == 0 ? 1 : 0))))
end

"""
    ζ(N::Int, M::Int, ρ::T, ϕ::T) -> Tuple{T,T}

Evaluate Zernike polynomial term ``Z_{n}^{m}(\\rho, \\phi)``.
"""
@inline function ζ(N::Int, M::Int, ρ::T, ϕ::T)::T where {T <: Real}
    aM = abs(M)
    if M < 0
        return normalisation(T, N, M) * R(N, aM, ρ) * sin(aM * ϕ)
    else
        return normalisation(T, N, M) * R(N, aM, ρ) * cos(aM * ϕ)
    end
end

"""
    δζ(N::Int, M::Int, ρ::T, ϕ::T) -> Tuple{T,T}

Evaluate partial derivatives of Zernike polynomial term ``Z_{n}^{m}(\\rho, \\phi)``.
"""
@inline function δζ(N::Int, M::Int, ρ::T, ϕ::T)::Tuple{T, T} where {T <: Real}
    n = normalisation(T, N, M)
    aM = abs(M)
    RNM = R(N, aM, ρ)
    ρ2 = ρ^2
    δRNM =
        (
            (2 * N * aM * (ρ2 - 1) + (N - aM) * (aM + N * (2 * ρ2 - 1))) * RNM -
            (N + aM) * (N - aM) * R(N - 2, aM, ρ)
        ) / (2 * N * ρ * (ρ2 - 1))
    Mϕ = aM * ϕ
    if M < 0
        δρ = n * δRNM * sin(Mϕ)
        δϕ = n * RNM * aM * cos(Mϕ)
    else
        δρ = n * δRNM * cos(Mϕ)
        δϕ = n * RNM * aM * -sin(Mϕ)
    end
    return δρ, δϕ
end

"""
    R(N::Int, M::Int, ρ::T) -> T

Evaluate radial polynomial ``R_{n}^{m}(\\rho)``.
"""
@inline function R(N::Int, M::Int, ρ::T)::T where {T <: Real}
    if (N - M) % 2 === 1
        return zero(T)
    end
    total = zero(T)
    @simd for k in 0:((N - M) ÷ 2)
        total +=
            ((-1)^k * factorial(N - k)) / (
                factorial(k) *
                factorial((N + M) ÷ 2 - k) *
                factorial((N - M) ÷ 2 - k)
            ) * ρ^(N - 2 * k)
    end
    return total
end

end # module Zernike

#########################################################################################################

"""
Either `ZernikeIndexingOSA` or `ZernikeIndexingNoll`, see [Zernike polynomials wikipedia entry](https://en.wikipedia.org/wiki/Zernike_polynomials) for details.
"""
@enum ZernikeIndexType ZernikeIndexingOSA ZernikeIndexingNoll

"""
    ZernikeSurface{T,N,P,Q,M} <: ParametricSurface{T,N}

Surface incorporating the Zernike polynomials - radius, conic and aspherics are defined relative to absolute semi-diameter, Zernike terms are normalized according to the `normradius` parameter.
`T` is the datatype, `N` is the dimensionality, `P` is the number of Zernike terms, `Q` is the number of aspheric terms and `M` is the Aspheric Type.

The surface is centered at the origin and treated as being the cap of an infinite cylinder, thus creating a true half-space.
Outside of `0 <= ρ <= 1` the height of the surface is not necessarily well defined, so NaN may be returned.

For convenience the input `zcoeff` can be indexed using either OSA or Noll convention, indicated using the `indexing` argument as either `ZernikeIndexingOSA` or `ZernikeIndexingNoll`.

```julia
ZernikeSurface(semidiameter, radius = Inf, conic = 0, zcoeff = nothing, aspherics = nothing, normradius = semidiameter, indexing = ZernikeIndexingOSA)
```

`zcoeff` and `aspherics` should be vectors containing tuples of the form `(i, v)` where `i` is either the index of the Zernike term
for the corresponding `indexing`, or the polynomial power of the aspheric term (may be even or odd) and
`v` is the corresponding coefficient ``A_i`` or ``\\alpha_i`` respectively..  `M` will be determined from the terms entered to optimize
the evaluation of the aspheric polynomial.

The sag is defined by the equation

```math
z(r,\\phi) = \\frac{cr^2}{1 + \\sqrt{1 - (1+k)c^2r^2}} + \\sum_{i}^{Q}\\alpha_ir^{2i} + \\sum_{i}^PA_iZ_i(\\rho, \\phi)
```

where ``\\rho = \\frac{r}{\\texttt{normradius}}``, ``c = \\frac{1}{\\texttt{radius}}``, ``k = \\texttt{conic}`` and ``Z_n`` is the nᵗʰ Zernike polynomial.
"""
struct ZernikeSurface{T, N, P, Q, M} <: ParametricSurface{T, N}
    asp::AsphericSurface{T, N, Q, M}
    coeffs::SVector{P, Tuple{Int, Int, T}}
    boundingcylinder::Cylinder{T, N}

    function ZernikeSurface(
        semidiameter::T;
        radius::T = typemax(T),
        conic::T = zero(T),
        zcoeff::Union{Nothing, Vector{Tuple{Int, T}}} = nothing,
        aspherics::Union{Nothing, Vector{Tuple{Int, T}}} = nothing,
        normradius::T = semidiameter,
        indexing::ZernikeIndexType = ZernikeIndexingOSA,
    ) where {T <: Real}
        asp =
            AsphericSurface(semidiameter; radius, conic, aspherics, normradius)
        Q = length(asp.aspherics) #this is not the same as the aspherics variable passed to the function!

        zcs = []
        if zcoeff !== nothing
            for (i, k) in zcoeff
                if abs(k) > zero(T)
                    if indexing === ZernikeIndexingOSA
                        R, S = Zernike.OSAtoNM(i)
                    else
                        R, S = Zernike.NolltoNM(i)
                    end
                    push!(zcs, (R, S, k))
                end
            end
        end
        P = length(zcs)
        M = asphericType(asp)
        return new{T, 3, P, Q, M}(
            asp::AsphericSurface{T, 3, Q, M},
            SVector{P, Tuple{Int, Int, T}}(zcs),
            Cylinder(semidiameter, interface = opaqueinterface(T)),
        ) # TODO!! incorrect interface on cylinder
    end
end

function uvrange(
    ::Type{ZernikeSurface{T, N, P, Q, M}},
) where {T <: Real, N, P, Q, M}
    return ((zero(T), one(T)), (-T(π), T(π)))
end # ρ and ϕ

semidiameter(z::ZernikeSurface{T}) where {T <: Real} = z.asp.semidiameter
halfsizeu(z::ZernikeSurface{T}) where {T <: Real} = semidiameter(z)
halfsizev(z::ZernikeSurface{T}) where {T <: Real} = semidiameter(z)

boundingobj(z::ZernikeSurface{T}) where {T <: Real} = z.boundingcylinder

function point(
    z::ZernikeSurface{T, 3, P, Q, M},
    ρ::T,
    ϕ::T,
)::SVector{3, T} where {T <: Real, P, Q, M}
    pnt = point(z.asp, ρ, ϕ)

    # sum zernike
    rad = semidiameter(z.asp)
    r = ρ * rad
    u = r / z.asp.normradius
    h = zero(T)
    @inbounds @simd for m in 1:P
        (R, S, k) = z.coeffs[m]
        h += k * Zernike.ζ(R, S, u, ϕ)
    end
    return SVector{3, T}(pnt[1], pnt[2], pnt[3] + h)
end

function partials(
    z::ZernikeSurface{T, 3, P, Q, M},
    ρ::T,
    ϕ::T,
)::Tuple{SVector{3, T}, SVector{3, T}} where {T <: Real, P, Q, M}
    pρ, pϕ = partials(z.asp, ρ, ϕ)
    # sum zernike partials
    rad = z.asp.semidiameter
    n = rad / z.asp.normradius
    u = ρ * n
    dhdρ = zero(T)
    dhdϕ = zero(T)
    @inbounds @simd for m in 1:P
        (R, S, k) = z.coeffs[m]
        du, dϕ = Zernike.δζ(R, S, u, ϕ)
        dhdρ += k * du * n # want the derivative wrt ρ, not u
        dhdϕ += k * dϕ
    end
    return SVector{3, T}(pρ[1], pρ[2], pρ[3] + dhdρ),
    SVector{3, T}(pϕ[1], pϕ[2], pϕ[3] + dhdϕ)
end

function normal(
    z::ZernikeSurface{T, 3, P, Q, M},
    ρ::T,
    ϕ::T,
)::SVector{3, T} where {T <: Real, P, Q, M}
    du, dv = partials(z, ρ, ϕ)
    if ρ == zero(T) && norm(dv) == zero(T)
        # in cases where there is no δϕ at ρ = 0 (i.e. anything which is rotationally symetric)
        # then we get some big problems, hardcoding this case solves the problems
        return SVector{3, T}(0, 0, 1)
    end
    return normalize(cross(du, dv))
end

function uv(
    z::ZernikeSurface{T, 3, P, Q, M},
    p::SVector{3, T},
) where {T <: Real, P, Q, M}
    # avoid divide by zero for ForwardDiff
    ϕ = NaNsafeatan(p[2], p[1])
    if p[1] == zero(T) && p[2] == zero(T)
        ρ = zero(T)
    else
        ρ = sqrt(p[1]^2 + p[2]^2) / semidiameter(z)
    end
    return SVector{2, T}(ρ, ϕ)
end

function onsurface(
    surf::ZernikeSurface{T, 3, P, Q, M},
    p::SVector{3, T},
) where {T <: Real, P, Q, M}
    ρ, ϕ = uv(surf, p)
    if ρ > one(T)
        return false
    else
        surfpoint = point(surf, ρ, ϕ)
        return samepoint(p[3], surfpoint[3])
    end
end

function inside(
    surf::ZernikeSurface{T, 3, P, Q, M},
    p::SVector{3, T},
) where {T <: Real, P, Q, M}
    ρ, ϕ = uv(surf, p)
    if ρ > one(T)
        return false
    else
        surfpoint = point(surf, ρ, ϕ)
        return p[3] < surfpoint[3]
    end
end

#########################################################################################################

# Assumes the ray has been transformed into the canonical coordinate frame which has the vertical axis passing through (0,0,0) and aligned with the z axis.
function surfaceintersection(
    surf::AcceleratedParametricSurface{T, 3, ZernikeSurface{T, 3, P, Q, M}},
    r::AbstractRay{T, 3},
) where {T <: Real, P, Q, M}
    cylint = surfaceintersection(surf.surface.boundingcylinder, r)
    if cylint isa EmptyInterval{T}
        return EmptyInterval(T)
    else
        if doesintersect(surf.triangles_bbox, r) ||
           inside(surf.triangles_bbox, origin(r))
            surfint = triangulatedintersection(surf, r)
            if !(surfint isa EmptyInterval{T})
                return intervalintersection(cylint, surfint)
            end
        end
        # hasn't hit the surface
        if lower(cylint) isa RayOrigin{T} && upper(cylint) isa Infinity{T}
            if inside(surf.surface, origin(r))
                return Interval(RayOrigin(T), Infinity(T))
            else
                return EmptyInterval(T)
            end
            # otherwise check that the intersection is underneath the surface
        else
            p = point(closestintersection(cylint, false))
            ρ, ϕ = uv(surf, p)
            surfpoint = point(surf.surface, ρ, ϕ)
            if p[3] < surfpoint[3]
                return cylint # TODO!! UV (and interface) issues?
            else
                return EmptyInterval(T)
            end
        end
    end
end

function AcceleratedParametricSurface(
    surf::T,
    numsamples::Int = 17;
    interface::NullOrFresnel{S} = NullInterface(S),
) where {S <: Real, N, T <: ZernikeSurface{S, N}}
    # Zernike uses ρ, ϕ uv space so need to modify extension of triangulation
    a = AcceleratedParametricSurface(
        surf,
        triangulate(surf, numsamples, true, false, true, false),
        interface = interface,
    )
    emptytrianglepool!(S)
    return a
end

function BoundingBox(
    surf::ZernikeSurface{T, 3, P, Q, M},
) where {T <: Real, P, Q, M}
    bb = BoundingBox(surf.asp)
    # zernike terms have condition than |Zᵢ| <= 1
    # so this gives us a (loose) bounding box
    ak = if P > 0
        sum(abs.(Zernike.normalisation(T, n, m) * k for (n, m, k) in surf.coeffs))
    else
        zero(T)
    end
    bb.zmin -= ak
    bb.zmax += ak
    return BoundingBox(bb.xmin, bb.xmax, bb.ymin, bb.ymax, bb.zmin, bb.zmax) #could just return bb, but this way is safer
end


@valueequals struct Circle
    x::Number
    y::Number

    r::Number

    Circle(x::Number, y::Number, r::Number) = new(x, y, r)
    Circle(x1::Number, y1::Number, x2::Number, y2::Number) = new(x1, y1, sqrt((x1 - x2)^2 + (y1 - y2)^2))
end

function circumference(circle::Circle)
    return 2 * pi * circle.r
end

# Efficient enough, checks might be too small for huge circles
function pointsOnCircle(circle::Circle; filled::Bool=false, checks::Int=max(1, ceil(Int, circumference(circle))))
    res = Set{Tuple{Number, Number}}()
    step = 2 * pi / checks

    if checks > 0
        for theta in 0:step:2 * pi
            ox = cos(theta) * circle.r
            oy = sin(theta) * circle.r

            push!(res, (round(Int, circle.x + ox), round(Int, circle.y + oy)))
            push!(res, (round(Int, circle.x - ox), round(Int, circle.y + oy)))

            if filled
                for i in floor(Int, -ox + 1):ceil(Int, ox - 1)
                    push!(res, (round(Int, circle.x + i), round(Int, circle.y + oy)))
                end
            end
        end
    end

    return collect(res)
end


# GMM estimation for a sample from Chi^2(theta)
# compare to two method of moments estimators (see chi2mm.m)
using Econometrics, Random, Distributions, Plots
function main()
n = 30
theta = 3.0
reps = 1000
results = zeros(reps,2)

y = zeros(n) # this is just a place holder to define the moments, below
# define the moment conditions
m = theta -> [(theta.- y) (theta .- 0.5*(y .- theta).^2.0)]
for i = 1:1000
    rand!(Chisq(theta), y) # this replaces the place holder, from above
    thetahat, junk, junk, ms, junk = gmm(m, [3.0], eye(2))
    results[i,1] = sqrt(n)*(thetahat[1]-theta)
    W = inv(cov(ms))
    thetahat, junk, ms, junk = gmm(m, thetahat, W)
    results[i,2] = sqrt(n)*(thetahat[1]-theta)
end
p1 = npdensity(results[:,1])
plot!(p1, title="Inefficient")
p2 = npdensity(results[:,2])
plot!(p2, title="Efficient")
p = plot(p1,p2,layout=(2,1),xlims=(-10,20))
#savefig("Efficient.png")
println("Monte Carlo covariance: (1,1) is inefficient GMM, (2,2) is efficient")
prettyprint(cov(results))
return p
end
main()


"""
    milca_weights(cij, a, B)

This function returns weights (a vector of the number of frequency channels) of the Modified ILC Algorith (MILCA) method.

*Reference*: Equation (15) of Hurier, Macías-Pérez, Hildebrandt, A&A, 558, A118 (2013)

# Arguments
- `cij::Array{<:AbstractFloat,2}`: symmetric covariance matrix with the dimention of `(nν, nν)` where `nν` is the number of frequency bands.
- `a::Array{<:AbstractFloat,1}`: vector of the frequency response, for the component to be extracted. E.g., `a = [1,...,1]` for CMB.
- `B::Array{<:AbstractFloat,2}`: `nν`-by-`nrc` matrix of the frequency response, for `nrc` components to be nulled.
"""
function milca_weights(
    cij::Array{T,2},
    a::Array{T,1},
    B::Array{T,2},
) where {T<:AbstractFloat}
    if size(cij)[1] ≠ size(cij)[2]
        throw(DimensionMismatch("covariance matrix must be a square matrix"))
    elseif size(cij)[1] ≠ length(a) || size(cij)[1] ≠ size(B)[1]
        throw(DimensionMismatch("dimensions of the covariance matrix and the frequency response vector or matrix do not match"))
    else
        M = Symmetric(cij)
    end
    e = zeros(size(B)[2] + 1)
    e[1] = 1
    F = [a B]
    w = inv(M) * F * inv(F' * inv(M) * F) * e
end

"""
    milca_weights(cijℓ, a, B[, ℓid=3])

This function returns weights (a `nν`-by-`nℓ` matrix) of the Modified ILC Algorith (MILCA) method.

Here, `nν` is the number of frequency channels and `nℓ` is the number of elements in the relevant domain, e.g., multipoles, band-power bins, pixels, etc.

*Reference*: Equation (15) of Hurier, Macías-Pérez, Hildebrandt, A&A, 558, A118 (2013)

# Arguments
- `cijℓ::Array{<:AbstractFloat,3}`: symmetric covariance matrix with the dimention of `(nℓ, nν, nν)`, `(nν, nℓ, nν)` or `(nν, nν, nℓ)` (default).
- `a::Array{<:AbstractFloat,1}`: vector of the frequency response, for the component to be extracted. E.g., `a = [1,...,1]` for CMB.
- `B::Array{<:AbstractFloat,2}`: `nν`-by-`nrc` matrix of the frequency response, for `nrc` components to be nulled.

# Optional Arguments
- `ℓid::Integer=3`: location of the index for the `nℓ` domain. `ℓid=1` if `cijℓ[nℓ,nν,nν]`, `ℓid=2` if `cijℓ[nν,nℓ,nν]`, and `ℓid=3` (the default value) if `cijℓ[nν,nν,nℓ]`.
"""
function milca_weights(
    cijℓ::Array{T,3},
    a::Array{T,1},
    B::Array{T,2},
    ℓid::Integer = 3,
) where {T<:AbstractFloat}
    if ℓid > 3 || ℓid < 1
        throw(DomainError(ℓid, "ℓid must be 1, 2, or 3"))
    end
    if (ℓid == 3 && size(cijℓ)[1] ≠ size(cijℓ)[2]) ||
       (ℓid == 2 && size(cijℓ)[1] ≠ size(cijℓ)[3]) ||
       (ℓid == 1 && size(cijℓ)[2] ≠ size(cijℓ)[3])
        throw(DimensionMismatch("covariance matrix must be a square matrix"))
    end
    nℓ = size(cijℓ)[ℓid]
    nν = ifelse(ℓid == 3, size(cijℓ)[1], size(cijℓ)[3])
    if length(a) ≠ nν || size(B)[1] ≠ nν
        throw(DimensionMismatch("dimensions of the covariance matrix and the frequency response vector/matrix do not match"))
    end
    wℓ = zeros(nν, nℓ) # ILC weights
    e = zeros(size(B)[2] + 1)
    e[1] = 1
    F = [a B]
    for iℓ = 1:nℓ
        if ℓid == 3
            M = Symmetric(cijℓ[:, :, iℓ])
        elseif ℓid == 2
            M = Symmetric(cijℓ[:, iℓ, :])
        else
            M = Symmetric(cijℓ[iℓ, :, :])
        end
        wℓ[:, iℓ] = inv(M) * F * inv(F' * inv(M) * F) * e
    end
    return wℓ
end


function linear_interpolation(initial_state, final_state, horizon)
    n = length(initial_state)
    X = [copy(Array(initial_state)) for t = 1:horizon]
    for t = 1:horizon
        for i = 1:n
            X[t][i] = (final_state[i] - initial_state[i]) / (horizon - 1) * (t - 1) + initial_state[i]
        end
    end
    return X
end


#   Unit examples of the FourierAnalysis Package for julia language
#
#   MIT License
#   Copyright (c) 2019-2022,
#   Marco Congedo, CNRS, Grenobe, France:
#   https://sites.google.com/site/marcocongedo/home

# ? CONTENTS :
#   This example shows how to compute and plot spectra
#   and how to extract inforation from them.

#   ~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~  #
#                                                                             #
#   ~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~~¤~  #

using FourierAnalysis, FFTW, LinearAlgebra, Statistics,
      Plots, Plots.PlotMeasures

# add module for reading the two EEG text files to be used ater
push!(LOAD_PATH, @__DIR__)
using IOtxt

# Get EEG file names with complete path (they have extension .txt)
S=getFilesInDir(@__DIR__; ext=(".txt",))

# read the two EEG data files and put them in a Matrix object
X1=readEEG(S[1])
X2=readEEG(S[2])

######## compute the FFT of one epoch of a time series ########
sr, t, f, a = 128, 128, 8, 2
v=sinusoidal(a, f, sr, t, 0.6)
plot(v) # plot the time series
tapering=hann

### The long way using FFTW.jl only:
# get a plan.  The normalization 2/t outputs correct peak amplitude
P=plan_rfft(v)*(2/t)

# get a tapering window
a=taper(tapering, t)

# plot the tapered signal
plot!(v.*a.y)

# do the FFT on tapered data
Ψ=P*(v.*a.y)

# get the amplitude spectrum
Σ=abs.(Ψ)
bar!(Σ[2:end])

# check the amplitude at frequency f and neighbooring bins
pos=f2b(f, sr, t, DC=true)

Σ[pos-1]
Σ[pos]
Σ[pos+1]

# The fast way using FourierAnalysis.jl:
Σ2=spectra(v, sr, t; tapering=tapering, func=√)
bar!(Σ2.y)

# equivalently:
Σ2=spectra(v, sr, t; tapering=taper(hann, t), func=√)
bar!(Σ2.y)

####################################################

# test spectra with odd samples and with DC
sr=65; t=65
f, a = fres(sr, t)*4, 2
v=sinusoidal(a, f, sr, t, 0.6; DC=3)
plot(v) # plot the time series
Σ2=spectra(v, sr, t; tapering=rectangular, DC=true, func=√)
bar!(Σ2.y)

####################################################
### Check amplitude spectra at all Fourier Frequency (rectangular taper) ###
# the amplitudes are in increment of 10 along frequencies
# NB when t is even, correct amplitude for the last frequency is obtained
# only if the sinusoidal has a phase=π/6
sr, t, = 16, 32
t½ = t÷2
V=Matrix{Float64}(undef, t, t½)
for i=1:t½ V[:, i]=sinusoidal(10*i, b2f(i, sr, t), sr, t, π/6) end

# plot the sinusoids in V
Z=deepcopy(V)
for i=2:size(Z, 2), j=1:size(Z, 1) Z[j, i]+=(10.0*i) end
plot(Z, label="")

# long way using FFTW only.jl
P=plan_rfft(V, 1)*(2/t)
Ψ=P*V
Ψ[end, :]*=0.5*2^0.25 # correction for Nyquist frequency
Σ=abs.(Ψ)
bar(Σ[brange(t, DC=true), :], labels="")

# fast way using FourierAnalysis.jl
Σ2=spectra(V, sr, t; tapering=rectangular, func=√, DC=true)
bar(Σ2.y[brange(t, DC=true), :], labels="")

#############################################################################
### Check amplitude spectra on long signals obtained by welch methods
# one sinusoidal is at an exact discrete Fourier Frequency and the other not
# Rectangular window
sr, t, f, a = 128, 128, 10, 0.5
v=sinusoidal(a, f, sr, t*16)+sinusoidal(a, f*3.5+0.5, sr, t*16)+randn(t*16)
Σ=spectra(v, sr, t; tapering=rectangular, func=√)
bar(Σ.y, labels="rectangular")

# harris4 window (default)
Σ2=spectra(v, sr, t; func=√)
bar!(Σ2.y, labels="harris4")

#smooth spectra
Σ3=smooth(blackmanSmoother, Σ2)
bar!(Σ3.y, labels="smoothed")

######## Check amplitude spectrum obtained with dpss (slepian tapers) ########
sr, t = 128, 128
t½ = t÷2
V=Matrix{Float64}(undef, t, t½)
for i=1:t½ V[:, i]=sinusoidal(1, b2f(i, sr, t), sr, t, 0 ) end
# plot half of the the first 7 sinusoids in V
plot(V[1:t½+1, 1:7], color=:gray, legend=false)

Σ=spectra(V, sr, t; tapering=taper(slepian, t; α=2, n=3), func=√)
plot(Σ.y, labels="")
#############################################################

#### Compute Welch amplitude spectra of EEG data ####
t, sr, slide, tapering = 1024, 128, 512, harris4

# gather some attributes to obtain nice spectra plots
spectraArgs=(fmax          = 48,
             left_margin   = 2mm,
             bottom_margin = 2mm,
             xtickfont     = font(10, "Times"),
             ytickfont     = font(10, "Times"))

S=spectra(X1, sr, t; tapering=tapering, slide=sr, func=sqrt)
plot(S; ytitle="Amplitude (\\muV)", spectraArgs...)

#smooth spectra
S2=smooth(blackmanSmoother, S)
plot(S2; ytitle="Amplitude (\\muV)", spectraArgs...)

# extract spectra in alpha range (8Hz to 12Hz) at all electrodes
e=extract(S, (8, 12))

# You can use any combination of integer and real numbers
e=extract(S, (8, 12.5))
e=extract(S, (8.5, 12))

# the following equivalently extract the spectra at 10Hz only (1 bin)
e=extract(S, 10)
e=extract(S, (10))
e=extract(S, (10, 10))

# extract spectra at all frequencies. Equivalent to S.y
e=extract(S, :)

# All these ways to indicate frequency ranges work with the mean function
# as well, e.g.,

# compute average spectra in alpha range for each time-series
bar(mean(S, (8, 12)))
# or
bar(mean(S, (8.0,12.0)))

# this actually extract the power at 10Hz for each time series
# and is equivalent to extract(S, 10)
mean(S, 10)
# thus this computes the average power at 10Hz across time-series
mean(extract(S, 10)) # or mean(mean(S, 10))
# and this computes the mean across time-series and across frequencies in range [8Hz, 12Hz]
mean(extract(S, (8, 12))) # or equivalently mean(mean(S, (8, 12)))

# average spectrum across all frequencies for each time-series
bar(mean(S, :))

# plot average spectra in 2Hz-band-pass regions for all time-series
plot(bands(S, 2))

# plot average spectra in in 2Hz-band-pass regions for series 15-19
plot(bands(S, 2)[:, 15:19])

# plot average spectra in in 2Hz-band-pass regions for series 19
plot(bands(S, 2)[:, 19])

# plot average spectra in alpha range for series 15-19
bar(mean(S, (8,12))[15:19])

# extract spectra in alpha range for series 15-19
e=extract(S, (8, 12))[:, 15:19]

# Compute several spectra altogether
𝐗=[X1, X2]
𝐒=spectra(𝐗, sr, t; tapering=tapering, slide=sr, func=√)
plot(𝐒[1]; ytitle="Amplitude (\\muV)", spectraArgs...)
plot(𝐒[2]; ytitle="Amplitude (\\muV)", spectraArgs...)

# do the same thing using a fast FFTW plan (wait up to 10s for computing the plan)
plan=Planner(plan_exhaustive, 10.0, t, eltype(𝐗[1]))
𝐒_fast=spectra(𝐗, sr, t; planner=plan, slide=sr, func=√)

# compute the average spectra in the alpha range for all time-series and all subjects
mean(𝐒, (8,12))

# compute the average spectra in the alpha range across subjects for all time-series
mean(mean(𝐒, (8,12)))
# compute and plot it
plot(mean(mean(𝐒, (8,12))))

# extract spectra in alpha range for all time-series and all subjects
extract(𝐒, (8, 12))

# if you enter en illegal range, nothing will be done and you will get
# an error in the REPL explaining what is wrong in your argument
extract(𝐒, (0, 12))
mean(𝐒, (0, 128))

# extract 4Hz-band-pass average spectra for all electrodes and all subjects
bands(𝐒, 4)

# Apply smoothing in the above spectra computations
𝐒=spectra(𝐗, sr, t;
    tapering=tapering, smoothing=blackmanSmoother, func=√)
plot(𝐒[1]; ytitle="Amplitude (\\muV)", spectraArgs...)
plot(𝐒[2]; ytitle="Amplitude (\\muV)", spectraArgs...)

# plot the average spectrum across all electrodes for the two files
# using Julia standard mean function
plot(mean(𝐒[1].y[:, i] for i=1:size(𝐒[1].y, 2)))
plot!(mean(𝐒[2].y[:, i] for i=1:size(𝐒[2].y, 2)))

# plot spectra in in 1Hz band-pass regions for all electrodes in 𝐒[1]
plot(bands(𝐒[1], 1))

# use slepian multi-tapering on EEG data
S_sl=spectra(X1, sr, t; tapering=slepians(sr, t, 1.25), func=√)
plot(S_sl; ytitle="Amplitude (\\muV)", spectraArgs...)

# use slepian multi-tapering and get smoothed spectra
S_sl=spectra(X1, sr, t;
    tapering=slepians(sr, t, 1.25), func=√, smoothing=blackmanSmoother)
plot(S_sl; xspace=8, ytitle="Amplitude (\\muV)", spectraArgs...)

# See how the variance of the spectra estimation is lower
# as compared to using a normal tapering window !
S=spectra(X1, sr, t;
    tapering=harris4, func=√, smoothing=blackmanSmoother)
plot(S; ytitle="Amplitude (\\muV)", spectraArgs...)


using DrWatson
@quickactivate "StatReth"

# %%
using CSV
using DataFrames
using Turing
using Plots

include(srcdir("quap.jl"))
include(srcdir("tools.jl"))

# %%
# Load the data, don't forget to drop everything missing.
d = DataFrame(CSV.File(datadir("exp_raw/cherry_blossoms.csv"), missingstring = "NA"))
d2 = dropmissing(d, :doy)

# %% 4.76
# Redo the model from the text ...
using BSplines: BSplineBasis, basismatrix

num_knots = 15
knot_list = quantile(d2.year, range(0, 1, length = num_knots))

basis = BSplineBasis(4, knot_list)
B = basismatrix(basis, d2.year)

@model function spline(doy, B = B)
    α ~ Normal(100, 10)
    w ~ filldist(Normal(0, 10), size(B, 2))
    σ ~ Exponential(1)
    μ = α .+ B * w
    doy ~ MvNormal(μ, σ)
end

m4_7 = quap(spline(d2.doy))

# %%
# ... and look at the splines
w_str = ["w[$i]" for i in 1:length(basis)]
post = DataFrame(rand(m4_7.distr, 1000)', ["α"; w_str; "σ"])

w = mean.(eachcol(post[:, w_str]))

plot(legend = false, xlabel = "year", ylabel = "basis * weight")
for y in eachcol(B .* w')
    plot!(d2.year, y)
end
plot!()

mu = B * Array(post[!, w_str])'
mu_m, mu_lower, mu_upper = meanlowerupper(mu)

plot!(d2.year, mu_m, ribbon = (mu_m .- mu_lower, mu_upper .- mu_m),
                    lw = 2, fa = 0.1, c = :black)

# %%
# Now do the entire thing with more knots. The model can stay the same, just need to replace
# B
num_knots = 30
knot_list = quantile(d2.year, range(0, 1, length = num_knots))

basis2 = BSplineBasis(4, knot_list)
B2 = basismatrix(basis2, d2.year)

m4_7_2 = quap(spline(d2.doy, B2))

# %%
# ... and look at the splines
w_str = ["w[$i]" for i in 1:length(basis2)]
post = DataFrame(rand(m4_7_2.distr, 1000)', ["α"; w_str; "σ"])

w = mean.(eachcol(post[:, w_str]))

plot(legend = false, xlabel = "year", ylabel = "basis * weight")
for y in eachcol(B2 .* w')
    plot!(d2.year, y)
end
plot!()

mu = B2 * Array(post[!, w_str])'
mu_m, mu_lower, mu_upper = meanlowerupper(mu)

plot!(d2.year, mu_m, ribbon = (mu_m .- mu_lower, mu_upper .- mu_m),
                    lw = 2, fa = 0.1, c = :black)

# %%
"""
More knots make it possible for the spline to have more wiggles. Since that improves the
fit, it is taken advantage of and the spline gets wigglier. This is pretty similar to how
things go when trying to fit a polynomial while increasing its order. The fit gets better
for the cost of losing generality.
I'm not going to do the prior on the weights thing, I already did the same thing in 4H6.
The amplitude goes up or down, depending in what direction I change the prior.
"""


"""
    pearson_bounds(dA::UnivariateDistribution, dB::UnivariateDistribution; n::Int=7)

Compute the theoretical lower and upper Pearson correlation values for a pair of
univariate distributions.

See also: [`pearson_match`](@ref)

# Examples
```jldoctest
julia> using Distributions

julia> A = Normal(78, 10); B = LogNormal(3, 1);

julia> pearson_bounds(A, B)
(lower = -0.7628739783668452, upper = 0.762873978367046)
```
"""
function pearson_bounds(dA::UD, dB::UD; n::Int=7)
    μA = mean(dA)
    σA = std(dA)
    μB = mean(dB)
    σB = std(dB)
    _pearson_bounds(dA, dB, μA, μB, σA, σB, n)
end
function _pearson_bounds(dA::UD, dB::UD, μA, μB, σA, σB, n)
    k = 0:1:n
    a = get_coefs(dA, n)
    b = get_coefs(dB, n)

    c1 = -μA * μB
    c2 = 1 / (σA * σB)
    kab = factorial.(k) .* a .* b
    ρ_l = c1 * c2 + c2 * sum((-1) .^ k .* kab)
    ρ_u = c1 * c2 + c2 * sum(kab)

    ρ_l, ρ_u = clampcor.((ρ_l, ρ_u))
    (lower = ρ_l, upper = ρ_u)
end

"""
    pearson_bounds(margins::Vector{<:UnivariateDistribution})

Compute the theoretical lower and upper Pearson correlation values for a set of
univariate distributions.

See also: [`pearson_match`](@ref)

# Examples
```jldoctest
julia> using Distributions

julia> m = [Normal(78, 10), LogNormal(3, 1)];

julia> b = pearson_bounds(m);

julia> b.lower
2×2 Array{Float64,2}:
  1.0       -0.762874
 -0.762874   1.0

julia> b.upper
2×2 Array{Float64,2}:
 1.0       0.762874
 0.762874  1.0
```
"""
function pearson_bounds(margins::Vector{<:UD})
    d = length(margins)
    lower, upper = zeros(Float64, d, d), zeros(Float64, d, d)

    @threads for i in collect(subsets(1:d, Val{2}()))
        l, u = pearson_bounds(margins[i[1]], margins[i[2]])
        lower[i...] = l
        upper[i...] = u
    end

    lower .= Symmetric(lower)
    cor_constrain!(lower)

    upper .= Symmetric(upper)
    cor_constrain!(upper)

    (lower = lower, upper = upper)
end


"

 Himmelblau function

 author: Atsushi Sakai
"

function himmelblau_function(x, y)
    return (x^2 + y - 11)^2 + (x + y^2 - 7)^2
end

x = parse(Float64,ARGS[1])
y = parse(Float64,ARGS[2])
println(himmelblau_function(x, y))


using LinearAlgebra, Random, Plots; pyplot()
Random.seed!(1)

N = 10^4
edges = [(1,2), (1,3), (2,4), (1,4), (3,4)]
L = maximum(maximum.(edges))
source, dest = 1, L

function adjMatrix(edges,L)
    R = zeros(Int, L, L)
    for e in edges
        R[ e[1], e[2] ], R[ e[2], e[1] ] = 1, 1
    end
    R
end

pathExists(R, source, destination) = sign.((I+R)^L)[source,destination]
randNet(p) = randsubseq(edges,1-p)

relEst(p) = sum([pathExists(adjMatrix(randNet(p),L),source,dest) for _ in 1:N])/N
relAnalytic(p) = 1-p^3*(p-2)^2

pGrid = 0:0.05:1
scatter(pGrid, relEst.(pGrid),
	c=:blue, ms=5, msw=0,label="Monte Carlo")
plot!(pGrid, relAnalytic.(pGrid),
	c=:red, label="Analytic", xlims=(0,1.05), ylims=(0,1.05),
	xlabel="p", ylabel="Reliability")


# Code to reproduce figure 1 numerical experiments for the paper "Understanding the Origins of Information-Seeking Exploratin in Probabilistic Objectives for Control"
# Essentially what we do here is demonstrate the difference between an evidence "maximizing" objective vs a divergence "probability matching" objective in a simple univariate scenario where we have two overlapping Gaussians
# as our "desire distribution". The evidence objective finds the peak of the distribution and concentrates all probability mass there. The divergence objective explicitly tries to match the desire distribution in all its modes.
# Even though both the desire and evidence objective can fit two Gaussians, the evidence objective nevertheless only utilizes one Gaussian and collapses the weight of the other to zero.

using Plots
using LinearAlgebra
using Statistics
using Distributions
using Optim

e = MathConstants.e

### Utility Function ###

function sigmoid(x)
    return 1 / (1 + exp(-x))
end
function Gaussian_cdf(x, mu, var)
    return (1 / sqrt(2 * pi * var)) * exp(-(x - mu)^2 / (2 * var))
end

function Gaussian_KL(mu1, mu2, var1, var2)
    log(sqrt(var2) / sqrt(var1)) + (var1 + (mu1 - mu2)^2) / (2 * var2) - 0.5
end

function Gaussian_entropy(mu, var)
    return log(sqrt(var * 2 * pi * e))
end

function approximate_KL(cdfs1, cdfs2)
    sum = 0
    for i in 1:length(cdfs1)
        sum += cdfs1[i] * log(cdfs1[i] / cdfs2[i])
    end
    return sum
end

function approximate_evidence(cdfs1, cdfs2)
    return sum(cdfs1 .* -log.(cdfs2))
end
function approximate_entropy(cdfs)
    return sum(cdfs .* log.(cdfs))
end

# Create our desire distribution cdf
function get_desired_cdf(mu_a=1, mu_b=4, var_1=1, var_b=0.4)
    xs = [i*0.01 for i in -500:1000]
    cdfs = (0.5 * Gaussian_cdf.(xs, mu_a, var_a)) + (0.5 * Gaussian_cdf.(xs, mu_b, var_b))
    return cdfs
end

### loss functions for evidence and divergence objective ###

function compute_divergence_loss(mu1, mu2, logvar1, logvar2, alpha,beta)
    lambda = 1e6
    desired_cdfs = get_desired_cdf()
    xs = [i*0.01 for i in -500:1000]
    cdfs = (sigmoid(alpha) .* Gaussian_cdf.(xs, mu1, exp(logvar1))) + (sigmoid(beta) .* Gaussian_cdf.(xs, mu2, exp(logvar2)))
    loss = approximate_KL(cdfs / sum(cdfs), desired_cdfs / sum(desired_cdfs))
    # ensure that constraint is respected
    loss += lambda * ((sigmoid(alpha) + sigmoid(beta)) - 1)^2
    return loss
end
function compute_evidence_loss(mu1, mu2, logvar1, logvar2, alpha,beta)
    lambda = 1e6
    desired_cdfs = get_desired_cdf()
    xs = [i*0.01 for i in -500:1000]
    cdfs = (sigmoid(alpha) .* Gaussian_cdf.(xs, mu1, exp(logvar1))) + (sigmoid(beta) .* Gaussian_cdf.(xs, mu2, exp(logvar2)))
    loss = approximate_evidence(cdfs / sum(cdfs), desired_cdfs / sum(desired_cdfs))
    # ensure that constraint is respected
    loss += lambda * ((sigmoid(alpha) + sigmoid(beta)) - 1)^2
    return loss
end

### Optimize using nonlinear optimization ###

function optimize_loss(loss_fn)
    init_params = zeros(6)
    init_params[1:4] = rand(4)
    print("$init_params \n")
    if loss_fn == "divergence"
        ll_func = compute_divergence_loss
    elseif loss_fn == "evidence"
        ll_func = compute_evidence_loss
    else
        error("Loss function not recognized")
    end
    opt = optimize(init_params -> ll_func(init_params...), init_params)
    output_params = Optim.minimizer(opt)
    final_loss = Optim.minimum(opt)
    return output_params,final_loss
end

### and plot ###

function plot_cdf(params,loss_fn)
    params, final_loss = optimize_loss(loss_fn)
    mu1, mu2, logvar1, logvar2, alpha,beta = params
    xs = [i*0.01 for i in -500:1000]
    desired_cdf = get_desired_cdf()
    cdfs = (sigmoid(alpha) .* Gaussian_cdf.(xs, mu1, exp(logvar1))) + (sigmoid(beta) .* Gaussian_cdf.(xs, mu2, exp(logvar2)))
    plot(xs ./ 10, cdfs/sum(cdfs), xaxis="X value", yaxis="Probability Density", label="Predicted Density")
    if loss_fn == "evidence"
        plot!(xs ./10, desired_cdf, title="Peak Finding with Evidence Objectives",linestyle=:dash,label="Desired Density")
    else
        plot!(xs ./10, desired_cdf / sum(desired_cdf),title="Probability Matching with Divergence objective",linestyle=:dash,label="Desired Density") # this is just for visibility
    end
    savefig("figures/$loss_fn")
    #return cdfs
end

plot_cdf(params,"evidence")
plot_cdf(params, "divergence")


# # Visualization of Frank-Wolfe running on a 2-dimensional polytope

# This example provides an intuitive view of the Frank-Wolfe algorithm
# by running it on a polyhedral set with a quadratic function.
# The Linear Minimization Oracle (LMO) corresponds to a call to a generic simplex solver from `MathOptInterface.jl` (MOI).

# ## Import and setup

# We first import the necessary packages, including Polyhedra to visualize the feasible set.

using LinearAlgebra
using FrankWolfe

import MathOptInterface
const MOI = MathOptInterface
using GLPK

using Polyhedra
using Plots

# We can then define the objective function,
# here the squared distance to a point in the place, and its in-place gradient.

n = 2
y = [3.2, 0.5]

function f(x)
    return 1 / 2 * norm(x - y)^2
end
function grad!(storage, x)
    @. storage = x - y
end

# ## Custom callback
#
# FrankWolfe.jl lets users define custom callbacks to record information about each iteration.
# In that case, the callback will copy the current iterate `x`, the current vertex `v`, and the current step size `gamma`
# to an array thanks to a closure.
# We then declare the array and the callback over this array.
# Each iteration will then push to this array.

function build_callback(trajectory_arr)
    return function callback(state)
        return push!(trajectory_arr, (copy(state.x), copy(state.v), state.gamma))
    end
end

iterates_information_vector = []
callback = build_callback(iterates_information_vector)

# ## Creating the Linear Minimization Oracle
# The LMO is defined as a call to a linear optimization solver, each iteration resets the objective and calls the solver.
# The linear constraints must be defined only once at the beginning and remain identical along iterations.
# We use here MathOptInterface directly but the constraints could also be defined with JuMP or Convex.jl.

o = GLPK.Optimizer()
x = MOI.add_variables(o, n)

## −x + y ≤ 2
c1 = MOI.add_constraint(
    o,
    -1.0x[1] + x[2],
    MOI.LessThan(2.0),
)

## x + 2 y ≤ 4
c2 = MOI.add_constraint(
    o,
    x[1] + 2.0x[2],
    MOI.LessThan(4.0),
)

## −2 x − y ≤ 1
c3 = MOI.add_constraint(
    o,
    -2.0x[1] - x[2],
    MOI.LessThan(1.0),
)

## x − 2 y ≤ 2
c4 = MOI.add_constraint(
    o,
    x[1] - 2.0x[2],
    MOI.LessThan(2.0),
)

## x ≤ 2
c5 = MOI.add_constraint(
    o,
    x[1] + 0.0x[2],
    MOI.LessThan(2.0),
)

# The LMO is then built by wrapping the current MOI optimizer

lmo_moi = FrankWolfe.MathOptLMO(o)

# ## Calling Frank-Wolfe
# We can now compute an initial starting point from any direction
# and call the Frank-Wolfe algorithm.
# Note that we copy `x0` before passing it to the algorithm because it is modified in-place by `frank_wolfe`.

x0 = FrankWolfe.compute_extreme_point(lmo_moi, zeros(n))

xfinal, vfinal, primal_value, dual_gap, traj_data = FrankWolfe.frank_wolfe(
    f,
    grad!,
    lmo_moi,
    copy(x0),
    line_search = FrankWolfe.Adaptive(),
    max_iteration = 10,
    epsilon=1e-8,
    callback=callback,
    verbose=true,
    print_iter=1,
)

# We now collect the iterates and vertices across iterations.

iterates = Vector{Vector{Float64}}()
push!(iterates, x0)
vertices = Vector{Vector{Float64}}()
for s in iterates_information_vector
    push!(iterates, s[1])
    push!(vertices, s[2])
end

# ## Plotting the algorithm run

# We define another method for `f` adapted to plot its contours.

function f(x1, x2)
    x = [x1, x2]
    return f(x)
end

xlist = collect(range(-1, 3, step = 0.2))
ylist = collect(range(-1, 3, step = 0.2))

X = repeat(reshape(xlist, 1, :), length(ylist), 1)
Y = repeat(ylist, 1, length(xlist))

# The feasible space is represented using Polyhedra.

h = HalfSpace([-1, 1], 2) ∩
    HalfSpace([1, 2], 4) ∩
    HalfSpace([-2, -1], 1) ∩
    HalfSpace([1, -2], 2) ∩
    HalfSpace([1, 0], 2)

p = polyhedron(h)

p1 = contour(xlist, ylist, f, fill = true, line_smoothing = 0.85)
plot(p1, opacity = 0.5)
plot!(p, ratio = :equal, opacity = 0.5, label = "feasible region", framestyle = :zerolines, legend = true, color=:blue);

# Finally, we add all iterates and vertices to the plot.
colors = ["gold", "purple", "darkorange2", "firebrick3"]
iterates = unique!(iterates)
for i in 1:3
    scatter!([iterates[i][1]], [iterates[i][2]], label = string("x_", i - 1), markersize = 6, color = colors[i])
end
scatter!([last(iterates)[1]], [last(iterates)[2]], label = string("x_", length(iterates) - 1), markersize = 6, color = last(colors))

# plot chosen vertices
scatter!([vertices[1][1]], [vertices[1][2]], m = :diamond, markersize = 6, color = colors[1], label = "v_1")
scatter!([vertices[2][1]], [vertices[2][2]], m = :diamond, markersize = 6, color = colors[2], label = "v_2", legend = :outerleft, colorbar = true)


#=

  Euler #25 in Julia.

  Problem 25:
  """
  The Fibonacci sequence is defined by the recurrence relation:

     Fn = Fn1 + Fn2, where F1 = 1 and F2 = 1.

  Hence the first 12 terms will be:

     F1 = 1
     F2 = 1
     F3 = 2
     F4 = 3
     F5 = 5
     F6 = 8
     F7 = 13
     F8 = 21
     F9 = 34
     F10 = 55
     F11 = 89
     F12 = 144

  The 12th term, F12, is the first term to contain three digits.

  What is the first term in the Fibonacci sequence to contain 1000 digits?")
  """

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

=#

using Memoization
include("Euler.jl")

# Is not faster
@memoize function fib_len_memo(n)
    return fib(BigInt(n))|>string|>length
end

function fib_len(n)
    return fibmemBigInt(n)|>string|>length
end

# Brute force: 0.09872957s
function euler25a()
    i = 0
    len = 0
    while len < 1000
        i += 1
        len = fibmemBigInt(i)|>string|>length
    end

    return i;
end

#
# Using some heuristics to find the upper limit
# (from my Picat solution).
# 0.08765257s
function euler25b()
    target = 1000;
    foundUpper = 0;
    i = 1;
    fibLen = 0;
    step = 43;
    # Get the upper limit
    while fibLen < target && foundUpper == 0
        fibLen = fib_len(step*i);
        if fibLen > target
            foundUpper = i
            break
        end
        i += 1
    end

    # Now check all numbers from Step*(FoundUpper-1) .. Step*FoundUpper
    # The target must be in that interval.
    f = step*(foundUpper-1)
    fibLen = fib_len(f);
    while fibLen < target && f <= step*foundUpper
        f += 1
        fibLen = fib_len(f)
    end

    return f
end

# run_euler(euler25a)
run_euler(euler25b)


struct FieldMatrix{T,N,Lt<:AbstractArray{T},Dt<:AbstractArray{T,N}} <: AbstractMatrix{T}
    L :: Lt # Lower diagonal coefficients
    D :: Dt # Diagonal coefficients
    R :: CartesianIndices{N,NTuple{N,UnitRange{Int}}}
end
Base.size(p::FieldMatrix) = (s = length(p.R); (s,s))
Base.IndexStyle(::Type{<:FieldMatrix}) = IndexCartesian()
function Base.getindex(p::FieldMatrix{T}, i::Int, j::Int) where T
    if i == j
        return p.D[p.R[i]]
    else
        N = length(size(p.R))
        for d ∈ 1:N
            p.R[i]==p.R[j]+δ(d,N) && return p.L[p.R[i],d]
            p.R[i]==p.R[j]-δ(d,N) && return p.L[p.R[j],d]
        end
        return zero(T)
    end
end

@inline δ(i,N::Int) = CartesianIndex(ntuple(j -> j==i ? 1 : 0, N))
@fastmath @inline multL(I::CartesianIndex{N},L,x) where {N} =
    sum(@inbounds(x[I-δ(i,N)]*L[I,i]) for i ∈ 1:N)
@fastmath @inline multU(I::CartesianIndex{N},L,x) where {N} =
    sum(@inbounds(x[I+δ(i,N)]*L[I+δ(i,N),i]) for i ∈ 1:N)
@fastmath @inline mult(I::CartesianIndex{N},L,D,x) where {N} =
    @inbounds(x[I]*D[I])+multL(I,L,x)+multU(I,L,x)

import LinearAlgebra: mul!,dot,diag
mul!(b::FieldVector,p::FieldMatrix,x::FieldVector) = (@loop b[I]=mult(I,p.L,p.D,x); b)
@fastmath function dot(b::FieldVector,p::FieldMatrix,x::FieldVector)
    s = zero(eltype(x))
    @inbounds @simd for I ∈ b.R
        s+= b[I]*mult(I,p.L,p.D,x)
    end
    s
end
diag(p::FieldMatrix) = FieldVector(p.D)

import Base: *
*(p::FieldMatrix,x::FieldVector) = mul!(zero(x),p,x)

"""
    Poisson(L)

Construct a symmetric `FieldMatrix` from lower diagonal coefficients `L` with zero-sum rows.
This is a requirement for conservative Poisson equations: ∫ ∇⋅β∇ϕ dv = ∮ β ∂ϕ/∂n da.
"""
function Poisson(L::AbstractArray{T}) where T
    D = zeros(T,Base.front(size(L)))
    R = inside(D)
    for I ∈ R; D[I] = calcdiag(I,L); end
    FieldMatrix(L,D,R)
end
@fastmath @inline calcdiag(I::CartesianIndex{N},L) where {N} =
    -sum(@inbounds(L[I,i]+L[I+δ(i,N),i]) for i ∈ 1:N)


# 1-D interpolation from irregularly sampled data.

import PyPlot
import Random
import FIRRemez

fig_num = 1
PyPlot.close("all")

Random.seed!(25)

# generate sampling locations.
N_samples = 25
X = convert( Vector{BigFloat}, randn(N_samples).*π )

# generate samples.
Y = sin.(X)

# query.
N_queries = 500
xq_BF = convert(Vector{BigFloat}, collect( LinRange(minimum(X), maximum(X), N_queries) ) )
yq_BF = collect( FIRRemez.Barycentric2nditp(xq_BF[i],X,Y) for i = 1:length(xq_BF) )

xq = convert(Vector{Float64}, xq_BF)
yq = convert(Vector{Float64}, yq_BF)

X_display = convert(Vector{Float64}, X)
Y_display = convert(Vector{Float64}, Y)

PyPlot.figure(fig_num)
fig_num += 1
PyPlot.plot(xq, yq, label = "Barycentric")
PyPlot.plot(X_display, Y_display, label = "samples", "x")
PyPlot.title("Barycentric interpolation")
PyPlot.legend()


# Is z within an isapprox of any x[i] ?
# Assumes that x is sorted from lowest to highest
function nearby(z::T, x::Vector{T}) where {T}
    top = length(x)
    bot = 1
    mid = div(top, 2)
    while bot < mid < top
        if z > x[mid]
            bot = mid
        elseif z < x[mid]
            top = mid
        else
            return mid
        end
        mid = bot + div(top-bot, 2)
    end
    isapprox(x[bot], z) && return bot
    isapprox(x[top], z) && return top
    return -1   # should error
end

# For general location points with precalculated weights
"""
    bary(z, f, x, w)

evaluate f(z)

# Arguments
- `z::Float64`:         the point at which to evaluate f
- `f::Vector{Float64}`: vector of function values at x
- `x:Vector{Float64}`:  vector of eval locations of f (sorted)
- `w:Vector{Float64}`:  weights for locations x
"""
function bary(z::T, f::Vector{T}, x::Vector{T}, w::Vector{T}) where {T}
    # assert(length(f) == length(x) == length(w))
    num = zero(T)
    den = zero(T)
    @inbounds for j = 1:length(f)
        t = w[j] / (z - x[j])
        num += t * f[j]
        den += t
    end
    fz = num / den
    fz = isnan(fz) ? f[nearby(z, x)] : fz
end

function bary(z, f)
    n = length(f)
    x = chebpts(n)
    num = den = zero(z)
    t = 1 / 2(z - x[1])
    num += t * f[1]
    den += t
    sgn = -1
    @inbounds for j = 2:n-1
        t = sgn / (z - x[j])
        num += t*f[j]
        den += t
        sgn = -sgn
    end
    t = sgn / 2(z - x[n])
    num += t * f[n]
    den += t
    fz = num / den
    fz = isnan(fz) ? f[nearby(z, x)] : fz
end