"""
Solves the incompressible Navier-Stokes equations in three dimensions using a
pseudo-spectral approach involving the FFT. Simulates the 3D Taylor-Green vortex
from vortex stretching to fully turbulent to dissipation. The approach follows
the "SpectralDNS in Python" paper by Mortensen and Langtangen:
https://arxiv.org/pdf/1602.03638v1.pdf

Momentum:           ∂u/∂t + (u ⋅ ∇) u = − 1/ρ ∇p + ν ∇²u

Incompressibility:  ∇ ⋅ u = 0

u :  Velocity (3d vector)
p :  Pressure
ν :  Kinematic Viscosity
ρ :  Density (here=1)
t :  Time
∇ :  Nabla operator (defining nonlinear convection, gradient and divergence)
∇²:  Laplace Operator

------

Scenario:

The flow is simulated on a Ω = (0, 2π)³ domain (in 3D). The initial condition is
according to

    u(t=0, x, y, z) =   sin(x)cos(y)cos(z)

    v(t=0, x, y, z) = − cos(x)sin(y)cos(z)

    w(t=0, x, y, z) = 0

No forcing is applied.

------

Expected Outcome:

(Following page 18 of
https://doc.nektar.info/tutorials/latest/incns/taylor-green-vortex/incns-taylor-green-vortex.pdf)

1. Laminar Vortex sheets

2. Vortex stretching

3. Rearrangement

4. Breakdown

5. Fully turbulent

6. Dissipation & Decay

------

Solution Strategy:

u  : The velocity field (3D vector) in spatial domain
λ  : The velocity field (3D vector) in Fourier domain

ω  : The vorticity field (3D vector) in spatial domain
η  : The vorticity field (3D vector) in Fourier domain

m  : The product of vorticity and velocity (3D vector) in spatial domain
ϕ  : The product of vorticity and velocity (3D vector) in Fourier domain

k  : The wavenumber vector (3D vector)

ψ  : The pressure in Fourier domain (1D scalar)

b  : The rhs of the ODE system in Fourier domain (3D vector)

i  : The imaginary unit (1D scalar)

0. Initialize the velocity vectors according to the IC and transform them into
   Fourier Domain

1. Compute the curl in Fourier Domain by means of a cross
   product with the wavenumber vector and imaginary unit

    η = ℱ(ω) = ℱ(∇ × u) = i k × λ

2. Transform the vorticity back to spatial domain

    ω = ℱ⁻¹(η)

3. Compute the "convection" by means of a cross product

    m = u × ω

4. Transform the "convection" back into Fourier domain

    ϕ = ℱ(m)

5. Perform a dealising on the convection in order to suppress
   unresolved wave numbers

6. Compute the (pseudo) "pressure" in Fourier domain

    ψ = (k ⋅ ϕ) / ||k||₂

7. Compute the rhs to the ODE system

    b = m - ν ||k||₂² λ - (k ψ) / ||k||₂

8. Advance the velocity in Fourier Domain by means of an
   Explicit Euler time step

   λ ← λ + Δt b

9. Transform the newly obtained velocity back into spatial
   domain

   u = ℱ⁻¹(λ)

10. [Optional] visualize the vorticity magnitude in spatial
   domain interactively

11. Repeat from (1.)

In total, it takes us three (three-dimensional) Fourier Transforms
per time iteration:
- Transformation of the curl to spatial domain
- Transformation of the "convection" to Fourier Domain
- Transformation of the velocity to spatial domain

It is called pseudo-spectral because some operations are performed
in Fourier Domain and some in the spatial domain.

-------

Stability Consideration

Choose the time step size carefully.

-----

Additional Hints:

Check out the Python Repo:
https://github.com/spectralDNS/spectralDNS

In particular this file:
https://github.com/spectralDNS/spectralDNS/blob/master/spectralDNS3D_short.py

"""

using FFTW
using GLMakie
using Statistics
using ProgressMeter

N_POINTS_P_AXIS = 50
KINEMATIC_VISCOSITY = 1.0 / 1_600
TIME_STEP_LENGTH = 0.02
N_TIME_STEPS = 1_700
PLOT_EVERY = 3

function cross_product!(
    res_x::Array{T, 3},
    res_y::Array{T, 3},
    res_z::Array{T, 3},
    a_x::Array{T, 3},
    a_y::Array{T, 3},
    a_z::Array{T, 3},
    b_x::Array{T, 3},
    b_y::Array{T, 3},
    b_z::Array{T, 3},
) where T<:Union{Float64, Complex{Float64}}
    res_x .= (
        a_y .* b_z
        .-
        a_z .* b_y
    )

    res_y .= (
        a_z .* b_x
        .-
        a_x .* b_z
    )

    res_z .= (
        a_x .* b_y
        .-
        a_y .* b_x
    )
end

function main()
    x_range = range(0.0, 2*pi, N_POINTS_P_AXIS)
    y_range = range(0.0, 2*pi, N_POINTS_P_AXIS)
    z_range = range(0.0, 2*pi, N_POINTS_P_AXIS)

    coordinates_x = [x for x in x_range, y in y_range, z in z_range]
    coordinates_y = [y for x in x_range, y in y_range, z in z_range]
    coordinates_z = [z for x in x_range, y in y_range, z in z_range]

    wavenumbers_1d_x = fftfreq(N_POINTS_P_AXIS) .* N_POINTS_P_AXIS
    wavenumbers_1d_y = fftfreq(N_POINTS_P_AXIS) .* N_POINTS_P_AXIS
    wavenumbers_1d_z = fftfreq(N_POINTS_P_AXIS) .* N_POINTS_P_AXIS

    wavenumbers_x = [k_x for k_x in wavenumbers_1d_x, k_y in wavenumbers_1d_y, k_z in wavenumbers_1d_z]
    wavenumbers_y = [k_y for k_x in wavenumbers_1d_x, k_y in wavenumbers_1d_y, k_z in wavenumbers_1d_z]
    wavenumbers_z = [k_z for k_x in wavenumbers_1d_x, k_y in wavenumbers_1d_y, k_z in wavenumbers_1d_z]

    wavenumbers_norm = sqrt.(
        wavenumbers_x.^2
        +
        wavenumbers_y.^2
        +
        wavenumbers_z.^2
    )

    wavenumbers_norm[iszero.(wavenumbers_norm)] .= 1.0
    normalized_wavenumbers_x = wavenumbers_x ./ wavenumbers_norm
    normalized_wavenumbers_y = wavenumbers_y ./ wavenumbers_norm
    normalized_wavenumbers_z = wavenumbers_z ./ wavenumbers_norm

    # Taylor-Green vortex initial condition
    velocity_x =   sin.(coordinates_x) .* cos.(coordinates_y) .* cos.(coordinates_z)
    velocity_y = - cos.(coordinates_x) .* sin.(coordinates_y) .* cos.(coordinates_z)
    velocity_z =   zeros((N_POINTS_P_AXIS, N_POINTS_P_AXIS, N_POINTS_P_AXIS))

    # Pre-Plan the FFT for additional performance
    fft_operator = plan_fft(velocity_x; flags=FFTW.MEASURE)

    velocity_x =   sin.(coordinates_x) .* cos.(coordinates_y) .* cos.(coordinates_z)

    # Pre-Allocatearrays
    velocity_x_fft = fft_operator * velocity_x
    velocity_y_fft = fft_operator * velocity_y
    velocity_z_fft = fft_operator * velocity_z

    curl_x = zeros(Float64, (N_POINTS_P_AXIS, N_POINTS_P_AXIS, N_POINTS_P_AXIS))
    curl_y = zeros(Float64, (N_POINTS_P_AXIS, N_POINTS_P_AXIS, N_POINTS_P_AXIS))
    curl_z = zeros(Float64, (N_POINTS_P_AXIS, N_POINTS_P_AXIS, N_POINTS_P_AXIS))

    curl_x_squared = zeros(Float64, (N_POINTS_P_AXIS, N_POINTS_P_AXIS, N_POINTS_P_AXIS))
    curl_y_squared = zeros(Float64, (N_POINTS_P_AXIS, N_POINTS_P_AXIS, N_POINTS_P_AXIS))
    curl_z_squared = zeros(Float64, (N_POINTS_P_AXIS, N_POINTS_P_AXIS, N_POINTS_P_AXIS))

    curl_magnitude = zeros(Float64, (N_POINTS_P_AXIS, N_POINTS_P_AXIS, N_POINTS_P_AXIS))

    curl_x_fft = zeros(Complex{Float64}, (N_POINTS_P_AXIS, N_POINTS_P_AXIS, N_POINTS_P_AXIS))
    curl_y_fft = zeros(Complex{Float64}, (N_POINTS_P_AXIS, N_POINTS_P_AXIS, N_POINTS_P_AXIS))
    curl_z_fft = zeros(Complex{Float64}, (N_POINTS_P_AXIS, N_POINTS_P_AXIS, N_POINTS_P_AXIS))

    convection_x = zeros(Float64, (N_POINTS_P_AXIS, N_POINTS_P_AXIS, N_POINTS_P_AXIS))
    convection_y = zeros(Float64, (N_POINTS_P_AXIS, N_POINTS_P_AXIS, N_POINTS_P_AXIS))
    convection_z = zeros(Float64, (N_POINTS_P_AXIS, N_POINTS_P_AXIS, N_POINTS_P_AXIS))

    convection_x_fft = zeros(Complex{Float64}, (N_POINTS_P_AXIS, N_POINTS_P_AXIS, N_POINTS_P_AXIS))
    convection_y_fft = zeros(Complex{Float64}, (N_POINTS_P_AXIS, N_POINTS_P_AXIS, N_POINTS_P_AXIS))
    convection_z_fft = zeros(Complex{Float64}, (N_POINTS_P_AXIS, N_POINTS_P_AXIS, N_POINTS_P_AXIS))

    # De-Aliasing is necessary to stabilize the simulation
    k_max_dealias = 2.0/3.0 * (N_POINTS_P_AXIS//2 + 1)
    dealias = (
        (abs.(wavenumbers_x) .< k_max_dealias)
        .*
        (abs.(wavenumbers_y) .< k_max_dealias)
        .*
        (abs.(wavenumbers_z) .< k_max_dealias)
    )

    # Prepare the visualization
    makie_plot_figure_window = Figure()
    display(makie_plot_figure_window)

    makie_plot_axis = Axis3(
        makie_plot_figure_window[1, 1],
        perspectiveness = 0.3,
    )

    makie_plot_3d_contour = contour!(
        makie_plot_axis,
        x_range,
        y_range,
        z_range,
        curl_magnitude,
        alpha = 0.05,
        levels = range(0.2, 1.0, length=4),
    )
    println("Hit enter when plot is ready")
    readline()

    @showprogress "Simulating & Animating ..." for t in 1:N_TIME_STEPS
        # (1) Compute the Curl in Fourier Domain
        cross_product!(
            curl_x_fft,
            curl_y_fft,
            curl_z_fft,
            im .* wavenumbers_x,
            im .* wavenumbers_y,
            im .* wavenumbers_z,
            velocity_x_fft,
            velocity_y_fft,
            velocity_z_fft,
        )

        # (2) Transform curl to spatial domain
        curl_x .= real(fft_operator \ curl_x_fft)
        curl_y .= real(fft_operator \ curl_y_fft)
        curl_z .= real(fft_operator \ curl_z_fft)

        # (3) Compute "Convection" in spatial domain
        cross_product!(
            convection_x,
            convection_y,
            convection_z,
            velocity_x,
            velocity_y,
            velocity_z,
            curl_x,
            curl_y,
            curl_z,
        )

        # (4) Transform "Convection" to Fourier Domain
        convection_x_fft .= fft_operator * convection_x
        convection_y_fft .= fft_operator * convection_y
        convection_z_fft .= fft_operator * convection_z

        # (5) Dealiasing on the higher wavenumbers
        convection_x_fft .*= dealias
        convection_y_fft .*= dealias
        convection_z_fft .*= dealias

        # (6) Compute the Pseudo-Pressure by a Divergence in Fourier Domain
        pressure_fft = (
            normalized_wavenumbers_x .* convection_x_fft
            +
            normalized_wavenumbers_y .* convection_y_fft
            +
            normalized_wavenumbers_z .* convection_z_fft
        )

        # (7) Assemble the rhs to the ODE system in Fourier Domain
        rhs_x_fft = (
            convection_x_fft
            -
            KINEMATIC_VISCOSITY * wavenumbers_norm.^2 .* velocity_x_fft
            -
            normalized_wavenumbers_x .* pressure_fft
        )
        rhs_y_fft = (
            convection_y_fft
            -
            KINEMATIC_VISCOSITY * wavenumbers_norm.^2 .* velocity_y_fft
            -
            normalized_wavenumbers_y .* pressure_fft
        )
        rhs_z_fft = (
            convection_z_fft
            -
            KINEMATIC_VISCOSITY * wavenumbers_norm.^2 .* velocity_z_fft
            -
            normalized_wavenumbers_z .* pressure_fft
        )

        # (8) Euler Step Update
        velocity_x_fft .+= rhs_x_fft * TIME_STEP_LENGTH
        velocity_y_fft .+= rhs_y_fft * TIME_STEP_LENGTH
        velocity_z_fft .+= rhs_z_fft * TIME_STEP_LENGTH

        # (9) Transform the velocities back to spatial domain
        velocity_x .= real(fft_operator \ velocity_x_fft)
        velocity_y .= real(fft_operator \ velocity_y_fft)
        velocity_z .= real(fft_operator \ velocity_z_fft)

        # VISUALIZATION
        if t % PLOT_EVERY == 0
            curl_x_squared .= curl_x.^2
            curl_y_squared .= curl_y.^2
            curl_z_squared .= curl_z.^2
            curl_magnitude .= sqrt.(
                curl_x_squared
                +
                curl_y_squared
                +
                curl_z_squared
            )

            delete!(makie_plot_3d_contour.parent, makie_plot_3d_contour)

            makie_plot_3d_contour = contour!(
                makie_plot_axis,
                x_range,
                y_range,
                z_range,
                curl_magnitude,
                alpha = 0.01,
                levels = range(1.0, 10.0, length=5),
            )

            sleep(0.02)

        end

    end

end

main()


### functions for posterior odds and Bayes p-value
#
"""
    todds(theta_hat,theta_hat_se,v)

Input:

        v = degrees of freedom

        theta_hat, theta_hat_se
        = estimate and standard error for theta

Optional argument:

        h0 = value in hull hypothesis (default = 0).

Returns:

        t-posterior odds ratio of evidence against h0 (default = 0).
"""
function todds(theta_hat, theta_se, v; h0 = 0.0)
  # compute posterior odds assuming a Student-t distribution
  # mcs = MC sample
  # v = degrees of freedom, n - k
  # h0 = value under null hypothesis
  that = abs.(theta_hat - h0)/theta_se
  odds = (1.0 + (that^2)/v)^((v+1)/2.0)
end

"""
    mcodds(theta_draws)

Input:

        theta_draws = MC sample from posterior for theta

Optional argument:

        h0 = value in hull hypothesis (default = 0).

Returns:

        posterior odds ratio of evidence against h0.
"""
function mcodds(mcs; h0=0.0)
  # dependency: KernelDensity
  # compute posterior odds from nonparametric density
  # mcs = MC sample

### NEED catches for when odds too small/large - out of bounds of sample
  postden = KernelDensity.kde(mcs)
  ax = abs.(postden.x .- h0)
  indx = findall(ax .== minimum(ax))  # find all x that = 4 in x
  numodds = maximum(postden.density)
  denodds = postden.density[indx[1]]
  odds = numodds/denodds[1]
  if odds > 5000000.0
    odds = 5000000.0
  end
  return odds
end

# Example: if bs[:,2] is vector of MCMC draws for beta
# to test if beta = 1.00:
#mcodds(bs[:,2],h0=1.00)

"""
    pdr_pval(theta_draws)

Input:

        theta_draws = MC sample from posterior for theta

Optional argument:

        h0 = value in hull hypothesis (default = 0).

Returns:

        posterior density ratio, one-tailed and two-tailed 'p-value' of evidence against h0.
        one-tailed 'p-value' = minimum[Prob(theta < 0), Prob(theta>0)]
"""
# x = MC sample, h0 = null hypothesis value
# usage: odds, p_val, pval_2 = post_odds_pval(x, h0 = -2.3)
function pdr_pval(x; h0 = 0.0)
    d = kde(x)
    p,ind = findmax(d.density)
    ph0 = ifelse(pdf(d,h0) == 0.0,0.0000001,pdf(d,h0))
    odds = pdf(d,d.x[ind])/ph0
    p_val = length(x[x .<= h0])
    p_val2 = length(x[x .>= h0])
    if p_val <= p_val2
        p_value = p_val/length(x)
    else
        p_value = p_val2/length(x)
    end
    p_value_2tail = 2*p_value
    return odds, p_value, p_value_2tail
end

"""
    bayespval(theta_draws)

Input:

        theta_draws = MC sample from posterior for theta

Optional argument:

        h0 = value in hull hypothesis (default = 0).

Returns:

        posterior 'p-value' of evidence against h0.
        pval = one tailed value minimum[(Prob(theta <0), Prob(theta>0)]
        pval2 = two-tailed value (2*pval)
"""
function bayespval(mcs; h0=0.0)
  p_val = length(mcs[mcs .<= h0])/length(mcs)
  if p_val == 0.0
    p_val = 0.00001
  end
  if p_val > 0.5
    p_val = (1.0 - p_val)
  end
  p_val2 = 2.0*p_val
  return p_val, p_val2
end

#= old deprecated todds assuming t
function todds(mcs,n,k; h0 = 0.0)
  # compute posterior odds assuming a Student-t distribution
  # mcs = MC sample
  # v = degrees of freedom, n - k
  # h0 = value under null hypothesis
  v = n-k
  unbiasedseb = std(mcs)*sqrt(v)/sqrt(n)
  that = abs.(mean(mcs) - h0)/unbiasedseb
  odds = (1 + (that^2)/v)^((v+1)/2)
end
=#


l_vertex = 2
center = 5
r_vertex = 7

mf = TriangularMF(l_vertex, center, r_vertex)

# Evaluation tests
@assert mf.eval(center) == 1
@assert mf.eval(l_vertex) == mf.eval(r_vertex) == 0
@assert mf.eval((l_vertex + center) / 2) == mf.eval((r_vertex + center) / 2) == 0.5

# Mean finding tests
@assert mf.mean_at(1) == center
@assert mf.mean_at(0) == (l_vertex + r_vertex) / 2


#=
Functions that automatically create Goppa codes.

Author: Ivan A. Moreno Soto
Last updated: 2019/July/15
=#

module GoppaCode

export
    generateGoppaCode,
    decode!

include("./Polynomials.jl")
using .Polynomials

include("./ExtensionField.jl")
using .ExtensionField

include("./FiniteFieldMatrix.jl")
using .FiniteFieldMatrix

include("./Misc.jl")
using .Misc

using StatsBase
using LinearAlgebra

using DelimitedFiles

"""
    selectCodeSupport(n::Int, m::Int)

Randomly selects and returns n distinct elements from the
extension field 2^m.
"""
function selectCodeSupport(n::Int, m::Int)
    return sample(Array{UInt32, 1}(0:(2^m - 1)), n, replace=false)
end

"""
    getY!(codeSupport, t, g)

Creates the Y matrix needed to compute the parity check matrix H. Basically
computes a Vandermonde matrix with the code support.
"""
function getY!(codeSupport, t, g)
    n = length(codeSupport)
    Y = zeros(UInt32, (t, n))

    for i in 1:n
        Y[1, i] = 1
    end

    for i in 1:n
        for j in 2:t
            Y[j, i] = power(codeSupport[i], (j-1), g)
        end
    end

    return Y
end

"""
    getZ!(goppaPolynomial, moduloPolynomial, codeSupport)

Creates the Z matrix needed to compute the parity check matrix H.
"""
function getZ!(goppaPolynomial, moduloPolynomial, codeSupport)
    n = length(codeSupport)
    Z = zeros(UInt32, (n, n))

    for i in 1:n
        Z[i, i] = inverse(evaluate!(codeSupport[i], goppaPolynomial, moduloPolynomial), moduloPolynomial)
    end

    return Z
end

"""
    toBits!(matrix, t, n, m)

Takes a matrix of dimension txn and returns a binary matrix
with m-bits representations of every element spread vertically.
"""
function toBits!(matrix, t, n, m)
    bits = falses(m*t, n)

    for i in 1:t
        for j in 1:n
            bits[((i - 1) * m + 1):(i * m), j] = digits(matrix[i, j], base=2, pad=m)
        end
    end

    return bits
end

"""
    generateGoppaCode(n::Int, t::Int, m::Int)

Generates a (n, k) irreducible linear Goppa code with minimum
distance d >= 2t + 1. Returns the generator matrix G, the goppa
polynomial, the code support, and the number of rows k of G.
"""
function generateGoppaCode(n::Int, t::Int, m::Int)
    # Get polynomials
    moduloPolynomial = primitivePolynomials[m]
    goppaPolynomial = randomPoly(t, m)
    # Get code support
    codeSupport = selectCodeSupport(n, m)

    # Compute H = YZ
    Y = getY!(codeSupport, t, moduloPolynomial)
    Z = getZ!(goppaPolynomial, moduloPolynomial, codeSupport)
    T = multiplyPolyMatrices!(Y, Z, moduloPolynomial)

    H = toBits!(T, t, n, m)
    gaussianEliminationColumnPivoting!(H, m*t, codeSupport)

    # Compute G
    k = n - m*t
    ATransposed = transpose(H[:, (m*t + 1):end])
    G = hcat(ATransposed, Matrix{UInt32}(I, k, k))

    return G, goppaPolynomial, codeSupport, k
end

"""
    syndrome!(msg, codeSupport, goppaPolynomial, m)

Computes the syndrome of a vector using a binary irreducible goppa code.
"""
function syndrome!(msg, codeSupport, goppaPolynomial, m)
    syndrome = Array{UInt32, 1}([0])
    for (c, γ) in zip(msg, codeSupport)
        if c == 1
            currentInverse = inverseRing!(Array{UInt32, 1}([γ, 1]), goppaPolynomial, m)
            syndrome = addRing!(syndrome, currentInverse)
        end
    end
    return syndrome
end

"""
    syndromeSquareRoot!(syndrome, goppaPolynomial, m)

Computes the square root of syndrome + x mod goppaPolynomial.
"""
function syndromeSquareRoot!(syndrome, goppaPolynomial, m)
    # Computing T(x) + x.
    syndromeInverse = inverseRing!(syndrome, goppaPolynomial, m)
    h = addRing!(syndromeInverse, PolynomialConstants['x'])

    # Computing the square root of x mod goppaPolynomial.
    goppaEven, goppaOdd = splitPolynomial!(goppaPolynomial, m)
    sx = multiplyRing!(goppaEven, inverseRing!(goppaOdd, goppaPolynomial, m), m)

    # Computing the odd and even parts of the square root of T(x) + x.
    evenh = [multiplyRing!(power(coeff, 2^(m-1), primitivePolynomials[m]), monomial(i-1), m) for (i, coeff) in enumerate(h[1:2:end])]
    oddh = [multiplyRing!(multiplyRing!([power(coeff, 2^(m-1), primitivePolynomials[m])], monomial(i-1), m), sx, m) for (i, coeff) in enumerate(h[2:2:end])]
    if length(oddh) == 0
        oddh = [copy(PolynomialConstants['0'])]
    end

    # Computing the square root by adding the even and odd parts.
    τ = copy(PolynomialConstants['0'])
    for term in vcat(evenh, oddh)
        τ = addRing!(τ, term)
    end

    return τ
end

"""
    halfwayXGCD!(goppaPolynomial, tau, t, m)

Runs the extended eucliedean algorithm halfway to compute the least
degree polynomials that can define the error locator polynomial for
a codeword with errors.
"""
function halfwayXGCD!(goppaPolynomial, tau, t, m)
    r = copy(goppaPolynomial)
    R = copy(tau)
    s = Array{UInt32, 1}([1])
    S = Array{UInt32, 1}([0])
    u = Array{UInt32, 1}([0])
    U = Array{UInt32, 1}([1])

    while length(R) - 1 > div(t, 2)
        q = divRing!(r, R, m)

        x = R; R = addRing!(r, multiplyRing!(q, R, m)); r = x
        x = S; S = addRing!(s, multiplyRing!(q, S, m)); s = x
        x = U; U = addRing!(u, multiplyRing!(q, U, m)); u = x
    end

    return R, U
end

"""
    findRootsErrorLocatorPolynomial!(sigma, codeSupport, m)

Finds the roots of sigma by evaluating every element of the code support.
"""
function findRootsErrorLocatorPolynomial!(sigma, codeSupport, m)
    totalErrors = 0
    errors = []
    for (i, gamma) in enumerate(codeSupport)
        if evaluate!(gamma, sigma, primitivePolynomials[m]) == 0
            push!(errors, i)
            totalErrors += 1
        end
    end

    return errors, totalErrors
end

"""
    decode!(msg, goppaPolynomial, codeSupport, m, t)

Given a msg with noise, finds a codeword in the goppa code defined
by the goppaPolynomial and codeSupport in GF(2^m).
"""
function decode!(msg, goppaPolynomial, codeSupport, m, t)
    # Compute syndrome
    msgSyndrome = syndrome!(msg, codeSupport, goppaPolynomial, m)
    if msgSyndrome == PolynomialConstants['0']
        throw(InvalidCodewordError())
    end

    tau = syndromeSquareRoot!(msgSyndrome, goppaPolynomial, m)

    # Using the extended euclidean algorithm to find the error locator polynomial.
    a, b = halfwayXGCD!(goppaPolynomial, tau, t, m)
    sigma = addRing!(multiplyRing!(a, a, m), multiplyRing!(PolynomialConstants['x'], multiplyRing!(b, b, m), m))
    sigma = makePolynomialMonic!(sigma, m)

    # Finding the roots of the error locator polynomial.
    errors, totalErrors = findRootsErrorLocatorPolynomial!(sigma, codeSupport, m)

    # Checking that it was a valid cyphertext.
    if totalErrors < t
        throw(InvalidCodewordError())
    end

    # Correcting the t errors.
    for position in errors
        msg[position] = (msg[position] + 1) % 2
    end

    return msg
end

end # GoppaCode


@node Categorical Stochastic [ out, p ]

@average_energy Categorical (q_out::Categorical, q_p::Any) = -sum(probvec(q_out) .* logmean(q_p))


using SymPy
using LinearAlgebra
λ = μ = 1.0

function gen_code(u)
    s = replace(replace(replace(sympy.julia_code(simplify(u)), ".*"=>"*"), ".^"=>"^"), "./"=>"/")
    "  (x,y)->begin;$s;end,\n"
end

function gen_code(u, v, g, h)
    s1 = gen_code(u)
    s2 = gen_code(v)
    s3 = gen_code(g)
    s4 = gen_code(h)
    "(\n" * s1 * s2 * s3 * s4 * ")"
end

x, y = @vars x y
# u = x*(1-x)*y*(1-y)
# v = x*(1-x)*y*(1-y)

# u = x^2 * (1-x) * y^2 * (1-y^2)
# v = u

u = x^2 * (1-x) * y^2 * (1-y^2)
v = x*(1-x)*y*(1-y)

ux = diff(u, x)
uy = diff(u, y)
vx = diff(v, x)
vy = diff(v, y)

ε = [
    ux (uy+vx)/2
    (uy+vx)/2 vy
]

σ = 2μ * ε + λ*(ux + vy) * diagm(0=>ones(2))

g = diff(σ[1,1], x) + diff(σ[1,2], y)
h = diff(σ[2,1], x) + diff(σ[2,2], y)

println(gen_code(u, v, g, h))


@defcomp neteconomy begin

    ABATECOST   = Variable(index=[time])    #Cost of emissions reductions  (trillions 2005 USD per year)
    C           = Variable(index=[time])    #Consumption (trillions 2005 US dollars per year)
    CPC         = Variable(index=[time])    #Per capita consumption (thousands 2005 USD per year)
    CPRICE      = Variable(index=[time])    #Carbon price (2005$ per ton of CO2)
    I           = Variable(index=[time])    #Investment (trillions 2005 USD per year)
    MCABATE     = Variable(index=[time])    #Marginal cost of abatement (2005$ per ton CO2)
    Y           = Variable(index=[time])    #Gross world product net of abatement and damages (trillions 2005 USD per year)
    YNET        = Variable(index=[time])    #Output net of damages equation (trillions 2005 USD per year)

    cost1       = Parameter(index=[time])   #Abatement cost function coefficient
    DAMAGES     = Parameter(index=[time])   #Damages (Trillion $)
    l           = Parameter(index=[time])   #Level of population and labor
    MIU         = Parameter(index=[time])   #Emission control rate GHGs
    partfract   = Parameter(index=[time])   #Fraction of emissions in control regime
    pbacktime   = Parameter(index=[time])   #Backstop price
    S           = Parameter(index=[time])   #Gross savings rate as fraction of gross world product
    YGROSS      = Parameter(index=[time])   #Gross world product GROSS of abatement and damages (trillions 2005 USD per year)
    expcost2    = Parameter()               #Exponent of control cost function

    function run_timestep(p, v, d, t)
        #Define function for YNET
        v.YNET[t] = p.YGROSS[t] - p.DAMAGES[t]

        #Define function for ABATECOST
        v.ABATECOST[t] = p.YGROSS[t] * p.cost1[t] * (p.MIU[t]^p.expcost2) * (p.partfract[t]^(1 - p.expcost2))

        #Define function for MCABATE (equation from GAMS version)
        v.MCABATE[t] = p.pbacktime[t] * p.MIU[t]^(p.expcost2 - 1)

        #Define function for Y
        v.Y[t] = v.YNET[t] - v.ABATECOST[t]

        #Define function for I
        v.I[t] = p.S[t] * v.Y[t]

        #Define function for C
        v.C[t] = v.Y[t] - v.I[t]

        #Define function for CPC
        v.CPC[t] = 1000 * v.C[t] / p.l[t]

        #Define function for CPRICE (equation from GAMS version of DICE2013)
        v.CPRICE[t] = p.pbacktime[t] * (p.MIU[t] / p.partfract[t])^(p.expcost2 - 1)
    end
end


using Random
using Plots
import Distributions
include("functions.jl")

mutable struct Net
    weights::Vector{Matrix}
    biases::Vector{Vector}
    f::Vector{Function}

    function Net(sizes::Int...)
        dist = Distributions.Uniform(-0.1, 0.1)

        weights = Vector{Matrix}()
        biases = Vector{Vector}()

        for i in 1:length(sizes)-1
            w = rand(dist, sizes[i+1], sizes[i])
            b = rand(dist, sizes[i+1])

            push!(weights, w)
            push!(biases, b)
        end

        new(weights, biases, [fill(vtanh, length(sizes) - 2); softmax])
    end
end

mean(x) = sum(x)/length(x)
relu(x::Vector) = max.(x, 0)

function softmax(x::Vector)
    m = maximum(x)
    s = exp.(x .- m)
    s/sum(s)
end

function chop(v::Vector, s::Int)
    res = Vector{Vector}()
    l = length(v)

    for i in 1:s:l
        stop = i+s-1
        stop = stop <= l ? stop : l

        push!(res, v[i:stop])
    end

    res
end

function forward(net::Net, x::AbstractVector)
    outputs = Vector{AbstractVector}()
    push!(outputs, x)

    for (W, b, f) in zip(net.weights, net.biases, net.f)
        raw = W*outputs[end] + b
        push!(outputs, f(raw))
    end

    outputs
end

predict(net::Net, x::AbstractVector) = argmax(forward(net, x)[end])

function accuracy(net, x::Vector, t::Vector)
    predictions = [predict(net, i) for i in x]
    cmp = isequal.(predictions, t)
    sum(cmp)/length(x)
end

function update_weights(net::Net, batch_x::Vector, batch_t::Vector, α::Number, regularization::Number)
    out_grad(o, t) = o - [Int(==(i, t)) for i in 1:length(o)]
    hid_grad(W, g, h) = dtanh(h) * (W' * g)
    dW(g, h) = g*h'

    outputs = map(x -> forward(net, x), batch_x)
    layers = collect(zip(outputs...))
    output_layer_gradient = out_grad.(layers[end], batch_t)

    net.biases[end] -= α*mean(output_layer_gradient)
    net.weights[end] -= α*(mean(dW.(output_layer_gradient, layers[end-1])) + regularization * net.weights[end])

    grad = output_layer_gradient
    for i in (length(net.weights)-1):-1:1
        grad = [hid_grad(net.weights[i+1], g, l) for (g, l) in zip(grad, layers[i+1])]

        net.biases[i] -= α*mean(grad)
        net.weights[i] -= α*(mean(dW.(grad, layers[i])) + regularization*net.weights[i])
    end
end

function cost(net, x, t)
    o = forward(net, x)[end]
    -log(o[t])
end

function train(net::Net, x::Vector, t::Vector, epochs::Int, batch_size::Int, rate::Number, decay::Number = 1; regularization::Number = 0)
    plotly()

    indices = 1:length(x)
    test_indices = shuffle(indices)[1:2*batch_size]
    train_indices = filter(i -> !(i in test_indices), indices)

    mc = Vector{Float64}()
    test_accs = Vector{Float64}()
    train_accs = Vector{Float64}()

    for epoch in 1:epochs
        batch_indices = chop(shuffle(train_indices), batch_size)

        costs = nothing
        for bi in batch_indices
            update_weights(net, x[bi], t[bi], rate, regularization)
            costs = [cost(net, x[bi][j], t[bi][j]) for j in 1:length(bi)]
        end

        mean_cost = mean(costs)
        test_acc = accuracy(net, x[test_indices], t[test_indices])
        train_acc = accuracy(net, x[train_indices], t[train_indices])

        println("Epoch $epoch/$epochs")
        println("\tTrain accuracy: $(train_acc * 100) %")
        println("\tTest accuracy: $(test_acc * 100) %")
        println("\tMean cost: $mean_cost")
        println("\tLearning rate: $rate")

        push!(mc, mean_cost)
        push!(test_accs, test_acc)
        push!(train_accs, train_acc)

        rate /= decay
    end

    plt = plot(1:epochs, mc, xlabel = "Epoch", ylabel = "Value", label = "Mean cost", line = [:scatter], grid = false) # dpi = 600, size = (1440, 810)
    plot!(plt, 1:epochs, test_accs, line = [:scatter], label = "Test accuracy")
    plot!(plt, 1:epochs, train_accs, line = [:scatter], label = "Train accuracy")
    display(plt)
end

using Serialization
using Mmap
data = deserialize("trainset/data2")
targets = deserialize("trainset/targets2")

#data = [Float64.(Gray.(d)) for d in data]

N = 310
l = 900
#data = mmap("trainset/trainst3", BitMatrix, (l, N))
#targets = mmap("trainset/targets3", BitMatrix, (3, N))

#data = [data[:, i] for i in 1:N]
#targets = [argmax(targets[:, i]) for i in 1:N]

net = Net(l, 10, 3)
train(net, data, targets, 750, 50, 2e-2, regularization = 1e-2)


using ProductizeBayes

bernoulli_model = "
data {
  int<lower=1> N;
  int<lower=0,upper=1> y[N];
}
parameters {
  real<lower=0,upper=1> theta;
}
model {
  theta ~ beta(1,1);
  y ~ bernoulli(theta);
}
";

data = Dict("N" => 10, "y" => [0, 1, 0, 1, 0, 0, 0, 0, 0, 1])

sm = VersionedModel("bernoulli", bernoulli_model)
rc = stan_sample(sm.stan_model; data, num_chains=2, seed=-1)


export checkDerivative

"""
checkDerivative(f,x0;out,tol,nSuccess)

Check if function f provides correct derivatives around point x0 using a Taylor expansion.

Inputs:

	f::Function  - fx,dfx= f(x) provides function value and derivative
	x::Array     - vector around which to test the derivatives
	out::Bool    - flag for activating printing (default: true)
	tol::Number  - tolerance to judge if error in linearization decreases quadratically (default: 1.9)
	nSuccess:Int - number of quadratically-decreasing steps needed to decide that derivative is correct

Outputs:
	pass, Error, Order

"""
function checkDerivative(f::Function,x0;out::Bool=true,tol::Number=1.9,nSuccess::Int=3)
	if out
		println(@sprintf("%9s\t%9s\t%9s\t%9s\t%9s\t%5s","h","E0","E1","O1","O2","OK?"))
	end
	v = randn(size(x0))
	if eltype(x0) <: Complex
		v += 1im*randn(size(v))
	end
	f0,dvf = f(x0,v)
	dvf   = real(dvf)
	Error = zeros(10,2)
	Order = zeros(10,2)
	Success = zeros(10)
	for j=1:10
		ft = f(x0+10.0^(-j)*v)            # function value
		Error[j,1] = norm(f0-ft)/norm(f0)           # Error TaylorPoly 0
		Error[j,2] = norm(f0 .+10.0^(-j)*dvf .- ft)/norm(f0) # Error TaylorPoly 1
		if j>1
			Order[j,:] = log10.(Error[j-1,:]./Error[j,:]);
		end
		if (Order[j,2]>tol) || (Error[j,1]/Error[j,2] > 100); Success[j]=1; end
		if out
			println(@sprintf("%1.3e\t%1.3e\t%1.3e\t%1.3e\t%1.3e\t%5d",
							10.0^(-j), Error[j,1],Error[j,2], Order[j,1],Order[j,2],Success[j]))
		end
	end
	pass = sum(Success) > nSuccess
	return  pass,Error,Order
end


### A Pluto.jl notebook ###
# v0.12.21

using Markdown
using InteractiveUtils

# ╔═╡ a29c2ea0-6ff4-11eb-2954-372c6727bee2
begin
	import Pkg
	Pkg.activate(mktempdir())
	Pkg.Registry.update()

	Pkg.add("Plots")
	using Plots
end

# ╔═╡ 724acb50-6db3-11eb-07a8-7d048b8e4b3e
md"# Introduction to Elementary Cellular Automata"

# ╔═╡ 9464f530-6db3-11eb-0629-9f05df1979bd
md"
Cellular automata are fascination \"mathematical machines\" which, despite normally being associated with having relatively simple rules, can still produce highly complex output. This point is stressed over and over in Stephen Wolfram's *A New Kind of Science*. This is how Wikipedia defines them:
> A cellular automaton is a discrete model of computation studied in automata theory.
In this notebook we will explore some simple cellular automata, called the *Elementary Cellular Automata*.
"

# ╔═╡ f3f9f3b0-6db3-11eb-1446-cfe467e42fee
md"## Defining Rules"

# ╔═╡ fdc546b0-6db3-11eb-1a4c-514e32ae4a27
md"Put simply, a cellular automata takes some simple rule and applies it over and over on an infinite string of numbers. We will be taking a look at 1 dimensional cellular automata, which is to say there is one spatial dimension and one temporal dimension. This can easily be visualized in a 2D image unlike higher-dimensional automata. Chances are you have heard of Conway's *Game of Life*, which is an example of a two dimensional cellular automaton."

# ╔═╡ 87bd67d0-6db4-11eb-0931-f9a3ec94cd70
md"Let's define some rules as a dictionary of associations"

# ╔═╡ b3fdfd00-6db4-11eb-3fa3-79d5911c567d
rule = Dict(
	0b000 => 0,
	0b001 => 1,
	0b010 => 1,
	0b011 => 1,
	0b100 => 1,
	0b101 => 0,
	0b110 => 0,
	0b111 => 0
)

# ╔═╡ f21ba14e-6db4-11eb-382a-87bc07250f65
md"For our rule, each cell can have a value of either a $0$ or a $1$. To determine what the cell's next value will be, our rule looks at the value of its left and right neighbors, where the left neighbor is the most significant digit and the right neighbor is the least significant digit, and associates that configuration, which has $2^3=8$ different combinations, to a new value of either $1$ or $0$."

# ╔═╡ 7895f0e0-6db6-11eb-1d6e-b12880ebd5b9
md"## Executing Rules

To start out we can define a recursive function which will return a list of lists, where each list represents the state of the automaton at a particular timestep."

# ╔═╡ 0ab21310-6db6-11eb-00b5-1f2bd2203c79
function value_at(state, pos)
	if pos < 1 || pos > length(state)
		return 0
	end
	return state[pos]
end

# ╔═╡ 656f4030-6db5-11eb-0360-bd9ac3ac6b45
function automaton(initial_state, steps, rule)
	# Base case
	if steps <= 1
		return [initial_state]
	end

	# Otherwise...
	previous_states = automaton(initial_state, steps - 1, rule)
	ps = previous_states[end]
	next_state = [rule[value_at(ps, i-1) << 2 | value_at(ps, i) << 1 | value_at(ps, i+1)] for i = 1:length(initial_state)]

	[previous_states..., next_state]
end

# ╔═╡ 232031c0-6db6-11eb-2f48-81afe98ccac3
md"Despite not being the most efficient it will do for smaller step counts. Let's try it out for our rule for 50 steps"

# ╔═╡ 71d85310-6db6-11eb-0728-f34323017fe1
history = automaton([i == 51 ? 1 : 0 for i = 1:101], 50, rule)

# ╔═╡ aefbdff0-6db6-11eb-0441-a924d0f567c1
history_mat = [history[i][j] for i=1:length(history), j=1:length(history[1])]

# ╔═╡ ea3c54f0-6db6-11eb-1312-3daa47f3fa91
md"## Rendering the Results"

# ╔═╡ f018e190-6db6-11eb-3eef-c984453d36d9
md"Now that we have a matrix of values we can make use of the Plots package to make a plot displaying it."

# ╔═╡ 06c775f0-6db7-11eb-2603-4942353b201e
heatmap(1 .- history_mat, color = :greys, yflip = true, size=(640, 320))

# ╔═╡ 8bb7e6a0-6db7-11eb-3491-039358de17d7
md"Wow! Interpretting this plot is fairly easy. Black is a $1$ and white is a $0$, and time flows downwards. It's rather remarkable how starting from just a single $1$ and the rest $0$'s at $t=1$ the behavior of this cellular automata still appears quite complex! In fact, the center column of this particular rule has been shown to exhibit many of the qualities of a pseudo-random number generator.

Just for fun let's run this automaton for 500 steps just to see what happens if we let it run for a little longer."

# ╔═╡ dfe4d210-6db7-11eb-1425-732f8753c0f5
longer_history = automaton([i == 501 ? 1 : 0 for i = 1:1001], 500, rule)

# ╔═╡ f4bd6df0-6db7-11eb-15a0-c52c636db094
longer_history_mat = [longer_history[i][j] for i=1:length(longer_history), j=1:length(longer_history[1])]

# ╔═╡ f9fe56d0-6db7-11eb-1a02-6133a37e6426
heatmap(1 .- longer_history_mat, color = :greys, yflip = true, size=(640, 320))

# ╔═╡ 065d25f2-6db8-11eb-0ab4-95151bc3554f
md"It's a little tough to see the details since our image isn't high enough resolution to see each and every state value, but it is still fairly clear that the pseudo-random behavior does not settle down in the first 500 steps. However when viewing this many steps it begins to become clear that patterns, or rather regularities, begin to appear moving through space and time at the speed of causality on the left side."

# ╔═╡ 4117f030-6db8-11eb-0df5-4d606d3faf6b
md"## So what??

Beyond being an interesting thing to look at, what else are these automata useful for? The Wikipedia definition of what a cellular automata are mentioned that they are a **model of computation**. Put simply this means that fundamentally cellular automata are performing computations over time. In fact it has been shown that 1 particular rule using only nearest-neighbors and two colors in one dimension (just like our rule) is **Turing-complete**, which means it can do any computation that the computer you're viewing this on can!

I encourage you now to go back and experiment with some of the rule values on the right-hand side. Don't change the left since all 8 of the possible 3 bit combinations have to be present for the automata to be guarenteed to have a valid next state. Also, try using the rule below, which has been proven to be Turing-complete!"

# ╔═╡ caf3afb0-6db8-11eb-39de-852bb7609cb3
turing_complete_rule = Dict(
	0b000 => 0,
	0b001 => 1,
	0b010 => 1,
	0b011 => 1,
	0b100 => 0,
	0b101 => 1,
	0b110 => 1,
	0b111 => 0
)

# ╔═╡ Cell order:
# ╟─724acb50-6db3-11eb-07a8-7d048b8e4b3e
# ╟─9464f530-6db3-11eb-0629-9f05df1979bd
# ╟─f3f9f3b0-6db3-11eb-1446-cfe467e42fee
# ╟─fdc546b0-6db3-11eb-1a4c-514e32ae4a27
# ╟─87bd67d0-6db4-11eb-0931-f9a3ec94cd70
# ╠═b3fdfd00-6db4-11eb-3fa3-79d5911c567d
# ╟─f21ba14e-6db4-11eb-382a-87bc07250f65
# ╟─7895f0e0-6db6-11eb-1d6e-b12880ebd5b9
# ╠═0ab21310-6db6-11eb-00b5-1f2bd2203c79
# ╠═656f4030-6db5-11eb-0360-bd9ac3ac6b45
# ╟─232031c0-6db6-11eb-2f48-81afe98ccac3
# ╠═71d85310-6db6-11eb-0728-f34323017fe1
# ╠═aefbdff0-6db6-11eb-0441-a924d0f567c1
# ╟─ea3c54f0-6db6-11eb-1312-3daa47f3fa91
# ╟─f018e190-6db6-11eb-3eef-c984453d36d9
# ╠═a29c2ea0-6ff4-11eb-2954-372c6727bee2
# ╠═06c775f0-6db7-11eb-2603-4942353b201e
# ╟─8bb7e6a0-6db7-11eb-3491-039358de17d7
# ╠═dfe4d210-6db7-11eb-1425-732f8753c0f5
# ╠═f4bd6df0-6db7-11eb-15a0-c52c636db094
# ╠═f9fe56d0-6db7-11eb-1a02-6133a37e6426
# ╟─065d25f2-6db8-11eb-0ab4-95151bc3554f
# ╟─4117f030-6db8-11eb-0df5-4d606d3faf6b
# ╠═caf3afb0-6db8-11eb-39de-852bb7609cb3


# This shows how to read in moho topography, plot the data as points and fit a surface through the data

# Spada, M., Bianchi, I., Kissling, E., Agostinetti, N.P., Wiemer, S., 2013. *Combining controlled-source seismology and receiver function information to derive 3-D Moho topography for Italy.* Geophysical Journal International 194, 1050–1068. [doi:10.1093/gji/ggt148](https://doi.org/10.1093/gji/ggt148)
#
# We have uploaded it here: https://seafile.rlp.net/d/a50881f45aa34cdeb3c0/
#
# Please download the files `Moho_Map_Data-WesternAlps-SpadaETAL2013_Moho1.txt`, `Moho_Map_Data-WesternAlps-SpadaETAL2013_Moho2.txt` and `Moho_Map_Data-WesternAlps-SpadaETAL2013_Moho3.txt`

using DelimitedFiles, GeophysicalModelGenerator

# read European Moho
data                =   readdlm("Moho_Map_Data-WesternAlps-SpadaETAL2013_Moho1.txt",' ',Float64,'\n', skipstart=38,header=false)
lon, lat, depth     =   data[:,1], data[:,2], -data[:,3];
data_Moho1          =   GeoData(lon,lat,depth,(MohoDepth=depth*km,))

# Write as data points:
Write_Paraview(data_Moho1, "Spada_Moho_Europe", PointsData=true)

# Do the same with the other Moho's:
data =readdlm("Moho_Map_Data-WesternAlps-SpadaETAL2013_Moho2.txt",' ',Float64,'\n', skipstart=38,header=false);
lon, lat, depth        = data[:,1], data[:,2], -data[:,3];
Tutorial_MohoSpada_LonLat_Paraview = GeoData(lon,lat,depth,(MohoDepth=depth*km,))
Write_Paraview(data_Moho2, "Spada_Moho_Adria", PointsData=true)
data =readdlm("Moho_Map_Data-WesternAlps-SpadaETAL2013_Moho3.txt",' ',Float64,'\n', skipstart=38,header=false);
lon, lat, depth        = data[:,1], data[:,2], -data[:,3];
data_Moho3 = GeoData(lon,lat,depth,(MohoDepth=depth*km,))
Write_Paraview(data_Moho3, "Spada_Moho_Tyrrhenia", PointsData=true)

# Combine all data points into one data set:
lon   = [data_Moho1.lon.val;   data_Moho2.lon.val;   data_Moho3.lon.val];
lat   = [data_Moho1.lat.val;   data_Moho2.lat.val;   data_Moho3.lat.val];
depth = [data_Moho1.depth.val; data_Moho2.depth.val; data_Moho3.depth.val];
data_Moho_combined = GeoData(lon, lat, depth, (MohoDepth=depth*km,))

# Define regular grid
Lon,Lat,Depth     =   LonLatDepthGrid(4.1:0.1:11.9,42.5:.1:49,-30km);

# Interpolate using NN:
using NearestNeighbors
kdtree = KDTree([lon'; lat'])
idxs, dists = knn(kdtree, [Lon[:]'; Lat[:]'], 1, true)
Depth = zeros(size(Lon))*km;
for i=1:length(idxs)
    Depth[i] = depth[idxs[i]][1]
end
data_Moho = GeoData(Lon, Lat, Depth, (MohoDepth=Depth,))
Write_Paraview(data_Moho, "Spada_Moho_combined")


using SparseArrays, LinearAlgebra

"Check whether a package is installed."
function pkg_installed(name::String, uuid::String)
	pkg_id = Base.PkgId(Base.UUID(uuid), name::String)
	return Base.locate_package(pkg_id) != nothing
end

# generate a random pos def matrix with eigenvalues between 0.1 and 2
generate_pos_def_matrix(rng::MersenneTwister, n::Integer, aMin::Float64 = 0.1, aMax::Float64 = 2.) = generate_pos_def_matrix(rng, n, aMin, aMax, MT = Float64)
function generate_pos_def_matrix(rng::MersenneTwister, n::Integer, aMin::Real = 0.1, aMax::Real = 2.; MT::Type{<:AbstractFloat} = Float64 )
	X = rand(rng, MT, n, n)
	# any real square matrix can be QP decomposed into a orthogonal matrix and an uppertriangular matrix R
	Q, R = qr(X)
	eigs = rand(rng ,MT, n) .* (MT(aMax) .- MT(aMin)) .+ MT(aMin)
	X = Q * Matrix(Diagonal(eigs)) * Q'
	X = MT(0.5) * (X + X')
	return X
end

function is_numerically_pos_sem_def(X, atol)
	X = X ./ 2
	X = X + X'

	F = eigfact(X)
	if size(find( x-> x < -atol, F[:values]), 1) == 0
		return true
	else
		return false
	end
end

function is_numerically_symmetric(X,atol)
	n = size(X, 2)
	for i = 1:n-1, j = i+1:n
		if abs(X[i, j] - X[j, i]) >= atol
			return false
		end
	end
	return true
end

function find_nonsymmetric_component(X)
	for i = 2:size(X, 1), j = 1:(i - 1)
		if abs(X[i, j] - X[j, i]) > 0.0
			return i, j, abs(X[i, j] - X[j, i])
		end
	end
end

function find_different_elements(A, B)
	if size(A) != size(B)
		error("Matrices are not the same size")
	end
	m, n = size(A)
	diff_el = Array[]
	for iii = 1:m, jjj = 1:n
		if A[iii, jjj] != B[iii, jjj]
			push!(diff_el, [iii, jjj])
		end
	end
	return diff_el
end

function duplicate_sparsity_pattern(A)
	m, n = size(A)
	B = zeros(m, n)
	for iii = 1:m, jjj = 1:n
		if A[iii, jjj] != 0
			B[iii, jjj] = 1
		end
	end
	return B
end

function apply_pattern!(A::AbstractMatrix, pattern::AbstractMatrix)
  @assert size(A) == size(pattern) "Matrix A and pattern must have same dimensions."
  m, n = size(A)
  for i = 1:m, j = 1:n
    if pattern[i, j] == 0
      A[i, j] = 0.
    end
  end
end

# create a feasible SDP with one PSDConeTriangle constraint
# min c' x
# s.t. At x + s == bt
#       s ∈ PSDConeTriangle
# choose At and bt in such a way that S has the provided sparsity pattern
function feasible_sdp_with_pattern(rng::MersenneTwister, pattern::AbstractMatrix)
  n = size(pattern, 1)
  d = div(n * (n + 1), 2)

  S = generate_pos_def_matrix(rng, n, 0.1, 2.)
  apply_pattern!(S, pattern)
  S = Symmetric(S, :U)

  A1 = rand(rng, n, n)
  apply_pattern!(A1, pattern)
  A1 = Symmetric(A1, :U)
  A = hcat(A1[:])
  s = S[:]
  x = rand(rng, 1)
  b = A * x + s
  B = reshape(b, n, n)

  Y = generate_pos_def_matrix(rng, n,  0.1, 1.)
  y = vec(Y)
  P = sparse(zeros(1, 1))
  q = -P * x - A' * y

  Ct = [COSMO.PsdConeTriangle(d)];
  At = zeros(d)
  bt = zeros(d)
  COSMO.extract_upper_triangle!(A1, At, sqrt(2))
  COSMO.extract_upper_triangle!(B, bt, sqrt(2))
  At = hcat(At)
  return P, q, At, bt, Ct
end

"Take a vector `svec` representing the `d` upper-triangular entries of a matrix `X`, and return `X`."
function matrixify(svec::AbstractVector)
  n = - 0.5 + sqrt(0.25 + 2 * length(svec))
  n = Int64(n)
  X = zeros(n, n)
  COSMO.populate_upper_triangle!(X, svec, 1 /sqrt(2))
  return Symmetric(X, :U)
end

function recreate_sparse_matrix(A::SparseMatrixCSC)
	rowInd = A.rowval
	colPtr = A.colptr
	val = A.nzval

	#compute column indices
	colInd = zeros(Int64, length(rowInd))
	cval = 1
	for iii = 2:length(colPtr)
		currentPtr = colPtr[iii]
		prevPtr = colPtr[iii - 1]
		colInd[prevPtr:currentPtr - 1] = cval
		cval += 1
	end

	# sort rowInd and vals
	p = sortperm(rowInd)
	return sparse(rowInd, colInd, val, size(A, 1), size(A, 2))
end

# Geometric mean from https://github.com/JuliaStats/StatsBase.jl
function gmean(a::AbstractArray{T}) where T<:Real
	s = 0.0
	n = length(a)
	for i in 1:n
		tmp = a[i]
		if tmp < 0.0
			throw(DomainError())
		elseif tmp == 0.0
			return 0.0
		else
			s += log(tmp)
		end
	end
	return exp(s / n)
end

function check_clique_tree(snd, sep, snd_par, snd_child)
	# c.t. should have one root
	ind = findall( x -> x == 0, snd_par)
	length(ind) == 1 || return false

	for (ch, p) in enumerate(snd_par)
		# removed
		if p == -1
			isempty(snd[ch]) || return false
			isempty(sep[ch]) || return false
		# root
		elseif p == 0
			isempty(sep[ch]) || return false
			# find all children of ch
			children = findall(x -> x == ch, snd_par)
			mapreduce(x -> x ∈ snd_child[ch], *,  children) || return false
		else
			intersect(union(snd[ch], sep[ch]), union(snd[p], sep[p])) == sep[ch] || return false
			children = findall(x -> x == ch, snd_par)
			if !isempty(children)
				mapreduce(x -> x ∈ snd_child[ch], *,  children) || return false
			end
		end
	end
	return true
end


module mrotationmatrixs2
using FinEtools
using LinearAlgebra: norm, I
# using BenchmarkTools
using Test
function rotmat3original!(Rmout, a)
    na = norm(a);
    thetatilde = zeros(3,3);
    skewmat!(thetatilde, a);
    a = a/na;
    ca = cos(na);
    sa = sin(na);
    aa = a*a';
    copyto!(Rmout, ca * (1.0*I-aa) + sa/na*thetatilde + aa);
    return Rmout
end
function test()
    for i in 1:10
        Rmout = rand(3, 3)
        a = rand(3)
        # Rmout = [0.056575544213388396 0.8235145288079604 0.4566335160952417; 0.6335215149040396 0.8852762577618685 0.803417762628938; 0.5940995203626245 0.35601280936101065 0.6825575362008556]
        # a = [0.11011521831193871, 0.5097478695998647, 0.8139760429477749]
        rotmat3!(Rmout, a)
        Rmout2 = rand(3, 3)
        rotmat3original!(Rmout2, a)
        # Rmout = [0.5736166081143759 -0.6670428981551146 0.47541325067375245; 0.7189364978350996 0.6881251218379588 0.0980516638109532; -0.3925484670406451
        # 0.2855478746485778 0.8742814834523946]
        @test norm(Rmout-Rmout2) <= 1.0e-6
        # @btime rotmat3!($Rmout, $a)
    end
    true
end
end
using .mrotationmatrixs2
mrotationmatrixs2.test()


using NestedMaps, Test

@testset "fallback" begin
    x = 1:10
    @test nested_map(identity, x) == x
    @test nested_map(sum, x) == sum(x)
end

@testset "tuple" begin
    t = [(i, -i) for i in 1:10]
    @test nested_map(identity, t) == (1:10, -(1:10))
    @test nested_map(sum, t) == (55, -55)
    @test_throws ArgumentError nested_map(identity, [(1, ), (2, 3)])
end

@testset "named tuple" begin
    nt = [(a = i, b = -i) for i in 1:10]
    @test nested_map(identity, nt) == (a = 1:10, b = -(1:10))
    @test nested_map(sum, nt) == (a = 55, b = -55)
    @test_throws ArgumentError nested_map(identity, [(a = 1, ), (b = 2, )])
end

@testset "array" begin
    r1 = 1:5
    r2 = 2:7
    A = [[i, j] for i in r1, j in r2]
    @test nested_map(identity, A) == [[i for i in r1, _ in r2],
                                      [j for _ in r1, j in r2]]
    @test nested_map(sum, A) == [sum(r1) * length(r2), sum(r2) * length(r1)]
    @test_throws ArgumentError nested_map(identity, [[1], [1, 2]])
end

@testset "recursive" begin
    r = 1:5
    A = [(a = i, b = (c = -i, d = [i^2, -i^2])) for i in r]
    @test nested_map_recursive(identity, A) == (a = r, b = (c = .-r, d = [r.^2, -r.^2]))
    s1, s2 = sum(r), sum(abs2, r)
    @test nested_map_recursive(sum, A) == (a = s1, b = (c = -s1, d = [s2, -s2]))
end


# Testing part 1
using Test, ReferenceTests

# dir to the part 1 script
include("../scripts-part1/Diffusion_3D.jl")

# Unit tests
Hτ = rand(5, 5, 5);
D_dx, D_dy, D_dz = 1.0, 1.0, 1.0;
@testset "Test heat flux macros" begin
    @test @qx(2,1,3) ≈ -(Hτ[3,2,4] - Hτ[2,2,4])
    @test @qx(4,3,1) ≈ -(Hτ[5,4,2] - Hτ[4,4,2])
    @test @qy(1,1,1) ≈ -(Hτ[2,2,2] - Hτ[2,1,2])
    @test @qy(3,2,3) ≈ -(Hτ[4,3,4] - Hτ[4,2,4])
    @test @qz(1,4,2) ≈ -(Hτ[2,5,3] - Hτ[2,5,2])
    @test @qz(2,3,4) ≈ -(Hτ[3,4,5] - Hτ[3,4,4])
end

# Reference test using ReferenceTests.jl
"Compare all dict entries"
comp(d1, d2) = keys(d1) == keys(d2) && all([ isapprox(v1, v2; atol = 1e-5) for (v1,v2) in zip(values(d1), values(d2))])
inds = Int.(ceil.(LinRange(1, length(xc_g), 12)))
d = Dict(:X=> xc_g[inds], :H=>H_g[inds, inds, 15])

@testset "Reference Test" begin
    @test_reference "reftest-files/test_part_1.bson" d by=comp
end


using Distances
using Random

# export generate_random_point_cloud,
#         generate_geometric_matrix,
#         generate_shuffled_matrix,
#         generate_random_matrix,
#         generate_matrix_ordering,
#         # generate_set_of_graphs,
#         # plot_betti_numbers,
#         # save_matrix_to_file;

"""
Returns a random matrix of size @number_of_points x @dimensions in which every
column is a point and every n-th row is a position in the n-th dimension.
"""
generate_random_point_cloud(number_of_points = 12, dimensions=2) =
                                    rand(Float64, dimensions, number_of_points)

"""
Return a matrix which stores the pariwise distances between every point in the
@random_points matrix.
"""
function generate_geometric_matrix(random_points)

    geometric_matrix = Distances.pairwise(Euclidean(), random_points, dims=2)
    return geometric_matrix
end

"""
Returns a ordered geometric matrix, which was generated by samping 'dims'
dimensional unit cuve with 'total_points' samples.
"""
function get_ordered_geom_matrix(dims::Integer, total_points::Integer)
    # TODO to be completed

end

"""
Returns a symetric matrix with randomly permuted valuse from the @input_matrix.
"""
function generate_shuffled_matrix(input_matrix)
    matrix_size = size(input_matrix,1)

    indicies_collection = findall(x->x>0, input_matrix)
    rand!(indicies_collection, indicies_collection)
    shuffeled_matrix = copy(input_matrix)

    # Swap the elements
    n=1
    for k in 1:matrix_size
        for m in k+1:matrix_size
            a = indicies_collection[n][1]
            b = indicies_collection[n][2]
            shuffeled_matrix[k,m] = input_matrix[a,b]
            shuffeled_matrix[m,k] = input_matrix[b,a]

            shuffeled_matrix[a,b] = input_matrix[k,m]
            shuffeled_matrix[b,a] = input_matrix[m,k]

            n +=1
        end
    end
    return shuffeled_matrix
end

"""
Returns matrix with random values which are symmetric accros the diagonal. The
matrix has @matrix_size rows and @matrix_size columns.
"""
function generate_random_matrix(matrix_size)
    elemnts_above_diagonal = Int((matrix_size^2-matrix_size)/2)
    random_matrix = zeros(matrix_size, matrix_size)
    set_of_random_numbers = rand(elemnts_above_diagonal)

    h = 1
    for k in 1:matrix_size
        for m in k+1:matrix_size
            random_matrix[k,m] = set_of_random_numbers[h]
            random_matrix[m,k] = set_of_random_numbers[h]
            h += 1
        end
    end

    return random_matrix
end

# function generate_set_of_graphs(matrix_size, matrix_ordering)
#     """
#     Returns set of graphs generated from the @matrix_ordering. In every succesive
#     graph, single connection between points is added.
#
#     NOTE: the function does not take the coordinates of the numbered vertices.
#     """
#     vetrices = matrix_size
#     edges = matrix_ordering
#     num_of_edges = size(edges)[2]
#
#     set_of_graphs = [a=Graph(vetrices) for a=1:num_of_edges]
#     edges_counter = zeros(Int, num_of_edges)
#     edge_density =  zeros(num_of_edges)
#
#     k=1
#     for k in range(1,stop=num_of_edges)~
#         add_edge!(set_of_graphs[k], edges[1,k], edges[2,k]);
#         edges_counter[k] = ne(set_of_graphs[k])
#         edge_density[k] = edges_counter[k]/binomial(matrix_size,2)
#
#         if k<num_of_edges # if is used to eliminate copying at last iteration
#             set_of_graphs[k+1] = copy(set_of_graphs[k])
#         end
#     end
#     return set_of_graphs, edge_density
# end

# function save_matrix_to_file(matrix, filename)
#     """
#     Saves given @matrix to the csv file with the name @filename. If there is no path
#     added to the @filename, then file saved is in local folder.
#     """
#     open(filename, "w") do io
#         writedlm(io,  matrix, ',')
#     end
# end

# =====
# Copied form Julia learning repo

"""
Returns ordering of the @geometric_matrix given as an input. If value @ascending
is set to true, the values are number from the lowest value, to the highest. If
false, the values are numbered from highest to the lowest.
"""
function generate_matrix_ordering(geometric_matrix, ascending = true)

    matrix_size = size(geometric_matrix, 2)
    elemnts_above_diagonal = Int((matrix_size^2-matrix_size)/2)
    matrix_ordering = zeros(Int, 2,elemnts_above_diagonal)

    A = copy(geometric_matrix)

    (ascending) ? (method=findmax) : (method=findmin)

    for element in 1:elemnts_above_diagonal
    #     Find maximal distance
        minimal_value = method(A)
    #     Get the coordinates (only 2 dimensions, because it is distance matrix)
        matrix_ordering[1,element] = Int(minimal_value[2][1])
        matrix_ordering[2,element] = Int(minimal_value[2][2])
    #
    # #     Zero minval in A (above and below diagonal) so next minval can be found
        A[matrix_ordering[1,element], matrix_ordering[2,element]] = 0.0
        A[matrix_ordering[2,element], matrix_ordering[1,element]] = 0.0
    end

    # change from min to max order to the max to min order (? necessary ?)
    if ascending
        matrix_ordering = matrix_ordering[:,end:-1:1]
    end

    return matrix_ordering
end

"""
    function get_geometric_matrix(points, dimensions; save_as_file=false)

Created a point cloud with 'points' number of points from 'dimension'
dimensional eucidean unit cube and computes distances between the points.

Distance matrix may be saved to csv file by setting 'save_as_file' to 'true'.
"""
function get_geometric_matrix(points, dimensions; save_as_file=false)
    point_cloud = generate_random_point_cloud(points,dimensions)
    geom_mat = generate_geometric_matrix(point_cloud)

    if save_as_file
        open("geometric_matrix_points$(points)_dims$(dimensions).csv", "w") do io
           writedlm(io, geom_mat, ',')
       end
   end
   return geom_mat
end


#=
VecE3: a 3d euclidean vector
=#

struct VecE3{R<:Real} <: VecE{3,R}
    x::R
    y::R
    z::R
end
VecE3() = VecE3(0.0,0.0,0.0)
VecE3(x::Integer, y::Integer, z::Integer) = VecE3(float(x), float(y), float(z))
VecE3(t::Tuple) = VecE3(promote(t...)...)

Base.show(io::IO, a::VecE3) = @printf(io, "VecE3(%.3f, %.3f, %.3f)", a.x, a.y, a.z)

function dist(a::VecE3, b::VecE3)
    Δx = a.x-b.x
    Δy = a.y-b.y
    Δz = a.z-b.z
    sqrt(Δx*Δx + Δy*Δy + Δz*Δz)
end
function dist2(a::VecE3, b::VecE3)
    Δx = a.x-b.x
    Δy = a.y-b.y
    Δz = a.z-b.z
    Δx*Δx + Δy*Δy + Δz*Δz
end

proj(a::VecE3, b::VecE3, ::Type{Float64}) = (a.x*b.x + a.y*b.y + a.z*b.z) / sqrt(b.x*b.x + b.y*b.y + b.z*b.z) # dot(a,b) / |b|
function proj(a::VecE3, b::VecE3, ::Type{VecE3})
    # dot(a,b) / dot(b,b) ⋅ b
    s = (a.x*b.x + a.y*b.y + a.z*b.z) / (b.x*b.x + b.y*b.y + b.z*b.z)
    VecE3(s*b.x, s*b.y, s*b.z)
end

lerp(a::VecE3, b::VecE3, t::Real) = VecE3(a.x + (b.x-a.x)*t, a.y + (b.y-a.y)*t, a.z + (b.z-a.z)*t)

function rot(a::VecE3, axis::VecE3, θ::Real)

    #=
    rotate the point a: (x,y,z) about the axis: (u,v,w) by the angle θ following the
    right hand rule
    =#

    x,y,z = a.x, a.y, a.z
    u,v,w = axis.x, axis.y, axis.z

    u² = u*u
    v² = v*v
    w² = w*w

    m = u² + v² + w²

    c = cos(θ)
    s = sin(θ)
    mc = 1.0 - c

    rtms = sqrt(m)*s

    magc = m*c
    ux_vy_wz_mc = (u*x + v*y + w*z)*mc

    new_x = (u*ux_vy_wz_mc + x*magc + rtms*(-w*y + v*z)) / m
    new_y = (v*ux_vy_wz_mc + y*magc + rtms*( w*x - u*z)) / m
    new_z = (w*ux_vy_wz_mc + z*magc + rtms*(-v*x + u*y)) / m

    VecE3(new_x, new_y, new_z)
end
function rot_normalized(a::VecE3, axis::VecE3, θ::Real)

    #=
    rotate the point a: (x,y,z) about the axis: (u,v,w) by the angle θ following the
    right hand rule

    Here we assume axis is normalized
    =#

    x,y,z = a.x, a.y, a.z
    u,v,w = axis.x, axis.y, axis.z

    u² = u*u
    v² = v*v
    w² = w*w

    m = u² + v² + w²

    c = cos(θ)
    s = sin(θ)
    mc = 1.0 - c

    ux_vy_wz_mc = (u*x + v*y + w*z)*mc

    new_x = u*ux_vy_wz_mc + x*c + (-w*y + v*z)*s
    new_y = v*ux_vy_wz_mc + y*c + ( w*x - u*z)*s
    new_z = w*ux_vy_wz_mc + z*c + (-v*x + u*y)*s

    VecE3(new_x, new_y, new_z)
end


import Base: ==

# https://github.com/JuliaLang/julia/issues/11618
signed_prod(lst) = length(lst) == 1 ? lst[1] : Base.prod(lst)
signed_sum(lst) = length(lst) == 1 ? lst[1] : Base.sum(lst)

abstract Node
==(n1::Node, n2::Node) = isequal(n1.val, n2.val)
calc(n::Float64) = n
type Inv<:Node val end
calc(i::Inv) = 1.0/calc(i.val)
type Prd<:Node val end
calc(p::Prd) = signed_prod(map(calc, p.val))
type Neg<:Node val end
calc(n::Neg) = -calc(n.val)
type Sum<:Node val end
calc(s::Sum) = signed_sum(map(calc, s.val))

@with_names begin

    spc = Drop(Star(Space()))

    @with_pre spc begin

        sum = Delayed()
        val = E"(" + spc + sum + spc + E")" | PFloat64()

        neg = Delayed()             # allow multiple negations (eg ---3)
        neg.matcher = Nullable{Matcher}(val | (E"-" + neg > Neg))

        mul = E"*" + neg
        div = E"/" + neg > Inv
        prd = neg + (mul | div)[0:end] |> Prd

        add = E"+" + prd
        sub = E"-" + prd > Neg
        sum.matcher = Nullable{Matcher}(prd + (add | sub)[0:end] |> Sum)

        all = sum + spc + Eos()

    end
end

parse_one(" 1 + 2 * 3 / 4 ", Trace(all); debug=true)

for (src, val) in [
                   (" 1 ", 1),
                   (" - 1 ", -1),
                   (" 1 + 1 ", 2),
                   (" 1 - 1 ", 0),
                   (" - 1 - 1 ", -2)
                   ]
#    @test_approx_eq calc(parse_dbg(src, Trace(all))[1]) val
    @test_approx_eq calc(parse_one(src, Trace(all))[1]) val
    println("$src = $val")
end

println("fix ok")


# TODO allow the lower and upper bounds to be changed, like uniform.

struct BetaUniformMixture <: Distribution{Float64} end

"""
    beta_uniform(theta::Real, alpha::Real, beta::Real)

Samples a `Float64` value from a mixture of a uniform distribution on [0, 1] with probability `1-theta` and a beta distribution with parameters `alpha` and `beta` with probability `theta`.
"""
const beta_uniform = BetaUniformMixture()

function logpdf(::BetaUniformMixture, x::Real, theta::Real, alpha::Real, beta::Real)
    lbeta = log(theta) + logpdf(Beta(), x, alpha, beta)
    luniform = log(1.0 - theta)
    logsumexp(lbeta, luniform)
end

function logpdf_grad(::BetaUniformMixture, x::Real, theta::Real, alpha::Real, beta::Real)
    beta_logpdf = logpdf(Beta(), x, alpha, beta)
    uniform_logpdf = logpdf(uniform_continuous, x, 0., 1.)
    beta_grad = logpdf_grad(Beta(), x, alpha, beta)
    uniform_grad = logpdf_grad(uniform_continuous, x, 0., 1.)
    w1 = 1. / (1. + exp(log(1. - theta) + uniform_logpdf - log(theta) - beta_logpdf))
    w2 = 1. - w1
    x_deriv = w1 * beta_grad[1] + w2 * uniform_grad[1]
    alpha_deriv = w1 * beta_grad[2]
    beta_deriv = w1 * beta_grad[3]
    theta_deriv = (exp(beta_logpdf) - exp(uniform_logpdf)) / (theta * exp(beta_logpdf) + (1. - theta) * exp(uniform_logpdf))
    (x_deriv, theta_deriv, alpha_deriv, beta_deriv)
end

function random(::BetaUniformMixture, theta::Real, alpha::Real, beta::Real)
    if bernoulli(theta)
        random(Beta(), alpha, beta)
    else
        random(uniform_continuous, 0., 1.)
    end
end

(::BetaUniformMixture)(theta, alpha, beta) = random(BetaUniformMixture(), theta, alpha, beta)

has_output_grad(::BetaUniformMixture) = true
has_argument_grads(::BetaUniformMixture) = (true, true, true)

export beta_uniform


using MatrixAlgebra
using LinearAlgebra
using Test

let A = [1.0 2 4 7 11; 2 3 5 8 12; 4 5 6 9 13; 7 8 9 10 14; 11 12 13 14 15], n = size(A, 1)
    @testset "argument checks" begin
        @test_throws ArgumentError tridiag_real(A, n+1)
        @test_throws ArgumentError tridiag_real(A, 0)
        @test_throws DimensionMismatch tridiag_real(A[:,2:n], 1)
    end
    @testset "real $(2d-1)-diagonal" for d in 2:n
        M = [abs(i-j) < d ? A[i,j] : zero(eltype(A)) for i in 1:n, j=1:n]
        E = tridiag_real(M, d)
        @test eigvals(E) ≈ eigvals(M)
    end
end
let A = [1.0 2 4+im 7 11; 2 3 5 8+im 12; 4-im 5 6 9 13; 7 8-im 9 10 14; 11 12 13 14 15], n = size(A, 1)
    @testset "complex $(2d-1)-diagonal" for d in 2:n
        M = [abs(i-j) < d ? A[i,j] : zero(eltype(A)) for i in 1:n, j=1:n]
        E = tridiag_real(M, d)
        @test eigvals(E) ≈ eigvals(M)
    end
end


##INVESTIGATE DISTRIBUTION OF RE-ABC LIKELIHOOD ESTIMATES
using RareEventABC
using ProgressMeter
using DataFrames
using Distributions
using Plots; pgfplots()
using StatPlots

##GENERATE OBSERVATIONS
srand(1);
d = 25;
yobs = randn(d) .* 3.0;
a = RareEventABC.IIDNormalDetails(yobs, d);

##GENERATE LIKELIHOOD ESTIMATES
schedule = doSMC(a, [3.0], 5.0, nparticles=200, quiet=true).schedule
llexperiment = Dict();
@showprogress for nparticles in 10:10:100
    llest = Float64[doSMC(a, [3.0], schedule, nparticles=nparticles, d=25, quiet=true).llest for i in 1:100]
    llexperiment[nparticles] = llest
end

##QQ PLOTS
theoretical_quantiles = quantile(Normal(), (1:100)/101)
qqdf = DataFrame(samp=Float64[], theoretical=Float64[], nparticles=Int[])
for nparticles in 20:10:50
    ll = llexperiment[nparticles]
    temp = DataFrame(samp=sort(ll), theoretical=theoretical_quantiles, particles=nparticles)
    qqdf = vcat(qqdf, temp)
end

qqdf2 = qqdf[qqdf[:samp] .> -Inf, :] ##Plots.jl doesn't seem to work with infinite values
default(guidefont=font(25), tickfont=font(25), legendfont=font(25), size=(500,500));
p1 = scatter(qqdf2, :theoretical, :samp, group=:particles,
             markershape=:auto,
             xlims=(-2.5,3.0), ##Hack so that legend isn't in the way
             xlabel="Sample quantile", ylabel="Theoretical quantile");
savefig(p1, "iidnormal_qq.pdf")

##COMPARE ADAPTIVE AND NON-ADAPTIVE ALGORITHMS
b = RareEventABC.IIDNormalDetails([0.0], 1);
eps = 0.2;
##trueL = (cdf(n, eps) - cdf(n,-eps)) / (2*eps)

nparticles = 10;

srand(1);
nreps = 10000;
likelihood_adaptive = Float64[doSMC(b, [1.0], eps, nparticles=nparticles, d=1, quiet=true).llest for i in 1:nreps] |> exp;

srand(2);
schedule = doSMC(b, [1.0], eps, nparticles=nparticles, quiet=true).schedule
likelihood_nonadaptive = Float64[doSMC(b, [1.0], schedule, nparticles=nparticles, d=1, quiet=true).llest for i in 1:nreps] |> exp;

mean(likelihood_adaptive), std(likelihood_adaptive)/sqrt(nreps)
mean(likelihood_nonadaptive), std(likelihood_nonadaptive)/sqrt(nreps)


# # Knapsack Problem - Single Expansion
# This first example demonstrates how a basic JuDGE model can be set up, with
# a 7-node tree, and a single expansion.
using JuMP, JuDGE, GLPK

# First we will define our tree with a depth of 2 and a degree of 2.
mytree = narytree(2,2)

# ## Data

# Now we specify data for each node of the tree. Here we define the data using
# functions, but you could use dictionaries, as shown in subsequent examples.
function invest_cost(node)
   if node == get_node(mytree,[1])
	  180.0
   elseif node == get_node(mytree,[1,1])
	  50.0
   elseif node == get_node(mytree,[1,2])
	  60.0
   elseif node == get_node(mytree,[1,1,1])
	  40.0
   elseif node == get_node(mytree,[1,1,2])
	  60.0
   elseif node == get_node(mytree,[1,2,1])
	  10.0
   elseif node == get_node(mytree,[1,2,2])
	  10.0
   end
end

function item_volume(node)
   if node == get_node(mytree,[1])
	  [6, 2, 1, 1, 1]
   elseif node == get_node(mytree,[1,1])
	  [8, 2, 2, 2, 1]
   elseif node == get_node(mytree,[1,2])
	  [8, 1, 1, 1, 3]
   elseif node == get_node(mytree,[1,1,1])
	  [4, 4, 3, 1, 2]
   elseif node == get_node(mytree,[1,1,2])
	  [1, 3, 1, 1, 2]
   elseif node == get_node(mytree,[1,2,1])
	  [7, 3, 1, 1, 1]
   elseif node == get_node(mytree,[1,2,2])
	  [2, 5, 2, 1, 2]
   end
end

function item_reward(node)
   if node == get_node(mytree,[1])
	  [60, 20, 10, 15, 10]
   elseif node == get_node(mytree,[1,1])
	  [8, 10, 20, 20, 10]
   elseif node == get_node(mytree,[1,2])
	  [8, 10, 15, 10, 30]
   elseif node == get_node(mytree,[1,1,1])
	  [40, 40, 35, 10, 20]
   elseif node == get_node(mytree,[1,1,2])
	  [15, 35, 15, 15, 20]
   elseif node == get_node(mytree,[1,2,1])
	  [70, 30, 15, 15, 10]
   elseif node == get_node(mytree,[1,2,2])
	  [25, 50, 25, 15, 20]
   end
end

# ## Subproblem definitions

# The JuDGE subproblems are defined through a function that takes a node as its
# single argument, returning a JuMP model. We use GLPK as our optimizer.
JuDGE_SP_Solver = optimizer_with_attributes(GLPK.Optimizer, "msg_lev" => 0, "mip_gap" => 0.0)
function sub_problems(node)
   model = Model(JuDGE_SP_Solver)
   @expansion(model, bag)
   @expansioncosts(model, bag*invest_cost(node))
   @variable(model, y[1:5], Bin)
   @constraint(model, BagExtension, sum(y[i]*item_volume(node)[i] for i in 1:5) <= 3 + 4 * bag)
   @objective(model, Min, sum(-item_reward(node)[i] * y[i] for i in 1:5))
   model
end

# Note that there are a few JuDGE-specific macros used to define a JuDGE subproblem

# This defines our single expansion variable
@expansion(model, bag)

# This defines the cost of buying this bag at particular node
@expansioncosts(model, bag*invest_cost(node))

# JuDGE subproblem objective is defined using the standard @objective macro:
@objective(model, Min, sum(-item_reward(node)[i] * y[i] for i in 1:5))

# ## Defining and solving the JuDGE model
# The JuDGEModel is defined based on a tree, probability distribution,
# sub_problems, and an optimizer using an interior point method.
JuDGE_MP_Solver = optimizer_with_attributes((method=GLPK.INTERIOR) -> GLPK.Optimizer(),
							"msg_lev" => 0, "mip_gap" => 0.0)
judy = JuDGEModel(mytree, ConditionallyUniformProbabilities, sub_problems, JuDGE_MP_Solver)

# Solve the model
JuDGE.solve(judy)

# ## Displaying the output
# Print the objective, and optimal expansions
println("Objective: "*string(judy.bounds.UB))
JuDGE.print_expansions(judy,onlynonzero=false)

# Re-solve the subproblems and print the objective
println("Re-solved Objective: " * string(resolve_subproblems(judy)))

# ## Deterministic equivalent
# Here we set up and solve the deterministic equivalent, this uses the same structure as a `JuDGEModel`.
# We set up our GLPK solver to show its progress.
JuDGE_DE_Solver = optimizer_with_attributes(GLPK.Optimizer, "msg_lev" => 2, "mip_gap" => 0.0)
deteq = DetEqModel(mytree, ConditionallyUniformProbabilities, sub_problems, JuDGE_DE_Solver)
JuDGE.solve(deteq)
println("Deterministic Equivalent Objective: " * string(objective_value(deteq.problem)))


using Test, Astro

# equinox tests from meeus chapter

print("27.a Approximate solstice")
eqa = equinox_approx(1962, "summer")
y, m , d = jd_to_cal(eqa)

@test y == 1962
@test m == 6
@test isapprox(d, 21, atol=1)

println(" passed")

print("27.a Finding nearest Equinox for 12 January 2007")
jd = cal_to_jd(2007,1,12)
eq = equinox(jd, "Spring", 1) # return jd of equinox to the nearest day

y, m , d = jd_to_cal(eq)
@test y == 2007
@test m == 3
@test isapprox(d, 21, atol=1)

println(" passed")


import Base: iterate
using Printf
export gmres, gmres!

struct ArnoldiDecomp{T, matT}
    A::matT
    V::Matrix{T} # Orthonormal basis vectors
    H::Matrix{T} # Hessenberg matrix
end

ArnoldiDecomp(A::matT, order::Int, T::Type) where {matT} = ArnoldiDecomp{T, matT}(
    A,
    zeros(T, size(A, 1), order + 1),
    zeros(T, order + 1, order)
)

mutable struct Residual{numT, resT}
    current::resT # Current, absolute, preconditioned residual
    accumulator::resT # Used to compute the residual on the go
    nullvec::Vector{numT} # Vector in the null space of H to compute residuals
    β::resT # the initial residual
end

Residual(order, T::Type) = Residual{T, real(T)}(
    one(real(T)),
    one(real(T)),
    ones(T, order + 1),
    one(real(T))
)

mutable struct GMRESIterable{preclT, precrT, solT, rhsT, vecT, arnoldiT <: ArnoldiDecomp, residualT <: Residual, resT <: Real, orthmethT}
    Pl::preclT
    Pr::precrT
    x::solT
    b::rhsT
    Ax::vecT # Some room to work in.

    arnoldi::arnoldiT
    residual::residualT

    mv_products::Int
    restart::Int
    k::Int
    maxiter::Int
    tol::resT
    β::resT

    orth_meth::orthmethT
end

converged(g::GMRESIterable) = g.residual.current ≤ g.tol

start(::GMRESIterable) = 0

done(g::GMRESIterable, iteration::Int) = iteration ≥ g.maxiter || converged(g)

function iterate(g::GMRESIterable, iteration::Int=start(g))
    # Check for termination first
    if done(g, iteration)
        return nothing
    end

    # Arnoldi step: expand
    expand!(g.arnoldi, g.Pl, g.Pr, g.k, g.Ax)
    g.mv_products += 1

    # Orthogonalize V[:, k + 1] w.r.t. V[:, 1 : k]
    g.arnoldi.H[g.k + 1, g.k] = orthogonalize_and_normalize!(
        view(g.arnoldi.V, :, 1 : g.k),
        view(g.arnoldi.V, :, g.k + 1),
        view(g.arnoldi.H, 1 : g.k, g.k),
        g.orth_meth
    )

    # Implicitly computes the residual
    update_residual!(g.residual, g.arnoldi, g.k)

    g.k += 1

    # Computation of x only at the end of the iterations
    # and at restart.
    if g.k == g.restart + 1 || done(g, iteration + 1)

        # Solve the projected problem Hy = β * e1 in the least-squares sense
        rhs = solve_least_squares!(g.arnoldi, g.β, g.k)

        # And improve the solution x ← x + Pr \ (V * y)
        update_solution!(g.x, view(rhs, 1 : g.k - 1), g.arnoldi, g.Pr, g.k, g.Ax)

        g.k = 1

        # Restart when not done.
        if !done(g, iteration)

            # Set the first basis vector
            g.β = init!(g.arnoldi, g.x, g.b, g.Pl, g.Ax)

            # And initialize the residual
            init_residual!(g.residual, g.β)

            g.mv_products += 1
        end
    end

    g.residual.current, iteration + 1
end

function gmres_iterable!(x, A, b;
                         Pl = Identity(),
                         Pr = Identity(),
                         abstol::Real = zero(real(eltype(b))),
                         reltol::Real = sqrt(eps(real(eltype(b)))),
                         restart::Int = min(20, size(A, 2)),
                         maxiter::Int = size(A, 2),
                         initially_zero::Bool = false,
                         orth_meth::OrthogonalizationMethod = ModifiedGramSchmidt())
    T = eltype(x)

    # Approximate solution
    arnoldi = ArnoldiDecomp(A, restart, T)
    residual = Residual(restart, T)
    mv_products = initially_zero ? 1 : 0

    # Workspace vector to reduce the # allocs.
    Ax = similar(x)
    residual.current = init!(arnoldi, x, b, Pl, Ax, initially_zero = initially_zero)
    init_residual!(residual, residual.current)

    tolerance = max(reltol * residual.current, abstol)

    GMRESIterable(Pl, Pr, x, b, Ax,
        arnoldi, residual,
        mv_products, restart, 1, maxiter, tolerance, residual.current,
        orth_meth
    )
end

"""
    gmres(A, b; kwargs...) -> x, [history]

Same as [`gmres!`](@ref), but allocates a solution vector `x` initialized with zeros.
"""
gmres(A, b; kwargs...) = gmres!(zerox(A, b), A, b; initially_zero = true, kwargs...)

"""
    gmres!(x, A, b; kwargs...) -> x, [history]

Solves the problem ``Ax = b`` with restarted GMRES.

# Arguments

- `x`: Initial guess, will be updated in-place;
- `A`: linear operator;
- `b`: right-hand side.

## Keywords

- `initially_zero::Bool`: If `true` assumes that `iszero(x)` so that one
  matrix-vector product can be saved when computing the initial
  residual vector;
- `abstol::Real = zero(real(eltype(b)))`,
  `reltol::Real = sqrt(eps(real(eltype(b))))`: absolute and relative
  tolerance for the stopping condition
  `|r_k| ≤ max(reltol * |r_0|, abstol)`, where `r_k = A * x_k - b`
- `restart::Int = min(20, size(A, 2))`: restarts GMRES after specified number of iterations;
- `maxiter::Int = size(A, 2)`: maximum number of inner iterations of GMRES;
- `Pl`: left preconditioner;
- `Pr`: right preconditioner;
- `log::Bool`: keep track of the residual norm in each iteration;
- `verbose::Bool`: print convergence information during the iterations.
- `orth_meth::OrthogonalizationMethod = ModifiedGramSchmidt()`: orthogonalization method (ModifiedGramSchmidt(), ClassicalGramSchmidt(), DGKS())

# Return values

**if `log` is `false`**

- `x`: approximate solution.

**if `log` is `true`**

- `x`: approximate solution;
- `history`: convergence history.
"""
function gmres!(x, A, b;
                Pl = Identity(),
                Pr = Identity(),
                abstol::Real = zero(real(eltype(b))),
                reltol::Real = sqrt(eps(real(eltype(b)))),
                restart::Int = min(20, size(A, 2)),
                maxiter::Int = size(A, 2),
                log::Bool = false,
                initially_zero::Bool = false,
                verbose::Bool = false,
                orth_meth::OrthogonalizationMethod = ModifiedGramSchmidt())
    history = ConvergenceHistory(partial = !log, restart = restart)
    history[:abstol] = abstol
    history[:reltol] = reltol
    log && reserve!(history, :resnorm, maxiter)

    iterable = gmres_iterable!(x, A, b; Pl = Pl, Pr = Pr,
                               abstol = abstol, reltol = reltol, maxiter = maxiter,
                               restart = restart, initially_zero = initially_zero,
                               orth_meth = orth_meth)

    verbose && @printf("=== gmres ===\n%4s\t%4s\t%7s\n","rest","iter","resnorm")

    for (iteration, residual) = enumerate(iterable)
        if log
            nextiter!(history)
            history.mvps = iterable.mv_products
            push!(history, :resnorm, residual)
        end

        verbose && @printf("%3d\t%3d\t%1.2e\n", 1 + div(iteration - 1, restart), 1 + mod(iteration - 1, restart), residual)
    end

    verbose && println()
    setconv(history, converged(iterable))
    log && shrink!(history)

    log ? (x, history) : x
end

function update_residual!(r::Residual, arnoldi::ArnoldiDecomp, k::Int)
    if iszero(arnoldi.H[k + 1, k])
        r.current = zero(r.current)
    else
        # Cheaply computes the current residual
        r.nullvec[k + 1] = -conj(dot(view(r.nullvec, 1 : k), view(arnoldi.H, 1 : k, k)) / arnoldi.H[k + 1, k])
        r.accumulator += abs2(r.nullvec[k + 1])
        r.current = r.β / √r.accumulator
    end
end

function init!(arnoldi::ArnoldiDecomp{T}, x, b, Pl, Ax; initially_zero::Bool = false) where {T}
    # Initialize the Krylov subspace with the initial residual vector
    # This basically does V[1] = Pl \ (b - A * x) and then normalize

    first_col = view(arnoldi.V, :, 1)

    copyto!(first_col, b)

    # Potentially save one MV product
    if !initially_zero
        mul!(Ax, arnoldi.A, x)
        first_col .-= Ax
    end

    ldiv!(Pl, first_col)

    # Normalize
    β = norm(first_col)
    first_col .*= inv(β)
    β
end

@inline function init_residual!(r::Residual{numT, resT}, β) where {numT,resT}
    r.accumulator = one(resT)
    r.β = β
end

function solve_least_squares!(arnoldi::ArnoldiDecomp{T}, β, k::Int) where {T}
    # Compute the least-squares solution to Hy = β e1 via Given's rotations
    rhs = zeros(T, k)
    rhs[1] = β

    H = FastHessenberg(view(arnoldi.H, 1 : k, 1 : k - 1))
    ldiv!(H, rhs)

    rhs
end

function update_solution!(x, y, arnoldi::ArnoldiDecomp{T}, Pr::Identity, k::Int, Ax) where {T}
    # Update x ← x + V * y
    mul!(x, view(arnoldi.V, :, 1 : k - 1), y, one(T), one(T))
end

function update_solution!(x, y, arnoldi::ArnoldiDecomp{T}, Pr, k::Int, Ax) where {T}
    # Computing x ← x + Pr \ (V * y) and use Ax as a work space
    mul!(Ax, view(arnoldi.V, :, 1 : k - 1), y)
    ldiv!(Pr, Ax)
    x .+= Ax
end

function expand!(arnoldi::ArnoldiDecomp, Pl::Identity, Pr::Identity, k::Int, Ax)
    # Simply expands by A * v without allocating
    mul!(view(arnoldi.V, :, k + 1), arnoldi.A, view(arnoldi.V, :, k))
end

function expand!(arnoldi::ArnoldiDecomp, Pl, Pr::Identity, k::Int, Ax)
    # Expands by Pl \ (A * v) without allocating
    nextV = view(arnoldi.V, :, k + 1)
    mul!(nextV, arnoldi.A, view(arnoldi.V, :, k))
    ldiv!(Pl, nextV)
end

function expand!(arnoldi::ArnoldiDecomp, Pl, Pr, k::Int, Ax)
    # Expands by Pl \ (A * (Pr \ v)). Avoids allocation by using Ax.
    nextV = view(arnoldi.V, :, k + 1)
    ldiv!(nextV, Pr, view(arnoldi.V, :, k))
    mul!(Ax, arnoldi.A, nextV)
    copyto!(nextV,  Ax)
    ldiv!(Pl, nextV)
end


function base_to_conservative(Qbase, cellxmax, cellymax, specific_heat_ratio, nval)
    """
    Qbase=[rho,u,v,p]
    Qcon=[rho,rhou,rhov,e]
    """
    Qcon = zeros(cellxmax, cellymax, nval)
    Threads.@threads for i in 1:cellxmax
        for j in 1:cellymax
            Qcon[i,j,1] = Qbase[i,j,1]
            Qcon[i,j,2] = Qbase[i,j,1]*Qbase[i,j,2]
            Qcon[i,j,3] = Qbase[i,j,1]*Qbase[i,j,3]
            Qcon[i,j,4] = Qbase[i,j,4]/(specific_heat_ratio-1)+Qbase[i,j,1]*(Qbase[i,j,2]^2+Qbase[i,j,3]^2)/2
        end
    end
    return Qcon
end

function conservative_to_base(Qcon, cellxmax, cellymax, specific_heat_ratio, nval)
    """
    Qbase=[rho,u,v,p]
    Qcon=[rho,rhou,rhov,e]
    """
    Qbase = zeros(cellxmax, cellymax, nval)
    Threads.@threads for i in 1:cellxmax
        for j in 1:cellymax
            Qbase[i,j,1] = Qcon[i,j,1]
            Qbase[i,j,2] = Qcon[i,j,2]/Qcon[i,j,1]
            Qbase[i,j,3] = Qcon[i,j,3]/Qcon[i,j,1]
            Qbase[i,j,4] = (Qcon[i,j,4]-Qcon[i,j,1]*(Qbase[i,j,2]^2+Qbase[i,j,3]^2)/2)*(specific_heat_ratio-1)
        end
    end
    return Qbase
end

function setup_Qcon_hat(Qcon, cellxmax, cellymax, volume, nval)
    Qcon_hat = zeros(cellxmax, cellymax, nval)
    Threads.@threads for i in 1:cellxmax
        for j in 1:cellymax
            for l in 1:nval
                Qcon_hat[i,j,l] = Qcon[i,j,l] * volume[i,j]
            end
        end
    end
    return Qcon_hat
end

function Qhat_to_Q(Qcon_hat, cellxmax, cellymax, volume, nval)
    Qcon = zeros(cellxmax, cellymax, nval)
    Threads.@threads for i in 1:cellxmax
        for j in 1:cellymax
            for l in 1:nval
                Qcon[i,j,l] = Qcon_hat[i,j,l] / volume[i,j]
            end
        end
    end
    return Qcon
end


using OrdinaryDiffEq, Test, ForwardDiff

# Test problem
# u' = p(t)u; p(t) = p if t <= tstop, 0 otherwise
# Target: du(1.0)/dtstop = exp(tstop)u0

u0 = 1.0
tspan = (0.0, 1.0)
p = 1.0

function stopping_cb(tstop)
    condition = (u,t,integrator) -> t == tstop
    affect! = integrator -> (println("Stopped!"); integrator.p = zero(integrator.p))
    DiscreteCallback(condition, affect!)
end

function test_fun(tstop)
    DualT = typeof(tstop)
    prob = ODEProblem((u,p,t) -> p*u, DualT(u0), DualT.(tspan), DualT(p))
    sol = solve(prob, Tsit5(), callback=stopping_cb(tstop), tstops=[tstop])
    sol(1.0)
end

@test ForwardDiff.derivative(test_fun, 0.5) ≈ exp(0.5)*u0 # Analytical solution: exp(tstop)*u0

function test_fun(tstop)
    DualT = typeof(tstop)
    prob = ODEProblem((u,p,t) -> p*u, DualT(u0), DualT.(tspan), DualT(p))
    sol = solve(prob, Tsit5(), callback=stopping_cb(tstop), tstops=[tstop], adaptive=false, dt=0.01)
    sol(1.0)
end

@test ForwardDiff.derivative(test_fun, 0.5) ≈ exp(0.5)*u0 # Analytical solution: exp(tstop)*u0


module Primeau_2x2x2
#=
This circulation and grid was created by François Primeau (fprimeau@uci.edu)
It was taken from his Notebook at
https://github.com/fprimeau/BIOGEOCHEM_TEACHING/blob/master/Intro2TransportOperators.ipynb

The simple box model we consider is embeded in a 2×2×2 "shoebox".
It has 5 wet boxes and 3 dry boxes.
See image at
https://github.com/fprimeau/BIOGEOCHEM_TEACHING/blob/master/boxmodel.png
=#

using LinearAlgebra, SparseArrays
using Unitful
using Reexport
@reexport using OceanGrids            # To store the grid
using ..CirculationGeneration
CG = CirculationGeneration

function build_wet3D()
    wet3D = trues(2, 2, 2)
    wet3D[[4,7,8]] .= false # land points
    return wet3D
end

function build_grid()
    elon = [0,180,360] * u"°"
    elat = [-90, 0, 90] * u"°"
    edepth = [0, 200, 3700] * u"m"
    return OceanGrid(elon, elat, edepth, build_wet3D())
end

function build_T(grid)
    # From Archer et al. [2000]
    v3D = array_of_volumes(grid)
    nb = length(v3D)

    # Antarctic Circumpoloar Current 1 -> 3 -> 1
    ACC = 100e6u"m^3/s"
    T  = CG.T_advection(ACC, [1, 3, 1], v3D, nb)
    # Meridional Overturning Circulation 1 -> 2 -> 6 -> 5 -> 1
    MOC = 15e6u"m^3/s"
    T += CG.T_advection(MOC, [1, 2, 6, 5, 1], v3D, nb)
    # vertical mixing at "high northern latitudes" 2 <-> 6
    MIX = 10e6u"m^3/s"
    T += CG.T_diffusion(MIX, 2, 6, v3D, nb)

    # Only keep wet points
    iwet = findall(vec(grid.wet3D))

    return T[iwet, iwet]
end

"""
    load

Returns `grd` and `T` (in that order).
"""
function load()
    print("Creating François Primeau's 2x2x2 model")
    grd = build_grid()
    T = build_T(grd)
    println(" ✔")
    return grd, ustrip.(T)
end

end

export Primeau_2x2x2


using IntervalArithmetic
using Test
using StaticArrays
using Random

@testset "rand tests" begin

    X = 3..4
    for i in 1:100
        @test rand(X) ∈ X
    end

    Y = IntervalBox(3..4, 5..6)
    for i in 1:100
        @test rand(Y) ∈ Y
    end

    X = 3..4
    for i in 1:100
        @test rand(X,4)[(i%4)+1] ∈ X
    end

#     Y = IntervalBox(3..4, 5..6)
#     for i in 1:100
#         @test rand(Y,2)[(i%2)+1][(i%2)+1] ∈ Y[(i%2)+1]
#     end

    for T in (Float32, Float64, BigFloat)
        X = Interval{T}(3, 4)
        @test rand(X) isa T

        Y = IntervalBox(X, X)
        @test rand(Y) isa SVector{2,T}
    end

    for T in (Float32, Float64, BigFloat)
        X = Interval{T}(3, 7)
        Y = rand(X, 7)
        @test Y isa Array{T, 1}
        for x in Y
            @test x isa T
        end
    end
end


######################
# Logit and Logistic #
######################
using StatsFuns: logit, logistic

struct Logit{N, T<:Real} <: Bijector{N}
    a::T
    b::T
end
function Logit(a, b)
    T = promote_type(typeof(a), typeof(b))
    Logit{0, T}(a, b)
end
up1(b::Logit{N, T}) where {N, T} = Logit{N + 1, T}(b.a, b.b)
# For equality of Logit with Float64 fields to one with Duals
Base.:(==)(b1::Logit, b2::Logit) = b1.a == b2.a && b1.b == b2.b

(b::Logit{0})(x::Real) = _logit(x, b.a, b.b)
(b::Logit{0})(x) = _logit.(x, b.a, b.b)
(b::Logit{1})(x::AbstractVector) = _logit.(x, b.a, b.b)
(b::Logit{1})(x::AbstractMatrix) = _logit.(x, b.a, b.b)
(b::Logit{2})(x::AbstractMatrix) = _logit.(x, b.a, b.b)
(b::Logit{2})(x::AbstractArray{<:AbstractMatrix}) = map(b, x)
_logit(x, a, b) = logit((x - a) / (b - a))

(ib::Inverse{<:Logit{0}})(y::Real) = _ilogit(y, ib.orig.a, ib.orig.b)
(ib::Inverse{<:Logit{0}})(y) = _ilogit.(y, ib.orig.a, ib.orig.b)
(ib::Inverse{<:Logit{1}})(x::AbstractVecOrMat) = _ilogit.(x, ib.orig.a, ib.orig.b)
(ib::Inverse{<:Logit{2}})(x::AbstractMatrix) = _ilogit.(x, ib.orig.a, ib.orig.b)
(ib::Inverse{<:Logit{2}})(x::AbstractArray{<:AbstractMatrix}) = map(ib, x)
_ilogit(y, a, b) = (b - a) * logistic(y) + a

logabsdetjac(b::Logit{0}, x::Real) = logit_logabsdetjac(x, b.a, b.b)
logabsdetjac(b::Logit{0}, x) = logit_logabsdetjac.(x, b.a, b.b)
logabsdetjac(b::Logit{1}, x::AbstractVector) = sum(logit_logabsdetjac.(x, b.a, b.b))
logabsdetjac(b::Logit{1}, x::AbstractMatrix) = vec(sum(logit_logabsdetjac.(x, b.a, b.b), dims = 1))
logabsdetjac(b::Logit{2}, x::AbstractMatrix) = sum(logit_logabsdetjac.(x, b.a, b.b))
logabsdetjac(b::Logit{2}, x::AbstractArray{<:AbstractMatrix}) = map(x) do x
    logabsdetjac(b, x)
end
logit_logabsdetjac(x, a, b) = -log((x - a) * (b - x) / (b - a))


#------------------------------------------------------------------------------
#  Example.jl
#
#
#
#
#
#
#
#
#
#  Paul.Soderlind@unisg.ch, April 2002, to Julia Nov 2015
#------------------------------------------------------------------------------

using Dates, Optim

include("jlFiles/BondPricePs.jl")
include("jlFiles/BondNSxPs.jl")
include("jlFiles/printmat.jl")

ytmLoss  = 0           #0: mimimize squared price errors: 1: minimize squared ytm errors
  weightLoss = 1       #1: 1/mat as weights for (fitted price-price)^2 in loss fn (only if ytmLoss==0)
#------------------------------------------------------------------------------

                                 #Swedish bond data for 29 Dec 1993

#coupons in %/yr
c  = [ 0;        0;         0;         0;         11.5;      10.75;
       11;       13;        10.25;     6;         9 ]

#time to maturity in year
tm = [ 0.00274;  0.21096;   0.46027;   0.88219;   1.67397;   3.06849;
       5.06301;  7.46027;   9.34795;   11.11507;  15.30685 ]

#interest rates.  Bills: simple rates in %/yr; bonds: yield to maturity in %/yr
y  = [ 7.75;     6.835;    6.655;     6.41;      6.215;     6.195;
       6.41;     6.755;    7.01;      7.21;      7.325 ]

n = length(y)       #number of bonds
#----------------------------------------------------------------------------

#transform the data

c  = c/100 # -> column vector, coupons and yields as 0.05 rather than 5
y  = y/100

vvc    = c .== 0                #if bill, change from simple to effective rate
y[vvc] = (1 .+ tm[vvc].*y[vvc]).^(1.0./tm[vvc]) .- 1

P = fill(NaN,length(y))         #calculate bond prices
for i = 1:n
  local ti
  ti   = mod(tm[i],1):tm[i]
  P[i] = BondPricePs(y[i],c[i],ti)
end
#----------------------------------------------------------------------------

#estimating parameters in extended Nelson&Siegel model, restricted so b0 + b1 = log(1+y[1])

parX0 = [0.1045;-0.03;-0.0562;1.2;0;0.5]       #starting guess

if ytmLoss == 1                                #loss(ytm)
  NSXbR = BondNSxEstPs(parX0,y,tm,c,log(1+y[1]),1)
else
  if weightLoss == 1                           #loss(P/maturity)
    NSXbR = BondNSxEstPs(parX0,P,tm,c,log(1+y[1]),0,1.0./tm)
  else                                         #loss(P)
    NSXbR = BondNSxEstPs(parX0,P,tm,c,log(1+y[1]))
  end
end

println("\nParameter estimates: ")
printTable(NSXbR,[""],["b0","b1","b2","tau","b3","tau2"])
#----------------------------------------------------------------------------

ytmx = fill(NaN,n)            #model implied ytm
for i = 1:n
  local ti,d,Qx
  ti      = mod(tm[i],1):tm[i]
  d       = BondNSxPs(ti,NSXbR...)[3]    #... expands into NSXbR[1],NSXbR[2],...
  Qx      = sum(d.*c[i]) + d[end]     #model implied bond price
  #println(ti)
  ytmx[i] = BondYieldToMatPs(Qx,c[i],ti,1,1,0.05,1e-7)[1]
end

println("\nActual and model ytm, %: ")
printTable([y ytmx]*100,["ytm actual","ytm model"],string.(tm),width=15,cell00="maturity")
#----------------------------------------------------------------------------

#calculate model implied rates (spot, forward, yield to maturity) to plot

tmFig      = [1e-8;0.1:0.1:16]           #maturities to plot
(shx,fhx,) = BondNSxPs(tmFig,NSXbR...)
shx        = exp.(shx) .- 1     #effective interest rate
fhx        = exp.(fhx) .- 1

println("\nmodel spot and forward rates, %: ")
printTable([shx fhx]*100,["spot","forward"],string.(tmFig),width=15,cell00="maturity")
#----------------------------------------------------------------------------

#Comment out this if you do not have PyPlot installed.

using PyPlot
#PyPlot.svg(true)           #for ipynb notebooks
close("all")
                                       #plotting
figure()
  plot(tm,c,"+",tm,y,"s",tmFig,shx,"b--",tmFig,fhx,"r-",tm,ytmx,".")
  title("Swedish Interest Rates 29 Dec 1993")
  xlabel("Years to Maturity")
  legend(["Coupon rate","Yield to maturity","Estimated spot rate",
         "Estimated forward rate","Estimated yield to maturity"],loc=1)
  ylim(0.04,0.14)
  #display(gcf())            #uncomment in Atom/Juno
#----------------------------------------------------------------------------


function my_f(x, y)
    return 2*x + 3*y
end;

function derivative_of_my_f(x, y)
    return ForwardDiff.derivative(t -> my_f(t, y), x)
end;


## Sample Space: Omega
## ==================

# Omega mathematically is defined as a unit n dimensional Euclidean hypercube.
# It is useful in different contexts to have types for
# 1. concrete elements ω in Omega - vectors or Real values
# 2. abstract values in Omega, which are subsets of Omega.
# 3. random values of Omega.
# Because we wan't to allow users to write moels without worrying about Omega,
# we use LazyBoxes to generate needed dimensions of Omega when required.

LazyOmega() = LazyBox(Float64)
LazyOmega{T<:Real}(T2::Type{T}) = LazyBox(T2)

"Create a Euclidean unit box with dimensions `dims`"
function unit_box{T<:Real}(::Type{LazyBox{T}}, dims::Set{Int})
  abox = LazyBox(T)
  for dim in dims
    abox[dim] = Interval{T}(zero(T), one(T))
  end
  abox
end

"All kinds of Omega"
typealias Omega{T} Union{Vector{T},HyperBox{T}, LazyBox{T}, LazyRandomVector{T}, Dict{Int,T}}

"Abstract representations of sample space - euclidean box"
typealias AbstractOmega{T} Union{HyperBox{T}, LazyBox{T}}

"A concrete (i.e. not abstract) element ω in sample space Ω"
typealias ConcreteOmega{T} Union{Vector{T}, LazyRandomVector{T}}

"Build an omega of `n` dimensions"
abstract_omega(n::Int) = HyperBox([Interval(0.0,1.0) for i = 1:n])


using Parameters
include("common.jl")

# Na-K ATPase (NKA) parameters
@with_kw struct NKAParams
    k_o = 5.4
    na_o = 140.0
    # Max Na-K ATPase current (μA/cm²)
    IMAX = 3.147
    # Na half-saturate constant of Na-K ATPase (mM)
    KM_NA = 10.0
    # K half-saturate constant of Na-K ATPase (mM)
    KM_K = 1.5
    # ATP half-saturate constant of Na-K ATPase (mM)
    KM_ATP = 8E-3
    # ADP half-saturate constant of Na-K ATPase (mM)
    KI_ADP = 0.1
    # Factor of extracellular potassium
    FAC_K = _mm(k_o, KM_K)
    # Factor of extracellular sodium
    FAC_NA = expm1(na_o / 67.3) / 7
end

# Na/K ATPase current (iNaK) (uA/cm^2), scalar version
function inak(na_i, atp_i, adp_i, vm, evfrtm1, KM_ATP, KI_ADP, KM_NA, IMAX, FAC_K, FAC_NA)
    evfrt = evfrtm1 + 1
    f_na = _hills(na_i, KM_NA, 1.5)
    f_nak_inv = 1.0 + 0.1245 * _ra(-0.1 * vm) + 0.0365 * FAC_NA / evfrt
    f_atp = _mm(atp_i * _mm(KI_ADP, adp_i), KM_ATP)
    iNak = IMAX * FAC_K * f_na * f_atp / f_nak_inv
end

function inak(na_i, atp_i, adp_i, vm, evfrtm1, pNKA::NKAParams)
    @unpack KM_ATP, KI_ADP, KM_NA, IMAX, FAC_K, FAC_NA = pNKA
    inak(na_i, atp_i, adp_i, vm, evfrtm1, KM_ATP, KI_ADP, KM_NA, IMAX, FAC_K, FAC_NA)
end

# Plasma membrane calcium ATPase (PMCA) parameters
@with_kw struct PMCAParams
    # Max PMCA current (μA/cm²)
    IMAX = 0.575
    # Ca half-saturation constant (mM)
    KM_CA = 5E-4
    # ATP 1st half-saturation constant (mM)
    KM1_ATP = 0.012
    # ATP 2nd half-saturation constant (mM)
    KM2_ATP = 0.23
    # ADP half-inhibition constant (mM)
    KI_ADP = 1.0
end

# PMCA current (iPCa) (uA/cm^2), scalar version
function ipca(ca_i, atp_i, adp_i, KM1_ATP, KI_ADP, KM2_ATP, KM_CA, IMAX)
    f_atp = _mm(atp_i * _mm(KI_ADP, adp_i), KM1_ATP) + _mm(atp_i, KM2_ATP)
    f_ca = _mm(ca_i, KM_CA)
    iPca = IMAX * f_atp * f_ca
end

function ipca(ca_i, atp_i, adp_i, pPMCA::PMCAParams)
    @unpack KM1_ATP, KI_ADP, KM2_ATP, KM_CA, IMAX = pPMCA
    ipca(ca_i, atp_i, adp_i, KM1_ATP, KI_ADP, KM2_ATP, KM_CA, IMAX)
end

# SER ca ATPase (SERCA) parameters
# Params from Li et al. (2015)
@with_kw struct SERCAParams
    # Max forward rate (mM/ms)
    VF = 2.989E-4
    # Max reverse rate (mM/ms)
    VR = 3.179E-4
    # Michaelis constant for Ca of forward reaction (mM)
    KMF_CA = 1.5E-4
    # Michaelis constant for Ca of reverse reaction (mM)
    KMR_CA = 3.3
    # Cooperativity of forward reaction
    NFB = 0.55
    # Cooperativity of reverse reaction
    NRB = 0.5
    # Michaelis constant for ATP (mM)
    KM_ATP = 0.01
    # 1st michaelis constant for ADP (mM)
    KI1_ADP = 0.14
    # 2nd michaelis constant for ADP (mM)
    KI2_ADP = 5.1
end

# SERCA calcium flux (jUp) (mM/ms), scalar version
function jup(ca_i, ca_nsr, atp_i, adp_i,
             NFB, NRB, KMF_CA, KMR_CA, KM_ATP, KI1_ADP, KI2_ADP, VF, VR)
    fb = (ca_i / KMF_CA) ^ NFB
    rb = (ca_nsr / KMR_CA) ^ NRB
    f_atp = inv(KM_ATP / atp_i * _mm_reci(KI1_ADP, adp_i) + _mm_reci(KI2_ADP, adp_i))
    jUp = f_atp * (VF * fb - VR * rb) / (1.0 + fb + rb)
end

function jup(ca_i, ca_nsr, atp_i, adp_i, pSERCA::SERCAParams)
    @unpack NFB, NRB, KMF_CA, KMR_CA, KM_ATP, KI1_ADP, KI2_ADP, VF, VR = pSERCA
    jup(ca_i, ca_nsr, atp_i, adp_i,
        NFB, NRB, KMF_CA, KMR_CA, KM_ATP, KI1_ADP, KI2_ADP, VF, VR)
end


function cd3_dissociation!(du,u,p,t,dens,cons)
    # unpack rates and constants
    nᵣ,nₓ,n₃ = u
    k₋₃ = p[1]
    k₋₁,k₋₂ = cons
    # mᵣ,mₗ,mₓ,A = dens

    # model
    du[1] = dnᵣ = - k₋₁*nᵣ
    du[2] = dnₓ = - k₋₂*nₓ
    du[3] = dn₃ = - k₋₃*n₃

end

function bimolecular_diss_bell_activated_state_w_cons!(du,u,p,t,dens,cons;f=0.0)
    # unpack rates and constants
    Pᵣ,Pₐ = u
    x₋ᵒf,
        k₋ᵒs,x₋ᵒs,
        kᵒa,xᵒa,
        k₋ᵒa,x₋ᵒa = p
    k₋ᵒf = cons[1] # zero force diss, generally obtained from thermal fluct.

    k₋f = bell_diss_kT(k₋ᵒf, x₋ᵒf, f) # fast diss form L-R
    k₋s = bell_diss_kT(k₋ᵒs, x₋ᵒs, f) # slow diss from L-R*
    ka  = bell_diss_kT(kᵒa, xᵒa, f)   # activation from L-R to L-R*
    k₋a = bell_ass_kT(k₋ᵒa, x₋ᵒa, f)  # deactivation from L-R* to L-R

    # model
    du[1] = dPᵣ = - k₋f.*Pᵣ - ka.*Pᵣ + k₋a.*Pₐ
    du[2] = dPₐ = + ka.*Pᵣ - k₋a.*Pₐ - k₋s.*Pₐ
end

function cd3_tri_dissociation_bell!(du,u,p,t,dens,cons;f=0.0)
    # unpack rates and constants
    Pᵣ₁,Pₐ₁,
        Pᵣ₂,Pₐ₂,
        Pᵣ₃,Pₐ₃ = u

    x₋₃ᵒf,
        k₋₃ᵒs, x₋₃ᵒs,
        k₃ᵒa, x₃ᵒa,
        k₋₃ᵒa, x₋₃ᵒa = p

    #  for now declare the kᵢ from bimolecular as global variables
    k₋₁ᵒf, x₋₁ᵒf,
        k₋₁ᵒs, x₋₁ᵒs,
        k₁ᵒa, x₁ᵒa,
        k₋₁ᵒa, x₋₁ᵒa,
        k₋₂ᵒf, x₋₂ᵒf,
        k₋₂ᵒs, x₋₂ᵒs,
        k₂ᵒa, x₂ᵒa,
        k₋₂ᵒa, x₋₂ᵒa,
        k₋₃ᵒf         = cons

    k₋₁f = bell_diss_kT(k₋₁ᵒf, x₋₁ᵒf, f) # fast diss form L-R
    k₋₁s = bell_diss_kT(k₋₁ᵒs, x₋₁ᵒs, f) # slow diss from L-R*
    k₁a  = bell_diss_kT(k₁ᵒa, x₁ᵒa, f)   # activation from L-R to L-R*
    k₋₁a = bell_ass_kT(k₋₁ᵒa, x₋₁ᵒa, f)  # deactivation from L-R* to L-R

    k₋₂f = bell_diss_kT(k₋₂ᵒf, x₋₂ᵒf, f) # fast diss form L-R
    k₋₂s = bell_diss_kT(k₋₂ᵒs, x₋₂ᵒs, f) # slow diss from L-R*
    k₂a  = bell_diss_kT(k₂ᵒa, x₂ᵒa, f)   # activation from L-R to L-R*
    k₋₂a = bell_ass_kT(k₋₂ᵒa, x₋₂ᵒa, f)  # deactivation from L-R* to L-R

    k₋₃f = bell_diss_kT(k₋₃ᵒf, x₋₃ᵒf, f) # fast diss form L-R
    k₋₃s = bell_diss_kT(k₋₃ᵒs, x₋₃ᵒs, f) # slow diss from L-R*
    k₃a  = bell_diss_kT(k₃ᵒa, x₃ᵒa, f)   # activation from L-R to L-R*
    k₋₃a = bell_ass_kT(k₋₃ᵒa, x₋₃ᵒa, f)  # deactivation from L-R* to L-R

    # model
    # γϵ
    du[1] = dPᵣ₁ = - k₋₁f*Pᵣ₁ - k₁a*Pᵣ₁ + k₋₁a*Pₐ₁
    du[2] = dPₐ₁ = + k₁a*Pᵣ₁ - k₋₁a*Pₐ₁ - k₋₁s*Pₐ₁
    # δϵ
    du[3] = dPᵣ₂ = - k₋₂f*Pᵣ₂ - k₂a*Pᵣ₂ + k₋₂a*Pₐ₂
    du[4] = dPₐ₂ = + k₂a*Pᵣ₂ - k₋₂a*Pₐ₂ - k₋₂s*Pₐ₂
    # tri
    du[5] = dPᵣ₃ = - k₋₃f*Pᵣ₃ - k₃a*Pᵣ₃ + k₋₃a*Pₐ₃
    du[6] = dPₐ₃ = + k₃a*Pᵣ₃ - k₋₃a*Pₐ₃ - k₋₃s*Pₐ₃

end

## functions to calculate initial bond distribution vector (at t=0)

#=
if the delay between binding and beginning of the measurement is sufficiently
long, the probability of occupying either state will reach an equilibrium. This
equilibrium is independent of force
=#
function initial_state(kᵒa,k₋ᵒa)
        u₀1 = k₋ᵒa / (kᵒa + k₋ᵒa)
        u₀ = [u₀1, (1.0 - u₀1)]
end

#=
assumes that the time between binding and measurement is short relative to the
transition rates between the two bound states, then the bond will stay in the
initial state it occupied upon binding. The probabilities of occupying each of
these states can be derived from the principle of detailed balance.
=#
function initial_state(k₋ᵒf,k₋ᵒs,kᵒa,k₋ᵒa)
        u₀1 = (k₋ᵒa * k₋ᵒf) / (k₋ᵒa * k₋ᵒf + kᵒa * k₋ᵒs)
        u₀ = [u₀1, (1.0 - u₀1)]
end

#=
assumes enough time to reach equilibrium, and that the bond experiences some
fraction a of the measured force.
IMPORTANT: must add a*f to function call: initial_state(p, a * f)
=#

function initial_state(p::Vector,f)
        k₋ᵒf,x₋ᵒf,
        k₋ᵒs,x₋ᵒs,
        kᵒa,xᵒa,
        k₋ᵒa,x₋ᵒa = p

        k₋f = bell_diss_kT(k₋ᵒf, x₋ᵒf, f) # fast diss form L-R
        k₋s = bell_diss_kT(k₋ᵒs, x₋ᵒs, f) # slow diss from L-R*
        ka  = bell_diss_kT(kᵒa, xᵒa, f)   # activation from L-R to L-R*
        k₋a = bell_ass_kT(k₋ᵒa, x₋ᵒa, f)

        u₀1 = k₋a / (ka + k₋a)
        u₀ = [u₀1, (1.0 - u₀1)]
end

# function
# u = []
# for f in Fs
#         u₀ = initial_state(rates,f)
#         push!(u,u₀)
# end
#
# plot(reduce(hcat,u))

function initial_u(k,M)
    if ~isempty(M.cons)
        k = vcat(M.cons,k)
    end

    if M.u₀_type == "equilibrium"
        k₋ᵒf,x₋ᵒf,k₋ᵒs,x₋ᵒs,kᵒa,xᵒa,k₋ᵒa,x₋ᵒa = k
        new_u₀ = initial_state(kᵒa,k₋ᵒa)

    elseif M.u₀_type == "detailed_balance"
        k₋ᵒf,x₋ᵒf,k₋ᵒs,x₋ᵒs,kᵒa,xᵒa,k₋ᵒa,x₋ᵒa = k
        new_u₀ = initial_state(k₋ᵒf,k₋ᵒs,kᵒa,k₋ᵒa)

    elseif M.u₀_type == "force"
        f = M.cons
        a = k[end]
        new_u₀ = initial_state(k,a*f)

    else
        new_u₀ = []
    end

    if isempty(M.u₀)
        u₀ = new_u₀
    else
        u₀ = M.u₀
        u₀ = convert.(eltype(new_u₀),u₀)
        u₀ = append!(u₀, new_u₀) # append new initial states to provided initial states
    end
    return u₀
end


# https://en.wikipedia.org/wiki/MU_puzzle#Solution

using Metatheory

miu = @theory x y z begin
  # Composition of the string monoid is associative
  x ⋅ (y ⋅ z) --> (x ⋅ y) ⋅ z
  # Add a uf to the end of any string ending in I
  x ⋅ :I ⋅ :END --> x ⋅ :I ⋅ :U ⋅ :END
  # Double the string after the M
  :M ⋅ x ⋅ :END --> :M ⋅ x ⋅ x ⋅ :END
  # Replace any III with a U
  :I ⋅ :I ⋅ :I --> :U
  # Remove any UU
  x ⋅ :U ⋅ :U ⋅ y --> x ⋅ y
end

@testset "MU puzzle" begin
  # no matter the timeout we set here,
  # MU is not a theorem of the MIU system
  params = SaturationParams(timeout = 12, eclasslimit = 8000)
  start = :(M ⋅ I ⋅ END)
  g = EGraph(start)
  saturate!(g, miu)
  @test false == areequal(g, miu, start, :(M ⋅ U ⋅ END); params = params)
end