export solve, Result, nresults, nsolutions, nfinite, nsingular, nat_infinity,
    nfailed, nnonsingular, nreal, ntracked, finite, results, mapresults,
    failed, at_infinity, singular, nonsingular, seed,
    solutions, real_solutions, multiplicities, statistics, multiplicities!

"""
    solve(args...; options...)::Result

The solve function takes many different arguments and options depending on your specific situation,
but in the it always returns a [`Result`](@ref) containing the result of the computations.
In the following we show the different inputs `solve` takes and at the end we list all possible options.

# Total Degree Homotopy

    solve(F; options...)

Solve the system `F` using a start system computed from the degrees of the entries of `F`. The number of paths to track is equal to the total degree `d1⋯dn`, where `di` is the degree of the `i`th entry of `F`. `F` can be
- `Vector{<:MultivariatePolynomials.AbstractPolynomial}` (e.g. constructed by `@polyvar`)
- A composition of polynomial systems constructed by [`compose`](@ref).
- [`AbstractSystem`](@ref) (the system has to represent a **homogeneous** polynomial system.)

### Example
We can solve the system ``F(x,y) = (x^2+y^2+1, 2x+3y-1)`` in the following way:
```julia
julia> @polyvar x y;
julia> solve([x^2+y^2+1, 2x+3y-1])
Result with 2 solutions
==================================
• 2 non-singular solutions (0 real)
• 0 singular solutions (0 real)
• 2 paths tracked
• random seed: 661766
```

# Polyhedral Homotopy

    solve(F; start_system = :polyhedral, only_torus=false, options...)

Solve the system `F` using a start system computed from the Newton Polytopes of the entries `F`. The number of paths to track is equal to the mixed volume of the Newton Polytopes of the entries of `F`. The mixed volume is at most the total degree of `F`. `F` can be
- `Vector{<:MultivariatePolynomials.AbstractPolynomial}` (e.g. constructed by `@polyvar`)
- A composition of polynomial systems constructed by [`compose`](@ref).
- [`AbstractSystem`](@ref) (the system has to represent a **homogeneous** polynomial system.)

If `only_torus == true` then only solutions in the algebraic torus ``(ℂ\\setminus \\{0\\})^n`` will be computed.

### Example
We can solve the system ``F(x,y) = (x^2+y^2+1, 2x+3y-1)`` in the following way:
```julia
julia> @polyvar x y;
julia> solve([x^2+y^2+1, 2x+3y-1]; start_system = :polyhedral)
Result with 2 solutions
==================================
• 2 non-singular solutions (0 real)
• 0 singular solutions (0 real)
• 2 paths tracked
• random seed: 222880
```

# Homogeneous Systems

If `F` has is homogeneous, we return results in projective space

### Examples
```julia
julia> @polyvar x y z;
julia> solve([x^2+y^2+z^2, 2x+3y-z])
Result with 2 solutions
==================================
• 2 non-singular solutions (0 real)
• 0 singular solutions (0 real)
• 2 paths tracked
• random seed: 291729
```

If your polynomial system is already homogeneous, but you would like to consider it as an affine system
you can do
```julia
@polyvar x y z
solve([x^2+y^2+z^2, 2x+3y-z], homvar=z)
```
This yields the same result as `solve([x^2+y^2+1, 2x+3y-1])`.

# Multihomogeneous Systems

By exploiting the multi-homogenous structure of a polynomial system it is possible
to decrease the number of paths necessary to track.
```julia
@polyvar x y
# Use variable groups to only track 2 paths instead of 4
solve([x*y - 6, x^2 - 5], variable_groups=[(x,), (y,)])
```
To check whether a certain variable grouping is beneficial you can use the [`bezout_number`](@ref)
function.

# Start Target Homotopy

    solve(G, F, start_solutions; options...)

This constructs the homotopy ``H(x,t) = tG(x)+(1-t)F(x)`` to compute solutions of the
system `F`. `start_solutions` is a list of solutions of `G` which are tracked to solutions
of `F`.

### Example
```julia
@polyvar x y
G = [x^2-1,y-1]
F = [x^2+y^2+z^2, 2x+3y-z]
solve(G, F, [[1, 1], [-1, 1]])
```

# Parameter Homotopy

    solve(F, startsolutions; parameters, start_parameters, target_parameters, start_gamma=nothing, target_gamma=nothing)

Solve the parameter homotopy
```math
H(x, t) = F(x, (tγ₁p₁+(1-t)γ₀p₀) / (tγ₁+(1-t)γ₀)),
```
where ``p₁`` (=`start_parameters`) and ``p₀`` (=`target_parameters`) are vectors of parameter values for ``F`` and
``γ₁`` (=`start_gamma`) and ``γ₀`` (=`target_gamma`) are complex numbers.
If `start_parameters` or `target_parameters` is `nothing`, it is assumed that `γ₁` and `γ₀` are ``1``.
The input `parameters` specifies the variables of `F` which should be considered as parameters.
Necessarily we have `length(parameters) == length(p₁) == length(p₀)`.

    solve(F, startsolutions; parameters, p₁, p₀, γ₁=nothing, γ₀=nothing)

This is a unicode variant where `γ₁=start_parameters`, `γ₀=target_parameters`,
    `γ₁=start_gamma`, γ₀=`target_gamma`.

### Example
We want to solve a parameter homotopy ``H(x,t) := F(x; t[1, 0]+(1-t)[2, 4])`` where
```math
F(x; a) := (x₁^2-a₁, x₁x₂-a₁+a₂)
```
and let's say we are only intersted in tracking of ``[1,1]``.
This can be accomplished as follows
```julia
@polyvar x[1:2] a[1:2]
F = [x[1]^2-a[1], x[1]*x[2]-a[1]+a[2]]
startsolutions = [[1, 1]]
p₁ = [1, 0]
p₀ = [3im, 0.5+2im]
solve(F, startsolutions; parameters=a, start_parameters=p₁, target_parameters=p₀)
# If you like unicode this is also possible
solve(F, startsolutions; parameters=a, p₁=p₁, p₀=p₀)
```

# Abstract Homotopy

    solve(H::AbstractHomotopy, start_solutions; options...)

Solve the homotopy `H` by tracking the each solution of
``H(⋅, t)`` (as provided by `start_solutions`) from ``t=1`` to ``t=0``.
Note that `H` has to be a homotopy between *homogeneous* polynomial systems.
If it should be considered as an affine system indicate which is the index
of the homogenization variable, e.g. `solve(H, startsolutions, homvar=3)`
if the third variable is the homogenization variable.

# Options
### General options:

* `seed::Int`: The random seed used during the computations.
* `show_progress=true`: Whether a progress bar should be printed to standard out.
* `threading=true`: Enable or disable multi-threading.
* `path_result_details=:default`: The amount of information computed in each path result. Possible values are `:minimal` (minimal details), `:default` (default) and `:extensive` (all information possible).
* `homvar::Union{Int,MultivariatePolynomials.AbstractVariable}`: This considers the *homogeneous* system `F` as an affine system which was homogenized by `homvar`. If `F` is an `AbstractSystem` `homvar` is the index (i.e. `Int`) of the homogenization variable. If `F` is an `AbstractVariables` (e.g. created by `@polyvar x`) `homvar` is the actual variable used in the system `F`.
* `system::AbstractSystem`: A constructor to assemble a [`AbstractSystem`](@ref). The default is [`SPSystem`](@ref). This constructor is only applied to the input of `solve`. The constructor is called with `system(polynomials, variables)` where `polynomials` is a vector of `MultivariatePolynomials.AbstractPolynomial`s and `variables` determines the variable ordering. If you experience significant compilation times, consider to change system to `FPSystem`.
* `homotopy::AbstractHomotopy`: A constructor to construct a [`AbstractHomotopy`](@ref) for the totaldegree and start target homotopy. The default is [`StraightLineHomotopy`](@ref). The constructor is called with `homotopy(start, target)` where `start` and `target` are homogeneous [`AbstractSystem`](@ref)s.
* `affine_tracking::Bool=true`: Indicate whether path tracking should happen in affine space.
* `projective_tracking::Bool=false`: Indicate whether path tracking should happen in projective space. The flag `affine_tracking` is dominant.
* `path_jumping_check::Bool=true`: Enable a check whether one of the paths jumped to another one.

### Path tracking specific options:
* `accuracy=1e-7`: The accuracy required during the path tracking.
* `corrector::AbstractCorrector`: The corrector used during in the predictor-corrector scheme. The default is [`NewtonCorrector`](@ref).
* `initial_step_size=0.1`: The size of the first step.
* `max_corrector_iters=2`: The maximal number of correction steps in a single step.
* `max_lost_digits::Real`: The tracking is terminated if we estimate that we loose more than `max_lost_digits` in the linear algebra steps. This threshold depends on the `precision` argument.
* `max_refinement_iters=5`: The maximal number of correction steps used to refine the final value.
* `max_steps=1_000`: The maximal number of iterations the path tracker has available. Note that this changes to `10_000` for parameter homotopies.
* `max_step_size=Inf`: The maximal step size.
* `min_step_size=1e-14`: The minimal step size.
* `precision::PrecisionOption=PRECISION_FIXED_64`: The precision used for evaluating the residual in Newton's method.
* `predictor::AbstractPredictor`: The predictor used during in the predictor-corrector scheme. The default is [`Heun`](@ref)()`.
* `refinement_accuracy=1e-8`: The precision required for the final value.
* `simple_step_size_alg=false`: Use a more simple step size algorithm.
* `steps_jacobian_info_update::Int=1`: Every n-th step a linear system will be solved using a QR factorization to obtain an estimate for the condition number of the Jacobian.
* `terminate_ill_conditioned::Bool=true`: Indicates whether the path tracking should be terminated for ill-conditioned paths. A path is considerd ill-conditioned if the condition number of the Jacobian is larger than ≈1e14 or if it is larger than 1e`max_lost_digits`.

### Endgame specific options:
* `accuracy_eg::Float64=min(accuracy, 1e-5))`: It is possible to change the accuracy during the path tracking. Usually you want lower the accuracy.
* `cond_eg_start::Float64=1e4`: The endgame is only started if the condition of the Jacobian is larger than this threshold.
* `max_winding_number::Int=12`: This limits the maximal number of loops taken in applying Cauchy's formula.
* `min_cond_at_infinity::Float64=1e7`: A path is declared as going to infinity only if it's Jacobian is also larger than this threshold.
* `samples_per_loop::Int=12`: To compute singular solutions Cauchy's integral formula is used. The accuracy of the solutions increases with the number of samples per loop.
* `t_eg_start::Float64=0.1`: The endgame starts only if `t` is smaller than this threshold.
* `tol_val_inf_accurate::Float64=1e-4`: A valuation which would result in a path declared as going to infinity is only accepted if the estimated accuracy of the valuation is less than this threshold.
* `tol_val_finite_accurate::Float64=1e-3`: A valuation which would result in a proper solution is only accepted if the estimated accuracy of the valuation is less than this threshold. This is only affects solutions where the path has at some point near 0 a condition number larger than `cond_eg_start`.
It is recommended to also take a look at the [`PathTracker`](@ref) documentation for some context.

### Overdetermined system specific options:
* `overdetermined_min_accuracy=1e-5`: The minimal accuracy a non-singular solution needs to have to be considered a solution of the original system.
* `overdetermined_min_residual=1e-3`: The minimal residual a singular solution needs to have to be considered a solution of the original system.
"""
function solve end

const solve_supported_keywords = [:threading, :show_progress,
            :path_result_details, :save_all_paths, :path_jumping_check,
            # deprecated
            :report_progress]

function solve(args...; kwargs...)
    solve_kwargs, rest = splitkwargs(kwargs, solve_supported_keywords)
    tracker, start_solutions = pathtracker_startsolutions(args...; rest...)
    solve(tracker, start_solutions; solve_kwargs...)
end

function solve(tracker::Union{<:PathTracker,<:PolyhedralTracker}, start_solutions;
        path_result_details=:default, save_all_paths=false,
        path_jumping_check=true, kwargs...)
    all_results, ntracked = track_paths(tracker, start_solutions; path_result_details=path_result_details, save_all_paths=save_all_paths, kwargs...)
    if path_jumping_check
        path_jumping_check!(all_results, tracker, path_result_details; finite_results_only=!save_all_paths)
    end

    Result(all_results, ntracked, seed(tracker); multiplicity_tol=100*accuracy(tracker))
end

accuracy(T::PathTracker) = T.options.accuracy
accuracy(T::PolyhedralTracker) = accuracy(T.generic_tracker)

mutable struct SolveStats
    regular::Int
    regular_real::Int
    singular::Int
    singular_real::Int
end
SolveStats() = SolveStats(0,0,0,0)

function track_paths(tracker, start_solutions;
                threading=true, show_progress=true,
                path_result_details::Symbol=:default, save_all_paths=false,
                # deprecated
                report_progress=nothing)
    results = Vector{result_type(tracker)}()
    n = length(start_solutions)

    @deprecatekwarg report_progress show_progress

    if show_progress
        progress = ProgressMeter.Progress(n; dt=0.1, desc="Tracking $n paths... ",
                                    clear_output_ijulia=true, delay=0.3)
    else
        progress = nothing
    end

    stats = SolveStats()
    ntracked = 0
    try
        nthreads = Threads.nthreads()
        if threading && nthreads > 1
            batch_size = 32 * nthreads
            batches = BatchIterator(start_solutions, batch_size)
            batch_tracker = BatchTracker(tracker, batches, path_result_details, save_all_paths)
            for batch in batches
                prepare_batch!(batch_tracker, batch, ntracked)
                ccall(:jl_threading_run, Ref{Cvoid}, (Any,), batch_tracker)

                for R in batch_tracker.results
                    if R !== nothing
                        push!(results, R)
                        is_success(R) && update!(stats, R)
                    end
                end
                ntracked += length(batch)
                if batch_tracker.interrupted
                    return results, ntracked
                end

                update_progress!(progress, ntracked, stats)
            end
        else
            for (k, s) in enumerate(start_solutions)
                return_code = track!(tracker, s)
                if save_all_paths ||
                   return_code == PathTrackerStatus.success ||
                   return_code == PathTrackerStatus.terminated_invalid_startvalue

                    R = PathResult(tracker, s, k; details=path_result_details)
                    push!(results, R)

                    if return_code == PathTrackerStatus.success
                        update!(stats, R)
                    end
                    ntracked = k
                end
                k % 32 == 0 && update_progress!(progress, k, stats)
            end
            # don't print if it already got printed above
            n % 32 != 0 && update_progress!(progress, n, stats)
        end
    catch e
        if isa(e, InterruptException)
            return results, ntracked
        else
            rethrow(e)
        end
    end
    results, n
end

function update!(stats::SolveStats, R::PathResult)
    if is_singular(R)
        stats.singular_real += is_real(R)
        stats.singular += 1
    else
        stats.regular_real += is_real(R)
        stats.regular += 1
    end
end

function update_progress!(progress, ntracked, stats::SolveStats; finished::Bool=false)
    progress === nothing && return nothing

    nsols = stats.regular + stats.singular
    nreal = stats.regular_real + stats.singular_real

    showvalues = (
        ("# paths tracked", ntracked),
        ("# non-singular solutions (real)", "$(stats.regular) ($(stats.regular_real))"),
        ("# singular solutions (real)", "$(stats.singular) ($(stats.singular_real))"),
        ("# total solutions (real)", "$nsols ($nreal)")
    )

    ProgressMeter.update!(progress, ntracked; showvalues=showvalues)
    nothing
end

mutable struct BatchTracker{Tracker<:PathTracker, V, R} <: Function
    trackers::Vector{Tracker}
    results::Vector{Union{Nothing, R}}
    ranges::Vector{UnitRange{Int}}
    batch::Vector{V}
    details::Symbol
    all_paths::Bool
    interrupted::Bool
    ntracked::Int
end

function BatchTracker(tracker::PathTracker, batches::BatchIterator, path_result_details::Symbol, save_all_paths::Bool)
    trackers = Threads.resize_nthreads!([tracker])
    batch_size = length(batches.batch)
    batch_results = Vector{Union{Nothing, result_type(tracker)}}(undef, batch_size)
    batch_results .= nothing
    ranges = partition_work(1:batch_size, length(trackers))
    BatchTracker(trackers, batch_results, ranges, batches.batch,
                 path_result_details, save_all_paths, false, 0)
end

function prepare_batch!(batch_tracker::BatchTracker, batch, ntracked)
    n = length(batch)
    if length(batch_tracker.results) ≠ n
        resize!(batch_tracker.results, n)
    end
    batch_tracker.results .= nothing
    batch_tracker.ntracked = ntracked
    partition_work!(batch_tracker.ranges, 1:n, length(batch_tracker.trackers))
end

function (batch::BatchTracker)()
    tid = Threads.threadid()
    try
        track_batch!(batch.results, batch.trackers[tid], batch.ranges[tid],
                 batch.batch, batch.details, batch.all_paths, batch.ntracked)
     catch e
         if isa(e, InterruptException)
             batch.interrupted = true
         else
             rethrow(e)
         end
     end
end
function track_batch!(results, pathtracker, range, starts, details, all_paths, ntracked)
    for k in range
        return_code = track!(pathtracker, starts[k])
        if all_paths ||
           return_code == PathTrackerStatus.success ||
           return_code == PathTrackerStatus.terminated_invalid_startvalue

            results[k] = PathResult(pathtracker, starts[k], ntracked+k; details=details)
        else
            results[k] = nothing
        end
    end
    nothing
end

"""
    path_jumping_check!(results, tracker, details)

Try to detect path jumping by comparing the winding numbers of finite results.
"""
function path_jumping_check!(results::Vector{<:PathResult}, tracker, details::Symbol; finite_results_only=false)
    if finite_results_only
        finite_results_indices = collect(1:length(results))
        finite_results = results
    else
        finite_results_indices = Int[]
        finite_results = Vector{eltype(results)}()
        for (i, r) in enumerate(results)
            if isfinite(r)
                push!(finite_results, r)
                push!(finite_results_indices, i)
            end
        end
    end
    !isempty(finite_results) || return results

    tol = tracker.core_tracker.options.refinement_accuracy
    # find cluster of multiple solutions
    clusters = multiplicities(solution, finite_results; tol=tol)
    while true
        rerun_paths = false
        for cluster in clusters
            m = length(cluster)
            # We have to consider a couple of cases
            # 1) All m solutions in the cluster also have winding number m
            #    -> Everything OK
            # 2) We find two non-singular solutions (winding_number === nothing)
            #    in the same cluster
            #    -> Path jumping happened
            # 3) There can be multiple paths going to the same solution, i.e.,
            #    winding number | m
            #    -> OK
            # 4) It can happen that paths going to a positive dimensional component
            #    have the same endpoint and still winding number = 1.
            #
            # This is quite complicated to tear apart. Therefore we focus for
            # now only on case 2)

            jumping = false
            for i in cluster
                cond = unpack(finite_results[i].condition_jacobian, 1.0)
                w = unpack(finite_results[i].winding_number, 0)
                if w ≤ 1 && cond < 1e10
                    jumping = true
                    break
                end
            end

            if jumping
                rerun_paths = true
                # rerun
                for i in cluster
                    rᵢ = finite_results[i]
                    new_rᵢ = track(tracker, start_solution(rᵢ);
                                            path_number=rᵢ.path_number,
                                            details=details,
                                            accuracy=min(1e-7, accuracy(tracker.core_tracker)),
                                            max_corrector_iters=1)
                    finite_results[i] = new_rᵢ
                    results[finite_results_indices[i]] = new_rᵢ
                end
            end
        end

        rerun_paths || break

        prev_clusters = clusters
        clusters = multiplicities(solution, finite_results; tol=tol)
        if clusters == prev_clusters
            break
        end
    end

    results
end
function path_jumping_check!(results::Vector{<:PathResult}, tracker::PolyhedralTracker, details::Symbol; kwargs...)
    path_jumping_check!(results, tracker.generic_tracker, details; kwargs...)
end

######################
## MultiplicityInfo ##
######################

"""
    MultiplicityInfo

This contains informations about the multiplicities of the solutions.
"""
struct MultiplicityInfo
    multiplicities::Dict{Int, Vector{Vector{Int}}}
    # This set stores all path numbers which we want to ignore if we don't show multiple results
    multiple_indicator::Set{Int32}
end

function MultiplicityInfo(pathresults::Vector{<:PathResult}; tol=1e-6)
    multiple_indicator = Set{Int32}()
    multiplicities = compute_multiplicities(pathresults; tol=float(tol))
    for clusters in values(multiplicities), cluster in clusters
        for i in 2:length(cluster)
            push!(multiple_indicator, pathresults[cluster[i]].path_number)
        end
    end
    MultiplicityInfo(multiplicities, multiple_indicator)
end

function compute_multiplicities(result::Vector{<:PathResult}; kwargs...) where T
    D = Dict{Int, Vector{Vector{Int}}}()
    compute_multiplicities!(D, result; kwargs...)
end
function compute_multiplicities!(D::Dict, result::Vector{<:PathResult}; tol::Float64 = 1e-6) where T
    for m in multiplicities(solution, result, tol = tol)
        if haskey(D, length(m))
            push!(D[length(m)], m)
        else
            D[length(m)] = [m]
        end
    end
    D
end

"""
    assign_multiplicities!(path_results::Vector{<:PathResult}, I::MultiplicityInfo)

Assign the path results their multiplicity according to the multiplicity info `I`.
"""
function assign_multiplicities!(path_results::Vector{<:PathResult}, I::MultiplicityInfo)
    #assign multiplicities
    for r in path_results
        r.multiplicity = 1
    end
    for (k, clusters) in I.multiplicities, cluster in clusters, vᵢ in cluster
        path_results[vᵢ].multiplicity = k
    end
    path_results
end

is_multiple_result(r::PathResult, I::MultiplicityInfo) = r.path_number ∈ I.multiple_indicator

"""
    Result{V<:AbstractVector}

The result of `solve`. This is a wrapper around the results of each single path ([`PathResult`](@ref)) and it contains some additional informations like
a random seed to replicate the result.
"""
mutable struct Result{V}
    pathresults::Vector{PathResult{V}}
    tracked_paths::Int
    seed::Int
    multiplicity_info::MultiplicityInfo
end

function Result(pathresults::Vector{PathResult{V}}, tracked_paths::Int, seed::Int;
                    multiplicity_tol=1e-6) where {V}
    multiplicity_info = MultiplicityInfo(pathresults; tol=multiplicity_tol)
    assign_multiplicities!(pathresults, multiplicity_info)
    Result(pathresults, tracked_paths, seed, multiplicity_info)
end

Base.length(r::Result) = length(r.pathresults)
Base.getindex(r::Result, I) = getindex(r.pathresults, I)

Base.iterate(r::Result) = iterate(r.pathresults)
Base.iterate(r::Result, state) = iterate(r.pathresults, state)
Base.lastindex(r::Result) = lastindex(r.pathresults)
Base.eltype(r::Type{Result{V}}) where {V} = PathResult{V}

is_multiple_result(r::PathResult, R::Result) = is_multiple_result(r, R.multiplicity_info)
is_multiple_result(r::PathResult, R::Vector{<:PathResult}) = false

"""
    multiplicities!(result::Result; tol=1e-6)

Compute the multiplicities of the solutions in `result` with respect to the given tolerance.
"""
function multiplicities!(result::Result; tol=1e-6)
    multiplicity_info = MultiplicityInfo(result.pathresults; tol=tol)
    result.multiplicity_info = multiplicity_info
    assign_multiplicities!(result.pathresults, multiplicity_info)
    result
end

const Results = Union{Result, Vector{<:PathResult}}
const ProjectiveResult = Result{<:PVector}

"""
    nresults(result; only_real=false, real_tol=1e-6, only_nonsingular=false, singular_tol=1e10, onlyfinite=true)

The number of solutions which satisfy the corresponding predicates.

## Example
```julia
result = solve(F)
# Get all non-singular results where all imaginary parts are smaller than 1e-8
nresults(result, only_real=true, real_tol=1e-8, only_nonsingular=true)
```
"""
function nresults(R::Results; only_real=false, real_tol=1e-6,
    only_nonsingular=false, only_singular=false, singular_tol=1e10, onlyfinite=true, multiple_results=false)
    count(R) do r
        (!only_real || is_real(r, real_tol)) &&
        (!only_nonsingular || is_nonsingular(r, singular_tol)) &&
        (!only_singular || is_singular(r, singular_tol)) &&
        (!onlyfinite || isfinite(r) || is_projective(r)) &&
        (multiple_results || !is_multiple_result(r, R))
    end
end

"""
    statistics(R::Result; only_real=false, real_tol=1e-6,
        only_nonsingular=false, only_singular=false, singular_tol=1e10)

Statistic about the number of (real) singular and non-singular solutions etc. Returns a named tuple with the statistics.

## Example
```julia
julia> statistics(solve([x^2+y^2-2, 2x+3y-1]))
(nonsingular = 2, singular = 0, real_nonsingular = 2, real_singular = 0, real = 2, at_infinity = 0, failed = 0, total = 2)
"""
function statistics(R::Results, only_real=false, real_tol=1e-6,
    only_nonsingular=false, only_singular=false, singular_tol=1e10)

    failed = at_infinity = nonsingular = singular = real_nonsingular = real_singular = 0
    singular_with_multiplicity = real_singular_with_multiplicity = 0
    for r in R
        is_multiple_result(r, R) && continue

        if is_failed(r)
            failed += 1
        elseif is_singular(r, singular_tol)
            if is_real(r, real_tol)
                real_singular += 1
                real_singular_with_multiplicity += unpack(multiplicity(r), 1)
            end
            singular += 1
            singular_with_multiplicity += unpack(multiplicity(r), 1)
        elseif !is_projective(r) && !isfinite(r)
            at_infinity += 1
        else # finite, nonsingular
            if is_real(r, real_tol)
                real_nonsingular += 1
            end
            nonsingular += 1
        end
    end
    (nonsingular = nonsingular,
    singular = singular,
    singular_with_multiplicity=singular_with_multiplicity,
    real_nonsingular = real_nonsingular,
    real_singular = real_singular,
    real_singular_with_multiplicity=real_singular_with_multiplicity,
    real = real_nonsingular + real_singular,
    at_infinity = at_infinity,
    failed = failed,
    total = R.tracked_paths)
end

"""
    nfinite(result)

The number of finite solutions.
"""
nfinite(R::Results) = count(isfinite, R)

"""
    nsolutions(result)

The number of solutions.
"""
nsolutions(R::Results) = nresults(R)

"""
    nsingular(result; singular_tol=1e10, multiplicitytol=1e-5, counting_multiplicities=false, kwargs...)

The number of singular solutions. A solution is considered singular if its windingnumber is larger than 1 or the condition number is larger than `tol`.
If `counting_multiplicities=true` the number of singular solutions times their multiplicities is returned.
"""
function nsingular(R::Results; singular_tol=1e10, counting_multiplicities=false, kwargs...)
    S = results(R; only_singular=true, multiple_results=false, singular_tol=singular_tol, kwargs...)
    isempty(S) && return 0
    counting_multiplicities && return sum(multiplicity, S)
    length(S)
end

"""
    nat_infinity(result)

The number of solutions at infinity.
"""
nat_infinity(R::Results) = count(is_at_infinity, R)

"""
    nafailed(result)

The number of failed paths.
"""
nfailed(R::Results) = count(is_failed, R)

"""
    nnonsingular(result; tol=1e-10)

The number of non-singular solutions.
"""
nnonsingular(R::Result; tol = 1e10) = count(r -> is_nonsingular(r, tol), R)

"""
    nreal(result; tol=1e-6)

The number of real solutions where all imaginary parts of each solution
are smaller than `tol`.
"""
nreal(R::Results; tol = 1e-6) = count(r -> is_real(r, tol), R)

"""
    ntracked(R::Result)

Returns the total number of paths tracked.
"""
ntracked(R::Result) = R.tracked_paths

"""
    seed(result)

The random seed used in the computation.
"""
seed(result::Result) = result.seed

# Filtering
"""
    results(result; only_real=false, real_tol=1e-6, only_nonsingular=false,
                onlysigular=false, singular_tol=1e10, onlyfinite=true, multiple_results=false)

Return all `PathResult`s for which the given conditions apply.

## Example

```julia
R = solve(F)

# This gives us all PathResults considered non-singular and real (but still as a complex vector).
real_solutions = results(R, only_real=true, only_nonsingular=true)
```
"""
results(R::Results; kwargs...) = mapresults(identity, R; kwargs...)
results(f::Function, R::Results; kwargs...) = mapresults(f, R; kwargs...)

"""
    mapresults(f::Function, result; conditions...)

Apply the function `f` to all `PathResult`s for which the given conditions apply. For the possible
conditions see [`results`](@ref).

## Example
```julia
# This gives us all solutions considered real (but still as a complex vector).
real_solutions = mapresults(solution, R, only_real=true)
```
"""
function mapresults(f::Function, R::Results;
    only_real=false, real_tol=1e-6, only_nonsingular=false, only_singular=false, singular_tol=1e10,
    onlyfinite=true, multiple_results=false)
    [f(r) for r in R if
        (!only_real || is_real(r, real_tol)) &&
        (!only_nonsingular || is_nonsingular(r, singular_tol)) &&
        (!only_singular || is_singular(r, singular_tol)) &&
        (!onlyfinite || isfinite(r) || is_projective(r)) &&
        (multiple_results || !is_multiple_result(r,R))]
end

"""
    solutions(result; conditions...)

Return all solution (as `Vector`s) for which the given conditions apply.
For the possible `conditions` see [`results`](@ref).

## Example
```julia
julia> @polyvar x y
julia> result = solve([(x-2)y, y+x+3]);
julia> solutions(result)
[[2.0+0.0im, -5.0+0.0im], [-3.0+0.0im, 0.0+0.0im]]
```
"""
function solutions(result::Results; kwargs...)
    mapresults(solution, result; kwargs...)
end

"""
    real_solutions(result; tol=1e-6, conditions...)

Return all real solution (as `Vector`s of reals) for which the given conditions apply.
For the possible `conditions` see [`results`](@ref). Note that `only_real` is always `true`
and `real_tol` is now `tol`.

## Example
```julia
julia> @polyvar x y
julia> result = solve([(x-2)y, y+x+3]);
julia> real_solutions(result)
[[2.0, -5.0], [-3.0, 0.0]]
```
"""
function real_solutions(result::Results; only_real=true, tol=1e-6, kwargs...)
    mapresults(r -> real.(solution(r)), result; only_real=true, real_tol=tol, kwargs...)
end
@deprecate realsolutions(result; kwargs...) real_solutions(result; kwargs...)

"""
    nonsingular(result::Results; conditions...)

Return all `PathResult`s for which the solution is non-singular. This is just a shorthand
for `results(R; only_nonsingular=true, conditions...)`. For the possible `conditions` see [`results`](@ref).
"""
nonsingular(R::Results; kwargs...) = results(R; only_nonsingular=true, kwargs...)

"""
    singular(R::Results; tol=1e10, multiple_results=false, kwargs...)

Return all `PathResult`s for which the solution is singular. A solution is labeled singular if the condition number is greater than `singular_tol`, or if the winding number is > 1.
If `multiple_results=false` only one point from each cluster of multiple solutions is returned.
If If `multiple_results=true` all singular solutions in `R` are returned.
For the possible `kwargs` see [`results`](@ref).
"""
function singular(R::Results; tol=1e10, kwargs...)
    results(R; only_singular = true, singular_tol=tol, kwargs...)
end

"""
    finite(result::AffineResults; conditions...)

Return all `PathResult`s for which the solution is finite. This is just a shorthand
for `results(R; onlyfinite=true, conditions...)`. For the possible `conditions` see [`results`](@ref).
"""
finite(R::Results; kwargs...) = results(R; onlyfinite=true, kwargs...)

"""
    real(result, tol=1e-6)

Get all results where the solutions are real with the given tolerance `tol`.
See [`is_real`](@ref) for details regarding the determination of 'realness'.
"""
Base.real(R::Results; tol=1e-6) = [r for r in R if is_real(r, tol)]

"""
    failed(result)

Get all results where the path tracking failed.
"""
failed(R::Results) = [r for r in R if is_failed(r)]

"""
    at_infinity(result::AffineResult)

Get all results where the solutions is at infinity.
"""
at_infinity(R::Results) = [r for r in R if is_at_infinity(r)]

"""
    multiplicities(V::Results; tol=1e-6)

Returns a `Vector` of `Vector{PathResult}`s grouping the `PathResult`s whose solutions appear with multiplicities *greater* 1 in 'V'.
Two solutions are regarded as equal, when their pairwise distance is less than 'tol'.
"""
function multiplicities(results::Results; tol=1e-6)
    map(i -> results[i], multiplicities(solution, results; tol = tol))
end

####Show function

function Base.show(io::IO, x::Result)
    s = statistics(x)
    println(io, "Result with $(s.nonsingular + s.singular) solutions")
    println(io, "==================================")
    println(io, "• $(s.nonsingular) non-singular $(plural("solution", s.nonsingular)) ($(s.real_nonsingular) real)")
    println(io, "• $(s.singular) singular $(plural("solution", s.singular)) ($(s.real_singular) real)")
    s.at_infinity > 0 &&
        println(io, "• $(s.at_infinity) $(plural("solution", s.at_infinity)) at infinity")
    s.failed > 0 &&
        println(io, "• $(s.failed) failed $(plural("path", s.failed))")
    println(io, "• $(ntracked(x)) paths tracked")
    println(io, "• random seed: $(seed(x))")
    if s.singular > 0
        println(io, "• multiplicity table of singular solutions:")
        singular_multiplicities_table(io, x, s)
    end
end

function Base.show(io::IO, x::ProjectiveResult)
    s = statistics(x)
    println(io, "Result with $(s.nonsingular + s.singular) solutions")
    println(io, "==================================")
    println(io, "• $(s.nonsingular) non-singular $(plural("solution", s.nonsingular)) ($(s.real_nonsingular) real)")
    println(io, "• $(s.singular) singular $(plural("solution", s.singular)) ($(s.real_singular) real)")
    s.failed > 0 &&
        println(io, "• $(s.failed) failed $(plural("path", s.failed))")
    println(io, "• $(ntracked(x)) paths tracked")
    println(io, "• random seed: $(seed(x))")
    if s.singular > 0
        println(io, "• multiplicity table of singular solutions:")
        singular_multiplicities_table(io, x, s)
    end
end

TreeViews.hastreeview(::Result) = true
TreeViews.hastreeview(::ProjectiveResult) = true
TreeViews.numberofnodes(::Result) = 8
TreeViews.numberofnodes(::ProjectiveResult) = 7
TreeViews.treelabel(io::IO, x::Result, ::MIME"application/prs.juno.inline") =
    print(io, "<span class=\"syntax--support syntax--type syntax--julia\">Result</span>")

function TreeViews.nodelabel(io::IO, x::Result, i::Int, ::MIME"application/prs.juno.inline")
    s = statistics(x)
    if i == 1
        print(io, "Paths tracked")
    elseif i == 2 && s.nonsingular > 0
        print(io, "$(s.nonsingular) non-singular ($(s.real_nonsingular) real)")
    elseif i == 3 && s.singular > 0
        print(io, "$(s.singular) singular ($(s.real_singular) real)")
    elseif i == 4 && (s.real_nonsingular+s.real_singular) > 0
        print(io, "$(s.real_nonsingular+s.real_singular) real")
    elseif i == 5 && s.at_infinity > 0
        print(io, "$(s.at_infinity) at_infinity")
    elseif i == 6 && s.failed > 0
        print(io, "$(s.failed) failed")
    elseif i == 7
        print(io, "Random seed used")
    elseif i == 8 && s.singular > 0
        print(io, "  multiplicity table of singular solutions: \n")
        singular_multiplicities_table(io, x, s)
    end
end

function TreeViews.treenode(r::Result, i::Integer)
    s = statistics(r)
    if i == 1
        return ntracked(r)
    elseif i == 2 && s.nonsingular > 0
        return finite(r, only_nonsingular=true)
    elseif i == 3 && s.singular > 0
        return finite(r, only_singular=true)
    elseif i == 4 && (s.real_nonsingular+s.real_singular) > 0
        return finite(r, only_real = true)
    elseif i == 5 && s.at_infinity > 0
        return at_infinity(r)
    elseif i == 6 && s.failed > 0
        return failed(r)
    elseif i == 7
        return seed(r)
    elseif i == 8 && s.singular > 0
        return missing
    end
    missing
end

function TreeViews.nodelabel(io::IO, x::ProjectiveResult, i::Int, ::MIME"application/prs.juno.inline")
    s = statistics(x)
    if i == 1
        print(io, "Paths tracked")
    elseif i == 2 && s.nonsingular > 0
        print(io, "$(s.nonsingular) non-singular ($(s.real_nonsingular) real)")
    elseif i == 3 && s.singular > 0
        print(io, "$(s.singular) singular ($(s.real_singular) real)")
    elseif i == 4 && (s.real_nonsingular + s.real_singular) > 0
        print(io, "$(s.real_nonsingular + s.real_singular) real solutions")
    elseif i == 5 && s.failed > 0
        print(io, "$(s.failed) failed")
    elseif i == 6
        print(io, "Random seed used")
    elseif i == 7 && s.singular > 0
        print(io, "  multiplicity table of singular solutions: \n")
        singular_multiplicities_table(io, x, s)
    end
end

function TreeViews.treenode(r::ProjectiveResult, i::Integer)
    s = statistics(r)
    if i == 1
        return length(r)
    elseif i == 2 && s.nonsingular > 0
        return finite(r, only_nonsingular=true)
    elseif i == 3 && s.singular > 0
        return finite(r, only_singular=true)
    elseif i == 4 && (s.real_nonsingular + s.real_singular) > 0
        return finite(r, only_real=true)
    elseif i == 5 && s.failed > 0
        return failed(r)
    elseif i == 6
        return seed(r)
    elseif i == 7 && s.singular > 0
        return missing
    end
    missing
end

plural(singularstr, n) = n == 1 ? singularstr : singularstr * "s"

function singular_multiplicities_table(io, result::Result, stats = statistics(result))
    M = result.multiplicity_info.multiplicities
    if isempty(M)
        n_higher_mults_total = 0
    else
        n_higher_mults_total = sum(length, values(M))
    end

    headers = ["mult.", "total", "# real", "# non-real"]
    mult_1_exists = n_higher_mults_total < stats.singular
    data = Matrix{Int}(undef, mult_1_exists+length(M), 4)

    n_higher_real = 0
    i = mult_1_exists + 1
    for k in sort(collect(keys(M)))
    	data[i, 1] = k
    	n_real_solsᵢ = count(Mᵢ -> is_real(result.pathresults[Mᵢ[1]]), M[k])
    	n_solsᵢ = length(M[k])
        data[i, 2] = n_solsᵢ
    	data[i, 3] = n_real_solsᵢ
    	data[i, 4] = n_solsᵢ - n_real_solsᵢ
    	n_higher_real += n_real_solsᵢ
        i += 1
    end

    if mult_1_exists
        data[1,1] = 1
        data[1,2] = stats.singular - n_higher_mults_total
    	data[1,3] = stats.real_singular - n_higher_real
    	data[1,4] = data[1,2] - data[1,3]
    end

    PrettyTables.pretty_table(io, data, headers,
            PrettyTables.unicode;
            alignment=:c,
            header_crayon=PrettyTables.Crayon(bold=false),
            border_crayon=PrettyTables.Crayon(faint=true))
end


using StatsBase
using Base.Test

# rle & inverse_rle

z = [1, 1, 2, 2, 2, 3, 1, 2, 2, 3, 3, 3, 3]
(vals, lens) = rle(z)
@test vals == [1, 2, 3, 1, 2, 3]
@test lens == [2, 3, 1, 1, 2, 4]
@test inverse_rle(vals, lens) == z

z = [true, true, false, false, true, false, true, true, true]
vals, lens = rle(z)
@test vals == [true, false, true, false, true]
@test lens == [2, 2, 1, 1, 3]
@test inverse_rle(vals, lens) == z

# levelsmap
a = [1, 1, 2, 2, 2, 3, 1, 2, 2, 3, 3, 3, 3, 2]
b = [true, false, false, true, false, true, true, false]

@test levelsmap(a) == Dict(2=>2, 3=>3, 1=>1)
@test levelsmap(b) == Dict(false=>2, true=>1)

# indicatormat

I = [false true  false false false;
     true  false false false true;
     false false true  true  false]

x = [2, 1, 3, 3, 2]
@test indicatormat(x, 3) == I
@test full(indicatormat(x, 3; sparse=true)) == I

x = ["b", "a", "c", "c", "b"]
@test indicatormat(x) == I
@test full(indicatormat(x; sparse=true)) == I


using Base.Test
using StrainRecon

@testset "randomM" begin

# check that the columns are unique
(m,n,N) = (3,2,100)
M = randomM(m,n,N=N)
@test size(M) == (m,n,N)
uniqcols = zeros(Int,N)
for idx = 1:N
	uniqcols[idx] = size(unique(M[:,:,idx],2),2)
end
@test all(uniqcols .== n)

# check that the column sums of each block are at most 1
(m,n,p) = (24,2,3)
M = randomM(m,n,p=p)
@test size(M) == ((p-1)*m,n)
blockcolsums = sum(reshape(M,(p-1,:)),1)
@test maximum(blockcolsums) <= 1

end


### A Pluto.jl notebook ###
# v0.14.5

using Markdown
using InteractiveUtils

# ╔═╡ aab4db99-8eaf-4b70-add7-d470e6b6b6e2
using Plots, LaTeXStrings, DifferentialEquations, Distributions

# ╔═╡ 460194e8-add2-11eb-0f92-27b8fffb5046
md"# Chapter 1 Figures. Statistics"

# ╔═╡ de27faec-4e2b-4520-9d3f-1ac15344bab1
random_count() = floor(Int, n_min + rand() * n_max)

# ╔═╡ a8fa0a9a-4fb7-4380-87dd-2534687409c4
begin
	affect1!(integrator) = integrator.u += 1
	affect2!(integrator) = integrator.u -= 1
	prod1Jump(A) = ConstantRateJump((u, p, t) -> A, affect1!)
	deg1Jump = ConstantRateJump((u, p, t) -> u, affect2!)
	u01 = 0
	t1 = (0.0, 15.0)
	dprob1 = DiscreteProblem(u01, t1)
	jprob1 = JumpProblem(dprob1, Direct(), prod1Jump(100), deg1Jump)
	continous(A) = t -> -A*ℯ^(-t) + A
end

# ╔═╡ 343e56a2-4837-49a6-aa69-5a285d362c03
sol1 = solve(jprob1, SSAStepper())

# ╔═╡ 238a13fa-6d94-448d-a236-fdde5b78ad98
begin
	plot(sol1, label="Stochastic", legend=:topleft)
	p1 = plot!(continous(100), sol1.t, label="Continuous", xlabel=L"t", ylabel=L"n", ylimit=(0,145))
end

# ╔═╡ 67a22053-294a-47d1-9793-2188398de91c
begin
	last_ones = []
	num_samples=10000
	for _ ∈ 1:num_samples
		push!(last_ones, solve(jprob1, SSAStepper()).u[end])
	end
	last_ones
end

# ╔═╡ 67f8f4d0-0e6d-4857-9ac6-ea3b9f1eca3f
begin
	# poisson(mean)
	histogram(last_ones, label="Distribution of final values", color=:orange, orientation=:h)
	p2 = plot!(rand(Poisson(100), num_samples), seriestype=:stephist, label="Sampled Poisson Distribution", color=:blue, orientation=:h, ylimit=(0,145), ticks=nothing, yaxis=false, xaxis=false)
end

# ╔═╡ 6c9d9245-317a-47ca-83c4-b00523eaf3fd
pf = plot(p1, p2)

# ╔═╡ 2be2984e-1315-41a5-b7c4-63e79fad18f2
savefig(pf, "ch1_stochastic.pdf")

# ╔═╡ Cell order:
# ╠═aab4db99-8eaf-4b70-add7-d470e6b6b6e2
# ╟─460194e8-add2-11eb-0f92-27b8fffb5046
# ╠═de27faec-4e2b-4520-9d3f-1ac15344bab1
# ╠═a8fa0a9a-4fb7-4380-87dd-2534687409c4
# ╠═343e56a2-4837-49a6-aa69-5a285d362c03
# ╠═238a13fa-6d94-448d-a236-fdde5b78ad98
# ╠═67a22053-294a-47d1-9793-2188398de91c
# ╠═67f8f4d0-0e6d-4857-9ac6-ea3b9f1eca3f
# ╠═6c9d9245-317a-47ca-83c4-b00523eaf3fd
# ╠═2be2984e-1315-41a5-b7c4-63e79fad18f2


### A Pluto.jl notebook ###
# v0.14.0

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ e46441c4-97bd-11eb-330c-97bd5ac41f9e
html"""

Section 2.6

 Random Walks II

Lecture Video

"""

# ╔═╡ 85b45a43-d7bf-4597-a1a6-329b41dce20d
begin
    import Pkg

    Pkg.activate(mktempdir())
    Pkg.add([
        Pkg.PackageSpec(name="PlutoUI", version="0.7"),
        Pkg.PackageSpec(name="Plots", version="1"),
    ])

    using PlutoUI, Plots, LinearAlgebra, SparseArrays
end

# ╔═╡ 85c26eb4-c258-4a8b-9415-7b5f7ddff02a
TableOfContents(aside=true)

# ╔═╡ 2d48bfc6-9617-11eb-3a85-279bebd21332
md"""
# Julia concepts

- Customised display of objects
- Structured matrices in Julia
- `cumsum`
- Vectors of vectors
- Concatenation of vectors: `hcat`

- `heatmap` (Plots.jl)
- `surface` (Plots.jl)

"""

# ╔═╡ 30162386-039f-4cd7-9121-a3382be3c294
md"""
# Pascal's triangle

"""

# ╔═╡ 4e7b163e-dfd0-457e-b1f3-8807a4d8060a
md"""
Let's start by thinking about Pascal's triangle. (Note that [Pascal was not the first person to study these numbers](https://en.wikipedia.org/wiki/Pascal%27s_triangle).)
"""

# ╔═╡ e8ceab7b-45db-4393-bb8e-e000ecf78d2c
pascal(N) = [binomial(n, k) for n = 0:N, k=0:N]

# ╔═╡ 2d4dffb9-39e4-48de-9688-980b96814c9f
pascal(10)

# ╔═╡ 8ff66523-bc2e-4c53-975b-8ba4f99eb1c6
md"""
The non-zero entries are the **binomial coefficients**: the $k$th entry on the $n$th row is the coefficient of $x^k$ in the expansion $(1 + x)^n$, starting from $n=0$ in the first row and $k=0$ in the first column.
"""

# ╔═╡ 2868dd57-7164-4162-8c5d-30628dedeb7a
md"""
Note that there are 0s above the main diagonal -- in other words the matrix is **lower triangular**.

Julia has special types to represent some classes of structured matrices in the standard library `LinearAlgebra` package:
"""

# ╔═╡ f6009473-d3c1-444f-88ae-814f770e811b
L = LowerTriangular(pascal(10))

# ╔═╡ 9a368602-acd3-43fb-9dff-e407a4bab930
md"""
We see that the display is special: the known "**structural**" zeros are shown as dots, instead of numbers.
"""

# ╔═╡ 67517333-175f-48c4-a915-76658cbf1304
md"""
As we have already seen, Julia also has a sparse matrix type that we could use for this, in the `SparseArrays` standard library:
"""

# ╔═╡ d6832372-d336-4a54-bbcf-d0bb70e4de64
sparse(pascal(10))

# ╔═╡ 35f14826-f1e4-4977-a31a-0f6148fe25ad
md"""
For fun, let's look at where the entries of Pascal's triangle are odd. We'll make a slider:
"""

# ╔═╡ 7468fc5d-7f35-45e2-b5fc-7e63b562bc8f
@bind n_pascal Slider(1:63, show_value=true, default=1)

# ╔═╡ 1ca8aa3b-b05d-40f6-a925-2f0248b79ca2
sparse(isodd.(pascal(n_pascal)))

# ╔═╡ 38d88b7c-3b4f-430b-8d3c-f672ab0c7a49
md"""
Note that the visual representation of a sparse matrix changes at a particular value of $n$. For larger values of $n$ the **sparsity pattern** is displayed using dots for non-zero values.
"""

# ╔═╡ f4c9b02b-738b-4de1-9e9d-05b1616bee0b
md"""
The pattern is quite striking!
"""

# ╔═╡ d1c47afa-ab7f-4543-a161-e3ceb6f11eb4
md"""
You may be interested to know that there is an alternative way to look at Pascal's triangle:
"""

# ╔═╡ bf78e00f-05d9-4a05-8512-4924ef9e25f7
[binomial(i + j, i) for i = 0:10, j = 0:10]

# ╔═╡ b948830f-ead1-4f36-a237-c998f2f7deb8
md"""
and that in fact this can be produced from the previous version using matrix multiplication!:
"""

# ╔═╡ 15223c51-8d31-4a50-a8ff-1cb7d35de454
pascal(10) * pascal(10)'

# ╔═╡ 0bd45c4a-3286-427a-a927-15869be2ebfe
md"""
## Convolutions build Pascal's triangle!
"""

# ╔═╡ 999fb608-fb1a-46cb-82ca-f3f31fe617e1
pascal(6)

# ╔═╡ 6509e69a-6e50-4816-a98f-67ba437383fb
md"""
Where do those binomial coefficients come from? That's not how we learnt to build Pascal's triangle at school!

We learned to build each number up by *summing two adjacent numbers from the previous row*. In the lower-triangular representation we we sum the number immediately above with the one to the left of it.
"""

# ╔═╡ e58976a9-1784-441e-bb76-3011538b8ad0
md"""
We can think of this as... a **convolution** with the vector $[1, 1]$ !
Recall that convolutions like this are used as **image filters**, where we want to apply the same *local* operation everywhere along a vector (i.e. one in which we modify a vector in a neighbourhood).
"""

# ╔═╡ 1efc2b68-9313-424f-9850-eb4496cc8486
md"""
# Random walks: Independent and identically-distributed random variables
"""

# ╔═╡ 6e93ffda-217b-4d46-86b5-534ddc1bae90
md"""
The discussion of Pascal's triangle and convolutions will surprisingly turn up below. For now let's go back to thinking about **random walks**.

Recall that in a random walk, at each tick of a clock we take a step in space, say left or right.

Each spatial step is random, so we can think of each step as being a **random variable** with a certain probability distribution, for example the random variable that takes the value $+1$ with probability $\frac{1}{2}$ and $-1$ also with probability $\frac{1}{2}$.

Often we will think about random walks in which each step is "the same". What do we mean by "the same"? We don't mean that the *outcome* is the same, but rather that each step has the same *probability distribution*, i.e. the same probability of taking each possible value. In other words, each step is given by *copies of the same random variable*.

Steps are also *independent* of each other (the choice of step direction right now *does not affect* the choice at the next step). Hence the steps form a collection of **independent and identically-distributed random variables**, or IID random variables.
"""

# ╔═╡ 396d2490-3cb9-4f68-8fdf-9209d2010e02
md"""
## Random walks as a cumulative sum
"""

# ╔═╡ dc1c22e8-1c7b-43b7-8421-c2ca708931a5
md"""
Let's call the random variable that describes the $i$th *step* $X_i$. Then the overall position $S_n$ of the random walk at time $n$ is given by

$$S_n = X_1 + \cdots + X_n = \sum_{i=1}^n X_i.$$

(Here we have taken the initial position $S_0$ equal to 0. If the initial position is not zero, then this instead gives the **displacement** $S_n - S_0$ at time $n$.)

Recall that by the sum of two random variables we mean a new random variable whose outcomes are sums of the outcomes, with probabilities given by taking the possible pairs of outcomes giving each new outcome.
"""

# ╔═╡ e7a52b56-322c-4478-a670-dec1013c9bd8
md"""
## Cumulative sum
"""

# ╔═╡ 5f136388-5573-4814-a088-a66278acdbbe
md"""
We previously looked at sums like this. The difference now is that we are interested in the entire **trajectory**, i.e. the sequence $S_1$, $S_2$, $\ldots$, $S_n$, $\ldots$. The trajectory is given by **partial sums**:
"""

# ╔═╡ e440dc3b-bafd-4e0c-9fe8-13fce8eea22d
md"""
$$\begin{align}
S_1 & = X_1 \\
S_2 &= X_1 + X_2 \\
S_3 &= X_1 + X_2 + X_3
\end{align}$$

etc.
"""

# ╔═╡ 89d3e90d-3685-473d-aea4-0b7c5b80d7f7
md"""
Note that $S_{n+1}$ is *not* independent of $S_n$ (or the other $S_m$). E.g. if $S_{100}$ happens to be large, then $S_{101}$ will be about as large.

Thinking about types, as we did a couple of lectures ago, we would need to define a *new* type to represent a random walk, since we cannot generate consecutive values as independent samples.
"""

# ╔═╡ 203eea14-1c68-4a9f-ab18-9a2e5f408a79
md"""
How could we calculate a trajectory of a walk? Suppose we generate steps $X_i$, e.g.
"""

# ╔═╡ b6d4b045-a39f-4236-ace2-9f631e853d1b
steps = rand( (-1, +1), 10 )

# ╔═╡ e61b56be-d334-4c8f-aa8e-887bb27c058c
md"""
The whole trajectory is given by the **cumulative sum** also called "prefix sum" or "scan". Julia has the function `cumsum` to calculate this:
"""

# ╔═╡ 782f0b9a-3793-4abb-826b-9e14d6eae690
cumsum([1, 2, 3, 4])

# ╔═╡ 1764b56a-f297-4a4e-a931-31aa987ec785
md"""
So the trajectory is given by
"""

# ╔═╡ 8082092b-b6bf-4619-8776-39fdd6c9b7c1
cumsum(steps)

# ╔═╡ 34c0b850-5e95-4eb9-8435-3aae8d124772
plot(cumsum(steps), m=:o, leg=false, size=(500, 300))

# ╔═╡ 4115f7cb-d45f-4cc2-86bf-316c074393f8
md"""
Note that in Julia this is just a convenience function; it will be no more performant than writing the `for` loop yourself (unlike in some other languages, where `for` loops are slow).
"""

# ╔═╡ 6bc2f20d-3d09-425b-a471-44090dc3161e
md"""
# Evolving probability distributions in time
"""

# ╔═╡ 98994cb9-45dc-48aa-b62d-2407f7184bee
md"""
So far we have looked at single trajectories of random walks. We can think of this as the equivalent of sampling using `rand`.

Suppose that we run many trajectories of a random walk. At a given time $t$ we can ask where all of the walks are by taking a histogram of the $S_t^{(k)}$, where the superscript $^{(k)}$ denotes the $k$th trajectory out of the collection.

Doing so gives us the *probability distribution* of the random variable $S_t$.
Let's call $p_i^t := \mathbb{P}(S_t = i)$ the probability of being at site $i$ at time $t$. Then the probability distribution at time $t$ is given by the collection of all the $p_i^t$; we can group these into a vector $\mathbf{p}^t$.

Now we can ask what happens at the *next* time step $t+1$. Taking histograms again gives the probabilities $p_j^{t+1}$ and the vector $\mathbf{p}^t$ of all of them.

But $\mathbf{p}^{t+1}$ and $\mathbf{p}^{t}$ are related in some way, since we go from $S_t$ to $S_{t+1}$ by just taking a single step.

Let's think about the case of a simple random walk in one dimension. To arrive at site $i$ at time $t+1$, we must have been in one of the *neighbouring* sites at time $t$ and jumped with probability $\frac{1}{2}$ in the direction of site $i$. Hence we have

$$p_i^{t+1} = \textstyle \frac{1}{2} (p_{i-1}^{t} + p_{i+1}^{t}).$$

This is sometimes called the **master equation** (even though that is rather a useless, non-descriptive name); it describes how the *probability distribution of the random walk evolves in time*.
"""

# ╔═╡ f11f8d7d-cd3f-4585-aab4-083b892c6d4c
md"""
## Implementing time evolution of probabilities

Let's write a function to **evolve** a probability vector to the next time step for a simple random walk:
"""

# ╔═╡ fb804fe2-58be-46c9-9200-ceb8863d052c
function evolve(p)
	p′ = similar(p)   # make a vector of the same length and type
	                  # to store the probability vector at the next time step

	for i in 2:length(p)-1   # iterate over the *bulk* of the system

		p′[i] = 0.5 * (p[i-1] + p[i+1])

	end

	# boundary conditions:
	p′[1] = 0
	p′[end] = 0

	return p′
end

# ╔═╡ 40e457b4-616c-4fab-9c8e-2e5063129597
md"""
Wait... Do you recognise this?

This is just a **convolution** again! The kernel is now $[\frac{1}{2}, 0, \frac{1}{2}]$. Apart from the extra $0$ and the $\frac{1}{2}$, this is the *same* as in Pascal's triangle... so probabilities in simple random walk behave like Pascal's triangle!
"""

# ╔═╡ 979c1fbd-c9f6-4e8b-a648-6a0210fc9e7f
md"""
Note that just as with images we have a problem at the edges of the system: We need to specify **boundary conditions** on the first and last cells. For now we have put 0s at the boundary, corresponding to probability *escaping* at the boundary: any probability that arrives at the boundary at a given time step does not stay in the system. We can think of the probability as analogous to a chemical moving through a system that leaves our system (and e.g. moves into the outside world) when it reaches an edge.
"""

# ╔═╡ 583e3a92-01e7-4b88-9be0-f1e3b3c95005
md"""
We also need to specify an *initial* condition $\mathbf{p}_0$. This tells us where our walker is at time $0$. Suppose that all walkers start at site $i=0$. We will place this in the middle of our vector. Then the *probability* of being at $0$ is 1 (certainty), while the probability at any other site is 0 (impossible):
"""

# ╔═╡ 0b26efab-4e93-4d53-9c4d-faea68d12174
function initial_condition(n)

	p₀ = zeros(n)
	p₀[n ÷ 2 + 1] = 1

	return p₀
end

# ╔═╡ b9ce5af1-84f7-4a2d-92c9-de2c5498a88d
md"""
Now let's evolve the probability vector in time by applying the `evolve` function repeatedly, starting from the initial probability distribution:
"""

# ╔═╡ b48e55b7-4b56-41aa-9796-674d04adf5df
function time_evolution(p0, N)
	ps = [p0]
	p = p0

	for i in 1:N
		p = evolve(p)
		push!(ps, copy(p))
	end

	return ps
end

# ╔═╡ 53a36c1a-0b8c-4099-8854-08d73c9f118e
md"""
Let's visualise this:
"""

# ╔═╡ 6b298184-32c6-412d-a900-b113d6bd3d53
begin
	grid_size = 101

	p0 = initial_condition(grid_size)
end

# ╔═╡ b84a7255-7b0a-4ba1-8c87-9f5d3fa32ef3
ps = time_evolution(p0, 100)

# ╔═╡ c430c4de-d9bf-44e1-aa40-6b823d718b04
md"""
Note that `ps` is a `Vector`, whose elements are each `Vector`s! I.e. we have a vector of vectors. This is often a convenient way to *build up* a matrix.
"""

# ╔═╡ 99717c6e-f713-49d5-8ee5-a08c4a464223
ps[2]

# ╔═╡ 242ea831-c421-4a76-b658-2a57fa924a4f
md"""
t = $(@bind tt Slider(1:length(ps), show_value=true, default=1))
"""

# ╔═╡ aeaef573-1e90-45f3-a7fe-31ec5e2808c4
bar(ps[tt], ylim=(0, 1), leg=false, size=(500, 300), alpha=0.5)

# ╔═╡ efe186da-3273-4767-aafc-fc8eae01dbd9
md"""
### Concatenating vectors into a matrix
"""

# ╔═╡ 61201091-b8b3-4776-9be9-4c23d5ba88ba
md"""
Now we want to visualise this in a different way, for which we must join up (**concatenate**) all the probability vectors into a matrix:
"""

# ╔═╡ 66c4aed3-a04b-4a09-b954-79e816d2a3f7
M = reduce(hcat, ps)'

# ╔═╡ b135f6be-5e82-4c72-af11-0eb0d4141dec
md"""
We can visualise the matrix with a **heat map**:
"""

# ╔═╡ e74e18e3-ad08-4a53-a803-cd53564dca65
heatmap(M, yflip=true)

# ╔═╡ ed02f00f-1bcd-43fa-a56c-7be9968614cc
md"""
We can also visualise this as a 3D surface:
"""

# ╔═╡ 8d453f89-4a4a-42d0-8a00-9b153a3f435e
plotly()

# ╔═╡ f7de29b5-2a51-45e4-a0a5-f7f602681303
surface(M)

# ╔═╡ 7e817bad-dc51-4c29-a4fc-f7a8bb3663ca
md"""
But this is not necessarily very clear, so let's draw it ourselves as stacked histograms:
"""

# ╔═╡ 403d607b-6171-431b-a058-0aad0909846f
gr()

# ╔═╡ 80e033c4-a862-443d-a198-932f5822a44e
ylabels = [-(grid_size÷2):grid_size÷2;]

# ╔═╡ c8c16c14-26b0-4f83-8135-4f862ed90686
begin
	plot(leg=false)

	for which in 1:15
		for i in 1:length(ps[which])
			plot!([which, which], [-grid_size÷2 + i, -grid_size÷2 + i], [0, ps[which][i] ], c=which, alpha=0.8, lw = 2)
		end
	end

	xlims!(1, 15)

	plot!()
end

# ╔═╡ 29612df6-203f-42bc-b53b-86af618d60ec
let
	color_list = [:red, RGB(0, 1, 0.1), :blue]
	xs = []
	ys = []
	zs = []
	cs = []
	cs2 = []

	for which in 1:15
		for i in 1:length(ps[which])
			push!(xs, which, which, NaN)
			push!(ys, ylabels[i], ylabels[i], NaN)
			push!(zs, 0, ps[which][i], NaN)
			# push!(zs, 0, 1, NaN)D
			push!(cs, color_list[mod1(which, 3)], color_list[mod1(which, 3)], color_list[mod1(which, 3)])
		end
		push!(cs2, which)
	end
				# plot(xs, ys, zs)

	plot(leg=false)

	plot!(1:15, [0; cumsum(sign.(randn(14)))], zeros(15), alpha=0.6, m=:o, lw=2, c=color_list)
	plot!(xs, ys, zs, c=cs, xlims=(1, 15), ylims=(-30, 30), zlims=(0, 1), lw=2.5, alpha=0.7, xticks=(1:15), yticks=(-20:10:20))

xlabel!("t")
	ylabel!("space")

end

# ╔═╡ Cell order:
# ╟─e46441c4-97bd-11eb-330c-97bd5ac41f9e
# ╠═85b45a43-d7bf-4597-a1a6-329b41dce20d
# ╠═85c26eb4-c258-4a8b-9415-7b5f7ddff02a
# ╟─2d48bfc6-9617-11eb-3a85-279bebd21332
# ╟─30162386-039f-4cd7-9121-a3382be3c294
# ╟─4e7b163e-dfd0-457e-b1f3-8807a4d8060a
# ╠═e8ceab7b-45db-4393-bb8e-e000ecf78d2c
# ╠═2d4dffb9-39e4-48de-9688-980b96814c9f
# ╟─8ff66523-bc2e-4c53-975b-8ba4f99eb1c6
# ╟─2868dd57-7164-4162-8c5d-30628dedeb7a
# ╠═f6009473-d3c1-444f-88ae-814f770e811b
# ╟─9a368602-acd3-43fb-9dff-e407a4bab930
# ╟─67517333-175f-48c4-a915-76658cbf1304
# ╠═d6832372-d336-4a54-bbcf-d0bb70e4de64
# ╟─35f14826-f1e4-4977-a31a-0f6148fe25ad
# ╠═7468fc5d-7f35-45e2-b5fc-7e63b562bc8f
# ╠═1ca8aa3b-b05d-40f6-a925-2f0248b79ca2
# ╟─38d88b7c-3b4f-430b-8d3c-f672ab0c7a49
# ╟─f4c9b02b-738b-4de1-9e9d-05b1616bee0b
# ╟─d1c47afa-ab7f-4543-a161-e3ceb6f11eb4
# ╠═bf78e00f-05d9-4a05-8512-4924ef9e25f7
# ╟─b948830f-ead1-4f36-a237-c998f2f7deb8
# ╠═15223c51-8d31-4a50-a8ff-1cb7d35de454
# ╟─0bd45c4a-3286-427a-a927-15869be2ebfe
# ╠═999fb608-fb1a-46cb-82ca-f3f31fe617e1
# ╟─6509e69a-6e50-4816-a98f-67ba437383fb
# ╟─e58976a9-1784-441e-bb76-3011538b8ad0
# ╟─1efc2b68-9313-424f-9850-eb4496cc8486
# ╟─6e93ffda-217b-4d46-86b5-534ddc1bae90
# ╟─396d2490-3cb9-4f68-8fdf-9209d2010e02
# ╟─dc1c22e8-1c7b-43b7-8421-c2ca708931a5
# ╟─e7a52b56-322c-4478-a670-dec1013c9bd8
# ╟─5f136388-5573-4814-a088-a66278acdbbe
# ╟─e440dc3b-bafd-4e0c-9fe8-13fce8eea22d
# ╟─89d3e90d-3685-473d-aea4-0b7c5b80d7f7
# ╟─203eea14-1c68-4a9f-ab18-9a2e5f408a79
# ╠═b6d4b045-a39f-4236-ace2-9f631e853d1b
# ╟─e61b56be-d334-4c8f-aa8e-887bb27c058c
# ╠═782f0b9a-3793-4abb-826b-9e14d6eae690
# ╟─1764b56a-f297-4a4e-a931-31aa987ec785
# ╠═8082092b-b6bf-4619-8776-39fdd6c9b7c1
# ╠═34c0b850-5e95-4eb9-8435-3aae8d124772
# ╟─4115f7cb-d45f-4cc2-86bf-316c074393f8
# ╟─6bc2f20d-3d09-425b-a471-44090dc3161e
# ╟─98994cb9-45dc-48aa-b62d-2407f7184bee
# ╟─f11f8d7d-cd3f-4585-aab4-083b892c6d4c
# ╠═fb804fe2-58be-46c9-9200-ceb8863d052c
# ╟─40e457b4-616c-4fab-9c8e-2e5063129597
# ╟─979c1fbd-c9f6-4e8b-a648-6a0210fc9e7f
# ╟─583e3a92-01e7-4b88-9be0-f1e3b3c95005
# ╠═0b26efab-4e93-4d53-9c4d-faea68d12174
# ╟─b9ce5af1-84f7-4a2d-92c9-de2c5498a88d
# ╠═b48e55b7-4b56-41aa-9796-674d04adf5df
# ╟─53a36c1a-0b8c-4099-8854-08d73c9f118e
# ╠═6b298184-32c6-412d-a900-b113d6bd3d53
# ╠═b84a7255-7b0a-4ba1-8c87-9f5d3fa32ef3
# ╟─c430c4de-d9bf-44e1-aa40-6b823d718b04
# ╠═99717c6e-f713-49d5-8ee5-a08c4a464223
# ╟─242ea831-c421-4a76-b658-2a57fa924a4f
# ╠═aeaef573-1e90-45f3-a7fe-31ec5e2808c4
# ╟─efe186da-3273-4767-aafc-fc8eae01dbd9
# ╟─61201091-b8b3-4776-9be9-4c23d5ba88ba
# ╠═66c4aed3-a04b-4a09-b954-79e816d2a3f7
# ╟─b135f6be-5e82-4c72-af11-0eb0d4141dec
# ╠═e74e18e3-ad08-4a53-a803-cd53564dca65
# ╟─ed02f00f-1bcd-43fa-a56c-7be9968614cc
# ╠═8d453f89-4a4a-42d0-8a00-9b153a3f435e
# ╠═f7de29b5-2a51-45e4-a0a5-f7f602681303
# ╟─7e817bad-dc51-4c29-a4fc-f7a8bb3663ca
# ╠═403d607b-6171-431b-a058-0aad0909846f
# ╠═80e033c4-a862-443d-a198-932f5822a44e
# ╟─c8c16c14-26b0-4f83-8135-4f862ed90686
# ╠═29612df6-203f-42bc-b53b-86af618d60ec


using Test
using ModelingToolkit
using ModelingToolkit: Equation
using ModelingToolkit.StructuralTransformations: SystemStructure, find_solvables!
using NonlinearSolve
using LinearAlgebra
using UnPack

###
### Nonlinear system
###
@parameters t
@variables u1(t) u2(t) u3(t) u4(t) u5(t)
eqs = [
    0 ~ u1 - sin(u5),
    0 ~ u2 - cos(u1),
    0 ~ u3 - hypot(u1, u2),
    0 ~ u4 - hypot(u2, u3),
    0 ~ u5 - hypot(u4, u1),
]
@named sys = NonlinearSystem(eqs, [u1, u2, u3, u4, u5], [])
state = TearingState(sys)
StructuralTransformations.find_solvables!(state)

io = IOBuffer()
show(io, MIME"text/plain"(), state.structure)
prt = String(take!(io))

if VERSION >= v"1.6"
    @test occursin("Incidence matrix:", prt)
    @test occursin("×", prt)
    @test occursin("⋅", prt)
end

# u1 = f1(u5)
# u2 = f2(u1)
# u3 = f3(u1, u2)
# u4 = f4(u2, u3)
# u5 = f5(u4, u1)
state = TearingState(sys)
find_solvables!(state)
@unpack structure, fullvars = state
@unpack graph, solvable_graph = state.structure
int2var = Dict(eachindex(fullvars) .=> fullvars)
graph2vars(graph) = map(is -> Set(map(i -> int2var[i], is)), graph.fadjlist)
@test graph2vars(graph) == [Set([u1, u5])
       Set([u1, u2])
       Set([u1, u3, u2])
       Set([u4, u3, u2])
       Set([u4, u1, u5])]
@test graph2vars(solvable_graph) == [Set([u1])
                                     Set([u2])
                                     Set([u3])
                                     Set([u4])
                                     Set([u5])]

state = TearingState(tearing(sys))
let sss = state.structure
    @unpack graph = sss
    @test graph2vars(graph) == [Set([u1, u2, u5])]
end

# Before:
#      u1  u2  u3  u4  u5
# e1 [  1               1 ]
# e2 [  1   1             ]
# e3 [  1   1   1         ]
# e4 [      1   1   1     ]
# e5 [  1           1   1 ]
# solvable_graphs:
#      u1  u2  u3  u4  u5
# e1 [  1                 ]
# e2 [      1             ]
# e3 [          1         ]
# e4 [              1     ]
# e5 [                  1 ]
#
# Optimal:
#      u2  u3  u4  u5  | u1
# e2 [  1              |  1 ]
# e3 [  1   1          |  1 ]
# e4 [  1   1   1      |    ]
# e5 [          1   1  |  1 ]
# ---------------------|-----
# e1 [          1   1  |    ]
#
# Or:
#      u1  u2  u3  u4 | u5
# e1 [  1             |  1 ]
# e2 [  1   1         |    ]
# e3 [  1   1   1     |    ]
# e4 [      1   1   1 |    ]
# --------------------|-----
# e5 [  1           1 |  1 ]

let state = TearingState(sys)
    torn_matching = tearing(state)
    S = StructuralTransformations.reordered_matrix(sys, torn_matching)
    @test S == [1 0 0 0 1
                1 1 0 0 0
                1 1 1 0 0
                0 1 1 1 0
                1 0 0 1 1]
end

# unknowns: u5
# u1 := sin(u5)
# u2 := cos(u1)
# u3 := hypot(u1, u2)
# u4 := hypot(u2, u3)
# solve for
# 0 = u5 - hypot(u1, u4)

# unknowns: u5
# solve for
# 0 = u5 - hypot(sin(u5), hypot(cos(sin(u5)), hypot(sin(u5), cos(sin(u5)))))
tornsys = tearing(sys)
@test isequal(equations(tornsys), [0 ~ u5 - hypot(u4, u1)])
prob = NonlinearProblem(tornsys, ones(1))
sol = solve(prob, NewtonRaphson())
@test norm(prob.f(sol.u, sol.prob.p)) < 1e-10

###
### Simple test (edge case)
###
@parameters t
@variables x(t) y(t) z(t)
eqs = [
    0 ~ x - y,
    0 ~ z + y,
    0 ~ x + z,
]
@named nlsys = NonlinearSystem(eqs, [x, y, z], [])
let (mm, _, _) = ModelingToolkit.aag_bareiss(nlsys)
    @test mm == [-1 1 0;
                 0 -1 -1;
                 0 0 0]
end

newsys = tearing(nlsys)
@test length(equations(newsys)) <= 1

###
### DAE system
###
using ModelingToolkit, OrdinaryDiffEq, BenchmarkTools
@parameters t p
@variables x(t) y(t) z(t)
D = Differential(t)
eqs = [D(x) ~ z
       0 ~ x - y
       0 ~ sin(z) + y - p * t]
@named daesys = ODESystem(eqs, t)
newdaesys = tearing(daesys)
@test equations(newdaesys) == [D(x) ~ z; 0 ~ y + sin(z) - p * t]
@test equations(tearing_substitution(newdaesys)) == [D(x) ~ z; 0 ~ x + sin(z) - p * t]
@test isequal(states(newdaesys), [x, z])
prob = ODAEProblem(newdaesys, [x => 1.0], (0, 1.0), [p => 0.2])
du = [0.0];
u = [1.0];
pr = 0.2;
tt = 0.1;
@test (@ballocated $(prob.f)($du, $u, $pr, $tt)) == 0
@test du≈[-asin(u[1] - pr * tt)] atol=1e-5

# test the initial guess is respected
@named sys = ODESystem(eqs, t, defaults = Dict(z => Inf))
infprob = ODAEProblem(tearing(sys), [x => 1.0], (0, 1.0), [p => 0.2])
@test_throws DomainError infprob.f(du, u, pr, tt)

sol1 = solve(prob, Tsit5())
sol2 = solve(ODEProblem{false}((u, p, t) -> [-asin(u[1] - pr * t)],
                               [1.0],
                               (0, 1.0),
                               0.2), Tsit5(), tstops = sol1.t, adaptive = false)
@test Array(sol1)≈Array(sol2) atol=1e-5

@test sol1[x] == first.(sol1.u)
@test sol1[y] == first.(sol1.u)
@test sin.(sol1[z]) .+ sol1[y]≈pr[1] * sol1.t atol=1e-5
@test sol1[sin(z) + y]≈sin.(sol1[z]) .+ sol1[y] rtol=1e-12

@test sol1[y, :] == sol1[x, :]
@test (@. sin(sol1[z, :]) + sol1[y, :])≈pr * sol1.t atol=1e-5

# 1426
function Translational_Mass(; name, m = 1.0)
    sts = @variables s(t) v(t) a(t)
    ps = @parameters m = m
    D = Differential(t)
    eqs = [D(s) ~ v
           D(v) ~ a
           m * a ~ 0.0]
    ODESystem(eqs, t, sts, ps; name = name)
end

m = 1.0
@named mass = Translational_Mass(m = m)

ms_eqs = []

@named _ms_model = ODESystem(ms_eqs, t)
@named ms_model = compose(_ms_model,
                          [mass])

calculate_jacobian(ms_model)
calculate_tgrad(ms_model)

# Mass starts with velocity = 1
u0 = [mass.s => 0.0
      mass.v => 1.0]

sys = structural_simplify(ms_model)
@test sys.jac[] === ModelingToolkit.EMPTY_JAC
@test sys.tgrad[] === ModelingToolkit.EMPTY_TGRAD
prob_complex = ODAEProblem(sys, u0, (0, 1.0))
sol = solve(prob_complex, Tsit5())
@test all(sol[mass.v] .== 1)


################################################################################
# this has to be an update, since the updating of the x variable does not occur in valdir
struct PolyakMomentum{T, D, S<:Real} <: Update{T}
    direction::D
    x::T # old state variable
    α::S # stepsize
    β::S # momentum variable
end
const Polyak = PolyakMomentum
objective(P::PolyakMomentum) = objective(P.d)
function valdir(P::PolyakMomentum, x)
    v, d = valdir(P.direction, x)
    @. d += P.β * (x - P.x)
    return v, d
end
function update!(P::PolyakMomentum, x)
    d = P.direction(x) # update direction
    d += P.β * (x - P.x)
    copy!(P.x, x) # update old value
    @. x += P.α * d
end

################################################################################
# Nesterov Accelerated Gradients, Nesterov Momentum
# e.g. https://jlmelville.github.io/mize/nesterov.html
# In conjunction with stochastic methods, variance reduction,
# look at https://arxiv.org/pdf/1910.07939.pdf
# and https://arxiv.org/pdf/1910.09494.pdf
struct NesterovMomentum{T, G, S<:Real} <: Update{T}
    direction::G # should be negative gradient
    x::T # keeping old θ vector around
    t::T # need temporary
    α::S # A
    μ::S # M # technically this should also be a step size
end
const Nesterov = NesterovMomentum
function Nesterov(D, x::AbstractArray, α::Real, μ::Real)
    Nesterov(D, copy(x), copy(x), α, μ)
end
function update!(N::Nesterov, x)
    d = N.direction(x)
    @. x = x + N.α * d # classical gradient step
    copy!(N.t, x)
    @. x = x + N.μ * (x - N.x)
    copy!(N.x, N.t) # update old θ vector
    return x
end

################################################################################
# ADAM implementation will work by setting step size selector to constant
# technically it's a stochastic descent direction, make f stochastic
# based on "ADAM: A METHOD FOR STOCHASTIC OPTIMIZATION"
# make this a descent direction
struct ADAM{T, F, V} <: Update{T} # think of it as direction
    d::F # direction
    α::T # basic step size
    β₁::T # "memory" parameter for m
    β₂::T # "memory" parameter for v
    m::V # first moment estimate
    v::V # second moment estimate
    ε::T
end
function ADAM(d, x::AbstractArray, α::Real, β₁::Real, β₂::Real, ε::Real = 1e-8)
    m, v = zero(x), zero(x)
    ADAM(d, α, β₁, β₂, m, v, ε)
end
# updates x in place
function update!(A::ADAM, x::AbstractArray, t::Int)
    dx = A.d(x)
    @. A.m = expave(A.m, dx, A.β₁) # update biased moment estimates
    @. A.v = expave(A.v, dx^2, A.β₂)
    α_t = stepsize(A, t)
    @. x += α_t * A.m / (√A.v + A.ε)
    x
end

# exponential average
function expave(a::Real, b::Real, γ::Real)
    γ * a + (1-γ) * b
end

# AMSGRAD https://openreview.net/pdf?id=ryQu7f-RZ
struct AMSGRAD{T, F, V} <: Update{T} # think of it as direction
    d::F # direction
    α::T # basic step size
    β₁::T # "memory" parameter for m
    β₂::T # "memory" parameter for v
    m::V # first moment estimate
    v::V # second moment estimate
    maxv::V # maximum second moment in optimization history
    ε::T
end
function AMSGRAD(d, x::AbstractArray, α::Real, β₁::Real, β₂::Real, ε::Real = 1e-8)
    m, v, maxv = zero(x), zero(x), zero(x)
    AMSGRAD(d, α, β₁, β₂, m, v, maxv, ε)
end
# updates x in place
function update!(A::AMSGRAD, x::AbstractArray, t::Int)
    dx = A.d(x)
    @. A.m = expave(A.m, dx, A.β₁) # update biased moment estimates
    @. A.v = expave(A.v, dx^2, A.β₂)
    @. A.maxv = max(A.maxv, A.v)
    α_t = stepsize(A, t)
    @. x += α_t * A.m / (√A.maxv + A.ε)
end

# stepsize to debias moments
function stepsize(A::Union{ADAM, AMSGRAD}, t::Int)
    t > 0 || throw(DomainError("time step has to be bigger than 0 but is equal to $t"))
    A.α * sqrt(1-A.β₂^t)/(1-A.β₁^t)
end


# bradi-tests.jl -
#
# Tests for BRADI algorithm.
#

module BraDiTests

using Test
import OptimPackNextGen: BraDi

# Simple parabola.  To be minimized over [-1,2].
parabola(x) = x*x

# Brent's 5th function.  To be minimized over [-10,10].
brent5(x) = (x - sin(x))*exp(-x*x)

# Michalewicz's 1st function.  To be minimized over [-1,2].
michalewicz1(x) = x*sin(10.0*x)

# Michalewicz's 2nd function.  To be maximized over [0,pi].
function michalewicz2(x)
    s = 0.0
    a = sin(x)
    b = x*x/pi
    for i in 1:10
        s += a*(sin(b*i)^20)
    end
    return s
end

function runtests(; quiet::Bool=false)
    tol = sqrt(eps(Float64))
    @testset "BraDi algorithm" begin
        quiet || println("\n# Simple parabola:")
        (xbest, fbest) = BraDi.minimize(parabola, range(-1, stop=2, length=2))
        println("x = $xbest, f(x) = $fbest")
        @test xbest ≈ 0.0 atol=tol

        quiet || println("\n# Brent's 5th function:")
        (xbest, fbest) = BraDi.minimize(brent5, range(-10, stop=10, length=5))
        quiet || println("x = $xbest, f(x) = $fbest")
        @test xbest ≈ -1.1951366418407416 rtol=tol

        quiet || println("\n# Michalewicz's 1st function:")
        (xbest, fbest) = BraDi.minimize(michalewicz1, range(-1, stop=2, length=21))
        quiet || println("x = $xbest, f(x) = $fbest")
        @test xbest ≈ 1.7336377815999973 rtol=tol

        quiet || println("\n# Michalewicz's 2nd function:")
        (xbest, fbest) = BraDi.maximize(michalewicz2, range(0, stop=pi, length=60))
        quiet || println("x = $xbest, f(x) = $fbest")
        @test xbest ≈ 2.2208651539586493 rtol=tol
    end
end
runtests()

end # module


export select_region_ft, select_region_rft
export rft_fix_after, rft_fix_before, ft_fix_after, ft_fix_before

"""
    select_region_ft(mat,new_size)

Performs the necessary Fourier-space operations of resampling
in the space of ft (meaning the already circshifted version of fft).

Note that in dependence of the size, fixes in Fourier-space are applied under the assumption
that the corresponding data in Real space is real-valued only. You should used this function
if you want to select a region (extract) in Fourier space and know that you are dealing with
a real-valued signal. For complex-valued data, the ordinary `select_region` function should be
used instead.
`new_size`.

The size of the array view after the operation finished. The Fourier-center
is always assumed to align before and after the padding operation.

 # Examples
```jldoctest
julia> x = [1 20 3; 4 500 6; -7 821 923]
3×3 Matrix{Int64}:
  1   20    3
  4  500    6
 -7  821  923

julia> ffts(x)
3×3 ShiftedArrays.CircShiftedArray{ComplexF64, 2, Matrix{ComplexF64}}:
   106.5+390.577im  -1099.5-1062.61im   1000.5+700.615im
 -1138.5+354.204im   2271.0+0.0im      -1138.5-354.204im
  1000.5-700.615im  -1099.5+1062.61im    106.5-390.577im

julia> select_region_ft(ffts(x), (4,4))
4×4 PaddedView(0.0 + 0.0im, OffsetArray(::ShiftedArrays.CircShiftedArray{ComplexF64, 2, Matrix{ComplexF64}}, 2:4, 2:4), (Base.OneTo(4), Base.OneTo(4))) with eltype ComplexF64:
 0.0+0.0im      0.0+0.0im          0.0+0.0im          0.0+0.0im
 0.0+0.0im    106.5+390.577im  -1099.5-1062.61im   1000.5+700.615im
 0.0+0.0im  -1138.5+354.204im   2271.0+0.0im      -1138.5-354.204im
 0.0+0.0im   1000.5-700.615im  -1099.5+1062.61im    106.5-390.577im

julia> x = [1 20; 4 500; -7 821; -2 2]
4×2 Matrix{Int64}:
  1   20
  4  500
 -7  821
 -2    2

julia> ffts(x)
4×2 ShiftedArrays.CircShiftedArray{ComplexF64, 2, Matrix{ComplexF64}}:
  -347.0+0.0im     331.0+0.0im
   809.0-492.0im  -793.0+504.0im
 -1347.0+0.0im    1339.0+0.0im
   809.0+492.0im  -793.0-504.0im

julia> select_region_ft(ffts(x), (5,3))
5×3 FourierTools.FourierSplit{ComplexF64, 2, FourierTools.FourierSplit{ComplexF64, 2, PaddedViews.PaddedView{ComplexF64, 2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}, OffsetArrays.OffsetMatrix{ComplexF64, ShiftedArrays.CircShiftedArray{ComplexF64, 2, Matrix{ComplexF64}}}}}}:
 -86.75+0.0im     165.5+0.0im    -86.75+0.0im
  404.5-246.0im  -793.0+504.0im   404.5-246.0im
 -673.5+0.0im    1339.0+0.0im    -673.5+0.0im
  404.5+246.0im  -793.0-504.0im   404.5+246.0im
 -86.75+0.0im     165.5+0.0im    -86.75+0.0im
```
"""
function select_region_ft(mat,new_size)
    old_size = size(mat)
    mat_fixed_before = ft_fix_before(mat,old_size,new_size)
    mat_pad = ft_pad(mat_fixed_before,new_size)
    # afterwards we add the highest pos. frequency to the highest lowest one
    return ft_fix_after(mat_pad ,old_size,new_size)
end

"""
    select_region_rft(mat,old_size, new_size)

performs the necessary Fourier-space operations of resampling
in the space of rft (meaning the already circshifted version of rfft).
Note that in dependence of the size, fixes in Fourier-space are applied, which is
why you should used this function if you want to select a region (extract) in rft space.
Since rfts always assume the corresponding Real-space data to be real, this rule should
have no exception for rft data.

`old_size`.
The size of the corresponding real-space array before it was rfted.

`new_size`.
The size of the corresponding real-space array view after the operation finished. The Fourier-center
is always assumed to align before and after the padding aperation.
"""
function select_region_rft(mat,old_size, new_size)
    rft_old_size = size(mat)
    rft_new_size = Base.setindex(new_size,new_size[1]÷2 +1, 1)
    return rft_fix_after(rft_pad(
        rft_fix_before(mat,old_size,new_size),
        rft_new_size),old_size,new_size)
end

"""
    select_region(mat,new_size)

performs the necessary Fourier-space operations of resampling
in the space of ft (meaning the already circshifted version of fft).

`new_size`.
The size of the array view after the operation finished.

`center`.
Specifies the center of the new view in coordinates of the old view. By default an alignment of the Fourier-centers is assumed.
# Examples
```jldoctest
julia> using FFTW, FourierTools

julia> select_region(ones(3,3),new_size=(7,7),center=(1,3))
7×7 PaddedView(0.0, OffsetArray(::Matrix{Float64}, 4:6, 2:4), (Base.OneTo(7), Base.OneTo(7))) with eltype Float64:
 0.0  0.0  0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0  0.0  0.0
 0.0  1.0  1.0  1.0  0.0  0.0  0.0
 0.0  1.0  1.0  1.0  0.0  0.0  0.0
 0.0  1.0  1.0  1.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0  0.0  0.0
```
"""
function select_region(mat; new_size=size(mat), center=ft_center_diff(size(mat)).+1, pad_value=zero(eltype(mat)))
    new_size = Tuple(expand_size(new_size, size(mat)))
    center = Tuple(expand_size(center, ft_center_diff(size(mat)).+1))
    oldcenter = ft_center_diff(new_size).+1
    PaddedView(pad_value, mat,new_size, oldcenter .- center.+1);
end

function ft_pad(mat, new_size)
    return select_region(mat;new_size=new_size)
end

function rft_pad(mat, new_size)
    c2 = rft_center_diff(size(mat))
    c2 = Base.setindex(c2, new_size[1].÷2, 1);
    return select_region(mat;new_size=new_size, center=c2.+1)
end

function ft_fix_before(mat, size_old, size_new; start_dim=1)
    for d = start_dim:ndims(mat)
        sn = size_new[d]
        so = size_old[d]
        if sn < so && iseven(sn)
            L1 = (size_old[d] -size_new[d] )÷2 +1
            mat = FourierJoin(mat, d, L1)
        end
    end
    return mat
end

function ft_fix_after(mat,size_old,size_new; start_dim=1)
    start_dim
    ndims(mat)
    for d=start_dim:ndims(mat)
        sn = size_new[d]
        so = size_old[d]
        if sn > so && iseven(so)
            L1 = (size_new[d]-size_old[d])÷2+1
            mat = FourierSplit(mat,d,L1)
        end
        # if equal do nothing
    end
    return mat
end

function rft_fix_first_dim_before(mat,size_old,size_new;dim=1)
    sn = size_new[dim] # Note that this dim is the corresponding real-space size
    so = size_old[dim] # Note that this dim is the corresponding real-space size
    if sn < so && iseven(sn) # result size is even upon cropping
        L1 = size_new[dim] ÷ 2 + 1
        mat = FourierJoin(mat, dim, L1, L1) # a hack to dublicate the value
    end
    return mat
end

function rft_fix_first_dim_after(mat,size_old,size_new;dim=1)
    sn = size_new[dim] # Note that this dim is the corresponding real-space size
    so = size_old[dim] # Note that this dim is the corresponding real-space size
    if sn > so && iseven(so) # source size is even upon padding
        L1 = size_old[dim] ÷ 2 + 1
        mat = FourierSplit(mat,dim,L1,-1) # This hack prevents a second position to be affected
    end
    # if equal do nothing
    return mat
end

function rft_fix_before(mat,size_old,size_new)
    mat=rft_fix_first_dim_before(mat,size_old,size_new;dim=1) # ignore the first dimension
    ft_fix_before(mat,size_old,size_new;start_dim=2) # ignore the first dimension
end

function rft_fix_after(mat,size_old,size_new)
    mat = rft_fix_first_dim_after(mat,size_old,size_new;dim=1) # ignore the first dimension
    ft_fix_after(mat,size_old,size_new;start_dim=2) # ignore the first dimension
end


immutable Poisson <: DiscreteUnivariateDistribution
    lambda::Float64
    function Poisson(l::Real)
    	l > zero(l) || error("lambda must be positive")
        new(float64(l))
    end
    Poisson() = new(1.0)
end

@_jl_dist_1p Poisson pois

function entropy(d::Poisson)
    λ = d.lambda
    if λ < 50.0
        s = 0.0
        for k in 1:100
            s += λ^k * lgamma(k + 1.0) / gamma(k + 1.0)
        end
        return λ * (1.0 - log(λ)) + exp(-λ) * s
    else
        return 0.5 * log(2 * pi * e * λ) -
               (1 / (12 * λ)) -
               (1 / (24 * λ * λ)) -
               (19 / (360 * λ * λ * λ))
    end
end

insupport(::Poisson, x::Real) = isinteger(x) && zero(x) <= x
insupport(::Type{Poisson}, x::Real) = isinteger(x) && zero(x) <= x

kurtosis(d::Poisson) = 1.0 / d.lambda

mean(d::Poisson) = d.lambda

median(d::Poisson) = quantile(d, 0.5)

function mgf(d::Poisson, t::Real)
    l = d.lambda
    return exp(l * (exp(t) - 1.0))
end

function cf(d::Poisson, t::Real)
    l = d.lambda
    return exp(l * (exp(im * t) - 1.0))
end

mode(d::Poisson) = ifloor(d.lambda)
modes(d::Poisson) = [mode(d)]

skewness(d::Poisson) = 1.0 / sqrt(d.lambda)

var(d::Poisson) = d.lambda

function fit_mle(::Type{Poisson}, x::Array)
    for i in 1:length(x)
        if !insupport(Poisson(), x[i])
            error("Poisson observations must be non-negative integers")
        end
    end
    Poisson(mean(x))
end


"""
Big class of general phase-retrieval-related problems.
"""
abstract type AbstractPRproblem <: Problem end

"""
  PRproblem(a, A)

Classical phase-retrieval problem of finding complex arrays x, X such that
  |x| = a, |X| = A, and X =F(x),
where F denotes the Fourier transform.
"""
struct PRproblem{T <: Real, N} <: AbstractPRproblem
  a::Array{T,N}
  A::Array{T,N}
end

# this problem can be converted to a feasibility problem
TwoSetsFP(pr::PRproblem) = TwoSetsFP(ConstrainedByAmplitude(pr.a), FourierTransformedSet(ConstrainedByAmplitude(pr.A)))

# methods to solve
"""
Gerchberg-Saxton, classical method to solve PR problem
"""
abstract type GS <: IterativeAlgorithm end

"""
    GSparam(; ϕ⁰, maxit, maxϵ, keephistory, snapshots)

Parameters for a concrete realisation of GS method.
"""
Base.@kwdef struct GSparam <: GS
    ϕ⁰ = missing
    maxϵ::Union{Float64, Missing} = missing
    maxit::Union{Missing,Int64} = missing
    keephistory::Bool = false
    snapshots::Array{Int64} = Int64[]
end

initial(alg::GSparam) = alg.ϕ⁰
tolerance(alg::GSparam) = alg.maxϵ
maxit(alg::GSparam) = alg.maxit
keephistory(alg::GSparam) = alg.keephistory
snapshots(alg::GSparam) = alg.snapshots

# GS is just alternating projection method
"""
    solve(pr::PRproblem, alg::GS, args...)

Solve PR problem using GS method. Any change of the GS parameters can be given
"""
solve(pr::PRproblem, ::GSparam, ϕ⁰, maxϵ, maxit, keephistory, snapshots) =
        solve(
            TwoSetsFP(pr),
            AlternatingProjections.APparam(),
            ismissing(ϕ⁰) ?  missing : pr.a .* exp.(2π .* im .* ϕ⁰),
            maxϵ,
            maxit,
            keephistory,
            snapshots)

# PRproblem(x, X) = TwoSetsFP(ConstrainedByAmplitude(x), FourierTransformedSet(ConstrainedByAmplitude(X)))


export postprocess
function postprocess(Events::Array{Event,1})
    set_context("poster")
    set_context("poster")
    close("all")
    pathlengths = zeros(0)
    distance = 0.0

    distancetofirst = true
    for i=1:length(Events)
        distance += Events[i].distance
        if Events[i].collision
            if distancetofirst
                distancetofirst = false # only store distances between collision
            else
                push!(pathlengths, distance)
            end
            distance = 0.0
        end
    end
    return postprocess(pathlengths)
end

function postprocess(pathlengths::Array{Float64,1},latticeflag = false)
    set_context("poster")
    set_context("poster")
    close("all")

    sort!(pathlengths)

    npl = length(pathlengths)

    fig, ax = plt.subplots(1,1,figsize=(40,20))

    logpathlengths = log10.(pathlengths)

    h = fit(Histogram, pathlengths, nbins = 200)

    centers = (h.edges[1][1:end-1] .+ h.edges[1][2:end])/2
    weights = h.weights
    weights = weights/sum(weights)/(centers[2]-centers[1])
    ax.plot(centers,weights,lw = 5,"-o",color=tuple([0,150,130]/255...))
    ax.set_xlim([0,1.05*maximum(pathlengths)])
    plt.yscale("log")

    d = fit(Exponential, pathlengths)
	if !latticeflag
		ax.plot(centers,pdf.(d,centers),lw=5,"--",color=tuple([163,16,124]/255...))
	end

    m(t, p) = p[1] .* t.^(-2.0)#p[2]
    p0 = [3.0]#, -2.0]
    lsqfit = curve_fit(m, centers, weights, p0)

    m(t) = lsqfit.param[1].*t.(-2.0)#^lsqfit.param[2]
    if latticeflag
		ax.plot(centers,m.(centers),lw=2,"--")
	end
	ax.set_title("Path-length distribution for n=$npl samples.", pad = 20)
    ax.set_xlabel("Distance s")
    ax.set_ylabel("Rel. frequency")
   # ax.set_ylim([1e-16,1])
    theta = params(d)[1]
    lambda = round(1.0/theta,digits = 3)
    sigma = round(lsqfit.param[1],digits = 3)
    #alpha = round(lsqfit.param[2],digits = 3)
	#if !latticeflag
#		ax.legend(["Data","Best fit, \$ \\lambda \\cdot \\exp( -\\lambda s)\$, \$\\lambda=\$ $lambda"])#
#	else
#		ax.legend(["Best fit, \$ \\sigma \\cdot s^\\alpha \$, \$\\sigma=\$ $sigma, \$\\alpha=\$ $alpha"])
#	end
    plt.tight_layout()
    plt.savefig("Histogram.png",dpi=200)

    return pathlengths
end


"""
    FVector(type::VectorSpaceType, data, basis::AbstractBasis)

Decorator indicating that the vector `data` contains coordinates of a vector from a fiber
of a vector bundle of type `type`. `basis` is an object describing the basis of that space
in which the coordinates are given.

Conversion between `FVector` representation and the default representation of an object
(for example a tangent vector) for a manifold should be done using [`get_coordinates`](@ref)
and [`get_vector`](@ref).

# Examples

```julia-repl
julia> using Manifolds

julia> M = Sphere(2)
Sphere(2, ℝ)

julia> p = [1.0, 0.0, 0.0]
3-element Vector{Float64}:
 1.0
 0.0
 0.0

julia> X = [0.0, 2.0, -1.0]
3-element Vector{Float64}:
  0.0
  2.0
 -1.0

julia> B = DefaultOrthonormalBasis()
DefaultOrthonormalBasis(ℝ)

julia> fX = TFVector(get_coordinates(M, p, X, B), B)
TFVector([2.0, -1.0], DefaultOrthonormalBasis(ℝ))

julia> X_back = get_vector(M, p, fX.data, fX.basis)
3-element Vector{Float64}:
 -0.0
  2.0
 -1.0
```
"""
struct FVector{TType<:VectorSpaceType,TData,TBasis<:AbstractBasis}
    type::TType
    data::TData
    basis::TBasis
end

const TFVector = FVector{TangentSpaceType}
const CoTFVector = FVector{CotangentSpaceType}

function TFVector(data, basis::AbstractBasis)
    return TFVector{typeof(data),typeof(basis)}(TangentSpace, data, basis)
end
function CoTFVector(data, basis::AbstractBasis)
    return CoTFVector{typeof(data),typeof(basis)}(CotangentSpace, data, basis)
end

function Base.show(io::IO, fX::TFVector)
    return print(io, "TFVector(", fX.data, ", ", fX.basis, ")")
end
function Base.show(io::IO, fX::CoTFVector)
    return print(io, "CoTFVector(", fX.data, ", ", fX.basis, ")")
end

"""
    AbstractFibreVector{TType<:VectorSpaceType}

Type for a vector from a vector space (fibre of a vector bundle) of type `TType` of a manifold.
While a [`AbstractManifold`](@ref) does not necessarily require this type, for example when it is
implemented for `Vector`s or `Matrix` type elements, this type can be used for more
complicated representations, semantic verification, or even dispatch for different
representations of tangent vectors and their types on a manifold.
"""
abstract type AbstractFibreVector{TType<:VectorSpaceType} end

"""
    TVector = AbstractFibreVector{TangentSpaceType}

Type for a tangent vector of a manifold. While a [`AbstractManifold`](@ref) does not necessarily
require this type, for example when it is implemented for `Vector`s or `Matrix` type
elements, this type can be used for more complicated representations, semantic verification,
or even dispatch for different representations of tangent vectors and their types on a
manifold.
"""
const TVector = AbstractFibreVector{TangentSpaceType}

"""
    CoTVector = AbstractFibreVector{CotangentSpaceType}

Type for a cotangent vector of a manifold. While a [`AbstractManifold`](@ref) does not necessarily
require this type, for example when it is implemented for `Vector`s or `Matrix` type
elements, this type can be used for more complicated representations, semantic verification,
or even dispatch for different representations of cotangent vectors and their types on a
manifold.
"""
const CoTVector = AbstractFibreVector{CotangentSpaceType}

Base.:+(X::FVector, Y::FVector) = FVector(X.type, X.data + Y.data, X.basis)

Base.:-(X::FVector, Y::FVector) = FVector(X.type, X.data - Y.data, X.basis)
Base.:-(X::FVector) = FVector(X.type, -X.data, X.basis)

Base.:*(a::Number, X::FVector) = FVector(X.type, a * X.data, X.basis)

allocate(x::FVector) = FVector(x.type, allocate(x.data), x.basis)
allocate(x::FVector, ::Type{T}) where {T} = FVector(x.type, allocate(x.data, T), x.basis)

function Base.copyto!(X::FVector, Y::FVector)
    copyto!(X.data, Y.data)
    return X
end

function number_eltype(
    ::Type{FVector{TType,TData,TBasis}},
) where {TType<:VectorSpaceType,TData,TBasis}
    return number_eltype(TData)
end
number_eltype(v::FVector) = number_eltype(v.data)

Base.show(io::IO, ::TangentSpaceType) = print(io, "TangentSpace")
Base.show(io::IO, ::CotangentSpaceType) = print(io, "CotangentSpace")

"""
    vector_space_dimension(M::AbstractManifold, V::VectorSpaceType)

Dimension of the vector space of type `V` on manifold `M`.
"""
vector_space_dimension(::AbstractManifold, ::VectorSpaceType)

function vector_space_dimension(M::AbstractManifold, ::TCoTSpaceType)
    return manifold_dimension(M)
end


function make_plot(R,r,T,t)
	years = 9
	data_pts = 900
	ti = LinRange(0,years*T,data_pts)
	thetaR = (2*pi/T) .* ti
	thetar = (2*pi/t) .* ti

	XR,YR = R.*cos.(thetaR),R.*sin.(thetaR)
	Xr,Yr = r.*cos.(thetar),r.*sin.(thetar)
	ti = LinRange(0,years*T,data_pts)
	X,Y = R .* cos.(0:0.01:2*pi),R .* sin.(0:0.01:2*pi)
	x,y = r .* cos.(0:0.01:2*pi),r .* sin.(0:0.01:2*pi)
	plot(X,Y,linecolor = :green,lw = 2,labels = false)
	plot!(x,y,linecolor = :blue,lw = 2,labels = false)
	plot!([XR[1],Xr[1]],[YR[1],Yr[1]],labels = false,linecolor = :red)

	for i in 2:length(ti)
		#scatter([XR[i]],[YR[i]],markercolor = :red,markersize = 8)
		#scatter!([XR[i]],[YR[i]],markercolor = :green,markersize = 5)
		plot!([XR[i],Xr[i]],[YR[i],Yr[i]],labels = false,linecolor = :red,aspect_ratio = :equal,grid = false,showaxis = false,ticks = false,xlim = (-R,R),ylim = (-R,R),linealpha = 0.4,dpi =  300)
	end
	savefig("earth_venus.png")
end
make_plot(1,0.72,365.25,224.7)


#
using Plots, Distributions, DataFrames, DelimitedFiles, ColorSchemes, LaTeXStrings, Roots, NLsolve, Random
#
include("toolbox.jl")
#
Random.seed!(12345)
global ee = 1e-8
## generate random sample from the qExponential distribution
xs = [0.1:0.001:1;]
plot(xs,dqexp(xs,0.9,10.))
#
data = rqexp(50,0.9,10.)
#
histogram(data,normalize=true,color="beige",bins=500,label="data")
xs = [minimum(data):0.001:maximum(data);]
plot!(xs,dqexp(xs,0.9,10.),linewidth=2,xlim=[0,1],label=L"\rho_1^{\dagger}(x)",color="red")
#
savefig("qexpdata.pdf") # Figure 1
#
########################################################################
## Calculate \beta-estimators via the Shannon and Renyi MaxEnt procedures
Random.seed!(4545)
global ee = 1e-09

qs = [0.55, 0.65, 0.75];#LinRange(0.55, 0.9, nsims)
betatrue = 10.;
qestimse = zeros(length(qs));
ss, nreps = 10000, 10000; # sample size and Monte Carlo repetitions
bestimse_sh = zeros(nreps,length(qs));
bestimse_r1 = zeros(nreps,length(qs));
bestimse_r2 = zeros(nreps,length(qs));

@time for i in 1:length(qs)
   display("$i/3")
   tt = monteCarlo_betas(ss, nreps, qs[i], betatrue);
   bestimse_sh[:,i] = tt[:,1];
   bestimse_r1[:,i] = tt[:,2];
   bestimse_r2[:,i] = tt[:,3];
end
##############
# writedlm("bestimseNe04_nsims10000_sh.txt",bestimse_sh)
# writedlm("bestimseNe04_nsims10000_r1.txt",bestimse_r1)
# writedlm("bestimseNe04_nsims10000_r2.txt",bestimse_r2)
####
# bestimse_sh = readdlm("bestimseNe03_nsims10000_sh.txt");
# bestimse_r1 = readdlm("bestimseNe03_nsims10000_r1.txt");
# bestimse_r2 = readdlm("bestimseNe03_nsims10000_r2.txt");
#
# bestimse_sh2 = readdlm("bestimseNe04_nsims10000_sh.txt");
# bestimse_r12 = readdlm("bestimseNe04_nsims10000_r1.txt");
# bestimse_r22 = readdlm("bestimseNe04_nsims10000_r2.txt");
####
pyplot()

a = kde_lscv(bestimse_sh[:,1]);
xvals, densityvals = a.x, a.density
p11 = plot(xvals, densityvals,label="Shannon",linewidth=2,title=L"(a)\,\, Q=0.55, \, n=10^3",titlefontsize=10,legendfontsize=7,legend=:topright,ylab="Density",xlim=[0,25],ylim=[0,1.5])
a = kde_lscv(bestimse_r1[:,1]);
xvals, densityvals = a.x, a.density
p11 = plot!(xvals, densityvals,label="Renyi (lin.)",linewidth=2)
a = kde_lscv(bestimse_r2[:,1]);
xvals, densityvals = a.x, a.density
p11 = plot!(xvals, densityvals,label="Renyi (esc.)",linewidth=2)
vline!([betatrue],label="",c="black",line=:dash)
#
a = kde_lscv(bestimse_sh[:,2]);
xvals, densityvals = a.x, a.density
p12 = plot(xvals, densityvals,linewidth=2,title=L"(b)\,\, Q=0.65, \, n=10^3",titlefontsize=10,legendfontsize=7,label="",xlim=[0,13],ylim=[0,2.5])
a = kde_lscv(bestimse_r1[:,2]);
xvals, densityvals = a.x, a.density
p12 = plot!(xvals, densityvals,linewidth=2,label="")
a = kde_lscv(bestimse_r2[:,2]);
xvals, densityvals = a.x, a.density
p12 = plot!(xvals, densityvals,linewidth=2,label="")
vline!([betatrue],label="",c="black",line=:dash)
#
a = kde_lscv(bestimse_sh[:,3]);
xvals, densityvals = a.x, a.density
p13 = plot(xvals, densityvals,linewidth=2,title=L"(c)\,\, Q=0.75,\, n=10^3",titlefontsize=10,legendfontsize=7,label="",xlim=[0,13],ylim=[0,3.5])
a = kde_lscv(bestimse_r1[:,3]);
xvals, densityvals = a.x, a.density
p13 = plot!(xvals, densityvals,linewidth=2,label="")
a = kde_lscv(bestimse_r2[:,3]);
xvals, densityvals = a.x, a.density
p13 = plot!(xvals, densityvals,linewidth=2,label="")
vline!([betatrue],label="",c="black",line=:dash)
#
a = kde_lscv(bestimse_sh2[:,1]);
xvals, densityvals = a.x, a.density
p21 = plot(xvals, densityvals,label="Shannon",linewidth=2,title=L"(d)\,\, Q=0.55, \, n=10^4",titlefontsize=10,legendfontsize=7,legend=:topright,ylab="Density",xlim=[0,25],ylim=[0,1.5])
a = kde_lscv(bestimse_r12[:,1]);
xvals, densityvals = a.x, a.density
p21 = plot!(xvals, densityvals,label="Renyi (lin.)",linewidth=2)
a = kde_lscv(bestimse_r22[:,1]);
xvals, densityvals = a.x, a.density
p21 = plot!(xvals, densityvals,label="Renyi (esc.)",linewidth=2)
vline!([betatrue],label="",c="black",line=:dash)
#
a = kde_lscv(bestimse_sh2[:,2]);
xvals, densityvals = a.x, a.density
p22 = plot(xvals, densityvals,linewidth=2,title=L"(e)\,\, Q=0.65, \, n=10^4",titlefontsize=10,legendfontsize=7,label="",xlab=L"\hat{\Lambda}",xlim=[0,13],ylim=[0,2.5])
a = kde_lscv(bestimse_r12[:,2]);
xvals, densityvals = a.x, a.density
p22 = plot!(xvals, densityvals,linewidth=2,label="")
a = kde_lscv(bestimse_r22[:,2]);
xvals, densityvals = a.x, a.density
p22 = plot!(xvals, densityvals,linewidth=2,label="")
vline!([betatrue],label="",c="black",line=:dash)
#
a = kde_lscv(bestimse_sh2[:,3]);
xvals, densityvals = a.x, a.density
p23 = plot(xvals, densityvals,linewidth=2,title=L"(f)\,\, Q=0.75,\, n=10^4",titlefontsize=10,legendfontsize=7,label="",xlim=[0,13],ylim=[0,3.5])
a = kde_lscv(bestimse_r12[:,3]);
xvals, densityvals = a.x, a.density
p23 = plot!(xvals, densityvals,linewidth=2,label="")
a = kde_lscv(bestimse_r22[:,3]);
xvals, densityvals = a.x, a.density
p23 = plot!(xvals, densityvals,linewidth=2,label="")
vline!([betatrue],label="",c="black",line=:dash)
#
plot(p11,p12,p13,p21,p22,p23,layout=(2,3)) # Figure 3
plot!(size=(1000,500))
# savefig("kdes_qs_ns2.pdf")

########################################################################
## Calculate \beta-estimators via the Shannon and Renyi MaxEnt procedures

nsims = 200 # number of simulations
qs = LinRange(0.51, 0.9, nsims) # grid of q values
betas = LinRange(1., 100., nsims) # grid of \beta values
# initialization
bestimse_sh = zeros(nsims, nsims)
bestimse_r1 = zeros(nsims, nsims)
bestimse_r2 = zeros(nsims, nsims)
qestimse = zeros(nsims, nsims)
strue = zeros(nsims, nsims)
ss, nreps = 1000, 1000 # sample size and Monte Carlo repetitions

#
@time for i in 1:nsims
   display("$i/$nsims")
   for j in 1:nsims
      tt = sqrt.(monteCarloMSE(ss, nreps, qs[i], betas[j])) # root of the MSE
      qestimse[i,j] = tt[1]
      bestimse_sh[i,j] = tt[2]
      bestimse_r1[i,j] = tt[3]
      bestimse_r2[i,j] = tt[4]
   end
end
###
# save results
# writedlm("q_estimations.txt",qestimse)
# writedlm("b_estimations_sh.txt",bestimse_sh)
# writedlm("b_estimations_linear.txt",bestimse_r1)
# writedlm("b_estimations_escort.txt",bestimse_r2)
###

# Calculate relative RMS errors
for i in 1:length(bestimse_sh[:,1])
   bestimse_sh[:,i] = bestimse_sh[:,i] / betas[i]
   bestimse_r1[:,i] = bestimse_r1[:,i] / betas[i]
   bestimse_r2[:,i] = bestimse_r2[:,i] / betas[i]
end

## Plot results
p1 = heatmap(qs[1:end],betas[1:end],transpose(log10.(bestimse_r1[1:end,1:end])),c=cgrad(ColorSchemes.jet1.colors),
   xlabel=L"q",ylabel=L"\beta",title=(L"\log{\hat{MSE}}"))
#
p2 = heatmap(qs[1:end],betas[1:end],transpose(log10.(bestimse_r2[1:end,1:end])),c=cgrad(ColorSchemes.jet1.colors),
   xlabel=L"q",ylabel=L"\beta",line_z=[-1.5,0.5])
#
p3 = heatmap(qs[1:end],betas[1:end],transpose(log10.(bestimse_sh[1:end,1:end])),c=cgrad(ColorSchemes.jet1.colors),
   xlabel=L"q",ylabel=L"\beta",line_z=[-1.5,0.5])
#
plot(p1,p2,p3,layout=(3,1))
##
savefig("colormaps.png") # Figure 5

#########################################################################################################################
#########################################################################################################################


@testset "688.knight-probability-in-chessboard.jl" begin
    @test knight_probability(3, 2, 0, 0) ≈ 0.0625
    @test isapprox(knight_probability(8, 10, 4, 5), 0.0773899; atol = 6)
end


export ARC_step_computation

function ARC_step_computation(h::Float64,
                              g::Float64,
                              Δ::Float64;
                              kwargs...)

    discr1 = (h)^2-4(g/Δ)
    discr2 = (h)^2+4(g/Δ)

    if h>0
      if g>0
        d=(-2*g)/(h+sqrt(discr2))
      else
        d=(-2*g)/(h+sqrt(discr1))
      end
    else
      if g<0
        d=(-h+sqrt(discr1))/(2/Δ)
      else
        d=(-h+sqrt(discr2))/(-2/Δ)
      end
    end

    return d
  end