using Revise
using NNFEM
using PyPlot
using ADCME

NT = 100
T = 1.0
Δt = T/NT

m, n =  15, 15
h = 1/m

ν = 0.35
E = 2.0
H = zeros(3,3)
H[1,1] = E*(1. -ν)/((1+ν)*(1. -2. *ν));
H[1,2] = H[1,1]*ν/(1-ν);
H[2,1] = H[1,2];
H[2,2] = H[1,1];
H[3,3] = H[1,1]*0.5*(1. -2. *ν)/(1. -ν);

# Create a very simple mesh
elements = []
prop = Dict("name"=> "PlaneStrain", "rho"=> 1.0, "E"=> 2.0, "nu"=> 0.35)
coords = zeros((m+1)*(n+1), 2)
for j = 1:n
    for i = 1:m
        idx = (m+1)*(j-1)+i
        elnodes = [idx; idx+1; idx+1+m+1; idx+m+1]
        ngp = 3
        nodes = [
            (i-1)*h (j-1)*h
            i*h (j-1)*h
            i*h j*h
            (i-1)*h j*h
        ]
        coords[elnodes, :] = nodes
        push!(elements, FiniteStrainContinuum(nodes, elnodes, prop, ngp))
    end
end

# fixed on the bottom, push on the right
EBC = zeros(Int64, (m+1)*(n+1), 2)
FBC = zeros(Int64, (m+1)*(n+1), 2)
g = zeros((m+1)*(n+1), 2)
f = zeros((m+1)*(n+1), 2)

for i = 2:m
    idx = n*(m+1)+i
    EBC[idx,:] .= -2
    EBC[i,:] .= -2
end

for j = 1:n+1
    idx = (j-1)*(m+1)+1
    EBC[idx,:] .= -2
    idx = (j-1)*(m+1)+m+1
    EBC[idx,:] .= -2
end

ndims = 2
domain = Domain(coords, elements, ndims, EBC, g, FBC, f)

EBC_idx = findall(EBC[:,1] .== -2)
EBC_X = domain.nodes[EBC_idx,:]
function gt(t)
    x, y = EBC_X[:,1], EBC_X[:,2]
    u1 = @. (x^2+y^2)*exp(-t) * 0.1
    u2 = @. (x^2-y^2)*exp(-t) * 0.1
    return [u1;u2], -[u1;u2], [u1;u2]
end

FBC_idx = findall(FBC[:,1] .== -2)
FBC_X = domain.nodes[FBC_idx,:]
# function ft(t)
#     x, y = FBC_X[:,1], FBC_X[:,2]
#     T1 = @. -1.48148148148148*x*exp(-t) - 1.48148148148148*y*exp(-t)
#     T2 = @. -3.45679012345679*x*exp(-t) - 6.41975308641975*y*exp(-t)
#     return [T1;T2] * h * 0.1
# end

function bt(x, y, t)
    f1 = @. 0.1*(x^2 + y^2)*exp(-t) - 0.790123456790123*exp(-t)
    f2 = @. 0.1*(x^2 - y^2)*exp(-t) + 0.493827160493827*exp(-t)
    [f1 f2]
end

c = domain.nodes[domain.dof_to_eq]
nn = div(length(c), 2)
x = c[1:nn]
y = c[nn+1:end]

Dstate =  @. [x^2+y^2;x^2-y^2]* 0.1
state = @. [x^2+y^2;x^2-y^2]* 0.1
velo =  @. -[x^2+y^2;x^2-y^2]* 0.1
acce =  @. [x^2+y^2;x^2-y^2]* 0.1
globdat = GlobalData(state, Dstate, velo, acce, domain.neqs, gt, nothing, bt)

assembleMassMatrix!(globdat, domain)
# updateTimeDependentEssentialBoundaryCondition!(domain,globdat)
# updateStates!(domain, globdat)

# for i = 1:NT
#     @info i
#     # global globdat, domain = GeneralizedAlphaSolverStep(globdat, domain, Δt)
#     global globdat, domain = ExplicitSolverStep(globdat, domain, Δt)
# end

# d_ = hcat(domain.history["state"]...)'
# i = 3
# j = 4
# plot(d_[:,(j-1)*(m+1)+i])
# x0 = (i-1)*h
# y0 = (j-1)*h
# ts = LinRange(0, 1, NT+1)
# plot((@. (x0^2+y0^2)*exp(-ts))*0.1,"--")

# error()
Hs = zeros(domain.neles*length(domain.elements[1].weights), 3, 3)
for i = 1:size(Hs,1)
    Hs[i,:,:] = H
end

# ts = GeneralizedAlphaSolverTime(Δt, NT)
ts = ExplicitSolverTime(Δt, NT)

ubd = zeros(NT,sum(domain.EBC[:].!=0))
abd = zeros(NT,sum(domain.EBC[:].!=0))
Fext = zeros(NT,domain.neqs)

idx = findall(domain.EBC[:,1].!=0)
x = domain.nodes[idx,1]
y = domain.nodes[idx,2]
for i = 1:NT
    t = ts[i]
    ubd[i,:] = [
        @. (x^2+y^2)*0.1*exp(-t)
        @. (x^2-y^2)*0.1*exp(-t)
    ]
    abd[i,:] = [
        @. (x^2+y^2)*0.1*exp(-t)
        @. (x^2-y^2)*0.1*exp(-t)
    ]
    globdat.time = t
    Fext[i,:] = getBodyForce(domain, globdat)
end

x = domain.nodes[:,1]
y = domain.nodes[:,2]
a0 = @. [x^2+y^2;x^2-y^2]* 0.1
v0 =  @. -[x^2+y^2;x^2-y^2]* 0.1
d0 =  @. [x^2+y^2;x^2-y^2]* 0.1

H = constant(Hs)
# d, v, a= GeneralizedAlphaSolver(globdat, domain, d0, v0, a0, Δt, NT, H, Fext, ubd, abd)
d, v, a= ExplicitSolver(globdat, domain, d0, v0, a0, Δt, NT, H, Fext, ubd, abd; strain_type="small")

sess = Session(); init(sess)
d_, v_, a_ = run(sess, [d,v,a])

for i = 1:5
    i = rand(1:m+1)
    j = rand(1:n+1)
    plot(d_[:,(j-1)*(m+1)+i], color = "C$i")
    x0 = (i-1)*h
    y0 = (j-1)*h
    plot((@. (x0^2+y0^2)*exp(-ts))*0.1,"--", color="C$i")
end


using Test, SummationByPartsOperators
using LinearAlgebra

function accuracy_test!(res, ufunc, dufunc, D)
    u = compute_coefficients(ufunc, D)
    du = compute_coefficients(dufunc, D)
    mul!(res, D, u)
    maximum(abs, du-res) < 5*length(res)*eps(eltype(res))
end

# Accuracy Tests
for T in (Float32, Float64)
    xmin = -one(T)
    xmax = one(T)

    for N in 2 .^ (3:6)
        D = fourier_derivative_operator(xmin, xmax, N)
        println(devnull, D)
        @test SummationByPartsOperators.derivative_order(D) == 1
        @test issymmetric(D) == false
        u = compute_coefficients(zero, D)
        res = D*u
        for k in 0:(N÷2)-1
            ufunc = x->sinpi(k*x)
            dufunc = x->typeof(x)(k*π)*cospi(k*x)
            @test accuracy_test!(res, ufunc, dufunc, D)
            xplot, duplot = evaluate_coefficients(res, D)
            @test maximum(abs, duplot - dufunc.(xplot)) < 5N*eps(T)
            @test abs(integrate(u, D)) < N*eps(T)

            ufunc = x->cospi(k*x)
            dufunc = x->-typeof(x)(k*π)*sinpi(k*x)
            @test accuracy_test!(res, ufunc, dufunc, D)
            xplot, duplot = evaluate_coefficients(res, D)
            @test maximum(abs, duplot - dufunc.(xplot)) < 5N*eps(T)
            @test abs(integrate(u, D)) < N*eps(T)
        end
    end
end

# (Super) Spectral Viscosity
source_SV = (Tadmor1989(), MadayTadmor1989(), TadmorWaagan2012Standard(), TadmorWaagan2012Convergent())
source_SSV = (Tadmor1993(),)

for T in (Float32, Float64), source in source_SV
    xmin = -one(T)
    xmax = one(T)

    for N in 2 .^ (3:6)
        D = fourier_derivative_operator(xmin, xmax, N)
        println(devnull, D)
        @test issymmetric(D) == false

        Di = dissipation_operator(source, D)
        println(devnull, Di)
        @test issymmetric(Di) == true
        Di_full = Matrix(Di)
        @test maximum(abs, Di_full-Di_full') < 80*eps(T)

        @test maximum(eigvals(Symmetric(Di_full))) < 10N*eps(T)
    end
end

for T in (Float32, Float64), source in source_SSV
    xmin = -one(T)
    xmax = one(T)

    for N in 2 .^ (3:6), order in 1:3
        D = fourier_derivative_operator(xmin, xmax, N)
        println(devnull, D)
        @test issymmetric(D) == false

        Di = dissipation_operator(source, D, order=order)
        println(devnull, Di)
        @test issymmetric(Di) == true
        Di_full = Matrix(Di)
        @test maximum(abs, Di_full-Di_full') < 80*eps(T)

        @test maximum(eigvals(Symmetric(Di_full))) < 15N*eps(T)
    end
end

# Modal Filtering
for T in (Float32, Float64), filter_type in (ExponentialFilter(),)
    xmin = -one(T)
    xmax = one(T)

    for N in 2 .^ (1:4)
        D = fourier_derivative_operator(xmin, xmax, N)
        filter! = ConstantFilter(D, filter_type)
        u = compute_coefficients(zero, D)
        res = D*u
        for k in 1:N-1
            compute_coefficients!(u, x->exp(sinpi(x)), D)
            norm2_u = integrate(u->u^2, u, D)
            filter!(u)
            @test integrate(u->u^2, u, D) <= norm2_u
        end
    end
end


using Distributed, StatsBase, Distributions, PyPlot
using TransitionalMCMC

lb  = -15
ub  = 15

fT(x) = logpdf(Uniform(lb,ub), x[1])
sample_fT(Nsamples) = rand(Uniform(lb,ub), Nsamples, 1)

log_fD_T(x) = log(pdf(Normal(0, 1), x[1]) + pdf(Normal(5, 0.2), x[1]))

Nsamples = 2000
samps, acc = tmcmc(log_fD_T, fT, sample_fT, Nsamples)

plt.hist(samps,50)


"""
    absolute_distance(c₁::AnyOptimalityCut, c₂::AnyOptimalityCut)

Absolute distance between two optimality cuts

"""
function absolute_distance(c₁::AnyOptimalityCut, c₂::AnyOptimalityCut)
    c̃₁ = vcat(c₁.δQ, c₁.q) ./ num_subproblems(c₁)
    c̃₂ = vcat(c₂.δQ, c₂.q) ./ num_subproblems(c₂)
    return norm(c̃₁-c̃₂) / max(norm(c̃₁), norm(c̃₂))
end
function absolute_distance(c₁::AnyOptimalityCut{T}, c₂::AggregatedOptimalityCut{T,ZeroVector{T}}) where T <: AbstractFloat
    return Inf
end
function absolute_distance(c₁::AggregatedOptimalityCut{T,ZeroVector{T}}, c₂::AnyOptimalityCut{T}) where T <: AbstractFloat
    return Inf
end

"""
    angular_distance(c₁::AnyOptimalityCut, c₂::AnyOptimalityCut)

Angular distance between two optimality cuts

"""
function angular_distance(c₁::AnyOptimalityCut, c₂::AnyOptimalityCut)
    return 1-abs(c₁.δQ ⋅ c₂.δQ) / (norm(c₁.δQ) * norm(c₂.δQ))
end
function angular_distance(c₁::AnyOptimalityCut{T}, c₂::AggregatedOptimalityCut{T,ZeroVector{T}}) where T <: AbstractFloat
    return Inf
end
function angular_distance(c₁::AggregatedOptimalityCut{T,ZeroVector{T}}, c₂::AnyOptimalityCut{T}) where T <: AbstractFloat
    return Inf
end

"""
    spatioangular_distance(c₁::AnyOptimalityCut, c₂::AnyOptimalityCut)

Spatioangular distance between two optimality cuts.

"""
function spatioangular_distance(c₁::AnyOptimalityCut, c₂::AnyOptimalityCut)
    return 1-abs(c₁.δQ ⋅ c₂.δQ) / (norm(c₁.δQ) * norm(c₂.δQ)) +
        (abs(c₁.q / num_subproblems(c₁) - c₂.q / num_subproblems(c₂)) /
        max(abs(c₁.q / num_subproblems(c₁)), abs(c₂.q / num_subproblems(c₂))))
end
function spatioangular_distance(c₁::AnyOptimalityCut{T}, c₂::AggregatedOptimalityCut{T,ZeroVector{T}}) where T <: AbstractFloat
    return Inf
end
function spatioangular_distance(c₁::AggregatedOptimalityCut{T,ZeroVector{T}}, c₂::AnyOptimalityCut{T}) where T <: AbstractFloat
    return Inf
end


import MittagLefflerFunctions.MLneg_integral
using Printf
using PyPlot

T = BigFloat
lenA = 20
A = Vector{T}(undef, lenA)
α = parse(T, "0.6")
β = one(T)
t = parse(T, "0.1")
λ = one(T)

for N = 1:lenA
    A[N] = MLneg_integral(α, β, λ, t, N)
    if N ≥ 2
        @printf("%5d  %14.4e\n", N, A[N]-A[N-1])
    end
end

figure(1)
semilogy(collect(1:lenA-5), abs.(A[1:end-5].-A[end]), "o")
grid(true)


# protein type that keeps centered MD data for a protein
# 2014 - murozturk@ku.edu.tr

immutable protein
	R::Matrix{Float64}
	index::Vector{Int}
	means::Vector{Float64}
	stdPerRes::Vector{Float64}
	covariance::Matrix{Float64}

	function protein(R::Matrix{Float64}, index=1:size(R,2))
		m = mean(R,1)
		R .-= m

		println ("Calculating covariance matrix...")
		gamma = cov(R, mean=0, corrected=false)

		println ("Calculating real fluctuations...")
		rF = realFluc(R, centered=true)

		new(R, index[:], m[:], std(rF,1)[:], gamma  )
	end
end

# Returns decentered coordinates
# recovers original md data

deCenter(p::protein) = p->p.R .+ p.means'

# Returns real fluctuations per residue

function realFluc(R::Matrix{Float64}; centered=false)
	!centered && (R .-= mean(R,1))
	sqrt(R[:,1:3:end].^2 + R[:,2:3:end].^2 + R[:,3:3:end].^2)
end
realFluc(p::protein) = realFluc(p.R, centered=true)


#
# @description Implements the DES (Data Encryption Standard) block cipher algorithm.
#
# @author Alejandro Claro (alejandro.claro@gmail.com)
#
# Copyright 2015 All rights reserved.
# Use of this source code is governed by a MIT-style license that can be found in the LICENSE file.
#
export DesCipher
export DES_BLOCK_SIZE, DES_LEGAL_KEY_SIZES
export block_size, key_size, encrypt, decrypt

# @description Defines the DES cipher data struture.
struct DesCipher <: BlockCipher
  subkeys::Vector{Vector{UInt8}}

  # @description Initialize the DES cipher.
  #
  # @param {Vector{UInt8}} key The symmetric key.
  function DesCipher(key::Vector{UInt8})
    return new(compute_des_subkeys(key))
  end
end

# @description The DES cipher block size in bytes.
const DES_BLOCK_SIZE = UInt64(8)

# @description The DES cipher legal key sizes in bytes.
const DES_LEGAL_KEY_SIZES = UInt64[ 8 ]

# @description Gets the cipher block size.
#
# @param {DesCipher} self The cipher data struture.
#
# @return {Integer} The block size in bytes.
function block_size(self::DesCipher)
  return DES_BLOCK_SIZE
end

# @description Gets the cipher key size.
#
# @param {DesCipher} self The cipher data struture.
#
# @return {Integer} The key size in bytes.
function key_size(self::DesCipher)
  return DES_LEGAL_KEY_SIZES[1]
end

# @description Encrypts the given plaintext block with the key set at initialization.
#
# @param {DesCipher}     self      The cipher data struture.
# @param {Vector{UInt8}} plaintext The message to encrypt.
#
# @return {Vector{UInt8}} The resulting ciphertext block.
function encrypt(self::DesCipher, plaintext::Vector{UInt8})
  if length(plaintext) != block_size(self)
    throw(ArgumentError("Invalid input plaintext length."))
  end

  feistel    = FeistelCipher(compute_des_round, self.subkeys, key_size(self), block_size(self))
  block      = permutate_bits(plaintext, DES_INITIAL_PERMUTATION)
  block      = encrypt(feistel, block)
  ciphertext = permutate_bits(block, DES_FINAL_PERMUTATION)

  return ciphertext
end

# @description Decrypts the given ciphertext block with the key set at initialization.
#
# @param {DesCipher}     self       The cipher data struture.
# @param {Vector{UInt8}} ciphertext The message to decrypt.
#
# @return {Vector{UInt8}} The resulting plaintext block.
function decrypt(self::DesCipher, ciphertext::Vector{UInt8})
  if length(ciphertext) != block_size(self)
    throw(ArgumentError("Invalid input ciphertext length."))
  end

  feistel   = FeistelCipher(compute_des_round, self.subkeys, key_size(self), block_size(self))
  block     = permutate_bits(ciphertext, DES_INITIAL_PERMUTATION)
  block     = decrypt(feistel, block)
  plaintext = permutate_bits(block, DES_FINAL_PERMUTATION)

  return plaintext
end

# PRIVATE IMPLEMENTATION #######################################################

# @description First permuted choice, used in the key schedule to select 56 bits from a 64-bit input.
const DES_PC1 = UInt8[
  57,   49,    41,   33,    25,    17,    9,
   1,   58,    50,   42,    34,    26,   18,
  10,    2,    59,   51,    43,    35,   27,
  19,   11,     3,   60,    52,    44,   36,
  63,   55,    47,   39,    31,    23,   15,
   7,   62,    54,   46,    38,    30,   22,
  14,    6,    61,   53,    45,    37,   29,
  21,   13,     5,   28,    20,    12,    4
]

# @description Schedule of "left shifts" for subkeys construction.
const DES_LEFT_SHIFTS = [ 1, 1, 2, 2, 2, 2, 2, 2, 1, 2, 2, 2, 2, 2, 2, 1 ]

# @description Second permuted choice, used in the key schedule to produce each subkey.
const DES_PC2 = UInt8[
  14,    17,   11,    24,     1,    5,
   3,    28,   15,     6,    21,   10,
  23,    19,   12,     4,    26,    8,
  16,     7,   27,    20,    13,    2,
  41,    52,   31,    37,    47,   55,
  30,    40,   51,    45,    33,   48,
  44,    49,   39,    56,    34,   53,
  46,    42,   50,    36,    29,   32
]

# @description Initial permutation for the 64-bit input block.
const DES_INITIAL_PERMUTATION = UInt8[
  58,    50,   42,    34,    26,   18,    10,    2,
  60,    52,   44,    36,    28,   20,    12,    4,
  62,    54,   46,    38,    30,   22,    14,    6,
  64,    56,   48,    40,    32,   24,    16,    8,
  57,    49,   41,    33,    25,   17,     9,    1,
  59,    51,   43,    35,    27,   19,    11,    3,
  61,    53,   45,    37,    29,   21,    13,    5,
  63,    55,   47,    39,    31,   23,    15,    7
]

# @description Expansion of input block of 32 bits, producing an output block of 48 bits.
const DES_EXPANSION_FUNCTION = UInt8[
  32,     1,    2,     3,     4,    5,
   4,     5,    6,     7,     8,    9,
   8,     9,   10,    11,    12,   13,
  12,    13,   14,    15,    16,   17,
  16,    17,   18,    19,    20,   21,
  20,    21,   22,    23,    24,   25,
  24,    25,   26,    27,    28,   29,
  28,    29,   30,    31,    32,    1
]

# @description The DES' S-box 1
const DES_SBOX1 = UInt8[
  14   4  13   1   2  15  11   8   3  10   6  12   5   9   0   7;
   0  15   7   4  14   2  13   1  10   6  12  11   9   5   3   8;
   4   1  14   8  13   6   2  11  15  12   9   7   3  10   5   0;
  15  12   8   2   4   9   1   7   5  11   3  14  10   0   6  13
]

# @description The DES' S-box 2
const DES_SBOX2 = UInt8[
  15   1   8  14   6  11   3   4   9   7   2  13  12   0   5  10;
   3  13   4   7  15   2   8  14  12   0   1  10   6   9  11   5;
   0  14   7  11  10   4  13   1   5   8  12   6   9   3   2  15;
  13   8  10   1   3  15   4   2  11   6   7  12   0   5  14   9
]

# @description The DES' S-box 3
const DES_SBOX3 = UInt8[
  10   0   9  14   6   3  15   5   1  13  12   7  11   4   2   8;
  13   7   0   9   3   4   6  10   2   8   5  14  12  11  15   1;
  13   6   4   9   8  15   3   0  11   1   2  12   5  10  14   7;
   1  10  13   0   6   9   8   7   4  15  14   3  11   5   2  12
]

# @description The DES' S-box 4
const DES_SBOX4 = UInt8[
   7  13  14   3   0   6   9  10   1   2   8   5  11  12   4  15;
  13   8  11   5   6  15   0   3   4   7   2  12   1  10  14   9;
  10   6   9   0  12  11   7  13  15   1   3  14   5   2   8   4;
   3  15   0   6  10   1  13   8   9   4   5  11  12   7   2  14
]

# @description The DES' S-box 5
const DES_SBOX5 = UInt8[
   2  12   4   1   7  10  11   6   8   5   3  15  13   0  14   9;
  14  11   2  12   4   7  13   1   5   0  15  10   3   9   8   6;
   4   2   1  11  10  13   7   8  15   9  12   5   6   3   0  14;
  11   8  12   7   1  14   2  13   6  15   0   9  10   4   5   3
]

# @description The DES' S-box 6
const DES_SBOX6 = UInt8[
  12   1  10  15   9   2   6   8   0  13   3   4  14   7   5  11;
  10  15   4   2   7  12   9   5   6   1  13  14   0  11   3   8;
   9  14  15   5   2   8  12   3   7   0   4  10   1  13  11   6;
   4   3   2  12   9   5  15  10  11  14   1   7   6   0   8  13
]

# @description The DES' S-box 7
const DES_SBOX7 = UInt8[
   4  11   2  14  15   0   8  13   3  12   9   7   5  10   6   1;
  13   0  11   7   4   9   1  10  14   3   5  12   2  15   8   6;
   1   4  11  13  12   3   7  14  10  15   6   8   0   5   9   2;
   6  11  13   8   1   4  10   7   9   5   0  15  14   2   3  12
]

# @description The DES' S-box 8
const DES_SBOX8 = UInt8[
  13   2   8   4   6  15  11   1  10   9   3  14   5   0  12   7;
   1  15  13   8  10   3   7   4  12   5   6  11   0  14   9   2;
   7  11   4   1   9  12  14   2   0   6  10  13  15   3   5   8;
   2   1  14   7   4  10   8  13  15  12   9   0   3   5   6  11
]

# @description Collection of DES' S-boxes.
const DES_SBOXES = [ DES_SBOX1, DES_SBOX2, DES_SBOX3, DES_SBOX4, DES_SBOX5, DES_SBOX6, DES_SBOX7, DES_SBOX8 ]

# @description Predefined permutation for S-boxes output.
const DES_SBOX_PERMUTATION = UInt8[
  16,     7,    20,    21,
  29,    12,    28,    17,
   1,    15,    23,    26,
   5,    18,    31,    10,
   2,     8,    24,    14,
  32,    27,     3,     9,
  19,    13,    30,     6,
  22,    11,     4,    25
]

# @description Final permutation for the 4-bit preoutput block.
const DES_FINAL_PERMUTATION = UInt8[
  40,     8,   48,    16,    56,   24,    64,   32,
  39,     7,   47,    15,    55,   23,    63,   31,
  38,     6,   46,    14,    54,   22,    62,   30,
  37,     5,   45,    13,    53,   21,    61,   29,
  36,     4,   44,    12,    52,   20,    60,   28,
  35,     3,   43,    11,    51,   19,    59,   27,
  34,     2,   42,    10,    50,   18,    58,   26,
  33,     1,   41,     9,    49,   17,    57,   25
]

# @description Computes the subkey for DES algorithm.
#
# @param key The symmetric key.
#
# @return {Vector{Vector{UInt8}}} The set of subkeys (16 subkeys).
function compute_des_subkeys(key::Vector{UInt8})
  if length(key) ∉ DES_LEGAL_KEY_SIZES
    throw(ArgumentError("Invalid key length. $(length(key) * 8) bits provided, but 64 bits expected."))
  end

  result = Vector{Vector{UInt8}}()
  key    = permutate(unpack_bits(key), DES_PC1)
  cn     = key[1 : 28]
  dn     = key[29 : end]

  for i in 1 : 16
    for j in 1:DES_LEFT_SHIFTS[i]
      cn = circshift(cn, -1)
      dn = circshift(dn, -1)
    end

    subkey = permutate(vcat(cn, dn), DES_PC2)
    push!(result, subkey)
  end

  return result
end

# @description Computes the DES' f(*) function for the given R and subkey.
#
# @param {Vector{UInt8}} r   The R (right) bytes.
# @param {Vector{UInt8}} key The subkey for the DES iteration.
#
# @return {Vector{UInt8}} The round function value.
function compute_des_round(r::Vector{UInt8}, key::Vector{UInt8})
  r = permutate(r, DES_EXPANSION_FUNCTION)
  r = map(xor, r, key)

  b = [ r[1 : 6], r[7 : 12], r[13 : 18], r[19 : 24], r[25 : 30], r[31 : 36], r[37 : 42], r[43 : end] ]

  result = zeros(UInt8, 32)
  index  = 0

  for j in 1 : length(DES_SBOXES)
    m, n  = compute_permutation_position(b[j])
    value = DES_SBOXES[j][m + 1, n + 1]

    result[index + 1] = (value & 8) >> 3
    result[index + 2] = (value & 4) >> 2
    result[index + 3] = (value & 2) >> 1
    result[index + 4] = (value & 1) >> 0
    index += 4
  end

  return permutate(result, DES_SBOX_PERMUTATION)
end

# @description Computes the permutation position in the s-box(n).
#
# @param Vector{UInt8} bn The s-box indices.
#
# @return {Integer, Intenger} The permutation position.
function compute_permutation_position(bn::Vector{UInt8})
  m = (bn[1] << 1) + bn[6]
  n = (bn[2] << 3) + (bn[3] << 2) + (bn[4] << 1) + bn[5]

  return m, n
end


include("findLargestInCol.jl")
include("swapRows.jl")

# we write a function that decomposes a square matrix A as A = P^{-1}LU
function PALUdecomp(A)

    # we get the size of A
    m, n = size(A)

    # we check that A is indeed a square matrix
    if m != n
        error("Input must be a square matrix")
    end

    # we initialize P to the indentity and L and U to zero
    P = one(A)
    L = zeros(n, n)
    U = zeros(n, n)

    # we make a copy of A to work with
    B = copy(A)

    for k = 1:n

        # we find the pivot
        maxk = findLargestInCol(B, k)
        if maxk != k
            B = swapRows(B, maxk, k)
            L = swapRows(L, maxk, k)
            P = swapRows(P, maxk, k)
        end

        # we proceed with the normal LU decomposition algorithm
        for j = k:n
            U[k,j] = B[k,j]
        end

        # we check that we don't divide by zero
        if U[k,k] == 0
            error("** A^(k-1)_{k,k} == 0 in PALU decomp")
        end

        for i = k:n
           L[i,k] = B[i,k] / U[k,k]
        end

        for i = k:n
            for j = k:n
                B[i,j] -= L[i,k] * U[k,j];
            end
        end
    end

    # we return P, L and U
    return P, L, U
end


using Test
using Random
using Statistics
using AdaptiveMCMC

Random.seed!(1234)

# Standard normal test target in 2d
log_p(x) = -(x[1]^2+x[2]^2)/2
x0 = zeros(2)
n = 100_000
test_means(X, atol=0.05) = all(abs.(mapslices(mean, X, dims=2)) .<= atol)
test_vars(X, atol=0.05) = all(abs.(mapslices(var, X, dims=2) .- 1.0) .<= atol)
test_stats(X, atol=0.05) = test_vars(X, atol) && test_means(X, atol)
test_acc(acc, acc_opt=0.234, atol=0.05) = all(abs.(acc .- acc_opt) .<= atol)

o_ram = adaptive_rwm(x0, log_p, n; algorithm=:ram)
@test test_acc(o_ram.accRWM)
@test test_stats(o_ram.X)
o_ram_32 = adaptive_rwm(Float32.(x0), log_p, n; algorithm=:ram)
@test test_acc(o_ram_32.accRWM)
@test test_stats(o_ram_32.X)

o_am = adaptive_rwm(x0, log_p, n; algorithm=:am)
@test o_am.S[1].L.L[1] ≈ 1.0 atol=0.05
@test test_stats(o_am.X)
o_am_32 = adaptive_rwm(x0, log_p, n; algorithm=:am)
@test o_am_32.S[1].L.L[1] ≈ 1.0 atol=0.05
@test test_stats(o_am_32.X)

o_asm = adaptive_rwm(x0, log_p, n; algorithm=:asm)
@test test_acc(o_asm.accRWM)
@test test_stats(o_asm.X)
o_asm_32 = adaptive_rwm(Float32.(x0), log_p, n; algorithm=:asm)
@test test_acc(o_asm_32.accRWM)
@test test_stats(o_asm_32.X)

o_aswam = adaptive_rwm(x0, log_p, n; algorithm=:aswam)
@test o_aswam.S[1].AM.L.L[1] ≈ 1.0 atol=0.05
@test test_acc(o_aswam.accRWM)
@test test_stats(o_aswam.X)
o_aswam_32 = adaptive_rwm(Float32.(x0), log_p, n; algorithm=:aswam)
@test o_aswam_32.S[1].AM.L.L[1] ≈ 1.0 atol=0.10
@test test_acc(o_aswam_32.accRWM)
@test test_stats(o_aswam_32.X)

o_apt = adaptive_rwm(x0, log_p, n; L=3, all_levels=true)
@test test_acc(o_apt.accRWM)
@test test_acc(o_apt.accSW)
@test test_stats(o_apt.X)
@test test_means(o_apt.allX[2], 0.10)
@test test_means(o_apt.allX[3], 0.15)
o_apt_32 = adaptive_rwm(Float32.(x0), log_p, n; L=3, all_levels=true)
@test test_acc(o_apt_32.accRWM)
@test test_acc(o_apt_32.accSW)
@test test_stats(o_apt_32.X)
@test test_means(o_apt_32.allX[2], 0.10)
@test test_means(o_apt_32.allX[3], 0.15)


module SmolyakInterpolation

export
    # chebyshev
    Chebyshev, set_length, set_range, gridpoint_set, evaluate, evaluate!,
    # basis
    HomogeneousBasis, degrees_of_freedom, basis_matrix_and_coordinates, interpolate,
    interpolated_basis!, interpolated_basis

using ArgCheck: @argcheck
using DocStringExtensions: FIELDS, SIGNATURES, TYPEDEF
using UnPack: @unpack
using StaticArrays: SVector

include("utilities.jl")
include("chebyshev.jl")
include("basis.jl")

end # module


@testset ExtendedTestSet "squared_euclidean" begin
    @testset ExtendedTestSet "squared_euclidean" begin
        x = CUDA.CuArray(
            boolean_indicator(
                [
                    0 1 1 1 0
                    1 1 1 1 1
                    1 0 0 0 1
                    1 0 0 0 1
                    1 0 0 0 1
                    1 1 1 1 1
                    0 1 1 1 0
                ]
            ),
        )
        dt = CuArray{Float32}(undef, size(x))
        v = CUDA.ones(Int64, size(x))
        z = CUDA.zeros(Float32, size(x) .+ 1)
        tfm = DistanceTransforms.SquaredEuclidean(x, dt, v, z)
        answer = CUDA.CuArray(
            [
                1.0 0.0 0.0 0.0 1.0
                0.0 0.0 0.0 0.0 0.0
                0.0 1.0 1.0 1.0 0.0
                0.0 4.0 4.0 4.0 0.0
                0.0 1.0 1.0 1.0 0.0
                0.0 0.0 0.0 0.0 0.0
                1.0 0.0 0.0 0.0 1.0
            ],
        )
        @test transform!(x, tfm) == answer
    end
end


InverseFunctions.inverse(::typeof(log1pexp)) = logexpm1
InverseFunctions.inverse(::typeof(logexpm1)) = log1pexp

InverseFunctions.inverse(::typeof(log1mexp)) = log1mexp

InverseFunctions.inverse(::typeof(log2mexp)) = log2mexp

InverseFunctions.inverse(::typeof(logit)) = logistic
InverseFunctions.inverse(::typeof(logistic)) = logit


# Example for basic 3D modeling:
# The receiver positions and the source wavelets are the same for each of the four experiments.
# Author: Philipp Witte, pwitte@eos.ubc.ca
# Date: January 2017
#

using PyCall, PyPlot, opesciSLIM.TimeModeling

## Set up model structure
n = (120,100,90)	# (x,y,z) or (x,z)
d = (10.,10.,10.)
o = (0.,0.,0.)

# Velocity [km/s]
v = ones(Float32,n) + 0.4f0
v0 = ones(Float32,n) + 0.4f0
v[:,:, Int(round(end/2)):end] = 4.0f0
rho = ones(Float32,n)

# Slowness squared [s^2/km^2]
m = (1./v).^2
m0 = (1./v0).^2
dm = vec(m - m0)

# Setup info and model structure
nsrc = 4
model = Model(n,d,o,m)	# to include density call Model(n,d,o,m,rho)
model0 = Model(n,d,o,m0)

## Set up 3D receiver geometry by defining one receiver vector in each x and y direction
nxrec = 120
nyrec = 100
xrec = linspace(50.,1150.,nxrec)
yrec = linspace(100., 900., nyrec)
zrec = 50

# Construct 3D grid from basis vectors
(xrec, yrec, zrec) = setup_3D_grid(xrec, yrec, zrec)

# receiver sampling and recording time
timeR = 1000.	# receiver recording time [ms]
dtR = 4.	# receiver sampling interval

# Set up receiver structure
recGeometry = Geometry(xrec,yrec,zrec;dt=dtR,t=timeR,nsrc=nsrc)

## Set up source geometry (cell array with source locations for each shot)
xsrc = convertToCell([250. 500. 750. 1000.])
ysrc = convertToCell([200. 400. 600. 800.])
zsrc = convertToCell([50. 60. 70. 80.])

# source sampling and number of time steps
timeS = 1000.	# source length in [ms]
dtS = 2.	# source sampling interval

# Set up source structure
srcGeometry = Geometry(xsrc,ysrc,zsrc;dt=dtS,t=timeS)

# setup wavelet
f0 = 0.01
wavelet = source(timeS,dtS,f0)

# Info structure for linear operators
ntComp = get_computational_nt(srcGeometry,recGeometry,model)	# no. of computational time steps
info = Info(prod(n),nsrc,ntComp)

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

# Setup operators
Pr = joProjection(info,recGeometry)
F = joModeling(info,model)
F0 = joModeling(info,model0)
Ps = joProjection(info,srcGeometry)
q = joData(srcGeometry,wavelet)

# Combined operator Pr*F*Ps'
Ffull = joModeling(info,model,srcGeometry,recGeometry)

# Nonlinear modeling
d = Pr*F*Ps'*q	# equivalent to d = Ffull*q
qad = Ps*F'*Pr'*d

# Linearized modeling
J = joJacobian(Pr*F0*Ps',q)	# equivalent to J = joSetupJacobian(Ffull,q)
dD = J*dm
rtm = J'*dD

# Plot results
figure(); imshow(d.data[1][:,1:500], vmin=-1e-1, vmax=1e-1)
rtm = reshape(rtm,model.n)
figure(); imshow(rtm[:,50,:]', ColorMap("gray"))


using Test
using Jecco
using Jecco.AdS5_3_1

sigmavars   = tuple(ones(58)...)

axx, ayy, axy, bx, by, cc= AdS5_3_1.AH_eq_coeff(sigmavars, AdS5_3_1.Outer())

S = AdS5_3_1.AH_eq_res(sigmavars, AdS5_3_1.Outer())

print("(axx, ayy, axy, bx, by, cc, S) =",(axx, ayy, axy, bx, by, cc, S),"\n")


using DifferentialEquations, LinearAlgebra
using Plots
gr()

const N0 =  2.5000013629442974e25
const M0 = 1.5264213305629408e16
# number of air molecules inside the fiber is 1.5264213305629408e16
#=
ozone from
http://www.columbia.edu/itc/chemistry/chem-c2407/hw/ozone_kinetics.pdf
equations 5 & 6
=#

function gaussianPulse(;τ= τ, FWHM=FWHM, t0=t0)
    res0 = sqrt.((exp.(-((τ.-t0).^2.0) ./ (FWHM^2.0))))
end

function createTimeGrid(time_window; npts = 2^14)
    dt = time_window/npts
    τ = LinRange(0, time_window, npts)
    τ = collect(τ)
end

t=createTimeGrid(200)
FWHM=10
gauss(τ) = gaussianPulse(τ= τ, FWHM=FWHM, t0=10) .+ gaussianPulse(τ= τ, FWHM=FWHM, t0=60) .+ gaussianPulse(τ= τ, FWHM=FWHM, t0=110) .+
            gaussianPulse(τ= τ, FWHM=FWHM, t0=160)
# plot(t, gauss.(t))

function fun(du,u,p,t)
    # k1,k2,k3,k4 = p
    k1 = 1.03E-14*gauss(t)
    k3 = 8.58E-2*gauss(t)

    # k1 = 1.03E-14
    k2 = 5.60E-34
    # k3 = 8.58E-2
    # k4 = 8.34E-15
    k4=8.34E-19#sice the reactions are slow compared to the time scale of the ultrafast pulses
    # A = N0*0.50*1e-6*1e-3
    M = M0
    # print(t)
    #=
    u[1] = O, u[2] = O3, u[3] = O2, du[4] = dM = 0
    =#
    du[1] = 2*k1*u[3]-k2*u[1]*u[3]*M+k3*u[2]-k4*u[1]*u[2]
    du[2] = k2*u[1]*u[3]*M - k3*u[2]-k4*u[1]*u[2]
    du[3] = -k1*u[3] - k2*u[3]*M*u[1] + k3*u[2] + 2*k4*u[1]*u[2]
end

# k1 = 1.03E-14
# k2 = 5.60E-34
# k3 = 8.58E-2
# k4 = 8.34E-15
# # A = N0*0.50*1e-6*1e-3
# M = N0*0.80*1e-6*1e-3
p = []

tspan = (0.0, 200)
u0 = [1e7,1e7, N0*0.20*1e-6*1e-3]

prob1 = ODEProblem(fun, u0, tspan, p)
sol = solve(prob1, Tsit5(),reltol = 1e-10, abstol=1e-16)

plot(sol,vars=(0,1))
plot!(sol,vars=(0,2))
# plot!(sol,vars=(0,3))


# basic test of multiplying big and small together should give small

using ApproxManifoldProducts
using Test

@testset "multiply big and small" begin

N = 100

X1 = randn(3,N)
X2 = 0.1.*randn(3,N)

p = manikde!(X1, (:Euclid, :Euclid, :Circular))
q = manikde!(X2, (:Euclid, :Euclid, :Circular))

pq = manifoldProduct([p;q],(:Euclid,:Euclid,:Circular))

X12 = getPoints(pq)

@test 0.7*N < sum(abs.(X12[1,:]) .< 0.3)
@test 0.7*N < sum(abs.(X12[2,:]) .< 0.3)
@test 0.7*N < sum(abs.(X12[3,:]) .< 0.3)

end


"""
    function destriper(N::Int, dataset_baselines; n_iter=10)
# Brief description
You have to use this function to clean a noisy-TOD (comphrensive
of white and one-over-f nosie).

Suggestion for raster scan:
In order to suppress the atmospheric noise you should use many
baselines, better if you have observed those maps with different
scanning strategies. Now you can use only Left-Right and Up-Down
scanning.

The descriper solve the noisy-tod linear system:
[Pᵀ N⁻² P] m = [Pᵀ N⁻²] d

We can reduce it at the tipical solution of
A x = b where:

A = [Pᵀ N⁻² P]
b = [Pᵀ N⁻²] d

m = [Pᵀ N⁻² P]⁻¹ [Pᵀ N⁻²] d
These matris are too big, so we can you the conjugate gradient method
in order to solve the system.
"""

function destriper(N::Int, dataset_baselines; tolerance=1e-6, n_iter=100)
    b = zeros(N*N)
    for i in dataset_baselines
            tod = denoise(i.time_order_data, i.noise)
            a, h = binned_map(N, tod, i.pointing)
            b .+= reshape(a, N*N, 1)[:, 1]
    end

    function A(x)
        res = zeros(N, N)
        for i in dataset_baselines
            tod = get_tod(reshape(x, N, N), i.pointing)
            tod = denoise(tod, i.noise)
            mat, hmat = binned_map(N, tod, i.pointing)
            res .+= mat
        end
        return reshape(res, N*N, 1)[:, 1]
    end

    return ConjGrad.cg(A, b, tol=1e-10, maxIter=100)
    #return conjgrad(A, b, tol=1e-10, maxiter=100)

end

"""
    function denoise(tod::Array{Float64, 1}, noise_s::Array{Float64, 1})

# Brief Description
This function apply the N⁻² variance matrix to the TOD working in the
fourier space.
"""
function denoise(tod::Array{Float64, 1}, noise_s::Array{Float64, 1})
        FFTW.set_num_threads(12)
        ftod = fft(tod)
        ftod ./= noise_s
        tod = real.(ifft(ftod))
        return tod
end

"""
    function conjgrad(A, b; x0=0, tol=1e-6, maxiter=100)
# Brief description
"""
genblas_dot(x, y) = dot(x,y)
genblas_scal!(a, x) = x .*= a
genblas_axpy!(a, x, y) = y .+= a.*x
genblas_nrm2(x) = norm(x)
function conjgrad(A, b; tol=1e-6, maxiter=100)

    x = b .* 0
    r = b

    n = length(b)
    z = copy(r)

    p = z
    err = Inf
    i = 0

    residual_0 = genblas_nrm2(r)

    for i = 1:maxiter
        Ap = A(p)
        gamma = genblas_dot(r,z)
        alpha = gamma / genblas_dot(p, Ap)

        if alpha == Inf || alpha < 0
            return -13, iter
        end

        genblas_axpy!(alpha, p, x)
        genblas_axpy!(-alpha, Ap, r)

        residual = genblas_nrm2(r) / residual_0

        if residual <= tol
            println("Resi: ",residual)
            println("Iter: ",i)
            return x
        end

        z = copy(r)

        beta = genblas_dot(z, r) / gamma
        genblas_scal!(beta, p)
        genblas_axpy!(1.0, z, p)

        #println("Err: ", i, " ", residual)
    end

    return x

end

function make_dataset(tod::Array{Float64, 1}, pix_idx::Array{Int64, 1},
                      Num_bas::Int64
                      )
    len_tod = length(tod)
    len2    = Int(round(len_tod / Num_bas))

    mat_tod = reshape(tod, (len2, Num_bas))
    mat_point = reshape(pix_idx, (len2, Num_bas))

    dataset_baselines = Array{Tpn, 1}(undef, Num_bas)

    noise_spec = tod_noise(len2,  0.0083, 0.1, 3.0, 20.76)
    # Ora e` tutto uguale per ogni baselines, l'ideale sarebbe
    # inserire un fit, con il modello scelto, per ogni baselines

    for i in 1:Num_bas
        dat = Tpn(mat_tod[:, i], mat_point[:, i], noise_spec)
        dataset_baselines[i] = dat
    end

    return dataset_baselines

end


module MultivariateOrthogonalPolynomials
using ClassicalOrthogonalPolynomials, FastTransforms, BlockBandedMatrices, BlockArrays, DomainSets,
      QuasiArrays, StaticArrays, ContinuumArrays, InfiniteArrays, InfiniteLinearAlgebra,
      LazyArrays, SpecialFunctions, LinearAlgebra, BandedMatrices, LazyBandedMatrices, ArrayLayouts,
      HarmonicOrthogonalPolynomials

import Base: axes, in, ==, *, ^, \, copy, OneTo, getindex, size, oneto
import DomainSets: boundary

import QuasiArrays: LazyQuasiMatrix, LazyQuasiArrayStyle
import ContinuumArrays: @simplify, Weight, grid, plotgrid, TransformFactorization, Expansion

import ArrayLayouts: MemoryLayout
import BlockArrays: block, blockindex, BlockSlice, viewblock
import BlockBandedMatrices: _BandedBlockBandedMatrix, AbstractBandedBlockBandedMatrix, _BandedMatrix, blockbandwidths, subblockbandwidths
import LinearAlgebra: factorize
import LazyArrays: arguments, paddeddata, LazyArrayStyle, LazyLayout
import LazyBandedMatrices: LazyBandedBlockBandedLayout
import InfiniteArrays: InfiniteCardinal

import ClassicalOrthogonalPolynomials: jacobimatrix, Weighted, orthogonalityweight, HalfWeighted
import HarmonicOrthogonalPolynomials: BivariateOrthogonalPolynomial, MultivariateOrthogonalPolynomial, Plan,
                                          PartialDerivative, BlockOneTo, BlockRange1, interlace

export UnitTriangle, UnitDisk, JacobiTriangle, TriangleWeight, WeightedTriangle, PartialDerivative, Laplacian,
      MultivariateOrthogonalPolynomial, BivariateOrthogonalPolynomial, Zernike, RadialCoordinate,
      zerniker, zernikez, Weighted, Block, ZernikeWeight, AbsLaplacianPower

include("disk.jl")
include("triangle.jl")

end # module


using ITensors
using ITensorMakie
using GLMakie
using Graphs: Graph
using Random

using ITensors.ContractionSequenceOptimization: optimal_contraction_sequence

Random.seed!(1234)

N = 5
g = Graph(N, N)
A = itensornetwork(g; linkspaces=5)
sequence = optimal_contraction_sequence(A)
edge_labels = (; tags=true)
R = @visualize_sequence fig ITensors.contract(A; sequence=sequence) edge_labels =
  edge_labels

fig


using JuMP
import Ipopt; IpoptOptimizer = Ipopt.Optimizer

@with_kw struct MMDNumerical <: AbstractMMD
    positivity::Bool=true
    normalisation::Bool=true
end

function _estimate_ratio(mmd::MMDNumerical, Kdede, Kdenu)
    n_de, n_nu = size(Kdenu)
    model = Model(with_optimizer(IpoptOptimizer; print_level=0))
    @variable(model, r[1:n_de])
    @objective(model, Min, 1 / n_de ^ 2 * sum(r[i] * Kdede[i,j] * r[j] for i = 1:n_de, j=1:n_de) - 2 / (n_de * n_nu) * sum(r[i] * Kdenu[i,j] for i = 1:n_de, j=1:n_nu))
    if mmd.positivity
        @constraint(model, r .>= 0)
    end
    if mmd.normalisation
        @constraint(model, 1 / n_de * sum(r) == 1)
    end
    optimize!(model)
    return value.(r)
end


# Thompson sampling.
# Based on https://arxiv.org/abs/1111.1797.

mutable struct ThompsonSampling <: Policy end

struct ThompsonSamplingDetails <: PolicyDetails
  solver_time::Float64
end

function choose_action(instance::CombinatorialInstance{T}, policy::ThompsonSampling, state::State{T}; with_trace::Bool=false) where T
  # Determine the weights for each arm and use them to solve the combinatorial problem.
  # println("Counts: $(state.arm_counts); rewards: $(state.arm_average_reward)")
  weights = Dict(arm => rand(Beta(state.arm_counts[arm] * state.arm_average_reward[arm] + 1,
                                  state.arm_counts[arm] * (1 - state.arm_average_reward[arm]) + 1)) for arm in keys(state.arm_counts))

  t0 = time_ns()
  sol = solve_linear(instance, weights)
  t1 = time_ns()

  if with_trace
    return sol, ThompsonSamplingDetails((t1 - t0) / 1_000_000_000)
  else
    return sol
  end
end


using TomoForward
using XfromProjections

# using MKLSparse # uncomment if you've installed MKLSparse, it will boost the performance

# test slice by slice
img = zeros(128, 128, 128)
img[40:70, 40:60, 50:70] .= 1.0

nslice = size(img, 3)

nangles = 90
detcount = Int(floor(size(img,1)*1.4))
proj_geom = ProjGeom(1.0, detcount, LinRange(0,pi,nangles+1)[1:nangles])

isdefined_A = @isdefined A
if isdefined_A == false
    A = fp_op_parallel2d_strip(proj_geom, size(img, 1), size(img, 2))
end

p = zeros(nangles, nslice, detcount)
for i=1:nslice
    p[:,i,:] = reshape(A * vec(img[:,:,i]), nangles, detcount);
end

p = permutedims(p, [2,3,1])

u0 = zeros(size(img));
@time u = recon2d_stack_tv_primaldual!(u0, A, p, 20, 0.1);

# using PyPlot
# imshow(u[:,:,60])


using JuMP
using Gurobi
using Plasmo
using PlasmoAlgorithms

##Place MP and SP into ModelGraph
mp = Model(solver = GurobiSolver())
sp = Model(solver = GurobiSolver())

@variable(mp,y>=1)
@objective(mp,Min,2y)

@variable(sp,x[1:2]>=0)
@variable(sp,y>=0)
@constraint(sp,x[1:2].<=1)
@constraint(sp,x[1]+2x[2]+y>=5)
@objective(sp,Min,2x[1]+3x[2])

## Plasmo Graph
g = ModelGraph()
setsolver(g, GurobiSolver())
n1 = add_node(g)
setmodel(n1,mp)
n2 = add_node(g)
setmodel(n2,sp)

##Set n2 as a child node of n1
edge = Plasmo.add_edge(g,n1,n2)

## Linking constraints between MP and SP
@linkconstraint(g, n1[:y] == n2[:y])

bendersolve(g, max_iterations=20)


import HypothesisTests.LjungBoxTest
import HypothesisTests.BoxPierceTest

#
#
# """
#     LjungBoxTest(d::UncertainDataset; n::Int = 10000,
#         lag::Int = 1, dof::Int = 0) -> LjungBoxTest
#
# Compute the Ljung-Box `Q` statistic to test the null hypothesis of
# independence in a data series represented by the uncertain dataset `d`.
#
# `lag` specifies the number of lags used in the construction of `Q`. When
# testing the residuals of an estimated model, `dof` has to be set to the number
# of estimated parameters. E.g., when testing the residuals of an ARIMA(p,0,q)
# model, set `dof=p+q`.
# """
# function LjungBoxTest(d::UncertainDataset; lag::Int = 1, dof::Int = 0)
#     LjungBoxTest(resample(d, 1)[1], lag, dof)
# end
#
# """
#     LjungBoxTest(d::UncertainDataset, n_tests::Int = 100; lag::Int = 1, dof::Int = 0) -> LjungBoxTest
#
# Compute the Ljung-Box `Q` statistic to test the null hypothesis of
# independence in a data series represented by the uncertain dataset `d`
# by performing the test on `n` independent draws of the dataset.
#
# `lag` specifies the number of lags used in the construction of `Q`. When
# testing the residuals of an estimated model, `dof` has to be set to the number
# of estimated parameters. E.g., when testing the residuals of an ARIMA(p,0,q)
# model, set `dof=p+q`.
# """
# function LjungBoxTest(d::UncertainDataset, n_tests::Int = 100; lag::Int = 1, dof::Int = 0)
#     [LjungBoxTest(resample(d)[1], lag, dof) for i = 1:n_tests]
# end
#
# """
#     BoxPierceTest(d::UncertainDataset;
#         lag::Int = 1, dof::Int = 0) -> Vector{BoxPierceTest}
#
# Compute the Box-Pierce `Q` statistic to test the null hypothesis of
# independence in a data series represented by the uncertain dataset `d`.
#
# `lag` specifies the number of lags used in the construction of `Q`. When
# testing the residuals of an estimated model, `dof` has to be set to the number
# of estimated parameters. E.g., when testing the residuals of an ARIMA(p,0,q)
# model, set `dof=p+q`.
# """
# function BoxPierceTest(d::UncertainDataset; lag::Int = 1, dof::Int = 0)
#     BoxPierceTest(resample(d, 1)[1], lag, dof)
# end
#
#
# """
#     BoxPierceTest(d::UncertainDataset, n_tests::Int = 100;
#         lag::Int = 1, dof::Int = 0) -> BoxPierceTest
#
# Compute the BoxPierceTest `Q` statistic to test the null hypothesis of
# independence in a data series represented by the uncertain dataset `d`
# by performing the test on `n` independent draws of the dataset.
#
# `lag` specifies the number of lags used in the construction of `Q`. When
# testing the residuals of an estimated model, `dof` has to be set to the number
# of estimated parameters. E.g., when testing the residuals of an ARIMA(p,0,q)
# model, set `dof=p+q`.
# """
# function BoxPierceTest(d::UncertainDataset, n_tests::Int = 100;
#         lag::Int = 1, dof::Int = 0)
#     [BoxPierceTest(resample(d, 1)[1], lag, dof) for i = 1:n_tests]
# end
#
#
# export
# LjungBoxTest,
# BoxPierceTest


using DataFrames
using Base.Test

include("$(Pkg.dir())/MultiDimEquations/src/MultiDimEquations.jl")

# TEST 1: Testing both defVars()and the @meq macro using a single IndexedTable
df = CSV.read(IOBuffer("""
reg	prod	var	value
us	banana	production	10
us	banana	transfCoef	0.6
us	banana	trValues	2
us	apples	production	7
us	apples	transfCoef	0.7
us	apples	trValues	5
us	juice	production	NA
us	juice	transfCoef	NA
us	juice	trValues	NA
eu	banana	production	5
eu	banana	transfCoef	0.7
eu	banana	trValues	1
eu	apples	production	8
eu	apples	transfCoef	0.8
eu	apples	trValues	4
eu	juice	production	NA
eu	juice	transfCoef	NA
eu	juice	trValues    NA
"""), delim=" ", ignorerepeated=true)
variables =  vcat(unique(DataFrames.dropmissing(df)[:var]),["consumption"])
#defVars(variables,df;dfName="df",varNameCol="var", valueCol="value")
data = defVars(variables, df, tableName="data", varNameCol="var", valueCol="value")
products = ["banana","apples","juice"]
primPr   = products[1:2]
secPr    = [products[3]]
reg      = ["us","eu"]
# equivalent to [production!(sum(trValues_(r,pp) * transfCoef_(r,pp)  for pp in primPr), r, sp) for r in reg, sp in secPr]
@meq production!(r in reg, sp in secPr)   = sum(trValues_(r,pp) * transfCoef_(r,pp)  for pp in primPr)
@meq consumption!(r in reg, pp in primPr) = production_(r,pp) - trValues_(r,pp)
@meq consumption!(r in reg, sp in secPr)  = production_(r, sp)
totalConsumption = sum(consumption_(r,p) for r in reg, p in products)

@test totalConsumption == 26.6

# # TEST 2: Testing both the "old" defVarsDf()and the @meq macro using a single DataFrame as base
# df2 = wsv"""
# reg	prod	var	value
# us	banana	production	10
# us	banana	transfCoef	0.6
# us	banana	trValues	2
# us	apples	production	7
# us	apples	transfCoef	0.7
# us	apples	trValues	5
# us	juice	production	NA
# us	juice	transfCoef	NA
# us	juice	trValues	NA
# eu	banana	production	5
# eu	banana	transfCoef	0.7
# eu	banana	trValues	1
# eu	apples	production	8
# eu	apples	transfCoef	0.8
# eu	apples	trValues	4
# eu	juice	production	NA
# eu	juice	transfCoef	NA
# eu	juice	trValues    NA
# """
# variables =  vcat(unique(dropna(df[:var])),["consumption"])
# defVarsDf(variables,df2;dfName="df2",varNameCol="var", valueCol="value")
# products = ["banana","apples","juice"]
# primPr   = products[1:2]
# secPr    = [products[3]]
# reg      = ["us","eu"]
# # equivalent to [production!(sum(trValues_(r,pp) * transfCoef_(r,pp)  for pp in primPr), r, sp) for r in reg, sp in secPr]
# @meq production!(r in reg, sp in secPr)   = sum(trValues_(r,pp) * transfCoef_(r,pp)  for pp in primPr)
# @meq consumption!(r in reg, pp in primPr) = production_(r,pp) - trValues_(r,pp)
# @meq consumption!(r in reg, sp in secPr)  = production_(r, sp)
# totalConsumption = sum(consumption_(r,p) for r in reg, p in products)
#
# @test totalConsumption == 26.6

# TEST : Testing the @meq macro with individual IndexedTables
a = IndexedTable([["a","a","b","b"],[1,2,1,2]]...,[1,2,3,4])
dim1 = ["a","b"]
dim2 = [1,2]

@meq a[d1 in dim1,2] = a[d1,1]+3
tot = sum(a[d1,d2] for d1 in dim1, d2 in dim2)
@test tot == 14

# Test n: Fake test, this should not pass
# @test 1 == 2


using NLSolversBase, Test
using Random
using LinearAlgebra: Diagonal, I
using SparseArrays
using OptimTestProblems
using RecursiveArrayTools
MVP = OptimTestProblems.MultivariateProblems

# TODO: Use OptimTestProblems (but it does not have exponential_gradient_hession etc.)
# TODO: MultivariateProblems.UnconstrainedProblems.examples["Exponential"]

# Test example
function exponential(x)
    return exp((2.0 - x[1])^2) + exp((3.0 - x[2])^2)
end

function exponential_gradient!(storage, x)
    storage[1] = -2.0 * (2.0 - x[1]) * exp((2.0 - x[1])^2)
    storage[2] = -2.0 * (3.0 - x[2]) * exp((3.0 - x[2])^2)
end

function exponential_gradient(x)
    storage = similar(x)
    storage[1] = -2.0 * (2.0 - x[1]) * exp((2.0 - x[1])^2)
    storage[2] = -2.0 * (3.0 - x[2]) * exp((3.0 - x[2])^2)
    storage
end

function exponential_value_gradient!(storage, x)
    storage[1] = -2.0 * (2.0 - x[1]) * exp((2.0 - x[1])^2)
    storage[2] = -2.0 * (3.0 - x[2]) * exp((3.0 - x[2])^2)
    return exp((2.0 - x[1])^2) + exp((3.0 - x[2])^2)
end

function exponential_value_gradient(x)
    return exponential(x), exponential_gradient(x)
end

function exponential_gradient_hessian(x)
    F = similar(x)
    F[1] = -2.0 * (2.0 - x[1]) * exp((2.0 - x[1])^2)
    F[2] = -2.0 * (3.0 - x[2]) * exp((3.0 - x[2])^2)

    nx = length(x)
    J = zeros(nx, nx)
    J[1, 1] = 2.0 * exp((2.0 - x[1])^2) * (2.0 * x[1]^2 - 8.0 * x[1] + 9)
    J[1, 2] = 0.0
    J[2, 1] = 0.0
    J[2, 2] = 2.0 * exp((3.0 - x[1])^2) * (2.0 * x[2]^2 - 12.0 * x[2] + 19)
    F, J
end

function exponential_hessian(x)
    nx = length(x)
    storage = zeros(nx, nx)
    storage[1, 1] = 2.0 * exp((2.0 - x[1])^2) * (2.0 * x[1]^2 - 8.0 * x[1] + 9)
    storage[1, 2] = 0.0
    storage[2, 1] = 0.0
    storage[2, 2] = 2.0 * exp((3.0 - x[1])^2) * (2.0 * x[2]^2 - 12.0 * x[2] + 19)
    storage
end

function exponential_hessian!(storage, x)
    storage[1, 1] = 2.0 * exp((2.0 - x[1])^2) * (2.0 * x[1]^2 - 8.0 * x[1] + 9)
    storage[1, 2] = 0.0
    storage[2, 1] = 0.0
    storage[2, 2] = 2.0 * exp((3.0 - x[1])^2) * (2.0 * x[2]^2 - 12.0 * x[2] + 19)
end

function exponential_hessian_product!(storage, x)
    storage[1, 1] = 2.0 * exp((2.0 - x[1])^2) * (2.0 * x[1]^2 - 8.0 * x[1] + 9)
    storage[1, 2] = 0.0
    storage[2, 1] = 0.0
    storage[2, 2] = 2.0 * exp((3.0 - x[1])^2) * (2.0 * x[2]^2 - 12.0 * x[2] + 19)
end

include("objective_types.jl")
include("interface.jl")
include("incomplete.jl")
include("constraints.jl")
include("abstractarrays.jl")
include("autodiff.jl")
include("sparse.jl")
include("kwargs.jl")


"""
    skipmissingpairwise(fn::Function)::Function
Given a function `fn` with method `fn(x::RealVector,y::RealVector)::Float64` returns a
function `fn_out(X,Y)` that calls `fn` iteratively on the columns of `X` and `Y` and where
missing elements in those columns are pairwise skipped.

`fn_out` has five methods implementing the single-matrix, two-matrix, matrix-vector,
vector-matrix and vector-vector cases.

For example, if `X` is of type `RealOrMissingMatrix` then `fn_out(X)` returns a matrix
`C` where `C[i,j] == fn(skipmissingpairwise(X[:,i],X[:,j])...)`

# Example
```
julia> x = [missing;1;2;5;4;missing];

julia> y = [1;3;missing;2;3;6];

julia> hcat(x,y)
6×2 Matrix{Union{Missing, Int64}}:
  missing  1
 1         3
 2          missing
 5         2
 4         3
  missing  6

julia> KendallTau.skipmissingpairwise(StatsBase.corkendall)(x,y)
-0.8164965809277261

julia> StatsBase.corkendall([1;5;4],[3;2;3])
-0.8164965809277261
```
"""
function skipmissingpairwise(fn::Function)::Function

    function fn_out(X::RealOrMissingMatrix)
        n = size(X, 2)
        #= TODO this hard-wires on diagonal elements to 1.0,
        OK for cor but wrong in general. =#
        C = Matrix{Float64}(I, n, n)
        for j = 2:n
            for i = 1:j - 1
                C[i,j] = C[j,i] = fn(skipmissingpairwise(X[:,j], X[:,i])...)
            end
        end
        return C
    end

    function fn_out(X::RealOrMissingMatrix,Y::RealOrMissingMatrix)
        nr = size(X, 2)
        nc = size(Y, 2)
        C = Matrix{Float64}(undef, nr, nc)
        for j = 1:nr
            for i = 1:nc
                C[j,i] = fn(skipmissingpairwise(X[:,j], Y[:,i])...)
            end
        end
        return C
    end

    #= Return matrix. Consistent with Statistics.cor and corspearman but not with corkendall,
    but maybe that doesn't matter? =#
    function fn_out(X::RealOrMissingMatrix,y::RealOrMissingVector)
        size(X, 1) == length(y) ||
                throw(DimensionMismatch("X and y have inconsistent dimensions"))
            n = size(X, 2)
            return(reshape([fn(skipmissingpairwise(X[:,i], y)...) for i in 1:n], n, 1))
        end

    function fn_out(x::RealOrMissingVector,Y::RealOrMissingMatrix)
        size(Y, 1) == length(x) ||
            throw(DimensionMismatch("x and Y have inconsistent dimensions"))
        n = size(Y, 2)
        return(reshape([fn(skipmissingpairwise(x, Y[:,i])...) for i in 1:n], 1, n))
    end

    function fn_out(x::RealOrMissingVector,y::RealOrMissingVector)
        length(x) == length(y) || throw(DimensionMismatch("Vectors must have same length"))
        fn(skipmissingpairwise(x, y)...)
    end

    return(fn_out)
end

function skipmissingpairwise(x::RealVector, y::RealVector)
    x, y
end

"""
    skipmissingpairwise(x::RealOrMissingVector, y::RealOrMissingVector)
Returns a pair `(a,b)`, filtered copies of `x` and `y`, in which elements `x[i]` and `y[i]`
are "skipped" (filtered out) if either `ismissing(x[i])` or `ismissing(y[i])`.
"""
function skipmissingpairwise(x::RealOrMissingVector{T}, y::RealOrMissingVector{U}) where T where U

    length(x) == length(y) || error("Vectors must have same length")

    T2 = x isa Vector{Missing} ? Missing : T
    U2 = y isa Vector{Missing} ? Missing : U

    nout::Int = 0
    @inbounds for i = 1:length(x)
        if !(ismissing(x[i]) || ismissing(y[i]))
            nout += 1
        end
    end

    res1 = Vector{T2}(undef, nout)
    res2 = Vector{U2}(undef, nout)
    j::Int = 0

    @inbounds for i = 1:length(x)
        if !(ismissing(x[i]) || ismissing(y[i]))
            j += 1
            res1[j] = x[i]
            res2[j] = y[i]
        end
    end

    return(res1, res2)
end


using Hotwire
using Polynomials
using CurveFit
using Test

@testset "Hotwire.jl" begin

    # Resistance
    r = Resistor(1e3, 0.0, 298.15)

    @test resistance(r) == 1e3
    @test resistance(r, 298.15) == 1e3

    r = Resistor(1e3, 0.01, 298.15)
    @test resistance(r, 298.15) == 1e3
    @test resistance(r, 299.15) == 1.01e3

    @test temperature(r, 1010.0) == 299.15

    # Thermistor
    r = Thermistor(1e3, 0.0, 298.15)
    @test resistance(r) == 1e3
    @test resistance(r, 303.15) == 1e3

    r = Thermistor(1e3, 3300.0, 298.15)
    @test temperature(r, resistance(r, 50.0+273.15)) ≈ 50.0+273.15
    @test temperature(r, resistance(r, 100.0+273.15)) ≈ 100.0+273.15

    @test 100.0 ≈ resistance(r, temperature(r, 100.0))

    # CTASensor

    r = Resistor(1.0, 0.01, 300.0)
    w = CTASensor(r, 2.0)
    @test r == resistor(w)
    @test optemperature(w) ≈ 400.0
    @test overtemp(w) ≈ 100.0
    @test overheat_ratio(w) ≈ 1.0

    w1 = Wire(1.0, 2.0, 300.0, 0.01)
    @test w == w1

    r = Thermistor(1e3, 3300, 25.0+273.15)
    Rw = 100.0
    w = CTASensor(r, Rw)
    Tw = temperature(r, Rw)
    a = Rw/1e3 - 1.0
    @test r == resistor(w)
    @test optemperature(w) ≈ Tw
    @test overtemp(w) ≈ Tw - 298.15
    @test overheat_ratio(w) ≈ a

    # CCASensor - not much to do
    r = Resistor(1.0, 0.01, 300.0)
    w = CCASensor(r, 1.0)
    @test resistor(w) == r
    @test current(w) == 1.0

    # tempcorr.jl
    # tempcorr(CTASensor)
    Too = 20.0+273.15
    r = Resistor(3.5, 0.0036, Too)
    w = CTASensor(r, 5.5)
    @test 1.0 ≈ tempcorr(w, 1.0, Too) # No temperature correction
    @test tempcorr(w, 1.0, 25.0+273.15) > 1.0 # Increased temperature -> lower anemometer output
    @test tempcorr(w, 1.0, 15.0+273.15) < 1.0 # Decreased temperature -> higher anemometer output
    E1 = 1.0
    Too1 = 25.0+273.15
    E = tempcorr(w, E1, Too1)
    Tw = optemperature(w)
    @test E/E1 ≈ sqrt( (Tw - Too) / (Tw - Too1) )

    r = Thermistor(1e3, 3300, Too)
    w = CTASensor(r, 10.0)
    @test 1.0 ≈ tempcorr(w, 1.0, 20.0+273.15) # No temperature correction
    @test tempcorr(w, 1.0, 25.0+273.15) > 1.0 # Increased temperature -> lower anemometer output
    @test tempcorr(w, 1.0, 15.0+273.15) < 1.0 # Decreased temperature -> higher anemometer output
    E1 = 1.0
    Too1 = 25.0+273.15
    E = tempcorr(w, E1, Too1)
    Tw = optemperature(w)
    @test E/E1 ≈ sqrt( (Tw - Too) / (Tw - Too1) )

    @test 1.0 ≈ tempcorr(w, 1.0, Too, Too)
    @test tempcorr(w, 1.0, 25.0+273.15, Too) > 1.0
    @test tempcorr(w, 1.0, 15.0+273.15, Too) < 1.0
    E = tempcorr(w, E1, Too1, Too)
    Tw = optemperature(w)
    @test E/E1 ≈ sqrt( (Tw - Too) / (Tw - Too1) )

    # tempcorr(CCASensor)
    To = 20.0+273.15
    r = Resistor(3.5, 0.0036, To)
    Two = 100.0+273.15

    Rwo = resistance(r, Two)
    I = 0.2
    Eo = Rwo * I
    hA = Eo*I / (Two - To)
    g = 1.0

    w = CCASensor(r, I, g)

    @test tempcorr(w, Eo, To) ≈ Eo
    @test tempcorr(w, Eo, To, To) ≈ Eo

    Te = 30.0+273.15
    E = Eo
    Eo = tempcorr(w, E, Te)
    Two = temperature(r, Eo/(g*I))
    Tw = temperature(r, E/(g*I))
    Qo = Eo/g * I
    Q = E/g * I
    hAo = Qo / (Two - To)
    hA = Q / (Tw - Te)
    @test Eo / E ≈ (Two - To) / (Tw - Te)
    @test hA ≈ hAo

    Te = 30.0+273.15
    E = Eo
    Eo = tempcorr(w, E, Te, temperature(w.R))
    Two = temperature(r, Eo/(g*I))
    Tw = temperature(r, E/(g*I))
    Qo = Eo/g * I
    Q = E/g * I
    hAo = Qo / (Two - To)
    hA = Q / (Tw - Te)
    @test Eo / E ≈ (Two - To) / (Tw - Te)
    @test hA ≈ hAo

    # Testing root finding equation
    @test Hotwire.funroot(x->(x^2-1.0), 0.5, 0.6) ≈ 1.0 atol=1e-4
    @test Hotwire.funroot(x->sin(x), π*1.1, π*1.02) ≈ π atol=1e-4

    # Testing tempcorr for CCASensor{Thermistor}
    To = 20.0+273.15
    r = Thermistor(1e3, 3300, To)
    g = 1.0
    I = 0.2
    w = CCASensor(r, I, g)
    @test tempcorr(w, 1.0, To) ≈ 1.0

    for Te in [15.0, 30.0] .+ 273.15
        for E in [0.5, 1.0, 2.0]
            Eo = tempcorr(w, E, Te)
            Two = temperature(r, Eo/(g*I))
            Tw = temperature(r, E/(g*I))
            @test Eo/E ≈ (Two-To) / (Tw - Te) atol=1e-6
            Qo = Eo/g * I
            Q = E/g * I
            hAo = Qo / (Two - To)
            hA = Q / (Tw - Te)
            @test hA ≈ hAo atol=1e-5
        end
    end

    for Te in [15.0, 30.0] .+ 273.15
        for E in [0.5, 1.0, 2.0]
            Eo = tempcorr(w, E, Te, To)
            Two = temperature(r, Eo/(g*I))
            Tw = temperature(r, E/(g*I))
            @test Eo/E ≈ (Two-To) / (Tw - Te) atol=1e-6
            Qo = Eo/g * I
            Q = E/g * I
            hAo = Qo / (Two - To)
            hA = Q / (Tw - Te)
            @test hA ≈ hAo atol=1e-5
        end
    end

    # Testing calibration curve
    Rw = Thermistor(5e3, 3500, 25.0+273.15);
    cta = CTASensor(Rw, 100.0, 20/120);

    E = [1.333, 1.385, 1.423, 1.48, 1.521, 1.551, 1.575, 1.624]
    T0 = 20.0+273.15
    coefs = [2527.0, -3916.0, -1548.0, 6110.0, -3945.0, 815.0]
    fit = Polynomial(coefs)
    U = fit.(E)

    cal = calibr_curve(cta, U, E, T0, 5, extrapolate=false)

    @test cal.T0 == T0
    @test cal.fit.coeffs ≈ coefs rtol=1e-4
    @test all(cal.(cta, E) .≈ U)

    U = 1.0:0.5:10.0
    A = 2.0
    B = 0.7
    n = 0.42
    E = sqrt.(A .+ B .* U .^ n)
    fitk = KingFit(E, U)
    efitk = ExtrapolateFit(fitk, sqrt(A), E[1])

    cal = CalibrCurve(E, collect(U), efitk, T0)

    @test efitk.A ≈ A
    @test efitk.B ≈ B
    @test efitk.n ≈ n

    @test all(cal.(cta, E) .≈ U)

    @test cal(cta, sqrt(2.0)) ≈ 0.0

    # Test Probe3d stuff
    k² = (0.0, 0.0, 0.0)
    h² = (1.0, 1.0, 1.0)

    U = collect(1.0:0.5:10.0)
    T0 = 20.0+273.15
    E1 = sqrt.(2.0 .+ 0.7 .* U .^ 0.42)
    E2 = sqrt.(2.2 .+ 0.65 .* U .^ 0.42)
    E3 = sqrt.(1.95 .+ 0.75 .* U .^ 0.41)
    cal = (CalibrCurve(E1, U, KingFit(E1, U), T0),
           CalibrCurve(E2, U, KingFit(E2, U), T0),
           CalibrCurve(E3, U, KingFit(E3, U), T0))

    wires = (Wire(5.0, 5.0*1.8, T0),
             Wire(5.2, 5.2*1.8, T0),
             Wire(5.3, 5.3*1.8, T0))

    probe = Probe3d(wires, cal, k², h²)
    for i in 1:length(U)
        Uⁱ = probe(E1[i], E2[i], E3[i], T0)
        @test Uⁱ[1] ≈ U[i] rtol=1e-4
    end

    k² = (0.03, 0.04, 0.05)
    h² = (1.05, 1.08, 1.02)

    probe = Probe3d(wires, cal, k², h²)
    for i in 1:length(U)
        Uⁱ = probe(E1[i], E2[i], E3[i], T0)
        @test Uⁱ[1] ≈ U[i] rtol=1e-4
    end

end


"""
    ScaleTransform(l::Real)

Multiply every element of the input by `l`
```
    l = 2.0
    tr = ScaleTransform(l)
```
"""
struct ScaleTransform{T<:Real} <: Transform
    s::Vector{T}
end

function ScaleTransform(s::T=1.0) where {T<:Real}
    ScaleTransform{T}([s])
end

set!(t::ScaleTransform,ρ::Real) = t.s .= [ρ]

(t::ScaleTransform)(x) = first(t.s) .* x

_map(t::ScaleTransform, x::AbstractVector{<:Real}) = first(t.s) .* x
_map(t::ScaleTransform, x::ColVecs) = ColVecs(first(t.s) .* x.X)
_map(t::ScaleTransform, x::RowVecs) = RowVecs(first(t.s) .* x.X)

Base.isequal(t::ScaleTransform,t2::ScaleTransform) = isequal(first(t.s),first(t2.s))

Base.show(io::IO,t::ScaleTransform) = print(io,"Scale Transform (s = ", first(t.s), ")")


################################################################################
### JULIA PACKAGE IMPORTS ###
using Plots
using Colors
using DelimitedFiles

################################################################################
### CONTROL FLOW ###

# Source data, if any, also the save location for the next run
Xfilename = "LastX.txt"
Yfilename = "LastY.txt"
GifOutputName = "Dynamics.gif"

# Is warm-start data available?
WarmStart = true

# Add noise to prey data?
ShakeUpX = false

# Set a grid size for simulation
GridSize = 512

# Choose the time length of the simulation
TMAX = 24

# Choose the time point to begin animating from (i.e., animate TSTART:TMAX; note min is 1, not 0)
TSTART = 1

################################################################################
### PARAMETERS ###

μ = 4
β = 5
D1 = 0.001
D2 = 0.2

δ1x = 0 # prey min initialization
δ2x = 1 # prey max initialization
δ1y = 0.3 # predator min initialization
δ2y = 0.4 # predator max initialization

################################################################################
### SET UP GRID ###

# Actual grid size
ShadowGridSize = GridSize+2 # +2 because julia is 1-indexed; the size of the grid including the PBC
GridEnd = GridSize+1 # +1 because julia is 1-indexed; the grid end point excluding the PBC

# Shadow grid with PBC elements
I1S = collect(1:ShadowGridSize) # Ghost elements for grid PBC
I2S = I1S # alias for second dimension

# True grid without PBC elements
I1 = collect(2:GridEnd) # This is one of those times 1-indexing is annoying; the grid points excluding the PBC
I2 = I1 # alias for second dimension

# Number of time points
T = collect(1:TMAX)
TANIM = collect(TSTART:TMAX)

# time points where a progress message is displayed
TPrint = []
TPrintAnim = []
UpdateInterval = 10
for i = 1:UpdateInterval
    push!(TPrint,Int(floor((i/UpdateInterval)*TMAX)))
    push!(TPrintAnim,Int(floor((i/UpdateInterval)*(TMAX-TSTART)+TSTART)))
end

################################################################################
### VARIABLES ###

X = zeros(ShadowGridSize,ShadowGridSize,TMAX)
Y = zeros(ShadowGridSize,ShadowGridSize,TMAX)
∇X = zeros(ShadowGridSize,ShadowGridSize,TMAX)
∇Y = zeros(ShadowGridSize,ShadowGridSize,TMAX)

################################################################################
### SIMULATION FUNCTIONS ###

function x_update(t)
    """
    Prey dynamics
    """
    for i1 in I1
        for i2 in I2
            X[i1,i2,t] = μ*X[i1,i2,t-1]*(1 - X[i1,i2,t-1])*exp(-β*Y[i1,i2,t-1]) + D1*∇X[i1,i2,t-1]
            if X[i1,i2,t] < 0
                X[i1,i2,t] = 0
            end
        end
    end
end

function y_update(t)
    """
    Predator dynamics
    """
    for i1 in I1
        for i2 in I2
            Y[i1,i2,t] = X[i1,i2,t-1]*(1 - exp(-β*Y[i1,i2,t-1])) + D2*∇Y[i1,i2,t-1]
            if Y[i1,i2,t] < 0
                Y[i1,i2,t] = 0
            end
        end
    end
end

function delx_update(t)
    """
    Spatial coupling
    """
    for i1 in I1
        for i2 in I2
            ∇X[i1,i2,t] = X[i1-1,i2,t] + X[i1+1,i2,t] + X[i1,i2-1,t] + X[i1,i2+1,t] - 4*X[i1,i2,t]
        end
    end
end

function dely_update(t)
    """
    Spatial coupling
    """
    for i1 in I1
        for i2 in I2
            ∇Y[i1,i2,t] = Y[i1-1,i2,t] + Y[i1+1,i2,t] + Y[i1,i2-1,t] + Y[i1,i2+1,t] - 4*Y[i1,i2,t]
        end
    end
end

function PBC(t)
    """
    Period Boundary Condition
    """
    for i1 in I1
        X[i1,1,t] = X[i1,GridEnd,t]
        X[i1,ShadowGridSize,t] = X[i1,2,t]
        Y[i1,1,t] = Y[i1,GridEnd,t]
        Y[i1,ShadowGridSize,t] = Y[i1,2,t]
    end
    for i2 in I2
        X[1,i2,t] = X[GridEnd,i2,t]
        X[ShadowGridSize,i2,t] = X[2,i2,t]
        Y[1,i2,t] = Y[GridEnd,i2,t]
        Y[ShadowGridSize,i2,t] = Y[2,i2,t]
    end
end

function XAddNoise(t)
    """
    Add noise to data; i.e., if importing results that are steady-state.
    Noise centred at 0, uniform in +/- δ2x (upper initialization limit)
    After noise addition, values >1 -> 1, value <0 -> 0
    """
    X[:,:,t] += 2*(rand(ShadowGridSize,ShadowGridSize)-0.5*ones(ShadowGridSize,ShadowGridSize))*δ2x
    for i1 in I1
        for i2 in I2
            if X[i1,i2,t] > 1
                X[i1,i2,t] = 1
            elseif X[i1,i2,t] < 0
                X[i1,i2,t] = 0
            end
        end
    end
end

################################################################################
### INITIALIZE FOR t=1 ###
if WarmStart
    # Assign warm start data as initial value
    X[:,:,1] = readdlm(Xfilename, '|', Any)
    Y[:,:,1] = readdlm(Yfilename, '|', Any)
else
    # Assign random values
    X[:,:,1] = rand(ShadowGridSize,ShadowGridSize)*(δ2x-δ1x) + ones(ShadowGridSize,ShadowGridSize)*δ1x
    Y[:,:,1] = rand(ShadowGridSize,ShadowGridSize)*(δ2y-δ1y) + ones(ShadowGridSize,ShadowGridSize)*δ1y
end

# Add noise to prey data; useful to add disturbance to steady-state warm-start results
if ShakeUpX
    XAddNoise(1)
end

################################################################################
### RUN SIMULATION ###
println("Starting simulation...")

# Iterate over time horizon starting from t=1
for t = 2:TMAX
    # Apply boundary conditions to previous time step
    PBC(t-1)

    # Calculate derivative analogues (using PBC)
    delx_update(t-1)
    dely_update(t-1)

    # Update X and Y for period t
    x_update(t)
    y_update(t)

    if t in TPrint
        println("Simulation progress: "*string(Int(floor(100*t/TMAX)))*"%")
    end
end
println("Done simulation!")

# Save last time point for warm start
file = open(Xfilename,"w")
    writedlm(file, X[:,:,TMAX], '|')
close(file)
file = open(Yfilename,"w")
    writedlm(file, Y[:,:,TMAX], '|')
close(file)

################################################################################
### ANIMATE ###
# Start prey animation loop
println("Animating dynamics...")
anim = @animate for t in TANIM
    # Takes the t-th slice and plot
    # for a different colorscheme: color=:halide or, color=:magma
    h1 = heatmap(X[:,:,t], clims=(0,1), xlabel="Prey")
    h2 = heatmap(Y[:,:,t], clims=(0,1), xlabel="Predators")
    plot(h1,h2, layout=(1,2), legend=false, wsize=(1200,600))

    # Add the number to the heatmap
    annotate!(0,0,string(lpad(t - (TSTART-1),ndigits(TMAX),"0")))

    # Progress message:
    if t in TPrintAnim
        println("Animation progress: "*string(Int(floor(100*(t-TSTART)/(TMAX-(TSTART-1)))))*"%")
    end
end
gif(anim, GifOutputName, fps = 7)
println("Done animation!")


export AbstractHomotopy

"""
    AbstractHomotopy

An abstract type representing a homotopy ``H(x,t)``.

The following homotopies are available:

* [`AffineChartHomotopy`](@ref)
* [`AffineSubspaceHomotopy`](@ref)
* [`ModelKitHomotopy`](@ref)
* [`ParameterHomotopy`](@ref)
* [`StraightLineHomotopy`](@ref)
"""
abstract type AbstractHomotopy end

Base.size(H::AbstractHomotopy, i::Integer) = size(H)[i]
function set_solution!(x::AbstractVector, H::AbstractHomotopy, y::AbstractVector, t)
    x .= y
end
get_solution(H::AbstractHomotopy, x::AbstractVector, t) = copy(x)

start_parameters!(H::AbstractHomotopy, p) = H
target_parameters!(H::AbstractHomotopy, p) = H
# internal only
include("homotopies/toric_homotopy.jl")

# public, these should be linked on the top
include("homotopies/affine_chart_homotopy.jl")
include("homotopies/affine_subspace_homotopy.jl")
include("homotopies/model_kit_homotopy.jl")
include("homotopies/parameter_homotopy.jl")
include("homotopies/coefficient_homotopy.jl")
include("homotopies/straight_line_homotopy.jl")


using CuArrays, CUDAnative, CUDAdrv

#This is the kernel. The plus and minus ones relative to the C version are due
#to the fact that ONE INDEXING SUCKS
function dot(a,b,c, N, threadsPerBlock, blocksPerGrid)

    #Set up shared memory cache for this current block. Has to be dynamic to be
    #able to take command line args - this will not work with
    #@cuStaticSharedMem
    cache = @cuDynamicSharedMem(Int64, threadsPerBlock)

    #Initialise some variables. Minus ones are to get the indexing consistent
    #with the C version
    tid = (threadIdx().x - 1) + (blockIdx().x - 1) * blockDim().x
    cacheIndex = threadIdx().x - 1
    temp::Int64 = 0

    #iterate over vector to do dot product in parallel way
    while tid < N
        temp += a[tid + 1] * b[tid + 1]
        tid += blockDim().x * gridDim().x
    end

    #set cache values
    cache[cacheIndex + 1] = temp

    #synchronise threads
    sync_threads()

    #In the step below, we add up all of the values stored in the cache
    i::Int = blockDim().x/2
    while i!=0
        if cacheIndex < i
            cache[cacheIndex + 1] += cache[cacheIndex + i + 1]
        end
        sync_threads()
        i/=2
    end

    #cache[1] now contains the sum of vector dot product calculations done in
    #this block, so we write it to c
    if cacheIndex == 0
        c[blockIdx().x] = cache[1]
    end

    return nothing
end

#Self explanatory
function sum_squares(x)
    return (x * (x + 1) * (2 * x + 1) / 6)
end

function main()

    #Can't type global variables in Julia, which messes things up.
    #So instead we locally define these and pass them explicitly as arguments to
    #the kernel
    N::Int64 = 33 * 1024
    threadsPerBlock::Int64 = 256
    blocksPerGrid::Int64 = min(32, (N + threadsPerBlock - 1) / threadsPerBlock)

    #Create a,b and c
    a = CuArrays.CuArray(fill(0, N))
    b = CuArrays.CuArray(fill(0, N))
    c = CuArrays.CuArray(fill(0, blocksPerGrid))

    #Fill and b
    for i in 1:N
        a[i] = i
        b[i] = 2*i
    end

    #Execute the kernel. Note the shmem argument - this is necessary to allocate
    #space for the cache we allocate on the gpu with @cuDynamicSharedMem
    @cuda blocks = blocksPerGrid threads = threadsPerBlock shmem =
    (threadsPerBlock * sizeof(Int64)) dot(a,b,c, N, threadsPerBlock, blocksPerGrid)

    #Copy c back from the gpu (device) to the host
    c = Array(c)

    local result = 0

    #Sum the values in c
    for i in 1:blocksPerGrid
        result += c[i]
    end

    #Check whether output is correct
    println("Does GPU value ", result, " = ", 2 * sum_squares(N - 1))
end

main()


using Test

include("insertion_sort.jl")
include("selection_sort.jl")
include("merge_sort.jl")
include("heap_sort.jl")
include("quick_sort.jl")
include("counting_sort.jl")

@testset "test sorting correctness" begin
    inputs = [
        [5 2 4 7 1 3 2 6],
        [4 1 3 2 16 9 10 14 8 7],
        [2 8 7 1 3 5 6 4],
        [2 5 3 0 2 3 0 3],
        [329 457 657 839 436 720 355],
    ]
    outputs = [
        [1 2 2 3 4 5 6 7],
        [1 2 3 4 7 8 9 10 14 16],
        [1 2 3 4 5 6 7 8],
        [0 0 2 2 3 3 3 5],
        [329 355 436 457 657 720 839],
    ]

    @testset "insertion sort" begin
        for (input, output) in zip(inputs, outputs)
            A = copy(input)
            insertion_sort!(A)
            @test A == output
        end
    end

    @testset "selection sort" begin
        for (input, output) in zip(inputs, outputs)
            A = copy(input)
            selection_sort!(A)
            @test A == output
        end
    end

    @testset "merge sort" begin
        for (input, output) in zip(inputs, outputs)
            A = copy(input)
            merge_sort!(A, 1, length(A))
            @test A == output
        end
    end

    @testset "heap sort" begin
        for (input, output) in zip(inputs, outputs)
            A = copy(input)
            heap_sort!(A)
            @test A == output
        end
    end

    @testset "quick sort" begin
        for (input, output) in zip(inputs, outputs)
            A = copy(input)
            quick_sort!(A, 1, length(A))
            @test A == output

            A = copy(input)
            randomized_quick_sort!(A, 1, length(A))
            @test A == output
        end
    end

    @testset "counting sort" begin
        for (input, output) in zip(inputs, outputs)
            A = copy(input)
            k = maximum(A)
            B = counting_sort(A, k)
            @test B == output
        end
    end

    @testset "radix sort" begin
        for (input, output) in zip(inputs, outputs)
            A = copy(input)
            radix_sort!(A)
            @test A == output
        end
    end
end


module StrategyExample

abstract type Algo end
struct Memoized <: Algo end
struct Iterative <: Algo end

# Memoization startegy
using Memoize
@memoize function _fib(n)
    n <= 2 ? 1 : _fib(n-1) + _fib(n-2)
end

function fib(::Memoized, n)
    println("Using memoization algorithm")
    _fib(n)
end

# Iterative strategy
function fib(algo::Iterative, n)
    println("Using iterative algorithm")
    n <= 2 && return 1
    prev1, prev2 = 1, 1
    local curr
    for i in 3:n
        curr = prev1 + prev2
        prev1, prev2 = curr, prev1
    end
    return curr
end

# From client's perpsective, it just needs to pass an algorithm
# that is implemented for the `fib` function.
# function find_golden_ratio(n; algo = Memoized())
#     return fib(algo, n) / fib(algo, n-1)
# end

# auto-selection
function fib(n)
    algo = n > 50 ? Memoized() : Iterative()
    return fib(algo, n)
end

function test()
    @show fib(30)
    @show fib(60)
end

end #module

using .StrategyExample
StrategyExample.test()


@testset ExtendedTestSet "euclidean" begin
    @testset ExtendedTestSet "euclidean" begin
        x = [
            1 1 0 0
            0 1 1 0
            0 1 0 1
            0 1 0 0
        ]
        answer = [
            1.0 1.0 0.0 0.0
            0.0 1.0 1.0 0.0
            0.0 1.0 0.0 1.0
            0.0 1.0 0.0 0.0
        ]
        @test euclidean(x) == answer
    end

    @testset ExtendedTestSet "euclidean" begin
        x = [
            1 1 1 0
            0 1 1 1
            0 1 1 1
            0 1 1 1
        ]
        answer = [
            1.0 1.4142135623730951 1.0 0.0
            0.0 1.0 1.4142135623730951 1.0
            0.0 1.0 2.0 2.0
            0.0 1.0 2.0 3.0
        ]
        @test euclidean(x) ≈ answer
    end

    @testset ExtendedTestSet "euclidean" begin
        x = Bool.([
            1 1 0 0
            0 1 1 0
            0 1 0 1
            0 1 0 0
        ])
        answer = [
            1.0 1.0 0.0 0.0
            0.0 1.0 1.0 0.0
            0.0 1.0 0.0 1.0
            0.0 1.0 0.0 0.0
        ]
        @test euclidean(x) == answer
    end
end


using Test
using LinearAlgebra
using PolynomialBasis
using ImplicitDomainQuadrature
# using Revise
using CutCell
include("useful_routines.jl")

function test_linear_cut_cell_assembly()
    polyorder = 1
    numqp = 2

    basis = TensorProductBasis(2, polyorder)
    levelset = InterpolatingPolynomial(1, basis)
    nf = CutCell.number_of_basis_functions(basis)
    mesh = CutCell.Mesh([0.0, 0.0], [2.0, 1.0], [1, 1], nf)

    normal = [1.0, 0.0]
    x0 = [0.5, 0.0]
    levelsetcoeffs =
        CutCell.levelset_coefficients(x -> plane_distance_function(x, normal, x0), mesh)

    cutmesh = CutCell.CutMesh(levelset, levelsetcoeffs, mesh)
    cellquads = CutCell.CellQuadratures(levelset, levelsetcoeffs, cutmesh, numqp)
    interfacequads = CutCell.InterfaceQuadratures(levelset, levelsetcoeffs, cutmesh, numqp)

    lambda, mu = (1.0, 2.0)
    penalty = 1.0
    dx = 0.1
    e11 = dx / 2.0
    e22 = -lambda / (lambda + 2mu) * e11
    dy = e22
    s11 = (lambda + 2mu) * e11 + lambda * e22

    stiffness = CutCell.HookeStiffness(lambda, mu, lambda, mu)

    bilinearforms = CutCell.BilinearForms(basis, cellquads, stiffness, cutmesh)
    interfacecondition = CutCell.coherent_interface_condition(
        basis,
        interfacequads,
        stiffness,
        cutmesh,
        penalty,
    )

    sysmatrix = CutCell.SystemMatrix()
    sysrhs = CutCell.SystemRHS()

    CutCell.assemble_bilinear_form!(sysmatrix, bilinearforms, cutmesh)
    CutCell.assemble_interface_condition!(sysmatrix, interfacecondition, cutmesh)

    matrix = CutCell.make_sparse(sysmatrix, cutmesh)
    rhs = CutCell.rhs(sysrhs, cutmesh)

    CutCell.apply_dirichlet_bc!(matrix, rhs, 9, 0.0)
    CutCell.apply_dirichlet_bc!(matrix, rhs, 10, 0.0)
    CutCell.apply_dirichlet_bc!(matrix, rhs, 11, 0.0)

    CutCell.apply_dirichlet_bc!(matrix, rhs, 5, dx)
    CutCell.apply_dirichlet_bc!(matrix, rhs, 7, dx)

    sol = matrix \ rhs
    disp = reshape(sol, 2, :)

    testdisp = [
        0.0 0.0 dx dx 0.0 0.0 dx dx
        0.0 dy 0.0 dy 0.0 dy 0.0 dy
    ]
    @test allapprox(disp, testdisp, 1e2eps())
end

function test_curved_interface_assembly()
    polyorder = 2
    numqp = 4

    basis = TensorProductBasis(2, polyorder)
    levelset = InterpolatingPolynomial(1, basis)
    nf = CutCell.number_of_basis_functions(basis)
    mesh = CutCell.Mesh([0.0, 0.0], [2.0, 1.0], [1, 1], nf)

    xc = [2.0, 0.5]
    rad = 1.0
    levelsetcoeffs =
        CutCell.levelset_coefficients(x -> circle_distance_function(x, xc, rad), mesh)

    cutmesh = CutCell.CutMesh(levelset, levelsetcoeffs, mesh)
    cellquads = CutCell.CellQuadratures(levelset, levelsetcoeffs, cutmesh, numqp)
    interfacequads = CutCell.InterfaceQuadratures(levelset, levelsetcoeffs, cutmesh, numqp)

    lambda, mu = (1.0, 2.0)
    penalty = 1.0
    dx = 0.1
    e11 = dx / 2.0
    e22 = -lambda / (lambda + 2mu) * e11
    dy = e22
    s11 = (lambda + 2mu) * e11 + lambda * e22

    stiffness = CutCell.HookeStiffness(lambda, mu, lambda, mu)

    bilinearforms = CutCell.BilinearForms(basis, cellquads, stiffness, cutmesh)
    interfacecondition = CutCell.coherent_interface_condition(
        basis,
        interfacequads,
        stiffness,
        cutmesh,
        penalty,
    )

    sysmatrix = CutCell.SystemMatrix()
    sysrhs = CutCell.SystemRHS()

    CutCell.assemble_bilinear_form!(sysmatrix, bilinearforms, cutmesh)
    CutCell.assemble_interface_condition!(sysmatrix, interfacecondition, cutmesh)

    matrix = CutCell.make_sparse(sysmatrix, cutmesh)
    rhs = CutCell.rhs(sysrhs, cutmesh)

    CutCell.apply_dirichlet_bc!(matrix, rhs, [10, 11, 12], 1, 0.0, 2)
    CutCell.apply_dirichlet_bc!(matrix, rhs, [10], 2, 0.0, 2)
    CutCell.apply_dirichlet_bc!(matrix, rhs, [7, 8, 9], 1, dx, 2)

    sol = matrix \ rhs
    disp = reshape(sol, 2, :)

    testdisp = [
        0.0 0.0 0.0 dx/2 dx/2 dx/2 dx dx dx
        0.0 dy/2 dy 0.0 dy/2 dy 0.0 dy/2 dy
    ]
    testdisp = repeat(testdisp, outer = (1, 2))

    @test allapprox(disp, testdisp, 1e2eps())
end

test_linear_cut_cell_assembly()
test_curved_interface_assembly()


function OneVsAllPredict(Θ;μ=0,σ=1,ξ=0.5)

  P=function (X)

      m=size(X,1)
      s=zeros(m)
      X=(X.-μ)./σ
      X=[ones(m) X]
      p=Sigmoid(X*Θ')
      for i=1:m
           s[i]=argmax(p[i,:])
        end
        return s
    end

        return P
end


using Flux, MLDatasets, Statistics, Random, BSON
using Flux.Optimise: update!
using Flux: logitbinarycrossentropy

batchSize, latentDim = 500, 100
epochs = 40
etaD, etaG = 0.0002, 0.0002

images, _ = MLDatasets.MNIST.traindata(Float32)
imageTensor = reshape(@.(2f0 * images - 1f0), 28, 28, 1, :)
data = [imageTensor[:, :, :, r] for r in Iterators.partition(1:60000, batchSize)]

dscr = Chain(Conv((4,4),1=>64;stride=2,pad=1),x->leakyrelu.(x,0.2f0),
        Dropout(0.25),Conv((4,4),64=>128;stride=2,pad=1),x->leakyrelu.(x,0.2f0),
        Dropout(0.25), x->reshape(x, 7 * 7 * 128, :), Dense(7 * 7 * 128, 1))
gen =  Chain(Dense(latentDim,7*7*256),BatchNorm(7*7*256,relu),
        x->reshape(x,7,7,256,:),ConvTranspose((5,5),256=>128;stride=1,pad=2),
        BatchNorm(128,relu),ConvTranspose((4,4),128=>64;stride=2,pad=1),
        BatchNorm(64,relu),ConvTranspose((4,4),64=>1,tanh;stride=2,pad=1))

dLoss(realOut,fakeOut) =    mean(logitbinarycrossentropy.(realOut,1f0)) +
                            mean(logitbinarycrossentropy.(fakeOut,0f0))
gLoss(u) = mean(logitbinarycrossentropy.(u, 1f0))

function updateD!(gen, dscr, x, opt_dscr)
    noise = randn!(similar(x, (latentDim, batchSize)))
    fakeInput = gen(noise)
    ps = Flux.params(dscr)
    loss, back = Flux.pullback(()->dLoss(dscr(x), dscr(fakeInput)), ps)
    grad = back(1f0)
    update!(opt_dscr, ps, grad)
    return loss
end

function updateG!(gen, dscr, x, optGen)
    noise = randn!(similar(x, (latentDim, batchSize)))
    ps = Flux.params(gen)
    loss, back = Flux.pullback(()->gLoss(dscr(gen(noise))),ps)
    grad = back(1f0)
    update!(optGen, ps, grad)
    return loss
end

optDscr, optGen = ADAM(etaD), ADAM(etaG)
cd(@__DIR__)
@time begin
    for ep in 1:epochs
        for (bi,x) in enumerate(data)
            lossD = updateD!(gen, dscr, x, optDscr)
            lossG = updateG!(gen, dscr, x, optGen)
            @info "Epoch $ep, batch $bi, D loss = $(lossD), G loss = $(lossG)"
        end
        @info "Saving generator for epcoh $ep"
        BSON.@save "../data/mnistGAN$(ep).bson" genParams=cpu.(params(gen))
    end
end


using SynchrotronKernel, Test

@testset "Synchrotron Kernel" begin

    @testset "F" begin
        @test SynchrotronKernel.F(1.e-8) ≈ 0.0046310000727714734
        @test SynchrotronKernel.F(  1.0) ≈ 0.6514228153553652
        @test SynchrotronKernel.F(  3.5) ≈ 0.08268719536694746
        @test SynchrotronKernel.F(  4.5) ≈ 0.03357441971502366
        @test SynchrotronKernel.F( 10.0) ≈ 0.00019223826430086885
        @test SynchrotronKernel.F( 50.0) ≈ 1.734785203976593e-21
        @test SynchrotronKernel.F(100.0) ≈ 4.697593665922202e-43
        @test SynchrotronKernel.F( 1.e6) == 0.0

        @test ℱ(1.0) == SynchrotronKernel.F(1.0)
    end

    @testset "G" begin
        @test SynchrotronKernel.G(1.0) ≈ 0.4944750621042115

        @test 𝒢(1.0) == SynchrotronKernel.G(1.0)
    end

    @testset "Synchrotron Kernel" begin
        K_tot, K_ort, K_par = synchrotron_kernel(1.0)

        @test K_tot == SynchrotronKernel.F(1.0)
        @test K_ort ≈ 0.5729489387297884
        @test K_par ≈ 0.07847387662557684
        @test K_tot == K_ort + K_par
    end

    @testset "Intensity" begin
        @test synchrotron_intensity(1.0) == SynchrotronKernel.F(1.0)
    end

    @testset "Polarisation" begin
        K = SynchrotronKernel.F(1.0)
        sk_ort, sk_par = synchrotron_polarisation(1.0)

        @test sk_ort ≈ 0.5729489387297884
        @test sk_par ≈ 0.07847387662557684
        @test K == sk_ort + sk_par
    end

end


type FdG2SwaptionEngine{Y <: TermStructure} <: PricingEngine
  model::G2
  tGrid::Int
  xGrid::Int
  yGrid::Int
  dampingSteps::Int
  invEps::Float64
  schemeDesc::FdmSchemeDesc
  ts::Y
end

FdG2SwaptionEngine(model::G2, tGrid::Int = 100, xGrid::Int = 50, yGrid::Int = 50, dampingSteps::Int = 0, invEps::Float64 = 1e-5,
                  schemeDesc::FdmSchemeDesc = FdmSchemeDesc(Hundsdorfer())) =
                  FdG2SwaptionEngine(model, tGrid, xGrid, yGrid, dampingSteps, invEps, schemeDesc, model.ts)

# methods #
function _calculate!(pe::FdG2SwaptionEngine, swaption::Swaption)
  # 1. Term structure
  ts = pe.ts

  # 2. Mesher
  dc = ts.dc
  refDate = reference_date(ts)
  maturity = year_fraction(dc, refDate, swaption.exercise.dates[end])

  process1 = OrnsteinUhlenbeckProcess(get_a(pe.model), get_sigma(pe.model))
  process2 = OrnsteinUhlenbeckProcess(get_b(pe.model), get_eta(pe.model))

  xMesher = FdmSimpleProcess1dMesher(pe.xGrid, process1, maturity, 1, pe.invEps)
  yMesher = FdmSimpleProcess1dMesher(pe.yGrid, process2, maturity, 1, pe.invEps)
  mesher = FdmMesherComposite(xMesher, yMesher)

  # 3. Inner Value calculator
  exerciseDates = swaption.exercise.dates
  t2d = Dict{Float64, Date}()
  for i = 1:length(exerciseDates)
    t = year_fraction(dc, refDate, exerciseDates[i])
    t2d[t] = exerciseDates[i]
  end

  disTs = pe.model.ts
  fwdTs = swaption.swap.iborIndex.ts

  fwdModel = G2(fwdTs, get_a(pe.model), get_sigma(pe.model), get_b(pe.model), get_eta(pe.model), get_rho(pe.model))
  calculator = FdmAffineModelSwapInnerValue(pe.model, fwdModel, swaption.swap, t2d, mesher, 1)

  # 4. Step Conditions
  conditions = vanilla_FdmStepConditionComposite(DividendSchedule(), swaption.exercise, mesher, calculator, refDate, dc)

  # 5. Boundary conditions
  boundaries = FdmBoundaryConditionSet()

  # 6. Solver
  solverDesc = FdmSolverDesc(mesher, boundaries, conditions, calculator, maturity, pe.tGrid, pe.dampingSteps)
  solver = FdmG2Solver(pe.model, solverDesc, pe.schemeDesc)
  swaption.results.value = value_at(solver, 0.0, 0.0)
end


struct DynamicBudgetedBipartiteMatchingSolution{T, U} <: BudgetedBipartiteMatchingSolution{T, U}
  instance::BudgetedBipartiteMatchingInstance{T, U}
  solution::Vector{Edge{T}}
  value::Float64
  state::Dict{Tuple{Int, Int, Int, Int}, Float64} # Data structure built by the dynamic-programming recursion.
  solutions::Dict{Tuple{Int, Int, Int, Int}, Vector{Edge{T}}} # From the indices of state to the corresponding solution.
end

function matching_dp_budgeted(instance::BudgetedBipartiteMatchingInstance{T, U}) where {T, U}
  # First, determine the size of a perfect matching.
  nr = Dict(k => 1.0 for (k, v) in instance.matching.reward)
  ni = BipartiteMatchingInstance(instance.matching, nr)
  ns = matching_dp_imperfect(ni)
  m = length(ns.solution)

  # Solve this budgeted matching to ensure the constraints are met.

  # Four indices:
  # - two vertices, like matching_dp_imperfect;
  # - the number of edges still to add to have a perfect matching, like matching_dp;
  # - the minimum budget to respect.
  state = Dict{Tuple{Int, Int, Int, Int}, Float64}((i, j, k, l) => -Inf for i in instance.matching.vertex_left, j in instance.matching.vertex_right, k in 0:m, l in 0:instance.budget)
  solutions = Dict{Tuple{Int, Int, Int, Int}, Vector{Edge{T}}}((i, j, m, zero(U)) => Edge{T}[] for i in instance.matching.vertex_left, j in instance.matching.vertex_right)

  # Initialise: k=m-1, l=0, i.e. take at most one edge, don't consider the budget yet.
  v1 = first(instance.matching.vertex_right)
  for i in instance.matching.vertex_left
    if has_edge(instance.matching.graph, i, v1)
      state[i, v1, m - 1, 0] = instance.matching.reward[Edge(i, v1)]
      solutions[i, v1, m - 1, 0] = [Edge(i, v1)]
    end
  end

  v1 = first(instance.matching.vertex_left)
  for j in instance.matching.vertex_right
    if has_edge(instance.matching.graph, v1, j)
      state[v1, j, m - 1, 0] = instance.matching.reward[Edge(v1, j)]
      solutions[v1, j, m - 1, 0] = [Edge(v1, j)]
    end
  end

  for (i, prev_i) in zip(instance.matching.vertex_left[2:end], instance.matching.vertex_left[1:end-1])
    for (j, prev_j) in zip(instance.matching.vertex_right[2:end], instance.matching.vertex_right[1:end-1])
      # Take (i, j).
      if has_edge(instance.matching.graph, i, j)
        state[i, j, m - 1, 0] = instance.matching.reward[Edge(i, j)]
        solutions[i, j, m - 1, 0] = [Edge(i, j)]
      end

      # Don't take (i, j): copy a previous solution. No mixing possible, as only one edge is considered.
      if state[prev_i, j, m - 1, 0] > state[i, j, m - 1, 0]
        state[i, j, m - 1, 0] = state[prev_i, j, m - 1, 0]
        solutions[i, j, m - 1, 0] = solutions[prev_i, j, m - 1, 0]
      end
      if state[i, prev_j, m - 1, 0] > state[i, j, m - 1, 0]
        state[i, j, m - 1, 0] = state[i, prev_j, m - 1, 0]
        solutions[i, j, m - 1, 0] = solutions[i, prev_j, m - 1, 0]
      end
    end
  end

  # Initialisation for l=0, k < m-1: create perfect matchings, regardless of budget.
  for k in (m - 2):-1:0
    # Don't handle the cases with just one vertex in one part of the graph:
    # this loop requires at least two edges to be taken, but the considered
    # subset of edges can lead to at most only one being taken.

    # Consider the other vertices.
    for (i, prev_i) in zip(instance.matching.vertex_left[2:end], instance.matching.vertex_left[1:end-1])
      for (j, prev_j) in zip(instance.matching.vertex_right[2:end], instance.matching.vertex_right[1:end-1])
        # Take (i, j).
        if has_edge(instance.matching.graph, i, j) && state[prev_i, prev_j, k + 1, 0] > -Inf
          state[i, j, k, 0] = instance.matching.reward[Edge(i, j)] + state[prev_i, prev_j, k + 1, 0]
          solutions[i, j, k, 0] = solutions[prev_i, prev_j, k + 1, 0]
          push!(solutions[i, j, k, 0], Edge(i, j))
        end

        # Don't take (i, j): mix previous solutions (as a special case, maybe take a previous solution completely).
        for k2 in 0:m
          for k3 in 0:m
            if state[prev_i, j, k2, 0] + state[i, prev_j, k3, 0] > state[i, j, k, 0]
              potential_sol = union(Set{Edge{T}}(solutions[prev_i, j, k2, 0]), Set{Edge{T}}(solutions[i, prev_j, k3, 0]))
              val = sum(instance.matching.reward[Edge(src(e), dst(e))] for e in potential_sol)

              if length(potential_sol) == m - k && val > state[i, j, k, 0]
                state[i, j, k, 0] = val
                solutions[i, j, k, 0] = collect(potential_sol)
              end
            end
          end
        end
      end
    end
  end

  # Main iteration, considering the first vertex of each side is already done.
  for l in 1:instance.budget
    # Repeat the loop as often as necessary to handle the corner case with zero-weight edges.
    while true
      changes = false

      # Deal with the first two vertices (one on each side). All constraints are imposed.
      # Just one edge to be taken, as one vertex is always fixed on one side:
      # no feasible solution for k < m-1.
      vl = first(instance.matching.vertex_left)
      vr = first(instance.matching.vertex_right)
      if has_edge(instance.matching.graph, vl, vr) && instance.weight[Edge(vl, vr)] >= l
        state[vl, vr, m - 1, l] = instance.matching.reward[Edge(vl, vr)]
        solutions[vl, vr, m - 1, l] = [Edge(vl, vr)]
      end

      v1 = first(instance.matching.vertex_right)
      for i in instance.matching.vertex_left
        if has_edge(instance.matching.graph, i, v1) && instance.weight[Edge(i, v1)] >= l
          state[i, v1, m - 1, l] = instance.matching.reward[Edge(i, v1)]
          solutions[i, v1, m - 1, l] = [Edge(i, v1)]
        end
      end

      v1 = first(instance.matching.vertex_left)
      for j in instance.matching.vertex_right
        if has_edge(instance.matching.graph, v1, j) && instance.weight[Edge(v1, j)] >= l
          state[v1, j, m - 1, l] = instance.matching.reward[Edge(v1, j)]
          solutions[v1, j, m - 1, l] = [Edge(v1, j)]
        end
      end

      # Do the rest.
      for k in (m - 1):-1:0
        for (i, prev_i) in zip(instance.matching.vertex_left[2:end], instance.matching.vertex_left[1:end-1])
          for (j, prev_j) in zip(instance.matching.vertex_right[2:end], instance.matching.vertex_right[1:end-1])

            # Take (i, j) if possible.
            if has_edge(instance.matching.graph, i, j)
              remaining_budget = max(0, l - instance.weight[Edge(i, j)])

              if state[prev_i, prev_j, k + 1, remaining_budget] > -Inf
                state[i, j, k, l] = instance.matching.reward[Edge(i, j)] + state[prev_i, prev_j, k + 1, remaining_budget]
                solutions[i, j, k, l] = solutions[prev_i, prev_j, k + 1, remaining_budget]
                push!(solutions[i, j, k, l], Edge(i, j))

                changes = true
              end
            end

            # Don't take (i, j): mix previous solutions (as a special case, maybe take a previous solution completely).
            for k2 in 0:m
              for k3 in 0:m
                if state[prev_i, j, k2, l] + state[i, prev_j, k3, l] > state[i, j, k, l]
                  potential_sol = union(Set{Edge{T}}(solutions[prev_i, j, k2, l]), Set{Edge{T}}(solutions[i, prev_j, k3, l]))
                  val = sum(instance.matching.reward[Edge(src(e), dst(e))] for e in potential_sol)
                  wei = sum(instance.weight[Edge(src(e), dst(e))] for e in potential_sol)

                  if length(potential_sol) == m - k && val > state[i, j, k, l] && wei >= l
                    state[i, j, k, l] = val
                    solutions[i, j, k, l] = collect(potential_sol)

                    changes = true
                  end
                end
              end
            end
          end
        end
      end

      # If this iteration made no change, stop and go on to the next budget.
      if ! changes
        break
      end
    end
  end

  val = state[last(instance.matching.vertex_left), last(instance.matching.vertex_right), 0, instance.budget]
  sol = solutions[last(instance.matching.vertex_left), last(instance.matching.vertex_right), 0, instance.budget]

  return DynamicBudgetedBipartiteMatchingSolution(instance, sol, val, state, solutions)
  # return SimpleBudgetedBipartiteMatchingSolution(instance, sol)
end


"""
    Chain(layers...)

Chain multiple layers / functions together, so that they are called in sequence
on a given input.

```julia
m = Chain(x -> x^2, x -> x+1)
m(5) == 26

m = Chain(Dense(10, 5), Dense(5, 2))
x = rand(10)
m(x) == m[2](m[1](x))
```

`Chain` also supports indexing and slicing, e.g. `m[2]` or `m[1:end-1]`.
`m[1:3](x)` will calculate the output of the first three layers.
"""
type Chain
  layers::Vector{Any}
  Chain(xs...) = new([xs...])
end

@forward Chain.layers Base.getindex, Base.first, Base.last, Base.endof, Base.push!
@forward Chain.layers Base.start, Base.next, Base.done

children(c::Chain) = c.layers
mapchildren(f, c::Chain) = Chain(f.(c.layers)...)
adapt(T, c::Chain) = Chain(map(x -> adapt(T, x), c.layers)...)

(c::Chain)(x) = foldl((x, m) -> m(x), x, c.layers)

Base.getindex(c::Chain, i::AbstractArray) = Chain(c.layers[i]...)

function Base.show(io::IO, c::Chain)
  print(io, "Chain(")
  join(io, c.layers, ", ")
  print(io, ")")
end

"""
    Dense(in::Integer, out::Integer, σ = identity)

Creates a traditional `Dense` layer with parameters `W` and `b`.

    y = σ.(W * x .+ b)

The input `x` must be a vector of length `in`, or a batch of vectors represented
as an `in × N` matrix. The out `y` will be a vector or batch of length `out`.

```julia
julia> d = Dense(5, 2)
Dense(5, 2)

julia> d(rand(5))
Tracked 2-element Array{Float64,1}:
  0.00257447
  -0.00449443
```
"""
struct Dense{F,S,T}
  W::S
  b::T
  σ::F
end

Dense(W, b) = Dense(W, b, identity)

function Dense(in::Integer, out::Integer, σ = identity;
               initW = glorot_uniform, initb = zeros)
  return Dense(param(initW(out, in)), param(initb(out)), σ)
end

treelike(Dense)

function (a::Dense)(x)
  W, b, σ = a.W, a.b, a.σ
  @fix σ.(W*x .+ b)
end

function Base.show(io::IO, l::Dense)
  print(io, "Dense(", size(l.W, 2), ", ", size(l.W, 1))
  l.σ == identity || print(io, ", ", l.σ)
  print(io, ")")
end

"""
    Diagonal(in::Integer)

Creates an element-wise linear transformation layer with learnable
vectors `α` and `β`:

    y = α .* x .+ β

The input `x` must be a array where `size(x, 1) == in`.
"""
struct Diagonal{T}
  α::T
  β::T
end

Diagonal(in::Integer; initα = ones, initβ = zeros) =
  Diagonal(param(initα(in)), param(initβ(in)))

treelike(Diagonal)

function (a::Diagonal)(x)
  α, β = a.α, a.β
  α.*x .+ β
end

function Base.show(io::IO, l::Diagonal)
  print(io, "Diagonal(", length(l.α), ")")
end


@testset "Laplace" begin
    test_interface(LaplaceLikelihood(3.0), Laplace)
    @test LaplaceLikelihood().β == 1
    test_auglik(LaplaceLikelihood(1.0); rng=MersenneTwister(42))
    # Test the custom kl divergence
    λ = rand()
    μ = rand()
    @test kldivergence(InverseGaussian(μ, 2λ), InverseGamma(1//2, λ)) ≈
        (log(2λ) / 2 - log(2π) / 2 - log(λ) / 2 + loggamma(1//2) + λ / μ)
end


"""
The module `Blocks` contains common input-output components, referred to as blocks.
"""
module Blocks
using ModelingToolkit, Symbolics, OrdinaryDiffEq
using IfElse: ifelse

@parameters t
D = Differential(t)

export RealInput, RealOutput, SISO
include("utils.jl")

export Gain, Sum, MatrixGain, Feedback, Add, Add3, Product, Division
export Abs, Sign, Sqrt, Sin, Cos, Tan, Asin, Acos, Atan, Atan2, Sinh, Cosh, Tanh, Exp
export Log, Log10
include("math.jl")

export Constant, Sine, Cosine, ContinuousClock, Ramp, Step, ExpSine, Square, Triangular
include("sources.jl")

export Limiter, DeadZone, SlewRateLimiter
include("nonlinear.jl")

export Integrator, Derivative, FirstOrder, SecondOrder, StateSpace
export PI, LimPI, PID, LimPID
include("continuous.jl")

end