export Category, FreeCategory, Ob, Hom, dom, codom, id, compose, ⋅,
  Copresheaf, FreeCopresheaf, El, ElExpr, ob, act,
  Presheaf, FreePresheaf, coact,
  DisplayedCategory, Act, hom,
  MCategory, FreeMCategory, Tight, reflexive, transitive

import Base: show

# Category
##########

""" Theory of *categories* (with no extra structure)

**Warning**: We compose functions from left to right, i.e., if f:A→B and g:B→C
then compose(f,g):A→C. Under this convention function are applied on the right,
e.g., if a∈A then af∈B.

We retain the usual meaning of the symbol ∘ (\\circ), i.e., g∘f = compose(f,g).
This usage is too entrenched to overturn, inconvenient though it may be.
We use symbol ⋅ (\\cdot) for diagrammatic composition: f⋅g = compose(f,g).
"""
@theory Category{Ob,Hom} begin
  # Unicode aliases.
  @op begin
    (→) := Hom
    (⋅) := compose
  end

  """ Object in a category """
  Ob::TYPE

  """ Morphism in a category """
  Hom(dom::Ob,codom::Ob)::TYPE

  id(A::Ob)::(A → A)
  compose(f::(A → B), g::(B → C))::(A → C) ⊣ (A::Ob, B::Ob, C::Ob)

  # Category axioms.
  ((f ⋅ g) ⋅ h == f ⋅ (g ⋅ h)
    ⊣ (A::Ob, B::Ob, C::Ob, D::Ob, f::(A → B), g::(B → C), h::(C → D)))
  f ⋅ id(B) == f ⊣ (A::Ob, B::Ob, f::(A → B))
  id(A) ⋅ f == f ⊣ (A::Ob, B::Ob, f::(A → B))
end

# Convenience constructors
compose(fs::AbstractVector) = foldl(compose, fs)
compose(f, g, h, fs...) = compose([f, g, h, fs...])

@syntax FreeCategory{ObExpr,HomExpr} Category begin
  compose(f::Hom, g::Hom) = associate_unit(new(f,g; strict=true), id)
end

function show_unicode(io::IO, expr::HomExpr{:compose}; kw...)
  Syntax.show_unicode_infix(io, expr, "⋅"; kw...)
end

function show_latex(io::IO, expr::HomExpr{:id}; kw...)
  Syntax.show_latex_script(io, expr, "\\mathrm{id}")
end
function show_latex(io::IO, expr::HomExpr{:compose}; paren::Bool=false, kw...)
  Syntax.show_latex_infix(io, expr, "\\cdot"; paren=paren)
end

function show(io::IO, ::MIME"text/plain", expr::HomExpr)
  show_unicode(io, expr)
  print(io, ": ")
  show_unicode(io, dom(expr))
  print(io, " → ")
  show_unicode(io, codom(expr))
end
function show(io::IO, ::MIME"text/latex", expr::HomExpr)
  print(io, "\$")
  show_latex(io, expr)
  print(io, " : ")
  show_latex(io, dom(expr))
  print(io, " \\to ")
  show_latex(io, codom(expr))
  print(io, "\$")
end

# (Co)presheaf
##############

""" Theory of *copresheaves*.

Axiomatized as a covariant category action.
"""
@theory Copresheaf{Ob,Hom,El} <: Category{Ob,Hom} begin
  # copresheaf = object-indexed family
  El(ob::Ob)::TYPE

  # functoriality = covariant action
  act(x::El(A), f::Hom(A,B))::El(B) ⊣ (A::Ob, B::Ob)
  @op (⋅) := act

  # action equations
  act(act(x, f), g) == act(x, (f ⋅ g)) ⊣
    (A::Ob, B::Ob, C::Ob, f::(A → B), g::(B → C), x::El(A))
  act(x, id(A)) == x ⊣ (A::Ob, x::El(A))
end

abstract type ElExpr{T} <: GATExpr{T} end

@syntax FreeCopresheaf{ObExpr,HomExpr,ElExpr} Copresheaf begin
  compose(f::Hom, g::Hom) = associate_unit(new(f,g; strict=true), id)
end

""" Theory of *presheaves*.

Axiomatized as a contravariant category action.
"""
@theory Presheaf{Ob,Hom,El} <: Category{Ob,Hom} begin
  # presheaf = object-indexed family
  El(ob::Ob)::TYPE

  # functoriality = contravariant action
  coact(f::Hom(A,B), x::El(B))::El(A) ⊣ (A::Ob, B::Ob)
  @op (⋅) := coact

  # action equations
  coact(f, coact(g, x)) == coact((f ⋅ g), x) ⊣
    (A::Ob, B::Ob, C::Ob, f::(A → B), g::(B → C), x::El(C))
  coact(id(A), x) == x ⊣ (A::Ob, x::El(A))
end

@syntax FreePresheaf{ObExpr,HomExpr,ElExpr} Presheaf begin
  compose(f::Hom, g::Hom) = associate_unit(new(f,g; strict=true), id)
end

function show(io::IO, ::MIME"text/plain", expr::ElExpr)
  show_unicode(io, expr)
  print(io, ": ")
  show_unicode(io, ob(expr))
end
function show(io::IO, ::MIME"text/latex", expr::ElExpr)
  print(io, "\$")
  show_latex(io, expr)
  print(io, " : ")
  show_latex(io, ob(expr))
  print(io, "\$")
end

function show_unicode(io::IO, expr::Union{ElExpr{:act},ElExpr{:coact}}; kw...)
  Syntax.show_unicode_infix(io, expr, "⋅"; kw...)
end
function show_latex(io::IO, expr::Union{ElExpr{:act},ElExpr{:coact}};
                    paren::Bool=false, kw...)
  Syntax.show_latex_infix(io, expr, "\\cdot"; paren=paren)
end

# Displayed category
####################

""" Theory of a *displayed category*.

More precisely, this is the theory of a base category ``C`` (`Ob`,`Hom`) and a
displayed category (`El`,`Act`) over ``C``. Displayed categories axiomatize lax
functors ``C → **Span**`` or equivalently objects in ``**Cat**/C`` in a
pleasant, generalized algebraic style.

Reference: Ahrens & Lumsdaine 2019, "Displayed categories", Definition 3.1.
"""
@theory DisplayedCategory{Ob,Hom,El,Act} <: Category{Ob,Hom} begin
  El(ob::Ob)::TYPE
  Act(hom::Hom(A,B), dom::El(A), codom::El(B))::TYPE ⊣ (A::Ob, B::Ob)

  id(x::El(A))::Act(id(A),x,x) ⊣ (A::Ob)
  compose(f̄::Act(f,x,y), ḡ::Act(g,y,z))::Act(compose(f,g),x,z) ⊣
    (A::Ob, B::Ob, C::Ob, f::(A → B), g::(B → C))

  # Displayed category axioms: category axioms over the base category.
  ((f̄ ⋅ ḡ) ⋅ h̄ == f̄ ⋅ (ḡ ⋅ h̄) ⊣
    (A::Ob, B::Ob, C::Ob, D::Ob, f::(A → B), g::(B → C), h::(C → D),
     w::El(A), x::El(B), y::El(C), z::El(D),
     f̄::Act(f,w,x), ḡ::Act(g,x,y), h̄::Act(h,y,z)))
  f̄ ⋅ id(y) == f̄ ⊣ (A::Ob, B::Ob, f::(A → B), x::El(A), y::El(B), f̄::Act(f,x,y))
  id(x) ⋅ f̄ == f̄ ⊣ (A::Ob, B::Ob, f::(A → B), x::El(A), y::El(B), f̄::Act(f,x,y))
end

# ℳ-category
############

""" Theory of an *ℳ-category*.

The term "ℳ-category", used on the
[nLab](https://ncatlab.org/nlab/show/M-category) is not very common, but the
concept itself shows up commonly. An *ℳ-category* is a category with a
distinguished wide subcategory, whose morphisms are suggestively called *tight*;
for contrast, a generic morphism is called *loose*. Equivalently, an ℳ-category
is a category enriched in the category ℳ of injections, the full subcategory of
the arrow category of Set spanned by injections.

In the following GAT, tightness is axiomatized as a property of morphisms: a
dependent family of sets over the hom-sets, each having at most one inhabitant.
"""
@theory MCategory{Ob,Hom,Tight} <: Category{Ob,Hom} begin
  Tight(hom::Hom(A,B))::TYPE ⊣ (A::Ob, B::Ob)

  # Tightness is a property.
  t == t′ ⊣ (A::Ob, B::Ob, f::Hom(A,B), t::Tight(f), t′::Tight(f))

  # Tight morphisms form a subcategory.
  reflexive(A::Ob)::Tight(id(A))
  transitive(t::Tight(f), u::Tight(g))::Tight(compose(f,g)) ⊣
    (A::Ob, B::Ob, C::Ob, f::Hom(A,B), g::Hom(B,C), t::Tight(f), u::Tight(g))
end

abstract type TightExpr{T} <: GATExpr{T} end

@syntax FreeMCategory{ObExpr,HomExpr,TightExpr} MCategory begin
  compose(f::Hom, g::Hom) = associate_unit(new(f,g; strict=true), id)
  transitive(t::Tight, u::Tight) = associate_unit(new(t,u; strict=true), reflexive)
end


export SymbolicAmplitudeAwarePermutation

"""
    SymbolicAmplitudeAwarePermutation(; τ = 1, m = 3, A = 0.5, lt = Entropies.isless_rand) <: PermutationProbabilityEstimator

See docstring for [`SymbolicPermutation`](@ref).
"""
struct SymbolicAmplitudeAwarePermutation{F} <: PermutationProbabilityEstimator
    τ::Int
    m::Int
    A::Float64
    lt::F
end
function SymbolicAmplitudeAwarePermutation(; τ::Int = 1, m::Int = 2, A::Real = 0.5,
        lt::F = isless_rand) where {F <: Function}
    2 ≤ m || error("Need m ≥ 2, otherwise no dynamical information is encoded in the symbols.")
    0 ≤ A ≤ 1 || error("Weighting factor A must be on interval [0, 1]. Got A=$A.")
    SymbolicAmplitudeAwarePermutation{F}(τ, m, A, lt)
end

"""
    AAPE(x, A::Real = 0.5, m::Int = length(a))

Encode relative amplitude information of the elements of `a`.
- `A = 1` emphasizes only average values.
- `A = 0` emphasizes changes in amplitude values.
- `A = 0.5` equally emphasizes average values and changes in the amplitude values.
"""
function AAPE(x; A::Real = 0.5, m::Int = length(x))
    (A/m)*sum(abs.(x)) + (1-A)/(m-1)*sum(abs.(diff(x)))
end

function probabilities(x::AbstractDataset{m, T}, est::SymbolicAmplitudeAwarePermutation) where {m, T}
    πs = symbolize(x, SymbolicPermutation(m = m, lt = est.lt)) # motif length controlled by dimension of input data
    wts = AAPE.(x.data, A = est.A, m = est.m)

    Probabilities(symprobs(πs, wts, normalize = true))
end

function probabilities(x::AbstractVector{T}, est::SymbolicAmplitudeAwarePermutation) where {T<:Real}
    τs = tuple([est.τ*i for i = 0:est.m-1]...)
    emb = genembed(x, τs)
    πs = symbolize(emb, SymbolicPermutation(m = est.m, lt = est.lt))  # motif length controlled by estimator m
    wts = AAPE.(emb.data, A = est.A, m = est.m)
    p = symprobs(πs, wts, normalize = true)
    Probabilities(p)
end


# basic figure with julia

@time using Plots  # 6.9s
@time pyplot()  # 5.2s

t = collect(0.0:0.01:2.0)
s = sin.(2*pi*t)
c = cos.(2*pi*t)

@time f = plot(t, s, label="sin") # 11s
plot!(t, c, label="cos")
plot!(xlab="time (s)", ylab="voltage (mV)")
plot!(legend=:bottomleft)

@time savefig("test.png") # 5.8s

display(f)

println("done. press <ENTER> to quit")
readline()


mutable struct NeuronResponse
    name::String
    get_spikes::Function
end

function get_spikes_from_distribution(d)
    N -> (filter(x->first(τ)<x<last(τ), sort(rand(d,2N)[1:N])))
end

# smooth ramp
function ramp(τ)
    mt = maximum(τ)
    f(t) = begin
        if t < 0
            return 0.1
        else
            return t
        end
    end
    [f(t/mt) for t in τ]
end

# oscillate after
function oscillate(τ)
    mt = maximum(τ)
    f(t) = begin
        if t < 0
            return 0.5
        else
            return sin(π*t)+0.5
        end
    end
    [f(t) for t in τ]
end

# short response and refractory period
function refractory(τ)
    mt = maximum(τ)
    f(t) = begin
        if t < 0
            return 0.5
        elseif t < 0.05
            return 0.95
        elseif t < 0.1
            return 0.05
        else
            return 0.5
        end
    end
    [f(t) for t in τ]
end


#=

Author: Shunsuke Hori

=#

module CES

using LS

function baseline(;
    beta = 0.9,
    sigma = 2.0,
    gamma = 2.0,
    Pi = 0.5 *ones(2,2),
    G = [0.1, 0.2],
    Theta = ones(Float64, 2),
    transfers = false
    )
    function U(c,n)
        if sigma == 1.0
            U = log(c)
        else
            U = (c.^(1.0-sigma)-1.0)/(1.0-sigma)
        end
        return U - n.^(1+gamma)/(1+gamma)
    end
    #derivatives of utiltiy function
    Uc(c,n) =  c.^(-sigma)
    Ucc(c,n) = -sigma*c.^(-sigma-1.0)
    Un(c,n) = -n.^gamma
    Unn(c,n) = -gamma*n.^(gamma-1.0)
    n_less_than_one = false
    return LS.Para(beta, Pi, G, Theta, transfers,
                    U, Uc, Ucc, Un, Unn, n_less_than_one)
end

#Model 1
M1 = baseline()

#Model 2

M2 = baseline(G=[0.15], Pi=ones(1,1), Theta=[1.0])

#Model 3 with time varying

M_time_example = baseline(G=[0.1, 0.1, 0.1, 0.2, 0.1, 0.1],
                            Theta = ones(6)# Theta can in principle be random
                            )

M_time_example.Pi = [0.0 1.0 0.0 0.0 0.0 0.0;
                     0.0 0.0 1.0 0.0 0.0 0.0;
                     0.0 0.0 0.0 0.5 0.5 0.0;
                     0.0 0.0 0.0 0.0 0.0 1.0;
                     0.0 0.0 0.0 0.0 0.0 1.0;
                     0.0 0.0 0.0 0.0 0.0 1.0]
end


using Turing
using Turing.RandomMeasures
using Test
using Random

dir = splitdir(splitdir(pathof(Turing))[1])[1]
include(dir*"/test/test_utils/AllUtils.jl")

@testset "RandomMeasures.jl" begin
    partitions = [
        [[1, 2, 3, 4]],
        [[1, 2, 3], [4]],
        [[1, 2, 4], [3]],
        [[1, 2], [3, 4]],
        [[1, 2], [3], [4]],
        [[1, 3, 4], [2]],
        [[1, 3], [2, 4]],
        [[1, 3], [2], [4]],
        [[1, 4], [2, 3]],
        [[1], [2, 3, 4]],
        [[1], [2, 3], [4]],
        [[1, 4], [2], [3]],
        [[1], [2, 4], [3]],
        [[1], [2], [3, 4]],
        [[1], [2], [3], [4]]]

    @turing_testset "chinese restaurant processes" begin
        # Data
        data = [-2,2,-1.5,1.5]

        # Base distribution
        mu_0 = mean(data)
        sigma_0 = 4
        sigma_1 = 0.5
        tau0 = 1/sigma_0^2
        tau1 = 1/sigma_1^2

        # DP parameters
        alpha = 0.25

        @model crpimm(y, rpm) = begin
            # Base distribution.
            H = Normal(mu_0, sigma_0)
            # Latent assignments.
            N = length(y)
            z = tzeros(Int, N)

            # Cluster counts.
            cluster_counts = tzeros(Int, N)

            # Cluster locations.
            x = tzeros(Float64, N)
            for i in 1:N
                # Draw assignments using a CRP.
                z[i] ~ ChineseRestaurantProcess(rpm, cluster_counts)
                if cluster_counts[z[i]] == 0
                    # Cluster is new, therefore, draw new location.
                    x[z[i]] ~ H
                end
                cluster_counts[z[i]] += 1
                # Draw observation.
                y[i] ~ Normal(x[z[i]], sigma_1)
            end
        end

        rpm = DirichletProcess(alpha)

        sampler = SMC(1000)
        mf = crpimm(data, rpm)

        # Compute empirical posterior distribution over partitions
        samples = sample(mf, sampler)

        # Check that there is no NaN value associated
        z_samples = Int.(samples[:z].value)
        @test all(!isnan, samples[:x].value[z_samples])
        @test all(!ismissing, samples[:x].value[z_samples])

        empirical_probs = zeros(length(partitions))
        w = map(x -> x.weight, samples.info[:samples])
        sum_weights = sum(w)
        z = z_samples

        for i in 1:size(z,1)
            partition = map(c -> findall(z[i,:,1] .== c), unique(z[i,:,1]))
            partition_idx = findfirst(p -> sort(p) == sort(partition), partitions)
            @test partition_idx != nothing
            empirical_probs[partition_idx] += sum_weights == 0 ? 1 : w[i]
        end

        if sum_weights == 0
            empirical_probs /= length(w)
        end

        l2, discr = correct_posterior(empirical_probs, data, partitions, tau0, tau1, alpha, 1e-7)

        @test l2 < 0.05
        @test discr < 0.2
    end
    @testset "distributions" begin
        @turing_testset "Representations" begin
            d = StickBreakingProcess(DirichletProcess(1.0))
            @test minimum(d) == 0
            @test maximum(d) == 1

            d = SizeBiasedSamplingProcess(DirichletProcess(1.0), 2.0)
            @test minimum(d) == 0
            @test maximum(d) == 2

            d = ChineseRestaurantProcess(DirichletProcess(1.0), [2, 1])
            @test minimum(d) == 1
            @test maximum(d) == 3
        end
        @turing_testset "Dirichlet Process" begin

            α = 0.1
            N = 10_000

            # test SB representation
            d = StickBreakingProcess(DirichletProcess(α))
            Ev = mapreduce(_ -> rand(d), +, 1:N) / N
            @test Ev ≈ mean(Beta(1, α)) atol=0.05

            # test SBS representation
            d = SizeBiasedSamplingProcess(DirichletProcess(α), 2.0)
            Ej = mapreduce(_ -> rand(d), +, 1:N) / N
            @test Ej ≈ mean(Beta(1, α)) * 2 atol=0.05

            # test CRP representation
            d = ChineseRestaurantProcess(DirichletProcess(α), [2, 1])
            ks = map(_ -> rand(d), 1:N)
            c = map(k -> sum(ks .== k), support(d))
            p = c ./ sum(c)

            q = [2, 1, α] ./ (2 + α)
            q ./= sum(q)

            @test p[1] ≈ q[1] atol=0.1
            @test p[2] ≈ q[2] atol=0.1
            @test p[3] ≈ q[3] atol=0.1
        end
        @turing_testset "Pitman-Yor Process" begin

            a = 0.5
            θ = 0.1
            t = 2
            N = 10_000

            # test SB representation
            d = StickBreakingProcess(PitmanYorProcess(a, θ, t))
            Ev = mapreduce(_ -> rand(d), +, 1:N) / N
            @test Ev ≈ mean(Beta(1 - a, θ + t*a)) atol=0.05

            # test SBS representation
            d = SizeBiasedSamplingProcess(PitmanYorProcess(a, θ, t), 2.0)
            Ej = mapreduce(_ -> rand(d), +, 1:N) / N
            @test Ej ≈ mean(Beta(1 - a, θ + t*a)) * 2 atol=0.05

            # test CRP representation
            d = ChineseRestaurantProcess(PitmanYorProcess(a, θ, t), [2, 1])
            ks = map(_ -> rand(d), 1:N)
            c = map(k -> sum(ks .== k), support(d))
            p = c ./ sum(c)

            q = [2 - a, 1 - a, θ + a*t] ./ (3 + θ)
            q ./= sum(q)

            @test p[1] ≈ q[1] atol=0.1
            @test p[2] ≈ q[2] atol=0.1
            @test p[3] ≈ q[3] atol=0.1
        end
    end
    @turing_testset "stick breaking" begin
        # Data
        data = [-2,2,-1.5,1.5]

        # Base distribution
        mu_0 = mean(data)
        sigma_0 = 4
        sigma_1 = 0.5
        tau0 = 1/sigma_0^2
        tau1 = 1/sigma_1^2

        # DP parameters
        alpha = 0.25

        # stick-breaking process based on Papaspiliopoulos and Roberts (2008).
        @model sbimm(y, rpm, trunc) = begin
            # Base distribution.
            H = Normal(mu_0, sigma_0)

            # Latent assignments.
            N = length(y)
            z = tzeros(Int, N)

            # Infinite (truncated) collection of breaking points on unit stick.
            v = tzeros(Float64, trunc)

            # Cluster locations.
            x = tzeros(Float64, trunc)

            # Draw weights and locations.
            for k in 1:trunc
                v[k] ~ StickBreakingProcess(rpm)
                x[k] ~ H
            end

            # Weights.
            w = vcat(v[1], v[2:end] .* cumprod(1 .- v[1:end-1]))

            # Normalize weights to ensure they sum exactly to one.
            # This is required by the Categorical distribution in Distributions.
            w ./= sum(w)

            for i in 1:N
                # Draw location
                z[i] ~ Categorical(w)

                # Draw observation.
                y[i] ~ Normal(x[z[i]], sigma_1)
            end
        end

        rpm = DirichletProcess(alpha)

        sampler = SMC(1000)
        mf = sbimm(data, rpm, 10)

        # Compute empirical posterior distribution over partitions
        samples = sample(mf, sampler)

        # Check that there is no NaN value associated
        z_samples = Int.(samples[:z].value)
        @test all(!isnan, samples[:x].value[z_samples])
        @test all(!ismissing, samples[:x].value[z_samples])

        empirical_probs = zeros(length(partitions))
        w = map(x -> x.weight, samples.info[:samples])
        sum_weights = sum(w)
        z = z_samples

        for i in 1:size(z,1)
            partition = map(c -> findall(z[i,:,1] .== c), unique(z[i,:,1]))
            partition_idx = findfirst(p -> sort(p) == sort(partition), partitions)
            @test partition_idx != nothing
            empirical_probs[partition_idx] += sum_weights == 0 ? 1 : w[i]
        end

        if sum_weights == 0
            empirical_probs /= length(w)
        end

        l2, discr = correct_posterior(empirical_probs, data, partitions, tau0, tau1, alpha, 1e-7)

        # Increased ranges due to truncation error.
        @test l2 < 0.1
        @test discr < 0.3
    end
    @turing_testset "size-based sampling" begin
        # Data
        data = [-2,2,-1.5,1.5]

        # Base distribution
        mu_0 = mean(data)
        sigma_0 = 4
        sigma_1 = 0.5
        tau0 = 1/sigma_0^2
        tau1 = 1/sigma_1^2

        # DP parameters
        alpha = 0.25

        # size-biased sampling process
        @model sbsimm(y, rpm, trunc) = begin
            # Base distribution.
            H = Normal(mu_0, sigma_0)

            # Latent assignments.
            N = length(y)
            z = tzeros(Int, N)

            x = tzeros(Float64, N)
            J = tzeros(Float64, N)
            z = tzeros(Int, N)

            k = 0
            surplus = 1.0

            for i in 1:N
                ps = vcat(J[1:k], surplus)
                z[i] ~ Categorical(ps)
                if z[i] > k
                    k = k + 1
                    J[k] ~ SizeBiasedSamplingProcess(rpm, surplus)
                    x[k] ~ H
                    surplus -= J[k]
                end
                y[i] ~ Normal(x[z[i]], sigma_1)
            end
        end

        rpm = DirichletProcess(alpha)

        sampler = SMC(1000)
        mf = sbsimm(data, rpm, 100)

        # Compute empirical posterior distribution over partitions
        samples = sample(mf, sampler)

        # Check that there is no NaN value associated
        z_samples = Int.(samples[:z].value)
        @test all(!isnan, samples[:x].value[z_samples])
        @test all(!ismissing, samples[:x].value[z_samples])

        empirical_probs = zeros(length(partitions))
        w = map(x -> x.weight, samples.info[:samples])
        sum_weights = sum(w)
        z = z_samples

        for i in 1:size(z,1)
            partition = map(c -> findall(z[i,:,1] .== c), unique(z[i,:,1]))
            partition_idx = findfirst(p -> sort(p) == sort(partition), partitions)
            @test partition_idx != nothing
            empirical_probs[partition_idx] += sum_weights == 0 ? 1 : w[i]
        end

        if sum_weights == 0
            empirical_probs /= length(w)
        end

        l2, discr = correct_posterior(empirical_probs, data, partitions, tau0, tau1, alpha, 1e-7)

        @test l2 < 0.05
        @test discr < 0.2
    end
end


"""
A Julia interface module for SciPy

- [GitHub repository: AtsushiSakai/SciPy.jl: Julia interface for SciPy](https://github.com/AtsushiSakai/SciPy.jl)
- [Stable Doc](https://atsushisakai.github.io/SciPy.jl/stable)
- [Dev Doc](https://atsushisakai.github.io/SciPy.jl/dev)

"""
module SciPy

using PyCall
using InteractiveUtils

export print_configulations
export cluster, constants, fft, integrate, interpolate, io, linalg, ndimage, odr
export optimize, signal, sparse, spatial, stats, special
# raw PyObject modules
export pystats, pyoptimize, pyinterpolate, pyspatial

const scipy = PyNULL()

include("util.jl")

"""
scipy.cluster module

- [Clustering package (scipy.cluster) Reference Guide](https://docs.scipy.org/doc/scipy/reference/cluster.html)

# Examples

You can do k-means clustering using this module:

```julia-repl
julia> features  = [[ 1.9  2.3];
                    [ 1.5 2.5];
                    [ 0.8 0.6];
                    [ 0.4 1.8];
                    [ 0.1 0.1];
                    [ 0.2 1.8];
                    [ 2.0 0.5];
                    [ 0.3 1.5];
                    [ 1.0 1.0]]
9×2 Array{Float64,2}:
 1.9  2.3
 1.5  2.5
 0.8  0.6
 0.4  1.8
 0.1  0.1
 0.2  1.8
 2.0  0.5
 0.3  1.5
 1.0  1.0

julia> whitened = SciPy.cluster.vq.whiten(features)
9×2 Array{Float64,2}:
 2.7396    2.91001
 2.16284   3.16306
 1.15351   0.759134
 0.576757  2.2774
 0.144189  0.126522
 0.288379  2.2774
 2.88379   0.632612
 0.432568  1.89784
 1.44189   1.26522

julia> SciPy.cluster.vq.kmeans(whitened, [whitened[1,:] whitened[3,:]] )
([1.1174670798453024 1.8345740800894272; 2.8837860125040065 0.6326117517549749], 1.073399
3090584457)

```

"""
const cluster = PyNULL()

"""
scipy.constants module

- [Constants (scipy.constants) Reference Guide](https://docs.scipy.org/doc/scipy/reference/constants.html)

# Examples

You can access each constants:

```julia-repl
julia> SciPy.constants.golden
1.618033988749895

julia> SciPy.constants.physical_constants["electron mass"]
(9.10938356e-31, "kg", 1.1e-38)

julia> SciPy.constants.convert_temperature([-40, 40.0], "Celsius", "Kelvin")

2-element Array{Float64,1}:
 233.14999999999998
 313.15

```
"""
const constants = PyNULL()

"""
scipy.fft module

- [Discrete Fourier transforms (scipy.fft) Reference Guide](https://docs.scipy.org/doc/scipy/reference/fft.html)

# Examples

You can use FFT (Fast Fourier Transform):

```julia-repl

julia> SciPy.fft.fft(exp.(π/8 * collect(1:8)))
8-element Array{Complex{Float64},1}:
   68.17385416403044 - 0.0im
   1.408601300061675 + 31.248171041435185im
 -10.268363617931092 + 15.207165888808841im
 -12.695027025520982 + 6.493878653648949im
 -13.216494113330363 - 0.0im
 -12.695027025520982 - 6.493878653648949im
 -10.268363617931092 - 15.207165888808841im
   1.408601300061675 - 31.248171041435185im

```
"""
const fft = PyNULL()

"""
scipy.integrate module

- [Integration and ODEs (scipy.integrate) Reference Guide](https://docs.scipy.org/doc/scipy/reference/integrate.html)

# Examples

You can compute a definite integral:

```julia-repl
julia> f(x) = x^2
f (generic function with 1 method)

julia> SciPy.integrate.quad(f, 0, 4)
(21.333333333333336, 2.368475785867001e-13)

```
"""
const integrate = PyNULL()

const pyinterpolate = PyNULL()
include("interpolate.jl")

"""
scipy.io module

- [Input and output (scipy.io) Reference Guide](https://docs.scipy.org/doc/scipy/reference/io.html)

# Examples

You can save a MATLAB-style .mat file:

```julia-repl
julia> mdic = Dict([("a", 100), ("label", "experiment")])
Dict{String,Any} with 2 entries:
  "label" => "experiment"
  "a"     => 100

julia> SciPy.io.savemat("sample_data.mat", mdic)
```
"""
const io = PyNULL()

"""
scipy.linalg module

- [Linear algebra (scipy.linalg) Reference Guide](https://docs.scipy.org/doc/scipy/reference/linalg.html)

# Examples

You can compute the matrix exponential using Pade approximation:

```julia-repl
julia> SciPy.linalg.expm(zeros((2,2)))
2×2 Array{Float64,2}:
 1.0  0.0
 0.0  1.0
```
"""
const linalg = PyNULL()

"""
scipy.ndimage module

- [Multi-dimensional image processing (scipy.ndimage) Reference Guide](https://docs.scipy.org/doc/scipy/reference/ndimage.html)

# Examples

You can compute a multidimensional convolution:

```julia-repl
julia> k = [[1 1 1];[1 1 0];[1 0 0]]
3×3 Array{Int64,2}:
 1  1  1
 1  1  0
 1  0  0

julia> a = [[1 2 0 0];
            [5 3 0 4];
            [0 0 0 7];
            [9 3 0 0]]
4×4 Array{Int64,2}:
 1  2  0  0
 5  3  0  4
 0  0  0  7
 9  3  0  0

julia> SciPy.ndimage.convolve(a, k, mode="constant", cval=0.0)
4×4 Array{Int64,2}:
 11  10   7   4
 10   3  11  11
 15  12  14   7
 12   3   7   0
```
"""
const ndimage = PyNULL()

"""
scipy.odr module

- [Orthogonal distance regression (scipy.odr) Reference Guide](https://docs.scipy.org/doc/scipy/reference/odr.html)

# Examples

You can calculate orthogonal distance regression with an exponential model:

```julia-repl
julia> x = collect(0.0:5.0);

julia> y = -10.0 .+ exp.(0.5*x);

julia> data = SciPy.odr.Data(x, y)
PyObject <scipy.odr.odrpack.Data object at 0x7fe5fda4ccc0>

julia> data = SciPy.odr.Data(x, y);

julia> odr_obj = SciPy.odr.ODR(data, SciPy.odr.exponential);

julia> output = odr_obj.run();

julia> println(output.beta)
[-10.0, 0.5]
```
"""
const odr = PyNULL()

const pyoptimize = PyNULL()
include("optimize.jl")

"""
scipy.signal module

- [Signal processing (scipy.signal) Reference Guide](https://docs.scipy.org/doc/scipy/reference/signal.html)

# Examples

You can compute the Kaiser parameter beta, given the attenuation a:

```julia-repl
julia> SciPy.signal.kaiser_beta(65)
6.20426

```
"""
const signal = PyNULL()

"""
scipy.sparse module

- [Sparse matrices (scipy.sparse) Reference Guide](https://docs.scipy.org/doc/scipy/reference/sparse.html)

# Examples

You can do sparse matrix calculation:

```julia-repl
julia> A = SciPy.sparse.csc_matrix([[1.0 0.0];
                                   [1.0 2.0]]);

julia> Ainv = SciPy.sparse.linalg.inv(A);

julia> A.dot(Ainv).todense()
2×2 Array{Float64,2}:
 1.0  0.0
 0.0  1.0

```
"""
const sparse = PyNULL()

const pyspatial = PyNULL()
include("spatial.jl")

"""
scipy.special module

- [Special functions (scipy.special) Reference Guide](https://docs.scipy.org/doc/scipy/reference/special.html)

"""
const special = PyNULL()

const pystats = PyNULL()
include("stats.jl")

"""
Module initialization function
"""
function __init__()
    copy!(scipy, pyimport_conda("scipy", "scipy"))
    copy!(cluster, pyimport_conda("scipy.cluster", "scipy"))
    copy!(constants, pyimport_conda("scipy.constants", "scipy"))
    copy!(fft, pyimport_conda("scipy.fft", "scipy"))
    copy!(integrate, pyimport_conda("scipy.integrate", "scipy"))
    copy!(io, pyimport_conda("scipy.io", "scipy"))
    copy!(linalg, pyimport_conda("scipy.linalg", "scipy"))
    copy!(ndimage, pyimport_conda("scipy.ndimage", "scipy"))
    copy!(odr, pyimport_conda("scipy.odr", "scipy"))
    copy!(signal, pyimport_conda("scipy.signal", "scipy"))
    copy!(sparse, pyimport_conda("scipy.sparse", "scipy"))
    copy!(special, pyimport_conda("scipy.special", "scipy"))

    copy!(pyspatial, pyimport_conda("scipy.spatial", "scipy"))
    copy!(pystats, pyimport_conda("scipy.stats", "scipy"))
    copy!(pyoptimize, pyimport_conda("scipy.optimize", "scipy"))
    copy!(pyinterpolate, pyimport_conda("scipy.interpolate", "scipy"))
end

"""
Print configulations:
- Julia version
- Python version
- Python path
- scipy version
"""
function print_configulations()
    println("==== SciPy.jl configulations ====")
    println("-- Julia settings --")
    versioninfo()
    println("-- PyCall.jl settings --")
    @show PyCall.pyversion
    @show PyCall.libpython
    @show PyCall.pyprogramname
    @show PyCall.conda
    println("-- SciPy settings --")
    @show scipy.__version__
    @show scipy.version.full_version
    @show scipy.version.git_revision
    return nothing
end

end


export loc2cs,cs2loc

function loc2cs(loc::Array{Int64,1},n::Array{Int64,1})
@inbounds cs = loc[1] + (loc[2]-1)*n[1];
return cs;
end

function loc2cs(loc1::Int64,loc2::Int64,n::Array{Int64,1})
@inbounds cs = loc1 + (loc2-1)*n[1];
return cs;
end

function cs2loc(loc::Array{Int64,1},cs_loc::Int64,n::Array{Int64,1})
@inbounds loc[1] = mod(cs_loc-1,n[1]) + 1;
@inbounds loc[2] = div(cs_loc-1,n[1])+1;
return;
end

function cs2loc(cs_loc::Int64,n::Array{Int64,1})
@inbounds loc1 = mod(cs_loc-1,n[1]) + 1;
@inbounds loc2 = div(cs_loc-1,n[1])+1;
return loc1,loc2;
end

function get2DNeighbors(loc::Array{Int64,1},n::Array{Int64,1},ans::Array{Int64,2},relevant::Array{Bool,1})
# This function fills the neighbours array ans with neighbours the "C" style, so it doesn't allocate.
@inbounds relevant[1]= loc[1]>1;
@inbounds relevant[2]= loc[1]<n[1];
@inbounds relevant[3]= loc[2]>1;
@inbounds relevant[4]= loc[2]<n[2];

@inbounds ans[1,1] = loc[1]-1;
@inbounds ans[1,2] = loc[2];
@inbounds ans[2,1] = loc[1]+1;
@inbounds ans[2,2] = loc[2];
@inbounds ans[3,1] = loc[1];
@inbounds ans[3,2] = loc[2]-1;
@inbounds ans[4,1] = loc[1];
@inbounds ans[4,2] = loc[2]+1;
return;
end

function setPaddedArr2D(arr::Array{Bool,2},loc1::Int64,loc2::Int64,val::Bool)
# this function assumes that the boundaries are OK and the value is placed in a legal location.
@inbounds arr[loc1+2,loc2+2] = val
end

function getPaddedArr2D(arr::Array{Bool,2},loc1::Int64,loc2::Int64)
# this function assumes that the boundaries are OK and the value is placed in a legal location.
@inbounds return arr[loc1+2,loc2+2]
end

function ddxLongDiff(h::Float64,n::Int64)
	dm1 = (-1/(2*h)*ones(n-1));
	dm1[n-1] = -1/h;
	d   = zeros(n);
	d[1] = -1/h;
	d[n] = 1/h;
	dp1 = (1/(2*h))*ones(n-1);
	dp1[1] = 1/h;
	D = spdiagm((dm1,d,dp1),(-1,0,1),n,n);
	return D;
end


### chessboard

tile_size = 50
board_size = 4 # must be even
mini_board = [zeros(tile_size, tile_size) ones(tile_size, tile_size);
              ones(tile_size, tile_size) zeros(tile_size, tile_size)]

chessboard = repeat(mini_board, outer=(convert(Integer,(board_size/2)), convert(Integer,(board_size/2))))
img = chessboard

img = gen_chess();

###

f(x :: Array{<:Real, 2}) = println(x)

Array{Float64, 2} <: Array{Real, 2}

Float64 <: Real

function g()
    ar::Array{Real, 2} = ones(1,1)
end

g()

f(ones(1, 1))

###

using TestImages
using PyPlot

img = testimage("mandril_gray");
img = reverse(img, dims=1)
img = chessboard

point_count = 100;
spacing = 30;

coverage = 1;

(pi * spacing^2 * point_count)/length(img)

r = sqrt(( * ))

points, mag, mag_points = LAP_julia.gradient_points.find_keypoints_from_gradients(img, sigma = 5, number = point_count, spacing = spacing)

imgshow(mag_points)
imgshow(mag)

pos_x = [point.pos[1] for point in points]
pos_y = [point.pos[2] for point in points]

imgshow(img)
PyPlot.scatter(pos_y, pos_x, marker = :x)
gcf()

### probability
using PyPlot

img = testimage("mandril_gray");
img = chessboard

# all(vec(reshape(img, (1, :))) == img[:])

points, mag, mag_points = LAP_julia.gradient_points.find_keypoints_from_gradients_p_field(img, sigma = 1, number = 100, spacing = 30)

wsample(Array(1:2), [-0, 1], 12)

imgshow(mag_points)
imgshow(mag)

pos_x = [point.pos[1] for point in points]
pos_y = [point.pos[2] for point in points]

imgshow(img)
PyPlot.scatter(pos_y, pos_x, marker = :x)
gcf()

###
using LAP_julia

img = gen_chess()

# genrate flow
flow = gen_tiled_flow(size(img), 20, 1000);

showflow(flow)

# generate warpped image
imgw = LAP_julia.interpolation.warp_img(img, real(flow), imag(flow));

# params:
#lap
image_1 = img;
image_2 = imgw;
filter_num = 3;
filter_half_size = 32;
window_size = [65, 65];
#points
point_count = 100;
spacing = 10;

# get points
points, mag, mag_points = find_keypoints_from_gradients(img, sigma = 1, number = point_count, spacing = spacing);

using PyPlot
# see points
pos_x = [point.pos[1] for point in points]
pos_y = [point.pos[2] for point in points]

pos = transpose([pos_x pos_y])
pos[:, 1]

imgshow(img)
PyPlot.scatter(pos_y, pos_x, marker = :x)
gcf()

u_est, coeffs = single_lap(image_1, image_2, filter_half_size, window_size, 3, pos);

showflow(u_est)

### try imfilter with indices:
img = chessboard;

filter_size = 21
filter_half_size = (filter_size - 1) / 2
# Calculate separable filters from basis:
sigma = (filter_half_size + 2) / 4
gaus = ImageFiltering.KernelFactors.gaussian(sigma, filter_size)

using PyPlot
using Plots
# 1
kernel = kernelfactors((gaus, gaus))
kernp1 = broadcast(*, kernel...)
surface([-filter_half_size:filter_half_size], -filter_half_size:filter_half_size, kernp1)

@time imgfilt1 = ImageFiltering.imfilter(chessboard, kernel, "symmetric")
imgshow(imgfilt1)
surface(imgfilt1)

# @time imgfilt1 = ImageFiltering.imfilter(chessboard, kernel, "symmetric", inds=[1, 1])

imgfilt = similar(img)

inds = CartesianIndex(50, 50)

imfilter!(CPU1(Algorithm.FIRTiled()), imgfilt, img, kernel, "symmetric", inds)

imfilter!(CPU1(Algorithm.FIR()), imgfilt, img, (), NoPad(), inds)

imgshow(imgfilt)
surface(imgfilt1)

### exapmle
using ImageFiltering, Images, ComputationalResources

# Create a sample image:
tile_size = 50
board_size = 4 # must be even
mini_board = [zeros(tile_size, tile_size) ones(tile_size, tile_size);
              ones(tile_size, tile_size) zeros(tile_size, tile_size)]
chessboard = repeat(mini_board, outer=(convert(Integer,(board_size/2)), convert(Integer,(board_size/2))))
img = chessboard

# Create a filter:
sigma = 3
filter_size = 21
gaus = ImageFiltering.KernelFactors.gaussian(sigma, filter_size)
kernel = kernelfactors((gaus, gaus))

typeof(kernel) <: ImageFiltering.ProcessedKernel
typeof(kernel) <: Union{ImageFiltering.ArrayLike, ImageFiltering.Laplacian}

# params of imfilter function
imgfilt = copy(img)
# inds = CartesianIndex(50, 50)

# The filtering that errors:
# imfilter!(CPU1(Algorithm.FIRTiled()), imgfilt, img, kernel, "symmetric", (30:40, 30:40))
imfilter!(CPU1(Algorithm.FIRTiled()), imgfilt, img, kernel, Pad(:Symmetric), (50:50, 50:50))

imgshow(imgfilt)

### alternative to imfilter inds
function f(img_pix)
    sum(img_pix .* [0 1/10 0;
                    1/10 6/10 1/10;
                    0 1/10 0])
end

mapped_img = mapwindow(f, img, (3, 3), "symmetric", (2:8, 2:8))

imgshow(mapped_img)

###

flow = gen_tiled_flow(size(img), 30, 1000)

one = [1 + 0im, 2 - im]
two = [0 + 6im, 3 + im]

one = [1 + 0im]
two = [0 + im]

import LAP_julia: helpers.angle_rmse

angle_rmse(one, two)

rad2deg.(angle.(one).-angle.(two))

### mse

function mse2(x, y)
    @assert (size(x) == size(y)) "sizes of $x and $y don't match"
    return 1/length(x) * sum((x .- y).^2)
end

function mse1(x, y)
    @assert (size(x) == size(y)) "sizes of $x and $y don't match"
    return mean((x .- y).^2)
end

function mean(x)
    sum(x)/length(x)
end

ran = rand(123)
dan = rand(123)

mse1(ran, dan)
mse2(ran, dan)
mse1(ran, dan) ≈ mse2(ran, dan)

####


function gen_newton_functions(solver::Solver)
    n,m,N = get_sizes(solver)
    p,pI,pE = get_num_constraints(solver)
    p_N,pI_N,pE_N = get_num_terminal_constraints(solver)

    c_function!, c_jacobian!, c_labels, cI!, cE! = generate_constraint_functions(solver.obj)
    f = solver.fd
    s_x = reshape(1:(N*n),n,N)
    s_u = reshape(N*n .+ (1:(m*(N-1))),m,N-1)
    s_λ = (N*n + (N-1)*m) .+ (1:((N-1)*pE + pE_N))
    s_μ = (N*n + (N-1)*m + length(s_λ)) .+ (1:((N-1)*pI + pI_N))
    s_ν = (N*n + (N-1)*m + length(s_λ) + length(s_μ)) .+ (1:N*n)
    s_s = (N*n + (N-1)*m + length(s_λ) + length(s_μ) + N*n) .+ (1:length(s_μ))
    sze = (X=s_x, U=s_u, λ=s_λ, μ=s_μ, ν=s_ν, s=s_s)

    function newton_cost(V::Vector)
        X,U,λ,μ,ν,s = unpackV(V)
        J = cost(solver,X,U) + λ'cE(V) + μ'*(cI(V) + 0.5*s.^2) + ν'd(V)
        return J
    end

    function cI(V)
        X,U,λ,μ,ν,s = unpackV(V)
        C = zeros(eltype(V),pI,N-1)
        for k = 1:N-1
            cI!(view(C,1:pI,k),X[:,k],U[:,k])
        end
        CN = zeros(eltype(V),pI_N)
        cI!(CN,X[:,N])
        CI = [vec(C); CN]
        return CI
    end

    function cE(V)
        X,U,λ,μ,ν,s = unpackV(V)
        C = zeros(eltype(V),pE,N-1)
        for k = 1:N-1
            cE!(view(C,1:pE,k),X[:,k],U[:,k])
        end
        CN = zeros(eltype(V),pE_N)
        cE!(CN,X[:,N])
        CE = [vec(C); CN]
        return CE
    end

    function d(V)
        X,U,λ,μ,ν,s = unpackV(V)
        D = zeros(eltype(V),n,N)
        D[:,1] = X[:,1] - solver.obj.x0
        for k = 2:N
            f(view(D,1:n,k),X[:,k-1],U[:,k-1])
            D[:,k] -= X[:,k]
        end
        return vec(D)
    end

    function active_set(V,eps=1e-2)
        X,U,λ,μ,ν,s = unpackV(V)
        a = ones(Bool,length(V))
        ci = cI(V)
        c_inds = ones(Bool,length(ci))
        c_inds = ci .>= -eps

        a[sze.μ] = c_inds
        a[sze.s] = c_inds
        return a
    end

    function packV(X,U,λ,μ,ν,s)
        V = [vec(X);vec(U);λ;μ;ν;s]
    end

    function unpackV(V)
        X = V[sze.X]
        U = V[sze.U]
        λ = V[sze.λ]
        μ = V[sze.μ]
        ν = V[sze.ν]
        s = V[sze.s]
        return X,U,λ,μ,ν,s
    end
    return newton_cost, packV, unpackV, cI, cE, d, active_set, sze
end

model,obj = Dynamics.dubinscar_parallelpark
solver = Solver(model,obj,N=5)
solver.opts.constraint_tolerance = 1e-4
n,m,N = get_sizes(solver)
U0 = rand(m,N-1)
res,stats = solve(solver,U0)
cost(solver,res)
λ_update_default!(res,solver);
@assert max_violation(res) < solver.opts.constraint_tolerance
max_violation(res)

newton_cost, packV, unpackV, cI, cE, d, active_set = gen_newton_functions(solver)
max_c(V) = max(maximum(cI(V)),norm(cE(V),Inf),norm(d(V),Inf))
p,pI,pE = get_num_constraints(solver)
pN,pI_N,pE_N = get_num_terminal_constraints(solver)

X = to_array(res.X)
U = to_array(res.U)
μ = vcat([k < N ? res.λ[k][1:pI] : res.λ[k][1:pI_N] for k = 1:N]...)
λ = vcat([k < N ? res.λ[k][pI .+ (1:pE)] : res.λ[k][pI_N .+ (1:pE_N)] for k = 1:N]...)
ν = vec(to_array(res.s))
s = vcat([k < N ? res.C[k][1:pI] : res.C[k][1:pI_N] for k = 1:N]...)
s = sqrt.(2*max.(0,-s))

V = packV(X,U,λ,μ,ν,s)

J1 = newton_cost(V)
g = ForwardDiff.gradient(newton_cost,V)
H = ForwardDiff.hessian(newton_cost,V)
# V = line_search(V,H,g)

a = active_set(V,1e-3)
rank(H)
size(H)
rank(H[a,a])
size(H[a,a])
ρ = 1.0
newton_cost(V - (H + ρ*I)\g)
max_c(V - (H + ρ*I)\g)

X,U,λ,μ,ν,s == unpackV(V)
ν
# newton_cost(V)
# d(V)
# norm(d(V),Inf)
# X,U,λ,μ,ν,s == unpackV(V)
# d(V)
#
# function line_search(V,H,g)
#     J0 = newton_cost(V)
#     α = 1.
#     V_ = V - α*(H\g)
#     J = newton_cost(V_)
#     iter = 1
#     while J > J0
#         α /= 2
#         V_ = V - α*(H\g)
#         J = newton_cost(V_)
#         iter += 1
#         @show J, α
#         iter > 10 ? break : nothing
#     end
#     @show J0-J
#     return V_
# end

max_c(V)
findmax(cI(V))
cI(V)
X,U,λ,μ,ν,s == unpackV(V)
μ[8]
norm(μ.*s,Inf)
d(V)

# Block move
opts = TrajectoryOptimization.SolverOptions()
opts.verbose = false
opts.cost_tolerance = 1e-4
opts.cost_tolerance_intermediate = 1e-4
opts.constraint_tolerance = 1e-4
opts.square_root = false
opts.active_constraint_tolerance = 0.0
opts.outer_loop_update_type = :default
opts.live_plotting = false

# Block move
model, obj = TrajectoryOptimization.Dynamics.double_integrator
u_min = -0.2
u_max = 0.2
obj_con = TrajectoryOptimization.ConstrainedObjective(obj, tf=5.0,use_xf_equality_constraint=true,u_min=u_min, u_max=u_max)#, x_max=x_max)

solver = TrajectoryOptimization.Solver(model,obj_con,integration=:rk4,N=5,opts=opts)
U = rand(solver.model.m, solver.N)

res, stats = TrajectoryOptimization.solve(solver,U)
@assert max_violation(res) < opts.constraint_tolerance
max_violation(res)

newton_cost, packV, unpackV, cI, cE, d, active_set = gen_newton_functions(solver)
max_c(V) = max(maximum(cI(V)),norm(cE(V),Inf),norm(d(V),Inf))
p,pI,pE = get_num_constraints(solver)
pN,pI_N,pE_N = get_num_terminal_constraints(solver)

X = to_array(res.X)
U = to_array(res.U)
μ = vec(to_array(res.λ[1:solver.N-1]))
λ = vec(res.λ[solver.N])
ν = vec(to_array(res.s))
s = vcat([k < N ? res.C[k][1:pI] : nothing for k = 1:solver.N-1]...)
s = sqrt.(2*max.(0,-s))
V = packV(X,U,λ,μ,ν,s)
cost(solver,res)
max_c(V)
newton_cost(V)
g = ForwardDiff.gradient(newton_cost,V)
H = ForwardDiff.hessian(newton_cost,V)
a = active_set(V,1.0)
rank(H[a,a])
size(H[a,a])
V[a] = V[a] - (H[a,a] + 0I)\g[a]
# V = V - H\g
newton_cost(V)
max_c(V)

X,U,λ,μ,ν,s == unpackV(V)
ν

d(V)


fitness(a::AbstractString, b::AbstractString) = count(l == t for (l, t) in zip(a, b))
function mutate(str::AbstractString, rate::Float64)
    L = collect(Char, " ABCDEFGHIJKLMNOPQRSTUVWXYZ")
    return map(str) do c
        if rand() < rate rand(L) else c end
    end
end

function evolve(parent::String, target::String, mutrate::Float64, nchild::Int)
    println("Initial parent is $parent, its fitness is $(fitness(parent, target))")
    gens = 0
    while parent != target
        children = collect(mutate(parent, mutrate) for i in 1:nchild)
        bestfit, best = findmax(fitness.(children, target))
        parent = children[best]
        gens += 1
        if gens % 10 == 0
            println("After $gens generations, the new parent is $parent and its fitness is $(fitness(parent, target))")
        end
    end
    println("After $gens generations, the parent evolved into the target $target")
end

evolve("IU RFSGJABGOLYWF XSMFXNIABKT", "METHINKS IT IS LIKE A WEASEL", 0.08998, 100)


"""
    AffineIntervalMatrix1{T, IT, MT0<:AbstractMatrix{T}, MT1<:AbstractMatrix{T}} <: AbstractIntervalMatrix{IT}

Interval matrix representing the matrix

```math
A₀ + λA₁,
```
where ``A₀`` and ``A₁`` are real (or complex) matrices, and ``λ`` is an interval.

### Fields

- `A0` -- matrix
- `A1` -- matrix
- `λ`  -- interval

### Examples

The matrix ``I + [1 1; -1 1] * (0 .. 1)`` is:

```jldoctest
julia> using LinearAlgebra

julia> P = AffineIntervalMatrix1(Matrix(1.0I, 2, 2), [1 1; -1 1.], 0 .. 1);

julia> P
2×2 AffineIntervalMatrix1{Float64, Interval{Float64}, Matrix{Float64}, Matrix{Float64}}:
  [1, 2]  [0, 1]
 [-1, 0]  [1, 2]
```
"""
struct AffineIntervalMatrix1{T, IT, MT0<:AbstractMatrix{T}, MT1<:AbstractMatrix{T}} <: AbstractIntervalMatrix{IT}
    A0::MT0
    A1::MT1
    λ::IT

    # inner constructor with dimension check
    function AffineIntervalMatrix1(A0::MT0, A1::MT1, λ::IT) where {T, IT, MT0<:AbstractMatrix{T}, MT1<:AbstractMatrix{T}}
        @assert checksquare(A0) == checksquare(A1) "the size of `A0` and `A1` should match, got $(size(A0)) and $(size(A1)) respectively"
        return new{T, IT, MT0, MT1}(A0, A1, λ)
    end
end

IndexStyle(::Type{<:AffineIntervalMatrix1}) = IndexLinear()
size(M::AffineIntervalMatrix1) = size(M.A0)
getindex(M::AffineIntervalMatrix1, i::Int) = getindex(M.A0, i) + M.λ * getindex(M.A1, i)
function setindex!(M::AffineIntervalMatrix1{T}, X::T, inds...) where {T}
    setindex!(M.A0, X, inds...)
    setindex!(M.A1, zero(T), inds...)
end
copy(M::AffineIntervalMatrix1) = AffineIntervalMatrix1(copy(M.A0), copy(M.A1), M.λ)

"""
    AffineIntervalMatrix{T, IT, MT0<:AbstractMatrix{T}, MT<:AbstractMatrix{T}, MTA<:AbstractVector{MT}} <: AbstractIntervalMatrix{IT}

Interval matrix representing the matrix

```math
A₀ + λ₁A₁ + λ₂A₂ + … + λₖAₖ,
```
where ``A₀`` and ``A₁, …, Aₖ`` are real (or complex) matrices, and ``λ₁, …, λₖ``
are intervals.

### Fields

- `A0` -- matrix
- `A`  -- vector of matrices
- `λ`  -- vector of intervals

### Notes

This type is the general case of the [`AffineIntervalMatrix1`](@ref), which only
contains one matrix proportional to an interval.

### Examples

The affine matrix ``I + [1 1; -1 1] * (0 .. 1) + [0 1; 1 0] * (2 .. 3)`` is:

```jldoctest
julia> using LinearAlgebra

julia> A0 = Matrix(1.0I, 2, 2);

julia> A1 = [1 1; -1 1.]; A2 = [0 1; 1 0];

julia> λ1 = 0 .. 1; λ2 = 2 .. 3;

julia> P = AffineIntervalMatrix(A0, [A1, A2], [λ1, λ1]);

julia> P[1, 1]
[1, 2]

julia> P[1, 2]
[0, 2]
```
"""
struct AffineIntervalMatrix{T, IT, MT0<:AbstractMatrix{T}, MT<:AbstractMatrix{T}, MTA<:AbstractVector{MT}, VIT<:AbstractVector{IT}} <: AbstractIntervalMatrix{IT}
    A0::MT0
    A::MTA
    λ::VIT

    # inner constructor with dimension check
    function AffineIntervalMatrix(A0::MT0, A::MTA, λ::VIT) where {T, IT, MT0<:AbstractMatrix{T}, MT<:AbstractMatrix{T}, MTA<:AbstractVector{MT}, VIT<:AbstractVector{IT}}
        k = length(A)
        @assert k == length(λ) "expected `A` and `λ` to have the same length, got $(length(A)) and $(length(λ)) respectively"
        n = checksquare(A0)
        @inbounds for i in 1:k
            @assert n == checksquare(A[i]) "each matrix should have the same size $n × $n"
        end
        return new{T, IT, MT0, MT, MTA, VIT}(A0, A, λ)
    end
end

IndexStyle(::Type{<:AffineIntervalMatrix}) = IndexLinear()
size(M::AffineIntervalMatrix) = size(M.A0)
getindex(M::AffineIntervalMatrix, i::Int) = getindex(M.A0, i) + sum(M.λ[k] * getindex(M.A[k], i) for k in eachindex(M.λ))
function setindex!(M::AffineIntervalMatrix{T}, X::T, inds...) where {T}
    setindex!(M.A0, X, inds...)
    @inbounds for k in 1:length(M.A)
        setindex!(M.A[k], zero(T), inds...)
    end
end
copy(M::AffineIntervalMatrix) = AffineIntervalMatrix(copy(M.A0), copy(M.A), M.λ)


using RandomTestMatrices
using Base.Test

d = rand(1:100)

let A = randsym(d)
    @Test.test isreal(A)
    @Test.test size(A) == (d,d)
    @Test.test issym(A)
end

let A = randantisym(d)
    i = rand(1:d)
    j = rand(1:d)
    @Test.test isreal(A)
    @Test.test size(A) == (d,d)
    @Test.test A[i,j] == -A[j,i]
end

let A = randposdef(d)
    @Test.test isreal(A)
    @Test.test size(A) == (d,d)
    @Test.test isposdef(A)
    @Test.test issym(A)
end

let A = randstable(d)
    @Test.test size(A) == (d,d)
    @Test.test isreal(A)
    @Test.test all(real(eig(A)[1]) .< 0.0)
end

let A = randunitary(d)
    @Test.test size(A) == (d,d)
    @Test.test !isreal(A)
    @Test.test norm(A*A' - eye(d)) < d*eps()
end

let A = randorth(d)
    @Test.test size(A) == (d,d)
    @Test.test isreal(A)
    @Test.test norm(A*A' - eye(d)) < d*eps()
end

let A = randnormal(d)
    @Test.test size(A) == (d,d)
    @Test.test isreal(A)
    @Test.test norm(A*A' - A'*A) < d*eps(100.)
end


# find the maximal end component in an MDP

# find the maximal accepting end components
# algorithm 47 from "Principle of model checking"
# an end component is defined as a set of states

function maximal_end_components(mdp::M; verbose=false) where {M <: Union{MDP, POMDP}}
    verbose ? println("Building graph from mdp ... \n") : nothing
    mdpg = mdp_to_graph(mdp)
    state_space = states(mdp)
    MEC = Vector{Vector{Int64}}(undef, 0)
    MECnew = Vector{Vector{Int64}}(undef, 0)
    push!(MECnew, 1:length(state_space))# initialize
    DEBUG_STEP = 1 #XXX
    DEBUG_MAX_STEPS = 100 # XXX
    verbose ? println("Computing Maximal End Components ... \n") : nothing
    while MEC != MECnew
        MEC = deepcopy(MECnew)
        MECnew = Vector{Vector{Int64}}(undef, 0)
        for sub_state_space in MEC
            # build subgraph that corresponds to sub_state_space, subgraph of mdpgraph, where all actions moving outside of ss are removed
            sub_g = sub_mdp(mdp, sub_state_space, state_space)
            scc = strongly_connected_components(sub_g)
            for component in scc
                # map bag to state indices
                stateindex_component = sub_state_space[component]
                # check that component is not trivial (is at least a cycle)
                if length(component) == 1 && !has_edge(sub_g, component[1], component[1])
                    continue # skip trivial components
                end
                push!(MECnew, stateindex_component)
            end
        end
        verbose ? println("finished $DEBUG_STEP step, old MEC $(length(MEC)) -> new MEC $(length(MECnew))") : nothing
        DEBUG_STEP += 1
    end
    verbose ? println("MECs length: ", length.(MEC)) : nothing
    verbose ? println("MECs computed. \n") : nothing
    return MEC
end

# convert MDP to graph
function mdp_to_graph(mdp::M) where {M <: Union{MDP, POMDP}}
    g = DiGraph(length(states(mdp)))
    for (i, s) in enumerate(ordered_states(mdp))
        si = stateindex(mdp, s)
        if isterminal(mdp, s)
            add_edge!(g, si, si)
        end
        for a in actions(mdp, s)
            d = transition(mdp, s, a)
            for (sp, p) in weighted_iterator(d)
                if !(p ≈ 0.)
                   spi = stateindex(mdp, sp)
                   add_edge!(g, si, spi)
                end
            end
        end
    end
    return g
end

# return a graph of the subMDP defined by ss
function sub_mdp(mdp::M, state_indices::Vector{Int64}, state_space) where {S, M<:Union{MDP,POMDP}}
    g = DiGraph(length(state_indices))
    for (i, si) in enumerate(state_indices)
        s = state_space[si]
        if isterminal(mdp, s)
            add_edge!(g, i, i)
        end
        for a in actions(mdp, s)
            d = transition(mdp, s, a)
            is_action_valid = true
            for (sp, p) in weighted_iterator(d)
                if !(p ≈ 0.)
                    spi = stateindex(mdp, sp)
                    if !in(spi, state_indices)
                        is_action_valid = false
                    end
                end
            end
            if is_action_valid
                for (sp, p) in weighted_iterator(d)
                    if !(p≈0.)
                        spi = stateindex(mdp, sp)
                        j = findfirst(x->x==spi, state_indices)
                        add_edge!(g, i, j)
                    end
                end
            end
        end
    end
    return g
end


__precompile__()

module IncoherentScatter

const ℜ(z)=Base.real(z)
const ℑ(z)=Base.imag(z)
const ℂ(r, i) = Base.complex(r, i)

using FFTW,SpecialFunctions,Polynomials,PhysicalConstants.CODATA2018,Unitful
using CUDA
using LoopVectorization,Tullio

export Ion,Plasma,CollisionalPlasma,ScatterVolume,pwrspec,autocor,autocor!

friedconte(z) = im*√(π)*erfcx(-im*z)

struct Faddeeva
    n::Int
    l::Float64
    a::Polynomial{Float64}
    function Faddeeva(n=16)
        m  = 2*n
        l  = √(n/√(2))

        theta = (-m+1:m-1)*π/m
        t = l*tan.(theta/2)
        f = [0; exp.(-t.^2).*(l*l .+ t.^2)]
        a = real.(fft(fftshift(f)))[2:n-1]/(2*m)   # no flip because of convention in polyval
        new(n, l, Polynomial(a))
    end
end

const deff = Faddeeva()

"""
    faddeeva(z, f=deff)

   The Faddeeva function is defined as:

   w(z) = exp(-z^2) * erfc(-im*z)

   where erfc(x) is the complex complementary error function. Also,

   friedconte(z) = im * √(π) * faddeeva(z)

   where friedconte(z) is the plasma dispersion function as defined in [2].

   w = faddeeva(z, f) reuses previously calculated coefficients in struct f
   where f is initialized like f=Faddeeva(n) with n specifying the number of terms
   to truncate the expansion (see (13) in [1]). n = 16 is used as default, i.e. for f=Faddeeva().

   Example:
       x = linspace(-10,10,1001); [X,Y] = meshgrid(x,x);
       W = faddeeva(complex(X,Y));

   Reference:
   [1] J.A.C. Weideman, "Computation of the Complex Error Function," SIAM
       J. Numerical Analysis, pp. 1497-1518, No. 5, Vol. 31, Oct., 1994
       Available Online: http://www.jstor.org/stable/2158232

   [2] B.D. Fried, S.D. Conte, The plasma dispersion function. New York Academic Press, 1961.
"""
function faddeeva(z, f::Faddeeva=deff)
    if ℜ(z)==0      # purely imaginary input
        exp.(-z.^2).*erfc(ℑ(z))
    else               # complex valued input
        # make sure input in the upper half-plane (positive imag. values)
        zz = ℑ(z)<0 ? conj(z) : z
        yy = (f.l + im*zz)/(f.l - im*zz)
        w  = 2*polyval(f.a, yy)/(f.l - im*zz)^2 + (1/√(π))/(f.l - im*zz)
        # convert the upper half-plane results to the lower half-plane if necessary
        ℑ(z)<0 ? conj(2*exp(-zz^2) - w) : w
    end
end

const c  = ustrip(u"m/s", SpeedOfLightInVacuum)
const me = ustrip(u"kg",  ElectronMass)
const e  = ustrip(u"C",   ElementaryCharge)
const kB = ustrip(u"J/K", BoltzmannConstant)
const ϵ0 = ustrip(u"F/m", VacuumElectricPermittivity)
const amu= ustrip(u"kg",  AtomicMassConstant)

const fnormf = c*√(me/(2*kB))/2
const debyef = (4*π/(e*c))^2*ϵ0*kB

const mp = ustrip(u"kg",  ProtonMass)
const mn = ustrip(u"kg",  NeutronMass)
const sqrtmpme = √((mp+mn)/me/2)

"gfe(B in Tesla) -> electron gyrofrequency in Hz"
gfe(B) = e*B/me

"""
    ScatterVolume(tfeff, system, b, alpha)

has parameters of the scattering volume that are constant in a fit.
"""
mutable struct ScatterVolume{T<:Real}
    tfeff::T    # effective transmitter frequency in Hz = tf*cos(scattering angle)
    fnormc::T   # factor for frequency normalization f -> theta
    debyec::T   # Debye length correction factor

    system::T   # spectrum scale factor (depends on tx power, volume size, ... )

    b::T        # magnetic field strength in Tesla
    gfe::T      # electron gyrofrequency
    phiel::T    # normalized electron gyrofrequency
    alpha::T    # angle between b and scattering vector in rad
    sinfac::T   # correction factor for magnetic field
    cosal::T    #

    function ScatterVolume{T}(tfeff, system, b, alpha) where {T<:Real}
        fnormc = fnormf/tfeff
        debyec = debyef*tfeff^2
        gf     = gfe(b)
        phiel  = fnormc*gf
        sinfac = 0.5*(sin(alpha)/phiel)^2
        cosal  = cos(alpha)
        new(tfeff, fnormc, debyec, system, b, gf, phiel, alpha, sinfac, cosal)
    end
end
# ScatterVolume(tfeff::T, system, b, alpha) where {T<:Real} = ScatterVolume{T}(tfeff, system, b, alpha)
ScatterVolume(;tfeff::T=230e6, system=1e-6, b=50000e-9, alpha=deg2rad(30)) where {T<:Real} =
    ScatterVolume{T}(tfeff, system, b, alpha)

struct Ion{T<:Real}
    mass::T   # ion mass in amu
    srmr::T   # sqrt of mass ratio ion/electron
end
Ion(m=15.999) = Ion(m, √(m*amu/me))

"""
Plasma(ne, te, ti, ::ScatterVolume, [ions, fractn]

has parameters of the plasma that are constant over all frequencies of the spectrum.
"""
struct Plasma{T<:Real}
    ne::T
    te::T
    ti::T
    fractn1::T
    fractn::NTuple{2, T}    # up to two more ions
    ions::NTuple{3, Ion{T}}
    nion::Int

    spectrsf::T
    debyen::T   # normalized Debye length correction factor
    sqte::T
    sqti::T
    thetc::T
    tex::T
    tefac::T
    dpsidc::T

    # parameter for spectrum at zero frequency:
    cdtine::T
    ym::T

    function Plasma{T}(ne::T, te::T, ti::T, s::ScatterVolume{T}, fractn, ions, nion) where {T<:Real}
        sqti    = √(ti)
        thetc   = s.fnormc/sqti
        tex     = exp(-te*s.sinfac)
        tefac   = 0.5/te + s.sinfac/tex
        fractn1 = one(ne) - fractn[1] - fractn[2]
        dpsidc  = thetc*ions[1].srmr
        cdtine  = ti*s.debyec + ne
        ym      = ti*ne + te*cdtine
        new(ne, te, ti, fractn1, fractn, ions, nion,
            s.system*ne, s.debyec/ne, √(te), sqti, thetc,
            tex, tefac, dpsidc, cdtine, ym)
    end
end
Plasma(ne::T, te::T, ti::T,
       s::ScatterVolume{T}, nion=1, ions=(Ion(16*one(ne)), Ion(31*one(ne)), Ion(one(ne))),
       fractn=(zero(ne),zero(ne))) where {T<:Real} =
    Plasma{T}(ne, te, ti, s, fractn, ions, nion)
Plasma(s::ScatterVolume{T}; ne::T, te::T, ti::T,
       nion=1, ions=(Ion(16*one(ne)), Ion(31*one(ne)), Ion(one(ne))),
       fractn=(zero(ne),zero(ne))) where {T<:Real} =
    Plasma{T}(ne, te, ti, s, fractn, ions, nion)

struct CollisionalPlasma{T<:Real}
    p::Plasma{T}
    ν::T
end
CollisionalPlasma(p; ν) = CollisionalPlasma(p, ν)

"""
    ionadm(theta, z=friedconte(theta))
-> unmagnetized ion admittance for real normalized frequency
"""
ionadm(theta::T, z::Complex{T}=friedconte(theta)) where {T<:Real} =
    ℂ(theta*ℑ(z), 1 + theta*ℜ(z))

"""
    diadt1(::Plasma, ionadm::Complex)
-> derivative with respect to Ti of the unmagnetized ion admittance, 1st ion, zero frequency
"""
diadt1(p::Plasma{T}, ia::Complex{T}) where {T<:Real} =
    ℂ(0, p.fractn1*p.dpsidc*(√(π)*ℜ(ia) - 2*p.ions[1].srmr))

"""
    diadt1(θ, ::Plasma, ionadm::Complex)
-> derivative with respect to Ti of the unmagnetized ion admittance, 1st ion, real normalized frequency
"""
function diadt1(θ::T, p::Plasma{T}, ia::Complex{T}) where {T<:Real}
    r5ti = T(0.5)/p.ti
    p.fractn1*((θ^2/p.ti - r5ti)*ia + im*r5ti)
end

function diadc(theta::T, ia::Complex{T}=ionadm(theta)) where {T<:Real}
    th =1/theta
    ℂ(ℜ(ia)^2 + ℑ(ia)*(1 - ℑ(ia)))*th + ℑ(ia)*(th - 2*theta) - th,
      ℜ(ia)*(theta + (ℑ(ia) - 1)*th)
end

"""
    ionadm(freq, ::Plasma)
-> unmagnetized ion admittance for real unnormalized frequency and plasma parameters
"""
function ionadm(freq::T, p::Plasma{T}, s::ScatterVolume{T}) where {T<:Real}
    fnormt = s.fnormc/p.sqti
    thetai = fnormt*p.ions[1].srmr*freq
    yi     = ionadm(thetai)
    for k in 2:p.nion
        if k==2
            yi *= p.fractn1
        end
        thetai = fnormt*p.ions[k].srmr*freq
        yi    += p.fractn[k-1]*ionadm(thetai)
    end
    yi
end

"""
    ionadm(theta, z=friedconte(theta))
-> unmagnetized ion admittance for complex normalized frequency.
"""
function ionadm(theta::T, z::T=friedconte(theta)) where {T<:Complex}
    psizm = ℑ(theta)*(ℜ(z)^2 + ℑ(z)^2)
    yd    = 1 - ℑ(theta)*(2*ℑ(z) - psizm)
    ℂ(ℜ(theta)*(ℑ(z) - psizm)/yd, 1 + ℜ(theta)*ℜ(z)/yd)
end

"""
    ionadm(freq, ::CollisionalPlasma)
-> unmagnetized ion admittance for unnormalized frequency and collisional plasma parameters
"""
function ionadm(freq::T, coll::CollisionalPlasma{T}, s::ScatterVolume{T}) where {T<:Real}
    fnormt = s.fnormc/coll.p.sqti
    thetai = fnormt*coll.p.ions[1].srmr*ℂ(freq, coll.ν)
    yi     = ionadm(thetai)
    for k in 2:coll.p.nion
        if k==2
            yi *= coll.p.fractn1
        end
        thetai = fnormt*coll.p.ions[k].srmr*ℂ(freq, coll.ν)
        yi    += coll.p.fractn[k-1]*ionadm(thetai)
    end
    yi
end

"""
    elecadm(theta, tex, z)
-> magnetized electron admittance for real normalized frequency
"""
function elecadm(theta::T, tex::T, z::Complex{T}=friedconte(theta)) where {T<:Real}
    tt = tex*theta
    ℂ(tt*imag(z), 1 + tt*real(z)), imag(z)
end

# helper function, this part of the spectrum code is the same for non-collisional and collisional
"""
    espec(freq, ::Plasma, ::ScatterVolume, yi::Complex)
-> power spectral density for non-zero real frequency
"""
function espec(freq::T, p::Plasma{T}, s::ScatterVolume{T}, yi::Complex{T}) where {T<:Real}
    yiti   = yi/p.ti + im*p.debyen

    # electron admittance:
    thetae = s.fnormc/(p.sqte*s.cosal)*freq
    ye,    = elecadm(thetae, p.tex)
    yete   = ye/p.te

    the2   = thetae^2
    texthe = p.tex*thetae

    ym     = abs2(yete+yiti)
    yn     = abs2(yete)*real(yi) + abs2(yiti)*real(ye)

    (p.spectrsf/freq)*(yn/ym)
end

"""
    pwrspec(freq, ::Plasma, ::ScatterVolume)
-> power spectral density for non-zero real frequency
"""
function pwrspec(freq::T, p::Plasma{T}, s::ScatterVolume{T}) where {T<:Real}
    yi = ionadm(freq, p, s)
    espec(freq, p, s, yi)
end

"""
    pwrspec(freq, ::CollisionalPlasma, ::ScatterVolume)
-> power spectral density for non-zero frequency and collisional plasma
"""
function pwrspec(freq::T, coll::CollisionalPlasma{T}, s::ScatterVolume{T}) where {T<:Real}
    yi = ionadm(freq, coll, s)
    espec(freq, coll.p, s, yi)
end

# helper function
function yn0(p::Plasma{T}, s::ScatterVolume{T}, yir::T) where {T<:Real}
    tineq  = p.sqti*p.ne^2*yir*p.ti
    # sqtec  = p.cdtine*p.sqte*√(π)/s.cosal2 #  divide by colal, not colal*colal:
    sqtec  = p.cdtine*p.sqte*√(π)/s.cosal # zi(1) in fortran code = sqrt(π) for real zero frequency
    stecx  = sqtec*p.tex*p.te
    tineq + stecx*p.cdtine
end

"""
    pwrspec(::Plasma)
-> power spectral density for zero real frequency
"""
function pwrspec(p::Plasma, s::ScatterVolume)
    yir = sqrt(π)*p.ions[1].srmr
    for k in 2:p.nion
        if k==2
            yir *= p.fractn1
        end
        yir += p.fractn[k-1]*sqrt(π)*p.ions[k].srmr
    end
    yn = yn0(p, s, yir)
    s.fnormc*p.spectrsf*yn/p.ym^2
end

"""
    pwrspec(::Plasma, ::ScatterVolume, freq_itr)
-> power spectral density for iterable frequencies (with length)
"""
function pwrspec(p::Union{Plasma{T}, CollisionalPlasma{T}}, s::ScatterVolume{T}, freqs;
                 threads=true) where {T<:Real}
    if threads
        pwr = Vector{T}(undef, length(freqs))
        if freqs[1]==0
            rng = 2:length(freqs)
            pwr[1] = pwrspec(p, s)
        else
            rng = 1:length(freqs)
        end
        Threads.@threads for k ∈ rng
            pwr[k] = pwrspec(freqs[k], p, s)
        end
    else
        if freqs[1]==0
            pwr = vcat(pwrspec(p, s), @inbounds pwrspec.(freqs[2:end], [p], [s]))
        else
            pwr = @inbounds pwrspec.(freqs, [p], [s])
        end
    end
    return pwr
end

function yirc(coll::CollisionalPlasma, kion::Int)
    psiion = coll.p.dpsidc*coll.ν
    z      = friedconte(ℂ(zero(coll.p.ne), psiion))
    psiz   = 1 - psiion*ℑ(z)
    psizm  = coll.p.ions[kion].srmr/psiz
    psizm*ℑ(z)
end

"""
    pwrspec(::CollisionalPlasma)
-> power spectral density for zero frequency and collisional plasma
"""
function pwrspec(coll::CollisionalPlasma, s::ScatterVolume)
    yir = yirc(coll, 1)
    for k in 2:coll.p.nion
        if k==2
            yir *= coll.p.fractn1
        end
        z    = friedconte(ℂ(zero(coll.p.ne), p.thetc*coll.p.ions[k].srmr*coll.ν))
        yir += coll.p.fractn[k-1]*yirc(coll, k)
    end

    yn = yn0(coll.p, s, yir)
    s.fnormc*coll.p.spectrsf*yn/coll.p.ym^2
end

"""
    cosine(freqs, τ)
-> matrix for the transform spectrum->ACF, method by Wes Swartz, ISCATSPE.FOR
"""
function cosine(freqs, τ)
    c           = Matrix{eltype(freqs)}(undef, length(freqs), length(τ))
    c[    1, 1] = freqs[2]/2
    for k ∈ 2:length(freqs)-1
        c[k, 1] = (freqs[k+1]-freqs[k-1])/2
    end
    c[end, 1] = (freqs[end] - freqs[end-1])/2
    df        = diff(freqs)
    for l ∈ 2:length(τ)
        τ2    = 2*π*τ[l]
        c[1, l]       = (1 - cos(τ2*freqs[2]))/(τ2*τ2*freqs[2])
        c2            =  cos.(τ2*freqs[2:end-1])
        c[2:end-1, l] = ((c2 - cos.(τ2*freqs[1:end-2]))./df[1:end-1] +
                         (c2 - cos.(τ2*freqs[3:end  ]))./df[2:end  ])/(τ2*τ2)
        cend = cos(τ2*freqs[end])
        fend = 2*freqs[end]-freqs[end-1]
        c[end, l]     =  (2*cend - cos(τ2*freqs[end-1]) - cos(τ2*fend))/(df[end]*τ2*τ2)
    end
    return c
end

struct FT{T<:Real}
    cft::Matrix{T}
end
FT(freqs, τ) = FT(cosine(freqs, τ))

function pwr2acf(pwr::Vector{T}, ftm::Matrix{T}) where {T<:Real}
    @tullio acf[l] := ftm.cft[k, l]*pwr[k]
end

function autocor(cft::Matrix{T}, pwr::Vector{T}) where {T<:Real}
    @tullio acf[l] := cft[k, l]*pwr[k]
end

function autocor!(acf::Vector{T}, cft::Matrix{T}, pwr::Vector{T}) where {T<:Real}
    @tullio acf[l] = cft[k, l]*pwr[k]
end

greet() = print("Hello World!")

end   # module IncoherentScatter


#=
# Compute the node locations for the reference element
=#
include("jacobi_gauss_quad.jl");
include("triangle_nodes.jl");
include("tet_nodes.jl");

function refel_nodes!(refel, nodetype)
    if refel.dim == 0
        # 0D is a point so some of this doesn't make sense
        refel.r1d = [0];
        refel.wr1d = [1];
        refel.g1d = [0];
        refel.wg1d = [1];
        refel.r = [0];
        refel.wr = [1];
        refel.g = [0];
        refel.wg = [1];

    elseif refel.dim == 1
        # 1D has line segments
        if nodetype == UNIFORM
            refel.r1d = Array(-1:(2/(refel.Np-1)):1);
            refel.wr1d = ones(length(refel.r1d)) ./ length(refel.r1d);
        elseif nodetype == GAUSS
            (r,w) = jacobi_gauss_quad(0,0,refel.N);
            refel.r1d = r;
            refel.wr1d = w;
        elseif nodetype == LOBATTO
            if refel.N == 1
                refel.r1d = [-1; 1];
                refel.wr1d = [1; 1];
            else
                # (r,w) = jacobi_gauss_quad(1,1,refel.N-2);
                # refel.r1d = [-1; r ; 1];

                # # compute the weights
                # w = jacobi_polynomial(refel.r1d, 0, 0, refel.N)
                # adgammaN = (2*refel.N + 1) / (refel.N * (refel.N + 1))
                # w = w.*w
                # w = adgammaN./w

                # refel.wr1d = w;
                (refel.r1d, refel.wr1d) = jacobi_LGL_quad(refel.N);
            end
        end
        # Then find Gauss points
        (g,w) = jacobi_gauss_quad(0,0,refel.N);
        refel.g1d = g;
        refel.wg1d = w;

        # 1D is easy
        refel.g = refel.g1d;
        refel.wg = refel.wg1d;
        refel.r = refel.r1d;
        refel.wr = refel.wr1d;

        # surface is just points
        refel.face2local = [[1], [refel.Np]];
        refel.surf_r = [[-1], [1]];
        refel.surf_wr = [[1], [1]];
        refel.surf_g = [[-1], [1]];
        refel.surf_wg = [[1], [1]];

    elseif refel.dim == 2
        # 2D has triangles and quads
        if refel.Nfaces == 3 # triangles
            triangle_refel_nodes!(refel);
        else # quads
            if nodetype == UNIFORM
                refel.r1d = Array(-1:(2/(refel.Np-1)):1);
                refel.wr1d = ones(length(refel.r1d)) ./ length(refel.r1d);
            elseif nodetype == GAUSS
                (r,w) = jacobi_gauss_quad(0,0,refel.N);
                refel.r1d = r;
                refel.wr1d = w;
            elseif nodetype == LOBATTO
                if refel.N == 1
                    refel.r1d = [-1; 1];
                    refel.wr1d = [1; 1];
                else
                    (r,w) = jacobi_gauss_quad(1,1,refel.N-2);
                    refel.r1d = [-1; r ; 1];

                    # compute the weights
                    w = jacobi_polynomial(refel.r1d, 0, 0, refel.N);
                    adgammaN = (2*refel.N + 1) / (refel.N * (refel.N + 1));
                    w = w.*w;
                    w = adgammaN./w;

                    refel.wr1d = w;
                end
            end
            # Then find Gauss quadrature points
            (g,w) = jacobi_gauss_quad(0,0,refel.N);
            refel.g1d = g;
            refel.wg1d = w;

            # r and w
            refel.r = zeros(refel.Np,2);
            refel.wr = zeros(refel.Np);
            refel.g = zeros(refel.Np,2);
            refel.wg = zeros(refel.Np);
            n1d = length(refel.r1d);
            for j=1:n1d
                for i=1:n1d
                    k = (j-1)*n1d + i;
                    refel.r[k,:] = [refel.r1d[i]; refel.r1d[j]];
                    refel.wr[k] = refel.wr1d[i] * refel.wr1d[j];
                    refel.g[k,:] = [refel.g1d[i]; refel.g1d[j]];
                    refel.wg[k] = refel.wg1d[i] * refel.wg1d[j];
                end
            end

            ### surface lines ################
            # First get the face2local maps
            tol = 1e-8;
            qf1(x) = abs(x[1] + 1) < tol;
            qf2(x) = abs(x[2] + 1) < tol;
            qf3(x) = abs(x[1] - 1) < tol;
            qf4(x) = abs(x[2] - 1) < tol;
            refel.face2local = [get_face2local_map(refel.r, qf1),
                                get_face2local_map(refel.r, qf2),
                                get_face2local_map(refel.r, qf3),
                                get_face2local_map(refel.r, qf4)];

            tmp = zeros(2,0);
            refel.surf_r = [tmp, tmp, tmp, tmp];
            refel.surf_g = [tmp, tmp, tmp, tmp];

            for fi=1:4
                refel.surf_r[fi] = refel.r[refel.face2local[fi], :];
                refel.surf_g[fi] = refel.g[refel.face2local[fi], :];
            end
            refel.surf_wr = [refel.wr1d, refel.wr1d, refel.wr1d, refel.wr1d];
            refel.surf_wg = [refel.wg1d, refel.wg1d, refel.wg1d, refel.wg1d];
        end

    elseif refel.dim == 3
        # 3D has tets, hexs and prisms
        if refel.Nfaces == 4 # tets
            tetrahedron_refel_nodes!(refel);
        else # hexs
            if nodetype == UNIFORM
                refel.r1d = Array(-1:(2/(refel.Np-1)):1);
                refel.wr1d = ones(length(refel.r1d)) ./ length(refel.r1d);
            elseif nodetype == GAUSS
                (r,w) = jacobi_gauss_quad(0,0,refel.N);
                refel.r1d = r;
                refel.wr1d = w;
            elseif nodetype == LOBATTO
                if refel.N == 1
                    refel.r1d = [-1; 1];
                    refel.wr1d = [1; 1];
                else
                    (r,w) = jacobi_gauss_quad(1,1,refel.N-2);
                    refel.r1d = [-1; r ; 1];

                    # compute the weights
                    w = jacobi_polynomial(refel.r1d, 0, 0, refel.N);
                    adgammaN = (2*refel.N + 1) / (refel.N * (refel.N + 1));
                    w = w.*w;
                    w = adgammaN./w;

                    refel.wr1d = w;
                end
            end
            # Then find Gauss points
            (g,w) = jacobi_gauss_quad(0,0,refel.N);
            refel.g1d = g;
            refel.wg1d = w;

            # r and w
            refel.r = zeros(refel.Np,3);
            refel.wr = zeros(refel.Np);
            refel.g = zeros(refel.Np,3);
            refel.wg = zeros(refel.Np);
            n1d = length(refel.r1d);
            for k=1:n1d
                for j=1:n1d
                    for i=1:n1d
                        ind = (k-1)*n1d*n1d + (j-1)*n1d + i;
                        refel.r[ind,:] = [refel.r1d[i]; refel.r1d[j]; refel.r1d[k]];
                        refel.wr[ind] = refel.wr1d[i] * refel.wr1d[j] * refel.wr1d[k];
                        refel.g[ind,:] = [refel.g1d[i]; refel.g1d[j]; refel.g1d[k]];
                        refel.wg[ind] = refel.wg1d[i] * refel.wg1d[j] * refel.wg1d[k];
                    end
                end
            end

            ### surface quads ################
            # First get the face2local maps
            tol = 1e-8;
            hf1(x) = abs(x[1] + 1) < tol;
            hf2(x) = abs(x[2] + 1) < tol;
            hf3(x) = abs(x[3] + 1) < tol;
            hf4(x) = abs(x[1] - 1) < tol;
            hf5(x) = abs(x[2] - 1) < tol;
            hf6(x) = abs(x[3] - 1) < tol;
            refel.face2local = [get_face2local_map(refel.r, hf1),
                                get_face2local_map(refel.r, hf2),
                                get_face2local_map(refel.r, hf3),
                                get_face2local_map(refel.r, hf4),
                                get_face2local_map(refel.r, hf5),
                                get_face2local_map(refel.r, hf6)];

            tmp = zeros(3,0);
            refel.surf_r = [tmp, tmp, tmp, tmp, tmp, tmp];
            refel.surf_g = [tmp, tmp, tmp, tmp, tmp, tmp];

            for fi=1:6
                refel.surf_r[fi] = refel.r[refel.face2local[fi], :];
                refel.surf_g[fi] = refel.g[refel.face2local[fi], :];
            end

            # The weights are a little more tricky
            fwr = zeros(refel.Nfp[1]);
            fwg = zeros(refel.Nfp[1]);
            n1d = length(refel.r1d);
            for j=1:n1d
                for i=1:n1d
                    k = (j-1)*n1d + i;
                    fwr[k] = refel.wr1d[i] * refel.wr1d[j];
                    fwg[k] = refel.wg1d[i] * refel.wg1d[j];
                end
            end
            refel.surf_wr = [fwr, fwr, fwr, fwr, fwr, fwr];
            refel.surf_wg = [fwg, fwg, fwg, fwg, fwg, fwg];
        end

    else
        # Not ready for 4D
    end
end

# Determines surface nodes using a supplied function.
# Returns a face to local map. The order of which will be omnipotent.
function get_face2local_map(r, compare)
    n = size(r,1);
    map = zeros(Int,n);
    nf = 0;
    for i=1:n
        if compare(r[i,:])
            nf = nf+1;
            map[nf] = i;
        end
    end

    return map[1:nf];
end

# Returns the global node locations for elemental nodes
# Has size Np,dim
function get_node_coords(vx, elnodes)
    x = [];
    if refel.dim == 1
        hx = vx[2] - vx[1];
        x = vx[1] .+ (elnodes .+ 1) .* (hx*0.5);
    elseif refel.dim == 2
        # for now assume rectangular quads
        x = zeros(size(elnodes));
        hx = max(vx[2,1]-vx[1,1], vx[3,1]-vx[1,1]);
        hy = max(vx[2,2]-vx[1,2], vx[3,2]-vx[1,2]);
        x[:,1] = vx[1,1] .+ (elnodes[:,1] .+ 1) .*(hx*0.5);
        x[:,2] = vx[1,2] .+ (elnodes[:,2] .+ 1) .*(hy*0.5);
    elseif refel.dim == 3
        # for now assume regular hexs
        x = zeros(size(elnodes));
        hx = max(vx[2,1]-vx[1,1], vx[3,1]-vx[1,1]);
        hy = max(vx[2,2]-vx[1,2], vx[3,2]-vx[1,2]);
        x[:,1] = vx[1,1] .+ (elnodes[:,1] .+ 1) .*(hx*0.5);
        x[:,2] = vx[1,2] .+ (elnodes[:,2] .+ 1) .*(hy*0.5);
        x[:,3] = vx[1,3] .+ (elnodes[:,3] .+ 1) .*(hy*0.5);
    end

    return x;
end

# Returns the volume of an element
function get_volume(refel, vx)
    vol = 0;
    if refel.dim == 1
        vol = abs(vx[2]-vx[1]);
    elseif refel.dim == 2
        # for now assume rectangular quads
        hx = max(vx[2,1]-vx[1,1], vx[3,1]-vx[1,1]);
        hy = max(vx[2,2]-vx[1,2], vx[3,2]-vx[1,2]);
        vol = abs(hx*hy);
    elseif refel.dim == 3

    end

    return vol;
end


using Smg2s, Test, SparseArrays, LinearAlgebra

#Define size of matrix to be generated
n = 20

#maximum number of continous `1` in the diagonal of generated nilpotent matrix
nbOne = 3
#offsets of lower diagonals between which the matrix will be initially filled
diag_l = -10
diag_u = -5

@testset "non-Herm SMG2S" begin
    @testset "Generation 1" begin
        #a function to generate the user-provided spectrum, `idx` is the indexing of eigenvalues
        function f(idx::Integer, n = n)
            return 10 * idx + 1 + idx * im
        end
        spec = zeros(ComplexF64, n)
        Spectrum!(spec, f, n)
        A = nonherm(nbOne, n, diag_l, diag_u, spec)
        @test spec ≈ eigen!(Matrix(A)).values atol=0.1
    end
end

@testset "non-Sym SMG2S" begin
    @testset "Generation 1" begin
        #a function to generate the user-provided spectrum, `idx` is the indexing of eigenvalues
        function f(idx::Integer, n = n)
            if idx % 2  == 1
                return  ((idx + 1) / 2)  + 1 - 0.1 * ((idx + 1) / 2) * im
            else
                return  (idx / 2) + 1 + 0.1 * (idx / 2) * im
            end
        end
        spec = zeros(ComplexF64, n)
        Spectrum!(spec, f, n)
        A = nonsym(nbOne, n, diag_l, diag_u, spec)
        @test spec ≈ eigen!(Matrix(A)).values atol=0.5
    end
end


s = open("input.txt") do file
    readlines(file);
end

function EvalExpression(expr)
    # println("=====")
    # println(expr);
    newExp = expr;
    m = match(r"(\d+) \+ (\d+)", expr)
    while m != nothing
        val = parse(Int, m.captures[1]) + parse(Int, m.captures[2]);
        newExp = replace(newExp, string(m.match)=>"$(val)", count=1);

        m = match(r"(\d+) \+ (\d+)", newExp)
    end

    # println(newExp)

    nums = [];
    for m in eachmatch(r"\d+", newExp)
        push!(nums, parse(Int, m.match));
    end
    # println(nums);
    return *(nums...);
end

function DoTheMath(expression)
    println(expression);
    expressionStack = [];
    currentExpr = "";
    for i = 1:length(expression)
        if expression[i] == '('
            push!(expressionStack, currentExpr);
            currentExpr = "";
        elseif expression[i] == ')'
            val = EvalExpression(currentExpr);
            currentExpr = pop!(expressionStack) * "$(val)";
        else
            currentExpr *= string(expression[i])
        end
    end
    # println(currentExpr);

    val = EvalExpression(currentExpr);
    # println(val)
    # println("----------")
    return Int128(val);
end

eachAns = map(DoTheMath, s);

answer = sum(eachAns)


"embedding function for the Gupta potential"
type GuptaEmbed <: SimpleFunction
    xi
end
evaluate(p::GuptaEmbed, r) = p.xi * sqrt(r)
evaluate_d(p::GuptaEmbed, r) = 0.5*p.xi ./ sqrt(r)

"""`GuptaPotential`:
    E_i = A ∑_{j ≠ i} v(r_ij) - ξ ∑_i √ ρ_i
        v(r_ij) = exp[ -p (r_ij/r0 - 1) ]
        ρ_i = ∑_{j ≠ i} exp[ -2q (r_ij / r0 - 1) ]
"""
GuptaPotential(A, xi, p, q, r0, TC::Type, TCargs...)  =
    EAMPotential( TC( SimpleExponential(A, p, r0), TCargs... ),      # V
                  TC( SimpleExponential(1.0, 2*q, r0), TCargs...),   # rho
                  GuptaEmbed( xi ) )                                 # embed

gupta_parameters = Dict(
   #          a       D       α      R₀      <<< Gupta notation
   "Cu" => (3.6147, 0.3446, 1.3921, 2.838),
   "Ag" => (4.0862, 0.3294, 1.3939, 3.096),
   "Au" => (4.0785, 0.4826, 1.6166, 3.004),
   "Ni" => (3.5238, 0.4279, 1.3917, 2.793)
)


# Implicit compressed sensing via the GenSPGL package, which allows an implicit
# DCT-II transform that significantly speeds things up.

export csimage, simulate_csimage

const dσ = 1e-1
const doptions = (verbosity=0,)

function csimage(cs; σ = dσ, options = doptions)
    BLAS.set_num_threads(16)
    s = spgl1(joMatrix(cs.patterns) * joDCT(size(cs.patterns)[2]),
              cs.measurements,
              tau     = 0.0,
              sigma   = σ,
              options = spgOptions(;stepMin=10.0*eps(), iterations=5000,
                                   options...))[1]
    renormalize(dct(s))
end

"Simulate compressed sensing of `img` with `m` measurements."
function simulate_csimage(img, m, σ = dσ, options = doptions)
    spgl_sense(measure(imvec(img), m), σ, options)
end


using MarkovTransitionMatrices

export makeSparseTransition

function makeSparseTransition(nP::Integer=31, nC::Integer=37, extendSD::Real = 1.5, minP::Real = 1e-6)
  n = nP * nC

  T = Float64

  σp = [0x1.269c263357e05p-4, 0x1.9bec08eb57abfp-5]
  σc = [0x1.2c9a3a5b8b5bap-5, 0x1.09d5d411d8a33p-7]

  Πk = [0x1.d7e13e2f38db1p-1 0x1.40f60e8638b53p-4;
        0x1.098244df0fddfp-4 0x1.decfb7641df78p-1]
  all(sum(Πk,2) .≈ 1.0) || throw(error("each row of π must sum to 1"))

  K = size(Πk,2)
  extrema_p = [0.8776572, 2.485073]
  extrema_c = [0.2437302, 1.529529]

  # state-spaces
  pspace = range(extrema_p[1]-log(extendSD), stop=extrema_p[2] + log(extendSD), length=nP)
  cspace = range(extrema_c[1]-log(extendSD), stop=extrema_c[2] + log(extendSD), length=nC)
  zspace = (pspace, cspace, 1:K,)

  # make deviations
  zrandwalk(x::Real, st::Real, σ::Real) = (x - st) / σ

  # allocate giant transition matrix
  P = zeros(T,n*K,n*K)

  # match moments for logp
  Pp1, JN, Λ, L_p1, approxErr = discreteNormalApprox(pspace, pspace, (x::Real,st::Real) -> zrandwalk(x,st,σp[1]), 15)
  Pp2, JN, Λ, L_p2, approxErr = discreteNormalApprox(pspace, pspace, (x::Real,st::Real) -> zrandwalk(x,st,σp[2]), 15)
  # plot(pspace, [L_p1,L_p2], yticks = 1:15, labels=["Hi vol", "Lo vol"], xlabel="Log p", ylabel="Moments matched")

  # match moments for logc
  Pc1, JN, Λ, L_c1, approxErr = discreteNormalApprox(cspace, cspace, (x::Real,st::Real) -> zrandwalk(x,st,σc[1]), 13)
  Pc2, JN, Λ, L_c2, approxErr = discreteNormalApprox(cspace, cspace, (x::Real,st::Real) -> zrandwalk(x,st,σc[2]), 5)
  # plot(cspace, [L_c1,L_c2], yticks = 1:15, labels=["Hi vol", "Lo vol"], xlabel="Log c", ylabel="Moments matched")

  # sparse versions
  sPp1 = MarkovTransitionMatrices.sparsify!(Pp1, minp)
  sPp2 = MarkovTransitionMatrices.sparsify!(Pp2, minp)
  sPc1 = MarkovTransitionMatrices.sparsify!(Pc1, minp)
  sPc2 = MarkovTransitionMatrices.sparsify!(Pc2, minp)

  # do the big matrix
  P12 = (kron(sPc1, sPp1), kron(sPc2, sPp2))
  for r2 in 1:K
    for r1 in 1:K
      P[(r1-1)*n+1:r1*n,  (r2-1)*n+1:r2*n] .= Πk[r1, r2] .* P12[r2]
    end
  end

  sP = sparse(P)

  return (zspace, sP)
end


function test_mapped_dicts()
    # b = Fourier(20)
    # x1 = 1.0+im
    # x2 = 2.0+3im
    # m1 = interval_map(0.0, 1.0, x1, x2)
    # m2 = embedding_map(Float64, Complex{Float64})
    # m = m1 ∘ m2
    # c = mapped_dict(b, m)
    #
    # @test !in_support(c, 10, 0.5)
    # @test in_support(c, 10, x1)
    # @test in_support(c, 10, (x1+x2)/2)
    #
    # u = approximate(c, exp)
    # z = 1/3*x1+2/3*x2
    # @test abs(u(z)-exp(z)) < 0.2

    d = UnitCircle()
    m = parameterization(d)
    b = Fourier(20)
    c = mapped_dict(b, m)
    f1(x,y) = exp(x+im*y)
    u = approximate(c, f1)
    @test in_support(c, 1, SVector(1.0, 0.0))
    @test abs(u(SVector(cos(10),sin(10)))-f1(cos(10),sin(10))) < 1e-6
end


__precompile__()

module PoincareInvariants

    using ApproxFun
    using HDF5
    using ProgressMeter
    using GeometricIntegrators
    using GeometricIntegrators.Utils

    export PoincareInvariant1st, PoincareInvariant1stCanonical,
           PoincareInvariant2nd, PoincareInvariant2ndCanonical,
           PoincareInvariant2ndTrapezoidal,
           evaluate_poincare_invariant, write_to_hdf5

    include("matrix_utils.jl")
    include("poincare_invariant_1st_common.jl")
    include("poincare_invariant_1st.jl")
    include("poincare_invariant_1st_canonical.jl")
    include("poincare_invariant_2nd.jl")
    include("poincare_invariant_2nd_canonical.jl")
    include("poincare_invariant_2nd_trapezoidal.jl")
    include("poincare_invariant_2nd_common.jl")

end


using JLD2
using DSGE
m = KrusellSmithCT()
T, R, ~, inverse_basis = solvect(m)

n_vars = get_setting(m, :n_vars)
n_expectation_errors = get_setting(m, :n_expectation_errors)
n_shocks = get_setting(m, :n_shocks)
x = zeros(Float64, n_vars)

nst = size(T,1) ###
C = zeros(nst)

n = 40 #100
s = zeros(nst,n) #s = zeros(n,406)
s[:,1] = inverse_basis' * x
e = randn(n)

dt = 1/2

for i = 2:n
    s[:, i] = s[:, i-1] + T*s[:,i-1]*dt + R*e[i]*(sqrt(dt))
end
data = zeros(1,Int(n/2))
data[1, :] = [s[nst - 2, 2 * i] for i = 1:Int(n/2)]

save("simulated_data.jld2", "data", data)


export Platte

struct Platte{T<:AbstractFloat, Ti<:Integer}
    xL::T
    yL::T
    nx::Ti
    ny::Ti
    rho::Matrix{T}
    lambda::Matrix{T}
    cp::Matrix{T}
    temp::Matrix{T}
end

function Platte(xL::T, yL::T, nx::Ti, ny::Ti, dens, heat, cap, tv::Union{Number,Function}) where {T,Ti}
    if nx <= 0
        throw(ArgumentError("nx muss > 0 sein, ist aber $nx"))
    end

        function setgrid(f::Number)
            return [f for ix in 1:nx, iy in 1:ny]
        end

        function setgrid(f::Function)
            return [f((iy - 0.5) * yL / ny, (ix - 0.5) * xL / nx) for ix in 1:nx, iy in 1:ny]
        end

    rho = setgrid(dens)
    temp = setgrid(tv)
    lambda= setgrid(heat)
    cp = setgrid(cap)

    return Platte{T, Ti}(xL, yL, nx, ny, rho, lambda, cp, temp)
end

function tempiter(p::Platte, deltat::T) where {T}

    dtemp = heattransfer(p)
    tempn = dtemp .* deltat .+ p.temp

    return Platte(p.xL, p.yL, p.nx, p.ny, p.rho, p.lambda, p.cp, tempn)
end

function heattransfer(p::Platte)

    produkt = p.lambda .* nabla(p.temp, p)
    tempdiff = nabla(produkt, p) ./ (p.cp .* p.rho)

    return tempdiff
end

"""
    nabla(∇m::Matrix{Vector}, p::Platte)

Calculate new matrix containing `∇(∇m)` given `∇m` at all positions of `p`.
"""
function nabla(vm::Matrix{<:Vector}, p::Platte)
    vmx = [vm[ix, iy][1] for ix in 1:p.nx, iy in 1:p.ny]
    vmy = [vm[ix, iy][2] for ix in 1:p.nx, iy in 1:p.ny]
    [nablax(vmx, ix, iy, p) + nablay(vmy, ix, iy, p) for ix in 1:p.nx, iy in 1:p.ny]
end

"""
    nabla(m::Matrix, p::Platte)

Calculate new matrix of vectors containing `∇m` at all postions of `p`.
"""
function nabla(m::Matrix{<:Number}, p::Platte)
    [[nablax(m, ix, iy, p), nablay(m, ix, iy, p)] for ix in 1:p.nx, iy in 1:p.ny]
end

"""
    nablax(m::Matrix, ix, iy, p::Platte)

Calculate finite x-difference at position `[ix,iy]` of matrix `p`.
Matrix dimensions must coincide with Platte.
"""
function nablax(m::Matrix, ix, iy, p::Platte)
    dx = p.xL / p.nx
    size(m) == (p.nx, p.ny) || throw(ArgumentException("wrong matrix size"))
    ix == 1 ? (m[ix + 1, iy] - m[ix, iy]) / dx :
    ix == p.nx ? (m[ix, iy] - m[ix - 1, iy]) / dx :
    (m[ix + 1, iy] - m[ix - 1, iy]) / (2 * dx)
end

"""
    nablay(m::Matrix, ix, iy, p::Platte)

Calculate finite y-difference at position `[ix,iy]` of matrix `m`.
Matrix dimensions must coincide with Platte.
"""
function nablay(m::Matrix, ix, iy, p::Platte)
    dy = p.yL / p.ny
    size(m) == (p.nx, p.ny) || throw(ArgumentException("wrong matrix size"))
    iy == 1 ? (m[ix, iy + 1] - m[ix, iy]) / dy :
    iy == p.ny ? (m[ix, iy] - m[ix, iy - 1]) / dy :
    (m[ix, iy + 1] - m[ix, iy - 1]) / (2 * dy)
end

nothing


using ValidatedNumerics
using IntervalConstraintProgramming

S = @separator 1 <= x^2 + y^2 <= 3
X = IntervalBox(-10..10, -10..10)

@time inner, boundary = pave(S, X, 0.125)

@show length(inner), length(boundary)

include("draw_boxes.jl")

draw_boxes(inner, "green", 0.5, 1)
draw_boxes(boundary, "grey", 0.2)
axis("image")  # set aspect ratio


using ATA
# For load data to input in custom obj_fun:
using CSV
using FileIO
using LinearAlgebra
cd("where your input files are")

# 1. Start ATA and add file with custom settings (Needed)
ata_model = start_ata(;
    settings_file = "settings_ata_custom.jl",
    bank_file = "data/bank.csv",
    bank_delim = ";",
)

# 2. Add categorical constraints (Optional)
add_constraints!(ata_model; constraints_file = "constraints.csv", constraints_delim = ";")

# 3. Add overlap maxima (Optional)
add_overlap!(ata_model; overlap_file = "overlap_matrix.csv", overlap_delim = ";")

# custom objective type, function and arguments
ata_model.obj.name = "custom"

ata_model.obj.fun = function (x::Matrix{Float64}, obj_args::NamedTuple)
    IIF = obj_args.IIF
    T = obj_args.T
    TIF = zeros(Float64, T)
    for t = 1:T
        K, I = size(IIF[t])
        # ungroup items
        xₜ = fs_to_items(x[:, t], I, obj_args.fs_items)
        if K > 1
            TIF[t] = Inf
            for k = 1:K
                TIF[t] = min(TIF[t], LinearAlgebra.dot(IIF[t][k, :], xₜ))
            end
        else
            TIF[t] = LinearAlgebra.dot(IIF[t][1, :], xₜ)[1]
        end
    end
    min_TIF = minimum(TIF)
    TIF = [min_TIF for t = 1:T]
    # Must return a vector of length T.
    # The resulting objective function is the minimum of all values in this vector.
    return TIF::Vector{Float64}
end

ata_model.obj.args = (
    T = ata_model.settings.T,
    IIF = FileIO.load("data/IIF.jld2", "IIF"),
    fs_items = ata_model.settings.fs.items,
)

print_infos(ata_model)
# Assembly settings

# SIMAN (Suggested for Large scale ATA):
# Set the solver, "siman" for simulated annealing, "jump" for MILP solver.
solver = "siman"

# SIMAN (Suggested for Large scale ATA):

start_temp = 0.0001
# Default: `0.1`. Values:  `[0, Inf]`.
# Starting temperature, set to minimum for short journeys (if 0 worse solutions will never be accepted).

geom_temp = 0.1
# Default: `0.1`. Values:  `[0, Inf)`.
# Decreasing geometric factor.

n_item_sample = Inf
# Default: 1. Values: `[1, Inf]`.
# Number of items to alter. Set to minimum for a shallow analysis,
# set to maximum for a deep analysis of the neighbourhoods.

n_test_sample = ata_model.settings.T
# Default: 1. Values: `[1, Inf]`.
# Number of tests to alter. Set to minimum for a shallow analysis, set to maximum for a deep analysis of the neighbourhoods.

opt_feas = 0.9
# Default: 0.0. Values: `[0, Inf)`.
# Optimality/feasibility balancer, if 0 only feasibility of solution is analysed. Viceversa, if 1, only optimality is considered (uncontrained model). All the other values are accepted but produce uninterpretable results.

n_fill = 1
# Default: 1. Values: `[0, Inf)`.
# Number of fill-up phases, usually 1 is sufficient, if start_temp is high it can be higher.
# If a starting_design is supplied, it should be set to 0.

verbosity = 1
# Default: 2. Values: `1` (minimal), `2` (detailed).
# Verbosity level. In the console '+' stands for improvement, '_' for accepting worse solution.
# The dots are the fill-up improvement steps.

#! Termination criteria:

max_time = 100.0
# Default: `1000.0`. Values: `[0, Inf)`.
# Time limit in seconds.

max_conv = 5
# Default: `2`. Values: `[1, Inf)`.
# Maximum convergence, stop when, after max_conv rounds no improvements have been found.
# Set to minimum for shallow analysis, increase it for deep analysis of neighbourhoods.

feas_nh = 1
# Default: `0`. Values: `[1, Inf)`.
# Maximum number of Feasibility neighbourhoods to explore, set to the minimum if the model is small or not highly constrained.

opt_nh = Inf
# Default: `5`. Values: `[1, Inf)`.
# Maximum number of Optimality neighbourhoods to explore, set to the minimum if the model is highly constrained.

# 4. assemble
assemble!(
    ata_model;
    solver = solver,
    max_time = max_time,
    start_temp = start_temp,
    geom_temp = geom_temp,
    n_item_sample = n_item_sample,
    n_test_sample = n_test_sample,
    verbosity = verbosity,
    max_conv = max_conv,
    opt_feas = opt_feas,
    n_fill = n_fill,
    feas_nh = feas_nh,
    opt_nh = opt_nh,# ,
    # optimizer_attributes = optimizer_constructor,
    # optimizer_constructor =optimizer_attributes
)

# All the settings and outputs from optimization are in ata_model object.
# See the struct in ATA.jl to understand how to retrieve all the information.
# A summary of the resulting tests is available in results_folder/results.txt
# If siman is chosen, the optimality and feasibility of the best neighbourhood
# is reported in "results/results_ata.jl"

ata_model.obj.name = "maximin"
print_results(ata_model; results_folder = "results")

#]add https://github.com/giadasp/ATAPlot.jl
using ATAPlot
plot_results(ata_model; plots_folder = "plots")


include("../src/qspice.jl")

using QSpice.Gates, QSpice.State, QSpice.BitOps

function pow2between(n1, n2)
    for i = n1:n2
        if (i & (i - 1)) == 0
            return i
        end
    end
    error("Cannot find appropriate power of 2 in the given interval")
end

# Two random integers have a ~60% chance of being coprimes, so just resamping
# until we find one is a viable algorithm
function randcoprime(n)
    cprm = rand(3:n - 1)
    while gcd(cprm, n) != 1
        cprm = rand(3:n - 1)
    end
    return cprm
end

function approxrational(n::Real, limit::Int)
    m = [1 0; 0 1]

    current = n
    cint = floor(Int, current)

    while m[2,1] * cint + m[2, 2] <= limit
        m = m * [cint 1; 1 0]
        if current == float(cint)
            break
        end
        current = 1.0 / (current - cint)
        cint = floor(Int, current)
    end

    e1 = n - m[1,1] / m[2,1]

    cint = floor(Int, (limit - m[2,2]) / m[2,1])
    n2 = m[1,1] * cint + m[1,2]
    d2 = m[2,1] * cint + m[2,2]
    e2 = n - n2 / d2

    if abs(e1) <= abs(e2)
        return m[1,1], m[2,1]
    else
        return n2, d2
    end
end

function factorize(n)
    if isprime(n)
        return [n]
    end

    if iseven(n)
        return [2, factorize(n / 2)...]
    end

    q = pow2between(n * n, 2 * n * n - 1)
    x = randcoprime(n)

    factor = 1
    while factor == 1
        m = findperiod(n, q, x)
        r = approxrational(m / q, q)[2]
        factor = max(gcd(x^div(r, 2) + 1, n) % n,
                     gcd(x^div(r, 2) - 1, n) % n)
    end

    return [factor, div(n, factor)]
end

function findperiod(n, q, x)
    # First initialize register 1 with an equally weighed superposition of integers
    # from 0 to q - 1
    reg1 = fromstates([QUBIT0 for _ in 1:floor(Int, log2(q))]...)
    for i = 1:reg1.bits
        reg1 = hadamard(reg1, i)
    end

    # This part is not yet done with actual quantum gates/functions
    bits = ceil(Int, log2(n))
    reg2 = QuantumState([0.0 + 0im for _ in 1:2^bits], bits)
    for a = 0:q - 1
        reg2[powermod(x, a, n) + 1] += reg1[a + 1]
    end
    n = sqrt(sumabs2(reg2.vector))
    reg2.vector = reg2.vector .* (1.0 / n)

    reg = Gates.join(reg1, reg2)
    post, k = partialmeasure(reg, reg1.bits + 1:reg.bits...)

    collapsed, m = measure(qft(post))
    return todecimal(m)
end


""" Dynamic """

function dynamic(ekf::ExtendedKalmanFilter, x::AbstractVector, u::AbstractVector)
    # Non linear function F has to have x and u as inputs
    return ekf.f(x, u) + cholesky(ekf.Q).L*randn(size(ekf.Q, 1))
end

function prediction(ekf::ExtendedKalmanFilter, s::State, u::AbstractVector)
    xp = ekf.f(s.x, u) # predicted state prior
    F = ForwardDiff.jacobian(μ -> ekf.f(μ, u), s.x)
    Pp = F*s.P*F' + ekf.Q # a priori state covariance
    return State(xp, Pp)
end

function param_dynamic(θ, parma_ekf::ExtendedParamKalmanFilter, x::AbstractVector, u::AbstractVector)
    # Non linear function F has to have x and u as inputs
    return param_ekf.f(θ, x, u) + cholesky(param_ekf.Q(θ)).L*randn(size(parma_ekf.Q(θ), 1))
end

function param_prediction(θ, param_ekf::ExtendedParamKalmanFilter, s::State, u::AbstractVector)
    xp = param_ekf.f(θ, s.x, u) # predicted state prior
    F = ForwardDiff.jacobian(μ -> param_ekf.f(θ, μ, u), s.x)
    Pp = F*s.P*F' + param_ekf.Q(θ) # a priori state covariance
    return State(xp, Pp)
end

""" Observation """

function observation(ekf::ExtendedKalmanFilter, x::AbstractVector)
    return ekf.h(x) + cholesky(ekf.R).L*randn(size(ekf.R, 1))
end

function correction(ekf::ExtendedKalmanFilter, s::State, y::AbstractVector)
    v = y - ekf.h(s.x) # measurement pre fit residual
    H = ForwardDiff.jacobian(μ -> ekf.h(μ), s.x)
    S = H*s.P*H' + ekf.R # pre fit residual covariance
    K = s.P*H'*inv(S) # Kalman gain
    x_hat = s.x + K*v # a posteriori state estiamate
    P_hat = (I - K*H)*s.P # a posteriori covariance estiamate
    return State(x_hat, P_hat)
end

function param_observation(θ, param_ekf::ExtendedParamKalmanFilter, x::AbstractVector)
    return param_ekf.h(θ, x) + cholesky(param_ekf.R(θ)).L*randn(size(param_ekf.R(θ), 1))
end

function param_correction(θ, param_ekf::ExtendedParamKalmanFilter, s::State, y::AbstractVector)
    v = y - param_ekf.h(θ, s.x) # measurement pre fit residual
    H = ForwardDiff.jacobian(μ -> param_ekf.h(θ, μ), s.x)
    S = H*s.P*H' + param_ekf.R(θ) # pre fit residual covariance
    K = s.P*H'*inv(S) # Kalman gain
    x_hat = s.x + K*v # a posteriori state estiamate
    P_hat = (I - K*H)*s.P # a posteriori covariance estiamate
    return State(x_hat, P_hat)
end

"""Loss """

function state_mse(filter::ExtendedKalmanFilter, s::State, u::AbstractVector, y::AbstractVector)
    ϵx = norm(s.x - filter.f(s.x))
    ϵy = norm(y - filter.h(s.x))
    return ϵx + ϵy
end


mutable struct xNES_state{T}
    ημ::T
    ησ::T
    samples::Int
    termination::Termination{T}
    function xNES_state{T}(d::Integer;P...) where T
        ημ=T(1)
        ησ=T( (9+3*log(d))/(5*d*√d) )
        samples=4 + ceil(Int, log(3*d))
        Term=Termination{T}()
        S=new{T}(ημ,ησ,samples,Term)
        for k in keys(P)
            hasfield(sNES_state{T},k) && setfield!(S,k,P[k])
            hasfield(Termination{T},k) && setfield!(Term,k,P[k])
        end
        S
    end
end

struct xNES <: OptMethod end

function δ(T,i,j)
    ifelse(i==j,one(T),zero(T))
end

function exponential_nes(f,μ0::AbstractVector{T},A::AbstractMatrix{T},params::xNES_state{T}) where T
    n=params.samples
    ημ=params.ημ
    ηB=ησ=params.ησ
    μ=copy(μ0)
    d=length(μ)
    Z=Array{T}(undef,d,n)
    σ=abs(det(A))^(1/d)
    F=fill(f(μ),n)
    B=A./σ
    idx=collect(1:n)
    u=utility_function{T}(n)
    Gδ=zeros(T,d)
    GM=zeros(T,d,d)
    tmp_μ=copy(μ)
    iterations=0
    @inbounds while !termination_criteria(params.termination,σ,iterations)
        randn!(Z)
        for i in 1:n
            #F[i] = f(μ .+ σ .* B * Z[:,i])
            mul!(tmp_μ,B,view(Z,:,i))
            tmp_μ .= μ .+ tmp_μ .* σ
            F[i] = f(tmp_μ)
        end
        sortperm!(idx,F)
        #Gδ=sum(u[i] .* Z[:,idx[i]] for i in 1:n)
        #GM=sum(u[i] .* (Z[:,idx[i]] * (Z[:,idx[i]]') - I) for i in 1:n)
        let i=1
            j = idx[i]
            for k2 in 1:d
                Gδ[k2] = u[i] * Z[k2,j]
                for k1 in 1:d
                    GM[k1,k2] = u[i] * (Z[k1,j] * Z[k2,j] - δ(T,k1,k2))
                end
            end
        end
        for i in 2:n
            j = idx[i]
            for k2 in 1:d
                Gδ[k2] += u[i] * Z[k2,j]
                for k1 in 1:d
                    GM[k1,k2] += u[i] * (Z[k1,j] * Z[k2,j] - δ(T,k1,k2))
                end
            end
        end
        #=
        Gσ = tr(GM) / d
        GB=GM - Gσ * I
        μ=μ .+ ημ * σ *B *Gδ
        σ=σ * exp(ησ/2 * Gσ)
        B=B * exp(ηB/2 .* GB)
        =#
        Gσ = tr(GM) / d
        for i in 1:d
            GM[i,i]-=Gσ
        end
        GM .*= ηB/2
        Gσ *= ησ/2
        mul!(tmp_μ,B,Gδ)
        μ .+= ημ .* σ .* tmp_μ
        σ *= exp(Gσ)
        mul!(GM,B,exp(GM))
        B.=GM
        iterations+=1
    end
    return (sol=μ, cost=f(μ))
end

function exponential_nes(f,μ::AbstractVector{T},σ::AbstractVector{T},params::xNES_state{T}) where T
    A=diagm(σ)
    exponential_nes(f,μ,A,params)
end

function exponential_nes(f,μ::AbstractVector{T},σ::T,params::xNES_state{T}) where T
    exponential_nes(f,μ,fill(σ,length(μ)),params)
end

function exponential_nes(f,μ::AbstractVector{T},σ; P...) where T
    exponential_nes(f,μ,σ,xNES_state{T}(length(μ);P...))
end

pickmethod(::Type{xNES}) = exponential_nes


function stderr_SparKLIE(Ψx, Ψy, θ, ω)
    nx = size(Ψx, 2)
    ny = size(Ψy, 2)
    supp = findall(!iszero, ω)
    S = (cov(transpose(Ψx[supp, :]), corrected=false) ./ nx) .+ (cov(rhat(θ, Ψy) .* transpose(Ψy[supp, :]), corrected=false) ./ ny)
    σ2 = 0.
    for k in 1:length(supp)
        for l in 1:length(supp)
            σ2 += S[k,l] * ω[supp][k] * ω[supp][l]
        end
    end
    sqrt(σ2)
end


function [ varp ] = lateralSponge(varp,varb,x_loc,x_dist_in,dx,nx,trigpi,latdmpcoef,u_switch)
%UNTITLED2 Summary of this function goes here

%  by default, make the sponging layer 10% of the total domain width
      xwidth=0.1*(nx-1)*dx;

      coef=0.0;

%  if we're within the west sponge zone, then apply damping
        if x_loc <= xwidth
          coef =  0.5*(1.-cos(trigpi*(xwidth-x_loc))/xwidth);
        end

%  if we're within the east sponge zone, then apply damping
        if x_loc >= ((nx-1)*dx-xwidth)
          coef =  0.5*(1.-cos(trigpi*         ...
                       (x_loc-(nx-1)*dx-xwidth))/xwidth);
        end

%  apply sponge to slowly remove perturbations
        if coef > 0.0
          varp = varp - coef*(varp-varb);
        end

end