using LinearAlgebra,Distributions,DataFrames,StatsBase,GLM,LsqFit,Optim,BlackBoxOptim,HypothesisTests,Colors,Images,StatsModels,CategoricalArrays,
ImageFiltering,SpecialFunctions,DSP,HCubature,Combinatorics,DataStructures,ANOVA,StatsFuns,Trapz,CircStats,ImageSegmentation,ProgressMeter,PyCall,
Dierckx,BandedMatrices,OffsetArrays
import Base: vec,range
import StatsBase: predict

include("NeuroDataType.jl")
include("Stats.jl")
include("Spike.jl")
include("LFP.jl")
include("Image.jl")
include("Condition.jl")
include("2P.jl")

vec(x::RGBA)=[x.r,x.g,x.b,x.alpha]
anscombe(x) = 2*sqrt(x+(3/8))

"`Gaussian` function"
gaussianf(x;a=1,μ=0,σ=1) = a*exp(-0.5((x-μ)/σ)^2)
function gaussianf(x,y;a=1,μ₁=0,σ₁=1,μ₂=0,σ₂=1,θ=0)
    sinv,cosv = sincos(θ)
    x₀ = x-μ₁
    y₀ = y-μ₂
    x′ = cosv * x₀ + sinv * y₀
    y′ = cosv * y₀ - sinv * x₀
    a*exp(-0.5((x′/σ₁)^2 + (y′/σ₂)^2))
end

"""
`von Mises` function [^1]

```math
f(α) =  βe^{κ(cos(n(α - μ)) - 1)}
```

[^1]

Swindale, N.V. (1998). Orientation tuning curves: empirical description and estimation of parameters. Biol Cybern 78, 45–56.

- β: amplitude at μ
- μ: angle of peak
- κ: width parameter
- n: frequency parameter
"""
vmf(α;β=1,μ=0,κ=1,n=1) = β*exp(κ*(cos(n*(α-μ))-1))
"""
`Generalized von Mises` function [^1]

```math
f(α) =  βe^{κ₁(cos(α - μ₁) - 1) + κ₂(cos2(α - μ₂) - 1)}
```

[^1]

Gatto, R., and Jammalamadaka, S.R. (2007). The generalized von Mises distribution. Statistical Methodology 4, 341–353.
"""
gvmf(α;β=1,μ₁=0,κ₁=1,μ₂=0,κ₂=1) = β*exp(κ₁*cos(α-μ₁) + κ₂*cos(2(α-μ₂)))

"""
`Difference of Gaussians` function
"""
dogf(x;aₑ=2,μₑ=0,σₑ=1,aᵢ=1,μᵢ=0,σᵢ=2) = aₑ*exp(-0.5((x-μₑ)/σₑ)^2) - aᵢ*exp(-0.5((x-μᵢ)/σᵢ)^2)
function dogf(x,y;aₑ=2,μₑ₁=0,σₑ₁=1,μₑ₂=0,σₑ₂=1,θₑ=0,aᵢ=1,μᵢ₁=0,σᵢ₁=2,μᵢ₂=0,σᵢ₂=2,θᵢ=0)
    sinvₑ,cosvₑ = sincos(θₑ)
    xₑ₀ = x-μₑ₁
    yₑ₀ = y-μₑ₂
    xₑ′ = cosvₑ * xₑ₀ + sinvₑ * yₑ₀
    yₑ′ = cosvₑ * yₑ₀ - sinvₑ * xₑ₀
    sinvᵢ,cosvᵢ = sincos(θᵢ)
    xᵢ₀ = x-μᵢ₁
    yᵢ₀ = y-μᵢ₂
    xᵢ′ = cosvᵢ * xᵢ₀ + sinvᵢ * yᵢ₀
    yᵢ′ = cosvᵢ * yᵢ₀ - sinvᵢ * xᵢ₀
    aₑ*exp(-0.5((xₑ′/σₑ₁)^2 + (yₑ′/σₑ₂)^2)) - aᵢ*exp(-0.5((xᵢ′/σᵢ₁)^2 + (yᵢ′/σᵢ₂)^2))
end

"""
`sin` grating function

- μ: x offset
- f: Frequency in cycle/unit_x
- phase: Phase of a cycle in [0, 1] scale
"""
gratingf(x;μ=0, f=1, phase=0) = sin(2π * (f * (x-μ) + phase))

"""
2D `sin` grating function

- μ₁: x offset
- μ₂: y offset
- θ: Orientation in radius, 0 is -, increase counter-clock wise
- f: Frequency in cycle/unit_x/y
- phase: Phase of a cycle in [0, 1] scale
"""
function gratingf(x,y;μ₁=0,μ₂=0,θ=0,f=1,phase=0)
    sinθ,cosθ = sincos(θ)
    y′ = cosθ * (y-μ₂) - sinθ * (x-μ₁)
    sin(2π * (f * y′ + phase))
end

"""
`cas` function defined as ``cas(x) = cos(x) + sin(x)``

- f: Frequency in cycle/unit_x
- phase: Phase of a cycle in [0, 1] scale
- isnorm: scale `cas` in [-√2, √2] to [-1, 1]
"""
function cas(x;f=1, phase=0, isnorm::Bool=true)
    r = sum(sincos(2π * (f * x + phase)))
    if isnorm
        r /=sqrt(2)
    end
    return r
end

"""
2D `cas` function defined as ``cas(x+y) = cos(x+y) + sin(x+y)``

- kx: Frequency in cycle/unit_x
- ky: Frequency in cycle/unit_y
- phase: Phase of a cycle in [0, 1] scale
- isnorm: scale `cas` in [-√2, √2] to [-1, 1]
"""
function cas(x,y;kx=1,ky=1, phase=0, isnorm::Bool=true)
    r = sum(sincos(2π * (kx * x + ky * y + phase)))
    isnorm && (r /=sqrt(2))
    return r
end

"""
`cas` phase to `sin` phase, phase is in [0, 1] scale

!!! note
    ``cas(x) = √2 sin(x + π/4)``
"""
cas2sin(phase) = phase + 0.125

"""
2D `cas` to 2D sin `gratingf`
"""
function cas2sin(kx,ky,phase)
    θ = atan(ky,kx) - π/2
    f = sqrt(kx*kx + ky*ky)
    return (θ=θ,f=f,phase=phase + 0.125)
end

"""
`sin` phase to `cas` phase, phase is in [0, 1] scale

!!! note
    ``cas(x) = √2 sin(x + π/4)``
"""
sin2cas(phase) = phase - 0.125

"""
2D sin `gratingf` to 2D `cas`
"""
function sin2cas(θ,f,phase)
    sinθ,cosθ = sincos(θ + π/2)
    kx = cosθ*f
    ky = sinθ*f
    return (kx=kx,ky=ky,phase=phase - 0.125)
end

"`Gabor` function"
gaborf(x;a=1,μ=0,σ=1,f=1,phase=0) = a*exp(-0.5((x-μ)/σ)^2)*sin(2π*(f*(x-μ)+phase))
function gaborf(x,y;a=1,μ₁=0,σ₁=1,μ₂=0,σ₂=1,θ=0,f=1,phase=0)
    sinv,cosv = sincos(θ)
    x₀ = x-μ₁
    y₀ = y-μ₂
    x′ = cosv * x₀ + sinv * y₀
    y′ = cosv * y₀ - sinv * x₀
    a*exp(-0.5((x′/σ₁)^2 + (y′/σ₂)^2)) * sin(2π*(f * y′ + phase))
end

"Binary mask for gabor envelope"
function gaborenvelopemask(x,y;fσ=2.5,μ₁=0,σ₁=1,μ₂=0,σ₂=1,θ=0)
    sinv,cosv = sincos(θ)
    x₀ = x-μ₁
    y₀ = y-μ₂
    x′ = cosv * x₀ + sinv * y₀
    y′ = cosv * y₀ - sinv * x₀
    (x′/fσ/σ₁)^2 + (y′/fσ/σ₂)^2 <= 1 ? true : false
end
"Binary mask for concentric circular dog envelope"
dogenvelopemask(x,y;fσ=2.5,μ₁=0,σₑ₁=1,μ₂=0,rσᵢₑ=2) = edogenvelopemask(x,y;fσ,μ₁,σₑ₁,rσ₂₁=1,μ₂,rσᵢₑ,θ=0)
"Binary mask for concentric orientated elliptical dog envelope"
function edogenvelopemask(x,y;fσ=2.5,μ₁=0,σₑ₁=1,rσ₂₁=1,μ₂=0,rσᵢₑ=2,θ=0)
    sinv,cosv = sincos(θ)
    x₀ = x-μ₁
    y₀ = y-μ₂
    x′ = cosv * x₀ + sinv * y₀
    y′ = cosv * y₀ - sinv * x₀
    σ₁ = max(σₑ₁,rσᵢₑ*σₑ₁)
    σ₂ = σ₁ * rσ₂₁
    (x′/fσ/σ₁)^2 + (y′/fσ/σ₂)^2 <= 1 ? true : false
end

"Fit 1D model to data"
function fitmodel(model,x,y)
    lb,ub = extrema(y)
    bm = (lb+ub)/2
    br = (ub-lb)/2
    alb,aub = abs.((lb,ub))
    ab = max(alb,aub)

    xlb,xub = extrema(x)
    xbm = (xlb+xub)/2
    xalb,xaub = abs.((xlb,xub))
    xab = max(xalb,xaub)

    rlt = fun = missing
    if model == :vmn2
        fun = (x,p) -> vmf.(x,β=p[1],μ=p[2],κ=p[3],n=2)
        ofun = (p;x=x,y=y) -> sum((y.-fun(x,p)).^2)

        ub=[1.8ab,   prevfloat(float(π)),   200]
        lb=[0.2ab,            0,              0]
        p0=[ab,               0,              1]
    elseif model == :gvm
        # fun = (x,p) -> gvmf.(x,β=p[1],μ₁=p[2],κ₁=p[3],μ₂=p[4],κ₂=p[5])
        fun = gvmff
        ofun = (p;x=x,y=y) -> sum((y.-fun(x,p)).^2)

        ub=[1.8ab,   prevfloat(float(2π)),   40,    prevfloat(float(π)),     40]
        lb=[nextfloat(0.0),   0,              0,             0,               0]
        p0=[ab,               0,              1,             0,               1]
    elseif model == :dog
        fun = (x,p) -> dogf.(x,aₑ=p[1],μₑ=p[2],σₑ=p[3],aᵢ=p[4],μᵢ=p[5],σᵢ=p[6])
        ofun = (p;x=x,y=y) -> sum((y.-fun(x,p)).^2)

        ub=[1.8ab,   10xab,            10xab,    1.8ab,     10xab,              10xab]
        lb=[0,      -10xab,   nextfloat(0.0),        0,    -10xab,     nextfloat(0.0)]
        p0=[ab,        0,                  1,       ab,         0,                  1]
    elseif model == :sfdog
        # fun = (x,p) -> dogf.(x,aₑ=p[1],μₑ=p[2],σₑ=p[3],aᵢ=p[4],μᵢ=p[5],σᵢ=p[6]) .+ p[7]
        fun = sfdogff
        ofun = (p;x=x,y=y) -> sum((y.-fun(x,p)).^2)

        ub=[1.5ab,    10,             10,          1.5ab,       10,                10,         bm+br/3]
        lb=[0,         0,      nextfloat(0.0),       0,          0,        nextfloat(0.0),        0]
        p0=[ab,     0.5xab,         0.5xab,          0,        0.5xab,           0.5xab,       bm-br/3]
    elseif model == :sfgaussian
        # fun = (x,p) -> gaussianf.(log2.(x),a=p[1],μ=p[2],σ=p[3]) .+ p[4]
        fun = sfgaussianff
        ofun = (p;x=x,y=y) -> sum((y.-fun(x,p)).^2)

        # limit gaussian center ~[-8 10], sigma ~[1.4 8]
        ub=[1.5ab,     3.3,          3,       bm]
        lb=[0,          -3,        0.5,        0]
        p0=[ab,          0,          1,        0]
    end
    if !ismissing(fun)
        # ofit = optimize(ofun,lb,ub,p0,SAMIN(rt=0.92),Optim.Options(iterations=220000))
        # param=ofit.minimizer
        ofit = bboptimize(ofun,p0;SearchRange=collect(zip(lb,ub)),Method=:adaptive_de_rand_1_bin_radiuslimited,MaxSteps=200000)
        param = best_candidate(ofit)

        rlt = (;model,fun,param, goodnessoffit(y,fun(x,param),k=length(param))...)
    end
    return rlt
end
gvmff(x,p) = gvmf.(x,β=p[1],μ₁=p[2],κ₁=p[3],μ₂=p[4],κ₂=p[5])
sfdogff(x,p) = dogf.(x,aₑ=p[1],μₑ=p[2],σₑ=p[3],aᵢ=p[4],μᵢ=p[5],σᵢ=p[6]) .+ p[7]
sfgaussianff(x,p) = gaussianf.(log2.(x),a=p[1],μ=p[2],σ=p[3]) .+ p[4]

"Fit 2D model to image"
function fitmodel2(model,data::Matrix,ppu;w=0.5)
    rspx = (size(data).-1)./2
    radii = rspx./ppu
    x = (mapreduce(i->[i[2] -i[1]],vcat,CartesianIndices(data)) .+ [-(rspx[2]+1) (rspx[1]+1)])/ppu
    y = vec(data)

    # try estimate solution
    roi = peakroi(localcontrast(data,round(Int,w*ppu)))
    alb,aub = abs.(extrema(data[roi.i]))
    ab = max(alb,aub)
    r = roi.radius/ppu
    c = [roi.center[2] - (rspx[2]+1), -roi.center[1] + (rspx[1]+1)]/ppu

    rlt = fun = missing
    if model == :dog
        if aub >= alb
            ai = 5alb
            ae = aub + ai
        else
            ae = 5aub
            ai = alb + ae
        end
        # concentric circular dog
        # fun = (x,y,p) -> dogf.(x,y,aₑ=p[1],μₑ₁=p[2],σₑ₁=p[3],μₑ₂=p[4],σₑ₂=p[3],θₑ=0,aᵢ=p[5],μᵢ₁=p[2],σᵢ₁=p[6]*p[3],μᵢ₂=p[4],σᵢ₂=p[6]*p[3],θᵢ=0)
        # mfun = (x,y,p) -> dogenvelopemask.(x,y;fσ=2.5,μ₁=p[2],μ₂=p[4],σₑ₁=p[3],rσᵢₑ=p[6])
        fun = dogff
        mfun = dogfmf
        ofun = (p;x=x,y=y) -> @views sum((y.-fun(x[:,1],x[:,2],p)).^2)
        ub=[5ae,    0.5r+c[1],    0.9r,    0.5r+c[2],     5ai,    4]
        lb=[0,     -0.5r+c[1],    0.1r,   -0.5r+c[2],     0,      0.25]
        p0=[ae,      c[1],        0.3r,     c[2],         ai,     1]
    elseif model == :edog
        if aub >= alb
            ai = 5alb
            ae = aub + ai
        else
            ae = 5aub
            ai = alb + ae
        end
        # concentric orientated elliptical dog
        # fun = (x,y,p) -> dogf.(x,y,aₑ=p[1],μₑ₁=p[2],σₑ₁=p[3],μₑ₂=p[4],σₑ₂=p[5]*p[3],θₑ=p[6],aᵢ=p[7],μᵢ₁=p[2],σᵢ₁=p[8]*p[3],μᵢ₂=p[4],σᵢ₂=p[5]*p[8]*p[3],θᵢ=p[6])
        # mfun = (x,y,p) -> edogenvelopemask.(x,y;fσ=2.5,μ₁=p[2],μ₂=p[4],σₑ₁=p[3],rσ₂₁=p[5],rσᵢₑ=p[8],θ=p[6])
        fun = edogff
        mfun = edogfmf
        ofun = (p;x=x,y=y) -> @views sum((y.-fun(x[:,1],x[:,2],p)).^2)
        ub=[5ae,    0.5r+c[1],    0.9r,    0.5r+c[2],       1,       prevfloat(float(π)),    5ai,     4]
        lb=[0,     -0.5r+c[1],    0.1r,   -0.5r+c[2],      0.5,              0,               0,      0.25]
        p0=[ae,      c[1],        0.3r,     c[2],           1,               0,               ai,     1]
    elseif model == :gabor
        # fun = (x,y,p) -> gaborf.(x,y,a=p[1],μ₁=p[2],σ₁=p[3],μ₂=p[4],σ₂=p[5]*p[3],θ=p[6],f=p[7],phase=p[8])
        # mfun = (x,y,p) -> gaborenvelopemask.(x,y;fσ=2.5,μ₁=p[2],μ₂=p[4],σ₁=p[3],σ₂=p[5]*p[3],θ=p[6])
        fun = gaborff
        mfun = gaborfmf
        ofun = (p;x=x,y=y) -> @views sum((y.-fun(x[:,1],x[:,2],p)).^2)

        ori,sf = f1orisf(powerspectrum2(data,ppu)...)
        ub=[5ab,   0.5r+c[1],   0.9r,    0.5r+c[2],    6.0,     prevfloat(float(π)),     12,            prevfloat(1.0)]
        lb=[0,    -0.5r+c[1],   0.1r,   -0.5r+c[2],    0.2,             0,              0.05,                  0]
        p0=[ab,    c[1],        0.3r,     c[2],          1,            ori,     clamp(sf,lb[7],ub[7]),       0.5]
    end
    if !ismissing(fun)
        # ofit = optimize(ofun,lb,ub,p0,SAMIN(rt=0.92),Optim.Options(iterations=220000))
        # param = ofit.minimizer
        ofit = bboptimize(ofun,p0;SearchRange=collect(zip(lb,ub)),Method=:adaptive_de_rand_1_bin_radiuslimited,MaxSteps=200000)
        param = best_candidate(ofit)

        @views rlt = (;model,fun,mfun,param,radii, goodnessoffit(y,fun(x[:,1],x[:,2],param),k=length(param))...)
    end
    return rlt
end
# workaround since JLD2 can not save/load anonymous functions yet
dogff(x,y,p) = dogf.(x,y,aₑ=p[1],μₑ₁=p[2],σₑ₁=p[3],μₑ₂=p[4],σₑ₂=p[3],θₑ=0,aᵢ=p[5],μᵢ₁=p[2],σᵢ₁=p[6]*p[3],μᵢ₂=p[4],σᵢ₂=p[6]*p[3],θᵢ=0)
edogff(x,y,p) = dogf.(x,y,aₑ=p[1],μₑ₁=p[2],σₑ₁=p[3],μₑ₂=p[4],σₑ₂=p[5]*p[3],θₑ=p[6],aᵢ=p[7],μᵢ₁=p[2],σᵢ₁=p[8]*p[3],μᵢ₂=p[4],σᵢ₂=p[5]*p[8]*p[3],θᵢ=p[6])
gaborff(x,y,p) = gaborf.(x,y,a=p[1],μ₁=p[2],σ₁=p[3],μ₂=p[4],σ₂=p[5]*p[3],θ=p[6],f=p[7],phase=p[8])
dogfmf(x,y,p) = dogenvelopemask.(x,y;fσ=2.5,μ₁=p[2],μ₂=p[4],σₑ₁=p[3],rσᵢₑ=p[6])
edogfmf(x,y,p) = edogenvelopemask.(x,y;fσ=2.5,μ₁=p[2],μ₂=p[4],σₑ₁=p[3],rσ₂₁=p[5],rσᵢₑ=p[8],θ=p[6])
gaborfmf(x,y,p) = gaborenvelopemask.(x,y;fσ=2.5,μ₁=p[2],μ₂=p[4],σ₁=p[3],σ₂=p[5]*p[3],θ=p[6])

predict(fit,x) = fit.fun(x,fit.param)
function predict(fit,x,y;xygrid=true,yflip=false)
    if xygrid
        z = [fit.fun(i,j,fit.param) for j in y, i in x]
        yflip && (z=reverse(z,dims=1))
    else
        z = fit.fun(x,y,fit.param)
    end
    z
end

"""
Goodness of Fit Metrics:

- r: Pearson Correlation Coefficient
- mae: Mean Absolute Error
- rmse: Root Mean Squared Error
- rae: Relative Absolute Error
- rse: Relative Squared Error
- r2: R Squared
- adjr2: Adjusted-R²
- s: Residual Standard Error
- aic: Akaike Information Criterion
- bic: Bayesian Information Criterion

1. y: responses
2. ŷ: model predictions

- n: sample size
- e: errors(y - ŷ)
- k: number of predictors
- df: degree of freedom(n - k - 1)
"""
function goodnessoffit(y,ŷ;n = length(y),e = y .- ŷ,k=missing,df = n-k-1)
    r = cor(y,ŷ)
    ae = abs.(e)
    e2 = e.^2
    ssᵣ = sum(e2)
    mae = mean(ae)
    rmse = sqrt(mean(e2))
    ydm = y .- mean(y)
    ssₜ = sum(ydm.^2)
    rae = sum(ae)/sum(abs.(ydm))
    fvu = ssᵣ/ssₜ
    rse = sqrt(fvu)
    r2 = 1 - fvu
    s = ssᵣ/df
    s = s < 0 ? missing : sqrt(s)
    adjr2 = 1 - fvu*(n-1)/df
    aic = n*log(ssᵣ) + 2k
    bic = n*log(ssᵣ/n) + k*log(n) # under iid normal error assumption
    (;r,mae,rmse,rae,rse,r2,adjr2,s,aic,bic)
end

function searchclosest(v,vs;start::Integer=1,step::Integer=1,circ=false)
    n=length(vs);ssign = sign(vs[start]-v)
    i = start
    for _ in 1:n
        if !circ
            i<1 && return -Inf
            n<i && return Inf
        end
        sign(vs[i]-v) != ssign && return i
        i += step
        if circ
            i<1 && (i+=n)
            n<i && (i-=n)
        end
    end
    Inf
end

"""
left and right half width of `v` relative to `y[start]`, -Inf/Inf when no `v` is found.

- start: index of `y` which is the center of the width
- v: value on which width is cutoff
- circ: whether `y` is defined on circular domain and width can wrap around
- x: domain of `y`, return width when `x` provided, otherwise return cutoff indices
"""
function halfwidth(y;start=argmax(y),v=y[start]/2,circ=false,x=nothing)
    li = searchclosest(v,y;start,step=-1,circ)
    ri = searchclosest(v,y;start,step=1,circ)
    if isnothing(x)
        return li,ri
    else
        if circ
            lw = isinf(li) ? li : li<=start ? abs(x[start]-x[li]) : abs(x[start]-x[1])+abs(x[end]-x[li])
            rw = isinf(ri) ? ri : ri>=start ? abs(x[ri]-x[start]) : abs(x[ri]-x[1])+abs(x[end]-x[start])
        else
            lw = isinf(li) ? li : abs(x[start]-x[li])
            rw = isinf(ri) ? ri : abs(x[ri]-x[start])
        end
        return lw,rw
    end
end

circtuningfeature(mfit;od=[π,0.5π],fn=:a,x = 0:0.002:2π) = circtuningfeature(x,predict(mfit,x);od,fn) # 0.002rad ≈ 0.11deg
"""
Properties of Circular Tuning

    - Prefered Angle with Peak Response
    - Selectivity Index
        - version 1: (ResponsePeak - ResponseOpposing)/ResponsePeak
        - version 2: (ResponsePeak - ResponseOpposing)/(ResponsePeak + ResponseOpposing)
    - Half Width at Half Peak-to-Trough

1. x: angles in radius
2. y: responses
- od: opposing angle distance to prefered angle, e.g. π for DSI, 0.5π for OSI
- fn: factor name
"""
function circtuningfeature(x,y;od=[π,0.5π],fn=:a)
    maxr,maxi = findmax(y)
    minr,mini = findmin(y)
    maxx = x[maxi]
    ox = maxx.+od
    _,oi = findclosestangle(x,ox)
    or = y[oi]

    si1 = 1 .- or./maxr
    si2 = (maxr .- or)./(maxr .+ or)
    hw = halfwidth(y,start=maxi,v=(maxr+minr)/2,circ=true,x=x)

    (;Symbol(:p,fn)=>rad2deg(mod2pi(maxx)),Symbol(fn,:hw)=>rad2deg.(hw),Symbol(fn,:si1)=>si1,Symbol(fn,:si2)=>si2,Symbol(fn,:od)=>od)
end

sftuningfeature(mfit;x = 0:0.003:10) = sftuningfeature(x,predict(mfit,x))
"""
Properties of Spatial Frequency Tuning

    - Prefered Spatial Frequency with Peak Response
    - Half Width at Half Peak-to-Trough
    - Freq Passing Type {A:All Pass, H:High Pass, L:Low Pass, B:Band Pass}
    - Bandwidth ``log2(H_cut/L_cut)``
    - Passwidth at Half Peak-to-Trough constrained by low freq lim and high freq lim

1. x: sf in cycle/degree
2. y: responses
"""
function sftuningfeature(x,y;low=minimum(x),high=maximum(x))
    maxr,maxi = findmax(y)
    minr,mini = findmin(y)
    maxx = x[maxi]

    hw = halfwidth(y,start=maxi,v=(maxr+minr)/2,circ=false,x=x)
    pt = all(isinf.(hw)) ? 'A' : isinf(hw[1]) ? 'L' : isinf(hw[2]) ? 'H' : 'B'
    bw = log2((maxx+hw[2])/(maxx-hw[1]))
    pw = pt == 'A' ? high-low : pt == 'L' ? maxx-low+hw[2] : pt == 'H' ? high-maxx+hw[1] : sum(hw)

    (;psf=maxx,sfhw=hw,sftype=pt,sfbw=bw,sfpw=pw)
end

"""
Tuning properties of factor response

1. fl: factor levels
2. fr: factor responses for each level

    Angle, Orientation and Direction follow the same convention such that 0 is -/→, then increase counter-clock wise.

    For cases where Orientation and Direction are interlocked(drifting grating):
        - when Orientation is -(0), then Direction is ↑(90)
        - when Direction is →(0), then Orientation is |(-90)
"""
function factorresponsefeature(fl,fr;fm=mean.(fr),factor=:Ori,isfit::Bool=true)
    i = .!ismissing.(fr)
    fl = fl[i];fr=fr[i];fm=fm[i]
    ls = mapreduce((l,r)->fill(l,length(r)),append!,fl,fr)
    rs = mapreduce(deepcopy,append!,fr)

    if factor in [:Ori,:Ori_Final]
        θ = mod2pi.(deg2rad.(ls))
        α = mod2pi.(deg2rad.(fl))
        d = mean(diff(sort(unique(α)))) # angle spacing
        up, = circ_otest(α,w=fm) # Omnibus test for non-uniformity
        # for orientation
        oα = mod.(α,π)
        ocv, = circ_var(2oα,w=fm,d=2d)
        ocm, = circ_mean(2oα,w=fm)
        ocm = rad2deg(mod2pi(ocm)/2)
        # for direction
        dcv, = circ_var(α;w=fm,d)
        dcm, = circ_mean(α,w=fm)
        dcm = rad2deg(mod2pi(dcm+0.5π))
        maxr,maxi = findmax(fm)
        maxl = fl[maxi]

        fit = ()
        if isfit
            try
                mfit = fitmodel(:gvm,θ,rs) # fit Generalized von Mises
                fit = (;circtuningfeature(mfit,od=[π,0.5π],fn=:o)...,mfit)
            catch
                display.(stacktrace(catch_backtrace()))
            end
        end

        return (;up,ocv,ocm,dcv,dcm,max=maxl=>maxr,fit)
    elseif factor == :Dir
        θ = deg2rad.(fl)
        d = mean(diff(sort(unique(θ)))) # angle spacing
        # for orientation
        oθ = mod.(θ.-0.5π,π)
        om, = circ_mean(2oθ,w=fr)
        oo = rad2deg(mod(angle(om),2π)/2)
        ocv = circ_var(2oθ,w=fr,d=2d)
        # for direction
        dm, = circ_mean(θ,w=fr)
        od = rad2deg(mod(angle(dm),2π))
        dcv = circ_var(θ;w=fr,d)
        # fit Generalized von Mises for direction
        fit = ()
        if isfit
            try
                mfit = fitmodel(:gvm,θ,fr)
                fit = (circtuningfeature(mfit,od=π,fn=:d)...,gvm=mfit)
            catch
            end
            # fit von Mises for orientation
            try
                mfit = fitmodel(:vmn2, θ.-0.5π,fr)
                fit =(fit...,circtuningfeature(mfit, od=0.5π,fn=:o)...,vmn2=mfit)
            catch
            end
        end

        return (;dm,od,dcv,om,oo,ocv,fit)
    elseif factor == :SpatialFreq
        up = PyOnewayANOVA.anova_oneway(fr,use_var="unequal").pvalue
        msf = 2^(sum(rs.*log2.(ls))/sum(rs)) # weighted average
        maxr,maxi = findmax(fm)
        maxl = fl[maxi]

        fit = ()
        if isfit
            try
                # mfit = fitmodel(:sfdog,ls,rs) # fit Difference of Gaussians
                mfit = fitmodel(:sfgaussian,ls,rs) # fit Gaussian of logarithm sf
                fit = (;sftuningfeature(mfit)...,mfit)
            catch
                display.(stacktrace(catch_backtrace()))
            end
        end

        return (;up,msf,max=maxl=>maxr,fit)
    elseif factor == :ColorID
        # transform colorId to hue angle
        ucid = sort(unique(fl))
        hstep = 2pi/length(ucid)
        ha = map(l->hstep*(findfirst(c->c==l,ucid)-1),fl)
        # for hue direction
        hm, = circ_mean(ha,w=fr)
        oh = mod(rad2deg(angle(hm)),360)
        hcv = circ_var(ha,w=fr,d=hstep)
        maxi = argmax(fr)
        maxh = rad2deg(ha[maxi])
        maxr = fr[maxi]

        fit = ()
        if isfit
            # fit Generalized von Mises for hue
            try
                mfit = fitmodel(:gvm,θ,fr)
                fit = (circtuningfeature(mfit,od=π,fn=:h)...,gvm=mfit)
            catch
            end
            # fit von Mises for hue axis
            try
                mfit = fitmodel(:vmn2,θ,fr)
                fit = (fit...,circtuningfeature(mfit,od=0.5π,fn=:ha)...,vmn2=mfit)
            catch
            end
        end

        return (;ham,oha,hacv,hm,oh,hcv,maxh,maxr,fit)
    elseif factor in [:HueAngle,:Angle]
        θ = mod2pi.(deg2rad.(ls))
        α = mod2pi.(deg2rad.(fl))
        d = mean(diff(sort(unique(α)))) # angle spacing
        up, = circ_otest(α,w=fm) # Omnibus test for non-uniformity
        # for axis
        aα = mod.(α,π)
        acv, = circ_var(2aα,w=fm,d=2d)
        acm, = circ_mean(2aα,w=fm)
        acm = rad2deg(mod2pi(acm)/2)
        # for angle
        cv, = circ_var(α;w=fm,d)
        cm, = circ_mean(α,w=fm)
        cm = rad2deg(mod2pi(cm))
        maxr,maxi = findmax(fm)
        maxl = fl[maxi]

        fit = ()
        if isfit
            try
                mfit = fitmodel(:gvm,θ,rs) # fit Generalized von Mises
                fit = (;circtuningfeature(mfit,od=[π,0.5π],fn=:a)...,mfit)
            catch
                display.(stacktrace(catch_backtrace()))
            end
        end

        return (;up,acv,acm,cv,cm,max=maxl=>maxr,fit)
    else
        return ()
    end
end

"""
Spike Triggered Average of Images

1. x: Matrix where each row is one image
2. y: Vector of image response

- norm: normalization factor, default no normalization.
        it could be ``sum(y)`` if y is number of spike or spike rate, then STA would be spiking probability.
- whiten: whiten factor, default no whiten.
        it could be ``(xᵀx)⁻¹`` or inverse of covariance matrix, that decorrelate STA.
"""
function sta(x::AbstractMatrix,y::AbstractVector;norm=nothing,whiten=nothing)
    r = x'*y
    isnothing(norm) || (r/=norm)
    isnothing(whiten) || (r=whiten*r)

    # r = x'*x\r
    # r=length(y)*inv(cov(x,dims=1))*r

    return r
end

function psthsts(xs::Vector,binedges::Vector,c;israte::Bool=true,normfun=nothing)
    m,se,x = psth(xs,binedges,israte=israte,normfun=normfun)
    df = DataFrame(x=x,m=m,se=se,c=fill(c,length(x)))
end
function psthsts(xs::Vector,binedges::Vector,cond::DataFrame;israte::Bool=true,normfun=nothing)
    fs = finalfactor(cond)
    vcat([psth(xs[r[:i]],binedges,condstring(r,fs),israte=israte,normfun=normfun) for r in eachrow(cond)]...)
end
function psthsts(xs::Vector,binedges::Vector,ctc::DataFrame,factor;israte::Bool=true,normfun=nothing)
    vf = intersect(names(ctc),factor)
    isempty(vf) && error("No Valid Factor Found.")
    psth(xs,binedges,condin(ctc[:,vf]),israte=israte,normfun=normfun)
end
function psth(ds::DataFrame,binedges::Vector,conds::Vector;normfun=nothing,spike=:spike,isse::Bool=true)
    is,ss = findcond(ds,conds)
    df = psth(map(x->ds[spike][x],is),binedges,ss,normfun=normfun)
    if isse
        df[:ymin] = df[:y]-df[:ysd]./sqrt(df[:n])
        df[:ymax] = df[:y]+df[:ysd]./sqrt(df[:n])
    end
    return df,ss
end
function psthstss(xss::Vector,binedges::Vector,conds;normfun=nothing)
    n = length(xss)
    n!=length(conds) && error("Length of xss and conds don't match.")
    dfs = [psth(xss[i],binedges,conds[i],normfun=normfun) for i=1:n]
    return cat(1,dfs)
end
function unitdensity(pos;w=ones(length(pos)),spacerange=extrema(pos),bw=0.01(last(spacerange)-first(spacerange)),
                    step=bw/2,r=nothing,wfun=sum,s=nothing)
    hbw = bw/2
    y = first(spacerange):step:last(spacerange)
    n = [wfun(w[i-hbw .<=pos.< i+hbw]) for i in y]
    if !isnothing(r)
        n = n/(bw*π*r^2)
    end
    i = isnan.(n) .| isinf.(n)
    if any(i)
        n = Spline1D(y[.!i],n[.!i],k=3,bc="extrapolate")(y)
    end
    if !isnothing(s)
        g = gaussianf.(-5:5,σ=s)
        g ./= sum(g)
        n = imfilter(n,centered(g),ImageFiltering.Fill(0))
    end
    return (;n,y)
end
function spacepsth(unitpsth,unitposition;w=ones(size(unitposition,1)),spacerange=extrema(unitposition),
                    bw=0.01(last(spacerange)-first(spacerange)),step=bw/2)
    hbw = bw/2
    x = unitpsth[1].x
    y = first(spacerange):step:last(spacerange)

    n = zeros(length(y))
    psth = zeros(length(y),length(x))
    for i in eachindex(y)
        @views j = y[i]-hbw .<=unitposition[:,2].< y[i]+hbw
        n[i] = sum(w[j])
        if n[i] > 0
            @views psth[i,:] = mapreduce(u->u.m,.+,unitpsth[j])/n[i]
        end
    end
    return (;psth,x,y,n)
end

"""
Shift(shuffle) corrected, normalized(coincidence/spike), condition/trial-averaged Cross-Correlogram of binary spike trains.
(Bair, W., Zohary, E., and Newsome, W.T. (2001). Correlated Firing in Macaque Visual Area MT: Time Scales and Relationship to Behavior. J. Neurosci. 21, 1676–1697.)
"""
function correlogram(bst1,bst2;lag=nothing,isnorm=true,shiftcorrection=true,condis=nothing)
    if !isnothing(condis)
        cccg=[];x=[]
        for ci in condis
            ccg,x = correlogram(bst1[:,ci],bst2[:,ci],lag=lag,isnorm=isnorm,shiftcorrection=shiftcorrection,condis=nothing)
            push!(cccg,ccg)
        end
        ccg=dropdims(mean(hcat(cccg...),dims=2),dims=2)
        return ccg,x
    end
    n,nepoch = size(bst1)
    lag = floor(Int,isnothing(lag) ? min(n-1, 10*log10(n)) : lag)
    x = -lag:lag;xn=2lag+1
    cc = Array{Float64}(undef,xn,nepoch)
    for k in 1:nepoch
        cc[:,k]=crosscov(bst1[:,k],bst2[:,k],x,demean=false)*n
    end
    ccg = dropdims(mean(cc,dims=2),dims=2)
    if isnorm
        λ1 = mean(mean(bst1,dims=1))
        λ2 = mean(mean(bst2,dims=1))
        gmsr = sqrt(λ1*λ2)
        Θ = n.-abs.(x)
        normfactor = 1 ./ Θ ./ gmsr
    end
    if shiftcorrection
        psth1 = dropdims(mean(bst1,dims=2),dims=2)
        psth2 = dropdims(mean(bst2,dims=2),dims=2)
        s = crosscov(psth1,psth2,x,demean=false)*n
        shiftccg = (nepoch*s .- ccg)/(nepoch-1)
        if isnorm
            ccg .*= normfactor
            shiftccg .*= normfactor
        end
        ccg .-= shiftccg
    elseif isnorm
        ccg .*= normfactor
    end
    ccg,x
end
function circuitestimate(unitbinspike;lag=nothing,maxprojlag=3,minepoch=5,minspike=10,esdfactor=5,isdfactor=3.5,unitid=[],condis=nothing)
    nunit=length(unitbinspike)
    n,nepoch = size(unitbinspike[1])
    lag = floor(Int,isnothing(lag) ? min(n-1, 10*log10(n)) : lag)
    x = -lag:lag;xn=2lag+1

    isenoughspike = (i,j;minepoch=5,minspike=10) -> begin
        vsi = sum(i,dims=1)[:] .>= minspike
        vsj = sum(j,dims=1)[:] .>= minspike
        count(vsi) >= minepoch && count(vsj) >= minepoch
    end

    ccgs=[];ccgis=[];projs=[];eunits=[];iunits=[];projweights=[]
    for (i,j) in combinations(1:nunit,2)
        if isnothing(condis)
            !isenoughspike(unitbinspike[i],unitbinspike[j],minepoch=minepoch,minspike=minspike) && continue
            vcondis=nothing
        else
            vci = map(ci->isenoughspike(unitbinspike[i][:,ci],unitbinspike[j][:,ci],minepoch=minepoch,minspike=minspike),condis)
            all(.!vci) && continue
            vcondis = condis[vci]
        end

        ccg,_ = correlogram(unitbinspike[i],unitbinspike[j],lag=lag,condis=vcondis)
        ps,es,is,pws = projectionfromcorrelogram(ccg,i,j,maxprojlag=maxprojlag,esdfactor=esdfactor,isdfactor=isdfactor)
        if !isempty(ps)
            push!(ccgs,ccg);push!(ccgis,(i,j))
            append!(projs,ps);append!(eunits,es);append!(iunits,is);append!(projweights,pws)
        end
    end
    unique!(eunits);unique!(iunits)
    if length(unitid)==nunit
        map!(t->(unitid[t[1]],unitid[t[2]]),ccgis,ccgis)
        map!(t->(unitid[t[1]],unitid[t[2]]),projs,projs)
        map!(t->unitid[t],eunits,eunits)
        map!(t->unitid[t],iunits,iunits)
    end
    return ccgs,x,ccgis,projs,eunits,iunits,projweights
end
function projectionfromcorrelogram(cc,i,j;maxprojlag=3,minbaselag=maxprojlag+1,esdfactor=5,isdfactor=5)
    midi = Int((length(cc)+1)/2)
    base = [cc[midi+minbaselag:end];cc[1:midi-minbaselag]]
    bm,bsd = mean_and_std(base);hl = bm + esdfactor*bsd;ll = bm - isdfactor*bsd
    forwardlags = midi+1:midi+maxprojlag
    backwardlags = midi-maxprojlag:midi-1
    fcc = cc[forwardlags]
    bcc = cc[backwardlags]
    ps=[];ei=[];ii=[];pws=[]

    if ll <= cc[midi] <= hl
        if any(fcc .> hl)
            push!(ps,(i,j));push!(ei,i);push!(pws,(maximum(fcc)-bm)/bsd)
        elseif any(fcc .< ll)
            push!(ps,(i,j));push!(ii,i);push!(pws,(minimum(fcc)-bm)/bsd)
        end
        if any(bcc .> hl)
            push!(ps,(j,i));push!(ei,j);push!(pws,(maximum(bcc)-bm)/bsd)
        elseif any(bcc .< ll)
            push!(ps,(j,i));push!(ii,j);push!(pws,(minimum(bcc)-bm)/bsd)
        end
    end
    return ps,ei,ii,pws
end

"Check Layer Boundaries"
function checklayer!(ls::Dict;ln=["1", "2", "3", "23", "4A", "4B", "4AB", "4Cα", "4Cβ", "4C", "5", "6", "56", "WM"])
    n = length(ln)
    for i in 1:n-1
        if haskey(ls,ln[i])
            for j in (i+1):n
                if haskey(ls,ln[j])
                    ls[ln[i]][1] = ls[ln[j]][2]
                    break
                end
            end
        end
    end
    return ls
end

"""
Try to locate cell layer.

1. y coordinate of cell postion
2. layer definition
"""
function assignlayer(y,layer::Dict)
    l = missing
    for k in keys(layer)
        if layer[k][1] <= y < layer[k][2]
            l=k;break
        end
    end
    return l
end

"Check circuit consistency and remove duplicates"
function checkcircuit(projs,eunits,iunits,projweights)
    ivu = intersect(eunits,iunits)
    veunits = setdiff(eunits,ivu)
    viunits = setdiff(iunits,ivu)
    ivp = map(p->p[1] in ivu,projs)
    vprojs = projs[.!ivp]
    vprojweights = projweights[.!ivp]
    ui = indexin(unique(vprojs),vprojs)
    return vprojs[ui],veunits,viunits,vprojweights[ui]
end


## Here is the code from unit 1

##
using DelimitedFiles
using SparseArrays
ei,ej = copy.(eachcol(Int.(readdlm("wiki-simple.edges"))))
xy = readdlm("wiki-simple.xy")
pages = readlines("wiki-simple.nodes")
n = size(xy,1)
A = sparse(ei,ej,1,n,n)

##  PageRank with coordinate relaxation
#
function relaxpagerank(A,α,v)
  @assert(0 ≤ α < 1, "needs probably α")
  @assert(all(vi -> vi ≥ 0, v), "needs non-negative v")
  v = v ./ sum(v) # we can normalize for them.
  d = vec(sum(A,dims=2)) # compute the degrees
  x = copy(v) # start of with v
  nsteps = 2*ceil(Int,log(eps(1.0))/log(α)) # upper bound on steps
  # we need row's of A', which are columns of A.
  # it's also easier to iterate over Di*x, instead of x
  d[d .== 0] .= 1 # switch 0's to 1's
  idx = x./d
  for i=1:nsteps
    for j=1:size(A,1)
      rhs = (1-α)*v[j]  + α*(A[:,j]'*idx)
      gamma = rhs/(1 - α*A[j,j])
      idx[j] = gamma/d[j]
    end
  end
  x =  (d.*idx)
  return x/sum(x) # renormalize to probability
end
x = relaxpagerank(A,0.85,ones(n)/n)

##
using MatrixNetworks
pr = pagerank(A,0.85)
norm(x-pr)


"""
    gamma_rng(α, β, shape=1; seed=nothing)

Generate a `shape` element array of random variables from a Gamma(`α`, `β`) distribution. Optionally you can set a specific seed.

# Notes

The Gamma distribution is given:

``f(x,α,β) = \\frac{β^α x^{α-1} e^{-βx}}{Γ(α)} \\quad x ≥ 0``

# Examples

```julia-repl
julia> gamma_rng(1,1)
1-element Vector{Float64}:
 0.5190236735858542

julia> gamma_rng(1,1,4)
4-element Vector{Float64}:
 0.3035517926878862
 0.5765419737109622
 0.44121996206333797
 0.7325887616559309

julia> gamma_rng(1,1,(2,2))
2×2 Matrix{Float64}:
 0.228818  0.88849
 0.665729  1.01668

```

# References

Law, A. Simulation modeling and analysis, 5th Ed. McGraw Hill Education, Tuscon, 2013.
"""
function gamma_rng(α::Real, β::Real=1, shape::Union{Int, Tuple{Vararg{Int}}}=1; seed::Union{Int, Nothing}=nothing)
    seed_setter(seed)
    X = β .* get_gamma_rv(α, shape)
    return X
end

function get_gamma_rv(α::Real, shape::Union{Int, Tuple{Vararg{Int}}}=1; seed::Union{Int, Nothing}=nothing)
    seed_setter(seed)
    if α < 0
        throw(ArgumentError("Not implemented for α < 0."))
    elseif α == 1
        return expon_rng(1, shape)
    elseif α < 1
        X = zeros(shape)
        X .= get_gamma_ad.(α)
        return X
    else
        X = zeros(shape)
        X .= get_gamma_gb.(α)
        return X
    end
end

function get_gamma_ad(α::Real; max_iter=300)
    b = (ℯ + α)/ℯ
    i = 0
    while i < max_iter
        U₁ = get_std_uniform()[1]
        P = b*U₁
        if P > 1
            Y = -log((b-P)/α)
            U₂ = get_std_uniform()[1]
            if U₂ ≤ Y^(α-1)
                return Y
            end
        else
            Y = P^(1/α)
            U₂ = get_std_uniform()[1]
            if U₂ ≤ ℯ^(-Y)
                return Y
            end
        end
        i +=1
    end
end

function get_gamma_gb(α::Real; max_iter=300)
    a = 1 / sqrt(2*α - 1)
    b = α - log(4)
    q = α + 1/a
    θ = 4.5
    d = 1 + log(θ)
    i = 0
    while i < max_iter
        U₁, U₂ = get_std_uniform(2)
        V = a * log(U₁/(1-U₁))
        Y = α * ℯ^V
        Z = U₁^2 * U₂
        W = b + q * V - Y
        if W + d - θ * Z ≥ 0
            return Y
        elseif W ≥ log(Z)
            return Y
        end
        i += 1
    end
end


# The following iterative sequence is defined for the set of positive integers:
#
# n → n/2 (n is even)
# n → 3n + 1 (n is odd)
#
# Using the rule above and starting with 13, we generate the following sequence:
#
# 13 → 40 → 20 → 10 → 5 → 16 → 8 → 4 → 2 → 1
#
# It can be seen that this sequence (starting at 13 and finishing at 1) contains
# 10 terms. Although it has not been proved yet (Collatz Problem), it is thought
# that all starting numbers finish at 1.
#
# Which starting number, under one million, produces the longest chain?
#
# NOTE: Once the chain starts the terms are allowed to go above one million.

using ProjectEulerSolutions

# Memoization type solution to store intermediate results without recursion.
# It does end up recalculating some solutions multiple times.
function p014solution_memoize(num::Integer=3)::Integer
    cache = Dict{Integer, Integer}(1 => 1)
    for k = 2:num
        n = k
        i = 1
        while !haskey(cache, n)
            if n % 2 == 0
                n = fld(n, 2)
            else
                n = 3 * n + 1
            end
            i += 1
        end
        @inbounds cache[k] = get(cache, n, 0) + i - 1
    end

    # Find max value
    return maxdictkey(cache)
end

# Return maximum value in dictionary
function maxdictkey(cache::Dict{Integer, Integer})::Integer
    # Find max value
    maxkey, maxvalue = 0, 0
    for (key, value) in cache
        if value > maxvalue
            maxkey, maxvalue = key, value
        end
    end
    return maxkey
end

# Dynamic programming type solution to store intermediate results
function collatz!(cache::Dict{Integer, Integer}, n::Integer)::Integer
    if haskey(cache, n)
        return cache[n]
    else
        if n % 2 == 0
            @inbounds cache[n] = collatz!(cache, fld(n, 2)) + 1
        else
            @inbounds cache[n] = collatz!(cache, 3 * n + 1) + 1
        end
        return cache[n]
    end
end

# Same type of answer, but with no caching.  Sometimes the "smartest" solution
# isn't always the fastest.
function p014solution_nocache(num::Integer=3)::Integer
    longest = 0
    max_terms = 0
    for i in 1:num
        n = i
        nterms = 1
        while n != 1
            nterms += 1
            if nterms > max_terms
                max_terms = nterms
                longest = i
            end
            if n % 2 == 0
                n = fld(n, 2)
            else
                n = 3 * n + 1
            end
        end
    end

    return longest
end

# Recursive type solution to store intermediate results.  Theoretically
# recalculates no numbers, but recursion adds an overhead.
function p014solution_recurse(num::Integer=3)::Integer
    cache = Dict{Integer, Integer}(1 => 1)
    for k = num:-1:2
        if !haskey(cache, k)
            collatz!(cache, k)
        end
    end

    # Find max value
    return maxdictkey(cache)
end

p014 = Problems.Problem(Dict("Recurse" => p014solution_recurse,
                             "Memoize" => p014solution_memoize,
                             "No Cache" => p014solution_nocache))

Problems.benchmark(p014, 1_000_000)


# sse(x) gives the L2 norm of x
sse(x) = (x'*x)[1]

function levenberg_marquardt(f::Function, g::Function, x0; tolX=1e-8, tolG=1e-12, maxIter=100, lambda=100.0, show_trace=false)
	# finds argmin sum(f(x).^2) using the Levenberg-Marquardt algorithm
	#          x
	# The function f should take an input vector of length n and return an output vector of length m
	# The function g is the Jacobian of f, and should be an m x n matrix
	# x0 is an initial guess for the solution
	# fargs is a tuple of additional arguments to pass to f
	# available options:
	#   tolX - search tolerance in x
	#   tolG - search tolerance in gradient
	#   maxIter - maximum number of iterations
	#   lambda - (inverse of) initial trust region radius
	#   show_trace - print a status summary on each iteration if true
	# returns: x, J
	#   x - least squares solution for x
	#   J - estimate of the Jacobian of f at x

	# other constants
	const MAX_LAMBDA = 1e16 # minimum trust region radius
	const MIN_LAMBDA = 1e-16 # maximum trust region radius
	const MIN_STEP_QUALITY = 1e-3
	const MIN_DIAGONAL = 1e-6 # lower bound on values of diagonal matrix used to regularize the trust region step

	converged = false
	iterCt = 0
	x = x0
	delta_x = copy(x0)
	f_calls = 0
	g_calls = 0

	fcur = f(x)
	f_calls += 1
	residual = sse(fcur)

	# Maintain a trace of the system.
	tr = OptimizationTrace()
	if show_trace
		d = {"lambda" => lambda}
		os = OptimizationState(x, sse(fcur), iterCt, d)
		push!(tr, os)
		println(os)
	end

	while ( ~converged && iterCt < maxIter )
		J = g(x)
		g_calls += 1
		# we want to solve:
		#    argmin 0.5*||J(x)*delta_x + f(x)||^2 + lambda*||diagm(J'*J)*delta_x||^2
		# Solving for the minimum gives:
		#    (J'*J + lambda*DtD) * delta_x == -J^T * f(x), where DtD = diagm(sum(J.^2,1))
		# Where we have used the equivalence: diagm(J'*J) = diagm(sum(J.^2, 1))
		# It is additionally useful to bound the elements of DtD below to help
		# prevent "parameter evaporation".
		DtD = diagm(Float64[max(x, MIN_DIAGONAL) for x in sum(J.^2,1)])
		delta_x = ( J'*J + sqrt(lambda)*DtD ) \ -J'*fcur
		# if the linear assumption is valid, our new residual should be:
		predicted_residual = sse(J*delta_x + fcur)
		# check for numerical problems in solving for delta_x by ensuring that the predicted residual is smaller
		# than the current residual
		if predicted_residual > residual
			error("Error solving for delta_x: predicted residual increase.")
		end
		# try the step and compute its quality
		trial_f = f(x + delta_x)
		f_calls += 1
		trial_residual = sse(trial_f)
		# step quality = residual change / predicted residual change
		rho = (trial_residual - residual) / (predicted_residual - residual)

		if rho > MIN_STEP_QUALITY
			x += delta_x
			fcur = trial_f
			residual = trial_residual
			# increase trust region radius
			lambda = max(0.1*lambda, MIN_LAMBDA)
		else
			# decrease trust region radius
			lambda = min(10*lambda, MAX_LAMBDA)
		end
		iterCt += 1

		# show state
		if show_trace
			d = {"g(x)" => norm(J'*fcur, Inf), "dx" => delta_x, "lambda" => lambda}
			os = OptimizationState(x, sse(fcur), iterCt, d)
			push!(tr, os)
			println(os)
		end

		# check convergence criteria:
		# 1. Small gradient: norm(J^T * fcur, Inf) < tolG
		# 2. Small step size: norm(delta_x) < tolX
		if norm(J' * fcur, Inf) < tolG
			converged = true
		elseif norm(delta_x) < tolX*(tolX + norm(x))
			converged = true
		end
	end

	# give the user info about the stopping condition
	if ~converged
		println("Exceeded maximum number of iterations")
	end

	OptimizationResults("Levenberg-Marquardt", x0, x, sse(fcur), iterCt, converged, tr, f_calls, g_calls)
end


"""
    Adam(;lr=0.001, gclip=0, beta1=0.9, beta2=0.999, eps=1e-8)
    update!(w,g,p::Adam)

Container for parameters of the Adam optimization algorithm used by
[`update!`](@ref).

Adam is one of the methods that compute the adaptive learning rate. It
stores accumulated gradients (first moment) and the sum of the squared
of gradients (second).  It scales the first and second moment as a
function of time. Here is the update formulas:

    m = beta1 * m + (1 - beta1) * g
    v = beta2 * v + (1 - beta2) * g .* g
    mhat = m ./ (1 - beta1 ^ t)
    vhat = v ./ (1 - beta2 ^ t)
    w = w - (lr / (sqrt(vhat) + eps)) * mhat

where `w` is the weight, `g` is the gradient of the objective function
w.r.t `w`, `lr` is the learning rate, `m` is an array with the same
size and type of `w` and holds the accumulated gradients. `v` is an
array with the same size and type of `w` and holds the sum of the
squares of the gradients. `eps` is a small constant to prevent a zero
denominator. `beta1` and `beta2` are the parameters to calculate bias
corrected first and second moments. `t` is the update count.

If `vecnorm(g) > gclip > 0`, `g` is scaled so that its norm is equal
to `gclip`.  If `gclip==0` no scaling takes place.

Reference: [Kingma, D. P., & Ba,
J. L. (2015)](https://arxiv.org/abs/1412.6980). Adam: a Method for
Stochastic Optimization. International Conference on Learning
Representations, 1–13.

"""
type Adam
    lr::AbstractFloat
    gclip::AbstractFloat
    beta1::AbstractFloat
    beta2::AbstractFloat
    eps::AbstractFloat
    t::Int
    fstm
    scndm
end

Adam(; lr=0.001, gclip=0, beta1=0.9, beta2=0.999, eps=1e-8)=Adam(lr, gclip, beta1, beta2, eps, 0, nothing, nothing)

function update!(w, g, p::Adam)
    gclip!(g, p.gclip)
    if p.fstm===nothing; p.fstm=zeros(w); p.scndm=zeros(w); end
    p.t += 1
    scale!(p.beta1, p.fstm)
    axpy!(1-p.beta1, g, p.fstm)
    scale!(p.beta2, p.scndm)
    axpy!(1-p.beta2, g .* g, p.scndm)
    fstm_corrected = p.fstm / (1 - p.beta1 ^ p.t)
    scndm_corrected = p.scndm / (1 - p.beta2 ^ p.t)
    axpy!(-p.lr, (fstm_corrected ./ (sqrt.(scndm_corrected) + p.eps)), w)
end


# Linear algebra functions not in the standard library.

export orthogonal, orthogonal!, pca, explained_variance, projection_distance

"""Return the angle between a and b"""
function Base.angle(a::AbstractVector, b::AbstractVector)
    length(a) == length(b) || throw(DimensionMismatch("a has dimension $(length(a)), but b has dimension $(length(b))"))
    acos(min(max(-1.0, dot(a, b) / norm(a) / norm(b)), 1.0))
end

"""return the angle between w and v, accounting for their sign"""
minangle(w, v) = min(angle(w, v), angle(w, -v))

"""
    pca(X::AbstractMatrix, k::Integer=min(size(X)...))

Compute the top `k` PCs of `X` via `LinearAlgebra.svd`. The columns of `X` correspond to samples and
the rows to features.
"""
function pca(X::AbstractMatrix, k::Integer=min(size(X)...))
    F = svd(X)
    return F.U[:, 1:k]
end

"""Orthogonalize and normalize the columns of A in-place."""
function orthogonal!(A::AbstractMatrix)
    m, n = size(A)
    for i in 1:n
        for j in 1:i-1
            l = dot(view(A, :, j), view(A, :, i))
            for k in 1:size(A, 1)
                A[k, i] -= l*A[k, j]
            end
        end
        g = norm(view(A, :, i))
        if g > sqrt(eps(eltype(A)) * size(A, 1))
            view(A, :, i) ./= g
        end
        replace!(view(A, :, i), NaN=>zero(eltype(A)))
    end
    return A
end

"""Orthogonalize and normalize the columns of A."""
orthogonal(A::AbstractMatrix) = orthogonal!(copy(A))

"""
    explained_variance(X, V, den)

Explained variance with pre-computed denominator.
"""
explained_variance(X, V, Xnorm::Real) = min((norm(X'*V) / Xnorm)^2, 1.0-eps(Float64))

"""
    explained_variance(X, V)

Return the fraction of variance explained by the principal components
in V, defined as tr(V'X'XV) / tr(X'X).

"""
explained_variance(X, V) = explained_variance(X, V, norm(X))

"""

Compute the right-singular vectors from the left-singular vectors `U` and the data matrix `X`.
"""
right_from_left_singular_vectors(X, U) = orthogonal!(U'*X)

"""

In-place version of the above.
"""
right_from_left_singular_vectors!(V, X, U) = orthogonal!(mul!(V, U', X))

projection_distance(X, V) = sqrt(norm(X .- (X*V)*V')) / reduce(*, size(X))

function projection_distance(X::SparseMatrixCSC, V)
    rv = 0.0
    L = X*V
    R = V'
    Is, Js, Vs = findnz(X)
    for (i, j, v) in zip(Is, Js, Vs)
        rv += (Float64(v) - dot(view(L, i, :), view(R, :, j)))^2
    end
    rv = sqrt(rv)
    rv /= length(V)
end

"""
    matrix_from_tensor(T)

Convert a tensor (i.e., a 3D matrix) to a matrix by flattening the first two dimensions. For
example, if the tensor corresponds to a vector of images, each image corresponds to a row of
the resulting matrix.
"""
function matrix_from_tensor(T)
    d1, d2, n = size(T)
    d = d1*d2
    Matrix(reshape(T, d, n)')
end

import SparseArrays.AbstractSparseMatrixCSC
import SparseArrays.permute, SparseArrays.permute!
import SparseArrays.getcolptr
import SparseArrays._checkargs_sourcecompatdest_permute!, SparseArrays._checkargs_sourcecompatperms_permute!, SparseArrays._ispermutationvalid_permute!

function perm!(X::AbstractSparseMatrixCSC{Tv,Ti}, A::AbstractSparseMatrixCSC{Tv,Ti}, q::AbstractVector{<:Integer}, f::Function=identity) where {Tv,TvA,Ti}
    Xi = 1
    getcolptr(X)[1] = Xi
    for Xj in 1:size(A, 2) # column of X, add @inbounds
        Aj = q[Xj] # corresponding column of A
        for Ai in nzrange(A, Aj) # row of A
            rowvals(X)[Xi] = rowvals(A)[Ai]
            nonzeros(X)[Xi] = f(nonzeros(A)[Ai])
            Xi += 1
        end
        getcolptr(X)[Xj+1] = Xi
    end
    return X
end

function permute!(X::AbstractSparseMatrixCSC{Tv,Ti}, A::AbstractSparseMatrixCSC{Tv,Ti}, q::AbstractVector{<:Integer}) where {Tv,Ti}
    _checkargs_sourcecompatdest_permute!(A, X)
    _checkargs_sourcecompatperms_permute!(A, 1:size(A, 1), q)
    if !_ispermutationvalid_permute!(q, getcolptr(X))
        throw(ArgumentError("column-permutation argument `q` must be a valid permutation"))
    end
    perm!(X, A, q)
end

function permute(A::AbstractSparseMatrixCSC{Tv,Ti}, q::AbstractVector{<:Integer}) where {Tv,Ti}
    _checkargs_sourcecompatperms_permute!(A, 1:size(A, 1), q)
    X = SparseMatrixCSC(size(A, 1), size(A, 2),
                        ones(Ti, size(A, 2) + 1),
                        Vector{Ti}(undef, nnz(A)),
                        Vector{Tv}(undef, nnz(A)))
    _checkargs_sourcecompatdest_permute!(A, X)
    if !_ispermutationvalid_permute!(q, getcolptr(X))
        throw(ArgumentError("column-permutation argument `q` must be a valid permutation"))
    end
    perm!(X, A, q)
end


### A Pluto.jl notebook ###
# v0.12.16

using Markdown
using InteractiveUtils

# ╔═╡ e077ab16-22e5-11eb-15ed-5f0654bf7328
using Pkg, DrWatson

# ╔═╡ e75a9a06-22e5-11eb-2e36-8d4f62b830ed
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Statistics, Distributions, DataFrames, CSV
	using HypothesisTests
end;

# ╔═╡ ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
md"## Listing6.10"

# ╔═╡ fa68607e-22e5-11eb-0558-c9a4d9f77426
begin
	data = CSV.read(datadir("machine1.csv"), DataFrame; delim=",")[:, 1]
end;

# ╔═╡ 1aa41898-3afa-11eb-2f49-570853ce6435
begin
	n, s, alpha = length(data), std(data), 0.1
	ci = (  (n-1)*s^2/quantile(Chisq(n-1),1-alpha/2),
			(n-1)*s^2/quantile(Chisq(n-1),alpha/2)  )

	Text("Point estimate for the variance: $(s^2)")
end

# ╔═╡ d0a56572-3b2a-11eb-383a-f94fe2f45a5b
Text("Confidence interval for the variance: $ci")

# ╔═╡ 475ff888-22e6-11eb-2354-09f8f40a8e12
md"## End of listing6.10"

# ╔═╡ Cell order:
# ╟─ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
# ╠═e077ab16-22e5-11eb-15ed-5f0654bf7328
# ╠═e75a9a06-22e5-11eb-2e36-8d4f62b830ed
# ╠═fa68607e-22e5-11eb-0558-c9a4d9f77426
# ╠═1aa41898-3afa-11eb-2f49-570853ce6435
# ╠═d0a56572-3b2a-11eb-383a-f94fe2f45a5b
# ╟─475ff888-22e6-11eb-2354-09f8f40a8e12


"""
In both digital filter design and spectral estimation, the choice of a windowing
function can play an important role in determining the quality of overall results.
The main role of the window is to damp out the effects of the Gibbs phenomenon that
results from truncation of an infinite series.

# Available Window Functions

    | Function       | Window Type                                    |
    | -------------- | ---------------------------------------------- |
    | barthann       | Modified Bartlett-Hann                         |
    | bartlett       | Bartlett                                       |
    | blackman       | Blackman                                       |
    | blackmanharris | Minimum four-term Blackman-Harris              |
    | bohman         | Bohman                                         |
    | flattop        | Flat top weighted                              |
    | hamming        | Hamming                                        |
    | hanning        | Hann (Hanning)                                 |
    | nuttall        | Nuttall-defined minimum 4-term Blackman-Harris |
    | parzen         | Parzen (de la Vallée Poussin)                  |
    | rectangular    | Rectangular                                    |
    | triangular     | Triangular                                     |
    | tukey          | Tukey (tapered cosine)                         |

"""
module WindowFunctions

export barthann
export bartlett
export blackman
export blackmanharris
export bohman
export flattop
export hamming
export hanning
export nuttall
export parzen
export rectangular
export triangular

"""
    barthann(N::Int) -> AbstractArray
    Sidelobe peak attenuation = 35.9dB
barthann window has a mainlobe at the origin and asymptotically decaying
sidelobes on both sides. It is a linear combination of weighted Bartlett
and Hanning windows with near sidelobes lower than both Bartlett and Hanning
and with far sidelobes lower than both Bartlett and Hamming windows. The
mainlobe width of the modified Bartlett-Hann window is not increased relative
to either Bartlett or Hann window mainlobes.
"""
function barthann(N::Int)
    n = (0:N-1)/(N-1)
    return @. 0.62 - 0.48*abs(n-0.5) + 0.38*cos(2*pi*n-pi)
end

"""
    bartlett(N::Int) -> AbstractArray
    Sidelobe peak attenuation = 26.5dB
Bartlett window is very similar to a triangular window as returned by the triang
function. However, the Bartlett window always has zeros at the first and last samples,
while the triangular window is nonzero at those points.
"""
function bartlett(L::Int)
    L==1 && return [1.0]
    N = L - 1
    w = zeros(L)
    for i = 1:L
        n = i - 1
        w[i] = 0<=n<=N/2 ? 2n/N : 2-2n/N
    end
    return w
end

"""
    blackman(N::Int) -> AbstractArray
    Sidelobe peak attenuation = 58.1dB
It is recommended to use blackman window to detect two
signals with similar frequency but different amplitude.
"""
function blackman(N::Int)
    n = 2 * pi .* (0:(N-1))/(N-1)
	return @. 0.42 - 0.5 * cos(n) + 0.08 * cos(2n)
end

"""
    blackmanharris(N::Int) -> AbstractArray
    Sidelobe peak attenuation = 92dB
Minimum four-term Blackman-Harris window.
"""
function blackmanharris(N::Int)
    n = 2 * pi .* (0:N-1)/(N-1)
    a₀ = 0.35875
    a₁ = 0.48829
    a₂ = 0.14128
    a₃ = 0.01168
    return @. a₀ - a₁*cos(n) + a₂*cos(2n) - a₃*cos(3n)
end

"""
    bohman(N::Int) -> AbstractArray
    Sidelobe peak attenuation = 46dB
Bohman window is the convolution of two half-duration cosine lobes.
In the time domain, it is the product of a triangular window and a
single cycle of a cosine with a term added to set the first derivative
to zero at the boundary. Bohman windows fall off as 1/ω⁴
"""
function bohman(N::Int)
    n = (0:N-1)/(N-1)
    n = abs.((2n) .- 1)
    return @. (1-n)*cos(pi*n) + 1/pi*sin(pi*n)
end

"""
    flattop(N::Int) -> AbstractArray
    Sidelobe peak attenuation = 93.6dB
flat-top window is summations of cosines with very small passband fluctuations.
"""
function flattop(N::Int)
    n = 2 * pi .* (0:N-1)/(N-1)
    a₀ = 0.215578950
    a₁ = 0.416631580
    a₂ = 0.277263158
    a₃ = 0.083578947
    a₄ = 0.006947368
    return @. a₀ - a₁*cos(n) + a₂*cos(2n) - a₃*cos(3n) + a₄*cos(4n)
end

"""
    hamming(N::Int) -> AbstractArray
    Sidelobe peak attenuation = 43.2dB
Similar to hanning but has lower Sidelobe.
"""
function hamming(N::Int)
	return @. 0.54 - 0.46 * cos( 2*pi * (0:(N-1))/(N-1) )
end

"""
    hanning(N::Int) -> AbstractArray
    Sidelobe peak attenuation = 31.5dB

If the signal is random or unknown, or has multiple frequency components, and
the test focuses on frequency points rather than energy, it is recommended to
select hanning window.
"""
function hanning(N::Int)
	return @. 0.5 - 0.5 * cos( 2*pi * (0:(N-1))/(N-1) )
end

"""
    nuttall(N::Int) -> AbstractArray
    Sidelobe peak attenuation = 93.8dB
The window is minimum in the sense that its maximum sidelobes are minimized.
The coefficients for this window differ from the Blackman-Harris window coefficients
computed with blackmanharris and produce slightly lower sidelobes.
"""
function nuttall(N::Int)
    n = 2 * pi .* (0:N-1)/(N-1)
    a₀ = 0.3635819
    a₁ = 0.4891775
    a₂ = 0.1365995
    a₃ = 0.0106411
    return @. a₀ - a₁*cos(n) + a₂*cos(2n) - a₃*cos(3n)
end

"""
    parzen(N::Int) -> AbstractArray
    Sidelobe peak attenuation = 53.1dB
Parzen windows are piecewise-cubic approximations of Gaussian windows.
Parzen window sidelobes fall off as 1/ω⁴
"""
function parzen(N::Int)
    w = zeros(N)
    for i = 1:N
        m = i - 1
        n = abs(m - (N-1)/2)
        w[i] = n<(N-1)/4 ? 1 - 6(2n/N)^2 + 6(2n/N)^3 : 2(1-2n/N)^3
    end
    return w
end

"""
    rectangular(N::Int) -> AbstractArray
    Sidelobe peak attenuation = 13.3dB
"""
function rectangular(N::Int)
    return ones(N)
end

"""
    triangular(N::Int) -> AbstractArray
    Sidelobe peak attenuation = 26.5dB
Compared with a rectangular window, the mainlobe is about twice wide, but the
sidelobes are smaller and have no negative sidelobes.
"""
function triangular(N::Int)
    N%2==0 && return triangulareven(N)
    N%2==1 && return triangularodd(N)
end

function triangularodd(N::Int)
    w = zeros(N)
    for n = 1:N
        w[n] = 1<=n<=(N+1)/2 ? 2n/(N+1) : 2-2n/(N+1)
    end
    return w
end

function triangulareven(N::Int)
    w = zeros(N)
    for n = 1:N
        w[n] = 1<=n<=N/2 ? (2n-1)/N : 2-(2n-1)/N
    end
    return w
end

end # WindowFunctions


using Random
using Distributions
"""
    vegas(f, st, en, kwargs...)

VEGAS is a Monte Carlo algorithm for
multidimensional integration based on
adaptive importance sampling. It divides
each dimension into bins and adaptively adjusts
bin widths so the points sampled from the
region where the function has highest magnitude.

Arguments:
----------
- st: Array of starting values in each dimension.
Defaults to zeros(2)
- end: Array of ending values in each dimension.
Defaults to ones(2)

Kwargs:
------
- nbins: Number of bins in each dimension.
Defaults to 100.
- ncalls: Number of function calls per iteration.
Defaults to 1000.
- maxiter: Maximum number of iterations.
Defaults to 100.
- rtol: Relative tolerance required.
Defaults to 1e-4.
- atol: Absolute tolerance required.
Defaults to 1e-2.
- debug: Prints `abs(sd/I)` every 100 iterations.
Defaults to false.
- batch: Whether `f` returns batches of function
evaluations. `f` is assumed to take one argument
`pts`, an `ncalls × `ndims` matrix. Each row
is a unique point and returns an `ncalls` length
vector of function evals. This argument defaults
to false.

Output:
------
- Estimate for the integral
- Standard deviation
- χ^2 / (numiter - 1): should be less than 1
otherwise integral estimate should not be trusted.

References:
-----------
- Lepage, G. Peter. "A new algorithm for adaptive
multidimensional integration." Journal of
Computational Physics 27.2 (1978): 192-203.
"""
function vegas(func,
               a = [0.,0.],
               b = [1.,1.];
               maxiter = 100,
               nbins = 100,
               ncalls = 1000,
               rtol = 1e-4,
               atol = 1e-4,
               debug = false,
               batch = false)

    N = nbins
    M = ncalls

    @assert length(a) == length(b)

    ndim = length(a)

    # Start out with uniform grid
    grid = zeros(N, ndim)
    for d = 1:ndim
        grid[:,d] .= fill((b[d] - a[d])/N, N)
    end
    cgrid = cumsum(grid, dims = 1)

    # Initialize all cumulative variables
    # for integral and sd estimation
    nevals = 0
    Itot = 0.
    sd = 0.
    integrals = Float64[]
    sigma_squares = Float64[]
    iter = 1

    while iter <= maxiter

        # Sample `M` points from this grid
        pts, bpts = generate_pts(grid, cgrid, M, a, b)

        # Estimate integral
        S, S², fevals = evaluate_at_samples(func,
                                            pts,
                                            bpts,
                                            M,
                                            N,
                                            grid,
                                            batch
                                           )

        σ² = (S² - S^2) / (M - 1) + eps() # When σ² = 0
        nevals += M
        push!(integrals, S)
        push!(sigma_squares, σ²)

        # Estimate sub-increments distribution
        m = calculate_m_dist(fevals,
                             bpts,
                             grid,
                             ndim
                            )

        # Update grid to reflect sub-inc dist
        update_grid!(grid, cgrid, N, M, m)

        # Update grid and generate new points
        cumsum!(cgrid, grid, dims=1)

        oldItot = Itot

        # Calculate integral and s.d upto this point
        Itot = sum((integrals.^3) ./ sigma_squares) /
                sum((integrals.^2) ./ sigma_squares)

        sd = Itot * sum((integrals.^2) ./ sigma_squares)^(-0.5)

        if debug
            iter % 100 == 0 && println("Iteration $iter, abs(sd/Itot) = $(abs(sd/Itot))")
        end

        if abs(sd/Itot) < rtol && abs(sd) < atol
            println("Converged in $nevals evaluations")
            break
        end

        iter += 1
        # M += Minc

    end
    χ² = sum(((integrals .- Itot).^2) ./ sigma_squares)
    @show nevals

    Itot, sd, χ²/(iter-1)
end

function evaluate_at_samples(f, pts, bpts, M, N, grid, batch)

    S = 0.
    S² = 0.
    fevals = zeros(M)
    dim = size(pts, 2)

    # Get all fevals in one shot
    if batch
        fevals = f(pts)
        @assert length(fevals) == M
    end

    for i = 1:M

        if !batch

            # Extract point
            p = vec(pts[i,:])

            # Eval function
            fp = f(p)
            fevals[i] = fp
        else
            fp = fevals[i]
        end

        # Get probability of that particular point
        prob = 1.
        for d = 1:dim
            prob *= (1/(N*grid[bpts[i,d],d]))
        end

        S += (fp / prob)
        S² += (fp / prob)^2

    end

    S/M, S²/M, fevals
end

"""
pts, bpts = generate_pts(grid, cumgrid, M)

Generate `M` points from `grid` which probabilities
inversely proportional to grid spacings
"""
function generate_pts(grid, cgrid, M, a, b)

    # Get bins and dimension
    dim = size(cgrid, 2)
    N = size(cgrid, 1)

    # Each dimension needs M points
    pts = zeros(M, dim)
    bpts = zeros(Int, M, dim)

    for d = 1:dim

        # Remember which bins they come from
        b = rand(1:N, M)

        bpts[:,d] .= b

        idx = 1
        for (i,bin) in enumerate(b)
            if bin == 1
                pts[i,d] = a[d] + rand(Uniform(0, cgrid[1,d]))
            else
                st = cgrid[bin-1,d]
                en = cgrid[bin,d]
                if st == en
                    continue
                end
                pts[i,d] = a[d] + rand(Uniform(cgrid[bin-1,d], cgrid[bin,d]))
            end
        end
    end

   pts, bpts
end

function calculate_m_dist(fevals, bpts, grid, dim)

    M = size(bpts, 1)
    N = size(grid, 1)
    m = zeros(N, dim)
    probs = zeros(size(grid)...)
    for d = 1:dim
        probs[:,d] .= (1 ./ grid[:,d]) ./ sum(1 ./ grid[:,d])
    end

    for d = 1:dim
        for i = 1:M
            f̄ = sqrt(fevals[i]^2 /
                                (
                                 (prod(probs[bpts[i,d],:]) /
                                 probs[bpts[i,d],d]
                                )
                                )
                    )
            m[bpts[i,d], d] += f̄ * grid[bpts[i,d],d]
        end
        m[:,d] .+= sum(m[:,d])
        m[:,d] .= m[:,d] ./ sum(m[:,d])
    end

    m
end

function update_grid!(grid, cgrid, N, M, m)

    dims = size(grid, 2)

    for d = 1:dims
        gridcol = grid[:,d]
        mdim = m[:,d]

        # Count zeros in m
        z = iszero.(mdim)
        c = sum(z)
        pos = findall(z)

        res = zeros(size(grid, 1))
        for i = 1:N
            mdim[i] == 0 && continue
            res[i] = gridcol[i] / mdim[i]
        end

        # Calculate optm as sum / # non-zero m's
        optm = sum(mdim)/(N - c)

        for i = 1:N

            # If there are no m's, grid[i] = 0
            if res[i] == 0
                gridcol[i] = 0
                continue
            end

            dist  = extract_from_bins!(mdim, optm, grid, res)
            gridcol[i] = dist

        end
        grid[:,d] .= gridcol
    end

    grid
end

function extract_from_bins!(m, optm, grid, res)

    N = size(m, 1)
    dist = 0.
    collected = 0
    for k = 1:N

        # First, calculate how many required
        required = optm - collected

        (m[k] == 0) && continue

        # If bin has what's required, take everything
        if required <= m[k]
            #res = grid[k] / morig[k]
            dist += (res[k] * required)
            m[k] -= required
            collected += required
        end

        # Update requirements
        required = optm - collected

        # If more is required, take the entire bin
        if required >= m[k]
            #res = grid[k] / morig[k]
            dist += (res[k] * m[k])
            collected += m[k]
            m[k] = 0
        end

        # If collected everything, exit
        if collected >= optm
            break
        end
    end

    dist
end


# # Power Network Flexibility Design
# Here we consider a stochastic flexibility design problem that is inspired
# from an example problem in Pulsipher JL, Zavala VM. *A scalable stochastic
# programming approach for the design of flexible systems*. Computers & Chemical
# Engineering. 2019 Sep 2;128:69-76. See the paper
# [here](https://doi.org/10.1016/j.compchemeng.2019.05.033).

# ## Modeling and Solution

using InfiniteOpt, Clp, Distributions

## Set the covariance matrix for the uncertain parameters
θ_nom = [0.; 60.; 10.]
covar = [80. 0 0; 0 80. 0; 0 0 120.]

## Set the dimensions
n_z = 3; n_θ = 3; n_d = 3

## Set the problem parameters
c = ones(n_d) / sqrt(n_d)
c_max = 5
U = 10000
num_samples = 100 # more should be used (kept small for doc generation)

## Initialize the model
m = InfiniteModel(Clp.Optimizer)
set_silent(m)

## Set the uncertainty parameters
@infinite_parameter(m, θ[i = 1:n_θ] ~ MvNormal(θ_nom, covar),
                    num_supports = num_samples)

## Initialize the variables
@variable(m, 0 <= y <= 1, Infinite(θ))
@variable(m, z[1:n_z], Infinite(θ))
@variable(m, d[1:n_d] >= 0)

## Set objective function
@objective(m, Max, expect(1 - y, θ))

## Set the line capacity constraints
@constraint(m, f1, -z[1] - 35 - d[1] <= y * U)
@constraint(m, f2, z[1] - 35 - d[1] <= y * U)
@constraint(m, f3, -z[2] - 50 - d[2] <= y * U)
@constraint(m, f4, z[1] - 50 - d[2] <= y * U)

## Set the generator capacity constraints
@constraint(m, f5, -z[3] <= y * U)
@constraint(m, f6, z[3] - 100 - d[3] <= y * U)

## Set the node balance constraints
@constraint(m, h1, z[1] - θ[1] == 0)
@constraint(m, h2, -z[1] -z[2] + z[3] - θ[2] == 0)
@constraint(m, h3, z[2] - θ[3] == 0)

## Enforce the minimum SF
@constraint(m, max_cost, sum(c[i] * d[i] for i = 1:n_d) <= c_max)

## Solve and and obtain results
optimize!(m)
if has_values(m)
    opt_y = value(y)
    opt_d = value.(d)
    opt_obj = objective_value(m)
end

## Estimate the value of SF
SF = 1 - sum(opt_y .>= 1e-8) / num_samples

## Print the results
print("------------------RESULTS------------------\n")
print("Optimal Objective:     ", opt_obj, "\n")
print("Optimal Cost:          ", sum(c[i] * opt_d[i] for i = 1:n_d), "\n")
print("Maximum Cost:          ", c_max, "\n")
print("Predicted SF:          ", 100 * SF, "%\n")
print("Optimal Design Values: ", opt_d, "\n\n")

# ### Maintenance Tests
# These are here to ensure this example stays up to date.
using Test
@test termination_status(m) == MOI.OPTIMAL
@test opt_y isa Vector{<:Real}


"""
    FBMKernel(; h::Real=0.5)

Fractional Brownian motion kernel with Hurst index `h`.

# Definition

For inputs ``x, x' \\in \\mathbb{R}^d``, the fractional Brownian motion kernel with
[Hurst index](https://en.wikipedia.org/wiki/Hurst_exponent#Generalized_exponent)
``h \\in [0,1]`` is defined as
```math
k(x, x'; h) =  \\frac{\\|x\\|_2^{2h} + \\|x'\\|_2^{2h} - \\|x - x'\\|^{2h}}{2}.
```
"""
struct FBMKernel{T<:Real} <: Kernel
    h::Vector{T}
    function FBMKernel(h::Real)
        @check_args(FBMKernel, h, zero(h) ≤ h ≤ one(h), "h ∈ [0, 1]")
        return new{typeof(h)}([h])
    end
end

FBMKernel(; h::Real=0.5) = FBMKernel(h)

@functor FBMKernel

function (κ::FBMKernel)(x::AbstractVector{<:Real}, y::AbstractVector{<:Real})
    modX = sum(abs2, x)
    modY = sum(abs2, y)
    modXY = sqeuclidean(x, y)
    h = first(κ.h)
    return (modX^h + modY^h - modXY^h) / 2
end

function (κ::FBMKernel)(x::Real, y::Real)
    return (abs2(x)^first(κ.h) + abs2(y)^first(κ.h) - abs2(x - y)^first(κ.h)) / 2
end

function Base.show(io::IO, κ::FBMKernel)
    return print(io, "Fractional Brownian Motion Kernel (h = ", first(κ.h), ")")
end

_fbm(modX, modY, modXY, h) = (modX^h + modY^h - modXY^h) / 2

_mod(x::AbstractVector{<:Real}) = abs2.(x)
_mod(x::ColVecs) = vec(sum(abs2, x.X; dims=1))
_mod(x::RowVecs) = vec(sum(abs2, x.X; dims=2))

function kernelmatrix(κ::FBMKernel, x::AbstractVector)
    modx = _mod(x)
    modxx = pairwise(SqEuclidean(), x)
    return _fbm.(modx, modx', modxx, κ.h)
end

function kernelmatrix!(K::AbstractMatrix, κ::FBMKernel, x::AbstractVector)
    modx = _mod(x)
    pairwise!(K, SqEuclidean(), x)
    K .= _fbm.(modx, modx', K, κ.h)
    return K
end

function kernelmatrix(κ::FBMKernel, x::AbstractVector, y::AbstractVector)
    modxy = pairwise(SqEuclidean(), x, y)
    return _fbm.(_mod(x), _mod(y)', modxy, κ.h)
end

function kernelmatrix!(
    K::AbstractMatrix, κ::FBMKernel, x::AbstractVector, y::AbstractVector
)
    pairwise!(K, SqEuclidean(), x, y)
    K .= _fbm.(_mod(x), _mod(y)', K, κ.h)
    return K
end

function kernelmatrix_diag(κ::FBMKernel, x::AbstractVector)
    modx = _mod(x)
    modxx = colwise(SqEuclidean(), x)
    return _fbm.(modx, modx, modxx, κ.h)
end

function kernelmatrix_diag(κ::FBMKernel, x::AbstractVector, y::AbstractVector)
    modxy = colwise(SqEuclidean(), x, y)
    return _fbm.(_mod(x), _mod(y), modxy, κ.h)
end


export
    isperfectsquare,
    factor, eulerphi,
    divisorcount, divisorsigma,
    isperfect, isdeficient, isabundant,
    primefactors, factors,
    least_number_with_d_divisors

isperfectsquare(n::Integer) = n == isqrt(n)^2

Nemo.factor(n::Int128) = (factor ∘ big)(n)
Nemo.factor(n::T) where {T<:Integer} = factor(n, Dict{T, Int})
Nemo.factor(n::Integer, out::Type{O}) where {OP<:Integer, OK<:Integer, O<:Associative{OP, OK}} = begin
    factorization = (factor ∘ fmpz)(n)
    pairs = (convert(OP, p) => convert(OK, k) for (p, k) in factorization)
    out(pairs)
end
Nemo.factor(n::Integer, out::Type{Nemo.Fac}) = (factor ∘ fmpz)(n)

##  Euler's Phi (or: totient) function
Nemo.eulerphi(n::T) where {T<:Integer} = convert(T, (eulerphi ∘ fmpz)(n))
Nemo.eulerphi(f::Associative{T, <:Integer}) where {T<:Integer} = begin
    phi = one(T)
    for (p, k) in f
        phi *= p^(k-1) * (p - 1)
    end
    phi
end

# https://oeis.org/wiki/Divisor_function
divisorcount(n::T) where {T<:Integer} = convert(T, sigma(fmpz(n), 0))
divisorcount(f::Associative{T, <:Integer}) where {T<:Integer} = begin
    c = one(T)
    for k in values(f)
        c *= k + 1
    end
    c
end

divisorsigma(n::T, s = 1) where {T<:Integer} = convert(T, sigma(fmpz(n), s))
divisorsigma(f::Associative{T, <:Integer}, s = 1) where {T<:Integer} = begin
    s >= 0 || error("Argument 's' must be an integer greater or equal 0")
    s == 0 && return divisorcount(f)

    σ = one(T)
    for (p, k) in f
        σ *= (p^((k + 1) * s) - 1) ÷ (p^s - 1)
    end
    σ
end

isperfect(n::Integer) = divisorsigma(n, 1) - n == n
isdeficient(n::Integer) = divisorsigma(n, 1) - n < n
isabundant(n::Integer) = divisorsigma(n, 1) - n > n

primefactors(n::Integer) = (sort! ∘ collect ∘ keys ∘ factor)(n)

# http://rosettacode.org/wiki/Factors_of_an_integer
factors(n::T, negative::Bool = false) where {T<:Integer} = begin
    n > 0 || error("Argument 'n' must be an integer greater 0")

    f = [one(n)]
    for (p, k) in factor(n)
        f = reduce(vcat, f, [f * p^j for j in 1:k])::Array{T,1}
    end

    if length(f) == 1
        f = [one(n), n]
    else
        sort!(f)
    end

    negative ? flatten([f -f]') : f
end

indexfactorization2number(x::Array{T,1}) where {T<:Integer} =
    prod(big(nthprime(i))^k for (i, k) in enumerate(x))

# http://www.primepuzzles.net/problems/prob_019.htm
least_number_with_d_divisors(d::Integer) =
    minimum(indexfactorization2number(e) for e in least_number_with_d_divisors_exponents(d))

least_number_with_d_divisors_exponents(d::T, i::Int = 1, prevn::T = zero(T)) where {T<:Integer} = begin
    d <= 1 && return Vector{T}[T[]]

    f = factor(d, SortedDict{T, Int})
    pmax = last(f)[1]
    k = sum(values(f))

    p = nthprime(k+i-1)
    p_i = nthprime(i)
    m = floor(Integer, log(p) / log(p_i))

    c = [pmax]
    for b in 2:m
        !(b in keys(f)) && continue

        for a in b:m
            a*b <= pmax && continue
            d % (a*b) != 0 && continue
            first(primefactors(a)) < b && continue
            i > 1 && a*b > prevn && continue

            push!(c, a*b)
        end
    end

    ans = Vector{T}[]
    for ni in c
        for tailn in least_number_with_d_divisors_exponents(d ÷ ni, i+1, ni)
            push!(tailn, ni-1)
            push!(ans, tailn)
        end
    end

    i == 1 ? [reverse(x) for x in ans] : ans
end


#############################################################################
# relative_entropy.jl
# relative entropy (ie, sum_i( x_i log (x_i/y_i) ) of expressions x and y
# All expressions and atoms are subtypes of AbstractExpr.
# Please read expressions.jl first.
#############################################################################

export relative_entropy, log_perspective
export sign, curvature, monotonicity, evaluate

# TODO: make this work for a *list* of inputs, rather than just for scalar/vector/matrix inputs

struct RelativeEntropyAtom <: AbstractExpr
    head::Symbol
    id_hash::UInt64
    children::Tuple{AbstractExpr,AbstractExpr}
    size::Tuple{Int, Int}

    function RelativeEntropyAtom(x::AbstractExpr, y::AbstractExpr)
        if sign(x) == ComplexSign() || sign(y) == ComplexSign()
            error("Both the arguments should be real but these are instead $(sign(x)) and $(sign(y))")
        else
            children = (x, y)
            return new(:entropy, hash(children), children, size(x))
        end
    end
end

function sign(x::RelativeEntropyAtom)
    return NoSign()
end

function monotonicity(x::RelativeEntropyAtom)
    return (NoMonotonicity(),NoMonotonicity())
end

function curvature(x::RelativeEntropyAtom)
    return ConvexVexity()
end

function evaluate(e::RelativeEntropyAtom)
    x = evaluate(e.children[1])
    y = evaluate(e.children[2])
    if any(isnan, y) return Inf end

    out = x.*log.(x./y)
    # fix value when x=0:
    # out will only be NaN if x=0, in which case the correct value is 0
    out[isnan.(out)] = 0
    return out
end

function conic_form!(e::RelativeEntropyAtom, unique_conic_forms::UniqueConicForms=UniqueConicForms())
    if !has_conic_form(unique_conic_forms, e)
        # transform to conic form:
        # x log x/y <= z
        # x log y/x >= -z
        # log y/x >= -z/x
        # y/x >= exp(-z/x)
        # y >= x exp(-z/x)
        # and cf the standard form for the exponential cone {(x,y,z): y*exp(x/y) <= z}
        z = Variable(e.size)
        x = e.children[1]
        y = e.children[2]
        objective = conic_form!(z, unique_conic_forms)
        for i=1:size(x,1)
            for j=1:size(x,2)
                conic_form!(ExpConstraint(-z[i,j], x[i,j], y[i,j]), unique_conic_forms)
            end
        end
        # need to constrain x>=0 and y>0.
        # x>=0 we get for free from the form of the exponential cone, so just add
        conic_form!(y>=0, unique_conic_forms) # nb we don't know how to ask for strict inequality
        cache_conic_form!(unique_conic_forms, e, objective)
    end
    return get_conic_form(unique_conic_forms, e)
end

relative_entropy(x::AbstractExpr, y::AbstractExpr) = sum(RelativeEntropyAtom(x, y))
# y*log(x/y)
log_perspective(x::AbstractExpr, y::AbstractExpr) = -relative_entropy(y, x)


struct Polytope{T<:Integer}
    points::Matrix{Rational{T}}
    rays::Matrix{T}
    bidrays::Matrix{T}
end

function find_generators(A::Matrix{T}, b::Vector{T}) where T<:Rational
    mat = [A b]
    imat = intmat(mat)
    Aint, bint = imat[:,1:end-1], imat[:,end]
    find_generators(Aint, bint)
end

"""
# Description
Find generators for the polytope
  P = { x : Ax ≥ b }
that is, calculate the following representation
  P = { ∑ᵢ λᵢUᵢ + ∑ⱼ μⱼVⱼ + ∑ₖ νₖWₖ : ∑ᵢλᵢ = 1, μ ≥ 0 }

This is done through the use of the generalised Chernikova
algorithm on the cone homogenization of the above system.

# Returns
(points, rays, bidrays):
  `points::Matrix{Rational{T}}`: matrix whose columns correspond extremal points of polytope
  `rays::Matrix{T}`: matrix whose columns correspond to unidirectional rays of polytope
  `bidrays::Matrix{T}`: matrix whose columns correspond to bidirectional rays of polytopes
"""
function find_generators(A::Matrix{T}, b::Vector{T}) where T<:Integer
    # Create homogenized system
    m, n = size(A)
    Ab = [[A -b]; zeros(T, 1, n+1)]
    Ab[m+1,n+1] = one(T)
    bid, uni = chernikova_general(Ab)

    n_bid = size(bid, 2)
    n_uni = size(uni, 2)

    # Count number of trailing zeros (rays of polytope)
    # Note that all bidirectional rays in homogenized system
    # must have a trailing zero
    # (due to positivity constraint on this coordinate)
    n_zeros_uni = 0
    for j in 1:n_uni
        if uni[n+1,j] == 0 n_zeros_uni += 1 end
    end

    # If no unidirectional rays intersect with {x : x[n+1] == 0}
    # then polytope is empty
    if n_zeros_uni == n_uni
        return Array{Rational{T}}(undef, n, 0), Array{T}(undef, n, 0), Array{T}(undef, n, 0)
    end

    # Allocate matrices and get points.
    # The vertices of the polytope correspond to the unidirectional rays with non-zero trailing coordinate
    points = Array{Rational{T}}(undef, n, (n_uni - n_zeros_uni))
    rays = Array{T}(undef, n, n_zeros_uni)
    p_count = 1
    r_count = 1
    for j in 1:n_uni
        if uni[n+1, j] != 0
            points[:, p_count] = uni[1:n,j]//uni[n+1,j]
            p_count+=1
        else
            rays[:,r_count] = uni[1:n,j]
            r_count+=1
        end
    end

    bidrays = bid[1:n,:]
    return points, rays, bidrays
end

"""
# Description
Calculates the minimum and maximum values of
the projection of the polytope onto
specified vectors:

  min_{x ∈ P} Tᵢx
  max_{x ∈ P} Tᵢx

where Tᵢ are the columns of a matrix T.
# Arguments
* `P::Polytope{I}`: Polytope to project
* `T::Matrix{I}`: Matrix onto whose columns we project polytope
# Returns
`(mins::Array{Rational{T}}, maxs::Array{Rational{T}})`
"""
function min_max_projections(P::Polytope{S}, T::Matrix{Rational{S}} = Array{Rational{S}}(I, length(P), length(P))) where S<:Integer
    n = size(T, 2)
    n_points, n_rays, n_bidrays = size(P.points, 2), size(P.rays, 2), size(P.bidrays, 2)
    mins, maxs = fill(typemax(Rational{S}), n), fill(typemin(Rational{S}), n)
    for i in 1:n
        for j in 1:n_points
            tp = dot(T[:,i], P.points[:,j])
            if tp < mins[i]
                mins[i] = tp
            end
            if tp > maxs[i]
                maxs[i] = tp
            end
        end
        for j in 1:n_rays
            tr = dot(T[:,i], P.rays[:,j])
            if tr < 0
                mins[i] = -Inf
                break
            elseif tr > 0
                maxs[i] = Inf
                break
            end
        end
        for j in 1:n_bidrays
            if dot(T[:,i], P.bidrays[:,j]) != 0
                maxs[i] = Inf
                mins[i] = -Inf
                break
            end
        end
    end
    return mins, maxs
end

length(P::Polytope{T} where T<:Integer) = size(P.points, 1)


using EconPDEs

Base.@kwdef  struct CampbellCochraneModel
    # consumption process parameters
    μ::Float64 = 0.0189
    σ::Float64 = 0.015
    # utility
    γ::Float64 = 2.0
    ρ::Float64 = 0.116
    # habit
    κs::Float64 = 0.138
    b::Float64 = 0.0
end
# I choose persistence so that monthly simulation of the model matches processes in CC (1999)
# ρ = 12 * (1 - 0.89^(1/12))
# κs = 12 * (1 - 0.87^(1/12))

function initialize_stategrid(m::CampbellCochraneModel; sn = 1000)
    μ = m.μ ; σ = m.σ ; γ = m.γ ; ρ = m.ρ ; κs = m.κs ; b = m.b
    Sbar = σ * sqrt(γ / (κs - b / γ))
    sbar = log.(Sbar)
    smax =  sbar + 0.5 * (1 - Sbar^2)
    # corresponds to Grid 3 in Wachter (2005)
    shigh = log.(range(0.0, exp(smax), length = div(sn, 10)))
    slow = range(-300.0, shigh[2], length = sn - div(sn, 10))
    OrderedDict(:s => vcat(slow[1:(end-1)], shigh[2:end]))
end

function (m::CampbellCochraneModel)(state::NamedTuple, y::NamedTuple)
    (; μ, σ, γ, ρ, κs, b) = m
    (; s) = state
    (; p, ps_up, ps_down, pss) = y

    # drift and volatility of  s and p
    Sbar = σ * sqrt(γ / (κs - b / γ))
    sbar = log(Sbar)
    λ = 1 / Sbar * sqrt(1 - 2 * (s - sbar)) - 1
    μs = - κs * (s - sbar)
    ps = (μs >= 0) ? ps_up : ps_down
    σs = λ * σ
    σp = ps / p * σs
    μp = ps / p * μs + 0.5 * pss / p * σs^2

    # market price of risk κ
    κ = γ * (σ + σs)

    # risk free rate  r
    r = ρ + γ * μ - (γ * κs - b) / 2 + b * (sbar - s)

    # PDE
    pt = - p * (1 / p + μ + μp + σp * σ - r - κ * (σ + σp))
    return (; pt)
end

# Campbell Cochrane (1999)
m = CampbellCochraneModel()
stategrid = initialize_stategrid(m)
yend = OrderedDict(:p => ones(length(stategrid[:s])))
result = pdesolve(m, stategrid, yend)
@assert result.residual_norm <= 1e-5

# Wachter (2005) calibration
# m = CampbellCochraneModel(μ = 0.022, σ = 0.0086, γ = 2.0, ρ = 0.073, κs = 0.116, b = 0.011)
# stategrid = initialize_stategrid(m)
# yend = OrderedDict(:p => ones(length(stategrid[:s])))
# y, result, distance = pdesolve(m, stategrid, yend)


module StatisticTools

export weightedmean

"""
weightedmean(values::Array{Float64,1}, errors::Array{Float64,1})

computes the weighted mean of values, using as weights w = 1/ sigma^2
"""
function weightedmean(x::Array{Float64,1}, errors::Array{Float64,1})
    @assert length(x) == length(errors)
    weights = 1./ (errors.^2)
    sigma = sqrt( 1./ sum( weights ) )
    mean = sigma^2 * sum(x .* weights)

    mean, sigma
end

end # module


module Epidemiology

using AlgebraicPetri
using Petri
using Catlab
using Catlab.Theories

export InfectiousDiseases, FunctorGenerators, F_epi, S, E, I, R, D, transmission, exposure, illness, recovery, death

ob = PetriCospanOb(1)
spontaneous_petri = PetriCospan([1], Petri.Model(1:2, [(Dict(1=>1), Dict(2=>1))]), [2])
transmission_petri = PetriCospan([1], Petri.Model(1:2, [(Dict(1=>1, 2=>1), Dict(2=>2))]), [2])
exposure_petri = PetriCospan([1, 2], Petri.Model(1:3, [(Dict(1=>1, 2=>1), Dict(3=>1, 2=>1))]), [3, 2])

""" InfectiousDiseases
"""
@present InfectiousDiseases(FreeBiproductCategory) begin
    S::Ob
    E::Ob
    I::Ob
    R::Ob
    D::Ob
    transmission::Hom(S⊗I, I)
    exposure::Hom(S⊗I, E⊗I)
    illness::Hom(E,I)
    recovery::Hom(I,R)
    death::Hom(I,D)
end

""" generators
"""
S,E,I,R,D,transmission,exposure,illness,recovery,death = generators(InfectiousDiseases);

""" FunctorGenerators
"""
const FunctorGenerators = Dict(S=>ob, E=>ob, I=>ob, R=>ob, D=>ob,
        transmission=>transmission_petri, exposure=>exposure_petri,
        illness=>spontaneous_petri, recovery=>spontaneous_petri, death=>spontaneous_petri)

""" F_epi
"""
F_epi(ex) = functor((PetriCospanOb, PetriCospan), ex, generators=FunctorGenerators)

end


function tvlqr(A, B, Q, R)
    T = length(Q)

    P = [zero(A[1]) for t = 1:T]
    K = [zero(B[1]') for t = 1:T-1]
    P[T] = Q[T]

    for t = T-1:-1:1
        K[t] = (R[t] + B[t]' * P[t+1] *  B[t]) \ (B[t]' * P[t+1] * A[t])
        P[t] = (Q[t] + K[t]' * R[t] * K[t]
                + (A[t] - B[t] * K[t])' * P[t+1] * (A[t] - B[t] * K[t]))
    end

    return K, P
end

function reference_gains(s::Simulation, traj::ContactTraj, obj::TrackingVelocityObjective;
		N::Int=10, κ=2e-4, U_scaling=100, V_scaling=100)

	model = s.model
	H = traj.H
	z = [[deepcopy(traj.z) for i = 1:N]...;]
	θ = [[deepcopy(traj.θ) for i = 1:N]...;]

	rz = [LinearizedStep(s, z[t], θ[t], κ).rz for t = 1:N*H]
	rθ = [LinearizedStep(s, z[t], θ[t], κ).rθ for t = 1:N*H]

	idx_z = indices_z(s)
	idx_θ = indices_θ(model, nf=1)

	∂z∂θ = [-rz[t] \ rθ[t] for t = 1:N*H]

	∂q2∂q1 = [zeros(model.nq, model.nq) for t = 1:N*H]
	∂q2∂q2 = [I(model.nq) for t = 1:N*H]
	∂q3∂q1 = [∂z∂θ[t][idx_z.q, idx_θ.q1] for t = 1:N*H]
	∂q3∂q2 = [∂z∂θ[t][idx_z.q, idx_θ.q2] for t = 1:N*H]

	∂q2∂u = [zeros(model.nq, model.nu) for t = 1:N*H]
	∂q3∂u = [∂z∂θ[t][idx_z.q, idx_θ.u] for t = 1:N*H]

	A = [[∂q2∂q1[t] ∂q2∂q2[t]; ∂q3∂q1[t] ∂q3∂q2[t]] for t = 1:N*H]
	B = [[∂q2∂u[t]; ∂q3∂u[t]] for t = 1:N*H]
	Q = [
		[obj.q[1]+V_scaling*obj.v[1]    -V_scaling*obj.v[1];
		 -V_scaling*obj.v[1]    obj.q[1]+V_scaling*obj.v[1]
		 ] for i = 1:N*H+1]
	R = [U_scaling * obj.u[1] for i = 1:N*H]

	K, P = tvlqr(A, B, Q, R)
	return K[1:H]
end


module Geometry
export volume, surfacearea, lateralarea

abstract type Solid end

# *******************************************
mutable struct Cone <: Solid
    height::Float32
    # one element implies start from one point
    # two elements for truncated cone volume
    radii::Array{Float32, 1}
    function Cone(depth, diameters)
        if length(diameters) > 2
            error("Too many diameter arguments in list")
        end
        if length(diameters) == 1
            push!(diameters,0.0)
        end
        new(depth, diameters / 2)
    end
end
mutable struct Cylinder <: Solid
    height::Float32
    radius::Float32
    function Cylinder(depth, diameter)
        new(depth, (diameter / 2.0))
    end
end
struct Cube <: Solid
    side::Float32
end
struct RectangularPrism <: Solid
    height::Float32
    width::Float32
    depth::Float32
end

function volume(c::Solid)
    tc = typeof(c)
    if tc == Cone
        rads = (c.radii[1]^2 + c.radii[1]*c.radii[2] + c.radii[2]^2)
        return (1/3)*(pi*rads*c.height)
    elseif tc == Cylinder
        return (1/1)*(pi*(c.radius^2)*c.height)
    elseif tc == Cube
        return (1/1)*(c.side^3)
    elseif tc == RectangularPrism
        return (1/1) * (c.height * c.width * c.depth)
    end
    return nothing
end

function lateralarea(c::Solid)
    tc = typeof(c)
    if tc == Cone
        return pi * (c.radii[1] + c.radii[2]) * (((c.radii[1] - c.radii[2])^2 + c.height^2)^(1/2))
    elseif tc == Cylinder
        return 2*pi*c.radius*c.height
    elseif tc == Cube
        return (1/1)*4*(c.side^2)
    elseif tc == RectangularPrism
        return (1/1) * (2 * (c.height * c.width) + 2 * (c.height * c.depth))
    end
    return nothing
end

function surfacearea(c::Solid)
    tc = typeof(c)
    if tc == Cone
        return lateralarea(c) + pi * (c.radii[1]^2 + c.radii[2]^2)
    elseif tc == Cylinder
        return lateralarea(c) + 2*pi*(c.radius^2)
    elseif tc == Cube
        return (1/1)*6*(c.side^2)
    elseif tc == RectangularPrism
        return lateralarea(c) + (1/1) * (2 * (c.width * c.depth))
    end
    return nothing

end

end


using FFTW
using Plots

include("naiveDFT.jl")
include("simpleDFT.jl")
include("benchmarkFunction.jl")

function main()
    #
    # Create test set
    #
    N = 2^8 # 2^8=256
    tmax = 30.0
    Δ = tmax / N

    t = collect(range(0.0, tmax, length = N))
    ω = collect(range(0, stop = N - 1))
    hk = sin.(t) .+ 0.5 .* sin.(5.0 .* t .+ 1.0) .+ sin.(4.0 .* t .- 1.5)

    #
    # Calculations
    #
    resNaive = naiveDFT(hk)                 #naive approach
    resNaiveIter = naiveDFTIterative(hk)    #naive approach iterative
    resSimple = simpleDFT(hk)               #simple recursive dft
    resFFTW = fft(hk)                       #fftw
    resRealFFTW = rfft(hk)                  #real fftw

    println(sum(abs.(resFFTW .- resNaive))) #differences to naive<->fftw
    println(sum(abs.(resFFTW .- resNaiveIter))) #differences to naive<->fftw
    println(sum(abs.(resFFTW .- resSimple))) #differences to simple<->fftw

    #
    # Benchmarking
    #
    benchmarkDFTFunction(fft, 1.0, 5, "FFTW.csv")
    benchmarkDFTFunction(simpleDFT, 1.0, 5, "simpleDFT.csv")
    benchmarkDFTFunction(naiveDFT, 1.0, 5, "naiveDFT.csv")
    benchmarkDFTFunction(naiveDFTIterative, 1.0, 5, "naiveDFTIter.csv")

    #
    # plotting
    #
    gr()
    p1re = plot(ω, real.(resFFTW), title = "FFTW Real", label = "");
    p1im = plot(ω, imag.(resFFTW), title = "FFTW Imag", label = "");
    p1abs = plot(ω, abs.(resFFTW) .^ 2 , title = "FFTW Abs^2", label = "");

    p2re = plot(ω, real.(resNaive), title = "Naive Real", label = "");
    p2im = plot(ω, imag.(resNaive), title = "Naive Imag", label = "");
    p2abs = plot(ω, abs.(resNaive) .^ 2 , title = "Naive Abs^2", label = "");

    p3re = plot(ω, real.(resSimple), title = "Simple Real", label = "");
    p3im = plot(ω, imag.(resSimple), title = "Simple Imag", label = "");
    p3abs = plot(ω, abs.(resSimple) .^ 2 , title = "Simple Abs^2", label = "");

    p = plot(p1re, p1im, p1abs, p2re, p2im, p2abs, p3re, p3im, p3abs, layout = (3, 3), size = (1024, 800))
    savefig(p, "fftPlot.png")

end

main()


"""
```
transition(m::PoolModel{T}) where {T<:AbstractFloat}
```

Assign transition equation

```
x_t = (1-ρ) μ + ρ x_{t-1} + sqrt{1 - ρ^2} σ ϵ_t
ϵ_t ∼ iid N(0,1), x_0 ∼ N(μ,σ^2)
λ_t = Φ(x_t)
```
where Φ(⋅) is the cdf of a N(0,1) random variable, F_ϵ is the distribution of ϵ_t,
and F_λ is the distribution of λ(x_0).
"""
function transition(m::PoolModel{T}) where {T<:AbstractFloat}
    @inline Φ(x::Vector{Float64}, ϵ::Vector{Float64}) = abs.([0;1] .-
                                                             (cdf.(Normal(), (1 - m[:ρ].value) .*
                                                                   m[:μ].value .+ m[:ρ].value .*
                                                                   quantile(Normal(),x[1]) .+
                                                                   sqrt(1 - m[:ρ].value^2) .*
                                                                   m[:σ].value .* ϵ)))
    F_ϵ = Normal(0.,1.)
    F_λ = Uniform(0.,1.)
    return Φ, F_ϵ, F_λ
end


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

using Markdown
using InteractiveUtils

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

# ╔═╡ 0507f8ed-8b64-48af-a6c4-ff7c4211b9e9
begin
	using Pkg
	Pkg.activate(".")
end

# ╔═╡ 3310d2f7-450a-4cd1-9c3a-46d02d23a7c6
using Revise

# ╔═╡ d27b2d72-d264-11eb-0be5-13dcacfd2adc
using DeconvOptim, TestImages, ImageShow, Plots, LinearAlgebra, IndexFunArrays, Noise, FourierTools, SpecialFunctions, FFTW, LaTeXStrings, PlutoUI, Images, Tullio

# ╔═╡ 952b251c-207b-4412-b6e1-268fce1647d9
begin
	img = Float32.(testimage("fabio_gray"));
	img_1D = img[:, 200]
end;

# ╔═╡ fa8cd9c9-a2fd-495c-8d22-ada7bb9c39f6
otf(x, Δx=1) = begin
	x = abs(x)
	if x <= Δx
		SpecialFunctions.jinc(x * Δx *(1-x/Δx)) .* 2 / π * (acos(x/Δx) - x/Δx * sqrt(1-(x/Δx)^2))
	else
		zero(x)
	end
end

# ╔═╡ 9a07bc88-be76-4531-bc91-df0d20c3221c
begin
	x = range(-1.5, 1.5, length=size(img, 1))
	freqs = fftshift(fftfreq(size(img_1D, 1), 1))
	psf = Float32.(DeconvOptim.generate_psf(size(img), 20))
	psf_1D = psf[1, :]
	psf_1D ./= sum(psf_1D)
	otf_1D = abs.(ffts(psf_1D))
end;

# ╔═╡ 24fc86da-0b7c-493c-8977-2a19ef6dc133
img_n = Float32.(poisson(DeconvOptim.conv(img, psf), 1000));

# ╔═╡ 3bda922d-552b-42ec-9055-33a141b5841a
blur(x, otf=otf_1D) = iffts(ffts(x) .* otf)

# ╔═╡ 308e2562-5a20-4082-8a6d-9fb135738c49
md"
Mathematically:

$(S * \text{PSF})(\mathbf r) = \int_{-\infty}^{\infty} S(\mathbf r - \mathbf x) \cdot \text{PSF}(\mathbf x) \, \mathrm d \mathbf x$
"

# ╔═╡ b542187a-3ae7-4430-a81d-f968d9000427
img_blurry = DeconvOptim.conv(img, psf);

# ╔═╡ 88d83629-f1b2-4d69-928f-d4acfbc76b70
reg_1D = TV(num_dims=1);

# ╔═╡ d6d0f436-48fd-4e8a-863c-3e9d3850cc91
begin
	reg_tik = Tikhonov()
	reg_TV = TV()
	reg_GR = DeconvOptim.GR()
end

# ╔═╡ 061faf49-662c-4508-83b4-ddcf0970ed0d
md"### DeconvOptim.jl: Microscopy Image Deconvolution
"

# ╔═╡ 23829201-9756-4ef8-90c9-3917b761fe4b
load("../docs/src/assets/logo.png")

# ╔═╡ 30f21bb8-6d09-4fce-9d2a-568bfaf3ff7a
md"
* **Felix Wechsler:** Master Student at the Leibniz Institute of Photonic Technology in  Jena, Germany
* https://github.com/roflmaostc/DeconvOptim.jl
* `]add DeconvOptim`
"

# ╔═╡ 7a0a44c9-fb07-44e6-9a8b-8f720b84e6f6
md"""### Image Convolution

* Typical description of isotropic blur of an image
* The blurring kernel describes blur
    * In optics/microscopy a finite sized dot called Point Spread Function (**PSF**)
    * often a Gaussian function used in image processing
    * cigarre shaped object for motion blur

Discrete version:

$(S * \text{PSF})[i] = \sum_{m} S[i-m] \cdot \text{PSF}[m]$
"""

# ╔═╡ 49686d9a-1683-428f-83ba-a9131c2ad432
[Gray.(img) Gray.(DeconvOptim.conv(img, psf))]

# ╔═╡ 491ec7cb-1663-4acf-b81f-9acafba2b63d
md"## Convolution Theorem

$(S * \text{PSF})(\mathbf r) = \mathcal{F}^{-1}\bigg[ \mathcal{F}[S] \cdot \mathcal{F}[\text{PSF}]  \bigg]$

* we can express the convolution with a Fast Fourier Transform (FFT) which only takes $\mathcal O(N \log(N))$ operations

* For large kernels (especially in 3D), sliding kernels are slower

*  $\mathcal{F}[\text{PSF}]$ is called the $\text{OTF}$

"

# ╔═╡ b6f9e42e-5a23-40ad-9e73-8f02da48f69c
md"### Optical System act as low pass filter

*  $\text{OTF}$ shows the frequency throughput

"

# ╔═╡ e31f4704-5bce-4c8d-b3a7-873a3460d9f8
plot(x, otf.(x), xlabel="frequency / maximum frequency", ylabel="contrast")

# ╔═╡ dc8cf4e2-a87c-47ea-8247-3ae772852241
md"## Frequency spectrum of blurred sample $Y(\mathbf r)$

Blurred sample:
$Y(\mathbf r) = (S * \text{PSF})(\mathbf r)$
"

# ╔═╡ f8033d13-faee-41f5-bdd6-ae8721e8b8a8
begin
	plot(freqs, abs.(ffts(img_blurry)[:, 128]), yaxis=:log, ylabel="real part of FFT output in AU", xlabel="frequency in 1/px", ylims=(1e-4, 1e2), label="blurred")
	plot!(freqs, abs.(ffts(img)[:, 128]), ylabel="abs of FFT output in AU",
				xlabel="frequency in 1/px", yaxis=:log,
				ylims=(1e-4, 1e4), label="ground truth")
	#plot!(freqs, abs.(ffts(DeconvOptim.conv(img_1D, psf_1D))))
end

# ╔═╡ fe7e0292-31f6-43b5-83a5-a38698a87563
md"## Deconvolution Pipeline

* based on:
    * Zygote.jl
    * Optim.jl
    * Tullio.jl
    * CUDA.jl
"

# ╔═╡ e9ef5ba4-56c0-4595-bc28-e04882f44a9a
load("../docs/src/assets/tex/pipeline.png")

# ╔═╡ 90e5708c-05a4-46e4-b1e4-9a61c96dae32
TV_by_hand(x) = @tullio r = sqrt(1f-8 + abs2(x[i, j] - x[i+1, j]) +
						 	     abs2(x[i, j] - x[i, j+1]))

# ╔═╡ 03139ac5-3525-4fad-abf1-84421492b763
DeconvOptim.generate_TV(4, [1,2, 3], [1,1, 1], 1, 0)[1]

# ╔═╡ 4e84e739-9c59-4939-8b04-aec7dc069d67
md"
### Deconvolve with DeconvOptim.jl

"

# ╔═╡ 9d9a5da2-14df-46e6-b7e6-5a33aade1754
@bind reg_list2 Select(["1" => ("Tikhonov"), "2" => ("Total Variation TV"), "3" => ("Good's Roughness GR")])

# ╔═╡ 5ae1a6a3-4505-4123-9f1e-8a1d4ac0b4e1
reg = [reg_tik, reg_TV, reg_GR][parse(Int, reg_list2)]

# ╔═╡ 15d54e5e-64f4-4a1d-8cc8-9334b2e3784f
md"
iterations =
$(@bind iter Slider(0:50, show_value=true))

λ = $(@bind λ Slider(0:0.001:0.3, show_value=true))

regularizer = $(@bind reg_bool CheckBox())"

# ╔═╡ 803368e6-53fd-4413-b3f5-ffe46ee8983e
img_deconv, res_img = deconvolution(img_blurry, psf, regularizer=reg_bool ? reg : nothing, iterations=iter, λ=λ);

# ╔═╡ e58e1f63-2c81-48fb-866a-4bb70bd428a6
Gray.(img_deconv)

# ╔═╡ 438a6639-bd35-464f-a81d-d98eb65e006e
res_1D, o = deconvolution(real(blur(img_1D)), psf_1D, iterations=iter, regularizer=reg_1D, λ=0.01);

# ╔═╡ 7a16df73-95ad-47f5-907c-6fd23c6000cf
[Gray.(img) Gray.(img_blurry) Gray.(img_deconv)]

# ╔═╡ 5c8030b7-8805-4771-9329-23abb2744544
begin
	plot(freqs, abs.(ffts(img_blurry)[:, 128]), yaxis=:log, ylabel="abs of FFT output in AU", xlabel="frequency in 1/px", ylims=(1e-4, 1e2), label="blurred")
	plot!(freqs, abs.(ffts(img)[:, 128]),  yaxis=:log,
				ylims=(1e-4, 1e4), label="ground truth")
	plot!(freqs, abs.(ffts(img_deconv)[:, 128]), label="deconvolved image")
end

# ╔═╡ ff9af79c-f06c-42ac-9c8d-6f09d2ff4056
[Gray.(img_1D); Gray.(res_1D)];

# ╔═╡ d8aee845-e922-41da-a17e-37ffa3e692f0
md"### Real Microscopy Data"

# ╔═╡ fc93fa3d-8599-463b-b9e6-8043b90e9d63
load("figures/real_data_large.png")

# ╔═╡ 97c45bfd-d1ef-49ad-908d-7360c03b0170
md"Image taken from [DeconvolutionLab2](http://bigwww.epfl.ch/deconvolution/deconvolutionlab2/)."

# ╔═╡ bf37664a-f726-46b6-9592-da419165af91
md"## Conclusion - DeconvOptim.jl
"

# ╔═╡ f8117250-bac8-43a1-aae4-5c8bab3a522d
[Gray.(ones(130, 012)) load("../docs/src/assets/logo.png")]

# ╔═╡ b5a70276-e9b9-46e8-8c67-0cad2cfa19da
md"* Flexible Image Deconvolution Software
* N-dimensional signal deconvolution
* Works both on CPU and GPUs
    * GPUs usually 5-15x speed improvement
"

# ╔═╡ Cell order:
# ╠═3310d2f7-450a-4cd1-9c3a-46d02d23a7c6
# ╠═0507f8ed-8b64-48af-a6c4-ff7c4211b9e9
# ╠═d27b2d72-d264-11eb-0be5-13dcacfd2adc
# ╠═952b251c-207b-4412-b6e1-268fce1647d9
# ╠═24fc86da-0b7c-493c-8977-2a19ef6dc133
# ╠═3bda922d-552b-42ec-9055-33a141b5841a
# ╟─fa8cd9c9-a2fd-495c-8d22-ada7bb9c39f6
# ╠═9a07bc88-be76-4531-bc91-df0d20c3221c
# ╟─308e2562-5a20-4082-8a6d-9fb135738c49
# ╠═b542187a-3ae7-4430-a81d-f968d9000427
# ╠═88d83629-f1b2-4d69-928f-d4acfbc76b70
# ╠═e58e1f63-2c81-48fb-866a-4bb70bd428a6
# ╠═d6d0f436-48fd-4e8a-863c-3e9d3850cc91
# ╠═803368e6-53fd-4413-b3f5-ffe46ee8983e
# ╠═5ae1a6a3-4505-4123-9f1e-8a1d4ac0b4e1
# ╠═438a6639-bd35-464f-a81d-d98eb65e006e
# ╟─061faf49-662c-4508-83b4-ddcf0970ed0d
# ╟─23829201-9756-4ef8-90c9-3917b761fe4b
# ╟─30f21bb8-6d09-4fce-9d2a-568bfaf3ff7a
# ╟─7a0a44c9-fb07-44e6-9a8b-8f720b84e6f6
# ╟─49686d9a-1683-428f-83ba-a9131c2ad432
# ╟─491ec7cb-1663-4acf-b81f-9acafba2b63d
# ╟─b6f9e42e-5a23-40ad-9e73-8f02da48f69c
# ╟─e31f4704-5bce-4c8d-b3a7-873a3460d9f8
# ╟─dc8cf4e2-a87c-47ea-8247-3ae772852241
# ╟─f8033d13-faee-41f5-bdd6-ae8721e8b8a8
# ╟─fe7e0292-31f6-43b5-83a5-a38698a87563
# ╟─e9ef5ba4-56c0-4595-bc28-e04882f44a9a
# ╠═90e5708c-05a4-46e4-b1e4-9a61c96dae32
# ╠═03139ac5-3525-4fad-abf1-84421492b763
# ╟─4e84e739-9c59-4939-8b04-aec7dc069d67
# ╟─9d9a5da2-14df-46e6-b7e6-5a33aade1754
# ╟─15d54e5e-64f4-4a1d-8cc8-9334b2e3784f
# ╟─7a16df73-95ad-47f5-907c-6fd23c6000cf
# ╟─5c8030b7-8805-4771-9329-23abb2744544
# ╟─ff9af79c-f06c-42ac-9c8d-6f09d2ff4056
# ╟─d8aee845-e922-41da-a17e-37ffa3e692f0
# ╟─fc93fa3d-8599-463b-b9e6-8043b90e9d63
# ╟─97c45bfd-d1ef-49ad-908d-7360c03b0170
# ╟─bf37664a-f726-46b6-9592-da419165af91
# ╟─f8117250-bac8-43a1-aae4-5c8bab3a522d
# ╟─b5a70276-e9b9-46e8-8c67-0cad2cfa19da


function simsys(sys,nt,seed)
	###
	# simulated systems
	###

	# random seed for the simulations
	srand(seed)

	if sys==1
		#
		# system 1

		# ref: "Subspace Identification for Linear Systems"
		# P. Van Overschee and B. De Moor, page 81
		#
		A = [0.6 0.6 0.;-0.6 0.6 0.;0. 0. 0.4]
		K = [0.1706 -0.1507 0.2772]'
		C = [0.7831 0.5351 0.9701]
		R = 6.3663
		Q = K * R * K'

		et = sqrt(R)*randn(1,nt)
		w = K*et

		ns = 3
 		x = zeros(ns,nt)
		x[:,1] = w[:,1]
		for t in 1:nt-1
			x[:,t+1] = A * x[:,t] + w[:,t]
		end
		y = C * x + et

		salida = y, A, K, C, Q, R
		#
	elseif sys==2
		###
		# ref: "Subspace Identification for Linear Systems"
		# P. Van Overschee and B. De Moor, page 115
		###
		A = [0.6 0.6 0.;-0.6 0.6 0.;0. 0. 0.7]
		B = [1.6161 -0.3481 2.6319]'
		K = [-1.1472 -1.5204 -3.1993]'
		C = [-0.4373 -0.5046 -0.0936]
		D = -0.7759
		R = 0.0432
		Q = K * R * K'

		et = sqrt(R)*randn(1,nt)
		w = K*et

		# input = AR(1)
		u = zeros(1,nt)
		a = randn(nt)
		for t in 1:nt-1
			u[t+1] = 0.3*u[t] + a[t]
		end

		# states
		ns = size(A,1)
		x = zeros(ns,nt)
		x[:,1] = w[:,1]
		for t in 1:nt-1
			x[:,t+1:t+1] = A * x[:,t:t] + B * u[t] + w[:,t:t]
		end
		# output
		y = C * x + D * u + et

		salida = y, u, A, B, K, C, D, Q, R
	end

	return salida

end