# This file is a part of EuclidianNormalizingFlows.jl, licensed under the MIT License (MIT).

using EuclidianNormalizingFlows
using Test

using LinearAlgebra
using InverseFunctions, ChangesOfVariables
using ForwardDiff, Zygote

using EuclidianNormalizingFlows: householder_trafo, chained_householder_trafo, HouseholderTrafo

@testset "householder_trafo" begin
    v = rand(5)
    x = rand(5)
    X = rand(5, 3)

    householder_matrix(v::AbstractVector{<:Real}) = I - 2 * (v*v') / (v'*v)

    @test @inferred(householder_trafo(v, x)) ≈ householder_matrix(v) * x
    @test @inferred(householder_trafo(v, householder_trafo(v, x))) ≈ x

    @test @inferred(householder_trafo(v, X)) ≈ hcat(householder_trafo.(Ref(v), eachcol(X))...)
    @test @inferred(householder_trafo(v, X)) ≈ householder_matrix(v) * X
    @test @inferred(householder_trafo(v, householder_trafo(v, X))) ≈ X

    @test Zygote.pullback(householder_trafo, v, x)[1] == householder_trafo(v, x)
    @test Zygote.jacobian(householder_trafo, v, x)[1] ≈ ForwardDiff.jacobian(v -> householder_matrix(v) * x, v)
    @test Zygote.jacobian(householder_trafo, v, x)[2] ≈ ForwardDiff.jacobian(x -> householder_matrix(v) * x, x)

    @test Zygote.pullback(householder_trafo, v, X)[1] == householder_trafo(v, X)
    @test Zygote.jacobian(householder_trafo, v, X)[1] ≈ ForwardDiff.jacobian(v -> householder_matrix(v) * X, v)
    @test Zygote.jacobian(householder_trafo, v, X)[2] ≈ ForwardDiff.jacobian(X -> householder_matrix(v) * X, X)

    V = rand(5, 3)

    chained_householder_matrix(v::AbstractMatrix{<:Real}) = *(reverse(householder_matrix.(eachcol(V)))...)

    @test chained_householder_trafo(V, x) ≈ *(reverse(householder_matrix.(eachcol(V)))...) * x
    @test chained_householder_trafo(reverse(V, dims = 2), chained_householder_trafo(V, x)) ≈ x

    @test @inferred(chained_householder_trafo(V, X)) ≈ hcat(chained_householder_trafo.(Ref(V), eachcol(X))...)

    # # Yields incorrect results:
    # ForwardDiff.jacobian(V -> chained_householder_matrix(V) * x, V)
    # Zygote.jacobian(V -> chained_householder_matrix(V) * x, V)

    @test Zygote.pullback(chained_householder_trafo, V, x)[1] == chained_householder_trafo(V, x)
    @test Zygote.jacobian(chained_householder_trafo, V, x)[1] ≈ ForwardDiff.jacobian(V -> chained_householder_trafo(V, x), V)
    @test Zygote.jacobian(chained_householder_trafo, V, x)[2] ≈ ForwardDiff.jacobian(x -> chained_householder_trafo(V, x), x)

    @test Zygote.pullback(chained_householder_trafo, V, X)[1] == chained_householder_trafo(V, X)
    @test Zygote.jacobian(chained_householder_trafo, V, X)[1] ≈ ForwardDiff.jacobian(V -> chained_householder_trafo(V, X), V)
    @test Zygote.jacobian(chained_householder_trafo, V, X)[2] ≈ ForwardDiff.jacobian(X -> chained_householder_trafo(V, X), X)

    @test HouseholderTrafo(V) == HouseholderTrafo(V)
    @test isequal(HouseholderTrafo(V), HouseholderTrafo(V))
    @test hash(HouseholderTrafo(V)) == hash(HouseholderTrafo(V))

    for arg in (x, X)
        InverseFunctions.test_inverse(HouseholderTrafo(V), arg)
    end

    @test @inferred(with_logabsdet_jacobian(HouseholderTrafo(V), x)) == (chained_householder_trafo(V, x), 0)
    @test @inferred(with_logabsdet_jacobian(HouseholderTrafo(V), X)) == (chained_householder_trafo(V, X), fill(0, size(X, 2))')
end


include("C:\\Users\\darien.shannon\\Documents\\Code\\Julia\\JuliaSE\\src\\2DBeam.jl")

COORD = [Node(0,0)
         Node(120, 0)
         Node(360, 0)
         Node(480, 0)]

MSUP = [Support(1, 1, 0)
        Support(2, 1, 0)
        Support(3, 1, 0)
        Support(4, 1, 1)]

EM = Dict( 1 => Material(29e3))

CP = Dict( 1 => Section(350),
           2 => Section(500))

MPRP = [Element(1, 2, 1, 1)
        Element(2, 3, 1, 1)
        Element(3, 4, 1, 2)]

memberLoads = [(2, DLoad(0.1667, 0.1667, 0, 120))
               (2, CLoad(25, 180))
               (3, DLoad(0.25, 0, 0, 120))]

nodeLoads = [(1, NLoad((0, -480)))]

forces, reactions = run_analysis(COORD, MSUP, EM, CP, MPRP, memberLoads, nodeLoads)
println(forces)
println(reactions)


abstract type
    AbstractIDSequence{S<:ValueSupport} <: MultivariateDistribution{S}
end

length(X::AbstractIDSequence) = length(distributions(X))

function insupport(X::AbstractIDSequence,x::AbstractVector)
  length(x)==length(X) && all(dx->insupport(dx[1],dx[2]),zip(distributions(X),x))
end

function _rand!{T<:Real}(X::AbstractIDSequence, x::AbstractVector{T})
  for (i,d_i) in enumerate(distributions(X))
    x[i] = rand(d_i)
  end
  return x
end

function _logpdf{T<:Real}(X::AbstractIDSequence, x::AbstractVector{T})
  return sum(arg->logpdf(arg[1],arg[2]),zip(distributions(X),x))
end

mean(X::AbstractIDSequence) = Float64[mean(d_i) for d_i in distributions(X)]
var(X::AbstractIDSequence) = Float64[var(d_i) for d_i in distributions(X)]
cov(X::AbstractIDSequence) = Diagonal(var(X))
entropy(X::AbstractIDSequence) = sum(entropy,distributions(X))

function mgf(X::AbstractIDSequence, t::AbstractVector)
  prod([mgf(d_i,t_i) for (d_i,t_i) in zip(distributions(X),t)])
end
function cf(X::AbstractIDSequence, t::AbstractVector)
  prod([cf(d_i,t_i) for (d_i,t_i) in zip(distributions(X),t)])
end


function findenaechstesBild(Aufloesung,prozent)

Phantom = echteCTDaten(Aufloesung);
nrows,ncols = size(Phantom);
f_nahBild = zeros(size(Phantom,1),size(Phantom,2));
PhantomVec = reshape(Phantom,(Int(nrows)*Int(ncols),1));

x0,y0,N,JJ,nx,ny,x1,y1,verschx,verschy,a,b,A,theta1,counter = Dateneinlesen(Aufloesung);
A2 = erstelleA(x0,x1,y0,y1,N,JJ,nx,ny,verschx,verschy);
Messwerte =  A2*PhantomVec;

MesswerteMitFehler = erzeugeFehlerbehafteteDaten(Messwerte,prozent);

f_nah = (transpose(A2)*A2)\(transpose(A2)*MesswerteMitFehler)
for i = 1:size(f_nah,1)
    if f_nah[i] < 0
        f_nah[i] = 0;
    elseif f_nah[i] > 1
        f_nah[i] = 1;
    end
end
ObjektNaiv = reshape(f_nah,(nrows,ncols));
    #print(ObjektNaiv)
f_nahBild = convert(Array{Gray{Normed{UInt8,8}},2},ObjektNaiv);

return PhantomVec,f_nah,Messwerte,MesswerteMitFehler,A2,f_nahBild;
end


export chebyshev, chebyshev!

chebyshev(A, b, λmin::Real, λmax::Real, Pr = 1, n = size(A,2);
          tol::Real = sqrt(eps(typeof(real(b[1])))), maxiter::Int = n^3) =
    chebyshev!(zerox(A, b), A, b, λmin, λmax, Pr, n; tol=tol,maxiter=maxiter)

function chebyshev!(x, A, b, λmin::Real, λmax::Real, Pr = 1, n = size(A,2);
	tol::Real = sqrt(eps(typeof(real(b[1])))), maxiter::Int = n^3
    )
    history = ConvergenceHistory()
    history[:tol] = tol
    reserve!(history,:resnorm,maxiter)

	K = KrylovSubspace(A, n, 1, Adivtype(A, b))
	init!(K, x)
	chebyshev_method!(history, x, K, b, λmin, λmax, Pr; tol = tol, maxiter = maxiter)
    x, history
end

function chebyshev_method!(
    log::ConvergenceHistory, x, K::KrylovSubspace, b, λmin::Real, λmax::Real,
    Pr = 1; tol::Real = sqrt(eps(typeof(real(b[1])))), maxiter::Int = K.n^3
    )

	local α, p
	K.order = 1
    tol = tol*norm(b)
	r = b - nextvec(K)
	d::eltype(b) = (λmax + λmin)/2
	c::eltype(b) = (λmax - λmin)/2
	for iter = 1:maxiter
        nextiter!(log, mvps=1)
		z = Pr\r
		if iter == 1
			p = z
			α = 2/d
		else
			β = (c*α/2)^2
			α = 1/(d - β)
			p = z + β*p
		end
		append!(K, p)
		update!(x, α, p)
		r -= α*nextvec(K)
		#Check convergence
		resnorm = norm(r)
        push!(log, :resnorm, resnorm)
        resnorm < tol && break
	end
    shrink!(log)
    setconv(log, 0<=norm(r)<tol)
	x
end


# A SearchSpace determines which candidate points can be
# considered in a search/optimization. This base class has very few restrictions
# and can allow varying number of dimensions etc.
abstract SearchSpace

# SearchSpace with a fixed number of dimensions. The vast majority of problems.
abstract FixedDimensionSearchSpace <: SearchSpace

# In a ContinuousSearchSpace each dimension has a continuous range of
# valid values.
abstract ContinuousSearchSpace <: FixedDimensionSearchSpace

ranges(css::ContinuousSearchSpace) = [range_for_dim(css, i) for i in 1:numdims(css)]

# Access individual range for a dim.
range_for_dim(css::ContinuousSearchSpace, i::Int) = (mins(css)[i], maxs(css)[i])

# Generate a number of individuals by random sampling in the search space.
function rand_individuals(css::ContinuousSearchSpace, numIndividuals)
  # Basically min + delta * rand(), but broadcast over the columns...
  broadcast(+, mins(css), broadcast(*, deltas(css), rand(numdims(css), numIndividuals)))
end

# Generate a number of individuals via latin hypercube sampling. This should
# be the default when creating the initial population.
function rand_individuals_lhs(css::ContinuousSearchSpace, numIndividuals)
  BlackBoxOptim.Utils.latin_hypercube_sampling(mins(css), maxs(css), numIndividuals)
end

# Generate one random candidate.
function rand_individual(css::ContinuousSearchSpace)
  rand_individuals(css, 1)[:,1:1]
end

# True iff ind is within the search space.
function isinspace(ind, css::ContinuousSearchSpace)
  all(mins(css) .<= ind .<= maxs(css))
end

# In a RangePerDimSearchSpace each dimension is specified as a range of valid
# values.
type RangePerDimSearchSpace <: ContinuousSearchSpace
  # We save the ranges as individual mins, maxs and deltas for faster access later.
  mins::Array{Float64,2}
  maxs::Array{Float64,2}
  deltas::Array{Float64,2}

  function RangePerDimSearchSpace(ranges)
    mins = map(t -> t[1], ranges)''
    maxs = map(t -> t[2], ranges)''
    new(mins, maxs, (maxs - mins))
  end
end
mins(rss::RangePerDimSearchSpace) = rss.mins
maxs(rss::RangePerDimSearchSpace) = rss.maxs
deltas(rss::RangePerDimSearchSpace) = rss.deltas
numdims(rss::RangePerDimSearchSpace) = size(mins(rss), 1)

diameters(rss::RangePerDimSearchSpace) = deltas(rss)

# Convenience function to create symmetric search spaces.
symmetric_search_space(numdims, range = (0.0, 1.0)) = RangePerDimSearchSpace([range for i in 1:numdims])

# Create a feasible point (i.e. within the search space) given one which is
# outside.
function feasible(v, ss::RangePerDimSearchSpace)
  minimum(hcat(maxs(ss), maximum(hcat(mins(ss), v), 2)), 2)
end


n_scen = 10_000       # number of Monte Carlo scenarios
costs_fixed = 20      # fixed costs
s = 0.03              # risk free rate of interest
α = 0.99              # confidence level of economic capital
seed = 2              # random seed for repeatable calculations

insurance_input =
  DataFrame(
  id =          [:fire, :liab, :theft],
  name =        ["Fire", "Liability", "Theft"],
  premium =     [600.,  300.,  100.],
  loss_ratio =  [0.75,  0.75,  0.75],
  cost_ratio =  [0.05,  0.05,  0.05],
  var_coeff =   [0.50,  0.60,  0.70],
  re_ceded =    [0.25,  0.20,  0.20],
  re_costs =    [-0.06, -0.06, -0.06])
## add a column with an explicit counter: the lines of business
insurance_input[:ctr] = collect(1:nrow(insurance_input))

invest_input =
  DataFrame(
  id =    [:invest],
  name = ["Investment"],
  init = [1400.],
  cost_ratio = [0.005],
  mean = [0.05],
  sd =[0.02])
## add a column with an explicit counter: investments
invest_input[:ctr] = [nrow(insurance_input)+1]

tau_kendall = Real[1.0  0.3  0.2  0.0;
                   0.3  1.0  0.6  0.0;
                   0.2  0.6  1.0  0.0;
                   0.0  0.0  0.0  1.0]

hurdle = 0.1          # hurdle rate for EVA calculation
n_cloud = 100         # number of points in cloud diagrams


@inline function upperbound(x::ArbComplex{P}, ::Type{ArbFloat}) where {P}
    xreal = real(x)
    ximag = imag(x)
    wreal = ArbFloat{P}()
    wimag = ArbFloat{P}()
    ccall(@libarb(arb_get_ubound_arf), Cvoid, (Ref{ArbFloat}, Ref{ArbReal}, Clong), wreal, xreal, P)
    ccall(@libarb(arb_get_ubound_arf), Cvoid, (Ref{ArbFloat}, Ref{ArbReal}, Clong), wimag, ximag, P)
    return ArbComplex{P}(wreal, wimag)
end

@inline function lowerbound(x::ArbComplex{P}, ::Type{ArbFloat}) where {P}
    xreal = real(x)
    ximag = imag(x)
    wreal = ArbFloat{P}()
    wimag = ArbFloat{P}()
    ccall(@libarb(arb_get_lbound_arf), Cvoid, (Ref{ArbFloat}, Ref{ArbReal}, Clong), wreal, xreal, P)
    ccall(@libarb(arb_get_lbound_arf), Cvoid, (Ref{ArbFloat}, Ref{ArbReal}, Clong), wimag, ximag, P)
    return ArbComplex{P}(wreal, wimag)
end

upperbound(x::ArbComplex{P}) where {P} = upperbound(x, ArbFloat)
lowerbound(x::ArbComplex{P}) where {P} = lowerbound(x, ArbFloat)

function interval(x::ArbComplex{P}, ::Type{ArbFloat}) where {P}
    lowerbound(x), upperbound(x)
end

interval(x::ArbComplex{P}) where {P} = interval(x, ArbFloat)

@inline function upperbound_abs(x::ArbComplex{P}, ::Type{ArbFloat}) where {P}
    xreal = real(x)
    ximag = imag(x)
    wreal = ArbFloat{P}()
    wimag = ArbFloat{P}()
    ccall(@libarb(arb_get_abs_ubound_arf), Cvoid, (Ref{ArbFloat}, Ref{ArbReal}, Clong), wreal, xreal, P)
    ccall(@libarb(arb_get_abs_ubound_arf), Cvoid, (Ref{ArbFloat}, Ref{ArbReal}, Clong), wimag, ximag, P)
    return ArbComplex{P}(wreal, wimag)
end

@inline function lowerbound_abs(x::ArbComplex{P}, ::Type{ArbFloat}) where {P}
    xreal = real(x)
    ximag = imag(x)
    wreal = ArbFloat{P}()
    wimag = ArbFloat{P}()
    ccall(@libarb(arb_get_abs_lbound_arf), Cvoid, (Ref{ArbFloat}, Ref{ArbReal}, Clong), wreal, xreal, P)
    ccall(@libarb(arb_get_abs_lbound_arf), Cvoid, (Ref{ArbFloat}, Ref{ArbReal}, Clong), wimag, ximag, P)
    return ArbComplex{P}(wreal, wimag)
end

upperbound_abs(x::ArbComplex{P}) where {P} = upperbound_abs(x, ArbFloat)
lowerbound_abs(x::ArbComplex{P}) where {P} = lowerbound_abs(x, ArbFloat)

function interval_abs(x::ArbComplex{P}, ::Type{ArbFloat}) where {P}
    lowerbound_abs(x), upperbound_abs(x)
end

interval_abs(x::ArbComplex{P}) where {P} = interval_abs(x, ArbFloat)

function ball(x::ArbComplex{P}) where {P}
    mid = midpoint(x)
    rad = radius(x)
    return mid, rad
end


## Used for solving explicit sindy problems

function normalize_theta!(scales::AbstractVector, theta::AbstractMatrix)
    map(1:length(scales)) do i
        scales[i] = norm(theta[i,:], 2)
        theta[i, :] .= theta[i,:]./scales[i]
    end
    return
end

function rescale_xi!(xi::AbstractMatrix, scales::AbstractVector, round_::Bool)
    digs = 10
    @inbounds for i in 1:length(scales), j in 1:size(xi, 2)
        iszero(xi[i,j]) ? continue : nothing
        round_ && (xi[i,j] % 1) != zero(xi[i,j]) ? digs = round(Int64,-log10(abs(xi[i,j]) % 1))+1 : nothing
        xi[i,j] = xi[i,j] / scales[i]
        round_ ? xi[i,j] = round(xi[i,j], digits = digs) : nothing
    end
    return
end

# Main
function DiffEqBase.solve(p::DataDrivenProblem{dType}, b::Basis, opt::Optimize.AbstractOptimizer;
    normalize::Bool = false, denoise::Bool = false, maxiter::Int = 0,
    round::Bool = true,
    eval_expression = false, kwargs...) where {dType <: Number}
    # Check the validity
    @assert is_valid(p) "The problem seems to be ill-defined. Please check the problem definition."

    # Evaluate the basis
    θ = b(DataDrivenDiffEq.get_oop_args(p)...)
    # Extract the target
    DX = DataDrivenDiffEq.get_target(p)

    maxiter = maxiter <= 0 ? maximum(size(θ)) : maxiter

    # Normalize via p norm
    scales = ones(dType, size(θ, 1))

    normalize ? normalize_theta!(scales, θ) : nothing

    # Denoise via optimal shrinkage
    denoise ? optimal_shrinkage!(θ') : nothing

    # Init the coefficient matrix
    Ξ = DataDrivenDiffEq.Optimize.init(opt, θ', DX')
    # Solve
    @views Optimize.sparse_regression!(Ξ, θ', DX', opt; kwargs...)

    normalize ? rescale_xi!(Ξ, scales, round) : nothing

    # Build solution Basis
    return build_solution(
        p, Ξ, opt, b, eval_expression = eval_expression
    )
end

function _isin(x::Num, y)
    vs = get_variables(y)
    any(isequal(x, yi) for yi in vs)
end

function _isin(x::Vector{Num}, y::Vector)
    [_isin(xi, yi) for xi in x, yi in y]
end

_isin(x, y) = _isin(Num.(Symbolics.unwrap(x)), y)

function _ind_matrix(x::Vector{Num}, y::Vector)
    isins = _isin(x, y)
    inds = ones(Bool, size(isins)) # We take all
    excludes = zeros(Bool, length(x))
    for i in 1:length(x)
        excludes .= true
        excludes[i] = false
        inds[i, :] .= inds[i, :] .* sum(eachrow(isins[excludes, :]))
    end
    return .~ inds
end

@views function DiffEqBase.solve(p::DataDrivenProblem{dType}, b::Basis,
    opt::Optimize.AbstractSubspaceOptimizer, implicits::Vector{Num} = Num[];
    normalize::Bool = false, denoise::Bool = false, maxiter::Int = 0,
    eval_expression = false,
    round::Bool = true, kwargs...) where {dType <: Number}
    # Check the validity
    @assert is_valid(p) "The problem seems to be ill-defined. Please check the problem definition."

    # Check for the variables
    @assert all(any.(eachrow(_isin(implicits, states(b)))))

    # Evaluate the basis
    θ = b(DataDrivenDiffEq.get_implicit_oop_args(p)...)
    # Extract the target
    DX = DataDrivenDiffEq.get_target(p)

    maxiter = maxiter <= 0 ? maximum(size(θ)) : maxiter

    # Normalize via p norm
    scales = ones(dType, size(θ, 1))

    normalize ? normalize_theta!(scales, θ) : nothing

    # Denoise via optimal shrinkage
    denoise ? optimal_shrinkage!(θ') : nothing

    # Init the coefficient matrix
    Ξ = DataDrivenDiffEq.Optimize.init(opt, θ', DX')
    # Find the implict variables in the equations and
    # eliminite duplictes
    if !isempty(implicits)
        inds = _ind_matrix(implicits, [eq.rhs for eq in equations(b)])
    else
        inds = ones(Bool, 1, length(b))
    end
    offset = 0
    # Solve for each implicit variable
    @views for i in 1:size(inds, 1)
        # Initial progress offset
        offset = i*length(get_threshold(opt))
        Optimize.sparse_regression!(Ξ[inds[i,:], i:i], θ[inds[i,:],:]', DX[i:i, :]', opt; maxiter = maxiter,
         progress_outer = size(inds, 1), progress_offset = offset, kwargs...)
    end

    normalize ? rescale_xi!(Ξ, scales, round) : nothing
    # Build solution Basis
    return build_solution(
        p, Ξ, opt, b, implicits, eval_expression = eval_expression
    )
end


function propagator(f)
  function (cells...)
    output = last(cells)
    inputs = cells[1:end-1]
    function activate!()
      inputvalues = map(getindex, inputs)
      if !any(isunusable, inputvalues)
        addcontent!(output, f(inputvalues...))
      end
    end
    return Propagator(activate!, inputs)
  end
end

function block(tobuild, inputs...)
  isbuilt = false
  function maybebuild()
    if !(isbuilt || all(isunusable, map(getindex, inputs)))
      tobuild()
      isbuilt = true
    end
  end
  return Propagator(maybebuild, inputs)
end

constant(x) = propagator(() -> x)

function conditional(p, t, f, out)
  Propagator([p, t, f]) do
    predicate = p[]
    predicate != nothing || return
    out[] = predicate ? t[] : f[]
  end
end

switch(p, t, out) = conditional(p, t, Cell(), out)

function constraint(::typeof(*))
  function (x, y, xy)
    propagator(*)(x, y, xy)
    propagator(/)(xy, x, y)
    propagator(/)(xy, y, x)
  end
end

function constraint(::typeof(+))
  function (x, y, xy)
    propagator(+)(x, y, xy)
    propagator(-)(xy, x, y)
    propagator(-)(xy, y, x)
  end
end

function quadratic(x, x²)
  propagator(x -> x^2)(x, x²)
  propagator(sqrt)(x², x)
end


using Test, Bridge

# test the tricky bridge sampling when endpoint is or is not in grid
#srand(10)
G =  GammaProcess(10.0, 1.5)
GB = GammaBridge(1.0, 2.0, G)
n = 1000
@test abs(mean(sample([0.0, 1.0, 3.0],  G).yy[end] for i in 1:1000) -  mean(Bridge.increment(3, G))) < 4*std(Bridge.increment(0.5, G))/sqrt(n)

X1 = sample([0.0, 0.5, 1.0], GB, 0.2)
X2 = sample([0.0, 0.5, 2.0], GB, 0.2)
X3 = sample([0.0, 0.5], GB, 0.2)

@test  X1.tt ≈ [0,1,2]/2
@test  X2.tt ≈ [0,1,4]/2
@test  X3.tt ≈ [0,1]/2

@test  X1.yy[1] == X1.yy[1] == X1.yy[1] == 0.2

@test  X1.yy[3] ≈ 2.0
@test  X2.yy[3] >= 2.0
@test  X3.yy[2] <= 2.0

γᵒ, γ = 0.75, 1.2
t = 10000.
α = 16.0
n = 10000
tt = range(0, stop=t, length=n)
P = GammaProcess(γ, α)
X = sample(tt, P)

Y = copy(X)
Bridge.uniform_thinning!(Y, P, γᵒ)

@test abs(γ*t/n/α - mean(diff(X.yy))) < 0.2/sqrt(n)
@test abs(γᵒ*t/n/α - mean(diff(Y.yy))) < 0.2/sqrt(n)

γ = 10.
α = 5.
P = GammaProcess(γ, α)
T = 5.0
tt = 0.0:T/n:T
m = 1000
sigma = (t, x = 0.0) -> exp(-cos(t))
L = SamplePath(tt, sample(tt, P).yy .- sample(tt, P).yy)
X = solve(EulerMaruyama(), 0.0, L, ((t,x)-> -0*x, sigma))


export rnnConnectGraph

function rnnConnectGraph(L,T)
  connectGraph=zeros(Int64,L*T,L*T)
  for i=1:L*T
    if i%L==1
      connectGraph[i,i+1]=1
    elseif i%L!=0
      connectGraph[i,i+1]=1
      if i+L<L*T
        connectGraph[i,i+L]=1
      end
    else connectGraph[i,i-L+1]=1
    end
  end
  connectGraph
end


@testset "OLS" begin
    T, N, K = 100, 2, 4
    X = randn(T, K)
    b = randn(K, N)
    Y = X * b
    m = OLS(Y, X)
    @test m.coef ≈ b
    @test modelmatrix(m) === X
    @test coef(m) === m.coef
    @test residuals(m) === m.resid
    @test sprint(show, m) == "OLS regression"
end

@testset "LeastSquaresLP" begin
    e = LeastSquaresLP()
    @test sprint(show, e) == "LeastSquaresLP"
    @test sprint(show, MIME("text/plain"), e) == "Ordinary least squares local projection"
end

@testset "LeastSquaresLPResult" begin
    er = LeastSquaresLPResult([OLS(rand(5), rand(5,2))])
    @test sprint(show, er) == "LeastSquaresLPResult"
end

@testset "LocalProjectionResult" begin
    B = cat(randn(5,2,1), randn(5,2,1); dims=3)
    V = cat(randn(10,10), randn(10,10); dims=3)
    r = LocalProjectionResult(B, V, [1,2], LeastSquaresLP(), nothing, HRVCE(), VarName[:y1, :y2], VarName[:x], VarName[:w1, :w2], Dict{VarName,Int}(:y1=>1, :y2=>2), Dict{VarName,Int}(:x=>1), Dict{VarName,Int}(:w1=>1,:w2=>2,:y1=>3,:y2=>4), 2, 0, nothing, nothing, nothing, nothing, nothing, nothing, false, true)
    @test coef(r, 1, :x) == B[1,1,1]
    @test coef(r, 2, :w1, yname=:y2, lag=1) == B[2,2,2]
    @test coef(r, 2, :w1, lag=2) == B[4,1,2]
    @test vcov(r, 1, :x) == V[1,1,1]
    @test vcov(r, 2, :x, yname1=:y2, xwname2=:w1, lag2=1) == V[2,4,2]
    @test vcov(r, 2, :x, yname1=:y2, xwname2=:w1, yname2=:y1, lag2=1) == V[2,3,2]
    @test vcov(r, 2, :w1, yname1=:y2, lag1=2, xwname2=:w1, yname2=:y1, lag2=1) == V[8,3,2]
    @test vcov(r, 2, :w2, yname1=:y2, lag1=2, xwname2=:w1, yname2=:y1) == V[10,7,2]

    @test sprint(show, r) == "LocalProjectionResult"
    @test sprint(show, MIME("text/plain"), r) == """
        LocalProjectionResult with 2 lags over 2 horizons:
        ──────────────────────────────────────────────────────────────────────────────
        Variable Specifications
        ──────────────────────────────────────────────────────────────────────────────
        Outcome variables:              y1 y2    Minimum horizon:                    0
        Regressor:                          x    Lagged controls:                w1 w2
        ──────────────────────────────────────────────────────────────────────────────"""
    r = LocalProjectionResult(ones(1,1,1), ones(1,1,1), [1], LeastSquaresLP(), nothing, HRVCE(), VarName[:y], VarName[], VarName[:w], Dict{VarName,Int}(:y=>1), Dict{VarName,Int}(), Dict{VarName,Int}(:w=>1), 1, 0, nothing, nothing, nothing, nothing, nothing, nothing, false, true)
    @test sprint(show, MIME("text/plain"), r) == """
        LocalProjectionResult with 1 lag over 1 horizon:
        ──────────────────────────────────────────────────────────────────────────────
        Variable Specifications
        ──────────────────────────────────────────────────────────────────────────────
        Outcome variable:                   y    Minimum horizon:                    0
        Regressor:                               Lagged control:                     w
        ──────────────────────────────────────────────────────────────────────────────"""
    r = LocalProjectionResult(ones(1,1,1), ones(1,1,1), [1], LeastSquaresLP(), nothing, HRVCE(), VarName[Cum(:y)], VarName[Cum(:x)], VarName[:w], Dict{VarName,Int}(Cum(:y)=>1), Dict{VarName,Int}(Cum(:x)=>1), Dict{VarName,Int}(:w=>1), 1, 0, nothing, nothing, nothing, nothing, VarName[Cum(:x)], VarName[:z1,:z2], false, true)
    @test sprint(show, MIME("text/plain"), r) == """
        LocalProjectionResult with 1 lag over 1 horizon:
        ──────────────────────────────────────────────────────────────────────────────
        Variable Specifications
        ──────────────────────────────────────────────────────────────────────────────
        Outcome variable:              Cum(y)    Minimum horizon:                    0
        Regressor:                     Cum(x)    Lagged control:                     w
        Endogenous variable:           Cum(x)    Instruments:                    z1 z2
        ──────────────────────────────────────────────────────────────────────────────"""
end

@testset "_makeYX" begin
    T = 100
    ys = [randn(T), randn(T)]
    xs = [randn(T), randn(T)]
    ws = ys
    Y, X, T1, eT = _makeYX(ys, xs, ws, nothing, 3, 5, nothing)
    @test T1 == T-8
    @test size(Y) == (T1, 2)
    @test size(X) == (T1, 8)
    @test Y[:,1] == ys[1][9:end]
    @test Y[:,2] == ys[2][9:end]
    @test X[:,1] == xs[1][4:end-5]
    @test X[:,3] == ws[1][3:end-6]
    @test X[:,4] == ws[2][3:end-6]
    @test X[:,8] == ws[2][1:end-8]
    @test all(eT)

    sts = (rand(T), rand(T))
    Y, X, T1, eT = _makeYX(ys, xs, ws, sts, 1, 0, nothing)
    @test size(X) == (99, 6)
    @test X[:,1] == xs[1][2:end,1]
    @test X[:,2] == xs[2][2:end,1]
    @test X[:,3] == ys[1][1:99,1].*sts[1][2:end]
    @test X[:,4] == ys[1][1:99,1].*sts[2][2:end]
    @test X[:,5] == ys[2][1:99,1].*sts[1][2:end]
    @test X[:,6] == ys[2][1:99,1].*sts[2][2:end]
    @test all(eT)

    ys = (randn(T),)
    xs = ()
    ws = ys
    Y, X, T1, eT = _makeYX(ys, xs, ws, nothing, 1, 0, nothing)
    @test T1 == 99
    @test size(Y) == (T1, 1)
    @test size(X) == (T1, 1)
    @test Y == reshape(ys[1][2:end], T1, 1)
    @test X == reshape(ws[1][1:end-1], T1, 1)
    @test all(eT)

    ys[1][2], ys[1][3] = NaN, Inf
    xs = (convert(Vector{Union{Float64, Missing}}, randn(T)),)
    xs[1][3] = missing
    Y, X, T1, eT = _makeYX(ys, xs, ws, nothing, 1, 0, nothing)
    @test T1 == 96
    @test size(Y) == (T1, 1)
    @test size(X) == (T1, 2)
    @test Y == reshape(ys[1][5:end], T1, 1)
    @test X[:,1] == xs[1][5:end]
    @test X[:,2] == ws[1][4:end-1]
    @test eT == ((1:99).>3)

    Y, X, T1, eT = _makeYX(ys, xs, ws, nothing, 1, 1, (1:100).<=90)
    @test T1 == 85
    @test size(Y) == (T1, 1)
    @test size(X) == (T1, 2)
    @test Y == reshape(ys[1][6:90], T1, 1)
    @test X[:,1] == xs[1][5:89]
    @test X[:,2] == ws[1][4:88]
    # eT covers 2:99 in full data
    @test eT == ((1:98).>3) .& ((1:98).<=88)

    Y, X, T1, eT = _makeYX(ys, xs, ws, nothing, 1, 1, ((1:100).<60).|((1:100).>=70))
    @test T1 == 83
    @test Y[:] == vcat(ys[1][6:59], ys[1][72:end])
    @test X[:,1] == vcat(xs[1][5:58], xs[1][71:end-1])
    @test X[:,2] == vcat(ws[1][4:57], ws[1][70:end-2])
    @test eT == ((1:98).>3) .& (((1:98).<58) .| ((1:98).>=70))

    ys = ()
    @test_throws ArgumentError _makeYX(ys, xs, ws, nothing, 1, 0, nothing)
    ws = ()
    @test_throws ArgumentError _makeYX(ys, xs, ws, nothing, 1, 0, nothing)
    @test_throws ArgumentError _makeYX(ys, xs, ws, nothing, 0, 0, nothing)

    ys = ([Inf, NaN, 1.0],)
    xs = ()
    ws = ys
    @test_throws ArgumentError _makeYX(ys, xs, ws, nothing, 1, 0, nothing)
end

@testset "lp" begin
    T = 100
    ys = [randn(T), randn(T)]
    xs = [randn(T), randn(T)]
    ws = ys
    b, v, t, m = _lp(ys, xs, ws, nothing, 3, 5, SimpleVCE(), nothing)
    @test size(b) == (8, 2)
    @test size(v) == (16, 16)
    @test t == 92

    # Reproduce estimates from example based on Gertler & Karadi (2015)
    # Specification is non-augmented local projection
    # EHW standard errors
    # https://github.com/jm4474/Lag-augmented_LocalProjections/blob/master/examples/gk.m
    # The intercept estimated will be the same if not detrend the level before estimation
    # The left division in Matlab code uses QR factorization to solve OLS
    # The use of Cholesky factorization results in some trivial discrepancy
    df = exampledata(:gk)
    ns = (:logcpi, :logip, :ff, :ebp, :ff4_tc)
    r = lp(df, :ebp, wnames=ns, nlag=12, nhorz=48, vce=HRVCE())
    @test coef(r, 1, :ff4_tc, lag=1) ≈ 0.394494467439933 atol=1e-8
    @test coef(r, 5, :ff4_tc, lag=1) ≈ 0.147288229090455 atol=1e-8
    @test sqrt(vcov(r, 1, :ff4_tc, lag1=1)) ≈ 0.473834526724786 atol=1e-8
    @test sqrt(vcov(r, 5, :ff4_tc, lag1=1)) ≈ 0.737092370022356 atol=1e-8
    # Compare additional estimates obtained from Matlab lp before being turned into irf
    # Run in gk.m: lp(Y, 11, 1, 4, false, false)
    @test coef(r, 1, :constant) ≈ -2.945365821418727 atol=1e-8
    @test coef(r, 1, :logip, lag=1) ≈ -0.060094199554036 atol=1e-8
    @test vcov(r, 1, :constant) ≈ 0.641933721888132 atol=1e-8
    @test vcov(r, 1, :logcpi, lag1=1) ≈ 0.007331238288153 atol=1e-8
    @test vcov(r, 1, :logcpi, lag1=12) ≈ 0.004780116443880 atol=1e-8

    f = irf(r, :ebp, :ff4_tc, lag=1)
    @test coef(f)[1] ≈ 0.394494467439933 atol=1e-8
    @test coef(f)[20] ≈ -0.552849402428304 atol=1e-8
    @test stderror(f)[1] ≈ 0.473834526724786 atol=1e-8
    @test stderror(f)[20] ≈ 1.005046789996172 atol=1e-8
    ci = confint(f)
    @test ci[1][1] ≈ -0.384893972486403 atol=1e-8
    @test ci[2][1] ≈ 1.173882907366730 atol=1e-8
    @test ci[1][20] ≈ -2.206004261810359 atol=1e-8
    @test ci[2][20] ≈ 1.100305456953579 atol=1e-8

    ns1 = [3, :logip, 5, :ff4_tc]
    r1 = lp(df, 11, wnames=ns1, nlag=12, nhorz=1, vce=HRVCE())
    @test r1.B[1:3,1,1] ≈ r.B[1:3,1,1]
    @test r1.V[1:3,1:3,1] ≈ r.V[1:3,1:3,1]

    rn = lp(df, :ebp, xnames=:ff4_tc, wnames=(:ff4_tc,), nlag=12, nhorz=2,
        normalize=:ff4_tc=>:ff, vce=HRVCE())
    riv = lp(df, :ebp, xnames=:ff, wnames=(:ff4_tc,), nlag=12, nhorz=2,
        iv=:ff=>:ff4_tc, vce=HRVCE())
    @test coef(rn, 1, :ff4_tc) ≈ coef(riv, 1, :ff)
    @test coef(rn, 2, :ff4_tc) ≈ coef(riv, 2, :ff)
    @test vcov(rn, 1, :ff4_tc) ≈ vcov(riv, 1, :ff)
    @test vcov(rn, 2, :ff4_tc) ≈ vcov(riv, 2, :ff)
    @test rn.normnames == [:ff4_tc]
    @test rn.normtars == [:ff]
    @test rn.normmults ≈ [-5.3028242533068495] atol=1e-8
    @test riv.endonames == VarName[:ff]
    @test riv.ivnames == VarName[:ff4_tc]

    @test_throws ArgumentError lp(df, :ebp, xnames=:ff4_tc, normalize=:no=>:ff)
    @test_throws ArgumentError lp(df, :ebp, xnames=:ff, iv=:no=>:ff4_tc)
    @test_throws ArgumentError lp(df, :ebp, xnames=:ff, iv=:ff=>())
    @test_throws ArgumentError lp(df, :ebp, xnames=:ff, iv=:ff=>1.0)

    # Compare standard errors based on EWC with Matlab results
    r2 = lp(r, HARVCE(EWC()))
    f2 = irf(r2, :ebp, :ff4_tc, lag=1)
    @test f2.B == f.B
    @test f2.T == f.T
    @test stderror(f2)[1] ≈ 0.408141982124863 atol=1e-8
    @test stderror(f2)[20] ≈ 1.037575254556690 atol=1e-8
    ci2 = confint(f2)
    @test ci2[1][1] ≈ -0.318073956754668 atol=1e-8
    @test ci2[2][1] ≈ 1.107062891634995 atol=1e-8
    @test ci2[1][20] ≈ -2.371771071486807 atol=1e-8
    @test ci2[2][20] ≈ 1.266072266630027 atol=1e-8

    # Reproduce point estimates in Table 1 from Ramey and Zubairy (2018)
    # Compare estimates produced by ivreg2 in Stata replication files
    # `jordagk.do` and `jordagk_twoinstruments.do`
    # Need to run `set type double` before executing the Stata do files for more precision
    # Use `bro h multlin1` to see the Stata results
    df = exampledata(:rz)

    # Baseline specifications
    r1 = lp(df, Cum(:y), xnames=Cum(:g), wnames=(:newsy, :y, :g), iv=Cum(:g)=>:newsy,
        nlag=4, nhorz=17, addylag=false, firststagebyhorz=true)
    f1 = irf(r1, Cum(:y), Cum(:g))
    @test coef(f1)[1] ≈ 1.306459420753 atol=1e-8
    @test coef(f1)[9] ≈ .6689611146745 atol=1e-9
    @test coef(f1)[17] ≈ .7096110638625 atol=1e-9
    # Standard errors based on EWC rather than Newey-West in original paper
    @test stderror(f1)[1] ≈ 0.4286854811143802 atol=1e-8
    @test stderror(f1)[9] ≈ 0.10921291172237432 atol=1e-8
    @test stderror(f1)[17] ≈ 0.20049776900175723 atol=1e-8

    r2 = lp(df, Cum(:y), xnames=Cum(:g), wnames=(:newsy, :y, :g), iv=Cum(:g)=>(:newsy, :g),
        nlag=4, nhorz=16, minhorz=1, addylag=false, firststagebyhorz=true)
    f2 = irf(r2, Cum(:y), Cum(:g))
    @test coef(f2)[1] ≈ .218045283661 atol=1e-9
    @test coef(f2)[8] ≈ .4509452247638 atol=1e-9
    @test coef(f2)[16] ≈ .5591002357465 atol=1e-9
    # Standard errors based on EWC rather than Newey-West in original paper
    @test stderror(f2)[1] ≈ 0.16423131478595118 atol=1e-8
    @test stderror(f2)[8] ≈ 0.10754779358937377 atol=1e-8
    @test stderror(f2)[16] ≈ 0.14263798115850218 atol=1e-8

    # Omit WWII
    r1 = lp(df, Cum(:y), xnames=Cum(:g), wnames=(:newsy, :y, :g), iv=Cum(:g)=>:newsy,
        nlag=4, nhorz=17, addylag=false, firststagebyhorz=true, subset=df.wwii.==0)
    f1 = irf(r1, Cum(:y), Cum(:g))
    @test coef(f1)[1] ≈ 1.7188411171 atol=1e-8
    @test coef(f1)[9] ≈ .7672214331391 atol=1e-9
    @test coef(f1)[17] ≈ .7167179685971 atol=1e-9

    r2 = lp(df, Cum(:y), xnames=Cum(:g), wnames=(:newsy, :y, :g), iv=Cum(:g)=>(:newsy, :g),
        nlag=4, nhorz=16, minhorz=1, addylag=false, firststagebyhorz=true, subset=df.wwii.==0)
    f2 = irf(r2, Cum(:y), Cum(:g))
    @test coef(f2)[1] ≈ -.0342369278527 atol=1e-10
    @test coef(f2)[8] ≈ .2570200212481 atol=1e-9
    @test coef(f2)[16] ≈ .2225592580744 atol=1e-9

    # State dependency
    # Compare estimates generated from `jordagk_twoinstruments.do`
    # Use `bro multexpbh* multrecbh*` to see the Stata results

    # Construct the lag of slack
    df.rec = similar(df.slack)
    df.rec[1] = missing
    # Fill the beginning missing values
    df.rec[2:5] .= true
    df.rec[6:end] .= view(df.slack, 5:length(df.rec)-1)
    df.exp = 1.0 .- df.rec
    df.recnewsy = df.rec.*df.newsy
    df.expnewsy = df.exp.*df.newsy
    df.recg = df.rec.*df.g
    df.expg = df.exp.*df.g

    # nomit
    r2 = lp(df, Cum(:y), xnames=(Cum(:g,:rec), Cum(:g,:exp), :rec), wnames=(:newsy, :y, :g),
        iv=(Cum(:g,:rec), Cum(:g,:exp))=>(:recnewsy, :expnewsy, :recg, :expg),
        states=(:rec, :exp), nlag=4, nhorz=16, minhorz=1, addylag=false, firststagebyhorz=true)
    f2rec = irf(r2, Cum(:y), Cum(:g,:rec))
    @test coef(f2rec)[1] ≈ .2718093668839 atol=1e-9
    @test coef(f2rec)[8] ≈ .6357587189621 atol=1e-9
    @test coef(f2rec)[16] ≈ .6784987089259 atol=1e-9
    f2exp = irf(r2, Cum(:y), Cum(:g,:exp))
    @test coef(f2exp)[1] ≈ .2660710337235 atol=1e-9
    @test coef(f2exp)[8] ≈ .3511868388287 atol=1e-9
    @test coef(f2exp)[16] ≈ .373442191223 atol=1e-9

    # wwii
    r2 = lp(df, Cum(:y), xnames=(Cum(:g,:rec), Cum(:g,:exp), :rec), wnames=(:newsy, :y, :g),
        iv=(Cum(:g,:rec), Cum(:g,:exp))=>(:recnewsy, :expnewsy, :recg, :expg),
        states=(:rec, :exp), nlag=4, nhorz=16, minhorz=1, addylag=false, firststagebyhorz=true,
        subset=df.wwii.==0)
    f2rec = irf(r2, Cum(:y), Cum(:g,:rec))
    @test coef(f2rec)[1] ≈ .345617820757 atol=1e-9
    @test coef(f2rec)[8] ≈ 1.350903038984 atol=1e-8
    @test coef(f2rec)[16] ≈ 1.365148452565 atol=1e-8
    f2exp = irf(r2, Cum(:y), Cum(:g,:exp))
    @test coef(f2exp)[1] ≈ .010365014606 atol=1e-10
    @test coef(f2exp)[8] ≈ .2171636455168 atol=1e-9
    @test coef(f2exp)[16] ≈ .2140385336746 atol=1e-9

    @test_logs (:warn, "firststagebyhorz=false while endogenous variables contain Cum")
        lp(df, Cum(:y), xnames=Cum(:g), wnames=(:y,), iv=Cum(:g)=>:newsy, addylag=false)
    @test_logs (:warn, "addylag=true while outcome variables contain Cum")
        lp(df, Cum(:y), xnames=Cum(:g))
end


"""
A crude method for finding integer-valued rows of an array
"""
function findintegerfeatures(X::AbstractArray, N=10, prec=3)
    mask = fill(false, size(X, 1))
    for (i, f) ∈ enumerate(eachrow(X))
        σ = std(f)
        fr = round.(f./σ, sigdigits=prec)
        mask[i] = length(unique(fr)) < N
    end
    return mask
end
export findintegerfeatures

_nointegerfeatures(X::AbstractArray, args...) = Y->Y[.!findintegerfeatures(X, args...), :]
nointegerfeatures(X, args...) = _nointegerfeatures(X, args...)(X)
export _nointegerfeatures, nointegerfeatures

"""
Can we remove any perfectly correlated features?
"""
function findredundantfeatures(X::AbstractArray)
    Σ = cor(X, dims=2)
    for r ∈ 1:size(Σ, 1)
        cfs = findall(r == 1.0) # Features perfectly correlated to this one
        cfs = cfs[cfs .!= r] # Don't count the one we're looking at
        for cf ∈ cfs
            Σ[cf, :] .= NaN
            Σ[:, cf] .= NaN
        end
    end
    idxs = isnan.(diag(Σ))
end
export findredundantfeatures

_noredundantfeatures(X::AbstractArray, args...) = Y->Y[.!findredundantfeatures(X, args...), :]
noredundantfeatures(X, args...) = _noredundantfeatures(X, args...)(X)
export _noredundantfeatures, noredundantfeatures


function separate_columns(data::Vector{Vector{T}}) where T <: Number
    if length(data) > 0
        columns_number = length(data[1])
        columns = Vector{Vector{Float64}}(undef, columns_number)
        for i = 1:columns_number
            columns[i] = getindex.(data, i)
        end
        return columns
    else
        error("Null data vector.")
    end
end

function separate_columns(data::Vector{Vector{T}}) where T <: Any
    if length(data) > 0
        columns_number = length(data[1])
        type = typeof(data[1])
        columns = Vector{type}(undef, columns_number)
        for i = 1:columns_number
            columns[i] = getindex.(data, i)
        end
        return columns
    else
        error("Null data vector.")
    end
end

function correct_upper_bounds!(vector, bounds)
    if length(vector) == length(bounds)
        for i = 1:length(vector)
            if vector[i] > bounds[i]
                vector[i] = bounds[i]
            end
        end
    else
        error("Bounds were not set to all elements.")
    end
end

function correct_upper_bounds(vector, bounds)
    if length(vector) == length(bounds)
        new_vector = typeof(vector)(undef, length(vector))
        for i = 1:length(vector)
            if vector[i] > bounds[i]
                new_vector[i] = bounds[i]
            else
                new_vector[i] = vector[i]
            end
        end
        return new_vector
    else
        error("Bounds were not set to all elements.")
    end
end

function get_bin_from_indices(vector, bounds)
    if length(vector) == length(bounds)
        index = 0
        len = length(vector)
        for i = 1:(length(vector)-1)
            index += (vector[i]-1)*prod(bounds[(i+1):end])
        end
        if len >= 1 index += vector[end] end
        return index
    else
        error("Vector and bound vector have different lengths.")
    end
end

get_weights_length(h::Histogram) = length(h.weights)

function flat_vector(vector; nbins = 5, lower_lim=-1, upper_lim=1)
    step = (upper_lim-lower_lim)/nbins
    columns = separate_columns(vector)
    hists = fit.(Histogram, columns, [lower_lim:step:upper_lim])
    bounds = (get_weights_length.(hists)...,)
    indices = [getindex.(get_bin.(hists, v),1) for v in vector]
    indices = [correct_upper_bounds(ind, bounds) for ind in indices]
    get_bin_from_indices.(indices, [bounds])
end

#=
function get_layer_bins_vector(vector; nbins = 5, lower_lim=-1, upper_lim=1)
    columns = separate_columns(vector)
    hists = create_bin_histogram.(columns, [nbins], [lower_lim], [upper_lim])
    bounds = (get_weights_length.(hists)...,)
    indices = [getindex.(get_bin.(hists, v),1) for v in vector]
    indices = [correct_upper_bounds(ind, bounds) for ind in indices]
    combination_values = get_bin_from_indices.(indices, [bounds])
    max_value = foldr(*, bounds)
    create_bin_vector(combination_values, nbins, 1, max_value)
end
=#

function get_layer_bins_vector(vector; neuron_nbins = 5, nbins = 5, lower_lim=-1, upper_lim=1)
    columns = separate_columns(vector)
    hists = create_bin_histogram.(columns, [neuron_nbins], [lower_lim], [upper_lim])
    bounds = (get_weights_length.(hists)...,)
    indices = [getindex.(get_bin.(hists, v),1) for v in vector]
    indices = [correct_upper_bounds(ind, bounds) for ind in indices]
    combination_values = get_bin_from_indices.(indices, [bounds])
    max_value = foldr(*, bounds)
    create_bin_vector(combination_values, nbins, 1, max_value)
end

#=
function fit_into_histogram(vector, nbins::Int, lower_lim, upper_lim; ϵ = 0.01)
    hist_step = (upper_lim-lower_lim)/nbins
    ϵ = 0.01
    edges = collect(lower_lim:hist_step:upper_lim)
    edges[1] -= ϵ; edges[end] += ϵ
    hist = fit(Histogram, vector, edges)
    [getindex.(get_bin(hist, v),1) for v in vector]
end
=#

function create_bin_histogram(vector::Vector{T}, nbins::Int, lower_lim, upper_lim; ϵ = 0.01) where T <: Number
    hist_step = abs(upper_lim-lower_lim)/nbins
    edges = collect(lower_lim:hist_step:upper_lim)
    edges[1] -= ϵ; edges[end] += ϵ
    hist = fit(Histogram, vector, edges)
end

function create_bin_vector(vector::Vector{T}, nbins::Int, lower_lim, upper_lim; ϵ = 0.01) where T <: Number
    hist = create_bin_histogram(vector, nbins, lower_lim, upper_lim, ϵ = ϵ)
    [getindex.(get_bin(hist, v),1) for v in vector]
end

function get_input_bin_vector(vector, nfeatures; nbins = 5, return_bins::Int=0)
    hists = Vector{Vector{Int}}(undef, nfeatures)
    columns = separate_columns(vector)
    for feature in 1:nfeatures
        feature_values = columns[feature]
        min = minimum(feature_values)
        max = maximum(feature_values)
        hists[feature] = create_bin_vector(feature_values, nbins, min, max)
    end
    bounds = repeat([nbins], nfeatures)
    indices = separate_columns(hists)
    combination_values = get_bin_from_indices.(indices, [bounds])
    if return_bins > 0
        max_value = foldr(*, bounds)
        return create_bin_vector(combination_values, return_bins, 1, max_value)
    else
        return floor.(Int, combination_values)
    end
end


using Newton
using Test
using ForwardDiff
using LinearAlgebra
using StaticArrays

multiinput_rf!(r::Vector, x::Vector, A::Matrix, b::Vector) = (r .= b .+ A*x)

function setup_cache(x, A, b)
    rf!(r, x) = multiinput_rf!(r, x, A, b)
    return NewtonCache(x, rf!)
end

function determine_solution(A, b)
    x = zero(b)
    rf!(r, x) = multiinput_rf!(r, x, A, b)
    cache=NewtonCache(x, rf!)
    newtonsolve!(x, get_drdx(cache), rf!, cache)
    return x
end

@testset "linsolve!" begin
    A = 2*I + rand(10,10)
    b = rand(10)
    A1, A2 = [copy(A) for _ in 1:2]
    b1, b2 = [copy(b) for _ in 1:2]
    rf!(r, x) = (r .= x)
    @test A1\b1 ≈ Newton.linsolve!(A2, b2, NewtonCache(b, rf!))
end

@testset "newtonsolve!" begin
    nsize = 4
    (a,b,x0) = [rand(nsize) for _ in 1:3]
    tol = 1.e-10

    # Basic functionality
    function rf_solution!(r, x)
        r .= - a + b.*x + exp.(x)
    end
    x = copy(x0)
    cache = NewtonCache(x, rf_solution!)
    drdx = get_drdx(cache)
    @test size(drdx,1) == nsize
    converged = newtonsolve!(x, drdx, rf_solution!; tol=tol)
    r_check = similar(x0)
    @test converged
    @test isapprox(rf_solution!(r_check, x), zero(r_check); atol=tol)
    @test drdx ≈ ForwardDiff.jacobian(rf_solution!, r_check, x)

    function rf_nosolution!(r, x)
        r .= a .+ b.*x.^2
    end
    x = copy(x0)
    converged = newtonsolve!(x, drdx, rf_nosolution!)
    @test !converged

    # Test function with different anynomous functions for cache and solving
    cache = setup_cache(x0, zeros(nsize,nsize), zeros(nsize))   # If these A and b inputs are used, solution is zero
    A = nsize*I+rand(nsize,nsize)
    A += transpose(A)   # Symmetrize to ensure invertible
    b = rand(nsize)

    rf_solve!(r, x) = multiinput_rf!(r, x, A, b)
    drdx = get_drdx(cache)
    @test newtonsolve!(x, drdx, rf_solve!, cache)
    @test isapprox(rf_solve!(r_check, x), zero(r_check); atol=tol)
    @test .!(isapprox(x, zero(x); atol=tol))

    # Test that error is thrown if we try to use automatic differentiation of the newtonsolve!
    diff_fun(y) = determine_solution(A, y)
    failed = false
    try
        df = ForwardDiff.jacobian(diff_fun, rand(nsize))
    catch err
        failed = true
    end
    @test failed

end

@testset "linsolve" begin
    dim = 5
    A = rand(SMatrix{dim,dim})
    x = rand(SVector{dim})
    b = A*x
    @test x ≈ Newton.linsolve(A, b)
end

@testset "newtonsolve" begin
    dim = 6
    (a,b,x0) = [rand(SVector{dim}) for _ in 1:3]

    rf_solution(x) = - a + b.*x + exp.(x)
    rf_nosolution(x) = a .+ b.*x.^2

    converged, x, drdx = newtonsolve(x0, rf_solution; tol=1.e-6)
    @test converged
    @test isapprox(norm(rf_solution(x)), 0.0, atol=1.e-6)
    @test drdx ≈ ForwardDiff.jacobian(rf_solution, x)

    converged, x, drdx = newtonsolve(x0, rf_nosolution; tol=1.e-6, maxiter=4)
    @test ~converged
    @test all(isnan.(x))
    @test all(isnan.(drdx))

end


export ToeplitzOperator, HankelOperator

mutable struct ToeplitzOperator{T<:Number} <: Operator{T}
    negative::Vector{T}
    nonnegative::Vector{T}
end

ToeplitzOperator(V::Vector{T},W::Vector{Q}) where {T<:Number,Q<:Number} =
    ToeplitzOperator{promote_type(T,Q)}(V,W)
ToeplitzOperator(V::AbstractVector,W::AbstractVector) =
    ToeplitzOperator(collect(V),collect(W))

convert(::Type{Operator{TT}},T::ToeplitzOperator) where {TT} =
    ToeplitzOperator(convert(Vector{TT},T.negative),convert(Vector{TT},T.nonnegative))

for op in (:(Base.real), :(Base.imag))
    @eval $op(T::ToeplitzOperator) = ToeplitzOperator($op(T.negative), $op(T.nonnegative))
end

function SymToeplitzOperator(V::Vector)
    W=V[2:end]
    V=copy(V)
    V[1]*=2
    ToeplitzOperator(W,V)
end

for OP in (:domainspace,:rangespace)
    @eval $OP(T::ToeplitzOperator) = ℓ⁰
end

getindex(T::ToeplitzOperator,k::Integer,j::Integer) =
    toeplitz_getindex(T.negative,T.nonnegative,k,j)

function toeplitz_getindex(negative::AbstractVector{T},nonnegative::AbstractVector{T},k::Integer,j::Integer) where T
    if 0<k-j≤length(negative)
        negative[k-j]
    elseif 0≤j-k≤length(nonnegative)-1
        nonnegative[j-k+1]
    else
        zero(T)
    end
end

function toeplitz_getindex(cfs::AbstractVector{T},k::Integer,j::Integer) where T
    if k==j && !isempty(cfs)
        2cfs[1]
    elseif 0<k-j≤length(cfs)-1
        cfs[k-j+1]
    elseif 0<j-k≤length(cfs)-1
        cfs[j-k+1]
    else
        zero(T)
    end
end

function BandedMatrix(S::SubOperator{T,ToeplitzOperator{T},Tuple{UnitRange{Int},UnitRange{Int}}}) where T
    ret = BandedMatrix(Zeros, S)

    kr,jr=parentindices(S)

    neg=parent(S).negative
    pos=parent(S).nonnegative

    toeplitz_axpy!(1.0,neg,pos,kr,jr,ret)
end

bandwidths(T::ToeplitzOperator)=(length(T.negative),length(T.nonnegative)-1)

# slice of a ToeplitzOPerator is a ToeplitzOperator

function Base.getindex(T::ToeplitzOperator,kr::InfRanges,jr::InfRanges)
    sh=first(jr)-first(kr)
    st=step(jr)
    @assert st==step(kr)
    if sh ≥0
        ToeplitzOperator([reverse!(T.nonnegative[1:sh]);T.negative],T.nonnegative[sh+1:st:end])
    else
        ToeplitzOperator(T.negative[-sh+1:st:end],[reverse!(T.negative[1:-sh]);T.nonnegative])
    end
end

## Hankel Operator

mutable struct HankelOperator{T<:Number} <: Operator{T}
    coefficients::Vector{T}
end

for OP in (:domainspace,:rangespace)
    @eval $OP(T::HankelOperator) = ℓ⁰
end

HankelOperator(V::AbstractVector)=HankelOperator(collect(V))

HankelOperator(f::Fun)=HankelOperator(f.coefficients)

@eval convert(::Type{Operator{TT}},T::HankelOperator) where {TT}=HankelOperator(convert(Vector{TT},T.coefficients))

function hankel_getindex(v::AbstractVector,k::Integer,j::Integer)
   if k+j-1 ≤ length(v)
        v[k+j-1]
    else
        zero(eltype(v))
    end
end

getindex(T::HankelOperator,k::Integer,j::Integer) =
    hankel_getindex(T.coefficients,k,j)

function BandedMatrix(S::SubOperator{T,HankelOperator{T},Tuple{UnitRange{Int},UnitRange{Int}}}) where T
    ret=BandedMatrix(Zeros, S)

    kr,jr=parentindices(S)
    cfs=parent(S).coefficients

    hankel_axpy!(1.0,cfs,kr,jr,ret)
end

bandwidths(T::HankelOperator) = (max(0,length(T.coefficients)-1),max(0,length(T.coefficients)-1))

## algebra

function Base.maximum(T::ToeplitzOperator)
    if isempty(T.negative)
        maximum(T.nonnegative)
    elseif isempty(T.nonnegative)
        maximum(T.negative)
    else
        max(maximum(T.negative),maximum(T.nonnegative))
    end
end

-(T::ToeplitzOperator)=ToeplitzOperator(-T.negative,-T.nonnegative)
*(c::Number,T::ToeplitzOperator)=ToeplitzOperator(c*T.negative,c*T.nonnegative)

-(H::HankelOperator)=HankelOperator(-H.coefficients)
*(c::Number,H::HankelOperator)=HankelOperator(c*H.coefficients)

## inv

function Base.inv(T::ToeplitzOperator)
    @assert length(T.nonnegative)==1
    ai=\(T,[1.0];maxlength=100000)
    ToeplitzOperator(ai[2:end],ai[1:1])
end

####
# Toeplitz/Hankel
####

# αn,α0,αp give the constant for the negative, diagonal and positive
# entries.  The usual case is 1,2,1
function toeplitz_axpy!(αn,α0,αp,neg,pos,kr,jr,ret)
    dat=ret.data
    m=size(dat,2)

    dg=diagindrow(ret,kr,jr)

    # diagonal
    if dg ≥ 1
        α0p=α0*pos[1]
        @simd for j=1:m
            @inbounds dat[dg,j]+=α0p
        end
    end

    # positive entries
    for k=2:min(length(pos),dg)
        αpp=αp*pos[k]
        @simd for j=1:m
            @inbounds dat[dg-k+1,j]+=αpp
        end
    end

    # negative entries
    for k=1:min(length(neg),size(dat,1)-dg)
        αnn=αn*neg[k]
        @simd for j=1:m
            @inbounds dat[dg+k,j]+=αnn
        end
    end

    ret
end

# this routine is for when we want a symmetric toeplitz
function sym_toeplitz_axpy!(αn,α0,αp,cfs,kr,jr,ret)
    dg=diagindrow(ret,kr,jr)

    dat=ret.data
    m=size(dat,2)
    # diagonal
    if dg ≥ 1
        α0p=α0*cfs[1]
        @simd for j=1:m
            @inbounds dat[dg,j]+=α0p
        end
    end

    # positive entries
    for k=2:min(length(cfs),dg)
        αpp=αp*cfs[k]
        @simd for j=1:m
            @inbounds dat[dg-k+1,j]+=αpp
        end
    end

    # negative entries
    for k=2:min(length(cfs),size(dat,1)-dg+1)
        αnn=αn*cfs[k]
        @simd for j=1:m
            @inbounds dat[dg+k-1,j]+=αnn
        end
    end

    ret
end

toeplitz_axpy!(α,neg,pos,kr,jr,ret) =
    toeplitz_axpy!(α,α,α,neg,pos,kr,jr,ret)

sym_toeplitz_axpy!(α0,αp,cfs,kr,jr,ret) =
    sym_toeplitz_axpy!(αp,α0,αp,cfs,kr,jr,ret)

function hankel_axpy!(α,cfs,kr,jr,ret)
    dat=ret.data

    st=stride(dat,2)-2
    mink=first(kr)+first(jr)-1

    N=length(dat)

    # dg gives the row corresponding to the diagonal of the original operator
    dg=diagindrow(ret,kr,jr)

    # we need the entry where the first entry is written to
    # this is going to be a shift of the diagonal of the true operator
    dg1=dg+first(kr)-first(jr)

    for k=mink:min(length(cfs),size(dat,1)+mink-dg1)
        dk=k-mink+dg1
        nk=k-mink
        αc=α*cfs[k]
        @simd for j=dk:st:min(dk+nk*st,N)
            @inbounds dat[j]+=αc
        end
    end

    ret
end


@doc """
    simple_estimator(model::Ising, T::Real, Js::AbstractArray)

Returns the following observables as `Dict{String, Any}`

# Observables
- `"Energy"`
    - energy density
- `"Energy^2"`
    - square of energy density
- `"Magnetization"`
    - magnetization density
- `"|Magnetization|"`
    - absolute value of magnetization density
- `"Magnetization^2"`
    - square of magnetization density
- `"Magnetization^4"`
    - quadruple of magnetization density
"""
function simple_estimator(model::Ising, T::Real, Js::AbstractArray, _=nothing)
    nsites = numsites(model)
    nbonds = numbonds(model)

    M = mean(model.spins)
    E = 0.0
    @inbounds for b in bonds(model)
        s1, s2 = source(b), target(b)
        E += ifelse(model.spins[s1] == model.spins[s2], -1.0, 1.0) * Js[bondtype(b)]
    end
    E /= nsites

    res = Measurement()
    res["Magnetization"] = M
    res["|Magnetization|"] = abs(M)
    res["Magnetization^2"] = M^2
    res["Magnetization^4"] = M^4
    res["Energy"] = E
    res["Energy^2"] = E^2
    return res
end

@doc raw"""
    improved_estimator(model::Ising, T::Real, Js::AbstractArray, sw::SWInfo)

Returns the following observables as `Dict{String, Any}` using cluster information `sw`

# Observables
- `"Energy"`
    - energy density
- `"Energy^2"`
    - square of energy density
- `"Magnetization"`
    - magnetization density
- `"|Magnetization|"`
    - absolute value of magnetization density
- `"Magnetization^2"`
    - square of magnetization density
- `"Magnetization^4"`
    - quadruple of magnetization density
- `"Clustersize^2"`
    - ``\sum_c r_c^2 ``, where ``r_c`` is the size density of ``c``-th cluster
- `"Clustersize^4"`
    - ``\sum_c r_c^4 ``
- `"Clustersize^2 Clustersize^2"`
    - ``\sum_{c\ne c'} r_c^2 r_{c'}^2``
"""
function improved_estimator(model::Ising, T::Real, Js::AbstractArray, sw::SWInfo)
    nsites = numsites(model)
    nbonds = numbonds(model)
    nc = numclusters(sw)
    invV = 1.0/nsites

    ## magnetization
    M = 0.0
    N2 = 0.0
    N2N2 = 0.0
    N4 = 0.0
    for (m,s) in zip(sw.clustersize, sw.clusterspin)
        M += m*invV*s
        m2 = (m*invV)^2
        N4 += m2*m2
        N2N2 += N2*m2
        N2 += m2
    end
    M4 = N4 + 6*N2N2
    M2 = N2

    # energy
    aJ = 2.0*abs.(Js)
    mbeta = -1.0/T
    ns = sw.activated_bonds
    As = -aJ ./ expm1.(mbeta.*aJ)
    Ans = ns.*As
    E0 = 0.0
    for b in 1:numbondtypes(model)
        E0 += Js[b] * numbonds(model,b)
    end
    E = 0.0
    E2 = 0.0
    for b in 1:numbondtypes(model)
        E2 += (aJ[b]-2.0*E0)*Ans[b]
        E2 += Ans[b] * As[b]*(ns[b]-1)
        E2 += 2.0*Ans[b]*E
        E += Ans[b]
    end
    E -= E0
    E2 += E0^2

    E *= -invV
    E2 *= invV*invV

    res = Measurement()
    res["Magnetization"] = M
    res["|Magnetization|"] = abs(M)
    res["Magnetization^2"] = M2
    res["Magnetization^4"] = M4
    res["Energy"] = E
    res["Energy^2"] = E2
    res["Clustersize^2"] = N2
    res["Clustersize^4"] = N4
    res["Clustersize^2 Clustersize^2"] = 2*N2N2

    return res
end

default_estimator(model::Ising, update) = ifelse(update==SW_update!, improved_estimator, simple_estimator)


export emsd_f, emsd_f_circuit, ρ_junct, ρ_junct_avg, ρ_junct_mean_std

using Statistics: mean, var, std

"""
```julia
emsd_f(f, params; N=32)
```
Computes the mean and standard deviation of an ensemble of `StickNetwork`s determined by `params` with ensemble size `N`.
"""
function emsd_f(f, params; N=32)
    e = ensemble(N, params...)
    res = f.(e)
    mean(res), std(res)
end

"""
```julia
emsd_f_circuit(f, params, M, args...; N=32)
```
Computes the mean and standard deviation of an ensemble of `N` nanowire networks
with model `M` and parameters 'params'. `args` may specify additonal parameters to
`ensemble_circuit` such as the `Val(:conducting)` property.
"""
function emsd_f_circuit(f, params, M, args...; N=32)
    e = ensemble_circuit(N, params..., M, args...)
    res = f.(e)
    mean(res), std(res)
end

"""
```julia
ρ_junct(net)
```
Computes junction density of `StickNetwork` `net`.
"""
function ρ_junct(net::StickNetwork{T,N}) where {T,N}
    n = length(net.lines)
    m = 0
    for i ∈ 2:(n)
        for j ∈ 1:(i-1)
            ps = intersection_params(net.lines[i],net.lines[j])
            if (0.0 ≤ ps[1] ≤ 1.0) && (0.0 ≤ ps[2] ≤ 1.0)
                p = line_segment_point(ps[1],net.lines[i])
                if sum( zeros(SVector{N,T}) .≤ p .≤ net.dims ) == N
                    m += 1
                end
            end
        end
    end
    m/prod(net.dims)
end

"""
```julia
ρ_junct_mean_std(params; N=32)
```
Computes the mean and standard deviation of the junction
density of an ensemble of `StickNetwork`s given by `params`.
"""
ρ_junct_mean_std(params; N=32) = emsd_f(ρ_junct, params; N=N)
ρ_junct_mean_std(params...; N=32) = ρ_junct_mean_std(params; N=N)

"""
Deprecated, use `ρ_junct_mean_std` instead.
"""
ρ_junct_avg(dims, ρₙ, l, ρ, D; N=32) = emsd_f(ρ_junct, [dims, ρₙ, l, ρ, D], N=N)[1]


doc"""
    typical(ohlc; h="High", l="Low", c="Close")

Typical Price

```math
    \text{Typical Price} = \frac{H + L + C}{3}
```
"""
function typical(ohlc::TimeArray{T,N}; h="High", l="Low", c="Close") where {T,N}
    val = (ohlc[h] .+ ohlc[l] .+ ohlc[c]) ./ 3
    TimeArray(val.timestamp, val.values, ["typical"], ohlc.meta)
end

function mean_abs_dev(a::Array{T,1}, scale::Bool=false) where T
    scale ? c = 1 / (-sqrt(2) * erfcinv(3 * 1 / 2)) : c =1

    res = ones(length(a))
    for i in 1:length(a)
        @inbounds res[i] = abs(a[i] - mean(a))
    end

    mean(res) * c
end

"""
Parameters:
    * orig: the original colnames
    * suffix: the suffix applied to `orig`

```jldoctest
julia> gen_colnames(["Open", "Close"], ["macd", "dif", "sig"])
6-element Array{String,1}:
 "Open_macd"
 "Close_macd"
 "Open_dif"
 "Close_dif"
 "Open_sig"
 "Close_sig"
```
"""
gen_colnames(orig::Vector{String}, suffix::Vector{String}) =
    vec(["$o\_$s" for o ∈ orig, s ∈ suffix])

relu(x) = max(x, 0)

relu(ta::TimeArray) =
    TimeArray(ta.timestamp, relu.(ta.values), ta.colnames, ta.meta)

"""
The soomth method used by ADX
"""
function wilder_smooth(ta::TimeArray, n::Integer;
                       padding::Bool=false, dx::Bool=false)
    val = similar(Array{Float64}, indices(ta.values))

    first_cal = (dx ? mean : sum)

    for i ∈ 1:size(val, 1)
        val[i, :] =
            if i < n
                NaN
            elseif i == n
                first_cal(ta[1:n].values, 1)
            elseif dx
                (val[i-1, :] .* (n - 1) .+ ta.values[i, :]) ./ n
            else
                (val[i-1, :] .* (n - 1) ./ n) .+ ta.values[i, :]
            end
    end

    ret = TimeArray(ta.timestamp, val, ta.colnames, ta.meta)

    if padding
        ret
    else
        dropnan(ret)
    end
end

safediv(x, y) = ifelse(iszero(x) && iszero(y), x, x / y)


@inbounds begin
function _vcmp(x1::Vector{U}, x2::Vector{U}) where {U<:BitUnsigned}
    l1, l2 = length(x1), length(x2)
    if l2 == l1
        for k = l1:-1:1
            if x1[k] != x2[k]
                return cmp(x1[k], x2[k])
            end
        end
        return 0
    else
        return cmp(l1, l2)
    end
end

function _vleftshift!(x::Vector{U}, n::UInt) where {U<:BitUnsigned}
    n==0 && return x
    k = 1
    while n >= fullshift(U)
        n -= fullshift(U)
        unshift!(x, zero(U))
        k += 1
    end
    carry = zero(U)
    @inbounds while k <= length(x)
        # nextcarry = x[k] >> (fullshift(U)-n)
        # x[k] = (x[k] << n) + carry
        # carry = nextcarry
        carry, x[k] = (x[k] >> (fullshift(U)-n)) , ((x[k] << n) + carry)
        k += 1
    end
    !iszero(carry) && push!(x, carry)
    return x
end
function _vrightshift!(x::Vector{U}, n::UInt) where {U<:BitUnsigned}
    n==0 && return x
    k = 1
    while n >= fullshift(U) && length(x) > 0
        n -= fullshift(U)
        shift!(x)
    end
    length(x) == 0 && return x
    while k < length(x)
        x[k] = (x[k] >> n) | (x[k+1] << (fullshift(U)-n))
        k += 1
    end
    x[k] = (x[k] >> n)

    iszero(x[k]) && pop!(x)
    return x
end
function _vand!(y::Vector{U}, x::Vector{U}) where {U<:BitUnsigned}
    while length(y) > length(x)
        pop!(y)
    end
    @inbounds for k in 1:length(y)
        y[k] &= x[k]
    end
    while length(y) > 0 && iszero(last(y))
        pop!(y)
    end
    return y
end
function _vor!(y::Vector{U}, x::Vector{U}) where {U<:BitUnsigned}
    while length(y) < length(x)
        push!(y, zero(U))
    end
    @inbounds for k = 1:length(x)
        y[k] |= x[k]
    end
    return y
end

function _vadd!(y::Vector{U}, x::Vector{U}) where {U<:BitUnsigned}
    while length(y) < length(x)
        push!(y, zero(U))
    end
    lx, ly = length(x), length(y)
    carry = zero(U)
    for k = 1:lx
        y[k], ncarry = _add(y[k], x[k])
        if iszero(carry)
            carry = ncarry
        elseif iszero(ncarry)
            y[k], carry = _add(y[k], carry)
        else # if ncarry is nonzero, y[k] < typemax and ...
            y[k] += carry # ... this cannot overflow since carry = 1
            carry = ncarry
        end
    end
    for k = (lx+1):ly
        iszero(carry) && break
        y[k], carry = _add(y[k], carry)
    end
    !iszero(carry) && push!(y, carry)
    return y
end

function _vsub!(y::Vector{U}, x::Vector{U}) where {U<:BitUnsigned}
    _vcmp(y,x) < 0 && throw(OverflowError())
    lx, ly = length(x), length(y)
    borrow = zero(U)
    for k = 1:lx
        y[k], nborrow = _sub(y[k], x[k])
        if iszero(borrow)
            borrow = nborrow
        elseif iszero(nborrow)
            y[k], borrow = _sub(y[k], borrow)
        else # if nborrow is nonzero, y[k] > 0 and ...
            y[k] -= borrow # ... this cannot overflow since borrow = 1
            borrow = nborrow
        end
    end
    for k=(lx+1):ly
        iszero(borrow) && break
        y[k], borrow = _sub(y[k], borrow)
    end
    while length(y) > 0 && iszero(last(y)) # remove zero words at beginning
        pop!(y)
    end
    return y
end

function _vmuladd!(y::Vector{U}, x1::Vector{U}, x2::Vector{U}) where {U<:BitUnsigned}
    while length(y) < length(x1)+length(x2)
        push!(y, zero(U))
    end
    l1, l2 = length(x1), length(x2)
    for l = 1:l2
        # if !iszero(x2[l]) # not worthwhile according to Knuth
            carry = zero(U)
            for k = 1:l1
                prod, carry1 = _mul(x1[k], x2[l]) # carry1 < typemax - 1
                y[k+l-1], carry2 = _add(y[k+l-1], prod) # carry2 <= 1
                y[k+l-1], carry3 = _add(y[k+l-1], carry) # carry3 <= 1 - carry2
                carry = carry1 + carry2 + carry3 # this should not overflow
            end
            for k = (l1+l):(l1+l2)
                iszero(carry) && break
                y[k], carry = _add(y[k], carry)
            end
        # end
    end
    while iszero(last(y))
        pop!(y)
    end
    return y
end
function _vdivrem!(num::Vector{U}, den::U) where {U<:BitUnsigned} # n will be destroyed in the process
    iszero(den) && throw(DivideError())
    length(num) == 0 && return(num, num)
    length(num) == 1 && return divrem(num[1], den)

    last(num) >= den && push!(num, zero(U))
    q = Vector{U}(length(num)-1)
    k = length(q)
    while length(num) >= 2
        q[k], r = _divrem(num[end-1], num[end], den)
        pop!(num)
        num[end] = r
        k -= 1
    end
    iszero(last(num)) && pop!(num)
    return q, num
end

function _vdivrem!(num::Vector{U}, den::Vector{U}) where {U<:BitUnsigned} # n will be destroyed in the process
    isempty(den) && throw(diverr)
    c = _vcmp(num,den)
    c < 0 && return (Vector{U}(0), num)
    c == 0 && return ([one(U)], Vector{U}(0))
    length(den) == 1 && return _vdivrem!(num, den[1])

    # we now know that length(den) >= 2, length(num) >= length(den)
    #normalize
    s = unsigned(leading_zeros(last(den)))
    u = _vleftshift!(num, s)
    v = _vleftshift!(den, s) # d is modified but will be restored at the end

    n = length(v)
    # compare v with first n words of u
    m = length(u) - n
    k = n
    while k > 0 && u[m+k] == v[k]
        k -= 1
    end
    if k > 0 && u[m+k] < v[k] # smaller, but u in total larger
        m -= 1
    else
        push!(u, zero(U))
    end
    # now m = length(u) - n -1
    q = fill(zero(U), m+1)

    for j = m:-1:0
        q̂ = _estimateq(u[j+n-1],u[j+n],u[j+n+1], v[n-1], v[n])
        borrow = zero(U)
        for k=1:n # combination of _vmul! and _vadd!
            u[k+j], borrow1 = _sub(u[k+j], borrow) # borrow1 <= 1
            p, borrow2 = _mul(v[k], q̂) # borrow2 <= typemax - 2
            u[k+j], borrow3 = _sub(u[k+j], p) # borrow3 <= 1 - borrow1
            borrow = borrow1 + borrow2 + borrow3 # cannot overflow
        end
        u[j+n+1], borrow = _sub(u[n+j+1], borrow)
        if !iszero(borrow) # should happen rarely according to Knuth
            q̂ -= one(q̂)
            carry = zero(U)
            for k=1:n # add back: see _vadd!
                u[k+j], ncarry = _add(u[k+j], v[k])
                if iszero(carry)
                    carry = ncarry
                elseif iszero(ncarry)
                    u[k+j], carry = _add(u[k+j], carry)
                else
                    u[k+j] += carry
                    carry = ncarry
                end
            end
            u[n+j+1], carry = _add(u[n+j+1], carry)
            # @assert borrow == carry
        end
        # @assert u[n+j+1] == zero(u[n+j+1])
        q[j+1] = q̂
    end
    while length(u) > 0 && iszero(last(u))
        pop!(u)
    end
    # restore shift
    den = _vrightshift!(v, s) # d is modified but will be restored at the end
    r = _vrightshift!(u, s) # d is modified but will be restored at the end
    return q, r
end
end


# --- PRIVATE METHODS --------------------------------------------------------------------------------------- #

# array based -
function _build_ternary_lattice_intrinsic_value_array(contractSet::Set{PSAbstractAsset}, ternaryPriceArray::Array{Float64,1})::PSResult

    # initialize -
    intrinsic_value_array = similar(ternaryPriceArray)

    # ok, so lets go through each price, compute the intrinsic value - and then store
    for (index,price) in enumerate(ternaryPriceArray)

        # ok, get the underlying price -
        underlying_price_value = ternaryPriceArray[index]

        # compute the intrinsic value -
        result = intrinsic_value(contractSet,underlying_price_value)
        if (isa(result.value,Exception) == true)
            return result
        end
        iv_value = result.value

        # add the iv to the array - index should work out
        intrinsic_value_array[index] = iv_value
    end

    # return -
    return PSResult{Array{Float64,1}}(intrinsic_value_array)
end

function _build_ternary_lattice_underlying_price_array(basePrice::Float64, volatility::Float64, timeToExercise::Int)::PSResult

    # compute up and down perturbations -
    numberOfLevels = timeToExercise + 1     # assumption: our time unit is 1 day *always* = is this legit?
    Δt = (1.0)
    U = exp(volatility * sqrt(2*Δt))
    D = 1.0 / U
    C = 1.0

    # compute price array -
    number_of_elements = Int64(((3^numberOfLevels) - 1)/2)
    priceArray = zeros(number_of_elements)
    priceArray[1] = basePrice

    # populate the ternary price tree -
    for index = 1:number_of_elements

        # get the basePrice -
        basePrice = priceArray[index]

        # compute the prices -
        down_price = basePrice*D
        up_price = basePrice*U
        center_price = basePrice*C

        # add the prices to the array -
        left_index = 3*(index - 1) + 2
        center_index = 3*(index - 1) + 3
        right_index = 3*(index - 1) + 4

        # note - we need to check that we don't write into the array past the end
        if (left_index<=number_of_elements)
            priceArray[left_index] = up_price
        end

        if (center_index<=number_of_elements)
            priceArray[center_index] = center_price
        end

        if (right_index<=number_of_elements)
            priceArray[right_index] = down_price
        end
    end

    # return the price array -
    return PSResult{Array{Float64,1}}(priceArray)
end

function _build_ternary_lattice_option_value_array(intrinsicValueArray::Array{Float64,1}, latticeModel::PSTernaryLatticeModel; earlyExcercise::Bool = false)::PSResult

    # initialize -
    contract_price_array = copy(intrinsicValueArray)

    # get stuff from the lattice model -
    σ = latticeModel.volatility
    timeToExercise = latticeModel.timeToExercise
    riskFreeRate = latticeModel.riskFreeRate
    dividendRate = latticeModel.dividendRate

    # compute the probability -
    numberOfLevels = timeToExercise + 1     # assumption: our time unit is 1 day *always* = is this legit?
    Δt = (1.0)
    T1 = exp((riskFreeRate - dividendRate)*(Δt/2))
    T2 = exp(-σ*sqrt(Δt/2))
    T3 = exp(σ*sqrt(Δt/2))
    pup = ((T1 - T2)/(T3 - T2))^2
    pdown = ((T3 - T1)/(T3 - T2))^2
    DF = exp(-riskFreeRate*Δt)
    U = exp(σ*sqrt(2*Δt))
    D = 1.0 / U
    C = 1.0

    # create a index table -
    # number_of_elements = Int64(((3^numberOfLevels) - 1)/2)
    number_of_elements = numberOfLevels
    index_table = zeros(numberOfLevels, 5)
    backwards_index_array = range(numberOfLevels, step=-1,stop=1) |> collect
    for (forward_index, backward_index) in enumerate(backwards_index_array)

        # add the prices to the array -
        left_index = 3*(backward_index - 1) + 2
        center_index = 3*(backward_index - 1) + 3
        right_index = 3*(backward_index - 1) + 4

        # populate the index table -
        index_table[forward_index,1] = backward_index
        index_table[forward_index,2] = left_index
        index_table[forward_index,3] = center_index
        index_table[forward_index,4] = right_index
    end

    # @show index_table

    # ok, so now lets compute the value for the nodes -
    for compute_index = 1:number_of_elements

        # get the indexs -
        parent_node_index = Int(index_table[compute_index,1])
        child_left_index = Int(index_table[compute_index,2])
        child_center_index = Int(index_table[compute_index,3])
        child_right_index = Int(index_table[compute_index,4])

        #@show (parent_node_index,child_left_index,child_center_index,child_right_index)

        # compute -
        contract_price = DF*(pup*contract_price_array[child_left_index]+(pdown)*contract_price_array[child_right_index]+(1-pup-pdown)*contract_price_array[child_center_index])
        if (earlyExcercise == false)
            contract_price_array[parent_node_index] = contract_price
        else
            excercise_value = contract_price_array[parent_node_index]
            contract_price_array[parent_node_index] = max(excercise_value,contract_price)
        end
    end

   # setup the results tuple -
   results_tuple = (option_index_table=index_table, option_contract_price_array=contract_price_array, U=U, D=D, PUP=pup, PDOWN=pdown, PCENTER=(1-pup-pdown), DF=DF)

   # return -
   return PSResult(results_tuple)
end
# ----------------------------------------------------------------------------------------------------------- #

# --- PUBLIC METHODS ---------------------------------------------------------------------------------------- #

"""
    option_contract_price(contractSet::Set{PSAbstractAsset}, latticeModel::PSTernaryLatticeModel, baseUnderlyingPrice::Float64;
        earlyExercise::Bool = false)::PooksoftBase.PSResult

Estimate the price of a contract using a ternary lattice pricing model.
"""
function option_contract_price(contractSet::Set{PSAbstractAsset}, latticeModel::PSTernaryLatticeModel, baseUnderlyingPrice::Float64;
    earlyExercise::Bool = false)::PSResult

    # initialize -
    option_contract_price = 0.0

    # we need to get some stuff from the lattice model -
    volatility = latticeModel.volatility
    timeToExercise = latticeModel.timeToExercise

    # compute the price array -
    result = _build_ternary_lattice_underlying_price_array(baseUnderlyingPrice, volatility, timeToExercise)
    if (isa(result.value,Exception) == true)
        return result
    end
    lattice_price_array = result.value

    # compute the intrinsic value array -
    result = _build_ternary_lattice_intrinsic_value_array(contractSet, lattice_price_array)
    if (isa(result.value,Exception) == true)
        return result
    end
    iv_array = result.value

    # ok, let's build the option value array -
    result = _build_ternary_lattice_option_value_array(iv_array, latticeModel; earlyExcercise = earlyExercise)
    if (isa(result.value,Exception) == true)
        return result
    end
    cost_calc_tuple = result.value

    # create a named tuple and return the results -
    results_tuple = (lattice_price_array=lattice_price_array, intrinsic_value_array = iv_array, cost_calculation_result=cost_calc_tuple)

    # return -
    return PSResult(results_tuple)
end
# ----------------------------------------------------------------------------------------------------------- #


# This file is a part of RunStatistics.jl, licensed under the MIT License (MIT).

__precompile__()

"""
This package implements the evaluation of the cumulative distribution function of the `Squares test statistic` originally defined in

Frederik Beaujean and Allen Caldwell. *A Test Statistic for Weighted Runs.*

https://arxiv.org/abs/1005.3233

The authors further derived an approximation to be able to compute the cumulative also for large numbers of observations in

Frederik Beaujean and Allen Caldwell. *Is the bump significant? An axion-search example*

https://arxiv.org/abs/1710.06642

Where they renamed the weighted-runs statistic to the `*Squares statistic*`.
This code is based on the original implementation by Frederik Beaujean in c++ and mathematica:

https://github.com/fredRos/runs
"""
module RunStatistics

using Distributions
using QuadGK
using ArgCheck
using Distributions
# using Threads

include("t_obs.jl")
include("partitions.jl")
include("squares.jl")
include("squares_approx.jl")

end # module


# Load modules
using Plots, LaTeXStrings
pyplot()
using ThinFilmsTools
using Optim

# Wavelength range [nm]
λ = 400:1000
# Angle of incidence [degrees]
θ = [5.]
# Polarisation (p->1.0, s->0.0, weighted->between 0.0 and 1.0)
pol = 0.5
beam = PlaneWave(λ, θ; p=pol)

# Indices of refraction
n1 = RIdb.air(beam.λ)
n2 = RIdb.silicon(beam.λ)
n3 = RIdb.glass(beam.λ)

# Define the RI models to use
layers = [
    LayerTMMO(n1), # 1
    ModelFit(:looyenga; N=(ninc=n1, nhost=n2)), # 2
    ModelFit(:looyenga; N=(ninc=n1, nhost=n2)), # 3
    LayerTMMO(n3), # 4
]

# Set the order of the layers (build the system) with the layers
order = [1, # incident medium
         2, 3, 2, 3, 2, 3, 2, 3, # top DBR of the MC
         2, 2, # defect of the MC
         3, 2, 3, 2, 3, 2, 3, 2, # bottom DBR of the MC
         4, # substrate
]

# Absolute reflectance spectrum to fit stored in Utils
Rexp = SpectraDB.fp_pectrum(beam.λ)
# plot(Spectrum1D(), beam.λ, Rexp)
# gui()

options = Optim.Options(
    g_abstol=1e-8, g_reltol=1e-8, iterations=10^5, show_trace=true, store_trace=true,
);

# Seeds for each ModelFit layer defined above (No alpha)
seed = [
    [119.0, 0.8], # layers[2]
    [76.0, 0.5], # layers[3]
]

solOptim = fit_tmm_optics(
    Reflectance(Rexp), seed, beam, layers; order=order, options=options, alg=LBFGS(),
)

plot(FitSpectrum(), solOptim.beam.λ, solOptim.spectrumExp, solOptim.spectrumFit)
gui()

# Seeds for each ModelFit layer defined above plus the alpha (with alpha)
seed2 = [
    [120, 0.8],
    [83.0, 0.5],
    [0.995], # alpha
]

solOptim2 = fit_tmm_optics(
    Reflectance(Rexp), seed2, beam, layers;
    order=order, options=options, alg=LBFGS(), alpha=true,
)

plot(FitSpectrum(), solOptim2.beam.λ, solOptim2.spectrumExp, solOptim2.spectrumFit)
gui()


# Type definition
mutable struct TaylorSeries{T, S<:Real} <: FunctionSeries{T}
    coeffs::Vector{T}
    offset::S
    function TaylorSeries(coeffs::Vector{T}, offset = 0.) where {T}
        @assert length(coeffs) >= 1
        S = typeof(offset)
        new{T,S}(coeffs, offset)
    end
end

# Basic properties
offset(f::TaylorSeries) = f.offset
domain(f::TaylorSeries) = (-Inf,+Inf)

degree(f::TaylorSeries) = length(f.coeffs) - 1
coefficients(f::TaylorSeries) = f.coeffs

function shift!(f::TaylorSeries{T}, b) where T
    a = offset(f)
    a == b && return f
    for k = 0:degree(f)
        for l = degree(f):-1:k+1
            if T <: AbstractArray
                axpy!((b-a)^(l-k)*binomial(l,k), f[l], f[k])
            else
                f[k] += f[l]*(b-a)^(l-k)*binomial(l,k)
            end
        end
    end
    f.offset = b
    return f
end
shift(f::TaylorSeries, b) = b == offset(f) ? f : shift!(copy(f), b)

function Base.:(==)(f1::TaylorSeries, f2::TaylorSeries)
    offset(f1) == offset(f2) || return (f1 == shift(f2, offset(f1)))
    K = max(degree(f1), degree(f1))
    for k = 0:K
        f1[k] == f2[k] || return false
    end
    return true
end

# Indexing, getting and setting coefficients
Base.eachindex(f::TaylorSeries) = 0:degree(f)
function Base.getindex(f::TaylorSeries, k)
    k < 0 && throw(BoundsError(f, k))
    @inbounds if k > degree(f)
        return zero(f.coeffs[1])
    else
        return f.coeffs[k+1]
    end
end
function Base.setindex!(f::TaylorSeries, v, k)
    k < 0 && return BoundsError(f, k)
    @inbounds if k > degree(f)
        f0 = f[0]
        while degree(f) < k-1
            push!(f.coeffs, zero(f0))
        end
        push!(f.coeffs, v)
        return v
    else
        setindex!(f.coeffs, v, k+1)
        return v
    end
end

# Use as function
function (f::TaylorSeries)(x::Real)
    k = degree(f)
    xa = x - offset(f)
    (xa == zero(xa) || k == 0) && return copy(f[0])
    v = copy(f[k])
    if v isa AbstractArray
        while k > 0
            v = LinearAlgebra.axpby!(one(xa), f[k-1], xa, v)
            k -= 1
        end
    else
        while k > 0
            v = muladd(xa, v, f[k-1])
            k -= 1
        end
    end
    return v
end

# Change number of coefficients
function truncate!(f::TaylorSeries; Kmax::Integer = degree(f), tol::Real = 0, dx::Real = 1.)
    ktol = findlast(k->(norm(f[k])*abs(dx)^k >= tol), eachindex(f))
    Ktol = ktol === nothing ? 0 : ktol-1
    Kmax = min(Kmax, Ktol)
    if Kmax < degree(f)
        resize!(f.coeffs, Kmax + 1)
    end
    return f
end

function setdegree!(f::TaylorSeries, K::Int)
    while degree(f) < K
        push!(f.coeffs, zero(f[0]))
    end
    if degree(f) > K
        resize!(f.coeffs, K+1)
    end
    return f
end

# special purpose constructor
Base.similar(f::TaylorSeries, ::Type{T} = scalartype(f)) where {T} =
    TaylorSeries([zero(T)*f[0]], offset(f))

Base.zero(f::TaylorSeries) = TaylorSeries([zero(f[0])], offset(f))
Base.one(f::TaylorSeries) = TaylorSeries([one(f[0])], offset(f))

# Arithmetic (out of place)
Base.copy(f::TaylorSeries) = TaylorSeries(copy.(coefficients(f)), offset(f))

Base.:-(f::TaylorSeries) = TaylorSeries(.-coefficients(f), offset(f))
Base.:+(f::TaylorSeries) = TaylorSeries(.+coefficients(f), offset(f))

Base.:*(f::TaylorSeries, a::Const) =
    TaylorSeries([c*a for c in coefficients(f)], offset(f))
Base.:*(a::Const, f::TaylorSeries) =
    TaylorSeries([a*c for c in coefficients(f)], offset(f))
Base.:/(f::TaylorSeries, a::Const) =
    TaylorSeries([c/a for c in coefficients(f)], offset(f))
Base.:\(a::Const, f::TaylorSeries) =
    TaylorSeries([a\c for c in coefficients(f)], offset(f))

function Base.:+(f1::TaylorSeries, f2::TaylorSeries)
    offset(f1) == offset(f2) || return (f1 + shift(f2, offset(f1)))
    K = min(degree(f1), degree(f2))
    coeffs = [f1[j]+f2[j] for j = 0:K]
    for j = K+1:degree(f1)
        push!(coeffs, copy(f1[j]))
    end
    for j = K+1:degree(f2)
        push!(coeffs, copy(f2[j]))
    end
    return TaylorSeries(coeffs, offset(f1))
end

function Base.:-(f1::TaylorSeries, f2::TaylorSeries)
    offset(f1) == offset(f2) || return (f1 - shift(f2, offset(f1)))
    K = min(degree(f1), degree(f2))
    coeffs = [f1[j]-f2[j] for j = 0:K]
    for j = K+1:degree(f1)
        push!(coeffs, copy(f1[j]))
    end
    for j = K+1:degree(f2)
        push!(coeffs, -(f2[j]))
    end
    return TaylorSeries(coeffs, offset(f1))
end

Base.:*(f1::TaylorSeries, f2::TaylorSeries) = truncmul(f1, f2)

function truncmul(f1::TaylorSeries, f2::TaylorSeries; kwargs...)
    if offset(f1) != offset(f2)
        _truncmul(f1, shift(f2, offset(f1)); kwargs...)
    else
        _truncmul(f1, f2; kwargs...)
    end
end

function _truncmul(f1::TaylorSeries, f2::TaylorSeries;
                    Kmax = degree(f1) + degree(f2), tol::Real = 0, dx = 1)
    K = min(Kmax, degree(f1) + degree(f2))
    f = TaylorSeries(sizehint!([zero(f1[0])*zero(f2[0])], K+1), offset(f1))
    return truncmul!(f, f1, f2, true, true; Kmax = Kmax, tol = tol, dx = dx)
end

Base.conj(f::TaylorSeries) = TaylorSeries(map(conj, coefficients(f)), offset(f))
Base.adjoint(f::TaylorSeries) = TaylorSeries(map(adjoint, coefficients(f)), offset(f))
Base.transpose(f::TaylorSeries) = TaylorSeries(map(transpose, coefficients(f)), offset(f))

LinearAlgebra.tr(f::TaylorSeries) = TaylorSeries(map(tr, coefficients(f)), offset(f))

Base.real(f::TaylorSeries) = TaylorSeries(map(real, coefficients(f)), offset(f))
Base.imag(f::TaylorSeries) = TaylorSeries(map(imag, coefficients(f)), offset(f))

# Arithmetic (in place / mutating methods)
function Base.copy!(fdst::TaylorSeries, fsrc::TaylorSeries)
    fdst.offset = offset(fsrc)
    setdegree!(fdst, degree(fsrc))
    for j = 0:degree(fsrc)
        fdst[j] = copy(fsrc[j])
    end
    return fdst
end
function LinearAlgebra.rmul!(f::TaylorSeries, α::Number)
    if eltype(f) <: Number
        rmul!(coefficients(f), α)
    else
        for k in eachindex(f)
            rmul!(f[k], α)
        end
    end
    return f
end
function LinearAlgebra.lmul!(α::Number, f::TaylorSeries)
    if eltype(f) <: Number
        lmul!(α, coefficients(f))
    else
        for k in eachindex(f)
            lmul!(α, f[k])
        end
    end
    return f
end
LinearAlgebra.axpy!(α::Number, fx::TaylorSeries, fy::TaylorSeries) =
    truncadd!(fy, fx, α)
LinearAlgebra.axpby!(α::Number, fx::TaylorSeries, β::Number, fy::TaylorSeries) =
    truncadd!(fy, fx, α, β)
LinearAlgebra.mul!(fy::TaylorSeries, s::Number, fx::TaylorSeries, α = true, β = false) =
    truncmul!(fy, s, fx, α, β)
LinearAlgebra.mul!(fy::TaylorSeries, fx::TaylorSeries, s::Number, α = true, β = false) =
    truncmul!(fy, fx, s, α, β)
LinearAlgebra.mul!(f::TaylorSeries, f1::TaylorSeries, f2::TaylorSeries,
                    α = true, β = false) = truncmul!(f, f1, f2, α, β)

function truncadd!(fy::TaylorSeries, fx::TaylorSeries, α = true, β = true;
                    Kmax::Integer = max(iszero(β) ? 0 : degree(fy), degree(fx)),
                    tol::Real = 0, dx = 1)
    offset(fx) == offset(fy) || throw(DomainMismatch())
    Kx = degree(fx)
    Ky = min(Kmax, iszero(β) ? Kx : max(Kx, degree(fy)))
    setdegree!(fy, Ky)
    # if offset(fx) == offset(fy)
        if eltype(fy) <: Number
            if Kx > Ky
                axpby!(α, view(coefficients(fx, 1:Ky+1)), β, coefficients(fy))
            else
                axpby!(α, coefficients(fx), β, view(coefficients(fy), 1:Kx+1))
                lmul!(β, view(coefficients(fy), (Kx+2):(Ky+1)))
            end
        else
            for k in 0:Kx
                axpby!(α, fx[k], β, fy[k])
            end
            for k in Kx+1:Ky
                lmul!(β, fy[k])
            end
        end
    # else
    #     dx = offset(fy) - offset(fx)
    #     lmul!(β, fy)
    #     for k = 0:Kx
    #         for l = 0:min(k, Ky)
    #             if eltype(fy) <: Number
    #                 fy[l] += α*fx[k]*dx^(l-k)*binomial(l,k)
    #             else
    #                 axpy!(α*dx^(l-k)*binomial(l,k), fx[k], fy[l])
    #             end
    #         end
    #     end
    # end
    return truncate!(fy; tol = tol, dx = dx)
end

truncmul!(fdst::TaylorSeries, α₁::Number, fsrc::TaylorSeries, α₂ = true, β = false;
            kwargs...) = truncadd!(fdst, fsrc, α₁ * α₂, β; kwargs...)
truncmul!(fdst::TaylorSeries, fsrc::TaylorSeries, α₁::Number, α₂ = true, β = false;
            kwargs...) = truncadd!(fdst, fsrc, α₁ * α₂, β; kwargs...)

function truncmul!(F::TaylorSeries, F1::TaylorSeries, F2::TaylorSeries,
                    α = true, β = false;
                    Kmax::Integer = max(iszero(β) ? 0 : degree(F), degree(F1)+degree(F2)),
                    tol::Real = 0,
                    dx = 1)
    offset(F) == offset(F1) == offset(F2) || throw(DomainMismatch())
    K1 = degree(F1)
    K2 = degree(F2)
    K = min(Kmax, max(iszero(β) ? 0 : degree(F), K1+K2))
    setdegree!(F, K)
    Threads.@threads for k = 0:K
        fk = F[k]
        if fk isa AbstractArray
            T = eltype(fk)
            if iszero(β)
                fill!(fk, T(β))
            elseif !isone(β)
                rmul!(fk, T(β))
            end
            for k1 = max(0, k-K2):min(K1, k)
                k2 = k - k1
                fk1 = F1[k1]
                fk2 = F2[k2]
                if fk1 isa AbstractArray && fk2 isa AbstractArray
                    mul!(fk, fk1, fk2, T(α), T(true))
                elseif fk1 isa AbstractArray
                    axpy!(T(fk2*α), fk1, fk)
                elseif fk2 isa AbstractArray
                    axpy!(T(fk1*α), fk2, fk)
                else
                    @warn "unexpected branch"
                    fk .= fk + (fk1*fk2*α)
                end
            end
        else
            if iszero(β)
                fk = zero(fk)
            elseif !isone(β)
                fk *= β
            end
            for k1 = max(0, k-K2):min(K1, k)
                k2 = k - k1
                fk1 = F1[k1]
                fk2 = F2[k2]
                fk += fk1*fk2*α
            end
            F[k] = fk
        end
    end
    if tol != 0
        truncate!(F; tol = tol, dx = dx)
    end
    return F
end

# Inner product and norm
function localdot(f1::TaylorSeries, f2::TaylorSeries)
    offset(f1) == offset(f2) || localdot(f1, shift(f2, offset(f1)))
    K1 = degree(f1)
    K2 = degree(f2)
    coeffs = let K = K1 + K2
        [sum(dot(f1[k1], f2[k-k1]) for k1 = max(0, k-K2):min(k, K1)) for k=0:K]
    end
    return TaylorSeries(coeffs, offset(f1))
end

function differentiate(f::TaylorSeries)
    if degree(f) >= 1
        return TaylorSeries([k*f[k] for k = 1:degree(f)], offset(f))
    else
        return zero(f)
    end
end
function integrate(f::TaylorSeries, (a,b)::Tuple{Real,Real})
    s = f[0]*(b-a)
    c = offset(f)
    db = b - c
    da = a - c
    for k = 1:degree(f)
        s += f[k]*(db^(k+1) - da^(k+1))/(k+1)
    end
    return s
end

# Simple Taylor series least squares fitting:
# typically not needed for large number of points,
# but should work with matrix valued functions etc
# => generality over efficiency
function fit(f, ::Type{TaylorSeries}, domain::Tuple{Real,Real};
                Kmax = 10, numpoints = Kmax+1, offset = (domain[1]+domain[2])/2)

    a, b = domain
    x = range(a, stop = b, length = numpoints)
    dx = x .- offset
    coeffs = similar(x, (numpoints, Kmax+1))
    coeffs[:, 1] .= 1
    for k = 1:Kmax
        coeffs[:, k+1] .= view(coeffs, :, k) .* dx
    end
    TaylorSeries(pinv(coeffs) * f.(x), offset)
end

# Inverse (and square root)
# Warning: experimental
function Base.inv(t::TaylorSeries; tol = 1e-10, Kmax = 10*degree(t))
    t0inv = inv(t[0])
    coeffs = [t0inv]
    K = degree(t)
    for k = 1:Kmax
        invtk = -t0inv*sum(coeffs[i+1]*t[k-i] for i = max(0,k-K):(k-1))
        if norm(invtk) < tol && norm(coeffs[end]) < tol
            break
        end
        push!(coeffs, invtk)
    end
    return TaylorSeries(coeffs, offset(t))
end