text:
"""
Handles WebSockets communication logic.
"""
module WebChannels

import HTTP, Distributed, Logging, JSON3
import Genie, Genie.Renderer

const ClientId = UInt # web socket hash
const ChannelName = String

struct ChannelNotFoundException <: Exception
  name::ChannelName
end

mutable struct ChannelClient
  client::HTTP.WebSockets.WebSocket
  channels::Vector{ChannelName}
end

const ChannelClientsCollection = Dict{ClientId,ChannelClient} # { id(ws) => { :client => ws, :channels => ["foo", "bar", "baz"] } }
const ChannelSubscriptionsCollection = Dict{ChannelName,Vector{ClientId}}  # { "foo" => ["4", "12"] }
const MessagePayload = Union{Nothing,Dict}

mutable struct ChannelMessage
  channel::ChannelName
  client::ClientId
  message::String
  payload::MessagePayload
end

function JSON3.StructTypes.StructType(::Type{T}) where {T<:ChannelMessage}
  JSON3.StructTypes.Struct()
end

const CLIENTS = ChannelClientsCollection()
const SUBSCRIPTIONS = ChannelSubscriptionsCollection()

clients() = collect(values(CLIENTS))
subscriptions() = SUBSCRIPTIONS
websockets() = map(c -> c.client, clients())
channels() = collect(keys(SUBSCRIPTIONS))

function connected_clients(channel::ChannelName) :: Vector{ChannelClient}
  clients = ChannelClient[]
  for client_id in SUBSCRIPTIONS[channel]
    ! (CLIENTS[client_id].client.txclosed && CLIENTS[client_id].client.rxclosed) && push!(clients, CLIENTS[client_id])
  end

  clients
end
function connected_clients() :: Vector{ChannelClient}
  clients = ChannelClient[]
  for ch in channels()
    clients = vcat(clients, connected_clients(ch))
  end

  clients
end

function disconnected_clients(channel::ChannelName) :: Vector{ChannelClient}
  clients = ChannelClient[]
  for client_id in SUBSCRIPTIONS[channel]
    (CLIENTS[client_id].client.txclosed || CLIENTS[client_id].client.rxclosed) && push!(clients, CLIENTS[client_id])
  end

  clients
end
function disconnected_clients() :: Vector{ChannelClient}
  clients = ChannelClient[]
  for ch in channels()
    clients = vcat(clients, disconnected_clients(ch))
  end

  clients
end

"""
Subscribes a web socket client `ws` to `channel`.
"""
function subscribe(ws::HTTP.WebSockets.WebSocket, channel::ChannelName) :: ChannelClientsCollection
  if haskey(CLIENTS, id(ws))
    in(channel, CLIENTS[id(ws)].channels) || push!(CLIENTS[id(ws)].channels, channel)
  else
    CLIENTS[id(ws)] = ChannelClient(ws, ChannelName[channel])
  end

  push_subscription(id(ws), channel)

  CLIENTS
end

function id(ws::HTTP.WebSockets.WebSocket) :: UInt
  hash(ws)
end

"""
Unsubscribes a web socket client `ws` from `channel`.
"""
function unsubscribe(ws::HTTP.WebSockets.WebSocket, channel::ChannelName) :: ChannelClientsCollection
  haskey(CLIENTS, id(ws)) && deleteat!(CLIENTS[id(ws)].channels, CLIENTS[id(ws)].channels .== channel)
  pop_subscription(id(ws), channel)

  CLIENTS
end
function unsubscribe(channel_client::ChannelClient, channel::ChannelName) :: ChannelClientsCollection
  unsubscribe(channel_client.client, channel)
end

"""
Unsubscribes a web socket client `ws` from all the channels.
"""
function unsubscribe_client(ws::HTTP.WebSockets.WebSocket) :: ChannelClientsCollection
  if haskey(CLIENTS, id(ws))
    for channel_id in CLIENTS[id(ws)].channels
      pop_subscription(id(ws), channel_id)
    end

    delete!(CLIENTS, id(ws))
  end

  CLIENTS
end
function unsubscribe_client(client_id::ClientId) :: ChannelClientsCollection
  unsubscribe_client(CLIENTS[client_id].client)

  CLIENTS
end
function unsubscribe_client(channel_client::ChannelClient) :: ChannelClientsCollection
  unsubscribe_client(channel_client.client)

  CLIENTS
end

"""
unsubscribe_disconnected_clients() :: ChannelClientsCollection

Unsubscribes clients which are no longer connected.
"""
function unsubscribe_disconnected_clients() :: ChannelClientsCollection
  for channel_client in disconnected_clients()
    unsubscribe_client(channel_client)
  end

  CLIENTS
end
function unsubscribe_disconnected_clients(channel::ChannelName) :: ChannelClientsCollection
  for channel_client in disconnected_clients(channel)
    unsubscribe(channel_client, channel)
  end

  CLIENTS
end

"""
Adds a new subscription for `client` to `channel`.
"""
function push_subscription(client_id::ClientId, channel::ChannelName) :: ChannelSubscriptionsCollection
  if haskey(SUBSCRIPTIONS, channel)
    ! in(client_id, SUBSCRIPTIONS[channel]) && push!(SUBSCRIPTIONS[channel], client_id)
  else
    SUBSCRIPTIONS[channel] = ClientId[client_id]
  end

  SUBSCRIPTIONS
end
function push_subscription(channel_client::ChannelClient, channel::ChannelName) :: ChannelSubscriptionsCollection
  push_subscription(id(channel_client.client), channel)
end

"""
Removes the subscription of `client` to `channel`.
"""
function pop_subscription(client::ClientId, channel::ChannelName) :: ChannelSubscriptionsCollection
  if haskey(SUBSCRIPTIONS, channel)
    filter!(SUBSCRIPTIONS[channel]) do (client_id)
      client_id != client
    end
    isempty(SUBSCRIPTIONS[channel]) && delete!(SUBSCRIPTIONS, channel)
  end

  SUBSCRIPTIONS
end
function pop_subscription(channel_client::ChannelClient, channel::ChannelName) :: ChannelSubscriptionsCollection
  pop_subscription(id(channel_client.client), channel)
end

"""
Removes all subscriptions of `client`.
"""
function pop_subscription(channel::ChannelName) :: ChannelSubscriptionsCollection
  if haskey(SUBSCRIPTIONS, channel)
    delete!(SUBSCRIPTIONS, channel)
  end

  SUBSCRIPTIONS
end

"""
Pushes `msg` (and `payload`) to all the clients subscribed to the channels in `channels`, with the exception of `except`.
"""
function broadcast(channels::Union{ChannelName,Vector{ChannelName}},
                    msg::String,
                    payload::Union{Dict,Nothing} = nothing;
                    except::Union{HTTP.WebSockets.WebSocket,Nothing,UInt} = nothing) :: Bool
  isa(channels, Array) || (channels = ChannelName[channels])

  if isempty(SUBSCRIPTIONS)
    @warn("No clients are subscribed to any channel.")
    return false
  end

  for channel in channels
    haskey(SUBSCRIPTIONS, channel) || throw(ChannelNotFoundException(channel))

    for client in SUBSCRIPTIONS[channel]
      except !== nothing && client == id(except) && continue
      CLIENTS[client].client.txclosed && CLIENTS[client].client.rxclosed && continue

      try
        payload !== nothing ?
          message(client, ChannelMessage(channel, client, msg, payload) |> Renderer.Json.JSONParser.json) :
          message(client, msg)
      catch ex
        @error ex
      end
    end
  end

  true
end

"""
Pushes `msg` (and `payload`) to all the clients subscribed to the channels in `channels`, with the exception of `except`.
"""
function broadcast(msg::String;
                    channels::Union{Union{ChannelName,Vector{ChannelName}},Nothing} = nothing,
                    payload::Union{Dict,Nothing} = nothing,
                    except::Union{HTTP.WebSockets.WebSocket,Nothing,UInt} = nothing) :: Bool
  channels === nothing && (channels = collect(keys(SUBSCRIPTIONS)))
  broadcast(channels, msg, payload; except = except)
end

"""
Pushes `js_code` (a JavaScript piece of code) to be executed by all the clients subscribed to the channels in `channels`,
with the exception of `except`.
"""
function jscomm(js_code::String, channels::Union{Union{ChannelName,Vector{ChannelName}},Nothing} = nothing;
            except::Union{HTTP.WebSockets.WebSocket,Nothing,UInt} = nothing)
  broadcast(string(Genie.config.webchannels_eval_command, js_code); channels = channels, except = except)
end

"""
Writes `msg` to web socket for `client`.
"""
function message(ws::HTTP.WebSockets.WebSocket, msg::String) :: Int
  write(ws, msg)
end
function message(client::ClientId, msg::String) :: Int
  message(CLIENTS[client].client, msg)
end
function message(client::ChannelClient, msg::String) :: Int
  message(client.client, msg)
end

end


text:
using StatsBase
using LinearAlgebra
using Images
using ImageSegmentation
export get_normal_image,
    get_normal_SVD,
    rgb2lab,
    rgb_plane2rgb_world,
    compute_local_planes,
    resize

function  get_projection_mask()
    mask = fill(false, 480, 640);
    mask[45:471, 41:601] .= true;
    sz = [427 561]
    return mask,sz
end
"""
 Projects the given depth image to world coordinates. Note that this 3D
 coordinate space is defined by a horizontal plane made from the X and Z
 axes and the Y axis points up.

 Args:
   imgDepthAbs - 480x640 depth image whose values indicate depth in
                meters.

 Returns:
   points3d - Nx3 matrix of 3D world points (X,Y,Z).

"""
function depth_plane2depth_world(imgDepthAbs, c_d, f_d)
    (H, W) = size(imgDepthAbs);
    X = (collect(1:W)' .- c_d[1]) .* imgDepthAbs / f_d[1];
    Y = (collect(1:H) .- c_d[2]) .* imgDepthAbs / f_d[2];
    Z = imgDepthAbs
    return cat(X,Y,Z, dims=3)
end
"""
```julia
function compute_local_planes(X, Y, Z, params)
```

 Computes surface normal information. Note that in this file, the Y
 coordinate points up, consistent with the image coordinate frame.

 Args:
   X - Nx1 column vector of 3D point cloud X-coordinates
   Y - Nx1 column vector of 3D point cloud Y-coordinates
   Z - Nx1 column vector of 3D point cloud Z-coordinates.
   params

 Returns:
   - `imgPlanes`. An 'image' of the plane parameters for each pixel.
   - `imgNormals`. HxWx3 matrix of surface normals at each pixel.
   - `imgConfs`. HxW image of confidences.
"""
function compute_local_planes(XYZ;window::Integer=3, relDepthThresh::Float64 = 0.8)
    ( a , sz) = get_projection_mask()
    #H = sz[1];
    #W = sz[2];
	H = size(XYZ,1)
    W = size(XYZ,2)
    N = H * W;

    imgPlanes = zeros(Float32,H,W,4)
    imgConfs  = zeros(Float32,H,W)
	imgNormals = zeros(Float32,H,W,3)
    for r=1:H, c=1:W
        wc = max(c-window,1):min(c+window,W)
        wr = max(r-window,1):min(r+window,H)
        idxs = [(rwr, cwc)  for rwr in wr, cwc in wc
                 if abs(XYZ[rwr,cwc,3] - XYZ[r, c, 3]) < XYZ[r,c,3] * relDepthThresh ]
        if length(idxs) < 3
	    continue;
	end
	A = ones(4,length(idxs))
	j = 1
	for (rwr, rwc) in idxs
	    A[1:3,j] = XYZ[rwr,rwc,:]
	    j=j+1
	end
    	(l, eigv) = eigen(A*A');
	v = eigv[:,1]
	len = norm(v[1:3])
	imgPlanes[r,c,:]= v ./ (len+eps())
	imgNormals[r,c,:]=imgPlanes[r,c,1:3]
	imgConfs[r,c] = 1 - sqrt(max(l[1],0) / l[2]);
    end
    for r=1:H, c=1:W
        if norm(imgNormals[r,c,:] - [0.0;0;0]) < 0.000001
            wc = max(c-window,1):min(c+window,W)
	    wr = max(r-window,1):min(r+window,H)
            v = zeros(3)
            for pos_wr in wr, pos_wc in wc
                v=v+imgNormals[pos_wr,pos_wc,:]
            end
            v=v/norm(v)
            imgNormals[r,c,:]=v
        end
    end
    return (imgPlanes, imgNormals, imgConfs)
end

function clean_normal_image(img)
    (z,nr,nc) = size(img)
    for c = 1:nc, c=1:nr
        if  any(.!(isfinite.(img[:,r,c])))
            wr = max(r-3,1):min(r+3,nr)
            wc = max(c-3,1):min(c+3,nc)
            imgw = img[:,wr,wc]
            imgw_b = sum(!isfinite.(imgw),1)
            cc = reshape(imgw_b,length(wr),length(wc))
            bb =     find(cc.== 0)
            (fs,cs) = ind2sub(size(cc),bb)
            v = zeros(3)
            for (i,j) in zip(fs,cs)
                v=v+imgw[:,i,j]
            end
            v=v/length(fs)
            v=v/norm(v)
            img[:,r,c]=v
        end
    end
    return img
end
"""
```
function fill_normal(N::Array{T,3}; w::Integer=2) where T
```
Fill the invalid values of a normal image with the mean direction of its vicinity of radius `w`
"""
function fill_normal!(N::Array{T,3}; w::Integer=2) where T
    nanN = isnan.(N)
    nanN = nanN[:,:,1] .| nanN[:,:,2] .| nanN[:,:,3]
    invalid_r, invalid_c = ind2sub(size(nanN),find(nanN))
    for j=1:length(invalid_r)
        r=invalid_r[j]
        c=invalid_c[j]
        wr = max(r-w,1):min(r+w,size(nanN,1))
        wc = max(c-w,1):min(c+w,size(nanN,2))
        n  = zeros(3)
        for nposr=1:wr, nposc=1:wc
            if all(.!(isnan.(N[nposr,nposc,:])))
                n  = n + N[nposr,nposc,:]
            end
        end
        n = n /norm(n)
        N[r,c,:] = vec(n)
    end
end
function get_normal_SVD(R, w::Integer, th::Float64)
    nr = size(R,1)
    nc = size(R,2)
    N  = zeros(Float32,nr,nc,3)
    for c = 1:nc
        for r=1:nr
            wc = max(c-w,1):min(c+w,nc)
            wr = max(r-w,1):min(r+w,nr)
            indexes    = Base.CartesianRange((wr,wc))
            x = Float32[]
            y = Float32[]
            z = Float32[]
            for V in indexes
                (wwr, wwc)   = (V[1], V[2])
                if abs(R[r,c,3]-R[wwr,wwc,3]) <= th
                    push!(x,R[wwr,wwc,1])
                    push!(y,R[wwr,wwc,2])
                    push!(z,R[wwr,wwc,3])
                end
            end
            if (length(x) > 3)
                push!(x,R[r,c,1])
                push!(y,R[r,c,2])
                push!(z,R[r,c,3])

                M = hcat(x,y,z)'
                M = M .- mean(M,2)

                (U, A2, V) = svd(M)
                N[r,c,:]= U[:,3] / norm(U[:,3])
            end
        end
    end
    return N
end
function get_normal_image(R, ttype::String; d::Integer = 3)
    w  = params["d"]
    if ttype=="SVD"
    end
end

#="""

"""
function rgb2lab(RGB_image::Array{T,3}) where T
	(nr, nc, c) = size(RGB_image)
	res         = zeros(Float16, (nr,nc,3))
	for i in CartesianRange((nr,nc))
		res[i[1],i[2],:] = rgb2lab(RGB_image[i[1],i[2],:])
	end
	return res
end=#
"""
```
function rgb2lab(RGB_image::Array{T,3}) where T
```
% RGB2Lab takes matrices corresponding to Red, Green, and Blue, and
% transforms them into CIELab.  This transform is based on ITU-R
% Recommendation  BT.709 using the D65 white point reference.
% The error in transforming RGB -> Lab -> RGB is approximately
% 10^-5.  RGB values can be either between 0 and 1 or between 0 and 255.
% By Mark Ruzon from C code by Yossi Rubner, 23 September 1997.
% Updated for MATLAB 5 28 January 1998.
"""
function  rgb2lab(RGB_image::Array{ElType,3}) where ElType

    B = float(RGB_image[:,:,3]);
    G = float(RGB_image[:,:,2]);
    R = float(RGB_image[:,:,1]);
    if ((maximum(R) > 1.0) | (maximum(G) > 1.0) | (maximum(B) > 1.0))
        R = R./255.0;
        G = G./255.0;
        B = B./255.0;
    end
    (M, N) = size(R);
    s = M*N;
    #Set a threshold
    T = 0.008856;

    RGB = hcat(reshape(R,s), reshape(G,s), reshape(B,s))';
    # RGB to XYZ
    MAT = [0.412453 0.357580 0.180423;
           0.212671 0.715160 0.072169;
           0.019334 0.119193 0.950227];
    XYZ = MAT * RGB;

    X = XYZ[1,:] / 0.950456;
    Y = XYZ[2,:]
    Z = XYZ[3,:] / 1.088754;

    XT = X .> T;
    YT = Y .> T;
    ZT = Z .> T;

    fX = XT .* X.^(1/3) + (.!XT) .* (7.787 .* X + 16/116);

    # Compute L
    Y3 = Y.^(1/3);
    fY = YT .* Y3 + (.!YT) .* (7.787 .* Y + 16/116);
    L  = YT .* (116 * Y3 - 16.0) + (.!YT) .* (903.3 * Y);

    fZ = ZT .* Z.^(1/3) + (.!ZT) .* (7.787 .* Z + 16/116);

    #    Compute a and b
    a =     500 * (fX - fY);
    b = 200 * (fY - fZ);

    L = reshape(L, M, N);
    a = reshape(a, M, N);
    b = reshape(b, M, N);

    return L = cat(3,L,a,b);
end
"""

"""
function rgb2lab(input_color::Vector)
    RGB = zeros(3)
    for num=1:3
        value = convert(Float64,input_color[num]) / 255.0
        if value > 0.04045
            value = ((value + 0.055) / 1.055) ^ 2.4
        else
            value = value / 12.92
		end
        RGB[num] = value * 100
    end
    XYZ = zeros(3)
    XYZ[1] = RGB[1] * 0.4124 + RGB[2] * 0.3576 + RGB[3] * 0.1805
    XYZ[2] = RGB[1] * 0.2126 + RGB[2] * 0.7152 + RGB[3] * 0.0722
    XYZ[3] = RGB[1] * 0.0193 + RGB[2] * 0.1192 + RGB[3] * 0.9505

    # Observer= 2°, Illuminant= D65
    XYZ[1] = XYZ[1] / 95.047         # ref_X =  95.047
    XYZ[2] = XYZ[2] / 100.0          # ref_Y = 100.000
    XYZ[3] = XYZ[3] / 108.883        # ref_Z = 108.883

    num = 0
    for num=1:3
	value = XYZ[num]
        if value > 0.008856
            value = value ^ (0.3333333333333333)
        else
            value = (7.787 * value) + (16 / 116.0)
	end
        XYZ[num] = value
    end
    Lab = zeros(3)

    L = (116 * XYZ[2]) - 16
    a = 500 * (XYZ[1] - XYZ[2])
    b = 200 * (XYZ[2] - XYZ[3])

    Lab[1] = L
    Lab[2] = a
    Lab[3] = b
    return Lab
end
function rgb2lab!(img::CDNImage{T}) where T
    for r=1:size(img,2), c=1:size(img,3)
        img[1:3,r,c] = rgb2lab(img[1:3,r,c])
    end
end
function resize(labels::Matrix{T}, scale::Float64) where T<:Integer
    dim = (round(Integer, size(labels, 1) * scale),
           round(Integer, size(labels, 2) * scale))
    return round.(Integer, imresize(labels, dim))
end
function resize(img::CDNImage, scale::Float64)
    D = distances(img)
    dim = (round(Integer, size(D, 2) * scale),
           round(Integer, size(D, 3) * scale))
    RGB = Array(channelview(imresize(colors(img), dim)))
    D = imresize(distances(img), (3, dim...))
    N = imresize(normals(img), (3, dim...))
    return CDNImage(RGB,D,N)
end
function resize(img::LabelledCDNImage, scale::Float64)
    resized_img = resize(img.image, scale)
    resized_labels = resize(labels_map(img.ground_truth), scale)
    return LabelledCDNImage(
        resized_img,
        resized_labels)
end

ext: jl


text:
module Inference

using ..Core, ..Core.RandomVariables, ..Utilities
using ..Core.RandomVariables: Metadata, _tail, TypedVarInfo,
    islinked, invlink!, getlogp, tonamedtuple
using Distributions, Libtask, Bijectors
using ProgressMeter, LinearAlgebra
using ..Turing: PROGRESS, CACHERESET, AbstractSampler
using ..Turing: Model, runmodel!, get_pvars, get_dvars,
    Sampler, SampleFromPrior, SampleFromUniform,
    Selector, AbstractSamplerState
using ..Turing: in_pvars, in_dvars, Turing
using StatsFuns: logsumexp
using Random: GLOBAL_RNG, AbstractRNG
using ..Turing.Interface

import MCMCChains: Chains
import AdvancedHMC; const AHMC = AdvancedHMC
import ..Turing: getspace
import Distributions: sample
import ..Core: getchunksize, getADtype
import ..Interface: AbstractTransition, sample, step!, sample_init!,
    transitions_init, sample_end!, AbstractSampler, transition_type,
    callback, init_callback, AbstractCallback, psample

export  InferenceAlgorithm,
        Hamiltonian,
        AbstractGibbs,
        GibbsComponent,
        StaticHamiltonian,
        AdaptiveHamiltonian,
        SampleFromUniform,
        SampleFromPrior,
        MH,
        Gibbs,      # classic sampling
        HMC,
        SGLD,
        SGHMC,
        HMCDA,
        NUTS,       # Hamiltonian-like sampling
        DynamicNUTS,
        IS,
        SMC,
        CSMC,
        PG,
        PIMH,
        PMMH,
        IPMCMC,  # particle-based sampling
        assume,
        observe,
        step,
        WelfordVar,
        WelfordCovar,
        NaiveCovar,
        get_var,
        get_covar,
        add_sample!,
        reset!,
        step!

#######################
# Sampler abstraction #
#######################
abstract type AbstractAdapter end
abstract type InferenceAlgorithm end
abstract type ParticleInference <: InferenceAlgorithm end
abstract type Hamiltonian{AD} <: InferenceAlgorithm end
abstract type StaticHamiltonian{AD} <: Hamiltonian{AD} end
abstract type AdaptiveHamiltonian{AD} <: Hamiltonian{AD} end

getchunksize(::T) where {T <: Hamiltonian} = getchunksize(T)
getchunksize(::Type{<:Hamiltonian{AD}}) where AD = getchunksize(AD)
getADtype(alg::Hamiltonian) = getADtype(typeof(alg))
getADtype(::Type{<:Hamiltonian{AD}}) where {AD} = AD

"""
    mh_accept(H::T, H_new::T, log_proposal_ratio::T) where {T<:Real}

Peform MH accept criteria with log acceptance ratio. Returns a `Bool` for acceptance.

Note: This function is only used in PMMH.
"""
function mh_accept(H::T, H_new::T, log_proposal_ratio::T) where {T<:Real}
    return log(rand()) + H_new < H + log_proposal_ratio, min(0, -(H_new - H))
end

######################
# Default Transition #
######################

struct Transition{T, F<:AbstractFloat} <: AbstractTransition
    θ  :: T
    lp :: F
end

function Transition(spl::Sampler, nt::NamedTuple=NamedTuple())
    theta = merge(tonamedtuple(spl.state.vi), nt)
    lp = getlogp(spl.state.vi)
    return Transition{typeof(theta), typeof(lp)}(theta, lp)
end

function additional_parameters(::Type{<:Transition})
    return [:lp]
end

Interface.transition_type(::Sampler{alg}) where alg = transition_type(alg)

##########################################
# Internal variable names for MCMCChains #
##########################################

const TURING_INTERNAL_VARS = (internals = [
    "elapsed",
    "eval_num",
    "lf_eps",
    "lp",
    "weight",
    "le",
    "acceptance_rate",
    "hamiltonian_energy",
    "hamiltonian_energy_error",
    "max_hamiltonian_energy_error",
    "is_accept",
    "log_density",
    "n_steps",
    "numerical_error",
    "step_size",
    "nom_step_size",
    "tree_depth",
    "is_adapt",
],)

#########################################
# Default definitions for the interface #
#########################################

function Interface.sample(
    rng::AbstractRNG,
    model::ModelType,
    alg::AlgType,
    N::Integer;
    kwargs...
) where {
    ModelType<:Sampleable,
    SamplerType<:AbstractSampler,
    AlgType<:InferenceAlgorithm
}
    return sample(rng, model, Sampler(alg, model), N; progress=PROGRESS[], kwargs...)
end

function Interface.sample(
    model::ModelType,
    alg::AlgType,
    N::Integer;
    resume_from=nothing,
    kwargs...
) where {
    ModelType<:Sampleable,
    SamplerType<:AbstractSampler,
    AlgType<:InferenceAlgorithm
}
    if resume_from === nothing
        return sample(model, Sampler(alg, model), N; progress=PROGRESS[], kwargs...)
    else
        return resume(resume_from, N)
    end
end

function Interface.psample(
    model::ModelType,
    alg::AlgType,
    N::Integer,
    n_chains::Integer;
    kwargs...
) where {
    ModelType<:Sampleable,
    SamplerType<:AbstractSampler,
    AlgType<:InferenceAlgorithm
}
    return psample(GLOBAL_RNG, model, alg, N, n_chains; progress=false, kwargs...)
end

function Interface.psample(
    rng::AbstractRNG,
    model::ModelType,
    alg::AlgType,
    N::Integer,
    n_chains::Integer;
    kwargs...
) where {
    ModelType<:Sampleable,
    SamplerType<:AbstractSampler,
    AlgType<:InferenceAlgorithm
}
    return psample(rng, model, Sampler(alg, model), N, n_chains; progress=false, kwargs...)
end

function Interface.sample_init!(
    ::AbstractRNG,
    model::Model,
    spl::Sampler,
    N::Integer;
    kwargs...
)
    # Resume the sampler.
    set_resume!(spl; kwargs...)

    # Set the parameters to a starting value.
    initialize_parameters!(spl; kwargs...)
end

function Interface.sample_end!(
    ::AbstractRNG,
    ::Model,
    spl::AbstractSampler,
    ::Integer,
    ::Vector{TransitionType};
    kwargs...
) where {TransitionType<:AbstractTransition}
    # Silence the default API function.
end

function initialize_parameters!(
    spl::AbstractSampler;
    init_theta::Union{Nothing,Array{<:Any,1}}=nothing,
    verbose::Bool=false,
    kwargs...
)
    # Get `init_theta`
    if init_theta !== nothing
        verbose && @info "Using passed-in initial variable values" init_theta
        # Convert individual numbers to length 1 vector; `ismissing(v)` is needed as `size(missing)` is undefined`
        init_theta = [ismissing(v) || size(v) == () ? [v] : v for v in init_theta]
        # Flatten `init_theta`
        init_theta_flat = foldl(vcat, map(vec, init_theta))
        # Create a mask to inidicate which values are not missing
        theta_mask = map(x -> !ismissing(x), init_theta_flat)
        # Get all values
        theta = spl.state.vi[spl]
        @assert length(theta) == length(init_theta_flat) "Provided initial value doesn't match the dimension of the model"
        # Update those which are provided (i.e. not missing)
        theta[theta_mask] .= init_theta_flat[theta_mask]
        # Update in `vi`
        spl.state.vi[spl] = theta
    end
end

##########################
# Chain making utilities #
##########################

function _params_to_array(ts::Vector{T}, spl::Sampler) where {T<:AbstractTransition}
    names = Set{String}()
    dicts = Vector{Dict{String, Any}}()

    # Extract the parameter names and values from each transition.
    for t in ts
        nms, vs = flatten_namedtuple(t.θ)
        push!(names, nms...)

        # Convert the names and values to a single dictionary.
        d = Dict{String, Any}()
        for (k, v) in zip(nms, vs)
            d[k] = v
        end
        push!(dicts, d)
    end

    # Convert the set to an ordered vector so the parameter ordering
    # is deterministic.
    ordered_names = collect(names)
    vals = Matrix{Union{Real, Missing}}(undef, length(ts), length(ordered_names))

    # Place each element of all dicts into the returned value matrix.
    for i in eachindex(dicts)
        for (j, key) in enumerate(ordered_names)
            vals[i,j] = get(dicts[i], key, missing)
        end
    end

    return ordered_names, vals
end

function flatten_namedtuple(nt::NamedTuple{pnames}) where {pnames}
    vals = Vector{Real}()
    names = Vector{AbstractString}()

    for k in pnames
        v = nt[k]
        if length(v) == 1
            flatten(names, vals, string(k), v)
        else
            for (vnval, vn) in zip(v[1], v[2])
                flatten(names, vals, vn, vnval)
            end
        end
    end

    return names, vals
end

function flatten(names, value :: AbstractArray, k :: String, v)
    if isa(v, Number)
        name = k
        push!(value, v)
        push!(names, name)
    elseif isa(v, Array)
        for i = eachindex(v)
            if isa(v[i], Number)
                name = string(ind2sub(size(v), i))
                name = replace(name, "(" => "[");
                name = replace(name, ",)" => "]");
                name = replace(name, ")" => "]");
                name = k * name
                isa(v[i], Nothing) && println(v, i, v[i])
                push!(value, v[i])
                push!(names, name)
            elseif isa(v[i], AbstractArray)
                name = k * string(ind2sub(size(v), i))
                flatten(names, value, name, v[i])
            else
                error("Unknown var type: typeof($v[i])=$(typeof(v[i]))")
            end
        end
    else
        error("Unknown var type: typeof($v)=$(typeof(v))")
    end
    return
end

ind2sub(v, i) = Tuple(CartesianIndices(v)[i])

function get_transition_extras(ts::Vector{T}) where T<:AbstractTransition
    # Get the extra field names from the sampler state type.
    # This handles things like :lp or :weight.
    extra_params = additional_parameters(T)

    # Get the values of the extra parameters.
    local extra_names
    all_vals = []

    # Iterate through each transition.
    for t in ts
        extra_names = String[]
        vals = []

        # Iterate through each of the additional field names
        # in the struct.
        for p in extra_params
            # Check whether the field contains a NamedTuple,
            # in which case we need to iterate through each
            # key/value pair.
            prop = getproperty(t, p)
            if prop isa NamedTuple
                for (k, v) in pairs(prop)
                    push!(extra_names, string(k))
                    push!(vals, v)
                end
            else
                push!(extra_names, string(p))
                push!(vals, prop)
            end
        end
        push!(all_vals, vals)
    end

    # Convert the vector-of-vectors to a matrix.
    valmat = [all_vals[i][j] for i in 1:length(ts), j in 1:length(all_vals[1])]

    return extra_names, valmat
end

# Default Chains constructor.
function Chains(
    rng::AbstractRNG,
    model::ModelType,
    spl::Sampler,
    N::Integer,
    ts::Vector{T};
    discard_adapt::Bool=true,
    save_state=false,
    kwargs...
) where {ModelType<:Sampleable, T<:AbstractTransition}
    # Check if we have adaptation samples.
    if discard_adapt && :n_adapts in fieldnames(typeof(spl.alg))
        ts = ts[(spl.alg.n_adapts+1):end]
    end

    # Convert transitions to array format.
    # Also retrieve the variable names.
    nms, vals = _params_to_array(ts, spl)

    # Get the values of the extra parameters in each Transition struct.
    extra_params, extra_values = get_transition_extras(ts)

    # Extract names & construct param array.
    nms = string.(vcat(nms..., string.(extra_params)...))
    parray = hcat(vals, extra_values)

    # If the state field has average_logevidence or final_logevidence, grab that.
    le = missing
    if :average_logevidence in fieldnames(typeof(spl.state))
        le = getproperty(spl.state, :average_logevidence)
    elseif :final_logevidence in fieldnames(typeof(spl.state))
        le = getproperty(spl.state, :final_logevidence)
    end

    # Check whether to invlink! the varinfo
    if islinked(spl.state.vi, spl)
        invlink!(spl.state.vi, spl)
    end

    # Set up the info tuple.
    info = if save_state
        (range = rng,
        model = model,
        spl = spl)
    else
        NamedTuple()
    end

    # Chain construction.
    return Chains(
        parray,
        string.(nms),
        deepcopy(TURING_INTERNAL_VARS);
        evidence=le,
        info=info
    )
end

function save(c::Chains, spl::AbstractSampler, model, vi, samples)
    nt = NamedTuple{(:spl, :model, :vi, :samples)}((spl, model, deepcopy(vi), samples))
    return setinfo(c, merge(nt, c.info))
end

function resume(c::Chains, n_iter::Int; kwargs...)
    @assert !isempty(c.info) "[Turing] cannot resume from a chain without state info"
    return sample(
        c.info[:range],
        c.info[:model],
        c.info[:spl],
        n_iter;    # this is actually not used
        resume_from=c,
        reuse_spl_n=n_iter,
        kwargs...
    )
end

function set_resume!(
    s::Sampler;
    resume_from::Union{Chains, Nothing}=nothing,
    kwargs...
)
    # If we're resuming, grab the sampler info.
    if resume_from !== nothing
        s = resume_from.info[:spl]
    end
end

function split_var_str(var_str)
    ind = findfirst(c -> c == '[', var_str)
    inds = Vector{String}[]
    if ind === nothing
        return var_str, inds
    end
    sym = var_str[1:ind-1]
    ind = length(sym)
    while ind < length(var_str)
        ind += 1
        @assert var_str[ind] == '['
        push!(inds, String[])
        while var_str[ind] != ']'
            ind += 1
            if var_str[ind] == '['
                ind2 = findnext(c -> c == ']', var_str, ind)
                push!(inds[end], strip(var_str[ind:ind2]))
                ind = ind2+1
            else
                ind2 = findnext(c -> c == ',' || c == ']', var_str, ind)
                push!(inds[end], strip(var_str[ind:ind2-1]))
                ind = ind2
            end
        end
    end
    return sym, inds
end

#########################
# Default sampler state #
#########################

"""
A blank `AbstractSamplerState` that contains only `VarInfo` information.
"""
mutable struct SamplerState{VIType<:VarInfo} <: AbstractSamplerState
    vi :: VIType
end

#######################################
# Concrete algorithm implementations. #
#######################################

include("hmc.jl")
include("mh.jl")
include("is.jl")
include("AdvancedSMC.jl")
include("gibbs.jl")
include("../contrib/inference/sghmc.jl")
include("../contrib/inference/AdvancedSMCExtensions.jl")

################
# Typing tools #
################

for alg in (:SMC, :PG, :PMMH, :IPMCMC, :MH, :IS)
    @eval getspace(::$alg{space}) where {space} = space
    @eval getspace(::Type{<:$alg{space}}) where {space} = space
end
for alg in (:HMC, :HMCDA, :NUTS, :SGLD, :SGHMC)
    @eval getspace(::$alg{<:Any, space}) where {space} = space
    @eval getspace(::Type{<:$alg{<:Any, space}}) where {space} = space
end
getspace(::Gibbs) = Tuple{}()
getspace(::Type{<:Gibbs}) = Tuple{}()

@inline floatof(::Type{T}) where {T <: Real} = typeof(one(T)/one(T))
@inline floatof(::Type) = Real

@inline Turing.Core.get_matching_type(spl::Turing.Sampler, vi::Turing.RandomVariables.VarInfo, ::Type{T}) where {T <: AbstractFloat} = floatof(eltype(vi, spl))
@inline Turing.Core.get_matching_type(spl::Turing.Sampler{<:Hamiltonian}, vi::Turing.RandomVariables.VarInfo, ::Type{TV}) where {T, N, TV <: Array{T, N}} = Array{Turing.Core.get_matching_type(spl, vi, T), N}
@inline Turing.Core.get_matching_type(spl::Turing.Sampler{<:Union{PG, SMC}}, vi::Turing.RandomVariables.VarInfo, ::Type{TV}) where {T, N, TV <: Array{T, N}} = TArray{T, N}

## Fallback functions

# utility funcs for querying sampler information
require_gradient(spl::Sampler) = false
require_particles(spl::Sampler) = false

assume(spl::Sampler, dist::Distribution) =
error("Turing.assume: unmanaged inference algorithm: $(typeof(spl))")

observe(spl::Sampler, weight::Float64) =
error("Turing.observe: unmanaged inference algorithm: $(typeof(spl))")

## Default definitions for assume, observe, when sampler = nothing.
function assume(spl::A,
    dist::Distribution,
    vn::VarName,
    vi::VarInfo) where {A<:Union{SampleFromPrior, SampleFromUniform}}

    if haskey(vi, vn)
        r = vi[vn]
    else
        r = isa(spl, SampleFromUniform) ? init(dist) : rand(dist)
        push!(vi, vn, r, dist, spl)
    end
    # NOTE: The importance weight is not correctly computed here because
    #       r is genereated from some uniform distribution which is different from the prior
    # acclogp!(vi, logpdf_with_trans(dist, r, istrans(vi, vn)))

    r, logpdf_with_trans(dist, r, istrans(vi, vn))

end

function assume(spl::A,
    dists::Vector{T},
    vn::VarName,
    var::Any,
    vi::VarInfo) where {T<:Distribution, A<:Union{SampleFromPrior, SampleFromUniform}}

    @assert length(dists) == 1 "Turing.assume only support vectorizing i.i.d distribution"
    dist = dists[1]
    n = size(var)[end]

    vns = map(i -> VarName(vn, "[$i]"), 1:n)

    if haskey(vi, vns[1])
        rs = vi[vns]
    else
        rs = isa(spl, SampleFromUniform) ? init(dist, n) : rand(dist, n)

        if isa(dist, UnivariateDistribution) || isa(dist, MatrixDistribution)
            for i = 1:n
                push!(vi, vns[i], rs[i], dist, spl)
            end
            @assert size(var) == size(rs) "Turing.assume: variable and random number dimension unmatched"
            var = rs
        elseif isa(dist, MultivariateDistribution)
            for i = 1:n
                push!(vi, vns[i], rs[:,i], dist, spl)
            end
            if isa(var, Vector)
                @assert length(var) == size(rs)[2] "Turing.assume: variable and random number dimension unmatched"
                for i = 1:n
                    var[i] = rs[:,i]
                end
            elseif isa(var, Matrix)
                @assert size(var) == size(rs) "Turing.assume: variable and random number dimension unmatched"
                var = rs
            else
                @error("Turing.assume: unsupported variable container"); error()
            end
        end
    end

    # acclogp!(vi, sum(logpdf_with_trans(dist, rs, istrans(vi, vns[1]))))

    var, sum(logpdf_with_trans(dist, rs, istrans(vi, vns[1])))

end

observe(::Nothing,
        dist::T,
        value::Any,
        vi::VarInfo) where T = observe(SampleFromPrior(), dist, value, vi)

function observe(spl::A,
    dist::Distribution,
    value::Any,
    vi::VarInfo) where {A<:Union{SampleFromPrior, SampleFromUniform}}

    vi.num_produce += one(vi.num_produce)
    Turing.DEBUG && @debug "dist = $dist"
    Turing.DEBUG && @debug "value = $value"

    # acclogp!(vi, logpdf(dist, value))
    logpdf(dist, value)
end

function observe(spl::A,
    dists::Vector{T},
    value::Any,
    vi::VarInfo) where {T<:Distribution, A<:Union{SampleFromPrior, SampleFromUniform}}

    @assert length(dists) == 1 "Turing.observe only support vectorizing i.i.d distribution"
    dist = dists[1]
    @assert isa(dist, UnivariateDistribution) ||
        isa(dist, MultivariateDistribution) "Turing.observe: vectorizing matrix distribution is not supported"
    if isa(dist, UnivariateDistribution)  # only univariate distributions support broadcast operation (logpdf.) by Distributions.jl
        # acclogp!(vi, sum(logpdf.(Ref(dist), value)))
        sum(logpdf.(Ref(dist), value))
    else
        # acclogp!(vi, sum(logpdf(dist, value)))
        sum(logpdf(dist, value))
    end

end

##############
# Utilities  #
##############

getspace(spl::Sampler) = getspace(typeof(spl))
getspace(::Type{<:Sampler{Talg}}) where {Talg} = getspace(Talg)

end # module


text:
# ------------ Constrain anywhere ------------ #

struct ConstrainedAnywhereSelection{T <: ConstrainedSelectQuery} <: ConstrainedSelection
    query::T
    ConstrainedAnywhereSelection(obs::Vector{Tuple{T, K}}) where {T <: Any, K} = new{ConstrainedByAddress}(ConstrainedByAddress(Dict{Address, Any}(obs)))
    ConstrainedAnywhereSelection(obs::Tuple{T, K}...) where {T <: Any, K} = new{ConstrainedByAddress}(ConstrainedByAddress(Dict{Address, Any}(collect(obs))))
end

has_top(cas::ConstrainedAnywhereSelection, addr::T) where T <: Address = has_top(cas.query, addr)
has_top(cas::ConstrainedAnywhereSelection, addr::T) where T <: Tuple = has_top(cas.query, addr[end])

dump_queries(cas::ConstrainedAnywhereSelection) = dump_queries(cas.query)

get_top(cas::ConstrainedAnywhereSelection, addr::T) where T <: Address = get_top(cas.query, addr)
get_top(cas::ConstrainedAnywhereSelection, addr::T) where T <: Tuple = get_top(cas.query, addr[end])

get_sub(cas::ConstrainedAnywhereSelection, addr) = cas

isempty(cas::ConstrainedAnywhereSelection) = isempty(cas.query)

# ------------ Pretty printing ------------ #

function Base.display(chs::ConstrainedAnywhereSelection)
    println("  __________________________________\n")
    println("              Selection\n")
    addrs = Union{Symbol, Pair}[]
    chd = Dict{Address, Any}()
    collect!(addrs, chd, chs.query)
    for a in addrs
        println(" (Anywhere)   $(a) : $(chd[a])")
    end
    println("  __________________________________\n")
end

# ------------ Unconstrained select anywhere ------------ #

struct UnconstrainedAnywhereSelection{T <: UnconstrainedSelectQuery} <: UnconstrainedSelection
    query::T
    UnconstrainedAnywhereSelection(obs::Vector{Tuple{T, K}}) where {T <: Any, K} = new{UnconstrainedByAddress}(UnconstrainedByAddress(Dict{Address, Any}(obs)))
    UnconstrainedAnywhereSelection(obs::Tuple{T, K}...) where {T <: Any, K} = new{UnconstrainedByAddress}(UnconstrainedByAddress(Dict{Address, Any}(collect(obs))))
end

has_top(cas::UnconstrainedAnywhereSelection, addr::T) where T <: Address = has_top(cas.query, addr)
has_top(cas::UnconstrainedAnywhereSelection, addr::T) where T <: Tuple = has_top(cas.query, addr[end])
dump_queries(cas::UnconstrainedAnywhereSelection) = dump_queries(cas.query)
get_sub(cas::UnconstrainedAnywhereSelection, addr) = cas
isempty(cas::UnconstrainedAnywhereSelection) = isempty(cas.query)

# ------------ Unconstrained select all ------------ #

struct UnconstrainedAllSelection <: UnconstrainedSelection end

has_top(uas::UnconstrainedAllSelection, addr) = true
get_sub(uas::UnconstrainedAllSelection, addr) = uas
isempty(uas::UnconstrainedAllSelection) = false


text:
abstract type RegressionLikelihood{T<:Real} <: AbstractLikelihood{T} end

include("gaussian.jl")
include("studentt.jl")
include("laplace.jl")
include("heteroscedastic.jl")
include("matern.jl")

### Return the labels in a vector of vectors for multiple outputs
function treat_labels!(
    y::AbstractVector{T}, likelihood::L
) where {T,L<:RegressionLikelihood}
    @assert T <: Real "For regression target(s) should be real valued"
    return y, 1, likelihood
end

predict_y(::RegressionLikelihood, μ::AbstractVector{<:Real}) = μ
predict_y(::RegressionLikelihood, μ::AbstractVector{<:AbstractVector}) = first(μ)


text:
export extract_amp, extract_vis, extract_lcamp, extract_cphase

function getvisfield(obs)
    obsamps = obs.data::PyObject
    u = deepcopy((get(obsamps, Vector{Float64}, "u")))
    v = deepcopy((get(obsamps, Vector{Float64}, "v")))
    err = deepcopy((get(obsamps, Vector{Float64}, "sigma")))
    vis = deepcopy((get(obsamps, Vector{Complex{Float64}}, "vis")))
    t1 = Symbol.(deepcopy((get(obsamps, Vector{String}, "t1"))))
    t2 = Symbol.(deepcopy((get(obsamps, Vector{String}, "t2"))))
    baseline = tuple.(t1, t2)
    time = deepcopy((get(obsamps, Vector{Float64}, "time")))
    freq = zeros(length(time))
    bw = zeros(length(time))

    return  StructArray{Comrade.EHTVisibilityDatum{Float64}}(
        visr = real.(vis),
        visi = imag.(vis),
        u = u,
        v = v,
        error = err,
        time = time,
        frequency = freq,
        bandwidth = bw,
        baseline = baseline
    )
end

function getampfield(obs)
    obsamps = obs.amp::PyObject
    uamp = deepcopy(get(obsamps, Vector{Float64}, "u"))
    vamp = deepcopy(get(obsamps, Vector{Float64}, "v"))
    erramp = deepcopy(get(obsamps, Vector{Float64}, "sigma"))
    amps = deepcopy(get(obsamps, Vector{Float64}, "amp"))
    t1 = Symbol.(deepcopy(get(obsamps, Vector{String}, "t1")))
    t2 = Symbol.(deepcopy(get(obsamps, Vector{String}, "t2")))
    baseline = tuple.(t1, t2)
    time = deepcopy(get(obsamps, Vector{Float64}, "time"))
    freq = zeros(length(time))
    bw = zeros(length(time))

    return  StructArray{Comrade.EHTVisibilityAmplitudeDatum{Float64}}(
        amp = amps,
        u = uamp,
        v = vamp,
        error = erramp,
        time = time,
        frequency = freq,
        bandwidth = bw,
        baseline = baseline
    )
end

function getcpfield(obs)
    obscp = obs.cphase::PyObject
    u1 = deepcopy((get(obscp, Vector{Float64}, "u1")))
    v1 = deepcopy((get(obscp, Vector{Float64}, "v1")))
    u2 = deepcopy((get(obscp, Vector{Float64}, "u2")))
    v2 = deepcopy((get(obscp, Vector{Float64}, "v2")))
    u3 = deepcopy((get(obscp, Vector{Float64}, "u3")))
    v3 = deepcopy((get(obscp, Vector{Float64}, "v3")))
    cp = deg2rad.(deepcopy((get(obscp, Vector{Float64}, "cphase"))))
    errcp = deg2rad.(deepcopy((get(obscp, Vector{Float64}, "sigmacp"))))

    t1 = Symbol.(deepcopy((get(obscp, Vector{String}, "t1"))))
    t2 = Symbol.(deepcopy((get(obscp, Vector{String}, "t2"))))
    t3 = Symbol.(deepcopy((get(obscp, Vector{String}, "t3"))))
    baseline = tuple.(t1, t2, t3)
    time = deepcopy(get(obscp, Vector{Float64}, "time"))
    freq = zeros(length(time))
    bw = zeros(length(time))

    return StructArray{Comrade.EHTClosurePhaseDatum{Float64}}(
        phase = cp,
        u1 = u1,
        v1 = v1,
        u2 = u2,
        v2 = v2,
        u3 = u3,
        v3 = v3,
        error = errcp,
        time = time,
        frequency = freq,
        bandwidth = bw,
        triangle = baseline
    )

end

function getlcampfield(obs)
    obslcamp = obs.logcamp::PyObject
    u1 = get(obslcamp, Vector{Float64}, "u1")
    v1 = get(obslcamp, Vector{Float64}, "v1")
    u2 = get(obslcamp, Vector{Float64}, "u2")
    v2 = get(obslcamp, Vector{Float64}, "v2")
    u3 = get(obslcamp, Vector{Float64}, "u3")
    v3 = get(obslcamp, Vector{Float64}, "v3")
    u4 = get(obslcamp, Vector{Float64}, "u4")
    v4 = get(obslcamp, Vector{Float64}, "v4")
    camp = ((get(obslcamp, Vector{Float64}, "camp")))
    errcamp = ((get(obslcamp, Vector{Float64}, "sigmaca")))

    t1 = Symbol.(((get(obslcamp, Vector{String}, "t1"))))
    t2 = Symbol.(((get(obslcamp, Vector{String}, "t2"))))
    t3 = Symbol.(((get(obslcamp, Vector{String}, "t3"))))
    t4 = Symbol.(((get(obslcamp, Vector{String}, "t4"))))
    baseline = tuple.(t1, t2, t3, t4)
    time = (get(obslcamp, Vector{Float64}, "time"))
    freq = zeros(length(time))
    bw = zeros(length(time))

    return StructArray{Comrade.EHTLogClosureAmplitudeDatum{Float64}}(
        amp = camp,
        u1 = u1,
        v1 = v1,
        u2 = u2,
        v2 = v2,
        u3 = u3,
        v3 = v3,
        u4 = u4,
        v4 = v4,
        error = errcamp,
        time = time,
        frequency = freq,
        bandwidth = bw,
        quadrangle = baseline
    )

end

function getradec(obs)::Tuple{Float64, Float64}
    return (float(obs.ra), float(obs.dec))
end

function getmjd(obs)::Int
    return Int(obs.mjd)
end

function getsource(obs)::Symbol
    return Symbol(obs.source)
end

"""
    extract_amp(obs)
Extracts the visibility amplitudes from an ehtim observation object.

Any valid keyword arguments to `add_amp` in ehtim can be passed through extract_amp.

Returns an EHTObservation with visibility amplitude data
"""
function extract_amp(obsc; kwargs...)
    obs = obsc.copy()
    obs.reorder_tarr_snr()
    obs.add_amp(;kwargs...)
    data = getampfield(obs)
    ra, dec = getradec(obs)
    mjd = getmjd(obs)
    source = getsource(obs)
    bw = obs.bw
    rf = obs.rf
    ac = _arrayconfig(data, bw, rf)
    return Comrade.EHTObservation(data = data, mjd = mjd,
                   ra = ra, dec= dec,
                   config = ac,
                   bandwidth=bw, frequency=rf,
                   source = source,
    )
end

"""
    extract_vis(obs)
Extracts the complex visibilities from an ehtim observation object

This grabs the raw `data` object from the obs object. Any keyword arguments are ignored.

Returns an EHTObservation with complex visibility data
"""
function extract_vis(obsc; kwargs...)
    obs = obsc.copy()
    obs.reorder_tarr_snr()

    data = getvisfield(obs)
    ra, dec = getradec(obs)
    mjd = getmjd(obs)
    source = getsource(obs)
    bw = obs.bw
    rf = obs.rf
    ac = _arrayconfig(data, bw, rf)
    return Comrade.EHTObservation(
                   data = data, mjd = mjd,
                   config=ac,
                   ra = ra, dec= dec,
                   bandwidth=bw, frequency=rf,
                   source = source,
    )
end

function closure_designmat(scancp::Scan{A,B,C}, scanvis) where {A,B,C<:StructArray{<:EHTClosurePhaseDatum}}
    ant1, ant2, ant3 = baselines(scancp)
    antvis1, antvis2 = baselines(scanvis)
    design_mat = zeros(length(scancp), length(scanvis))
    #throw("here")
    # fill in each row of the design matrix
    for i in axes(design_mat, 2), j in axes(design_mat,1)
        a1 = ant1[j]
        a2 = ant2[j]
        a3 = ant3[j]
        # check leg 1
        ((antvis1[i] == a1) & (antvis2[i] == a2)) && (design_mat[j,i] = 1.0)
        ((antvis1[i] == a2) & (antvis2[i] == a1)) && (design_mat[j,i] = -1.0)

        #check leg 2
        ((antvis1[i] == a2) & (antvis2[i] == a3)) && (design_mat[j,i] = 1.0)
        ((antvis1[i] == a3) & (antvis2[i] == a2)) && (design_mat[j,i] = -1.0)

        #check leg 3
        ((antvis1[i] == a3) & (antvis2[i] == a1)) && (design_mat[j,i] = 1.0)
        ((antvis1[i] == a1) & (antvis2[i] == a3)) && (design_mat[j,i] = -1.0)
    end
    return design_mat
end

function closure_designmat(scanlca::Scan{A,B,C}, scanvis) where {A,B,C<:StructArray{<:EHTLogClosureAmplitudeDatum}}
    ant1, ant2, ant3, ant4 = baselines(scanlca)
    antvis1, antvis2 = baselines(scanvis)
    design_mat = zeros(length(scanlca), length(scanvis))
    # fill in each row of the design matrix
    for i in axes(design_mat, 2), j in axes(design_mat,1)
        a1 = ant1[j]
        a2 = ant2[j]
        a3 = ant3[j]
        a4 = ant4[j]

        av1 = antvis1[i]
        av2 = antvis2[i]
        # check leg 1
        (((av1 == a1)&(av2 == a2)) || (av1 == a2)&(av2 == a1))&& (design_mat[j,i] = 1.0)

        #check leg 2
        (((av1 == a3)&(av2 == a4)) || (av1 == a4)&(av2 == a3))&& (design_mat[j,i] = 1.0)

        #check leg 3
        (((av1 == a1)&(av2 == a4)) || (av1 == a4)&(av2 == a1))&& (design_mat[j,i] = -1.0)

        #check leg 4
        (((av1 == a2)&(av2 == a3)) || (av1 == a3)&(av2 == a2))&& (design_mat[j,i] = -1.0)

    end
    return design_mat
end

function minimal_lcamp(obsc; kwargs...)

    obs = obsc.copy()
    # reorder to maximize the snr
    obs.reorder_tarr_snr()

    lcamp = _ehtim_lcamp(obs; count="max", kwargs...)
    stlca = scantable(lcamp)

    #Now make the vis obs
    dvis = extract_vis(obsc)
    st = scantable(dvis)

    minset, dmat = _minimal_closure(stlca, st)

    # now create EHTObservation
    ra, dec = getradec(obs)
    mjd = getmjd(obs)
    source = getsource(obs)
    bw = obs.bw
    rf = obs.rf
    ac = arrayconfig(dvis)
    clac = ClosureConfig(ac, dmat)
    return EHTObservation(data = minset, mjd = mjd,
                          config=clac,
                          ra = ra, dec= dec,
                          bandwidth=bw, frequency=rf,
                          source = source,
                        )

end

function _ehtim_cphase(obsc; count="max", cut_trivial=false, uvmin=0.1e9, kwargs...)
    obs = obsc.copy()

    # cut 0 baselines since these are trivial triangles
    if cut_trivial
        obs = obs.flag_uvdist(uv_min=uvmin)
    end

    obs.reorder_tarr_snr()

    obs.add_cphase(;count=count, kwargs...)
    data = getcpfield(obs)
    ra, dec = getradec(obs)
    mjd = getmjd(obs)
    source = getsource(obs)
    bw = obs.bw
    rf = obs.rf

    cphase = EHTObservation(data = data, mjd = mjd,
                   config=nothing,
                   ra = ra, dec= dec,
                   bandwidth=bw, frequency=rf,
                   source = source,
    )

    stcp = scantable(cphase)

    #Now make the vis obs
    dvis = extract_vis(obsc)
    st = scantable(dvis)
    S = eltype(dvis[:visr])

    dmat = Matrix{S}[]
    for i in 1:length(st)
        scanvis = st[i]
        ind = findfirst(==(scanvis.time), stcp.times)
        if isnothing(ind)
            inow = length(dmat)
            # I want to construct block diagonal matrices for efficienty but we need
            # to be careful with scan that can't form closures. This hack moves these
            # scans into the preceding scan but fills it with zeros
            if i > 1
                dmat[inow] = hcat(dmat[inow], zeros(S, size(dmat[inow],1), length(scanvis.scan)))
            else
                push!(dmat, zeros(S, 1, length(scanvis.scan)))
            end
            continue
        end
        dmatscan = closure_designmat(stcp[ind], st[i])
        push!(dmat, dmatscan)
    end

    if iszero(dmat[1])
        dmat[2] = hcat(zeros(S, size(dmat[2],1), size(dmat[1],2)), dmat[2])
        dmat = dmat[2:end]
    end

    ac = arrayconfig(dvis)
    clac = ClosureConfig(ac, dmat)
    return  EHTObservation(data = data, mjd = mjd,
                           config=clac,
                           ra = ra, dec= dec,
                           bandwidth=bw, frequency=rf,
                           source = source
                          )
end

function _make_lcamp(obsc, count="max"; kwargs...)
    obs = obsc.copy()

    obs.reorder_tarr_snr()

    obs.add_logcamp(;count=count, kwargs...)
    data = getlcampfield(obs)
    ra, dec = getradec(obs)
    mjd = getmjd(obs)
    source = getsource(obs)
    bw = obs.bw
    rf = obs.rf

    return EHTObservation(data = data, mjd = mjd,
                   config=nothing,
                   ra = ra, dec= dec,
                   bandwidth=bw, frequency=rf,
                   source = source,
    )

end

function _ehtim_lcamp(obsc; count="max", kwargs...)

    lcamp = _make_lcamp(obsc; count, kwargs...)

    stlca = scantable(lcamp)

    #Now make the vis obs
    dvis = extract_vis(obsc)
    st = scantable(dvis)
    S = eltype(dvis[:visr])

    dmat = Matrix{S}[]
    for i in 1:length(st)
        scanvis = st[i]
        ind = findfirst(==(scanvis.time), stlca.times)
        if isnothing(ind)
            inow = length(dmat)
            # I want to construct block diagonal matrices for efficienty but we need
            # to be careful with scan that can't form closures. This hack moves these
            # scans into the preceding scan but fills it with zeros
            if i > 1
                dmat[inow] = hcat(dmat[inow], zeros(S, size(dmat[inow],1), length(scanvis.scan)))
            else
                push!(dmat, zeros(S, 1, length(scanvis.scan)))
            end
            continue
        end
        dmatscan = closure_designmat(stlca[ind], st[i])
        push!(dmat, dmatscan)
    end

    if iszero(dmat[1])
        dmat[2] = hcat(zeros(S, size(dmat[2],1), size(dmat[1],2)), dmat[2])
        dmat = dmat[2:end]
    end

    ac = arrayconfig(dvis)
    clac = ClosureConfig(ac, dmat)
    return  EHTObservation(data = lcamp.data, mjd = lcamp.mjd,
                           config=clac,
                           ra = lcamp.ra, dec= lcamp.dec,
                           bandwidth=lcamp.bandwidth, frequency=lcamp.frequency,
                           source = lcamp.source
                          )
end

function minimal_cphase(obsc; kwargs...)
    # compute a maximum set of closure phases
    obs = obsc.copy()

    # reorder to maximize the snr
    obs.reorder_tarr_snr()

    cphase = _ehtim_cphase(obs; count="max", kwargs...)
    stcp = scantable(cphase)

    #Now make the vis obs
    dvis = extract_vis(obsc)
    st = scantable(dvis)

    minset, dmat = _minimal_closure(stcp, st)
    # now create EHTObservation
    ra, dec = getradec(obs)
    mjd = getmjd(obs)
    source = getsource(obs)
    bw = obs.bw
    rf = obs.rf
    ac = arrayconfig(dvis)
    clac = ClosureConfig(ac, dmat)
    return EHTObservation(data = minset, mjd = mjd,
                          config=clac,
                          ra = ra, dec= dec,
                          bandwidth=bw, frequency=rf,
                          source = source,
                        )
end

function _minimal_closure(stcl, st)

    # Determine the number of timestamps containing a closure triangle
    T = typeof(stcl[1][1])
    S = typeof(stcl[1][1].u1)
    # loop over all timestamps
    minset = StructVector{T}(undef, 0)
    dmat = Matrix{S}[]
    for i in 1:length(st.times)

        scanvis = st[i]
        # find closure scan that matches time
        ind = findfirst(==(scanvis.time), stcl.times)
        if isnothing(ind)
            inow = length(dmat)
            # I want to construct block diagonal matrices for efficiency but we need
            # to be careful with scans that can't form closures. This hack moves these
            # scans into the preceding scan but fills it with zeros
            if i > 1
                dmat[inow] = hcat(dmat[inow], zeros(S, size(dmat[inow],1), length(scanvis.scan)))
            else
                push!(dmat, zeros(S, 1, length(scanvis.scan)))
            end
            continue
        end

        # Get our cphase scan
        scancl = stcl[ind]

        # sort by closure SNR so we form a nice minimal set
        snr = inv.(scancl.scan.error)
        ind_snr = sortperm(snr)
        scancl = scancl[ind_snr]
        snr = snr[ind_snr]

        # now find the visibility scan that matches the closure one based on timestamps

        # initialize the design matrix
        design_mat = closure_designmat(scancl, scanvis)

        # determine the expected size of the minimal set
        # this is needed to make sure we aren't killing too many triangles
        nmin = rank(design_mat)

        # print some info
        #println("For timestamp $(stcl.times[i]):")

        # get the current stations
        #println("Observing stations are $(stations(scancl))")

        #println("Size of maximal set of closure products = $(length(scancl))")
        #println("Size of minimal set of closure products = $(nmin)")
        #println("...")

        ##########################################################
        # start of loop to recover minimal set
        ##########################################################
        minset_scan, dmat_min = minimal_closure_scan(scancl, scanvis, nmin)
        minset = append!!(minset, minset_scan.scan)
        push!(dmat, dmat_min)
    end

    # Now if the first scan can't form a closure we will have an extra row of zeros
    # kill this
    if iszero(dmat[1])
        dmat[2] = hcat(zeros(S, size(dmat[2],1), size(dmat[1],2)), dmat[2])
        dmat = dmat[2:end]
    end
    return minset, dmat
end

function minimal_closure_scan(scancl0, scanvis, nmin::Int)
    # make a mask to keep track of which clhases will stick around
    scancl = deepcopy(scancl0)
    keep = fill(true, length(scancl))

    # remember the original minimal set size
    nmin0 = nmin

    # perform the loop
    count = 1
    keep = fill(true, length(scancl))
    good = true
    while good
        # recreate the mask each time
        scankeep = scancl[keep]

        design_mat = closure_designmat(scankeep, scanvis)

        # determine the size of the minimal set
        nmin = rank(design_mat)
        #println(nmin0, " ", nmin, " ", size(design_mat, 1))

        # Now decide whether to continue pruning
        if (sum(keep) == nmin0) && (nmin == nmin0)
            good = false
        else
            if nmin == nmin0
                keep[count] = false
            else
                keep[count-1] = true
                count -= 1
            end
            #(nmin == nmin0) && (keep[count] = false)
        end
        count += 1
        count+1 > length(scancl) && break
    end

    # print out the size of the recovered set for double-checking
    scankeep = scancl[keep]
    dmat = closure_designmat(scankeep, scanvis)
    if length(scankeep) != nmin
        @error "minimal set not found $(length(scankeep)) $(nmin)"
        throw("No minimal set found at time $(scancl.time)")
    end
    return scankeep, dmat
end

"""
    extract_cphase(obs)
Extracts the closure phases from an ehtim observation object

Any valid keyword arguments to `add_cphase` in ehtim can be passed through extract_cphase.

Returns an EHTObservation with closure phases datums

# Keyword arguments:
 - count: How the closures are formed, the available options are "min-correct", "min", "max"
 - cut_trivial: Cut the trivial triangles from the closures
 - uvmin: The flag to decide what are trivial triangles. Any baseline with ||(u,v)|| < uvmin
          are removed.

# Warning

The `count` keyword argument is treated specially in `Comrade`. The default option
is "min-correct" and should almost always be used.
This option construct a minimal set of closure phases that is valid even when
the array isn't fully connected. For testing and legacy reasons we `ehtim` other count
options are also included. However, the current `ehtim` count="min" option is broken
and does construct proper minimal sets of closure quantities if the array isn't fully connected.

"""
function extract_cphase(obs; count="min-correct", cut_trivial=false, uvmin=0.1e9,  kwargs...)
    if count == "min-correct"
        return minimal_cphase(obs; cut_trivial, uvmin, kwargs...)
    else
        return _ehtim_cphase(obs; count=count, cut_trivial=cut_trivial, uvmin=uvmin, kwargs...)
    end
end

"""
    extract_lcamp(obs)
Extracts the log-closure amp. from an ehtim observation object

Any valid keyword arguments to `add_logcamp` in ehtim can be passed through extract_lcamp.

Returns an EHTObservation with log-closure amp. datums

# Warning
The `count` keyword argument is treated specially in `Comrade`. The default option
is "min-correct" and should almost always be used.
This option construct a minimal set of closure phases that is valid even when
the array isn't fully connected. For testing and legacy reasons we `ehtim` other count
options are also included. However, the current `ehtim` count="min" option is broken
and does construct proper minimal sets of closure quantities if the array isn't fully connected.

"""
function extract_lcamp(obs; count="min-correct", kwargs...)
    if count == "min-correct"
        return minimal_lcamp(obs; kwargs...)
    else
        return _ehtim_lcamp(obs; count=count, kwargs...)
    end
end


text:
using LinearAlgebra
using Statistics

mutable struct LinearRegression
    X
    Y
    n_data
    input_dim
    w
    b
end

function LinearRegression(X, Y)
    n_data, input_dim = size(X)
    return LinearRegression(X, Y, n_data, input_dim, nothing, nothing)
end

# 2. 最小二乗法を用いてモデルパラメーターを最適化
function train!(model::LinearRegression)
    # 行列Xに「1」の要素を追加
    ndim = model.n_data
    Z = hcat(model.X, ones(ndim))

    # 分母の計算
    ZZ = 1 / ndim * Z' * Z

    # 分子の計算
    ZY = 1 / ndim * Z' * model.Y

    # パラメーターvの最適化
    # ZZ * v = ZY を解く
    v = ZZ \ ZY
    model.w = v[1:end-1]
    model.b = last(v)
    model
end

# 3. 予測
# X: 入力データ(データ数 x 次元数のArray)
predict(model::LinearRegression, x) = x * model.w .+ model.b

# 4. 平方平均二乗誤差(Root Mean Squared Error)
# X: 入力データ(データ数 x 次元数のArray)
# Y: 出力データ(データ数 x 1のArray)
rmse(model::LinearRegression, x, y) = sqrt.(mean((predict(model, x) - y).^2))

# 5. 決定係数の計算
# X: 入力データ(データ数 x 次元数のArray)
# Y: 出力データ(データ数 x 1のArray)
function r2(model::LinearRegression, x, y)
    sq = sum((predict(model, x) - y).^2)
    vr = sum((y .- mean(y; dims=1)).^2)
    return 1 - sq / vr
end

# L2ノルム正則化最小二乗法を用いてモデルパラメーターを最適化
# lamb: 正則化の重み係数(実数スカラー)
function train_regularized!(model::LinearRegression, lamb=0.1)
    # 行列Xに「1」の要素を追加
    ndim = model.n_data
    Z = hcat(model.X, ones(ndim))

    # 分母の計算
    ZZ = 1 / ndim * Z' * Z + lamb * I

    # 分子の計算
    ZY = 1 / ndim * Z' * model.Y

    # パラメーターvの最適化
    # ZZ * v = ZY を解く
    v = ZZ \ ZY
    model.w = v[1:end-1]
    model.b = last(v)
    model
end


text:
using sb

sbp = OhmicSBParams(No = 5000, ωm = 300)
sbm = SBModel(sbp)

omegarange = range(0,length = 1001, stop=sbp.ωm*1.2)

specden = sb.spectraldensity(sbm,omegarange)
teospecden = sb.ohmicJ0(omegarange,sbp)

dos = sb.densityofstates(sbm,omegarange)
teodos = sb.ρosc(convert(Vector{Float64},omegarange),sbp)

@test abs(findmax(teospecden)[1] - findmax(specden)[1]) < 0.01
@test sum(abs.(specden[200:600]-teospecden[200:600])) < 0.1
@test sum(abs.(dos[200:600]-teodos[200:600])) < 2.0


text:
const bb_score = gensym("score")
const bb_weight = gensym("weight")
const bb_new_trace = gensym("trace")

struct BBUpdateState
    marked::Set{ValueNode}
    stmts::Vector{Expr}
    schema::Union{StaticAddressSchema,EmptyAddressSchema}
    addr_visited::Set{Symbol}
    discard_leaf_nodes::Dict{Symbol, Symbol}
    discard_internal_nodes::Dict{Symbol, Symbol}
end

function BBUpdateState(stmts::Vector{Expr}, ir::BasicBlockIR, schema::Union{StaticAddressSchema,EmptyAddressSchema}, args_change_type)
    addr_visited = Set{Symbol}()
    marked = Set{ValueNode}()
    mark_arguments!(marked, ir, args_change_type)
    mark_input_change_nodes!(marked, ir)
    discard_leaf_node = Dict{Symbol,Symbol}()
    discard_internal_node = Dict{Symbol,Symbol}()
    BBUpdateState(marked, stmts, schema, addr_visited, discard_leaf_node, discard_internal_node)
end

struct BBFixUpdateState
    marked::Set{ValueNode}
    stmts::Vector{Expr}
    schema::Union{StaticAddressSchema,EmptyAddressSchema}
    addr_visited::Set{Symbol}
    discard_leaf_nodes::Dict{Symbol, Symbol}
    discard_internal_nodes::Dict{Symbol, Symbol}
end

function BBFixUpdateState(stmts::Vector{Expr}, ir::BasicBlockIR, schema::Union{StaticAddressSchema,EmptyAddressSchema}, args_change_type)
    addr_visited = Set{Symbol}()
    marked = Set{ValueNode}()
    mark_arguments!(marked, ir, args_change_type)
    mark_input_change_nodes!(marked, ir)
    discard_leaf_node = Dict{Symbol,Symbol}()
    discard_internal_node = Dict{Symbol,Symbol}()
    BBFixUpdateState(marked, stmts, schema, addr_visited, discard_leaf_node, discard_internal_node)
end

struct BBExtendState
    marked::Set{ValueNode}
    stmts::Vector{Expr}
    schema::Union{StaticAddressSchema,EmptyAddressSchema}
    addr_visited::Set{Symbol}
end

function BBExtendState(stmts::Vector{Expr}, ir::BasicBlockIR, schema::Union{StaticAddressSchema,EmptyAddressSchema}, args_change_type)
    addr_visited = Set{Symbol}()
    marked = Set{ValueNode}()
    mark_arguments!(marked, ir, args_change_type)
    mark_input_change_nodes!(marked, ir)
    BBExtendState(marked, stmts, schema, addr_visited)
end

function mark_input_change_nodes!(marked::Set{ValueNode}, ir::BasicBlockIR)
    for node in ir.generator_input_change_nodes
        # for now, mark every input change node to a generator
        # TODO we should only mark the node if the corresponding generator is
        # either constrained or marked, however, this requires two passes.
        # postponed for simplicity.
        push!(marked, node)
    end
end

function mark_arguments!(marked::Set{ValueNode}, ir::BasicBlockIR, args_change::Type{Nothing})
    for arg_node in ir.arg_nodes
        push!(marked, arg_node)
    end
end

function mark_arguments!(marked::Set{ValueNode}, ir::BasicBlockIR, args_change::Type{NoChange}) end

function mark_arguments!(marked::Set{ValueNode}, ir::BasicBlockIR, args_change::Type{T}) where {T <: MaskedArgChange}
    mask = args_change.parameters[1].parameters
    for (arg_node, maybe_changed_val) in zip(ir.arg_nodes, mask)
        if maybe_changed_val.parameters[1]
            push!(marked, arg_node)
        end
    end
end

function process!(ir::BasicBlockIR, state::Union{BBUpdateState,BBFixUpdateState,BBExtendState}, node::JuliaNode)

    # if any input nodes are marked, mark the output node
    if any([input in state.marked for input in node.input_nodes])
        push!(state.marked, node.output)
    end

    # set the value in the new trace based on other values in the new trace
    (typ, trace_field) = get_value_info(node)
    if node.output in state.marked
        push!(state.stmts, quote
            $bb_new_trace.$trace_field = $(expr_read_from_trace(node, bb_new_trace))
        end)
    end
end

function process!(ir::BasicBlockIR, state::Union{BBUpdateState,BBFixUpdateState,BBExtendState}, node::ArgsChangeNode)
    # always mark
    push!(state.marked, node.output)

    # set the value in the new trace (in the future, for performance
    # optimization, we can avoid tracing this value). we trace it for
    # simplicity and uniformity of implementation.
    (typ, trace_field) = get_value_info(node)
    push!(state.stmts, quote
        $bb_new_trace.$trace_field = args_change
    end)
end

const addr_change_prefix = gensym("addrchange")

function addr_change_variable(addr::Symbol)
    Symbol("$(addr_change_prefix)_$(addr)")
end

function process!(ir::BasicBlockIR, state::Union{BBUpdateState,BBFixUpdateState,BBExtendState}, node::AddrChangeNode)
    # always mark
    push!(state.marked, node.output)

    (typ, trace_field) = get_value_info(node)
    addr = node.address
    @assert addr in state.addr_visited
    if haskey(ir.addr_dist_nodes, addr)
        dist_node = ir.addr_dist_nodes[addr]
        # TODO: this implies we cannot access @change for addresses that don't have outputs?
        constrained = dist_node.output in state.marked
        # return whether the value changed and the previous value
        push!(state.stmts, quote
            $bb_new_trace.$trace_field = ($(QuoteNode(constrained)), trace.$addr)
        end)
    else
        if !haskey(ir.addr_gen_nodes, addr)
            # it is neither the address of a distribution or a generator
            error("Unknown address: $addr")
        end
        push!(state.stmts, quote
            $bb_new_trace.$trace_field = $(addr_change_variable(addr))
        end)
    end
end

function process_no_discard!(ir::BasicBlockIR, state::Union{BBUpdateState,BBFixUpdateState,BBExtendState}, node::AddrDistNode)
    addr = node.address
    push!(state.addr_visited, addr)
    typ = get_return_type(node.dist)
    dist = QuoteNode(node.dist)
    args = get_args(bb_new_trace, node)
    prev_args = get_args(:trace, node)
    decrement = gensym("decrement")
    increment = gensym("increment")
    input_nodes_marked = any([input in state.marked for input in node.input_nodes])
    if isa(state.schema, StaticAddressSchema) && addr in leaf_node_keys(state.schema)
        # constrained to a new value (mark the output)
        if has_output(node)
            push!(state.marked, node.output)
        end
        prev_value = gensym("prev_value")
        push!(state.stmts, quote
            $bb_new_trace.$addr = static_get_leaf_node(constraints, Val($(QuoteNode(addr))))
            $prev_value::$typ = trace.$addr
            $increment = logpdf($dist, $bb_new_trace.$addr, $(args...))
            $decrement = logpdf($dist, $prev_value, $(prev_args...))
            $bb_score += $increment - $decrement
            $bb_weight += $increment - $decrement
        end)
        state.discard_leaf_nodes[addr] = prev_value
        if has_output(node)
            (_, trace_field) = get_value_info(node)
            # TODO redundant with addr field, for use by other later statements:
            push!(state.stmts, quote
                $bb_new_trace.$trace_field = $bb_new_trace.$addr
            end)
        end
    elseif input_nodes_marked
         push!(state.stmts, quote
            $increment = logpdf($dist, trace.$addr, $(args...))
            $decrement = logpdf($dist, trace.$addr, $(prev_args...))
            $bb_score += $increment - $decrement
            $bb_weight += $increment - $decrement
        end)
    end
end

function process!(ir::BasicBlockIR, state::Union{BBUpdateState,BBFixUpdateState}, node::AddrDistNode)
    process_no_discard!(ir, state, node)
    if isa(state.schema, StaticAddressSchema) && addr in leaf_node_keys(state.schema)
        state.discard_leaf_nodes[addr] = prev_value
    end
end

function process!(ir::BasicBlockIR, state::BBExtendState, node::AddrDistNode)
    process_no_discard!(ir, state, node)
end

function get_constraints(schema::Union{StaticAddressSchema,EmptyAddressSchema}, addr::Symbol)
    if isa(schema, StaticAddressSchema) && addr in internal_node_keys(schema)
        constraints = :(static_get_internal_node(constraints, Val($(QuoteNode(addr)))))
        constrained = true
    else
        constrained = false
        constraints = :(EmptyAssignment())
    end
    (constrained, constraints)
end

function generate_generator_output_statement!(stmts::Vector{Expr}, node::AddrGeneratorNode, addr::Symbol)
    if has_output(node)
        (_, output_value_field) = get_value_info(node)
        push!(stmts, quote
            $bb_new_trace.$output_value_field = get_call_record($bb_new_trace.$addr).retval
        end)
    end
end

function generate_generator_call_statement!(state::BBUpdateState, addr::Symbol, node::AddrGeneratorNode, constraints)
    args = get_args(bb_new_trace, node)
    prev_args = get_args(:trace, node)
    change_value_ref = :($bb_new_trace.$(value_field(node.change_node)))
    discard = gensym("discard")
    weight = gensym("weight")
    push!(state.stmts, quote
        ($bb_new_trace.$addr, $weight, $discard, $(addr_change_variable(addr))) = update(
            $(QuoteNode(node.gen)), $(Expr(:tuple, args...)),
            $change_value_ref, trace.$addr, $constraints)
    end)
    push!(state.stmts, quote
        $bb_weight += $weight
    end)
    state.discard_internal_nodes[addr] = discard
end

function generate_generator_call_statement!(state::BBFixUpdateState, addr::Symbol, node::AddrGeneratorNode, constraints)
    args = get_args(bb_new_trace, node)
    prev_args = get_args(:trace, node)
    change_value_ref = :($bb_new_trace.$(value_field(node.change_node)))
    discard = gensym("discard")
    weight = gensym("weight")
    push!(state.stmts, quote
        ($bb_new_trace.$addr, $weight, $discard, $(addr_change_variable(addr))) = fix_update(
            $(QuoteNode(node.gen)), $(Expr(:tuple, args...)),
            $change_value_ref, trace.$addr, $constraints)
    end)
    push!(state.stmts, quote
        $bb_weight += $weight
    end)
    state.discard_internal_nodes[addr] = discard
end

function generate_generator_call_statement!(state::BBExtendState, addr::Symbol, node::AddrGeneratorNode, constraints)
    args = get_args(bb_new_trace, node)
    prev_args = get_args(:trace, node)
    change_value_ref = :($bb_new_trace.$(value_field(node.change_node)))
    weight = gensym("weight")
    push!(state.stmts, quote
        ($bb_new_trace.$addr, $weight, $(addr_change_variable(addr))) = extend(
            $(QuoteNode(node.gen)), $(Expr(:tuple, args...)),
            $change_value_ref, trace.$addr, $constraints)
    end)
    push!(state.stmts, quote
        $bb_weight += $weight
    end)
end

function generate_generator_score_statements!(stmts::Vector{Expr}, addr::Symbol)
    decrement = gensym("decrement")
    increment = gensym("increment")
    push!(stmts, quote
        $decrement = get_call_record(trace.$addr).score
        $increment = get_call_record($bb_new_trace.$addr).score
        $bb_score += $increment - $decrement
    end)
end

function process_generator_update_marked!(state::Union{BBUpdateState,BBFixUpdateState,BBExtendState}, node::AddrGeneratorNode)
    # return value could change (even if just the input nodes are marked,
    # we don't currently statically identify a generator that can absorb
    # arbitrary changes to its arguments)
    if has_output(node)
        push!(state.marked, node.output)
    end
end

function process!(ir::BasicBlockIR, state::Union{BBUpdateState,BBFixUpdateState}, node::AddrGeneratorNode)
    addr = node.address
    push!(state.addr_visited, addr)
    input_nodes_marked = any([input in state.marked for input in node.input_nodes])
    (constrained, constraints) = get_constraints(state.schema, addr)
    if constrained || input_nodes_marked
        process_generator_update_marked!(state, node)
        generate_generator_call_statement!(state, addr, node, constraints)
        generate_generator_score_statements!(state.stmts, addr)
        generate_generator_output_statement!(state.stmts, node, addr)
    else
        push!(state.stmts, quote
            $(addr_change_variable(addr)) = NoChange()
        end)
    end
end

function process!(ir::BasicBlockIR, state::BBExtendState, node::AddrGeneratorNode)
    addr = node.address
    push!(state.addr_visited, addr)
    input_nodes_marked = any([input in state.marked for input in node.input_nodes])
    (constrained, constraints) = get_constraints(state.schema, addr)
    if constrained || input_nodes_marked
        process_generator_update_marked!(state, node)
        generate_generator_call_statement!(state, addr, node, constraints)
        generate_generator_score_statements!(state.stmts, addr)
        generate_generator_output_statement!(state.stmts, node, addr)
    else
        push!(state.stmts, quote
            $(addr_change_variable(addr)) = NoChange()
        end)
    end
end

function generate_init_statements!(stmts::Vector{Expr})
    push!(stmts, quote
        $bb_new_trace = copy(trace)
        $bb_score = trace.$call_record_field.score
        $bb_weight = 0.
    end)
end

function generate_arg_statements!(stmts::Vector{Expr}, ir::BasicBlockIR)

    # unpack arguments into variables
    arg_names = Symbol[arg_node.name for arg_node in ir.arg_nodes]
    push!(stmts, Expr(:(=), Expr(:tuple, arg_names...), :new_args))

    # record arguments in trace
    for arg_node in ir.arg_nodes
        push!(stmts, quote $bb_new_trace.$(value_field(arg_node)) = $(arg_node.name) end)
    end
end

function generate_expr_node_statements!(state::Union{BBUpdateState,BBFixUpdateState,BBExtendState}, ir::BasicBlockIR)
    # visit statements in topological order, generating code for each one
    for node in ir.expr_nodes_sorted
        process!(ir, state, node)
    end
end

function generate_is_empty!(stmts::Vector{Expr}, ir::BasicBlockIR)
    # NOTE: this is still O(N) where N is the number of generator calls,
    # including non-visited calls
    if !isempty(ir.addr_dist_nodes)
        push!(stmts, quote
            $bb_new_trace.$is_empty_field = false
        end)
    else
        for (addr, node::AddrGeneratorNode) in ir.addr_gen_nodes
            push!(stmts, quote
                $bb_new_trace.$is_empty_field = $bb_new_trace.$is_empty_field && !has_choices($bb_new_trace.$addr)
            end)
        end
    end
end

function generate_discard!(stmts::Vector{Expr}, discard_leaf_nodes, discard_internal_nodes)
    if length(discard_leaf_nodes) > 0
        (leaf_keys, leaf_nodes) = collect(zip(discard_leaf_nodes...))
    else
        (leaf_keys, leaf_nodes) = ((), ())
    end
    if length(discard_internal_nodes) > 0
        (internal_keys, internal_nodes) = collect(zip(discard_internal_nodes...))
    else
        (internal_keys, internal_nodes) = ((), ())
    end
    leaf_keys = map((k) -> QuoteNode(k), leaf_keys)
    internal_keys = map((k) -> QuoteNode(k), internal_keys)
    push!(stmts, quote
        discard = StaticAssignment(
            NamedTuple{($(leaf_keys...),)}(($(leaf_nodes...),)),
            NamedTuple{($(internal_keys...),)}(($(internal_nodes...),)))
    end)
end

function check_no_extra_constraints(schema::StaticAddressSchema, ir::BasicBlockIR)
    addresses = union(keys(ir.addr_dist_nodes), keys(ir.addr_gen_nodes))
    for addr in union(leaf_node_keys(schema), internal_node_keys(schema))
        if !(addr in addresses)
            error("Update did not consume all constraints")
        end
    end
end

function check_no_extra_constraints(schema::EmptyAddressSchema, ir::BasicBlockIR)
end

function generate_call_record!(stmts::Vector{Expr}, ir::BasicBlockIR, marked::Set{ValueNode})

    # return value
    if ir.output_node === nothing
        retval = :nothing
    else
        if ir.output_node in marked
            retval = quote $bb_new_trace.$(value_field(ir.output_node)) end
        else
            retval = quote trace.$call_record_field.retval end
        end
    end

    # construct new call record
    # TODO move the change to a returnvalue of update, not part of the call record
    # if we do that, then we will need separate fields in which to store the
    # retchange values (one for each geneator)
    push!(stmts, quote
        $bb_new_trace.$call_record_field = CallRecord($bb_score, $retval, new_args)
    end)
end

function generate_update_return_statement!(stmts::Vector{Expr}, ir::BasicBlockIR)
    if ir.retchange_node === nothing
        retchange = :(nothing)
    else
        retchange = Expr(:(.), bb_new_trace, QuoteNode(value_field(ir.retchange_node)))
    end
    push!(stmts, quote return ($bb_new_trace, $bb_weight, discard, $retchange) end)
end

function generate_extend_return_statement!(stmts::Vector{Expr}, ir::BasicBlockIR)
    if ir.retchange_node === nothing
        retchange = :(nothing)
    else
        retchange = Expr(:(.), bb_new_trace, QuoteNode(value_field(ir.retchange_node)))
    end
    push!(stmts, quote return ($bb_new_trace, $bb_weight, $retchange) end)
end

function codegen_update(gen_type::Type{T}, new_args_type, args_change_type, trace_type, constraints_type) where {T <: BasicGenFunction}
    schema = get_address_schema(constraints_type)
    ir = get_ir(gen_type)
    stmts = Expr[]
    generate_init_statements!(stmts)
    generate_arg_statements!(stmts, ir)
    state = BBUpdateState(stmts, ir, schema, args_change_type)
    generate_expr_node_statements!(state, ir)
    generate_is_empty!(stmts, ir)
    generate_discard!(stmts, state.discard_leaf_nodes, state.discard_internal_nodes)
    generate_call_record!(stmts, ir, state.marked)
    generate_update_return_statement!(stmts, ir)
    return Expr(:block, stmts...)
end

function codegen_fix_update(gen_type::Type{T}, new_args_type, args_change_type, trace_type, constraints_type) where {T <: BasicGenFunction}
    schema = get_address_schema(constraints_type)
    ir = get_ir(gen_type)
    stmts = Expr[]
    generate_init_statements!(stmts)
    generate_arg_statements!(stmts, ir)
    state = BBFixUpdateState(stmts, ir, schema, args_change_type)
    generate_expr_node_statements!(state, ir)
    generate_is_empty!(stmts, ir)
    generate_discard!(stmts, state.discard_leaf_nodes, state.discard_internal_nodes)
    generate_call_record!(stmts, ir, state.marked)
    generate_update_return_statement!(stmts, ir)
    return Expr(:block, stmts...)
end

function codegen_extend(gen_type::Type{T}, new_args_type, args_change_type, trace_type, constraints_type) where {T <: BasicGenFunction}
    schema = get_address_schema(constraints_type)
    ir = get_ir(gen_type)
    stmts = Expr[]
    generate_init_statements!(stmts)
    generate_arg_statements!(stmts, ir)
    state = BBExtendState(stmts, ir, schema, args_change_type)
    generate_expr_node_statements!(state, ir)
    generate_is_empty!(stmts, ir)
    generate_call_record!(stmts, ir, state.marked)
    generate_extend_return_statement!(stmts, ir)
    return Expr(:block, stmts...)
end

push!(Gen.generated_functions, quote
@generated function Gen.update(gen::Gen.BasicGenFunction{T,U}, new_args, args_change, trace::U, constraints) where {T,U}
    schema = get_address_schema(constraints)
    if !(isa(schema, StaticAddressSchema) || isa(schema, EmptyAddressSchema))
        # try to convert it to a static assignment
        return quote update(gen, new_args, args_change, trace, StaticAssignment(constraints)) end
    end
    Gen.codegen_update(gen, new_args, args_change, trace, constraints)
end
end)

push!(Gen.generated_functions, quote
@generated function Gen.fix_update(gen::Gen.BasicGenFunction{T,U}, new_args, args_change, trace::U, constraints) where {T,U}
    schema = get_address_schema(constraints)
    if !(isa(schema, StaticAddressSchema) || isa(schema, EmptyAddressSchema))
        # try to convert it to a static assignment
        return quote fix_update(gen, new_args, args_change, trace, StaticAssignment(constraints)) end
    end
    Gen.codegen_fix_update(gen, new_args, args_change, trace, constraints)
end
end)

push!(Gen.generated_functions, quote
@generated function Gen.extend(gen::Gen.BasicGenFunction{T,U}, new_args, args_change, trace::U, constraints) where {T,U}
    schema = get_address_schema(constraints)
    if !(isa(schema, StaticAddressSchema) || isa(schema, EmptyAddressSchema))
        # try to convert it to a static assignment
        return quote extend(gen, new_args, args_change, trace, StaticAssignment(constraints)) end
    end
    Gen.codegen_extend(gen, new_args, args_change, trace, constraints)
end
end)

max_stars_repo_name: BoltzmannBrain/Gen

max_issues_repo_name: BoltzmannBrain/Gen

max_forks_repo_name: BoltzmannBrain/Gen


text:
module BestBasis
export
    # cost functions
    CostFunction,
    LSDBCost,
    JBBCost,
    BBCost,
    LoglpCost,
    NormCost,
    DifferentialEntropyCost,
    ShannonEntropyCost,
    LogEnergyEntropyCost,
    # tree cost
    tree_costs,
    # best basis types
    BestBasisType,
    LSDB,
    JBB,
    BB,
    SIBB,
    # best basis tree
    bestbasistreeall

using
    Wavelets,
    LinearAlgebra,
    Statistics,
    AverageShiftedHistograms,
    Parameters

using
    ..Utils,
    ..DWT,
    ..SIWPD

include("bestbasis_costs.jl")
include("bestbasis_tree.jl")

## ----- BEST TREE SELECTION -----
# Tree selection for 1D signals
"""
    bestbasis_treeselection(costs, n[, type])
    bestbasis_treeselection(costs, n, m[, type])

Computes the best basis tree based on the given cost vector.

# Arguments
- `costs::AbstractVector{T}`: Vector containing costs for each node.
- `n::Integer`: Length of signals (for 1D cases) or vertical length of signals (for 2D
  cases).
- `m::Integer`: Horizontal length of signals (for 2D cases).
- `type::Symbol`: (Default: `:min`) Criterion used to select the best tree. Supported types
  are `:min` and `:max`. Eg. Setting `type = :min` results in a basis tree with the lowest
  cost to be selected.

# Returns
- `::BitVector`: Best basis tree selected based on cost.

**See also:** [`bestbasistree`](@ref), [`tree_costs`](@ref)
"""
function bestbasis_treeselection(costs::AbstractVector{T},
                                 n::Integer,
                                 type::Symbol = :min) where T<:AbstractFloat
    k = length(costs)
    @assert k ≤ gettreelength(2*n)
    @assert type ∈ [:min, :max] || throw(ArgumentError("Unsupported type $type."))
    L = getdepth(k,:binary)
    tree = maketree(n, L, :full)
    for i in reverse(eachindex(tree))
        if tree[i]                  # Check costs if node i exists
            # Get parent and children costs
            @inbounds pc = costs[i]
            @inbounds cc = costs[getchildindex(i,:left)] + costs[getchildindex(i,:right)]
            if type == :min && cc < pc          # Child cost < parent cost
                @inbounds costs[i] = cc
            elseif type == :max && cc > pc      # Child cost > parent cost
                @inbounds costs[i] = cc
            else
                delete_subtree!(tree, i, :binary)
            end
        end
    end
    @assert isvalidtree(zeros(n), tree)
    return tree
end
# Tree selection for 2D signals
function bestbasis_treeselection(costs::AbstractVector{T},
                                 n::Integer, m::Integer,
                                 type::Symbol = :min) where T<:AbstractFloat
    k = length(costs)
    @assert k ≤ gettreelength(2*n,2*m)
    @assert type ∈ [:min, :max] || throw(ArgumentError("Unsupported type $type."))
    L = getdepth(k,:quad)
    tree = maketree(n, m, L, :full)
    for i in reverse(eachindex(tree))
        if tree[i]
            # Get parent and children costs
            @inbounds pc = costs[i]
            @inbounds cc = costs[getchildindex(i,:topleft)] + costs[getchildindex(i,:topright)] +
                        costs[getchildindex(i,:bottomleft)] + costs[getchildindex(i,:bottomright)]
            if type == :min && cc < pc
                @inbounds costs[i] = cc
            elseif type == :max && cc > pc
                @inbounds costs[i] = cc
            else
                delete_subtree!(tree, i, :quad)
            end
        end
    end
    @assert isvalidtree(zeros(n,m), tree)
    return tree
end

# SIWPD tree selection
"""
    bestbasis_treeselection(costs, tree)

Best basis tree selection on SIWPD.

# Arguments
- `costs::AbstractVector{Tc} where Tc<:AbstractVector{<:Union{Number, Nothing}}`: Cost of
  each node.
- `tree::AbstractVector{Tt} where Tt<:BitVector`: SIWPD tree.

!!! warning
    Current implementation works but is unstable, ie. we are still working on better
    syntax/more optimized computations/better data structure.
"""
function bestbasis_treeselection(costs::AbstractVector{Tc},
        tree::AbstractVector{Tt}) where
        {Tc<:AbstractVector{<:Union{Number,Nothing}}, Tt<:BitVector}

    @assert length(costs) == length(tree)
    bt = deepcopy(tree)
    bc = deepcopy(costs)
    nn = length(tree)
    for i in reverse(eachindex(bt))
        if getchildindex(i,:left) > nn                        # current node is at bottom level
            continue
        end
        level = floor(Int, log2(i))
        for j in eachindex(bt[i])       # iterate through all available shifts
            if !bt[i][j]                # node of current shift does not exist
                continue
            end
            @assert bt[getchildindex(i,:left)][j] ==
                    bt[getchildindex(i,:right)][j] ==
                    bt[getchildindex(i,:left)][j+1<<level] ==
                    bt[getchildindex(i,:right)][j+1<<level] == true || continue
            # child cost of current shift of current node
            cc = bc[getchildindex(i,:left)][j] + bc[getchildindex(i,:right)][j]
            # child cost of circshift of current node
            scc = bc[getchildindex(i,:left)][j+1<<level] +
                  bc[getchildindex(i,:right)][j+1<<level]
            mincost = min(cc, scc, bc[i][j])
            # mincost=parent, delete subtrees of both cc & scc
            if mincost == bc[i][j]
                delete_subtree!(bt, getchildindex(i,:left), j)
                delete_subtree!(bt, getchildindex(i,:right), j)
                delete_subtree!(bt, getchildindex(i,:left), j+1<<level)
                delete_subtree!(bt, getchildindex(i,:right), j+1<<level)
            # mincost=cc, delete subtrees of scc
            elseif mincost == cc
                bc[i][j] = mincost
                delete_subtree!(bt, getchildindex(i,:left), j+1<<level)
                delete_subtree!(bt, getchildindex(i,:right), j+1<<level)
            # mincost=scc, delete subtrees of cc
            else
                bc[i][j] = mincost
                delete_subtree!(bt, getchildindex(i,:left), j)
                delete_subtree!(bt, getchildindex(i,:right), j)
            end
        end
    end
    # ensure only 1 node selected from each Ω(i,j)
    @assert all(map(node -> sum(node), bt) .<= 1)
    return bt
end

# Deletes subtree due to inferior cost
"""
    delete_subtree!(bt, i, tree_type)

Deletes a subtree of the entire tree due to children's inferior costs.

# Arguments
- `bt::BitVector`: Tree.
- `i::Integer`: Root of the subtree to be deleted.
- `tree_type::Symbol`: Type of tree (`:binary` or `:quad`).

# Returns
`bt::BitVector`: Tree with deleted subtree.

**See also:** [`bestbasistree`](@ref), [`bestbasis_treeselection`](@ref)
"""
function delete_subtree!(bt::BitVector, i::Integer, tree_type::Symbol)
    @assert 1 ≤ i ≤ length(bt)
    @assert tree_type ∈ [:binary, :quad]
    children = tree_type == :binary ? [:left, :right] : [:topleft, :topright, :bottomleft, :bottomright]
    bt[i] = false
    for c in children
        # Delete child subtree if exist
        if (getchildindex(i,c) ≤ length(bt)) && bt[getchildindex(i,c)]
            delete_subtree!(bt, getchildindex(i,c), tree_type)
        end
    end
    return bt
end

function delete_subtree!(bt::AbstractVector{BitVector}, i::Integer, j::Integer)
    @assert 1 <= i <= length(bt)
    level = floor(Int, log2(i))
    bt[i][j] = false
    if (getchildindex(i,:left)) < length(bt)          # current node can contain subtrees
        if bt[getchildindex(i,:left)][j]              # left child of current shift
            delete_subtree!(bt, getchildindex(i,:left), j)
        end
        if bt[getchildindex(i,:right)][j]            # right child of current shift
            delete_subtree!(bt, getchildindex(i,:right), j)
        end
        if bt[getchildindex(i,:left)][j+1<<level]     # left child of added shift
            delete_subtree!(bt, getchildindex(i,:left), j+1<<level)
        end
        if bt[getchildindex(i,:right)][j+1<<level]   # right child of added shift
            delete_subtree!(bt, getchildindex(i,:right), j+1<<level)
        end
    end
    return nothing
end

## BEST BASIS TREES
# Best basis tree for LSDB
"""
    bestbasistree(X[, method])

Extension to the best basis tree function from Wavelets.jl. Given a set of decomposed
signals, returns different types of best basis trees based on the methods specified.
Available methods are the joint best basis ([`JBB`](@ref)), least statistically dependent
basis ([`LSDB`](@ref)), individual regular best basis ([`BB`](@ref)), and shift-invariant
best basis ([`SIBB`](@ref)).

# Arguments
- `X::AbstractArray{T} where T<:AbstractFloat`: A set of decomposed signals, of sizes
  `(n,L,k)` for 1D signals or `(n,m,L,k)` for 2D signals, where:
    - `n`: Length of signal (1D) or vertical length of signal (2D).
    - `m`: Horizontal length of signal (2D).
    - `L`: Number of decomposition levels plus 1 (for standard wavelet decomposition) or
      number of nodes in the tree (for redundant transforms such as ACWT and SWT).
    - `k`: Number of signals.
- `method::BestBasisType`: Type of best basis, ie. `BB()`, `JBB()` or `LSDB()`.

!!! tip
    For standard best basis (`BB()`), this current function can only process one signal at a
    time, ie. the input `X` should have dimensions `(n,L)` or `(n,m,L)`. To process multiple
    signals using one function, see [`bestbasistreeall`](@ref).

# Returns
- `::BitVector`: Best basis tree.

# Examples
```julia
using Wavelets, WaveletsExt

X = generatesignals(:heavisine, 6) |> x -> duplicatesignals(x, 5, 2, true)
wt = wavelet(WT.db4)
Xw = wpdall(X, wt)

bestbasistree(Xw, JBB())
bestbasistree(Xw, LSDB())
```

**See also:** [`getbasiscoef`](@ref), [`getbasiscoefall`](@ref), [`tree_costs`](@ref),
[`delete_subtree!`](@ref)
"""
function Wavelets.Threshold.bestbasistree(X::AbstractArray{T}, method::LSDB) where
                                          T<:AbstractFloat
    @assert 3 ≤ ndims(X) ≤ 4                # Compatible for 1D and 2D decomposed signals
    sz = size(X)[1:end-2]                   # Signal size
    costs = tree_costs(X, method)
    tree = bestbasis_treeselection(costs, sz...)
    return tree
end
# Best basis tree for JBB
function Wavelets.Threshold.bestbasistree(X::AbstractArray{T}, method::JBB) where
                                          T<:AbstractFloat
    @assert 3 ≤ ndims(X) ≤ 4                # Compatible for 1D and 2D decomposed signals
    sz = size(X)[1:end-2]                   # Signal size
    costs = tree_costs(X, method)
    tree = bestbasis_treeselection(costs, sz...)
    return tree
end
# Standard best basis tree
function Wavelets.Threshold.bestbasistree(X::AbstractArray{T}, method::BB) where
    T<:AbstractFloat
    @assert 2 ≤ ndims(X) ≤ 3                # Compatible for 1D and 2D decomposed signal
    sz = size(X)[1:end-1]                   # Signal size
    costs = tree_costs(X, method)
    tree = bestbasis_treeselection(costs, sz...)
    return tree
end
# SIWPD Best basis tree
# TODO: find a way to compute bestbasis_tree without input d
"""
    bestbasistree(y, d, method)

Computes the best basis tree for the shift invariant wavelet packet decomposition (SIWPD).

# Arguments
- `y::AbstractArray{T,2} where T<:Number`: A SIWPD decomposed signal.
- `d::Integer`: The number of depth computed for the decomposition.
- `method::SIBB`: The `SIBB()` method.

# Returns
- `Vector{BitVector}`: SIWPD best basis tree.

!!! warning
    Current implementation works but is unstable, ie. we are still working on better
    syntax/more optimized computations/better data structure.
"""
function Wavelets.Threshold.bestbasistree(y::AbstractArray{T,2},
                                          d::Integer,
                                          method::SIBB) where T<:Number

    nn = size(y,2)
    L = maxtransformlevels((nn+1)÷2)
    ns = size(y,1)
    tree = makesiwpdtree(ns, L, d)
    costs = tree_costs(y, tree, method)
    besttree = bestbasis_treeselection(costs, tree)
    return besttree
end
# Default best basis tree search
function Wavelets.Threshold.bestbasistree(X::AbstractArray{T},
                                          method::BestBasisType = JBB()) where
                                          T<:AbstractFloat
    return bestbasistree(X, method)
end

"""
    bestbasistreeall(X, method)

Compute the standard best basis tree of a set of signals.

# Arguments
- `X::AbstractArray{T} where T<:AbstractFloat`: A set of decomposed signals, of sizes
  `(n,L,k)` for 1D signals or `(n,m,L,k)` for 2D signals, where:
    - `n`: Length of signal (1D) or vertical length of signal (2D).
    - `m`: Horizontal length of signal (2D).
    - `L`: Number of decomposition levels plus 1 (for standard wavelet decomposition) or
        number of nodes in the tree (for redundant transforms such as ACWT and SWT).
    - `k`: Number of signals.
- `method::BB`: Standard best basis method, eg. `BB()`.

# Returns
- `::BitMatrix`: `(nₜ,k)` matrix where each column corresponds to a tree.

# Examples
```julia
using Wavelets, WaveletsExt

X = generatesignals(:heavisine, 6) |> x -> duplicatesignals(x, 5, 2, true)
wt = wavelet(WT.db4)

Xw = wpdall(X, wt)
bestbasistreeall(Xw, BB())

Xw = swpdall(X, wt)
bestbasistree(Xw, BB(redundant=true))
```

**See also:** [`bestbasistree`](@ref)
"""
function bestbasistreeall(X::AbstractArray{T}, method::BB) where T<:AbstractFloat
    @assert 3 ≤ ndims(X) ≤ 4                # Compatible for 1D and 2D decomposed signals
    sz = size(X)[1:end-2]                   # Signal size
    k = size(X)[end]                        # Number of signals
    trees = falses(gettreelength(sz...), k) # Allocate trees
    @views for (i, Xᵢ) in enumerate(eachslice(X, dims=ndims(X)))
        @inbounds trees[:,i] = bestbasistree(Xᵢ, method)
    end
    return trees
end

end # end of module


text:
using ITensors
using ITensors.NDTensors
using Test

@testset "SymmetryStyle trait" begin
  i = Index(2)
  iqn = Index([QN(0) => 1, QN(1) => 2])

  @test @inferred(ITensors.symmetrystyle(i)) == ITensors.NonQN()
  @test @inferred(ITensors.symmetrystyle((i,))) == ITensors.NonQN()
  @test @inferred(ITensors.symmetrystyle([i])) == ITensors.NonQN()
  @test @inferred(ITensors.symmetrystyle(i', i)) == ITensors.NonQN()
  @test @inferred(ITensors.symmetrystyle((i', i))) == ITensors.NonQN()
  @test @inferred(ITensors.symmetrystyle([i', i])) == ITensors.NonQN()
  @test @inferred(ITensors.symmetrystyle(i'', i', i)) == ITensors.NonQN()
  @test @inferred(ITensors.symmetrystyle((i'', i', i))) == ITensors.NonQN()
  @test @inferred(ITensors.symmetrystyle([i'', i', i])) == ITensors.NonQN()
  @test @inferred(ITensors.symmetrystyle(i''', i'', i', i)) == ITensors.NonQN()
  @test @inferred(ITensors.symmetrystyle((i''', i'', i', i))) == ITensors.NonQN()
  @test @inferred(ITensors.symmetrystyle([i''', i'', i', i])) == ITensors.NonQN()

  @test @inferred(ITensors.symmetrystyle(iqn)) == ITensors.HasQNs()
  @test @inferred(ITensors.symmetrystyle((iqn,))) == ITensors.HasQNs()
  @test @inferred(ITensors.symmetrystyle([iqn])) == ITensors.HasQNs()
  @test @inferred(ITensors.symmetrystyle(iqn', iqn)) == ITensors.HasQNs()
  @test @inferred(ITensors.symmetrystyle((iqn', iqn))) == ITensors.HasQNs()
  @test @inferred(ITensors.symmetrystyle([iqn', iqn])) == ITensors.HasQNs()
  @test @inferred(ITensors.symmetrystyle(iqn'', iqn', iqn)) == ITensors.HasQNs()
  @test @inferred(ITensors.symmetrystyle((iqn'', iqn', iqn))) == ITensors.HasQNs()
  @test @inferred(ITensors.symmetrystyle([iqn'', iqn', iqn])) == ITensors.HasQNs()
  @test @inferred(ITensors.symmetrystyle(iqn', i)) == ITensors.HasQNs()
  @test @inferred(ITensors.symmetrystyle(i', i, iqn)) == ITensors.HasQNs()
  @test @inferred(ITensors.symmetrystyle((i', i, iqn))) == ITensors.HasQNs()
  @test @inferred(ITensors.SymmetryStyle, ITensors.symmetrystyle([i', i, iqn])) ==
    ITensors.HasQNs()

  A = randomITensor(i', dag(i))
  Aqn = randomITensor(iqn', dag(iqn))

  @test @inferred(ITensors.SymmetryStyle, ITensors.symmetrystyle(A)) == ITensors.NonQN()
  @test @inferred(ITensors.SymmetryStyle, ITensors.symmetrystyle(Aqn)) == ITensors.HasQNs()

  T = Tensor(A)
  Tqn = Tensor(Aqn)

  @test @inferred(ITensors.symmetrystyle(T)) == ITensors.NonQN()
  @test @inferred(ITensors.symmetrystyle(Tqn)) == ITensors.HasQNs()

  sqn = siteinds("S=1/2", 10; conserve_qns=true)
  s = removeqns(sqn)
  psi = MPS(s)
  psiqn = MPS(sqn)
  @test @inferred(ITensors.SymmetryStyle, ITensors.symmetrystyle(psi)) == ITensors.NonQN()
  @test @inferred(ITensors.SymmetryStyle, ITensors.symmetrystyle(psiqn)) ==
    ITensors.HasQNs()
end