"""
This module will be deleted shortly, when we beginning using a rasterized PSF.
In the future, PSF code will appear only in `SDSSIO.jl` (to load the eigen-PSF)
and `model/psf.jl` (functions for computing parts of the log probability
relating to the PSF).
"""
module PSF

using Celeste
using Celeste: Const, @aliasscope, @unroll_loop
using ..Model
using ..Model: GalaxySigmaDerivs, BivariateNormalDerivatives, get_bvn_derivs!,
               eval_bvn_pdf!, transform_bvn_derivs!, BvnComponent, get_bvn_cov
using ..Transform
using ..SensitiveFloats.SensitiveFloat
import ..SensitiveFloats

import Optim
import WCS

using ForwardDiff
using StaticArrays

const ID_MAT_2D = eye(Float64, 2)

"""
A data type to store functions related to optimizing a PSF fit. Initialize
using the transform and number of components, and then call fit_psf
to perform a fit to a specificed psf and initial parameters.
"""
mutable struct PsfOptimizer
    psf_transform::DataTransform
    ftol::Float64
    grtol::Float64
    num_iters::Int
    verbose::Bool
    K::Int

    # Variable that will be allocated in optimization:
    bvn_derivs::BivariateNormalDerivatives{Float64}

    log_pdf::SensitiveFloat{Float64}
    pdf::SensitiveFloat{Float64}
    pixel_value::SensitiveFloat{Float64}
    squared_error::SensitiveFloat{Float64}
    sf_free::SensitiveFloat{Float64}

    psf_params_free_vec_cache::Vector{Float64}

    function PsfOptimizer(psf_transform::DataTransform, K::Int;
                          verbose::Bool=false,
                          ftol::Float64=1e-9,
                          grtol::Float64=1e-9)
        num_iters = 50

        bvn_derivs = BivariateNormalDerivatives{Float64}()

        log_pdf = SensitiveFloat{Float64}(length(PsfParams), 1, true, true)
        pdf = SensitiveFloat{Float64}(length(PsfParams), 1, true, true)
        pixel_value = SensitiveFloat{Float64}(length(PsfParams), K, true, true)
        squared_error = SensitiveFloat{Float64}(length(PsfParams), K, true, true)
        sf_free = SensitiveFloat{Float64}(length(PsfParams), K, true, true)

        psf_params_free_vec_cache = fill(NaN, K * length(PsfParams))

        new(psf_transform, ftol, grtol, num_iters, verbose, K,
            bvn_derivs, log_pdf, pdf, pixel_value, squared_error,
            sf_free, psf_params_free_vec_cache)
    end
end

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

function psf_fit_for_optim(psf_optimizer::PsfOptimizer, raw_psf, x_mat, psf_params_free_vec::Vector)
    if psf_params_free_vec == psf_optimizer.psf_params_free_vec_cache
        return psf_optimizer.sf_free
    else
        copy!(psf_optimizer.psf_params_free_vec_cache, psf_params_free_vec)
    end

    psf_params_free = unwrap_psf_params(psf_params_free_vec)
    psf_params = constrain_psf_params(psf_params_free, psf_optimizer.psf_transform)

    # Update squared_error in place.
    evaluate_psf_fit!(psf_params, raw_psf, x_mat, psf_optimizer.bvn_derivs,
                      psf_optimizer.log_pdf, psf_optimizer.pdf,
                      psf_optimizer.pixel_value, psf_optimizer.squared_error, true)

    # Update sf_free in place.
    transform_psf_sensitive_float!(psf_params, psf_optimizer.psf_transform,
                                   psf_optimizer.squared_error,
                                   psf_optimizer.sf_free, true)

    psf_optimizer.sf_free
end

function fit_psf(psf_optimizer::PsfOptimizer, raw_psf::Matrix{Float64}, initial_params::Vector{Vector{Float64}})
    x_mat = get_x_matrix_from_psf(raw_psf)
    psf_params_free = unconstrain_psf_params(initial_params, psf_optimizer.psf_transform)
    psf_params_free_vec = vec(wrap_psf_params(psf_params_free))

    function psf_fit_value(psf_params_free_vec::Vector)
        psf_fit_for_optim(psf_optimizer, raw_psf, x_mat, psf_params_free_vec).v[]
    end

    function psf_fit_grad!(psf_params_free_vec::Vector{Float64}, grad::Vector{Float64})
        grad[:] = psf_fit_for_optim(psf_optimizer, raw_psf, x_mat, psf_params_free_vec).d[:]
    end

    function psf_fit_hess!(psf_params_free_vec::Vector{Float64}, hess::Matrix{Float64})
        hess[:] = psf_fit_for_optim(psf_optimizer, raw_psf, x_mat, psf_params_free_vec).h
        hess[:] = 0.5 * (hess + hess')
    end

    tr_method = Optim.NewtonTrustRegion(initial_delta=10.0, delta_hat=1e9, eta=0.1,
                                        rho_lower=0.2, rho_upper=0.75)

    options = Optim.Options(;
                            x_tol = 0.0, # Don't allow convergence in params
                            f_tol = psf_optimizer.ftol,
                            g_tol = psf_optimizer.grtol,
                            iterations = psf_optimizer.num_iters,
                            store_trace = psf_optimizer.verbose,
                            show_trace = false,
                            extended_trace = psf_optimizer.verbose)

    return Optim.optimize(psf_fit_value, psf_fit_grad!, psf_fit_hess!,
                          psf_params_free_vec, tr_method, options)
end

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

"""
Return an image of a Celeste GMM PSF evaluated at rows, cols.

Args:
- psf_array: The PSF to be evaluated as an array of PsfComponent
- rows: The rows in the image (usually in pixel coordinates)
- cols: The column in the image (usually in pixel coordinates)

Returns:
 - The PSF values at rows and cols.    The default size is the same as
     that returned by get_psf_at_point applied to FITS header values.

Note that the point in the image at which the PSF is evaluated --
that is, the center of the image returned by this function -- is
already implicit in the value of psf_array.
"""
function get_psf_at_point(psf_array::Vector{PsfComponent};
                          rows = -25:25, cols = -25:25)
    function get_psf_value(psf::PsfComponent, row::Float64, col::Float64)
         x = @SVector(Float64[row, col]) - psf.xiBar
         exp_term = exp(-0.5 * dot(x, psf.tauBarInv * x) - 0.5 * psf.tauBarLd)
         psf.alphaBar * exp_term / (2 * pi)
    end

    return Float64[
     sum(get_psf_value(psf, float(row), float(col)) for psf in psf_array)
         for row in rows, col in cols ]
end

"""
Get the PSF located at a particular world location in an image.

Args:
 - world_loc: A location in world coordinates.
 - img: An Image

Returns:
 - An array of PsfComponent objects that represents the PSF as a mixture
     of Gaussians.
"""
function get_source_psf(world_loc::Vector{Float64}, img::Image, psf_K::Int)
    # Some stamps or simulated data have no raw psf information.    In that case,
    # just use the psf from the image.
    if size(img.raw_psf_comp.rrows) == (0, 0)
        # Also return a vector of empty psf params
        return img.psf, fill(fill(NaN, length(PsfParams)), psf_K)
    else
        pixel_loc = WCS.world_to_pix(img.wcs, world_loc)
        psfstamp = Model.eval_psf(img.raw_psf_comp, pixel_loc[1], pixel_loc[2])
        return PSF.fit_raw_psf_for_celeste(psfstamp, psf_K)
    end
end

"""
Get the PSF located at a particular world location in an image.

Args:
 - world_loc: A location in world coordinates.
 - img: An Image

Returns:
 - An array of PsfComponent objects that represents the PSF as a mixture
     of Gaussians.
"""
function get_source_psf(world_loc::Vector{Float64},
                        img::Image,
                        psf_optimizer::PsfOptimizer,
                        initial_psf_params::Vector{Vector{Float64}})
    # Some stamps or simulated data have no raw psf information. In that case,
    # just use the psf from the image.
    if size(img.raw_psf_comp.rrows) == (0, 0)
        return img.psf
    else
        pixel_loc = WCS.world_to_pix(img.wcs, world_loc)
        psfstamp = Model.eval_psf(img.raw_psf_comp, pixel_loc[1], pixel_loc[2])
        return PSF.fit_raw_psf_for_celeste(psfstamp, psf_optimizer, initial_psf_params)
    end
end

"""
Get initial values for a set of PSF parameters.

Args:
    - K: The number of components.
    - for_test: If true, return a set of slightly asymmetric parameters so that
                            derivatives are interesting for testing.

Returns:
    - A vector of PSF parameters, one for each component.
"""
function initialize_psf_params(K::Int; for_test::Bool=false)
    psf_params = Vector{Vector{Float64}}(K)
    for k=1:K
        if for_test
            # Choose asymmetric values for testing.
            psf_params[k] = zeros(length(PsfParams))
            psf_params[k][psf_ids.mu] = [0.1, 0.2]
            psf_params[k][psf_ids.gal_ab] = 0.8
            psf_params[k][psf_ids.gal_angle] = pi / 4
            psf_params[k][psf_ids.gal_scale] = sqrt(2 * k)
            psf_params[k][psf_ids.weight] = 1 / K + k / 10
        else
            psf_params[k] = zeros(length(PsfParams))
            psf_params[k][psf_ids.mu] = [0.0, 0.0]
            psf_params[k][psf_ids.gal_ab] = 0.95
            psf_params[k][psf_ids.gal_angle] = 0.0
            psf_params[k][psf_ids.gal_scale] = sqrt(2 * k)
            psf_params[k][psf_ids.weight] = 1 / K
        end
    end

    psf_params
end

"""
Given PSF parameters, return a DataTransform to transfrom them to and from
an unconstrained parameterization.

Args:
    - psf_params: The parameters to be transformed
    - scale: A vector as long as PsfParams to define a linear rescaling.

Returns:
    - A DataTransform object.
"""
function get_psf_transform(
        psf_params::Vector{Vector{Float64}};
        scale::Vector{Float64}=ones(length(PsfParams)))
    K = length(psf_params)
    bounds = Vector{ParamBounds}(length(psf_params))
    # Note that, for numerical reasons, the bounds must be on the scale
    # of reasonably meaningful changes.
    for k in 1:K
        bounds[k] = ParamBounds()
        bounds[k][:mu] = ParamBox[ ParamBox(-5.0, 5.0, scale[psf_ids.mu[1]]),
                                   ParamBox(-5.0, 5.0, scale[psf_ids.mu[2]]) ]
        bounds[k][:gal_ab] = ParamBox[ ParamBox(0.1, 1.0, scale[psf_ids.gal_ab] ) ]
        bounds[k][:gal_angle] =
            ParamBox[ ParamBox(-4 * pi, 4 * pi, scale[psf_ids.gal_angle] ) ]
        bounds[k][:gal_scale] =
            ParamBox[ ParamBox(0.05, 10.0, scale[psf_ids.gal_scale] ) ]

        # Note that the weights do not need to sum to one.
        bounds[k][:weight] = ParamBox[ ParamBox(0.05, 2.0, scale[psf_ids.weight] ) ]
    end
    DataTransform(bounds, collect(1:K), K)
end

"""
Transform psf_params to and from psf_params_free in place.

Args:
    - psf_params, psf_params_free: The constrained and unconstrained parameters,
            respectively.
    - psf_transform: The DataTransform to be applied.
    - to_unconstrained: If true, then update psf_params_free with the unconstrained
            representation of psf_params.    If false, update psf_params with the
            constrained version of psf_params_free.

Returns:
    - Updates psf_params or psf_params_free in place.
"""
function transform_psf_params!{NumType<:Number}(
        psf_params::Vector{Vector{NumType}}, psf_params_free::Vector{Vector{NumType}},
        psf_transform::DataTransform, to_unconstrained::Bool)
    for k=1:length(psf_params)
        for (param, constraint_vec) in psf_transform.bounds[k]
            for ind in 1:length(getfield(psf_ids, param))
                param_ind = getfield(psf_ids, param)[ind]
                constraint = constraint_vec[ind]
                to_unconstrained ?
                    psf_params_free[k][param_ind] =
                        Transform.unbox_parameter(psf_params[k][param_ind], constraint):
                    psf_params[k][param_ind] =
                        Transform.box_parameter(psf_params_free[k][param_ind], constraint)
            end
        end
    end

    true # return type
end

"""
Allocate memory for and return a constrained parameter set.
"""
function constrain_psf_params{NumType<:Number}(
        psf_params_free::Vector{Vector{NumType}}, psf_transform::DataTransform)
    K = length(psf_params_free)
    psf_params = Vector{Vector{NumType}}(K)
    for k=1:K
        psf_params[k] = zeros(NumType, length(PsfParams))
    end

    transform_psf_params!(psf_params, psf_params_free, psf_transform, false)

    psf_params
end

"""
Allocate memory for and return an unconstrained parameter set.
"""
function unconstrain_psf_params{NumType<:Number}(
        psf_params::Vector{Vector{NumType}}, psf_transform::DataTransform)
    K = length(psf_params)
    psf_params_free = Vector{Vector{NumType}}(K)
    for k=1:K
        psf_params_free[k] = zeros(NumType, length(PsfParams))
    end

    transform_psf_params!(psf_params, psf_params_free, psf_transform, true)

    psf_params_free
end

"""
Convert a single vector of psf parameters to a vector of vectors.
"""
function unwrap_psf_params{NumType <: Number}(psf_param_vec::Vector{NumType})
    @assert length(psf_param_vec) % length(PsfParams) == 0
    K = round(Int, length(psf_param_vec) / length(PsfParams))
    psf_param_mat = reshape(psf_param_vec, length(PsfParams), K)
    psf_params = Vector{Vector{NumType}}(K)
    for k = 1:K
        psf_params[k] = psf_param_mat[:, k]
    end
    psf_params
end

"""
Convert a vector of vectors of psf parameters to a single vector.
"""
function wrap_psf_params{NumType <: Number}(psf_params::Vector{Vector{NumType}})
    psf_params_mat = zeros(NumType, length(PsfParams), length(psf_params))
    for k=1:length(psf_params)
        psf_params_mat[:, k] = psf_params[k]
    end
    psf_params_mat[:]
end

"""
Return a sensitive float representing the value of the psf at pixel x
with all its associated derivatives (with respect to the constrained
parameterization).    Note that "fit" is a bad name -- it is just the pixel value.
TODO: fix the name.

Args:
    - x: The 2d location at which to evaluate the pdf
    - calculate_gradient: If false, only update the value of pixel_value
    - other values: Pre-allocated memory for intermediate calculations

Returns:
    - Updates pixel_value in place (and all the other placeholder values as well)
"""
function evaluate_psf_pixel_fit!{NumType <: Number}(
        x::SVector{2,Float64}, psf_params::Vector{Vector{NumType}},
        sig_sf_vec::Vector{GalaxySigmaDerivs{NumType}},
        bvn_vec::Vector{BvnComponent{NumType}},
        bvn_derivs::BivariateNormalDerivatives{NumType},
        log_pdf::SensitiveFloat{NumType},
        pdf::SensitiveFloat{NumType},
        pixel_value::SensitiveFloat{NumType},
        calculate_gradient::Bool)
    SensitiveFloats.zero!(pixel_value)

    K = length(psf_params)
    sigma_ids = (psf_ids.gal_ab, psf_ids.gal_angle, psf_ids.gal_scale)
    @inbounds for k = 1:K
        # I will put in the weights later so that the log pdf sensitive float
        # is accurate.
        bvn = bvn_vec[k]
        eval_bvn_pdf!(bvn_derivs, bvn, x)
        get_bvn_derivs!(bvn_derivs, bvn, true, true)
        transform_bvn_derivs!(bvn_derivs, sig_sf_vec[k], I, true)

        SensitiveFloats.zero!(log_pdf)
        SensitiveFloats.zero!(pdf)

        # This is redundant, but it's what eval_bvn_pdf returns.
        log_pdf.v[] = log(bvn_derivs.f_pre[1])

        if calculate_gradient
            for ind=1:2
                log_pdf.d[psf_ids.mu[ind]] = bvn_derivs.bvn_u_d[ind]
            end
            for ind=1:3
                log_pdf.d[sigma_ids[ind]] = bvn_derivs.bvn_s_d[ind]
            end
            log_pdf.d[psf_ids.weight] = 0

            for ind1 = 1:2, ind2 = 1:2
                log_pdf.h[psf_ids.mu[ind1], psf_ids.mu[ind2]] =
                    bvn_derivs.bvn_uu_h[ind1, ind2]
            end
            for mu_ind = 1:2, sig_ind = 1:3
                log_pdf.h[psf_ids.mu[mu_ind], sigma_ids[sig_ind]] =
                log_pdf.h[sigma_ids[sig_ind], psf_ids.mu[mu_ind]] =
                    bvn_derivs.bvn_us_h[mu_ind, sig_ind]
            end
            for ind1 = 1:3, ind2 = 1:3
                log_pdf.h[sigma_ids[ind1], sigma_ids[ind2]] =
                    bvn_derivs.bvn_ss_h[ind1, ind2]
            end
        end

        pdf_val = exp(log_pdf.v[])
        pdf.v[] = pdf_val

        if calculate_gradient
            for ind1 = 1:length(PsfParams)
                if ind1 == psf_ids.weight
                    pdf.d[ind1] = pdf_val
                else
                    pdf.d[ind1] = psf_params[k][psf_ids.weight] * pdf_val * log_pdf.d[ind1]
                end

                for ind2 = 1:ind1
                    pdf.h[ind1, ind2] = pdf.h[ind2, ind1] =
                        psf_params[k][psf_ids.weight] * pdf_val *
                        (log_pdf.h[ind1, ind2] + log_pdf.d[ind1] * log_pdf.d[ind2])
                end
            end

            # Weight hessian terms.
            for ind1 = 1:length(PsfParams)
                pdf.h[psf_ids.weight, ind1] = pdf.h[ind1, psf_ids.weight] =
                    pdf_val * log_pdf.d[ind1]
            end

        end

        pdf.v[] *= psf_params[k][psf_ids.weight]

        SensitiveFloats.add_sources_sf!(pixel_value, pdf, k)
    end

    true # Set return type
end

"""
Convert PSF parameters to covariance matrices and derivatives and BvnComponents.
"""
function get_sigma_from_params{NumType <: Number}(psf_params::Vector{Vector{NumType}})
    K = length(psf_params)
    sigma_vec = Vector{SMatrix{2,2,NumType,4}}(K)
    sig_sf_vec = Vector{GalaxySigmaDerivs{NumType}}(K)
    bvn_vec = Vector{BvnComponent{NumType}}(K)
    for k = 1:K
        sigma_vec[k] = get_bvn_cov(psf_params[k][psf_ids.gal_ab],
            psf_params[k][psf_ids.gal_angle],
            psf_params[k][psf_ids.gal_scale])
        sig_sf_vec[k] = GalaxySigmaDerivs(
            psf_params[k][psf_ids.gal_angle],
            psf_params[k][psf_ids.gal_ab],
            psf_params[k][psf_ids.gal_scale], sigma_vec[k], 1.0, true)

        bvn_vec[k] =
            BvnComponent(SVector{2,NumType}(psf_params[k][psf_ids.mu]), sigma_vec[k], 1.0)
    end
    sigma_vec, sig_sf_vec, bvn_vec
end

"""
evaluate_psf_fit but with pre-allocated memory for intermediate calculations.
"""
function evaluate_psf_fit!{NumType <: Number}(
        psf_params::Vector{Vector{NumType}}, raw_psf::Matrix{Float64},
        x_mat::Matrix{SVector{2,Float64}},
        bvn_derivs::BivariateNormalDerivatives{NumType},
        log_pdf::SensitiveFloat{NumType},
        pdf::SensitiveFloat{NumType},
        pixel_value::SensitiveFloat{NumType},
        squared_error::SensitiveFloat{NumType},
        calculate_gradient::Bool)
    K = length(psf_params)
    sigma_vec, sig_sf_vec, bvn_vec = get_sigma_from_params(psf_params)
    SensitiveFloats.zero!(squared_error)

    @inbounds for x_ind in 1:length(x_mat)
        SensitiveFloats.zero!(pixel_value)
        evaluate_psf_pixel_fit!(
                x_mat[x_ind], psf_params, sig_sf_vec, bvn_vec,
                bvn_derivs, log_pdf, pdf, pixel_value, calculate_gradient)

        diff = (pixel_value.v[] - raw_psf[x_ind])
        squared_error.v[] +=    diff ^ 2
        if calculate_gradient
            for ind1 = 1:length(squared_error.d)
                squared_error.d[ind1] += 2 * diff * pixel_value.d[ind1]
                for ind2 = 1:ind1
                    squared_error.h[ind1, ind2] +=
                        2 * (diff * pixel_value.h[ind1, ind2] +
                                 pixel_value.d[ind1] * pixel_value.d[ind2]')
                    squared_error.h[ind2, ind1] = squared_error.h[ind1, ind2]
                end
            end
        end # if calculate_gradient
    end

    squared_error
end

"""
Transform a sf, a SensitiveFloat that has derivatives with respect to
constrianed parameters, into sf_free, a SensitiveFloat that has derivatives
with respect to unconstrained parameteres.

Args:
    - psf_params: The constrianed parameters
    - psf_transform: The data transform to be applied
    - sf: The SensitiveFloat with derivatives with respect to the constrained params
    - sf_free: Updated in place.    The SensitiveFloat with derivatives with respect
                         to the unconstrained parameters.
    - calculate_gradient: If false, only calculate the value.

Returns:
    - Updates sf_free in place.
"""
function transform_psf_sensitive_float!{NumType<:Number}(
        psf_params::Vector{Vector{NumType}}, psf_transform::DataTransform,
        sf::SensitiveFloat{NumType}, sf_free::SensitiveFloat{NumType},
        calculate_gradient::Bool)
    sf_free.v[] = sf.v[]
    if calculate_gradient
        K = length(psf_params)

        # These are the hessians of each individual parameter's transform.    We
        # can represent it this way since each parameter's transform only depends on
        # its own value and not on others.

        # This is the diagonal of the Jacobian transform.
        jacobian_diag = zeros(length(PsfParams) * K)

        # These are hte Hessians of the each parameter's transform.
        hessian_values = zeros(length(PsfParams) * K)

        for k = 1:K
            offset = length(PsfParams) * (k - 1)
            for ind = 1:2
                mu_ind = psf_ids.mu[ind]
                jac, hess = Transform.box_derivatives(
                    psf_params[k][mu_ind], psf_transform.bounds[k][:mu][ind])
                jacobian_diag[offset + mu_ind] = jac
                hessian_values[offset + mu_ind] = hess
            end

            # The rest are one-dimensional.
            for field in setdiff(fieldnames(PsfParams), [ :mu ])
                ind = getfield(psf_ids, field)
                jac, hess = Transform.box_derivatives(
                    psf_params[k][ind], psf_transform.bounds[k][field][1])
                jacobian_diag[offset + ind] = jac
                hessian_values[offset + ind] = hess
            end
        end

        # Apply the transformations.
        for ind1 = 1:length(sf_free.d)
            sf_free.d[ind1] = jacobian_diag[ind1] * sf.d[ind1]

            for ind2 = 1:ind1
                # Calculate the Hessian
                sf_free.h[ind1, ind2] = sf_free.h[ind2, ind1] =
                    (jacobian_diag[ind1] * jacobian_diag[ind2]) * sf.h[ind1, ind2]
                if ind1 == ind2
                    sf_free.h[ind1, ind2] +=    hessian_values[ind1] * sf.d[ind1]
                end
            end
        end
    end

    true # return type
end

"""
Given a psf matrix, return a matrix of the same size, where each element of
the matrix is a 2-length vector of the [x1, x2] location of that matrix cell.
The locations are chosen so that the scale is pixels, but the center of the
psf is [0, 0].
"""
function get_x_matrix_from_psf(psf::Matrix{Float64})
    psf_center1, psf_center2 = (size(psf, 1) - 1) / 2 + 1, (size(psf, 2) - 1) / 2 + 1
    return [ @SVector Float64[i - psf_center1, j - psf_center2] for i = 1:size(psf, 1), j = 1:size(psf, 2) ]
end

"""
Fit a mixture of 2d Gaussians to a PSF image (evaluated at a single point).

Args:
    - raw_psf: An matrix image of the point spread function.
    - psf_optimizer: A PsfOptimizer object for the fit.

Returns:
    - A vector of PsfComponents fit to the raw_psf.
"""
function fit_raw_psf_for_celeste{P<:PsfOptimizer}(
        raw_psf::Array{Float64, 2}, psf_optimizer::P,
        initial_psf_params::Vector{Vector{Float64}})
    K = length(initial_psf_params)
    @assert K == psf_optimizer.K
    optim_result = fit_psf(psf_optimizer, raw_psf, initial_psf_params)
    psf_params_fit =
        constrain_psf_params(
            unwrap_psf_params(Optim.minimizer(optim_result)), psf_optimizer.psf_transform)

    sigma_vec = get_sigma_from_params(psf_params_fit)[1]
    celeste_psf = Vector{PsfComponent}(K)
    for k=1:K
        mu = psf_params_fit[k][psf_ids.mu]
        weight = psf_params_fit[k][psf_ids.weight]
        celeste_psf[k] = PsfComponent(weight, SVector{2,eltype(mu)}(mu), SMatrix{2,2,eltype(mu),4}(sigma_vec[k]))
    end

    celeste_psf, psf_params_fit
end

"""
Fit a mixture of 2d Gaussians to a PSF image (evaluated at a single point).

Args:
 - raw_psf: An matrix image of the point spread function.
 - K: The number of components to fit.

 Returns:
     - A vector of PsfComponents fit to the raw_psf.
"""
function fit_raw_psf_for_celeste(raw_psf::Array{Float64, 2}, K::Integer; ftol=1e-9)
    psf_params = PSF.initialize_psf_params(K, for_test=false)
    psf_transform = PSF.get_psf_transform(psf_params)
    psf_optimizer = PsfOptimizer(psf_transform, K)
    psf_optimizer.ftol = ftol
    fit_raw_psf_for_celeste(raw_psf, psf_optimizer, psf_params)
end

function trim_psf(raw_psf::Array{Float64, 2}; trim_percent=0.999)
    h_mid = cld(size(raw_psf, 1), 2)
    w_mid = cld(size(raw_psf, 2), 2)

    width = 1
    function get_trimmed_psf_image()
        raw_psf[(h_mid - width):(h_mid + width),
                (w_mid - width):(w_mid + width)]
    end

    psf_tot = sum(abs, raw_psf)
    while sum(abs, get_trimmed_psf_image()) < trim_percent * psf_tot
        width += 1
    end

    deepcopy(get_trimmed_psf_image())
end

end # module


import LinearAlgebra.LAPACK: syev!, @blasfunc, BlasInt, chkstride1, checksquare, chklapackerror, liblapack
import LinearAlgebra.BLAS: libblas
import LinearAlgebra: eigen, eigen!
import Base: @propagate_inbounds

abstract type Spin end
struct Up <: Spin end
struct Down <: Spin end

"Represents a magnetic Hamiltonian matrix with the block structure [up updown;downup down]"
abstract type AbstractMagneticMatrix{T} <: AbstractMatrix{T} end

data(m::AbstractMatrix) = m
data(m::AbstractMagneticMatrix) = m.data

Base.similar(::Type{M}, i::NTuple{2,Int}) where {M <: AbstractMagneticMatrix} =
    M(Matrix{M.parameters[1]}(undef, i))

for f in (:length, :size, :setindex!, :elsize)
    @eval @inline @propagate_inbounds Base.$f(c::AbstractMagneticMatrix, args...) =
    Base.$f(c.data, args...)
end

Base.pointer(c::AbstractMagneticMatrix, i::Integer) = pointer(c.data, i)

"Magnetic block dimensions"
blockdim(c::AbstractMatrix) = div(size(data(c), 2), 2)

up(c::AbstractMatrix) =   (d = blockdim(c); view(data(c), 1:d, 1:d))
down(c::AbstractMatrix) = (d = blockdim(c); r = d + 1:2 * d; view(data(c), r, r))

# Standard getindex behavior

for f in (:view, :getindex)
    @eval @inline @propagate_inbounds Base.$f(c::AbstractMagneticMatrix, args...) =
        getindex(c.data, args...)
    @eval @inline @propagate_inbounds Base.$f(c::AbstractMagneticMatrix, r::Union{Colon, AbstractUnitRange}, i::Int) =
        MagneticVector(Base.$f(c.data, r, i))
    @eval @inline @propagate_inbounds Base.$f(c::AbstractMagneticMatrix, i::Int, r::Union{Colon, AbstractUnitRange}) =
        MagneticVector(Base.$f(c.data, i, r))
end

Base.similar(c::M, args::AbstractUnitRange...) where {M <: AbstractMagneticMatrix} =
    M(similar(c.data), args...)

"""
    ColinMatrix{T, M <: AbstractMatrix{T}} <: AbstractMagneticMatrix{T}

Defines a Hamiltonian Matrix with [up zeros; zeros down] structure.
It is internally only storing the up and down block.
"""
struct ColinMatrix{T,M <: AbstractMatrix{T}} <: AbstractMagneticMatrix{T}
    data::M
end

function ColinMatrix(up::AbstractMatrix, down::AbstractMatrix)
    @assert size(up) == size(down)
    return ColinMatrix([up down])
end

Base.Array(c::ColinMatrix{T}) where T =
    (d = blockdim(c); [c[Up()] zeros(T, d, d); zeros(T, d, d) c[Down()]])

down(c::ColinMatrix) = (d = blockdim(c); view(c.data, 1:d, d + 1:2 * d))
blockdim(c::ColinMatrix) = size(c.data, 1)

function LinearAlgebra.diag(c::ColinMatrix)
    d = blockdim(c)
    r = LinearAlgebra.diagind(d, d)
    [c[r];c[r.+last(r)]]
end

"""
    NonColinMatrix{T, M <: AbstractMatrix{T}} <: AbstractMagneticMatrix{T}

Defines a Hamiltonian Matrix with [up updown;downup down] structure.
Since atomic projections w.r.t spins are defined rather awkwardly in Wannier90 for exchange calculations,
i.e. storing the up-down parts of an atom sequentially,
a NonColinMatrix reshuffles the entries of a matrix such that it follows the above structure.
"""
struct NonColinMatrix{T,M <: AbstractMatrix{T}} <: AbstractMagneticMatrix{T}
    data::M
end

"Reshuffles standard Wannier90 up-down indices to the ones for the structure of a NonColinMatrix."
function Base.convert(::Type{NonColinMatrix}, m::M) where {M <: AbstractMatrix}
    @assert iseven(size(m, 1)) "Error, dimension of the supplied matrix is odd, i.e. it does not contain both spin components."
    data = similar(m)
    d    = blockdim(m)
    for i in 1:2:size(m, 1), j in 1:2:size(m, 2)
        up_id1 = div1(i, 2)
        up_id2 = div1(j, 2)
        data[up_id1, up_id2] = m[i, j]
        data[up_id1 + d, up_id2] = m[i + 1, j]
        data[up_id1, up_id2 + d] = m[i, j + 1]
        data[up_id1 + d, up_id2 + d] = m[i + 1, j + 1]
    end
    return NonColinMatrix(data)
end

function NonColinMatrix(up::AbstractMatrix{T}, down::AbstractMatrix{T}) where {T}
    @assert size(up) == size(down)
    return NonColinMatrix([up zeros(T, size(up));zeros(T, size(up)) down])
end

Base.Array(c::NonColinMatrix) = copy(c.data)

function uprange(a::DFC.Structures.Projection)
    projrange = range(a)
    if length(projrange) > a.orbital.size
        return range(div1(first(projrange), 2), length = div(length(projrange), 2))
    else
        return projrange
    end
end
uprange(a::DFC.Structures.Atom) = vcat(uprange.(a.projections)...)

## Indexing ##
for f in (:view, :getindex)
    @eval function Base.$f(c::ColinMatrix, a1::T, a2::T) where {T <: Union{DFC.Structures.Projection,DFC.Structures.Atom}}
        projrange1 = range(a1)
        projrange2 = range(a2)

        return ColinMatrix($f(c, projrange1, projrange2), $f(c, projrange1, projrange2 .+ blockdim(c)))
    end
    @eval function Base.$f(c::NonColinMatrix, a1::T, a2::T) where {T <: Union{DFC.Structures.Projection,DFC.Structures.Atom}}
        up_range1 = uprange(a1)
        up_range2 = uprange(a2)
        d = blockdim(c)
        dn_range1 = up_range1 .+ d
        dn_range2 = up_range2 .+ d
        return NonColinMatrix([$f(c, up_range1, up_range2) $f(c, up_range1, dn_range2)
                               $f(c, dn_range1, up_range2) $f(c, dn_range1, dn_range2)])
    end

    @eval Base.$f(c::AbstractMagneticMatrix, a1::T) where {T <: Union{DFC.Structures.Projection,DFC.Structures.Atom}} =
        $f(c, a1, a1)

    @eval Base.$f(c::ColinMatrix, a1::T, a2::T, ::Up) where {T <: Union{DFC.Structures.Projection,DFC.Structures.Atom}} =
        $f(c, range(a1), range(a2))

    @eval Base.$f(c::NonColinMatrix, a1::T, a2::T, ::Up) where {T <: Union{DFC.Structures.Projection,DFC.Structures.Atom}} =
        $f(c, uprange(a1), uprange(a2))

    @eval Base.$f(c::ColinMatrix, a1::T, a2::T, ::Down) where {T <: Union{DFC.Structures.Projection,DFC.Structures.Atom}} =
        $f(c, range(a1), range(a2) .+ blockdim(c))

    @eval Base.$f(c::NonColinMatrix, a1::T, a2::T, ::Down) where {T <: Union{DFC.Structures.Projection,DFC.Structures.Atom}} =
        $f(c, uprange(a1) .+ blockdim(c), uprange(a2) .+ blockdim(c))

    @eval Base.$f(c::NonColinMatrix, a1::T, a2::T, ::Up, ::Down) where {T <: Union{DFC.Structures.Projection,DFC.Structures.Atom}} =
        $f(c, uprange(a1), uprange(a2) .+ blockdim(c))

    @eval Base.$f(c::NonColinMatrix, a1::T, a2::T, ::Down, ::Up) where {T <: Union{DFC.Structures.Projection,DFC.Structures.Atom}} =
        $f(c, uprange(a1) .+ blockdim(c), uprange(a2))

    @eval Base.$f(c::NonColinMatrix, ::Up, ::Down) where {T <: Union{DFC.Structures.Projection,DFC.Structures.Atom}} =
        (s = size(c,1); $f(c, 1:div(s, 2), div(s, 2)+1:s))

    @eval Base.$f(c::NonColinMatrix, ::Down, ::Up) where {T <: Union{DFC.Structures.Projection,DFC.Structures.Atom}} =
        (s = size(c,1); $f(c, div(s, 2)+1:s, 1:div(s, 2)))

    @eval Base.$f(c::AbstractMatrix, a1::T, a2::T, ::Up) where {T<:Union{DFC.Structures.Projection, DFC.Structures.Atom}} =
        $f(c, range(a1), range(a2))

    @eval Base.$f(c::AbstractMatrix, a1::T, a2::T, ::Down) where {T<:Union{DFC.Structures.Projection, DFC.Structures.Atom}} =
        $f(c, range(a1), range(a2) .+ blockdim(c))

    @eval Base.$f(c::AbstractMagneticMatrix, ::Up) =
        (r = 1:blockdim(c); $f(c, r, r))

    @eval Base.$f(c::ColinMatrix, ::Down) =
        (d = blockdim(c); r = 1:d; $f(c, r, r .+ d))

    @eval Base.$f(c::NonColinMatrix, ::Down) =
        (d = blockdim(c); r = d+1 : 2*d; $f(c, r, r))

    @eval Base.$f(c::AbstractMatrix, ::Up) =
        (r = 1:blockdim(c); $f(c, r, r))

    @eval Base.$f(c::AbstractMatrix, ::Down) =
        (d = blockdim(c); r = d + 1:2 * d; $f(c, r, r))

end

for op in (:*, :-, :+, :/)
    @eval @inline Base.$op(c1::ColinMatrix, c2::ColinMatrix) =
        ColinMatrix($op(c1[Up()], c2[Up()]), $op(c1[Down()], c2[Down()]))
    @eval @inline Base.$op(c1::NonColinMatrix, c2::NonColinMatrix) =
        NonColinMatrix($op(c1.data, c2.data))
end

    # BROADCASTING
Base.BroadcastStyle(::Type{<:AbstractMagneticMatrix}) =
    Broadcast.ArrayStyle{AbstractMagneticMatrix}()

Base.ndims(::Type{<:AbstractMagneticMatrix}) =
    2

Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{AbstractMagneticMatrix}}, ::Type{ElType}) where {ElType} =
    Base.similar(AbstractMagneticMatrix{ElType}, axes(bc))

Base.axes(c::AbstractMagneticMatrix) =
    Base.axes(c.data)

@inline @propagate_inbounds Base.broadcastable(c::AbstractMagneticMatrix) =
    c

@inline @propagate_inbounds Base.unsafe_convert(::Type{Ptr{T}}, c::AbstractMagneticMatrix{T}) where {T} =
    Base.unsafe_convert(Ptr{T}, c.data)

for (elty, cfunc) in zip((:ComplexF32, :ComplexF64), (:cgemm_, :zgemm_))
    @eval @inline function LinearAlgebra.mul!(C::ColinMatrix{$elty}, A::ColinMatrix{$elty}, B::ColinMatrix{$elty})
        dim = blockdim(C)
        dim2 = dim * dim
        ccall((@blasfunc($(cfunc)), libblas), Cvoid,
                        (Ref{UInt8}, Ref{UInt8}, Ref{BlasInt}, Ref{BlasInt},
                         Ref{BlasInt}, Ref{$elty}, Ptr{$elty}, Ref{BlasInt},
                         Ptr{$elty}, Ref{BlasInt}, Ref{$elty}, Ptr{$elty},
                         Ref{BlasInt}),
                         'N', 'N', dim, dim,
                         dim, one($elty), A, dim,
                         B, dim, zero($elty), C, dim)
        ccall((@blasfunc($(cfunc)), libblas), Cvoid,
                        (Ref{UInt8}, Ref{UInt8}, Ref{BlasInt}, Ref{BlasInt},
                         Ref{BlasInt}, Ref{$elty}, Ptr{$elty}, Ref{BlasInt},
                         Ptr{$elty}, Ref{BlasInt}, Ref{$elty}, Ptr{$elty},
                         Ref{BlasInt}),
                         'N', 'N', dim, dim,
                         dim, one($elty), pointer(A, dim2 + 1), dim,
                         pointer(B, dim2 + 1), dim, zero($elty), pointer(C, dim2 + 1), dim)

        return C
    end
end

@inline function LinearAlgebra.adjoint(c::AbstractMagneticMatrix)
    out = similar(c)
    adjoint!(out, c)
end

@inline @inbounds function LinearAlgebra.adjoint!(out::ColinMatrix, in1::ColinMatrix)
    dim = blockdim(out)
    for i in 1:dim, j in 1:dim
        out[j, i] = in1[i, j]'
        out[j, i + dim] = in1[i, j + dim]'
    end
    return out
end

@inline LinearAlgebra.adjoint!(out::NonColinMatrix, in1::NonColinMatrix) =
    adjoint!(out.data, in1.data)

for (elty, relty, cfunc) in zip((:ComplexF32, :ComplexF64), (:Float32, :Float64), (:cheev_, :zheev_))
# We use Upper Triangular blas for everything! And Eigvals are always all calculated
    @eval @inline function blas_eig_ccall(A::AbstractMatrix{$elty},
                                          W::AbstractVector{$relty},
                                          work::Vector{$elty},
                                          lwork::BlasInt,
                                          rwork::Vector{$relty},
                                          n::Int,
                                          info::Ref{BlasInt})

    ccall((@blasfunc($cfunc), liblapack), Cvoid,
           (Ref{UInt8}, Ref{UInt8}, Ref{BlasInt}, Ptr{$elty}, Ref{BlasInt},
           Ptr{$relty}, Ptr{$elty}, Ref{BlasInt}, Ptr{$relty}, Ptr{BlasInt}),
           'V', 'U', n, A, n, W, work, lwork, rwork, info)
    chklapackerror(info[])
    end
end

"Vector following the same convention as the in AbstractMagneticMatrix, i.e. first half of the indices contain the up part, second the down part"
struct MagneticVector{T} <: AbstractVector{T}
    data::Vector{T}
end

for f in (:length, :size, :setindex!, :elsize)
    @eval @inline @propagate_inbounds Base.$f(c::MagneticVector, args...) =
    Base.$f(c.data, args...)
end

up(c::MagneticVector) = view(c.data, 1:div(length(c), 2))
down(c::MagneticVector) = (lc = length(c); view(c.data, div(lc, 2):lc))

# Standard getindex behavior
@inline @propagate_inbounds Base.getindex(c::MagneticVector, args...) =
    getindex(c.data, args...)

for f in (:view, :getindex)
    @eval @inline @propagate_inbounds Base.$f(c::MagneticVector, args::AbstractUnitRange...) =
    Base.$f(c.data, args...)
end

Base.similar(v::MagneticVector) = MagneticVector(similar(v.data))

"Reshuffles standard Wannier90 up-down indices to the ones for the structure of a MagneticVector."
function Base.convert(::Type{MagneticVector}, v::V) where {V <: AbstractVector}
    @assert iseven(length(v)) "Error, dimension of the supplied matrix is odd, i.e. it does not contain both spin components."
    data = similar(v)
    vl = length(v)
    d    =div(vl, 2)

    for i in 1:2:vl
        up_id1 = div1(i, 2)
        data[up_id1] = v[i]
        data[up_id1 + d] = v[i + 1]
    end
    return NonColinMatrix(data)
end

## Indexing with atoms and spins ##
for f in (:view, :getindex)
    @eval function Base.$f(c::MagneticVector, a1::T) where {T <: Union{DFC.Structures.Projection,DFC.Structures.Atom}}
        projrange1 = uprange(a1)
        return MagneticVector([$f(c, projrange1); $f(c, projrange1 .+ div(length(c), 2))])
    end
    @eval Base.$f(c::MagneticVector, a1::T, ::Up) where {T <: Union{DFC.Structures.Projection,DFC.Structures.Atom}} =
        $f(c, uprange(a1))

    @eval Base.$f(c::MagneticVector, a1::T, ::Down) where {T <: Union{DFC.Structures.Projection,DFC.Structures.Atom}} =
        $f(c, range(a1), uprange(a2) + div(length(c), 2))

    @eval Base.$f(c::MagneticVector, ::Up) =
        $f(c, 1:div(length(c), 2))

    @eval Base.$f(c::MagneticVector, ::Down) =
        (lc = length(c); $f(c, div(lc, 2) + 1:lc))
end

for op in (:*, :-, :+, :/)
    @eval @inline Base.$op(c1::MagneticVector, c2::MagneticVector) =
        MagneticVector($op(c1.data, c2.data))
end

struct EigCache{T <: AbstractFloat}
    work::Vector{Complex{T}}
    lwork::BlasInt
    rwork::Vector{T}
    n::Int
    info::Ref{BlasInt}
    function EigCache(A::AbstractMatrix{Complex{relty}}) where {relty}
        chkstride1(A)
        elty = Complex{relty}
        n = checksquare(A)
        W     = similar(A, relty, n)
        work  = Vector{elty}(undef, 1)
        lwork = BlasInt(-1)
        rwork = Vector{relty}(undef, max(1, 3n - 2))
        info  = Ref{BlasInt}()

        blas_eig_ccall(A, W, work, lwork, rwork, n, info)
        lwork = BlasInt(real(work[1]))
        resize!(work, lwork)
        return new{relty}(work, lwork, rwork, n, info)
    end
end
EigCache(A::ColinMatrix) = EigCache(up(A))
EigCache(A::NonColinMatrix) = EigCache(A.data)

@inline function eigen!(vals::AbstractVector{T}, vecs::AbstractMatrix{Complex{T}}, c::EigCache{T}) where {T}
    blas_eig_ccall(vecs, vals, c.work, c.lwork, c.rwork, c.n, c.info)
    return Eigen(vals, vecs)
end

@inline function eigen!(vals::AbstractVector{T}, vecs::ColinMatrix{Complex{T}}, c::EigCache{T}) where {T}
    eigen!(view(vals.data, 1:c.n), up(vecs), c)
    eigen!(view(vals.data, c.n + 1:2 * c.n), down(vecs), c)
    return Eigen(vals, vecs)
end

@inline function eigen!(vals::AbstractVector{T}, vecs::NonColinMatrix{Complex{T}}, c::EigCache{T}) where {T}
    eigen!(vals.data, vecs.data, c)
    return Eigen(vals, vecs)
end

@inline function eigen(vecs::AbstractMatrix{Complex{T}}, c::EigCache{T}) where {T}
    out  = copy(vecs)
    vals = similar(out, T, size(out, 2))
    return eigen!(vals, out, c)
end

@inline function eigen(vecs::AbstractMagneticMatrix{Complex{T}}, c::EigCache{T}) where {T}
    out  = copy(vecs)
    vals = MagneticVector(similar(out, T, size(out, 2)))
    return eigen!(vals, out, c)
end

# @inline function eigen!(orig::M) where {M<:AbstractMagneticMatrix}
#     vals, vecs = eigen!(orig.data)
#     return Eigen(vals, M(vecs))
# end

@inline function eigen(h::AbstractMagneticMatrix)
    return eigen(h, EigCache(h))
end
    # Very unoptimized
function Base.Array(e::Eigen{Complex{T},T,<:ColinMatrix{Complex{T}}}) where {T}
    d = size(e.vectors, 1)
    return [e.vectors[1:d, 1:d] * diagm(0 => e.values[1:d]) * e.vectors[1:d, 1:d]' e.vectors[1:d, d + 1:2d] * diagm(0 => e.values[d + 1:2d]) * e.vectors[1:d, d + 1:2d]']
end


@inline function save_noise!(W::NoiseProcess)
  if W.t != W.curt
    push!(W.W,copy(W.curW))
    push!(W.t,copy(W.curt))
    if W.Z != nothing
      push!(W.Z,copy(W.curZ))
    end
  end
end

@inline function accept_step!(W::NoiseProcess,dt,setup_next=true)

  W.curt += W.dt
  W.iter += 1

  if isinplace(W)
    @.. W.curW += W.dW
    if W.Z != nothing
      @.. W.curZ += W.dZ
    end
  else
    W.curW += W.dW
    if W.Z != nothing
      W.curZ += W.dZ
    end
  end

  if W.save_everystep
    push!(W.W,copy(W.curW))
    push!(W.t,copy(W.curt))
    if W.Z != nothing
      push!(W.Z,copy(W.curZ))
    end
  end

  W.dt = dt #dtpropose
  # Setup next step
  if setup_next
    setup_next_step!(W::NoiseProcess)
  end
end

@inline function setup_next_step!(W::NoiseProcess)
  if adaptive_alg(W)==:RSwM3
    ResettableStacks.reset!(W.S₂) #Empty W.S₂
  end
  if adaptive_alg(W)==:RSwM1
    if !isempty(W.S₁)
      W.dt,W.dW,W.dZ = pop!(W.S₁)
    else # Stack is empty
      calculate_step!(W,W.dt)
    end
  elseif adaptive_alg(W)==:RSwM2 || adaptive_alg(W)==:RSwM3
    if !isinplace(W)
      dttmp = zero(W.dt); W.dW = zero(W.dW)
      if W.Z != nothing
        W.dZ = zero(W.dZ)
      end
    else
      dttmp = zero(W.dt); fill!(W.dW,zero(eltype(W.dW)))
      if W.Z != nothing
        fill!(W.dZ,zero(eltype(W.dZ)))
      end
    end
    while !isempty(W.S₁)
      L₁,L₂,L₃ = pop!(W.S₁)
      qtmp = (W.dt-dttmp)/L₁
      if qtmp>1
        dttmp+=L₁
        if isinplace(W)
          @.. W.dW+=L₂
          if W.Z != nothing
            @.. W.dZ+=L₃
          end
        else
          W.dW+=L₂
          if W.Z != nothing
            W.dZ+=L₃
          end
        end
        if adaptive_alg(W)==:RSwM3
          if L₁ > W.rswm.discard_length
            push!(W.S₂,(L₁,L₂,L₃))
          end
        end
      else #Popped too far
        # Generate numbers to bridge and step perfectly
        dttmp += qtmp*L₁
        if isinplace(W)
          W.bridge(W.dWtilde,W,W.curW,L₂,qtmp,L₁,W.rng)
          if W.Z != nothing
            W.bridge(W.dZtilde,W,W.curZ,L₃,qtmp,L₁,W.rng)
          end
        else
          W.dWtilde = W.bridge(W,W.curW,L₂,qtmp,L₁,W.rng)
          if W.Z != nothing
            W.dZtilde = W.bridge(W,W.curZ,L₃,qtmp,L₁,W.rng)
          end
        end
        if isinplace(W)
          @.. W.dW += W.dWtilde
          if W.Z != nothing
            @.. W.dZ += W.dZtilde
          end
        else
          W.dW += W.dWtilde
          if W.Z != nothing
            W.dZ += W.dZtilde
          end
        end
        if (1-qtmp)*L₁ > W.rswm.discard_length
          if isinplace(W)
            @.. L₂ -= W.dWtilde
            if W.Z != nothing
              @.. L₃ -= W.dZtilde
            end
          else
            L₂ -= W.dWtilde
            if W.Z != nothing
              L₃ -= W.dZtilde
            end
          end
          if W.Z == nothing
            push!(W.S₁,((1-qtmp)*L₁,L₂,nothing))
          else
            push!(W.S₁,((1-qtmp)*L₁,L₂,L₃))
          end
          if adaptive_alg(W)==:RSwM3 && qtmp*L₁ > W.rswm.discard_length
            if W.Z == nothing
              ResettableStacks.copyat_or_push!(W.S₂,(qtmp*L₁,W.dWtilde,nothing))
            else
              ResettableStacks.copyat_or_push!(W.S₂,(qtmp*L₁,W.dWtilde,W.dZtilde))
            end
          end
        end
        break
      end
    end #end while empty
    # This is a control variable so do not diff through it
    dtleft = DiffEqBase.ODE_DEFAULT_NORM(W.dt - dttmp,W.curt)
    if dtleft > W.rswm.discard_length #Stack emptied
      if isinplace(W)
        W.dist(W.dWtilde,W,dtleft,W.rng)
        if W.Z != nothing
          W.dist(W.dZtilde,W,dtleft,W.rng)
        end
      else
        W.dWtilde = W.dist(W,dtleft,W.rng)
        if W.Z != nothing
          W.dZtilde = W.dist(W,dtleft,W.rng)
        end
      end
      if isinplace(W)
        @.. W.dW += W.dWtilde
        if W.Z != nothing
          @.. W.dZ += W.dZtilde
        end
      else
        W.dW += W.dWtilde
        if W.Z != nothing
          W.dZ += W.dZtilde
        end
      end
      if adaptive_alg(W)==:RSwM3
        if W.Z == nothing
          ResettableStacks.copyat_or_push!(W.S₂,(dtleft,W.dWtilde,nothing))
        else
          ResettableStacks.copyat_or_push!(W.S₂,(dtleft,W.dWtilde,W.dZtilde))
        end
      end
    end
  end # End RSwM2 and RSwM3
end

@inline function calculate_step!(W::NoiseProcess,dt)
  if isinplace(W)
    W.dist(W.dW,W,dt,W.rng)
    if W.Z != nothing
      W.dist(W.dZ,W,dt,W.rng)
    end
  else
    W.dW = W.dist(W,dt,W.rng)
    if W.Z != nothing
      W.dZ = W.dist(W,dt,W.rng)
    end
  end
  W.dt = dt
end

@inline function reject_step!(W::NoiseProcess,dtnew)
  q = dtnew/W.dt
  if adaptive_alg(W)==:RSwM1 || adaptive_alg(W)==:RSwM2
    if isinplace(W)
      W.bridge(W.dWtilde,W,0,W.dW,q,dtnew,W.rng)
      if W.Z != nothing
        W.bridge(W.dZtilde,W,0,W.dZ,q,dtnew,W.rng)
      end
    else
      W.dWtilde = W.bridge(W,0,W.dW,q,dtnew,W.rng)
      if W.Z != nothing
        W.dZtilde=  W.bridge(W,0,W.dZ,q,dtnew,W.rng)
      end
    end
    cutLength = W.dt-dtnew
    if cutLength > W.rswm.discard_length
      if W.Z == nothing
        push!(W.S₁,(cutLength,W.dW-W.dWtilde,nothing))
      else
        push!(W.S₁,(cutLength,W.dW-W.dWtilde,W.dZ-W.dZtilde))
      end
    end
    if length(W.S₁) > W.maxstacksize
        W.maxstacksize = length(W.S₁)
    end
    if isinplace(W)
      copyto!(W.dW,W.dWtilde)
      if W.Z!=nothing
        copyto!(W.dZ,W.dZtilde)
      end
    else
      W.dW = W.dWtilde
      if W.Z != nothing
        W.dZ = W.dZtilde
      end
    end
    W.dt = dtnew
  else # RSwM3
    if !isinplace(W)
      dttmp = zero(W.dt); W.dWtmp = zero(W.dW)
      if W.Z != nothing
        W.dZtmp = zero(W.dZtmp)
      end
    else
      dttmp = zero(W.dt); fill!(W.dWtmp,zero(eltype(W.dWtmp)))
      if W.Z!= nothing
        fill!(W.dZtmp,zero(eltype(W.dZtmp)))
      end
    end
    if length(W.S₂) > W.maxstacksize2
      W.maxstacksize2= length(W.S₂)
    end
    while !isempty(W.S₂)
      L₁,L₂,L₃ = pop!(W.S₂)
      if dttmp + L₁ < (1-q)*W.dt #while the backwards movement is less than chop off
        dttmp += L₁
        if isinplace(W)
          @.. W.dWtmp += L₂
          if W.Z != nothing
            @.. W.dZtmp += L₃
          end
        else
          W.dWtmp += L₂
          if W.Z != nothing
            W.dZtmp += L₃
          end
        end
        push!(W.S₁,(L₁,L₂,L₃))
      else
        push!(W.S₂,(L₁,L₂,L₃))
        break
      end
    end # end while
    dtK = W.dt - dttmp
    qK = q*W.dt/dtK
    if isinplace(W)
      @.. W.dWtmp = W.dW - W.dWtmp
      if W.Z != nothing
        @.. W.dZtmp = W.dZ - W.dZtmp
      end
    else
      W.dWtmp = W.dW - W.dWtmp
      if W.Z != nothing
        W.dZtmp = W.dZ - W.dZtmp
      end
    end
    if isinplace(W)
      W.bridge(W.dWtilde,W,0,W.dWtmp,qK,dtK,W.rng)
      #W.dWtilde .-= W.curW
      if W.Z != nothing
        W.bridge(W.dZtilde,W,0,W.dZtmp,qK,dtK,W.rng)
        #W.dZtilde .-= W.curZ
      end
    else
      W.dWtilde = W.bridge(W,0,W.dWtmp,qK,dtK,W.rng)# - W.curW
      if W.Z != nothing
        W.dZtilde = W.bridge(W,0,W.dZtmp,qK,dtK,W.rng)# - W.curZ
      end
    end
    # This is a control variable so do not diff through it
    cutLength = DiffEqBase.ODE_DEFAULT_NORM((1-qK)*dtK,W.curt)
    if cutLength > W.rswm.discard_length
      if W.Z == nothing
        push!(W.S₁,(cutLength,W.dWtmp-W.dWtilde,nothing))
      else
        push!(W.S₁,(cutLength,W.dWtmp-W.dWtilde,W.dZtmp-W.dZtilde))
      end
    end
    if length(W.S₁) > W.maxstacksize
        W.maxstacksize = length(W.S₁)
    end
    W.dt = dtnew
    if isinplace(W)
      copyto!(W.dW,W.dWtilde)
      if W.Z != nothing
        copyto!(W.dZ,W.dZtilde)
      end
    else
      W.dW = W.dWtilde
      if W.Z != nothing
        W.dZ = W.dZtilde
      end
    end
  end
end

@inline function interpolate!(W::NoiseProcess,t)
  if t > W.t[end] # Steps past W
    dt = t - W.t[end]
    if isinplace(W)
      W.dist(W.dW,W,dt,W.rng)
      W.curW .+= W.dW
      if W.Z != nothing
        W.dist(W.dZ,W,dt,W.rng)
        W.curZ .+= W.dZ
      end
    else
      W.dW = W.dist(W,dt,W.rng)
      W.curW += W.dW
      if W.Z != nothing
        W.dZ = W.dist(W,dt,W.rng)
        W.curZ += W.dZ
      end
    end
    out1 = copy(W.curW)
    if W.save_everystep
      push!(W.t,t)
      push!(W.W,out1)
    end
    if W.Z != nothing
      out2= copy(W.curZ)
      if W.save_everystep
        push!(W.Z,out2)
      end
    else
      out2 = nothing
    end
    return out1,out2
  else # Bridge
    i = searchsortedfirst(W.t,t)
    if t == W.t[i]
      if isinplace(W)
        W.curW .= W.W[i]
      else
        W.curW = W.W[i]
      end
      if W.Z != nothing
        if isinplace(W)
          W.curZ .= W.Z[i]
        else
          W.curZ = W.Z[i]
        end
        return copy(W.curW),copy(W.curZ)
      else
        return copy(W.curW),nothing
      end
    else
      W0,Wh = W.W[i-1],W.W[i]
      if W.Z != nothing
        Z0,Zh = W.Z[i-1],W.Z[i]
      end
      h = W.t[i]-W.t[i-1]
      q = (t-W.t[i-1])/h
      if isinplace(W)
        new_curW = similar(W.dW)
        W.bridge(new_curW,W,W0,Wh,q,h,W.rng)
        new_curW .+= (1-q)*W0
        if W.Z != nothing
          new_curZ = similar(W.dZ)
          W.bridge(new_curZ,W,Z0,Zh,q,h,W.rng)
          new_curZ .+= (1-q)*Z0
        else
          new_curZ = nothing
        end
      else
        new_curW = W.bridge(W,W0,Wh,q,h,W.rng)
        new_curW += (1-q)*W0
        if W.Z != nothing
          new_curZ = W.bridge(W,Z0,Zh,q,h,W.rng)
          new_curZ += (1-q)*Z0
        else
          new_curZ = nothing
        end
      end
      W.curW = new_curW
      if W.save_everystep
        insert!(W.W,i,new_curW)
        insert!(W.t,i,t)
      end
      if W.Z != nothing
        W.curZ = new_curZ
        if W.save_everystep
          insert!(W.Z,i,new_curZ)
        end
      end
      return new_curW,new_curZ
    end
  end
end

@inline function interpolate!(out1,out2,W::NoiseProcess,t)
  if t > W.t[end] # Steps past W
    dt = t - W.t[end]
    W.dist(W.dW,W,dt,W.rng)
    out1 .+= W.dW
    if W.Z != nothing
      W.dist(W.dZ,W,dt,W.rng)
      out2 .+= W.dZ
    end
    if W.save_everystep
      push!(W.t,t)
      push!(W.W,copy(out1))
      if W.Z != nothing
        push!(W.Z,copy(out2))
      end
    end
  else # Bridge
    i = searchsortedfirst(W.t,t)
    if t == W.t[i]
      out1 .= W.W[i]
      if W.Z != nothing
        out2 .= W.Z[i]
      end
    else
      W0,Wh = W.W[i-1],W.W[i]
      if W.Z != nothing
        Z0,Zh = W.Z[i-1],W.Z[i]
      end
      h = W.t[i]-W.t[i-1]
      q = (t-W.t[i-1])/h
      W.bridge(out1,W,W0,Wh,q,h,W.rng)
      out1 .+= (1-q)*W0
      if W.Z != nothing
        W.bridge(out2,W,Z0,Zh,q,h,W.rng)
        out2 .+= (1-q)*Z0
      end
      W.curW .= out1
      if W.save_everystep
        insert!(W.W,i,copy(out1))
        insert!(W.t,i,t)
      end
      if W.Z != nothing
        W.curZ .= out2
        if W.save_everystep
          insert!(W.Z,i,copy(out2))
        end
      end
    end
  end
end

function resize_stack!(W::NoiseProcess,i)
  for j in eachindex(W.S₂.data)
    resize!(W.S₂.data[j][2],i)
    W.S₂.data[j][3] != nothing && resize!(W.S₂.data[j][3],i)
  end
end

function deleteat_stack!(W::NoiseProcess,i)
  for j in eachindex(W.S₂.data)
    deleteat!(W.S₂.data[j][2],i)
    W.S₂.data[j][3] != nothing && deleteat!(W.S₂.data[j][3],i)
  end
end

#=
function addat_stack!(W::NoiseProcess,i)
  for j in eachindex(W.S₂.data)
    resize!(W.S₂.data[j],i)
    W.S₂.data[j][3] != nothing && resize!(W.S₂.data[j][3],i)
  end
end
=#


module LinAlgHeaders

using Clang

include("wrap_utils.jl")
const topdir = splitdir(splitdir(JULIA_HOME)[1])[1]
const depsdir = joinpath(topdir, "deps")
const basedir = joinpath(topdir, "base")

### CHOLMOD
function cholmod()
    cmdir = finddir(depsdir, ("SuiteSparse", "CHOLMOD"))
    includedir = "Include"
    headers = ["cholmod_core.h"]
    headers = [joinpath(depsdir, cmdir, includedir, h) for h in headers]
    outdir = joinpath(basedir, "sparse", "gen")
    isdir(outdir) || mkdir(outdir)
    context=wrap_c.init(output_dir=outdir,
                        output_file="cholmod_decls.jl",
                        common_file="cholmod_defs.jl",
                        header_library=x->:libcholmod,
                        headers = headers,
                        clang_diagnostics=true,
                        header_wrapped=(x,y)->true)

    context.options.wrap_structs = true
#    context.options.immutable_structs = true
    run(context)
end

end # module


module SWIM

using DataStructures
using Distances
using Rectangle
using Distributions
using Dates
using Base

export Canvas
export initialize
export generate_nodes
export Cell
export Node
export move
export update_seen
export simulate
export close_writer

include("Canvas.jl")
include("EventWriter.jl")
include("Node.jl")
include("Simulation.jl")

end


# This is a copy of Knet/data/imdb.jl and Knet/tutorial/07.imdb.ipynb cleaned up

# Based on https://github.com/fchollet/keras/raw/master/keras/datasets/imdb.py
# Also see https://github.com/fchollet/keras/raw/master/examples/imdb_lstm.py
# Also see https://github.com/ilkarman/DeepLearningFrameworks/raw/master/common/utils.py

using Pkg; for p in ("PyCall","JSON","JLD2"); haskey(Pkg.installed(),p) || Pkg.add(p); end
using PyCall,JSON,JLD2,Random
np = pyimport("numpy")

"""

    imdb()

Load the IMDB Movie reviews sentiment classification dataset from
https://keras.io/datasets and return (xtrn,ytrn,xtst,ytst,dict) tuple.

# Keyword Arguments:
- url=https://s3.amazonaws.com/text-datasets: where to download the data (imdb.npz) from.
- dir=Pkg.dir("Knet/data"): where to cache the data.
- maxval=nothing: max number of token values to include. Words are ranked by how often they occur (in the training set) and only the most frequent words are kept. nothing means keep all, equivalent to maxval = vocabSize + pad + stoken.
- maxlen=nothing: truncate sequences after this length. nothing means do not truncate.
- seed=0: random seed for sample shuffling. Use system seed if 0.
- pad=true: whether to pad short sequences (padding is done at the beginning of sequences). pad_token = maxval.
- stoken=true: whether to add a start token to the beginning of each sequence. start_token = maxval - pad.
- oov=true: whether to replace words >= oov_token with oov_token (the alternative is to skip them). oov_token = maxval - pad - stoken.

"""
function imdb(;
              url = "https://s3.amazonaws.com/text-datasets",
              dir = "./", # joinpath(@__DIR__, "imdb"),
              data="imdb.npz",
              dict="imdb_word_index.json",
              jld2="imdb.jld2",
              maxval=nothing,
              maxlen=nothing,
              seed=0, oov=true, stoken=true, pad=true
              )
    global _imdb_xtrn,_imdb_ytrn,_imdb_xtst,_imdb_ytst,_imdb_dict
    if !(@isdefined _imdb_xtrn)
        isdir(dir) || mkpath(dir)
        jld2path = joinpath(dir,jld2)
        if !isfile(jld2path)
            @info("Downloading IMDB...")
            datapath = joinpath(dir,data)
            dictpath = joinpath(dir,dict)
            isfile(datapath) || download("$url/$data",datapath)
            isfile(dictpath) || download("$url/$dict",dictpath)
            d = np.load(datapath)
            _imdb_xtrn = map(a->np.asarray(a,dtype=np.int32), get(d, "x_train"))
            _imdb_ytrn = Array{Int8}(get(d, "y_train") .+ 1)
            _imdb_xtst = map(a->np.asarray(a,dtype=np.int32), get(d, "x_test"))
            _imdb_ytst = Array{Int8}(get(d, "y_test") .+ 1)
            _imdb_dict = Dict{String,Int32}(JSON.parsefile(dictpath))
            JLD2.@save jld2path _imdb_xtrn _imdb_ytrn _imdb_xtst _imdb_ytst _imdb_dict
            #rm(datapath)
            #rm(dictpath)
        end
        @info("Loading IMDB...")
        JLD2.@load jld2path _imdb_xtrn _imdb_ytrn _imdb_xtst _imdb_ytst _imdb_dict
    end
    if seed != 0; Random.seed!(seed); end
    xs = [_imdb_xtrn;_imdb_xtst]
    if maxlen == nothing; maxlen = maximum(map(length,xs)); end
    if maxval == nothing; maxval = maximum(map(maximum,xs)) + pad + stoken; end
    if pad; pad_token = maxval; maxval -= 1; end
    if stoken; start_token = maxval; maxval -= 1; end
    if oov; oov_token = maxval; end
    function _imdb_helper(x,y)
        rp = randperm(length(x))
        newy = y[rp]
        newx = similar(x)
        for i in 1:length(x)
            xi = x[rp[i]]
            if oov
                xi = map(w->(w<=oov_token ? w : oov_token), xi)
            else
                xi = filter(w->(w<=oov_token), xi)
            end
            if stoken
                xi = [ start_token; xi ]
            end
            if length(xi) > maxlen
                xi = xi[end-maxlen+1:end]
            end
            if pad && length(xi) < maxlen
                xi = append!(repeat([pad_token], maxlen-length(xi)), xi)
            end
            newx[i] = xi
        end
        newx,newy
    end
    xtrn,ytrn = _imdb_helper(_imdb_xtrn,_imdb_ytrn)
    xtst,ytst = _imdb_helper(_imdb_xtst,_imdb_ytst)
    return xtrn,ytrn,xtst,ytst,_imdb_dict
end

### From 07.imdb.ipynb

using Pkg
for p in ("Knet",) #"ProgressMeter")
    haskey(Pkg.installed(),p) || Pkg.add(p)
end

EPOCHS=3          # Number of training epochs
BATCHSIZE=64      # Number of instances in a minibatch
EMBEDSIZE=125     # Word embedding size
NUMHIDDEN=100     # Hidden layer size
MAXLEN=150        # maximum size of the word sequence, pad shorter sequences, truncate longer ones
VOCABSIZE=30000   # maximum vocabulary size, keep the most frequent 30K, map the rest to UNK token
NUMCLASS=2        # number of output classes
DROPOUT=0.0       # Dropout rate
LR=0.001          # Learning rate
BETA_1=0.9        # Adam optimization parameter
BETA_2=0.999      # Adam optimization parameter
EPS=1e-08         # Adam optimization parameter

using Knet # : Knet
#ENV["COLUMNS"]=92                     # column width for array printing
#include(Knet.dir("data","imdb.jl"))   # defines imdb loader

#@doc imdb

#@time
(xtrn,ytrn,xtst,ytst,imdbdict)=imdb(maxlen=MAXLEN,maxval=VOCABSIZE);

#summary.((xtrn,ytrn,xtst,ytst,imdbdict))

# Words are encoded with integers
#rand(xtrn)'

# Each word sequence is padded or truncated to length 150
#length.(xtrn)'

# Define a function that can print the actual words:
imdbvocab = Array{String}(undef,length(imdbdict))
for (k,v) in imdbdict; imdbvocab[v]=k; end
imdbvocab[VOCABSIZE-2:VOCABSIZE] = ["<unk>","<s>","<pad>"]
function reviewstring(x,y=0)
    x = x[x.!=VOCABSIZE] # remove pads
    """$(("Sample","Negative","Positive")[y+1]) review:\n$(join(imdbvocab[x]," "))"""
end

# Hit Ctrl-Enter to see random reviews:
#r = rand(1:length(xtrn))
#println(reviewstring(xtrn[r],ytrn[r]))

# Here are the labels: 1=negative, 2=positive
#ytrn'

using Knet: param, dropout, RNN

struct SequenceClassifier; input; rnn; output; pdrop; end

SequenceClassifier(input::Int, embed::Int, hidden::Int, output::Int; pdrop=0) =
    SequenceClassifier(param(embed,input), RNN(embed,hidden,rnnType=:gru), param(output,hidden), pdrop)

function (sc::SequenceClassifier)(input)
    embed = sc.input[:, permutedims(hcat(input...))]
    embed = dropout(embed,sc.pdrop)
    hidden = sc.rnn(embed)
    hidden = dropout(hidden,sc.pdrop)
    return sc.output * hidden[:,:,end]
end

using Knet: minibatch
dtrn = minibatch(xtrn,ytrn,BATCHSIZE;shuffle=true)
dtst = minibatch(xtst,ytst,BATCHSIZE)
#length.((dtrn,dtst))

# For running experiments
#using Knet: train!, Adam
#import ProgressMeter

function trainresults(file,model)
    # if (print("Train from scratch? ");readline()[1]=='y')
    #     updates = 0; prog = ProgressMeter.Progress(EPOCHS * length(dtrn))
    #     function callback(J)
    #         ProgressMeter.update!(prog, updates)
    #         return (updates += 1) <= prog.n
    #     end
    #     opt = Adam(lr=LR, beta1=BETA_1, beta2=BETA_2, eps=EPS)
    #     train!(model, dtrn; callback=callback, optimizer=opt, pdrop=DROPOUT)
    #     Knet.gc()
    #     Knet.save(file,"model",model)
    # else
        isfile(file) || download("http://people.csail.mit.edu/deniz/models/tutorial/$file",file)
        model = Knet.load(file,"model")
    # end
    return model
end

#using Knet: nll, accuracy
model = SequenceClassifier(VOCABSIZE,EMBEDSIZE,NUMHIDDEN,NUMCLASS)
#nll(model,dtrn), nll(model,dtst), accuracy(model,dtrn), accuracy(model,dtst)

model = trainresults("imdbmodel113.jld2",model);

# 33s (0.059155148f0, 0.3877507f0, 0.9846153846153847, 0.8583733974358975)
#nll(model,dtrn), nll(model,dtst), accuracy(model,dtrn), accuracy(model,dtst)

predictstring(x)="\nPrediction: " * ("Negative","Positive")[argmax(Array(vec(model([x]))))]
UNK = VOCABSIZE-2
str2ids(s::String)=[(i=get(imdbdict,w,UNK); i>=UNK ? UNK : i) for w in split(lowercase(s))]

# Here we can see predictions for random reviews from the test set; hit Ctrl-Enter to sample:
# r = rand(1:length(xtst))
# println(reviewstring(xtst[r],ytst[r]))
# println(predictstring(xtst[r]))

# Here the user can enter their own reviews and classify them:
# println(predictstring(str2ids(readline(stdin))))


function writemesh(mesh,filename,ifile)

nbe = size(mesh.eblks,2);
nbf = size(mesh.fblks,2);
neb = maximum(eblks[2,:] .- eblks[1,:])+1;
nfb = maximum(fblks[2,:] .- fblks[1,:])+1;
# nfe = size(mesh.t2f,1);
# nv = maximum(mesh.t[:]));

ndims = zeros(20,1);
ndims[1] = mesh.nd;
ndims[2] = mesh.ne;
ndims[3] = mesh.nf;
ndims[4] = mesh.nv;
ndims[5] = mesh.nfe;
ndims[6] = nbe;
ndims[7] = neb;
ndims[8] = nbf;
ndims[9] = nfb;

nsize = zeros(30,1);
nsize[1] = length(ndims);
nsize[2] = length(mesh.facecon);
nsize[3] = length(mesh.eblks);
nsize[4] = length(mesh.fblks);
nsize[5] = length(mesh.nbsd);
nsize[6] = length(mesh.elemsend);
nsize[7] = length(mesh.elemrecv);
nsize[8] = length(mesh.elemsendpts);
nsize[9] = length(mesh.elemrecvpts);
nsize[10] = length(mesh.elempart);
nsize[11] = length(mesh.elempartpts);
nsize[12] = length(mesh.cgelcon);
nsize[13] = length(mesh.rowent2elem);
nsize[14] = length(mesh.cgent2dgent);
nsize[15] = length(mesh.colent2elem);
nsize[16] = length(mesh.rowe2f1);
nsize[17] = length(mesh.cole2f1);
nsize[18] = length(mesh.ent2ind1);
nsize[19] = length(mesh.rowe2f2);
nsize[20] = length(mesh.cole2f2);
nsize[21] = length(mesh.ent2ind2);

if ifile>0
    display(["Writing mesh into file " * string(ifile)]);
    fileID = open([filename * string(ifile) * ".bin"],"w");
else
    display(["Writing mesh into file "]);
    fileID = open([filename * ".bin"],"w");
end

write(fileID,Float64(length(nsize[:])));
write(fileID,Float64.(nsize[:]));

if nsize[1]>0
    write(fileID,Float64.(ndims[:]));
end
#reshape(permute(mesh.facecon,[2 1 3]),[2*master.npf(1)*ndims[3) 1])
if nsize[2]>0
    write(fileID,Float64.(mesh.facecon[:]));
end
if nsize[3]>0
    write(fileID,Float64.(mesh.eblks[:]));
end
if nsize[4]>0
    write(fileID,Float64.(mesh.fblks[:]));
end
if nsize[5]>0
    write(fileID,Float64.(mesh.nbsd[:]));
end
if nsize[6]>0
    write(fileID,Float64.(mesh.elemsend[:]));
end
if nsize[7]>0
    write(fileID,Float64.(mesh.elemrecv[:]));
end
if nsize[8]>0
    write(fileID,Float64.(mesh.elemsendpts[:]));
end
if nsize[9]>0
    write(fileID,Float64.(mesh.elemrecvpts[:]));
end
if nsize[10]>0
    write(fileID,Float64.(mesh.elempart[:]));
end
if nsize[11]>0
    write(fileID,Float64.(mesh.elempartpts[:]));
end
if nsize[12]>0
    write(fileID,Float64.(mesh.cgelcon[:]-1));
end
if nsize[13]>0
    write(fileID,Float64.(mesh.rowent2elem[:]));
end
if nsize[14]>0
    write(fileID,Float64.(mesh.cgent2dgent[:]-1));
end
if nsize[15]>0
    write(fileID,Float64.(mesh.colent2elem[:]-1));
end
if nsize[16]>0
    write(fileID,Float64.(mesh.rowe2f1[:]));
end
if nsize[17]>0
    write(fileID,Float64.(mesh.cole2f1[:]-1));
end
if nsize[18]>0
    write(fileID,Float64.(mesh.ent2ind1[:]-1));
end
if nsize[19]>0
    write(fileID,Float64.(mesh.rowe2f2[:]));
end
if nsize[20]>0
    write(fileID,Float64.(mesh.cole2f2[:]-1));
end
if nsize[21]>0
    write(fileID,Float64.(mesh.ent2ind2[:]-1));
end

close(fileID);

end


@testset "fsm_active_close.jl" begin
    base_seq = WrappingInt32(1 << 31)
    DEFAULT_CAPACITY = 64000
    TIMEOUT_DFLT = 1000
    @testset "START -> SYN_SENT -> SYN/ACK -> ACK" begin
        conn = TCPConnection()
        connect!(conn)
        tick!(conn, 1)
        seg1 = expect_one_seg(conn; ack=false, syn=true)
        expect_state(conn, JLSponge.SYN_SENT)

        isn = WrappingInt32(rand(UInt32))
        send_syn!(conn, isn, seg1.header.seqno + 1)
        tick!(conn, 1)
        expect_state(conn, JLSponge.ESTABLISHED)
        expect_one_seg(conn; ack=true, syn=false, ackno=isn+1)
        @test bytes_in_flight(conn) == 0
    end

    @testset "START -> SYN_SENT -> SYN/ACK -> ACK" begin
        conn = TCPConnection()
        connect!(conn)
        tick!(conn, 1)
        seg1 = expect_one_seg(conn; ack=false, syn=true)
        seg_hdr = seg1.header
        expect_state(conn, JLSponge.SYN_SENT)

        isn = WrappingInt32(rand(UInt32))
        send_syn!(conn, isn)
        tick!(conn, 1)
        expect_one_seg(conn, syn=false, ack=true, ackno=isn + 1)

        expect_state(conn, JLSponge.SYN_RCVD)
        send_ack!(conn, isn, seg_hdr.seqno + 1)
        tick!(conn, 1)
        expect_no_seg(conn)
        expect_state(conn, JLSponge.ESTABLISHED)
    end

    @testset "START -> SYN_SENT -> SYN/ACK -> ESTABLISHED" begin
        conn = TCPConnection()
        connect!(conn)
        tick!(conn, 1)
        seg1 = expect_one_seg(conn; ack=false, syn=true)
        seg_hdr = seg1.header
        expect_state(conn, JLSponge.SYN_SENT)

        isn = WrappingInt32(rand(UInt32))
        send_syn!(conn, isn, seg_hdr.seqno + 1)
        tick!(conn, 1)
        expect_one_seg(conn; ack=true, ackno=isn + 1, syn=false)
        expect_state(conn, JLSponge.ESTABLISHED)
    end
end


module DoubleStrain

using Random
using ProgressMeter
using CovidPOMDPs
using ..CovidPOMDPs:
    INFECTION_HORIZON,
    INF_DIST,
    VIRAL_LOADS,
    POS_TEST_PROBS,
    SYMPTOMATIC_ISOLATION_PROB,
    ASYMPTOMATIC_PROB,
    SYMPTOM_PROBS,
    DEFAULT_PARAMS,
    InfParams,
    CovidState,
    CovidAction,
    simplex_sample,
    RVsum

include("typedef.jl")
export DoubleCovidPOMDP, DoubleCovidState

include("simulate.jl")

include("POMDPsInterface.jl")

include("updater.jl")

end


"""
    FENodeSetModule

Module for the finite element node set.
"""
module FENodeSetModule

using FinEtools.FTypesModule: FInt, FFlt, FCplxFlt, FFltVec, FIntVec, FFltMat, FIntMat, FMat, FVec, FDataDict
import Base.count

"""
    Finite element node set type.

xyz =
    Array of node locations.
    Array of coordinates, the number of rows corresponds to the number of
    nodes in the set and the columns corresponds to the space dimensions.
    The location of node j is given by x[j,:].

"""
mutable struct FENodeSet
  xyz::Array{FFlt, 2}

  function FENodeSet(xyz::FFltMat)
    self = new(deepcopy(xyz)); # Need to make a COPY of the input array!
    return self
  end
end

"""
    spacedim(self::FENodeSet)

Number of dimensions of the space in which the node lives, 1, 2, or 3.
"""
spacedim(self::FENodeSet) =size(self.xyz,2)

"""
    xyz3(self::FENodeSet)

Get the  3-D coordinate that define the location  of the node.
Even if the nodes  were specified in  lower dimension (1-D, 2-D)
this function returns  a 3-D coordinate  by padding with zeros.
"""
function xyz3(self::FENodeSet)
    if (size(self.xyz,2)==1)
        val = [self.xyz zeros(size(self.xyz,1),2)];
    elseif (size(self.xyz,2)==2)
        val = [self.xyz zeros(size(self.xyz,1),1)];
    else
        val = deepcopy(self.xyz);
    end
end

"""
    count(self::FENodeSet)

Get the number of finite element nodes in the node set.
"""
function count(self::FENodeSet)
    return size(self.xyz,1)::FInt
end

end


"""
    `creat_cp_rs_struct(RS_REF)`

    Function that returns a RegStruct object RS with fields taken from RS_REF.

    # Examples
    ```julia-repl
    julia> CpelNano.creat_cp_rs_struct(rs_ref)
    ```
"""
function cp_rs_flds(rs_ref::RegStruct)::RegStruct

    # Create structure
    rs = deepcopy(rs_ref)

    # Reset fields
    rs.m = 0
    rs.nme = []
    rs.mml = []
    rs.ϕhat = []
    rs.calls = []
    rs.logZ = NaN
    rs.tm = TransferMat()
    rs.exps = Expectations()
    rs.nls_rgs = AnalysisRegions()

    # Return
    return rs

end
"""
    `pmap_two_samp_null_stats(MOD_REF,CALLS_S1,CALLS_S2,PERM,CONFIG)`

    Function that computes set of null statistics for a given data permutation.

    # Examples
    ```julia-repl
    julia> CpelNano.pmap_two_samp_null_stats(mod_ref,calls_s1,calls_s2,perm,config)
    ```
"""
function pmap_two_samp_null_stats(rs_ref::RegStruct, calls_s1::Vector{Vector{MethCallCpgGrp}}, calls_s2::Vector{Vector{MethCallCpgGrp}}, perm_ids::Vector{Int64}, config::CpelNanoConfig)::NTuple{3,Vector{Float64}}

    # Create two aux structs
    aux_rs_s1 = cp_rs_flds(rs_ref)
    aux_rs_s2 = cp_rs_flds(rs_ref)
    aux_rs_s1.m = length(calls_s1)
    aux_rs_s2.m = length(calls_s2)

    # Assign calls
    aux_calls = vcat(calls_s1, calls_s2)
    aux_rs_s1.calls = aux_calls[perm_ids]
    deleteat!(aux_calls, perm_ids)
    aux_rs_s2.calls = aux_calls

    # Init output
    tmml = fill(NaN, rs_ref.nls_rgs.num)
    tnme = fill(NaN, rs_ref.nls_rgs.num)
    tcmd = fill(NaN, rs_ref.nls_rgs.num)

    # Do estimation
    get_ϕhat!(aux_rs_s1, config)
    get_ϕhat!(aux_rs_s2, config)

    # Leave if we couldn't estimate ϕ
    (length(aux_rs_s1.ϕhat) > 0 && length(aux_rs_s2.ϕhat) > 0) || return tmml, tnme, tcmd
    config.verbose && print_log("ϕhat1=$(aux_rs_s1.ϕhat); ϕhat2=$(aux_rs_s2.ϕhat)")

    # Re-scale to per cpg site resolution
    rscle_grp_mod!(aux_rs_s1)
    rscle_grp_mod!(aux_rs_s2)

    # Statistical summaries
    get_stat_sums!(aux_rs_s1, config)
    get_stat_sums!(aux_rs_s2, config)

    # Compute stats
    tmml = aux_rs_s1.mml - aux_rs_s2.mml
    tnme = aux_rs_s1.nme - aux_rs_s2.nme
    tcmd = comp_cmd(aux_rs_s1, aux_rs_s2)
    config.verbose && print_log("tmml=$(tmml)")
    config.verbose && print_log("tnme=$(tnme)")
    config.verbose && print_log("tcmd=$(tcmd)")

    # Return tuple
    return tmml, tnme, tcmd

end
"""
    `pmap_two_samp_null_stats_test_fun(MOD_REF,CALLS_S1,CALLS_S2,PERM,CONFIG)`

    Function that computes set of null statistics for a given data permutation.

    # Examples
    ```julia-repl
    julia> CpelNano.pmap_two_samp_null_stats_test_fun(mod_ref,calls_s1,calls_s2,perm,config)
    ```
"""
function pmap_two_samp_null_stats_test_fun(rs_ref::RegStruct, calls_s1::Vector{Vector{MethCallCpgGrp}}, calls_s2::Vector{Vector{MethCallCpgGrp}}, perm_ids::Vector{Int64}, config::CpelNanoConfig)::NTuple{3,Vector{Float64}}

    # Create two aux structs
    aux_rs_s1 = cp_rs_flds(rs_ref)
    aux_rs_s2 = cp_rs_flds(rs_ref)
    aux_rs_s1.m = length(calls_s1)
    aux_rs_s2.m = length(calls_s2)

    # Assign calls
    aux_calls = vcat(calls_s1, calls_s2)
    aux_rs_s1.calls = aux_calls[perm_ids]
    deleteat!(aux_calls, perm_ids)
    aux_rs_s2.calls = aux_calls

    # Init output
    tmml = fill(NaN, rs_ref.nls_rgs.num)
    tnme = fill(NaN, rs_ref.nls_rgs.num)
    tcmd = fill(NaN, rs_ref.nls_rgs.num)

    # Do estimation
    get_ϕhat!(aux_rs_s1, config)
    get_ϕhat!(aux_rs_s2, config)

    # Leave if we couldn't estimate ϕ
    (length(aux_rs_s1.ϕhat) > 0 && length(aux_rs_s2.ϕhat) > 0) || return tmml, tnme, tcmd
    print_log("ϕhat1=$(aux_rs_s1.ϕhat); ϕhat2=$(aux_rs_s2.ϕhat)")

    # Statistical summaries
    get_stat_sums!(aux_rs_s1, config)
    get_stat_sums!(aux_rs_s2, config)

    # Compute stats
    tmml = aux_rs_s1.mml - aux_rs_s2.mml
    tnme = aux_rs_s1.nme - aux_rs_s2.nme
    tcmd = comp_cmd(aux_rs_s1, aux_rs_s2)
    print_log("μ1=$(aux_rs_s1.mml)")
    print_log("μ2=$(aux_rs_s2.mml)")
    print_log("tmml=$(tmml)")
    print_log("h1=$(aux_rs_s1.nme)")
    print_log("h2=$(aux_rs_s2.nme)")
    print_log("tnme=$(tnme)")
    print_log("tcmd=$(tcmd)")
    print_log("********************")

    # Return tuple
    return tmml, tnme, tcmd

end
"""
    `pmap_diff_two_samp_comp(MODS_CHR_S1,MODS_CHR_S2,CONFIG)`

    Function that performs differential methylation analysis between two samples given two
    CpelNano models.

    # Examples
    ```julia-repl
    julia> CpelNano.pmap_diff_two_samp_comp(mods_chr_s1,mods_chr_s2,config)
    ```
"""
function pmap_diff_two_samp_comp(mod_s1::RegStruct, mod_s2::RegStruct, nano_s1::String, nano_s2::String, fa_rec::FASTA.Record, config::CpelNanoConfig)::RegStatTestStruct

    print_log("Working on: $(mod_s1.chr)-$(mod_s1.chrst)-$(mod_s1.chrend)")

    # Init output struct
    test_struct = RegStatTestStruct(mod_s1)

    # Observed statistics
    tmml_obs = mod_s1.mml - mod_s2.mml
    tnme_obs = mod_s1.nme - mod_s2.nme
    tcmd_obs = comp_cmd(mod_s1, mod_s2)

    # Init p-values
    tmml_pvals = fill(NaN, length(tmml_obs))
    tnme_pvals = fill(NaN, length(tnme_obs))
    tcmd_pvals = fill(NaN, length(tcmd_obs))

    # Compute p-values if need be
    if config.pval_comp

        # Get group calls
        get_calls!(mod_s1, nano_s1, config)
        get_calls!(mod_s2, nano_s2, config)

        # Number of reads
        n_rds_s1 = length(mod_s1.calls)
        n_rds_s2 = length(mod_s2.calls)

        # Compute number of possible randomizations
        L = binomial(BigInt(n_rds_s1 + n_rds_s2), BigInt(n_rds_s1))

        # Create iteratable object with all combinations
        comb_iter = combinations(1:(n_rds_s1 + n_rds_s2), n_rds_s1)

        # If enough permutations
        if L > 20

            #  Check if exact test
            exact = L < config.LMAX_TWO_SAMP

            # Go back to group scale
            get_grp_info!(mod_s1, fa_rec, config.min_grp_dist)
            get_grp_info!(mod_s2, fa_rec, config.min_grp_dist)

            # Store calls in aux vars
            calls_s1 = deepcopy(mod_s1.calls)
            calls_s2 = deepcopy(mod_s2.calls)

            # Get read combinations to use
            if exact
                # Use all sample assignments
                comb_iter_used = comb_iter
            else
                # Sample assignments to group 1
                comb_iter_used = fill(Vector{Int64}(), config.LMAX_TWO_SAMP)
                @inbounds for i = 1:config.LMAX_TWO_SAMP
                    comb_iter_used[i] = sort!(StatsBase.sample(1:(n_rds_s1 + n_rds_s2), n_rds_s1;replace=false))
                end
            end

            ## Null statistics

            # Compute null statistics
            pmap_out = pmap(perm -> pmap_two_samp_null_stats(mod_s1, calls_s1, calls_s2, perm, config), comb_iter_used)

            # Distribute statistics
            tmml_perms = [x[1] for x in pmap_out]
            tnme_perms = [x[2] for x in pmap_out]
            tcmd_perms = [x[3] for x in pmap_out]

            ## P-value computation

            # Loop over analysis regions
            @inbounds for k in 1:length(tmml_obs)

                # Check if data
                isnan(tmml_obs[k]) && continue

                # Get permutation stats from k-th analysis region
                tmml_perms_k = [perm[k] for perm in tmml_perms]
                tnme_perms_k = [perm[k] for perm in tnme_perms]
                tcmd_perms_k = [perm[k] for perm in tcmd_perms]

                # Clean NaNs in null stats
                tmml_perms_k = tmml_perms_k[.!isnan.(tmml_perms_k)]
                tnme_perms_k = tnme_perms_k[.!isnan.(tnme_perms_k)]
                tcmd_perms_k = tcmd_perms_k[.!isnan.(tcmd_perms_k)]

                # Check enough null stats after filtering NaNs
                length(tmml_perms_k) > 20 || continue

                # Get number of permutation stats equal or above observed
                tmml_pval_k = sum(abs.(tmml_perms_k) .>= abs(tmml_obs[k]))
                tnme_pval_k = sum(abs.(tnme_perms_k) .>= abs(tnme_obs[k]))
                tcmd_pval_k = sum(tcmd_perms_k .>= tcmd_obs[k])

                # Compute p-values
                tmml_pvals[k] = exact ? tmml_pval_k / length(tmml_perms_k) : (1.0 + tmml_pval_k) / (1.0 + length(tmml_perms_k))
                tnme_pvals[k] = exact ? tnme_pval_k / length(tnme_perms_k) : (1.0 + tnme_pval_k) / (1.0 + length(tnme_perms_k))
                tcmd_pvals[k] = exact ? tcmd_pval_k / length(tcmd_perms_k) : (1.0 + tcmd_pval_k) / (1.0 + length(tcmd_perms_k))

            end

        end
    end

    # Fill return object
    @inbounds for k in 1:length(tmml_obs)
        test_struct.tests.tmml_test[k] = (tmml_obs[k], tmml_pvals[k])
        test_struct.tests.tnme_test[k] = (tnme_obs[k], tnme_pvals[k])
        test_struct.tests.tcmd_test[k] = (tcmd_obs[k], tcmd_pvals[k])
    end

    # Return test struct
    return test_struct

end
"""
    `pmap_diff_two_samp_comp_test(MODS_CHR_S1,MODS_CHR_S2,CONFIG)`

    Function that performs differential methylation analysis between two samples given two
    CpelNano models. This function is for testing purposes.

    # Examples
    ```julia-repl
    julia> CpelNano.pmap_diff_two_samp_comp_test(mods_chr_s1,mods_chr_s2,config)
    ```
"""
function pmap_diff_two_samp_comp_test(mod_s1::RegStruct, mod_s2::RegStruct, config::CpelNanoConfig)::RegStatTestStruct

    # Init output struct
    test_struct = RegStatTestStruct(mod_s1)

    # Observed statistics
    tmml_obs = mod_s1.mml - mod_s2.mml
    tnme_obs = mod_s1.nme - mod_s2.nme
    tcmd_obs = comp_cmd(mod_s1, mod_s2)

    # Init p-values
    tmml_pvals = fill(NaN, length(tmml_obs))
    tnme_pvals = fill(NaN, length(tnme_obs))
    tcmd_pvals = fill(NaN, length(tcmd_obs))

    # Compute p-values if need be
    if config.pval_comp

        # Number of reads
        n_rds_s1 = length(mod_s1.calls)
        n_rds_s2 = length(mod_s2.calls)

        # Compute number of possible randomizations
        L = binomial(BigInt(n_rds_s1 + n_rds_s2), BigInt(n_rds_s1))

        # Create iteratable object with all combinations
        comb_iter = combinations(1:(n_rds_s1 + n_rds_s2), n_rds_s1)

        # If enough permutations
        if L > 20

            #  Check if exact test
            exact = L < config.LMAX_TWO_SAMP

            # # Go back to group scale
            # get_grp_info!(mod_s1, fa_rec, config.min_grp_dist)
            # get_grp_info!(mod_s2, fa_rec, config.min_grp_dist)

            # Store calls in aux vars
            calls_s1 = deepcopy(mod_s1.calls)
            calls_s2 = deepcopy(mod_s2.calls)

            # Get read combinations to use
            if exact
                # Use all sample assignments
                comb_iter_used = comb_iter
            else
                # Sample assignments to group 1
                comb_iter_used = fill(Vector{Int64}(), config.LMAX_TWO_SAMP)
                @inbounds for i = 1:config.LMAX_TWO_SAMP
                    comb_iter_used[i] = sort!(StatsBase.sample(1:(n_rds_s1 + n_rds_s2), n_rds_s1;replace=false))
                    print_log("$(comb_iter_used[i])")
                end
            end

            ## Null statistics

            # Compute null statistics
            pmap_out = pmap(perm -> pmap_two_samp_null_stats_test_fun(mod_s1, calls_s1, calls_s2, perm, config), comb_iter_used)

            # Distribute statistics
            tmml_perms = [x[1] for x in pmap_out]
            tnme_perms = [x[2] for x in pmap_out]
            tcmd_perms = [x[3] for x in pmap_out]

            ## P-value computation

            # Loop over analysis regions
            @inbounds for k in 1:length(tmml_obs)

                # Check if data
                isnan(tmml_obs[k]) && continue

                # Get permutation stats from k-th analysis region
                tmml_perms_k = [perm[k] for perm in tmml_perms]
                tnme_perms_k = [perm[k] for perm in tnme_perms]
                tcmd_perms_k = [perm[k] for perm in tcmd_perms]

                # Clean NaNs in null stats
                tmml_perms_k = tmml_perms_k[.!isnan.(tmml_perms_k)]
                tnme_perms_k = tnme_perms_k[.!isnan.(tnme_perms_k)]
                tcmd_perms_k = tcmd_perms_k[.!isnan.(tcmd_perms_k)]

                # Check enough null stats after filtering NaNs
                length(tmml_perms_k) > 20 || continue

                # Get number of permutation stats equal or above observed
                tmml_pval_k = sum(abs.(tmml_perms_k) .>= abs(tmml_obs[k]))
                tnme_pval_k = sum(abs.(tnme_perms_k) .>= abs(tnme_obs[k]))
                tcmd_pval_k = sum(tcmd_perms_k .>= tcmd_obs[k])

                # Compute p-values
                tmml_pvals[k] = exact ? tmml_pval_k / length(tmml_perms_k) : (1.0 + tmml_pval_k) / (1.0 + length(tmml_perms_k))
                tnme_pvals[k] = exact ? tnme_pval_k / length(tnme_perms_k) : (1.0 + tnme_pval_k) / (1.0 + length(tnme_perms_k))
                tcmd_pvals[k] = exact ? tcmd_pval_k / length(tcmd_perms_k) : (1.0 + tcmd_pval_k) / (1.0 + length(tcmd_perms_k))

            end

        end
    end

    # Fill return object
    @inbounds for k in 1:length(tmml_obs)
        test_struct.tests.tmml_test[k] = (tmml_obs[k], tmml_pvals[k])
        test_struct.tests.tnme_test[k] = (tnme_obs[k], tnme_pvals[k])
        test_struct.tests.tcmd_test[k] = (tcmd_obs[k], tcmd_pvals[k])
    end

    # Return test struct
    return test_struct

end
"""
    `diff_two_samp_comp(MODEL_FILE_G1,MODEL_FILE_G2,FASTA,CONFIG)`

    Function that performs differential methylation analysis between two samples given two
    CpelNano model files.

    # Examples
    ```julia-repl
    julia> CpelNano.diff_two_samp_comp(mods_path_s1,mods_path_s2,fasta,config)
    ```
"""
function diff_two_samp_comp(mods_path_s1::String, mods_path_s2::String, nano_s1::String, nano_s2::String, fasta::String, config::CpelNanoConfig)::Nothing

    # Find chromosomes
    chr_names, chr_sizes = get_genome_info(fasta)

    # Create output directory if not existant
    isdir(config.out_dir) || mkdir(config.out_dir)

    # Get output file names
    config.out_diff_files = OutputDiffFiles(config.out_dir, config.out_prefix)

    # Check if output file exists
    check_diff_output_exists(config) && return nothing

    # Loop over chromosomes
    for (i, chr) in enumerate(chr_names)

        print_log("Processing chromosome $(chr)")

        ## Get fasta record

        fa_rec = get_chr_fa_rec(chr, fasta)

        ## Read in models

        mods_chr_s1 = read_mod_file_chr(mods_path_s1, chr, fa_rec, config)
        mods_chr_s2 = read_mod_file_chr(mods_path_s2, chr, fa_rec, config)

        ## Find regions with at least one sample per group

        comm_regs = find_comm_regs([mods_chr_s1], [mods_chr_s2])

        # If no data in chromosome continue
        length(comm_regs) > 0 || continue
        config.verbose && print_log("Comm regions: $(comm_regs)")

        ## Process each region
        pmap_out = map(x -> pmap_diff_two_samp_comp(mods_chr_s1[x], mods_chr_s2[x], nano_s1, nano_s2, fa_rec, config), comm_regs)

        ## Write chromosome output
        write_diff_output(pmap_out, config)

    end

    # Perform multiple hypothesis testing
    if config.pval_comp
        print_log("Multiple hypothesis testing")
        mult_hyp_corr(config)
    end

    # Return nothing
    return nothing

end


using RasterDataSources, Test, Dates
using RasterDataSources: rasterurl, rastername, rasterpath

@testset "CHELSA Future BioClim CMIP5" begin
    @test rastername(CHELSA{Future{BioClim,CMIP5,CCSM4,RCP26}}, 5; date=Date(2050)) == "CHELSA_bio_mon_CCSM4_rcp26_r1i1p1_g025.nc_5_2041-2060_V1.2.tif"

    bioclim_path = joinpath(ENV["RASTERDATASOURCES_PATH"], "CHELSA", "Future", "BioClim", "RCP26", "CCSM4")
    @test rasterpath(CHELSA{Future{BioClim,CMIP5,CCSM4,RCP26}}) == bioclim_path

    raster_path = joinpath(bioclim_path, "CHELSA_bio_mon_CCSM4_rcp26_r1i1p1_g025.nc_5_2041-2060_V1.2.tif")

    @test getraster(CHELSA{Future{BioClim,CMIP5,CCSM4,RCP26}}, 5; date=Date(2050)) == raster_path
    @test getraster(CHELSA{Future{BioClim,CMIP5,CCSM4,RCP26}}, (5,); date=Date(2050)) == (bio5=raster_path,)
    @test getraster(CHELSA{Future{BioClim,CMIP5,CCSM4,RCP26}}, [5]; date=Date(2050)) == (bio5=raster_path,)
    @test isfile(raster_path)
end

@testset "CHELSA Future BioClim CMIP6" begin
    bioclim_name = "CHELSA_bio5_2041-2070_mri-esm2-0_ssp126_V.2.1.tif"
    @test rastername(CHELSA{Future{BioClim,CMIP6,MRIESM2,SSP126}}, 5; date=Date(2050)) == bioclim_name

    bioclim_path = joinpath(ENV["RASTERDATASOURCES_PATH"], "CHELSA", "Future", "BioClim", "SSP126", "MRIESM2")
    @test rasterpath(CHELSA{Future{BioClim,CMIP6,MRIESM2,SSP126}}) == bioclim_path

    raster_path = joinpath(bioclim_path, bioclim_name)
    raster_path2 = joinpath(bioclim_path, "CHELSA_bio5_2071-2100_mri-esm2-0_ssp126_V.2.1.tif")
    @test getraster(CHELSA{Future{BioClim,CMIP6,MRIESM2,SSP126}}, 5; date=Date(2050)) == raster_path
    @test getraster(CHELSA{Future{BioClim,CMIP6,MRIESM2,SSP126}}, (5,); date=Date(2050)) == (bio5=raster_path,)
    @test getraster(CHELSA{Future{BioClim,CMIP6,MRIESM2,SSP126}}, [5]; date=Date(2050)) == (bio5=raster_path,)
    @test getraster(CHELSA{Future{BioClim,CMIP6,MRIESM2,SSP126}}, (5,); date=[Date(2050)]) ==
        [(bio5=raster_path,)]
    @test isfile(raster_path)
end

@testset "CHELSA Future Climate CMIP5" begin
    tmax_name = "CHELSA_tasmax_mon_CCSM4_rcp60_r1i1p1_g025.nc_7_2041-2060_V1.2.tif"
    @test rastername(CHELSA{Future{Climate,CMIP5,CCSM4,RCP60}}, :tmax; date=Date(2050), month=7) == tmax_name

    climate_path = joinpath(ENV["RASTERDATASOURCES_PATH"], "CHELSA", "Future", "Climate", "RCP60", "CCSM4")
    @test rasterpath(CHELSA{Future{Climate,CMIP5,CCSM4,RCP60}}) == climate_path

    raster_path = joinpath(climate_path, tmax_name)
    @test rasterpath(CHELSA{Future{Climate,CMIP5,CCSM4,RCP60}}, :tmax; date=Date(2050), month=7) == raster_path
    @test rasterurl(CHELSA{Future{Climate,CMIP5,CCSM4,RCP45}}, :prec; date=Date(2050), month=6) |> string ==
        "https://os.zhdk.cloud.switch.ch/envicloud/chelsa/chelsa_V1/cmip5/2041-2060/prec/CHELSA_pr_mon_CCSM4_rcp45_r1i1p1_g025.nc_6_2041-2060.tif"
    @test rasterurl(CHELSA{Future{Climate,CMIP5,CCSM4,RCP45}}, :tmin; date=Date(2050), month=1) |> string ==
        "https://os.zhdk.cloud.switch.ch/envicloud/chelsa/chelsa_V1/cmip5/2041-2060/tmin/CHELSA_tasmin_mon_CCSM4_rcp45_r1i1p1_g025.nc_1_2041-2060_V1.2.tif"
    @test getraster(CHELSA{Future{Climate,CMIP5,CCSM4,RCP60}}, :tmax; date=Date(2050), month=7) == raster_path
    @test getraster(CHELSA{Future{Climate,CMIP5,CCSM4,RCP60}}, [:tmax]; date=Date(2050), month=7) == (tmax=raster_path,)
    @test getraster(CHELSA{Future{Climate,CMIP5,CCSM4,RCP60}}, (:tmax,); date=Date(2050), month=7:7) == [(tmax=raster_path,)]
    @test getraster(CHELSA{Future{Climate,CMIP5,CCSM4,RCP60}}, :tmax; date=[Date(2050)], month=7:7) == [[raster_path]]
    @test isfile(raster_path)
end

@testset "CHELSA Future Climate CMIP6" begin
    date_name =  "CHELSA_gfdl-esm4_r1i1p1f1_w5e5_ssp585_tas_01_2011_2040_norm.tif"
    date_name2 = "CHELSA_gfdl-esm4_r1i1p1f1_w5e5_ssp585_tas_01_2041_2070_norm.tif"
    date_name3 = "CHELSA_gfdl-esm4_r1i1p1f1_w5e5_ssp585_tas_01_2071_2100_norm.tif"
    month_name =  "CHELSA_gfdl-esm4_r1i1p1f1_w5e5_ssp585_tas_01_2011_2040_norm.tif"
    month_name2 = "CHELSA_gfdl-esm4_r1i1p1f1_w5e5_ssp585_tas_02_2011_2040_norm.tif"
    @test rastername(CHELSA{Future{Climate,CMIP6,GFDLESM4,SSP585}}, :temp; date=Date(2030), month=1) == date_name

    climate_path = joinpath(ENV["RASTERDATASOURCES_PATH"], "CHELSA", "Future", "Climate", "SSP585", "GFDLESM4")
    @test rasterpath(CHELSA{Future{Climate,CMIP6,GFDLESM4,SSP585}}) == climate_path
    date_path = joinpath(climate_path, date_name)
    date_path2 = joinpath(climate_path, date_name2)
    date_path3 = joinpath(climate_path, date_name3)
    month_path = joinpath(climate_path, month_name)
    month_path2 = joinpath(climate_path, month_name2)
    @test rasterpath(CHELSA{Future{Climate,CMIP6,GFDLESM4,SSP585}}, :temp; date=Date(2030), month=1) == date_path
    @test rasterpath(CHELSA{Future{Climate,CMIP6,GFDLESM4,SSP585}}, :temp; date=Date(2030), month=1) == date_path
    date_url = "https://os.zhdk.cloud.switch.ch/envicloud/chelsa/chelsa_V2/GLOBAL/climatologies/2011-2040/GFDL-ESM4/ssp585/tas/" * date_name
    @test rasterurl(CHELSA{Future{Climate,CMIP6,GFDLESM4,SSP585}}, :temp; date=Date(2030), month=1) |> string == date_url
    @test getraster(CHELSA{Future{Climate,CMIP6,GFDLESM4,SSP585}}, :temp; date=Date(2030), month=1) == date_path
    @test getraster(CHELSA{Future{Climate,CMIP6,GFDLESM4,SSP585}}, (:temp,); date=Date(2030), month=1) == (temp=date_path,)
    @test getraster(CHELSA{Future{Climate,CMIP6,GFDLESM4,SSP585}}, :temp; date=Date(2030), month=1) == date_path
    # Month is the inner vector
    @test getraster(CHELSA{Future{Climate,CMIP6,GFDLESM4,SSP585}}, [:temp]; date=[Date(2030)], month=1:2) ==
        [[(temp=month_path,), (temp=month_path2,)]]
    @test getraster(CHELSA{Future{Climate,CMIP6,GFDLESM4,SSP585}}, [:temp]; date=(Date(2030), Date(2090)), month=1:1) ==
        [[(temp=date_path,)], [(temp=date_path2,)], [(temp=date_path3,)]]

    @test isfile(date_path)
    @test isfile(date_path2)
    @test isfile(date_path3)
    @test isfile(month_path)
    @test isfile(month_path2)
end


"""
Run the PIMDP synthesis
"""
function run_pimdp_synthesis(pimdp::PIMDP, pimdp_filename::String; add_opt=false, horizon=-1)
    write_pimdp_to_file(pimdp, pimdp_filename)

    @info "Synthesizing a controller... (maximize pessimistic)"
    res_mat = run_unbounded_imdp_synthesis(pimdp_filename)

    if add_opt
        @info "Synthesizing a controller... (maximize optimistic)"
        basename = pimdp_filename[1:end-3]
        opt_pimdp_filename = "$basename-optimistic.txt"
        res_mat_opt = run_imdp_synthesis(pimdp_filename, horizon, mode2="optimistic", save_mats=false)
        for j in 1:length(res_mat[:,1])
            if res_mat[j, 3] == res_mat_opt[j,3] && res_mat_opt[j,4] > res_mat[j,4]
                res_mat[j, 2] = res_mat_opt[j,2]
                res_mat[j, 4] = res_mat_opt[j,4]
            end
        end
    end

    return res_mat
end


module VITA2D

export Model, Config, Gradient, Thermo, StochasticMatrix, SGDParams
export half, third
export warmup!, mc, pmc
export SGDParams
export sgd

using Parameters
using DelimitedFiles
using Dates
using Random
using LinearAlgebra
using Statistics
using Distributed

set_zero_subnormals(true)

# @everywhere using Pkg
# @everywhere Pkg.activate("/home/mbeach/VITA2D.jl")
# @everywhere Pkg.instantiate()
# @everywhere Pkg.update()
# @everywhere Pkg.resolve()

include("types.jl")
include("utils.jl")
include("sr.jl")
include("metro.jl")
include("update.jl")
include("mc.jl")
include("parallel.jl")
include("parallel_optimise.jl")

end # module


function erarawregion(
    emod::Dict, epar::Dict, ereg::Dict,
    start::TimeType, finish::TimeType;
    path::AbstractString="",
    region::AbstractString="GLB"
)

    if path == ""; eroot = eraroot(actionID); else; eroot = eraroot(path,actionID); end

    @info "$(Dates.now()) - Extracting $(uppercase(emod["dataset"])) $(epar["name"]) data for the entire $(gregionfullname(region)) region ..."

    if !isdir(eraregfolder(ereg,eroot))
        error("$(Dates.now()) - No data has been downloaded from $(uppercase(emod["dataset"])) $(epar["name"]) in the $(gregionfullname(region))")
    end

    dvec,dys,dyf,ndy = extractdate(start,finish); ndates = length(dvec);
    if emod["datasetID"] == 1; nt = 24; else; nt = 4; end

    lon = ereg["lon"]; lat = ereg["lat"]; rlon,rlat,rinfo = gregiongridvec(region,lon,lat);
    nlon = length(rlon); nlat = length(rlat); data = zeros(Float32,nlon,nlat,ndy*nt);

    for ii = 1 : ndates; dateii = dvec[ii];

        fol = erarawfolder(epar,ereg,eroot);
        if !isdir(fol)
            error("$(Dates.now()) - There is no data for $(uppercase(emod["dataset"])) $(epar["name"]) for $(yrmo2dir(dateii)).")
        end

        fnc = joinpath(fol,erarawname(emod,epar,ereg,dateii));
        ds  = Dataset(fnc,"r"); vds = ds[varname];

        if     ii == 1 && ii != ndates
            moday = daysinmonth(dateii);
            ibeg = (dys-1)*nt + 1; iend = (moday+1-dys)*nt;
            data[:,:,1:iend] = vds[:,:,ibeg:end];
        elseif ii != 1 && ii == ndates
            ibeg = iend+1; iend = dyf*nt;
            data[:,:,ibeg:end] = vds[:,:,1:iend];
        elseif ii == 1 && ii == ndates
            ibeg = (dys-1)*nt + 1; iend = dyf*nt;
            data = vds[:,:,ibeg:iend];
        else
            moday = daysinmonth(dateii);
            ibeg = iend+1; iend = ibeg-1 + moday*nt;
            data[:,:,ibeg:iend] = vds[:];
        end

        close(ds);

    end

    @info "$(Dates.now()) - $(uppercase(emod["dataset"])) $(epar["name"]) data for the entire $(gregionfullname(region)) region has been extracted."

    return datavec,info,[rlon,rlat]

end

function erarawpoint(
    emod::Dict, epar::Dict, ereg::Dict,
    start::TimeType, finish::TimeType;
    coord::Array{<:Real,1},
    path::AbstractString="",
    region::AbstractString="GLB"
)

    if path == ""; eroot = eraroot(actionID); else; eroot = eraroot(path,actionID); end
    if length(coord) != 2
        error("$(Dates.now()) - Coordinate vector must be in the form [lon,lat]")
    end

    @info "$(Dates.now()) - Extracting $(uppercase(emod["dataset"])) $(epar["name"]) data at coordinates $(coord) ..."

    if !isdir(eraregfolder(ereg,eroot))
        error("$(Dates.now()) - No data has been downloaded from $(uppercase(emod["dataset"])) $(epar["name"]) in the $(gregionfullname(region))")
    end

    dvec,dys,dyf,ndy = extractdate(start,finish); ndates = length(dvec);
    if emod["datasetID"] == 1; nt = 24; else; nt = 4; end

    lon = ereg["lon"]; lat = ereg["lat"]; rlon,rlat,rinfo = gregiongridvec(region,lon,lat);
    plon,plat = coord; ilon,ilat = regionpoint(plon,plat,rlon,rlat);
    data = zeros(Float32,ndy*nt);

    for ii = 1 : ndates; dateii = dvec[ii];

        fol = erarawfolder(epar,ereg,eroot);
        if !isdir(fol)
            error("$(Dates.now()) - There is no data for $(uppercase(emod["dataset"])) $(epar["name"]) for $(yrmo2dir(dateii)).")
        end

        fnc = joinpath(fol,erarawname(emod,epar,ereg,dateii));
        ds  = Dataset(fnc,"r"); vds = ds[varname];

        if     ii == 1 && ii != ndates
            moday = daysinmonth(dateii);
            ibeg = (dys-1)*nt + 1; iend = (moday+1-dys)*nt;
            data[1:iend] = vds[ilon,ilat,ibeg:end];
        elseif ii != 1 && ii == ndates
            ibeg = iend+1; iend = dyf*nt;
            data[ibeg:end] = vds[ilon,ilat,1:iend];
        elseif ii == 1 && ii == ndates
            ibeg = (dys-1)*nt + 1; iend = dyf*nt;
            data = vds[ilon,ilat,ibeg:iend];
        else
            moday = daysinmonth(dateii);
            ibeg = iend+1; iend = ibeg-1 + moday*nt;
            data[ibeg:iend] = vds[ilon,ilat,:];
        end

        close(ds);

    end

    @info "$(Dates.now()) - $(uppercase(emod["dataset"])) $(epar["name"]) data for the coordinates $(coord) has been extracted."

    return datavec,info

end

function erarawgrid(
    emod::Dict, epar::Dict, ereg::Dict,
    start::TimeType, finish::TimeType;
    grid::Array{<:Real,1},
    path::AbstractString="",
    region::AbstractString="GLB"
)

    if path == ""; eroot = eraroot(actionID); else; eroot = eraroot(path,actionID); end
    if length(grid) != 4
        error("$(Dates.now()) - Grid vector must be in the form [N,S,E,W]")
    end

    @info "$(Dates.now()) - Extracting $(uppercase(emod["dataset"])) $(epar["name"]) data for the [N,S,E,W] bounds $(grid) ..."

    if !isdir(eraregfolder(ereg,eroot))
        error("$(Dates.now()) - No data has been downloaded from $(uppercase(emod["dataset"])) $(epar["name"]) in the $(gregionfullname(region))")
    end

    dvec,dys,dyf,ndy = extractdate(start,finish); ndates = length(dvec);
    if emod["datasetID"] == 1; nt = 24; else; nt = 4; end

    lon = ereg["lon"]; lat = ereg["lat"]; isgridinregion(grid,region);
    rlon,rlat,rinfo = gregiongridvec(region,lon,lat);
    glon,glat,ginfo = regiongridvec(grid,rlon,rlat); iWE,iNS = ginfo["IDvec"];
    nlon = length(glon); nlat = length(glat);
    data = zeros(Float32,nlon,nlat,ndy*nt);

    for ii = 1 : ndates; dateii = dvec[ii];

        fol = erarawfolder(epar,ereg,eroot);
        if !isdir(fol)
            error("$(Dates.now()) - There is no data for $(uppercase(emod["dataset"])) $(epar["name"]) for $(yrmo2dir(dateii)).")
        end

        fnc = joinpath(fol,erarawname(emod,epar,ereg,dateii));
        ds  = Dataset(fnc,"r"); vds = ds[varname];

        if     ii == 1 && ii != ndates
            moday = daysinmonth(dateii);
            ibeg = (dys-1)*nt + 1; iend = (moday+1-dys)*nt;
            data[:,:,1:iend] = vds[iWE,iNS,ibeg:end];
        elseif ii != 1 && ii == ndates
            ibeg = iend+1; iend = dyf*nt;
            data[:,:,ibeg:end] = vds[iWE,iNS,1:iend];
        elseif ii == 1 && ii == ndates
            ibeg = (dys-1)*nt + 1; iend = dyf*nt;
            data = vds[iWE,iNS,ibeg:iend];
        else
            moday = daysinmonth(dateii);
            ibeg = iend+1; iend = ibeg-1 + moday*nt;
            data[:,:,ibeg:iend] = vds[iWE,iNS,:];
        end

        close(ds);

    end

    @info "$(Dates.now()) - $(uppercase(emod["dataset"])) $(epar["name"]) data within the [N,S,E,W] bounds $(grid) has been extracted."

    return datavec,info,[glon,glat]

end