######### Stan program example  ###########

using AxisKeys, DataFrames
using StanSample
using Test

binorm_model = "
  transformed data {
      matrix[2,2] Sigma;
      vector[2] mu;

      mu[1] <- 0.0;
      mu[2] <- 0.0;
      Sigma[1,1] <- 1.0;
      Sigma[2,2] <- 1.0;
      Sigma[1,2] <- 0.10;
      Sigma[2,1] <- 0.10;
  }
  parameters {
      vector[2] y;
  }
  model {
        y ~ multi_normal(mu,Sigma);
  }
  "

sm = SampleModel("binormal", binorm_model);

rc = stan_sample(sm)

if success(rc)
  samples = read_samples(sm)

  # Fetch cmdstan summary data frame
  df = read_summary(sm)
  @test df[df.parameters .== Symbol("y[1]"), :mean][1] ≈ 0.0 atol=2.0
  @test df[df.parameters .== Symbol("y[2]"), :mean][1] ≈ 0.0 atol=2.0

end


include("algorithm.jl")

graphNodes = 3
#              A  B   C   D    E
adjMatrix = [ -1  500 400 500 800;
              500 -1  200 200 300;
              400 200 -1  300 500;
              500 200 300 -1  200;
              800 300 500 200 -1  ]
foodNode = 3
maxIteration = 200
antsCounts = 20
pheromoneExponent = 1
desirabilityExponent = 0.1
pheromoneDepositFactor = 1
evaporationRate = 0.1

graph = generateGraph(graphNodes,adjMatrix,foodNode)

bestPath, bestFitness, graph = algorithmAntColonyOptimization(graph, maxIteration, antsCounts, pheromoneExponent,
                                            desirabilityExponent, pheromoneDepositFactor, evaporationRate)

print("Best solution $(bestPath) with length $(bestFitness)")


using EconPDEs, Distributions

struct HaddadModel
  # consumption process parameters
  μbar::Float64
  vbar::Float64
  κμ::Float64
  νμ::Float64
  κv::Float64
  νv::Float64

  # active capital
  αbar::Float64
  λ::Float64

  # utility parameters
  ρ::Float64
  γ::Float64
  ψ::Float64
end

function HaddadModel(;μbar = 0.018, vbar = 0.00052, κμ = 0.3, νμ = 0.456, κv = 0.012, νv = 0.00472, αbar = 1.2, λ = 0.018, ρ = 0.0132, γ = 10.0, ψ = 1.5)
  HaddadModel(μbar, vbar, κμ, νμ, κv, νv, αbar, λ, ρ, γ, ψ)
end

function initialize_stategrid(m::HaddadModel; μn = 30, vn = 30)
  μbar = m.μbar; vbar = m.vbar ; κμ = m.κμ ; νμ = m.νμ ; κv = m.κv ; νv = m.νv ; αbar = m.αbar ; λ = m.λ ; ρ = m.ρ ; γ = m.γ ; ψ = m.ψ

  σ = sqrt(νμ^2 * vbar / (2 * κμ))
  μmin = quantile(Normal(μbar, σ), 0.001)
  μmax = quantile(Normal(μbar, σ), 0.999)
  μs = range(μmin, stop = μmax, length = μn)

  α = 2 * κv * vbar / νv^2
  β = νv^2 / (2 * κv)
  vmin = quantile(Gamma(α, β), 0.001)
  vmax = quantile(Gamma(α, β), 0.999)
  vs = range(vmin, stop = vmax, length = vn)

  OrderedDict(:μ => μs, :v => vs)
end

function initialize_y(m::HaddadModel, stategrid::OrderedDict)
  OrderedDict(:p => ones(length(stategrid[:μ]), length(stategrid[:v])))
end

function (m::HaddadModel)(state::NamedTuple, y::NamedTuple)
  μbar = m.μbar ; vbar = m.vbar ; κμ = m.κμ ; νμ = m.νμ ; κv = m.κv ; νv = m.νv ; αbar = m.αbar ; λ = m.λ ; ρ = m.ρ ; γ = m.γ ; ψ = m.ψ
  μ, v = state.μ, state.v
  p, pμ, pv, pμμ, pμv, pvv = y.p, y.pμ, y.pv, y.pμμ, y.pμv, y.pvv

  # drift and volatility of μ, ν, p
  μc = μ
  σc = sqrt(v)
  μμ = κμ * (μbar - μ)
  σμ = νμ * sqrt(v)
  μv = κv * (vbar - v)
  σv = νv * sqrt(v)
  σp_Zμ = pμ / p * σμ
  σp_Zv = pv / p * σv
  σp2 = σp_Zμ^2 + σp_Zv^2
  μp = pμ / p * μμ + pv / p * μv + 0.5 * pμμ / p * σμ^2 + 0.5 * pvv / p * σv^2

  # Market Price of Risk κ
  αstar = αbar - sqrt(λ * αbar / (0.5 * γ * (σc^2 + σp_Zμ^2 + σp_Zv^2)))
  αstar = clamp(αstar, 0.0, 1.0)
  κ_Zc = αstar * γ * σc
  κ_Zμ = (αstar - (1 / γ - 1) / (ψ - 1)) * γ * σp_Zμ
  κ_Zv = (αstar - (1 / γ - 1) / (ψ - 1)) * γ * σp_Zv
  κ2 = κ_Zc^2 + κ_Zμ^2 + κ_Zv^2

  # Interest rate r
  #r = ρ + μc / ψ - (1 - 1 / ψ) / (2 * γ) * κ2 - (1 / ψ - γ) / (2 * γ * (ψ - 1)) * σp2 - (1 - γ) / (γ * ψ) * (κμ * σμ + κv * σv) + 1 / ψ  * (κ_Zc * σc + κμ * σμ + κv * σv)

  # Market pricing
  #out = p * (1 / p + μc + μp - r - κ_Zc * σc - κ_Zμ * σp_Zμ - κ_Zv * σp_Zv)

  out = p * (1 / p - ρ + (1 - 1 / ψ) * (μc - 0.5 * γ * σc^2 * (1 - (αstar - 1)^2)) + μp + (0.5 * (1 / ψ - γ) / (1 - 1 / ψ) + 0.5 * γ * (1 - 1 / ψ) * (αstar - 1)^2) * σp2)

  return (out,) , (μμ, μv), (p = p, κ_Zc = κ_Zc, κ_Zμ = κ_Zμ, κ_Zv = κ_Zv, αstar = αstar)
end

m = HaddadModel()
stategrid = initialize_stategrid(m)
y0 = initialize_y(m, stategrid)
y, result, distance = pdesolve(m, stategrid, y0)


# We have n particles here.
# For each particle i (i ∈ N⁺)  :
# We have Mₘᵢₙ ≤ mᵢ ≤ Mₘₐₓ
# We have an three dimentional vector v⃗ᵢ for i's velocity.
# And an three dimentional point Pᵢ for i's locationn.
# We have Gravitational Constant G ≈ 6.67408*10⁻¹¹ (m³kg⁻¹s⁻²) .
# And a small time step Δt→0 (actually, it's a very small value).
# At the beginning, every particle will spread in an area randomly.
# At the beginning, every particle will have a velocity which direction and size are all given randomly but in a range.
#

# Modules
using LinearAlgebra

# Constants
n = 1000
G = 6.67408e-11
Δt = 60*12
Δtₕ = Δt/2
Mₘₐₓ = 1e22
Mₘᵢₙ = 1e20
Areaₘᵢₙ = [-1e11 -1e11 -1e11]
Areaₘₐₓ = [1e11 1e11 1e11]
Areaₜₑₘₚ = Areaₘₐₓ - Areaₘᵢₙ
v⃗ₘᵢₙ  = [-5e2 -5e2 -5e2]
v⃗ₘₐₓ  = [5e2 5e2 5e2]
v⃗ₜₑₘₚ  = v⃗ₘₐₓ  - v⃗ₘᵢₙ
SavePer = 5
Total = 1024*16

# Generate
MGΔt = (rand(n)*(Mₘₐₓ-Mₘᵢₙ) + ones(n)*Mₘᵢₙ)*G*Δt
P₀ = zeros(3, n)
P₁ = zeros(3, n)
v⃗  = zeros(3, n)
for i=1:3
  P₀[i, :] = rand(n)*Areaₜₑₘₚ[i] + ones(n)*Areaₘᵢₙ[i]
  v⃗[i,  :] = rand(n)*v⃗ₜₑₘₚ[i]  + ones(n)*v⃗ₘᵢₙ[i]
end

#Set an blackhole
P₀[:, 1] = [0.0 0.0 0.0]
v⃗[:, 1] = [0.0 0.0 0.0]
MGΔt[1] = 1e30*G*Δt

# Functions
#rsqrt(x) = 1/sqrt(x);
function update(P₀, P₁, v⃗)
  for i=1:n
    Δta⃗  = [0.0, 0.0, 0.0]
    for j=1:(i-1)
      Δta⃗  += MGΔt[j]/(norm(P₀[:, j] - P₀[:, i])^3)*(P₀[:, j] - P₀[:, i])
    end
    for j=(i+1):n
      Δta⃗  += MGΔt[j]/(norm(P₀[:, j] - P₀[:, i])^3)*(P₀[:, j] - P₀[:, i])
    end
    P₁[:, i] = P₀[:, i] + v⃗[:, i]*Δt + Δta⃗ *Δtₕ
    v⃗[:, i] = v⃗[:, i] + Δta⃗
  end
end

# Run
for j=0:Total-1
  @time if mod(j, 2)==0
    update(P₀, P₁, v⃗)
  else
    update(P₁, P₀, v⃗)
  end
  if mod(j, SavePer)==0
    println(P₀[:, 1])
  end
end


using Distributions, NLopt, ChaosCommunications

carrier = LogisticCarrier(4)

noise_var = ebn0db2var(5., system=:UTURUniChannelCSK, carrier=carrier)

println("    Testing correlation decoder of coherent single-user CSK...")

system = UTURUniChannelCSK(true, carrier, noise=Normal(0., sqrt(noise_var)))

sim_ber(system, 1, 100000)
# psim_ber(system, 1, 100000)

println("    Testing correlation decoder of non-coherent single-user CSK...")

system = UTURUniChannelCSK(false, carrier, noise=Normal(0., sqrt(noise_var)))

sim_ber(system, 1, 100000)

println("    Testing MCML decoder of non-coherent single user-CSK...")

decoder = mcml_decoder(system=:UTURUniChannelCSK, carrier=carrier)

system = UTURUniChannelCSK(false, carrier, noise=Normal(0., sqrt(noise_var)), decoder=decoder)

sim_ber(system, 1, 1000)


module HexFFT

import Base: fft, ifft

export hfft2, ihfft2, fft, ifft, OffsetHexData

# HELPER FUNCTIONS
pad_zeros(arr) = hcat(arr, zeros(arr))
decimate(data, offset) = view(data, :, offset + 1:2:size(data, 2))
div2(n::Int) = n >> 1
function fold_half(data, op)
    n = size(data, 1)
    op.(view(data, 1:div2(n), :), view(data, div2(n) + 1:n, :))
end

# FORWARD TRANSFORM
function dft_nst1(data)
    out = fft(pad_zeros(data), 2)
    return decimate(out, 0), decimate(out, 1)
end

nst2(data) = repmat(fft(fold_half(data, (+)), 1), 2, 1)

coeff_inner(i, n) = exp(-2pi * im * i / n)
nst3_coeffs(n, m) = reshape(coeff_inner.(0:div2(n) - 1, n), div2(n), 1)
nst3(data) = repmat(fft(nst3_coeffs(size(data)...) .*
                        fold_half(data, (-)), 1), 2, 1)

w_inner(b, s, d, n, m) = exp(-pi * im * ((b + 2d) / 2m + (b + 2s) / n))
w(b, n, m) = [w_inner(b, s, d, n, m) for s in 0:n - 1, d in 0:m - 1]

function hfft2(data0, data1)
    g00, g01 = dft_nst1(data0)
    g10, g11 = dft_nst1(data1)
    return nst2(g00) .+ w(0, size(g10)...) .* nst2(g10),
           nst3(g01) .+ w(1, size(g11)...) .* nst3(g11)
end

# INVERSE TRANSFORM
function idft_inst1(data)
    out = 2ifft(pad_zeros(data), 2)
    return decimate(out, 0), decimate(out, 1)
end

inst2(data) = repmat(0.5ifft(fold_half(data, (+)), 1), 2, 1)

icoeff_inner(i, n) = exp(2pi * im * i / n)
inst3_coeffs(n, m) = reshape(icoeff_inner.(0:div2(n) - 1, n), div2(n), 1)
inst3(data) = repmat(0.5ifft(inst3_coeffs(size(data)...) .*
                             fold_half(data, (-)), 1), 2, 1)

iw_inner(a, r, c, n, m) = exp(pi * im * ((a + 2c) / 2m + (a + 2r) / n))
iw(a, n, m) = [iw_inner(a, r, c, n, m) for r in 0:n - 1, c in 0:m - 1]

function ihfft2(data0, data1)
    g00, g01 = idft_inst1(data0)
    g10, g11 = idft_inst1(data1)
    return 0.5(inst2(g00) .+ iw(0, size(g10)...) .* inst2(g10)),
           0.5(inst3(g01) .+ iw(1, size(g11)...) .* inst3(g11))
end

# Hex Array type
immutable OffsetHexData{T, S <: AbstractArray{T, 2}}
    odd_rows::S
    even_rows::S
end
function OffsetHexData{S}(odd_rows::S, even_rows::S)
    @assert size(odd_rows) == size(even_rows)
    OffsetHexData{eltype(odd_rows), S}(odd_rows, even_rows)
end

fft(H::OffsetHexData) = OffsetHexData(hfft2(H.odd_rows, H.even_rows)...)
ifft(H::OffsetHexData) = OffsetHexData(ihfft2(H.odd_rows, H.even_rows)...)

end # module


using StateSpaceModels, Distributions, LinearAlgebra, Plots, Random

# Fix seed
Random.seed!(1)

# Number of observations
n = 400
# State dimension (2d position + 2d speed)
m = 4
# Measurements dimension (2d noisy position)
p = 2
# Control dimension (2d acceleration)
q = 2

# Damping ratio
ρ = 0.05
# Time delta
Δ = 1.0

# State transition matrix
T = kron(Matrix{Float64}(I, p, p), [1.0 (1.0 - ρ * Δ / 2.0) * Δ; 0.0 (1.0 - ρ * Δ)])
# Input matrix
R = kron(Matrix{Float64}(I, p, p), [0.5 * Δ^2; Δ])
# Output (measurement) matrix
Z = kron(Matrix{Float64}(I, p, p), [1.0 0.0])

# Generate random actuators
Q = 0.5 * Matrix{Float64}(I, q, q)
η = MvNormal(Q)

# Generate random measurement noise
H = 2.0 * Matrix{Float64}(I, p, p)
ε = MvNormal(H)

# Simulate vehicle trajectory
α = zeros(n + 1, m)
y = zeros(n, p)
for t in 1:n
    y[t, :] = Z * α[t, :] + rand(ε)
    α[t + 1, :] = T * α[t, :] + R * rand(η)
end
α = α[1:n, :]

# User defined model
model = StateSpaceModel(y, Z, T, R)
# Estimate vehicle speed and position
ss = statespace(model)

ss.model.H
ss.model.Q

anim = @animate for i in 1:n
    plot(y[1:i, 1], y[1:i, 2], label="Measured position", line=:scatter, lw=2, markeralpha=0.2, color=:black, title="Vehicle tracking")
    plot!(α[1:i, 1], α[1:i, 3], label="True position", lw=3, color=:indianred)
    plot!(ss.filter.a[2:i+1, 1], ss.filter.a[2:i+1, 3], label="Estimated position", lw=2, color=:forestgreen)
end

gif(anim, "vehicle_tracking.gif", fps = 15)


###########################################################################
# Least-squares
###########################################################################
function bearing_ls(angles, gains, ref_angles, ref_gains)
	@assert length(angles) == length(gains)
	num_gains = length(gains)
	best_err = Inf
	best_shift = 0
	for shift = 0:359
		err = 0.0
		for j = 1:num_gains
			if validgain(gains[j])
				ref_gain = interp_linear(angles[j], shift, ref_angles, ref_gains)
				err += (gains[j] - ref_gain)^2
			end
		end
		if err < best_err
			best_err = err
			best_shift = shift
		end
	end
	return 360 - best_shift
end

# Uses least-square method
# samples is how many samples you want to use (maybe you dont want whole rotation)
# sense is direction of rotation: -1 for negative sense, 1 for positive sense
function bearing_ls(dir_file::AbstractString, omni_file::AbstractString, ref_file::AbstractString; samples::Int64=0, sense::Int64=1)
	test_gp = GainPattern(dir_file) - GainPattern(omni_file) + 5
	ref_gp = GainPattern(ref_file)

	num_samples = length(test_gp.angles)
	num_exp = length(test_gp.samples[1])

	# Allowing user to specify count
	if (samples > 0) && (samples < num_samples)
		num_samples = samples
	end

	aoa_array = zeros(num_exp)
	angles = zeros(num_samples)
	gains = zeros(num_samples)
	for i = 1:num_exp
		if sense > 0
			for j = 1:num_samples
				# Loop through all samples in this experiment...
				angles[j] = test_gp.angles[j]
				gains[j] = test_gp.samples[j][i]
			end
		elseif sense == -1
			# Backwards loop, starting from last measurement
			for j = 1:num_samples
				angles[j] = test_gp.angles[end - j + 1]
				gains[j] = test_gp.samples[end - j + 1][i]
			end
		elseif sense == -2
			# Backwards loop, starting from origin
			for j = 2:num_samples
				angles[j] = test_gp.angles[end - j + 2]
				gains[j] = test_gp.samples[end - j + 2][i]
			end
			angles[1] = test_gp.angles[1]
			gains[1] = test_gp.samples[1][i]
		end
		aoa_array[i] = bearing_ls(angles, gains, ref_gp.angles, ref_gp.meangains)
	end
	return aoa_array
end


using HMMs
using Random
using Brobdingnag

import HMMs: GenTDist

Random.seed!(1)
(mu1,sigma1,nu1) = HMMs.randomparameters( GenTDist )
t1 = GenTDist( mu1, sigma1, nu1 )
y1 = rand( t1, 1_000_000 );
w1 = ones(length(y1));

parameters = HMMs.randomparameters( GenTDist )
HMMs.fit_mle!( GenTDist, parameters, y1, w1, Dict{Symbol,Any}() )
@assert( maximum(abs.(parameters - [mu1,sigma1,nu1])) < 1e-2 )

(mu2,sigma2,nu2) = HMMs.randomparameters( GenTDist )
t2 = GenTDist( mu2, sigma2, nu2 )
y2 = fill(NaN,1_000_000);
w2 = fill(NaN,length(y2));
for i = 1:length(y2)
    w2[i] = rand()
    y2[i] = rand() < w2[i] ? rand( t1 ) : rand( t2 )
end

parameters = HMMs.randomparameters( GenTDist )
HMMs.fit_mle!( GenTDist, parameters, y2, w2, Dict{Symbol,Any}() )
# this doesn't seem to work but I'm not sure it's well-founded
#@assert( maximum(abs.(parameters - [mu1,sigma1,nu1])) < 1e-2 )

graph = HMMs.fullyconnected(1)
hmm1 = HMMs.randomhmm( graph, dist=GenTDist, calc=Brob, seed=1 )
y1 = rand( hmm1, 10_000 );

hmm2 = HMMs.randomhmm( graph, dist=GenTDist, calc=Brob, seed=2 )
HMMs.setobservations( hmm2, y1 );
HMMs.em( hmm2, debug=2 )

graph = HMMs.fullyconnected(2)
hmm1 = HMMs.randomhmm( graph, dist=GenTDist, calc=Brob, seed=1 )
y1 = rand( hmm1, 100_000 );

hmm2 = HMMs.randomhmm( graph, dist=GenTDist, calc=Brob, seed=2 )
HMMs.setobservations( hmm2, y1 );
HMMs.em( hmm2, debug=2 )

hmm1 = HMMs.randomhmm( graph, dist=GenTDist, calc=Brob, seed=1 )
y1 = rand( hmm1, 10_000 );

hmm2 = HMMs.randomhmm( graph, dist=GenTDist, calc=Brob, seed=2 )
HMMs.setobservations( hmm2, y1 );
HMMs.em( hmm2, debug=2 )

hmm2 = HMMs.randomhmm( graph, dist=GenTDist, calc=Brob, seed=3 )
HMMs.setobservations( hmm2, y1 );
HMMs.em( hmm2, debug=2 )


function parseInput(input)
    chars = split(input, ",")
    chars = parse.(Int, chars)
    return chars
end

function part1(codes)
    i = 1
    while codes[i] != 99
        operation = codes[i]
        pos1 = codes[i+1] + 1
        pos2 = codes[i+2] + 1
        dest = codes[i+3] + 1
        if operation == 1
            codes[dest] = codes[pos1] + codes[pos2]
        elseif operation == 2
            codes[dest] = codes[pos1] * codes[pos2]
        end
        i += 4
    end
    return codes
end

function run(s)
    input = parseInput(split(s, "\n")[1])
    input[2] = 12
    input[3] = 2
    return part1(input)[1]
end

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

main()


function ∂ end

function ∂(::GridType{:Primal},grid::Grid{Dim},dir::Direction{D}, F::AbstractArray) where {Dim,D}
    shifts = -1 .* (ntuple(identity,Val(Dim)) .== D)
	if D > Dim return zeros(CFloat, size(grid)) end
    (circshift(F,shifts) - F) / spacing(grid,dir)
end

function ∂(::GridType{:Dual},grid::Grid{Dim},dir::Direction{D}, F::AbstractArray) where {Dim,D}
    shifts = 1 .* (ntuple(identity,Val(Dim)) .== D)
	 if D > Dim return zeros(CFloat, size(grid)) end
    (circshift(F,shifts) - F) / spacing(grid,dir)
end


using BrainFlow

data = Array{Float64,3}(undef, 2, 4, 2)
data[:, :, 1] = [[6,3] [3,0] [1,5] [5,1]]
data[:, :, 2] = [[1,5] [5,1] [6,2] [2,2]]
labels = [0.0, 1.0]

filters, eigvals = BrainFlow.get_csp(data, labels)

println("filters = ")
println(filters)


include("testsystem1.jl")
srand(11)
Y, X = sample(20, 100, M)
kalmanrts(Y, M)


for order in range(0, length=20)
    for degree in range(0, length=order+1)
        print(binomial(order, degree), ' ')
    end
    print('\n')
end


using ForwardDiff

function check_derivatives(g, gx, gy)
    p = rand(2)
    display(ForwardDiff.jacobian(g, p))
    display([gx(p) gy(p)])
    return ForwardDiff.jacobian(g, p) ≈ [gx(p) gy(p)]
end

function check_second_derivatives(f, fx, fy, fxx, fxy, fyy)
    p = rand(2)
    display([fxx(p) fxy(p) fyy(p)])
    display(ForwardDiff.jacobian(fx, p))
    display(ForwardDiff.jacobian(fy, p))
end


module Final

import Evolve

L = 2                       # width of the box
xmin = -1                   # leftmost point in the box
N = 20000                   # number of subintervals
ϵ = L/N                     # width of each subinterval
k = 20                      # wavenumbers
σ = 0.2                     # wavepacket width
x0 = 0.0                    # wavepacket starting position

V(xs) = map(x->0.0, xs)     # zero potential energy
AV = Evolve.A(V)            # evolution matrix constructor for square well
ϕV = Evolve.Φ(k, σ, x0)     # initial state vector constructor for Gaussian packet

x = (L/N .* [0:N] + xmin)   # x positions

end


@i function ⊖(exp)(out!::GVar, x::GVar{T}) where T
    value(out!) -= exp(value(x))
    @anc anc1 = zero(T)
    anc1 += exp(value(x))
    grad(x) += grad(out!) * anc1
    anc1 -= exp(value(x))
end

@i function ⊖(log)(out!::GVar, x::GVar{T}) where T
    value(out!) -= log(value(x))
    grad(x) += grad(out!) / value(x)
end

@i function ⊖(sin)(out!::GVar, x::GVar{T}) where T
    value(out!) -= sin(value(x))
    @anc anc1 = zero(T)
    anc1 += cos(value(x))
    grad(x) += grad(out!) * anc1
    anc1 -= cos(value(x))
end

@i function ⊖(cos)(out!::GVar, x::GVar{T}) where T
    value(out!) -= cos(value(x))
    @anc anc1 = zero(T)
    anc1 -= sin(value(x))
    grad(x) += grad(out!) * anc1
    anc1 += sin(value(x))
end

for op in [:exp, :log, :sin, :cos]
    @eval @nograd ⊖($op)(out!, x::GVar)
    @eval @nograd ⊖($op)(out!::GVar, x)
end

@i function IROT(a!::GVar, b!::GVar, θ::GVar)
    IROT(value(a!), value(b!), value(θ))
    NEG(value(θ))
    value(θ) ⊖ π/2
    ROT(grad(a!), grad(b!), value(θ))
    grad(θ) += value(a!) * grad(a!)
    grad(θ) += value(b!) * grad(b!)
    value(θ) ⊕ π/2
    NEG(value(θ))
    ROT(grad(a!), grad(b!), π/2)
end

@i function IROT(a!::GVar, b!::GVar, θ)
    IROT(value(a!), value(b!), θ)
    NEG(θ)
    θ ⊖ π/2
    ROT(grad(a!), grad(b!), θ)
    θ ⊕ π/2
    NEG(θ)
    ROT(grad(a!), grad(b!), π/2)
end

@nograd IROT(a!, b!, θ::GVar)


using Test

function parseInput(input)
    chars = split(input, ",")
    chars = parse.(Int, chars)
    return chars
end

function part1(codes)
    i = 1
    while codes[i] != 99
        operation = codes[i]
        pos1 = codes[i+1] + 1
        pos2 = codes[i+2] + 1
        dest = codes[i+3] + 1
        # pos1, pos2, dest = codes[i+1:i+3]
        if operation == 1
            codes[dest] = codes[pos1] + codes[pos2]
        elseif operation == 2
            codes[dest] = codes[pos1] * codes[pos2]
        end
        i += 4
    end
    return codes
end

function run(s)
    for i in 1:99
        for j in 1:99
            input = parseInput(split(s, "\n")[1])
            input[2] = i
            input[3] = j
            output = part1(input)
            if output[1] == 19690720
                return 100 * i + j
            end
        end
    end
end

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

main()


# use SGD for linear regression

using SGDOptim

function linreg_sgd(θ_g::Vector{Float64}, n::Int, σ::Float64)

    # prepare experimental data
    d = length(θ_g) - 1
    X = randn(d, n)
    y = vec(θ_g[1:d]'X) + θ_g[d+1] + σ * randn(n)

    # initialize solution
    θ_0 = zeros(d + 1)

    # optimize
    rmodel = riskmodel(AffinePred(d), SqrLoss())

    θ = sgd(rmodel, θ_0,
        minibatch_seq(X, y, 10),          # configure the way data are supplied
        reg = SqrL2Reg(1.0e-4),           # regularization
        lrate = t->1.0 / (100.0 + t),     # learing rate policy
        cbinterval = 5,                   # how frequently callback is invoked
        callback = gtcompare_trace(θ_g))  # the callback function

    # compare solution with initial guess
    println()
    @printf("Initial:  deviation = %.4e | avg.risk = %.4e\n",
        vecnorm(θ_0 - θ_g), value(rmodel, θ_0, X, y) / n)
    @printf("Solution: deviation = %.4e | avg.risk = %.4e\n",
        vecnorm(θ - θ_g), value(rmodel, θ, X, y) / n)
end

linreg_sgd([3.0, 5.0, 2.0], 10000, 0.1)


using StochasticAiry
using Test
using Random: seed!

@testset "StochasticAiry.jl" begin
    seed!(1)

    n_re_lambda = 1000
    n_im_lambda = 200

    re_grid = collect(range(-8, 15, length=n_re_lambda))
    im_grid = collect(range(-8, 8, length=n_im_lambda))

    lambda_grid = re_grid' .+ im*im_grid;

    ##

    t_grid_len = 3*10^3

    t = -20
    s = 10
    t_grid = range(s, t, length=t_grid_len);

    dt = (t - s)/t_grid_len;

    for beta in [1, Inf]
        phi, phi_dash= rand(SAi(beta), :space, 0, lambda_grid)

        @test size(phi) == size(lambda_grid)
        @test size(phi_dash) == size(lambda_grid)

        phi_re, phi_dash_re = rand(SAi(beta), :space, 0, -8:8, brownian_gridlen = 100)

        @test eltype(phi_re) <: Real
        @test eltype(phi_dash_re) <: Real

        sai, sai_dash = rand(SAi(beta), :time, 1 + im, t_grid)

        @test ndims(sai) == 1
        @test ndims(sai_dash) == 1

        sai_re, sai_dash_re = rand(SAi(beta), :time, 0, t_grid)

        @test eltype(sai_re) <: Real
        @test eltype(sai_dash_re) <: Real
    end

end


function kkrebook2(omega, imchi, alpha)
    if size(omega,1)>size(omega,2)
        omega = omega'
    end
    if size(imchi,1)>size(imchi,2)
        imchi=imchi'
    end
    g = size(omega,2)
    rechi = zeros(size(imchi))
    a = zeros(size(imchi))
    b = zeros(size(imchi))
    deltaomega = omega[2]-omega[1]
    j = 1
    beta1 = 0

    for k=2:g
        b[1] = beta1+imchi[k]*omega[k]^(2*alpha+1)/(omega[k]^2-omega[1]^2)
        beta1 = b[1]
    end
    rechi[1] = 2/pi*deltaomega*b[1]*omega[1]^(-2*alpha)

    j=g
    alpha1 = 0
    for k=1:g-1
        a[g] = alpha1+imchi[k]*omega[k]^(2*alpha+1)/(omega[k]^2-omega[g]^2)
        alpha1 = a[g]
    end
    rechi[g] = 2/pi*deltaomega*a[g]*omega[g]^(-2*alpha)

    for j=2:g-1
        alpha1 = 0
        beta1 = 0
        for k=1:j-1
            a[j] = alpha1+imchi[k]*omega[k]^(2*alpha+1)/(omega[k]^2-omega[j]^2)
            alpha1 = a[j]
        end
        for k=j+1:g
            b[j] = beta1+imchi[k]*omega[k]^(2*alpha+1)/(omega[k]^2-omega[j]^2)
            beta1 = b[j]
        end
        rechi[j] = 2/pi*deltaomega*(a[j]+b[j])*omega[j]^(-2*alpha)
        print("\r$j")
    end
    return rechi
end

imchi = readdlm(raw"C:\Users\ms246\Dropbox (Heriot-Watt University Team)\RES_EPS_Lupo\Projects\Mohammed\phd\ozone absorption spectrum\imchi_0.124nm.txt")
omega = readdlm(raw"C:\Users\ms246\Dropbox (Heriot-Watt University Team)\RES_EPS_Lupo\Projects\Mohammed\phd\ozone absorption spectrum\freqdata_0.124nm.txt")
kk = kkrebook2(omega[2:end], imchi, 0)

dir = raw"C:\Users\ms246\Dropbox (Heriot-Watt University Team)\RES_EPS_Lupo\Projects\Mohammed\phd\ozone absorption spectrum\\"

open(dir*"julia_kk.txt", "w") do io
    writedlm(io, kk)
end

open(dir*"freqJulia.txt", "w") do io
    writedlm(io, omega)
end


function stratification(r,Nmax,rbottom,iopt)

# defines stratification on radial grid.

# r       - radial grid points
# Nmax    - stratification at CMB (r=1)
# Rbottom - radius of base of stratified layer
# iopt =  - constant stratification (=0); linear stratification (=1)

#   constructs density stratification
   n = length(r);
   N2 = zeros(n);

   N2_peak = Nmax * Nmax;

#  stratified layer
   if (iopt == 1)
       N2[1:n] .= (r[1:n] .- r[1]) * N2_peak / (r[n] - r[1]);
   end

   if (iopt == 0 )
       N2[1:n] .= N2_peak;
   end

   return N2

end


# from the classical Annals paper..

# To do this properly with KernelDensity.jl later,
# also figure out generic interface

struct BrownGreenshtein
    m::Int64
    bandwidth::Float64
end

BrownGreenshtein(m) = BrownGreenshtein(m, 1/sqrt(log(m)))

function estimate(Xs, ::Type{BrownGreenshtein}, target::PosteriorTarget{PosteriorMeanNumerator})
    bg = BrownGreenshtein(length(Xs))
    estimate(Xs, bg, target)
end

function estimate(Xs, bg::BrownGreenshtein, target::PosteriorTarget{PosteriorMeanNumerator})
    x = target.num.x
    h = bg.bandwidth
    Z = Normal(0, 1)
    denom = mean( pdf.(Z, (Xs.-x)./h))/h
    num = mean( (Xs.-x).*pdf.(Z, (Xs.-x)./h) )/(h^3)
    x + num/denom
end


using MinimaxCalibratedEBayes
using RCall
using Plots
using StatsPlots
using LaTeXStrings
using ExponentialFamilies
using MosekTools
using Random

const MCEB = MinimaxCalibratedEBayes
pgfplotsx()

# # Prostate dáta analysis

# ## Preliminaries

# Let us make some analysis choices from the beginning

# ### Empirical Bayes Targets
# Our goal is to estimate the following two targets:

target_grid = -3:0.2:3
post_means = PosteriorMean.(StandardNormalSample.(target_grid))
lfsrs =  LFSR.(StandardNormalSample.(target_grid))

# ### Empirical Bayes Prior Class
# For the prior class we consider a mixture of Gaussians
# that have standard deviation 0.2 and with mixing measure
# supported on ... Note that we also specify the solver
# to be used for solving the convex programming problems;
# here we use Mosek.
gcal = GaussianMixturePriorClass(0.2, -3:0.02:3, Mosek.Optimizer)

# ### Discretization of marginal observations
marginal_grid = -6:0.02:6

# ## Load the data
prostz_file = MinimaxCalibratedEBayes.prostate_data_file
R"load($prostz_file)"
@rget prostz;
Zs = StandardNormalSample.(prostz)

# We will implement the preprocessing steps manually here.
# Afterwards, we will use wrapper that automate all of these steps.

# First we split our data into two folds
Random.seed!(200)
Zs_train, Zs_test, _, _ = MCEB._split_data(Zs)

# let us quickly check the range of values in the folds
extrema(response(Zs_train)), extrema(response(Zs_test))

# Now we use the first fold to estimate an $L_\infty$ neighborhood
# of the marginal density. We also fit a logspline model of the
# prior density to get our pilot estimators.

# ### Neighborhood construction:

infinity_band_options = KDEInfinityBandOptions(a_min=-6.0, a_max=6.0)
fkde_train = fit(infinity_band_options, response.(Zs_train))

# Let us visualize the result by showing a histogram of the datapoints
# the estimated density function and the $L_\infty$ bands.
prostate_marginal_plot = histogram(response.(Zs_train), bins=50, normalize=true,
          label="", alpha=0.6, fillcolor=:lightgray)
plot!(fkde_train, label="Kernel estimate \n  (and bands)",
      legend=:topright)

# ### G-model fitting

# We choose the following class
exp_family = ContinuousExponentialFamilyModel(Uniform(-3.6,3.6),
                         collect(-3.6:0.1:3.6); df=5, scale=true)
exp_family_solver = MCEB.ExponentialFamilyDeconvolutionMLE(cefm = exp_family, c0=0.001)

# Let us fit to the training data
exp_family_fit = fit(exp_family_solver, Zs_train)

# Let us look at the pilot estimates we got
lfsr_logspline_estimate = estimate.(lfsrs , exp_family_fit)
post_mean_logspline_estimate = estimate.(post_means , exp_family_fit)

expfamily_fits_plot = plot(plot(target_grid, lfsr_logspline_estimate),
                           plot(target_grid, post_mean_logspline_estimate))

# discretize our test samples to implement minimax calibration methodology
Zs_test_discr = DiscretizedStandardNormalSamples(Zs_test, marginal_grid)
Zs_test_discr = set_neighborhood(Zs_test_discr, fkde_train)

# add all of the above in one struct
mceb_setup = MinimaxCalibratorSetup(Zs_train = Zs_train,
                                    Zs_test = Zs_test,
						            prior_class = gcal,
						            fkde_train = fkde_train,
						            Zs_test_discr = Zs_test_discr,
						            pilot_method = exp_family_fit,
									delta_tuner =  MCEB.HalfCIWidth)#MCEB.RMSE)

post_means_fits = fit.(mceb_setup, post_means)
lfsrs_fits = fit.(mceb_setup, lfsrs)

gr()
prostate_postmeans_plot = plot(post_means_fits, ylims=(-2.5,2.5))
prostate_lfsr_plot = plot(lfsrs_fits, label=nothing)
plot(prostate_lfsr_plot, prostate_postmeans_plot, layout=(1,2))

# # The impact of Neighborhoods: Moving to opportunity

# We next seek to repeat the analysis in the second dataset. Here we will also
# see how most of the steps manually conducted for the prostate dataset
# can be automated.

# ## Loading the dataset
# Let us start by loading the dataset and filtering out missing observations.

using CSV
nbhood_file = MinimaxCalibratedEBayes.chetty_hendren_file
nbhoods_df = CSV.read(nbhood_file, types=Dict(3=>String));
missing_dx = [!ismissing(x) for x in nbhoods_df[!,:p25_coef]]
β_hat = [nbhoods_df[i,:p25_coef] for i in findall(missing_dx)]
se_hat = [nbhoods_df[i,:p25_se] for i in findall(missing_dx)]
nbhood_Zs = StandardNormalSample.(β_hat./se_hat);
length(nbhood_Zs)

# ## Setup the minimax calibrator
# The above steps repeat all the setup we conducted for the prostate data
# to compute the `mceb_setup` object.
Random.seed!(34)
mceb_options = MinimaxCalibratorOptions(prior_class = gcal,
                               marginal_grid = marginal_grid,
							   pilot_options = exp_family_solver)
nbhood_mceb_setup = fit(mceb_options, nbhood_Zs)

extrema(response.(nbhood_mceb_setup.Zs_train)),extrema(response.(nbhood_mceb_setup.Zs_test))

GR.settextfontprec(232, 3)
gr()
nbhood_post_means_fits = fit.(nbhood_mceb_setup, post_means)
nbhood_lfsrs_fits = fit.(nbhood_mceb_setup, lfsrs)

prostate_postmeans_plot = plot(post_means_fits, ylims=(-3.2,3.2),
                            fg_legend=:transparent, bg_legend=:transparent,
							legend = (0.35,0.84))
prostate_lfsrs_plot = plot(lfsrs_fits, label=nothing)

nbhood_postmeans_plot = plot(nbhood_post_means_fits,  ylims=(-3.2,3.2), label=nothing)
nbhood_lfsrs_plot = plot(nbhood_lfsrs_fits,  label=nothing)

prostate_marginal_plot = histogram(response.(Zs_train), bins=20, normalize=true,
          label="", alpha=0.4, fillcolor=:lightgray)
plot!(prostate_marginal_plot, fkde_train, fillalpha=0.45, #L"\bar{f}(x)\pm c_m",
      linewidth=1, label=nothing, #L"\bar{f}(x)\pm c_m",
	  legend=:topleft,  bg_legend=:transparent,
	  fg_legend=:transparent)

nbhood_marginal_plot = histogram(response.(nbhood_mceb_setup.Zs_train), bins=20, normalize=true,
	            label=nothing, alpha=0.4, fillcolor=:lightgray)
plot!(nbhood_marginal_plot, nbhood_mceb_setup.fkde_train, fillalpha=0.45,
                  label=nothing, linewidth=1)

plot()
using Plots.Measures
dataset_plot = plot(prostate_marginal_plot, prostate_postmeans_plot, prostate_lfsrs_plot,
	 nbhood_marginal_plot, nbhood_postmeans_plot, nbhood_lfsrs_plot,
	 title = ["a) Prostate" "b)" "c)" " d) Neighborhoods" "e)" "f)"],
	 layout= (2,3), size=(750,500))
dataset_plot

savefig(dataset_plot, "dataset_plots.pdf")
#savefig(dataset_plot, "mceb_datasets.pdf")


function powerlevel(x, y)
    serial = 7511
    rackid = x + 10
    power = rackid * y
    power += serial
    power *= rackid
    power = convert(Int, floor(power / 100)) % 10
    return power - 5
end

function foo()
    max_power = -100000000
    max_x = 0
    max_y = 0
    max_size = 0

    data = zeros(Int, 300, 300)

    for x in 1:300
        for y in 1:300
            data[x, y] = powerlevel(x, y)
        end
    end

    for x in 1:300
        for y in 1:300
            for size in 0:(300-max(x,y))
                power = 0
                for i in 0:size
                    for j in 0:size
                        power += data[x + i, y + j]
                    end
                end
                if power > max_power
                    max_power = power
                    max_x = x
                    max_y = y
                    max_size = size + 1
                end
            end
        end
    end

    return max_x, max_y, max_size
end

z = foo()
println(z[1])
println(z[2])
println(z[3])


struct FourBlockMatrix <: AbstractArray{Float64,2}
  A::BandedMatrix
  B::SparseMatrixCSC
  C::SparseMatrixCSC
  D::Matrix

  n::Int
  m::Int

  l::Int
  u::Int
end

function FourBlockMatrix(n::Int,m::Int,l::Int,u::Int)

  A = BandedMatrix(Zeros(n,n), (l,u))

  B = spzeros(n,m)
  C = spzeros(m,n)

  D = zeros(m,m)

  cur_matrix = FourBlockMatrix(
    A, B, C, D,
    n, m, l, u
  )

end


abstract type LinearProcess{T} <: ContinuousTimeProcess{T}
end

b(s, x, P::LinearProcess) = B(s, P)*x + β(s, P)
σ(s, x, P::LinearProcess) = σ(s, P)
a(s, P::LinearProcess) = outer(σ(s, P))

symmetrize(A) = (A + A')/2
#####################

mutable struct Ptilde{T} <: ContinuousTimeProcess{T}
    cs::CSpline{T}
    σ
    a
    Γ
    Ptilde{T}(cs, σ) where T = new(cs, σ, σ*σ', inv(σ*σ'))
end
b(t::Float64, x, P::Ptilde) = P.cs(t)

B(t::Float64, P::Ptilde) = 0.0
β(t::Float64, P::Ptilde) = P.cs(t)

mu(s, x, t, P::Ptilde) = x + integrate(P.cs, s, t)
σ(t, x, P::Ptilde) = P.σ
a(t, x, P::Ptilde) = P.a
a(t, P::Ptilde) = P.a
gamma(t, x, P::Ptilde) = P.Γ
constdiff(::Ptilde) = true

"""
    Ptilde(cs::CSpline, σ)

Affine diffusion ``dX = cs(t) dt + σdW``
with cs a cubic spline `::CSpline`.

"""
Ptilde(cs::CSpline{T}, σ) where {T} = Ptilde{T}(cs, σ)

function lp(s, x, t, y, P::Ptilde{T}) where T
    logpdfnormal(y - mu(s,x,t,P), (t-s)*P.a)
end

transitionprob(s,x,t,P::Ptilde) = Gaussian(mu(s,x,t,P), (t-s)*P.a)

function V(t, T, v, P::Ptilde)
    v - integrate(P.cs, t, T)
end

function dotV(t, T, v, P::Ptilde)
     P.cs(t)
end

function H(t, T, P::Ptilde)
    P.Γ/(T-t)
end
function H(t, T, P::Ptilde, x)
    P.Γ*x/(T-t)
end

mutable struct LinPro{S,T,U} <: ContinuousTimeProcess{T}
    B::S
    μ::T
    σ::U
    a::U
    Γ::U
    lambda::S # stationary covariance
    function LinPro(B::S, μ::T, σ::U) where {S,T,U}
        a = σ*σ'
        return new{S,T,U}(B, μ, σ, a, inv(a), symmetrize(lyap(B, a)))
    end
end

B(t, P::LinPro) = P.B
β(t, P::LinPro) = -P.B*P.μ
b(t::Float64, x, P::LinPro) = P.B*(x .- P.μ)
σ(t, x, P::LinPro) = P.σ
bderiv(t, x, P::LinPro) = P.B
σderiv(t, x, P::LinPro) = 0*(P.σ)

a(t, x, P::LinPro) = P.a
a(t, P::LinPro) = P.a
constdiff(::LinPro) = true

"""
    LinPro(B, μ::T, σ)

Linear diffusion ``dX = B(X - μ)dt + σdW``.

"""
LinPro

function lp(s, x, t, y, P::LinPro{T}) where T
    logpdfnormal(y - mu(s,x,t,P), K(s, t, P::LinPro))
end

transitionprob(s,x,t,P::LinPro) = Gaussian(mu(s,x,t,P), K(s, t, P::LinPro))

function Phi(t, T, P::LinPro)
    exp((T-t)*P.B)
end

function mu(t, x, T, P::LinPro)
    phi = exp((T-t)*P.B)
    phi*(x - P.μ) + P.μ
end

function K(t, T, P::LinPro)
    phi = exp((T-t)*P.B)
    P.lambda - phi*P.lambda*phi'
end

function H(t, T, P::LinPro, x)
     phim = exp(-(T-t)*P.B)
     (phim*P.lambda*phim'-P.lambda)\x
end
function H(t, T, P::LinPro)
     phim = exp(-(T-t)*P.B)
     inv(phim*P.lambda*phim'-P.lambda)
end

function V(t, T, v, P::LinPro)
    phim = exp(-(T-t)*P.B)
    phim*(v - P.μ) + P.μ
end

function dotV(t, T, v, P::LinPro)
    exp(-(T-t)*P.B)*P.B*(v - P.μ)
end

"""
    LinProBridge

Bridge process of `P::LinPro` with `μ == 0` conditional on ending in `v` at time `t`.
"""
struct LinProBridge{T,S<:LinPro} <: ContinuousTimeProcess{T}
    t::Float64  # end time
    v::T        # end point
    P::S
    LinProBridge(t, v::T, P::S) where {T,S<:LinPro} = !iszero(P.μ) ? throw(ArgumentError("μ ≠ 0")) : new{T,S}(t,v,P)
end

b(s, x, P::LinProBridge) = P.P.B * (x - P.P.μ) + P.P.a * H(s, P.t, P.P, V(s, P.t, P.v, P.P) - x)
B(s, P::LinProBridge) = P.P.B - P.P.a * H(s, P.t, P.P)

σ(t, x, P::LinProBridge) = P.P.σ
bderiv(t, x, P::LinProBridge) = B(t, P)
σderiv(t, x, P::LinProBridge) = 0*(P.P.σ)

a(t, x, P::LinProBridge) = P.P.a
a(t, P::LinProBridge) = P.P.a
constdiff(::LinProBridge) = true

Phi(t, P::LinProBridge) = P.P.lambda*exp((P.t - t)*P.P.B)' - exp(-(P.t - t)*P.P.B)*P.P.lambda

Phi(t, T, Ps::LinProBridge) = Phi(T, Ps)*inv(Phi(t, Ps))

mu(s, x, t, P::LinProBridge{Float64}) = inv(sinh(P.P.B*(P.t - s)))*( P.v*sinh(P.P.B*(t - s)) + x*sinh(P.P.B*(P.t - t)) )
K(s, t, P::LinProBridge{Float64}) = 2 * P.P.lambda * csch(-P.P.B*(P.t - s)) * sinh(-P.P.B*(t-s))*sinh(-P.P.B*(P.t - t))
# Wolfram alpha: integral_(1/e)^ gamma (sinh^2(-(π - gamma )))/(sinh^2(-(π - s))) ds = sinh(1/e - gamma ) sinh( gamma - π) (-csch(1/e - π))≈0.170692

#function mu(s, x, t, Ps::LinProBridge)
#    T = Ps.t
#    P = Ps.P
#    inv(P.a - expm(P.B*(T-s))*P.a*expm(P.B*(T-s))')*(P.a - expm(P.B*(T-t))*P.a*expm(P.B*(T-t))')*x
#end

transitionprob(s, x, t, P::LinProBridge{Float64}) = Gaussian(mu(s, x, t, P), K(s, t, P))

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

"""
    LinearAppr(tt, B, β, a)
"""
struct LinearAppr{R,S,T} <: ContinuousTimeProcess{T}
    tt::Vector{Float64}
    xx::Vector{T}
    B::Vector{S}
    b::Vector{T}
    Σ::Vector{R}
end

_b((i,s)::IndexedTime, x, P::LinearAppr) = P.B[i]*(x - P.xx[i]) + P.b[i]
B((i,s)::IndexedTime, P::LinearAppr) = P.B[i]
β((i,s)::IndexedTime, P::LinearAppr) = P.b[i] - P.B[i]*P.xx[i]
a((i,s)::IndexedTime, x, P::LinearAppr) = outer(P.Σ[i])
a((i,s)::IndexedTime, P::LinearAppr) = outer(P.Σ[i])
constdiff(::LinearAppr) = false

linearappr(Y, P) = LinearAppr(Y.tt, Y.yy, map((t,x) -> Bridge.bderiv(t, x, P), Y.tt, Y.yy), map((t,x) -> Bridge.b(t, x, P), Y.tt, Y.yy), map((t,x) -> Bridge.σ(t, x, P), Y.tt, Y.yy))
function linearappr!(Pt::LinearAppr, Y, P)
    Pt.tt[:] = Y.tt
    Pt.xx[:] = Y.yy
    Pt.B[:] = map((t,x) -> Bridge.bderiv(t, x, P), Y.tt, Y.yy)
    Pt.b[:] = map((t,x) -> Bridge.b(t, x, P), Y.tt, Y.yy)
    Pt.Σ[:] = map((t,x) -> Bridge.σ(t, x, P), Y.tt, Y.yy)
    Pt
end


#
#  The following functions allow us to more slickly do probmods (not including any Gen functions!)
#
#	flip(p=0.5)
# 	flip(output1::Any, output2::Any, p::Float64)
# 	inverse_sample(ys::Vector{Float64}; consideration_interval=[0.,1.])
# 	get_beta_params(desired_mode)#mode, spread
#
#  NOT INCLUDING ANY GEN FILES
#
using Distributions

function logmeanexp(scores::Vector{Float64})
    logsumexp(scores) - log(length(scores))
end;

function flip(p::Float64=0.5)
    rand() < p;
end;

# @gen function flip(output1::Any, output2::Any, p::Float64)
#     if ({:flip} ~ bernoulli(p))
#         return output1
#     else
#         return output2
#     end
# end;

function inverse_sample(ys::Vector{Float64}; consideration_interval=[0.,1.])
    ys_idxs_sort = sortperm(ys)
    ys_cdf_fx = StatsBase.ecdf(ys)
    ys_cdf = ys_cdf_fx(ys[ys_idxs_sort])
    xi = uniform(consideration_interval[1],consideration_interval[2])
    yi = findfirst(y -> y >= xi, ys_cdf)
    y0 = findfirst(y -> y >= consideration_interval[1], ys_cdf)
    y1 = findfirst(y -> y >= consideration_interval[2], ys_cdf)
    consideration_interval_std = StatsBase.std(ys[ys_idxs_sort[y0:y1]])
    selected_ys_index = ys_idxs_sort[yi]
    return (selected_ys_index, consideration_interval_std)
end

function get_beta_params(desired_mode)#mode, spread
    m = desired_mode
    a = 10
    b = (a*(1-m) + 2*m - 1)/m
    return (a, b)
end


using CUDNN: CD, cudnnConvolutionForward, CUDNN_CONVOLUTION, CUDNN_CROSS_CORRELATION, cudnnGetConvolutionNdDescriptor

function cudnnConvolutionForward(x,w; cd=nothing)
    nd, pads, strides, upscale, mode, dtype = cd != nothing ? cudnnGetConvolutionNdDescriptor(cd) : (2, (0,0), (1,1), (1,1), CUDNN_CONVOLUTION, Float32)
    pad, stride = pads[1], strides[1]
end

# outputDim = 1 + (inputDim + 2*pad - filterDim)/convolutionStride
###
# cudnnConvolutionForward(x, w, y;
#    padding=c.padding, stride=c.stride, upscale=c.upscale, mode=c.mode,
#                            algorithm=c.algorithm, workSpace=c.workSpace, workSpaceSizeInBytes=c.workSpaceSizeInBytes,
#                                                        alpha=c.alpha, beta=c.beta)
###
function kudnnConvolutionForward(x,w; cd=nothing)
    # x: (W,H,C,N)
    # w: (W,H,C,K)
    # y: (W,H,K,N)
    nd, pads, strides, upscale, mode, dtype = cd != nothing ? cudnnGetConvolutionNdDescriptor(cd) : (2, (0,0), (1,1), (1,1), CUDNN_CONVOLUTION, Float32)
    pad, stride = pads[1], strides[1]
    Wx,Hx,Cx,N = size(x)
    Ww,Hw,Cw,K = size(w)
    @assert Cx==Cw
    @assert mode == CUDNN_CONVOLUTION
    Wy,Hy = floor(Int, 1 + (Int[Wx,Hx] + 2*pad - Int[Ww,Hw]) / stride)
    # Wy,Hy = 1 + div(Int[Wx,Hx] + 2*pad - Int[Ww,Hw], stride)
    y = zeros(dtype, (Wy,Hy,K,N))
    for n in 1:N, k in 1:K, c in 1:Cx
        y[:,:,k,n] += _conv2(x[:,:,c,n], w[:,:,c,k]; pad=pad, stride=stride, xcorr=CUDNN_CONVOLUTION!=mode)
    end
    return y
end

# dw = rot180(xcorr(x,dy))
function kudnnConvolutionBackwardFilter(x,dy,dw; cd=nothing)
    # x:    (Wx,Hx,Cx,N)
    # dy:   (Wy,Hy,K,N)
    # dw:    (Ww,Hw,Cw,K)
    nd, pads, strides, upscale, mode, dtype = cd != nothing ? cudnnGetConvolutionNdDescriptor(cd) : (2, (0,0), (1,1), (1,1), CUDNN_CONVOLUTION, Float32)
    pad, stride = pads[1], strides[1]
    @assert pad==0 && stride==1
    @assert mode == CUDNN_CONVOLUTION
    Wx,Hx,C,Nx = size(x)
    Wy,Hy,K,Ny = size(dy)
    dw2 = zeros(dw)
    for c in 1:C, k in 1:K, n in 1:Ny
        dw2[:,:,c,k] += rot180(_conv2(x[:,:,c,n], dy[:,:,k,n]; pad=pad, stride=stride, xcorr=true))
    end
    return dw2
end

# dx = xcorr(dy, w, 'full')
function kudnnConvolutionBackwardData(dy,w,dx; cd=nothing) # TODO change the order of parameters
    nd, pads, strides, upscale, mode, dtype = cd != nothing ? cudnnGetConvolutionNdDescriptor(cd) : (2, (0,0), (1,1), (1,1), CUDNN_CONVOLUTION, Float32)
    pad, stride = pads[1], strides[1]
    @assert pad==0 && stride==1
    @assert mode == CUDNN_CONVOLUTION
    Wy,Hy,Ky,N = size(dy)
    Ww,Hw,C,Kw = size(w)
    @assert Ky==Kw
    dx2 = zeros(dx)
    for n in 1:N, c in 1:C, k in 1:Kw
        t = conv2(dy[:,:,k,n], rot180(w[:,:,c,k]))
        # t = _conv2(dy[:,:,k,n], w[:,:,c,k]; pad=Ww-1, stride=stride, xcorr=true)
        dx2[:,:,c,n] += t
    end
    return dx2
end


const Detection2D = Tuple{Float64, Float64, Int64}

"""
    get_line_segment(camera_params::Camera_Params, params::Video_Params, detection::Detection2D)

Returns a line segment of the possible 3D positions of the object that could have caused a 2D detection
"""
function get_line_segment(camera_params::Camera_Params, params::Video_Params, detection::Detection2D)

    x = detection[1]
    y = detection[2]

    #now (0,0) is the center of the image
    x = x - params.image_dim_x/2
    y = y - params.image_dim_y/2

    angle_from_vertical = x / (params.image_dim_x/2) * deg2rad(params.horizontal_FoV/2)
    angle_from_horizontal = -1 * y / (params.image_dim_y/2) * deg2rad(params.vertical_FoV/2)

    #println("angle_from_vertical ", rad2deg(angle_from_vertical))
    #println("angle_from_horizontal ", rad2deg(angle_from_horizontal))

    #something wrong

    (a_vertical, b_vertical, c_vertical) = get_vertical_plane(camera_params)
    # println((a_vertical, b_vertical, c_vertical))
    (a_horizontal, b_horizontal, c_horizontal) = get_horizontal_plane(camera_params, a_vertical, b_vertical, c_vertical)
    # println((a_horizontal, b_horizontal, c_horizontal))
    # normal_to_image = (camera_params.camera_focus.x - camera_params.camera_location.x,
    # camera_params.camera_focus.y - camera_params.camera_location.y,
    # camera_params.camera_focus.z - camera_params.camera_location.z)
    # #get one more point on the ray. don't care which point, but could be on same plane as camera's focus
    # l = get_distance(camera_params.camera_location, camera_params.camera_focus)
    # println("l ", l)
    # w = 2*l*tan(deg2rad(params.horizontal_FoV/2)) #total image width in meters (non-pixel space)
    # #x = angle_from_vertical / (deg2rad(params.horizontal_FoV/2)) * w / 2
    # x = x * (w / 2) / (params.image_dim_x/2)
    # h = 2*l*tan(deg2rad(params.vertical_FoV/2)) #total image heigth in meters (non-pixel space)
    # #y = angle_from_horizontal / (deg2rad(params.vertical_FoV/2)) * h / 2
    # y = -1 * y * (h / 2) / (params.image_dim_y/2)

    l = get_distance(camera_params.camera_location, camera_params.camera_focus)
    x = l*tan(angle_from_vertical)
    y = l*tan(angle_from_horizontal)

    #println("x ", x)
    #println("y ", y)

    normalized_v = [a_vertical, b_vertical, c_vertical]./norm([a_vertical, b_vertical, c_vertical])
    normalized_h = [a_horizontal, b_horizontal, c_horizontal]./norm([a_horizontal, b_horizontal, c_horizontal])
    #println("normalized_v ", normalized_v)
    #println("normalized_h ", normalized_h)
    focus = [camera_params.camera_focus.x, camera_params.camera_focus.y, camera_params.camera_focus.z]
    #println("focus ", focus)
    #println("dot that should be 0 ", dot(x .* normalized_v + y .* normalized_h, normal_to_image))
    point = focus + x .* normalized_v + y .* normalized_h
    point = Coordinate(point[1], point[2], point[3])

    #println("point ", point)

    #check each wall for an intersection within limits
    (endpoint_x, endpoint_y, endpoint_z, a, b, c) = check_walls(point, camera_params.camera_location, params)

    #return Line_Segment(camera_params.camera_location, Coordinate(endpoint_x, endpoint_y, endpoint_z), a, b, c)
    #Flip line segment to LHR coordinate system
    l = Line_Segment(Coordinate(camera_params.camera_location.x, camera_params.camera_location.y, -camera_params.camera_location.z),
        Coordinate(endpoint_x, endpoint_y, -endpoint_z), a, b, -c)
    return l
end

function get_distance(p1::Coordinate, p2::Coordinate)
    return sqrt((p1.x-p2.x)^2 + (p1.y-p2.y)^2 + (p1.z-p2.z)^2)
end

"""
    check_walls(point::Coordinate, camera::Coordinate, params::Video_Params)

Calculates parameters for the line along which the object could appear, alongside
the coordinates corresponding to the farthest the object could be from the
camera while still being within the boundary of the scene.

Arguments:
- point: possible 3d location of the object
- camera: location of the camera
- params:: Video_Params

Returns:
- (x, y, z): farthest the location from the camera the object could be while
still being within the boundary of the scene.
- (a, b, c): parametrization of the line along which the object could appear.
when t = 1, it gives the farthest point (x, y, z)
"""
function check_walls(point::Coordinate, camera::Coordinate, params::Video_Params)
    a = point.x - camera.x
    b = point.y - camera.y
    c = point.z - camera.z

    #get intersection between line and plane

    #check x = x_min
    x = params.x_min
    t = (x - camera.x)/a
    y = camera.y + b*t
    z = camera.z + c*t
    if y < params.y_max && y > params.y_min && z < params.z_max && z > params.z_min && t > 0 #t>0 checks right side
        #redoing a,b,c so when t=1, at this fartest point
        a = x - camera.x
        b = y - camera.y
        c = z - camera.z
        return (x, y, z, a, b, c)
    end

    #check x = x_max
    x = params.x_max
    t = (x - camera.x)/a
    y = camera.y + b*t
    z = camera.z + c*t
    if y < params.y_max && y > params.y_min && z < params.z_max && z > params.z_min && t > 0
        a = x - camera.x
        b = y - camera.y
        c = z - camera.z
        return (x, y, z, a, b, c)
    end

    #check y = y_min
    y = params.y_min
    t = (y - camera.y)/b
    x = camera.x + a*t
    z = camera.z + c*t
    if x < params.x_max && x > params.x_min && z < params.z_max && z > params.z_min && t > 0
        a = x - camera.x
        b = y - camera.y
        c = z - camera.z
        return (x, y, z, a, b, c)
    end

    #check y = y_max
    y = params.y_max
    t = (y - camera.y)/b
    x = camera.x + a*t
    z = camera.z + c*t
    if x < params.x_max && x > params.x_min && z < params.z_max && z > params.z_min && t > 0
        a = x - camera.x
        b = y - camera.y
        c = z - camera.z
        return (x, y, z, a, b, c)
    end

    #check z = z_min
    z = params.z_min
    t = (z - camera.z)/c
    x = camera.x + a*t
    y = camera.y + b*t
    if x < params.x_max && x > params.x_min && y < params.y_max && y > params.y_min && t > 0
        a = x - camera.x
        b = y - camera.y
        c = z - camera.z
        return (x, y, z, a, b, c)
    end

    #check z = z_max
    z = params.z_max
    t = (z - camera.z)/c
    x = camera.x + a*t
    y = camera.y + b*t
    if x < params.x_max && x > params.x_min && y < params.y_max && y > params.y_min && t > 0
        a = x - camera.x
        b = y - camera.y
        c = z - camera.z
        return (x, y, z, a, b, c)
    end
end

export get_line_segment


# Online Outlier Detection
"""
    Euclidean_distance!{tp}(d::Array{tp, 1}, x::AbstractArray{tp, 2}, i::Int, j::Int, dim::Int = 1)

compute the Euclidean distance between x[i,:] and x[j,:] and write the result to d. Memory efficient. dim is the dimension of i and j.
"""
function Euclidean_distance!(d::Array{tp, 1}, x::AbstractArray{tp, 2}, i::Int, j::Int, dim::Int = 1) where {tp}
  d[1] = 0.0
  if dim == 2
    for v = 1:size(x, 1)
      d[1] = d[1] + (x[v, i] - x[v, j]) * (x[v, i] - x[v, j])
    end
  else
    for v = 1:size(x, 2)
      d[1] = d[1] + (x[i, v] - x[j, v]) * (x[i, v] - x[j, v])
    end
  end
  d[1] = sqrt(d[1])
  return d
end

"""
    Mahalanobis_distance!{tp}(d::Array{tp, 1}, x::AbstractArray{tp, 2}, Q::AbstractArray{tp, 2}, i::Int, j::Int, dim::Int = 1)

compute the Euclidean distance between x[i,:] and x[j,:] and write the result to d. Memory efficient. dim is the dimension of i and j.
"""
function Mahalanobis_distance!(d::Array{tp, 1}, x::AbstractArray{tp, 2}, Q::AbstractArray{tp, 2},i::Int, j::Int, dim::Int = 1) where {tp}
         d[1] = 0.0
         if dim == 2
         for v2 = 1:size(x, 1)
           for v = 1:size(x, 1)
             d[1] = d[1] + (x[v, i] - x[v, j]) * Q[v, v2] *  (x[v2, i] - x[v2, j])
           end
         end
         else
         for v = 1:size(x, 2)
           for v2 = 1:size(x, 2)
             d[1] = d[1] + (x[i, v] - x[j, v]) * Q[v, v2] *  (x[i, v2] - x[j, v2])
           end
         end
         end
         d[1] = sqrt(d[1])
         return d
end

"""
    KDEonline!{tp}(kdescores::AbstractArray{tp, 1}, x::AbstractArray{tp, 2} [, Q::AbstractArray{tp, 2}], σ::tp, dim::Int = 1)

compute (1.0 - Kernel Density Estimates) from x and write it to kdescores with dim being the dimension of the observations.
If teh covariance matrix Q is given, the Mahalanobis distance is used instead of the Euclidean distance.
"""
function KDEonline!(k::AbstractArray{tp, 1}, d::AbstractArray{tp, 1}, x::AbstractArray{tp, 2}, i::Int, σ::tp, dim::Int = 1) where {tp}
  k[1] = 0.0
  for j = 1:size(x, dim)
    Euclidean_distance!(d, x, i, j, dim)
    innerKDEonline!(k, d, σ)
  end
  finalizeKDEonline!(k, x, dim)
  return k
end

function KDEonline!(k::AbstractArray{tp, 1}, d::AbstractArray{tp, 1}, x::AbstractArray{tp, 2}, Q::AbstractArray{tp, 2}, i::Int, σ::tp, dim::Int = 1) where {tp}
  k[1] = 0.0
  for j = 1:size(x, dim)
    Mahalanobis_distance!(d, x, Q, i, j, dim)
    innerKDEonline!(k, d, σ)
  end
  finalizeKDEonline!(k, x, dim)
  return k
end

function innerKDEonline!(k::AbstractArray{tp, 1}, d::AbstractArray{tp, 1}, σ::tp) where {tp}
  k[1] = k[1] + exp(-0.5 * d[1] / σ^2)
  return k
end

function finalizeKDEonline!(k::AbstractArray{tp, 1}, x::AbstractArray{tp, 2}, dim::Int) where {tp}
  k[1] = 1.0 - (k[1] / size(x, dim))
  return k
end

function KDEonline!(kdescores::AbstractArray{tp, 1}, x::AbstractArray{tp, 2}, σ::tp, dim::Int = 1) where {tp}
  d = zeros(tp, 1)
  k = zeros(tp, 1)
  for i = 1:size(x, dim)
    KDEonline!(k, d, x, i, σ, dim)
    kdescores[i] = k[1]
  end
  return kdescores
end

function KDEonline!(kdescores::AbstractArray{tp, 1}, x::AbstractArray{tp, 2}, Q::AbstractArray{tp, 2}, σ::tp, dim::Int = 1) where {tp}
  d = zeros(tp, 1)
  k = zeros(tp, 1)
  for i = 1:size(x, dim)
    KDEonline!(k, d, x, Q, i, σ, dim)
    kdescores[i] = k[1]
  end
  return kdescores
end

"""
    REConline!{tp}(recscores::AbstractArray{tp, 1}, x::AbstractArray{tp, 2} [, Q::AbstractArray{tp, 2}], ɛ::tp, dim::Int = 1)

compute recurrence scores from x and write it to recscores with dim being the dimension of the observations.
If the covariance matrix Q is given, the mahalanobis distance is used instead of the euclidean distance.
"""
function REConline!(r::AbstractArray{tp, 1}, d::AbstractArray{tp, 1}, x::AbstractArray{tp, 2}, i::Int, ɛ::tp, dim::Int = 1) where {tp}
    r[1] = 0.0
    for j = 1:size(x, dim)
      Euclidean_distance!(d, x, i, j, dim)
      innerREConline!(r, d, ɛ)
    end
    finalizeREConline!(r, x, dim)
    return r
end

function REConline!(r::AbstractArray{tp, 1}, d::AbstractArray{tp, 1}, x::AbstractArray{tp, 2}, Q::AbstractArray{tp, 2}, i::Int, ɛ::tp, dim::Int = 1) where {tp}
    r[1] = 0.0
    for j = 1:size(x, dim)
      Mahalanobis_distance!(d, x, Q, i, j, dim)
      innerREConline!(r, d, ɛ)
    end
    finalizeREConline!(r, x, dim)
    return r
end

function innerREConline!(r::AbstractArray{tp, 1}, d::AbstractArray{tp, 1}, ɛ::tp) where {tp}
  if d[1] <= ɛ
    r[1] = r[1] + 1.0
  end
  return r
end

function finalizeREConline!(r::AbstractArray{tp, 1}, x::AbstractArray{tp, 2}, dim::Int) where {tp}
  r[1] = 1.0 - ((r[1] - 1.0) / size(x, dim))
  return r
end

function REConline!(recscores::AbstractArray{tp, 1}, x::AbstractArray{tp, 2}, ɛ::tp, dim::Int = 1) where {tp}
  d = zeros(tp, 1)
  r = zeros(tp, 1)
  for i = 1:size(x, dim)
    REConline!(r, d, x, i, ɛ, dim)
    recscores[i] = r[1]
  end
  return recscores
end

function REConline!(recscores::AbstractArray{tp, 1}, x::AbstractArray{tp, 2}, Q::AbstractArray{tp, 2}, ɛ::tp, dim::Int = 1) where {tp}
  d = zeros(tp, 1)
  r = zeros(tp, 1)
  for i = 1:size(x, dim)
    REConline!(r, d, x, Q, i, ɛ, dim)
    recscores[i] = r[1]
  end
  return recscores
end

"""
    KNNonline!{tp}(knnscores::AbstractArray{tp, 1}, x::AbstractArray{tp, 2}, [Q::AbstractArray{tp, 2},] k::Int, dim::Int = 1)

compute k-nearest neighbor (gamma) scores from x and write it to knnscores with dim being the dimension of the observations.
If the covariance matrix Q is given, the mahalanobis distance is used instead of the euclidean distance.
"""
function KNNonline!(knn::AbstractArray{tp, 1}, d::AbstractArray{tp, 1}, x::AbstractArray{tp, 2}, i::Int, k::Int, dim::Int = 1) where {tp}
  if size(knn, 1) != k+1 error("size knn != k+1") end
  for l = 1:length(knn)
    knn[l] = 1.0e30#typemax(tp)
  end
  for j = 1:size(x, dim)
    Euclidean_distance!(d, x, i, j, dim)
    innerKNNonline!(knn, d, k)
  end
  finalizeKNNonline!(knn, k)
  return knn
end

function KNNonline!(knn::AbstractArray{tp, 1}, d::AbstractArray{tp, 1}, x::AbstractArray{tp, 2}, Q::AbstractArray{tp, 2}, i::Int, k::Int, dim::Int = 1) where {tp}
  if size(knn, 1) != k+1 error("size knn != k+1") end
  for l = 1:length(knn)
    knn[l] = 1.0e30#typemax(tp)
  end
  for j = 1:size(x, dim)
    Mahalanobis_distance!(d, x, Q, i, j, dim)
    innerKNNonline!(knn, d, k)
  end
  finalizeKNNonline!(knn, k)
  return knn
end

function KNNonline!(knnscores::AbstractArray{tp, 1}, x::AbstractArray{tp, 2}, k::Int, dim::Int = 1) where {tp}
  d = zeros(tp, 1)
  knn = fill(typemax(tp), k+1)
  for i = 1:size(x, dim)
    KNNonline!(knn, d, x, i, k, dim)
    knnscores[i] = knn[1] #mean(knn[2:k+1]) --> in finalize
  end
  return knnscores
end

function KNNonline!(knnscores::AbstractArray{tp, 1}, x::AbstractArray{tp, 2}, Q::AbstractArray{tp, 2}, k::Int, dim::Int = 1) where {tp}
  d = zeros(tp, 1)
  knn = fill(typemax(tp), k+1)
  for i = 1:size(x, dim)
    KNNonline!(knn, d, x, Q, i, k, dim)
    knnscores[i] = knn[1] #mean(knn[2:k+1]) --> in finalize
  end
  return knnscores
end

function innerKNNonline!(knn::AbstractArray{tp, 1}, d::AbstractArray{tp, 1}, k::Int) where {tp}
    if d[1] <= knn[1]
      for kshift = (k+1):-1:2
        knn[kshift] = knn[kshift-1] # shift old one to k+1
      end
      knn[1] = d[1]   # replace by new distance
    else
      for j = 2:(k+1)
        if d[1] > knn[j-1] && d[1] <= knn[j]
          for kshift = k+1:-1:j
            knn[kshift] = knn[kshift-1] # shift old one to k+1
          end
          knn[j] = d[1]     # replace by new distance
        end
      end
    end
  return knn
end

function finalizeKNNonline!(knn::AbstractArray{tp, 1}, k::Int) where {tp}
  knn[1] = 0.0
  for i = 2:(k+1)
    knn[1] = knn[1] + knn[i]
  end
  knn[1] = knn[1] / k
  return knn
end

"""
    SigmaOnline!{tp}(sigma::Array{tp, 1}, x::AbstractArray{tp, 2}, [Q::AbstractArray{tp, 2}], samplesize::Int = 250, dim::Int = 1)

compute `sigma` parameter as mean of the distances of `samplesize` randomly sampled points along `dim`.
If Q is given the Mahalanobis distance is used instead of Euclidean.
"""
function SigmaOnline!(sigma::Array{tp, 1}, x::AbstractArray{tp, 2}, samplesize::Int = 250, dim::Int = 1) where {tp}
  sigma[1] = 0.0
  d = zeros(tp, 1)
  trainsample1 = zeros(Int, 1)
  trainsample2 = zeros(Int, 1)
  for i = 1:samplesize
    StatsBase.rand!(trainsample1, 1:size(x, dim))[1]
    for j = 1:samplesize
      if j > i
        Euclidean_distance!(d, x, trainsample1[1], StatsBase.rand!(trainsample2, 1:size(x, dim))[1], dim)
        # sum of distances
        sigma[1] = sigma[1] + d[1]
      end
    end
  end
  # get mean
  sigma[1] = sigma[1] / (samplesize^2 / 2 - samplesize / 2)
  return sigma
end

function SigmaOnline!(sigma::Array{tp, 1}, x::AbstractArray{tp, 2}, Q::AbstractArray{tp, 2}, samplesize::Int = 250, dim::Int = 1) where {tp}
  sigma[1] = 0.0
  d = zeros(tp, 1)
  trainsample1 = zeros(Int, 1)
  trainsample2 = zeros(Int, 1)
  for i = 1:samplesize
    StatsBase.rand!(trainsample1, 1:size(x, dim))[1]
    for j = 1:samplesize
      if j > i
        Mahalanobis_distance!(d, x, Q, trainsample1[1], StatsBase.rand!(trainsample2, 1:size(x, dim))[1], dim)
        # sum of distances
        sigma[1] = sigma[1] + d[1]
      end
    end
  end
  # get mean
  sigma[1] = sigma[1] / (samplesize^2 / 2 - samplesize / 2)
  return sigma
end


using GR
using LaTeXStrings

formulas = (
    L"- \frac{{\hbar ^2}}{{2m}}\frac{{\partial ^2 \psi (x,t)}}{{\partial x^2 }} + U(x)\psi (x,t) = i\hbar \frac{{\partial \psi (x,t)}}{{\partial t}}",
    L"\zeta \left({s}\right) := \sum_{n=1}^\infty \frac{1}{n^s} \quad \sigma = \Re(s) > 1",
    L"\zeta \left({s}\right) := \frac{1}{\Gamma(s)} \int_{0}^\infty \frac{x^{s-1}}{e^x-1} dx" )

selntran(0)

settextfontprec(232, 3)
settextalign(2, 3)
chh = 0.036

for phi in LinRange(0, 2pi, 360)
    clearws()
    setcharheight(chh)
    setcharup(sin(phi), cos(phi))
    y = 0.6
    for s in formulas
        mathtex(0.5, y, s)
        tbx, tby = inqmathtex(0.5, y, s)
        fillarea(tbx, tby)
        y -= 0.2
    end
    updatews()
end


import Pkg
Pkg.activate("..")

using CrimsonSkyline
using Distributions: Beta, Geometric
using Logging
using StatsBase: mean, std
using Random: seed!
using Plots

seed!(2021)

num_constraint(x, num, bound) = (x > num - bound) && (x < num + bound)

function constrained(t::Trace, evidence::Vector{Int64}, bound::Int64)
    prob = sample(t, "prob", Beta(2.0, 2.0))
    num = sample(t, "num", Geometric(prob))
    plate(t, observe, "constraint", SoftDelta(x -> num_constraint(x, num, bound), -2.0), evidence)
end

function unconstrained(t::Trace, evidence::Vector{Int64})
    prob = sample(t, "prob", Beta(2.0, 2.0))
    plate(t, observe, "num", Geometric(prob), evidence)
end

function main()
    data = [1, 5, 2, 9, 3, 2, 1, 3, 2, 1, 1, 4, 3, 2]
    bound = 3
    @time constrained_results = mh(constrained; params = (data, bound))
    @time unconstrained_results = mh(unconstrained; params = (data,))
    aic_constrained = aic(constrained_results)
    aic_unconstrained = aic(unconstrained_results)
    @info "AIC of constrained model = $aic_constrained"
    @info "AIC of unconstrained model = $aic_unconstrained"
    hpdi_con = hpdi(constrained_results, 0.5, ["prob"])
    hpdi_uncon = hpdi(unconstrained_results, 0.5, ["prob"])
    @info "50% HPDI of prob under constrained model: $hpdi_con"
    @info "50% HPDI of prob under unconstrained model: $hpdi_uncon"
end

if abspath(PROGRAM_FILE) == @__FILE__
    main()
end


module CoulombExchange

#LearnHatreeFock.jl Modules
using TypesBasis,TypesParticles

export getcoulexchange

@doc raw"""
function getcoulexchange()

"""
function getcoulexchange(density::Array,elecelecrepul::Array)

    nbasisfunc = size(density)[1];
    coulexchange = zeros(Float64,nbasisfunc,nbasisfunc);

    for l=1:nbasisfunc
        for m=1:nbasisfunc
            for n=1:nbasisfunc
                for o=1:nbasisfunc
                    diff = 2.00e0*elecelecrepul[l,m,n,o]-elecelecrepul[l,n,m,o];
                    coulexchange[l,m] += density[n,o]*(diff);
                end
            end #n=1:nbasisfunc
        end
    end #l=1:nbasisfunc

    return coulexchange
end #getcoulexchange

end #CoulombExchange


using OceanTurb

@use_pyplot_utils # add utilities for plotting OceanTurb Fields

     N = 128        # Model resolution
     L = 128        # Vertical extent of the model domain
    Fb = 1e-7       # Surface buoyancy flux (positive implies cooling)
  dTdz = 1e-3       # Interior/initial temperature gradient
    Δt = 10minute   # Time step size
tfinal = 8hour      # Final time

# Build the model with a Backward Euler timestepper
model = KPP.Model(N=N, L=L, stepper=:BackwardEuler)

# Set initial condition
T₀(z) = 20 + dTdz * z
model.solution.T = T₀

# Set boundary conditions
model.bcs.T.top = FluxBoundaryCondition(Fb / (model.constants.α * model.constants.g))
model.bcs.T.bottom = GradientBoundaryCondition(dTdz)

# Run the model
run_until!(model, Δt, tfinal)

plot(model.solution.T)
removespines("top", "right")
xlabel("Temperature (\$ {}^\\circ \\mathrm{C} \$)")
ylabel(L"z \, \mathrm{(m)}")

gcf()
savefig("figs/kpp_free_convection.png", dpi=480)


"""
A shape where anything inside could cause a disturbance at the listener position
from a point source wavefront starting at the listener. Also everything inside
has a positive `x` coordinate. It is equivalent to a segment of a circle.

More precisely, if the listener is at (l_x,l_y) then the interior of the shape
is defined as
sqrt((x-l_x)^2+(y-l_y)^2)<time and x>0
"""
struct TimeOfFlightFromPoint{T <: AbstractFloat} <: Shape{T,2}
    listener_position::Vector{T}
    time::T
end

name(shape::TimeOfFlightFromPoint) = "Time of flight from point source"

function volume(shape::TimeOfFlightFromPoint)
    θ = 2acos(-shape.listener_position[1] / shape.time)
    return shape.time^2 * (θ - sin(θ))/2
end

import Base.issubset
function issubset(circle::Sphere{T,2}, shape::TimeOfFlightFromPoint) where T
    (origin(circle)[1] - circle.radius) > 0 &&
    norm(origin(circle) - shape.listener_position) < (shape.time - 2circle.radius)
end

function bounding_box(shape::TimeOfFlightFromPoint{T}) where T <: AbstractFloat
    box_height = 2sqrt(shape.time^2 - shape.listener_position[1]^2)
    box_width = max(shape.time + shape.listener_position[1], zero(T))
    return Box([SVector(zero(T), -box_height/2), SVector(box_width, box_height/2)])
end

function boundary_functions(shape::TimeOfFlightFromPoint)

    function x(t)
        check_boundary_coord_range(t)
        if t <= 1//2
            θ = acos(-shape.listener_position[1] / shape.time)
            return shape.time*cos(θ*(4t-1)) + shape.listener_position[1]
        else
            return zero(t)
        end
    end

    function y(t)
        check_boundary_coord_range(t)
        if t <= 1//2
            θ = acos(-shape.listener_position[1] / shape.time)
            return shape.time*sin(θ*(4t-1)) + shape.listener_position[2]
        else
            return 2*(3//4-t)*2*sqrt(shape.time^2 - shape.listener_position[1]^2)
        end
    end

    return x, y
end


module TBfromk2r
    using Dierckx
    using SparseArrays
    const eps = 1e-7

    const σ0 = [
    1 0
    0 1
    ]

    const σx = [
    0 1
    1 0
    ]

    const σy = [
    0 -im
    im 0
    ]

    const σz = [
    1 0
    0 -1
    ]

    export surfaceHamiltonian,σx,σy,σz,σ0

    function surfaceHamiltonian(Hk,norb;numhop=3,L=32,kpara="kx",BC="OBC")
        function calc_H(k)
            if kpara == "kx"
                fx(kperp) = Hk([k,kperp])[:,:]
                mat_H = make_Hamiltonian(fx,[numhop],L,norb,BC)
            elseif kpara == "ky"
                fy(kperp) = Hk([kperp,k])[:,:]
                mat_H = make_Hamiltonian(fy,[numhop],L,norb,BC)
            else
                println("Error! not supported.")
                exit()
            end

            return mat_H
        end
        k ->  calc_H(k)
    end

    function make_Hamiltonian(Hk,numhop,L,norb,BC)
        Vol = prod(L)
        dim = length(L)
        mat_H = spzeros(ComplexF64,Vol*norb,Vol*norb)
        #println(Vol*norb)

        mat_t = make_hop(Hk,numhop,norb)

        for i = 1:Vol
            for jorb = 1:norb
                for iorb = 1:norb
                    ii = (iorb-1)*Vol + i
                    #println(ii,"\t",mat_t[iorb,jorb])
                    isite = calc_i2site(i,L)

                    for hop in mat_t[iorb,jorb]
                        jsite = isite .+ hop[1]
                        if BC == "PBC"
                            for il=1:dim
                                jsite[il] += ifelse(jsite[il] < 1,L[il],0)
                                jsite[il] += ifelse(jsite[il] > L[il],-L[il],0)
                            end
                        end
                        flag = false

                        if dim ==1
                            if 1 <= jsite[1] <= L[1]
                                flag = true
                            end
                        elseif dim == 2
                            if 1 <= jsite[1] <= L[1] && 1 <= jsite[2] <= L[2]
                                flag = true
                            end
                        end

                        if flag
                            j = calc_r2j(jsite,L)

                            jj = (jorb-1)*Vol + j

                            mat_H[ii,jj] = hop[2]
                        end
                    end
                end
            end
        end
        return mat_H

    end

    function calc_i2site(i,L)
        dim = length(L)
        if dim == 1
            isite = [i]
        elseif dim ==2
            ix = (i-1) % L[1] + 1
            iy = div(i-ix,L[1])+ 1
            isite = [ix,iy]
        end
        return  isite
    end

    function calc_r2j(hopping,L)
        dim = length(L)
        if dim == 1
            j = hopping[1]
        elseif dim == 2
            j = (hopping[2]-1)*L[1]+hopping[1]
        end
        return j
    end

    function make_hop(Hk,n::Array{Int,1},norb)
        mat_t  = Array{Any}(undef, norb, norb)

        for i=1:norb
            for j=1:norb
                f(k) = Hk(k)[i,j]
                hop_t = calc_coeff(f,n)
                mat_t[i,j] = hop_t
                #println("hoppings matrix at ($i,$j) element")
                #for hop in hop_t
                #    println(hop[1],"\t",hop[2])
                #end
            end
        end

        #println(mat_t)
        return mat_t

    end

    function calc_coeff(fk,n::Array{Int,1})
        dim = length(n)

        #println(dim)
        if dim == 1
            ck = calc_coeff(fk,n[1],1)
        elseif dim == 2
            ck = calc_coeff(fk,n[1],n[2],1,1)
        else
            println("not supported yet: dimension is $dim")
            return
        end
        return ck
    end

    function calc_coeff(fk,nkx::Int,ax::Number;L=256)
        numhop = 2nkx + 1
        ck = zeros(ComplexF64,numhop)

        hopping = range(-nkx, stop=nkx)
        kmax_x = 2π/ax
        kx = range(0.0,stop=kmax_x ,length=L)

        for i=1:numhop
            f(k) = fk(k)*exp(im*hopping[i]*k)
            mat_f = make_f(f,L,kmax_x)./(2π)
            spline_re = Spline1D(kx, real.(mat_f))
            spline_im = Spline1D(kx, imag.(mat_f))
            integ_re = integrate(spline_re, 0, kmax_x)
            integ_im = integrate(spline_im, 0, kmax_x)

            if abs(integ_re) > eps
                ck[i] += integ_re
            end
            if abs(integ_im) > eps
                ck[i] += im*integ_im
            end

        end

        hop_t = []
        for ix=1:numhop

            if abs(ck[ix]) > eps
                push!(hop_t,
                    (((hopping[ix]),ck[ix]))
                )
            end

        end

        return hop_t #ck ./(2π)
    end

    function make_f(f,L,kmax_x,kmax_y)
        mat_f = zeros(ComplexF64,L,L)
        kx = range(0.0,stop=kmax_x,length=L)
        ky = range(0.0,stop=kmax_y,length=L)

        for ikx =1:L
            for iky=1:L
                mat_f[ikx,iky] = f([kx[ikx],ky[iky]])#f(kx[ikx],ky[iky])
            end
        end

        return mat_f
    end

    function make_f(f,L,kmax_x)
        mat_f = zeros(ComplexF64,L)
        kx = range(0.0,stop=kmax_x,length=L)

        for ikx =1:L
            mat_f[ikx] = f(kx[ikx])#f(kx[ikx],ky[iky])
        end

        return mat_f
    end

    function calc_coeff(fk,nkx::Int,nky::Int,ax::Number,ay::Number;L=256)
        numhop_x = 2nkx + 1
        numhop_y = 2nky + 1
        ck = zeros(ComplexF64,numhop_x,numhop_y)
        hopping_x = range(-nkx, stop=nkx)
        hopping_y = range(-nky, stop=nky)
        kmax_x = 2π/ax
        kmax_y = 2π/ay
        kx = range(0.0,stop=kmax_x ,length=L)
        ky = range(0.0,stop=kmax_y,length=L)
        hop_t = []

        for ix=1:numhop_x
            for iy=1:numhop_y
                f(k) = fk(k)*exp(im*hopping_x[ix]*k[1])*exp(im*hopping_y[iy]*k[2])
                mat_f = make_f(f,L,kmax_x,kmax_y)./(2π)^2
                spline_re = Spline2D(kx, ky, real.(mat_f))
                spline_im = Spline2D(kx, ky, imag.(mat_f))
                integ_re = integrate(spline_re, 0, kmax_x, 0, kmax_y)
                integ_im = integrate(spline_im, 0, kmax_x, 0, kmax_y)
                if abs(integ_re) > eps
                    ck[ix,iy] += integ_re
                end
                if abs(integ_im) > eps
                    ck[ix,iy] += im*integ_im
                end
                #ck[ix,iy] = integ_re + im*integ_im
                #println("$(hopping_x[ix]), $(hopping_y[iy]), $(ck[ix,iy])")
            end
        end

        hop_t = []
        for ix=1:numhop_x
            for iy=1:numhop_y
                if abs(ck[ix,iy]) > eps
                    push!(hop_t,
                        (((hopping_x[ix],hopping_y[iy]),ck[ix,iy]))
                    )
                end
            end
        end

        return hop_t

    end

end


using JuMP, Ipopt

function tcm_estimateMNL(N, assortments, transaction_counts)

	M = size(assortments, 1)

	total_trans_by_assortment = sum(transaction_counts, 2)

	m_MNL = Model(solver = IpoptSolver(print_level=1))

	@variable(m_MNL, -20 <= u[1:N-1] <= 20)

	@NLobjective(m_MNL, Max, sum(sum( transaction_counts[m,p] * u[p] for p in 1:N-1) for m in 1:M) - sum(total_trans_by_assortment[m]*log(1+sum(assortments[m,p]* exp(u[p]) for p in 1:N-1)) for m in 1:M) )

	start_time = time()
	solve(m_MNL)
	end_time = time()
	elapsed_time = end_time - start_time

	u_val = getvalue(u)

	obj = getobjectivevalue(m_MNL)

	return u_val, obj, elapsed_time
end


using Evolutionary
using Test
using Random
using LinearAlgebra
using Statistics
Random.seed!(9874984737486)

for tests in [
    "interface.jl",
    "selections.jl",
    "recombinations.jl",
    "mutations.jl",
    "sphere.jl",
    "rosenbrock.jl",
    "schwefel.jl",
    "rastrigin.jl",
    "n-queens.jl",
    "knapsack.jl",
    "onemax.jl",
    "regression.jl",
    "gp.jl"
]
    include(tests)
end


# Idea: Use LASSO for screening the active factors in an optimization
# problem.

using GLMNet

# Say that we want to optimize a function which has many factors but where
# only very few of them are active in calculating the outcome.
N = 1000
K = 3
active_factors = shuffle(collect(1:N))[1:K]
beta_intercept = 10.0
beta_slope = 1.0
betas = beta_intercept + beta_slope * randn(K, 1)
ofunc(x) = sum(x[active_factors] .* betas)

# Now lets sample and calc the function values for samples
S = int(10*K*log10(N));
xs = randn(N, S);
y = zeros(S)
[(y[i] = ofunc(xs[:,i])) for i in 1:S]

# Now do lasso regression:
#path = glmnet(xs', y; standardize = false)
# cv = glmnetcv(xs', y; standardize = false)

function active_factors_from_lasso(xs, y; alpha = 1.0)
  n = size(xs, 2)
  # Find the best lambda using cross-validation. We limit the size of the predictors.
  #cv = glmnetcv(xs, y) # ; constraints = hcat(-1000.0*ones(n), 1000.0*ones(n)))
  cv = glmnetcv(xs, y; standardize = false, alpha = alpha)
  num_steps = size(cv.path.betas, 2)
  res = Any[]
  for(i in 1:num_steps)
    pbetas = cv.path.betas[:, i]
    pbwi = [(pbetas[i], i) for i in 1:length(pbetas)]
    push!(res, filter((p) -> abs(p[1]) > 0.001, pbwi))
  end
  return res
end

# Count how many times each factor was found by the lasso.
function most_active_factors_from_lasso(xs, y; alpha = 1.0)
  lasso_factors = active_factors_from_lasso(xs, y; alpha = alpha)
  factor_counts = Dict{Int64, Int64}()
  for(i in 1:length(lasso_factors))
    factors = map((t) -> t[2], lasso_factors[i])
    for(f in factors)
      if haskey(factor_counts, f)
        factor_counts[f] += 1
      else
        factor_counts[f] = 1
      end
    end
  end
  return factor_counts, lasso_factors
end

function eval_if_lasso_finds_active_factors(N, K, S, reps = 100;
  beta_intercept = 0.0, beta_slope = 1.0, alpha = 1.0)

  active_factors = shuffle(collect(1:N))[1:K]
  print("Active factors:"); show(active_factors)
  betas = beta_intercept + beta_slope * randn(K, 1)
  print("\nBetas:"); show(betas)

  objective_func(x) = begin
    active = x[active_factors]
    sum(active .* betas)
  end

  # We want to count how many times the predicted factors (ranked based on how
  # often they were in the lasso factor set) are among the actually active factors.
  factor_found_counts = Dict{Int64, Int64}()

  for(rep in 1:reps)
    # Now lets sample
    xs = randn(N, S)

    # And calc the function values for them:
    y = zeros(S)
    for(i in 1:S)
      y[i] = objective_func(xs[:,i])
    end

    # Get the active factors from lasso reg
    factor_counts, factor_path = most_active_factors_from_lasso(xs', y; alpha = alpha)
    print("\nFactor counts:"); show(factor_counts)

    # Rank factors based on how often they were predicted as active
    factors_ranked = map((t) -> t[1], sort(collect(factor_counts), by = (t) -> t[2], rev = true))

    for(i in 1:length(factors_ranked))
      if in(factors_ranked[i], active_factors)
        if haskey(factor_found_counts, i)
          factor_found_counts[i] += 1
        else
          factor_found_counts[i] = 1
        end
      end
    end
  end

  return factor_found_counts
end

# It seems that 5*K*log10(N) samples is often always enough:
N = 10000
K = 3
@time eval_if_lasso_finds_active_factors(N, K, int(5*K*log10(N)), 10)
@time eval_if_lasso_finds_active_factors(N, K, int(5*K*log10(N)), 10; beta_intercept = 0.0)

# To make this even more systematic we need to be able to generate new objective functions
# involving a variety of factors, interactions and so forth.
type RandObjectiveFunction
  factors
  betas

  # Constructor creates a random function where we can control number of factors
  # per order, max order, total num vars as well as if effect heritage is active:
  #  N     = total number of variables
  #  maxOrder = maximum number of interaction terms in one factor (a product) of the sum that is the objective func
  #  sparsitiesPerOrder = array that gives the number of factors for that order.
  #  effect_heritage = true iff the factors of higher order can only involve vars used for order 1 (main effects)
  RandObjectiveFunction(N, maxOrder = 1, factorsPerOrder = [2];
    effect_heritage = true) = begin
    all_vars = 1:N

    # One factor is represented by the indices of vars that should be multiplied to create it.

    # Create main vars first since we might need them for heritance.
    main_vars = shuffle(collect(all_vars))[1:factorsPerOrder[1]]
    factors = map((v) -> [v], main_vars)

    # Add factors for higher orders
    for(order in 2:length(factorsPerOrder))
      for(i in 1:factorsPerOrder[order])
        factor = []
        for(j in 1:order)
          if effect_heritage
            vars_to_sample = main_vars
          else
            vars_to_sample = all_vars
          end
          push!(factor, shuffle(vars_to_sample)[1])
        end
        push!(factors, factor)
      end
    end

    # Create random betas for each factor, i.e. their weight.
    betas = randn(length(factors), 1)

    new(factors, betas)
  end
end

# Apply function f to an array a.
function eval(f::RandObjectiveFunction, a)
  sum = 0.0
  index = 1
  for(factor in f.factors)
    sum += f.betas[index] * a
    index += 1
  end
  s = size(xs, 2)
  y = zeros(s)
  for(i in 1:s)
    y[i] =
  end
end


# Return the index of the maximum entry of a given vector transformed by f.
Base.argmax(f::Any, v::AbsVec) = argmax(f, v, 1:length(v))

# Among specified indices, return the index of the maximum entry of a given vector
# transformed by f.
function Base.argmax(f::Any, v::AbsVec, indv::AbsVecInteger)
    ind = 0  # return 0 if indv is empty
    val = -Inf
    for n = indv
        if f(v[n]) ≥ val  # ≥ rather than > to prevent returning 0 when v has -Inf
            val = f(v[n])
            ind = n
        end
    end

    return ind
end

# Return the index of the minimum entry of a given vector transformed by f
Base.argmin(f::Any, v::AbsVec) = argmin(f, v, 1:length(v))

# Among specified indices, return the index of the minimum entry of a given vector
# transformed by f.
function Base.argmin(f::Any, v::AbsVec, indv::AbsVecInteger)
    ind = 0  # return 0 if v is empty
    val = Inf
    for n = indv
        if f(v[n]) ≤ val  # ≤ rather than < to prevent returning 0 when v has -Inf
            val = f(v[n])
            ind = n
        end
    end

    return ind
end


struct ConfusionMatrix
  tp::Int
  fp::Int
  fn::Int
  tn::Int
end

function getConfusionMatrix(Δ::AbstractMatrix{T}, hΔ::AbstractMatrix{T}, thr=1e-4) where T
  p = size(Δ, 1)

  tp = fp = fn = tn = 0
  for c=1:p
    for r=c:p
      if Δ[r, c] != zero(T)
        if abs(hΔ[r, c]) > thr
          # not zero estimated as not zero
          tp += 1
        else
          # not zero estimated as zero
          fn += 1
        end
      else
        if abs(hΔ[r, c]) > thr
          # zero estimated as not zero
          fp += 1
        else
          # zero estimated as zero
          tn += 1
        end
      end
    end
  end
  ConfusionMatrix(tp, fp, fn, tn)
end

function getConfusionMatrix(Δ::AbstractMatrix{T}, hΔ::SparseIterate{T}, thr=1e-4) where T
  p = size(Δ, 1)

  tp = fp = fn = tn = 0
  for c=1:p
    for r=c:p
      if Δ[r, c] != zero(T)
        if abs(hΔ[sub2indLowerTriangular(p, r, c)]) > thr
          # not zero estimated as not zero
          tp += 1
        else
          # not zero estimated as zero
          fn += 1
        end
      else
        if abs(hΔ[sub2indLowerTriangular(p, r, c)]) > thr
          # zero estimated as not zero
          fp += 1
        else
          # zero estimated as zero
          tn += 1
        end
      end
    end
  end
  ConfusionMatrix(tp, fp, fn, tn)
end

function tpr(t::ConfusionMatrix)
  dn = t.tp + t.fn
  dn > 0 ? t.tp / dn : 0.
end
function fpr(t::ConfusionMatrix)
  dn = t.fp + t.tn
  dn > 0 ? t.fp / dn : 0.
end
function precision(t::ConfusionMatrix)
  dn = t.tp + t.fp
  dn > 0 ? t.tp / dn : 1.
end
recall(t::ConfusionMatrix) = tpr(t)


abstract type HuffmanTree end

struct HuffmanLeaf <: HuffmanTree
    ch::Char
    freq::Int
end

struct HuffmanNode <: HuffmanTree
    freq::Int
    left::HuffmanTree
    right::HuffmanTree
end

function makefreqdict(s::String)
    d = Dict{Char, Int}()
    for c in s
        if !haskey(d, c)
            d[c] = 1
        else
            d[c] += 1
        end
    end
    d
end

function huffmantree(ftable::Dict)
    trees::Vector{HuffmanTree} = [HuffmanLeaf(ch, fq) for (ch, fq) in ftable]
    while length(trees) > 1
        sort!(trees, lt = (x, y) -> x.freq < y.freq, rev = true)
        least = pop!(trees)
        nextleast = pop!(trees)
        push!(trees, HuffmanNode(least.freq + nextleast.freq, least, nextleast))
    end
    trees[1]
end

printencoding(lf::HuffmanLeaf, code) = println(lf.ch == ' ' ? "space" : lf.ch, "\t", lf.freq, "\t", code)

function printencoding(nd::HuffmanNode, code)
    code *= '0'
    printencoding(nd.left, code)
    code = code[1:end-1]

    code *= '1'
    printencoding(nd.right, code)
    code = code[1:end-1]
end

const msg = "this is an example for huffman encoding"

println("Char\tFreq\tHuffman code")

printencoding(huffmantree(makefreqdict(msg)), "")


get_vars = Symbolics.get_variables

function isin(var::Num, vars)
    for varible in vars
        if isequal(varible, var)
            return true
        end
    end
    return false
end

function get_Num(vars::Dict)
    return collect(keys(vars))
end

function check_eqs(eqs)
    if typeof(eqs) == Equation
        return toexpr([eqs])
    elseif typeof(eqs) == Vector{Equation}
        return toexpr(eqs)
    else
        error("Error: type of eqs should not be $(typeof(eqs)).It must be Equation or Vector{Equation}")
    end
end

function toexpr(eqs)
    expr = Vector{Num}([])
    for eq in eqs
        push!(expr,eq.lhs - eq.rhs)
    end
    return expr
end

function get_all_vars(eqs_expr)
    vars = Vector{Num}([])
    for expr in eqs_expr
        vars = union(vars,get_vars(expr))
    end
    return vars
end

function check_vars(eqs, vars)
    copy_vars = copy(vars)
    eqs_expr = eqs
    eqs_var = get_all_vars(eqs_expr)
    if length(eqs_var) != length(collect(keys(vars)))
        error("Error: number of varialbes in vars don't equal to
        number of variables in eqs")
    end
    for var in eqs_var
        try
            pop!(copy_vars, var)
        catch
            error("Error: $(var) not in vars")
        end
    end
    rest_var = collect(keys(copy_vars))
    if length(rest_var) > 0
        error("$(rest_var) not in eqs")
    end
    return true
end

function islinear(eqs, vars)
    for eq in eqs
        if !Symbolics.islinear(eq.lhs, get_Num(vars))
            return false
        end
    end
    return true
end


# used in shortest path calculations

struct DefaultDistance<:AbstractMatrix{Int}
    nv::Int
    DefaultDistance(nv::Int=typemax(Int)) = new(nv)
end

DefaultDistance(nv::Integer) = DefaultDistance(Int(nv))

show(io::IO, x::DefaultDistance) = print(io, "$(x.nv) × $(x.nv) default distance matrix (value = 1)")
show(io::IO, z::MIME"text/plain", x::DefaultDistance) = show(io, x)

getindex(::DefaultDistance, s::Integer, d::Integer) = 1
getindex(::DefaultDistance, s::UnitRange, d::UnitRange) = DefaultDistance(length(s))
size(d::DefaultDistance) = (d.nv, d.nv)
transpose(d::DefaultDistance) = d
ctranspose(d::DefaultDistance) = d

"""
    eccentricity(g[, v][, distmx])

Return the eccentricity[ies] of a vertex / vertex list `v` or the
entire graph. An optional matrix of edge distances may be supplied; if missing,
edge distances default to `1`.

The eccentricity of a vertex is the maximum shortest-path distance between it
and all other vertices in the graph.

The output is either a single float (when a single vertex is provided) or a
vector of floats corresponding to the vertex vector. If no vertex vector is
provided, the vector returned corresponds to each vertex in the graph.

### Performance
Because this function must calculate shortest paths for all vertices supplied
in the argument list, it may take a long time.

### Implementation Notes
The eccentricity vector returned by `eccentricity()` may be used as input
for the rest of the distance measures below. If an eccentricity vector is
provided, it will be used. Otherwise, an eccentricity vector will be calculated
for each call to the function. It may therefore be more efficient to calculate,
store, and pass the eccentricities if multiple distance measures are desired.
"""
function eccentricity(
    g::AbstractGraph,
    v::Integer,
    distmx::AbstractMatrix{T} = weights(g)
) where T
    e = maximum(dijkstra_shortest_paths(g, v, distmx).dists)
    e == typemax(T) && error("Infinite path length detected")

    return e
end

eccentricity(
    g::AbstractGraph,
    vs::AbstractVector = vertices(g),
    distmx::AbstractMatrix = weights(g)
) = [eccentricity(g, v, distmx) for v in vs]

eccentricity(g::AbstractGraph, distmx::AbstractMatrix) =
    eccentricity(g, vertices(g), distmx)

"""
    diameter(g, distmx=weights(g))
    diameter(eccentricities)

Given a graph and optional distance matrix, or a vector of precomputed
eccentricities, return the maximum eccentricity of the graph.
"""
diameter(eccentricities::Vector) = maximum(eccentricities)
diameter(g::AbstractGraph, distmx::AbstractMatrix = weights(g)) =
    maximum(eccentricity(g, distmx))

"""
    periphery(g, distmx=weights(g))
    periphery(eccentricities)

Given a graph and optional distance matrix, or a vector of precomputed
eccentricities, return the set of all vertices whose eccentricity is
equal to the graph's diameter (that is, the set of vertices with the
largest eccentricity).
"""
function periphery(eccentricities::Vector)
    diam = maximum(eccentricities)
    return filter(x -> eccentricities[x] == diam, 1:length(eccentricities))
end

periphery(g::AbstractGraph, distmx::AbstractMatrix = weights(g)) =
    periphery(eccentricity(g, distmx))

"""
    radius(g, distmx=weights(g))
    radius(eccentricities)

Given a graph and optional distance matrix, or a vector of precomputed
eccentricities, return the minimum eccentricity of the graph.
"""
radius(eccentricities::Vector) = minimum(eccentricities)
radius(g::AbstractGraph, distmx::AbstractMatrix = weights(g)) =
    minimum(eccentricity(g, distmx))

"""
    center(g, distmx=weights(g))
    center(eccentricities)

Given a graph and optional distance matrix, or a vector of precomputed
eccentricities, return the set of all vertices whose eccentricity is equal
to the graph's radius (that is, the set of vertices with the smallest eccentricity).
"""
function center(eccentricities::Vector)
    rad = radius(eccentricities)
    return filter(x -> eccentricities[x] == rad, 1:length(eccentricities))
end

center(g::AbstractGraph, distmx::AbstractMatrix = weights(g)) =
    center(eccentricity(g, distmx))


# Disjoint-Set

############################################################
#
#   A forest of disjoint sets of integers
#
#   Since each element is an integer, we can use arrays
#   instead of dictionary (for efficiency)
#
#   Disjoint sets over other key types can be implemented
#   based on an IntDisjointSet through a map from the key
#   to an integer index
#
############################################################
_intdisjointset_bounds_err_msg(T) = "the maximum number of elements in IntDisjointSet{$T} is $(typemax(T))"

abstract type Scalarness end
struct Scalar <: Scalarness end
struct NotScalar <: Scalarness end

isscalar(::Type{Any}) = NotScalar() # if we don't know the type we can't really know if scalar or not
isscalar(::Type{<:AbstractString}) = NotScalar() # We consider strings to be nonscalar
isscalar(::Type{<:Number}) = Scalar() # We consider Numbers to be scalar
isscalar(::Type{Char}) = Scalar() # We consider Character to be scalar
isscalar(::Type{Symbol}) = Scalar()
isscalar(::Type{T}) where T = hasmethod(iterate, (T,)) ? NotScalar() : Scalar()
"""
    IntDisjointSet{T<:Integer}(n::Integer)
A forest of disjoint sets of integers, which is a data structure
(also called a union–find data structure or merge–find set)
that tracks a set of elements partitioned
into a number of disjoint (non-overlapping) subsets.
"""
mutable struct IntDisjointSet{T<:Integer}
    parents::Vector{T}
    ranks::Vector{T}
    ngroups::T
end

IntDisjointSet(n::T) where {T<:Integer} = IntDisjointSet{T}(collect(Base.OneTo(n)), zeros(T, n), n)
IntDisjointSet{T}(n::Integer) where {T<:Integer} = IntDisjointSet{T}(collect(Base.OneTo(T(n))), zeros(T, T(n)), T(n))
Base.length(s::IntDisjointSet) = length(s.parents)

"""
    num_groups(s::IntDisjointSet)
Get a number of groups.
"""
num_groups(s::IntDisjointSet) = s.ngroups
Base.eltype(::Type{IntDisjointSet{T}}) where {T<:Integer} = T

# find the root element of the subset that contains x
# path compression is implemented here
function find_root_impl!(parents::Vector{T}, x::Integer) where {T<:Integer}
    p = parents[x]
    @inbounds if parents[p] != p
        parents[x] = p = _find_root_impl!(parents, p)
    end
    return p
end

# unsafe version of the above
function _find_root_impl!(parents::Vector{T}, x::Integer) where {T<:Integer}
    @inbounds p = parents[x]
    @inbounds if parents[p] != p
        parents[x] = p = _find_root_impl!(parents, p)
    end
    return p
end

"""
    find_root!(s::IntDisjointSet{T}, x::T)
Find the root element of the subset that contains an member x.
Path compression is implemented here.
"""
find_root!(s::IntDisjointSet{T}, x::T) where {T<:Integer} = find_root_impl!(s.parents, x)

"""
    in_same_set(s::IntDisjointSet{T}, x::T, y::T)
Returns `true` if `x` and `y` belong to the same subset in `s` and `false` otherwise.
"""
in_same_set(s::IntDisjointSet{T}, x::T, y::T) where {T<:Integer} = find_root!(s, x) == find_root!(s, y)

"""
    union!(s::IntDisjointSet{T}, x::T, y::T)
Merge the subset containing x and that containing y into one
and return the root of the new set.
"""
function Base.union!(s::IntDisjointSet{T}, x::T, y::T) where {T<:Integer}
    parents = s.parents
    xroot = find_root_impl!(parents, x)
    yroot = find_root_impl!(parents, y)
    xroot != yroot ? root_union!(s, xroot, yroot) : xroot
end

"""
    root_union!(s::IntDisjointSet{T}, x::T, y::T)
Form a new set that is the union of the two sets whose root elements are
x and y and return the root of the new set.
Assume x ≠ y (unsafe).
"""
function root_union!(s::IntDisjointSet{T}, x::T, y::T) where {T<:Integer}
    parents = s.parents
    rks = s.ranks
    @inbounds xrank = rks[x]
    @inbounds yrank = rks[y]

    if xrank < yrank
        x, y = y, x
    elseif xrank == yrank
        rks[x] += one(T)
    end
    @inbounds parents[y] = x
    s.ngroups -= one(T)
    return x
end

"""
    push!(s::IntDisjointSet{T})
Make a new subset with an automatically chosen new element x.
Returns the new element. Throw an `ArgumentError` if the
capacity of the set would be exceeded.
"""
function Base.push!(s::IntDisjointSet{T}) where {T<:Integer}
    l = length(s)
    l < typemax(T) || throw(ArgumentError(_intdisjointset_bounds_err_msg(T)))
    x = l + one(T)
    push!(s.parents, x)
    push!(s.ranks, zero(T))
    s.ngroups += one(T)
    return x
end

"""
    DisjointSet{T}(xs)
A forest of disjoint sets of arbitrary value type T.
It is a wrapper of IntDisjointSet{Int}, which uses a
dictionary to map the input value to an internal index.
"""
mutable struct DisjointSet{T} <: AbstractSet{T}
    intmap::Dict{T,Int}
    revmap::Vector{T}
    internal::IntDisjointSet{Int}
    DisjointSet{T}() where T = new{T}(Dict{T,Int}(), Vector{T}(), IntDisjointSet(0))

    function DisjointSet{T}(xs) where T    # xs must be iterable
        imap = Dict{T,Int}()
        rmap = Vector{T}()
        if(typeof(isscalar(typeof(xs)))!=Scalar)
            n = length(xs)
            sizehint!(imap, n)
            sizehint!(rmap, n)
            id = 0
            for x in xs
                imap[x] = (id += 1)
                push!(rmap,x)
            end
            new{T}(imap, rmap, IntDisjointSet(n))
        else
            id = 0
            imap[xs] = (id += 1)
            push!(rmap,xs)
            new{T}(imap, rmap, IntDisjointSet(1))

        end
    end
end

DisjointSet() = DisjointSet{Any}()
DisjointSet(xs) = _DisjointSet(xs, Base.IteratorEltype(xs))
_DisjointSet(xs, ::Base.HasEltype) = DisjointSet{eltype(xs)}(xs)
function _DisjointSet(xs, ::Base.EltypeUnknown)
    T = Base.@default_eltype(xs)
    (isconcretetype(T) || T === Union{}) || return Base.grow_to!(DisjointSet{T}(), xs)
    return DisjointSet{T}(xs)
end

Base.iterate(s::DisjointSet) = iterate(s.revmap)
Base.iterate(s::DisjointSet, i) = iterate(s.revmap, i)

Base.length(s::DisjointSet) = length(s.internal)

"""
    num_groups(s::DisjointSet)
Get a number of groups.
"""
num_groups(s::DisjointSet) = num_groups(s.internal)
Base.eltype(::Type{DisjointSet{T}}) where T = T
Base.empty(s::DisjointSet{T}, ::Type{U}=T) where {T,U} = DisjointSet{U}()
function Base.sizehint!(s::DisjointSet, n::Integer)
    sizehint!(s.intmap, n)
    sizehint!(s.revmap, n)
    return s
end

"""
    find_root!{T}(s::DisjointSet{T}, x::T)
Finds the root element of the subset in `s` which has the element `x` as a member.
"""
find_root!(s::DisjointSet{T}, x::T) where {T} = s.revmap[find_root!(s.internal, s.intmap[x])]

"""
    in_same_set(s::DisjointSet{T}, x::T, y::T)
Returns `true` if `x` and `y` belong to the same subset in `s` and `false` otherwise.
"""
in_same_set(s::DisjointSet{T}, x::T, y::T) where {T} = in_same_set(s.internal, s.intmap[x], s.intmap[y])

"""
    union!(s::DisjointSet{T}, x::T, y::T)
Merge the subset containing x and that containing y into one
and return the root of the new set.
"""
Base.union!(s::DisjointSet{T}, x::T, y::T) where {T} = s.revmap[union!(s.internal, s.intmap[x], s.intmap[y])]

"""
    root_union!(s::DisjointSet{T}, x::T, y::T)
Form a new set that is the union of the two sets whose root elements are
x and y and return the root of the new set.
Assume x ≠ y (unsafe).
"""
root_union!(s::DisjointSet{T}, x::T, y::T) where {T} = s.revmap[root_union!(s.internal, s.intmap[x], s.intmap[y])]

"""

    push!(s::DisjointSet{T}, x::T)
Make a new subset with an automatically chosen new element x.
Returns the new element.

"""
function Base.push!(s::DisjointSet{T}, x::T) where T
    haskey(s.intmap, x) && return x
    id = push!(s.internal)
    s.intmap[x] = id
    push!(s.revmap, x) # Note, this assumes invariant: length(s.revmap) == id
    return x
end