Base.:(==)(i1::CompositeInterpolation, i2::CompositeInterpolation) = i1.before_first == i2.before_first && i1.inner == i2.inner && i1.after_last == i2.after_last
Base.hash(i::CompositeInterpolation) = 1 + hash(i.before_first) + hash(i.inner) + hash(i.after_last)

# Curve methods implementation

function _zero_rate(method::METHOD, curve::AbstractIRCurve, maturity::T) where {METHOD<:RateInterpolation, T<:Union{Date, YearFraction}}
    _zero_rate(method, curve_get_dtm(curve), curve_get_zero_rates(curve), days_to_maturity(curve, maturity))
end

function _zero_rate(comp::CompositeInterpolation, curve::AbstractIRCurve, maturity::T) where {T<:Union{Date, YearFraction}}
    dtm = days_to_maturity(curve, maturity)

    if dtm < curve_get_dtm(curve)[1]
        return _zero_rate(comp.before_first, curve, maturity)
    elseif dtm > curve_get_dtm(curve)[end]
        return _zero_rate(comp.after_last, curve, maturity)
    else
        return _zero_rate(comp.inner, curve, maturity)
    end
end

# Returns tuple (index_a, index_b) for input vector x
# for interpolands on linear interpolation on point x_out
function _interpolationpoints(x::Vector{T}, x_out::Number) where {T<:Number}
    local index_a::Int
    local index_b::Int

    if x_out <= x[1]
        # Interpolation point is before first vertice
        # Slope will be determined by the 1st and 2nd vertices
        index_a = 1
        index_b = 2
    elseif x_out >= x[end]
        # Interpolation point is after last vertice
        # Slope will be determined by the last and last-1 vertices
        index_b = length(x)
        index_a = index_b - 1
    else
        # Inner point
        index_a =  findlast(a -> a < x_out, x) # last element before x_out on x
        index_b = findfirst(b -> b >=  x_out, x) # first element after x_out on x
    end

    return index_a, index_b
end

# Perform Linear interpolation. Slope is determined by points (Xa, Ya) and (Xb, Yb).
# Interpolation occurs on point (x_out, returnvalue)
_linearinterp(Xa::Number, Ya::Number, Xb::Number, Yb::Number, x_out::Number) = (x_out - Xa) * (Yb - Ya) / (Xb - Xa) + Ya

# Linear interpolation of zero_rates
function _zero_rate(::Linear, x::Vector{N1}, y::Vector{N2}, x_out::Number) where {N1<:Number, N2<:Number}
    # If this curve has only 1 vertice, this will be a flat curve
    if length(x) == 1
        return y[1]
    end

    index_a, index_b = _interpolationpoints(x, x_out)
    return _linearinterp(x[index_a], y[index_a], x[index_b], y[index_b], x_out)
end

# Step Function
function _zero_rate(::StepFunction, x::Vector{N1}, y::Vector{N2}, x_out::Number) where {N1<:Number, N2<:Number}
    # If this curve has only 1 vertice, this will be a flat curve
    if length(x) == 1
        return y[1]
    end

    if x_out <= x[1]
        # Interpolation point is before first vertice
        # The result will be extrapolated using the first vertice zero_rate
        return y[1]
    elseif x_out >= x[end]
        # Interpolation point is after last vertice
        # The result will be extrapolated using the last vertice zero_rate
        return y[end]
    else
        # Inner point
        return y[findlast(a -> a <= x_out, x)] # last element before x_out on x
    end
end

# Flat Forward is linear interpolation on the log of discountfactors
# Maybe not useful for SimpleCompounding curves.
function _zero_rate(::FlatForward, curve::AbstractIRCurve, maturity::T) where {T<:Union{Date, YearFraction}}

    # If this curve has only 1 vertice, this will be a flat curve
    if length(curve_get_zero_rates(curve)) == 1
        return curve_get_zero_rates(curve)[1]
    end

    x_out = days_to_maturity(curve, maturity)
    curve_dtm = curve_get_dtm(curve)
    curve_zero_rates = curve_get_zero_rates(curve)
    index_a, index_b = _interpolationpoints(curve_dtm, x_out)
    Xa = curve_dtm[index_a]
    Ya = curve_zero_rates[index_a]
    Xb = curve_dtm[index_b]
    Yb = curve_zero_rates[index_b]

    year_fraction_a = yearfraction(curve_get_daycount(curve), Xa)
    logPa = log(discountfactor(curve_get_compounding(curve), Ya, year_fraction_a))

    year_fraction_b = yearfraction(curve_get_daycount(curve), Xb)
    logPb = log(discountfactor(curve_get_compounding(curve), Yb, year_fraction_b))

    year_fraction_x = yearfraction(curve_get_daycount(curve), x_out)
    logPx = _linearinterp(value(year_fraction_a), logPa, value(year_fraction_b), logPb, value(year_fraction_x))

    return discountfactor_to_rate(curve_get_compounding(curve), exp(logPx), year_fraction_x)
end

function _zero_rate(::NelsonSiegel, curve::AbstractIRCurve, maturity::YearFraction)

    # beta1 = param[1]
    # beta2 = param[2]
    # beta3 = param[3]
    # lambda = param[4]

    param = curve_get_model_parameters(curve)
    t = value(maturity)
    _exp_lambda_t_ = exp(-param[4]*t)
    F_beta2 = (1.0 - _exp_lambda_t_) / (param[4]*t)

    return param[1] + param[2]*F_beta2 + param[3]*(F_beta2 - _exp_lambda_t_)
end

_zero_rate(method::NelsonSiegel, curve::AbstractIRCurve, maturity::Date) = _zero_rate(method, curve, yearfraction(curve, maturity))

function _zero_rate(::Svensson, curve::AbstractIRCurve, maturity::YearFraction)

    # beta1 = param[1]
    # beta2 = param[2]
    # beta3 = param[3]
    # beta4 = param[4]
    # lambda1 = param[5]
    # lambda2 = param[6]

    param = curve_get_model_parameters(curve)
    t = value(maturity)
    _exp_lambda1_t_ = exp(-param[5]*t)
    _exp_lambda2_t_ = exp(-param[6]*t)
    F_beta2 = (1.0 - _exp_lambda1_t_) / (param[5]*t)

    return param[1] + param[2]*F_beta2 + param[3]*(F_beta2 - _exp_lambda1_t_) +
            param[4]*( (1.0 - _exp_lambda2_t_)/(param[6]*t) - _exp_lambda2_t_)
end

_zero_rate(method::Svensson, curve::AbstractIRCurve, maturity::Date) = _zero_rate(method, curve, yearfraction(curve, maturity))

function _zero_rate(method::CubicSplineOnRates, curve::AbstractIRCurve, maturity::T) where {T<:Union{Date, YearFraction}}
    sp = curve_get_spline_fit_on_rates(curve)
    return splineint(sp, days_to_maturity(curve, maturity))
end

function _zero_rate(::CubicSplineOnRates, curve::AbstractIRCurve, maturity_vec::Vector{Date})
    sp = curve_get_spline_fit_on_rates(curve)

    l = length(maturity_vec)
    rates = Vector{Float64}(undef, l)

    for i in 1:l
        @inbounds dtm = days_to_maturity(curve, maturity_vec[i])
        @inbounds rates[i] = splineint(sp, dtm)
    end
    return rates
end

function _zero_rate(::CubicSplineOnDiscountFactors, curve::AbstractIRCurve, maturity::YearFraction)
    sp = curve_get_spline_fit_on_discount_factors(curve)
    result_discount_factor = splineint(sp, value(maturity))
    return discountfactor_to_rate(curve_get_compounding(curve), result_discount_factor, maturity)
end

function _zero_rate(method::CubicSplineOnDiscountFactors, curve::AbstractIRCurve, maturity::Date)
    return _zero_rate(method, curve, yearfraction(curve, maturity))
end

# to make splineint work in the following method _zero_rate(::CubicSplineOnDiscountFactors, ::AbstractIRCurve, ::Vector{Date})
function splineint(s::Spline{T1}, x_out::YearFraction{T2}) where {T1<:Number, T2<:Number}
    splineint(s, value(x_out))
end

function _zero_rate(::CubicSplineOnDiscountFactors, curve::AbstractIRCurve, maturity_vec::Vector{Date})
    sp = curve_get_spline_fit_on_discount_factors(curve)
    mat_vec_len = length(maturity_vec)

    yf_maturity_vec = Vector{YearFraction}(undef, mat_vec_len)
    for i in 1:mat_vec_len
        @inbounds yf_maturity_vec[i] = yearfraction(curve_get_daycount(curve), curve_get_date(curve), maturity_vec[i])
    end

    return discountfactor_to_rate(curve_get_compounding(curve), splineint(sp, yf_maturity_vec), yf_maturity_vec)
end

# Generate vector functions
for elty in (:FlatForward, :CompositeInterpolation, :StepFunction, :Linear, :NelsonSiegel, :Svensson)
    @eval begin
        function _zero_rate(m::$elty, curve::AbstractIRCurve, maturity_vec::Vector{Date})
            l = length(maturity_vec)
            rates = Vector{Float64}(undef, l)
            for i = 1:l
                @inbounds rates[i] = _zero_rate(m, curve, maturity_vec[i])
            end
        return rates
        end
    end
end


using CarlemanLinearization

"""
    CARLIN <: AbstractContinuousPost

Implementation of the reachability method using Carleman linearization from [1].

[1] Forets, Marcelo, and Christian Schilling. "Reachability of weakly nonlinear systems using Carleman linearization."
    International Conference on Reachability Problems. Springer, Cham, 2021.
"""
Base.@kwdef struct CARLIN <: AbstractContinuousPost
    N::Int=2 # order of the algorithm
    compress::Bool=false # choose to use compressed Kronecker form
    δ::Float64=0.1 # step size for the linear reachability solver
    bloat::Bool=true # choose to include the error estimate in the result
    resets::Union{Int,Vector{Float64}}=0 # choose the number of resets (equal spacing) or a vector specifying the reset times within the time interval [0, T]
end

include("kronecker.jl")
include("post.jl")
include("reach.jl")


"""
Tableau for the 1-stage SRKw1 method due to Wang, Hong & Xu
  Method cited in
  Wang, Hong, Xu, "Construction of Symplectic Runge-Kutta Methods for Stochastic Hamiltonian Systems",
  Commun. Comput. Phys. 21(1), 2017
  According to the paper, the method has weak order 1.0.
"""
function getTableauSRKw1(x::Number=0.0)

    A0 = 0.5*ones(typeof(x),1,1)

    A1 = (1. - x)*ones(typeof(x),1,1)

    B0 = x*ones(typeof(x),1,1)

    B1 = 0.5*ones(typeof(x),1,1)

    B3 = 0.5*ones(typeof(x),1,1)

    α  =  [1.0]
    β1 =  [1.0]

    c0 =  [0.5]
    c1 =  [1. - x]

    TableauWIRK(:SRKw1, A0, A1, B0, B1, B3, α, β1, c0, c1)
end

"""
Tableau for the 4-stage SRKw2 method due to Wang, Hong & Xu
  Method cited in
  Wang, Hong, Xu, "Construction of Symplectic Runge-Kutta Methods for Stochastic Hamiltonian Systems",
  Commun. Comput. Phys. 21(1), 2017
  According to the paper, the method has weak order 2.0 when applied to systems
  driven by one-dimensional noise.
"""
function getTableauSRKw2(x1::Number=0.0, x2::Number=0.0, x3::Number=0.0)

    A0 = [[1. / 8.   0.0        0.0       0.0]
          [1. / 4.   1. / 8.    0.0       0.0]
          [1. / 4.   1. / 4.    1. / 8.   0.0]
          [1. / 4.   1. / 4.    1. / 4.   1. / 8.]]

    A1 = [[-1. / 6. + sqrt(3)/6.   1. / 3. - sqrt(3)/6.   0.0   1. / 3.]
          [ 1. / 2.                0.0                    0.0   0.0]
          [0.0                     0.0                    0.0   0.0]
          [0.0                     0.0                    0.0   0.0]]

    B0 = [[ 5. / 6. - sqrt(3)/3.     -1. / 2.    0.0   0.0]
          [-1. / 6. + sqrt(3)/3.      1. / 2.    0.0   0.0]
          [ 1. / 2.                   1. / 2.    0.0   0.0]
          [-1. / 6.                   1. / 2.    0.0   0.0]]

    B1 = [[1. / 4.                1. / 4. - sqrt(3)/6.   0.0   0.0]
          [1. / 4. + sqrt(3)/6.   1. / 4.                0.0   0.0]
          [x1                     x2                     0.0   x3]
          [0.0                    -1. / 2.               0.0   0.0]]

    B3 = zeros(typeof(x1),4,4)

    α  =  [0.25, 0.25, 0.25, 0.25]
    β1 =  [0.5,  0.5,  0.0,  0.0]

    c0 =  [1. / 8., 3. / 8., 5. / 8., 7. / 8.]
    c1 =  [0.5, 0.5, 0.0, 0.0]

    TableauWIRK(:SRKw2, A0, A1, B0, B1, B3, α, β1, c0, c1)
end


"""
Struct for ShiftedLogNormalMixture

		lnmModel=ShiftedLogNormalMixture(η::Array{num1},λ::Array{num2},num3) where {num1,num2,num3<: Number}

Where:\n
		η  =	Array of volatilities.
		λ  = 	Array of weights.
		α  = 	shift.
"""
mutable struct ShiftedLogNormalMixture{num <: Number, num2 <: Number, num3 <: Number, abstrUnderlying <: AbstractUnderlying, numtype <: Number} <: ItoProcess{numtype}
    η::Array{num, 1}
    λ::Array{num2, 1}
    α::num3
    underlying::abstrUnderlying
    function ShiftedLogNormalMixture(η::Array{num, 1}, λ::Array{num2, 1}, α::num3, underlying::abstrUnderlying) where {num <: Number, num2 <: Number, num3 <: Number, abstrUnderlying <: AbstractUnderlying}
        if minimum(η) <= 0
            error("Volatilities must be positive")
        elseif minimum(λ) <= 0
            error("weights must be positive")
        elseif sum(λ) > 1.0
            error("λs must be weights")
        elseif length(λ) != length(η) - 1
            error("Check vector lengths")
        else
            zero_typed = zero(num) + zero(num2) + zero(num3)
            return new{num, num2, num3, abstrUnderlying, typeof(zero_typed)}(η, λ, α, underlying)
        end
    end
end

export ShiftedLogNormalMixture;

function simulate!(X, mcProcess::ShiftedLogNormalMixture, rfCurve::AbstractZeroRateCurve, mcBaseData::AbstractMonteCarloConfiguration, T::Number)
    @assert T > 0
    r = rfCurve.r
    S0 = mcProcess.underlying.S0
    d = dividend(mcProcess)
    Nstep = mcBaseData.Nstep
    η_gbm = copy(mcProcess.η)
    λ_gmb = copy(mcProcess.λ)
    mu_gbm = r - d
    dt = T / Nstep
    A0 = S0 * (1 - mcProcess.α)
    simulate!(X, LogNormalMixture(η_gbm, λ_gmb, Underlying(A0, d)), rfCurve, mcBaseData, T)
    tt = collect(0:dt:T)
    X .= X .+ mcProcess.α .* S0 .* exp.(integral(mu_gbm, t_) for t_ in tt)'
    return nothing
end


# # Multinomial logistic regression

# Import the necessary packages:

using DataFrames
using Distributions
using MLJBase
using NNlib
using RDatasets
using Soss
using SossMLJ
using Statistics

# In this example, we fit a Bayesian multinomial logistic regression model
# with the canonical link function.

# Suppose that we are given a matrix of features `X` and a column vector of
# labels `y`. `X` has `n` rows and `p` columns. `y` has `n` elements. We assume
# that our observation vector `y` is a realization of a random variable `Y`.
# We define `μ` (mu) as the expected value of `Y`, i.e. `μ := E[Y]`. Our model
# comprises three components:
#
# 1. The probability distribution of `Y`. We assume that each `Yᵢ` follows a multinomial distribution with `k` categories, mean `μᵢ`, and one trial.
# 2. The systematic component, which consists of linear predictor `η` (eta), which we define as `η := Xβ`, where `β` is the column vector of `p` coefficients.
# 3. The link function `g`, which provides the following relationship: `g(E[Y]) = g(μ) = η = Xβ`. It follows that `μ = g⁻¹(η)`, where `g⁻¹` denotes the inverse of `g`. In multinomial logistic regression, the canonical link function is the generalized logit function. The inverse of the generalized logit function is the softmax function. Therefore, when using the canonical link function, `μ = g⁻¹(η) = softmax(η)`.

# A multinomial distribution with one trial is equivalent to the categorical distribution.
# Therefore, the following two statements are equivalent:
# - `Yᵢ` follows a multinomial distribution with `k` categories, mean `μᵢ`, and one trial.
# - `Yᵢ` follows a categorical distribution with `k` categories and mean `μᵢ`.
#
# Observe that the logistic regression model is a special case of the
# multinomial logistic regression model where `k = 2`.
#
# In this model, the parameters that we want to estimate are the coefficients `β`.
# We need to select prior distributions for these parameters. For each `βᵢ`
# we choose a normal distribution with zero mean and unit variance. Here, `βᵢ`
# denotes the `i`th component of `β`.

# We define this model using the Soss probabilistic programming library:

m = @model X,pool begin
    n = size(X, 1) # number of observations
    p = size(X, 2) # number of features
    k = length(pool.levels) # number of classes
    β ~ Normal(0.0, 1.0) |> iid(p, k) # coefficients
    η = X * β # linear predictor
    μ = NNlib.softmax(η; dims=2) # μ = g⁻¹(η) = softmax(η)
    y_dists = UnivariateFinite(pool.levels, μ; pool=pool) # `UnivariateFinite` is mathematically equivalent to `Categorical`
    y ~ For(j -> y_dists[j], n) # `Yᵢ ~ UnivariateFinite(mean=μᵢ, categories=k)`, which is mathematically equivalent to `Yᵢ ~ Categorical(mean=μᵢ, categories=k)`
end;

# Import the Iris flower data set:

iris = dataset("datasets", "iris");

# Define our feature columns:

feature_columns = [
    :PetalLength,
    :PetalWidth,
    :SepalLength,
    :SepalWidth,
]

# Define our label column:

label_column = :Species

# Convert the Soss model into a `SossMLJModel`:

model = SossMLJModel(;
    model       = m,
    predictor   = MLJBase.UnivariateFinite,
    hyperparams = (pool=iris.Species.pool,),
    infer       = dynamicHMC,
    response    = :y,
);

# Create an MLJ machine for fitting our model:

mach = MLJBase.machine(model, iris[!, feature_columns], iris[!, :Species])

# Fit the machine. This may take several minutes.

MLJBase.fit!(mach)

# Construct the joint posterior:

predictor_joint = MLJBase.predict_joint(mach, iris[!, feature_columns])
typeof(predictor_joint)

# Draw a single sample from the joint posterior:

single_sample = rand(predictor_joint)

# For each row in the dataset, construct the marginal posterior predictive distribution

predictor_marginal = MLJBase.predict(mach, iris[!, feature_columns])

# predictor_marginal is a `UnivariateFiniteVector`

typeof(predictor_marginal)

# `predictor_marginal` has one element for each row in the data set

@show size(predictor_marginal); @show size(iris, 1);

# Use cross-validation to evaluate the model with respect to the Brier score:

evaluate!(mach, resampling=CV(; nfolds = 4, shuffle = true), measure=brier_score, operation=MLJBase.predict)

# Use cross-validation to evaluate the model with respect to accuracy:

evaluate!(mach, resampling=CV(; nfolds = 4, shuffle = true), measure=accuracy, operation=MLJBase.predict_mode)

# The cross-validated accuracy is greater than 90%, which is pretty good!


# We use nHidden as a vector, containing the number of hidden units in each layer
using DataFrames
# Function that returns total number of parameters
function NeuralNet_nParams_class(d,nHidden,k)

	# Connections from inputs to first hidden layer
	nParams = d*nHidden[1]

	# Connections between hidden layers
	for h in 2:length(nHidden)
		nParams += nHidden[h-1]*nHidden[h]
	end

	# Connections from last hidden layer to output
	nParams += nHidden[end]*k

end

# Compute squared error and gradient
# for a single training example (x,y)
# (x is assumed to be a column-vector)
function NeuralNet_backprop_class(bigW,x,y,nHidden,k)
	d = length(x)
	nLayers = length(nHidden)

	#### Reshape 'bigW' into vectors/matrices
	W1 = reshape(bigW[1:nHidden[1]*d],nHidden[1],d)
	ind = nHidden[1]*d
	Wm = fill([],(nLayers-1))
	for layer in 2:nLayers
		Wm[layer-1] = reshape(bigW[ind+1:ind+nHidden[layer]*nHidden[layer-1]],nHidden[layer],nHidden[layer-1])
		ind += nHidden[layer]*nHidden[layer-1]
	end
	w = reshape(bigW[ind+1:end],k,nHidden[nLayers])

	### Bin y variables

	#### Define activation function and its derivative
	h(z) = tanh.(z)
	dh(z) = (sech.(z)).^2

	#### Forward propagation
	z = fill([],(nLayers))
	z[1] = W1*x
	for layer in 2:nLayers
		z[layer] = Wm[layer-1]*h(z[layer-1])
	end
	yhat,_ = findmax(w*h(z[end]))
	yhat=yhat-1

	f,r = softmax(w,x,y,k)

	#### Backpropagation
	dr = r
	err = dr

	# Output weights
	Gout = err*h(z[end])

	Gm = fill([],(nLayers-1))
	if nLayers > 1
		# Last Layer of Hidden Weights
		backprop = err*(dh(z[end]).*w)
		Gm[end] = backprop*h(z[end-1])'

		# Other Hidden Layers
		for layer in nLayers-2:-1:1
			backprop = (Wm[layer+1]'*backprop).*dh(z[layer+1])
			Gm[layer] = backprop*h(z[layer])'
		end

		# Input Weights
		backprop = (Wm[1]'*backprop).*dh(z[1])
		G1 = backprop*x'
	else
		# Input weights
		print(size(err),size(x'))
		G1 = err*(dh(z[1]).*w)*x'
	end

	#### Put gradients into vector
	g = zeros(size(bigW))
	g[1:nHidden[1]*d] = G1
	ind = nHidden[1]*d
	for layer in 2:nLayers
		g[ind+1:ind+nHidden[layer]*nHidden[layer-1]] = Gm[layer-1]
		ind += nHidden[layer]*nHidden[layer-1]
	end
	g[ind+1:end] = Gout

	return (f,g)
end

# Computes predictions for a set of examples X

function NeuralNet_predict_class(bigW,Xhat,nHidden,k)
	(t,d) = size(Xhat)
	nLayers = length(nHidden)

	#### Reshape 'bigW' into vectors/matrices
	W1 = reshape(bigW[1:nHidden[1]*d],nHidden[1],d)
	ind = nHidden[1]*d
	Wm = fill([],(nLayers-1))
	for layer in 2:nLayers
		Wm[layer-1] = reshape(bigW[ind+1:ind+nHidden[layer]*nHidden[layer-1]],nHidden[layer],nHidden[layer-1])
		ind += nHidden[layer]*nHidden[layer-1]
	end
	w = reshape(bigW[ind+1:end],k,nHidden[nLayers])

	#### Define activation function and its derivative
	h(z) = tanh.(z)
	dh(z) = (sech.(z)).^2

	#### Forward propagation on each example to make predictions
	yhat = zeros(t,1)
	for i in 1:t
		# Forward propagation
		z = fill([],(nLayers))
		z[1] = W1*Xhat[i,:]
		for layer in 2:nLayers
			z[layer] = Wm[layer-1]*h(z[layer-1])
		end
		yhat[i] = w'*h(z[end])
	end
	return yhat
end

function softmax(w,X,y,k)
	d=length(X)

	yunique=unique(y)
	f=0.0
	g=zeros(k,d)
	for i in 1:n
		print(yunique.==y[i])
		for j in length(yunique)
			if yunique[j]==y[i]
				s=j
			end
		end
		s=firstindex((yunique.==y[i]),1)
		print(s)
        wyi=transpose(w[s,:])
        f+=log(sum(exp.((w*X[i,:]))))-wyi*X[i,:]

        for r in 1:k
            if r==s
                indic=1.0
            else
                indic=0.0
			end

            for j in 1:d
                g[r,j]+=X[i,j]*((exp.((transpose(w[r,:])*X[i,:]))/sum(exp.((w*X[i,:]))))-indic)
	      	end
	  	end
	end

    return f,g
end


import Pkg
Pkg.activate("..")

using CrimsonSkyline
using Distributions: Normal, Gamma, LogNormal
using Logging
using StatsBase: mean, std
using Random: seed!

seed!(2021)

function random_walk(t :: Trace, data :: Array{Float64, 2})
    loc = sample(t, :loc, Normal())
    scale = sample(t, :scale, Gamma(3.0, 5.0))
    shape = size(data)
    for n in 1:shape[1]
        x_n = sample(t, (:ic, n), Normal())
        for t_ix in 1:shape[2]
            x_n = observe(t, (:data, n, t_ix), Normal(loc + x_n, scale), data[n, t_ix])
        end
    end
end

function global_trend(t :: Trace, data :: Array{Float64, 2})
    α = sample(t, :α, Normal())
    β = sample(t, :β, Normal())
    σ = sample(t, :σ, LogNormal())
    shape = size(data)
    for n in 1:shape[1]
        for t_ix in 1:shape[2]
            observe(t, (:data, n, t_ix), Normal(α + β * t_ix, σ), data[n, t_ix])
        end
    end
end

function main()
    n = 10
    t = 20
    t_pred = 10

    @info "Random walk inference"
    loc = 1.0
    scale = 1.5
    noise = randn(n, t + t_pred)
    data = cumsum(loc .+ scale .* noise, dims=2)
    test_data = data[:, t:end]
    data = data[:, 1:t]
    @info "True loc = $loc, true scale = $scale"
    @time results = mh(random_walk; params=(data,), burn=1000, thin=50, num_iterations=50000)
    @info "Estimated loc = ($(mean(results, :loc)) ± $(2.0 * std(results, :loc)))"
    @info "Estimated scale = ($(mean(results, :scale)) ± $(2.0 * std(results, :scale)))"

    @info begin
        """
        Time series model comparison
        Comparing (prior) evidence for random walk model vs global trend
        """
    end
    nsamples = 5000
    rw_results = likelihood_weighting(random_walk, data; nsamples = nsamples)
    gt_results = likelihood_weighting(global_trend, data; nsamples = nsamples)
    log_evidence_rw = log_evidence(rw_results)
    log_evidence_gt = log_evidence(gt_results)
    @info begin
        """
        log p(x | random walk) ≈ $log_evidence_rw
        log p(x | global trend) ≈ $log_evidence_gt
        """
    end
    @info "Comparing model quality using AIC"
    aic_rw = aic(rw_results)
    aic_gt = aic(gt_results)
    @info "AIC(random walk) ≈ $aic_rw, AIC(global trend) ≈ $aic_gt"

end

main()


module BasisFunctionsModule
export PrimitiveGaussianBasisFunction, ContractedGaussianBasisFunction, identityPGF, identityCGF
import Base.display
import QuantumLab.BaseModule.evaluateFunction
using ..BaseModule

"""
A PrimitiveGaussianBasisFunction is a function of the form x^a y^b z^c exp(-ζ r^2) centered at a position in space. (a,b,c) here are by definition the cartesian MQuantumNumber tuple.
"""
type PrimitiveGaussianBasisFunction
  center::Position
  exponent::Float64
  mqn::MQuantumNumber
end

"""
A ContractedGaussianBasisFunction is defined as a sum of PrimitiveGaussianBasisFunctions with fixed relative coefficients. Different from other Quantum Chemistry codes we do not require
an additional scaling for the total ContractedGaussianBasisFunction - scaling the whole ContractedGaussianBasisFunction happens by modifying the contraction coefficients accordingly.
"""
type ContractedGaussianBasisFunction
  coefficients::Array{Float64,1}
  primitiveBFs::Array{PrimitiveGaussianBasisFunction,1}
end

function display(cgbf::ContractedGaussianBasisFunction,indent="")
  print(indent); dump(typeof(cgbf))
  for idx in 1:length(cgbf.coefficients)
    print(indent * "  $(cgbf.coefficients[idx]) × "); dump(typeof(cgbf.primitiveBFs[idx]))
    print(indent * "    center:   "); println(cgbf.primitiveBFs[idx].center)
    print(indent * "    exponent: "); println(cgbf.primitiveBFs[idx].exponent)
    print(indent * "    mqn:      "); println(cgbf.primitiveBFs[idx].mqn)
  end
end

function evaluateFunction(x::Position,pgbf::PrimitiveGaussianBasisFunction)
  α = pgbf.exponent
  O = pgbf.center
  r = distance(x,O)
  return (x.x - O.x)^(pgbf.mqn.x) * (x.y - O.y)^(pgbf.mqn.y) * (x.z - O.z)^(pgbf.mqn.z) * exp(-α*r^2)
end

function evaluateFunction(x::Position,cgbf::ContractedGaussianBasisFunction)
  result = 0.::Float64
  for (coeff,pgbf) in zip(cgbf.coefficients,cgbf.primitiveBFs)
    result += coeff * evaluateFunction(x,pgbf)
  end
  return result
end

const identityPGF = PrimitiveGaussianBasisFunction(origin, 0., MQuantumNumber(0,0,0))
const identityCGF = ContractedGaussianBasisFunction([1.],[identityPGF])

end


#############################################################################
#############################################################################
#
# A script that runs all unit tests in the project.
#
#############################################################################
#############################################################################

include("../poly_factorization_project.jl")

####
# Execute unit tests for integers
###
include("mod_integers_test.jl")
mod_test_euclid_ints()
mod_test_ext_euclid_ints()

####
# Execute unit tests for polynomials
####
include("mod_polynomials_test.jl")
for i in [3,11,17,19,67,101]
    println("")
    println("testing p = $i")
    mod_sum_test_poly(i)
    mod_prod_test_poly(i)
    mod_prod_derivative_test_poly(i)
    mod_ext_euclid_test_poly(i)
    mod_division_test_poly(i)

end

####
# Execute unit tests for polynomial factorization
####
include("mod_factorization_test.jl")
mod_factor_test_poly()


using Plots

# Plot of linked scalar sequences
x = 0:0.001:5
s1 = sin.(x)'
s2 = sin.(log.(range(1;length=5001, stop=exp(5))))'
l, m = linkseq(s1', s2', 1000, 50)

p = plot(x, [s1[:] s2[:]])
plot!(p, l, m; legend=false)
savefig(p, "docs/assets/2d_linked_seqs.svg")

# Plot of linked vector sequences
x = 0:0.001:5
s1 = [0.3; -0.7] * sin.(x)'
s2 = [0.7; 0.3] * sin.(log.(range(1;length=5001, stop=exp(5))))'
l, m = linkseq(s1, s2, 1000, 50)

p = plot(x, [s1[1,:] s2[1,:]], [s1[2,:] s2[2,:]])
plot!(p, l, m[1,:,:], m[2,:,:]; legend=false)
savefig(p, "docs/assets/3d_linked_seqs.svg")

# Plot different linked sequences
s1 = [0, 1, 2, 4, 5, 2, 3, 1, 0, 0]
s2 = [2, 4, 5, 6, 6, 5, 5, 4, 3, 1, 0, 0]
l, m = linkseq(s1, s2)
p = plot(s1)
plot!(p, s2)
plot!(p, l.+1, m; color=:black, legend=false)
savefig(p, "docs/assets/different_seqs.svg")


using ArgParse,CairoMakie

CairoMakie.activate!(type="pdf")

function poisson_prob_desity(λ,x)
    ((λ^x)/factorial(x))*(ℯ^(-λ))
end

function main()
    graph_name = "poisson_cumulative_dist"
    xs_max = 20
    xs = 0:xs_max

    graph_params = [Dict("lambda"=>1,  "color"=>:gray0,  "marker"=>:circle,  "label"=>"λ=1"),
                    Dict("lambda"=>5,  "color"=>:gray30, "marker"=>:diamond, "label"=>"λ=5"),
                    Dict("lambda"=>10, "color"=>:gray50, "marker"=>:rect,    "label"=>"λ=15")]

    fig = Figure(resolution = (800, 600))

    axis = Axis(fig[1,1], title = "poisson cumulative distribution function")

    for params in graph_params
        c_s::Array{Float64} = []
        accum::Float64 = 0
        for x in xs
            accum += poisson_prob_desity(params["lambda"],x)
            push!(c_s,accum)
        end
        stem!(fig[1,1], xs, c_s, color = params["color"], marker = params["marker"], label=params["label"])
    end

    axislegend(axis, position = :rb)

    save("$graph_name.pdf", fig)
end

main()


@testset "cheb_diff" begin
    # Test the Chebyshev differentiation matrix implementation
    x2 = [0, 1//2, 1]
    D2 = [-3 4 -1; -1 0 1; 1 -4 3]
    x3 = [0, 1//4, 3//4, 1]
    D3 = [-38//6 8 -8//3 1; -2 2//3 2 -2//3; 2//3 -2 -2//3 2; -1 8//3 -8 38//6]
    # x4 = [1, 1/sqrt(2), 0, -1/sqrt(2), -1] # could work out symbolically if want to test BigFloat accuracy
    @test all(cheb_diff(Float64, 2) .≈ (x2, D2))
    @test all(cheb_diff(Float64, 3) .≈ (x3, D3))
    @test all(eltype.(cheb_diff(Float32, 3)) .=== (Float32, Float32))
end


abstract AbstractKronMat{T} <: AbstractMatrix{T}
abstract AbstractKronVec{T} <: AbstractVector{T}
typealias AbstractKronVecOrMat Union(AbstractKronMat,AbstractKronVec)

#sizes, indexing, etc
sizes(C::AbstractKronVecOrMat) = [size(M) for M in terms(C)]
size(C::AbstractKronVecOrMat) = map(*,sizes(C)...)

lengths(C::AbstractKronVec) = [length(term) for term in terms(C)]
length(C::AbstractKronVec) = prod(lengths(C))

#function hasmixedproductproperty(C::KronMat,D::KronMat)
#    (MCO,NCO),(MCI,NCI) = sizes(C)
#    (MDO,NDO),(MDI,NDI) = sizes(D)
#    return NCO==MDO && NCI==MDI
#end
#hasmpp = hasmixedproductproperty

#do I need to merely reverse something?
function hindexes(i::Int64,widths)
    widths = reverse(widths)
    return reverse([(mod(div(i-1,u),w)+1)::Int64 for (u,w) in zip(cumprod([1, widths...]),widths)])
end

function kronindexes(C::AbstractKronVec,i::Int64)
    widths = lengths(C)
    return hindexes(i,widths)
end

function kronindexes(C::AbstractKronMat,i::Int64,j::Int64)
    iwidths,jwidths = zip(sizes(C)...)
    hi,hj = hindexes(i,iwidths),hindexes(j,jwidths)
    return (hi,hj)
end

function convert{T}(::Type{Matrix{T}},C::AbstractKronMat)
    return kronprod(terms(C)...)
end

function full{T}(C::AbstractKronMat{T})
    M,N = size(C)
    return T[C[i,j] for i=1:M, j = 1:N]
end

function full{T}(C::AbstractKronVec{T})
    N = length(C)
    return T[C[i] for i=1:N]
end


#=
""" Give the boundary nodes """
function bdy_nodes(dims::Tuple)
    D = length(dims)
    nbdy = prod(dims) - prod(dims .- 2)
    bdy = CartesianIndex{D}[]
    for (i, d) in enumerate(dims)
        # This only works with julia v 1.6
        push!(bdy, CartesianIndices((dims[1:(i-1)]..., 1:(d-1):d, dims[(i+1):D]...))[:]...)
    end
    bdy = unique(bdy)
    neighbors = zeros(length(bdy), D)
    for (j, b) in enumerate(bdy)
        for i in 1:D
          neighbors[j, i] = ((Tuple(b)[i] == 1)||(Tuple(b)[i] == dims[i])) ? 1 : 2
        end
    end
    return bdy, neighbors
end
=#

"""
  bdy_nodes(dims)

Boundary node computation, for arbitrary dimension

...
# Arguments

  - `dims::Tuple` number of points in each direction
...

...
# Outputs
  - `Vector{Int64}`: vector containing the indices of coordinates
  on the boundary of the hyperrectangle volume
...

"""
function bdy_nodes(dims::Tuple)
    D = length(dims)
    bdy = CartesianIndex[]
    for d in CartesianIndices(dims)
        if (sum(Tuple(d) .== 1) > 0) || (sum(Tuple(d) .== dims) > 0)
            push!(bdy, d)
        end
    end
    neighbors = zeros(Int64, length(bdy), D)
    for (j, b) in enumerate(bdy)
        for i in 1:D
          neighbors[j, i] = ((Tuple(b)[i] == 1)||(Tuple(b)[i] == dims[i])) ? 1 : 2
        end
    end
    return bdy, neighbors
end

""" Write the matrix """
function nablasq_arb(dims, aspect_ratios)
  D = length(dims)
  oh = map(i -> ones(dims[i])/aspect_ratios[i], 1:length(dims))
  del = Vector{SparseArrays.SparseMatrixCSC{Float64, Int64}}(undef, D)
  for (i, d) in enumerate(dims)
    del[i] = spdiagm_nonsquare(d + 1, d, -1 => -oh[i], 0 => oh[i])
  end
  A = zeros(prod(dims), prod(dims))
  for i in 1:D
      A += kron([sparse(I, dims[j], dims[j]) for j in 1:(i-1)]..., del[i]'*del[i],
                [sparse(I, dims[j], dims[j]) for j in (i+1):D]...)
  end
  bdy, neighbors = bdy_nodes(dims)
  for (count, b) in enumerate(bdy)
    i = LinearIndices(zeros(dims...))[b]
    A[i, i] = 0.0
    for j in 1:D
      A[i, i] += neighbors[count, j] / aspect_ratios[j] ^ 2
    end
  end
  return A
end

""" Helper function to give the Matern matrix in arbitrary dimensions """
function _Matern_matrix(dims, m, eps, aspect_ratios)
    A = nablasq_arb(dims, aspect_ratios)
    sizeA = size(A, 1)
    for i = 1:sizeA
        A[i, i] = A[i, i] + eps^2
    end
    return A^m
end

""" Compute the interpolation """
function interp(data, ind, m, eps, aspect_ratios)
    dims = size(data)
    A = (eps == 0.0)&&(m == 1) ?
                nablasq_arb(dims, aspect_ratios) :
                _Matern_matrix(dims, m, eps, aspect_ratios)
    totalsize = prod(dims)
    C = sparse(I, totalsize, totalsize)
    rhs_a = data[:]
    for i in ind
        j = (typeof(i) <: CartesianIndex) ? LinearIndices(data)[i] : i
        C[j, j] = 0.0
        rhs_a[j] = 0.0
    end
    Id = sparse(I, totalsize, totalsize)
    u = ((C - (Id - C) * A)) \ rhs_a
    return reshape(u, dims...)
end


using Statistics
using Plots
using BenchmarkTools
using Profile
using FFTW
using Statistics
using Optim
using Images, FileIO, ImageIO
using Printf
using Revise
using Profile
using LinearAlgebra
using Random
using StatsBase
using Distributions
using FITSIO
using SparseArrays

push!(LOAD_PATH, pwd()*"/main")
using DHC_2DUtils
#include("Deriv_Utils_New.jl")
#import Deriv_Utils_New
push!(LOAD_PATH, pwd()*"/scratch_NM")
using Deriv_Utils_New
using Data_Utils
using Visualization
using ReconFuncs

##How does full cov in the log-im case make a difference? vs Diag (Getting the high density and the other artefacts)
fname = "scratch_NM/StandardizedExp/Nx64/noisy_stdtrue/SFDTargSFDCov/log_apd_noiso/1000_triu"
gttarget = load(fname*".jld2")
#Calc 1dps log
Nx = size(gttarget["true_img"])[1]
kbins = collect(1:Nx/2.0)

true_img = gttarget["true_img"]
init = gttarget["init"]
recon_img =gttarget["recon"]

println("Mean Abs Res: Init-True = ", mean(abs.(init - true_img)), " Recon-True = ", mean(abs.(recon_img - true_img)))
println("Mean Abs Frac Res", mean(abs.((init - true_img)./true_img)), " Recon-True=", mean(abs.((recon_img - true_img)./true_img)))
println("RMSE Init-True= ", (mean((init - true_img).^2)).^0.5, " Recon-True=", (mean((recon_img - true_img).^2)).^0.5)

true_lps = calc_1dps(apodizer(log.(gttarget["true_img"])), kbins)
recon_lps = calc_1dps(apodizer(log.(gttarget["recon"])), kbins)
init_lps = calc_1dps(apodizer(log.(gttarget["init"])), kbins)
p = plot(log.(kbins), log.(true_lps), label="True")
plot!(log.(kbins), log.(recon_lps), label="Recon")
plot!(log.(kbins), log.(init_lps), label="Init")
plot!(title="DiagCov Noiso: P(k) of Log Image")
xlabel!("lnk")
ylabel!("lnP(k)")

true_ps = calc_1dps(apodizer(gttarget["true_img"]), kbins)
recon_ps = calc_1dps(apodizer(gttarget["recon"]), kbins)
init_ps = calc_1dps(apodizer(gttarget["init"]), kbins)
p = plot(log.(kbins), log.(true_ps), label="True")
plot!(log.(kbins), log.(recon_ps), label="Recon")
plot!(log.(kbins), log.(init_ps), label="Init")
plot!(title="DiagCov Noiso: P(k) of Image")
xlabel!("lnk")
ylabel!("lnP(k)")

##With the Full Cov:
fname = "scratch_NM/StandardizedExp/Nx64/noisy_stdtrue/SFDTargSFDCov/log_apd_noiso/1000_fullcov_triu"
gttarget = load(fname*".jld2")
#Calc 1dps log
Nx = size(gttarget["true_img"])[1]
kbins = collect(1:Nx/2.0)

true_img = gttarget["true_img"]
init = gttarget["init"]
recon_img =gttarget["recon"]

println("Mean Abs Res: Init-True = ", mean(abs.(init - true_img)), " Recon-True = ", mean(abs.(recon_img - true_img)))
println("Mean Abs Frac Res", mean(abs.((init - true_img)./true_img)), " Recon-True=", mean(abs.((recon_img - true_img)./true_img)))
println("RMSE Init-True= ", (mean((init - true_img).^2)).^0.5, " Recon-True=", (mean((recon_img - true_img).^2)).^0.5)

true_lps = calc_1dps(apodizer(log.(gttarget["true_img"])), kbins)
recon_lps = calc_1dps(apodizer(log.(gttarget["recon"])), kbins)
init_lps = calc_1dps(apodizer(log.(gttarget["init"])), kbins)
p = plot(log.(kbins), log.(true_lps), label="True")
plot!(log.(kbins), log.(recon_lps), label="Recon")
plot!(log.(kbins), log.(init_lps), label="Init")
plot!(title="FullCov Noiso: P(k) of Log Image")
xlabel!("lnk")
ylabel!("lnP(k)")

true_ps = calc_1dps(apodizer(gttarget["true_img"]), kbins)
recon_ps = calc_1dps(apodizer(gttarget["recon"]), kbins)
init_ps = calc_1dps(apodizer(gttarget["init"]), kbins)
p = plot(log.(kbins), log.(true_ps), label="True")
plot!(log.(kbins), log.(recon_ps), label="Recon")
plot!(log.(kbins), log.(init_ps), label="Init")
plot!(title="FullCov Noiso: P(k) of Image")
xlabel!("lnk")
ylabel!("lnP(k)")

##Why Don't Reg plots move away?? Actually it ends up closer to the target than the true image.
#Maybe loss in reg doesn't differentiate noisy from not noisy?
fname = "scratch_NM/StandardizedExp/Nx64/noisy_stdtrue/SFDTargSFDCov/reg_apd_noiso/1000_fullcov_triu"
s_targ_mean, s_targ_invcov = gttarget["dict"]["s_targ_mean"], gttarget["dict"]["s_invcov"]
dhc_args = Dict(:doS2=>false, :doS12=>false, :doS20=>true, :apodize=>true, :iso=>false)

altrue = apodizer(gttarget["true_img"])
alinit = apodizer(gttarget["init"])
alrecon = apodizer(gttarget["recon"])

coeff_mask = gttarget["coeff_mask"]
s_true = DHC_compute_wrapper(gttarget["true_img"], filter_hash, norm=false; dhc_args...)[coeff_mask]
s_init = DHC_compute_wrapper(gttarget["init"], filter_hash, norm=false; dhc_args...)[coeff_mask]
s_recon = DHC_compute_wrapper(gttarget["recon"], filter_hash, norm=false; dhc_args...)[coeff_mask]
ltruecheck = ( 0.5 .* (s_true - s_targ_mean)' * s_targ_invcov * (s_true - s_targ_mean))[1]
l2true = 0.5*gttarget["dict"]["lambda"]*sum((alinit - altrue).^2)
#7.82, 2.32

l1init = ( 0.5 .* (s_init - s_targ_mean)' * s_targ_invcov * (s_init - s_targ_mean))[1]
l2init = 0.5*gttarget["dict"]["lambda"]*sum((alinit - alinit).^2)
#11.2, 0

l1recon = ( 0.5 .* (s_recon - s_targ_mean)' * s_targ_invcov * (s_recon - s_targ_mean))[1]
l2recon = 0.5*gttarget["dict"]["lambda"]*sum((alrecon - alinit).^2)
#7.61, 3.57
ltruecheck + l2true
l1init + l2init
l1recon + l2recon

##Log
fname = "scratch_NM/StandardizedExp/Nx64/noisy_stdtrue/SFDTargSFDCov/log_apd_noiso/1000_fullcov_triu"
s_targ_mean, s_targ_invcov = gttarget["dict"]["s_targ_mean"], gttarget["dict"]["s_invcov"]
dhc_args = Dict(:doS2=>false, :doS12=>false, :doS20=>true, :apodize=>true, :iso=>false)

altrue = apodizer(log.(gttarget["true_img"]))
alinit = apodizer(log.(gttarget["init"]))
alrecon = apodizer(log.(gttarget["recon"]))

coeff_mask = gttarget["coeff_mask"]
s_true = DHC_compute_wrapper(log.(gttarget["true_img"]), filter_hash, norm=false; dhc_args...)[coeff_mask]
s_init = DHC_compute_wrapper(log.(gttarget["init"]), filter_hash, norm=false; dhc_args...)[coeff_mask]
s_recon = DHC_compute_wrapper(log.(gttarget["recon"]), filter_hash, norm=false; dhc_args...)[coeff_mask]
ltruecheck = ( 0.5 .* (s_true - s_targ_mean)' * s_targ_invcov * (s_true - s_targ_mean))[1]
l2true = 0.5*gttarget["dict"]["lambda"]*sum((alinit - altrue).^2)
#446, 84503

l1init = ( 0.5 .* (s_init - s_targ_mean)' * s_targ_invcov * (s_init - s_targ_mean))[1]
l2init = 0.5*gttarget["dict"]["lambda"]*sum((alinit - alinit).^2)
#6e9, 0

l1recon = ( 0.5 .* (s_recon - s_targ_mean)' * s_targ_invcov * (s_recon - s_targ_mean))[1]
l2recon = 0.5*gttarget["dict"]["lambda"]*sum((alrecon - alinit).^2)
#2.5e3, 6e4

ztrue_recon = (s_recon - s_true) .* diag(s_targ_invcov)
ztrue_target = (s_true - s_targ_mean) .* diag(s_targ_invcov)
zinit_target = (s_init - s_targ_mean) .* diag(s_targ_invcov)

scatter(log.(abs.(zinit_target)), log.(abs.(ztrue_recon)))
xlabel!("ln Abs Z score: Coeffs Target-Init")
ylabel!("ln Abs Z score: Coeffs Recon-True")

zscoremat = zeros(34, 34)
zscoremat[reshape(gttarget["coeff_mask"][3+34:end], (34, 34))] .= log.(abs.(ztrue_target))
heatmap(zscoremat, title="Im 1000: Ln Abs Z-Score b/w True & Target")

gttarget["fhash"]

zscoremat
#basically the J=3 scales that are problematic

filename = normpath("scratch_NM/NewWrapper/Weave", "4-7.jmd")
weave(filename, out_path = :pwd)
using Pkg
Pkg.add("IJulia")

using IJulia
IJulia.notebook(dir="/home/nayantara/Desktop/GalacticDustProject/WST/DHC/scratch_NM/NewWrapper/4-7/")

##ANALYZING Z SCORES AND COVARIANCES
#Log coeff, log im
Nx=64
dbnimg = readsfd(Nx, logbool=false)
fname = "scratch_NM/StandardizedExp/Nx64/noisy_stdtrue/SFDTargSFDCov/reg_apd_noiso/LogCoeff/1000_0-01_full_triu"
gttarget = load(fname*".jld2")
dhc_args = gttarget["dhc_args"]
coeff_mask = gttarget["coeff_mask"]
fhash = gttarget["fhash"]
dbn1 = log.(get_dbn_coeffs(dbnimg[:, :, 1:5000], fhash, dhc_args, coeff_mask = coeff_mask))
dbn2 = log.(get_dbn_coeffs(dbnimg[:, :, 5001:end], fhash, dhc_args, coeff_mask= coeff_mask))

#Mean dia, cov
Ncovsamp = 5000
s_targ_mean = mean(dbn1, dims=1)
scov  = (dbn1 .- s_targ_mean)' * (dbn1 .- s_targ_mean) ./(Ncovsamp-1)
smean1, sdiag1, scov1 = s_targ_mean, diag(scov), scov
s_targ_mean = mean(dbn2, dims=1)
scov  = (dbn2 .- s_targ_mean)' * (dbn2 .- s_targ_mean) ./(Ncovsamp-1)
smean2, sdiag2, scov2 = s_targ_mean, diag(scov), scov

scov1inv = invert_covmat(scov1, 1e-8)
scov2inv = invert_covmat(scov2, 1e-8)

diff = dbn1 .- reshape(smean1, (1, 595))
chisq1 = sum((diff * scov1inv) .* diff, dims=2)

diff = dbn2 .- reshape(smean2, (1, 595))
chisq2 = sum((diff * scov2inv) .* diff, dims=2)
zscores = (dbn1 .- s_targ_mean).* reshape(diag(scov1inv).^0.5, (1, 595))
p = histogram(chisq1[:], label="Dbn1")
histogram!(sum(zscores, dims=2)./595, label="Z_Dbn1")
histogram!(chisq2[:], label="Dbn2")
chisqsamp = rand(Distributions.Chisq(595), 5000)
histogram!(chisqsamp, label="ChiSq")

fname = "scratch_NM/StandardizedExp/Nx64/noisy_stdtrue/SFDTargSFDCov/reg_apd_noiso/LogCoeff/1000_0-01_full_triu"
gttarget = load(fname*".jld2")
filter_hash = gttarget["fhash"]
s_targ_mean, s_targ_invcov = gttarget["dict"]["fs_targ_mean"], gttarget["dict"]["fs_invcov"]
dhc_args = Dict(:doS2=>false, :doS12=>false, :doS20=>true, :apodize=>true, :iso=>false)

altrue = apodizer(gttarget["true_img"])
alinit = apodizer(gttarget["init"])
alrecon = apodizer(gttarget["recon"])

coeff_mask = gttarget["coeff_mask"]
s_true = DHC_compute_wrapper(gttarget["true_img"], filter_hash, norm=false; dhc_args...)[coeff_mask])
s_init = DHC_compute_wrapper(gttarget["init"], filter_hash, norm=false; dhc_args...)[coeff_mask]
s_recon = DHC_compute_wrapper(gttarget["recon"], filter_hash, norm=false; dhc_args...)[coeff_mask]

struesq = zeros(34, 34)
struesq[reshape(coeff_mask[37:end], (34, 34))] .= s_true
sinitsq = zeros(34, 34)
sinitsq[reshape(coeff_mask[37:end], (34, 34))] .= s_init
sreconsq = zeros(34, 34)
sreconsq[reshape(coeff_mask[37:end], (34, 34))] .= s_recon
stargsq = zeros(34, 34)
stargsq[reshape(coeff_mask[37:end], (34, 34))] .= exp.(gttarget["dict"]["fs_targ_mean"])

s_true = log.(s_true)
s_init  = log.(s_init)
s_recon = log.(s_recon)
ltruecheck = ( 0.5 .* (s_true - s_targ_mean)' * s_targ_invcov * (s_true - s_targ_mean))[1]
l2true = 0.5*gttarget["dict"]["lambda"]*sum((alinit - altrue).^2)
println("True", ltruecheck)
println(l2true)

l1init = ( 0.5 .* (s_init - s_targ_mean)' * s_targ_invcov * (s_init - s_targ_mean))[1]
l2init = 0.5*gttarget["dict"]["lambda"]*sum((alinit - alinit).^2)
println("Init", l1init)
println(l2init)
println("From true", ( 0.5 .* (s_init - s_true)' * s_targ_invcov * (s_init - s_true))[1])

l1recon = ( 0.5 .* (s_recon - s_targ_mean)' * s_targ_invcov * (s_recon - s_targ_mean))[1]
l2recon = 0.5*gttarget["dict"]["lambda"]*sum((alrecon - alinit).^2)
println("Recon", l1recon)
println(l2recon)
println("From true", ( 0.5 .* (s_recon - s_true)' * s_targ_invcov * (s_recon - s_true))[1])

##Z-Score for mv gaussians test case
covmat = reshape([30, -15, -1, -15, 21, 3, -1, 3, 2], (3, 3))
mvg = Distributions.MultivariateNormal([-100, 3, 0.05], covmat)
samps = rand(mvg, 100)
samean = mean(samps, dims=2)
sacov = (samps .- samean) * (samps .- samean)'
sacov = sacov ./(99)
zscores = (samps .- samean) ./ (diag(sacov).^0.5)
#these do look like z scores

##Are the larger J value coefficients close to the true??
Nx=64
dbnimg = readsfd(Nx, logbool=false)
fname = "scratch_NM/StandardizedExp/Nx64/noisy_stdtrue/SFDTargSFDCov/reg_apd_noiso/LogCoeff/1000_0-01_full_triu"
gttarget = load(fname*".jld2")
dhc_args = gttarget["dhc_args"]
coeff_mask = gttarget["coeff_mask"]
fhash = gttarget["fhash"]
dbn = get_dbn_coeffs(dbnimg, fhash, dhc_args, coeff_mask = nothing)

s_true = DHC_compute_wrapper(gttarget["true_img"], fhash, norm=false; dhc_args...)
s_init = DHC_compute_wrapper(gttarget["init"], fhash, norm=false; dhc_args...)
s_recon = DHC_compute_wrapper(gttarget["recon"], fhash, norm=false; dhc_args...)
ffb, ffinfo = fink_filter_bank(1, 8, nx=64, Omega=true)

function J_hashindices(J_values)
    jindlist = []
    for jval in J_values
        push!(jindlist, findall(fhash["J_L"][:, 1].==jval))
    end
    return vcat(jindlist'...)
end

function J_S1indices(J_values)
    #Assumes this is applied to an object of length 2+Nf+Nf^2 or 2+Nf
    return J_hashindices(J_values) .+ 2
end

function distribution_percentiles(dbn, idx, perc)
    idxperc = zeros(size(idx))
    idxperc .= (x-> percentile(dbn[:, x], perc)).(idx)
    return idxperc
end

Jfiltidx = J_hashindices([0, 1, 2, 3])
JS1ind = J_S1indices([0, 1, 2, 3])
p25 = distribution_percentiles(dbn, JS1ind, 25)
p75 = distribution_percentiles(dbn, JS1ind, 75)
strueselect = (x->s_true[x]).(JS1ind)
sinitselect = (x->s_init[x]).(JS1ind)
sreconselect = (x->s_recon[x]).(JS1ind)
truebool = (strueselect .> p25) .& (strueselect .<p75)
initbool = (sinitselect .> p25) .& (sinitselect .<p75)
reconbool = (sreconselect .> p25) .& (sreconselect .<p75)

frac_devn_init_true = (sinitselect .- strueselect)./strueselect
frac_devn_recon_true = (sreconselect .- strueselect)./strueselect
dbn_mean = mean(dbn, dims=1)[:]
meanselect = (x->dbn_mean[x]).(JS1ind)
frac_devn_sfd_true = (meanselect .- strueselect)./strueselect
frac_devn_75_25 = (p75 .- p25)./p25

dbn_median = distribution_percentiles(dbn, JS1ind, 50)
frac_devn_medsfd_true = (dbn_median .- strueselect)./strueselect
clim = (minimum(frac_devn_medsfd_true), maximum(frac_devn_medsfd_true))
heatmap(frac_devn_init_true, title="Frac Devn Init-True", xlabel="L", ylabel="J", xticks=collect(1:8))
heatmap(frac_devn_sfd_true, title="Frac Devn SFDMean-True", xlabel="L", ylabel="J", xticks=collect(1:8))
heatmap(frac_devn_75_25, title="Frac Devn 75-25", xlabel="L", ylabel="J", xticks=collect(1:8))
heatmap(frac_devn_medsfd_true, title="Frac Devn SFDMedian-True", xlabel="L", ylabel="J", xticks=collect(1:8))
frac_devn_mean_med = (meanselect .- dbn_median)./dbn_median
heatmap(frac_devn_mean_med, title="Frac Devn Mean-Median", xlabel="L", ylabel="J", xticks=collect(1:8))

heatmap(frac_devn_medsfd_true, title="Frac Devn SFDMedian-True", xlabel="L", ylabel="J", xticks=collect(1:8), clim=clim)
heatmap(frac_devn_init_true, title="Frac Devn Init-True", xlabel="L", ylabel="J", xticks=collect(1:8), clim=clim)

##Does mean-subtraction make a difference?
meanim = mean(dbnimg, dims=[1, 2])
dbnimg_ms = dbnimg .- meanim
fname = "scratch_NM/StandardizedExp/Nx64/noisy_stdtrue/SFDTargSFDCov/reg_apd_noiso/LogCoeff/1000_0-01_full_triu"
gttarget = load(fname*".jld2")
dhc_args = gttarget["dhc_args"]
coeff_mask = gttarget["coeff_mask"]
fhash = gttarget["fhash"]
dbn = get_dbn_coeffs(dbnimg_ms, fhash, dhc_args, coeff_mask = nothing)

s_true = DHC_compute_wrapper(gttarget["true_img"] .- mean(gttarget["true_img"]), fhash, norm=false; dhc_args...)
s_init = DHC_compute_wrapper(gttarget["init"] .- mean(gttarget["init"]), fhash, norm=false; dhc_args...)
s_recon = DHC_compute_wrapper(gttarget["recon"] .- mean(gttarget["recon"]), fhash, norm=false; dhc_args...)
ffb, ffinfo = fink_filter_bank(1, 8, nx=64, Omega=true)

Jfiltidx = J_hashindices([0, 1, 2, 3])
JS1ind = J_S1indices([0, 1, 2, 3])
p25 = distribution_percentiles(dbn, JS1ind, 25)
p75 = distribution_percentiles(dbn, JS1ind, 75)
strueselect = (x->s_true[x]).(JS1ind)
sinitselect = (x->s_init[x]).(JS1ind)
sreconselect = (x->s_recon[x]).(JS1ind)
truebool = (strueselect .> p25) .& (strueselect .<p75)
initbool = (sinitselect .> p25) .& (sinitselect .<p75)
reconbool = (sreconselect .> p25) .& (sreconselect .<p75)

frac_devn_init_true = (sinitselect .- strueselect)./strueselect
frac_devn_recon_true = (sreconselect .- strueselect)./strueselect
dbn_mean = mean(dbn, dims=1)[:]
meanselect = (x->dbn_mean[x]).(JS1ind)
frac_devn_sfd_true = (meanselect .- strueselect)./strueselect
frac_devn_75_25 = (p75 .- p25)./p25

dbn_median = distribution_percentiles(dbn, JS1ind, 50)
frac_devn_medsfd_true = (dbn_median .- strueselect)./strueselect
heatmap(frac_devn_init_true, title="Frac Devn Init-True MeanSub", xlabel="L", ylabel="J", xticks=collect(1:8))
heatmap(frac_devn_sfd_true, title="Frac Devn SFDMean-True MeanSub", xlabel="L", ylabel="J", xticks=collect(1:8))
heatmap(frac_devn_75_25, title="Frac Devn 75-25 MeanSub", xlabel="L", ylabel="J", xticks=collect(1:8))
heatmap(frac_devn_medsfd_true, title="Frac Devn SFDMedian-True MeanSub", xlabel="L", ylabel="J", xticks=collect(1:8))
frac_devn_mean_med = (meanselect .- dbn_median)./dbn_median
heatmap(frac_devn_mean_med, title="Frac Devn Mean-Median MeanSub", xlabel="L", ylabel="J", xticks=collect(1:8))

##Making sure the log of coefficients didn't have to do with the results
Nx=64
dbnimg = readsfd(Nx, logbool=false)
fname = "scratch_NM/StandardizedExp/Nx64/noisy_stdtrue/SFDTargSFDCov/reg_apd_noiso/LambdaVary/1000_tunedlam_0-01_fullcov_triu"
gttarget = load(fname*".jld2")
dhc_args = gttarget["dhc_args"]
coeff_mask = gttarget["coeff_mask"]
fhash = gttarget["fhash"]
dbn = get_dbn_coeffs(dbnimg, fhash, dhc_args, coeff_mask = nothing)

s_true = DHC_compute_wrapper(gttarget["true_img"], fhash, norm=false; dhc_args...)
s_init = DHC_compute_wrapper(gttarget["init"], fhash, norm=false; dhc_args...)
s_recon = DHC_compute_wrapper(gttarget["recon"], fhash, norm=false; dhc_args...)
ffb, ffinfo = fink_filter_bank(1, 8, nx=64, Omega=true)

Jfiltidx = J_hashindices([0, 1, 2, 3])
JS1ind = J_S1indices([0, 1, 2, 3])
p25 = distribution_percentiles(dbn, JS1ind, 25)
p75 = distribution_percentiles(dbn, JS1ind, 75)
strueselect = (x->s_true[x]).(JS1ind)
sinitselect = (x->s_init[x]).(JS1ind)
sreconselect = (x->s_recon[x]).(JS1ind)
truebool = (strueselect .> p25) .& (strueselect .<p75)
initbool = (sinitselect .> p25) .& (sinitselect .<p75)
reconbool = (sreconselect .> p25) .& (sreconselect .<p75)

frac_devn_init_true = (sinitselect .- strueselect)./strueselect
frac_devn_recon_true = (sreconselect .- strueselect)./strueselect
dbn_mean = mean(dbn, dims=1)[:]
meanselect = (x->dbn_mean[x]).(JS1ind)
frac_devn_sfd_true = (meanselect .- strueselect)./strueselect
frac_devn_75_25 = (p75 .- p25)./p25

dbn_median = distribution_percentiles(dbn, JS1ind, 50)
frac_devn_medsfd_true = (dbn_median .- strueselect)./strueselect
p= heatmap(frac_devn_init_true, title="Frac Devn Init-True", xlabel="L", ylabel="J", xticks=collect(1:8))
savefig(p, "scratch_NM/NewWrapper/4-10/fdevn-init-true.png")
p=heatmap(frac_devn_sfd_true, title="Frac Devn SFDMean-True", xlabel="L", ylabel="J", xticks=collect(1:8))
savefig(p, "scratch_NM/NewWrapper/4-10/fdevn-sfdmean-true.png")
p= heatmap(frac_devn_75_25, title="Frac Devn 75-25", xlabel="L", ylabel="J", xticks=collect(1:8))
p= heatmap(frac_devn_medsfd_true, title="Frac Devn SFDMedian-True", xlabel="L", ylabel="J", xticks=collect(1:8))
savefig(p, "scratch_NM/NewWrapper/4-10/fdevn-sfdmedian-true.png")
frac_devn_mean_med = (meanselect .- dbn_median)./dbn_median
heatmap(frac_devn_mean_med, title="Frac Devn Mean-Median", xlabel="L", ylabel="J", xticks=collect(1:8))

p99 = distribution_percentiles(dbn, JS1ind, 90)
meanselect .> p99

##WHEN you used the log
fname = "scratch_NM/StandardizedExp/Nx64/noisy_stdtrue/SFDTargSFDCov/reg_apd_noiso/LambdaVary/1000_tunedlam_0-01_fullcov_triu"
gttarget = load(fname*".jld2")
dhc_args = gttarget["dhc_args"]
coeff_mask = gttarget["coeff_mask"]
fhash = gttarget["fhash"]
dbn = log.(get_dbn_coeffs(dbnimg, fhash, dhc_args, coeff_mask = nothing))

s_true = log.(DHC_compute_wrapper(gttarget["true_img"], fhash, norm=false; dhc_args...))
s_init = log.(DHC_compute_wrapper(gttarget["init"], fhash, norm=false; dhc_args...))
s_recon = log.(DHC_compute_wrapper(gttarget["recon"], fhash, norm=false; dhc_args...))
ffb, ffinfo = fink_filter_bank(1, 8, nx=64, Omega=true)

Jfiltidx = J_hashindices([0, 1, 2, 3])
JS1ind = J_S1indices([0, 1, 2, 3])
p25 = distribution_percentiles(dbn, JS1ind, 25)
p75 = distribution_percentiles(dbn, JS1ind, 75)
strueselect = (x->s_true[x]).(JS1ind)
sinitselect = (x->s_init[x]).(JS1ind)
sreconselect = (x->s_recon[x]).(JS1ind)
truebool = (strueselect .> p25) .& (strueselect .<p75)
initbool = (sinitselect .> p25) .& (sinitselect .<p75)
reconbool = (sreconselect .> p25) .& (sreconselect .<p75)

frac_devn_init_true = (sinitselect .- strueselect)./abs.(strueselect)
frac_devn_recon_true = (sreconselect .- strueselect)./abs.(strueselect)
dbn_mean = mean(dbn, dims=1)[:]
meanselect = (x->dbn_mean[x]).(JS1ind)
frac_devn_sfd_true = (meanselect .- strueselect)./abs.(strueselect)
frac_devn_75_25 = (p75 .- p25)./abs.(p25)

dbn_median = distribution_percentiles(dbn, JS1ind, 50)
frac_devn_medsfd_true = (dbn_median .- strueselect)./abs.(strueselect)
p= heatmap(frac_devn_init_true, title="Frac Devn Init-True", xlabel="L", ylabel="J", xticks=collect(1:8))
savefig(p, "scratch_NM/NewWrapper/4-10/LogCoeff/fdevn-init-true.png")
p=heatmap(frac_devn_sfd_true, title="Frac Devn SFDMean-True", xlabel="L", ylabel="J", xticks=collect(1:8))
savefig(p, "scratch_NM/NewWrapper/4-10/LogCoeff/fdevn-sfdmean-true.png")
p= heatmap(frac_devn_75_25, title="Frac Devn 75-25", xlabel="L", ylabel="J", xticks=collect(1:8))
p= heatmap(frac_devn_medsfd_true, title="Frac Devn SFDMedian-True", xlabel="L", ylabel="J", xticks=collect(1:8))
savefig(p, "scratch_NM/NewWrapper/4-10/LogCoeff/fdevn-sfdmedian-true.png")
frac_devn_mean_med = (meanselect .- dbn_median)./abs.(dbn_median)
heatmap(frac_devn_mean_med, title="Frac Devn Mean-Median", xlabel="L", ylabel="J", xticks=collect(1:8))

p99 = distribution_percentiles(dbn, JS1ind, 90)
meanselect .> p99

##Z-scores
Nx=64
dbnimg = readsfd(Nx, logbool=false)
fname = "scratch_NM/StandardizedExp/Nx64/noisy_stdtrue/SFDTargSFDCov/reg_apd_noiso/LambdaVary/1000_tunedlam_0-01_fullcov_triu"
gttarget = load(fname*".jld2")
dhc_args = gttarget["dhc_args"]
coeff_mask = gttarget["coeff_mask"]
fhash = gttarget["fhash"]
dbn = get_dbn_coeffs(dbnimg, fhash, dhc_args, coeff_mask = nothing)

sigdbn = std(dbn, dims=1)

s_true = DHC_compute_wrapper(gttarget["true_img"], fhash, norm=false; dhc_args...)
s_init = DHC_compute_wrapper(gttarget["init"], fhash, norm=false; dhc_args...)
s_recon = DHC_compute_wrapper(gttarget["recon"], fhash, norm=false; dhc_args...)
ffb, ffinfo = fink_filter_bank(1, 8, nx=64, Omega=true)

function J_hashindices(J_values)
    jindlist = []
    for jval in J_values
        push!(jindlist, findall(fhash["J_L"][:, 1].==jval))
    end
    return vcat(jindlist'...)
end

function J_S1indices(J_values)
    #Assumes this is applied to an object of length 2+Nf+Nf^2 or 2+Nf
    return J_hashindices(J_values) .+ 2
end

function zscores(a, b, std, idx)
    return (a[idx] .- b[idx])./std[idx]
end

Jfiltidx = J_hashindices([0, 1, 2, 3])
JS1ind = J_S1indices([0, 1, 2, 3])
strueselect = (x->s_true[x]).(JS1ind)
sinitselect = (x->s_init[x]).(JS1ind)
sreconselect = (x->s_recon[x]).(JS1ind)
stdselect = (x->sigdbn[x]).(JS1ind)

z_init_true = (sinitselect .- strueselect)./stdselect
z_recon_true = (sreconselect .- strueselect)./stdselect
dbn_mean = mean(dbn, dims=1)[:]
meanselect = (x->dbn_mean[x]).(JS1ind)
z_sfdmean_true = (meanselect .- strueselect)./stdselect

function distribution_percentiles(dbn, idx, perc)
    idxperc = zeros(size(idx))
    idxperc .= (x-> percentile(dbn[:, x], perc)).(idx)
    return idxperc
end

dbn_median = distribution_percentiles(dbn, JS1ind, 50)
z_sfdmed_true = (dbn_median .- strueselect)./stdselect

heatmap(z_init_true, title="Z Init-True", xlabel="L", ylabel="J", xticks=collect(1:8))
heatmap(z_sfdmean_true, title="Z SFDMean-True", xlabel="L", ylabel="J", xticks=collect(1:8))
heatmap(z_sfdmed_true, title="Z SFDMedian-True", xlabel="L", ylabel="J", xticks=collect(1:8))
z_mean_med = (meanselect .- dbn_median)./stdselect
heatmap(z_mean_med, title="Z Mean-Median", xlabel="L", ylabel="J", xticks=collect(1:8))

clim = (minimum(z_sfdmed_true), maximum(z_sfdmed_true))
heatmap(z_sfdmed_true, title="Z SFDMedian-True", xlabel="L", ylabel="J", xticks=collect(1:8), clim=clim)
heatmap(z_init_true, title="Z Init-True", xlabel="L", ylabel="J", xticks=collect(1:8), clim=clim)

badz_init_true = z_init_true./stdselect

##Log Coeff-Zscores

logdbn = log.(dbn)
sigdbn = std(logdbn, dims=1)

s_true = log.(DHC_compute_wrapper(gttarget["true_img"], fhash, norm=false; dhc_args...))
s_init = log.(DHC_compute_wrapper(gttarget["init"], fhash, norm=false; dhc_args...))
s_recon = log.(DHC_compute_wrapper(gttarget["recon"], fhash, norm=false; dhc_args...))
ffb, ffinfo = fink_filter_bank(1, 8, nx=64, Omega=true)

Jfiltidx = J_hashindices([0, 1, 2, 3])
JS1ind = J_S1indices([0, 1, 2, 3])
strueselect = (x->s_true[x]).(JS1ind)
sinitselect = (x->s_init[x]).(JS1ind)
sreconselect = (x->s_recon[x]).(JS1ind)
stdselect = (x->sigdbn[x]).(JS1ind)

z_init_true = (sinitselect .- strueselect)./stdselect
z_recon_true = (sreconselect .- strueselect)./stdselect
dbn_mean = mean(logdbn, dims=1)[:]
meanselect = (x->dbn_mean[x]).(JS1ind)
z_sfdmean_true = (meanselect .- strueselect)./stdselect

function distribution_percentiles(dbn, idx, perc)
    idxperc = zeros(size(idx))
    idxperc .= (x-> percentile(dbn[:, x], perc)).(idx)
    return idxperc
end

dbn_median = distribution_percentiles(logdbn, JS1ind, 50)
z_sfdmed_true = (dbn_median .- strueselect)./stdselect

heatmap(z_init_true, title="Z Init-True", xlabel="L", ylabel="J", xticks=collect(1:8))
heatmap(z_sfdmean_true, title="Z SFDMean-True", xlabel="L", ylabel="J", xticks=collect(1:8))
heatmap(z_sfdmed_true, title="Z SFDMedian-True", xlabel="L", ylabel="J", xticks=collect(1:8))
z_mean_med = (meanselect .- dbn_median)./stdselect
heatmap(z_mean_med, title="Z Mean-Median", xlabel="L", ylabel="J", xticks=collect(1:8))

clim = (minimum(z_sfdmed_true), maximum(z_sfdmed_true))
heatmap(z_sfdmed_true, title="Z SFDMedian-True", xlabel="L", ylabel="J", xticks=collect(1:8), clim=clim)
heatmap(z_init_true, title="Z Init-True", xlabel="L", ylabel="J", xticks=collect(1:8), clim=clim)

badz_init_true = z_init_true./stdselect

fnamelog = "scratch_NM/StandardizedExp/Nx64/noisy_stdtrue/SFDTargSFDCov/reg_apd_noiso/LogCoeff/1000_0-01_full_triu"
gtlog = load(fnamelog*".jld2")
s_invcov = gtlog["dict"]["fs_invcov"]
diag(s_invcov)
sigdbn[coeff_mask].^(-2)

##SOME DISCREPANCY BETWEEN THE FSTARG_INVCOV AND THE ONE YOU CALCULATE HERE FOR BOTH S_TARG_INVCOV AND LOGS_TARG_INVCOV:
#Resolved in check_invcov. Numerical instability issues.

Nx=64
dbnimg = readsfd(Nx, logbool=false)
fname = "scratch_NM/StandardizedExp/Nx64/noisy_stdtrue/SFDTargSFDCov/reg_apd_noiso/LambdaVary/1000_tunedlam_0-01_fullcov_triu"
gttarget = load(fname*".jld2")
dhc_args = gttarget["dhc_args"]
coeff_mask = gtlog["coeff_mask"]
fhash = gttarget["fhash"]
dbn = get_dbn_coeffs(dbnimg, fhash, dhc_args, coeff_mask = nothing)

sigdbn = std(dbn, dims=1)
saved = gttarget["dict"]["s_invcov"]
diag(saved)
sigdbn[coeff_mask].^(-2)
(dbn[:, coeff_mask]' * dbn[:, coeff_mask]) ./9999 #roughly consistent with sigdbn.^2

#INCONSISTENT WITH DIAG(SAVED)
datfile = "scratch_NM/StandardizedExp/Nx64/Data_10000.jld2" #Replace w SLURM array
loaddf = load(datfile)
true_img = loaddf["true_img"]
init = loaddf["init"]

filter_hash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)
(S1iso, Nf) = size(filter_hash["S1_iso_mat"])
(S2iso, Nfsq) = size(filter_hash["S2_iso_mat"])

dhc_args = Dict(:doS2=>false, :doS12=>false, :doS20=>true, :apodize=>true, :iso=>false) #Iso #CHANGE: Change sig for the sfd data since the noise model is super high and the tiny values make sense
if dhc_args[:iso]
    coeff_mask = falses(2+S1iso+S2iso)
    coeff_mask[S1iso+3:end] .= true
else #Not iso
    coeff_mask = falses(2+Nf+Nf^2)
    coeff_mask[Nf+3:end] .= true
end

white_noise_args = Dict(:loc=>0.0, :sig=>std(true_img), :Nsam=>1000, :norm=>false, :smooth=>false, :smoothval=>0.8) #Iso #only if you're using noise based covar
optim_settings = Dict([("iterations", 1000), ("norm", false), ("minmethod", ConjugateGradient())])
recon_settings = Dict([("target_type", "sfd_dbn"), ("covar_type", "sfd_dbn"), ("log", false), ("GaussianLoss", true), ("Invcov_matrix", "Diagonal+Eps"),
  ("optim_settings", optim_settings), ("white_noise_args", white_noise_args)]) #Add constraints

if recon_settings["log"] & dhc_args[:apodize]
    lval = std(apodizer(log.(true_img))).^(-2)
elseif !recon_settings["log"] & dhc_args[:apodize]
    lval = std(apodizer(true_img)).^(-2)
elseif recon_settings["log"] & !dhc_args[:apodize]
    lval = std(log.(true_img)).^(-2)
else
    lval = std(true_img).^(-2)
end

recon_settings["lambda"] = lval
coeff_mask = gtlog["coeff_mask"]
regs_true = DHC_compute_wrapper(log.(true_img), filter_hash; norm=false, dhc_args...)[coeff_mask]
s2mean, s2icov = meancov_generator(true_img, filter_hash, dhc_args, coeff_mask, recon_settings, safety=nothing)
diag(s2icov)
##THIS Is consistent but the saved s_targ isn't???? HOW


# it take Z as Gauss DCA
# and returns
# - fi_tens (q,N), fij_tens(q,q,N,N)
# - fi (q*N),fij (q*N*q*N)
# - Meff (effective numb of seq)
# - W (the vector of M weights)

# [normal gaussDCA returns N*(q-1), frequencies]

q = 21

function compute_new_fi_fij(Z::Matrix{Int8}, q, theta)

	#G.C
	#q+1 to get sequences which are 21X21

    W, Meff = GaussDCA.compute_weights(Z, q+1, theta)
    Pi_matr, Pij_matr = compute_all_freqs(Z, W, Meff,q+1)
    Pi_tens, Pij_tens = freq_tens_form(Pi_matr, Pij_matr, Z)

    return Pi_tens, Pij_tens, Pi_matr, Pij_matr, Meff, W
end

function compute_all_freqs(Z::Matrix{Int8}, W::Vector{Float64}, Meff::Float64,q::Int)
    N, M = size(Z)
    s = q - 1

    Ns = N * s

    Pij = zeros(Ns, Ns)
    Pi = zeros(Ns)

    ZZ = Vector{Int8}[vec(Z[i,:]) for i = 1:N]

    i0 = 0
    for i = 1:N
        Zi = ZZ[i]
        for k = 1:M
            a = Zi[k]
            a == q && continue
            Pi[i0 + a] += W[k]
        end
        i0 += s
    end
    Pi /= Meff

    i0 = 0
    for i = 1:N
        Zi = ZZ[i]
        j0 = i0
        for j = i:N
            Zj = ZZ[j]
            for k = 1:M
                a = Zi[k]
                b = Zj[k]
                (a == q || b == q) && continue
                Pij[i0+a, j0+b] += W[k]
            end
            j0 += s
        end
        i0 += s
    end
    for i = 1:Ns
        Pij[i,i] /= Meff
        for j = i+1:Ns
            Pij[i,j] /= Meff
            Pij[j,i] = Pij[i,j]
        end
    end

    return Pi, Pij
end

function freq_tens_form(fi, fij, Z::Matrix{Int8})

    N, M = size(Z)

	#adapt the frequencies
	fi_tens = zeros(q,N)

	for i = 1:N
		fi_tens[:,i] = fi[(1+(q)*(i-1)):((q)+(q)*(i-1))]
	end

	fij_tens= zeros(q,q,N,N);
	for i = 1:N
		for j = 1:N
			fij_tens[:,:,i,j] = fij[(1+(q)*(i-1)):((q)+(q)*(i-1)),(1+(q)*(j-1)):((q)+(q)*(j-1))]
		end
	end

	return fi_tens, fij_tens
end


using ArnoldiMethod

function test_fd_derivatives(a, b, N, ::Type{B}, f::Function, g::Function, h::Function) where {B<:AbstractQuasiMatrix}
    L = b-a
    Δx = L/(N+1)
    j = (1:N) .+ round(Int, a/Δx)

    R = if B == RadialDifferences
        RadialDifferences(N, Δx, 1.0, 0.0)
    else
        B(j, Δx)
    end

    D = Derivative(axes(R,1))

    ∇ = R'D*R
    ∇² = R'D'D*R

    r = axes(R,1)

    fv = R \ f.(r)
    gv = similar(fv)
    hv = similar(fv)

    mul!(gv, ∇, fv)
    mul!(hv, ∇², fv)

    δg = gv - R \ g.(r)
    δh = hv - R \ h.(r)

    r,fv,gv,hv,δg,δh,step(R)
end

function error_slope(loghs,ϵ)
    # To avoid the effect of round-off errors on the order
    # estimation.
    i = argmin(abs.(ϵ)) - 1

    ([loghs[1:i] ones(i)] \ log10.(abs.(ϵ[1:i])))[1]
end

function compute_derivative_errors(a, b, Ns, ::Type{B}, f::Function, g::Function, h::Function) where {B<:AbstractQuasiMatrix}
    errors = map(Ns) do N
        r,fv,gv,hv,δg,δh,Δx = test_fd_derivatives(a, b, N, B, f, g, h)

        [norm(δg)*Δx norm(δh)*Δx]
    end |> e -> vcat(e...)

    ϵg = errors[:,1]
    ϵh = errors[:,2]

    loghs = log10.(1.0 ./ Ns)
    pg = error_slope(loghs, ϵg)
    ph = error_slope(loghs, ϵh)

    ϵg,ϵh,pg,ph
end

function derivative_test_functions(d)
    a,b = √d*[-1,1]

    # The functions need to vanish at the boundaries, for the
    # derivative approximation to be valid (only Dirichlet0 boundary
    # conditions implemented).
    f = x -> exp(-1/(d-x^2))
    g = x -> -2*exp(-1/(d-x^2))*x/((d-x^2)^2)
    h = x -> -2*exp(-1/(d-x^2))*(d^2 + 2*(d-1)*x^2-3x^4)/((d-x^2)^4)

    f,g,h,a,b
end

function norm_rot!(v)
    normalize!(v)
    vc = v.args[2]
    vc[:] *= sign(vc[1])
    v
end

struct ShiftInvert{M}
    A⁻¹::M
end

Base.size(S::ShiftInvert, args...) = size(S.A⁻¹, args...)
Base.eltype(S::ShiftInvert) = eltype(S.A⁻¹)

LinearAlgebra.mul!(y, S::ShiftInvert, x) =
    ldiv!(y, S.A⁻¹, x)

function diagonalize_fd_hamiltonian(H, R::B, nev, σ; method=:arnoldi_shift_invert) where {B<:AbstractQuasiMatrix}
    A,target = if method == :arnoldi_shift_invert
        ShiftInvert(factorize(H - σ*I)),LR()
    else
        H, SR()
    end
    schur,history = partialschur(A, nev=nev, which=target)
    println(history)

    ϕ = [norm_rot!(R*schur.Q[:,j]) for j = 1:nev]

    θ = diag(schur.R)
    λ = if method == :arnoldi_shift_invert
        σ .+ inv.(θ)
    else
        θ
    end

    λ,ϕ
end

function get_kinetic_operator(R::B) where {B<:AbstractQuasiMatrix}
    D = Derivative(Base.axes(R,1))
    apply(*, R', D', D, R) / -2
end

function test_fd_particle_in_a_box(::Type{B}, N, L, nev; kwargs...) where {B<:AbstractQuasiMatrix}
    R = B(1:N, L/(N+1))
    display(R)
    println()

    Tm = get_kinetic_operator(R)
    λ,ϕ = diagonalize_fd_hamiltonian(Tm, R, nev, 0.0; kwargs...)

    r̃ = FiniteDifferencesQuasi.locs(R)

    n = 1:nev
    δλ = λ - n.^2*π^2/(2L^2)

    # Could/should also test eigenvectors

    λ,ϕ,r̃,R,δλ
end

function test_singular_fd_scheme(::Type{B}, N, rₘₐₓ, Z, ℓ, nev; kwargs...) where {B<:AbstractQuasiMatrix}
    R = if B == RadialDifferences
        ρ = rₘₐₓ/(N+0.5)
        RadialDifferences(N, ρ, Z)
    else
        ρ = rₘₐₓ/(N+1)
        NumerovFiniteDifferences(1:N, ρ, true, Z)
    end
    display(R)
    println()

    r = axes(R,1)

    n = 1:nev
    λₐ = -inv.(2(n .+ ℓ).^2)

    Tm = get_kinetic_operator(R) + R'*QuasiDiagonal(ℓ*(ℓ+1)./2r.^2)*R
    V = R'*QuasiDiagonal(-inv.(r))*R
    H = Tm + V
    λ,ϕ = diagonalize_fd_hamiltonian(H, R, nev, 1.1λₐ[1]; kwargs...)

    r̃ = FiniteDifferencesQuasi.locs(R)

    δλ = λ - λₐ

    # Could/should also test eigenvectors

    λ,ϕ,r̃,R,δλ
end

function compute_diagonalization_errors(f::Function, ::Type{B}, Ns, args...; kwargs...) where B
    errors = map(Ns) do N
        println("N = $N")
        t = time()
        λ,ϕ,r̃,R,δλ = f(B, N, args...; kwargs...)
        elapsed = time()-t
        vcat(δλ,elapsed)'
    end |> e -> vcat(e...)

    errors,elapsed = errors[:,1:end-1],errors[:,end]

    loghs = log10.(inv.(Ns))
    slopes = map(1:size(errors,2)) do j
        error_slope(loghs, errors[:,j])
    end

    errors,slopes,elapsed
end


"""
    parse_node(node)

take a node and parse into Symbolics form
"""
function parse_node(node)
    tagmap[node.name](node)
end

function parse_str(str)
    doc = parsexml(str)
    parse_doc(doc)
end

function parse_doc(doc)
    node = doc.root
    parse_node(node)
end

function parse_file(fn)
    node = readxml(fn).root
    parse_node(node)
end

"""
    parse_cn(node)

parse a <cn> node
"""
function parse_cn(node)
    # elements(node) != EzXML.Node[] && error("node cant have elements rn, check if </sep> is in the node")
    # this isnt great might want to check that `sep` actually shows up
    if haskey(node, "type") && elements(node) != EzXML.Node[]
        parse_cn_w_sep(node)
    else
        Float64(Meta.parse(node.content))   # convert to Float64 for CellML compatibility
    end
end

"""
    parse_cn_w_sep(node)

parse a <cn type=".."> node

where type ∈ ["e-notation", "rational", "complex-cartesian", "complex-polar"]
"""
function parse_cn_w_sep(node)
    # node = clean_attributes(node)
    # txts = findall("text()", node)
    txts = [n for n in eachnode(node) if istext(n)]
    length(txts) != 2 && error("stop, collaborate, and listen!, problem with <cn>")
    x1, x2 = map(x -> Meta.parse(x.content), txts)
    t = node["type"]
    if t == "e-notation"
        x1 * exp10(x2)
    elseif t == "rational"
        Rational(x1, x2)
    elseif t == "complex-cartesian"
        Complex(x1, x2)
    elseif t == "complex-polar"
        x1 * exp(x2 * im)
    else
        error("$t in parse_cn_w_sep, somethings wrong")
    end
end

"""
    parse_ci(node)

parse a <ci> node
"""
function parse_ci(node)
    c = Meta.parse(strip(node.content))
    Num(Variable(c))
end

########## Parse piecewise ###################################################

function parse_piecewise(node)
    ns = elements(node)
    pieces = filter(x -> nodename(x) == "piece", ns)
    others = filter(x -> nodename(x) == "otherwise", ns)

    if length(others) > 0
        otherwise = parse_node(firstelement(others[1]))
    else
        otherwise = 0.0
    end

    return process_pieces(pieces, otherwise)
end

function process_pieces(pieces, otherwise)
    node = pieces[1]
    c = parse_node.(elements(node))
    return IfElse.ifelse(c[2] > 0.5, c[1],
                length(pieces)==1 ? otherwise : process_pieces(pieces[2:end], otherwise))
end

"""
    parse_apply(node)

parse an <apply> node into Symbolics form

how to deal w apply within apply, need to ensure we've hit bottom
"""
function parse_apply(node)
    node.name != "apply" && error("calling parse_apply requires the name of the element to be `apply`")
    elms = elements(node)
    cs = parse_node.(elms[2:end])
    applymap[elms[1].name](cs)
end

"""
    parse_bvar(node)

parse a <bvar> node
"""
function parse_bvar(node)
    es = elements(node)
    length(es) == 1 ? (parse_node(es[1]), 1) : Tuple(parse_node.(es))
end

"""
    parse_diff(x)

parse a <diff>

"""
function parse_diff(a)
    (iv, deg), x = a
    deg = trunc(Int, deg)
    # num = Num(Symbolics.Sym{Symbolics.FnType{Tuple{Real},Real}}(Symbol(x))(iv))
    D = Differential(iv)^deg
    D(x)
end

"""
    parse_lambda(node)

parse a <lambda> node

```xml
<lambda>
  <bvar> x1 </bvar><bvar> xn </bvar>
   expression-in-x1-xn
</lambda>
```
"""
function parse_lambda(node)
    es = elements(node)
    vars = findall("//x:bvar | //bvar", node, ["x" => MathML.mathml_ns])
    # vars2 = findall("//bvar", node) # works in tests
    # vars = union(vars, vars2) # FIX

    args = first.(parse_bvar.(vars))
    num = parse_apply(es[end])
    eval(build_function([num], args...)[1])
end


"""
    MushroomTools
Module containing many functions helpful in simulating (perfect) mushroom
billiards, see [`billiard_mushroom`](@ref).
Contains stuff like initializing efficiently regular or chaotic particles
and functions that return the corresponding chaotic or regular phase-space volumes
or portions.
The functions [`V_3D_tot`](@ref) and [`V_3D_reg`](@results) use equations derived
in ref. [1].

Made by Lukas Hupe.

## References
[1] A. Barnett & T. Betcke, [Chaos **17**, 043125 (2007)](https://doi.org/10.1063/1.2816946).
"""
module MushroomTools

using DynamicalBilliards
using DynamicalBilliards: SV, cossin
using LinearAlgebra

#="""
    is_regular(p::Particle, o::Semicircle, w)
checks whether the particle `p` is on an integrable orbit in a mushroom billiard
with cap `o` and stem width `w`
"""=#
function is_regular(p::Particle{T}, l, w, r) where {T}
    δ = SV{T}(0.0, l) - p.pos
    vperp = [-p.vel[2],p.vel[1]]
    if norm(δ) > r
        #debug && println("not inside circle: pos = $(p.pos)")
        return false
    elseif p.pos[2] < l
        #debug && println("wrong half of circle")
        return false
    elseif abs(dot(vperp, δ)) < w/2
        #debug && println("not outside critical radius")
        return false
    else
        #debug && println("Perfectly happy with $(p.pos) and $(p.vel)")
        return true
    end
end

#Formulae from quantum mushroom paper
"""
    V_3D_tot(l,w,r)
Return the total phasespace volume (3D) of a [`billiard_mushroom`](@ref)
parameterized by `(l,w,r)`.
"""
V_3D_tot(l,w,r) = 2π*(l*w + (1/2)*π*r^2)
V_3D_reg(l,w,r) = 2π*r^2*(acos(w/(2r)) - w/(2r)*sqrt(1 - (w^2)/(4r^2)))
V_3D_cha(l,w,r) = V_3D_tot(l,w,r) - V_3D_reg(l,w,r)

"""
    g_r_3D(l, w, r)
Return the regular phasespace portion of the full (3D) phase-space of a
[`billiard_mushroom`](@ref) with stem length `l`, stem width `w` and cap radious `r`.
This result is known analytically, see [`MushroomTools`](@ref) for references.
"""
g_r_3D(l,w,r) = V_3D_reg(l,w,r)/V_3D_tot(l,w,r)

"""
    g_c_3D(l, w, r)
Return the chaotic phasespace portion of the full (3D) phase-space of a
[`billiard_mushroom`](@ref) with stem length `l`, stem width `w` and cap radious `r`.
This result is known analytically, see [`MushroomTools`](@ref) for references.
"""
g_c_3D(l,w,r) = 1 - g_r_3D(l,w,r)

"""
    V_2D_tot(l,w,r)
Return the total boundary map volume (2D) of a [`billiard_mushroom`](@ref)
parameterized by `(l,w,r)`.
"""
V_2D_tot(l, w, r) = 2(π*r + 2r + 2l)
V_2D_reg(l, w, r) = 2π*r*(1 - w/2r) + 2sqrt(4r*r - w*w) -2w*acos(w/2r)
V_2D_cha(l,w,r) = V_2D_tot(l,w,r) - V_2D_reg(l,w,r)

"""
    g_r_2D(l, w, r)
Return the regular phasespace portion of the boundary map (2D) of a
[`billiard_mushroom`](@ref) with stem length `l`, stem width `w` and cap radious `r`.
This result is known analytically, see [`MushroomTools`](@ref) for references.
"""
g_r_2D(l,w,r) = V_2D_reg(l,w,r)/V_2D_tot(l,w,r)

"""
    g_c_2D(l, w, r)
Return the chaotic phasespace portion of the boundary map (2D) of a
[`billiard_mushroom`](@ref) with stem length `l`, stem width `w` and cap radious `r`.
This result is known analytically, see [`MushroomTools`](@ref) for references.
"""
g_c_2D(l,w,r) = 1 - g_r_2D(l,w,r)

#applying Kac's Lemma to the stem
τ_r_2D_kac(l,w,r) = (V_2D_tot(l,w,r) - V_2D_reg(l,w,r)) / (2w + 4l)

#="""
    insidemushroom(pos::StaticVector, l, w, r)
Returns `true` if pos is within the mushroom parameterised by `l`,  `w` and `r`
"""=#
function insidemushroom(pos::SV{T}, l::T, w::T, r::T) where {T <: AbstractFloat}
    if pos[2] > 0
        if pos[2] <= l
            if abs(pos[1]) <= w/2
                return true
            end
        elseif pos[2] <= l + r
            if pos[1]^2 + (pos[2] - l)^2 <= r^2
                return true
            end
        end
    end
    return false
end

"""
    randin_mushroom(l, w, r [, ω])
Generate a random particle within the [`billiard_mushroom`](@ref)
parameterised by `l`, `w` and `r`.
If `ω` is given the particle is magnetic instead.

This function is much more efficient than [`randominside`](@ref).
"""
function randin_mushroom(l::T = 1.0, w::T = 0.2, r::T = 1.0) where {T <: AbstractFloat}
    pos, vel = _randin_mushroom(l, w, r)
    return Particle(pos, vel, zero(SV{T}))
end
function randin_mushroom(l::T, w::T, r::T, ω::T) where {T <: AbstractFloat}
    pos, vel = _randin_mushroom(l, w, r)
    return MagneticParticle(pos, vel, zero(SV{T}), T(ω))
end
function _randin_mushroom(l::T = 1.0, w::T = 0.2, r::T = 1.0) where {T <: AbstractFloat}
    x = (-r, r, 2r)
    y = (0, l + r, l + r)
    pos = SV{T}(rand(T)*x[3] + x[1], rand(T)*y[3] + y[1])
    while ! insidemushroom(pos, l, w, r)
        pos = SV{T}(rand(T)*x[3] + x[1], rand(T)*y[3] + y[1])
    end

    φ = T(2π * rand(T))
    vel = SV{T}(cossin(φ)...)

    return pos, vel
end

"""
    randomchaotic(l, w, r)
Generate a chaotic particle, i.e. not trapped in the cap.
"""
function randomchaotic(l, w, r)
    p = randin_mushroom(l, w, r)
    while is_regular(p, l, w, r)
        p = randin_mushroom(l, w, r)
    end
    return p
end

"""
    randomregular(l, w, r)
Generate a regular particle (i.e. trapped in the cap).
"""
function randomregular(l, w, r)
    p = randin_mushroom(l, w, r)
    while !is_regular(p, l, w, r)
        p = randin_mushroom(l, w, r)
    end
    return p
end

end # module


using Plots, Printf, LinearAlgebra

# enable plotting by default
if !@isdefined do_visu; do_visu = true end

# finite-difference support functions
@views av_xi(A) = 0.5*(A[1:end-1,2:end-1].+A[2:end,2:end-1])
@views av_yi(A) = 0.5*(A[2:end-1,1:end-1].+A[2:end-1,2:end])
@views   inn(A) = A[2:end-1,2:end-1]

@views function diffusion_2D_impl(; do_visu=true, save_fig=false)
    # Physics
    lx, ly = 10.0, 10.0   # domain size
    npow   = 3            # power-law exponent
    ttot   = 1.0          # total simulation time
    dt     = 0.2          # physical time step
    # Numerics
    nx, ny = 128, 128     # numerical grid resolution
    tol    = 1e-6         # tolerance
    itMax  = 1e5          # max number of iterations
    # Derived numerics
    dx, dy = lx/nx, ly/ny # grid size
    xc, yc = LinRange(dx/2, lx-dx/2, nx), LinRange(dy/2, ly-dy/2, ny)
    # Array allocation
    qHx    = zeros(nx-1, ny-2) # on staggered grid
    qHy    = zeros(nx-2, ny-1) # on staggered grid
    ResH   = zeros(nx-2, ny-2) # normal grid, without boundary points
    dHdtau = zeros(nx-2, ny-2) # normal grid, without boundary points
    dtau   = zeros(nx-2, ny-2) # normal grid, without boundary points
    # Initial condition
    H      = exp.(.-(xc.-lx/2).^2 .-(yc.-ly/2)'.^2)
    Hold   = copy(H)
    t = 0.0; it = 0; ittot = 0
    # Physical time loop
    while t<ttot
        iter = 0; err = 2*tol
        # Picard-type iteration
        while err>tol && iter<itMax
            qHx    .= -av_xi(H).^npow.*diff(H[:,2:end-1], dims=1)/dx  # flux
            qHy    .= -av_yi(H).^npow.*diff(H[2:end-1,:], dims=2)/dy  # flux
            ResH   .= -(inn(H) - inn(Hold))/dt +
                       (-diff(qHx, dims=1)/dx -diff(qHy, dims=2)/dy)  # residual of the PDE
            dHdtau .= ResH                                            # rate of change
            dtau   .= (1.0./(min(dx, dy)^2 ./inn(H).^npow./4.1) .+ 1.0/dt).^-1  # time step (obeys ~CFL condition)
            H[2:end-1,2:end-1] .= inn(H) .+ dtau.*dHdtau              # update rule, sets the BC as H[1]=H[end]=0
            iter += 1; err = norm(ResH)/length(ResH)
        end
        ittot += iter; it += 1; t += dt
        Hold .= H
    end
    @printf("Total time = %1.2f, time steps = %d, iterations tot = %d \n", round(ttot, sigdigits=2), it, ittot)
    # Visualize
    if do_visu
        fontsize = 12
        opts = (aspect_ratio=1, yaxis=font(fontsize, "Courier"), xaxis=font(fontsize, "Courier"),
                ticks=nothing, framestyle=:box, titlefontsize=fontsize, titlefont="Courier", colorbar_title="",
                xlabel="Lx", ylabel="Ly", xlims=(xc[1],xc[end]), ylims=(yc[1],yc[end]), clims=(0.,1.))
        display(heatmap(xc, yc, H'; c=:davos, title="implicit diffusion (nt=$it, iters=$ittot)", opts...))
        if save_fig savefig("diff2D_impl.png") end
    end
    return xc, yc, H
end

diffusion_2D_impl(; do_visu=do_visu);


module classopath_test

using Test, ConstrainedLasso, ECOS, Random

@info("Test lsq_classopath: sum-to-zero constraint")

# set up
Random.seed!(456)
n, p = 100, 20
# truth with sum constraint sum(β) = 0
β = zeros(p)
β[1:round(Int, p / 4)] .= 0
β[(round(Int, p / 4) + 1):round(Int, p / 2)] .= 1
β[(round(Int, p / 2) + 1):round(Int, 3p / 4)] .= 0
β[(round(Int, 3p / 4) + 1):p] .= -1
# generate data
X = randn(n, p)
y = X * β + randn(n)
# equality constraints
Aeq = ones(1, p)
beq = [0.0]
penwt  = ones(p)

β̂path1, ρpath1, objpath, = lsq_classopath(X, y; Aeq = Aeq, beq = beq)
# print(β̂path1)
# print(size(β̂path1))
# print(sum(β̂path1,1))
@test all(abs.(sum(β̂path1, 1)) .< 1.0e-6)

@info("Test lsq_classopath: non-negativity constraint")

# set up
n, p = 20, 100
# truth with a few non-zero coefficients
β = zeros(p)
β[1:10] = 1:10
# generate data
Random.seed!(123)
X = randn(n, p)
y = X * β + randn(n)
# inequality constraints
Aineq = - eye(p)
bineq = zeros(p)

solver = ECOSSolver(verbose=0, maxit=1e8)

β̂path2, ρpath2, = lsq_classopath(X, y; Aineq = Aineq, bineq = bineq, solver=solver)

@testset "non-negativity" begin
for i in reshape(β̂path2, p * length(ρpath2), 1)
  @test i >= 0
end
end

end # classopath_test module


#=
# 1D advection matching the example in the book: Nodal Discontinuous Galerkin Method.
=#

### If the Finch package has already been added, use this line #########
using Finch # Note: to add the package, first do: ]add "https://github.com/paralab/Finch.git"

### If not, use these four lines (working from the examples directory) ###
# if !@isdefined(Finch)
#     include("../Finch.jl");
#     using .Finch
# end
##########################################################################

init_finch("advection1d");

useLog("dgnewadvection1dlog", level=3)

domain(1)
solverType(DG)
functionSpace(order=3)
timeStepper(LSRK4, cfl=0.01)

mesh(LINEMESH, elsperdim=10, bids=2, interval=[0,2])

u = variable("u")
testSymbol("v")

T = 1;
timeInterval(T)
initial(u, "sin(x)")

boundary(u, 1, DIRICHLET, "-sin(2*pi*t)")
boundary(u, 2, NO_BC)

coefficient("a", 2*pi)
coefficient("beta", 0.2)

###################################################
# Uncomment one of the following two groups of code(or both)

### To generate the Julia code and export it to a set of files, use these two lines.
weakForm(u, "Dt(u*v) + a*grad(u)*v + surface(-a*normal()*u*v + a*normal()*ave(u)*v - 0.5*(1-beta)*a*jump(u)*v)")
exportCode("advec1d")
###

### To import previously generated or modified code, use this line
#importCode("advec1d")
###

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

solve(u);

# Uncomment to plot
# using Plots
# pyplot();
# display(plot(Finch.grid_data.allnodes[1,:], u.values[1,:], marker=:circle, reuse=false))

# @finalize()


using Test
using Combinatorics
using Assignment

find_best(M, k) = size(M, 1) > size(M, 2) ? find_best(M', k) : partialsort([
    (sum(M[CartesianIndex.(axes(M, 1), p)]), p) for p in permutations(axes(M, 2), size(M, 1))], 1:k)

function test_kbest(M, k, sols)
    for ((cost, perm), sol) in zip(find_best(M, k), sols)
        @test cost ≈ sol.cost
        @test perm == (sol.use_row ? sol.row4col : sol.col4row)
    end
end

@testset "find best assignment" begin
      @testset "Float64" begin
          A = [ 0.891171  0.0320582   0.564188  0.8999    0.620615;
                0.166402  0.861136    0.201398  0.911772  0.0796335;
                0.77272   0.782759    0.905982  0.800239  0.297333;
                0.561423  0.170607    0.615941  0.960503  0.981906;
                0.748248  0.00799335  0.554215  0.745299  0.42637]

          sol = @inferred find_best_assignment(A)
          @test sol.row4col == [2, 3, 5, 1, 4]
          @test sol.cost ≈ 1.8375112
          @test sum(A[sol]) ≈ 1.8375112
      end

      @testset "Int64" begin
          B = [ 24     1     8;
                 5     7    14;
                 6    13    20;
                12    19    21;
                18    25     2]

          sol = @inferred find_best_assignment(B)
          @test sol.row4col == [2, 1, 0, 0, 3]
          @test sol.col4row == [2, 1, 5]
          @test sol.cost == 8
          @test sum(B[sol]) == 8

          sol = @inferred find_best_assignment(B')
          @test sol.col4row == [2, 1, 0, 0, 3]
          @test sol.row4col == [2, 1, 5]
          @test sol.cost == 8
          @test sum(B'[sol]) == 8

          sol = @inferred find_best_assignment([ifelse(i-j==0,0,i*j) for i in 1:5, j in 1:5])
          @test sol.col4row == [1, 2, 3, 4, 5]
          @test sol.row4col == [1, 2, 3, 4, 5]
          @test sol.cost == 0
      end

      @testset "UInt16" begin
        M=UInt16[28092 44837 19882 39481 59139;
                 26039 46258 38932 51057 9;
                 11527 59487 61993 29072 8734;
                 10691 16977 12796 16370 14266;
                 5199  42319 34194 41332 16472]
        sol = @inferred find_best_assignment(M)
        @test sol.row4col == [3, 5, 4, 2, 1]
        @test sum(M[sol]) == 71139
        @test sol.cost    == 71139 % 2^16
    end

    @testset "UInt8" begin
        M=UInt8[67  228 135 197 244;
                112 44  84  206 31;
                225 103 231 225 227;
                170 37  135 9   130;
                110 22  133 77  96]

        sol = @inferred find_best_assignment(M)
        @test sol.row4col == [1, 5, 2, 4, 3]
        @test sum(M[sol]) == 343
        @test sol.cost   == 343 % 2^8
    end
end

@testset "find kbest assignments" begin
    @testset "Float64" begin
        A = [ 0.891171  0.0320582   0.564188  0.8999    0.620615;
              0.166402  0.861136    0.201398  0.911772  0.0796335;
              0.77272   0.782759    0.905982  0.800239  0.297333;
              0.561423  0.170607    0.615941  0.960503  0.981906;
              0.748248  0.00799335  0.554215  0.745299  0.42637]

        sols = @inferred find_kbest_assignments(A, 10)
        sol = sols[1]
        @test sol.row4col == [2, 3, 5, 1, 4]
        @test sol.cost ≈ 1.8375112
        @test sum(A[sol]) ≈ 1.8375112
        test_kbest(A, 10, sols)
    end

    @testset "Int64" begin
        B = [ 24     1     8;
               5     7    14;
               6    13    20;
              12    19    21;
              18    25     2]

        sols = @inferred find_kbest_assignments(B, 10)
        sol = sols[1]
        @test sol.row4col == [2, 1, 0, 0, 3]
        @test sol.col4row == [2, 1, 5]
        @test sol.cost == 8
        @test sum(B[sol]) == 8
        test_kbest(B, 10, sort(sols, by=x -> (x.cost, x.col4row)))

        sols = @inferred find_kbest_assignments(B', 10)
        sol = sols[1]
        @test sol.col4row == [2, 1, 0, 0, 3]
        @test sol.row4col == [2, 1, 5]
        @test sol.cost == 8
        @test sum(B'[sol]) == 8
        test_kbest(B', 10, sort(sols, by=x -> (x.cost, x.row4col)))

        C = [ifelse(i-j==0,0,i*j) for i in 1:5, j in 1:5]
        sols = @inferred find_kbest_assignments(C, 10)
        sol = sols[1]
        @test sol.col4row == [1, 2, 3, 4, 5]
        @test sol.row4col == [1, 2, 3, 4, 5]
        @test sol.cost == 0
        test_kbest(C, 10, sort(sols, by=x -> (x.cost, x.row4col)))
    end

    @testset "UInt16" begin
      M=UInt16[28092 44837 19882 39481 59139;
               26039 46258 38932 51057 9;
               11527 59487 61993 29072 8734;
               10691 16977 12796 16370 14266;
               5199  42319 34194 41332 16472]
      sols = @inferred find_kbest_assignments(M, 10)
      sol = sols[1]
      @test sol.row4col == [3, 5, 4, 2, 1]
      @test sum(M[sol]) == 71139
      @test sol.cost    == 71139 % 2^16
  end

  @testset "UInt8" begin
      M=UInt8[67  228 135 197 244;
              112 44  84  206 31;
              225 103 231 225 227;
              170 37  135 9   130;
              110 22  133 77  96]

      sols = @inferred find_kbest_assignments(M, 10)
      sol = sols[1]
      @test sol.row4col == [1, 5, 2, 4, 3]
      @test sum(M[sol]) == 343
      @test sol.cost   == 343 % 2^8
  end
end


#! /usr/bin/julia

# Rosetta Code, AKS test for primes

function polycoefs(n::Int64)
    pc = typeof(n)[]
    if n < 0
        return pc
    end
    sgn = one(n)
    for k in n:-1:0
        push!(pc, sgn*binomial(n, k))
        sgn = -sgn
    end
    return pc
end

function stringpoly(n::Int64)
    if n < 0
        return ""
    end
    st = @sprintf "(x - 1)^{%d} & = & " n
    for (i, c) in enumerate(polycoefs(n))
        if i == 1
            op = ""
            ac = c
        elseif c < 0
            op = "-"
            ac = abs(c)
        else
            op = "+"
            ac = abs(c)
        end
        p = n + 1 - i
        if p == 0
            st *= @sprintf " %s %d\\\\" op ac
        elseif ac == 1
            st *= @sprintf " %s x^{%d}" op p
        else
            st *= @sprintf " %s %dx^{%d}" op ac p
        end
    end
    return st
end

function isaksprime(n::Int64)
    if n < 2
        return false
    end
    for c in polycoefs(n)[2:(end-1)]
        if c%n != 0
            return false
        end
    end
    return true
end

println("<math>")
println("\\begin{array}{lcl}")
for i in 0:10
    println(stringpoly(i))
end
println("\\end{array}")
println("</math>\n")

L = 50
print("AKS primes less than ", L, ":  ")
sep = ""
for i in 1:L
    if isaksprime(i)
        print(sep, i)
        sep = ", "
    end
end
println()


"""
    simplify(expr::Expr)

Simplify expression.
"""
function simplify(ex::Expr)
    @assert ex.head == :call
    return simplify(Val{ex.args[1]}, ex)
end

"""
    simplify(f::Union{Number, Symbol})
"""
function simplify(f::Union{Number, Symbol})
    return f
end

"""
    simplify(Val{:*}, ex::Expr)
"""
function simplify(::Type{Val{:*}}, ex::Expr)
    args = simplify.(ex.args[2:end])
    0 in args && return 0
    filter!(k -> !(isa(k, Number) && k == 1), args)
    length(args) == 1 && return first(args)
    return Expr(:call, :*, args...)
end

"""
    simplify(Val{:+}, ex::Expr)
"""
function simplify(::Type{Val{:+}}, ex::Expr)
    args = simplify.(ex.args[2:end])
    filter!(k -> !isa(k, Number) || k != 0, args)
    length(args) == 1 && return first(args)
    return Expr(:call, :+, args...)
end

"""
    simplify(Val{:-}, ex::Expr)
"""
function simplify(::Type{Val{:-}}, ex::Expr)
    args = simplify.(ex.args[2:end])
    filter!(k -> !isa(k, Number) || k != 0, args)
    length(args) == 1 && return first(args)
    return Expr(:call, :-, args...)
end

"""
    simplify(Val{:^}, ex::Expr)
"""
function simplify(::Type{Val{:^}}, ex::Expr)
    args = simplify.(ex.args[2:end])
    return Expr(:call, :^, args...)
end

"""
    simplify(Val{:/}, ex::Expr)
"""
function simplify(::Type{Val{:/}}, ex::Expr)
    args = simplify.(ex.args[2:end])
    return Expr(:call, :/, args...)
end


# Test external function calls / semantic attachments
path = joinpath(dirname(pathof(PDDL)), "..", "test", "functions")

# Load domain and define external functions
domain = load_domain(joinpath(path, "domain.pddl"))
throw_range(v, θ) = v^2*sin(2*θ*π/180) / 9.81
domain.funcdefs[:range] = throw_range
throw_height(v, θ, x) = tan(θ*π/180)*x - 9.81*x^2 / (2*v^2 * cos(θ*π/180)^2)
domain.funcdefs[:height] = throw_height

# Load problem
problem = load_problem(joinpath(path, "problem.pddl"))
state = init_state(problem)

# Check that evaluation with external functions works correctly
@test state[domain, pddl"(range 20 45)"] == throw_range(20, 45)
@test state[domain, pddl"(height 20 45 10)"] == throw_height(20, 45, 10)

# Execute plan
state = init_state(problem)
state = execute(pddl"(pick ball1)", state, domain)
state = execute(pddl"(throw ball1 85)", state, domain)
state = execute(pddl"(pick ball2)", state, domain)
state = execute(pddl"(throw ball2 75)", state, domain)

# Check if goal is satisfied
@test state[pddl"(< (loc ball1) 10)"]
@test state[pddl"(> (loc ball2) 15)"]
@test satisfy(problem.goal, state, domain)[1] == true


using DelimitedFiles

function move(pos, cmd, dir)
	w = cmd[1]
	n = parse(Int, cmd[2:end])
	if w == 'N'
		pos += [n,0]
	elseif w == 'S'
		pos -= [n,0]
	elseif w == 'E'
		pos += [0,n]
	elseif w == 'W'
		pos -= [0,n]
	elseif w == 'L'
		dir = [cosd(n) sind(n);-sind(n) cosd(n)]*dir
	elseif w == 'R'
		dir = [cosd(n) -sind(n);sind(n) cosd(n)]*dir
	elseif w == 'F'
		pos += n*dir
	end
	pos = round.(pos)
	dir = round.(dir)
	(Int.(pos),Int.(dir))
end

# returns solution of part 1
function x()
	t = readdlm("12.txt", String)
	pos = [0,0]
	dir = [0,1]
	for c in t
		(pos, dir) = move(pos, c, dir)
	end
	abs(pos[1])+abs(pos[2])
end

function move2(spos, wpos, cmd)
	w = cmd[1]
	n = parse(Int, cmd[2:end])
	if w == 'N'
		wpos += [n,0]
	elseif w == 'S'
		wpos -= [n,0]
	elseif w == 'E'
		wpos += [0,n]
	elseif w == 'W'
		wpos -= [0,n]
	elseif w == 'L'
		diff = wpos - spos
		diff = [cosd(n) sind(n);-sind(n) cosd(n)]*diff
		wpos = spos + diff
	elseif w == 'R'
		diff = wpos - spos
		diff = [cosd(n) -sind(n);sind(n) cosd(n)]*diff
		wpos = spos + diff
	elseif w == 'F'
		mv = wpos - spos
		spos += n*mv
		wpos += n*mv
	end
	spos = round.(spos)
	wpos = round.(wpos)
	(Int.(spos),Int.(wpos))
end

# returns solution of part 2
function y()
	t = readdlm("12.txt", String)
	spos = [0,0]
	wpos = [1,10]
	for c in t
		(spos, wpos) = move2(spos, wpos, c)
	end
	abs(spos[1])+abs(spos[2])
end


export BlackHole,
    compute_eddington_luminosity,
    compute_bolometric_luminosity,
    compute_mass_accretion_rate

struct BlackHole{T<:AbstractFloat}
    M::T
    mdot::T
    spin::T
    isco::T
    efficiency::T
    Rg::T
    function BlackHole(M, mdot, spin)
        Rg = G * M / C^2
        isco = compute_isco(spin)
        eff = compute_efficiency(spin)
        new{typeof(M)}(M, mdot, spin, isco, eff, Rg)
    end
end

function BlackHole(parameters::Parameters)
    M = parameters.M * M_SUN
    return BlackHole(M, parameters.mdot, parameters.spin)
end

"""
Computes the Eddington luminosity for the given black hole mass.
"""
function compute_eddington_luminosity(bh_mass; mu_electron = 1.17)
    return 4π * G * bh_mass * C * mu_electron / SIGMA_E
end
compute_eddington_luminosity(bh::BlackHole; mu_electron = 1.17) =
    compute_eddington_luminosity(bh.M, mu_electron = mu_electron)

"""
Computes the bolometric luminosity of the AGN given the black
hole mass and the normalized eddington rate.
"""
compute_bolometric_luminosity(bh_mass, eddington_rate; mu_electron = 1.17) =
    eddington_rate * compute_eddington_luminosity(bh_mass; mu_electron = mu_electron)
compute_bolometric_luminosity(bh::BlackHole; mu_electron = 1.17) =
    compute_bolometric_luminosity(bh.M, bh.mdot, mu_electron = 1.17)

"""
Computes the Innermost Stable Circular Orbit (ISCO) in units of Rg.

# Parameters
- spin : normalized spin parameter ∈ (-1, 1)
"""
function compute_isco(spin)
    z1 = 1 + (1 - spin^2)^(1 / 3) * ((1 + abs(spin))^(1 / 3) + (1 - abs(spin))^(1 / 3))
    z2 = sqrt(3 * spin^2 + z1^2)
    rms = 3 + z2 - sign(spin) * sqrt((3 - z1) * (3 + z1 + 2 * z2))
    return rms
end

"""
Computes the black hole accretion efficiency parameter from the isco.

# Parameters
- spin : black hole spin parameter ∈ (-1, 1)
"""
function compute_efficiency(spin)
    isco = compute_isco(spin)
    return 1 - sqrt(1 - 2 / (3 * isco))
end

"""
Computes the mass accretion rate to the black hole,
assuming no wind losses.

# Parameters
- M : black hole mass [in grams]
- mdot : normalized eddington rate ∈ [0,1]
- spin : black hole spin parameter ∈ (-1, 1)
"""
function compute_mass_accretion_rate(M, mdot, spin; mu_electron = 1.17)
    efficiency = compute_efficiency(spin)
    return compute_bolometric_luminosity(M, mdot, mu_electron = mu_electron) /
           (efficiency * C^2)
end

function compute_mass_accretion_rate(bh; mu_electron = 1.17)
    return compute_mass_accretion_rate(bh.M, bh.mdot, bh.spin, mu_electron = mu_electron)
end


"""
Quadrature information

Most spectral-element operations require calculation of integrals and derivatives.
The `QuadType` structure stores the required information when calculating integrals
and derivatives
"""
type QuadType{T<:Number}
    Q::Int
    z::Vector{T}
    w::Vector{T}
    D::Matrix{T}
end

"""
Constructs quadrature rules for Gauss-type quadrature.
"""
function QuadType{T<:Number,QT<:QUADRATURE_TYPE}(q::Int, ::Type{QT}=GLJ, ::Type{T}=Float64, a=0, b=0)
    z = Jacobi.qzeros(QT, q, a, b, T)
    w = Jacobi.qweights(QT, z, a, b)
    D = Jacobi.qdiff(QT, z, a, b)
    QuadType(q, z, w, D)
end

"""
Creates the ith Lagrange polynomial from a vector of nodes

Explicit construction of Lagrange polynomials from the nodes.
This function builds a `Poly` object.

 * `i` Index specifying which Lagrange polynomial should be returned.
 * `z` Vector containing the nodes of the Lagrange polynomials.
"""
function lagrange_poly{T<:Number}(i::Integer, z::AbstractVector{T})
    np = length(z)

    y = zeros(T, np-1)

    for j = 1:(i-1)
        y[j] = z[j]
    end
    for j = (i+1):np
        y[j-1] = z[j]
    end
    den = one(T)
    for j = 1:(np-1)
        den = den * (z[i] - y[j])
    end
    p = poly(y)
    return p/den
end

"""
Build Lagrange basis from any nodes

Probably not the fastest algorithm. But this won't be seriously used.
"""
function QuadType{T<:Number}(x::AbstractVector{T})  # Generic numbers:
    np = length(x)
    z = zeros(T,np)
    for i = 1:np
        z[i] = x[i]
    end

    w = zeros(T, np)
    D = zeros(T, np, np)

    for i = 1:np
        p = lagrange_poly(i, z)
        pint = polyint(p)
        pder = polyder(p)
        w[i] = polyval(pint, one(T)) - polyval(pint, -one(T))
        # Compute the rows of the derivative matrix
        for k = 1:np
            D[k,i] = polyval(pder, z[k])
        end
    end
    QuadType(np, z, w, D)
end


let
    x, y = variables(["x", "y"])
    @test (x*y).alpha == [1 1]
    @test (x*y).c == [1]

    @test (1 + 2*x*y).alpha == [0 0; 1 1]
    @test (1 + 2*x*y).c == [1; 2]

    @test ( (x+y)*(x-y) ).alpha == [2 0; 0 2]
    @test ( (x+y)*(x-y) ).c == [1; -1]

    @test ( (1+x+y)^2 * y ).deg == 3
    @test ( (1+x+y)^2 * y ).alpha == [0 1; 1 1; 0 2; 2 1; 1 2; 0 3]
    @test ( (1+x+y)^2 * y ).c == [1; 2; 2; 1; 2; 1]
end

let
    x, y, z = variables(["x", "y", "z"])
    @test Polyopt.monomials(3,[x,y,z]) == [1,x,y,z,x^2,x*y,x*z,y^2,y*z,z^2,x^3,x^2*y,x^2*z,x*y^2,x*y*z,x*z^2,y^3,y^2*z,y*z^2,z^3]
end