import Distributions: logpdf, DiscreteUnivariateDistribution

"""
Numerically stable Poisson log likelihood function. Accepts log of rate parameter.
"""
struct LogPoisson{T<:Real} <: DiscreteUnivariateDistribution
    logλ::T
end

function logpdf(lp::LogPoisson, k::Int)
    return k * lp.logλ - exp(lp.logλ) - loggamma(k + 1)
end

@model AHMCpoisson(y, x, idx, N, Ns) = begin
  a0 ~ Normal(0, 10)
  a1 ~ Normal(0, 1)
  a0_sig ~ Truncated(Cauchy(0, 1), 0.0, Inf)
  a0s ~ MvNormal(zeros(Ns), a0_sig)
  for i ∈ 1:N
    λ = a0 + a0s[idx[i]] + a1 * x[i]
    y[i] ~ LogPoisson(λ)
  end
end

AHMCconfig = Turing.NUTS(1000, .80)

function simulatePoisson(; Nd=1, Ns=10, a0=1.0, a1=.5, a0_sig=.3, kwargs...)
  N = Nd * Ns
  y = fill(0, N)
  x = fill(0.0, N)
  idx = similar(y)
  i = 0
  for s in 1:Ns
    a0s = rand(Normal(0, a0_sig))
    logpop = rand(Normal(9, 1.5))
    λ = exp(a0 + a0s + a1 * logpop)
    for nd in 1:Nd
      i += 1
      x[i] = logpop
      idx[i] = s
      y[i] = rand(Poisson(λ))
    end
  end
  return (y=y, x=x, idx=idx, N=N, Ns=Ns)
 end

CmdStanPoisson = "
data {
  int N;
  int y[N];
  int Ns;
  int idx[N];
  real x[N];
}
parameters {
  real a0;
  vector[Ns] a0s;
  real a1;
  real<lower=0> a0_sig;
}
model {
  vector[N] mu;
  a0 ~ normal(0, 10);
  a1 ~ normal(0, 1);
  a0_sig ~ cauchy(0, 1);
  a0s ~ normal(0, a0_sig);
  for(i in 1:N) mu[i] = exp(a0 + a0s[idx[i]] + a1 * x[i]);
  y ~ poisson(mu);
}
"

CmdStanConfig = Stanmodel(
  name="CmdStanPoisson", model=CmdStanPoisson, nchains=1, output_format=:mcmcchains,
  Sample(num_samples=1000, num_warmup=1000, adapt=CmdStan.Adapt(delta=0.8), save_warmup=true)
)

struct PoissonProb
  y::Array{Int64,1}
  x::Array{Float64,1}
  idx::Array{Int64,1}
  N::Int64
  Ns::Int64
end

function (problem::PoissonProb)(θ)
  @unpack y, x, idx, N, Ns = problem   # extract the data
  @unpack a0, a1, a0s, a0_sig = θ
  LL = 0.0
  LL += logpdf(Truncated(Cauchy(0, 1), 0.0, Inf), a0_sig)
  LL += sum(logpdf(MvNormal(zeros(Ns), a0_sig), a0s))
  LL += logpdf.(Normal(0, 10), a0)
  LL += logpdf.(Normal(0, 1), a1)
  for i in 1:N
    λ = a0 + a0s[idx[i]] + a1 * x[i]
    LL += logpdf(LogPoisson(λ), y[i])
  end
  return LL
end

# Define problem with data and inits.
function sampleDHMC(y, x, idx, N, Ns, nsamples, autodiff)
  p = PoissonProb(y, x, idx, N, Ns)
  p((a0=0.0, a1=0.0, a0s=fill(0.0, Ns), a0_sig=.3))
  # Write a function to return properly dimensioned transformation.
  trans = as((a0=asℝ, a1=asℝ, a0s=as(Array, Ns), a0_sig=asℝ₊))
  P = TransformedLogDensity(trans, p)
  ∇P = ADgradient(autodiff, P)
  # Sample from the posterior.
  results = mcmc_with_warmup(Random.GLOBAL_RNG, ∇P, nsamples; reporter = NoProgressReport())
  # Undo the transformation to obtain the posterior from the chain.
  posterior = TransformVariables.transform.(trans, results.chain)
  chns = nptochain(results, posterior)
  return chns
end


"""
    kompsimp(f,a,b,M)
adalah fungsi yang digunakan untuk menghitung nilai integral `f` secara numerik
pada interval `[a,b]` menggunakan komposit Simpson dengan `M` sub-interval.

# Example
```jl
julia> a = 1; b = 6;

julia> f(x) = 2+sin(2*sqrt(x));

julia> y = kompsimp(f,a,b,10)
8.183447496636239
```
return solusi `y`
"""
function kompsimp(f,a,b,M)
  h = (b-a)/(2*M);
  s1 = 0;   s2 = 0;
  for k = 1:M
    x = a + (2*k-1)*h;
    s1 = s1 + f(x);
  end
  for k = 1:M-1
    x = a + (2*k)*h;
    s2 = s2 + f(x);
  end
  y = h/3*(f(a)+f(b)+4*s1+2*s2)
end


"""
    function length_conversion(value, from_type, to_type)

A function that converts a value from a measurement unit to another one

Accepted units are: millimeter(s), centimeter(s), meter(s), kilometer(s),
inch(es), feet, foot, yard(s), mile(s). Abbreviations are also supported.

# Examples/Tests (optional but recommended)
```julia
length_conversion(10, "METERS", "cm")
length_conversion(12345, "yards", "FEET")
```

Because the algorithm converts values to meters, and from meters to the final type,
some imperial system units may have errors:
length_conversion(1, "yards", "FEET") returns 3.000000096, intead of 3.

# Contributed by:- [Fernanda Kawasaki](https://github.com/fernandakawasaki)
"""

# Lookup table that returns conversion of 1 unit of type to meters
METER_CONVERSION = Dict{String,Float64}(
    "mm" => 0.001,
    "cm" => 0.01,
    "m" => 1,
    "km" => 1000,
    "in" => 0.0254,
    "ft" => 0.3048,
    "yd" => 0.9144,
    "mi" => 1609.34,
)

# Lookup table that returns the conversion of 1 meter to type
TYPE_CONVERSION = Dict{String,Float64}(
    "mm" => 1000,
    "cm" => 100,
    "m" => 1,
    "km" => 0.001,
    "in" => 39.3701,
    "ft" => 3.28084,
    "yd" => 1.09361,
    "mi" => 0.000621371,
)

NAME_CONVERSION = Dict{String,String}(
    "millimeter" => "mm",
    "millimeters" => "mm",
    "centimeter" => "cm",
    "centimeters" => "cm",
    "meter" => "m",
    "meters" => "m",
    "kilometer" => "km",
    "kilometers" => "km",
    "inch" => "in",
    "inches" => "in",
    "feet" => "ft",
    "foot" => "ft",
    "yard" => "yd",
    "yards" => "yd",
    "mile" => "mi",
    "miles" => "mi",
)

function normalize_type(type)
    l_type = lowercase(type)
    if !haskey(METER_CONVERSION, l_type)
        if !haskey(NAME_CONVERSION, l_type)
            throw(
                error(
                    "Invalid 'type' value: $(type)\n",
                    "Supported values are: $(keys(NAME_CONVERSION))\n",
                    "Supported abbreviations are: $(keys(METER_CONVERSION))\n",
                ),
            )
        end
        return NAME_CONVERSION[l_type]
    end
    return l_type
end

function length_conversion(value, from_type, to_type)
    from_type_norm = normalize_type(from_type)
    to_type_norm = normalize_type(to_type)
    value_in_meters = value * METER_CONVERSION[from_type_norm]
    return value_in_meters * TYPE_CONVERSION[to_type_norm]
end


### Error functions.

function error_function_selector(error::ASCIIString)
	if error == "absolute_error"
		return "absolute_error", absolute_error!, absolute_error_prime!
	elseif error == "cross_entropy"
		return "cross_entropy", cross_entropy!, cross_entropy_prime!
	else
		return "squared_error", squared_error!, squared_error_prime!
	end
end

function absolute_error!{T<:FloatingPoint}(YH::AbstractVector{T}, Y::AbstractVector{T}, E::AbstractVector{T}; eta::T=1e-10)
	@inbounds begin
		for i = 1:length(E)
			if YH[i] > Y[i]
				E[i] = YH[i] - Y[i]
			else
				E[i] = Y[i] - YH[i]
			end
		end
	end
end

function absolute_error_prime!{T<:FloatingPoint}(YH::AbstractVector{T}, Y::AbstractVector{T}, DE_DYH::AbstractVector{T}; eta::T=1e-10)
	@inbounds begin
		for i = 1:length(Y)
			if YH[i] > Y[i]
				DE_DYH[i] = 1.0
			else
				DE_DYH[i] = -1.0
			end
		end
	end
end

function cross_entropy!{T<:FloatingPoint}(YH::AbstractVector{T}, Y::AbstractVector{T}, E::AbstractVector{T}; eta::T=1e-10)
	@inbounds begin
		for i = 1:length(E)
			E[i] = -((Y[i] * log(YH[i] + eta)) + ((1.0 - Y[i]) * log(1.0 - YH[i] + eta)))
		end
	end
end

function cross_entropy_prime!{T<:FloatingPoint}(YH::AbstractVector{T}, Y::AbstractVector{T}, DE_DYH::AbstractVector{T}; eta::T=1e-10)
	@inbounds begin
		for i = 1:length(Y)
			DE_DYH[i] = ((1.0 - Y[i]) / (1.0 - YH[i] + eta)) - (Y[i] / (YH[i] + eta))
		end
	end
end

function squared_error!{T<:FloatingPoint}(YH::AbstractVector{T}, Y::AbstractVector{T}, E::AbstractVector{T})
	@inbounds begin
		for i = 1:length(E)
			E[i] = 0.5 * abs2(YH[i] - Y[i])
		end
	end
end

function squared_error_prime!{T<:FloatingPoint}(YH::AbstractVector{T}, Y::AbstractVector{T}, DE_DYH::AbstractVector{T})
	@inbounds begin
		for i = 1:length(Y)
			DE_DYH[i] = YH[i] - Y[i]
		end
	end
end


"""
```
initialize_stategrid(method::Symbol, dims::Vector{Int})
```
constructs a state grid (using the implementation in EconPDEs)
using various methods, such as Chebyshev points.
"""
function initialize_stategrid(method::Symbol, grid_info::OrderedDict{Symbol, Tuple{T, T, Int}};
                              get_stategrid::Bool = true) where {T <: Real}
    stategrid_init = OrderedDict{Symbol, Vector{T}}()
    if method == :uniform
        for (k, v) in grid_info
            stategrid_init[k] = range(v[1], stop = v[2], length = v[3])
        end
    elseif method == :chebyshev
        for (k, v) in grid_info
            stategrid_init[k] = v[1] .+ 1. / 2. * (v[2] - v[1]) .* (1. .- cos(pi * (0:(v[3] - 1))' / (v[3] - 1)))
        end
    elseif method == :exponential
        for (k, v) in grid_info
            stategrid_init[k] = exp.(range(log(v[1]), stop = log(v[2]), length = v[3]))
        end
    elseif method == :smolyak
        error("Construction of a Smolyak interpolation grid has not been implemented yet.")
        # This should make a call to BasisMatrices from QuantEcon, as they have a nice user-friendly implementation of Smolyak.
        # Alternatively, we can try SmolyakApprox, tho QuantEcon seems more likely to be well-maintained.
    else
        error("Grid construction method $method has not been implemented.")
    end

    if get_stategrid
        return StateGrid(stategrid_init)
    else
        return stategrid_init
    end
end


"""
nonparametric regression and quantile regression
* uses pre-generated kernel weights
* can be local constant, local linear, or local quadratic

y is n X p

x is n X dimx

xeval is neval X dimx

weights is n X neval (different weights for each eval. point)

weights should sum to one by columns, they are typical
nonparametric weights from a kernel.

keyword order=0,1 or 2 for local consant, local linear, or local quadratic
(default local linear)

keyword do_median (default false)
keyword do_ci to compute .05 and 0.95 quantiles

execute npreg() for an example
"""

using Random, Econometrics, Plots
function npreg(bandwidth=-1)
    println("npreg(), with no arguments, runs a simple example")
    println("npreg(bw) will run the example with your chosen bandwidth (Float64)")
    println("execute edit(npreg,()) to see the code")
    k = 1 # number of regressors
    Random.seed!(1) # set seed to
    n = 1000
    bandwidth == -1 ? bandwidth = 0.25*n^(-1.0/(4 + k)) : nothing
    neval = 100
    x = rand(n)*pi*2.0 # from 0 to 2π
    xeval = collect(range(pi/2., stop=pi*1.5, length=neval)) # from π/2 to 3π/2
    y = cos.(x) .+ 0.25 .* cos.(3.0.*x) + 0.1*randn(n)
    # npreg wants args to be matrices, not vectors, so convert them
    y = reshape(y,n,1)
    x = reshape(x,n,1)
    xeval = reshape(xeval,neval,1)
    ytrue = cos.(xeval) .+ 0.25.*cos.(3.0*xeval)
    weights = kernelweights(x, xeval, bandwidth, true, "knngaussian", 200)
    yhat, y50, y05, y95 = npreg(y, x, xeval, weights, order=1, do_median=true, do_ci=true)
    labels = ["true" "mean" "median" "0.05 quantile" "0.95 quantile"]
    title = "Kernel regression and quantiles"
    p = Plots.plot(xeval, [ytrue yhat y50 y05 y95], labels=labels, title=title)
    display(p)
    nothing
end

function npreg(y, x, xeval, weights; order=1, do_median=false, do_ci=false)
    weights = sqrt.(weights)
    neval, dimx = size(xeval)
    n, dimy = size(y)
    # local constant
    if order==0
        X = ones(n,1)
        Xeval = ones(neval,1)
    elseif order==1
    # local linear
        X = [ones(n,1) x]
        Xeval = [ones(neval,1) xeval]
    else
    # local quadratic
        stacked = [x; xeval]
        # cross products
        CP = zeros(n+neval, Int((dimx-1)*dimx/2))
        cpind = 0
        for i = 1:dimx-1
            for j = (i+1):dimx
                cpind += 1
                CP[:,cpind] = stacked[:,[i]].*stacked[:,[j]]
            end
        end
        ZZ = [ones(n+neval,1) stacked CP]
        X = view(ZZ,1:n,:)
        Xeval = view(ZZ,(n+1):n+neval,:)
    end
    # do the fit
    yhat = zeros(neval, dimy)
    y50 = nothing
    y05 = nothing
    y95 = nothing
    if do_median y50 = zeros(neval, dimy) end
    if do_ci
        y05 = zeros(neval, dimy)
        y95 = zeros(neval, dimy)
    end
    for i = 1:neval
        WX = weights[:,i] .* X
        Wy = weights[:,i] .* y
        b = WX\Wy
        yhat[i,:] = Xeval[i:i,:]*b
        if do_median
            for j = 1:dimy
                y50[i,j] = (Xeval[i:i,:]*qreg_coef(Wy[:,j], WX, 0.5))[1,1]
            end
        end
        if do_ci
            for j = 1:dimy
                y05[i,j] = (Xeval[i:i,:]*qreg_coef(Wy[:,j], WX, 0.05))[1,1]
                y95[i,j] = (Xeval[i:i,:]*qreg_coef(Wy[:,j], WX, 0.95))[1,1]
            end
        end
    end
    if do_median==false
        return yhat
    elseif do_ci==false
        return yhat, y50
    else
        return yhat, y50, y05, y95
    end
end


function harmonic(n::Integer)
    H = 0.0
    for i in n:-1:1
        H += 1 / i
    end
    return H
end

function harmonic(n::Integer, m::Integer)
    H = 0.0
    for i in n:-1:1
        H += 1 / (i ^ m)
    end
    return H
end

td_a1(n::Integer) = harmonic(n - 1)

td_a2(n::Integer) = harmonic(n - 1, 2)

td_b1(n::Integer) = (n + 1) / (3 * (n - 1))

function td_b2(n::Integer)
    return (2 * ((n ^ 2) + n + 3)) / ((9 * n) * (n - 1))
end

td_c1(a1::AbstractFloat, b1::AbstractFloat) = b1 - 1 / a1

function td_c2(n::Integer, a1::AbstractFloat, a2::AbstractFloat, b2::AbstractFloat)
    return b2 - ((n + 2) / (a1 * n)) + a2 / (a1 ^ 2)
end

td_e1(a1::AbstractFloat, c1::AbstractFloat) = c1 / a1

td_e2(a1::AbstractFloat, a2::AbstractFloat, c2::AbstractFloat) = c2 / (a1 ^ 2 + a2)

function tajimad(π::AbstractFloat, S::Integer, a1::AbstractFloat, e1::AbstractFloat, e2::AbstractFloat)
    return (π - S / a1) / sqrt(e1 * S + e2 * S * (S - 1))
end

"""
    tajimad(π::AbstractFloat, S::Integer, n::Integer)

Compute Tajima's D from:

* `π`: The average number of SNPs found in (n choose 2) pairwise comparisons of
       a sample of sequences.

* `S`: The number of segregating sites in a sample of sequences.

* `n`: The number of sequences in your sample.

*Example*

```julia
tajimad(3.88888, 16, 10)
```

"""
function tajimad(π::AbstractFloat, S::Integer, n::Integer)
    a1 = td_a1(n)
    a2 = td_a2(n)
    b1 = td_b1(n)
    b2 = td_b2(n)
    c1 = td_c1(a1, b1)
    c2 = td_c2(n, a1, a2, b2)
    e1 = td_e1(a1, c1)
    e2 = td_e2(a1, a2, c2)
    return tajimad(π, S, a1, e1, e2)
end

"""
    tajimad(seqs)

Compute Tajima's D from a collection of BioSequences{DNAAlphabet{n}} (n = 2 or 4).

This will estimate the `π`, `S`, and `n` parameters from the sequences and use
those parameters to estimate Tajima's D.

*Example*

```julia

sample = [dna"ATAATAAAAAAATAATAAAAAAATAAAAAAAATAAAAAAAA",
          dna"AAAAAAAATAAATAATAAAAAAATAAAAAAAAAAAAAAAAA",
          dna"AAAATAAAAATATAATAAAAAAATATAAAAAAAAAAAAAAA",
          dna"AAAAAAAAAAAATAATAAAAAAATAAATAAATAAAAAAAAA",
          dna"AAAATAAAAAAAATATAAAAAAATAAAAAAAAAAAAAAAAA",
          dna"AAAATAAAAAAAAAATAAAAAAAAAAAAAAAAAAATAAAAA",
          dna"AAAAAATAAAAATAATAAAAAAATAAAAAAAAAAAAAAAAA",
          dna"AAAAAAAAAAAAAAATAAAAAAATAAAAAAAAAAAAAAATA",
          dna"AAAAAAAAAAAAAAAAAAAAAAATAAAAAAAAAAAAAAAAA",
          dna"AAAAAAAAAAAAAAATAAAAAAATAATAAAAAAAAAAAAAA"]

tajimad(sample)
```

"""
function tajimad(seqs)
    π = avg_mut(seqs)
    S = count(Segregating, seqs)
    n = length(seqs)
    return tajimad(π, S[1], n)
end


module Excercise_4

# Exercise 4 (JULIA)
# 1. Create two matrices of the same layout and test if addition and subtraction of the matrix works as expected: C = A + B
# 2. Now compare matrix multiplication either this way A * B and this way A .* B. Whats the difference?!
# 3. What about matrix division with "/" or "\"?!
# 4. Create a 3x3 integer matrix A with useful numbers. Now try A+1, A-1, A*2, A/2.
# 5. Now multiply a 3x4 matrix with a suitable (4)vector.

# Part 1
matrix_a = rand(0:10, 4,4)
matrix_b = rand(0:10, 4,4)
println("Matrix A:")
println(matrix_a)
println("Matrix B:")
println(matrix_b)
println("A + B:")
addition = matrix_a + matrix_b
println(addition)
println("A - B:")
subtraction = matrix_a -matrix_b
println(subtraction)

# Part 2

# Matrix multiplication
println("A * B")
multi_1 = matrix_a * matrix_b
println(multi_1)

# Element-wise multiplication
println("A .* B:")
multi_2 = matrix_a .* matrix_b
println(multi_2)
# Part 3

# Part 4

matrix_c = [1 2 3; 3 2 1; 2 3 1]
println("Matrix C:")
println(matrix_c)
println("C+1:")
# println(matrix_c-1)
println("C*2:")
println(matrix_c*2)
println("C/2:")
println(matrix_c/2)

# Part 5

matrix_d = rand(0:10, 3,4)

println("Matrix D:")
println(matrix_d)

end


# Energy

function NbodyEnergy(u,Gm)
"""
     Nbody problem Hamiltonian (Cartesian Coordinates)
"""

    dim=2
    nbody=length(Gm)

 @inbounds begin
    x = view(u,1:7)   # x
    y = view(u,8:14)  # y
    v = view(u,15:21) # x′
    w = view(u,22:28) # y′

    H=zero(eltype(u))
    P=zero(eltype(u))

    for i in 1:nbody
        H+=Gm[i]*(v[i]*v[i]+w[i]*w[i])
        for j in i+1:nbody
            r = ((x[i]-x[j])^2+(y[i]-y[j])^2)^(1/2)
            P+=(Gm[i]/r)*Gm[j]
        end
    end

    return(H/2-P)
    end

end

# OdeProblem

function f(du,u,p,t)
  @inbounds begin
  x = view(u,1:7)   # x
  y = view(u,8:14)  # y
  v = view(u,15:21) # x′
  w = view(u,22:28) # y′
  du[1:7] .= v
  du[8:14].= w
  for i in 15:28
    du[i] = zero(u[1])
  end
  for i=1:7,j=1:7
    if i != j
      r = ((x[i]-x[j])^2 + (y[i] - y[j])^2)^(3/2)
      du[14+i] += j*(x[j] - x[i])/r
      du[21+i] += j*(y[j] - y[i])/r
    end
  end
  end
end

# DynamicalOdeProblem

function dotv(dv,q,v,par,t)
@inbounds begin
  x = view(q,1:7)   # x
  y = view(q,8:14)  # y
  vx = view(v,1:7)   # x′
  vy = view(v,8:14)  # y′
  for i in 1:14
    dv[i] = zero(x[1])
  end
  for i=1:7,j=1:7
    if i != j
      r = ((x[i]-x[j])^2 + (y[i] - y[j])^2)^(3/2)
      dv[i] += j*(x[j] - x[i])/r
      dv[7+i] += j*(y[j] - y[i])/r
    end
  end
  end
end

function dotq(dq,q,v,par,t)
@inbounds begin
  x = view(q,1:7)   # x
  y = view(q,8:14)  # y
  vx = view(v,1:7)   # x′
  vy = view(v,8:14)  # y′
  dq[1:7] .= vx
  dq[8:14].= vy
  end
end

#
#  Second Order Problem
#

function f2nd!(ddu,du,u,p,t)
  @inbounds begin
  x = view(u,1:7)   # x
  y = view(u,8:14)  # y
  vx = view(du,1:7) # x′
  vy = view(du,8:14) # y′

  for i in 1:14
    ddu[i] = zero(x[1])
  end

  for i=1:7,j=1:7
    if i != j
      r = ((x[i]-x[j])^2 + (y[i] - y[j])^2)^(3/2)
      ddu[i] += j*(x[j] - x[i])/r
      ddu[7+i] += j*(y[j] - y[i])/r
    end
  end
  end
end


module PassiveCable

using Unitful: uconvert, NoUnits

export cable_normalize, cable, rallpack1

# Compute solution g(x, t) to
#
#     ∂²g/∂x² - g - ∂g/∂t = 0
#
# on [0, L] × [0,∞), subject to:
#
#     g(x, 0) = 0
#     ∂g/∂x (0, t) = 1
#     ∂g/∂x (L, t) = 0
#
# Parameters:
#     x, t, L:  as described above
#   tol:  absolute error tolerance in result
#
# Return:
#     g(x, t)
#
# TODO: verify correctness when L≠1

function cable_normalized(x::Float64, t::Float64, L::Float64; tol=1e-8)
    if t<=0
        return 0.0
    else
        ginf = -cosh(L-x)/sinh(L)
        sum = exp(-t/L)
        Ltol = L*tol

        for k = Iterators.countfrom(1)
            a = k*pi/L
            b = exp(-t*(1+a^2))

            sum += 2*b*cos(a*x)/(1+a^2)
            resid_ub = b/(a^3*t)

            if resid_ub<Ltol
                break
            end
        end
        return ginf+sum/L;
     end
end

# Compute solution f(x, t) to
#
#     λ²∂²f/∂x² - f - τ∂f/∂t = 0
#
# on [0, L] x [0, ∞), subject to:
#
#     f(x, 0) = V
#     ∂f/∂x (0, t) = I·r
#     ∂f/∂x (L, t) = 0
#
# where:
#
#     λ² = 1/(r·g)   length constant
#     τ  = r·c       time constant
#
# In the physical model, the parameters correspond to the following:
#
#     L:  length of cable
#     r:  linear axial resistivity
#     g:  linear membrane conductivity
#     c:  linear membrane capacitance
#     V:  membrane reversal potential
#     I:  injected axial current on the left end (x = 0) of the cable.
#
# Note that r, g and c are specific 1-d quantities that differ from
# the cable resistivity r_L, specific membrane conductivity ḡ and
# specific membrane capacitance c_m as used elsewhere. If the
# cross-sectional area is A and cable circumference is f, then
# these quantities are related by:
#
#     r = r_L/A
#     g = ḡ·f
#     c = c_m·f
#
# Parameters:
#     x:  displacement along cable
#     t:  time
#     L, lambda, tau, r, V, I:  as described above
#   tol:  absolute error tolerance in result
#
# Return:
#     computed potential at (x,t) on cable.

function cable(x, t, L, lambda, tau, r, V, I; tol=1e-8)
    scale = I*r*lambda;
    if scale == 0
        return V
    else
        tol_n = abs(tol/scale)
        return scale*cable_normalized(uconvert(NoUnits, x/lambda), uconvert(NoUnits, t/tau), uconvert(NoUnits, L/lambda), tol=tol_n) + V
    end
end

# Rallpack 1 test
#
# One sided cable model with the following parameters:
#
#     RA = 1 Ω·m    bulk axial resistivity
#     RM = 4 Ω·m²   areal membrane resistivity
#     CM = 0.01 F/m²  areal membrane capacitance
#     d  = 1 µm     cable diameter
#     EM = -65 mV   reversal potential
#     I  = 0.1 nA   injected current
#     L  = 1 mm     cable length.
#
# (This notation aligns with that used in the Rallpacks paper.)
#
# Note that the injected current as described in the Rallpack model
# is trans-membrane, not axial. Consequently we need to swap the
# sign on I when passing to the cable function.
#
# Parameters:
#     x:  displacement along cable [m]
#     t:  time [s]
#   tol:  absolute tolerance for reported potential.
#
# Return:
#     computed potential at (x,t) on cable.

function rallpack1(x, t; tol=1e-8)
    RA = 1
    RM = 4
    CM = 1e-2
    d  = 1e-6
    EM = -65e-3
    I  = 0.1e-9
    L  = 1e-3

    r = 4*RA/(pi*d*d)
    lambda = sqrt(d/4 * RM/RA)
    tau = CM*RM

    return cable(x, t, L, lambda, tau, r, EM, -I, tol=tol)
end

end # module PassiveCable


using ViewerGL, LinearAlgebraicRepresentation
GL = ViewerGL
Lar = LinearAlgebraicRepresentation

# input of primitive shape
V, (VV,EV,FV) = Lar.cuboid([1,1], true)
square = (V,EV)
model = Lar.Struct([ square,
			Lar.r(π), square ])
V,EV = Lar.struct2lar(model)
VV = [[k] for k=1:size(V,2)]

GL.VIEW( GL.numbering(.4)((V,[VV, EV]),GL.COLORS[1]) );

# arrangement of input data
W = convert(Lar.Points, V')
cop_EV = Lar.coboundary_0(EV)
V, copEV, copFE = Lar.planar_arrangement(W, cop_EV)

# compute containment graph of components
bicon_comps = Lar.Arrangement.biconnected_components(copEV)

# compute euler characteristic
χ = Lar.euler_characteristic(V, copEV, copFE)
println("χ = $χ ; bicon_comps = $(length(bicon_comps))");

EW = Lar.cop2lar(copEV)
W = convert(Lar.Points, V')
comps = [ GL.GLLines(W,EW[comp],GL.COLORS[(k-1)%12+1]) for (k,comp) in enumerate(bicon_comps) ];
GL.VIEW(comps);

# final solid visualization
triangulated_faces = Lar.triangulate2D(V, [copEV, copFE])
V = convert(Lar.Points, V')
FVs = convert(Array{Lar.Cells}, triangulated_faces)
GL.VIEW(GL.GLExplode(V,FVs,1.2,1.2,1.2,99,1));

# polygonal face boundaries
EVs = Lar.FV2EVs(copEV, copFE)
EVs = convert(Array{Array{Array{Int64,1},1},1}, EVs)
GL.VIEW(GL.GLExplode(V,EVs,1.2,1.2,1.2,1,1));


function circHohmann(r₁,r₂,μ)
#[dV1,dV2,T] = circHohmann(r1,r2,μ)
#   calculates circular Hohmann transfer details given two radii
#   input and output are the same units
#
# Alex Hoffman
# 10/31/2020

aₜ   = 0.5(r₁ + r₂) #transfer semi-major axis
v1 = sqrt(μ/r₁)
v2 = sqrt(μ/r₂)
vt1 = sqrt(μ*(2/r₁-1/aₜ))
vt2 = sqrt(μ*(2/r₂-1/aₜ))
dV1 = vt1 - v1
dV2 = v2 - vt2
T = 2pi*sqrt(aₜ^3/μ)

return [dV1, dV2, T]
end


module GaussianProcesses

export k, K, GP_Regressor

using LinearAlgebra, Distributions

# exponential quadratic kernel
k(a::Array, b::Array; h::Real, d::Function) = exp(-d(a,b)^2 / (2 * h^2))
k(a::Real, b::Real; h::Real, d::Function) = k([a], [b]; h=h, d=d)

# kernel matrix
K(A::Vector, B::Vector; h::Real, d::Function) = [k(a, b; h = h, d = d) for a in A, b in B]

# prior distribution: mean 0, covariance K
mutable struct GP_Regressor
	"MVN distribution of prior"
	prior::Distributions.MvNormal
	"MVN distribution of posterior"
	posterior::Distributions.MvNormal
	"equivalent smoothing kernel weights (rows of matrix)"
	equiv_weights::Array
	"input points for sampling and plotting"
	x_grid::Vector
	"sample observations"
	x_train::Vector
	"sample outcomes"
	y_train::Vector
	"standard deviation of iid Gaussian noise in observations"
	σ::Real
	"characteristic length scale"
	h::Real
	"distance function"
	d::Function

	# inner constructor: call this function to instantiate
	function GP_Regressor(X_train::Vector, y_train::Vector, X::Vector; σ::Real, h::Real, d::Function)

		# prior
		Kpred = K(X, X; h=h, d=d)  # covariance
		prior = MvNormal(Kpred)

		# posterior terms
		Kinv = inv(K(X_train, X_train; h=h, d=d) + σ^2 * I)
		Kcross = K(X, X_train; h=h, d=d)
		Kequiv = Kcross * Kinv  # equivalent smoothing kernel weights (each row)

		mean_posterior = Kequiv * y_train
		cov_posterior = Kpred .- Kequiv * Kcross'

		posterior = MvNormal(mean_posterior, Matrix(Hermitian(cov_posterior)))

		new(prior, posterior, Kequiv, X, X_train, y_train, σ, h, d)
	end
end

end  # module


function rk4(f::Function, x₀::Float64, y₀::Float64, x₁::Float64, n)
    vx = Vector{Float64}(n + 1)
    vy = Vector{Float64}(n + 1)
    vx[1] = x = x₀
    vy[1] = y = y₀
    h = (x₁ - x₀) / n
    for i in 1:n
        k₁ = h * f(x, y)
        k₂ = h * f(x + 0.5h, y + 0.5k₁)
        k₃ = h * f(x + 0.5h, y + 0.5k₂)
        k₄ = h * f(x + h, y + k₃)
        vx[i + 1] = x = x₀ + i * h
        vy[i + 1] = y = y + (k₁ + 2k₂ + 2k₃ + k₄) / 6
    end
    return vx, vy
end

vx, vy = rk4(f, 0.0, 1.0, 10.0, 100)
for (x, y) in Iterators.take(zip(vx, vy), 10)
    @printf("%4.1f %10.5f %+12.4e\n", x, y, y - theoric(x))
end


# Solves the SVM classification problem using the barrier method.

# The main function.
# This implementation is deterministic for illustration purpose.
# In practice we may want to use random starting point.
function svmbarrier(X::AbstractMatrix,
                    Y::AbstractVector,
                    C::Real,
                    ϵ::Real;
                    t0::Real=1.,
                    μ::Real=100.,
                    α::Real=0.01, β::Real=0.05)
  w = zeros(size(X, 2))
  z = ones(Y)*2
  x0 = [w; z]
  A = [-Y.*X (-eye(size(Y, 1))); zeros(X) (-eye(size(Y, 1)))]
  b = [-1.*ones(Y); zeros(Y)]
  wsize = size(w, 1)
  f(x) = sum(x[1: wsize].^2)/2 + C*sum(x[wsize+1: end])
  ∇f(x) = [x[1: size(w, 1)]; C*ones(Y)]
  Hf(x) = Diagonal([ones(wsize); zeros(Y)])
  x, αdual, numstepsarray = barriermethod(
    f, ∇f, Hf, A, b, x0; ϵ=ϵ, t0=t0, μ=μ, α=α, β=β)
  x[1: wsize], αdual[1: size(Y, 1)], numstepsarray
end

# Implements the barrier method given the objective function and its
# first and second derivatives.
# The problem is supposed to have a linear constraint Ax ≤ b.
function barriermethod(f, ∇f, Hf,
                       A::AbstractMatrix,
                       b::AbstractVector,
                       x0::AbstractVector;
                       ϵ::Real=1e-3,
                       t0::Real=1.,
                       μ::Real=100.,
                       α::Real=0.01, β::Real=0.05)
  (all(A*x0.<b)
   || throw(DomainError("x0 must be strictly feasible.")))
  m = float(size(Y, 1))
  x, t = x0, t0
  numstepsarray::Vector{Int} = []
  while true
    x, numsteps = centeringstep(f, ∇f, Hf, A, b, x, t; α=α, β=β)
    push!(numstepsarray, numsteps)
    m/t < ϵ && break
    t *= μ
  end
  αdual = 1./(t*(b-A*x))
  x, αdual, numstepsarray
end

# Carries out a single centering step of the barrier method.
# The problem is supposed to have a linear constraint Ax ≤ b.
function centeringstep(f, ∇f, Hf,
                       A::AbstractMatrix,
                       b::AbstractVector,
                       x0::AbstractVector,
                       t::Real;
                       ϵ::Real=1e-8,
                       α::Real=0.01, β::Real=0.05)
  x = copy(x0)
  numsteps = 0
  obj(x) = t*f(x)-sum(log.(b-A*x))
  while true
    d = 1./(b-A*x)
    ∇ = t*∇f(x)+A'*d
    H = t*Hf(x)+A'*Diagonal(d)^2*A
    Δxnt = -H\∇
    λ² = -∇⋅Δxnt
    λ²/2 < ϵ && break
    tbacktrace = 1
    xtest = x + tbacktrace*Δxnt
    while(any(A*xtest .≥ b) || obj(xtest) > obj(x)-tbacktrace*α*λ²)
      tbacktrace *= β
      xtest = x + tbacktrace*Δxnt
    end
    # numeric limit
    obj(xtest) == obj(x) && break
    x = xtest
    numsteps += 1
  end
  x, numsteps
end


#############################################################################
#############################################################################
#
# This file implements polynomial addition
#
#############################################################################
#############################################################################

"""
Add two polynomials.
"""
function +(p1::Polynomial, p2::Polynomial)::Polynomial
    p1, p2 = deepcopy(p1), deepcopy(p2)
    p3 = Polynomial()
    while !iszero(p1) && !iszero(p2)
        t1, t2 = trailing(p1), trailing(p2)
        if t1.degree == t2.degree
            push!(p3, popfirst!(p1)+popfirst!(p2))
        elseif t1.degree < t2.degree
            push!(p3,popfirst!(p1))
        else
            push!(p3,popfirst!(p2))
        end
    end
    while !iszero(p1)
        push!(p3,popfirst!(p1))
    end
    while !iszero(p2)
        push!(p3,popfirst!(p2))
    end
    return p3
end

"""
Add a polynomial and a term.
"""
+(p::Polynomial, t::Term) = p + Polynomial(t)
+(t::Term, p::Polynomial) = p + t

"""
Add a polynomial and an integer.
"""
+(p::Polynomial, n::Int) = p + Term(n,0)
+(n::Int, p::Polynomial) = p + Term(n,0)


include("/home/jefter66/Workspace/lablib/src/calcs/Uncertainty.jl")

function I1()
    Δr = 0.01
    r1 = 0.0335
    r2 = 0.06

    m = 2.2429
    Δm = 0.0001

    I1 = (1/2)*m*(r1^2 + r2^2)
    ΔI1 = (1/2)*((Δm*(r1^2 + r2^2) + 2*m*(r1*Δr + r2*Δr)))
    println(" Ig ± Δ Ig = ", I1, " ± ", ΔI1)
#    return [I1,ΔI1]
end

function I2(ω1, ω)
   Δω=1e-3
   I2 = I1()[1]*( ω1/(ω1 - ω))
   ΔI2   =  (I1()[2]  + (I1()[2]*ω +  Δω*I1()[1])/(ω^2))
    println("I2 ± Δ I2 = ", I2, " ± ", ΔI2)
    return [I2, ΔI2]
end

function equivalencia(v1,v2,σ1, σ2)
    abs(v1-v2) < 2*(σ1 + σ2)
end

function ΔE(ω1, ω,I1, I2)
    Δω = 0.001
    Ef = (1/2)*(I1[1] + I2[1])*ω^2
    Ei = (1/2)*(I1[1])*ω1^2

    δE = ((1/2)*(I1[2] + I2[2])*ω^2 + (I1[1] + I2[1])*ω*Δω) + (1/2)*(I1[2]*ω1^2 + 2*ω1*Δω*I1[1])

    println("E ± ΔE = ", Ef - Ei, " ± ", δE)
end

function results()
ω1 = [14.444, 22.765,24.353]
ω   = [7.453,12.777,13.629]
print("(1) : ")
I2_1 = I2(ω1[1], ω[1])
print("(2) : ")
I2_2 = I2(ω1[2],ω[2])
print("(3) : ")
I2_3 = I2(ω1[3],ω[3])

println("==============================================")

println("(1) & (2) : ", equivalencia(I2_1[1], I2_1[2], I2_2[1] , I2_2[2]))
println("(1) & (3) : ", equivalencia(I2_1[1],I2_1[2], I2_3[1],I2_3[2]))
println("(2) & (3) : ", equivalencia(I2_2[1],I2_2[2], I2_3[1], I2_3[2]))

println("==========================")

E1 = ΔE(ω1[1], ω[1], I1(), I2_1)
E2 = ΔE(ω1[2], ω[2], I1(), I2_2)
E3 = ΔE(ω1[3], ω[3], I1(), I2_3)
#println("ΔE1 ± δ E1 = ", E1[1], " ± ", E1[2])
#println("ΔE2 ± δ E2 =",  E2[1], " ± ", E2[2])
#println("ΔE3 ± δ E3 = ", E3[1] , " ± ", E3[2])

end


g = [
  [08 02 22 97 38 15 00 40 00 75 04 05 07 78 52 12 50 77 91 08]
  [49 49 99 40 17 81 18 57 60 87 17 40 98 43 69 48 04 56 62 00]
  [81 49 31 73 55 79 14 29 93 71 40 67 53 88 30 03 49 13 36 65]
  [52 70 95 23 04 60 11 42 69 24 68 56 01 32 56 71 37 02 36 91]
  [22 31 16 71 51 67 63 89 41 92 36 54 22 40 40 28 66 33 13 80]
  [24 47 32 60 99 03 45 02 44 75 33 53 78 36 84 20 35 17 12 50]
  [32 98 81 28 64 23 67 10 26 38 40 67 59 54 70 66 18 38 64 70]
  [67 26 20 68 02 62 12 20 95 63 94 39 63 08 40 91 66 49 94 21]
  [24 55 58 05 66 73 99 26 97 17 78 78 96 83 14 88 34 89 63 72]
  [21 36 23 09 75 00 76 44 20 45 35 14 00 61 33 97 34 31 33 95]
  [78 17 53 28 22 75 31 67 15 94 03 80 04 62 16 14 09 53 56 92]
  [16 39 05 42 96 35 31 47 55 58 88 24 00 17 54 24 36 29 85 57]
  [86 56 00 48 35 71 89 07 05 44 44 37 44 60 21 58 51 54 17 58]
  [19 80 81 68 05 94 47 69 28 73 92 13 86 52 17 77 04 89 55 40]
  [04 52 08 83 97 35 99 16 07 97 57 32 16 26 26 79 33 27 98 66]
  [88 36 68 87 57 62 20 72 03 46 33 67 46 55 12 32 63 93 53 69]
  [04 42 16 73 38 25 39 11 24 94 72 18 08 46 29 32 40 62 76 36]
  [20 69 36 41 72 30 23 88 34 62 99 69 82 67 59 85 74 04 36 16]
  [20 73 35 29 78 31 90 01 74 31 49 71 48 86 81 16 23 57 05 54]
  [01 70 54 71 83 51 54 69 16 92 33 48 61 43 52 01 89 19 67 48]
]

l = 4
m, n = size(g)
maxsofar = 0
for i = 1:m-(l-1)
  global maxsofar
  for j = 1:n-(l-1)
    p_col = prod(g[i:i+(l-1), j])
    p_row = prod(g[i,j:j+(l-1)])
    p_dia1 = g[i,j] * g[i+1,j+1] * g[i+2,j+2] * g[i+3,j+3]
    p_dia2 = g[i,j+3] * g[i+1,j+2] * g[i+2,j+1] * g[i+3,j]
    maxsofar = max(maxsofar, p_col)
    maxsofar = max(maxsofar, p_row)
    maxsofar = max(maxsofar, p_dia1)
    maxsofar = max(maxsofar, p_dia2)
  end
end
println(maxsofar)


# Copyright (c) 2021 Idiap Research Institute, http://www.idiap.ch/
#  Niccolò Antonello <nantonel@idiap.ch>

export TropicalWeight

"""
`TropicalWeight(x)`

| Set                                 | ``\\oplus``|   ``\\otimes`` | ``\\bar{0}`` | ``\\bar{1}`` |
|:-----------------------------------:|:----------:|:--------------:|:------------:|:------------:|
|``\\mathbb{R}\\cup\\{\\pm\\infty\\}``|  ``\\min`` |     ``+``      |``\\infty``   |   ``0``      |
"""
struct TropicalWeight{T <: AbstractFloat} <: Semiring
  x::T
end

zero(::Type{TropicalWeight{T}}) where T = TropicalWeight{T}(T(Inf))
one(::Type{TropicalWeight{T}}) where T = TropicalWeight{T}(zero(T))

*(a::TropicalWeight{T}, b::TropicalWeight{T}) where {T <: AbstractFloat} = TropicalWeight{T}(a.x + b.x)
+(a::TropicalWeight{T}, b::TropicalWeight{T}) where {T <: AbstractFloat} = TropicalWeight{T}( min(a.x,b.x) )
/(a::TropicalWeight{T}, b::TropicalWeight{T}) where {T <: AbstractFloat} = TropicalWeight{T}( a.x-b.x )

# parsing
parse(::Type{S},str) where {T, S <: TropicalWeight{T}} = S(parse(T,str))

#properties
isidempotent(::Type{W}) where {W <: TropicalWeight} = true
iscommulative(::Type{W}) where {W <: TropicalWeight} = true
isleft(::Type{W}) where {W <: TropicalWeight}= true
isright(::Type{W}) where {W <: TropicalWeight}= true
isweaklydivisible(::Type{W}) where {W <: TropicalWeight}= true
iscomplete(::Type{W}) where {W <: TropicalWeight}= true
ispath(::Type{W}) where {W <: TropicalWeight}= true


"""
    bresenham(start::Tuple{<:T, <:T}, stop::Tuple{<:T, <:T}) where T<:Number

Implements the [Bresenham's line algorithm](https://en.wikipedia.org/wiki/Bresenham%27s_line_algorithm).

# Examples

```jldoctest
julia> bresenham()

```
"""
function bresenham(start::Tuple{T, T}, stop::Tuple{T, T}) where T<:Integer
    x0, y0 = start
    x1, y1 = stop
    if abs(y1 - y0) < abs(x1 - x0)
        if x0 > x1
            reverse!(low_line(stop, start))
        else
            low_line(start, stop)
        end
    else
        if y0 > y1
            reverse!(high_line(stop, start))
        else
            high_line(start, stop)
        end
    end
end

"""
    low_line(start::Tuple{T, T}, stop::Tuple{T, T}) where T<:Integer

# Examples

```jldoctest
julia> low_line()

```
"""
function low_line(start::Tuple{T, T}, stop::Tuple{T, T}) where T<:Integer
    x0, y0 = start
    x1, y1 = stop
    line = fill((zero(T), zero(T)), length(x0:x1))
    dx = x1 - x0
    dy = y1 - y0
    yi = 1
    if dy < 0
        yi = -1
        dy = -dy
    end
    D = 2*dy - dx
    y = y0

    for x in x0:x1
        line[x - x0 + 1] =  (x, y)
        if D > 0
            y = y + yi
            D = D - 2*dx
        end
        D = D + 2*dy
    end

    line
end

"""
    high_line(start::Tuple{T, T}, stop::Tuple{T, T}) where T<:Integer

# Examples

```jldoctest
julia> high_line()

```
"""
function high_line(start::Tuple{T, T}, stop::Tuple{T, T}) where T<:Integer
    x0, y0 = start
    x1, y1 = stop
    line = fill((zero(T), zero(T)), length(y0:y1))
    dx = x1 - x0
    dy = y1 - y0
    xi = 1
    if dx < 0
        xi = -1
        dx = -dx
    end
    D = 2*dx - dy
    x = x0

    for y in y0:y1
        line[y - y0 + 1] =  (x, y)
        if D > 0
            x = x + xi
            D = D - 2*dy
        end
        D = D + 2*dx
    end

    line
end


#=
    Cheating model.

    From
    https://discourse.julialang.org/t/fitting-a-observed-value-to-a-binomial-distribution-turing/66619
    """
    Fitting a observed value to a Binomial Distribution Turing

    I am new to Turing and trying to learn it by trying to replicate the Chapters from the book
    https://github.com/CamDavidsonPilon/Probabilistic-Programming-and-Bayesian-Methods-for-Hackers 2 .

    Here is my problem (from Chapter 2) :

    1.) Have a probability p ~ Uniform(0,1) which I have to infer the “number_of_cheaters”
    2.) p_skewed = p*(0.5) + (0.25) (deterministic value)
    3.) with model:
    yes_responses = pm.Binomial(“number_of_cheaters”, 100, p_skewed, observed=35)

    How do I write this in Turing ?

    ...

    @model function model(yes_responses,N)
        p ~ Uniform(0,1)
        p_skewed = p*(0.5)+0.25
        rv_yes_responses ~ Binomial(N,p_skewed)

        // This line on how to include yes_responses i am unable to figure out

        return p_skewed

    end
   """

   The model refered to is in
   https://github.com/CamDavidsonPilon/Probabilistic-Programming-and-Bayesian-Methods-for-Hackers/blob/master/Chapter2_MorePyMC/Ch2_MorePyMC_PyMC3.ipynb

    Output of this Turing.jl model:
        Summary Statistics
    parameters      mean       std   naive_se      mcse        ess      rhat   ess_per_sec
        Symbol   Float64   Float64    Float64   Float64    Float64   Float64       Float64

               p    0.2042    0.0993     0.0031    0.0040   446.4728    0.9990      394.4106
        p_skewed    0.3521    0.0497     0.0016    0.0020   446.4728    0.9990      394.4106

    Quantiles
    parameters      2.5%     25.0%     50.0%     75.0%     97.5%
        Symbol   Float64   Float64   Float64   Float64   Float64

               p    0.0349    0.1325    0.2020    0.2749    0.3981
        p_skewed    0.2674    0.3162    0.3510    0.3875    0.4490

    Summary Stats:
    Length:         1000
    Missing Count:  0
    Mean:           0.352105
    Minimum:        0.250395
    1st Quartile:   0.316228
    Median:         0.350993
    3rd Quartile:   0.387453
    Maximum:        0.494718

=#

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

@model function cheating_model2(yes_responses,N)
    p ~ Uniform(0,1)
    # Using Dirac(.) makes p_skewed in the chain
    p_skewed ~ Dirac(p*(0.5)+0.25)
    # p_skewed = p*(0.5)+0.25
    yes_responses ~ Binomial(N,p_skewed)

    # An alternative is to return the variable (see below how to handle this)
    return p_skewed
end

yes_response = 35
N = 100
model = cheating_model2(yes_response,N)

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 10_000)
chns = sample(model, PG(5),  1_000)
# chns = sample(model, PG(5),  MCMCThreads(), 10_000, 4)
# 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)
# display(plot(chns))

chains_params = Turing.MCMCChains.get_sections(chns, :parameters)
genq = generated_quantities(model, chains_params)
display(summarystats(vcat(genq...)))


#    Problem : GROUP A
#    *********
#	 The extended Woods problem.
#	 This problem is a sum of n/4 sets of 6 terms, each of which is
#    assigned its own group.  For a given set i, the groups are
#    A(i), B(i), C(i), D(i), E(i) and F(i). Groups A(i) and C(i) contain 1
#    nonlinear element each, denoted Y(i) and Z(i).
#
#    The problem dimension is defined from the number of these sets.
#    The number of problem variables is then 4 times larger.
#
#	 This version uses a slightly unorthodox expression of Woods
#    function as a sum of squares (see Buckley)
#
#    Origonal SIF Source: problem 14 in
#    J.J. More', B.S. Garbow and K.E. Hillstrom,
#    "Testing Unconstrained Optimization Software",
#    ACM Transactions on Mathematical Software, vol. 7(1), pp. 17-41, 1981.
#
#    See also Toint#27, Buckley#17 (p. 101), Conn, Gould, Toint#7
#
#    WOODS.SIF classification SUR2-AN-V-0
#
#	 This problem is decomposed in n linear groups, the last n-1 of which
#    are 2 x 2 and singular.
#
#    NS is the number of sets (= n/4)
#
# Daniel Henderson, 08/2021

f = (x) -> begin
	fx = 0.0
	for i in 1:Int(lastindex(x)/4)
		fx += (100*(x[4*i-2]-x[4*i-3]^2)^2 + (1-x[4*i-3])^2 + 90*(x[4*i]-x[4*i-1]^2)^2 + (1-x[4*i-1])^2 + 10*(x[4*i-2]+x[4*i]-2)^2 + 0.1*(x[4*i-2]-x[4*i])^2)
	end
    return fx
end

g! = (g, x) -> begin
	for i in 1:Int(lastindex(x)/4)
		g[4i-3] = -2(1-x[4i-3]) - 400*x[4i-3]*(x[4i-2]-x[4i-3]^2)
		g[4i-2] = 200(x[4i-2]-x[4i-3]^2) + 0.2(x[4i-2]-x[4i]) + 20(x[4i-2] + x[4i]-2)
		g[4i-1] = -2(1-x[4i-1]) - 360*x[4i-1]*(x[4i]-x[4i-1]^2)
		g[4i] = -0.2(x[4i-2]-x[4i]) + 20(x[4i-2] + x[4i]-2) + 180(x[4i]-x[4i-1]^2)
	end
    return g
end

fg! = (g, x) -> begin
	fx = 0.0
	for i in 1:Int(lastindex(x)/4)
		fx += (100*(x[4*i-2]-x[4*i-3]^2)^2 + (1-x[4*i-3])^2 + 90*(x[4*i]-x[4*i-1]^2)^2 + (1-x[4*i-1])^2 + 10*(x[4*i-2]+x[4*i]-2)^2 + 0.1*(x[4*i-2]-x[4*i])^2)
		g[4i-3] = -2(1-x[4i-3]) - 400*x[4i-3]*(x[4i-2]-x[4i-3]^2)
		g[4i-2] = 200(x[4i-2]-x[4i-3]^2) + 0.2(x[4i-2]-x[4i]) + 20(x[4i-2] + x[4i]-2)
		g[4i-1] = -2(1-x[4i-1]) - 360*x[4i-1]*(x[4i]-x[4i-1]^2)
		g[4i] = -0.2(x[4i-2]-x[4i]) + 20(x[4i-2] + x[4i]-2) + 180(x[4i]-x[4i-1]^2)
	end
    return fx, g
end

init = (n::Int=4000) -> begin
	mod(n, 4) > 0 && @warn "WOODS: number of variables must be divisible by 4"
	q = max(1, div(n, 4))
	n = 4*q

	x0 = [j % 2 == 1 ? -3.0 : -1.0 for j in 1:n]
    return n, x0
end

TestSet["WOODS"] = UncProgram("WOODS", f, g!, fg!, init)


using OrdinaryDiffEq, ParameterizedFunctions, ForwardDiff

f = @ode_def begin
  dx = a*x - b*x*y
  dy = -c*y + x*y
end a b c

p = [1.5,1.0,3.0]
prob = ODEProblem(f,[1.0;1.0],(0.0,10.0),p)
t = 0.0:0.5:10.0

function G(p)
  tmp_prob = remake(prob,u0=convert.(eltype(p),prob.u0),p=p)
  sol = solve(tmp_prob,Vern9(),abstol=1e-14,reltol=1e-14,saveat=t)
  A = convert(Array,sol)
  sum(((1 .- A).^2)./2)
end
# G([1.5,1.0,3.0])
# res2 = ForwardDiff.gradient(G,[1.5,1.0,3.0])

@info("Running forwarddiff")
# G([1.5,1.0,3.0])
es2 = ForwardDiff.gradient(G,[1.5,1.0,3.0])
# using Zygote
# Zygote.gradient(G,[1.5,1.0,3.0])


using Revise
using DrWatson
@quickactivate "SICMModelingToolkit"
DrWatson.greet()
using ModelingToolkit, Unitful, DifferentialEquations, Optim
using Latexify
using BenchmarkTools, LinearAlgebra, ForwardDiff, Zygote
using Quadrature, StaticArrays
using SICMModelingToolkit
using Plots
using Measures

#### section 1.4

@parameters t m k
D = Differential(t)

locals1 = generate_generic_localtuple(1)
locals2 = generate_generic_localtuple(2)
locals3 = generate_generic_localtuple(3)

velocity(locals3)

L_free_mass = L_free_particle(m)

Lagrange_equations(L_free_mass, locals3)

L_harmonic_spring_apply = L_harmonic_spring(m, k)
Lagrange_equations(L_harmonic_spring_apply, locals1)
Lagrange_equations(L_harmonic_spring_apply, locals3)

ode_generator = ode_problem_generator(L_free_particle, locals3, [m])
#Lagrangian_pre_param, locals, params = L_free_particle, locals3, [m]

#ode_func = ODEFunction(ode_generator[1])

#du = [0., 0., 0., 0., 0., 0.]
#u0 = @SVector [1., 0.5, 0.25, 0., 0., 0.]
#p0 = @SVector [1.0]
#ode_func(du, u0, p0, 0.0)

#@show du

#dyn_ode_func(du, u, p, t) = DynamicalODEFunction{true}(ode_func)

du = [0., 0., 0.]
q0 = [0., 0., 0.]
qv0 = [1.0, 0.5, 0.25]
#u0 = [1., 0.5, 0.25, 0., 0., 0.]
u0 = ArrayPartition(qv0, q0)
p0 = [1.0]
t0 = 0.0

tspan = (0.0, 2.0)

prob = ode_generator(qv0, q0, tspan, p0; saveat=0.05)

sol = solve(prob, McAte5(); dt=0.05)

plot(sol, vars=(4,5,6); dpi=350)

### kepler

kepler_ode_generator = ode_problem_generator(L_kepler, locals2, [m, k])

begin
    q0_kep = [0.0, 1.0]
    qv0_kep = [-1.0, 0.0]
	tspan_kep = (0.0, 5.0)
    p_kep = [3.0, 10.0]
    dt = 0.01

    kepler_prob = kepler_ode_generator(qv0_kep, q0_kep, tspan_kep, p_kep)
end

kepler_sol = solve(kepler_prob, McAte5(); dt=dt)
kepler_sol_nosymp = solve(kepler_prob, Tsit5(); dt=dt)

curxlims = (-1.1, 1.1)
curylims = (9.0/16) .* curxlims
cur_left_margin = Measures.Length{:mm, Float64}(12)
cur_bottom_margin = Measures.Length{:mm, Float64}(7)
gr()
ellipseplot = plot(kepler_sol, vars=(4,3); dpi=500, ylims=curylims, xlims=curxlims,
    label="orbiter", xlabel="x", ylabel="y", title="Kepler Orbital Problem Symplectic",
    left_margin=cur_left_margin, bottom_margin=cur_bottom_margin)
scatter!(ellipseplot, [0.0], [0.0]; label="planet")
save("kepler_symplectic.png", ellipseplot)
ellipseplot_nosymp = plot(kepler_sol_nosymp, vars=(4,3); dpi=500, ylims=curylims, xlims=curxlims,
    label="orbiter", xlabel="x", ylabel="y", title="Kepler Orbital Problem Nonsymplectic",
    left_margin=cur_left_margin, bottom_margin=cur_bottom_margin)
scatter!(ellipseplot_nosymp, [0.0], [0.0]; label="planet")
save("kepler_nosymplectic.png", ellipseplot_nosymp)
plot(kepler_sol, vars=(0,3,4); dpi=350)

#### harmonic

harmonic_ode_generator = ode_problem_generator(L_harmonic_spring, locals1, [m, k])

begin
    q0_harm = [1.0]
    qv0_harm = [0.0]
	tspan_harm = (0.0, 5.0)
    p_harm = [1.0, 1.0]
    dt = 0.01

    harmonic_prob = harmonic_ode_generator(qv0_harm, q0_harm, tspan_harm, p_harm)
end

harmonic_sol = solve(harmonic_prob, McAte5(); dt=dt)
plot(harmonic_sol; dpi=500, left_margin=cur_left_margin, bottom_margin=cur_bottom_margin)


# A Julia-native gamma function. Pretty darn accurate and fast, and using this
# instead of SpecialFunctions.gamma would mean that the dependency can be
# removed. That doesn't seem worth it to me at this point in time, but putting
# this here in case.

const _G  = 607/128
const _C1 = 0.99999999999999709182
const _C  = (57.156235665862923517,
            -59.597960355475491248,
             14.136097974741747174,
            -0.49191381609762019978,
            .33994649984811888699e-4,
            .46523628927048575665e-4,
           -.98374475304879564677e-4,
            .15808870322491248884e-3,
           -.21026444172410488319e-3,
            .21743961811521264320e-3,
           -.16431810653676389022e-3,
            .84418223983852743293e-4,
           -.26190838401581408670e-4,
            .36899182659531622704e-5)
const _Cl = length(_C)
const SQRT_2_PI = sqrt(2*pi)

# Lanczos approximation.
@inline function _gamma_pos(_z::T) where{T}
  T_one    = one(T)
  T_half   = T_one/(T_one+T_one)
  z_half   = _z-T_half
  z_g_half = z_half+_G
  out = 0.99999999999999709182
  zpj = _z
  @inbounds for c in _C
    out += c/zpj
    zpj += T_one
  end
  out*SQRT_2_PI*exp(z_half*log(z_g_half) - z_g_half)
end

function _gamma(x)
  iszero(x) && return Inf
  (isinteger(x) && x < zero(x)) && return NaN
  x > zero(x) && return _gamma_pos(x)
  flx = floor(x)
  xf  = x - flx
  g   = _gamma_pos(xf)
  _x  = xf-one(x)
  for _ in 1:Int(-flx)
    g  /= _x
    _x -= one(x)
  end
  g
end


# Hello!
# This is the second Project Euler question where I find the sum of the
# even-valued terms of the Fibonacci Sequence under 4 million!

# I am using Binet's Formula to calculate the terms in the Fibonacci Sequence!

# The Golden Ratio!
φ = (1 + sqrt(5))/2

# The Binet's Formula is :
# the golden ratio (~ 1.618) raised to the nth power and divide it by sqrt of 5
# and rounding the answer to the nearest integer. The error approaches 0 as you
# approach the infinitely big terms!

fibonacci_term = 0
a = (φ ^ fibonacci_term)/sqrt(5)
b = round(Int, a)

# I can calculate many Fibonacci Numbers at once!
# Binet's Formula in ACTION!!!

c = 100
d = 1 : c
e = φ .^ d
f = e/sqrt(5)
fibonacci_numbers = round.(f)

# I can just find the even-valued Fibonacci terms!
# Every third Fibonacci term is even!

g = 100
h = 1 : g
# 3 * h is too skip 2 terms and go to the next even-valued Fibonacci Number.
i = φ .^ (3 * h)
j = i / sqrt(5)
even_valued_fibonacci_numbers = round.(j)

# I control the size of the Fibonacci Numbers.
# I am just setting the domain for the Fibonacci Numbers because I do not want
# them to be too big (The even-valued Fibonacci terms cannot be greater than
# 4 million).

k = 1
sum = 0
while even_valued_fibonacci_numbers[k] < 4000000
	k = k + 1
	sum = sum + even_valued_fibonacci_numbers[k - 1]
end
println(sum)


using LinearAlgebraicRepresentation
Lar = LinearAlgebraicRepresentation
using Plasm

VOID = [[]], [[1]]

# example 1

V = [[0,0] [1,0] [2,0] [0,1] [1,1] [2,1] [0,2] [1,2] [2,2]];
FV = [[1,2,4],[2,3,5],[3,5,6],[4,5,7],[5,7,8],[6,8,9]];
pattern = repeat([1,2,-3],outer=4);
model = (V,FV);
W,FW = Lar.extrudeSimplicial(model, pattern);
Plasm.view(W,FW)

# example 2

model = Lar.extrudeSimplicial( VOID, ones(10) )
Plasm.view(model)
model = Lar.extrudeSimplicial( model, ones(10) )
Plasm.view(model)
model = Lar.extrudeSimplicial( model, ones(10) )
Plasm.view(model)

# example 3

model = Lar.extrudeSimplicial( VOID, repeat([1,-1],outer=10) )
Plasm.view(model)
model = Lar.extrudeSimplicial( model, repeat([1,-1],outer=10) )
Plasm.view(model)

# example 4

grid_2d = Lar.simplexGrid([3,3])
Plasm.view(grid_2d)
grid_3d = Lar.simplexGrid([2,3,4])
Plasm.view(grid_3d)
V,CV = Lar.simplexGrid([1,1,1])
Plasm.view(V,CV)

# example 5

SK2 = Lar.simplexFacets(CV)
Plasm.view(V, SK2)
SK1 = Lar.simplexFacets(SK2)
Plasm.view(V, SK1)


using Pkg
Pkg.activate(pwd())

# # Constrained Optimization
# ## Numerical method

using Plots, Random

function create_anim(
    f,
    path,
    xlims,
    ylims,
    file_name = joinpath(pwd(), randstring(12) * ".gif");
    xbounds = xlims,
    ybounds = ylims,
    fps = 15,
)
    xs = range(xlims...; length = 100)
    ys = range(ylims...; length = 100)
    plt = contourf(xs, ys, f; color = :jet)

    # add constraints if provided
    if !(xbounds == xlims && ybounds == ylims)
        x_rect = [xbounds[1]; xbounds[2]; xbounds[2]; xbounds[1]; xbounds[1]]
        y_rect = [ybounds[1]; ybounds[1]; ybounds[2]; ybounds[2]; ybounds[1]]

        plot!(x_rect, y_rect; line = (2, :dash, :red), label="")
    end

    # add an empty plot
    plot!(Float64[], Float64[]; line = (4, :arrow, :black), label = "")

    # extract the last plot series
    plt_path = plt.series_list[end]

    # create the animation and save it
    anim = Animation()
    for x in eachcol(path)
        push!(plt_path, x[1], x[2]) # add a new point
        frame(anim)
    end
    gif(anim, file_name; fps = fps, show_msg = false)
    return nothing
end

f(x) = sin(x[1] + x[2]) + cos(x[1])^2
f(x1,x2) = f([x1;x2])
g(x) = [cos(x[1] + x[2]) - 2*cos(x[1])*sin(x[1]); cos(x[1] + x[2])]

#+

function optim(f, g, P, x, α; max_iter=100)
    xs = zeros(length(x), max_iter+1)
    ys = zeros(length(x), max_iter)
    xs[:,1] = x
    for i in 1:max_iter
        ys[:,i] = xs[:,i] - α*g(xs[:,i])
        xs[:,i+1] = P(ys[:,i])
    end
    return xs, ys
end

#+

P(x, x_min, x_max) = min.(max.(x, x_min), x_max)

x_min = [-1; -1]
x_max = [0; 0]

xs, ys = optim(f, g, x -> P(x,x_min,x_max), [0;-1], 0.1)

xlims = (-3, 1)
ylims = (-2, 1)

#+

create_anim(f, xs, xlims, ylims, joinpath(pwd(), "lecture_08", "anim6.gif");
    xbounds=(x_min[1], x_max[1]),
    ybounds=(x_min[2], x_max[2]),
)

#+

xys = hcat(reshape([xs[:,1:end-1]; ys][:], 2, :), xs[:,end])

create_anim(f, xys, xlims, ylims, joinpath(pwd(), "lecture_08", "anim7.gif");
    xbounds=(x_min[1], x_max[1]),
    ybounds=(x_min[2], x_max[2]),
)


module NumericError

import Base: convert, promote_rule, +, -, *, /, ^

export Measure

type Measure <: Number
  x::Float64
  σ::Float64
  Measure(x::Real) = new(Float64(x), 0)
  Measure(x::Real, σ::Real) = new(Float64(x), σ)
end

Base.show(io::IO, x::Measure) = print(io, string(x.x, " ± ", x.σ))

Base.convert(::Type{Measure}, x::Real) = Measure(Float64(x), 0.0)
Base.promote_rule(::Type{Float64}, ::Type{Measure}) = Measure
Base.promote_rule(::Type{Int64}, ::Type{Measure}) = Measure

+(a::Measure, b::Measure) = Measure(a.x + b.x, sqrt(a.σ ^ 2 + b.σ ^ 2))
-(a::Measure, b::Measure) = Measure(a.x - b.x, sqrt(a.σ ^ 2 + b.σ ^ 2))
-(a::Measure) = Measure(-a.x, a.σ)

*(a::Measure, b::Measure) = begin
  x = a.x * b.x
  σ = sqrt(x ^ 2 * ((a.σ / a.x) ^ 2 + (b.σ / b.x) ^ 2))
  Measure(x, σ)
end
/(a::Measure, b::Measure) = begin
  x = a.x / b.x
  σ = sqrt(x ^ 2 * ((a.σ / a.x) ^ 2 + (b.σ / b.x) ^ 2))
  Measure(x, σ)
end

^(a::Measure, b::Float64) = begin
  x = a.x ^ b
  σ = abs(x * b * a.σ / a.x)
  Measure(x, σ)
end

Base.sqrt(a::Measure) = a ^ .5

end  # module NumericError

using NumericError

# x1 = 100 ± 1.1
# y1 = 50 ± 1.2
# x2 = 200 ± 2.2
# y2 = 100 ± 2.3

x1 = Measure(100, 1.1)
y1 = Measure(50,  1.2)
x2 = Measure(200, 2.2)
y2 = Measure(100, 2.3)

d = sqrt((x1 - x2) ^ 2 + (y1 - y2) ^ 2)

@show x1 y1 x2 y2 sqrt((x1 - x2) ^ 2 + (y1 - y2) ^ 2)


using QuantEcon.compute_fixed_point, QuantEcon.DiscreteRV, QuantEcon.draw, QuantEcon.LinInterp

"""
For a given probability return expected loss of choosing model 0
"""
expect_loss_choose_0(p::Real, L0::Real) = (1 - p) * L0

"""
For a given probability return expected loss of choosing model 1
"""
expect_loss_choose_1(p::Real, L1::Real) = p * L1

"""
We will need to be able to evaluate the expectation of our Bellman
equation J. In order to do this, we need the current probability
that model 0 is correct (p), the distributions (f0, f1), and a
function that can evaluate the Bellman equation
"""
function EJ(p::Real, f0::AbstractVector, f1::AbstractVector, J::LinInterp)
    # Get the current distribution we believe (p * f0 + (1 - p) * f1)
    curr_dist = p * f0 + (1 - p) * f1

    # Get tomorrow's expected distribution through Bayes law
    tp1_dist = clamp.((p * f0) ./ (p * f0 + (1 - p) * f1), 0, 1)

    # Evaluate the expectation
    EJ = dot(curr_dist, J.(tp1_dist))

    return EJ
end

expect_loss_cont(p::Real, c::Real,
                 f0::AbstractVector, f1::AbstractVector, J::LinInterp) =
    c + EJ(p, f0, f1, J)

"""
Evaluates the value function for a given continuation value
function; that is, evaluates

    J(p) = min(pL0, (1-p)L1, c + E[J(p')])

Uses linear interpolation between points
"""
function bellman_operator(pgrid::AbstractVector,
                          c::Real,
                          f0::AbstractVector,
                          f1::AbstractVector,
                          L0::Real,
                          L1::Real,
                          J::AbstractVector)
    m = length(pgrid)
    @assert m == length(J)

    J_out = zeros(m)
    J_interp = LinInterp(pgrid, J)

    for (p_ind, p) in enumerate(pgrid)
        # Payoff of choosing model 0
        p_c_0 = expect_loss_choose_0(p, L0)
        p_c_1 = expect_loss_choose_1(p, L1)
        p_con = expect_loss_cont(p, c, f0, f1, J_interp)

        J_out[p_ind] = min(p_c_0, p_c_1, p_con)
    end

    return J_out
end

# Create two distributions over 50 values for k
# We are using a discretized beta distribution

p_m1 = linspace(0, 1, 50)
f0 = clamp.(pdf.(Beta(1, 1), p_m1), 1e-8, Inf)
f0 = f0 / sum(f0)
f1 = clamp.(pdf.(Beta(9, 9), p_m1), 1e-8, Inf)
f1 = f1 / sum(f1)

# To solve
pg = linspace(0, 1, 251)
J1 = compute_fixed_point(x -> bellman_operator(pg, 0.5, f0, f1, 5.0, 5.0, x),
    zeros(length(pg)), err_tol=1e-6, print_skip=5);


function log_taylor(x::Number)
    if x <= 0
        return undef
    end

    e = float(Base.MathConstants.e)
    lb = 0.546
    ub = lb * e

    res = 0

    while x < lb
        x *= e
        res -= 1
    end

    while x > ub
        x /= e
        res += 1
    end

    s = 1

    for i in 1 : 40
        res += s * (x - 1)^i / i
        s *= -1
    end

    return res
end


using DrWatson
@quickactivate "StatReth"

# %%
using CSV
using DataFrames
using Turing
using Plots

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

# %%
# This is from the chapter
d = DataFrame(CSV.File(datadir("exp_raw/Howell_1.csv")))
sort!(d, :weight)  # to make the plots look nice

z_weight, unz_weight = zscore_transform(d.weight)
d.weight_s = z_weight(d.weight)

f_parabola(weight_s, a, b1, b2) = a + b1 * weight_s + b2 * weight_s^2

@model function parabola(weight_s, height)
    a ~ Normal(178, 20)
    b1 ~ LogNormal(0, 1)
    b2 ~ Normal(0, 1)
    σ ~ Uniform(0, 50)
    μ = f_parabola.(weight_s, a, b1, b2)
    height ~ MvNormal(μ, σ)
end

mparab = parabola(d.weight_s, d.height)
prior = sample(mparab, Prior(), 1_000) |> DataFrame

# %%
plot(legend = false)
for c in eachrow(prior[1:100, :])
    p = f_parabola.(d.weight_s, c.a, c.b1, c.b2)
    plot!(d.weight_s, p, color = :black, alpha = 0.2)
end
plot!()

# %%
"""
Thinking about the cbrt explanation, b2 should be negative so I get a right turn here as
well. Since I don't want the extremum to be in the range of my data I need the maximum to
be at like 3 or 4 or something like that. Over the range of ~4 stddev I want the height to
change 100 cm or so.
f(x) = b₂ * x² + b₁ * x + a
f'(x) = 2b₂ * x + b₁

f(2) - f(-2) = 100
f'(3) = 0
b₂ < 0

f(2) - f(-2) = (b₂ * 4 + b₁ * 2 + a) - (b₂ * 4 + b₁ * -2 + a)
             = 4b₁ = 100
            =>  b₁ =  25
f'(3) = 2b₂ * 3 + b₁ = 0 => -b₁ = 6b₂
     => b₂ = -25 / 6 ≈ -4

I'll move a by hand until it seams reasonable.
"""

@model function parabola(weight_s, height)
    a ~ Normal(130, 20)
    b1 ~ Normal(25, 5)
    b2 ~ truncated(Normal(-4, 2), -Inf, 0)
    σ ~ Exponential(5)
    μ = f_parabola.(weight_s, a, b1, b2)
    height ~ MvNormal(μ, σ)
end

mparab = parabola(d.weight_s, d.height)
prior = sample(mparab, Prior(), 100) |> DataFrame

plot(legend = false)
for c in eachrow(prior)
    p = f_parabola.(d.weight_s, c.a, c.b1, c.b2)
    plot!(d.weight_s, p, color = :black, alpha = 0.2)
end
plot!(ylims = (0, 200))


"""
    generate_conditioned_dist(covar::CovarianceAtDate, conditioning_draws::Array{Symbol,T}) where T<:Real
Given some subset of known stochastic integral values this generates a conditional multivariate normal distribution.
"""
function generate_conditioned_distribution(covar::CovarianceAtDate, conditioning_draws::Dict{Symbol,T}) where T<:Real
    # https://stats.stackexchange.com/questions/30588/deriving-the-conditional-distributions-of-a-multivariate-normal-distribution
    what_conditioned_on = collect(keys(conditioning_draws))
    len = length(what_conditioned_on)
    if len < 1 error("Nothing that was input was in the covariance matrix. So nothing to condition on.") end
    conditioning_indices = Array{typeof(Integer(1)),1}()
    for i in 1:len
        append!(conditioning_indices, findall(what_conditioned_on[i] .==  covar.covariance_labels_))
    end
    other_indices = setdiff(1:length(covar.covariance_labels_), conditioning_indices)
    labels = covar.covariance_labels_[other_indices]
    # Segmenting the covariance matrix.
    sigma11 = covar.covariance_[other_indices,other_indices]
    sigma12 = covar.covariance_[other_indices,conditioning_indices]
    sigma21 = covar.covariance_[conditioning_indices,other_indices]
    sigma22 = covar.covariance_[conditioning_indices,conditioning_indices]
    mu1     = zeros(length(other_indices))
    mu2     = zeros(length(conditioning_indices))
    conditioned_values = map( x -> conditioning_draws[x], what_conditioned_on)
    sigma12_invsigma22 = sigma12 * inv(sigma22)
    conditional_mu = mu1 + sigma12_invsigma22 * (conditioned_values - mu2)
    conditional_sigma = sigma11 -  sigma12_invsigma22 * sigma21
    return conditional_mu, conditional_sigma, labels
end


# Pkg.add("ForwardDiff")
using ForwardDiff

"""
```
f(x) = x^2 - 2;
x0 = 1;
x = newton(f, x0);
```
"""
function newton(f, x0)
  # 求目标函数导数
  Δf = x -> ForwardDiff.derivative(f, x);

  xk = Base.float(x0);
  xk_1 = Base.float(x0);

  # 比机器精度少一个数量级
  epsilon = Base.eps(xk) * 10;

  while true
    xk_1 = xk - f(xk)/Δf(xk);
    if Base.abs(xk_1 - xk) < epsilon || Base.abs(f(xk_1)) < epsilon
      break;
    end
    xk = xk_1;
  end

  return xk_1;
end


"""
    bernoulli(n)

 Calculates the first 35 Bernoulli numbers ``B_n``  (of the first-kind or NIST type)
 e.g., see

 + [http://mathworld.wolfram.com/BernoulliNumber.html](http://mathworld.wolfram.com/BernoulliNumber.html)
 + [https://en.wikipedia.org/wiki/Bernoulli_number](https://en.wikipedia.org/wiki/Bernoulli_number)
 + [http://dlmf.nist.gov/24](http://dlmf.nist.gov/24)

 N.B. Bernoulli numbers of second kind only seem to differ in that ``B_1 = + 1/2`` (instead of -1/2)

## Arguments
* ``n`` `::Integer`: the index into the series, ``n=0,1,2,3,...,35`` (for larger ``n`` use `bernoulli(n,0.0)` )

 We only provide the 1st 36 values as beyond this, we can't
 return Int64 rationals, so best to compute the real approximation
 using `bernoulli(n,0.0)`.

 Odd values for ``n>1`` are all zero.

## Examples
```jldoctest; setup = :(using Polylogarithms)
julia> bernoulli(6)
1//42
```
"""
function bernoulli(n::Integer)
    # this just does a lookup -- seemed like it would be easier to code and faster
    # for the size of numbers I am working with
    if n<0
        throw(DomainError(n))
    elseif n > 1 && isodd(n)
        return 0 // 1
    elseif n > 35
        throw(DomainError(n, "If n > 35, then the numerator needs Int128 at least, and worse, so this code is not the code you want. Try using bernoulli(n, 0.0) to get a floating point approximation to the result."))
    end

    # Denominator of Bernoulli number B_n
    #   http://oeis.org/A027642
    D = [2, 6, 1, 30, 1, 42, 1, 30, 1, 66, 1, 2730, 1, 6, 1, 510, 1, 798, 1, 330, 1, 138, 1, 2730, 1, 6, 1, 870, 1, 14322, 1, 510, 1, 6, 1, 1919190, 1, 6, 1, 13530, 1, 1806, 1, 690, 1, 282, 1, 46410, 1, 66, 1, 1590, 1, 798, 1, 870, 1, 354, 1, 56786730]

    # Numerator of Bernoulli number B_n (storing 62 of these because they are easy)
    #   http://oeis.org/A027641
    N = [-1, 1, 0, -1, 0, 1, 0, -1, 0, 5, 0, -691, 0, 7, 0, -3617, 0, 43867, 0, -174611, 0, 854513, 0, -236364091, 0, 8553103, 0, -23749461029, 0, 8615841276005, 0, -7709321041217, 0, 2577687858367, 1]

    if n==0
        return 1
    else
        return N[n] // D[n]
    end
end


module DataGenerator

import Distributions; D = Distributions
using LinearAlgebra

export linregdata, logregdata, poissonregdata

eye(p) = Matrix(1.0 * I, p, p)

defaultβ(p) = collect(range(-1, stop=1, length=p))

function linregdata(n, p; β = defaultβ(p), V = eye(p), μ = zeros(p), σ = 1.0)
	x = rand(D.MvNormal(μ, V), n)'
	y = x * β + σ * randn(n)
	return x, y, β
end

function logregdata(n, p, sgn = true; β = defaultβ(p), V = eye(p), μ = zeros(p))
	x = rand(D.MvNormal(μ, V), n)'
	y = Float64[rand(D.Bernoulli(1 / (1 + exp(-η)))) for η in x * β]
	if sgn
		y = 2y .- 1
	end
	return x, y, β
end

function poissonregdata(n, p, β = defaultβ(p), V = eye(p), μ = zeros(p))
	x = rand(D.MvNormal(μ, V), n)'
	y = Float64[rand(D.Poisson(exp(η))) for η in x * β]
	return x, y, β
end

end


"""
    randSphere(K::Int)

Returns a random K-vector from the unit sphere
"""
randSphere(K::Int) = normalize(randn(K))

"Computes `γ` from `ρ`"
gammasoc(ρ::Float64) = sqrt(1 - ρ^2) / ρ

"Computes `ρ` from `γ`"
rhosoc(γ::Float64)   = 1/ sqrt(1 + γ^2)

"Computes `ϵ` from `μ` and `s^2`"
epssoc(μ::Float64, s2::Float64) = phi(μ/sqrt(1+s2))

"The cumulative function of the Normal distribution"
phi(x::Real) = cdf(Normal(), x)

"The density function of the Normal distribution"
G(x::Real) = pdf(Normal(), x)

"Returns the two modulation functions: Fw and Fϵ"
function modfunc(hσ1γ::Float64, μ1sqs2::Float64)
    phiϵ = phi(μ1sqs2)
    phiw = phi(hσ1γ)
    Z    = phiϵ + phiw - 2*phiϵ*phiw

    Fw = (1 - 2phiϵ) * G(hσ1γ) / Z
    Fϵ = (1 - 2phiw) * G(μ1sqs2) / Z

    return (Fw, Fϵ)
end


# L0 pseudo-norm (times a constant)

export NormL0

"""
**``L_0`` pseudo-norm**

    NormL0(λ=1)

Returns the function
```math
f(x) = λ\\cdot\\mathrm{nnz}(x)
```
for a nonnegative parameter `λ`.
"""
struct NormL0{R <: Real} <: ProximableFunction
    lambda::R
    function NormL0{R}(lambda::R) where {R <: Real}
        if lambda < 0
            error("parameter λ must be nonnegative")
        else
            new(lambda)
        end
    end
end

NormL0(lambda::R=1) where {R <: Real} = NormL0{R}(lambda)

function (f::NormL0)(x::AbstractArray{T}) where {R, T <: RealOrComplex{R}}
    return f.lambda * R(count(v -> v != 0, x))
end

function prox!(y::AbstractArray{T}, f::NormL0, x::AbstractArray{T}, gamma::Real=1) where {R, T <: RealOrComplex{R}}
    countnzy = R(0)
    gl = gamma * f.lambda
    for i in eachindex(x)
        over = abs(x[i]) > sqrt(2 * gl)
        y[i] = over * x[i]
        countnzy += over
    end
    return f.lambda * countnzy
end

fun_name(f::NormL0) = "weighted L0 pseudo-norm"
fun_dom(f::NormL0) = "AbstractArray{Real}, AbstractArray{Complex}"
fun_expr(f::NormL0) = "x ↦ λ countnz(x)"
fun_params(f::NormL0) = "λ = $(f.lambda)"

function prox_naive(f::NormL0, x::AbstractArray{T}, gamma::Real=1) where {R, T <: RealOrComplex{R}}
    over = abs.(x) .> sqrt(2 * gamma * f.lambda)
    y = x.*over
    return y, f.lambda * R(count(v -> v != 0, y))
end


#! /usr/bin/julia

# Rosetta Code, Fibonacci n-step number sequences

type NFib{T<:Integer}
    n::T
    klim::T
    seeder::Function
end

type FState
    a::Array{BigInt,1}
    adex::Integer
    k::Integer
end

function Base.start{T<:Integer}(nf::NFib{T})
    a = nf.seeder(nf.n)
    adex = 1
    k = 1
    return FState(a, adex, k)
end

function Base.done{T<:Integer}(nf::NFib{T}, fs::FState)
    fs.k > nf.klim
end

function Base.next{T<:Integer}(nf::NFib{T}, fs::FState)
    f = sum(fs.a)
    fs.a[fs.adex] = f
    fs.adex = rem1(fs.adex+1, nf.n)
    fs.k += 1
    return (f, fs)
end

function fib_seeder{T<:Integer}(n::T)
    a = zeros(BigInt, n)
    a[1] = one(BigInt)
    return a
end

function fib{T<:Integer}(n::T, k::T)
    NFib(n, k, fib_seeder)
end

function luc_rc_seeder{T<:Integer}(n::T)
    a = zeros(BigInt, n)
    a[1] = 3
    a[2] = -1
    return a
end

function luc_rc{T<:Integer}(n::T, k::T)
    NFib(n, k, luc_rc_seeder)
end

function luc_seeder{T<:Integer}(n::T)
    a = -ones(BigInt, n)
    a[end] = big(n)
    return a
end

function luc{T<:Integer}(n::T, k::T)
    NFib(n, k, luc_seeder)
end

lo = 2
hi = 10
klim = 16

print("n-step Fibonacci for n = (", lo, ",", hi)
println(") up to k = ", klim, ":")
for i in 2:10
    print(@sprintf("%5d => ", i))
    for j in fib(i, klim)
        print(j, " ")
    end
    println()
end

println()
print("n-step Rosetta Code Lucas for n = (", lo, ",", hi)
println(") up to k = ", klim, ":")
for i in 2:10
    print(@sprintf("%5d => ", i))
    for j in luc_rc(i, klim)
        print(j, " ")
    end
    println()
end

println()
print("n-step MathWorld Lucas for n = (", lo, ",", hi)
println(") up to k = ", klim, ":")
for i in 2:10
    print(@sprintf("%5d => ", i))
    for j in luc(i, klim)
        print(j, " ")
    end
    println()
end


# # Comparison with MathOptInterface on a Probability Simplex

# In this example, we project a random point onto a probability simplex with the Frank-Wolfe algorithm using
# either the specialized LMO defined in the package or a generic LP formulation using `MathOptInterface.jl` (MOI) and
# `GLPK` as underlying LP solver.
# It can be found as Example 4.4 [in the paper](https://arxiv.org/abs/2104.06675).

using FrankWolfe

using LinearAlgebra
using LaTeXStrings

using Plots

using JuMP
const MOI = JuMP.MOI

import GLPK

n = Int(1e3)
k = 10000

xpi = rand(n);
total = sum(xpi);
const xp = xpi ./ total;

f(x) = norm(x - xp)^2
function grad!(storage, x)
    @. storage = 2 * (x - xp)
    return nothing
end

lmo_radius = 2.5
lmo = FrankWolfe.FrankWolfe.ProbabilitySimplexOracle(lmo_radius)

x00 = FrankWolfe.compute_extreme_point(lmo, zeros(n))
gradient = collect(x00)

x_lmo, v, primal, dual_gap, trajectory_lmo = FrankWolfe.frank_wolfe(
    f,
    grad!,
    lmo,
    collect(copy(x00)),
    max_iteration=k,
    line_search=FrankWolfe.Shortstep(),
    L=2,
    print_iter=k / 10,
    emphasis=FrankWolfe.memory,
    verbose=false,
    trajectory=true,
);

# Create a MathOptInterface Optimizer and build the same linear constraints:

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

for xi in x
    MOI.add_constraint(o, xi, MOI.GreaterThan(0.0))
end

MOI.add_constraint(
    o,
    MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(1.0, x), 0.0),
    MOI.EqualTo(lmo_radius),
)

lmo_moi = FrankWolfe.MathOptLMO(o)

x, v, primal, dual_gap, trajectory_moi = FrankWolfe.frank_wolfe(
    f,
    grad!,
    lmo_moi,
    collect(copy(x00)),
    max_iteration=k,
    line_search=FrankWolfe.Shortstep(),
    L=2,
    print_iter=k / 10,
    emphasis=FrankWolfe.memory,
    verbose=false,
    trajectory=true,
);

# Alternatively, we can use one of the modelling interfaces based on `MOI` to formulate the LP.
# The following example builds the same set of constraints using `JuMP`:

m = JuMP.Model(GLPK.Optimizer)
@variable(m, y[1:n] ≥ 0)

@constraint(m, sum(y) == lmo_radius)

lmo_jump = FrankWolfe.MathOptLMO(m.moi_backend)

x, v, primal, dual_gap, trajectory_jump = FrankWolfe.frank_wolfe(
    f,
    grad!,
    lmo_jump,
    collect(copy(x00)),
    max_iteration=k,
    line_search=FrankWolfe.Shortstep(),
    L=2,
    print_iter=k / 10,
    emphasis=FrankWolfe.memory,
    verbose=false,
    trajectory=true,
);

x_lmo, v, primal, dual_gap, trajectory_lmo_blas = FrankWolfe.frank_wolfe(
    f,
    grad!,
    lmo,
    x00,
    max_iteration=k,
    line_search=FrankWolfe.Shortstep(),
    L=2,
    print_iter=k / 10,
    emphasis=FrankWolfe.blas,
    verbose=false,
    trajectory=true,
);

x, v, primal, dual_gap, trajectory_jump_blas = FrankWolfe.frank_wolfe(
    f,
    grad!,
    lmo_jump,
    x00,
    max_iteration=k,
    line_search=FrankWolfe.Shortstep(),
    L=2,
    print_iter=k / 10,
    emphasis=FrankWolfe.blas,
    verbose=false,
    trajectory=true,
);

iteration_list = [[x[1] + 1 for x in trajectory_lmo], [x[1] + 1 for x in trajectory_moi]]
time_list = [[x[5] for x in trajectory_lmo], [x[5] for x in trajectory_moi]]
primal_gap_list = [[x[2] for x in trajectory_lmo], [x[2] for x in trajectory_moi]]
dual_gap_list = [[x[4] for x in trajectory_lmo], [x[4] for x in trajectory_moi]]

label = [L"\textrm{Closed-form LMO}", L"\textrm{MOI LMO}"]

FrankWolfe.plot_results(
    [primal_gap_list, primal_gap_list, dual_gap_list, dual_gap_list],
    [iteration_list, time_list, iteration_list, time_list],
    label,
    ["", "", L"\textrm{Iteration}", L"\textrm{Time}"],
    [L"\textrm{Primal Gap}", "", L"\textrm{Dual Gap}", ""],
    xscalelog=[:log, :identity, :log, :identity],
    yscalelog=[:log, :log, :log, :log],
    legend_position=[:bottomleft, nothing, nothing, nothing]
)


# mathematical constants related to statistics

@irrational twoπ   6.2831853071795864769 big(π) * 2.0
@irrational fourπ  12.566370614359172954 big(π) * 4.0
@irrational halfπ  1.5707963267948966192 big(π) * 0.5
@irrational quartπ 0.7853981633974483096 big(π) * 0.25

@irrational invπ     0.31830988618379067154 inv(big(π))
@irrational twoinvπ  0.63661977236758134308 big(invπ) * 2.0
@irrational fourinvπ 1.27323954473516268615 big(invπ) * 4.0
@irrational inv2π    0.159154943091895335769 big(invπ) * 0.5
@irrational inv4π    0.079577471545947667884 big(invπ) * 0.25

@irrational sqrt2     1.4142135623730950488 sqrt(big(2.0))
@irrational sqrt3     1.7320508075688772935 sqrt(big(3.0))
@irrational sqrtπ     1.7724538509055160273 sqrt(big(π))
@irrational sqrt2π    2.5066282746310005024 sqrt(big(π) * 2.0)
@irrational sqrt4π    3.5449077018110320546 sqrt(big(π) * 4.0)
@irrational sqrthalfπ 1.2533141373155002512 sqrt(big(0.5)*π)

@irrational invsqrt2  0.7071067811865475244 inv(big(sqrt2))
@irrational invsqrt2π 0.3989422804014326779 inv(big(sqrt2π))


# 1D diffusion problem

# TODO: Add more complex tests.

# Packages and inclusions
using ModelingToolkit,DiffEqOperators,DiffEqBase,LinearAlgebra,Test

# Tests
@testset "Test 00: Dt(u(t,x)) ~ Dxx(u(t,x))" begin
    # Parameters, variables, and derivatives
    @parameters t x
    @variables u(..)
    @derivatives Dt'~t
    @derivatives Dxx''~x

    # 1D PDE and boundary conditions
    eq  = Dt(u(t,x)) ~ Dxx(u(t,x))
    bcs = [u(0,x) ~ -x*(x-1)*sin(x),
           u(t,0) ~ 0.0,
           u(t,1) ~ 0.0]

    # Space and time domains
    domains = [t ∈ IntervalDomain(0.0,1.0),
               x ∈ IntervalDomain(0.0,1.0)]

    # PDE system
    pdesys = PDESystem(eq,bcs,domains,[t,x],[u])

    # Method of lines discretization
    dx = 0.1
    order = 2
    discretization = MOLFiniteDifference(dx,order)

    # Convert the PDE problem into an ODE problem
    prob = discretize(pdesys,discretization)

    # Solve ODE problem
    using OrdinaryDiffEq
    sol = solve(prob,Tsit5(),saveat=0.1)

    #Plot and save results
    # using Plots
    # plot(prob.space[2],Array(prob.extrapolation[1]*sol[:,1,1]))
    # plot!(prob.space[2],Array(prob.extrapolation[1]*sol[:,1,2]))
    # plot!(prob.space[2],Array(prob.extrapolation[1]*sol[:,1,3]))
    # plot!(prob.space[2],Array(prob.extrapolation[1]*sol[:,1,4]))
    # savefig("MOL_1D_Linear_Diffusion_Test00.png")

    # Test
    n = size(sol,1)
    t_f = size(sol,3)

    @test sol[:,1,t_f] ≈ zeros(n) atol = 0.001;
end

@testset "Test 01: Dt(u(t,x)) ~ D*Dxx(u(t,x))" begin
    # Parameters, variables, and derivatives
    @parameters t x D
    @variables u(..)
    @derivatives Dt'~t
    @derivatives Dxx''~x

    D = 1.1

    # 1D PDE and boundary conditions
    eq  = Dt(u(t,x)) ~ D*Dxx(u(t,x))
    bcs = [u(0,x) ~ -x*(x-1)*sin(x),
           u(t,0) ~ 0.0,
           u(t,1) ~ 0.0]

    # Space and time domains
    domains = [t ∈ IntervalDomain(0.0,1.0),
               x ∈ IntervalDomain(0.0,1.0)]

    # PDE system
    pdesys = PDESystem(eq,bcs,domains,[t,x,D],[u])

    # Method of lines discretization
    dx = 0.1
    order = 2
    discretization = MOLFiniteDifference(dx,order)

    # Convert the PDE problem into an ODE problem
    prob = discretize(pdesys,discretization)

    # Solve ODE problem
    using OrdinaryDiffEq
    sol = solve(prob,Tsit5(),saveat=0.1)

    # Plot and save results
    # using Plots
    # plot(prob.space[2],Array(prob.extrapolation[1]*sol[:,1,1]))
    # plot!(prob.space[2],Array(prob.extrapolation[1]*sol[:,1,2]))
    # plot!(prob.space[2],Array(prob.extrapolation[1]*sol[:,1,3]))
    # plot!(prob.space[2],Array(prob.extrapolation[1]*sol[:,1,4]))
    # savefig("MOL_1D_Linear_Diffusion_Test01.png")

    # Test
    n = size(sol,1)
    t_f = size(sol,3)
    @test sol[:,1,t_f] ≈ zeros(n) atol = 0.001;
end

@testset "Test 02: Dt(u(t,x)) ~ Dx(D(t,x))*Dx(u(t,x))+D(t,x)*Dxx(u(t,x))" begin
    # Parameters, variables, and derivatives
    @parameters t x
    @variables u(..) D(..)
    @derivatives Dt'~t
    @derivatives Dx'~x
    @derivatives Dxx''~x

    # 1D PDE and boundary conditions

    eq  = [ Dt(u(t,x)) ~ Dx(D(t,x))*Dx(u(t,x))+D(t,x)*Dxx(u(t,x)),
            D(t,x) ~ 0.999 + 0.001 * t * x  ]

    bcs = [u(0,x) ~ -x*(x-1)*sin(x),
           u(t,0) ~ 0.0,
           u(t,1) ~ 0.0,
           D(0,x) ~ 0.999,
           D(t,0) ~ 0.999,
           D(t,1) ~ 0.999 + 0.001 * t ]

    # Space and time domains
    domains = [t ∈ IntervalDomain(0.0,1.0),
               x ∈ IntervalDomain(0.0,1.0)]

    # PDE system
    pdesys = PDESystem(eq,bcs,domains,[t,x],[u,D])

    # Method of lines discretization
    dx = 0.1
    order = 2
    discretization = MOLFiniteDifference(dx,order)

    # Convert the PDE problem into an ODE problem
    prob = discretize(pdesys,discretization)

    # Solve ODE problem
    using OrdinaryDiffEq
    sol = solve(prob,Tsit5(),saveat=0.1)

    # Plot and save results
    # using Plots
    # plot(prob.space[2],Array(prob.extrapolation[1]*sol[:,1,1]))
    # plot!(prob.space[2],Array(prob.extrapolation[1]*sol[:,1,2]))
    # plot!(prob.space[2],Array(prob.extrapolation[1]*sol[:,1,3]))
    # plot!(prob.space[2],Array(prob.extrapolation[1]*sol[:,1,4]))
    # savefig("MOL_1D_Linear_Diffusion_Test02.png")

    # Test
    n = size(sol,1)
    t_f = size(sol,3)
    @test sol[:,1,t_f] ≈ zeros(n) atol = 0.1;
end


#############################################################################
#############################################################################
#
# This file implements polynomial division
#
#############################################################################
#############################################################################

"""  Modular algorithm.
f divide by g

f = q*g + r

p is a prime
"""

function divide(num::Polynomial, den::Polynomial)
    function division_function(p::Int)
        f, g = mod(num,p), mod(den,p)
        degree(f) < degree(num) && return nothing
        iszero(g) && throw(DivideError())
        q = Polynomial()
        prev_degree = degree(f)
        while degree(f) ≥ degree(g)
            h = Polynomial( (leading(f) ÷ leading(g))(p) )  #syzergy
            f = mod((f - h*g), p)
            q = mod((q + h), p)
            prev_degree == degree(f) && break
            prev_degree = degree(f)
        end

        @assert iszero( mod((num  - (q*g + f)),p))
        return q, f
    end
    return division_function
end

function divide(num::PolynomialModP, den::PolynomialModP)
    p = num.mod
    @assert num.mod == den.mod
    f, g = num, den
    if degree(g) > degree(f)
        output =  [zero(PolynomialModP,p) , f]
        return output
    end

    iszero(g) && throw(DivideError())
    iszero(f) && return zero(PolynomialModP,p), zero(PolynomialModP,p)

    q = PolynomialModP(f.mod)
    prev_degree = degree(f)
    while degree(f) ≥ degree(g)
        h = (leading(f) ÷ leading(g))(p) #syzergy
        hg = - h*g
        f = f+hg
        q = q+h
        prev_degree == degree(f) && break
        prev_degree = degree(f)
    end

    @assert iszero(num - (q*g + f))

    return q,f
end

"""
The quotient from polynomial division. Returns a function of an integer.
"""
÷(num::Polynomial, den::Polynomial)  = (p::Int) -> first(divide(num,den)(p))
÷(num::PolynomialModP, den::PolynomialModP) = first(divide(num,den))

"""
The remainder from polynomial division. Returns a function of an integer.
"""
rem(num::Polynomial, den::Polynomial)  = (p::Int) -> last(divide(num,den)(p))
rem(num::PolynomialModP, den::PolynomialModP)  = last(divide(num,den))


using Symata

@sym begin
  ClearAll(f, y0, a, l, x)
  y0(x_) := (e*x)*(l^3 - 2*l*x^2 + x^3)

  # Stodola-Vianello method, e.g. see Den Hartog, Advanced Strength of Materials, pg 268
  f(maxiters_, y0_) := Module([i, yder, res=ConstantArray(0.0,[maxiters,6])],
    (
      # Initialize the first row of the res matrix
      res[1,2] = y0(x),
      For(i=1, i<=4, Increment(i),
        [
          yder(x_) = Simplify(D(y0, x, i))
          res[1, 2+i] = yder(x)
        ]
      ),
      Println("N[res[1, 1]]: ", res[1,1]),
      # Start the loop to improve accuracy stored in subsequent rows,
      # res[i, 1] contains the Euler ratio.
      For(i=2, i<=maxiters, Increment(i),
        begin
          res[i, 6] = Simplify(w/(EI)*((x-l)*res[i-1,4] + res[i-1,3]))
          # Above the basic Euler equation for a pinned-pinned flagpole
          ClearAll(K1, K11, tmpy5, tmpy4, tmp)
          tmpy5(x_) = Simplify(Integrate(res[i,6], x) + K1)
          tmpy4(x_) = Simplify(Integrate(tmpy5(x), x))
          # Use first pinned constraint at x=l to resolve K1
          K11 = Solve(-tmpy4(l), K1)[1]
          res[i,5] = Simplify( tmpy5(x) ./ K11 )
          res[i,4] = Simplify( tmpy4(x) ./ K11 )
          ClearAll(K2, K22, tmpy3, tmpy2, tmp)
          tmpy3(x_) = Simplify(Integrate(res[i,4], x) + K2)
          tmpy2(x_) = Simplify(Integrate(tmpy3(x), x))
          # Use second pinned constraint at x=l to resolve K2
          K22 = Solve(-tmpy2(l), K2)[1]
          res[i,3] = Simplify( tmpy3(x) ./ K22 )
          res[i,2] = Simplify( tmpy2(x) ./ K22 )
          tmp = Simplify( (((l^3)*w)/EI)*(res[i-1,3]/res[i, 3]) )
          res[i,1] = N(Simplify(ReplaceAll(tmp, x => l)))
          Println("N[res[$i, 1]]: ", N(res[i,1], 4))
        end
      ),
      Println(),
      Return(res)
    )
  )

  r = f(8, y0(x))
  Println(N(Transpose(r)[1], 4))
end


module FirstOrder

using LinearAlgebra

import ..LocalDescent:
search,
LineSearch, StrongBacktracking

struct Termination
    max_iter
    ϵ_abs
    ϵ_rel
    ϵ_grad
    Termination(;
        max_iter = 10_000,
        ϵ_abs = eps(),
        ϵ_rel = eps(),
        ϵ_grad = eps()) =
        new(max_iter, ϵ_abs, ϵ_rel, ϵ_grad)
    Termination(ϵ;
        max_iter = 10_000) =
        new(max_iter, ϵ, ϵ, ϵ)
end

mutable struct TerminationConditions
    abs
    rel
    grad
    TerminationConditions(term::Termination) =
    new((fx, fx_next)->fx - fx_next < term.ϵ_abs,
        (fx, fx_next)->fx - fx_next < term.ϵ_abs * abs(fx),
        ∇fx_next->norm(∇fx_next) < term.ϵ_grad)
end

function descent_until(term::Termination, descent_method, f, ∇f, X, trace)
    term_cond = TerminationConditions(term)
    for _ = 1:term.max_iter
        x = X[:, end]
        x_next = descent_method(f, ∇f, x)
        if trace
            X = hcat(X, x_next)
        else
            X = x_next
        end
        fx, fx_next, ∇fx_next = f(x), f(x_next), ∇f(x_next)
        if term_cond.abs(fx, fx_next) ||
            term_cond.rel(fx, fx_next) ||
            term_cond.grad(∇fx_next)
            return X
        end
    end
    @warn "descent_until: max number of iterations reached"
    X
end

abstract type FirstOrderMethods end

# Maximum Gradient Descent

struct MaximumGradientDescent <: FirstOrderMethods end

function search(::MaximumGradientDescent, f, ∇f, x_0; term = Termination(), trace = false)
    descent_until(term, descent_step_maxgd, f, ∇f, x_0, trace)
end

function descent_step_maxgd(f, ∇f, x)
    """
    Maximum Gradient Descent

    A step of maximum gradient descent produces the value of x(k+1)
    where f(x) in the direction of maximum descent is at its local
    minimum.
    """
    g = ∇f(x)
    d = direction_maxgd(g)
    search(LineSearch(), f, x, d)
end

function direction_maxgd(g)
    -g / norm(g)
end

# Gradient Descent

struct GradientDescent <: FirstOrderMethods
    α
    GradientDescent(; α = 0.001) = new(α)
end

function search(params::GradientDescent, f, ∇f, x_0; term = Termination(), trace = false)
    descent_step = (f, ∇f, x)->descent_step_gd(params.α, f, ∇f, x)
    descent_until(term, descent_step, f, ∇f, x_0, trace)
end

function descent_step_gd(α, f, ∇f, x)
    """
    Gradient Descent
    Algorithm 5.1

    A step of a gradient descent algorithm with fixed learning rate α.
    """
    g = ∇f(x)
    x - α * g
end

# Conjugate Gradient Descent

struct ConjugateGradientDescent <: FirstOrderMethods end

mutable struct CGDProblemState
    g
    d
    CGDProblemState(dim) = new(ones(dim), zeros(dim))
end

function search(params::ConjugateGradientDescent, f, ∇f, x_0; term = Termination(), trace = false)
    state = CGDProblemState(size(x_0))
    descent_step = (f, ∇f, x)->descent_step_cgd!(state, f, ∇f, x)
    descent_until(term, descent_step, f, ∇f, x_0, trace)
end

function descent_step_cgd!(state, f, ∇f, x)
    """
    Conjugate Gradient Descent
    Algorithm 5.2

    https://en.wikipedia.org/wiki/Nonlinear_conjugate_gradient_method
    """
    g = ∇f(x)
    d = direction_cgd(g, state.g, state.d)
    state.g = g
    state.d = d
    search(StrongBacktracking(), f, ∇f, x, d)
end

function direction_cgd(g, gm1, dm1)
    """
    Polak-Ribiere update
    Notation:
        g is for g(k)
        dm1 is for d(k-1)
        gm1 indicates g(k-1)
    """
    β_PR = (g' * (g - gm1)) / (gm1'gm1)
    β = max(β_PR, 0)
    -g + β * dm1
end

end # module


include("kruskal.jl")

"""
    count_spanning_trees(g::AGraph)

Returns the number of spanning trees of `g`, computed through
[Kirchhoff's theorem](https://en.wikipedia.org/wiki/Kirchhoff%27s_theorem).
The return type is a float, since the number can be very large.
"""
function count_spanning_trees(g::AGraph)
    @assert is_connected(g) "The graph has to be connected"
    nv(g) <= 1 && return 1.
    Δ = laplacian_matrix(g)[2:nv(g),2:nv(g)]
    return round(det(Δ))
end


"""
    gpw(ζ, φ; k=1)

Closure that defines a generalized plane wave in polar coordinates.

The closure captures the values of the evanescence parameter ``ζ``, the
angle ``φ`` and the wavenumber ``k``.
It does not include any normalization.
The function can then be evaluated at any ``(r,θ)`` point.

The generalized plane wave in polar coordinates is defined by
```math
ϕ := (r,θ) ↦ e^{ı k \\mathbf{d} ⋅ \\mathbf{x}},
```
where
```math
\\mathbf{d} = (\\cos[φ+ıζ], \\sin[φ+ıζ]),
\\qquad\\text{and}\\qquad
\\mathbf{x} = (r \\cos θ, r\\sin θ).
```
"""
function gpw(ζ,φ;k=1)
    d = [cos(φ+im*ζ), sin(φ+im*ζ)] # Direction
    return (r,θ) -> exp(im * k * [r*cos(θ), r*sin(θ)] ⋅ d)
end

"""
    approximation_set(Y, W; k=1)

Generate a set of generalized plane waves.

The parameters `Y` and `W` are respectively the sets of complex angles
``(ζ_j,φ_j)`` and associated weights ``ω_j`` corresponding to the generalized
plane waves, see [`gpw`](@ref)
```math
ϕ_j := (r,θ) ↦ e^{ı k \\mathbf{d}_j ⋅ \\mathbf{x}},
```
where
```math
\\mathbf{d}_j = (\\cos[φ_j+ıζ_j], \\sin[φ_j+ıζ_j]),
\\qquad\\text{and}\\qquad
\\mathbf{x} = (r \\cos θ, r\\sin θ).
```
Approximations are constructed in the form of
```math
(r, θ) ↦ \\sum_j ξ_j ω_j ϕ_j(r, θ),
```
where ``(ξ_j)_j`` is the set of unknown coefficients.

The value of the wavenumber `k` defaults to ``1`` and can be provided as an
optional parameter.
"""
function approximation_set(Y, W; k=1)
    @assert length(Y) == length(W)
    return [(r, θ) -> w * gpw(y...;k=k)(r, θ) for (y, w) in zip(Y, W)]
end

"""
    approximation_set(N; k=1)

Generate a set of ``N`` propagative plane waves with equispaced angles.

The equispaced angles are defined by
```math
θ_n := 2π s / R, \\qquad n = 0,…,N-1.
```

The value of the wavenumber `k` defaults to ``1`` and can be provided as an
optional parameter.
"""
function approximation_set(N; k=1)
    Y = [(0, 2π * n / N) for n in 0:N-1]
    W = ones(Int64, N) ./ sqrt(N)
    return approximation_set(Y, W; k=k)
end

"""
    approximation_set(N, qs, smpl_type; k=1)
    approximation_set(N, Q::Integer, smpl_type; k=1)

Generate a set of ``N`` evanescent plane waves according to some sampling type.

The second parameter can be an integer ``Q`` which controls the truncation
parameter of the [`TruncKernel`](@ref), or a set of integers `qs` to specify
only some specific terms in the expansion.
The parameter `smpl_type` should be a function name: see
[`uniform_sampling`](@ref), [`sobol_sampling`](@ref) and
[`random_sampling`](@ref).

The value of the wavenumber `k` defaults to ``1`` and can be provided as an
optional parameter.
"""
function approximation_set(N, qs, smpl_type; k=1)
    K = TruncKernel(Tuple(qs), logwz(;k=k));
    smpl = Sampler(K)
    Y, W = smpl_type(smpl, N)
    W ./= sqrt(N)
    return approximation_set(Y, W; k=k)
end
function approximation_set(N, Q::Int64, smpl_type; k=1)
    return approximation_set(N, -Q:Q, smpl_type; k=k)
end


# # Geometry

# The `WaveProp.Geometry` module handles everything related to the geometrical
# descriptions of curves/surfaces. Its main interface is exposed through
# `AbstractElement`, which represents an `M`-dimensional geometical object
# embedded in ℝᴺ. In abstract terms, an *element* τ is a subset of ℝᴺ described
# by
# * a reference element ``\tau \subset \mathbb{R}^N``
# * a *push-forward map* χ : τ̂ → τ ⊂ ℝᴺ.

# Concrete implementations of `AbstractElement`s are responsible for
# implementing a few crucial methods used to derive almost all functionality of
# an `el::AbstractElement` used throughout this package. In the [next section](#interface) we
# describe this interface.

# ## Interface

# The interface is composed of the following methods:
# - `domain(el)` : return the reference element `R::AbstractReferenceElement`
# - `el(u)` : evaluate the parametrization of the element at (parametric).
#   coordinate `u ∈ domain(el)`.
# - `jacobian(u)` : evaluate the jacobian of the parametrization at (parametric)
#   coordinate `u ∈ domain(el)`.

# ## Lagrangian elements

# An concrete implementation of `AbstractElement` is the parametric type
# `LagrangeElement{D,Np,N,T}`. The type parameters represent:
# - `D` : the domain of the element. Thus, `domain(el)` return an instance of
#   the singleton type `R`.
# - `Np` : the number of points/DOF which describe the element. For example, for
#   a flat triangle, `Np=3`, while for a flat quadrilateral element `Np=4`.
# - `N` : the dimension of the ambient/embedding space.
# - `T` : the data type used in the presentation of `el`. Defaults to `Float64`.


# Fourier demo with gtk-server
# #
# Note: requires package Gaston for plotting, or adjust the code to work with your
# plotting front-end.

using GtkServer, Gaston, FFTW

# plot
function doplots(t, A, f₀, ϕ, fs)
    s = A.*sin.(2π*f₀*t .+ ϕ)
    S = fftshift(fft(s))./length(s)
    F = fftshift(fftfreq(length(s), fs))

    p1 = plot(t, s, handle = 1,
              Axes(title = "'Sinusoidal in time domain'", xlabel = "'Time (s)'",
                   ylabel = "'Amplitude'", yrange = (-5.5, 5.5), ytics = -6:2:6))
    p2 = plot(F, abs.(S), handle = 2,
              Axes(title = "'Magnitude spectrum'", xlabel = "'Frequency (Hz)'",
                   yrange = (0, 3)))
    p3 = plot(F, angle.(S), handle = 3,
              Axes(title = "'Phase spectrum'", xlabel = "'Frequency (Hz)'",
                   yrange = (-3, 3)))
    P = plot([p1, p2, p3], handle = 4)
    display(P)
end

function demo()
    # Set initial values for amplitude, frequency and phase, and display
    # the sine function.
    fs = 100
    t = 0.0:1.0/fs:1.0
    A = 1.0
    f₀ = 1.0
    ϕ = 0.0

    # Do initial plot with default values
    doplots(t, A, f₀, ϕ, fs)

    # Main program
    # This program interactively plots a sine function in time, and its
    # magnitude and phase spectra. The amplitude, frequency and phase of
    # sine function are selected through three sliders.
    gtk_init()
    # 'win' is a handle (ID number) to the main GUI window.
    gtk_send("gtk_init NULL NULL")
    win = gtk_send("gtk_window_new 0")
    # Set window's title, size and position.
    # Note that the message sent to gtk-server is always a string.
    gtk_send("gtk_window_set_title $win 'Fourier Series Demo'")
    gtk_send("gtk_window_set_default_size $win 300 200")
    gtk_send("gtk_window_set_position $win 1")
    # We'll set our widgets in a table, one widget per row-column combination.
    tbl = gtk_send("gtk_table_new 8 10 0")
    gtk_send("gtk_container_add $win $tbl")
    # Create three labels...
    lblA = gtk_send("gtk_label_new Amplitude:")
    lblF = gtk_send("gtk_label_new Frequency:")
    lblP = gtk_send("gtk_label_new Phase:")
    # and put them in specific position on our table.
    gtk_send("gtk_table_attach_defaults $tbl $lblA 1 3 1 2")
    gtk_send("gtk_table_attach_defaults $tbl $lblF 1 3 3 4")
    gtk_send("gtk_table_attach_defaults $tbl $lblP 1 3 5 6")
    # horizontal separator
    hsep = gtk_send("gtk_hseparator_new")
    gtk_send("gtk_table_attach_defaults $tbl $hsep 1 9 6 7")
    # Some Julia advertisement
    adv = gtk_send("gtk_label_new 'Created with Julia!'")
    gtk_send("gtk_table_attach_defaults $tbl $adv 1 6 7 8")
    # Create (and put on the table) the quit button.
    but = gtk_send("gtk_button_new_from_stock 'gtk-quit'")
    gtk_send("gtk_table_attach_defaults $tbl $but 7 9 7 8")
    # Create three sliders...
    sliA = gtk_send("gtk_hscale_new_with_range 1 5 0.5")
    sliF = gtk_send("gtk_hscale_new_with_range 1 40 1")
    sliP = gtk_send("gtk_hscale_new_with_range -1.6 1.6 0.1")
    # and put them on the table.
    gtk_send("gtk_table_attach_defaults $tbl $sliA 4 9 1 2")
    gtk_send("gtk_table_attach_defaults $tbl $sliF 4 9 3 4")
    gtk_send("gtk_table_attach_defaults $tbl $sliP 4 9 5 6")
    # Now we need to actually display all the widgets we've created.
    gtk_send("gtk_widget_show_all $win")

    # endless loop; we only break out of it when the quit button is clicked
    while true
        # Up to this point, we have sent commands to gtk-server in order
        # to set up the GUI and display it. However, once the GUI is
        # created, gtk-server will send messages (events) to us whenever
        # the user clicks something on the GUI. In order to read these
        # events, we send gtk-server the message "gtk_server_callback wait".
        # So, if the user has clicked on a widget on the GUI, then
        # 'event' contains the widget's handle (ID number).
        event = gtk_send("gtk_server_callback wait")
        # The user clicked on the windows 'close' button.
        if event == win
            break
        end
        # The user clicked on the quit button.
        if event == but
            break
        end
        # The user clicked on the 'amplitude' slider. Now we need
        # to go to gtk-server again to ask for the new slider's value.
        if event == sliA
            # here we get the new value
            A_g = gtk_send("gtk_range_get_value $sliA")
            A = parse(Float32, A_g)
            # plot again with the new sine values
            P = doplots(t, A, f₀, ϕ, fs)
        end
        # Frequency slider.
        if event == sliF
            f₀_g = gtk_send("gtk_range_get_value $sliF")
            f₀ = parse(Float32, f₀_g)
            P = doplots(t, A, f₀, ϕ, fs)
        end
        # ϕ slider.
        if event == sliP
            ϕ_g = gtk_send("gtk_range_get_value $sliP")
            ϕ = parse(Float32, ϕ_g)
            P = doplots(t, A, f₀, ϕ, fs)
        end
    end

    # Clean up -- close all figures and tell gtk-server to quit.
    closeall()
    gtk_exit()
end


# Practical Machine learning
# Regression Analysis - Linear Regression example
# Chapter 10

using Gadfly

# Normalizes the features in x
# The mean value of each feature is 0 and the standard deviation is 1
# Returns normalized x, μ and σ
function featureNormalize(x)
rows = size(x,1)
cols = size(x,2)

μ = mean(x,1)
σ = std(x,1)
xNorm = zeros(x)

# normalize
for i in 1:cols
	for j in 1:rows
		xNorm[j,i] = (x[j,i] - μ[i]) / σ[i];
	end
end

(xNorm, μ, σ)
end

println("Loading data ... ")
data = readdlm("data.txt",',')
x = data[:,1:2]
y = data[:, 3]
m = length(y)

@printf("First 10 examples from the dataset: \n");
t = [x[1:10,:] y[1:10,:]]'
for i in 1:10
  @printf(" x = [%.0f %.0f], y = %.0f \n", t[1,i], t[2,i], t[3,i]);
end

# Scale features and set them to zero mean
(x, μ, σ) = featureNormalize(x);

# Add intercept term to x
x = [ones(m,1) x]

#### Run Gradient Descent
α = 0.001
numIter = 4000
θ = zeros(3,1)
jHist = zeros(numIter, 1)

for i in 1:numIter
  # next theta
  θ = θ - (α/m) * (x' * ((x*θ)-y))
  # compute cost
  jHist[i] = sum((x*θ-y).^2)/(2m)
end

# plot convergence graph
pl = plot(
  x=collect(1:numIter),
  y=jHist,
  Guide.xlabel("Iterations"),
  Guide.ylabel("Error"),
  Guide.title("Convergence Graph"),
  Geom.line
  )
draw(SVGJS("jHist.js.svg", 6inch, 6inch), pl)

# Estimate the price of a 1650 sq-ft, 3 br house
price = [1, (1650-μ[1])/σ[1], (3-μ[2])/σ[2]]' * θ
println("Estimated price for a 1650 sq-ft, 3 br house: $price")

println("done!")


using LinearAlgebra, SparseArrays, Random, Statistics
using Plots, Plots.PlotMeasures

function gomp(y0::AbstractVector, Ψ::AbstractMatrix,
    K::Int = 2, S::Int = 1;
    max_iter::Int = 100, ϵ::Real = eps())

	# Get the dimensions
	m = length(y0)
	m_psi, n = size(Ψ)
	# Assert the dimensionality
	@assert m == m_psi "Please provide consistent input sizes"
	# Assert the selector
	@assert S <=  min(K, m/K) "S <= min(K, m/K)"

	# Normalize
	ψ = deepcopy(Ψ)
    # Inplace
	normalize!.(eachcol(ψ), 2)

	# Preconditioning
	P = ψ'pinv(ψ*ψ')
	# New matrix
	ψ = P*ψ
	y = P*y0
	# Iteration
	iters = 0

	# Support
	Λ = zeros(Bool, n)
	u = zeros(eltype(y), n)
	r = y
	amps = zeros(eltype(y), n)

	# Find the magnitudes
	for i in 1:max_iter
		# Compute the similarities via magnitude
		amps .= ψ'r
		# Get the largest entry
		idx = sortperm(abs.(amps), rev = true)[1:S]
		# Update the support
		Λ[idx] .= true
		# Update the coefficients
		u[Λ] .= (y' / ψ')[1, Λ]
		# Update r
		r .= y - ψ*u
		# Convergence
		(norm(r,2) < ϵ || sum(Λ) >= K) && break
	end

    # Last time to get the right coefficients
	u[Λ] .= (y0' / Ψ')[1, Λ]

	return u
end

## Test data
Random.seed!(1111) # Lucky
x̂ = sprandn(Float64, 100, 0.1)
k_opt = Int(norm(x̂, 0))
A = randn(Float64, 300, 100)
y = A*x̂

# Evaluate
x = gomp(y, A, k_opt)

# Check for different sparsties levels
sp_n = map(0.01:0.01:0.2) do i
	x̃ = sprandn(Float64, 100, i)
	ỹ = A*x̃
	norm(gomp(ỹ, A, 50, 1, ϵ = 0.01), 0), norm(x̃, 0)
end

gr()
pl_1 = scatter(x̂, label = "Ground Truth", markeralpha = 0.5, xlabel = "Index", ylabel = "x", legend = :bottomright)
scatter!(x, label = "Estimated", markershape = :cross, markersize = 12, markerstrokewidth = 5)
savefig(joinpath(pwd(), "figures", "gomp.pdf"))

xs, ys = first.(sp_n), last.(sp_n)
pl_2 = scatter(ys, xlabel = "Run", ylabel = "Sparsity", label = "Ground Truth", legend = :topleft)
scatter!(xs, label = "Estimated", markershape = :cross, markersize = 12, markerstrokewidth = 5)
savefig(joinpath(pwd(), "figures", "gomp_vary.pdf"))

plot(pl_1, pl_2, layout = (1,2), size = (900, 300), margins = 5mm)
savefig(joinpath(pwd(), "figures", "merged.pdf"))


using ApproxFun

#The following solves the Airy ODE with dirichlet boundary conditions

x=Fun(identity,-1000..15)   # Fun corresponding to multiplication by x, on [-100,15]
d=domain(x)
D=Derivative()             # The derivative operator
B=Dirichlet()              # Dirichlet boundary conditions, [u(-100),u(15)]

#Construct operator

A=[B;D^2-x]                # This is dirichlet conditions and u'' - x u
b=[airyai(leftendpoint(d)),0]     # We want it to equal airyai(-100) at -100, and 0 at
                           # 10, with 0 rhs
#Solve ODE

u=A\[b,0]                  # u satisfies A*u = b, or in other words,
                           # B*u = [airyai(-100.),0.] and (D2 - x)*u = 0.

# Check the accuracy
norm(u - Fun(airyai,d))

## We now solve with Neumann conditions

B=Neumann()
A=[B;D^2-x]
b=[[airyaiprime(leftendpoint(d)),0.],0.]

u=A\b

# Check the accuracy
norm(u - Fun(airyai,d))