"""
    jacobi(x, p::Integer, α, β)

Evaluate the Legendre polynomial with parameters `α`, `β` of degree `p` at `x`
using the three term recursion [Karniadakis and Sherwin, Spectral/hp Element
Methods for CFD, Appendix A].
"""
function jacobi(x, p::Integer, α, β)
    T = typeof( (2+α+β)*x / 2 )
    a = one(T)
    b = ((2+α+β)*x + α - β) / 2

    if p <= 0
        return a
    elseif p == 1
        return b
    end

    for n in 2:p
        a1 = 2n*(n+α+β)*(2n-2+α+β)
        a2 = (2n-1+α+β)*(α+β)*(α-β)
        a3 = (2n-2+α+β)*(2n-1+α+β)*(2n+α+β)
        a4 = 2*(n-1+α)*(n-1+β)*(2n+α+β)
        a, b = b, ( (a2+a3*x)*b - a4*a ) / a1
    end

    b
end

"""
    jacobi_and_derivative(x, p::Integer, α, β)

Evaluate the Jacobi polynomial with parameters `α`, `β` of degree `p` and its
derivative at `x` using the three term recursion [Karniadakis and Sherwin,
Spectral/hp Element Methods for CFD, Appendix A].
"""
function jacobi_and_derivative(x, p::Integer, α, β)
    T = typeof( (2+α+β)*x / 2 )
    # Coefficients for the polynomial...
    a = one(T)
    b = ((2+α+β)*x + α - β) / 2
    # ... and for the derivative. The derivative is computed via eq. (A.1.8)
    # of Karniadakis and Sherwin as Jacobi polynomial with parameters α+1, β+1.
    aa = one(T)
    bb = ((4+α+β)*x + α - β) / 2

    if p <= 0
        return a, zero(T)
    elseif p == 1
        return b, (2+α+β)/2
    end

    for n in 2:p
        a1 = 2n*(n+α+β)*(2n-2+α+β)
        a2 = (2n-1+α+β)*(α+β)*(α-β)
        a3 = (2n-2+α+β)*(2n-1+α+β)*(2n+α+β)
        a4 = 2*(n-1+α)*(n-1+β)*(2n+α+β)
        a, b = b, ( (a2+a3*x)*b - a4*a ) / a1

        b1 = 2n*(n+2+α+β)*(2n+α+β)
        b2 = (2n+1+α+β)*(α+β+2)*(α-β)
        b3 = (2n+α+β)*(2n+1+α+β)*(2n+2+α+β)
        b4 = 2*(n+α)*(n+β)*(2n+2+α+β)
        aa, bb = bb, ( (b2+b3*x)*bb - b4*aa ) / b1
    end

    b, (α+β+p+1)*aa/2
end

"""
    gauss_jacobi_nodes_and_weights(p, α, β, T=Float64::Type, tol=4*eps(T), maxit=100)

Compute the Gauss-Jacobi nodes and weights for polynomials of degree `p` with
parameters `α`, `β` using the scalar type `T`, tolerance `tol` and maximal number
of Newton iterations `maxit` [Karniadakis and Sherwin, Spectral/hp Element
Methods for CFD, Appendix B].
"""
function gauss_jacobi_nodes_and_weights(p, α, β, T=Float64::Type, tol=4*eps(T), maxit=1000)
    T = promote_type(typeof(α), typeof(β), T)
    nodes = Vector{T}(undef, p+1)
    weights = Vector{T}(undef, p+1)
    for j in 0:p
        x = -cospi(T(2j+1)/(2p+2))
        if j > 0
            x = (x + nodes[j]) / 2
        end
        for k in 1:maxit
            s = zero(T)
            for l in 1:j
                s += 1 / (x - nodes[l])
            end
            pol, der = jacobi_and_derivative(x, p+1, α, β)
            Δ = -pol / (der - s*pol)
            x = x + Δ
            abs(Δ) <= tol*abs(x) && break
        end
        pol, der = jacobi_and_derivative(x, p+1, α, β)
        nodes[j+1] = x
        weights[j+1] = 2^(α+β+1) * gamma(α+p+2) * gamma(β+p+2) /
                        (gamma(p+2) * gamma(α+β+p+2) * (1-x^2) * der^2)
    end

    return nodes, weights
end

"""
    jacobi_vandermonde(nodes, α, β)

Computes the Vandermonde matrix with respect to the Jacobi polynomials with
parameters `α`, `β` and the nodal basis on `nodes`.
The Vandermonde matrix `V` is the transformation matrix from the modal Jacobi
basis to the nodal Lagrange basis associated with `nodes`.
"""
function jacobi_vandermonde(nodes::AbstractVector, α, β)
    T = eltype(nodes)
    pp1 = length(nodes)
    V = Array{T}(undef, pp1, pp1)
    for j in 1:pp1, (i,x) in enumerate(nodes)
      V[i, j] = jacobi(x, j-1, α, β)
    end
    V
end

jacobi_vandermonde(basis::NodalBasis{Line}, α, β) = jacobi_vandermonde(basis.nodes, α, β)


function sum_n_below(factor, max)
    # Return the sum of all the natural numbers of factor below max
    max_nb = div(max-1,factor)
    ((factor / 2 * max_nb * (max_nb + 1)))
end

function n_below(n)
    # Return all the natural number of 3 and 5 multipliers
    return sum_n_below(3, n) + sum_n_below(5, n) - sum_n_below(3*5, n)
end

@assert n_below(10) == 23
println(n_below(1000))


module JuliaLambdaExample

# Calculate π using Monte Carlo simulation
function monte_carlo_pi(n)
    inside((x, y)) = x * x + y * y < 1.0
    points = rand(2, n)
    hits = count(inside(points[:, v]) for v in 1:n)
    return hits / n * 4
end

# Parsing input (this is more commonly structured as JSON string)
parse_data(event_data) = parse(Int, event_data)

# Formatting result (this is more commonly structured as JSON string)
format_result(value) = "$value"

# Define a handler function that is called by the lambda runtime
function handle_event(event_data, headers)
    @info "Handling request" event_data headers
    simulations = parse_data(event_data)
    my_pi = monte_carlo_pi(simulations)
    return format_result(my_pi)
end

end


"""
    newton(f::Function, x0::Real; niters=10^2, tol=1e-8)

Compute roots of scalar `f` by Newton's method

"""

function newton(f::Function, x0::Real; niters=10^2, tol=1e-8)
    x = xnew = x0
    for it in 1:niters
        df = ForwardDiff.derivative(f, x)
        xnew = x - f(x)/df
        abs(x - xnew) < tol && return xnew
        x, xnew = xnew, x
    end
    xnew
end


"""
Computes entropy of a sequence.
"""
function entropy(X :: AbstractVector)
  counts = Dict{eltype(X), Int}()
  n = length(X)
  for i = 1:length(X)
    counts[X[i]] = get(counts, X[i], 0) + 1
  end
  sum(kv -> log2(n / kv[2]) * kv[2], counts) / n
end

"""
Computes conditional entropy `H(Y | X)`.
"""
function conditional_entropy(X :: AbstractVector, Y :: AbstractVector)
  cnts_for_pref = Dict{eltype(X), Dict{eltype(Y), Int}}()
  n = length(X)
  for i = 1:length(X)
    if !haskey(cnts_for_pref, X[i])
      cnts_for_pref[X[i]] = Dict{eltype(Y), Int}()
    end
    counts = cnts_for_pref[X[i]]
    counts[Y[i]] = get(counts, Y[i], 0) + 1
  end

  # this helper function computes "unnormalized entropy" for given counts dict
  # i.e. entropy multiplied by sum of counts
  function _prefix_helper(counts :: Dict{eltype(Y), Int})
    counts_sum = sum(kv -> kv[2], counts)
    sum(kv -> log2(counts_sum / kv[2]) * kv[2], counts)
  end

  sum(kv -> _prefix_helper(kv[2]), cnts_for_pref) / n
end


#simple DFT function
function DFT(x)
    N = length(x)

    # We want two vectors here for real space (n) and frequency space (k)
    n = 0:N-1
    k = n'
    transform_matrix = exp.(-2im*pi*n*k/N)
    return transform_matrix*x

end

# Implementing the Cooley-Tukey Algorithm
function cooley_tukey(x)
    N = length(x)

    if (N > 2)
        x_odd = cooley_tukey(x[1:2:N])
        x_even = cooley_tukey(x[2:2:N])
    else
        x_odd = x[1]
        x_even = x[2]
    end
    n = 0:N-1
    half = div(N,2)
    factor = exp.(-2im*pi*n/N)
    return vcat(x_odd + x_even .* factor[1:half],
                x_odd - x_even .* factor[1:half])

end

function bitreverse(a::Array)
    # First, we need to find the necessary number of bits
    digits = convert(Int,ceil(log2(length(a))))

    indices = [i for i = 0:length(a)-1]

    bit_indices = []
    for i = 1:length(indices)
        push!(bit_indices, bits(indices[i]))
    end

    # Now stripping the unnecessary numbers
    for i = 1:length(bit_indices)
        bit_indices[i] = bit_indices[i][end-digits:end]
    end

    # Flipping the bits
    for i =1:length(bit_indices)
        bit_indices[i] = reverse(bit_indices[i])
    end

    #replacing indices
    for i = 1:length(indices)
        indices[i] = 0
        for j = 1:digits
            indices[i] += 2^(j-1) * parse(string(bit_indices[i][end-j]))
        end
       indices[i] += 1
    end

    b = [float(i) for i = 1:length(a)]
    for i = 1:length(indices)
        b[i] = a[indices[i]]
    end

    return b
end

function iterative_cooley_tukey(x)
    N = length(x)
    logN = convert(Int,ceil(log2(length(x))))
    bnum = div(N,2)
    stride = 0;

    x = bitreverse(x)

    z = [Complex(x[i]) for i = 1:length(x)]
    for i = 1:logN
       stride = div(N, bnum)
       for j = 0:bnum-1
           start_index = j*stride + 1
           y = butterfly(z[start_index:start_index + stride - 1])
           for k = 1:length(y)
               z[start_index+k-1] = y[k]
           end
       end
       bnum = div(bnum,2)
    end

    return z
end

function butterfly(x)
    N = length(x)
    half = div(N,2)
    n = [i for i = 0:N-1]
    half = div(N,2)
    factor = exp.(-2im*pi*n/N)

    y = [0 + 0.0im for i = 1:length(x)]

    for i = 1:half
        y[i] = x[i] + x[half+i]*factor[i]
        y[half+i] = x[i] - x[half+i]*factor[i]
    end

    return y
end

function approx(x, y)
    val = true
    for i = 1:length(x)
        if (abs(x[i]) - abs(y[i]) > 1e-5)
            val = false
        end
    end
    println(val)
end

function main()
    x = rand(128)
    y = cooley_tukey(x)
    z = iterative_cooley_tukey(x)
    w = fft(x)
    approx(y, w)
    approx(z, w)
end

main()


using TuringModels

Turing.setadbackend(:reversediff)

d = DataFrame(CSV.read(joinpath(@__DIR__, "..", "..", "data", "UCBadmit.csv"),
delim=';'));

size(d) |> display # Should be 12x5

# Turing model

@model m11_5(admit, applications) = begin
    N=length(applications)
    θ ~ truncated(Exponential(1), 0, Inf)
    α ~ Normal(0,2)

    for i ∈ 1:N
        prob = logistic(α)

        # alpha and beta for the BetaBinomial must be provided.
        # The two parameterizations are related by
        # alpha = prob * theta, and beta = (1-prob) * theta.
        # See https://github.com/rmcelreath/rethinking/blob/master/man/dbetabinom.Rd
        alpha = prob * θ
        beta = (1 - prob) * θ

        admit[i] ~ BetaBinomial(applications[i], alpha, beta)
    end
end

# Sample

chns = sample(m11_5(d[:, :admit],d[:, :applications]), Turing.NUTS(0.65), 1000);

# Result rethinking

m115rethinking = "
         mean   sd  5.5% 94.5% n_eff Rhat
theta  2.74 0.96  1.43  4.37  3583    1
a       -0.37 0.31 -0.87  0.12  3210    1
";

# Show summary

describe(chns)

# End of `11/m811.5t.jl`


using RootsNewton
using Test

#test to see if returns nothing for functions without real roots
f(x)=x^2+2
fp(x)=2*x
@test newtonroot(f,fp;iv=1.0, tolerance=1E-7, maxiter=1000).value==nothing

#Tests of functions with known solutions
f(x)=x-2
fp(x)=1
@test newtonroot(f,fp;iv=1.0, tolerance=1E-7, maxiter=1000).value ≈ 2

f(x)=x^3-1
fp(x)=3*x^2
@test newtonroot(f,fp;iv=1.0, tolerance=1E-7, maxiter=1000).value ≈ 1

f(x)=log(x)
fp(x)=1/x
@test newtonroot(f,fp;iv=1.0, tolerance=1E-7, maxiter=1000).value ≈ 1

f(x)=cos(x)
fp(x)=-sin(x)
@test newtonroot(f,fp;iv=1.0, tolerance=1E-7, maxiter=1000).value ≈ pi/2

#Tests using automatic differentiation

f(x)=x-2
@test newtonroot(f;iv=1.0, tolerance=1E-7, maxiter=1000).value ≈ 2

f(x)=x^3-1
@test newtonroot(f;iv=2.0, tolerance=1E-7, maxiter=1000).value ≈ 1

f(x)=log(x)
@test newtonroot(f;iv=2.0, tolerance=1E-7, maxiter=1000).value ≈ 1

f(x)=cos(x)
@test newtonroot(f;iv=1.0, tolerance=1E-7, maxiter=1000).value ≈ pi/2

#Tests with BigFloat
f(x)=x-2
@test newtonroot(f;iv=BigFloat(1.0), tolerance=1E-7, maxiter=1000).value ≈ 2

#Tests that maxiter is working
f(x)=x^3-1
@test newtonroot(f;iv=100, tolerance=1E-7, maxiter=3).iter <= 3

#Tests that tolerance is working
f(x)=x^3-1
@test 1-2 <= newtonroot(f;iv=10, tolerance=2, maxiter=1000).value <= 1+2 || newtonroot(f;iv=10, tolerance=2, maxiter=1000).iter==1000

#@testset "RootsNewton.jl" begin
    # Write your own tests here.
#end


"""
barychebinit(n)

Get Chebyshev nodes of first kind and corresponding barycentric Lagrange interpolation weights.
Reference: Berrut, J.P. and Trefethen, L.N., 2004. Barycentric lagrange interpolation. SIAM review, 46(3), pp.501-517.

# Arguments
- `n`: order of the Chebyshev interpolation

# Returns
- Chebyshev nodes
- Barycentric Lagrange interpolation weights
"""
function barychebinit(n)
    x = zeros(Float64, n)
    w = similar(x)
    for i in 1:n
        c = (2i - 1)π / (2n)
        x[n - i + 1] = cos(c)
        w[n - i + 1] = (-1)^(i - 1) * sin(c)
    end
    return x, w
end

"""
function barycheb(n, x, f, wc, xc)

Barycentric Lagrange interpolation at Chebyshev nodes
Reference: Berrut, J.P. and Trefethen, L.N., 2004. Barycentric lagrange interpolation. SIAM review, 46(3), pp.501-517.

# Arguments
- `n`: order of the Chebyshev interpolation
- `x`: coordinate to interpolate
- `f`: array of size n, function at the Chebyshev nodes
- `wc`: array of size n, Barycentric Lagrange interpolation weights
- `xc`: array of size n, coordinates of Chebyshev nodes

# Returns
- Interpolation result
"""
function barycheb(n, x, f, wc, xc)
    for j in 1:n
        if x == xc[j]
            return f[j]
        end
    end

    num, den = 0.0, 0.0
    for j in 1:n
        q = wc[j] / (x - xc[j])
        num += q * f[j]
        den += q
    end
    return num / den
end

function barycheb2(n, x, f, wc, xc)
    for j in 1:n
        if x == xc[j]
            return f[j, :]
        end
    end

    den = 0.0
    num = zeros(eltype(f), size(f)[2])
    for j in 1:n
        q = wc[j] / (x - xc[j])
        num += q .* f[j, :]
        den += q
    end
    return num ./ den
end


using Polynomials: Polynomial,fit,coeffs
using Plots
using Distributions: Normal
using Random
using Statistics: mean, std

p = Polynomial([1,2,3])
@show coeffs(p)

errstd = 3
errNormal = Normal(0,errstd)

x = 0:10
y = 2 .+ 3 .* x .+rand(errNormal,11)

scatter(x,y,yerror=errstd,label=false)

fit_line = fit(x,y,1)
@show coeffs(fit_line)
plot!(x,fit_line.(x),label=false)

b_list = []
m_list = []
N = 1000
for i in 1:N
    y = 2 .+ 3 .* x .+rand(errNormal,11)
    fit_c = coeffs(fit(x,y,1))
    push!(b_list,fit_c[1])
    push!(m_list,fit_c[2])
end

@show mean(b_list)
@show mean(m_list)
@show std(b_list)
@show std(m_list)

histogram(b_list,bins=30)
histogram(m_list,bins=30)


#=~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~=#
# Problem set 7 solutions
# Written by Tyler Ransom
# Commented by Giuseppe Grasso
#=~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~=#

using Random
using LinearAlgebra
using Statistics
using Distributions
using Optim
using DataFrames
using CSV
using HTTP
using GLM
using FreqTables
using ForwardDiff
using LineSearches
using SMM
cd("/Users/peppegrass/Documents/GitHub/fall-2020/ProblemSets/PS6-ccp/") # GG: sets working directory
pwd() ## GG: prints working directory
readdir() # GG: equivalent to -ls- to see elements of working directory

function allwrap()

    #=
    GG: In this problem set, we will practice estimating models by
    Generalized Method of Moments (GMM) and Simulated Method of Moments (SMM).
    =#

    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # Question 1 - GMM for simple regression model
    # data & model are SAME AS PS2, QUESTION 2
    #
    # 1. Estimate the linear regression model from Question 2 of Problem Set 2 by GMM.
    # Write down the moment function as in slide #8 of the Lecture 9 slide deck and use Optim for estimation.
    # Use the N × N Identity matrix as your weighting matrix. Check your answer using the closed-form matrix formula for the OLS estimator.
    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    url = "https://raw.githubusercontent.com/OU-PhD-Econometrics/fall-2020/master/ProblemSets/PS1-julia-intro/nlsw88.csv"
    df = CSV.read(HTTP.get(url).body)
    X = [ones(size(df,1),1) df.age df.race.==1 df.collgrad.==1]
    y = log.(df.wage)

    function ols_gmm(β, X, y) # GG: GMM function
        K = size(X,2) # GG: don't really needed for GMM, unless for computing RMSE in step below
        N = size(y,1) # GG: don't really needed for GMM, unless for computing RMSE in step below
        ŷ = X*β
        g = y .- ŷ # GG: force y and y_hat to be as close as possible
        J = g'*I*g # GG: GMM objective function
        return J
    end

    # we can also estimate the RMSE of the regression by ADDING ANOTHER MOMENT (and assuming homoskedasticity)
    function ols_gmm_with_σ(θ, X, y)
        K = size(X,2)
        N = size(y,1)
        β = θ[1:end-1]
        σ = θ[end]
        g = y .- X*β
        g = vcat(g,( (N-1)/(N-K) )*var(g) .- σ^2) # this builds on the assumption that MSE: SSR/(N-K) = σ² # GG: performing DOFs corrections
        J = g'*I*g
        return J
    end

    β_hat_gmm = optimize(b -> ols_gmm(b, X, y), rand(size(X,2)), LBFGS(), Optim.Options(g_tol=1e-6, iterations=100_000, show_trace=true))
    println(β_hat_gmm.minimizer)
    println("ans: ",X\y) # GG: X\y gives the least squares; cool uh!?
    β_hat_gmm = optimize(b -> ols_gmm_with_σ(b, X, y), rand(size(X,2)+1), LBFGS(), Optim.Options(g_tol=1e-6, iterations=100_000, show_trace=true))
    println(β_hat_gmm.minimizer)
    println("ans: ",vcat(X\y,sqrt( sum( (y .- X*β_hat_gmm.minimizer[1:end-1]).^2 )/( size(y,1)-size(X,2) ) )))

    # GG: of course ols_gmm(b, X, y) and ols_gmm_with_σ(b, X, y) yield same results but latter with MSE

    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # Question 2 - GMM for multinomial logit model
    # data & model are SAME AS PS2, QUESTION 5
    #
    # 2. Estimate the multinomial logit model from Question 5 of Problem Set 2 by the following means:
    #   (a) Maximum likelihood (i.e. re-run your code [or mine] from Question 5 of Problem Set 2)
    #   (b) GMM with the MLE estimates as starting values. Your g object should be a vector of dimension N ×J where
    #   N is the number of rows of the X matrix and J is the dimension of the choice set. Each element, g should equal d − P,
    #   where d and P are “stacked” vectors of dimension N×J
    #   (c) GMM with random starting values
    # Compare your estimates from part (b) and (c). Is the objective function globally concave?
    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    df = dropmissing(df, :occupation)
    df[df.occupation.==8 ,:occupation] .= 7
    df[df.occupation.==9 ,:occupation] .= 7
    df[df.occupation.==10,:occupation] .= 7
    df[df.occupation.==11,:occupation] .= 7
    df[df.occupation.==12,:occupation] .= 7
    df[df.occupation.==13,:occupation] .= 7

    #df = @transform(df, white = race.==1)
    df[:,:white] = df.race .== 1
    X = [ones(size(df,1),1) df.age df.white df.collgrad]
    y = df.occupation

    # get starting values for mlogit from series of binary logits
    svals = zeros(size(X,2),7)
    for j=1:7
        tempname = Symbol(string("occ",j))
        df[:,tempname] = df.occupation.==j
    end
    # GG: this is to derive some "reasonable starting values" for a multinomial logit model
    # GG: what do we mean for "reasonable starting values". You don't wanna have neither a vector of 0 nor a vector of random numbers
    # GG: So what starting values for a Multinomial logit. Estimate a bunch of binary logits and then...
    svals[:,1] = coef(lm(@formula(occ1 ~ age + white + collgrad), df))
    svals[:,2] = coef(lm(@formula(occ2 ~ age + white + collgrad), df))
    svals[:,3] = coef(lm(@formula(occ3 ~ age + white + collgrad), df))
    svals[:,4] = coef(lm(@formula(occ4 ~ age + white + collgrad), df))
    svals[:,5] = coef(lm(@formula(occ5 ~ age + white + collgrad), df))
    svals[:,6] = coef(lm(@formula(occ6 ~ age + white + collgrad), df))
    svals[:,7] = coef(lm(@formula(occ7 ~ age + white + collgrad), df))
    svals = svals[:,1:6] .- svals[:,7] # GG: ... impose the identification normalization (i.e. you subtract for the first 6 estimates the one of the 7th)
    svals = svals[:] # GG: then vectorize it and use it as a starting value

    function mlogit_mle(α, X, y)

        K = size(X,2)
        J = length(unique(y))
        N = length(y)
        bigY = zeros(N,J)
        for j=1:J
            bigY[:,j] = y.==j
        end
        bigα = [reshape(α,K,J-1) zeros(K)]

        P = exp.(X*bigα) ./ sum.(eachrow(exp.(X*bigα)))

        loglike = -sum( bigY.*log.(P) ) # GG: MLE objective function

        return loglike
    end

    α_zero = zeros(6*size(X,2))
    α_rand = rand(6*size(X,2))
    # α_true = [-.5120867,-.6950393,.3165616,-1.646016,-2.300697,-.7656569]
    α_true = [.1910213,-.0335262,.5963968,.4165052,-.1698368,-.0359784,1.30684,-.430997,.6894727,-.0104578,.5231634,-1.492475,-2.26748,-.0053001,1.391402,-.9849661,-1.398468,-.0142969,-.0176531,-1.495123,.2454891,-.0067267,-.5382892,-3.78975]
    α_start = α_true.*rand(size(α_true))
    println(size(α_true))
    α_hat_optim = optimize(a -> mlogit_mle(a, X, y), svals, LBFGS(), Optim.Options(g_tol = 1e-5, iterations=100_000, show_trace=true, show_every=50))
    α_hat_mle = α_hat_optim.minimizer
    println("mlogit MLE estimates: ",α_hat_mle)

    function mlogit_gmm(α, X, y)
        K = size(X,2)
        J = length(unique(y))
        N = length(y)
        bigY = zeros(N,J)
        for j=1:J
            bigY[:,j] = y.==j
        end
        bigα = [reshape(α,K,J-1) zeros(K)]

        P = exp.(X*bigα) ./ sum.(eachrow(exp.(X*bigα)))

        g = zeros((J-1)*K)
        for j=1:(J-1)
            for k = 1:K
                g[(j-1)*K+k] = mean( (bigY[:,j] .- P[:,j]) .* X[:,k] ) # GG: GMM (exactly-identified) objective function
                # GG: (j-1)*K moment conditions; (j-1) equations and K betas; same number of elements as number non-zero elements in bigαlpha
            end
        end
        J = N*g⋅g #GG: exactly identified; no weighting matrix
        return J
    end

    function mlogit_gmm_overid(α, X, y)
        K = size(X,2)
        J = length(unique(y))
        N = length(y)
        bigY = zeros(N,J)
        for j=1:J
            bigY[:,j] = y.==j
        end
        bigα = [reshape(α,K,J-1) zeros(K)]

        P = exp.(X*bigα) ./ sum.(eachrow(exp.(X*bigα))) # GG: same as before P = exp.(X*bigα) ./ sum.(eachrow(exp.(X*bigα)))

        g = bigY[:] .- P[:] # GG: now in over-identified case; we vectorize big Y and P with the [:] operator and subtract the two
        # GG: this will be equivalent to NONLINEAR LEAST SQUARES

        J = g'*I*g # GG: now we have weighting matrix
        return J
    end
    #α_zero = zeros(6*size(X,2))
    #α_rand = rand(6*size(X,2))
    α_true = [.1910213,-.0335262,.5963968,.4165052,-.1698368,-.0359784,1.30684,-.430997,.6894727,-.0104578,.5231634,-1.492475,-2.26748,-.0053001,1.391402,-.9849661,-1.398468,-.0142969,-.0176531,-1.495123,.2454891,-.0067267,-.5382892,-3.78975]
    #α_start = α_true.*rand(size(α_true))
    #println("starting values: ",α_rand)

    # GG: here we're playing around with different optimization algorithms, starting values, and later using autodiff

    # LBFGS
    α_hat_optim = optimize(a -> mlogit_gmm_overid(a, X, y), α_true .+ .0001*rand(size(α_true)), LBFGS(), Optim.Options(g_tol = 1e-5, iterations=1_000, show_trace=true, show_every=100))
    α_hat_gmm = α_hat_optim.minimizer
    println("LBFGS GMM minimum: ",α_hat_optim.minimum)
    println("LBFGS GMM estimates: ",α_hat_gmm)

    # use autodiff with LBFGS
    td = TwiceDifferentiable(b -> mlogit_gmm_overid(b, X, y), svals; autodiff = :forward)
    α_hat_optim_ad = optimize(td, α_true .+ .0001*rand(size(α_true)), LBFGS(), Optim.Options(g_tol = 1e-5, iterations=10_000, show_trace=true, show_every=100))
    α_hat_gmm_ad = α_hat_optim_ad.minimizer
    println("LBFGS GMM autodiff minimum: ",α_hat_optim_ad.minimum)
    println("LBFGS GMM autodiff estimates: ",α_hat_gmm_ad)

    # use autodiff with Newton
    α_hat_optim_ad_newton = optimize(td, α_true .+ .0001*rand(size(α_true)), Newton(linesearch = BackTracking()), Optim.Options(g_tol = 1e-5, iterations=10_000, show_trace=true, show_every=100))
    α_hat_gmm_ad_newton = α_hat_optim_ad_newton.minimizer
    println("Newton GMM autodiff minimum: ",α_hat_optim_ad_newton.minimum)
    println("Newton GMM autodiff estimates: ",α_hat_gmm_ad_newton)

    # autodiff from bad starting values
    α_hat_optim_ad_fr_bad = optimize(td, α_rand.*sign.(α_true), LBFGS(), Optim.Options(g_tol = 1e-10, f_tol = 1e-10, x_tol = 1e-10, iterations=10_000, show_trace=true, show_every=50))
    α_hat_gmm_ad_fr_bad = α_hat_optim_ad_fr_bad.minimizer
    println("LBFGS GMM autodiff minimum: ",α_hat_optim_ad_fr_bad.minimum)
    println("LBFGS GMM autodiff estimates: ",α_hat_gmm_ad_fr_bad)

    # GG: putting different answers in a DataFrame and comparing them
    # GG: found out close answers; the one using bad starting value didn't get close to other solutions; some times GMM tolerance values needed to be played around with
    compare = DataFrame(mle = α_hat_mle, gmm = α_hat_gmm, gmm_ad = α_hat_gmm_ad, gmm_ad_newton = α_hat_gmm_ad_newton, gmm_ad_fr_bad = α_hat_gmm_ad_fr_bad)
    println(compare)
    # GG: objective function globally concave, but couldn't get to global maximum from random starting values, although pretty close

    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # Question 3 - SIMULATE DATA from a logit model
    #
    # 3. SIMULATE A DATASET FROM A MULTINOMIAL LOGIT MODEL, and then estimate its parameter values and verify that the estimates are close to the
    # parameter values you set. That is, for a given sample size N, choice set dimension J and parameter vector β, write a function that outputs
    # data X and Y. I will let you choose N, J, β and the number of covariates in X (K), but J should be larger than 2 and K should be larger than 1.
    # If you haven’t done this before, you may want to follow these steps:
    #   (a) Generate X using a random number generator—rand() or randn().
    #   (b) Set values for β such that conformability with X and J is satisﬁed
    #   (c) Generate the N×J matrix of choice probabilities P
    #   (d) Draw the preference shocks ε as a N×1 vector of U[0,1] random numbers
    #   (e) Generate Y as follows:
    #     • Initialize Y as an N ×1 vector of 0s
    #     • Update Y i = ∑ j=1 J 1 P ik > εi  ∑k=j J  [see text of problem set]
    #   (f) An alternative way to generate choices would be to draw a N × J matrix of ε’s from a T1EV distribution.
    #   This distribution is already deﬁned in the Distributions package. Then Y i = argmax j X i β j + ε i j .
    #   I’ll show you an example of how to do that in the solutions code for this problem set.
    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    function sim_logit(N=100_000,J=4) # putting "=" after the arguments here sets default values for these arguments
        # generate X, β and P
        X = hcat(ones(N),randn(N),2 .+ 2 .* randn(N))
        if J==4
            β = hcat([1, -1, 0.5],[-2, 0.5, 0.3],[0, -0.5, 2],zeros(3)) # GG: numbers that square nicely with 4
        else
            β = 2 .* rand(size(X,2),J) .- 1
        end
        P = exp.(X*β) ./ sum.(eachrow(exp.(X*β))) # GG: usual one-liner to generate multinomial ogit choice probabilities; eachrow is summation
        # draw choices
        draw = rand(N)
        Y = zeros(N) # Initializing Y as an Nx1 vector of 0s
        for j=1:J
            Ytemp = sum(P[:,j:J];dims=2) .> draw # GG: then updating it (summing across j) based on whether choice probability is above or below draw [see formula of problem set]
            Y += Ytemp # GG: for each iteration of the loop adding Ytemp to Y
        end
        return Y,X
    end
    ySim,XSim = sim_logit() # GG: it will use "default values" for N and J set in line 244 through s"et to the default operator =""
    α_hat_sim_optim = optimize(a -> mlogit_mle(a, XSim, ySim), rand(9), LBFGS(), Optim.Options(g_tol = 1e-5, iterations=100_000, show_trace=true, show_every=50))
    α_hat_sim_mle = α_hat_sim_optim.minimizer
    println("simulated mlogit MLE estimates: ",α_hat_sim_mle) # GG: and here are our simulated estimates; values pretty close because sample size big and SEs low

    # now with drawing from Gumbel instead (#GG: i.e T1EV)
    # GG: python does not have the nice Distributions package as in Julia
    function sim_logit_w_gumbel(N=100_000,J=4) # putting = here sets default values for these arguments
        # generate X and β (No P needed)
        X = hcat(ones(N),randn(N),2 .+ 2 .* randn(N))
        if J==4
            β = hcat([1, -1, 0.5],[-2, 0.5, 0.3],[0, -0.5, 2],zeros(3))
        else
            β = 2 .* rand(size(X,2),J) .- 1
        end
        # draw choices
        ϵ = rand(Gumbel(0,1),N,J) # GG: using Julia's Distributions package; Gumbel(0,1) is "standard" T1EV
        Y = argmax.(eachrow(X*β .+ ϵ)) # GG: getting argmax of eachrow; eachrow allows avoiding to have to write a loop; tells Julia to do argmax computation row-wise and find column j maximizer
        # GG: that's just another way of deriving the Y; the loop in the previous function is just as fine
        return Y,X
    end
    ySim,XSim = sim_logit_w_gumbel()
    α_hat_sim_optim = optimize(a -> mlogit_mle(a, XSim, ySim), rand(9), LBFGS(), Optim.Options(g_tol = 1e-5, iterations=100_000, show_trace=true, show_every=50))
    α_hat_sim_mle = α_hat_sim_optim.minimizer
    println("simulated mlogit MLE estimates: ",α_hat_sim_mle)

    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # Question 4 - run SMM.parallelNormal()
    #
    # 4. Use SMM.jl to run the example code on slide #21 of the Lecture 9 slide deck.
    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # GG: SMM.jl has a nice count-down that Optim does not have
    MA = SMM.parallelNormal()
    dc = SMM.history(MA.chains[1])
    dc = dc[dc[:accepted].==true, :]
    println(describe(dc)) # GG: can verify you get the mean of the MVN distribution

    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # Question 5 - estimate mlogit by SMM
    # GG: Wanna use SIMULATED METHOD OF MOMENTS (using code of Q3) to ESTIMATE RESULT OF Q2
    #
    # 5. Use your code from Question 3 to estimate the multinomial logit model from Question 2 using SMM and the code example from slide #18 of
    # the Lecture 9 slide deck.
    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # GG: X and y of Q2
    X = [ones(size(df,1),1) df.age df.white df.collgrad]
    y = df.occupation
    # GG: chunk below was commented out by Tyler beacuse the set of moment conditions used were just not working out for him
    #function mlogit_smm(α, X, y, D)
    #    K = size(X,2)
    #    J = length(unique(y))
    #    N = length(y)
    #    bigY = zeros(N,J)
    #    bigỸ = zeros(N,J)
    #    for j=1:J
    #        bigY[:,j] = y.==j
    #    end
    #    bigα = [reshape(α,K,J-1) zeros(K)]

    #    T = promote_type(eltype(α),eltype(X))
    #    gmodel = zeros(T,(J-1)*K,D)
    #    gdata  = zeros(T,(J-1)*K)

    #    # data moments
    #    for j=1:(J-1)
    #        gdata[(j-1)*K+1] = mean( (bigY[:,j]) )
    #        for k = 2:K
    #            gdata[(j-1)*K+k] = cov( bigY[:,j], X[:,k] )
    #        end
    #    end

    #    Random.seed!(1234)
    #    # simulated model moments
    #    for d=1:D
    #        # draw choices
    #        ε = rand(Gumbel(0,1),N,J)
    #        ỹ = argmax.(eachrow(X*bigα .+ ε))
    #        for j=1:J
    #            bigỸ[:,j] = ỹ.==j
    #        end
    #        # data moments
    #        for j=1:(J-1)
    #            gmodel[(j-1)*K+1,d] = mean( (bigỸ[:,j]) )
    #            for k = 2:K
    #                gmodel[(j-1)*K+k,d] = cov( bigỸ[:,j], X[:,k] )
    #            end
    #        end
    #    end

    #    # minimize squared difference between data and moments
    #    g = mean(gmodel .- gdata; dims=2)
    #    J = N*g⋅g
    #    return J
    #

    # GG: the function below for the SIMULATED METHOD OF MOMENTS is the one from the SLIDES OF LECTURE 9 (mind the SCROLLER; gotta use HTML and not PDF version)
    function mlogit_smm_overid(α, X, y, D) # GG: D is the number of draws
        K = size(X,2)
        J = length(unique(y))
        N = length(y)
        bigY = zeros(N,J)
        bigỸ = zeros(N,J)
        for j=1:J
            bigY[:,j] = y.==j # GG: bigY is our real data that we're trying to MATCH
        end
        bigα = [reshape(α,K,J-1) zeros(K)] # GG: bigα is our parameter vector that we're gonna try to guess value of, so that we get values so that simulated data match actual ones as closely as possible

        Random.seed!(1234) # GG: cardinal rule, set up the seed; don't forget that!
        for d=1:D
            # draw choices
            ε = rand(Gumbel(0,1),N,J) # GG: if we change the distribution of ε to a multivariate normal, we can do MNP without having to do any Monte Carlo integration. We don't need to have a closed form solution for this.
            # GG: and if we want to do GEV models, we need to specify nest-specific and alternative-specific component; just two errors but same principle.
            ỹ = argmax.(eachrow(X*bigα .+ ε)) # GG: we're goona draw ỹ through ε
            for j=1:J
                bigỸ[:,j] .+= (ỹ.==j)*(1/D) # GG: bigỹ is our SIMULATED data
                # GG: we set the jth column of bigỸ to a dummy if ỹ is equal to j
            end
        end

        g = bigY[:] .- bigỸ[:] # GG: exactly like g in Question 2; very same principle

        J = g'*I*g
        return J
    end
    println(size(X,2))
    println(length(unique(y))-1)
    println("size of svals: ",size(α_true .+ .0001*rand(size(α_true))))
    bigα = [reshape(α_true,size(X,2),length(unique(y))-1) zeros(size(X,2))]
    println("size of bigα: ",size(bigα))
    td = TwiceDifferentiable(th -> mlogit_smm_overid(th, X, y, 2_000), α_true .+ .0001*rand(size(α_true)); autodiff = :forward)
    α_hat_smm_overid = optimize(td, α_true .+ .0001*rand(size(α_true)), LBFGS(), Optim.Options(g_tol=1e-8, x_tol=1e-8, f_tol=1e-8, iterations=100_000, show_trace=true, show_every=5))
    println(α_hat_smm_overid.minimizer)
    return nothing
end

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# Question 6
#:::::::::::::::::::::::::::::::::::::::::::::::::::
@time allwrap()


# A variety of means

# Geometric mean
function geomean(a::RealArray)
    s = 0.0
    n = length(a)
    for i = 1 : n
        @inbounds s += log(a[i])
    end
    return exp(s / n)
end

# Harmonic mean
function harmmean(a::RealArray)
    s = 0.0
    n = length(a)
    for i in 1 : n
        @inbounds s += inv(a[i])
    end
    return n / s
end

# Trimmed mean

function trimmean(x::RealArray, p::FloatingPoint)
    n = length(x)
    n > 0 || error("x can not be empty.")
    rn = min(iround(n * p), n-1)
    if rn == n - 1
        return median(x)
    else
        sx = sort(x)
        nl = rn >> 1
        nh = (rn - nl)
        s = 0.0
        @inbounds for i = (1+nl) : (n-nh)
            s += x[i]
        end
        return s / (n - rn)
    end
end

# Weighted mean
function wmean{T<:Number,W<:Real}(v::AbstractArray{T}, w::AbstractArray{W}; wsum::W=NaN)
	# wsum is the pre-computed sum of weights

	n = length(v)
	length(w) == n || raise_dimerror()
    sv = zero(T)

    if isnan(wsum)
    	sw = zero(W)
    	for i = 1 : n
    		@inbounds wi = w[i]
        	@inbounds sv += v[i] * wi
        	sw += wi
    	end
    else
    	for i = 1 : n
    		@inbounds sv += v[i] * w[i]
    	end
    	sw = wsum
    end

    return sv / sw
end

# Weighted mean (fast version for contiguous arrays using BLAS)
function wmean{T<:BlasReal}(v::Array{T}, w::Array{T}; wsum::T=NaN)
	sw = (isnan(wsum) ? sum(w) : wsum)
	dot(v, w) / sw
end

mean{T<:Number}(v::AbstractArray{T}, w::WeightVec) = wmean(v, values(w); wsum=sum(w))


@doc raw"""
    ExponentiatedKernel([a=1])

The exponentiated kernel is a Mercer kernel given by:

```math
\kappa(\mathbf{x},\mathbf{y}) = \exp\left(a \mathbf{x}^\intercal \mathbf{y} \right)
\qquad a > 0
```
"""
struct ExponentiatedKernel{T<:AbstractFloat} <: MercerKernel{T}
    alpha::HyperParameter{T}
    ExponentiatedKernel{T}(α::Real) where {T<:AbstractFloat} = new{T}(
        HyperParameter(convert(T,α), interval(OpenBound(zero(T)), nothing))
    )
end
ExponentiatedKernel(α::T1 = 1.0) where {T1<:Real} = ExponentiatedKernel{floattype(T1)}(α)

@inline exponentiatedkernel(z::T, α::T) where {T<:AbstractFloat} = exp(α*z)

@inline basefunction(::ExponentiatedKernel) = ScalarProduct()
@inline kappa(κ::ExponentiatedKernel{T}, z::T) where {T} = exponentiatedkernel(z, getvalue(κ.alpha))


"""
    M = time_to_mean(t, T, Mo=0., epoch=0.)

Computes the mean anomaly of an orbit with period `T` from a time `t`, the orbit's mean anomaly
at epoch `Mo`, and the epoch.
"""
function time_to_mean(t, T, Mo=0., epoch=0.)
    return Mo + 2π*(t - epoch)/T
end

function time_to_mean(t, orb::Orbit)
    M = time_to_mean(t, orb.period, orb.Mo, orb.epoch)
    if orb.e < 1
        return wrap_angle(M)
    else
        return M
    end
end

"""
    t = mean_to_time(M, T, Mo=0, epoch=0, tmin=nothing)

Computes the time at which an orbit with period `T`, mean anomaly
at epoch `Mo`, and epoch `epoch` has the specified mean anomaly `M`.

Because this time is not uniquely defined for elliptical orbits, `tmin` can be specified
to ensure that the returned time falls between `tmin` and `tmin + T`.
"""
function mean_to_time(M, T, Mo=0, epoch=0, tmin=nothing)
    t = epoch + (M-Mo)*T/(2π)
    if tmin === nothing
        return t
    else
        return t + T*ceil((tmin-t)/T)
    end
end

mean_to_time(M, orb::Orbit, tmin=nothing) = mean_to_time(M, orb.period, orb.Mo, orb.epoch, orb.e<1 ? tmin : nothing)

"""
    θ = mean_to_true(M, e)

Computes the true anomaly of an orbit with eccentricity `e` at the specified mean anomaly `M`.
"""
function mean_to_true(M, e)
    if e == 1       # Parabolic case
        throw(ArgumentError("Parabolic case (e=1) not implemented"))
    elseif e < 1    # Elliptical case
        E = kepler(M, e)
        return 2*atan(√(1+e)*sin(E/2),  √(1-e)*cos(E/2))
    else            # Hyperbolic case
        H = kepler(M, e)
        return 2*atan(√(e+1)*sinh(H/2), √(e-1)*cosh(H/2))
    end
end

"""
    θ = time_to_true(t, e, T, Mo=0., epoch=0.)

Computes the mean anomaly of an orbit with eccentricity `e` and period `T` from a time `t`, the orbit's mean anomaly
at epoch `Mo`, and the epoch.
"""
function time_to_true(t, e, T, Mo=0, epoch=0)
    M = time_to_mean(t, T, Mo, epoch)
    return mean_to_true(M, e)
end

time_to_true(t, orb::Orbit) = time_to_true(t, orb.e, orb.period, orb.Mo, orb.epoch)

"""
    M = true_to_mean(θ, e)

Computes the mean anomaly of an orbit with eccentricity `e` at the specified true anomaly `θ`.
"""
function true_to_mean(θ, e)
    if e == 1       # Parabolic case
        throw(ArgumentError("Parabolic case (e=1) not implemented"))
    elseif e < 1    # Elliptical case
        E = 2*atan(sin(θ/2)*√(1-e), cos(θ/2)*√(1+e))
        return wrap_angle(E - e*sin(E))
    else            # Hyperbolic case
        H = 2*atanh(tan(θ/2)*√((e-1)/(e+1)))
        # H = copysign(acosh((e+cos(θ))/(1+e*cos(θ))), θ)
        return e*sinh(H) - H
    end
end

true_to_mean(θ, orb::Orbit) = true_to_mean(θ, orb.e)

"""
    t = true_to_time(θ, e, T, Mo=0, epoch=0, tmin=nothing)

Computes the time at which an orbit with eccentricity `e`, period `T`, mean anomaly
at epoch `Mo`, and epoch `epoch` has the specified true anomaly `θ`.

Because this time is not uniquely defined for elliptical orbits, `tmin` can be specified
to ensure that the returned time falls between `tmin` and `tmin + T`.
"""
function true_to_time(θ, e, T, Mo=0, epoch=0, tmin=nothing)
    M = true_to_mean(θ, e)
    return mean_to_time(M, T, Mo, epoch, e<1 ? tmin : nothing)
end

# true_to_time(θ, e, a, μ, Mo=0, epoch=0, tmin=nothing) = true_to_time(θ, e, period(a, μ), Mo, epoch, tmin)
true_to_time(θ, orb::Orbit, tmin=nothing) = true_to_time(θ, orb.e, orb.period, orb.Mo, orb.epoch, tmin)


@testset "Gaussian Quadrature" begin
    # NOTE: It is assumed p = 1:5 is sufficient for testing
    @testset "RefLine" begin
        for T ∈ Floats
            for p = 1:5
                weights, points = gauss_quadrature(Val(:legendre), RefLine(), Val(p), T)
                @test abs(sum(weights) - 1) < 10*eps(T)
                # Test the integration of a polynomial of degree 2p-1, which should
                # be exact for p integration points.
                # 1
                # ∫xᵖdx = 1/(p+1) ≈ ∑ wᵢxᵢᵖ
                # 0
                exact = 1/big(p+1)
                approx = zero(T)
                for i = 1:p
                    approx += weights[i]*(points[i][1]^p)
                end
                rel_err = T((approx - exact)/exact)
                @test rel_err < 10*eps(T)
            end
        end
    end

    @testset "RefSquare" begin
        for T ∈ Floats
            for p = 1:5
                weights, points = gauss_quadrature(Val(:legendre), RefSquare(), Val(p), T)
                @test abs(sum(weights) - 1) < 10*eps(T)
                #
                # q = 2p-1
                #
                # 11                                   5 + (-1)^q + 2q
                # ∫∫(r^q + s^q + (1 - r - s)^q ds dr = --------------
                # 00                                     (q+1)(q+2)
                q = 2p-1
                exact = big((5 + (-1)^q + 2q))/big((q+1)*(q+2))
                approx = zero(T)
                for i = 1:p^2
                    r, s = points[i]
                    approx += weights[i]*(r^q + s^q + (1 - r - s)^q)
                end
                rel_err = T((approx - exact)/exact)
                @test rel_err < 10*eps(T)
            end
        end
    end

    @testset "RefCube" begin
        for T ∈ Floats
            for p = 1:5
                weights, points = gauss_quadrature(Val(:legendre), RefCube(), Val(p), T)
                @test abs(sum(weights) - 1) < 10*eps(T)
                #
                # q = 2p-1
                #
                # 111
                # ∫∫∫(r^q + s^q + t^q) dt ds dr = 3/(q+1)
                # 000
                q = 2p-1
                exact = 3/big(q + 1)
                approx = zero(T)
                for i = 1:p^3
                    r, s, t = points[i]
                    approx += weights[i]*(r^q + s^q + t^q)
                end
                rel_err = T((approx - exact)/exact)
                @test rel_err < 10*eps(T)
            end
        end
    end

    @testset  "RefTriangle" begin
        for T ∈ Floats
            for p = 1:5
                weights, points = gauss_quadrature(Val(:legendre), RefTriangle(), Val(p), T)
                @test abs(sum(weights) - 1//2) < 10*eps(T)
                # This quadrature is only exact to polynomial degree p, not 2p-1!
                #
                # 1 1-r
                # ∫  ∫ (r^p + s^p + (1 - r - s)^p) ds dr = 3/((p+1)(p+2))
                # 0  0
                #
                exact = 3/big((p+1)*(p+2))
                approx = zero(T)
                for i = 1:length(weights)
                    r, s = points[i]
                    approx += weights[i]*(r^p + s^p + (1 - r - s)^p)
                end
                rel_err = T((approx - exact)/exact)
                @test rel_err < 10*eps(T)
            end
        end
    end
end


using PyPlot
using OffsetArrays

function bary_grid(n, a)
    x = OffsetArray{Float64}(undef, 0:n, 0:n)
    y = OffsetArray{Float64}(undef, 0:n, 0:n)
    pts = OffsetArray(range(0, 1, length=n+1), 0:n)
    for i = 0:n
        ξ₁ = pts[i]
        for j = 0:n-i
            ξ₂ = pts[j]
            ξ₃ = 1 - ξ₁ - ξ₂
            x[i,j] = ξ₁ * a[1,1] + ξ₂ * a[1,2] + ξ₃ * a[1,3]
            y[i,j] = ξ₁ * a[2,1] + ξ₂ * a[2,2] + ξ₃ * a[2,3]
        end
        for j = n-i+1:n
            x[i,j] = NaN
            y[i,j] = NaN
        end
    end
    return x, y
end

function bary_func(f, n, a)
    z = OffsetArray{Float64}(undef, 0:n, 0:n)
    pts = OffsetArray(range(0, 1, length=n+1), 0:n)
    ξ = Vector{Float64}(undef, 3)
    for i = 0:n
        ξ[1] = pts[i]
        for j = 0:n-i
            ξ[2] = pts[j]
            ξ[3] = 1 - ξ[1] - ξ[2]
            z[i,j] = f(ξ)
        end
    end
    return z
end

function bary_plot(x, y, z)
    n = size(x, 1) - 1
    xd = [ x[i,n-i] for i = 0:n ]
    yd = [ y[i,n-i] for i = 0:n ]
    zd = [ z[i,n-i] for i = 0:n ]
    idx = [1, 2, 3, 1]
    plot3D(a[1,idx], a[2,idx], zeros(4), "k")
    for i = 0:n
        plot3D(x[i,0:n-i], y[i,0:n-i], z[i,0:n-i], "c")
    end
    for j = 0:n
        plot3D(x[0:n-j,j], y[0:n-j,j], z[0:n-j,j], "c")
    end
    for i = 0:n
        plot3D(xd, yd, zd, "c")
    end
    axis("off")
#    xlabel(L"$x$")
#    ylabel(L"$y$")
end

function midpoints(a)
    prev = [ 3, 1, 2 ]
    succ = [ 2, 3, 1 ]
    m = Matrix{Float64}(undef, 2, 3)
    for i = 1:3
        m[:,i] = ( a[:,prev[i]] + a[:,succ[i]] ) / 2
    end
    return m
end

ψ₁(ξ) = 2 * ξ[1] * ( ξ[1] - 0.5 )
ψ₂(ξ) = 2 * ξ[2] * ( ξ[2] - 0.5 )
ψ₃(ξ) = 2 * ξ[3] * ( ξ[3] - 0.5 )
ψ₄(ξ) = 4 * ξ[2] * ξ[3]
ψ₅(ξ) = 4 * ξ[3] * ξ[1]
ψ₆(ξ) = 4 * ξ[1] * ξ[2]

a = Float64[ 0  1  2
             2  0  3 ]
m = midpoints(a)
n = 20
x, y = bary_grid(n, a)

figure(1)
z = bary_func(ψ₁, n, a)
bary_plot(x, y, z)
plot3D(a[1,1], a[2,1], "ro")
savefig("psi1.pdf")

figure(2)
z = bary_func(ψ₂, n, a)
bary_plot(x, y, z)
plot3D(a[1,2], a[2,2], "ro")
savefig("psi2.pdf")

figure(3)
z = bary_func(ψ₃, n, a)
bary_plot(x, y, z)
plot3D(a[1,3], a[2,3], "ro")
savefig("psi3.pdf")

figure(4)
z = bary_func(ψ₄, n, a)
bary_plot(x, y, z)
plot3D(m[1,1], m[2,1], "ro")
savefig("psi4.pdf")

figure(5)
z = bary_func(ψ₅, n, a)
bary_plot(x, y, z)
plot3D(m[1,2], m[2,2], "ro")
savefig("psi5.pdf")

figure(6)
z = bary_func(ψ₆, n, a)
bary_plot(x, y, z)
plot3D(m[1,3], m[2,3], "ro")
savefig("psi6.pdf")


doc"""
    second_derivative(f::Function, x::AbstractFloat)

# Description

Evaluate the second derivative of `f` at `x` using finite differences. In
mathematical notation, we calculate,

```math
\frac{f(x + \epsilon) - 2 f(x) + f(x - \epsilon)}{\epsilon^2},
```

where ``\epsilon`` is chosen to be small enough to approximate the second
derivative, but not so small as to suffer from extreme numerical inaccuracy.

# Arguments

* `f::Function`: The function to be differentiated.
* `x::AbstractFloat`: The point at which to evaluate the derivative of `f`. Its
    type must implement `eps`.

# Returns

* `y::Real`: The second derivative of `f` evaluated at `x`.

# Examples

```jl
import FiniteDiff: second_derivative
y = second_derivative(sin, 1.0)
```
"""
@inline function second_derivative(f::Function, x::AbstractFloat)::Real
    ϵ = step_size(HessianMode(), x)
    return (f(x + ϵ) - 2 * f(x) + f(x - ϵ)) / (ϵ * ϵ)
end

doc"""
    second_derivative!(
        output::AbstractArray,
        f::Function,
        x::AbstractFloat,
    )

# Description

Evaluate the second derivative of `f` at `x` using finite differences. In
mathematical notation, we calculate,

```math
\frac{f(x + \epsilon) - 2 f(x) + f(x - \epsilon)}{\epsilon^2},
```

where ``\epsilon`` is chosen to be small enough to approximate the second
derivative, but not so small as to suffer from extreme numerical inaccuracy.

The first argument, `output`, will be mutated so that the value of the second
derivative is its first element.

# Arguments

* `output::AbstractArray`: An array whose first element will be mutated.
* `f::Function`: The function to be differentiated.
* `x::AbstractFloat`: The point at which to evaluate the derivative of `f`. Its
    type must implement `eps`.

# Returns

* `void::Nothing`: This function is called for its side-effects.

# Examples

```jl
import FiniteDiff: second_derivative
y = Array(Float64, 1)
second_derivative!(y, sin, 1.0)
```
"""
function second_derivative!(
    output::AbstractArray,
    f::Function,
    x::AbstractFloat,
)::Void
    output[1] = second_derivative(f, x)
    return
end

"""
    second_derivative(f::Function; mutates::Bool=false)

# Description

Construct a new function that will evaluate the second derivative of `f` at
any point `x`. A keyword argument specifies whether the resulting function
should be mutating or non-mutating.

# Arguments

* `f::Function`: The function to be differentiated.

# Keyword Arguments

* `mutates::Bool = false`: Determine whether the resulting function will mutate
    its inputs or will be a pure function.

# Returns

* `void::Nothing`: This function is called for its side-effects.

# Examples

```jl
import FiniteDiff: second_derivative
f′′ = second_derivative(sin)
f′′(1.0)
```
"""
function second_derivative(f::Function; mutates::Bool=false)::Function
    if mutates
        return (output, x) -> second_derivative!(output, f, x)
    else
        return (x, ) -> second_derivative(f, x)
    end
end


###############################################################################
#
# Linear regression of Y~X
#
###############################################################################
"""
    linear_df_xy(df::DataFrame, intercept::Bool)

Linear fit the X and Y in a given DataFrame, given
- `df` DataFrame that use `X` and `Y` to store values
- `intercept` Whether using intercept when making the regression
"""
function linear_df_xy(df::DataFrame, intercept::Bool)
    # run the fitting
    if intercept
        lr = lm(@formula(Y ~ X), df);
    else
        lr = lm(@formula(Y ~ 0 + X), df);
    end

    return lr
end

"""
    line_regress(
                xs::Array,
                ys::Array;
                intercept::Bool = true,
                sorting::Bool = true
    )

Make linear regression and return the fitted results, given
- `xs` Array of x, can be NaN
- `ys` Array of y, can be NaN
- `intercept` Optional: if true use intercept in the fitting
- `sorting` Optional: if true, sort the values
"""
function line_regress(
            xs::Array,
            ys::Array;
            intercept::Bool = true,
            sorting::Bool = true
)
    # filter out NaN from the lists
    new_x = Float64[];
    new_y = Float64[];
    for (x,y) in zip(xs, ys)
        if ~isnan(x) && ~isnan(y)
            push!(new_x, x);
            push!(new_y, y);
        end
    end

    # put new_x and new_y into a DataFrame
    df = DataFrame(X=new_x, Y=new_y);
    sorting ? sort!(df) : nothing;

    # run the fitting
    lr  = linear_df_xy(df, intercept);
    _r2 = adjr2(lr);

    # create items for the struct
    if intercept
        _inter    = coef(lr)[1];
        _slope    = coef(lr)[2];
        _inter_p  = coeftable(lr).cols[4][1];
        _slope_p  = coeftable(lr).cols[4][2];
        _inter_ci = confint(lr)[1,:];
        _slope_ci = confint(lr)[2,:];
    else
        _inter    = Inf;
        _slope    = coef(lr)[1];
        _inter_p  = Inf;
        _slope_p  = coeftable(lr).cols[4][1];
        _inter_ci = [-Inf, Inf];
        _slope_ci = confint(lr)[1,:];
    end

    # calculate the predictions
    if intercept
        mt = [ones(length((df).X)) (df).X];
    else
        mt = (df).X[:,:];
    end
    pred = predict(lr, mt; interval=:confidence);
    df[!,"predY" ] = pred.prediction;
    df[!,"lowerY"] = pred.lower[:,1];
    df[!,"upperY"] = pred.upper[:,1];

    return LinearRegressionResult(
                lr,
                _r2,
                _inter,
                _slope,
                _inter_p,
                _slope_p,
                _inter_ci,
                _slope_ci,
                df
    )
end


### A Pluto.jl notebook ###
# v0.16.1

using Markdown
using InteractiveUtils

# ╔═╡ 987c5c4d-3257-41cf-b8a9-7b3e814ba61b
begin
	using Pkg
	Pkg.add("Plots")
	using Plots
end

# ╔═╡ 9215d698-38ed-11ec-1e52-45e96b1bc12e
md"""
# Chapter 1
## Section 1

Taylor series in the neighborhood of x = a

$$f(x) = \sum_{n=0}^{\infty}\frac{f^{(n)}(a)}{n!}(x-a)^n$$

Near a = 0 and for x = $\pm$h we can write at the third order

$$f(h) = f_0 + hf' + \frac{h^2}{2}f''$$
$$f(-h) = f_0 - hf' + \frac{h^2}{2}f''$$

If we sub

$$f(h) - f(-h) = 2hf'$$
$$\Rightarrow f' = \frac{f(h) - f(-h)}{2h}$$

Centered difference
"""

# ╔═╡ c7732e76-19b9-4aa6-a61e-5b142dae1d3e
function funcdiff(f, x::T, dx::T)::T where T <: Number
	(f(x + dx) - f(x - dx)) / (2 * dx)
end

# ╔═╡ 5137159c-90a8-49a5-9aa3-5356f3137662
begin
	x = collect(0:0.01:7)
	y = funcdiff.(sin, x, 0.01)
	plot(x, y)
end

# ╔═╡ 72affc38-4df9-48af-98b7-0b1b92d9d980
function arraydiff(x::Vector{T}, dx::T)::Vector{T} where T <: Number
	(x[2:end] - x[1:end-1]) / (2 * dx)
end

# ╔═╡ 42d3c309-bb3a-479b-91fd-3d26adda0555
begin
	u = collect(0:0.01:7)
	v = arraydiff(sin.(x), 0.01)
	plot(u[2:end], v)
end

# ╔═╡ cb996ce0-a6f1-4176-b226-b1e4b322fabb
md"""
We can apply this method to other orders of derivations and order of the Taylor series.

Increasing the order of the series will increase the accuracy but it will also increase the calculation time.

The value of h also influences the precision

Exercice: For the same function and interval, evaluate the mean error with the analytic derivation.
"""

# ╔═╡ 7b6412e5-2f4f-409a-b959-6c9c4eb29a9e
begin
	allh = collect(0.01:0.01:1)
	alleps = Float64[]
	for h in allh
		x = 0:h:7
		cosine = cos.(x)
		sinus = funcdiff.(sin, x, h)
		eps = sum(sqrt.(cosine.^2 + sinus.^2)) / length(x)
		append!(alleps, eps)
	end
	plot(allh, alleps)
end

# ╔═╡ Cell order:
# ╠═987c5c4d-3257-41cf-b8a9-7b3e814ba61b
# ╟─9215d698-38ed-11ec-1e52-45e96b1bc12e
# ╠═c7732e76-19b9-4aa6-a61e-5b142dae1d3e
# ╠═5137159c-90a8-49a5-9aa3-5356f3137662
# ╠═72affc38-4df9-48af-98b7-0b1b92d9d980
# ╠═42d3c309-bb3a-479b-91fd-3d26adda0555
# ╟─cb996ce0-a6f1-4176-b226-b1e4b322fabb
# ╠═7b6412e5-2f4f-409a-b959-6c9c4eb29a9e


using SpecialFunctions: gamma
using PyCall
np = pyimport("numpy")
plt = pyimport("matplotlib.pyplot")
tri = pyimport("matplotlib.tri")

function tri_area(xy, pair)
    """The area of the triangle formed by point xy and another pair or points."""
    return 0.5 * np.linalg.norm(np.cross(pair[1,:] .- xy, pair[2,:] .- xy))
end

function xy2bc(xy, pair_corners; tol=1.e-4, AREA=0.5*1*0.75^0.5)
    """Converts 2D Cartesian coordinates to barycentric."""
    coords = np.array([tri_area(xy, p) for p in pair_corners]) / AREA
    return np.clip(coords, tol, 1.0 - tol)
end

function pdf_Dir(θ, α)
    """Plots pdf of Dirichlet distribution"""
    return prod(gamma.(α))/gamma(prod(α)) * prod(θ.^α)
end

function pdf_contours_simplex(concentration_params; nlevels=200, subdiv=8)
    """Generate pdf contours."""

    corners = np.array([[0, 0], [1, 0], [0.5, 0.75^0.5]])
    AREA = 0.5 * 1 * 0.75^0.5
    triangle = tri.Triangulation(corners[:, 1], corners[:, 2])

    pair_corners = [corners[np.roll(range(1, length=3), -i)[2:end],:] for i in range(0, length=3)]

    refiner = tri.UniformTriRefiner(triangle)
    trimesh = refiner.refine_triangulation(subdiv=subdiv)
    pvals = [pdf_Dir(xy2bc(xy, pair_corners, AREA=AREA), concentration_params) for xy in zip(trimesh.x, trimesh.y)]

    return trimesh, pvals
end


using Polynomials

"""
Calculate Chebyshev polynomials coefficients

Instead of computing the polynomial at a specific point,
use `Polynomials` package to compute the actual polynomial
"""
function poly_chebyshev{T<:Number}(n, ::Type{T}=Int, var=:x)

    if n==0
        return Poly{T}([one(T)], var)
    elseif n==1
        return Poly{T}([zero(T), one(T)], var)
    end

    px = Poly{T}([zero(T), one(T)], var)
    T0 = Poly{T}([one(T)], var)
    T1 = px

    for i = 2:n
        T2 = 2*px*T1 - T0
        T0 = T1
        T1 = T2
    end

    return T1

end

"""
Calculate derivative of Chebyshev polynomials coefficients

Instead of computing the polynomial at a specific point,
use `Polynomials` package to compute the actual polynomial
"""
poly_dchebyshev{T<:Number}(n, ::Type{T}=Int, var=:x) = polyder(poly_chebyshev(n,T, var))

"""
Calculate the coefficients of Chebyshev polynomials of the second kind

Instead of computing the polynomial at a specific point,
use `Polynomials` package to compute the actual polynomial
"""
function poly_chebyshev2{T<:Number}(n, ::Type{T}=Int, var=:x)

    if n==0
        return Poly{T}([one(T)], var)
    elseif n==1
        return Poly{T}([zero(T), 2*one(T)], var)
    end

    px = Poly{T}([zero(T), one(T)], var)
    U0 = Poly{T}([one(T)], var)
    U1 = Poly{T}([zero(T), 2*one(T)], var)

    for i = 2:n
        U2 = 2*px*U1 - U0
        U0 = U1
        U1 = U2
    end

    return U1

end

"""
Calculate the coefficients of the derivative of Chebyshev polynomials of the second kind

Instead of computing the polynomial at a specific point,
use `Polynomials` package to compute the actual polynomial
"""
poly_dchebyshev2{T<:Number}(n, ::Type{T}=Int, var=:x) = polyder(poly_chebyshev2(n,T, var))

"""
Calculate the coefficients of Legendre polynomials

Instead of computing the polynomial at a specific point,
use `Polynomials` package to compute the actual polynomial
"""
function poly_legendre{T<:Number}(n, ::Type{T}=Float64, var=:x)

    if n==0
        return Poly{T}([one(T)], var)
    elseif n==1
        return Poly{T}([zero(T), one(T)], var)
    end

    px = Poly{T}([zero(T), one(T)], var)
    p0 = Poly{T}([one(T)], var)
    p1 = px

    for i = 2:n
        p2 = ( (2i-1)*px*p1 - (i-1)*p0 ) / i
        p0 = p1
        p1 = p2
    end

    return p1
end

"""
Calculate the coefficients of the derivativ of Legendre polynomials

Instead of computing the polynomial at a specific point,
use `Polynomials` package to compute the actual polynomial
"""
poly_dlegendre{T<:Number}(n, ::Type{T}=Int, var=:x) = polyder(poly_legendre(n,T, var))

"""
Calculate the coefficients of Jacobi polynomials

Instead of computing the polynomial at a specific point,
use `Polynomials` package to compute the actual polynomial
"""
function poly_jacobi{T<:Number}(n, a, b, ::Type{T}=Float64, var=:x)

    ox = one(T)
    zx = zero(T)
    if n==0
        return Poly{T}([one(T)], var)
    elseif n==1
        return Poly{T}([(a-b)/2, (a+b+2)/2], var)
    end
    p0 = Poly{T}([one(T)], var)
    p1 = Poly{T}([(a-b)/2, (a+b+2)/2], var)
    px = Poly{T}([zero(T), one(T)], var)

    for i = 1:(n-1)
	a1 = 2*(i+1)*(i+a+b+1)*(2*i+a+b);
	a2 = (2*i+a+b+1)*(a*a-b*b);
	a3 = (2*i+a+b)*(2*i+a+b+1)*(2*i+a+b+2);
	a4 = 2*(i+a)*(i+b)*(2*i+a+b+2);
	p2 = ox/a1*( (a2 + a3*px)*p1 - a4*p0);

        p0 = p1
        p1 = p2
    end

    return p1
end

"""
Calculate the coefficients of the derivative of Jacobi polynomials

Instead of computing the polynomial at a specific point,
use `Polynomials` package to compute the actual polynomial
"""
poly_djacobi{T<:Number}(n, a, b, ::Type{T}=Float64, var=:x) = polyder(poly_jacobi(n,a,b,T,var))


using Polynomials
using LsqFit

using Plots

function main()

    # nonlinear fit stuff
    p0 = [0.5, 0.5]                         # guess
    model(t, p) = p[1] * exp.(-p[2] * t)    # model trying to fit
    model1(t, p) = p[1] * sin.(p[2] * t) + p[3]

    xdata = range(0, stop=10, length=20)
    ydata = model(xdata, [1.0 2.0]) + 0.01*randn(length(xdata))

    # poly_fit(xdata, ydata)
    nlin_fit(model, xdata, ydata, p0)

end

function poly_fit(xdata, ydata)

    pfit1 = fit(xdata, ydata, 1)
    pfit2 = fit(xdata, ydata, 2)
    pfit3 = fit(xdata, ydata, 3)

    xlin = range(xdata[1], xdata[end], length=200)

    # plotting
    scatter(xdata, ydata, markersize=3, legend=:topright, label="data")
    # plot!(xlin, pfit1.(xlin), linestyle=:dash, label="f(x) = x")
    # plot!(xlin, pfit2.(xlin), label="f(x) = x²")
    # plot!(xlin, pfit3.(xlin), label="f(x) = x³", dpi=200)

    pfit1 = fit(xdata[1:5], ydata[1:5], 3)
    pfit2 = fit(xdata[5:end], ydata[5:end], 1)

    plot!(xlin[1:50], pfit1.(xlin[1:50]), label="f(x) = x³")
    plot!(xlin[50:end], pfit2.(xlin[50:end]), label="f(x) = x", dpi=200)

    xaxis!("x")
    yaxis!("y")
    title!("poly fit")

end

function nlin_fit(model, xdata, ydata, p0)

    nlinfit = curve_fit(model, xdata, ydata, p0)
    pfit = nlinfit.param
    print(pfit)
    xlin = range(xdata[1], xdata[end], length=200)

    scatter(xdata, ydata, markersize=3, legend=:topright, label="data")
    plot!(xlin, model(xlin, [p0[1], p0[2]]), label="initial model")
    plot!(xlin, model(xlin, [pfit[1], pfit[2]]), linestyle=:dash, label="fitted model", dpi=200)

    xaxis!("x")
    yaxis!("y")
    title!("nonlinear fit")

end

main()


########################################
###### Vectorized "dot" operators ######
########################################

# For every binary operation like ^, there is a corresponding "dot" operation .^ that is
# automatically defined to perform ^ element-by-element on arrays. For example, [1,2,3] ^ 3
# is not defined, since there is no standard mathematical meaning to "cubing" a (non-square) array,
# but [1,2,3] .^ 3 is defined as computing the elementwise (or "vectorized") result [1^3, 2^3, 3^3].
# Similarly for unary operators like ! or √, there is a corresponding .√ that applies the operator elementwise.

# [1,2,3] .^ 3
println([1,2,3] .^ 3)  # [1, 8, 27]


# Rank-based correlations
#
# - Spearman's correlation
# - Kendall's correlation
#

#######################################
#
#   Spearman correlation
#
#######################################

corspearman(x::RealVector, y::RealVector) = cor(tiedrank(x), tiedrank(y))

corspearman(X::RealMatrix, Y::RealMatrix) = cor(mapslices(tiedrank, X, 1), mapslices(tiedrank, Y, 1))
corspearman(X::RealMatrix, y::RealVector) = cor(mapslices(tiedrank, X, 1), tiedrank(y))
corspearman(x::RealVector, Y::RealMatrix) = cor(tiedrank(x), mapslices(tiedrank, Y, 1))

corspearman(X::RealMatrix) = (Z = mapslices(tiedrank, X, 1); cor(Z, Z))

#######################################
#
#   Kendall correlation
#
#######################################

# Knigh JASA (1966)

function corkendall!(x::RealVector, y::RealVector)
    if any(isnan(x)) || any(isnan(y)) return NaN end
    n = length(x)
    if n != length(y) error("Vectors must have same length") end

    # Initial sorting
    pm = sortperm(y)
    x[:] = x[pm]
    y[:] = y[pm]
    pm[:] = sortperm(x)
    x[:] = x[pm]

    # Counting ties in x and y
    iT = 1
    nT = 0
    iU = 1
    nU = 0
    for i = 2:n
        if x[i] == x[i-1]
            iT += 1
        else
            nT += iT*(iT - 1)
            iT = 1
        end
        if y[i] == y[i-1]
            iU += 1
        else
            nU += iU*(iU - 1)
            iU = 1
        end
    end
    if iT > 1 nT += iT*(iT - 1) end
    nT = div(nT,2)
    if iU > 1 nU += iU*(iU - 1) end
    nU = div(nU,2)

    # Sort y after x
    y[:] = y[pm]

    # Calculate double ties
    iV = 1
    nV = 0
    jV = 1
    for i = 2:n
        if x[i] == x[i-1] && y[i] == y[i-1]
            iV += 1
        else
            nV += iV*(iV - 1)
            iV = 1
        end
    end
    if iV > 1 nV += iV*(iV - 1) end
    nV = div(nV,2)

    nD = div(n*(n - 1),2)
    return (nD - nT - nU + nV - 2swaps!(y)) / (sqrt(nD - nT) * sqrt(nD - nU))
end

corkendall(x::RealVector, y::RealVector) = corkendall!(float(copy(x)), float(copy(y)))

corkendall(X::RealMatrix, y::RealVector) = Float64[corkendall!(float(X[:,i]), float(copy(y))) for i in 1:size(X, 2)]

corkendall(x::RealVector, Y::RealMatrix) = (n = size(Y,2); reshape(Float64[corkendall!(float(copy(x)), float(Y[:,i])) for i in 1:n], 1, n))

corkendall(X::RealMatrix, Y::RealMatrix) = Float64[corkendall!(float(X[:,i]), float(Y[:,j])) for i in 1:size(X, 2), j in 1:size(Y, 2)]

function corkendall(X::RealMatrix)
    n = size(X, 2)
    C = eye(n)
    for j = 2:n
        for i = 1:j-1
            C[i,j] = corkendall!(X[:,i],X[:,j])
            C[j,i] = C[i,j]
        end
    end
    return C
end

# Auxilliary functions for Kendall's rank correlation

function swaps!(x::RealVector)
    n = length(x)
    if n == 1 return 0 end
    n2 = div(n, 2)
    xl = sub(x, 1:n2)
    xr = sub(x, n2+1:n)
    nsl = swaps!(xl)
    nsr = swaps!(xr)
    sort!(xl)
    sort!(xr)
    return nsl + nsr + mswaps(xl,xr)
end

function mswaps(x::RealVector, y::RealVector)
    i = 1
    j = 1
    nSwaps = 0
    n = length(x)
    while i <= n && j <= length(y)
        if y[j] < x[i]
            nSwaps += n - i + 1
            j += 1
        else
            i += 1
        end
    end
    return nSwaps
end


using NumericalMethodsforEngineers
using Test

@sym begin
  ClearAll(xi, yi, Y, R, sol, ytilde)
  xi = [0, 1//2, 1]
  yi = [0, a, 1]
  Y(x_) := LagrangePolynomial(xi, yi)
  #
  # Can be formulated as ytile(x) = F(x) + C1(a) * Ψ(x)
  #
  # F(x_) := 2*x^2 - x
  # C(a_) := -4a
  # Ψ(x_) := x^2 - x
  #
  #Y(x_) := F(x) + C(a)*Ψ(x)
  #
  R(x_) = Simplify(D(Y(x), x, 2) - 3*x - 4*Y(x))
  sol = Solve(Integrate(R(x), [x, 0, 1]), a)
  ytilde(x_) = Y(x) ./ Flatten(sol)
  SetJ(r, ToString(Simplify(R(x))))
  SetJ(s, "$(sol[1][1][2])")
  SetJ(t, ToString(Simplify(Expand(ytilde(x)))));
end

println("\n\nExample 7.16: y'' = 3x + 4y, y(0)=0, y(1)=1")
println("by 1-point subdomain Weighted Residual Method")
@sym Println("\nY(x) = ", Y(x), "\n")
@sym Println("R(x) = ", R(x), "\n")
@sym Println("C1 = $(-4*sol[1][1][2])\n")
@sym Println("ytilde_1pt_subdomain(x) = ", Simplify(ytilde(x)), "\n")
println("( Example 7.16 gives: ytilde = 1/16*x*(21x - 5) )")
println()

@eval rf_1pt_subdomain(x, a) = $(Meta.parse(r))
@eval a = $(Meta.parse(s))
@eval ytilde_1pt_subdomain(x) = $(Meta.parse(t))

rf_1pt_subdomain_1(x) = rf_1pt_subdomain(x, a)
@test r == "4 - 8a - 3x + 4x*(1 - 2x + 4a*(-1 + x))"
@test t == "(1/16)*x*(-5 + 21x)"
@test (quadgk(rf_1pt_subdomain_1, 0, 1))[1] < eps()