text:
export
    QWSearch,
    marked,
    penalty,
    QSearchState,
    state,
    probability,
    runtime

"""
    QWSearch(model, parameters, marked, penalty[; check])

Simulates quantum search on `model` with `marked` vertices and additional `parameters`.
`penalty` represents the cost of initial state creation and measurement, which should
be included for better optimization, see documentation of `maximizing_function`.
Note that marked vertices needs to be between `1` and `nv(graph(model))`. Furthermore
penalty needs to be nonnegative.

`check_qwdynamics` is executed if and only iff `check` is true.

Needs implementation of
* `initial_state(qws::QWSearch)`
* `evolve(qws::QWSearch{<:QWModelDiscr}, state)` or `evolve(qws::QWSearch{<:QWModelCont}, state, time::Real)`
* `measure(qws::QWSearch, state[, vertices])`
* `check_qwdynamics(QWSearch, model::QWModel, parameters::Dict{Symbol}, penalty::Real)`
* proper constructors.

Offers functions
* `execute`
* `execute_single`
* `execute_single_measured`
* `execute_all`
* `execute_all_measured`
* `maximize_quantum_search`.

It is encoureged to implement constructor, which changes the `penalty` and/or `marked`
vertices only, as their are usually simple to adapt.
"""
struct QWSearch{T,W<:Real} <: QWDynamics{T}
  model::T
  parameters::Dict{Symbol}
  marked::Vector{Int}
  penalty::W

  function QWSearch(model::T,
                    parameters::Dict{Symbol},
                    marked::Vector{Int},
                    penalty::W;
                    check::Bool=false) where {T<:QWModel, W<:Real}
    @assert all(1 <= v <= nv(model.graph) for v=marked) && marked != [] "marked vertices needs to be non-empty subset of graph vertices set"
    @assert penalty >= 0 "Penalty needs to be nonnegative"

    if check
        check_qwdynamics(QWSearch, model, parameters, marked)
    end
    new{T,W}(model, parameters, marked, penalty)
  end
end

function check_qwdynamics(qws::QWSearch)
    check_qwdynamics(QWSearch, model(qws), parameters(qws), marked(qws))
end

"""
    marked(qws)

Returns `marked` vertices element of `qws`.
"""
marked(qws::QWSearch) = qws.marked

"""
    penalty(qws)

Returns `penalty` element of `qws`.
"""
penalty(qws::QWSearch) = qws.penalty

"""
    QSearchState(state, probability, runtime, penalty)
    QSearchState(qws, state, runtime)

Creates container which consists of `state`, success probability `probability`,
running time `runtime` and penalty `penalty`.

In second case `state` is measured according to `qws`.

# Example
```jldoctest; setup = :(using QuantumWalk,LightGraphs)
julia> qws = QWSearch(Szegedy(complete_graph(4)), [1]);

julia> result = QSearchState(qws, initial_state(qws), 0)
QSearchState{Array{Float64,1},Int64,Int64}([0.0, 0.28867513459481287, 0.28867513459481287, 0.28867513459481287, 0.28867513459481287, 0.0, 0.28867513459481287, 0.28867513459481287, 0.28867513459481287, 0.28867513459481287, 0.0, 0.28867513459481287, 0.28867513459481287, 0.28867513459481287, 0.28867513459481287, 0.0], [0.25], 0, 0)
```
"""
struct QSearchState{S,Y<:Real,T<:Real}
  state::S
  probability::Vector{Float64}
  runtime::Y
  penalty::T

  function QSearchState(state::S,
                        probability::AbstractVector{Float64},
                        runtime::Y,
                        penalty::T=0.) where {S,Y<:Real,T<:Real}
     @assert runtime >= zero(Y) "runtime needs to be nonnegative"
     @assert penalty >= zero(T) "penalty needs to be nonnegative"
     @assert all(0. .<= probability .<= 1.001) && sum(probability) <= 1.001 "probability is not a subprobability vector"

     new{S,Y,T}(state, Vector(probability), runtime, penalty)
  end
end

function QSearchState(qws::QWSearch, state, runtime::Real)
   QSearchState(state, measure(qws, state, qws.marked), runtime, penalty(qws))
end

# Following for resolving method ambiguity error
function QSearchState(qws::QWSearch, state::Vector{Float64}, runtime::Real)
   QSearchState(state, measure(qws, state, qws.marked), runtime, penalty(qws))
end

"""
    state(qsearchstate)

Returns the state of qsearchstate.
"""
state(qsearchstate::QSearchState) = qsearchstate.state

"""
    probability(qsearchstate)

Returns the list of probabilities of finding marked vertices.
"""
probability(qsearchstate::QSearchState) = qsearchstate.probability

"""
    runtime(qsearchstate)

Returns the time for which the state was calulated.
"""
runtime(qsearchstate::QSearchState) = qsearchstate.runtime

"""
    penalty(qsearchstate)

Returns the penalty time for which the state was calulated.
"""
penalty(qsearchstate::QSearchState) = qsearchstate.penalty

"""
    initial_state(qws)

Generates initial state for `qws`.
"""
initial_state

include("qwsearch_dynamics.jl")
include("qwsearch_util.jl")
include("maximizing_function.jl")


text:
abstract type AbstractGrid{N} end

abstract type AbstractGridVar{T, N} <: AbstractArray{T, N} end

struct CellVar{T, N} <: AbstractGridVar{T, N}
    data::Array{T, N}
    CellVar(data::Array{T, N}) where{T, N} = new{T, N}(data)
end
struct FaceVar{T, N} <: AbstractGridVar{T, N}
    data::Array{T, N}
end

import Base.size
@inline size(type::Type{CellVar}, grid::AbstractGrid) = grid.nr_cells
@inline size(type::Type{FaceVar}, grid::AbstractGrid) = grid.nr_faces

import Base.eltype
@inline eltype(var::AbstractGridVar) = eltype(var.data)

import Base: getindex, setindex!, firstindex, lastindex, size
@inline getindex(var::AbstractGridVar, i::Int) = getindex(var.data, i)
@inline getindex(var::AbstractGridVar, I::Vararg{Int, N}) where N = getindex(var.data, I...)
@inline getindex(var::AbstractGridVar, c::Colon) = getindex(var.data, c::Colon)
@inline getindex(var::AbstractGridVar, i::Base.AbstractCartesianIndex) = getindex(var.data, i)
@inline getindex(var::AbstractGridVar, i::AbstractRange) = getindex(var.data, i)
@inline getindex(var::AbstractGridVar, i) = getindex(var.data, index(i))

@inline setindex!(var::AbstractGridVar, v, i::Int) = setindex!(var.data, v, i)
@inline setindex!(var::AbstractGridVar, v, I::Vararg{Int, N}) where N = setindex!(var.data, v, I...)
@inline setindex!(var::AbstractGridVar, v, i::Colon) = setindex!(var.data, v, i)
@inline setindex!(var::AbstractGridVar, v, i::Base.AbstractCartesianIndex) = setindex!(var.data, v, i)
@inline setindex!(var::AbstractGridVar, v, i::AbstractRange) = setindex!(var.data, v, i)
@inline setindex!(var::AbstractGridVar, v, i) = setindex!(var.data, v, index(i))

@inline firstindex(var::AbstractGridVar) = firstindex(var.data)
@inline lastindex(var::AbstractGridVar) = lastindex(var.data)
@inline size(var::AbstractGridVar, args...) = size(var.data, args...)

function CellVar(fun::Function, grid::AbstractGrid; filter::Function = _ -> true)
    var = CellVar(zeros(size(CellVar, grid)))
    @inbounds for cell ∈ cells(grid, filter = filter)
        var[cell] = fun(centroid(cell))
    end
    return var
end
function FaceVar(fun::Function, grid::AbstractGrid; filter::Function = _ -> true)
    var = FaceVar(zeros(size(FaceVar, grid)))
    @inbounds for face ∈ faces(grid, filter = filter)
        vec = fun(centroid(face))
        var[face] = vec[direction(face)]
    end
    return var
end

ext: jl


text:
using BlastLCA
using Test

@testset "BlastLCA.jl" begin
    # Write your tests here.
end


text:
# # Optimal advertising

# This example is taken from <https://web.stanford.edu/~boyd/papers/pdf/cvx_applications.pdf>.

# Setup:
#
# * We have $m$ adverts and $n$ timeslots
# * The total traffic in time slot $t$ is $T_t$
# * The number of ad $i$ displayed in period $t$ is $D_{it} \geq 0$
# * We require $\sum_{i=1}^m D_{it} \leq T_t$ since we cannot show more than $T_t$ ads during time slot $t$.
# * We require $\sum_{t=1}^n D_{it} \geq c_i$ to fulfill a contract to show advertisement $i$ at least $c_i$ times.
#
# Goal: Choose $D_{it}$.
#
# For some empirical $P_{it}$ with $0 \leq P_{it} \leq 1$, we obtain $C_{it} = P_{it}D_{it}$ clicks for ad $i$, which pays us some number $R_i > 0$ up to a budget $B_i$.
# The ad revenue for ad $i$ is $S_i = \min( R_i \sum_t C_{it}, B_i )$ which is concave in $D$.
# We aim to maximize $\sum_i S_i$.

using Random
using Distributions: LogNormal
Random.seed!(1);

m = 5; # number of adverts
n = 24; # number of timeslots
SCALE = 10000;
B = rand(LogNormal(8), m) .+ 10000;
B = round.(B, digits=3); # Budget

P_ad = rand(m);
P_time = rand(1,n);
P = P_ad * P_time;

T = sin.(range(-2*pi/2, stop=2*pi-2*pi/2, length=n)) * SCALE;
T .+= -minimum(T) + SCALE; # traffic
c = rand(m); # contractual minimum
c *= 0.6*sum(T)/sum(c);
c = round.(c, digits=3);
R = [rand(LogNormal(minimum(c)/c[i]), 1) for i=1:m]; # revenue

#-

## Form and solve the optimal advertising problem.
using Convex, SCS;
D = Variable(m, n);
Si = [min(R[i]*dot(P[i,:], D[i,:]'), B[i]) for i=1:m];
problem = maximize(sum(Si),
               [D >= 0, sum(D, dims=1)' <= T, sum(D, dims=2) >= c]);
solve!(problem, () -> SCS.Optimizer(verbose=0));

#-

# Plot traffic.
using Plots
plot(1:length(T), T, xlabel="hour", ylabel="Traffic")

#-

# Plot P.
heatmap(P)

#-

# Plot optimal D.
heatmap(evaluate(D))