"""
Data module for preparing data for analysis with
    - GaussianProcess (src/gpr/GaussianProcess.jl)
"""

module Data

using OceanTurb
using OrderedCollections

include("../LES/custom_avg.jl")
export custom_avg

# harvesting Oceananigans data
include("../LES/read_les_output.jl")
export get_les_data

# running OceanTurb KPP simulations based on OceananigansData conditions
include("../KPP/run.jl")
export closure_free_convection_kpp_full_evolution,
       closure_free_convection_kpp

# running OceanTurb TKE simulations based on OceananigansData conditions
include("../TKE/run.jl")
export closure_free_convection_tke_full_evolution,
       closure_free_convection_tke

# normalization
include("scalings.jl")
export  min_max_scaling,
        scale, # normalize
        unscale # un-normalize

# pre- and post-processing on the normalized data
include("problems.jl")
export  Problem,
        Sequential,
        Residual,
        Slack,
        SequentialProblem,
        ResidualProblem,
        SlackProblem,
        get_problem_v,
        predict # important!

export Sequential_KPP,
       Sequential_TKE

include("sequential.jl")
include("residual.jl")
include("slack.jl")
export  get_predictors_targets,
        postprocess_prediction

include("modify_predictor_fns.jl")
export append_tke,
       partial_temp_profile

include("convective_adjust.jl")
export convective_adjust!

# ProfileData struct
export  ProfileData,
        data

function approx_initial_buoyancy_stratification(T_initial,z)
    # b_initial = 𝒟.T[:,1] .* α*g
    α = 2e-4
    g = 9.80665
    N² = (T_initial[1] - T_initial[end])*α*g / z[1] # approximate initial buoyancy gradient N², where b = N²z + 20*α*g and T = N²z/(αg) + 20
    return N²
end

struct StateVariables
    tke_avg # turbulent kinetic energy
end

"""
ProfileData
# Description
- data structure for preparing profile data for analysis with gpr or nn.
# Data Structure and Description
    v::Array,           Nz x Nt array of T or wT values directly from the LES simulation, not preprocessed.
    vavg::Array,        Nt-length array of Nz-length vectors from LES simulation, scaled and pre-processed
    x::Array,           all simulation inputs, scaled and pre-processed
    y::Array,           all simulation inputs, scaled and pre-processed
    x_train::Array,     training inputs (predictors; array of states). (length-n array of D-length vectors, where D is the length of each input n is the number of training points)
    y_train::Array,     training outputs (predictions) (length-n array of D-length vectors).
    validation_set::Array, vector of indices corresponding to validation data
    z::Array,           Nz-length vector of depth values
    zavg::Array,        length-D vector; depth values averaged to D gridpoints
    t::Array,           timeseries [seconds]
    Nt::Int64,          length(timeseries)
    problem::Problem,   what mapping you wish to evaluate with the model. (Sequential("T"), Sequential("wT"), Residual("T"), Residual("KPP"), or Residual("TKE"))
    all_problems::Array{Array{Any,1},1}
    state_variables::StateVariables
    modify_predictor_fn::Function
    convective_adjust::Function
"""
struct ProfileData
    v       ::Array{Float64,2}
    vavg    ::Array{Array{Float64,1},1}
    x       ::Array{Array{Float64,1},1}
    y       ::Array{Array{Float64,1},1}
    x_train ::Array{Array{Float64,1},1}
    y_train ::Array{Array{Float64,1},1}
    validation_set::Vector{Int64}
    z       ::Vector{Float64}
    zavg    ::Vector{Float64}
    t       ::Vector{Float64}
    Nt      ::Int64
    problem ::Problem
    all_problems::Array{Array{Any,1},1}
    state_variables::StateVariables
    modify_predictor_fn::Function
    convective_adjust::Function
end

"""
data(filename, problem; D=16, N=4)

# Description
Returns a ProfileData object based on data from `filename`

# Arguments
- `filename`: (string)  Name of the NetCDF or JLD2 file containing the data from the Oceananigans simulation.
- `problem`: (Problem). What mapping you wish to evaluate with the model.

# Keyword Arguments
- `D`: (integer)        Number of gridpoints in the z direction to average the data to for training and prediction.
- `N`: (integer)        Interval between the timesteps to be reserved for training data (default 4).
                        If N=4, the profile data for every 4 timesteps will be reserved for training (~25% training data);
                        the rest will be used in the validation set.
"""
function data(filename::String, problem::Problem; D=16, N=4)

    # collect data from Oceananigans simulation output file
    les = get_les_data(filename) # <: OceananigansData

    # eddy diffusivity
    κₑ = les.κₑ

    # depth values
    z = les.z
    zavg = custom_avg(z, D)

    # get v (variable array, Nz x Nt) and cut out the first 2 hours
    start = floor(Int64, 7200 / (les.t[2] - les.t[1]))+1
    s = start:length(les.t)
    v = get_v(problem, les)[:,s]

    # timeseries [s]
    t = les.t[s]
    Nt = length(t)

    # approximate buoyancy stratification at the initial timestep
    N² = approx_initial_buoyancy_stratification(les.T[:,1],z)

    # state variables
    array_to_avg(array) = [custom_avg(array[:,j], D) for j in s]
    state_variables = StateVariables(array_to_avg(les.tke))

    # modify_predictor_fn
    modify_predictor_fn(state, time_index) = problem.modify_predictor_fn(state, time_index, state_variables)

    # convective_adjust
    if problem.convective_adjust
        convective_adjust = convective_adjust!
    else
        convective_adjust = x -> x
    end

    # get problem (sets how the data will be pre- and post-processed)
    specific_problem = get_problem(problem, les, v, N², D, Nt)

    # compress variable array to D gridpoints to get an Nt-length array of D-length vectors
    vavg = [custom_avg(v[:,j], D) for j in 1:Nt]

    # divide up the data
    total_set = 1:(Nt-1)
    training_set = 1:N:(Nt-1)
    validation_set = setdiff(total_set, training_set)
    n_train = length(training_set)

    # get the predictors and target predictions to train the model on
    x, y = get_predictors_targets(vavg, specific_problem)

    x_train = [modify_predictor_fn(x[i], i) for i in training_set]
    y_train = y[training_set]

    # if the data contains multiple files, we'll need a way to postprocess the data separately for each one
    # (preprocessing is already taken care of before the data is merged)
    all_problems = [[specific_problem, length(x)]]

    return ProfileData(v, vavg, x, y, x_train, y_train, validation_set, z, zavg, t, Nt, specific_problem, all_problems, state_variables, modify_predictor_fn, convective_adjust)
end

"""
data(filename, D; N=4)

# Description
Returns an instance of ProfileData containing training data from multiple simulations.
*** Important:
    ONLY v, x_train and y_train contain data from all filenames;
    the remaining attributes are taken from the first filename in filenames

# Arguments
- `filenames`: (string)  Vector of filenames (.nc or .jld2) to collect data from.
- `problem`: (Problem). What mapping you wish to evaluate with the model. (Sequential("T"), Sequential("wT"), Residual("T"), Residual("KPP"), or Residual("TKE"))

# Keyword Arguments
- `D`: (integer)        Number of gridpoints in the z direction to average the data to for training and prediction.
- `N`: (integer)        Interval between the timesteps to be reserved for training data (default 4).
                        If N=4, the profile data for every 4 timesteps will be reserved for training (~25% training data);
                        the rest will be used in the validation set.
"""
function data(filenames::Vector{String}, problem::Problem; D=16, N=4)

    # combines data from multiple files
    𝒟 = data(filenames[1], problem; D=D, N=N)

    v = 𝒟.v
    vavg = 𝒟.vavg
    x = 𝒟.x
    y = 𝒟.y
    x_train = 𝒟.x_train
    y_train = 𝒟.y_train
    tke_avg = 𝒟.state_variables.tke_avg
    validation_set = 𝒟.validation_set
    all_problems = 𝒟.all_problems

    t = 𝒟.t
    Nt = 𝒟.Nt

    for filename in filenames[2:end]
        𝒟2 = data(filename, problem; D=D, N=N)

        v = hcat(v, 𝒟2.v) # unscaled
        vavg = vcat(vavg, 𝒟2.vavg)
        x = vcat(x, 𝒟2.x)
        y = vcat(y, 𝒟2.y)
        validation_set = vcat(validation_set, 𝒟2.validation_set)
        x_train = vcat(x_train, 𝒟2.x_train)
        y_train = vcat(y_train, 𝒟2.y_train)
        tke_avg = vcat(tke_avg, 𝒟2.state_variables.tke_avg)
        append!(all_problems, 𝒟2.all_problems)
        t = vcat(t, 𝒟2.t)
        Nt += 𝒟2.Nt
    end

    state_variables = StateVariables(tke_avg)
    modify_predictor_fn(state, time_index) = problem.modify_predictor_fn(state, time_index, state_variables)

    # Note the problem is that from the first file in filenames. This is only included so that the problem type can be determined easily.
    return ProfileData(v, vavg, x, y, x_train, y_train, validation_set, 𝒟.z, 𝒟.zavg, t, Nt, 𝒟.problem, all_problems, state_variables, modify_predictor_fn, 𝒟.convective_adjust)
end

end # module


function init_state(model::AbstractGPModel)
    state = init_local_vars((;), model)
    state = init_opt_state(state, model)
    state = init_hyperopt_state(state, model)
    if n_output(model) > 1
        state = init_state_A(state, model)
    end
    return state
end

@traitfn function init_local_vars(
    state, model::TGP
) where {TGP <: AbstractGPModel; !IsMultiOutput{TGP}}
    if inference(model) isa Union{Analytic,AnalyticVI,GibbsSampling}
        local_vars = init_local_vars(likelihood(model), batchsize(model))
        return merge(state, (; local_vars))
    else
        return state
    end
end

@traitfn function init_local_vars(
    state, model::TGP
) where {TGP <: AbstractGPModel; IsMultiOutput{TGP}}
    if inference(model) isa Union{Analytic,AnalyticVI,GibbsSampling}
        local_vars = init_local_vars.(likelihood(model), batchsize(model))
        return merge(state, (; local_vars))
    else
        return state
    end
end

function init_opt_state(state, model::AbstractGPModel)
    if inference(model) isa VariationalInference
        opt_state = map(model.f) do gp
            init_opt_state(gp, inference(model))
        end
        return merge(state, (; opt_state))
    else
        return state
    end
end

function init_opt_state(
    ::Union{VarLatent{T},TVarLatent{T}}, vi::VariationalInference
) where {T}
    return (;)
end

function init_opt_state(gp::Union{VarLatent{T},TVarLatent{T}}, vi::NumericalVI) where {T}
    return (;
        ν=zeros(T, batchsize(vi)), # Derivative -<dv/dx>_qn
        λ=zeros(T, batchsize(vi)), # Derivative  <d²V/dx²>_qm
        state_μ=Optimisers.state(opt(vi).optimiser, mean(gp)),
        state_Σ=Optimisers.state(opt(vi), cov(gp).data),
        ∇η₁=zero(mean(gp)),
        ∇η₂=zero(cov(gp).data),
    )
end

function init_opt_state(gp::SparseVarLatent{T}, vi::VariationalInference) where {T}
    state = (; ∇η₁=zero(mean(gp)), ∇η₂=zero(cov(gp).data))
    if vi isa AnalyticVI && is_stochastic(vi)
        state = merge(
            state,
            (;
                state_η₁=Optimisers.state(opt(vi).optimiser, nat1(gp)),
                state_η₂=Optimisers.state(opt(vi), nat2(gp).data),
            ),
        )
    end
    if vi isa NumericalVI
        state = merge(
            state,
            (;
                ν=zeros(T, batchsize(vi)), # Derivative -<dv/dx>_qn
                λ=zeros(T, batchsize(vi)), # Derivative  <d²V/dx²>_qm
                state_μ=Optimisers.state(opt(vi).optimiser, mean(gp)),
                state_Σ=Optimisers.state(opt(vi), cov(gp).data),
            ),
        )
    end
    return state
end

function init_opt_state(gp::OnlineVarLatent{T}, vi::VariationalInference) where {T}
    state = (; ∇η₁=zero(mean(gp)), ∇η₂=zero(cov(gp).data))
    if is_stochastic(vi)
        state_η₁ = state(opt(vi), nat1(gp))
        state_η₂ = state(opt(vi), nat2(gp).data)
        state = merge(state, (; state_η₁, state_η₂))
    end
    k = dim(gp)
    prev𝓛ₐ = zero(T)
    invDₐ = Symmetric(Matrix{T}(I(k)))
    prevη₁ = zeros(T, k)
    return merge(state, (; previous_gp=(; prev𝓛ₐ, invDₐ, prevη₁)))
end

function init_state_A(state, model::AbstractGPModel)
    A_state = [
        [Optimisers.init(model.A_opt, model.A[i][j]) for j in 1:model.nf_per_task[i]] for
        i in 1:n_output(model)
    ]
    return merge(state, (; A_state))
end

function init_hyperopt_state(state, model::GP)
    hyperopt_state = init_hyperopt_state(model.f)
    return merge(state, (; hyperopt_state))
end

function init_hyperopt_state(state, model::AbstractGPModel)
    hyperopt_state = map(model.f) do gp
        init_hyperopt_state(gp)
    end
    return merge(state, (; hyperopt_state))
end

@traitfn function init_hyperopt_state(gp::TGP) where {TGP <: AbstractLatent; IsFull{TGP}}
    hyperopt_state = (;)
    if !isnothing(opt(gp))
        k = kernel(gp)
        state_k = Optimisers.state(opt(gp), k)
        hyperopt_state = merge(hyperopt_state, (; state_k))
    end
    hyperopt_state = init_priormean_state(hyperopt_state, pr_mean(gp))
    return hyperopt_state
end

@traitfn function init_hyperopt_state(gp::TGP) where {TGP <: AbstractLatent; !IsFull{TGP}}
    hyperopt_state = (;)
    if !isnothing(opt(gp))
        k = kernel(gp)
        state_k = Optimisers.state(opt(gp), k)
        hyperopt_state = merge(hyperopt_state, (; state_k))
    end
    if !isnothing(Zopt(gp))
        Z = Zview(gp)
        state_Z = Optimisers.state(opt(gp), Z)
        hyperopt_state = merge(hyperopt_state, (; state_Z))
    end
    hyperopt_state = init_priormean_state(hyperopt_state, pr_mean(gp))
    return hyperopt_state
end

function Optimisers.state(opt, Z::Union{ColVecs,RowVecs})
    return Optimisers.state(opt, Z.X)
end


using Distributed
using StatsBase: mad, iqr, zscore
using Printf
import Glob
import Dierckx
import SparseArrays: spdiagm

include("./blocks.jl")

"""Helper to make a group path similar to `dset`, with new base"""
_match_dset_path(dset, newgroup) = join([newgroup; split(dset, '/')[2:end]], '/')

# E.g. `counts/1` when passed `stack/1`
_count_dset(dset) = _match_dset_path(dset, "counts")
_excl_dset(dset) = _match_dset_path(dset, "excluded")
# Used if we want to track the standard deviaition
# _stddev_dset(dset) = _match_dset_path(dset, "stddev")
# _stddev_raw_dset(dset) = _match_dset_path(dset, "stddev_raw")

# I'll be passing the included, clean geos, and want to record the bad missing ones
# is_excluded(total, included) = isnothing.(indexin(total, included))  # Too much space :(

function proc_pixel_linear(
    unw_stack_file,
    in_dset,
    valid_igram_indices,
    unmasked_idx_set,
    outfile,
    outdset,
    geolist,
    intlist,
    alpha,
    L1 = false;
    prune_outliers = true,
    sigma = 4,
    row_slice = nothing,
    col_slice = nothing,
)
    # in_buf = zeros(Float32, (length(col_slice), length(row_slice), length(valid_igram_indices)))
    in_buf = zeros(Float32, (length(row_slice), length(col_slice), length(valid_igram_indices)))
    try
        # in_buf .= h5read(unw_stack_file, in_dset, (col_slice, row_slice, :))[:, :, valid_igram_indices]
        in_buf .= h5read(unw_stack_file, in_dset, (row_slice, col_slice, :))[:, :, valid_igram_indices]
    catch e
        println(e, " ", "fail ", row_slice, " ", col_slice)
        return
    end
    # out_buf = Array{Float32, 2}(undef, (length(row_slice), length(col_slice)))
    # out_buf = Array{Float32, 2}(undef, (length(col_slice), length(row_slice)))
    # out_buf = zeros(Float32, (length(col_slice), length(row_slice)))
    out_buf = zeros(Float32, (length(row_slice), length(col_slice)))
    # println(row_slice, col_slice, length(row_slice), length(col_slice))

    for (i,row) in enumerate(row_slice), (j,col) in enumerate(col_slice)
    # seems to be transposed
    # for (i,row) in enumerate(col_slice), (j,col) in enumerate(row_slice)
        !((row, col) in unmasked_idx_set) && continue

        unw_pixel_raw = in_buf[i, j, :]
        # unw_pixel_raw = h5read(unw_stack_file, in_dset, (row, col, :))[1, 1, valid_igram_indices]
        if any(isnan.(unw_pixel_raw))
            println("$row, $col: nan")
            return
        end

        # Also load correlations for cutoff
        # cor_pixel = h5read(CC_FILENAME, "stack", (row, col, :))[1, 1, valid_igram_indices]
        # cor_thresh = 0.05
        # cor_pixel, cor_thresh = nothing, 0.0

        linear = true
        soln_phase, igram_count, geo_clean, unw_clean = calc_soln(
            unw_pixel_raw,
            geolist,
            intlist,
            alpha,
            linear;
            L1 = L1,
            prune_outliers = prune_outliers,
            sigma = sigma,
        )

        # println(size(soln_phase), size(out_buf))
        # println(soln_phase, " $i $j")
        out_buf[i, j] = soln_phase[1] * P2MM
    end
    # println(row_slice, " ", col_slice)
    # println(out_buf)

    dist_outfile = string(Distributed.myid()) * outfile
    h5open(dist_outfile, "r+") do f
        f[outdset][row_slice, col_slice] = out_buf
        # f[outdset][col_slice, row_slice] = out_buf
        # f[outdset][row, col] = P2MM * soln_phase
        # f[_count_dset(outdset)][row, col] = igram_count
        # f[_excl_dset(outdset)][row, col] = encode_missing(geolist, geo_clean)
        # f[_stddev_dset(outdset)][row, col] = std(unw_clean) .* abs(PHASE_TO_CM)
        # f[_stddev_raw_dset(outdset)][row, col] = std(unw_pixel_raw) .* abs(PHASE_TO_CM)
    end
    return
end

function proc_pixel_daily(
    unw_stack_file,
    in_dset,
    valid_igram_indices,
    unmasked_idx_set,
    outfile,
    outdset,
    geolist,
    intlist,
    alpha,
    L1 = false;
    prune_outliers = true,
    sigma = 4,
    row_slice = nothing,
    col_slice = nothing,
)
    in_buf = zeros(Float32, (length(row_slice), length(col_slice), length(valid_igram_indices)))
    try
        # in_buf .= h5read(unw_stack_file, in_dset, (col_slice, row_slice, :))[:, :, valid_igram_indices]
        in_buf .= h5read(unw_stack_file, in_dset, (row_slice, col_slice, :))[:, :, valid_igram_indices]
    catch e
        println(e, " ", "fail ", row_slice, " ", col_slice)
        return
    end
    out_buf = zeros(Float32, (length(row_slice), length(col_slice), length(geolist)))

    for (i,row) in enumerate(row_slice), (j,col) in enumerate(col_slice)
        !((row, col) in unmasked_idx_set) && continue

        unw_pixel_raw = in_buf[i, j, :]
        if any(isnan.(unw_pixel_raw))
            println("$row, $col: nan")
            return
        end
        linear = false
        soln_velos, igram_count, geo_clean, unw_clean = calc_soln(
            unw_pixel_raw,
            geolist,
            intlist,
            alpha,
            linear;
            L1 = L1,
            prune_outliers = prune_outliers,
            sigma = sigma,
        )
        timeseries_cm = _unreg_to_cm(soln_velos, geo_clean, geolist)
        out_buf[i, j, :] = timeseries_cm
    end

    dist_outfile = string(Distributed.myid()) * outfile
    h5open(dist_outfile, "r+") do f
        f[outdset][row_slice, col_slice, :] = out_buf
        # f[outdset][row, col, :] = phi_arr
        # f[_count_dset(outdset)][row, col] = igram_count
        # f[_excl_dset(outdset)][row, col] = encode_missing(geolist, geo_clean)
        # f[_stddev_dset(outdset)][row, col] = std(unw_clean) .* abs(PHASE_TO_CM)
        # f[_stddev_raw_dset(outdset)][row, col] = std(unw_pixel_raw) .* abs(PHASE_TO_CM)
    end
    return
end

function _unreg_to_cm(soln_velos, geolist_short, geolist_full)
    # First find the phase for only the clean dates we solved for
    timediffs = day_diffs(geolist_short)
    phi_clean = integrate_velocities(soln_velos, timediffs)

    # Then linearly interpolate between these for the removed phases
    return PHASE_TO_CM * interpolate_phase(geolist_short, phi_clean, geolist_full)
end

function calc_soln(
    unw_pixel,
    geolist,
    intlist,
    alpha,
    constant_velocity;
    L1 = true,
    cor_pixel = nothing,
    cor_thresh = 0.0,
    prune_outliers = true,
    sigma = 4,
)::Tuple{Array{Float32,1},Int64,Array,Array}
    geo_clean, intlist_clean, unw_clean = geolist, intlist, unw_pixel
    if prune_outliers
        geo_clean, intlist_clean, unw_clean =
            remove_outliers(geo_clean, intlist_clean, unw_clean, mean_sigma_cutoff = sigma)
    end
    if cor_thresh > 0 && !isnothing(cor_pixel)
        intlist_clean, unw_clean = prune_cor(
            intlist_clean,
            unw_clean,
            cor_pixel = cor_pixel,
            cor_thresh = cor_thresh,
        )
    end

    igram_count = length(unw_clean)

    B = prepB(geo_clean, intlist_clean, constant_velocity, alpha)
    # Last, pad with zeros if doing Tikh. regularization
    unw_final = alpha > 0 ? augment_zeros(B, unw_clean) : unw_clean

    if L1
        soln = invert_pixel_l1(unw_final, B)
    elseif constant_velocity
        # TODO: Do i ever want the backslash for this, which weights dt like 1/dt^2 ?
        soln = [sum(unw_final) / sum(B)]
    else
        soln = B \ unw_final
    end
    soln_phase = Float32.(soln)
    # end
    return soln_phase, igram_count, geo_clean, unw_clean
end

"""Run sbas on multiple processes"""
function run_sbas(
    unw_stack_file::String,
    dset::String,
    outfile::String,
    outdset::String,
    geolist,
    intlist,
    valid_igram_indices,
    constant_velocity::Bool,
    alpha::Real,
    L1::Bool = false,
    prune_outliers = true,
    sigma = 4,
)

    L1 ? println("Using L1 penalty for fitting") :
    println("Using least squares for fitting")
    prune_outliers ? println("Pruning .geo dates for outliers") :
    println("Not pruning outliers.")
    alpha > 0 ? println("Regularizing solution with alpha = $alpha") :
    println("No regularization")

    # TODO: do I ever really care about these to change as variable?
    # rho, alpha, abstol = 1.0, 1.6, 1e-3
    nrows, ncols, _ = size(unw_stack_file, dset)

    # Choose the correct processing function and size based on if we choose constant
    # TODO: is using multiple dispatch here better in any way (clarity/speed)?
    if constant_velocity
        println("Using constant velocity for inversion solution")
        proc_func = proc_pixel_linear
        outsize = (nrows, ncols)
        chunksize = (500, 500)
    else
        println("Not using linear: Finding solution for each day")
        proc_func = proc_pixel_daily
        outsize = (nrows, ncols, length(geolist))
        chunksize = (10, 10, length(geolist))
    end

    println("Making new writing file for each of $(length(Distributed.workers()))")
    @time @sync @distributed for id in Distributed.workers()
        h5open(string(id) * outfile, "w") do fout
            d_create(fout, outdset, datatype(Float32), dataspace(outsize)) #, "chunk", chunksize)
            d_create(fout, _count_dset(outdset), datatype(Int32), dataspace(nrows, ncols))
            d_create(fout, _excl_dset(outdset), datatype(Int64), dataspace(nrows, ncols))
            # d_create(fout, _stddev_dset(outdset), datatype(Float64), dataspace(nrows, ncols))
            # d_create(fout, _stddev_raw_dset(outdset), datatype(Float64), dataspace(nrows, ncols))
        end
    end

    # Get blocks to load and process
    unmasked_idx_set = Set(get_unmasked_idxs(geolist))  # Set of tuples
    # windows returns (row_slice, col_slice)
    w = collect(windows(unw_stack_file, dset))
    # print(w[1:10])

    @time @sync @distributed for (row_slice, col_slice) in w
    # @time @sync @distributed for (row_slice, col_slice) in w[1:3]
        # @time @sync @distributed for (row, col) in collect(Iterators.product(300:400, 300:400))
        proc_func(
            unw_stack_file,
            dset,
            valid_igram_indices,
            unmasked_idx_set,
            outfile,
            outdset,
            geolist,
            intlist,
            alpha,
            L1;
            prune_outliers = prune_outliers,
            sigma = sigma,
            row_slice = row_slice,
            col_slice = col_slice,
        )
    end
    println("Merging files into $outfile")
    @time merge_partial_files(outfile, outdset, _count_dset(outdset), _excl_dset(outdset))
    # @time merge_partial_files(outfile, outdset)
    # _stddev_dset(outdset), _stddev_raw_dset(outdset))

    # _save_std_attrs(outfile, outdset)
    return outfile, outdset
end

# Need the .I so we can use to load from h5
get_unmasked_idxs(geolist, do_permute = false) = [
    cart_idx.I for cart_idx in findall(.!Sario.load_mask(geolist, do_permute = do_permute))
]

function _save_std_attrs(outfile, outdset)
    attr_dict1 = Dict(
        "description" => "original unwrapped (shifted) std. dev of pixels",
        "units" => "cm",
    )
    h5writeattr(outfile, _stddev_raw_dset(outdset), attr_dict1)
    attr_dict2 =
        Dict("description" => "Final, outlier-removed std. dev of pixels", "units" => "cm")
    h5writeattr(outfile, _stddev_dset(outdset), attr_dict2)
end

"""Linearly interpolate the values between dates that we didn't solve for"""
function interpolate_phase(geo_clean, phis, geolist_full)
    return interp1d(_get_day_nums(geo_clean), phis, _get_day_nums(geolist_full))
end

function interp1d(x, y, xv, k = 1)
    sp = Dierckx.Spline1D(x, y, k = k)
    return sp(xv)
end

function prepB(geolist, intlist, constant_velocity = false, alpha = 0)
    if constant_velocity
        return Float32.(reshape(temporal_baseline(intlist), :, 1))
    end
    # Prepare A and B matrix used for each pixel inversion
    # A = build_A_matrix(geolist, intlist)
    B = build_B_matrix(geolist, intlist)

    if alpha > 0
        B = augment_B(B, alpha)
    end
    return B
end

# In case we make each pixel inversion more complicated (extra penalty functions, etc.)
# TODO: does a Val type for the case with no extra zeros help here?
function invert_pixel(pixel::Array{Float32,1}, pB::Array{Float32,2}, extra_zeros = 0)
    if extra_zeros == 0
        return pB * pixel
    else
        return pB * vcat(pixel, zeros(extra_zeros))
    end
end

# Defined in optimize.jl
function invert_pixel_l1(
    pixel::AbstractArray{<:AbstractFloat},
    B::AbstractArray{<:AbstractFloat};
    rho = 1.0,
    alpha = 1.6,
    lu_tuple = nothing,
    abstol = 1e-3,
)
    return Float32.(huber_fit(
        Float64.(B),
        Float64.(pixel),
        rho,
        alpha,
        lu_tuple = lu_tuple,
        abstol = abstol,
    ))
end

"""Takes the list of igram dates and builds the SBAS B (velocity coeff) matrix

Each row corresponds to an igram, each column to a .geo
Values will be t_k+1 - t_k for columns after the -1 in A,
up to and including the +1 entry
"""
function build_B_matrix(geolist::Array{Date,1}, intlist::Array{Igram,1})

    timediffs = day_diffs(geolist)
    # We take the first .geo to be time 0, leave out of matrix
    # Match on date (not time) to find indices
    geolist2 = @view geolist[2:end]

    M = length(intlist)  # Number of igrams,
    N = length(geolist2)

    B = zeros(Float32, (M, N))

    for i = 1:M
        # row = A[i, :]
        early_date, late_date = intlist[i]
        idx_early = findfirst(isequal(early_date), geolist2)
        if isnothing(idx_early)  # The first SLC will not be in the matrix
            start_idx = 1
        else
            start_idx = idx_early + 1
        end
        idx_late = findfirst(isequal(late_date), geolist2)
        end_idx = idx_late

        # Now fill in the time diffs in the index range
        B[i, start_idx:end_idx] .= @view timediffs[start_idx:end_idx]
    end

    return B
end

"""
    build_A_matrix(geolist::Array{Date, 1}, intlist::Array{Igram, 1}) -> Array{Int8, 2}

Takes the list of igram dates and builds the SBAS A matrix

Returns the incident-like matrix from the SBAS paper: A*phi = dphi
    Each row corresponds to an igram, each column to a .geo
    value will be -1 on the early (slave) igrams, +1 on later (master)
"""
function build_A_matrix(geolist::Array{Date,1}, intlist::Array{Igram,1})
    # We take the first .geo to be time 0, leave out of matrix
    # Match on date (not time) to find indices
    geolist = geolist[2:end]

    M = length(intlist)  # Number of igrams,
    N = length(geolist)
    A = zeros(Int8, M, N)
    for i = 1:M
        early_date, late_date = intlist[i]
        idx_early = findfirst(isequal(early_date), geolist)
        if ~isnothing(idx_early)  # The first SLC will not be in the matrix
            A[i, idx_early] = -1
        end
        idx_late = findfirst(isequal(late_date), geolist)
        A[i, idx_late] = 1

    end
    return A
end

_D(n) = spdiagm(0 => -ones(n - 1), 1 => ones(n - 1))[1:end-1, :]

"""For Tikhonov regularization, pad the B matrix with alpha*D"""
function augment_B(B::Array{Float32,2}, alpha::Real)
    extra = alpha * I
    n = size(B, 2)
    return Float64.(vcat(B, alpha * _D(n)))  # Note: as of 9/29/19, qr has bug for sparse float32 matrices
end

augment_zeros(B, unw) = vcat(unw, zeros(size(B, 1) - length(unw)))

# function augment_B(B::Array{Float32, 2}, alpha::Float32)
#     extra = alpha*I
#     return Float32.(vcat(B, ))
# end

# TODO: this can't work for huge stacks
function augment_matrices(B::Array{Float32,2}, unw_stack::Array{Float32,3}, alpha::Real)
    B = Float32.(vcat(B, alpha * I))
    # Now make num rows match
    nrows, ncols, nlayers = size(unw_stack)
    zeros_shape = (nrows, ncols, size(B, 1) - nlayers)
    unw_stack = Float32.(cat(unw_stack, zeros(zeros_shape), dims = 3))
    return B, unw_stack
end

partial_files(outfile) = Glob.glob("[0-9]*" * outfile)

function merge_partial_files(outfile, dsets...)
    f1 = partial_files(outfile)[1]
    for dset in dsets
        tmp = zeros(eltype(f1, dset), size(f1, dset))

        for f in partial_files(outfile)
            cur = h5read(f, dset)
            tmp += cur
        end
        h5open(outfile, "cw") do fout
            fout[dset] = tmp
        end
    end
    for f in partial_files(outfile)
        rm(f)
    end
end

"""Finds the number of days between successive .geo files"""
function day_diffs(geolist::Array{Date,1})
    [difference.value for difference in diff(geolist)]
end

"""Takes velocity solution output and finds phases

Arguments:
    velocities come from invert_sbas
    timediffs are the days between each SAR acquisitions
        length will be 1 less than num SAR acquisitions
"""
function integrate_velocities(
    velocities::AbstractArray{<:AbstractFloat,1},
    timediffs::Array{Int,1},
)
    phi_diffs = velocities .* timediffs
    phi_arr = cumsum(coalesce.(phi_diffs, 0))
    pushfirst!(phi_arr, 0)
    return phi_arr
end

### Outlier functions ######
# 3 Reasons for removing igrams:
# 1. remove all from .geo dates with huge averages (whole day is garbage)
# 2. remove very low correlation
# 3. remove longer igrams w/ longer baseline than is possible
#   e.g. if we have 2.5 cm/year velocity, anything longer than ~1 year
#   will be all noise- we can't reliably sense such quickly moving ground
#
# TODO: figure out which should go in separate file for cleanliness
function remove_outliers(geolist, intlist, unw_pixel; mean_sigma_cutoff = 4)
    # TODO: verify this third criteria?

    # @show mean_sigma_cutoff
    # @show size(intlist), size(geolist), std(unw_pixel)
    # 1. find outliers in this pixels' values and remove them
    geo_clean, intlist_clean, unw_clean =
        peel_nsigma(geolist, intlist, unw_pixel, nsigma = mean_sigma_cutoff)
    # @show size(intlist_clean), size(geo_clean), std(unw_clean)
    return geo_clean, intlist_clean, unw_clean
end

# 2. low correlation cleaning
function prune_cor(intlist, unw_pixel, cor_pixel, cor_thresh = 0.0)
    low_cor_igrams = intlist[cor_pixel.<cor_thresh]
    intlist_clean, unw_clean = remove_igrams(intlist, unw_pixel, low_cor_igrams)
    return intlist_clean, unw_clean
end

function remove_igrams(
    intlist,
    unw_vals,
    bad_items::Union{Date,AbstractArray{Date},AbstractArray{Igram}}...,
)
    good_idxs = trues(size(intlist))
    for item in bad_items
        good_idxs .&= _good_idxs(item, intlist)
    end
    return remove_by_idx(good_idxs, intlist, unw_vals)
end

"""Pass directly the indexes of good ones igrams (removes all BUT the good_idxs)"""
function remove_by_idx(good_idxs::AbstractArray, intlist, unw_vals)
    unw_clean = unw_vals[good_idxs]
    intlist_clean = intlist[good_idxs]
    return intlist_clean, unw_clean
end

# If we pass the B, also remove rows for bad idxs
function remove_by_idx(good_idxs::AbstractArray, intlist, unw_vals, B)
    B_clean = B[good_idxs, :]
    unw_clean = unw_vals[good_idxs]
    intlist_clean = intlist[good_idxs]
    return intlist_clean, unw_clean, B_clean
end

function _good_idxs(bad_date_arr::AbstractArray{Date}, intlist)
    return .!reduce(
        (x, y) -> x .| y,
        (b in intlist for b in bad_date_arr),
        init = falses(size(intlist)),
    )
end

function _good_idxs(bad_igram_arr::AbstractArray{Igram}, intlist)
    return [!(ig in bad_igram_arr) for ig in intlist]
end

_good_idxs(bad_date::Date, intlist) = .!(bad_date in intlist)

""" 4* the sigma valud as calculated using MAD
Used for robust variance est. (as a cutoff for outliers)
See https://en.wikipedia.org/wiki/Robust_measures_of_scale#IQR_and_MAD
Normalize makes it equal to stddev for normally dist. data"""
mednsigma(arr, n = 4) = n * mad(arr, normalize = true)

modzscore(arr) = (arr .- median(arr)) ./ mad(arr, normalize = true)

# min_spread added in case `mad` produces a very low number,
# so if you want to avoid removing lots, set based on your data
function two_way_cutoff(arr, nsigma, min_spread = 0)
    # min_spread = abs(2 * median(arr))
    # spread = 1.5 * iqr(arr)
    spread = max(min_spread, mednsigma(arr, nsigma))
    #@show spread
    # spread = max(min_spread, std(arr)*nsigma)
    #@show std(arr)*nsigma

    # Make we dont cut off very low var points
    low = min(0, median(arr) - spread)
    high = median(arr) + spread
    return (low, high)
end

function two_way_outliers(arr, nsigma, min_spread = 0)
    # return abs.(modzscore(arr)) .> nsigma
    low, high = two_way_cutoff(arr, nsigma, min_spread)
    return (arr .< low) .| (arr .> high)
end

"""Simpler means by day: just abs. value of phases, either mean or median"""
mean_abs_val(geolist, intlist, unw_vals) =
    [mean(abs.(vals_by_date(d, intlist, unw_vals))) for d in geolist];
median_abs_val(geolist, intlist, unw_vals) =
    [median(abs.(vals_by_date(d, intlist, unw_vals))) for d in geolist];
# valsq(geolist, intlist, unw_vals) = [mean(vals_by_date(d, intlist, unw_vals) .^ 2) for d in geolist];

"""Get the values (unw, cc, etc.) of one date"""
vals_by_date(date::Date, intlist, vals) = vals[date in intlist]

"""Get the values (unw, cc, etc.) grouped by each date"""
vals_by_date(date_arr::Array{Date}, intlist, vals) = [vals[d in intlist] for d in date_arr]

"""Remove all igrams corresponding to the dates with means falling outside an `nsigma` interval"""
function peel_nsigma(geo, int, val; nsigma = 4)
    dates_to_remove = nsigma_days(geo, int, val, nsigma)
    int2, val2 = remove_igrams(int, val, dates_to_remove)
    geo2 = [g for g in geo if !(g in dates_to_remove)]
    return geo2, int2, val2
end

"""Return the days of `geo` which are more than nsigma away from mean"""
function nsigma_days(geo, int, val, nsigma = 4, min_spread = 2)
    means = mean_abs_val(geo, int, val)
    out_idxs = two_way_outliers(means, nsigma, min_spread)

    # FOR PRINTING ONLY
    # low, high = two_way_cutoff(means, nsigma, min_spread)
    # println("Using cutoff around $(median(means)), spread $(max(min_spread, mednsigma(means, nsigma))) : ($low, $high) ")
    # println("Number removed: $(sum(out_idxs))")
    # @show std(means), mad(means, normalize=true)
    return geo[out_idxs]
end

# TODO: this is gonna fail for more than 63 dates ( log2(typemax(Int64)) = 63)...
function encode_missing(total, included)
    out = Int64(0)
    for (i, t) in enumerate(total)
        if !(t in included)
            out += 2^(i - 1)
        end
    end
    return out
end

function decode_excluded(excl_int::Int64, geos::Array{Date,1})
    excl_bitstring = reverse(bitstring(excl_int)[end-length(geos)+1:end])
    return [geos[i] for (i, g) in enumerate(geos) if excl_bitstring[i] == '1']
end

function decode_excluded(excl_int_map::AbstractArray{Int64,2}, geos::Array{Date,1})
    out = Array{Array{Date,1},2}(undef, size(excl_int_map))
    for idx in eachindex(excl_int_map)
        out[idx] = decode_excluded(excl_int_map[idx], geos)
    end
    return out
end


# An array type for colorized overlays of grayscale images
"""
```
A = Overlay(channels, colors)
A = Overlay(channels, colors, clim)
A = Overlay(channels, colors, mapi)
```

Create an `Overlay` array from grayscale channels.  `channels = (channel1,
channel2, ...)`, `colors` is a vector or tuple of `Color`s, and `clim` is a
vector or tuple of min/max values, e.g., `clim = ((min1,max1),(min2,max2),...)`.
Alternatively, you can supply a list of `MapInfo` objects.

See also: `OverlayImage`.
"""
immutable Overlay{T,N,NC,AT<:Tuple{Vararg{AbstractArray}},MITypes<:Tuple{Vararg{MapInfo}}} <: AbstractArray{RGB{T},N}
    channels::AT   # this holds the grayscale arrays
    colors::NTuple{NC,RGB{T}}
    mapi::MITypes

    function Overlay(channels::Tuple{Vararg{AbstractArray}}, colors, mapi::Tuple{Vararg{MapInfo}})
        length(channels) == NC || error("Number of channels must match number of colors")
        length(mapi) == NC   || error("Number of ConvertInfo objects must match number of colors")
        for i = 2:NC
            size(channels[i]) == size(channels[1]) || error("All arrays must have the same size")
        end
        new(channels, colors, mapi)
    end
end
Overlay(channels::Tuple{Vararg{AbstractArray}}, colors::AbstractVector, mapi::Tuple{Vararg{MapInfo}}) =
    Overlay(channels,tuple(colors...),mapi)
Overlay{NC,T}(channels::Tuple{Vararg{AbstractArray}}, colors::NTuple{NC,RGB{T}}, mapi::Tuple{Vararg{MapInfo}}) =
    Overlay{T,ndims(channels[1]),NC,typeof(channels),typeof(mapi)}(channels,colors,mapi)

function Overlay(channels::Tuple{Vararg{AbstractArray}}, colors,
                 clim = ntuple(i->(zero(eltype(channels[i])), one(eltype(channels[i]))), length(channels)))
    n = length(channels)
    for i = 1:n
        if length(clim[i]) != 2
            error("clim must be a 2-vector")
        end
    end
    mapi = ntuple(i->ScaleMinMax(Float32, channels[i], clim[i][1], clim[i][2]), n)
    Overlay(channels, colors, mapi)
end

# Returns the overlay as an image, if possible
"`OverlayImage` is identical to `Overlay`, except that it returns an Image."
function OverlayImage(channels::Tuple{Vararg{AbstractArray}}, colors::Tuple{Vararg{Colorant}},
            arg = ntuple(i->(zero(eltype(channels[i])), one(eltype(channels[i]))), length(channels)))
    ovr = Overlay(channels, colors, arg)
    local prop
    haveprop = false
    for i = 1:length(channels)
        if isa(channels[i], AbstractImage)
            prop = copy(properties(channels[i]))
            haveprop = true
            break
        end
    end
    if !haveprop
        prop = properties(channels[1])
    end
    haskey(prop, "colorspace") && delete!(prop, "colorspace")
    haskey(prop, "colordim")   && delete!(prop, "colordim")
    haskey(prop, "limits")     && delete!(prop, "limits")
    Image(ovr, prop)
end

for NC = 1:3
    NCm = NC-1
    for K = 1:4
        indexargs = Symbol[symbol(string("i_",d)) for d = 1:K]
        sigargs = Expr[:($a::Integer) for a in indexargs]
        @eval begin
            function getindex{T,N,AT,MITypes}(O::Overlay{T,N,$NC,AT,MITypes}, $(sigargs...))
                # one of them needs a bounds-check
                out = map(O.mapi[$NC], getindex(O.channels[$NC], $(indexargs...))) * O.colors[$NC]
                @inbounds @nexprs $NCm c->begin
                    out += map(O.mapi[c], getindex(O.channels[c], $(indexargs...))) * O.colors[c]
                end
                clamp01(eltype(O), out)
            end
        end
    end
end

setindex!(O::Overlay, val, I::Real...) = error("Overlays are read-only. Convert to Image{RGB} to adjust values.")

#### Other Overlay support functions ####
eltype{T}(o::Overlay{T}) = RGB{T}
length(o::Overlay) = isempty(o.channels) ? 0 : length(o.channels[1])
size(o::Overlay) = isempty(o.channels) ? (0,) : size(o.channels[1])
size(o::Overlay, d::Integer) = isempty(o.channels) ? 0 : size(o.channels[1],d)
nchannels(o::Overlay) = length(o.channels)

similar(o::Overlay) = Array(eltype(o), size(o))
similar(o::Overlay, ::NTuple{0}) = Array(eltype(o), size(o))
similar{T}(o::Overlay, ::Type{T}) = Array(T, size(o))
similar{T}(o::Overlay, ::Type{T}, sz::Int64) = Array(T, sz)
similar{T}(o::Overlay, ::Type{T}, sz::Int64...) = Array(T, sz)
similar{T}(o::Overlay, ::Type{T}, sz) = Array(T, sz)

showcompact(io::IO, o::Overlay) = print(io, summary(o), " with colors ", o.colors)


@enum TimeCorrectionType begin
    tctUnknown
    tctGPUT  # /< GPS  to UTC using A0, A1
    tctGAUT  # /< GAL  to UTC using A0, A1
    tctSBUT  # /< SBAS to UTC using A0, A1, incl. provider and UTC ID
    tctGLUT  # /< GLO  to UTC using A0 = -TauC , A1 = 0
    tctGPGA  # /< GPS  to GAL using A0 = A0G   , A1 = A1G
    tctGAGP  # /< GAL  to GPS using A0 = A0G   , A1 = A1G
    tctGLGP  # /< GLO  to GPS using A0 = -TauGPS, A1 = 0
    tctQZGP  # /< QZS  to GPS using A0, A1
    tctQZUT  # /< QZS  to UTC using A0, A1
    tctBDUT  # /< BDT  to UTC using A0, A1
    tctBDGP  # /< BDT  to GPS using A0, A1  !! not in RINEX
    tctIRUT  # /< IRN  to UTC using A0, A1
    tctIRGP  # /< IRN  to GPS using A0, A1
end

const str2tct = Dict{String,TimeCorrectionType}(
    "GPUT" => tctGPUT,
    "GAUT" => tctGAUT,
    "SBUT" => tctSBUT,
    "GLUT" => tctGLUT,
    "GPGA" => tctGPGA,
    "GAGP" => tctGAGP,
    "GLGP" => tctGLGP,
    "QZGP" => tctQZGP,
    "QZUT" => tctQZUT,
    "BDUT" => tctBDUT,
    "BDGP" => tctBDGP,
    "IRUT" => tctIRUT,
    "IRGP" => tctIRGP,

)
function str2TimeCorrectionType(str::String,
    str2tct::Dict{String,TimeCorrectionType}=str2tct)
    haskey(str2tct, str) || throw(error("no such TimeSystemType: $str"))
    str2tct[str]
end

const tct2tsts = Dict(
    tctGPUT => (tstGPS, tstUTC),
    tctGAUT => (tstGAL, tstUTC),
    tctSBUT => (tstGPS, tstUTC),
    tctGLUT => (tstGLO, tstUTC),
    tctGPGA => (tstGPS, tstGAL),
    tctGAGP => (tstGAL, tstGPS),
    tctGLGP => (tstGLO, tstGPS),
    tctQZGP => (tstQZS, tstGPS),
    tctQZUT => (tstQZS, tstUTC),
    tctBDUT => (tstBDT, tstUTC),
    tctBDGP => (tstBDT, tstGPS),
    tctIRUT => (tstIRN, tstUTC),
    tctIRGP => (tstIRN, tstGPS),
)
function tct2TimeSystemTypes(tct::TimeCorrectionType,
    tct2tsts::Dict{TimeCorrectionType,Tuple{TimeSystemType,TimeSystemType}}=tct2tsts,
)::Tuple{TimeSystemType,TimeSystemType}
    haskey(tct2tsts, tct) ||
        throw(error("no such TimeCorrectionType: $tct"))
    tct2tsts[tct]
end
@with_kw struct TimeSystemCorrection
    type::TimeCorrectionType
    frTST::TimeSystemType
    toTST::TimeSystemType
    A0::Float64
    A1::Float64
        # /< reference time for polynominal (week,sow) - MUST BE GPS TIME
    refWeek::Int64
    refSOW::Int64
        # /< reference time (yr,mon,day) for RINEX ver 2 GLO
    refYr::Int64
    refMon::Int64
    refDay::Int64
        # /< string 'EGNOS' 'WAAS' or 'MSAS'
    geoProvider::String
        # /< UTC Identifier [0 unknown, 1=UTC(NIST),
        # /<  2=UTC(USNO), 3=UTC(SU), 4=UTC(BIPM),
        # /<  5=UTC(Europe), 6=UTC(CRL)]
    geoUTCid::String

end

function getTimeSystemCorrection(tp::TimeCorrectionType,
    A0::Float64,A1::Float64,
    refWeek::Int64, refSOW::Int64,
    geoProvider::String, geoUTCid::String)
    frTST, toTST = tct2TimeSystemTypes(tp)

    dt = gpsws2datetime(refWeek, refSOW)

    TimeSystemCorrection(tp, frTST, toTST,
                    A0, A1, refWeek, refSOW,
                    year(dt), month(dt), day(dt),
                    geoProvider, geoUTCid,)
    # if in(tp, Set([tctGLGP, tctGLUT, tctBDUT,
    #         tctGPUT, tctGPGA, tctQZGP, tctQZUT, ]))
    # end
    #  if(tc.type == TimeSystemCorrection::GLGP ||
    #     tc.type == TimeSystemCorrection::GLUT ||        // TD ?
    #     tc.type == TimeSystemCorrection::BDUT ||        // TD ?
    #     tc.type == TimeSystemCorrection::GPUT ||
    #     tc.type == TimeSystemCorrection::GPGA ||
    #     tc.type == TimeSystemCorrection::QZGP ||
    #     tc.type == TimeSystemCorrection::QZUT)

    #  {
    #     GPSWeekSecond gws(tc.refWeek,tc.refSOW);
    #     CivilTime ct(gws);
    #     tc.refYr = ct.year;
    #     tc.refMon = ct.month;
    #     tc.refDay = ct.day;
    #  }

    #  if(tc.type == TimeSystemCorrection::GAUT)
    #  {
    #     GALWeekSecond gws(tc.refWeek,tc.refSOW);
    #     CivilTime ct(gws);
    #     tc.refYr = ct.year;
    #     tc.refMon = ct.month;
    #     tc.refDay = ct.day;
    #  }

    #     //if(tc.type == TimeSystemCorrection::GLUT)
    #     // {
    #     //   tc.refYr =  1980;
    #     //   tc.refMon = 1;
    #     //   tc.refDay = 6;
    #     //   tc.refWeek = 0;
    #     //   tc.refSOW = 0;
    #     //}

end


# NE - Run carbon cases
push!(LOAD_PATH, joinpath(ENV["REPO_PATH"], "Sesi.jl/src"))

using CSV
using JuMP
using IterTools
using JSON
using DataFrames

import Sesi # import vs using so I can reload this from the REPL

reload("Sesi")

include("jl_utils.jl")

save = false
verb = true

if verb
    println(repeat("#",80))
    println("Starting carbon analysis")
    println(repeat("#",80))
end

folderIn = joinpath(ENV["REPO_PATH"], "dataR1/jl_in/")

results = Dict()
prices = Dict()
costs = Dict()
reductions = Dict()
#fileNms = ["jlInput_x0.csv","jlInput_2018.csv", "jlInput_x2.csv"]
fileNms = ["jlInput_x0.csv", "jlInput_x1.csv", "jlInput_x2.csv",
            "jlInput_x4.csv","jlInput_x6.csv", "jlInput_2018.csv"]

for dc_ratio in [1., 0.5, 0.1]
if verb
println("Dealing with demand ratio " * string(dc_ratio))
end

folderOut = joinpath(ENV["REPO_PATH"],
                     "dataR1/jl_out/demand_" * string(dc_ratio))

for sc in fileNms
    prices[sc] = []
    costs[sc] = []
    reductions[sc] = []

    println("Dealing with " * sc)
    results[sc] = Dict()
    fileNm = folderIn * sc

    carbonPrice_lst = [0., 50., 75., 100., 150., 200., 300., 500., 1000.,
        10000., 100000.]
    for carbonPrice in carbonPrice_lst
        sesi, p, f = build_model(fileNm; verb=false, startDate=Date(2016,1,1),
            endDate=Date(2016,12,31), carbonPrice=carbonPrice, dc_ratio=dc_ratio)
        status = solve(sesi)
        println(status)
        results[sc][carbonPrice] = Sesi.campus_bill(sesi, f, p, verb=0)
        if verb
            @printf("carbon price: %d - %.2f\n",carbonPrice,
                carbon_cost(results, sc, 0., carbonPrice))
            @printf("carbon reduction: %d - %.2f\n",carbonPrice,
                carbon_reduction(results, sc, 0., carbonPrice))
            @printf("sanity : %d - %.2f\n",carbonPrice,
                sanity(results, sc, 0., carbonPrice))
        end
        push!(prices[sc], carbonPrice)
        push!(costs[sc], carbon_cost(results, sc, 0., carbonPrice))
        push!(reductions[sc], carbon_reduction(results, sc, 0., carbonPrice))
        if save
            Sesi.save_data(sesi, f, p, folderOut * sc[1:end-4] * "_"
                * string(carbonPrice))
        end
    end
end

json_data = json([results, prices, costs, reductions]);
open(folderOut * "ne_carbon_json_data", "w") do fw
    write(fw, json_data)
end
end


## HSP family of heuristics ##
export HSPHeuristic, HAdd, HMax
export HSPRHeuristic, HAddR, HMaxR

"Precomputed domain information for HSP heuristic."
struct HSPCache
    domain::Domain # Preprocessed domain
    axioms::Vector{Clause} # Preprocessed axioms
    preconds::Dict{Symbol,Vector{Vector{Term}}} # Preconditions in DNF
    additions::Dict{Symbol,Vector{Term}} # Action add lists
end

"HSP family of delete-relaxation heuristics."
mutable struct HSPHeuristic <: Heuristic
    op::Function # Aggregator (e.g. maximum, sum) for fact costs
    cache::HSPCache # Precomputed domain information
    pre_key::UInt64 # Key to check if information needs to be precomputed again
    HSPHeuristic(op) = new(op)
end

Base.hash(heuristic::HSPHeuristic, h::UInt) =
    hash(heuristic.op, hash(HSPHeuristic, h))

function precompute!(h::HSPHeuristic,
                     domain::Domain, state::State, spec::Specification)
    # Check if cache has already been computed
    if is_precomputed(h, domain, state, spec) return h end
    h.pre_key = objectid(domain) # Precomputed data is unique to each domain
    domain = copy(domain) # Make a local copy of the domain
    # Preprocess axioms
    axioms = regularize_clauses(domain.axioms) # Regularize domain axioms
    axioms = [Clause(ax.head, [t for t in ax.body if t.name != :not])
              for ax in axioms] # Remove negative literals
    domain.axioms = Clause[] # Remove axioms so they do not affect execution
    # Preprocess actions
    preconds = Dict{Symbol,Vector{Vector{Term}}}()
    additions = Dict{Symbol,Vector{Term}}()
    for (act_name, act_def) in domain.actions
        # Convert preconditions to DNF without negated literals
        conds = get_preconditions(act_def; converter=to_dnf)
        conds = [c.args for c in conds.args]
        for c in conds filter!(t -> t.name != :not, c) end
        preconds[act_name] = conds
        # Extract additions from each effect
        diff = effect_diff(act_def.effect)
        additions[act_name] = diff.add
    end
    h.cache = HSPCache(domain, axioms, preconds, additions)
    return h
end

function is_precomputed(h::HSPHeuristic,
                        domain::Domain, state::State, spec::Specification)
    return (isdefined(h, :cache) && objectid(domain) == h.pre_key)
end

function compute(h::HSPHeuristic,
                 domain::Domain, state::State, spec::Specification)
    # Precompute if necessary
    if !is_precomputed(h, domain, state, spec)
        precompute!(h, domain, state, spec) end
    @unpack op, cache = h
    @unpack domain = cache
    @unpack types, facts = state
    goals = get_goal_terms(spec)
    # Initialize fact costs in a GraphPlan-style graph
    fact_costs = Dict{Term,Float64}(f => 0 for f in facts)
    while true
        facts = Set(keys(fact_costs))
        state = State(types, facts, Dict{Symbol,Any}())
        if is_goal(spec, domain, state)
            return op([0; [fact_costs[g] for g in goals]]) end
        # Compute costs of one-step derivations of domain axioms
        for ax in cache.axioms
            _, subst = resolve(ax.body, [Clause(f, []) for f in facts])
            for s in subst
                body = [substitute(t, s) for t in ax.body]
                cost = op([0; [get(fact_costs, f, 0) for f in body]])
                derived = substitute(ax.head, s)
                if cost < get(fact_costs, derived, Inf)
                    fact_costs[derived] = cost end
            end
        end
        # Compute costs of all effects of available actions
        actions = available(state, domain)
        for act in actions
            act_args = domain.actions[act.name].args
            subst = Subst(var => val for (var, val) in zip(act_args, act.args))
            # Look-up preconds and substitute vars
            preconds = cache.preconds[act.name]
            preconds = [[substitute(t, subst) for t in c] for c in preconds]
            # Compute cost of reaching each action
            cost = minimum([[op([0; [get(fact_costs, f, 0) for f in conj]])
                             for conj in preconds]; Inf])
            # Compute cost of reaching each added fact
            additions = [substitute(a, subst) for a in cache.additions[act.name]]
            cost = cost + 1 # TODO: Handle arbitrary action costs
            for fact in additions
                if cost < get(fact_costs, fact, Inf)
                    fact_costs[fact] = cost end
            end
        end
        # Terminate if there's no change to the number of facts
        if length(fact_costs) == length(facts) && keys(fact_costs) == facts
            return Inf end
    end
end

"HSP heuristic where a fact's cost is the maximum cost of its dependencies."
HMax(args...) = HSPHeuristic(maximum, args...)

"HSP heuristic where a fact's cost is the summed cost of its dependencies."
HAdd(args...) = HSPHeuristic(sum, args...)

"HSPr family of delete-relaxation heuristics for regression search."
mutable struct HSPRHeuristic <: Heuristic
    op::Function
    fact_costs::Dict{Term,Float64} # Est. cost of reaching each fact from goal
    pre_key::NTuple{3,UInt64} # Key to check if precomputation needs to re-run
    HSPRHeuristic(op) = new(op)
end

Base.hash(heuristic::HSPRHeuristic, h::UInt) =
    hash(heuristic.op, hash(HSPRHeuristic, h))

function precompute!(h::HSPRHeuristic,
                     domain::Domain, state::State, spec::Specification)
    # Check if data has already been computed
    if is_precomputed(h, domain, state, spec) return h end
    # Precomputed data is tied to all three inputs
    h.pre_key = (objectid(domain), hash(state), hash(spec))
    @unpack op = h
    @unpack types, facts = state
    goals = get_goal_terms(spec)
    # Preprocess domain and axioms
    domain = copy(domain)
    axioms = regularize_clauses(domain.axioms)
    axioms = [Clause(ax.head, [t for t in ax.body if t.name != :not])
              for ax in axioms]
    domain.axioms = Clause[]
    # Preprocess actions
    preconds = Dict{Symbol,Vector{Vector{Term}}}()
    additions = Dict{Symbol,Vector{Term}}()
    for (act_name, act_def) in domain.actions
        # Convert preconditions to DNF without negated literals
        conds = get_preconditions(act_def; converter=to_dnf)
        conds = [c.args for c in conds.args]
        for c in conds filter!(t -> t.name != :not, c) end
        preconds[act_name] = conds
        # Extract additions from each effect
        diff = effect_diff(act_def.effect)
        additions[act_name] = diff.add
    end
    # Initialize fact costs in a GraphPlan-style graph
    fact_costs = Dict{Term,Float64}(f => 0 for f in PDDL.get_facts(state))
    while true
        facts = Set(keys(fact_costs))
        state = State(types, facts, Dict{Symbol,Any}())
        # Compute costs of one-step derivations of domain axioms
        for ax in axioms
            _, subst = resolve(ax.body, [Clause(f, []) for f in facts])
            for s in subst
                body = [substitute(t, s) for t in ax.body]
                cost = op([0; [get(fact_costs, f, 0) for f in body]])
                derived = substitute(ax.head, s)
                if cost < get(fact_costs, derived, Inf)
                    fact_costs[derived] = cost end
            end
        end
        # Compute costs of all effects of available actions
        actions = available(state, domain)
        for act in actions
            act_args = domain.actions[act.name].args
            subst = Subst(var => val for (var, val) in zip(act_args, act.args))
            # Look-up preconds and substitute vars
            conds = preconds[act.name]
            conds = [[substitute(t, subst) for t in c] for c in conds]
            # Compute cost of reaching each action
            cost = minimum([[op([0; [get(fact_costs, f, 0) for f in conj]])
                             for conj in conds]; Inf])
            # Compute cost of reaching each added fact
            added = [substitute(a, subst) for a in additions[act.name]]
            cost = cost + 1 # TODO: Handle arbitrary action costs
            for fact in added
                if cost < get(fact_costs, fact, Inf)
                    fact_costs[fact] = cost end
            end
        end
        # Terminate when there's no change to the number of facts
        if length(fact_costs) == length(facts) && keys(fact_costs) == facts
            break end
    end
    h.fact_costs = fact_costs
    return h
end

function is_precomputed(h::HSPRHeuristic,
                        domain::Domain, state::State, spec::Specification)
    return (isdefined(h, :fact_costs) && objectid(domain) == h.pre_key[1]
            && hash(state) == h.pre_key[2] && hash(spec) == h.pre_key[3])
end

function compute(h::HSPRHeuristic,
                 domain::Domain, state::State, spec::Specification)
    # Compute cost of achieving all facts in current state
    return h.op([0; [get(h.fact_costs, f, 0) for f in PDDL.get_facts(state)]])
end

"HSPr heuristic where a fact's cost is the maximum cost of its dependencies."
HMaxR(args...) = HSPRHeuristic(maximum, args...)

"HSPr heuristic where a fact's cost is the summed cost of its dependencies."
HAddR(args...) = HSPRHeuristic(sum, args...)


__precompile__(false)

module SCRIPT

using DataStructures #OrderedDict

using GLFW, ModernGL, SharedArrays, StaticArrays, Quaternions

using RessourceManager
#using WindowManager
using GraphicsManager
using DefaultModelData
using CameraManager
using FrustumManager
using ChunkManager
using MeshManager
using TextureManager
using ShaderManager
#using ..ScriptManager
using TimeManager
using LogManager
using MathManager

const GPU = GraphicsManager

mutable struct Script
  args::Dict{Symbol,Any}
  vars::Dict{Symbol,Any}
  shaderManager::JShaderManager
  Script() = new(Dict(),Dict(),JShaderManager())
end

this = Script()

function reset()
  this.vars = typeof(this.vars)()
  delPrograms(this.shaderManager)
end

""" TODO """
function main(args::Dict{Symbol,Any})
  this.args = args
  println("Script: $(basename(@__FILE__)), $(this.args), time: $(mtime(@__FILE__))")
  reset()
end

""" TODO """
function OnDestroy()
  println("OnDestroy")
  reset()
  GraphicsManager.cleanUp() #remove all
  GC.gc()
end

function resizeWindow(window, width, height)
  this.vars[:RESOLUTION] = (
    width = width,
    height = height,
    ratio = width / (height*1f0),
    size = width * height
  )

  GLFW.SetWindowSize(window, width, height)
  glViewport(0, 0, width, height)
end

""" TODO """
function OnInit()
  SS = this.vars[:SECTION_SHADER] = Dict{Symbol,Any}()
  SM = this.vars[:SECTION_MESH] = Dict{Symbol,Any}()
  ST = this.vars[:SECTION_TEXTURES] = Dict{Symbol,Any}()

  UF = SS[:UNIFORMS] = OrderedDict{Symbol,Any}()
  #UF = SS[:UNIFORMS] = OrderedDict{Symbol,Any}()
  #TS = ST[:TEXTURES] = OrderedDict{Symbol,Any}()

  resizeWindow(this.args[:WINDOW], 800, 600)

  glEnable(GL_DEPTH_TEST)
  glEnable(GL_BLEND)
  glEnable(GL_CULL_FACE)
  #glEnable(GL_MULTISAMPLE)

  glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA)
  #glBlendEquation(GL_FUNC_ADD)
  #glFrontFace(GL_CCW)

  glCullFace(GL_BACK) #default: GL_BACK
  glDepthMask(GL_TRUE) #default: GL_TRUE
  glDepthFunc(GL_LEQUAL) #default: GL_LESS

  #glPolygonMode(GL_FRONT_AND_BACK, GL_LINE)#GL_FILL,GL_LINE
  glClearColor(0.5, 0.5, 0.5, 0.0)

  resolution = values(this.vars[:RESOLUTION])[1:2]

  UF[:iResolution] = Uniform(value=Float32[resolution...,1])
  UF[:iTime] = Uniform(value=Float32(0),default=0)
  UF[:iFrame] = Uniform(value=Float32(0),default=0)
  UF[:iTimeDelta] = Uniform(value=Float32(0),default=0)
  UF[:iFrameRate] = Uniform(value=Float32(0),default=0)
  UF[:iChannelTime] = Uniform(value=zeros(Float32,4))
  UF[:iMouse] = Uniform(value=zeros(Float32,4))
  UF[:iDate] = Uniform(value=zeros(Float32,4))
  UF[:iSampleRate] = Uniform(value=Float32(0))
  UF[:iChannelResolution] = Uniform(value=[zeros(Float32,3),zeros(Float32,3),zeros(Float32,3),zeros(Float32,3)])
  UF[:iChannel0] = Uniform(value=Int32(0),typ=GL_SAMPLER_2D,layout=Dict(:binding=>0))
  UF[:iChannel1] = Uniform(value=Int32(1),typ=GL_SAMPLER_2D,layout=Dict(:binding=>1))
  UF[:iChannel2] = Uniform(value=Int32(2),typ=GL_SAMPLER_2D,layout=Dict(:binding=>2))
  UF[:iChannel3] = Uniform(value=Int32(3),typ=GL_SAMPLER_2D,layout=Dict(:binding=>3))

  this.vars[:renderTime] = this.vars[:startTime] = time()

  SM[:SCREEN] = MeshData()
  linkData(SM[:SCREEN], :vertices=>(DATA_PLANE2D_VERTEX_STRIP,2))

  size = (UInt32(resolution[1]), UInt32(resolution[2]))

  txp(path::String) = RessourceManager.getPath(:TEXTURES, path)

  global iChannel0=createTexture(:iChannel0, size) #txp("A.png"))
  global iChannel1=createTexture(:iChannel1, size) #txp("B.png"))
  global iChannel2=createTexture(:iChannel2, size) #txp("C.png"))
  global iChannel3=createTexture(:iChannel3, size) #txp("D.png"))

  glBindImageTexture(0, iChannel0, 0, GL_FALSE, 0, GL_WRITE_ONLY, GL_RGBA32F)
  glBindImageTexture(1, iChannel1, 0, GL_FALSE, 0, GL_WRITE_ONLY, GL_RGBA32F)
  glBindImageTexture(2, iChannel2, 0, GL_FALSE, 0, GL_WRITE_ONLY, GL_RGBA32F)
  glBindImageTexture(3, iChannel3, 0, GL_FALSE, 0, GL_WRITE_ONLY, GL_RGBA32F)

  SS[:PROGRAMS] = OrderedDict(
  :SCREEN => (shaders=[:SCREEN_VSH, :MC_FSH],attributes=[:SCREEN],),
  :MC_BA => (shaders=[:MC_BA_CSH],),
  :MC_BB => (shaders=[:MC_BB_CSH],),
  :MC_BC => (shaders=[:MC_BC_CSH],),
  :MC_BD => (shaders=[:MC_BD_CSH],),
  )

  dispatch = (x=round(UInt32, resolution[1]/16f0), y=round(UInt32, resolution[2]/16f0), z=UInt32(1))
  dispatchInFile = (x=16,y=16,z=1)

  SS[:DISPATCH] = [
  (:MC_BA, dispatch),
  (:MC_BB, dispatch),
  (:MC_BC, dispatch),
  (:MC_BD, dispatch),
  ]

  SS[:GLOBALS] = Dict(
  :iDispatchX=>dispatchInFile.x,
  :iDispatchY=>dispatchInFile.y,
  :iDispatchZ=>dispatchInFile.z,
  )

  ST[:TEXTURES] = OrderedDict(
  :iChannel0 => (location=0,),
  :iChannel1 => (location=1,),
  :iChannel2 => (location=2,),
  :iChannel3 => (location=3,),
  )

  reloadShaderPrograms()
end

""" TODO """
function OnReload()
end

""" TODO """
function OnUpdate()
  checkForUpdate()

  UF = this.vars[:SECTION_SHADER][:UNIFORMS]

  deltaTime = Float32(this.vars[:renderTime] - this.vars[:startTime]) #now timeIntervalSinceDate:time()
  if deltaTime > 1f0/20f0
      deltaTime = 1f0/20f0
  end

  UF[:iTimeDelta].value = Float32(deltaTime)
  UF[:iFrameRate].value = Float32(1f0/deltaTime)
end

function setShaderProperties()
  for (k,v) in this.vars[:SECTION_SHADER][:UNIFORMS]
    setShaderProperty(string(k), v.value)
  end
end

function bindTextures()
  for (k,v) in this.vars[:SECTION_TEXTURES][:TEXTURES]
    bindTexture(k, v.location)
  end
end

function dispatchPrograms(list)
  for (p,d) in list
    useShaderProgram(p)
    setShaderProperties()
    glDispatchCompute(d.x, d.y, d.z)
    glMemoryBarrier(GL_SHADER_IMAGE_ACCESS_BARRIER_BIT)
  end
end

""" TODO """
function OnRender()
  SS = this.vars[:SECTION_SHADER]
  SM = this.vars[:SECTION_MESH]
  UF = SS[:UNIFORMS]

  iframe = UF[:iFrame]
  UF[:iTime].value = Float32(programTime())

  bindTextures()

  dispatchPrograms(SS[:DISPATCH])

  mesh = SM[:SCREEN]

  glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT | GL_STENCIL_BUFFER_BIT)
  useShaderProgram(:SCREEN)
  glBindVertexArray(mesh.vao)
  setShaderProperties()
  glDrawArrays(GL_TRIANGLE_STRIP, 0, mesh.draw.count)

  UF[:iFrame].value += 1
  this.vars[:renderTime] = time()
end

""" TODO """
function checkForUpdate()
  keyValue, keyPressed = getKey()

  if keyPressed
    if keyValue == 45 #ß
    elseif keyValue == 48 #0
    elseif keyValue >= 49 && keyValue <= 57 #1-9

    elseif keyValue == 65 #a
    elseif keyValue == 66 #b
    elseif keyValue == 67 #c
    elseif keyValue == 68 #d
    elseif keyValue == 69 #e
    elseif keyValue == 70 #f
    elseif keyValue == 71 #g
    elseif keyValue == 72 #h
    elseif keyValue == 73 #i
    elseif keyValue == 74 #j
    elseif keyValue == 75 #k
    elseif keyValue == 76 #l
    elseif keyValue == 77 #m
    elseif keyValue == 78 #n
    elseif keyValue == 79 #o
    elseif keyValue == 80 #p
    elseif keyValue == 81 #q
    elseif keyValue == 82 #r
      reloadShaderPrograms()
    elseif keyValue == 83 #s
    elseif keyValue == 84 #t
    elseif keyValue == 85 #u
    elseif keyValue == 86 #v
    elseif keyValue == 86 #w
    elseif keyValue == 87 #x
    elseif keyValue == 88 #y
    elseif keyValue == 89 #z

    elseif (keyValue >= 290 && keyValue <= 301) # F1 - F12
    end
  end
end

useShaderProgram(program::Union{Nothing, Symbol}) = ShaderManager.useProgram(this.shaderManager, program)
getShaderProgram(program::Symbol) = ShaderManager.getProgram(this.shaderManager, program).id
setShaderProperty(name::String, value::Any; program::Union{Nothing,Symbol}=getActiveProgram(this.shaderManager)) = ShaderManager.setProperty(this.shaderManager, name::String, value::Any; program=program)
#if mode != "" println("MODE(",stringColor(mode;color=:yellow),"): ",stringColor(value;color=:yellow)) end

function reloadShaderPrograms()
  println("Load shader files...")
  SS = this.vars[:SECTION_SHADER]
  SM = this.vars[:SECTION_MESH]

  shaders = loadShaders(SS)

  for (p,t) in SS[:PROGRAMS]
    ShaderManager.reloadProgram(this.shaderManager, p, (x->shaders[x]).(t.shaders))
    programID = getShaderProgram(p)
    if haskey(t,:attributes)
      for a in t.attributes
        setAttributes(SM[a], programID)
      end
    end
  end

end

end # SCRIPT


"""
    stem(xs, ys, [zs]; kwargs...)

Plots markers at the given positions extending from `offset` along stem lines.

`offset` can be a number, in which case it sets y for 2D, and z for 3D stems.
It can be a Point2 for 2D plots, as well as a Point3 for 3D plots.
It can also be an iterable of any of these at the same length as xs, ys, zs.

The conversion trait of stem is `PointBased`.

## Attributes
$(ATTRIBUTES)
"""
@recipe(Stem) do scene
    Attributes(
        stemcolor = theme(scene, :linecolor),
        stemcolormap = theme(scene, :colormap),
        stemcolorrange = automatic,
        stemwidth = theme(scene, :linewidth),
        stemlinestyle = nothing,
        trunkwidth = theme(scene, :linewidth),
        trunklinestyle = nothing,
        trunkcolor = theme(scene, :linecolor),
        trunkcolormap = theme(scene, :colormap),
        trunkcolorrange = automatic,
        offset = 0,
        marker = :circle,
        markersize = theme(scene, :markersize),
        color = theme(scene, :markercolor),
        colormap = theme(scene, :colormap),
        colorrange = automatic,
        strokecolor = theme(scene, :markerstrokecolor),
        strokewidth = theme(scene, :markerstrokewidth),
        visible = true,
        inspectable = theme(scene, :inspectable),
        cycle = [[:stemcolor, :color, :trunkcolor] => :color],
    )
end

conversion_trait(::Type{<:Stem}) = PointBased()

trunkpoint(stempoint::P, offset::Number) where P <: Point2 = P(stempoint[1], offset)
trunkpoint(stempoint::P, offset::Point2) where P <: Point2 = P(offset...)
trunkpoint(stempoint::P, offset::Number) where P <: Point3 = P(stempoint[1], stempoint[2], offset)
trunkpoint(stempoint::P, offset::Point3) where P <: Point3 = P(offset...)

function plot!(s::Stem{<:Tuple{<:AbstractVector{<:Point}}})
    points = s[1]

    stemtuples = lift(points, s.offset) do ps, to
        tuple.(ps, trunkpoint.(ps, to))
    end

    trunkpoints = @lift(last.($stemtuples))

    lines!(s, trunkpoints,
        linewidth = s.trunkwidth,
        color = s.trunkcolor,
        colormap = s.trunkcolormap,
        colorrange = s.trunkcolorrange,
        visible = s.visible,
        linestyle = s.trunklinestyle,
        inspectable = s.inspectable)
    linesegments!(s, stemtuples,
        linewidth = s.stemwidth,
        color = s.stemcolor,
        colormap = s.stemcolormap,
        colorrange = s.stemcolorrange,
        visible = s.visible,
        linestyle = s.stemlinestyle,
        inspectable = s.inspectable)
    scatter!(s, s[1],
        color = s.color,
        colormap = s.colormap,
        colorrange = s.colorrange,
        markersize = s.markersize,
        marker = s.marker,
        strokecolor = s.strokecolor,
        strokewidth = s.strokewidth,
        visible = s.visible,
        inspectable = s.inspectable)
    s
end


export iauAtco13
"""
ICRS RA,Dec to observed place.  The caller supplies UTC, site
coordinates, ambient air conditions and observing wavelength.

SOFA models are used for the Earth ephemeris, bias-precession-
nutation, Earth orientation and refraction.

This function is part of the International Astronomical Union's
SOFA (Standards of Fundamental Astronomy) software collection.

Status:  support function.

Given:
   rc,dc  double   ICRS right ascension at J2000.0 (radians, Note 1)
   pr     double   RA proper motion (radians/year; Note 2)
   pd     double   Dec proper motion (radians/year)
   px     double   parallax (arcsec)
   rv     double   radial velocity (km/s, +ve if receding)
   utc1   double   UTC as a 2-part...
   utc2   double   ...quasi Julian Date (Notes 3-4)
   dut1   double   UT1-UTC (seconds, Note 5)
   elong  double   longitude (radians, east +ve, Note 6)
   phi    double   latitude (geodetic, radians, Note 6)
   hm     double   height above ellipsoid (m, geodetic, Notes 6,8)
   xp,yp  double   polar motion coordinates (radians, Note 7)
   phpa   double   pressure at the observer (hPa = mB, Note 8)
   tc     double   ambient temperature at the observer (deg C)
   rh     double   relative humidity at the observer (range 0-1)
   wl     double   wavelength (micrometers, Note 9)

Returned:
   aob    double*  observed azimuth (radians: N=0,E=90)
   zob    double*  observed zenith distance (radians)
   hob    double*  observed hour angle (radians)
   dob    double*  observed declination (radians)
   rob    double*  observed right ascension (CIO-based, radians)
   eo     double*  equation of the origins (ERA-GST)

Returned (function value):
         int      status: +1 = dubious year (Note 4)
                           0 = OK
                           -1 = unacceptable date

Notes:

   1.  Star data for an epoch other than J2000.0 (for example from the
   Hipparcos catalog, which has an epoch of J1991.25) will require
   a preliminary call to iauPmsafe before use.

   2.  The proper motion in RA is dRA/dt rather than cos(Dec)*dRA/dt.

   3.  utc1+utc2 is quasi Julian Date (see Note 2), apportioned in any
   convenient way between the two arguments, for example where utc1
   is the Julian Day Number and utc2 is the fraction of a day.

   However, JD cannot unambiguously represent UTC during a leap
   second unless special measures are taken.  The convention in the
   present function is that the JD day represents UTC days whether
   the length is 86399, 86400 or 86401 SI seconds.

   Applications should use the function iauDtf2d to convert from
   calendar date and time of day into 2-part quasi Julian Date, as
   it implements the leap-second-ambiguity convention just
   described.

   4.  The warning status "dubious year" flags UTCs that predate the
   introduction of the time scale or that are too far in the
   future to be trusted.  See iauDat for further details.

   5.  UT1-UTC is tabulated in IERS bulletins.  It increases by exactly
   one second at the end of each positive UTC leap second,
   introduced in order to keep UT1-UTC within +/- 0.9s.  n.b. This
   practice is under review, and in the future UT1-UTC may grow
   essentially without limit.

   6.  The geographical coordinates are with respect to the WGS84
   reference ellipsoid.  TAKE CARE WITH THE LONGITUDE SIGN:  the
   longitude required by the present function is east-positive
   (i.e. right-handed), in accordance with geographical convention.

   7.  The polar motion xp,yp can be obtained from IERS bulletins.  The
   values are the coordinates (in radians) of the Celestial
   Intermediate Pole with respect to the International Terrestrial
   Reference System (see IERS Conventions 2003), measured along the
   meridians 0 and 90 deg west respectively.  For many
   applications, xp and yp can be set to zero.

   8.  If hm, the height above the ellipsoid of the observing station
   in meters, is not known but phpa, the pressure in hPa (=mB),
   is available, an adequate estimate of hm can be obtained from
   the expression

         hm = -29.3 * tsl * log ( phpa / 1013.25 );

   where tsl is the approximate sea-level air temperature in K
   (See Astrophysical Quantities, C.W.Allen, 3rd edition, section
   52).  Similarly, if the pressure phpa is not known, it can be
   estimated from the height of the observing station, hm, as
   follows:

         phpa = 1013.25 * exp ( -hm / ( 29.3 * tsl ) );

   Note, however, that the refraction is nearly proportional to
   the pressure and that an accurate phpa value is important for
   precise work.

   9.  The argument wl specifies the observing wavelength in
   micrometers.  The transition from optical to radio is assumed to
   occur at 100 micrometers (about 3000 GHz).

   10. The accuracy of the result is limited by the corrections for
   refraction, which use a simple A*tan(z) + B*tan^3(z) model.
   Providing the meteorological parameters are known accurately and
   there are no gross local effects, the predicted observed
   coordinates should be within 0.05 arcsec (optical) or 1 arcsec
   (radio) for a zenith distance of less than 70 degrees, better
   than 30 arcsec (optical or radio) at 85 degrees and better
   than 20 arcmin (optical) or 30 arcmin (radio) at the horizon.

   Without refraction, the complementary functions iauAtco13 and
   iauAtoc13 are self-consistent to better than 1 microarcsecond
   all over the celestial sphere.  With refraction included,
   consistency falls off at high zenith distances, but is still
   better than 0.05 arcsec at 85 degrees.

   11. "Observed" Az,ZD means the position that would be seen by a
   perfect geodetically aligned theodolite.  (Zenith distance is
   used rather than altitude in order to reflect the fact that no
   allowance is made for depression of the horizon.)  This is
   related to the observed HA,Dec via the standard rotation, using
   the geodetic latitude (corrected for polar motion), while the
   observed HA and RA are related simply through the Earth rotation
   angle and the site longitude.  "Observed" RA,Dec or HA,Dec thus
   means the position that would be seen by a perfect equatorial
   with its polar axis aligned to the Earth's axis of rotation.

   12. It is advisable to take great care with units, as even unlikely
   values of the input parameters are accepted and processed in
   accordance with the models used.

Called:
   iauApco13    astrometry parameters, ICRS-observed, 2013
   iauAtciq     quick ICRS to CIRS
   iauAtioq     quick CIRS to observed

This revision:   2016 February 2

SOFA release 2018-01-30

Copyright (C) 2018 IAU SOFA Board.  See notes at end.
"""
# int iauAtco13(double rc, double dc,
#               double pr, double pd, double px, double rv,
#               double utc1, double utc2, double dut1,
#               double elong, double phi, double hm, double xp, double yp,
#               double phpa, double tc, double rh, double wl,
#               double *aob, double *zob, double *hob,
#               double *dob, double *rob, double *eo)

function iauAtco13(rc::Real, dc::Real, pr::Real, pd::Real,
                   px::Real, rv::Real, utc1::Real, utc2::Real,
                   dut1::Real, elong::Real, phi::Real, hm::Real,
                   xp::Real, yp::Real, phpa::Real, tc::Real,
                   rh::Real, wl::Real)

   # Allocate return values
   ref_aob = Ref{Float64}(0.0)
   ref_zob = Ref{Float64}(0.0)
   ref_hob = Ref{Float64}(0.0)
   ref_dob = Ref{Float64}(0.0)
   ref_rob = Ref{Float64}(0.0)
   ref_eo  = Ref{Float64}(0.0)

   status = ccall((:iauAtco13, libsofa_c), Cint,
            (Cdouble, Cdouble, Cdouble, Cdouble,
            Cdouble, Cdouble, Cdouble, Cdouble,
            Cdouble, Cdouble, Cdouble, Cdouble,
            Cdouble, Cdouble, Cdouble, Cdouble,
            Cdouble, Cdouble,
            Ref{Cdouble}, Ref{Cdouble}, Ref{Cdouble},
            Ref{Cdouble}, Ref{Cdouble}, Ref{Cdouble}),
            convert(Float64, rc), convert(Float64, dc),
            convert(Float64, pr), convert(Float64, pd),
            convert(Float64, px), convert(Float64, rv),
            convert(Float64, utc1), convert(Float64, utc2),
            convert(Float64, dut1), convert(Float64, elong),
            convert(Float64, phi), convert(Float64, hm),
            convert(Float64, xp), convert(Float64, yp),
            convert(Float64, phpa), convert(Float64, tc),
            convert(Float64, rh), convert(Float64, wl),
            ref_aob, ref_zob, ref_hob, ref_dob, ref_rob, ref_eo)

   return status, ref_aob[], ref_zob[], ref_hob[], ref_dob[], ref_rob[], ref_eo[]
end


using TracedRandom, Random, Test

@testset "Addressed function calls" begin

gaussian_mixture(μs) = randn(:z) + μs[rand(:k, 1:length(μs))]
Random.seed!(0)
traced = rand(:y, gaussian_mixture, [1, 10, 100])
Random.seed!(0)
untraced = gaussian_mixture([1, 10, 100])
@test traced == untraced

Random.seed!(0)
traced = rand(here, gaussian_mixture, [1, 10, 100])
Random.seed!(0)
untraced = gaussian_mixture([1, 10, 100])
@test traced == untraced

end

@testset "Addressed calls with global RNG" begin

for fn in (rand, randn, randexp, bitrand, randstring)
   Random.seed!(0)
   traced = fn(:addr)
   Random.seed!(0)
   untraced = fn()
   @test traced == untraced
end

for fn in (rand!, randn!, randexp!)
   Random.seed!(0)
   traced = fn(:addr => :subaddr, zeros(5))
   Random.seed!(0)
   untraced = fn(zeros(5))
   @test traced == untraced
end

for fn in (randperm, randcycle)
   Random.seed!(0)
   traced = fn(:addr, 5)
   Random.seed!(0)
   untraced = fn(5)
   @test traced == untraced
end

for fn in (randperm!, randcycle!)
   Random.seed!(0)
   traced = fn(:addr => 0, Vector{Int}(undef, 5))
   Random.seed!(0)
   untraced = fn(Vector{Int}(undef, 5))
   @test traced == untraced
end

Random.seed!(0)
traced = shuffle(:addr, [1, 2, 3, 4, 5])
Random.seed!(0)
untraced = shuffle([1, 2, 3, 4, 5])
@test traced == untraced

Random.seed!(0)
traced = shuffle!(:addr, [1, 2, 3, 4, 5])
Random.seed!(0)
untraced = shuffle!([1, 2, 3, 4, 5])
@test traced == untraced

Random.seed!(0)
traced = randsubseq(:addr, [1, 2, 3, 4, 5], 0.2)
Random.seed!(0)
untraced = randsubseq([1, 2, 3, 4, 5], 0.2)
@test traced == untraced

Random.seed!(0)
traced = randsubseq!(:addr, [], [1, 2, 3, 4, 5], 0.2)
Random.seed!(0)
untraced = randsubseq!([], [1, 2, 3, 4, 5], 0.2)
@test traced == untraced

end

@testset "Addressed calls with custom RNG" begin

rng = MersenneTwister()

for fn in (rand, randn, randexp, bitrand, randstring)
   Random.seed!(rng, 42)
   traced = fn(rng, :addr)
   Random.seed!(rng, 42)
   untraced = fn(rng)
   @test traced == untraced
end

for fn in (rand!, randn!, randexp!)
   Random.seed!(rng, 42)
   traced = fn(rng, :addr => :subaddr, zeros(5))
   Random.seed!(rng, 42)
   untraced = fn(rng, zeros(5))
   @test traced == untraced
end

for fn in (randperm, randcycle)
   Random.seed!(rng, 42)
   traced = fn(rng, :addr, 5)
   Random.seed!(rng, 42)
   untraced = fn(rng, 5)
   @test traced == untraced
end

for fn in (randperm!, randcycle!)
   Random.seed!(rng, 42)
   traced = fn(rng, :addr => 0, Vector{Int}(undef, 5))
   Random.seed!(rng, 42)
   untraced = fn(rng, Vector{Int}(undef, 5))
   @test traced == untraced
end

Random.seed!(rng, 42)
traced = shuffle(rng, :addr, [1, 2, 3, 4, 5])
Random.seed!(rng, 42)
untraced = shuffle(rng, [1, 2, 3, 4, 5])
@test traced == untraced

Random.seed!(rng, 42)
traced = shuffle!(rng, :addr, [1, 2, 3, 4, 5])
Random.seed!(rng, 42)
untraced = shuffle!(rng, [1, 2, 3, 4, 5])
@test traced == untraced

Random.seed!(rng, 42)
traced = randsubseq(rng, :addr, [1, 2, 3, 4, 5], 0.2)
Random.seed!(rng, 42)
untraced = randsubseq(rng, [1, 2, 3, 4, 5], 0.2)
@test traced == untraced

Random.seed!(rng, 42)
traced = randsubseq!(rng, :addr, [], [1, 2, 3, 4, 5], 0.2)
Random.seed!(rng, 42)
untraced = randsubseq!(rng, [], [1, 2, 3, 4, 5], 0.2)
@test traced == untraced

end


export ParticleType
export photon, electron, positron, neutron, proton

@enum ParticleType photon=1 electron=2 positron=3 neutron=4 proton=5

for pt in instances(ParticleType)
    fname = Symbol("is", pt)
    @eval $fname(p) = p.typ == $pt
    eval(Expr(:export, fname))
end

function load(path::AbstractString, T)
    open(path) do io
        load(io, T)
    end
end

function compute_u_v_w(u::Float32, v::Float32, sign_w::Float32)
    tmp = Float64(u)^2 + Float64(v)^2
    if tmp <= 1
        w = Float32(sign_w) * Float32(√(1 - tmp))
    else
        w = Float32(0)
        tmp = √(tmp)
        u = Float32(u/tmp)
        v = Float32(v/tmp)
    end
    u,v,w
end

@noinline function call_fenced(f::F, arg::A) where {F,A}
    f(arg)
end

function kwshow(io::IO, o; calle=nameof(typeof(o)))
    print(io, calle, "(")
    for pname in propertynames(o)
        pval = getproperty(o, pname)
        print(io, string(pname), "=")
        show(io, pval)
        print(io, ", ")
    end
    print(io, ")")
end


type Interpreter <: AbstractPlanOutMapper
    serialization::Dict
    inputs::Dict
    env::Assignment
    salt::AbstractString
    in_experiment::Bool
    evaluated::Bool
end

default_salt = "global_salt"
Interpreter(serialization::Dict) = Interpreter(serialization, Dict(), Assignment(default_salt), default_salt, false, false)
Interpreter(serialization::Dict, inputs::Dict) = Interpreter(serialization, inputs, Assignment(default_salt), default_salt, false, false)
Interpreter(serialization::Dict, salt::AbstractString) = Interpreter(serialization, Dict(), Assignment(salt), salt, false, false)
Interpreter(serialization::Dict, salt::AbstractString, inputs::Dict) = Interpreter(serialization, inputs, Assignment(salt), salt, false, false)

function str_to_op(s::AbstractString)
  if s == AbstractString("<")
    return AbstractString("lessthan")
  elseif s == AbstractString(">")
    return AbstractString("greaterthan")
  elseif s == AbstractString(">=")
    return AbstractString("greaterthanorequalto")
  elseif s == AbstractString("<=")
    return AbstractString("lessthanorequalto")
  elseif s == AbstractString("%")
    return AbstractString("mod")
  elseif s == AbstractString("/")
    return AbstractString("div")
  elseif s == AbstractString("length")
    return AbstractString("len")
  elseif s == AbstractString("return")
    return AbstractString("ret")
  end
  return s
end

setindex!(i::Interpreter, value, key) = setindex!(i.env, value, key)
getindex(i::Interpreter, key) = getindex!(i.env, key)
delete!(i::Interpreter, key) = delete!(i.env, key)

function get_params(i::Interpreter)
  if !i.evaluated
    try
      evaluate(i, i.serialization)
    catch e
      if isa(e, StopPlanOutException)
        i.in_experiment = e.in_experiment
      end
    end
    i.evaluated = true
  end
  return i.env
end

function evaluate(i::Interpreter, planout_code::Any)
  return planout_code
end

function evaluate(i::Interpreter, planout_code::AbstractArray)
  return [evaluate(i, j) for j in planout_code]
end

function evaluate(i::Interpreter, planout_code::Associative)
  if haskey(planout_code, "op")
    op = str_to_op(getindex(planout_code, "op"))
    fxn_call = "PlanOut.$op($planout_code)"
    return execute(eval(parse(fxn_call)), i)
  else
    return planout_code
  end
end


"""
    Cell

Represents a single cell.
"""
mutable struct Cell
    mutations::Array{Int64,1}
    clonetype::Int64
end

mutable struct CloneTracker
    parenttype::Int64
    parentmodule::Int64
    time::Float64
    mutations::Array{Int64, 1}
    N0::Int64
    Ndivisions::Int64
    avdivisions::Float64
    size::Int64
end

struct Clone
    parenttype::Int64
    parentmodule::Int64
    time::Float64
    mutations::Int64
    N0::Int64
    Ndivisions::Int64
    avdivisions::Float64
    freq::Float64
    freqp::Float64
end
struct ModuleTracker
    Nvec::Array{Int64, 1}
    tvec::Array{Float64, 1}
    cells::Array{Cell, 1}
    subclones::Array{CloneTracker, 1}
    id::Int64
    parentid::Int64
end

Base.length(moduletracker::ModuleTracker) = moduletracker.Nvec[end]

abstract type SimulationInput end

struct BranchingInput <: SimulationInput
    numclones::Int64
    Nmax::Int64
    clonalmutations::Int64
    selection::Array{Float64,1}
    μ::Float64
    b::Float64
    d::Float64
    tevent::Array{Float64,1}
    fixedmu::Bool
    maxclonesize::Union{Int64, Nothing}
    ploidy::Int64
end

struct MoranInput <: SimulationInput
    N::Int64
    numclones::Int64
    tmax::Float64
    clonalmutations::Int64
    selection::Array{Float64,1}
    μ::Float64
    bdrate::Float64
    tevent::Array{Float64,1}
    fixedmu::Bool
    ploidy::Int64
end

struct BranchingMoranInput <: SimulationInput
    numclones::Int64
    Nmax::Int64
    tmax::Float64
    clonalmutations::Int64
    selection::Array{Float64,1}
    μ::Float64
    bdrate::Float64
    b::Float64
    d::Float64
    tevent::Array{Float64,1}
    fixedmu::Bool
    ploidy::Int64
end
struct MultilevelInput <: SimulationInput
    modulesize::Int64
    maxtime::Float64
    maxmodules::Int64
    clonalmutations::Int64
    μ::Float64
    bdrate::Float64
    b::Float64
    d::Float64
    fixedmu::Bool
    branchrate::Float64
    branchinitsize::Int64
    ploidy::Int64
end

struct SimulationResult
    subclones::Array{Clone, 1}
    tend::Float64
    trueVAF::Array{Float64,1}
    cells::Array{Cell, 1}
    Nvec::Array{Int64, 1}
    tvec::Array{Float64, 1}
end

struct SampledData
    VAF::Array{Float64,1}
    depth::Array{Int64,1}
end

struct Simulation{T<:SimulationInput}
    input::T
    output::ModuleTracker
end

struct VAFResult{T<:SimulationInput}
    read_depth::Float64
    cellularity::Float64
    detectionlimit::Float64
    input::T
    trueVAF::Array{Float64,1}
    sampledVAF::Array{Float64,1}
    subclonefreq::Array{Float64, 1}
    subclonefreqp::Array{Float64, 1}

end

struct MultiSimulation{T<:SimulationInput}
    input::T
    output::Array{ModuleTracker, 1}
    # mutations::Array{Int64, 1}
end

function MultiSimulation{T}(input::T) where T<:SimulationInput
    return MultiSimulation(
        input,
        ModuleTracker[]
    )
end

Base.length(multisim::MultiSimulation) = length(multisim.output)
Base.iterate(multisim::MultiSimulation) = iterate(multisim.output)
Base.iterate(multisim::MultiSimulation, state) = iterate(multisim.output, state)
Base.getindex(multisim::MultiSimulation, i) = getindex(multisim.output, i)
Base.setindex(multisim::MultiSimulation, v, i) = getindex(multisim.output, v, i)
Base.firstindex(multisim::MultiSimulation) = firstindex(multisim.output)
Base.lastindex(multisim::MultiSimulation) = lastindex(multisim.output)

const Population = MultiSimulation{MultilevelInput}

function get_simulation(multsim, i)
    return Simulation(multsim.input, multsim.output[i])
end

function BranchingInput(;numclones = 1, Nmax = 10000, ploidy = 2, μ = 10.0,
    clonalmutations = μ, selection = fill(0.0,numclones), b = log(2.0), d = 0.0,
    tevent = collect(1.0:0.5:(1+numclones)/2), fixedmu = false,
    maxclonesize = nothing)

    return BranchingInput(
            numclones,
            Nmax,
            clonalmutations,
            selection,
            μ,
            b,
            d,
            tevent,
            fixedmu,
            maxclonesize,
            ploidy
    )
end

function MoranInput(;numclones = 1, N = 10000, ploidy = 2, μ = 10.0, clonalmutations = μ,
    selection = fill(0.0,numclones), bdrate = log(2.0), tmax = 15.0,
    tevent = collect(1.0:0.5:(1+numclones)/2), fixedmu = false)

    return MoranInput(
            N,
            numclones,
            tmax,
            clonalmutations,
            selection,
            μ,
            bdrate,
            tevent,
            fixedmu,
            ploidy
    )
end

function BranchingMoranInput(;numclones = 1, Nmax = 10000, ploidy = 2, μ = 10.0,
    clonalmutations = μ, selection = fill(0.0,numclones), bdrate = log(2.0), b = log(2),
    d = 0, tmax = 15.0, tevent = collect(1.0:0.5:(1+numclones)/2), fixedmu = false)

    return BranchingMoranInput(
            numclones,
            Nmax,
            tmax,
            clonalmutations,
            selection,
            μ,
            bdrate,
            b,
            d,
            tevent,
            fixedmu,
            ploidy
    )

end

function MultilevelInput(;modulesize=200, ploidy=2, μ=10.0, clonalmutations=0,
    bdrate=log(2.0), b=log(2), d=0, maxtime=15, maxmodules=10000, fixedmu=false, branchrate=5,
    branchfraction=0.1, branchinitsize=nothing)

    return MultilevelInput(
            modulesize,
            maxtime,
            maxmodules,
            clonalmutations,
            μ,
            bdrate,
            b,
            d,
            fixedmu,
            branchrate,
            branchinitsize !== nothing ? branchinitsize : ceil(modulesize * branchfraction),
            ploidy,
    )
end

age(moduletracker::ModuleTracker) = moduletracker.tvec[end]
age(populationtracker::Vector{ModuleTracker}) = maximum(map(age, populationtracker))
age(simulation::Simulation) = age(simulation.output)
age(multisim::MultiSimulation) = maximum(age(output) for output in multisim.output)


module RipsererBenchmarks

using BenchmarkTools
using CSV
using DataFrames
using Eirene
using Logging
using ProgressMeter
using Ripserer
using SparseArrays
using TerminalLoggers

export Benchmark, ripser_benchmark, ripserer_benchmark, eirene_benchmark,
    nv, load_data, setup_all, grind_ripser, grind_ripserer, run_and_save, save_csv

include("utils.jl")
include("benchmarks.jl")
include("ripserer.jl")
include("ripser.jl")
include("eirene.jl")
include("benchmarksuites.jl")
include("valgrind.jl")
include("report.jl")

function run_and_save(selected=nothing)
    get_ripsers()

    start_time = time_ns()
    if isnothing(selected)
        suite = setup_all()
        file = "all_results.json"
        @info "running all"
    else
        suite = setup_all()[selected]
        file = "$selected.json"
        @info "running $selected"
    end
    result = with_logger(TerminalLogger()) do
        run(suite)
    end
    BenchmarkTools.save(joinpath(@__DIR__, "..", "results", file), result)
    elapsed = (time_ns() - start_time) / 1e9
    @info "Took $(ProgressMeter.durationstring(elapsed))"
    return result
end

end


using Cairo

function common_header()
    ## header to provide surface and context
    c = CairoRGBSurface(256,256);
    cr = CairoContext(c);

    save(cr);
    set_source_rgb(cr,0.8,0.8,0.8);    # light gray
    rectangle(cr,0.0,0.0,256.0,256.0); # background
    fill(cr);
    restore(cr);

    return c,cr
end

function write_out_picture(filename::AbstractString,c::CairoSurface)
    ## mark picture with current date
    move_to(cr,0.0,12.0);
    set_source_rgb(cr, 0,0,0);
    show_text(cr,Libc.strftime(time()));
    ##
    write_to_png(c,filename);
    nothing
end

function example_current_point(cr)
    save(cr);

    x=25.6;  y=128.0;
    x1=102.4; y1=30.4;
    x2=123.6; y2=45.6;
    x3=130.4; y3=128.0;

    move_to(cr, x, y);
    curve_to(cr, x1, y1, x2, y2, x3, y3);

    if has_current_point(cr)
        x,y = get_current_point(cr);
        save(cr)
        move_to(cr,x,y)
        set_source_rgb(cr, 0,0,1.0);
        show_text(cr,"current point")
        restore(cr)
    end

set_line_width(cr, 10.0);
stroke(cr);

    restore(cr);
    save(cr);

    nothing
end

c,cr = common_header();
save(cr);
example_current_point(cr);
restore(cr);
write_out_picture("sample_current_point.png",c);


#
# InfiniteMPS
#

# TODO: store the cell 1 as an MPS
# Implement `getcell(::InfiniteMPS, n::Integer) -> MPS`
mutable struct InfiniteMPS <: AbstractInfiniteMPS
  data::CelledVector{ITensor}
  llim::Int #RealInfinity
  rlim::Int #RealInfinity
  reverse::Bool
end

translator(ψ::InfiniteMPS) = ψ.data.translator
#
# InfiniteCanonicalMPS
#

# L and R are the orthogonalized MPS tensors.
# C are the center bond matrices (singular values of C
# are the singular values of the MPS at the specified
# cut)
struct InfiniteCanonicalMPS <: AbstractInfiniteMPS
  AL::InfiniteMPS
  C::InfiniteMPS
  AR::InfiniteMPS
end

# TODO: check if `isempty(ψ.AL)`, if so use `ψ.AR`
nsites(ψ::InfiniteCanonicalMPS) = nsites(ψ.AL)
isreversed(ψ::InfiniteCanonicalMPS) = isreversed(ψ.AL)
translator(ψ::InfiniteCanonicalMPS) = translator(ψ.AL)
#ITensors.data(ψ::InfiniteCanonicalMPS) = data(ψ.AL)
ITensors.data(ψ::InfiniteCanonicalMPS) = ψ.AL.data

Base.copy(ψ::InfiniteCanonicalMPS) = InfiniteCanonicalMPS(copy(ψ.AL), copy(ψ.C), copy(ψ.AR))

function fmap(f, ψ::InfiniteCanonicalMPS)
  return InfiniteCanonicalMPS(f(ψ.AL), f(ψ.C), f(ψ.AR))
end

function ITensors.prime(ψ::InfiniteCanonicalMPS, args...; kwargs...)
  return fmap(x -> prime(x, args...; kwargs...), ψ)
end

function ITensors.prime(f::typeof(linkinds), ψ::InfiniteCanonicalMPS, args...; kwargs...)
  # Doesn't prime ψ.C
  #return fmap(x -> prime(f, x, args...; kwargs...), ψ)
  AL = prime(linkinds, ψ.AL, args...; kwargs...)
  C = prime(ψ.C, args...; kwargs...)
  AR = prime(linkinds, ψ.AR, args...; kwargs...)
  return InfiniteCanonicalMPS(AL, C, AR)
end

function ITensors.dag(ψ::InfiniteCanonicalMPS, args...; kwargs...)
  return fmap(x -> dag(x, args...; kwargs...), ψ)
end

ITensors.siteinds(f::typeof(only), ψ::InfiniteCanonicalMPS) = siteinds(f, ψ.AL)

getcell(i::Index) = ITensorInfiniteMPS.getcell(tags(i))
getsite(i::Index) = getsite(tags(i))

#Before, for a two site Hamiltonian, findsites(ψ, H[3]) would return [1, 2]
#Appeared unsafe for index purpose
function ITensors.findfirstsiteind(ψ::InfiniteMPS, i::Index)
  c = ITensorInfiniteMPS.getcell(i)
  n1 = getsite(i)
  # if translator(ψ) == translatecelltags
  return (c - 1) * nsites(ψ) + n1
  # else
  #   s = siteinds(only, ψ)
  #   n = nsites(ψ)
  #   j = 1
  #   index = 0
  #   while c > ITensorInfiniteMPS.getcell(s[j]) && index < 1000
  #     j += n
  #     index += 1
  #   end
  #   while c < ITensorInfiniteMPS.getcell(s[j]) && index < 1000
  #     j -= n
  #     index += 1
  #   end
  #   if index == 1000 || c != ITensorInfiniteMPS.getcell(s[j])
  #     error("Index not found")
  #   end
  #   for k in 1:n
  #     if getsite(s[j]) == n1
  #       return j
  #     end
  #     j += 1
  #   end
  #   error("Index not found")
  #   return 0
  # end
end
function ITensors.findfirstsiteind(ψ::InfiniteCanonicalMPS, i::Index)
  return ITensors.findfirstsiteind(ψ.AL, i)
end

function ITensors.findsites(ψ::InfiniteCanonicalMPS, is::Union{<:Tuple,<:Vector})
  return sort([ITensors.findfirstsiteind(ψ, i) for i in is])
end
function ITensors.findsites(ψ::InfiniteMPS, is::Union{<:Tuple,<:Vector})
  return sort([ITensors.findfirstsiteind(ψ, i) for i in is])
end
function ITensors.findsites(ψ::InfiniteMPS, T::MPO)
  s = [noprime(filterinds(T[x]; plev=1)[1]) for x in 1:length(T)]
  return sort([ITensors.findfirstsiteind(ψ, i) for i in s])
end
ITensors.findsites(ψ::InfiniteCanonicalMPS, T::MPO) = findsites(ψ.AL, T)

#Kept for historical reason
function ITensors.findsites(ψ::InfiniteMPS, T::ITensor)
  s = filterinds(T; plev=0)
  return sort([ITensors.findfirstsiteind(ψ, i) for i in s])
end
ITensors.findsites(ψ::InfiniteCanonicalMPS, T::ITensor) = findsites(ψ.AL, T)

# For now, only represents nearest neighbor interactions
# on a linear chain
struct InfiniteSum{T}
  data::CelledVector{T}
end
InfiniteSum{T}(N::Int) where {T} = InfiniteSum{T}(Vector{T}(undef, N))
InfiniteSum{T}(data::Vector{T}) where {T} = InfiniteSum{T}(CelledVector(data))
function InfiniteSum{T}(data::Vector{T}, translator::Function) where {T}
  return InfiniteSum{T}(CelledVector(data, translator))
end
translator(h::InfiniteSum{T}) where {T} = translator(h.data)
# Automatically converts from ITensor to MPO.
# XXX: check this conversion is correct.
#InfiniteSum{T}(data::Vector{ITensor}) where {T} = InfiniteSum{T}(CelledVector(T.(data)))
#This crashed for me when T is ITensor
function InfiniteSum{MPO}(data::Vector{ITensor})
  N = length(data)
  temp_inds = [filterinds(data[n]; plev=0) for n in 1:N]
  return InfiniteSum{MPO}([
    MPO(
      data[n],
      [(temp_inds[n][j], dag(prime(temp_inds[n][j]))) for j in 1:length(temp_inds[n])],
    ) for n in 1:N
  ])
end

function InfiniteSum{MPO}(data::Vector{ITensor}, translator::Function)
  N = length(data)
  temp_inds = [filterinds(data[n]; plev=0) for n in 1:N]
  return InfiniteSum{MPO}(
    [
      MPO(
        data[n],
        [(temp_inds[n][j], dag(prime(temp_inds[n][j]))) for j in 1:length(temp_inds[n])],
      ) for n in 1:N
    ],
    translator,
  )
end

function InfiniteSum{MPO}(infsum::InfiniteSum{ITensor})
  N = nsites(infsum)
  temp_inds = [filterinds(infsum[n]; plev=0) for n in 1:N]
  return InfiniteSum{MPO}(
    [
      MPO(
        infsum[n],
        [(temp_inds[n][j], dag(prime(temp_inds[n][j]))) for j in 1:length(temp_inds[n])],
      ) for n in 1:N
    ],
    translator(infsum),
  )
end

function Base.getindex(l::InfiniteSum, n1n2::Tuple{Int,Int})
  n1, n2 = n1n2
  @assert n2 == n1 + 1
  return l.data[n1]
end
function Base.getindex(l::InfiniteSum, n1::Int)
  return l.data[n1]
end
nsites(h::InfiniteSum) = length(h.data)
#Gives the range of the Hamiltonian. Useful for better optimized contraction in VUMPS
nsites_support(h::InfiniteSum) = length.(h.data)
nsites_support(h::InfiniteSum, n::Int64) = length(h.data[n])

nrange(h::InfiniteSum) = nrange.(h.data, ncell=nsites(h))
nrange(h::InfiniteSum, n::Int64) = nrange(h.data[n]; ncell=nsites(h))

function nrange(h::MPO; ncell=1)
  ns = findsites(h; ncell=ncell)
  return ns[end] - ns[1] + 1
end

ITensors.findsites(h::InfiniteSum) = [findsites(h, n) for n in 1:nsites(h)]
ITensors.findsites(h::InfiniteSum, n::Int64) = findsites(h.data[n]; ncell=nsites(h))
#ITensors.findsites(h::InfiniteSum, is::Union{<:Tuple,<:Vector}) = [findsites(h.data[n], is) for n in 1:nsites(h)]
#ITensors.findsites(h::InfiniteSum, i::Index) = [findsites(h.data[n], i) for n in 1:nsites(h)]
#TODO improve the findsites routines for Infinite Sum

function ITensors.findfirstsiteind(i::Index, ncell::Int64)
  c = ITensorInfiniteMPS.getcell(i)
  n1 = getsite(i)
  return (c - 1) * ncell + n1
end
function ITensors.findsites(h::MPO; ncell::Int64=1)
  s = [filterinds(h[x]; plev=1)[1] for x in 1:length(h)]
  return sort([ITensors.findfirstsiteind(i, ncell) for i in s])
end

#Kept for historical reasons
function nrange(h::ITensor; ncell=0)
  ns = findsites(h; ncell=ncell)
  return ns[end] - ns[1] + 1
end
function nrange(ψ::InfiniteCanonicalMPS, h::ITensor)
  ns = findsites(ψ, h)
  return ns[end] - ns[1] + 1
end
function nrange(ψ::InfiniteCanonicalMPS, h::MPO)
  ns = findsites(ψ, h)
  return ns[end] - ns[1] + 1
end

function ITensors.findfirstsiteind(h::ITensor, i::Index, ncell::Int64)
  c = getcell(i)
  n1 = getsite(i)
  return (c - 1) * ncell + n1
end
function ITensors.findsites(h::ITensor; ncell::Int64=1)
  s = filterinds(h; plev=0)
  return sort([ITensors.findfirstsiteind(h, i, ncell) for i in s])
end

## HDF5 support for the InfiniteCanonicalMPS type

function HDF5.write(
  parent::Union{HDF5.File,HDF5.Group}, name::AbstractString, ψ::InfiniteCanonicalMPS
)
  g = create_group(parent, name)
  attributes(g)["type"] = "InfiniteCanonicalMPS"
  attributes(g)["version"] = 1
  N = nsites(ψ)
  write(g, "length", N)
  for n in 1:N
    write(g, "AL[$(n)]", ψ.AL[n])
    write(g, "AR[$(n)]", ψ.AR[n])
    write(g, "C[$(n)]", ψ.C[n])
  end
end

function HDF5.read(
  parent::Union{HDF5.File,HDF5.Group}, name::AbstractString, ::Type{InfiniteCanonicalMPS}
)
  g = open_group(parent, name)
  if read(attributes(g)["type"]) != "InfiniteCanonicalMPS"
    error("HDF5 group or file does not contain InfiniteCanonicalMPS data")
  end
  N = read(g, "length")
  AL = InfiniteMPS([read(g, "AL[$(i)]", ITensor) for i in 1:N])
  AR = InfiniteMPS([read(g, "AR[$(i)]", ITensor) for i in 1:N])
  C = InfiniteMPS([read(g, "C[$(i)]", ITensor) for i in 1:N])
  return InfiniteCanonicalMPS(AL, C, AR)
end