# In case more abstraction is needed
abstract type AbstractSLM{T} end

"""
```
SLM
```

is a port of Shape Language Modeling (SLM) by John D'Errico,
who implements least squares spline modeling for curve fitting.

See https://www.mathworks.com/matlabcentral/fileexchange/24443-slm-shape-language-modeling for details about the SLM toolbox.

Note that only the features described in the Keywords section have been ported for now.
Other features may be added on as-needed basis.

The main constructor has the form

```
SLM(x::AbstractVector{T}, y::AbstractVector{T}; calculate_stats::Bool = false,
    verbose::Symbol = :low, kwargs...) where {T <: Real}
```

### Inputs
- The vectors `x` and `y` are the points with which the least-squares spline will be fit.

### Keywords
- `degree::Int = 3`: Degree of the spline. A degree 3 spline is a piecewise cubic Hermite spline.
- `scaling::Bool = true`: If true, scale the `y` values and coefficients to avoid numerical issues
- `knots::Int = 6`: Number of knots for the spline
- `C2::Bool = true`: If true, enforce that the spline satisfies C2 continuity
- `λ::T = 1e-4`: Regularization parameter for how smooth the spline is (smaller λ ⇒ smoother spline)
- `left_value::T = NaN`: Controls the spline's value at the left end point. If it is a `NaN`, then the
    spline's value at the left end point is not controlled, i.e. it can be any number.
- `right_value::T = NaN`: Controls the spline's value at the right end point. If it is a `NaN`, then the
    spline's value at the right end point is not controlled, i.e. it can be any number.
- `min_value::T = NaN`: Sets the spline's global minimum value. If it is a `NaN`, then
    no minimum value is imposed.
- `max_value::T = NaN`: Set the spline's global maximum value. If it is a `NaN`, then the
    no maximum value is imposed.
- `min_max_sample_points::AbstractVector{T} =
    [.017037, .066987, .14645, .25, .37059, .5, .62941, .75, .85355, .93301, .98296]`:
    Sample points used to enforce global min/max values. The default points are Chebyshev nodes.
- `increasing::Bool = false`: If true, the spline will monotonically increase over the interpolation domain.
- `decreasing::Bool = false`: If true, the spline will monotonically decrease over the interpolation domain.
- `increasing_intervals::AbstractMatrix{T} = []`: If nonempty, each row of the n × 2 matrix `increasing_intervals` specifies
    the left and right end points of an interval over which the spline monotonically increases.
- `decreasing_intervals::AbstractMatrix{T} = []`: If nonempty, each row of the n × 2 matrix `decreasing_intervals` specifies
    the left and right end points of an interval over which the spline monotonically decreases.
- `concave_up_intervals::AbstractMatrix{T} = []`: If nonempty, each row of the n × 2 matrix `concave_up_intervals` specifies
    the left and right end points of an interval over which the spline is concave up (i.e. convex in the normal sense).
- `concave_down_intervals::AbstractMatrix{T} = []`: If nonempty, each row of the n × 2 matrix `concave_down_intervals` specifies
    the left and right end points of an interval over which the spline is concave down (i.e. concave in the normal sense).
- `init::AbstractVector{T} = []`: If nonempty, `init` is just as an initial guess for the coefficients of the spline.
- `calculate_stats::Bool = false`: If true, calculate statistics about the spline regression (e.g. R²). Currently,
    no statistics are calculated, so setting `calculate_stats` to true does nothing.
- `verbose::Symbol`: Verbosity of information printed during construction of an SLM object.
"""
mutable struct SLM{T} <: AbstractSLM{T}
    stats::NamedTuple
    type::Symbol
    x::AbstractVector{T}
    y::AbstractVector{T}
    knots::AbstractVector{T}
    coef::AbstractArray{T}
    extrapolation::Symbol
end

function Base.show(io::IO, slm::AbstractSLM{T}) where {T <: Real}
    @printf io "SLM with element type %s\n" string(T)
    @printf io "degree: %i\n" get_stats(slm)[:degree]
    @printf io "knots:  %i\n" length(get_stats(slm)[:knots])
end

# Access functions
get_stats(slm::AbstractSLM) = slm.stats
get_knots(slm::AbstractSLM) = slm.knots
get_type(slm::AbstractSLM) = slm.type
get_x(slm::AbstractSLM) = slm.x
get_y(slm::AbstractSLM) = slm.y
get_coef(slm::AbstractSLM) = slm.coef
get_extrapolation(slm::AbstractSLM) = slm.extrapolation
eltype(slm::AbstractSLM) = slm.stats

function getindex(slm::AbstractSLM, x::Symbol)
    if x == :stats
        get_stats(slm)
    elseif x == :knots
        get_knots(slm)
    elseif x == :x
        get_x(slm)
    elseif x == :y
        get_y(slm)
    elseif x == :coef
        get_coef(slm)
    elseif x == :type
        get_type(slm)
    elseif x == :extrapolation || x == :extrap
        get_extrapolation(slm)
    else
        error("type " * typeof(slm) * " has no field " * string(x))
    end
end

# Main user interface for constructing an SLM object
function SLM(x::AbstractVector{T}, y::AbstractVector{T}; calculate_stats::Bool = false,
             verbose::Symbol = :low, kwargs...) where {T <: Real}

    @assert length(x) == length(y) "x and y must be the same size"
    if verbose == :high
        calculate_stats = true # Statistics will be calculated if verbose is high
    end
    kwargs = Dict{Symbol, Any}(kwargs)

    if haskey(kwargs, :increasing_intervals)
        @assert size(kwargs[:increasing_intervals], 2) == 2 "The matrix for the keyword increasing_intervals must be n × 2."
    end
    if haskey(kwargs, :decreasing_intervals)
        @assert size(kwargs[:decreasing_intervals], 2) == 2 "The matrix for the keyword decreasing_intervals must be n × 2."
    end
    if haskey(kwargs, :concave_up_intervals)
        @assert size(kwargs[:concave_up_intervals], 2) == 2 "The matrix for the keyword concave_up_intervals must be n × 2."
    end
    if haskey(kwargs, :concave_down_intervals)
        @assert size(kwargs[:concave_down_intervals], 2) == 2 "The matrix for the keyword concave_down_intervals must be n × 2."
    end

    # Remove nans
    to_remove = isnan.(x) .| isnan.(y)
    if any(to_remove)
        x = x[to_remove]
        y = y[to_remove]

        if haskey(kwargs, :weights)
            error("Weights are not implemented currently.")
            kwargs[:weights] = kwargs[:weights][to_remove]
        end
    end

    # Additional checks
    if haskey(kwargs, :weights)
        @assert length(kwargs[:weights]) == length(x)
    end

    # Add default keyword arguments
    default_slm_kwargs!(kwargs, T)

    # Scale y. This updates the kwargs
    ŷ = scale_problem!(x, y, kwargs)

    slm = if kwargs[:degree] == 0
        error("degree 0 has not been implemented")
    elseif kwargs[:degree] == 1
        error("degree 1 has not been implemented")
    elseif kwargs[:degree] == 3
        SLM_cubic(x, ŷ, kwargs[:y_scale], kwargs[:y_shift]; init = kwargs[:init],
                  nk = kwargs[:knots], C2 = kwargs[:C2], λ = kwargs[:λ], increasing = kwargs[:increasing],
                  decreasing = kwargs[:decreasing], increasing_intervals = kwargs[:increasing_intervals],
                  decreasing_intervals = kwargs[:decreasing_intervals],
                  concave_up = kwargs[:concave_up], concave_down = kwargs[:concave_down],
                  concave_up_intervals = kwargs[:concave_up_intervals],
                  concave_down_intervals = kwargs[:concave_down_intervals],
                  left_value = kwargs[:left_value], right_value = kwargs[:right_value],
                  min_value = kwargs[:min_value], max_value = kwargs[:max_value],
                  min_max_sample_points = kwargs[:min_max_sample_points],
                  use_sparse = kwargs[:use_sparse], use_lls = kwargs[:use_lls])
    else
        error("degree $(kwargs[:degree]) has not been implemented")
    end

    if calculate_stats
        # Need to write a function that calculates stats like total df or R²
    end

    if kwargs[:scaling]
        coef = get_coef(slm)
        coef[:, 1] .-= kwargs[:y_shift]
        coef       ./= kwargs[:y_scale]
    end

    if verbose == :high
        @info "Model Statistics Report"
        println("Number of data points:      $(length(y))")
        println("Scale factor applied to y   $(y_scale)")
        println("Shift applied to y          $(y_shift)")
        # println("Total degrees of freedom:   $(get_stats(slm)[:total_df])")
        # println("Net degrees of freedom:     $(get_stats(slm)[:net_df])")
        # println("R-squared:                  $(get_stats(slm)[:R2])")
        # println("Adjusted R-squared:         $(get_stats(slm)[:R2_adj])")
        # println("RMSE:                       $(get_stats(slm)[:RMSE])")
        # println("Range of prediction errors: $(get_stats(slm)[:error_range])")
        # println("Error quartiles (25%, 75%): $(get_stats(slm)[:quartiles])")
    end

    return slm
end

function SLM_cubic(x::AbstractVector{T}, y::AbstractVector{T}, y_scale::T, y_shift::T;
                   init::AbstractVector{T} = Vector{T}(undef, 0),
                   nk::Int = 6, C2::Bool = true, λ::T = 1e-4, increasing::Bool = false, decreasing::Bool = false,
                   increasing_intervals::AbstractMatrix{T} = Matrix{T}(undef, 0, 0),
                   decreasing_intervals::AbstractMatrix{T} = Matrix{T}(undef, 0, 0),
                   concave_up::Bool = false, concave_down::Bool = false,
                   concave_up_intervals::AbstractMatrix{T} = Matrix{T}(undef, 0, 0),
                   concave_down_intervals::AbstractMatrix{T} = Matrix{T}(undef, 0, 0),
                   left_value::T = NaN, right_value::T = NaN,
                   min_value::T = NaN, max_value::T = NaN,
                   min_max_sample_points::AbstractVector{T} = [.017037, .066987, .1465, .25, .37059,
                                                               .5, .62941, .75, .85355, .93301, .98296],
                   use_sparse::Bool = false, use_lls::Bool = true) where {T <: Real}

    nₓ = length(x)

    # Choose knots
    knots = choose_knots(nk, x)
    dknots = diff(knots)
    if any(dknots .== 0.)
        error("Knots must be distinct.")
    end

    ### Calculate coefficients

    ## Set up
    nc = 2 * nk
    Mineq = zeros(0, nc)
    rhsineq = Vector{T}(undef, 0)
    Meq = zeros(0, nc)
    rhseq = Vector{T}(undef, 0)

    ## Build design matrix

    # Bin data so that xbin has nₓ and xbin specifies into which bin each x value falls
    xbin = bin_sort(x, knots)

    # design matrix
    Mdes = construct_design_matrix(x, knots, dknots, xbin, nₓ, nk, nc; use_sparse = use_sparse)
    rhs = y

    ## Regularizer
    Mreg = construct_regularizer(dknots, nk; use_sparse = use_sparse)
    rhsreg = zeros(T, nk)

    ## C2 continuity across knots
    if C2
        Meq, rhseq = C2_matrix(nk, nc, dknots, Meq, rhseq; use_sparse = use_sparse)
    end

    ## Left and right values
    if !isnan(left_value)
        Meq, rhseq = set_left_value(left_value, nc, Meq, rhseq; use_sparse = use_sparse)
    end

    if !isnan(right_value)
        Meq, rhseq = set_right_value(right_value, nc, nk, Meq, rhseq; use_sparse = use_sparse)
    end

    # Global minimum and maximum values
    if !isnan(min_value)
        Mineq, rhsineq = set_min_value(min_value, nk, nc, dknots, Mineq, rhsineq; sample_points = min_max_sample_points,
                                       use_sparse = use_sparse)
    end

    if !isnan(max_value)
        Mineq, rhsineq = set_max_value(max_value, nk, nc, dknots, Mineq, rhsineq; sample_points = min_max_sample_points,
                                       use_sparse = use_sparse)
    end

    ## Monotonicity restrictions
    @assert !(increasing && decreasing) "Only one of increasing and decreasing can be true"

    monotone_settings = Vector{NamedTuple{(:knotlist, :increasing), Tuple{Tuple{Int, Int}, Bool}}}(undef, 0)
    total_monotone_intervals = 0

    # Monotone increasing
    if increasing
        @assert isempty(increasing_intervals) && isempty(decreasing_intervals) "The spline cannot be monotone increasing and " *
            "have nonempty entries for the keywords increasing_intervals and/or decreasing_intervals"
        total_monotone_intervals += monotone_increasing!(monotone_settings, nk)
    end

    # Increasing intervals: nx2 array where each row is an interval over which we have increasing
    if !isempty(increasing_intervals)
        total_monotone_intervals += increasing_intervals_info!(monotone_settings, knots, increasing_intervals, nk)
    end

    # Monotone decreasing
    if decreasing
        @assert isempty(increasing_intervals) && isempty(decreasing_intervals) "The spline cannot be monotone decreasing and " *
            "have nonempty entries for the keywords increasing_intervals and/or decreasing_intervals"
        total_monotone_intervals += monotone_decreasing!(monotone_settings, nk)
    end

    # Decreasing intervals
    if !isempty(decreasing_intervals)
        total_monotone_intervals += decreasing_intervals_info!(monotone_settings, knots, decreasing_intervals, nk)
    end

    # Add inequalities enforcing monotonicity
    if !isempty(monotone_settings)
        Mineq, rhsineq = construct_monotonicity_matrix(monotone_settings, nc, nk, dknots, total_monotone_intervals, Mineq, rhsineq;
                                                       use_sparse = use_sparse)
    end

    ## Concavity
    @assert !(concave_up && concave_down) "Only one of concave_up and concave_down can be true"

    curvature_settings = Vector{NamedTuple{(:concave_up, :range), Tuple{Bool, Vector{T}}}}(undef, 0)

    if concave_up
        @assert isempty(concave_up_intervals) && isempty(concave_down_intervals) "The spline cannot be concave up and " *
            "have nonempty entries for the keywords concave_up_intervals and/or concave_down_intervals"
        concave_up_info!(curvature_settings)
    end

    if concave_down
        @assert isempty(concave_up_intervals) && isempty(concave_down_intervals) "The spline cannot be concave down and " *
            "have nonempty entries for the keywords concave_up_intervals and/or concave_down_intervals"
        concave_down_info!(curvature_settings)
    end

    if !isempty(concave_up_intervals)
        concave_up_intervals_info!(curvature_settings, concave_up_intervals)
    end

    if !isempty(concave_down_intervals)
        concave_down_intervals_info!(curvature_settings, concave_down_intervals)
    end

    # Add inequalities enforcing curvature
    if !isempty(curvature_settings)
        Mineq, rhsineq = construct_curvature_matrix(curvature_settings, nc, nk, knots, dknots, Mineq, rhsineq; use_sparse = use_sparse)
    end

    # scale equalities for unit absolute row sum
    if !isempty(Meq)
        rs = Diagonal(vec(1. ./ sum(abs.(Meq), dims = 2)))
        Meq = rs * Meq
        rhseq = rs * rhseq
    end

    # scale equalities for unit absolute row sum
    if !isempty(Mineq)
        rs = Diagonal(vec(1. ./ sum(abs.(Mineq), dims = 2)))
        Mineq = rs * Mineq
        rhsineq = rs * rhsineq
    end

    # Currently, we only implement the standard regularizer parameter, while
    # SLM allows matching a specified RMSE and cross-validiation, too.
    coef = solve_slm_system(Mdes, rhs, Mreg, rhsreg, λ,
                            Meq, rhseq, Mineq, rhsineq; init = init,
                            use_lls = use_lls)
    coef = reshape(coef, nk, 2)

    ## Calculate model statistics
    # Currently, we just add degree and knots to the statistics NamedTuple
    stats = (degree = 3, knots = knots,
             y_scale = y_scale, y_shift = y_shift)

    # Unpack coefficients into the result structure
    return SLM{T}(stats, :cubic, x, y, knots, coef, :none)
end


module mpec

    using Ipopt
    using JuMP
    using DataFrames
    using DataFramesMeta

    """
        build_trans(p::Vector{Float64},n::Int)

    Build transition matrix for each of `n` states where each element `p[j]` of `p` is the probability of moving `j-1` states up in the grid defined by `1:n`.
    """
    function build_trans(p::Vector{Float64},n::Int)
        P = spzeros(n,n)
        p = vcat(p,(1-sum(p)))
        N = length(p)
        for i=0:(N-1)
            P =   P + sparse(1:(n-i),(1+i):n,ones(n-i)*p[i+1],n,n)
            P[n-i,n] = 1-sum(p[1:i])
        end
        return P
    end

    type Param
        n :: Int  # num of grid points
        gmax :: Int # upper boudn of mileage grid

        # structural params
        pr :: Vector{Float64}  # transition probs
        RC :: Float64         # replacement cost
        c :: Float64  # cost parameter
        beta :: Float64 # discount factor

        # computation objects
        grid :: LinSpace{Float64}
        P :: SparseMatrixCSC{Float64,Int64}
        # EV0 :: Vector{Float64}   # old value function
        # EV1 :: Vector{Float64}   # new
        pk :: Vector{Float64}  # probability of keeping engine
        cost :: Vector{Float64}   #operating cost

        """
            create default param type
        """
        function Param()
            this = new()
            this.n = 175
            this.gmax = 450
            this.grid = 0.0:(this.n-1)

            this.pr = [0.0937; 0.4475; 0.4459; 0.0127]
            this.RC = 11.7257
            this.c = 2.45569
            this.beta = 0.9999
            this.P = build_trans(this.pr,this.n)
            this.cost = 0.001*this.c*collect(this.grid)
            # this.EV0 = zeros(this.n)
            # this.EV1 = zeros(this.n)
            this.pk = zeros(this.n)
            return this
        end
    end

    function read_busses(p::Param,num_types=4)
        # load data
        d = readtable(joinpath(dirname(@__FILE__),"buses.csv"),header=false)
        # setup data
        rename!(d,[:x1,:x2,:x5,:x7,:x9],[:id,:bus_type,:d1,:x,:dx1])

        # subset to bus type
        d = @where(d,:bus_type .<= num_types)

        # discretize odometer data
        d[:x] = ceil(Int,p.n/(p.gmax*1000) * d[:x])
        d[:dx1] = d[:x] .- vcat(0,d[:x][1:(end-1)])
        d[:dx1] = d[:dx1].*(1-d[:d1]) .+d[:d1].*d[:x]  # replace first diff of x by x if replaced

        # get replacement dummy
        d[:d] = vcat(d[:d1][2:end],0)

        # remove obs with missing lagged mileage
        # this is the first row for each obs
        remove = d[:id] .- vcat(0,d[:id][1:(end-1)])
        d = d[remove.==0,:]

        d = @select(d,:id,:bus_type,:d,:x,:dx1)

        return d
    end

end


using ForwardDiff
using LinearAlgebra
include("poisson_linear_system.jl")
include("gradient_descent.jl")
include("learned_jacobi.jl")

function optimize_parameters(
    numelements,
    wavenumber,
    guessparams;
    lr = 1.0,
    gditers = 20,
)
    numiter = numelements^2

    xL, xR = -1.0, 1.0
    bcleft, bcright = 0.0, 0.0

    stepsize = (xR - xL) / numelements
    numpts = numelements + 1
    xrange = range(xL, stop = xR, length = numpts)

    rhs = cos.(wavenumber * xrange)

    op = -poisson_linear_system(numpts)
    exactsol = solve_direct(op, rhs, bcleft, bcright, stepsize)

    func(x) = LearnedJacobi.error_for_parameters(
        bcleft,
        bcright,
        rhs,
        stepsize,
        numiter,
        exactsol,
        x,
    )
    grad(x) = ForwardDiff.gradient(func, x)

    optimumparams = run_gradient_descent_iterations(func, grad, p0, lr, gditers)

    itersol = LearnedJacobi.run_parametric_iterations(
        bcleft,
        bcright,
        rhs,
        stepsize,
        numiter,
        optimumparams,
    )

    return optimumparams, itersol, exactsol
end

numelements = 32
# jacobiparams = [0.5, 0.0, 0.5, 0.5]
p0 = [0.0, 0.0, 0.0, 0.0]
wavenumber = pi / 2

optimumparams, itersol, exactsol = optimize_parameters(numelements,wavenumber,p0)


#=
testODE:
- Julia version:
- Author: ymocquar
- Date: 2019-11-13
=#
include("../src/interface.jl")
include("../src/henon_heiles.jl")
using LinearAlgebra
using Plots
using Plots.PlotMeasures
function fctmain(n_tau, prec)
    setprecision(prec)
    u0 = BigFloat.([90, -44, 83, 13] // 100)
    tab_eps = zeros(BigFloat, 8)
    epsilon = big"0.19"
    for i = 1:size(tab_eps, 1)
        tab_eps[i] = epsilon
        epsilon /= 10
    end
    nbmaxtest = 12
    order = 6
    t_max = big"1.0"
    y = ones(Float64, nbmaxtest, size(tab_eps, 1))
    x = zeros(Float64, nbmaxtest)
    ind = 1
    A = [0 0 1 0; 0 0 0 0; -1 0 0 0; 0 0 0 0]
    for epsilon in tab_eps
        prob = HiOscODEProblem(henon_heiles, u0, (big"0.0", t_max), missing, A, epsilon)
        tabsol = Array{Array{BigFloat,1},1}(undef, 0)
        eps_v = convert(Float32, epsilon)
        nb = 60 * 2^nbmaxtest
        res = solve(prob, nb_tau = n_tau, order = order, nb_t = nb, dense = false)
        par_u0 = res.par_u0
        solRef = res[end]
        nb = 100
        indc = 1
        labels = Array{String,2}(undef, 1, ind)
        while indc <= nbmaxtest
            res = solve(
                prob,
                nb_tau = n_tau,
                order = order,
                nb_t = nb,
                par_u0 = par_u0,
                dense = false,
            )
            sol = res[end]
            nm = norm(sol - solRef, Inf)
            y[indc, ind] = (nm < 1) ? nm : NaN
            x[indc] = 1.0 / nb
            nb *= 2
            indc += 1
        end
        for i = 1:ind
            labels[1, i] = "  ε = $(convert(Float32,tab_eps[i])) "
        end
        p = Plots.plot(
            x,
            view(y, :, 1:ind),
            xlabel = "Δt",
            xaxis = :log,
            ylabel = "error",
            yaxis = :log,
            legend = :bottomright,
            label = labels,
            marker = 2,
        )
        pp = Plots.plot(
            x,
            view(y, :, 1:ind),
            xlabel = "Δt",
            xaxis = :log,
            ylabel = "error",
            yaxis = :log,
            legend = :bottomright,
            label = labels,
            marker = 2,
            bottom_margin = 30px,
            left_margin = 60px,
        )
        prec_v = precision(BigFloat)
        Plots.savefig(p, "out/r2p_$(prec_v)_$(eps_v)_$(order)_$(n_tau)_epsilon_hh_2.pdf")
        Plots.savefig(p, "out/r2p_$(prec_v)_$(eps_v)_$(order)_$(n_tau)_epsilon_hh_2.png")
        Plots.savefig(pp, "out/r2pp_$(prec_v)_$(eps_v)_$(order)_$(n_tau)_epsilon_hh_2.pdf")
        Plots.savefig(pp, "out/r2pp_$(prec_v)_$(eps_v)_$(order)_$(n_tau)_epsilon_hh_2.png")
        ind += 1
    end
end

fctmain(32, 256)


include("ChebyshevT.jl")
include("ChebyshevU.jl")
include("diff.jl")
include("orth.jl")

using DomainSets: WrappedMap

const CosineMap{T} = WrappedMap{T,typeof(cos)}
const SineMap{T} = WrappedMap{T,typeof(sin)}
const ArcCosineMap{T} = WrappedMap{T,typeof(acos)}
const ArcSineMap{T} = WrappedMap{T,typeof(asin)}

cosinemap(::Type{T} = Float64) where {T} = convert(Map{T}, cos)
sinemap(::Type{T} = Float64) where {T} = convert(Map{T}, sin)
arccosinemap(::Type{T} = Float64) where {T} = convert(Map{T}, acos)
arcsinemap(::Type{T} = Float64) where {T} = convert(Map{T}, asin)

inv(::CosineMap{T}) where {T} = arccosinemap(T)
inv(::ArcCosineMap{T}) where {T} = cosinemap(T)
inv(::SineMap{T}) where {T} = arcsinemap(T)
inv(::ArcSineMap{T}) where {T} = sinemap(T)

jacobian(::CosineMap, x) = -sin(x)
jacobian(::SineMap{T}) where {T} = cosinemap(T)


"""Functions for constructing parameter profiles"""
constantParameter(offset::Real=0.0) = x -> offset
export constantParameter
constant = constantParameter
export constant

function heaviside(x::Real, stepOpt::Real=1.0)
    if x < 0
        y = 0
    elseif x > 0
        y = 1
    elseif x == 0
        y = stepOpt
    else
        return nothing
    end
    y = convert(typeof(x), y)
end
heaviside(x::AbstractArray, stepOpt::Real=1.0) = heaviside.(x, stepOpt)

sigmoid(x::Real, width=1) = 1/(1+exp(-x/width))

"""
    unitStep(threshold::Real=0.0, baseline::Real=0.0, stepHeight::Real=1.0, stepOpt::Real=1.0)

Construct a constant function of 'x' that undergoes a step of size 'stepHeight' from a 'baseline' at 'threshold'. 'stepOpt' sets the value at 0. A Discontinuous type is returned, which can be called like a normal function but also contains the x-value for which there is a discontinuity.

# Examples
```julia-repl
julia> D = unitStep();
D([-1, 0, 1])
```
"""
function unitStep(threshold::Real=0.0, baseline::Real=0.0, stepHeight::Real=1.0, stepOpt::Real=1.0)
    Discontinuous(x -> (heaviside(x-threshold, stepOpt).*stepHeight)+baseline, Set([threshold]))
end
unitStep(T::Tuple, args...) = unitStep(T[1], args...) # For compatibility with other parameter profile syntax. If T is tuple, use T[1] as the threshold
export unitStep

# * A unit step that is smooth as butter
function unitsigmoid(threshold::Real=0.0, baseline::Real=0.0, stepHeight::Real=1.0, width::Real=0.1)
    Discontinuous(x -> (sigmoid(x-threshold, width).*stepHeight)+baseline, Set([threshold]))
end
unitsigmoid(T::Tuple, args...) = unitsigmoid(T[1], args...) # For compatibility with other parameter profile syntax. If T is tuple, use T[1] as the threshold
export unitsigmoid

# Same as unitStep, but come back down
function unitBump(T::Tuple, baseline::Real=0.0, bumpHeight::Real=1.0, stepOpt::Real=1.0)
    D = unitStep(T[1], baseline, bumpHeight, stepOpt) - unitStep(T[2], 0.0, bumpHeight, stepOpt)
end
export unitBump

function compositeStep(stepThrs::Vector, baselines::Vector, stepHeights::Vector)
    D = map(unitStep, stepThrs, baselines, stepHeights)
    D = sum(D)
end
compositeStep(T::Tuple, args...) = compositeStep(T[1], args...)
export compositeStep

"""
    stepNoise()

Compose a discontinuous function of many square bumps and troughs, following a white noise process.
    'stepHeight':   the standard deviation of the step height
    'T':            a tuple, (t0, tmax)

# Examples
```julia-repl
julia> D = stepNoise();
D([-1, 0, 1])
```
"""
function stepNoise(T::Tuple, stepWidth::Real=100, stepHeight::Real=1, baseline::Real=0)
    stepIdxs = T[1]:stepWidth:T[2]-stepWidth
    steps = stepHeight.*randn(Float64, (1, length(stepIdxs)))
    ps = sum(map((x, y) -> unitBump((x, x+stepWidth), baseline, y), stepIdxs, steps))
end
export stepNoise

function stepRandomWalk(T::Tuple, stepWidth::Real=100, stepHeight::Real=1, baseline::Real=0)
    stepIdxs = T[1]:stepWidth:T[2]
    steps = stepHeight.*randn(Float64, (1, length(stepIdxs)))
    ps = sum(map((x, y) -> unitStep((x, x+stepWidth), baseline, y), stepIdxs, steps))
end
export stepRandomWalk

function ramp(gradient::Real=1, p0::Real=0, t0::Real=0)
    x -> gradient.*(x.-t0) .+ p0
end
function ramp(p1::Real, p2::Real, t1::Real, t2::Real)
    gradient = (p2 - p1)/(t2 - t1)
    ramp(gradient, p1, t1)
end
export ramp

"""
Define a line over an interval. Saturate before and after this interval
"""
function rampInterval(p1::Real, p2::Real, t1::Real, t2::Real)
    r = ramp(p1, p2, t1, t2)
    b = unitBump((t1, t2), 0, 1)
    s1 = unitStep(t1, p1, -p1)
    s2 = unitStep(t2, 0, p2)
    return r*b + s1 + s2
end
export rampInterval

"""
Define a line over an interval. Before this interval, saturate, but after, extrapolate.
"""
function rampOn(p1::Real, p2::Real, t1::Real, t2::Real)
    r = ramp(p1, p2, t1, t2)
    b = unitStep(t1, 0, 1)
    s1 = unitStep(t1, p1, -p1)
    return r*b + s1
end
export rampOn

"""
Define a line over an interval. After this interval, saturate, but before, extrapolate.
"""
function rampOff(p1::Real, p2::Real, t1::Real, t2::Real)
    r = ramp(p1, p2, t1, t2)
    b = unitStep(t2, 1, -1)
    s2 = unitStep(t2, 0, p2)
    return r*b + s2
end
export rampOff

function sineWave(period::Real=1, amplitude::Real=1, t0::Real=0, baseline::Real=5*amplitude)
    x -> amplitude.*sin.((2π/period).*(x.-t0)) + baseline
end
export sineWave

function triangleWave(period::Real=1, amplitude::Real=1, t0::Real=0, tmax=100, baseline::Real=5*amplitude)
    # This is a discontinuous wave, so we can't (yet) define it over all x
    # Be careful to set the t0 and tmax to match the simulation time
    f = x -> (2amplitude/π).*asin.(sin.((2π/period).*(x.-t0))) + baseline
    d = Set((t0 + π/2):π:tmax) # The set of maxima and minima of the triangle wave
    Discontinuous(f, d)
end
export triangleWave

function lorentzian(A=1.0, Γ=1.0, x₀=0.0, y₀=0.0)
    # Γ is F.W.H.M, x₀ is centre, A is height and y₀ is the baseline
    # (function is not mass normalised, as would be usual)
    γ = Γ/2
    x -> (A.*γ.^2)./((x.-x₀).^2 .+ γ.^2) + y₀
end
export lorentzian


doc"""
    Weibull(α,θ)

The *Weibull distribution* with shape `α` and scale `θ` has probability density function

$f(x; \alpha, \theta) = \frac{\alpha}{\theta} \left( \frac{x}{\theta} \right)^{\alpha-1} e^{-(x/\theta)^\alpha},
    \quad x \ge 0$

```julia
Weibull()        # Weibull distribution with unit shape and unit scale, i.e. Weibull(1.0, 1.0)
Weibull(a)       # Weibull distribution with shape a and unit scale, i.e. Weibull(a, 1.0)
Weibull(a, b)    # Weibull distribution with shape a and scale b

params(d)        # Get the parameters, i.e. (a, b)
shape(d)         # Get the shape parameter, i.e. a
scale(d)         # Get the scale parameter, i.e. b
```

External links

* [Weibull distribution on Wikipedia](http://en.wikipedia.org/wiki/Weibull_distribution)

"""
immutable Weibull <: ContinuousUnivariateDistribution
    α::Float64   # shape
    θ::Float64   # scale

    function Weibull(α::Real, θ::Real)
    	@check_args(Weibull, α > zero(α) && θ > zero(θ))
    	new(α, θ)
    end
    Weibull(α::Real) = Weibull(α, 1.0)
    Weibull() = new(1.0, 1.0)
end

@distr_support Weibull 0.0 Inf

#### Parameters

shape(d::Weibull) = d.α
scale(d::Weibull) = d.θ

params(d::Weibull) = (d.α, d.θ)

#### Statistics

mean(d::Weibull) = d.θ * gamma(1.0 + 1.0 / d.α)
median(d::Weibull) = d.θ * logtwo ^ (1.0 / d.α)
mode(d::Weibull) = d.α > 1.0 ? (iα = 1.0 / d.α; d.θ * (1.0 - iα) ^ iα) : 0.0

var(d::Weibull) = d.θ^2 * gamma(1.0 + 2.0 / d.α) - mean(d)^2

function skewness(d::Weibull)
    μ = mean(d)
    σ2 = var(d)
    σ = sqrt(σ2)
    r = μ / σ
    gamma(1.0 + 3.0 / d.α) * (d.θ / σ)^3 - 3.0 * r - r^3
end

function kurtosis(d::Weibull)
    α, θ = params(d)
    μ = mean(d)
    σ = std(d)
    γ = skewness(d)
    r = μ / σ
    r2 = r^2
    r4 = r2^2
    (θ / σ)^4 * gamma(1.0 + 4.0 / α) - 4.0 * γ * r - 6.0 * r2 - r4 - 3.0
end

function entropy(d::Weibull)
    α, θ = params(d)
    0.5772156649015328606 * (1.0 - 1.0 / α) + log(θ / α) + 1.0
end

#### Evaluation

function pdf(d::Weibull, x::Float64)
    if x >= 0.0
        α, θ = params(d)
        z = x / θ
        (α / θ) * z^(α - 1.0) * exp(-z^α)
    else
        0.0
    end
end

function logpdf(d::Weibull, x::Float64)
    if x >= 0.0
        α, θ = params(d)
        z = x / θ
        log(α / θ) + (α - 1.0) * log(z) - z^α
    else
        -Inf
    end
end

zv(d::Weibull, x::Float64) = (x / d.θ) ^ d.α
xv(d::Weibull, z::Float64) = d.θ * z ^ (1.0 / d.α)

cdf(d::Weibull, x::Float64) = x > 0.0 ? -expm1(-zv(d, x)) : 0.0
ccdf(d::Weibull, x::Float64) = x > 0.0 ? exp(-zv(d, x)) : 1.0
logcdf(d::Weibull, x::Float64) = x > 0.0 ? log1mexp(-zv(d, x)) : -Inf
logccdf(d::Weibull, x::Float64) = x > 0.0 ? -zv(d, x) : 0.0

quantile(d::Weibull, p::Float64) = xv(d, -log1p(-p))
cquantile(d::Weibull, p::Float64) = xv(d, -log(p))
invlogcdf(d::Weibull, lp::Float64) = xv(d, -log1mexp(lp))
invlogccdf(d::Weibull, lp::Float64) = xv(d, -lp)

function gradlogpdf(d::Weibull, x::Float64)
    if insupport(Weibull, x)
        α, θ = params(d)
        (α - 1.0) / x - α * x^(α - 1.0) / (θ^α)
    else
        0.0
    end
end

#### Sampling

rand(d::Weibull) = xv(d, randexp())


"""
Estimation of the covariance matrix in the standard simple way given returns.
https://en.wikipedia.org/wiki/Sample_mean_and_covariance
"""
function simple_covariance_given_returns(returns::Array{R,2}) where R<:Real
    N = size(returns)[2]
    mat = zeros(N,N)
    for i in 1:N
        for j in i:N
            if i == j
                mat[i,i] = var(returns[:,i])
            else
                mat[i,j] = cov(returns[:,i], returns[:,j])
            end
        end
    end
    return Hermitian(mat)
end

"""
Estimation of the covariance matrix in the standard simple way given a time grid.
"""
function simple_covariance_given_time_grid(ts::SortedDataFrame, assets::Vector{Symbol}, time_grid::Vector; regularisation::Union{Missing,Symbol} = :covariance_default,
                                           regularisation_params::Dict = Dict(), only_regulise_if_not_PSD::Bool = false)
    dd_compiled = latest_value(ts, time_grid; assets = assets)
    dd = get_returns(dd_compiled; rescale_for_duration = false)

    if nrow(dd) < 1 return make_nan_covariance_matrix(assets, ts.time_period_per_unit) end

    returns = Matrix(dd[:, assets])
    covariance = simple_covariance_given_returns(returns)

    # Regularisation
    dont_regulise = ismissing(regularisation) || (only_regulise_if_not_PSD && is_psd_matrix(covariance))
    covariance = dont_regulise ? covariance : regularise(covariance, ts, assets, regularisation; regularisation_params... )

    # av_spacing
    N = length(time_grid)
    spacing = safe_multiply_period(Statistics.mean(time_grid[2:N] .- time_grid[1:(N-1)]), ts.time_period_per_unit)

    # Packing into a CovarianceMatrix and returning.
    cor, vols = cov2cor_and_vol(covariance, spacing, ts.time_period_per_unit)
    return CovarianceMatrix(cor, vols, assets, ts.time_period_per_unit)
end

"""
    simple_covariance(ts::SortedDataFrame, assets::Vector{Symbol} = get_assets(ts);
                      regularisation::Union{Missing,Symbol} = :covariance_default,
                      regularisation_params::Dict = Dict(), only_regulise_if_not_PSD::Bool = false,
                      time_grid::Union{Missing,Vector} = missing, fixed_spacing::Union{Missing,<:Real} = missing,
                      refresh_times::Bool = false, rough_guess_number_of_intervals::Integer = 5)

Estimation of the covariance matrix in the standard textbook way.
### Inputs
* `ts` - The tick data.
* `assets` - The assets you want to estimate volatilities for.
* `regularisation` - A symbol representing what regularisation technique should be used. If missing no regularisation is performed.
* `regularisation_params` - keyword arguments to be consumed in the regularisation algorithm.
* `only_regulise_if_not_PSD` - Should regularisation only be attempted if the matrix is not psd already.
* `time_grid` - The grid with which to calculate returns.
* `fixed_spacing` - A spacing used to calculate a time grid. Not used if `refresh_times=true`.
* `refresh_times` - Should refresh times be used to estimate covariance.
* `rough_guess_number_of_intervals` - A rough number of intervals to calculate a default spacing. Not used if a `time_grid` or `fixed_spacing` is provided or if `refresh_times=true`.
### Returns
* A `CovarianceMatrix`.

### References
"""
function simple_covariance(ts::SortedDataFrame, assets::Vector{Symbol} = get_assets(ts);
                           regularisation::Union{Missing,Symbol} = :covariance_default,
                           regularisation_params::Dict = Dict(), only_regulise_if_not_PSD::Bool = false,
                           time_grid::Union{Missing,Vector} = missing, fixed_spacing::Union{Missing,<:Real} = missing,
                           refresh_times::Bool = false, rough_guess_number_of_intervals::Integer = 5)
   if ismissing(time_grid)
       time_grid = Vector{eltype(ts.df[:,ts.time])}()
       if refresh_times
           time_grid = get_all_refresh_times(ts, assets)
       elseif !ismissing(fixed_spacing)
            time_grid = collect(minimum(ts.df[:,ts.time]):fixed_spacing:maximum(ts.df[:,ts.time]))
       else
           n_grid = default_spacing(ts; rough_guess_number_of_intervals = rough_guess_number_of_intervals)
           vals = collect(values(n_grid))
           spacing = min( mean(vals[(isnan.(vals) .== false) .& (isinf.(vals) .== false)]) , duration(ts; in_dates_period = false)/20   )
           spacing = isnan(spacing) ? duration(ts; in_dates_period = false)/20 : spacing
           time_grid = collect(minimum(ts.df[:,ts.time]):spacing:maximum(ts.df[:,ts.time]))
       end
   end
   return simple_covariance_given_time_grid(ts, assets, time_grid; regularisation = regularisation, regularisation_params = regularisation_params, only_regulise_if_not_PSD = only_regulise_if_not_PSD)
end


using BenchmarkTools, MicroscopyTools

function bin_3(arr::AbstractArray{T, N}, factors=ntuple(x -> 2, Val(N))) where {T, N}
    # old size
    sz = size(arr)
    # new size
    sn = sz .÷ factors
    s2 = ntuple(i -> isodd(i) ? factors[(i+1)/2] : sn[i/2], Val(2 * N))

    dims = ntuple(i -> 2 * i - 1, Val(N))
    reshape(sum(reshape(arr,s2),dims=dims),sn)
end

function f(x)
    @info "Random data with size $(size(x))"
    print("Tullio: ")
    @btime bin2($x)
    print("CartesianIndices: ")
    @btime bin($x, (2,2))
    print("RH: ")
    @btime bin_3($x, (2,2))
end

function main()
    @info "Binning by (2,2)"
    x = randn(Float32, (256, 512))
    f(x)

    x = randn(Float32, (8192, 7113))
    f(x)

    x = randn((97, 32))
    f(x)

    x = randn((3097, 113, 32))
    f(x)

    x = randn((23, 31, 8))
    f(x)

    x = randn(Float32, (3097, 4096, 32))
    f(x)

    return
end

main()


using MosekTools, PyPlot, LaTeXStrings, MarkovBounds
# use Catalyst to define the reaction network under consideration
# in this case, we investigate Schlögel's system
rn = @reaction_network begin
    0.01*2, 2X --> X
    1.0, ∅ --> X
end
X = species(rn)[1]
# specify initial molecular counts in system
x0 = Dict(X => 1)
x_scale =  Dict(X => 10.0)

# compute bounds on mean molecular counts of the product counts at steady state
# for different truncation orders 2:5
orders = 2:2:6
sol = Dict()
for m in orders
    sol[m] = stationary_mean(rn, [X], x0, m, Mosek.Optimizer, auto_scaling = false, scales = x_scale)
end

fig, ax = subplots()
ax.plot(orders, [sol[m].bounds[X][1] for m in orders], label="lower bound", color="blue")
ax.plot(orders, [sol[m].bounds[X][2] for m in orders], label="upper bound", color="red")
ax.legend()
ax.set_xlabel("truncation order")
ax.set_ylabel(L"⟨ X_{\infty} ⟩")
display(fig)

# same for variance
sol = Dict()
for m in orders
    sol[m] = stationary_variance(rn, [X], x0, m, Mosek.Optimizer, auto_scaling = false, scales = x_scale)
end
fig, ax = subplots()
ax.plot(orders, [sol[m].bounds[X] for m in orders], label="upper bound", color="red")
ax.legend()
ax.set_xlabel("truncation order")
ax.set_ylabel(L"V(X_{\infty})")
display(fig)

# transient moments
nT = 10
traj = []
Tf_range = [0.0, 0.5, 0.75, 1.0, 1.5, 2.0, 2.5, 5.0, 7.5,
            10.0, 12.5, 15.0, 20.0, 25.0, 30.0, 40.0, 50.0]
order = 4
for Tf in Tf_range
    trange = range(Tf/nT, stop=Tf, length=nT)
    sol = transient_mean(rn, [X], x0, order, trange, Mosek.Optimizer, scales = x_scale)
    push!(traj, sol)
end
fig, ax = subplots()
ax.plot(Tf_range, [sol.bounds[X][1] for sol in traj], color="red", label="lower bound")
ax.plot(Tf_range, [sol.bounds[X][2] for sol in traj], color="blue", label="upper bound")
ax.legend()
ax.set_xlabel("time")
ax.set_ylabel(L"⟨ X ⟩")
display(fig)

# variance
traj = []
for Tf in Tf_range
    trange = range(Tf/nT, stop=Tf, length=nT)
    sol = transient_variance(rn, [X], x0, order, trange, Mosek.Optimizer, scales =  x_scale)
    push!(traj, sol)
end
fig, ax = subplots()
ax.plot(Tf_range, [sol.bounds[X][1] for sol in traj], color="red", label="lower bound")
ax.legend()
ax.set_xlabel("time")
ax.set_ylabel(L"Var(X)")
display(fig)


### A Pluto.jl notebook ###
# v0.12.16

using Markdown
using InteractiveUtils

# ╔═╡ e077ab16-22e5-11eb-15ed-5f0654bf7328
using Pkg, DrWatson

# ╔═╡ e75a9a06-22e5-11eb-2e36-8d4f62b830ed
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Random, Distributions, Statistics, Plots
	Random.seed!(2)
end

# ╔═╡ ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
md"## Listing5.14"

# ╔═╡ fa68607e-22e5-11eb-0558-c9a4d9f77426
begin
	n, N, alpha = 10, 10^7, 0.05
	mActual = 0.75
	dist0, dist1 = Uniform(0,1), Uniform(0,mActual)

	ts(sample) = maximum(sample) - minimum(sample)
end

# ╔═╡ d8d55f04-3af8-11eb-2251-75efd46ad5a2
begin
	empiricalDistUnderH0 = [ts(rand(dist0,n)) for _ in 1:N]
	rejectionValue = quantile(empiricalDistUnderH0,alpha)
end;

# ╔═╡ d8d5a0b8-3af8-11eb-1cc3-4db334734491
begin
	sample = rand(dist1,n)
	testStat = ts(sample)
	pValue = sum(empiricalDistUnderH0 .<= testStat)/N

	if testStat > rejectionValue
		print("Didn't reject: ", round(testStat,digits=4))
		print(" > ", round(rejectionValue,digits=4))
	else
		print("Reject: ", round(testStat,digits=4))
		print(" <= ", round(rejectionValue,digits=4))
	end
	Text("p-value = $(round(pValue,digits=4))")
end

# ╔═╡ d8d6283a-3af8-11eb-37cd-5fa24b45e40f
begin
	stephist(empiricalDistUnderH0, bins=100, c=:blue, normed=true, label="")
	plot!([testStat, testStat], [0,4], c=:red, label="Observed test statistic")
	plot!([rejectionValue, rejectionValue], [0,4], c=:black, ls=:dash,
		label="Critical value boundary", legend=:topleft, ylims=(0,4),
			xlabel = "x", ylabel = "Density")
end

# ╔═╡ 475ff888-22e6-11eb-2354-09f8f40a8e12
md"## End of listing5.14"

# ╔═╡ Cell order:
# ╟─ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
# ╠═e077ab16-22e5-11eb-15ed-5f0654bf7328
# ╠═e75a9a06-22e5-11eb-2e36-8d4f62b830ed
# ╠═fa68607e-22e5-11eb-0558-c9a4d9f77426
# ╠═d8d55f04-3af8-11eb-2251-75efd46ad5a2
# ╠═d8d5a0b8-3af8-11eb-1cc3-4db334734491
# ╠═d8d6283a-3af8-11eb-37cd-5fa24b45e40f
# ╟─475ff888-22e6-11eb-2354-09f8f40a8e12


using Pkg
Pkg.activate(joinpath(Pkg.devdir(), "MLCourse"))
using MLCourse, CSV, MLJ, DataFrames, MLJLinearModels, Random, Plots
include("./Data_Processing.jl")

#Load the data
drop = generate(option = "drop", std = "false", valid = "false", test = "true");
med = generate(option = "med", std = "false", valid = "false", test = "true");

#Test if the machine has a greater AUC on the drop data set or the med data set
mach_drop = machine(LogisticClassifier(), select(drop.train[:,:], Not(:precipitation_nextday)), drop.train.precipitation_nextday) |> fit!;
mach_med = machine(LogisticClassifier(), select(med.train[:,:], Not(:precipitation_nextday)), med.train.precipitation_nextday) |> fit!;
losses(mach_drop, select(drop.test[:,:], Not(:precipitation_nextday)), drop.test.precipitation_nextday) # 0.79
losses(mach_med, select(med.test[:,:], Not(:precipitation_nextday)), med.test.precipitation_nextday) # 0.88
# AUC_drop = 0.786, AUC_med = 0.880

#Function to tune the model with L1 and L2 regularization with log scaled lambda values
function TunedModel_L2(data)
	Random.seed!(2711)
	model = LogisticClassifier(penalty = :l2)
	self_tuning_model = TunedModel(model = model,
	                         resampling = CV(nfolds = 10),
	                         tuning = Grid(goal = 100),
	                         range = range(model, :lambda,
									    lower = 1e1, upper = 1e3, scale = :log),
	                         measure = auc)
    self_tuning_mach = machine(self_tuning_model,
							select(data[:,:], Not(:precipitation_nextday)),
							data.precipitation_nextday) |> fit!;
	self_tuning_mach
end

function TunedModel_L1(data)
	Random.seed!(2711)
	model = LogisticClassifier(penalty = :l1)
	self_tuning_model = TunedModel(model = model,
	                         resampling = CV(nfolds = 10),
	                         tuning = Grid(goal = 100),
	                         range = range(model, :lambda,
									    lower = 1e-1, upper = 1e2, scale = :log),
	                         measure = auc)
    self_tuning_mach = machine(self_tuning_model,
							select(data[:,:], Not(:precipitation_nextday)),
							data.precipitation_nextday) |> fit!;
	self_tuning_mach
end

#Load the standardized data
drop_std = generate(option = "drop", std = "true", valid = "false", test = "true");
med_std = generate(option = "med", std = "true", valid = "false", test = "true");

#Test de best model on the drop and med data set with L1 or L2 regularization
rep_drop_L2 = report(TunedModel_L2(drop_std.train));
scatter(reshape(rep_drop_L2.plotting.parameter_values, :), rep_drop_L2.plotting.measurements, xlabel = "Lamba", ylabel = "AUC")
rep_drop_L1 = report(TunedModel_L1(drop_std.train));
scatter(reshape(rep_drop_L1.plotting.parameter_values, :), rep_drop_L1.plotting.measurements, xlabel = "Lamba", ylabel = "AUC")
rep_med_L2 = report(TunedModel_L2(med_std.train));
scatter(reshape(rep_med_L2.plotting.parameter_values, :), rep_med_L2.plotting.measurements, xlabel = "Lamba", ylabel = "AUC")
rep_med_L1 = report(TunedModel_L1(med_std.train));
scatter(reshape(rep_med_L1.plotting.parameter_values, :), rep_med_L1.plotting.measurements, xlabel = "Lamba", ylabel = "AUC")

#Test the best models over the test set with the best hyper-parameters to pick the best one
best_mach_drop_L2 = machine(LogisticClassifier(penalty = :l2, lambda = rep_drop_L2.best_model.lambda), select(drop_std.train[:,:], Not(:precipitation_nextday)), drop_std.train.precipitation_nextday)|> fit!
best_mach_drop_L1 = machine(LogisticClassifier(penalty = :l1, lambda = rep_drop_L1.best_model.lambda), select(drop_std.train[:,:], Not(:precipitation_nextday)), drop_std.train.precipitation_nextday)|> fit!
best_mach_med_L2 = machine(LogisticClassifier(penalty = :l2, lambda = rep_med_L2.best_model.lambda), select(med_std.train[:,:], Not(:precipitation_nextday)), med_std.train.precipitation_nextday)|> fit!
best_mach_med_L1 = machine(LogisticClassifier(penalty = :l1, lambda = rep_med_L1.best_model.lambda), select(med_std.train[:,:], Not(:precipitation_nextday)), med_std.train.precipitation_nextday)|> fit!
losses(best_mach_drop_L2, select(drop_std.test, Not(:precipitation_nextday)), drop_std.test.precipitation_nextday)
losses(best_mach_drop_L1, select(drop_std.test, Not(:precipitation_nextday)), drop_std.test.precipitation_nextday)
losses(best_mach_med_L2, select(med_std.test, Not(:precipitation_nextday)), med_std.test.precipitation_nextday)
losses(best_mach_med_L1, select(med_std.test, Not(:precipitation_nextday)), med_std.test.precipitation_nextday)
# AUC_DROP_L2 = 0.917, AUC_DROP_L1 = 0.914, AUC_MED_L2 = 0.914, AUC_MED_L1 = 0.912

#Write in the submission file with the best machine trained on all the data
train, test = generate(option = "drop", std = "true", valid = "false", test = "false");
best_mach = machine(LogisticClassifier(penalty = :l2, lambda = rep_drop_L2.best_model.lambda), select(train[:,:], Not(:precipitation_nextday)), train.precipitation_nextday)|> fit!;
sample = CSV.read(joinpath(@__DIR__, "data", "sample_submission.csv"), DataFrame);
pred = pdf.(predict(best_mach, test), true);
sample.precipitation_nextday = pred;
CSV.write(joinpath(@__DIR__, "results", "Linear_submission.csv"), sample);


"""
    multinomial_hyper(α::AbstractArray{Float32,1})

[Dirichlet Distribution](https://en.wikipedia.org/wiki/Dirichlet_distribution)
"""
struct multinomial_hyper <: distribution_hyper_params
    α::AbstractArray{Float32,1}
end

mutable struct multinomial_sufficient_statistics <: sufficient_statistics
    N::Float32
    points_sum::AbstractArray{Float32,1}
end

function calc_posterior(
    prior::multinomial_hyper,
    suff_statistics::multinomial_sufficient_statistics,
)
    if suff_statistics.N == 0
        return prior
    end
    return multinomial_hyper(prior.α + suff_statistics.points_sum)
end

function sample_distribution(hyperparams::multinomial_hyper)
    return multinomial_dist(log.(rand(Dirichlet(Float64.(hyperparams.α)))))
end

function create_sufficient_statistics(
    hyper::multinomial_hyper,
    posterior::multinomial_hyper,
    points::AbstractArray{Float32,2},
    pts_to_group = 0,
)
    pts = copy(points)
    points_sum = sum(pts, dims = 2)[:]
    S = pts * pts'
    return multinomial_sufficient_statistics(size(points, 2), points_sum)
end

function log_marginal_likelihood(
    hyper::multinomial_hyper,
    posterior_hyper::multinomial_hyper,
    suff_stats::multinomial_sufficient_statistics,
)
    D = length(suff_stats.points_sum)
    logpi = log(pi)
    val =
        logabsgamma(sum(hyper.α))[1] - logabsgamma(sum(posterior_hyper.α))[1] + sum(
            (x -> logabsgamma(x)[1]).(posterior_hyper.α) -
            (x -> logabsgamma(x)[1]).(hyper.α),
        )
    return val
end

function aggregate_suff_stats(
    suff_l::multinomial_sufficient_statistics,
    suff_r::multinomial_sufficient_statistics,
)
    return multinomial_sufficient_statistics(
        suff_l.N + suff_r.N,
        suff_l.points_sum + suff_r.points_sum,
    )
end


using Atum
using Atum.Advection

using Test
using Printf
using LinearAlgebra: norm
using StaticArrays: SVector

function square(law, x⃗)
  FT = eltype(law)
  ρ = @inbounds abs(x⃗[1]) < FT(1 // 2) ? FT(2) : FT(1)
  SVector(ρ)
end

function run(A, FT, N, K; volume_form=WeakForm())
  Nq = N + 1

  law = AdvectionLaw{FT, 1}()

  cell = LobattoCell{FT, A}(Nq)
  v1d = range(FT(-1), stop=FT(1), length=K+1)
  grid = brickgrid(cell, (v1d,); periodic=(true,))

  dg = DGSEM(; law, grid, volume_form,
               surface_numericalflux = CentralFlux())

  cfl = FT(1 // 4)
  dt = cfl * min_node_distance(grid) / norm(constants(law).u⃗)
  timeend = FT(5.0)

  q = square.(Ref(law), points(grid))
  η0 = entropyintegral(dg, q)

  @info @sprintf """Starting
  N           = %d
  K           = %d
  volume_form = %s
  norm(q)     = %.16e
  η0          = %.16e
  """ N K volume_form weightednorm(dg, q) η0

  odesolver = RLSRK54(dg, q, dt)

  solve!(q, timeend, odesolver)
  ηf = entropyintegral(dg, q)

  Δη = (ηf - η0) / abs(η0)

  @info @sprintf """Finished
  norm(q) = %.16e
  ηf      = %.16e
  Δη      = %.16e
  """ weightednorm(dg, q) ηf Δη

  Δη
end

let
  A = Array
  FT = Float64
  N = 4
  K = 5

  volume_form = FluxDifferencingForm(CentralFlux())
  Δη = run(A, FT, N, K; volume_form)

  @test abs(Δη) <= 10eps(FT)
end


using CompScienceMeshes, BEAST

o, x, y, z = euclidianbasis(3)

# Γ = meshsphere(1.0, 0.11)
Γ = readmesh(joinpath(@__DIR__,"sphere2.in"))
# Γ = readmesh("/Users/Benjamin/Documents/sphere.in")
# Γ = readmesh(joinpath(@__DIR__,"sphere_subd1.in"))
X = lagrangecxd0(Γ)
# X = subdsurface(Γ)
# X = raviartthomas(Γ)
@show numfunctions(X)

κ = 1.0; γ = im*κ
a = Helmholtz3D.singlelayer(wavenumber=κ)
b = Helmholtz3D.doublelayer_transposed(gamma=κ*im) +0.5Identity()

uⁱ = Helmholtz3D.planewave(wavenumber=κ, direction=z)
f = strace(uⁱ,Γ)
g = ∂n(uⁱ)

@hilbertspace u
@hilbertspace v

eq1 = @discretise a[v,u] == f[v] u∈X v∈X
eq2 = @discretise b[v,u] == g[v] u∈X v∈X

x1 = gmres(eq1)
x2 = gmres(eq2)

fcr1, geo1 = facecurrents(x1, X)
fcr2, geo2 = facecurrents(x2, X)

# include(Pkg.dir("CompScienceMeshes","examples","plotlyjs_patches.jl"))
using LinearAlgebra
using Plotly
pt1 = Plotly.plot(patch(geo1, real.(norm.(fcr1))));
pt2 = Plotly.plot(patch(geo2, real.(norm.(fcr2))));
display([pt1 pt2])

# using Test

# ## test the results
# Z = assemble(a,X,X);
# m1, m2 = 1, numfunctions(X)
# chm, chn = chart(Γ,cells(Γ)[m1]), chart(Γ,cells(Γ)[m2])
# ctm, ctn = CompScienceMeshes.center(chm), CompScienceMeshes.center(chn)
# R = norm(cartesian(ctm)-cartesian(ctn))
# G = exp(-im*κ*R)/(4π*R)
# Wmn = volume(chm) * volume(chn) * G
# @show abs(Wmn-Z[m1,m2]) / abs(Z[m1,m2])
# @test abs(Wmn-Z[m1,m2]) / abs(Z[m1,m2]) < 2.0e-3

# r = assemble(f,X)
# m1 = 1
# chm = chart(Γ,cells(Γ)[m1])
# ctm = CompScienceMeshes.center(chm)
# sm = volume(chm) * f(ctm)
# r[m1]
# @show abs(sm - r[m1]) / abs(r[m1])
# @test abs(sm - r[m1]) / abs(r[m1]) < 1e-3


mutable struct Submarine
  depth::Int
  position::Int
  aim::Int
end

function move(self::Submarine, instruction::String)
  action, n = split(strip(instruction))
  n = parse(Int, n)
  if action == "forward"
    self.position += n
  elseif action == "down"
    self.depth += n
  elseif action == "up"
    self.depth -= n
  else
    throw(RuntimeError("Wrong instruction"))
  end
end

function updated_movement(self::Submarine, instruction::String)
  action, n = split(strip(instruction))
  n = parse(Int, n)
  if action == "forward"
    self.position += n
    self.depth += self.aim * n
  elseif action == "down"
    self.aim += n
  elseif action == "up"
    self.aim-= n
  else
    throw(RuntimeError("Wrong instruction"))
  end
end

function codify_position(self::Submarine)
  return self.depth * self.position
end

function run(instructions, mode, init_depth=0, init_pos=0, init_aim=0)
  sub = Submarine(init_depth, init_pos, init_aim)
  for instruction in instructions
    if mode == 1
      move(sub, instruction)
    else
      updated_movement(sub, instruction)
    end
  end
  result = codify_position(sub)
  println("The result (in mode $mode) is $result")
end

test_input = ["forward 5", "down 5", "forward 8", "up 3", "down 8", "forward 2"]
test_output = run(test_input, 1)
test_output = run(test_input, 2)

open("input.txt", "r") do f
  data = [s for s in readlines(f)]
  run(data, 1)
  run(data, 2)
end


# Description: see script_mse_N2 but now we vary N and N2 while keeping
# their ratio fixed s.t. N2 = 1/alpha *N where alpha is the subsampling
# factor.
#
#
# Parameters:
# Nrange: vector of number of measurements
# N2: ambient dimensions
# alpha: subsampling factor
# K: number of sinusoids
# L: number of snapshots
# Niters: number of Monte Carlo iterations
# SNR_db: signal-to-noise-ratio [dB] of the AWGN model
# f_true: location of the frequencies of the sinusoids

include("MUSIC.jl")
include("rootMUSIC.jl")
include("ESPRIT.jl")
include("TLSesprit.jl")
include("MPM.jl")
include("ANM.jl")
include("EMAC.jl")
include("periodogram.jl")
include("findMinWrapDist.jl")
include("findMinWrapDistUnique.jl")
include("WrapDist.jl")

using Plots, Random

Nrange = ceil.(Int,range(6,30,length=5))
L = 1
SNR_db = 30
Niters = 5
alpha = 0.5
f_true = [0.3 0.32 0.34]
K = length(f_true)

MSE = zeros(length(Nrange),3)

A(f,N) = exp.(2*pi*1im*kron(f,(0:N-1)))
SNR = 10 .^(SNR_db./10)

for ii = 1:length(Nrange)
    N = Nrange[ii]
    N2::Int = ceil(1/alpha*N)
    for jj = 1:Niters

        noise = (randn(N,L)+1im*randn(N,L))/ (sqrt(2)*SNR)
        omega = sort(randperm(N2)[1:N])
        S = exp.(2*pi*rand(K,L)*1im)

        Y1 = A(f_true,N)*S + noise
        Y2 = A(f_true,N2)[omega,:]*S + noise

        t0 = time_ns()
        f_MUSIC = MPM(Y1,K)
        t1 = (time_ns()-t0)/1e9
        MinWrapDist = findMinWrapDistUnique(f_true,f_MUSIC)
        MSE[ii,1] += norm(MinWrapDist)^2

        f_EMAC = EMAC(Y2,N2,K,omega,1/SNR)
        MinWrapDist = findMinWrapDistUnique(f_true,f_EMAC)
        MSE[ii,2] += norm(MinWrapDist)^2

        f_ANM = ANM(Y2,N2,K,omega,1/SNR)
        MinWrapDist = findMinWrapDistUnique(f_true,f_ANM)
        MSE[ii,3] += norm(MinWrapDist)^2

    end
end

MSE /= Niters

plot(Nrange, MSE[:,1],yaxis=:log,tickfontsize = 6,
label="Matrix Pencil Method",
#xlims = (0,1),
#xticks = round.(f_range.-f1; digits=2),
xlabel = "N",
ylabel = "MSE"
)
plot!(Nrange,MSE[:,2],label="EMAC")
plot!(Nrange,MSE[:,3],label="ANM")
#savefig("/home/thomas/atom-amd64/usr/share/atom/example/mse_N_N2.png")


include("utils.jl")

""" Nominal facility location model. """
function facility_model(c::Matrix, f::Vector)
    n, m = size(c)
    @assert length(f) == n
    model = Model(Gurobi.Optimizer)
    set_optimizer_attribute(model, "OutputFlag", 0)

    # VARIABLES
    @variable(model, x[1:n], Bin)      # Facility locations
    @variable(model, y[1:n, 1:m] >= 0) # Flow variables
    # CONSTRAINTS
    @constraint(model, [j=1:m], sum(y[i, j] for i = 1:n) >= d[j])      # demand constraint
    @constraint(model, [i=1:n], sum(y[i, j] for j = 1:m) <= s[i]*x[i]) # capacity constraint
    # OBJECTIVE
    @objective(model, Min,
            sum(f[j] * x[j] for j = 1:n) + sum(c[i, j] * y[i, j] for i=1:n, j=1:m))
    return model, x, y
end

""" Robust facility location model, with JuMP using the robust counterpart. """
function robust_facility_model(c::Matrix, f::Vector, rho::Real, Gamm::Real)
    n, m = size(c)
    @assert length(f) == n
    model = Model(Gurobi.Optimizer)
    set_optimizer_attribute(model, "OutputFlag", 0)

    # VARIABLES
    @variable(model, x[1:n], Bin)                 # Facility locations
    @variable(model, u[1:n, 1:m] >= 0)            # Flow variables

    # CONSTRAINTS
    for j=1:m
        infdummy = @variable(model)
        normdummy = @variable(model, [1:m])
        y = @variable(model, [1:m])
        @constraint(model, [i=1:m], normdummy[i] >= y[i])
        @constraint(model, [i=1:m], normdummy[i] >= -y[i])
        @constraint(model, [i=1:m], infdummy >= -y[i] - P[j,i])
        @constraint(model, [i=1:m], infdummy >= y[i] + P[j,i])
        @constraint(model, sum(u[i, j] for i = 1:n) >= d[j] + rho*sum(normdummy) + Gamm * infdummy)
    end
    @constraint(model, [i=1:n], sum(u[i, j] for j = 1:m) <= s[i] * x[i]) # capacity constraint
    # OBJECTIVE
    @objective(model, Min,
            sum(f[j] * x[j] for j = 1:n) + sum(c[i, j] * u[i, j] for i=1:n, j=1:m))
    return model, x, u
end

# Nominal model (3.1)
model, x, y = facility_model(c, f)
optimize!(model)
plt = plot_solution(model, x, y)
@show plt

# Robust model (3.2)
rho = 1
Gamm = 5
model, x, y = robust_facility_model(c, f, rho, Gamm)
optimize!(model)
plt = plot_solution(model, x, y)
@show plt

# Different radii of influence
R_Ds = collect(0:0.05:0.75)
R_Ds[1] = 1e-5

# Unscaled P (3.3)
df = DataFrame("R_D" => [], "sumP" => [], "Gamma" => [], "f+c" => [], "f" => [], "c" => [], "nP" => [], "nx" => [])
for k=1:length(R_Ds)
    global P = [0.2*exp(-1/R_Ds[k] * LinearAlgebra.norm(customers[i, :] .- customers[j, :])[1]) for j=1:m, i=1:m];
    global P = (P .>= 0.2*exp(-1/R_Ds[k] .* R_Ds[k])) .* P
    println("Nonzero Ps: $(sum(P .> 0))")
    model, x, y = robust_facility_model(c, f, rho, Gamm)
    optimize!(model)
    push!(df, Dict("R_D" => R_Ds[k],
                     "sumP" => sum(P),
                     "Gamma" => Gamm,
                     "f+c" => objective_value(model),
                     "f" => value(sum(f[j] * x[j] for j = 1:n)),
                     "c" => value(sum(c[i, j] * y[i, j] for i=1:n, j=1:m)),
                     "nP" => sum(P .> 0),
                     "nx" => sum(value.(x))))
end

# Column-wise scaled P (3.7)
df = DataFrame("R_D" => [], "sumP" => [], "Gamma" => [], "f+c" => [], "f" => [], "c" => [], "nP" => [], "nx" => [])
P_nominal = [0.2*exp(-1/R_D .*LinearAlgebra.norm(customers[i, :] .- customers[j, :])[1]) for j=1:m, i=1:m];
P_nominal = (P_nominal .>= 0.2*exp(-1/R_D .* R_D)) .* P_nominal
for k=1:length(R_Ds)
    global P = [0.2*exp(-1/R_Ds[k] * LinearAlgebra.norm(customers[i, :] .- customers[j, :])[1]) for j=1:m, i=1:m];
    global P = (P .>= 0.2*exp(-1/R_Ds[k] .* R_Ds[k])) .* P
    for i=1:m
        global P[i,:] = P[i,:] .* sum(P_nominal[i,:]) ./ sum(P[i,:])
    end
    println("Nonzero Ps: $(sum(P .> 0))")
    model, x, y = robust_facility_model(c, f, rho, Gamm)
    optimize!(model)
    push!(df, Dict("R_D" => R_Ds[k],
                     "sumP" => sum(P),
                     "Gamma" => Gamm,
                     "f+c" => objective_value(model),
                     "f" => value(sum(f[j] * x[j] for j = 1:n)),
                     "c" => value(sum(c[i, j] * y[i, j] for i=1:n, j=1:m)),
                     "nP" => sum(P .> 0),
                     "nx" => sum(value.(x))))
end


"""
Structure implementing lists of intervals.
These structures are extremely useful when working with stream graphs.
"""
mutable struct Intervals
    list::Array{Tuple{Float64,Float64},1}
end

Intervals()=Intervals([])

"""
Equality between two lists of intervals
NOTE: Use approx here because of rounding errors in intervals.
"""
==(i1::Intervals,i2::Intervals)=clean(i1).list≈clean(i2).list

"""
In operator.
Returns true if the given element is inside one of the intervals of the list.
"""
function ∈(x::Float64, i::Intervals)
    for ii in i.list
        if ii[1] <= x <= ii[2]
            return true
        end
    end
    return false
end

"""Returns the number of intervals in the list."""
count(i::Intervals) = length(i.list)

"""Compute the length of a list of intervals as the sum of the interval lengths."""
length(i::Intervals)=sum([ii[2]-ii[1] for ii in i.list])

"""Merge a new interval within a list of intervals."""
merge(i::Intervals, x::Tuple{Float64,Float64})=union(i, Intervals([x]))

"""
Push a new interval at the end of the list.
WARNING: No sorting is done here...
"""
function push(i::Intervals, x::Tuple{Float64,Float64})
    push!(i.list, x)
    return i
end

"""Intersection between two list of intervals."""
function ∩(i1::Intervals, i2::Intervals)
    i1=clean(i1)
    i2=clean(i2)
    l = []
    components = vcat(i1.list,i2.list)
    if length(components)==0
        return Intervals()
    end
    sort!(components)
    push!(l, popfirst!(components))
    while length(components)>0
        modif = false
        c = popfirst!(components)
        if length(l) == 0
            push!(l,c)
            continue
        end
        if c[2] < l[end][2]
            push!(components,(c[2],l[end][2]))
            l[end] = (l[end][1],c[2])
            modif = true
        end
        if l[end][2] >= c[1] >= l[end][1]
            l[end] = (c[1], l[end][2])
            modif = true
        end
        if !modif
            pop!(l)
        else
            sort!(components)
        end
        push!(l,c)
    end
    pop!(l)
    return clean(Intervals(l))
end

"""Union between two lists of intervals."""
function ∪(i1::Intervals, i2::Intervals)
    l = []
    components = union(i1.list,i2.list)
    sort!(components)
    for c in components
        if length(l) == 0 || c[1] > l[end][2]
            push!(l,c)
        elseif c[2] > l[end][2]
            l[end] = (l[end][1], c[2])
        end
    end
    return clean(Intervals(l))
end

function clean(i::Intervals)
    l = []
    sort!(i.list)
    for c in i.list
        if length(l) == 0
            push!(l,c)
            continue
        end
        if c[1] == l[end][2]
            l[end] = (l[end][1],c[2])
            continue
        end
        push!(l,c)
    end
    return Intervals(l)
end

"""Is subset returns True if i2 is included within i1."""
⊆(i1::Intervals,i2::Intervals)=(i1 ∩ i2)==clean(i1)
⊆(i1::Tuple{Float64,Float64},i2::Intervals)=Intervals([i1]) ⊆ i2


# -----------------------------------------------------------
#
# Implementing a nodal basis for DG sparse grids,
# Multidimensional case
#
# -----------------------------------------------------------

#=

	We continue the work started in 1D_Nodal_Basis.jl by applying the
	tensor product construction together with the sparse cutoff
	to constract a way to translate between the modal basis used in
	linear wave evolution and the nodal/point basis necessary for
	efficiently handling nonlinearities in the time evolution equations

=#

# Check whether (level1, cell1) overlaps with (level2, cell2)
function relevant_cell_1D(k::Int, level1::Int, cell1::Int, level2::Int, cell2::Int)
    left1  = (cell1-1)/(1<<max(0, level1-1))
    right1 = cell1/(1<<max(0, level1-1))
    left2  = (cell2-1)/(1<<max(0, level2-1))
    right2 = cell2/(1<<max(0, level2-1))
    return (left1 <= left2 && right1 >= right2) || (left1 >= left2 && right1 <= right2)
end

# Multidimensional version of the above
function relevant_cell{D}(k::Int, level1::CartesianIndex{D}, cell1::CartesianIndex{D},
								  level2::CartesianIndex{D}, cell2::CartesianIndex{D})
	bool = true
	for d in 1:D
		bool &= relevant_cell_1D(k, level1[d]-1, cell1[d], level2[d]-1, cell2[d])
	end
	return bool
end

# Takes an input (level, cell, mode) and yielding the
# corresponding index in the vector of coefficients for the nodal/point basis
function get_point_index(k::Int, l::Int, c::Int, m::Int)
	if l == 0
		return m
	end
	return k + (k-1) * (1<<(l-1) - 1) + (k-1) * (c-1) + m
end

function get_modal_index(k::Int, l::Int, c::Int, m::Int)
	if l == 0
		return m
	end
	return k * (1<<(l-1) + (c-1) ) + m
end

function inner_loop_to_points{D, T<:Real}(c2::Int, k::Int, l::CartesianIndex{D},
									 level::CartesianIndex{D},
									 cell::CartesianIndex{D},
									 mode::CartesianIndex{D},
									 I::Array{Int, 1},
									 V::Array{T, 1},
									 base_indices::Array{Int, 1},
									 point_mat_1D::Array{T, 2})
	cells::NTuple{D, Int} = ntuple(i -> 1<<max(0, l[i]-2), D)
	modes::NTuple{D, Int} = ntuple(i->l[i]==1?k:k-1,D)
	cellrange = CartesianRange(cells)
	moderange = CartesianRange(modes)
	for c in cellrange
		!relevant_cell(k, l, c, level, cell) && (c2 += prod(modes); continue)

		for m in moderange
			val = one(T)
			for d in 1:D
				index = get_point_index(k, l[d]-1, c[d], m[d])
				val *= point_mat_1D[index, base_indices[d]]
			end
			(abs(val) < eps(T)) && (c2 += 1; continue)
			push!(I, c2)
			push!(V, val)
			c2 += 1
		end
	end
	return c2
end

function inner_loop_to_modal{D, T<:Real}(c2::Int, k::Int, l::CartesianIndex{D},
									 level::CartesianIndex{D},
									 cell::CartesianIndex{D},
									 mode::CartesianIndex{D},
									 I::Array{Int, 1},
									 V::Array{T, 1},
									 base_indices::Array{Int, 1},
									 point_mat_1D::Array{T, 2})
	cells::NTuple{D, Int} = ntuple(i -> 1<<max(0, l[i]-2), D)
	modes::NTuple{D, Int} = ntuple(i -> k, D)
	cellrange = CartesianRange(cells)
	moderange = CartesianRange(modes)
	for c in cellrange
		!relevant_cell(k, l, c, level, cell) && (c2 += prod(modes); continue)

		for m in moderange
			val = one(T)
			for d in 1:D
				index = get_modal_index(k, l[d]-1, c[d], m[d])
				val *= point_mat_1D[index, base_indices[d]]
			end
			(abs(val) < eps(T)) && (c2 += 1; continue)
			push!(I, c2)
			push!(V, val)
			c2 += 1
		end
	end
	return c2
end

# Evaluates a given multidimensional basis element in the modal (DG Sparse) basis
# on all the points of the corresponding sparse grid
function eval_points{D, T<:Real}(k::Int, n::Int,
								level::CartesianIndex{D},
								cell::CartesianIndex{D},
								mode::CartesianIndex{D},
								modal_indices::Array{Int, 1},
								point_mat_1D::Array{T, 2};
								scheme="sparse")
	cutoff = get_cutoff(scheme, D, n)
	ls::NTuple{D, Int} = ntuple(i->(n+1),D)
	I  = Int[]
	V  = T[]
	c2 = 1

	for l in CartesianRange(ls)
		cutoff(l) && continue

		c2 = inner_loop_to_points(c2, k, l, level, cell, mode, I, V, modal_indices, point_mat_1D)
	end
	return sparsevec(I, V)
end

# Using the above method, builds the transformation matrix
# From the modal sparse basis to
# the "point basis" of values of a given function at the points on the sparse grid
function modal2points{D}(k::Int, n::Int,
							srefVD::Array{NTuple{3,CartesianIndex{D}},1};
							scheme="sparse")
	cutoff = get_cutoff(scheme, D, n)
	len = get_size(D, k, n; scheme=scheme)
	modal_ref = D2Vref(1, k, n)
	point_mat_1D = full(modal2points_1D(k, n))
	I = Int[]
	J = Int[]
	V = Real[]
	for c1 in 1:len
		lcm   = srefVD[c1]
		level = lcm[1]
		cell  = lcm[2]
		mode  = lcm[3]
		modal_indices = [modal_ref[(CartesianIndex{1}(level[d],),
								  CartesianIndex{1}(cell[d],),
								  CartesianIndex{1}(mode[d],))] for d in 1:D]
		basis_eval = eval_points(k, n, level, cell, mode,
									modal_indices,
									point_mat_1D;
									scheme=scheme)
		for (c2, val) in zip(findnz(basis_eval)...)
			push!(I, c2)
			push!(J, c1)
			push!(V, val)
		end
	end
	return sparse(I, J, V)
end

# Version of the above function with simpler arguments
function modal2points(D::Int, k::Int, n::Int; scheme="sparse")
	VD = V2Dref(D, k, n; scheme=scheme)
	return modal2points(k, n, VD; scheme=scheme)
end

# Given a specific collocation point, this constructs the corresponding
# combination of nodal basis functions that vanish on all collcation points but that one
function eval_nodal{D, T<:Real}(k::Int, n::Int, level::CartesianIndex{D},
								cell::CartesianIndex{D},
								mode::CartesianIndex{D},
								inv_mat_1D::Array{T, 2};
								scheme="sparse")
	cutoff = get_cutoff(scheme, D, n)
	ls::NTuple{D, Int} = ntuple(i->(n+1),D)
	point_indices = [get_point_index(k, level[d]-1, cell[d], mode[d]) for d in 1:D]

	I = Int[]
	V = T[]
	c2 = 1
	for l in CartesianRange(ls)
		cutoff(l) && continue

		c2 = inner_loop_to_points(c2, k, l, level, cell, mode, I, V, point_indices, inv_mat_1D)
	end
	return sparsevec(I, V)
end

# Using the above, this gives the transformation to go from the point basis
# to the nodal basis of functions
function points2nodal(D::Int, k::Int, n::Int; scheme="sparse")
	cutoff = get_cutoff(scheme, D, n)
	ls	   = ntuple(i->(n+1),D)
	hier_ref = D2Vref(1, k, n)
	inv_mat_1D = full(points2nodal_1D(k, n))
	I = Int[]
	J = Int[]
	V = Real[]

	c1 = 1
	for l in CartesianRange(ls)
		cutoff(l) && continue

		cells = ntuple(i -> 1<<max(0, l[i]-2), D)
		modes = ntuple(i->l[i]==1?k:(k-1),D)
		for c in CartesianRange(cells)
			for m in CartesianRange(modes)
				coeffs = eval_nodal(k, n, l, c, m,
									inv_mat_1D;
									scheme=scheme)
				for (c2, val) in zip(findnz(coeffs)...)
					push!(I, c2)
					push!(J, c1)
					push!(V, val)
				end
				c1 += 1
			end
		end
	end
	return dropzeros!(sparse(I, J, V))
end

# Given a nodal basis function, this expresses it in terms of the (sparse) DG Basis
function eval_DG{D, T<:Real}(k::Int, n::Int,
								level::CartesianIndex{D},
								cell::CartesianIndex{D},
								mode::CartesianIndex{D},
								hier_ref::Dict{NTuple{3, CartesianIndex{1}}, Int},
								inv_nodal_mat_1D::Array{T, 2};
								scheme="sparse")
	nodal_indices = [get_point_index(k, level[d]-1, cell[d], mode[d]) for d in 1:D]
	cutoff = get_cutoff(scheme, D, n)
	#len = get_size(D, k, n; scheme=scheme)
	I = Int[]
	V = T[]
	ls::NTuple{D, Int} = ntuple(i->(n+1),D)

	c2 = 1
	for l in CartesianRange(ls)
		cutoff(l) && continue

		c2 = inner_loop_to_modal(c2, k, l, level, cell, mode, I, V, nodal_indices, inv_nodal_mat_1D)
	end
	return sparsevec(I, V)
end

# This gives the transformation matrix between the (sparse) nodal basis
# and the (sparse) DG basis already developed in the previous paper
function nodal2modal{D}(k::Int, n::Int,
							srefVD::Array{NTuple{3, CartesianIndex{D}},1};
							scheme="sparse")
	cutoff = get_cutoff(scheme, D, n)
	ls     = ntuple(i->(n+1), D)
	hier_ref = D2Vref(1, k, n)
	inv_point_mat_1D = full(nodal2modal_1D(k, n))
	I = Int[]
	J = Int[]
	V = Real[]

	c1 = 1
	for l in CartesianRange(ls)
		cutoff(l) && continue

		cells = ntuple(i -> 1<<max(0, l[i]-2), D)
		modes = ntuple(i->l[i]==1?k:k-1,D)
		for c in CartesianRange(cells)
			for m in CartesianRange(modes)
				coeffs = eval_DG(k, n, l, c, m,
									hier_ref,
									# srefVD,
									inv_point_mat_1D;
									scheme=scheme)
				for (c2, val) in zip(findnz(coeffs)...)
					push!(I, c2)
					push!(J, c1)
					push!(V, val)
				end
				c1 += 1
			end
		end
	end

	return dropzeros!(sparse(I, J, V))
end

# Version of the above function with simpler arguments
function nodal2modal(D::Int, k::Int, n::Int; scheme="sparse")
	VD = V2Dref(D, k, n; scheme=scheme)
	return nodal2modal(k, n, VD; scheme=scheme)
end

function points2modal(D::Int, k::Int, n::Int)
	return nodal2modal(D, k, n)*points2nodal(D, k, n)
end


"""
    Reluplex(optimizer, eager::Bool)

Reluplex uses binary tree search to find an activation pattern that maps a feasible input to an infeasible output.

# Problem requirement
1. Network: any depth, ReLU activation
2. Input: hyperrectangle
3. Output: PolytopeComplement

# Return
`CounterExampleResult`

# Method
Binary search of activations (0/1) and pruning by optimization.

# Property
Sound and complete.

# Reference
[G. Katz, C. Barrett, D. L. Dill, K. Julian, and M. J. Kochenderfer,
"Reluplex: An Efficient SMT Solver for Verifying Deep Neural Networks," in
*International Conference on Computer Aided Verification*, 2017.](https://arxiv.org/abs/1702.01135)
"""
@with_kw struct Reluplex <: Solver
    optimizer = GLPK.Optimizer
end

function solve(solver::Reluplex, problem::Problem)
    model = Model(solver)
    encode(solver, model, problem)
    return reluplex_step(solver, model)
end

function find_relu_to_fix(ẑ, z)
    for i in 1:length(z), j in 1:length(z[i])
        ẑᵢⱼ = value(ẑ[i][j])
        zᵢⱼ = value(z[i][j])

        if type_one_broken(ẑᵢⱼ, zᵢⱼ) ||
           type_two_broken(ẑᵢⱼ, zᵢⱼ)
            return (i, j)
        end
    end
    return (0, 0)
end

 # TODO consider renaming to `inactive_broken` and `active_broken`
type_one_broken(ẑᵢⱼ, zᵢⱼ) = (zᵢⱼ > 0.0 + TOL[])  && (!(-TOL[] < ẑᵢⱼ - zᵢⱼ < TOL[])) # (zᵢⱼ > 0) && (ẑᵢⱼ != zᵢⱼ)
type_two_broken(ẑᵢⱼ, zᵢⱼ) = (-TOL[] < zᵢⱼ < TOL[]) && (ẑᵢⱼ > 0.0 + TOL[]) # (zᵢⱼ == 0) && (ẑᵢⱼ > 0)

# Corresponds to a ReLU that shouldn't be active but is
function type_one_repair!(model, ẑᵢⱼ, zᵢⱼ)
    con_one = @constraint(model, ẑᵢⱼ == zᵢⱼ)
    con_two = @constraint(model, ẑᵢⱼ >= 0.0)
    return con_one, con_two
end
# Corresponds to a ReLU that should be active but isn't
function type_two_repair!(model, ẑᵢⱼ, zᵢⱼ)
    con_one = @constraint(model, ẑᵢⱼ <= 0.0)
    con_two = @constraint(model, zᵢⱼ == 0.0)
    return con_one, con_two
end

function activation_constraint!(model, ẑᵢ, zᵢ, act::ReLU)
    # ReLU ensures that the variable after activation is always
    # greater than before activation and also ≥0
    @constraint(model, zᵢ .>= ẑᵢ)
    @constraint(model, zᵢ .>= 0.0)
end

function activation_constraint!(model, ẑᵢ, zᵢ, act::Id)
    @constraint(model, zᵢ .== ẑᵢ)
end

function encode(solver::Reluplex, model::Model,  problem::Problem)
    layers = problem.network.layers
    ẑ = init_vars(model, layers, :ẑ, with_input=true) # before activation
    z = init_vars(model, layers, :z, with_input=true) # after activation

    # Each layer has an input set constraint associated with it based on the bounds.
    # Additionally, consective variables zᵢ, ẑᵢ₊₁ are related by a constraint given
    # by the affine map encoded in the layer Lᵢ.
    # Finally, the before-activation-variables and after-activation-variables are
    # related by the activation function. Since the input layer has no activation,
    # its variables are related implicitly by identity.
    activation_constraint!(model, ẑ[1], z[1], Id())
    bounds = get_bounds(problem)
    for (i, L) in enumerate(layers)
        @constraint(model, affine_map(L, z[i]) .== ẑ[i+1])
        add_set_constraint!(model, bounds[i], z[i])
        activation_constraint!(model, ẑ[i+1], z[i+1], L.activation)
    end
    # Add the bounds on your output layer
    add_set_constraint!(model, last(bounds), last(z))
    # Add the complementary set defind as part of the problem
    add_complementary_set_constraint!(model, problem.output, last(z))
    feasibility_problem!(model)

    return nothing
end

function reluplex_step(solver::Reluplex, model::Model)

    optimize!(model)

    ẑ, z = model[:ẑ], model[:z]
    # If the problem is optimally solved, this could potentially be a counterexample.
    # Branch by repair type (inactive or active) and redetermine if this is a valid
    # counterexample. If the problem is infeasible or unbounded. The property holds.
    if termination_status(model) == OPTIMAL
        i, j = find_relu_to_fix(ẑ, z)

        # In case no broken relus could be found, return the input as a counterexample
        i == 0 && return CounterExampleResult(:violated, value(first(ẑ)))

        for repair! in (type_one_repair!, type_two_repair!)
            # Add the constraints associated with the ReLU being fixed
            new_constraints = repair!(model, ẑ[i][j], z[i][j])

            # Recurse with the ReLU i, j fixed to active or inactive
            result = reluplex_step(solver, model)

            # Return (all the way to top level) if violated.
            result.status == :violated && return result

            # Reset the model when we're done with this ReLU
            delete.(model, new_constraints)
        end
    end
    return CounterExampleResult(:holds)
end


# # Symmetry reduction

#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/generated/Symmetry reduction.ipynb)
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/generated/Symmetry reduction.ipynb)
# **Adapted from**: [SymbolicWedderburn example](https://github.com/kalmarek/SymbolicWedderburn.jl/blob/tw/ex_sos/examples/ex_C4.jl)

import MutableArithmetics
const MA = MutableArithmetics
using MultivariatePolynomials
const MP = MultivariatePolynomials
using MultivariateBases
const MB = MultivariateBases

using Test #src
using DynamicPolynomials
@polyvar x[1:4]

# We would like to find the minimum value of the polynomial

poly = sum(x) + sum(x.^2)

# As we can decouple the problem for each `x[i]` for which `x[i] + x[i]^2` has
# minimum value 0.25, we would expect to get `-1` as answer.
# Can this decoupling be exploited by SumOfSquares as well ?
# For this, we need to use a certificate that can exploit the permutation symmetry of the polynomial.
# Symmetry reduction is still a work in progress in SumOfSquares, so we include the following files that will be incorporated into SumOfSquares.jl once SymbolicWedderburn.jl is released:
using SumOfSquares
include(joinpath(dirname(dirname(pathof(SumOfSquares))), "examples", "symmetry.jl"))

# We define the symmetry group as a permutation group in the variables.
# In order to do that, we define the action of a permutation on a monomial
# as the monomial obtained after permuting the variables.

using PermutationGroups
function action(mono::AbstractMonomial, p::Permutation)
    v = variables(mono)
    MP.substitute(MP.Eval(), mono, v => [v[i^p] for i in eachindex(v)])
end
function action(term::MP.AbstractTerm, el::Permutation)
    return MP.coefficient(term) * action(MP.monomial(term), el)
end
function action(poly::MP.AbstractPolynomial, el::Permutation)
    return MP.polynomial([action(term, el) for term in MP.terms(poly)])
end

G = PermGroup([perm"(1,2,3,4)"])

# We can use this certificate as follows:

import CSDP
solver = CSDP.Optimizer
model = Model(solver)
@variable(model, t)
@objective(model, Max, t)
certificate = SymmetricIdeal(Certificate.MaxDegree(SOSCone(), MonomialBasis, maxdegree(poly)), G, action)
con_ref = @constraint(model, poly - t in SOSCone(), ideal_certificate = certificate)
optimize!(model)
@test value(t) ≈ -1 #src
value(t)

# We indeed find `-1`, let's verify that symmetry was exploited:

g = gram_matrix(con_ref).sub_gram_matrices #src
@test length(g) == 4                       #src
@test g[1].basis.polynomials == [sum(x), 1.0] #src
@test size(g[1].Q) == (2, 2)                  #src
@test g[1].Q[1, 1] ≈ 1/4 atol=1e-6            #src
@test g[1].Q[1, 2] ≈ 1/2 atol=1e-6            #src
@test g[1].Q[2, 2] ≈ 1.0 atol=1e-6            #src
@test g[2].basis.polynomials == [x[1] - x[3]] #src
@test size(g[2].Q) == (1, 1)                  #src
@test g[2].Q[1, 1] ≈ 1/2 atol=1e-6            #src
@test g[3].basis.polynomials == [x[2] - x[4]] #src
@test size(g[3].Q) == (1, 1)                  #src
@test g[3].Q[1, 1] ≈ 1/2 atol=1e-6            #src
@test g[4].basis.polynomials == [x[1] - x[2] + x[3] - x[4]] #src
@test size(g[4].Q) == (1, 1)                  #src
@test g[4].Q[1, 1] ≈ 1/4 atol=1e-6            #src
for g in gram_matrix(con_ref).sub_gram_matrices
    println(g.basis.polynomials)
end


using LinearAlgebraicRepresentation
Lar = LinearAlgebraicRepresentation
using ViewerGL
GL = ViewerGL
using LinearAlgebra

# input of primitive shapes
V,(VV,EV,FV) = Lar.simplex(2, true)
W,(WW,EW,FW) = Lar.cuboid([1,1], true)
triangle = (V,EV)
square = (W,EW)
# hierarchical assembly
model = Lar.Struct([
            Lar.Struct([
                square,
                Lar.t(.175,.175), Lar.s(.5,.5),
                triangle
            ]),
            Lar.t(.825,.825), Lar.s(-.5,-.5),
            triangle
        ])
# visualization of generated wire-frame model
V,EV = Lar.struct2lar(model)
GL.VIEW(GL.numbering(.4)((V,[VV, EV])));

triangles = Lar.triangulate2d(V, EV)
GL.VIEW([ GL.GLGrid(V,triangles) ]);

# generate edges
ev = map(sort,cat([[[u,v], [v,w], [w,u]] for (u,v,w) in triangles]))
# remove duplicated edges from triangulation
nodupev = collect(Set(ev))

GL.VIEW(GL.numbering(.3)((V,[VV, nodupev])));
GL.VIEW([ GL.GLGrid(V,triangles) ]);
GL.VIEW(GL.numbering(.25)((V,[VV, nodupev, triangles]),GL.COLORS[1],0.2));


module AlgorithmTests

using Test
using HDF5
using FFTW

using PicoQuant
include("test_utils.jl")

@testset "Algorithm tests" begin
    @testset "Test state preparation algorithm" begin
        @test begin
            qubits = 1
            circ = create_simple_preparation_circuit(qubits, 4, 42)

            ψ = get_statevector_using_picoquant(circ,
                                                big_endian=false)
            ψ′ = get_statevector_using_qiskit(circ, big_endian=false)

            overlap = (a, b) -> abs(transpose(conj.(a)) * b)

            overlap(ψ, ψ′) ≈ 1.
        end

        @test begin
            qubits = 3
            circ = create_simple_preparation_circuit(qubits, 4, 42)

            ψ = get_statevector_using_picoquant(circ,
                                                big_endian=false)
            ψ′ = get_statevector_using_qiskit(circ, big_endian=false)

            overlap = (a, b) -> abs(transpose(conj.(a)) * b)

            overlap(ψ, ψ′) ≈ 1.
        end
    end

    @testset "Small QFT circuit test" begin
        @test begin
            n = 3

            prep_circ = create_simple_preparation_circuit(n, 3, 43)
            qft_circ = create_qft_circuit(n)

            full_circ = prep_circ.combine(qft_circ)

            qft_input = get_statevector_using_qiskit(prep_circ, big_endian=false)
            ref_output = ifft(qft_input)
            norm = x -> x ./ sqrt(sum(x .* conj.(x)))
            ref_output = norm(ref_output)

            ψ = get_statevector_using_picoquant(full_circ)

            overlap = (a, b) -> abs(transpose(conj.(a)) * b)

            overlap(ψ, ref_output) ≈ 1.
        end
    end

    @testset "Larger QFT circuit test" begin
        @test begin
            n = 8

            prep_circ = create_simple_preparation_circuit(n, 3, 43)
            qft_circ = create_qft_circuit(n)

            full_circ = prep_circ.combine(qft_circ)

            qft_input = get_statevector_using_qiskit(prep_circ, big_endian=false)
            ref_output = ifft(qft_input)
            norm = x -> x ./ sqrt(sum(x .* conj.(x)))
            ref_output = norm(ref_output)

            ψ = get_statevector_using_picoquant(full_circ)

            overlap = (a, b) -> abs(transpose(conj.(a)) * b)

            overlap(ψ, ref_output) ≈ 1.
        end
    end

    @testset "Small RQC circuit test" begin
        @test begin
            n = 3; m = 3; depth = 8

            rqc = create_RQC(n, m, depth)

            qiskit_Ψ = get_statevector_using_qiskit(rqc, big_endian=false)

            picoquant_ψ = get_statevector_using_picoquant(rqc)

            overlap = (a, b) -> abs(transpose(conj.(a)) * b)

            overlap(picoquant_ψ, qiskit_Ψ) ≈ 1.
        end
    end

    @testset "iSWAP RQC circuit test" begin
        @test begin
            n = 3; m = 3; depth = 8

            rqc = create_RQC(n, m, depth;
                            use_iswap=true, final_Hadamard_layer=true)

            qiskit_Ψ = get_statevector_using_qiskit(rqc, big_endian=false)

            picoquant_ψ = get_statevector_using_picoquant(rqc)

            overlap = (a, b) -> abs(transpose(conj.(a)) * b)

            overlap(picoquant_ψ, qiskit_Ψ) ≈ 1.
        end
    end
end
end


Pkg.add("VennDiagrams")

using VennDiagrams

c1 = [0:5] # Can be any iterable
c2 = [3:10]

p = venn(c1, c2)
draw(PDF("beaut1.pdf", 8cm, 10cm), p)

c3 = [0:2:12]

p = venn(c1, c2, c3, proportional=false)
draw(PDF("beaut3.pdf", 8cm, 10cm), p)


# This file constains code copied from https://github.com/jeanpauphilet/SubsetSelection.jl
# as at commit e6b3258ba394f6152853498d81d3cecb9d67e097

const sgtol = 1e-6

mutable struct PolyakCache
  best_inds::Vector{Int}
  best_upper::Float64
  best_lower::Float64
end
PolyakCache(inds::Vector{Int}) = PolyakCache(inds, Inf, -Inf)

abstract type SteppingRule end
struct ConstantStepping <: SteppingRule
  stepsize::Float64
end
mutable struct PolyakStepping <: SteppingRule
  initial_factor::Float64        # initial factor for scaling (ub-lb)/norm^2
  wait::Int64                    # currently unused TODO
  function PolyakStepping(initial_factor, wait)
      if 0.0 <= initial_factor <= 2.0
          new(initial_factor, wait)
      else
          error("Initial stepping parameter must be between 0 and 2. Current
                    choice is $initial_factor.")
      end
  end
end
PolyakStepping() = PolyakStepping(2.0, 30)

function resetindices(::PolyakStepping)
  true
end
function resetindices(::SteppingRule)
  false
end

function getdelta!(sr::SteppingRule, ::Any, ::Array{Float64,2}, ::Vector{Float64}, ::Vector{Float64}, ::Vector{Int}, ::Int, ::Float64, ::SubsetSelection.Cache, ::PolyakCache)
  error("Need to define `getdelta` for stepping rule $(sr).")
end
function getdelta!(sr::ConstantStepping, ::Any, ::Array{Float64,2}, ::Vector{Float64}, ::Vector{Float64}, ::Vector{Int}, ::Int, ::Float64, ::SubsetSelection.Cache, ::PolyakCache)
  sr.stepsize
end
function getdelta!(sr::PolyakStepping,
      Y::YVector, X::Array{Float64,2},
      α::Vector{Float64}, ∇::Vector{Float64},
      indices::Vector{Int}, n_indices::Int, γ::Float64,
      cache::SubsetSelection.Cache, pc::PolyakCache
     )
  @timeit TimeStats "Dual Bound" lower_bound = dual_bound(SubsetSelection.OLS(), Y, X, α, indices, n_indices, γ, cache)
  @timeit TimeStats "Primal Bound" upper_bound = primal_bound(SubsetSelection.OLS(), Y, X, γ, indices, n_indices)
  if upper_bound < pc.best_upper
    pc.best_upper = upper_bound
    pc.best_inds .= indices
  end
  (lower_bound > pc.best_lower) && (pc.best_lower = lower_bound)
  if lower_bound - sgtol > upper_bound + sgtol
      error("Bounds overlap. LB = $lowerbound and UB = $upper_bound.")
  end
  sr.initial_factor * (pc.best_upper - lower_bound) / sum(abs2.(∇))
end

##############################################
##DUAL SUB-GRADIENT ALGORITHM
##############################################
""" Function to compute a sparse regressor/classifier. Solve an optimization problem of the form
        min_s max_α f(α, s)
by gradient ascent on α:        α_{t+1} ← α_t + δ ∇_α f(α_t, s_t)
and partial minimization on s:  s_{t+1} ← argmin_s f(α_t, s)
INPUTS
- ℓ           Loss function used
- Card        Model to enforce sparsity (constraint or penalty)
- Y (n×1)     Vector of outputs. For classification, use ±1 labels
- X (n×p)     Array of inputs.
- indInit     (optional) Initial subset of features s
- αInit       (optional) Initial dual variable α
- γ           (optional) ℓ2 regularization penalty
- intercept   (optional) Boolean. If true, an intercept term is computed as well
- maxIter     (optional) Total number of Iterations
- δ           (optional) Gradient stepsize
- gradUp      (optional) Number of gradient updates of dual variable α performed per update of primal variable s
- anticycling (optional) Boolean. If true, the algorithm stops as soon as the support is not unchanged from one iteration to another
- averaging   (optional) Boolean. If true, the dual solution is averaged over past iterates
OUTPUT
- SparseEstimator """
function subsetSelection_bm(ℓ::LossFunction, Card::Sparsity,
    Y::YVector,
    X::Array{Float64,2};
    indInit = SubsetSelection.ind_init(Card, size(X,2)),
    αInit = SubsetSelection.alpha_init(ℓ, Y),
    γ = 1/sqrt(size(X,1)),  intercept = false,
    maxIter = 100, sr::SteppingRule = ConstantStepping(1e-3), gradUp = 1,
    anticycling = false, averaging = true)

  @timeit TimeStats "in subgradient descent" begin

  n,p = size(X)
  cache = SubsetSelection.SubsetSelection.Cache(n, p)

  #Add sanity checks
  if size(Y,1) != n
    throw(DimensionMismatch("X and Y must have the same number of rows"))
  end
  if isa(ℓ, SubsetSelection.Classification)
    levels = sort(unique(Y))
    if length(levels) != 2
      throw(ArgumentError("subsetSelection only supports two-class classification"))
    elseif (levels[1] != -1) || (levels[2] != 1)
      throw(ArgumentError("Class labels must be ±1's"))
    end
  end

  indices = indInit #Support
  n_indices = length(indices)
  best_inds = copy(indices)

  n_indices_max = SubsetSelection.max_index_size(Card, p)
  resize!(indices, n_indices_max)

  indices_old = Vector{Int}(n_indices_max)
  α = αInit[:]  #Dual variable α
  a = αInit[:]  #Past average of α

  # We will compute the dual bound incorrectly unless indicies match initial α
  if resetindices(sr)
    n_indices = SubsetSelection.partial_min!(indices, Card, X, α, γ, cache)
  end

  pc = PolyakCache(copy(indices))

  bounds_match = false

  ##Dual Sub-gradient Algorithm
  for iter in 2:maxIter

    #Gradient ascent on α
    for _ in 1:min(gradUp, div(p, n_indices))
      ∇ = SubsetSelection.grad_dual(ℓ, Y, X, α, indices, n_indices, γ, cache)
      δ = getdelta!(sr, Y, X, α, ∇, indices, n_indices, γ, cache, pc)
      α .+= δ*∇
      α = SubsetSelection.proj_dual(ℓ, Y, α)
      α = SubsetSelection.proj_intercept(intercept, α)
    end

    if !all(isfinite.(α))
        warn("Algorithm diverges! Did you normalize your data? Otherwise, try reducing stepsize δ.")
    end
    #Update average a
    @__dot__ a = (iter - 1) / iter * a + 1 / iter * α
    # a *= (iter - 1)/iter; a .+= α/iter

    #Minimization w.r.t. s
    indices_old[1:n_indices] = indices[1:n_indices]
    n_indices = SubsetSelection.partial_min!(indices, Card, X, α, γ, cache)

    #Anticycling rule: Stop if indices_old == indices
    if anticycling && SubsetSelection.indices_same(indices, indices_old, n_indices)
      averaging = false #If the algorithm stops because of cycling, averaging is not needed
      break
    end
    # If lower and upper bounds have closed, skip to the end
    if abs(pc.best_upper - pc.best_lower) / (1 + abs(pc.best_upper)) < 1e-4
      indices .= pc.best_inds
      break
    end
  end

  ##Compute sparse estimator
  #Subset of relevant features
  n_indices = SubsetSelection.partial_min!(indices, Card, X, averaging ? a : α, γ, cache)
  #Regressor
  # w = [-γ * dot(X[:, indices[j]], a) for j in 1:n_indices]
  w = SubsetSelection.recover_primal(ℓ, Y, X[:,indices], γ)
  #Bias
  b = SubsetSelection.compute_bias(ℓ, Y, X, a, indices, n_indices, γ, intercept, cache)

  #Resize final indices vector to only have relevant entries
  resize!(indices, n_indices)

   end # timing

  return SubsetSelection.SparseEstimator(ℓ, Card, indices, w, a, b, maxIter)
end

function ax_squared(X, α::Vector{Float64}, indices::Vector{Int}, n_indices::Int)
  # TODO update ax in cache and use it rather than recomputing
  axsum = 0.0
  for j = 1:n_indices
    x = @view(X[:, indices[j]])
    axsum += dot(α, x)^2
  end
  axsum
end

function primal_bound(ℓ::SubsetSelection.OLS, Y::YVector, X::Array{Float64,2}, γ::Float64, indices::Vector{Int}, n_indices::Int)
  @timeit TimeStats "Inverting matrix" αstar = SubsetSelectionCIO.sparse_inverse(ℓ, Y, X[:, indices], γ) # TODO could do this less often. also don't return vector after function.
  axsum = ax_squared(X, αstar, indices, n_indices)
  bound = -0.5 * dot(αstar, αstar) - dot(Y, αstar) - γ * 0.5 * axsum
  # Normalize TODO normalize this and grad dual
  # bound / size(X, 1)
end

function dual_bound(ℓ::SubsetSelection.OLS, Y::YVector, X::Array{Float64,2}, α::Vector{Float64}, indices::Vector{Int}, n_indices::Int, γ, cache::SubsetSelection.Cache)
  axsum = ax_squared(X, α, indices, n_indices)
  bound = -0.5 * dot(α, α) - dot(Y, α) - γ * 0.5 * axsum
  # Normalize
  # bound / size(X, 1)
end


correct = ( pi/4 + atan(3.0) ) / 2
for endpt in [neither, both, left, right]
    x, w = legendre(40, endpt)
    if endpt in [left, both]
        @test x[1] == -1
    end
    if endpt in [right, both]
        @test x[end] == 1
    end
    Q = sum(w .* 1 ./ ( 1 .+ (2x .- 1).^2))
    @test Q ≈ correct
end


using NURBS
using Base.Test

@testset "Base Generator" begin

    @testset "basis" begin
        x = knot(4,4)
        @test typeof(basis(4,4,0.5,x)) == Array{Float64,1}
        @test isapprox(basis(4,4,0.5,x),[0.125, 0.375, 0.375, 0.125],atol=1e-5)
    end

    @testset "dbasis" begin
        x = knot(4,4)
        @test typeof(dbasis(4,4,0.5,x)) == Tuple{Array{Float64,1}, Array{Float64,1}, Array{Float64,1}}
        @test isapprox(dbasis(4,4,0.5,x)[1],[0.125, 0.375, 0.375, 0.125],atol=1e-5)
        @test isapprox(dbasis(4,4,0.5,x)[2],[-0.75, -0.75, 0.75, 0.75],atol =1e-5)
        @test isapprox(dbasis(4,4,0.5,x)[3],[3.0, -3.0, -3.0, 3.0],atol =1e-5)
    end

    @testset "dbasisu" begin
        x = knotu(5,3)
        @test typeof(dbasisu(5,3,0.5,x)) == Tuple{Array{Float64,1}, Array{Float64,1}, Array{Float64,1}}
        @test isapprox(dbasisu(5,3,0.5,x)[1],[0.125, 0.0, 0.0, 0.0, 0.0],atol=1e-5)
        @test isapprox(dbasisu(5,3,0.5,x)[2],[0.5, 0.0, 0.0, 0.0, 0.0],atol=1e-5)
        @test isapprox(dbasisu(5,3,0.5,x)[3],[1.0, 0.0, 0.0, 0.0, 0.0],atol=1e-5)
    end
end


using OmegaCore
using Test
using Distributions
using OmegaTest

function samplemodel()
  x = 1 ~ Normal(0, 1)
  y = 2 ~ Normal(0, 1)
  z(ω) = (x(ω), y(ω), x(ω))
  randsample(z)
end

function test_ciid()
  @test isinferred(samplemodel)
  z_ = samplemodel()
  @test z_[1] != z_[2]
  @test z_[1] == z_[3]
end

function test_shared_parent()
  parent(ω) = Bool(Bernoulli(0.5)((1,), ω)) ? -100 : 100
  x(id, ω) = parent(ω) + Uniform(0, 1)(id, ω)
  x1 = 2 ~ x
  x2 = 3 ~ x
  z = ω -> (x1(ω), x2(ω))
  samples = randsample(z, 10)
  t1 = [(x1_ != x2_) for (x1_, x2_) in samples]
  t2 = [(abs(x1_ - x2_) <= 2) for (x1_, x2_) in samples]
  @test all(t1)
  @test all(t2)
end

@testset "ciid" begin
  test_ciid()
  test_shared_parent()
end


q = Quaternion(rand() * 2pi, normalize(ℝ³(rand(3))))
cp = ComplexPlane(q)
s = SU2(cp)
θ = rand() * pi
u = normalize(ℝ³(rand(3)))
g = Quaternion(θ, u)

@test isapprox(cp, s)
@test isapprox(cp, q)
@test isapprox(g * q, g * Quaternion(s))


"""
    eigHAD_Distance(𝚽, 𝛌; indexEigs = 1:size(𝚽,2))

compute HAD "distance" (not really a distance) between pairwise graph Laplacian
eigenvectors, i.e., d_HAD(𝜙ᵢ₋₁, 𝜙ⱼ₋₁) = √(1 - a_HAD(𝜙ᵢ₋₁, 𝜙ⱼ₋₁)²).

# Input Arguments
- `𝚽::Matrix{Float64}`: matrix of graph Laplacian eigenvectors, 𝜙ⱼ₋₁ (j = 1,...,size(𝚽,1)).
- `𝛌::Array{Float64}`: array of eigenvalues. (ascending order)
- `indexEigs::Int`: default is all eigenvectors, indices of eigenvectors considered.

# Output Argument
- `dis::Matrix{Float64}`: the HAD distance matrix, dis[i,j] = d_HAD(𝜙ᵢ₋₁, 𝜙ⱼ₋₁).
"""
function eigHAD_Distance(𝚽, 𝛌; indexEigs = 1:size(𝚽,2))
    n = length(indexEigs)
    A = eigHAD_Affinity(𝚽, 𝛌; indexEigs = indexEigs)
    dis = sqrt.(ones(n, n) - A.^2)
    dis[diagind(dis)] .= 0
    return dis
end

function eigHAD_Distance_neglog(𝚽, 𝛌; indexEigs = 1:size(𝚽,2))
    A = eigHAD_Affinity(𝚽, 𝛌; indexEigs = indexEigs)
    n = size(A,1)
    dis = zeros(n,n)
    for i = 1:n, j = 1:n
        if A[i,j] == 0
            dis[i,j] = 1e9
        else
            dis[i,j] = -log(A[i,j])
        end
    end
    return dis
end

"""
    eigHAD_Affinity(𝚽, 𝛌; indexEigs = 1:size(𝚽,2))

EIGHAD_AFFINITY compute Hadamard (HAD) affinity between pairwise graph Laplacian eigenvectors.

# Input Arguments
- `𝚽::Matrix{Float64}`: matrix of graph Laplacian eigenvectors, 𝜙ⱼ₋₁ (j = 1,...,size(𝚽,1)).
- `𝛌::Array{Float64}`: array of eigenvalues. (ascending order)
- `indexEigs::Int`: default is all eigenvectors, indices of eigenvectors considered.

# Output Argument
- `A::Matrix{Float64}`: a numEigs x numEigs affinity matrix, A[i,j] = a_HAD(𝜙ᵢ₋₁, 𝜙ⱼ₋₁).
"""
function eigHAD_Affinity(𝚽, 𝛌; indexEigs = 1:size(𝚽,2))
    N, numEigs = size(𝚽,1), length(indexEigs)
    indNoDC = setdiff(indexEigs, 1) # get rid of DC component
    J = length(indNoDC)
    A = zeros(J, J)
    for a in 1:J, b in a:J
        i, j = indNoDC[a], indNoDC[b]
        hadamardProd = 𝚽[:,i] .* 𝚽[:,j]
        if norm(hadamardProd,2) < 0.01/sqrt(N)
            continue
        end
        λ, μ = 𝛌[i], 𝛌[j]
        x₀ = 1 ./ (max(λ, μ))
        # Find minimizer t
        result = optimize(t -> abs(exp(-t[1]*λ) + exp(-t[1]*μ) - 1), [x₀], BFGS());
        t = Optim.minimizer(result)[1]
        # Compute Hadamard affinity
        heatEvolution = 𝚽 * Diagonal(exp.(-t .* 𝛌)) * 𝚽' * hadamardProd
        A[a,b] = norm(heatEvolution,2) / (norm(hadamardProd,2) + 1e-6)
    end
    A = A + A'; for i in 1:J; A[i,i] /= 2; end

    if 1 in indexEigs
        # Set affinity measure of 𝜙₀ with itself to be the maximum and equals to 1.
        A = matrix_augmentation(A)
        A[1,1] = maximum(A)
    end
    return A ./ maximum(A)
end

function matrix_augmentation(A)
    m, n = size(A)
    B = zeros(m+1, n+1)
    B[2:end, 2:end] = A
    return B
end


function getGaussianKernel(s::Int, width::Float64)
    σ = width / 4.3 # full width at tenth maximum = 4.3 * σ
    k = 1./(σ .* sqrt(2*π)).*exp.(-0.5.*(   (0.5:s)   .-s/2).^2/(σ.^2))
    #k = g(0.5:s)
    k = k ./ sum(k)
    return k
end

function smoothArray(a::Array{Float64,1}, kernelSize::Int, width::Float64; kernel::Array{Float64,1}=Array{Float64,1}(0))
    if length(kernel)==0
        kernel = getGaussianKernel(kernelSize, width)
    end

    halfKernelSize = Int((kernelSize-1)/2)
    return conv(a, kernel)[halfKernelSize+1:end-halfKernelSize]
end

function smoothArray(a::Array{Float64,2}, kernelSize::Int, width::Float64; kernel::Array{Float64,1}=Array{Float64,1}(0))
    if length(kernel)==0
        kernel = getGaussianKernel(kernelSize, width)
    end
    #kernel=[1/6, 2/3, 1/6]
    a_smooth = zeros(size(a))
    for i = 1:size(a,1)
        a_smooth[i,:] = smoothArray(a[i,:], kernelSize, width, kernel=kernel)
    end
    for i = 1:size(a,2)
        a_smooth[:,i] = smoothArray(a_smooth[:,i], kernelSize, width, kernel=kernel)
    end

    return a_smooth
end

function smoothArray(a::scaledArray, kernelSize::Int, width::Float64; kernel::Array{Float64,1}=Array{Float64,1}(0))
    if length(kernel)==0
        kernel = getGaussianKernel(kernelSize, width)
    else
        x = reshape(a.data[1:Int(end/2)], a.dimensions)
        xs = smoothArray(x, kernelSize, width)[:]
        y = reshape(a.data[1+Int(end/2):end], a.dimensions)
        ys = smoothArray(y, kernelSize, width)[:]
    end
    return scaledArray(vcat(xs,ys), a.dimensions, a.voxelsize, a.shift)
end


using AlgebraicPetri
using AlgebraicPetri.Epidemiology

using Catlab
using Catlab.Graphics
using Catlab.WiringDiagrams
using Catlab.CategoricalAlgebra
using Catlab.Programs.RelationalPrograms

using LabelledArrays
using OrdinaryDiffEq
using Random
using Plots

# helper function to visualize categorical representation
display_uwd(ex) = to_graphviz(ex, box_labels=:name, junction_labels=:variable, edge_attrs=Dict(:len=>".75"));

# population x spontaneously moves to population y
spontaneous_petri(x::Symbol, y::Symbol, transition::Symbol) =
    Open(LabelledPetriNet(unique([x,y]), transition=>(x, y)))
# population y causes population x to move to population z
exposure_petri(x::Symbol, y::Symbol, z::Symbol, transition::Symbol) =
    Open(LabelledPetriNet(unique([x,y,z]), transition=>((x,y)=>(z,y))))

infection = exposure_petri(:S, :I, :I, :inf)
exposure = exposure_petri(:S, :I, :E, :exp)
illness = spontaneous_petri(:E,:I,:ill)
recovery = spontaneous_petri(:I,:R,:rec)
death = spontaneous_petri(:I,:D,:death)

epi_dict = Dict(:infection=>infection,
                :exposure=>exposure,
                :illness=>illness,
                :recovery=>recovery,
                :death=>death)

oapply_epi(ex, args...) = oapply(ex, epi_dict, args...)

Graph(infection)

Graph(recovery)

sir_wiring_diagram = @relation (s, i, r) begin
    infection(s, i)
    recovery(i, r)
end
display_uwd(sir_wiring_diagram)

sir_model = apex(oapply_epi(sir_wiring_diagram));
Graph(sir_model)

tmax = 40.0
tspan = (0.0,tmax);

u0 = LVector(S=990.0, I=10.0, R=0.0)

p = LVector(inf=0.05*10.0/sum(u0), rec=0.25); # β*c/N,γ

Random.seed!(1234);

prob_ode = ODEProblem(vectorfield(sir_model),u0,tspan,p)
sol_ode = solve(prob_ode, Tsit5());
plot(sol_ode)