text:
"""
Functions to estimate Q|_{U^0}. See our paper for details.
"""

using LinearAlgebra
using SparseArrays
using FastRounding
using ValidatedNumerics
using FFTW

export norms_of_powers_noise

"""
Estimates the norms ||Q||, ||Q^2||, ... ||Q^m|| on U^0.

U is the matrix [ones(1,n-1); -I_(n-1,n-1)]. It is currently assumed that
f*U==0 (i.e., all elements of f are equal).

f must be an interval vector.

The following constants may be specified as keyword arguments:

normQ, normE, normv0, normEF, normIEF, normN

otherwise they are computed (which may be slower).

e and f must be specified in case is_integral_preserving==false
In case is_integral_preserving is true, they may be specified but they are then ignored.
"""
function norms_of_powers_noise(N::Type{L1},
        m::Integer,
        Q::DiscretizedOperator,
        MK::NoiseKernel,
        f::AbstractArray;
        normv0::Real=-1., #used as "missing" value
        normQ::Real=-1.,
        normE::Real=-1.,
        normEF::Real=-1.,
        normIEF::Real=-1.,
        normN::Real=-1.,
        normρ::Real=-1.)

    @assert eltype(f) <: Interval
    T = typeof(zero(eltype(Q.L)).hi) # gets "Float64" from Q.L
    n = size(Q.L, 1)
    M = mid.(Q.L)

    R = radius.(Q.L)
    δ = opnormbound(N, R)

    # this is the operator radius of the noise operator,
    # we use this function since the operator may be defined implictly

    γz = gamma(T, max_nonzeros_per_row(Q.L))
    γn = gamma(T, n+3) # not n+2 like in the paper, because we wish to allow for f to be the result of rounding

    ϵ = zero(T)

    nrmM = opnormbound(N, M)

    # these need to be implemented for each kernel
    δₖ = opradius(N, MK) #this is the opnorm of the radius matrix

    # essentially nonzero_per_row(MK) is
    # the  number of mul_add needed to obtain one
    # entry of MK*v
    γₖ = gamma(T, nonzero_per_row(MK))

    nrm_MK = opnormbound(N, MK)
    normMK = nrm_MK ⊕₊ δₖ

    # precompute norms
    if !is_integral_preserving(Q)
        if normE == -1.
            normE = opnormbound(N, Q.e)
        end
        if normEF == -1.
            normEF = opnormbound(N, Q.e*f)
        end
        if normIEF == -1.
            normIEF =  opnormbound(N, [Matrix(UniformScaling{Float64}(1),n,n) Q.e*f])
        end
        if normN == -1.
            normN = opnormbound(N, Matrix(UniformScaling{Float64}(1),n,n) - Q.e*f)
        end
    end

    if normQ == -1.
        if is_integral_preserving(Q)
            normQ = nrmM ⊕₊ δ
        else
            defect = opnormbound(N, Q.w)
            normQ = nrmM ⊕₊ δ ⊕₊ normE ⊗₊ defect
        end
    end

    # initialize normcachers
    normcachers = [NormCacher{N}(n) for j in 1:m]

    midf = map(mid, f)

    # main loop

    v = zeros(T, n)

    @showprogress for j = 1:n-1
        v .= zero(T) # TODO: check for type stability in cases with unusual types
        v[1] = one(T) # TODO: in full generality, this should contain entries of f rather than ±1
        v[j+1] = -one(T)
        if normv0 == -1.
            nrmv = opnormbound(N, v)
        else
            nrmv = normv0
        end
        ϵ = 0.
        nrmw = nrmv # we assume that initial vectors are already integral-preserving
        for k = 1:m
            w = M * v
            if is_integral_preserving(Q)
                v = w
                ϵ = (γz ⊗₊ nrmM ⊕₊ δ) ⊗₊ nrmv ⊕₊ normQ ⊗₊ ϵ
            else
                v = w - Q.e * (midf*w)  # TODO: we are currently assuming that f is not too large, to estimate the error (the result of only one floating point operation)
                new_nrmw = opnormbound(N, w)
                ϵ = γn ⊗₊ normIEF ⊗₊ (new_nrmw ⊕₊ normEF ⊗₊ nrmw) ⊕₊ normN ⊗₊ (γz ⊗₊ nrmM ⊕₊ δ) ⊗₊ nrmv ⊕₊ normQ ⊗₊ ϵ
                nrmw = new_nrmw
            end

            # the noise step
            nrmv = opnormbound(N, v)

            v = w
            w = MK * v
            v = w

            ϵ = (γₖ ⊗₊nrm_MK ⊕₊δₖ) ⊗₊ nrmv ⊕₊ normMK ⊗₊ ϵ

            nrmv = opnormbound(N, v)
            add_column!(normcachers[k], v, ϵ) #TODO: Could pass and reuse nrmv in the case of norm-1
        end
    end
    return map(get_norm, normcachers)
end

"""
Array of "trivial" bounds for the powers of a DiscretizedOperator (on the whole space)
coming from from ||Q^k|| ≤ ||Q||^k
"""
function norms_of_powers_trivial_noise(N::Type{<:NormKind},
                                       Q::DiscretizedOperator,
                                       MK::NoiseKernel,
                                       m::Integer)
    norms = fill(NaN, m)
    δₖ = opradius(N, MK)
    nrm_MK = opnormbound(N, MK)
    norms[1] = opnormbound(N, Q) ⊗₊ nrm_MK⊕₊δₖ
    for i = 2:m
        norms[i] = norms[i-1] ⊗₀ norms[1]
    end
    return norms
end

"""
Arrays of bounds to ||Q^k||_{w → s} = sup_{||f||_w=1} ||Q^k f||_s
and to ||Q^k||_{w}
coming theoretically from iterated DFLY inequalities (the "small matrix method").

Returns two arrays (strongs, norms) of length m:
strongs[k] bounds ||Q^k f||_s, norms[k] bounds ||Q^k f||)
"""
function norms_of_powers_abstract_noise(Bas::Basis, N::NoiseKernel, m)
    A, B = dfly(strong_norm(Bas), aux_norm(Bas), N)
    Eh = BasisDefinition.aux_normalized_projection_error(Bas)
    M₁n = BasisDefinition.strong_weak_bound(Bas)
    M₂ = BasisDefinition.aux_weak_bound(Bas)
    S₁, S₂ = BasisDefinition.weak_by_strong_and_aux_bound(Bas)

    norms = fill(NaN, m)
    strongs = fill(NaN, m)

    v = Array{Float64}([M₁n; M₂])
    # We evaluate [S₁ S₂] * ([1 0; Eh 1]*[A B; 0 1])^k * [M₁n; M₂] (with correct rounding)
    for k = 1:m
        # invariant: v[1] bounds ||Q^kf||_s for ||f||_w=1
        # v[2] bounds |||Q^kf||| for ||f||_w=1
        v[1] = A ⊗₊ v[1] ⊕₊ B ⊗₊ v[2]
        v[2] = Eh ⊗₊ v[1] ⊕₊ v[2]
        strongs[k] = v[1]
        norms[k] = S₁ ⊗₊ v[1] ⊕₊ S₂ ⊗₊ v[2]
    end
    return strongs, norms
end

"""
Estimate norms of powers from those on a coarser grid (see paper for details)
TODO: Check if it works for other basis types
"""
function norms_of_powers_from_coarser_grid_noise(fine_basis::Ulam,
                                            coarse_basis::Ulam,
                                            Q::DiscretizedOperator,
                                            NK::NoiseKernel,
                                            coarse_norms::Vector)
    if !BasisDefinition.is_refinement(fine_basis, coarse_basis)
        @error "The fine basis is not a refinement of the coarse basis"
    end
    m = length(coarse_norms)
    fine_norms = fill(NaN, m+1)
    trivial_norms = norms_of_powers_trivial_noise(weak_norm(fine_basis), Q, NK, m+1)

    A, B = dfly(strong_norm(fine_basis), aux_norm(fine_basis), NK)

    # adds a 0th element to strongs
    trivial_norms0(k::Integer) = k==0 ? 1. : trivial_norms[k]
    coarse_norms0(k::Integer) = k==0 ? 1. : coarse_norms[k]

    Kh =  BasisDefinition.weak_projection_error(coarse_basis)

    fine_norms[1] = trivial_norms0(1)

    for k in 1:m
		temp = 0.
		for i in 0:k-1
			temp = temp ⊕₊ coarse_norms0(i) ⊗₊ (trivial_norms0(k-i) ⊕₊ trivial_norms0(k-i-1))
		end
        fine_norms[k+1] = coarse_norms0(k) ⊕₊ B ⊗₊ Kh ⊗₊ (temp ⊕₊ 1.0)
	end
    return fine_norms
end

function norms_of_powers_from_coarser_grid_noise_abstract(coarse_basis::Ulam,
                                                        NK::NoiseKernel,
                                                        coarse_norms::Vector)

    m = length(coarse_norms)
    abstract_norms = fill(NaN, m+1)
    trivial_norms = ones(m+1)

    A, B = dfly(strong_norm(coarse_basis), aux_norm(coarse_basis), NK)

    # adds a 0th element to strongs
    trivial_norms0(k::Integer) = k==0 ? 1. : trivial_norms[k]
    coarse_norms0(k::Integer) = k==0 ? 1. : coarse_norms[k]

    Kh =  BasisDefinition.weak_projection_error(coarse_basis)

    fine_norms = fill(NaN, m+1)
    fine_norms[1] = trivial_norms0(1)

    for k in 1:m
        temp = 0.
        for i in 0:k-1
            temp = temp ⊕₊ coarse_norms0(i) ⊗₊ (trivial_norms0(k-i) ⊕₊ trivial_norms0(k-i-1))
        end
    fine_norms[k+1] = coarse_norms0(k) ⊕₊ B ⊗₊ Kh ⊗₊ (temp ⊕₊ 1.0)
    end
    return fine_norms
end

function powernormboundsnoise(B; Q=DiscretizedOperator(B, D), NK = NK::NoiseKernel)
	m = 8
	computed_norms = []
	while true
		computed_norms = norms_of_powers_noise(weak_norm(B), m, Q, NK, integral_covector(B))
		if any(computed_norms .< 0.1)
			break
		end
		m = 2*m
	end
	trivial_norms = norms_of_powers_trivial_noise(weak_norm(B), Q, NK, m)
	norms = min.(trivial_norms, computed_norms)

	m_extend = 2*m
	better_norms = []
	while true
		better_norms = refine_norms_of_powers(norms, m_extend)
		if better_norms[end] < 1e-8
			break
		end
		m_extend = 2*m_extend
	end

	return better_norms

end

"""
Uses power norm bounds already computed for a coarse operator to estimate
the same norms for a finer operator
"""
function finepowernormboundsnoise(B,
                                  Bfine,
                                  coarse_norms;
                                  Qfine::DiscretizedOperator,
                                  NKfine::NoiseKernel)
	                              m = length(coarse_norms)

	trivial_norms_fine = norms_of_powers_trivial_noise(weak_norm(B), Qfine, NKfine, m+1)
	twogrid_norms_fine = norms_of_powers_from_coarser_grid_noise(Bfine,
                                                                 B,
                                                                 Qfine,
                                                                 NKfine,
                                                                 coarse_norms)

	norms_fine = min.(trivial_norms_fine, twogrid_norms_fine)

	better_norms_fine = refine_norms_of_powers(norms_fine, m+1)
	return better_norms_fine
end

function abstractpowernormboundsnoise(B, NK, coarse_norms; m = length(coarse_norms))
    trivial_norms_fine = ones(m+1)
    abstract_norms_from_coarse = norms_of_powers_from_coarser_grid_noise_abstract(B, NK, coarse_norms)

    norms_fine = min.(trivial_norms_fine, abstract_norms_from_coarse)
    better_norms_abstract = refine_norms_of_powers(norms_fine, m+1)
    return better_norms_abstract
end

"""
Return a numerical approximation to the (hopefully unique) invariant vector
of the dynamic with discretized operator Q.

The vector is normalized so that integral_covector(B)*w ≈ 1
"""
function invariant_vector_noise(B::Basis, Q::DiscretizedOperator, NK::NoiseKernel; tol = 0.0, iter =10)
	mQ = mid(Q)
	v = one_vector(B)
    for _ in 1:iter
        v = mQ*v
        v = NK*v
    end
    v = v ./ (mid.(integral_covector(B))*v)
    return v
end

"""
Return an upper bound to Q_h*w - w in the given norm
"""
function residualboundnoise(N::Type{<:NormKind}, Q::DiscretizedOperator, NK::NoiseKernel, w::AbstractVector)
	v = Q*w
    v = NK*v # actually, a priori estimate... TODO modify mult NK*v
    return normbound(N, v - w)
end

"""
Bounds rigorously the distance of w from the fixed point of Q (normalized with integral = 1),
using a vector of bounds norms[k] ≥ ||Q_h^k|_{U_h^0}||.
"""
function distance_from_invariant_noise(B::Basis,
                                 Q::DiscretizedOperator,
                                 NK::NoiseKernel,
                                 w::AbstractVector,
                                 norms::Vector;
                                 ε₁::Float64 = residualboundnoise(weak_norm(B), Q, NK, w),
                                 ε₂::Float64 = mag(integral_covector(B) * w - 1),
                                 normQ::Float64 = opnormbound(weak_norm(B), Q))
	if ε₂ > 1e-8
		@error "w does not seem normalized correctly"
	end
    _ , us = dfly(strong_norm(B), aux_norm(B), NK)
	Cs = infinite_sum_norms(norms)
	Kh =  BasisDefinition.weak_projection_error(B)
	normw = normbound(weak_norm(B), w)

	return Cs ⊗₊ (2. ⊗₊ Kh ⊗₊ (1. ⊕₊ normQ) ⊗₊ us ⊕₊ ε₁ ⊘₊ (1. ⊖₋ ε₂)) ⊕₊ ε₂ ⊘₊ (1. ⊖₋ ε₂) ⊗₊ normw
end


text:
const UCR_DATASETS = """
50words,
Adiac,
ArrowHead,
Beef,
BeetleFly,
BirdChicken,
CBF,
Car,
ChlorineConcentration,
CinC_ECG_torso,
Coffee,
Computers,
Cricket_X,
Cricket_Y,
Cricket_Z,
DiatomSizeReduction,
DistalPhalanxOutlineAgeGroup,
DistalPhalanxOutlineCorrect,
DistalPhalanxTW,
ECG200,
ECG5000,
ECGFiveDays,
Earthquakes,
ElectricDevices,
FISH,
FaceAll,
FaceFour,
FacesUCR,
FordA,
FordB,
Gun_Point,
Ham,
HandOutlines,
Haptics,
Herring,
InlineSkate,
InsectWingbeatSound,
ItalyPowerDemand,
LargeKitchenAppliances,
Lighting2,
Lighting7,
MALLAT,
Meat,
MedicalImages,
MiddlePhalanxOutlineAgeGroup,
MiddlePhalanxOutlineCorrect,
MiddlePhalanxTW,
MoteStrain,
NonInvasiveFatalECG_Thorax1,
NonInvasiveFatalECG_Thorax2,
OSULeaf,
OliveOil,
PhalangesOutlinesCorrect,
Phoneme,
Plane,
ProximalPhalanxOutlineAgeGroup,
ProximalPhalanxOutlineCorrect,
ProximalPhalanxTW,
RefrigerationDevices,
ScreenType,
ShapeletSim,
ShapesAll,
SmallKitchenAppliances,
SonyAIBORobotSurface,
SonyAIBORobotSurfaceII,
StarLightCurves,
Strawberry,
SwedishLeaf,
Symbols,
ToeSegmentation1,
ToeSegmentation2,
Trace,
TwoLeadECG,
Two_Patterns,
UWaveGestureLibraryAll,
Wine,
WordsSynonyms,
Worms,
WormsTwoClass,
synthetic_control,
uWaveGestureLibrary_X,
uWaveGestureLibrary_Y,
uWaveGestureLibrary_Z,
wafer,
yoga
"""

"""
    TimeWarp.download_data([path="/TimeWarp/data/"])

Downloads the UC Riverside Time Series Classification Archive to
the specified path.
"""
function download_ucr()

    UCR_URL = "http://www.cs.ucr.edu/~eamonn/time_series_data/UCR_TS_Archive_2015.zip"

    # the password is "attempttoclassify"
    info(
        """
        A simple password is required to extract the UCR Time Series Classification. While the
        data is downloading, please read the following paper

        Bing Hu, Yanping Chen, Eamonn J. Keogh: Time Series Classification under More Realistic
        Assumptions. SDM 2013: 578-586.
        http://dx.doi.org/10.1137/1.9781611972832.64

        Find the following sentence (it's near the end of the first page):

        “Every item that we ******* ## @@@@@@@ belongs to exactly one of our well defined classes”

        Enter the three redacted words (without spaces) to extract the data. The purpose of this
        exercise is to encourage you to read the associated paper, see:

        http://www.cs.ucr.edu/~eamonn/time_series_data/

        """
    )

    download(UCR_URL, "$DATAPATH/UCR.zip")

    # should unzip the file correctly in platform-specific manner
    run(unpack_cmd("UCR", DATAPATH , ".zip", ""))

    # delete the zip file when we're done
    rm("$DATAPATH/UCR.zip")

    info("Download and extraction successful!")

end

"""
    data,labels = TimeWarp.traindata(name)

Loads the training set of the specified dataset. Returns a matrix `data` where each column
holds a 1-dimensional time series. The class label for each column is held `labels`
which is a vector of length `size(data,2)`.

Available datasets:

$UCR_DATASETS
"""
function ucr_traindata(name::AbstractString)
    try
        Y = readcsv(DATAPATH*"/UCR_TS_Archive_2015/"*name*"/"*name*"_TRAIN")
        labels = round(Int,vec(Y[:,1]))
        data = transpose(Y[:,2:end])
        return data,labels
    catch err
        showerror(STDOUT, err, backtrace());println()
        info("You may have recieved this error because you haven't downloaded the database yet.")
        info("Try running download_ucr() first.")
        info("You may also have mispelled the name of the dataset.")
    end
end

"""
    data,labels = TimeWarp.traindata(name)

Loads the test set of the specified dataset. Returns a matrix `data` where each column
holds a 1-dimensional time series. The class label for each column is held `labels`
which is a vector of length `size(data,2)`.

Available datasets:

$UCR_DATASETS
"""
function ucr_testdata(name::AbstractString)
    try
        Y = return readcsv(DATAPATH*"/UCR_TS_Archive_2015/"*name*"/"*name*"_TEST")
        labels = round(Int,vec(Y[:,1]))
        data = transpose(Y[:,2:end])
        return data,labels
    catch err
        showerror(STDOUT, err, backtrace());println()
        info("You may have recieved this error because you haven't downloaded the database yet.")
        info("Try running download_ucr() first.")
        info("You may also have mispelled the name of the dataset.")
    end
end

ext: jl


text:
"""
	BafflingDriver
Driver that randomly changes lanes and speeds.

# Constructors
	BafflingDriver(timestep::Float64;mlon::LaneFollowingDriver=IntelligentDriverModel(), mlat::LateralDriverModel=ProportionalLaneTracker(), mlane::LaneChangeModel=RandLaneChanger(timestep),rec::SceneRecord = SceneRecord(1, timestep))

# Fields
- `rec::SceneRecord` A record that will hold the resulting simulation results
- `mlon::LaneFollowingDriver = IntelligentDriverModel()` Longitudinal driving model
- `mlat::LateralDriverModel = ProportionalLaneTracker()` Lateral driving model
- `mlane::LaneChangeModel =RandLaneChanger` Lane change model (randomly)
"""
mutable struct BafflingDriver <: DriverModel{LatLonAccel}
    rec::SceneRecord
    mlon::LaneFollowingDriver
    mlat::LateralDriverModel
    mlane::LaneChangeModel
    ΔT::Float64

    function BafflingDriver(
        timestep::Float64;
        mlon::LaneFollowingDriver=BafflingLongitudinalTracker(),
        mlat::LateralDriverModel=ProportionalLaneTracker(),
        mlane::LaneChangeModel=BafflingLaneChanger(timestep),
        rec::SceneRecord = SceneRecord(1, timestep)
        )

        retval = new()

        retval.rec = rec
        retval.mlon = mlon
        retval.mlat = mlat
        retval.mlane = mlane
        retval.ΔT = timestep

        retval
    end
end
get_name(::BafflingDriver) = "BafflingDriver"
function AutomotiveDrivingModels.set_desired_speed!(model::BafflingDriver, v_des::Float64)
    set_desired_speed!(model.mlon, v_des)
    set_desired_speed!(model.mlane, v_des)
    model
end
# function track_longitudinal!(driver::LaneFollowingDriver, scene::Scene, roadway::Roadway, vehicle_index::Int, fore::NeighborLongitudinalResult)
#     v_ego = scene[vehicle_index].state.v
#     if fore.ind != nothing
#         headway, v_oth = fore.Δs, scene[fore.ind].state.v
#     else
#         headway, v_oth = NaN, NaN
#     end
#     return track_longitudinal!(driver, v_ego, v_oth, headway)
# end
function AutomotiveDrivingModels.observe!(driver::BafflingDriver, scene::Scene, roadway::Roadway, egoid::Int)

    update!(driver.rec, scene)
    observe!(driver.mlane, scene, roadway, egoid) # receive action from the lane change controller

    vehicle_index = findfirst(egoid, scene)
    lane_change_action = rand(driver.mlane)
    laneoffset = get_lane_offset(lane_change_action, driver.rec, roadway, vehicle_index)
    lateral_speed = convert(Float64, get(VELFT, driver.rec, roadway, vehicle_index))

    if lane_change_action.dir == DIR_MIDDLE
        fore = get_neighbor_fore_along_lane(scene, vehicle_index, roadway, VehicleTargetPointFront(), VehicleTargetPointRear(), VehicleTargetPointFront())
    elseif lane_change_action.dir == DIR_LEFT
        fore = get_neighbor_fore_along_left_lane(scene, vehicle_index, roadway, VehicleTargetPointFront(), VehicleTargetPointRear(), VehicleTargetPointFront())
    else
        @assert(lane_change_action.dir == DIR_RIGHT)
        fore = get_neighbor_fore_along_right_lane(scene, vehicle_index, roadway, VehicleTargetPointFront(), VehicleTargetPointRear(), VehicleTargetPointFront())
    end

    track_longitudinal!(driver.mlon, scene, roadway, vehicle_index, fore) # receive acceleration from the longitudinal controller
    track_lateral!(driver.mlat, laneoffset, lateral_speed) # receive acceleration from the lateral controller

    driver
end
Base.rand(driver::BafflingDriver) = LatLonAccel(rand(driver.mlat), rand(driver.mlon).a)
Distributions.pdf(driver::BafflingDriver, a::LatLonAccel) = pdf(driver.mlat, a.a_lat) * pdf(driver.mlon, a.a_lon)
Distributions.logpdf(driver::BafflingDriver, a::LatLonAccel) = logpdf(driver.mlat, a.a_lat) * logpdf(driver.mlon, a.a_lon)


text:
"""
    ZeroInflatedPoisson(λ, p)
A *Zero-Inflated Poisson distribution* is a mixture distribution in which data arise from two processes. The first process is is a Poisson distribution, with mean λ, that descibes the number of independent events occurring within a unit time interval:
```math
P(X = k) = (1 - p) \\frac{\\lambda^k}{k!} e^{-\\lambda}, \\quad \\text{ for } k = 0,1,2,\\ldots.
```
Zeros may arise from this process, an additional Bernoulli process, where the probability of observing an excess zero is given as p:
```math
P(X = 0) = p + (1 - p) e^{-\\lambda}
```
As p approaches 0, the distribution tends toward Poisson(λ).
```julia
ZeroInflatedPoisson()        # Zero-Inflated Poisson distribution with rate parameter 1, and probability of observing a zero 0.5
ZeroInflatedPoisson(λ)       # ZeroInflatedPoisson distribution with rate parameter λ, and probability of observing a zero 0.5
params(d)                    # Get the parameters, i.e. (λ, p)
mean(d)                      # Get the mean of the mixture distribution
var(d)                       # Get the variance of the mixture distribution
```
External links:
* [Zero-inflated Poisson Regression on UCLA IDRE Statistical Consulting](https://stats.idre.ucla.edu/stata/dae/zero-inflated-poisson-regression/)
* [Zero-inflated model on Wikipedia](https://en.wikipedia.org/wiki/Zero-inflated_model)
* McElreath, R. (2020). Statistical Rethinking: A Bayesian Course with Examples in R and Stan (2nd ed.). Chapman and Hall/CRC. https://doi.org/10.1201/9780429029608

"""
struct ZeroInflatedPoisson{T<:Real} <: DiscreteUnivariateDistribution
  λ::T
  p::T

  function ZeroInflatedPoisson{T}(λ::T, p::T) where {T <: Real}
      return new{T}(λ, p)
    end
end

function ZeroInflatedPoisson(λ::T, p::T; check_args = true) where {T <: Real}
  if check_args
    @check_args(Poisson, λ >= zero(λ))
    @check_args(ZeroInflatedPoisson, zero(p) <= p <= one(p))
  end
  return ZeroInflatedPoisson{T}(λ, p)
end

ZeroInflatedPoisson(λ::Real, p::Real) = ZeroInflatedPoisson(promote(λ, p)...)
ZeroInflatedPoisson(λ::Integer, p::Integer) = ZeroInflatedPoisson(float(λ), float(p))
ZeroInflatedPoisson(λ::Real) = ZeroInflatedPoisson(λ, 0.0)
ZeroInflatedPoisson() = ZeroInflatedPoisson(1.0, 0.0, check_args = false)

@distr_support ZeroInflatedPoisson 0 (d.λ == zero(typeof(d.λ)) ? 0 : Inf)

### Statistics

mean(d::ZeroInflatedPoisson) = (1 - d.p) * d.λ

var(d::ZeroInflatedPoisson) = d.λ * (1 - d.p) * (1 + d.p * d.λ)

#### Conversions

function convert(::Type{ZeroInflatedPoisson{T}}, λ::Real, p::Real) where {T<:Real}
  return ZeroInflatedPoisson(T(λ), T(p))
end

function convert(::Type{ZeroInflatedPoisson{T}}, d::ZeroInflatedPoisson{S}) where {T <: Real, S <: Real}
  return ZeroInflatedPoisson(T(d.λ), T(d.p), check_args = false)
end

#### Parameters

params(d::ZeroInflatedPoisson) = (d.λ, d.p,)
partype(::ZeroInflatedPoisson{T}) where {T} = T

rate(d::ZeroInflatedPoisson) = d.λ
excessprob(d::ZeroInflatedPoisson) = d.p

#### Evaluation

function logpdf(d::ZeroInflatedPoisson, y::Real)
    lp = if iszero(y)
        logaddexp(log(d.p), log1p(-d.p) - d.λ)
    else
        log1p(-d.p) + logpdf(Poisson(d.λ), y)
    end
    return lp
end

function cdf(d::ZeroInflatedPoisson, x::Real)
  pd = Poisson(d.λ)

  deflat_limit = -1.0 / expm1(d.λ)

  if x < 0
    out = 0.0
  elseif d.p < deflat_limit
    out = NaN
  else
    out = d.p + (1 - d.p) * cdf(pd, x)
  end
  return out
end

# quantile
function quantile(d::ZeroInflatedPoisson, q::Real)

  deflat_limit = -1.0 / expm1(d.λ)

  if (q <= d.p)
    out = 0
  elseif (d.p < deflat_limit)
    out = convert(Int64, NaN)
  elseif (d.p < q) & (deflat_limit <= d.p) & (q < 1.0)
    qp = (q - d.p) / (1.0 - d.p)
    pd = Poisson(d.λ)
    out = quantile(pd, qp) # handles d.p == 1 as InexactError(Inf)
  end
  return out
end

#### Fitting

struct ZeroInflatedPoissonStats <: SufficientStats
  sx::Float64   # (weighted) sum of x
  p0::Float64   # observed proportion of zeros
  tw::Float64   # total sample weight
end

suffstats(::Type{<:ZeroInflatedPoisson}, x::AbstractArray{T}) where {T<:Integer} = ZeroInflatedPoissonStats(
    sum(x),
    mean(iszero, x),
    length(x)
  )

# weighted
function suffstats(::Type{<:ZeroInflatedPoisson}, x::AbstractArray{T}, w::AbstractArray{Float64}) where T<:Integer
    n = length(x)
    n == length(w) || throw(DimensionMismatch("Inconsistent array lengths."))
    sx = 0.
    tw = 0.
    p0 = 0.
    for i = 1 : n
        @inbounds wi = w[i]
        @inbounds sx += x[i] * wi
        tw += wi
        @inbounds p0i = (x[i] == 0) * wi
        p0 += p0i
    end
    return ZeroInflatedPoissonStats(sx, p0, tw)
end

function fit_mle(::Type{<:ZeroInflatedPoisson}, ss::ZeroInflatedPoissonStats)
  m = ss.sx / ss.tw
  s = m / (1 - ss.p0)

  λhat = lambertw(-s * exp(-s)) + s
  phat = 1 - (m / λhat)

  return ZeroInflatedPoisson(λhat, phat)
end

function fit_mle(::Type{<:ZeroInflatedPoisson}, x::AbstractArray{T}) where T<:Real
  pstat = suffstats(ZeroInflatedPoisson, x)
  return fit_mle(ZeroInflatedPoisson, pstat)
end

function fit_mle(::Type{<:ZeroInflatedPoisson}, x::AbstractArray{T}, w::AbstractArray{Float64}) where T<:Real
  pstat = suffstats(ZeroInflatedPoisson, x, w)
  return fit_mle(ZeroInflatedPoisson, pstat)
end


text:
using StructLH, Test

@common_fields set1 begin
    x :: Int
    y :: Float64
end

struct Foo
    @set1
    z
end

function common_fields_test()
    @testset "Common fields macro" begin
        f = Foo(1, 2.2, "abc");
        @test f.x == 1
        @test f.y == 2.2
        @test f.z == "abc"
	end
end

@testset "Helpers" begin
    common_fields_test();
end

# -------------


text:
using Distributed, Test

@testset "AdvancedHMC" begin
    tests = [
        "adaptation/precond",
        "trajectory",
        "hamiltonian",
        "integrator",
        "demo",
        "models",
        "sampler",
        "sampler_mat",
    ]

    res = map(tests) do t
        @eval module $(Symbol("Test_", t))
            include($t * ".jl")
        end
        return
    end
end


text:
#!/usr/bin/env julia

using SimradRaw

# Prints the survey name for one or more RAW files.

function main(args)
    for arg in args
        for datagram in datagrams(arg)
            if isa(datagram, SimradRaw.ConfigurationDatagram)
                println(datagram.configurationheader.surveyname)
                break
            end
        end
    end

end

main(ARGS)


text:
using ReduxImGui
using ReduxImGui.Redux
using Test

@testset "Buttons | Button" begin
    name = "Button_test"
    store = create_store(Buttons.reducer, Buttons.State(name))
    state = get_state(store)

    # test initial state
    @test state.label == name
    @test state.size[1] == 0
    @test state.size[2] == 0
    @test state.is_triggered == false

    # test actions
    new_name = "Button_new_name"
    state = get_state(store)
    dispatch!(store, Buttons.Rename(state.label, new_name))
    @test get_state(store).label == new_name
    dispatch!(store, Buttons.Rename(state.label, name))
    @test Buttons.get_label(get_state(store)) == name

    new_size = (100, 200)
    state = get_state(store)
    dispatch!(store, Buttons.Resize(state.label, new_size...))
    @test get_state(store).size[1] == new_size[1]
    @test get_state(store).size[2] == new_size[2]

    dispatch!(store, Buttons.ChangeWidth(state.label, 7))
    @test get_state(store).size[1] == 7

    dispatch!(store, Buttons.ChangeHeight(state.label, 77))
    @test get_state(store).size[2] == 77

    dispatch!(store, Buttons.SetTriggeredTo(state.label, true))
    @test get_state(store).is_triggered == true

    # test vectors
    store = create_store(Buttons.reducer, [Buttons.State("b1"),Buttons.State("b2")])
    state = get_state(store)
    dispatch!(store, Buttons.AddButton("b3"))
    dispatch!(store, Buttons.DeleteButton("b1"))
    state = get_state(store)
    @test state[1] == Buttons.State("b2")
    @test state[2] == Buttons.State("b3")
    dispatch!(store, Buttons.SetTriggeredTo("b2", false))
    @test Buttons.is_triggered(get_state(store)[1]) == false

    # test dicts
    store = create_store(Buttons.reducer, Dict("b1"=>Buttons.State("b1"),"b2"=>Buttons.State("b2")))
    state = get_state(store)
    dispatch!(store, Buttons.Resize("b2", 10, 20))
    state = get_state(store)
    @test Buttons.get_size(state["b2"]) == (10,20)
end

@testset "Buttons | OnOffButton" begin
    name = "OnOffButton_test"
    store = create_store(OnOffButtons.reducer, OnOffButtons.State(name))
    state = get_state(store)

    # test initial state
    @test OnOffButtons.get_label(state) == name
    @test OnOffButtons.get_size(state) == (0,0)
    @test OnOffButtons.is_triggered(state) == false
    @test OnOffButtons.is_on(state) == false

    # test actions
    new_size = (100, 200)
    state = get_state(store)
    dispatch!(store, OnOffButtons.Toggle(OnOffButtons.get_label(state)))
    state = get_state(store)
    @test OnOffButtons.is_on(state) == true
    dispatch!(store, OnOffButtons.Toggle(OnOffButtons.get_label(state)))
    state = get_state(store)
    @test OnOffButtons.is_on(state) == false

    # test vectors
    store = create_store(OnOffButtons.reducer, [OnOffButtons.State("b1"),OnOffButtons.State("b2")])
    state = get_state(store)
    dispatch!(store, OnOffButtons.AddButton("b3"))
    dispatch!(store, OnOffButtons.DeleteButton("b1"))
    state = get_state(store)
    @test state[1] == OnOffButtons.State("b2")
    @test state[2] == OnOffButtons.State("b3")
    dispatch!(store, OnOffButtons.SetTriggeredTo("b2", false))
    @test OnOffButtons.is_triggered(get_state(store)[1]) == false

    # test dicts
    store = create_store(OnOffButtons.reducer, Dict("b1"=>OnOffButtons.State("b1"),"b2"=>OnOffButtons.State("b2")))
    state = get_state(store)
    dispatch!(store, OnOffButtons.Resize("b2", 10, 20))
    state = get_state(store)
    @test OnOffButtons.get_size(state["b2"]) == (10,20)
end

@testset "Buttons | ColorButton" begin
    name = "ColorButton_test"
    store = create_store(ColorButtons.reducer, ColorButtons.State(name))
    state = get_state(store)

    # test initial state
    @test ColorButtons.get_label(state) == name
    @test ColorButtons.get_size(state) == (0,0)
    @test ColorButtons.is_triggered(state) == false

    # test actions
    new_size = (100, 200)
    state = get_state(store)
    dispatch!(store, ColorButtons.SetButtonColorTo(ColorButtons.get_label(state), 0.0, 0.0, 1.0))
    @test ColorButtons.get_color(get_state(store)) == ReduxImGui.CImGui.HSV(0.0, 0.0, 1.0)
    dispatch!(store, ColorButtons.SetHoveredColorTo(ColorButtons.get_label(state), 1.0, 0.0, 1.0))
    @test ColorButtons.get_hovered_color(get_state(store)) == ReduxImGui.CImGui.HSV(1.0, 0.0, 1.0)
    dispatch!(store, ColorButtons.SetActiveColorTo(ColorButtons.get_label(state), 0.5, 1.0, 0.0))
    @test ColorButtons.get_active_color(get_state(store)) == ReduxImGui.CImGui.HSV(0.5, 1.0, 0.0)

    # test vectors
    store = create_store(ColorButtons.reducer, [ColorButtons.State("b1"),ColorButtons.State("b2")])
    state = get_state(store)
    dispatch!(store, ColorButtons.AddButton("b3"))
    dispatch!(store, ColorButtons.DeleteButton("b1"))
    state = get_state(store)
    @test state[1] == ColorButtons.State("b2")
    @test state[2] == ColorButtons.State("b3")
    dispatch!(store, ColorButtons.SetTriggeredTo("b2", false))
    @test ColorButtons.is_triggered(get_state(store)[1]) == false

    # test dicts
    store = create_store(ColorButtons.reducer, Dict("b1"=>ColorButtons.State("b1"),"b2"=>ColorButtons.State("b2")))
    state = get_state(store)
    dispatch!(store, ColorButtons.Resize("b2", 10, 20))
    state = get_state(store)
    @test ColorButtons.get_size(state["b2"]) == (10,20)
end

@testset "Buttons | RadioButton" begin
    name = "Radio_button_test"
    store = create_store(RadioButtons.reducer, RadioButtons.State(name))
    state = get_state(store)

    # test initial state
    @test state.label == name
    @test state.is_active == false
    @test state.is_triggered == false

    # test actions
    new_name = "Radio_button_new_name"
    state = get_state(store)
    dispatch!(store, RadioButtons.Rename(state.label, new_name))
    @test get_state(store).label == new_name
    dispatch!(store, RadioButtons.Rename(state.label, name))
    @test RadioButtons.get_label(get_state(store)) == name

    dispatch!(store, RadioButtons.SetActiveTo(state.label, true))
    @test get_state(store).is_active == true

    dispatch!(store, RadioButtons.SetTriggeredTo(state.label, true))
    @test get_state(store).is_triggered == true

    # test vectors
    store = create_store(RadioButtons.reducer, [RadioButtons.State("b1"),RadioButtons.State("b2")])
    state = get_state(store)
    dispatch!(store, RadioButtons.AddButton("b3"))
    dispatch!(store, RadioButtons.DeleteButton("b1"))
    state = get_state(store)
    @test state[1] == RadioButtons.State("b2")
    @test state[2] == RadioButtons.State("b3")
    dispatch!(store, RadioButtons.SetTriggeredTo("b2", true))
    @test RadioButtons.is_triggered(get_state(store)[1]) == true

    dispatch!(store, RadioButtons.OnlySetThisOneToActive("b3"))
    @test RadioButtons.is_active(get_state(store)[1]) == false
    @test RadioButtons.is_active(get_state(store)[2]) == true

    # test dicts
    store = create_store(RadioButtons.reducer, Dict("b1"=>RadioButtons.State("b1"),"b2"=>RadioButtons.State("b2")))
    state = get_state(store)
    dispatch!(store, RadioButtons.SetActiveTo("b2", true))
    state = get_state(store)
    @test RadioButtons.is_active(state["b2"]) == true
end

@testset "Buttons | RadioButtonGroup" begin
    RBG = RadioButtonGroups
    name = "Radio_button_group_test"
    store = create_store(RBG.reducer, RBG.State(name))
    state = get_state(store)

    # test initial state
    @test state.label == name
    @test isempty(state.buttons) == true
    @test state.is_triggered == false

    # test actions
    new_name = "Radio_button_group_new_name"
    state = get_state(store)
    dispatch!(store, RBG.Rename(state.label, new_name))
    @test get_state(store).label == new_name
    dispatch!(store, RBG.Rename(state.label, name))
    @test RBG.get_label(get_state(store)) == name

    dispatch!(store, RBG.SetTriggeredTo(state.label, true))
    @test RBG.is_triggered(get_state(store)) == true

    dispatch!(store, RBG.EditButtons(state.label, RadioButtons.AddButton("radio button1")))
    dispatch!(store, RBG.EditButtons(state.label, RadioButtons.SetActiveTo("radio button1", true)))
    @test get_state(store).buttons[1].is_active == true

    # test dicts
    store = create_store(RBG.reducer, Dict("b1"=>RBG.State("b1"),"b2"=>RBG.State("b2")))
    state = get_state(store)
    dispatch!(store, RBG.EditButtons(RBG.get_label(state["b1"]), RadioButtons.AddButton("radio button1")))
    state = get_state(store)
    @test state["b1"].buttons[1].label == "radio button1"
end


text:
module Sigmoid
    using LoopVectorization

    function func!(output_matrix::Array{Float32}, value_matrix::Array{Float32})
        @avxt for i in eachindex(value_matrix)
            value_matrix[i] = ifelse(value_matrix[i]>3.0f38, 3.0f38, value_matrix[i])
            value_matrix[i] = ifelse(value_matrix[i]<-3.0f38, -3.0f38, value_matrix[i])
            output_matrix[i] = ifelse(value_matrix[i]>=0, 1/(1+exp(-value_matrix[i])), exp(value_matrix[i])/(1+exp(value_matrix[i])))
        end
    end

    function get_∇biases!(∇biases::Array{Float32}, value_matrix::Array{Float32}, δ::Array{Float32})
        @avxt for i in eachindex(value_matrix)
            ∇biases[i] = exp(-value_matrix[i])/(1+exp(-value_matrix[i]))^2*δ[i]
        end
    end

    function get_name()
        return "Sigmoid"
    end
end

# Source: https://timvieira.github.io/blog/post/2014/02/11/exp-normalize-trick/


text:
module temp

using DifferentialEquations
using BoundaryValueDiffEq
using OrdinaryDiffEq
using RecursiveArrayTools

using BenchmarkTools
using InteractiveUtils
using RecursiveArrayTools

struct Param{S}
    state::S
    dtstate2::S
end

@inline function Base.zeros(::Type{ArrayPartition},::Type{T},n::Int) where T
    return ArrayPartition(
        zeros(n),zeros(n),zeros(n),zeros(n),zeros(n),zeros(n),
        zeros(n),zeros(n),zeros(n),zeros(n),zeros(n),zeros(n)
    )
end

function rhs!(dtstate::S, regstate::S, param::Param{S}, t) where S

    #Unpack the parameters

    dtstate2 = param.dtstate2
    state = param.state

    ## Store the current state in the container
    ## so it can be modified during the calculations
    # for i in 1:12
    #     state.x[i] .= regstate.x[i]
    # end
    state .= regstate

    # Unpack the states so I can refer to individual things
    a,b,c,d,e,f,g,h,i,j,k,l = state.x

    A,B,C,D,E,F,G,H,I,J,K,L = dtstate.x

    ## Do all your calculations here....
    ## here is an example

    @. A = a + b + c + d + e + f + g + h + i + j + k + l
    @. B = a + b + c + d + e + f + g + h + i + j + k + l
    @. C = a + b + c + d + e + f + g + h + i + j + k + l
    @. D = a + b + c + d + e + f + g + h + i + j + k + l
    @. E = a + b + c + d + e + f + g + h + i + j + k + l
    @. F = a + b + c + d + e + f + g + h + i + j + k + l
    @. G = a + b + c + d + e + f + g + h + i + j + k + l
    @. H = a + b + c + d + e + f + g + h + i + j + k + l
    @. I = a + b + c + d + e + f + g + h + i + j + k + l
    @. J = a + b + c + d + e + f + g + h + i + j + k + l
    @. K = a + b + c + d + e + f + g + h + i + j + k + l
    @. L = a + b + c + d + e + f + g + h + i + j + k + l

    ## Store the current time derivatives in the container
    # for use elsewhere
    # for i in 1:12
    #     dtstate2.x[i] .= dtstate.x[i]
    # end
    dtstate2 .= dtstate

end

function main()

    n = 4000

    numvar = 12

    T = Float64

    tspan = T[0.,1.]

    cont = ArrayPartition{T, NTuple{numvar, Vector{T}}}

    regstate = zeros(ArrayPartition,T,n)::cont
    state = zeros(ArrayPartition,T,n)::cont
    dtstate = zeros(ArrayPartition,T,n)::cont

    param = Param(state,dtstate)

    prob = ODEProblem(rhs!, regstate, tspan, param)

    atol = eps(T)^(T(3) / 4)

    alg = RK4()

    # @btime sol = solve(
    #     $prob, $alg,
    #     abstol = $atol,
    #     dt = 0.1,
    #     adaptive = false,
    #     saveat = 0.1,
    #     alias_u0 = true
    #     # progress = true,
    #     # progress_steps = custom_progress_step,
    #     # progress_message = custom_progress_message
    # )

    @code_warntype rhs!(dtstate,regstate,param,0.)

    return

end

end


text:
using Base: func_for_method_checked, structdiff, Const
using GLMakie
using Statistics
using Colors
using Distributions
include("load_eeglab.jl")

##

# Ticks
"""
Create 10 x-ticks from 10% to 90% of the x-axis as time in seconds.
time_idx should be the lower bound of the plotted data.
"""
function set_time_ticks(axis, time_idx, plot_range)
    xtick_pos = LinRange(round(plot_range*0.1),round_int(plot_range*0.9), 10)
    # Convert x_pos to time
    xtick_pos_label = map(x->round((time_idx + x)/srate, digits=2), xtick_pos)
    axis.xticks = (xtick_pos, ["$(x)s" for x in xtick_pos_label])
end

# Events
"""
Check what events lines are in the current plot range and make them visible.
Make events outside of this range invisible.
"""
function redraw_event_markers(lines, lower_bound, upper_bound)
    # Check for visibility
    vis_events = lower_bound .< event_df[:,:latency] .< upper_bound

    for (event_idx, event) in enumerate(lines[vis_events])
        event.visible = true
    end
    for (event_idx, event) in enumerate(lines[.!vis_events])
        event.visible = false
    end
end

"""
Check what events tags are in the current plot range and make them visible.
Make events outside of this range invisible.
"""
function redraw_event_text(tags, event_df, lower_bound, upper_bound, tag_height)
    # Check for visibility
    vis_events = lower_bound .< event_df[:,:latency] .< upper_bound
    for tag in tags[vis_events]
        # Index of tag in unfiltered array
        idx_unfiltered = findfirst((x)-> x==tag, tags)
        # Observable position doesn't work for text -> manually update position
        corrected_lat = event_df[idx_unfiltered,:latency] - lower_bound
        tag.attributes.attributes[:position][] = (corrected_lat, tag_height)
        tag.visible = true
    end
    for tag in tags[.!vis_events]
        tag.visible = false
    end
end

# Reject regions mark time regions for rejection
"""
Create a new reject region at the passed region if it does not overlap with an exisiting one.
Stores the plot in the reject_plots array and its limits in reject_limits array.
"""
function create_new_reject_region!(axis, new_rect, reject_limits, reject_plots, slider_time, nsample_to_show_obs, lower_bound, upper_bound)
    start_idx = deepcopy(new_rect.origin[1]+max(1, slider_time-nsample_to_show_obs/2))
    end_idx = deepcopy(start_idx + new_rect.widths[1])
    overlaps = false
    for limit in reject_limits
        if limit[1] < start_idx < limit[2] || limit[1] < end_idx < limit[2]
            overlaps = true
        end
    end
    if !overlaps
        append!(reject_limits, [[start_idx, end_idx]])

        plot_start = @lift(max(0, start_idx-$lower_bound))
        plot_end = @lift(min(end_idx - $lower_bound, $upper_bound))
        y_lim = @lift([$(axis.finallimits).origin[2], $(axis.finallimits).widths[2]])
        y_height = @lift(abs($y_lim[1]-$y_lim[2]))
        draw_rect = @lift(FRect($plot_start, $y_lim[1], abs($plot_end-$plot_start), $y_height))
        append!(reject_plots, [draw_reject_region(axis, draw_rect)])
    end
end

"""
Check if reject region is under passed x-position.
Delete region from axis, plot_array, and region limits.
"""
function delete_reject_region(axis, position, reject_limits, reject_plots, slider_time, nsample_to_show_obs)
    for (region_idx, limit) in enumerate(reject_limits)
        if limit[1] < position+max(1, slider_time-nsample_to_show_obs/2) < limit[2]
            deleteat!(reject_limits, region_idx)
            delete!(axis, reject_plots[region_idx])
            deleteat!(reject_plots, region_idx)
        end
    end
end

"""
Draw a transparent rectangle that represents a region to be rejected
"""
function draw_reject_region(axis, rect)
    poly!(axis, rect, raw = true, visible = true, color = RGBAf0(1.0, 0.08, 0.58, 0.3), strokewidth = 0)
end

"""
Check what reject regions are in the current plot range and make them visible.
Make reject regions outside of this range invisible.
"""
function redraw_reject_regions(limits, plots, lower_bound, upper_bound)
    # Create dataframe with lower and upper bounds of the regions
    limits_df = DataFrame( lower=[], upper=[])
    for region in limits
        push!(limits_df, region)
    end
    # Check what regions are visible
    vis_regions = ((lower_bound .< limits_df[:,:lower] .< upper_bound)
                .| (lower_bound .< limits_df[:,:upper] .< upper_bound))

    for region in plots[vis_regions]
        region.visible = true
    end
    for region in plots[.!vis_regions]
        region.visible = false
    end
end

# Channel rejection
"""
Find channel that is closest to mouseposition at click and color it red.
If channel is marked already, unmark it.
"""
function mark_channel(axis, mouse_pos, channel_plots)
    mark_plot = Makie.pick(axis.scene, mouse_pos, 500)[1]
    # Lineplot is an actual EEG plot

    if mark_plot in channel_plots
        # If marked, unmark it
        if mark_plot.attributes.color[] == :red
            mark_plot.attributes.color = :grey0
        else
            mark_plot.attributes.color = :red
        end
    end

    findall(x -> x==mark_plot, channel_plots)
end

# Helper functions
"""
Check if mouse is inside scene.
"""
function is_mouseinside(scene)
    return Vec(scene.events.mouseposition[]) in pixelarea(scene)[]
end

"""
Take a rectangle with origin and widths and return the same rectangle but with
the origin in the lower left corner.
This ensures that widths are positive.
"""
function absrect(rectangle)
    origin = rectangle.origin
    widths = rectangle.widths
    opposite_corner = origin + widths
    new_origin = min.(origin, opposite_corner)
    new_widths = abs.(widths)
    new_rectangle = FRect(new_origin[1], new_origin[2], new_widths[1], new_widths[2])
end

"""
Select a time region in the plot when Left-Shift is hold.
Return the selected region.
"""
function select_rectangle(axis; blocking = false, priority = 2, strokewidth = 3.0, kwargs...)
    scene = axis.scene
    key = Keyboard.left_shift
    waspressed = Node(false)
    rect = Node(FRect(0, 0, 1, 1)) # plotted rectangle
    rect_ret = Node(FRect(0, 0, 1, 1)) # returned rectangle
    y_lim = @lift([$(axis.finallimits).origin[2], $(axis.finallimits).widths[2]])
    y_height = @lift(abs($y_lim[1]-$y_lim[2]))
    # Create an initially hidden rectangle
    plotted_rect = poly!(
        scene, rect, raw = true, visible = false, color = RGBAf0(1.0, 0.08, 0.58, 0.3), strokewidth = 0, transparency = false, kwargs...,
    )

    # Control key press and hold
    on(events(scene).keyboardbutton , priority=priority) do event
        # Start tracking on key press
        if event.key == key
            if event.action == Keyboard.press && is_mouseinside(scene)
                mp = mouseposition(scene)
                waspressed[] = true
                plotted_rect[:visible] = true # start displaying
                rect[] = FRect(mp[1], y_lim[][1], 0.0, y_height[])
                return Consume(blocking)
            end
        end
        # End if key is released
        if !(event.key == key && event.action == Keyboard.repeat)
            if waspressed[] # User has selected the rectangle
                waspressed[] = false
                r = absrect(rect[])
                w, h = widths(r)
                if w > 0.0 && h > 0.0 # Ensure that the rectangle has non0 size.
                    rect_ret[] = r
                end
            end
            # always hide if not the right key is pressed
            plotted_rect[:visible] = false # make the plotted rectangle invisible
            return Consume(blocking)
        end

        return Consume(false)
    end
    # Adapt rectangle depending on mouse position; only if currently tracking.
    on(events(scene).mouseposition, priority=priority) do event
        if waspressed[]
            mp = mouseposition(scene)
            mini = minimum(rect[])
            rect[] = FRect(mini[1], y_lim[][1], mp[1] - mini[1], y_height[])
            return Consume(blocking)
        end
        return Consume(false)
    end

    return rect_ret
end

"""
Disable Makie-native scroll zoom, pan scroll and rectangle zoom.
"""
function disable_native_zoom(axis)
    axis.yzoomlock = true
    axis.xzoomlock = true
    axis.xpanlock = true
    axis.ypanlock = true
    axis.xrectzoom = false
    axis.yrectzoom = false
end

"""
Make channels in mask visible and channels not in it invisible.
"""
function toggle_channel_visibility(channel_plots, mask)
    for (plot_idx, plot) in enumerate(channel_plots)
        if plot_idx in mask
            plot.visible = true
        else
            plot.visible = false
        end
    end
end

"""
Find index of element in array. If element is not in array, return 0.
"""
function findfirst_zero(element, array)
    idx = findfirst(x-> x==element, array)
    if isnothing(idx)
        idx = 0
    end
    idx
end

"""
Round number to neares Integer.
"""
function round_int(x)
    round(Int, x)
end

"""
Check what event regions are in the current plot range and make them visible.
Make event regions outside of this range invisible.
"""
function redraw_event_regions(events, plots, lower_bound, upper_bound)
    # Create dataframe with lower and upper bounds of the regions
    # limits_df = DataFrame( lower=[], upper=[])
    # for event in events
    #     push!(limits_df, region)
    # end
    # Check what regions are visible
    vis_regions = ((lower_bound .< events[:,:latency] .< upper_bound)
                .| (lower_bound .< events[:,:latency] + events[:,:duration] .< upper_bound))

    for region in plots[vis_regions]
        region.visible = true
    end
    for region in plots[.!vis_regions]
        region.visible = false
    end
end

"""
Control a reasonable scale bar size.
"""
function get_scale_factor(scale)
    1/round(scale; digits=1)
end

"""
Plot the data and set up the various interactions.
"""
function data_plot(data,srate,event_df,chanlabels,title, path_to_output, output_file_name)

    # Output file
    io = open(path_to_output * output_file_name, "w")
    reject_chans = []

    # Data size
    nchan = size(data,1)
    nsamples = size(data,2)

    # Parameters to set
    offset = 5.0
    y_estate = offset*nchan
    nsample_to_show = 1000 # only for initial plotting

    tag_height = (nchan+5)*offset
    plot_high = (nchan+10)*offset
    plot_low = -offset
    event_upper_lim = tag_height/plot_high-0.01

    # Set up figure
    fig = Figure(resolution = (2560, 1440))
    axis = fig[1, 1] = Axis(fig, title = title)

    # disable native zoom
    disable_native_zoom(axis)

    # Set up time slider
    time_slider = Slider(fig[2, 1], range = 1:1:nsamples-nsample_to_show, startvalue = 1)
    slider_time = time_slider.value

    # Channel scroll
    channel_slider = IntervalSlider(fig[1, 2], range = 1:1:nchan, startvalues = (1,nchan), horizontal = false)
    channel_interval = channel_slider.interval

    # Time range to plot
    nsample_to_show_obs = Node(nsample_to_show)
    lower_bound = @lift(max(1, round_int($slider_time)-round_int($nsample_to_show_obs/2)))
    upper_bound = @lift(min(nsamples, $lower_bound+round_int($nsample_to_show_obs)))
    range = @lift([$lower_bound,$upper_bound])

    # Scale Observables
    scale_obs = Node(0.4)
    amp_scalable = Node(false)

    # Channel Observables
    offset_obs = Node(offset)

    # Plot EEG
    channel_plots = []
    for chan in 1:nchan
        # Get data chunk currently to be plottet
        data_chunk = @lift(data[chan,$range[1]:$range[2]])
        chunk_mean = @lift(mean($data_chunk))
        chan_pos = @lift(findfirst_zero(chan, ($channel_interval[1]:$channel_interval[2])))
        # Mean center, scale and shift to plot all channels above each other in y-direction
        data_chunk_plot = @lift((($data_chunk .-$chunk_mean)  .* $scale_obs) .+ $chan_pos*$offset_obs)
        append!(channel_plots, [lines!(axis, data_chunk_plot, color = :grey0)])
    end

    # Plot events as vertical lines
    event_lines = map(eachrow(event_df)) do event
        plot_pos = @lift(event.latency-$lower_bound)
        vlines!(axis, plot_pos, ymax=event_upper_lim, visible = false, color = :steelblue1)
    end

    # event_regions = map(eachrow(event_df)) do event
    #     plot_start = @lift(event.latency-$lower_bound)
    #     plot_end = @lift(min(event.latency + event.duration - $lower_bound, $upper_bound))
    #     y_lim = @lift([$(axis.finallimits).origin[2], $(axis.finallimits).widths[2]])
    #     y_height = @lift(abs($y_lim[1]-$y_lim[2]))
    #     draw_rect = @lift(FRect($plot_start, $y_lim[1], abs($plot_end-$plot_start), $y_height))

    #     poly!(axis, draw_rect, raw = true, visible = false, color = :steelblue1, strokewidth = 0)
    # end

    #TW SPECIFIC
    level_1 = ["experimentStart", "longBaselineStart",   "longBaselineWaypoint", "longBaselineTargetpoint", "intertrialBaselineStart", "intertrialBaselineTransitpoint", "intertrialBaselineTargetpoint", "intertrialBaselineStop", "triangleStart", "triangleTaskFirstpoint", "triangleTaskSecondpoint", "triangleTaskHomingtaskStart", "triangleTaskHomingtaskLocationConfirmed", "triangleTaskHomingtaskDistanceerror", "triangleStop", "longBaselineStop", "experimentStop"]
    level_2 = ["rotationStart", "rotationStop", "translationStart", "translationStop","triangleTaskPoleVisible", "triangleTaskPoleSpotted", "longBaselineSphereSpotted", "intertrialBaselineSphereSpotted"]
    level_3 = ["miniBaselineStart", "miniBaselineStop", "subtaskFinished"]
    # TW SPECIFIC

    # Plot event tag above event line
    event_tags = map(eachrow(event_df)) do event
        plot_pos = @lift(event.latency-$lower_bound)
        # TW SPECIFIC
        if event.event ∈ level_1
            text!(axis.scene, event.event, position=(plot_pos,tag_height), offset=(0,30), align=(:center, :center), visible=false, color = :green)
        elseif event.event ∈ level_2
            text!(axis.scene, event.event, position=(plot_pos,tag_height), offset=(0,10), align=(:center, :center), visible=false, color = :orange)
        else
            text!(axis.scene, event.event, position=(plot_pos,tag_height), offset=(0,-10), align=(:center, :center), visible=false, color = :red)
        end
        # TW SPECIFIC
        #text!(axis.scene, event.event, position=(plot_pos,tag_height), align=(:center, :center), visible=false, color = :black)
    end

    # Plot range changes (i.e. zoom or scroll)
    on(range) do value
        redraw_event_markers(event_lines, lower_bound[], upper_bound[])
        redraw_event_text(event_tags, event_df, lower_bound[], upper_bound[], tag_height)
        redraw_reject_regions(reject_limits, reject_plots, lower_bound[], upper_bound[])
        # redraw_event_regions(event_df, event_regions)
        set_time_ticks(axis, lower_bound[], nsample_to_show_obs[])
    end

    # Set initial ticks
    axis.yticks = (offset:offset:(nchan)*offset, chanlabels)
    axis.yticklabelsize = 10.0
    set_time_ticks(axis, 0, nsample_to_show_obs[])

    # Scrolling time/amplitude-zoom
    on(events(fig).scroll, priority = 100) do scroll
        if amp_scalable[]
            scale_obs[] = max(0, scale_obs[] + 0.01*scroll[2])
        else
            nsample_to_show_obs[] = clamp(nsample_to_show_obs[]-floor(scroll[2]*30), 150, nsamples)
        end
    end

    # Enable amplitude-zoom on left-ctrl hold
    on(events(fig).keyboardbutton) do event
        if event.key == Keyboard.left_control && event.action in (Keyboard.repeat, Keyboard.press)
            amp_scalable[] = true
        else
            amp_scalable[] = false
        end
    end

    # Color channel on click
    on(events(fig).mousebutton) do event
        if event.action == Mouse.press && event.button == Mouse.left
            mouse_pos = events(fig).mouseposition[]
            channel_id = mark_channel(axis, mouse_pos, channel_plots)
            if !isempty(channel_id)
                if !(channel_id[1] in reject_chans)
                    append!(reject_chans, channel_id[1])
                else
                    deleteat!(reject_chans, findall(x->x==channel_id[1], reject_chans))
                end
            end
        end
        return Consume(false)
    end

    # Keyboard input
    on(events(fig).keyboardbutton) do event
        if event.action in (Keyboard.press, Keyboard.repeat)
            # Time-scroll on left/right arrow key
            if event.key == Keyboard.left
                set_close_to!(time_slider, max(1, slider_time[]-100))
            end
            if event.key == Keyboard.right
                set_close_to!(time_slider, min(slider_time[]+100, nsamples-nsample_to_show))
            end

            # Close file
            if event.key == Keyboard.p
                for chan_id in reject_chans
                    println(io, chan_id)
                end
                close(io)
            end
        end
        return Consume(false)
    end

    # Mark time windows
    rect = select_rectangle(axis)
    reject_limits = Array{Float32}[]
    reject_plots = []
    on(rect) do new_rect
        create_new_reject_region!(axis, new_rect, reject_limits, reject_plots, slider_time[], nsample_to_show_obs[], lower_bound, upper_bound)
    end

    # Delete time window on right click
    on(events(fig).mousebutton) do event
        if event.button == Mouse.right && event.action == Mouse.press
            mp_x = mouseposition(axis.scene)[1]
            delete_reject_region(axis, mp_x, reject_limits, reject_plots, slider_time[], nsample_to_show_obs[])
        end
    end

    # Time input
    time_input = Textbox(fig[2, 2], width = 100, validator = Float64, placeholder=string(nsample_to_show_obs[]/srate))
    time_value = time_input.stored_string

    on(time_value) do time_range
        new_range = min(parse(Float64, time_range), nsamples/srate)
        nsample_to_show_obs[] = round_int(new_range * srate)
    end

    # Update time zoom variables
    on(nsample_to_show_obs) do nsample
        xlims!(0,nsample)
        time_input.displayed_string = string(nsample/srate)
    end

    # Channel scroll
    on(channel_interval) do interval
        toggle_channel_visibility(channel_plots, (interval[1]:interval[2]))
        offset_obs[] = y_estate/(interval[2]-interval[1]+1)
        new_nchan = interval[2]-interval[1]+1
        new_chanlabels = chanlabels[interval[1]:interval[2]]
        axis.yticks = (offset_obs[]:offset_obs[]:(new_nchan)*offset_obs[], new_chanlabels)
    end

    # Inital event rendering
    redraw_event_markers(event_lines, lower_bound[], upper_bound[])
    redraw_event_text(event_tags, event_df, lower_bound[], upper_bound[], tag_height)

    # Set visible limits (doesn't affect plotted data)
    tightlimits!(axis, Left(), Right()) # No left/right margin
    ylims!(plot_low, plot_high)

    # Scalebar
    scale_x_pos = @lift([round_int($nsample_to_show_obs*0.98)])
    scale_y_pos = [round_int((plot_low + plot_high)/2)]
    scale_size = @lift($scale_obs*100/2) # Divide by 2 because it extends in both directions
    errorbars!(scale_x_pos, scale_y_pos, scale_size, color = :red, linewidth = 3, whiskerwidth = 10)

    # Scalelabel
    label_pos = @lift(($scale_x_pos[1], scale_y_pos[1]))
    label_offset = @lift((-round_int($nsample_to_show_obs*0.02), 0))
    label_text = string(100) * " μV"

    text!(axis.scene, label_text, position = label_pos, color = :red, align=(:right, :center), offset = (-30, 0))

    fig
end
##

subject = 21

path = "/Users/timotheusberg/Documents/Work/BPN/triangle_walking/04_data/3_EEG-preprocessing/s$subject/"

input_file = "s$(subject)_preprocessed.set"

output_file = "s$(subject)_reject.txt"

data,srate,event_df,chanlocs_df,EEG = import_eeglab(path * input_file)

# Test data
time = 1:5000
sin_time = sin.(time./10).*50
test_data = repeat(transpose(sin_time), 128, 1)

chanlabels = chanlocs_df[:, :labels]

##

fig = data_plot(data,srate,event_df,chanlabels,input_file,path, output_file)

##
# PROBLEMS
# - make sure that nsample to display
#   doesn't exceed data bounds and scroll when end is reached

# TODO
# - show amplitude scale
# - read in reject stuff
# - export reject regions

#
# OPTIONAL
# - show mini plot under time-scroll that visualizes position and zoom level


text:
"""
    setCrossoverOperator(method::Function, varArgs...; genj::GeneticJuliaStructure = GenJ,
                         probability::AbstractFloat = 0.7, nParents::Integer = 2,
                         nChildren::Integer = -1)

Sets the crossover operator of parents for `genj`, receiving a function that
crosses multiples individual representations and returns one child or multiple
children. This function constructs a structure of type `CrossoverOperator` and
adds it to the main structure.

# Arguments
- `method::Function`: function that receives, first, a list of parents as
    arguments which are individual representations and `"varArgs"`, which is a
    list of arguments, as own arguments, and returns a number between 1 and
    nParents of children in an array, where each one is an individual's genotype.
    (For information about genotypes look [`Genotype`](@ref)). See `"Examples"`.
- `varArgs`: list of arguments for crossover method. See `"Examples"`.

# Keyword Arguments
- `genj::GeneticJuliaStructure = GenJ`: the main structure. For `code user`: don't modify
    the default value unless multiple experiments are going to be run.
- `probability::AbstractFloat = 0.7`: a floating point number, between `0` and `1`,
    that represents the probability of doing each single crossover operation.
- `nParents::Integer = 2`: the number of parents that are going to be crossed. It
    must be a number greater than 0.
- `nChildren::Integer = -1`: the number of children as result of the crossover method.
    It can be any number greater than 0, while -1 can be used to express that the number
    of children generated is the same as the number of parents.

# Examples
```jldoctest
julia> #Function with 2 parents, no more arguments and 1 child
julia> function ExampleCrossParents1(Parent1::Genotype, Parent2::Genotype)
    #···Code···
    return children::Array{Genotype} # Type expressed here for user awareness
end

julia> setCrossoverOperator(ExampleCrossParents1, probability = 0.4, nChildren = 1)
```
```jldoctest
julia> #Function with more than 2 parents, no more arguments and the same number of children as parents.
julia> function ExampleCrossParents2(Parent1::Genotype, Parent2::Genotype, Parent3::Genotype)
    #···Code···
    return children::Array{Genotype} # Type expressed here for user awareness
end

julia> setCrossoverOperator(ExampleCrossParents2, probability = 0.6, nParents = 3)
```
```jldoctest
julia> #Function with more arguments and the same number of children as parents
julia> function ExampleCrossParents3(Parent1::Genotype, Parent2::Genotype, arg1, arg2)
    #···Code···
    return children::Array{Genotype}(undef, 2) # Type expressed here for user awareness
end

julia> setCrossoverOperator(ExampleCrossParents3, arg1, arg2, probability = 0.6)
```

See also: [`CrossoverOperator`](@ref)
"""
function setCrossoverOperator(method::Function, varArgs...; genj::GeneticJuliaStructure = GenJ,
                              probability::AbstractFloat = 0.7, nParents::Integer = 2,
                              nChildren::Integer = -1)

    if !(0 <= probability <= 1)
        error("Crossover probability ($probability) must be within 0 and 1")
    end

    if nParents < 1
        error("Number of parents ($nParents) given must be greater than 0")
    end

    if nChildren != -1 && nChildren < 1
        error("Number of children ($nChildren) per cross must be either -1 or greater than 0")
    end

    if nChildren == -1
        nChildren = nParents
    end

    varArgs = [i for i in varArgs]
    genj._crossoverOp = CrossoverOperator(method, probability, nParents, nChildren, varArgs)
end # function

"""
    setNullCrossover()

Makes the crossover non-existent for the evolutionary proccess.
This function constructs a structure of type `CrossoverOperator` and adds it to the main structure.

At least, one between crossover and mutation (or both) must exist in the experiment.
"""
function setNullCrossover(;genj::GeneticJuliaStructure = GenJ)
    noFunction() = nothing
    genj._crossoverOp = CrossoverOperator(noFunction, 0.0, 2, -1, [])
end

max_stars_repo_name: SergioGmezM/GeneticJulia

max_issues_repo_name: SergioGmezM/GeneticJulia

max_forks_repo_name: SergioGmezM/GeneticJulia


text:
# domain.jl

using Dates

# ??? customer id
struct JournalEntry
    id::String
    period::Int64
    date::DateTime
    customer_id::String
    invoice_nbr::String
    from::Int64
    to::Int64
    debit::Float64
    credit::Float64
    vat::Float64
    descr::String
    # constructors
    JournalEntry(id, period, date, customer_id, invoice_nbr, from, to, debit, credit, vat, descr) =
        new(id, period, date, customer_id, invoice_nbr, from, to, debit, credit, vat, descr)
end # JournalEntry

struct Record
    id::String # was stm_nbr
    accountid::Int64
    date::DateTime
    customerid::String
    invoice_nbr::String
    debit::Float64
    credit::Float64
    descr::String
end # Record

struct Account
    id::String
    name::String
    group::String
    # constructors
    Account(id, name) = new(id, name, "")
    Account(id, name, group) = new(id, name, group)
end # Account


text:
using BinaryProvider

# Parse some basic command-line arguments
const verbose = "--verbose" in ARGS
const prefix = Prefix(get([a for a in ARGS if a != "--verbose"], 1, joinpath(@__DIR__, "usr")))

# These are the several binary objects we care about
products = Product[
    LibraryProduct(prefix, "libopk",    :libopk),
    LibraryProduct(prefix, "libbobyqa", :libbobyqa),
    LibraryProduct(prefix, "libcobyla", :libcobyla),
    LibraryProduct(prefix, "libnewuoa", :libnewuoa),
]

# Download binaries from hosted location
bin_prefix = "https://github.com/emmt/OptimPackBuilder/releases/download/v3.1.0"
download_info = Dict(
    Linux(:aarch64; libc=:glibc) => ("$bin_prefix/OptimPack.v3.1.0.aarch64-linux-gnu.tar.gz", "71e6a481776e9b5f2fed5e7d0ab8915d6acb0f1ad44562ab43f9931ea295c2f6"),
    Linux(:aarch64; libc=:musl) => ("$bin_prefix/OptimPack.v3.1.0.aarch64-linux-musl.tar.gz", "1a61e86921c530252ca40960b4c2441f0b8511e1c7e9e3adbccbd5509e1944f1"),
    Linux(:armv7l; libc=:glibc) => ("$bin_prefix/OptimPack.v3.1.0.armv7l-linux-gnueabihf.tar.gz", "b1dae85bde6686db7dad23c1533df2e4e33c9f84cdaa54573bd2d78e96a2f18a"),
    Linux(:armv7l; libc=:musl) => ("$bin_prefix/OptimPack.v3.1.0.armv7l-linux-musleabihf.tar.gz", "7f964d05429fecd7c20713d3672a15d4019f1f390f94c65282c713d96a0f5bc2"),
    Linux(:i686; libc=:glibc) => ("$bin_prefix/OptimPack.v3.1.0.i686-linux-gnu.tar.gz", "b2a548e7bc570ae493ce2bc9213d083675aaca28a34c429f4e0c2bf060e474d0"),
    Linux(:i686; libc=:musl) => ("$bin_prefix/OptimPack.v3.1.0.i686-linux-musl.tar.gz", "a8e830e01b5772e778111c2151755498fead4bd1e8c51281bce47c245857658a"),
    Linux(:x86_64; libc=:glibc) => ("$bin_prefix/OptimPack.v3.1.0.x86_64-linux-gnu.tar.gz", "f18bb702463eddc1e9219f92d1c8e0db184e8987d664d16bd7e7bfd2dfca0ebd"),
    Linux(:x86_64; libc=:musl) => ("$bin_prefix/OptimPack.v3.1.0.x86_64-linux-musl.tar.gz", "c3c02dde5ef26216706997312d01f80452c9a29b73698347ed7582ee86eabb47"),
    Linux(:powerpc64le; libc=:glibc) => ("$bin_prefix/OptimPack.v3.1.0.powerpc64le-linux-gnu.tar.gz", "1f2f9ea391488ee3362fad35c9ff9d33336c6b3091bff4882ba5a8291603e5bc"),
    FreeBSD(:x86_64) => ("$bin_prefix/OptimPack.v3.1.0.x86_64-unknown-freebsd11.1.tar.gz", "9ddfd26f6295828097f7eeeb38add123f309715f7f59aeca4c8cb737812b4e99"),
    MacOS(:x86_64) => ("$bin_prefix/OptimPack.v3.1.0.x86_64-apple-darwin14.tar.gz", "99dfdf40ad334f430981af2aa5df44e8efd25a475ca2c4d97f38a5066238909c"),
    Windows(:i686) => ("$bin_prefix/OptimPack.v3.1.0.i686-w64-mingw32.tar.gz", "9c78796fe2b58693b610e440aa5390b5fce670dad4250ba4542d79e72636a86c"),
    Windows(:x86_64) => ("$bin_prefix/OptimPack.v3.1.0.x86_64-w64-mingw32.tar.gz", "f839ab4f8d42f398bf089917a4581e1a7e37be5519e5516aec8154da09f6643c"),
)
# First, check to see if we're all satisfied
if any(!satisfied(p; verbose=verbose) for p in products)
    try
        # Download and install binaries
        url, tarball_hash = choose_download(download_info)
        install(url, tarball_hash; prefix=prefix, force=true, verbose=true)
    catch e
        if typeof(e) <: ArgumentError
            error("Your platform $(Sys.MACHINE) is not supported by this package!")
        else
            rethrow(e)
        end
    end

    # Finally, write out a deps.jl file
    write_deps_file(joinpath(@__DIR__, "deps.jl"), products, verbose=true)
end


text:
disabled() = get(ENV, "R123_DISABLE_AESNI", "") != ""

const deps_dir = dirname(@__FILE__)

has_aesni() = try
    cmd = Base.julia_cmd()
    push!(cmd.exec, joinpath(deps_dir, "has_aesni.jl"))
    run(cmd)
    true
catch
    false
end

const filename = joinpath(deps_dir, "aes-ni")
isfile(filename) && rm(filename)

if disabled()
    @info "AES-NI is disabled."
elseif has_aesni()
    @info "AES-NI is enabled."
    touch(filename)
else
    @info "AES-NI is not supported."
end


text:
@doc raw"""
`HDAE`: Hamiltonian Differential Algebraic Equation

Defines a Hamiltonian differential algebraic initial value problem, that is
a canonical Hamiltonian system of equations subject to Dirac constraints,
```math
\begin{aligned}
\dot{q} (t) &= v(t, q(t), p(t)) + u(t, q(t), p(t), \lambda(t)) + \bar{g}(t, q(t), p(t), \lambda(t), \gamma(t)) , & q(t_{0}) &= q_{0} , \\
\dot{p} (t) &= f(t, q(t), p(t)) + g(t, q(t), p(t), \lambda(t)) + \bar{f}(t, q(t), p(t), \lambda(t), \gamma(t)) , & p(t_{0}) &= p_{0} , \\
0 &= \phi (t, q(t), p(t)) , \\
0 &= \psi (t, q(t), p(t), \dot{q}(t), \dot{p}(t)) ,
\end{aligned}
```
with vector fields ``v``, ``u``, ``\bar{u}`` and ``f``, ``g``, ``\bar{g}``,
primary constraint ``\phi(q,p)=0`` and secondary constraint ``\psi(q,p,\lambda)=0``,
initial conditions ``(q_{0}, p_{0})``, the dynamical variables ``(q,p)``
taking values in ``\mathbb{R}^{d} \times \mathbb{R}^{d}`` and
the algebraic variables ``(\lambda, \gamma)`` taking values in
``\mathbb{R}^{m} \times \mathbb{R}^{m}``.

### Parameters

* `DT <: Number`: data type
* `TT <: Real`: time step type
* `AT <: AbstractArray{DT}`: array type
* `vType <: Function`: type of `v`
* `fType <: Function`: type of `f`
* `uType <: Function`: type of `u`
* `gType <: Function`: type of `g`
* `ϕType <: Function`: type of `ϕ`
* `ūType <: Function`: type of `ū`
* `ḡType <: Function`: type of `ḡ`
* `ψType <: Function`: type of `ψ`
* `PType <: Function`: type of `P`
* `v̄Type <: Function`: type of `v̄`
* `f̄Type <: Function`: type of `f̄`
* `hamType <: Function`: Hamiltonian type
* `invType <: OptionalNamedTuple`: invariants type
* `parType <: OptionalNamedTuple`: parameters type
* `perType <: OptionalArray{AT}`: periodicity type

### Fields

* `d`: dimension of dynamical variables ``q`` and ``p`` as well as the vector fields ``v`` and ``f``
* `m`: dimension of algebraic variables ``\lambda`` and ``\gamma`` and the constraints ``\phi`` and ``\psi``
* `v`: function computing the Hamiltonian vector field ``v``
* `f`: function computing the Hamiltonian vector field ``f``
* `u`: function computing the projection for ``q``
* `g`: function computing the primary projection field ``g``
* `ϕ`: primary constraints
* `ū`: function computing the secondary projection field ``\bar{u}`` (optional)
* `ḡ`: function computing the secondary projection field ``\bar{g}`` (optional)
* `ψ`: secondary constraints (optional)
* `P`: function computing the Poisson matrix ``P``
* `v̄`: function computing an initial guess for the velocity field ``v``` (optional, defaults to `v`)
* `f̄`: function computing an initial guess for the force field ``f`` (optional, defaults to `f`)
* `t₀`: initial time (optional)
* `q₀`: initial condition for dynamical variable ``q``
* `p₀`: initial condition for dynamical variable ``p``
* `λ₀`: initial condition for algebraic variable ``λ``
* `μ₀`: initial condition for algebraic variable ``μ`` (optional)
* `hamiltonian`: function computing the Hamiltonian ``H``
* `invariants`: either a `NamedTuple` containing the equation's invariants or `nothing`
* `parameters`: either a `NamedTuple` containing the equation's parameters or `nothing`
* `periodicity`: determines the periodicity of the state vector `q` for cutting periodic solutions

### Constructors

```julia
HDAE(v, f, u, g, ϕ, ū, ḡ, ψ, v̄, f̄, P, t₀, q₀, p₀, λ₀, hamiltonian, invariants, parameters, periodicity)

HDAE(v, f, u, g, ϕ, h, t₀, q₀::StateVector, p₀::StateVector, λ₀::StateVector; kwargs...)
HDAE(v, f, u, g, ϕ, h, q₀::StateVector, p₀::StateVector, λ₀::StateVector; kwargs...)
HDAE(v, f, u, g, ϕ, h, t₀, q₀::State, p₀::State, λ₀::State; kwargs...)
HDAE(v, f, u, g, ϕ, h, q₀::State, p₀::State, λ₀::State; kwargs...)

HDAE(v, f, u, g, ϕ, ū, ḡ, ψ, h, t₀, q₀::StateVector, p₀::StateVector, λ₀::StateVector; kwargs...)
HDAE(v, f, u, g, ϕ, ū, ḡ, ψ, h, q₀::StateVector, p₀::StateVector, λ₀::StateVector; kwargs...)
HDAE(v, f, u, g, ϕ, ū, ḡ, ψ, h, t₀, q₀::State, p₀::State, λ₀::State; kwargs...)
HDAE(v, f, u, g, ϕ, ū, ḡ, ψ, h, q₀::State, p₀::State, λ₀::State; kwargs...)
```

"""
struct HDAE{vType <: Callable, fType <: Callable,
            uType <: Callable, gType <: Callable, ϕType <: Callable,
            ūType <: OptionalCallable, ḡType <: OptionalCallable, ψType <: OptionalCallable,
            v̄Type <: Callable, f̄Type <: Callable,
            poiType <: Callable,
            hamType <: Callable,
            invType <: OptionalInvariants,
            parType <: OptionalParameters,
            perType <: OptionalPeriodicity} <: AbstractEquationPDAE{invType,parType,perType,ψType}

    v::vType
    f::fType
    u::uType
    g::gType
    ϕ::ϕType
    ū::ūType
    ḡ::ḡType
    ψ::ψType
    v̄::v̄Type
    f̄::f̄Type

    poisson::poiType
    hamiltonian::hamType
    invariants::invType
    parameters::parType
    periodicity::perType

    function HDAE(v, f, u, g, ϕ, ū, ḡ, ψ, v̄, f̄, poisson, hamiltonian, invariants, parameters, periodicity)
        @assert !isempty(methods(v))
        @assert !isempty(methods(f))
        @assert !isempty(methods(u))
        @assert !isempty(methods(g))
        @assert !isempty(methods(ϕ))
        @assert !isempty(methods(ū)) || ū === nothing
        @assert !isempty(methods(ḡ)) || ḡ === nothing
        @assert !isempty(methods(ψ)) || ψ === nothing
        @assert !isempty(methods(v̄))
        @assert !isempty(methods(f̄))
        @assert !isempty(methods(poisson))
        @assert !isempty(methods(hamiltonian))

        new{typeof(v), typeof(f),
            typeof(u), typeof(g), typeof(ϕ),
            typeof(ū), typeof(ḡ), typeof(ψ),
            typeof(v̄), typeof(f̄),
            typeof(poisson), typeof(hamiltonian), typeof(invariants), typeof(parameters), typeof(periodicity)}(
                v, f, u, g, ϕ, ū, ḡ, ψ, v̄, f̄, poisson, hamiltonian, invariants, parameters, periodicity)
    end
end

HDAE(v, f, u, g, ϕ, ū, ḡ, ψ, poisson, hamiltonian; v̄=v, f̄=f, invariants=NullInvariants(), parameters=NullParameters(), periodicity=NullPeriodicity()) = HDAE(v, f, u, g, ϕ, ū, ḡ, ψ, v̄, f̄, poisson, hamiltonian, invariants, parameters, periodicity)
HDAE(v, f, u, g, ϕ, poisson, hamiltonian; kwargs...) = HDAE(v, f, u, g, ϕ, nothing, nothing, nothing, poisson, hamiltonian; kwargs...)

GeometricBase.invariants(equation::HDAE) = equation.invariants
GeometricBase.parameters(equation::HDAE) = equation.parameters
GeometricBase.periodicity(equation::HDAE) = equation.periodicity

hasvectorfield(::HDAE) = true
hashamiltonian(::HDAE) = true

function check_initial_conditions(equ::HDAE, ics::NamedTuple)
    haskey(ics, :q) || return false
    haskey(ics, :p) || return false
    haskey(ics, :λ) || return false
    if hassecondary(equ)
        haskey(ics, :μ) || return false
        eltype(ics.λ) == eltype(ics.μ) || return false
        typeof(ics.λ) == typeof(ics.μ) || return false
        axes(ics.λ) == axes(ics.μ) || return false
    end
    return true
end

function check_methods(equ::HDAE, tspan, ics::NamedTuple, params)
    applicable(equ.v, tspan[begin], ics.q, ics.p, zero(ics.q), params) || return false
    applicable(equ.f, tspan[begin], ics.q, ics.p, zero(ics.p), params) || return false
    applicable(equ.ϕ, tspan[begin], ics.q, ics.p, ics.λ, zero(ics.λ), params) || return false
    # TODO add missing methods
    return true
end

function datatype(equ::HDAE, ics::NamedTuple)
    @assert check_initial_conditions(equ, ics)
    return eltype(ics.q)
end

function arrtype(equ::HDAE, ics::NamedTuple)
    @assert check_initial_conditions(equ, ics)
    return typeof(ics.q)
end

_get_v(equ::HDAE, params) = (t,q,p,v)     -> equ.v(t, q, p, v, params)
_get_f(equ::HDAE, params) = (t,q,p,f)     -> equ.f(t, q, p, f, params)
_get_u(equ::HDAE, params) = (t,q,p,λ,u)   -> equ.u(t, q, p, λ, u, params)
_get_g(equ::HDAE, params) = (t,q,p,λ,g)   -> equ.g(t, q, p, λ, g, params)
_get_ϕ(equ::HDAE, params) = (t,q,p,ϕ)     -> equ.ϕ(t, q, p, ϕ, params)
_get_ū(equ::HDAE, params) = (t,q,p,λ,u)   -> equ.ū(t, q, p, λ, u, params)
_get_ḡ(equ::HDAE, params) = (t,q,p,λ,g)   -> equ.ḡ(t, q, p, λ, g, params)
_get_ψ(equ::HDAE, params) = (t,q,p,v,f,ψ) -> equ.ψ(t, q, p, v, f, ψ, params)
_get_v̄(equ::HDAE, params) = (t,q,p,v)     -> equ.v̄(t, q, p, v, params)
_get_f̄(equ::HDAE, params) = (t,q,p,f)     -> equ.f̄(t, q, p, f, params)
_get_h(equ::HDAE, params) = (t,q,p) -> equ.hamiltonian(t, q, p, params)
_get_poisson(equ::HDAE, params) = (t,q,p,ω) -> equ.poisson(t, q, p, ω, params)

function _functions(equ::HDAE)
    if hassecondary(equ)
        (v = equ.v, f = equ.f, u = equ.u, g = equ.g, ϕ = equ.ϕ, ū = equ.ū, ḡ = equ.ḡ, ψ = equ.ψ, v̄ = equ.v̄, f̄ = equ.f̄, poisson = equ.poisson, h = equ.hamiltonian)
    else
        (v = equ.v, f = equ.f, u = equ.u, g = equ.g, ϕ = equ.ϕ, v̄ = equ.v̄, f̄ = equ.f̄, poisson = equ.poisson, h = equ.hamiltonian)
    end
end

function _functions(equ::HDAE, params::OptionalParameters)
    if hassecondary(equ)
        (v = _get_v(equ, params), f = _get_f(equ, params), u = _get_u(equ, params), g = _get_g(equ, params), ϕ = _get_ϕ(equ, params), ū = _get_ū(equ, params), ḡ = _get_ḡ(equ, params), ψ = _get_ψ(equ, params), v̄ = _get_v̄(equ, params), f̄ = _get_f̄(equ, params), poisson = _get_poisson(equ, params), h = _get_h(equ, params))
    else
        (v = _get_v(equ, params), f = _get_f(equ, params), u = _get_u(equ, params), g = _get_g(equ, params), ϕ = _get_ϕ(equ, params), v̄ = _get_v̄(equ, params), f̄ = _get_f̄(equ, params), poisson = _get_poisson(equ, params), h = _get_h(equ, params))
    end
end


text:
##
using Plots
backend(:plotly)
theme(:juno)

##
function finklet(j, l)
    # -------- set filters
    jrad = 7-j
    dθ = π/8        # 8 angular bins hardwired
    θ_l = dθ*l
    # -------- define coordinates
    nx = 256
    xbox = LinRange(-nx/2, nx/2-1 , nx)
    # make a 256x256 grid of X
    sx = xbox' .* ones(nx)
    sy = ones(nx)' .* xbox
    r  = sqrt.((sx).^2 + (sy).^2)
    θ  = mod.(atan.(sy, sx).+π .-θ_l,2*π)
    nozeros = r .> 0
    logr = log2.(r[nozeros])
    r[nozeros] = logr
    # -------- in Fourier plane, envelope of psi_j,l
    mask = (abs.(θ.-π).<= dθ) .& (abs.(r.-jrad) .<= 1)
    # -------- angular part
    ang = cos.((θ.-π).*4)
    # -------- radial part
    rad = cos.((r.-jrad).*π./2)
    psi = mask.*ang.*rad             #mask times angular part times radial part
    return psi
end

##
temp = finklet(1,0)
Plots.heatmap(temp, aspectratio=1)

##
using FFTW

##
J = 8
L = 16

fink_filter = Array{Float64, 4}(undef, J, L, 256, 256)
for j = 1:J
    for l = 1:L
        fink_filter[j,l,:,:]=fftshift(finklet(j-1,l-1))
    end
end

Plots.heatmap(fftshift(fink_filter[3,16,:,:]), aspectratio=1)

fink_filter_sum = zeros(256, 256)
for j = 1:J
    for l = 1:L
        fink_filter_sum+=(finklet(j-1,l-1)).^2
    end
end

Plots.heatmap(fink_filter_sum, aspectratio=1)

fink_filter_sum[127:129,127:129]

##
Plots.heatmap(fftshift(fink_filter[8,4,:,:]), aspectratio=1)

(fftshift(fink_filter[8,4,:,:]))[126:130,126:130]

fink_filter_1 = Array{Float64, 4}(undef, J, L, 256, 256)
for j = 1:J
    for l = 1:L
        fink_filter_1[j,l,:,:]=finklet(j-1,l-1)
    end
end

fink_filter_1[8,3,127:131,127:131]

##
temp2 = [1 2; 3 4]
fftshift(temp2)

blarg = rand(4,4)
fftshift(blarg)

fink_filter_sum = zeros(256, 256)
for j = J
    for l = 1:L
        fink_filter_sum+=(fink_filter_1[j,l,:,:]).^2
    end
end

##
Plots.heatmap(fink_filter_sum, aspectratio=1)

fink_filter_sum[127:131,127:131]

phi = sqrt.(fink_filter_sum)

##
phi_real = ifft(fftshift(phi))

Plots.heatmap(phi, aspectratio=1)

Plots.heatmap(abs.(phi_real), aspectratio=1)

Plots.heatmap(real.(phi_real), aspectratio=1)

Plots.heatmap(imag.(phi_real), aspectratio=1)

function fink_filter_bank(J,L)
    fink_filter = Array{Float64, 4}(undef, J, L, 256, 256)
    for j = 1:J-1
        for l = 1:L
            fink_filter[j,l,:,:]=fftshift(finklet(j-1,l-1))
        end
    end
end

@time fink_filter_bank(8,16)

@time fink_filter_bank(8,16)

##
using ColorSchemes

Plots.heatmap(abs.(phi_real), aspectratio=1,c=:cubehelix)

colorschemes[:cubehelix]

findcolorscheme("cube")

##

Plots.heatmap(abs.(phi_real),
    xlims=(1,256),
    ylims=(1,256),
    aspectratio=1,
    axis=nothing,
    border=:none,
    c=:cubehelix)

using Profile

##
Profile.clear()
@profile fink_filter_bank(8,16)

Juno.profiler()

Profile.print()

@time fink_filter_bank(8,16)

max_stars_repo_name: andrew-saydjari/DHC

max_issues_repo_name: andrew-saydjari/DHC

max_forks_repo_name: andrew-saydjari/DHC


text:
#=------------------------------------------------------------------------------
   Numerical Routines that are useful for tensors but are not a part of the
   core functionality. Though one could create a seperate package implementing
   these, it would be more effective to store common functions in this package
   that can be optimized for the COOTensor package specifically.

  Operations Glossary
  -------------------
    * Dynamical
------------------------------------------------------------------------------=#

"""-----------------------------------------------------------------------------
    Dynamical_System_Solver(A,x0,h,tol,m=1,update=0)

  This function takes in a sparse symmetric tensor and computes the tensor
eigenvector by solving the dynamical system formed by contracting the tensor
A into a vector, and then computing the mth largest eigenvector of the matrix
Ax^{k-1}. With this we form the dynamical system
					dxdt = :Lambda_m(Ax^{k+1})
which we solve with a forward Euler scheme with a step size of h, starting at
the point x0. We solve this until the norm of dxdt reaches a specified
tolerance.

Input:
------
* A - (SSSTensor or Array{AbstractFloat,k > 2})

    The symmetric tensor to compute the eigenvector of. Dense tensors have
    limited support, primary support is for constructors to the sparse cases.

* x0 - (Array{AbstractFloat,1})

    The initial starting point for solving the dynamical system.
* h - (Float64)

    The step size for running the forward Euler scheme.
* tol - (Float64)

    The tolerance to solve the dynamical system up to, stops when
    norm(dxdt) < tol.
* m - (Int)

    The eigenvector to find when computing the dynamical system. Default is the
    largest eigenvector of the matrix.
* update - (Int)

    Indicates how many steps the program should update the user by, default
    value is 0.

Output:
-------
* x - (Array{AbstractFloat,1}

    The resulting eigenvector computed by the dynamical system method.
* lambda - (AbstractFloat)

    The resulting eigenvalue computed by the dynamical system method.
-----------------------------------------------------------------------------"""
function Dynamical_System_Solver(A::SSSTensor,x0::Array{N,1},h::Float64,
                                 tol::Float64,m::Int64 = 1, update::Int = 0) where
								 N <: AbstractFloat
                                 #start with strictly positive random vector
                                 #check the monotonic dereasing property
  n = A.cubical_dimension
  @assert m <= n
  @assert length(x0) == n
  k = order(A)
  x = copy(x0)/norm(x0)
  step = 1

  while true
    A_x_k_2 = contract(A,x,k-2)
    _,V,_ = eigs(A_x_k_2,nev=m) # check eigs
    dxdt = sign(V[1,m])*V[:,m] - x
	#x /= norm(x)

	if norm(dxdt) <= tol
      return x, x'*A_x_k_2*x , step
    else
      x += h*dxdt
    end

	if update > 0
	  if step % update == 0
	    z = A_x_k_2*x
		lambda = dot(x,z)
		residual = z - lambda*x
	    @printf("step %5d: norm(dxdt) = %.16f | lambda = %.16f | res norm = %.16f \n",
		       step, norm(dxdt),lambda, norm(residual))
	  end
	end
    step += 1
  end
end

#figure out if it's faster to do an internal check than two overloaded functions

function Dynamical_System_Solver(A::Array{N,k},x0::Array{N,1},h::Float64,
                                 tol::Float64,m::Int64 = 1) where {N <: AbstractFloat,k}
  #k = length(size(A))
  n = size(A)[1]
  @assert m <= n
  @assert length(x0) == n
  x = copy(x0)

  while true
    _,V = eigen(contract(A,x,k-2))
    dxdt = sign(real(V[1,m]))*real(V[:,m]) - x
	x /= norm(x)

	if norm(dxdt) <= tol
      return x
    else
      x += h*dxdt
    end
  end
end

"""-----------------------------------------------------------------------------
    SSHOPM(A,x_0,shift,max_iter,tol)

  This function runs the shifts symmetric higher order power method for a super
symmetric tensor with the passed in shift, up to a tolerance or up until a
maximum iteration.

Input:
------
* A - (SSSTensor):

    An instance of the super symmetric tensor class.
* x_0 - (Array{AbstractFloat,1}):

    An initial vector to start the algorithm with.
* shift - (AbstractFloat)

    The shift for the algorithm, can be predetermined to ensure convergence of
    the method.
* max_iter - (Int)

    The maximum number of iterations to run the routine for, prints a warning if
    the method hasn't converged by then.
* tol - (Float)

    The tolerance in difference between subsequent approximate eigenvalues to
    solve the routine up to.

* display - (optional Int)

    Indicates how many iterations in between printing to the user.

Output:
-------
* z - (Array{AbstractFloat,1})

    The final vector produced by the SSHOPM routine.
* lambda_k - (AbstractFloat)

    The final approximate eigenvalue at the last iteration.
* iterations - (Integer)

    The number of iterations the algorithm ran for.
-----------------------------------------------------------------------------"""
function SSHOPM(A::Ten, x_0::Array{N,1},shift::N,max_iter::Int,tol::N,
                display::Int=0) where
				{N <: AbstractFloat, Ten <: AbstractSSTen}
    @assert A.cubical_dimension == length(x_0)

	z = Array{N,1}(undef,A.cubical_dimension)
    x = x_0/norm(x_0)
    iterations = 0
    lambda_k_1 = Inf

    while true

		contract_k_1!(A,x,z)
        if shift != 0
			z .+= shift*x
			if shift < 0
				z .*= -1
			end
        end

        lambda_k = x'*z

        #normalize
		residual = z - lambda_k*x
        z ./= norm(z)

        iterations += 1

		if display != 0
			if iterations % display == 0
				@printf("step = %-3d -- λ_k:% 0.12f -- |λ_k - λ_{k-1}| :%0.12f -- res_norm:%0.12f\n",
						iterations,lambda_k,abs(lambda_k - lambda_k_1),norm(residual))
			end
		end
        if norm(residual) < tol || iterations >= max_iter
            if iterations >= max_iter
                @warn("maximum iterations reached")
            end
            return z, lambda_k, iterations
        else
            lambda_k_1 = lambda_k
            x = copy(z)
        end
    end
end

"""-----------------------------------------------------------------------------
    SSHOPM(indices,vals,n,x_0,shift,max_iter,tol)

  This function runs a low memory usage shifted symmetric higher order power
method for a super symmetric tensor with the passed in shift, up to a tolerance
or up until a maximum iteration.

TODO: update to new SSSTensor class
Input:
------
* indices - (Array{Int,2}):

    Indices associated with
* x_0 - (Array{AbstractFloat,1}):

    An initial vector to start the algorithm with.
* shift - (AbstractFloat)

    The shift for the algorithm, can be predetermined to ensure convergence of
    the method.
* max_iter - (Int)

    The maximum number of iterations to run the routine for, prints a warning if
    the method hasn't converged by then.
* tol - (Float)

    The tolerance in difference between subsequent approximate eigenvalues to
    solve the routine up to.

Output:
-------
* z - (Array{AbstractFloat,1})

    The final vector produced by the SSHOPM routine.
* lambda_k - (AbstractFloat)

    The final approximate eigenvalue at the last iteration.
* iterations - (Integer)

    The number of iterations the algorithm ran for.
-----------------------------------------------------------------------------"""
function SSHOPM(indices::Array{Int,2},vals::Array{N,1},n::Int, x_0::Array{N,1},
                shift::N,max_iter,tol) where N <: AbstractFloat

    x = x_0/norm(x_0)
    iterations = 0
	z = Array{N,1}(undef,n)
    lambda_k_1 = Inf

    while true

		contract_k_1!(indices,vals,x,z)
        if shift != 0

            z += shift*x

			if shift < 0
				x *= -1
			end
        end

        lambda_k = x'*z

		residual = z - lambda_k*x

        #normalize
        z /= norm(z)
        iterations += 1

        if norm(residual) < tol || iterations >= max_iter
            if iterations >= max_iter
                @warn("maximum iterations reached")
            end
            return z, lambda_k, iterations
        else
            lambda_k_1 = lambda_k
            x = z
        end
    end
end

"""-----------------------------------------------------------------------------
    find_shift_for_convergence(A,use_fro)

  This function takes in a tensor and computes a shift bound to ensure
convergence of the SSHOPM. If the contracted tensor is large, this can be
computed with the frobenius norm.

Input:
------
* A - (SSSTensor):

    An instance of a super symmetric tensor class.
* use_fro - (bool):

    Indicates whether or not to use the frobenius norm.

Output:
-------
* shift_bound - (Float)

    A float indicating the lower bound for which the method is guaranteed to
    converge.
-----------------------------------------------------------------------------"""
function find_shift_for_convergence(A::SSSTensor)
    error("unfinished")
    shift_bound = (order(A) -1)
    if use_fro
        shift_bound *= contract(A)
    end

end

"""-----------------------------------------------------------------------------
    HOSVD(A,k)

  Returns the top k singular vectors and values of the Higher Order SVD.

 Notes:
 ------
 Currently outputs the left singular vectors and values from the flattened
 tensor. More functionality can be added in, but the current use case is for the
 starting vectors for the SSHOPM.

TODO: would be better to create a adjoint matvec operation from some variant of
  the contraction

-----------------------------------------------------------------------------"""
function HOSVD(A::SSSTensor,k::Int)

	(U,S,_),_ = svds(flatten(A),nsv=k)
	return U,S

end


text:
@testset "Permutation" begin
    @testset "permutedims" begin
        @testset "Space" begin
            S = Space(4, 3, Symmetry(2, 2))

            @test (@inferred permutedims(S, Val((1,2,3,4)))) == Space(4, 3, Symmetry(2, 2))
            @test (@inferred permutedims(S, Val((1,2,4,3)))) == Space(4, 3, Symmetry(2, 2))
            @test (@inferred permutedims(S, Val((1,3,2,4)))) == Space(4, 2, 3, 2)
            @test (@inferred permutedims(S, Val((1,3,4,2)))) == Space(4, Symmetry(2, 2), 3)
            @test (@inferred permutedims(S, Val((1,4,2,3)))) == Space(4, 2, 3, 2)
            @test (@inferred permutedims(S, Val((1,4,3,2)))) == Space(4, Symmetry(2, 2), 3)

            @test (@inferred permutedims(S, Val((2,1,3,4)))) == Space(3, 4, Symmetry(2, 2))
            @test (@inferred permutedims(S, Val((2,1,4,3)))) == Space(3, 4, Symmetry(2, 2))
            @test (@inferred permutedims(S, Val((2,3,1,4)))) == Space(3, 2, 4, 2)
            @test (@inferred permutedims(S, Val((2,3,4,1)))) == Space(3, Symmetry(2, 2), 4)
            @test (@inferred permutedims(S, Val((2,4,1,3)))) == Space(3, 2, 4, 2)
            @test (@inferred permutedims(S, Val((2,4,3,1)))) == Space(3, Symmetry(2, 2), 4)

            @test (@inferred permutedims(S, Val((3,1,2,4)))) == Space(2, 4, 3, 2)
            @test (@inferred permutedims(S, Val((3,1,4,2)))) == Space(2, 4, 2, 3)
            @test (@inferred permutedims(S, Val((3,2,1,4)))) == Space(2, 3, 4, 2)
            @test (@inferred permutedims(S, Val((3,2,4,1)))) == Space(2, 3, 2, 4)
            @test (@inferred permutedims(S, Val((3,4,1,2)))) == Space(Symmetry(2, 2), 4, 3)
            @test (@inferred permutedims(S, Val((3,4,2,1)))) == Space(Symmetry(2, 2), 3, 4)

            @test (@inferred permutedims(S, Val((4,1,2,3)))) == Space(2, 4, 3, 2)
            @test (@inferred permutedims(S, Val((4,1,3,2)))) == Space(2, 4, 2, 3)
            @test (@inferred permutedims(S, Val((4,2,1,3)))) == Space(2, 3, 4, 2)
            @test (@inferred permutedims(S, Val((4,2,3,1)))) == Space(2, 3, 2, 4)
            @test (@inferred permutedims(S, Val((4,3,1,2)))) == Space(Symmetry(2, 2), 4, 3)
            @test (@inferred permutedims(S, Val((4,3,2,1)))) == Space(Symmetry(2, 2), 3, 4)
        end
        @testset "Tensor" begin
            for T in (Float32, Float64)
                x = rand(Tensorial.tensortype(Space(4, 3, Symmetry(2, 2))){T})

                @test (@inferred permutedims(x, Val((1,2,3,4))))::Tensor{Tuple{4, 3, @Symmetry{2, 2}}} ≈ permutedims(x, [1,2,3,4])
                @test (@inferred permutedims(x, Val((1,2,4,3))))::Tensor{Tuple{4, 3, @Symmetry{2, 2}}} ≈ permutedims(x, [1,2,4,3])
                @test (@inferred permutedims(x, Val((1,3,2,4))))::Tensor{Tuple{4, 2, 3, 2}} ≈ permutedims(x, [1,3,2,4])
                @test (@inferred permutedims(x, Val((1,3,4,2))))::Tensor{Tuple{4, @Symmetry{2, 2}, 3}} ≈ permutedims(x, [1,3,4,2])
                @test (@inferred permutedims(x, Val((1,4,2,3))))::Tensor{Tuple{4, 2, 3, 2}} ≈ permutedims(x, [1,4,2,3])
                @test (@inferred permutedims(x, Val((1,4,3,2))))::Tensor{Tuple{4, @Symmetry{2, 2}, 3}} ≈ permutedims(x, [1,4,3,2])

                @test (@inferred permutedims(x, Val((2,1,3,4))))::Tensor{Tuple{3, 4, @Symmetry{2, 2}}} ≈ permutedims(x, [2,1,3,4])
                @test (@inferred permutedims(x, Val((2,1,4,3))))::Tensor{Tuple{3, 4, @Symmetry{2, 2}}} ≈ permutedims(x, [2,1,4,3])
                @test (@inferred permutedims(x, Val((2,3,1,4))))::Tensor{Tuple{3, 2, 4, 2}} ≈ permutedims(x, [2,3,1,4])
                @test (@inferred permutedims(x, Val((2,3,4,1))))::Tensor{Tuple{3, @Symmetry{2, 2}, 4}} ≈ permutedims(x, [2,3,4,1])
                @test (@inferred permutedims(x, Val((2,4,1,3))))::Tensor{Tuple{3, 2, 4, 2}} ≈ permutedims(x, [2,4,1,3])
                @test (@inferred permutedims(x, Val((2,4,3,1))))::Tensor{Tuple{3, @Symmetry{2, 2}, 4}} ≈ permutedims(x, [2,4,3,1])

                @test (@inferred permutedims(x, Val((3,1,2,4))))::Tensor{Tuple{2, 4, 3, 2}} ≈ permutedims(x, [3,1,2,4])
                @test (@inferred permutedims(x, Val((3,1,4,2))))::Tensor{Tuple{2, 4, 2, 3}} ≈ permutedims(x, [3,1,4,2])
                @test (@inferred permutedims(x, Val((3,2,1,4))))::Tensor{Tuple{2, 3, 4, 2}} ≈ permutedims(x, [3,2,1,4])
                @test (@inferred permutedims(x, Val((3,2,4,1))))::Tensor{Tuple{2, 3, 2, 4}} ≈ permutedims(x, [3,2,4,1])
                @test (@inferred permutedims(x, Val((3,4,1,2))))::Tensor{Tuple{@Symmetry{2, 2}, 4, 3}} ≈ permutedims(x, [3,4,1,2])
                @test (@inferred permutedims(x, Val((3,4,2,1))))::Tensor{Tuple{@Symmetry{2, 2}, 3, 4}} ≈ permutedims(x, [3,4,2,1])

                @test (@inferred permutedims(x, Val((4,1,2,3))))::Tensor{Tuple{2, 4, 3, 2}} ≈ permutedims(x, [4,1,2,3])
                @test (@inferred permutedims(x, Val((4,1,3,2))))::Tensor{Tuple{2, 4, 2, 3}} ≈ permutedims(x, [4,1,3,2])
                @test (@inferred permutedims(x, Val((4,2,1,3))))::Tensor{Tuple{2, 3, 4, 2}} ≈ permutedims(x, [4,2,1,3])
                @test (@inferred permutedims(x, Val((4,2,3,1))))::Tensor{Tuple{2, 3, 2, 4}} ≈ permutedims(x, [4,2,3,1])
                @test (@inferred permutedims(x, Val((4,3,1,2))))::Tensor{Tuple{@Symmetry{2, 2}, 4, 3}} ≈ permutedims(x, [4,3,1,2])
                @test (@inferred permutedims(x, Val((4,3,2,1))))::Tensor{Tuple{@Symmetry{2, 2}, 3, 4}} ≈ permutedims(x, [4,3,2,1])
            end
        end
    end
end


text:
module problem111

description = """
Primes with runs

Considering 4-digit primes containing repeated digits it is clear that they cannot all be the same: 1111 is divisible by 11, 2222 is divisible by 22, and so on. But there are nine 4-digit primes containing three ones:
1117, 1151, 1171, 1181, 1511, 1811, 2111, 4111, 8111

We shall say that M(n, d) represents the maximum number of repeated digits for an n-digit prime where d is the repeated digit, N(n, d) represents the number of such primes, and S(n, d) represents the sum of these primes.
So M(4, 1) = 3 is the maximum number of repeated digits for a 4-digit prime where one is the repeated digit, there are N(4, 1) = 9 such primes, and the sum of these primes is S(4, 1) = 22275. It turns out that for d = 0, it is only possible to have M(4, 0) = 2 repeated digits, but there are N(4, 0) = 13 such cases.

In the same way we obtain the following results for 4-digit primes.
Digit, d 	M(4, d) 	N(4, d) 	S(4, d)
0 	2 	13 	67061
1 	3 	9 	22275
2 	3 	1 	2221
3 	3 	12 	46214
4 	3 	2 	8888
5 	3 	1 	5557
6 	3 	1 	6661
7 	3 	9 	57863
8 	3 	1 	8887
9 	3 	7 	48073

For d = 0 to 9, the sum of all S(4, d) is 273700.

Find the sum of all S(10, d).
"""

function solve(n :: Integer)
  return sum([S(n, d) for d = 0:9])
end

# search for the first r from (n-1) to 2 producing some primes having the digit d repeat r times in n-digit numbers.
function S(n :: Integer, d :: Integer)
  for r = n-1 : -1 : 2
    result = S(n, d, r)
    if result > 0
      return result
    end
  end
end

# Find the sum of primes with n digits, having digit d repeat r times.
function S(n :: Integer, d :: Integer, r :: Integer)
  total = 0
  repeatCombinations = collect(combinations([1 : (d == 0 ? n-1 : n)], r)) # no leading zeros
  for repeated in repeatCombinations
    digits = zeros(n)
    digits[repeated] = d
    changing = setdiff([1:n], repeated)
    total += SPrimes(digits, changing, 1)
  end
  return total
end

function SPrimes(digits :: Vector, changing :: Vector, i :: Integer)

  # Base case: all digits are set, check if result is prime
  if i > length(changing)
    n = sum([int(10^(x-1) * digits[x]) for x in 1:length(digits)])
    return isprime(n) ? n : 0
  end

  # recursive case, set digit at index changing[i], and call with i+1
  total = 0
  for d = (changing[i] == length(digits) ? 1 : 0) : 9
    digits[changing[i]] = d
    total += SPrimes(digits, changing, i+1)
  end
  return total
end

using Base.Test

@test solve(4) == 273700

end

max_stars_repo_name: mbuhot/mbuhot-euler-solutions

max_issues_repo_name: mbuhot/mbuhot-euler-solutions

max_forks_repo_name: mbuhot/mbuhot-euler-solutions