using RobustAndOptimalControl, ControlSystems, LinearAlgebra
# NOTE: Glad, Ljung chap 9 contains several numerical examples that can be used as test cases

approxin(el,col;kwargs...) = any(colel -> isapprox(el, colel; kwargs...), col)
approxsetequal(s1,s2;kwargs...) = all(approxin(p,s1;kwargs...) for p in s2) && all(approxin(p,s2;kwargs...) for p in s1)

w = exp10.(range(-5, stop=5, length=1000))
s = tf("s")
P = [1/(s+1) 2/(s+3); 1/(s+1) 1/(s+1)]
sys = ss(P)
sysmin = minreal(sys)
A,B,C,D = sys.A,sys.B,sys.C,sys.D
Am,Bm,Cm,Dm = sysmin.A,sysmin.B,sysmin.C,sysmin.D

@test approxsetequal(eigvals(Am), [-3,-1,-1])

Q1 = 100ones(4)
Q2 = 1I(2)
R1 = 100I(4)
R2 = 1I(2)
G = LQGProblem(sys, Q1, Q2, R1, R2)
gangoffourplot(G) # Test that it at least does not error
@test approxsetequal(eigvals(observer_controller(G).A), [ -31.6209+0.0im, -1.40629+0.0im, -15.9993+0.911174im, -15.9993-0.911174im, ], rtol = 1e-3)

# Test some utility functions
C = observer_controller(G)
@test C.A == sys.A - sys.B*lqr(G) - kalman(G)*sys.C
@test C.C == lqr(G)
@test C.B == kalman(G)
@test all(iszero, C.D)

Ce = extended_controller(G)
@test Ce.ny == C.ny
@test Ce.nu == C.nu
@test Ce.nw == size(Q1, 1)
@test system_mapping(Ce) == -C
Cp = performance_mapping(Ce)
@test Cp.ny == 0
@test Cp.A == C.A
@test all(iszero, Cp.B)
@test Cp.nu == Ce.nw

# testing with integral action
# qQ = 1
# qR = 1
# Q1 = 1000I(4)
# Q2 = 1I(2)
# R1 = 1I(6)
# R2 = 1I(2)
# N = I(6)
# Gi = LQGProblem(sys, Q1, Q2, R1, R2, qQ=qQ, qR=qR, integrator=true, ϵ=0.001, N=N)
# gangoffourplot(Gi) # Test that it at least does not error
# @test approxsetequal(eigvals(Gi.sysc.A), [-0.001, -0.001, -47.4832, -44.3442, -3.40255, -1.15355 ], rtol = 1e-2)

# @test approxsetequal(eigvals(closedloop(G).A), [-1.0, -14.1774, -2.21811, -14.3206, -1.60615, -22.526, -1.0, -14.1774], rtol=1e-3)
# @test approxsetequal(eigvals(G.T.A), [-22.526, -2.21811, -1.60615, -14.3206, -14.1774, -1.0, -1.0, -14.1774], rtol=1e-3)
# @test approxsetequal(eigvals(G.S.A), [-22.526, -2.21811, -1.60615, -14.3206, -14.1774, -1.0, -1.0, -14.1774], rtol=1e-3)
# @test approxsetequal(eigvals(G.CS.A), [-31.6209+0.0im, -1.40629+0.0im, -15.9993+0.911174im, -15.9993-0.911174im, -22.526+0.0im, -2.21811+0.0im, -1.60615+0.0im, -14.3206+0.0im, -14.1774+0.0im, -1.0+0.0im, -1.0+0.0im, -14.1774+0.0im], rtol=1e-3)
# @test approxsetequal(eigvals(G.PS.A), [-1.0, -1.0, -3.0, -1.0, -22.526, -2.21811, -1.60615, -14.3206, -14.1774, -1.0, -1.0, -14.1774], rtol=1e-3)

# @test approxsetequal(eigvals(Gi.cl.A), [-1.0, -44.7425, -44.8455, -2.23294, -4.28574, -2.06662, -0.109432, -1.31779, -0.78293, -1.0, -0.001, -0.001], rtol=1e-3)
# @test approxsetequal(eigvals(Gi.T.A), [-44.7425, -44.8455, -4.28574, -0.109432, -2.23294, -2.06662, -1.31779, -0.78293, -1.0, -1.0], rtol=1e-3)
# @test approxsetequal(eigvals(Gi.S.A), [-44.7425, -44.8455, -4.28574, -0.109432, -2.23294, -2.06662, -1.31779, -0.78293, -1.0, -1.0], rtol=1e-3)

# @test eigvals(Gi.sysc.A) == eigvals(Gi.Fr.A)

## Test cases from Glad Ljung chap 9

# Aircrapt control

A = [-0.292 8.13 -201 9.77 0 -12.5 17.1
    -0.152  -2.54  0.561  -0.0004  0  107  7.68
    0.0364  -0.0678  -0.481  0.0012  0  4.67  -7.98
    0 1 0.0401 0 0 0 0
    0 0 1 0 0 0 0
    0 0 0 0 0 -20 0
    0 0 0 0 0 0 -20
]

B = [
    0 -2.15
    -31.7 0.0274
    0 1.48
    0 0
    0 0
    20 0
    0 20
]
# Example 9.1
M = [0 0 0 1 0 0 0
     0 0 0 0 1 0 0]
C = M

N = I(7)
nx,ny,nu = size(A,1), size(C, 1), size(B,2)

Q1 = I(ny)
Q2 = I(nu)
R1 = I(nx)
R2 = I(ny)

sys = ss(A,N,B,M,C)
G = LQGProblem(sys, Q1, Q2, R1, R2)
@test lqr(G) ≈ [-0.0022 0.17 0.12 0.98 0.31 0.76 0.018
            -0.0038 0.028  -0.25 0.16  -0.95 0.072 0.10] rtol=0.01

Q2 = 0.1*I(nu)
G = LQGProblem(sys, Q1, Q2, R1, R2)
@test lqr(G) ≈ [-0.0036 0.39 0.21 3.11 0.63 1.54 0.046
            -0.0073 0.085 -0.78 0.57 -3.10 0.20 0.30] rtol=0.01

@test dcgain(RobustAndOptimalControl.closedloop(G)) ≈ I

# Room temperature control
K1 = 1
K2 = 1
K3 = 1/8
K4 = 0.2
A = [
    -K3-K4 K4 0 0 0
    K1 -K1-K2 0 K2 0
    0 0 -0.01 -0.068 0
    0 0 1 0 0
    0 0 0 0 -0.0001
]
# @warn "I added the last -0.0001"

B2 = [
    K3
    0
    0
    0
    0
]

C1 = [0 1 0 0 0]
C2 = [
    0 1 0 0 1
    0 1 0 0 0
    0 0 0 1 0
]
B1 = [
    0 0
    0 0
    1 0
    0 0
    0 1
]
sys = ss(A,B1,B2,C1,C2)

Q1 = 1.0I(1)
Q2 = 0.01I(1)
R1 = diagm([10.0, 1])
# R1 = diagm([0, 0, 10.0, 1, 0])
R2 = diagm([0.0001, 0.1, 0.001])
G = LQGProblem(sys, Q1, Q2, R1, R2, qQ = 0.000001)

# NOTE
# the kalman filter qutions will not work out until the noise model is incorporated

@test kalman(G) ≈ [-0.0039  0.0003  0.0035
            0.0463  0.0027  0.9676
            8.5224  0.0148  99.7543
            1.1426  0.0097  14.1198
            99.9634  -0.0026 -0.8533] rtol=1e-2
@test lqr(G) ≈ [2.8106 1.4071 6.8099 3.7874 0] rtol=1e-2

@test dcgain(RobustAndOptimalControl.closedloop(G))[] ≈ 1

## example 9.3

A = diagm([-1.0, -1, -2])
B = [
    1.0 0
    0.0 1
    0.0 1
]

C = [
    2.0 0 3
    1.0 1 0
]
M = C
sys = ss(A,I(3), B, M, C)

Q1 = I(2)
Q2 = 0.1I(2)
R1 = I(2)
R2 = 0.001I(2)

# test with integrator
# G = LQGProblem(sys, Q1, Q2, R1, R2, integrator=true, M = M, measurement=true)

# Lr = pinv(M * ((B * lqr(G)[:, 1:3] - A) \ B))

# @test Lr ≈ [
#     4.24 -3.43
#     0.12 4.82
# ] rtol = 0.01

# # NOTE: Glad ljung does not have identity in the end here, but the eigenvalues of A-BL are identical for this L and their L
# @test lqr(G) ≈ [
#     4.05 0.81 4.23 1 0
#     5.05 1.01 5.96 0 1
# ] rtol = 0.01

# @test kalman(G) ≈ [
#     0 0
#     0 0
#     0 0
#     31.62 0
#     0 31.62
# ] rtol = 0.01

# @test dcgain(closedloop(G)) ≈ I

# Random system, compare with ML output

A = [  0.154388  1.42524    0.914886  -0.313071
1.36189   0.30188   -0.767355   0.313058
-0.485766  0.926861  -0.755906  -0.309824
2.10783   0.146366  -1.65582    0.115643]

B = [
  0.625203  -1.87059
 -0.399363   0.516783
 -0.150984   0.161573
 -0.432394  -0.162178
]

C = [
  0.99517   -0.26967    0.255771  -0.0356366
 -0.358348   0.114812  -0.417828   0.983806
  0.45019   -0.166284   1.47825    2.21589
]

K = [ 1.22237  -0.114471  -1.10722
1.00941   0.445568  -1.91614]

P = ss(A,B,C,0)
C = ss(K)

SiA = [
 -0.3534     1.707    -3.451    -5.871
   1.674    0.1881    0.1495     1.257
 -0.3683    0.8871   -0.5026  -0.08282
   2.433   0.09358    -2.655    -1.636
]

SiB = [
  0.6252   -1.871
 -0.3994   0.5168
  -0.151   0.1616
 -0.4324  -0.1622
]

SiC = [
   -0.759    0.1587     1.276      2.61
  0.01776  -0.09757     2.761     3.844
]

SiD = [
    1      0
    0      1]

Si = ss(SiA, SiB, SiC, SiD)

@test isapprox(input_sensitivity(P, C), Si, rtol=1e-3)

SoA = [
    -0.3534     1.707    -3.451    -5.871
      1.674    0.1881    0.1495     1.257
    -0.3683    0.8871   -0.5026  -0.08282
      2.433   0.09358    -2.655    -1.636
]

SoB = [
    1.124     0.905    -2.892
 -0.03348    -0.276     0.548
  0.02146  -0.08928    0.1424
   0.6922   0.02276   -0.7895
]

SoC = [
   0.9952   -0.2697    0.2558  -0.03564
  -0.3583    0.1148   -0.4178    0.9838
   0.4502   -0.1663     1.478     2.216
]

SoD = I(3)

So = ss(SoA, SoB, SoC, SoD)

@test isapprox(output_sensitivity(P, C), So, rtol=1e-3)

ATo  = [
      -0.3534     1.707    -3.451    -5.871
        1.674    0.1881    0.1495     1.257
      -0.3683    0.8871   -0.5026  -0.08282
        2.433   0.09358    -2.655    -1.636
]

BTo = [
    1.124     0.905    -2.892
    -0.03348    -0.276     0.548
    0.02146  -0.08928    0.1424
    0.6922   0.02276   -0.7895
    ]

CTo = [
    -0.9952   0.2697  -0.2558  0.03564
    0.3583  -0.1148   0.4178  -0.9838
    -0.4502   0.1663   -1.478   -2.216
]
To = ss(ATo, BTo, CTo, 0)

@test isapprox(output_comp_sensitivity(P, C), To, rtol=1e-3)

ATi = [
      -0.3534     1.707    -3.451    -5.871
        1.674    0.1881    0.1495     1.257
      -0.3683    0.8871   -0.5026  -0.08282
        2.433   0.09358    -2.655    -1.636
]

BTi = [
    0.6252   -1.871
    -0.3994   0.5168
    -0.151   0.1616
    -0.4324  -0.1622
    ]

CTi = [
0.759   -0.1587    -1.276     -2.61
    -0.01776   0.09757    -2.761    -3.844
]

Ti = ss(ATi, BTi, CTi, 0)

@test isapprox(input_comp_sensitivity(P, C), Ti, rtol=1e-3)


###
# Helper Tests
###
include("../src/helpers.jl")

facts("bootstrapping tests") do
	#single variate data
	srand(8765309)
	data = randn(1000)
	@fact boot(data, mean, .9, 10000) --> roughly(0.108, 1e-3)

	#multivariate data
	srand(8765309)
	data = randn(1000, 3)
	f= x->mean(minimum(x, 1)) # a nonsense function
	@fact boot(data, f, .9, 10000) --> roughly(-2.8777884239388682, 1e-8)
end

facts("t-approximations") do
	srand(8765309)
	data = randn(1000, 3)
	out = calcMeansT(data, .1)
	@fact out[1] --> roughly(0.004451891728022329, 1e-8)
	@fact out[2] --> roughly(0.06592683689569365, 1e-8)

	out = calcMeansT(data, .1, joint=false)
	@fact out[1] --> roughly(0.011243060792593858, 1e-8)
	@fact out[2] --> roughly(0.05913566783112214, 1e-8)
end

facts("fwdBackSigsTest") do
	srand(8675309); data = randn(100)
	sigfwd, sigback = calcSigsBoot(data, .1, 10000)
	@fact sigfwd  --> roughly(1.09, 1e-2)
	@fact sigback --> roughly(1.09, 1e-2)
end

facts("KSGammaTest") do
	@fact KSGamma(.1, 1000) --> roughly(0.0385517413380297)
end

facts("boot_mu_test") do
	srand(8675309); data = randn(100)
	@fact boot_mu(data, .1, 100) --> roughly(0.15,1e-2)
end

facts("boot_sigma_test") do
	srand(8675309); data = randn(100)
	@fact boot_sigma(data, .1, 100) --> roughly(0.20,1e-2)
end

facts("ab_thresh_test") do
	srand(8675309); data = randn(100, 3)
	@fact calc_ab_thresh(data, .2, 100, 100) --> roughly(0.177,1e-3)
end


import Base.transpose, Base.ctranspose
export transpose, ctranspose

transpose(x::Constant) = Constant(transpose(x.value))
ctranspose(x::Constant) = transpose(x)

function transpose(x::AffineExpr)
  vec_sz = x.size[1] * x.size[2]
  num_rows = x.size[1]
  num_cols = x.size[2]

  I = Array{Int64}(vec_sz)
  J = Array{Int64}(vec_sz)

  k = 1
  for r = 1:num_rows
    for c = 1:num_cols
      J[k] = (c - 1) * num_rows + r
      I[k] = (r - 1) * num_cols + c
      k += 1
    end
  end

  perm_constant = Constant(sparse(I, J, 1.0))

  vars_to_coeffs_map = Dict{UInt64, Constant}()
  for (v, c) in x.vars_to_coeffs_map
    vars_to_coeffs_map[v] = perm_constant * c
  end
  constant = perm_constant * x.constant

  this = AffineExpr(:transpose, (x,), vars_to_coeffs_map, constant, (num_cols, num_rows))
  this.evaluate = ()->x.evaluate()'
  return this
end

ctranspose(x::AffineExpr) = transpose(x::AffineExpr)


"""
    Triweight(μ, σ)
"""
struct Triweight{T<:Real} <: ContinuousUnivariateDistribution
    μ::T
    σ::T
    Triweight{T}(µ::T, σ::T) where {T} = new{T}(µ, σ)
end

function Triweight(μ::T, σ::T; check_args::Bool=true) where {T <: Real}
    @check_args Triweight (σ, σ > zero(σ))
    return Triweight{T}(μ, σ)
end

Triweight(μ::Real, σ::Real; check_args::Bool=true) = Triweight(promote(μ, σ)...; check_args=check_args)
Triweight(μ::Integer, σ::Integer; check_args::Bool=true) = Triweight(float(μ), float(σ); check_args=check_args)
Triweight(μ::Real=0.0) = Triweight(μ, one(μ); check_args=false)

@distr_support Triweight d.μ - d.σ d.μ + d.σ

## Conversions

convert(::Type{Triweight{T}}, μ::Real, σ::Real) where {T<:Real} = Triweight(T(μ), T(σ))
Base.convert(::Type{Triweight{T}}, d::Triweight) where {T<:Real} = Triweight{T}(T(d.μ), T(d.σ))
Base.convert(::Type{Triweight{T}}, d::Triweight{T}) where {T<:Real} = d

## Parameters

location(d::Triweight) = d.μ
scale(d::Triweight) = d.σ
params(d::Triweight) = (d.μ, d.σ)
@inline partype(d::Triweight{T}) where {T<:Real} = T

## Properties
mean(d::Triweight) = d.μ
median(d::Triweight) = d.μ
mode(d::Triweight) = d.μ

var(d::Triweight) = d.σ^2 / 9
skewness(d::Triweight{T}) where {T<:Real} = zero(T)
kurtosis(d::Triweight{T}) where {T<:Real} = T(1)/33 - 3

## Functions
function pdf(d::Triweight{T}, x::Real) where T<:Real
    u = abs(x - d.μ)/d.σ
    u >= 1 ? zero(T) : 1.09375*(1 - u*u)^3/d.σ
end
logpdf(d::Triweight, x::Real) = log(pdf(d, x))

function cdf(d::Triweight{T}, x::Real) where T<:Real
    u = (x - d.μ)/d.σ
    u ≤ -1 ? zero(T) : u ≥ 1 ? one(T) : 0.03125*(1 + u)^4*@horner(u,16,-29,20,-5)
end

function ccdf(d::Triweight{T}, x::Real) where T<:Real
    u = (d.μ - x)/d.σ
    u ≤ -1 ? zero(T) : u ≥ 1 ? one(T) : 0.03125*(1 + u)^4*@horner(u,16,-29,20,-5)
end

@quantile_newton Triweight

function mgf(d::Triweight{T}, t::Float64) where T<:Real
    a = d.σ*t
    a2 = a*a
    a == 0 ? one(T) : 105*exp(d.μ*t)*((15/a2+1)*cosh(a)-(15/a2-6)/a*sinh(a))/(a2*a2)
end

function cf(d::Triweight{T}, t::Float64) where T<:Real
    a = d.σ*t
    a2 = a*a
    a == 0 ? complex(one(T)) : 105*cis(d.μ*t)*((1-15/a2)*cos(a)+(15/a2-6)/a*sin(a))/(a2*a2)
end


using ChaosTools
using Test, StaticArrays

println("\nTesting orbit diagrams...")

@testset "Discrete Orbit Diagram" begin
    @testset "1D" begin
        ds = Systems.logistic(0.5112661)
        i = 1
        parameter = 1
        pvalues = 2:0.01:4
        n = 50
        Ttr = 5000
        output = orbitdiagram(ds, i, parameter, pvalues; n = n, Ttr = Ttr)
        @test output[1][1] ≈ 0.5
        @test output[2][1] ≈ 0.502487562
        @test output[end][1] != output[end][2] != output[end][3]

        @testset "specific range" begin
        ulims = (0.0, 0.5)
        ics = zeros(length(pvalues)); ics[end] = 0.5112661
        output = orbitdiagram(
            ds, i, parameter, pvalues;
            n = n, Ttr = Ttr, u0 = ics, ulims = ulims
        )
        @test all(iszero, output[1])
        @test !all(iszero, output[end])
        @test all(≤(0.5), output[end])
        end
    end

    @testset "2D" begin
        ds = Systems.standardmap()
        i = 2
        parameter = 1
        pvalues = 0:0.05:2
        ics = [0.001rand(2) for m in pvalues]
        n = 50
        Ttr = 5000
        output = orbitdiagram(ds, i, parameter, pvalues; n = n, Ttr = Ttr)
        @test length(output[1]) == n
        output = orbitdiagram(ds, i, parameter, pvalues; n = n, Ttr = Ttr, u0=ics)
        @test length(output[1]) == n
    end

    @testset "Vector values collection" begin
        ds = Systems.coupledstandardmaps(4)
        i = SVector(2, 3)
        parameter = 1
        pvalues = 0:0.05:2
        ics = [0.001rand(8) for m in pvalues]
        n = 50
        Ttr = 500
        output = orbitdiagram(ds, i, parameter, pvalues; n = n, Ttr = Ttr)
        @test length(output) == length(pvalues)
        @test length(output[1]) == n
        @test length(output[1][1]) == 2
    end

    @testset "IIP" begin
        ds = Systems.henon_iip()
        i = 2
        parameter = 1
        pvalues = [1.4, 1.5]
        ics = [rand(2), rand(2)]
        n = 50
        Ttr = 5000
        output = orbitdiagram(ds, i, parameter, pvalues; n = n, Ttr = Ttr)
        @test length(output[1]) == n
    end
end


# Min cost max flow

# This solves the min-cost max flow problem using an LP.
# There are far more efficient algorithms out there (e.g. network
# simplex) but most implementations are only for floats.
#
# Args:
#   X - n x p matrix of points in R^d
#   E - the edge set for the rows of X
#   q1 - the pmf for the source distribution
#   q2 - the pmf for the target distribution
#   solver - the LP solver used to compute the flow.
# Returns
#   mincost - the minimum cost of the max flow
#   flow - the flows to achieve the min cost
#   status - the status of the LP solver.

function mincostflow(X::AbstractArray{T,2},
                     E::Array{Tuple{Int,Int},1},
                     q1::AbstractArray{T,1},
                     q2::AbstractArray{T,1},
                     solver::AbstractMathProgSolver) where {T<:Number}
    # number of points
    (n, d) = size(X)
    # size of supports
    numedges = length(E)
    # construct distances
    dists = zeros(numedges)
    for ii = 1:numedges
        dists[ii] = norm(X[E[ii][1],:] - X[E[ii][2],:], 1)
    end
    # setup model
    m = Model(solver=solver)
    # flow on edges, twice the number for both directions
    @variable(m, f[i=1:(2*numedges)] >= 0)
    # construct objective
    costobjective = f[1:(2*numedges)]'[dists; dists]
    @objective(m, Min, costobjective)
    # first we index vertices by incident edges
    inedges_for_node = Dict{Int,Array{Int,1}}()
    outedges_for_node = Dict{Int,Array{Int,1}}()
    for (edgeid, edge) in enumerate(E)
        (from, to) = edge
        _addedgetonodeindex!(inedges_for_node, outedges_for_node,
                             edgeid, from, to)
        _addedgetonodeindex!(inedges_for_node, outedges_for_node,
                             edgeid + numedges, to, from)
    end
    # impose flow constraints for every pair
    for ii=1:n
        in_edgeids = inedges_for_node[ii]
        out_edgeids = outedges_for_node[ii]
        @constraint(m, sum(f[in_edgeids]) - sum(f[out_edgeids]) == q2[ii] - q1[ii])
    end
    # Solve the problem
    @time status = JuMP.solve(m)
    mincost = getobjectivevalue(m)
    directedflows = getvalue(f)[1:(2*numedges)]
    flows = directedflows[1:numedges] - directedflows[numedges+(1:numedges)]
    # return the values
    return (mincost, flows, status)
end

function _addedgetonodeindex!(
    inedges_for_node::Dict{Int,Array{Int,1}},
    outedges_for_node::Dict{Int,Array{Int,1}},
    edgeid::Int,
    from::Int,
    to::Int)
    # first do the normal orientation
    if !haskey(inedges_for_node, to)
        inedges_for_node[to] = Int[]
    end
    if !haskey(outedges_for_node, from)
        outedges_for_node[from] = Int[]
    end
    push!(inedges_for_node[to], edgeid)
    push!(outedges_for_node[from], edgeid)
end


using TimedAutomata
using Test, LightGraphs, GraphPlot, Compose, IntervalArithmetic

# Clock related test

@testset "All Tests" begin
    op1 = PartialOrder(3, (1,2)=>:<, (2,3)=>:(<=))
    op2 = PartialOrder(3, (2,1)=>:(>=), (1,3)=>:<)
    op3 = PartialOrder(3, (2,3)=>:<)
    ot1 = TotalOrder([1,2,3])
    ot2 = TotalOrder([1,1,2])
    ot3 = TotalOrder([1,2,2])

    t1 = Clk"0"
    t2 = Clk"1"
    t3 = Clk"3"
    t4 = Clk"(0,1)"
    t5 = Clk"[0,2)"
    t6 = Clk"(0,∞)"
    t7 = Clk"(∞,1]"
    t8 = Clk"[0,∞)"
    t9 = Clk""

    c1 = ClockRegion{3,Int}(1=>Clk"(0,1)", 2=>Clk"[0,1]", 3=>Clk"1"; order=op1)
    c2 = ClockRegion{3,Int}(1=>Clk"(0,1)", 2=>Clk"[0,1]", 3=>Clk"1"; order=op2)
    c3 = ClockRegion{3,Int}(1=>Clk"(0,1)", 2=>Clk"[0,1]"; order=op2)
    c4 = ClockRegion{3,Int}(1=>Clk"(0,1)", 2=>Clk"[0,1)"; order=op2)
    c5 = ClockRegion{3,Int}(1=>Clk"(0,1)", 2=>Clk"1"; order=op2)

    @testset "Equality tests" begin
        @test Clk"1" == Clk"1"
        @test Clk"1" != Clk"0"
        @test Clk"(0,1]" == Clk"(0,1]"
        @test Clk"(0,1]" != Clk"(0,1)"

        @test ClockRegion{1,Int64}(TimeSlice{Int64}[TimePoint{Int64}(0)], TotalOrder{1}([0])) == ClockRegion{1,Int64}(TimeSlice{Int64}[TimePoint{Int64}(0)], TotalOrder{1}([0]))
    end

    @testset "Order Relation tests" begin
        expected_results = [
            true  true  false false false false;
            false true  false false false false;
            false false true  false false false;
            true  true  true  true  false false;
            false true  true  false true  false;
            true  true  false false false true;
        ]
        results = zeros(Bool, size(expected_results))

        for (i,r1) ∈ enumerate((op1,op2,op3,ot1,ot2,ot3))
            for (j,r2) ∈ enumerate((op1,op2,op3,ot1,ot2,ot3))
                results[i,j] = (r1 ⊆ r2)
                @test results[i,j] == expected_results[i,j]
            end
        end

        failures = findall(results .!= expected_results)

        if ~isempty(failures)
            println("Failed for: $(failures)")
        end
    end

    @testset "Time slice tests" begin
        expected_results = [
            true  false false false true  false true  true  true;
            false true  false false true  true  true  true  true;
            false false true  false false true  false true  true;
            false false false true  true  true  true  true  true;
            false false false false true  false false true  true;
            false false false false false true  false true  true;
            false false false false false false true  false true;
            false false false false false false false true  true;
            false false false false false false false false true;
        ]
        results = zeros(Bool, size(expected_results))

        for (i,r1) ∈ enumerate((t1,t2,t3,t4,t5,t6,t7,t8,t9))
            for (j,r2) ∈ enumerate((t1,t2,t3,t4,t5,t6,t7,t8,t9))
                results[i,j] = (r1 ⊆ r2)
                @test results[i,j] == expected_results[i,j]
            end
        end

        failures = findall(results .!= expected_results)

        if ~isempty(failures)
            println("Failed for: $(failures)")
        end
    end

    @testset "Clock Region tests" begin
        expected_results = [
            true  true  true false false;
            false true  true false false;
            false false true false false;
            false false true true  false;
            false false true false true
        ]
        results = zeros(Bool, size(expected_results))

        for (i,r1) ∈ enumerate((c1,c2,c3,c4,c5))
            for (j,r2) ∈ enumerate((c1,c2,c3,c4,c5))
                results[i,j] = (r1 ⊆ r2)
                @test results[i,j] == expected_results[i,j]
            end
        end

        failures = findall(results .!= expected_results)

        if ~isempty(failures)
            println("Failed for: $(failures)")
        end
    end

end

#=
tma = TMA(
    [
        [0,0],
        [0,1],
        [1,0],
        [1,1]
    ],
    Vector{Int}[],
    2,
    Dict(
        ([0,0],:a)=>[TMATransition([0,1],[1],@condition true)],
        ([0,1],:b)=>[TMATransition([1,0],[2],@condition true)],
        ([1,0],:c)=>[TMATransition([1,1],Int[],@condition c[1] < 1.0), TMATransition([0,0],Int[],@condition c[1] == 1.0)],
        ([1,1],:d)=>[TMATransition([0,0],Int[],@condition c[2] > 2.0)],
    ),
    Dict([1,0]=>@condition c[1] < 1.0),
    [0,0]
)
@test collect_constants(tma) == [Set((1//1,)),Set((2//1,))]
=#

ma = MA(
    [
        0,
        1,
        2
    ],
    [0],
    Dict(
        0=>[MATransition(1,:a),MATransition(2,:b)],
        1=>[MATransition(2,:a),MATransition(0,:b)],
        2=>[MATransition(0,:a),MATransition(1,:b)],
    ),
    0
)

L = language(ma)
println("Language of simple, timeless automaton:")
println(repr(L))

tma = TMA(
    [
        0,
        1,
        2,
        3
    ],
    [3],
    1,
    Dict(
        0=>[TMATransition(1,:msg,[1],ClockRegion{1,Int}())],
        1=>[TMATransition(1,:msg,[1],ClockRegion{1,Int}(1=>Clk"1")),TMATransition(1,:msg,[1],ClockRegion{1,Int}(1=>Clk"[1,∞)")),TMATransition(2,:msg,[1],ClockRegion{1,Int}(1=>Clk"(0,1)"))],
        2=>[TMATransition(3,:alarm,Int[],ClockRegion{1,Int}())]
    ),
    #Dict{Int,ClockRegion{1,Int}}(),
    Dict(2=>ClockRegion{1,Int}(1=>Clk"(0,1)")),
    0
)

@test collect_constants(tma) == [Set([0,1])]

ma = untime(tma)
draw(SVG("test_untimed.svg", 16cm, 16cm), gplot(ma))

L = language(ma)
println("Language of untimed automaton")
println(repr(L))

draw(SVG("test_timed.svg", 16cm, 16cm), gplot(tma))

#=
for s ∈ ma.Σ
    println("State $(s):")
    for ((s′,a),ee) ∈ pairs(ma.E)
        if s′==s
            for e ∈ ee
                println("\t$(a)=>$(e.s)")
            end
        end
    end
end
 =#


"""
    calc_resid(X, Y, Z, B)

Calculate residuals

# Arguments

- X = 2d array of floats consisting of the row covariates, standardized as
  necessary
- Y = 2d array of floats consisting of the multivariate response
  observations, standardized as necessary
- Z = 2d array of floats consisting of the column covariates, standardized
  as necessary
- B = 2d array of floats consisting of coefficient estimates

# Value

2d array of floats

"""
function calc_resid(X::AbstractArray{Float64,2}, Y::AbstractArray{Float64,2},
                    Z::AbstractArray{Float64,2}, B::AbstractArray{Float64,2})

    # Obtain fitted values
    resid = calc_preds(X, Z, B)

    # Compute residuals over fitted values
    resid .= Y .- resid

    return resid
end

"""
    calc_resid!(resid, X, Y, Z, B)

Calculate residuals in place

# Arguments

- resid = 2d array of floats consisting of the residuals, to be updated in
  place
- X = 2d array of floats consisting of the row covariates, standardized as
  necessary
- Y = 2d array of floats consisting of the multivariate response
  observations, standardized as necessary
- Z = 2d array of floats consisting of the column covariates, standardized
  as necessary
- B = 2d array of floats consisting of coefficient estimates

# Value

None; updates residuals in place.

"""
function calc_resid!(resid::AbstractArray{Float64,2},
                     X::AbstractArray{Float64,2},
                     Y::AbstractArray{Float64,2},
                     Z::AbstractArray{Float64,2},
                     B::AbstractArray{Float64,2})

    # Obtain fitted values in place
    calc_preds!(resid, X, Z, B)

    # Compute residuals in place over fitted values
    resid .= Y .- resid
end


module Trajectories

using ..UpdateMethods

import Base:push!, length

export Trajectory, heun_fill, heun_fill!
mutable struct Trajectory
    t::Vector{<:Real}
    x::Vector{<:Real}
    v::Union{Vector{<:Real}, Nothing}
end
Trajectory(t::Vector{<:Real}, x::Vector{<:Real}) = Trajectory(t, x, nothing)
Trajectory(T::Type) = Trajectory(Vector{T}(), Vector{T}(), nothing)
Trajectory() = Trajectory(Float64)

function push!(tj::Trajectory, t::Real, x::Real)
    push!(tj.t, t)
    push!(tj.x, x)
    tj
end

function genrate_v!(tj::Trajectory)
    tj.v = (tj.x[3:end] - tj.x[1:end-2]) ./ (tj.t[3:end] - tj.t[1:end-2])
end

lastx(tj::Trajectory)::Real = tj.x[end]

function length(tj::Trajectory)
    length(tj.x)
end

function heun_fill!(tj::Trajectory, h::Heun, f::Function, trange::Tuple{T, T})::Trajectory where T<:Real
    t0, tend = trange
    for t in t0:h.step:tend
        push!(tj, heun_update(f, h, lastx(tj), t)...)
    end
    tj
end

function heun_fill(x0::Real, h::Heun, f::Function, trange::Tuple{T, T})::Trajectory where T<:Real
    heun_fill!(Trajectory([trange[1]],[x0]), h, f, trange)
end

end


H = Hilbert(10)
@test Matrix(H)[10,10] == 1//(10+10-1)


import CPLEX
import Plots
import PolyhedralRelaxations
using JuMP

const PR = PolyhedralRelaxations

opt = CPLEX.Optimizer

# This line changes the Plots backend to GR, which speeds up plotting.
# An alternative is `Plots.plotly()`, which uses Python's plotly library as the backend.
Plots.gr()

function plot_mip(f::Function, base_partition::Vector{<:Real}; f_dash::Union{Function,Nothing}=nothing, out_file="plots/mip.pdf")
    if isa(f_dash, Nothing)
        milp_data, function_data = PR.construct_milp_relaxation(f, base_partition)
    else
        milp_data, function_data = PR.construct_milp_relaxation(f, f_dash, base_partition)
    end
    lb, ub = PR.get_variable_bounds(milp_data)
    x_min, x_max = lb[milp_data.x_index], ub[milp_data.x_index]
    @info "x bounds: $x_min, $x_max"

    # Add variables
    milp = Model(opt)
    set_optimizer_attribute(milp, "CPXPARAM_ScreenOutput", 0)
    num_vars = PR.get_num_variables(milp_data)
    binary_indices = PR.get_binary_indices(milp_data)
    @variable(milp, lb[i] <= x[i=1:num_vars] <= ub[i])
    for k in binary_indices
        set_binary(x[k])
    end
    for k in 1:num_vars
        set_name(x[k], PR.get_variable_names(milp_data)[k])
    end

    # Add constraints
    A, b = PR.get_eq_constraint_matrices(milp_data)
    @constraint(milp, A * x .== b)
    A, b = PR.get_leq_constraint_matrices(milp_data)
    @constraint(milp, A * x .<= b)

    # Add objective.
    set_objective_function(milp, x[milp_data.y_index])

    # Collect solutions from different linear objectives.
    xs = collect(x_min:0.01:x_max)
    max_ys = Float64[]
    min_ys = Float64[]

    for p in xs
        set_lower_bound(x[milp_data.x_index], p)
        set_upper_bound(x[milp_data.x_index], p)

        set_objective_sense(milp, MOI.MAX_SENSE)
        optimize!(milp)
        @assert termination_status(milp) == MOI.OPTIMAL
        yval = value.(x[milp_data.y_index])
        obj = objective_value(milp)
        @assert isapprox(yval, obj, atol = 1e-5) "obj $obj and y $yval different in max problem"
        push!(max_ys, yval)

        set_objective_sense(milp, MOI.MIN_SENSE)
        optimize!(milp)
        @assert termination_status(milp) == MOI.OPTIMAL
        yval = value.(x[milp_data.y_index])
        obj = objective_value(milp)
        @assert isapprox(yval, obj, atol = 1e-5) "obj $obj and y $yval different in min problem"
        push!(min_ys, yval)
    end

    fs = [f(x) for x in xs]
    p = Plots.plot(xs, hcat(max_ys,fs,min_ys), legend=false)
    Plots.savefig(p, out_file)
    @info "finished plotting mip"
end

function get_lp_cut(m::Float64, c::Float64, x_min::Float64, x_max::Float64, y_min::Float64,
        y_max::Float64)::Pair{Vector{Float64},Vector{Float64}}
    xs, ys = Float64[], Float64[]
    for a in [x_min, x_max]
        y = (m * a) + c
        if y >= y_max
            x = (y_max - c) / m
        elseif y >= y_min
            x = a
        else
            x = (y_min - c) / m
        end
        push!(xs, x)
        push!(ys, (m*x)+c)
    end
    return Pair(xs,ys)
end

function plot_lp(f::Function, base_partition::Vector{<:Real}; f_dash::Union{Nothing,Function}=nothing, out_file="plots/lp.pdf")
    if isa(f_dash, Nothing)
        lp_data, fn_data = PR.construct_lp_relaxation(f, base_partition)
    else
        lp_data, fn_data = PR.construct_lp_relaxation(f, f_dash, base_partition)
    end
    lb, ub = PR.get_variable_bounds(lp_data)

    x_min, x_max = lb[lp_data.x_index], ub[lp_data.x_index]
    @info "x bounds: $x_min, $x_max"

    y_min, y_max = lb[lp_data.y_index], ub[lp_data.y_index]
    @info "y bounds: $y_min, $y_max"

    xs = collect(x_min:0.01:x_max)
    fs = [f(x) for x in xs]
    p = Plots.plot(xs,fs,legend=false,color=:orange)

    lp = Model(opt)
    set_optimizer_attribute(lp, "CPXPARAM_ScreenOutput", 0)

    num_vars = PR.get_num_variables(lp_data)
    @variable(lp, lb[i] <= y[i=1:num_vars] <= ub[i])
    A, b = PR.get_eq_constraint_matrices(lp_data)
    @constraint(lp, A * y .== b)

    for k in 1:num_vars
        set_name(y[k], PR.get_variable_names(lp_data)[k])
    end

    res = π/100
    for α ∈ collect(0:res:π)
        xvar = y[lp_data.x_index]
        yvar = y[lp_data.y_index]

        if (isapprox(α, π/2, atol=1e-5) || isapprox(α, -π/2, atol=1e-5))
            continue
        end

        if isapprox(α, 0, atol=1e-5)
            @objective(lp, Max, yvar)
        # elseif (isapprox(α, π/2, atol=1e-5) || isapprox(α, -π/2, atol=1e-5))
        #     @objective(lp, Max, xvar)
        else
            m = tan(α)
            @objective(lp, Max, yvar - (xvar * tan(α)))
        end

        optimize!(lp)
        @assert termination_status(lp) == MOI.OPTIMAL
        c = objective_value(lp)

        if isapprox(α, 0, atol=1e-5)
            line_xs, line_ys = [x_min, x_max], [c, c]
        # elseif (isapprox(α, π/2, atol=1e-5) || isapprox(α, -π/2, atol=1e-5))
        #     line_xs, line_ys = [c, c], [y_min, y_max]
        else
            line_xs, line_ys = get_lp_cut(m,c,x_min,x_max,y_min,y_max)
        end

        if !isempty(line_xs)
            Plots.plot!(line_xs, line_ys, color=:blue)
        end

        set_objective_sense(lp, MOI.MIN_SENSE)
        optimize!(lp)
        @assert termination_status(lp) == MOI.OPTIMAL
        c = objective_value(lp)
        if isapprox(α, 0, atol=1e-5)
            line_xs, line_ys = [x_min, x_max], [c, c]
        # elseif (isapprox(α, π/2, atol=1e-5) || isapprox(α, -π/2, atol=1e-5))
        #     line_xs, line_ys = [c, c], [y_min, y_max]
        else
            line_xs,line_ys = get_lp_cut(m,c,x_min,x_max,y_min,y_max)
        end
        Plots.plot!(line_xs, line_ys, color=:green)
    end

    Plots.savefig(p, out_file)
    @info "finished plotting lp"
end

function generate_plots()
    # f, bp = x->x^3, collect(-1.0:0.25:1.0)
    # f, bp = sin, collect(-(2*π):π/8:(2*π))
    # f, bp = cos, collect(-(2*π):π/8:(2*π))
    # f, bp = log, collect(1.0:1.0:5.0)
    # f, bp = exp, collect(0.0:0.5:2.0)
    # f, bp = cot, collect(π/16:π/32:(π-π/16))
    @info "plotting sin(x)"
    f, bp = sin, collect(0:π/8:(2*π))
    plot_mip(f, bp, out_file="plots/sin-mip.pdf")
    plot_lp(f, bp, out_file="plots/sin-lp.pdf")
    @info "plotting x ⋅ abs(x)"
    f, bp = x -> x * abs(x), collect(-1.0:0.25:1.0)
    plot_mip(f, bp, out_file="plots/xabsx-mip.pdf", f_dash=x->2*abs(x))
    plot_lp(f, bp, out_file="plots/xabsx-lp.pdf", f_dash=x->2*abs(x))
    @info "plotting x^4 - x^3"
    f, bp = x -> x^4 - x^3, collect(-0.5:0.1:1.0)
    plot_mip(f, bp, out_file="plots/poly-mip.pdf")
    plot_lp(f, bp, out_file="plots/poly-lp.pdf")
end


using Test
using JuMP
using MIPVerify
@isdefined(TestHelpers) || include("../../TestHelpers.jl")

@testset "pool.jl" begin

    @testset "Pool" begin
        strides = (1, 2, 2, 1)
        p = Pool(strides, MIPVerify.maximum)
        @test p.strides == strides
        @testset "Base.show" begin
            io = IOBuffer()
            Base.show(io, p)
            @test String(take!(io)) == "max pooling with a 2x2 filter and a stride of (2, 2)"
        end
    end

    @testset "getsliceindex" begin
        @testset "inbounds" begin
            @test MIPVerify.getsliceindex(10, 2, 3)==[5, 6]
            @test MIPVerify.getsliceindex(10, 3, 4)==[10]
        end
        @testset "out of bounds" begin
            @test MIPVerify.getsliceindex(10, 5, 4)==[]
        end
    end

    input_size = (6, 6)
    input_array = reshape(1:*(input_size...), input_size)
    @testset "getpoolview" begin
        @testset "inbounds" begin
            @test MIPVerify.getpoolview(input_array, (2, 2), (3, 3)) == [29 35; 30 36]
            @test MIPVerify.getpoolview(input_array, (1, 1), (3, 3)) == reshape([15], (1,1))
        end
        @testset "out of bounds" begin
            @test length(MIPVerify.getpoolview(input_array, (1, 1), (7, 7))) == 0
        end
    end

    @testset "maxpool" begin
        true_output = [
            8 20 32;
            10 22 34;
            12 24 36
        ]
        @testset "Numerical Input" begin
            @test MIPVerify.pool(input_array, MaxPool((2, 2))) == true_output
        end
        @testset "Variable Input" begin
            m = TestHelpers.get_new_model()
            input_array_v = map(
                i -> @variable(m, lowerbound=i-2, upperbound=i),
                input_array
            )
            pool_v = MIPVerify.pool(input_array_v, MaxPool((2, 2)))
            # elements of the input array are made to take their maximum value
            @objective(m, Max, sum(input_array_v))
            solve(m)

            solve_output = getvalue.(pool_v)
            @test solve_output≈true_output
        end
    end

end


using MCMCChains
using Distributions
using Random
using Test

Random.seed!(20)

@testset "get tests" begin
    # Consider the model
    # ```math
    # m ~ Normal([1, 2, ..., 10], 0.1^2)
    # s ~ IG(2, 3)
    # x ~ Normal(0, s)
    # ```
    # for variable ``x``. Given observation ``x = 2``, sample one chain with 1000 samples
    # from the posterior
    # ```math
    # m ~ Normal([1, 2, ..., 10], 0.1^2)
    # s ∼ IG(2.5, 5)
    # ```
    n_samples = 1000
    vals = hcat(rand(MvNormal(1:10, 0.1), n_samples)',
                rand(InverseGamma(2.5, 5), n_samples))
    chn = Chains(vals, [("m[$i]" for i in 1:10)..., "s"])

    get1 = get(chn, :m)
    get2 = get(chn, [:m, :s])

    @test length(get1.m) == 10
    @test length(get1.m[5]) == n_samples
    @test round(mean(get1.m[1])) ≈ 1.0
    @test round(mean(get1.m[10])) ≈ 10.0
    @test length(get2.s) == n_samples
    @test collect(keys(get2)) == [:m, :s]

    getall = get_params(chn)
    @test getall == get(chn, section = [:parameters])
    @test length(keys(getall)) == 2
end


#This example has been adopted from https://github.com/IntelLabs/ParallelAccelerator.jl/blob/master/examples/black-scholes/black-scholes.jl

function blackscholes_serial(sptprice,
                             strike,
                             rate,
                             volatility,
                             time)
    logterm = log10( sptprice / strike)
    powterm = .5f0 * volatility * volatility
    den = volatility * sqrt(time)
    d1 = (((rate + powterm) * time) + logterm) / den
    d2 = d1 - den
    NofXd1 = cndf2(d1)
    NofXd2 = cndf2(d2)
    futureValue = strike * exp(- rate * time)
    c1 = futureValue * NofXd2
    call_ = sptprice * NofXd1 - c1
    put  = call_ - futureValue + sptprice
end

@inline function cndf2(in)
    out = 0.5f0 + 0.5f0 * erf(0.707106781f0 * in)
    return out
end

function runs(iterations)
    sptprice   = Float32[ 42.0 for i = 1:iterations ]
    initStrike = Float32[ 40.0 + (i / iterations) for i = 1:iterations ]
    rate       = Float32[ 0.5 for i = 1:iterations ]
    volatility = Float32[ 0.2 for i = 1:iterations ]
    time       = Float32[ 0.5 for i = 1:iterations ]

    sptprice_gpu = AFArray(sptprice)
    initStrike_gpu = AFArray(initStrike)
    rate_gpu = AFArray(rate)
    volatility_gpu = AFArray(volatility)
    time_gpu = AFArray(time)

    put1 = blackscholes_serial.(sptprice, initStrike, rate, volatility, time)
    put2 = blackscholes_serial.(sptprice_gpu, initStrike_gpu, rate_gpu, volatility_gpu, time_gpu)
    sync(put2)
    @test sum(put1) ≈ sum(put2)
    put2
end

@inferred runs(10^3)


include("../src/Tau.jl")
using Main.Tau
VERSION < v"0.7.0-beta2.199" ? using Base.Test : using Test

@testset "self-identity" begin
    @test isa(tau, Irrational)
    @test τ == τ
    @test τ == tau
end

@testset "tau vs. 2pi" begin

    @testset "symbols" begin
        @test τ ≠ 2π # tau is Irrational, can't be equal to an AbstractFloat
        @test 2π ≠ τ
        @test π ≠ τ/2
        @test τ/2 ≠ π
    end

    @testset "ascii" begin
        @test tau ≠ 2*pi # tau is Irrational, can't be equal to an AbstractFloat
        @test 2*pi ≠ tau
        @test pi ≠ tau/2 # pi is Irrational, can't be equal to an AbstractFloat
        @test tau/2 ≠ pi
    end

    @testset "explicit type conversions" begin
        @test Float32(tau) == 2 * Float32(pi)
        @test Float64(Float32(tau)) == Float64(2 * Float32(pi))
        @test BigFloat(tau) == 2 * BigFloat(pi)
    end

end

@testset "sintau/costau" begin

    @testset "approximate results for fractional inputs" begin

        @testset "real values" begin
            for T = (Float32, Float64), x = -3:0.1:3.0
                @test @inferred(sintau(T(x))) ≈ T(sinpi(2 * parse(BigFloat, string(x))))
                @test @inferred(costau(T(x))) ≈ T(cospi(2 * parse(BigFloat, string(x))))
            end
        end

        @testset "complex values" begin
            for x in -1.0:0.1:1.0, y in -1.0:0.1:1.0
                z = complex(x, y)
                @test @inferred(sintau(z)) ≈ sinpi(2 * z)
                @test @inferred(costau(z)) ≈ cospi(2 * z)
            end
        end

    end

    @testset "exact results for integer inputs" begin

        @testset "real and complex values passed as integer types" begin
            for T = (Int, Complex), x = -3:3
                @test @inferred(sintau(T(x))) == zero(T)
                @test @inferred(costau(T(x))) == one(T)
            end
        end

        @testset "real and complex values passed as floating point types" begin
            for T = (Float32, Float64, BigFloat, Complex), x = -3.0:3.0
                @test @inferred(sintau(T(x))) == sign(x) * zero(T)
                @test @inferred(costau(T(x))) == one(T)
            end
        end

    end

    @testset "corner cases for abnormal inputs" begin

        @testset "real values" begin
            for x in (Inf, -Inf)
                @test_throws DomainError sintau(x)
                @test_throws DomainError costau(x)
            end
            @test isequal(@inferred(sintau(NaN)), sinpi(NaN))
            @test isequal(@inferred(costau(NaN)), cospi(NaN))
        end

        @testset "complex values" begin
            for x in (Inf, NaN, 0), y in (Inf, NaN, 0)
                z = complex(x, y)
                @test isequal(@inferred(sintau(z)), sinpi(2 * z))
                @test isequal(@inferred(costau(z)), cospi(2 * z))
            end
        end

    end

    # Adapted from julia/test/math.jl
    @testset "type stability" begin
        for T = (Int, Float32, Float64, BigFloat), f = (sintau, costau)
            @test Base.return_types(f, Tuple{T}) == [T]
            @test Base.return_types(f, Tuple{Complex{T}}) == [Complex{float(T)}]
        end
    end
end

# Adapted from julia/test/mod2pi.jl
@testset "modtau" begin
    @test_throws ArgumentError modtau(Int64(2)^60 - 1)
    @test modtau(10) ≈ mod(10, tau)
    @test modtau(-10) ≈ mod(-10, tau)
    @test modtau(355) ≈ 3.1416227979431572
    @test modtau(Int32(355)) ≈ 3.1416227979431572
    @test modtau(355.0) ≈ 3.1416227979431572
    @test modtau(355.0f0) ≈ 3.1416228f0
    @test modtau(Int64(2)^60) == modtau(2.0^60)
end


#=

Authors: Shunsuke Hori

=#

using QuantEcon

"""
The approximate Coleman operator, updated using the endogenous grid
method.

Parameters
----------
g : Function
    The current guess of the policy function
k_grid : AbstractVector
    The set of *exogenous* grid points, for capital k = y - c
β : AbstractFloat
    The discount factor
u_prime : Function
    The derivative u'(c) of the utility function
u_prime_inv : Function
    The inverse of u' (which exists by assumption)
f : Function
    The production function f(k)
f_prime : Function
    The derivative f'(k)
shocks : AbstractVector
    An array of draws from the shock, for Monte Carlo integration (to
    compute expectations).
"""

function coleman_egm(g::Function,
                     k_grid::AbstractVector,
                     β::AbstractFloat,
                     u_prime::Function,
                     u_prime_inv::Function,
                     f::Function,
                     f_prime::Function,
                     shocks::AbstractVector)

    # Allocate memory for value of consumption on endogenous grid points
    c = similar(k_grid)

    # Solve for updated consumption value
    for (i, k) in enumerate(k_grid)
        vals = u_prime.(g.(f(k) * shocks)) .* f_prime(k) .* shocks
        c[i] = u_prime_inv(β * mean(vals))
    end

    # Determine endogenous grid
    y = k_grid + c  # y_i = k_i + c_i

    # Update policy function and return
    Kg = LinInterp(y,c)
    Kg_f(x)=Kg(x)
    return Kg_f
end


@testset verbose = true "Markov" begin

    @testset "Discrete" begin
        dmc = DiscreteMarkovChain([0.3, 0.7], [0.9 0.1; 0.2 0.8])
        states = rand(dmc, 1000)
        dmc_est = fit(DiscreteMarkovChain, states)
        error = transition_matrix(dmc_est) - transition_matrix(dmc)
        @test maximum(error) < 0.2
    end

    @testset "Continuous" begin
        cmc = ContinuousMarkovChain([0.3, 0.7], [-1. 1.; 2. -2.])
        history = rand(cmc, 0., 1000.)
        cmc_est = fit(ContinuousMarkovChain, history)
        error = rate_matrix(cmc_est) - rate_matrix(cmc)
        @test maximum(error) < 0.2
    end

end


using InteractiveDynamics, GLMakie, DynamicalSystems

# the second range is a convenience for intermittency example of logistic
rrange = 1:0.001:4.0
# rrange = (rc = 1 + sqrt(8); [rc, rc - 1e-5, rc - 1e-3])

lo = Systems.logistic(0.4; r=rrange[1])

interactive_cobweb(lo, rrange, 5)


using SymPy
using Test
using LinearAlgebra
using Plots
using Calculus

# using SymbolicUtils

include("../../lib/KTBC.jl")
using .KTBC: CreateKTBC

using Pkg
Pkg.add("Pluto")
Pluto.run()

# 1 Symbolic path
@vars u₁ u₂ x₁ y₁

u⁺ = [u₁ u₂]
u = [
  u₁
  u₂
]

T₂ = convert(Array{Int}, CreateKTBC(2)[2])

u⁺T₂u = u⁺ * T₂ * u

result₁ = [u₁ * (u₁ - u₂) + u₂ * (-u₁ + 2 * u₂)]

@test u⁺T₂u[1] - result₁[1] == 0

# 3
A = [
  1 -1 0
  0 1 -1
  -1 0 1
]

plot3d(A)

C = (A') * A

plot!(C)

@test_throws PosDefException cholesky(C)

@test A[1:3] + A[4:6] + A[7:9] == [0; 0; 0]

# or
u = ones(Int, 3, 1)
z = zeros(Int, 3, 1)

@test A * u == z


export pdfnegbin_truncated

cdfnegbin(k; p, r) = one(k) - cdf(Beta(k+one(k),r),p)
pdfnegbin(k;p,r) = binomial( k+r - one(k), k) * (one(p) - p)^r * (p^k)
meannegbin(p,r) = (p*r) / (1 - p)

function pdfnegbin_truncated(k; p, r, maxwins)
    if k < maxwins
        pdfnegbin(k; p = p, r = r)
    elseif k == maxwins
        pdfnegbin(maxwins; p = p, r = r) + (1 - cdfnegbin(maxwins; p = p, r = r))
    else
        zero(pdfnegbin(k; p = p, r = r))
    end
end


###############################################################################
#
# Atmospheric pressure and partial pressure
#
###############################################################################
GM2RT25(FT) = GRAVITY(FT) * FT(0.02896968) / RK_25(FT);

"""
    atmospheric_pressure(h::FT)

Calculate the atmospheric pressure, given
- `h` elevation in `[m]`
"""
function atmospheric_pressure_ratio(
            h::FT
) where {FT<:AbstractFloat}
    return exp( -h * GM2RT25(FT) );
end

"""
    atmospheric_pressure(h::FT)

Calculate the atmospheric pressure, given
- `h` elevation in `[m]`
"""
function atmospheric_pressure(
            h::FT
) where {FT<:AbstractFloat}
    return P_ATM(FT) * atmospheric_pressure_ratio(h);
end

"""
    ppm_to_Pa(h::FT)

Convert ppm to Pa, given
- `h` elevation in `[m]`
"""
function ppm_to_Pa(
            h::FT
) where {FT<:AbstractFloat}
    return atmospheric_pressure(h) * FT(1e-6)
end


function [d_time]=GH_greens_funct(rd,dt,nt,c,src,opt)

% Inputs:
% f: frequency [Hz]
% rd: zeros(ns,nr);src and rev distance [m].
% src: zeros(1,nt);usually Ricker.
%
% Outputs:
% d_time: data (Greens function) in the time domain.
%
%  Copyright (C) 2019, Exploration and Environmental Geophysics (EEG) at
%  ETH Zurich

src=reshape(src,[1,nt]);

ff = 1/dt;    % [Hz] Sampling rate
fs = linspace(0,ff,nt); % [Hz] frequency vector

k = 2*pi*fs/c;
if strcmp(opt,'greens1d')
    greens = @(r,f) -1i./(2*k) .* exp(-1i*r*k);
elseif strcmp(opt,'greens2d_1st')
    greens = @(r,f) 1i/4 * besselh(0,2, r*k) .* (-2*pi*fs*1i); % first-order system
elseif strcmp(opt,'greens2d_2nd')
    greens = @(r,f) 1i/4 * besselh(0,2, r*k); % second-order system
elseif strcmp(opt,'greens3d')
    greens = @(r,f) exp(-1i*r*k)./(4*pi*r);
else
    error('opt not supported.\n');
end

%% ---------
SRC=fft(src);

d_f = zeros(length(fs),length(rd));
for l=1:length(rd)
    d_f(:,l) = SRC .* sum( greens(rd(:,l),fs) , 1 );
end
d_f( isnan(d_f) ) = 0;

%% 6. Display in time domain
d_time = ifft( d_f, [], 1, 'symmetric' );

end


using Test
using Distances, LinearAlgebra
using KernelFunctions

A = rand(10,5)
B = rand(20,5)
@testset "Distance" begin
    @testset "Dot Product" begin
        d = KernelFunctions.DotProduct()
        @test diag(pairwise(d,A,dims=2)) == [dot(A[:,i],A[:,i]) for i in 1:size(A,2)]
        @test_throws DimensionMismatch d(rand(3),rand(4))
        @test d(3.0,2.0) == 6.0
    end
    @testset "Delta" begin
        d = KernelFunctions.Delta()
        @test pairwise(d,A,dims=1) == Matrix(I,size(A,1),size(A,1))
        @test pairwise(d,A,B,dims=1) == zeros(size(A,1),size(B,1))
        @test d(1,2) == 0
        @test d(1,1) == 1
        @test_throws DimensionMismatch d(rand(3),rand(4))
    end
end


include("graph.jl")
include("zadatak73DFS.jl")

M = [0 1 1 0 0 0 1;
     1 0 0 0 0 0 0;
     1 0 0 1 1 0 0;
     0 0 1 0 0 0 1;
     0 0 1 0 0 1 0;
     0 0 0 0 1 0 1;
     1 0 0 1 0 1 0]

G = nonEmptyGraph(M)
DFS!(G, 1, 6)

printGraph(G)


using Luxor, Random, Test

function test_tangents(fname)
    Drawing(500, 500, fname)
    background("white")
    origin()
    setline(0.5)
    sethue("black")
    for i in 1:30
        circlecenter = polar(rand(100:190), 2π * i/30)
        circleradius = rand(4:10)
        point = polar(rand(170:240), 2π * i/20)
        circle(circlecenter, 2, :fill)
        circle(point, 2, :fill)
        circle(circlecenter, circleradius, :fill)
        pt1, pt2 = pointcircletangent(point, circlecenter, circleradius)
        if pt1 != pt2
            line(point, pt1, :stroke)
            line(point, pt2, :stroke)
        end
    end
    @test finish() == true
end

fname = "point-circle-tangent.png"
test_tangents(fname)
println("...finished test: output in $(fname)")


@everywhere begin
    using StochasticSearch
    function rosenbrock(x::Configuration, parameters::Dict{Symbol, Any})
        return (1.0 - x["i0"].value)^2 + 100.0 * (x["i1"].value - x["i0"].value^2)^2
    end
end

configuration = Configuration([FloatParameter(-2.0, 2.0, 0.0,"i0"),
                               FloatParameter(-2.0, 2.0, 0.0,"i1")],
                               "rosenbrock_config")

tuning_run = Run(cost               = rosenbrock,
                 starting_point     = configuration,
                 methods            = [[:simulated_annealing 1];
                                       [:iterative_first_improvement 1];
                                       [:randomized_first_improvement 1];
                                       [:iterative_greedy_construction 1];
                                       [:iterative_probabilistic_improvement 1];])

search_task = @task optimize(tuning_run)
result = consume(search_task)
print(result)
while result.is_final == false
    result = consume(search_task)
    print(result)
end


module ContinuedFractions
"""
Calculate the continued fraction until convergence against specified tolerance
∑ᵢ uᵢ/vᵢ = v(0) + u(1) / (v(1) + u(2) / (v(2) + u(3)/(v(3) + ...)))
"""

export continuedfraction

"""
    continuedfraction(v, u; atol=0, rtol=sqrt(eps(typeof(real(v(1) / u(1)))))
Calculate the continued fraction until convergence against specified tolerance
∑ᵢ uᵢ/vᵢ = v(0) + u(1) / (v(1) + u(2) / (v(2) + u(3)/(v(3) + ...)))
"""
function continuedfraction(v::V, u::U; atol=0,
    rtol=sqrt(eps(typeof(real(v(1) / u(1)))))
    ) where {V<:Function, U<:Function}
  n = 8
  a, b = continuedfraction(v, u, n, 2n)
  n *= 2
  converged = isapprox(a, b, rtol=rtol, atol=atol)
  while !converged && n <= 2^10
    a = deepcopy(b)
    n *= 2
    b = continuedfraction(v, u, n)
    converged = isapprox(a, b, rtol=rtol, atol=atol)
  end
  return b, converged
end
function continuedfraction(v::V, u::U, n::Int, m::Int
    ) where {V<:Function, U<:Function}
  @assert m > n
  output_m = u(m) / v(m)
  for i ∈ reverse(n+1:m-1)
    output_m = u(i) / (v(i) + output_m)
  end
  ui, vi = u(n), v(n)
  output_m = ui / (vi + output_m)
  output_n = ui / vi
  for i ∈ reverse(1:n-1)
    ui, vi = u(i), v(i)
    output_m = ui / (vi + output_m)
    output_n = ui / (vi + output_n)
  end
  vi = v(0)
  return vi + output_n, vi + output_m
end
function continuedfraction(v::V, u::U, n::Int) where {V<:Function, U<:Function}
  output = u(n) / v(n)
  for i ∈ reverse(1:n-1)
    output = u(i) / (v(i) + output)
  end
  return v(0) + output
end

end


using RecursiveArrayTools: ArrayPartition
using ProximalCore

struct ComposeAffine
    f
    A
    b
end

function (g::ComposeAffine)(x)
    res = g.A * x .+ g.b
    return g.f(res)
end

function compose_affine_gradient!(y, g::ComposeAffine, x)
    res = g.A * x .+ g.b
    gradres, v = gradient(g.f, res)
    mul!(y, adjoint(g.A), gradres)
    return v
end

ProximalCore.gradient!(y, g::ComposeAffine, x) = compose_affine_gradient!(y, g, x)
ProximalCore.gradient!(y::ArrayPartition, g::ComposeAffine, x::ArrayPartition) = compose_affine_gradient!(y, g, x)

function ProximalCore.gradient(h::ComposeAffine, x::ArrayPartition)
    grad_fx = similar(x)
    fx = ProximalCore.gradient!(grad_fx, h, x)
    return grad_fx, fx
end


using SolveDSGE, Plots;
if 1==1
    open("Det_INV.txt", "w") do io
        println(io,"# inv\n")
        println(io,"states:\nk, z \nend\n")
        println(io,"jumps:\ni \nend\n")
        println(io,"shocks:\nε \nend\n")
        println(io,"parameters:\nδ=0.10, ρ=0.15 \nend\n")
        println(io,"equations:")
        println(io,"k(+1) = (1-δ)*k + i")
        println(io,"i(+1) = ((ρ+δ- (0.5 + 0.5*z) )/(1-δ)) +((1+ρ)/(1-δ))*i")
        println(io,"z(+1) = 0.5 + 0.5*z + 1*ε")
        println(io,"end")
    end;

    ###########################################################################
    path = joinpath(@__DIR__,"Det_INV.txt")
    process_model(path)
    ngm = retrieve_processed_model(path)
    ###########################################################################
    # Steady State
    ###########################################################################
    tol = 1e-8
    maxiters = 1000
    x = .25*ones(ngm.number_variables) #x = [0.25,0.25] #x = [0.25,0.25, .25] #x = [0.25,0.25, .25, .25]
    ss = compute_steady_state(ngm, x, tol, maxiters)
    ss = ss.zero
    #ss ≈ [K_SS, I_SS]
    sss = ss[1:ngm.number_states] #ss for states only
    ###########################################################################
    # Perturbation, Order = 1,2,3
    # soln_fo/soln_so/soln_to
    ###########################################################################
    if 1==1
        g0 = ss;          # point about which to perturb the model (generally the steady state)
        cutoff = 1.0      # param that separates unstable from stable eigenvalues
        # order = "first" # order of the perturbation
        N   = PerturbationScheme(g0,cutoff,"first")   # Klein (2000)
        NN  = PerturbationScheme(g0,cutoff,"second")  # Gomme and Klein (2011)
        NNN = PerturbationScheme(g0,cutoff,"third")   # Binning (2013) w refinement from Levintal (2017)

        soln_fo = solve_model(ngm, N)
        soln_so = solve_model(ngm, NN)
        soln_to = solve_model(ngm, NNN)
    end

    # To Do: Stochastic Projection
    # ProjectionSchemes: ChebyshevSchemes/SmolyakSchemes/PiecewiseLinearSchemes

    # Simulation
    if 1 == 1
        T = 100 # num Sim pds
        ω = 1.0 - 0.2; # distance from steady-state
        #ω = 1.0 - 0.75;

        sim_d1   = simulate(soln_fo, sss*ω,  T)
        sim_d2   = simulate(soln_so, sss*ω,  T)
        sim_d3   = simulate(soln_to, sss*ω,  T)
        # sim_d4   = simulate(soln_cha, sss*ω, T)
        # sim_d5   = simulate(soln_sa, sss*ω,  T)
        # sim_d6   = simulate(soln_pl, sss*ω,  T)
        # sim_d7   = simulate(soln_c1, sss*ω,  T)

        # Closed Form Sol

        lv = reshape(ngm.variables,1,ngm.number_variables)
        p  = plot(legend=:bottomright);
        #p0 = plot!(1:T,cf,      title="Closed Form",    lab="");
        p0 = plot!()
        p1 = plot(p0, 1:T,sim_d1', title="Pert. 1st Order", lab=lv)
        p2 = plot(p0, 1:T,sim_d2', title="Pert. 2nd Order", lab=lv)
        p3 = plot(p0, 1:T,sim_d3', title="Pert. 3rd Order", lab=lv)
        # p4 = plot(p0, 1:T,sim_d4', title="Proj. Cheby 1",   lab=lv)
        # p5 = plot(p0, 1:T,sim_d5', title="Proj. Smol  1",   lab=lv)
        # p6 = plot(p0, 1:T,sim_d6', title="Proj. PL    1",   lab=lv)
        # p7 = plot(p0, 1:T,sim_d7', title="Proj. & Pert",    lab=lv)
        #p8 = plot!(1:T,sim_d8', title="Proj. & Pert",   lab=lv)
        plot(p1,p2,p3,size=1.25 .* (600, 400))
    end
    #
end


using DSP
using PyPlot

# Generate a signal u of length m = 50, with each entry a random value that is either −1 or +1
u = rand([1, -1], 50)

c = [1, 0.7, -0.3]
y = conv(c, u)

# Plot u and y
plot(u) # Blue Graph
plot(y) # Orange Graph

h = [0.9, -0.5, 0.5, -0.4, 0.3, -0.3, 0.2, -0.1]
z = conv(h , y)

# Plot for z = h * y
plot(z)


using Devectorize
include("helper_functions.jl")

function fprop(input_batch,
              word_embedding_weights,
              embed_to_hid_weights,
              hid_to_output_weights,
              hid_bias,
              output_bias)
#=
To forward propagate through the neural network.
---
Inputs:
  input_batch :
  The input data as a matrix of size numwords X batchsize where,
  numwords is the number of words, batchsize is the number of data points.
  So, if input_batch(i, j) = k then the ith word in data point j is word
  index k of the vocabulary.
=#

# Compute the state of the word embedding layer
# Here simply look up the word_embedding_weights and choose the corresponding weights
numwords, batchsize = size(input_batch)
vocab_size, numofhid1units = size(word_embedding_weights)
numofhid2units = size(embed_to_hid_weights, 2)

# Compute the stae of the WORD EMBEDDING LAYER
# Simple look up in inputs word indices in the word_embedding_weights matrix.
embedding_layer_state = modified_reshape(
                                  word_embedding_weights[
                                    vec(modified_reshape(input_batch, row_number=1)), :]',
                                        row_number=numofhid1units * numwords)
# Compute the state of the HIDDEN LAYER

# First Compute inputs to hidden units
inputs_to_hidden_units = embed_to_hid_weights' * embedding_layer_state + repmat(hid_bias, 1, batchsize)

# Next step is to apply the logistic function to find activations
@devec hidden_layer_state = 1./ (1 + exp(-inputs_to_hidden_units))

# Compute the state of the OUTPUT LAYER

# First we need to compute the inputs for softmax units
inputs_to_softmax = hid_to_output_weights' * hidden_layer_state + repmat(output_bias, 1, batchsize)

# This trick is taken from Hinton's homeworks
inputs_to_softmax  = inputs_to_softmax - repmat(findmaxcols(inputs_to_softmax), vocab_size, 1)

# Apply the softmax approach
@devec output_layer_state = exp(inputs_to_softmax)

# Normalize the output_layer_state
output_layer_state = output_layer_state ./ repmat(sum_columns(output_layer_state), vocab_size, 1)

return (embedding_layer_state, hidden_layer_state, output_layer_state)
end


function test_measures()
    m1 = Lebesgue()
    @test !isdiscrete(m1)
    @test iscontinuous(m1)
    @test !isnormalized(m1)
    @test domaintype(m1) == Float64
    @test DomainIntegrals.unsafe_weightfun(m1, 0.4) == 1

    m2 = DomainIntegrals.LebesgueUnit()
    @test !isdiscrete(m2)
    @test iscontinuous(m2)
    @test isnormalized(m2)
    @test domaintype(m2) == Float64
    @test support(m2) == UnitInterval()
    @test weightfun(m2, 0.4) == 1
    @test weightfun(m2, 1.4) == 0
    @test m2 == dx(UnitInterval())

    m3 = LegendreWeight()
    @test !isdiscrete(m3)
    @test iscontinuous(m3)
    @test !isnormalized(m3)
    @test domaintype(m3) == Float64
    @test support(m3) == ChebyshevInterval()
    @test weightfun(m3, 0.4) == 1
    @test weightfun(m3, -0.4) == 1
    @test weightfun(m3, 1.4) == 0
    @test weightfun(m3, -1.4) == 0
    @test jacobi_α(m3) == 0
    @test jacobi_β(m3) == 0

    x = 0.5
    m4 = DiracWeight(x)
    @test !isdiscrete(m4)
    @test iscontinuous(m4)
    @test isnormalized(m4)
    @test domaintype(m4) == Float64
    @test support(m4) == Point(x)
    @test point(m4) == x
    @test weightfun(m4, x) == Inf
    @test weightfun(m4, x+1) == 0
    @test weightfun(m4, big(x)) == Inf
    @test weightfun(m4, big(x+1)) == 0

    m5 = GaussianWeight{SVector{2,Float64}}()
    @test !isdiscrete(m5)
    @test iscontinuous(m5)
    @test isnormalized(m5)
    @test domaintype(m5) == SVector{2,Float64}
    @test support(m5) == FullSpace{SVector{2,Float64}}()
    @test weightfun(m5, SVector(0.0,0.0)) ≈ 1/(2pi)
    @test weightfun(m5, SVector(0,0)) ≈ 1/(2pi)
    @test weightfun(m5, SVector(big(0),big(0))) ≈ 1/(2pi)
    @test weightfun(m5, SVector(big(0.0),big(0.0))) ≈ 1/(2pi)

    m6 = LaguerreWeight(0.0)
    @test !isdiscrete(m6)
    @test iscontinuous(m6)
    @test isnormalized(m6)
    @test !isnormalized(LaguerreWeight(0.1))
    @test domaintype(m6) == Float64
    @test support(m6) == HalfLine()
    @test weightfun(m6, 0.4) == exp(-0.4)
    @test weightfun(m6, -0.4) == 0
    @test weightfun(m6, big(0.4)) == exp(-big(0.4))

    m7 = ChebyshevTWeight()
    @test support(m7) isa ChebyshevInterval
    @test !isdiscrete(m7)
    @test iscontinuous(m7)
    @test !isnormalized(m7)
    @test domaintype(m7) == Float64
    @test jacobi_α(m7) == -1/2
    @test jacobi_β(m7) == -1/2
    @test m7 == JacobiWeight(-1/2, -1/2)

    m8 = ChebyshevUWeight()
    @test support(m8) isa ChebyshevInterval
    @test !isdiscrete(m8)
    @test iscontinuous(m8)
    @test !isnormalized(m8)
    @test domaintype(m8) == Float64
    @test jacobi_α(m8) == 1/2
    @test jacobi_β(m8) == 1/2
    @test m8 == JacobiWeight(1/2, 1/2)
end


ipa = Dict()
ipb = Dict()
ops = Dict()

# Input membership functions
ipa["small"] = TriangularMF(1, 2, 3)
ipa["large"] = TriangularMF(5, 6, 7)
ipb["small"] = TriangularMF(1, 2, 3)
ips = [ipa, ipb]

# Output membership functions
ops["small"] = TriangularMF(1, 2, 3)
ops["large"] = TriangularMF(5, 6, 7)

# Simple rule
rule1 = Rule(["small", "small"], "small")
rule2 = Rule(["large", "small"], "large")
rules = [rule1, rule2]

fis = FISMamdani(ips, ops, rules)

# Input values
ival1 = [1.2, 2.3]
ival2 = [5.7, 2.0]

@assert eval_fis(fis, ival1) == 2.0
@assert eval_fis(fis, ival2) == 6.0

service = Dict( "poor"      => TrapezoidalMF(0.0, 0.0, 2.0, 4.0),
                "good"      => TriangularMF(3., 5., 7.),
                "excellent" => TrapezoidalMF(6., 8., 10., 10.))

food = Dict("rancid"    => TrapezoidalMF(0., 0., 3., 6.),
            "delicious" => TrapezoidalMF(4., 7., 10., 10.))

inputs = [service, food]

tip  = Dict("cheap"     => TrapezoidalMF(10., 10., 20., 30.),
            "average"   => TriangularMF(20., 30., 40.),
            "generous"  => TrapezoidalMF(30., 40, 50., 50.))

rule1 = Rule(["poor", "rancid"], "cheap", "MAX")

# How to not include food input ?
rule2 = Rule(["good", ""], "average", "MAX")

rule3 = Rule(["excellent", "delicious"], "generous", "MAX")
rules = [rule1, rule2, rule3]

fis_tips = FISMamdani(inputs, tip, rules)

in_vals = [7., 8,]

@assert eval_fis(fis_tips, in_vals) == 45.0


#==
Example taken from
Indraneel Das and John E Dennis.
Normal-boundary intersection: A new method for generating the pareto surface in
nonlinear multicriteria optimization problems.
SIAM Journal on Optimization, 8(3):631–657, 1998.
==#

using MultiJuMP, JuMP
using Ipopt

m = multi_model(Ipopt.Optimizer)
@variable(m, x[i=1:5])
@NLexpression(m, f1, sum(x[i]^2 for i=1:5))
@NLexpression(m, f2, 3x[1]+2x[2]-x[3]/3+0.01*(x[4]-x[5])^3)
@NLconstraint(m, x[1]+2x[2]-x[3]-0.5x[4]+x[5]==2)
@NLconstraint(m, 4x[1]-2x[2]+0.8x[3]+0.6x[4]+0.5x[5]^2 == 0)
@NLconstraint(m, sum(x[i]^2 for i=1:5) <= 10)

obj1 = SingleObjective(f1)
obj2 = SingleObjective(f2)

md = get_multidata(m)
md.objectives = [obj1, obj2]
md.pointsperdim = 20
optimize!(m, method = NBI())

using Plots
plot(md)


"""
Matrix of a 2D piecewise linear FEM discretization
=======================

*Input options:*

+ n: size of the matrix
"""
function fem2d(n::Int; nnzrow=4)
    A = ExtendableSparseMatrix(n, n)
    sprand_sdd!(A, nnzrow=nnzrow)
    flush!(A)
    return A.cscmatrix
end

"""
Matrix for a mock finite difference operator for a diffusion
problem with random coefficients on a unit hypercube
=======================

*Input options:*

+ n: number of unknowns in x, y, z direction
"""
function fdm3d(n::Int)
	nx = floor(Int, n^(1/3))
	ny = floor(Int, n^(1/3))
	nz = floor(Int, n^(1/3))
    A = ExtendableSparseMatrix(nx*ny*nz, nx*ny*nz)
    fdrand!(A, nx, ny, nz)
    flush!(A)
    return  A.cscmatrix
end


"""
    xbrs(RR::Vector{<:Real}, SBP::Vector{<:Real};...)

Calculate the xBRS index for the assesment of baroreflex sensitivity. Based on Westerhof, B. E. et al. (2004). Time-domain cross-correlation baroreflex sensitivity: performance on the EUROBAVAR data set. Journal of hypertension, 22(7), 1371-1380

# Keyword Arguments

* minCor: Minimal significant correlation, dafaults to 0.632 (p = 0.05, two-sided for 10 s window)
* tExcerpt: Length of the sliding window in seconds, defaults to 10 s
* delays: Which delays, shift of the RR window shall be considered, defaults to 0:5

# Return

Returns an xBRS structure, that can be be plotted for visual inspection.
"""
function xbrs(RR::Vector{<:Real}, SBP::Vector{<:Real}; minCor::Real = 0.632, tExcerpt::Int = 10, delays::Union{Vector{Int}, UnitRange} = 0:5)
    1 > minCor > 0 || throw(DomainError("The minimal correlation hast to be postive and ≤1"))

    # Resample signals at 1 Hz through cubic splines
    time = cumsum(RR./1000) # ms to s
    time1Hz = 1:round(Int, time[end])
    S1Hz = Spline1D(time, SBP; k = 3)(time1Hz)
    RR1Hz = Spline1D(time, RR; k = 3)(time1Hz)

    # Correlation and xBRS
    len = time1Hz[end] - tExcerpt - maximum(delays)
    maxCors = zeros(Float64, len)
    delayMaxCors = zeros(Int, len)
    xBRSs = zeros(Float64, len)
    xBRSs2 = zeros(Float64, len)
    for i in 1:len
        maxCor, delayMaxCor = findmax(map(delay -> cor(S1Hz[i:i+tExcerpt-1], RR1Hz[(i:i+tExcerpt-1) .+ delay]), delays))
        maxCors[i] = maxCor
        delayMaxCors[i] = delayMaxCor - 1
        xBRSs[i] = std(RR1Hz[(i:i+tExcerpt-1) .+ delayMaxCors[i]]) / std(S1Hz[i:i+tExcerpt-1]) # 2017
        xBRSs2[i] = linReg(S1Hz[i:i+tExcerpt-1], RR1Hz[(i:i+tExcerpt-1) .+ delayMaxCors[i]])[1] / maxCors[i] # 2004
    end

    valid = findall(x -> x >= minCor, maxCors)
    xBRSg = gmean(xBRSs[valid])
    xBRSg2 = gmean(xBRSs2[valid])

    return xBRS(value = xBRSg,
    xBRSg2 = xBRSg2,
    xBRSs = xBRSs,
    xBRSs2 = xBRSs2,
    valid = valid,
    nValid = length(valid),
    maxCors = maxCors,
    delayMaxCors = delayMaxCors,
    RR = RR,
    SBP = SBP)
end

"""
Struct that stores all information regarting the xBRS etimation. The final result is stored in 'xBRSg'. It can be plotted for visual inspection.
"""
@with_kw mutable struct xBRS
    value::Real = 0
    xBRSg2::Real = 0
    nValid::Int = 0
    xBRSs::Vector{<:Real} = Float64[]
    xBRSs2::Vector{<:Real} = Float64[]
    maxCors::Vector{<:Real} = Float64[]
    delayMaxCors::Vector{Int} = Int[]
    valid::Vector{Int} = Int[]
    RR::Vector{<:Real} = Float64[]
    SBP::Vector{<:Real} = Float64[]
end

"""
linReg(x, y)

Solves the linear regression problem y = mx + b using least squares method through '\\'. Returns (m, b)
"""
linReg(x::Vector{<:Real}, y::Vector{<:Real}) = ([x ones(length(x))] \ y)

"""
    gmean(x::Vector{<:Real})

Compute the geometric mean for the positive Vector x.

"""
gmean(x::Vector{<:Real}) = exp(1/length(x) * sum(log.(x)))

@recipe function f(res::xBRS)
    layout := (4,1)

    @series begin
        subplot := 1
        yguide := "SBP [mmHg]"
        label := ""
        seriescolor := :grey

        res.SBP
    end

    @series begin
        subplot := 2
        yguide := "RR [ms]"
        label := ""
        seriescolor := :grey

        res.RR
    end

    @series begin
        subplot := 3
        marker := :d
        yguide := "correlation"
        label := ""
        seriescolor := :grey

        res.valid, res.maxCors[res.valid]
    end

    @series begin
        subplot := 3
        seriestype := :hline
        yguide := "correlation"
        label := "mean valid cor: $(round(mean(res.maxCors[res.valid]), digits = 2))"
        legend := :topright
        seriescolor := :red

        [mean(res.maxCors[res.valid])]
    end

    @series begin
        subplot := 4
        # seriestype := :scatter
        marker := :d
        yguide := "xBRS [ms/mmHg]"
        label := ""
        seriescolor := :grey

        res.valid, res.xBRSs[res.valid]
    end

    @series begin
        subplot := 4
        seriestype := :hline
        yguide := "xBRS [ms/mmHg]"
        xguide := "beat #"
        label := "xBRS: $(round(res.value, digits = 2)) ms/mmHg"
        legend := :topright
        seriescolor := :steelblue

        [res.xBRSg]
    end
end


# Distributions with compound variates
#
#	NormalGamma
#	NormalInverseGamma
#	NormalWishart
#	NormalInverseWishart
#
# and their applications for conjugates & posteriors
#

using Base.Test
using Distributions

### NormalGamma

n = 100

mu0 = 2.
nu0 = 3.
shape0 = 5.
rate0 = 2.
ng = NormalGamma(mu0, nu0, shape0, rate0)

# Random number generation
mu, tau2 = rand(ng)

# Did it generate something valid?
@test insupport(NormalGamma, mu, tau2)

mu = 2.5
tau2 = 3.

# pdf
npdf = pdf(Normal(mu0, 1./sqrt(nu0*tau2)), mu)
gpdf = pdf(Gamma(shape0, 1./rate0), tau2)
lnpdf = logpdf(Normal(mu0, 1./sqrt(nu0*tau2)), mu)
lgpdf = logpdf(Gamma(shape0, 1./rate0), tau2)

@test_approx_eq_eps pdf(ng, mu, tau2) (npdf*gpdf) 1e-8
@test_approx_eq_eps logpdf(ng, mu, tau2) (lnpdf+lgpdf) 1e-8

### NormalInverseGamma

n = 100

mu0 = 2.
v0 = 3.
shape0 = 5.
scale0 = 2.
nig = NormalInverseGamma(mu0, v0, shape0, scale0)

# Random number generation
mu, sig2 = rand(nig)

# Did it generate something valid?
@test insupport(NormalInverseGamma, mu, sig2)

mu = 2.5
sig2 = 3.

# pdf
npdf = pdf(Normal(mu0, sqrt(v0*sig2)), mu)
gpdf = pdf(InverseGamma(shape0, scale0), sig2)
lnpdf = logpdf(Normal(mu0, sqrt(v0*sig2)), mu)
lgpdf = logpdf(InverseGamma(shape0, scale0), sig2)

@test_approx_eq_eps pdf(nig, mu, sig2) (npdf*gpdf) 1e-8
@test_approx_eq_eps logpdf(nig, mu, sig2) (lnpdf+lgpdf) 1e-8

### NormalWishart

n = 100

mu0 = [2., 3.]
kappa0 = 3.
nu0 = 4.
T0 = eye(2)
T0[1,2] = T0[2,1] = .5
nw = NormalWishart(mu0, kappa0, T0, nu0)

# Random number generation
mu, Lam = rand(nw)

# Did it generate something valid?
@test insupport(NormalWishart, mu, Lam)

mu = [1.5, 2.5]
T = 0.75*eye(2)
T[1,2] = T[2,1] = 0.6

# pdf
npdf = pdf(MultivariateNormal(mu0, inv(kappa0*T)), mu)
wpdf = pdf(Wishart(nu0, T0), T)
lnpdf = logpdf(MultivariateNormal(mu0, inv(kappa0*T)), mu)
lwpdf = logpdf(Wishart(nu0, T0), T)

@test_approx_eq_eps pdf(nw, mu, T) (npdf*wpdf) 1e-8
@test_approx_eq_eps logpdf(nw, mu, T) (lnpdf+lwpdf) 1e-8

### NormalInverseWishart

n = 100

mu0 = [2., 3.]
kappa0 = 3.
nu0 = 4.
T0 = eye(2)
T0[1,2] = T0[2,1] = .5
niw = NormalInverseWishart(mu0, kappa0, T0, nu0)

# Random number generation
mu, Sig = rand(niw)

# Did it generate something valid?
@test insupport(NormalInverseWishart, mu, Sig)

mu = [1.5, 2.5]
T = 0.75*eye(2)
T[1,2] = T[2,1] = 0.6

# pdf
npdf = pdf(MultivariateNormal(mu0, 1/kappa0*T), mu)
wpdf = pdf(InverseWishart(nu0, T0), T)
lnpdf = logpdf(MultivariateNormal(mu0, 1/kappa0*T), mu)
lwpdf = logpdf(InverseWishart(nu0, T0), T)

@test_approx_eq_eps pdf(niw, mu, T) (npdf*wpdf) 1e-8
@test_approx_eq_eps logpdf(niw, mu, T) (lnpdf+lwpdf) 1e-8


function getTableauHPARK(name, q::Tableau{T}, p::Tableau{T}, d=[]) where {T}

    @assert q.s == p.s

    o = min(q.o, p.o)

    a_q = q.a
    a_p = p.a
    b_q = q.b
    b_p = p.b
    c_q = q.c
    c_p = p.c

    if length(d) == 0
        return TableauHPARK(name, o,
                            a_q, a_p, a_q, a_p,
                            a_q, a_p, a_q, a_p,
                            b_q, b_p, b_q, b_p,
                            c_q, c_p, c_q, c_p)
    else
        @assert length(d) == q.s == p.s

        return TableauHPARK(name, o,
                            a_q, a_p, a_q, a_p,
                            a_q, a_p, a_q, a_p,
                            b_q, b_p, b_q, b_p,
                            c_q, c_p, c_q, c_p,
                            d)
    end
end

"SPARK tableau for Gauss-Lobatto IIIA-IIIB HPARK method with s stages."
function TableauHPARKLobattoIIIAIIIB(s)
    getTableauHPARK(Symbol("HPARKLobattoIIIAIIIB($s)"), TableauLobattoIIIA(s), TableauLobattoIIIB(s))
end

"SPARK tableau for Gauss-Lobatto IIIB-IIIA  method with s stages."
function TableauHPARKLobattoIIIBIIIA(s)
    getTableauHPARK(Symbol("HPARKLobattoIIIBIIIA($s)"), TableauLobattoIIIB(s), TableauLobattoIIIA(s))
end

"Tableau for Gauss-Legendre HPARK method with s stages."
function TableauHPARKGLRK(s)
    glrk = TableauGauss(s)
    getTableauHPARK(Symbol("HPARKGLRK", s), glrk, glrk)
end


using Test
using QuantumOptics
using FFTW, LinearAlgebra, Random

@testset "particle" begin

Random.seed!(0)

D(op1::AbstractOperator, op2::AbstractOperator) = abs(tracedistance_nh(dense(op1), dense(op2)))

N = 200
xmin = -32.5
xmax = 20.1

basis_position = PositionBasis(xmin, xmax, N)
basis_momentum = MomentumBasis(basis_position)

b2 = PositionBasis(basis_momentum)
@test basis_position.xmax - basis_position.xmin ≈ b2.xmax - b2.xmin
@test basis_position.N == b2.N

# Test Gaussian wave-packet in both bases
x0 = 5.1
p0 = -3.2
sigma = 1.
sigma_x = sigma/sqrt(2)
sigma_p = 1.0/(sigma*sqrt(2))

psi0_bx = gaussianstate(basis_position, x0, p0, sigma)
psi0_bp = gaussianstate(basis_momentum, x0, p0, sigma)

@test 1 ≈ norm(psi0_bx)
@test 1 ≈ norm(psi0_bp)

p_bx = momentum(basis_position)
x_bx = position(basis_position)

@test 1e-10 > D(p_bx, transform(basis_position, basis_momentum)*dense(momentum(basis_momentum))*transform(basis_momentum, basis_position))

p_bp = momentum(basis_momentum)
x_bp = position(basis_momentum)

@test x0 ≈ expect(x_bx, psi0_bx)
@test 0.1 > abs(p0 - expect(p_bx, psi0_bx))
@test p0 ≈ expect(p_bp, psi0_bp)
@test 0.1 > abs(x0 - expect(x_bp, psi0_bp))

@test 1e-13 > abs(variance(x_bx, psi0_bx) - sigma^2/2)
@test 1e-13 > abs(variance(x_bp, psi0_bp) - sigma^2/2)
@test 1e-13 > abs(variance(p_bx, psi0_bx) - 1/(2*sigma^2))
@test 1e-13 > abs(variance(p_bp, psi0_bp) - 1/(2*sigma^2))

# Test potentialoperator
V(x) = x^2
V_bx = potentialoperator(basis_position, V)
V_bp = potentialoperator(basis_momentum, V)

@test expect(V_bp, psi0_bp) ≈ expect(V_bx, psi0_bx)

# Test FFT transformation
function transformation(b1::MomentumBasis, b2::PositionBasis, psi::Ket)
    Lp = (b1.pmax - b1.pmin)
    dx = particle.spacing(b2)
    if b1.N != b2.N || abs(2*pi/dx - Lp)/Lp > 1e-12
        throw(IncompatibleBases())
    end
    N = b1.N
    psi_shifted = exp.(1im*b2.xmin*(particle.samplepoints(b1) .- b1.pmin)).*psi.data
    psi_fft = exp.(1im*b1.pmin*particle.samplepoints(b2)).*ifft(psi_shifted)*sqrt(N)
    return Ket(b2, psi_fft)
end

function transformation(b1::PositionBasis, b2::MomentumBasis, psi::Ket)
    Lx = (b1.xmax - b1.xmin)
    dp = particle.spacing(b2)
    if b1.N != b2.N || abs(2*pi/dp - Lx)/Lx > 1e-12
        throw(IncompatibleBases())
    end
    N = b1.N
    psi_shifted = exp.(-1im*b2.pmin*(particle.samplepoints(b1) .- b1.xmin)).*psi.data
    psi_fft = exp.(-1im*b1.xmin*particle.samplepoints(b2)).*fft(psi_shifted)/sqrt(N)
    return Ket(b2, psi_fft)
end

psi0_bx_fft = transformation(basis_position, basis_momentum, psi0_bx)
psi0_bp_fft = transformation(basis_momentum, basis_position, psi0_bp)

@test 0.1 > abs(x0 - expect(x_bp, psi0_bx_fft))
@test p0 ≈ expect(p_bp, psi0_bx_fft)
@test 0.1 > abs(p0 - expect(p_bx, psi0_bp_fft))
@test x0 ≈ expect(x_bx, psi0_bp_fft)

@test 1e-12 > norm(psi0_bx_fft - psi0_bp)
@test 1e-12 > norm(psi0_bp_fft - psi0_bx)

Tpx = transform(basis_momentum, basis_position)
Txp = transform(basis_position, basis_momentum)
psi0_bx_fft = Tpx*psi0_bx
psi0_bp_fft = Txp*psi0_bp

@test 0.1 > abs(x0 - expect(x_bp, psi0_bx_fft))
@test p0 ≈ expect(p_bp, psi0_bx_fft)
@test 0.1 > abs(p0 - expect(p_bx, psi0_bp_fft))
@test x0 ≈ expect(x_bx, psi0_bp_fft)

@test 1e-12 > norm(psi0_bx_fft - psi0_bp)
@test 1e-12 > norm(psi0_bp_fft - psi0_bx)

@test 1e-12 > norm(dagger(Tpx*psi0_bx) - dagger(psi0_bx)*dagger(Tpx))
@test 1e-12 > norm(dagger(Txp*psi0_bp) - dagger(psi0_bp)*dagger(Txp))

# Test gemv!
psi_ = deepcopy(psi0_bp)
operators.gemv!(Complex(1.), Tpx, psi0_bx, Complex(0.), psi_)
@test 1e-12 > norm(psi_ - psi0_bp)
operators.gemv!(Complex(1.), Tpx, psi0_bx, Complex(1.), psi_)
@test 1e-12 > norm(psi_ - 2*psi0_bp)

psi_ = deepcopy(psi0_bx)
operators.gemv!(Complex(1.), Txp, psi0_bp, Complex(0.), psi_)
@test 1e-12 > norm(psi_ - psi0_bx)
operators.gemv!(Complex(1.), Txp, psi0_bp, Complex(1.), psi_)
@test 1e-12 > norm(psi_ - 2*psi0_bx)

alpha = complex(3.2)
beta = complex(-1.2)
randdata1 = rand(ComplexF64, N)
randdata2 = rand(ComplexF64, N)

state = Ket(basis_position, randdata1)
result_ = Ket(basis_momentum, copy(randdata2))
result0 = alpha*dense(Tpx)*state + beta*result_
operators.gemv!(alpha, Tpx, state, beta, result_)
@test 1e-11 > norm(result0 - result_)

state = Bra(basis_position, randdata1)
result_ = Bra(basis_momentum, copy(randdata2))
result0 = alpha*state*dense(Txp) + beta*result_
operators.gemv!(alpha, state, Txp, beta, result_)
@test 1e-11 > norm(result0 - result_)

state = Ket(basis_momentum, randdata1)
result_ = Ket(basis_position, copy(randdata2))
result0 = alpha*dense(Txp)*state + beta*result_
operators.gemv!(alpha, Txp, state, beta, result_)
@test 1e-11 > norm(result0 - result_)

state = Bra(basis_momentum, randdata1)
result_ = Bra(basis_position, copy(randdata2))
result0 = alpha*state*dense(Tpx) + beta*result_
operators.gemv!(alpha, state, Tpx, beta, result_)
@test 1e-11 > norm(result0 - result_)

# Test gemm!
rho0_xx = tensor(psi0_bx, dagger(psi0_bx))
rho0_xp = tensor(psi0_bx, dagger(psi0_bp))
rho0_px = tensor(psi0_bp, dagger(psi0_bx))
rho0_pp = tensor(psi0_bp, dagger(psi0_bp))

rho_ = DenseOperator(basis_momentum, basis_position)
operators.gemm!(Complex(1.), Tpx, rho0_xx, Complex(0.), rho_)
@test 1e-12 > D(rho_, rho0_px)
@test 1e-12 > D(Tpx*rho0_xx, rho0_px)

rho_ = DenseOperator(basis_position, basis_momentum)
operators.gemm!(Complex(1.), rho0_xx, Txp, Complex(0.), rho_)
@test 1e-12 > D(rho_, rho0_xp)
@test 1e-12 > D(rho0_xx*Txp, rho0_xp)

rho_ = DenseOperator(basis_momentum, basis_momentum)
operators.gemm!(Complex(1.), Tpx, rho0_xp, Complex(0.), rho_)
@test 1e-12 > D(rho_, rho0_pp)
@test 1e-12 > D(Tpx*rho0_xx*Txp, rho0_pp)

rho_ = DenseOperator(basis_momentum, basis_momentum)
operators.gemm!(Complex(1.), rho0_px, Txp, Complex(0.), rho_)
@test 1e-12 > D(rho_, rho0_pp)
@test 1e-12 > D(Txp*rho0_pp*Tpx, rho0_xx)

alpha = complex(3.2)
beta = complex(-1.2)
randdata1 = rand(ComplexF64, N, N)
randdata2 = rand(ComplexF64, N, N)

op = DenseOperator(basis_position, basis_position, randdata1)
result_ = DenseOperator(basis_momentum, basis_position, copy(randdata2))
result0 = alpha*dense(Tpx)*op + beta*result_
operators.gemm!(alpha, Tpx, op, beta, result_)
@test 1e-11 > D(result0, result_)

result_ = DenseOperator(basis_position, basis_momentum, copy(randdata2))
result0 = alpha*op*dense(Txp) + beta*result_
operators.gemm!(alpha, op, Txp, beta, result_)
@test 1e-11 > D(result0, result_)

op = DenseOperator(basis_momentum, basis_momentum, randdata1)
result_ = DenseOperator(basis_position, basis_momentum, copy(randdata2))
result0 = alpha*dense(Txp)*op + beta*result_
operators.gemm!(alpha, Txp, op, beta, result_)
@test 1e-11 > D(result0, result_)

result_ = DenseOperator(basis_momentum, basis_position, copy(randdata2))
result0 = alpha*op*dense(Tpx) + beta*result_
operators.gemm!(alpha, op, Tpx, beta, result_)
@test 1e-11 > D(result0, result_)

# Test FFT with lazy product
psi_ = deepcopy(psi0_bx)
operators.gemv!(Complex(1.), LazyProduct(Txp, Tpx), psi0_bx, Complex(0.), psi_)
@test 1e-12 > norm(psi_ - psi0_bx)
@test 1e-12 > norm(Txp*(Tpx*psi0_bx) - psi0_bx)

psi_ = deepcopy(psi0_bx)
id = dense(identityoperator(basis_momentum))
operators.gemv!(Complex(1.), LazyProduct(Txp, id, Tpx), psi0_bx, Complex(0.), psi_)
@test 1e-12 > norm(psi_ - psi0_bx)
@test 1e-12 > norm(Txp*id*(Tpx*psi0_bx) - psi0_bx)

# Test dense FFT operator
Txp_dense = DenseOperator(Txp)
Tpx_dense = DenseOperator(Tpx)
@test isa(Txp_dense, DenseOperator)
@test isa(Tpx_dense, DenseOperator)
@test 1e-5 > D(Txp_dense*rho0_pp*Tpx_dense, rho0_xx)

# Test FFT in 2D
N = [21, 18]
xmin = [-32.5, -10π]
xmax = [24.1, 9π]

basis_position = [PositionBasis(xmin[i], xmax[i], N[i]) for i=1:2]
basis_momentum = MomentumBasis.(basis_position)

x0 = [5.1, -0.2]
p0 = [-3.2, 1.33]
sigma = [1., 0.9]
sigma_x = sigma./sqrt(2)
sigma_p = 1.0/(sigma.*sqrt(2))

Txp = transform(tensor(basis_position...), tensor(basis_momentum...))
Tpx = transform(tensor(basis_momentum...), tensor(basis_position...))

Txp_sub = [transform(basis_position[i], basis_momentum[i]) for i=1:2]
Tpx_sub = dagger.(Txp_sub)
Txp_dense = dense.(Txp_sub)
Txp_comp = tensor(Txp_dense...)

difference = (dense(Txp) - Txp_comp).data
@test isapprox(difference, zero(difference); atol=1e-12)

psi0_x = [gaussianstate(basis_position[i], x0[i], p0[i], sigma_x[i]) for i=1:2]
psi0_p = [gaussianstate(basis_momentum[i], x0[i], p0[i], sigma_p[i]) for i=1:2]

psi_p_fft = Tpx*tensor(psi0_x...)
psi_p_fft2 = tensor((Tpx_sub.*psi0_x)...)
@test norm(psi_p_fft - psi_p_fft2) < 1e-15

psi_x_fft = Txp*tensor(psi0_p...)
psi_x_fft2 = tensor((Txp_sub.*psi0_p)...)
@test norm(psi_p_fft - psi_p_fft2) < 1e-15

psi_p_fft = dagger(tensor(psi0_x...))*Txp
psi_p_fft2 = tensor((dagger.(psi0_x).*Txp_sub)...)
@test norm(psi_p_fft - psi_p_fft2) < 1e-15

psi_x_fft = dagger(tensor(psi0_p...))*Tpx
psi_x_fft2 = tensor((dagger.(psi0_p).*Tpx_sub)...)
@test norm(psi_p_fft - psi_p_fft2) < 1e-15

difference = (dense(Txp) - identityoperator(DenseOperator, Txp.basis_l)*Txp).data
@test isapprox(difference, zero(difference); atol=1e-12)
@test_throws AssertionError transform(tensor(basis_position...), tensor(basis_position...))
@test_throws particle.IncompatibleBases transform(SpinBasis(1//2)^2, SpinBasis(1//2)^2)

@test dense(Txp) == dense(Txp_sub[1] ⊗ Txp_sub[2])

# Test ket only FFTs
Txp = transform(tensor(basis_position...), tensor(basis_momentum...); ket_only=true)
Tpx = transform(tensor(basis_momentum...), tensor(basis_position...); ket_only=true)

Txp_sub = [transform(basis_position[i], basis_momentum[i]; ket_only=true) for i=1:2]
Tpx_sub = dagger.(Txp_sub)

psi_p_fft = Tpx*tensor(psi0_x...)
psi_p_fft2 = tensor((Tpx_sub.*psi0_x)...)
@test norm(psi_p_fft - psi_p_fft2) < 1e-15

psi_x_fft = Txp*tensor(psi0_p...)
psi_x_fft2 = tensor((Txp_sub.*psi0_p)...)
@test norm(psi_p_fft - psi_p_fft2) < 1e-15

psi_p_fft = dagger(tensor(psi0_x...))*Txp
psi_p_fft2 = tensor((dagger.(psi0_x).*Txp_sub)...)
@test norm(psi_p_fft - psi_p_fft2) < 1e-15

psi_x_fft = dagger(tensor(psi0_p...))*Tpx
psi_x_fft2 = tensor((dagger.(psi0_p).*Tpx_sub)...)
@test norm(psi_p_fft - psi_p_fft2) < 1e-15

psi_x_fft = Txp*tensor(psi0_p...)
psi_x_fft2 = tensor(Txp_sub...)*tensor(psi0_p...)
@test norm(psi_x_fft - psi_x_fft2) < 1e-15

# Test composite basis of mixed type
bc = FockBasis(2)
psi_fock = fockstate(FockBasis(2), 1)
psi1 = tensor(psi0_p[1], psi_fock, psi0_p[2])
psi2 = tensor(psi0_x[1], psi_fock, psi0_x[2])

basis_l = tensor(basis_position[1], bc, basis_position[2])
basis_r = tensor(basis_momentum[1], bc, basis_momentum[2])
Txp = transform(basis_l, basis_r; ket_only=true)
Tpx = transform(basis_r, basis_l; ket_only=true)

psi1_fft = Txp*psi1
psi1_fft2 = tensor(Txp_sub[1]*psi0_p[1], psi_fock, Txp_sub[2]*psi0_p[2])
@test norm(psi1_fft - psi1_fft2) < 1e-15

psi2_fft = Tpx*psi2
psi2_fft2 = tensor(Tpx_sub[1]*psi0_x[1], psi_fock, Tpx_sub[2]*psi0_x[2])
@test norm(psi2_fft - psi2_fft2) < 1e-15

Txp = transform(basis_l, basis_r)
Txp_sub = [transform(basis_position[i], basis_momentum[i]) for i=1:2]
difference = (dense(Txp) - tensor(dense(Txp_sub[1]), dense(one(bc)), dense(Txp_sub[2]))).data
@test isapprox(difference, zero(difference); atol=1e-12)

basis_l = tensor(bc, basis_position[1], basis_position[2])
basis_r = tensor(bc, basis_momentum[1], basis_momentum[2])
Txp2 = transform(basis_l, basis_r)
Tpx2 = transform(basis_r, basis_l)
difference = (dense(Txp) - permutesystems(dense(Txp2), [2, 1, 3])).data
@test isapprox(difference, zero(difference); atol=1e-13)
difference = (dense(dagger(Txp)) - permutesystems(dense(Tpx2), [2, 1, 3])).data
@test isapprox(difference, zero(difference); atol=1e-13)

# Test potentialoperator in more than 1D
N = [21, 18]
xmin = [-32.5, -10π]
xmax = [24.1, 9π]

basis_position = [PositionBasis(xmin[i], xmax[i], N[i]) for i=1:2]
basis_momentum = MomentumBasis.(basis_position)

bcomp_pos = tensor(basis_position...)
bcomp_mom = tensor(basis_momentum...)
V(x, y) = sin(x*y) + cos(x)
xsample, ysample = samplepoints.(basis_position)
V_op = diagonaloperator(bcomp_pos, [V(x, y) for y in ysample for x in xsample])
V_op2 = potentialoperator(bcomp_pos, V)
@test V_op == V_op2

basis_position = PositionBasis.(basis_momentum)
bcomp_pos = tensor(basis_position...)
Txp = transform(bcomp_pos, bcomp_mom)
Tpx = transform(bcomp_mom, bcomp_pos)
xsample, ysample = samplepoints.(basis_position)
V_op = Tpx*dense(diagonaloperator(bcomp_pos, [V(x, y) for y in ysample for x in xsample]))*Txp
V_op2 = potentialoperator(bcomp_mom, V)
@test V_op == V_op2

N = [17, 12, 9]
xmin = [-32.5, -10π, -0.1]
xmax = [24.1, 9π, 22.0]

basis_position = [PositionBasis(xmin[i], xmax[i], N[i]) for i=1:3]
basis_momentum = MomentumBasis.(basis_position)

bcomp_pos = tensor(basis_position...)
bcomp_mom = tensor(basis_momentum...)
V(x, y, z) = exp(-z^2) + sin(x*y) + cos(x)
xsample, ysample, zsample = samplepoints.(basis_position)
V_op = diagonaloperator(bcomp_pos, [V(x, y, z) for z in zsample for y in ysample for x in xsample])
V_op2 = potentialoperator(bcomp_pos, V)
@test V_op == V_op2

# Test error messages
b1 = PositionBasis(-1, 1, 50)
b2 = MomentumBasis(-1, 1, 30)
@test_throws bases.IncompatibleBases transform(b1, b2)
@test_throws bases.IncompatibleBases transform(b2, b1)

bc1 = b1 ⊗ bc
bc2 = b2 ⊗ bc
@test_throws bases.IncompatibleBases transform(bc1, bc2)
@test_throws bases.IncompatibleBases transform(bc2, bc1)

b1 = PositionBasis(-1, 1, 50)
b2 = MomentumBasis(-1, 1, 50)
bc1 = b1 ⊗ bc
bc2 = b2 ⊗ bc
@test_throws bases.IncompatibleBases transform(bc1, bc2)
@test_throws bases.IncompatibleBases transform(bc2, bc1)
@test_throws bases.IncompatibleBases transform(bc1, bc2; index=[2])

bc1 = b1 ⊗ b2
bc2 = b1 ⊗ b2
@test_throws bases.IncompatibleBases transform(bc1, bc2)
@test_throws bases.IncompatibleBases transform(bc2, bc1)

@test_throws bases.IncompatibleBases potentialoperator(bc ⊗ bc, V)

end # testset


## ---

### Rate calculation used in mingpt
function cosRate( tokens, lr=6e-4, warmuptoken=512*20, len_train_dataset=1115394, blocksize=128)
    finalTok=2*len_train_dataset*blocksize
    progress = (tokens - warmuptoken) / max(1, finalTok - warmuptoken)
    lr_mult = max(0.1, 0.5 * (1.0 + cos(pi * progress)))
    lr_mult*lr
end

@show cosRate(60*512*128)
@show cosRate(125*512*128)
@show cosRate(235*512*128)
@show cosRate(350*512*128)


push!(LOAD_PATH, pwd()*"/../../src/")
using Distributions
using PyPlot
using BVB

# example 0
true_d = MixtureModel(Normal, [(5.0, 2), (10.0, 2), (20.0, 1), (-2, 1)], [1.0/3, 1.0/4, 1.0/4, 1.0/6])
# true_d = Gamma(2.0, 3.0)
ll(x) = logpdf(true_d, x)

init_approx_d =  Normal(0, 40)
x = linspace(-10,30,500)
max_iter = 50

figure()
subplot(3, 2, 1)
plot(x, exp(ll(x)), label="true")
plot(x, pdf(init_approx_d, x), label="init approx")
for (col, method) in enumerate(["inclusive", "exclusive"])
    approx_d, trace = BVB.infer(ll, init_approx_d, method=method, max_iter=max_iter,
                                eval_cross_entropy=true, eval_ELBO=true, show_plot=true)
    subplot(3, 2, 1)
    plot(x, pdf(approx_d, x), label=method); legend(loc=0)

    subplot(3, 2, 2)
    plot(trace["iter"], trace["cross_entropy"], ".-", label=method)
    xlabel("x"); ylabel("cross entropy"); legend(loc=0); grid("on")

    subplot(3, 2, 3)
    plot(trace["iter"], trace["ELBO"], ".-", label=method)
    xlabel("x"); ylabel("ELBO"); legend(loc=0); grid("on")

    subplot(3, 2, 4)
    vlines(collect(1:max_iter) + 0.5 * (col-1), 0., trace["α"], label=method, colors=(col==1 ? "blue" : "green"))
    xlabel("iter"); ylabel("alpha"); legend(loc=0); grid("on")

    subplot(3, 2, 5)
    plot(1:length(probs(approx_d)), sort(probs(approx_d), rev=true), ".-", label=method)
    xlabel("component"); ylabel("weight"); legend(loc=0); grid("on")

    subplot(3, 2, 6)
    plot(x, ll(x) - logpdf(approx_d,x), "-", label=method)
    axhline(y=0., ls="--", color="red")
    xlabel("x"); ylabel("log L - log p"); legend(loc=0); grid("on")
end

show()


# Basic structures

export SampleData,SampleArray,SampleVector

export duration
export stoi,itos
export channel,times

"""
    stoi(freq::Real,s::{Real,Second,Millisecond})

Convert `s` seconds to data index if data is sampled at `freq`"""
stoi(freq::K,s::T) where {K<:Real,T<:Real} = round(Int,s*freq)
stoi(freq::K,s::Second) where {K<:Real} = stoi(freq,Dates.value(s))
stoi(freq::K,s::Millisecond) where {K<:Real} = stoi(freq,Dates.value(s)/1000)
stoi(freq::K,s::AbstractArray{T,1}) where {K<:Real,T<:Union{Real,Second,Millisecond}} = [stoi(freq,ss) for ss in s]

"""
    itos(freq::Real,i::Integer)

Convert data index `i` to `Millisecond` if data is sampled at `freq`"""
itos(freq::K,i::T) where {K<:Real,T<:Integer} = Millisecond(round(Int,1000*i/freq))

"""
    itos(freq::Real,i::Integer,start::DateTime)
    itos(freq::Real,i::Vector{Integer},start::DateTime)

Convert data index `i` to `DateTime` if data is sampled at `freq` and started at `start`"""
itos(freq::K,i::T,start::DateTime) where {K<:Real,T<:Integer} = itos(freq,i) + start
itos(freq::K,i::AbstractArray{T,1},start::DateTime) where {K<:Real,T<:Integer} = [itos(freq,ii,start) for ii in i]

"""
    times(num::Int,freq::Real,start::DateTime=DateTime(0))

Create a DateTime vector the length `num`, assuming the data was collected at `freq` Hz and started at `start`"""
times(num::Int,freq::K,start::DateTime=DateTime(0)) where K<:Real = [itos(freq,i-1,start) for i in 1:num]

"Abstract type containing data sampled at a certain frequency, along with timing and metadata"
abstract type SampleData{T,K} end

"""
    SampleVector(data::AbstractArray{Number,1}, [times::AbstractArray{{Float,DateTime},1},] freq::Float64; name::String="", code::Symbol=:unknown, start::DateTime=DateTime(0))

Struct of a single `Vector` of sampled data.
* `data`: raw data points
* `times`: corresponding times of raw data points, either as a `Number` or `DateTime`.

   If omitted, it will be calculated as `DateTime` using the given `freq` starting at `start`
* `freq`: data frequency
* `name`: arbitrary name as `String` (default = `""`)
* `code`: arbitrary data type as `Symbol` (default = `:unknown`)

Data can be accessed using []-notation
"""
mutable struct SampleVector{T<:Number,K<:Union{AbstractFloat,DateTime}} <: SampleData{T,K}
    data::AbstractArray{T,1}
    times::AbstractArray{K,1}
    freq::Float64
    name::String
    code::Symbol

    function SampleVector(data::AbstractArray{T,1},times::AbstractArray{K,1},freq::F,name::String,code::Symbol) where {T<:Number,K<:Union{AbstractFloat,DateTime},F<:Real}
        if length(data) != length(times)
            error("""Trying to create a SampleVector with data length $(length(data)) and times length $(length(times))""")
        end
        return new{T,K}(data,times,Float64(freq),name,code)
    end
end

SampleVector(data::AbstractArray{T,1},times::AbstractArray{K,1},freq::F;name::String="",code::Symbol=:unknown) where {T<:Number,K<:Union{AbstractFloat,DateTime},F<:Real} = SampleVector(data,times,Float64(freq),name,code)
SampleVector(data::AbstractArray{T,1},freq::F;name::String="",code::Symbol=:unknown,start::DateTime=DateTime(0)) where {T<:Number,K<:Union{AbstractFloat,DateTime},F<:Real} = SampleVector(data,times(length(data),freq,start),Float64(freq),name,code)

"""
    SampleArray(data::AbstractArray{Number,2}, [times::AbstractArray{{Float,DateTime},1},] freq::Float64; names::Vector{String}=[], code::Symbol=:unknown, start::DateTime=DateTime(0))

Struct of an `Array` frequency data. Makes the assumption `data` was collected at same frequency, synced in time.

* `data`: raw data points; arranged as (time, channels), column-major
* `times`: corresponding times of raw data points, either as a `Number` or `DateTime`.

   If omitted, it will be calculated as `DateTime` using the given `freq` starting at `start`
* `freq`: data frequency
* `names`: `Vector{String}` of names (default = ["","" ...])
* `code`: arbitrary data type as `Symbol` (default = `:unknown`)

Data can be accessed by []-notation, or using the helper function `channel` to return a single `SampleVector`
"""
mutable struct SampleArray{T<:Number,K<:Union{AbstractFloat,DateTime}} <: SampleData{T,K}
    data::AbstractArray{T,2}
    times::AbstractArray{K,1}
    freq::Float64
    names::Vector{String}
    code::Symbol

    function SampleArray(data::AbstractArray{T,2},times::AbstractArray{K,1},freq::F,names::Vector{String},code::Symbol) where {T<:Number,K<:Union{AbstractFloat,DateTime},F<:Real}
        if size(data,1) != length(times)
            error("""Trying to create a SampleArray with data length $(size(data,1)) and times length $(length(times))
                    Data must be arranged in columns""")
        end
        if size(data,2) != length(names)
            error("""Trying to create a SampleArray with data size $(size(data)) and names length $(length(names))
                    Data must be arranged in columns""")
        end
        return new{T,K}(data,times,Float64(freq),names,code)
    end
end

function SampleArray(data::AbstractArray{T,2},times::AbstractArray{K,1},freq::F;names::Vector{String}=[],code::Symbol=:unknown) where {T<:Number,K<:Union{AbstractFloat,DateTime},F<:Real}
    if length(names)==0
        names = ["" for i=1:size(data,2)]
    end
    SampleArray(data,times,freq,names,code)
end

function SampleArray(data::AbstractArray{T,2},freq::F;names::Vector{String}=[],code::Symbol=:unknown,start::DateTime=DateTime(0)) where {T<:Number,K<:Union{AbstractFloat,DateTime},F<:Real}
    SampleArray(data,times(size(data,1),freq,start),freq,names=names,code=code)
end

"""
    SampleArray(v::SampleVector...)

Combine multiple `SampleVector`s into a single `SampleArray`. Is only valid for combining `SampleVector`s that have
the same `code` and `freq`, and that each contain the same number of samples.
"""
function SampleArray(v::SampleVector{T,K}...) where {T,K}
    codes = collect(Set([vv.code for vv in v]))
    freqs = collect(Set([vv.freq for vv in v]))
    length(codes)!=1 && error("Trying to combine SampleVectors containing different codes")
    length(freqs)!=1 && error("Trying to combine SampleVectors with different frequencies")

    lengths = collect(Set([length(vv.data) for vv in v]))
    length(lengths)!=1 && error("Trying to combine SampleVectors of different lengths (stacking together into a SampleArray requires the same length SampleVector's)")

    names = [vv.name for vv in v]
    data = Array{T,2}(undef,(lengths[1],length(names)))
    times = Array{K,2}(undef,(lengths[1],length(names)))

    for i=1:length(names)
        data[:,i] = v[i].data
        times[:,i] = v[i].times
    end

    SampleArray(data,times,freqs[1],names,codes[1])
end

import Base.length,Base.getindex,Base.lastindex,Base.==

itype = Union{Vector{Int},UnitRange,Colon}

length(v::SampleVector) = length(v.data)
length(a::SampleArray) = size(a.data,1)

lastindex(v::SampleData) = length(v)
lastindex(a::SampleArray,d::Int) = size(a.data,d)

function ==(a::SampleVector,b::SampleVector)
    ((a.code != b.code) || (a.freq != b.freq) || (a.name != b.name)) && return false
    return a.data == b.data
end

function ==(a::SampleArray,b::SampleArray)
    ((a.code != b.code) || (a.freq != b.freq) || (a.names != b.names)) && return false
    return a.data == b.data
end

"""length of `SampleData` in `Millisecond`s"""
duration(v::SampleData) = v.times[end] - v.times[1] + Millisecond(round(Int,1000/v.freq))

getindex(v::SampleVector,I::itype) = SampleVector(view(v.data,I),view(v.times,I),v.freq,v.name,v.code)
getindex(a::SampleArray,I::itype) = SampleArray(view(a.data,I,Colon()),view(a.times,I),a.freq,a.names,a.code)

getindex(v::SampleVector,I::StepRange) = SampleVector(view(v.data,I),view(v.times,I),v.freq/I.step,v.name,v.code)
getindex(a::SampleArray,I::StepRange) = SampleArray(view(a.data,I,Colon()),view(a.times,I),a.freq/I.step,a.names,a.code)

getindex(v::SampleVector,I::Int) = v.data[I]
getindex(a::SampleArray,I::Int) = SampleArray(view(a.data,I:I,Colon()),view(a.times,I:I),a.freq,a.names,a.code)

# If passing a Float or DateTime as the index (and the times are corresponding), return the first index at or after given time
getindex(a::SampleData{T,K},f::F) where {T,K<:AbstractFloat,F<:AbstractFloat} = getindex(a,searchsortedfirst(a.times,f))
getindex(a::SampleData{T,K},f::F) where {T,K<:DateTime,F<:DateTime} = getindex(a,searchsortedfirst(a.times,f))
getindex(a::SampleData{T,K},f::F) where {T,K<:DateTime,F<:TimePeriod} = getindex(a,searchsortedfirst(a.times,f+a.times[1]))

"""Extract a specific channel from `SampleArray` and returns a single `SampleVector`"""
channel(a::SampleArray,c::Int,I::itype) = SampleVector(view(a.data,I,c),view(a.times,I),a.freq,a.names[c],a.code)
channel(a::SampleArray,c::Int,I::Int) = a.data[I,c]
channel(a::SampleArray,c::Int) = channel(a,c,Colon())
channel(a::SampleArray,c::String,I::itype) = channel(a,findfirst(isequal(c),a.names),I)
channel(a::SampleArray,c::String,I::Int) = a.data[I,findfirst(isequal(c),a.names)]
channel(a::SampleArray,c::String) = channel(a,c,Colon())

"""Extract a set of channels from `SampleArray` and returns another `SampleArray`"""
channel(a::SampleArray,c::itype,I::itype) = SampleArray(view(a.data,I,c),view(a.times,I),a.freq,a.names[c],a.code)
channel(a::SampleArray,c::itype) = channel(a,c,Colon())
channel(a::SampleArray,c::Vector{String}) = channel(a,[findfirst(isequal(cc),a.names) for cc in c])


# See http://arxiv.org/pdf/1204.6448v3.pdf for more details
#
# The SteinMaternTensorizedKernel is given by
#
# k(x, y) = (1/(8 sigma^3)) * (1 + sigma * |x - y|) * exp{-sigma * |x-y|}
#
# where the norm is
#
# |u|_H^2 = int_R sigma^4 f(u)^2 + 2*sigma^2 f'(u)^2 + f''(u)^2 du

type SteinMaternTensorizedKernel <: SteinTensorizedKernel
    sigma::Float64
end

# default is use sigma value of 1
SteinMaternTensorizedKernel() = SteinMaternTensorizedKernel(1.0)

function getc(ker::SteinMaternTensorizedKernel)
    1 / (8.0 * ker.sigma^3)
end

function ki(ker::SteinMaternTensorizedKernel, x::Float64, y::Float64)
    c = getc(ker)
    d = abs(x-y)
    sigma = ker.sigma
    c * (1 + sigma * d) * exp(-sigma * d)
end

function gradxki(ker::SteinMaternTensorizedKernel, x::Float64, y::Float64)
    c = getc(ker)
    d = abs(x-y)
    sigma = ker.sigma

    c * sigma * (y - x) * exp(-sigma * d)
end

function gradxyki(ker::SteinMaternTensorizedKernel, x::Float64, y::Float64)
    c = getc(ker)
    d = abs(x-y)
    sigma = ker.sigma

    c * sigma * (1 - sigma * d) * exp(-sigma * d)
end


module grid
using LinearAlgebra
export Spacegrid,grid_n,grid_k,grid_xy_square
struct Spacegrid
    x::Array{Float64,2}
    y::Array{Float64,2}
end

function grid_n(Nx,Ny)
    #reciprocal space coordinate indices
    ny=[c  for r in -Nx:Nx, c in -Ny:Ny]
    nx=[r  for r in -Nx:Nx, c in -Ny:Ny]
    nx=vec(nx)
    ny=vec(ny)
    #difference matrix for 2dft
    dnx=[nx[a]-nx[b] for a in 1:length(nx),b in 1:length(nx)]
    dny=[ny[a]-ny[b] for a in 1:length(ny),b in 1:length(ny)]
    return nx,ny,dnx,dny
end

function grid_k(nx,ny,theta,alpha,lambda,ax,ay,epsilon_ref)
    #straightforward
    #all k vectors are generally normalized to k0 here
    k0=2*pi/real(lambda)
    #The incoming wave, transformed from spherical coordinates to normalized cartesian coordinates, |kin|=1
    kin=[sin(theta*pi/180)*cos(alpha*pi/180),sin(theta*pi/180)*sin(alpha*pi/180),cos(theta*pi/180)]*real(sqrt(epsilon_ref))
    #the spatial vectors are the sum of the incoming vector and the reciprocal lattice vectors
    kx=kin[1].+(2*pi/ax)*nx/k0;
    ky=kin[2].+(2*pi/ay)*ny/k0;
    #need matrix for later computations
    Kx=Diagonal(kx)
    Ky=Diagonal(ky)
    return k0,Kx,Ky,kin
end

function grid_xy_square(acc)
    #creating a square meshgrid
    x=[r  for r in -acc/2+.5:acc/2-.5, c in -acc/2+.5:acc/2-.5]/acc
    y=[c  for r in -acc/2+.5:acc/2-.5, c in -acc/2+.5:acc/2-.5]/acc
    return Spacegrid(x,y)
end

end