#= PROCESSES
Examples of Levy measures (functions ν_p, ν_n and σ_κ)
Tempered stable, truncated stable and Watanabe class
(compatible with the second versions of the functions above)
=#

using SpecialFunctions, GSL, Base.Threads, PoissonRandom

#= TemperedStable (tempered stable), which includes the models: CGMY and KoBoL

Levy measure:
ν(dx)/dx = c_±|x|^{-1-α_±} e^{-λ_±|x|}

Input:
1. b: drift
2-3. c_p, c_n: nonnegative factors such that c_p + c_n > 0
4-5. α_p, α_n: nonnegative factors in [0,2)
6-7. λ_p, λ_n: positive factors

Output:
Functions (b_ν,sν_p,sν_n,sum_sν_p,sum_sν_n,σ_ν), where
1. b_ν(x) = b-∫_{(-1,1)∖(-x,x)}y ν(dy),
2-5. sν_p, sν_n, sum_sν_p and sum_sν_n are as above and
6. σ_ν(x) = √(∫_{(-x,x)}y^2 ν(dy)).
=#
function TemperedStable_init(b::Number,c_p::Number,c_n::Number,α_p::Number,α_n::Number,λ_p::Number,λ_n::Number)
    b_ν(x) = b - ( c_p*λ_p^(α_p-1)*(sf_gamma_inc(1-α_p, x*λ_p)-sf_gamma_inc(1-α_p, λ_p)) -
        c_n*λ_n^(α_n-1)*(sf_gamma_inc(1-α_n, x*λ_n)-sf_gamma_inc(1-α_n, λ_n)) )
    σ_ν(x) = sqrt(c_p*λ_p^(α_p-2)*gamma(2-α_p) * gamma_inc(2-α_p, x*λ_p, 0)[1] +
        c_n*λ_n^(α_n-2)*gamma(2-α_n) * gamma_inc(2-α_n, x*λ_n, 0)[1])
    (sν_p,sum_sν_p) = tempered_power_sampler(α_p,λ_p,c_p)
    (sν_n,sum_sν_n) = tempered_power_sampler(α_n,λ_n,c_n)

    return (b_ν,sν_p,sν_n,sum_sν_p,sum_sν_n,σ_ν)
end

# (Rosinski's series representation)
# Samples the jumps of the measure y ↦ c y^{-α-1}e^{-λy}1_{y>x}dy
function tempered_power_sampler(α::Number,λ::Number,c::Number)
    function s(x::Number,t::Number)
        n = pois_rand(c*t*x^(-α)/α)
        X = min.( rand(n) .^ (-1/α) .* x , rand(Exponential(1/λ),n) .* rand(n) .^ (1/α))
        return X[X .> x]
    end
    function sum_s(x::Number,t::Number)
        n = pois_rand(c*t*x^(-α)/α)
        ap = 1/α
        an = -ap
        dist = Exponential(1/λ)

        if n > 64
            X = zeros(nthreads())
            Y = zeros(nthreads())
            @threads for i = 1:n
                Y[threadid()] = min(rand() ^ an * x , rand(dist) * rand() ^ ap)
                if Y[threadid()] > x
                    X[threadid()] += Y[threadid()]
                end
            end
            return sum(X)
        else
            X = min.(rand(n) .^ an .* x , rand(dist,n) .* rand(n) .^ ap)
            return sum(X[X .> x])
        end
    end
    function sum_s(x::Number,y::Number,t::Number)
        n = pois_rand(c*t*y^(-α)/α)
        ap = 1/α
        an = -ap
        dist = Exponential(1/λ)

        if n > 64
            X1 = zeros(nthreads())
            X2 = zeros(nthreads())
            Y = zeros(nthreads())
            @threads for i = 1:n
                Y[threadid()] = min(rand() ^ an * y , rand(dist) * rand() ^ ap)
                if Y[threadid()] > x
                    X1[threadid()] += Y[threadid()]
                    X2[threadid()] += Y[threadid()]
                elseif Y[threadid()] > y
                    X2[threadid()] += Y[threadid()]
                end
            end
            return (sum(X1),sum(X2))
        else
            X = min.(rand(n) .^ an .* y , rand(dist,n) .* rand(n) .^ ap)
            return (sum(X[X .> x]), sum(X[X .> y]))
        end
    end
    return (s,sum_s)
end

#= Stable process

Levy measure:
ν(dx)/dx = c_±|x|^{-1-α_±}

Input:
1. b: drift
2-3. c_p, c_n: nonnegative factors such that c_p + c_n > 0
4-5. α_p, α_n: nonnegative factors in [0,2)
6-7. λ_p, λ_n: positive factors

Output:
Functions (b_ν,sν_p,sν_n,sum_sν_p,sum_sν_n,σ_ν), where
1. b_ν(x) = b-∫_{(-1,1)∖(-x,x)}y ν(dy),
2-5. sν_p, sν_n, sum_sν_p and sum_sν_n are as above and
6. σ_ν(x) = √(∫_{(-x,x)}y^2 ν(dy)).
=#
function Stable_init(b::Number,c_p::Number,c_n::Number,α_p::Number,α_n::Number)
    b_ν(x) = b + c_p * (α_p != 1 ? (x^(1 - α_p) - 1) / (1-α_p) : log(x)) + c_p * (α_n != 1 ? (x^(1 - α_n) - 1) / (1-α_n) : log(x))
    σ_ν(x) = sqrt(c_p * x^(2 - α_p) / (2-α_p) + c_n * x^(2 - α_n) / (2-α_n))
    (sν_p,sum_sν_p) = power_sampler(α_p,c_p)
    (sν_n,sum_sν_n) = power_sampler(α_n,c_n)

    return (b_ν,sν_p,sν_n,sum_sν_p,sum_sν_n,σ_ν)
end

# Inversion of tail Levy measure
# Samples the jumps of the measure y↦cy^{-α-1}dy
function power_sampler(α::Number,c::Number)
    function s(x::Number,t::Number)
        return rand(pois_rand(c*t*x^(-α)/α)) .^ (-1/α) .* x
    end
    function sum_s(x::Number,t::Number)
        n = pois_rand(c*t*x^(-α)/α)
        a = -1/α
        if n > 32
            X = zeros(nthreads())
            @threads for i = 1:n
                X[threadid()] += rand() ^ a
            end
            return sum(X) * x
        else
            return sum(rand(n) .^ a) * x
        end
    end
    function sum_s(x::Number,y::Number,t::Number)
        n = pois_rand(c*t*y^(-α)/α)
        sc = x/y
        a = -1/α

        if n > 32
            X1 = zeros(nthreads())
            X2 = zeros(nthreads())
            Y = zeros(nthreads())
            @threads for i = 1:n
                Y[threadid()] = rand() ^ a
                X2[threadid()] += Y[threadid()]
                if Y[threadid()] > sc
                    X1[threadid()] += Y[threadid()]
                end
            end
            return (sum(X1) * x, sum(X2) * x)
        else
            X = rand(n) .^ a
            return (sum(X[X .> sc]) * x, sum(X) * x)
        end
    end
    return (s,sum_s)
end

#= (subclass of) Watanabe (continuous singular)

General Levy measure: for some m∈ℕ, m≥2,
ν = ∑_{n∈ℤ}c_{n,+}δ_{m^n}+c_{n,-}δ_{-m^n}
With ∑_{n∈ℕ}(c_{-n,+}+c_{-n,-})b^(-n)+∑_{n∈ℕ}c_{n,+}+c_{n,-}<∞
and max_{n∈ℤ}(c_+ + c_-)<∞.

Our class:
ν = ∑_{n∈ℕ}c_+ δ_{m^(-n)} + c_- δ_{-m^(-n)}.

Note:
Infinite activity iff min(r_+,r_-)≤1
Blumenthal-Getoor index = max(0,-log(max(r_+,r_-))/log(b)))

Input:
1. b: drift
2. m: integer ≥2
3-4. c_p, c_n: nonnegative numbers

Output:
Functions (b_ν,sν_p,sν_n,sum_sν_p,sum_sν_n,σ_ν), where
1. b_ν(x) = b-∫_{(-1,1)∖(-x,x)}y ν(dy),
2-5. sν_p, sν_n, sum_sν_p and sum_sν_n are as above and
6. σ_ν(x) = √(∫_{(-x,x)}y^2 ν(dy)).
=#
function Watanabe_init(b::Number,m::Integer,c_p::Number,c_n::Number)
    function b_ν(x::Number)
        return b - (c_p - c_n) * m^(-( x >= 1 ? 0. : floor(-log(x)/log(m))+1 )) / (1-1/m)
    end
    function σ_ν(x::Number)
        return sqrt((c_p + c_n) * m^(-2*( x >= 1 ? 0. : floor(-log(x)/log(m))+1 )) / (1-1/m^2))
    end
    (sν_p,sum_sν_p) = wat_sampler(c_p,m)
    (sν_n,sum_sν_n) = wat_sampler(c_n,m)

    return (b_ν,sν_p,sν_n,sum_sν_p,sum_sν_n,σ_ν)
end

# Computes the sum c * |{n∈ℕ:m^n≤1/x}|
function wat(c::Number,m::Integer)
    local f(x) = x >= 1 ? 0. :
        c * ceil(-log(x)/log(m))
    return f
end

# Samples from the density p(k) = 1 / |{n∈ℕ:m^n≤1/x}|
function wat_sampler(c::Number,m::Integer)
    function s(x::Number,t::Number)
        if x >= 1
            return  Float64[]
        else
            return [1/m^ceil(rand() * floor(-log(x)/log(m))) for i = 1:Int(pois_rand(c * t * ceil(-log(x)/log(m))))]
        end
    end
    function sum_s(x::Number,t::Number)
        return sum(s(x,t))
    end
    function sum_s(x::Number,y::Number,t::Number)
        aux = s(y,t)
        return (sum(aux .* (aux .> x)), sum(aux))
    end
    return (s,sum_s)
end


using Test
using GradientBoost.Util
using Random

@testset "utils" begin
    @testset "weighted_median" begin
        weights = [1, 1, 1, 1, 1]
        values = [1, 2, 3, 4, 5]
        @test weighted_median(weights, values) == 3

        weights = [1, 1, 1, 1, 1]
        values = [5, 4, 3, 2, 1]
        @test weighted_median(weights, values) == 3

        weights = [1, 1, 1, 1, 1]
        values = [4, 5, 2, 3, 1]
        @test weighted_median(weights, values) == 3

        weights = [1, 3, 1, 1, 1]
        values = [1, 2, 3, 4, 5]
        @test weighted_median(weights, values) == 2

        weights = [1, 1, 1, 1, 100]
        values = [1, 2, 3, 4, 5]
        @test weighted_median(weights, values) == 5

        weights = [100, 1, 1, 1, 1]
        values = [1, 2, 3, 4, 5]
        @test weighted_median(weights, values) == 1

        weights = [0.5, 0.5, 1, 1, 1]
        values = [1, 2, 3, 4, 5]
        @test weighted_median(weights, values) == 3
    end

    @testset "holdout returns proportional partitions" begin
        n = 10
        right_prop = 0.3
        (left, right) = holdout(n, right_prop)

        @test size(left, 1) == n - (n * right_prop)
        @test size(right, 1) == n * right_prop
        @test isempty(intersect(left, right))
        @test size(union(left, right), 1) == n
    end

end # module


# simple program to test MCL clustering

using Test
using Clustering

@testset "MCL" begin

Random.seed!(34568)

@testset "nonsquare matrices not supported" begin
    @test_throws DimensionMismatch mcl(zeros(Float64, 4, 3))
end

# initialize adjacency matrix of a weighted graph
nodes = [:bat, :bit, :cat, :fit, :hat, :hit]
edges = Tuple{Symbol, Symbol, Float64}[(:cat, :hat, 0.2), (:hat, :bat, 0.16),
                                      (:bat, :cat, 1.0), (:bat, :bit, 0.125),
                                      (:bit, :fit, 0.25), (:fit, :hit, 0.5),
                                      (:hit, :bit, 0.16)]
adj_matrix = zeros(Float64, length(nodes), length(nodes))
for edge in edges
    n1 = findfirst(isequal(edge[1]), nodes)
    n2 = findfirst(isequal(edge[2]), nodes)
    adj_matrix[n1, n2] = adj_matrix[n2, n1] = edge[3]
end
@assert issymmetric(adj_matrix)

@testset "fractional inflation param (1.8)" begin
    res = mcl(adj_matrix, display=:none, inflation=1.8)
    @test isa(res, MCLResult)
    local k = length(res.counts)
    # @show k
    @test k == 2
    @test all(a -> 1 <= a <= k, res.assignments)
    @test length(res.assignments) == length(nodes)
    @test length(res.counts) == k
    local c
    for c in 1:k
        @test count(==(c), res.assignments) == res.counts[c]
    end
    @test res.nunassigned == 0
    @test res.assignments == [1, 2, 1, 2, 1, 2]
end

@testset "integer inflation param (2)" begin
    res = mcl(adj_matrix, display=:none, inflation=2)
    @test isa(res, MCLResult)
    @test length(res.assignments) == length(nodes)
    @test res.nunassigned == 0
end

@testset "test non-integral expansion" begin
    # should not raise an exception
    res = mcl(adj_matrix, display=:none, inflation=1.5, expansion=1.5, save_final_matrix=true)
    @test isa(res, MCLResult)
    @test length(res.assignments) == length(nodes)
    @test size(res.mcl_adj) == size(adj_matrix) # test that the matrix is returned
end

@testset "allow_singles option" begin
    res = mcl(diagm(0 => [1.0, 1.0]), display=:none, allow_singles=true)
    @test length(res.counts) == 2
    @test res.assignments == [1, 2]
    @test res.counts == [1, 1]
    @test res.nunassigned == 0

    res = mcl(diagm(0 => [1.0, 1.0]), display=:none, allow_singles=false)
    @test length(res.counts) == 0
    @test res.assignments == [0, 0]
    @test res.nunassigned == 2
end

@testset "sparse input matrix" begin
    res = mcl(sparse(adj_matrix), display=:none, expansion=2)
    @test isa(res, MCLResult)
    @test length(res.assignments) == length(nodes)
    @test res.nunassigned == 0
    @test eltype(res.mcl_adj) === Float64

    # fractional powers not supported for sparse matrices
    @test_broken mcl(sparse(adj_matrix), display=:none, expansion=2.1)
end

@testset "use Float32 input" begin
    res = mcl(convert(Matrix{Float32},adj_matrix), display=:none, expansion=2)
    @test isa(res, MCLResult)
    @test eltype(res.mcl_adj) === Float32
end

end


include("preamble.jl")

# -----------------------
# Figures with VTK
# -----------------------

using LazySets, WriteVTK
import Polyhedra

# propagate sets
rot(θ) = [cos(θ) -sin(θ) 0; sin(θ) cos(θ) 0; 0 0 1]
H = Hyperrectangle(ones(3), fill(0.1, 3))
adom = range(0, 2pi, length=20)
X = [rot(θ) * H for θ in adom]

# convexify
R = [CH(X[i], X[i+1]) for i in 1:length(X)-1]
push!(R, CH(X[end], X[1]))
RCH = concretize.(R)

# save
writevtk(X, file="rotations")
writevtk(RCH, file="rotations_CH")

using ReachabilityAnalysis

@taylorize function lorenz!(dx, x, p, t)
    local σ = 10.0
    local β = 8.0 / 3.0
    local ρ = 28.0
    dx[1] = σ * (x[2] - x[1])
    dx[2] = x[1] * (ρ - x[3]) - x[2]
    dx[3] = x[1] * x[2] - β * x[3]
    return dx
end

X0 = Hyperrectangle(low=[0.9, 0.0, 0.0], high=[1.1, 0.0, 0.0])
prob = @ivp(x' = lorenz!(x), dim=3, x(0) ∈ X0);
alg = TMJets(abstol=1e-15, orderT=10, orderQ=2, maxsteps=50_000);

sol = solve(prob, T=20.0, alg=alg);

solZ = set.(overapproximate(sol, Zonotope))
writevtk(solZ, file="lorenz") # generates lorenz.vtu


using Plots
using RecipesBase

function get_width_bounds(mid_point_x, interval_width)
    half_width = 0.5 * interval_width
    return mid_point_x .+ [-half_width, half_width]
end

function get_interval_vertices(width_bounds, length_bounds)
    lower_bound_width, upper_bound_width = width_bounds
    lower_bound_length, upper_bound_length = length_bounds
    vertices_x = [
        lower_bound_width,
        lower_bound_width,
        upper_bound_width,
        upper_bound_width
    ]
    vertices_y = [
        lower_bound_length,
        upper_bound_length,
        upper_bound_length,
        lower_bound_length
    ]
    return vertices_x, vertices_y
end

@shorthands confidenceinterval

@recipe function f(
        ::Type{Val{:confidenceinterval}},
        args...;
        interval_width=0.75,
        interval_color=:skyblue,
        line_color=:black
)
    interval_data = args[2]
    mid_point_x = plotattributes[:series_plotindex]
    width_bounds = get_width_bounds(mid_point_x, interval_width)

    legend := false
    @series begin
        # Confidence Interval Bar
        seriestype := :shape
        seriesalpha --> 0.5
        fillcolor := interval_color
        length_bounds = [interval_data[1], interval_data[3]]
        vertices = get_interval_vertices(width_bounds, length_bounds)
        x := vertices[1]
        y := vertices[2]
        ()
    end

    @series begin
        seriestype := :shape
        seriescolor := line_color
        line_value = interval_data[2]
        x := width_bounds
        y := [line_value, line_value]
        ()
    end
end


export bm_solver

"""
`history_check(hist::Dict, g::Int)` sees if the guess `g` is consistent
with the previous answers stored in `hist`.
"""
function history_check(history::Dict{Int,Tuple{Int,Int}}, g::Int)::Bool
    for d in keys(history)
        if history[d] != bm_count(g, d)
            return false
        end
    end
    return true
end

"""
`history_count(h)` returns the number of possible codes that
are consistent with the history dictionary `h`.
"""
function history_count(history::Dict{Int,Tuple{Int,Int}})::Int
    length([c for c in 0:9999 if history_check(history, c)])
end

"""
`bm_solver(c::Int, verbose::Bool=true, delay::Int=0)` playes Bullets and Maggots to try
to find code `c`. Returns the number of guesses. It doesn't cheat!

* `verbose` controls printing during execution.
* `delay` indicates how many steps we should wait before using the history; `0` means no delay
"""
function bm_solver(c::Int, verbose::Bool = true, first_guess = guess_any)::Int
    if !code_check(c)
        throw(bad_code_message(c))
        return 0
    end
    if verbose
        println("Trying to find secret code $(string4(c))")
    end

    history = Dict{Int,Tuple{Int,Int}}()   # place to hold past guesses/results
    result_server(g::Int) = bm_count(g, c)
    bm_solver_engine(result_server, verbose, history, first_guess)
end

function bm_solver_engine(
    result_server::Function,
    verbose::Bool,
    history::Dict{Int,Tuple{Int,Int}},
    first_guess::Function = guess_any,
)::Int
    all_codes = collect(0:9999)            # randomized list of possible answers
    shuffle!(all_codes)

    steps = 0

    g = first_guess()
    cnt = result_server(g)
    history[g] = cnt
    steps += 1
    if verbose
        println(
            "Guess $steps:\t",
            string4(g),
            "\t",
            cnt,
            "\tsearch space size is now ",
            history_count(history),
        )
    end

    if cnt == (4, 0)
        if verbose
            println("Solved! Code is $(string4(g))")
        end
        return steps
    end

    for g in all_codes  # current guess (g) code

        if !history_check(history, g)
            continue
        end

        cnt = result_server(g)

        steps += 1
        history[g] = cnt

        if verbose
            println(
                "Guess $steps:\t",
                string4(g),
                "\t",
                cnt,
                "\tsearch space size is now ",
                history_count(history),
            )
        end

        if cnt == (4, 0)
            if verbose
                println("Solved! Code is $(string4(g))")
            end
            return steps
        end
    end
    return -1  # signal failure
end


function insertionsort!(data::Array)
	for i in 2 : length(data)
		current_value = data[i]
		for j in i - 1 : -1 : 1
			if current_value < data[j]
				data[j + 1] = data[j]
				if j == 1
					data[j] = current_value
				end
			else
				data[j + 1] = current_value
				break
			end
		end
	end
end


@testset "Concrete Accuracy" begin
    zero_accurate = Double64(0.0)
    one_accurate  = Double64(1.0)

    pi_accurate   = Double64(pi)
    phi_accurate  = Double64(phi)

    @test pi_accurate   == DoubleFloat{Float64}(pi)

    @test abs(inv(inv(pi_accurate)) - pi_accurate) <= eps(LO(pi_accurate))
    @test abs(phi_accurate - (phi_accurate*phi_accurate - 1.0)) <= eps(LO(phi_accurate))

    a = Base.TwicePrecision(3.0) / Base.TwicePrecision(7.0)
    b = Double64(3.0) / Double64(7.0)
    c = Base.TwicePrecision(17.0) / Base.TwicePrecision(5.0)
    d = Double64(17.0) / Double64(5.0)

    @test a.hi == b.hi
    @test a.lo == b.lo
    @test (a+c).hi == (b+d).hi
    @test (a+c).lo == (b+d).lo
    @test (a-c).hi == (b-d).hi
    @test abs((a-c).lo - (b-d).lo) <= eps((a-c).lo)
    @test (a*c).hi == (b*d).hi
    @test (a*c).lo == (b*d).lo
    @test (a/c).hi == (b/d).hi
    @test (a/c).lo == (b/d).lo

    @test zero(Double64) == Double64(0.0, 0.0)
    @test one(Double64) == Double64(1.0, 0.0)
    @test zero(Double32) == Double32(0.0f0, 0.0f0)
    @test one(Double32) == Double32(1.0f0, 0.0f0)

    @test round(Double64(123456.0, 1.0e-17), RoundUp) == Double64(123457.0, 0.0)
    @test round(Double64(123456.0, -1.0e-17), RoundUp) == Double64(123456.0, 0.0)
    @test round(Double64(123456.0, -1.0e-17), RoundDown) == Double64(123455.0, 0.0)

    @test typemax(Double64) == Double64(typemax(Float64))
    @test floatmin(Double32) == Double32(floatmin(Float32))

    @test isnan( nan(Double64) )
    @test isinf( inf(Double16) )

    @test typeof(rand(Double64)) == Double64
    @test typeof(rand(Double32)) == Double32
    @test typeof(rand(Double16)) == Double16
end


n = rand(1:10)
@test matrixdepot("parter", n) == matrixdepot("parter", Float64, n)

A = matrixdepot("cauchy", [1.:n;] + 0.5, -[1.:n;])

@test_approx_eq matrixdepot("parter", n) A
println("'parter' passed test...")


struct PrisonersDilemma end

n_agents(simpleGame::PrisonersDilemma) = 2

ordered_actions(simpleGame::PrisonersDilemma, i::Int) = [:cooperate, :defect]
ordered_joint_actions(simpleGame::PrisonersDilemma) = vec(collect(Iterators.product([ordered_actions(simpleGame, i) for i in 1:n_agents(simpleGame)]...)))

n_joint_actions(simpleGame::PrisonersDilemma) = length(ordered_joint_actions(simpleGame))
n_actions(simpleGame::PrisonersDilemma, i::Int) = length(ordered_actions(simpleGame, i))

function reward(simpleGame::PrisonersDilemma, i::Int, a)
    if i == 1
        noti = 2
    else
        noti = 1
    end

    if a[i] == :cooperate && a[noti] == :cooperate
        return -1.0
    elseif a[i] == :defect && a[noti] == :cooperate
        return 0.0
    elseif a[i] == :cooperate && a[noti] == :defect
        return -4.0
    elseif a[i] == :defect && a[noti] == :defect
        return -3.0
    end
end

function joint_reward(simpleGame::PrisonersDilemma, a)
    return [reward(simpleGame, i, a) for i in 1:n_agents(simpleGame)]
end

function SimpleGame(simpleGame::PrisonersDilemma)
    return SimpleGame(
        0.9,
        vec(collect(1:n_agents(simpleGame))),
        [ordered_actions(simpleGame, i) for i in 1:n_agents(simpleGame)],
        (a) -> joint_reward(simpleGame, a)
    )
end


d = readdlm("input.txt", ':', Int)
d = Dict(mapslices(x -> (x[1], x[2]), d, 2))
# Part 1
severity = 0
for i in keys(d)
    if i % (2 * d[i] - 2) == 0
        severity += i * d[i]
    end
end
println(severity)

# Part 2
cleanRun = false
delay = 1
while !cleanRun
    cleanRun = true
    for i in keys(d)
        if (i + delay) % (2 * d[i] - 2) == 0
            cleanRun = false
            delay += 1
            break
        end
    end
end
println(delay)


### A Pluto.jl notebook ###
# v0.16.4

using Markdown
using InteractiveUtils

# ╔═╡ 50c6aa0a-38b8-11ec-36c5-7784eaef9e4c
begin
	using Pkg, Revise
	Pkg.activate("../../")
end

# ╔═╡ 69dd34c7-ac62-443c-8908-eba042438f07
function add_gauss_noise_fl!(img, σ)
	for i in eachindex(img)
		img[i] = img[i] + randn(eltype(img)) * eltype(img)(σ)
		# img[i] = img[i] + randn() * σ # this also works
	end
	return img
end

# ╔═╡ bb684353-9d66-4ecd-b2bc-b24aeadcf972
randn() # returns normal distribution noise

# ╔═╡ 9f6dbe12-a0f1-47a1-9b81-0f0b417b43b3
typeof(randn())

# ╔═╡ a6be7e47-6837-487d-9490-3b4689c1a3fc
begin
	arr = ones(Float32, (10, 10))
	add_gauss_noise_fl!(arr, 123213)
end

# ╔═╡ d82c4907-8de1-47f9-9959-2f562b544770
eltype(arr)

# ╔═╡ 6104149f-c39c-45d1-a9fd-50c1c8b0d71e
arr

# ╔═╡ 8092c8a3-2b09-442a-b8c0-6fce4e8e0889

# ╔═╡ 2ab1d4fc-e191-4261-9bb3-d189db987b55

# ╔═╡ Cell order:
# ╠═50c6aa0a-38b8-11ec-36c5-7784eaef9e4c
# ╠═69dd34c7-ac62-443c-8908-eba042438f07
# ╠═d82c4907-8de1-47f9-9959-2f562b544770
# ╠═bb684353-9d66-4ecd-b2bc-b24aeadcf972
# ╠═9f6dbe12-a0f1-47a1-9b81-0f0b417b43b3
# ╠═a6be7e47-6837-487d-9490-3b4689c1a3fc
# ╠═6104149f-c39c-45d1-a9fd-50c1c8b0d71e
# ╠═8092c8a3-2b09-442a-b8c0-6fce4e8e0889
# ╠═2ab1d4fc-e191-4261-9bb3-d189db987b55


# Spatial calculations

export
    vector1D,
    vector,
    wrap_coords,
    wrap_coords_vec,
    maxwell_boltzmann,
    random_velocities,
    random_velocities!,
    bond_angle,
    torsion_angle

"""
    vector1D(c1, c2, side_length)

Displacement between two 1D coordinate values from c1 to c2, accounting for
the bounding box.
The minimum image convention is used, so the displacement is to the closest
version of the coordinate accounting for the periodic boundaries.
"""
function vector1D(c1, c2, side_length)
    if c1 < c2
        return (c2 - c1) < (c1 - c2 + side_length) ? (c2 - c1) : (c2 - c1 - side_length)
    else
        return (c1 - c2) < (c2 - c1 + side_length) ? (c2 - c1) : (c2 - c1 + side_length)
    end
end

"""
    vector(c1, c2, box_size)

Displacement between two coordinate values from c1 to c2, accounting for
the bounding box.
The minimum image convention is used, so the displacement is to the closest
version of the coordinates accounting for the periodic boundaries.
"""
vector(c1, c2, box_size) = vector1D.(c1, c2, box_size)

@generated function vector(c1::SVector{N}, c2::SVector{N}, box_size) where N
    quote
        Base.Cartesian.@ncall $N SVector{$N} i->vector1D(c1[i], c2[i], box_size[i])
    end
end

square_distance(i, j, coords, box_size) = sum(abs2, vector(coords[i], coords[j], box_size))

# Pad a vector to 3D to allow operations such as the cross product
function vector_pad3D(c1::SVector{2, T}, c2::SVector{2, T}, box_size::SVector{2, T}) where T
    SVector{3, T}(
        vector1D(c1[1], c2[1], box_size[1]),
        vector1D(c1[2], c2[2], box_size[2]),
        zero(T),
    )
end

vector_pad3D(c1::SVector{3}, c2::SVector{3}, box_size::SVector{3}) = vector(c1, c2, box_size)

# Trim a vector back to 2D if required
trim3D(v::SVector{3, T}, box_size::SVector{2}) where T = SVector{2, T}(v[1], v[2])
trim3D(v::SVector{3}, box_size::SVector{3}) = v

"""
    wrap_coords(c, side_length)

Ensure a 1D coordinate is within the simulation box and return the coordinate.
"""
wrap_coords(c, side_length) = c - floor(c / side_length) * side_length

"""
    wrap_coords_vec(c, box_size)

Ensure a coordinate is within the simulation box and return the coordinate.
"""
wrap_coords_vec(v, box_size) = wrap_coords.(v, box_size)

"""
    velocity(mass, temperature; dims=3)

Generate a random velocity from the Maxwell-Boltzmann distribution.
"""
function AtomsBase.velocity(mass, temp; dims::Integer=3)
    return SVector([maxwell_boltzmann(mass, temp) for i in 1:dims]...)
end

const mb_conversion_factor = uconvert(u"u * nm^2 * ps^-2 * K^-1", Unitful.k)

"""
    maxwell_boltzmann(mass, temperature)

Generate a random speed along one dimension from the Maxwell-Boltzmann distribution.
"""
function maxwell_boltzmann(mass, temp)
    T = typeof(convert(AbstractFloat, ustrip(temp)))
    k = unit(temp) == NoUnits ? T(ustrip(mb_conversion_factor)) : T(mb_conversion_factor)
    σ = sqrt(k * temp / mass)
    return rand(Normal(zero(T), T(ustrip(σ)))) * unit(σ)
end

"""
    random_velocities(sys, temp)

Generate random velocities from the Maxwell-Boltzmann distribution
for a `System`.
"""
function random_velocities(sys::System{D}, temp) where D
    [velocity(a.mass, temp; dims=D) for a in Array(sys.atoms)]
end

"""
    random_velocities!(sys, temp)

Set the velocities of a `System` to random velocities generated from the
Maxwell-Boltzmann distribution.
"""
function random_velocities!(sys::System, temp)
    sys.velocities = random_velocities(sys, temp)
    return sys
end

# Sometimes domain error occurs for acos if the value is > 1.0 or < -1.0
acosbound(x::Real) = acos(clamp(x, -1, 1))

"""
    bond_angle(coord_i, coord_j, coord_k, box_size)
    bond_angle(vec_ba, vec_bc)

Calculate the bond or pseudo-bond angle in radians between three
coordinates or two vectors.
The angle between B→A and B→C is returned in the range 0 to π.
"""
function bond_angle(coords_i, coords_j, coords_k, box_size)
    vec_ba = vector(coords_j, coords_i, box_size)
    vec_bc = vector(coords_j, coords_k, box_size)
    return bond_angle(vec_ba, vec_bc)
end

function bond_angle(vec_ba, vec_bc)
    acosbound(dot(vec_ba, vec_bc) / (norm(vec_ba) * norm(vec_bc)))
end

"""
    torsion_angle(coord_i, coord_j, coord_k, coord_l, box_size)
    torsion_angle(vec_ab, vec_bc, vec_cd)

Calculate the torsion angle in radians defined by four coordinates or
three vectors.
The angle between the planes defined by atoms (i, j, k) and (j, k, l) is
returned in the range -π to π.
"""
function torsion_angle(coords_i, coords_j, coords_k, coords_l, box_size)
    vec_ab = vector(coords_i, coords_j, box_size)
    vec_bc = vector(coords_j, coords_k, box_size)
    vec_cd = vector(coords_k, coords_l, box_size)
    return torsion_angle(vec_ab, vec_bc, vec_cd)
end

function torsion_angle(vec_ab, vec_bc, vec_cd)
    cross_ab_bc = vec_ab × vec_bc
    cross_bc_cd = vec_bc × vec_cd
    θ = atan(
        ustrip(dot(cross_ab_bc × cross_bc_cd, normalize(vec_bc))),
        ustrip(dot(cross_ab_bc, cross_bc_cd)),
    )
    return θ
end


using ForwardDiff
using ForwardDiff: Dual

seed(x::Real, ::Val) = Dual(x, true)

function seed(x, ::Val{N}, offset = 0) where N
  map(x, reshape(1:length(x), size(x))) do x, i
    Dual(x, ntuple(j -> j+offset == i, Val(N)))
  end
end

extract(x::ForwardDiff.Dual) = x.value, [x.partials...]

function extract(xs::AbstractArray{ForwardDiff.Dual{T,V,N}}) where {T,V,N}
  J = similar(xs, V, N, length(xs))
  for i = 1:length(xs), j = 1:N
    J[j, i] = xs[i].partials.values[j]
  end
  return map(x -> x.value, xs), J
end

function forward_jacobian(f, x, ::Val{N}) where N
  y, _J = extract(f(seed(x, Val(N))))
  J = similar(_J, length(x), length(y))
  J[1:N,:] = _J
  offset = 0
  while offset + N < length(x)
    offset += N
    _, _J = extract(f(seed(x, Val(N), offset)))
    range = (1+offset):min(N+offset,length(x))
    J[range,:] = @view _J[range.-offset,:]
  end
  return y, J
end

function forward_jacobian(f, x)
  len = min(length(x), ForwardDiff.DEFAULT_CHUNK_THRESHOLD)
  forward_jacobian(f, x, Val(len))
end

vec_scalar(x) = vec(x)
vec_scalar(x::Real) = [x]
reshape_scalar(x, y) = reshape(y, size(x))
reshape_scalar(x::Real, y) = y[]

"""
    forwarddiff(f, x) -> f(x)

Runs `f(x)` as usual, but instructs Zygote to differentiate `f` using forward
mode, rather than the usual reverse mode.

Forward mode takes time linear in `length(x)` but only has constant memory
overhead, and is very efficient for scalars, so in some cases this can be a
useful optimisation.

```julia
julia> function pow(x, n)
         r = one(x)
         for i = 1:n
           r *= x
         end
         return r
       end
pow (generic function with 1 method)

julia> gradient(5) do x
         forwarddiff(x) do x
           pow(x, 2)
         end
       end
(10,)
```

Note that the function `f` will *drop gradients* for any closed-over values.

```julia
julia> gradient(2, 3) do a, b
         forwarddiff(a) do a
           a*b
         end
       end
(3, nothing)
```

This can be rewritten by explicitly passing through `b`, i.e.

```julia
gradient(2, 3) do a, b
  forwarddiff([a, b]) do (a, b)
    a*b
  end
end
```
"""
forwarddiff(f, x) = f(x)

@adjoint function forwarddiff(f, x)
  y, J = forward_jacobian(f, x)
  return y, ȳ -> (nothing, reshape_scalar(x, J*vec_scalar(ȳ)))
end


apply_covariance = (x,M) -> sum(dot.(x,M*x))

optional_posZ(x::NTuple{1,T}, offset::NTuple{1,T}) where {T,N} = 1
optional_posZ(x::NTuple{2,T}, offset::NTuple{2,T}) where {T,N} = 1
optional_posZ(x::NTuple{N,T}, offset::NTuple{N,T}) where {T,N} = x[3]-offset[3]

# List of functions and names we want to predefine
function generate_functions_expr()
    # offset and scale is already wrapped in the generator function
    x_expr = :(scale .* (x .- offset))
    x_expr1 = :(scale[1] .* (x[1] .- offset[1]))
    x_expr2 = :(scale[2] .* (x[2] .- offset[2]))
    x_expr3 = :(scale[3] .* (x[3] .- offset[3]))
    x_expr4 = :(scale[4] .* (x[4] .- offset[4]))
    x_expr5 = :(scale[5] .* (x[5] .- offset[5]))
    x_expr14 = :(maximum(abs.(scale .* (x .- offset)))) # useful for rectangular things
    x_expr15 = :(minimum(abs.(scale .* (x .- offset))))
    x_expr6 = :(prod(x .== offset))
    x_expr7 = :(cis(dot((x .- offset), scale)))
    x_expr8 = :(exp(.- sum(abs2.(x .- offset).*scale))) # scale is 1/(2 sigma^2)
    x_expr9 = :(exp(.- sum(abs2.(x .- offset).*scale)) ./ prod(sqrt.(pi ./ scale)))
    x_expr10 = :(exp(.- apply_covariance((x .- offset), scale) ))
    x_expr11 = :(exp(.- apply_covariance((x .- offset), scale)) ./ prod(sqrt.(pi ./ abs2(scale))))
    x_expr12 = :(abs2(scale[1] .* (x[1] .- offset[1]))+abs2(scale[2] .* (x[2] .- offset[2])))
    x_expr13 = :((x[1] .- offset[1]).* scale[1] .+ (x[2] .- offset[2]) .* scale[2])

    functions = [
        (:(rr2), :(x -> T(sum(abs2.($x_expr))))),
        (:(rr),  :(x -> T(sqrt.(sum(abs2.($x_expr)))))),
        (:(xx),  :(x -> T($x_expr1))),
        (:(yy),  :(x -> T($x_expr2))),
        (:(zz),  :(x -> T($x_expr3))),
        (:(ee),  :(x -> T($x_expr4))),
        (:(tt),  :(x -> T($x_expr5))),
        (:(idx_min),  :(x -> T($x_expr15))),
        (:(idx_max),  :(x -> T($x_expr14))),
        (:(delta),  :(x -> T($x_expr6))),
        (:(phiphi), :(x -> T(atan.($x_expr2, $x_expr1)))),  # this is the arcus tangens of y/x yielding a spiral phase ramp
        (:(phase_kz), :(x -> T(optional_posZ(x,offset).*sqrt.(max(1 .- $x_expr12,0))))),  # can be used for constucting a free-space propagator in optics
        (:(phase_kxy), :(x -> T($x_expr13))),  # useful for xy shifting
        (:(exp_is),  :(x -> T($x_expr7))),  # exp(2pi i s (x-o)) # by modifying s, this becomes exp(i kx)
        (:(exp_sqr),  :(x -> T($x_expr8))),  # maximum-normalized Gaussian
        (:(exp_sqr_norm),  :(x -> T($x_expr9))),  # integral-normalized Gaussian (over infinite ROI)
        (:(exp_sqr_cov),  :(x -> T($x_expr10))),  # Gaussian based on the covariance matrix
        (:(exp_sqr_cov_norm),  :(x -> T($x_expr11))),  # normalized Gaussian based on the covariance matrix
    ]
    return functions
end

function wrap_fkt(f)  # encapsulates the expression in a proper function with index, offset and scale argument
    return (idx, offset, scale, weight) -> f(idx)  # Not that f contains offset and scale
end

function wrap_zipped(g) # encapsulates a wrapped function as called with a sinle tuple of offset and scale
    return (idx, mytuple) -> g(idx,mytuple...)
end

curry(f, x) = (xs...) -> f(x, xs...)   # just a shorthand to remove x

## These functions ensure that also numbers can be iterated and zipped
function cast_iter(vals::Matrix)
    Tuple(Tuple(vals[:,n]) for n in 1:size(vals,2))
    # Tuple(vals[:,axes(vals, 2)])
end

function cast_iter(vals::Vector{<:Number})
    Tuple(vals)   # makes the matrix iterable, interpreting it as a series of vectors
end

function cast_iter(vals::IterType)
    vals
end

function cast_iter(vals)
    repeated(vals) # during the zip operation this type always yields the same value
end

function cast_number_iter(vals::Vector{<:Number})
    Tuple(vals)   # makes the matrix iterable, interpreting it as a series of vectors
end

function cast_number_iter(vals::NTuple{N,<:Number} where N)
    vals
end

function cast_number_iter(vals::Number)
    repeated(vals) # during the zip operation this type always yields the same value
end

function optional_mat_to_iter(vals)  # only for matrices
    vals
end

function optional_mat_to_iter(vals::Matrix)
    cast_iter(vals)
end

mat_to_tvec(v) = [Tuple(v[:,n]) for n in 1:size(v,2)] # converts a 2d matrix to a Vector of Tuples

function apply_dims(scale2, dims, N)  # replaces scale entries not in dims with zeros
    ntuple(i -> i ∈ dims ? scale2[i] : zero(eltype(scale2)), N)  # zero(scale2[1])
end

function apply_dims(scales::IterType, dims, N)  # replaces scale entries not in dims with zeros
    Tuple(ntuple(i -> i ∈ dims ? scale3[i] : zero(eltype(scale3)), N)  for scale3 in  scales)
end

# we automatically generate the functions for rr2, rr, ...
# We set the types for the arguments correctly in the default cases
for F in generate_functions_expr()

    # default functions with offset and scaling behavior. This version allows no list of numbers or tuples
    @eval function $(Symbol(:_, F[1]))(::Type{T}, size::NTuple{N, Int},
                           offset::Union{Type{<:Ctr}, Number, NTuple{N,Number}, Vector{<:Number}},
                           scale::Union{Type{<:Sca}, Number, NTuple{N,Number}, Vector{<:Number}},
                           dims,
                           accumulator,
                           weight::Number) where{N, T}
        offset = get_offset(size, offset)
        scale_init = get_scale(size, scale)
        scale = apply_dims(scale_init, dims, N) # replaces scale entries not in dims with zeros

        g(x) = weight .* ($(F[2])(x)) #
        IndexFunArray(T,g, size)
    end

    @eval function $(Symbol(:_, F[1]))(::Type{T}, size::NTuple{N, Int},
        offset,   # a version that supports an iterable collection of tuples for offset and scale
        scale,
        dims=ntuple(+, N),
        accumulator = sum,
        weight=1) where{N, T}

        # its important to use different variable names due to a julia bug
        offset_ = get_offset(size, offset)
        scale_init = get_scale(size, scale)
        scale_ = apply_dims(scale_init, dims, N)

        offset2 = cast_iter(offset_)
        # @show offset2
        scale2 = cast_iter(scale_)
        weight2 = cast_number_iter(weight)
        w = (x, offset, scale, weight3) -> weight3 .* $(F[2])(x) # adds offset and scale as requires parameters to the expression
        # the line below makes a function only depending on the index position but iterating over both, offsets and scales
        fkt = (idx) -> accumulator(curry(wrap_zipped(w),idx),zip(offset2,scale2,weight2))

        IndexFunArray(T, fkt, size)
    end
    # default functions with offset and scaling behavior
    @eval function $(F[1])(::Type{T}, size::NTuple{N, Int};
                           offset=CtrFT,
                           scale=ScaUnit,
                           dims=ntuple(+, N), accumulator=sum, weight=1) where{N, T}
        # $(F[1])(T, size, offset, scale, dims)
        # print("Dispatch T3\n")
        $(Symbol(:_, F[1]))(T, size, offset, scale, dims, accumulator, weight)
    end

    # change order of offset and scale
    @eval function $(F[1])(size::NTuple{N, Int};
                           offset=CtrFT, scale=ScaUnit, dims=ntuple(+, N), accumulator=sum, weight=1) where {N}
        T = $DEFAULT_T
        $(F[1])(T, size, scale=scale, offset=offset, dims=dims, accumulator=accumulator, weight=weight)
    end

    # convenient wrapper to provide an array as input
    @eval function $(F[1])(arr::AbstractArray{T, N};
                           offset=CtrFT, scale=ScaUnit, dims=ntuple(+, N), accumulator=sum, weight=1) where {T, N}
        T2 = default_type(T, $DEFAULT_T)
        $(F[1])(T2, size(arr), scale=scale, offset=offset, dims=dims, accumulator=accumulator, weight=weight)
    end

    @eval export $(F[1])
end

export axes1d

function axes1d(::Type{T}, size::NTuple{N, Int}; offset=CtrFT, scale=ScaUnit, dims=ntuple(+, N), keepdims=false) where {T,N}
    offset = get_offset(size, offset)
    scale_init = get_scale(size, scale)
    scale = apply_dims(scale_init, dims, N) # replaces scale entries not in dims with zeros
    if keepdims
        @show offset
        (ramp(T, d, size[d]; offset=offset[d],scale=scale[d]) for d in dims)
    else
        (xx(T, (size[d],); offset=offset[d], scale=scale[d]) for d in dims)
    end
end

function axes1d(size::NTuple{N, Int}; offset=CtrFT, scale=ScaUnit, dims=ntuple(+, N), keepdims=false) where {N}
    T = DEFAULT_T
    axes1d(T, size; offset=offset, scale=scale, dims=dims, keepdims=keepdims)
end

function axes1d(arr::AbstractArray{T, N}; offset=CtrFT, scale=ScaUnit, dims=ntuple(+, N), keepdims=false) where {T, N}
    axes1d(size(arr); offset=offset, scale=scale, dims=dims, keepdims=keepdims)
end


using Zygote
using RayTracer: get_params, set_params!

function loss_fn(θ, color)
    rendered_color = raytrace(origin, direction, θ, light, eye_pos, 0)
    loss = sum(abs2.(rendered_color.x .- color.x) .+
               abs2.(rendered_color.y .- color.y) .+
               abs2.(rendered_color.z .- color.z))
    return loss
end

# Define the Scene Parameters

screen_size = (w = 400, h = 300)

light = PointLight(
    Vec3(1.0),
    1000.0,
    Vec3(0.15, 0.5, -110.5)
)

eye_pos = Vec3(0.0, 0.0, -5.0)

cam = Camera(eye_pos, Vec3(0.0), Vec3(0.0, 1.0, 0.0), 45.0, 1.0, screen_size...)

origin, direction = get_primary_rays(cam);

@testset "Triangle" begin
    scene = [
        Triangle(Vec3(-1.7, 1.0, 0.0), Vec3(1.0, 1.0, 0.0), Vec3(1.0, -1.0, 0.0),
                 Material(color_ambient = Vec3(1.0),
                          color_diffuse = Vec3(1.0),
                          color_specular = Vec3(1.0),
                          specular_exponent = 50.0,
                          reflection = 1.0))
    ]

    scene_new = [
        Triangle(Vec3(-1.9, 1.3, 0.1), Vec3(1.2, 1.1, 0.3), Vec3(0.8, -1.2, -0.15),
                 Material(color_ambient = Vec3(1.0),
                          color_diffuse = Vec3(1.0),
                          color_specular = Vec3(1.0),
                          specular_exponent = 50.0,
                          reflection = 1.0))
    ]

    color = raytrace(origin, direction, scene, light, eye_pos, 0)

    zygote_grads = get_params(gradient(x -> loss_fn(x, color), scene_new)[1][1])

    numerical_grads = get_params(numderiv(x -> loss_fn([x], color), scene_new[1]))

    @test isapprox(numerical_grads, zygote_grads, rtol = 1.0e-3)
end

@testset "Sphere" begin
    scene = [
        Sphere(Vec3(-1.7, 1.0, 0.0), [0.6],
               Material(color_ambient = Vec3(1.0),
                        color_diffuse = Vec3(1.0),
                        color_specular = Vec3(1.0),
                        specular_exponent = 50.0,
                        reflection = 1.0))
    ]

    scene_new = [
        Sphere(Vec3(-1.9, 1.3, 0.1), [0.9],
	       Material(color_ambient = Vec3(1.0),
                        color_diffuse = Vec3(1.0),
                        color_specular = Vec3(1.0),
                        specular_exponent = 50.0,
                        reflection = 1.0))
    ]

    color = raytrace(origin, direction, scene, light, eye_pos, 0)

    zygote_grads = get_params(gradient(x -> loss_fn(x, color), scene_new)[1][1])

    numerical_grads = get_params(numderiv(x -> loss_fn([x], color), scene_new[1]))

    @test isapprox(numerical_grads[1:4], zygote_grads[1:4], rtol = 1.0e-3)
end


module ExtendedFESpacesTests

using Test
using Gridap.Arrays
using Gridap.TensorValues
using Gridap.Algebra
using Gridap.ReferenceFEs
using Gridap.Geometry
using Gridap.FESpaces
using Gridap.Fields
using Gridap.CellData

n = 10
mesh = (n,n)
domain = (0,1,0,1) .- 1
order = 1
model = CartesianDiscreteModel(domain, mesh)

Ω = Triangulation(model)

const R = 0.7

function is_in(coords)
  n = length(coords)
  x = (1/n)*sum(coords)
  d = x[1]^2 + x[2]^2 - R^2
  d < 0
end

oldcell_to_coods = get_cell_coordinates(Ω)
oldcell_to_is_in = collect1d(lazy_map(is_in,oldcell_to_coods))

Ω_in = Triangulation(model,oldcell_to_is_in)
@test isa(Ω_in,BodyFittedTriangulation)

@test model === get_background_model(Ω_in)

V = TestFESpace(Ω_in,ReferenceFE(lagrangian,Float64,order),conformity=:H1)
test_single_field_fe_space(V)

U = TrialFESpace(V)
test_single_field_fe_space(U)

oldcell_to_is_out = lazy_map(!,oldcell_to_is_in)
Ω_out = Triangulation(model,oldcell_to_is_out)

Γ = InterfaceTriangulation(Ω_in,Ω_out)

degree = 2
dΩ = Measure(Ω,degree)
dΩ_in = Measure(Ω_in,degree)
dΩ_out = Measure(Ω_out,degree)
dΓ = Measure(Γ,degree)

u(x) = x[1]+1

dv = get_fe_basis(V)
du = get_trial_fe_basis(U)
uh = interpolate(u,U)

x = get_cell_points(Ω)
x_in = get_cell_points(Ω_in)
x_out = get_cell_points(Ω_out)
x_Γ = get_cell_points(Γ)

r = dv(x)
test_array(r,collect(r))
r = dv(x_in)
test_array(r,collect(r))
r = dv(x_out)
test_array(r,collect(r))
r = dv.⁺(x_Γ)
test_array(r,collect(r))

r = ∇(dv)(x)
test_array(r,collect(r))
r = ∇(dv)(x_in)
test_array(r,collect(r))
r = ∇(dv)(x_out)
test_array(r,collect(r))
r = ∇(dv).⁺(x_Γ)
test_array(r,collect(r))

r = du(x)
test_array(r,collect(r))
r = du(x_in)
test_array(r,collect(r))
r = du(x_out)
test_array(r,collect(r))
r = du.⁺(x_Γ)
test_array(r,collect(r))

r = ∇(du)(x)
test_array(r,collect(r))
r = ∇(du)(x_in)
test_array(r,collect(r))
r = ∇(du)(x_out)
test_array(r,collect(r))
r = ∇(du).⁺(x_Γ)
test_array(r,collect(r))

r = uh(x)
test_array(r,collect(r))
r = uh(x_in)
test_array(r,collect(r))
r = uh(x_out)
test_array(r,collect(r))
r = uh.⁺(x_Γ)
test_array(r,collect(r))

r = ∇(uh)(x)
test_array(r,collect(r))
r = ∇(uh)(x_in)
test_array(r,collect(r))
r = ∇(uh)(x_out)
test_array(r,collect(r))
r = ∇(uh).⁺(x_Γ)
test_array(r,collect(r))

#using Gridap.Visualization
#writevtk(Ω,"Omega",cellfields=["uh"=>uh])
#writevtk(Ω_in,"Omega_in",cellfields=["uh"=>uh])
#writevtk(Ω_out,"Omega_out",cellfields=["uh"=>uh])
#writevtk(Γ,"gamma",cellfields=["uh"=>uh.plus])

a(u,v) =
  ∫(v⊙u + ∇(v)⊙∇(u))*dΩ +
  ∫(v⊙u + ∇(v)⊙∇(u))*dΩ_in +
  ∫(v⊙u + ∇(v)⊙∇(u))*dΩ_out +
  ∫(jump(v)⊙jump(u) + jump(∇(v))⊙jump(∇(u)))*dΓ

l(v) =
  ∫(v⊙u + ∇(v)⊙∇(u))*dΩ +
  ∫(v⊙u + ∇(v)⊙∇(u))*dΩ_in +
  ∫(v⊙u + ∇(v)⊙∇(u))*dΩ_out +
  ∫(jump(v))*dΓ

op = AffineFEOperator(a,l,U,V)

V = TestFESpace(Ω_in,ReferenceFE(lagrangian,Float64,2),conformity=:H1)

V = TestFESpace(Ω_in,ReferenceFE(lagrangian,Float64,2),conformity=:H1,constraint=:zeromean)
uh = FEFunction(V,rand(num_free_dofs(V)))
@test sum(∫(uh)*dΩ_in) + 1 ≈ 1

V = TestFESpace(Ω,ReferenceFE(lagrangian,Float64,2),conformity=:H1)

V_in = TestFESpace(Ω_in,ReferenceFE(lagrangian,Float64,2),conformity=:H1)
V = TestFESpace(Ω,ReferenceFE(lagrangian,Float64,2),conformity=:H1)

vh_in = interpolate(V_in) do x
    x[1]
end
vh_in = interpolate(vh_in, V_in)
vh = interpolate(vh_in, V)

#using Gridap.Visualization
#writevtk(Ω,"Ω",cellfields=["vh"=>vh,"vh_in"=>vh_in])

end # module


import Base.rand

export BallInf,
       volume

"""
    BallInf{N, VN<:AbstractVector{N}} <: AbstractHyperrectangle{N}

Type that represents a ball in the infinity norm.

### Fields

- `center` -- center of the ball as a real vector
- `radius` -- radius of the ball as a real scalar (``≥ 0``)

### Notes

Mathematically, a ball in the infinity norm is defined as the set

```math
\\mathcal{B}_∞^n(c, r) = \\{ x ∈ \\mathbb{R}^n : ‖ x - c ‖_∞ ≤ r \\},
```
where ``c ∈ \\mathbb{R}^n`` is its center and ``r ∈ \\mathbb{R}_+`` its radius.
Here ``‖ ⋅ ‖_∞`` denotes the infinity norm, defined as
``‖ x ‖_∞ = \\max\\limits_{i=1,…,n} \\vert x_i \\vert`` for any
``x ∈ \\mathbb{R}^n``.

### Examples

Create the two-dimensional unit ball and compute its support function along the
positive ``x=y`` direction:

```jldoctest
julia> B = BallInf(zeros(2), 1.0)
BallInf{Float64, Vector{Float64}}([0.0, 0.0], 1.0)

julia> dim(B)
2

julia> ρ([1., 1.], B)
2.0
```
"""
struct BallInf{N, VN<:AbstractVector{N}} <: AbstractHyperrectangle{N}
    center::VN
    radius::N

    # default constructor with domain constraint for radius
    function BallInf(center::VN, radius::N) where {N, VN<:AbstractVector{N}}
        @assert radius >= zero(N) "radius must not be negative"
        return new{N, VN}(center, radius)
    end
end

isoperationtype(::Type{<:BallInf}) = false
isconvextype(::Type{<:BallInf}) = true

# --- AbstractHyperrectangle interface functions ---

"""
    radius_hyperrectangle(B::BallInf, i::Int)

Return the box radius of a ball in the infinity norm in a given dimension.

### Input

- `B` -- ball in the infinity norm
- `i` -- dimension of interest

### Output

The box radius of the ball in the infinity norm in the given dimension.
"""
function radius_hyperrectangle(B::BallInf, i::Int)
    return B.radius
end

"""
    radius_hyperrectangle(B::BallInf)

Return the box radius of a ball in the infinity norm, which is the same in every
dimension.

### Input

- `B` -- ball in the infinity norm

### Output

The box radius of the ball in the infinity norm.
"""
function radius_hyperrectangle(B::BallInf)
    return fill(B.radius, dim(B))
end

"""
    isflat(B::BallInf)

Determine whether a ball in the infinity norm is flat, i.e. whether its radius
is zero.

### Input

- `B` -- ball in the infinity norm

### Output

`true` iff the ball is flat.

### Notes

For robustness with respect to floating-point inputs, this function relies on
the result of `isapproxzero` when applied to the radius of the ball.
Hence, this function depends on the absolute zero tolerance `ABSZTOL`.
"""
function isflat(B::BallInf)
    return isapproxzero(B.radius)
end

function load_genmat_ballinf_static()
return quote
    function genmat(B::BallInf{N, SVector{L, N}}) where {L, N}
        r = B.radius
        if isapproxzero(B.radius)
            return SMatrix{L, 0, N, 0}()
        else
            gens = zeros(MMatrix{L, L})
            @inbounds for i in 1:L
                gens[i, i] = r
            end
            return SMatrix(gens)
        end
    end
end
end

# --- AbstractCentrallySymmetric interface functions ---

"""
    center(B::BallInf)

Return the center of a ball in the infinity norm.

### Input

- `B` -- ball in the infinity norm

### Output

The center of the ball in the infinity norm.
"""
function center(B::BallInf)
    return B.center
end

# --- LazySet interface functions ---

"""
    σ(d::AbstractVector, B::BallInf)

Return the support vector of a ball in the infinity norm in a given direction.

### Input

- `d` -- direction
- `B` -- ball in the infinity norm

### Output

The support vector in the given direction.
If the direction has norm zero, the vertex with biggest values is returned.
"""
function σ(d::AbstractVector, B::BallInf)
    @assert length(d) == dim(B) "a $(length(d))-dimensional vector is " *
                                "incompatible with a $(dim(B))-dimensional set"
    return center(B) .+ sign.(d) .* B.radius
end

# Particular dispatch for SingleEntryVector
function σ(d::SingleEntryVector, B::BallInf)
    return _σ_sev_hyperrectangle(d, B)
end

"""
    ρ(d::AbstractVector, B::BallInf)

Evaluate the support function of a ball in the infinity norm in a given
direction.

### Input

- `d` -- direction
- `B` -- ball in the infinity norm

### Output

Evaluation of the support function in the given direction.

### Algorithm

Let ``B`` be a ball in the infinity norm with center ``c`` and radius ``r`` and
let `d` be the direction of interest.
For balls with dimensions less than 30 we use the implementation for
`AbstractHyperrectangle`, taylored to a `BallInf`, which computes

```math
    ∑_{i=1}^n d_i * (c_i + \\textrm{sgn}(d_i) * r)
```

where ``\\textrm{sgn}(α) = 1`` if ``α ≥ 0`` and ``\\textrm{sgn}(α) = -1`` if ``α < 0``.

For balls of higher dimension, we instead exploit that for a support vector
``v = σ(d, B) = c + \\textrm{sgn}(d) * (r, …, r)ᵀ`` we have

```math
    ρ(d, B) = ⟨d, v⟩ = ⟨d, c⟩ + ⟨d, \\textrm{sgn}(d) * (r, …, r)ᵀ⟩ = ⟨d, c⟩ + r · ∑_{i=1}^n |d_i|
```

where ``⟨·, ·⟩`` denotes the dot product.
"""
function ρ(d::AbstractVector, B::BallInf)
    @assert length(d) == dim(B) "a $(length(d))-dimensional vector is " *
                                "incompatible with a $(dim(B))-dimensional set"
    c = center(B)
    r = B.radius
    if length(d) > 30
        # more efficient for higher dimensions
        return dot(d, c) + r * sum(abs, d)
    end
    N = promote_type(eltype(d), eltype(B))
    res = zero(N)
    @inbounds for (i, di) in enumerate(d)
        if di < zero(N)
            res += di * (c[i] - r)
        elseif di > zero(N)
            res += di * (c[i] + r)
        end
    end
    return res
end

# Particular dispatch for SingleEntryVector
function ρ(d::SingleEntryVector, B::BallInf)
    return _ρ_sev_hyperrectangle(d, B)
end

"""
    radius(B::BallInf, [p]::Real=Inf)

Return the radius of a ball in the infinity norm.

### Input

- `B` -- ball in the infinity norm
- `p` -- (optional, default: `Inf`) norm

### Output

A real number representing the radius.

### Notes

The radius is defined as the radius of the enclosing ball of the given
``p``-norm of minimal volume with the same center.
"""
function radius(B::BallInf, p::Real=Inf)
    return (p == Inf) ? B.radius : norm(fill(B.radius, dim(B)), p)
end

"""
    rand(::Type{BallInf}; [N]::Type{<:Real}=Float64, [dim]::Int=2,
         [rng]::AbstractRNG=GLOBAL_RNG, [seed]::Union{Int, Nothing}=nothing)

Create a random ball in the infinity norm.

### Input

- `BallInf` -- type for dispatch
- `N`       -- (optional, default: `Float64`) numeric type
- `dim`     -- (optional, default: 2) dimension
- `rng`     -- (optional, default: `GLOBAL_RNG`) random number generator
- `seed`    -- (optional, default: `nothing`) seed for reseeding

### Output

A random ball in the infinity norm.

### Algorithm

All numbers are normally distributed with mean 0 and standard deviation 1.
Additionally, the radius is nonnegative.
"""
function rand(::Type{BallInf};
              N::Type{<:Real}=Float64,
              dim::Int=2,
              rng::AbstractRNG=GLOBAL_RNG,
              seed::Union{Int, Nothing}=nothing)
    rng = reseed(rng, seed)
    center = randn(rng, N, dim)
    radius = abs(randn(rng, N))
    return BallInf(center, radius)
end

"""
    translate(B::BallInf, v::AbstractVector)

Translate (i.e., shift) a ball in the infinity norm by a given vector.

### Input

- `B` -- ball in the infinity norm
- `v` -- translation vector

### Output

A translated ball in the infinity norm.

### Algorithm

We add the vector to the center of the ball.

### Notes

See also [`translate!(::BallInf, ::AbstractVector)`](@ref) for the in-place version.
"""
function translate(B::BallInf, v::AbstractVector)
    return translate!(copy(B), v)
end

"""
    translate!(B::BallInf, v::AbstractVector)

Translate (i.e., shift) a ball in the infinity norm by a given vector, in-place.

### Input

- `B` -- ball in the infinity norm
- `v` -- translation vector

### Output

The ball `B` translated by `v`.

### Notes

See also [`translate(::BallInf, ::AbstractVector)`](@ref) for the out-of-place version.
"""
function translate!(B::BallInf, v::AbstractVector)
    @assert length(v) == dim(B) "cannot translate a $(dim(B))-dimensional " *
                                "set by a $(length(v))-dimensional vector"
    c = center(B)
    c .+= v
    return B
end

@inline function _vol_prod(B::BallInf{N}, n) where {N}
    vol = one(N)
    diam = 2 * B.radius
    for i in 1:n
        vol *= diam
    end
    return vol
end

# computation for floating-point results
function volume(B::BallInf{N}) where {N<:AbstractFloat}
    n = dim(B)
    if n < 50
        vol = _vol_prod(B, n)
    else
        r = B.radius
        vol = exp(n*log(2r))
    end
    return vol
end

function volume(B::BallInf)
    n = dim(B)
    vol = _vol_prod(B, n)
    return vol
end

function project(B::BallInf, block::AbstractVector{Int}; kwargs...)
    return BallInf(B.center[block], B.radius)
end


module CellFieldsTests

using Test
using FillArrays
using Gridap.Arrays
using Gridap.TensorValues
using Gridap.Fields
using Gridap.ReferenceFEs
using Gridap.Geometry
using Gridap.CellData
#using Gridap.FESpaces
using Random

domain = (0,1,0,1)
cells = (3,3)
model = CartesianDiscreteModel(domain,cells)

trian = Triangulation(model)
trian_N =BoundaryTriangulation(model)
trian_D =BoundaryTriangulation(model,tags="tag_8")
trian_S =SkeletonTriangulation(model)
trian_0 =Triangulation(model,Int32[])

trian_sv = view(trian_S,[1,5,4,2])
nv = get_normal_vector(trian_sv)
@test isa(nv,SkeletonPair)
@test isa(nv.plus,CellField)
@test isa(nv.minus,CellField)

ϕ = GenericCellField(get_cell_map(trian),trian,ReferenceDomain())
ϕinv = GenericCellField(lazy_map(inverse_map,get_cell_map(trian)),trian,PhysicalDomain())
xref = get_cell_points(trian)
xphy = CellPoint(get_array(xref),trian,PhysicalDomain())
test_array(ϕ(xref),collect1d(ϕ(xphy)),(i,j)->all(map(≈,i,j)))
test_array(ϕinv(xref),collect1d(ϕinv(xphy)),(i,j)->all(map(≈,i,j)))

x = get_cell_points(trian)
@test DomainStyle(x) == ReferenceDomain()
@test get_array(x) == get_cell_coordinates(trian)
@test get_data(x) == get_cell_ref_coordinates(trian)

px = get_physical_coordinate(trian)
test_array(px(x),collect1d(get_array(x)))

_x = change_domain(x,PhysicalDomain())
@test DomainStyle(_x) == PhysicalDomain()
@test get_array(_x) == get_cell_coordinates(trian)
@test get_data(_x) == get_cell_coordinates(trian)

_x = change_domain(x,ReferenceDomain())
@test DomainStyle(_x) == ReferenceDomain()
@test get_array(x) == get_cell_coordinates(trian)
@test get_data(x) == get_cell_ref_coordinates(trian)

ffun(x) = 2*x[1]
f = CellField(ffun,trian)
fx = f(x)
r = map(xs->ffun.(xs),get_array(x))
r = reshape(r,length(r))
test_array(fx,r,≈)

x_0 = get_cell_points(trian_0)
fx_0 = f(x_0)
test_array(fx_0,collect(fx_0))

n_S = get_normal_vector(trian_S)
x_S = get_cell_points(trian_S)

ts = get_triangulation(∇(f))
tt = get_triangulation(n_S.plus)

@test is_change_possible(ts,tt) == true

nf_S = n_S⋅∇(f)

jnf_S = jump(n_S⋅∇(f))
jnfx_S = jnf_S(x_S)
test_array(jnfx_S,0*collect(jnfx_S))

h = CellField(rand(num_cells(trian_S)),trian_S)*jump(∇(f))
hx_S = h(x_S)
test_array(hx_S,collect(hx_S))

h = 3*mean(f)⋅jump(n_S⋅∇(f))
hx_S = h(x_S)
test_array(hx_S,0*collect(hx_S))

aa(f) = 4*f
test_array((aa∘f)(x),4*r,≈)

f1 = f
f2 = 2*f
b(f1,f2) = f1+f2
test_array((b∘(f1,f2))(x),3*r,≈)

f2 = ones(num_cells(trian))
test_array((b∘(f1,f2))(x),map(i->i.+1,r),≈)

v = GenericCellField(get_cell_shapefuns(trian),trian,ReferenceDomain())
vx = v(x)
test_array(vx,collect(vx))

u = GenericCellField(lazy_map(transpose,get_data(v)),v.trian,v.domain_style)
m = v*u
test_array(m(x),collect(m(x)))
m = ∇(v)⋅∇(u)
test_array(m(x),collect(m(x)))

∇vx = ∇(v)(x)
test_array(∇vx,collect(∇vx))

∇fx = ∇(f)(x)
test_array(∇fx,collect(∇fx))

k = VectorValue(1.0,2.0)
∇kfx = ((∇+k)(f))(x)
test_array(∇kfx,collect(∇kfx))

∇kvx = ((∇+k)(v))(x)
test_array(∇kvx,collect(∇kvx))

β(x) = 2*x[1]
α = CellField(x->2*x,trian)
ax = ((∇+k)(β*α))(x)
test_array(ax,collect(ax))

ν = CellField(x->2*x,trian)
ax =((∇-k)⋅ν)(x)
test_array(ax,collect(ax))

ax =((∇-k)×ν)(x)
test_array(ax,collect(ax))

ax =((∇-k)⊗ν)(x)
test_array(ax,collect(ax))

ax =(∇.*ν)(x)
test_array(ax,collect(ax))

ax =(ν.*ν)(x)
test_array(ax,collect(ax))

ax =((∇-k).*ν)(x)
test_array(ax,collect(ax))

ax =(ν⊗(∇-k))(x)
test_array(ax,collect(ax))

σ(x) = diagonal_tensor(VectorValue(1*x[1],2*x[2]))
Fields.gradient(::typeof(σ)) = x-> ThirdOrderTensorValue{2,2,2,Float64}(1,0,0,0,0,0,0,2)
ax = ((∇+k)(σ⋅α))(x)
test_array(ax,collect(ax))

h = Operation(*)(2,f)
hx = h(x)
test_array(hx,2*fx)

a = fill(2,num_cells(trian))
h = Operation(*)(a,f)
hx = h(x)
test_array(hx,2*fx)

fx = evaluate(ffun,x)
test_array(fx,r,≈)

f_N = CellField(ffun,trian_N)
x_N = get_cell_points(trian_N)
fx_N = f_N(x_N)
test_array(fx_N,collect(fx_N))

n_N = get_normal_vector(trian_N)
nx_N = n_N(x_N)
test_array(nx_N,collect(nx_N))

h = f*n_N
hx = h(x_N)
test_array(hx,collect(hx))

gfun(x) = 3*x
g = CellField(gfun,trian)

h = Operation(*)(f,g)
gx = g(x)
hx = h(x)
r = map((i,j)->broadcast(*,i,j),fx,gx)
test_array(hx,r)

h_N = Operation(*)(f_N,g)
gx_N = g(x_N)
r = map((i,j)->broadcast(*,i,j),fx_N,gx_N)
hx_N = h_N(x_N)
test_array(hx_N,r)

g_D = CellField(gfun,trian_D)

cell_h = rand(num_cells(trian))
h = CellField(cell_h,trian)
test_array(h(x),collect(h(x)))
test_array(h(x_N),collect(h(x_N)))

h_N = (2*f_N+g)⋅g
hx_N = h_N(x_N)
test_array(hx_N,collect(hx_N))

end # module


module FP

using Distributions
using Printf

# Show floats to 2 digits after the decimal point
Base.show(io::IO, f::Float64) = @printf(io, "%2.2f", f)

Δ = 0.0     # propagation delay upper bound
h = 0.12    # honest mining rate
a = 0.08    # adversarial mining rate
λ = h + a   # total mining rate

K = 30 # confirmation depth

println("Simulation begins.\n")
println(
    "confirmation depth=",
    K,
    "  max delay=",
    Δ,
    "  honest mining rate=",
    h,
    "  adversarial mining rate=",
    a,
)

nmax = round(Int, K * 3 * λ / a)   # In principle the first violation can occur in the arbitrary future.
# However, the probability that it first occurs at a height much larger than K is extremely small.  Here we only simulate up to about 3K adversarial blocks.

ntrials = 1000

violation = 0.0

for nt = 1:ntrials
    xx = rand(Exponential(1.0 / λ), nmax)   # exponential inter-mining times for all blocks
    tt = cumsum(xx)   # the (cumulative) mining times
    bb = rand(Bernoulli(h / λ), nmax)   # each block is honest with probability h/λ
    bb[1] = true    # We assume block 1 is always honest - which is the target of the attack.

    tth = tt[bb]        # honest block mining times
    tta = tt[.!bb]      # adversarial block mining times

    len = min(length(tth), length(tta))

    if any(tth[K:len] .> tta[K:len])
        global violation += 1
    end

    if nt == 1  # print an example of
        println("\nThe following are the mining times in trial #1")
        for n = 1:nmax
            print(bb[n] ? "honest" : "advers")
            println(" block mined at ", tt[n], "    after interval ", xx[n])
            # println( bb[n], " ", xx[n], " ", tt[n] )
        end
        println(
            "\nThe mining times of the first ",
            len,
            " honest blocks:\n",
            tth[1:len],
        )
        println(
            "The mining times of the first ",
            len,
            " adversarial blocks:\n",
            tta[1:len],
        )
        println(
            "When the adversarial is ahead in this trial (1)?  ",
            tth[1:len] .> tta[1:len],
        )
    end
end

println(
    "\nThe probability of block 1's safety violation is about ",
    100.0 * violation / ntrials,
    "% based on ",
    ntrials,
    " trials.",
)
println("Simulation ends.")

end


img = SquareImage(rows=4, cols=4, l=2, calibration_value=1, background=0)
@testset "Testing FBP algorithms" begin
    @testset "Testing FBP FFT algorithm" begin
        sinog = radon(img, Radon())
        tomog = iradon(sinog, FBP())
        @test tomog ≈ [
            [-0.642772  0.374742  0.374742  -0.642772];
            [0.033042  1.05056   1.05056    0.0330419];
            [0.033042  1.05056   1.05056    0.0330422];
            [-0.642772  0.374742  0.374742  -0.642772]
        ]
    end
    @testset "Testing FBP FFT square algorithm" begin
        sinog = radon(img, RadonSquare())
        tomog = iradon(sinog, FBPFFTSquare())
        @test tomog ≈ [
            [0.177991   0.55048  0.554279  -0.260271];
            [0.0837881  1.22652  1.21234   -0.185625];
            [-0.148054   1.21234  1.22652   -0.132705];
            [-0.0358718  0.55428  0.55048    0.140421]
        ]
    end
end


using Test, LinearMaps, LinearAlgebra, BenchmarkTools

@testset "linear combinations" begin
    CS! = LinearMap{ComplexF64}(cumsum!,
                                (y, x) -> (copyto!(y, x); reverse!(y); cumsum!(y, y)), 10;
                                ismutating=true)
    v = rand(ComplexF64, 10)
    u = similar(v)
    b = @benchmarkable mul!($u, $CS!, $v)
    @test run(b, samples=3).allocs == 0
    n = 10
    L = sum(fill(CS!, n))
    @test mul!(u, L, v) ≈ n * cumsum(v)
    b = @benchmarkable mul!($u, $L, $v)
    @test run(b, samples=5).allocs <= 1
    for α in (false, true, rand(ComplexF64)), β in (false, true, rand(ComplexF64))
        @test mul!(copy(u), L, v, α, β) ≈ Matrix(L)*v*α + u*β
    end

    A = 2 * rand(ComplexF64, (10, 10)) .- 1
    B = rand(ComplexF64, size(A)...)
    M = @inferred LinearMap(A)
    N = @inferred LinearMap(B)
    @test @inferred(LinearMaps.MulStyle(M)) === matrixstyle
    @test @inferred(LinearMaps.MulStyle(N)) === matrixstyle
    LC = @inferred M + N
    @test @inferred(LinearMaps.MulStyle(LC)) === matrixstyle
    @test @inferred(LinearMaps.MulStyle(LC + I)) === matrixstyle
    @test @inferred(LinearMaps.MulStyle(LC + 2.0*I)) === matrixstyle
    v = rand(ComplexF64, 10)
    w = similar(v)
    b = @benchmarkable mul!($w, $M, $v)
    @test run(b, samples=3).allocs == 0
    if VERSION >= v"1.3.0-alpha.115"
        b = @benchmarkable mul!($w, $LC, $v)
        @test run(b, samples=3).allocs == 0
        for α in (false, true, rand(ComplexF64)), β in (false, true, rand(ComplexF64))
            if testallocs
                b = @benchmarkable mul!($w, $LC, $v, $α, $β)
                @test run(b, samples=3).allocs == 0
                b = @benchmarkable mul!($w, $(I + LC), $v, $α, $β)
                @test run(b, samples=3).allocs == 0
                b = @benchmarkable mul!($w, $(LC + I), $v, $α, $β)
                @test run(b, samples=3).allocs == 0
            end
            y = rand(ComplexF64, size(v))
            @test mul!(copy(y), LC, v, α, β) ≈ Matrix(LC)*v*α + y*β
            @test mul!(copy(y), LC+I, v, α, β) ≈ Matrix(LC + I)*v*α + y*β
        end
    end
    # @test_throws ErrorException LinearMaps.LinearCombination{ComplexF64}((M, N), (1, 2, 3))
    @test @inferred size(3M + 2.0N) == size(A)
    # addition
    @test @inferred convert(Matrix, LC) == A + B
    @test @inferred convert(Matrix, LC + LC) ≈ 2A + 2B
    @test @inferred convert(Matrix, M + LC) ≈ 2A + B
    @test @inferred convert(Matrix, M + M) ≈ 2A
    # subtraction
    @test (@inferred Matrix(LC - LC)) ≈ zeros(eltype(LC), size(LC)) atol=10eps()
    @test (@inferred Matrix(LC - M)) ≈ B
    @test (@inferred Matrix(N - LC)) ≈ -A
    @test (@inferred Matrix(M - M)) ≈ zeros(size(M)) atol=10eps()
    # scalar multiplication
    @test @inferred Matrix(-M) == -A
    @test @inferred Matrix(-LC) == -A - B
    @test @inferred Matrix(3 * M) == 3 * A
    @test @inferred Matrix(M * 3) == 3 * A
    @test Matrix(3.0 * LC) ≈ Matrix(LC * 3) ≈ 3A + 3B
    @test @inferred Matrix(3 \ M) ≈ A/3
    @test @inferred Matrix(M / 3) ≈ A/3
    @test @inferred Matrix(3 \ LC) ≈ (A + B) / 3
    @test @inferred Matrix(LC / 3) ≈ (A + B) / 3
    @test @inferred Array(3 * M' - CS!) == 3 * A' - Array(CS!)
    @test (3M - 1im*CS!)' == (3M + ((-1im)*CS!))' == M'*3 + CS!'*1im
    @test @inferred Array(M + A) == 2 * A
    @test @inferred Array(A + M) == 2 * A
    @test @inferred Array(M - A) == 0 * A
    @test Array(A - M) == 0 * A
end


# methods/inplace.jl
#
# Method-based access to tensor operations using inplace definitions.

tensorcopy!(A, IA, C, IC) = tensoradd!(1, A, IA, 0, C, IC)

function tensoradd!(α, A, IA, β, C, IC)
    checkindices(A, IA)
    checkindices(C, IC)
    indCinA = add_indices(IA, IC)
    add!(α, A, Val{:N}, β, C, indCinA)
end

function tensortrace!(α, A, IA, β, C, IC)
    checkindices(A, IA)
    checkindices(C, IC)
    indCinA, cindA1, cindA2 = trace_indices(IA, IC)
    trace!(α, A, Val{:N}, β, C, indCinA, cindA1, cindA2)
    return C
end

function tensorcontract!(α, A, IA, conjA, B, IB, conjB, β, C, IC; method::Symbol = :BLAS)
    # Updates C as β*C+α*contract(A, B), whereby the contraction pattern
    # is specified by IA, IB and IC. The iterables IA(B, C)
    # should contain a unique label for every index of array A(B, C), such that
    # common I of A and B correspond to indices that will be contracted.
    # Common I between A and C or B and C indicate the position of the
    # uncontracted indices of A and B with respect to the indices of C, such
    # that the output array of the contraction can be added to C. Every label
    # should thus appear exactly twice in the union of IA, IB and
    # IC and the associated indices of the tensors should have identical
    # size.
    # Array A and/or B can be also conjugated by setting conjA and/or conjB
    # equal  to 'C' instead of 'N'.
    # The parametric argument method can be specified to choose between two
    # different contraction strategies:
    # -> method = :BLAS : permutes tensors (requires extra memory) and then
    #                   calls built-in (typically BLAS) multiplication
    # -> method = :native : memory-free native julia tensor contraction

    checkindices(A, IA)
    checkindices(B, IB)
    checkindices(C, IC)

    conjA == 'N' || conjA == 'C' || throw(ArgumentError("Value of conjA should be 'N' or 'C' instead of $conjA"))
    conjB == 'N' || conjB == 'C' || throw(ArgumentError("Value of conjB should be 'N' or 'C' instead of $conjB"))
    CA = conjA == 'N' ? :N : :C
    CB = conjB == 'N' ? :N : :C

    oindA, cindA, oindB, cindB, indCinoAB = contract_indices(IA, IB, IC)

    if method == :BLAS
        contract!(α, A, Val{CA}, B, Val{CB}, β, C, oindA, cindA, oindB, cindB, indCinoAB,Val{:BLAS})
    elseif method == :native
        contract!(α, A, Val{CA}, B, Val{CB}, β, C, oindA, cindA, oindB, cindB, indCinoAB,Val{:native})
    else
        throw(ArgumentError("unknown contraction method"))
    end
    return C
end

function tensorproduct!(α, A, IA, B, IB, β, C, IC)
    isempty(intersect(IA, IB)) || throw(LabelError("not a valid tensor product"))
    tensorcontract!(α, A, IA, 'N', B, IB, 'N', β, C, IC; method = :native)
end


# Hamiltonian operator
h(φ::Float64, Δφ::Float64, v::Float64) = -Δφ/2 + v*φ

function h!(dst::AbstractArray, φ::AbstractArray, Δφ::AbstractArray, v::AbstractArray)
    @assert length(dst) == length(φ) == length(Δφ) == length(v)
    for i in 1:length(dst)
        dst[i] = h(φ[i],Δφ[i],v[i])
    end
end

function h!(dst::AbstractMatrix, φ::AbstractMatrix, Δφ::AbstractMatrix, v::AbstractArray)
    @assert size(dst) == size(φ) == size(Δφ)
    @assert size(dst,1) == length(v)
    for i in 1:size(dst,2)
        for j in 1:size(dst,1)
            dst[j,i] = h(φ[j,i],Δφ[j,i],v[j])
        end
    end
end

function h(φ::AbstractArray, Δφ::AbstractArray, v::AbstractArray)
    dst = similar(φ)
    h!(dst,φ,Δφ,v)
    dst
end

function h!(dst::SparseMatrixCSC, φ::SparseMatrixCSC, Δφ::SparseMatrixCSC, v::AbstractArray)
    @assert size(dst) == size(φ) == size(Δφ)
    @assert size(dst,1) == length(v)
    for i in 1:length(φ.nzval)
        j = φ.rowval[i]
        dst.nzval[i] = h(φ.nzval[i],Δφ.nzval[i],v[j])
    end
end

function h(φ::SparseMatrixCSC, Δφ::SparseMatrixCSC, v::AbstractArray)
    dst = SparseMatrixCSC(φ.m,φ.n,φ.colptr,φ.rowval,similar(φ.nzval))
    h!(dst,φ,Δφ,v)
    dst
end


@registerN Meixner AbstractCDOP2 γ μ
export Meixner

"""
    Meixner{γ,μ}

References: [Koekoek and Swarttouw §1.9](https://arxiv.org/pdf/math/9602214.pdf)
"""
Meixner

basis_symbol(::Type{<:Meixner{γ,μ}}) where {γ,μ} = "m⁽ᵞᵝ⁾"
Polynomials.domain(::Type{<:Meixner{γ,μ}}) where {γ,μ} = Polynomials.Interval(0, Inf)
abcde(::Type{<:Meixner{γ,μ}}) where {γ,μ} =
    NamedTuple{(:a, :b, :c, :d, :e)}((0, 1, 0, μ - 1, γ * μ))

function kn(P::Type{<:Meixner{γ,μ}}, n::Int) where {γ,μ}
    (1 - 1 / μ)^n
end

function k1k0(P::Type{<:Meixner{γ,μ}}, n::Int) where {γ,μ}
    (1 - 1 / μ)
end
function k1k_1(P::Type{<:Meixner{γ,μ}}, n::Int) where {γ,μ}
    (1 - 1 / μ)^(n > 0 ? 2 : 1)
end

function classical_hypergeometric(P::Type{<:Meixner{γ,μ}}, n::Int, x) where {γ,μ}
    Pochhammer(γ, n) * pFq((-n, -x), γ, 1 - 1 / μ)
end

# Overrides


#######################################################################
function plotresults(flux, idx, randpcaline, S)
#######################################################################

    figure(2222)

    cla()

    #-------------------------------------------------
    # Plot data
    #-------------------------------------------------
    F = [flux[i .== idx] for i in 1:3]

    plot3D(F[1], F[2], F[3], "bo", label = "observed data in u,g,r")

    x = reduce(hcat, F)

    minT, maxT = Inf, -Inf

    proj(v, x) = dot(v, x) / dot(v, v)

    for _  in 1:S

        # T, s = randgalaxyline()
        #
        # minT, maxT = min(minT, T), max(maxT, T)

        w, μ = randpcaline()

        z = [proj(w, x[i,:] - μ) for i=1:size(x,1)]

        z = sort(z)

        rangeline = maximum(z) - minimum(z)

        zgrid = collect(LinRange(minimum(z)-rangeline*0.0, maximum(z)+rangeline*0.0, 100))

        fwd(z::Real) = w*z + μ

        PCAline = reduce(hcat, fwd.(zgrid))

        @show size(PCAline)

        plot3D(PCAline[1,:], PCAline[2,:], PCAline[3,:], "-")

        # plot3D(gline[1,:], gline[2,:], gline[3,:], "r-")

        # x0 = s * G

        # plot3D([x0[1]],[x0[2]],[x0[3]], "co")

    end

    # title(@sprintf("temperatures range from %f to %f", minT, maxT))

    xlabel("u"); ylabel("g"); zlabel("r")

end


# File: riskreturn.jl
# Author:
# Date: 29-June-2020

# import XL data
[names,~,returns] = getXLData('hfma_matlab_data','10 Hedge Funds');

am1 = fMean(returns,12); # annual mean
am2 = fStd(returns,0,12); # annual std. dev.

# settings for scatter plot
scatter(am2,am1,'b*')
xlabel('Risk (# )');
ylabel('Return (# )');


Base.@kwdef struct LinearScaling{T} <: AbstractParametersScaling{T}
    xrange::Vector{T}
    xmin::Vector{T}
end

Base.@kwdef struct SigmoidScaling{T} <: AbstractParametersScaling{T}
    xrange::Vector{T}
    xmin::Vector{T}
    xsigm::Vector{T}
end

function LinearScaling{T}(df::AbstractDataFrame) where {T}
    df_view = getparams(df, Val(:tanks), Val(:var))
    xrange = @with(df_view, :Max_value .- :Min_value)
    xmin = @with(df_view, @. -:Min_value / xrange)
    LinearScaling{T}(; xrange, xmin)
end

function SigmoidScaling{T}(df::AbstractDataFrame) where {T}
    df_view = getparams(df, Val(:tanks), Val(:var))
    xrange = @with(df_view, :Max_value .- :Min_value)
    xmin = @with(df_view, @. -:Min_value / xrange)
    xsigm = Array{T}(undef, length(xrange))
    SigmoidScaling{T}(; xrange, xmin, xsigm)
end

function scalex!(x, y, scale::LinearScaling{T}) where {T}
    @unpack xrange, xmin = scale
    @turbo for i = 1:length(x)
        x[i] = clamp(xmin[i] + y[i] / xrange[i], zero(T), one(T))
    end
end

function scalex!(x, y, scale::SigmoidScaling{T}) where {T}
    @unpack xrange, xmin = scale
    @turbo for i = 1:length(x)
        zmin = T(0.0001) / xrange[i]
        z = clamp(xmin[i] + y[i] / xrange[i], zmin, one(T) - zmin)
        x[i] = log(z / (one(T) - z))
    end
end

function unscalex!(y, x, scale::LinearScaling)
    @unpack xrange, xmin = scale
    @turbo for i = 1:length(y)
        y[i] = (x[i] - xmin[i]) * xrange[i]
    end
end

function unscalex!(y, x, scale::SigmoidScaling{T}) where {T}
    @unpack xrange, xmin, xsigm = scale
    @turbo for i = 1:length(y)
        xsigm[i] = one(T) / (one(T) + exp(-x[i]))
        y[i] = xrange[i] * (xsigm[i] - xmin[i])
    end
end

function unscaleg!(g, scale::LinearScaling)
    @unpack xrange = scale
    @turbo for i = 1:length(g)
        g[i] *= xrange[i]
    end
end

function unscaleg!(g, scale::SigmoidScaling{T}) where {T}
    @unpack xrange, xsigm = scale
    @turbo for i = 1:length(g)
        g[i] *= xrange[i] * xsigm[i] * (one(T) - xsigm[i])
    end
end


function spikeTriggeredAvg(stimulus, response, windowSize, filter=[1], cond=1)
	spikeIndices = findall(x-> x == cond, conv(response, filter))
	temp = zeros(Int(length(spikeIndices)), windowSize)
	for i in 1:1:length(spikeIndices)
		if spikeIndices[i][1] <= windowSize
			subArr = stimulus[1:spikeIndices[i][1]]
			prepend!(subArr, zeros(windowSize-spikeIndices[i][1]))
			temp[i,:] = subArr
		else
			temp[i,:] = stimulus[spikeIndices[i][1]-(windowSize-1):spikeIndices[i][1]]
		end
	end

	return mean(temp, dims=1), sqrt.(var(temp, dims=1)./size(temp)[1])
end

function spikeTriggeredAvgNDimFilter(stimulus, response, windowSize, filter=[1], cond=1)
	spikeIndices = findall(x-> x == cond, sum(conv(response, filter), dims=2))
	temp = zeros(Int(length(spikeIndices)), windowSize)
	for i in 1:1:length(spikeIndices)
		if spikeIndices[i][1] <= windowSize
			subArr = stimulus[1:spikeIndices[i][1]]
			prepend!(subArr, zeros(windowSize-spikeIndices[i][1]))
			temp[i,:] = subArr
		else
			temp[i,:] = stimulus[spikeIndices[i][1]-(windowSize-1):spikeIndices[i][1]]
		end
	end

	return mean(temp, dims=1), sqrt.(var(temp, dims=1)./size(temp)[1])
end


cd(@__DIR__)
using Pkg
Pkg.activate(".")
##

include("sp_model.jl")
include("plot_utils.jl")

##

using DataFrames
using CSV

timesteps = 1:365*24

pv = CSV.read("timeseries/basic_example.csv", DataFrame)[timesteps, 3]
wind = CSV.read("timeseries/basic_example.csv", DataFrame)[timesteps, 4]
demand = CSV.read("timeseries/basic_example.csv", DataFrame)[timesteps, 2]

##

pars = copy(default_es_pars)

##

using Statistics
average_hourly_demand = mean(demand)

pars[:flexible_demand] = average_hourly_demand * 100. # Over the year, we have 100 demand hours that are shiftable freely...
pars[:recovery_time] = 24
pars[:c_storage] = 100.
pars[:c_pv] = 300.
pars[:c_wind] = 800.
pars[:c_sto_op] = 0.00001

es = define_energy_system(pv, wind, demand; p = pars)

##
n = 100
F_max = average_hourly_demand * 0.1 # Have ancillary services equal to 10% of our typical demand
t_max = length(pv) - es.parameters[2].defaults[:recovery_time]
scens = simple_flex_sampler(n, F_max, t_max)

##

using Cbc

sp = instantiate(es, scens, optimizer = Cbc.Optimizer)

# using GLPK

# sp = instantiate(es, scens, optimizer = GLPK.Optimizer)

##

optimize!(sp)

##

println("Objective value: $(objective_value(sp))")

println("PV: $(round(value.(sp[1, :u_pv]); digits = 2)) -- Wind: $(round(value.(sp[1, :u_wind]); digits = 2)) -- Battery: $(round(value.(sp[1, :u_storage]); digits = 2))")

# for s in 1:minimum((n, 10))
#     println("Second stage objective value in scenario $s: $(objective_value(sp, s))")
# end

##

plot_results(sp, pv, wind, demand, s = 8)

##


"""
    LinearNormalization{T}(a::T,b::T)

This is a functor that normalizes a vector x as ``(x - a) * b``. This is the standard interface for all
linear normalizations such as mapping to ``[0,1]``, ``[-1,1]`` and mean zero standard deviation one.
LinearNormalization also supports the functions Base.lenght and Base.eltype.

See also: [`PosNegNormalization`](@ref), [`GaussianNormalization`](@ref)

# Examples
```jldoctest
julia> nrm = LinearNormalization([0.1, 2.0], [1.0, 0.5]);

julia> x = [1.0, 2.0];

julia> nrm(x)
2-element Vector{Float64}:
 0.9
 0.0

julia> nrm = LinearNormalization(2);  # no scaling to input vector

julia> nrm(x)
2-element Vector{Float64}:
 1.0
 2.0

julia> low = [0.0, -1.0];

julia> high = [3.0, 0.5];

julia> nrm = ZeroOneNormalization(low, high);  # normalize each entry to [0,1]

julia> x = [1.0, 0.0];

julia> feats = nrm(x)
2-element Array{Float64,1}:
0.3333333333333333
0.6666666666666666

julia> y = zero(x);  # create buffer to prevent allocations

julia> feats = nrm(y, x);  # no allocation return

julia> nrm = PosNegNormalization(low, high);  # normalize each entry to [0,1]

julia> nrm(x)
2-element Vector{Float64}:
 -0.3333333333333333
  0.3333333333333333

julia> μ = [0.0, 1.0];  # vector of means

julia> σ = [1.0, 2.0];  # vector of standard deviations

julia> nrm = GaussianNormalization(μ, σ);  # normalize x to be mean 0 and standard deviation 1

julia> nrm(x)
2-element Vector{Float64}:
  1.0
 -0.5
```
"""
struct LinearNormalization{T} <: Any where {T}
    a::T
    b::T

    function LinearNormalization(a::T, b::T) where {T}
        return new{T}(deepcopy(a), deepcopy(b))
    end

    function LinearNormalization(::Type{T}, num_features::Int) where {T}
        a = zeros(T, num_features)
        b = ones(T, num_features)
        return new{typeof(a)}(a,b)
    end

    function LinearNormalization(num_features::Int) where {T}
        LinearNormalization(Float64, num_features)
    end
end

function ZeroOneNormalization(low, high)
    a, b = zero(low), zero(low)
    @. b = one(eltype(low)) / high - low
    @. a = low
    return LinearNormalization(a,b)
end

function ZeroOneNormalization(ranges)
    return ZeroOneNormalization(ranges[:, 1], ranges[:, 2])
end

function PosNegNormalization(low, high)
    a, b = zero(low), zero(low)
    @. a = posneg_norm_a(eltype(low), low, high)
    @. b = posneg_norm_b(eltype(low), low, high)
    return LinearNormalization(a, b)
end

function PosNegNormalization(ranges)
    return PosNegNormalization(ranges[:, 1], ranges[:, 2])
end

function GaussianNormalization(μ, σ)
    a = μ
    b = one(eltype(a)) ./ σ
    return LinearNormalization(a, b)
end

function length(f::LinearNormalization)
    return length(f.a)
end

function eltype(f::LinearNormalization)
    return eltype(f.a)
end

function (f::LinearNormalization)(x)
    y = zero(x)
    @. y = (x - f.a) * f.b
    return y
end

function (f::LinearNormalization)(buff::T, x::T) where {T}
    @. buff = (x - f.a) * f.b
    return buff
end

function update!(f::LinearNormalization, a, b)
    @. f.a = a
    @. f.b = b
    return nothing
end

function update!(f::LinearNormalization, st::Group, mode=:default) where {TS<:Union{KahanVariance, Variance}}
    T = eltype(f)
    for i in 1:length(f)
        f.a[i], f.b[i] = update_linearnorm(T, st[i], mode)
    end
end

function update_linearnorm(::Type{T}, st::TS, mode=:default) where {T, TS<:Union{KahanVariance, Variance}}
    a = convert(T, mean(st))
    b = zero(T)
    if mode == :default
        b = one(T) / convert(T, std(st))
    elseif mode == :var
        b = one(T) / convert(T, var(st))
    else
        throw(DomainError(mode, "Not a valid mode for normalization with $TS"))
    end
    return a, b
end

function update_linearnorm(::Type{T}, st::TS, mode=:default) where {T,TS<:Extrema}
    a,b = zero(T), zero(T)
    if mode == :default
        a = convert(T, minimum(st))
        b = convert(T, maximum(st) - minimum(st))
    elseif mode == :posneg
        a = posneg_norm_a(T, minimum(st), maximum(st))
        b = posneg_norm_b(T, minimum(st), maximum(st))
    else
        throw(DomainError(mode, "Not a valid mode for normalization"))
    end
    return a,b
end

function posneg_norm_a(::Type{T}, low, high) where {T}
    return convert(T, 0.5 * (low + high))
end

function posneg_norm_b(::Type{T}, low, high) where {T}
    return convert(T, 2 / (high - low))
end


# Ordinary metrics

###########################################################
#
#   Metric types
#
###########################################################

type Euclidean <: Metric end
type SqEuclidean <: SemiMetric end
type Chebyshev <: Metric end
type Cityblock <: Metric end

type Minkowski <: Metric
    p::Real
end

type Hamming <: Metric end
result_type(::Hamming, T1::Type, T2::Type) = Int

type CosineDist <: SemiMetric end
type CorrDist <: SemiMetric end

type ChiSqDist <: SemiMetric end
type KLDivergence <: PreMetric end
type JSDivergence <: SemiMetric end

type SpanNormDist <: SemiMetric end

###########################################################
#
#   Specialized distances
#
###########################################################

# SqEuclidean

function sumsqdiff(a::AbstractVector, b::AbstractVector)
    n = get_common_len(a, b)::Int
    s = 0.
    for i = 1:n
        @inbounds s += abs2(a[i] - b[i])
    end
    s
end

evaluate(dist::SqEuclidean, a::AbstractVector, b::AbstractVector) = sumsqdiff(a, b)
sqeuclidean(a::AbstractVector, b::AbstractVector) = evaluate(SqEuclidean(), a, b)
evaluate{T <: Number}(dist::SqEuclidean, a::T, b::T) = abs2(a-b)
sqeuclidean{T <: Number}(a::T, b::T) = evaluate(SqEuclidean(), a, b)

function pairwise!(r::AbstractMatrix, dist::SqEuclidean, a::AbstractMatrix, b::AbstractMatrix)
    m::Int, na::Int, nb::Int = get_pairwise_dims(r, a, b)
    At_mul_B!(r, a, b)
    sa2 = sumsq_percol(a)
    sb2 = sumsq_percol(b)
    for j = 1 : nb
        for i = 1 : na
            @inbounds r[i,j] = sa2[i] + sb2[j] - 2 * r[i,j]
        end
    end
    r
end

function pairwise!(r::AbstractMatrix, dist::SqEuclidean, a::AbstractMatrix)
    m::Int, n::Int = get_pairwise_dims(r, a)
    At_mul_B!(r, a, a)
    sa2 = sumsq_percol(a)
    for j = 1 : n
        for i = 1 : j-1
            @inbounds r[i,j] = r[j,i]
        end
        @inbounds r[j,j] = 0
        for i = j+1 : n
            @inbounds r[i,j] = sa2[i] + sa2[j] - 2 * r[i,j]
        end
    end
    r
end

# Euclidean

evaluate(dist::Euclidean, a::AbstractVector, b::AbstractVector) = sqrt(sumsqdiff(a, b))
euclidean(a::AbstractVector, b::AbstractVector) = evaluate(Euclidean(), a, b)
evaluate{T <: Number}(dist::Euclidean, a::T, b::T) = abs(a-b)
euclidean{T <: Number}(a::T, b::T) = evaluate(Euclidean(), a, b)

function pairwise!(r::AbstractMatrix, dist::Euclidean, a::AbstractMatrix, b::AbstractMatrix)
    m::Int, na::Int, nb::Int = get_pairwise_dims(r, a, b)
    At_mul_B!(r, a, b)
    sa2 = sumsq_percol(a)
    sb2 = sumsq_percol(b)
    for j = 1 : nb
        for i = 1 : na
            @inbounds v = sa2[i] + sb2[j] - 2 * r[i,j]
            @inbounds r[i,j] = isnan(v) ? NaN : sqrt(max(v, 0.))
        end
    end
    r
end

function pairwise!(r::AbstractMatrix, dist::Euclidean, a::AbstractMatrix)
    m::Int, n::Int = get_pairwise_dims(r, a)
    At_mul_B!(r, a, a)
    sa2 = sumsq_percol(a)
    for j = 1 : n
        for i = 1 : j-1
            @inbounds r[i,j] = r[j,i]
        end
        @inbounds r[j,j] = 0
        for i = j+1 : n
            @inbounds v = sa2[i] + sa2[j] - 2 * r[i,j]
            @inbounds r[i,j] = isnan(v) ? NaN : sqrt(max(v, 0.))
        end
    end
    r
end

# Cityblock

function evaluate(dist::Cityblock, a::AbstractVector, b::AbstractVector)
    n = get_common_len(a, b)::Int
    s = 0.
    for i = 1:n
        @inbounds s += abs(a[i] - b[i])
    end
    s
end
cityblock(a::AbstractVector, b::AbstractVector) = evaluate(Cityblock(), a, b)
evaluate{T <: Number}(dist::Cityblock, a::T, b::T) = abs(a-b)
cityblock{T <: Number}(a::T, b::T) = evaluate(Cityblock(), a, b)

# Chebyshev

function evaluate(dist::Chebyshev, a::AbstractVector, b::AbstractVector)
    n = get_common_len(a, b)
    s = 0.
    for i = 1:n
        @inbounds ai = abs(a[i] - b[i])
        if ai > s
            s = ai
        end
    end
    s
end
chebyshev(a::AbstractVector, b::AbstractVector) = evaluate(Chebyshev(), a, b)
evaluate{T <: Number}(dist::Chebyshev, a::T, b::T) = abs(a-b)
chebyshev{T <: Number}(a::T, b::T) = evaluate(Chebyshev(), a, b)

# Minkowski

function evaluate(dist::Minkowski, a::AbstractVector, b::AbstractVector)
    n = get_common_len(a, b)
    s = 0.
    p = dist.p
    for i = 1:n
        @inbounds s += abs(a[i] - b[i]) .^ p
    end
    s .^ inv(p)
end
minkowski(a::AbstractVector, b::AbstractVector, p::Real) = evaluate(Minkowski(p), a, b)
evaluate{T <: Number}(dist::Minkowski, a::T, b::T) = abs(a-b)
minkowski{T <: Number}(a::T, b::T, p::Real) = evaluate(Minkowski(p), a, b)

# Hamming

function evaluate(dist::Hamming, a::AbstractVector, b::AbstractVector)
    n = length(a)
    r = 0
    for i = 1 : n
        @inbounds if a[i] != b[i]
            r += 1
        end
    end
    r
end

hamming(a::AbstractVector, b::AbstractVector) = evaluate(Hamming(), a, b)
evaluate{T <: Number}(dist::Hamming, a::T, b::T) = a != b ? 1 : 0
hamming{T <: Number}(a::T, b::T) = evaluate(Hamming(), a, b)

# Cosine dist

function evaluate{T<:FloatingPoint}(dist::CosineDist, a::AbstractVector{T}, b::AbstractVector{T})
    n = length(a)
    ab = zero(T)
    a2 = zero(T)
    b2 = zero(T)

    for i = 1:n
        @inbounds ai = a[i]
        @inbounds bi = b[i]
        ab += ai * bi
        a2 += ai * ai
        b2 += bi * bi
    end
    max(1 - ab / (sqrt(a2) * sqrt(b2)), 0)
end

cosine_dist(a::AbstractVector, b::AbstractVector) = evaluate(CosineDist(), a, b)

function pairwise!(r::AbstractMatrix, dist::CosineDist, a::AbstractMatrix, b::AbstractMatrix)
    m::Int, na::Int, nb::Int = get_pairwise_dims(r, a, b)
    At_mul_B!(r, a, b)
    ra = sqrt!(sumsq_percol(a))
    rb = sqrt!(sumsq_percol(b))
    for j = 1 : nb
        for i = 1 : na
            @inbounds r[i,j] = max(1 - r[i,j] / (ra[i] * rb[j]), 0)
        end
    end
    r
end

function pairwise!(r::AbstractMatrix, dist::CosineDist, a::AbstractMatrix)
    m, n = get_pairwise_dims(r, a)
    At_mul_B!(r, a, a)
    ra = sqrt!(sumsq_percol(a))
    for j = 1 : n
        for i = j+1 : n
            @inbounds r[i,j] = max(1 - r[i,j] / (ra[i] * ra[j]), 0)
        end
        @inbounds r[j,j] = 0
        for i = 1 : j-1
            @inbounds r[i,j] = r[j,i]
        end
    end
    r
end

# Correlation Dist

_centralize(x::AbstractVector) = x .- mean(x)
_centralize(x::AbstractMatrix) = x .- mean(x, 1)

evaluate(dist::CorrDist, a::AbstractVector, b::AbstractVector) = cosine_dist(_centralize(a), _centralize(b))
corr_dist(a::AbstractVector, b::AbstractVector) = evaluate(CorrDist(), a, b)

function colwise!(r::AbstractArray, dist::CorrDist, a::AbstractMatrix, b::AbstractMatrix)
    colwise!(r, CosineDist(), _centralize(a), _centralize(b))
end

function colwise!(r::AbstractArray, dist::CorrDist, a::AbstractVector, b::AbstractMatrix)
    colwise!(r, CosineDist(), _centralize(a), _centralize(b))
end

function pairwise!(r::AbstractMatrix, dist::CorrDist, a::AbstractMatrix, b::AbstractMatrix)
    pairwise!(r, CosineDist(), _centralize(a), _centralize(b))
end

function pairwise!(r::AbstractMatrix, dist::CorrDist, a::AbstractMatrix)
    pairwise!(r, CosineDist(), _centralize(a))
end

# Chi-square distance

function evaluate{T<:FloatingPoint}(dist::ChiSqDist, a::AbstractVector{T}, b::AbstractVector{T})
    r = zero(T)
    for i = 1 : length(a)
        @inbounds ai = a[i]
        @inbounds bi = b[i]
        r += abs2(ai - bi) / (ai + bi)
    end
    r
end

chisq_dist(a::AbstractVector, b::AbstractVector) = evaluate(ChiSqDist(), a, b)

# KL divergence

function evaluate{T<:FloatingPoint}(dist::KLDivergence, a::AbstractVector{T}, b::AbstractVector{T})
    r = zero(T)
    for i = 1 : length(a)
        @inbounds ai = a[i]
        @inbounds bi = b[i]
        if ai > 0
            r += ai * log(ai / bi)
        end
    end
    r
end

kl_divergence(a::AbstractVector, b::AbstractVector) = evaluate(KLDivergence(), a, b)

# JS divergence

function evaluate{T<:FloatingPoint}(dist::JSDivergence, a::AbstractVector{T}, b::AbstractVector{T})
    r = zero(T)
    n = length(a)
    for i = 1 : n
        @inbounds ai = a[i]
        @inbounds bi = b[i]
        u = (ai + bi) / 2
        ta = ai > 0 ? ai * log(ai) / 2 : 0
        tb = bi > 0 ? bi * log(bi) / 2 : 0
        tu = u > 0 ? u * log(u) : 0
        r += (ta + tb - tu)
    end
    r
end

js_divergence(a::AbstractVector, b::AbstractVector) = evaluate(JSDivergence(), a, b)

# SpanNorm

function evaluate(dist::SpanNormDist, a::AbstractVector, b::AbstractVector)
    max_d = min_d = a[1] - b[1]
    for i = 1 : length(a)
        @inbounds di = a[i] - b[i]
        if di > max_d
            max_d = di
        elseif di < min_d
            min_d = di
        end
    end
    max_d - min_d
end

spannorm_dist(a::AbstractVector, b::AbstractVector) = evaluate(SpanNormDist(), a, b)


# Graph to matrices

using Graphs
using Base.Test

# data

eds = [Edge(1,1,2), Edge(2,1,3), Edge(3,2,4), Edge(4,3,4)]

a0 = fill(false, 4, 4)
a0[1, 2] = a0[1, 3] = a0[2, 4] = a0[3, 4] = true

a0u = copy(a0)
a0u[2, 1] = a0u[3, 1] = a0u[4, 2] = a0u[4, 3] = true

a0_s = sparse(a0)
a0u_s = sparse(a0u)

# graphs

gd_elst = simple_edgelist(4, eds)
gu_elst = simple_edgelist(4, eds; is_directed=false)

gd_ilst = simple_inclist(4)
gu_ilst = simple_inclist(4, is_directed=false)

for e in eds
    add_edge!(gd_ilst, e.source, e.target)
    add_edge!(gu_ilst, e.source, e.target)
end

# adjacency matrix

@test adjacency_matrix(gd_elst) == a0
@test adjacency_matrix(gu_elst) == a0u

@test adjacency_matrix_sparse(gd_elst) == a0_s
@test adjacency_matrix_sparse(gu_elst) == a0u_s

@test adjacency_matrix(gd_ilst) == a0
@test adjacency_matrix(gu_ilst) == a0u

@test adjacency_matrix_sparse(gd_ilst) == a0_s
@test adjacency_matrix_sparse(gu_ilst) == a0u_s

# weight matrix

eweights = [1., 2., 3., 4.]
wm0  = [0. 1. 2. 0.; 0. 0. 0. 3.; 0. 0. 0. 4.; 0. 0. 0. 0.]
wm0u = [0. 1. 2. 0.; 1. 0. 0. 3.; 2. 0. 0. 4.; 0. 3. 4. 0.]
wm0_s = sparse(wm0)
wm0u_s = sparse(wm0u)

@test weight_matrix(gd_elst, eweights) == wm0
@test weight_matrix(gu_elst, eweights) == wm0u

@test weight_matrix_sparse(gd_elst, eweights) == wm0_s
@test weight_matrix_sparse(gu_elst, eweights) == wm0u_s

@test weight_matrix(gd_ilst, eweights) == wm0
@test weight_matrix(gu_ilst, eweights) == wm0u

@test weight_matrix_sparse(gd_ilst, eweights) == wm0_s
@test weight_matrix_sparse(gu_ilst, eweights) == wm0u_s

# distance matrix

dm0  = [0. 1. 2. Inf; Inf 0. Inf 3.; Inf Inf 0. 4.; Inf Inf Inf 0.]
dm0u = [0. 1. 2. Inf; 1. 0. Inf 3.; 2. Inf 0. 4.; Inf 3. 4. 0.]

@test distance_matrix(gd_elst, eweights) == dm0
@test distance_matrix(gu_elst, eweights) == dm0u

@test distance_matrix(gd_ilst, eweights) == dm0
@test distance_matrix(gu_ilst, eweights) == dm0u

# Laplacian matrix

L0 = [2. -1. -1. 0.; -1. 2. 0. -1.; -1. 0. 2. -1.; 0. -1. -1. 2.]
L0_s = sparse(L0)

@test laplacian_matrix(gu_elst) == L0
@test laplacian_matrix(gu_ilst) == L0

@test laplacian_matrix_sparse(gu_elst) == L0_s
@test laplacian_matrix_sparse(gu_ilst) == L0_s

Lw = [3. -1. -2. 0.; -1. 4. 0. -3.; -2. 0. 6. -4.; 0. -3. -4. 7.]
Lw_s = sparse(Lw)

@test laplacian_matrix(gu_elst, eweights) == Lw
@test laplacian_matrix(gu_ilst, eweights) == Lw

@test laplacian_matrix_sparse(gu_elst, eweights) == Lw_s
@test laplacian_matrix_sparse(gu_ilst, eweights) == Lw_s


export Conv1d

doc"""
    Conv1d(T, ksize, inchannel, outchannel, [padding=0, stride=1, dilation=1])

1-dimensional convolution function.

```julia
T = Float32
x = Var(rand(T,10,5))
f = Conv1d(T, 5, 10, 3, padding=2)
y = f(x)
```
"""
mutable struct Conv1d
    ksize::Int
    padding::Int
    stride::Int
    dilation::Int
    W::Var
    b::Var
end

function Conv1d(::Type{T}, ksize::Int, inchannel::Int, outchannel::Int;
    padding=0, stride=1, dilation=1, init_W=Xavier(), init_b=Fill(0)) where T

    W = init_W(T, ksize*inchannel, outchannel)
    b = init_b(T, outchannel)
    Conv1d(ksize, padding, stride, dilation, parameter(W), parameter(b))
end

function (f::Conv1d)(x, dims)
    h = window1d(x, dims, f.ksize, f.padding, f.stride, f.dilation)
    linear(h, f.W, f.b)
end


module GraphingCalculator

using Alexya, Luxor

@init "Graphing Calculator" 800 600

@layout aside(:v, 250)

curves = Vector{Point}()

@create @options begin
    variables = (
        a = (-10:.1:10, 1),
        b = (-10:.1:10, 1),
        c = (-10:.1:10, 1)
    )

    input = @nolabel "f(x) = a*cos(x) + b*sin(c*x)"

    settings = (
        zoom = (1:300, 100),
        x_min = (-50, -100:-1),
        x_max = (50, 1:100),
        points = (10:300, 300)
    )

    style = (
        color = colorant"yellowgreen",
        line_width = (1, 1:10)
    )

    actions = @nolabel (
        plot = Button("Plot", @hexpand),
        delete = Button("Delete", @hexpand, onclicked = () -> empty!(curves))
    )
end

function eventplot()
    ex = Meta.parse(input[])
    block = last(ex.args)

    @eval begin
        $ex
        plot() do x
            a, b, c = getindex.($variables)
            return $block
        end
    end
end

onevent(eventplot, :clicked, actions.plot)

function plot(f::Function)
    global curves = [Point[]]
    i = 1
    xmin = settings.x_min[]
    xmax = settings.x_max[]
    N = settings.points[]
    Δx = (xmax - xmin)/N
    for x in xmin:Δx:xmax
        local segment = curves[i]
        try
            y = -f(x)
            isinf(y) && throw(DivideError())
            push!(segment, Point(x, f(x)))
        catch e
            if !isempty(segment)
                i += 1
                push!(curves, Point[])
            end
        end
    end
end

@use function setup()
    for (key, option) in pairs(settings)
        key ≠ :zoom && onevent(:value_changed, option) do
            eventplot()
        end
    end

    for option in variables
        onevent(eventplot, :value_changed, option)
    end

    eventplot()
end

@use function update(Δt)
    background("black")
    origin()
    scale(settings.zoom[]/10)
    drawaxis()
    sethue(style.color[])
    setline(style.line_width[])
    setopacity(1)
    for curve in filter(c -> length(c) > 0, curves)
        poly(simplify(curve), :stroke)
    end
end

function drawaxis()
    w, h = @width, @height
    sethue("white")
    setline(1)
    setopacity(1)
    fontsize(0.85)

    line(Point(0, -h/2), Point(0, h/2), :stroke)
    line(Point(-w/2, 0), Point(w/2, 0), :stroke)

    xmin = settings.x_min[]
    xmax = settings.x_max[]

    for y in -h÷2:5:h÷2
        y == 0 && continue
        line( Point(-0.5, y), Point(0.5, y), :stroke )
        label(string(-y), :W, Point(0, y); offset = 1)
        setopacity(.1)
        line( Point(xmin, y), Point(xmax, y), :stroke)
    end

    for x in floor.(Int, [xmin:5:xmax...])
        align = x == 0 ? :SW : :S
        line( Point(x, -0.5), Point(x, 0.5), :stroke )
        label(string(x), align, Point(x, 0); offset = 1)
        setopacity(.1)
        line( Point(x, -h/2), Point(x, h/2), :stroke)
    end
end

start()

end


using QuasiGeostrophy
using LinearAlgebra, Test, FFTW, BenchmarkTools, Plots
include(pwd() * "/sandbox/qg_utils.jl") # timestepping and wrappers here
# Domain
Ω = S¹(0, 2π) × S¹(0, 2π)
Nx = Ny = 2^5;
array = Array # to switch between gpu and cpu
fourier_grid = FourierGrid((Nx, Ny), Ω, arraytype = array)
# Fields
fieldnames = ("q¹", "q²", "ψ¹", "ψ²", "f")
q¹, q², ψ¹, ψ², f = create_fields(names = fieldnames, grid = fourier_grid, arraytype = array)
# Operators
∂x, ∂y = create_operators(fourier_grid, arraytype = array)
Δ = ∂x*∂x + ∂y*∂y  # multiplication for GPU purposes
Δ⁴ = -Δ  #* -Δ * -Δ * -Δ  # multiplication for GPU purposes
# Parameters (λ/L)² = 0.01 => λ = 0.1 L
const λ  = 0.01*2π # Rossby deformation radius, √(gravity * depth) / (coriolis force)
const U  = 1e-4    # background velocity in the top layer, - in bottom
const κ  = 1.0     # bottom drag
const Q  = 0.5     # surface temperature flux

Δtᵐᵃˣ = 1e-1 / Nx     # Largest timestep size
ν = 1e-0/maximum(abs.((Δ⁴.op)))/Δtᵐᵃˣ  # largest stable damping
# initialize stream functions in layers
x, y = fourier_grid.grid
ψ¹( 1.0 * (sin.(x) .* sin.(y) .- 1.0 * sin.(y)) )
ψ²( 1.0 *( (cos.(x) .* sin.(y)) .+ 0.1 * sin.(8*y) .+ 1.0 * sin.(y)  ))
# initialize forcing
f(Q * sin.(y))
# stream function to potential vorticity
q¹.data .= (Δ(ψ¹) + 1 /(2*λ^2) * (ψ² - ψ¹)).data
q².data .= (Δ(ψ²) + 1 /(2*λ^2) * (ψ¹ - ψ²)).data
# potential vorticity to stream function
G = inv( (Δ + -1 /(2*λ^2) ) * (Δ + -1 /(2*λ^2) ) + -(1/(2*λ^2))^2 )
ψ¹.data .= G(Δ(q¹) + 1 /(2*λ^2) * (-q² - q¹)).data
ψ².data .= G(Δ(q²) + 1 /(2*λ^2) * (-q¹ - q²)).data

# Impero
@wrapper q¹=q¹ ψ¹=ψ¹ q²=q² ψ²=ψ² f=f
Δ = Operator(Δ); ∂x = Operator(∂x); ∂y = Operator(∂y)
Δ⁴ = Operator(Δ⁴);
G = Operator(nothing, OperatorMetaData(G, "G"))
∂t = Operator(nothing, OperatorMetaData(nothing, "∂t"))
drag = (-2 * κ) * Δ(ψ²)

# create cutoff filter
fmd = FourierOperatorMetaData("Filter")
dx, dy = create_operators(fourier_grid, arraytype = array)
wavenums = sqrt.(abs.((dx^2 + dy^2).op))
maxk = maximum(wavenums)
op = cutoff.(wavenums, 1/2 * maxk) #max k is about \sqrt(2) larger than max kx
cutoff_filter = FourierOperator(op, fmd)
ℱ = Operator(nothing, OperatorMetaData(cutoff_filter, "ℱ"))

# more dissipative QG
@pde_system qg_system = (
    ψ¹ = G(Δ(q¹) + 1/(2*λ^2) * (-q²-q¹)),
    ψ² = G(Δ(q²) + 1/(2*λ^2) * (-q¹-q²)),
    ∂t(q¹) = -ν*Δ⁴(q¹) + (ℱ(∂y(ψ¹))*ℱ(∂x(q¹)) - ℱ(∂x(ψ¹))*ℱ(∂y(q¹))) - (U/λ^2) * ∂x(ψ¹) - U * ∂x(q¹) + f,
    ∂t(q²) = -ν*Δ⁴(q²) + (ℱ(∂y(ψ²))*ℱ(∂x(q²)) - ℱ(∂x(ψ²))*ℱ(∂y(q²))) + (U/λ^2) * ∂x(ψ²) + U * ∂x(q²) - f + drag
)

## Timestepping in qg_utils
vort1 = []
vort2 = []
for i in 1:10000
    # adaptive timestepping
    Δt = 0.3 * cfl(ψ¹, ψ², ∂x, ∂y)
    Δt = minimum([Δt, Δtᵐᵃˣ])
    evolve_system(qg_system, Δt)

    if i%100==0
        println("hi at "* string(i))

        p1 = plot(q¹.data)
        p2 = plot(q².data)
        display(plot(p1,p2))
        sleep(0.01)
        println("hi at "* string(i))
        println("The norm is ")
        println(norm(compute(∂y(ψ¹)).data)/length(ψ¹.data.data))
        println("The norm is ")
        println(norm(compute(∂y(ψ²)).data)/length(ψ².data.data))
        println("The Δt is ")
        println(Δt)
        println("The CFL is ")
        println(Δt/cfl(ψ¹, ψ², ∂x, ∂y) )

        P = q¹.data.metadata.transform.backward
        push!( vort1 , real.(P * q¹.data.data))
        push!( vort2 , real.(P * q².data.data))
    end
end
@save "qg_stuff.jld2" vort1 vort2

plot(q¹.data)
plot(q².data)

spectrum(q¹.data-q².data)
pyplot()
tmp = q¹.data.metadata.transform.backward * (q¹.data.data - q².data.data)
contourf(real.(tmp'), color = :thermometer, linewidth = 0, contourlevlels = 30)
y = sum(real.(tmp'), dims = 2) ./Nx
plot(x, y, label = "horizontal average")
plot!(x,sin.(x) * maximum(y), label = "sin")