include("util.jl")
include("graph.jl")
include("opinion.jl")

# function convex_approach(Q::Array{WEIGHT, 2}, y::OPINIONS, budget::WEIGHT)
#     z = Q * y
#     N = size(y)
#     N = N[1]
#     δ = Convex.Variable(N)
#     problem = minimize(sumsquares(z + (Q * δ)))
#     problem.constraints += (z + δ) <= 1
#     problem.constraints += (z + δ) >= -1
#     problem.constraints += sum((dot(^)(δ,2))) <= budget
#     solve!(problem, SCS.Optimizer)
#     @show problem.status
#     @show (z + (Q * δ.value))
#     return δ.value
# end

cost_of_removing(x) = x * x

function convex_approach(Q::Matrix{WEIGHT}, y::OPINIONS,
                        α::PSYCHOLOGICAL_FACTOR, budget::WEIGHT)
    z = Q * (α .* y)
    N = size(y)
    N = N[1]
    δ = Variable(N)
    problem = minimize(sumsquares(z + (Q * (α .* δ))))
    problem.constraints += (z + δ) <= 1
    problem.constraints += (z + δ) >= -1
    problem.constraints += sum((dot(^)(δ,2))) <= budget
    solve!(problem, SCS.Optimizer)
    @show problem.status
    @show Q * (α .* (y + δ))
    return δ.value
end

function mult_but_exclude(M::Matrix{WEIGHT}, y::OPINIONS,
    exclude::Array{Int64, 1})

    temp_M = M[:, setdiff(1:end, exclude)]
    temp_y = y[setdiff(1:end, exclude)]

    return (temp_M * temp_y)

end

function greedy_approach(Q::Matrix{WEIGHT}, y::OPINIONS, α::PSYCHOLOGICAL_FACTOR, budget::WEIGHT)

    (n, m) = size(Q)
    T = Q * diagm(α)
    total_cost = 0
    prev_len = 0
    to_remove = fill(0, 0)

    # while we can still potentially remove a node's influence
    while total_cost < budget
        current_cost = total_cost
        best_so_far = WEIGHT(Inf)
        current_to_remove = -1

        # consider each node
        for i = 1:n
            # get cost of adding node
            cost_to_add_i = (cost_of_removing(y[i]) + current_cost)

            # if we have not considered adding that node and
            # our budget would allow us to add them
            if i ∉ to_remove && (cost_to_add_i <= budget)
                # consider them
                push!(to_remove, i)
                # get change to the norm
                marginal_norm = norm(mult_but_exclude(T, y, to_remove))

                # if the norm obtained is lower than our current best in this
                # iteration
                if marginal_norm < best_so_far
                    # flag it as the best to remove
                    current_to_remove = i
                    best_so_far = marginal_norm
                end
                # pop off to remove temporarily
                pop!(to_remove)
            end
        end

        # if we found a viable node to remove
        if current_to_remove != -1
            # add it to remove list
            push!(to_remove, current_to_remove)
             #@show current_to_remove
        end

        # if (current_cost == total_cost) || (length(to_remove) == n)
        #     break
        if (length(to_remove) == n)
            break
        else
            total_cost = current_cost + cost_of_removing(y[to_remove[end]])
        end
    end

    δ = WEIGHT.(zeros(n))
    δ[to_remove] = -y[to_remove]

    return δ, total_cost, (T * (y + δ))
end

function select_using_centrality(Q::Matrix{WEIGHT}, y::OPINIONS,
    α::PSYCHOLOGICAL_FACTOR, budget::WEIGHT,
    G::SimpleGraph{Int64}, centrality::String)

    centralities = CENTRALITIES[centrality](G)
    (n, m) = size(Q)
    T = Q * diagm(α)

    value = fill(WEIGHT(0.0), n)
    cost = fill(WEIGHT(0.0), n)
    z_inf = T * y
    original = norm(z_inf)
    mean_y = mean(y)
    for i = 1:n
        value[i] = 1 * centralities[i]
        cost[i] = cost_of_removing(y[i])
    end
    perm = sortperm(-value)
    δ = fill(WEIGHT(0.0), n)
    total_cost = 0
    for i ∈ perm
        if (cost[i] + total_cost) <= budget
            total_cost += cost[i]
            δ[i] = -y[i]
        else
            remaining = budget - total_cost
            fraction = sqrt(remaining)
            δ[i] = -(sign(y[i])) * fraction
            break
        end
    end
    z = T * (y + δ)
    return (δ, total_cost, z)

end

function select_using_pagerank(Q::Matrix{WEIGHT}, y::OPINIONS,
    α::PSYCHOLOGICAL_FACTOR, budget::WEIGHT, A)

    g = construct_graph(A)
    centralities = closeness_centrality(g)

    (n, m) = size(Q)
    T = Q * diagm(α)
    value = fill(WEIGHT(0.0), n)
    cost = fill(WEIGHT(0.0), n)
    z_inf = T * y
    original = norm(z_inf)
    for i = 1:n
        value[i] = centralities[i]
        cost[i] = cost_of_removing(y[i])
    end

    perm = sortperm(-value)
    δ = fill(WEIGHT(0.0), n)
    total_cost = 0
    for i ∈ perm
        if (cost[i] + total_cost) <= budget
            total_cost += cost[i]
            δ[i] = -y[i]
        else
            remaining = budget - total_cost
            fraction = sqrt(remaining)
            δ[i] = -(sign(y[i])) * fraction
            break
        end
    end
    z = T * (y + δ)
    return (δ, total_cost, z)

end

function fractional_knapsack(Q::Matrix{WEIGHT}, y::OPINIONS,
    α::PSYCHOLOGICAL_FACTOR, budget::WEIGHT)
    (n, m) = size(Q)
    T = Q * diagm(α)
    value = fill(WEIGHT(0.0), n)
    cost = fill(WEIGHT(0.0), n)
    z_inf = T * y
    original = norm(z_inf)
    for i = 1:n
        v = mult_but_exclude(T, y, [i])
        v = norm(z_inf) - norm(v)
        value[i] = 0
        if value[i] > 0
            value[i] = v
        end
        cost[i] = cost_of_removing(y[i])
    end

    perm = sortperm(-value)
    δ = fill(WEIGHT(0.0), n)
    total_cost = 0
    for i ∈ perm
        if (cost[i] + total_cost) <= budget
            total_cost += cost[i]
            δ[i] = -y[i]
        else
            remaining = budget - total_cost
            fraction = sqrt(remaining)
            δ[i] = -(sign(y[i])) * fraction
            break
        end
    end
    z = T * (y + δ)
    return (δ, total_cost, z)
end

function plot_multiple_series(series)
    p = plot()
    n, m = size(series)
    x = 1:m
    for i = 1:n
        plot!(p, x, series[i, :])
    end
end

function approx_equlibrium_vector(M::SparseMatrixCSC, y::OPINIONS,
     α::PSYCHOLOGICAL_FACTOR, k::Int64)
    #k = Int64(ceil(compute_num_iters(p.y, p.M, epsilon, p.α)))
    z = y
    for i = 1:k
        z = (α .* y) + (M * z)
    end
    return z
 end

 function greedy_approach_approx(M::SparseMatrixCSC{WEIGHT}, y::OPINIONS,
     α::PSYCHOLOGICAL_FACTOR, budget::WEIGHT, epsilon::WEIGHT)

     k = Int64(ceil(compute_num_iters(y, M, epsilon, α)))

     (n, m) = size(M)
     #T = Q * diagm(α)
     total_cost = 0
     to_remove = fill(0, 0)

     # while we can still potentially remove a node's influence
     while (total_cost < budget) && (length(to_remove) < n)
         current_cost = total_cost
         best_so_far = WEIGHT(Inf)
         current_to_remove = -1

         # consider each node
         for i = 1:n
             # get cost of adding node
             cost_to_add_i = (cost_of_removing(y[i]) + current_cost)

             # if we have not considered adding that node and
             # our budget would allow us to add them
             if i ∉ to_remove && (cost_to_add_i <= budget)
                 # consider them
                 push!(to_remove, i)
                 # get change to the norm
                 temp_y = copy(y)
                 temp_y[to_remove] .= 0
                 temp_eqilibrium = approx_equlibrium_vector(M, temp_y, α, k)
                 marginal_norm = norm(temp_eqilibrium)

                 # if the norm obtained is lower than our current best in this
                 # iteration
                 if marginal_norm < best_so_far
                     # flag it as the best to remove
                     current_to_remove = i
                     best_so_far = marginal_norm
                 end
                 # pop off to remove temporarily
                 pop!(to_remove)
             end
         end

         # if we found a viable node to remove
         if current_to_remove != -1
             # add it to remove list
             push!(to_remove, current_to_remove)
             #@show current_to_remove
         end

         # if (current_cost == total_cost) || (length(to_remove) == n)
         #     break
         total_cost = current_cost + cost_of_removing(y[to_remove[end]])
     end

     δ = WEIGHT.(zeros(n))
     δ[to_remove] = -y[to_remove]
     # @show M
     # @show y + δ
     z = approx_equlibrium_vector(M, y + δ, α, k)
     return δ, total_cost, z
 end

function get_approx_mat(M::SparseMatrixCSC{WEIGHT}, α::PSYCHOLOGICAL_FACTOR, k::Int64)

    p1 = M ^ 0
    powers = Array{SparseMatrixCSC{WEIGHT}, 1}()
    push!(powers, p1)
    for i = 1:k
        temp = powers[i] * M
        push!(powers, temp)
    end

    p2 = sum(powers[1:end-1]) * diagm(α)
    p2 += powers[end]
    return p2

end

 function fractional_knapsack_approx(M::SparseMatrixCSC{WEIGHT}, y::OPINIONS,
     α::PSYCHOLOGICAL_FACTOR, budget::WEIGHT, epsilon::WEIGHT)

     k = Int64(ceil(compute_num_iters(y, M, epsilon, α)))

     (n, m) = size(M)
     value = fill(WEIGHT(0.0), n)
     cost = fill(WEIGHT(0.0), n)
     # T = get_approx_mat(M, α, k)
     # @show T
     # z_inf = norm(T * y)
     z_inf = approx_equlibrium_vector(M, y, α, k)
     original = norm(z_inf)
     for i = 1:n
         temp_y = copy(y)
         temp_y[i] = 0
         v = approx_equlibrium_vector(M, temp_y, α, k)
         #v = T * temp_y
         v = norm(z_inf) - norm(v)
         value[i] = 0
         if value[i] > 0
             value[i] = v
         end
         cost[i] = cost_of_removing(y[i])
     end

     perm = sortperm(-value)
     δ = fill(WEIGHT(0.0), n)
     total_cost = 0
     for i ∈ perm
         if (cost[i] + total_cost) <= budget
             total_cost += cost[i]
             δ[i] = -y[i]
         else
             remaining = budget - total_cost
             fraction = sqrt(remaining)
             δ[i] = -(sign(y[i])) * fraction
             break
         end
     end
     z = approx_equlibrium_vector(M, y + δ, α, k)
     return (δ, total_cost, z)
 end

# function main()
#
#     DATASET = "congress"
#     BUDGET = WEIGHT(3.0)
#     filepath = "./data/raw/$(DATASET).edgelist"
#
#
#     @show filepath
#     (W, U, n) = read_edgelist(filepath)
#     normalize_edgelist!(W)
#     normalize_edgelist!(U)
#
#     y = generate_opinions(n)
#     α, β, γ = generate_factors(n)
#     M = (β .* W) - (γ .* U)
#     G = construct_graph(M)
#
#     M = Matrix(M)
#     Q = inv(I - M)
#
#     z = Q * (α .* y)
#
#
#
#     (δ, total_cost, z_prime) =  greedy_approach(Q, y, α,  BUDGET)
#
#
#
#     δ1 = convex_approach(Q, y, α,  BUDGET)
#     res = Q * (α .* (y + δ1))
#
#     (δ, total_cost, z_knap) = fractional_knapsack(Q, y, α,  BUDGET)
#
#     (δ, total_cost, z_page_rank) = select_using_pagerank(Q, y, α,  BUDGET, M)
#
#     @show norm(y)
#     @show norm(z)
#     @show norm(z_prime)
#     @show norm(z_knap)
#     @show norm(z_page_rank)
#     @show norm(res)
#
#     times = n * 2
#
#     trace = fill(WEIGHT(0.0), n, times)
#     trace[:,1] = y
#     for i = 2:times
#         next_opinion = get_next_opinion(α, β, γ,W, U, y, trace[:,i - 1])
#         trace[:,i] = next_opinion
#     end
#
#
#     # plot_multiple_series(trace)
#     # plot(1:times)
#     return (y, Q, α, δ, total_cost, z_prime, G, trace)
#     # return y, z, res
# end
#
# #y, z, res = main()
# (y, Q, α, δ, total_cost, z_prime, G, trace) = main()


"return `x+1` with probability `p` or `x-1` with `1-p`."
function update(x::Real,p::Real=0.5)
    ifelse(rand() < p, x+one(x), x-one(x))
end


#!/usr/bin/env julia

## brillouin.jl
#
#    Author: Leandro Seixas <seixasle@gmail.com>
#            Centre for Advanced 2D Materials (CA2DM),
#            National University of Singapore (NUS)
#            Singapore
#

using YAML
#using DataFrames
#using JulieTest
#using Gadfly

include("src/monkhorst-pack.jl")
include("src/reciprocal.jl")
include("src/firstBZ.jl")

lattice_constant = 2.47

bravais = [ sqrt(3)/2  1/2 0.0;
            sqrt(3)/2 -1/2 0.0;
            0.0  0.0 10.0 ]

bravais_lattice = lattice_constant*bravais

reciprocal_lattice = 2π*inv(bravais_lattice)'

G( n::Integer ) = reciprocal_lattice[n,:]

g_neighbors = collect( [ x*G(1)+y*G(2)+z*G(3) for x=-1:1:1, y=-1:1:1, z=-1:1:1 ] )

#for i in 1:length(kvec)
#  @printf("%14.8f%14.8f%14.8f\n", kvec[i]...)
#end


struct JointDistribution <: RandomUQInput
    marginals::Array{RandomVariable}
    copula::Copula

    function JointDistribution(marginals::Array{RandomVariable}, copula::Copula)
        length(marginals) == dimensions(copula) ||
            error("Dimension mismatch between copula and marginals")

        return new(marginals, copula)
    end
end

function sample(jd::JointDistribution, n::Integer=1)
    u = sample(jd.copula, n)

    samples = DataFrame()

    for (i, rv) in enumerate(jd.marginals)
        samples[!, rv.name] = quantile.(rv.dist, u[:, i])
    end

    return samples
end

function to_physical_space!(jd::JointDistribution, x::DataFrame)
    correlated_cdf = to_copula_space(jd.copula, Matrix{Float64}(x[:, names(jd)]))
    for (i, rv) in enumerate(jd.marginals)
        x[!, rv.name] = quantile.(rv.dist, correlated_cdf[:, i])
    end
    return nothing
end

function to_standard_normal_space!(jd::JointDistribution, x::DataFrame)
    for rv in jd.marginals
        x[!, rv.name] = cdf.(rv.dist, x[:, rv.name])
    end
    uncorrelated_stdnorm = to_standard_normal_space(
        jd.copula, Matrix{Float64}(x[:, names(jd)])
    )
    for (i, rv) in enumerate(jd.marginals)
        x[!, rv.name] = uncorrelated_stdnorm[:, i]
    end
    return nothing
end

names(jd::JointDistribution) = vec(map(x -> x.name, jd.marginals))

mean(jd::JointDistribution) = mean(jd.marginals)

dimensions(jd::JointDistribution) = dimensions(jd.copula)


type DiscreteMMSE
    rho # snr

    W
    bias

    DiscreteMMSE() = new()
end

function DiscreteMMSE(snr,get_cov;nSamples=1)
    est = DiscreteMMSE()
    est.rho = 10^(0.1*snr)

    est.W = Array{Any}(nSamples)
    est.bias = Array{Any}(nSamples)

    for i in 1:nSamples
        C = get_cov() # get random sample from cov. prior
        nAntennas   = size(C,1)
        est.W[i]    = C/(C + 1/est.rho * eye(nAntennas)) # MMSE filter at sample i
        est.bias[i] = real(logdet(eye(nAntennas) - est.W[i])) # bias term at sample i
    end

    return est
end

"""
    estimate(est::DiscreteMMSE, y)

For each batch `b`, estimate channels `h[:,t,b]` from
observations `y[:,1], ..., y[:,T,b]` by formula

h[:,t,b] = W_est[b] * y[:,t,b]

with

W_est[b] = sum_i{ weights_i[b] W_i }

and

weights_i[b] = softmax( rho*sum_t{ y[:,t,b]' W[i] * y[:,t,b] } + nCoherence*bias[i] )
"""
function estimate(est::DiscreteMMSE, y)
    nAntennas  = size(y,1)
    nCoherence = size(y,2)
    nBatches   = size(y,3);
    nSamples   = length(est.W);

    # calculate weights for all batches
    exps = zeros(nSamples,nBatches)
    for b in 1:nBatches, t in 1:nCoherence, i in 1:nSamples
        exps[i,b] += real(y[:,t,b]'*est.W[i]*y[:,t,b])[1] * est.rho + est.bias[i]
    end
    exps = exps .- maximum(exps,1)

    weights = exp.(exps)
    weights = weights./sum(weights,1)

    # Calculate estimated MMSE filters and channel estimates
    hest = zeros(nAntennas, nCoherence, nBatches) + 1im*zeros(nAntennas, nCoherence, nBatches)
    for b = 1:nBatches
        W_est = zeros(nAntennas, nAntennas)
        for i=1:nSamples
            W_est += est.W[i]*weights[i,b];
        end
        hest[:,:,b] = W_est*y[:,:,b]
    end

    return hest
end


#Written By Johnathan Bizzano

include("../NeuralNetTest.jl")

const sample_size = 100
const iterations = 100
const scalefactor = .8
const crossoverrate = .6
const population_size = 30

const IsVerbose = true
const InputType = Float64
const StorageType = Float64

f(x) = MathConstants.e ^ sin(x)

function collect_data()
    inputs = zeros(Double, sample_size)
    outputs = zeros(Double, sample_size)
    for i in 1:sample_size
        inputs[i] = i
        outputs[i] = f(i)
    end
    TrainAndTestSplit(Data.ReadArrayToDataSet(inputs, outputs); trainFrac=.8, sort_randomly=true)
end

function test()
    designer = NeuralNet.ModelDesigner(input_shape=(1), error=Error.Meanabserr())
    rng_init = Initializer.RNG(-1.0:.01:1.0)
    push!(designer, Layers.Dense(100, rng_init))
    push!(designer, Layers.Dense(100, rng_init))
    push!(designer, Layers.Dense(1, rng_init))

    neuralnet = designer()
    nettrainer = trainer(neuralnet, optimizer=Optimizer.de_rand_1_bin(; population_size = population_size,
                     iterations = iterations, scalefactor = scalefactor, crossoverrate = crossoverrate, IsVerbose = IsVerbose))

    data = collect_data()
    println("Pre Training Test: f(3.0) = $(f(3.0)). Neural Net:$(neuralnet([3.0])[1])")
    train!(nettrainer, data[1]; epochs=1, testset=data[2], IsVerbose = IsVerbose)
    println("After Training Test: f(3.0) = $(f(3.0)). Neural Net:$(neuralnet([3.0])[1])")
end

test()


# RSA protocol (RFC-3447) implementation in Julia
# Copyright (C): Jaka Smrekar (vinctux) <vinctux@outlook.com>, 2016
#
# PyCrypto (https://github.com/dlitz/pycrypto) as Pythonic reference

using Krypto
using SHA

abstract RSAKey

type RSAPubKey <: RSAKey
    n::BigInt             # RSA modulus
    e::BigInt             # RSA public component
    o::Array{UInt8, 1}    # Octet array representation of RSA key (n, e)
    hex::AbstractString   # Hex representation of RSA key (n, e)
end
RSAPubKey() = RSAPubKey(0, 65537, Array{UInt8, 1}(), "")

type RSAPrivKey <: RSAKey
    n::BigInt             # RSA modulus
    d::BigInt             # RSA private component
    o::Array{UInt8, 1}    # Octet array representation of RSA key (n, e)
    hex::AbstractString   # Hex representation of RSA key (n, e)
end
RSAPrivKey() = RSAPrivKey(0, 65537, Array{UInt8, 1}(), "")

# RFC-3447/4.1
function IP2OS{T<:Integer, L<:Integer}(x::T, xLen::L)
    #if x >= 256^BigInt(xLen) error("Integer too large.") end
    X = Array{UInt8, 1}()
    while x > 256 append!(X, x % 256); x = div(x, 256) end
    while length(X) < xLen append!(X, 0o0) end
    return reverse(X)
end
IP2OS{T<:Integer}(x::T) = IP2OS(x, div(bitsize(x), 8))

# RFC-3447/4.2
function OS2IP(X::Array{UInt8, 1})
    x = BigInt(0)
    for i in 1:length(X) x += X[i] * 256^BigInt(length(X) - i) end
    return x
end

function OS2STR(X::Array{UInt8, 1})
    r = ""
    for i in 1:length(X) r *= string(Char('\0' + d)) end
    return r
end

# RFC-3447/5.1.1
function RSAEP(pubkey::RSAPubKey, m::Integer)
    if m > pubkey.n - 1 || m < 0 error("Message representative out of range.") end
    return powermod(m, pubkey.e, pubkey.n)
end

# RFC-3447/5.1.2
function RSADP(privkey::RSAPrivKey, m::Integer)
    if m > privkey.n - 1 || m < 0 error("Message representative out of range.") end
    return powermod(m, privkey.d, privkey.n)
end

# RFC-3447/5.2.1
function RSASP1(privkey::RSAPrivKey, m::Integer)
    if m > privkey.n - 1 || m < 0 error("Message representative out of range.") end
    return powermod(m, privkey.d, privkey.n)
end

# RFC-3447/5.2.2
function RSAVP1(pubkey::RSAPubKey, m::Integer)
    if m > pubkey.n - 1 || m < 0 error("Signature representative out of range.") end
    return powermod(m, pubkey.e, pubkey.n)
end

# FIXME: RSAES_PKCS1_V1_5_ENCRYPT not working properly.
# RFC-3447/7.2.1
function RSAES_PKCS1_V1_5_ENCRYPT(pubkey::RSAPubKey, M::Array{UInt8, 1})
    k, mLen = length(IP2OS(pubkey.n)), length(M)
    mLen > k - 11 && error("Message too long.")
    PS = [csrand(0o1:0o255) for i in 1:(k - mLen - 3)]
    EM = vcat([0o0, 0o2], PS, [0o0], M)
    m, c, C = OS2IP(EM), RSAEP(pubkey, m), IP2OS(c, k)
    return C
end

# FIXME: RSAES_PKCS1_V1_5_DECRYPT not working properly.
# RFC-3447/7.2.2
function RSAES_PKCS1_V1_5_DECRYPT(privkey::RSAPrivKey, C::Array{UInt8, 1})
    k = length(IP2OS(privkey.n))
    length(C) != k && error("Decryption error.")
    EM = IP2OS(RSADP(privkey, OS2IP(C)), k)
    (EM[1] != 0o0 || EM[2] != 0o2) && error("Decryption error.")
    i, PS = 3, Array{UInt8, 1}()
    while EM[i] != 0o0 vcat(PS, EM[i]); i += 1 end
    i < 8 && error("Decryption error.")
    M = EM[(i + 1):end]
    return M
end

# FIXME: RSASSA_PKCS1_V1_5_SIGN not working properly.
# RFC-3447/8.2.1
function RSASSA_PKCS1_V1_5_SIGN(privkey::RSAPrivKey, M::Array{UInt8, 1})
    EM = EMSA_PKCS1_V1_5_ENCODE(M, length(IP2OS(privkey.n)))
    return IP2OS(RSASP1(privkey, OS2IP(EM)))
end

# FIXME: RSASSA_PKCS1_V1_5_VERIFY not working properly.
# RFC-3447/8.2.2
function RSASSA_PKCS1_V1_5_VERIFY(pubkey::RSAPubKey, M::Array{UInt8, 1}, S::Array{UInt8, 1})
    k = length(IP2OS(pubkey.n))
    length(S) != k && error("Invalid signature.")
    EM1, EM2 = IP2OS(RSAVP1(pubkey, OS2IP(S)), k), EMSA_PKCS1_V1_5_ENCODE(M, k)
    return EM1 == EM2
end

# TODO: Use DER/BER encoding standard before concating.
# RFC-3447/9.1
EMSA_PKCS1_V1_5_ENCODE{C<:SHA.SHA_CTX}(M::Array{UInt8, 1}, k::Integer, H::C = SHA.SHA3_CTX) =
    vcat([0o0, 0o1], [0o377 for i in 1:(k - length(M) - 3)], [0o0], update!(H, M))

function RSAKeyGen{T<:Integer}(b::T, e::T = 65537, newe::Bool = false, ERR::Int8 = Int8(75))
    if (b - 1024) % 256 != 0 error("Number of bits should be in form b = 1024 + 256x.") end
    p = q = BigInt(1)
    while bitsize(BigInt(p * q)) < b
        p = BigInt(csprime(div(b, 2), ERR))
        q = BigInt(csprime(div(b, 2), ERR))
    end
    N = BigInt(p * q)
    phi = BigInt(N - (p + q - 1))
    if newe
        e = BigInt(csrand(1:phi))
        while gcd(e, phi) != 1 e = csrand(1:phi) end
    end
    PUB  = RSAPubKey(N, e, Array{UInt8, 1}(), "")
    PRIV = RSAPrivKey(N, invmod(e, phi), Array{UInt8, 1}(), "")
    return PUB, PRIV
end

## SHORTHANDS ##
RSAEncrypt(pubkey::RSAPubKey,   M::Array{UInt8, 1}) = RSAES_PKCS1_V1_5_ENCRYPT(pubkey, M)
RSADecrypt(privkey::RSAPrivKey, C::Array{UInt8, 1}) = RSAES_PKCS1_V1_5_DECRYPT(privkey, C)
RSASign(privkey::RSAPrivKey,    M::Array{UInt8, 1}) = RSASSA_PKCS1_V1_5_SIGN(privkey, M)
RSAVerify(pubkey::RSAPubKey,    M::Array{UInt8, 1}, S::Array{UInt8, 1}) = RSASSA_PKCS1_V1_5_VERIFY(pubkey, M, S)

encrypt(::Type{Krypto.RSA}, K::RSAPubKey,  M::Array{UInt8, 1}) = RSAEncrypt(K, M)
decrypt(::Type{Krypto.RSA}, K::RSAPrivKey, C::Array{UInt8, 1}) = RSADecrypt(K, C)
sign(::Type{Krypto.RSA},    K::RSAPrivKey, M::Array{UInt8, 1}) = RSASign(K, M)
verify(::Type{Krypto.RSA},  K::RSAPubKey,  C::Array{UInt8, 1}) = RSAVerify(K, C)


using MathLink

e = mlinitialize()

l = mlopen(e)
mlactivate(l)

# skip initial input packet
if mlnextpacket(l) != MathLink.INPUTNAMEPKT
    error("Unexpected MathLink packet")
end
mlnewpacket(l)

# 1) log(2.0)
mlput(l,MLFunction("EvaluatePacket",1))
mlput(l,MLFunction("Log",1))
mlput(l,2.0)
mlendpacket(l)

# check for next packet
p = mlnextpacket(l)

t = mlgetnext(l) # a real
mlget(l,Float64)

# 2) using expressions
@mlput l EvaluatePacket(Log(2.0))
mlendpacket(l)

p = mlnextpacket(l)

t = mlgetnext(l) # a real
mlget(l,Float64)

@mlput l EvaluatePacket(Sin(:Pi))
mlendpacket(l)

p = mlnextpacket(l)

t = mlgetnext(l) # a real
mlget(l,Float64)

# 3) log(2)
mlput(l,MLFunction("EvaluatePacket",1))
mlput(l,MLFunction("Log",1))
mlput(l,2)
mlendpacket(l)

# check for next packet
p = mlnextpacket(l)

t = mlgetnext(l) # a function
mlgetfunction(l)

t = mlgetnext(l) # an integer
mlget(l,Int)

# 4) we can pass expressions
mlput(l,:(EvaluatePacket(Log(2))))
mlendpacket(l)

p = mlnextpacket(l)
# automatically figure out types
mlget(l)

# 5) or use strings representing Mathematica expressions
mlput(l,:(EvaluatePacket(ToExpression("Factorial[30]"))))
mlendpacket(l)

p = mlnextpacket(l)
# occasionally need to specify return types.
mlget(l,BigInt)

mlput(l,:(EvaluatePacket(Integrate(Abs(x), x, Assumptions => x > 0))))
mlendpacket(l)

p = mlnextpacket(l)
mlget(l)

bg = big(1e-30)
@mlput l EvaluatePacket(Log(CDF(GammaDistribution(1,1),bg)))
mlendpacket(l)

p = mlnextpacket(l)
mlget(l)

# done
mlclose(l)
mldeinitialize(e)


lines = readlines("input.txt")
octomap = parse.(Int, reduce(hcat, collect.(lines)))
idx = CartesianIndices(octomap)
minval = typemin(Int)

o = CartesianIndex(1,1)
V = CartesianIndices((-1:1, -1:1))
padded = fill(minval, size(octomap) .+ 2)
padded[idx .+ o] = octomap[idx]

for step = 1:1000
    padded .+= 1
    flashmap = fill(false, size(padded))
    while (yxs = findall(padded .>= 10)) |> length > 0
        flashmap[yxs] .= true
        padded[yxs] .= minval # Prevent any further flashing for this step
        [padded[yx .+ V] .+= 1 for yx = yxs]
    end
    padded[findall(flashmap)] .= 0

    if sum(flashmap) == 100
        println(step)
        break
    end
end


"""
LinearSystem + SingleIntegrator
"""
function LinearSystem_SingleIntegrator(args_linearsystem)
    linearsystem = LinearSystem(args_linearsystem...)
    integ = SingleIntegrator()
    linearsystem, integ
end

"""
# Notes
- `r` denotes integrand of `integ`.
"""
function State(linearsystem::LinearSystem, integ::SingleIntegrator)
    return function (x0_linearsystem, x0_integrator=[0.0])
        x_linearsystem = State(linearsystem)(x0_linearsystem)
        x_integrator = State(integ)(x0_integrator)
        ComponentArray(x=x_linearsystem, ∫r=x_integrator)
    end
end

"""
# Notes
running_cost: integrand; function of (x, u)
"""
function Dynamics!(linearsystem::LinearSystem, integ::SingleIntegrator, running_cost::Function)
    @Loggable function dynamics!(dX, X, p, t; u)
        @unpack x = X
        @nested_log :linearsystem Dynamics!(linearsystem)(dX.x, X.x, nothing, t; u=u)
        r = running_cost(x, u)
        @nested_log :integ Dynamics!(integ)(dX.∫r, X.∫r, nothing, t; u=r)
    end
end


using PyPlot
using LinearAlgebra
using PolynomialBasis
using ImplicitDomainQuadrature
using Revise
using CutCell
include("../../../test/useful_routines.jl")

function bulk_modulus(l, m)
    return l + 2m / 3
end

function lame_lambda(k, m)
    return k - 2m / 3
end

function displacement(alpha, x)
    u1 = alpha * x[2] * sin(pi * x[1])
    u2 = alpha * (x[1]^3 + cos(pi * x[2]))
    return [u1, u2]
end

function body_force(lambda, mu, alpha, x)
    b1 = alpha * (lambda + 2mu) * pi^2 * x[2] * sin(pi * x[1])
    b2 =
        -alpha * (6mu * x[1] + (lambda + mu) * pi * cos(pi * x[1])) +
        alpha * (lambda + 2mu) * pi^2 * cos(pi * x[2])
    return [b1, b2]
end

function stress_field(lambda, mu, alpha, x)
    s11 =
        (lambda + 2mu) * alpha * pi * x[2] * cos(pi * x[1]) -
        lambda * alpha * pi * sin(pi * x[2])
    s22 =
        -(lambda + 2mu) * alpha * pi * sin(pi * x[2]) +
        lambda * alpha * pi * x[2] * cos(pi * x[1])
    s12 = alpha * mu * (3x[1]^2 + sin(pi * x[1]))
    return [s11, s22, s12]
end

function onboundary(x, L, W)
    return x[2] ≈ 0.0 || x[1] ≈ L || x[2] ≈ W || x[1] ≈ 0.0
end

function solve_for_displacement(
    basis,
    cellquads,
    interfacequads,
    facequads,
    cutmesh,
    levelset,
    levelsetcoeffs,
    stiffness,
    penalty,
    alpha,
)

    L, W = CutCell.widths(cutmesh)
    lambda, mu = CutCell.lame_coefficients(stiffness, +1)

    bilinearforms = CutCell.BilinearForms(basis, cellquads, stiffness, cutmesh)
    interfacecondition = CutCell.coherent_interface_condition(
        basis,
        interfacequads,
        stiffness,
        cutmesh,
        penalty,
    )
    displacementbc = CutCell.DisplacementCondition(
        x -> displacement(alpha, x),
        basis,
        facequads,
        stiffness,
        cutmesh,
        x -> onboundary(x, L, W),
        penalty,
    )

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

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

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

    nodaldisplacement = matrix \ rhs

    return nodaldisplacement
end

function spatial_closest_points(refclosestpoints, refclosestcellids, mesh)
    dim, npts = size(refclosestpoints)
    spclosestpoints = zeros(dim, npts)
    for i = 1:npts
        cellmap = CutCell.cell_map(mesh, refclosestcellids[i])
        spclosestpoints[:, i] .= cellmap(refclosestpoints[:, i])
    end
    return spclosestpoints
end

function angular_position(points)
    cpoints = points[1, :] + im * points[2, :]
    return rad2deg.(angle.(cpoints))
end

function traction_force(stressvector, normal)
    return [
        stressvector[1] * normal[1] + stressvector[3] * normal[2],
        stressvector[3] * normal[1] + stressvector[2] * normal[2],
    ]
end

function traction_force_at_points(stresses, normals)
    npts = size(stresses)[2]
    @assert size(normals)[2] == npts

    return hcat(
        [traction_force(stresses[:, i], normals[:, i]) for i = 1:npts]...,
    )
end

K1, K2 = 247.0, 192.0    # Pa
mu1, mu2 = 126.0, 87.0   # Pa
lambda1 = lame_lambda(K1, mu1)
lambda2 = lame_lambda(K2, mu2)
stiffness = CutCell.HookeStiffness(lambda1, mu1, lambda2, mu2)

rho1 = 1.0
rho2 = 1.0
V1 = inv(rho1)
V2 = inv(rho2)

theta0 = 0.0
transfstress = CutCell.plane_strain_transformation_stress(lambda1, mu1, theta0)

width = 1.0
displacementscale = 0.01 * width
penaltyfactor = 1e1

polyorder = 3
numqp = required_quadrature_order(polyorder) + 2
nelmts = 9
center = [width / 2, width / 2]
inradius = width / 4

dx = width / nelmts
meanmoduli = 0.5 * (lambda1 + lambda2 + mu1 + mu2)
penalty = penaltyfactor / dx * meanmoduli

basis = TensorProductBasis(2, polyorder)
mesh = CutCell.Mesh([0.0, 0.0], [width, width], [nelmts, nelmts], basis)
levelset = InterpolatingPolynomial(1, basis)
levelsetcoeffs = CutCell.levelset_coefficients(
    x -> -circle_distance_function(x, center, inradius),
    mesh,
)

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

nodaldisplacement = solve_for_displacement(
    basis,
    cellquads,
    interfacequads,
    facequads,
    cutmesh,
    levelset,
    levelsetcoeffs,
    stiffness,
    penalty,
    displacementscale,
)

refseedpoints, spatialseedpoints, seedcellids =
    CutCell.seed_zero_levelset_with_interfacequads(interfacequads, cutmesh)
normals = CutCell.collect_interface_normals(interfacequads, cutmesh)

spatialpoints = spatialseedpoints
referencepoints = refseedpoints
referencecellids = seedcellids

relspatialpoints = spatialpoints .- center
angularposition = angular_position(relspatialpoints)
sortidx = sortperm(angularposition)

angularposition = angularposition[sortidx]
referencepoints = referencepoints[:, sortidx]
spatialpoints = spatialpoints[:, sortidx]
referencecellids = referencecellids[sortidx]
normals = normals[:, sortidx]

productdisplacement = CutCell.displacement_at_reference_points(
    referencepoints,
    referencecellids,
    +1,
    basis,
    nodaldisplacement,
    cutmesh,
)
parentdisplacement = CutCell.displacement_at_reference_points(
    referencepoints,
    referencecellids,
    -1,
    basis,
    nodaldisplacement,
    cutmesh,
)
exactdisplacement = hcat(
    [
        displacement(displacementscale, spatialpoints[:, i]) for
        i = 1:size(spatialpoints)[2]
    ]...,
)

# fig, ax = PyPlot.subplots(2, 1)
# ax[1].plot(angularposition, productdisplacement[1, :], label = "product u1")
# ax[1].plot(angularposition, parentdisplacement[1, :], label = "parent u1")
# ax[1].grid()
# ax[1].legend()
# ax[2].plot(angularposition, productdisplacement[2, :], label = "product u2")
# ax[2].plot(angularposition, parentdisplacement[2, :], label = "parent u2")
# ax[2].grid()
# ax[2].legend()
# fig.tight_layout()
# fig
#

productstress = CutCell.product_stress_at_reference_points(
    referencepoints,
    referencecellids,
    basis,
    stiffness,
    transfstress,
    theta0,
    nodaldisplacement,
    cutmesh,
)
parentstress = CutCell.parent_stress_at_reference_points(
    referencepoints,
    referencecellids,
    basis,
    stiffness,
    nodaldisplacement,
    cutmesh,
)

producttraction = traction_force_at_points(productstress,normals)
parenttraction = traction_force_at_points(parentstress,normals)

# fig, ax = PyPlot.subplots(2, 1)
# ax[1].plot(angularposition, producttraction[1, :], label = "product")
# ax[1].plot(angularposition, parenttraction[1, :], label = "parent")
# ax[1].set_title("t1")
# ax[1].legend()
# ax[1].grid()
# ax[2].plot(angularposition, producttraction[2, :], label = "product")
# ax[2].plot(angularposition, parenttraction[2, :], label = "parent")
# ax[2].set_title("t2")
# ax[2].legend()
# ax[2].grid()
# fig.tight_layout()
# fig

productpressure = CutCell.pressure_at_points(productstress)
parentpressure = CutCell.pressure_at_points(parentstress)

fig,ax = PyPlot.subplots()
ax.plot(angularposition,productpressure,label="product")
ax.plot(angularposition,parentpressure,label="parent")
ax.legend()
ax.grid()
ax.set_title("Pressure on interface")
fig

# productdevstress = CutCell.deviatoric_stress_at_points(productstress,productpressure)
# parentdevstress = CutCell.deviatoric_stress_at_points(parentstress,parentpressure)
#
# productdevnorm = CutCell.stress_inner_product_over_points(productdevstress)
# parentdevnorm = CutCell.stress_inner_product_over_points(parentdevstress)
#
# productnormaldevstress = CutCell.normal_stress_component_over_points(productdevstress,normals)
# parentnormaldevstress = CutCell.normal_stress_component_over_points(parentdevstress,normals)
#
# productspecificvolume = V1*(1.0 .- productpressure/K1)
# parentspecificvolume = V2*(1.0 .- parentpressure/K2)
#
#
# productp1 = productspecificvolume .* productpressure
# productp2 = V1/(2K1)*(productpressure.^2)
# productp3 = -(productspecificvolume .* productnormaldevstress)
# productp4 = V1/(4mu1)*productdevnorm
# productpotential = productp1+productp2+productp3+productp4
#
# parentp1 = parentspecificvolume .* parentpressure
# parentp2 = V2/(2K2)*(parentpressure.^2)
# parentp3 = -(parentspecificvolume .* parentnormaldevstress)
# parentp4 = V2/(4mu2)*parentdevnorm
# parentpotential = parentp1+parentp2+parentp3+parentp4
#
# potentialdifference = productpotential - parentpotential


using PyPlot

using SubmodularMaximization

pygui(false)
fig_path = "./fig/coverage_test"
mkpath(fig_path)

num_agents = 5
num_sensors = 5
nominal_area = 2.0

sensor_radius = sqrt(nominal_area / (num_agents * pi))
station_radius = sensor_radius

agent_specification = CircleAgentSpecification(sensor_radius, station_radius,
                                         num_sensors)

agents = generate_agents(agent_specification, num_agents)

f(x) = mean_area_coverage(x, 100)
problem = ExplicitPartitionProblem(f, agents)

function evaluate_solver(solver, name)
  println("$name solver running")
  @time solution = solver(problem)

  figure()
  xlim([0, 1])
  ylim([0, 1])
  colors = generate_colors(agents)
  visualize_agents(agents, colors)
  visualize_solution(problem, solution, colors)

  gca().set_aspect("equal")
  savefig("$(fig_path)/$(to_file(name)).png", pad_inches=0.00, bbox_inches="tight")

  coverage = solution.value

  title("$name Solver ($coverage)")

  @show coverage
end

#evaluate_solver(solve_optimal, "Optimal")
#evaluate_solver(solve_worst, "Worst-case")
#evaluate_solver(solve_myopic, "Myopic")
#evaluate_solver(solve_random, "Random")
evaluate_solver(solve_sequential, "Sequential")

for num_partitions in [2, 4, 8]
  solve_n(p) = solve_n_partitions(num_partitions, p)
  evaluate_solver(solve_n, "Partition-$num_partitions")
end

#@show mean_weight(problem)
#@show total_weight(problem)


using GeneralizedDynamicsFromData
using Plots
using HDF5
using OrderedCollections
using Dates

experiment_name = "selkov_grid_$(Dates.month(today()))_$(Dates.day(today()))"

reps = 100
neurons = [2^i for i in 1:10]
layers = [i for i in 1:10]

grid_config = [OrderedDict([
    :inputs => 2,
    :outputs => 1,
    :neurons => N,
    :layers => L,
    :non_lin => tanh,
    :initialization => Flux.glorot_normal
]) for (N, L) in Iterators.product(neurons, layers)]

problem = Dict([
    :equation => selkov,
    :parameters => [0.1 0.6],
    :u0 => Float32[1.0, 1.0],
    :tE => 50.0f0,
    :tspan => (0.0f0, 50.0f0),
    :ts => 0.1,
    :solver => Vern7,
    :loss => mse_loss
])

result = grid_experiment(problem, grid_config, reps)


module LTVModels
using LinearAlgebra, Statistics, Printf, Random, Distributions
using LinearTimeVaryingModelsBase, ControlSystemIdentification
import LinearTimeVaryingModelsBase: AbstractModel, AbstractCost, ModelAndCost, f,
dc,calculate_cost,calculate_final_cost,
predict, simulate, df,costfun, LTVStateSpaceModel,
SimpleLTVModel, rms, sse, nrmse, modelfit, aic

import ControlSystemIdentification: OutputData, InputOutputData, InputOutputStateData, AnyInput, AbstractIdData, oftype, time1, input, output, state
import ControlSystems: ninputs, noutputs, nstates

export predict, simulate, iddata, input, output, state, ninputs, noutputs, nstates

export SimpleLTVModel, KalmanModel, LTVAutoRegressive, rootspectrogram

using ControlSystems, DSP, Plots
using Test

using DiffResults, Optim
using ReverseDiff: GradientTape, GradientConfig, gradient!

"""
Struct representing a statespace model fit using a Kalman smoother
# Fields
- `At, Bt, Pt` all of type `Array{T,3}` with time in the last dimension
- `extended::Bool` indicates whether or not the model has been extended by one time step to match the length of the input data.
"""
mutable struct KalmanModel{T} <: LTVStateSpaceModel
    At::Array{T,3}
    Bt::Array{T,3}
    Pt::Array{T,3}
    extended::Bool
    ll::Float64
end
function KalmanModel(At::Array{T,3},Bt::Array{T,3},Pt::Array{T,3},extend::Bool) where T
    if extend
        At = cat(At,At[:,:,end], dims=3)
        Bt = cat(Bt,Bt[:,:,end], dims=3)
        Pt = cat(Pt,Pt[:,:,end], dims=3)
    end
    return KalmanModel{T}(At,Bt,Pt,extend,0.0)
end

KalmanModel(At,Bt,Pt,extend::Bool=false) = KalmanModel{eltype(At)}(At,Bt,Pt,extend,0.0)

mutable struct GMMModel <: AbstractModel
    M
    dynamics
    T
end

mutable struct LTVAutoRegressive{T,PT<:Union{Array{T,3},Nothing}} <: LinearTimeVaryingModelsBase.LTVModel
    na::Int
    θ::Array{T,2}
    Pt::PT
    extended::Bool
    ll::Float64
end

Base.length(m::LTVAutoRegressive) = size(m.θ,2)

include("utilities.jl")
include("peakdetection.jl")
include("statespace_fit.jl")
include("seg_bellman.jl")
include("kalmanmodel.jl")
# include("gmmmodel.jl") # Not yet working with julia v0.6
include("statespace_utils.jl")
include("wrappers.jl")
end # module


import MitosisStochasticDiffEq as MSDE
using Mitosis, Statistics
using LinearAlgebra, Test
using Random
Random.seed!(126)
using Mitosis: AffineMap
atol = 0.015
# set true model parameters
par = [-0.2, 0.1, 0.9]

# Samples
K = 5000

# define SDE function
f(u,p,t) = p[1]*u .+ p[2]
g(u,p,t) = p[3] .+ 0u # needs to preserve type of u for forwardguided

# time span
tstart = 0.0
tend = 1.0
dt = 0.01
trange = tstart:dt:tend

# initial condition
u0 = 1.1

# set of linear parameters Eq.~(2.2)
plin = copy(par)
sdekernel = MSDE.SDEKernel(f,g,trange,plin)

plin2 = par2 = [-0.1, -0.05, 1.0]
sdekernel2 = MSDE.SDEKernel(f,g,trange,plin)
sdekerneltilde2 = MSDE.SDEKernel(f,g,trange,plin2)

y_, (ts, sol, _)  = MSDE.sample(sdekernel, u0, save_noise=true)

y_ = 1.39
y = vcat(y_)
samples_ = MSDE.sample(sdekernel, u0, K, save_noise=true)

samples = vcat.(samples_)

# Compute transition density. See Ludvig Arnold (ISBN 9780486482361 ), also Proposition 3.5. in [1]
B = fill(par[1], 1, 1)
Φ = exp(B*(tend - tstart))
β = (Φ - I)*inv(B)*fill(par[2], 1)
Σ = fill(par[3]^2, 1, 1)
Λ = lyap(B, Σ)
Q = Λ - Φ*Λ*Φ'

gkernel = Mitosis.kernel(Gaussian; μ=AffineMap(Φ, β), Σ=ConstantMap(Q))

B̃ = fill(par2[1], 1, 1)
Φ̃ = exp(B̃*(tend - tstart))
β̃ = (Φ̃ - I)*inv(B̃)*fill(par2[2], 1)
Σ̃ = fill(par2[3]^2, 1, 1)
Λ̃ = lyap(B̃, Σ̃)
Q̃ = Λ̃ - Φ̃*Λ̃*Φ̃'

gkernel2 = Mitosis.kernel(Gaussian; μ=AffineMap(Φ̃, β̃), Σ=ConstantMap(Q̃))

p = gkernel([u0])

p̂ = Gaussian{(:μ, :Σ)}(mean(samples), cov(samples))
@test norm(p.μ - p̂.μ) < 5/sqrt(K)
@test norm(p.Σ - p̂.Σ) < 5/sqrt(K)

m, l = Mitosis.backwardfilter(gkernel, y) # returns a "Leaf"
p = l.y
# initial values for ODE
WG = WGaussian{(:μ,:Σ,:c)}(y_, 0.0, 0.0) #

message, solend = MSDE.backwardfilter(sdekernel, WG)
@testset "Mitosis backward" begin
    @test (p.c)[] ≈ solend.c
    @test (p.Γ\p.F)[] ≈ solend.μ atol=atol
    @test inv(p.Γ)[] ≈ solend.Σ atol=atol
end

# Try forward
V = WGaussian{(:F, :Γ, :c)}(y*5.0, 5.0I, 0.0)
m, p = Mitosis.backwardfilter(gkernel, V)
kᵒ = Mitosis.left′(BFFG(), gkernel, m)

pT = kᵒ([u0])

# initial values for ODE
WG = WGaussian{(:μ,:Σ,:c)}(y_, 0.2, 0.0) # start with observation μ=y with uncertainty Σ=0.1
message, solend = MSDE.backwardfilter(sdekernel, WG)

(uend1, ll), (ts1, u1, noise1) = MSDE.forwardguiding(sdekernel, message, (u0, 0.0), Z=nothing; save_noise=true)

@test ll == 0

samples = [MSDE.forwardguiding(sdekernel, message, (u0, 0.0), Z=nothing; save_noise=true)[1][1] for k in 1:K]

@testset "Mitosis forward" begin
    @test pT.μ ≈ mean(samples) atol=atol
    @test pT.Σ ≈ cov(samples) atol=atol
end

# try tilted forward

m, p2 = Mitosis.backwardfilter(gkernel2, V)
gᵒ = Mitosis.forward(BFFG(), gkernel, m, [u0])

message, solend = MSDE.backwardfilter(sdekerneltilde2, WG)
@testset "Mitosis backward tilted" begin
    @test (p2.c)[] ≈ solend.c
    @test (p2.Γ\p2.F)[] ≈ solend.μ atol=atol
    @test inv(p2.Γ)[] ≈ solend.Σ atol=atol
end

(uend1, ll), (ts1, u1, noise1) = MSDE.forwardguiding(sdekernel2, message, (u0, 0.0))

we((x,c)) = x*exp(c)
samples2 = []
for k=1:K
    (uend1, ll), _ = MSDE.forwardguiding(sdekernel2, message, (u0, 0.0))
    push!(samples2,[uend1, ll])
end

samples = we.(samples2)

@show std(samples)
ptrue = Mitosis.density(p, [u0])
p̃ = Mitosis.density(WGaussian{(:F,:Γ,:c)}(solend.Σ\solend.μ, inv(solend.Σ), solend.c), u0)
@testset "Mitosis tilted forward" begin
    @test pT.μ ≈ mean(samples)*p̃/ptrue atol=atol
    @test gᵒ.μ ≈ mean(first.(samples2)) atol=2atol
    @test gᵒ.Σ ≈ cov(first.(samples2)) atol=2atol
    @test gᵒ.c ≈ mean(last.(samples2)) atol=atol
end


# # Gross-Pitaevskii equation with magnetism

# We solve the 2D Gross-Pitaevskii equation with a magnetic field.
# This is similar to the
# previous example ([Gross-Pitaevskii equation in one dimension](@ref)),
# but with an extra term for the magnetic field.

using DFTK
using StaticArrays
using Plots

# Unit cell. Having one of the lattice vectors as zero means a 2D system
a = 10
lattice = a .* [[1 0 0.]; [0 1 0]; [0 0 0]];

# Confining scalar potential, and magnetic vector potential
pot(x, y, z) = (x - a/2)^2 + (y - a/2)^2
Apot(x, y, z) = .2 * @SVector [y - a/2, -(x - a/2), 0]
Apot(X) = Apot(X...);

# Parameters
Ecut = 20  # Increase this for production
C = 500.0
α = 2
n_electrons = 1;  # Increase this for fun

# Collect all the terms, build and run the model
terms = [Kinetic(),
         ExternalFromReal(X -> pot(X...)),
         PowerNonlinearity(C, α),
         Magnetic(Apot),
]
model = Model(lattice; n_electrons=n_electrons,
              terms=terms, spin_polarization=:spinless)  # "spinless electrons"
basis = PlaneWaveBasis(model, Ecut)
scfres = direct_minimization(basis, tol=1e-5)  # Reduce tol for production
heatmap(scfres.ρ.real[:, :, 1], c=:blues)


using LinearAlgebraicRepresentation
Lar = LinearAlgebraicRepresentation
using ViewerGL
GL = ViewerGL

VOID = [[]], [[1]]

# example 1

V = [[0,0] [1,0] [2,0] [0,1] [1,1] [2,1] [0,2] [1,2] [2,2]];
FV = [[1,2,4],[2,3,5],[3,5,6],[4,5,7],[5,7,8],[6,8,9]];
pattern = repeat([1,2,-3],outer=4);
model = (V,FV);
W,FW = Lar.extrudeSimplicial(model, pattern);
GL.VIEW([ GL.GLGrid(W,FW, GL.COLORS[7], 0.5) ]);

# example 2

model = Lar.extrudeSimplicial( VOID, ones(10) )
GL.VIEW([ GL.GLGrid(model..., GL.COLORS[7], 0.5) ]); #TODO: check

model = Lar.extrudeSimplicial( model, ones(10) )
GL.VIEW([ GL.GLLines(model..., GL.COLORS[7]),
        GL.GLFrame ]);

model = Lar.extrudeSimplicial( model, ones(10) )
GL.VIEW([ GL.GLGrid(model..., GL.COLORS[7], 0.5) ]);

# example 3

model = Lar.extrudeSimplicial( VOID, repeat([1,-1],outer=10) )
GL.VIEW([ GL.GLGrid(model..., GL.COLORS[7], 0.5) ]); #TODO: check

model = Lar.extrudeSimplicial( model, repeat([1,-1],outer=10) )
GL.VIEW([ GL.GLGrid(model..., GL.COLORS[7],0.7),
        GL.GLFrame ]);

# example 4

grid_2d = Lar.simplexGrid([3,3])
GL.VIEW([ GL.GLGrid(grid_2d..., GL.COLORS[7],0.7),
        GL.GLLines(grid_2d...),    # TODO:  check lines
        GL.GLFrame ]);

grid_3d = Lar.simplexGrid([2,3,4])
GL.VIEW([ GL.GLGrid(grid_3d..., GL.COLORS[7]),
        GL.GLFrame ]);  # TODO  check


function getallergens(lines::Vector{<:AbstractString})
    ingredients = Vector{Vector{String}}()
    allergens = Dict{String,Vector{String}}()
    for line in lines
        ingr, aller = split(line[1:end - 1], " (contains ")
        ingr_list = split(ingr, " ")
        push!(ingredients, ingr_list)
        aller_list = split(aller, ", ")
        for al in aller_list
            allergens[al] = intersect(get(allergens, al, ingr_list), ingr_list)
        end
    end
    return allergens, ingredients
end

function countsafe(
    allergens::Dict{String,Vector{String}},
    ingredients::Vector{Vector{String}})
    unsafe = union(values(allergens)...)
    all_ingredients = union(ingredients...)
    safe = setdiff(all_ingredients, unsafe)
    total = 0
    for ingr_list in ingredients
        total += count(x -> x ∈ safe, ingr_list)
    end
    return total
end

const data = readlines("day21.txt")
const allergens, ingredients = getallergens(data)

println("Part 1: $(countsafe(allergens, ingredients))")

function next_identified(allergens::Dict{String,Vector{String}})
    sure = Tuple{String,String}[]
    for (aller, ing) in filter(x -> length(x[2]) == 1, collect(allergens))
        push!(sure, (aller, ing[1]))
        for (k, v) in allergens
            allergens[k] = setdiff(v, ing)
        end
    end
    return sure
end

function findmap(allergens::Dict{String,Vector{String}})
    known = []
    m = next_identified(allergens)
    while !(isempty(m))
        push!(known, m)
        m = next_identified(allergens)
    end
    return collect(Iterators.flatten(known))
end

function part2(allergens)
    known = findmap(allergens)
    return join((x[2] for x in sort(known)), ",")
end

println("Part 2: $(part2(allergens))")


"""
Compute preimages of monotonic sequences
"""

using IntervalArithmetic
using .Contractors

"""
Type used to represent a "branch" of a dynamic. The branch is represented by a monotonic map `f` with domain `X=(a,b)` with a≤b (where typically a,b are intervals).
`Y=(f(a),f(b))` and `increasing` may be provided (for instance if we know that `Y=(0,1)`), otherwise they are computed automatically.
"""
struct Branch{T,S}
    f::T
    X::Tuple{S, S}
    Y::Tuple{S, S}
    increasing::Bool
end
Branch(f, X, Y=(f(Interval(X[1])), f(Interval(X[2]))), increasing=unique_increasing(Y[1], Y[2])) = Branch{typeof(f), typeof(interval(X[1]))}(f, X, Y, increasing)

"""
Return Branches for a given PwMap, in an iterable.

TODO: in future, maybe it is a better idea to replace the type PwMap directly with an array of branches, since that's all we need
"""
function branches(D::PwMap)
    return [Branch(D.Ts[k], (D.endpoints[k], D.endpoints[k+1]), (D.y_endpoints[k,1], D.y_endpoints[k,2]), D.increasing[k]) for k in 1:length(D.Ts)]
end

"""
Smallest possible i such that a is in the semi-open interval [y[i], y[i+1]).

This should work properly even if `a, y` are intervals; in this case it returns the *smallest* possible value of i over all possible "assignments" of a, y inside those intervals.
Assumes y is sorted, i.e., map(y, x->Interval(x).lo) and map(y, x->Interval(x).hi) are sorted.
"""
function first_overlapping(y, a)
    if iszero(a) # avoids -0 crap
        a = zero(a)
    end
    searchsortedlast(y, Interval(a).lo, by=x->Interval(x).hi)
end

"""
Largest possible j such that a-ε is in the semi-open interval [y[j], y[j+1]).

This should work properly even if `a, y` are intervals; in this case it returns the *largest* possible value of i over all possible "assignments" of a, y inside those intervals.
Assumes y is sorted, i.e., map(y, x->Interval(x).lo) and map(y, x->Interval(x).hi) are sorted.
"""
function last_overlapping(y, a)
    if iszero(a) # avoids -0 crap
        a = zero(a)
    end
    searchsortedfirst(y, Interval(a).hi, by=x->Interval(x).lo) - 1
end

"""
Construct preimages of an increasing array y under a monotonic branch defined on X = (a, b), propagating additional labels `ylabel`

The sequence y subdivides the y-axis into semi-open intervals [y[l], y[l+1]); each of them is identified by the label `ylabel[l]`. We construct an increasing sequence
x that splits X (in the x-axis) into semi-open intervals, each of them with f([x[k], x[k+1]) ⊂ [y[l], y[l+1]) for a certain l.
We set xlabel[k] = ylabel[l], and return the pair (x, xlabel).

In the simplest case where D is full-branch, the points in x are preimages of the points in y, but in the general case they can also include D.endpoints:
in general, there may be a certain number of points in y that have no preimage at the beginning and the end of the sequence, because
they fall out of the range R = [f(a), f(b)]. In the worst case, no point has a preimage, because y[i] < R < y[i+1] for some
i (or vice versa with orientations), and in this case we just return the 1-element vectors x = [branch.X[1]] and xlabel = [i].

x[begin] always coincides with branch.X[1], while branch.X[2] is "the point after x[end]", and is not stored explicitly in x, for easier composing.
In this way x and xlabel have the same length.

This function fills the array by using a bisection strategy to save computations: if y ∈ [a,b], then f⁻¹(y) ∈ [f⁻¹(a),f⁻¹(b)] (paying attention to orientation).
So we can fill v by filling in first entries `v[k+1]` with higher dyadic valuation of k.

For a dynamic with multiple branches, preimages(y, D) is simply the concatenation of x, xlabel for b in all branches. These values still form an increasing sequence that
splits X into intervals, each of which is mapped into a different semi-open interval [y[k], y[k+1]).
"""
function preimages(y, br::Branch, ylabel = 1:length(y), ϵ = 0.0)

    if br.increasing
        i = first_overlapping(y, br.Y[1])  # smallest possible i such that a is in the semi-open interval [y[i], y[i+1]).
        j = last_overlapping(y, br.Y[2]) # largest possible j such that b-ε is in the semi-open interval [y[j], y[j+1]).
        n = j - i + 1
        x = fill((-∞..∞)::typeof(Interval(br.X[1])), n)
        xlabel = collect(ylabel[i:j]) # we collect to avoid potential type instability, since this may be an UnitRange while in the other branch we could have a StepRange
        x[1] = br.X[1]
        if n == 1
            return (x, xlabel)
        end
        # the bisection strategy: fill the array in "strides" of length `stride`, then halve the stride and repeat
        # for instance, if the array is 1..13 (with x[1] filled in already), we first take stride=8 and fill in x[9],
        # then stride=4 and fill in x[5], x[13], (note that the distance is `2stride`, since x[9], and in general all the even multiples of `stride`, is already filled in)
        # then stride=2 and fill in x[3], x[7], x[11],
        # then stride=1 and fill in x[2], x[4], x[6], x[8], x[10], x[12]
        # at each step we have bracketed the preimage in a "search range" given by already-computed preimages x[k-stride] and x[k+stride].
        stride = prevpow(2, n-1)
        while stride >= 1
            # fill in v[i] using x[i-stride].lo and x[i+stride].hi as range for the preimage search
            for k = 1+stride:2*stride:n
                search_range = Interval(x[k-stride].lo, (k+stride <= n ? x[k+stride] : Interval(br.X[2])).hi)
                x[k] = preimage(y[i-1+k], br.f, search_range, ϵ)
            end
            stride = stride ÷ 2
        end
    else # branch decreasing
        i = last_overlapping(y, br.Y[1]) # largest possible j such that b-ε is in the semi-open interval [y[j], y[j+1]).
        j = first_overlapping(y, br.Y[2]) # smallest possible i such that a is in the semi-open interval [y[i], y[i+1]).
        n = i - j + 1
        x = fill((-∞..∞)::typeof(Interval(br.X[1])), n)
        xlabel = collect(ylabel[i:-1:j])
        x[1] = br.X[1]
        if n == 1
            return (x, xlabel)
        end
        stride = prevpow(2, n-1)
        while stride >= 1
            # fill in v[i] using x[i-stride].lo and x[i+stride].hi as range for the preimage search
            for k = 1+stride:2*stride:n
                search_range = Interval(x[k-stride].lo, (k+stride <= n ? x[k+stride] : Interval(br.X[2])).hi)
                x[k] = preimage(y[i+2-k], br.f, search_range, ϵ)
            end
            stride = stride ÷ 2
        end
    end
    return (x, xlabel)
end

function preimages(y, D::Dynamic, ylabel = 1:length(y), ϵ = 0.0)
    results = collect(preimages(y, b, ylabel, ϵ) for b in branches(D))
    x = vcat((result[1] for result in results)...)
    xlabel = vcat((result[2] for result in results)...)
    return x, xlabel
end

"""
Compute preimages of D *and* the derivatives f'(x) in each point.

Returns: x, xlabel, x′

Assumes that the dynamic is full-branch, because otherwise things may compose the wrong way.
This is not restrictive because we'll need it only for the Hat assembler (at the moment)

We combine them in a single function because there are avenues to optimize by recycling some computations (not completely exploited for now)
"""
function preimages_and_derivatives(y, br::Branch, ylabel = 1:length(y), ϵ = 0.0)
    x, xlabel = preimages(y, br, ylabel, ϵ)
    f′ = Contractors.derivative(br.f)
    x′ = f′.(x)
    return x, xlabel, x′
end
function preimages_and_derivatives(y, D::Dynamic, ylabel = 1:length(y), ϵ = 0.0)
    @assert is_full_branch(D)
    results = collect(preimages_and_derivatives(y, b, ylabel, ϵ) for b in branches(D))
    x = vcat((result[1] for result in results)...)
    xlabel = vcat((result[2] for result in results)...)
    x′ = vcat((result[3] for result in results)...)
    return x, xlabel, x′
end

"""
Composed map D1 ∘ D2 ∘ D3. We store with [D1, D2, D3] in this order.

We overwrite ∘ in base, so one can simply write D1 ∘ D2 or ∘(D1, D2, D3) to construct them.
"""
struct ComposedDynamic <: Dynamic
    dyns::Tuple{Vararg{Dynamic}}
end
Base.:∘(d::Dynamic...) = ComposedDynamic(d)

"""
Utility function to return the domain of a dynamic
"""
domain(D::PwMap) = (D.endpoints[begin], D.endpoints[end])
domain(D::ComposedDynamic) = domain(D.dyns[end])

function preimages(z, Ds::ComposedDynamic, zlabel = 1:length(z), ϵ = 0.0)
    for d in Ds.dyns
        z, zlabel = preimages(z, d, zlabel, ϵ)
    end
    return z, zlabel
end
function preimages_and_derivatives(z, Ds::ComposedDynamic, zlabel = 1:length(z), ϵ = 0.0)
    derivatives = fill(1, length(z))
    for d in Ds.dyns
        z, zindex, z′ = preimages_and_derivatives(z, d, 1:length(z), ϵ)
        zlabel = zlabel[zindex]
        derivatives = derivatives[zindex] .* z′
    end
    return z, zlabel, derivatives
end

"""
Replacement of DualComposedWithDynamic.
"""
abstract type Dual end

struct UlamDual <: Dual
    x::Vector{Interval} #TODO: a more generic type may be needed in future
    xlabel::Vector{Int}
    lastpoint::Interval
end
Dual(B::Ulam, D, ϵ) = UlamDual(preimages(B.p, D, 1:length(B.p)-1, ϵ)..., domain(D)[end])

Base.length(dual::UlamDual) = length(dual.x)
Base.eltype(dual::UlamDual) = Tuple{eltype(dual.xlabel), Tuple{eltype(dual.x), eltype(dual.x)}}
function iterate(dual::UlamDual, state = 1)
    n = length(dual.x)
    if state < n
        return (dual.xlabel[state], (dual.x[state], dual.x[state+1])), state+1
    elseif state == n
        return (dual.xlabel[n], (dual.x[n], dual.lastpoint)), state+1
    else
        return nothing
    end
end

struct HatDual <: Dual
    x::Vector{Interval} #TODO: a more generic type may be needed in future
    xlabel::Vector{Int}
    x′::Vector{Interval}
end

Dual(B::Hat, D, ϵ) = HatDual(preimages_and_derivatives(B.p, D, 1:length(B.p)-1, ϵ)...)
Base.length(dual::HatDual) = length(dual.x)
Base.eltype(dual::HatDual) = Tuple{eltype(dual.xlabel), Tuple{eltype(dual.x), eltype(dual.x′)}}
function iterate(dual::HatDual, state=1)
    if state <= length(dual.x)
        return ((dual.xlabel[state], (dual.x[state], abs(dual.x′[state]))), state+1)
    else
        return nothing
    end
end

# Variants of assemble and DiscretizedOperator; the code is repeated here for easier comparison with the older algorithm
function assemble2(B, D, ϵ=0.0; T = Float64)
	I = Int64[]
	J = Int64[]
	nzvals = Interval{T}[]
	n = length(B)

	# TODO: reasonable size hint?

	for (i, dual_element) in Dual(B, D, ϵ)
		if !is_dual_element_empty(B, dual_element)
			for (j, x) in ProjectDualElement(B, dual_element)
				push!(I, i)
				push!(J, mod(j,1:n))
				push!(nzvals, x)
			end
		end
	end

	return sparse(I, J, nzvals, n, n)
end

function DiscretizedOperator2(B, D, ϵ=0.0; T = Float64)
	L = assemble2(B, D, ϵ; T)
	if is_integral_preserving(B)
		return IntegralPreservingDiscretizedOperator(L)
	else
		f = integral_covector(B)
		e = one_vector(B)
		w = f - f*L #will use interval arithmetic when L is an interval matrix
		return NonIntegralPreservingDiscretizedOperator(L, e, w)
	end
end


# For CI, make sure the downloads do not hang
ENV["DATADEPS_ALWAYS_ACCEPT"] = true

using Test

@testset "quasi-Rosenbrock derivative" begin
    include("rosenbrock.jl")
end

@testset "AIBECS test" begin
    include("AIBECS_test.jl")
end


include("utilities.jl")

using SEAL
using Printf

"""
    example_bfv_basics()

Perform some basic operations such as encryption/decryption, multiplication, addition etc. using the
BFV scheme. This routine is based on the file `native/examples/1_bfv_basics.cpp` of the original
SEAL library and should yield the exact same output.

* [SEAL](https://github.com/microsoft/SEAL)
* [native/examples/1_bfv_basics.cpp](https://github.com/microsoft/SEAL/blob/master/native/examples/1_bfv_basics.cpp)

See also: [`example_ckks_basics`](@ref)
"""
function example_bfv_basics()
  print_example_banner("Example: BFV Basics")

  parms = EncryptionParameters(SchemeType.bfv)

  poly_modulus_degree = 4096
  set_poly_modulus_degree!(parms, poly_modulus_degree)

  set_coeff_modulus!(parms, coeff_modulus_bfv_default(poly_modulus_degree))

  set_plain_modulus!(parms, 1024)

  context = SEALContext(parms)

  print_line(@__LINE__)
  println("Set encryption parameters and print")
  print_parameters(context)

  println("Parameter validation (success): ", parameter_error_message(context))

  println()
  println("~~~~~~ A naive way to calculate 4(x^2+1)(x+1)^2. ~~~~~~")

  keygen = KeyGenerator(context)
  secret_key_ = secret_key(keygen)
  public_key_ = PublicKey()
  create_public_key!(public_key_, keygen)

  encryptor = Encryptor(context, public_key_)

  evaluator = Evaluator(context)

  decryptor = Decryptor(context, secret_key_)

  print_line(@__LINE__)
  x = 6
  x_plain = Plaintext(string(x))
  println("Express x = " * string(x) * " as a plaintext polynomial 0x" * to_string(x_plain) * ".")

  print_line(@__LINE__)
  x_encrypted = Ciphertext()
  println("Encrypt x_plain to x_encrypted.")
  encrypt!(x_encrypted, x_plain, encryptor)

  println("    + size of freshly encrypted x: ", length(x_encrypted))

  println("    + noise budget in freshly encrypted x: ",
          invariant_noise_budget(x_encrypted, decryptor),
          " bits")

  x_decrypted = Plaintext()
  decrypt!(x_decrypted, x_encrypted, decryptor)
  println("    + decryption of x_encrypted: 0x", to_string(x_decrypted), " ...... Correct.")

  print_line(@__LINE__)
  println("Compute x_sq_plus_one (x^2+1).")
  x_sq_plus_one = Ciphertext()
  square!(x_sq_plus_one, x_encrypted, evaluator)
  plain_one = Plaintext("1")
  add_plain_inplace!(x_sq_plus_one, plain_one, evaluator)

  println("    + size of x_sq_plus_one: ", length(x_sq_plus_one))
  println("    + noise budget in x_sq_plus_one: ",
          invariant_noise_budget(x_sq_plus_one, decryptor),
          " bits")

  decrypted_result = Plaintext()
  decrypt!(decrypted_result, x_sq_plus_one, decryptor)
  println("    + decryption of x_sq_plus_one: 0x", to_string(decrypted_result), " ...... Correct.")

  print_line(@__LINE__)
  println("Compute x_plus_one_sq ((x+1)^2).")
  x_plus_one_sq = Ciphertext()
  add_plain!(x_plus_one_sq, x_encrypted, plain_one, evaluator)
  square_inplace!(x_plus_one_sq, evaluator)
  println("    + size of x_plus_one_sq: ", length(x_plus_one_sq))
  println("    + noise budget in x_plus_one_sq: ",
          invariant_noise_budget(x_plus_one_sq, decryptor),
          " bits")
  decrypt!(decrypted_result, x_plus_one_sq, decryptor)
  println("    + decryption of x_plus_one_sq: 0x", to_string(decrypted_result),
          " ...... Correct.")

  print_line(@__LINE__)
  println("Compute encrypted_result (4(x^2+1)(x+1)^2).")
  encrypted_result = Ciphertext()
  plain_four = Plaintext("4")
  multiply_plain_inplace!(x_sq_plus_one, plain_four, evaluator)
  multiply!(encrypted_result, x_sq_plus_one, x_plus_one_sq, evaluator)
  println("    + size of encrypted_result: ", length(encrypted_result))
  println("    + noise budget in encrypted_result: ",
          invariant_noise_budget(encrypted_result, decryptor),
          " bits")
  println("NOTE: Decryption can be incorrect if noise budget is zero.")

  println()
  println("~~~~~~ A better way to calculate 4(x^2+1)(x+1)^2. ~~~~~~")

  print_line(@__LINE__)
  println("Generate relinearization keys.")
  relin_keys_ = RelinKeys()
  create_relin_keys!(relin_keys_, keygen)

  print_line(@__LINE__)
  println("Compute and relinearize x_squared (x^2),")
  println(" "^13, "then compute x_sq_plus_one (x^2+1)")
  x_squared = Ciphertext()
  square!(x_squared, x_encrypted, evaluator)
  println("    + size of x_squared: ", length(x_squared))
  relinearize_inplace!(x_squared, relin_keys_, evaluator)
  println("    + size of x_squared (after relinearization): ", length(x_squared))
  add_plain!(x_sq_plus_one, x_squared, plain_one, evaluator)
  println("    + noise budget in x_sq_plus_one: ",
          invariant_noise_budget(x_sq_plus_one, decryptor),
          " bits")
  decrypt!(decrypted_result, x_sq_plus_one, decryptor)
  println("    + decryption of x_sq_plus_one: 0x", to_string(decrypted_result), " ...... Correct.")

  print_line(@__LINE__)
  x_plus_one = Ciphertext()
  println("Compute x_plus_one (x+1),")
  println(" "^13, "then compute and relinearize x_plus_one_sq ((x+1)^2).")
  add_plain!(x_plus_one, x_encrypted, plain_one, evaluator)
  square!(x_plus_one_sq, x_plus_one, evaluator)
  println("    + size of x_plus_one_sq: ", length(x_plus_one_sq))
  relinearize_inplace!(x_plus_one_sq, relin_keys_, evaluator)
  println("    + noise budget in x_plus_one_sq: ",
          invariant_noise_budget(x_plus_one_sq, decryptor),
          " bits")
  decrypt!(decrypted_result, x_plus_one_sq, decryptor)
  println("    + decryption of x_plus_one_sq: 0x", to_string(decrypted_result),
          " ...... Correct.")

  print_line(@__LINE__)
  println("Compute and relinearize encrypted_result (4(x^2+1)(x+1)^2).")
  multiply_plain_inplace!(x_sq_plus_one, plain_four, evaluator)
  multiply!(encrypted_result, x_sq_plus_one, x_plus_one_sq, evaluator)
  println("    + size of encrypted_result: ", length(encrypted_result))
  relinearize_inplace!(encrypted_result, relin_keys_, evaluator)
  println("    + size of encrypted_result (after relinearization): ", length(encrypted_result))
  println("    + noise budget in encrypted_result: ",
          invariant_noise_budget(encrypted_result, decryptor),
          " bits")
  println()
  println("NOTE: Notice the increase in remaining noise budget.")

  print_line(@__LINE__)
  println("Decrypt encrypted_result (4(x^2+1)(x+1)^2).")
  decrypt!(decrypted_result, encrypted_result, decryptor)
  println("    + decryption of 4(x^2+1)(x+1)^2 = 0x", to_string(decrypted_result),
          " ...... Correct.")
  println()

  print_line(@__LINE__)
  println("An example of invalid parameters")
  set_poly_modulus_degree!(parms, 2048)
  context = SEALContext(parms)
  print_parameters(context)
  println("Parameter validation (failed): ", parameter_error_message(context))
  println()

  return
end


# Differentiation for Hackers
# ===========================
#
# These notebooks are an exploration of various approaches to analytical
# differentiation. Differentiation is something you learned in school; we start
# with an expression like $y = 3x^2 + 2x + 1$ and find an expression for the
# derivative like $\frac{dy}{dx} = 6x + 2$. Once we have such an expression, we
# can *evaluate* it by plugging in a specific value for $x$ (say 0.5) to find
# the derivative at that point (in this case $\frac{dy}{dx} = 5$).
#
# Despite its surface simplicity, this technique lies at the core of all modern
# machine learning and deep learning, alongside many other parts of statistics,
# mathematical optimisation and engineering. There has recently been an
# explosion in automatic differentiation (AD) tools, all with different designs
# and tradeoffs, and it can be difficult to understand how they relate to each
# other.
#
# We aim to fix this by beginning with the "calculus 101" rules that you are
# familiar with and implementing simple symbolic differentiators over mathematical
# expressions. Then we show how tweaks to this basic framework generalise from
# expressions to programming languages, leading us to modern automatic
# differentiation tools and machine learning frameworks like TensorFlow and
# PyTorch, and giving us a unified view across the AD landscape.

# Symbolic Differentiation
# ------------------------
#
# To talk about derivatives, we need to talk about *expressions*, which are
# symbolic forms like $x^2 + 1$ (as opposed to numbers like $5$). Normal Julia
# programs only work with numbers; we can write down $x^2 + 1$ but this only
# lets us calculate its value for a specific $x$.

x = 2
y = x^2 + 1

# However, Julia also offers a *quotation operator* which lets us talk about the
# expression itself, without needing to know what $x$ is.

y = :(x^2 + 1)
#-
typeof(y)

# Expressions are a tree data structure. They have a `head` which tells us what
# kind of expression they are (say, a function call or if statement). They have
# `args`, their children, which may be further sub-expressions. For example,
# $x^2 + 1$ is a call to $+$, and one of its children is the expression $x^2$.

y.head
#-
y.args

# We could have built this expression by hand rather than using quotation. It's
# just a bog-standard tree data structure that happens to have nice printing.

x2 = Expr(:call, :^, :x, 2)
#-
y = Expr(:call, :+, x2, 1)

# We can evaluate our expression to get a number out.

eval(y)

# When we differentiate something, we'll start by manipulating an expression
# like this, and then we can optionally evaluate it with numbers to get a
# numerical derivative. I'll call these the "symbolic phase" and the "numeric
# phase" of differentiation, respectively.

# How might we differentiate an expression like $x^2 + 1$? We can start by
# looking at the basic rules in differential calculus.
#
# $$
# \begin{align}
# \frac{d}{dx} x &= 1 \\
# \frac{d}{dx} (-u) &= - \frac{du}{dx} \\
# \frac{d}{dx} (u + v) &= \frac{du}{dx} + \frac{dv}{dx} \\
# \frac{d}{dx} (u * v) &= v \frac{du}{dx} + u \frac{dv}{dx} \\
# \frac{d}{dx} (u / v) &= (v \frac{du}{dx} - u \frac{dv}{dx}) / v^2 \\
# \frac{d}{dx} u^n &= n u^{n-1} \\
# \end{align}
# $$
#
# Seeing $\frac{d}{dx}(u)$ as a function, these rules look a lot like a
# recursive algorithm. To differentiate something like `y = a + b`, we
# differentiate `a` and `b` and combine them together. To differentiate `a` we
# do the same thing, and so on; eventually we'll hit something like `x` or `3`
# which has a trivial derivative ($1$ or $0$).

# Let's start by handling the obvious cases, $y = x$ and $y = 1$.

function derive(ex, x)
  ex == x ? 1 :
  ex isa Union{Number,Symbol} ? 0 :
  error("$ex is not differentiable")
end
#-
y = :(x)
derive(y, :x)
#-
y = :(1)
derive(y, :x)

# We can look for expressions of the form `y = a + b` using pattern matching,
# with a package called
# [MacroTools](https://github.com/MikeInnes/MacroTools.jl). If `@capture`
# returns true, then we can work with the sub-expressions `a` and `b`.

using MacroTools

y = :(x + 1)
#-
@capture(y, a_ * b_)
#-
@capture(y, a_ + b_)
#-
a, b

# Let's use this to add a rule to `derive`, following the chain rule above.

function derive(ex, x)
  ex == x ? 1 :
  ex isa Union{Number,Symbol} ? 0 :
  @capture(ex, a_ + b_) ? :($(derive(a, x)) + $(derive(b, x))) :
  error("$ex is not differentiable")
end
#-
y = :(x + 1)
derive(y, :x)
#-
y = :(x + (1 + (x + 1)))
derive(y, :x)

# These are the correct derivatives, even if they could be simplified a bit. We
# can go on to add the rest of the rules similarly.

function derive(ex, x)
  ex == x ? 1 :
  ex isa Union{Number,Symbol} ? 0 :
  @capture(ex, a_ + b_) ? :($(derive(a, x)) + $(derive(b, x))) :
  @capture(ex, a_ * b_) ? :($a * $(derive(b, x)) + $b * $(derive(a, x))) :
  @capture(ex, a_^n_Number) ? :($(derive(a, x)) * ($n * $a^$(n-1))) :
  @capture(ex, a_ / b_) ? :($b * $(derive(a, x)) - $a * $(derive(b, x)) / $b^2) :
  error("$ex is not differentiable")
end

# This is enough to get us a slightly more difficult derivative.

y = :(3x^2 + (2x + 1))
dy = derive(y, :x)

# This is correct – it's equivalent to $6x + 2$ – but it's also a bit noisy, with a
# lot of redundant terms like $x + 0$. We can clean this up by creating some
# smarter functions to do our symbolic addition and multiplication. They'll just
# avoid actually doing anything if the input is redundant.

addm(a, b) = a == 0 ? b : b == 0 ? a : :($a + $b)
mulm(a, b) = 0 in (a, b) ? 0 : a == 1 ? b : b == 1 ? a : :($a * $b)
mulm(a, b, c...) = mulm(mulm(a, b), c...)
#-
addm(:a, :b)
#-
addm(:a, 0)
#-
mulm(:b, 1)

# Our tweaked `derive` function:

function derive(ex, x)
  ex == x ? 1 :
  ex isa Union{Number,Symbol} ? 0 :
  @capture(ex, a_ + b_) ? addm(derive(a, x), derive(b, x)) :
  @capture(ex, a_ * b_) ? addm(mulm(a, derive(b, x)), mulm(b, derive(a, x))) :
  @capture(ex, a_^n_Number) ? mulm(derive(a, x),n,:($a^$(n-1))) :
  @capture(ex, a_ / b_) ? :($(mulm(b, derive(a, x))) - $(mulm(a, derive(b, x))) / $b^2) :
  error("$ex is not differentiable")
end

# And the output is much cleaner.

y = :(3x^2 + (2x + 1))
dy = derive(y, :x)

# Having done this, we can also calculate a nested derivative
# $\frac{d^2y}{dx^2}$, and so on.

ddy = derive(dy, :x)
#-
derive(ddy, :x)

# There is a deeper problem with our differentiation algorithm, though. Look at
# how big this derivative is.

derive(:(x / (1 + x^2)), :x)

# Adding an extra `* x` makes it even bigger! There's a bunch of redundant work
# here, repeating the expression $1 + x^2$ three times over.

derive(:(x / (1 + x^2) * x), :x)

# This happens because our rules look like,
# $\frac{d(u*v)}{dx} = u*\frac{dv}{dx} + v*\frac{du}{dx}$.
# Every multiplication repeats the whole sub-expression and its derivative,
# making the output exponentially large in the size of its input.
#
# This seems to be an achilles heel for our little differentiator, since it will
# make it impractical to run on any realistically-sized program. But wait!
# Things are not quite as simple as they seem, because this expression is not
# *actually* as big as it looks.
#
# Imagine we write down:

y1 = :(1 * 2)
y2 = :($y1 + $y1 + $y1 + $y1)

# This looks like a large expression, but in actual fact it does not contain
# $1*2$ four times over, just four pointers to $y1$; it is not really a tree but
# a graph that gets printed as a tree. We can show this by explicitly printing
# the expression in a way that preserves structure.
#
# (The definition of `printstructure` is not important to understand, but is
# here for reference.)

printstructure(x, _, _) = x

function printstructure(ex::Expr, cache = IdDict(), n = Ref(0))
  haskey(cache, ex) && return cache[ex]
  args = map(x -> printstructure(x, cache, n), ex.args)
  cache[ex] = sym = Symbol(:y, n[] += 1)
  println(:($sym = $(Expr(ex.head, args...))))
  return sym
end

printstructure(y2);

# Note that this is *not* the same as running common subexpression elimination
# to simplify the tree, which would have an $O(n^2)$ computational cost. If
# there is real duplication in the expression, it'll show up.

:(1*2 + 1*2) |> printstructure;

# This is effectively a change in notation: we were previously using a kind of
# "calculator notation" in which any computation used more than once had to be
# repeated in full. Now we are allowed to use variable bindings to get the same
# effect.

# If we try `printstructure` on our differentiated code, we'll see that the
# output is not so bad after all:

:(x / (1 + x^2)) |> printstructure;
#-
derive(:(x / (1 + x^2)), :x)
#-
derive(:(x / (1 + x^2)), :x) |> printstructure;

# The expression $x^2 + 1$ is now defined once and reused rather than being
# repeated, and adding the extra `* x` now adds a couple of instructions to our
# derivative, rather than doubling its size. It turns out that our "naive"
# symbolic differentiator actually preserves structure in a very sensible way,
# and we just needed the right program representation to exploit that.

derive(:(x / (1 + x^2) * x), :x)
#-
derive(:(x / (1 + x^2) * x), :x) |> printstructure;

# Calculator notation – expressions without variable bindings – is a terrible
# format for anything, and will tend to blow up in size whether you
# differentiate it or not. Symbolic differentiation is commonly criticised for
# its susceptability to "expression swell", but in fact has nothing to do with
# the differentiation algorithm itself, and we need not change it to get better
# results.
#
# Conversely, the way we have used `Expr` objects to represent variable bindings
# is perfectly sound, if a little unusual. This format could happily be used to
# illustrate all of the concepts in this handbook, and the recursive algorithms
# used to do so are elegant. However, it will clarify some things if they are
# written a little more explicitly; for this we'll introduce a new, equivalent
# representation for expressions.

# The Wengert List
# ----------------
#
# The output of `printstructure` above is known as a "Wengert List", an explicit
# list of instructions that's a bit like writing assembly code. Really, Wengert
# lists are nothing more or less than mathematical expressions written out
# verbosely, and we can easily convert to and from equivalent `Expr` objects.

include("utils.jl");
#-
y = :(3x^2 + (2x + 1))
#-
wy = Wengert(y)
#-
Expr(wy)

# Inside, we can see that it really is just a list of function calls, where
# $y_n$ refers to the result of the $n^{th}$.

wy.instructions

# Like `Expr`s, we can also build Wengert lists by hand.

w = Wengert()
tmp = push!(w, :(x^2))
w
#-
push!(w, :($tmp + 1))
w

# Armed with this, we can quite straightforwardly port our recursive symbolic
# differentiation algorithm to the Wengert list.

function derive(ex, x, w)
  ex isa Variable && (ex = w[ex])
  ex == x ? 1 :
  ex isa Union{Number,Symbol} ? 0 :
  @capture(ex, a_ + b_) ? push!(w, addm(derive(a, x, w), derive(b, x, w))) :
  @capture(ex, a_ * b_) ? push!(w, addm(mulm(a, derive(b, x, w)), mulm(b, derive(a, x, w)))) :
  @capture(ex, a_^n_Number) ? push!(w, mulm(derive(a, x, w),n,:($a^$(n-1)))) :
  @capture(ex, a_ / b_) ? push!(w, :($(mulm(b, derive(a, x, w))) - $(mulm(a, derive(b, x, w))) / $b^2)) :
  error("$ex is not differentiable")
end

derive(w::Wengert, x) = (derive(w[end], x, w); w)

# It behaves identically to what we wrote before; we have only changed the
# underlying representation.

derive(Wengert(:(3x^2 + (2x + 1))), :x) |> Expr

# In fact, we can compare them directly using the `printstructure` function we
# wrote earlier.

derive(:(x / (1 + x^2)), :x) |> printstructure
#-
derive(Wengert(:(x / (1 + x^2))), :x)

# They are *almost* identical; the only difference is the unused variable `y3`
# in the Wengert version. This happens because our `Expr` format effectively
# removes dead code for us automatically. We'll see the same thing happen if
# we convert the Wengert list back into an `Expr`.

derive(Wengert(:(x / (1 + x^2))), :x) |> Expr

function derive(w::Wengert, x)
  ds = Dict()
  ds[x] = 1
  d(x) = get(ds, x, 0)
  for v in keys(w)
    ex = w[v]
    Δ = @capture(ex, a_ + b_) ? addm(d(a), d(b)) :
        @capture(ex, a_ * b_) ? addm(mulm(a, d(b)), mulm(b, d(a))) :
        @capture(ex, a_^n_Number) ? mulm(d(a),n,:($a^$(n-1))) :
        @capture(ex, a_ / b_) ? :($(mulm(b, d(a))) - $(mulm(a, d(b))) / $b^2) :
        error("$ex is not differentiable")
    ds[v] = push!(w, Δ)
  end
  return w
end

derive(Wengert(:(x / (1 + x^2))), :x) |> Expr

# One more thing. The astute reader may notice that our differentiation
# algorithm begins with $\frac{dx}{dx}=1$ and propagates this forward to the
# output; in other words it does [forward-mode
# differentiation](./backandforth.ipynb). We can tweak our code a little to do
# reverse mode instead.

function derive_r(w::Wengert, x)
  ds = Dict()
  d(x) = get(ds, x, 0)
  d(x, Δ) = ds[x] = haskey(ds, x) ? addm(ds[x],Δ) : Δ
  d(lastindex(w), 1)
  for v in reverse(collect(keys(w)))
    ex = w[v]
    Δ = d(v)
    if @capture(ex, a_ + b_)
      d(a, Δ)
      d(b, Δ)
    elseif @capture(ex, a_ * b_)
      d(a, push!(w, mulm(Δ, b)))
      d(b, push!(w, mulm(Δ, a)))
    elseif @capture(ex, a_^n_Number)
      d(a, mulm(Δ, n, :($a^$(n-1))))
    elseif @capture(ex, a_ / b_)
      d(a, push!(w, mulm(Δ, b)))
      d(b, push!(w, :(-$(mulm(Δ, a))/$b^2)))
    else
      error("$ex is not differentiable")
    end
  end
  push!(w, d(x))
  return w
end

# There are only two distinct algorithms in this handbook, and this is the
# second! It's quite similar to forward mode, with the difference that we
# walk backwards over the list, and each time we see a usage of a variable
# $y_i$ we accumulate a gradient for that variable.

derive_r(Wengert(:(x / (1 + x^2))), :x) |> Expr

# For now, the output looks pretty similar to that of forward mode; we'll
# explain why the [distinction makes a difference](./backandforth.ipynb) in future
# notebooks.


# Arnoldi/Lanczos iteration algorithms (with custom IOP)

#######################################
# Output type/cache
"""
    KrylovSubspace{T}(n,[maxiter=30]) -> Ks

Constructs an uninitialized Krylov subspace, which can be filled by `arnoldi!`.

The dimension of the subspace, `Ks.m`, can be dynamically altered but should
be smaller than `maxiter`, the maximum allowed arnoldi iterations.

    getV(Ks) -> V
    getH(Ks) -> H

Access methods for the (extended) orthonormal basis `V` and the (extended)
Gram-Schmidt coefficients `H`. Both methods return a view into the storage
arrays and has the correct dimensions as indicated by `Ks.m`.

    resize!(Ks, maxiter) -> Ks

Resize `Ks` to a different `maxiter`, destroying its contents.

This is an expensive operation and should be used scarcely.
"""
mutable struct KrylovSubspace{B, T, U}
    m::Int        # subspace dimension
    maxiter::Int  # maximum allowed subspace size
    augmented::Int# length of the augmented part
    beta::B       # norm(b,2)
    V::Matrix{T}  # orthonormal bases
    H::Matrix{U}  # Gram-Schmidt coefficients (real for Hermitian matrices)
    KrylovSubspace{T,U}(n::Integer, maxiter::Integer=30, augmented::Integer=false) where {T,U} =
        new{real(T), T, U}(maxiter, maxiter, augmented, zero(real(T)), Matrix{T}(undef, n + augmented, maxiter + 1),
                           fill(zero(U), maxiter + 1, maxiter + !iszero(augmented)))
    KrylovSubspace{T}(args...) where {T} = KrylovSubspace{T,T}(args...)
end
getV(Ks::KrylovSubspace) = @view(Ks.V[:, 1:Ks.m + 1])
getH(Ks::KrylovSubspace) = @view(Ks.H[1:Ks.m + 1, 1:Ks.m+!iszero(Ks.augmented)])
function Base.resize!(Ks::KrylovSubspace{B,T,U}, maxiter::Integer) where {B,T,U}
    isaugmented = !iszero(Ks.augmented)
    V = Matrix{T}(undef, size(Ks.V, 1), maxiter + 1)
    H = fill(zero(U), maxiter + 1, maxiter + isaugmented)
    if isaugmented
      copyto!(@view(V[axes(Ks.V)...]), Ks.V)
      copyto!(@view(H[axes(Ks.H)...]), Ks.H)
    end
    Ks.V = V; Ks.H = H
    Ks.m = Ks.maxiter = maxiter
    return Ks
end
function Base.show(io::IO, ::MIME"text/plain", Ks::KrylovSubspace)
    io′ = IOContext(io, :limit => true, :displaysize => 3 .* displaysize(io) .÷ 7)
    println(io, "$(Ks.m)-dimensional Krylov subspace with fields")
    println(io, "beta: $(Ks.beta)")
    println(io, "V:")
    show(io′, "text/plain", getV(Ks))
    println(io)
    println(io, "H:")
    show(io′, "text/plain", getH(Ks))
end

#######################################
# Arnoldi/Lanczos with custom IOP
## High-level interface
"""
    arnoldi(A,b[;m,tol,opnorm,iop]) -> Ks

Performs `m` anoldi iterations to obtain the Krylov subspace K_m(A,b).

The n x (m + 1) basis vectors `getV(Ks)` and the (m + 1) x m upper Hessenberg
matrix `getH(Ks)` are related by the recurrence formula

```
v_1=b,\\quad Av_j = \\sum_{i=1}^{j+1}h_{ij}v_i\\quad(j = 1,2,\\ldots,m)
```

`iop` determines the length of the incomplete orthogonalization procedure [^1].
The default value of 0 indicates full Arnoldi. For symmetric/Hermitian `A`,
`iop` will be ignored and the Lanczos algorithm will be used instead.

Refer to `KrylovSubspace` for more information regarding the output.

Happy-breakdown occurs whenver `norm(v_j) < tol * opnorm`, in this case
the dimension of `Ks` is smaller than `m`.

[^1]: Koskela, A. (2015). Approximating the matrix exponential of an
advection-diffusion operator using the incomplete orthogonalization method. In
Numerical Mathematics and Advanced Applications-ENUMATH 2013 (pp. 345-353).
Springer, Cham.
"""
function arnoldi(A, b; m=min(30, size(A, 1)), ishermitian=LinearAlgebra.ishermitian(A), kwargs...)
    TA, Tb = eltype(A), eltype(b)
    T = promote_type(TA, Tb)
    Ks = KrylovSubspace{T, ishermitian ? real(T) : T}(length(b), m)
    arnoldi!(Ks, A, b; m=m, ishermitian=ishermitian, kwargs...)
end

## Low-level interface

@inline function applyA!(y, A::AT, x, V, j, n, p) where AT
    # We cannot add `@inbounds` to `mul!`, because it is provided by the user.
    AT <: Tuple #= augmented =# || (mul!(y, A, x); return)
    A, B = A
    @inbounds begin
        # V[1:n, j + 1] = A * @view(V[1:n, j]) + B * @view(V[n+1:n+p, j])
        mul!(@view(V[1:n, j + 1]), A, @view(V[1:n, j]))
        BLAS.gemm!('N', 'N', 1.0, B, @view(V[n+1:n+p, j]), 1.0, @view(V[1:n, j + 1]))
        copyto!(@view(V[n+1:n+p-1, j + 1]), @view(V[n+2:n+p, j]))
        V[end, j + 1] = 0
    end
    return
end

function checkdims(A, b, V)
    isaugmented = b isa Tuple # isaugmented
    if isaugmented
        b′, b_aug = b
        n, p = length(b′), length(b_aug)
        _A = first(A)
    else
        n, p = size(V, 1), 0
        _A, b′, b_aug = A, b, nothing
    end
    length(b′) == size(_A,1) == size(_A,2) == size(V, 1)-p || throw(DimensionMismatch("length(b′) [$(length(b′))] == size(_A,1) [$(size(_A,1))] == size(_A,2) [$(size(_A,2))] == size(V, 1)-p [$(size(V, 1)-p)] doesn't hold"))
    return b′, b_aug, n, p
end

##############################
# Utilities
##############################
"""
    firststep!(Ks, V, H, b) -> nothing

Compute the first step of Arnoldi or Lanczos iteration.
"""
function firststep!(Ks::KrylovSubspace, V, H, b)
    @inbounds begin
        fill!(H, zero(eltype(H)))
        Ks.beta = norm(b)
        @. V[:, 1] = b / Ks.beta
    end
    return
end

"""
    firststep!(Ks, V, H, b, b_aug, t, mu, l) -> nothing

Compute the first step of Arnoldi or Lanczos iteration of augmented system.
"""
function firststep!(Ks::KrylovSubspace, V, H, b, b_aug, t, mu, l)
    @inbounds begin
        n, p = length(b), length(b_aug)
        for k=1:p-1
            i = p - k
            b_aug[k] = t^i/factorial(i) * mu
        end
        b_aug[p] = mu

        # Initialize the matrices V and H
        fill!(H, 0)

        # Normalize initial vector (this norm is nonzero)
        bl = @view b[:, l]
        Ks.beta = beta = sqrt(bl'bl + b_aug'b_aug)

        # The first Krylov basis vector
        @. V[1:n, 1]     = bl / beta
        @. V[n+1:n+p, 1] = b_aug / beta
    end
    return
end

##############################
# Arnoldi
##############################
"""
    arnoldi_step!(j, iop, n, A, V, H)

Take the `j`:th step of the Lanczos iteration.
"""
function arnoldi_step!(j::Integer, iop::Integer, A::AT,
                       V::AbstractMatrix{T}, H::AbstractMatrix{U},
                       n::Int=-1, p::Int=-1) where {AT,T,U}
    x, y = @view(V[:, j]), @view(V[:, j+1])
    applyA!(y, A, x, V, j, n, p)
    @inbounds for i = max(1, j - iop + 1):j
        α = H[i, j] = coeff(U, dot(@view(V[:, i]), y))
        axpy!(-α, @view(V[:, i]), y)
    end
    β = H[j+1, j] = norm(y)
    @. y /= β
    return β
end

"""
    arnoldi!(Ks,A,b[;tol,m,opnorm,iop,init]) -> Ks

Non-allocating version of `arnoldi`.
"""
function arnoldi!(Ks::KrylovSubspace{B, T1, U}, A::AT, b;
                  tol::Real=1e-7, m::Int=min(Ks.maxiter, size(A, 1)),
                  ishermitian::Bool=LinearAlgebra.ishermitian(A isa Tuple ? first(A) : A),
                  opnorm=nothing, iop::Int=0,
                  init::Int=0, t::Number=NaN, mu::Number=NaN, l::Int=-1) where {B, T1 <: Number, U <: Number, AT}
    ishermitian && return lanczos!(Ks, A, b; tol=tol, m=m, init=init, t=t, mu=mu, l=l)
    m > Ks.maxiter ? resize!(Ks, m) : Ks.m = m # might change if happy-breakdown occurs
    @inbounds V, H = getV(Ks), getH(Ks)
    b′, b_aug, n, p = checkdims(A, b, V)
    if iszero(init)
        isaugmented = AT <: Tuple
        isaugmented ? firststep!(Ks::KrylovSubspace, V, H, b′, b_aug, t, mu, l) : firststep!(Ks::KrylovSubspace, V, H, b)
        init = 1
    end
    iszero(iop) && (iop = m)
    for j = init:m
        beta = arnoldi_step!(j, iop, A, V, H, n, p)
        if beta < tol # happy-breakdown
            Ks.m = j
            break
        end
    end
    return Ks
end

##############################
# Lanczos
##############################

"""
    lanczos_step!(j, m, n, A, V, H)

Take the `j`:th step of the Lanczos iteration.
"""
function lanczos_step!(j::Integer, A,
                       V::AbstractMatrix{T},
                       u::AbstractVector{U},
                       v::AbstractVector{B},
                       n::Int=-1, p::Int=-1) where {B,T,U}
    x, y = @view(V[:, j]),@view(V[:, j+1])
    applyA!(y, A, x, V, j, n, p)
    α = u[j] = coeff(U, dot(x, y))
    axpy!(-α, x, y)
    j > 1 && axpy!(-v[j-1], @view(V[:, j-1]), y)
    β = v[j] = norm(y)
    @. y /= β
    return β
end

"""
    coeff(::Type,α)

Helper functions that returns the real part if that is what is
required (for Hermitian matrices), otherwise returns the value
as-is.
"""
coeff(::Type{T},α::T) where {T} = α
coeff(::Type{U},α::T) where {U<:Real,T<:Complex} = real(α)

"""
    realview(::Type, V) -> real view of `V`
"""
realview(::Type{R}, V::AbstractVector{C}) where {R,C<:Complex} =
    @view(reinterpret(R, V)[1:2:end])
realview(::Type{R}, V::AbstractVector{R}) where {R} = V

"""
    lanczos!(Ks,A,b[;tol,m,opnorm]) -> Ks

A variation of `arnoldi!` that uses the Lanczos algorithm for
Hermitian matrices.
"""
function lanczos!(Ks::KrylovSubspace{B, T1, U}, A::AT, b;
                  tol=1e-7, m=min(Ks.maxiter, size(A, 1)),
                  opnorm=nothing,
                  init::Int=0, t::Number=NaN, mu::Number=NaN, l::Int=-1) where {B, T1 <: Number, U <: Number, AT}
    m > Ks.maxiter ? resize!(Ks, m) : Ks.m = m # might change if happy-breakdown occurs
    @inbounds V, H = getV(Ks), getH(Ks)
    b′, b_aug, n, p = checkdims(A, b, V)
    if iszero(init)
        isaugmented = AT <: Tuple
        isaugmented ? firststep!(Ks::KrylovSubspace, V, H, b′, b_aug, t, mu, l) : firststep!(Ks::KrylovSubspace, V, H, b)
        init = 1
    end
    @inbounds begin
        u = @diagview(H)
        # `v` is always real, even though `u` may (in general) be complex.
        v = realview(B, @diagview(H,-1))
    end
    for j = 1:m
        if tol > lanczos_step!(j, A, V, u, v, n, p)
            # happy-breakdown
            Ks.m = j
            break
        end
    end
    @inbounds copyto!(@diagview(H,1), v[1:end-1])
    return Ks
end


# ---
# title: 415. Add Strings
# id: problem415
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: String
# link: <https://leetcode.com/problems/add-strings/description/>
# hidden: true
# ---
#
# Given two non-negative integers `num1` and `num2` represented as string,
# return the sum of `num1` and `num2`.
#
# **Note:**
#
#   1. The length of both `num1` and `num2` is < 5100.
#   2. Both `num1` and `num2` contains only digits `0-9`.
#   3. Both `num1` and `num2` does not contain any leading zero.
#   4. You **must not use any built-in BigInteger library** or **convert the inputs to integer** directly.
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


using PQEcolaPoint
using Dates

# what is the contract?
short_put_contract = PutContractModel()
short_put_contract.ticker = "XYZ"
short_put_contract.expiration_date = Date(2022, 04, 14)
short_put_contract.strike_price = 105.0
short_put_contract.premium = 0.0
short_put_contract.number_of_contracts = 1
short_put_contract.direction = 1

# other version -
val_2 = θ(short_put_contract; Sₒ=111.69, number_of_levels=80, σ=0.5486, T=(28.0 / 365), μ=0.0045)


using Oceananigans.Utils: prettysummary

struct VariableBiharmonicDiffusionCoefficient
    time_scale :: Float64
end

@inline function (vbd::VariableBiharmonicDiffusionCoefficient)(i, j, k, grid, lx, ly, lz)
    Δh⁻² = 1 / Δx(i, j, k, grid, lx, ly, lz)^2 + 1 / Δy(i, j, k, grid, lx, ly, lz)^2
    Δh⁴ = 1 / Δh⁻²^2
    return Δh⁴ / vbd.time_scale
end

Base.summary(vbdc::VariableBiharmonicDiffusionCoefficient) = string("VariableBiharmonicDiffusionCoefficient(time_scale=",
                                                                    prettytime(vbdc.time_scale), ")")