include("Conserved.jl")
include("HydroSnapshot.jl")

mutable struct HDSim
    time::Float64 # Time
    cycle::Int64 # Cycle number
    state::HydroSnapshot # Primitive variables and grid
    extensives::Array{Conserved, 1} # Conserved quantities
    eos # Equation of state
    tsf # Time step function
    vm # Vertex motion
    fc # Flux calculator
    pg # Physical geometry
    st # Source term
    cu # Cell updater
end

function primitive2extensive(p::PrimitiveVars,
                             volume::Float64,
                             eos)::Conserved
    mass = p.density*volume
    momentum = mass*p.velocity
    kinetic_specific_energy = 0.5*p.velocity^2
    thermal_specific_energy =
     calcSpecificThermalEnergy(eos,
                        p.density,
                        p.pressure)
    specific_energy =
     kinetic_specific_energy+
     thermal_specific_energy
    energy = specific_energy*mass
    return Conserved(mass,
                     momentum,
                     energy)
end

function initConserved(state::HydroSnapshot, pg, eos)
    cumulative_volume = [calcVolume(pg, r) for r in state.grid]
    volumes = diff(cumulative_volume)
    conserved_list = [primitive2extensive(p,V,eos)
                      for (p,V) in zip(state.cells, volumes)]
    return conserved_list
end

function initHDSim(init_cond::HydroSnapshot,
                   eos,
                   tsf,
                   vm,
                   fc,
                   pg,
                   st,
                   cu)
    conserved_list = initConserved(init_cond, pg, eos)
    return HDSim(0,
                 0,
                 init_cond,
                 conserved_list,
                 eos,
                 tsf,
                 vm,
                 fc,
                 pg,
                 st,
                 cu)
end

function updateConserved!(conserved_list::Array{Conserved,1},
                          fluxes::Array{Conserved,1},
                          grid::Array{Float64,1},
                          pg,
                          dt::Float64)
    areas = [calcArea(pg,r) for r in grid]
    currents = dt.*areas.*fluxes
    conserved_list .-= diff(currents)
end

function updateSourceContribution!(st,
                                  extensives::Array{Conserved,1},
                                  state::HydroSnapshot,
                                  t::Float64,
                                  dt::Float64)
    extensives .+= dt*calcSource(st,state,t,dt)
end

#function updatePositions!(grid::Array{Float64, 1},
#                          velocities::Array{Float64, 1},
#                          dt::Float64)
#    grid .+= dt*velocities
#end

function timeAdvance(hdsim::HDSim)
    cached = Dict("sound_speeds"=>[calcSoundSpeed(hdsim.eos,
                                                  c.density,
                                                  c.pressure)
                                   for c in hdsim.state.cells],
                   "energies"=>[calcSpecificThermalEnergy(hdsim.eos,
                                                          c.density,
                                                          c.pressure)
                                for c in hdsim.state.cells])
    dt = calcTimeStep(hdsim.tsf, hdsim.state, hdsim.eos, cached)
    edge_velocities = calcEdgeVelocities(hdsim.vm, hdsim.state)
    fluxes = calcFluxes(hdsim.fc,
                        hdsim.state,
                        edge_velocities,
                        hdsim.eos,
                        cached)
    updateConserved!(hdsim.extensives,
                     fluxes,
                     hdsim.state.grid,
                     hdsim.pg,
                     dt)
    updateSourceContribution!(hdsim.st,
                              hdsim.extensives,
                              hdsim.state,
                              hdsim.time,
                              dt)
    #updatePositions!(hdsim.state.grid,
    #                 edge_velocities,
    #                 dt)
    #hdsim.state.grid .+= dt*edge_velocities
    for n in 1:length(edge_velocities)
        hdsim.state.grid[n] += dt*edge_velocities[n]
    end

    updateCells!(hdsim.cu,
                 hdsim.state,
                 hdsim.extensives,
                 hdsim.pg,
                 hdsim.eos)

    hdsim.time += dt
    hdsim.cycle += 1
end


function Resistor(; name, R = 1.0)
    # 2 pin component
    @named oneport = OnePort()
    @unpack v, i = oneport
    ps = @parameters R=R
    eqs = [
        v ~ i * R
    ]
    extend(ODESystem(eqs, t, [], ps; name), oneport)
end

function Capacitor(; name, C = 1.0)
    @named oneport = OnePort()
    @unpack v, i = oneport
    ps = @parameters C=C
    D = Differential(t)
    eqs = [
        D(v) ~ i / C
    ]
    extend(ODESystem(eqs, t, [], ps; name), oneport)
end

function Inductor(; name, L = 1.0)
    @named oneport = OnePort()
    @unpack v, i = oneport
    ps = @parameters L=L
    D = Differential(t)
    eqs = [
        D(i) ~ v / L
    ]
    extend(ODESystem(eqs, t, [], ps; name), oneport)
end


## (Conditional) Probability Queries
## =================================
function prob(Y::SymbolicRandVar{Bool};
              partition_alg::Type{BFSPartition} = BFSPartition,
              args...)
  # init_box = unit_box(LazyBox{Float64}, dims(Y))
  partition = pre_partition(Y, partition_alg; args...)
  res = measure(partition)
  AbstractDomains.Interval(max(0.0, res.l), min(1.0, res.u))
end

function prob_conj(XY::SymbolicRandVar{Bool},
              Y::RandVar{Bool};
              partition_alg::Type{BFSPartition} = BFSPartition,
              args...)
  # init_box = unit_box(LazyBox{Float64}, dims(Y))
  Y_partition = pre_partition(Y, partition_alg; args...)
  Y_measure = measure(Y_partition)
  isequal(Y_measure, zero(Y_measure)) && error("Cannot condition on unsatisfiable events")

  XY_partition = pre_partition(XY, partition_alg; args...)
  res = measure(XY_partition) / Y_measure
  AbstractDomains.Interval(max(0.0, res.l), min(1.0, res.u))
end

# "Lower and upper bounds for marginal probability that 'Y' is true"
# function prob(Y::SymbolicRandVar{Bool}; RandVarType = default_randvar(), args...)
#   executable_Y = convert(RandVarType{Bool}, Y)
#   prob(executable_Y; args...)
# end

"Lower/upper bounds for conditional probability that `X` is true given `Y` is true"
function prob(X::SymbolicRandVar{Bool},
              Y::SymbolicRandVar{Bool};
              args...)
  # executable_Y = convert(RandVarType{Bool}, Y)
  # # We do & here because after conversion & may not be defined on
  # # compiled RandVar types
  # executable_XY = convert(RandVarType{Bool}, X & Y)
  prob_conj(X & Y, Y; args...)
end


include("JuliaDe3DFitting.jl")

X = [1.2 1 1; 2 2.3 2; 3 3 3.1; 4 4.1 4; 5 5 5]
Y = planeFit(X)
@show X
@show Y

X = [1 0 0; 0 1 0 ; 0 0 1; -1 0 0; 0 -1 0; 0 0 -1]
Y = sphereFit(X, 100)
@show X
@show Y

X = [1 0 0; 0 1 0 ; -1 0 0; 0 -1 0; 1 0 1; 0 1 1; -1 0 1; 0 -1 1; 1 0 2; 0 1 2; -1 0 2; 0 -1 2]
Y = cylinderFit(X)
@show X
@show Y


using StatsBase
using RDatasets
using Distributions
d = Dirichlet([1.0, 3.0, 5.0])
arr=rand(d)
entropy(arr,2)
entropy(arr)


macro goldensectiontrace()
    quote
        if tracing
            dt = Dict()
            if extended_trace
                dt["x_minimum"] = x_minimum
                dt["x_lower"] = x_lower
                dt["x_upper"] = x_upper
            end
            update!(tr,
                    iteration,
                    f_minimum,
                    NaN,
                    dt,
                    store_trace,
                    show_trace,
                    show_every,
                    callback)
        end
    end
end

immutable GoldenSection <: Optimizer end

function optimize{T <: AbstractFloat}(f::Function, x_lower::T, x_upper::T,
                                      mo::GoldenSection;
                                      rel_tol::T = sqrt(eps(T)),
                                      abs_tol::T = eps(T),
                                      iterations::Integer = 1_000,
                                      store_trace::Bool = false,
                                      show_trace::Bool = false,
                                      callback = nothing,
                                      show_every = 1,
                                      extended_trace::Bool = false,
                                      nargs...)
    if !(x_lower < x_upper)
        error("x_lower must be less than x_upper")
    end

    # Save for later
    initial_lower = x_lower
    initial_upper = x_upper

    const golden_ratio::T = 0.5 * (3.0 - sqrt(5.0))

    x_minimum = x_lower + golden_ratio*(x_upper-x_lower)
    f_minimum = f(x_minimum)
    f_calls = 1 # Number of calls to f

    iteration = 0
    converged = false

    # Trace the history of states visited
    tr = OptimizationTrace{typeof(mo)}()
    tracing = store_trace || show_trace || extended_trace || callback != nothing
    @goldensectiontrace

    while iteration < iterations

        tolx = rel_tol * abs(x_minimum) + abs_tol

        x_midpoint = (x_upper+x_lower)/2

        if abs(x_minimum - x_midpoint) <= 2*tolx - (x_upper-x_lower)/2
            converged = true
            break
        end

        iteration += 1

        if x_upper - x_minimum > x_minimum - x_lower
            x_new = x_minimum + golden_ratio*(x_upper - x_minimum)
            f_new = f(x_new)
            f_calls += 1
            if f_new < f_minimum
                x_lower = x_minimum
                x_minimum = x_new
                f_minimum = f_new
            else
                x_upper = x_new
            end
        else
            x_new = x_minimum - golden_ratio*(x_minimum - x_lower)
            f_new = f(x_new)
            f_calls += 1
            if f_new < f_minimum
                x_upper = x_minimum
                x_minimum = x_new
                f_minimum = f_new
            else
                x_lower = x_new
            end
        end

        @goldensectiontrace
    end

    return UnivariateOptimizationResults("Golden Section Search",
                                         initial_lower,
                                         initial_upper,
                                         x_minimum,
                                         Float64(f_minimum),
                                         iteration,
                                         iteration == iterations,
                                         converged,
                                         rel_tol,
                                         abs_tol,
                                         tr,
                                         f_calls)
end


using BlockArrays, Test

@testset "block product" begin
    @testset "Khatri Rao Product (size)" begin
        #A size
        m = 5
        n = 6

        #A block size
        mi = [3, 2]
        ni = [4, 1, 1]
        #B size
        p = 8
        q = 7

        #B block size
        pi = [5, 3]
        qi = [3, 2, 2]

        @assert sum(mi) == m
        @assert sum(ni) == n
        @assert sum(qi) == q
        @assert sum(pi) == p

        A = BlockArray(ones(m, n), mi, ni)
        B = BlockArray(ones(p, q), pi, qi)

        AB = khatri_rao(A, B)

        @test blocksize(AB) == blocksize(A)
        @test blocksize(AB) == blocksize(B)

        #Test: Size of resulting blocks
        for i in blockaxes(AB,1)
            for j in blockaxes(AB,2)
                @test size(AB[i, j]) == (mi[Int(i)]*pi[Int(i)], ni[Int(j)]*qi[Int(j)])
            end
        end
    end

    @testset "Khatri Rao Product (constant blocks)" begin
        #A size
        m = 5
        n = 6

        #A block size
        mi = [3, 2]
        ni = [4, 1, 1]
        #B size
        p = 8
        q = 7

        #B block size
        pi = [5, 3]
        qi = [3, 2, 2]

        @assert sum(mi) == m
        @assert sum(ni) == n
        @assert sum(qi) == q
        @assert sum(pi) == p

        A = BlockArray(ones(m, n), mi, ni)
        B = BlockArray(ones(p, q), pi, qi)

        #Test: Resulting values for a matrix of constant sub-blocks
        for i in blockaxes(A,1)
            for j in blockaxes(A,2)
                A[i, j] .*= Int(i)+Int(j)
            end
        end

        for i in blockaxes(B,1)
            for j in blockaxes(B,2)
                B[i, j] .*= Int(i)+Int(j)+10
            end
        end

        AB = khatri_rao(A, B)

        for i in blockaxes(AB,1)
            for j in blockaxes(AB,2)
                @test AB[i, j] ≈ (Int(i)+Int(j))*(Int(i)+Int(j)+10)*ones(mi[Int(i)]*pi[Int(i)], ni[Int(j)]*qi[Int(j)])
            end
        end

    end

    @testset "Khatri Rao Product (wrong blocksize)" begin

        A = BlockArray(ones(1, 2), [1], [1,1])
        B = BlockArray(ones(2, 2), [1,1], [1,1])

        @test_throws AssertionError khatri_rao(A, B)
    end

    @testset "Khatri Rao Product (Matrix)" begin

        A = ones(1, 2)
        B = ones(2, 2)

        @test khatri_rao(A, B) ≈ kron(A, B)
    end

    @testset "blockkron" begin
        a = [1,2]
        b = [3,4,5]
        k̃ = BlockKron(a,b) # lazy version of blockkron
        k = blockkron(a,b)
        @test k == k̃ == kron(a,b)
        @test blocksize(k) == blocksize(k̃) == size(a)

        @test k[Block(1)] == k̃[Block(1)] == a[1]*b
        @test k[Block(2)] == k̃[Block(2)] == a[2]*b
        c = 6:8
        k̄ = BlockKron(a,b,c)
        @test k̄ == blockkron(a,b,c) == kron(a,b,c)
        @test k̄[Block(1)][Block(1)] == a[1]*b[1]*c
        @test k̄[Block(1)][Block(2)] == a[1]*b[2]*c
        @test k̄[Block(2)][Block(3)] == a[2]*b[3]*c

        A = randn(2,3)
        B = randn(3,4)
        K = blockkron(A,B)
        K̃ = BlockKron(A,B)
        @test K == K̃ == kron(A,B)
        @test blocksize(K) == blocksize(K̃) == size(A)
        @test K[Block(1,1)] == K̃[Block(1),Block(1)] == A[1,1]*B
        @test K[Block(2,3)] == K̃[Block(2),Block(3)] == A[2,3]*B
        C = randn(2,5)
        K̄ = BlockKron(A,B,C)
        @test K̄ == blockkron(A,B,C) == kron(A,B,C)
        @test K̄[Block(1,1)][Block(1,1)] ≈ A[1,1]*B[1,1]*C
        @test K̄[Block(2,3)][Block(3,4)] ≈ A[2,3]*B[3,4]*C

        @test blockkron(a,B) == kron(a,B)
        @test blockkron(A,b) == kron(A,b)
        @test blockkron(A,b,c) == kron(A,b,c)
        @test blockkron(A,b,C) == kron(A,b,C)

        @test_throws MethodError BlockKron()
        @test_throws MethodError BlockKron(a)
    end
end


"""
    outneighbors_looped(graph, vertex)

Return lists of outneighbors of `vertex` in `graph`. If outdegree of `vertex` is
zero, then [v] is returned.
"""
function outneighbors_looped(graph::AbstractGraph, vertex::Int)
   outns = outneighbors(graph, vertex)
   if outns == []
      return [vertex]
   end
   outns
end

"""
    isstochastic(stochastic)

Checks whether `stochastic` is column-stochastic matrix. Returns nothing.
"""
function stochasticcheck(stochastic::SparseMatrixCSC{<:Number})
   @assert size(stochastic, 1) == size(stochastic, 2) "Matrix is not square"
   @assert all(sum(stochastic[:,i]) ≈ 1 for i=1:size(stochastic, 1)) "Columns sums do not equal one"
   @assert all(findnz(stochastic)[3] .> 0 ) "Matrix elements are not nonnegative"
   nothing
end

"""
    stochastic_preserves_graph_check(graph, stochastic)

Checks whether `stochastic` matrix presreves graph structure. Returns nothing.
"""
function graphpreservationcheck(graph::AbstractGraph,
                                stochastic::SparseMatrixCSC{<:Number})
   @assert size(stochastic, 1) == nv(graph) "Orders of matrix and graph do not match"
   @assert all([ i ∈ outneighbors_looped(graph, j) for (i, j, _)=zip(findnz(stochastic)...)]) "Nonzero elements of stochastic do not coincide with graph edges"
   nothing
end

"""
    graphstochasticcheck(graph, stochastic)

Check, whether `stochastic` is a column-stochatsic matrix preserving `graph`
structure. Three properties are verified:
* `stochastic` is column-stochastic matrix,
* nonzero element of `stochastic` are located on edges of `graph`.

Function returns `nothing` if `stochastic` satisfies properties above, otherwise
raise error.
"""
function graphstochasticcheck(graph::AbstractGraph,
                              stochastic::SparseMatrixCSC{<:Number})
   stochasticcheck(stochastic)
   graphpreservationcheck(graph, stochastic)
   nothing
end

"""
    default_stochastic(graph)

Generates default column-stochastic matrix, which represents random walk with
uniform spreading. If outdegree of vertex is zero, additional 1 is added on the
diagonal.
"""
function default_stochastic(graph::AbstractGraph)
   result = adjacency_matrix(graph, Float64, dir=:in)
   outdegrees = outdegree(graph)
   for i = 1:size(result, 1)
      if outdegrees[i] == 0
         result[i,i] = 1.
      else
         result[:,i] /= outdegrees[i]
      end
   end
   result
end


function test_tied_inner_product_layer(backend::Backend, n_input, T, eps)
  println("-- Testing TiedInnerProductLayer on $(typeof(backend)){$T}...")

  batch_size = abs(rand(Int)) % 10 + 5
  orig_dim = 20
  hidden_dim = 10

  inputs = Array[rand(T, orig_dim, batch_size) for i=1:n_input]
  b1 = rand(T, hidden_dim, 1)
  b2 = rand(T, orig_dim, 1)
  param_key = "shared-ip"

  println("    > Setup")
  registry_reset(backend)
  layer_data = MemoryDataLayer(tops=[symbol("i-$i") for i = 1:n_input],
      batch_size=batch_size, data=inputs)
  layer_ip = InnerProductLayer(name="ip1", param_key=param_key,
      tops=[symbol("ip1-$i") for i=1:n_input], bottoms=[symbol("i-$i") for i=1:n_input],
      output_dim=hidden_dim)
  layer_ip2 = TiedInnerProductLayer(name="ip2", tied_param_key=param_key,
      tops=[symbol("ip2-$i") for i=1:n_input], bottoms=[symbol("ip1-$i") for i=1:n_input])

  net = Net("test-tied-ip", backend, [layer_data,layer_ip,layer_ip2])
  init(net)
  copy!(net.states[2].b, b1)
  copy!(net.states[3].b, b2)
  W = to_array(net.states[2].W)

  println("    > Forward")
  forward(net)
  for i = 1:n_input
    expected_output = W * (W' * inputs[i] .+ b1) .+ b2
    got_output = to_array(net.states[3].blobs[i])
    @test all(abs(expected_output - got_output) .< eps)
  end

  println("    > Backward")
  # This network is invalid for back-propagation because there is
  # no top layer that compute loss
  @test_throws TopologyError check_bp_topology(net)

  # however, we will manually set the diffs for the top-most layer
  # to do a fake back-propagate test
  top_diffs = Array[rand(T, orig_dim, batch_size) for i=1:n_input]
  for i = 1:n_input
    copy!(net.states[3].blobs_diff[i], top_diffs[i])
  end
  backward(net)

  bias_grad = to_array(net.states[3].∇b)
  bias_grad_expected = sum([sum(top_diffs[i],2) for i = 1:n_input])
  @test all(abs(bias_grad - bias_grad_expected) .< eps)

  for i = 1:n_input
    back_grad = to_array(net.states[2].blobs_diff[i])
    back_grad_expected = W' * top_diffs[i]
    @test all(abs(back_grad - back_grad_expected) .< eps)
  end

  destroy(net)
end

function test_tied_inner_product_layer(backend::Backend)
  test_tied_inner_product_layer(backend, 3, Float32, 1e-3)
  test_tied_inner_product_layer(backend, 3, Float64, 1e-10)
end

if test_cpu
  test_tied_inner_product_layer(backend_cpu)
end
if test_gpu
  test_tied_inner_product_layer(backend_gpu)
end


#md # # Plot recipes examples
#nb # # BAT.jl Plot recipes examples
#jl # # BAT.jl Plot recipes examples

using BAT

# ## Generate samples
# We generate some generic multimodal samples to be plotted afterwards
using Distributions, IntervalSets

##Model definition to generate samples
struct GaussianShellDensity<:AbstractDensity
    r::Vector{Float64}
    sigma::Vector{Float64}
end

#define likelihood
function BAT.density_logval(target::GaussianShellDensity, params::Any)
    loglikelihood::Float64 = 0.
    for i in 1:length(params)

        result = -0.5 * (params[i][1]-target.r[i])^2/target.sigma[i]^2
        l1 = result - log(√2π * target.sigma[i])

        result2 = -0.5 * (params[i][1]+target.r[i])^2/target.sigma[i]^2
        l2 = result2 - log(√2π * target.sigma[i])^2

        loglikelihood += log(exp(l1) + 2*exp(l2))
    end

    return loglikelihood
end

likelihood = GaussianShellDensity([5.0, 5., 3.], [2, 2.4, 1.5])

prior = NamedTupleDist(
    λ1 = Normal(-3, 4.5),
    λ2 = -30.0..30.0,
    λ3 = Normal(3, 3.5)
)

posterior = PosteriorDensity(likelihood, prior);

#generate samples
samples, stats, chains = bat_sample(posterior, (10^5, 4), MetropolisHastings())

# ## Set up plotting
#  We set up plotting by using the ["Plots"](https://github.com/JuliaPlots/Plots.jl) package with the pyplot backend.
using Plots

# The plotting recipes have been tested with `pyplot`, `gr` and `plotly` backends.
# All available plotting options work correctly with `pyplot`.
# Only the colored 2D contour plots are currently not correctly supported with `gr` or `plotly` backends.

pyplot()

# ## Examples for 1D plots
# Below, all available seriestypes and plotting features for 1D representations of samples or prior are shown
# ### default 1D plot style for samples:
plot(samples, 2)
# The default seriestype for samples is `:smallest_intervals`.
# By default, the local mode of the histogram is indicated as a dotted black line.

# ### default 1D plot style for prior:
plot(prior, 2)
# The default seriestype for sprior is `:stephist`.

# Samples can either be plotted by their index (as shown above) or by using the parameter names given in NamedTupleDist. This can be done by passing either the posterior
plot(posterior, samples, :λ2)
# or the value shape(s) of the prior:
using ValueShapes
parshapes = valshape(prior)
plot(parshapes, samples, :λ2)

# Prior can also be plotted by their index or by using the parameter names given in NamedTupleDist:
plot(prior, 1)
# or
plot(prior, :λ1)

# Plot prior + posterior
plot(posterior, samples, :λ1)
plot!(prior, :λ1)

# ### Further seriestypes available:
# #### central intervals
plot(samples, 2, seriestype = :central_intervals)

# ### Customizing smallest/central interval plots:
# The probability intervals to be highlighted can be specified using the `intervals` keyword. Their colors (in same order) need to be specified using the `colors` keyword argument.
plot(samples, 2, seriestype=:smallest_intervals, intervals=[0.5, 0.1, 0.3, 0.99], colors=[:grey, :red, :blue, :orange])

# #### histogram
plot(samples, 2, seriestype = :histogram)
# alias: `:hist == :histogram`

# #### step histogram
plot(samples, 2, seriestype = :stephist)

# ### Keyword arguments for [attributes supported by *Plots.jl*](https://docs.juliaplots.org/latest/attributes/#magic-arguments) can be passed:
plot(samples, 2, seriestype = :stephist, nbins=50, linecolor = :red, linewidth = 5, linealpha=0.3, xlim=(0,10))

# ### Plotting point estimators in 1D plots:
plot(samples, 2, globalmode=true, localmode=true, mean=true, std_dev=true)

# #### It is possible to customize the style of estimators by passing a dict with the respective attributes:
# By passing a boolean, the point estimators are plotted using their default styles shown above.
# The style of the point estimators can be modified by passing a dictionary specifying `linestyle`, `linecolor`, `linewidth` and `alpha` for *mean*, *globalmode* and *localmode*.
# The style of the standard deviation can be modified by specifying `fillcolor` and `fillalpha`.
plot(samples, 1, localmode=false, mean=Dict("linestyle" => :dot, "linecolor"=> :red, "linewidth"=>2, "alpha" => 0.7), std_dev=Dict("fillcolor" => :red, "fillalpha"=>0.2))

# ## Examples for 2D plots of the samples
# Below, all available seriestypes and plotting features for 2D representations of samples are shown
# ### default 2D plot style:
plot(samples, (1,2), nbins=200)
# The default 2D plotting style is a 3-color heatmap showing the smallest intervals containing 68.3, 95.5 and 99.7 perecent of the posterior probability.

# Samples can either be plotted by their index (as shown above) or by using the parameter names given in NamedTupleDist. This can be done by passing either the posterior
plot(posterior, samples, (:λ1, :λ2))
# or the value shape(s) of the prior:
parshapes = valshape(prior)
plot(parshapes, samples, (:λ1, :λ2))

# Prior can also be plotted by their index or by using the parameter names given in NamedTupleDist.
plot(prior, (1, 2))
# or
plot(prior, (:λ1, :λ2))

# Plot prior + posterior
plot(posterior, samples, (:λ1, :λ3))
plot!(prior, (:λ1, :λ3))

# ### Further seriestypes available:
# #### central intervals
plot(samples, (1,2), seriestype = :histogram, nbins=200)
# alias: `:histogram2d == :histogram`

# #### smallest intervals as colored contour lines:
# (currently only correctly supported with `pyplot()` backend)
plot(samples, (1,2), seriestype=:smallest_intervals_contour, bins=50)

# #### smallest intervals as filled contours:
# (currently only correctly supported with `pyplot()` backend)
plot(samples, (1,2), seriestype=:smallest_intervals_contourf, bins=50)

# ### Customizing smallest interval plots:
# The probability intervals to be highlighted can be specified using the `intervals` keyword. Their colors (in same order) need to be specified using the `colors` keyword argument.
plot(samples, (1,2), seriestype=:smallest_intervals, nbins=200, intervals=[0.7, 0.2], colors=[:blue, :red])

# #### scatter
# for large number of samples, this plotting style needs a lot of time and resources
#plot(samples, (1,2), seriestype = :scatter)

# ### Plotting point estimators in 2D plots:
plot(samples, (1,2), seriestype=:smallest_intervals, nbins=200, mean=true, std_dev=true, localmode=true, globalmode=true)

# #### It is possible to customize the style of estimators by passing a dict with the respective attributes:
# By passing a boolean, the point estimators are plotted using their default styles shown above.
# The style of the point estimators can be modified by passing a dictionary specifying `markershape`, `markercolor`, `markersize`, `markeralpha`, `markerstrokecolor`, `markerstrokestyle`, `markerstrokewidth` and `markerstrokealpha` for *mean*, *globalmode* and *localmode*.
# If `std_dev==true`, the standard deviation of the mean value will be displayed as x- and y-errorbars.
plot(samples, (1,2), seriestype=:smallest_intervals, nbins=200, localmode=Dict("markershape"=> :diamond, "markeralpha"=>1, "markercolor"=>:red, "markersize"=>5))

# ### Marginal plots
plot(samples, (1,2), seriestype = :marginal)

# ### Customizing marginal plots:
# The marginal plots can be modified by passing dictionaries to the keyword arguments `upper`, `right` and `diagonal`.
# The dictionaries for `upper` and `right` can contain the 1D seriestypes and plot options shown above.
# The dictionary for `diagonal` can use the 2D seriestypes and plot options shown above.
# Nested dictonaries are possible (e.g. for modifying point estimators)
plot(samples, (1,2), seriestype = :marginal, diagonal = Dict("seriestype"=>:histogram, "mean"=>Dict("markercolor"=>:green)), upper=Dict("mean" => true, "seriestype" => :smallest_intervals, "colors"=>[:blue, :grey, :orange]), right=Dict("seriestype" => :stephist, "mean"=>true))

# ### Overview plot
plot(samples)
#
plot(prior)
#
plot(samples)
plot!(prior)

# By default the 1D and 2D plots for the first 5 parameters are shown.
# The keyword argument `params` allows to specify the considered parameters.
plot(samples, params=[1, 3])

# ### Customizing overview plots:
# The overview plots can be modified by passing dictionaries to the keyword arguments `upper`, `lower` and `diagonal`.
# The dictionaries for `upper` and `lower` can contain the 2D seriestypes and plot options shown above.
# The dictionary for `diagonal` can use the 1D seriestypes and plot options shown above.
# Nested dictonaries are possible (e.g. for modifying point estimators)

plot(samples, mean=true, globalmode=true, legend=true, diagonal=Dict("seriestype"=>:stephist, "mean"=>Dict("linecolor" => :green, "linewidth" => 8)), lower = Dict("mean" => false, "colors"=>[:orange, :green, :grey]))

# ## Plots for MCMC diagnostics
# Plots histograms of the samples, the trace, a kernel density estimate and the autocorrelation function for each parameter per chain.
diagnostics = BAT.MCMCDiagnostics(samples, chains)
plot(diagnostics, params=[1])

# ### Customizing diagnostics plots:
plot(diagnostics,
    params=[1, 2],
    chains=[1, 2],
    diagnostics = [:histogram, :kde, :trace, :acf],
    histogram = Dict("seriestype" => :smallest_intervals, "legend" => :false),
    trace = Dict("linecolor" => :red),
    acf = Dict("lags" => collect(1:20), "title"=> "Autocorrelation"),
    description = true)

# ### available keyword arguments:
# * `params` - list of parameters to be plotted
# * `chains` - list of chains to be plotted
# * `diagnostics` - list of MCMC diagnostics to be plotted
# * `:histogram` - 1D histograms of samples
# * `:kde` - Kernel density estimate (using [*KernelDensity.jl*](https://github.com/JuliaStats/KernelDensity.jl))
# * `:trace` - Trace plot
# * `:acf` - Autocorrelation function (using [*StatsBase.autocor*](http://juliastats.github.io/StatsBase.jl/stable/signalcorr/#StatsBase.autocor))
# * `description::Bool = true` - show description (current chain, parameter, number of samples) as first column of plots
# * `histogram::Dict` - options for histogram plots (supports all arguments for 1D plots for samples)
# * `kde::Dict` - options for kde plots
# * `trace::Dict` - options for trace plots
# * `acf::Dict` - options for acf plots

# ### special options arguments for `:kde` (see [*KernelDensity.jl*](https://github.com/JuliaStats/KernelDensity.jl))
# * `npoints`: number of interpolation points to use (default: npoints = 2048)
# * `boundary`: lower and upper limits of the kde as a tuple
# * `kernel`: the distributional family from [*Distributions.jl*](https://github.com/JuliaStats/Distributions.jl) to use as the kernel (default = Distributions.Normal)
# * `bandwidth`: bandwidth of the kernel

# ### special keyword arguments for `:acf` (see [*StatsBase.autocor*](http://juliastats.github.io/StatsBase.jl/stable/signalcorr/#StatsBase.autocor))
# * `lags` - list of lags to be considered for ACF plots
# * `demean` - denotes whether the mean should be subtracted before computing the ACF


"""
    encode_motif(x, m::Int = length(x)) → s::Int

Encode the length-`m` motif `x` (a vector of indices that would sort some vector `v`
in ascending order) into its unique integer symbol ``s \\in \\{1, 2, \\ldots, m - 1 \\}``,
using Algorithm 1 in Berger et al. (2019)[^Berger2019].

## Example

```julia
v = rand(5)

# The indices that would sort `v` in ascending order. This is now a permutation
# of the index permutation (1, 2, ..., 5)
x = sortperm(v)

# Encode this permutation as an integer.
encode_motif(x)
```
[^Berger2019]: Berger, Sebastian, et al. "Teaching Ordinal Patterns to a Computer: Efficient Encoding Algorithms Based on the Lehmer Code." Entropy 21.10 (2019): 1023.
"""
function encode_motif(x, m::Int = length(x))
    n = 0
    for i = 1:m-1
        for j = i+1:m
            n += x[i] > x[j] ? 1 : 0
        end
        n = (m-i)*n
    end

    return n
end

function isless_rand(a, b)
    if a == b
        rand(Bool)
    elseif a < b
        true
    else
        false
    end
end

""" Compute probabilities of symbols `Π`, given weights `wts`. """
function symprobs(Π::AbstractVector, wts::AbstractVector; normalize = true)
    length(Π) == length(wts) || error("Need length(Π) == length(wts)")
    N = length(Π)
    idxs = sortperm(Π)
    sΠ = Π[idxs]   # sorted symbols
    sw = wts[idxs] # sorted weights

    i = 1   # symbol counter
    W = 0.0 # Initialize weight
    ps = Float64[]

    prev_sym = sΠ[1]

    while i <= length(sΠ)
        symᵢ = sΠ[i]
        wtᵢ = sw[i]
        if symᵢ == prev_sym
            W += wtᵢ
        else
            # Finished counting weights for the previous symbol, so push
            # the summed weights (normalization happens later).
            push!(ps, W)

            # We are at a new symbol, so refresh sum with the first weight
            # of the new symbol.
            W = wtᵢ
        end
        prev_sym = symᵢ
        i += 1
    end
    push!(ps, W) # last entry

    # Normalize
    Σ = sum(sw)
    if normalize
        return ps ./ Σ
    else
        return ps
    end
end


using DiffEqFlux, OrdinaryDiffEq, DiffEqSensitivity
using CUDA, Test, Zygote, Random, LinearAlgebra

CUDA.allowscalar(false)

H = CuArray(rand(Float32, 2, 2))
ann = FastChain(FastDense(1, 4, tanh))
p = initial_params(ann)

function func(x, p, t)
    ann([t],p)[1]*H*x
end

x0 = CuArray(rand(Float32, 2))
x1 = CuArray(rand(Float32, 2))

prob = ODEProblem(func, x0, (0.0f0, 1.0f0))

function evolve(p)
    solve(prob, Tsit5(), p=p, save_start=false,
          save_everystep=false, abstol=1e-4, reltol=1e-4,
          sensealg=QuadratureAdjoint(autojacvec=ZygoteVJP())).u[1]
end

function cost(p)
    x = evolve(p)
    c = sum(abs,x - x1)
    #println(c)
    c
end

grad = Zygote.gradient(cost,p)[1]
@test !iszero(grad[1])
@test iszero(grad[2:4])
@test !iszero(grad[5])
@test iszero(grad[6:end])

###
# https://github.com/SciML/DiffEqFlux.jl/issues/632
###

rng = MersenneTwister(1234)
m = 32
n = 16
Z = randn(rng, Float32, (n,m)) |> gpu
𝒯 = 2.0
Δτ = 0.1
ca_init = [zeros(1) ; ones(m)] |> gpu

function f(ca, Z, t)
  a = ca[2:end]

  a_unit = a / sum(a)
  w_unit = Z*a_unit
  Ka_unit = Z'*w_unit
  z_unit = dot(abs.(Ka_unit), a_unit)
  aKa_over_z = a .* Ka_unit / z_unit
  [sum(aKa_over_z) / m; -abs.(aKa_over_z)] |> gpu
end

function c(Z)
  prob = ODEProblem(f, ca_init, (0.,𝒯), Z, saveat=Δτ)
  sol = solve(prob, Tsit5(), sensealg=BacksolveAdjoint(), saveat=Δτ)
  sum(last(sol.u))
end

println("forward:", c(Z))
println("backward: ", Zygote.gradient(c, Z))


module PhasePlots

using RecipesBase, ColorTypes

@userplot PhasePlot

@recipe function f(p::PhasePlot; clim=nothing, γ=2)
    x, y, z = p.args
    isnothing(clim) && (clim = maximum(abs, z))
    extend(x, n) = (h = (x[end]-x[begin])/2(n-1); [x[begin]-h, x[end]+h])
    x = extend(x, size(z,1))
    xlims --> x
    y = extend(y, size(z,2))
    ylims --> y
    # The docs say :yflip defaults to false, Plots.default agrees, but the plot with yflip omitted is the same as when it is set to true.
    xflip --> !issorted(x)
    yflip --> !issorted(y)
    aspect_ratio --> 1
    grid --> false
    tick_direction --> :out

    N = transpose(z)
    x, y, @. HSVA(360/2π*angle(N), 1, !iszero(N), (abs(N)/clim)^γ)
end

@doc """
    phaseplot(x, y, z; clim=maximum(abs.(z)))

Plot a complex function with phase as hue

The fourth roots of unity are red, green, cyan and purple.  Modulus is mapped to the alpha channel, so the classic complex phase plot is `phaseplot(..., bg=:black)`.  Alpha becomes opaque at the modulus clim; this is useful when multiple plots should have the same brightness scale.

Plots.jl is dreadfully inconsistent at mapping points in the plane to raster indices.  For consistency, phaseplot has a different default to every existing function.  The first dimension is the real/horizontal axis, the second the imaginary/vertical one, and by default x increases rightward and y upward.

When Plots.jl draws a raster, the extrema of x and y are the coordinates of points half a pixel outside the edge of the grid.  This function moves them to the edges, which is almost always what you will want.
""" phaseplot

end # module


"""
GTH algorithm
"""

"""
gth!(Q::Matrix{Tv})
gth(Q::Matrix{Tv})

Compute the stationary vector for the CTMC with kernel Q with GTH algorithm.
Both of gth! and gth return a vector.
In gth!, the matrix Q is used as a workspace.
Note that Q does not have any absorbing state.
"""

function gth!(Q::Matrix{Tv})::Vector{Tv} where {Tv}
    @inbounds begin
        m, n = size(Q)
        @assert m == n
        for l = n:-1:2
            tmp::Tv = 0
            for u = 1:l-1
                tmp += Q[l,u]
            end
            for j = 1:l-1
                for i = 1:l-1
                    if i != j
                        Q[i,j] += Q[l,j] * Q[i,l] / tmp
                    end
                end
            end
            for i = 1:l-1
                Q[i,l] /= tmp
            end
            for i = 1:l-1
                Q[l,i] = 0
            end
            Q[l,l] = -1
        end
        x = Vector{Tv}(undef, n)
        x[1] = 1.0
        for l = 2:n
            x[l] = 0.0
            for i = 1:l-1
                x[l] += x[i] * Q[i,l]
            end
        end
        x /= sum(x)
    end
end

function gth(Q::Matrix{Tv}) where {Tv}
    gth!(copy(Q))
end

"""
gth!(Q::Matrix{Tv}, index::Vector{Ti})
gth(Q::Matrix{Tv}, index::Vector{Ti})

Compute the stationary vector for the CTMC with kernel Q with GTH algorithm.
Both of gth! and gth return a vector.
The CTMC states are permutated with an index vector.
In gth!, the matrix Q is used as a workspace.
Note that Q does not have any absorbing state.
"""

function gth!(Q::Matrix{Tv}, index::Vector{Ti})::Vector{Tv} where {Tv,Ti}
    m, n = size(Q)
    @assert m == n
    for l = n:-1:2
        tmp::Tv = 0
        for u = 1:l-1
            tmp += Q[index[l],index[u]]
        end
        for j = 1:l-1
            for i = 1:l-1
                if i != j
                    Q[index[i],index[j]] += Q[index[l],index[j]] * Q[index[i],index[l]] / tmp
                end
            end
        end
        for i = 1:l-1
            Q[index[i],index[l]] /= tmp
        end
        for i = 1:l-1
            Q[index[l],index[i]] = 0
        end
        Q[index[l],index[l]] = -1
    end
    x = Vector{Tv}(undef, n)
    x[index[1]] = 1.0
    for l = 2:n
        x[index[l]] = 0.0
        for i = 1:l-1
            x[index[l]] += x[index[i]] * Q[index[i],index[l]]
        end
    end
    x /= sum(x)
end

function gth(Q::Matrix{Tv}, index::Vector{Ti}) where {Tv, Ti}
    gth!(copy(Q), index)
end


using Surrogates
using Distributions
using Test

#1D
lb = 0.0
ub = 5.0
n = 5
d = 1
Surrogates.sample(n,lb,ub,GridSample(0.1))
Surrogates.sample(n,lb,ub,UniformSample())
Surrogates.sample(n,lb,ub,SobolSample())
Surrogates.sample(n,lb,ub,LatinHypercubeSample())
Surrogates.sample(20,lb,ub,LowDiscrepancySample(10))
Surrogates.sample(5,d,Cauchy())
Surrogates.sample(5,d,Normal(0,4))

#ND
lb = [0.1,-0.5]
ub = [1.0,20.0]
n = 5
d = 2

#GridSample{T}
s = Surrogates.sample(n,lb,ub,GridSample([0.1,0.5]))
@test isa(s,Array{Tuple{typeof(s[1][1]),typeof(s[1][1])},1}) == true

#UniformSample()
s = Surrogates.sample(n,lb,ub,UniformSample())
@test isa(s,Array{Tuple{typeof(s[1][1]),typeof(s[1][1])},1}) == true

#SobolSample()
s = Surrogates.sample(n,lb,ub,SobolSample())
@test isa(s,Array{Tuple{typeof(s[1][1]),typeof(s[1][1])},1}) == true

#LHS
s = Surrogates.sample(n,lb,ub,LatinHypercubeSample())
@test isa(s,Array{Tuple{typeof(s[1][1]),typeof(s[1][1])},1}) == true

#LDS
s = Surrogates.sample(n,lb,ub,LowDiscrepancySample([10,3]))
@test isa(s,Array{Tuple{typeof(s[1][1]),typeof(s[1][1])},1}) == true

#Distribution 1
s = Surrogates.sample(n,d,Cauchy())
@test isa(s,Array{Tuple{typeof(s[1][1]),typeof(s[1][1])},1}) == true

#Distribution 2
s = Surrogates.sample(n,d,Normal(3,5))
@test isa(s,Array{Tuple{typeof(s[1][1]),typeof(s[1][1])},1}) == true


using Random

"""

Arguments:
    - `f`: Function to be optimized
    - `g`: Gradient function for `f`
    - `c`: Constraints function. Evaluates all the constraint equations and return a Vector of values
    - `x0`: (Vector) Initial position to start from
    - `n`: (Int) Number of evaluations allowed. Remember `g` costs twice of `f`
    - `prob`: (String) Name of the problem. So you can use a different strategy for each problem
"""
function optimize(f, g, c, x0, n, prob)
    p = squared_sum_penalty_function(c)
    x_best, xlist = penalty_method(f, p, x0, n)
    return x_best
end

function sum_penalty_function(c)
    f = function(x)
        ceval = c(x)
        sum(ceval[ceval .> 0])
    end
    return f
end

function squared_sum_penalty_function(c)
    f = function(x)
        ceval = c(x)
        sum(ceval[ceval .> 0].^2)
    end
    return f
end

function penalty_method(f, p, x, n; ρ=1, γ=2)
    numEvalsLeft = n
    hj_max_eval = 10;
    xlist = Array{Float64}[]
    push!(xlist, x)
    while numEvalsLeft >= hj_max_eval
        x, evalsUsed, feasFound = hooke_jeeves(f, ρ, p, x, hj_max_eval)
        push!(xlist, x)
        numEvalsLeft = numEvalsLeft - evalsUsed
        ρ *= γ
    end
    return x,xlist
end

function last_feasible_x(currXFeas, x, feasFound, last_x_feas)
    if currXFeas
        return x
    else
        if !feasFound
            return x
        else
            return last_x_feas
        end
    end
end

basis(i, n) = [k == i ? 1.0 : 0.0 for k in 1 : n]

function hooke_jeeves(f, ρ, p, x, maxEval; α=1, γ=0.5)
    xf = f(x)
    xp = p(x)
    y, m = xf + ρ*xp, length(x)
    evalsUsed = 1
    last_x_feas = x
    if xp==0
        feasFound = true
        currXFeas = true
    else
        feasFound = false
        currXFeas = false
    end
    while evalsUsed < maxEval
        improved = false
        for i in 1 : m
            x′ = x + α*basis(i, m)
            if evalsUsed >= maxEval
                return last_feasible_x(currXFeas, x, feasFound, last_x_feas), evalsUsed, feasFound
            end
            xf′ = f(x′)
            xp′ = p(x′)
            y′ = xf′ + ρ*xp′
            evalsUsed = evalsUsed + 1
            if y′ < y
                x, y, improved = x′, y′, true
                if xp′==0
                    last_x_feas = x
                    feasFound = true
                    currXFeas = true
                else
                    currXFeas = false
                end
            else
                x′ = x - α*basis(i, m)
                if evalsUsed >= maxEval
                    return last_feasible_x(currXFeas, x, feasFound, last_x_feas), evalsUsed, feasFound
                end
                xf′ = f(x′)
                xp′ = p(x′)
                y′ = xf′ + ρ*xp′
                evalsUsed = evalsUsed + 1
                if y′ < y
                    x, y, improved = x′, y′, true
                    if xp′==0
                        last_x_feas = x
                        feasFound = true
                        currXFeas = true
                    else
                        currXFeas = false
                    end
                end
            end
        end
        if !improved
            α *= γ
        end
    end
    return last_x_feas, evalsUsed, feasFound
end


"Any radial basis function."
abstract type RBF{T}
end

"A smooth radial basis function with a shape parameter `ε`."
abstract type SmoothRBF{T} <: RBF{T}
end

shape_parameter(rbf::SmoothRBF) = rbf.ε

"A piecewise smooth radial basis function."
abstract type PiecewiseSmoothRBF{T} <: RBF{T}
end

degree(rbf::PiecewiseSmoothRBF) = rbf.p

struct Multiquadric{T} <: SmoothRBF{T}
    ε   :: T
end

const MQ{T} = Multiquadric{T}

(rbf::Multiquadric)(r) = sqrt(1 + (rbf.ε*r)^2)

eval_derivative1(rbf::Multiquadric, r) = 1/rbf(r)*rbf.ε*r

similar(r::Multiquadric{T}, ε) where {T} = Multiquadric{T}(ε)

struct InverseMultiquadric{T} <: SmoothRBF{T}
    ε   :: T
end

const IMQ{T} = InverseMultiquadric{T}

(rbf::InverseMultiquadric)(r) = 1/sqrt(1 + (rbf.ε*r)^2)

eval_derivative1(rbf::InverseMultiquadric{T}, r) where {T} = (-one(T)/2)*rbf(r)^3*2*rbf.ε*r

similar(r::InverseMultiquadric{T}, ε) where {T} = InverseMultiquadric{T}(ε)

struct InverseQuadratic{T} <: SmoothRBF{T}
    ε   :: T
end

const IQ{T} = InverseQuadratic{T}

(rbf::InverseQuadratic)(r) = 1/(1 + (rbf.ε*r)^2)

similar(r::InverseQuadratic{T}, ε) where {T} = InverseQuadratic{T}(ε)

struct Gaussian{T} <: SmoothRBF{T}
    ε   :: T
end

const GA{T} = Gaussian{T}

(rbf::Gaussian)(r) = exp(-(rbf.ε*r)^2)

eval_derivative1(rbf::Gaussian, r) = -2*rbf.ε*r * rbf(r)

eval_derivative2(rbf::Gaussian, r) = -2*rbf.ε * rbf(r) - 2*rbf.ε*r*eval_derivative1(rbf,r)

similar(r::Gaussian{T}, ε) where {T} = Gaussian{T}(ε)

approximate_support(rbf::Gaussian{T}, threshold = eps(T)) where {T} = sqrt(-log(threshold))/rbf.ε

struct PolyharmonicSpline{T} <: PiecewiseSmoothRBF{T}
    p   ::  Int
end

const SPH{T} = PolyharmonicSpline{T}

PolyharmonicSpline(p::Int) = PolyharmonicSpline{Float64}(p)

(rbf::PolyharmonicSpline)(r) = r^(2*rbf.p+1)

similar(r::PolyharmonicSpline{T}, p) where {T} = PolyharmonicSpline{T}(p)

struct ThinPlateSpline{T} <: PiecewiseSmoothRBF{T}
    p   ::  Int
end

const TPS{T} = ThinPlateSpline{T}

ThinPlateSpline(p::Int) = ThinPlateSpline{Float64}(p)

(rbf::ThinPlateSpline)(r) = r^(2*rbf.p) * log(r)

similar(r::ThinPlateSpline{T}, p) where {T} = ThinPlateSpline{T}(p)


# Poly type manipulations

module Polynomial
#todo: sparse polynomials?

export Poly, polyval, polyint, polyder, poly, roots, degree
export polydir #Deprecated

import Base: length, endof, getindex, setindex!, copy, zero, one, convert
import Base: show, print, *, /, //, -, +, ==, divrem, rem, eltype

eps{T}(::Type{T}) = convert(T,0)
eps{F<:FloatingPoint}(x::Type{F}) = Base.eps(F)
eps{T}(x::Type{Complex{T}}) = eps(T)

immutable Poly{T<:Number}
    a::Vector{T}
    nzfirst::Int #for effiencicy, track the first non-zero index
    var::Symbol
    function Poly(a::Vector{T}, var)
        nzfirst = 0 #find and chop leading zeros
        for i = 1:length(a)
            if abs(a[i]) > 2*eps(T)
                break
            end
            nzfirst = i
        end
        new(a, nzfirst, symbol(var))
    end
end

Poly{T<:Number}(a::Vector{T}, var::Union(String,Symbol,Char)=:x) = Poly{T}(a, var)

convert{T}(::Type{Poly{T}}, p::Poly) = Poly(convert(Vector{T}, p.a), p.var)
promote_rule{T, S}(::Type{Poly{T}}, ::Type{Poly{S}}) = Poly{promote_type(T, S)}
eltype{T}(::Poly{T}) = T

length(p::Poly) = length(p.a)-p.nzfirst
endof(p::Poly) = length(p)
deg(p::Poly) = length(p) - 1

getindex(p::Poly, i) = p.a[i+p.nzfirst]
setindex!(p::Poly, v, i) = (p.a[i+p.nzfirst] = v)

copy(p::Poly) = Poly(copy(p.a[1+p.nzfirst:end]), p.var)

zero{T}(p::Poly{T}) = Poly([zero(T)], p.var)
zero{T}(::Type{Poly{T}}) = Poly([zero(T)])
one{T}(p::Poly{T}) = Poly([one(T)], p.var)
one{T}(::Type{Poly{T}}) = Poly([one(T)])

function show(io::IO, p::Poly)
    print(io,"Poly(")
    print(io,p)
    print(io,")")
end

function print{T<:Real}(io::IO, p::Poly{T})
    n = length(p)
    if n <= 0
        print(io,"0")
    else
        for j = 1:n
            pj = p[j]
            magpj = abs(pj)
            if magpj > 2*eps(T)
                if j == 1
                    pj < 0 && print(io, "-")    #Prepend - if first and negative
                else
                    pj < 0 ? print(io," - ") : print(" + ")
                end
                #Print pj if pj is the last coefficient, or pj is not identically 1
                if j == n || abs(magpj - 1) > 2*eps(T)
                    print(io, magpj)
                end
                exp = n-j
                if exp > 0
                    print(io, p.var)
                    if exp > 1
                        print(io, '^', exp)
                    end
                end
            end
        end
    end
end

function print{T<:ComplexPair}(io::IO, p::Poly{T})
    n = length(p)
    if n <= 0
        print(io,"0")
    else
        for j = 1:n
            pj = p[j]
            abs_repj = abs(real(pj))
            abs_impj = abs(imag(pj))
            if abs(pj) > 2*eps(T)
                if !(abs_impj > 2*eps(T))
                    if j > 1
                        real(pj) < 0 ? print(io," - ") : print(" + ")
                    else
                        real(pj) < 0 && print(io, "-")    #Prepend - if first and negative
                    end
                else
                    j > 1 && print(io, " + ")
                end
                if abs_repj > 2*eps(T)    #Real part is not 0
                    if abs_impj > 2*eps(T)    #Imag part is not 0
                        print(io,'(',pj,')')
                    elseif abs(abs_repj - 1) > 2*eps(T) || j == n
                        print(io,abs_repj)
                    end
                else
                    if abs_impj > 2*eps(T)
                        print(io,'(', imag(pj),"im)")
                    end
                end
                exp = n-j
                if exp > 0
                    print(io, p.var)
                    if exp > 1
                        print(io, '^', exp)
                    end
                end
            end
        end
    end
end

*(c::Number, p::Poly) = Poly(c * p.a[1+p.nzfirst:end], p.var)
*(p::Poly, c::Number) = Poly(c * p.a[1+p.nzfirst:end], p.var)
/(p::Poly, c::Number) = Poly(p.a[1+p.nzfirst:end] / c, p.var)
-(p::Poly) = Poly(-p.a[1+p.nzfirst:end], p.var)

-(p::Poly, c::Number) = +(p, -c)
+(c::Number, p::Poly) = +(p, c)
function +(p::Poly, c::Number)
    if length(p) < 1
        return Poly([c,], p.var)
    else
        p2 = copy(p)
        p2.a[end] += c;
        return p2;
    end
end
function -(c::Number, p::Poly)
    if length(p) < 1
        return Poly([c,], p.var)
    else
        p2 = -p;
        p2.a[end] += c;
        return p2;
    end
end

function +{T,S}(p1::Poly{T}, p2::Poly{S})
    if p1.var != p2.var
        error("Polynomials must have same variable")
    end
    R = promote_type(T,S)
    n = length(p1)
    m = length(p2)
    if n > m
        a = Array(R, n)
        for i = 1:m
            a[n-m+i] = p1[n-m+i] + p2[i]
        end
        for i = 1:n-m
            a[i] = p1[i]
        end
    else
        a = Array(R, m)
        for i = 1:n
            a[m-n+i] = p1[i] + p2[m-n+i]
        end
        for i = 1:m-n
            a[i] = p2[i]
        end
    end
    Poly(a, p1.var)
end

function -{T,S}(p1::Poly{T}, p2::Poly{S})
    if p1.var != p2.var
        error("Polynomials must have same variable")
    end
    R = promote_type(T,S)
    n = length(p1)
    m = length(p2)
    if n > m
        a = Array(R, n)
        for i = 1:m
            a[n-m+i] = p1[n-m+i] - p2[i]
        end
        for i = 1:n-m
            a[i] = p1[i]
        end
    else
        a = Array(R, m)
        for i = 1:n
            a[m-n+i] = p1[i] - p2[m-n+i]
        end
        for i = 1:m-n
            a[i] = -p2[i]
        end
    end
    Poly(a, p1.var)
end

function *{T,S}(p1::Poly{T}, p2::Poly{S})
    if p1.var != p2.var
        error("Polynomials must have same variable")
    end
    R = promote_type(T,S)
    n = length(p1)
    m = length(p2)
    if n == 0 || m == 0
        return Poly(R[], p1.var)
    end
    a = zeros(R, n+m-1)
    for i = 1:length(p1)
        for j = 1:length(p2)
            a[i+j-1] += p1[i] * p2[j]
        end
    end
    Poly(a, p1.var)
end

function divrem{T, S}(num::Poly{T}, den::Poly{S})
    if num.var != den.var
        error("Polynomials must have same variable")
    end
    m = length(den)
    if m == 0
        throw(DivideError())
    end
    R = typeof(one(T)/one(S))
    n = length(num)
    deg = n-m+1
    if deg <= 0
        return zero(Poly{R}), convert(Poly{R}, num)
    end
    d = zeros(R, n)
    q = zeros(R, deg)
    r = zeros(R, n)
    r[:] = num.a[1+num.nzfirst:end]
    for i = 1:deg
        quot = r[i] / den[1]
        q[i] = quot
        if i > 1
            d[i-1] = 0
            r[i-1] = 0
        end
        for j = 1:m
            k = i+j-1
            elem = den[j]*quot
            d[k] = elem
            r[k] -= elem
        end
    end
    return Poly(q, num.var), Poly(r, num.var)
end
/(num::Poly, den::Poly) = divrem(num, den)[1]
rem(num::Poly, den::Poly) = divrem(num, den)[2]

function ==(p1::Poly, p2::Poly)
    if length(p1) != length(p2)
        return false
    elseif p1.var != p2.var
        return false
    else
        return p1.a[1+p1.nzfirst:end] == p2.a[1+p2.nzfirst:end]
    end
end

function polyval{T}(p::Poly{T}, x::Number)
    R = promote_type(T, typeof(x))
    lenp = length(p)
    if lenp == 0
        return zero(R)
    else
        y = convert(R, p[1])
        for i = 2:lenp
            y = p[i] + x.*y
        end
        return y
    end
end

polyval(p::Poly, v::AbstractVector) = map(x->polyval(p, x), v)

function polyint{T}(p::Poly{T}, k::Number=0)
    n = length(p)
    R = typeof(one(T)/1)
    a2 = Array(R, n+1)
    for i = 1:n
        a2[i] = p[i] / (n-i+1)
    end
    a2[end] = k
    Poly(a2, p.var)
end

@deprecate polydir polyder

function polyder{T}(p::Poly{T})
    n = length(p)
    if n > 0
        a2 = Array(T, n-1)
        for i = 1:n-1
            a2[i] = p[i] * (n-i)
        end
    else
        a2 = zeros(T, 0)
    end
    Poly(a2, p.var)
end

# create a Poly object from its roots
function poly{T}(r::AbstractVector{T}, var=:x)
    n = length(r)
    c = zeros(T, n+1)
    c[1] = 1
    for j = 1:n
        c[2:j+1] = c[2:j+1]-r[j]*c[1:j]
    end
    return Poly(c, var)
end
poly(A::Matrix, var=:x) = poly(eig(A)[1], var)
poly(A::Matrix, var::String) = poly(eig(A)[1], symbol(var))
poly(A::Matrix, var::Char) = poly(eig(A)[1], symbol(var))

roots{T}(p::Poly{Rational{T}}) = roots(convert(Poly{promote_type(T, Float64)}, p))

# compute the roots of a polynomial
function roots{T}(p::Poly{T})
    R = promote_type(T, Float64)
    num_zeros = 0
    if length(p) == 0
        return zeros(R, 0)
    end
    while abs(p[end-num_zeros]) <= 2*eps(T)
        if num_zeros == length(p)-1
            return zeros(R, 0)
        end
        num_zeros += 1
    end
    n = length(p)-num_zeros-1
    if n < 1
        return zeros(R, length(p)-1)
    end
    companion = zeros(R, n, n)
    a0 = p[end-num_zeros]
    for i = 1:n-1
        companion[1,i] = -p[end-num_zeros-i] / a0
        companion[i+1,i] = 1;
    end
    companion[1,end] = -p[1] / a0
    D,V = eig(companion)
    r = zeros(eltype(D),length(p)-1)
    r[1:n] = 1./D
    return r
end

end # module Poly


kolmogorov_smirnov(F1::Function,F2::Function,grid) =
    maximum(abs.(F1.(grid)-F2.(grid)))

function kolmogorov_smirnov(d::SFMDistribution,F2::Function)
    grid = range(d.dq.mesh.nodes[1],d.dq.mesh.nodes[end],length=2001)
    F1 = cdf(d)
    M = 0.0
    for i in 1:n_phases(d.dq)
        m = maximum(abs.(F1.(grid,i)-F2.(grid,i)))
        (m>M)&&(M=m)
    end
    return M
end
kolmogorov_smirnov(d::SFMDistribution,s::Simulation) = kolmogorov_smirnov(d,cdf(s))
kolmogorov_smirnov(d::SFMDistribution,s::SFMDistribution) = kolmogorov_smirnov(d,cdf(s))

function Lp_cell_probs(d::SFMDistribution,F2::Function, p=1)
    grid = (d.dq.mesh.nodes[1:end-1]+d.dq.mesh.nodes[2:end])./2.0
    F1 = cell_probs(d)
    M = 0.0
    for i in 1:n_phases(d.dq)
        m = sum(abs.(F1.(grid,i)-F2.(grid,i)).^p)
        M += m
    end
    return M
end
Lp_cell_probs(d::SFMDistribution,s::Simulation,p=1) = Lp_cell_probs(d,cell_probs(s),p)
Lp_cell_probs(d::SFMDistribution,s::SFMDistribution,p=1) = Lp_cell_probs(d,cell_probs(s),p)

function Lp_pdf(d::SFMDistribution,f2::Function, p=1)
    grid = range(d.dq.mesh.nodes[1],d.dq.mesh.nodes[end],length=2001)
    f1 = pdf(d)
    M = 0.0
    for i in 1:n_phases(d.dq)
        fn_evals = abs.(f1.(grid,i)-f2.(grid,i)).^p
        m = sum((fn_evals[1:end-1]+fn_evals[2:end])./2.0 .* diff(grid))
        M += m
    end
    return M
end
Lp_pdf(d::SFMDistribution,s::SFMDistribution,p=1) = Lp_pdf(d,pdf(s),p)

function Lp(f1::Function, f2::Function, grid, p=1)
    fn_evals = abs.(f1.(grid)-f2.(grid)).^p
    m = sum((fn_evals[1:end-1]+fn_evals[2:end])./2.0 .* diff(grid))
    return m
end


@doc doc"""
    latexalign()
Generate a ``LaTeX`` align environment from an input.

# Examples
## use with arrays

```julia
lhs = [:(dx/dt), :(dy/dt), :(dz/dt)]
rhs = [:(y-x), :(x*z-y), :(-z)]
latexalign(lhs, rhs)
```

```LaTeX
\begin{align}
\frac{dx}{dt} =& y - x \\\\
\frac{dy}{dt} =& x \cdot z - y \\\\
\frac{dz}{dt} =& - z \\\\
\end{align}
```

## use with ParameterizedFunction

```julia-repl
julia> using DifferentialEquations
julia> ode = @ode_def foldChangeDetection begin
    dm = r_m * (i - m)
    dy = r_y * (p_y * i/m - y)
end i r_m r_y p_y

julia> latexalign(ode)
```
```LaTeX
\begin{align}
\frac{dm}{dt} =& r_{m} \cdot \left( i - m \right) \\\\
\frac{dy}{dt} =& r_{y} \cdot \left( \frac{p_{y} \cdot i}{m} - y \right) \\\\
\end{align}
```

"""
function latexalign end

function latexalign(arr::AbstractMatrix; separator=" =& ", double_linebreak=false, starred=false, kwargs...)
    (rows, columns) = size(arr)
    eol = double_linebreak ? " \\\\\\\\\n" : " \\\\\n"
    arr = latexraw(arr; kwargs...)

    str = "\\begin{align$(starred ? "*" : "")}\n"
    for i in 1:(rows-1)
        str *= join(arr[i,:], separator) * eol
    end
    str *= join(arr[end,:], separator) * "\n"  # No EOL on the last line.
    str *= "\\end{align$(starred ? "*" : "")}\n"
    latexstr = LaTeXString(str)
    COPY_TO_CLIPBOARD && clipboard(latexstr)
    return latexstr
end

function latexalign(lhs::AbstractArray, rhs::AbstractArray; kwargs...)
    return latexalign(hcat(lhs, rhs); kwargs...)
end

function latexalign(lhs::Tuple, rhs::Tuple; kwargs...)
    return latexalign(hcat(collect(lhs), collect(rhs)); kwargs...)
end

latexalign(args::Tuple...; kwargs...) = latexalign(hcat([collect(i) for i in args]...); kwargs...)

latexalign(arg::Tuple; kwargs...) = latexalign(hcat([collect(i) for i in arg]...); kwargs...)

function latexalign(nested::AbstractVector{AbstractVector}; kwargs...)
    return latexalign(hcat(nested...); kwargs...)
end

function latexalign(d::AbstractDict; kwargs...)
    latexalign(collect(keys(d)), collect(values(d)); kwargs...)
end

"""
    latexalign(vec::AbstractVector)

Go through the elements, split at any = sign, pass on as a matrix.
"""
function latexalign(vec::AbstractVector; kwargs...)
    lvec = latexraw(vec; kwargs...)
    ## turn the array into a matrix
    lmat = hcat(split.(lvec, " = ")...)
    ## turn the matrix ito arrays of left-hand-side, right-hand-side.
    larr = [lmat[i,:] for i in 1:size(lmat, 1)]
    length(larr) < 2 && error("Invalid intput to latexalign().")
    return latexalign( hcat(larr...) ; kwargs...)
end


################################################################################
# Douglas-Rachford splitting iterator
#
# See:
# [1] Eckstein, Bertsekas "On the Douglas-Rachford Splitting Method and the Proximal Point Algorithm for Maximal Monotone Operators*", Mathematical Programming, vol. 55, no. 1, pp. 293-318 (1989).
#

struct DRSIterator{I <: Integer, R <: Real, T <: AbstractArray{R}} <: ProximalAlgorithm{I,T}
    x::T
    f
    g
    gamma::R
    maxit::I
    tol::R
    verbose::I
    verbose_freq::I
    y::T
    r::T
    z::T
    FPR_x::T
end

################################################################################
# Constructor(s)

function DRSIterator(x0::AbstractArray{R}; f=Zero(), g=Zero(), gamma::R=1.0, maxit::I=10000, tol::R=1e-4, verbose=1, verbose_freq = 100) where {I, R}
    y = copy(x0)
    r = copy(x0)
    z = copy(x0)
    FPR_x = copy(x0)
    FPR_x .= Inf
    return DRSIterator{I, R, typeof(x0)}(x0, f, g, gamma, maxit, tol, verbose, verbose_freq, y, r, z, FPR_x)
end

################################################################################
# Utility methods

maxit(sol::DRSIterator) = sol.maxit

converged(sol::DRSIterator, it) = it > 0 && maximum(abs,sol.FPR_x)/sol.gamma <= sol.tol

verbose(sol::DRSIterator) = sol.verbose > 0
verbose(sol::DRSIterator, it) = sol.verbose > 0 && (sol.verbose == 2 ? true : (it == 1 || it%sol.verbose_freq == 0))

function display(sol::DRSIterator)
    @printf("%6s | %10s | %10s |\n ", "it", "gamma", "fpr")
    @printf("------|------------|------------|\n")
end

function display(sol::DRSIterator, it)
    @printf("%6d | %7.4e | %7.4e |\n", it, sol.gamma, maximum(abs,sol.FPR_x)/sol.gamma)
end

function Base.show(io::IO, sol::DRSIterator)
    println(io, "Douglas-Rachford Splitting" )
    println(io, "fpr        : $(maximum(abs,sol.FPR_x))")
    print(  io, "gamma      : $(sol.gamma)")
end

################################################################################
# Initialization

function initialize!(sol::DRSIterator)
    return sol.z
end

################################################################################
# Iteration

function iterate!(sol::DRSIterator{I, T}, it::I) where {I, T}
    prox!(sol.y, sol.f, sol.x, sol.gamma)
    sol.r .= 2 .*sol.y .- sol.x
    prox!(sol.z, sol.g, sol.r, sol.gamma)
    sol.FPR_x .= sol.y .- sol.z
    sol.x .-= sol.FPR_x
    return sol.z
end

################################################################################
# Solver interface

function DRS(x0; kwargs...)
    sol = DRSIterator(x0; kwargs...)
    it, point = run!(sol)
    return (it, point, sol)
end


### A Pluto.jl notebook ###
# v0.12.21

using Markdown
using InteractiveUtils

# ╔═╡ e704feae-919c-11eb-35ea-0932dbb413dd
begin
	# using Pkg;
	using StatsBase
	using BenchmarkTools
	using Test
	using Random
	using PlutoUI
end

# ╔═╡ 45d72214-919c-11eb-2e69-23c0786baacb
md"""
### Fast matrix multiplication

Src: [Julia 1.0 Cookbook -  Bogumił Kamiński , Przemysław Szufel - 2018](https://www.packtpub.com/product/julia-1-0-programming-cookbook/9781788998369)

**Context**
Performing computations on matrices are one of the fundamental operations in numerical computing. If we want to multiply an $m × n$ matrix by $n × p$ matrix the cost of this operattion is proportional to the product of these dimensions, namely $O(nmp)$, this  then produces and $n × p$ matrix.

If we want to multiply several matrices, the total cost will depend on the order in which the multiplication are performed.
For example if we have 3 matrices $A$, $B$ and  $C$ with respective dimensions $10×40$, $40×10$ and $10×50$ and we want to conpute $A × B × C$ we can do:

 $(A×B)×C$, cost: $10×40×10 + 10×10×50 ≡  9000$ or

 $A×(B×C)$, cost: $10×40×50 + 40×10×50 ≡ 40000$
"""

# ╔═╡ 2321e634-9267-11eb-1c6f-af2e0a29f7e7
md"""
We will use a technique called dynamic programming (DP) to find the optimal order of multiplications.

Assume we want to multiply the following matrices: A₁, A₂, ...Aₙ.

Let $f(i, j)$ be the minimal cost of multiplaying Aᵢ by Aⱼ.

Our objective is to find: $f(1, n)$.

Performing 0 multiplication has a zero cost, $f(i, i) = 0$

The minimal cost $(f(i, j)$ can be found by considering all possible outermost points of multiplication:

$$f(i, j) = min_{i \le k < j}\{f(i, k) + f(k-1, j) + d(i, 1)×d(k, 2)×d(j,2)\}  \tag{α}$$

where $d(i, s)$ is the $s$-th dimention of matrix $i$

Note: in order to perfomr the actual multiplication, we need to know the value of $f(i, j)$ and the value of $k$ for which the expression in the recurrence was minimal.
"""

# ╔═╡ 61f7bad8-926b-11eb-01c1-1310ca7f9e50
begin
	const TI = Tuple{Int, Int}
	const VTI = Vector{TI}
	const D_TI = Dict{TI, TI}
end

# ╔═╡ 08f24660-919e-11eb-37c9-155375b48cf4
function solve_mult(sizes::NTuple, part_cost::D_TI, from::Int, to::Int)
	if from == to
		part_cost[(from, to)] = (0, from)
		return
	end

	min_cost, min_j = typemax(Int), -1
	for j ∈ from:(to - 1)
		# using memoization, each tuple (f, t) is calc. once
		haskey(part_cost, (from, j)) || solve_mult(sizes, part_cost, from, j)
		haskey(part_cost, (j+1, to)) || solve_mult(sizes, part_cost, j+1, to)

		# using equation (α) defined above
		curr_cost = part_cost[(from, j)][1] + part_cost[(j+1, to)][1] +
			sizes[from][1] * sizes[j][2] * sizes[to][2]

		curr_cost < min_cost && ((min_cost, min_j) = (curr_cost, j))
	end
	# @assert min_j > 0  "min_j must be 0 but: $(min_j)"
	part_cost[(from, to)] = (min_cost, min_j)
	# return
end

# ╔═╡ 09071fe0-926b-11eb-0f26-9f4a7538bdcf
function do_mult(matrices::NTuple, part_cost::D_TI,
		from::Int, to::Int)
	from == to && (return matrices[from])
	from + 1 == to && (return matrices[from] * matrices[to])

	j = part_cost[(from, to)][2]
	do_mult(matrices, part_cost, from, j) * do_mult(matrices, part_cost, j+1, to)
end

# ╔═╡ e849191c-919c-11eb-0d04-b9ec7a4ac478
function fast_matmult(matrices::AbstractMatrix...)# ::AbstractMatrix
	length(matrices) ≤ 1 && (return *(matrices...))

	sizes = size.(matrices)
	!all(ix -> sizes[ix+1][1] == sizes[ix][2], 1:length(sizes) - 1) &&
		throw(ArgumentError("Matrices dimensions mismatch"))

	from, to = 1, length(sizes)
	part_cost = D_TI()
	solve_mult(sizes, part_cost, from, to)
	do_mult(matrices, part_cost, from, to)
end

# ╔═╡ 5b5c8830-926f-11eb-0654-854d9567db7c
begin
	n, p = 5000, 5
	m₁ = rand(Float64, (p, n))
	m₂ = rand(Float64, (n, p))
	m₃ = rand(Float64, (p, 4000))
	m₄ = rand(Float64, (4000, p))

	with_terminal() do
		@btime *(repeat([m₁, m₂, m₃, m₄], outer=10)...)
	end
end

# ╔═╡ 3c7301d4-9270-11eb-2a8a-bf0f3a69f0f9
with_terminal() do
	@btime fast_matmult(repeat([m₁, m₂, m₃, m₄], outer=10)...)
end

# ╔═╡ 204f8bd6-9277-11eb-2b15-1199f4f503a6
md"""
In practice, we migth want to define a macro to turn a matrix multiplication into a fast matrix multiplication:
"""

# ╔═╡ 4082f32a-9277-11eb-06f8-939e0ff5560e
macro fast_matmult(expr::Expr)
	expr.head == :call || throw(ArgumentError("expression must be a call"))
	expr.args[1] == :(*) || throw(ArgumentError("only for multiplication!"))
	#
	new_expr = deepcopy(expr)         ## make a defensive copy
	new_expr.args[1] = :fast_matmult  ## overwrite '*' operator with ours
	#
	esc(new_expr)                     ## macor hygiene
 end

# ╔═╡ 1d773e30-9278-11eb-1d3d-2bdd47ef091b
@fast_matmult rand(Float64, (2, 3)) * rand(Float64, (3,4)) * ones(4, 5)
#
# expr = :(rand(Float64, (2, 3)) * rand(Float64, (3, 4)) * ones(4, 5))
# new_expr = :(fast_matmult(rand(Float64, (2, 3)), rand(Float64, (3, 4)), ones(4, 5)))
#

# ╔═╡ Cell order:
# ╟─45d72214-919c-11eb-2e69-23c0786baacb
# ╟─2321e634-9267-11eb-1c6f-af2e0a29f7e7
# ╠═e704feae-919c-11eb-35ea-0932dbb413dd
# ╠═61f7bad8-926b-11eb-01c1-1310ca7f9e50
# ╠═e849191c-919c-11eb-0d04-b9ec7a4ac478
# ╠═08f24660-919e-11eb-37c9-155375b48cf4
# ╠═09071fe0-926b-11eb-0f26-9f4a7538bdcf
# ╠═5b5c8830-926f-11eb-0654-854d9567db7c
# ╠═3c7301d4-9270-11eb-2a8a-bf0f3a69f0f9
# ╟─204f8bd6-9277-11eb-2b15-1199f4f503a6
# ╠═4082f32a-9277-11eb-06f8-939e0ff5560e
# ╠═1d773e30-9278-11eb-1d3d-2bdd47ef091b


### A Pluto.jl notebook ###
# v0.12.18

using Markdown
using InteractiveUtils

# ╔═╡ c2efdfae-502c-11eb-24f3-ef331e609508
md"""
# Set

Set is a collection without an order and doesn't allow repeatition of an element. And similar to Mathematical set we can perform set operations on these data structures
"""

# ╔═╡ 08c92a08-502d-11eb-056b-794d9230ea3f
begin
	A = Set([1,2,3])
	B = Set([2,3,4])
	A ∪ B
end

# ╔═╡ 1b8e7af8-502d-11eb-269c-d1f3e80a2b0e

# ╔═╡ d61a4b96-502c-11eb-0265-ef3964d04fd1
md"""
# Map (Dictionaries)

Associative Data Structure that associates a Key with a Value. Keys in a map is a set, i.e each key exist only once.

"""

# ╔═╡ 27bdcffc-502d-11eb-25ba-fbc0720d76fb
dict1 = ["a" => 1, "b"=>2, "c"=>3]

# ╔═╡ dc9ec154-502c-11eb-30ac-61ee85f7dc15
md"""
# Hashing allows O(1) lookups

- In arrays, list we need to look through every element hence in lookups are in linear time.
- In stacks and queues, we can lookup the start or end almost immediately
- priority queues allows to access priority element fast but linear time for other elements

Hash functions coverts a value to a hash value that is used as index in an array. How ?
- One easiest approach is to divide the last few digits by certain number and use remainder as index to store the elements in as an array. And access elements by indexing in constant time.
-

Collision (same hash value for multiple item) can occur. There are two main ways to handle collision
1. Change hash function completely, or change the number for division to get enough slots for all values
    - using bigger number need a larger space to maintain a constant time lookup. Changing numbers requires moving data to a new arrays
2. Change the structure of arrays to buckets: rather than store one value in the slot store collection of values in the buckets
    - need to iterate through a list or link list. on average has a constant time look up but in worst case it can become linear time O(n).
    - Can use second hash function inside the bucket to further split the data

"""

# ╔═╡ e0e1b918-502c-11eb-1506-819b3261e35e
md"""
# Hash Maps

using Keys to create Hash values.

Strings can be converted using ASCII values for first few laters as hash values.

"""

# ╔═╡ e69324b6-502c-11eb-0307-13492268de72

# ╔═╡ Cell order:
# ╠═c2efdfae-502c-11eb-24f3-ef331e609508
# ╠═08c92a08-502d-11eb-056b-794d9230ea3f
# ╠═1b8e7af8-502d-11eb-269c-d1f3e80a2b0e
# ╠═d61a4b96-502c-11eb-0265-ef3964d04fd1
# ╠═27bdcffc-502d-11eb-25ba-fbc0720d76fb
# ╠═dc9ec154-502c-11eb-30ac-61ee85f7dc15
# ╠═e0e1b918-502c-11eb-1506-819b3261e35e
# ╠═e69324b6-502c-11eb-0307-13492268de72


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

using Markdown
using InteractiveUtils

# ╔═╡ dab179ae-4a5a-11eb-317c-c7fa9d9ccf8f
begin
	import Pkg
	Pkg.activate(Base.current_project())
	using ReinforcementLearning
end

# ╔═╡ 109c4fb2-4a5b-11eb-08d5-bd6b1eb0ebe9
using Plots

# ╔═╡ 1fcd93f0-4a5c-11eb-252d-9da5bc78b08b
using StatsPlots

# ╔═╡ 1fbc2952-4b1b-11eb-3b65-75c1058a9537
using Flux

# ╔═╡ db64341a-4b1b-11eb-3f7b-f11b26f442f4
using Statistics

# ╔═╡ 538080e4-4a5a-11eb-0570-65614c5797f0
md"""
# Ten Armed Bandits Environment

In this chapter, we'll use the `MultiArmBanditsEnv` to study two main concepts in reinforcement learning: **exploration** and **exploitation**.

Let's take a look at the environment first.
"""

# ╔═╡ dfe484d4-4a5a-11eb-0224-573d091b3d08
env = MultiArmBanditsEnv()

# ╔═╡ 1427b132-4a5b-11eb-3506-e744a9a5595c
violin(
	[
		[
			begin
				reset!(env)
				env(a)
				reward(env)
			end
			for _ in 1:100
		]
		for a in action_space(env)
	],
	leg=false
)

# ╔═╡ 69bc9e66-4a5c-11eb-0288-1930cdb31d9d
md"""
The above figure shows the reward distribution of each action. （Figure 2.1）
"""

# ╔═╡ c0ca4172-4aac-11eb-255d-8b0005441fb0
md"""
Now we create a testbed to calulate the average reward and perfect action percentage.
"""

# ╔═╡ 4bf0f782-4aad-11eb-291c-afa853f150a3
"""
A customized hook to record whether the action to take is the best action or not.
"""
begin
	Base.@kwdef struct CollectBestActions <: AbstractHook
		best_action::Int
		isbest::Vector{Bool} = []
	end

	function (h::CollectBestActions)(::PreActStage, agent, env, action)
		push!(h.isbest, h.best_action==action)
	end
end

# ╔═╡ 1ff8d726-4aad-11eb-0d88-c7f6080c4072
md"""
Writing a customized hook is easy.

1. Define your `struct` and make it inherit from `AbstractHook` (optional).
1. Write your customized runtime logic by overwriting some of the following functions. By default, they will do nothing if your hook inherits from `AbstractHook`.
    - `(h::YourHook)(::PreActStage, agent, env, action)`
    - `(h::YourHook)(::PostActStage, agent, env)`
    - `(h::YourHook)(::PreEpisodeStage, agent, env)`
    - `(h::YourHook)(::PostEpisodeStage, agent, env)`
"""

# ╔═╡ 1dfcd040-4b1a-11eb-248e-990c0e029f43
function bandit_testbed(
	;explorer=EpsilonGreedyExplorer(0.1),
	true_reward=0.0,
	init=0.,
	opt=InvDecay(1.0)
)
   env = MultiArmBanditsEnv(;true_reward=true_reward)
   agent = Agent(
	   policy=QBasedPolicy(
		   learner = TDLearner(
			   approximator = TabularQApproximator(
				   n_state=length(state_space(env)),
				   n_action=length(action_space(env)),
				   init=init,
				   opt = opt
			   ),
			   γ = 1.0,
			   method=:SARSA,
			   n = 0,
		   ),
		   explorer = explorer
	   ),
	   trajectory=VectorSARTTrajectory()
	)
	h1 = CollectBestActions(;best_action=findmax(env.true_values)[2])
	h2 = TotalRewardPerEpisode(;is_display_on_exit=false)
	run(agent, env, StopAfterStep(1000), ComposedHook(h1, h2))
    h1.isbest, h2.rewards
end

# ╔═╡ 15e65f7e-4b1b-11eb-26b7-ef85168b5112
begin
	p = plot(layout=(2, 1))
	for ϵ in [0.1, 0.01, 0.0]
    	stats = [
			bandit_testbed(;explorer=EpsilonGreedyExplorer(ϵ))
			for _ in 1:2000
		]
    	plot!(p, mean(x[2] for x in stats);
			subplot=1, legend=:bottomright, label="epsilon=$ϵ")
    	plot!(p, mean(x[1] for x in stats);
			subplot=2, legend=:bottomright, label="epsilon=$ϵ")
	end
	p
end

# ╔═╡ b302e1e6-4b1c-11eb-1cb8-91be2a24d5a7
begin
	# now compare the effect of setting init value
	p_2_3 = plot(legend=:bottomright)

	v1 = mean(
		bandit_testbed(
			;explorer=EpsilonGreedyExplorer(0.),
			init=5.,
			opt=Descent(0.1)
		)[1]
		for _ in 1:2000
	)
	plot!(p_2_3, v1, label="Q_1=5, epsilon=0.")

	v2 = mean(
		bandit_testbed(
			;explorer=EpsilonGreedyExplorer(0.1),
			init=0.,
			opt=Descent(0.1)
		)[1]
		for _ in 1:2000
	)
	plot!(p_2_3, v2, label="Q_1=0, epsilon=0.1")
	p_2_3
end

# ╔═╡ 21c988d8-4b80-11eb-15aa-550e091315f5
begin
	p_2_4 = plot(legend=:bottomright)
	plot!(p_2_4, mean(bandit_testbed(;explorer=UCBExplorer(10), opt=Descent(0.1))[2] for _ in 1:5000), label="UpperConfidenceBound, c=2")
	plot!(p_2_4, mean(bandit_testbed(;explorer=EpsilonGreedyExplorer(0.1), opt=Descent(0.1))[2] for _ in 1:5000), label="epsilon-greedy, epsilon=0.1")

	p_2_4
end

# ╔═╡ 5b3f7b36-4b1e-11eb-0035-2bced9aead26
md"""
Similar to the `bandit_testbed` function, we'll create a new function to test the performance of `GradientBanditLearner`.
"""

# ╔═╡ 6d93c3d0-4b1e-11eb-2b41-af6689ba72f4
function gb_bandit_testbed(
	;baseline=0.,
	explorer=WeightedExplorer(is_normalized=true),
	true_reward=0.0,
	init=0.,
	opt=InvDecay(1.0)
)
    env = MultiArmBanditsEnv(;true_reward=true_reward)
	agent = Agent(
	   policy=QBasedPolicy(
		   learner = GradientBanditLearner(
			   approximator = TabularQApproximator(
				   n_state=length(state_space(env)),
				   n_action=length(action_space(env)),
				   init=init,
				   opt = opt
			   ),
			   baseline=baseline
		   ),
		   explorer = explorer
	   ),
	   trajectory=VectorSARTTrajectory()
	)

	h1 = CollectBestActions(;best_action=findmax(env.true_values)[2])
	h2 = TotalRewardPerEpisode(;is_display_on_exit=false)
	run(agent, env, StopAfterStep(1000), ComposedHook(h1, h2))
    h1.isbest, h2.rewards
end

# ╔═╡ 04e8320c-4b1f-11eb-3340-47f7392a8282
md"""
Note that there's a keyword argument named `baseline` in the `GradientBanditLearner`. It can be either a number or a callable function (`reward -> value`). One of such functions mentioned in the book is to calculate the average of seen rewards.
"""

# ╔═╡ b291cb0c-4b1f-11eb-3ee5-cfcdfdcae00b
begin
	Base.@kwdef mutable struct SampleAvg
		t::Int = 0
		avg::Float64 = 0.0
	end

	function (s::SampleAvg)(x)
		s.t += 1
		s.avg += (x - s.avg) / s.t
		s.avg
	end
end

# ╔═╡ e0e72a60-4b1f-11eb-1001-89777fd3d0f7
begin
	baseline = SampleAvg()
	[baseline(x) for x in 1:10]
end

# ╔═╡ 42525d24-4b20-11eb-099c-b10c90af166e
begin
	true_reward = 4.0

	p_2_5 = plot(legend=:bottomright)

	plot!(p_2_5, mean(gb_bandit_testbed(;opt=Descent(0.1), baseline=SampleAvg(), true_reward=true_reward)[1] for _ in 1:2000), label="alpha = 0.1, with baseline")
	plot!(p_2_5, mean(gb_bandit_testbed(;opt=Descent(0.4), baseline=SampleAvg(), true_reward=true_reward)[1] for _ in 1:2000), label="alpha = 0.4, with baseline")
	plot!(p_2_5, mean(gb_bandit_testbed(;opt=Descent(0.1), true_reward=true_reward)[1] for _ in 1:2000), label="alpha = 0.1, without baseline")
	plot!(p_2_5, mean(gb_bandit_testbed(;opt=Descent(0.4), true_reward=true_reward)[1] for _ in 1:2000), label="alpha = 0.4, without baseline")

	p_2_5
end

# ╔═╡ aad675d2-4b80-11eb-3d78-a1ef731d7d8b
begin

	p_2_6 = plot(legend=:topleft)

	plot!(p_2_6, -7:-2, [mean(mean(bandit_testbed(;explorer=EpsilonGreedyExplorer(2.0^i))[2] for _ in 1:2000)) for i in -7:-2], label="epsilon greedy")
	plot!(p_2_6, -5:1, [mean(mean(gb_bandit_testbed(;explorer=WeightedExplorer(is_normalized=true), opt=Descent(2.0^i))[2] for _ in 1:2000)) for i in -5:1], label="gradient")
	plot!(p_2_6, -4:2, [mean(mean(bandit_testbed(;explorer=UCBExplorer(10; c=2.0^i))[2] for _ in 1:2000)) for i in -4:2], label="UCB")
	plot!(p_2_6, -2:2, [mean(mean(bandit_testbed(;explorer=EpsilonGreedyExplorer(0.), init=(2.0^i))[2] for _ in 1:2000)) for i in -2:2], label="greedy with initialization")

	p_2_6
end

# ╔═╡ Cell order:
# ╟─538080e4-4a5a-11eb-0570-65614c5797f0
# ╠═dab179ae-4a5a-11eb-317c-c7fa9d9ccf8f
# ╠═dfe484d4-4a5a-11eb-0224-573d091b3d08
# ╠═109c4fb2-4a5b-11eb-08d5-bd6b1eb0ebe9
# ╠═1fcd93f0-4a5c-11eb-252d-9da5bc78b08b
# ╠═1427b132-4a5b-11eb-3506-e744a9a5595c
# ╟─69bc9e66-4a5c-11eb-0288-1930cdb31d9d
# ╟─c0ca4172-4aac-11eb-255d-8b0005441fb0
# ╠═4bf0f782-4aad-11eb-291c-afa853f150a3
# ╟─1ff8d726-4aad-11eb-0d88-c7f6080c4072
# ╠═1fbc2952-4b1b-11eb-3b65-75c1058a9537
# ╠═db64341a-4b1b-11eb-3f7b-f11b26f442f4
# ╠═1dfcd040-4b1a-11eb-248e-990c0e029f43
# ╠═15e65f7e-4b1b-11eb-26b7-ef85168b5112
# ╠═b302e1e6-4b1c-11eb-1cb8-91be2a24d5a7
# ╠═21c988d8-4b80-11eb-15aa-550e091315f5
# ╠═5b3f7b36-4b1e-11eb-0035-2bced9aead26
# ╠═6d93c3d0-4b1e-11eb-2b41-af6689ba72f4
# ╟─04e8320c-4b1f-11eb-3340-47f7392a8282
# ╠═b291cb0c-4b1f-11eb-3ee5-cfcdfdcae00b
# ╠═e0e72a60-4b1f-11eb-1001-89777fd3d0f7
# ╠═42525d24-4b20-11eb-099c-b10c90af166e
# ╠═aad675d2-4b80-11eb-3d78-a1ef731d7d8b


# statefuliterator
it = ShapleyValues.statefuliterator(1:3)
@test length(it) == 3
@test done(it) == false
@test next(it) == 1
@test next(it) == 2
@test done(it) == false
tmp = next(it)
@test tmp == 3
@test done(it) == true
@test typeof(tmp) == Int64

# allocate_samples
@assert ShapleyValues.allocate_samples([1,1,1,1], 4, ones(Int64, 4)) == ones(4)
@assert ShapleyValues.allocate_samples([1,1,1,1], 8, ones(Int64, 4)) == ones(4)
@assert ShapleyValues.allocate_samples([1,2,1,1], 10, 4*ones(Int64, 4)) == [2,4,2,2]
@assert ShapleyValues.allocate_samples([1,2,1,1], 11, [4,4,4,1]) == [3,4,3,1]
@assert sum(ShapleyValues.allocate_samples([1,2,1,1], 10, 3*ones(Int64, 4))) == 10
@assert maximum(ShapleyValues.allocate_samples([1,2,1,1], 10, 3*ones(Int64, 4))) == 3
@assert minimum(ShapleyValues.allocate_samples([1,2,1,1], 10, 3*ones(Int64, 4))) == 2
@assert sum(ShapleyValues.allocate_samples([1,2,1,1], 21, 10*ones(Int64, 4))) == 21

# MeanVarianceAccumulator
v = [1,2,3.0]
mv = ShapleyValues.MeanVarianceAccumulator()
@test isnan(std(mv))
ShapleyValues.observe!(mv, v[1], 1)
@test mean(mv) == v[1]
@test isnan(std(mv))
ShapleyValues.observe!(mv, v[2], 1)
@test mean(mv) == mean(v[1:2])
@test std(mv) ≈ std(v[1:2])
ShapleyValues.observe!(mv, v[3], 1)
@test mean(mv) == mean(v)
@test std(mv) ≈ std(v)
@test length(mv) == 3

w = [0.1,0.4,0.5]
v = [1,2,3.0]
mv = ShapleyValues.MeanVarianceAccumulator()
@test isnan(std(mv))
ShapleyValues.observe!(mv, v[1], w[1])
@test mean(mv) == v[1]
@test isnan(std(mv))
ShapleyValues.observe!(mv, v[2], w[2])
@test mean(mv) == dot(v[1:2], w[1:2])/sum(w[1:2])
@test sum(((v[1:2] - mean(mv)).^2).*w[1:2])/sum(w[1:2]) * (2/1) ≈ var(mv)
ShapleyValues.observe!(mv, v[3], w[3])
@test mean(mv) ==  dot(v, w)/sum(w)
@test sum(((v - mean(mv)).^2).*w)/sum(w) * (3/2) ≈ var(mv)

# varying_groups
@test ShapleyValues.varying_groups(zeros(3), [0 0 1; 0 1 1; 0 0 1]', Array{Int64,1}[[1,2],[3]]) == [1,2]
@test ShapleyValues.varying_groups(zeros(3), [0 0 1; 0 1 1; 0 0 1]', Array{Int64,1}[[1],[2,3]]) == [2]
@test ShapleyValues.varying_groups(zeros(3), zeros(3,10), Array{Int64,1}[[1],[2,3]]) == Int64[]


# Before running this, please make sure to activate and instantiate the environment
# corresponding to [this `Project.toml`](https://raw.githubusercontent.com/alan-turing-institute/DataScienceTutorials.jl/master/Project.toml) and [this `Manifest.toml`](https://raw.githubusercontent.com/alan-turing-institute/DataScienceTutorials.jl/master/Manifest.toml)
# so that you get an environment which matches the one used to generate the tutorials:
#
# ```julia
# cd("DataScienceTutorials") # cd to folder with the *.toml
# using Pkg; Pkg.activate("."); Pkg.instantiate()
# ```

using MLJ

@load LinearRegressor pkg=MLJLinearModels

import RDatasets: dataset
import DataFrames: describe, select, Not, rename!
boston = dataset("MASS", "Boston")
first(boston, 3)

describe(boston, :mean, :std, :eltype)

data = coerce(boston, autotype(boston, :discrete_to_continuous));

y = data.MedV
X = select(data, Not(:MedV));

mdl = LinearRegressor()

X_uni = select(X, :LStat) # only a single feature
mach_uni = machine(mdl, X_uni, y)
fit!(mach_uni)

fp = fitted_params(mach_uni)
@show fp.coefs
@show fp.intercept

using PyPlot

figure(figsize=(8,6))
plot(X.LStat, y, ls="none", marker="o")
Xnew = (LStat = collect(range(extrema(X.LStat)..., length=100)),)
plot(Xnew.LStat, predict(mach_uni, Xnew))

mach = machine(mdl, X, y)
fit!(mach)

fp = fitted_params(mach)
coefs = fp.coefs
intercept = fp.intercept
for (name, val) in coefs
    println("$(rpad(name, 8)):  $(round(val, sigdigits=3))")
end
println("Intercept: $(round(intercept, sigdigits=3))")

ŷ = predict(mach, X)
round(rms(ŷ, y), sigdigits=4)

figure(figsize=(8,6))
res = ŷ .- y
stem(res)

figure(figsize=(8,6))
hist(res, density=true)

X2 = hcat(X, X.LStat .* X.Age);

rename!(X2, :x1 => :interaction);

mach = machine(mdl, X2, y)
fit!(mach)
ŷ = predict(mach, X2)
round(rms(ŷ, y), sigdigits=4)

X3 = hcat(X.LStat, X.LStat.^2)
mach = machine(mdl, X3, y)
fit!(mach)
ŷ = predict(mach, X3)
round(rms(ŷ, y), sigdigits=4)

Xnew = (LStat = Xnew.LStat, LStat2 = Xnew.LStat.^2)

figure(figsize=(8,6))
plot(X.LStat, y, ls="none", marker="o")
plot(Xnew.LStat, predict(mach, Xnew))

# This file was generated using Literate.jl, https://github.com/fredrikekre/Literate.jl


# ---
# title: 403. Frog Jump
# id: problem403
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Dynamic Programming
# link: <https://leetcode.com/problems/frog-jump/description/>
# hidden: true
# ---
#
# A frog is crossing a river. The river is divided into x units and at each unit
# there may or may not exist a stone. The frog can jump on a stone, but it must
# not jump into the water.
#
# Given a list of stones' positions (in units) in sorted ascending order,
# determine if the frog is able to cross the river by landing on the last stone.
# Initially, the frog is on the first stone and assume the first jump must be 1
# unit.
#
# If the frog's last jump was _k_ units, then its next jump must be either _k_
# \- 1, _k_ , or _k_ \+ 1 units. Note that the frog can only jump in the forward
# direction.
#
# **Note:**
#
#   * The number of stones is ≥ 2 and is < 1,100.
#   * Each stone's position will be a non-negative integer < 231.
#   * The first stone's position is always 0.
#
# **Example 1:**
#
#
#
#     **[0,1,3,5,6,8,12,17]**
#
#     There are a total of 8 stones.
#     The first stone at the 0th unit, second stone at the 1st unit,
#     third stone at the 3rd unit, and so on...
#     The last stone at the 17th unit.
#
#     **Return true**. The frog can jump to the last stone by jumping
#     1 unit to the 2nd stone, then 2 units to the 3rd stone, then
#     2 units to the 4th stone, then 3 units to the 6th stone,
#     4 units to the 7th stone, and 5 units to the 8th stone.
#
#
# **Example 2:**
#
#
#
#     **[0,1,2,3,4,8,9,11]**
#
#     **Return false**. There is no way to jump to the last stone as
#     the gap between the 5th and 6th stone is too large.
#
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


# julia> include("/Users/ian/jd/10-19 Code/12 Julia/12.03 PE/pe3.jl")

function sieve_e(n)
    ints = 2:n; p = 2; i = 1
    while ints[end] != p
        ints = setdiff(ints,[j*p for j=2:Int(ceil(n/p))])
        p = ints[i + 1]; i += 1
    end
    return ints
end

function gpd(n)
    test = []
    p = sieve_e(Int(ceil(sqrt(n))))
    for i in p
        if n%i==0
            push!(test,i)
        end
    end
    return test[end]
end

println("ans = $(gpd(600851475143))")


using RDatasets

function col_normalize(X::Array{Float64})
    xmin = minimum(X, 1)
    xmax = maximum(X, 1)
    for i in 1:size(X, 2)
        X[:, i] = (X[:, i] .- xmin[i]) ./ (xmax[i] - xmin[i])
    end
    X
end

function get_iris()
    iris = dataset("datasets", "iris")

    iris_data = Array{Float64}(hcat(iris[:SepalLength], iris[:SepalWidth],
                                    iris[:PetalWidth], iris[:PetalLength]));
    X = col_normalize(iris_data)'

    sps = unique(iris[:Species])
    iris[:label] = indexin(iris[:Species], sps)
    Y = Array{Int64}(iris[:label])
    X, Y
end

function get_spirals()
    spirals = readtable("data/spiral.txt", header=false, separator=' ',
                        names=[:x, :y, :label])
    X = col_normalize(Array{Float64}(hcat(spirals[:x], spirals[:y])))'
    Y = Array{Int64}(spirals[:label])
    X, Y
end

function get_yeast()
    yeast = readtable("data/yeast.data", header=false, separator=' ',
                      names=[:seq, :mcg, :gvh, :alm, :mit, :erl, :pox, :vac,
                             :nuc, :local])
    yeast_data = Array{Float64}(hcat(
        map(i->yeast[i],[:mcg, :gvh, :alm, :mit, :erl, :pox, :vac, :nuc])...))
    X = col_normalize(yeast_data)'
    locs = unique(yeast[:local])
    yeast[:label] = indexin(yeast[:local], locs)
    Y = Array{Int64}(yeast[:label])
    X, Y
end

function get_data(problem::String)
    if problem == "iris"
        return get_iris()
    elseif problem == "spirals"
        return get_spirals()
    elseif problem == "yeast"
        return get_yeast()
    end
    error("Problem not found")
end

function train_test_split(X::Array{Float64}, Y::Array{Int64}, trainp::Float64=0.75)
    train = rand(length(Y)) .< trainp
    test = .~(train)
    X[:, train], Y[train], X[:, test], Y[test]
end

function onehot(Y::Array{Int64})
    oneh = zeros(Int64, length(Y), length(unique(Y)))
    for i in eachindex(Y)
        oneh[i, Y[i]] = 1
    end
    oneh
end


include("../src/prrqr2.jl")
using BenchmarkTools
using LinearAlgebra
using LowRankApprox
using Random

Random.seed!(0)

tol = 1e-3

n = 1000; k = 10; m = 1000
s = zeros(k)
s[1:k] = @. exp(-(1:k));
U = randn(n,k)
V = randn(m,k)
A = U*diagm(s)*V';

#println("Benchmarking rank-revealing QR...")
#@btime _ = prrqr!(randn(n,100), tol; reltol=true)

println("Testing custom implementation...")
X = copy(A)
Q, R = _compress_block!(X; tol, reltol=false)
X = copy(A)
@time Q, R = _compress_block!(X; tol, reltol=false)
println("error in the truncated QR decomposition: ", norm(A - Q*R)/norm(A))

println("Testing LowRankApprox.jl implementation...")
# LowRankApprox.jl implementation
X = copy(A)
F = pqrfact(X; atol = tol)
X = copy(A)
@time F = pqrfact(X; atol = tol)
println("error in the LowRankApprox.jl PQR decomposition: ", norm(A - F.Q*F.R)/norm(A))

#_, σ, _ = svd(A);
#plot(1:n, σ, yaxis=:log, label="singular values")


# f(x)=mean(x)

function Fmean(x...)
    tmp=0.0
    for i in 1:length(x)
        tmp+=mean(x[i])
    end
    return ([tmp/length(x)],nothing)
end

function Fmean_inplace(value::Array,auxvalue,x...) # inplace
    tmp=0.0
    for i in 1:length(x)
        tmp+=mean(x[i])
    end
    fill!(value,tmp/length(x))
end

function Dmean(derivativeIDX,f_c,faux_c,grad_c,grad_n,x...)
    axpy!(grad_c[1],ones(size(grad_n))/(length(x)*length(x[derivativeIDX])),grad_n)
end

import Base.mean

if PROC=="GPU"

    function mean(A::CudaArray{Float64})
        return flatten(Float64,CUBLAS.gemv('T',1./length(A),flatten(Float64,A),CudaArray(ones(length(A)))))
    end

    function mean(A::CudaArray{Float32})
        return flatten(Float32,CUBLAS.gemv('T',Float32(1./length(A)),flatten(Float32,A),CudaArray(Float32,ones(Float32,length(A)))))
    end
    export mean

    function mean!(A::CudaArray{Float64},Out::CudaArray{Float64})
         CUBLAS.gemv!('T',1./length(A),flatten(Float64,A),CudaArray(ones(length(A))),0.0,Out)
    end

    function mean!(A::CudaArray{Float32},Out::CudaArray{Float32})
         CUBLAS.gemv!('T',Float32(1./length(A)),flatten(Float32,A),CudaArray(Float32,ones(length(A))),Float32(0.0),Out)
    end

    export mean!

    function Fmean()
    end

    function Fmean(x::CudaArray{Float64}...)
        tmp=CudaArray(zeros(1,1))
        for i in 1:length(x)
            axpy!(1.0/length(x),mean(x[i]),tmp)
        end
        return (tmp,nothing)
    end

    function Fmean(x::CudaArray{Float32}...)
        tmp=CudaArray(Float32,zeros(1,1))
        for i in 1:length(x)
            axpy!(Float32(1.0/length(x)),mean(x[i]),tmp)
        end
        return (tmp,nothing)
    end

    function Fmean_inplace(value::CudaArray{Float64},auxvalue,x::CudaArray{Float64}...) # inplace
        fill!(value,0.0)
        for i in 1:length(x)
            axpy!(1.0/length(x),mean(x[i]),value)
        end
    end

    function Fmean_inplace(value::CudaArray{Float32},auxvalue,x::CudaArray{Float32}...) # inplace
        fill!(value,Float32(0.0))
        for i in 1:length(x)
            axpy!(Float32(1.0/length(x)),mean(x[i]),value)
        end
    end

    function Dmean(derivativeIDX,f_c,faux_c,grad_c::CudaArray{Float64},grad_n::CudaArray{Float64},x::CudaArray{Float64}...)
        tmp=CudaArray(Float64,size(grad_n))
        fill!(tmp,1.0/(length(x)*length(x[derivativeIDX])))
        axpy!(grad_c,tmp,grad_n)
        free(tmp)
    end

    function Dmean(derivativeIDX,f_c,faux_c,grad_c::CudaArray{Float32},grad_n::CudaArray{Float32},x::CudaArray{Float32}...)
        tmp=CudaArray(Float32,size(grad_n))
        fill!(tmp,Float32(1.0/(length(x)*length(x[derivativeIDX]))))
        axpy!(grad_c,tmp,grad_n)
        free(tmp)
    end

end

Derivative[Fmean]=Dmean # Define dictionary lookup
Inplace[Fmean]=Fmean_inplace

mean(n::ADnode)=ADnode(Fmean,n)

function mean(n::ArrayADnode)
    return ADnode(Fmean,n)
end

#mean(A::ADtrans)=ADnode(Fmean, ftranspose(node[A.parent]))
mean(A::ADtrans)=ADnode(Fmean, node[A.parent]) # mean(A')=mean(A)

export mean


using Distributions, Statistics, Random
rng = MersenneTwister(01012019);

"""
    kernel_scenarios(data::Union{Array{Int64,2},Array{Float64,2}}, kernelDistribution = Logistic; Markovian::Bool = true)

Returns an instance of a closure of the conditional density estimation method.

Args:
- data - a matrix of data in 2 dimension (NxT) where N is the number of trajectories and T is the number of stages.
- kernelDistribution - the distribution of the kernel functions. Default is Logistic distribution.

An optional keyword `Markovian` is passed to handle weights
for Markovian trajectories or non-Markovian ones.

To get a Markov trajectory from above,
`KernelScenarios(data,kernelDistribution;Markovian=true)()` gives the result for
Markovian samples, which is used to generate a scenario lattice.
"""
function kernel_scenarios(data::Union{Array{Int64,2},Array{Float64,2}}, kernelDistribution = Logistic; Markovian::Bool = true)
    function closure()
        N , T = size(data)                             # Dimensions of the data
        d = 1
        w = fill(1.0, N)                               # initialize the weights
        x = Array{Float64,2}(undef, T , 1)             # trajectory to be created
        for t = 1 : T
            w = w / sum(w)                             # normalized weights
            Nt = sum(w)^2 / sum(w.^2)                  # effective sample size
            σt = sqrt(sum(w .* ((data[:,t] .- sum(w .* data[:,t])).^2)))  # effective standard devistion
            ht = σt * Nt^(-1 / (d + 4)) + eps()
            # Composition method comes in here
            u = rand(rng, Uniform(0,1))
            _,jstar = findmin(cumsum(w) .< u * sum(w))
            x[t] = rand(rng, kernelDistribution(data[jstar,t], ht))
            if t<T
                if Markovian      # Markovian is used for scenario lattices
                    for j = 1:N
                        w[j] = pdf(kernelDistribution(data[j,t], ht), x[t])
                    end
                else             # Non-Markovian is used for scenario trees
                    for j = 1:N
                        w[j] *= pdf(kernelDistribution(data[j,t], ht), x[t])
                    end
                end
            end
        end
        return x        # The length of x is equal to the number stages = T
    end
end


# constructor for the vecchia factorization approxiamtion

"""
Vecchia Factorization approximation:
```
vecchia(Σ, blk_sizes) -> R⁻¹ M R⁻ᴴ
```
Pivoted Vecchia Factorization approximation:
```
vecchia(Σ, blk_sizes, perm) -> Pᵀ R⁻¹ M R⁻ᴴ P
```
where the argument where `Σ` is an `AbstractMatrix` or a `Function` taking indices `(i,j)` and
returning `Σ[i,j]`.
"""
function vecchia end

"""
Construct individual factors in the Pivoted Vecchia approximation `Pᵀ R⁻¹ M R⁻ᴴ P`.
```
R_M_P(Σ, blk_sizes, perm) -> R, M, P
```
where the argument where `Σ` is an `AbstractMatrix` or a `Function` taking indices `(i,j)` and
returning `Σ[i,j]`. The return values have the following types:
```
R <: Ridiagonal
M <: Midiagonal
P <: Piv
```
"""
function R_M_P end

function vecchia(Σ::Union{AbstractMatrix, Function}, blk_sizes::AbstractVector{<:Integer}, perm::AbstractVector{<:Integer})
	R, M, P = R_M_P(Σ, blk_sizes, perm)
    # println("Vecchia Factorization approximation: Pᵀ R⁻¹ M R⁻ᴴ P")
	return P' * inv(R) * M * inv(R)' * P
end

function vecchia(Σ::Union{AbstractMatrix, Function}, blk_sizes::AbstractVector{<:Integer})
	R, M, P = R_M_P(Σ, blk_sizes)
    # println("Vecchia Factorization approximation: R⁻¹ M R⁻ᴴ")
	return inv(R) * M * inv(R)'
end

function R_M_P(Σ::AbstractMatrix{T}, blk_sizes::AbstractVector{<:Integer}, perm::AbstractVector{<:Integer}=1:sum(blk_sizes)) where T
	LinearAlgebra.checksquare(Σ)
	@assert isperm(perm)

	blk_indices = blocks(PseudoBlockArray(perm, blk_sizes))
	N = length(blk_sizes)
	M = Vector{Matrix{T}}(undef, N)
	R = Vector{Matrix{T}}(undef, N-1)
	for ic in 1:N # loops over the column block index
		if ic == 1
			M[ic] = Σ[blk_indices[ic], blk_indices[ic]]
		else
			U 		= cholesky(Σ[blk_indices[ic-1], blk_indices[ic-1]]).U
			C 		= Σ[blk_indices[ic], blk_indices[ic-1]] / U
			R[ic-1] = - C / U'
			M[ic]   = Σ[blk_indices[ic], blk_indices[ic]] - C*C'
		end
	end

    return Ridiagonal(R), Midiagonal(M), Piv(perm)
end

# QUESTION: it seems redundant to have two of these R_M_P methods. Perhaps we should just have this one then use dispatch
# on vecchia(...) above to pass the correct argument
function R_M_P(Σfun::Function, blk_sizes::AbstractVector{<:Integer}, perm::AbstractVector{<:Integer}=1:sum(blk_sizes))
	@assert isperm(perm)

	blk_indices = blocks(PseudoBlockArray(perm, blk_sizes))
	N = length(blk_sizes)
	T = typeof(Σfun(1,1)) # This seems brittle. To do it right, check how `map` does it
	M = Vector{Matrix{T}}(undef, N)
	R = Vector{Matrix{T}}(undef, N-1)
	for ic in 1:N # loops over the column block index
		if ic == 1
			M[ic] = Σfun.(blk_indices[ic], blk_indices[ic]')
		else
			U 		= cholesky(Σfun.(blk_indices[ic-1], blk_indices[ic-1]')).U
			C 		= Σfun.(blk_indices[ic], blk_indices[ic-1]') / U
			R[ic-1] = - C / U'
			M[ic]   = Σfun.(blk_indices[ic], blk_indices[ic]') - C*C'
		end
	end

    return Ridiagonal(R), Midiagonal(M), Piv(perm)
end


# These examples are included in README.jl

using GFlops
x = rand(1000);
@count_ops sum($x)
@gflops sum($x);

function mixed_dot(x, y)
    acc = 0.0
    @inbounds @simd for i in eachindex(x, y)
        acc += x[i] * y[i]
    end
    acc
end
x = rand(Float32, 1000); y = rand(Float32, 1000);
cnt = @count_ops mixed_dot($x, $y)
fieldnames(GFlops.Counter)
cnt.add64
@gflops mixed_dot($x, $y);

x = 0.5; coeffs = rand(10);
println("# 9 MulAdds but 18 flop")
cnt = @count_ops evalpoly($x, $coeffs)
@gflops evalpoly($x, $coeffs);

using LinearAlgebra
x = rand(1000);
@count_ops dot($x, $y)


# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
# ==============================================================================
#
#   Dipole model for the Earth geomagnetic field.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# References
# ==============================================================================
#
#   [1] http://helios.fmi.fi/~juusolal/geomagnetism/Lectures/Chapter3_dipole.pdf
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

export geomag_dipole

"""
    geomag_dipole(r_e::AbstractVector{T}, year::Number = 2019) where T

Compute the geomagnetic field [nT] using the simplified dipole model at position
`r_e` (ECEF reference frame). This function considers that the latitude of the
South magnetic pole (which lies in the North hemisphere) is `pole_lat` [rad] and
the longitude is `pole_lon` [rad]. Furthermore, the dipole moment is considered
to be `m` [A.m²].

    geomag_dipole(r_e::AbstractVector, year::Number = 2019)

Compute the geomagnetic field [nT] using the simplified dipole model at position
`r_e` (ECEF reference frame). This function uses the year `year` to obtain the
position of the South magnetic pole (which lies in the North hemisphere) and the
dipole moment. If `year` is omitted, then it will be considered as 2019.

# Remarks

1. In both functions, the output vector will be represented in the ECEF
    reference frame.
2. The returned vector will have the same type `T` of the input vector.
"""
function geomag_dipole(r_e::AbstractVector{T}, year::Number = 2019) where T
    d2r = T(π / 180)

    # Obtain the pole location and the dipole moment.
    pole_lat = T(_itp_dipole_lat(year)) * d2r
    pole_lon = T(_itp_dipole_lon(year)) * d2r
    m        = T(_itp_dipole_mag(year)) * T(1e22)

    return geomag_dipole(r_e, pole_lat, pole_lon, m)
end

function geomag_dipole(
    r_e::AbstractVector{T},
    pole_lat::Number,
    pole_lon::Number,
    m::Number
) where T
    # DCM that converts the ECEF into the geomagnetic coordinates.
    Dge = angle_to_dcm(pole_lon, T(π / 2) - pole_lat, 0, :ZYX)

    # Compute the dipole momentum represented in the ECEF reference frame.
    k₀_e = T(1e-7) * m * (Dge' * SVector{3, T}(0, 0, -1))

    # Compute the distance from the Earth center of the desired point.
    r = norm(r_e)

    # Compute the unitary vector that points to the desired direction.
    eᵣ_e = SVector{3, T}(r_e) / r

    # Compute the geomagnetic field vector [nT].
    B_e = 1 / r^3 * (3 * eᵣ_e*eᵣ_e' - I) * k₀_e * T(1e9)

    return B_e
end