using Test, LinearAlgebra, BenchmarkTools, Parameters, Revise

N = 64
U0 = randn(N,N,N)

U1 = similar(U0)
U2 = similar(U0)

x = LinRange(-1,1,N)
y = x'

@time z = reshape(x,1,1,length(x))

V(x,y,z,t) = x^2 + y^2 + z^2

@time @. U0 += V(x,y,z,0.)

@time @. U0 = V(x,y,z,0.)

@time @. U0 += V(x,y,z,0.)

@time @. U0 = U0 + U1

function Vadd!(V1,V0,V,t)
    @. V1 = V0 + V(x,y,z,t)
    return nothing
end

@time Vadd!(U1,U0,V,0.)

# function Valloc!(V1,V,t)
#     @. V1 = V(x,y,z,t)
#     return nothing
# end

# @btime Valloc!(V1,V,0.)

# macro

using FourierGPE
import FourierGPE.V

# static potential array
struct VStatic{N} <: UserParams
    V0::Array{Complex{Float64},N}
end

V0(x,y,z) = 0.25*(x^2 + y^2 + 4*z^2)
p0 = VStatic(V0.(x,y,z) .|> complex)

# dynamical potential function
@with_kw mutable struct Potential <: UserParams
    V::Expr = :( 0.5*(x^2 + y^2 + 4*z^2) )
end

macro potential(V,p)
    @eval $p = Potential($V)
end

macro staticpotential(V,p)
    @eval $p = VStatic($V.(x,y,z) .|> complex)
end

#TODO change sim to have V0, V, Vt fields
macro potentials(sim,p,q)
    @eval Sim($sim,V0 = $p,V = $q)
end

pot = Potential()
ex = pot.V

@potential :(0.25*(x^2 + y^2 + 4*z^2)) p1

@staticpotential V0 p2

#TODO V0 should just be a regular array
@test p2 == p0

p0 .≈ p2


### Fetch required packages and reaction networks ###
using Catalyst, DiffEqBase, Random, Test

### Checks that the jacobian is correct for networks without parameters ###
jacobian_network_1 = @reaction_network begin
    (2.0,1.0),   ∅ ↔ X
    (3.0,1.0),   ∅ ↔ Y
    (5.0,2.0),   X + Y ↔ XY
end
test_jac = ODEFunction(convert(ODESystem,jacobian_network_1),jac=true).jac([1.,1.,1.],[],0.)
real_jac = [-6. -5. 2.; -5. -6. 2.; 5. 5. -2.]
@test test_jac == real_jac

### Checks that the jacobian is correct for networks with parameters ###
jacobian_network_2 = @reaction_network begin
    (p1,1.0),   ∅ ↔ X
    (p2,1.0),   ∅ ↔ Y
    (p3*X,1.0), X + Y ↔ XY
end p1 p2 p3
for factor in [1e-2, 1e-1, 1e0, 1e1, 1e2], repeat = 1:10
    u = factor*rand(3)
    p = factor*rand(3)
    test_jac = ODEFunction(convert(ODESystem,jacobian_network_2),jac=true).jac(u,p,0.)
    real_jac = [-1-2*p[3]*u[2]*u[1] -p[3]*u[1]*u[1] 1.; -2*p[3]*u[2]*u[1] -1-p[3]*u[1]*u[1] 1; 2*p[3]*u[2]*u[1] p[3]*u[1]*u[1] -1.]
    @test all(abs.(test_jac .- real_jac) .< 1e-9)
end

jacobian_network_3 = @reaction_network begin
    k1, 2A → B
    k2, B → 2A
    k3, A + B → C
    k4, C → A + B
    k5, 3C → 3A
    k6, 0 → 2B
    hill(A,k7,k8,2), ∅ → B
end k1 k2 k3 k4 k5 k6 k7 k8
for factor in [1e-2, 1e-1, 1e0, 1e1, 1e2], repeat = 1:10
    u = factor*rand(3); p = factor*rand(8);
    A,B,C = u
    k1,k2,k3,k4,k5,k6,k7,k8 = p
    test_jac = ODEFunction(convert(ODESystem,jacobian_network_3),jac=true).jac(u,p,0.)
    real_jac = [-2*k1*A-k3*B 2*k2-k3*A k4+3*k5*C^2/2; k1*A-k3*B+2*k7*k8^2*A/(k8^2+A^2)^2 -k2-k3*A k4; k3*B k3*A -k4-3*k5*C^2/2]
    @test all(abs.(test_jac .- real_jac) .< 10e-9)
end


"""
Fast (non-allocating) version of norm(x-y,2)^2
"""
function normdiff2(x::AbstractArray{C}, y::AbstractArray{C}) where {
    R <: Real, C <: Union{R, Complex{R}}
}
    s = R(0)
    for i in eachindex(x)
        s += abs2(x[i]-y[i])
    end
    return s
end

"""
Fast (non-allocating) version of norm(x-y,2)
"""
normdiff(x, y) = sqrt(normdiff2(x, y))


using Test
using LinearAlgebra
using LinearAlgebraicRepresentation
Lar = LinearAlgebraicRepresentation

@testset "submanifold_mapping test" begin
    V = rand(3, 3)
    m = Lar.Arrangement.submanifold_mapping(V)
    err = 1e-10

    @test any(map((x,y)->-err<x-y<err, m*inv(m), Matrix{Float64}(LinearAlgebra.I, 4, 4)))
    @test any(x->-err<x<err, ([V [1; 1; 1]]*m)[:, 3])
end


import Base.isempty

export AbstractPolytope,
       vertices_list,
       isempty,
       minkowski_sum

"""
    AbstractPolytope{N} <: AbstractPolyhedron{N}

Abstract type for compact convex polytopic sets.

### Notes

Every concrete `AbstractPolytope` must define the following functions:
- `vertices_list(::AbstractPolytope{N})` -- return a list of all vertices

```jldoctest; setup = :(using LazySets: subtypes)
julia> subtypes(AbstractPolytope)
4-element Array{Any,1}:
 AbstractCentrallySymmetricPolytope
 AbstractPolygon
 HPolytope
 VPolytope
```

A polytope is a bounded polyhedron (see [`AbstractPolyhedron`](@ref)).
Polytopes are compact convex sets with either of the following equivalent
properties:
1. They are the intersection of a finite number of closed half-spaces.
2. They are the convex hull of finitely many vertices.
"""
abstract type AbstractPolytope{N} <: AbstractPolyhedron{N} end

isconvextype(::Type{<:AbstractPolytope}) = true

# =============================================
# Common AbstractPolytope functions
# =============================================

"""
    isbounded(P::AbstractPolytope)

Determine whether a polytopic set is bounded.

### Input

- `P` -- polytopic set

### Output

`true` (since a polytope must be bounded).
"""
function isbounded(::AbstractPolytope)
    return true
end

"""
    isempty(P::AbstractPolytope)

Determine whether a polytope is empty.

### Input

- `P` -- abstract polytope

### Output

`true` if the given polytope contains no vertices, and `false` otherwise.

### Algorithm

This algorithm checks whether the `vertices_list` of the given polytope is empty
or not.
"""
function isempty(P::AbstractPolytope)
    return isempty(vertices_list(P))
end

"""
    isuniversal(P::AbstractPolytope{N}, [witness]::Bool=false) where {N}

Check whether a polyhedron is universal.

### Input

- `P`       -- polyhedron
- `witness` -- (optional, default: `false`) compute a witness if activated

### Output

* If `witness` option is deactivated: `false`
* If `witness` option is activated: `(false, v)` where ``v ∉ P``

### Algorithm

A witness is produced using `isuniversal(H)` where `H` is the first linear
constraint of `P`.
"""
function isuniversal(P::AbstractPolytope{N}, witness::Bool=false) where {N}
    if witness
        constraints = constraints_list(P)
        if isempty(constraints)
            return (true, N[])  # special case for polytopes without constraints
        end
        return isuniversal(constraints[1], true)
    else
        return false
    end
end

# given a polytope P, apply the linear map P to each vertex of P
# it is assumed that the interface function `vertices_list(P)` is available
@inline function _linear_map_vrep(M::AbstractMatrix, P::AbstractPolytope)
    vertices = broadcast(v -> M * v, vertices_list(P))
    m = size(M, 1) # output dimension
    if m == 1
        # TODO: substitute with Interval(convex_hull(vertices)...): convex hull 1D
        return Interval(minimum(vertices)[1], maximum(vertices)[1])
    elseif m == 2
        return VPolygon(vertices)
    else
        return VPolytope(vertices)
    end
end

function _linear_map_hrep_helper(M::AbstractMatrix, P::AbstractPolytope,
                                 algo::AbstractLinearMapAlgorithm)
    constraints = _linear_map_hrep(M, P, algo)
    m = size(M, 1) # output dimension
    if m == 1
        # TODO: create interval directly ?
        return convert(Interval, HPolytope(constraints))
    elseif m == 2
        return HPolygon(constraints)
    else
        return HPolytope(constraints)
    end
end

# the "backend" argument is ignored, used for dispatch
function _vertices_list(P::AbstractPolytope, backend)
    return vertices_list(P)
end


using Clapeyron, Test

const SOL = Clapeyron.Solvers

quadratic(x) = x*x - 4
rosenbrock(x) = (1.0 - x[1])^2 + 100.0 * (x[2] - x[1]^2)^2
golden_number_fixpoint(x) = one(x) + one(x)/x
@testset "Solvers Module" begin
    @testset "ad_newton" begin
        x0 = 2+rand()
        fdf = SOL.f∂f(quadratic,x0)
        @test fdf[1] ≈ quadratic(x0)
        @test fdf[2] ≈ 2*x0
        @test SOL.ad_newton(quadratic,x0) ≈ 2.0
    end

    @testset "box optimize" begin
        #example of FMinBox in Optim.jl
        lower = [1.25, -2.1]
        upper = [Inf, Inf]
        solution = [1.25, 1.5625]
        initial_x = [2.0, 2.0]
        res = SOL.box_optimize(rosenbrock,initial_x,lower,upper)
        @test all(res.info.minimizer .≈ solution)
        @test res.info.minimum ≈ 0.0625
    end

    @testset "fixpoint" begin
        #example of FMinBox in Optim.jl
        x0 = 2+ rand()
        ϕ = (1+sqrt(5))/2
        @test SOL.fixpoint(golden_number_fixpoint,x0) ≈ ϕ
    end

    function f_diffmcp!(fvec, x)
        fvec[1] = (1-x[1])^2-1.01
    end

    # A difficult MCP.
    #
    # Presented in Miranda and Fackler (2002): "Applied Computational Economics and
    # Finance", p. 51
    #does not converge in NLSolve.jl converges here with NLSolvers.jl
    @testset "nlsolve" begin
        res = SOL.nlsolve(f_diffmcp!,[0.0])
        solution = SOL.x_sol(res)
        v = [1.0]
        f_diffmcp!(v,solution)
        @test v[1] <= sqrt(eps(Float64))
    end
end


Base.:(==)(u::Point, v::Point) = u.x == v.x && u.y == v.y
Base.:-(u::Point) = Point(-u.x, -u.y)
Base.:+(u::Point, v::Point) = Point(u.x + v.x, u.y + v.y)
Base.:-(u::Point, v::Point) = u + (-v)


# This file is a part of Julia. License is MIT: https://julialang.org/license

"""
    Statistics

Standard library module for basic statistics functionality.
"""
module Statistics

using LinearAlgebra, SparseArrays

using Base: has_offset_axes

export cor, cov, std, stdm, var, varm, mean!, mean,
    median!, median, middle, quantile!, quantile

##### mean #####

"""
    mean(itr)

Compute the mean of all elements in a collection.

!!! note
    If `itr` contains `NaN` or [`missing`](@ref) values, the result is also
    `NaN` or `missing` (`missing` takes precedence if array contains both).
    Use the [`skipmissing`](@ref) function to omit `missing` entries and compute the
    mean of non-missing values.

# Examples
```jldoctest
julia> mean(1:20)
10.5

julia> mean([1, missing, 3])
missing

julia> mean(skipmissing([1, missing, 3]))
2.0
```
"""
mean(itr) = mean(identity, itr)

"""
    mean(f::Function, itr)

Apply the function `f` to each element of collection `itr` and take the mean.

```jldoctest
julia> mean(√, [1, 2, 3])
1.3820881233139908

julia> mean([√1, √2, √3])
1.3820881233139908
```
"""
function mean(f, itr)
    y = iterate(itr)
    if y === nothing
        return Base.mapreduce_empty_iter(f, Base.add_sum, itr,
                                         Base.IteratorEltype(itr)) / 0
    end
    count = 1
    value, state = y
    f_value = f(value)
    total = Base.reduce_first(Base.add_sum, f_value)
    y = iterate(itr, state)
    while y !== nothing
        value, state = y
        total += f(value)
        count += 1
        y = iterate(itr, state)
    end
    return total/count
end
mean(f, A::AbstractArray) = sum(f, A) / length(A)

"""
    mean!(r, v)

Compute the mean of `v` over the singleton dimensions of `r`, and write results to `r`.

# Examples
```jldoctest
julia> v = [1 2; 3 4]
2×2 Array{Int64,2}:
 1  2
 3  4

julia> mean!([1., 1.], v)
2-element Array{Float64,1}:
 1.5
 3.5

julia> mean!([1. 1.], v)
1×2 Array{Float64,2}:
 2.0  3.0
```
"""
function mean!(R::AbstractArray, A::AbstractArray)
    sum!(R, A; init=true)
    x = max(1, length(R)) // length(A)
    R .= R .* x
    return R
end

"""
    mean(A::AbstractArray; dims)

Compute the mean of an array over the given dimensions.

# Examples
```jldoctest
julia> A = [1 2; 3 4]
2×2 Array{Int64,2}:
 1  2
 3  4

julia> mean(A, dims=1)
1×2 Array{Float64,2}:
 2.0  3.0

julia> mean(A, dims=2)
2×1 Array{Float64,2}:
 1.5
 3.5
```
"""
mean(A::AbstractArray; dims=:) = _mean(A, dims)

_mean(A::AbstractArray{T}, region) where {T} = mean!(Base.reducedim_init(t -> t/2, +, A, region), A)
_mean(A::AbstractArray, ::Colon) = sum(A) / length(A)

function mean(r::AbstractRange{<:Real})
    isempty(r) && return oftype((first(r) + last(r)) / 2, NaN)
    (first(r) + last(r)) / 2
end

median(r::AbstractRange{<:Real}) = mean(r)

##### variances #####

# faster computation of real(conj(x)*y)
realXcY(x::Real, y::Real) = x*y
realXcY(x::Complex, y::Complex) = real(x)*real(y) + imag(x)*imag(y)

var(iterable; corrected::Bool=true, mean=nothing) = _var(iterable, corrected, mean)

function _var(iterable, corrected::Bool, mean)
    y = iterate(iterable)
    if y === nothing
        T = eltype(iterable)
        return oftype((abs2(zero(T)) + abs2(zero(T)))/2, NaN)
    end
    count = 1
    value, state = y
    y = iterate(iterable, state)
    if mean === nothing
        # Use Welford algorithm as seen in (among other places)
        # Knuth's TAOCP, Vol 2, page 232, 3rd edition.
        M = value / 1
        S = real(zero(M))
        while y !== nothing
            value, state = y
            y = iterate(iterable, state)
            count += 1
            new_M = M + (value - M) / count
            S = S + realXcY(value - M, value - new_M)
            M = new_M
        end
        return S / (count - Int(corrected))
    elseif isa(mean, Number) # mean provided
        # Cannot use a compensated version, e.g. the one from
        # "Updating Formulae and a Pairwise Algorithm for Computing Sample Variances."
        # by Chan, Golub, and LeVeque, Technical Report STAN-CS-79-773,
        # Department of Computer Science, Stanford University,
        # because user can provide mean value that is different to mean(iterable)
        sum2 = abs2(value - mean::Number)
        while y !== nothing
            value, state = y
            y = iterate(iterable, state)
            count += 1
            sum2 += abs2(value - mean)
        end
        return sum2 / (count - Int(corrected))
    else
        throw(ArgumentError("invalid value of mean, $(mean)::$(typeof(mean))"))
    end
end

centralizedabs2fun(m) = x -> abs2.(x - m)
centralize_sumabs2(A::AbstractArray, m) =
    mapreduce(centralizedabs2fun(m), +, A)
centralize_sumabs2(A::AbstractArray, m, ifirst::Int, ilast::Int) =
    Base.mapreduce_impl(centralizedabs2fun(m), +, A, ifirst, ilast)

function centralize_sumabs2!(R::AbstractArray{S}, A::AbstractArray, means::AbstractArray) where S
    # following the implementation of _mapreducedim! at base/reducedim.jl
    lsiz = Base.check_reducedims(R,A)
    isempty(R) || fill!(R, zero(S))
    isempty(A) && return R

    if Base.has_fast_linear_indexing(A) && lsiz > 16 && !has_offset_axes(R, means)
        nslices = div(length(A), lsiz)
        ibase = first(LinearIndices(A))-1
        for i = 1:nslices
            @inbounds R[i] = centralize_sumabs2(A, means[i], ibase+1, ibase+lsiz)
            ibase += lsiz
        end
        return R
    end
    indsAt, indsRt = Base.safe_tail(axes(A)), Base.safe_tail(axes(R)) # handle d=1 manually
    keep, Idefault = Broadcast.shapeindexer(indsRt)
    if Base.reducedim1(R, A)
        i1 = first(Base.axes1(R))
        @inbounds for IA in CartesianIndices(indsAt)
            IR = Broadcast.newindex(IA, keep, Idefault)
            r = R[i1,IR]
            m = means[i1,IR]
            @simd for i in axes(A, 1)
                r += abs2(A[i,IA] - m)
            end
            R[i1,IR] = r
        end
    else
        @inbounds for IA in CartesianIndices(indsAt)
            IR = Broadcast.newindex(IA, keep, Idefault)
            @simd for i in axes(A, 1)
                R[i,IR] += abs2(A[i,IA] - means[i,IR])
            end
        end
    end
    return R
end

function varm!(R::AbstractArray{S}, A::AbstractArray, m::AbstractArray; corrected::Bool=true) where S
    if isempty(A)
        fill!(R, convert(S, NaN))
    else
        rn = div(length(A), length(R)) - Int(corrected)
        centralize_sumabs2!(R, A, m)
        R .= R .* (1 // rn)
    end
    return R
end

"""
    varm(v, m; dims, corrected::Bool=true)

Compute the sample variance of a collection `v` with known mean(s) `m`,
optionally over the given dimensions. `m` may contain means for each dimension of
`v`. If `corrected` is `true`, then the sum is scaled with `n-1`,
whereas the sum is scaled with `n` if `corrected` is `false` where `n = length(v)`.

!!! note
    If array contains `NaN` or [`missing`](@ref) values, the result is also
    `NaN` or `missing` (`missing` takes precedence if array contains both).
    Use the [`skipmissing`](@ref) function to omit `missing` entries and compute the
    variance of non-missing values.
"""
varm(A::AbstractArray, m::AbstractArray; corrected::Bool=true, dims=:) = _varm(A, m, corrected, dims)

_varm(A::AbstractArray{T}, m, corrected::Bool, region) where {T} =
    varm!(Base.reducedim_init(t -> abs2(t)/2, +, A, region), A, m; corrected=corrected)

varm(A::AbstractArray, m; corrected::Bool=true) = _varm(A, m, corrected, :)

function _varm(A::AbstractArray{T}, m, corrected::Bool, ::Colon) where T
    n = length(A)
    n == 0 && return oftype((abs2(zero(T)) + abs2(zero(T)))/2, NaN)
    return centralize_sumabs2(A, m) / (n - Int(corrected))
end

"""
    var(v; dims, corrected::Bool=true, mean=nothing)

Compute the sample variance of a vector or array `v`, optionally along the given dimensions.
The algorithm will return an estimator of the generative distribution's variance
under the assumption that each entry of `v` is an IID drawn from that generative
distribution. This computation is equivalent to calculating `sum(abs2, v - mean(v)) /
(length(v) - 1)`. If `corrected` is `true`, then the sum is scaled with `n-1`,
whereas the sum is scaled with `n` if `corrected` is `false` where `n = length(v)`.
The mean `mean` over the region may be provided.

!!! note
    If array contains `NaN` or [`missing`](@ref) values, the result is also
    `NaN` or `missing` (`missing` takes precedence if array contains both).
    Use the [`skipmissing`](@ref) function to omit `missing` entries and compute the
    variance of non-missing values.
"""
var(A::AbstractArray; corrected::Bool=true, mean=nothing, dims=:) = _var(A, corrected, mean, dims)

_var(A::AbstractArray, corrected::Bool, mean, dims) =
    varm(A, something(mean, Statistics.mean(A, dims=dims)); corrected=corrected, dims=dims)

_var(A::AbstractArray, corrected::Bool, mean, ::Colon) =
    real(varm(A, something(mean, Statistics.mean(A)); corrected=corrected))

varm(iterable, m; corrected::Bool=true) = _var(iterable, corrected, m)

## variances over ranges

varm(v::AbstractRange, m::AbstractArray) = range_varm(v, m)
varm(v::AbstractRange, m) = range_varm(v, m)

function range_varm(v::AbstractRange, m)
    f  = first(v) - m
    s  = step(v)
    l  = length(v)
    vv = f^2 * l / (l - 1) + f * s * l + s^2 * l * (2 * l - 1) / 6
    if l == 0 || l == 1
        return typeof(vv)(NaN)
    end
    return vv
end

function var(v::AbstractRange)
    s  = step(v)
    l  = length(v)
    vv = abs2(s) * (l + 1) * l / 12
    if l == 0 || l == 1
        return typeof(vv)(NaN)
    end
    return vv
end

##### standard deviation #####

function sqrt!(A::AbstractArray)
    for i in eachindex(A)
        @inbounds A[i] = sqrt(A[i])
    end
    A
end

stdm(A::AbstractArray, m; corrected::Bool=true) =
    sqrt.(varm(A, m; corrected=corrected))

"""
    std(v; corrected::Bool=true, mean=nothing, dims)

Compute the sample standard deviation of a vector or array `v`, optionally along the given
dimensions. The algorithm returns an estimator of the generative distribution's standard
deviation under the assumption that each entry of `v` is an IID drawn from that generative
distribution. This computation is equivalent to calculating `sqrt(sum((v - mean(v)).^2) /
(length(v) - 1))`. A pre-computed `mean` may be provided. If `corrected` is `true`,
then the sum is scaled with `n-1`, whereas the sum is scaled with `n` if `corrected` is
`false` where `n = length(v)`.

!!! note
    If array contains `NaN` or [`missing`](@ref) values, the result is also
    `NaN` or `missing` (`missing` takes precedence if array contains both).
    Use the [`skipmissing`](@ref) function to omit `missing` entries and compute the
    standard deviation of non-missing values.
"""
std(A::AbstractArray; corrected::Bool=true, mean=nothing, dims=:) = _std(A, corrected, mean, dims)

_std(A::AbstractArray, corrected::Bool, mean, dims) =
    sqrt.(var(A; corrected=corrected, mean=mean, dims=dims))

_std(A::AbstractArray, corrected::Bool, mean, ::Colon) =
    sqrt.(var(A; corrected=corrected, mean=mean))

_std(A::AbstractArray{<:AbstractFloat}, corrected::Bool, mean, dims) =
    sqrt!(var(A; corrected=corrected, mean=mean, dims=dims))

_std(A::AbstractArray{<:AbstractFloat}, corrected::Bool, mean, ::Colon) =
    sqrt.(var(A; corrected=corrected, mean=mean))

std(iterable; corrected::Bool=true, mean=nothing) =
    sqrt(var(iterable, corrected=corrected, mean=mean))

"""
    stdm(v, m; corrected::Bool=true)

Compute the sample standard deviation of a vector `v`
with known mean `m`. If `corrected` is `true`,
then the sum is scaled with `n-1`, whereas the sum is
scaled with `n` if `corrected` is `false` where `n = length(v)`.

!!! note
    If array contains `NaN` or [`missing`](@ref) values, the result is also
    `NaN` or `missing` (`missing` takes precedence if array contains both).
    Use the [`skipmissing`](@ref) function to omit `missing` entries and compute the
    standard deviation of non-missing values.
"""
stdm(iterable, m; corrected::Bool=true) =
    std(iterable, corrected=corrected, mean=m)

###### covariance ######

# auxiliary functions

_conj(x::AbstractArray{<:Real}) = x
_conj(x::AbstractArray) = conj(x)

_getnobs(x::AbstractVector, vardim::Int) = length(x)
_getnobs(x::AbstractMatrix, vardim::Int) = size(x, vardim)

function _getnobs(x::AbstractVecOrMat, y::AbstractVecOrMat, vardim::Int)
    n = _getnobs(x, vardim)
    _getnobs(y, vardim) == n || throw(DimensionMismatch("dimensions of x and y mismatch"))
    return n
end

_vmean(x::AbstractVector, vardim::Int) = mean(x)
_vmean(x::AbstractMatrix, vardim::Int) = mean(x, dims=vardim)

# core functions

unscaled_covzm(x::AbstractVector{<:Number})    = sum(abs2, x)
unscaled_covzm(x::AbstractVector)              = sum(t -> t*t', x)
unscaled_covzm(x::AbstractMatrix, vardim::Int) = (vardim == 1 ? _conj(x'x) : x * x')

unscaled_covzm(x::AbstractVector, y::AbstractVector) = sum(conj(y[i])*x[i] for i in eachindex(y, x))
unscaled_covzm(x::AbstractVector, y::AbstractMatrix, vardim::Int) =
    (vardim == 1 ? *(transpose(x), _conj(y)) : *(transpose(x), transpose(_conj(y))))
unscaled_covzm(x::AbstractMatrix, y::AbstractVector, vardim::Int) =
    (c = vardim == 1 ? *(transpose(x), _conj(y)) :  x * _conj(y); reshape(c, length(c), 1))
unscaled_covzm(x::AbstractMatrix, y::AbstractMatrix, vardim::Int) =
    (vardim == 1 ? *(transpose(x), _conj(y)) : *(x, adjoint(y)))

# covzm (with centered data)

covzm(x::AbstractVector; corrected::Bool=true) = unscaled_covzm(x) / (length(x) - Int(corrected))
function covzm(x::AbstractMatrix, vardim::Int=1; corrected::Bool=true)
    C = unscaled_covzm(x, vardim)
    T = promote_type(typeof(first(C) / 1), eltype(C))
    A = convert(AbstractMatrix{T}, C)
    b = 1//(size(x, vardim) - corrected)
    A .= A .* b
    return A
end
covzm(x::AbstractVector, y::AbstractVector; corrected::Bool=true) =
    unscaled_covzm(x, y) / (length(x) - Int(corrected))
function covzm(x::AbstractVecOrMat, y::AbstractVecOrMat, vardim::Int=1; corrected::Bool=true)
    C = unscaled_covzm(x, y, vardim)
    T = promote_type(typeof(first(C) / 1), eltype(C))
    A = convert(AbstractArray{T}, C)
    b = 1//(_getnobs(x, y, vardim) - corrected)
    A .= A .* b
    return A
end

# covm (with provided mean)
## Use map(t -> t - xmean, x) instead of x .- xmean to allow for Vector{Vector}
## which can't be handled by broadcast
covm(x::AbstractVector, xmean; corrected::Bool=true) =
    covzm(map(t -> t - xmean, x); corrected=corrected)
covm(x::AbstractMatrix, xmean, vardim::Int=1; corrected::Bool=true) =
    covzm(x .- xmean, vardim; corrected=corrected)
covm(x::AbstractVector, xmean, y::AbstractVector, ymean; corrected::Bool=true) =
    covzm(map(t -> t - xmean, x), map(t -> t - ymean, y); corrected=corrected)
covm(x::AbstractVecOrMat, xmean, y::AbstractVecOrMat, ymean, vardim::Int=1; corrected::Bool=true) =
    covzm(x .- xmean, y .- ymean, vardim; corrected=corrected)

# cov (API)
"""
    cov(x::AbstractVector; corrected::Bool=true)

Compute the variance of the vector `x`. If `corrected` is `true` (the default) then the sum
is scaled with `n-1`, whereas the sum is scaled with `n` if `corrected` is `false` where `n = length(x)`.
"""
cov(x::AbstractVector; corrected::Bool=true) = covm(x, mean(x); corrected=corrected)

"""
    cov(X::AbstractMatrix; dims::Int=1, corrected::Bool=true)

Compute the covariance matrix of the matrix `X` along the dimension `dims`. If `corrected`
is `true` (the default) then the sum is scaled with `n-1`, whereas the sum is scaled with `n`
if `corrected` is `false` where `n = size(X, dims)`.
"""
cov(X::AbstractMatrix; dims::Int=1, corrected::Bool=true) =
    covm(X, _vmean(X, dims), dims; corrected=corrected)

"""
    cov(x::AbstractVector, y::AbstractVector; corrected::Bool=true)

Compute the covariance between the vectors `x` and `y`. If `corrected` is `true` (the
default), computes ``\\frac{1}{n-1}\\sum_{i=1}^n (x_i-\\bar x) (y_i-\\bar y)^*`` where
``*`` denotes the complex conjugate and `n = length(x) = length(y)`. If `corrected` is
`false`, computes ``\\frac{1}{n}\\sum_{i=1}^n (x_i-\\bar x) (y_i-\\bar y)^*``.
"""
cov(x::AbstractVector, y::AbstractVector; corrected::Bool=true) =
    covm(x, mean(x), y, mean(y); corrected=corrected)

"""
    cov(X::AbstractVecOrMat, Y::AbstractVecOrMat; dims::Int=1, corrected::Bool=true)

Compute the covariance between the vectors or matrices `X` and `Y` along the dimension
`dims`. If `corrected` is `true` (the default) then the sum is scaled with `n-1`, whereas
the sum is scaled with `n` if `corrected` is `false` where `n = size(X, dims) = size(Y, dims)`.
"""
cov(X::AbstractVecOrMat, Y::AbstractVecOrMat; dims::Int=1, corrected::Bool=true) =
    covm(X, _vmean(X, dims), Y, _vmean(Y, dims), dims; corrected=corrected)

##### correlation #####

"""
    clampcor(x)

Clamp a real correlation to between -1 and 1, leaving complex correlations unchanged
"""
clampcor(x::Real) = clamp(x, -1, 1)
clampcor(x) = x

# cov2cor!

function cov2cor!(C::AbstractMatrix{T}, xsd::AbstractArray) where T
    @assert !has_offset_axes(C, xsd)
    nx = length(xsd)
    size(C) == (nx, nx) || throw(DimensionMismatch("inconsistent dimensions"))
    for j = 1:nx
        for i = 1:j-1
            C[i,j] = adjoint(C[j,i])
        end
        C[j,j] = oneunit(T)
        for i = j+1:nx
            C[i,j] = clampcor(C[i,j] / (xsd[i] * xsd[j]))
        end
    end
    return C
end
function cov2cor!(C::AbstractMatrix, xsd, ysd::AbstractArray)
    @assert !has_offset_axes(C, ysd)
    nx, ny = size(C)
    length(ysd) == ny || throw(DimensionMismatch("inconsistent dimensions"))
    for (j, y) in enumerate(ysd)   # fixme (iter): here and in all `cov2cor!` we assume that `C` is efficiently indexed by integers
        for i in 1:nx
            C[i,j] = clampcor(C[i, j] / (xsd * y))
        end
    end
    return C
end
function cov2cor!(C::AbstractMatrix, xsd::AbstractArray, ysd)
    @assert !has_offset_axes(C, xsd)
    nx, ny = size(C)
    length(xsd) == nx || throw(DimensionMismatch("inconsistent dimensions"))
    for j in 1:ny
        for (i, x) in enumerate(xsd)
            C[i,j] = clampcor(C[i,j] / (x * ysd))
        end
    end
    return C
end
function cov2cor!(C::AbstractMatrix, xsd::AbstractArray, ysd::AbstractArray)
    @assert !has_offset_axes(C, xsd, ysd)
    nx, ny = size(C)
    (length(xsd) == nx && length(ysd) == ny) ||
        throw(DimensionMismatch("inconsistent dimensions"))
    for (i, x) in enumerate(xsd)
        for (j, y) in enumerate(ysd)
            C[i,j] = clampcor(C[i,j] / (x * y))
        end
    end
    return C
end

# corzm (non-exported, with centered data)

corzm(x::AbstractVector{T}) where {T} = one(real(T))
function corzm(x::AbstractMatrix, vardim::Int=1)
    c = unscaled_covzm(x, vardim)
    return cov2cor!(c, collect(sqrt(c[i,i]) for i in 1:min(size(c)...)))
end
corzm(x::AbstractVector, y::AbstractMatrix, vardim::Int=1) =
    cov2cor!(unscaled_covzm(x, y, vardim), sqrt(sum(abs2, x)), sqrt!(sum(abs2, y, dims=vardim)))
corzm(x::AbstractMatrix, y::AbstractVector, vardim::Int=1) =
    cov2cor!(unscaled_covzm(x, y, vardim), sqrt!(sum(abs2, x, dims=vardim)), sqrt(sum(abs2, y)))
corzm(x::AbstractMatrix, y::AbstractMatrix, vardim::Int=1) =
    cov2cor!(unscaled_covzm(x, y, vardim), sqrt!(sum(abs2, x, dims=vardim)), sqrt!(sum(abs2, y, dims=vardim)))

# corm

corm(x::AbstractVector{T}, xmean) where {T} = one(real(T))
corm(x::AbstractMatrix, xmean, vardim::Int=1) = corzm(x .- xmean, vardim)
function corm(x::AbstractVector, mx, y::AbstractVector, my)
    @assert !has_offset_axes(x, y)
    n = length(x)
    length(y) == n || throw(DimensionMismatch("inconsistent lengths"))
    n > 0 || throw(ArgumentError("correlation only defined for non-empty vectors"))

    @inbounds begin
        # Initialize the accumulators
        xx = zero(sqrt(abs2(x[1])))
        yy = zero(sqrt(abs2(y[1])))
        xy = zero(x[1] * y[1]')

        @simd for i in eachindex(x, y)
            xi = x[i] - mx
            yi = y[i] - my
            xx += abs2(xi)
            yy += abs2(yi)
            xy += xi * yi'
        end
    end
    return clampcor(xy / max(xx, yy) / sqrt(min(xx, yy) / max(xx, yy)))
end

corm(x::AbstractVecOrMat, xmean, y::AbstractVecOrMat, ymean, vardim::Int=1) =
    corzm(x .- xmean, y .- ymean, vardim)

# cor
"""
    cor(x::AbstractVector)

Return the number one.
"""
cor(x::AbstractVector) = one(real(eltype(x)))

"""
    cor(X::AbstractMatrix; dims::Int=1)

Compute the Pearson correlation matrix of the matrix `X` along the dimension `dims`.
"""
cor(X::AbstractMatrix; dims::Int=1) = corm(X, _vmean(X, dims), dims)

"""
    cor(x::AbstractVector, y::AbstractVector)

Compute the Pearson correlation between the vectors `x` and `y`.
"""
cor(x::AbstractVector, y::AbstractVector) = corm(x, mean(x), y, mean(y))

"""
    cor(X::AbstractVecOrMat, Y::AbstractVecOrMat; dims=1)

Compute the Pearson correlation between the vectors or matrices `X` and `Y` along the dimension `dims`.
"""
cor(x::AbstractVecOrMat, y::AbstractVecOrMat; dims::Int=1) =
    corm(x, _vmean(x, dims), y, _vmean(y, dims), dims)

##### median & quantiles #####

"""
    middle(x)

Compute the middle of a scalar value, which is equivalent to `x` itself, but of the type of `middle(x, x)` for consistency.
"""
middle(x::Union{Bool,Int8,Int16,Int32,Int64,Int128,UInt8,UInt16,UInt32,UInt64,UInt128}) = Float64(x)
# Specialized functions for real types allow for improved performance
middle(x::AbstractFloat) = x
middle(x::Real) = (x + zero(x)) / 1

"""
    middle(x, y)

Compute the middle of two reals `x` and `y`, which is
equivalent in both value and type to computing their mean (`(x + y) / 2`).
"""
middle(x::Real, y::Real) = x/2 + y/2

"""
    middle(range)

Compute the middle of a range, which consists of computing the mean of its extrema.
Since a range is sorted, the mean is performed with the first and last element.

```jldoctest
julia> middle(1:10)
5.5
```
"""
middle(a::AbstractRange) = middle(a[1], a[end])

"""
    middle(a)

Compute the middle of an array `a`, which consists of finding its
extrema and then computing their mean.

```jldoctest
julia> a = [1,2,3.6,10.9]
4-element Array{Float64,1}:
  1.0
  2.0
  3.6
 10.9

julia> middle(a)
5.95
```
"""
middle(a::AbstractArray) = ((v1, v2) = extrema(a); middle(v1, v2))

"""
    median!(v)

Like [`median`](@ref), but may overwrite the input vector.
"""
function median!(v::AbstractVector)
    isempty(v) && throw(ArgumentError("median of an empty array is undefined, $(repr(v))"))
    eltype(v)>:Missing && any(ismissing, v) && return missing
    (eltype(v)<:AbstractFloat || eltype(v)>:AbstractFloat) && any(isnan, v) && return convert(eltype(v), NaN)
    inds = axes(v, 1)
    n = length(inds)
    mid = div(first(inds)+last(inds),2)
    if isodd(n)
        return middle(partialsort!(v,mid))
    else
        m = partialsort!(v, mid:mid+1)
        return middle(m[1], m[2])
    end
end
median!(v::AbstractArray) = median!(vec(v))

"""
    median(itr)

Compute the median of all elements in a collection.
For an even number of elements no exact median element exists, so the result is
equivalent to calculating mean of two median elements.

!!! note
    If `itr` contains `NaN` or [`missing`](@ref) values, the result is also
    `NaN` or `missing` (`missing` takes precedence if `itr` contains both).
    Use the [`skipmissing`](@ref) function to omit `missing` entries and compute the
    median of non-missing values.

# Examples
```jldoctest
julia> median([1, 2, 3])
2.0

julia> median([1, 2, 3, 4])
2.5

julia> median([1, 2, missing, 4])
missing

julia> median(skipmissing([1, 2, missing, 4]))
2.0
```
"""
median(itr) = median!(collect(itr))

"""
    median(A::AbstractArray; dims)

Compute the median of an array along the given dimensions.

# Examples
```jldoctest
julia> median([1 2; 3 4], dims=1)
1×2 Array{Float64,2}:
 2.0  3.0
```
"""
median(v::AbstractArray; dims=:) = _median(v, dims)

_median(v::AbstractArray, dims) = mapslices(median!, v, dims = dims)

_median(v::AbstractArray{T}, ::Colon) where {T} = median!(copyto!(Array{T,1}(undef, length(v)), v))

# for now, use the R/S definition of quantile; may want variants later
# see ?quantile in R -- this is type 7
"""
    quantile!([q::AbstractArray, ] v::AbstractVector, p; sorted=false)

Compute the quantile(s) of a vector `v` at a specified probability or vector or tuple of
probabilities `p` on the interval [0,1]. If `p` is a vector, an optional
output array `q` may also be specified. (If not provided, a new output array is created.)
The keyword argument `sorted` indicates whether `v` can be assumed to be sorted; if
`false` (the default), then the elements of `v` will be partially sorted in-place.

Quantiles are computed via linear interpolation between the points `((k-1)/(n-1), v[k])`,
for `k = 1:n` where `n = length(v)`. This corresponds to Definition 7 of Hyndman and Fan
(1996), and is the same as the R default.

!!! note
    An `ArgumentError` is thrown if `v` contains `NaN` or [`missing`](@ref) values.

* Hyndman, R.J and Fan, Y. (1996) "Sample Quantiles in Statistical Packages",
  *The American Statistician*, Vol. 50, No. 4, pp. 361-365

# Examples
```jldoctest
julia> x = [3, 2, 1];

julia> quantile!(x, 0.5)
2.0

julia> x
3-element Array{Int64,1}:
 1
 2
 3

julia> y = zeros(3);

julia> quantile!(y, x, [0.1, 0.5, 0.9]) === y
true

julia> y
3-element Array{Float64,1}:
 1.2
 2.0
 2.8
```
"""
function quantile!(q::AbstractArray, v::AbstractVector, p::AbstractArray;
                   sorted::Bool=false)
    @assert !has_offset_axes(q, v, p)
    if size(p) != size(q)
        throw(DimensionMismatch("size of p, $(size(p)), must equal size of q, $(size(q))"))
    end
    isempty(q) && return q

    minp, maxp = extrema(p)
    _quantilesort!(v, sorted, minp, maxp)

    for (i, j) in zip(eachindex(p), eachindex(q))
        @inbounds q[j] = _quantile(v,p[i])
    end
    return q
end

quantile!(v::AbstractVector, p::AbstractArray; sorted::Bool=false) =
    quantile!(similar(p,float(eltype(v))), v, p; sorted=sorted)

quantile!(v::AbstractVector, p::Real; sorted::Bool=false) =
    _quantile(_quantilesort!(v, sorted, p, p), p)

function quantile!(v::AbstractVector, p::Tuple{Vararg{Real}}; sorted::Bool=false)
    isempty(p) && return ()
    minp, maxp = extrema(p)
    _quantilesort!(v, sorted, minp, maxp)
    return map(x->_quantile(v, x), p)
end

# Function to perform partial sort of v for quantiles in given range
function _quantilesort!(v::AbstractArray, sorted::Bool, minp::Real, maxp::Real)
    isempty(v) && throw(ArgumentError("empty data vector"))
    @assert !has_offset_axes(v)

    if !sorted
        lv = length(v)
        lo = floor(Int,1+minp*(lv-1))
        hi = ceil(Int,1+maxp*(lv-1))

        # only need to perform partial sort
        sort!(v, 1, lv, Base.Sort.PartialQuickSort(lo:hi), Base.Sort.Forward)
    end
    ismissing(v[end]) && throw(ArgumentError("quantiles are undefined in presence of missing values"))
    isnan(v[end]) && throw(ArgumentError("quantiles are undefined in presence of NaNs"))
    return v
end

# Core quantile lookup function: assumes `v` sorted
@inline function _quantile(v::AbstractVector, p::Real)
    0 <= p <= 1 || throw(ArgumentError("input probability out of [0,1] range"))
    @assert !has_offset_axes(v)

    lv = length(v)
    f0 = (lv - 1)*p # 0-based interpolated index
    t0 = trunc(f0)
    h  = f0 - t0
    i  = trunc(Int,t0) + 1

    T  = promote_type(eltype(v), typeof(v[1]*h))

    if h == 0
        return convert(T, v[i])
    else
        a = v[i]
        b = v[i+1]
        if isfinite(a) && isfinite(b)
            return convert(T, a + h*(b-a))
        else
            return convert(T, (1-h)*a + h*b)
        end
    end
end

"""
    quantile(itr, p; sorted=false)

Compute the quantile(s) of a collection `itr` at a specified probability or vector or tuple of
probabilities `p` on the interval [0,1]. The keyword argument `sorted` indicates whether
`itr` can be assumed to be sorted.

Quantiles are computed via linear interpolation between the points `((k-1)/(n-1), v[k])`,
for `k = 1:n` where `n = length(itr)`. This corresponds to Definition 7 of Hyndman and Fan
(1996), and is the same as the R default.

!!! note
    An `ArgumentError` is thrown if `itr` contains `NaN` or [`missing`](@ref) values.
    Use the [`skipmissing`](@ref) function to omit `missing` entries and compute the
    quantiles of non-missing values.

- Hyndman, R.J and Fan, Y. (1996) "Sample Quantiles in Statistical Packages",
  *The American Statistician*, Vol. 50, No. 4, pp. 361-365

# Examples
```jldoctest
julia> quantile(0:20, 0.5)
10.0

julia> quantile(0:20, [0.1, 0.5, 0.9])
3-element Array{Float64,1}:
  2.0
 10.0
 18.0

julia> quantile(skipmissing([1, 10, missing]), 0.5)
5.5
```
"""
quantile(itr, p; sorted::Bool=false) = quantile!(collect(itr), p, sorted=sorted)

quantile(v::AbstractVector, p; sorted::Bool=false) =
    quantile!(sorted ? v : Base.copymutable(v), p; sorted=sorted)

##### SparseArrays optimizations #####

function cov(X::SparseMatrixCSC; dims::Int=1, corrected::Bool=true)
    vardim = dims
    a, b = size(X)
    n, p = vardim == 1 ? (a, b) : (b, a)

    # The covariance can be decomposed into two terms
    # 1/(n - 1) ∑ (x_i - x̄)*(x_i - x̄)' = 1/(n - 1) (∑ x_i*x_i' - n*x̄*x̄')
    # which can be evaluated via a sparse matrix-matrix product

    # Compute ∑ x_i*x_i' = X'X using sparse matrix-matrix product
    out = Matrix(unscaled_covzm(X, vardim))

    # Compute x̄
    x̄ᵀ = mean(X, dims=vardim)

    # Subtract n*x̄*x̄' from X'X
    @inbounds for j in 1:p, i in 1:p
        out[i,j] -= x̄ᵀ[i] * x̄ᵀ[j]' * n
    end

    # scale with the sample size n or the corrected sample size n - 1
    return rmul!(out, inv(n - corrected))
end

# This is the function that does the reduction underlying var/std
function centralize_sumabs2!(R::AbstractArray{S}, A::SparseMatrixCSC{Tv,Ti}, means::AbstractArray) where {S,Tv,Ti}
    @assert !has_offset_axes(R, A, means)
    lsiz = Base.check_reducedims(R,A)
    size(means) == size(R) || error("size of means must match size of R")
    isempty(R) || fill!(R, zero(S))
    isempty(A) && return R

    colptr = A.colptr
    rowval = A.rowval
    nzval = A.nzval
    m = size(A, 1)
    n = size(A, 2)

    if size(R, 1) == size(R, 2) == 1
        # Reduction along both columns and rows
        R[1, 1] = centralize_sumabs2(A, means[1])
    elseif size(R, 1) == 1
        # Reduction along rows
        @inbounds for col = 1:n
            mu = means[col]
            r = convert(S, (m-colptr[col+1]+colptr[col])*abs2(mu))
            @simd for j = colptr[col]:colptr[col+1]-1
                r += abs2(nzval[j] - mu)
            end
            R[1, col] = r
        end
    elseif size(R, 2) == 1
        # Reduction along columns
        rownz = fill(convert(Ti, n), m)
        @inbounds for col = 1:n
            @simd for j = colptr[col]:colptr[col+1]-1
                row = rowval[j]
                R[row, 1] += abs2(nzval[j] - means[row])
                rownz[row] -= 1
            end
        end
        for i = 1:m
            R[i, 1] += rownz[i]*abs2(means[i])
        end
    else
        # Reduction along a dimension > 2
        @inbounds for col = 1:n
            lastrow = 0
            @simd for j = colptr[col]:colptr[col+1]-1
                row = rowval[j]
                for i = lastrow+1:row-1
                    R[i, col] = abs2(means[i, col])
                end
                R[row, col] = abs2(nzval[j] - means[row, col])
                lastrow = row
            end
            for i = lastrow+1:m
                R[i, col] = abs2(means[i, col])
            end
        end
    end
    return R
end

end # module


using Plots

func1(R) = R/2
func2(R) = R/2 + [1 ; 0]
func3(R) = R/2 + [0 ; 1]

function RandomSPT(P, num)
    x = []
    y = []
    for i in 1:num
        point = [rand() ; rand()]
        for j in 1:P
            point = rand([func1(point), func2(point), func3(point)])
        end
        push!(x, point[1])
        push!(y, point[2])
    end
    return x, y
end

for i in [5, 500]
    for j in [10000,100000]
        x,y = RandomSPT(i, j)

        scatter(x,y,markersize = 10^4/j, legend = false, border=:none, dpi=200, label = "Points number = $j Operations number = $i")
        savefig("C:\\Users\\Yaghoub\\Documents\\GitHub\\ComputationalPhysics-Fall2021\\ProblemSet1\\Figs\\Q4\\RSPTp$(i)num$(j).png")
    end
end


"""
    WBS( segment_cost, n[,th_const = 1.3, sigma = 1.0, M = 5000, do_seeded = false, shrink = 1/sqrt(2)])

Runs the Wild Binary segmentation algorithm for the CUSUM cost function `segment_cost` for a time series of length `n`.
    The test threshold is determined by `th_const` and the known or estimated standard deviation `sigma`.
    `do_seeded` determined whether to use seeded intervals. If true, `shrink` is the decay factor for interval length; if not, `M` random intervals are drawn, and returns the position of found changepoints, and
the cost of this segmentation.

See also: [`@BS`](@ref), [`@segment_cost`](@ref)

# Returns
* `result::Array{NamedTuple{s::Int64, e::Int64, cpt::Int64, CUSUM::Float64, min_th::Float64, scale::Int64}}`:
    * `s::Int64`: Start of window in which maximum CUSUM was located
    * `e::Int64`: End of window in which maximum CUSUM was located
    * `cpt::Int64`: Location of maximum CUSUM
    * `CUSUM::Float64`: Value of maximum CUSUM
    * `min_th::Float64`: Threshold used
    * `scale::Int64`: Number of recursion iterations before the change point was detected
# Example

```julia-repl
# Sample Normal time series with changing mean
n = 1000
λ = 100
μ, σ = Normal(0.0, 10.0), 1.0
sample, cps = @changepoint_sampler n λ Normal(μ, σ)
# Run wild binary segmentation
seg_cost = CUSUM(sample, σ)
WBS_return = WBS(seg_cost, n)
```

# References
Fryzlewicz, P. (2014) Wild binary segmentation for multiple change-point detection, Annals of Statistics 42(6), 2243-2281
Kovács, S., Li, H., Bühlmann, P., & Munk, A. (2020). Seeded Binary Segmentation: A general methodology for fast and optimal change point detection. arXiv preprint arXiv:2002.06633.
"""
#import StatsBase
#using StatsBase
# sigma = StatsBase.mad(x)
result_type = @NamedTuple{s::Int64, e::Int64, cpt::Int64, CUSUM::Float64, min_th::Float64, scale::Int64}

function WBS( segment_cost::Function , n::Int64, th_const::Float64 = 1.3, sigma::Float64 = 1.0,
    M::Int64 = 5000, do_seeded::Bool = false, shrink::Float64 = 1/sqrt(2))
    result = Array{result_type}(undef,0)

    # thresulthold
    th = sigma * th_const * sqrt(2 * log(n))
    scale = 0
    # call WBS
    result = WBS_RECUR(segment_cost, 1, n, th, result, scale, M, do_seeded, shrink)

    return result
end

function WBS_RECUR(segment_cost::Function, s::Int64, e::Int64, th::Float64,
    result::Array{result_type}, scale::Int64, M::Int64 = 5000, do_seeded::Bool = false, shrink::Float64 = 1/sqrt(2))
    #k = length(CP) - 2 #current number of change points
    minlength = 20 #minimum interval length
    scale += 1 #update scale

    # draw intervals
    if do_seeded
        tau = get_seeded_intervals(s,e, shrink, minlength)
        M = length(tau)
    else
        tau = Tuple{Int, Int}[] # Segmentations to test
        push!(tau, (s, e))
        for m in 2:M
            s_m = rand(s:(e-minlength))
            e_m = rand( (s_m+minlength):e)
            push!(tau, (s_m, e_m))
        end
    end

    # compute CUSUMs find candidate change point
    x = fill(0.0,M, e-s+1)#zeros(Float64, M, e-s+1)#Array{Float64}(NaN, M, e-s+1 ) #CUSUM array
    for m in 1:M
        a, b = pop!(tau)
        for j in (a + 1):(b - 2)
            x[m,j-s+1] = segment_cost(s,e,j)
        end
    end
    #x = replace(x, 0.0 => maximum(x))
    maxval , ind = findmax(x)

    # segment here
    if maxval > th #&& scale < 10
        cpt = ind[2] + s -1
        newresult = (s=s,e=e, cpt=cpt, CUSUM = maxval, min_th=th, scale=scale )
        push!(result, newresult)
        #push!(CP,cpt )
        #push!(CUSUM, minval)
        if cpt - s > 2*minlength
            result = WBS_RECUR(segment_cost, s, cpt, th, result, scale, M, do_seeded, shrink)
        end
        if e - cpt > 2*minlength
            result = WBS_RECUR(segment_cost, cpt+1, e, th, result,  scale, M, do_seeded, shrink)
        end
    end
    #CP = sort(CP)

    return result #CUSUM, CP
end

"""
    get_WBS_changepoints(segment_cost, object, Kmax = 1,  alpha = 1.01)

Obtains changepoints from an `object` (output by the WBS function) by minimising an information criterion.
`segment_cost` is a cost function, normally sSIC, and `alpha` is the penalty exponent. `Kmax` determines the maximum number of changepoints to detect.

See also: [`@WBS`](@ref), [`@segment_cost`](@ref)

# Returns
* `out_cps`: Vector of changepoint locations under optimal segmentation
* `optimal_number`: Integer number of changepoints which optimises the criterion
* `cost`: Vector of segmentations cost for 0, ..., Kmax  changes
* `object[1:optimal]`: WBS object truncated at optimal number
# Example

```julia-repl
# Sample Normal time series with changing mean
n = 1000
λ = 100
μ, σ = Normal(0.0, 10.0), 1.0
sample, cps = @changepoint_sampler n λ Normal(μ, σ)

# Run wild binary segmentation
seg_cost_CUSUM = CUSUM(sample, σ)
WBS_return = WBS(seg_cost_CUSUM, n)

# Obtain changepoints
seg_cost_sSIC = sSIC(sample)
get_WBS_changepoints(seg_cost_sSIC, WBS_return, 10)
```

# References
Fryzlewicz, P. (2014) Wild binary segmentation for multiple change-point detection, Annals of Statistics 42(6), 2243-2281
"""

function get_WBS_changepoints(segment_cost::Function, object::Array{NamedTuple{(:s, :e, :cpt, :CUSUM, :min_th, :scale),Tuple{Int64,Int64,Int64,Float64,Float64,Int64}},1}, Kmax::Int64 = 1,  alpha::Float64 = 1.01)
    #result = Array{result_type}(undef,0)
    Kmax = min(Kmax,length(object))
    if length(object) > 0
        cost = zeros(Kmax+1)
        sorted = sort(object, lt = (a, b) -> ( a[4] < b[4]) ) #a[6] < b[6] &&
        cps_from_sorted = Vector{Int64}()
        for i in 1:Kmax
            push!(cps_from_sorted, sorted[i][3])
        end

        cost[1] = segment_cost(1, segment_cost.n)
        for i in 1:Kmax
            current_cps = sort(union([0, segment_cost.n], cps_from_sorted[1:i])) #select cps at current threshold
            for j in 1:(i+1)
                cost[i+1] = cost[i+1] + segment_cost(current_cps[j]+1, current_cps[j+1])
            end
        end

        pen = log(segment_cost.n)^alpha
        penvec = collect(0:Kmax) * pen
        cost += penvec

        optimal_number = findmin(cost)[2] -1

        out_cps = Vector{Int64}()
        for i in 1:optimal_number
            push!(out_cps, sorted[i][3])
        end
        out_cps = sort(out_cps)

        return  out_cps, optimal_number, cost, object[1:optimal_number]
    else
         println("no changepoints detected")
    end

end

function get_seeded_intervals(s::Int64, e::Int64,  shrink::Float64 = 1/sqrt(2), minlength = 10)
    n = e-s+1;
    max_k = Int(ceil(log(1/shrink, float(n) )))
    out = Tuple{Int, Int}[]

    for k = 1:max_k
        n_k = Int(2 * ceil(shrink^(1-k) ) - 1)
        l_k = n*shrink^(k-1)
        s_k = (n-l_k)/(n_k -1)
        if !isnan(s_k) && l_k >= minlength
            for i in 1:n_k
                left =  s + Int( floor((i-1)*s_k) )
                right = s + Int( ceil((i-1)*s_k + l_k) ) -1
                push!(out, (left, right) )
            end
        end
    end
    return out
end


module TestMatrices
using Test
using SparseArrays

using Catlab.CategoricalAlgebra.Matrices

# Dense matrices
################

const IntMat = Matrix{Int}

A, B, C = [1 2 3; 4 5 6], [1 -1; -1 2], [0 1; 1 0]
@test dom(A) == MatrixDom{IntMat}(3)
@test codom(A) == MatrixDom{IntMat}(2)
@test A⋅B == B*A
@test id(dom(A))⋅A == A
@test A⋅id(codom(A)) == A

@test A⊕(B⊕C) == (A⊕B)⊕C
@test swap(dom(A),dom(B))⋅(B⊕A)⋅swap(codom(B),codom(A)) == A⊕B
@test pair(B,C) == mcopy(dom(B))⋅(B⊕C)
@test copair(A,B) == (A⊕B)⋅plus(codom(A))

@test A⊗(B⊗C) == (A⊗B)⊗C
@test braid(dom(A),dom(B))⋅(B⊗A)⋅braid(codom(B),codom(A)) == A⊗B

# Sparse matrices
#################

const SparseIntMat = SparseMatrixCSC{Int,Int}

A, B, C = sparse(A), sparse(B), sparse(C)
@test dom(A) == MatrixDom{SparseIntMat}(3)
@test codom(A) == MatrixDom{SparseIntMat}(2)
@test A⋅B isa SparseIntMat
@test A⊕B isa SparseIntMat
@test A⊗B isa SparseIntMat
@test A⋅B == B*A
@test id(dom(A))⋅A == A
@test A⋅id(codom(A)) == A

end


iscoprime(x::Integer, y::Integer) = gcd(x,y) == 1


module Electre

import GreyDecision.GreyNumbers: GreyNumber
import GreyDecision.Utility:
    unitize, normalize, weightize, apply_columns, reverseminmax, rowmaxs

function electre(
    decisionMat::Array{GreyNumber{T},2},
    weights::Array{Float64,1},
    fns::Array{Function,1},
) where {T<:Real}

    w = unitize(weights)

    nalternatives, ncriteria = size(decisionMat)

    normalizedMat = normalize(decisionMat)
    weightednormalizedMat = weightize(normalizedMat, w)

    fitnessTable = []
    nonfitnessTable = []

    for i = 1:nalternatives
        for j = 1:nalternatives
            if i != j
                betterlist = []
                worselist = []
                for h = 1:ncriteria
                    if fns[h] == maximum
                        if weightednormalizedMat[i, h] >= weightednormalizedMat[j, h]
                            push!(betterlist, h)
                        else
                            push!(worselist, h)
                        end
                    else
                        if weightednormalizedMat[i, h] <= weightednormalizedMat[j, h]
                            push!(betterlist, h)
                        else
                            push!(worselist, h)
                        end
                    end
                end
                fitnesstableelement = Dict(:i => i, :j => j, :set => betterlist)
                push!(fitnessTable, fitnesstableelement)

                nonfitnesstableelement = Dict(:i => i, :j => j, :set => worselist)
                push!(nonfitnessTable, nonfitnesstableelement)
            end
        end
    end

    fitnessmatrix = Array{GreyNumber{Float64}, 2}(undef, nalternatives, nalternatives)
    nonfitnessmatrix = Array{GreyNumber{Float64}, 2}(undef, nalternatives, nalternatives)

    for elements in fitnessTable
        i = elements[:i]
        j = elements[:j]
        elemlist = elements[:set]

        CC = sum(w[elemlist])
        fitnessmatrix[i, j] = GreyNumber(CC, CC)
    end

    nonfitnessmatrix = Array{GreyNumber{Float64}, 2}(undef, nalternatives, nalternatives)
    for elements in nonfitnessTable
        i = elements[:i]
        j = elements[:j]
        elemlist = elements[:set]

        r_ik = weightednormalizedMat[i, elemlist]
        r_jk = weightednormalizedMat[j, elemlist]
        r_ik_full = weightednormalizedMat[i, :]
        r_jk_full = weightednormalizedMat[j, :]

        nom = maximum(abs.(r_ik - r_jk))
        dom = maximum(abs.(r_ik_full - r_jk_full))
        nonfitnessmatrix[i, j] = nom / dom
    end

    C = Array{GreyNumber{Float64}, 1}(undef, nalternatives)
    D = Array{GreyNumber{Float64}, 1}(undef, nalternatives)
    for i = 1:nalternatives
        C[i] = sum(fitnessmatrix[i, :]) - sum(fitnessmatrix[:, i])
        D[i] = sum(nonfitnessmatrix[i, :]) - sum(nonfitnessmatrix[:, i])
    end

    best_C_index = sortperm(C) |> last
    best_D_index = sortperm(D) |> first
    best = nothing

    if best_C_index == best_D_index
        best = (best_C_index,)
    else
        best = (best_C_index, best_D_index)
    end

    result = Dict(
        :C => C,
        :D => D,
        :bestindex => best
    )

    return result
end
end # end of module Electre


@testset "fit moments" begin
    D = LogNormal(1,0.6)
    M = Moments(mean(D), var(D))
    Dfit = fit(LogNormal, M)
    @test D ≈ Dfit
    # handle not giving variance
    @test_throws Exception fit(LogNormal, Moments(3.2))
end;
@testset "fit two quantiles" begin
    qpl = @qp_m(3)
    qpu = @qp_u(5)
    d = fit(LogNormal, qpl, qpu);
    @test quantile.(d, [qpl.p, qpu.p]) ≈ [qpl.q, qpu.q]
    d = fit(LogNormal, qpu, qpl) # sort
    @test quantile.(d, [qpl.p, qpu.p]) ≈ [qpl.q, qpu.q]
end;
@testset "fit to quantilepoint and mean" begin
    d = LogNormal(1,1)
    m = log(mean(d))
    qp = @qp(quantile(d,0.95),0.95)
    dfit = fit_mean_quantile(LogNormal, mean(d), qp)
    @test dfit ≈ d
    dfit = fit(LogNormal, mean(d), qp, Val(:mean))
    @test dfit ≈ d
    # with lower quantile
    qp = @qp(quantile(d,0.05),0.05)
    dfit = fit_mean_quantile(LogNormal, mean(d), qp)
    @test dfit ≈ d
    # very close to mean can give very different results:
    qp = @qp(mean(d)-1e-4,0.95)
    dfit = fit_mean_quantile(LogNormal, mean(d), qp)
    @test mean(dfit) ≈ mean(d) && quantile(dfit, qp.p) ≈ qp.q
end;
@testset "fit to quantilepoint and mode" begin
    d = LogNormal(1,1)
    m = log(mode(d))
    qp = @qp(quantile(d,0.95),0.95)
    dfit = fit_mode_quantile(LogNormal, mode(d), qp)
    @test dfit ≈ d
    dfit = fit(LogNormal, mode(d), qp, Val(:mode))
    @test dfit ≈ d
    # with lower quantile
    qp = @qp(quantile(d,0.025),0.025)
    dfit = fit_mode_quantile(LogNormal, mode(d), qp)
    @test mode(dfit) ≈ mode(d) && quantile(dfit, qp.p) ≈ qp.q
end;
@testset "fit to quantilepoint and median" begin
    d = LogNormal(1,1)
    qp = @qp(quantile(d,0.95),0.95)
    dfit = fit(LogNormal, median(d), qp, Val(:median))
    @test dfit ≈ d
end;

@testset "Σstar" begin
    ss = Σstar(4.5)
    @test ss() == 4.5
end;
@testset "fit to mean and Σstar" begin
    d = LogNormal(1, log(1.2))
    dfit = fit(LogNormal, mean(d), Σstar(1.2))
    @test d == dfit
    @test σstar(d) ≈ 1.2
end;

@testset "fit_mean_relerror" begin
    d = fit_mean_relerror(LogNormal, 10.0, 0.03);
    @test all((mean(d), std(d)/mean(d)) .≈ (10.0, 0.03))
end;


"""
There are also identification robust versions of likelihood ratio and
lagrange multiplier test. Moreire (2003)[@moreira2003] proposed a conditional
likelihood ratio (CLR) test for weakly identified linear IV
models. Kleibergen (2005)[@kleibergen2005] developed a Lagrange multiplier (often called
the KLM) test and extended Moreira's CLR test to weakly identified GMM
models.  More recently, Andrews and Guggenberge (2015) [@andrews2015]
and Andrews and Guggenberge (2017) [@andrews2017] showed the
validity of these tests under more general conditions. These tests are
somewhat more complicated than the AR test, but they have the
advantage that they are often more powerful. The AR test statistic has
a \\chi^2_{m} distribution, where m is the number of moments. The
CLR and KLM statistics under strong identification have \\chi^2_k
distributions (as does the Wald statistic), where k is the number of
parameters. Consequently, when the model is overidentified, the CLR
and LM tests are more powerful than the AR test.

Here is an implementation of the KLM and CLR statistics. The names of
variables roughly follows the notation of Andrews and
Guggenberge(2017

"""

"""
    statparts(gi::Function)

Return a function.

compute common components of klm, rk, & clr stats
follows notation of Andrews & Guggenberger 2017, section 3.1
"""
function statparts(gi::Function)
  # compute common components of klm, rk, & clr stats
  # follows notation of Andrews & Guggenberger 2017, section 3.1
    function P(A::AbstractMatrix) # projection matrix
    A*pinv(A'*A)*A'
    end
    function(θ)
        giθ = gi(θ)
        p = length(θ)
        (n, k) = size(giθ)
        Ω = cov(giθ)
        gn=mean(gi(θ), dims=1)'
        #G = ForwardDiff.jacobian(θ->mean(gi(θ),dims=1),θ)
        Gi= ForwardDiff.jacobian(gi,θ)
        Gi = reshape(Gi, n , k, p)
        G = mean(Gi, dims=1)
        Γ = zeros(eltype(G),p,k,k)
        D = zeros(eltype(G),k, p)
        for j in 1:p
        for i in 1:n
          Γ[j,:,:] += (Gi[i,:,j] .- G[1,:,j]) * giθ[i,:]'
        end
        Γ[j,:,:] ./= n
        D[:,j] = G[1,:,j] - Γ[j,:,:]*inv(Ω)*gn
        end
        return(n,k,p,gn, Ω, D, P)
    end
end

"""
    klm(gi::Function)

Return test statistic for KLM method.
"""

function klm(gi::Function)
    SP = statparts(gi)
    function(θ)
        (n,k,p,gn, Ω, D, P) = SP(θ)
        return n*(gn'*Ω^(-1/2)*P(Ω^(-1/2)*D)*Ω^(-1/2)*gn)[1]
    end
end

"""
    clr(gi::Function)

Return test statistic for CLR method.
"""
function clr(gi::Function)
  SP = statparts(gi)
  function(θ)
        (n,k,p,gn, Ω, D, P) = SP(θ)
        rk = eigmin(n*D'*inv(Ω)*D)
        AR  = (n*gn'*inv(Ω)*gn)[1]
        lm = (n*(gn'*Ω^(-1/2)*P(Ω^(-1/2)*D)*Ω^(-1/2)*gn))[1]
        lr = 1/2*(AR - rk + sqrt( (AR-rk)^2 + 4*lm*rk))

        # simulate to find p-value
        S = 5000
        function randc(k,p,r,S)
        χp = rand(Chisq(p),S)
        χkp = rand(Chisq(k-p),S)
        0.5.*(χp .+ χkp .- r .+
              sqrt.((χp .+ χkp .- r).^2 .+ 4 .* χp.*r))
        end
        csim = randc(k,p,rk,S)
        return mean(csim.<=lr)
    end
end


# Computation of MPO from given block matrices that together from a sparse matrix containing those blocks in specified locations
function sparseblockMPO(blocks::Vector{Matrix{Float64}},blocklocations::Matrix{Int64},middlesizes::Matrix{Int64},maxrank::Int64,acc::Float64)

    rowsizes = middlesizes[1,:];
    colsizes = middlesizes[2,:];
    cores0 = [zeros(1, rowsizes[e], colsizes[e], 1) for e in 1:length(rowsizes)];
    sparseblockmpo = MPT(cores0);
    # compute cores
    for a = 1:length(blocks)
        cores = Vector{Array{Float64,4}}(undef,length(rowsizes));
        # first core: block matrix
        cores[1] = reshape(blocks[a],(1, rowsizes[1], colsizes[1], 1));
        # second to last core: vectorized E matrix:
        # Zooming in from largest block to smallest
        xloc = blocklocations[a,1]
        yloc = blocklocations[a,2]
        remainingrows = prod(rowsizes[2:end-1])
        remainingcols = prod(colsizes[2:end-1])
        for e = length(rowsizes):-1:2
            xlocindiv              = Int(ceil(xloc/remainingrows)) # x location in division
            ylocindiv              = Int(ceil(yloc/remainingcols)) # y location in division

            E                      = zeros(rowsizes[e],colsizes[e]);
            E[xlocindiv,ylocindiv] = 1;
            cores[e]               = reshape(E,(1, rowsizes[e], colsizes[e], 1));

            xloc                   = xloc-(xlocindiv-1)*remainingrows
            yloc                   = yloc-(ylocindiv-1)*remainingcols
            remainingrows          = remainingrows/rowsizes[e-1]
            remainingcols          = remainingcols/colsizes[e-1]
        end

        sparseblockmpo = sparseblockmpo + MPT(cores);
        if maximum(maximum(rank(sparseblockmpo))) > maxrank
            sparseblockmpo = mpo2mps(sparseblockmpo);
            sparseblockmpo = roundTT(sparseblockmpo,acc);
            sparseblockmpo = mps2mpo(sparseblockmpo,middlesizes);
        end
    end
    return sparseblockmpo
end

## test ##
# blocks = [[1 2; 3 4],[7 7; 6 1],[1 2; 5 5.0],[5 3; 7 4]];
# blocklocations = [1 2; 2 2; 3 4; 4 1];
# middleindices = [2 2 2; 2 2 2];
# maxrank = 5;
# acc = 1e-10;
# Kmpo = sparseblockMPO(blocks,blocklocations,middleindices,maxrank,acc)
# K    = sparse(mpo2mat(Kmpo));
# print(K)


struct DirectSummand{T, M<:AbstractMatrix{T}} <: AbstractMatrix{T}
    basis::M
    multiplicity::Int
    degree::Int
    simple::Bool
end

image_basis(ds::DirectSummand) = ds.basis
issimple(ds::DirectSummand) = ds.simple
PermutationGroups.degree(ds::DirectSummand) = ds.degree
multiplicity(ds::DirectSummand) = ds.multiplicity

Base.size(ds::DirectSummand) = size(ds.basis)
Base.@propagate_inbounds Base.getindex(ds::DirectSummand, i...) =
    image_basis(ds)[i...]

function SparseArrays.sparse(ds::DirectSummand)
    sp = sparse(image_basis(ds))
    return DirectSummand(sp, multiplicity(ds), degree(ds), issimple(ds))
end

function SparseArrays.droptol!(
    ds::DirectSummand{T, <:AbstractSparseArray},
    tol
) where T
    droptol!(image_basis(ds), tol)
    return ds
end


# Copyright (c) 2021 Idiap Research Institute, http://www.idiap.ch/
#  Niccolò Antonello <nantonel@idiap.ch>

export RightStringWeight

"""
`RightStringWeight(x)`

| Set                     |        ``\\oplus``     |   ``\\otimes`` | ``\\bar{0}`` | ``\\bar{1}`` |
|:-----------------------:|:----------------------:|:--------------:|:------------:|:------------:|
|``L^*\\cup\\{\\infty\\}``|  longest common suffix |     ``\\cdot`` |``\\infty``   |``\\epsilon`` |

where ``L^*`` is Kleene closure of the set of characters ``L`` and ``\\epsilon`` the empty string.
"""
struct RightStringWeight <: Semiring
  x::String
  iszero::Bool
end
RightStringWeight(x::String) = RightStringWeight(x,false)
reversetype(::Type{S}) where {S <: RightStringWeight} = LeftStringWeight

zero(::Type{RightStringWeight}) = RightStringWeight("",true)
one(::Type{RightStringWeight})  = RightStringWeight("",false)

# longest common suffix
function lcs(a::RightStringWeight,b::RightStringWeight)
  if a.iszero
    return b
  elseif b.iszero
    return a
  else
    cnt = 0
    for x in zip(reverse(a.x),reverse(b.x))
      if x[1] == x[2]
        cnt +=1
      else
        break
      end
    end
    return RightStringWeight(a.x[end-cnt+1:end])
  end
end

*(a::T, b::T) where {T<:RightStringWeight}=
((a.iszero) | (b.iszero)) ? zero(T) : T(a.x * b.x)
+(a::RightStringWeight, b::RightStringWeight) = lcs(a,b)
function /(a::RightStringWeight, b::RightStringWeight)
  if b == zero(RightStringWeight)
    throw(ErrorException("Cannot divide by zero"))
  elseif a == zero(RightStringWeight)
    return a
  end
  return RightStringWeight(a.x[1:end-length(b.x)])
end
reverse(a::RightStringWeight) = LeftStringWeight(reverse(a.x),a.iszero)

#properties
isidempotent(::Type{W}) where {W <: RightStringWeight} = true
isright(::Type{W}) where {W <: RightStringWeight}= true
isweaklydivisible(::Type{W}) where {W <: RightStringWeight}= true


using LazySets, Polyhedra, Makie, LinearAlgebra

inputSet = Hyperrectangle(low=zeros(3), high=ones(3))

W = [1.8 -4 -2; 1.7 6 -8]

eigenvector = [1.0 0.25 0.4]
#big box
Cb = vcat(Array(Diagonal(ones(3))), Array(-Diagonal(ones(3))))
db = vcat(2*ones(3), 1*ones(3))

#small box / input LazySets
di = vcat(ones(3), zeros(3))

#project plane
C = vcat(Cb, eigenvector, -eigenvector)
d = vcat(db, 0.825, -0.825)
proj_plane = HPolytope(C, d)
#inputSet in front of project plane
infr = HPolytope(vcat(Cb, -eigenvector), vcat(di, -0.825)) #lightpink

#inputSet behind project plane
inbe = HPolytope(vcat(Cb, eigenvector), vcat(di, 0.825)) #lightpink3

#inter proj input
d1 = vcat(di, 0.825, -0.825)
inter = HPolytope(C, d1)

#q1
Cq = vcat(Array(Diagonal(ones(2))), Array(-Diagonal(ones(2))))
dq1 = [1, 6, 0, -5]

#q2
dq2 = [1, 4, 0, -3]

#q3
dq3 = [1, 1, 0, 0]

#p1
Cp = Cq * W

p1if = HPolytope(vcat(Cb, Cp, -eigenvector), vcat(db, dq1, -0.825)) #cyan
p1be = HPolytope(vcat(Cb, Cp, eigenvector), vcat(db, dq1, 0.825)) #cyan3
p1inter = HPolytope(vcat(Cb, Cp, eigenvector, -eigenvector), vcat(db, dq1, 0.825, -0.825))

p2if = HPolytope(vcat(Cb, Cp, -eigenvector), vcat(db, dq2, -0.825))
p2be = HPolytope(vcat(Cb, Cp, eigenvector), vcat(db, dq2, 0.825))
p2inter = HPolytope(vcat(Cb, Cp, eigenvector, -eigenvector), vcat(db, dq2, 0.825, -0.825))
p2in = HPolytope(vcat(Cb, Cp), vcat(di, dq2))

p3if = HPolytope(vcat(Cb, Cp, -eigenvector), vcat(db, dq3, -0.825))
p3be = HPolytope(vcat(Cb, Cp, eigenvector), vcat(db, dq3, 0.825))
p3inter = HPolytope(vcat(Cb, Cp, eigenvector, -eigenvector), vcat(db, dq3, 0.825, -0.825))
p3inf = HPolytope(vcat(Cb, Cp, -eigenvector), vcat(di, dq3, -0.825))
p3inb = HPolytope(vcat(Cb, Cp, eigenvector), vcat(di, dq3, 0.825))

scene1 = plot3d(proj_plane, color=:dodgerblue)

scene1 = plot3d!(inbe, color=:lightpink3)
scene1 = plot3d!(infr, color=:lightpink)
scene1 = plot3d!(inter, color=:dodgerblue4)
scene1 = plot3d!(p1be, color=:cyan3)
scene1 = plot3d!(p1if, color=:cyan)
scene1 = plot3d!(p1inter, color=:red)

scene1 = plot3d(proj_plane, color=:dodgerblue)

scene1 = plot3d!(inbe, color=:lightpink3)
scene1 = plot3d!(infr, color=:lightpink)
scene1 = plot3d!(inter, color=:dodgerblue4)
scene1 = plot3d!(p2be, color=:cyan3)
scene1 = plot3d!(p2if, color=:cyan)
scene1 = plot3d!(p2inter, color=:red)
scene1 = plot3d!(p2in, color=:yellow)

#

scene1 = plot3d(proj_plane, color=:dodgerblue)

scene1 = plot3d!(inbe, color=:lightpink3)
scene1 = plot3d!(p3be, color=:cyan3)
scene1 = plot3d!(p3inb, color=:yellow3)

scene1 = plot3d!(infr, color=:lightpink)
scene1 = plot3d!(inter, color=:dodgerblue4)
scene1 = plot3d!(p3if, color=:cyan)
scene1 = plot3d!(p3inf, color=:yellow)
scene1 = plot3d!(p3inter, color=:red)

#save("plot.png", scene)


using IVPTestSuite  # imports ODE, DASSL and Sundials
const S = Solvers # for convenience

# choose test-case (all of them are in IVPTestSuite.tc_all)
tc_name = [:vdpol
           :rober
           :threebody
           :bruss1d
           :chemakzo
           :hires][3]
tc = IVPTestSuite.tc_all[tc_name]

# Pick a solver.  S.allsolvers is the list of all solvers.  Generally
# each family of solvers also has a list.
solver = S.allsolvers[ODE.ode45_dp]

###
# make a TestRun which combines a TestCase with a Solver + some extras:
abstols = 10.0.^(-5:-1:-10)
reltols = abstols
dt0 = NaN # size of first step
suite = TestSuite(tc, solver, abstols, reltols, [NaN])

# running it will display some results:
res = run_ode_testsuite(suite)


using LinearAlgebra

using DelimitedFiles
fdata = "C:\\Users\\jcai.AUTH\\Downloads\\SCODE-master\\SCODE-master\\data2\\data.txt" # ARGS[1]
ftime = "C:\\Users\\jcai.AUTH\\Downloads\\SCODE-master\\SCODE-master\\data2\\time.txt" # ARGS[2]
X=readdlm(fdata,'\t')
pseudotime=readdlm(ftime,'\t')[:,2]
pseudotime = pseudotime/maximum(pseudotime)

cnum=length(pseudotime)

#using CSV
#using DataFrames
#pseudotime = CSV.read(ftime, header=false)
#
#[1:cnum,2]
#pseudotime = pseudotime/maximum(pseudotime)

#println("a")

#dir = ARGS[3]
tfnum = 100   #parse(Int32,ARGS[4])
pnum = 4      #parse(Int32,ARGS[5])
#cnum = parse(Int32,ARGS[6])
maxite = 10

maxB = 2.0
minB = -10.0

#if !(isdir(dir))
#    mkdir(dir)
#end

#X = readtable(fdata, separator='\t', header=false)[1:tfnum,1:cnum]
#X = CSV.read(fdata, header=false)
#X = convert(Matrix, X)

W = zeros(tfnum, pnum)
Z = zeros(pnum, cnum)
WZ = zeros(tfnum, cnum)

#pseudotime = CSV.read(ftime, header=false)[1:cnum,2]
#pseudotime = pseudotime/maximum(pseudotime)

new_B = zeros(pnum)
old_B = zeros(pnum)

#initialization
RSS = Inf
for i in 1:pnum
    new_B[i] = rand()*(maxB-minB)+minB
    old_B[i] = new_B[i]
end

function sample_Z()
    for i in 1:pnum
        for j in 1:cnum
            Z[i,j] = exp(new_B[i]*pseudotime[j]) + (rand()*0.002-0.001)
        end
    end
end

#optimization
for ite in 1:maxite
    #sampling B
    target = rand(1:pnum)
    global new_B[target] = rand()*(maxB-minB)+minB

    #for last calc
    if ite==maxite
        global new_B = copy(old_B)
    end

    #sample Z from new B
    sample_Z()

    #regression
    for i in 1:tfnum
        W[i,:] = (Z * Z')^-1 * (Z * X[i,:])
    end
    #prediction
    WZ = W * Z

    #RSS
    tmp_RSS = 0
    for i in 1:tfnum
        tmp_RSS += sum((X[i,:]-WZ[i,:]).*(X[i,:]-WZ[i,:]))
    end

    if tmp_RSS < RSS
        global RSS = tmp_RSS
        global old_B[target] = new_B[target]
    else
        global new_B[target] = old_B[target]
    end
end

#output RSS
open(dir*"/RSS.txt", "w") do f
    println(f, RSS)
end

#output W
#writetable(dir*"/W.txt", convert(DataFrame, W), separator = '\t', header = false)
#DataFrame(W) |> CSV.write(dir*"/W.txt", delim='\t', writeheader=false)

#infer A
B = zeros(pnum, pnum)
for i in 1:pnum
    B[i,i] = new_B[i]
end
invW = pinv(W)
A = W * B * invW

#write A and B
#writetable(dir*"/A.txt", convert(DataFrame, A), separator = '\t', header = false)
#writetable(dir*"/B.txt", convert(DataFrame, B), separator = '\t', header = false)
#DataFrame(A) |> CSV.write(dir*"/A.txt", delim='\t', writeheader=false)
#DataFrame(B) |> CSV.write(dir*"/B.txt", delim='\t', writeheader=false)


using LinearAlgebra, GPUifyLoops, StaticArrays, Random, Test

@static if Base.find_package("CuArrays") !== nothing
    using CuArrays, CUDAnative
end

Random.seed!(4)

"""
Implementation of the 2 x 2 eignvalue and eigenvector solver found in:
  @misc{Borges2017,
    title = {Numerically sane solution of the 2 x 2 real symmetric eigenvalue problem},
    author = {Carlos F. Borges},
    year = {2017},
    url = {https://calhoun.nps.edu/handle/10945/51506},
  }
"""
@inline function cfbeigen(A::SHermitianCompact{2,T,3}) where {T <: Real}
    @inbounds α₁ = A[1, 1]
    @inbounds α₂ = A[2, 2]
    @inbounds β = A[2, 1]

    δ = (α₁ - α₂)/2
    h = hypot(δ, β)

    if h == 0
        c = one(T)
        s = zero(T)
    elseif δ ≥ 0
        c = δ + h
        s = β
    else
        c = β
        s = h - δ
    end

    ρ = 1/hypot(s, c)

    c = ρ*c
    s = ρ*s

    V = @SMatrix [c -s;
                  s  c]
    λ = @SVector [c^2*α₁ + s*(s*α₂ + 2c*β),
                  c^2*α₂ + s*(s*α₁ - 2c*β)]

    return Eigen(λ, V)
end

# initialize with random symmetric matrices
N = 2^25
T = Float64

As = (rand(T, N), rand(T, N), rand(T, N))
Vs = ntuple(_->similar(As[1]), 4)
λs = ntuple(_->similar(As[1]), 2)

function kernel(As, Vs, λs)
    @inbounds @loop for i in (1:size(As[1],1);
                              (blockIdx().x-1)*blockDim().x + threadIdx().x)
        A = SHermitianCompact(@SVector([As[1][i], As[2][i], As[3][i]]))

        F = cfbeigen(A)

        @unroll for j = 1:length(F.vectors)
            Vs[j][i] = F.vectors[j]
        end

        @unroll for j = 1:length(F.values)
            λs[j][i] = F.values[j]
        end
    end
    nothing
end

@launch CPU() kernel(As, Vs, λs)

@static if Base.find_package("CuArrays") !== nothing
    cAs = CuArray.(As)
    cVs = ntuple(_->similar(cAs[1]), length(Vs))
    cλs = ntuple(_->similar(cAs[1]), length(λs))

    threads = 1024
    blocks = cld(size(cAs[1],1), threads)
    @launch(CUDA(), threads=threads, blocks=blocks, kernel(cAs, cVs, cλs))

    for j = 1:length(As)
      @test As[j] ≈ Array(cAs[j])
    end

    for j = 1:length(Vs)
      @test Vs[j] ≈ Array(cVs[j])
    end

    for j = 1:length(λs)
      @test λs[j] ≈ Array(cλs[j])
    end
end


using Flux, Flux.Tracker
using Distributions

struct VP
  b0
  b1

  VP() = new(param(randn(2)), param(randn(2)))
end

sigmoid(x) = 1.0 / (1.0 + exp(-x))

function loglike(y, x, params, N)
  b0 = params[:b0]
  b1 = params[:b1]
  prob = clamp.(sigmoid.(b0 .+ b1 .* x), .001, .999)
  # return sum(y .* log.(prob) .+ (1 .- y) .* log1p.(-prob)) * (N / length(y))
  return sum(logpdf.(Bernoulli.(prob), y)) * (N / length(y))
end

function log_p(params)
  # return 0.0
  b0 = params[:b0]
  b1 = params[:b1]
  return logpdf(Normal(0, 10), b0) + logpdf(Normal(0, 10), b1)
end

lpdf_vp(p, x) = logpdf(Normal(p[1], exp(p[2])), x)
function log_q(params, vp)
  return lpdf_vp(vp.b0, params[:b0]) + lpdf_vp(vp.b1, params[:b1])
end

rsample(p) = randn() * exp(p[2]) + p[1]
rsample(vp::VP) = Dict(:b0 => rsample(vp.b0), :b1 => rsample(vp.b1))

struct Metrics
  ll::Float64
  lp::Float64
  lq::Float64
end

function elbo(y, x, vp, N, metrics::Vector{Metrics})
  params = rsample(vp)

  ll = loglike(y, x, params, N)
  lp = log_p(params)
  lq = log_q(params, vp)

  m = Metrics(ll.data, lp.data, lq.data)
  push!(metrics, m)

  return ll + lp - lq
end


using CompScienceMeshes, BEAST

h = 2π / 51; Γ = meshcircle(1.0, h)
X, Y = lagrangecxd0(Γ), lagrangec0d1(Γ)

κ = 1.0
S, Dᵀ = SingleLayer(κ), DoubleLayerTransposed(κ)
D, N  = DoubleLayer(κ), HyperSingular(κ)
I  = Identity()

# TM scattering
E = PlaneWaveDirichlet(κ, point(1.0,0.0))
H = PlaneWaveNeumann(  κ, point(1.0,0.0))

@hilbertspace u; @hilbertspace u′
@hilbertspace t; @hilbertspace t′

tm_efie = @discretise           S[u′,u] == E[u′] u∈X u′∈X
tm_mfie = @discretise (0.5I + Dᵀ)[t′,u] == H[t′] u∈X t′∈X

u1, u2 = solve(tm_efie), solve(tm_mfie)

# TE scattering
E = PlaneWaveNeumann(  κ, point(1.0, 0.0))
H = PlaneWaveDirichlet(κ, point(1.0, 0.0))

te_efie = @discretise          N[t,t′] == E[t′] t∈Y t′∈Y
te_mfie = @discretise (0.5I - D)[t,u′] == H[u′] t∈Y u′∈Y

t1, t2 = solve(te_efie), solve(te_mfie)

using Plots
plotlyjs()

nx = numfunctions(X);
Δα = 2π/nx; α = (collect(1:nx) - 0.5) * Δα
plt1 = Plots.plot(α, real(u1), c=:blue, label="TM-EFIE")
Plots.scatter!(α, real(u2), c=:red, m=:circle, label="TM-MFIE")
Plots.title!("current vs. angle")

ny = numfunctions(Y);
Δα = 2π/ny; α = collect(0:ny-1) * Δα
plt2 = Plots.plot(α, real(t1), c=:blue, label="TE-EFIE")
Plots.scatter!(α, real(t2), c=:red, m=:circle, label="TE-MFIE")
Plots.title!("current vs. angle")

Plots.plot(plt1,plt2)


using Distributed

# Number of ITensorNetworkMaps in the sum
nmats = 2

addprocs(nmats)

# Set up for the calculation. Define everything
# in the `@everywhere` block that you need
# on all of the processes to perform the
# calculation.
@everywhere begin
  using ITensors
  using ITensors.ITensorNetworkMaps
  using LinearAlgebra
  using KrylovKit

  BLAS.set_num_threads(1)
  ITensors.Strided.disable_threads()
  ITensors.disable_threaded_blocksparse()

  include(joinpath("src", "DistributedSums.jl"))

  # Function that defines the tensor network
  # that will map the input vector (i.e.
  # a transfer matrix or effective Hamiltonian).
  function itensor_map(i, j)
    A1 = randomITensor(i', dag(i), dag(l))
    A2 = randomITensor(j', dag(j), l)
    return ITensorNetworkMap([A1, A2])
  end

  d = 2
  χ = 3
  i = Index(d, "i")
  j = Index(d, "j")
  l = Index(χ, "l")

  v0 = randomITensor(i, j)
end

# Defines an object representing a sum of
# linear maps (generalized matrices) built
# from contractions of ITensors that are
# distributed across `nmats` workers
# (each linear map lives on a worker).
A = DistributedSum(_ -> itensor_map(i, j), nmats)

# Apply the linear map, which internally
# is performed in parallel where the linear
# map is created.
Av0 = A(v0)

@show norm(Av0 - sum(fetch.(A))(v0))

# Solve an eigenvector equation in
# parallel.
λ⃗, v⃗ = eigsolve(A, v0, 1, :LM)
λ = λ⃗[1]
v = v⃗[1]

@show norm(A(v) - λ * v)


@diffusion_process FitzHughDiffusion begin
    :dimensions
    process --> 2
    wiener --> 1

    :parameters
    (ϵ, s, γ, β, σ) --> Float64

    :additional
    constdiff --> true
end

function b(t, x, P::FitzHughNagumo)
    ℝ{2}(
        (x[1] - x[2] - x[1]^3 + P.s)/P.ϵ,
        P.γ*x[1] - x[2] + P.β
    )
end

σ(t, x, P::FitzHughNagumo) = ℝ{2}(0.0, P.σ)


using Revise # reduce need for recompile
using Plots
using LinearAlgebra

push!(LOAD_PATH, "./src") # user defined modules
using CommonUtils, Basis1D

"Approximation parameters"
N   = 3 # The order of approximation
K1D = 16
CFL = 2
T   = 10 # endtime

"Mesh related variables"
VX = LinRange(-1,1,K1D+1)
EToV = transpose(reshape(sort([1:K1D; 2:K1D+1]),2,K1D))

"Construct matrices on reference elements"
r,w = gauss_lobatto_quad(0,0,N)
V = vandermonde_1D(N, r)
Dr = grad_vandermonde_1D(N, r)/V
M = inv(V*V')

"Nodes on faces, and face node coordinate"
wf = [1;1]
Vf = vandermonde_1D(N,[-1;1])/V
LIFT = M\(transpose(Vf)*diagm(wf)) # lift matrix

"Construct global coordinates"
V1 = vandermonde_1D(1,r)/vandermonde_1D(1,[-1;1])
x = V1*VX[transpose(EToV)]

"Connectivity maps"
xf = Vf*x
mapM = reshape(1:2*K1D,2,K1D)
mapP = copy(mapM)
mapP[1,2:end] .= mapM[2,1:end-1]
mapP[2,1:end-1] .= mapM[1,2:end]

"Make periodic"
mapP[1] = mapM[end]
mapP[end] = mapM[1]

"Geometric factors and surface normals"
J = repeat(transpose(diff(VX)/2),N+1,1)
nxJ = repeat([-1;1],1,K1D)
rxJ = 1

"initial conditions"
u = @. exp(-25*x^2)

"Time integration"
rk4a,rk4b,rk4c = rk45_coeffs()
CN = (N+1)*(N+2)/2  # estimated trace constant
dt = CFL * 2 / (CN*K1D)
Nsteps = convert(Int,ceil(T/dt))
dt = T/Nsteps

"pack arguments into tuples - will "
ops = (Dr,LIFT,Vf)
vgeo = (rxJ,J)
fgeo = (nxJ)

function rhs(u,ops,vgeo,fgeo,mapP)
    # unpack args
    Dr,LIFT,Vf = ops
    rxJ,J = vgeo
    nxJ = fgeo

    uM = Vf*u # can replace with nodal extraction
    du = uM[mapP]-uM

    ux = rxJ*(Dr*u)
    tau = 1 # upwind penalty parameter
    rhsu = ux + .5*LIFT*(@. du*nxJ - tau*abs(nxJ)*du)

    return -rhsu./J
end

"plotting nodes"
Vp = vandermonde_1D(N,LinRange(-1,1,10))/V
gr(size=(300,300),legend=false,markerstrokewidth=1,markersize=2)
plt = plot(Vp*x,Vp*u)

resu = zeros(size(x))
@gif for i = 1:Nsteps
    for INTRK = 1:5
        rhsu = rhs(u,ops,vgeo,fgeo,mapP)
        @. resu = rk4a[INTRK]*resu + dt*rhsu
        @. u   += rk4b[INTRK]*resu
    end

    if i%10==0 || i==Nsteps
        println("Number of time steps $i out of $Nsteps")
        # display(plot(Vp*x,Vp*u,ylims=(-.1,1.1)))
        # push!(plt, Vp*x,Vp*u,ylims=(-.1,1.1))
        plot(Vp*x,Vp*u,ylims=(-.1,1.1),title="Timestep $i out of $Nsteps",lw=2)
        scatter!(x,u)
        # sleep(.0)
    end
end every 5

# plot(Vp*x,Vp*u,ylims=(-.1,1.1))


module TestDiagrams
using Test

using Catlab.Theories, Catlab.Graphs, Catlab.CategoricalAlgebra
using Catlab.Graphs.BasicGraphs: TheoryGraph, TheorySymmetricGraph

const SchSGraph = TheorySymmetricGraph

# Diagrams
##########

# Diagram for paths of length 2.
C = FinCat(@acset Graph begin
  V = 3
  E = 2
  src = [1,2]
  tgt = [3,3]
end)
D = FinDomFunctor([:E,:E,:V], [:tgt,:src], C, FinCat(SchSGraph))
d = Diagram(D)
@test shape(d) == C
@test ob_map(d, 3) == SchSGraph[:V]
@test hom_map(d, 1) == SchSGraph[:tgt]
@test startswith(sprint(show, d), "Diagram{id}(")

# Diagram morphisms
###################

f = DiagramHom([(2,:inv), (1,:inv), 3], [2,1], d, d)
@test dom(f) == d
@test codom(f) == d
@test is_functorial(shape_map(f))
@test shape_map(f) == FinFunctor([2,1,3], [2,1], C, C)
ϕ = diagram_map(f)
@test is_natural(ϕ, check_equations=false)
@test ϕ[1] == SchSGraph[:inv]
@test ϕ[3] == id(SchSGraph[:V])
@test ob_map(f, 2) == (1, SchSGraph[:inv])
@test hom_map(f, 2) == Path(graph(C), 1)
@test collect_ob(f) == [(2, ϕ[1]), (1, ϕ[2]), (3, ϕ[3])]
@test collect_hom(f) == [Path(graph(C), 2), Path(graph(C), 1)]
f² = f⋅f
@test shape_map(f²) == FinFunctor(1:3, 1:2, C, C)

f = DiagramHom{op}([(2,:inv), (1,:inv), 3], [2,1], D, D)
ιV = FinDomFunctor([:V], FinCat(1), FinCat(SchSGraph))
g = DiagramHom{op}([(1,:src)], D, ιV)
@test dom(g) == Diagram{op}(D)
@test codom(g) == Diagram{op}(ιV)
@test compose(id(dom(g)), g) == g
@test compose(g, id(codom(g))) == g
@test is_natural(diagram_map(g))
@test startswith(sprint(show, f), "DiagramHom{op}(")

fg = f⋅g
@test ob_map(fg, 1) == (2, SchSGraph[:inv]⋅SchSGraph[:src])

d = dom(f)
@test op(op(d)) == d
@test op(op(f)) == f
@test dom(op(g)) == Diagram{co}(ιV)
@test codom(op(g)) == Diagram{co}(D)
@test op(g) == DiagramHom{co}([(1,:src)], ιV, D)
@test op(g)⋅op(f) == op(f⋅g)

# Monads of diagrams
####################

C = FinCat(TheoryGraph)
d = munit(Diagram{id}, C, :V)
@test is_discrete(shape(d))
@test only(collect_ob(diagram(d))) == TheoryGraph[:V]
f = munit(DiagramHom{id}, C, :src)
@test only(components(diagram_map(f))) == TheoryGraph[:src]

end


function part1(input)
    program = parse_program(input)
    accumulator, halts = run_program(program)
    @assert !halts
    return accumulator
end

function part2(input)
    program = parse_program(input)
    for i = 1:length(program)
        opcode = program[i].opcode
        if opcode == :nop
            program[i].opcode = :jmp
        elseif opcode == :jmp
            program[i].opcode = :nop
        else
            continue
        end
        accumulator, halts = run_program(program)
        halts && return accumulator
        program[i].opcode = opcode
    end
end

mutable struct Instruction
    opcode::Symbol
    value::Int
end

function parse_program(input)
    program = Instruction[]
    for line in readlines(input)
        name, value = split(line, " ")
        push!(program, Instruction(Symbol(name), parse(Int, value)))
    end
    return program
end

function run_program(program)
    accumulator = 0
    program_counter = 1
    visited_instructions = Set{Int}()
    while true
        program_counter == length(program) + 1 && return accumulator, true
        program_counter in visited_instructions && return accumulator, false
        push!(visited_instructions, program_counter)
        instruction = program[program_counter]
        if instruction.opcode == :acc
            accumulator += instruction.value
            program_counter += 1
        elseif instruction.opcode == :jmp
            program_counter += instruction.value
        elseif instruction.opcode == :nop
            program_counter += 1
        end
    end
end