#  Copyright 2017, Iain Dunning, Joey Huchette, Miles Lubin, and contributors
#  This Source Code Form is subject to the terms of the Mozilla Public
#  License, v. 2.0. If a copy of the MPL was not distributed with this
#  file, You can obtain one at http://mozilla.org/MPL/2.0/.
#############################################################################
# JuMP
# An algebraic modeling language for Julia
# See http://github.com/JuliaOpt/JuMP.jl
#############################################################################

using JuMP, GLPK, Test

"""
    example_sudoku(filepath)
A sudoku solver that uses a MIP to find solutions.
We have binary variables x[i, j, k] which, if = 1, say that cell (i, j) contains
the number k. The constraints are:
 1 - Each cell has one value only
 2 - Each row contains each number exactly once
 3 - Each column contains each number exactly once
 4 - Each 3x3 subgrid contains each number exactly once
We will take the initial grid as a CSV file at `filepath`, where 0s are blanks.
"""
function example_sudoku(filepath)
    initial_grid = zeros(Int, 9, 9)
    open(filepath, "r") do fp
        for row in 1:9
            line = readline(fp)
            initial_grid[row, :] .= parse.(Int, split(line, ","))
        end
    end

    model = Model(GLPK.Optimizer)

    @variable(model, x[1:9, 1:9, 1:9], Bin)

    @constraints(model, begin
        # Constraint 1 - Only one value appears in each cell
        cell[i in 1:9, j in 1:9], sum(x[i, j, :]) == 1
        # Constraint 2 - Each value appears in each row once only
        row[i in 1:9, k in 1:9], sum(x[i, :, k]) == 1
        # Constraint 3 - Each value appears in each column once only
        col[j in 1:9, k in 1:9], sum(x[:, j, k]) == 1
        # Constraint 4 - Each value appears in each 3x3 subgrid once only
        subgrid[i=1:3:7, j=1:3:7, val=1:9], sum(x[i:i + 2, j:j + 2, val]) == 1
    end)

    # Initial solution
    for row in 1:9, col in 1:9
        if initial_grid[row, col] != 0
            @constraint(model, x[row, col, initial_grid[row, col]] == 1)
        end
    end

    # Solve it
    JuMP.optimize!(model)

    term_status = JuMP.termination_status(model)
    primal_status = JuMP.primal_status(model)
    is_optimal = term_status == MOI.OPTIMAL

    # Check solution
    if is_optimal
        mip_solution = JuMP.value.(x)
        sol = zeros(Int, 9, 9)
        for row in 1:9, col in 1:9, val in 1:9
            if mip_solution[row, col, val] >= 0.9
                sol[row, col] = val
            end
        end
        return sol
    else
        error("The solver did not find an optimal solution.")
    end
end

function print_sudoku_solution(solution)
    println("Solution:")
    println("[-----------------------]")
    for row in 1:9
        print("[ ")
        for col in 1:9
            print(solution[row, col], " ")
            if col % 3 == 0 && col < 9
                print("| ")
            end
        end
        println("]")
        if row % 3 == 0
            println("[-----------------------]")
        end
    end
end

if length(ARGS) == 1
    sol = example_sudoku(ARGS[1])
    print_sudoku_solution(sol)
else
    solution = example_sudoku(joinpath(@__DIR__, "sudoku.csv"))
    @test solution == [
        3 1 7 9 5 8 2 6 4;
        4 6 9 3 2 7 8 1 5;
        8 2 5 1 6 4 7 9 3;
        2 7 1 6 4 5 3 8 9;
        9 3 8 7 1 2 5 4 6;
        5 4 6 8 3 9 1 7 2;
        1 8 4 2 9 3 6 5 7;
        6 9 2 5 7 1 4 3 8;
        7 5 3 4 8 6 9 2 1]
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...)


# Test pair-wise correlation and covariance functions
x, y = rand(1000), rand(1000)

@test nancov(x,y) ≈ cov(x,y)
@test nancov(x,y, corrected=false) ≈ cov(x,y, corrected=false)
@test nancor(x,y) ≈ cor(x,y)

# Test correlation and covariance functions as applied to matrices
X = rand(100,10)

@test nancov(X) ≈ cov(X)
@test nancov(X, dims=1) ≈ cov(X, dims=1)
@test nancov(X, dims=2) ≈ cov(X, dims=2)
@test nancov(X, corrected=false) ≈ cov(X, corrected=false)

@test nancor(X) ≈ cor(X)
@test nancor(X, dims=1) ≈ cor(X, dims=1)
@test nancor(X, dims=2) ≈ cor(X, dims=2)

# Test NaN-ignoring
xn = [1,2,3,NaN]
xnn = [1,2,3]
@test nancov(xn,xn) ≈ cov(xnn,xnn) ≈ 1
@test nancor(xn,xn) ≈ cor(xnn,xnn) ≈ 1


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)


function _partition(x::AbstractArray, p::Int, r::Int)::Int
    @inbounds y = x[r]
    i = p-1
    for j ∈ p:(r-1)
        @inbounds if x[j] ≤ y
            i += 1
            @inbounds tmp  = x[i]
            @inbounds x[i] = x[j]
            @inbounds x[j] = tmp
        end
    end
    @inbounds tmp    = x[i+1]
    @inbounds x[i+1] = x[r]
    @inbounds x[r]   = tmp
    return i+1
end

function quicksort!(x::AbstractArray, p::Int, r::Int)
    if p ≥ r
        return nothing
    end
    q::Int = _partition(x, p, r)
    quicksort!(x, p, q-1)
    quicksort!(x, q+1, r)
end


### A Pluto.jl notebook ###
# v0.14.3

using Markdown
using InteractiveUtils

# ╔═╡ 7dfe963f-3601-4515-a636-4930816a4e6c
using Pkg, DrWatson

# ╔═╡ 73d28c7f-8426-4c5a-b166-23dcac5c16ce
begin
	@quickactivate "NumericalMethodsforEngineers"
	using NumericalMethodsforEngineers
	using Test
end

# ╔═╡ ac233fe5-a00f-412f-ad4a-a1d39ee20517
md" ### f-01-01.jl"

# ╔═╡ d82838e5-1591-41d9-b5e0-0c49f48d7814
begin
	a = [10. 5. 6.; 5. 20. 4.; 6. 4. 30.]
	x0 = [1., 1., 1.]

	tol = 1.0e-5
	limit = 100
end;

# ╔═╡ 570d9517-14d8-4d40-8e29-039269a5e37d
(iters, eigval, eigvec) = nmex(a, x0; tol=tol, limit=limit, show=8)

# ╔═╡ 1614c806-1bdc-40bb-ae8c-e02d29d128fe
jeigval, jeigvec = eigen(a)

# ╔═╡ 90c91c6a-1b9b-4d34-858a-7535b1090fee
@test round(maximum(jeigval), digits=3) == round(eigval, digits=3)

# ╔═╡ 47256683-e6b6-4fae-916c-f80ba57a232f
md" ### End of f-01-01.jl"

# ╔═╡ Cell order:
# ╟─ac233fe5-a00f-412f-ad4a-a1d39ee20517
# ╠═7dfe963f-3601-4515-a636-4930816a4e6c
# ╠═73d28c7f-8426-4c5a-b166-23dcac5c16ce
# ╠═d82838e5-1591-41d9-b5e0-0c49f48d7814
# ╠═570d9517-14d8-4d40-8e29-039269a5e37d
# ╠═1614c806-1bdc-40bb-ae8c-e02d29d128fe
# ╠═90c91c6a-1b9b-4d34-858a-7535b1090fee
# ╟─47256683-e6b6-4fae-916c-f80ba57a232f


"""
Package: DataWrangler

normalize\\[!\\]([x]; method)

Normalize values in vector `x` either in-place or returning a copy of `x` with the imputed values

# Parameters

`x`: Vector of type Real with missing values to be imputed

`method`: There are four valid values: "z-score", "min-max", "softmax", "sigmoid"

# Description

The normalization is applied ignoring any missing values in the Array, and it takes place in all the available dimensions of the array. If normalization is required in just some specific dimensions the function `mapslices` can be used to select those dimensions.

# Example
```julia
x = [1.,2,3,4,5]
normalize!(x; method = "min-max")

println(x)
[0.0, 0.25, 0.5, 0.75, 1.0]
```
"""
function normalize!(x::AbstractArray{<:Union{Missing,T}};
                 method::String = "z-score") where T<:Real

    @assert !(T <: Integer) " Method `$(typeof(x))` cannot be normalized in-place, use `normalize(x)` instead"

    @assert method ∈ ["z-score","min-max","softmax","sigmoid"]

    smx = skipmissing(x)

    if method == "z-score"
        μ = mean(smx)
        σ = std(smx)
        for i in CartesianIndices(x)
            x[i] = (x[i] - μ) / σ
        end
    end

    if method == "min-max"
        m,M = extrema(smx)
        Mm = M-m
        for i in CartesianIndices(x)
            x[i] = (x[i] - m) / Mm
        end
    end

    if method == "softmax"
        Σe = sum(exp.(smx))
        for i in CartesianIndices(x)
            x[i] = exp(x[i]) / Σe
        end
    end

    if method == "sigmoid"
        for i in CartesianIndices(x)
            x[i] = 1/(1+exp(-x[i]))
        end
    end

end

"""
Package: DataWrangler

Check 'normalize! for further information.
"""
function normalize(x::AbstractArray{<:Union{Missing,T}};
                method::String = "z-score") where T<:Real

    P = any(ismissing,x) ? Union{Missing, Base.promote_op(/,T,T)} : Base.promote_op(/,T,T)
    cx = Array{P}(undef,size(x))
    cx[:] = x
    normalize!(cx; method)
    return cx

end


#=

  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


module VectorHeatEquationTests

using Gridap
using ForwardDiff
using LinearAlgebra
using Test
using GridapODEs.ODETools
using GridapODEs.TransientFETools
using Gridap.FESpaces: get_algebraic_operator

# using GridapODEs.ODETools: ThetaMethodLinear

import Gridap: ∇
import GridapODEs.TransientFETools: ∂t

θ = 1.0

u(x,t) = (1.0-x[1])*x[1]*(1.0-x[2])*x[2]*t
u(t::Real) = x -> u(x,t)
v(x) = t -> u(x,t)
f(t) = x -> ∂t(u)(t)(x)-Δ(u(t))(x)

domain = (0,1,0,1)
partition = (2,2)
model = CartesianDiscreteModel(domain,partition)

order = 2

reffe = ReferenceFE(lagrangian,Float64,order)
V0 = FESpace(
  model,
  reffe,
  conformity=:H1,
  dirichlet_tags="boundary"
)
U = TransientTrialFESpace(V0,u)

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

#
a(u,v) = ∫(∇(v)⊙∇(u))dΩ
b(v,t) = ∫(v⋅f(t))dΩ
m(ut,v) = ∫(ut⋅v)dΩ

X = TransientMultiFieldFESpace([U,U])
Y = MultiFieldFESpace([V0,V0])

_res(t,u,ut,v) = a(u,v) + m(ut,v) - b(v,t)

res(t,(u1,u2),(u1t,u2t),(v1,v2)) = _res(t,u1,u1t,v1) + _res(t,u2,u2t,v2)
jac(t,x,xt,(du1,du2),(v1,v2)) = a(du1,v1) + a(du2,v2)
jac_t(t,x,xt,(du1t,du2t),(v1,v2)) = m(du1t,v1) + m(du2t,v2)

op = TransientFEOperator(res,jac,jac_t,X,Y)

t0 = 0.0
tF = 1.0
dt = 0.1

U0 = U(0.0)
X0 = X(0.0)
uh0 = interpolate_everywhere(u(0.0),U0)
xh0 = interpolate_everywhere([uh0,uh0],X0)

ls = LUSolver()
# using Gridap.Algebra: NewtonRaphsonSolver
# nls = NLSolver(ls;show_trace=true,method=:newton) #linesearch=BackTracking())
odes = ThetaMethod(ls,dt,θ)
solver = TransientFESolver(odes)

sol_t = solve(solver,op,xh0,t0,tF)

l2(w) = w⋅w

tol = 1.0e-6
_t_n = t0

result = Base.iterate(sol_t)

for (xh_tn, tn) in sol_t
  global _t_n
  _t_n += dt
  uh_tn = xh_tn[1]
  e = u(tn) - uh_tn
  el2 = sqrt(sum( ∫(l2(e))dΩ ))
  @test el2 < tol
end

end #module


module Testing

import ..NodalDG

using Test

# Simple test of JacobiP
@test NodalDG.JacobiP(0.0, 0, 0, 0) == [1.0 /sqrt(2.0)]

# Test Gauss-Lobatto points

r, w = NodalDG.gausslobatto(3)
@test isapprox(r, [-1, 0, 1])
@test isapprox(w, [1/3, 4/3, 1/3])

r,w = NodalDG.gausslobatto(4)
@test isapprox(r, [-1, -sqrt(1/5), sqrt(1/5), 1])
@test isapprox(w, [1/6, 5/6, 5/6, 1/6])

# Tests based on p48 in textbook, Fig 3.1
# This tests JacobiP and LGL points

r, w = NodalDG.gausslobatto(7)

# Fig 3.1d
Vref = [0.707107 -1.22474 1.58114 -1.87083 2.12132 -2.34521 2.54951;
        0.707107 -1.01681 0.844182 -0.346644 -0.278373 0.807539 -1.05741;
        0.707107 -0.57422 -0.269222 0.833675 -0.504704 -0.365166 0.846707;
        0.707107 0.0 -0.790569 -0.0 0.795495 0.0 -0.796722;
        0.707107 0.57422 -0.269222 -0.833675 -0.504704 0.365166 0.846707;
        0.707107 1.01681 0.844182 0.346644 -0.278373 -0.807539 -1.05741;
        0.707107 1.22474 1.58114 1.87083 2.12132 2.34521 2.54951]

@test isapprox( NodalDG.Vandermonde1D(6,r), Vref, atol=1e-5 )

# Tests based on p54, Fig 3.4

# order N = 1
r, w = NodalDG.gausslobatto(2)
V = NodalDG.Vandermonde1D(1, r)

Dref1 = [-0.5 0.5;
         -0.5 0.5]
@test isapprox( NodalDG.Dmatrix1D(1, r, V), Dref1 )

# order N = 2
r, w = NodalDG.gausslobatto(3)
V = NodalDG.Vandermonde1D(2, r)

Dref2 = [-1.5 2.0 -0.5;
         -0.5 0.0 0.5;
         0.5 -2.0 1.5]

@test isapprox( NodalDG.Dmatrix1D(2, r, V), Dref2 )

# order N = 4
r, w = NodalDG.gausslobatto(5)
V = NodalDG.Vandermonde1D(4, r)

Dref4 = [-5.0 6.7565 -2.66667 1.41016 -0.5;
         -1.24099 2.51054e-16 1.74574 -0.763763 0.25901;
         0.375 -1.33658 -0.0 1.33658 -0.375;
         -0.25901 0.763763 -1.74574 -6.56798e-16 1.24099;
         0.5 -1.41016 2.66667 -6.7565 5.0]

@test maximum(abs, NodalDG.Dmatrix1D(4, r, V) - Dref4) < 1e-5

end


#The m-by-n Hilbert matrix has matrix elements
# H_{ij} = 1/(i+j-1)
export Hilbert, InverseHilbert

struct Hilbert{T} <: AbstractMatrix{T}
    m :: Int
    n :: Int
end
Hilbert(T::Type, m::Integer, n::Integer) = Hilbert{T}(m, n)
Hilbert(T::Type{<:Integer}, m::Integer, n::Integer) = Hilbert(Rational{T}, m, n)
Hilbert(m::Integer, n::Integer) = Hilbert(Int, m, n)
Hilbert(n::Integer) = Hilbert(n, n)

# Define its size
size(H::Hilbert, dim::Integer) = dim==1 ? H.m : dim==2 ? H.n : 1
size(H::Hilbert)= size(H,1), size(H,2)

# Index into a Hilbert matrix
function getindex(H::Hilbert{T}, i::Integer, j::Integer) where {T}
    return one(T)/(i+j-1)
end

# Dense version
Matrix(H::Hilbert) = [H[i,j] for i=1:size(H,1), j=1:size(H,2)]

# Some properties
ishermitian(H::Hilbert) = H.m==H.n
isposdef(H::Hilbert) = H.m==H.n
det(H::Hilbert) = inv(det(inv(H)))

# Inverse of a Hilbert matrix
struct InverseHilbert{T} <: AbstractMatrix{T}
    n :: Int
end
InverseHilbert(T::Type, n::Integer) = InverseHilbert{T}(n)
InverseHilbert(n::Integer) = InverseHilbert(Int, n)

# Define its size
size(A::InverseHilbert, dim::Integer) = dim==1 || dim==2 ? A.n : 1
size(A::InverseHilbert) = size(A,1), size(A,2)

# Index into a inverse Hilbert matrix
function getindex(A::InverseHilbert{T}, i::Integer, j::Integer) where {T}
    return T((-1)^(i+j)*(i+j-1)*binomial(A.n+i-1,A.n-j)*
              binomial(A.n+j-1,A.n-i)*binomial(i+j-2,i-1)^2)
end
# Use bigger numbers in case of bigger types
function getindex(A::InverseHilbert{T}, i::Integer, j::Integer) where {T<:Union{BigInt,BigFloat}}
    N = big(A.n)
    return T((-1)^(i+j)*(i+j-1)*binomial(N+i-1,N-j)*
              binomial(N+j-1,N-i)*binomial(big(i+j-2),i-1)^2)
end

# Explicit formula for the determinant
det(A::InverseHilbert{T}) where {T} = prod(T,(2k+1)*binomial(2k,k)^2 for k=1:A.n-1)

# Dense version
Matrix(A::InverseHilbert) = [A[i,j] for i=1:size(A,1), j=1:size(A,2)]

# Define Base.inv
function inv(H::Hilbert{T}) where {T}
    H.m == H.n || throw(ArgumentError("Works only for square Hilbert matrices."))
    return InverseHilbert(T,H.n)
end
function inv(A::InverseHilbert{T}) where {T}
    HT = promote_type(T,typeof(Rational(one(T))))
    return Hilbert(HT,A.n)
end


function getIja(Ia::Matrix{T}, j::Int, N::Int) where T

    Id = [one(T) zero(T); zero(T) one(T)]

    Ias = collect( Id for i = 1:N )
    Ias[j] = Ia
    return Kronecker.kron(Ias...)
end

function productIaIb(Ixs, Iys, Izs, a::Int, b::Int)
    return Ixs[a]*Ixs[b] + Iys[a]*Iys[b] + Izs[a]*Izs[b]
end

"""
For operator A.

Example:
N = 3
Ixs = collect( getIja(Ix, i, N) for i = 1:N )
Iys = collect( getIja(Iy, i, N) for i = 1:N )
Izs = collect( getIja(Iz, i, N) for i = 1:N )

# sanity check.
out1 = productIaIb(Ixs, Iys, Izs, 1, 3)
out2 = productIjIk(Ixs, Iys, Izs, [1;], [3;])
norm(out1-out2)
"""
function productIjIk(Ixs, Iys, Izs, inds_j::Vector{Int}, inds_k::Vector{Int})

    #

    out = sum(Ixs[j] for j in inds_j) * sum(Ixs[k] for k in inds_k) +
        sum(Iys[j] for j in inds_j) * sum(Iys[k] for k in inds_k) +
        sum(Izs[j] for j in inds_j) * sum(Izs[k] for k in inds_k)

    return out
end

"""
for operator B.

Example:
N = 3
Ixs = collect( getIja(Ix, i, N) for i = 1:N )
Iys = collect( getIja(Iy, i, N) for i = 1:N )
Izs = collect( getIja(Iz, i, N) for i = 1:N )

# # sanity check.
# out1 = productIaIb(Ixs, Iys, Izs, 1, 3)
# out2 = productIjIk(Ixs, Iys, Izs, [1;], [3;])

out1 = pairwiseproductIa(Ixs, Iys, Izs, [1; 2; 3])
out2 = productIaIb(Ixs, Iys, Izs, 1, 3) + productIaIb(Ixs, Iys, Izs, 2, 3) + productIaIb(Ixs, Iys, Izs, 1, 2)
norm(out1-out2)
"""
function pairwiseproductIa(Ixs, Iys, Izs, inds::Vector{Int})

    out = zeros(Float64, size(Ixs[1]))
    for i = 1:length(inds)

        ind1 = inds[i]
        for j = i+1:length(inds)
            ind2 = inds[j]

            out += productIaIb(Ixs, Iys, Izs, ind1, ind2)
        end
    end

    return out
end


function LC(p, var)
    d = maxdegree(p, var)
    d == 0 && return zero(p)
    return polynomial([subs(t, var => 1) for t in terms(p) if degree(t, var) == d])
end

function pseudorem(f, g, y)
    r = f
    lc = LC(g, y)
    m = maxdegree(g, y)
    d = maxdegree(r, y)

    while d ≥ m && !iszero(r)
        r = lc * r - LC(r, y) * g * y ^ (d - m)
        d = maxdegree(r, y)
    end

    return r
end

function triangulize(H, vars)
    n = length(H)

    T = similar(H)
    @inbounds for i = n:-1:1
        degrees = Int[] # degrees of polynomials in C
        indeces = Int[] # indeces in H of polynomials in C

        # initialize
        @inbounds for (idx, p) in enumerate(H)
            deg = maxdegree(p, vars[i])
            if deg > 0
                push!(degrees, deg)
                push!(indeces, idx)
            end
        end

        # If there are several polynomials in C but none has degree 1 in x, transform the
        # set until there is a polynomial of degree 1 or only one polynomial is left
        C = H[indeces]
        while minimum(degrees) > 1 && length(C) > 1
            _, idx_max = findmax(degrees)
            _, idx_min = findmin(degrees)

            # special case when all  have same degree
            if idx_max == idx_min
                idx_min = 1
                idx_max = 2
            end

            r = pseudorem(C[idx_max], C[idx_min], vars[i])
            d = maxdegree(r, vars[i])
            if iszero(d)
                H[indeces[idx_max]] = r
                deleteat!(C, idx_max)
                deleteat!(degrees, idx_max)
                deleteat!(indeces, idx_max)
            else
                C[idx_max] = r
                degrees[idx_max] = d
            end
        end

        # if only one polynomial in C, add to triangular set and go to the next variable
        if length(indeces) == 1
            T[i] = H[indeces[1]]
            deleteat!(H, indeces[1])
            continue
        end

        # if multiple polynomials but at least one has degree 1, add this to triangular set
        # and replace the others in H with the pseudoremainder
        idx = findlast(==(1), degrees)
        if !isnothing(idx)
            pivot = H[indeces[idx]]
            T[i] = pivot
            #pushfirst!(T, pivot)

            @inbounds for j in indeces
                if j != indeces[idx]
                    H[j] = pseudorem(H[j], pivot, vars[i])
                end
            end
            deleteat!(H, indeces[idx])
            continue
        end

    end
    return T
end


# # Even reduction

#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/generated/Symmetry/even_reduction.ipynb)
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/generated/Symmetry/even_reduction.ipynb)

using Test #src
using DynamicPolynomials
@polyvar x

# We would like to find the minimum value of the following polynomial:

poly = x^4 - 2x^2

using SumOfSquares

# We define the custom action as follows:

struct OnSign <: Symmetry.OnMonomials end
using PermutationGroups
import SymbolicWedderburn
SymbolicWedderburn.coeff_type(::OnSign) = Float64
function SymbolicWedderburn.action(::OnSign, p::Permutation, mono::AbstractMonomial)
    if isone(p) || iseven(DynamicPolynomials.degree(mono))
        return 1 * mono
    else
        @assert p.perm == perm"(1,2)"
        return -1 * mono
    end
end
G = PermGroup([perm"(1,2)"])

# We can exploit the symmetry as follows:

import CSDP
solver = CSDP.Optimizer
model = Model(solver)
@variable(model, t)
@objective(model, Max, t)
pattern = Symmetry.Pattern(G, OnSign())
con_ref = @constraint(model, poly - t in SOSCone(), symmetry = pattern)
optimize!(model)
@test value(t) ≈ -1 #src
value(t)

# We indeed find `-1`, let's verify that symmetry was exploited:

@test length(gram_matrix(con_ref).sub_gram_matrices) == 2 #src
@test gram_matrix(con_ref).sub_gram_matrices[1].basis.polynomials == [x^2, 1] #src
@test gram_matrix(con_ref).sub_gram_matrices[2].basis.polynomials == [x] #src
gram_matrix(con_ref)


using Dates

effective_date = Date(2020, 01, 15)

days = [369, 462, 553, 735, 1100, 1465]
df = [0.97678403, 0.97153083, 0.96655364, 0.95695995, 0.93829085, 0.91954718]
test_date = effective_date + Day(900)

const dc_act_360 = QuantLib.Time.Actual360()

# test log-linear interpolation
test_curve_loglinear = QuantLib.InterpolatedDiscountCurve(effective_date .+ Day.(days), df, dc_act_360, QuantLib.Math.LogLinear())
disc_factor = QuantLib.discount(test_curve_loglinear, test_date)
x1 = (test_date - test_curve_loglinear.dates[4]).value
dx = (test_curve_loglinear.dates[5]-test_curve_loglinear.dates[4]).value
dy = log(test_curve_loglinear.data[5]) - log(test_curve_loglinear.data[4])
df_recalc = exp(log(test_curve_loglinear.data[4]) + x1*dy/dx)
@test disc_factor == df_recalc

test_curve_linear = QuantLib.InterpolatedDiscountCurve(effective_date .+ Day.(days), df, dc_act_360, QuantLib.Math.LinearInterpolation())
disc_factor = QuantLib.discount(test_curve_linear, test_date)
x1 = (test_date - test_curve_linear.dates[4]).value
dx = (test_curve_linear.dates[5]-test_curve_linear.dates[4]).value
dy = (test_curve_linear.data[5]) - (test_curve_linear.data[4])
df_recalc = ((test_curve_linear.data[4]) + x1*dy/dx)
@test disc_factor == df_recalc

# test joint calendars
const calendar_nyc = QuantLib.Time.USNYSECalendar()
const calendar_target = QuantLib.Time.TargetCalendar()
const calendar_lon = QuantLib.Time.UKLSECalendar()

calendar_nyc_target = QuantLib.Time.JointCalendar(calendar_nyc, calendar_target) # standard constructor
@test calendar_nyc_target isa QuantLib.Time.BusinessCalendar
calendar_nyc_target_lon = QuantLib.Time.JointCalendar(calendar_nyc, calendar_target, calendar_lon) # new constructor
@test calendar_nyc_target_lon isa QuantLib.Time.BusinessCalendar
# at least 2 calendars must be specified for JointCalendar
@test_throws ErrorException QuantLib.Time.JointCalendar(calendar_nyc)


function binary(n::Int)
	ans=Int[]
	while n>=1

		push!(ans,n%2)
		n=floor(Int,n/2)

	end

	return reverse(ans)
end

function is_decimal_binary_palindrome(n::Int)
	s=string(n)
	if s!=reverse(s)
		return false
	end

	bin=binary(n)

	return bin==reverse(bin)
end

function compute()

	ans= sum(i for i=1:1000000 if is_decimal_binary_palindrome(i))

	return ans
end

println(compute())


################################################
## Specialized routines for tridiagonal matrices
################################################

export numnegevals

"""
Computes the number of negative eigenvalues of T - σI, a.k.a. spectrum slicing

Inputs:
    T: A SymTridiagonal{<:Real} matrix
    σ: The shift parameter

Outputs:
    ν: The number of negative eigenvalues

Reference:
    B. N. Parlett, "The symmetric eigenvalue problem", Section 3.3.1, p. 52.
"""
function numnegevals(T::SymTridiagonal{S}, σ::S = zero(S)) where S
    α = T.dv
    β = T.ev
    n = length(α)
    ϵ = eps(S)
    δ = α[1] - σ
    ν = δ < 0 ? 1 : 0
    for k = 1:n - 1
        if δ == 0
        info("zero in iteration $k")
        δ = ϵ * (β[k]+ϵ) #Parlett prefers adjusting σ and starting again
    end
        δ = (α[k + 1] - σ) - β[k]*(β[k]/δ)
    ν += (δ < 0)
    end
    ν
end


# ---
# title: 461. Hamming Distance
# id: problem461
# author: Indigo
# date: 2021-02-04
# difficulty: Easy
# categories: Bit Manipulation
# link: <https://leetcode.com/problems/hamming-distance/description/>
# hidden: true
# ---
#
# The [Hamming distance](https://en.wikipedia.org/wiki/Hamming_distance) between
# two integers is the number of positions at which the corresponding bits are
# different.
#
# Given two integers `x` and `y`, calculate the Hamming distance.
#
# **Note:**
# 0 ≤ `x`, `y` < 231.
#
# **Example:**
#
#
#
#     Input: x = 1, y = 4
#
#     Output: 2
#
#     Explanation:
#     1   (0 0 0 1)
#     4   (0 1 0 0)
#            ↑   ↑
#
#     The above arrows point to positions where the corresponding bits are different.
#
#
#
## @lc code=start
using LeetCode

function hamming_distance(x::Int, y::Int)
    res = 0
    while x != 0 || y != 0
        res += (x & 1) ⊻ (y & 1)
        x >>= 1
        y >>= 1
    end
    res
end
## @lc code=end


using JuMP
using Cbc

include("../OptMaps.jl")
using .OptMaps

function solve_model(solver, model_log, solver_options=Dict())
    model = Model(solver)

    for (k, v) in solver_options
        set_optimizer_attribute(model, k, v)
    end

    n_variables = 7
    n_constraints = 4

    @variable(model, 0 <= x[1:6] <= 10, Int)
    @variable(model, 0 <= y <= 1, Int)

    ocoeffs = [60 70 40 70 20 90]
    ccoeffs = [30 20 30 90 30 70]

    @constraint(model, sum((x[i]*ccoeffs[i] for i in 1:6)) <= 2000)
    @constraint(model, x[3] <= 10*x[4])
    @constraint(model, x[1] + x[2] >= 4*y)
    @constraint(model, x[5] + x[6] >= 4*(1 - y))

    constraints_f = []
    push!(constraints_f, (v) -> sum((v[i]*ccoeffs[i] for i in 1:6)) <= 2000)
    push!(constraints_f, (v) -> v[3] <= 10*v[4])
    push!(constraints_f, (v) -> v[1] + v[2] >= 4*v[7])
    push!(constraints_f, (v) -> v[5] + v[6] >= 4*(1 - v[7]))

    @objective(model, Max, sum(x[i]*ocoeffs[i] for i in 1:6))
    obj_f(v) = sum(v[i]*ocoeffs[i] for i in 1:6)

    starting_point = zeros(Float64, n_variables)
    solution_point = zeros(Float64, n_variables)

    ###############################################################
    #save model textual description
    open(model_log, "w") do f
        print(f, model)

        ###############################################################
        #solve model
        elapsed_time = @elapsed optimize!(model)

        println(f, "-----------------------------------------------")
        println(f, "Elapsed time: $elapsed_time")
        println(f, "-----------------------------------------------")
        println(f, "Status:    $(termination_status(model))")
        println(f, "Objective: $(objective_value(model))")
        println(f, "-----------------------------------------------")
        println(f, "Variables:")

        for i in 1:(n_variables-1)
            println(f, "$(x[i]): $(value(x[i]))")
            solution_point[i] = value(x[i])
        end

        println(f, "$(y): $(value(y))")
        solution_point[n_variables] = value(y)
    end

    variable_names = vcat(["x[$(string(i))]" for i in 1:(n_variables-1)], ["y"])
    constraint_names = ["c$(string(c))" for c in 1:n_constraints]

    obj_val = objective_value(model)

    return obj_f, constraints_f, variable_names, constraint_names, starting_point, solution_point, obj_val
end

#The datatype used in the tuples is used to determine if the variable is "real" or integer
variable_domains = [(0, 12), (0, 12), (0, 12), (0, 12), (0, 12), (0, 12), (0, 1)]

optimizers = [Cbc.Optimizer]
opt_names = ["cbc"]

problems = Dict()

for (opt, opt_name) in zip(optimizers, opt_names)
    objective_f, constraints_f, variable_names, constraint_names, starting_point, solution_point, obj_val = solve_model(opt, "knapsack_$(opt_name).log")
    problems[opt_name] = (objective_f, constraints_f, variable_names, constraint_names, starting_point, solution_point, obj_val)
end

#All functions are equal, get first to initialize OptMap
objective_f, constraints_f, variable_names, constraint_names, _, _, obj_val = problems[opt_names[1]]
O = OptMap(variable_names, variable_domains, objective_f; constraint_names=constraint_names, constraints_f=constraints_f, verbose=true)
create_map!(O)

for opt_name in opt_names
    _, _, _, _, starting_point, solution_point, _ = problems[opt_name]

    path_to_solution = zeros(Float64, length(variable_domains), 2)
    path_to_solution[:,1] = starting_point
    path_to_solution[:,2] = solution_point

    draw_solution!(O, path_to_solution)
    save_all(O, "knapsack_$(opt_name)", ".")
end


function eval_metric(::Val{:mse}, p::AbstractMatrix{T}, y::AbstractVector{T}, α=0.0) where {T <: AbstractFloat}
    eval = zero(T)
    @inbounds for i in eachindex(y)
        eval += (p[1,i] - y[i])^2
    end
    eval /= length(y)
    return eval
end

function eval_metric(::Val{:rmse}, p::AbstractMatrix{T}, y::AbstractVector{T}, α=0.0) where {T <: AbstractFloat}
    eval = zero(T)
    @inbounds for i in eachindex(y)
        eval += (p[1,i] - y[i])^2
    end
    eval = sqrt(eval / length(y))
    return eval
end

function eval_metric(::Val{:mae}, p::AbstractMatrix{T}, y::AbstractVector{T}, α=0.0) where {T <: AbstractFloat}
    eval = zero(T)
    @inbounds for i in eachindex(y)
        eval += abs(p[1,i] - y[i])
    end
    eval /= length(y)
    return eval
end

function eval_metric(::Val{:logloss}, p::AbstractMatrix{T}, y::AbstractVector{T}, α=0.0) where {T <: AbstractFloat}
    eval = zero(T)
    @inbounds for i in eachindex(y)
        pred = sigmoid(p[1,i])
        eval -= y[i] * log(pred) + (1 - y[i]) * log(1 - pred)
    end
    eval /= length(y)
    return eval
end

function eval_metric(::Val{:mlogloss}, p::AbstractMatrix{T}, y::AbstractVector{S}, α=0.0) where {T <: AbstractFloat,S <: Integer}
    eval = zero(T)
    p_prob = exp.(p) ./ sum(exp.(p), dims=1)
    @inbounds for i in eachindex(y)
        eval -= log(p_prob[y[i], i])
    end
    eval /= length(y)
    return eval
end

function eval_metric(::Val{:poisson}, p::AbstractMatrix{T}, y::AbstractVector{T}, α=0.0) where {T <: AbstractFloat}
    eval = zero(T)
    @inbounds for i in eachindex(y)
        eval += exp(p[1,i]) * (1 - y[i]) + log(factorial(y[i]))
    end
    eval /= length(y)
    return eval
end

function eval_metric(::Val{:gaussian}, p::AbstractMatrix{T}, y::AbstractVector{T}, α=0.0) where {T <: AbstractFloat}
    eval = zero(T)
    @inbounds for i in eachindex(y)
        eval += p[2,i] + (y[i] - p[1,i])^2 / (2 * exp(2 * p[2,i]))
    end
    eval /= length(y)
    return eval
end

function eval_metric(::Val{:quantile}, p::AbstractMatrix{T}, y::AbstractVector{T}, α=0.0) where {T <: AbstractFloat}
    eval = zero(T)
    for i in eachindex(y)
        eval += α * max(y[i] - p[1,i], zero(T)) + (1 - α) * max(p[1,i] - y[i], zero(T))
    end
    eval /= length(y)
    return eval
end

function gini_raw(y::T, p::S) where {T,S}
    if length(y) < 2
        return 0.0
    end
    random = (1:length(p)) ./ length(p)
    l_sort = y[sortperm(p)]
    l_cum_w = cumsum(l_sort) ./ sum(y)
    gini = sum(l_cum_w .- random)
    return gini
end

function gini_norm(y::T, p::S) where {T,S}
    if length(y) < 2
        return 0.0
    end
    return gini_raw(y, p) / gini_raw(y, y)
end

function eval_metric(::Val{:gini}, p::AbstractMatrix{T}, y::AbstractVector{T}, α=0.0) where {T <: AbstractFloat}
    return -gini_norm(y, view(p, 1, :))
end


using Printf
K = 0; N = 10000000;
@time begin
for i in 1:N
  if (rand()^2 + rand()^2) < 1.0
    global K += 1
  end
end
end
p = (4*K)/N
@printf "Estimation of PI is %.6f\n" p


# # Symmetry reduction

#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/generated/Symmetry/symmetry_reduction.ipynb)
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/generated/Symmetry/symmetry_reduction.ipynb)
# **Adapted from**: [SymbolicWedderburn example](https://github.com/kalmarek/SymbolicWedderburn.jl/blob/tw/ex_sos/examples/ex_C4.jl)

import MutableArithmetics
const MA = MutableArithmetics
using MultivariatePolynomials
const MP = MultivariatePolynomials
using MultivariateBases
const MB = MultivariateBases

using Test #src
using DynamicPolynomials
@polyvar x[1:4]

# We would like to find the minimum value of the polynomial

poly = sum(x) + sum(x.^2)

# As we can decouple the problem for each `x[i]` for which `x[i] + x[i]^2` has
# minimum value 0.25, we would expect to get `-1` as answer.
# Can this decoupling be exploited by SumOfSquares as well ?
# For this, we need to use a certificate that can exploit the permutation symmetry of the polynomial.

using SumOfSquares

# We define the symmetry group as a permutation group in the variables.
# In order to do that, we define the action of a permutation on a monomial
# as the monomial obtained after permuting the variables.

using PermutationGroups
G = PermGroup([perm"(1,2,3,4)"])

# We can use this certificate as follows:

import CSDP
solver = CSDP.Optimizer
model = Model(solver)
@variable(model, t)
@objective(model, Max, t)
pattern = Symmetry.Pattern(G, Symmetry.VariablePermutation())
con_ref = @constraint(model, poly - t in SOSCone(), symmetry = pattern)
optimize!(model)
@test value(t) ≈ -1 rtol=1e-6 #src
value(t)

# We indeed find `-1`, let's verify that symmetry was exploited:

g = gram_matrix(con_ref).sub_gram_matrices    #src
@test length(g) == 4                          #src
@test length(g[1].basis.polynomials) == 2     #src
@test g[1].basis.polynomials[1] == 1.0        #src
@test g[1].basis.polynomials[2] ≈ -0.5 * sum(x) #src
@test size(g[1].Q) == (2, 2)                  #src
@test g[1].Q[1, 1] ≈ 1.0 atol=1e-6            #src
@test g[1].Q[1, 2] ≈ -1.0 atol=1e-6            #src
@test g[1].Q[2, 2] ≈ 1.0 atol=1e-6            #src
@test length(g[2].basis.polynomials) == 1     #src
@test g[2].basis.polynomials[1] ≈ (x[3] - x[1]) / √2 #src
@test size(g[2].Q) == (1, 1)                  #src
@test g[2].Q[1, 1] ≈ 1.0 atol=1e-6            #src
@test length(g[3].basis.polynomials) == 1     #src
@test g[3].basis.polynomials[1] ≈ (x[4] - x[2]) / √2 #src
@test size(g[3].Q) == (1, 1)                  #src
@test g[3].Q[1, 1] ≈ 1.0 atol=1e-6            #src
@test length(g[4].basis.polynomials) == 1     #src
@test g[4].basis.polynomials[1] ≈ (-x[1] + x[2] - x[3] + x[4]) / 2 #src
@test size(g[4].Q) == (1, 1)                  #src
@test g[4].Q[1, 1] ≈ 1.0 atol=1e-6            #src
for g in gram_matrix(con_ref).sub_gram_matrices
    println(g.basis.polynomials)
end


# Single run of model with 2 parameters
using Turing, Statistics, Distributions, Random, LinearAlgebra, StatsPlots
using AdvancedMH

# the data generating process
function dgp(θ, n)
    [rand.(Exponential(θ[1]),n) rand.(Poisson(θ[2]),n)]
end

# summary statistics for estimation
function moments(y)
    n = size(y,1)
    sqrt(n) .* [mean(y, dims=1)[:]; std(y, dims=1)[:]]
end

@model function abc(z, S, n, calibrated)
    # create the prior: the product of the following array of marginal priors
    θ  ~ arraydist([LogNormal(1.,1.); LogNormal(1.,1.)])
    # sample from the model, at the trial parameter value, and compute statistics
    y = zeros(n,2)
    zs = zeros(S, size(z,1))
    @inbounds for i = 1:S
        y .= dgp(θ, n)
        zs[i,:] .= moments(y) # simulated summary statistics
    end
    # the asymptotic Gaussian distribution of the statistics
    m = mean(zs, dims=1)[:]
    calibrated ? Σ = Symmetric((1. + 1/S)*cov(zs)) : Σ = I
    z ~ MvNormal(m, Σ)
end;

function main(calibrated=true)
    θ⁰ = [2.; 3.] # true parameters
    n = 100 # sample size
    S = 100 # number of simulation draws
    # get data and statistics
    y = dgp(θ⁰, n)
    z = moments(y)
    # sample chains
    length = 5000
    burnin = 500
    chain = sample(abc(z, S, n, calibrated),
        MH(:θ => AdvancedMH.RandomWalkProposal(MvNormal(zeros(2), 0.25*I))),
        MCMCThreads(), length+burnin, 4)
    chain = chain[burnin+1:end,:,:]
    @show display(chain)
    p = plot(chain)
    chain = Array(chain)
    q1 = quantile(chain[:,1], [0.025, 0.975])
    q2 = quantile(chain[:,2], [0.025, 0.975])
    println()
    println("is first true parameter inside 95% CI?")
    display(q1[1] <= θ⁰[1] && q1[2] >= θ⁰[1])
    println()
    println("is second true parameter inside 95% CI?")
    display(q2[1] <= θ⁰[2] && q2[2] >= θ⁰[2])
    p
end
println("first a well-calibrated run")
p1 = main(true)
println("now a poorly calibrated run")
p2 = main(false)
plot(p1, p2)


"""
    wdist(d; h = 2, cri = 4, squared = false)
Compute weights from distances, using a decreasing exponential function.
* `d` : A vector of distances.
* `h` : A scaling positive scalar defining the shape of the function.
* `cri` : A positive scalar defining outliers in the distances vector.
* `squared`: If true, distances are replaced by the squared distances;
    the weight function is then a Gaussian (RBF) kernel function.

Weights are computed by exp(-d / (h * MAD(d))), or are set to 0 for
distances > Median(d) + cri * MAD(d).

The weights decrease with increasing distances. Lower is h, sharper is the decreasing function.
Weights are set to 0 for outliers (extreme distances).
"""
function wdist(d; h = 2, cri = 4, squared = false)
    w = copy(d)
    wdist!(w; h = h, cri = cri, squared = squared)
    w
end

function wdist!(d; h = 2, cri = 4, squared = false)
    # d, out : (n,)
    squared ? d = d.^2 : nothing
    zmed =  Statistics.median(d)
    zmad = Jchemo.mad(d)
    cutoff = zmed + cri * zmad
    d .= map(x -> ifelse(x <= cutoff, exp(-x / (h * zmad)), zero(eltype(d))), d)
    d .= d / maximum(d)
    d[isnan.(d)] .= 1
    return
end


using Colors
using Luxor

struct Pos
    x::Float64
    y::Float64
end

function jarvis_cross(point1::Pos, point2::Pos, point3::Pos)
    vec1 = Pos(point2.x - point1.x, point2.y - point1.y)
    vec2 = Pos(point3.x - point2.x, point3.y - point2.y)
    ret_cross = vec1.x*vec2.y - vec1.y*vec2.x
    return ret_cross*ret_cross
end

function jarvis_march(points::Vector{Pos})
    hull = Vector{Pos}()

    # sorting array based on leftmost point
    sort!(points, by = item -> item.x)
    push!(hull, points[1])

    i = 1
    curr_point = points[2]

    # Find cross product between points
    curr_product = jarvis_cross(Pos(0,0), hull[1], curr_point)
    while (curr_point != hull[1])
        for point in points
                product = 0.0
            if (i == 1)
                if (hull[i] != point)
                    product = jarvis_cross(Pos(0,0), hull[i], point)
                end
            else
                if (hull[i] != point && hull[i-1] != point)
                    product = jarvis_cross(hull[i-1], hull[i], point)
                end
            end
            if (product > curr_product)
                curr_point = point
                curr_product = product
            end
        end
        push!(hull, curr_point)
        curr_product = 0
        i += 1
    end

    return hull
end

function main(pts)
    hull = jarvis_march(pts)
end

# Auxiliary drawing functions
function make_drawing(width, height, img_path, bkg_color, origin_p)
    d = Drawing(width, height, img_path)
    background(bkg_color)
    origin(origin_p)
    return d
end

map2luxor(p) = Point(p.x, -p.y)
map2luxor(p) = Point(p...)

## Canvas constants
width = 500
height = 500
path = "arcane_march.png"
color = RGBA(0, 0, 0)
op = [width / 2, height / 2]

# Setting up drawing
my_draw = make_drawing(width, height, path, color, Point(op...))

# pts_list = [Pos(200, 15), Pos(100, 100), Pos(-20, -200), Pos(-130, 100)]
pts_list = [Pos(rand(-250:250), rand(-250:250)) for _ in 1:10]

hull = main(pts_list)
println(hull)
hull = [Point(p.x, -p.y) for p in hull]
println(hull)
pts_list = [Point(p.x, -p.y) for p in pts_list]

sethue(RGBA(1.0, 1.0, 0.0))
circle.(pts_list, 15; action = :fill)

sethue(RGBA(1.0, 1.0, 1.0))
poly(hull, action = :stroke; close = true)

finish()


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

"""
    gpol1(ct::Vector{<:Real})::Function
return a degree 1 polynomial function with parameters ct.
"""
function gpol1(ct::Vector{<:Real})::Function
    function f1(z::Real)
        return ct[1] + ct[2] * z
    end
    return f1
end

"""
    gpol2(ct::Vector{<:Real})::Function
return a degree 2 polynomial function with parameters ct.
"""
function gpol2(ct::Vector{<:Real})::Function
    function f2(z::Real)
        return ct[1] + ct[2] * z + ct[3] * z^2
    end
    return f2
end

"""
    gpol2(ct::Vector{<:Real})::Function
return a degree 3 polynomial function with parameters ct.
"""
function gpol3(ct::Vector{<:Real})::Function
    function f3(z::Real)
        return ct[1] + ct[2] * z + ct[3] * z^2 + ct[4] * z^3
    end
    return f3
end

"""
    func1dfit(ffit::Function, x::Vector{<:Real},
              y::Vector{<:Real}, p0::Vector{<:Real},
              lb::Vector{Float64}, ub::Vector{Float64})

Fit a function to the data x, y with start prediction p0
and return coefficients and errors.
"""
function func1dfit(ffit::Function, x::Vector{<:Real},
                   y::Vector{<:Real}, p0::Vector{<:Real},
                   lb::Vector{Float64}, ub::Vector{Float64})
    fq = curve_fit(ffit, x, y, p0, lower=lb, upper=ub)
    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

"""
    func1dfit(ffit::Function, x::Vector{<:Real}, y::Vector{<:Real},
              yerr::Vector{<:Real}, p0::Vector{<:Real},
              lb::Vector{<:Real}, ub::Vector{<:Real})
Fit function ffit to data x,y taking into account the errors
on the y (yerr). Assumes weights of 1/sigma^2 as in standard least squares
and returns fit result.
"""
function func1dfit(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

"""
    fit_pol1(x::Vector{<:Real}, y::Vector{<:Real}, ci=0.1)
Fit a 1st degree polynomial to the data x, y and return
a RFit object with the result.
"""
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

"""
    fit_pol2(x::Vector{<:Real}, y::Vector{<:Real}, ci=0.1)
Fit a 2nd degree polynomial to the data x, y and return
a RFit object with the result.
"""
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

"""
    fit_pol3(x::Vector{<:Real}, y::Vector{<:Real}, ci=0.1)
Fit a 3rd degree polynomial to the data x, y and return
a RFit object with the result.
"""
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

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

"""
    gausg(μ::Real, σ::Real, C::Real)::Function
Return a Gaussian function with parameters μ, σ and normalisation C.
"""
function gausg(μ::Real, σ::Real, C::Real)::Function
	function gausx(x::Real)
		return C * pdf(Normal(μ, σ,), x)
	end
	return gausx
end

"""
    gausg2(μ1::Real, σ1::Real, C1::Real, μ2::Real, σ2::Real, C2::Real)
Return a function for the sum of two Gaussians with parameters
    μ1, σ1 and μ2, σ2 and normalisations C1 and C2.
"""
function gausg2(μ1::Real, σ1::Real, C1::Real, μ2::Real, σ2::Real, C2::Real)::Function
	function gausx(x::Real)
		return C1 * pdf(Normal(μ1, σ1,), x) + C2 * pdf(Normal(μ2, σ2,), x)
	end
	return gausx
end

@. 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)
@. gauss2cm(x, p) = p[1] * pdf(Normal(p[2], p[3]), x) + p[4] * pdf(Normal(p[2], p[5]), x)

"""
    gaussfm(mu::Real)::Function
Return a Gaussian Fitting function with a fixed mean mu.
"""
function gaussfm(mu::Real)::Function
	function gauss(x::Vector{<:Real}, p::Vector{<:Real})
		return @. p[1]* pdf(Normal(mu, p[2]), x)
	end
	return gauss
end

"""
    gauss2fm(mu::Real)::Function
Return a function of the sum of two Gaussians with a fixed common mean mu.
"""
function gauss2fm(mu::Real)::Function
	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

"""
	fit_gauss(h::Histogram)

Fit a normal distribution to a Histogram and return FGauss object.
"""
function fit_gauss(h::Histogram)
	c = centers(h)
	w = eltype(c).(h.weights)
	@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.0, mu - 100.0 * sigma, sigma / 100.0]
    ub = [100 * sum(w), mu + 100.0 * sigma, 100.0 * sigma]
    p0 = [      sum(w), mu                ,         sigma]
	(CC, μ, σ), _ = func1dfit(gauss1, c, w, p0, 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 a normal distribution to data x,y and return FGauss object.
"""
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(x::Vector{Float64}, xmin::Float64, xmax::Float64, bins=50, fm=0.0)

Fit a gaussian with a fixed mean fm and return an FGauss object.
"""
function fit_gauss_fm(x::Vector{Float64}, xmin::Float64, xmax::Float64;
	                  bins=50, norm=false, fm=0.0)

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

	# bin distribution
    h = hist1d(xx, bins, xmin, xmax, norm)
    c = centers(h)
    w = eltype(c).(h.weights)
    @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 = [         sum(w), σ       ]

	g1 = gaussfm(fm)
	(CC, sigma), _ = func1dfit(g1, c, w, p0, lb, ub)

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

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

"""
    plot_fit_gauss(x::Vector{Float64}, xs::String, ys::String,
                   bins::Integer, xmin::Float64, xmax::Float64;
                   xgmin::Float64, xgmax::Float64, gbins::Integer=50)
Fit a Gaussian to the histogram of data x between xmin and xmax with bins
and plot the results.
"""
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, _ = 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::Vector{<:Real}, xs::String, nbins::Int64, xmin::Real, xmax::Real;
	      xgmin::Real, xgmax::Real, fbins::Int64=100, norm::Bool=true,
          fm::Real=0.0, flex_mean=false)

Histograms data x and returns a plot and a Gaussian fit to the histogram.
Option flex_mean allows the mean to float or be fixed to fm.
"""
function fitg1(x::Vector{<:Real}, xs::String, nbins::Int64, xmin::Real, xmax::Real;
	           xgmin::Real, xgmax::Real, fbins::Int64=100, norm::Bool=true,
               fm::Real=0.0, flex_mean=false)

	h, p = hist1d(x, xs, nbins, 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 fixed mean, 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 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 JJ Review!!!
"""
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

"""
    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. JJ Review!!
"""
function fitg2(x::Vector{<:Real}, xs::String, bins::Int64, xmin::Real, xmax::Real;
	           xg1min::Real, xg1max::Real, xg2min::Real, xg2max::Real, xgmin::Real, xgmax::Real, cm=0.0,
         	   g1bins::Int64=100, g2bins::Int64=100, gbins::Int64=100, norm::Bool=true)

	_, 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


#The lattice
using LinearAlgebra

#print("Enter dim value: "); d = parse(Int, readline(stdin)) # use it for IO input of d! else use fixed d!
d = 4;
v = Matrix{Any}(zeros(d,d));
h = Matrix{Any}(zeros(d,d));
p = Matrix{Float64}(undef,d,d);
x = Matrix{Float64}(undef,d,d);
y = Matrix{Float64}(undef,d,d);
Φ = Matrix{Float64}(undef,d,d);

for i = 1:d
    h[i,d] = Matrix{Float64}(I,d,d)[i,:];
end

#=for k = d:-1:1
    v[k,k] = (x[k,k] + y[k,k]im)*h[k,k];

    for j = k-1:-1:1
        v[j,k] = (x[j,k] + y[j,k]im) * sqrt(complex(p[j,k])) * h[j,k] + sqrt(complex(1-p[j,k])) * v[j+1,k];
        h[j,k-1] = (x[j,k] + y[j,k]im) * sqrt(complex(1-p[j,k])) * h[j,k] - sqrt(complex(p[j,k])) * v[j+1,k];
        println("k:",k,"\t","j:",j)
    end
end=#

function H(m::Int, n::Int)
    for k = d:-1:1
        v[k,k] = (x[k,k] + y[k,k]im)*h[k,k];

        for j = k-1:-1:1
            v[j,k] = (x[j,k] + y[j,k]im) * sqrt(complex(p[j,k])) * h[j,k] + sqrt(complex(1-p[j,k])) * v[j+1,k];
            h[j,k-1] = (x[j,k] + y[j,k]im) * sqrt(complex(1-p[j,k])) * h[j,k] - sqrt(complex(p[j,k])) * v[j+1,k];
        end
    end
    return (h[m,n])
end

#Ist = transpose(v);

#-----------------------------------------------------------------------------------------------------
#DFT

m(n::Int) = 1/sqrt(n)*[exp(2*pi*im*(r-1)*(s-1)/n) for r in 1:n, s in 1:n];

#Soll = m(d);

#------------------------------------------------------------------------------------------------------
#The solution of the algorithm

Denom(j::Int,k::Int) = dot(Matrix{Float64}(I,d,d)[j,:], H(j,k)) * (j==1 ? 1.0 : prod(sqrt(complex(1-p[l,k])) for l = 1:j-1));

Num(j::Int,k::Int) = j==1 ? 0.0 : sum((x[i,k] + y[i,k]im) * dot(Matrix{Float64}(I,d,d)[j,:], H(i,k)) * sqrt(complex(p[i,k])) * (i==1 ? 1.0 : prod(sqrt(complex(1-p[l,k])) for l = 1:i-1)) for i = 1:j-1);

Z(j::Int,k::Int) = (m(d)[j,k] - Num(j,k)) / Denom(j,k);

#=function Angle(Re,Im,j,k)
    if abs(exp(asin(Float32(Im))*im) - (Re + Im*im)) < 10.0^-12
        Φ[j,k] = asin(Float32(Im));
    else
        Φ[j,k] = -asin(Float32(Im)) - pi;
    end
end=#
#Angle(Re,Im,j,k) = abs(exp(asin(Float32(Im))*im) - (Re + Float32(Im)*im)) < 10.0^-12 ? Φ[j,k] = asin(Float32(Im)) : Φ[j,k] = -asin(Float32(Im)) - pi;
Angle(Re::Float16,Im::Float16,j::Int,k::Int) = abs2(exp(asin(Im)*im) - (Re + Im*im)) < 10.0^-12 ? Φ[j,k] = asin(Im) : Φ[j,k] = -asin(Im) - pi;
# the results vary depending on the precision of the 'Float' used!

#-----------------------------------------------------------------------------------------------------
#Evaluation

@time for k = d:-1:1
    for j = 1:k-1
        p[j,k] = abs2(Z(j,k));
        x[j,k] = real(Z(j,k)/sqrt(complex(p[j,k])));
        y[j,k] = imag(Z(j,k)/sqrt(complex(p[j,k])));
        Angle(Float16(x[j,k]),Float16(y[j,k]),j,k);
        println("p(", j, ",", k, ") = ", p[j,k], "\tΦ(", j, ",", k, ") = ", Φ[j,k]*180/pi);
    end

    x[k,k] = real(Z(k,k));
    y[k,k] = imag(Z(k,k));
    Angle(Float16(x[k,k]),Float16(y[k,k]),k,k);
    println("Φ(", k, ",", k, ") = ", Φ[k,k]*180/pi);
end

#------------------------------------------------------------------------------------------------------


# example with constraint in struct
# struct Vector2D{T<:Number}
#    x::T
#    y::T
# end

# function dot(u::Vector2D{T}, v::Vector2D{T}) where T
#     u.x*v.x + u.y*v.y
# end

# function cross(u::Vector2D{T}, v::Vector2D{T}) where T
#     u.x*v.y - u.y*v.x
# end

# julia> a = Vector2D(1, 2)
# Vector2D{Int64}(1, 2)

# julia> b = Vector2D(3, 4)
# Vector2D{Int64}(3, 4)

# julia> c =Vector2D{Char}('A', 'B')
# ERROR: TypeError: in Vector2D, in T, expected T<:Number, got Type{Char}
# Stacktrace:
#  [1] top-level scope
#    @ REPL[11]:1

# julia> dot(a,b)
# 11

# julia> cross(a,b)
# -2

# dot(a,c)

# same example with simpler method signature
# in this version  u and v can be of different Number subtype and still works
# struct Vector2D{T<:Number}
#    x::T
#    y::T
# end

# function dot(u::Vector2D, v::Vector2D)
#     u.x*v.x + u.y*v.y
# end

# function cross(u::Vector2D, v::Vector2D)
#     u.x*v.y - u.y*v.x
# end

# julia> a = Vector2D(1, 2)
# Vector2D{Int64}(1, 2)

# julia> b = Vector2D(3.0, 4.0)
# Vector2D{Float64}(3.0, 4.0)

# julia> c =Vector2D{Char}('A', 'B')
# ERROR: TypeError: in Vector2D, in T, expected T<:Number, got Type{Char}
# Stacktrace:
#  [1] top-level scope
#    @ REPL[7]:1

# julia> dot(a,b)
# 11.0

# julia> cross(a,b)
# -2.0

# example with constraint in method - allows non Number Vecttor2
# but methods are type protected
# struct Vector2D{T}
#    x::T
#    y::T
# end

# function dot(u::Vector2D{T}, v::Vector2D{T}) where T <: Number
#     u.x*v.x + u.y*v.y
# end

# function cross(u::Vector2D{T}, v::Vector2D{T}) where T <: Number
#     u.x*v.y - u.y*v.x
# end

# julia> a = Vector2D(1, 2)
# Vector2D{Int64}(1, 2)

# julia> b = Vector2D(3, 4)
# Vector2D{Int64}(3, 4)

# julia> c =Vector2D{Char}('A', 'B')
# Vector2D{Char}('A', 'B')

# julia> a = Vector2D(1, 2)
# Vector2D{Int64}(1, 2)

# julia> b = Vector2D(3, 4)
# Vector2D{Int64}(3, 4)

# julia> c =Vector2D{Char}('A', 'B')
# Vector2D{Char}('A', 'B')

# julia> dot(a,b)
# 11

# julia> cross(a,b)
# -2

# julia> dot(a,c)
# ERROR: MethodError: no method matching dot(::Vector2D{Int64}, ::Vector2D{Char})
# Closest candidates are:
#   dot(::Vector2D{T}, ::Vector2D{T}) where T<:Number at REPL[4]:1
# Stacktrace:
#  [1] top-level scope
#    @ REPL[14]:1

# julia> cross(a,c)
# ERROR: MethodError: no method matching cross(::Vector2D{Int64}, ::Vector2D{Char})
# Closest candidates are:
#   cross(::Vector2D{T}, ::Vector2D{T}) where T<:Number at REPL[5]:1
# Stacktrace:
#  [1] top-level scope
#    @ REPL[15]:1


#=

    This is a port of the SPPL model election.pynb

    This model:
    Summary Statistics
    parameters        mean        std   naive_se      mcse         ess      rhat   ess_per_sec
        Symbol     Float64    Float64    Float64   Float64     Float64   Float64       Float64

        param    299.7774    28.9717     0.2897    0.3653   6479.7518    0.9999      866.9724
            p      0.4814     0.0318     0.0003    0.0004   6508.3665    1.0002      870.8010
        votes   1925.2743   131.9134     1.3191    1.7207   6442.0098    1.0003      861.9226
          win      0.2902     0.4539     0.0045    0.0056   6597.9243    0.9999      882.7836

=#

using Turing, StatsPlots, Distributions, StatsBase
using CSV
include("jl_utils.jl")

#=
from sppl.sym_util import binspace
n = 4000
param ~= randint(low=250, high=350)
switch (param) cases (b in range(250, 350)):
    p ~= beta(a=277, b=b)
switch (p) cases (x in binspace(0, 1, 20)):
    votes ~= binom(n=n, p=(x.left + x.right)/2)
win ~= votes > 0.5*n

=#
@model function election(n=4000)
    param ~ DiscreteUniform(250,350)
    p ~ Beta(277,param)
    votes ~ Binomial(n,p)

    win ~ Dirac(votes > 0.5*n)

end

n = 4000
model = election(n)

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 10_000)
chns = sample(model, PG(5), 10_000)
# chns = sample(model, SMC(), 10_000)
# chns = sample(model, IS(), 10_000)

# chns = sample(model, HMC(0.1,6), 1_000)
# chns = sample(model, NUTS(), 1_000)

display(chns)