using UniqueVectors: UniqueVector

export Blackbox
"""
    Blackbox{T, U}(ω::Vector{Float64}, Y::Vector{U}, P::U, Q::U,
        ports::UniqueVector{T}) where {T, U}
    Blackbox(ω::Vector{<:Real}, Y::Vector{U}, P::U, Q::U,
        ports::AbstractVector{T}) where {T, U}

An admittance model with constant `P` and `Q` but with varying `Y`, indexed by a real
variable `ω`.
"""
struct Blackbox{T, U<:AbstractMatrix{<:Number}} <: AdmittanceModel{T}
    ω::Vector{Float64}
    Y::Vector{U}
    P::U
    Q::U
    ports::UniqueVector{T}
    function Blackbox{T, U}(ω::Vector{Float64}, Y::Vector{U}, P::U, Q::U,
        ports::UniqueVector{T}) where {T, U}
        if length(Y) > 0
            l = size(Y[1], 1)
            @assert size(P) == (l, length(ports))
            @assert size(Q) == (l, length(ports))
            @assert all([size(y) == (l, l) for y in Y])
        end
        @assert length(ω) == length(Y)
        return new{T, U}(ω, Y, P, Q, ports)
    end
end

Blackbox(ω::Vector{<:Real}, Y::Vector{U}, P::U, Q::U,
    ports::AbstractVector{T}) where {T, U} = Blackbox{T, U}(ω, Y, P, Q, UniqueVector(ports))

get_Y(bbox::Blackbox) = bbox.Y

get_P(bbox::Blackbox) = bbox.P

get_Q(bbox::Blackbox) = bbox.Q

get_ports(bbox::Blackbox) = bbox.ports

function partial_copy(bbox::Blackbox{T, U};
    Y::Union{Vector{V}, Nothing}=nothing,
    P::Union{V, Nothing}=nothing,
    Q::Union{V, Nothing}=nothing,
    ports::Union{AbstractVector{W}, Nothing}=nothing) where {T, U, V, W}
    Y = Y == nothing ? get_Y(bbox) : Y
    P = P == nothing ? get_P(bbox) : P
    Q = Q == nothing ? get_Q(bbox) : Q
    ports = ports == nothing ? get_ports(bbox) : ports
    return Blackbox(bbox.ω, Y, P, Q, ports)
end

function compatible(bboxes::AbstractVector{Blackbox{T, U}}) where {T, U}
    if length(bboxes) == 0
        return true
    end
    ω = bboxes[1].ω
    return all([ω == bbox.ω for bbox in bboxes[2:end]])
end

# TODO for sparse matrices use sparse linear solver instead of \ and /

function impedance_to_scattering(Z::AbstractMatrix{<:Number}, Z0::AbstractVector{<:Real})
    U = diagm(0 => Z0)
    return (Z + U) \ (Z - U)
end

function scattering_to_impedance(S::AbstractMatrix{<:Number}, Z0::AbstractVector{<:Real})
    U = diagm(0 => Z0)
    return U * (I + S) / (I - S)
end

function admittance_to_scattering(Y::AbstractMatrix{<:Number}, Z0::AbstractVector{<:Real})
    U = diagm(0 => Z0)
    return (I + Y * U) \ (I - Y * U)
end

function scattering_to_admittance(S::AbstractMatrix{<:Number}, Z0::AbstractVector{<:Real})
    V = diagm(0 => 1 ./ Z0)
    return (I - S) * ((I + S) \ V)
end

export impedance_matrices
"""
    impedance_matrices(bbox::Blackbox)

Find the impedance matrices `Z` with `y = Zx` for each `ω`.
"""
impedance_matrices(bbox::Blackbox) = [transpose(bbox.Q) * (Y \ collect(bbox.P))
    for Y in bbox.Y]

export admittance_matrices
"""
    admittance_matrices(bbox::Blackbox)

Find the admittance matrices `Y` with `Yy = x` for each `ω`.
"""
function admittance_matrices(bbox::Blackbox)
    if bbox.P == I # avoid unnecessary matrix inversion
        return map(collect, bbox.Y)
    else
        return [inv(Z) for Z in impedance_matrices(bbox)]
    end
end

export scattering_matrices
"""
    scattering_matrices(bbox::Blackbox, Z0::AbstractVector{<:Real})

Find the scattering matrices that map incoming waves to outgoing waves on transmission
lines with characteristic impedance `Z0` for each `ω`.
"""
function scattering_matrices(bbox::Blackbox, Z0::AbstractVector{<:Real})
    if bbox.P == I # avoid unnecessary matrix inversion
        return [admittance_to_scattering(Y, Z0) for Y in admittance_matrices(bbox)]
    else
        return [impedance_to_scattering(Z, Z0) for Z in impedance_matrices(bbox)]
    end
end

function canonical_gauge(bbox::Blackbox)
    n = length(bbox.ports)
    Y = admittance_matrices(bbox)
    P = Matrix{eltype(eltype(Y))}(I, n, n)
    Q = Matrix{eltype(eltype(Y))}(I, n, n)
    return Blackbox(bbox.ω, Y, P, Q, bbox.ports)
end

"""
    Blackbox(ω::Vector{<:Real}, pso::PSOModel)

Create a Blackbox for a PSOModel where `ω` represents angular frequency.
"""
function Blackbox(ω::Vector{<:Real}, pso::PSOModel)
    K, G, C = get_Y(pso)
    Y = [K ./ s + G + C .* s for s in 1im * ω]
    P = get_P(pso) * (1.0 + 0im)
    Q = get_Q(pso) * (1.0 + 0im)
    return Blackbox(ω, Y, P, Q, get_ports(pso))
end


function solve_str8ts(
    grid,
    white;
    backtrack = true,
    all_solutions = false,
    keep_logs = false,
    logging = [:Info, :Table],
)
    straight_tables = Vector{Array{Int,2}}(undef, 9)
    for i in 1:9
        straight_tables[i] = get_straights(collect(1:9), i)
    end

    m = Model(optimizer_with_attributes(
        CS.Optimizer,
        "backtrack" => backtrack,
        "keep_logs" => keep_logs,
        "all_solutions" => all_solutions,
        "logging" => logging,
    ))
    @variable(m, 0 <= x[1:9, 1:9] <= 9, Int)

    # set variables
    for r in 1:9, c in 1:9
        if grid[r, c] != 0
            @constraint(m, x[r, c] == grid[r, c])
        elseif white[r, c] == 1
            @constraint(m, x[r, c] >= 1)
        else # black ones without a number
            @constraint(m, x[r, c] == 0)
        end
    end

    straights = []
    for r in 1:9
        found_straight = false
        vec = Vector{Tuple{Int,Int}}()
        for c in 1:9
            if white[r, c] == 1
                found_straight = true
                push!(vec, (r, c))
            elseif found_straight
                push!(straights, copy(vec))
                empty!(vec)
                found_straight = false
            end
        end
        if found_straight
            push!(straights, copy(vec))
        end
    end

    for c in 1:9
        found_straight = false
        vec = Vector{Tuple{Int,Int}}()
        for r in 1:9
            if white[r, c] == 1
                found_straight = true
                push!(vec, (r, c))
            elseif found_straight
                push!(straights, copy(vec))
                empty!(vec)
                found_straight = false
            end
        end
        if found_straight
            push!(straights, copy(vec))
        end
    end

    for r in 1:9
        variables = [x[r, c] for c = 1:9 if white[r, c] == 1 || grid[r, c] != 0]
        @constraint(m, variables in CS.AllDifferent())
    end
    for c in 1:9
        variables = [x[r, c] for r = 1:9 if white[r, c] == 1 || grid[r, c] != 0]
        @constraint(m, variables in CS.AllDifferent())
    end

    for straight in straights
        len = length(straight)
        variables = [x[s[1], s[2]] for s in straight]
        @constraint(m, variables in CS.TableSet(straight_tables[len]))
    end

    optimize!(m)
    status = JuMP.termination_status(m)
    return status, m, x
end

function get_straights(numbers, len)
    sort!(numbers)
    nrows = (length(numbers) - len + 1) * factorial(len)
    table = Array{Int64}(undef, (nrows, len))
    i = 1
    for j in 1:(length(numbers) - len + 1)
        l = numbers[j:(j + len - 1)]
        for row in permutations(l, len)
            table[i, :] = row
            i += 1
        end
    end
    return table
end

@testset "Str8ts no backtrack" begin
    grid = zeros(Int, (9, 9))
    grid[1, :] = [0 0 0 0 0 0 0 3 0]
    grid[2, :] = [0 0 0 0 8 0 0 0 0]
    grid[3, :] = [0 0 0 0 0 0 1 0 0]
    grid[4, :] = [0 0 0 0 0 6 0 0 0]
    grid[5, :] = [0 0 1 0 0 0 0 0 8]
    grid[6, :] = [0 0 0 4 0 0 9 0 0]
    grid[7, :] = [0 0 3 5 0 0 0 0 0]
    grid[8, :] = [0 0 0 0 0 0 0 0 2]
    grid[9, :] = [5 0 9 0 0 0 0 0 0]

    white = zeros(Int, (9, 9))
    white[1, :] = [1 1 1 1 0 0 1 1 0]
    white[2, :] = [0 1 1 1 0 1 1 1 0]
    white[3, :] = [0 1 1 0 1 1 0 1 1]
    white[4, :] = [1 1 0 1 1 0 1 1 1]
    white[5, :] = [1 1 1 1 1 1 1 1 1]
    white[6, :] = [1 1 1 0 1 1 0 1 1]
    white[7, :] = [1 1 0 1 1 0 1 1 0]
    white[8, :] = [0 1 1 1 0 1 1 1 0]
    white[9, :] = [0 1 1 0 0 1 1 1 1]

    status, m, x = solve_str8ts(grid, white; all_solutions = true, logging = [])
    com = CS.get_inner_model(m)
    @test JuMP.termination_status(m) == MOI.OPTIMAL
    @test MOI.get(m, MOI.ResultCount()) == 1
    @test convert.(Int, JuMP.value.(x[1, :])) == [7, 9, 6, 8, 0, 0, 2, 3, 0]
    @test convert.(Int, JuMP.value.(x[2, :])) == [0, 6, 5, 7, 8, 4, 3, 2, 0]
    @test convert.(Int, JuMP.value.(x[3, :])) == [0, 3, 4, 0, 6, 5, 1, 8, 9]
    @test convert.(Int, JuMP.value.(x[4, :])) == [4, 5, 0, 2, 3, 6, 8, 9, 7]
    @test convert.(Int, JuMP.value.(x[5, :])) == [2, 4, 1, 3, 5, 9, 7, 6, 8]
    @test convert.(Int, JuMP.value.(x[6, :])) == [3, 1, 2, 4, 7, 8, 9, 5, 6]
    @test convert.(Int, JuMP.value.(x[7, :])) == [1, 2, 3, 5, 4, 0, 6, 7, 0]
    @test convert.(Int, JuMP.value.(x[8, :])) == [0, 7, 8, 6, 0, 3, 5, 4, 2]
    @test convert.(Int, JuMP.value.(x[9, :])) == [5, 8, 9, 0, 0, 2, 4, 1, 3]
end


module KineticEnergy

using TypesBasis
using UtilityFunctions

export buildkineticenergy

@doc raw"""
    Use: buildkineticenergy(basisfunc::Basis)

Description: Generate the Ne*N-basis x Ne*N-basis kinetic energy matrix for our contracted basis set. The following expression captures this:

```
$ T_{n,m}$
```

Input:
Results:

"""
function buildkineticenergy(natoms::Int,basisfunc::Array{GaussOrbitals},basis::Basis)

    numbasisfunc = basis.nbasisfunc; #number of basis functions
    basissize = natoms*basis.nbasisfunc; #overlap size
    ke = zeros(Float64,basissize,basissize); #overlap matrix Ne*Gaussians x Ne*Gaussians

    flatbasisfunc = flattenbasisfunc(natoms,numbasisfunc,basisfunc);

    for n=1:basissize
        for m=1:basissize
            basisn = flatbasisfunc[n];
            basism = flatbasisfunc[m];

            #Gaussian factors
            an = basisn.alphas;
            am = basism.alphas;

            #Contraction coeffs.
            cn = basisn.coefs;
            cm = basism.coefs;

            xn,yn,zn = basisn.ax, basisn.ay, basisn.az;
            xm,ym,zm = basism.ax, basism.ay, basism.az;

            #Number of primitive Gaussians
            primsn = length(an);
            primsm = length(am);

            for np=1:primsn
                for mp=1:primsm
                    #Get overlap integral
                    overlap = getgaussoverlap(np,mp,basisn,basism);

                    #Contraction coeff prod.
                    cnm = cn[np]*cm[mp];

                    #Gaussian product theorem
                    p = an[np]+am[mp];
                    px = (an[np]*xn + am[mp]*xm) / p;
                    py = (an[np]*yn + am[mp]*ym) / p;
                    pz = (an[np]*zn + am[mp]*zm) / p;

                    #Build K.E. based on notes
                    ke[n,m] += 3 * cnm * am[mp] * overlap;
                    ke[n,m] -= 2 * cnm * am[mp]^2 * ((px-xm)^2 + 1/(2*p)) * overlap;
                    ke[n,m] -= 2 * cnm * am[mp]^2 * ((py-ym)^2 + 1/(2*p)) * overlap;
                    ke[n,m] -= 2 * cnm * am[mp]^2 * ((pz-zm)^2 + 1/(2*p)) * overlap;
                end
            end # nprims
        end
    end #basissize
    return ke
end #buildkineticenergy

end #KineticEnergy


module CircleFinding

using CircleFit
using DelimitedFiles
using Plots
using StatsBase
using Random
using Statistics
using Distributions
using Optim
using ForwardDiff

function read_profiles(csv_file)
    a = readdlm(csv_file,',')
    b = map(v -> typeof(v) == Int64 ? v / 1000.0 : NaN,a)[:,2:end]
    x = b[:,1:2:end]
    y = b[:,2:2:end]
    return (x,y)
end

function read_apex_profile(csv_file)
    a = readdlm(csv_file,',')
    x = (a[:,1]./1000.0)'
    y = (a[:,2]./1000.0)'
    return (x,y)
end

function read_single_profile(csv_file)
    a = readdlm(csv_file,',')
    x = a[:,2:2:end]./1000.0
    y = a[:,3:2:end]./1000.0
end

"""

"""
function circle_ransac(data, radius; threshold = 0.4, max_iters = 2000, number_points = 8)
    data_length = size(data)[2]
    best_xc = 0
    best_yc = 0
    best_rfit = 0;
    best_err = 10000000;
    for i = 1:max_iters
        ri = randperm(data_length)
        maybe_in = data[:,ri[1:3]]
        xc,yc,rfit = circfit(maybe_in)
        also_in = hcat(filter(xy -> abs(sqrt((xy[1]-xc)^2 + (xy[2]-yc)^2) - radius) < threshold, collect(eachcol(data[:,ri[4:end]])))...)

        if size(also_in)[1] == 2 && size(also_in)[2] > number_points && abs(rfit - radius) < threshold
            current_data = hcat(maybe_in,also_in)
            better_xc, better_yc, better_rfit = circfit(current_data)

            err = circle_error(current_data,better_xc,better_yc,radius)

            if err < best_err
                best_xc = better_xc
                best_yc = better_yc
                best_rfit = better_rfit
                best_err = err
            end
        end
    end
    return (best_xc, best_yc, best_rfit, best_err)
end

circfit(xy) = circfit(xy[1,:],xy[2,:])

circle_error(data, xc, yc, radius) = sqrt(mean((radius .- sqrt.(sum((data .- [xc;yc]).^2,dims=1))).^2))

function filterdata(x,y)
    data = vcat(x[:]',y[:]')
    return hcat(filter(xy -> isfinite(xy[1]) && isfinite(xy[2]),collect(eachcol(data)))...)
end

function fit_circle(x1,y1,x2,y2,radius)
    x3 = (x1+x2) ./ 2
    y3 = (y1+y2) ./ 2
    q = sqrt((x1-x2)^2 + (y1-y2)^2)

    xc = x3 - sqrt(radius^2 - (q/2)^2) * (y1 - y2)/q
    yc = y3 - sqrt(radius^2 - (q/2)^2) * (x2 - x2)/q
    return(xc,yc)
end

function joe_circle(data,radius; xbins = (-15,15,500), ybins = (-30, 30, 1000))
    bins = zeros(ybins[3],xbins[3])

    d1 = data[:,1:end-1]
    d2 = data[:,2:end]

end

function circle_hough_map(data,radius; xlims = (-15,15), ylims = (-30, 30), step_size=0.05)

    bx = range(xlims...,step=step_size)
    len_x = length(bx)
    by = range(ylims...,step=step_size)
    len_y = length(by)

    bins = zeros(len_y,len_x)

    #d = Normal(radius, step_size)
    d = SymTriangularDist(radius,step_size)
    for xy = eachcol(data)
        xstart = findindex(xy[1] - radius - step_size, xlims[1], step_size, len_x)
        xend = findindex(xy[1] + radius + step_size, xlims[1], step_size, len_x)
        ystart = findindex(xy[2] - radius - step_size, ylims[1], step_size, len_y)
        yend = findindex(xy[2] + step_size, ylims[1], step_size, len_y)
        bins[ystart:yend,xstart:xend] .+= [pdf(d, sqrt((xy[1] - x)^2 + (xy[2] - y)^2)) for y = by[ystart:yend], x = bx[xstart:xend]]
    end
    return bins
end

function circle_hough(data,radius; xlims = (-15,15), ylims = (-30, 30), step_size = 0.05)

    bx = range(xlims...,step=step_size)
    by = range(ylims...,step=step_size)

    bins = circle_hough_map(data, radius, xlims=xlims, ylims=ylims, step_size=step_size)
    max_tuple = findmax(bins)
    return (bx[max_tuple[2][2]], by[max_tuple[2][1]],max_tuple[1])
end

function circle_hybrid(data,radius)
    x,y,w = circle_hough(data,radius)

    good_data = hcat(filter(xy -> radius-0.1 < sqrt( (xy[1] - x)^2 + (xy[2] - y)^2) < radius+0.1,collect(eachcol(data)))...)'
    cir_fit_err(p, x, y) = sum((sqrt.((x .- p[1]).^2 .+ (y .- p[2]).^2) .- p[3]).^2)

    res = optimize(b -> cir_fit_err(b,good_data[:,1],good_data[:,2]),[x,y,radius], LBFGS(), autodiff=:forward)
    return vcat(Optim.minimizer(res),length(good_data),Optim.minimum(res))
end

function findindex(x, initial, step_size, count)
    index = round(Integer, (x - initial) / step_size) + 1
    return clamp(index, 1, count)
end

end


@testset "kernelkroneckermat" begin
    rng = MersenneTwister(123456)
    k = SqExponentialKernel()
    x = range(0, 1; length=10)
    X = vcat(collect.(Iterators.product(x,x))'...)

    @test all(collect(kernelkronmat(k,collect(x),2)).≈kernelmatrix(k,X,obsdim=1))
    @test all(collect(kernelkronmat(k,[x,x])).≈kernelmatrix(k,X,obsdim=1))
    @test_throws AssertionError kernelkronmat(LinearKernel(),collect(x),2)
end


# Checks if an expression contains variable var
is_const(ex::SymbolicUtils.Pow, var::SymbolicUtils.Sym) = is_const(ex.base, var) && is_const(ex.exp, var)
is_const(ex::SymbolicUtils.Mul, var::SymbolicUtils.Sym) = all(is_const.(keys(ex.dict), Ref(var)))
is_const(ex::SymbolicUtils.Add, var::SymbolicUtils.Sym) = all(is_const.(keys(ex.dict), Ref(var)))
is_const(ex::SymbolicUtils.Sym, var::SymbolicUtils.Sym) = ex !== var
is_const(::Number, var::SymbolicUtils.Sym) = true

# Checks if an expression is an elementry function in the list of symbols/terms vars
is_elf(ex::SymbolicUtils.Pows) = is_elf(ex.base) & is_elf(ex.exp)
is_elf(ex::SymbolicUtils.Mul) = all(is_elf, keys(ex.dict))
is_elf(ex::SymbolicUtils.Add) = all(is_elf, keys(ex.dict))
is_elf(::SymbolicUtils.Sym) = true
is_elf(::Number) = true
is_elf(ex::SymbolicUtils.Term) = ex.f in [sin, cos, tan, csc, sec, cot, asin, acos, atan, acsc, asec, acot, log, log10]
is_elf(::Any) = false


# Mahalanobis distances

struct Mahalanobis{M<:AbstractMatrix} <: Metric
    qmat::M
end

struct SqMahalanobis{M<:AbstractMatrix} <: SemiMetric
    qmat::M
end

function result_type(d::Mahalanobis, ::Type{T1}, ::Type{T2}) where {T1,T2}
    z = zero(T1) - zero(T2)
    return typeof(sqrt(z * zero(eltype(d.qmat)) * z))
end

function result_type(d::SqMahalanobis, ::Type{T1}, ::Type{T2}) where {T1,T2}
    z = zero(T1) - zero(T2)
    return typeof(z * zero(eltype(d.qmat)) * z)
end

# SqMahalanobis

function (dist::SqMahalanobis)(a::AbstractVector, b::AbstractVector)
    if length(a) != length(b)
        throw(DimensionMismatch("first array has length $(length(a)) which does not match the length of the second, $(length(b))."))
    end

    Q = dist.qmat
    z = a - b
    return dot(z, Q * z)
end

sqmahalanobis(a::AbstractVector, b::AbstractVector, Q::AbstractMatrix) = SqMahalanobis(Q)(a, b)

function colwise!(r::AbstractArray, dist::SqMahalanobis, a::AbstractMatrix, b::AbstractMatrix)
    Q = dist.qmat
    get_colwise_dims(size(Q, 1), r, a, b)
    z = a - b
    dot_percol!(r, Q * z, z)
end

function colwise!(r::AbstractArray, dist::SqMahalanobis, a::AbstractVector, b::AbstractMatrix)
    Q = dist.qmat
    get_colwise_dims(size(Q, 1), r, a, b)
    z = a .- b
    dot_percol!(r, Q * z, z)
end

function _pairwise!(r::AbstractMatrix, dist::SqMahalanobis,
                    a::AbstractMatrix, b::AbstractMatrix)
    Q = dist.qmat
    m, na, nb = get_pairwise_dims(size(Q, 1), r, a, b)

    Qa = Q * a
    Qb = Q * b
    sa2 = dot_percol(a, Qa)
    sb2 = dot_percol(b, Qb)
    mul!(r, a', Qb)

    for j = 1:nb
        @simd for i = 1:na
            @inbounds r[i, j] = max(sa2[i] + sb2[j] - 2 * r[i, j], 0)
        end
    end
    r
end

function _pairwise!(r::AbstractMatrix, dist::SqMahalanobis,
                    a::AbstractMatrix)
    Q = dist.qmat
    m, n = get_pairwise_dims(size(Q, 1), r, a)

    Qa = Q * a
    sa2 = dot_percol(a, Qa)
    mul!(r, a', Qa)

    for j = 1:n
        for i = 1:(j - 1)
            @inbounds r[i, j] = r[j, i]
        end
        r[j, j] = 0
        for i = (j + 1):n
            @inbounds r[i, j] = max(sa2[i] + sa2[j] - 2 * r[i, j], 0)
        end
    end
    r
end

# Mahalanobis

function (dist::Mahalanobis)(a::AbstractVector, b::AbstractVector)
    sqrt(SqMahalanobis(dist.qmat)(a, b))
end

mahalanobis(a::AbstractVector, b::AbstractVector, Q::AbstractMatrix) = Mahalanobis(Q)(a, b)

function colwise!(r::AbstractArray, dist::Mahalanobis, a::AbstractMatrix, b::AbstractMatrix)
    sqrt!(colwise!(r, SqMahalanobis(dist.qmat), a, b))
end

function colwise!(r::AbstractArray, dist::Mahalanobis, a::AbstractVector, b::AbstractMatrix)
    sqrt!(colwise!(r, SqMahalanobis(dist.qmat), a, b))
end

function _pairwise!(r::AbstractMatrix, dist::Mahalanobis,
                    a::AbstractMatrix, b::AbstractMatrix)
    sqrt!(_pairwise!(r, SqMahalanobis(dist.qmat), a, b))
end

function _pairwise!(r::AbstractMatrix, dist::Mahalanobis,
                    a::AbstractMatrix)
    sqrt!(_pairwise!(r, SqMahalanobis(dist.qmat), a))
end


# Gives generic rules for generating the strictly new vertices when splitting
#  an E-dimensional simplex with a size reducing factor of k. Subsimplices containing
#  one of the original vertices is not considered.
#
# splits the decomposition part of a SimplexSplitting output. This is a
# (k^E * (E+1) x k) array of integers. (so, ignore the orientations for now)
#
# @return
#  NewVert An array of the vertices of the splitted simplex.
function simplicial_subdivision_single(k, E)
    ##
    # OUTCOME:
    # NewVert is an array of dimension R x k. Each row contains the indices of
    # the original vertices defining the new vertices of the simplicial subdivision.
    # R = (k+E)!/(k! E!)-E-1
    # Subtriangulation is an array of dimension (E+1)k^E x 1. It consists of a
    # stack of k^E columns of dimension E+1, each one of them containing
    # the indices of the vertices generating the corresponding simplex of the subdivision in the order
    # given by [rows(OriginalVertices);rows(NewVert)]
    ##

    # Get simplex splitting rules, ignore orientations for now.
    splits = simplex_split(k, E, orientations = false)

    # Retrieve integer labels from the tensor decomposition. A way of uniquely identifying
    # the rows of splits. Each integer corresponds to a row in
    # splits (but are not necessarily in order).
    integer_labels = (splits .- 1) * ((E+1) .^ collect(0:(k-1)))

    uniquelabels = unique(integer_labels)
    n_rows = size(uniquelabels, 1)

    J = indexin(integer_labels, uniquelabels)

    # IndicesOfUnqiueIntegersinOriginal
    I = Vector{Int}(length(uniquelabels))
    for i = 1:length(uniquelabels)
        I[i] = findfirst(integer_labels, uniquelabels[i])
    end

    Aux = splits[I, :]
    V = repeat(uniquelabels, 1, E + 1)
    aux = repeat((((E+1)^k - 1) ./ E) * transpose(collect(0:E)), size(uniquelabels, 1), 1)

    aux = heaviside0(-abs.(V - aux)) .* repeat(collect(1:n_rows), 1, E + 1)
    aux = aux[find(aux)]

    aux = round.(Int, aux)
    Caux = round.(Int, complementary(aux, n_rows))
    Ipp = [aux; Caux]

    Ip = zeros(Int, size(I))

    # for any i\in Ipp, Ip(i) is the number j\in {1,...,size(Ipp,1)} such that Ipp(j)=i
    Ip[Ipp] = transpose(collect(1:size(I, 1)))

    # Indices of the non-repeated vertices that are not part of the original simplex.
    # These are the truly new vertices generated by the splitting.
    NewVert = Aux[Caux, :]

    allverts = vcat(repeat(collect(1:(E+1)), 1, k), NewVert)

    # Subarrays of dimension E+1 x 1 vertically concatenated. Each subarray contains the
    # indices in the rows of NewVert corresponding to the vertices furnishing the corresponding
    # simplex (analogous to the output of SimplexSplitting, check notes).
    SubTriangulation = round.(Int, Ip[J])
    SubTriangulation = copy(transpose(reshape(SubTriangulation, E + 1, k^E)))

    return allverts, SubTriangulation
end


export ExponentialDist, nc

"""
    ExponentialDist{FT} <: AbstractSizeDist{FT}

Exponential size distribution defined by two moments: total water content, L̅ and
mean droplet mass, x̅
"""
struct ExponentialDist{FT} <: AbstractSizeDist{FT}
    L̅::FT # Total water content, kg/m3
    x̅::FT # Mean initial droplet mass, kg

    # TODO: Add special constructor to enforce positive args
    function ExponentialDist(;L̅::FT, x̅::FT) where {FT <: Real}
        new{FT}(L̅, x̅)
    end
end

function Base.show(io::IO, d::ExponentialDist{FT}) where {FT <: Real}
    print(io, "ExponentalDist(L̅=$(d.L̅ * 1e3) g/m³, x̅=$(d.x̅) kg)")
end

"""
    nc(dist::ExponentialDist{FT}, x::FT) where {FT <: Real}

Compute the number concentration for droplets of the given mass for a
specific size distribution.
"""
function nc(dist::ExponentialDist{FT}, x::FT) where {FT <: Real}
    return (dist.L̅ / dist.x̅^2) * exp(-x / dist.x̅)
end


using Distributions
using DelimitedFiles
using DataFrames

# Generate n simulations under different models and record their summary statistics

ncell = parse(UInt32, ARGS[1])
model = parse(UInt32, ARGS[2])
nsims = parse(UInt32, ARGS[3])
only_mut = 0
if size(ARGS, 1) > 3
    only_mut = parse(UInt32, ARGS[4])
end
# ncell = 20
# model = 1
# nsims = 2

# change directory accordingly
dir = "../../"
progSim = dir * "/code/simpdo"
simdir = dir * "/data/power_analysis/"

skip = 3

min_mu = 0
max_mu = 0.5
# mean_mu = 0.1
# sigma_mu = 0.01

min_brate = 0.3
max_brate = 0.5
# mean_brate = 0.4
# sigma_brate = 0.02

min_fitness = -1
max_fitness = 0
# mean_fitness = -0.5
# if size(ARGS, 1) > 3
#     mean_fitness = parse(UInt32, ARGS[4])
# end
# sigma_fitness = 0.5

function gen_simulations_neutral(fout, ncell, nsims, dim_stat)
    res_sim = zeros(Float64, nsims, dim_stat)
    for i in 1:nsims
        mu1 = rand(Uniform(0,max_mu), 1)[1]
        mu2 = rand(Uniform(0,max_mu), 1)[1]
        birth_rate = rand(Uniform(min_brate, max_brate), 1)[1]
        # mu1 = rand(truncated(Normal(mean_mu, sigma_mu), min_mu, max_mu), 1)[1]
        # mu2 = rand(truncated(Normal(mean_mu, sigma_mu), min_mu, max_mu), 1)[1]
        # birth_rate = rand(truncated(Normal(mean_brate, sigma_brate), min_brate, max_brate), 1)[1]
        # birth_rate = brate_sim
        death_rate = 0

        seed=rand(UInt32, 1)[1]

        cmdSim=`$progSim -o $dir -e $ncell --chr_prob $mu1 --arm_prob $mu2 --birth_rate $birth_rate --death_rate $death_rate --skip $skip --seed $seed --only_mut $only_mut`

        res = chomp(read(cmdSim, String))
        arr = split(res,"\t")
        simdata = map(x->parse(Float64,x), arr)

        params = [mu1, mu2, birth_rate, seed]
        res_sim[i,:] = vcat(params, simdata)
    end
    writedlm(fout, res_sim)
end

function gen_simulations_withsel(fout, ncell, nsims, dim_stat)
    res_sim = zeros(Float64, nsims, dim_stat)
    for i in 1:nsims
        mu1 = rand(Uniform(0,max_mu), 1)[1]
        mu2 = rand(Uniform(0,max_mu), 1)[1]
        birth_rate = rand(Uniform(min_brate, max_brate), 1)[1]
        # mu1 = rand(truncated(Normal(mean_mu, sigma_mu), min_mu, max_mu), 1)[1]
        # mu2 = rand(truncated(Normal(mean_mu, sigma_mu), min_mu, max_mu), 1)[1]
        # birth_rate = rand(truncated(Normal(mean_brate, sigma_brate), min_brate, max_brate), 1)[1]
        # birth_rate = brate_sim
        fitness = rand(Uniform(-1, 0), 1)[1]
        # fitness = rand(truncated(Normal(mean_fitness, sigma_fitness), min_fitness, max_fitness), 1)[1]
        death_rate = 0

        seed=rand(UInt32, 1)[1]

        fitval = "1 " * string(fitness)
        model = 1

        cmdSim=`$progSim -o $dir -e $ncell --model $model --fitness "$fitval" --chr_prob $mu1 --arm_prob $mu2 --birth_rate $birth_rate --death_rate $death_rate --skip $skip --seed $seed --only_mut $only_mut`

        res = chomp(read(cmdSim, String))
        arr = split(res,"\t")
        simdata = map(x->parse(Float64,x), arr)

        params = [mu1, mu2, birth_rate, fitness, seed]
        res_sim[i,:] = vcat(params, simdata)
    end
    writedlm(fout, res_sim)
end

# adding average of branch ratios
if only_mut == 1
    suffix = string(ncell) * "_n" * string(nsims) * "_maxmu" * string(max_mu) * "_minb" * string(min_brate) * "_maxb" * string(max_brate) * "_withbratio_mut.tsv"
else
    suffix = string(ncell) * "_n" * string(nsims) * "_maxmu" * string(max_mu) * "_minb" * string(min_brate) * "_maxb" * string(max_brate) * "_witheblen.tsv"
end

# Change dim_stat according to the number of summary statistics
if model == 0
    println("simulate data under neutral evolution")
    fout = simdir * "sim_data_neutral_ncell" * suffix
    # 8 summary statistics
    dim_stat=12
    gen_simulations_neutral(fout, ncell, nsims, dim_stat)
else
    println("simulate data with selection")
    fout = simdir * "sim_data_selection_ncell" * suffix
    # 8 summary statistics, one more column (fitness) than neutral model
    dim_stat=13
    gen_simulations_withsel(fout, ncell, nsims, dim_stat)
end


function get_VAR_lag_length(D::Array, pbar::Integer, ic::ASCIIString, i::Bool=true)
        IC   = zeros(pbar,1)
        for p = 1:pbar
            i==true ? V = VAR(D,p,true) : V = VAR(D,p,false)
            n,m = size(V.X)
            t = n-p
            sig = V.Σ/t
        if ic == "aic"
            IC[p] = log(det(sig))+2*p*m^2/t
        elseif ic == "bic"
            IC[p] = log(det(sig))+(m^2*p)*log(t)/t
        elseif ic == "hqc"
            IC[p] = log(det(sig))+2*log(log(t))*m^2*p/t
        elseif ic == "aicc"
            b = t/(t-(p*m+m+1))
            IC[p] = t*(log(det(sig))+m)+2*b*(m^2*p+m*(m+1)/2)
        elseif error("'ic' must be aic, bic, hqc or aicc")
        end
     end
     length_ic = indmin(IC)
     println("Using $ic the best lag-length is $length_ic")
     return length_ic
end


using Base.Test
using StochasticProcesses

let p = BrownianBridge(), t=linspace(0, 1, 10)
  cumsim(p, t, 3)
end


function adjoint_Jacobi_field(
    M::NONMUTATINGMANIFOLDS,
    p::Real,
    q::Real,
    t::Real,
    X::Real,
    β=βdifferential_geodesic_startpoint,
)
    x = shortest_geodesic(M, p, q, t)
    B = get_basis(M, p, DiagonalizingOrthonormalBasis(log(M, p, q)))
    V = get_vectors(M, p, B)
    Θ = vector_transport_to.(Ref(M), Ref(p), V, Ref(x), Ref(ParallelTransport()))
    # Decompose wrt. frame, multiply with the weights from w and recompose with Θ.
    Y = sum(
        (inner.(Ref(M), Ref(x), Ref(X), Θ)) .*
        (β.(B.data.eigenvalues, Ref(t), Ref(distance(M, p, q)))) .* V,
    )[1]
    return Y
end
function jacobi_field(
    M::NONMUTATINGMANIFOLDS,
    p::Real,
    q::Real,
    t::Real,
    X::Real,
    β=βdifferential_geodesic_startpoint,
)
    x = shortest_geodesic(M, p, q, t)
    B = get_basis(M, p, DiagonalizingOrthonormalBasis(log(M, p, q)))
    V = get_vectors(M, p, B)
    Θ = vector_transport_to.(Ref(M), Ref(p), V, Ref(x), Ref(ParallelTransport()))
    Y = zero_tangent_vector(M, p)
    # Decompose wrt. frame, multiply with the weights from w and recompose with Θ.
    Y = sum(
        (inner.(Ref(M), Ref(p), Ref(X), V)) .*
        (β.(B.data.eigenvalues, Ref(t), Ref(distance(M, p, q)))) .* Θ,
    )[1]
    return Y
end
function grad_TV2(M::NONMUTATINGMANIFOLDS, q, p::Int=1)
    c = mid_point(M, q[1], q[3], q[2]) # nearest mid point of x and z to y
    d = distance(M, q[2], c)
    innerLog = -log(M, c, q[2])
    X = [zero_tangent_vector(M, q[i]) for i in 1:3]
    if p == 2
        X[1] = adjoint_differential_geodesic_startpoint(M, q[1], q[3], 1 / 2, innerLog)
        X[2] = -log(M, q[2], c)
        X[3] = adjoint_differential_geodesic_endpoint(M, q[1], q[3], 1 / 2, innerLog)
    else
        if d > 0 # gradient case (subdifferential contains zero, see above)
            X[1] = adjoint_differential_geodesic_startpoint(
                M, q[1], q[3], 1 / 2, innerLog / (d^(2 - p))
            )
            X[2] = -log(M, q[2], c) / (d^(2 - p))
            X[3] = adjoint_differential_geodesic_endpoint(
                M, q[1], q[3], 1 / 2, innerLog / (d^(2 - p))
            )
        end
    end
    return X
end
function prox_TV2(::NONMUTATINGMANIFOLDS, λ, pointTuple::Tuple{T,T,T}, p::Int=1) where {T}
    w = @SVector [1.0, -2.0, 1.0]
    x = SVector(pointTuple)
    if p == 1 # Theorem 3.5 in Bergmann, Laus, Steidl, Weinmann, 2014.
        sr_dot_xw = sym_rem(sum(x .* w))
        m = min(λ, abs(sr_dot_xw) / (dot(w, w)))
        s = sign(sr_dot_xw)
        return sym_rem.(x .- m .* s .* w)
    elseif p == 2 # Theorem 3.6 ibd.
        t = λ * sym_rem(sum(x .* w)) / (1 + λ * dot(w, w))
        return sym_rem.(x - t .* w)
    else
        throw(
            ErrorException(
                "Proximal Map of TV2(Circle,λ,pT,p) not implemented for p=$(p) (requires p=1 or 2)",
            ),
        )
    end
end
function prox_TV2(M::PowerManifold{𝔽,N}, λ, x, p::Int=1) where {𝔽,N<:NONMUTATINGMANIFOLDS}
    power_size = power_dimensions(M)
    R = CartesianIndices(power_size)
    d = length(size(x))
    minInd = first(R).I
    maxInd = last(R).I
    y = deepcopy(x)
    for k in 1:d # for all directions
        ek = CartesianIndex(ntuple(i -> (i == k) ? 1 : 0, d)) #k th unit vector
        for l in 0:2
            for i in R # iterate over all pixel
                if (i[k] % 3) == l
                    JForward = i.I .+ ek.I #i + e_k
                    JBackward = i.I .- ek.I # i - e_k
                    if all(JForward .<= maxInd) && all(JBackward .>= minInd)
                        (y[M, JBackward...], y[M, i.I...], y[M, JForward...]) = prox_TV2(
                            M.manifold,
                            λ,
                            (y[M, JBackward...], y[M, i.I...], y[M, JForward...]),
                            p,
                        )
                    end
                end # if mod 3
            end # i in R
        end # for mod 3
    end # directions
    return y
end


module JuliaSets

export escape_time, JuliaSet, call

type JuliaSet
   Q::Function
end

function call{T<:Real}(J::JuliaSet, xs::Array{T}, ys::Array{T}, n_iter::Int64=100, escape_tol::T=100.0)
  @assert size(xs) == size(ys)
  ans = xs + 1im * ys
  for i = 1:n_iter
    ans = J.Q(ans)
  end
  return abs(ans) .< T(escape_tol)
end

function escape_time{T <: Real}(Q::Function, xs::Array{T,1}, ys::Array{T,1}, n_iter=100, escape_tol=100)
  A = zeros(length(xs), length(ys))
  for i in 1:length(xs)
    for j in 1:length(ys)
      z = xs[i] + 1im * ys[j]
      for k = 1:n_iter
        z = Q(z)
        if abs(z) > escape_tol # diverged
          A[i,j] = k
          break
        end
      end
      if A[i,j] == 0 # didn't diverge
        A[i,j] = escape_tol + 1
      end
    end
  end
  return A
end

end # module


using DiscO
using Random
data = db()

meta(data, "name", "(1+1) EA")
meta(data, "description", "This is for test.")

# Size of the bitstring.
n = 100

# Instance contains data relating to the problem, e.g. Graph
# Here, we don't have any specific information regarding the problem just the size
instance = Dict(["size"=> n])

# We want to optimize OneMax funciton, OneMax := sum(candidate)
fitness = onemax

# Using bitrand we create a string with n bits chosen
# unifromly random from 0 and 1
start = bitrand(n)
candidate = (point=start, fitness=fitness(instance, start))

# The algorithm must be stop when the fitness becomes n
stop(instance, best, iter) = (best.fitness == instance["size"] || iter == 100000)

# We use static bit mutation with probability 1/n.
mutation = staticBitMutation(1/n)

# Run algorithm (1+1) EA with standard mutation probability
optimum = ea1p1(instance, candidate, mutation, fitness, stop, data)


const BigFloats = Union{BigFloat,Complex{BigFloat}}

if VERSION < v"0.7-"
    import Base.FFTW: fft, fft!, rfft, irfft, ifft, conv, dct, idct, dct!, idct!,
                        plan_fft!, plan_ifft!, plan_dct!, plan_idct!,
                        plan_fft, plan_ifft, plan_rfft, plan_irfft, plan_dct, plan_idct
else
    import FFTW: fft, fft!, rfft, irfft, ifft, dct, idct, dct!, idct!,
                        plan_fft!, plan_ifft!, plan_dct!, plan_idct!,
                        plan_fft, plan_ifft, plan_rfft, plan_irfft, plan_dct, plan_idct
    import DSP: conv
end

# The following implements Bluestein's algorithm, following http://www.dsprelated.com/dspbooks/mdft/Bluestein_s_FFT_Algorithm.html
# To add more types, add them in the union of the function's signature.
function fft(x::Vector{T}) where T<:BigFloats
    n = length(x)
    ispow2(n) && return fft_pow2(x)
    ks = range(zero(real(T)),stop=n-one(real(T)),length=n)
    Wks = exp.((-im).*convert(T,π).*ks.^2 ./ n)
    xq, wq = x.*Wks, conj([exp(-im*convert(T,π)*n);reverse(Wks);Wks[2:end]])
    return Wks.*conv(xq,wq)[n+1:2n]
end

function fft!(x::Vector{T}) where T<:BigFloats
    x[:] = fft(x)
    return x
end

# add rfft for BigFloat, by calling fft
#  this creates ToeplitzMatrices.rfft, so avoids changing rfft

rfft(v::Vector{T}) where T<:BigFloats = fft(v)[1:div(length(v),2)+1]
function irfft(v::Vector{T},n::Integer) where T<:BigFloats
    @assert n==2length(v)-1
    r = Vector{Complex{BigFloat}}(undef, n)
    r[1:length(v)]=v
    r[length(v)+1:end]=reverse(conj(v[2:end]))
    real(ifft(r))
end

ifft(x::Vector{T}) where {T<:BigFloats} = conj!(fft(conj(x)))/length(x)
function ifft!(x::Vector{T}) where T<:BigFloats
    x[:] = ifft(x)
    return x
end

function conv(u::StridedVector{T}, v::StridedVector{T}) where T<:BigFloats
    nu,nv = length(u),length(v)
    n = nu + nv - 1
    np2 = nextpow(2,n)
    append!(u,zeros(T,np2-nu)),append!(v,zeros(T,np2-nv))
    y = ifft_pow2(fft_pow2(u).*fft_pow2(v))
    #TODO This would not handle Dual/ComplexDual numbers correctly
    y = T<:Real ? real(y[1:n]) : y[1:n]
end

# This is a Cooley-Tukey FFT algorithm inspired by many widely available algorithms including:
# c_radix2.c in the GNU Scientific Library and four1 in the Numerical Recipes in C.
# However, the trigonometric recurrence is improved for greater efficiency.
# The algorithm starts with bit-reversal, then divides and conquers in-place.
function fft_pow2!(x::Vector{T}) where T<:BigFloat
    n,big2=length(x),2one(T)
    nn,j=n÷2,1
    for i=1:2:n-1
        if j>i
            x[j], x[i] = x[i], x[j]
            x[j+1], x[i+1] = x[i+1], x[j+1]
        end
        m = nn
        while m ≥ 2 && j > m
            j -= m
            m = m÷2
        end
        j += m
    end
    logn = 2
    while logn < n
        θ=-big2/logn
        wtemp = sinpi(θ/2)
        wpr, wpi = -2wtemp^2, sinpi(θ)
        wr, wi = one(T), zero(T)
        for m=1:2:logn-1
            for i=m:2logn:n
                j=i+logn
                mixr, mixi = wr*x[j]-wi*x[j+1], wr*x[j+1]+wi*x[j]
                x[j], x[j+1] = x[i]-mixr, x[i+1]-mixi
                x[i], x[i+1] = x[i]+mixr, x[i+1]+mixi
            end
            wr = (wtemp=wr)*wpr-wi*wpi+wr
            wi = wi*wpr+wtemp*wpi+wi
        end
        logn = logn << 1
    end
    return x
end

function fft_pow2(x::Vector{Complex{T}}) where T<:BigFloat
    y = interlace(real(x),imag(x))
    fft_pow2!(y)
    return complex.(y[1:2:end],y[2:2:end])
end
fft_pow2(x::Vector{T}) where {T<:BigFloat} = fft_pow2(complex(x))

function ifft_pow2(x::Vector{Complex{T}}) where T<:BigFloat
    y = interlace(real(x),-imag(x))
    fft_pow2!(y)
    return complex.(y[1:2:end],-y[2:2:end])/length(x)
end

function dct(a::AbstractVector{Complex{BigFloat}})
	N = big(length(a))
    c = fft([a; flipdim(a,1)])
    d = c[1:N]
    d .*= exp.((-im*big(pi)).*(0:N-1)./(2*N))
    d[1] = d[1] / sqrt(big(2))
    lmul!(inv(sqrt(2N)), d)
end

dct(a::AbstractArray{BigFloat}) = real(dct(complex(a)))

function idct(a::AbstractVector{Complex{BigFloat}})
	N = big(length(a))
    b = a * sqrt(2*N)
    b[1] = b[1] * sqrt(big(2))
    shift = exp.(-im * 2 * big(pi) * (N - big(1)/2) * (0:(2N-1)) / (2 * N))
    b = [b; 0; -flipdim(b[2:end],1)] .* shift
    c = ifft(b)
    flipdim(c[1:N],1)
end

idct(a::AbstractArray{BigFloat}) = real(idct(complex(a)))

dct!(a::AbstractArray{T}) where {T<:BigFloats} = (b = dct(a); a[:] = b)
idct!(a::AbstractArray{T}) where {T<:BigFloats} = (b = idct(a); a[:] = b)

# dummy plans
struct DummyFFTPlan{T,inplace} <: Plan{T} end
struct DummyiFFTPlan{T,inplace} <: Plan{T} end
struct DummyrFFTPlan{T,inplace} <: Plan{T} end
struct DummyirFFTPlan{T,inplace} <: Plan{T} end
struct DummyDCTPlan{T,inplace} <: Plan{T} end
struct DummyiDCTPlan{T,inplace} <: Plan{T} end

for (Plan,iPlan) in ((:DummyFFTPlan,:DummyiFFTPlan),
                     (:DummyrFFTPlan,:DummyirFFTPlan),
                     (:DummyDCTPlan,:DummyiDCTPlan))
   @eval begin
       Base.inv(::$Plan{T,inplace}) where {T,inplace} = $iPlan{T,inplace}()
       Base.inv(::$iPlan{T,inplace}) where {T,inplace} = $Plan{T,inplace}()
    end
end

for (Plan,ff,ff!) in ((:DummyFFTPlan,:fft,:fft!),
                      (:DummyiFFTPlan,:ifft,:ifft!),
                      (:DummyrFFTPlan,:rfft,:rfft!),
                      (:DummyirFFTPlan,:irfft,:irfft!),
                      (:DummyDCTPlan,:dct,:dct!),
                      (:DummyiDCTPlan,:idct,:idct!))
    @eval begin
        *(p::$Plan{T,true}, x::StridedArray{T,N}) where {T,N} = $ff!(x)
        *(p::$Plan{T,false}, x::StridedArray{T,N}) where {T,N} = $ff(x)
        function LAmul!(C::StridedVector, p::$Plan, x::StridedVector)
            C[:] = $ff(x)
            C
        end
    end
end

plan_fft!(x::Vector{T}) where {T<:BigFloats} = DummyFFTPlan{Complex{BigFloat},true}()
plan_ifft!(x::Vector{T}) where {T<:BigFloats} = DummyiFFTPlan{Complex{BigFloat},true}()

#plan_rfft!{T<:BigFloats}(x::Vector{T}) = DummyrFFTPlan{Complex{BigFloat},true}()
#plan_irfft!{T<:BigFloats}(x::Vector{T},n::Integer) = DummyirFFTPlan{Complex{BigFloat},true}()
plan_dct!(x::Vector{T}) where {T<:BigFloats} = DummyDCTPlan{T,true}()
plan_idct!(x::Vector{T}) where {T<:BigFloats} = DummyiDCTPlan{T,true}()

plan_fft(x::Vector{T}) where {T<:BigFloats} = DummyFFTPlan{Complex{BigFloat},false}()
plan_ifft(x::Vector{T}) where {T<:BigFloats} = DummyiFFTPlan{Complex{BigFloat},false}()
plan_rfft(x::Vector{T}) where {T<:BigFloats} = DummyrFFTPlan{Complex{BigFloat},false}()
plan_irfft(x::Vector{T},n::Integer) where {T<:BigFloats} = DummyirFFTPlan{Complex{BigFloat},false}()
plan_dct(x::Vector{T}) where {T<:BigFloats} = DummyDCTPlan{T,false}()
plan_idct(x::Vector{T}) where {T<:BigFloats} = DummyiDCTPlan{T,false}()

function interlace(a::Vector{S},b::Vector{V}) where {S<:Number,V<:Number}
    na=length(a);nb=length(b)
    T=promote_type(S,V)
    if nb≥na
        ret=zeros(T,2nb)
        ret[1:2:1+2*(na-1)]=a
        ret[2:2:end]=b
        ret
    else
        ret=zeros(T,2na-1)
        ret[1:2:end]=a
        if !isempty(b)
            ret[2:2:2+2*(nb-1)]=b
        end
        ret
    end
end


# Step 2 Run DICE
using MimiDICE2010

m = MimiDICE2010.get_model()
run(m)

# Step 3: Altering Parameters

using Mimi

# update the forcings of equilibrium CO2
update_param!(m, :fco22x, 3.000)
run(m)

# update the time horizon
const ts = 10 # timestep 10 years
const years = collect(1995:ts:2505)
nyears = length(years)
set_dimension!(m, :time, years)


function mediansurvival_analysis(df,variable, grouping)
    dd1 = combine(groupby(df,[:MouseID,grouping]), variable => median => variable)
    dd2 = combine(groupby(dd1,grouping), variable => (t-> (Mean = mean(t),Sem = sem(t))) => AsTable)
    plt = @df dd2 scatter(string.(cols(grouping)), :Mean, yerror = :Sem,
        xlabel = "Group", ylabel = "Median survival time", label = "")
    return dd2, plt
end

function survivalrate_algorythm(variable; step = 0.05, xaxis = nothing)
    isnothing(xaxis) && (xaxis = range(extrema(variable)..., step = step))
    survival = 1 .- ecdf(variable).(xaxis)
    return (Xaxis = collect(xaxis), fy = survival)
end

function hazardrate_algorythm(variable; step = 0.05, xaxis = nothing)
    isnothing(xaxis) && (xaxis = range(extrema(variable)..., step = step))
    survival = 1 .- ecdf(variable).(xaxis)
    hazard = -pushfirst!(diff(survival),0)./survival
    return (Xaxis = collect(xaxis), fy = hazard)
end

function function_analysis(df,variable, f; grouping = nothing, step =0.05)
    subgroups = isnothing(grouping) ? [:MouseID] : vcat(:MouseID,grouping)
    xaxis = range(extrema(df[:, variable])..., step = step)
    dd1 = combine(groupby(df,subgroups), variable => (t-> f(t,xaxis = xaxis)) => AsTable)
    rename!(dd1, Dict(:Xaxis => variable))
    sort!(dd1,[:MouseID,variable])
    dd2 = combine(groupby(dd1,[grouping,variable]), :fy =>(t-> (Mean = mean(t),Sem = sem(t))) => AsTable)
    sort!(dd2,variable)
end

function survival_analysis(df,times,events;grouping = nothing, step =0.05)
    subgroups = isnothing(grouping) ? [:MouseID] : vcat(:MouseID,grouping)
    xaxis = range(extrema(df[:, times])..., step = step)
    dd1 = combine(groupby(df,subgroups), [times, events] => ((t,e)-> _survival(t,e; axis = xaxis)) => AsTable)
    rename!(dd1, Dict(:Xaxis => times))
    sort!(dd1,[:MouseID,times])
    subgroups2 = isnothing(grouping) ? [times] : vcat(times,grouping)
    dd2 = combine(groupby(dd1,subgroups2), :SurvRate =>(t-> (Mean = mean(t),Sem = sem(t))) => AsTable)
    sort!(dd2,times)
    return dd2
end

function _survival(times, events; axis = extrema(time), kwargs...)
    km = fit(KaplanMeier, times, events)
    surv = zeros(length(axis))
    for (i, ax) in enumerate(axis)
        if ax < km.times[1]
            surv[i] = 1
        elseif ax > km.times[end]
            surv[i] = km.survival[end]
        else
            surv[i] = km.survival[searchsortedfirst(km.times, ax)]
        end
    end
    return (Xaxis = collect(axis), SurvRate = surv)
end

function KM_median(KM::T) where T <:KaplanMeier
    idx = findfirst(KM.survival .< 0.5)
    isnothing(idx) ? missing : KM.times[idx-1]
end


# create the input signal
T = 5000;
t = linspace(0, 2*π, T);

x1 = sin(t)+cos(11*t).^2;
x2 = cos(11*t);
x = hcat(x1, x2);

# plot the original signal
using Winston
plot(x2, x1, "b")

# get slow signals
using SFA
y = linsfa(x);

y
plot(y, "r")


export Transition

struct Transition end

@node Transition Stochastic [ out, in, a ]

conjugate_type(::Type{ <: Transition }, ::Type{ Val{ :out } }) = Categorical
conjugate_type(::Type{ <: Transition }, ::Type{ Val{ :in} })   = Categorical
conjugate_type(::Type{ <: Transition }, ::Type{ Val{ :a } })   = MatrixDirichlet

@average_energy Transition (q_out::Any, q_in::Any, q_a::MatrixDirichlet) = begin
    return -probvec(q_out)' * logmean(q_a) * probvec(q_in)
end

@average_energy Transition (q_out_in::Contingency, q_a::MatrixDirichlet) = begin
    return -tr(contingency_matrix(q_out_in)' * logmean(q_a))
end


using ParallelAccelerator

#ParallelAccelerator.ParallelIR.set_debug_level(3)
#ParallelAccelerator.CGen.set_debug_level(3)

@acc function opt_At_mul_B!(U, X, W)
	 #At_mul_B!(U, X, W)  # u <- w'x
	 U = X' * W
end

function main(m::Int, k::Int, n::Int)
    U = Array{Float64}(n, k)
    W = Array{Float64}(m, k)
    X = Array{Float64}(m, n)
    fill!(W, 3)
    fill!(X, 5)
    opt_At_mul_B!(U, X, W)
    println("done")

end

main(100, 100, 200)


#############################################################################
# entropy.jl
# entropy (ie, sum_i( -x_i log (x_i) ) of an expression x
# All expressions and atoms are subtypes of AbstractExpr.
# Please read expressions.jl first.
#############################################################################

export entropy
export sign, curvature, monotonicity, evaluate

### Entropy: sum_i -x_i log (x_i)

# TODO: make this work for a *list* of inputs, rather than just for scalar/vector/matrix inputs

# Entropy atom: -xlogx entrywise
mutable struct EntropyAtom <: AbstractExpr
    head::Symbol
    id_hash::UInt64
    children::Tuple{AbstractExpr}
    size::Tuple{Int, Int}

    function EntropyAtom(x::AbstractExpr)
        if sign(x) == ComplexSign()
            error("The argument should be real but it's instead complex")
        else
            children = (x,)
            # TODO check positivity or enforce it
            return new(:entropy, hash(children), children, size(x))
        end
    end
end

function sign(x::EntropyAtom)
    return NoSign()
end

function monotonicity(x::EntropyAtom)
    return (NoMonotonicity(),)
end

function curvature(x::EntropyAtom)
    return ConcaveVexity()
end

function evaluate(x::EntropyAtom)
    c = evaluate(x.children[1])
    return -c .* log.(c)
end

function conic_form!(e::EntropyAtom, unique_conic_forms::UniqueConicForms=UniqueConicForms())
    if !has_conic_form(unique_conic_forms, e)
        # -x log x >= t  <=>  x exp(t/x) <= 1  <==>  (t,x,1) in exp cone
        t = Variable(e.size)
        x = e.children[1]
        objective = conic_form!(t, unique_conic_forms)
        for i = 1:size(x, 1)
            for j = 1:size(x, 2)
                conic_form!(ExpConstraint(t[i, j], x[i, j], 1), unique_conic_forms)
            end
        end
        cache_conic_form!(unique_conic_forms, e, objective)
    end
    return get_conic_form(unique_conic_forms, e)
end

entropy(x::AbstractExpr) = sum(EntropyAtom(x))


using LinearAlgebra

a = [1, 2, 3, 5, 5, 5, 6]
b = [1, 1, 2, 2, 3, 5]

function DTW(a, b)
    table = fill(Inf, size(a)[1]+1, size(b)[1]+1)
    table[1, 1] = 0
    for i in 2:size(table, 1)
        for j in 2:size(table, 2)
            table[i, j] = abs(a[i-1] - b[j-1])  # Calculate distance of current signal indices
            table[i, j] += min(Inf, table[i-1, j], table[i-1, j-1], table[i, j-1])
        end
    end
    return table[end, end], minimum_warp_path(table)
end

function minimum_warp_path(table)
    current_pos = size(table) |> collect
    at_start_index = false
    shortest_path = []
    push!(shortest_path, copy(current_pos))  # We start at the end and work back

    while at_start_index == false
        current_value = table[current_pos...]
        println(current_pos)
        # We search the direct left, left and up, and upward directions
        # for the smallest distance.
        indexes = ((current_pos[1] - 1, current_pos[2]),      # Left
                   (current_pos[1] - 1, current_pos[2] - 1),  # Up and Left
                   (current_pos[1], current_pos[2] - 1))      # Up

        values = [table[idx...] for idx in indexes]
        min_arg = argmin(values)  # Get direction had the smallest distance

        if min_arg == 1         # Left
            current_pos[1] -= 1
        elseif min_arg == 2     # Up and Left
            current_pos[1] -= 1
            current_pos[2] -= 1
        else                    # Up
            current_pos[2] -= 1
        end

        # Stop once you're at the beginning
        if current_pos == [1, 1]
            at_start_index = true
        end
        pushfirst!(shortest_path, copy(current_pos))
    end
    shortest_path = vcat(shortest_path'...)  # reshape
    return shortest_path
end


facts("PARAFAC2") do

srand(1)
r = 2
A = randn(r, 10)
T = Array{Float64, r}[(randn(4 * (i + 1), r) * A) for i in 1:5]
@time factors = parafac2(T, 2, compute_error=true)
@fact map(size, factors.factors) --> ([(size(t, 1), r) for t in T]...)
@fact length(factors.D) --> length(T)
@fact map(size, factors.D) --> fill((1, r), length(T))
@fact size(factors.A) --> (size(T[1], 2), r)
@fact rel_residue(factors) --> less_than(0.05)

end


export aperture_photometry

@doc raw"""
    aperture_photometry(channel::Array,
                                 obj::Tuple,
                                 ref1::Tuple,
                                 ref2::Tuple,
                                 radius::Tuple
                                )

Aperture photometry of a given object
"""
function aperture_photometry(channel::Array,
                             obj::Tuple,
                             ref1::Tuple,
                             ref2::Tuple,
                             radius::Tuple
                            )

    f_obj  = Array{Float64, 1}(undef, 0)
    f_ref1 = Array{Float64, 1}(undef, 0)
    f_ref2 = Array{Float64, 1}(undef, 0)

    for i in 1:length(channel)

        flux_o_in = Float64(sum(channel[i][obj[1]-radius[1]:obj[1]+radius[1], obj[2]-radius[1]:obj[2]+radius[1]]))
        flux_o_out = Float64(sum(channel[i][obj[1]-radius[2]:obj[1]+radius[2], obj[2]-radius[2]:obj[2]+radius[2]]))

        bg_o = (flux_o_out - flux_o_in) / (radius[2]^2) # bg / pix
        flux_o = flux_o_in - (bg_o * (radius[1]^2))

        flux_r1_in = Float64(sum(channel[i][ref1[1]-radius[1]:ref1[1]+radius[1], ref1[2]-radius[1]:ref1[2]+radius[1]]))
        flux_r1_out = Float64(sum(channel[i][ref1[1]-radius[2]:ref1[1]+radius[2], ref1[2]-radius[2]:ref1[2]+radius[2]]))

        bg_r1 = (flux_r1_out - flux_r1_in) / (radius[2]^2)
        flux_r1 = flux_r1_in - (bg_r1 * (radius[1]^2))

        flux_r2_in = Float64(sum(channel[i][ref2[1]-radius[1]:ref2[1]+radius[1], ref2[2]-radius[1]:ref2[2]+radius[1]]))
        flux_r2_out = Float64(sum(channel[i][ref2[1]-radius[2]:ref2[1]+radius[2], ref2[2]-radius[2]:ref2[2]+radius[2]]))

        bg_r2 = (flux_r2_out - flux_r2_in) / (radius[2]^2)
        flux_r2 = flux_r2_in - (bg_r2 * (radius[1]^2))

        f_obj  = append!(f_obj, flux_o)
        f_ref1 = append!(f_ref1, flux_r1)
        f_ref2 = append!(f_ref2, flux_r2)
    end

    return f_obj, f_ref1, f_ref2

end


function grey(decisionMat::DataFrame, weights::Array{Float64,1}, fs::Array{Function,1}, zeta::Float64=0.5)::GreyResult

    mat = convert(Matrix, decisionMat)

    nrows, ncols = size(mat)

    w = unitize(weights)

    referenceRow = apply_columns(fs, mat)

    normalizedReferenceRow = zeros(Float64, ncols)

    normalizedMat = similar(mat)
    for col in 1:ncols
        mmax = maximum(mat[:, col])
        mmin = minimum(mat[:, col])
        mrange = mmax - mmin
        for row in 1:nrows
            if fs[col] == maximum
                normalizedMat[row, col] = (mat[row, col] - mmin) / mrange
            elseif fs[col] == minimum
                normalizedMat[row, col] = (mmax - mat[row, col]) / mrange
            else
                @error fs[col]
                error("Function not defined")
            end
        end
        if fs[col] == maximum
            normalizedReferenceRow[col] = (referenceRow[col] - mmin) / mrange
        elseif fs[col] == minimum
            normalizedReferenceRow[col] = (mmax - referenceRow[col]) / mrange
        end
    end

    absoluteValueMat = similar(mat)
    for row in 1:nrows
        absoluteValueMat[row,:] = normalizedReferenceRow .- normalizedMat[row, :]  .|> abs
    end

    deltamin = absoluteValueMat |> minimum
    deltamax = absoluteValueMat |> maximum
    zetadeltamax = zeta * deltamax

    greytable = similar(absoluteValueMat)
    for row in 1:nrows
        for col in 1:ncols
            greytable[row, col] = (deltamin + zetadeltamax) / (absoluteValueMat[row, col] + zetadeltamax)
        end
    end

    weightedsums = zeros(Float64, nrows)
    for i in 1:nrows
        weightedsums[i] = w .* greytable[i, :] |> sum
    end

    orderings = sortperm(weightedsums)
    bestIndex = orderings |> last

    result = GreyResult(
        referenceRow,
        normalizedMat,
        absoluteValueMat,
        greytable,
        weightedsums,
        orderings,
        bestIndex
    )

    return result
end


function BlackScholes(CallPutFlag, S, X, T, r, v)

  d1 = (log(S / X) + (r + v^2 / 2) * T) / (v * sqrt(T))
  d2 = d1 - v * sqrt(T)

  if CallPutFlag == "c"
    BS = S * CND(d1) - X * exp(-r * T) * CND(d2)
  elseif CallPutFlag == "p"
    BS = X * exp(-r * T) * CND(-d2) - S * CND(-d1)
  end

end


function ground_reflections(settings::PyObject, data::PyObject, x_obs::Array{Float64}, r, beta, c_bar, rho_0)

    # Size of trajectory
    n_t = size(r)[1]

    # Calculate difference in direct and reflected distance between source and observer
    # Source: Berton - lateral attenuation paper (2019)
    r_r = sqrt.(r.^2 .+ 4 * x_obs[3]^2 .+ 4 * r .* x_obs[3] .* sin.(beta*π/180.))
    dr = r_r - r

    # Source: Zorumski report 1982 part 1. Chapter 3.2 page 1
    k = 2 * π * reshape(data.f_sb, (1, settings.N_f*settings.N_b)) ./ c_bar

    # Calculate dimensionless frequency eta (note: for acoustically hard surface: eta = 0)
    # Source: Zorumski report 1982 part 1. Chapter 3.2 page 2
    eta = 2 * π * rho_0 .* reshape(data.f_sb, (1, settings.N_f*settings.N_b)) / settings.sigma

    # Calculate the cosine of the incidence angle
    cos_theta = (r .* sin.(beta * π / 180.) .+ 2 * x_obs[3]) ./ r_r

    # Complex specific ground admittance nu
    # Source: Zorumski report 1982 part 1. Chapter 3.2 Equation 13 / adapted through Berton lateral attenuation paper
    nu_re = (1 .+ (6.86 * eta).^(-0.75)) ./ ((1 .+ (6.86 * eta).^(-0.75)).^2 .+ ((4.36 * eta).^(-0.73)).^2)
    nu_im = -((4.36 * eta).^(-0.73)) ./ ((1 .+ (6.86 * eta).^(-0.75)).^2 .+ ((4.36 * eta).^(-0.73)).^2)

    # Calculate Gamma
    # Source: Zorumski report 1982 part 1. Chapter 3.2 Equation 5
    Gamma_re = (cos_theta .- nu_re) .* ((cos_theta .+ nu_re) ./ ((cos_theta .+ nu_re).^2 .+ nu_im.^2)) .- (-nu_im) .* (-nu_im ./ ((cos_theta .+ nu_re).^2 .+ nu_im.^2))
    Gamma_im = (-nu_im) .* ((cos_theta .+ nu_re) ./ ((cos_theta .+ nu_re).^2 .+ nu_im.^2)) .+ (cos_theta .- nu_re) .* (-nu_im ./ ((cos_theta .+ nu_re).^2 .+ nu_im.^2))

    # Calculate tau
    # Source: Zorumski report 1982 part 1. Chapter 3.2 Equation 9
    tau_re = (sqrt.(k .* r_r / 2) * cos(-π / 4.)) .* (cos_theta .+ nu_re) .- (sqrt.(k .* r_r / 2) * sin(-π / 4.)) .* nu_im
    tau_im = (sqrt.(k .* r_r / 2) * sin(-π / 4.)) .* (cos_theta .+ nu_re) .+ (sqrt.(k .* r_r / 2) * cos(-π / 4.)) .* nu_im

    # Calculate complex spherical wave reflection coefficient
    # Source: Zorumski report 1982 part 1. Chapter 3.2 Equation 12
    U = zeros(eltype(r), (n_t, settings.N_b*settings.N_f))
    U[findall(-tau_re.>0)] = (tau_re.^0)[findall(-tau_re.>0)]
    U[findall(-tau_re .== 0)] = 0.5*(tau_re.^0)[findall(-tau_re .== 0)]

    # Calculate F
    # Source: Zorumski report 1982 part 1. Chapter 3.2 Equation 11
    tau_sq_re = (tau_re.^2 + tau_im.^2) .* (cos.(2 * atan.(tau_im ./ tau_re)))
    tau_sq_im = (tau_re.^2 + tau_im.^2) .* (sin.(2 * atan.(tau_im ./ tau_re)))
    tau_4th_re = (tau_sq_re.^2 .+ tau_sq_im.^2) .* (cos.(2 * atan.(tau_sq_im ./ tau_sq_re)))
    tau_4th_im = (tau_sq_re.^2 .+ tau_sq_im.^2) .* (sin.(2 * atan.(tau_sq_im ./ tau_sq_re)))
    term1_re = -2 * sqrt(π) * U .* (tau_re .* (exp.(tau_sq_re) .* cos.(tau_sq_im)) .- tau_im .* (exp.(tau_sq_re) .* sin.(tau_sq_im)))
    term1_im = -2 * sqrt(π) * U .* (tau_im .* (exp.(tau_sq_re) .* cos.(tau_sq_im)) .+ tau_re .* (exp.(tau_sq_re) .* sin.(tau_sq_im)))

    F_re = term1_re .+ 1/2. .* ( tau_sq_re ./ (tau_sq_re.^2 .+ tau_sq_im.^2)) .+ (-3/4 .* ( tau_4th_re ./ (tau_4th_re.^2 + tau_4th_im.^2)))
    F_im = term1_im .+ 1/2. .* (-tau_sq_im ./ (tau_sq_re.^2 .+ tau_sq_im.^2)) .+ (-3/4 .* (-tau_4th_im ./ (tau_4th_re.^2 + tau_4th_im.^2)))

    # Correct F for small tau
    itau_re = -tau_im
    itau_im = tau_re

    f_faddeeva_real = LinearInterpolation((data.Faddeeva_itau_im, data.Faddeeva_itau_re), data.Faddeeva_real, extrapolation_bc=Flat())
    f_faddeeva_imag = LinearInterpolation((data.Faddeeva_itau_im, data.Faddeeva_itau_re), data.Faddeeva_imag, extrapolation_bc=Flat())
    fad_re = f_faddeeva_real.(itau_im, itau_re)
    fad_im = f_faddeeva_imag.(itau_im, itau_re)
    F_correction_re = 1 .- sqrt(π) .* (tau_re .* fad_re .- tau_im .* fad_im)
    F_correction_im =    - sqrt(π) .* (tau_im .* fad_re .+ tau_re .* fad_im)
    tau_absolute = sqrt.(tau_re.^2 .+ tau_im.^2)

    F_re[findall(tau_absolute .< 10.)] = F_correction_re[findall(tau_absolute .< 10)]
    F_im[findall(tau_absolute .< 10.)] = F_correction_im[findall(tau_absolute .< 10)]

    # Calculate Z_cswfc
    # Source: Zorumski report 1982 part 1. Chapter 3.2 Equation 6
    Z_cswfc_re = Gamma_re .+ ((1 .- Gamma_re) .* F_re .- (    - Gamma_im) .* F_im)
    Z_cswfc_im = Gamma_im .+ ((   - Gamma_im) .* F_re .+ (1. .- Gamma_re) .* F_im)

    # Calculate R and alpha
    R = sqrt.(Z_cswfc_re.^2 .+ Z_cswfc_im.^2)
    alpha = 2 * atan.((R - Z_cswfc_re) ./ Z_cswfc_im)

    # Calculate the constant K and constant epsilon
    # Source: Zorumski report 1982 part 1. Chapter 3.2 Equation 16-17
    K = 2^(1 / (6 * settings.N_b))
    eps = K - 1

    # Calculate G
    # Source: Zorumski report 1982 part 1. Chapter 3.2 Equation 18
    G = zeros(eltype(r), (n_t, settings.N_b*settings.N_f))

    for i in range(1, n_t, step=1)
        if dr[i] == 0
            G[i,:] = 1 .+ R[i,:].^2 .+ 2 * R[i,:] .* exp.(-(settings.a_coh * k[i,:] .* dr[i]).^2) .* cos.(alpha[i,:] + k[i,:] .* dr[i])
        else
            G[i,:] = 1 .+ R[i,:].^2 .+ 2 * R[i,:] .* exp.(-(settings.a_coh * k[i,:] .* dr[i]).^2) .* cos.(alpha[i,:] + k[i,:] .* dr[i]) .* sin.(eps .* k[i,:] .* dr[i]) ./ (eps .* k[i,:] .* dr[i])
        end
    end

    return G
end


using QuantileRegressions, Distributions, CSV
mycyle = CSV.read("/home/pkofod/mcycle.csv")
x = mycyle.times
y = mycyle.accel

for t = range(0.01,0.99,length=99)
res = QuantileRegressions.npqreg(Array(y), Array(x), t; h=2)
plot!(res[1], res[2]; label="tau=$t")
end


##### convert the coordinates into fractional indices at grid cell center and face for a individual at (x,y,z)
##### 1-based indexing for grid faces (excluding halo regions)
@inline get_xf_index(x) = x + 1
@inline get_yf_index(y) = y + 1
@inline get_zf_index(z) = z + 1

##### 1-based indexing for grid centers (excluding halo regions)
@inline get_xc_index(x) = x + 0.5
@inline get_yc_index(y) = y + 0.5
@inline get_zc_index(z) = z + 0.5

##### trilinear interpolation
@inline ψ₀₀₀(xd, yd, zd) = (1 - xd) * (1 - yd) * (1 - zd)
@inline ψ₀₀₁(xd, yd, zd) = (1 - xd) * (1 - yd) *      zd
@inline ψ₀₁₀(xd, yd, zd) = (1 - xd) *      yd  * (1 - zd)
@inline ψ₀₁₁(xd, yd, zd) = (1 - xd) *      yd  *      zd
@inline ψ₁₀₀(xd, yd, zd) =      xd  * (1 - yd) * (1 - zd)
@inline ψ₁₀₁(xd, yd, zd) =      xd  * (1 - yd) *      zd
@inline ψ₁₁₀(xd, yd, zd) =      xd  *      yd  * (1 - zd)
@inline ψ₁₁₁(xd, yd, zd) =      xd  *      yd  *      zd

@inline normalized_trilinear_interpolation_u(u, xd, yd, zd, xi, yi, zi, g::AbstractGrid) =
        @inbounds (ψ₀₀₀(xd, yd, zd) * u[xi,   yi,   zi  ] / ΔxC(xi,   yi,   zi,   g) +
                   ψ₀₀₁(xd, yd, zd) * u[xi,   yi,   zi+1] / ΔxC(xi,   yi,   zi+1, g) +
                   ψ₀₁₀(xd, yd, zd) * u[xi,   yi+1, zi  ] / ΔxC(xi,   yi+1, zi,   g) +
                   ψ₀₁₁(xd, yd, zd) * u[xi,   yi+1, zi+1] / ΔxC(xi,   yi+1, zi+1, g) +
                   ψ₁₀₀(xd, yd, zd) * u[xi+1, yi,   zi  ] / ΔxC(xi+1, yi,   zi,   g) +
                   ψ₁₀₁(xd, yd, zd) * u[xi+1, yi,   zi+1] / ΔxC(xi+1, yi,   zi+1, g) +
                   ψ₁₁₀(xd, yd, zd) * u[xi+1, yi+1, zi  ] / ΔxC(xi+1, yi+1, zi,   g) +
                   ψ₁₁₁(xd, yd, zd) * u[xi+1, yi+1, zi+1] / ΔxC(xi+1, yi+1, zi+1, g))

@inline normalized_trilinear_interpolation_v(v, xd, yd, zd, xi, yi, zi, g::AbstractGrid) =
        @inbounds (ψ₀₀₀(xd, yd, zd) * v[xi,   yi,   zi  ] / ΔyC(xi,   yi,   zi,   g) +
                   ψ₀₀₁(xd, yd, zd) * v[xi,   yi,   zi+1] / ΔyC(xi,   yi,   zi+1, g) +
                   ψ₀₁₀(xd, yd, zd) * v[xi,   yi+1, zi  ] / ΔyC(xi,   yi+1, zi,   g) +
                   ψ₀₁₁(xd, yd, zd) * v[xi,   yi+1, zi+1] / ΔyC(xi,   yi+1, zi+1, g) +
                   ψ₁₀₀(xd, yd, zd) * v[xi+1, yi,   zi  ] / ΔyC(xi+1, yi,   zi,   g) +
                   ψ₁₀₁(xd, yd, zd) * v[xi+1, yi,   zi+1] / ΔyC(xi+1, yi,   zi+1, g) +
                   ψ₁₁₀(xd, yd, zd) * v[xi+1, yi+1, zi  ] / ΔyC(xi+1, yi+1, zi,   g) +
                   ψ₁₁₁(xd, yd, zd) * v[xi+1, yi+1, zi+1] / ΔyC(xi+1, yi+1, zi+1, g))

@inline normalized_trilinear_interpolation_w(w, xd, yd, zd, xi, yi, zi, g::AbstractGrid) =
        @inbounds (ψ₀₀₀(xd, yd, zd) * w[xi,   yi,   zi  ] / ΔzC(xi,   yi,   zi,   g) +
                   ψ₀₀₁(xd, yd, zd) * w[xi,   yi,   zi+1] / ΔzC(xi,   yi,   zi+1, g) +
                   ψ₀₁₀(xd, yd, zd) * w[xi,   yi+1, zi  ] / ΔzC(xi,   yi+1, zi,   g) +
                   ψ₀₁₁(xd, yd, zd) * w[xi,   yi+1, zi+1] / ΔzC(xi,   yi+1, zi+1, g) +
                   ψ₁₀₀(xd, yd, zd) * w[xi+1, yi,   zi  ] / ΔzC(xi+1, yi,   zi,   g) +
                   ψ₁₀₁(xd, yd, zd) * w[xi+1, yi,   zi+1] / ΔzC(xi+1, yi,   zi+1, g) +
                   ψ₁₁₀(xd, yd, zd) * w[xi+1, yi+1, zi  ] / ΔzC(xi+1, yi+1, zi,   g) +
                   ψ₁₁₁(xd, yd, zd) * w[xi+1, yi+1, zi+1] / ΔzC(xi+1, yi+1, zi+1, g))

@inline function u_itpl(u, x, y, z, ac, g::AbstractGrid)
    xi = get_xf_index(x) * ac
    yi = get_yc_index(y) * ac
    zi = get_zc_index(z) * ac
    xd, xi = mod(xi, 1), unsafe_trunc(Int, xi)
    yd, yi = mod(yi, 1), unsafe_trunc(Int, yi)
    zd, zi = mod(zi, 1), unsafe_trunc(Int, zi)
    ##### return fraction of grid spacing
    return normalized_trilinear_interpolation_u(u, xd, yd, zd, xi+g.Hx, yi+g.Hy, zi+g.Hz, g)
end

@inline function v_itpl(v, x, y, z, ac, g::AbstractGrid)
    xi = get_xc_index(x) * ac
    yi = get_yf_index(y) * ac
    zi = get_zc_index(z) * ac
    xd, xi = mod(xi, 1), unsafe_trunc(Int, xi)
    yd, yi = mod(yi, 1), unsafe_trunc(Int, yi)
    zd, zi = mod(zi, 1), unsafe_trunc(Int, zi)
    ##### return fraction of grid spacing
    return normalized_trilinear_interpolation_v(v, xd, yd, zd, xi+g.Hx, yi+g.Hy, zi+g.Hz, g)
end

@inline function w_itpl(w, x, y, z, ac, g::AbstractGrid)
    xi = get_xc_index(x) * ac
    yi = get_yc_index(y) * ac
    zi = get_zf_index(z) * ac
    xd, xi = mod(xi, 1), unsafe_trunc(Int, xi)
    yd, yi = mod(yi, 1), unsafe_trunc(Int, yi)
    zd, zi = mod(zi, 1), unsafe_trunc(Int, zi)
    ##### return fraction of grid spacing
    return normalized_trilinear_interpolation_w(w, xd, yd, zd, xi+g.Hx, yi+g.Hy, zi+g.Hz, g)
end


# Test taken from http://www.stanford.edu/~clint/bench/wilk.txt

using Base.Test
using DataFrames
using StatsBase
using GLM

testeps = sqrt(eps())

nasty = DataFrame(    label = ["One", "Two", "Three", "Four", "Five", "Six", "Seven", "Eight", "Nine"],
                    x = collect(1.:9),
                    zero = fill(0.0,9),
                    miss = fill(NA, 9),
                    big = 99999990.0 + collect(1:9),
                    little = (99999990.0 + collect(1:9))/10^8,
                    huge = collect(1.:9)*1e12,
                    tiny = collect(1.:9)*1e-12,
                    round = collect(0.5:8.5))

println(nasty)
println("\nII Real Numbers:\nII A")
print("Test rounding: ")
@test [@sprintf("%1.0f", x) for x in nasty[:round]] == ["1","2","3","4","5","6","7","8","9"]
println("OK")
print("Test math: ")
@test round(Int, 2.6*7 - 0.2) == 18
@test 2 - round(Int, exp(log(sqrt(2)*sqrt(2)))) == 0
@test round(Int, 3 - exp(log(sqrt(2)*sqrt(2)))) == 1
println("OK")

print("Test means: ")
for vars in names(nasty)[2:end]
    if vars == :miss
    @test isna(mean(nasty[vars]))
    else
    @test mean(nasty[vars]) ≈ nasty[vars][5]
    end
end
println("OK")

print("Test standard deviation: ")
for vars in names(nasty)[[2;5:9]]
#    @test (@sprintf("%.9e", std(nasty[vars])))[1:10] == "2.73861278"
    @test repr(std(nasty[vars]))[1:10] == "2.73861278"
end
println("OK")
# Note: Original had one more digit but little fails then. R gives the same as Julia. Stata and SAS less precise returning 2.738612714381 and 2.738612780003404")

println("\nII D")
print("Test correlation: ")
cn = names(nasty)[[2;5:9]]
for i in 1:5
    for j = i+1:6
    @test cor(nasty[cn[i]], nasty[cn[j]]) ≈ 1
    end
end
println("OK")

print("Test spearman correlation: ")
cn = names(nasty)[[2;5:9]]
for i in 1:5
    for j = i+1:6
    @test corspearman(nasty[cn[i]], nasty[cn[j]]) ≈ 1
    end
end
println("OK")

println("\nII F")
print("Testing regression: ")
ctable = coeftable(lm(big ~ x, nasty))
@test typeof(ctable) == CoefTable
# conversion to Vector{Float64} can be removed when 0.4 is no longer supported
@test Vector{Float64}(ctable.cols[1]) ≈ [99999990, 1]

@test sprint(show, ctable) == """\
             Estimate Std.Error t value Pr(>|t|)
(Intercept)     1.0e8       0.0     Inf   <1e-99
x                 1.0       0.0     Inf   <1e-99
"""

println("OK")

println("\nIV Regression:\nIV A")
nasty[:x1] = nasty[:x]
nasty[:x2] = nasty[:x].^2
nasty[:x3] = nasty[:x].^3
nasty[:x4] = nasty[:x].^4
nasty[:x5] = nasty[:x].^5
nasty[:x6] = nasty[:x].^6
nasty[:x7] = nasty[:x].^7
nasty[:x8] = nasty[:x].^8
nasty[:x9] = nasty[:x].^9

## Is it intended that the least squares problem be overdetermined in the lm fit?
## n = 9 and p = 10 because of the implicit intercept.
lm(x1~x2+x3+x4+x5+x6+x7+x8+x9, nasty)
@test coef(lm(x~x, nasty)) ≈ [0,1]
println("OK")


using ModelingToolkit, DiffEqBase, LinearAlgebra

# Define some variables
@parameters t x
@variables u(..)
Dt = Differential(t)
Dxx = Differential(x)^2
eq = Dt(u(t, x)) ~ Dxx(u(t, x))
bcs = [u(0, x) ~ -x * (x - 1) * sin(x),
    u(t, 0) ~ 0, u(t, 1) ~ 0]

domains = [t ∈ (0.0, 1.0),
    x ∈ (0.0, 1.0)]

@named pdesys = PDESystem(eq, bcs, domains, [t, x], [u])
@show pdesys

@test all(isequal.(independent_variables(pdesys), [t, x]))


# # Internal wave example
#
# In this example, we initialize an internal wave packet in two-dimensions
# and watch it propagate.

using Oceananigans, Oceananigans.Grids, Plots, Printf

# ## Numerical, domain, and internal wave parameters
#
# First, we pick a resolution and domain size,

Nx = 128 # resolution
Lx = 2π  # domain extent
nothing # hide

# We set up an internal wave with the pressure field
#
# $ p(x, y, z, t) = a(x, z) \, \cos(kx + mz - ω t) $.
#
# where $m$ is the vertical wavenumber, $k$ is the horizontal wavenumber,
# $ω$ is the wave frequncy, and $a(x, z)$ is a Gaussian envelope.

## Non-dimensional internal wave parameters
m = 16      # vertical wavenumber
k = 8       # horizontal wavenumber
N = 1       # buoyancy frequency
f = 0.2     # inertial frequency
nothing # hide

# ## A Gaussian wavepacket
#
# Next, we set up an initial condition corresponding to a propagating
# wave packet with a Gaussian envelope. The internal wave dispersion relation yields

ω² = (N^2 * k^2 + f^2 * m^2) / (k^2 + m^2)

## and thus
ω = sqrt(ω²)
nothing # hide

# The internal wave polarization relations follow from the linearized
# Boussinesq equations,

U = k * ω   / (ω^2 - f^2)
V = k * f   / (ω^2 - f^2)
W = m * ω   / (ω^2 - N^2)
B = m * N^2 / (ω^2 - N^2)
nothing # hide

# Finally, we set-up a small-amplitude, Gaussian envelope for the wave packet

## Some Gaussian parameters
A, x₀, z₀, δ = 1e-9, Lx/2, -Lx/2, Lx/15

## A Gaussian envelope
a(x, z) = A * exp( -( (x - x₀)^2 + (z - z₀)^2 ) / 2δ^2 )
nothing # hide

# Create initial condition functions
u₀(x, y, z) = a(x, z) * U * cos(k*x + m*z)
v₀(x, y, z) = a(x, z) * V * sin(k*x + m*z)
w₀(x, y, z) = a(x, z) * W * cos(k*x + m*z)
b₀(x, y, z) = a(x, z) * B * sin(k*x + m*z) + N^2 * z
nothing # hide

# We are now ready to instantiate our model on a uniform grid.
# We give the model a constant rotation rate with background vorticity `f`,
# use temperature as a buoyancy tracer, and use a small constant viscosity
# and diffusivity to stabilize the model.

model = IncompressibleModel(    grid = RegularCartesianGrid(size=(Nx, 1, Nx), extent=(Lx, Lx, Lx)),
                             closure = IsotropicDiffusivity(ν=1e-6, κ=1e-6),
                            coriolis = FPlane(f=f),
                             tracers = :b,
                            buoyancy = BuoyancyTracer())
nothing # hide

# We initialize the velocity and buoyancy fields
# with our internal wave initial condition.

set!(model, u=u₀, v=v₀, w=w₀, b=b₀)

# ## A wave packet on the loose
#
# Finally, we release the packet and watch it go!

simulation = Simulation(model, Δt = 0.001 * 2π/ω, stop_iteration = 0,
                        iteration_interval = 20)

anim = @animate for i=0:100
    x, z = xnodes(Cell, model.grid)[:], znodes(Face, model.grid)[:]
    w = model.velocities.w

    contourf(x, z, w.data[1:Nx, 1, 1:Nx+1]',
                   title = @sprintf("ωt = %.2f", ω * model.clock.time),
                  levels = range(-1e-8, stop=1e-8, length=10),
                   clims = (-1e-8, 1e-8),
                  xlabel = "x",
                  ylabel = "z",
                   xlims = (0, Lx),
                   ylims = (-Lx, 0),
               linewidth = 0,
                   color = :balance,
                  legend = false,
             aspectratio = :equal)

    simulation.stop_iteration += 20
    run!(simulation)
end

mp4(anim, "internal_wave.mp4", fps = 15) # hide


solve(i::MSetInstance, ::LinearProgramming; kwargs...) = msets_lp(i; kwargs...)

function msets_lp(i::MSetInstance; solver=nothing)
  model = Model(solver)
  @variable(model, x[1:length(values(i))], Bin)
  @objective(model, Max, dot(x, values(i)))
  @constraint(model, sum(x) <= m(i))

  set_silent(model)
  optimize!(model)

  return MSetSolution(i, findall(JuMP.value.(x) .>= 0.5))
end


using ChaosTools
using Test
using OrdinaryDiffEq: Tsit5

test_value = (val, vmin, vmax) -> @test vmin <= val <= vmax

println("\nTesting GALI...")

@testset "GALI discrete" begin

    @testset "Chaotic - towel map" begin
        ds = Systems.towel()
        model(x,p)= @. exp(-p[1]*x)
        ls = lyapunovspectrum(ds, 10000)
        threshold = 1e-16
        @testset "k=$k" for k in [2,3]
            ex = sum(ls[1] - ls[j] for j in 2:k)
            g, t = gali(ds, 1000, k; threshold = threshold)
            @test g[end] < threshold
            @test t[end] < 1000
            if k == 3
                fite = -linreg(t, log.(g))[2]
                @test isapprox(fite, ex, rtol=1)
            end
        end
    end

    @testset "Standard Map" begin
        ds = Systems.standardmap()
        k = 2
        @testset "chaotic" begin
            g, t = gali(ds, 1000, k)
            @test g[end] ≤ 1e-12
            @test t[end] < 1000
        end
        @testset "regular" begin
            g, t = ChaosTools.gali(ds, 1000, k; u0 =  [π, rand()])
            @test t[end] == 1000
            @test g[end] > 1/1000^2
        end
    end

    @testset "Coupled Standard Maps" begin
        M = 3; ks = 3ones(M); Γ = 0.1;
        stable = [π, π, π, 0.01, 0, 0] .+ 0.1
        chaotic = rand(2M)

        ds = Systems.coupledstandardmaps(M, stable; ks=ks, Γ = Γ)

        @testset "regular k=$k" for k in [2,3,4, 5, 6]
            g, t = gali(ds, 1000, k; threshold=1e-12)
            @test t[end] == 1000
            @test g[end] > 1e-12
        end

        @testset "chaotic k=$k" for k in [2,3,4, 5, 6]
            g, t = gali(ds, 1000, k; threshold=1e-12, u0 = chaotic)
            @test t[end] < 1000
            @test g[end] ≤ 1e-12
        end

    end
end

@testset "GALI continuous" begin
    @testset "Henon-Helies" begin
        sp = [0, .295456, .407308431, 0] #stable periodic orbit: 1D torus
        qp = [0, .483000, .278980390, 0] #quasiperiodic orbit: 2D torus
        ch = [0, -0.25, 0.42081, 0] # chaotic orbit
        tt = 1000
        ds = Systems.henonheiles()
        diffeq = Dict(:abstol=>1e-9, :reltol=>1e-9, :solver => Tsit5())
        @testset "regular" begin
            for k in [2,4]
                g, t = gali(ds, tt, k; diff_eq_kwargs = diffeq, u0 = sp)
                @test t[end] ≥ tt
            end
            for k in [2,4]
                g, t = gali(ds, tt, k; diff_eq_kwargs = diffeq, u0 = qp)
                @test t[end] ≥ tt
            end
        end
        @testset "chaotic" begin
            for k in [2,4]
                g, t = gali(ds, tt, k; diff_eq_kwargs = diffeq, u0 = ch)
                @test t[end] < tt
                @test g[end] ≤ 1e-12
            end
        end
    end

    # Test henon helies out of place & autodiffed
    function hhoop(u, p, t)
        du1 = u[3]
        du2 = u[4]
        du3 = -u[1] - 2u[1]*u[2]
        du4 = -u[2] - (u[1]^2 - u[2]^2)
        return SVector{4}(du1, du2, du3, du4)
    end
    @testset "Henon-Helies OOP" begin
        sp = [0, .295456, .407308431, 0] #stable periodic orbit: 1D torus
        qp = [0, .483000, .278980390, 0] #quasiperiodic orbit: 2D torus
        ch = [0, -0.25, 0.42081, 0] # chaotic orbit
        tt = 1000
        ds = Systems.ContinuousDynamicalSystem(hhoop, sp, nothing)
        diffeq = Dict(:abstol=>1e-9, :reltol=>1e-9, :solver => Tsit5())
        @testset "regular" begin
            for k in [2,4]
                g, t = gali(ds, tt, k; diff_eq_kwargs = diffeq, u0 = sp)
                @test t[end] ≥ tt
            end
            for k in [2,4]
                g, t = gali(ds, tt, k; diff_eq_kwargs = diffeq, u0 = qp)
                @test t[end] ≥ tt
            end
        end
        @testset "chaotic" begin
            for k in [2,4]
                g, t = gali(ds, tt, k; diff_eq_kwargs = diffeq, u0 = ch)
                @test t[end] < tt
                @test g[end] ≤ 1e-12
            end
        end
    end

end


#simple MNIST example with training on gpu
# once we build model , we simply push
# the loss function, and datasets to gpu to achieve this.
using Base64
using Flux,Statistics,BenchmarkTools, MLBase, Flux.Data.MNIST
using Flux: onehotbatch, onecold
using Base.Iterators: repeated
using ColorTypes
using MLJBase
using CUDA

#load training dataset
imgs = Flux.Data.MNIST.images()
print(size(imgs))
# flatten all thsimple_mniste 28x28 into 764 input [reshape(image,:) for image in imgs ]
# Reorder the layout of the data for the ANN
imagestrip(image::Matrix{<:ColorTypes.Gray}) = Float32.(reshape(image, :))
X = hcat(imagestrip.(imgs)...)
#Target to one-hot
labels = Flux.Data.MNIST.labels()
Y = onehotbatch(labels,0:9)
# defining the model ( a neural network )
m = Chain(
                Dense(784,32,relu),
                Dense(32,10),
                softmax
         ) |> gpu

loss(x, y) = Flux.Losses.crossentropy(m(X),Y)
data = Flux.Data.DataLoader((X, Y), batchsize=128) |> gpu
#data = Base.Iterators.repeated(data,200)
opt = Flux.Optimise.ADAM(0.1) |> gpu

# test dataset
tX = hcat(float.(reshape.(Flux.Data.MNIST.images(:test), :))...) |> gpu
tY = onehotbatch(Flux.Data.MNIST.labels(:test), 0:9) |> gpu
# train for 10 epochs
using IterTools: ncycle
@btime Flux.Optimise.train!(loss, Flux.params(m), ncycle(data, 2), opt, cb =  () -> println(loss(tX,tY)))
acc(x, y) = mean(onecold((m(x))) .== onecold(y))
@show acc(tX, tY)

#@show MLJBase.confusion_matrix(onecold((m(tX))),onecold(tY))


# m5.2s.jl

using Pkg, DrWatson

using StanQuap
using StatisticalRethinking

df = CSV.read(sr_datadir("WaffleDivorce.csv"), DataFrame);
scale!(df, [:Marriage, :MedianAgeMarriage, :Divorce])

stan5_2 = "
data {
      int N;
      vector[N] divorce_s;
      vector[N] marriage_s;
}
parameters {
      real a;
      real bM;
      real<lower=0> sigma;
}
model {
      vector[N] mu = a + bM * marriage_s;
      a ~ normal( 0 , 0.2 );
      bM ~ normal( 0 , 0.5 );
      sigma ~ exponential( 1 );
      divorce_s ~ normal( mu , sigma );
}
";

data = (N=size(df, 1), divorce_s=df.Divorce_s, marriage_s=df.Marriage_s)
init = (a=1.0, bM=1.0, sigma=10.0)
q5_2s, m5_2s, o5_2s = stan_quap("m5.2s", stan5_2; data, init);

if !isnothing(q5_2s)
      quap5_2s_df = sample(q5_2s)
      precis(quap5_2s_df)
end

# End m5.2s.jl


struct PSHestonAssetPricingModelParameters <: PSAbstractAssetReturnModel

    # data -
    μ::Float64
    𝝹::Float64
    ϴ::Float64
    σ::Float64
    𝜌::Float64

    function PSHestonAssetPricingModelParameters(μ::Float64, 𝝹::Float64, ϴ::Float64, σ::Float64, 𝜌::Float64)
        this = new(μ, 𝝹, ϴ, σ, 𝜌)
    end
end

struct PSGeometricBrownianMotionModelParameters <: PSAbstractAssetReturnModel

    # model parameters -
    μ::Float64
    σ::Float64

    function PSGeometricBrownianMotionModelParameters(μ=0.0, σ=0.0)
        this = new(μ, σ)
    end
end

struct PSSingleIndexModelParameters <: PSAbstractAssetReturnModel

    # model parameters -
    ⍺::Float64
    β::Float64
    riskFreeRate::Float64
    𝝐::ContinuousUnivariateDistribution

    function PSSingleIndexModelParameters(⍺, β, riskFreeRate, 𝝐)
        this = new(⍺, β, riskFreeRate, 𝝐)
    end
end

struct PSRandomWalkModelParameters <: PSAbstractAssetReturnModel

    # model parameters -
    𝝙::UnivariateKDE

    function PSRandomWalkModelParameters(𝝙::UnivariateKDE)
        this = new(𝝙)
    end
end


using SimpleSolvers
using Test
using GeometricIntegrators
using GeometricIntegrators.Utils
using GeometricProblems.LotkaVolterra3d

SimpleSolvers.set_config(:nls_atol, 8eps())
SimpleSolvers.set_config(:nls_rtol, 2eps())

const Δt = 0.01
const nt = 1000

const ref = [0.39947308320241187, 1.9479527336244262, 2.570183075433086]

@testset "$(rpad("Lotka-Volterra 3D",80))" begin
    ode  = lotka_volterra_3d_ode()

    int = Integrator(ode, TableauGauss(1), Δt)
    sol = integrate(ode, int, nt)
    H, ΔH = compute_energy_error(sol.t, sol.q)
    C, ΔC = compute_casimir_error(sol.t, sol.q)
    @test relative_maximum_error(sol.q, ref) < 5E-4
    @test ΔH[end] < 4E-6
    @test ΔC[end] < 8E-6

    int = Integrator(ode, TableauGauss(2), Δt)
    sol = integrate(ode, int, nt)
    H, ΔH = compute_energy_error(sol.t, sol.q)
    C, ΔC = compute_casimir_error(sol.t, sol.q)
    @test relative_maximum_error(sol.q, ref) < 2E-9
    @test ΔH[end] < 5E-11
    @test ΔC[end] < 2E-11

end


f = x-> 0;
u_exact = x-> x;

a = 0;
b = 1;
n = 9;

problem = SparseFiniteDifferenceBVPProblem(0, 1, n, f, NeumannBC(1), NeumannBC(1));
assemble_system!(problem);
u = solve_bvp(problem);
norm(@. u - problem.x - u[1] )<1e-13


######
# This file is part of the BivariateCopulas.jl package.
#
#   Module definition
#
#           University of Liverpool, Institute for Risk and Uncertainty
#
#                                           Authors: Ander Gray & Liam T. Henry
#                                           Email:  ander.gray@liverpool.ac.uk
######

module BivariateCopulas

using Reexport

@reexport using Distributions             #   https://github.com/JuliaStats/Distributions.jl
using LinearAlgebra                       #   Main Library
using Interpolations                      #   https://github.com/JuliaMath/Interpolations.jl

import Distributions: cdf, quantile, rand

import Base: rand

global  n  = 200;                       # Number of descretizations
global bOt = 0.0001;                    # smallest quantile to use if left tail is unbounded. Should we use these to define the copula ranges?
global tOp = 0.9999;                    # largest quantile to use if right tail is unbounded
#global bOt = 0;
#global tOp = 1;

abstract type AbstractCopula <: Real end
abstract type AbstractJoint <: Real end
abstract type AbstractMarginal <: ContinuousUnivariateDistribution end

include("copulas.jl")
include("NormalDistribution.jl")

export
    copula, density, cdf, sample, rand, CholeskyGaussian,
    conditional, calcCdf, calcDensity, M, W, Pi, Frank, Clayton, Gaussian,
    norm_cdf, bivariate_cdf, mvNormCdf, Joint, invCdf, marginal, quantile,
    ρCopula, SpearmanCopula, KendalCopula, τCopula

end #BivariateCopulas


using JuMP
using Ipopt
using CSV, DataFrames, DataStructures, MathOptFormat
using Statistics,Distributions, StatsFuns, StatsBase
using PyPlot

function LRO_opt(df,gamma_step,moment_info, x_bounds=nothing, recover_SAA=nothing)
    # calculate optimal stocking quantity x_var using LRO framework (convexification of problem (14) in the paper)
    # df - dataframe, historical observed demand data
    #gamma_step - int, amount to vary gamma by (for experments)
    #moment_info - boolean, true if using moment information, false else
    #x_bounds - int or nothing, value to fix x to (for experiments), free otherwise

    b = 1 #underage cost
    h = 1 #overage cost

    N_i = get_occurences(df)  #construct N_i observation data
    N = sum(N_i[:,"Occurence"])

    gamma_temp = get_gamma_val(N_i, N)  #get gamma value

    gamma = gamma_temp += gamma_step

    model = Model(with_optimizer(Ipopt.Optimizer,print_level=0))

    @variable(model, lambda_var >=0)

    @variable(model, y_vars[1:length(N_i[:,"Occurence"])] >=0)

    if isnothing(x_bounds)
        @variable(model, x_var)
    else
        @variable(model, x_bounds <= x_var <= x_bounds)
    end

    if moment_info
        @variable(model, mu_vars[1:3])

    else
        @variable(model, mu_var)
    end

    z = sum(N_i[:,"Occurence"] .* log.(N_i[:,"Occurence"])) #to store log likelihood

    var_len = length(N_i[:,"Occurence"])

    if !isnothing(recover_SAA) #only if testing if we can recover SAA for sanity
        data_size = length(df[:,"Column1"])

        @variable(model, z_vars[1:data_size]) #define variable to linearize max behavior of objective function

        @constraint(model,con3[i = 1:data_size], z_vars[i] >= x_var - df[:,"Column1"][i])
        @constraint(model, con4[i = 1:data_size], z_vars[i] >= df[:,"Column1"][i] - x_var )

        @constraint(model, lambda_var == 1/N *(sum(z_vars[i]  for i in 1:var_len) - mu_var)) # analytically solve for lambda according to Proposition 1, and p_i is empirical distribution
    end

    if moment_info
        mu_hat = mean(df[:,"Column1"])
        sigma_sq_hat = var(df[:,"Column1"])
        b_vec = [1, mu_hat, mu_hat^2 + sigma_sq_hat]

        @constraint(model,con1[i = 1:var_len], b * (N_i[:,"Demand"][i] - x_var) + y_vars[i] <= mu_vars[1]
        + mu_vars[2] * N_i[:,"Demand"][i] + mu_vars[3] * N_i[:,"Demand"][i]^2)

        @constraint(model, con2[i = 1:var_len], h * (x_var - N_i[:,"Demand"][i]) + y_vars[i] <= mu_vars[1]
        + mu_vars[2] * N_i[:,"Demand"][i] + mu_vars[3] * N_i[:,"Demand"][i]^2)
    else
        @constraint(model,con1[i = 1:var_len], b * (N_i[:,"Demand"][i] - x_var) + y_vars[i] <= mu_var)
        @constraint(model, con2[i = 1:var_len], h * (x_var - N_i[:,"Demand"][i]) + y_vars[i] <= mu_var)
    end

    if moment_info
        @NLobjective(model, Min, sum(b_vec[i] * mu_vars[i] for i in 1:3) + lambda_var * (z - N - gamma) + N * lambda_var * log(lambda_var)
         - lambda_var * sum(N_i[:,"Occurence"][j] * log(y_vars[j]) for j in 1:var_len))
    else
        @NLobjective(model, Min, mu_var + lambda_var * (z - N - gamma) + N * lambda_var * log(lambda_var)
         - lambda_var * sum(N_i[:,"Occurence"][j] * log(y_vars[j]) for j in 1:var_len))
    end

    write_LP(model,"model")

    optimize!(model)

    if isnothing(x_bounds)
        return value(x_var) # return x* (optimal stocking quantity) for experiments
    else
        return objective_value(model) # return objective value for experiments
    end
end


using Test
using AeroAcoustics
import DSP

# Generate a test signal, a 2 Vrms sine wave at 1234 Hz, corrupted by
# 0.001 V**2/Hz of white noise sampled at 10 kHz.
# From scipy.signal.welch docs:
# https://github.com/scipy/scipy/blob/f2ec91c4908f9d67b5445fbfacce7f47518b35d1/scipy/signal/spectral.py#L223

@testset "CSM test:" begin
    fs = 10e3
    N = 1e5
    amp = 2*sqrt(2)
    fr = 1234.0
    noise_power = 0.001*fs/2
    time = (1:N)/fs
    x = amp*sin.(2*pi*fr.*time)
    x .+= sqrt.(noise_power).*rand(length(time))
    n = 1024
    Pxx_d = csm(x;n=n,noverlap=div(n,2),win=DSP.hanning(n),fs=fs,scaling="density")
    fc = Pxx_d.fc
    Pxx_d1 = real.(Pxx_d[1,1,:])
    Pxx_w = DSP.welch_pgram(x,n;onesided=true,fs=fs,window=DSP.hanning(n))
    @test Pxx_d1 ≈ Pxx_w.power
    @test Pxx_d == AeroAcoustics.csm_slow(x;n=n,noverlap=div(n,2),win=DSP.hanning(n),fs=fs,scaling="density")

    Pxx_s = csm(x;n=n,noverlap=div(n,2),win=DSP.hanning(n),fs=fs,scaling="spectrum")
    fc = Pxx_s.fc
    Pxx_s1 = real.(Pxx_s[1,1,:])
    ENBW = enbw(fs,DSP.hanning(1024))
    @test Pxx_s1 ≈ Pxx_w.power*ENBW
    @test Pxx_s == AeroAcoustics.csm_slow(x;n=n,noverlap=div(n,2),win=DSP.hanning(n),fs=fs,scaling="spectrum")
end


export rasterize

# --------- #
# Constants #
# --------- #

const film_aperture = (0.980f0, 0.735f0)

# --------- #
# Utilities #
# --------- #
"""
    edge_function(pt1::Vec3, pt2::Vec3, point::Vec3)

Checks on which side of the line formed by `pt1` and `pt2` does
`point` lie.
"""
edge_function(pt1::Vec3, pt2::Vec3, point::Vec3) = edge_function_vector(pt1, pt2, point)[]

"""
    edge_function_vector(pt1::Vec3, pt2::Vec3, point::Vec3)

Vectoried form of the [`edge_function`](@ref)
"""
edge_function_vector(pt1::Vec3, pt2::Vec3, point::Vec3) =
    ((point.x .- pt1.x) .* (pt2.y - pt1.y) .- (point.y .- pt1.y) .* (pt2.x .- pt1.x))

"""
    convert2raster(vertex_world::Vec3, world_to_camera, left::Real, right::Real,
                   top::Real, bottom::Real, width::Int, height::Int)
    convert2raster(vertex_camera::Vec3{T}, left::Real, right::Real, top::Real, bottom::Real,
                   width::Int, height::Int) where {T}

Converts a Point in 3D world space to the 3D raster space. The conversion is done by the
following steps:

```math
V_{camera} = World2CameraTransform(V_{world})
```
```math
V_{screen_x} = -\\frac{V_{camera_x}}{V_{camera_z}}
```
```math
V_{screen_y} = -\\frac{V_{camera_y}}{V_{camera_z}}
```
```math
V_{NDC_x} = \\frac{2 \\times V_{screen_x} - right - left}{right - left}
```
```math
V_{NDC_y} = \\frac{2 \\times V_{screen_y} - top - bottom}{top - bottom}
```
```math
V_{raster_x} = \\frac{V_{NDC_x} + 1}{2 \\times width}
```
```math
V_{raster_y} = \\frac{1 - V_{NDC_y}}{2 \\times height}
```
```math
V_{raster_z} = - V_{camera_z}
```
"""
function convert2raster(vertex_world::Vec3, world_to_camera, left::Real, right::Real,
                        top::Real, bottom::Real, width::Int, height::Int)
    vertex_camera = world2camera(vertex_world, world_to_camera)

    return convert2raster(vertex_camera, left, right, top, bottom, width, height)
end

function convert2raster(vertex_camera::Vec3{T}, left::Real, right::Real, top::Real, bottom::Real,
                        width::Int, height::Int) where {T}
    outtype = eltype(T)

    vertex_screen = (x = vertex_camera.x[] / -vertex_camera.z[],
                     y = vertex_camera.y[] / -vertex_camera.z[])

    vertex_NDC = (x = outtype((2 * vertex_screen.x - right - left) / (right - left)),
                  y = outtype((2 * vertex_screen.y - top - bottom) / (top - bottom)))

    vertex_raster = Vec3([(vertex_NDC.x + 1) / 2 * outtype(width)],
                         [(1 - vertex_NDC.y) / 2 * outtype(height)],
                         -vertex_camera.z)

    return vertex_raster
end

# ---------- #
# Rasterizer #
# ---------- #

"""
   rasterize(cam::Camera, scene::Vector)
   rasterize(cam::Camera, scene::Vector, camera_to_world,
             world_to_camera, top, right, bottom, left)

Implements the raterization algorithm. This is extremely fast when compared
to the [`raytrace`](@ref) function. However, the image generated is much less
photorealistic with no lighting effects.
"""
function rasterize(cam::Camera, scene::Vector)
    top, right, bottom, left = compute_screen_coordinates(cam, film_aperture)
    camera_to_world = get_transformation_matrix(cam)
    world_to_camera = inv(camera_to_world)
    return rasterize(cam, scene, camera_to_world, world_to_camera, top,
                     right, bottom, left)
end

function rasterize(cam::Camera{T}, scene::Vector, camera_to_world,
                   world_to_camera, top, right, bottom, left) where {T}
    width = cam.fixedparams.width
    height = cam.fixedparams.height

    frame_buffer = Vec3(zeros(eltype(T), width * height))
    depth_buffer = fill(eltype(T)(Inf), width, height)

    for triangle in scene
        v1_camera = world2camera(triangle.v1, world_to_camera)
        v2_camera = world2camera(triangle.v2, world_to_camera)
        v3_camera = world2camera(triangle.v3, world_to_camera)

        normal = normalize(cross(v2_camera - v1_camera, v3_camera - v1_camera))

        v1_raster = convert2raster(v1_camera, left, right, top, bottom, width, height)
        v2_raster = convert2raster(v2_camera, left, right, top, bottom, width, height)
        v3_raster = convert2raster(v3_camera, left, right, top, bottom, width, height)

        # Bounding Box
        xmin, xmax = extrema([v1_raster.x[], v2_raster.x[], v3_raster.x[]])
        ymin, ymax = extrema([v1_raster.y[], v2_raster.y[], v3_raster.y[]])

        (xmin > width || xmax < 1 || ymin > height || ymax < 1) && continue

        area = edge_function(v1_raster, v2_raster, v3_raster)

        # Loop over only the covered pixels
        x₁ = max(     1, Int(ceil(xmin)))
        x₂ = min( width, Int(ceil(xmax)))
        y₁ = max(     1, Int(ceil(ymin)))
        y₂ = min(height, Int(ceil(ymax)))

        y = y₁:y₂
        x = x₁:x₂

        y_space = repeat(collect(y), inner = length(x))
        x_space = repeat(collect(x), outer = length(y))

        w1_arr = Float32[]
        w2_arr = Float32[]
        w3_arr = Float32[]
        depth  = Float32[]
        x_arr  = Int[]
        y_arr  = Int[]

        y_vec = y_space .+ 0.5f0
        x_vec = x_space .+ 0.5f0

        pixel = Vec3(x_vec, y_vec, zeros(eltype(x_vec), length(x) * length(y)))
        w1 = edge_function_vector(v2_raster, v3_raster, pixel)
        w2 = edge_function_vector(v3_raster, v1_raster, pixel)
        w3 = edge_function_vector(v1_raster, v2_raster, pixel)

        for (w1_val, w2_val, w3_val, x_val, y_val) in zip(w1, w2, w3, x_space, y_space)
            if w1_val >= 0 && w2_val >= 0 && w3_val >= 0
                w1_val = w1_val / area
                w2_val = w2_val / area
                w3_val = w3_val / area

                depth_val = 1 / (w1_val / v1_raster.z[] + w2_val / v2_raster.z[] +
                                 w3_val / v3_raster.z[])

                if depth_val < depth_buffer[x_val, y_val]
                    update_index!(depth_buffer, x_val, y_val, depth_val)
                    push!(w1_arr, w1_val)
                    push!(w2_arr, w2_val)
                    push!(w3_arr, w3_val)
                    push!(depth, depth_val)
                    push!(x_arr, x_val)
                    push!(y_arr, y_val)
                end
            end
        end

        length(w1_arr) == 0 && continue

        px = (v1_camera.x[] / -v1_camera.z[]) .* w1_arr .+
             (v2_camera.x[] / -v2_camera.z[]) .* w2_arr .+
             (v3_camera.x[] / -v3_camera.z[]) .* w3_arr

        py = (v1_camera.y[] / -v1_camera.z[]) .* w1_arr .+
             (v2_camera.y[] / -v2_camera.z[]) .* w2_arr .+
             (v3_camera.y[] / -v3_camera.z[]) .* w3_arr

        # Passing these gradients as 1.0f0 is incorrect
        pt = Zygote.hook(Δ -> Vec3([1.0f0 for _ in pt.x]),
                         camera2world(Vec3(px, py, ones(Float32, length(px)) * -1) * depth,
                                      camera_to_world))

        col = get_color(triangle, pt, Val(:diffuse))

        idx = x_arr .+ (y_arr .- 1) .* height

        frame_buffer = place_idx!(frame_buffer, col, idx)
    end

    return frame_buffer
end


include("plots.jl")

function get_max(coordinates, data)
    max, max_coor = findmax(data)
    return max, coordinates[:, max_coor]
end

function get_metrics(scan::ScanData, params::ScanParameters)
    voltage = scan.waveform
    pressure = voltage * params.calibration_factor

    isppa = calc_intensity_sppa(pressure, params.medium, params.excitation)
    ispta = calc_intensity_spta(pressure, params.medium, params.excitation)
    mi = calc_mechanical_index(pressure, params.excitation)

    isppa_max =  get_max(scan.coordinates, isppa.val)
    ispta_max = get_max(scan.coordinates, ispta.val)
    mi_max = get_max(scan.coordinates, mi.val)

    return ScanMetric(
        params,
        pressure,
        isppa,
        isppa_max,
        ispta,
        ispta_max,
        mi,
        mi_max,
     )
end

function compute_metrics(scan::ScanData, params::ScanParameters)
    return ScanData(
        scan.version,
        scan.axes,
        scan.scope_info,
        scan.time,
        scan.waveform,
        scan.coordinates,
        get_metrics(scan, params),
    )
end

function calc_intensity_sppa(pressure::PressureArray{N}, medium, excitation) where {N}
    return Metric{ISPPA, N-1}(IntensityMetrics.intensity_sppa(pressure, medium, excitation))
end

function calc_intensity_spta(pressure::PressureArray{N}, medium, excitation) where {N}
    return Metric{ISPTA, N-1}(IntensityMetrics.intensity_spta(pressure, medium, excitation))
end

function calc_mechanical_index(pressure::PressureArray{N}, excitation) where {N}
    return Metric{MI, N-1}(IntensityMetrics.mechanical_index(pressure, excitation))
end


import DielectricUpdateTechique as DUTCH
using Test

x = collect(0:1e-3:3e-3)
y = collect(0:1e-3:3e-3)
z = collect(0:1e-3:3e-3)

vσ = rand(3, 3, 3);
vϵ = rand(3, 3, 3);

yeeTup = DUTCH.cell_to_yee_dielectric(vσ, vϵ, x, y, z)

@test size(yeeTup[:σˣ]) == (3, 4, 4)
@test size(yeeTup[:σʸ]) == (4, 3, 4)
@test size(yeeTup[:σᶻ]) == (4, 4, 3)

@test round(yeeTup[:σˣ][2,2,2], digits=5) == round(sum(vec(vσ[2,1:2,1:2]))/4, digits=5)
@test round(yeeTup[:σˣ][2,1,3], digits=5) == round(sum(vec(vσ[2,1,2:3]))/2, digits=5)
@test round(yeeTup[:σˣ][2,1,4], digits=5) == round(vσ[2,1,3], digits=5)

@test round(yeeTup[:ϵˣ][2,2,2], digits=5) == round(sum(vec(vϵ[2,1:2,1:2]))/4, digits=5)
@test round(yeeTup[:ϵˣ][2,1,3], digits=5) == round(sum(vec(vϵ[2,1,2:3]))/2, digits=5)
@test round(yeeTup[:ϵˣ][2,1,4], digits=5) == round(vϵ[2,1,3], digits=5)

@test round(yeeTup[:σʸ][2,2,2], digits=5) == round(sum(vec(vσ[1:2,2,1:2]))/4, digits=5)
@test round(yeeTup[:σʸ][1,2,3], digits=5) == round(sum(vec(vσ[1,2,2:3]))/2, digits=5)
@test round(yeeTup[:σʸ][1,2,4], digits=5) == round(vσ[1,2,3], digits=5)

@test round(yeeTup[:ϵʸ][2,2,2], digits=5) == round(sum(vec(vϵ[1:2,2,1:2]))/4, digits=5)
@test round(yeeTup[:ϵʸ][1,2,3], digits=5) == round(sum(vec(vϵ[1,2,2:3]))/2, digits=5)
@test round(yeeTup[:ϵʸ][1,2,4], digits=5) == round(vϵ[1,2,3], digits=5)

@test round(yeeTup[:σᶻ][2,2,2], digits=5) == round(sum(vec(vσ[1:2,1:2,2]))/4, digits=5)
@test round(yeeTup[:σᶻ][3,1,2], digits=5) == round(sum(vec(vσ[2:3,1,2]))/2, digits=5)
@test round(yeeTup[:σᶻ][4,1,2], digits=5) == round(vσ[3,1,2], digits=5)

@test round(yeeTup[:ϵᶻ][2,2,2], digits=5) == round(sum(vec(vϵ[1:2,1:2,2]))/4, digits=5)
@test round(yeeTup[:ϵᶻ][3,1,2], digits=5) == round(sum(vec(vϵ[2:3,1,2]))/2, digits=5)
@test round(yeeTup[:ϵᶻ][4,1,2], digits=5) == round(vϵ[3,1,2], digits=5)


import FractionalTimeDG
using OffsetArrays
using PyPlot

include("../FPDE_problem.jl")

Nx = 8
Nt = 4
rx = 4
rt = 3
Mt = 2rt
ppImax = 10
store = FractionalTimeDG.Store(α, rt, Mt, ppImax)

x = OffsetVector(range(0, L, length=Nx+1), 0:Nx)
grid = FractionalTimeDG.SpatialGrid(x, rx, ("essential", "essential"), rx)

U0 = FractionalTimeDG.spatial_L2_projection(u0, grid)
t, U = FractionalTimeDG.FPDEDG(1.0, f, U0, grid, max_t, Nt, rt, Mt, store)

#pts_per_time_interval = 5
#pts_per_space_interval = 5
#xvals, pcwise_t, pcwise_U = FractionalTimeDG.evaluate_pcwise_poly(
#    U, t, grid, pts_per_time_interval, pts_per_space_interval)
τ = [ -1.0, -0.5, 0.0, 0.5, 1.0 ]
pts_per_time_interval = length(τ)
ξ = [ -0.9, -0.5, 0.0, 0.5, 0.9 ]
xvals, pcwise_t, pcwise_U = FractionalTimeDG.evaluate_pcwise_poly(
    U, t, grid, τ, ξ)
τ = [ -1.0, -0.5, 0.0, 0.5, 1.0 ]

figure(1)
Nxvals = length(xvals)
Ntvals = pts_per_time_interval
X = Float64[ xvals[i] for i = 1:Nxvals, j = 1:Ntvals ]
for n = 1:Nt
    T = Float64[ pcwise_t[j,n] for i = 1:Nxvals, j = 1:Ntvals ]
    Z = Float64[ pcwise_U[i,j,n] for i = 1:Nxvals, j = 1:Ntvals ]
    surf(X, T, Z)
end
xlabel(L"$x$")
ylabel(L"$t$")