"""
Data structure
--------------
mutable struct: SELF.recording

Functions
---------
SELF.recording.filter()

SELF.recording.resample()

SELF.recording.detect_spikes()
"""
module Neurons

export recording, normalize_trace, filter_trace, resample, detect_spikes, detect_spikes1, get_PSD, extract_spikes, extract_features, cluster_spikes, find_redundant, extract_isih,
		plot_detection, plot_waveforms

using DSP
using Images
using Statistics
using MultivariateStats
using GaussianMixtures
using Clustering
using LinearAlgebra
using Distances
using PyPlot

using ..MyUtils

# main type
"""
Fields
------
+ file::String
+ sf::Real
+ sf_preprocessed::Real
+ n_channels::Real
+ trace::AbstractArray # supposed to contain sample X channel
+ trial_marker::AbstractArray
+ trace_preprocessed::AbstractArray
+ preprocessing_steps::Vector{String}
+ spike_idx::AbstractArray
+ detection_threshold::AbstractArray
+ isih::AbstractArray
+ spikeforms::AbstractArray # channel (val_source) X channel (t_source) X time of spike/spikeform X sample
+ features::AbstractArray # channel X feature space representation of spikes
+ spike_class::AbstractArray # channel X clusters (as spike indices)
+ BIC::AbstractArray
"""
mutable struct recording
    ### FIELDS
    file::String
    sf::Real
    sf_preprocessed::Real
    n_channels::Real
    trace::Array{Array{Float64,2}}#AbstractArray # supposed to contain sample X channel
    trial_marker::AbstractArray
    trace_preprocessed::Array{Array{Float64,2}}#AbstractArray
    preprocessing_steps::Vector{String}
    spike_idx::AbstractArray
	detection_threshold::AbstractArray
    isih::AbstractArray
    spikeforms::AbstractArray # channel (val_source) X channel (t_source) X time of spike/spikeform X sample
    features::AbstractArray # channel X feature space representation of spikes
    spike_class::AbstractArray # channel X clusters (as spike indices)
    BIC::AbstractArray

    ### CONSTRUCTOR
    # recording(file, sf, trace, trial_marker) =
    #    new(file, sf, 0, trace, trial_marker, Array{Float64,2}(undef,0,0), Vector{String}(undef,0), Array{Float64,3}Array{Float64,2}(undef,0,0,0))
    # recording() = new("", 0, 0, Array{Float64,2}(undef,0,0), Array{Float64,2}(undef,0,0), Array{Float64,2}(undef,0,0), Vector{String}(undef,0), Array{Float64,3}(undef,0,0,0))
    recording() = new("", 0, 0, 0, [], [], [], Vector{String}(undef,0), [], [], [], [], [], [])
end

# functions

function normalize_trace(n::Neurons.recording, what="trace")

	x = deepcopy( getfield( n, Symbol(what) ) )
	sd = std.(x)
	mn = mean.(x)
	[x[ii]=(x[ii] .- mn[ii]) ./ sd[ii] for ii in 1:n.n_channels] ### normalization
	n.trace_preprocessed = x
	n.preprocessing_steps = ["normalized"];

end

"""
Documentation ...
"""
function filter_trace(n::Neurons.recording, filter_type::String, freqs::AbstractArray, filter_design::AbstractArray, filter_order::AbstractArray, bandwidth::AbstractArray, from_chnnls::AbstractArray, what::String="trace")
"""
    filter(n::recording, band::Tuple{Real}, filter_type)
    filter_type: lowpass, highpass, bandpass, notch
    freqs = highpass: freq | lowpass: freq | passband (Wn1,Wn2) frequency | Notch: array of frequencies | CommonAverage: channels used for averaging
    what: spcifies whether filter works on trace or trace_preprocessed. If the latter trace_preprocessed is replaced, respectively.
"""

    ######################
    ### raw trace
    ######################
    if what == "trace"
        nyq = n.sf/2;
        freqs_n = freqs./nyq; # normalized to nquist frequncy
        n.trace_preprocessed = []
        n.sf_preprocessed = n.sf # restore original sampling frequency

        if filter_type == "Bandpass"
            rt = Bandpass( freqs_n[1], freqs_n[2] );
            dm = getfield( Main, Symbol(filter_design[1]) )(filter_order[1]);
            f = digitalfilter(rt, dm);
            ### loop over channels and filter trace
            for trc in n.trace
                #trc_padded = flatten([[.0 for ii in 1:pads], trc, [.0 for ii in 1:pads]])
                #push!(n.trace_preprocessed, filtfilt(f, trc_padded)[pads+1:end-pads]);
                push!(n.trace_preprocessed, filtfilt(f, trc));
            end
        elseif filter_type == "Lowpass"
            rt = Lowpass( freqs_n[1] );
            dm = getfield( Main,Symbol(filter_design[1]) )(filter_order[1]);
            f = digitalfilter(rt, dm);
            for trc in n.trace
                # trc_padded = flatten([[.0 for ii in 1:pads], trc, [.0 for ii in 1:pads]])
                # push!(n.trace_preprocessed, filtfilt(f, trc_padded)[pads+1:end-pads]);
                push!(n.trace_preprocessed, filtfilt(f, trc));
            end
        elseif filter_type == "Highpass"
            rt = Highpass( freqs_n[1] );
            dm = getfield( Main,Symbol(filter_design[1]) )(filter_order[1]);
            f = digitalfilter(rt, dm);
            for trc in n.trace
                # trc_padded = flatten([[.0 for ii in 1:pads], trc, [.0 for ii in 1:pads]])
                # push!(n.trace_preprocessed, filt(f, trc_padded)[pads+1:end-pads]);
                push!(n.trace_preprocessed, filtfilt(f, trc));
            end
        elseif filter_type == "Notch"
            for trc in n.trace
                trc = flatten(trc)
                for ff in freqs_n.*nyq
                    nn = iirnotch(ff, bandwidth[1]; fs=n.sf)
                    trc = filtfilt(nn, trc);
                end
                push!(n.trace_preprocessed, trc);
            end
        elseif filter_type == "CommonAverage"
            chnnls = freqs
            average_trace = [median([n.trace[ch][idx] for ch in chnnls]) for idx in 1:length(n.trace[1])]
            # n.trace_preprocessed = [channel .- average_trace for channel in n.trace[from_chnnls]]
			[n.trace_preprocessed[ch] = n.trace_preprocessed[ch] .- average_trace for ch in from_chnnls]
            #n.trace_preprocessed = [ch .- mean(n.trace) for ch in n.trace]
        end
        n.preprocessing_steps = [string( filter_type, ": ",freqs )];

    ######################
    ### preprocessed trace
    ######################
    elseif what == "trace_preprocessed" && !isempty(n.preprocessing_steps)
        nyq = n.sf_preprocessed/2;
        freqs_n = freqs./nyq; # normalized to nquist frequncy

        if filter_type == "Bandpass"
            rt = Bandpass( freqs_n[1], freqs_n[2] );
            dm = getfield( Main,Symbol(filter_design[1]) )(filter_order[1]);
            f = digitalfilter(rt, dm);
            for (itrc, trc) in enumerate(n.trace_preprocessed)
                # trc_padded = flatten([[.0 for ii in 1:pads], trc, [.0 for ii in 1:pads]])
                # n.trace_preprocessed[itrc] = filt(f, trc_padded)[pads+1:end-pads];
                n.trace_preprocessed[itrc] = filtfilt(f, trc);
            end
        elseif  filter_type == "Lowpass"
            rt = Lowpass( freqs_n[1] );
            dm = getfield( Main,Symbol(filter_design[1]) )(filter_order[1]);
            f = digitalfilter(rt, dm);
            for (itrc, trc) in enumerate(n.trace_preprocessed)
                # trc_padded = flatten([[.0 for ii in 1:pads], trc, [.0 for ii in 1:pads]])
                # n.trace_preprocessed[itrc] = filt(f, trc_padded)[pads+1:end-pads];
                n.trace_preprocessed[itrc] = filtfilt(f, trc);
            end
        elseif filter_type == "Highpass"
            rt = Highpass( freqs_n[1] );
            dm = getfield( Main,Symbol(filter_design[1]) )(filter_order[1]);
            f = digitalfilter(rt, dm);
            for (itrc, trc) in enumerate(n.trace_preprocessed)
                # trc_padded = flatten([[.0 for ii in 1:pads], trc, [.0 for ii in 1:pads]])
                # n.trace_preprocessed[itrc] = filt(f, trc_padded)[pads+1:end-pads];
                n.trace_preprocessed[itrc] = filtfilt(f, trc);
            end
        elseif filter_type == "Notch"
            for (itrc, trc) in enumerate(n.trace_preprocessed)
                trc = flatten(trc)
                for ff in freqs_n.*nyq
                    nn = iirnotch(ff, bandwidth[1]; fs=n.sf_preprocessed)
                    trc = filtfilt(nn, trc);
                end
                n.trace_preprocessed[itrc] = trc;
            end
        elseif filter_type == "CommonAverage"
            chnnls = freqs
            average_trace = [median([n.trace_preprocessed[ch][idx] for ch in chnnls]) for idx in 1:length(n.trace_preprocessed[1])]
            # n.trace_preprocessed = [channel .- average_trace for channel in n.trace_preprocessed]
			[n.trace_preprocessed[ch] = n.trace_preprocessed[ch] .- average_trace for ch in from_chnnls]
			#n.trace_preprocessed = [ch .- mean(n.trace_preprocessed) for ch in n.trace_preprocessed]
        end
        push!(n.preprocessing_steps, string( filter_type, ": ", freqs ));
    else
        print(n.preprocessing_steps)
        @warn "SELF.trace_preprocessed is empty. Process SELF.trace first"
    end
end

"""
Arguments:

+ n::Neurons.recording
+ chans::Array{Int64,1}
+ SNR::Array{Float64,1}
+ lockout::Int64=30;
+ polarity::String="negative",
+ ww::Array{Int64,1}=[5,15],
+ stp::Int64=1,
+ w_lockout::Array{Int64,1}=[10],
+ method::String="petrantonakis"

"""
function detect_spikes(
	n::Neurons.recording,
	chans::Array{Int64,1},
	SNR::Array{Float64,1},
	lockout::Int64=30;
	#
	polarity::String="negative",
	ww::Array{Int64,1}=[5,15],
	stp::Int64=1,
	w_lockout::Array{Int64,1}=[10],
	method::String="petrantonakis")
#
#   spikes are stored as:s
#   n.spike_idx: channel X w X Dw/indices of peaks (Dw only if method is "petrantonakis")
#
#   w: window length if method is "petrantonakis"
#   (for consistency if method is "threshold" the same structure is stored )

### WAVEFORM/THRESHOLD BASED SPIKE DETECITON
    if method == "petrantonakis" ### ref: https://www.frontiersin.org/articles/10.3389/fnins.2015.00452/full
        # Dw = [ [[] for ii in 1:n.n_channels] for jj in eachindex(ww) ] # ww X channel
        # Spks = [ [CartesianIndex{1}[] for ii in 1:n.n_channels] for jj in eachindex(ww) ] # ww X channel
        Dw_spk = [ [[[],[]] for ii in eachindex(ww)] for jj in 1:n.n_channels ] ### ch X w X Dw,spk_idx
		n.detection_threshold = []

        # for (ii,ch) in enumerate([n.trace_preprocessed[2]]) ### loop over channels
		for ii in chans ### loop over channels

			ch = n.trace_preprocessed[ii]

			for (jj,w) in enumerate(ww) ### loop over w

                ### compute the Dw signal (euklidean distance signal)
                f = ch[collect(0:w-1).+collect(1:stp:length(ch)-w+1)'] ### cut signal in overlapping chunks
				Dw_spk[ii][jj][1] = flatten(sqrt.(sum(diff(f, dims=2).^2, dims=1))); ### compute euclidean distance between consecutive chunks

                ### normalization of Dw
				sig = median( Dw_spk[ii][jj][1] ./ 0.6745 )
				thr = SNR[jj]*sig
				append!(n.detection_threshold, thr)
                # Dw_spk[ii][jj][1] = Dw_spk[ii][jj][1] ./ th ### store the standardized Dw signal
                ### find all local mixima in Dw
                loc_maxima = findlocalmaxima(Float64.(Dw_spk[ii][jj][1])) ### find all local maxima
                ### remove local mixima < thr
                Dw_spk[ii][jj][2] = loc_maxima[findall(Dw_spk[ii][jj][1][loc_maxima] .> thr)] ### extract all local maxima > SNR, SNR depends on w

                ### remove redundant triggered spikes and keep the maximum of each
                if length(Dw_spk[ii][jj][2]) > 1

					Dw_spk[ii][jj][2] = find_redundant( Dw_spk[ii][jj][1], Dw_spk[ii][jj][2], lockout );

				end

			end #for w: jj

            ### remove spikes that are detected only by one value of w
			loop_cntr = 0
			while length(Dw_spk[ii][1][2]) != length(Dw_spk[ii][2][2]) ### necessary if there a spikes in value of w with less spikes that are not in w with more spikes

				loop_cntr += 1

				if loop_cntr == 10 ### 2 should be enough

					break;

				end

				if !any( [(length(Dw_spk[ii][1][2])==0), (length(Dw_spk[ii][2][2])==0)] )

					if length(Dw_spk[ii][1][2]) > length(Dw_spk[ii][2][2])

						dd = pairwise(Euclidean(), cartesian2matrix(Dw_spk[ii][2][2])', cartesian2matrix(Dw_spk[ii][1][2])', dims=2)
					    deleteat!(Dw_spk[ii][1][2], findall(flatten(.!any(dd .< w_lockout, dims=1)))) ### delete all spikes in Dw_spk[ii][1][2] that do not have a partner spike in Dw_spk[ii][2][2]

					elseif length(Dw_spk[ii][1][2]) < length(Dw_spk[ii][2][2])

						dd = pairwise(Euclidean(), cartesian2matrix(Dw_spk[ii][1][2])', cartesian2matrix(Dw_spk[ii][2][2])', dims=2)
					    deleteat!(Dw_spk[ii][2][2], findall(flatten(.!any(dd .< w_lockout, dims=1))))

					end

				else ### if one has no spikes at all

					Dw_spk[ii][1][2] = [];
					Dw_spk[ii][2][2] = [];
					# Dw_spk[ii][1][1] = []
					# Dw_spk[ii][2][1] = []

				end

			end

			if length(Dw_spk[ii][1][2]) != length(Dw_spk[ii][2][2])

				@info "@Channel ", ii, ": different spike numbers for values of w"

			end

			### scatter plot of w1 X w2
	        if (length(spk_idx_w1)-length(spk_idx_w1) == 0) && length(recording.spike_idx[channel][1][2]) != 0

	            ax1 = subplot(1,3,1)
	            w1 = recording.spike_idx[channel][1][1][spk_idx_w1]
	            w2 = recording.spike_idx[channel][2][1][spk_idx_w2]
	            pols = cartesian2polar(w1, w2)
	            r = pols[:,1]
	            phi = pols[:,2]
	            ax1.scatter( r, phi, color=pcols[2])

	        end

        end #for channels: ii

    ### THRESHOLD BASED SPIKE DETECTION
    elseif method == "threshold"

		Dw_spk = [ [[[],[]] for ii in [1]] for jj in 1:n.n_channels ] ### ch X w X DW,spk_idx
		n.detection_threshold = ones(n.n_channels) .* NaN
        x = n.trace_preprocessed[chans]

        sig = [median(abs.(ch) ./ 0.6745 ) for ch in x]
		thr = SNR[1] .* sig
		n.detection_threshold[chans] = thr

        # x = .-x ./ th

        if polarity == "both"

			ind_max = findlocalmaxima.(x)
            ind_max = [[lm[1] for lm in ch] for ch in ind_max] # cartesian coordinates to indices
            ind_min = findlocalminima.(x)
            ind_min = [[lm[1] for lm in ch] for ch in ind_min]
            ind_max = [ind_max[ii][findall(x[ii][ind_max[ii]] .> thr[ii])] for ii in eachindex(chans)]    # local minima with minimal SNR
            ind_min = [ind_min[ii][findall(x[ii][ind_min[ii]] .< -thr[ii])] for ii in eachindex(chans)]    # local maxima with minimal SNR
            ind = [sort(flatten([ind_max[ii], ind_min[ii]])) for ii in 1:length(ind_max)] # channel X index_of_spikes

		elseif polarity == "positive"

			ind = findlocalmaxima.(x)
            ind = [[lm[1] for lm in ch] for ch in ind] # cartesian coordinates to indices
            ind = [ind[ii][findall(x[ii][ind[ii]] .> thr[ii])] for ii in eachindex(chans)]    # local minima with minimal SNR

		elseif polarity == "negative"

			ind = findlocalminima.(x)
            ind = [[lm[1] for lm in ch] for ch in ind] # cartesian coordinates to indices
            ind = [ind[ii][findall(x[ii][ind[ii]] .< -thr[ii])] for ii in eachindex(chans)]

		end

        ### remove redundantly triggered spikes and store the maximum respectively
        for (ii,ch) in enumerate(chans)

			chans_spike_idxs = ind[ii]
    		Dw_spk[ch][1][2] = CartesianIndex.( find_redundant( x[ii], chans_spike_idxs, lockout ) )

		end

    end # method

    ### store clean indices in struct
    # n.spike_idx = ind_clean#[ind_clean[ii][d[ii]] for ii in 1:length(ind_clean)]
    n.spike_idx = Dw_spk; # channel X w X Dw/index of peak (Dw only if method is "petrantonakis")
end

"""
Documentation ...
"""
function extract_isih(n::Neurons.recording)
### Extracts the inter-spike intervals
    dt = 1/n.sf_preprocessed
    isih = []
    for ch in  eachindex(n.spike_idx)
        push!(isih, [])
        for w in eachindex(n.spike_idx[1])
            push!(isih[ch], [])
            if length(n.spike_idx[ch][w][2]) > 1
                isih[ch][w] = cartesian2matrix(diff(n.spike_idx[ch][w][2])) .* dt
            end
        end
    end
    n.isih = isih # value_source X w X time_source X time/values
end

"""
Documentation ...
"""
function extract_spikes(n::Neurons.recording, mp, w::Int64=1)
### Extract spike spikeforms and time stamps in ms.
#   extracts the spikeforms at indices n.spike_idx (normalized around zero)
#   from the filtered signal x using a fixed window around the
#   window +-mp of the spikes.
#	ONLY FOR ONE VALUE OF w
#   structure: n.spikeforms[val_source][time_source][spike_time/value]
#

    dt = 1/n.sf_preprocessed
	window = collect(mp[1]:1:mp[2]);
	spkfrms = []

	for ii in 1:n.n_channels # loop over channels as value source

		push!(spkfrms, [])
	    # for w in 1:size(n.spike_idx[1],1)
	    #     push!(spkfrms[ii], [])

			for (jj,spk_idxs) in enumerate(n.spike_idx) # loop over channels for time source

				push!(spkfrms[ii], [[],[]]) ### t, s

				for spk in spk_idxs[w][2] # loop over individual spikes

				    # skip spikes that have negative indices
	                spike_window_idxs = spk.I .+ window # generate window for each spike
	                spike_window_times = spike_window_idxs .* dt .* 1000 # spike times in ms

					if spike_window_idxs[1] > 0 && spike_window_idxs[end] <= length(n.trace_preprocessed[ii])# skip spikes that are not complete in window

						push!( spkfrms[ii][jj][1], spike_window_times ) # store spike times
	                    push!( spkfrms[ii][jj][2], n.trace_preprocessed[ii][spike_window_idxs] .- mean( n.trace_preprocessed[ii][spike_window_idxs]) ) #store spike values

					end
	            end
	        end
	    # end
	end
	n.spikeforms = spkfrms # value_source X w X time_source X time/values

end

"""
Documentation ...
"""
function extract_features(n::Neurons.recording)
    n.features = []
    for ch in 1:n.n_channels
        push!(n.features, [])
        if !isempty(n.spikeforms[ch][ch][2])
            ### make matrix for PCA
            m = zeros(length(n.spikeforms[ch][ch][2]), length(n.spikeforms[ch][ch][2][1])); # #spikes X #samples per spike
            [m[ii,:] = n.spikeforms[ch][ch][2][ii] for ii = 1:length(n.spikeforms[ch][ch][2])]; # spike X sample
            ### fit PCA
            M = fit(PCA, m)
            n.features[ch] = projection(M) # projection(M): spike X PC
        end
    end
end

"""
Documentation ...
"""
function cluster_spikes(n::Neurons.recording, algorithm::AbstractArray, init_method::AbstractArray, feats::AbstractArray, n_clusters::AbstractArray, nIter_EM::AbstractArray, nIter_kmeans_init::AbstractArray, cov_kind::AbstractArray)

############
# cluster given features using _algorithm_
# BIC optimization for cluster number
#
########
# algorithm: kmeans
#
########
# algorithm: gmm
# init_method: :split or :kmeans

    n.spike_class = []
    ###
    if algorithm[1] == "kmeans"
        for ch in 1:n.n_channels
            #push!(n.spike_class, [])
            if size(n.features[ch], 2) > feats[end] # if there are enough features
                result = kmeans(n.features[ch][:,feats]', n_clusters[1]);
                clusters = []
                for cl = 1:n_clusters[1]
                    push!(clusters, findall(result.assignments .== cl))
                end
                push!(n.spike_class, clusters)
            else ### push empty array for that channel
                push!(n.spike_class, [])
            end
        end
    ###
    elseif algorithm[1] == "gmm"
        n.BIC = []
        ### for each channel fit gmm with n components specified in
        ### n_clusters, compute BIC fear each value, find minimum BIC and use
        ### for clustering
        for ch in 1:n.n_channels
            if size(n.features[ch],2) > length(feats) ### if enough features
                N = size(n.features[ch][:,feats], 1)
                BICS = Array{Float64,2}(undef, 0, 2)
                for cmpnts in n_clusters
                    if size(n.features[ch], 1) > cmpnts
                        M = GMM(cmpnts, n.features[ch][:,feats];
                            method=init_method[1], kind=cov_kind[1], nInit=nIter_kmeans_init[1], nIter=nIter_EM[1], nFinal=nIter_EM[1])
                        LL = sum(llpg(M, n.features[ch][:,feats]))
                        par = cmpnts + cmpnts*length(feats) + cmpnts* length(feats)*(length(feats)+1)/2 # number of parameters
                        bic = -2*LL + log.(N)*par
                        BICS = vcat(BICS, [cmpnts bic])
                    end
                end
                push!(n.BIC, BICS)
                ### fit model with optimal number of components and cluster spikes
                opt_cmpnts = Int(BICS[findmin(BICS[:,2])[2]])
                push!(n.spike_class, [])
                M = GMM(opt_cmpnts, n.features[ch][:,feats];
                    method=init_method[1], kind=cov_kind[1], nInit=nIter_kmeans_init[1], nIter=nIter_EM[1], nFinal=nIter_EM[1])
                posteriors = gmmposterior( M, n.features[ch][:,feats] )
                classes = [findmax(posteriors[1], dims=2)[2][spk][2] for spk in 1:length(findmax(posteriors[1], dims=2)[2])] ### extract cluster index for each spike from CartesianCoordinates
                n.spike_class[ch] = [[] for ii in 1:opt_cmpnts]
                for (spk_idx,cl) in enumerate(classes)
                    push!(n.spike_class[ch][cl], spk_idx)
                end
            else
                push!(n.spike_class, [])
                push!(n.BIC, [])
            end
        end
        # for ch in 1:n.n_channels
        #     print(ch)
        #     push!(n.spike_class, [])
        #     if size(n.features[ch],2) > length(feats) ### if enough features
        #         M = GMM(n_clusters[1], n.features[ch][:,feats];
        #             method=:kmeans, kind=cov_kind[1], nInit=nIter_kmeans_init[1], nIter=nIter_EM[1], nFinal=nIter_EM[1])
        #         posteriors = gmmposterior( M, n.features[ch][:,feats] )
        #         classes = [findmax(posteriors[1], dims=2)[2][spk][2] for spk in 1:length(findmax(posteriors[1], dims=2)[2])] ### extract cluster index for each spike from CartesianCoordinates
        #         n.spike_class[ch] = [[] for ii in 1:n_clusters[1]]
        #         for (spk_idx,cl) in enumerate(classes)
        #             push!(n.spike_class[ch][cl], spk_idx)
        #         end
        #     else
        #         n.spike_class[ch] = [[] for ii in 1:n_clusters[1]]
        #     end
        # end
    end
end

"""
Up- or downsampling
Acts on either n.trace or n.trace_preprocessed
"""
function resample(n::Neurons.recording, frequency::Real)
    #idx = 1:factor:size(n.trace)[1]
    #n.trace_preprocessed = n.trace[idx];
    #n.sf_decimated = n.sf/factor;
    if isempty(n.preprocessing_steps)
        n.trace_preprocessed = DSP.resample( n.trace,frequency//convert(Int64, n.sf) );
    else
        n.trace_preprocessed = DSP.resample( n.trace_preprocessed,frequency//convert(Int64, n.sf_preprocessed) );
    end
    n.sf_preprocessed = frequency;
    push!(n.preprocessing_steps, string( "Decimated to: ",frequency ))
end

"""
Up- or downsampling
Acts on either n.trace or n.trace_preprocessed
"""
function get_PSD(n::Neurons.recording, which, psd_type)
    if which == "trace"
        return [getfield( Main, Symbol(psd_type) )( flatten(ch), fs=n.sf ) for ch in n.trace]
    elseif which == "trace_preprocessed"
        return [getfield( Main, Symbol(psd_type) )( flatten(ch), fs=n.sf_preprocessed ) for ch in n.trace_preprocessed]
    end
end

function find_redundant(ref::AbstractArray, ind::AbstractArray, lockout::Int)
    ### ref: 1D array containing amplitude like values
    ### ind: 1D array containing indices of identified events, e.g. spikes
    ### lockout: Integer defining minimum distance of evvents in samples that are treated as separate events

    ### make sure it's cartesian!
    ind = CartesianIndex.(ind)
    ### matrix with all pairs of spike indices
    # mm = Iterators.product(flatten(cartesian2matrix(ind)), flatten(cartesian2matrix(ind)))
	### differences between all pairs of spike indices
	# dd = [ float(abs(diff([x for x in m])[1][1])) for m in mm ]
	dd = pairwise(Euclidean(), flatten(cartesian2matrix(ind))', flatten(cartesian2matrix(ind))', dims=2)
	dd[diagind(dd)] .= NaN # set diagonal to NaN
	sup_diag = dd[diagind(dd, 1)]
	gg = findall(x -> x > lockout, sup_diag )

	if all(sup_diag .> lockout) ### no double triggered spikes

		groups = [[spk_idx] for spk_idx in ind]

	elseif all(sup_diag .< lockout) ### only one double triggered spike

		groups = [ind]

	else ### single and double triggered spikes

		groups = []
	    current_group = []

		for (ij,spk_id) in enumerate(ind)

			push!(current_group, spk_id)

			if in(ij, gg) || ij == length(ind) ### ij == length(ind) required for last element

				push!(groups, current_group)
	            current_group = []

			end
		end
	end
	### find maxima of groups and store indices
    ind_clean = []
    for gr in groups

		max_id = findmax( abs.(ref[CartesianIndex.( gr )]) )
		# max_id = findmax(ref[gr])
    	push!(ind_clean, gr[max_id[2]])

	end

    #################
    ### DELETE ME ###a
    #################
    # ind_clean = []
    # ind_ = deepcopy(ind)
    # if length(ind_) > 1
    #     groups = []
    #     current_group = []
    #     push!(current_group, ind_[1])
    #     deleteat!(ind_, 1)
    #
    #     ### extract groups of detected event that are too close to each other to represent seperate spikes
    #     while !isempty(ind_)
    #         ### while next data point is to close to preceeding one, add it to current group
    #         while ind_[1][1]-current_group[end][1] < lockout ### double index because CartesianIndex
    #             push!(current_group, ind_[1])
    #             deleteat!(ind_, 1)
    #             if isempty(ind_) ### if already at the last spike index
    #                 break;
    #             end
    #         end
    #
    #         ### add current_group to group
    #         push!(groups, current_group)
    #         current_group = []
    #         if !isempty(ind_) ### deals with if last index was already deleted in while loop
    #             push!(current_group, ind_[1])
    #             deleteat!(ind_, 1)
    #         end
    #     end
    #
    #     ### for each group extract the maximum as representation
    #     ind_c = []
    #     for gr in groups
    #         if length(gr) > 1 ### if the group is a "group"
    #             gg = [gr[ii] for ii in eachindex(gr)]
    #             push!( ind_c, gg[findmax(ref[gg])[2]] )
    #         else ### if the group has only one entry
    #             push!(ind_c, gr[1])
    #         end
    #     end
    #     push!(ind_clean, ind_c)
    # else ### if no spikes in channel, add empty list to kepp size consistent
    #     push!(ind_clean, [])
    # end

    return ind_clean
end

################
### PLOTTING ###
################

function plot_detection(
        recording::Neurons.recording;
        channel::Int64,
        method::String="petrantonakis",
        polarity::String="negative")

	colors = myColors() # from MyHelpers.jl
	pcols = map(col -> (red(col), green(col), blue(col)), colors)

    if method == "petrantonakis"

        fig = figure(figsize=(13,5))
        rc("xtick",labelsize=8)
        rc("ytick",labelsize=8)

		### time course of Dw signal
        Dw_w1 = recording.spike_idx[channel][1][1]
        Dw_w2 = recording.spike_idx[channel][2][1]
		### indices of spikes
		spk_idx_w1 = recording.spike_idx[channel][1][2]
        spk_idx_w2 = recording.spike_idx[channel][2][2]
		### spike indices from cartesian to matrix
		if !isempty(spk_idx_w1) && !isempty(spk_idx_w2)
			spk_idx_w1 = cartesian2matrix(spk_idx_w1)
			spk_idx_w2 = cartesian2matrix(spk_idx_w2)
		end
        w1_threshold = recording.detection_threshold[1]
        w2_threshold = recording.detection_threshold[2]

        ### scatter plot of w1 X w2
        if (length(spk_idx_w1)-length(spk_idx_w1) == 0) && length(recording.spike_idx[channel][1][2]) != 0

            ax1 = subplot(1,3,1)
            w1 = recording.spike_idx[channel][1][1][spk_idx_w1]
            w2 = recording.spike_idx[channel][2][1][spk_idx_w2]
            pols = cartesian2polar(w1, w2)
            r = pols[:,1]
            phi = pols[:,2]
            ax1.scatter( r, phi, color=pcols[2])

        end

        ax2 = subplot(1,3,2)
        ax2.plot( Dw_w1, color=pcols[1], zorder=1 )
        ax2.plot( 1:length(Dw_w1), ones(length(Dw_w1))*w1_threshold, color=:red, linestyle="--" )
        if length(recording.spike_idx[channel][1][2]) > 0

            spk_idx_w1 = cartesian2matrix(recording.spike_idx[channel][1][2])
            ax2.scatter(spk_idx_w1, Dw_w1[spk_idx_w1], color=pcols[2], zorder=2)

        end

        ax3 = subplot(1,3,3)
        ax3.plot(Dw_w2, color=pcols[1], zorder=1)
        ax3.plot( 1:length(Dw_w2), ones(length(Dw_w2))*w2_threshold, color=:red, linestyle="--" )
		if length(recording.spike_idx[channel][2][2]) > 0

            spk_idx_w2 = cartesian2matrix(recording.spike_idx[channel][2][2])
            ax3.scatter(spk_idx_w2, Dw_w2[spk_idx_w2], color=pcols[2], zorder=2)

        end

    elseif method == "threshold"

        fig = figure(figsize=(13,5))
        rc("xtick",labelsize=8)
        rc("ytick",labelsize=8)

        x = recording.trace_preprocessed[channel]
        spk_idx = cartesian2matrix(recording.spike_idx[channel][1][2])

		if polarity == "negative"

			thr = -recording.detection_threshold[channel]

		elseif polarity == "positive"

			thr = recording.detection_threshold[channel]

		elseif polarity == "both"

			thr = [-recording.detection_threshold[channel], recording.detection_threshold[channel]]
        end

        plot( x )
        scatter( spk_idx, x[spk_idx], color=pcols[2], zorder=2 )

        if length(thr) != 2

			plot( 1:length(x), ones(length(x))*thr, color=:red, linestyle="--" )

		else

			plot( 1:length(x), ones(length(x))*thr[1], color=:red, linestyle="--" )
            plot( 1:length(x), ones(length(x))*thr[2], color=:red, linestyle="--" )
        end
    end
end

function plot_waveforms(
	recording::Neurons.recording,
	t::Array{Float64,1},
	ylims::Array{Int64,2},
	w::Int64=1;
	#
	t_source )
	##########################
	##########################
	##########################

	colors = myColors() # from MyHelpers.jl
	pcols = map(col -> (red(col), green(col), blue(col)), colors)
	half_window = abs(minimum(t))/1000

	fig = figure(figsize=(15,7), dpi=300)
	rc("xtick",labelsize=4)
	rc("ytick",labelsize=4)

	# ylims = [-1000 1000]

	if t_source == "val_source"
		corresponding = true
	else
		corresponding = false
	end

	for val_source = 1:recording.n_channels
	    #### t_cource determines whose channels spike times are used to cut out the the waveforms: if t_cource = val_source then channel and spiketimes are matched, i.e. 1<->1, 2<->2 etc
		if corresponding
			t_source = val_source
		end
	#     t_source = val_source

	    ### sort spikes according to amplitude
	    sorted_idx = sortperm([maximum(spk) for spk in recording.spikeforms[val_source][t_source][2]] .- [minimum(spk) for spk in recording.spikeforms[val_source][t_source][2]], rev=true )
	    ### how many spikes to plot
	    # n = length(sorted_idx)
	    n = 10

	    if !isempty(sorted_idx) ### if there are spikes in sorted_idx, i.e. if channel has spikes

	        spks_to_plot = []

	        ### loop over clusters of current channel
	        if !isempty(recording.spike_class)
	            for cl in recording.spike_class[t_source] ### plot the clusters based on the t_source channel
	                push!(spks_to_plot, intersect(cl, sorted_idx))
	            end
	        else
	            push!(spks_to_plot, sorted_idx) ### i.e. one "cluster" if not yet clustered
	        end

	        ### check if there are enough spikes if number of spikes to plot is specified
	        if length(sorted_idx) < n
	            n = length(sorted_idx)
	        end

	        ### plot spikes belonging to clusters
			offset = 0
	        for (ii,spk_cl) in enumerate(spks_to_plot)
	            if !isempty(spk_cl) ### cluster is not empty
	                subplot(2,5,val_source)
				   ### plot individual traces
				   plot(t, reduce(hcat, recording.spikeforms[val_source][t_source][2][spk_cl]),
					  color=pcols[ii], alpha=0.3, linewidth=0.1, zorder=1) ### plot individual spikes
				   ### plot mean waveform +- std
				   clusters_mean = mean(reduce(hcat, recording.spikeforms[val_source][t_source][2][spk_cl]), dims=2) .+ offset
				   clusters_std = std(reduce(hcat, recording.spikeforms[val_source][t_source][2][spk_cl]), dims=2)
				   plot(t, clusters_mean, color=pcols[ii], alpha=0.5, linewidth=0.5, zorder=2) ### plot mean spike waveform of cluster
				   fill_between(t, flatten(clusters_mean+clusters_std), flatten(clusters_mean-clusters_std), color=pcols[ii], alpha=0.2, linewidth= 0.1)
				   ###
				   title(string("Channel ", val_source, " (", n, ")"), fontsize=5)
				   ylim(ylims[1], ylims[2])
	# 			   legend()
	            end
				offset = offset +0
	        end

	        ### add x- and ylabel
	        if val_source in collect(6:10)
	            subplot(2,5,val_source)
	            xlabel("ms", fontsize=5)
	        elseif val_source in [1,6]
	            subplot(2,5,val_source)
	#             ylabel("amplitude")
	        end
			if val_source in collect(1:5)
			   xticks(-half_window*1000:1:half_window*1000, []) ### half_window comes from waveform extraction step
			end
			if val_source in [2:5 7:10]
			   yticks(ylims[1]:100:ylims[2], [])
			end

	    else ### plot empty axes of channel has no spikes
	        subplot(2,5,val_source)
	        plot(t, t,
	            linewidth=0)
	        title(string("Channel ", val_source, " (", n, ")"))
	        ### add x- and ylabel
	        ### add x- and ylabel
	        if val_source in collect(6:10)
	            subplot(2,5,val_source)
	            xlabel("ms", fontsize=5)
	        elseif val_source in [1,6]
	            subplot(2,5,val_source)
	#             ylabel("amplitude")
	        end
			if val_source in collect(1:5)
			   xticks(-4:2:4, [])
			end
			if val_source in [2:5 7:10]
			   yticks(ylims[1]:100:ylims[2], [])
			end

	    end

	end

end

###
end


module NAG
export
    nag_licence_query, nag_license_query,
    nag_complex_polygamma,
    nag_opt_read!,
    nag_opt_nlp!,
    nag_1d_quad_inf_1,
    nag_zero_cont_func_brent,
    nag_opt_lp!,
    last_nag_error

nag_licence_query() = ccall((:a00acc, :libnagc_nag), Cint, ()) == 1
const nag_license_query = nag_licence_query # US vs UK spelling
nag_license_query() || warn("Cannot acquire a NAG license.")

const NagInt = Int32
const NagComplex = Complex128

type NagError <: Exception
    code::Int
    name::ASCIIString
    message::UTF8String
end

Base.showerror(io::IO, e::NagError) =
    print(io, "NAG function \"$(e.name)\" [$(e.code)] – $(e.message)")

const NAG_ERROR = zeros(Uint8,544)
const nag_errors = Array(NagError)
nag_errors[] = NagError(0, "", "NO_ERROR")

cstr_to_array(p::Ptr{Uint8}, own::Bool = false) =
    pointer_to_array(p, int(ccall(:strlen, Csize_t, (Ptr{Uint8},), p)), own)

function error_handler(msg::Ptr{Uint8}, code::Cint, name::Ptr{Uint8})
    code == 0 && return
    msg = UTF8String(copy(cstr_to_array(msg)))
    name = ASCIIString(copy(cstr_to_array(name)))
    nag_errors[] = NagError(int(code), name, msg)
    throw(nag_errors[])
end
const ptr_error_handler = cfunction(error_handler, Void, (Ptr{Uint8}, Cint, Ptr{Uint8}))

function reset_nag_error()
    fill!(NAG_ERROR, 0)
    unsafe_store!(
        convert(Ptr{Ptr{Void}}, pointer(NAG_ERROR)) + 2*sizeof(Cint) + 512,
        ptr_error_handler
    )
end
reset_nag_error()

last_nag_error() = nag_errors[]

const objfunref = Array(Function)
const confunref = Array(Function)

function objfun_wrapper(
    n::NagInt, x::Ptr{Float64}, objf::Ptr{Float64}, g::Ptr{Float64}, comm::Ptr{NagInt}
)
    x = pointer_to_array(x,int(n),false)
    g = pointer_to_array(g,int(n),false)
    f = unsafe_load(comm)
    unsafe_store!(objf,objfunref[1](x,g,f))
    return nothing
end
const c_objfun = cfunction(objfun_wrapper, Void,
    (NagInt, Ptr{Float64}, Ptr{Float64}, Ptr{Float64}, Ptr{NagInt}))

function confun_wrapper(
    n::NagInt, ncnlin::NagInt, needc::Ptr{NagInt}, x::Ptr{Float64},
    conf::Ptr{Float64}, conjac::Ptr{Float64}, comm::Ptr{NagInt}
)
    needc  = pointer_to_array(needc,int(ncnlin),false)
    x      = pointer_to_array(x,int(n),false)
    conf   = pointer_to_array(conf,int(ncnlin),false)
    conjac = pointer_to_array(conjac,(int(n),int(ncnlin)),false)
    flag   = unsafe_load(comm)
    confunref[1](needc,x,conf,conjac,flag)
    return nothing
end
const c_confun = cfunction(confun_wrapper, Void,
    (NagInt, NagInt, Ptr{NagInt}, Ptr{Float64}, Ptr{Float64}, Ptr{Float64}, Ptr{NagInt}))

function nag_complex_polygamma(z::Number, k::Integer)
    ccall((:s14afc, :libnagc_nag), NagComplex,
          (NagComplex, NagInt, Ptr{Void}),
          z, k, NAG_ERROR)
end

# if Base.polygamma doesn't support complex args, use NAG
if !applicable(Base.polygamma, 1, 1+2im)
    Base.polygamma(k::Int, z::Complex) = nag_complex_polygamma(z, k)
    Base.digamma(z::Complex) = polygamma(0, z)
    Base.trigamma(z::Complex) = polygamma(1, z)
end

function nag_opt_read!(name::ByteString, optfile::ByteString, print::Bool = false)
    options = zeros(Uint8,1440)
    ccall((:e04xxc, :libnagc_nag), Void, (Ptr{Void},), options)
    ccall((:e04xyc, :libnagc_nag), Void,
          (Ptr{Uint8}, Ptr{Uint8}, Ptr{Void}, Cint, Ptr{Uint8}, Ptr{Void}),
          name, optfile, options, print, "stdout", NAG_ERROR)
    return options
end

function nag_opt_lp!(
    A         :: Matrix{Float64},
    bl        :: Vector{Float64},
    bu        :: Vector{Float64},
    c         :: Vector{Float64},
    x         :: Vector{Float64};
    optfile   :: ByteString = "",
    transpose :: Bool = true,
)
    # since NAG is row-major
    #   when transpose == true the rows are linear constraints
    #   when transpose == false the columns are linear constraints
    transpose && (A = A')

    n = length(c)
    tda, nclin = size(A)

    nclin > 0 && tda >= n   ||
        error("bad linear constraint dimensions (tda ≱ n): $tda ≱ $n")
    length(bl) == n + nclin ||
        error("wrong number of lower bounds: $(length(bl)) ≠ $(n + nclin)")
    length(bu) == n + nclin ||
        error("wrong number of upper bounds: $(length(bu)) ≠ $(n + nclin)")
    length(c) == n          ||
        error("wrong number of coefficients: $(length(c)) ≠ $n")
    length(x) == n          ||
        error("wrong data vector size: $(length(x)) ≠ $n")

    objf = zeros()
    options = isempty(optfile) ? C_NULL :
        convert(Ptr{Void}, nag_opt_read!("e04mfc", optfile))

    reset_nag_error()
    ccall((:e04mfc,:libnagc_nag), Void,
           (NagInt, NagInt, Ptr{Float64}, NagInt, Ptr{Float64}, Ptr{Float64},
            Ptr{Float64}, Ptr{Float64}, Ptr{Float64}, Ptr{Void}, Ptr{Void}, Ptr{Void}),
           n, nclin, A, tda, bl, bu, c, x, objf, options, C_NULL, NAG_ERROR)

    return x, objf[]
end

function nag_opt_nlp!(
    A  :: Matrix{Float64},
    bl :: Vector{Float64},
    bu :: Vector{Float64},
    x  :: Vector{Float64},
    objfun! :: Function,
    confun! :: Function,
    optfile :: ByteString = "",
    transpose :: Bool = true,
)
    # since NAG is row-major
    #   when transpose == true the rows are linear constraints
    #   when transpose == false the columns are linear constraints
    transpose && (A = A')

    # extract various dimensions
    n = length(x)
    tda, nclin = size(A)
    ncnlin = length(bl) - n - nclin

    # check for usage problems
    length(bl) == length(bu) ||
        error("bounds vectors must have matching length")
    lexcmp(bl,bu) <= 0       ||
        error("lower bounds cannot be greater than upper bounds")
    0 <= ncnlin              ||
        error("as many bounds as variables and linear constraints must be given")
    nclin == 0 || n <= tda   ||
        error("second dimension of linear coefficients too small")

    # allocate output variables
    objf = zeros()
    g = zeros(n)

    # save callbacks in globals
    objfunref[1] = objfun!
    confunref[1] = confun!

    options = isempty(optfile) ? C_NULL :
        convert(Ptr{Void}, nag_opt_read!("e04ucc", optfile))

    reset_nag_error()
    ccall((:e04ucc, :libnagc_nag), Void,
          (NagInt, NagInt, NagInt, Ptr{Float64},
           NagInt, Ptr{Float64}, Ptr{Float64},
           Ptr{Void}, Ptr{Void},
           Ptr{Float64}, Ptr{Float64}, Ptr{Float64},
           Ptr{Void}, Ptr{Void}, Ptr{Void}),
          n, nclin, ncnlin, A,
          tda, bl, bu,
          c_objfun, c_confun,
          x, objf, g,
          options, C_NULL, NAG_ERROR)

    return x, objf[], g
end

const quadfunref = Array(Function)
quad_fun_wrapper(x::Float64, comm::Ptr{NagInt}) = quadfunref[1](x)::Float64
const c_quadfun = cfunction(quad_fun_wrapper, Float64, (Float64, Ptr{NagInt}))

function nag_1d_quad_inf_1(
    f :: Function,
    boundinf :: Symbol = :Infinite,
    bound :: Float64 = 0.0;
    epsabs :: Float64 = 0.0,
    epsrel :: Float64 = sqrt(eps(Float64)),
    max_num_subint :: NagInt = int32(1e7),
)
    max_num_subint > 0 || error("max num subint must be > 0")
    boundinf_val =
        boundinf == :UpperSemiInfinite ? int32(1076 + 0) :
        boundinf == :LowerSemiInfinite ? int32(1076 + 1) :
        boundinf == :Infinite          ? int32(1076 + 2) :
        error("""
        invalid boundinf symbol: $boundinf
        must be one of: UpperSemiInfinite, LowerSemiInfinite, Infinite
        """)

    # allocate output variables
    result = zeros()
    abserr = zeros()
    qp = zeros(Uint8, 48)
    comm = zeros(Uint8, 8)

    quadfunref[] = f

    reset_nag_error()
    ccall((:d01smc, :libnagc_nag), Void,
        (Ptr{Void}, NagInt, Float64, Float64,
         Float64, NagInt, Ptr{Float64}, Ptr{Float64},
         Ptr{Void}, Ptr{Void}, Ptr{Void}),
        c_quadfun, boundinf_val, bound, epsabs,
        epsrel, max_num_subint, result, abserr,
        qp, comm, NAG_ERROR)

    return result[], abserr[]
end

const quadfunref_brent = Array(Function)
quad_fun_wrapper_brent(x::Float64, comm::Ptr{NagInt}) = quadfunref_brent[1](x)::Float64
const c_quadfun_brent = cfunction(quad_fun_wrapper_brent, Float64, (Float64, Ptr{NagInt}))

function nag_zero_cont_func_brent(
    a :: Float64,
    b :: Float64,
    f :: Function;
    eps :: Float64 = 1e-5,
    eta :: Float64 = 0.0,
)
    f(a)*f(b) <= 0 || error("f(a)*f(b) must be <= 0")
    quadfunref_brent[1] = f

    x = zeros()
    comm = zeros(Uint8, 8)

    reset_nag_error()
    ccall((:c05ayc, :libnagc_nag), Void,
        (Float64, Float64, Float64, Float64,
         Ptr{Void}, Ptr{Float64}, Ptr{Void}, Ptr{Void}),
        a, b, eps, eta, NAG.c_quadfun_brent, x, comm, NAG_ERROR)

    return x[]
end

end # module


function tabulated_properties()
    # Various properties. Taken from CoolProp.
    data = Dict{String, MolecularProperty}()
    data["R507A"] = MolecularProperty(0.0988592, 3704900, 343.765, 0.0002014492399, 0.286)
    data["RC318"] = MolecularProperty(0.2000312, 2777500, 388.38, 0.0003226451742, 0.355345)
    data["SES36"] = MolecularProperty(0.18485, 2849000, 450.7, 0.0003571428571, 0.352)
    data["SulfurDioxide"] = MolecularProperty(0.0640638, 7886587.6, 430.64, 0.0001220256254, 0.2561299362)
    data["SulfurHexafluoride"] = MolecularProperty(0.1460554192, 3754983, 318.7232, 0.0001967606348, 0.21)
    data["Toluene"] = MolecularProperty(0.09213842, 4126000, 591.75, 0.000315556958, 0.2657)
    data["trans-2-Butene"] = MolecularProperty(0.05610632, 4027300, 428.61, 0.0002373605507, 0.2100768344)
    data["Water"] = MolecularProperty(0.018015268, 22064000, 647.096, 5.594803727e-05, 0.3442920843)
    data["Xenon"] = MolecularProperty(0.131293, 5842000, 289.733, 0.000119047619, 0.00363)
    data["1-Butene"] = MolecularProperty(0.05610632, 4005100, 419.29, 0.0002358490566, 0.1918606474)
    data["Acetone"] = MolecularProperty(0.05807914, 4700000, 508.1, 0.0002127659574, 0.3071)
    data["Air"] = MolecularProperty(0.02896546, 3786000, 132.5306, 8.452513778e-05, 0.0335)
    data["Ammonia"] = MolecularProperty(0.01703026, 11333000, 405.4, 7.569004444e-05, 0.25601)
    data["Argon"] = MolecularProperty(0.039948, 4863000, 150.687, 7.458551158e-05, -0.00219)
    data["Benzene"] = MolecularProperty(0.0781118, 4894000, 562.02, 0.0002562788314, 0.2108369733)
    data["CarbonDioxide"] = MolecularProperty(0.0440098, 7377300, 304.1282, 9.411847707e-05, 0.22394)
    data["CarbonMonoxide"] = MolecularProperty(0.0280101, 3494000, 132.86, 9.216589862e-05, 0.0497)
    data["CarbonylSulfide"] = MolecularProperty(0.0600751, 6370000, 378.77, 0.0001349527665, 0.0978)
    data["cis-2-Butene"] = MolecularProperty(0.05610632, 4225500, 435.75, 0.0002356267672, 0.2023595859)
    data["CycloHexane"] = MolecularProperty(0.08415948, 4082400, 553.6, 0.0003101736973, 0.20926)
    data["Cyclopentane"] = MolecularProperty(0.0701329, 4571200, 511.72, 0.0002617801047, 0.2011125734)
    data["CycloPropane"] = MolecularProperty(0.042081, 5579700, 398.3, 0.0001627891671, 0.1305495664)
    data["D4"] = MolecularProperty(0.29661576, 1347215.356, 586.5, 0.0009587727709, 0.5981248268)
    data["D5"] = MolecularProperty(0.3707697, 1160000, 619.15, 0.001216000031, 0.658)
    data["D6"] = MolecularProperty(0.444924, 961000, 645.78, 0.001594162692, 0.736)
    data["Deuterium"] = MolecularProperty(0.0040282, 1679600, 38.34, 5.803830528e-05, -0.1362902741)
    data["Dichloroethane"] = MolecularProperty(0.098959, 5227585.103, 561.6, 0.0002309468822, 0.2685785009)
    data["DiethylEther"] = MolecularProperty(0.0741216, 3649016.897, 466.7, 0.0002807636367, 0.2819293312)
    data["DimethylCarbonate"] = MolecularProperty(0.0900779, 4908800, 557, 0.00025, 0.346)
    data["DimethylEther"] = MolecularProperty(0.04606844, 5336800, 400.378, 0.0001683501684, 0.196)
    data["Ethane"] = MolecularProperty(0.03006904, 4872200, 305.322, 0.0001458387816, 0.099)
    data["Ethanol"] = MolecularProperty(0.04606844, 6268000, 514.71, 0.0001686340641, 0.644)
    data["EthylBenzene"] = MolecularProperty(0.106165, 3622400, 617.12, 0.0003648282243, 0.304)
    data["Ethylene"] = MolecularProperty(0.02805376, 5041800, 282.35, 0.0001309454817, 0.0866)
    data["EthyleneOxide"] = MolecularProperty(0.04405256, 7304686.16, 468.92, 0.0001394700139, 0.210195489)
    data["Fluorine"] = MolecularProperty(0.03799681, 5172400, 144.414, 6.409023906e-05, 0.0449)
    data["HeavyWater"] = MolecularProperty(0.020027508, 21672051.48, 643.847, 5.625704971e-05, 0.3643005688)
    data["Helium"] = MolecularProperty(0.004002602, 227600, 5.1953, 5.515719801e-05, -0.385)
    data["HFE143m"] = MolecularProperty(0.10004, 3635000, 377.921, 0.0002151397849, 0.2888713657)
    data["Hydrogen"] = MolecularProperty(0.00201588, 1296400, 33.145, 6.448284756e-05, -0.219)
    data["HydrogenChloride"] = MolecularProperty(0.0364609, 8288160.904, 324.55, 8.474576271e-05, 0.1281892174)
    data["HydrogenSulfide"] = MolecularProperty(0.03408088, 9000000, 373.1, 9.813542689e-05, 0.1005)
    data["IsoButane"] = MolecularProperty(0.0581222, 3629000, 407.817, 0.0002577481153, 0.1835317832)
    data["IsoButene"] = MolecularProperty(0.05610632, 4009800, 418.09, 0.0002398081535, 0.1925934522)
    data["Isohexane"] = MolecularProperty(0.08617536, 3040000, 497.7, 0.0003683241252, 0.2797)
    data["Isopentane"] = MolecularProperty(0.07214878, 3378000, 460.35, 0.0003057169061, 0.2274)
    data["Krypton"] = MolecularProperty(0.083798, 5525000, 209.48, 9.216589862e-05, -0.00089)
    data["m-Xylene"] = MolecularProperty(0.106165, 3534600, 616.89, 0.0003752345216, 0.326)
    data["MD2M"] = MolecularProperty(0.310685, 1227000, 599.4, 0.001093300515, 0.668)
    data["MD3M"] = MolecularProperty(0.384839, 945000, 628.36, 0.001458154971, 0.722)
    data["MD4M"] = MolecularProperty(0.45899328, 877000, 653.2, 0.001650000016, 0.8246471473)
    data["MDM"] = MolecularProperty(0.23653146, 1410044.756, 564.09, 0.000921288245, 0.5280658491)
    data["Methane"] = MolecularProperty(0.0160428, 4599200, 190.564, 9.862781099e-05, 0.01142)
    data["Methanol"] = MolecularProperty(0.03204216, 8215850, 512.5, 0.0001173705495, 0.5720322)
    data["MethylLinoleate"] = MolecularProperty(0.29447206, 1341000, 799, 0.001237011381, 0.8054063871)
    data["MethylLinolenate"] = MolecularProperty(0.29245618, 1369000, 772, 0.001180219521, 1.142605259)
    data["MethylOleate"] = MolecularProperty(0.29648794, 1246000, 782, 0.001230239282, 0.90584936)
    data["MethylPalmitate"] = MolecularProperty(0.27045066, 1350000, 755, 0.001114827202, 0.910316178)
    data["MethylStearate"] = MolecularProperty(0.29850382, 1239000, 775, 0.001258970162, 1.01756)
    data["MM"] = MolecularProperty(0.16237752, 1939000, 518.75, 0.0006290000472, 0.418)
    data["n-Butane"] = MolecularProperty(0.0581222, 3796000, 425.125, 0.0002549219298, 0.2008100946)
    data["n-Decane"] = MolecularProperty(0.14228168, 2103000, 617.7, 0.0006097560976, 0.4884)
    data["n-Dodecane"] = MolecularProperty(0.17033484, 1817000, 658.1, 0.0007518796992, 0.5741822212)
    data["n-Heptane"] = MolecularProperty(0.100202, 2736000, 540.13, 0.0004319051724, 0.349)
    data["n-Hexane"] = MolecularProperty(0.08617536, 3034000, 507.82, 0.000369565829, 0.299)
    data["n-Nonane"] = MolecularProperty(0.1282551, 2281000, 594.55, 0.0005524861878, 0.4433)
    data["n-Octane"] = MolecularProperty(0.1142285, 2497000, 569.32, 0.0004862867146, 0.395)
    data["n-Pentane"] = MolecularProperty(0.07214878, 3370000, 469.7, 0.0003109861207, 0.251)
    data["n-Propane"] = MolecularProperty(0.04409562, 4251200, 369.89, 0.0002, 0.1521)
    data["n-Undecane"] = MolecularProperty(0.15630826, 1990400, 638.8, 0.00066010223, 0.5390371014)
    data["Neon"] = MolecularProperty(0.020179, 2680000, 44.4918, 4.187253999e-05, -0.03844929927)
    data["Neopentane"] = MolecularProperty(0.07214878, 3196000, 433.74, 0.0003058103976, 0.1961)
    data["Nitrogen"] = MolecularProperty(0.02801348, 3395800, 126.192, 8.941423556e-05, 0.0372)
    data["NitrousOxide"] = MolecularProperty(0.0440128, 7245000, 309.52, 9.737098345e-05, 0.1613)
    data["Novec649"] = MolecularProperty(0.3160438, 1869026.583, 441.81, 0.0005208333333, 0.4710235936)
    data["o-Xylene"] = MolecularProperty(0.106165, 3737500, 630.259, 0.0003725088471, 0.312)
    data["OrthoDeuterium"] = MolecularProperty(0.0040282, 1679600, 38.34, 5.803830528e-05, -0.1362902741)
    data["OrthoHydrogen"] = MolecularProperty(0.00201594, 1310650, 33.22, 6.474779953e-05, -0.219)
    data["Oxygen"] = MolecularProperty(0.0319988, 5043000, 154.581, 7.336757153e-05, 0.0222)
    data["p-Xylene"] = MolecularProperty(0.106165, 3531500, 616.168, 0.0003712062719, 0.324)
    data["ParaDeuterium"] = MolecularProperty(0.0040282, 1679600, 38.34, 5.803830528e-05, -0.1362902741)
    data["ParaHydrogen"] = MolecularProperty(0.00201588, 1285800, 32.938, 6.435834728e-05, -0.219)
    data["Propylene"] = MolecularProperty(0.04207974, 4555000, 364.211, 0.0001832508704, 0.146)
    data["Propyne"] = MolecularProperty(0.04006, 5626000, 402.38, 0.0001635769703, 0.204)
    data["R11"] = MolecularProperty(0.137368, 4394000, 471.06, 0.0002431292035, 0.1887506483)
    data["R113"] = MolecularProperty(0.187375, 3392200, 487.21, 0.0003345982143, 0.252535)
    data["R114"] = MolecularProperty(0.170921, 3257000, 418.83, 0.0002947070612, 0.2523)
    data["R115"] = MolecularProperty(0.154466416, 3129036.968, 353.1, 0.0002512562814, 0.2484349931)
    data["R116"] = MolecularProperty(0.13801182, 3048000, 293.03, 0.0002250225023, 0.2566)
    data["R12"] = MolecularProperty(0.120913, 4136100, 385.12, 0.0002140053097, 0.1794783173)
    data["R123"] = MolecularProperty(0.152931, 3672000, 456.831, 0.0002780545193, 0.281922497)
    data["R1233zd(E)"] = MolecularProperty(0.1304944, 3623637.776, 439.6, 0.0002717391304, 0.3024522936)
    data["R1234yf"] = MolecularProperty(0.1140415928, 3382200, 367.85, 0.0002398081535, 0.276)
    data["R1234ze(E)"] = MolecularProperty(0.1140415928, 3636250, 382.52, 0.0002331002331, 0.313)
    data["R1234ze(Z)"] = MolecularProperty(0.1140415928, 3533000, 423.27, 0.0002426416868, 0.3274)
    data["R124"] = MolecularProperty(0.1364762, 3624295, 395.425, 0.0002437075, 0.2880950842)
    data["R125"] = MolecularProperty(0.1200214, 3617700, 339.173, 0.0002092487968, 0.3052)
    data["R13"] = MolecularProperty(0.104459, 3879000, 301.88, 0.0001792114695, 0.1745863278)
    data["R134a"] = MolecularProperty(0.102032, 4059280, 374.21, 0.0001993201985, 0.32684)
    data["R13I1"] = MolecularProperty(0.1959104, 3952566.072, 396.44, 0.000225703065, 0.1761811178)
    data["R14"] = MolecularProperty(0.0880046, 3750000, 227.51, 0.0001406584622, 0.1785)
    data["R141b"] = MolecularProperty(0.11694962, 4212000, 477.5, 0.0002550369804, 0.2195)
    data["R142b"] = MolecularProperty(0.10049503, 4055000, 410.26, 0.0002253267237, 0.2321)
    data["R143a"] = MolecularProperty(0.084041, 3761000, 345.857, 0.0001949906892, 0.2614896462)
    data["R152A"] = MolecularProperty(0.066051, 4520000, 386.411, 0.000179486413, 0.2752171145)
    data["R161"] = MolecularProperty(0.0480595, 5010000, 375.25, 0.0001592356688, 0.2162428411)
    data["R21"] = MolecularProperty(0.1029227, 5181200, 451.48, 0.0001956654009, 0.2061)
    data["R218"] = MolecularProperty(0.18801933, 2640000, 345.02, 0.0002994011976, 0.3172)
    data["R22"] = MolecularProperty(0.086468, 4990000, 369.295, 0.0001650649861, 0.22082)
    data["R227EA"] = MolecularProperty(0.17002886, 2925242.234, 374.9, 0.0002861230329, 0.357641242)
    data["R23"] = MolecularProperty(0.07001385, 4832000, 299.293, 0.0001329787234, 0.2629648925)
    data["R236EA"] = MolecularProperty(0.1520384, 3420000, 412.44, 0.0002691065662, 0.3687823804)
    data["R236FA"] = MolecularProperty(0.1520384, 3200000, 398.07, 0.0002757859901, 0.3769117976)
    data["R245ca"] = MolecularProperty(0.13404794, 3940740.489, 447.57, 0.0002551020408, 0.3545654191)
    data["R245fa"] = MolecularProperty(0.13404794, 3651000, 427.01, 0.0002580645161, 0.3776)
    data["R32"] = MolecularProperty(0.052024, 5782000, 351.255, 0.0001226981129, 0.2769)
    data["R365MFC"] = MolecularProperty(0.14807452, 3266207.068, 460, 0.0003125, 0.3774450972)
    data["R40"] = MolecularProperty(0.05048752, 6671686.977, 416.3, 0.0001390002242, 0.1500684964)
    data["R404A"] = MolecularProperty(0.0976038, 3734800, 345.27, 0.0002024291498, 0.293)
    data["R407C"] = MolecularProperty(0.0862036, 4631700, 359.345, 0.0001901140668, 0.363)
    data["R41"] = MolecularProperty(0.03403292, 5897000, 317.28, 0.0001075268817, 0.2004)
    data["R410A"] = MolecularProperty(0.0725854, 4901200, 344.494, 0.0001581277672, 0.296)
    return data
end

"""
    cubic_benchmark(name)

Get a benchmark equation-of-state instance together with some data for testing.
"""
function cubic_benchmark(name)
    psi = 6.894757293168360e+03
    Rankine = 0.555555555555556
    bar = 1e5
    data = Dict()
    eos_type = PengRobinson()

    bic = nothing
    if name == "spe5"
        T_c = [343, 665.7, 913.4, 1111.8, 1270.0, 1380.0]*Rankine;
        p_c = [667.8, 616.3, 436.9, 304.0, 200.0, 162.0]*psi;
        mw = [16.040, 44.100, 86.180, 142.290, 206.000, 282.000]/1000;
        acc = [0.0130, 0.1524, 0.3007, 0.4885, 0.6500, 0.8500];
        Z_c = [0.290, 0.277, 0.264, 0.257, 0.245, 0.235];
        # Convert critical volumes to critical densities
        V_c = Z_c.*8.314.*T_c./p_c;
        names = ["C1", "C3", "C6", "C10", "C15", "C20"];
        ncomp = length(names);
        bic = zeros(ncomp, ncomp);
        bic[1, 5] = 0.05;
        bic[1, 6] = 0.05;
        bic[2, 5] = 0.005;
        bic[2, 6] = 0.005;
        bic = Symmetric(bic)
        props = MolecularProperty.(mw, p_c, T_c, V_c, acc)

        z0 = [0.77, 0.20, 0.03, 0, 0, 0]
        z = [0.5, 0.03, 0.07, 0.20, 0.15, 0.05]
        p = 4000*psi
        T = 344.26
        p0 = p
        T0 = T

        initial_cond = (z = z0, p = p0, T = T0)
        displ_cond = (z = z, p = p, T = T)
        data["initial_condition"] = initial_cond
        data["injection_condition"] = displ_cond
    elseif name == "simple"
        names = ["Methane", "CarbonDioxide", "n-Decane"]
        props = MolecularProperty.(names)

        p = 75*bar
        T = 273.15 + 150
        z0 = [0.3, 0.1, 0.6]
        z = [0.1, 0.9, 0.0]
        initial_cond = (z = z0, p = p, T = T)
        displ_cond = (z = z, p = p, T = T)
        data["initial_condition"] = initial_cond
        data["injection_condition"] = displ_cond
    else
        error("Unknown benchmark $name")
    end
    mixture = MultiComponentMixture(props, names = names, A_ij = bic)
    eos = GenericCubicEOS(mixture, eos_type)

    return (eos, data)
end


# include("Move.jl")
# include("Wavefunction.jl")
# include("SwapWavefunction.jl")

abstract Walk

is_nontrivial_transition(walk::Walk) = walk.transition_is_nontrivial

function get_probability_ratio(walk::Walk)
    if !walk.transition_is_nontrivial
        warn("Asking for the probability ratio ( = 1) of a trivial transition.")
        return 1.
    else
        return walk.probability_ratio
    end
end

# Lots of code copying below since Julia does't really have abstract base classes.  There has to be a better way..
# I could have one class with finish_move! and cancel_move! functions passed to the constructor and held as members,
# but that seems less safe and less explicit.

# standard walk for ordinary wavefunctions / measurements
type StandardWalk <: Walk
    wf::Wavefunction
    transition_function::Function
    probability_ratio::Float64
    transition_is_nontrivial::Bool
    # transition_in_progress::Bool
    StandardWalk(wf::Wavefunction, transition_function::Function) = new(wf, transition_function)
end

function propose_and_assess_random_transition!(walk::StandardWalk)
    move = walk.transition_function(walk.wf)

    if !is_nontrivial(move)
        walk.transition_is_nontrivial = false
    else
        probrat = assess_move!(walk.wf, move)
        walk.transition_is_nontrivial = true
        walk.probability_ratio = probrat
    end
end

accept_transition!(walk::StandardWalk) = finish_move!(walk.wf)
reject_transition!(walk::StandardWalk) = cancel_move!(walk.wf)

check_for_numerical_error!(walk::StandardWalk) = check_for_numerical_error!(walk.wf)

# walk for Renyi mod measurement
type RenyiModWalk <: Walk
    wf::Union{SwapWavefunction, SwapWavefunctionNS}
    transition_function::Function
    probability_ratio::Float64
    transition_is_nontrivial::Bool
    # transition_in_progress::Bool
    RenyiModWalk(wf::Union{SwapWavefunction, SwapWavefunctionNS}, transition_function::Function) = new(wf, transition_function)
end

function propose_and_assess_random_transition!(walk::RenyiModWalk)
    move = walk.transition_function(walk.wf)
    # move here is either a 2-tuple of moves for the different copies (SwapWavefunction) or a Tuple{ParticlesMove{LatticeParticleMove},Int64} (SwapWavefunctionNS),
    # where the Int here is the copy_to_move
    @assert length(move) == 2

    if isa(move[2], Int) # for SwapWavefunctionNS
        trivial_condition = !is_nontrivial(move[1])
    else # for SwapWavefunction
        trivial_condition = all(map(!, map(is_nontrivial, move)))
    end

    if trivial_condition
        walk.transition_is_nontrivial = false
    else
        probrat = assess_move_for_mod!(walk.wf, move...)
        walk.transition_is_nontrivial = true
        walk.probability_ratio = probrat
    end
end

accept_transition!(walk::RenyiModWalk) = finish_move_for_mod!(walk.wf)
reject_transition!(walk::RenyiModWalk) = cancel_move_for_mod!(walk.wf)

check_for_numerical_error!(walk::RenyiModWalk) = check_for_numerical_error_for_mod!(walk.wf)

# walk for Renyi sign measurement
type RenyiSignWalk <: Walk
    wf::Union{SwapWavefunction, SwapWavefunctionNS}
    transition_function::Function
    probability_ratio::Float64
    transition_is_nontrivial::Bool
    # transition_in_progress::Bool
    RenyiSignWalk(wf::Union{SwapWavefunction, SwapWavefunctionNS}, transition_function::Function) = new(wf, transition_function)
end

function propose_and_assess_random_transition!(walk::RenyiSignWalk)
    move = walk.transition_function(walk.wf)
    @assert length(move) == 2 # a 2-tuple of moves for the different copies

    if isa(move[2], Int) # for SwapWavefunctionNS
        trivial_condition = !is_nontrivial(move[1])
    else # for SwapWavefunction
        trivial_condition = all(map(!, map(is_nontrivial, move)))
    end

    if trivial_condition
        walk.transition_is_nontrivial = false
    else
        probrat = assess_move_for_sign!(walk.wf, move...)
        walk.transition_is_nontrivial = true
        walk.probability_ratio = probrat
    end
end

accept_transition!(walk::RenyiSignWalk) = finish_move_for_sign!(walk.wf)
reject_transition!(walk::RenyiSignWalk) = cancel_move_for_sign!(walk.wf)

check_for_numerical_error!(walk::RenyiSignWalk) = check_for_numerical_error_for_sign!(walk.wf)

# type RenyiWalk <: Walk
#     wf::SwapWavefunction
#     transition_function::Function
#     probability_ratio::Float64
#     transition_is_nontrivial::Bool
#     # transition_in_progress::Bool
#     RenyiWalk(wf::SwapWavefunction, transition_function::Function) = new(wf, transition_function)
# end

# type RenyiModWalk
#     walk::RenyiWalk
#     RenyiModWalk(wf::SwapWavefunction, transition_function::Function) = new(RenyiWalk(wf, transition_function))
# end

# How to then make, e.g., RenyiModWalk.wf invoke RenyiModWalk.walk.wf?


"""
Subtype of `LDScoreRegression`.
"""

mutable struct Hsq <: LDScoreRegression
    y; x; w; N; M;
    n_blocks; intercept; slow; step1_ii; old_weights;
    __null_intercept__

    constrain_intercept; n_annot; intercept_se; twostep_filtered;

    function Hsq(
        y, x, w, N, M; n_blocks = 200, slow = false, old_weights = false,
        intercept = nothing, step1_ii = nothing,
    )
        hsq = new(
            y, x, w, N, M, n_blocks, intercept, slow, step1_ii, old_weights, 1,
        )
        return ld_score_regression(
            hsq,
            y, x, w, N, M, n_blocks, intercept, slow, step1_ii, old_weights,
        )
    end
end

function hsq_weights(ld, w_ld, N, M, hsq; intercept=nothing)
    if intercept == nothing intercept = 1.0 end
    M = M[1][1]
    hsq = max(maximum(hsq), 0.0)
    hsq = min(hsq, 1.0)
    ld = fmax(ld, 1.0)
    w_ld = fmax(w_ld, 1.0)
    c = hsq .* N ./ M

    het_w = 1.0 ./ (2 .* (c .* ld .+ intercept).^2)
    oc_w = 1.0 ./ w_ld
    w = het_w .* oc_w
    return w
end

function update_weights(reg::Hsq, ld, w_ld, N, M, hsq, intercept)
    if intercept == nothing intercept = reg.__null_intercept__ end

    return hsq_weights(ld, w_ld, N, M, hsq; intercept=intercept)
end

function overlap_output(
    hsq::Hsq,
    category_names, overlap_matrix, M_annot, M_tot, print_coefficients,
)
    overlap_matrix_prop = zeros([hsq.n_annot, hsq.n_annot])
    for i in range(hsq.n_annot)
        overlap_matrix_prop[i, :] = overlap_matrix[i, :] / M_annot
    end

    # TODO missing code here

    #=
    df = DataFrames.DataFrame({
        "Category": category_names,
        "Prop._SNPs": one_d_convert(prop_M_overlap),
        "Prop._h2": one_d_convert(prop_hsq_overlap),
        "Prop._h2_std_error": one_d_convert(prop_hsq_overlap_se),
        "Enrichment": one_d_convert(enrichment),
        "Enrichment_std_error": one_d_convert(enrichment_se),
        "Enrichment_p":diff_p,
        "Coefficient": one_d_convert(hsq.coef),
        "Coefficient_std_error": hsq.coef_se,
        "Coefficient_z-score": one_d_convert(hsq.coef) / one_d_convert(hsq.coef_se)
    })
    if print_coefficients
        return df[
            !,
            [
                "Category", "Prop._SNPs", "Prop._h2", "Prop._h2_std_error",
                "Enrichment","Enrichment_std_error", "Enrichment_p",
                "Coefficient", "Coefficient_std_error","Coefficient_z-score",
            ],
        ]
    else
        return df[
            !,
            [
                "Category", "Prop._SNPs", "Prop._h2", "Prop._h2_std_error",
                "Enrichment","Enrichment_std_error", "Enrichment_p",
            ],
        ]
    end
    =#
end


#### thermodynamics

export new_thermo_state_anelastic, recover_thermo_state_anelastic

"""
    new_thermo_state_anelastic(atmos::AtmosModel, state::Vars, aux::Vars)

Create a new thermodynamic state, based on the `state`, and _not_
the `aux` state.

!!! note
    This method calls the iterative saturation adjustment
    procedure for EquilMoist models.
"""
new_thermo_state_anelastic(atmos::AtmosModel, state::Vars, aux::Vars) =
    new_thermo_state_anelastic(
        atmos,
        energy_model(atmos),
        moisture_model(atmos),
        state,
        aux,
    )

"""
    recover_thermo_state_anelastic(atmos::AtmosModel, state::Vars, aux::Vars)

An atmospheric thermodynamic state.

!!! warn
    For now, we are directly calling new_thermo_state_anelastic to avoid
    inconsistent aux states in kernels where the aux states are
    out of sync with the boundary state.

# TODO: Define/call `recover_thermo_state_anelastic` when it's safely implemented
  (see https://github.com/CliMA/ClimateMachine.jl/issues/1648)
"""
recover_thermo_state_anelastic(atmos::AtmosModel, state::Vars, aux::Vars) =
    new_thermo_state_anelastic(
        atmos,
        energy_model(atmos),
        moisture_model(atmos),
        state,
        aux,
    )

function new_thermo_state_anelastic(
    atmos::AtmosModel,
    energy::TotalEnergyModel,
    moist::DryModel,
    state::Vars,
    aux::Vars,
)
    param_set = parameter_set(atmos)
    e_int = internal_energy(atmos, state, aux)
    p = aux.ref_state.p
    return PhaseDry_pe(param_set, p, e_int)
end

function new_thermo_state_anelastic(
    atmos::AtmosModel,
    energy::TotalEnergyModel,
    moist::EquilMoist,
    state::Vars,
    aux::Vars,
)
    param_set = parameter_set(atmos)
    e_int = internal_energy(atmos, state, aux)
    p = aux.ref_state.p
    ρ = density(atmos, state, aux)
    return PhaseEquil_peq(
        param_set,
        p,
        e_int,
        state.moisture.ρq_tot / ρ,
        moist.maxiter,
        moist.tolerance,
    )
end

function new_thermo_state_anelastic(
    atmos::AtmosModel,
    energy::TotalEnergyModel,
    moist::NonEquilMoist,
    state::Vars,
    aux::Vars,
)
    param_set = parameter_set(atmos)
    e_int = internal_energy(atmos, state, aux)
    p = aux.ref_state.p
    ρ = density(atmos, state, aux)
    q = PhasePartition(
        state.moisture.ρq_tot / ρ,
        state.moisture.ρq_liq / ρ,
        state.moisture.ρq_ice / ρ,
    )

    return PhaseNonEquil_peq(param_set, p, e_int, q)
end