text:
module ReducedTest

using Test
using Random, LinearAlgebra, Statistics, Distributions, SparseArrays
using Convex, SCS
using ProximalBase
using DiffPrecTest
using DiffPrecTest: TwoSampleUstatistics, svec, skron

function generateData(p)
  n = 100
  # generate model
  Ωx = Matrix{Float64}(I, p, p)
  for l=1:p-1
    Ωx[l  , l+1] = 0.6
    Ωx[l+1, l  ] = 0.6
  end
  for l=1:p-2
    Ωx[l  , l+2] = 0.3
    Ωx[l+2, l  ] = 0.3
  end
  Δ = zeros(p, p)
  # generate Delta
  for j=1:p
    Δ[j, j] = rand(Uniform(0.1, 0.2))
  end
  for j=1:p-1
    v = rand(Uniform(0.2, 0.5))
    Δ[j  , j+1] = v
    Δ[j+1, j  ] = v
  end
  Ωy = Ωx - Δ
  d = Vector{Float64}(undef, p)
  rand!(Uniform(0.5, 2.5), d)
  d .= sqrt.(d)
  D = Diagonal(d)
  Σx = inv(Symmetric(D * Ωx * D))
  Σy = inv(Symmetric(D * Ωy * D))

  dist_X = MvNormal(convert(Matrix, Σx))
  dist_Y = MvNormal(convert(Matrix, Σy))

  X = rand(dist_X, n)'
  Y = rand(dist_Y, n)'

  X, Y
end

function solve_pen_quad(A, b, λ, ω)
  p = size(A, 1)
  x = Variable(p)
  expr = 0.
  for i=1:p
    expr += abs(x[i]) * ω[i]
  end
  problem = minimize(quadform(x, A / 2.) + dot(x, b) + expr * λ)
  solve!(problem, () -> SCS.Optimizer(verbose = false))
  x.value
end

@testset "_computeVarStep1!" begin
  function createVarianceGrad1(Δ::SparseIterate, rowIndex::Int)
    @inline function h(x, y)::Float64
      # skron(vx_j, vy_k)[ri, :] * Δ
      out = 0.
      for inz = 1:nnz(Δ)
        indColumn = Δ.nzval2ind[inz]
        out += skron(x, y, rowIndex, indColumn) * Δ.nzval[inz]
      end
      out + svec(x, x, rowIndex) - svec(y, y, rowIndex)
    end
    TwoSampleUstatistics{1, 1}(h)
  end

  function varStep1Alt!(out, X, Y, θ::SparseIterate)
    for i=1:length(out)
      s = createVarianceGrad1(θ, i)
      out[i] = variance(s, X, Y)
    end
    out
  end

  p = 10
  m = div(p*(p+1), 2)
  X, Y = generateData(p)
  n = size(X, 1)
  Sx = Symmetric( X'X / n )
  Sy = Symmetric( Y'Y / n )
  out1 = zeros(m)
  out2 = zeros(m)

  θ = SparseIterate(m)
  DiffPrecTest._computeVarStep1!(out1, Sx, Sy, X, Y, θ);
  varStep1Alt!(out2, X, Y, θ);
  @test out1 ≈ sqrt.(out2)

  θ .= sprandn(m, 0.5)
  DiffPrecTest._computeVarStep1!(out1, Sx, Sy, X, Y, θ);
  varStep1Alt!(out2, X, Y, θ);
  @test out1 ≈ sqrt.(out2)
end

@testset "reducedDiffEstimation" begin

  function reducedDiffEstimationAlt(H::Symmetric, b::Vector, Sx, Sy, X, Y, λ)
    x = SparseIterate(size(H, 1))

    ω = Array{eltype(H)}(undef, length(x))
    DiffPrecTest._computeVarStep1!(ω, Sx, Sy, X, Y, x)

    xtmp = solve_pen_quad(H, b, λ, ω)
    for i=1:length(x)
      x[i] = xtmp[i]
    end

    DiffPrecTest._computeVarStep1!(ω, Sx, Sy, X, Y, x)
    solve_pen_quad(H, b, λ, ω)
  end

  p = 5
  m = div(p*(p+1), 2)
  X, Y = generateData(p)
  n = size(X, 1)
  Sx = Symmetric( X'X / n )
  Sy = Symmetric( Y'Y / n )
  H = Symmetric( skron(Sx, Sy) )
  b = svec(Sx) - svec(Sy)

  λ = 1.01 * quantile( Normal(), 1. - 0.1 / (p*(p+1)) )

  out1 = reducedDiffEstimationAlt(H, b, Sx, Sy, X, Y, λ)
  out2 = DiffPrecTest.reducedDiffEstimation(H, b, Sx, Sy, X, Y, λ)

  @test out1 ≈ Vector(out2) atol=1e-4
end

@testset "_computeVarStep2!" begin
  function createVarianceGrad2(Δ::SparseIterate, rowIndex::Int, indElem::Int)
    @inline function h(x, y)::Float64
      # skron(vx_j, vy_k)[ri, :] * Δ
      out = 0.
      for inz = 1:nnz(Δ)
        indColumn = Δ.nzval2ind[inz]
        out += skron(x, y, rowIndex, indColumn) * Δ.nzval[inz]
      end
      out + (rowIndex == indElem ? 1. : 0.)
    end
    TwoSampleUstatistics{1, 1}(h)
  end

  function varStep2Alt!(out, indElem, X, Y, θ::SparseIterate)
    for i=1:length(out)
      s = createVarianceGrad2(θ, i, indElem)
      out[i] = variance(s, X, Y)
    end
    out
  end

  p = 10
  m = div(p*(p+1), 2)
  X, Y = generateData(p)
  n = size(X, 1)
  Sx = Symmetric( X'X / n )
  Sy = Symmetric( Y'Y / n )
  out1 = zeros(m)
  out2 = zeros(m)

  for j=1:m
    θ = SparseIterate(m)
    DiffPrecTest._computeVarStep2!(out1, j, Sx, Sy, X, Y, θ);
    varStep2Alt!(out2, j, X, Y, θ);
    @test out1 ≈ sqrt.(out2)

    θ .= sprandn(m, 0.5)
    DiffPrecTest._computeVarStep2!(out1, j, Sx, Sy, X, Y, θ);
    varStep2Alt!(out2, j, X, Y, θ);
    @test out1 ≈ sqrt.(out2)
  end
end

@testset "invHessianReduced" begin
  function invHessianReducedAlt(H::Symmetric, indRow, Sx, Sy, X, Y, λ)
    x = SparseIterate(size(H, 1))
    x[indRow] = 1.

    b = zeros(size(H, 1))
    b[indRow] = 1.

    ω = Array{eltype(H)}(undef, length(x))
    DiffPrecTest._computeVarStep2!(ω, indRow, Sx, Sy, X, Y, x)
    ω[indRow] = 0.

    xtmp = solve_pen_quad(H, b, λ, ω)
    for i=1:length(x)
      x[i] = xtmp[i]
    end

    DiffPrecTest._computeVarStep2!(ω, indRow, Sx, Sy, X, Y, x)
    ω[indRow] = 0.
    solve_pen_quad(H, b, λ, ω)
  end

  p = 5
  m = div(p*(p+1), 2)
  X, Y = generateData(p)
  n = size(X, 1)
  Sx = Symmetric( X'X / n )
  Sy = Symmetric( Y'Y / n )
  H = Symmetric( skron(Sx, Sy) )
  b = svec(Sx) - svec(Sy)

  λ = 1.01 * quantile( Normal(), 1. - 0.1 / (p*(p+1)) )

  for j=1:m
    out1 = invHessianReducedAlt(H, j, Sx, Sy, X, Y, λ)
    out2 = DiffPrecTest.invHessianReduced(H, j, Sx, Sy, X, Y, λ)

    @test out1 ≈ Vector(out2) atol=1e-4
  end
end

@testset "varianceOracle" begin

  function createVarianceReducedEstim(ω, I, Δ, J)
    function h(x, y)::Float64
      dot( (skron(x, y, I, J) * Δ + svec(x, x, I) - svec(y, y, I)), ω )
    end
    TwoSampleUstatistics{1, 1}(h)
  end

  for rep=1:10
    p = 5
    m = div(p*(p+1), 2)
    X, Y = generateData(p)
    n = size(X, 1)
    Sx = Symmetric( X'X / n )
    Sy = Symmetric( Y'Y / n )

    indS = 3:9
    C = skron(Sx, Sy, indS)
    b = svec(Sy, indS) - svec(Sx, indS)

    tmp = zeros(length(indS))
    tmp[1] = 1.
    ω = C \ tmp
    Δ = C \ b
    v = variance(ReducedOracleNormal(), Sx, Sy, X, Y, ω, collect(indS), Δ, collect(indS))

    v2 = variance(createVarianceReducedEstim(ω, indS, Δ, indS), X, Y)
    @test v ≈ v2
  end
end

end


text:
using Econometrics, Test, Random
function main()
    # dstats
    @testset "dstats" begin
        @test dstats()
    end
    # vech
    @testset "vech" begin
        @test vech()
    end
    # OLS
    x = [ones(10,1) (1:10)]
    y = [10.,9.,8.,7.,6.,5.,4.,3.,2.,1.]
    b = ols(y,x,silent=true)[1]
    @testset "ols" begin
        @test b[1,1] / 11.0 ≈ 1.0 atol=1e-5
    end
    # ML
    @testset "mle" begin
        b = mle()
        @test b[1,1] / 0.49466773000112035 ≈ 1.0 atol = 1e-5
    end
    # GMM
    @testset "gmm" begin
        b = gmm()
        @test b[1,1] / 0.4946717833359474 ≈ 1.0  atol = 1e-5
    end
    # samin
    xopt, fopt = samin(1)
    @testset "samin" begin
    @test xopt[1] ≈ 0.0 atol=1e-5
    @test fopt ≈ 2.0 atol=1e-5
    end
    #=
    # this is ill-advised, need to test with non-random data,
    # so it doesn't break when julia changes
    # npreg
    @testset "npreg" begin
        yhat = npreg(1);
        @test yhat[1,1] ≈ 0.02080984791302587
    end
    =#
end
main()

ext: jl

max_stars_repo_name: Hiroakiyusheng/Econometrics

max_issues_repo_name: Hiroakiyusheng/Econometrics

max_forks_repo_name: Hiroakiyusheng/Econometrics


text:
struct FourVector{T, A <: StaticVector{4, T}} <: StaticVector{4, T}
    fv::A # px,py,pz,E
    FourVector(fv::A) where {T, A <: StaticVector{4, T}} = new{T, A}(fv)
end
FourVector{T, A}(x::Tuple) where {T, A} = FourVector(A(x))
FourVector(x, y, z; t) = FourVector(MVector(x, y, z, t))
Particle(; E, p) = FourVector(p...; t=E)
SParticle(; E, p) = FourVector(SVector(p..., E))

function Particle(px,py,pz; E=nothing, msq=nothing)
    if (msq === nothing) == (E === nothing)
        throw(ArgumentError("Must provide exactly one of either `msq` or `E`."))
    end
    E !== nothing && return FourVector(px, py, pz; t=E)
    return FourVector(px, py, pz; t=sqrt(hypot(px, py, pz)^2 + msq))
end
function SParticle(px,py,pz; E=nothing, msq=nothing)
    if (msq === nothing) == (E === nothing)
        throw(ArgumentError("Must provide exactly one of either `msq` or `E`."))
    end
    E !== nothing && return FourVector(SVector(px, py, pz, E))
    return FourVector(SVector(px, py, pz, sqrt(hypot(px, py, pz)^2 + msq)))
end

function StaticArrays.similar_type(
    ::Type{FourVector{S, A}}, ::Type{T}, ::Size{(4,)},
) where {S, A, T}
    return FourVector{T, similar_type(A, T)}
end

Base.@propagate_inbounds function Base.getindex(p::FourVector, i::Int)
    return getindex(p.fv, i == 0 ? 4 : i)
end
Base.@propagate_inbounds function Base.setindex!(p::FourVector, v, i::Int)
    return setindex!(p.fv, v, i == 0 ? 4 : i)
end

_view(x::MVector, i) = view(x, i)
_view(x::StaticVector, i) = x[i]

function Base.getproperty(p::FourVector, sym::Symbol)
    (sym in fieldnames(FourVector)) && return getfield(p, sym)
    (sym == :Px || sym == :px || sym == :x || sym == :X) && return getfield(p, :fv)[1]
    (sym == :Py || sym == :py || sym == :y || sym == :Y) && return getfield(p, :fv)[2]
    (sym == :Pz || sym == :pz || sym == :z || sym == :Z) && return getfield(p, :fv)[3]
    (sym == :E  || sym == :p0 || sym == :t || sym == :T) && return getfield(p, :fv)[4]
    (sym == :p  || sym == :P) && return _view(getfield(p, :fv), SOneTo(3))
    throw(ArgumentError("type FourVector has no property $(sym)"))
end

LinearAlgebra.dot(p1::FourVector, p2::FourVector) = p1.E * p2.E - dot(p1.P, p2.P)

# properties
psq(p::FourVector) = sum(abs2, p.P)
invmasssq(p::FourVector) = p ⋅ p
mass(p::FourVector) = sqrt(invmasssq(p))

sphericalangles(p::FourVector) = p.P[3] / norm(p.P), atan(p.P[2], p.P[1])
boostfactor(p::FourVector) = p.E / mass(p)


text:
# -*- coding: utf-8 -*-
function generalized_phase_sensitivity_func(
    F, Imin, Imax, dI, dims::Int, Nθ::Int, initX=nothing, dt=1e-4, alg=Tsit5(),
    origin_val_idx::Int=1, origin_thr=nothing, itp="cubic", extrap=false,
    print_progress=true)
    """
    Compute angular frequency ω(I)
    and generalized phase sensitivity function ζ(θ, I), ξ(θ, I) with the adjoint method.

    Args:
    - F(X, I) : Vector field function of the oscillator model; (list of functions)
    - Imin : Minimum external force strength; (Scalar)
    - Imax : Maximum external force strength; (Scalar)
    - dI  : step of I
    - dims   : Dimension of states; (Int)
    - Nθ  : The number of splitting phase for one period; (Int)
    - initX : Array of initial ∈ R^(D x NI)
    - itp : interpolation mode. "linear" or "cubic"

    Returns: Interpolated object
    - ω(I)    : Angular frequency ∈ R^NI
    - ζ(θ, I) : Phase sensitivity function ∈ R^(Nθ x NI)
    - ξ(θ, I) : Phase sensitivity function ∈ R^(Nθ x NI)
    - Xs(θ, I): stable periodic solution ∈ R^(D x Nθ x NI)
    """
    Irange = Imin:dI:Imax
    NI = size(Irange)[1]

    flag_arraydt = dt isa Array
    @assert (flag_arraydt || (dt isa Number)) "'dt' must be a scalar or an array."
    if flag_arraydt
        @assert size(dt)[1] == NI "The size of the array of 'dt' must be the same as the number of conditions for the external force 'I'."
        dtarray = copy(dt)
    end

    if initX == nothing
        initX = rand(dims)
    end

    flag_initX = size(initX) == (NI, dims)
    if flag_initX
        initXarray = copy(initX)
    else
        @assert size(initX) == (dims, ) "The size of 'initX' must be the same as (dims, )."
    end

    θrange = range(0, 2π, length=Nθ)

    ZθI = zeros(dims, Nθ, NI)
    ωI = zeros(NI)
    XsI = zeros(dims, Nθ, NI)
    @showprogress "[1/3] Computing Xs(θ, I) and Z(θ, I)..." for i in 1:NI
        if flag_arraydt
            dt = dtarray[i]
        end

        if flag_initX
            initX = initXarray[i, :]
        end

        T, ω, Xs = find_stable_periodic_solution(F, Irange[i], dims, Nθ, initX, dt, alg, origin_val_idx, origin_thr, print_progress=false)
        Zθ = phase_sensitivity_func(F, Irange[i], dims, Nθ, T, Xs)

        ωI[i] = ω
        ZθI[:, :, i] = Zθ
        XsI[:, :, i] = Xs
    end

    ζθI = zeros(Nθ, NI)
    ξθI = zeros(Nθ, NI)
    JFI(X, I) = ForwardDiff.derivative(I -> F(X, I), I)

    # Compute ζ(θ, I)
    if print_progress println("[2/3] Computing ζ(θ, I)...") end

    for i in 1:NI
        for th in 1:Nθ
            ζθI[th, i] = JFI(XsI[:, th, i], Irange[i])' * ZθI[:, th, i]
        end
    end

    # numerical integration
    dθ = 2π / Nθ
    ζI = sum(ζθI, dims=1)[1, :] / Nθ # dθ/2π = 1/Nθ

    # scaled interpolated variables
    if itp == "cubic"
        sitpXsI = [scale(interpolate(XsI[i, :, :], BSpline(Cubic(Line(OnGrid())))), θrange, Irange) for i in 1:dims]
    else
        sitpXsI = [scale(interpolate(XsI[i, :, :], BSpline(Linear())), θrange, Irange) for i in 1:dims]
    end

    # Compute ξ(θ, I)
    if print_progress println("[3/3] Computing ξ(θ, I)...") end
    for i in 1:NI
        ξθI[1, i] = - [Interpolations.gradient(sitpXsI[j], 0, Irange[i])[2] for j in 1:dims]' * ZθI[:, 1, i]
    end

    for i in 1:NI
        for th in 2:Nθ
            ξθI[th, i] = ξθI[th-1, i] - (ζθI[th, i] - ζI[i]) * dθ / ωI[i]
        end
    end

    # scaled interpolated variables
    if itp == "cubic"
        sitpωI = scale(interpolate(ωI, BSpline(Cubic(Line(OnGrid())))), Irange)
        sitpζθI = scale(interpolate(ζθI, BSpline(Cubic(Line(OnGrid())))), θrange, Irange)
        sitpξθI = scale(interpolate(ξθI, BSpline(Cubic(Line(OnGrid())))), θrange, Irange)
    else
        sitpωI = scale(interpolate(ωI, BSpline(Linear())), Irange)
        sitpζθI = scale(interpolate(ζθI, BSpline(Linear())), θrange, Irange)
        sitpξθI = scale(interpolate(ξθI, BSpline(Linear())), θrange, Irange)
    end

    if extrap
        etpωI = extrapolate(sitpωI, Flat())
        etpζθI = extrapolate(sitpζθI, Flat())
        etpξθI = extrapolate(sitpξθI, Flat())
        etpXsI = [extrapolate(sitpXsI[i], Flat()) for i in 1:dims]
        return etpωI, etpζθI, etpξθI, etpXsI
    else
        return sitpωI, sitpζθI, sitpξθI, sitpXsI
    end
end


text:
"""
Forecast types include Perfect and Imperfect forecasts
"""
abstract type Forecast end

"""
Perfect forecasts.
"""
struct Perfect <: Forecast
    scenario_data::Vector{Scenario}
end

"""
Imperfect forecasts with Kalman Filter based updates.
"""
struct Imperfect <: Forecast
    kalman_filter::Union{Nothing, KalmanFilter}
    scenario_data::Vector{Scenario}
end

get_scenario_data(forecast::F) where F<: Forecast = forecast.scenario_data
get_kalman_filter(forecast::Imperfect) = forecast.kalman_filter


text:
using VarianceComponentsHDFE
using Test
using DataDeps

register(DataDep("VarianceComponentsHDFE",
       "Downloading medium_nocontrols_pruned.csv, medium_controls_pruned.csv, large_nocontrols_pruned.csv, and large_controls_pruned.csv, huge_pruned_main.csv, large_controls_main.csv",
       ["https://vchdfe.s3-us-west-2.amazonaws.com/medium_pruned_main.tar.gz",
       "https://vchdfe.s3-us-west-2.amazonaws.com/medium_controls_pruned_main.tar.gz",
       "https://vchdfe.s3-us-west-2.amazonaws.com/large_pruned_main.tar.gz",
       "https://vchdfe.s3-us-west-2.amazonaws.com/large_controls_pruned_main.tar.gz",
       "https://vchdfe.s3-us-west-2.amazonaws.com/huge_pruned_main.tar.gz",
       "https://vchdfe.s3-us-west-2.amazonaws.com/large_controls_main.tar.gz"],
       ["e01474ebd8d7df6d645b0c865079a108c9fd33e052b7ab163bb2f5749bc58387",
       "a896b114eb952e60212196ecd1dffe2ea8bded263590a5b93e3b30b069feb0e1",
       "583ec55959b16ef93db210736c93daa8c8eab9097d4492347fe442237f6b3c2c",
       "e6ca435c7ac3bf7a8e425ab58581f36877845c099ad4da086d1990a3c270ac88",
       "5743ed199dd6e6180f6fd721d3232bec0573a05d33f358c9a228a930e517ec12",
       "82a95b368b2f6be381f10228026ebf2ffffc7e0d80bf5580ee905adf29b26128"];
       post_fetch_method= [unpack, unpack, unpack, unpack, unpack, unpack]
   ))

run_full_test = get(ENV, "VCHDFE_FULL_TEST", "1") == "1" ? true : false
include("prepare_test_data.jl")

@testset "VarianceComponentsHDFE.jl" begin
    # solver.jl has not been updated
    # include("solver.jl")
    include("test_matrices.jl")
    run_full_test && include("fulltest.jl")
end


text:
# UnitCommitment.jl: Optimization Package for Security-Constrained Unit Commitment
# Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
# Released under the modified BSD license. See COPYING.md for more details.

using UnitCommitment
using JuMP
using Gurobi
using JSON
using Logging
using Printf
using LinearAlgebra

function main()
    NUM_THREADS = 4
    time_limit = 60 * 20
    BLAS.set_num_threads(NUM_THREADS)

    if length(ARGS) >= 2
      mode = string("_", ARGS[2])
    else
      mode = "_default"
    end
    if length(ARGS) >= 3 && !isempty(strip(ARGS[3]))
      results_dir = ARGS[3]
    else
      results_dir = string(pwd(),"/","results$mode")
    end

    # Validate mode and set formulation
    if mode == "_default"
      formulation = UnitCommitment.DefaultFormulation
    elseif mode == "_tight"
      formulation = UnitCommitment.TightFormulation
    elseif mode == "_sparse"
      formulation = UnitCommitment.SparseDefaultFormulation
    else
      error("Unknown formulation requested: ", ARGS[2])
    end

    # Filename is instance_name.sample_number.sol.gz
    # Parse out the instance + sample parts to create output files
    basename, suffix = split(ARGS[1], ".") # will not work if suffix part is not present
    model_filename_stub = string(results_dir,"/$basename.$suffix")
    solution_filename = string("$model_filename_stub.sol.json")

    # Choose logging options
    logname, logfile = nothing, nothing
    #logname = string("$model_filename_stub.out")
    if isa(logname, String) && !isempty(logname)
      logfile = open(logname, "w")
      global_logger(TimeLogger(initial_time = time(), file = logfile))
    else
      global_logger(TimeLogger(initial_time = time()))
    end

    # Check that results_dir exists
    if !isdir(model_filename_stub)
      @info("Creating $model_filename_stub")
      mkpath(model_filename_stub)
    end

    total_time = @elapsed begin
        @info "Reading: $basename"
        time_read = @elapsed begin
            instance = UnitCommitment.read_benchmark(basename)
        end
        @info @sprintf("Read problem in %.2f seconds", time_read)

        time_model = @elapsed begin
            optimizer=optimizer_with_attributes(Gurobi.Optimizer,
                                                "Threads" => NUM_THREADS,
                                                "Seed" => rand(1:1000))
            model = build_model(instance=instance, optimizer=optimizer, formulation=formulation)
        end
    end

    @info "Setting names..."
    UnitCommitment.set_variable_names!(model)

    model_filename = string(model_filename_stub,".init",".mps.gz")
    @info string("Exporting initial model without transmission constraints to ", model_filename)
    JuMP.write_to_file(model.mip, model_filename)

    total_time += @elapsed begin
        @info "Optimizing..."
        BLAS.set_num_threads(1)
        UnitCommitment.optimize!(model, time_limit=time_limit, gap_limit=1e-3)
    end

    @info @sprintf("Total time was %.2f seconds", total_time)

    @info "Writing: $solution_filename"
    solution = UnitCommitment.get_solution(model)
    open(solution_filename, "w") do file
        JSON.print(file, solution, 2)
    end

    @info "Verifying solution..."
    UnitCommitment.validate(instance, solution)

    model_filename = string(model_filename_stub,".final",".mps.gz")
    @info string("Exporting final model to ", model_filename)
    JuMP.write_to_file(model.mip, model_filename)

    if !isnothing(logfile)
      close(logfile)
    end
end # main

main()


text:
# -*- mode: julia; -*-
VERSION >= v"0.4.0-dev+6521" && Base.__precompile__(true)
module WAV
export wavread, wavwrite, wavappend, wavplay
export WAVArray, WAVFormatExtension, WAVFormat
export isextensible, isformat, bits_per_sample
export WAVE_FORMAT_PCM, WAVE_FORMAT_IEEE_FLOAT, WAVE_FORMAT_ALAW, WAVE_FORMAT_MULAW
using FileIO
using Compat

function __init__()
    module_dir = dirname(@__FILE__)
    if Libdl.find_library(["libpulse-simple"]) != ""
        include(joinpath(module_dir, "wavplay-pulse.jl"))
    elseif Libdl.find_library(["AudioToolbox"],
                              ["/System/Library/Frameworks/AudioToolbox.framework/Versions/A"]) != ""
        include(joinpath(module_dir, "wavplay-audioqueue.jl"))
    else
        include(joinpath(module_dir, "wavplay-unsupported.jl"))
    end
end

include("AudioDisplay.jl")
wavplay(fname) = wavplay(wavread(fname)[1:2]...)

# The WAV specification states that numbers are written to disk in little endian form.
write_le(stream::IO, value) = write(stream, htol(value))
read_le(stream::IO, x::Type) = ltoh(read(stream, x))

# used by WAVE_FORMAT_EXTENSIBLE
immutable WAVFormatExtension
    nbits::UInt16 # overrides nbits in WAVFormat type
    channel_mask::UInt32
    sub_format::Array{UInt8, 1} # 16 byte GUID
    WAVFormatExtension() = new(0, 0, Array(UInt8, 0))
    WAVFormatExtension(nb, cm, sb) = new(nb, cm, sb)
end

# Required WAV Chunk; The format chunk describes how the waveform data is stored
immutable WAVFormat
    compression_code::UInt16
    nchannels::UInt16
    sample_rate::UInt32
    bytes_per_second::UInt32 # average bytes per second
    block_align::UInt16
    nbits::UInt16
    ext::WAVFormatExtension
end

const WAVE_FORMAT_PCM        = 0x0001 # PCM
const WAVE_FORMAT_IEEE_FLOAT = 0x0003 # IEEE float
const WAVE_FORMAT_ALAW       = 0x0006 # A-Law
const WAVE_FORMAT_MULAW      = 0x0007 # Mu-Law
const WAVE_FORMAT_EXTENSIBLE = 0xfffe # Extension!

isextensible(fmt::WAVFormat) = (fmt.compression_code == WAVE_FORMAT_EXTENSIBLE)
bits_per_sample(fmt::WAVFormat) = isextensible(fmt) ? fmt.ext.nbits : fmt.nbits

# DEFINE_GUIDSTRUCT("00000001-0000-0010-8000-00aa00389b71", KSDATAFORMAT_SUBTYPE_PCM);
const KSDATAFORMAT_SUBTYPE_PCM = [
0x01, 0x00, 0x00, 0x00,
0x00, 0x00,
0x10, 0x00,
0x80, 0x00,
0x00, 0xaa, 0x00, 0x38, 0x9b, 0x71
                                  ]
# DEFINE_GUIDSTRUCT("00000003-0000-0010-8000-00aa00389b71", KSDATAFORMAT_SUBTYPE_IEEE_FLOAT);
const KSDATAFORMAT_SUBTYPE_IEEE_FLOAT = [
0x03, 0x00, 0x00, 0x00,
0x00, 0x00,
0x10, 0x00,
0x80, 0x00,
0x00, 0xaa, 0x00, 0x38, 0x9b, 0x71
                                         ]
# DEFINE_GUID(KSDATAFORMAT_SUBTYPE_MULAW, 0x00000007, 0x0000, 0x0010, 0x80, 0x00, 0x00, 0xaa, 0x00, 0x38, 0x9b, 0x71);
const KSDATAFORMAT_SUBTYPE_MULAW = [
0x07, 0x00, 0x00, 0x00,
0x00, 0x00,
0x10, 0x00,
0x80, 0x00,
0x00, 0xaa, 0x00, 0x38, 0x9b, 0x71
                                    ]
#DEFINE_GUID(KSDATAFORMAT_SUBTYPE_ALAW, 0x00000006, 0x0000, 0x0010, 0x80, 0x00, 0x00, 0xaa, 0x00, 0x38, 0x9b, 0x71);
const KSDATAFORMAT_SUBTYPE_ALAW = [
0x06, 0x00, 0x00, 0x00,
0x00, 0x00,
0x10, 0x00,
0x80, 0x00,
0x00, 0xaa, 0x00, 0x38, 0x9b, 0x71
                                   ]

function isformat(fmt::WAVFormat, code)
    if code != WAVE_FORMAT_EXTENSIBLE && isextensible(fmt)
        subtype = Array(UInt8, 0)
        if code == WAVE_FORMAT_PCM
            subtype = KSDATAFORMAT_SUBTYPE_PCM
        elseif code == WAVE_FORMAT_IEEE_FLOAT
            subtype = KSDATAFORMAT_SUBTYPE_IEEE_FLOAT
        elseif code == WAVE_FORMAT_ALAW
            subtype = KSDATAFORMAT_SUBTYPE_ALAW
        elseif code == WAVE_FORMAT_MULAW
            subtype = KSDATAFORMAT_SUBTYPE_MULAW
        else
            return false
        end
        return subtype == fmt.ext.sub_format
    end
    return fmt.compression_code == code
end

function WAVFormatExtension(bytes)
    if isempty(bytes)
        return WAVFormatExtension()
    end
    # split bytes into valid_bits_per_sample, channel_mask, and sub_format
    valid_bits_per_sample = (convert(UInt16, bytes[2]) << 8) | convert(UInt16, bytes[1])
    channel_mask = (convert(UInt32, bytes[6]) << 24) | (convert(UInt32, bytes[5]) << 16) | (convert(UInt32, bytes[4]) << 8) | convert(UInt32, bytes[3])
    sub_format = bytes[7:end]
    return WAVFormatExtension(valid_bits_per_sample, channel_mask, sub_format)
end

function read_header(io::IO)
    # check if the given file has a valid RIFF header
    riff = read(io, UInt8, 4)
    if riff !=  b"RIFF"
        error("Invalid WAV file: The RIFF header is invalid")
    end

    chunk_size = read_le(io, UInt32)

    # check if this is a WAV file
    format = read(io, UInt8, 4)
    if format != b"WAVE"
        error("Invalid WAV file: the format is not WAVE")
    end
    return chunk_size
end

function write_header(io::IO, data_length::UInt32)
    write(io, b"RIFF") # RIFF header
    write_le(io, data_length) # chunk_size
    write(io, b"WAVE")
end
write_standard_header(io, data_length) = write_header(io, @compat UInt32(data_length + 36))
write_extended_header(io, data_length) = write_header(io, @compat UInt32(data_length + 60))

function read_format(io::IO, chunk_size::UInt32)
    # can I read in all of the fields at once?
    orig_chunk_size = convert(Int, chunk_size)
    if chunk_size < 16
        error("The WAVE Format chunk must be at least 16 bytes")
    end
    const compression_code = read_le(io, UInt16)
    const nchannels = read_le(io, UInt16)
    const sample_rate = read_le(io, UInt32)
    const bytes_per_second = read_le(io, UInt32)
    const block_align = read_le(io, UInt16)
    const nbits = read_le(io, UInt16)
    ext = Array(UInt8, 0)
    chunk_size -= 16
    if chunk_size > 0
        const extra_bytes_length = read_le(io, UInt16)
        if extra_bytes_length == 22
            ext = read(io, UInt8, extra_bytes_length)
        end
    end
    return WAVFormat(compression_code,
                     nchannels,
                     sample_rate,
                     bytes_per_second,
                     block_align,
                     nbits,
                     WAVFormatExtension(ext))
end

function write_format(io::IO, fmt::WAVFormat)
    len = 16 # 16 is size of base format chunk
    if isextensible(fmt)
        len += 24 # 24 is the added length needed to encode the extension
    end
    # write the fmt subchunk header
    write(io, b"fmt ")
    write_le(io, convert(UInt32, len)) # subchunk length

    write_le(io, fmt.compression_code) # audio format (UInt16)
    write_le(io, fmt.nchannels) # number of channels (UInt16)
    write_le(io, fmt.sample_rate) # sample rate (UInt32)
    write_le(io, fmt.bytes_per_second) # byte rate (UInt32)
    write_le(io, fmt.block_align) # byte align (UInt16)
    write_le(io, fmt.nbits) # number of bits per sample (UInt16)

    if isextensible(fmt)
        write_le(io, convert(UInt16, 22))
        write_le(io, fmt.ext.nbits)
        write_le(io, fmt.ext.channel_mask)
        @assert length(fmt.ext.sub_format) == 16
        write(io, fmt.ext.sub_format)
    end
end

function pcm_container_type(nbits::Unsigned)
    if nbits > 32
        return Int64
    elseif nbits > 16
        return Int32
    elseif nbits > 8
        return Int16
    end
    return  UInt8
end

ieee_float_container_type(nbits) = (nbits == 32 ? Float32 : (nbits == 64 ? Float64 : error("$nbits bits is not supported for WAVE_FORMAT_IEEE_FLOAT.")))

function read_pcm_samples(io::IO, fmt::WAVFormat, subrange)
    const nbits = bits_per_sample(fmt)
    if isempty(subrange)
        return Array(pcm_container_type(nbits), 0, fmt.nchannels)
    end
    samples = Array(pcm_container_type(nbits), length(subrange), fmt.nchannels)
    sample_type = eltype(samples)
    const nbytes = ceil(Integer, nbits / 8)
    const bitshift = [0x0, 0x8, 0x10, 0x18, 0x20, 0x28, 0x30, 0x38, 0x40]
    mask = @compat UInt64(0x1) << (nbits - 1)
    if nbits <= 8
        mask = @compat UInt64(0)
    end
    skip(io, convert(UInt, (first(subrange) - 1) * nbytes * fmt.nchannels))
    for i = 1:size(samples, 1)
        for j = 1:size(samples, 2)
            raw_sample = read(io, UInt8, nbytes)
            my_sample = @compat UInt64(0)
            for k = 1:nbytes
                my_sample |= convert(UInt64, raw_sample[k]) << bitshift[k]
            end
            my_sample >>= nbytes * 8 - nbits
            # sign extend negative values
            my_sample = (my_sample $ mask) - mask
            samples[i, j] = convert(sample_type, signed(my_sample))
        end
    end
    samples
end

function read_ieee_float_samples(io::IO, fmt::WAVFormat, subrange, floatType)
    if isempty(subrange)
        return Array(floatType, 0, fmt.nchannels)
    end
    const nblocks = length(subrange)
    samples = Array(floatType, nblocks, fmt.nchannels)
    const nbits = bits_per_sample(fmt)
    skip(io, convert(UInt, (first(subrange) - 1) * (nbits / 8) * fmt.nchannels))
    for i = 1:nblocks
        for j = 1:fmt.nchannels
            samples[i, j] = read_le(io, floatType)
        end
    end
    samples
end

# take the loop variable type out of the loop
function read_ieee_float_samples(io::IO, fmt::WAVFormat, subrange)
    const floatType = ieee_float_container_type(bits_per_sample(fmt))
    read_ieee_float_samples(io, fmt, subrange, floatType)
end

function read_companded_samples(io::IO, fmt::WAVFormat, subrange, table)
    if isempty(subrange)
        return Array(eltype(table), 0, fmt.nchannels)
    end
    const nblocks = length(subrange)
    samples = Array(eltype(table), nblocks, fmt.nchannels)
    skip(io, convert(UInt, (first(subrange) - 1) * fmt.nchannels))
    for i = 1:nblocks
        for j = 1:fmt.nchannels
            # add one to value from blocks because A-law stores values from 0 to 255.
            const compressedByte::UInt8 = read(io, UInt8)
            # Julia indexing is 1-based; I need a value from 1 to 256
            samples[i, j] = table[compressedByte + 1]
        end
    end
    return samples
end

function read_mulaw_samples(io::IO, fmt::WAVFormat, subrange)
    # Quantized μ-law algorithm -- Use a look up table to convert
    # From Wikipedia, ITU-T Recommendation G.711 and G.191 specify the following intervals:
    #
    # ---------------------------------------+--------------------------------
    #  14 bit Binary Linear input code       | 8 bit Compressed code
    # ---------------------------------------+--------------------------------
    # +8158 to +4063 in 16 intervals of 256  |  0x80 + interval number
    # +4062 to +2015 in 16 intervals of 128  |  0x90 + interval number
    # +2014 to +991 in 16 intervals of 64    |  0xA0 + interval number
    # +990 to +479 in 16 intervals of 32     |  0xB0 + interval number
    # +478 to +223 in 16 intervals of 16     |  0xC0 + interval number
    # +222 to +95 in 16 intervals of 8       |  0xD0 + interval number
    # +94 to +31 in 16 intervals of 4        |  0xE0 + interval number
    # +30 to +1 in 15 intervals of 2         |  0xF0 + interval number
    # 0                                      |  0xFF
    # −1                                     |  0x7F
    # −31 to −2 in 15 intervals of 2         |  0x70 + interval number
    # −95 to −32 in 16 intervals of 4        |  0x60 + interval number
    # −223 to −96 in 16 intervals of 8       |  0x50 + interval number
    # −479 to −224 in 16 intervals of 16     |  0x40 + interval number
    # −991 to −480 in 16 intervals of 32     |  0x30 + interval number
    # −2015 to −992 in 16 intervals of 64    |  0x20 + interval number
    # −4063 to −2016 in 16 intervals of 128  |  0x10 + interval number
    # −8159 to −4064 in 16 intervals of 256  |  0x00 + interval number
    # ---------------------------------------+--------------------------------
    const MuLawDecompressTable =
    [
    -32124,-31100,-30076,-29052,-28028,-27004,-25980,-24956,
    -23932,-22908,-21884,-20860,-19836,-18812,-17788,-16764,
    -15996,-15484,-14972,-14460,-13948,-13436,-12924,-12412,
    -11900,-11388,-10876,-10364, -9852, -9340, -8828, -8316,
    -7932, -7676, -7420, -7164, -6908, -6652, -6396, -6140,
    -5884, -5628, -5372, -5116, -4860, -4604, -4348, -4092,
    -3900, -3772, -3644, -3516, -3388, -3260, -3132, -3004,
    -2876, -2748, -2620, -2492, -2364, -2236, -2108, -1980,
    -1884, -1820, -1756, -1692, -1628, -1564, -1500, -1436,
    -1372, -1308, -1244, -1180, -1116, -1052,  -988,  -924,
    -876,  -844,  -812,  -780,  -748,  -716,  -684,  -652,
    -620,  -588,  -556,  -524,  -492,  -460,  -428,  -396,
    -372,  -356,  -340,  -324,  -308,  -292,  -276,  -260,
    -244,  -228,  -212,  -196,  -180,  -164,  -148,  -132,
    -120,  -112,  -104,   -96,   -88,   -80,   -72,   -64,
    -56,   -48,   -40,   -32,   -24,   -16,    -8,     -1,
    32124, 31100, 30076, 29052, 28028, 27004, 25980, 24956,
    23932, 22908, 21884, 20860, 19836, 18812, 17788, 16764,
    15996, 15484, 14972, 14460, 13948, 13436, 12924, 12412,
    11900, 11388, 10876, 10364,  9852,  9340,  8828,  8316,
    7932,  7676,  7420,  7164,  6908,  6652,  6396,  6140,
    5884,  5628,  5372,  5116,  4860,  4604,  4348,  4092,
    3900,  3772,  3644,  3516,  3388,  3260,  3132,  3004,
    2876,  2748,  2620,  2492,  2364,  2236,  2108,  1980,
    1884,  1820,  1756,  1692,  1628,  1564,  1500,  1436,
    1372,  1308,  1244,  1180,  1116,  1052,   988,   924,
    876,   844,   812,   780,   748,   716,   684,   652,
    620,   588,   556,   524,   492,   460,   428,   396,
    372,   356,   340,   324,   308,   292,   276,   260,
    244,   228,   212,   196,   180,   164,   148,   132,
    120,   112,   104,    96,    88,    80,    72,    64,
    56,    48,    40,    32,    24,    16,     8,     0
     ]
    @assert length(MuLawDecompressTable) == 256
    return read_companded_samples(io, fmt, subrange, MuLawDecompressTable)
end

function read_alaw_samples(io::IO, fmt::WAVFormat, subrange)
    # Quantized A-law algorithm -- Use a look up table to convert
    const ALawDecompressTable =
    [
    -5504, -5248, -6016, -5760, -4480, -4224, -4992, -4736,
    -7552, -7296, -8064, -7808, -6528, -6272, -7040, -6784,
    -2752, -2624, -3008, -2880, -2240, -2112, -2496, -2368,
    -3776, -3648, -4032, -3904, -3264, -3136, -3520, -3392,
    -22016,-20992,-24064,-23040,-17920,-16896,-19968,-18944,
    -30208,-29184,-32256,-31232,-26112,-25088,-28160,-27136,
    -11008,-10496,-12032,-11520,-8960, -8448, -9984, -9472,
    -15104,-14592,-16128,-15616,-13056,-12544,-14080,-13568,
    -344,  -328,  -376,  -360,  -280,  -264,  -312,  -296,
    -472,  -456,  -504,  -488,  -408,  -392,  -440,  -424,
    -88,   -72,   -120,  -104,  -24,   -8,    -56,   -40,
    -216,  -200,  -248,  -232,  -152,  -136,  -184,  -168,
    -1376, -1312, -1504, -1440, -1120, -1056, -1248, -1184,
    -1888, -1824, -2016, -1952, -1632, -1568, -1760, -1696,
    -688,  -656,  -752,  -720,  -560,  -528,  -624,  -592,
    -944,  -912,  -1008, -976,  -816,  -784,  -880,  -848,
    5504,  5248,  6016,  5760,  4480,  4224,  4992,  4736,
    7552,  7296,  8064,  7808,  6528,  6272,  7040,  6784,
    2752,  2624,  3008,  2880,  2240,  2112,  2496,  2368,
    3776,  3648,  4032,  3904,  3264,  3136,  3520,  3392,
    22016, 20992, 24064, 23040, 17920, 16896, 19968, 18944,
    30208, 29184, 32256, 31232, 26112, 25088, 28160, 27136,
    11008, 10496, 12032, 11520, 8960,  8448,  9984,  9472,
    15104, 14592, 16128, 15616, 13056, 12544, 14080, 13568,
    344,   328,   376,   360,   280,   264,   312,   296,
    472,   456,   504,   488,   408,   392,   440,   424,
    88,    72,   120,   104,    24,     8,    56,    40,
    216,   200,   248,   232,   152,   136,   184,   168,
    1376,  1312,  1504,  1440,  1120,  1056,  1248,  1184,
    1888,  1824,  2016,  1952,  1632,  1568,  1760,  1696,
    688,   656,   752,   720,   560,   528,   624,   592,
    944,   912,  1008,   976,   816,   784,   880,   848
     ]
    @assert length(ALawDecompressTable) == 256
    return read_companded_samples(io, fmt, subrange, ALawDecompressTable)
end

function compress_sample_mulaw(sample)
    const MuLawCompressTable =
    [
    0,0,1,1,2,2,2,2,3,3,3,3,3,3,3,3,
    4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,
    5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,
    5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,
    6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,
    6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,
    6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,
    6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,6,
    7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,
    7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,
    7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,
    7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,
    7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,
    7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,
    7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,
    7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7
     ]
    @assert length(MuLawCompressTable) == 256
    const cBias = 0x84
    const cClip = 32635

    const sampleSign = (sample >>> 8) & 0x80
    if sampleSign != 0
        sample = -sample
    end
    if sample > cClip
        sample = cClip
    end
    sample = sample + cBias
    const sampleExponent = MuLawCompressTable[(sample >>> 7) + 1]
    const mantissa = (sample >> (sampleExponent+3)) & 0x0F
    @compat UInt8((~ (sampleSign | (sampleExponent << 4) | mantissa)) & 0xff)
end

function compress_sample_alaw(sample)
    const ALawCompressTable =
    [
    1,1,2,2,3,3,3,3,
    4,4,4,4,4,4,4,4,
    5,5,5,5,5,5,5,5,
    5,5,5,5,5,5,5,5,
    6,6,6,6,6,6,6,6,
    6,6,6,6,6,6,6,6,
    6,6,6,6,6,6,6,6,
    6,6,6,6,6,6,6,6,
    7,7,7,7,7,7,7,7,
    7,7,7,7,7,7,7,7,
    7,7,7,7,7,7,7,7,
    7,7,7,7,7,7,7,7,
    7,7,7,7,7,7,7,7,
    7,7,7,7,7,7,7,7,
    7,7,7,7,7,7,7,7,
    7,7,7,7,7,7,7,7
     ]
    @assert length(ALawCompressTable) == 128
    const cBias = 0x84
    const cClip = 32635
    const sampleSign = ((~sample >>> 8) & 0x80)
    if sampleSign == 0
        sample = -sample
    end
    if sample > cClip
        sample = cClip
    end
    compressedByte = 0
    if sample >= 256
        const sampleExponent = ALawCompressTable[((sample >>> 8) & 0x7f) + 1]
        const mantissa = (sample >>> (sampleExponent + 3) ) & 0x0f
        compressedByte = ((sampleExponent << 4) | mantissa) & 0xff
    else
        compressedByte = (sample >>> 4) & 0xff
    end
    compressedByte $= (sampleSign $ 0x55)
    @compat UInt8(compressedByte & 0xff)
end

function write_companded_samples{T<:Integer}(io::IO, samples::AbstractArray{T}, compander::Function)
    for i = 1:size(samples, 1)
        for j = 1:size(samples, 2)
            write_le(io, compander(samples[i, j]))
        end
    end
end

function write_companded_samples{T<:AbstractFloat}(io::IO, samples::AbstractArray{T}, compander::Function)
    samples = convert(Array{Int16}, round(samples * typemax(Int16)))
    write_companded_samples(io, samples, compander)
end

# PCM data is two's-complement except for resolutions of 1-8 bits, which are represented as offset binary.

# support every bit width from 1 to 8 bits
convert_pcm_to_double(samples::AbstractArray{UInt8}, nbits::Integer) = convert(Array{Float64}, samples) ./ (2.0^nbits - 1) .* 2.0 .- 1.0
convert_pcm_to_double(::AbstractArray{Int8}, ::Integer) = error("WAV files use offset binary for less than 9 bits")
# support every bit width from 9 to 64 bits
convert_pcm_to_double{T<:Signed}(samples::AbstractArray{T}, nbits::Integer) = convert(Array{Float64}, samples) / (2.0^(nbits - 1) - 1)

function read_data(io::IO, chunk_size, fmt::WAVFormat, format, subrange)
    # "format" is the format of values, while "fmt" is the WAV file level format
    convert_to_double = x -> convert(Array{Float64}, x)

    if subrange === Void
        # each block stores fmt.nchannels channels
        subrange = 1:convert(UInt, chunk_size / fmt.block_align)
    end
    if isformat(fmt, WAVE_FORMAT_PCM)
        samples = read_pcm_samples(io, fmt, subrange)
        convert_to_double = x -> convert_pcm_to_double(x, bits_per_sample(fmt))
    elseif isformat(fmt, WAVE_FORMAT_IEEE_FLOAT)
        samples = read_ieee_float_samples(io, fmt, subrange)
    elseif isformat(fmt, WAVE_FORMAT_MULAW)
        samples = read_mulaw_samples(io, fmt, subrange)
        convert_to_double = x -> convert_pcm_to_double(x, 16)
    elseif isformat(fmt, WAVE_FORMAT_ALAW)
        samples = read_alaw_samples(io, fmt, subrange)
        convert_to_double = x -> convert_pcm_to_double(x, 16)
    else
        error("$(fmt.compression_code) is an unsupported compression code!")
    end
    if format == "double"
        samples = convert_to_double(samples)
    end
    samples
end

function write_pcm_samples{T<:Integer}(io::IO, fmt::WAVFormat, samples::AbstractArray{T})
    const nbits = bits_per_sample(fmt)
    # number of bytes per sample
    const nbytes = ceil(Integer, nbits / 8)
    for i = 1:size(samples, 1)
        for j = 1:size(samples, 2)
            my_sample = samples[i, j]
            # shift my_sample into the N most significant bits
            my_sample <<= nbytes * 8 - nbits
            for k = 1:nbytes
                write_le(io, convert(UInt8, my_sample & 0xff))
                my_sample = my_sample >> 8
            end
        end
    end
end

function write_pcm_samples{T<:AbstractFloat}(io::IO, fmt::WAVFormat, samples::AbstractArray{T})
    const nbits = bits_per_sample(fmt)
    # Scale the floating point values to the PCM range
    if nbits > 8
        # two's complement
        samples = convert(Array{pcm_container_type(nbits)}, round(samples * (2.0^(nbits - 1) - 1)))
    else
        # offset binary
        samples = convert(Array{UInt8}, round((samples .+ 1.0) / 2.0 * (2.0^nbits - 1)))
    end
    return write_pcm_samples(io, fmt, samples)
end

function write_ieee_float_samples(io::IO, samples)
    # Interleave the channel samples before writing to the stream.
    for i = 1:size(samples, 1) # for each sample
        for j = 1:size(samples, 2) # for each channel
            write_le(io, samples[i, j])
        end
    end
end

# take the loop variable type out of the loop
function write_ieee_float_samples(io::IO, fmt::WAVFormat, samples)
    const floatType = ieee_float_container_type(bits_per_sample(fmt))
    write_ieee_float_samples(io, convert(Array{floatType}, samples))
end

function write_data(io::IO, fmt::WAVFormat, samples::AbstractArray)
    if isformat(fmt, WAVE_FORMAT_PCM)
        return write_pcm_samples(io, fmt, samples)
    elseif isformat(fmt, WAVE_FORMAT_IEEE_FLOAT)
        return write_ieee_float_samples(io, fmt, samples)
    elseif isformat(fmt, WAVE_FORMAT_MULAW)
        return write_companded_samples(io, samples, compress_sample_mulaw)
    elseif isformat(fmt, WAVE_FORMAT_ALAW)
        return write_companded_samples(io, samples, compress_sample_alaw)
    else
        error("$(fmt.compression_code) is an unsupported compression code.")
    end
end

make_range(subrange) = subrange
make_range(subrange::Number) = 1:convert(Int, subrange)

function wavread(io::IO; subrange=Void, format="double")
    chunk_size = read_header(io)
    samples = Array(Float64)
    nbits = 0
    sample_rate = @compat Float32(0.0)
    opt = Dict{Symbol, Any}()

    # Note: This assumes that the format chunk is written in the file before the data chunk. The
    # specification does not require this assumption, but most real files are written that way.

    # Subtract the size of the format field from chunk_size; now it holds the size
    # of all the sub-chunks
    chunk_size -= 4
    # GitHub Issue #18: Check if there is enough data to read another chunk
    const subchunk_header_size = 4 + sizeof(UInt32)
    while chunk_size >= subchunk_header_size
        # Read subchunk ID and size
        subchunk_id = read(io, UInt8, 4)
        subchunk_size = read_le(io, UInt32)
        if subchunk_size > chunk_size
            chunk_size = 0
            break
        end
        chunk_size -= subchunk_header_size + subchunk_size
        # check the subchunk ID
        if subchunk_id == b"fmt "
            fmt = read_format(io, subchunk_size)
            sample_rate = @compat Float32(fmt.sample_rate)
            nbits = bits_per_sample(fmt)
            opt[:fmt] = fmt
        elseif subchunk_id == b"data"
            if format == "size"
                return convert(Int, subchunk_size / fmt.block_align), convert(Int, fmt.nchannels)
            end
            samples = read_data(io, subchunk_size, fmt, format, make_range(subrange))
        else
            opt[@compat Symbol(subchunk_id)] = read(io, UInt8, subchunk_size)
        end
    end
    return samples, sample_rate, nbits, opt
end

function wavread(filename::AbstractString; subrange=Void, format="double")
    open(filename, "r") do io
        wavread(io, subrange=subrange, format=format)
    end
end

# These are the MATLAB compatible signatures
wavread(filename::AbstractString, fmt::AbstractString) = wavread(filename, format=fmt)
wavread(filename::AbstractString, n) = wavread(filename, subrange=n)
wavread(filename::AbstractString, n, fmt) = wavread(filename, subrange=n, format=fmt)

get_default_compression{T<:Integer}(::AbstractArray{T}) = WAVE_FORMAT_PCM
get_default_compression{T<:AbstractFloat}(::AbstractArray{T}) = WAVE_FORMAT_IEEE_FLOAT
get_default_pcm_precision(::AbstractArray{UInt8}) = 8
get_default_pcm_precision(::AbstractArray{Int16}) = 16
get_default_pcm_precision(::Any) = 24

function get_default_precision(samples, compression)
    if compression == WAVE_FORMAT_ALAW || compression == WAVE_FORMAT_MULAW
        return 8
    elseif compression == WAVE_FORMAT_IEEE_FLOAT
        return 32
    end
    get_default_pcm_precision(samples)
end

function wavwrite(samples::AbstractArray, io::IO; Fs=8000, nbits=0, compression=0,
                  chunks::Dict{Symbol, Array{UInt8,1}}=Dict{Symbol, Array{UInt8,1}}())
    if compression == 0
        compression = get_default_compression(samples)
    elseif compression == WAVE_FORMAT_ALAW || compression == WAVE_FORMAT_MULAW
        nbits = 8
    end
    if nbits == 0
        nbits = get_default_precision(samples, compression)
    end
    compression_code = compression
    const nchannels = size(samples, 2)
    const sample_rate = Fs
    const my_nbits = ceil(Integer, nbits / 8) * 8
    const block_align = my_nbits / 8 * nchannels
    const bps = sample_rate * block_align
    const data_length::UInt32 = size(samples, 1) * block_align
    ext = WAVFormatExtension()

    if nchannels > 2 || my_nbits > 16 || my_nbits != nbits
        compression_code = WAVE_FORMAT_EXTENSIBLE
        const valid_bits_per_sample = nbits
        const channel_mask = 0
        sub_format = Array(UInt8, 0)
        if compression == WAVE_FORMAT_PCM
            sub_format = KSDATAFORMAT_SUBTYPE_PCM
        elseif compression == WAVE_FORMAT_IEEE_FLOAT
            sub_format = KSDATAFORMAT_SUBTYPE_IEEE_FLOAT
        elseif compression == WAVE_FORMAT_ALAW
            sub_format = KSDATAFORMAT_SUBTYPE_ALAW
        elseif compression == WAVE_FORMAT_MULAW
            sub_format = KSDATAFORMAT_SUBTYPE_MULAW
        else
            error("Unsupported extension sub format: $compression")
        end
        ext = WAVFormatExtension(valid_bits_per_sample, channel_mask, sub_format)
        write_extended_header(io, data_length)
    else
        write_standard_header(io, data_length)
    end
    fmt = WAVFormat(compression_code,
                    nchannels,
                    sample_rate,
                    bps,
                    block_align,
                    my_nbits,
                    ext)
    write_format(io, fmt)

    for eachchunk in chunks
        write(io, eachchunk[1])
        write_le(io, @compat UInt32(length(eachchunk[2])))
        for eachbyte in eachchunk[2]
            write(io, eachbyte)
        end
    end

    # write the data subchunk header
    write(io, b"data")
    write_le(io, data_length) # UInt32
    write_data(io, fmt, samples)
end

function wavwrite(samples::AbstractArray, filename::AbstractString; Fs=8000, nbits=0, compression=0,
                  chunks::Dict{Symbol, Array{UInt8,1}}=Dict{Symbol, Array{UInt8,1}}())
    open(filename, "w") do io
        wavwrite(samples, io, Fs=Fs, nbits=nbits, compression=compression, chunks=chunks)
    end
end

function wavappend(samples::AbstractArray, io::IO)
    seekstart(io)
    chunk_size = read_header(io)
    subchunk_id = read(io, UInt8, 4)
    subchunk_size = read_le(io, UInt32)
    if subchunk_id != b"fmt "
        error("First chunk is not the format")
    end
    fmt = read_format(io, subchunk_size)

    if fmt.nchannels != size(samples,2)
        error("Number of channels do not match")
    end

    const data_length = size(samples, 1) * fmt.block_align

    seek(io,4)
    write_le(io, convert(UInt32, chunk_size + data_length))

    seek(io,64)
    subchunk_size = read_le(io, UInt32)
    seek(io,64)
    write_le(io, convert(UInt32, subchunk_size + data_length))

    seekend(io)
    write_data(io, fmt, samples)
end

function wavappend(samples::AbstractArray, filename::AbstractString)
    open(filename, "a+") do io
        wavappend(samples,io)
    end
end

wavwrite(y::AbstractArray, f::Real, filename::AbstractString) = wavwrite(y, filename, Fs=f)
wavwrite(y::AbstractArray, f::Real, n::Real, filename::AbstractString) = wavwrite(y, filename, Fs=f, nbits=n)

# support for writing native arrays...
wavwrite{T<:Integer}(y::AbstractArray{T}, io::IO) = wavwrite(y, io, nbits=sizeof(T)*8)
wavwrite{T<:Integer}(y::AbstractArray{T}, filename::AbstractString) = wavwrite(y, filename, nbits=sizeof(T)*8)
wavwrite(y::AbstractArray{Int32}, io::IO) = wavwrite(y, io, nbits=24)
wavwrite(y::AbstractArray{Int32}, filename::AbstractString) = wavwrite(y, filename, nbits=24)
wavwrite{T<:AbstractFloat}(y::AbstractArray{T}, io::IO) = wavwrite(y, io, nbits=sizeof(T)*8, compression=WAVE_FORMAT_IEEE_FLOAT)
wavwrite{T<:AbstractFloat}(y::AbstractArray{T}, filename::AbstractString) = wavwrite(y, filename, nbits=sizeof(T)*8, compression=WAVE_FORMAT_IEEE_FLOAT)

# FileIO integration support
load(s::Stream{format"WAV"}; kwargs...) = wavread(s.io; kwargs...)
save(s::Stream{format"WAV"}, data; kwargs...) = wavwrite(data, s.io; kwargs...)

load(f::File{format"WAV"}; kwargs...) = wavread(f.filename; kwargs...)
save(f::File{format"WAV"}, data; kwargs...) = wavwrite(data, f.filename; kwargs...)

end # module


text:
import LinearAlgebra
import SparseArrays
#using Images
import Images
import MAT

LA = LinearAlgebra
SA = SparseArrays
Img = Images

# Matrix
dim = 4
A = rand(dim,dim)
AT = A'
A2 = A*AT

@show A[11] == A[1,2]
b = rand(dim)
# Solve for x where Ax = b
x = A\b
@show LA.norm(A*x - b)

@show typeof(A)
@show typeof(b)
@show typeof(rand(1,dim))
@show typeof(AT)

#A = [1 2;3 4]
sizeof(A) # => byte count
sizeof(AT)
B = copy(AT)
sizeof(B)

# \(x, y) => x^(-1) * y
#x = \(A,b) # => A^(-1)*b
@show \(3,6) # => 1/3 * 6
@show 3\6    # => 1/3 * 6

1//3 + 1//2

#A = sparse([1,2,4], [1,1,1], [1.0,1.0,1.0], 4, 2)
#F = qr(A)

luA = LA.lu(A)
# LU = PA
LA.norm(luA.L*luA.U - luA.P*A)

# QR = A
qrA = LA.qr(A)
Q = qrA.Q
R = qrA.R
LA.norm(Q*R - A)

# Skip Cholesky factorization

# Skip Sparse Linear Algebra

# Images as matrices
X1 = Images.load("data/khiam-small.jpg")
# Gray scale
Xgray = Images.Gray.(X1)
# Extract RGB values
R = map(i -> X1[i].r, 1:length(X1))    # => Vector
R2 = Float64.(reshape(R, size(X1)...)) # (80089,) => (283,283)
G = map(i -> X1[i].g, 1:length(X1))
G2 = Float64.(reshape(G, size(X1)...)) # (80089,) => (283,283)
B = map(i -> X1[i].b, 1:length(X1))
B2 = Float64.(reshape(B, size(X1)...))

Z = zeros(size(G2)...)
Images.RGB.(Z,G2,Z)

Xgrayvalues = Float64.(Xgray)
# SVD: A = U*Sigma*V'
usv =LA.svd(Xgrayvalues)
# usv.U, S, V, Vt
U = usv.U
S = usv.S
V = usv.V
Vt = usv.Vt

LA.norm(U*LA.diagm(S)*Vt - Xgrayvalues)

# Use the top 4 singular vectors/values to form a new image
i = 1
u1 = usv.U[:,i]
v1 = usv.V[:,i]
img1 = usv.S[i]*u1*v1'

i = 2
u1 = usv.U[:,i]
v1 = usv.V[:,i]
img1 += usv.S[i]*u1*v1'

i = 3
u1 = usv.U[:,i]
v1 = usv.V[:,i]
img1 += usv.S[i]*u1*v1'

i = 4
u1 = usv.U[:,i]
v1 = usv.V[:,i]
img1 += usv.S[i]*u1*v1'

Images.Gray.(img1)

i = 1:100
u1 = usv.U[:,i]
s1 = SA.spdiagm(0=>usv.S[i])
v1 = usv.V[:,i]
img1 = u1 * s1 * v1'
Img.Gray.(img1)
# Close but not the same
LA.norm(Xgrayvalues - img1)

# Face recognition
M = MAT.matread("data/face_recog_qr.mat")
# V2 contains 490 faces
# each column is a face and need to reshape into 192,168
V2 = M["V2"]
q = reshape(V2[:,200], 192, 168)
Img.Gray.(q)
b = q[:]

A = V2[:,2:end]
x = A\b
Img.Gray.(reshape(A*x, 192,168))
LA.norm(A*x - b)

# Skip


text:
@testset "Costs" begin
    trials = 100

    for m = 1:100, K = 1:4
        n = m
        A = dropzeros!(sprand(Int, m, n, 0.125))
        Π = MapPartition(K, rand(1:K, m))
        Φ = SplitPartition(K, [1, sort(rand(1:(n + 1), K - 1))..., n + 1])

        models = [
            AffineWorkModel(α = 0, β_vertex = 0, β_pin = 1),
            AffineWorkModel(α = 0, β_vertex = 1, β_pin = 0),
            AffineWorkModel(α = 0, β_vertex = 1, β_pin = 1),
            AffineConnectivityModel(α = 0, β_vertex = 0, β_pin = 0, β_net = 1),
            AffineConnectivityModel(α = 0, β_vertex = 0, β_pin = 1, β_net = 0),
            AffineConnectivityModel(α = 0, β_vertex = 0, β_pin = 1, β_net = 1),
            AffineConnectivityModel(α = 0, β_vertex = 1, β_pin = 0, β_net = 0),
            AffineConnectivityModel(α = 0, β_vertex = 1, β_pin = 0, β_net = 1),
            AffineConnectivityModel(α = 0, β_vertex = 1, β_pin = 1, β_net = 0),
            AffineConnectivityModel(α = 0, β_vertex = 1, β_pin = 1, β_net = 1),
            AffineMonotonizedSymmetricConnectivityModel(α = 10, β_vertex = 10, β_over_pin = 10, β_dia_net = 10, Δ_pins = 0),
            AffineMonotonizedSymmetricConnectivityModel(α = 10, β_vertex = 10, β_over_pin = 10, β_dia_net = 100, Δ_pins = 8),
        ]
        for mdl in models
            ocl = oracle_stripe(mdl, A, Φ)
            @test bottleneck_value(A, Π, Φ, mdl) == bottleneck_value(A, Π, Φ, ocl)
            @test bound_stripe(A, K, Π, mdl) == bound_stripe(A, K, Π, ocl)
            c_lo, c_hi = bound_stripe(A, K, Π, mdl)
            @test 0 <= c_lo <= bottleneck_value(A, Π, Φ, mdl) <= c_hi
        end

        models = [
            AffineHyperedgeCutModel(α = 0, β_vertex = 0, β_pin = 0, β_self_net = 0, β_cut_net = 1),
            AffineHyperedgeCutModel(α = 0, β_vertex = 0, β_pin = 0, β_self_net = 1, β_cut_net = 0),
            AffineHyperedgeCutModel(α = 0, β_vertex = 0, β_pin = 0, β_self_net = 1, β_cut_net = 1),
            AffineHyperedgeCutModel(α = 0, β_vertex = 0, β_pin = 0, β_self_net = -1),
            AffineSymmetricEdgeCutModel(α = 10, β_vertex = 10, β_self_pin = 10, β_cut_pin = 100),
            AffineSymmetricConnectivityModel(α = 0, β_vertex = 0, β_pin = 0, β_local_net = 0, β_remote_net = 1),
            AffineSymmetricConnectivityModel(α = 0, β_vertex = 0, β_pin = 0, β_local_net = 1, β_remote_net = 1),
            AffineSymmetricConnectivityModel(α = 1, β_vertex = 1, β_pin = 1, β_local_net = 1, β_remote_net = 1),
            AffineSymmetricEdgeCutModel(α = 0, β_vertex = 0, β_self_pin = 0, β_cut_pin = 1),
            AffineSymmetricEdgeCutModel(α = 0, β_vertex = 0, β_self_pin = 1, β_cut_pin = 1),
            AffineSymmetricEdgeCutModel(α = 1, β_vertex = 1, β_self_pin = 1, β_cut_pin = 1),
        ]
        for mdl in models
            ocl = oracle_stripe(mdl, A, Φ)
            @test bottleneck_value(A, Π, Φ, mdl) == bottleneck_value(A, Π, Φ, ocl)
        end

        Π = SplitPartition(K, [1, sort(rand(1:(m + 1), K - 1))..., m + 1])
        models = [
            (AffinePrimaryConnectivityModel(α = 0, β_vertex = 0, β_pin = 0, β_local_net = 0, β_remote_net = 1),
            AffineSecondaryConnectivityModel(α = 0, β_vertex = 0, β_pin = 0, β_local_net = 0, β_remote_net = 1)),
            (AffinePrimaryConnectivityModel(α = 0, β_vertex = 0, β_pin = 0, β_local_net = 1, β_remote_net = 1),
            AffineSecondaryConnectivityModel(α = 0, β_vertex = 0, β_pin = 0, β_local_net = 1, β_remote_net = 1)),
            (AffinePrimaryConnectivityModel(α = 1, β_vertex = 1, β_pin = 1, β_local_net = 1, β_remote_net = 1),
            AffineSecondaryConnectivityModel(α = 1, β_vertex = 1, β_pin = 1, β_local_net = 1, β_remote_net = 1)),
            (AffinePrimaryEdgeCutModel(α = 0, β_vertex = 0, β_self_pin = 0, β_cut_pin = 1),
            AffineSecondaryEdgeCutModel(α = 0, β_vertex = 0, β_self_pin = 0, β_cut_pin = 1)),
            (AffinePrimaryEdgeCutModel(α = 0, β_vertex = 0, β_self_pin = 1, β_cut_pin = 1),
            AffineSecondaryEdgeCutModel(α = 0, β_vertex = 0, β_self_pin = 1, β_cut_pin = 1)),
            (AffinePrimaryEdgeCutModel(α = 1, β_vertex = 1, β_self_pin = 1, β_cut_pin = 1),
            AffineSecondaryEdgeCutModel(α = 1, β_vertex = 1, β_self_pin = 1, β_cut_pin = 1)),
        ]
        adj_A = permutedims(A)
        for (comm_mdl, local_mdl) in models
            comm_ocl = oracle_stripe(comm_mdl, A, Π)
            local_ocl = oracle_stripe(local_mdl, adj_A, Φ)
            @test bottleneck_value(A, Π, Φ, comm_mdl) == bottleneck_value(A, Π, Φ, comm_ocl)
            @test bottleneck_value(adj_A, Φ, Π, local_mdl) == bottleneck_value(adj_A, Φ, Π, local_ocl)
            @test bound_stripe(A, K, Π, comm_mdl) == bound_stripe(A, K, Π, comm_ocl)
            @test bound_stripe(adj_A, K, Φ, local_mdl) == bound_stripe(adj_A, K, Φ, local_ocl)
            c_lo, c_hi = bound_stripe(A, K, Π, comm_mdl)
            @test 0 <= c_lo <= bottleneck_value(A, Π, Φ, comm_mdl) <= c_hi
            c_lo, c_hi = bound_stripe(adj_A, K, Φ, local_mdl)
            @test 0 <= c_lo <= bottleneck_value(adj_A, Φ, Π, local_mdl) <= c_hi
            @test bottleneck_value(adj_A, Φ, Π, local_mdl) == bottleneck_value(A, Π, Φ, comm_mdl)
            @test bottleneck_value(adj_A, Φ, Π, local_ocl) == bottleneck_value(A, Π, Φ, comm_ocl)
        end

        models = [
            (AffinePrimaryConnectivityModel(α = 0, β_vertex = 0, β_pin = 0, β_local_net = 0, β_remote_net = 1),
            AffineSymmetricConnectivityModel(α = 0, β_vertex = 0, β_pin = 0, β_local_net = 0, β_remote_net = 1)),
            (AffinePrimaryConnectivityModel(α = 0, β_vertex = 0, β_pin = 0, β_local_net = 1, β_remote_net = 1),
            AffineSymmetricConnectivityModel(α = 0, β_vertex = 0, β_pin = 0, β_local_net = 1, β_remote_net = 1)),
            (AffinePrimaryConnectivityModel(α = 1, β_vertex = 1, β_pin = 1, β_local_net = 1, β_remote_net = 1),
            AffineSymmetricConnectivityModel(α = 1, β_vertex = 1, β_pin = 1, β_local_net = 1, β_remote_net = 1)),
            (AffinePrimaryEdgeCutModel(α = 0, β_vertex = 0, β_self_pin = 0, β_cut_pin = 1),
            AffineSymmetricEdgeCutModel(α = 0, β_vertex = 0, β_self_pin = 0, β_cut_pin = 1)),
            (AffinePrimaryEdgeCutModel(α = 0, β_vertex = 0, β_self_pin = 1, β_cut_pin = 1),
            AffineSymmetricEdgeCutModel(α = 0, β_vertex = 0, β_self_pin = 1, β_cut_pin = 1)),
            (AffinePrimaryEdgeCutModel(α = 1, β_vertex = 1, β_self_pin = 1, β_cut_pin = 1),
            AffineSymmetricEdgeCutModel(α = 1, β_vertex = 1, β_self_pin = 1, β_cut_pin = 1)),
        ]

        for (comm_mdl, sym_mdl) in models
            comm_ocl = oracle_stripe(comm_mdl, A, Π)
            sym_ocl = oracle_stripe(sym_mdl, A)
            @test bottleneck_value(A, Π, Π, comm_mdl) == bottleneck_value(A, Π, Π, comm_ocl)
            @test bottleneck_value(A, Π, sym_mdl) == bottleneck_value(A, Π, sym_ocl)
            @test bottleneck_value(A, Π, sym_mdl) == bottleneck_value(A, Π, Π, comm_mdl)
            @test bottleneck_value(A, Π, sym_ocl) == bottleneck_value(A, Π, Π, comm_ocl)
        end
    end

    for m = 1:100, u = 1:4, w = 1:4
        n = m
        A = dropzeros!(sprand(Int, m, n, 0.125))
        adj_A = permutedims(A)
        Π = pack_stripe(A, EquiChunker(u))
        Φ = pack_stripe(A, EquiChunker(w))
        models = (
            BlockComponentCostModel{Int64}(β_row = (2, identity), β_col = (2, x->2x)),
            BlockComponentCostModel{Int64}(α_row = identity, α_col = x->3x, β_row = (2, identity), β_col = (2, x->2x)),
            )
        for mdl in models
            ocl = oracle_stripe(mdl, A, Π)
            adj_ocl = oracle_stripe(mdl, adj_A, Φ)
            @test total_value(A, Π, Φ, mdl) == total_value(A, Π, Φ, ocl)
            @test total_value(adj_A, Φ, Π, mdl) == total_value(adj_A, Φ, Π, adj_ocl)
        end
    end

    #TODO add wrapper cost functions to test that oracle versions of partitioners, bounds, and cost evaluators agree with direct versions.

    #=
    for m = 1:100, h = 1:4
        n = m
        k = 1000
        A = dropzeros!(sprand(UInt, m, n, 0.5))
        c = CommunicationCostFamily((l, w, d) -> d)(A)
        @testset "h = $h, m = $m" begin
            for l = 1:k
                j = rand(1:n)
                j′ = rand(j:n) + 1
                if !((@test c(j, j′) == length(union(A.rowval[A.colptr[j]:A.colptr[j′] - 1]))) isa Test.Pass)
                    @info "debugging stuff:" A, j, j′
                    break
                end
            end
        end
    end
    =#
end


text:
struct RamanTensor{T, T1, T2, T3}
    tensor::T
    tau1::T1
    tau2::T2
    raman_fraction::T3
end
struct ElectronicTensor{T, T1}
    tensor::T
    nl_index::T1
end

abstract type AbstractDielectric end
struct DielectricCoefficient <: AbstractDielectric end
struct DielectricTensor <: AbstractDielectric end

struct Material
    electronic::ElectronicTensor
    raman::RamanTensor
end

mutable struct OpticalAttr{TF<:Unitful.Frequency, T, F}
    source::Vector{TF}
    property::Vector{T}
    label::String
    fit_func::F
    OpticalAttr{TF, T, F}(s, p, label) where {TF, T, F} = fit(new(s, p, label))
    OpticalAttr{TF, T, F}(s, p, label, f) where {TF, T, F} = new(s, p, label, f)
end

function OpticalAttr(source, property, label, fit_func)
    OpticalAttr{eltype(source), eltype(property), typeof(fit_func)}(source, property, label, fit_func)
end
function OpticalAttr(source::AbstractVector{TF}, property, label) where {TF<:Unitful.Frequency}
    T = eltype(property)
    OpticalAttr{TF, T, ScaledFit{TF, Poly{T}}}(source, property, label)
end
function OpticalAttr(source::AbstractVector{TL}, property, label) where {TL<:Unitful.Length}
    source_infrequency = frequency.(source)
    TF = eltype(source_infrequency)
    T = eltype(property)
    OpticalAttr{TF, T, ScaledFit{TF, Poly{T}}}(source_infrequency, property, label)
end
function OpticalAttr(property::Number, label)
    TF = typeof(1.0THz2π)
    T = typeof(property)
    OpticalAttr{TF, T, Poly{T}}([TF(0.0)], [property], label, Poly([property]))
end

abstract type AbstractMode end
mutable struct Mode <: AbstractMode
    effectiveindex::OpticalAttr
    effectivearea::OpticalAttr
    linearloss::OpticalAttr
    corefraction::OpticalAttr
    coupling::OpticalAttr
    polarization::Symbol
    has_shock::Bool
    has_raman::Bool
    has_interaction::Bool
end
Mode(n_eff, area_eff) = Mode(n_eff, area_eff,
                             OpticalAttr(0.0, "Linear Loss"),
                             OpticalAttr(1.0, "Core Fraction"),
                             OpticalAttr(1.0, "Coupling"),
                             :TM,
                             false,
                             false,
                             false,
                             )
Mode(n_eff, area_eff, lin_loss, corefrac, coupling, pol, shock, raman) = Mode(
     n_eff, area_eff, lin_loss, corefrac, coupling, pol, shock, raman, false)

struct ToyMode <: AbstractMode
    beta::Vector{OpticalAttr}
    effectivearea::OpticalAttr
    linearloss::OpticalAttr
    corefraction::OpticalAttr
    coupling::OpticalAttr
end
ToyMode(beta, area_eff) = ToyMode(beta, area_eff,
                          OpticalAttr(0.0, "Linear Loss (1/m)"),
                          OpticalAttr(1.0, "Core Fraction"),
                          OpticalAttr(1.0, "Coupling"))

abstract type AbstractLaser end
mutable struct CWLaser{TF<:Unitful.Frequency, T1, TP<:Unitful.Power} <: AbstractLaser
    frequency::typeof(1.0THz2π)
    detuning::T1
    power::TP
end
CWLaser(f, δ, P) = CWLaser{typeof(δ), typeof(P)}(f, δ, P)
mutable struct PulsedLaser{TF<:Unitful.Frequency,
                           TP<:Unitful.Power,
                           TT<:Unitful.Time} <: AbstractLaser
    frequency::TF
    power_in::TP
    pulsetime::TT
end


text:
using Documenter
using ResumableFunctions
using SimJulia

makedocs(
  format   = :html,
  sitename = "SimJulia",
  authors = "Ben Lauwens",
  pages    = [
    "Home" => "index.md",
    "Tutorial" => "tutorial.md",
    "Topical Guides" => ["Basics" => "guides/basics.md",
                         "Environments" => "guides/environments.md",
                         "Events" => "guides/events.md",],
    "Examples" => ["Ross" => "examples/ross.md",],
    "API" => "api.md"
  ]
)

deploydocs(
  repo = "github.com/BenLauwens/SimJulia.jl.git",
  julia  = "1.0",
  osname = "linux",
  target = "build",
  deps = nothing,
  make = nothing,
)


text:
function crandn(dims...)
    sqrt(0.5) * ( randn(dims) + 1im*randn(dims) )
end
# Define "Q" matrix transformations:
# Q*x  = trans(y)
# Q'*x = trans(y,:transp)

# Q = unitary DFT
function circ_trans(x,tp)
    if tp == :transp
        y = ifft(x,1)*sqrt(size(x,1))
    else
        y = fft(x,1)/sqrt(size(x,1))
    end
    y
end
# Q = first M columns of 2Mx2M unitary DFT
function toep_trans(x,tp)
    if tp == :transp
        y = ifft(x,1)[1:Int(end/2),:,:]*sqrt(size(x,1))
        if length(size(x)) == 2
            y = y[:,:,1]
        end
    else
        y = fft([x;zeros(x)],1)/sqrt(2*size(x,1))
    end
    y
end

function train!(nn_est; snr = 0, nBatches = 1, get_channel = () -> 0.0, verbose = false)
    verbose && @printf "Learning: "
    for b in 1:nBatches
        verbose && mod(b,ceil(Int,nBatches/10))==0 && @printf " ... %.0f%%" b/nBatches*100

        (h,h_cov) = get_channel()
        y         = h + 10^(-snr/20) * crandn(size(h)...)
        for (_,nn) in nn_est
            cntf.train!(nn,y,h)
        end
    end
    verbose && @printf "\n"
end

function init_hier!(nn_est, init_params; verbose = false)
    for i in 1:length(init_params[:nBatches])
        verbose && @printf "Hierarchical Learning %.0f/%.0f\n" i length(init_params[:nBatches])
        nAntennas  = init_params[:nAntennas][i]
        nBatches   = init_params[:nBatches][i]
        nBatchSize = init_params[:nBatchSize][i]
        snr        = init_params[:snrs][i]
        train!(nn_est, snr = snr, nBatches = nBatches, get_channel = () -> init_params[:get_channel](nAntennas, nBatchSize), verbose = verbose)
        for (alg,nn) in nn_est
            nn_est[alg] = cntf.resize(nn, init_params[:nAntennas][i+1], learning_rate = init_params[:learning_rates][alg][i+1])
        end
    end
end

function evaluate(algs; snr = 0, nBatches = 1, get_channel = () -> 0.0, verbose = false)
    errs  = Dict{Symbol,Any}()
    rates = Dict{Symbol,Any}()

    rho = 10^(0.1*snr);

    for alg in keys(algs)
        errs[alg]  = 0.0
        rates[alg] = 0.0
    end

    # Generate channels, calculate errors and achievable rates
    verbose && @printf "Simulating: "
    for bb in 1:nBatches
        verbose && mod(bb,ceil(Int,nBatches/10))==0 && @printf " ... %.0f%%" bb/nBatches*100
        (h,h_cov) = get_channel()
        y   = h + 10^(-snr/20) * crandn(size(h)...)
        (nAntennas,nCoherence,nBatchSize) = size(h)
        for (alg,est) in algs
            hest = est(y,h,h_cov)
            errs[alg] += sum(abs2,h-hest)/length(h)/nBatches
            # Achievable rates
            for b in 1:nBatchSize, t in 1:nCoherence
                if sum(abs2,hest[:,t,b]) < 1e-8
                    verbose && warn(string(alg) * ": channel estimate is zero")
                    continue
                end

                rates[alg] += log2(1 + rho*abs2(dot(h[:,t,b],hest[:,t,b]))/max(1e-8,sum(abs2,hest[:,t,b])))/length(h[1,:,:])/nBatches
            end
        end
    end
    verbose && @printf "\n"
    (errs, rates)
end

max_stars_repo_name: tum-msv/learning-mmse-est


text:
path = dirname(@__FILE__)

# Initialize arguments to function
@load "$path/reference/kalman_filter_args.jld2" y T R C Q Z D E z0 P0

file  = h5open("$path/reference/kalman_filter_out.h5", "r")
pred  = read(file, "pred")
vpred = read(file, "vpred")
filt  = read(file, "filt")
vfilt = read(file, "vfilt")
close(file)

# Run smoothers
states = Dict{Symbol, Matrix{Float64}}()
shocks = Dict{Symbol, Matrix{Float64}}()

states[:hamilton], shocks[:hamilton] = hamilton_smoother(y, T, R, C, Q, Z, D, E, z0, P0)
states[:koopman], shocks[:koopman] = koopman_smoother(y, T, R, C, Q, Z, D, E, z0, P0, pred, vpred)
states[:carter_kohn], shocks[:carter_kohn] = carter_kohn_smoother(y, T, R, C, Q, Z, D, E, z0, P0; draw_states = false)
states[:durbin_koopman], shocks[:durbin_koopman] = durbin_koopman_smoother(y, T, R, C, Q, Z, D, E, z0, P0; draw_states = false)

# Check that last-period smoothed states equal last-period filtered states
@testset "Ensure last smoothed state = the last filtered state" begin
    for smoother in [:hamilton, :koopman, :carter_kohn, :durbin_koopman]
        @test filt[:, end] ≈ states[smoother][:, end]
    end
end

# Compare to expected output
exp_states, exp_shocks = h5open("$path/reference/smoothers_out.h5", "r") do file
    read(file, "exp_states"), read(file, "exp_shocks")
end

@testset "Check expected output for all smoothers" begin
    for smoother in [:hamilton, :koopman, :carter_kohn, :durbin_koopman]
        @test exp_states ≈ states[smoother]
        @test exp_shocks ≈ shocks[smoother]
    end
end

# Make sure that simulation smoothers run with `draw_states` on
carter_kohn_smoother(y, T, R, C, Q, Z, D, E, z0, P0; draw_states = true)
durbin_koopman_smoother(y, T, R, C, Q, Z, D, E, z0, P0; draw_states = true)

nothing


text:
using GeometryTypes, StaticArrays, ModernGL
using Transpiler, Visualize

vertex_main(vertex) = vertex

function geometry_main(emit!, geom_in)
    pos = geom_in[1]
    quad = Vec4f0(pos[1], pos[2], pos[1] + 0.4f0, pos[2] + 0.4f0)
    v1 = quad[Vec(1, 2)]
    v2 = quad[Vec(1, 4)]
    v3 = quad[Vec(3, 2)]
    v4 = quad[Vec(3, 4)]
    emit!(Vec4f0(v1[1], v1[2], 0f0, 1f0), (Vec4f0(1, 0, 0, 1),))
    emit!(Vec4f0(v2[1], v2[2], 0f0, 1f0), (Vec4f0(1, 0, 0, 1),))
    emit!(Vec4f0(v3[1], v3[2], 0f0, 1f0), (Vec4f0(1, 0, 0, 1),))
    emit!(Vec4f0(v4[1], v4[2], 0f0, 1f0), (Vec4f0(1, 0, 0, 1),))
    return
end
fragment_main(fragment_in) = (fragment_in[1],)

N = 20
vertices = [(Point2f0(sin(2pi * (i / N)) , cos(2pi * (i / N))), ) for i = 1:N]
canvas = JLCanvas(Area => SimpleRectangle(0, 0, 500, 500))

draw = JLRasterizer(
    vertices, (),
    vertex_main, fragment_main,
    geometryshader = geometry_main
)
@time draw(canvas, vertices, ())

save("test.png", canvas[Visualize.ColorBuffer][1])


text:
# Handers for execute_request and related messages, which are
# the core of the Jupyter protocol: execution of Julia code and
# returning results.

import Base.Libc: flush_cstdio
import Pkg

# global variable so that display can be done in the correct Msg context
execute_msg = Msg(["julia"], Dict("username"=>"jlkernel", "session"=>uuid4()), Dict())
# global variable tracking the number of bytes written in the current execution
# request
const stdio_bytes = Ref(0)

import REPL: helpmode

# use a global array to accumulate "payloads" for the execute_reply message
const execute_payloads = Dict[]

function execute_request(socket, msg)
    code = msg.content["code"]
    @vprintln("EXECUTING ", code)
    global execute_msg = msg
    global n, In, Out, ans
    stdio_bytes[] = 0
    silent = msg.content["silent"]
    store_history = get(msg.content, "store_history", !silent)
    empty!(execute_payloads)

    if !silent
        n += 1
        send_ipython(publish[],
                     msg_pub(msg, "execute_input",
                             Dict("execution_count" => n,
                                          "code" => code)))
    end

    silent = silent || REPL.ends_with_semicolon(code)
    if store_history
        In[n] = code
    end

    # "; ..." cells are interpreted as shell commands for run
    code = replace(code, r"^\s*;.*$" =>
                   m -> string(replace(m, r"^\s*;" => "Base.repl_cmd(`"),
                               "`, stdout)"))

    # "] ..." cells are interpreted as pkg shell commands
    if occursin(r"^\].*$", code)
        code = "IJulia.Pkg.REPLMode.do_cmd(IJulia.minirepl[], \"" *
            escape_string(code[2:end]) * "\"; do_rethrow=true)"
    end

    # a cell beginning with "? ..." is interpreted as a help request
    hcode = replace(code, r"^\s*\?" => "")

    try
        for hook in preexecute_hooks
            invokelatest(hook)
        end

        ans = result = if hcode != code # help request
            Core.eval(Main, helpmode(hcode))
        else
            #run the code!
            occursin(magics_regex, code) ? magics_help(code) :
                SOFTSCOPE[] ? softscope_include_string(current_module[], code, "In[$n]") :
                include_string(current_module[], code, "In[$n]")
        end

        if silent
            result = nothing
        elseif result !== nothing
            if store_history
                Out[n] = result
            end
        end

        user_expressions = Dict()
        for (v,ex) in msg.content["user_expressions"]
            user_expressions[v] = invokelatest(parse, ex)
        end

        for hook in postexecute_hooks
            invokelatest(hook)
        end

        # flush pending stdio
        flush_all()

        undisplay(result) # dequeue if needed, since we display result in pyout
        invokelatest(display) # flush pending display requests

        if result !== nothing
            result_metadata = invokelatest(metadata, result)
            result_data = invokelatest(display_dict, result)
            send_ipython(publish[],
                         msg_pub(msg, "execute_result",
                                 Dict("execution_count" => n,
                                              "metadata" => result_metadata,
                                              "data" => result_data)))

        end
        send_ipython(requests[],
                     msg_reply(msg, "execute_reply",
                               Dict("status" => "ok",
                                            "payload" => execute_payloads,
                                            "execution_count" => n,
                                            "user_expressions" => user_expressions)))
        empty!(execute_payloads)
    catch e
        bt = catch_backtrace()
        try
            # flush pending stdio
            flush_all()
            for hook in posterror_hooks
                invokelatest(hook)
            end
        catch
        end
        empty!(displayqueue) # discard pending display requests on an error
        content = error_content(e,bt)
        send_ipython(publish[], msg_pub(msg, "error", content))
        content["status"] = "error"
        content["execution_count"] = n
        send_ipython(requests[], msg_reply(msg, "execute_reply", content))
    end
end


text:
module GridFunctions

using ADIOS2
using Random
using StaticArrays
using Test
using openPMD

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

include("defs.jl")

include("domain.jl")
include("bbox.jl")
include("bboxset.jl")
include("bboxarray.jl")

include("abstract.jl")
include("mesh.jl")
include("component.jl")

include("gridarray.jl")

end


text:
using StrayCopses
using CSV
using BenchmarkTools

# Glass
d_train = CSV.read(joinpath(@__DIR__, "data", "uci_glass_train.csv"))
d_test = CSV.read(joinpath(@__DIR__, "data", "uci_glass_test.csv"))

X_train = convert(Matrix, d_train[:, 1:size(d_train, 2) - 1])
y_train = convert(Vector, d_train[!, size(d_train, 2)])
X_test = convert(Matrix, d_test[:, 1:size(d_test, 2) - 1])
y_test = convert(Vector, d_test[!, size(d_test, 2)])

sc = StrayCopse()
@btime md = fit($sc, $X_train, $y_train, max_depth = 1_000_000)
# 202.178 ms (1846793 allocations: 210.01 MiB)
md = fit(sc, X_train, y_train, max_depth = 1_000_000)
# 174.363 ms (1637322 allocations: 187.40 MiB)
@btime phat = predict($md, $X_test)
# 111.559 μs (2 allocations: 336 bytes)
phat = predict(md, X_test)

begin
    md = fit(sc, X_train, y_train, max_depth = 1_000_000)
    phat = predict(md, X_test)
    sum(y_test .== phat)/length(phat)
end

# Sonar
d_train = CSV.read(joinpath(@__DIR__, "data", "uci_sonar_train.csv"))
d_test = CSV.read(joinpath(@__DIR__, "data", "uci_sonar_test.csv"))

X_train = convert(Matrix, d_train[:, 1:size(d_train, 2) - 1])
y_train = convert(Vector, d_train[!, size(d_train, 2)])
X_test = convert(Matrix, d_test[:, 1:size(d_test, 2) - 1])
y_test = convert(Vector, d_test[!, size(d_test, 2)])

sc = StrayCopse()
@btime md = fit($sc, $X_train, $y_train)
# 419.147 ms (3499117 allocations: 375.54 MiB)
@btime phat = predict($md, $X_test)
# 111.559 μs (2 allocations: 336 bytes)

sum(y_test .== phat)/length(phat)

size(X_train)
md = fit(sc, X_train, y_train; max_depth = 1000000, n_trees = 500, n_features_per_node = 5)
phat = predict(md, X_test)

sum(y_test .== phat)/length(phat)

phat = predict(md, X_train)
sum(y_train .== phat)/length(phat)

r = 1:100
md = fit(sc, X_train[r, :], y_train[r]; max_depth = 10000, n_trees = 1, n_features_per_node = 60)
phat = predict(md, X_train[r, :])
sum(y_train[r] .== phat)/length(phat)


text:
# [
#  [ MPIStateArray Name, Field Name, Maximum, Minimum, Mean, Standard Deviation ],
#  [         :                :          :        :      :          :           ],
# ]

#! format: off
first_order = (
[
 [ "state",     "ρ",   9.88641470218911577739e-01,  1.00735927139044068035e+00,  1.00000000000000244249e+00,  1.32618819574356725140e-03 ],
 [ "state", "ρu[1]",  -2.42850976605026747102e-01,  5.73364756714398238202e-01,  1.59153832831308655882e-01,  9.67668639083398146594e-02 ],
 [ "state", "ρu[2]",  -4.12933704690272018745e-01,  4.50172304099278386413e-01, -4.55511721449193612529e-14,  1.48850741581273815495e-01 ],
 [ "state", "ρu[3]",  -3.70076433121673098459e-01,  3.26009261713208653433e-01, -6.46953234297288930041e-14,  8.65065044705036617634e-02 ],
 [ "state",    "ρθ",  -1.71139241110591289186e+00,  1.79353274817685659492e+00, -7.54540463032282362914e-16,  3.87893456476833764501e-01 ],
],
[
 ["state",     "ρ", 12, 12, 12, 12],
 ["state", "ρu[1]", 11, 11, 12, 12],
 ["state", "ρu[2]", 12, 11,  0, 12],
 ["state", "ρu[3]", 12, 11,  0, 12],
 ["state",    "ρθ", 11, 10,  0, 12],
],
)

fourth_order = (
[
 [ "state",     "ρ",   9.79756567265035793746e-01,  1.00986882291859325633e+00,  9.99972252557182250676e-01,  1.27738464370501453651e-03 ],
 [ "state", "ρu[1]",  -3.85746589538976891731e-01,  7.12308328222086784010e-01,  1.58403896799433618892e-01,  1.03266735858658170732e-01 ],
 [ "state", "ρu[2]",  -5.86525244955121705104e-01,  6.33258737050533038193e-01,  2.62244534853688972004e-04,  1.37491363204364197559e-01 ],
 [ "state", "ρu[3]",  -4.58727280284793481613e-01,  4.82564240905019259387e-01, -2.81636420608788126414e-04,  8.90227260571099660025e-02 ],
 [ "state",    "ρθ",  -4.51696345879803118351e+00,  4.14217705122501378412e+00, -5.08780751729477078403e-04,  3.32659156438805059253e-01 ],
],
[
 ["state",     "ρ", 8, 8, 10, 7],
 ["state", "ρu[1]", 5, 5,  8, 7],
 ["state", "ρu[2]", 6, 5,  0, 6],
 ["state", "ρu[3]", 6, 5,  0, 7],
 ["state",    "ρθ", 5, 4,  0, 6],
],
)

#! format: on

refVals = (; first_order, fourth_order)


text:
@testset "MixMatRep" begin
    change_fulldim_test(MixedMatHRep)
    change_fulldim_test(MixedMatVRep)

    isempty_vtest(vrep(zeros(0, 3)), 0, 0)
    isempty_vtest(vrep(zeros(1, 2)), 0, 1)
    isempty_vtest(vrep(zeros(1, 4), ones(2, 4)), 2, 1)
    isempty_vtest(vrep(zeros(2, 1), ones(1, 1)), 1, 2)

    isempty_htest(hrep(ones(0, 3), zeros(0)), 0, 0)
    isempty_htest(hrep(ones(1, 2), zeros(1)), 0, 1)
    isempty_htest(hrep(ones(2, 4), zeros(2), BitSet(1:2)), 2, 0)
    isempty_htest(hrep(ones(3, 1), zeros(3), BitSet([2])), 1, 2)

    @testset "with bad arguments" begin
        A = [1 1; -1 0; 0 -1]
        b = [1, 0, 0]
        linset = BitSet([1])

        @test_throws ErrorException hrep(A, [0, 0], linset)
        @test_throws ErrorException hrep(A, b, BitSet([4]))
        ine = hrep(A, b, linset)
        @test fulldim(ine) == 2
        @test (@inferred Polyhedra.FullDim(ine)) == 2
        @test Polyhedra.coefficient_type(ine) == Int
        @test translate(ine, [1, 0]).b == [2, -1, 0]

        V = [0 1; 1 0]
        @test_throws ErrorException vrep(zeros(0, 2), [1 0]) # V-consistency
        @test_throws ErrorException vrep(V, [1 0 0], BitSet())
        @test_throws ErrorException vrep(V, [1 1], BitSet([2]))
        ext = vrep(V)
        @test fulldim(ext) == 2
        @test (@inferred Polyhedra.FullDim(ine)) == 2
        @test Polyhedra.coefficient_type(ext) == Int
        @test translate(ext, [1, 0]).V == [1 1; 2 0]
    end

    @testset "Cartesian product" begin
        A = [1 2; 3 4; 5 6]
        a = [7, 8, 9]
        B = [10 11 12; 13 14 15]
        b = [16, 17]
        p1 = hrep(A, a, BitSet([2]))
        p2 = hrep(B, b, BitSet([1]))
        p = p1 * p2
        @test p.A == [A[2,:]' zeros(1, 3)
                      zeros(1, 2) B[1, :]'
                      A[[1,3],:] zeros(2, 3)
                      zeros(1, 2) B[2, :]']
        @test p.b == [a[2]; b[1]; a[[1,3]]; b[2]]
        @test p.linset == BitSet([1, 2])
    end
end


text:
################
# Exponentiation
################

"""
    Base.:^(a::AbstractUncertainValue, b::AbstractUncertainValue) -> UncertainValue

Exponentiation operator for pairs of uncertain values.

The operator first computes a default of `n = 10000` realizations of both `a` and `b`,
which gives `10000` base-exponent pairs. Then, it exponentiates element-wise over these
base-exponent pairs, and finally returns an uncertain value based on a kernel density
estimate to the distribution of numbers resulting from the exponentiation operations.

Use the `^(a, b, n)` syntax to tune the number (`n`) of draws.
"""
function Base.:^(a::AbstractUncertainValue, b::AbstractUncertainValue; n::Int = 30000)
    UncertainValue(resample(a, n) .^ resample(b, n))
end

"""
    Base.:^(a::Real, b::AbstractUncertainValue) -> UncertainValue

Exponentiation operator for between scalars and uncertain values.

The operator first computes a default of `n = 10000` realizations of `b`, which are used
as exponents. Then, it makes `10000` base-exponent pairs, where the base is always `a` and
the exponents are the exponents generated from `b`. Next, it exponentiates element-wise
over these base-exponent pairs, and finally returns an uncertain value based on a kernel
density estimate to the distribution of numbers resulting from the exponentiation operations.

Use the `^(a, b, n)` syntax to tune the number (`n`) of draws.
"""
Base.:^(a::Real, b::AbstractUncertainValue; n::Int = 30000) =
    UncertainValue(a .^ resample(b, n))

"""
    Base.:^(a::AbstractUncertainValue, b::Real) -> UncertainValue

Exponentiation operator for between uncertain values and scalars.

The operator first computes a default of `n = 10000` realizations of `a`, which are used
as exponents. Then, it makes `10000` base-exponent pairs, where the base is always `b` and
the exponents are the exponents generated from `a`. Next, it exponentiates element-wise
over these base-exponent pairs, and finally returns an uncertain value based on a kernel
density estimate to the distribution of numbers resulting from the exponentiation operations.

Use the `^(a, b, n)` syntax to tune the number (`n`) of draws.
"""
Base.:^(a::AbstractUncertainValue, b::Real; n::Int = 30000) =
    UncertainValue(resample(a, n) .^ b)

"""
    Base.:^(a::AbstractUncertainValue, b::AbstractUncertainValue, n::Int) -> UncertainValue

Exponentiation operator for pairs of uncertain values.

The operator first computes `n` realizations of both `a` and `b`,
which gives `n` base-exponent pairs. Then, it exponentiates element-wise over these
base-exponent pairs, and finally returns an uncertain value based on a kernel density
estimate to the distribution of numbers resulting from the exponentiation operations.

Call this function using the `^(a, b, n)` syntax.
"""
function Base.:^(a::AbstractUncertainValue, b::AbstractUncertainValue, n::Int)
    UncertainValue(resample(a, n) .^ resample(b, n))
end

"""
    Base.:^(a::Real, b::AbstractUncertainValue, n::Int) -> UncertainValue

Exponentiation operator for between scalars and uncertain values.

The operator first computes `n = 10000` realizations of `b`, which are used
as exponents. Then, it makes `n` base-exponent pairs, where the base is always `a` and
the exponents are the exponents generated from `b`. Next, it exponentiates element-wise
over these base-exponent pairs, and finally returns an uncertain value based on a kernel
density estimate to the distribution of numbers resulting from the exponentiation operations.

Call this function using the `^(a, b, n)` syntax.
"""
Base.:^(a::Real, b::AbstractUncertainValue, n::Int) =
    UncertainValue(a .^ resample(b, n))

"""
    Base.:^(a::AbstractUncertainValue, b::Real, n::Int) -> UncertainValue

Exponentiation operator for between scalars and uncertain values.

The operator first computes `n = 10000` realizations of `a`, which are used
as exponents. Then, it makes `n` base-exponent pairs, where the base is always `b` and
the exponents are the exponents generated from `a`. Next, it exponentiates element-wise
over these base-exponent pairs, and finally returns an uncertain value based on a kernel
density estimate to the distribution of numbers resulting from the exponentiation operations.

Call this function using the `^(a, b, n)` syntax.
"""
Base.:^(a::AbstractUncertainValue, b::Real, n::Int) =
    UncertainValue(resample(a, n) .^ b)

#####################################################################################
# Special cases
#####################################################################################

import ..UncertainValues: CertainValue

##################
# `CertainValue`s
#################
"""
    Base.:^(a::Union{CertainValue, Real}, b::Union{CertainValue, Real})

Exponentiation of certain values with themselves or scalars acts as regular exponentiation,
but returns the result wrapped in a `CertainValue` instance.
"""
Base.:^(a::Union{CertainValue, Real}, b::Union{CertainValue, Real})

Base.:^(a::CertainValue, b::CertainValue) = CertainValue(a.value ^ b.value)
Base.:^(a::CertainValue, b::Real) = CertainValue(a.value ^ b)
Base.:^(a::Real, b::CertainValue) = CertainValue(a ^ b.value)


text:
function stable_principal_component_pursuit(M, sigma;
    threshold = 1e-6, solver_output=0)
    """
    This function solves the stable principal component pursuit problem.

    :param M: An arbitrary n-by-n matrix.
    :param sigma: A parameter that controls the tradeoff between nuclear norm
                  term and the L1 norm term in the objective function
    :param threshold: Singular values in the output low rank matrix below
                      threshold are set to 0. Entries in the output sparse
                      matrix below threshold are set to 0 (Float64).
    :param solver_output: Solver_output param to be passed to SCS (Int64).

    :return: This function returns two values. The first value is a tuple of 2
             n-by-n arrays that correspond to the solution of stable PCP (the
             first element in the tuple is the matrix X and the second element
             is the matrix Y). The second value is the optimal objective value
             of the optimization problem (Float64).
    """
    n = size(M)[1]
    mu = (2 * n) ^ 0.5 * sigma

    # Build stable PCP formulation using JuMP
    m = Model(SCS.Optimizer)
    set_optimizer_attribute(m, "verbose", solver_output)

    @variable(m, L[i=1:n, j=1:n])
    @variable(m, W_1[i=1:n, j=1:n])
    @variable(m, W_2[i=1:n, j=1:n])
    @variable(m, S[i=1:n, j=1:n])
    @variable(m, S_abs[i=1:n, j=1:n])
    @variable(m, error[i=1:n, j=1:n])

    @constraint(m, [W_1 L; L' W_2] in PSDCone())
    @constraint(m, [i=1:n, j=1:n], S_abs[i, j] >= S[i, j])
    @constraint(m, [i=1:n, j=1:n], S_abs[i, j] >= -S[i, j])

    @constraint(m, [i=1:n, j=1:n], (M[i, j] - L[i, j] - S[i, j])^2 <= error[i, j])

    @objective(m, Min, 0.5 * sum(W_1[i, i] for i=1:n) +
        0.5 * sum(W_2[i, i] for i=1:n) +
        n^(-0.5) * sum(S_abs[i, j] for i=1:n, j=1:n) +
        sum(error[i, j] for i=1:n, j=1:n)/(2*mu))

    # Solve stable PCP
    optimize!(m)

    L_opt = value.(L)
    S_opt = value.(S)

    L_svd = svd(L_opt)
    L_U = L_svd.U
    L_S = L_svd.S
    L_Vt = L_svd.Vt

    for i = 1:n
        if abs(L_S[i]) < threshold
            L_S[i] = 0
        end
    end

    L_opt = L_U * Diagonal(L_S) * L_Vt

    for i = 1:n
        for j = 1:n
            if abs(S_opt[i, j]) < threshold
                S_opt[i, j] = 0
            end
        end
    end

    return (L_opt, S_opt), objective_value(m)

end;


text:
using naumer_Dimensionality_2022
using Random: MersenneTwister
using LinearAlgebra: I, tr, norm, svd
using CSV

dim_lst = 1:50
μ = 1.0
σ² = 1.0
τ = 0.05
stepcount = 1000
T = 10

stepsize = 1e4

crlb_trace = zeros(length(dim_lst))

Threads.@threads for i in 1:length(dim_lst)
    global crlb_trace

    # Initialize system
    system = VanDerPolSystem_expanded(μ,dim_lst[i])
    crlb = I(dimension(system))

    actionseed = MersenneTwister(1234)

    state = ones(dimension(system))

    # Run 1D approximate optimal sampling procedure
    for j in 1:stepcount

        _jacobian = flowJacobian(state, T, system)
        ~, ~, V = svd(_jacobian)
        limitvector = Array(V[:,1])

        # Decide action
        action = limitvector + .01 * randn(length(limitvector))
        action ./= norm(action)
        for step in 1:1000
            action .= actiongradientstep_1DApprox(action, limitvector, crlb, σ², stepsize*dimension(system) / step^(2.0/3.0))
        end

        # Update System
        jacobian = flowJacobian(state, τ, system)
        crlb = updateCRLB(crlb, action, jacobian, σ²)
        state = flow(state, τ, system)
    end

    # Now record final crlb trace
    crlb_trace[i] = tr(crlb)
end

# Now random smapling for comparison
random_crlb_trace = zeros(length(crlb_trace))
for i in 1:length(dim_lst)
    @info(i)
    global random_crlb_trace

    # Initialize system
    system = VanDerPolSystem_expanded(μ,dim_lst[i])
    SystemMat = -1*I(dim_lst[i] + 1)
    SystemMat[1,1] = 0
    system = LinearSystem(SystemMat)
    crlb = I(dimension(system))

    state = 2*ones(dimension(system))

    # Run 1D approximate optimal sampling procedure
    for j in 1:stepcount
        # Decide action
        action = randn(dimension(system))
        action ./= norm(action)

        # Update System
        jacobian = flowJacobian(state, τ, system)
        crlb = updateCRLB(crlb, action, jacobian, σ²)
        state = flow(state, τ, system)
    end

    # Now record final crlb trace
    random_crlb_trace[i] = tr(crlb)
end

println(random_crlb_trace ./ crlb_trace)

CSV.write("DimensionScaling.csv", ( crlb_trace = crlb_trace,
                                    random_crlb_trace = random_crlb_trace))


text:
"""
    rank_to_sortperm(::NTuple{N,Int}) -> NTuple{N,Int}
Returns the `sortperm` of the stride ranks.
"""
function rank_to_sortperm(R::NTuple{N,Int}) where {N}
  sp = ntuple(zero, Val{N}())
  r = ntuple(n -> sum(R[n] .≥ R), Val{N}())
  @inbounds for n = 1:N
    sp = Base.setindex(sp, n, r[n])
  end
  sp
end
rank_to_sortperm(R) = sortperm(R)

function view_quote(i, K, S, D, T, N, C, B, R, X, O, zero_offsets::Bool = false)
  @assert ((K == N) || isone(K))

  inds = Expr(:tuple)
  Nnew = 0
  s = Expr(:tuple)
  x = Expr(:tuple)
  o = Expr(:tuple)
  Rnew = Expr(:tuple)
  Dnew = Expr(:tuple)
  Cnew = -1
  Bnew = ifelse(iszero(B), 0, -1)
  sortp = rank_to_sortperm(R)
  still_dense = true
  densev = Vector{Bool}(undef, K)
  for k ∈ 1:K
    iₖ = Expr(:ref, :i, k)
    if i[k] === Colon
      Nnew += 1
      push!(inds.args, Expr(:ref, :o, k))
      push!(s.args, Expr(:ref, :s, k))
      push!(x.args, Expr(:ref, :x, k))
      push!(o.args, zero_offsets ? :(Zero()) : :(One()))
      if k == C
        Cnew = Nnew
      end
      if k == B
        Bnew = Nnew
      end
      push!(Rnew.args, R[k])
    else
      push!(inds.args, Expr(:call, :first, iₖ))
      if i[k] <: AbstractRange
        Nnew += 1
        push!(s.args, Expr(:call, :static_length, iₖ))
        push!(x.args, Expr(:ref, :x, k))
        push!(o.args, zero_offsets ? :(Zero()) : :(One()))
        if k == C
          Cnew = Nnew
        end
        if k == B
          Bnew = Nnew
        end
        push!(Rnew.args, R[k])
      end
    end
    spₙ = sortp[k]
    still_dense &= D[spₙ]
    if still_dense# && (D[spₙ])
      ispₙ = i[spₙ]
      still_dense = (ispₙ <: AbstractUnitRange) || (ispₙ === Colon)
      densev[spₙ] = still_dense
      if still_dense
        still_dense = if ((ispₙ === Colon)::Bool || (ispₙ <: Base.Slice)::Bool)
          true
        else
          ispₙ_len = ArrayInterface.known_length(ispₙ)
          if ispₙ_len !== nothing
            _sz = getfield(S, :parameters)[spₙ]
            if _sz <: StaticInt
              ispₙ_len == getfield(_sz, :parameters)[1]
            else
              false
            end
          else
            false
          end
          false
        end
      end
    else
      densev[spₙ] = false
    end
  end
  for k ∈ 1:K
    iₖt = i[k]
    if (iₖt === Colon) || (iₖt <: AbstractVector)
      push!(Dnew.args, densev[k])
    end
  end
  quote
    $(Expr(:meta,:inline))
    sp = stridedpointer(A)
    s = size(A)
    x = strides(sp)
    o = offsets(sp)
    si = StrideIndex{$Nnew,$Rnew,$Cnew}($x, $o)
    new_sp = stridedpointer(_offset_ptr(sp, $inds), si, StaticInt{$Bnew}())
    PtrArray(new_sp, $s, Val{$Dnew}())
  end
end

@generated function Base.view(A::PtrArray{S,D,T,N,C,B,R,X,O}, i::Vararg{Union{Integer,AbstractRange,Colon},K}) where {K,S,D,T,N,C,B,R,X,O}
  view_quote(i, K, S, D, T, N, C, B, R, X, O)
end
@generated function zview(A::PtrArray{S,D,T,N,C,B,R,X,O}, i::Vararg{Union{Integer,AbstractRange,Colon},K}) where {K,S,D,T,N,C,B,R,X,O}
  view_quote(i, K, S, D, T, N, C, B, R, X, O, true)
end
@inline Base.SubArray(A::AbstractStrideArray, i::Tuple{Vararg{Union{Integer,AbstractRange,Colon},K}}) where {K} = view(A, i...)

@inline Base.vec(A::PtrArray{S,D,T,1,C,0}) where {S,D,T,C} = A
@inline function Base.vec(A::PtrArray{S,D,T,N,C,0}) where {S,D,T,N,C}
  @assert all(D) "All dimensions must be dense for a vec view. Try `vec(copy(A))` instead."
  si = StrideIndex{1,(1,),1}((static_sizeof(T),), (One(),))
  sp = stridedpointer(pointer(A), si)
  PtrArray(sp, (static_length(A),), Val((true,)))
end
@inline Base.vec(A::StaticStrideArray{S,D,T,1}) where {S,D,T} = A
@inline Base.vec(A::StaticStrideArray{S,D,T,N}) where {S,D,T,N} = StrideArray(vec(PtrArray(A)), A)


text:
### A Pluto.jl notebook ###
# v0.12.21

using Markdown
using InteractiveUtils

# ╔═╡ 1585e592-fc0e-11ea-09b3-85bd83e3893e
using Pkg, DrWatson

# ╔═╡ 15862766-fc0e-11ea-160b-09995fb8e8fc
begin
	#@quickactivate "StatisticalRethinkingStan"
	using StanQuap
	using StatisticalRethinking
end

# ╔═╡ 6fdbe2ac-fc0d-11ea-0253-5fcd59eb669f
md"## Fig5.1s.jl"

# ╔═╡ 33c69ba4-01af-11eb-337e-19ecbe1fe80b
begin
	df = CSV.read(sr_datadir("WaffleDivorce.csv"), DataFrame)

	# Use WaffleHouses/Populations = Waffle houses per million

	df.WaffleHouses = df.WaffleHouses ./ df.Population
	scale!(df, [:WaffleHouses, :Divorce])
end;

# ╔═╡ db17d86a-fc0e-11ea-2df9-716728b916e6
	PRECIS(df[:, 3:15])

# ╔═╡ 15937fa4-fc0e-11ea-0aaf-e7049b6392bf
stan5_0 = "
data {
 int < lower = 1 > N; // Sample size
 vector[N] D; // Outcome (Divorce rate)
 vector[N] W; // Predictor ()
}

parameters {
 real a; // Intercept
 real bA; // Slope (regression coefficients)
 real < lower = 0 > sigma;    // Error SD
}

model {
  vector[N] mu;               // mu is a vector
  a ~ normal(0, 0.2);         //Priors
  bA ~ normal(0, 0.5);
  sigma ~ exponential(1);
  mu = a + bA * W;
  D ~ normal(mu , sigma);     // Likelihood
}
";

# ╔═╡ 15a36108-fc0e-11ea-310b-d50e8725f62c
md"##### Input data."

# ╔═╡ 15afbdfe-fc0e-11ea-15a5-9bf6506abf5d
begin
	data = (N = size(df, 1), D = df.Divorce_s, W = df.WaffleHouses_s)
	init = (a=0.0, bA=1.0, sigma=1.0)
end;

# ╔═╡ 15b07488-fc0e-11ea-2bdd-558c18e223f5
md"##### Sample using StanSample."

# ╔═╡ 15bce25e-fc0e-11ea-3c0d-a761765fd79f
begin
	q5_0s, m5_0s, o5_0s = stan_quap("m5.0s", stan5_0; data, init)
	if !isnothing(m5_0s)
		post5_0s_df = read_samples(m5_0s, :dataframe)
		part5_0s = Particles(post5_0s_df)
	end
	PRECIS(post5_0s_df)
end

# ╔═╡ 15ca02e2-fc0e-11ea-281e-194a9b83623d
if !isnothing(q5_0s)
	quap5_0s_df = sample(q5_0s)
	PRECIS(quap5_0s_df)
end

# ╔═╡ 3e6621a0-fc42-11ea-0790-6fdef3820273
md"##### Figure out which annotations we want."

# ╔═╡ 53b10c8c-fc3d-11ea-0cfa-f9a16db84264
df[[1, 4, 11, 20, 30, 40], [1, 2, 7, 9]]

# ╔═╡ da595b90-fc3d-11ea-252e-1d4c44f1545f
df[:,1]

# ╔═╡ 15d19672-fc0e-11ea-093b-bff8eaf427d2
if !isnothing(m5_0s)
	begin
		p2 = plotbounds(
			df, :WaffleHouses, :Divorce,
			post5_0s_df, [:a, :bA, :sigma];
			bounds=[:predicted, :hpdi],
			title="Divorce rate vs. waffle houses per million" * "\nshowing predicted and hpd range",
			xlab="Waffle houses per million",
			ylab="Divorce rate"
		)
		for i in [1, 4, 11, 20, 30, 40]
			annotate!(p2, [(df[i, :WaffleHouses]+01.5, df[i, :Divorce], Plots.text(df[i, 2],
				6, :red, :right))])
		end
	end
	plot(p2)
end

# ╔═╡ 15df0b0e-fc0e-11ea-0ca7-459b93ac79d5
md"## End of Fig5.1s.jl"

# ╔═╡ Cell order:
# ╟─6fdbe2ac-fc0d-11ea-0253-5fcd59eb669f
# ╠═1585e592-fc0e-11ea-09b3-85bd83e3893e
# ╠═15862766-fc0e-11ea-160b-09995fb8e8fc
# ╠═33c69ba4-01af-11eb-337e-19ecbe1fe80b
# ╠═db17d86a-fc0e-11ea-2df9-716728b916e6
# ╠═15937fa4-fc0e-11ea-0aaf-e7049b6392bf
# ╟─15a36108-fc0e-11ea-310b-d50e8725f62c
# ╠═15afbdfe-fc0e-11ea-15a5-9bf6506abf5d
# ╟─15b07488-fc0e-11ea-2bdd-558c18e223f5
# ╠═15bce25e-fc0e-11ea-3c0d-a761765fd79f
# ╠═15ca02e2-fc0e-11ea-281e-194a9b83623d
# ╟─3e6621a0-fc42-11ea-0790-6fdef3820273
# ╠═53b10c8c-fc3d-11ea-0cfa-f9a16db84264
# ╠═da595b90-fc3d-11ea-252e-1d4c44f1545f
# ╠═15d19672-fc0e-11ea-093b-bff8eaf427d2
# ╟─15df0b0e-fc0e-11ea-0ca7-459b93ac79d5