####################
####### DEMO #######
####################

## import required packages ##
using Distributions
import LightGraphs
import Gadfly
import Plots
Plots.gr()
using LaTeXStrings

include("rw_smc.jl")
include("utils.jl")
include("rand_rw.jl")
include("gibbs_updates.jl")

# set seed
srand(0)

########################
######### DATA #########
########################

######### SYNTHETIC #########
    syn_data = true
    # set synthetic data parameters
    n_edges_data = 25 # number of edges
    α = 0.25 # parameter for new vertex probability
    λ = 4.0 # paramter for random walk length distribution
    ld = Poisson(λ) # random walk length distribution (automatically shifted by +1 if minimum is zero)
    sb = true # degree-biased distribution for first vertex at each step (false for uniform distribution)

    # sample simple graph (returns edge list)
    g = randomWalkSimpleGraph(n_edges=n_edges_data,alpha_prob=α,length_distribution=ld,sizeBias=sb)

    # convert edge list to adjacency matrix
    A = edgelist2adj(g)
    Gadfly.spy(A)

######### REAL DATA #########
    syn_data = false
    # SJS
    # g = readcsv("data/sjsAdj.csv",Int64)
    # NIPS
    g = readcsv("data/nipsAdj_02_03.csv",Int64)
    # PPI
    # g = readcsv("data/interactomeAdj.csv",Int64)

    n_edges_data = size(g,1)
    # convert edge list to adjacency matrix
    A = edgelist2adj(g)
    Gadfly.spy(A)

######################################
# INTERFACE WITH LIGHTGRAPHS PACKAGE #
######################################

# create LightGraphs.Graph object for use with LightGraphs functions
    lg = LightGraphs.Graph(A)

    # Examples:
    lg_conn = LightGraphs.is_connected(lg)
    lg_diam = LightGraphs.diameter(lg)

#######################################
#### RUN SMC WITH FIXED PARAMETERS ####
#######################################

n_particles = 100

α_fixed = α
λ_fixed = λ

particle_paths,log_marginal = SMC(g,n_particles,α_fixed,λ_fixed,sb)

############################
#### RUN PARTICLE GIBBS ####
############################

# set sampler parameters
n_particles = 100
n_mcmc_iter = 500 # total number of iterations; includes burn-in
n_burn = 100 # burn-in
n_collect = 10 # collect a sample every n_collect iterations
k_trunc = 10*lg_diam # truncation for support of latent K distribution

# set starting values for sampling (can be anything)
α_start = convert(Float64,0.5*(maximum(g)-1)/(n_edges_data-1))
λ_start = convert(Float64,sqrt(lg_diam))

# set prior parameters
a_α = 1.0
b_α = 1.0
a_λ = 1.0
b_λ = 0.25

n_print = 1 # print progress updates every n_print iterations

samples = ParticleGibbs(g,n_particles,
                  a_α,b_α,a_λ,b_λ,
                  α_start,λ_start,sb,k_trunc,
                  n_mcmc_iter,n_burn,n_collect,n_print)

# plot samples
Plots.scatter(samples["lambda"],samples["alpha"],legend=false,
          ylims=[0,1],xlims=[0,15],
          markershape=:circle,markersize=5.5,markeralpha=0.7,markercolor=:grey,markerstrokewidth=0.0,
          xlabel=L"\lambda",ylabel=L"\alpha")
if syn_data
    Plots.scatter!([λ],[α],legend=false,marker=(:star4,12.0,1.0,:black))
end

# plot edge sequence
Plots.plot(samples["edge_sequence"]',legend=false,xlabel="Sampled arrival order",ylabel="Edge index")


"""
Hypercube iterates over a hypercube of radius γ around xᶜ, bounded inside the space {1..sz}ᴺ.

# Examples
```jldoctest
julia> hc= Hypercube((3,3),1,(10,10));

julia> foreach(println, hc)
(2, 2)
(3, 2)
(4, 2)
(2, 3)
(3, 3)
(4, 3)
(2, 4)
(3, 4)
(4, 4)
```
"""
struct Hypercube{N}
  "xᶜ are the center coordinates"
  xᶜ::NTuple{N,Int}
  "γ is the hypercube radius. Points < γ away from xᶜ belong to the hypercube"
  γ::Int
  "sz is the upper bound of the enclosing space"
  sz::NTuple{N,Int}
  indices::CartesianIndices{N}
end
Hypercube(xᶜ,γ,sz)= Hypercube(xᶜ,γ,sz, start(xᶜ,γ,sz))

# Hypercube implements Iterator interface
start(xᶜ,γ,sz)= CartesianIndices(map( (a,b)-> a:b,
                    max.(xᶜ .- γ, 1),
                    min.(xᶜ .+ γ, sz) ))
Base.iterate(hc::Hypercube)= begin
  i= iterate(hc.indices)
  !isnothing(i) ? (i[1].I,i[2]) : nothing
end
Base.iterate(hc::Hypercube{N}, state) where N= begin
  i= Base.iterate(hc.indices, state)::Union{Nothing, Tuple{CartesianIndex{N}, CartesianIndex{N}}}
  !isnothing(i) ? (i[1].I,i[2]) : nothing
end
#Base.eltype(hc::Hypercube)= eltype(hc.indices)
Base.length(hc::Hypercube)= length(hc.indices)
Base.size(hc::Hypercube)= size(hc.indices)

# Good first approximation!
"""
Hypersphere is approximated with a Hypercube for simplicity

See also: [`Hypercube`](@ref)
"""
const Hypersphere{N}= Hypercube{N}

# Create a hypercube that wraps around the edges of the enclosing space
function wrapping_hypercube(xᶜ::NTuple{N,T},radius::Int,dims::NTuple{N,T}) where {N,T}
  xᶜ= convert(NTuple{N,UIntSP},xᶜ)
  dims= convert(NTuple{N,UIntSP},dims)
  wrapping_hypercube(xᶜ,radius,dims, startWrapping(xᶜ,radius,dims))
end

# BUG Wrapping can't work with a single CartesianIndices! It needs 2/4 of them!
#   from UP -> DIM, 0 -> DOWN !!!
startWrapping(xᶜ::NTuple{N,T},radius::T,dims::NTuple{N,T}) where {N,T}= begin
    CartesianIndices( map((a,b)-> a:b,
      wrapUnder.(xᶜ .- Int(radius), dims),   # BUG wrapUnder does smt funky, should return 10
      wrapOver.(xᶜ .+ Int(radius), dims)))
    end

wrapOver(a::T, lim::T) where T= a>lim ? a-lim : a
wrapUnder(a::T, lim::T) where T= a>lim ? lim-(typemax(T)-a)-T(1) : a


function permutation(n1::S, n2::S) where S <: Int

  nn = n1*n2
  p = zeros(S,nn,nn)

  i = 1
  j = 1
  for r = 1:nn
    if j > n2
      j = 1
      i += 1
    end
    c = i+n1*(j-1)
    p[r,c] = 1

    j += 1
  end

  return p

end


using DelimitedFiles

filename = "./log-conv-CNS-1D-smooth-nodewise_CFL.5"
file = readdlm(filename)

num_scheme = 1
Narr = [2;3;4]
Karr = [50;100;200;400]
block_size = 4
num_lines = size(file,1)
num_block = div(num_lines,block_size)
L1_err_dict = Dict()
L2_err_dict = Dict()
for b in 1:num_block
    l = 1+block_size*(b-1)
    N = parse(Int,file[l,3][1])
    K = file[l,6]
    L1_err = file[l+1,4]
    L2_err = file[l+2,4]
    if !haskey(L1_err_dict,(N,K))
        L1_err_dict[(N,K)] = [L1_err]
    else
        push!(L1_err_dict[(N,K)],L1_err)
    end

    if !haskey(L2_err_dict,(N,K))
        L2_err_dict[(N,K)] = [L2_err]
    else
        push!(L2_err_dict[(N,K)],L2_err)
    end
end

num_row = length(Karr)-1
num_col = length(Narr)
conv_table_low_L1 = zeros(num_row,num_col)
conv_table_low_L2 = zeros(num_row,num_col)
conv_table_lim_L1 = zeros(num_row,num_col)
conv_table_lim_L2 = zeros(num_row,num_col)
for i = 1:num_row
    for j = 1:num_col
        K = Karr[i]
        N = Narr[j]
        conv_table_low_L1[i,j] = log2(L1_err_dict[(N,K)][1]/L1_err_dict[(N,2*K)][1])
        conv_table_low_L2[i,j] = log2(L2_err_dict[(N,K)][1]/L2_err_dict[(N,2*K)][1])
        conv_table_lim_L1[i,j] = log2(L1_err_dict[(N,K)][2]/L1_err_dict[(N,2*K)][2])
        conv_table_lim_L2[i,j] = log2(L2_err_dict[(N,K)][2]/L2_err_dict[(N,2*K)][2])
    end
end

err_table = zeros(length(Karr),2*length(Narr))
for i = 1:length(Karr)
    for j = 1:2*length(Narr)
        K = Karr[i]
        N = Narr[div(j-1,2)+1]
        is_L1 = mod1(j,2) == 1
        if is_L1
            err_table[i,j] = L1_err_dict[(N,K)][1]
        else
            err_table[i,j] = L2_err_dict[(N,K)][1]
        end
    end
end

conv_table = [conv_table_low_L1[:,1] conv_table_low_L2[:,1] conv_table_low_L1[:,2] conv_table_low_L2[:,2] conv_table_low_L1[:,3] conv_table_low_L2[:,3]]


@testset "Paths" begin
  grid = RegularGrid{Float64}(100,100)

  for path in [LinearPath(), RandomPath(),
               ShiftedPath(LinearPath(), 0),
               SourcePath(1:3)]
    p = traverse(grid, path)
    @test length(p) == 100*100
  end

  @testset "LinearPath" begin
    p = traverse(grid, LinearPath())
    @test p == 1:100*100
  end

  @testset "RandomPath" begin
    p = traverse(grid, RandomPath())
    @test all(1 .≤ collect(p) .≤ 100*100)
  end

  @testset "SourcePath" begin
    grid = RegularGrid{Float64}(3,3)
    pset = PointSet(coordinates(grid))

    for sdomain in [grid, pset]
      p = traverse(sdomain, SourcePath([1,9]))
      @test collect(p) == [1,9,2,4,6,8,5,3,7]

      p = traverse(sdomain, SourcePath([1]))
      @test collect(p) == [1,2,4,5,3,7,6,8,9]
    end
  end

  @testset "ShiftedPath" begin
    grid = RegularGrid{Float64}(3,3)
    path = LinearPath()
    for offset in [0,1,-1]
      spath = ShiftedPath(path, offset)
      p  = traverse(grid, path)
      sp = traverse(grid, spath)
      @test length(sp) == 9
      @test collect(sp) == circshift(p, -offset)
    end
  end
end


# should be constriants not projection, projection is a function on the contrsrinats
abstract type Constraints end
struct BoxConstraint{T} <: Constraints
    lower::T
    upper::T
end
abstract type Scaling end
struct NoScaling <: Scaling end
scale(S::NoScaling, gx) = gx
struct ConstantPDScaling{T} <: Scaling
    S::T
end
scale(S::ConstantPDScaling, gx) = S.S\gx

function update_xλ(prob::OptimizationProblem, S::Scaling, λ, x, gx)
    bounds = bounds(prob)
    min.(max.(bounds.lower, x .- λ.*scale(S, gx)), bounds.upper)
end
#update_xλ(p::BoxConstraint, λ, x, gx) = min.(max.(p.lower, x .- λ.*gx), p.upper)

# 5.13 in Kelley IMO
# what's it in Bertsekas?
function sufficient_decrease(fλ, fx, λ, α, x, xλ, gx)
    lhs = fλ
    # We use (22) from Bersekas (1976) to avoid dividing by λ
    rhs = fx - dot(gx, x.-xλ)*α
    lhs <= rhs
end
function sufficient_decrease(fλ, fx, λ, α, gx)
    lhs = fλ
    rhs = fx - norm(gx, 2)*α*λ
    lhs <= rhs
end

function projected_gradient(prob::OptimizationProblem, f, g, x0, sc::Union{Scaling, NoScaling}=NoScaling(); itermax, ls_itermax, β=0.5, α=1e-6, print_trace=true)

    f0, g0 = f(x0), g(x0)
    if norm(x0.-update_xλ(prob, sc, 1, x0, g0), Inf) <= 1e-6
      return x0, f0, g0
    end

    x = x0

	# 5.12 in Kelley IMO
	# what's it in Bertsekas?
	local fx, gx
    for i = 1:itermax
        fx, gx = f(x), g(x)
        if print_trace
	        println("i = $i    fx: $(fx)    |gx|: $(norm(gx,Inf))")
            println("x1 = $(x[1])")
            println("x2 = $(x[2])")
	    end
        if norm(x.-update_xλ(prob, sc, 1, x, gx)) <= 1e-6
            return x, fx, gx
        end
        for j = 1:ls_itermax
        	λ = β^j
        	xλ = update_xλ(P, sc, λ, x, gx)
            fλ = f(xλ)
            if sufficient_decrease(fλ, fx, λ, α, x, xλ, gx)
            	x = xλ
            	break
            end
        end
    end
    x, fx, gx
end


module RootedTrees

using LinearAlgebra

import Base: show, isless, ==, iterate, copy

export rootedtree, RootedTreeIterator

export α, β, γ, σ, order, residual_order_condition, elementary_weight, derivative_weight

export count_trees

"""
    RootedTree

Represents a rooted tree using its level sequence.

Reference:
  Beyer, Terry, and Sandra Mitchell Hedetniemi.
  "Constant time generation of rooted trees."
  SIAM Journal on Computing 9.4 (1980): 706-712.
"""
mutable struct RootedTree{T<:Integer, V<:AbstractVector}
  level_sequence::V
  iscanonical::Bool
end

function RootedTree(level_sequence::AbstractVector, iscanonical=false)
  T = eltype(level_sequence)
  V = typeof(level_sequence)
  RootedTree{T,V}(level_sequence, iscanonical)
end

"""
    rootedtree

Construct a canonical `RootedTree` object from a level sequence.

Reference:
  Beyer, Terry, and Sandra Mitchell Hedetniemi.
  "Constant time generation of rooted trees."
  SIAM Journal on Computing 9.4 (1980): 706-712.
"""
rootedtree(level_sequence::AbstractVector) = canonical_representation(RootedTree(level_sequence))

iscanonical(t::RootedTree) = t.iscanonical
#TODO: Validate rooted tree in constructor?

copy(t::RootedTree) = RootedTree(copy(t.level_sequence), t.iscanonical)

#  #function RootedTree(sequence::Vector{T}, valid::Bool)
#  function RootedTree(sequence::Array{T,1})
#    length(sequence) < 1 && throw(ArgumentError("Rooted trees must have a root, in particular at least one element!"))#
#
#    ## If there is only one element, the sequence must be valid.
#    #if !valid && length(sequence) > 1
#    #  # Test, whether there is exactly one root element at the beginning of sequence, if necessary.
#    #  root = sequence[1]
#    #  for level in sequence[2:end]
#    #    level <= root && throw(ArgumentError("Rooted trees must have exactly one element at root-level at the beginning."))
#    #  end
#    #end
#    # If there is only one element, the sequence must be valid.
#    if length(sequence) > 1
#      # Test, whether there is exactly one root element at the beginning of sequence, if necessary.
#      root = sequence[1]
#      for level in sequence[2:end]
#        level <= root && throw(ArgumentError("Rooted trees must have exactly one element at root-level at the beginning."))
#      end
#    end
#
#    new(sequence)
#  end
#RootedTree{T<:Integer}(sequence::Vector{T}) = RootedTree{T}(sequence, false)

function show(io::IO, t::RootedTree{T}) where {T}
  print(io, "RootedTree{", T, "}: ")
  show(io, t.level_sequence)
end

# comparison

"""
    isless(t1::RootedTree, t2::RootedTree)

Compares two rooted trees using a lexicographical comparison of their level sequences.
"""
function isless(t1::RootedTree, t2::RootedTree)
  isless(t1.level_sequence, t2.level_sequence)
end

function ==(t1::RootedTree, t2::RootedTree)
  t1.level_sequence == t2.level_sequence
end

# generation and canonical representation
"""
    canonical_representation!(t::RootedTree)

Change the representation of the rooted tree `t` to the canonical one, i.e. the
one with lexicographically biggest level sequence.
"""
function canonical_representation!(t::RootedTree)
  subtr = subtrees(t)
  for i in eachindex(subtr)
    canonical_representation!(subtr[i])
  end
  sort!(subtr, rev=true)

  i = 2
  for τ in subtr
    t.level_sequence[i:i+order(τ)-1] = τ.level_sequence[:]
    i += order(τ)
  end
  t.iscanonical = true

  t
end

"""
    canonical_representation(t::RootedTree)

Returns a new tree using the canonical representation of the rooted tree `t`,
i.e. the one with lexicographically biggest level sequence.
"""
function canonical_representation(t::RootedTree)
  canonical_representation!(copy(t))
end

"""
    RootedTreeIterator{T<:Integer}

Iterator over all rooted trees of given `order`. The returned trees are views to
an internal tree modified during the iteration. If the returned trees shall be
stored or modified during the iteration, a `copy` has to be made.
"""
struct RootedTreeIterator{T<:Integer}
  order::T
  t::RootedTree{T,Vector{T}}

  function RootedTreeIterator(order::T) where {T<:Integer}
    new{T}(order, RootedTree(Vector{T}(one(T):order), true))
  end
end

function iterate(iter::RootedTreeIterator{T}) where {T}
  iter.t.level_sequence[:] = one(T):iter.order
  (iter.t, false)
end

function iterate(iter::RootedTreeIterator{T}, state) where {T}
  state && return nothing

  two = iter.t.level_sequence[1] + one(T)
  p = 1
  q = 1
  @inbounds for i in 2:length(iter.t.level_sequence)
    if iter.t.level_sequence[i] > two
      p = i
    end
  end
  p == 1 && return nothing

  level_q = iter.t.level_sequence[p] - one(T)
  @inbounds for i in 1:p
    if iter.t.level_sequence[i] == level_q
      q = i
    end
  end

  @inbounds for i in p:length(iter.t.level_sequence)
    iter.t.level_sequence[i] = iter.t.level_sequence[i - (p-q)]
  end

  (iter.t, false)
end

"""
    count_trees(order)

Counts all rooted trees of `order`.
"""
function count_trees(order)
  order < 1 && throw(ArgumentError("The `order` must be at least one."))

  num = 0
  for _ in RootedTreeIterator(order)
    num += 1
  end
  num
end

# subtrees
struct Subtrees{T<:Integer} <: AbstractVector{RootedTree{T}}
  level_sequence::Vector{T}
  indices::Vector{T}
  iscanonical::Bool

  function Subtrees(t::RootedTree{T}) where {T}
    level_sequence = t.level_sequence
    indices = Vector{T}()

    start = 2
    i = 3
    while i <= length(level_sequence)
      if level_sequence[i] <= level_sequence[start]
        push!(indices, start)
        start = i
      end
      i += 1
    end
    push!(indices, start)

    # in order to get the stopping index for the last subtree
    push!(indices, length(level_sequence)+1)

    new{T}(level_sequence, indices, iscanonical(t))
  end
end

Base.size(s::Subtrees) = (length(s.indices)-1, )
Base.getindex(s::Subtrees, i::Int) = RootedTree(view(s.level_sequence, s.indices[i]:s.indices[i+1]-1), s.iscanonical)

"""
    subtrees(t::RootedTree)

Returns a vector of all subtrees of `t`.
"""
function subtrees(t::RootedTree{T}) where {T}
  subtr = typeof(t)[]

  if length(t.level_sequence) < 2
    return subtr
  end

  start = 2
  i = 3
  while i <= length(t.level_sequence)
    if t.level_sequence[i] <= t.level_sequence[start]
      push!(subtr, RootedTree(t.level_sequence[start:i-1]))
      start = i
    end
    i += 1
  end
  push!(subtr, RootedTree(t.level_sequence[start:end]))
end

# functions on trees

"""
    order(t::RootedTree)

The `order` of a rooted tree, i.e. the length of it's level sequence.
"""
order(t::RootedTree) = length(t.level_sequence)

"""
    σ(t::RootedTree)

The symmetry `σ` of a rooted tree `t`, i.e. the order of the group of automorphisms
on a particular labelling (of the vertices) of `t`.

Reference: Section 301 of
  Butcher, John Charles.
  Numerical methods for ordinary differential equations.
  John Wiley & Sons, 2008.
"""
function σ(t::RootedTree)
  if order(t) == 1 || order(t) == 2
    return 1
  end

  if !iscanonical(t)
    t = canonical_representation(t)
  end

  subtr = Subtrees(t)
  sym = 1
  num = 1

  @inbounds for i in 2:length(subtr)
    if subtr[i] == subtr[i-1]
      num += 1
    else
      sym *= factorial(num) * σ(subtr[i-1])^num
      num = 1
    end
  end
  sym *= factorial(num) * σ(subtr[end])^num
end

"""
    γ(t::RootedTree)

The density `γ(t)` of a rooted tree, i.e. the product over all vertices of `t`
of the order of the subtree rooted at that vertex.

Reference: Section 301 of
  Butcher, John Charles.
  Numerical methods for ordinary differential equations.
  John Wiley & Sons, 2008.
"""
function γ(t::RootedTree)
  if order(t) == 1
    return 1
  elseif order(t) == 2
    return 2
  end

  subtr = Subtrees(t)
  den = order(t)
  for τ in subtr
    den *= γ(τ)
  end
  den
end

"""
    α(t::RootedTree)

The number of monotonic labellings of `t` not equivalent under the symmetry group.

Reference: Section 302 of
  Butcher, John Charles.
  Numerical methods for ordinary differential equations.
  John Wiley & Sons, 2008.
"""
function α(t::RootedTree)
  div(factorial(order(t)), σ(t)*γ(t))
end

"""
    β(t::RootedTree)

The total number of labellings of `t` not equivalent under the symmetry group.

Reference: Section 302 of
  Butcher, John Charles.
  Numerical methods for ordinary differential equations.
  John Wiley & Sons, 2008.
"""
function β(t::RootedTree)
  div(factorial(order(t)), σ(t))
end

"""
    elementary_weight(t::RootedTree, A::AbstractMatrix, b::AbstractVector, c::AbstractVector)

Compute the elementary weight Φ(`t`) of `t` for the Butcher coefficients
`A, b, c` of a Runge-Kutta method.

Reference: Section 312 of
  Butcher, John Charles.
  Numerical methods for ordinary differential equations.
  John Wiley & Sons, 2008.
"""
function elementary_weight(t::RootedTree, A::AbstractMatrix, b::AbstractVector, c::AbstractVector)
  T = promote_type(promote_type(eltype(A), eltype(b)), eltype(c))
  elementary_weight(t, Matrix{T}(A), Vector{T}(b), Vector{T}(c))
end

function elementary_weight(t::RootedTree, A::AbstractMatrix{T}, b::AbstractVector{T}, c::AbstractVector{T}) where {T}
  dot(b, derivative_weight(t, A, b, c))
end

"""
    derivative_weight(t::RootedTree, A, b, c)

Compute the derivative weight (ΦᵢD)(`t`) of `t` for the Butcher coefficients
`A, b, c` of a Runge-Kutta method.

Reference: Section 312 of
  Butcher, John Charles.
  Numerical methods for ordinary differential equations.
  John Wiley & Sons, 2008.
"""
function derivative_weight(t::RootedTree, A, b, c)
  if order(t) == 1
    return zero(c) .+ one(eltype(c))
  end

  subtr = Subtrees(t)
  result = A * derivative_weight(subtr[1], A, b, c)
  for i in 2:length(subtr)
    tmp = A * derivative_weight(subtr[i], A, b, c)
    result = result .* tmp
  end
  return result
end

"""
    residual_order_condition(t::RootedTree, A, b, c)

The residual of the order condition
  `(Φ(t) - 1/γ(t)) / σ(t)`
with elementary weight `Φ(t)`, density `γ(t)`, and symmetry `σ(t)` of the
rooted tree `t` for the Runge-Kutta method with Butcher coefficients
`A, b, c`.

Reference: Section 315 of
  Butcher, John Charles.
  Numerical methods for ordinary differential equations.
  John Wiley & Sons, 2008.
"""
function residual_order_condition(t::RootedTree, A, b, c)
  ew = elementary_weight(t, A, b, c)
  T = typeof(ew)

  (ew - one(T) / γ(t)) / σ(t)
end

end # module


const number_of_chebyshevs_to_compile_into_binaries = 20

function next_chebyshev(previous::Sum_Of_Functions, two_previous::Sum_Of_Functions)
    return PE_Function(2.0,0.0,0.0,1) * previous - two_previous
end

first_kind_chebyshevs = Array{Sum_Of_Functions}(undef, number_of_chebyshevs_to_compile_into_binaries)
first_kind_chebyshevs[1] = Sum_Of_Functions([PE_Function(1.0,0.0,0.0,0)])
first_kind_chebyshevs[2] = Sum_Of_Functions([PE_Function(1.0,0.0,0.0,1)])
for i in 3:number_of_chebyshevs_to_compile_into_binaries
    first_kind_chebyshevs[i] = next_chebyshev(first_kind_chebyshevs[i-1], first_kind_chebyshevs[i-2])
end

second_kind_chebyshevs = Array{Sum_Of_Functions}(undef, number_of_chebyshevs_to_compile_into_binaries)
second_kind_chebyshevs[1] = Sum_Of_Functions([PE_Function(1.0,0.0,0.0,0)])
second_kind_chebyshevs[2] = Sum_Of_Functions([PE_Function(2.0,0.0,0.0,1)])
for i in 3:number_of_chebyshevs_to_compile_into_binaries
    second_kind_chebyshevs[i] = next_chebyshev(second_kind_chebyshevs[i-1], second_kind_chebyshevs[i-2])
end

function get_chevyshevs_up_to(num::Integer, first_kind::Bool = true)
    chebyshevs = Array{Sum_Of_Functions}(undef, num)
    if num >= number_of_chebyshevs_to_compile_into_binaries
        if first_kind
            chebyshevs[1:number_of_chebyshevs_to_compile_into_binaries] = first_kind_chebyshevs
        else
            chebyshevs[1:number_of_chebyshevs_to_compile_into_binaries] = second_kind_chebyshevs
        end
        for i in (number_of_chebyshevs_to_compile_into_binaries+1):num
            chebyshevs[i] = next_chebyshev(chebyshevs[i-1], chebyshevs[i-2])
        end
    else
        if first_kind
            chebyshevs[1:num] = first_kind_chebyshevs[1:num]
        else
            chebyshevs[1:num] = second_kind_chebyshevs[1:num]
        end
    end
    return chebyshevs
end

function get_chebyshev(num::Integer, first_kind::Bool = true)
    chebyshevs = get_chevyshev_up_to(num, first_kind)
    return chebyshevs[num]
end


include("dynamics.jl")
using Plots
using ReferenceFrameRotations
#dynamics.jlのテスト
#Runge-Kutta法のテスト

len=63000
a=[zeros(len) zeros(len)]
a[1,2]=1.0
dt=0.01
#理論解
b=[sin.(range(0,(len-1)*dt,step=dt)) cos.(range(0,(len-1)*dt,step=dt))]
#微分する関数
function dif(v::Array{Float64,1},A)
    θ=atan(v[1],v[2])
    cos(θ),-sin(θ)
end
#積分
for  i in 2:len
    a[i,:]=RK4((a[i-1,:]),0,dt)
    a[i,:]./=norm(a[i,:])
end


module TestGenerateJuliaPrograms

using Test

using Catlab, Catlab.Theories
using Catlab.Programs.GenerateJuliaPrograms

# Compilation
#############

# Note: Further tests of compilation are provided by the roundtripping
# integration tests in `TestParseJuliaPrograms`.

W, X, Y, Z = Ob(FreeCartesianCategory, :W, :X, :Y, :Z)
f, g, h = Hom(:f, X, Y), Hom(:g, Y, Z), Hom(:h, Z, W)

@test compile(f) isa Function
@test compile(compose(f,g)) isa Function
@test compile(otimes(f,h)) isa Function
@test compile(compose(f, mcopy(Y), otimes(g,g))) isa Function

# Evaluation
############

ℝ = Ob(FreeCartesianCategory, :ℝ)

plus, cos_hom, sin_hom = Hom(:+, ℝ⊗ℝ, ℝ), Hom(:cos, ℝ, ℝ), Hom(:sin, ℝ, ℝ)
generators = Dict{Symbol,Any}(:+ => +, :cos => cos, :sin => sin)
ev = (args...; kw...) -> evaluate(args...; generators=generators, kw...)

# Generators.
x = collect(range(-2,stop=2,length=50))
@test ev(sin_hom, π/2) == sin(π/2)
@test ev(cos_hom, x; broadcast=true) == cos.(x)

# Composition.
@test ev(compose(cos_hom, sin_hom), 1) == sin(cos(1))
@test ev(compose(cos_hom, sin_hom), x; broadcast=true) == sin.(cos.(x))
@test ev(id(ℝ), 1) == 1

# Symmetric monoidal product.
@test ev(otimes(cos_hom, sin_hom), 1, 2) == (cos(1), sin(2))
@test ev(otimes(cos_hom, sin_hom), x, x; broadcast=true) == (cos.(x), sin.(x))
@test ev(braid(ℝ,ℝ), 1, 2) == (2, 1)

# Copying and deleting.
@test ev(mcopy(ℝ), 1) == (1, 1)
@test ev(mcopy(ℝ⊗ℝ), 1, 2) == (1, 2, 1, 2)
@test ev(delete(ℝ), 1) == nothing
@test ev(delete(ℝ⊗ℝ), 1, 2) == nothing

f = compose(mcopy(ℝ), otimes(cos_hom, sin_hom), plus)
@test ev(f, 1) == cos(1) + sin(1)
@test ev(f, x; broadcast=true) == cos.(x) + sin.(x)

end


module Examples_sep
using LinearAlgebra
using Random

srcDir = dirname(@__FILE__)*"\\"
include(srcDir *"Utils_states.jl")
using .Utils_states

export get_sep_examples

## Separable States
"""
Separable state examples:
"""
function get_sep_example()
    ρ = Dict()
##
    ρ["Ex3-8DNY20"]   = get_ρ_DNY20()

##  CD12
    ρ["T1CD12i"]    = get_ρ_T1CD12(1)
    ρ["T1CD12ii"]   = get_ρ_T1CD12(2)
    # ρ["T1CD12iii"]  = get_ρ_T1CD12(3)

    ρ["T2CD12i"]    = get_ρ_T2CD12(1)
    ρ["T2CD12ii"]   = get_ρ_T2CD12(2)
    # ρ["T2CD12iii"]  = get_ρ_T2CD12(3)
    # ρ["T2CD12iv"]  = get_ρ_T2CD12(4)
    ρ["T2CD12v"]    = get_ρ_T2CD12(5)

    # ρ["Ex26CD12"]   = get_ρ_Ex26CD12() # Wrong shape

## Random batch
 ## Real
    ρ["RANDa1"]  = gen_ρ_RAND(3, 5)
    ρ["RANDa2"]  = gen_ρ_RAND(3, 10)
    ρ["RANDa3"]  = gen_ρ_RAND(3, 15)
    ρ["RANDa4"]  = gen_ρ_RAND(3, 20)
    ρ["RANDa5"]  = gen_ρ_RAND(3, 25)

    ρ["RANDb1"]  = gen_ρ_RAND(4, 5)
    ρ["RANDb2"]  = gen_ρ_RAND(4, 10)
    ρ["RANDb3"]  = gen_ρ_RAND(4, 15)
    ρ["RANDb4"]  = gen_ρ_RAND(4, 20)
    ρ["RANDb5"]  = gen_ρ_RAND(4, 25)
 ## Complex
    ρ["RANDℂa1"]  = gen_ρ_RAND(3,  5,true)
    ρ["RANDℂa2"]  = gen_ρ_RAND(3, 10,true)
    ρ["RANDℂa3"]  = gen_ρ_RAND(3, 15,true)
    ρ["RANDℂa4"]  = gen_ρ_RAND(3, 20,true)
    ρ["RANDℂa5"]  = gen_ρ_RAND(3, 25,true)

    ρ["RANDℂb1"]  = gen_ρ_RAND(4,  5,true)
    ρ["RANDℂb2"]  = gen_ρ_RAND(4, 10,true)
    ρ["RANDℂb3"]  = gen_ρ_RAND(4, 15,true)
    ρ["RANDℂb4"]  = gen_ρ_RAND(4, 20,true)
    ρ["RANDℂb5"]  = gen_ρ_RAND(4, 25,true)

    for key in keys(ρ)
        ρ[key] = Utils_states.maketraceone(ρ[key])
    end
    return ρ
end

## Examples
"""generates a matrix of the form ∑ʳaᵀa⊗bᵀb, with a,b ∈ uniform random entry-wise.  d ∈ {2,3,4} , r ∈ [9]"""
function gen_ρ_RAND(d::Integer, r::Integer, isℂ = false)
    seed = 343
    Random.seed!(seed)
    if isℂ
        a = [randn(1,d) + im*randn(1,d) for i in 1:r ]
        b = [randn(1,d) + im*randn(1,d) for i in 1:r ]
    else
        a = [randn(1,d) for i in 1:r ]
        b = [randn(1,d) for i in 1:r ]
    end
    ρ = zeros(d^2,d^2)
    for i in 1:r
        A = transpose(a[i])*conj(a[i])
        B = transpose(b[i])*conj(b[i])
        ρ = kron(A,B) + ρ
    end
    return ρ
end

""" http://arxiv.org/abs/1210.0111v2"""
function get_ρ_T1CD12(i)
    ρ = Dict()
    # |00><00| + |11><11|
    ρ[1] =  psep(1,1,2) + psep(2,2,2)

    # |00><00| + |11><11| + (|0> + |1>)⊗(|0> + |1>)(<0| + <1|)⊗(<0| + <1|)
    ρ[2] = psep(1,1,2) + psep(2,2,2) + sq(ψ(2))

    return ρ[i]
end

""" http://arxiv.org/abs/1210.0111v2"""
function get_ρ_T2CD12(i)
    ρ = Dict()

    # |00><00| + |11><11| + |12><12|
    ρ[1] = psep(1,1,3) + psep(2,2,3) + psep(2,3,3)

    # |00><00| + |01><01| + |11><11| + |12><12|
    ρ[2] = psep(1,1,3) + psep(1,2,3) + psep(2,2,3) + psep(2,3,3)

    # |00><00| + |01><01| + |02><02| + |11><11| + |12><12|
    ρ[5] = psep(1,1,3) + psep(1,2,3) +  psep(1,3,3) + psep(2,2,3) + psep(2,3,3)

    return ρ[i]
end

""" http://arxiv.org/abs/1210.0111v2"""
function get_ρ_Ex26CD12()

    ρ = [2 0 0 0 0 0
         0 4 2 0 0 2
         0 2 2 1 -1 0
         0 0 1 2 1 -1
         0 0 -1 1 5 1
         0 2 0 0 1 2]

    return ρ
end

"""https://arxiv.org/abs/2011.08132v1"""
function get_ρ_DNY20()
    H_flat = [i₁*j₁ + i₂*j₂ for i₁ in 1:3 for i₂ in 1:3 for j₁ in 1:3   for j₂ in 1:3 ]
    return  reshape(H_flat,9,9)
end

# """ http://arxiv.org/abs/quant-ph/9605038v2 """
# function get_ρ_HHH1(p  = 0.5)
#     # Random.seed!(343)
#     a  = 0.5; b = 0.8; # arbitary possitive numbers
#     ψ₁ = a*Utils_states.ψ(2,2,2) + b*Utils_states.ψ(1,1,2)
#     ψ₂ = a*Utils_states.ψ(2,1,2) + b*Utils_states.ψ(1,2,2)
#     ρ_temp = p*Utils_states.sq(ψ₁) + (1 - p)*Utils_states.sq(ψ₂)
#     return ρ_temp
# end

end


"""
    ParameterizedForcing(func, parameters)

Construct a forcing function with parameters. The forcing function, which is
applied at grid point `i, j, k`, is called with the signature

    `func(i, j, k, grid, clock, state, parameters)`

Example
=======

function cool_forcing_function(i, j, k, grid, clock, state, parameters) =
    return @inbounds - parameters.μ * exp(grid.zC[k] / parameters.λ) * state.velocities.u[i, j, k]
end

parameters = (μ=42, λ=π)

cool_forcing = ParameterizedForcing(cool_forcing_function, parameters)
"""
struct ParameterizedForcing{F, P}
    func :: F
    parameters :: P
end

@inline (F::ParameterizedForcing)(i, j, k, grid, clock, state) =
    F.func(i, j, k, grid, clock, state, F.parameters)


using StatsBase: mode, mean

abstract type TreeNode end

mutable struct ClassificationNode <: TreeNode
    depth::Int
    isleaf::Bool
    regionid::Union{Nothing, Int}
    prediction::Union{Nothing, Int}
    featureindex::Union{Nothing, Int}
    thresh::Union{Nothing, Float64}
    leftnode::Union{Nothing, TreeNode}
    rightnode::Union{Nothing, TreeNode}

    function ClassificationNode(depth)
        new(depth, false, nothing, nothing, nothing, nothing, nothing, nothing)
    end

end

struct TreeStoppingRule
    maxdepth::Int
    minsamplesplit::Int
end

"""
Tree node for regression task.
"""
mutable struct RegressionNode <: TreeNode
    depth::Int
    isleaf::Bool
    regionid::Union{Nothing, Int}
    prediction::Union{Nothing, Float64}
    featureindex::Union{Nothing, Int}
    thresh::Union{Nothing, Float64}
    leftnode::Union{Nothing, TreeNode}
    rightnode::Union{Nothing, TreeNode}

    function RegressionNode(depth)
        new(depth, false, nothing, nothing, nothing, nothing, nothing, nothing)
    end

end

abstract type DecisionTree end

struct ClassificationTree <: DecisionTree
    rootnode::TreeNode
    lossfunc::Function
    stoppingrule::TreeStoppingRule
    numregions::Int

    function ClassificationTree(X, y, lossfunc, stoppingrule)
        rootnode, lastregionid = fittree!(ClassificationNode(1), X, y, lossfunc, stoppingrule)
        new(rootnode, lossfunc, stoppingrule, lastregionid)
    end
end

struct RegressionTree <: DecisionTree
    rootnode::TreeNode
    lossfunc::Function
    stoppingrule::TreeStoppingRule
    numregions::Int

    function RegressionTree(X, y, lossfunc, stoppingrule)
        rootnode, lastregionid = fittree!(RegressionNode(1), X, y, lossfunc, stoppingrule)
        new(rootnode, lossfunc, stoppingrule, lastregionid)
    end
end

function RegressionTree(X, y, stoppingrule)
    return RegressionTree(X, y, squaredloss, stoppingrule)
end

function worthsplitting(y, leftindices, currentdepth, stoppingrule)
    rightindices = .!leftindices
    if (currentdepth > stoppingrule.maxdepth ||
        !checkminsamplesplit(sum(leftindices), sum(rightindices),
                             stoppingrule.minsamplesplit))
        return false
    else
        return true
    end
end

function checkminsamplesplit(leftnum, rightnum, minsamplesplit)
    return leftnum >= minsamplesplit && rightnum >= minsamplesplit
end

function findbestsplit(X, y, lossfunc, minsamplesplit)
    # no split case, loss value is evluated on all observations
    bestfeature = 1
    bestthresh = -Inf
    lossval = lossfunc(y)

    for (featureindex, featurearray) in enumerate(eachcol(X))
        sortedindices = sortperm(featurearray)
        sortedX = X[sortedindices, :]
        sortedy = y[sortedindices]
        for i in 1:(length(sortedy) - 1)
            templossval = lossfunc(sortedy[1:i], sortedy[i+1:end])
            if (templossval < lossval && checkminsamplesplit(i, length(y)-i, minsamplesplit))
                lossval = templossval
                bestfeature = featureindex
                bestthresh = (sortedX[i, featureindex] + sortedX[i+1, featureindex])/2
            end
        end
    end

    return bestfeature, bestthresh
end

function split(X, y, lossfunc, minsamplesplit)
    featureindex, thresh = findbestsplit(X, y, lossfunc, minsamplesplit)
    leftindices = X[:, featureindex] .< thresh
    return featureindex, thresh, leftindices
end

function _buildbranch(node::N, X, y, indices, depth) where N<:TreeNode
    selectedX = X[indices, :]
    selectedy = y[indices]
    node = N(depth)
    return selectedX, selectedy, node
end

function predict!(node::ClassificationNode, y)
    node.prediction = mode(y)
    nothing
end

function predict!(node::RegressionNode, y)
    node.prediction = mean(y)
    nothing
end

function _createyhat(model::RegressionTree, X)
    zeros(Float64, size(X)[1])
end

function _createyhat(model::ClassificationTree, X)
    zeros(Int, size(X)[1])
end

function predict(model::DecisionTree, X)
    yhat = _createyhat(model, X)
    for (i, x) in enumerate(eachrow(X))
        node = model.rootnode
        while !node.isleaf
            if x[node.featureindex] < node.thresh
                node = node.leftnode
            else
                node = node.rightnode
            end
        end
        yhat[i] = node.prediction
    end
    return yhat
end

function fittree!(node::TreeNode, X, y, lossfunc, stoppingrule)
    function fit!(node::TreeNode, X, y)
        featureindex, thresh, leftindices = split(X, y, lossfunc, stoppingrule.minsamplesplit)
        nextdepth = node.depth + 1

        if !worthsplitting(y, leftindices, nextdepth, stoppingrule)
            predict!(node, y)
            regionid += 1
            node.regionid = regionid
            node.isleaf = true
            return node
        else
            node.featureindex = featureindex
            node.thresh = thresh
            rightindices = .!leftindices

            leftX, lefty, leftnode = _buildbranch(node, X, y, leftindices, nextdepth)
            node.leftnode = fit!(leftnode, leftX, lefty)

            rightX, righty, rightnode = _buildbranch(node, X, y, rightindices, nextdepth)
            node.rightnode = fit!(rightnode, rightX, righty)
            return node
        end
    end

    regionid = 0
    fittednode = fit!(node::TreeNode, X, y)
    return fittednode, regionid
end

# Loss functions
function squaredloss(ys...)
    totalloss = 0.0
    for y in ys
        yhat = mean(y)
        totalloss += sum((y.-yhat).^2)
    end
    return totalloss
end

function _giniindex(y)
    num_ones = float(sum(y.==1))
    total_num = length(y)
    p = num_ones / total_num
    return 2*p*(1-p)
end

function giniindex(ys...)
    totalval = 0.0
    totalnum = sum([length(y) for y in ys])
    for y in ys
        n = float(length(y))
        totalval += n / totalnum * _giniindex(y)
    end
    return totalval
end


using IntervalArithmetic
using IntervalConstraintProgramming
using ModelingToolkit

using Test

@testset "Utilities" begin
    @test IntervalConstraintProgramming.unify_variables([:a, :c], [:c, :b]) == ([:a,:b,:c], [1,3], [3,2], [1,0,2], [0,2,1])
end
@testset "BasicContractor" begin
    @variables x y
    C = BasicContractor(x^2 + y^2)

    @test C(-∞..1, IntervalBox(-∞..∞,2)) == IntervalBox(-1..1, -1..1)

    X =IntervalBox(-1..1,2)
    @test C(X) == 0..2

    @test C((1,2)) == 5
end
@testset "Contractor" begin
    x = y = -∞..∞
    X = IntervalBox(x, y)

    C = @contractor x^2 + y^2

    @test C(-∞..1, X) == IntervalBox(-1..1, -1..1)

    X =IntervalBox(-1..1,2)
    @test C(X) == 0..2

    @test C((1,2)) == 5

end

@testset "Contractor without using macro" begin
    @variables x y
    C = Contractor(x^2 + y^2)
    @test  C(-∞..1, IntervalBox(-∞..∞, 2)) == IntervalBox(-1..1, -1..1)

end

@testset "Contractor (with macros and without macros) specifying variables explicitly" begin
    X =IntervalBox(0.5..1.5,3)
    A=-Inf..1

    C1 = @contractor(x+y, [x,y,z])
    @test C1(A,X) == IntervalBox(0.5..0.5, 0.5..0.5, 0.5..1.5)

    C2 = @contractor(y+z, [x,y,z])
    @test C2(A,X) == IntervalBox(0.5..1.5, 0.5..0.5, 0.5..0.5)

    vars = @variables x y z

    C1 = Contractor(vars, x+y)
    @test C1(A,X) == IntervalBox(0.5..0.5, 0.5..0.5, 0.5..1.5)

    C2 = Contractor(vars, y+z)
    @test C2(A,X) == IntervalBox(0.5..1.5, 0.5..0.5, 0.5..0.5)

    C1 = Contractor([x, y, z], x+y)
    @test C1(A,X) == IntervalBox(0.5..0.5, 0.5..0.5, 0.5..1.5)

    C2 = Contractor([x, y, z], y+z)
    @test C2(A,X) == IntervalBox(0.5..1.5, 0.5..0.5, 0.5..0.5)

    vars = @variables x1 x3 x2

    C = Contractor(vars, x1+x2)
    @test C(A, X) == IntervalBox(0.5..0.5, 0.5..1.5, 0.5..0.5)

end

@testset "Contractor is created by function name " begin
  vars = @variables x y
  g(x, y) = x + y
  C = Contractor(vars, g)

  @test C(-Inf..1, IntervalBox(0.5..1.5, 2)) == IntervalBox(0.5..0.5, 2)

end

@testset "Separators" begin
    II = -100..100
    X = IntervalBox(II, II)
    S = @constraint x^2 + y^2 <= 1

    @test typeof(S) <: IntervalConstraintProgramming.ConstraintSeparator

    inner, outer = S(X)
    @test inner == IntervalBox(-1..1, -1..1)
    @test outer == IntervalBox(II, II)

    X = IntervalBox(-∞..∞, -∞..∞)
    inner, outer = S(X)
    @test inner == IntervalBox(-1..1, -1..1)
    @test outer == IntervalBox(-∞..∞, -∞..∞)

end

@testset "Separators without using macros" begin
    II = -100..100
    X = IntervalBox(II, II)
    vars = @variables x y

    S = Separator(vars, x^2 + y^2 < 1)

    inner, outer = S(X)
    @test inner == IntervalBox(-1..1, -1..1)
    @test outer == IntervalBox(II, II)

    X = IntervalBox(-∞..∞, -∞..∞)
    inner, outer = S(X)
    @test inner == IntervalBox(-1..1, -1..1)
    @test outer == IntervalBox(-∞..∞, -∞..∞)

end

@testset "Separator is created by function name " begin
  vars = @variables x y
  g(x, y) = x + y < 1
  S = Separator(vars, g)

  @test S(IntervalBox(0.5..1.5, 2)) == (IntervalBox(0.5..0.5, 2), IntervalBox(0.5..1.5, 2))

end

@testset "pave" begin
    S1a = @constraint(x > 0 , [x, y])
    S1b = @constraint(y > 0 , [x, y])

    S1 = S1a ∩ S1b
    paving = pave(S1, IntervalBox(-3..3, -3..3), 2.0, 0.5)

    @test Set(paving.inner) == Set([IntervalBox(1.5..3, 0..3), IntervalBox(0..1.5, 1.5..3)])
    @test length(paving.inner) == 2
    @test isempty(paving.boundary) == false

    S2 = S1a ∪ S1b
    paving = pave(S2, IntervalBox(-3..3, -3..3), 0.1)
    @test Set(paving.inner) == Set([IntervalBox(-3..0, 0..3), IntervalBox(0..3, -3..3)])
    @test length(paving.inner) == 2
    @test isempty(paving.boundary) == false

    S3 = @constraint x^2 + y^2 <= 1
    X = IntervalBox(-∞..∞, -∞..∞)
    paving = pave(S3, X, 1.0, 0.5)

    @test Set(paving.inner) == Set([IntervalBox(Interval(0.0, 0.5), Interval(0.0, 0.8660254037844386)),
                            IntervalBox(Interval(0.0, 0.5), Interval(-0.8660254037844386, 0.0)),
                            IntervalBox(Interval(-0.5, 0.0), Interval(0.0, 0.8660254037844386)),
                            IntervalBox(Interval(-0.5, 0.0), Interval(-0.8660254037844386, 0.0))])

    @test Set(paving.boundary) == Set([ IntervalBox(Interval(0.5, 1.0), Interval(0.0, 0.8660254037844387)),
                                IntervalBox(Interval(0.0, 0.5), Interval(0.8660254037844386, 1.0)),
                                IntervalBox(Interval(0.5, 1.0), Interval(-0.8660254037844387, 0.0)),
                                IntervalBox(Interval(0.0, 0.5), Interval(-1.0, -0.8660254037844386)),
                                IntervalBox(Interval(-0.5, 0.0), Interval(0.8660254037844386, 1.0)),
                                IntervalBox(Interval(-1.0, -0.5), Interval(0.0, 0.8660254037844387)),
                                IntervalBox(Interval(-0.5, 0.0), Interval(-1.0, -0.8660254037844386)),
                                IntervalBox(Interval(-1.0, -0.5), Interval(-0.8660254037844387, 0.0))])

    @test length(paving.inner) == 4
    @test length(paving.boundary) == 8
end

# @testset "Constants" begin
#     x = y = -∞..∞
#     X = IntervalBox(x, y)

#     a = 3
#     S4 = @constraint x^2 + y^2 - $a <= 0
#     paving = pave(S4, X)

#     @test paving.ϵ == 0.01
#     @test length(paving.inner) == 1532
#     length(paving.boundary) == 1536

# end

@testset "Paving a 3D torus" begin
    S5 = @constraint (3 - sqrt(x^2 + y^2))^2 + z^2 <= 1
    x = y = z = -∞..∞
    X = IntervalBox(x, y, z)

    paving = pave(S5, X, 1.)

    @test typeof(paving) == IntervalConstraintProgramming.Paving{3, Float64}

end

@testset "Volume" begin
    x = 3..5
    @test Vol(x).bounds == 2

    V = Vol(IntervalBox(-1..1.5, 2..3.5))
    @test typeof(V) == IntervalConstraintProgramming.Vol{2, Float64}
    @test V.bounds == 3.75

end

@testset "Functions" begin
    @function f(x) = 4x;
    C1 = @contractor f(x);
    A = IntervalBox(0.5..1);
    x = IntervalBox(0..1);

    @test C1(A, x) == IntervalBox(0.125..0.25)   # x such that 4x ∈ A=[0.5, 1]

    C2 = @constraint f(x) ∈ [0.5, 0.6]
    X = IntervalBox(0..1)

    paving = pave(C2, X)
    @test length(paving.inner) == 2
    @test length(paving.boundary) == 2

    C3 = @constraint f(f(x)) ∈ [0.4, 0.8]
    @test length(paving.inner) == 2
    @test length(paving.boundary) == 2

end

@testset "Nested functions" begin
    @function f(x) = 2x
    @function g(x) = ( a = f(x); a^2 )
    @function g2(x) = ( a = f(f(x)); a^2 )

    C = @contractor g(x)
    C2 = @contractor g2(x)

    A = IntervalBox(0.5..1)
    x = IntervalBox(0..1)

    @test C(A, x) == IntervalBox(sqrt(A[1] / 4))
    @test C2(A, x) == IntervalBox(sqrt(A[1] / 16))

end

@testset "Iterated functions" begin

    @function f(x) = 2x

    A = IntervalBox(0.5..1)
    x = IntervalBox(0..1)

    @function g3(x) = ( a = (f↑2)(x); a^2 )
    C3 = @contractor g3(x)

    @test C3(A, x) == IntervalBox(sqrt(A[1] / 16))
end


##Differentiation and integration

Base.sum{DD}(f::Fun{Laurent{DD}})=fouriersum(domain(f),f.coefficients)
Base.sum{DD}(f::Fun{Fourier{DD}})=fouriersum(domain(f),f.coefficients)

function linesum{DD}(f::Fun{Laurent{DD}})
    d=domain(f)
    if isa(d,Circle)
        sum(setcanonicaldomain(f))*d.radius
    else
        sum(f)
    end
end

linesum{DD<:Circle}(f::Fun{Fourier{DD}})=sum(setcanonicaldomain(f))*d.radius
linesum{DD<:PeriodicInterval}(f::Fun{Fourier{DD}})=sum(f) #TODO: Complex periodic interval

differentiate{DD<:PeriodicInterval}(f::Fun{Taylor{DD}}) = Fun(im*tocanonicalD(f,0)*taylor_diff(f.coefficients),f.space)
differentiate{DD<:PeriodicInterval}(f::Fun{Hardy{false,DD}}) = Fun(im*tocanonicalD(f,0)*hardyfalse_diff(f.coefficients),f.space)
differentiate{DD<:PeriodicInterval}(f::Fun{Laurent{DD}}) = Fun(im*tocanonicalD(f,0)*laurentdiff(f.coefficients),f.space)

differentiate{DD<:PeriodicInterval}(f::Fun{CosSpace{DD}}) = Fun(tocanonicalD(f,0)*cosspacediff(f.coefficients),SinSpace(domain(f)))
differentiate{DD<:PeriodicInterval}(f::Fun{SinSpace{DD}}) = Fun(tocanonicalD(f,0)*sinspacediff(f.coefficients),CosSpace(domain(f)))
differentiate{DD<:PeriodicInterval}(f::Fun{Fourier{DD}}) = Fun(tocanonicalD(f,0)*fourierdiff(f.coefficients),f.space)

differentiate{DD}(f::Fun{Laurent{DD}}) = Derivative(space(f))*f
differentiate{DD}(f::Fun{Fourier{DD}}) = Derivative(space(f))*f

function integrate{D}(f::Fun{Hardy{false,D}})
    if isa(domain(f),Circle) # drop -1 term if zero and try again
        @assert ncoefficients(f)==0 || abs(f.coefficients[1])<100eps()
        integrate(Fun(f,space(f)|(2:∞)))
    else  # Probably periodic itnerval
        Integral(space(f))*f
    end
end

function integrate{D}(f::Fun{Taylor{D}})
    if isa(domain(f),Circle)
        Integral(space(f))*f
    else  # Probably periodic itnerval  drop constant term if zero
        @assert ncoefficients(f)==0 || abs(f.coefficients[1])<100eps()
        Fun(integrate(Fun(f,space(f)|(2:∞))),space(f))
    end
end

function integrate{CS<:CosSpace}(f::Fun{CS})
    if isa(domain(f),Circle)
        error("Integrate not implemented for CosSpace on Circle")
    else  # Probably periodic itnerval, drop constant term if zero
        tol=1E-14 #TODO: smart tolerance.  Here relative is a bit tricky
                  # since this is called by Fourier integrate
        if abs(f.coefficients[1])<tol
            integrate(Fun(f,space(f)|(2:∞)))
        else
            d=domain(f)
            @assert isa(d,PeriodicInterval)
            x=Fun(identity,[first(d),last(d)])
            (f.coefficients[1]*x)⊕integrate(Fun(f,space(f)|(2:∞)))
        end
    end
end

function integrate{SS<:SinSpace}(f::Fun{SS})
    if isa(domain(f),Circle) # drop term containing z^(-1)
        integrate(Fun(f,space(f)|(2:∞)))
    else  # Probably periodic itnerval\
        Integral(space(f))*f
    end
end

#TODO: This is a hack to make sure Fourier maps to Fourier
# we don't have banded differentiate from CosSpace/SinSpace on a circle
for OP in (:differentiate,:integrate)
    @eval begin
        $OP{T,D<:PeriodicInterval}(f::Fun{Fourier{D},T})=$OP(f[2])⊕$OP(f[1])
        $OP{T,D<:Circle}(f::Fun{Fourier{D},T})=$OP(Fun(f,Laurent))
    end
end

fouriersum(d::PeriodicInterval,cfs)=cfs[1].*arclength(d)

function fouriersum{T}(d::Circle,cfs::Vector{T})
    if length(cfs)≥2
        2π*im*cfs[2]*d.radius
    else
        im*zero(T)
    end
end

# O(min(m,n)) Laurent line integral

function linebilinearform{T,D<:Circle}(f::Fun{Laurent{D},T},g::Fun{Laurent{D},T})
    @assert domain(f) == domain(g)
    u,v,mn = f.coefficients,g.coefficients,min(ncoefficients(f),ncoefficients(g))
    if mn > 1
        ret = u[1]*v[1]
        for i=2:2:mn-1
            ret += u[i]*v[i+1] + u[i+1]*v[i]
        end
        return arclength(domain(f))*ret
    elseif mn > 0
        return arclength(domain(f))*u[1]*v[1]
    else
        return zero(T)
    end
end

function bilinearform{T,D<:Circle}(f::Fun{Laurent{D},T},g::Fun{Laurent{D},T})
    @assert domain(f) == domain(g)
    u,v,mn = f.coefficients,g.coefficients,min(ncoefficients(f),ncoefficients(g))
    if mn > 2
        ret = u[1]*v[2] + u[2]*v[1]
        for i=3:2:mn-1
            ret += u[i]*v[i+1] + u[i+1]*v[i]
        end
        return complexlength(domain(f))*ret
    else
        return zero(T)
    end
end


module DiodeClipper

using Interpolations
using NLsolve: nlsolve
using NumericalIntegration: cumul_integrate

export runsim, runsimaaout, runsimaa

# component values
const R = 1e3
const C = 33e-9

# nonlinear function mapping p = x+u to y (depending on used sampling rate)
function nlfunction(fS, p)
    Is = 1e-15
    Vt = 25e-3
    function f!(F, x)
        vD=x[1]
        F[1] = p - (2*R*C*fS + 1)*vD - 2*R*Is*sinh(vD/Vt)
    end
    return nlsolve(f!, [0.0]).zero[1]
end

# precompute antiderivate table for given sampling rate
function makeantiderivtable(fS)
    pmax = 1.5 * (4 + 8R*C*fS) # guesstimate of relevant range
    ps = range(-pmax, pmax, length=1024)
    zs = [nlfunction(fS, p) for p in ps]
    zsint = cumul_integrate(ps, zs)
    return scale(interpolate(zsint, BSpline(Cubic(Line(OnGrid())))), ps)
end

# simulation without any antialiasing
function runsim(fS, u)
    y = similar(u) # preallocate output
    x = 0.0 # initial state
    for n in eachindex(u)
        p = x + u[n]
        y[n] = nlfunction(fS, p)
        x = -x + 4*C*R*fS*y[n]
    end
    return y
end

# simulation with antialiasing of the output only
function runsimaaout(fS, u)
    y = similar(u)
    adtab = makeantiderivtable(fS)
    x = 0.0
    last_p = 0.0
    last_ad = adtab(last_p)
    for n in eachindex(u)
        p = R*x + u[n]
        ad = adtab(p)
        if abs(p-last_p) > 1e-15
            y[n] = (ad - last_ad) / (p - last_p)
        else
            y[n] = nlfunction(fS, 0.5 * (p + last_p))
        end
        last_p = p
        last_ad = ad
        x = -x + 4*C*fS*nlfunction(fS, p)
    end
    return y
end

# simulation with full antialiasing
function runsimaa(fS, u)
    fS = fS/1.5
    y = similar(u)
    adtab = makeantiderivtable(fS)
    x = 0.0
    last_p = 0.0
    last_ad = adtab(last_p)
    prev_x = 0.0
    for n in eachindex(u)
        p = R*x + u[n]
        ad = adtab(p)
        if abs(p-last_p) > 1e-15
            y[n] = (ad - last_ad) / (p - last_p)
        else
            y[n] = nlfunction(fS, 0.5 * (p + last_p))
        end
        last_p = p
        last_ad = ad
        prev_x, x = x, -0.5*(x + prev_x) + 4fS*C*y[n]
    end
    return y
end

end


# ----------------------------------------------------------------------------------- #
#  estimate_steady_state: estimates the steady-state vector for the dynamic system
#
# Type: GRN-JULIA
# Version: 1.0
#
# Input arguments:
# problem  - dictionary instance (holds model parameters)
#
# Return arguments:
# XSS - steady state vector
# ----------------------------------------------------------------------------------- #
function estimate_steady_state(problem::Dict{String,Any})

    # Get the initial conditions -
    initial_condition_array = problem["initial_condition_array"];

    # calculate the steady-state soln -
    steady_state_prob = SteadyStateProblem(Balances, initial_condition_array, problem)
    steady_state_soln  = solve(steady_state_prob, DynamicSS(AutoTsit5(Rosenbrock23(autodiff=false));abstol=1e-8,reltol=1e-6,tspan=Inf))

    # return -
    return abs.(steady_state_soln.u)
end


lines = readlines("inputs/6-10/09.txt")
nums = [[parse(Int, x) for x in line] for line in lines]

# gets position (r,c) of 2D vec grid, or 9 if out of bounds
function getat(r,c,grid,fallback=9)
  h = length(grid)
  w = length(grid[1])
  if r < 1 || c < 1 || r > h || c > w
    return fallback
  end
  return nums[r][c]
end

function part1(nums)
  t = 0
  for r in 1:length(nums), c in 1:length(nums[1])
    val = nums[r][c]
    if val < getat(r-1,c, nums) &&  val < getat(r+1,c, nums) &&  val < getat(r,c-1, nums) &&  val < getat(r,c+1, nums)
      t += val + 1
    end
  end
  t
end

# unique identifier for an x,y pair
f(x, y) = (100 * x) + y

function part2(nums)
  g = SimpleGraph(10000)

  for r in 1:length(nums)
    for c in 1:length(nums[1])
      if nums[r][c] == 9
        continue
      end
      for (dc, dr) in [(0, -1), (0, 1), (-1, 0), (1, 0)]
        # gets nums[r + dr][c + dc] or 9 if out of bounds
        if getat(r + dr, c + dc, nums, 9)  != 9
          add_edge!(g, f(r, c), f(r + dr, c + dc))
        end
      end
    end
  end

  lens = [length(cc) for cc in connected_components(g)] |> sort
  lens[end] * lens[end-1] * lens[end-2]
end


function SeisSmooth1(A;L=7)
	a = isodd(L) ? 0 : 1
	kernel = ones(eltype(A),L)/L;
	A = conv(A,kernel)
	A = A[floor(Int,L/2) + 1:end-floor(Int,L/2) + a]
	return A
end

function SeisSmooth1(A,h::Array{Header,1};L=7)
	for itrace = 1 : size(A[:,:],2)
		A[:,itrace] = SeisSmooth1(A[:,itrace],L=L)
	end
	return A,h
end

function SeisSmooth1(in::AbstractString,out::AbstractString;L=7)
	@compat parameters = Dict(:L=>L)
	SeisProcess(in,out,[SeisSmooth1],[parameters];group="some",ntrace=100000)
end


using NNlib

# Activation functions
@cufunc softplus(x::Real) = ifelse(x > 0, x + log1p(exp(-x)), log1p(exp(x)))

@cufunc logσ(x::Real) = -softplus(-x)

@cufunc function gelu(x::Real)
    p = oftype(x / 1, π)
    λ = oftype(x / 1, √(2 / p))
    α = oftype(x / 1, 0.044715)
    h = oftype(x / 1, 0.5)
    h * x * (one(x) + tanh(λ * (x + α * x^3)))
end

@cufunc lisht(x::Real) = x * tanh(x)

@cufunc logcosh(x::Real) = x + softplus(-2x) - log(oftype(x, 2))

@cufunc mish(x::Real) = x * tanh(softplus(x))

@cufunc tanhshrink(x::Real) = x - tanh(x)

# Batched matrix multiplication
# 1st argument is produced by NNlib.storage_type(A)
NNlib._batched_gemm!(::Type{<:CuArray}, transA::Char, transB::Char, α::Number, A, B, β::Number, C) =
     CUBLAS.gemm_strided_batched!(transA, transB, α, A, B, β, C)

Base.unsafe_convert(::Type{CuPtr{T}}, A::NNlib.BatchedAdjOrTrans{T}) where {T} =
    Base.unsafe_convert(CuPtr{T}, parent(A))


struct Discretization{DecompositionType}
  lu_decomposition::DecompositionType   # LU decomposition of system matrix
  num_steps_year::Int64                 # Number of time steps per astronomical year
  mesh::Mesh
  model::Model                          # Physical model
  annual_temperature::Array{Float64, 2} # Annual temperature at every time step in an astronomical year
  rhs::Vector{Float64}
  last_rhs::Vector{Float64}
end

function Discretization(mesh, model, num_steps_year; run_garbage_collector = true)
    lu_decomposition = compute_lu_matrices(mesh, model, num_steps_year)

    if run_garbage_collector
      GC.gc()
    end

    annual_temperature = fill(5.0, mesh.dof, num_steps_year) # Magic initialization
    rhs     = zeros(mesh.dof)  # TODO: The EBM Fortran code initializes the RHS to zero... Maybe we want to initialize it differently
    last_rhs = zeros(mesh.dof)

    Discretization(lu_decomposition, num_steps_year, mesh, model, annual_temperature, rhs, last_rhs)
end

Base.size(discretization::Discretization) = size(discretization.mesh)

function Base.show(io::IO, discretization::Discretization)
  nx, ny = size(discretization)
  print(io, "Discretization() with ", nx, "×", ny, " degrees of freedom")
end

function compute_lu_matrices(mesh, model, num_steps_year)
  mat = compute_matrix(mesh,num_steps_year,model)
  lu_decomposition = lu(mat)
  return lu_decomposition
end


@with_gap @testset "Lattices for symmetric group" begin
  # Look at the hook partition [2, 1^(n-1)] of S_{n + 1}.
  # Will have dimension n

  # n = 4
  n = 4
  g1 = matrix(QQ, 4, 4, [ 0, -1, 0, 0, 0, 0, -1, 0, 0, 0, 0, -1, 1, -1, 1, -1 ])
  g2 = matrix(QQ, 4, 4, [ -1, 0, 0, 0, 0, -1, 0, 0, 0, 0, -1, 0, -1, 1, -1, 1 ])
  A = matrix_algebra(QQ, 4)
  O = Order(A, [g1, g2])
  L = Hecke.natural_lattice(O)
  for p in [2, 3, 5, 7]
    S = Hecke.sublattice_classes(L, p)
    @test length(S) == valuation(4 + 1, p) + 1
  end

  @test_throws ErrorException isirreducible(L.V)

  n = 5
  g1 = matrix(QQ, 5, 5, [ 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 1, -1, 1, -1, 1])
  g2 = matrix(QQ, 5, 5, [ -1, 0, 0, 0, 0, 0, -1, 0, 0, 0, 0, 0, -1, 0, 0, 0, 0, 0, -1, 0, 1, -1, 1, -1, 1 ])
  A = matrix_algebra(QQ, 5)
  O = Order(A, [g1, g2])
  L = Hecke.natural_lattice(O)
  for p in [2, 3, 5, 7]
    S = Hecke.sublattice_classes(L, p)
    @test length(S) == valuation(5 + 1, p) + 1
  end
end

@with_gap @testset "Lattices for matrix order" begin
  # provided by Gabi Nebe
  A = matrix(QQ, 3, 3, [1,3^4,0,0,1,0,0,0,1]);
  B = matrix(QQ, 3, 3, [1,0,0,27,1,0,0,0,1]);
  C = matrix(QQ, 3, 3, [125,3^4,0,3*18,125-90,0,0,0,1]);
  D = matrix(QQ, 3, 3, [1,0,0,0,1,0,0,0,-1]);
  E = matrix(QQ, 3, 3, [1,0,3^4,0,1,0,0,0,1]);
  F = matrix(QQ, 3, 3, [1,0,0,0,1,3^5,0,0,1]);
  G = matrix(QQ, 3, 3, [1,0,0,0,1,0,3^3,0,1]);
  H = matrix(QQ, 3, 3, [1,0,0,0,1,0,0,9,1]);
  l = [A, B, C, D, E, F, G, H]
  A = matrix_algebra(QQ, 3)
  O = Order(A, l)
  L = Hecke.natural_lattice(O)
  S = Hecke.sublattices(L, 3)
  @time length(S) == 216
end


using Mamba
using Distributions

include("../src/simulation.jl")
include("../gali_bayesian.jl")
include("../hmc/hmc_aux.jl")

yy,shocks = simulate_dsge(GAMMA_0,GAMMA_1,PSI,PI,500)

mm = Model(
    y = Stochastic(1,
    (mu)-> MvNormal(mu,1),
    false
    ),

    mu = Logical(1,
    (par,x) -> dsge_fit(par,x),
    false
    ),

    # alfa = Stochastic(
    # ()-> Beta(3,3),
    # false
    # ),

    bet = Stochastic(
    ()-> Beta(4,2),
    false
    ),

    epsilon = Stochastic(
    () -> Gamma(7,1),
    false
    ),

    sig = Stochastic(
    () -> Gamma(2,1),
    false
    ),

    phi = Stochastic(
    () -> Gamma(2,1),
    false
    ),

    phi_pi = Stochastic(
    () -> Gamma(2,1),
    false
    ),

    phi_y = Stochastic(
    () -> Gamma(2,1),
    false
    ),

    rho_v = Stochastic(
    () -> Beta(1,1),
    false
    ),

    theta = Stochastic(
    () -> Beta(3,3),
    false
    ),

    # s2 = Stochastic(
    # () -> Gamma(2,2),
    # false
    # ),

    par = Logical(1,
    (bet,epsilon,theta,sig,phi,phi_pi,phi_y,rho_v) -> [2/3,bet,epsilon,theta,sig,1,phi,phi_pi,phi_y,rho_v]#() -> [0.3,0.99,6,2/3,2,1,2,1.5,1.5,0.4]#
    )
)

sampling1 = [NUTS([:par])]

sampling2 = [RWM([:par],10)]

dados = Dict{Symbol, Any}(:y => yy[1:499,2], :x => yy[:,2])

inits = [
    Dict(:y => dados[:y],
    #:alfa => rand(Beta(3,3)),
    :bet => rand(Beta(4,2)),
    :epsilon => rand(Gamma(7,1)),
    :sig => rand(Gamma(2,1)),
    :phi => rand(Gamma(2,1)),
    :phi_pi => rand(Gamma(2,1)),
    :phi_y => rand(Gamma(2,1)),
    :rho_v => rand(Beta(1,1)),
    :theta => rand(Beta(3,3)),
    #:s2 => rand(Gamma(2,2))
    )
    for i in 1:2
]

setsamplers!(mm,sampling1)

sim1 = mcmc(mm,dados,inits,5000, burnin=1000, thin=2, chains=2)

describe(sim1)

p = Gadfly.plot(sim1)

Gadfly.draw(p, filename = "plots.svg")

pp = [2/3,inits[1][:bet],inits[1][:epsilon],inits[1][:theta],inits[1][:sig],inits[1][:s2],inits[1][:phi],inits[1][:phi_pi],inits[1][:phi_y],inits[1][:rho_v]]

pp[6] = 0.05

dsge_fit(pp,yy[:,2])


export SpiralTrajectory, spiralNodes, spiralDensity

"""
    SpiralTrajectory(numProfiles, numSamplingPerProfile
                  ; TE::Float64=0.0
                  , AQ::Float64=1.e-3
                  , windings::Real= 6.25
                  , angleOffset= :equispaced
                  , kargs...)

returns a 2d spiral trajectory.

# Arguments
* `numProfiles::Int64`            - number of profiles
* `numSamplingPerProfile::Int64`  - number of sampling points per profile
* (`TE::Float64=0.0`)             - echo time in s
* (`AQ::Float64=1.e-3`)           - readout duration in s (per profile)
* (`windings::Real= 6.25`)        - number of windings of the spiral profiles
* (`angleOffset= :equispaced`)    - spacing of profile angles (`:equispaced` sampling, `:golden` angle sampling or `:random` sampling)
"""
function SpiralTrajectory(::Type{T}, numProfiles, numSamplingPerProfile
                  ; TE=0.0
                  , AQ=1.e-3
                  , windings::Real= 6.25
                  , kmax::Real=0.5
                  , angleOffset::String="equispaced"
                  , kargs...) where T
  nodes = spiralNodes(T, numProfiles, numSamplingPerProfile; windings=windings,
                      angleOffset=angleOffset, kmax=kmax)
  times = readoutTimes(T, numProfiles, numSamplingPerProfile; TE=TE, AQ=AQ)
  return  Trajectory("Spiral", nodes, times, TE, AQ, numProfiles, numSamplingPerProfile, 1, false, true)
end

function spiralNodes(::Type{T}, numProfiles::Int64
                            , numSamplingPerProfile::Int64
                            ; windings::Real= 6.25
                            , kmax::Real=0.5
                            , angleOffset::String="equispaced"
                            , kargs...) where T

  nodes = zeros(T, 2, numSamplingPerProfile, numProfiles)
  if angleOffset == "golden"
      angles = [i*(3-sqrt(5))/2  for i=0:numProfiles-1 ]
  elseif angleOffset == "random"
      angles = sort(rand(numProfiles))
  elseif angleOffset == "equispaced"
      angles = collect((0:numProfiles-1)/numProfiles)
  end

  w = windings # 8/64 * 50 = 6.25 which means we have 6.25 turns aka windings
  for l = 1:numProfiles
    for k = 1:numSamplingPerProfile
      # For the special case that τ(t)=t, the amount of time spent for each
      # winding is constant, regardless of whether the acquired
      # winding is near the center or in the outer part of the spiral. I
      # other words, the readout gradients reach their maximum
      # performance at the end of the acquisition.
      t = sqrt((k-1)/(numSamplingPerProfile-1)) #
      nodes[1,k,l] = kmax*t*cos(2*pi*( w*t + angles[l] ))
      nodes[2,k,l] = kmax*t*sin(2*pi*( w*t + angles[l] ))
    end
  end
  return reshape(nodes, 2, numSamplingPerProfile*numProfiles)
end

"""
Calculating the density using VoronoiCells
"""
function spiralDensity(::Type{T}, numProfiles::Int64, numSamplingPerProfile::Int64) where T
  @error "Not implemented!"
end


using RedBlackGraph
using Test

function test_floyd_warshall(T::Type)
	r = red_one(T)

	M = [
	r 2 3 0 0;
	0 r 0 2 0;
	0 0 1 0 0;
	0 0 0 r 0;
	2 0 0 0 1
	]
    W = floyd_warshall_transitive_closure(M)
    W_expected = [
	r 2 3 4 0;
	0 r 0 2 0;
	0 0 1 0 0;
	0 0 0 r 0;
	2 4 5 8 1
	]
	@test M == [
	r 2 3 0 0;
	0 r 0 2 0;
	0 0 1 0 0;
	0 0 0 r 0;
	2 0 0 0 1
	]
	@test W == W_expected

	W = floyd_warshall_transitive_closure!(M)
	@test M == W_expected
	@test W == W_expected

end

@testset "shortestpaths" begin
	test_floyd_warshall(AInt8)
	test_floyd_warshall(AInt16)
	test_floyd_warshall(AInt32)
	test_floyd_warshall(AInt64)
	test_floyd_warshall(AInt128)
end


function CJS_empirical(F::Matrix,ndim::Int64,nsamples::Int64,nsteps::Int64)

    results = zeros(ndim,ndim)

    for k in 1:ndim
        for l in k:ndim
            lbound = minimum(F[:,k])
            rbound = maximum(F[:,k])

            #print(edf(F[:,k],rbound))

            int=0.0

            for u in 1:nsteps

                x = lbound+(u-1)*(rbound-lbound)/nsteps
                p=edf(F[:,k],x)
                q=edf(F[:,l],x)

                int+=(p*log2(2*p/(p+q))+1/2 * 1/log(2) * (q-p))/nsteps

            end

            results[k,l] = int

        end
    end

    return(results+results')
end

function edf(data::Array{Float64},x::Float64)
    return(length(data[data .<= x])/length(data))
end


using BenchmarkTools, AutoGrad

function bench_tr_mul_yaad(x1, x2)
    z = tr(x1 * x2)
    backward(z)
    x1.grad, x2.grad
end

function bench_tr_mul_autograd(x1, x2)
    z = @diff tr(x1 * x2)
    grad(z, x1), grad(z, x2)
end

x = Variable(rand(30, 30))
y = Variable(rand(30, 30))

@benchmark bench_tr_mul_yaad(x, y)

x = Param(rand(30, 30))
y = Param(rand(30, 30))

@benchmark bench_tr_mul_autograd(x, y)


n = 10000000
randnstuple = [(x, y) for (x, y) in zip(rand(n), rand(n))]

num_inside = length(filter(x -> x[1]^2 + x[2]^2 <= 1, randnstuple))
println(num_inside / n * 4)


# A maximally fast upper branch incomplete gamma function when the first
# argument is a non-positive integer.
#
# (_gamma_upper_negative_integer)
@inline function _g_u_n_i(s::Int64, x, g0, expnx)
  n    = -s
  ser  = zero(x)
  _x   = one(x)
  _sgn = copy(_x)
  facn = convert(typeof(x), factorial(n))
  fac  = facn/n
  for k in 0:(n-1)
    term  = _sgn*fac*_x
    ser  += term
    _x   *= x
    _sgn *= -one(x)
    fac  /= (n-k-1)
  end
  ((expnx/_x)*ser + _sgn*g0)/facn
end

# The best speed I was able to accomplish with this thing was just to
# compartmentalize it into its own function. Definitely not ideal, but so it is.
#
# TODO (cg 2021/10/26 15:58): get rid of all the remaining factorial calls so
# that this won't literally break for order greater than, like, 19.
function exponential_improvement(v, x, l, m)
  onex = one(x)
  twox = onex+onex
  fv     = 4*v*v
  ser    = zero(x)
  _z     = x
  floatj = onex
  ak_num = fv - floatj
  factj  = onex
  twofloatj = onex
  eightj  = 8
  expx    = exp(2*x)
  expnx   = exp(-2*x)
  expintx = -expinti(-2*x)
  ser     = expx*factorial(l-1)*_g_u_n_i(1-l, 2*x, expintx, expnx)/(2*pi)
  for j in 1:(m-1)
    # add to the series:
    s   = Int(l-j)
    _g  = expx*factorial(Int(s-1))*_g_u_n_i(1-s, 2*x, expintx, expnx)/(2*pi)
    term       = ak_num/(factj*_z*eightj)*_g
    ser       += term
    # update ak and _z:
    floatj    += onex
    twofloatj += twox
    factj     *= floatj
    ak_num    *= (fv - twofloatj^2)
    _z        *= x
    eightj    *= 8
  end
  _sgn = isodd(l) ? -one(x) : one(x)
  ser*_sgn*twox*cospi(v)
end

function _besselk_as(v, x, order, use_remainder=true, modify=false)
  onex = one(x)
  twox = onex+onex
  fv     = 4*v*v
  ser    = zero(x)
  _z     = x
  ser    = onex #zero(x) #onex/_z
  floatj = onex
  ak_num = fv - floatj
  factj  = onex
  twofloatj = onex
  eightj = 8
  for j in 1:order
    # add to the series:
    term       = ak_num/(factj*_z*eightj)
    ser       += term
    # update ak and _z:
    floatj    += onex
    twofloatj += twox
    factj     *= floatj
    ak_num    *= (fv - twofloatj^2)
    _z        *= x
    eightj    *= 8
  end
  if use_remainder
    _rem = exponential_improvement(v, x, Int(order+1), Int(order))
    ser += _rem
  end
  # if you're modifying as (x^v)*besselk(v,x), since the series part is pretty
  # stable numerically, what we want to deal with is the (x^v)*exp(-x). That's
  # the problem of potentially huge*tiny.
  if modify
    mulval = exp(v*log(x)-x)
  else
    mulval = exp(-x)
  end
  return sqrt(pi/(x*twox))*mulval*ser
end


export PDataMA57, AInv

mutable struct PDataMA57{T} <: PDataFact{T}
    L::SparseMatrixCSC{T,Int64} # sparse L
    D::SparseMatrixCSC{T,Int64} # block diagonal 1X1 and 2X2
    #P::Array{T,2}  # permutation matrix: should be vector
    # LDL' = PHP'
    pp::Array{Int64,1}   # permutation vector LDL' = H[pp,pp]
    # P*v = v[pp]
    # P'*v = v[invperm(pp)]
    s::Array{T,1}  # Scaling matrix S=diagm(s)
    Δ::Array{T,1}  # diagonal, eigenvalues of D
    Q::SparseMatrixCSC{T,Int64}  # orthogonal matrix, eigenvectors of D:  should be sparse
    # QΔQ'  =  D
    g̃::Array{T,1}  # transformed gradient
    λ::T
    success::Bool        # previous iteration was successfull
    OK::Bool             # preprocess success

    PDataMA57() = new{Nothing}()
    PDataMA57(L, D, pp, s, Δ, Q, g̃, l, success, OK) =
        new{eltype(g̃)}(L, D, pp, s, Δ, Q, g̃, l, success, OK)
end

function preprocess(PData::PDataMA57, H, g, n1, n2)
    M = Ma57
    L = SparseMatrixCSC{Float64,Int64}
    D57 = SparseMatrixCSC{Float64,Int64}
    pp = Array{Int64,1}
    s = Array{Float64}
    ρ = Float64
    ncomp = Int64

    H57 = convert(SparseMatrixCSC{Float64,Int64}, H)  #  Hard coded Int
    try
        M = Ma57(H57)#, print_level = -1)
        ma57_factorize(M)
    catch
        println("*******   Problem in MA57_0")
        M = Ma57(H57, print_level = -1)
        ma57_factorize(M)
        res = PDataMA57()
        res.OK = false
        return res
    end
    ###### Tried in an alternative version ####
    # M.control.cntl[1] = 1.0 # diff
    # M.control.icntl[7] = 0 # diff
    ######################
    try
        (L, D57, s, pp) = ma57_get_factors(M)
    catch
        println("*******   Problem after MA57_1")
        # println(" Cond(H) = $(cond(full(H)))")
        res = PDataMA57_0()
        res.OK = false
        return res
    end

    #################  Future object BlockDiag operator?
    vD1 = diag(D57)       # create internal representation for block diagonal D
    vD2 = diag(D57, 1)     #
    vQ1 = ones(length(vD1))       # vector representation of orthogonal matrix Q
    vQ2 = zeros(length(vD2))      #
    vQ2m = zeros(length(vD2))     #
    veig = copy(vD1)      # vector of eigenvalues of D, initialized to diagonal of D
    # if D diagonal, nothing more will be computed

    i = 1
    while i < length(vD1)
        if vD2[i] == 0.0
            i += 1
        else
            mA = [vD1[i] vD2[i]; vD2[i] vD1[i+1]] #  2X2 submatrix
            # DiagmA, Qma = eig(mA)                   #  spectral decomposition of mA
            X = eigen(mA)
            DiagmA = X.values
            Qma = X.vectors
            veig[i] = DiagmA[1]
            vQ1[i] = Qma[1, 1]
            vQ2[i] = Qma[1, 2]
            vQ2m[i] = Qma[2, 1]
            vQ1[i+1] = Qma[2, 2]
            veig[i+1] = DiagmA[2]
            i += 2
        end
    end

    Q = sparse(SparseArrays.spdiagm(0 => vQ1, 1 => vQ2m, -1 => vQ2))
    Δ = veig

    l_m, = findmin(Δ)
    sg = s .* g
    L = SparseMatrixCSC{Float64,Int64}(L)  #### very important, the \ command doesn't work with Int32
    ĝ = L \ (sg[pp])
    g̃ = Q' * ĝ

    n_g = norm(g)
    λ = max(-l_m, 0.0)
    return PDataMA57(L, D57, pp, s, Δ, Q, g̃, λ, true, true)
end

function AInv(X::PDataMA57, d̃::Array{T,1}) where {T}
    d̂ = X.Q * d̃
    u = X.L' \ d̂
    return u[invperm(X.pp)] .* X.s
end


"""
Water flooding with BL assumptions
"""
function water_flood_numeric(core_props, fluids, rel_perms, core_flood; Nx = 30)
    pv_inj = core_flood.injected_pore_volume

    # Nx = 100 # number of cells in x direction
    H = core_props.length # [m] length of the domain in y direction
    D_core = core_props.diameter
    A_core=π*D_core^2/4
    V_core=H*A_core
    m = createMesh1D(Nx, H)

    ## define the physical parametrs
    # all the Corey-type relperm parametrs are defined for the oil-wet and water-wet
    # cases
    k0 = core_props.permeability # [m^2] average reservoir permeability
    phi0 = core_props.porosity # average porosity
    mu_oil = fluids.oil_viscosity # [Pa.s] oil viscosity
    mu_water = fluids.water_viscosity # [Pa.s] water viscosity
    krw0 = rel_perms.krw0
    kro0 = rel_perms.kro0
    nw = rel_perms.nw
    no = rel_perms.no
    sor=rel_perms.sor
    swc=rel_perms.swc

    p0 = core_flood.back_pressure # [bar] pressure
    u_in= core_flood.injection_velocity # [m/s] equal to 1 m/day
    sw0 = core_flood.initial_water_saturation

    sw_in = core_flood.injected_water_saturation

    eps1=1e-5
    perm_val= k0 #field2d(Nx,Ny,k0,V_dp,clx,cly)

    k=createCellVariable(m, perm_val)
    phi=createCellVariable(m, phi0)

    lw = geometricMean(k)/mu_water
    lo = geometricMean(k)/mu_oil

    ## Define the boundaries: all fixed Sw=1, fixed pressure everywhere(?)
    BCp = createBC(m) # Neumann BC for pressure
    BCs = createBC(m) # Neumann BC for saturation
    BCc = createBC(m) # Neumann BC for salinity

    BCp.left.a[:].=-k0/mu_water
    BCp.left.b[:].=0
    BCp.left.c[:].=u_in

    BCp.right.a[:].=0.0
    BCp.right.b[:].=1.0
    BCp.right.c[:].=p0

    BCs.left.a[:].=0.0
    BCs.left.b[:].=1.0
    BCs.left.c[:].=sw_in

    BCc.left.a[:].=0.0
    BCc.left.b[:].=1.0
    BCc.left.c[:].=1.0 # injection concentration

    ## define the time step and solver properties
    t_end = pv_inj*V_core*phi0/(u_in*A_core) # [s] final time
    dt=H/Nx/u_in/10
    SW0 = sw0
    eps_p = 1e-3 # pressure accuracy
    eps_sw = 1e-6 # saturation accuracy
    ## define the variables
    sw_old = createCellVariable(m, SW0, BCs)
    sw_new=createCellVariable(m, SW0, BCs)
    p_old = createCellVariable(m, p0, BCp)
    p_new= createCellVariable(m, p0, BCp)
    c_old = createCellVariable(m, 0.0, BCc)
    c_face=arithmeticMean(c_old) # just to initialize

    # define the rel perm functions
    KRW = sw -> CF.krw(sw, krw0, sor, swc, nw)
    KRO = sw -> CF.kro(sw, kro0, sor, swc, no)
    dKRW = sw -> CF.dkrwdsw(sw, krw0, sor, swc, nw)
    dKRO = sw -> CF.dkrodsw(sw, kro0, sor, swc, no)

    sw = copyCell(sw_old)
    oil_init=domainInt(1-sw_old)
    p = copyCell(p_old)
    uw = -gradientTerm(p_old) # an estimation of the water velocity
    ## start the main loop
    # generate intial pressure profile (necessary to initialize the fully
    # implicit solver)
    dp_core = zeros(1)
    dp_core[1]=H*u_in/(k0*(CF.kro(sw0, kro0, sor, swc, no)/mu_oil+
        CF.krw(sw0, krw0, sor, swc, nw)/mu_water))
    rec_fact=zeros(1)
    c_out_sal=zeros(1)
    c_out_sal[1]=c_face.xvalue[end]
    t_sec=zeros(1)
    t = 0.0
    dt0=dt
    dsw_alwd= 0.005
    dp_alwd= 100.0 # Pa
    dc_alwd= 0.005
    labdaw=createFaceVariable(m,1.0) # initialize
    FL=fluxLimiter("SUPERBEE")

    prog_1=Progress(100, 1)

    while (t<t_end)
        # Implicit loop
        error_p = 1e5
        error_sw = 1e5
        # implicit solver
        loop_count=0
        while ((error_sw>eps_sw))
            loop_count+=1
            if loop_count>10
                break
            end
            # calculate parameters
            pgrad = gradientTerm(p)
            uw = -labdaw*pgrad
            sw_face = upwindMean(sw, uw) # average value of water saturation
            # sw_ave=arithmeticMean(sw)
            # solve for pressure at known Sw
            labdao = lo*faceEval(KRO, sw_face)
            labdaw = lw*faceEval(KRW, sw_face)
            dlabdaodsw = lo*faceEval(dKRO, sw_face)
            dlabdawdsw = lw*faceEval(dKRW, sw_face)
            labda = labdao+labdaw
            dlabdadsw = dlabdaodsw+dlabdawdsw
            # compute [Jacobian] matrices
            Mdiffp1 = diffusionTerm(-labda)
            Mdiffp2 = diffusionTerm(-labdaw)
            # Mconvsw1 = convectionUpwindTerm(-dlabdadsw.pgrad, uw)
            # Mconvsw2 = convectionUpwindTerm(-dlabdawdsw*pgrad, uw)
            Mconvsw1 = convectionUpwindTerm(-dlabdadsw*pgrad, uw)
            Mconvsw2 = convectionUpwindTerm(-dlabdawdsw*pgrad, uw)
            Mtranssw2, RHStrans2 = transientTerm(sw_old, dt, phi)
            # Compute RHS values
            RHS1 = Mconvsw1*sw.value # divergenceTerm(-dlabdadsw*sw_face*pgrad)
            RHS2 = Mconvsw2*sw.value # divergenceTerm(-dlabdawdsw*sw_face*pgrad)
            # include boundary conditions
            Mbcp, RHSbcp = boundaryConditionTerm(BCp)
            Mbcsw, RHSbcsw = boundaryConditionTerm(BCs)
            # Couple the equations; BC goes into the block on the main diagonal
            M = [Mdiffp1+Mbcp Mconvsw1; Mdiffp2 Mconvsw2+Mtranssw2+Mbcsw]
            RHS = [RHS1+RHSbcp; RHS2+RHStrans2+RHSbcsw]
            x = M\RHS
            p_new.value[:] = (x[1:(Nx+2)])
            sw_new.value[:] = (x[(Nx+2)+1:end])

            error_p = maximum(abs.((internalCells(p_new)-internalCells(p))./internalCells(p_new)))
            error_sw = maximum(abs.(internalCells(sw_new)-internalCells(sw)))
            # dt_new=dt*min(dp_alwd/error_p, dsw_alwd/error_sw)
            # print("sw_error = $error_sw \n")
            # assign new values of p and sw
            p.value[:]=p_new.value[:]
            sw.value[:]=sw_new.value[:]
        end # end of implicit loop
        if loop_count>10
            p=copyCell(p_old)
            sw=copyCell(sw_old)
            dt=dt/5
            continue
        end
        # solve the advection equation
        # ϕ d/dt(c sw)=[sw dc/dt + c dsw/dt]
        dswdt=(sw_new-sw_old)/dt
        Msc=linearSourceTerm(phi*dswdt)
        sw_face = upwindMean(sw, -labdaw*pgrad) # average value of water saturation
        pgrad = gradientTerm(p_new)
        labdaw = lw*faceEval(KRW, sw_face)
        uw=-labdaw*pgrad
        Mtc, RHStc=transientTerm(c_old, dt, phi*sw)
        Mbcc, RHSbcc=boundaryConditionTerm(BCc)
        c_temp=copyCell(c_old)
        for j in 1:4
            Mconvc, RHSconv=convectionTvdTerm(uw, c_temp, FL)
            c_temp=solveLinearPDE(m, Msc+Mtc+Mconvc+Mbcc, RHStc+RHSbcc+RHSconv)
        end
        c_new=copyCell(c_temp)
        c_face=tvdMean(c_new, uw, FL)

        # adaptive time step
        dp = maximum(abs.((internalCells(p_new)-internalCells(p_old))./internalCells(p_new)))
        dsw = maximum(abs.(internalCells(sw_new)-internalCells(sw_old)))
        dc = maximum(abs.(internalCells(c_new)-internalCells(c_old)))

        # update the variables
        t+=dt
        ProgressMeter.update!(prog_1, floor(Int, t/t_end*100))
        c_old=copyCell(c_new)
        p_old = copyCell(p)
        sw_old = copyCell(sw)

        dt=min(dt*(dsw_alwd/dsw), 15*dt, t_end-t-eps())

        # calculate the recovery factor
        push!(rec_fact, (oil_init-domainInt(1-sw))/oil_init)
        push!(t_sec, t)
        push!(c_out_sal, c_face.xvalue[end])
        push!(dp_core, 0.5*sum(p.value[1:2])-p0)

        # print(t)
        #GR.plot(sw.value[2:end-1])
        # GR.plot(c_new.value[2:end-1])
        #GR.plot(SF.xvalue)
        # GR.imshow(sw.value[2:end-1,2:end-1])
        # visualizeCells(1-sw)
        # GR.plot(t_sec/3600/24, rec_fact)
        # xlabel('time [day]')
        # ylabel('recovery factor')
        # title([num2str(t/3600/24) ' day']) drawnow
    end
    pv = t_sec*u_in/(H*phi0)
    sw_face = linearMean(sw)
    c_face  = linearMean(c_old)
    x = m.facecenters.x
    t_sec, pv, rec_fact, dp_core, x, sw_face.xvalue, c_face.xvalue, c_out_sal
end # imb_impes


module Excercise_3
using Random

# Exercise 3 (JULIA)
# Create a 2x4 two dimensional matrix with random floats in it and in the next step determine the biggest element.

# create Matrix
M = randn(2,4)
# determine Maximum (biggest element)
maximum_M = maximum(M)

println("Matrix 2 x 4 is:")
println(M)
println("Maximum of Matrix is:")
println(maximum_M)

end


## Minimum diagonal adjustment to a matrix to make it
## positive-definite, preserving the off-diagonal values.
#
# The problem is to factor Q into C*C', where each column of C
# has at most two non-zero entries. Consequently we have the
# equation
#
#     C_{ik}*C_{jk} = Q_{ij}
#
# where k is the column of C corresponding to (i,j).
# In addition to these paired columns, we also allow single-entry
# columns with the value c_i (note lowercase). Then we can write
#
#    (C*C')_{ii} = c_i^2 + \sum_k C_{ik}^2
#
# and we seek to minimize
#
#  P = \sum_i (C*C')_{ii}/|Q_{ii}|
#
# over the i with Q_{ii} != 0, subject to the constraint
#
#  (C*C')_{ii} >= |Q_{ii}|.
#
# Unlike the straight numerator, the scaled version used in P is
# invariant with respect to physical units.
#
# This minimization can be done analytically by introducing a ratio
#
#     r_k = C_{ik}/C_{jk}
#
# and computing the gradient of P wrt r.

"""
    Qp = possemidef(Q)

Return a variant of Q preserving the off-diagonal values and adjusting
the diagonals, as needed, to make the matrix positive semi-definite. Any
NaN entries of Q are treated as 0. Q is assumed to be symmetric and
only the lower triangle is used.
"""
function possemidef(Q::AbstractMatrix)
    ind1, ind2 = axes(Q)
    ind1 == ind2 || error("Matrix must be square")
    Qp = copy(Q)
    for i in ind1
        Qp[i,i] = 0
    end
    for j in ind2
        csum = zero(eltype(Q))
        qjj = abs(Q[j,j])
        for i = j+1:last(ind1)
            q = Q[i,j]
            ((q == 0) | isnan(q)) && continue
            qii = abs(Q[i,i])
            r = sign(q) * sqrt(qjj/qii)
            if r == 0 || !isfinite(r)
                r = oftype(r, sign(q))
            end
            Qp[j,j] += q * r
            Qp[i,i] += q / r
        end
    end
    for i in ind1
        # this is equivalent to adding c_i^2
        Qp[i,i] = max(Qp[i,i], abs(Q[i,i]))
    end
    return Qp
end


@with_kw struct ContinuousBanditMDP <: MDP{Float32, Float32}
    goal::Float32 = 0
    σ::Float32 = 1
end

POMDPs.reward(mdp::ContinuousBanditMDP, s, a) = exp(-(a[1] - mdp.goal)^2 / mdp.σ) - 0.1*abs(a[1] - mdp.goal)
POMDPs.gen(mdp::ContinuousBanditMDP, s, a, rng::AbstractRNG = Random.GLOBAL_RNG; info=Dict()) = (sp = a[1], r = reward(mdp, s, a))
POMDPs.initialstate(mdp::ContinuousBanditMDP) = Deterministic(Float32(randn()))
POMDPs.isterminal(mdp::ContinuousBanditMDP, s) = true
isfailure(mdp::ContinuousBanditMDP, s) = s[1] > mdp.goal
POMDPs.discount(mdp::ContinuousBanditMDP) = 0.99
POMDPs.convert_s(::Type{V}, s::Float32, mdp::ContinuousBanditMDP) where {V<:AbstractArray} = [s]


# This script provides an example of simulating changepoint data,
# using PELT to find these changepoints, and plotting the results
#
# Requires Gadfly plotting package

using Distributions, Gadfly
using Changepoints

n = 1000        # Sample size
λ = 100         # Frequencey of changes

μ, σ = Normal(0.0, 10.0), 1.0
sample, cps = @changepoint_sampler n λ Normal(μ, σ)

# Find changepoints via PELT
pelt_cps, pelt_cost = @PELT sample Normal(?, σ)

# Find changepoints via BS
bs_cps = @BS sample Normal(?, σ)

# Plot data with true changepoints
p = plot(sample, cps)


using FourierCore, FourierCore.Grid, FourierCore.Domain
using FFTW, LinearAlgebra, BenchmarkTools, Random, JLD2
using GLMakie
rng = MersenneTwister(1234)
# Random.seed!(123456789)
Random.seed!(12)
# jld_name = "high_order_timestep_spatial_tracer_"
jld_name = "blocky"
include("transform.jl")
include("random_phase_kernel.jl")
# using GLMakie
using CUDA
arraytype = CuArray
Ω = S¹(4π)^2
N = 2^7 # number of gridpoints
grid = FourierGrid(N, Ω, arraytype=arraytype)
nodes, wavenumbers = grid.nodes, grid.wavenumbers

x = nodes[1]
y = nodes[2]
kˣ = wavenumbers[1]
kʸ = wavenumbers[2]
# construct filter
kxmax = maximum(kˣ)
kymax = maximum(kˣ)
filter = @. (kˣ)^2 + (kʸ)^2 ≤ ((kxmax / 2)^2 + (kymax / 2)^2)

# now define the random field
wavemax = 3
𝓀 = arraytype(collect(-wavemax:0.5:wavemax))
𝓀ˣ = reshape(𝓀, (length(𝓀), 1))
𝓀ʸ = reshape(𝓀, (1, length(𝓀)))
A = @. 0.1 * (𝓀ˣ * 𝓀ˣ + 𝓀ʸ * 𝓀ʸ)^(-11 / 12)
A[A.==Inf] .= 0.0
φ = arraytype(2π * rand(size(A)...))
field = arraytype(zeros(N, N))

##
# Fields
# velocity
ψ = arraytype(zeros(ComplexF64, N, N))
u = similar(ψ)
v = similar(ψ)

# theta
θ = similar(ψ)
∂ˣθ = similar(ψ)
∂ʸθ = similar(ψ)
κΔθ = similar(ψ)
θ̇ = similar(ψ)
s = similar(ψ)
θ̅ = similar(ψ)
k₁ = similar(ψ)
k₂ = similar(ψ)
k₃ = similar(ψ)
k₄ = similar(ψ)
θ̃ = similar(ψ)

# source
s = similar(ψ)
@. s = cos(kˣ[5] * x)

# phase
φ̇ = similar(A)

# operators
∂x = im * kˣ
∂y = im * kʸ
Δ = @. ∂x^2 + ∂y^2

# plan ffts
P = plan_fft!(ψ)
P⁻¹ = plan_ifft!(ψ)

# number of gridpoints in transition is about λ * N / 2
bump(x; λ=10 / N, width = π/2) = 0.5 * (tanh((x + width / 2) / λ) - tanh((x - width / 2) / λ))
bumps(x; λ=20 / N, width = 1.0) = 0.25 * (bump(x, λ = λ, width = width) + bump(x, λ = λ, width = 2.0*width) + bump(x, λ = λ, width = 3.0*width) + bump(x, λ = λ, width = 4.0*width))

##
Δx = x[2] - x[1]
κ = Δx^2
# κ = 2 / 2^8 # fixed diffusivity
# κ = 2e-4
Δt = (Δx) / (4π) * 5

# take the initial condition as negative of the source
tic = Base.time()

# save some snapshots
ψ_save = typeof(real.(Array(ψ)))[]
θ_save = typeof(real.(Array(ψ)))[]

r_A = Array(@. sqrt((x - 2π)^2 + (y - 2π)^2))
θ_A = [bumps(r_A[i, j]) for i in 1:N, j in 1:N]
θ .= CuArray(θ_A)
# @. θ = bump(sqrt(x^2 + y^2)) # scaling so that source is order 1
θclims = extrema(Array(real.(θ))[:])
P * θ # in place fft
@. κΔθ = κ * Δ * θ
P⁻¹ * κΔθ # in place fft
s .= -κΔθ * 0.0
P⁻¹ * θ # in place fft
θ̅ .= 0.0

t = [0.0]
tend = 5

iend = ceil(Int, tend / Δt)

params = (; ψ, A, 𝓀ˣ, 𝓀ʸ, x, y, φ, u, v, ∂ˣθ, ∂ʸθ, s, P, P⁻¹, filter)

size_of_A = size(A)

realizations = 100
for j in 1:realizations
    # new realization of flow
    rand!(rng, φ) # between 0, 1
    φ .*= 2π # to make it a random phase
    event = stream_function!(ψ, A, 𝓀ˣ, 𝓀ʸ, x, y, φ)
    wait(event)
    t[1] = 0

    θ .= CuArray(θ_A)
    for i = 1:iend
        # fourth order runge kutta on deterministic part
        # keep ψ frozen is the correct way to do it here

        # the below assumes that φ is just a function of time
        θ_rhs_new!(k₁, θ, params)
        @. θ̃ = θ + Δt * k₁ * 0.5

        φ_rhs_normal!(φ̇, φ, rng)
        @. φ += sqrt(Δt / 2) * φ̇

        θ_rhs_new!(k₂, θ̃, params)
        @. θ̃ = θ + Δt * k₂ * 0.5
        θ_rhs_new!(k₃, θ̃, params)
        @. θ̃ = θ + Δt * k₃

        φ_rhs_normal!(φ̇, φ, rng)
        @. φ += sqrt(Δt / 2) * φ̇

        θ_rhs_new!(k₄, θ̃, params)
        @. θ += Δt / 6 * (k₁ + 2 * k₂ + 2 * k₃ + k₄)

        # update stochastic part
        # φ_rhs_normal!(φ̇, φ, rng)
        # @. φ += sqrt(Δt) * φ̇

        t[1] += Δt
        # save output

        #=
        if i % div(iend, 10) == 0
            println("Saving at i=", i)
            push!(ψ_save, Array(real.(ψ)))
            push!(θ_save, Array(real.(θ)))
            println("extrema are ", extrema(θ_save[end]))
            println("time is t = ", t[1])
        end

        if t[1] >= tstart
            θ̅ .+= Δt * θ
        end

        if i % div(iend, 100) == 0
            println("time is t = ", t[1])
            local toc = Base.time()
            println("the time for the simulation is ", toc - tic, " seconds")
            println("extrema are ", extrema(real.(θ)))
            println("on wavenumber index ", index_choice)
        end
        =#

    end
    println("finished realization ", j)
    @. θ̅ += θ / realizations
end

toc = Base.time()
println("the time for the simulation was ", toc - tic, " seconds")

x_A = Array(x)[:] .- 2π
θ_F = Array(real.(θ))
θ̅_F = Array(real.(θ̅))

fig = Figure(resolution=(2048, 512))
ax1 = Axis(fig[1, 1], title="t = 0")
ax2 = Axis(fig[1, 2], title="instantaneous t = " * string(tend))
ax3 = Axis(fig[1, 4], title="ensemble average t = " * string(tend))
println("the extrema of the end field is ", extrema(θ_F))
println("the extrema of the ensemble average is ", extrema(θ̅_F))
colormap = :bone_1
colormap = :nipy_spectral
heatmap!(ax1, x_A, x_A, θ_A, colormap=colormap, colorrange=(0.0, 1.0), interpolate=true)
hm = heatmap!(ax2, x_A, x_A, θ_F, colormap=colormap, colorrange=(0.0, 1.0), interpolate=true)
hm_e = heatmap!(ax3, x_A, x_A, θ̅_F, colormap=colormap, colorrange=(0.0, 0.2), interpolate=true)
Colorbar(fig[1, 3], hm, height=Relative(3 / 4), width=25, ticklabelsize=30, labelsize=30, ticksize=25, tickalign=1,)
Colorbar(fig[1, 5], hm_e, height=Relative(3 / 4), width=25, ticklabelsize=30, labelsize=30, ticksize=25, tickalign=1,)
display(fig)

# indices = abs.(θ_F[:]) .> 1e-3;
# hist(θ_F[indices])
#=
println("saving ", jld_name * string(index_choice) * ".jld2")
θ̅a = Array(real.(θ̅))
xnodes = Array(x)[:]
ynodes = Array(y)[:]
kˣ_wavenumbers = Array(kˣ)[:]
kʸ_wavenumbers = Array(kˣ)[:]
source = Array(s)
jldsave(jld_name * string(index_choice) * ".jld2"; ψ_save, θ_save, θ̅a, κ, xnodes, ynodes, kˣ_wavenumbers, kʸ_wavenumbers, source)
=#


## Load everything
using Pkg
Pkg.activate(".")
Pkg.instantiate()
include("src/EasyStocks.jl")

## Setup
using CairoMakie # For plotting
mp = ModPar(q=0.005,ψ=0.015,xstar=2.0)
np = NumPar(mp)
MS = ModelSolution(mp,np)
## Run-code

@time SolveModel!(MS)

fig = Figure()
lines(fig[1,1],MS.np.xgrd, MS.α[:,1])
ylims!(-0.05,1.05)
lines(fig[1,2],MS.np.xgrd, MS.V[:,1])
ylims!(-6,1)
lines(fig[2,2],MS.np.xgrd, MS.V[:,2])
ylims!(-1.5,0)
lines(fig[2,1],MS.np.xgrd, MS.b[:,1] + MS.s[:,1])
ylims!(-0.05,5)
linkxaxes!(fig.content...)
fig


module NaiveResamplers
export naivefilt

# Naive rational resampler
function naivefilt( h::Vector, x::Vector, resamplerate::Rational = 1//1 )

    upfactor     = numerator( resamplerate )
    downfactor   = denominator( resamplerate )
    xLen         = length( x )
    xZeroStuffed = zeros( eltype(x), length( x ) * upfactor )

    for n in 0:length(x)-1
        xZeroStuffed[ n*upfactor+1 ] = x[ n+1 ]
    end

    y = Base.filt( h, 1.0, xZeroStuffed )
    y = [ y[n] for n = 1:downfactor:length( y ) ]
end

# Naive arbitrary resampler
function naivefilt( h::Vector, x::Vector, resamplerate::AbstractFloat, numfilters::Integer = 32 )
    xLen          = length( x )
    xInterpolated = naivefilt( h, x, numfilters//1 )
    xLen          = length( xInterpolated )
    yLen          = ceil( Int, xLen * resamplerate )
    y             = similar( x, yLen )
    yIdx          = 1
    xIdx          = 1
    α             = 0.0
    (δ, 𝜙Stride)  = modf( numfilters/resamplerate )
    𝜙Stride       = Int( 𝜙Stride )

    while xIdx < xLen
        yLower  = xInterpolated[xIdx]
        yUpper  = xInterpolated[xIdx+1]
        y[yIdx] = yLower + α*( yUpper - yLower )
        yIdx   += 1
        α      += δ
        xIdx   += floor( Int, α ) + 𝜙Stride
        α       = mod( α, 1.0 )
    end

    resize!( y, yIdx-1 )

    return y
end

end # module


function FODEdG!(λ::T, f::Function, u0::T, t::OffsetVector{T},
                 r::Integer, M::Integer,
                 store::Store{T}) where T <: AbstractFloat
    @argcheck r ≤ store.rmax
    @argcheck M ≤ store.Mmax
    N = length(t) - 1
    U = Vector{Vector{T}}(undef, N)
    G = coef_G(T, r)
    K = coef_K(T, r, r)
    Hn = OffsetArray{Matrix{T}}(undef, 0:N-1)
    for n = 0:N-1
        Hn[n] = Array{T}(undef, r, r)
    end
    τ, wτ = rule(store.legendre[M])
    Ψ = view(store.Ψ, 1:r, 1:M)
    legendre_polys!(Ψ, τ)
    Fn = Array{T}(undef, r)
    load_vector!(Fn, (t[0],t[1]), f, τ, wτ, Ψ)
    Ψ_at_minus1 = Vector{T}(undef, r)
    legendre_polys!(Ψ_at_minus1, -one(T))
    coef_Hn!(Hn, 1, 0, t, M, store)
    A = G + λ * Hn[0]
    b = copy(Fn) + Ψ_at_minus1 * u0
    Fact = LinearAlgebra.lu(A)
    U[1] = Fact \ b
    for n = 2:N
        coef_Hn!(Hn, n, n-1, t, M, store)
        load_vector!(Fn, (t[n-1],t[n]), f, τ, wτ, Ψ)
        fill!(b, zero(T))
        for ℓ = 1:n-1
            b .= b .+ ( Hn[n-ℓ] * U[ℓ] )
        end
        b .= Fn - λ * b + K * U[n-1]
        A = G + λ * Hn[0]
        Fact = LinearAlgebra.lu(A)
        U[n] = Fact \ b
    end
    return U
end

function FODEdG!(λ::T, tmax::T, f::Function, u0::T,
                 N::Integer, r::Integer, M::Integer,
                 store::Store{T}) where T <: AbstractFloat
    @argcheck r ≤ store.rmax
    @argcheck M ≤ store.Mmax
    G = coef_G(T, r)
    K = coef_K(T, r, r)
    H = coef_H_uniform!(r, N, M, store)
    U = Vector{Vector{T}}(undef, N)
    k = tmax / N
    α = store.α
    t_ = range(zero(T), tmax, length=N+1)
    t = OffsetArray(t_, 0:N)
    τ, wτ = rule(store.legendre[M])
    Ψ = view(store.Ψ, 1:r, 1:M)
    legendre_polys!(Ψ, τ)
    Fn = Array{T}(undef, r)
    load_vector!(Fn, (t[0],t[1]), f, τ, wτ, Ψ)
    Ψ_at_minus1 = Vector{T}(undef, r)
    legendre_polys!(Ψ_at_minus1, -one(T))
    A = G + ( λ * k^α ) * H[0]
    b = copy(Fn) + Ψ_at_minus1 * u0
    Fact = LinearAlgebra.lu(A)
    U[1] = Fact \ b
    for n = 2:N
        load_vector!(Fn, (t[n-1],t[n]), f, τ, wτ, Ψ)
        fill!(b, zero(T))
        for ℓ = 1:n-1
            b .= b .+ ( H[n-ℓ] * U[ℓ] )
        end
        b .= Fn - ( λ * k^α ) .* b + K * U[n-1]
        U[n] = Fact \ b
    end
    return t, U
end

function load_vector!(Fn::Vector{T}, In::Tuple{T,T}, f::Function,
                      τ::Vector{T}, wτ::Vector{T}, Ψ::AbstractMatrix{T}
                     ) where T <: AbstractFloat
    r, M = size(Ψ)
    tnm1, tn = In
    fill!(Fn, zero(T))
    for m = 1:M
        tm = ( (1-τ[m]) * tnm1 + (1+τ[m]) * tn ) / 2
        for i = 1:r
            Fn[i] += wτ[m] * f(tm) * Ψ[i,m]
        end
    end
    kn = tn - tnm1
    for i = 1:r
        Fn[i] *= kn / 2
    end
end

function graded_mesh(max_t::T, cutoff_t::T,
                     N_outer::Integer, mg::T) where { T <: AbstractFloat }
    k = ( max_t - cutoff_t ) / N_outer
    N_inner = ceil(Integer, -1/expm1(log1p(-k/cutoff_t)/mg))
    N = N_inner + N_outer
    t = OffsetVector{T}(undef, 0:N)
    t[0] = zero(T)
    for n = 1:N_inner
        t[n] = (n/N_inner)^mg * cutoff_t
    end
    for n = N_inner+1:N
        t[n] = cutoff_t + (n-N_inner)*k
    end
    return t
end

function ODEdG!(λ::T, tmax::T, f::Function, u0::T,
                N::Integer, r::Integer, M::Integer,
                ) where T <: AbstractFloat
    G = coef_G(T, r)
    K = coef_K(T, r, r)
    H = coef_H(T, r)
    U = Vector{Vector{T}}(undef, N)
    k = tmax / N
    t_ = range(zero(T), tmax, length=N+1)
    t = OffsetArray(t_, 0:N)
    τ, wτ = GaussQuadrature.legendre(M)
    Ψ = Matrix{T}(undef, r, M)
    legendre_polys!(Ψ, τ)
    Fn = Array{T}(undef, r)
    load_vector!(Fn, (t[0],t[1]), f, τ, wτ, Ψ)
    Ψ_at_minus1 = Vector{T}(undef, r)
    legendre_polys!(Ψ_at_minus1, -one(T))
    A = G + ( λ * k ) * H
    b = copy(Fn) + Ψ_at_minus1 * u0
    Fact = LinearAlgebra.lu(A)
    U[1] = Fact \ b
    for n = 2:N
        load_vector!(Fn, (t[n-1],t[n]), f, τ, wτ, Ψ)
        b .= Fn + K * U[n-1]
        U[n] = Fact \ b
    end
    return t, U
end


using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:3] <= 1)

                     @variable(m, -5.118547195289687 <= q <= 7.5851414059354285)

                     add_NL_constraint(m, :(sigmoid(0.14385330157018572 + -0.4061277131705876*$(x[1]) + 0.38692435149352855*$(x[2]) + 0.23112211611402778*$(x[3])) + sigmoid(0.5778882819806599 + 0.8181171649048586*$(x[1]) + -0.37932684839331854*$(x[2]) + 0.7567763507538823*$(x[3])) + sigmoid(0.9224951510136457 + -0.8061937532687207*$(x[1]) + 0.8183596504529116*$(x[2]) + -0.17766965347054509*$(x[3])) + sigmoid(-0.41093962924162053 + -0.25704902180989464*$(x[1]) + 0.523731990676354*$(x[2]) + 0.7904456861039284*$(x[3])) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m


using JuMP

using LinearAlgebra

function formulate_nearest_good_box_qp(
    Q::Matrix,
    c::Vector,
    y_max::Float64,
    x_target::Vector{Float64},
)
    n = length(c)
    @assert n == size(Q, 1) == size(Q, 2)
    model = JuMP.Model()
    JuMP.@variable(model, 0 ≤ x[1:n] ≤ 1)
    JuMP.@variable(model, y <= y_max)
    JuMP.@constraint(model, y ≥ -0.5dot(x, Q * x) - dot(c, x))

    JuMP.@variable(m, t[1:n] >= 0)
    for i = 1:n
        JuMP.@constraint(m, t[i] >= x[i] - x_target[i])
        JuMP.@constraint(m, t[i] >= -x[i] + x_target[i])
    end
    JuMP.@objective(model, Min, sum(t))
    return JuMP.backend(model)
end