using SummationByParts

function nodecalc(sbp::TetSBP, isDG::Bool)
  vtx = [0.0 0 0; 1 0 0; 0 1 0; 0 0 1]
  r1 = vtx[1, :]
  r2 = vtx[2, :]
  r3 = vtx[3, :]
  r4 = vtx[4, :]

  T = zeros(3,3)
  T[:, 1] = r2 - r1
  T[:, 2] = r3 - r1
  T[:, 3] = r4 - r1

  # create operator
  if isDG
    coords = SummationByParts.SymCubatures.calcnodes(sbp.cub, vtx)
  else
    coords = calcnodes(sbp, vtx)
  end

  xi = zeros(coords)

  for i=1:size(coords,2)
    xi[:, i] = T\(coords[:, i] - r1.')
  end

  return xi, coords
end

function printCaseStatement(xi)
  nnodes = size(xi, 2)

  f = open("caseout.txt", "w")
  for i=1:nnodes
    xi1 = xi[1, i]
    xi2 = xi[2, i]
    if size(xi, 1) == 2
      xi3 = 0.0
    else
      xi3 = xi[3, i]
    end
    println(f, "case $(i-1):")
    println(f, "  {xi = Vector3($xi1, $xi2, $xi3); break; }")
  end

  println(f, "default:")
  println(f, "  {xi = Vector3(0, 0, 0); break; }")
  close(f)
end

function printBaryCoords(xi, coords)

  (tmp, nnodes) = size(xi)

  for i=1:nnodes
    println("node ", i, " coords: ", coords[:,i])
  end

  print("\n")

  for i=1:nnodes
    xi4 = 1 - sum(xi[:,i])
    println("node ", i, " coords: ", xi[:,i], " , xi4 = ", xi4)
  end
end

#=
sbp = TetSBP{Float64}(degree=4, reorder=false, internal=false)
xi, coords = nodecalc(sbp, true)
printBaryCoords(xi, coords)
#printCaseStatement(xi)
=#


#=
Provides a type called LQ for solving linear quadratic control
problems.

@author : Chase Coleman <ccoleman@stern.nyu.edu>
@author : Spencer Lyon <spencer.lyon@nyu.edu>

@date : 2014-08-26

=#

"""
Compute the limit of a Nash linear quadratic dynamic game.

Player `i` minimizes

    sum_{t=1}^{inf}(x_t' r_i x_t + 2 x_t' w_i
    u_{it} +u_{it}' q_i u_{it} + u_{jt}' s_i u_{jt} + 2 u_{jt}'
    m_i u_{it})

subject to the law of motion

    x_{t+1} = A x_t + b_1 u_{1t} + b_2 u_{2t}

and a perceived control law :math:`u_j(t) = - f_j x_t` for the other player.

The solution computed in this routine is the `f_i` and `p_i` of the associated
double optimal linear regulator problem.

##### Arguments

- `A` : Corresponds to the above equation, should be of size (n, n)
- `B1` : As above, size (n, k_1)
- `B2` : As above, size (n, k_2)
- `R1` : As above, size (n, n)
- `R2` : As above, size (n, n)
- `Q1` : As above, size (k_1, k_1)
- `Q2` : As above, size (k_2, k_2)
- `S1` : As above, size (k_1, k_1)
- `S2` : As above, size (k_2, k_2)
- `W1` : As above, size (n, k_1)
- `W2` : As above, size (n, k_2)
- `M1` : As above, size (k_2, k_1)
- `M2` : As above, size (k_1, k_2)
- `;beta::Float64(1.0)` Discount rate
- `;tol::Float64(1e-8)` : Tolerance level for convergence
- `;max_iter::Int(1000)` : Maximum number of iterations allowed

##### Returns

- `F1::Matrix{Float64}`: (k_1, n) matrix representing feedback law for agent 1
- `F2::Matrix{Float64}`: (k_2, n) matrix representing feedback law for agent 2
- `P1::Matrix{Float64}`: (n, n) matrix representing the steady-state solution to the associated discrete matrix ticcati equation for agent 1
- `P2::Matrix{Float64}`: (n, n) matrix representing the steady-state solution to the associated discrete matrix riccati equation for agent 2

"""
function nnash(a, b1, b2, r1, r2, q1, q2, s1, s2, w1, w2, m1, m2;
               beta::Float64=1.0, tol::Float64=1e-8, max_iter::Int=1000)

    # Apply discounting
    a, b1, b2 = map(x->sqrt(beta) * x, Any[a, b1, b2])
    dd = 10
    its = 0
    n = size(a, 1)

    # NOTE: if b1/b2 has 2 dimensions, this is exactly what we want.
    #       if b1/b2 has 1 dimension size(b, 2) returns a 1, so it is also
    #       what we want
    k_1 = size(b1, 2)
    k_2 = size(b2, 2)

    # initial values
    v1 = eye(k_1)
    v2 = eye(k_2)
    p1 = zeros(n, n)
    p2 = zeros(n, n)
    f1 = randn(k_1, n)
    f2 = randn(k_2, n)

    while dd > tol
        # update
        f10 = f1
        f20 = f2

        g2 = (b2'*p2*b2 .+ q2)\v2
        g1 = (b1'*p1*b1 .+ q1)\v1
        h2 = g2*b2'*p2
        h1 = g1*b1'*p1

        # Break up computation of f1 and f2
        f_1_left = v1 .- (h1*b2 .+ g1*m1')*(h2*b1 .+ g2*m2')
        f_1_right = h1*a .+ g1*w1' .- (h1*b2 .+ g1*m1')*(h2*a .+ g2*w2')

        f1 =  f_1_left\f_1_right
        f2 = h2*a .+ g2*w2' .- (h2*b1 .+ g2*m2')*f1

        a2 = a .- b2*f2
        a1 = a .- b1*f1

        p1 = (a2'*p1*a2) .+ r1 .+ (f2'*s1*f2) .- (a2'*p1*b1 .+ w1 .- f2'*m1)*f1
        p2 = (a1'*p2*a1) .+ r2 .+ (f1'*s2*f1) .- (a1'*p2*b2 .+ w2 .- f1'*m2)*f2

        dd = maximum(abs(f10 .- f1) + abs(f20 .- f2))
        its = its + 1
        if its > max_iter
            error("Reached max iterations, no convergence")
        end

    end

    return f1, f2, p1, p2
end


using ModelingToolkit, StaticArrays, LinearAlgebra
using DiffEqBase
using Test

# Define some variables
@parameters t σ ρ β
@variables x(t) y(t) z(t)
@derivatives D'~t

function _clean(O::Operation)
    @assert isa(O.op, Variable)
    return O.op
end
_clean(x::Variable) = x
function test_diffeq_inference(name, sys, iv, dvs, ps)
    @testset "ODESystem construction: $name" begin
        @test independent_variables(sys) == Set([_clean(iv)])
        @test dependent_variables(sys)   == Set(_clean.(dvs))
        @test parameters(sys)            == Set(_clean.(ps))
    end
end
function test_nlsys_inference(name, sys, vs, ps)
    @testset "NonlinearSystem construction: $name" begin
        @test dependent_variables(sys) == Set(vs)
        @test parameters(sys)          == Set(_clean.(ps))
    end
end

# Define a differential equation
eqs = [D(x) ~ σ*(y-x),
       D(y) ~ x*(ρ-z)-y,
       D(z) ~ x*y - β*z]

ModelingToolkit.simplified_expr.(eqs)[1]
:(derivative(x(t), t) = σ * (y(t) - x(t))).args
de = ODESystem(eqs)
test_diffeq_inference("standard", de, t, (x, y, z), (σ, ρ, β))
generate_function(de, [x,y,z], [σ,ρ,β])
jac_expr = generate_jacobian(de)
jac = calculate_jacobian(de)
jacfun = eval(jac_expr)
# iip
f = ODEFunction(de, [x,y,z], [σ,ρ,β])
Wfact, Wfact_t = ModelingToolkit.calculate_factorized_W(de)
fw, fwt = map(eval, ModelingToolkit.generate_factorized_W(de))
du = zeros(3)
u  = collect(1:3)
p  = collect(4:6)
f(du, u, p, 0.1)
@test du == [4, 0, -16]
J = zeros(3, 3)
jacfun(J, u, p, t)
FW = zeros(3, 3)
FWt = zeros(3, 3)
fw(FW, u, p, 0.2, 0.1)
fwt(FWt, u, p, 0.2, 0.1)
# oop
f = ODEFunction(de, [x,y,z], [σ,ρ,β])
fw, fwt = map(eval, ModelingToolkit.generate_factorized_W(de))
du = @SArray zeros(3)
u  = SVector(1:3...)
p  = SVector(4:6...)
@test f(u, p, 0.1) === @SArray [4, 0, -16]
Sfw = fw(u, p, 0.2, 0.1)
@test Sfw.L ≈ UnitLowerTriangular(FW)
@test Sfw.U ≈ UpperTriangular(FW)
sol = Sfw \ @SArray ones(3)
@test sol isa SArray
@test sol ≈ -(I - 0.2*J)\ones(3)
Sfw_t = fwt(u, p, 0.2, 0.1)
@test Sfw_t.L ≈ UnitLowerTriangular(FWt)
@test Sfw_t.U ≈ UpperTriangular(FWt)
sol = Sfw_t \ @SArray ones(3)
@test sol isa SArray
@test sol ≈ -(I/0.2 - J)\ones(3)

@testset "time-varying parameters" begin
    @parameters σ′(t-1)
    eqs = [D(x) ~ σ′*(y-x),
           D(y) ~ x*(ρ-z)-y,
           D(z) ~ x*y - β*z]
    de = ODESystem(eqs)
    test_diffeq_inference("global iv-varying", de, t, (x, y, z), (σ′, ρ, β))
    @test begin
        f = eval(generate_function(de, [x,y,z], [σ′,ρ,β]))
        du = [0.0,0.0,0.0]
        f(du, [1.0,2.0,3.0], [x->x+7,2,3], 5.0)
        du ≈ [11, -3, -7]
    end

    @parameters σ(..)
    eqs = [D(x) ~ σ(t-1)*(y-x),
           D(y) ~ x*(ρ-z)-y,
           D(z) ~ x*y - β*z]
    de = ODESystem(eqs)
    test_diffeq_inference("single internal iv-varying", de, t, (x, y, z), (σ, ρ, β))
    @test begin
        f = eval(generate_function(de, [x,y,z], [σ,ρ,β]))
        du = [0.0,0.0,0.0]
        f(du, [1.0,2.0,3.0], [x->x+7,2,3], 5.0)
        du ≈ [11, -3, -7]
    end

    eqs = [D(x) ~ x + 10σ(t-1) + 100σ(t-2) + 1000σ(t^2)]
    de = ODESystem(eqs)
    test_diffeq_inference("many internal iv-varying", de, t, (x,), (σ,))
    @test begin
        f = eval(generate_function(de, [x], [σ]))
        du = [0.0]
        f(du, [1.0], [t -> t + 2], 5.0)
        du ≈ [27561]
    end
end

# Conversion to first-order ODEs #17
@derivatives D3'''~t
@derivatives D2''~t
@variables u(t) u_tt(t) u_t(t) x_t(t)
eqs = [D3(u) ~ 2(D2(u)) + D(u) + D(x) + 1
       D2(x) ~ D(x) + 2]
de = ODESystem(eqs)
de1 = ode_order_lowering(de)
lowered_eqs = [D(u_tt) ~ 2u_tt + u_t + x_t + 1
               D(x_t)  ~ x_t + 2
               D(u_t)  ~ u_tt
               D(u)    ~ u_t
               D(x)    ~ x_t]
@test de1 == ODESystem(lowered_eqs)
test_diffeq_inference("first-order transform", de1, t, [u_tt, x_t, u_t, u, x], [])
du = zeros(5)
ODEFunction(de1, [u_tt, x_t, u_t, u, x], [])(du, ones(5), nothing, 0.1)
@test du == [5.0, 3.0, 1.0, 1.0, 1.0]

# Internal calculations
a = y - x
eqs = [D(x) ~ σ*a,
       D(y) ~ x*(ρ-z)-y,
       D(z) ~ x*y - β*z]
de = ODESystem(eqs)
generate_function(de, [x,y,z], [σ,ρ,β])
jac = calculate_jacobian(de)
f = ODEFunction(de, [x,y,z], [σ,ρ,β])

# Define a nonlinear system
eqs = [0 ~ σ*(y-x),
       0 ~ x*(ρ-z)-y,
       0 ~ x*y - β*z]
ns = NonlinearSystem(eqs, [x,y,z])
test_nlsys_inference("standard", ns, (x, y, z), (σ, ρ, β))
@test begin
    f = eval(generate_function(ns, [x,y,z], [σ,ρ,β]))
    du = [0.0, 0.0, 0.0]
    f(du, [1,2,3], [1,2,3])
    du ≈ [1, -3, -7]
end

@derivatives D'~t
@parameters A B C
_x = y / C
eqs = [D(x) ~ -A*x,
       D(y) ~ A*x - B*_x]
de = ODESystem(eqs)
@test begin
    f = eval(generate_function(de, [x,y], [A,B,C]))
    du = [0.0,0.0]
    f(du, [1.0,2.0], [1,2,3], 0.0)
    du ≈ [-1, -1/3]
    du ≈ f([1.0,2.0], [1,2,3], 0.0)
end

# Now nonlinear system with only variables
@variables x y z
@parameters σ ρ β

# Define a nonlinear system
eqs = [0 ~ σ*(y-x),
       0 ~ x*(ρ-z)-y,
       0 ~ x*y - β*z]
ns = NonlinearSystem(eqs, [x,y,z])
jac = calculate_jacobian(ns)
@testset "nlsys jacobian" begin
    @test isequal(jac[1,1], σ * -1)
    @test isequal(jac[1,2], σ)
    @test isequal(jac[1,3], 0)
    @test isequal(jac[2,1], ρ - z)
    @test isequal(jac[2,2], -1)
    @test isequal(jac[2,3], x * -1)
    @test isequal(jac[3,1], y)
    @test isequal(jac[3,2], x)
    @test isequal(jac[3,3], -1 * β)
end
nlsys_func = generate_function(ns, [x,y,z], [σ,ρ,β])
jac_func = generate_jacobian(ns)
f = @eval eval(nlsys_func)

# Intermediate calculations
# Define a nonlinear system
eqs = [0 ~ σ*a,
       0 ~ x*(ρ-z)-y,
       0 ~ x*y - β*z]
ns = NonlinearSystem(eqs, [x,y,z])
nlsys_func = generate_function(ns, [x,y,z], [σ,ρ,β])
jac = calculate_jacobian(ns)
jac = generate_jacobian(ns)

function lotka(u,p,t)
  x = u[1]
  y = u[2]
  [p[1]*x - p[2]*x*y,
  -p[3]*y + p[4]*x*y]
end

prob = ODEProblem(ODEFunction{false}(lotka),[1.0,1.0],(0.0,1.0),[1.5,1.0,3.0,1.0])
de, vars, params = modelingtoolkitize(prob)
ODEFunction(de, vars, params)(similar(prob.u0), prob.u0, prob.p, 0.1)

function lotka(du,u,p,t)
  x = u[1]
  y = u[2]
  du[1] = p[1]*x - p[2]*x*y
  du[2] = -p[3]*y + p[4]*x*y
end

prob = ODEProblem(lotka,[1.0,1.0],(0.0,1.0),[1.5,1.0,3.0,1.0])

de, vars, params = modelingtoolkitize(prob)
ODEFunction(de, vars, params)(similar(prob.u0), prob.u0, prob.p, 0.1)


export cholinv, cholsqrt

using LinearAlgebra
using PositiveFactorizations

import LinearAlgebra: BlasInt

cholinv(x)           = inv(fastcholesky(x))
cholinv(x::Diagonal) = Diagonal(inv.(diag(x)))
cholinv(x::Real)     = inv(x)

function cholinv(x::AbstractMatrix{T}) where { T <: LinearAlgebra.BlasFloat }
    y = fastcholesky(x)
    LinearAlgebra.inv!(y)
    return y.factors
end

cholsqrt(x)           = Matrix(fastcholesky(x).L)
cholsqrt(x::Diagonal) = Diagonal(sqrt.(diag(x)))
cholsqrt(x::Real)     = sqrt(x)

chollogdet(x)           = logdet(fastcholesky(x))
chollogdet(x::Diagonal) = logdet(x)
chollogdet(x::Real)     = logdet(x)

function cholinv_logdet(x)
    # calculate cholesky decomposition
    y = fastcholesky(x)

    # return inverse and log-determinant
    return inv(y), logdet(y)
end
function cholinv_logdet(x::AbstractMatrix{T}) where { T <: LinearAlgebra.BlasFloat }
    # calculate cholesky decomposition
    y = fastcholesky(x)

    # calculate logdeterminant of cholesky decomposition
    ly = logdet(y)

    # calculate inplace inverse of A and store in y.factors
    LinearAlgebra.inv!(y)

    # return inverse and log-determinant
    return y.factors, ly
end
cholinv_logdet(x::Diagonal) = Diagonal(inv.(diag(x))), mapreduce(z -> log(z), +, diag(x))
cholinv_logdet(x::Real)     = inv(x), log(abs(x))

function fastcholesky(mat::AbstractMatrix)
    A = copy(mat)
    C = fastcholesky!(A)
    if !isposdef(C)
        return cholesky(PositiveFactorizations.Positive, Hermitian(mat))
    else
        return C
    end
end

function fastcholesky!(A::AbstractMatrix)
    n = LinearAlgebra.checksquare(A)
    @inbounds for col=1:n
        @inbounds @simd for idx in 1:col-1
            A[col, col] -= A[col, idx]^2;
        end
        if A[col,col] <= 0
            return Cholesky(A, 'L', convert(BlasInt, -1))
        end
        A[col, col] = sqrt(A[col, col])

        @inbounds for row in col+1: n
            @inbounds @simd for idx in 1:col-1
                A[row, col] -= A[row, idx]*A[col, idx]
            end
            A[row, col] /= A[col, col]
        end
    end
    return Cholesky(A, 'L', convert(BlasInt, 0))
end

function fastcholesky!(A::AbstractMatrix{T}) where { T <: LinearAlgebra.BlasFloat }
    # blocked version (https://arxiv.org/pdf/1812.02056.pdf)

    # step size
    s = 250

    n = LinearAlgebra.checksquare(A)
    z = 1
    @inbounds for c in 1:n

        if c == z + s
            BLAS.gemm!('N', 'T', -one(T), view(A, c:n, z:c-1), view(A, c:n, z:c-1), one(T), view(A, c:n, c:n)) # replace with syrk once julia bug has been fixed
            z = c
        end

        @inbounds for k in z:c-1
            A[c,c] -= A[c,k]^2
        end
        if A[c,c] <= 0
            return Cholesky(A, 'L', convert(BlasInt, -1))
        end
        A[c,c] = sqrt(A[c,c])
        @inbounds for i in c+1:n
            @inbounds for k in z:c-1
                A[i,c] -= A[i,k]*A[c,k]
            end
            A[i,c] /= A[c,c]
        end
    end

    return Cholesky(A, 'L', convert(BlasInt, 0))
end


##########################################################
# Rafał Nowak
#
# Instytut Informatyki
# Uniwersytet Wrocławski
#
#
# Pomocniczne materiały do wykładu z analizy numerycznej
#
##########################################################

f(x) = ( 1.0 - cos(x) ) ./ x.^2

# Zadanie: obliczyć f(1.2e-8)

# Rozwiązanie
x = 1.2e-8
y = f(x)    # 7.709882115452e-01

@printf("f(%.1e)=%.5e\n", x,y);

# Trochę dziwny ten wynik, bo łatwo udowodnić, że f(x) <= 0.5


"""
    BCDiag

A structure of diagnostic properties of a Binary Classifier,
facilitates summary plots and tables.
"""
struct BCDiag
    n::Int                  # Number of observations
    n1::Int                 # Number of Class 1
    n0::Int                 # Number of Class 0
    baserate::Float64       # Class 1 incidence rate
    ks::Float64             # Maximum separation
    ksarg::Float64          # max sep occurs at this predictor value
    ksdep::Float64          # max sep occurs at this depth

    concordant::Int         # concordant pairs
    tied::Int               # tied pairs
    discordant::Int         # discordant pairs
    auc::Float64            # Area under ROC = (concordant + 0.5tied)/ total pairs
    gini::Float64           # Gini = 2C - 1

    grp::Vector             # group number [0-groups]
    depth::Vector           # Depth [0, 1]
    cdf1::Vector            # cum1 / n1 = True Positive Rate
    cdf0::Vector            # cum0 / n0 = False Positive Rate

    cntg::Vector            # count this group
    cnt1::Vector            # count class 1 this group
    prd1::Vector            # predicted class 1 this group
    rrObs::Vector           # response rate actual
    rrPrd::Vector           # response rate Predicted

    cumg::Vector            # cum count up to this group
    cum1::Vector            # cum count class 1 up to this group
    cpr1::Vector            # cum predicted class 1 up to this group
    crObs::Vector           # cum response rate actual
    crPrd::Vector           # cum response rate predicted
end

function Base.show(io::IO, ::MIME"text/plain", x::BCDiag)
    print(io,
        "Base rate: ", round(x.baserate, digits=4),
        "   n: ", x.n,
        "   n1: ", x.n1,
        "   n0: ", x.n0,
        "\nks:        ", round(x.ks, digits=4),
        "   occurs at value of ", x.ksarg,
        " depth of ", x.ksdep,
        "\nroc:       ", round(x.auc, digits=4),
        "   concordant pairs: ", x.concordant,
        "   tied pairs: ", x.tied,
        "   discordant pairs: ", x.discordant,
        "\nGini:      ", round(x.gini, digits=4),
    )
end

function Base.show(io::IO, x::BCDiag)
    print(io, round(x.ks, digits=4), "   ", round(x.auc, digits=4))
end

"""
    bcdiag(target, pred; groups = 100, rev = true, tie = 1e-6)

Perform diagnostics of a binary classifier.\\
`target` is a 2 level categorical variable, `pred` is probability of class 1.\\
`groups` is the number of bins to use for plotting/printing.\\
`rev` = true orders `pred` from high to low.\\
`tie` is the tolerance of `pred` where values are considered tied.\\

Returns a BCDiag struct which can be used for plotting or printing:
- `biasplot` is calibration plot of `target` response rate vs. `pred` response rate
- `ksplot` produces ksplot of cumulative distributions
- `rocplot` plots the Receiver Operating Characteristics curve
- `accuracyplot` plots the accuracy curve with adjustable utility
- `liftcurve` is the lift curve
- `cumliftsurve` is the cumulative lift surve
- `liftable` is the lift table as a DataFrame
- `cumliftable` is the cumulative lift table as a DataFrame
"""
function bcdiag(target::BitVector, pred::Vector; groups=100, rev=true, tie=1e-6)
    ks = kstest(target, pred; rev=rev)
    auroc = concordance(target, pred, tie)

    grpid = ranks(pred, groups=groups, rev=rev)
    frqtbl = freqtable(grpid, target)
    frq = frqtbl.array
    grp = NamedArrays.names(frqtbl, 1)

    n0, n1 = sum(frq, dims=1)
    n = n0 + n1

    cntg = vec(sum(frq, dims=2))            # bin count
    cnt1 = @view frq[:, 2]                  # bin count class 1
    cnt0 = @view frq[:, 1]                  # bin count class 0
    cumg = cumsum(cntg)                     # cum bin count
    cum1 = cumsum(cnt1)                     # cum class 1
    cum0 = cumsum(cnt0)                     # cum class 0

    depth = cumg ./ n                       # Depth including this group
    cdf1 = cum1 ./ n1                       # TPR: True Positive Rate
    cdf0 = cum0 ./ n0                       # FPR: False Positive Rate

    rrObs = cnt1 ./ cntg                    # actual response rate
    crObs = cum1 ./ cumg                    # cum actual response rate

    prd1 = Vector{Float64}(undef, length(grp))  # predicted class 1
    for (i, g) in enumerate(grp)
        prd1[i] = sum(pred[grpid .== g])
    end
    cpr1 = cumsum(prd1)                     # cum predicted class 1

    rrPrd = prd1 ./ cntg                    # predicted response rate
    crPrd = cpr1 ./ cumg                    # cum predicted response rate

    BCDiag(
        ks...,
        auroc...,
        grp, depth, cdf1, cdf0,
        cntg, cnt1, prd1, rrObs, rrPrd,
        cumg, cum1, cpr1, crObs, crPrd,
    )
end

function bcdiag(target::Vector, pred::Vector; groups=100, rev=true, tie=1e-6)
    uc = sort!(unique(target))
    length(uc) == 2 || error(ArgumentError("target must be 2 levels"))

    bcdiag(target .== uc[2], pred; rev=rev)
end

"""
    biasplot(x::BCDiag)

return bias calibration plot of `x` - actual response vs. predicted response
"""
function biasplot(x::BCDiag)
    dmin = min(minimum(x.rrPrd), minimum(x.rrObs))
    dmax = max(maximum(x.rrPrd), maximum(x.rrObs))

    plt = plot(
        size=(500, 500),
        aspect_ratio=:equal,
        legend=:bottomright,
        xguidefontsize=10,
        yguidefontsize=10,
        titlefontsize=11,
    )
    title!("Model Calibration Graph", titlefontsize=11)
    xlabel!("Predicted Response Rate", xguidefontsize=10)
    ylabel!("Observed Response Rate", yguidefontsize=10)

    plot!(x.rrPrd, x.rrObs, label=nothing, seriescolor=:blue)
    plot!([dmin, dmax], [dmin, dmax], label=nothing, width=0.5, seriescolor=:black)

    plt
end

"""
    ksplot(x::BCDiag)

return KS plot of `x` - CDF1 (True Positive) and CDF0 (False Positive) versus depth
"""
function ksplot(x::BCDiag)
    plt = plot(
        size=(500, 500),
        aspect_ratio=:equal,
        xlims=(0.0, 1.0),
        ylims=(0.0, 1.0),
        legend=:bottomright,
        xguidefontsize=10,
        yguidefontsize=10,
        titlefontsize=11,
    )
    title!("KS Graph", titlefontsize=11)
    xlabel!("Depth", xguidefontsize=10)
    ylabel!("Empirical CDF", yguidefontsize=10)

    plot!(
        [0.0; x.depth],
        [0.0; x.cdf1],
        label="Class 1",
        width=1.5,
        seriescolor=:blue,
    )

    plot!(
        [0.0; x.depth],
        [0.0; x.cdf0],
        label="Class 0",
        width=0.5,
        seriescolor=:black,
    )

    plot!(
        [0.0, x.baserate, 1.0],
        [0.0, 1.0, 1.0],
        label="Perfect Class 1",
        width=0.5,
        seriescolor=:blue,
        linestyle=:dash,
    )

    plot!(
        [0.0, x.baserate, 1.0],
        [0.0, 0.0, 1.0],
        label="Perfect Class 0",
        width=0.5,
        seriescolor=:black,
        linestyle=:dash,
    )
    ksidx = findlast(x.depth .<= x.ksdep)

    plot!(
        [x.depth[ksidx], x.depth[ksidx]],
        [x.cdf0[ksidx], x.cdf1[ksidx]],
        label=nothing,
        width=0.8,
        linestyle=:dash,
    )

    annotate!((
        0.03,
        0.97,
        text("Base rate = $(round(x.baserate, digits=4))", :black, :left, 9),
    ))
    annotate!((0.03, 0.93, text("Depth = $(round(x.ksdep, digits=4))", :black, :left, 9)))
    annotate!((0.03, 0.89, text("KSarg = $(round(x.ksarg, digits=4))", :black, :left, 9)))
    annotate!((0.03, 0.85, text("KS = $(round(x.ks, digits=4))", :black, :left, 9)))

    plt
end

"""
    rocplot(x::BCDiag)

return ROC plot of `x` - CDF1 (True Positive) vs. CDF0 (False Positive)
"""
function rocplot(x::BCDiag)
    plt = plot(
        size=(500, 500),
        aspect_ratio=:equal,
        xlims=(0.0, 1.0),
        ylims=(0.0, 1.0),
        legend=:bottomright,
        xguidefontsize=10,
        yguidefontsize=10,
        titlefontsize=11,
    )
    title!("ROC Graph", titlefontsize=11)
    xlabel!("False Positive Rate", xguidefontsize=10)
    ylabel!("True Positive Rate", yguidefontsize=10)

    plot!([0.0; x.cdf0], [0.0; x.cdf1], label="Model", seriescolor=:blue, width=1.5)
    plot!(
        [0.0, 1.0],
        [0.0, 1.0],
        label="Random",
        seriescolor=:black,
        width=0.3,
        linestyle=:dash,
    )

    annotate!((0.75, 0.22, text("auc = $(round(x.auc, digits=4))", :black, :left, 9)))
    annotate!((
        0.75,
        0.18,
        text("Gini = $(round(2. * x.auc - 1., digits=4))", :black, :left, 9),
    ))

    plt
end

"""
    accuracyplot(x::BCDiag; util=[1, 0, 0, 1])

Using `util` values for [TP, FN, FP, TN], produce accuracy plot and its [max, argmax, argdep].\\
Default `util` values of [1, 0, 0, 1] gives the standard accuracy value of (TP+TN)/N.
"""
function accuracyplot(x::BCDiag; util=[1, 0, 0, 1])
    p = x.baserate
    q = 1.0 - p

    # perfect classifier
    up = dot([p, 0.0, 0.0, q], util)                  # TP, FN, FP, TN of perfect classifier

    # 0.5 / 0.5 random classifier
    # ur5 = dot([0.5p, 0.5p, 0.5q, 0.5q], util)

    # baserate random classifier
    urb = dot([p * p, p * q, p * q, q * q], util)

    # model utility value [TP, FN, FP, TN]
    uc = [p .* x.cdf1 p .* (1.0 .- x.cdf1) q .* x.cdf0 q .* (1.0 .- x.cdf0)] * util
    ucmax = findmax(uc)

    plt = plot(
        size=(500, 500),
        xlims=(0.0, 1.0),
        legend=:best,
        xguidefontsize=10,
        yguidefontsize=10,
        titlefontsize=11,
    )
    title!("Accuracy Plot", titlefontsize=11)
    xlabel!("Depth", xguidefontsize=10)
    ylabel!("Utility", yguidefontsize=10)

    hline!([up], label="Perfect $(round(up, sigdigits=3))", linestyle=:dash)

    plot!(
        x.depth,
        uc,
        width=1.5,
        seriescolor=:blue,
        label="Model $(round(ucmax[1], sigdigits=3))",
    )
    vline!([x.depth[ucmax[2]]], label="Depth $(x.depth[ucmax[2]])")

    hline!([urb], label="Base Random $(round(urb, sigdigits=3))", linestyle=:dash)

    plt
end

"""
    liftcurve(x::BCDiag)

return lift curve plot of `x` - actual and predicted versus depth
"""
function liftcurve(x::BCDiag)
    plt = Plots.plot(
        size=(500, 500),
        xlims=(0.0, 1.0),
        ylims=(0.0, Inf), widen=true,
        legend=:topright,
        xguidefontsize=10,
        yguidefontsize=10,
        titlefontsize=11,
    )
    title!("Lift Curve", titlefontsize=11)
    xlabel!("Depth", xguidefontsize=10)
    ylabel!("Response Rate", yguidefontsize=10)

    plot!(x.depth, x.rrObs, label="Actual", width=1.0, seriescolor=:red)
    plot!(x.depth, x.rrPrd, label="Predicted", width=1.5, seriescolor=:blue)
    hline!([x.baserate], label="Base rate", linestyle=:dash)

    plt
end

"""
    cumliftcurve(x::BCDiag)

return cumulative lift curve plot of `x` - cumulative actual and predicted vs. depth
"""
function cumliftcurve(x::BCDiag)
    plt = Plots.plot(
        size=(500, 500),
        xlims=(0.0, 1.0),
        ylims=(0.0, Inf), widen=true,
        legend=:topright,
        xguidefontsize=10,
        yguidefontsize=10,
        titlefontsize=11,
        )
        title!("Cumulative Lift Curve", titlefontsize=11)
        xlabel!("Depth", xguidefontsize=10)
        ylabel!("Response Rate", yguidefontsize=10)

        plot!(x.depth, x.crObs, label="Actual", width=1.0, seriescolor=:red)
        plot!(x.depth, x.crPrd, label="Predicted", width=1.5, seriescolor=:blue)
        hline!([x.baserate], label="Base rate", linestyle=:dash)

    plt
end

"""
    liftable(x::BCDiag)

return lift table of `x` as a DataFrame
"""
function liftable(x::BCDiag)
    out = DataFrame(
        grp=x.grp,
        depth=x.depth,
        count=x.cntg,
        cntObs=x.cnt1,
        cntPrd=x.prd1,
        rrObs=x.rrObs,
        rrPred=x.rrPrd,
        liftObs=x.rrObs ./ x.baserate,
        liftPrd=x.rrPrd ./ x.baserate,
    )

    out
end

"""
    cumliftable(x::BCDiag)

return cumulative lift table of `x` as a DataFrame
"""
function cumliftable(x::BCDiag)
    out = DataFrame(
        grp=x.grp,
        depth=x.depth,
        count=x.cumg,
        cumObs=x.cum1,
        cumPrd=x.cpr1,
        crObs=x.cum1 ./ x.cumg,
        crPrd=x.cpr1 ./ x.cumg,
    )
    out.liftObs = out.crObs ./ x.baserate
    out.liftPrd = out.crPrd ./ x.baserate

    out
end


function reset_seed(seed)
    Random.seed!(seed)
end

function basetype_string(set::LazySet)
    return string(basetype(set))
end

function flatten_dyn(dyn)
    return dyn.A(), dyn.b()
end

function flatten_dyn(dyn::AffDyn)
    return dyn.A, dyn.b
end

function flatten_Interval(I)
    I = convert(Interval, I)
    return min(I), max(I)
end

function flatten_BallInf(B)
    B = convert(BallInf, B)
    return B.center, B.radius
end

function flatten_Hyperrectangle(H)
    H = convert(Hyperrectangle, H)
    return H.center, H.radius
end

function flatten_HalfSpace(hs)
    hs = convert(HalfSpace, hs)
    return hs.a, hs.b
end

function flatten_HPolyhedron(P)
    P = convert(HPolyhedron, P)
    return P.constraints
end

function flatten_VPolygon(P)
    P = convert(VPolygon, P)
    return P.vertices
end

function successor(b::AbstractVector, t, x0)
    result = x0 + b * t
    return result
end

function successor(A::AbstractMatrix, t, x0)
    result = exp(A * t) * x0
    return result
end

function successor(dyn::AffDyn, t, x0)
    if iszero(dyn.A)
        return successor(dyn.b, t, x0)
    elseif iszero(dyn.b)
        return successor(dyn.A, t, x0)
    end
    # affine solution that only works for invertible matrix
#    try
#        inv_A = inv(dyn.A)
#        id = Diagonal(ones(n))
#        result = exp(dyn.A * t) * x0 + (exp(dyn.A * t) - id) * inv_A * dyn.b
#        return result
#    catch e
#        # A not invertible
#        # return nothing
#    end
    # make affine system linear by extending the dimension
    A_high = [dyn.A dyn.b; zeros(length(dyn.b) + 1)']
    x1_high = successor(A_high, t, vcat(x0, 1.0))
    x1 = x1_high[1:end-1]
    return x1
end

function successor(b::AbstractVector, time, P0::LazySet)
    offset = b * time
    P1 = translate(P0, offset)
    return P1
end

function successor(A::AbstractMatrix, time, P0::LazySet)
    B = exp(A * time)
    P1 = linear_map(B, P0)
    return P1
end

function successor(dyn::AffDyn, time, P0::LazySet)
    # make affine system linear by extending the dimension
    n = length(dyn.b)
    A_high = [dyn.A dyn.b; zeros(n + 1)']
    if n == 1
        # 1D case can be solved more efficiently and with numerical stability
        P0_itv = convert(Interval, P0)
        l = low(P0_itv)[1]
        u = high(P0_itv)[1]
        P0_high = Hyperrectangle(low=[l, 1.0], high=[u, 1.0])
        P1_high = successor(A_high, time, P0_high)
        P1 = LazySets.project(P1_high, 1:n)
        P1 = convert(Interval, P1)
    else
        one = Interval(1.0 - 1e-5, 1.0 + 1e-5)
        # one = Singleton([1.0])
        P0_high = CartesianProduct(P0, one)
        P1_high = successor(A_high, time, P0_high)
        P1_high = LazySets.normalize(P1_high)
        P1 = LazySets.project(P1_high, 1:n)
        P1 = LazySets.normalize(P1)
    end
    return P1
end

function is_similar_flow(P0, dyn, P1, time; return_set::Bool=false)
    Q1 = successor(dyn, time, P0)
    if return_set
        return !isdisjoint(P1, Q1), Q1
    else
        return !isdisjoint(P1, Q1)
    end
end

function contains(X::LazySet, point::AbstractVector)
    return point ∈ X
end

function polytope_hull(X, Y)
    return convex_hull(X, Y)
end

function get_ivp(b::AbstractVector, X0)
    problem = @ivp(x' = b, x(0) ∈ X0)
    return problem
end

function get_ivp(A::AbstractMatrix, X0)
    problem = @ivp(x' = Ax, x(0) ∈ X0)
    return problem
end

function get_ivp(dyn::AffDyn, X0)
    A = dyn.A
    b = dyn.b
    problem = @ivp(x' = Ax + b, x(0) ∈ X0)
    return problem
end

function compute_reach_tube_hull(dyn, X0::LazySet, time_horizon,
                                 time_step::Number=1e-3, oa=Hyperrectangle)
    problem = get_ivp(dyn, X0)
    if USE_REACHABILITY
        sets = compute_reach_tube_hull_Reachability(problem, time_horizon,
            time_step)
    else
        n = dim(X0)
        sets = compute_reach_tube_hull_ReachabilityAnalysis(problem,
            time_horizon, time_step, n)
    end
    hull = overapproximate(ConvexHullArray(sets), oa)
    return hull
end

function compute_reach_tube_hull_Reachability(problem, time_horizon, time_step)
    options = Options(:T => time_horizon)
    algorithm = BFFPSV18(:δ => time_step)
    solution = Reachability.solve(problem, options; op=algorithm)
    flowpipe = solution.flowpipes[1]
    sets = [Reachability.set(rs) for rs in flowpipe.reachsets]
    return sets
end

function compute_reach_tube_hull_ReachabilityAnalysis(problem, time_horizon,
                                                      time_step, n)
    algorithm = BOX(δ=time_step)
    solution = ReachabilityAnalysis.solve(problem, T=time_horizon, algorithm)
    flowpipe = solution.F
    sets = [ReachabilityAnalysis.set(rs) for rs in ReachabilityAnalysis.array(flowpipe)]
    return sets
end

function getInterval(lo, hi)
    return Interval(lo, hi)
end

function getBallInf(center, radius)
    return BallInf(center, radius)
end

function getHyperrectangle(center, radius)
    return Hyperrectangle(center, radius)
end

function getHalfSpace(a, b)
    return HalfSpace(a, b)
end

function getHPolyhedron(constraints)
    return HPolyhedron(constraints)
end

function getUniverse(dimension)
    return Universe(dimension)
end


# constructs a semantic theory about a commutative monoid
# A monoid whose operation is commutative is called a
# commutative monoid (or, less commonly, an abelian monoid).

include("docstrings.jl")

module Library

using Metatheory.Patterns
using Metatheory.Rules

macro associativity(op)
	quote
		[
			(@left_associative $op),
			(@right_associative $op)
		]
	end
end

macro monoid(op, id)
	quote
		[
			(@left_associative(op)),
			(@right_associative(op)),
			(@identity_left(op, id)),
			(@identity_right(op, id))
		]
	end
end

macro commutative_monoid(op, id)
	quote
		[
			(@commutativity $op),
			(@left_associative $op),
			(@right_associative $op),
			(@identity_left $op $id)
		]
	end
end

# constructs a semantic theory about a an abelian group
# The definition of a group does not require that a ⋅ b = b ⋅ a
# for all elements a and b in G. If this additional condition holds,
# then the operation is said to be commutative, and the group is called an abelian group.
macro commutative_group(op, id, invop)
	# @assert Base.isbinaryoperator(op)
	# @assert Base.isunaryoperator(invop)
	quote
		(@commutative_monoid $op $id) ∪ [@inverse_right $op $id $invop]
	end
end

macro distrib(outop, inop)
	quote
		[
			(@distrib_left $outop $inop),
			(@distrib_right $outop $inop),
		]
	end
end

macro commutativity(op)
	RewriteRule(
		PatTerm(:call, op, [PatVar(:a), PatVar(:b)], __module__),
		PatTerm(:call, op, [PatVar(:b), PatVar(:a)], __module__))
end

macro right_associative(op)
	RewriteRule(
		PatTerm(:call, op, [PatVar(:a),
			PatTerm(:call, op, [PatVar(:b), PatVar(:c)], __module__)], __module__),
		PatTerm(:call, op, [
			PatTerm(:call, op, [PatVar(:a), PatVar(:b)], __module__),
			PatVar(:c),
			], __module__))
end
macro left_associative(op)
	RewriteRule(
		PatTerm(:call, op, [
			PatTerm(:call, op, [PatVar(:a), PatVar(:b)], __module__),
			PatVar(:c),
			], __module__),
		PatTerm(:call, op, [PatVar(:a),
			PatTerm(:call, op, [PatVar(:b), PatVar(:c)], __module__)], __module__))
end

macro identity_left(op, id)
	RewriteRule(PatTerm(:call, op, [id, PatVar(:a)], __module__), PatVar(:a))
end

macro identity_right(op, id)
	RewriteRule(PatTerm(:call, op, [PatVar(:a), id], __module__), PatVar(:a))
end

macro inverse_left(op, id, invop)
	RewriteRule(PatTerm(:call, op, [
		PatTerm(:call, invop, [PatVar(:a)], __module__), PatVar(:a)], __module__), id)
end
macro inverse_right(op, id, invop)
	RewriteRule(PatTerm(:call, op, [
		PatVar(:a),
		PatTerm(:call, invop, [PatVar(:a)], __module__)], __module__), id)
end

# distributivity of two operations
# example: `@distrib (⋅) (⊕)`
macro distrib_left(outop, inop)
	EqualityRule(
		# left
		PatTerm(:call, outop, [
			PatVar(:a),
			PatTerm(:call, inop, [PatVar(:b), PatVar(:c)], __module__)
		], __module__),
		# right
		PatTerm(:call, inop, [
			PatTerm(:call, outop, [PatVar(:a), PatVar(:b)], __module__),
			PatTerm(:call, outop, [PatVar(:a), PatVar(:c)], __module__),
		], __module__))

end

macro distrib_right(outop, inop)
	EqualityRule(
		# left
		PatTerm(:call, outop, [
			PatTerm(:call, inop, [PatVar(:a), PatVar(:b)], __module__),
			PatVar(:c)
		], __module__),
		# right
		PatTerm(:call, inop, [
			PatTerm(:call, outop, [PatVar(:a), PatVar(:c)], __module__),
			PatTerm(:call, outop, [PatVar(:b), PatVar(:c)], __module__),
		], __module__))
end

# theory generation macros
export @commutativity
export @associativity
export @identity_left
export @identity_right
export @distrib_left
export @distrib_right
export @distrib
export @monoid
export @commutative_monoid
export @commutative_group

end


using LinearAlgebra
using Test

include("../egrss.jl")

include("test_generators.jl")
include("test_potrf.jl")
include("test_ldl.jl")
include("test_full.jl")
include("test_full_tril.jl")
include("test_gemv.jl")
include("test_symv.jl")
include("test_trmv.jl")
include("test_trsv.jl")
include("test_trtri.jl")
include("test_trnrms.jl")
include("test_bigfloat.jl")
include("test_rational.jl")


# Load Julia packages (libraries) needed  for the snippets in chapter 0

using DynamicHMCModels

ProjDir = @__DIR__
cd(ProjDir)

# Read in data

df = DataFrame(CSV.read(joinpath("..", "..", "data", "rugged.csv"), delim=';'))
df = filter(row -> !(ismissing(row[:rgdppc_2000])), df)
df[!, :log_gdp] = log.(df[:, :rgdppc_2000])
df[!, :cont_africa] = Array{Float64}(convert(Array{Int}, df[:, :cont_africa]))

Base.@kwdef mutable struct RuggedModel{Ty <: AbstractVector,
  Tx <: AbstractMatrix}
    "Observations."
    y::Ty
    "Covariates"
    x::Tx
end

# Write a function to return properly dimensioned transformation.

function make_transformation(model::RuggedModel)
  as((β = as(Array, size(model.x, 2)), σ = asℝ₊))
end

# Instantiate the model with data and inits.

x = hcat(ones(size(df, 1)), df[:, :rugged], df[:, :cont_africa],
  df[:, :rugged] .* df[:, :cont_africa]);
model = RuggedModel(;y=df[:, :log_gdp], x=x)

# Model callable with *a single argument*.

function (problem::RuggedModel)(θ)
    @unpack y, x = problem   # extract the data
    @unpack β, σ = θ            # works on the named tuple too
    ll = 0.0
    ll += logpdf(Normal(0, 100), x[1])
    ll += logpdf(Normal(0, 10), x[2])
    ll += logpdf(Normal(0, 10), x[3])
    ll += logpdf(Normal(0, 10), x[4])
    ll += logpdf(TDist(1.0), σ)
    ll += loglikelihood(Normal(0, σ), y .- x*β)
    ll
end

println()
model((β = [1.0, 2.0, 1.0, 2.0], σ = 1.0))
println()

# Wrap the problem with a transformation, then use Flux for the gradient.

P = TransformedLogDensity(make_transformation(model), model)
∇P = ADgradient(:ForwardDiff, P);

# Tune and sample.

results = mcmc_with_warmup(Random.GLOBAL_RNG, ∇P, 1000)
posterior = P.transformation.(results.chain)

println()
DynamicHMC.Diagnostics.EBFMI(results.tree_statistics) |> display
println()

DynamicHMC.Diagnostics.summarize_tree_statistics(results.tree_statistics) |> display
println()

a3d = Array{Float64, 3}(undef, 1000, 5, 1);
for j in 1:1
  for i in 1:1000
    a3d[i, 1:4, j] = values(posterior[i].β)
    a3d[i, 5, j] = values(posterior[i].σ)
  end
end

pnames = ["β[1]", "β[2]", "β[3]", "β[4]", "σ"]
sections =   Dict(
  :parameters => pnames,
)
chns = create_mcmcchains(a3d, pnames, sections, start=1);
chns = set_names(chns, Dict("β[1]" => "α", "β[2]" => "β[1]",
  "β[3]" => "β[2]", "β[4]" => "β[3]"))

# Result rethinking

rethinking = "
       mean   sd  5.5% 94.5% n_eff Rhat
a      9.22 0.14  9.00  9.46   282    1
bR    -0.21 0.08 -0.33 -0.08   275    1
bA    -1.94 0.24 -2.33 -1.59   268    1
bAR    0.40 0.14  0.18  0.62   271    1
sigma  0.96 0.05  0.87  1.04   339    1
"

# Summary

describe(chns)

# End of `08/m8.1s.jl`


function simulate_NDARMA(
    N::Int,
    p::Int,
    q::Int,
    prob::Vector{Float64},
    phi::Vector{Float64},
)
    X = zeros(Int64, N)
    X[1:p] = rand(Categorical(prob), p)
    d1 = Multinomial(1, phi)
    d2 = Categorical(prob)
    for t = (p+1):N
        alphabeta = rand(d1)
        eps = rand(d2, q + 1)
        X[t] = sum([X[(t-p):(t-1)]; eps] .* alphabeta)
    end
    return X
end

@testset "discretediag" begin
    Random.seed!(123)

    n = 10000
    x1 = simulate_NDARMA(n, 1, 0, [0.25, 0.5, 0.05, 0.2], [0.8, 0.2])
    x2 = simulate_NDARMA(n, 1, 0, [0.25, 0.5, 0.05, 0.2], [0.8, 0.2])
    x3 = simulate_NDARMA(n, 1, 0, [0.25, 0.5, 0.05, 0.2], [0.8, 0.2])

    sim = Chains(cat(x1, x2, x3, dims = 3))
    discretediag(sim)

    res = MCPhylo.weiss(hcat(x1, x2, x3))
    @test res[3] > 0.05

    x1 = simulate_NDARMA(n, 1, 0, [0.25, 0.5, 0.05, 0.2], [0.8, 0.2])
    x2 = simulate_NDARMA(n, 1, 0, [0.75, 0.15, 0.05, 0.05], [0.8, 0.2])
    x3 = simulate_NDARMA(n, 1, 0, [0.25, 0.5, 0.05, 0.2], [0.8, 0.2])
    res = MCPhylo.weiss(hcat(x1, x2, x3))
    @test res[3] < 0.05
end


"""
All sampling functions have three versions:
1. new* : function to create a new variable, return params: 1) the new variable (through return value or reference) 2) the proposal probability
2. del* : function to remove the old variable, return the proposal probability
3. shift*  : shift the old variable to a new one, return params: 1) the new variable (through return value or reference) 2) the proposal probability
"""

# Before using the sampling functions, call setRand function to reset the random number generator
# function setRand(_rand)
#     global const rand=_rand
# end
import Random:GLOBAL_RNG

@inline newIdx(size, rng=GLOBAL_RNG) = rand(rng, 1:size), Float(size)
@inline delIdx(oldIdx, size) = Float(1.0) / size
@inline shiftIdx(oldIdx, size, rng=GLOBAL_RNG) = rand(rng, 1:size), Float(1.0)

@inline newTau(β, rng=GLOBAL_RNG) = rand(rng, Float) * β, β
@inline delTau(oldTau, β) = Float(1.0) / β
@inline function shiftTau(oldTau, β, TauScale=β/4.0, rng=GLOBAL_RNG)
    x = rand(rng, Float)
    newTau::Float = 0.0
    if x < 1.0 / 3
        newTau = oldTau + TauScale * (rand(rng, Float) - Float(0.5))
    elseif x < 2.0 / 3
        newTau = -oldTau
    else
        newTau = rand(rng, Float) * β
    end

    if newTau < Float(0.0)
        newTau += β
    elseif newTau > β
        newTau -= β
    end
    return newTau, Float(1.0)
end

@inline function newK!(newK, Kf, Kscale=Kf/2.0, rng=GLOBAL_RNG)::Float
    ############ Simple Way ########################
    # for i in 1:DIM
    #     newK[i] = Kf * (rand(rng) - 0.5) * 2.0
    # end
    # return (2.0 * Kf)^DIM
    ################################################

    Kamp = Kf + (2.0*rand(rng) - 1.0) * Kscale
    (Kamp <= 0.0) && (return 0.0)
    # Kf-dK<Kamp<Kf+dK
    ϕ = 2π * rand(rng, Float)
    if length(newK) == 3
        θ = π * rand(rng, Float)
        newK[1] = Kamp * cos(ϕ) * sin(θ)
        newK[2] = Kamp * sin(ϕ) * sin(θ)
        newK[3] = Kamp * cos(θ)
        return 2Kscale * 2π * π * (sin(θ) * Kamp^2)
        # prop density of Kamp in [Kf-dK, Kf+dK), prop density of ϕ
        # prop density of θ, Jacobian
    elseif length(newK)==2
        newK[1] = Kamp * cos(ϕ)
        newK[2] = Kamp * sin(ϕ)
        return 2Kscale * 2π * Kamp
        # prop density of Kamp in [Kf-dK, Kf+dK), prop density of ϕ, Jacobian
    else
        throw(DimensionMismatch("Dimension $(length(newK)) is not supported!"))
    end
end

@inline function delK(oldK, Kf, Kscale=Kf/2.0)::Float
    ############## Simple Way #########################
    # for i in 1:DIM
    #     if abs(oldK[i]) > Kf
    #         return 0.0
    #     end
    # end
    # return 1.0 / (2.0 * Kf)^DIM
    ####################################################

    Kamp = norm(oldK)
    if !(Kf - Kscale < Kamp < Kf + Kscale)
        return 0.0
    end
    # (Kamp < Kf - dK || Kamp > Kf + dK) && return 0.0
    if length(oldK) == 3
        sinθ = sqrt(oldK[1]^2 + oldK[2]^2) / Kamp
        (sinθ < eps(Float)) && (return 0.0)
        return 1.0 / (2Kscale * 2π * π * sinθ * Kamp^2)
    elseif length(oldK)==2
        return 1.0 / (2Kscale * 2π * Kamp)
    else
        throw(DimensionMismatch("Dimension $(length(newK)) is not supported!"))
    end
end

@inline function shiftK!(oldK, newK, Kf, Kscale=Kf/3.0, λ=1.5, rng=GLOBAL_RNG)::Float
    x = rand(rng)
    D=length(newK)
    if x < 1.0 / 3
        # dK = β > 1.0 ? Kf / β * 3.0 : Kf
        # newK .= oldK .+ (rand(rng, DIM) .- 0.5) .* dK
        for i in 1:D
            newK[i] = oldK[i] + (rand(rng) - 0.5) * Kscale
        end
        return 1.0
    elseif x < 2.0 / 3
        ratio = 1.0 / λ + rand(rng) * (λ - 1.0 / λ)
        newK .= oldK .* ratio
        if D==3
            return ratio
        elseif D==2
            return 1.0
        else
            throw(DimensionMismatch("Dimension $(length(newK)) is not supported!"))
        end
    else
        newK .= oldK .* (-1.0)
        return 1.0
    end
end


########################################################################
# Expected item info -  for 1PL pars and 1-D latent                    #
########################################################################
"""
```julia
_item_expected_information(parameters::Parameters1PL, latent::Latent1D)
```

# Description

It computes the expected information (-second derivative of the likelihood) with respect to the difficulty parameter of the 1PL model.
It follows the parametrization ``a(θ - b)``.

# Arguments

- **`parameters::Parameters1PL`** : Required. A 1-parameter logistic parameters object.
- **`latent::Latent1D`** : Required. A 1-dimensional `Latent1D` latent variable.

# Output

A `Float64` scalar.
"""
function _item_expected_information(parameters::Parameters1PL, latent::Latent1D)
    p = _probability(latent, parameters)
    return p * (1 - p)::Float64
end

########################################################################
# Expected info -  2PL pars and 1-D latent                             #
########################################################################
"""
```julia
_item_expected_information(parameters::Parameters2PL, latent::Latent1D)
```

# Description

It computes the expected information (-second derivative of the likelihood) with respect to the 2 parameters of the 2PL model.
It follows the parametrization ``a(θ - b)``.

# Arguments

- **`parameters::Parameters1PL`** : Required. A 2-parameter logistic parameters object.
- **`latent::Latent1D`** : Required. A 1-dimensional `Latent1D` latent variable.

# Output

A ``2 \times 2`` matrix of the expected informations.
"""
function _item_expected_information(parameters::Parameters2PL, latent::Latent1D)
    p = _probability(latent, parameters)
    p_1_p = (1 - p) * p
    i_aa = p_1_p * (latent.val - parameters.b)^2
    i_ab = -parameters.a * p_1_p * (latent.val - parameters.b)
    i_bb = parameters.a^2 * p_1_p
    return [i_aa i_ab; i_ab i_bb]::Matrix{Float64}
end

########################################################################
# Expected info -  for 3PL pars and 1-D latent                         #
########################################################################
"""
```julia
_item_expected_information(parameters::Parameters3PL, latent::Latent1D)
```

# Description

It computes the expected information (-second derivative of the likelihood) with respect to the 3 parameters of the 3PL model.
It follows the parametrization ``a(θ - b)``.
It should be always positive.

# Arguments
- **`parameters::Parameters1PL`** : Required. A 3-parameter logistic parameters object.
- **`latent::Latent1D`** : Required. A 1-dimensional `Latent1D` latent variable.

# Output
A ``3 \times 3`` matrix of the expected informations.
"""
function _item_expected_information(parameters::Parameters3PL, latent::Latent1D)
    p = _probability(latent, parameters)
    den = p * (1 - parameters.c)^2
    theta_b = latent.val - parameters.b
    p_c = p - parameters.c
    i_cc = (1 - p) / den # vv p. 51 Kim, Baker, 2008 : - E(L_33)
    i_aa = i_cc * p_c^2 * theta_b^2  # v p. 51 Kim, Baker, 2008
    i_ab = - i_aa / theta_b * parameters.a
    i_ac = i_cc * p_c * theta_b
    i_bc = i_cc * p_c * parameters.a
    i_bb = (i_bc ^ 2) / i_cc #v p. 51 Kim, Baker, 2008
    return [i_aa i_ab i_ac; i_ab i_bb i_bc; i_ac i_bc i_cc]::Matrix{Float64}
end


using CyclicArrays
using Plots

nx=ny=100
nfaces=1;

faces=zeros(nfaces,2,2,4);
faces[1,1,1,:]=[1,1,2,0];
faces[1,1,2,:]=[1,1,1,0];
faces[1,2,1,:]=[1,2,2,0];
faces[1,2,2,:]=[1,2,1,0];

g=CyclicArray(faces);

ind=range(1, 360, length = 360)
x = y = ind/180*pi
f(y,x) = sin(x) + sin(y)

X = repeat(reshape(x, 1, :), length(y), 1)
Y = repeat(y, 1, length(x))

F=f.(Y,X)

#f = sin.(2*x .+ 10 .* sin.(x)) .+ sin.(3*y)

phi=CyclicArray(F,g)
contourf(phi.data, aspect_ratio=1,xlim=(0,360),ylim=(0,360))
xg=CyclicArray(repeat(ind, 1, length(x))',g)
yg=CyclicArray(repeat(ind, 1, length(x)),g)
u=diff(phi,dims=2)
v=-diff(phi,dims=1)

N=1500
np=25
xi=zeros(N+1,np)
yi=zeros(N+1,np)

x0=collect(0:20:490).%100 .+ 170
y0=reshape(repeat(collect(0:20:80),1,5)',25) .+ 160
xi[1,:]=x0;
yi[1,:]=y0;

dt=100
#for t in 1:N
plt=contourf(yg.data[:,1],xg.data[1,:],phi.data, aspect_ratio=1,xlim=(0,360),ylim=(0,360))
scatter!(xi[1,:],yi[1,:],leg=false)
anim = @animate for t ∈ 1:N
  println(t)
  i=Int.(floor.(xg[1,Int.(floor.(xi[t,:]))]))
  j=Int.(floor.(yg[Int.(floor.(yi[t,:])),1]))
  U=[u[i[ii],j[ii]] for ii in 1:np]
  V=[v[i[ii],j[ii]] for ii in 1:np]
  #println(t,' ',U,' ',V)
  xi[t+1,:]=xi[t,:]+U*dt
  yi[t+1,:]=yi[t,:]+V*dt
  xi[t+1,:]=floor.(xg[1,Int.(sign.(xi[t+1,:]).*floor.(abs.(xi[t+1,:])))])+rem.(xi[t+1,:],1)
  yi[t+1,:]=floor.(yg[Int.(sign.(yi[t+1,:]).*floor.(abs.(yi[t+1,:]))),1])+rem.(yi[t+1,:],1)
  plt=contourf(yg.data[:,1],xg.data[1,:],phi.data, aspect_ratio=1,xlim=(0,360),ylim=(0,360))
  scatter!([-1,xi[t+1,:]],[-1,yi[t+1,:]],leg=false, c=:green)
  #display(plt)
end
gif(anim, "RandomFlow_2D.gif", fps = 30)


using ODEInterfaceDiffEq, Test

callback_f = function (du,u,p,t)
  du[1] = u[2]
  du[2] = -9.81
end

condtion= function (u,t,integrator) # Event when event_f(u,t,k) == 0
  u[1]
end

affect! = nothing
affect_neg! = function (integrator)
  integrator.u[2] = -integrator.u[2]
end

callback = ContinuousCallback(condtion,affect!,affect_neg!)

u0 = [50.0,0.0]
tspan = (0.0,25.0)
prob = ODEProblem(callback_f,u0,tspan)

sol = solve(prob,dopri5(),callback=callback,dtmax=0.5)
@test sol(4.0)[1] > 0
sol = solve(prob,dopri5(),callback=callback,save_everystep=true)
@test sol(4.0)[1] > -1e-12


function run(s)
    # Your code here
    code = readlines(IOBuffer(s))
    seen = Set{Int}()
    curr = 1
    acc = 0
    while !(curr in seen)
        push!(seen, curr)
        line = code[curr]
        op = line[1:3]
        if op == "acc"
            acc += parse(Int, line[5:end])
            curr += 1
        elseif op == "jmp"
            curr += parse(Int, line[5:end])
        else
            curr += 1
        end
    end
    return acc
end

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

function main()
    res, time, memory = @timed run(ARGS[1])
    println("_duration:$(time * 1000)")
    println(res)
end

main()


"""
Three bouncing balls
q = (z1,z2,z3), u = (z1dot, z2dot, z3dot),
g(q) = (q[1],q[2],q[3])
W = del g / del q = I (3 x 3 identity)
M = [m1 0 0; 0 m2 0; 0 0 m3]
h = [-m1*g; -m2*g; -m3*g]
"""

using LinearAlgebra
using PyPlot
using TimeStepping

const m1 = m2 = m3 = 1.0
const g = 9.81
final_time = 2.0
q0 = [0.8; 1.0; 1.2]
u0 = [0.5; 2.0; -0.2]

function gap(q::AbstractArray, u::AbstractArray)
    return q, Matrix(1.0f0*I, 3, 3)
end

function dynamics(q::AbstractArray, u::AbstractArray)
    return convert.(Float32, [m1 0 0; 0 m2 0; 0 0 m3]), convert.(Float32, [-m1*g; -m2*g; -m3*g])
end

three_bb = Integrator(gap, dynamics, convert.(Float32, q0), convert.(Float32, u0); Δt=convert(Float32, 1e-3), ε=0.5f0)
integrate(three_bb, final_time)

fig = figure(1, clear=true)
fig.add_subplot(1,1,1)
fig.axes[1].clear()
fig.axes[1].plot(three_bb.t, getindex.(three_bb.q,1), "b", linewidth=2, label="ball 1")
fig.axes[1].plot(three_bb.t, getindex.(three_bb.q,2), "k", linewidth=2, label="ball 2")
fig.axes[1].plot(three_bb.t, getindex.(three_bb.q,3), "m", linewidth=2, label="ball 3")

fig.axes[1].set_title("Three bouncing balls", fontsize=18)
fig.axes[1].set_xlabel(L"$t$", fontsize=18)
fig.axes[1].set_ylabel(L"$z(t)$", fontsize=18)
fig.axes[1].legend(loc="best", fontsize="x-large")

fig.axes[1].autoscale_view()
fig.tight_layout()


using DelimitedFiles: readdlm
using ProgressBars

function load_edges(filename, delim, comments)
    return readdlm(filename, delim, Int, comments=comments)
end

function kcore(nodes::Vector{Int64}, edges::Matrix{Int64})
    if minimum(edges) == 0
        edges = edges .+ 1
        nodes = Vector(1:length(nodes)+1)
    end
    n = length(nodes)
    i = n
    c = 0
    # Compute degrees
    degrees = zeros(Int64, n)
    for edge in eachrow(edges)
        degrees[edge[1]] += 1
        degrees[edge[2]] += 1
    end

    for i in ProgressBar(1:n)
        index = argmin(degrees)
        node = nodes[index]
        c = max(c, degrees[index])
        # Update graph
        deleteat!(nodes, index)

        neighbor_edges_indices = Vector{Int64}()

        for (k, edge) in enumerate(eachrow(edges))
            if edge[1] == node || edge[2] == node
                push!(neighbor_edges_indices, k)
            end
        end

        edges = edges[findall(x->!(x in neighbor_edges_indices), 1:size(edges, 1)),:]

        degrees = zeros(Int64, length(nodes))
        for edge in eachrow(edges)
            index1 = findall(x->x==edge[1], nodes)[1]
            index2 = findall(x->x==edge[2], nodes)[1]
            degrees[index1] += 1
            degrees[index2] += 1
        end
        i = i - 1
    end

    return c
end

filename = "com-amazon.ungraph.txt"
edges = load_edges(filename, '\t', true)
nodes = Vector(1:maximum(edges))
core_value = kcore(nodes, edges)


"""
Uses `hunting-circular-orbits.jl` and the Page & Thorne (1974) eq. (12)
to simulate a flux model.

Works for a < 0.97, then the orbit finder starts to panic close to the ISCO.
"""

# using Revise

using GeodesicTracer
using ComputedGeodesicEquations
using GeodesicBase

import CarterBoyerLindquist # for rms function
import LinearAlgebra

using StaticArrays
using Plots
gr()

# hoist other functions
include("hunting-circular-orbits.jl")

# for Kerr, this is just `u[2]`, but to illustrate
# how you could do this for other metrics
function sqrt_g_tilde(m, u)
    metric = GeodesicBase.metric(m, u)
    metric_no_θ = metric[setdiff(1:4, 3), setdiff(1:4, 3)]

    sqrt(-LinearAlgebra.det(metric_no_θ))
end

# gradient function
#
# we just compute discrete derivatives, since this is
# sufficient when we have small `r`.
# using interpolations can make the resolution better
# and increase the accuracy, but this is a good quick-and-dirty way
# to see if our results are going in the right direction
function ∇r(array, r)
    deriv = diff(array) ./ diff(r)
    # ensure same dimension using a lazy extension
    push!(deriv, deriv[end])
    deriv
end

"""
    f_approx(Lz, ∇rwt, ∇rLz_array, wt_array)

A discretized version of `f` from Page & Thorne (1974) eq. (12):

```math
f(r_n) \\approx - \\left. \\left[
   \\frac{1}{L_z} \\left(\\frac{\\partial \\dot{u}^t}{\\partial r} \\right)
\\right] \\right\\rvert_{r = r_n}
\\sum_{i=0}^{n} \\left\\{ \\left. \\left[
   \\frac{1}{\\dot{u}^t} \\left( \\frac{\\partial L_z}{\\partial r}\\right)
\\right] \\right\\rvert_{r = r_i}
\\right\\}
```

"""
function f_approx(Lz, ∇rwt, ∇rLz_array, wt_array)
    -(∇rwt / Lz) * sum(∇rLz_array ./ wt_array)
end

# use this to see where things break down / if they break down
function plot_derivative_test(r_range, res)
    lzs = res[:, 2]
    # naive derivatives
    derivs = ∇r(lzs, r_range)

    plot(r_range, lzs, label = "Lz", c = :black, legend = false)
    plot!(r_range, derivs, label = "derivative", c = :black, style = :dot, lw = 1.2)
end

# test plot
function f_approx_test(r_range, res)
    wts = res[:, 4]
    Lzs = res[:, 2]

    ∇rwt_array = ∇r(wts, r_range)
    ∇rLz_array = ∇r(Lzs, r_range)

    fs = map(2:length(r_range)) do i
        @views f_approx(Lzs[i], ∇rwt_array[i], ∇rLz_array[1:i], wts[1:i])
    end

    # we rescale here for a better comparison
    fs = fs ./ maximum(fs)

    # we lose a point, so select all but last `r_range`
    # since our derivative function needs at least 2 points to work
    plot(r_range[2:end], fs, c = :black, xlabel = "r", ylabel = "f", label = "simulated f")
end

# this works up until about 0.97 then, the orbit search
# starts to break down
m = BoyerLindquist(M = 1.0, a = 0.97)
u = @SVector [0.0, 13.0, deg2rad(90), 0.0]

println("Starting approximation...")

# making the steps too big over-estimates the gradient
r_range = CarterBoyerLindquist.rms(m.M, m.a):0.01:30.0
res = find_orbit_range(; r_range = r_range, a = m.a)

# use this to check if the derivatives are okay
# plot_derivative_test(r_range, res)

f_approx_test(r_range, res)
println("Approximation done!")

# compare to "true" value
import AccretionFormulae
f_true = AccretionFormulae.flux.(r_range, m.a, m.M)
# rescale as well
f_true = f_true ./ maximum(f_true)
plot!(r_range, f_true, label = "true f")

# savefig("simulated-f-not-log.svg")


# ---
# title: 972. Equal Rational Numbers
# id: problem972
# author: Indigo
# date: 2021-02-20
# difficulty: Hard
# categories: Math
# link: <https://leetcode.com/problems/equal-rational-numbers/description/>
# hidden: true
# ---
#
# Given two strings `S` and `T`, each of which represents a non-negative
# rational number, return **True** if and only if they represent the same
# number. The strings may use parentheses to denote the repeating part of the
# rational number.
#
# In general a rational number can be represented using up to three parts: an
# _integer part_ , a  _non-repeating part,_ and a  _repeating part_. The number
# will be represented in one of the following three ways:
#
#   * `<IntegerPart>` (e.g. 0, 12, 123)
#   * `<IntegerPart><.><NonRepeatingPart>`  (e.g. 0.5, 1., 2.12, 2.0001)
#   * `<IntegerPart><.><NonRepeatingPart><(><RepeatingPart><)>` (e.g. 0.1(6), 0.9(9), 0.00(1212))
#
# The repeating portion of a decimal expansion is conventionally denoted within
# a pair of round brackets.  For example:
#
# 1 / 6 = 0.16666666... = 0.1(6) = 0.1666(6) = 0.166(66)
#
# Both 0.1(6) or 0.1666(6) or 0.166(66) are correct representations of 1 / 6.
#
#
#
# **Example 1:**
#
#
#
#     Input: S = "0.(52)", T = "0.5(25)"
#     Output: true
#     Explanation: Because "0.(52)" represents 0.52525252..., and "0.5(25)" represents 0.52525252525..... , the strings represent the same number.
#
#
# **Example 2:**
#
#
#
#     Input: S = "0.1666(6)", T = "0.166(66)"
#     Output: true
#
#
# **Example 3:**
#
#
#
#     Input: S = "0.9(9)", T = "1."
#     Output: true
#     Explanation:
#     "0.9(9)" represents 0.999999999... repeated forever, which equals 1.  [[See this link for an explanation.](https://en.wikipedia.org/wiki/0.999...)]
#     "1." represents the number 1, which is formed correctly: (IntegerPart) = "1" and (NonRepeatingPart) = "".
#
#
#
# **Note:**
#
#   1. Each part consists only of digits.
#   2. The `<IntegerPart>` will not begin with 2 or more zeros.  (There is no other restriction on the digits of each part.)
#   3. `1 <= <IntegerPart>.length <= 4 `
#   4. `0 <= <NonRepeatingPart>.length <= 4 `
#   5. `1 <= <RepeatingPart>.length <= 4`
#
#
## @lc code=start
using LeetCode

function is_rational_equal(s::String, t::String)
    function to_rational(s::String)::Rational
        pos1, pos2 = findfirst('.', s), findfirst('(', s)
        (pos1 === nothing) && return parse(Int, s)
        (pos1 == length(s)) && return parse(Int, s[1:end-1])
        pos2 === nothing && return parse(Int, s[1:pos1-1]) + parse(Int, s[pos1 + 1:end]) // 10 ^ (length(s) - pos1)
        pos3 = findlast(')', s)
        int_part = parse(Int, s[1:pos1-1]) |> Rational
        nr_part = pos1 == pos2 - 1 ? 0 : parse(Int, s[pos1+1:pos2-1]) // 10 ^ (pos2 - pos1 - 1)
        rp_part = parse(Int, s[pos2+1:pos3-1]) // (10 ^ (pos3 - pos2 - 1) - 1) // 10 ^ (pos2 - pos1 - 1)
        int_part + nr_part + rp_part
    end
    to_rational(s) == to_rational(t)
end
## @lc code=end
is_rational_equal("0.9(9)", "1.")


using Combinatorics
n = 4
m = 3
for i in combinations(0:n,m)
    println(i')
end


using Unitful

export ploss_coeff, ploss_coeff_nice, isunitless

function nounit(value)
    try
        uconvert(Unitful.NoUnits, value)
        true
    catch
        false
    end
end

# Source for notation and formulas:
# "Evaluating Gas Network Capacities", T. Koch et al., SIAM MO21, 2015
#
# Some constants are taken from that book, or the related source code of
# Lamatto++, see: http://www.mso.math.fau.de/edom/projects/lamatto.html

# pipe roughness, from REKO data (4390 * 0.012mm, 926 * 0.006mm, 166 * 0.002mm)
const k = 0.012u"mm"

# pipe (mean) diameter for friction factor
const D_m = 1.0u"m"

"friction factor by Nikoradse, for turbulent flow"
function nikoradse(diameter, roughness)
    (log10(diameter/roughness) + 1.138)^(-2)
end

# friction factor
const λ = nikoradse(D_m, k)

# universal gas constant (https://en.wikipedia.org/wiki/Gas_constant)
const R = 8.3144598u"J/K/mol"

# molar mass of methane
const molmass = 16.04u"g/mol"

# specific gas constant
const R_s = R/molmass

# critical pressure (default in Lamatto++)
const p_c = 46.4512u"bar"

# critical temperature (default in Lamatto++)
const T_c = 192.033u"K"

# mean pressure
const p_m = 60u"bar"

# mean temperature (default in Lamatto++)
const T_m = 283.15u"K"

"compressibility factor by AGA8"
function aga8(pressure, temperature)
    # "reduced" pressure and temperature
    p_r = pressure / p_c
    T_r = temperature / T_c
    z = 1.0 + 0.257 * p_r - 0.533 * p_r/T_r
end

# z_m: mean compressibility factor
const z_m = aga8(p_m, T_m)

# Pressure Loss equation ("Weymouth"), when we care about length and diameter
#   p_o^2 - p_i^2 = L/D^5 C q|q|
const weymouth = (π/4)^(-2) * λ * R_s * z_m * T_m
@assert !nounit(weymouth)

# alternative coefficient with same order of magnitude for nice, round numbers
const weymouth_nice = 3000u"m^2/s^2"
@assert nounit(weymouth/weymouth_nice) # same unit

# Unit-less Weymouth coefficient: Add types for length, diameters and flow, then
# divide by the expected type to get just a float
const ploss_coeff_unit = 1u"(km*m^(-5)*(kg/s)^2) / (bar^2)"
const ploss_coeff = weymouth * ploss_coeff_unit
const ploss_coeff_nice = weymouth_nice * ploss_coeff_unit
@assert nounit(ploss_coeff)
@assert nounit(ploss_coeff_nice)


# exact k-medoids modeled similar as in Kotzur et al, 2017

"Holds results of kmedoids run"
mutable struct kmedoidsResult
    medoids::Array{Float64}
    assignments::Array{Int}
    totalcost::Float64
end

"""
    kmedoids_exact(
     data::Array{Float64},
     nclust::Int,
     _dist::SemiMetric = SqEuclidean(),
     env::Any;
     )
   results = kmedoids_exact()
   data { HOURS,DAYS }
Performs the exact kmedoids algorithm as in Kotzur et al, 2017
optimizer=Gurobi.Optimizer
"""
function kmedoids_exact(
   data::Array{Float64},
   nclust::Int,
   optimizer::DataType;
   _dist::SemiMetric = SqEuclidean(),
   )
   N_i = size(data,2)

  # calculate distance matrix
  d_mat=pairwise(_dist,data, dims=2)

  # create jump model
  m = JuMP.Model(with_optimizer(optimizer)) # GurobiSolver(env,OutputFlag=0)
  @variable(m,z[1:N_i,1:N_i],Bin)
  @variable(m,y[1:N_i],Bin)
  @objective(m,Min,sum(d_mat[i,j]*z[i,j] for i=1:N_i, j=1:N_i))
  for j=1:N_i
    @constraint(m,sum(z[i,j] for i=1:N_i)==1)
  end
  for i=1:N_i
    for j=1:N_i
    @constraint(m,z[i,j]<=y[i])
    end
  end
  @constraint(m,sum(y[i] for i=1:N_i) == nclust)

  # solve jump model
  optimize!(m)
  status=Symbol(termination_status(m))

  println("status: ",status)
  y_opt=round.(Integer,JuMP.value.(y))
  z_opt=round.(Integer,JuMP.value.(z))
   #println("y ",y_opt, " z ",z_opt)
  # determine centers and cluster mappings
  id = zeros(Int,N_i)
  ii=0
  for i=1:N_i
    if y_opt[i]==1
      ii +=1
      id[i]=ii
    end
  end
  nz = findall(!iszero,z_opt)
  centerids=Int[]
  for i=1:length(nz) # just take the first dimension of each of the elements of the array of cartesianCoordinates
    push!(centerids, nz[i][1])
  end
  clustids = zeros(Int,N_i)
  for i=1:N_i
    clustids[i] = id[centerids[i]]
  end
  centers = data[:,findall(id.!=0.0)]
  tot_dist = objective_value(m)
   # output centers
  results = kmedoidsResult(centers, clustids, tot_dist)

  return results

end #kmedoids_exact


#=====================================================#
#== MATHEMATICS AND ARITHMETIC TEMPLATES FOR HUMANS ==#
#=====================================================#
module MATH

using LinearAlgebra

#===========================================#
#== first use case: function manipulation ==#
#===========================================#
#== general function sum ==#
macro ∑(bounds::UnitRange{Int64}, op::Function)
  return :( general_driver(+, $(esc(bounds)) , $(esc(op)) ) )
end

#== general function product ==#
macro ∏(bounds::UnitRange{Int64}, op::Function)
  return :( general_driver(*, $(esc(bounds)) , $(esc(op)) ) )
end

#== common driver ==#
function general_driver(funct::Function, bounds::UnitRange{Int64}, op::Function)
  #GC.enable(false)
  return reduce(funct, map(op, collect(bounds)))
  #GC.enable(true)
  #GC.gc()
end

#================================================#
#== second use case: array/matrix manipulation ==#
#================================================#
#== summations involving single arrays ==#
macro ∑(array)
  return :( sum_over_one( $(esc(array)) ) )
end

@inline function sum_over_one(array::T) where {T<:AbstractArray}
  return sum(array)
end

@inline function sum_over_one(range::UnitRange{T}) where {T<:Number}
  return sum(collect(range))
end

#== summations involving two arrays ==#
macro ∑(array_1, array_2)
  return :( sum_over_two( $(esc(array_1)), $(esc(array_2)) ) )
end

@inline function sum_over_two(array_1::T,
  array_2::T) where {T<:AbstractArray}

  return LinearAlgebra.dot(array_1,array_2)
end

@inline function sum_over_two(array_1::Array{T,1},
  array_2::Array{T,1}) where {T<:Number}

  return LinearAlgebra.dot(array_1,array_2)
end

@inline function sum_over_two(array_1::Array{T,2},
  array_2::Array{T,2}) where {T<:Number}

  return LinearAlgebra.dot(array_1,array_2)
end

#== quadratic transformation summation ==#
macro ∑(array_1,array_2,array_center)
  return :( quadratic_transform( $(esc(array_1)), $(esc(array_2)) ,
    $(esc(array_center)) ) )
end

@inline function quadratic_transform(array_1::Array{T},
  array_2::Array{T}, array_center::Array{T,2}) where {T<:Number}

  return transpose(array_2)*array_center*array_1
end

#== products involving single arrays ==#
macro ∏(array)
  return :( prod_over_one( $(esc(array)) ) )
end

@inline function prod_over_one(array::T) where {T<:AbstractArray}
  return prod(array)
end

@inline function prod_over_one(range::UnitRange{T}) where {T<:Number}
  return prod(collect(range))
end

export @∑
export @∏

end


# Machine Learning
# Data Preprocessing Array functions

using CSV
using Statistics
using Missings

"""
```
train_test_split(x::Array{T}; split_ratio::Float64=.7)
```
Split Data Into Train and Test Sets
"""
function train_test_split(x::Array{T}; split_ratio::Float64=.7)::Array{T} where {T<:Real}
    let len = length(len), split_no = Int64.(round(split_ratio * len))
    end
    train_set = data[1:split_no]
    test_set = data[split_no+1:len]
    return train_set,test_set
end

"""
```
standardization(x::Array{T})
```
Feature Scaling
standardization = x - mean(x) / standard deviation(x)
"""
function standardization(x::Array{T})::Array{T} where {T<:Real}
    let out = fill(NaN, size(x,1)), avg = mean(x), stdev = std(x)
    @inbounds for i = 1:size(x,1)
        out[i] = (x[i] - avg) / stdev
end
    return out
end
end

"""
```
normalization(x::Array{T})
```
Feature Scaling
normalization = x - min(x) / max(x) - min(x)
"""
function normalization(x::Array{T})::Array{T} where {T<:Real}
    let out = fill(NaN, size(x,1)), min = minimum(x), max = maximum(x)
    @inbounds for i = 1:size(x,1)
        out[i] = ((x[i] - min) / (max - min))
end
    return out
end
end

"""
```
data_fill_nan_inf(x::Array{T})
```
Convert NaN to Inf
"""
function data_fill_nan_inf(x::Array{T})::Array{T} where {T<:Real}
        @inbounds for i = 1:size(x,1)
        if (isnan(x[i]))
                x[i] = Inf
            end
end
return x
end

x = data_fill_nan_inf(x)

"""
```
fill_missing(x::Array{T}; method::String="mean")
```
Run function data_fill_nan_inf() this function thens replace Inf with the mean, median or previous values
"""
function fill_missing(x::Array{T}; method::String="mean")::Array{T} where {T<:Real}
    y = x[x .!= Inf] # Subset to remove InF, NaN
        @inbounds for i = 1:size(x,1)
        if method == "mean"
            avg = mean(y)
            if x[i] == Inf
                x[i] = avg
            end
        elseif method == "median"
            med = median(y)
            if x[i] == Inf
                x[i] = med
            end
        elseif method == "prev" && i >= 1
            prev = x[i-1]
            if x[i] == Inf
                x[i] = prev
            end
        elseif method == "prev" && i == 1
            prev = x[i+1]
            if x[i] == Inf
                x[i] = prev
            end
end
end
return x
end

# Data preprocessing - working with dataframes

# Read data
dataset = CSV.read("C:/Users/Andrew.Bannerman/Desktop/Machine Learning/Machine-Learning/Machine-Learning/data_preprocessing/Data.csv", header=true,types=[String, Float64, Float64, String])

# Replace missing values with the mean
using Missings
dataset[:Age] = collect(Missings.replace(dataset[:Age], mean(skipmissing(dataset[:Age]))))
dataset[:Salary] = collect(Missings.replace(dataset[:Salary], mean(skipmissing(dataset[:Salary]))))

# One hot encoding categroial variables
for i in unique(dataset[:Country])
    dataset[Symbol(i)] = ifelse.(dataset[:Country] .== i, 1, 0)
end

dataset[:Purchased] = ifelse.(dataset[:Purchased] .== "Yes", 1, 0)

# Feature Scaling - note in order to not have multiple variables over powering each other in the model
dataset[:Age] = standardization(dataset[:Age])
dataset[:Salary] = standardization(dataset[:Salary])

# Split dataframe to train and test sets
split_ratio = .7
len = size(dataset,1)
split_no = Int64.(round(split_ratio * len))
train_set = dataset[1:split_no,:]
test_set = dataset[split_no+1:len,:]


# m6.4s.jl

using Pkg, DrWatson

using MonteCarloMeasurements
using StanSample
using StatisticalRethinking

df = CSV.read(sr_datadir("milk.csv"), DataFrame; delim=';');
scale!(df, [:kcal_per_g, :perc_fat, :perc_lactose])

stan6_4 = "
data{
    int <lower=1> N;              // Sample size
    vector[N] K;
    vector[N] L;
}
parameters{
    real a;
    real bL;
    real<lower=0> sigma;
}
model{
    vector[N] mu;
    sigma ~ exponential( 1 );
    a ~ normal( 0 , 0.2 );
    bL ~ normal( 0 , 0.5 );
    mu = a + bL * L;
    K ~ normal( mu , sigma );
}
";

# Define the SampleModel and set the output format to :mcmcchains.

m6_4s = SampleModel("m6.4s", stan6_4);
m6_4_data = Dict("N" => size(df, 1), "L" => df[:, :perc_lactose_s],
    "K" => df[!, :kcal_per_g_s]);
rc6_4s = stan_sample(m6_4s, data=m6_4_data);
if success(rc6_4s)
    part6_4s = read_samples(m6_4s, :particles)
    part6_4s |> display
end

# End of m6.4s.jl


@testset "Cylindrical" begin
  f = ToCartesian(r->[r[1] * cos(r[2]), r[1] * sin(r[2]), r[3]], 3)
  f⁻¹ = FromCartesian(x->[sqrt(x[1]^2 + x[2]^2), atan(x[2], x[1]), x[3]], 3)
  numiters = 1

  @testset "inversions" begin
    r = [π, exp(1), √2] # r, θ, z
    @test f⁻¹(f(r)) ≈ r
    @test f(f⁻¹(r)) ≈ r

    @test invert(f)(r) ≈ f⁻¹(r)
    @test invert(f⁻¹)(r) ≈ f(r)

    v1 = ContravariantVector(z->[prod(z), prod(z)^2, prod(z)^3], f, true)
    v2 = ContravariantVector(z->[prod(z), prod(z)^2, prod(z)^3], f⁻¹, true)
    @test v1(r) ≈ v2(r)
    v1 = CovariantVector(z->[prod(z), prod(z)^2, prod(z)^3], f, true)
    v2 = CovariantVector(z->[prod(z), prod(z)^2, prod(z)^3], f⁻¹, true)
    @test v1(r) ≈ v2(r)
  end

  @testset "metrics, jacobians, line elements" begin
    r = [π, exp(1), √2] # r, θ, z
    @test gᵢⱼ(f, r) ≈ inv(gⁱʲ(f, r))
    @test gᵢⱼ(f⁻¹, r) ≈ inv(gⁱʲ(f⁻¹, r))
    @test gᵢⱼ(f⁻¹, r) ≈ gᵢⱼ(f, r)
    @test gⁱʲ(f⁻¹, r) ≈ gⁱʲ(f, r)
    @test J(f⁻¹, r) ≈ J(f, r)
    @test J⃗(f⁻¹, r) ≈ J⃗(f, r)
    @test dl(f⁻¹, r, Covariant) ≈ dl(f, r, Covariant)
    @test dl(f⁻¹, r, Contravariant) ≈ dl(f, r, Contravariant)
  end

  @testset "unit basis / physical units expressions" begin
    r = [π, exp(1), √2] # r, θ, z
    x = f(r) # x, y, z
    g(z) = [prod(z), prod(z)^2, prod(z)^3]
    v1 = ContravariantVector(z->g(z), f, true)
    v2 = ContravariantVector(z->g(f⁻¹(z)), f⁻¹, true)
    @test v1(r, true) ≈ v2(x, true)

    v1 = CovariantVector(z->g(z), f, true)
    v2 = CovariantVector(z->g(f⁻¹(z)), f⁻¹, true)
    @test v1(r, true) ≈ v2(x, true)
  end

  function getcoordsandtransforms(ct)
    if typeof(ct) <: ToCartesian
      R = [rand(), (rand() - 0.5) * 2π , randn()]
      X = f(R)
    else
      X = randn(3)
      R = f⁻¹(X)
    end
    fR(x) = typeof(ct) <: ToCartesian ? x : f⁻¹(x)
    # assert logic is correct or the tests are meaningless
    @assert f(R) ≈ X
    @assert f⁻¹(X) ≈ R
    @assert fR(X) ≈ R
    return (R,X,fR)
  end

  @testset "metrics, jacobians, line elements" begin
    for ct ∈ (f, f⁻¹), _ in 1:numiters
      (R,X,fR) = getcoordsandtransforms(ct)

      @test gᵢⱼ(ct, R) ≈ [1 0 0; 0 R[1]^2 0; 0 0 1]
      @test gⁱʲ(ct, R) ≈ [1 0 0; 0 1/R[1]^2 0; 0 0 1]
      @test dl(ct, R, Covariant) ≈ [1, 1/R[1], 1]
      @test dl(ct, R, Contravariant) ≈ [1, R[1], 1]

      @test J(ct, R) ≈ R[1]
    end
  end

  @testset "divergence" begin
    for ct ∈ (f, f⁻¹), _ in 1:numiters
      (R,X,fR) = getcoordsandtransforms(ct)

      div10r(r) = [2 * r[1]/2, 3 * r[1] * r[2], 5 * r[3]]
      v = ContravariantVector(z->div10r(fR(z)), ct, true)
      @test div(v)(R) ≈ 10
      if !(div(v)(R) ≈ 10)
        @show ct
      end

      p, m, a = rand(2:5), rand(2:5), rand(2:5)
      fr = r -> r[1]^p
      fθ = r -> sin(m * r[2])
      fz = r -> r[3]^a
      dfr(r) = p * r[1]^(p - 1)
      dfθ(r) = m * cos(m * r[2])
      dfz(r) = a * r[3]^(a - 1)
      divf(r) = dfr(r) + fr(r)/r[1] + dfθ(r)/r[1] + dfz(r)
      v = CovariantVector(r->[fr(r), fθ(r), fz(r)], f, true)
      @test div(v)(R) ≈ divf(R)
      v = ContravariantVector(r->[fr(r), fθ(r), fz(r)], f, true)
      @test div(v)(R) ≈ divf(R)
    end
  end

  @testset "curl" begin
    for ct ∈ (f, f⁻¹), _ in 1:numiters
      (R,X,fR) = getcoordsandtransforms(ct)

      v = CovariantVector(z->[z[3] - 2z[2] * z[1], 0, 0], ct, true)
      @test curl(v)(R, true) ≈ [0, 1, 2]

      a, b = rand(3), rand(2:5, 3, 3)
      fv(z) = [a[i] .* prod(z .^ b[:, i]) for i in 1:3]
      r = [2, 2pi, 2] .* rand(3) .- [0, pi, 1]
      j = ForwardDiff.jacobian(fv, r)
      jr = ForwardDiff.jacobian(x->fv(x) * x[1], r)
      for v ∈ (CovariantVector(fv, ct, true), ContravariantVector(fv, ct, true))
        @test curl(v)(r, true) ≈ [j[3, 2]/r[1] - j[2, 3],
                                  j[1, 3] - j[3, 1],
                                  (jr[2, 1] - j[1, 2]) / r[1]]
      end
    end
  end

  @testset "div curl, curl grad" begin
    for ct ∈ (f, f⁻¹), _ in 1:numiters
      (R,X,fR) = getcoordsandtransforms(ct)
      v = ContravariantVector(z->[prod(z), prod(z)^2, prod(z)^3], ct, true)
      @test div(curl(v))(R) ≈ 0 atol=1e-12
      @test ContravariantVector(CovariantVector(v))(R) ≈ v(R)
      @test CovariantVector(v)(R, true) ≈ v(R, true)
      v = CovariantVector(z->[prod(z), prod(z)^2, prod(z)^3], ct, true)
      @test div(curl(v))(R) ≈ 0 atol=1e-12
      @test CovariantVector(ContravariantVector(v))(R) ≈ v(R)
      @test ContravariantVector(v)(R, true) ≈ v(R, true)

      a, b = rand(3), rand(2:5, 3, 3)
      s(x) = mapreduce(i->a[i] * prod(x.^b[:, i]), +, 1:3)
      @test curl(grad(s, ct))(R) ≈ zeros(3) atol=1e-12
    end
  end

  @testset "endless operations" begin
    (R,X,fR) = getcoordsandtransforms(f)
    a, b = rand(3), rand(4:7, 3, 3)
    fv(z) = [a[i] .* prod(z .^ b[:, i]) for i in 1:3]
    for T1 ∈ (CovariantVector, ContravariantVector)
      for T2 ∈ (CovariantVector, ContravariantVector)
        A = T1(fv, f, true)
        B = T2(fv, f⁻¹, true)
        @test A(R, true) ≈ B(R, true)
        cA = curl(A)(R, true)
        ccA = curl(curl(A))(R, true)
        #cccA = curl(curl(curl(A)))(R, true)
        #ccccA = curl(curl(curl(curl(A))))(R, true)
        cB = curl(B)(R, true)
        ccB = curl(curl(B))(R, true)
        #cccB = curl(curl(curl(B)))(R, true)
        #ccccB = curl(curl(curl(curl(B))))(R, true)
        @test all(x->!iszero(x), cA)
        @test all(x->!iszero(x), ccA)
        #@test all(x->!iszero(x), cccA)
        #@test all(x->!iszero(x), ccccA)
        @test cA    ≈ cB
        @test ccA   ≈ ccB
        #@test cccA  ≈ cccB
        #@test ccccA ≈ ccccB
      end
    end
  end
end


export dvv_lstsq

"""
    dvv_lstsq(time_axis, dt, w)

Linear regression of phase shifts to time to give dt/t=-dv/v

# Arguments
- `time_axis::AbstractArray`: Array of lag times
- `dt::AbstractArray`: Time shifts for each lag time
- `w::AbstractArray`: Weights for the linear regression

# Returns
- `m::Float64`: dt/t for current correlation
- `em::Float64`: Error for calculation of `m`
- `a::Float64`: Intercept for regression calculation
- `ea::Float64`: Error on intercept
- `m0::Float64`: dt/t for current correlation with no intercept
- `em0::Float64`: Error for calculation of `m0`
"""
function dvv_lstsq(time_axis::AbstractArray, dt::AbstractArray; w::AbstractArray=ones(length(time_axis)))
    # regress data using least squares
    # force line through origin
    model0 = glm(@formula(Y ~0 + X),DataFrame(X=time_axis,Y=dt),Normal(),
                IdentityLink(),wts=w)
    # allow nonzero intercept
    model = glm(@formula(Y ~ X),DataFrame(X=time_axis,Y=dt),Normal(),
                IdentityLink(),wts=w)

    # parameters for the regression with nonzero intercept
    a,m = -coef(model).*100
    ea, em = stderror(model).*100
    # parameters for the regression through the origin
    m0 = -coef(model0)[1]*100
    em0 = stderror(model0)[1]*100

    return m, em, a, ea, m0, em0
end


#==========================================================================================#

# TYPE

mutable struct Poisson <: MLE

    sample::Microdata
    β::Vector{Float64}
    V::AbstractMatrix{Float64}

    Poisson() = new()
end

#==========================================================================================#

# CONSTRUCTOR

function Poisson(MD::Microdata)
    obj        = Poisson()
    obj.sample = MD
    return obj
end

#==========================================================================================#

# ESTIMATION

function _fit!(obj::Poisson, ::UnitWeights)

    has_offset = haskey(obj.sample.mapping, :offset)

    y  = getvector(obj, :response)
    x  = getmatrix(obj, :control)
    μ  = Array{Float64}(undef, length(y))
    r  = Array{Float64}(undef, length(y))
    xx = Array{Float64}(undef, size(x)...)

    β₀ = vcat(fill(0.0, size(x, 2) - 1), log(mean(y)))

    has_offset && (xo = getmatrix(obj, :offset, :control))

    function L(β::Vector)
        has_offset ? mul!(μ, xo, vcat(1.0, β)) : mul!(μ, x, β)
        return sum(exp.(μ) .- y .* μ)
    end

    function G!(g::Vector, β::Vector)
        has_offset ? mul!(μ, xo, vcat(1.0, β)) : mul!(μ, x, β)
        r .= exp.(μ) .- y
        mul!(g, transpose(x), r)
    end

    function LG!(g::Vector, β::Vector)

        has_offset ? mul!(μ, xo, vcat(1.0, β)) : mul!(μ, x, β)
        ll = 0.0

        @inbounds for (i, (yi, μi)) in enumerate(zip(y, μ))
            ηi   = exp(μi)
            r[i] = ηi - yi
            ll  += ηi - yi * μi
        end

        mul!(g, transpose(x), r)

        return ll
    end

    function H!(h::AbstractMatrix, β::Vector)
        has_offset ? mul!(μ, xo, vcat(1.0, β)) : mul!(μ, x, β)
        xx .= x .* exp.(μ)
        mul!(h, transpose(x), xx)
    end

    res = optimize(TwiceDifferentiable(L, G!, LG!, H!, β₀), β₀, Newton())

    if Optim.converged(res)
        obj.β = Optim.minimizer(res)
    else
        throw("likelihood maximization did not converge")
    end
end

function _fit!(obj::Poisson, w::AbstractWeights)

    has_offset = haskey(obj.sample.mapping, :offset)

    y  = getvector(obj, :response)
    x  = getmatrix(obj, :control)
    μ  = Array{Float64}(undef, length(y))
    r  = Array{Float64}(undef, length(y))
    xx = Array{Float64}(undef, size(x)...)

    β₀ = vcat(fill(0.0, size(x, 2) - 1), log(mean(y, w)))

    has_offset && (xo = getmatrix(obj, :offset, :control))

    function L(β::Vector)
        has_offset ? mul!(μ, xo, vcat(1.0, β)) : mul!(μ, x, β)
        return sum(exp.(μ) .- y .* μ, w)
    end

    function G!(g::Vector, β::Vector)
        has_offset ? mul!(μ, xo, vcat(1.0, β)) : mul!(μ, x, β)
        r .= w .* (exp.(μ) .- y)
        mul!(g, transpose(x), r)
    end

    function LG!(g::Vector, β::Vector)

        has_offset ? mul!(μ, xo, vcat(1.0, β)) : mul!(μ, x, β)
        ll = 0.0

        @inbounds for (i, (yi, μi, wi)) in enumerate(zip(y, μ, w))
            ηi   = exp(μi)
            r[i] = wi * (ηi - yi)
            ll  += wi * (ηi - yi * μi)
        end

        mul!(g, transpose(x), r)

        return ll
    end

    function H!(h::AbstractMatrix, β::Vector)
        has_offset ? mul!(μ, xo, vcat(1.0, β)) : mul!(μ, x, β)
        r  .= w .* exp.(μi)
        xx .= x .* r
        mul!(h, transpose(x), xx)
    end

    res = optimize(TwiceDifferentiable(L, G!, LG!, H!, β₀), β₀, Newton())

    if Optim.converged(res)
        obj.β = Optim.minimizer(res)
    else
        throw("likelihood maximization did not converge")
    end
end

#==========================================================================================#

# SCORE (DERIVATIVE OF THE LIKELIHOOD FUNCTION)

score(obj::Poisson) = Diagonal(residuals(obj)) * getmatrix(obj, :control)

# EXPECTED JACOBIAN OF SCORE × NUMBER OF OBSERVATIONS

function jacobian(obj::Poisson, ::UnitWeights)

    x = getmatrix(obj, :control)

    if haskey(obj.sample.mapping, :offset)
        xo = getmatrix(obj, :offset, :control)
        v  = xo * vcat(1.0, obj.β)
    else
        v  = x * obj.β
    end

    v .= exp.(v)

    return - crossprod(x, v)
end

function jacobian(obj::Poisson, w::AbstractWeights)

    x = getmatrix(obj, :control)

    if haskey(obj.sample.mapping, :offset)
        xo = getmatrix(obj, :offset, :control)
        v  = xo * vcat(1.0, obj.β)
    else
        v  = x * obj.β
    end

    v .= w .* exp.(v)

    return - crossprod(x, v)
end

#==========================================================================================#

# LINEAR PREDICTOR

function linear_predictor(obj::Poisson, MD::Microdata)

    (getnames(obj, :control) != getnames(MD, :control)) && throw("missing variables")

    if haskey(obj.sample.mapping, :offset)
        return getmatrix(MD, :offset, :control) * vcat(1.0, obj.β)
    else
        return getmatrix(MD, :control) * obj.β
    end
end

# FITTED VALUES

predict(obj::Poisson, MD::Microdata) = exp.(linear_predictor(obj, MD))

# DERIVATIVE OF FITTED VALUES

function jacobexp(obj::Poisson)

    if haskey(obj.sample.mapping, :offset)
        v = getmatrix(obj, :offset, :control) * vcat(1.0, obj.β)
    else
        v = getmatrix(obj, :control) * obj.β
    end

    v .= exp.(v)

    return Diag(v) * getmatrix(obj, :control)
end

#==========================================================================================#

# UTILITIES

coefnames(obj::Poisson) = getnames(obj, :control)

# LIKELIHOOD FUNCTION

function _loglikelihood(obj::Poisson, w::AbstractWeights)
    y = getvector(obj, :response)
    μ = linear_predictor(obj)
    return sum(y .* μ .- exp.(μ) .- loggamma.(1.0 .+ y), w)
end

# LIKELIHOOD FUNCTION UNDER NULL MODEL

function _nullloglikelihood(obj::Poisson, w::AbstractWeights)
    if haskey(obj.sample.mapping, :offset)
        y = getvector(obj, :response)
        o = getmatrix(obj, :offset)
        β = log(mean(y, w) / mean(exp.(o), w))
        return sum(y .* (o .+ β) .- exp.(o .+ β) .- loggamma.(1.0 .+ y), w)
    else
        y = getvector(obj, :response)
        η = mean(y, w)
        μ = log(η)
        return sum(y .* μ .- η .- loggamma.(1.0 .+ y), w)
    end
end

# DEVIANCE

function _deviance(obj::Poisson, w::AbstractWeights)
    y = getvector(obj, :response)
    μ = linear_predictor(obj)
    return 2 * sum(xlogx.(y) .- y .* (1.0 .+ μ) .+ exp.(μ), w)
end

# DEVIANCE UNDER NULL MODEL

function _nulldeviance(obj::Poisson, w::AbstractWeights)
    y = getvector(obj, :response)
    μ = mean(y, w)
    return 2 * sum(xlogx.(y) .- y .* (1.0 .+ log(μ)) .+ μ, w)
end


"""
    pdi{T<:Number}(x::Array{T, 1})

This function will range the values of each row, so that the strongest link has
a value of one. This works for deterministic and quantitative networks.

#### References

Poisot, T., Bever, J.D., Nemri, A., Thrall, P.H., Hochberg, M.E., 2011. A
conceptual framework for the evolution of ecological specialisation. Ecol. Lett.
14, 841–851. https://doi.org/10.1111/j.1461-0248.2011.01645.x

Poisot, T., Canard, E., Mouquet, N., Hochberg, M.E., 2012. A comparative study
of ecological specialization estimators. Methods in Ecology and Evolution 3,
537–544. https://doi.org/10.1111/j.2041-210X.2011.00174.x
"""
function pdi(x::Vector{T}) where {T<:Number}
    s = reverse(sort(x))
    s = s ./ maximum(s)
    p = s[1] .- s
    return sum(p[2:end])/(length(s)-1)
end

"""
    specificity(N::DeterministicNetwork)

Measure of specificity in a deterministic network. This returns a value between
0 and 1, where 1 indicates maximal specificity.

#### References

Poisot, T., Bever, J.D., Nemri, A., Thrall, P.H., Hochberg, M.E., 2011. A
conceptual framework for the evolution of ecological specialisation. Ecol. Lett.
14, 841–851. https://doi.org/10.1111/j.1461-0248.2011.01645.x

Poisot, T., Canard, E., Mouquet, N., Hochberg, M.E., 2012. A comparative study
of ecological specialization estimators. Methods in Ecology and Evolution 3,
537–544. https://doi.org/10.1111/j.2041-210X.2011.00174.x
"""
function specificity(N::BinaryNetwork)
    A = map(Int64, N.A)
    return Dict(zip(species(N; dims=1), vec(mapslices(pdi, A; dims=2))))
end

"""
    specificity(N::QuantitativeNetwork)

Measure of specificity in a quantitative network. This returns a value between 0
and 1, where 1 indicates maximal specificity. Note that the PDI is measured
species-wise, and the maximal interaction strength of every species is set to 1.

#### References

Poisot, T., Bever, J.D., Nemri, A., Thrall, P.H., Hochberg, M.E., 2011. A
conceptual framework for the evolution of ecological specialisation. Ecol. Lett.
14, 841–851. https://doi.org/10.1111/j.1461-0248.2011.01645.x

Poisot, T., Canard, E., Mouquet, N., Hochberg, M.E., 2012. A comparative study
of ecological specialization estimators. Methods in Ecology and Evolution 3,
537–544. https://doi.org/10.1111/j.2041-210X.2011.00174.x
"""
function specificity(N::QuantitativeNetwork)
    return Dict(zip(species(N; dims=1), vec(mapslices(pdi, N.A; dims=2))))
end


addprocs()
using MAT
include("include.jl")

TopN = 10

# Read Data
DATA = matread("yahoo.mat")
TrainSet = DATA["TrainSet"]
Holdout = DATA["Holdout"]
UW = DATA["SampledUnwatched"]
n, m = size(TrainSet)

# Read Item Model
W = readItemModel("example.model",m)

# Build RecWalk Model
P = RecWalk(TrainSet,W,0.005)

NDCG = SharedArray{Float64}(n)
RR = SharedArray{Float64}(n)
HR = SharedArray{Float64}(n)

#Base Item Model
PI = TrainSet*W
@sync @parallel for user = 1:n
	HR[user], RR[user], NDCG[user] = Single_HR_RR_NDCG(PI[user,:], vcat(Holdout[user], UW[:,user]), TopN)
end
println("Base Item Model:  HR = $(mean(HR))  ARHR=$(mean(RR))  NDCG=$(mean(NDCG))")

# RecWalk - K-Step
K = 7
@sync @parallel for user = 1:n
    ru  = sparse(reshape(P[user,:], 1, m+n))
    [ru  *= P for step=2:K]
    HR[user], RR[user], NDCG[user] = Single_HR_RR_NDCG(ru[n+1:end], vcat(Holdout[user], UW[:,user]), TopN)
end
println("RecWalk K-Step:   HR = $(mean(HR))  ARHR=$(mean(RR))  NDCG=$(mean(NDCG))")

# RecWalk - PR
eta = 0.7
PI = inv(full(I-eta*P)) # due to the small size of the example data the recwalk ppr vectors can be computed in batch.
PI = PI[1:n,n+1:end]
@sync @parallel for user = 1:n
	HR[user], RR[user], NDCG[user] = Single_HR_RR_NDCG(PI[user,:], vcat(Holdout[user], UW[:,user]), TopN)
end
println("RecWalk PR:       HR = $(mean(HR))  ARHR=$(mean(RR))  NDCG=$(mean(NDCG))")


const A = 1.0
const B = 1.5
const B1 = B + 1

@taylorize function brusselator!(du, u, p, t)
    x, y = u
    x² = x * x
    aux = x² * y
    du[1] = A + aux - B1*x
    du[2] = B*x - aux
    return du
end

U0(r) = Singleton([1.0, 1.0]) ⊕ BallInf(zeros(2), r)

bruss(r) = @ivp(u' = brusselator!(u), u(0) ∈ U0(r), dim: 2)

sol_1 = solve(bruss(0.1), T=30.0, alg=TMJets20(orderT=6, orderQ=2))


module MultivariateFunctions

using Dates
using SchumakerSpline
using GLM
using DataFrames: DataFrame
using DataStructures: OrderedDict
using Combinatorics: permutations
using LinearAlgebra: Symmetric, inv, det
using Optim

# This includes functions for use of MultivariateFunctions with dates.
include("date_conversions.jl")
export years_between, years_from_global_base, period_length
# The abstract type MultivariateFunction and all structs are implemented here.
# In addition we have operator reversals and some supporting functions.
include("0_structs_and_generic_reversals.jl")
export MultivariateFunction, PE_Function, Sum_Of_Functions, Piecewise_Function
export Sum_Of_Piecewise_Functions, PE_Unit
export change_base, trim_piecewise_function, sort, convert_to_linearly_rescale_inputs
export evaluate, rebadge, underlying_dimensions, convert, ≂
include("1_algebra.jl")
export +, -, *, /, ^
include("2_calculus.jl")
export derivative, all_derivatives, Hessian, jacobian, uniroot, find_local_optima
export integral
include("3_1D_splines_and_interpolation.jl")
export create_quadratic_spline, create_constant_interpolation_to_right
export create_constant_interpolation_to_left, create_linear_interpolation
include("chebyshevs.jl")
export get_chevyshevs_up_to
include("4_chebyshev_approximation.jl")
export create_chebyshev_approximation
include("5_ols_regression.jl")
export create_ols_approximation, create_saturated_ols_approximation
include("6_HighDimensionalApproximation.jl")
export create_recursive_partitioning, create_mars_spline, trim_mars_spline
end


Random.seed!(1)

T = BigFloat
setprecision(BigFloat, 512) do

  @testset "transform_Hess $n" for n in (3, 4, 5)
    AB = BigFloat.(randn(n, n)) + BigFloat.(randn(n,n))*im
    A = copy(AB)
    Q = Complex{BigFloat}.(I(n))
    ev = eig(Complex128.(AB))[1]
    s = ev[n:n]
    LinAlgBigFloat.transform_Hess!(A, 1, n, Q, s, n, n)
    @test checkfactors(AB, Q, A, rtol = eps(T) * n * n)
  end

  @testset "schurfact $n" for n in (1, 2, 3, 4, 5, 9, 20)
    AB = BigFloat.(randn(n, n))
    HF = hessfact(AB)
    @test checkfactors(AB, HF[:Q], HF[:H]) rtol = eps(T) * n * 2
    SF = schurfact(AB)
    @test checkfactors(AB, SF[:vectors], SF[:T]) rtol = eps(T) * n * 10

    AB = BigFloat.(randn(n, n)) + BigFloat.(randn(n,n)) * im
    HF = hessfact(AB)
    @test checkfactors(AB, HF[:Q], HF[:H]) rtol = eps(T) * n * 2
    SF = schurfact(AB)
    @test checkfactors(AB, SF[:vectors], SF[:T]) rtol = eps(T) * n * 20
  end
end


function draw_line(vent)
    x1, y1, x2, y2 = vent
    xs = x1 == x2 ? fill(x1, abs(y2 - y1) + 1) : x1:sign(x2 - x1):x2
    ys = y1 == y2 ? fill(y1, abs(x2 - x1) + 1) : y1:sign(y2 - y1):y2
    [zip(xs, ys)...]
end

function overlap_lines(lines)
    overlaps = Dict{Tuple{Int, Int}, Int}()
    points = reduce(vcat, lines)
    for point in points
        overlaps[point] = get(overlaps, point, 0) + 1
    end
    overlaps
end

function is_hv(vent)
    x1, y1, x2, y2 = vent
    x1 == x2 || y1 == y2
end

vents = [
    parse.(Int, split(replace(vent, " -> " => ','), ','))
    for vent in readlines("day5/vents.txt")
]

hv_lines = draw_line.(filter(is_hv, vents))
hv_overlaps = overlap_lines(hv_lines)

p1 = sum(values(hv_overlaps) .>= 2)

@show p1

lines = draw_line.(vents)
overlaps = overlap_lines(lines)

p2 = sum(values(overlaps) .>= 2)

@show p2


## Residual Constrained Alternating Minimization
# Function Definitions

function pAltMinChol{ET<:Number}(b::AbstractArray{ET}, r::Int, K::Int)
# Performs RCAM on a single matricised monochromatic slice
# Calls: minFrobQR

    # Generate L0
    m,n = size(b)
    L = collect(ET, randn(m,r))

    indn = (1:n)::UnitRange{Int64}
    local R

    # Reduce the data down to nonzero elements of each column
    ### change to do in one loop ###
    indnz = [find(x -> x != 0, b[:,i]) for i in indn]
    bsub = [b[indnz[i],i] for i in indn]

    # This needs to be cleaned up for non-square entries
    indnz_tran = [find(x -> x != 0, b[i,:]) for i in indn]
    bsub_tran = [b[i,indnz_tran[i]] for i in indn]

    # Distribute the reduced data
    bsubD = distribute(bsub)
    bsubD_tran = distribute(bsub_tran)

    # Begin Alterations
    for k = 1:K

        # Solve for R
        R = pminFrobChol(L,bsubD,indn,indnz)

        # Solve for L
        L = pminFrobChol(R,bsubD_tran,indn,indnz_tran)

    end

    return L, R
end
export pAltMinChol

function pminFrobChol{ET<:Number}(L::AbstractArray{ET}, bsubD::AbstractArray, indi::UnitRange{Int64},
                    indnz::AbstractArray)
# Solves for each row of input
# Calls: minFrobRsub

    # Init L2 and bsub
    L2 = [L[indnz[i],:] for i in indi]

    # Distribute L2 and bsub
    LD = distribute(L2)

    # Map solving function to distributed inputs
    Rcol = map(pminFrobRsubChol, LD, bsubD)

    # Reform R to matrix
    R = Array{ET}(size(L))
    for i = 1:length(Rcol)
        R[i,:] = Rcol[i]
    end

    R = conj(R)

    return R
end

function pminFrobRsubChol{ET<:Number}(L2::AbstractArray{ET},bsub::AbstractArray{ET})
# Solves a single row, meant to be mapped over a DArray

    bsub = transpose(L2)*bsub

    # Chol fact
    L2 = L2'*L2

    # Faster to do in w/o mem alloc, even though it does chol twice
    v = backsub(chol(L2),fwdsub((chol(L2))',bsub))

    return v

end

function backsub{ET<:Number}(A::AbstractArray{ET}, y::AbstractArray{ET})

    m = size(A,2)
    x = zeros(ET,m,1)

    x[m] = y[m]/A[m,m]

    for i = m-1:-1:1
        x[i] = (y[i]-dot(A[i,i+1:m], x[i+1:m]))/A[i,i]
    end

    return x
end

function fwdsub{ET<:Number}(A::AbstractArray{ET},y::AbstractArray{ET})

    m = size(A,2)
    x = zeros(ET,m,1)
    x[1] = y[1]/A[1,1]

    for i = 2:1:m
        x[i] = (y[i]-dot(A[i,i-1:-1:1], x[i-1:-1:1]))/A[i,i]
    end

    return x
end


#####################################
#InnerProduct/OuterProduct###########
#####################################

immutable OuterProduct{K<:Ket,B<:Bra} <: AbstractOperator
	ket::State{K}
	bra::State{B}
end

immutable InnerProduct <: AbstractScalar
	bra::State
	ket::State
	InnerProduct{B<:Bra,K<:Ket}(b::State{B}, k::State{K}) = new(b,k)
end

promote_rule{N<:Number, I<:InnerProduct}(::Type{I}, ::Type{N}) = ScalarExpr
promote_rule{N<:Number, I<:InnerProduct, T}(::Type{Array{I,T}}, ::Type{Array{N,T}}) = Array{ScalarExpr,T}
promote_rule{N<:Number, I<:InnerProduct, T}(::Type{SparseMatrixCSC{I,T}}, ::Type{SparseMatrixCSC{N,T}}) = SparseMatrixCSC{ScalarExpr,T}

#####################################
#Misc Functions######################
#####################################

=={O<:OuterProduct}(a::O, b::O) = a.ket==b.ket && a.bra==b.bra
isequal{O<:OuterProduct}(a::O, b::O) = isequal(a.ket,b.ket) && isequal(a.bra,b.bra)

=={I<:InnerProduct}(a::I, b::I) = a.ket==b.ket && a.bra==b.bra
isequal{I<:InnerProduct}(a::I, b::I) = isequal(a.ket,b.ket) && isequal(a.bra,b.bra)

bsym(o::OuterProduct) = [bsym(o.ket), bsym(o.bra)]
label(o::OuterProduct) = [label(o.ket), label(o.bra)]
bsym(i::InnerProduct) = [bsym(i.bra), bsym(i.ket)]
label(i::InnerProduct) = [label(i.bra), label(o.ket)]
conj(i::InnerProduct) = InnerProduct(i.ket', i.bra')
ctranspose(o::OuterProduct) = OuterProduct(o.bra', o.ket')
ctranspose(i::InnerProduct) = conj(i)

isdual(a::OuterProduct, b::OuterProduct) = isdual(a.ket, b.bra) && isdual(b.ket, a.bra)
isdual(a::InnerProduct, b::InnerProduct) = isdual(a.ket, b.bra) && isdual(b.ket, a.bra)

#####################################
#Show Functions######################
#####################################

show(io::IO, o::OuterProduct) = print(io, "$(repr(o.ket))$(repr(o.bra))");
show(io::IO, i::InnerProduct) = print(io, "$(repr(i.bra))$(repr(i.ket)[2:end])");

#####################################
#Arithmetic Operations###############
#####################################

*{K<:Ket}(o::OuterProduct, s::State{K}) = dvec([inner(o.bra,s)], basis(o.ket))
*{B<:Bra}(s::State{B}, o::OuterProduct) = dvec([inner(s,o.ket)], basis(o.bra))

kron{B<:Bra}(o::OuterProduct, s::State{B}) = OuterProduct(o.ket, tensor(o.bra,s))
kron{B<:Bra}(s::State{B}, o::OuterProduct) = OuterProduct(o.ket, tensor(s,o.bra))
kron{K<:Ket}(o::OuterProduct, s::State{K}) = OuterProduct(tensor(o.ket,s), o.bra)
kron{K<:Ket}(s::State{K}, o::OuterProduct) = OuterProduct(tensor(s,o.ket), o.bra)
kron(c::DiracCoeff, o::OuterProduct) = dmat([c], basis(o.ket), basis(o.bra))
kron(o::OuterProduct, c::DiracCoeff) = kron(c,o)

kron(a::OuterProduct, b::OuterProduct) = OuterProduct(tensor(a.ket, b.ket), tensor(a.bra, b.bra))

*(a::OuterProduct, b::OuterProduct) = dmat([a.bra*b.ket], basis(a.ket), basis(b.bra))

*(c::DiracCoeff, o::OuterProduct) = kron(c,o)
*(o::OuterProduct, c::DiracCoeff) = kron(c,o)
-(o::OuterProduct) = kron(-1,o)

function +{K<:Ket, B<:Bra}(a::OuterProduct{K,B}, b::OuterProduct{K,B})
	if a==b
		return dmat([2.0], basis(a.ket), basis(b.bra))
	else
		rowb = basis(a.ket, b.ket)
		colb = basis(a.bra, b.bra)
		res = dmat(zeros(2,2)[1:length(rowb), 1:length(colb)], rowb, colb)
		res[getpos(res, a)...] = 1.0
		res[getpos(res, b)...] = 1.0
		return res
	end
end

-(a::OuterProduct, b::OuterProduct) = a+(-b)