# Author: AJ Tyler
# Date: 24/01/22
module Comparison # This module contains functions which compare configurations using their CNA profiles.
include("io.jl") # This file contains the IO module.
using .IO # This module contains functions that read in and write out files.
using Statistics # This module allows easy mean and standard deviation.
export compareConfigs,groupConfigs,configStats # These functions ared used to compare configurations.

"""
Finds the most similar configuration to every configuration over all rCut values based on their CNA profile.

Inputs:
rCutRange: (Vector{Float64}) A range of cut-off radii to test, in units of equilibriumBondLength.
energies: (Vector{Float64}) The energy of each configuration in each the file.
totalProfiles: Array{Dict{String,Int}} 2D array containg the totalCNA profiles of each configuration, for each rCut value.
atomicProfiles: Array{Dict{String,Int}} 3D array containing the CNA profile of each configuration, for each rCut value, for each atom.
configDir: (String) The path to (or name of if in same location as this program) the directory containing the .xyz files.
filename: (String) The prefix of the file.
similarityMeasure: (String) The method of comparing configurations, either "atomic" or "total" CNA.
M: (Int64) The number of rCut values to test.
L: (Int64) The number of configurations in the file.
N: (Int64) The number of atoms in each configuration.

Outputs:
similarConfigs: Vector{Dict{Int64,Vector{Float64}}} A vector of dictionaries which map the most similar configurations to the rCut values
													for which it is maximised, for each configuraiton.
maxSims: Vector{Float64} A Vector of the maximum similarity score for each configuration when compared to every other configuration
						 over all rCut values.
"""
function compareConfigs(rCutRange,energies,totalProfiles,atomicProfiles,configDir,filename,similarityMeasure,M,L,N)
	# Initialise vector to contain most similar configurations for each configuration
	similarConfigs = Vector{Dict{Int64,Vector{Float64}}}(undef,L)
	maxSims = zeros(Float64,1,L) # Initialise the maximum similarity score for each configuration
	for i in 1:M # For all rCut values
		writeProfiles(rCutRange,configDir,filename,energies,totalProfiles,i,L) # Write out the total CNA profiles
		if L == 1
			continue
		end
		atomicProfilesDict = convertAtomicProfiles(atomicProfiles,i,L) # Convert atomic CNA profile format
		# For each pair of configurations
		for configA in 1:L-1
			for configB in configA+1:L
				# Compute the similarity for the cluster pair
				similarity = similarityScore(configA, configB, totalProfiles,atomicProfilesDict,similarityMeasure,N,i)
				# Update the most similar configurations to configA
				if (similarity > maxSims[configA]) # If new similarity score is greater than previous highest
					similarConfigs[configA] = Dict{Int64,Vector{Float64}}(configB=>Vector{Float64}([rCutRange[i]])) # Create new vector, containing configB
					maxSims[configA] = similarity # Update maximum similarity score
				elseif (similarity == maxSims[configA]) # If the new similarity is the same as the previous highest
					if (haskey(similarConfigs[configA],configB)) # If the same configuration pair is already most similar
						# Add the current rCut value to the dictionary
						push!(similarConfigs[configA], configB => push!(get!(similarConfigs[configA],configB,0),rCutRange[i]))
					else
						push!(similarConfigs[configA], configB => Vector{Float64}([rCutRange[i]])) # Add the configuration to the dictionary
					end
				end

				# Same for configB
				if (similarity > maxSims[configB]) # If new similarity score is greater than previous highest
					similarConfigs[configB] = Dict{Int64,Vector{Float64}}(configA=>Vector{Float64}([rCutRange[i]])) # Create new vector, containing configB
					maxSims[configB] = similarity # Update maximum similarity score
				elseif (similarity == maxSims[configB]) # If the new similarity is the same as the previous highest
					if (haskey(similarConfigs[configB],configA))
						push!(similarConfigs[configB], configA => push!(get!(similarConfigs[configB],configA,0),rCutRange[i]))
					else
						push!(similarConfigs[configB], configA => Vector{Float64}([rCutRange[i]]))
					end
				end
			end
		end
	end
	return similarConfigs, maxSims # Return outputs
end

"""
Converts a 2D array of containing the CNA profile of each configuration, for each atom, to a vector of dictionaries,
mapping atom CNA profiles to their frequency, for each configuration.

Inputs:
atomicProfiles: Array{Dict{String,Int}} 3D array containing the CNA profile of each configuration, for each rCut value, for each atom.
i: (Int64) The index of the current rCut value.
L: (Int64) The number of configurations in the file.

Outputs:
atomicProfileDict: Vector{Dict{String,Int64}} ector of dictionaries, mapping atom CNA profiles to their frequency, for each configuration.
"""
function convertAtomicProfiles(atomicProfiles,i,L)
	# Convert atomicProfiles into format better for computing similarity
	atomicProfilesDict = [Dict{String,Int64}() for j in 1:L] # Initialise new format
	for j in 1:L # For all configurations
		atomicProfile = atomicProfiles[i,j,:] # Get atom profile
		for atomProfile in atomicProfile # For each atom
			# Increments atomic CNA profile frequency by 1
			push!(atomicProfilesDict[j], "$atomProfile" => get!(atomicProfilesDict[j],"$atomProfile",0)+1)
		end
	end

	return atomicProfilesDict # Return output
end

"""
This function returns the similarity score of two configurations based on their CNA profiles.

Inputs:
configA: (Int64) The configuration number of one of the configurations to compare.
configB: (Int64) The configuration number of the other of the configurations to compare.
totalProfiles: Array{Dict{String,Int}} 2D array containg the totalCNA profiles of each configuration, for each rCut value.
atomicProfilesDict: Vector{Dict{String,Int64}} ector of dictionaries, mapping atom CNA profiles to their frequency, for each configuration.
similarityMeasure: (String) The method of comparing configurations, either "atomic" or "total" CNA.
N: (Int64) The number of atoms in each configuration.
i: (Int64) The index of the current rCut value.

Outputs:
similarity: (Float64) The similarity score between the two configurations.
"""
function similarityScore(configA, configB, totalProfiles,atomicProfilesDict,similarityMeasure,N,i)
	similarity = 0.0 # Initialise similarity score
	if (similarityMeasure == "atomic") # If want to compare using atomic CNA profile
		# Compute similarity of atomic CNA profiles
		for key in keys(atomicProfilesDict[configA]) # For each atom CNA profile
			# Add the number of atoms with the same atom CNA profile
			similarity += min(get(atomicProfilesDict[configA],key,0),get(atomicProfilesDict[configB],key,0))
		end
		similarity /= N # Normalise the similarity score

	elseif (similarityMeasure == "total") # If want to compare using total CNA profile
		# Compute Jacard index of their total CNA profiles
		similarity = JacardIndex(totalProfiles[i,configA],totalProfiles[i,configB])
	else # If none of the valid similarity measures were used
		throw(DomainError(similarityMeasure, "Invalid similarity measure")) # Throw an error
	end

	return round(similarity; digits = 3) # Return rounded similarity score
end

"""
This function calculates the Jacard Index for two dictionaries,
with their keys being the unique elements of a multiset and the values are their frequencies

Inputs:
A: (Dictionary) Multiset A
B: (Dictionary) Multiset B

Outputs:
similarity: (Float64) The Jacard Index of multsets A & B
"""
function JacardIndex(A,B)
	intersection = 0.0 # Initialise the number of elements in the intersection of A & B
	union = 0.0 # Initialise the number of elements in the union of A & B
	for key in keys(A) # For all the unique elements of A
		intersection += min(get(A,key,0),get(B,key,0)) # Add the minimum number of occurances of the key across both A & B
		union += get(A,key,0) # Add all the elements of A
	end

	for key in keys(B) # For all the unique elements of B
		union += get(B,key,0) # Add all the elements of B
	end

	union -= intersection # union = A + B - intersection
	similarity = intersection/union # Comput the Jacard index
	return similarity # Return the Jacard Index
end

"""
This function groups configurations into distinguisable sets based on their similarity scores.

Inputs:
energies: (Vector{Float64}) The energy of each configuration in each the file.
sortingArray: Vector{Int64} A vector of configuration numbers, sorted by maximum similarity score.
maxSims: Vector{Float64} A Vector of the maximum similarity score for each configuration when compared to every other configuration
						 over all rCut values, in decreasing order.
similarConfigs: Vector{Dict{Int64,Vector{Float64}}} A vector of dictionaries which map the most similar configurations to the rCut values
						 							for which it is maximised, for each configuraiton, sorted by similarity.
similarityThreshold: (Float64) Similarity threshold above which two configurations are considered identical.
M: (Int64) The number of rCut values tested.

Keyword Arguments:
rCutThreshold: (Float64) Proportion of rCut values that have to meet the similarity threshold for two configs to be considered identical.

Outputs:
uniqueConfigs: Vector{Vector{Int64}} Vector of configuration numbers of the distinguisable sets of configurations.
uniqueEnergies: Vector{Vector{Float64}} Vector of configuration energies of the distinguisable sets of configurations.
uniqueSimilarities: Vector{Float64} Vector of similarity scores of each of the distinguiable sets.
"""
function groupConfigs(energies,sortingArray,maxSims,similarConfigs,similarityThreshold,M;rCutThreshold=0.9)
	# Initialise arrays
	uniqueConfigs = Vector{Vector{Int64}}()
	uniqueEnergies = Vector{Vector{Float64}}()
	uniqueSimilarities = Vector{Float64}()
	classified = Vector{Int64}() # Vector of configuration numbers that have already been grouped
	for (j,config) in enumerate(sortingArray) # For all configurations
		if (!(config in classified)) # If the configuration has not already been grouped
			same = Vector{Int64}([config])
			if (maxSims[j] > similarityThreshold) # If the maximum similarity of the configuration exceeds a threshold
				for key in keys(similarConfigs[j])
					if (length(get!(similarConfigs[j],key,0)) > rCutThreshold*M && !(key in classified))
						# Get vector of configuration numbers of configurations that are similar
						push!(same,key)
					end
				end
			end
			push!(uniqueConfigs,same) # Add the set of configurations to the list
			push!(uniqueEnergies,energies[same]) # Add the energies of the configurations to the list
			push!(uniqueSimilarities,maxSims[j]) # Add the similarity score to the list
			classified = vcat(classified,same) # Add the configurations to the list of already classified
		end
	end

	return uniqueConfigs, uniqueEnergies, uniqueSimilarities # Return outputs
end

"""
This function returns some basic statistics of the energies of each of the distinguisbale configuration sets.

Inputs:
uniqueEnergies: (Vector{Vector{Float64}}) Vector of configuration energies of the distinguisable sets of configurations.
numUnique (Int64) Number of distinguiable sets of configurations.

Outputs:
uniqueEnergyStats: Array{Float64} 2D array of energy statistics. First row is the mean, second row is standard deviation.
"""
function configStats(uniqueEnergies,numUnique)
	uniqueEnergyStats = Array{Float64}(undef,2,numUnique) # Initialse array
	for i in 1:numUnique # For each grouped set of configurations
		uniqueEnergyStats[1,i] = Statistics.mean(uniqueEnergies[i]) # Compute mean
		uniqueEnergyStats[2,i] = Statistics.std(uniqueEnergies[i]) # Compute standard deviation
	end

	return uniqueEnergyStats # Return statistics
end

end # End of Module


"""
This module provides utilities to build neural networks with Flux,
along with a library of standard architectures.
"""
module FluxNets

export SimpleNet, SimpleNetHP, ResNet, ResNetHP

using ..Network
using Base: @kwdef
import ..GameInterface, ..Util, ..CyclicSchedule

using CUDA

import Flux
import Functors

CUDA.allowscalar(false)
array_on_gpu(::Array) = false
array_on_gpu(::CuArray) = true
array_on_gpu(arr) = error("Usupported array type: ", typeof(arr))

using Flux: relu, softmax, flatten
using Flux: Chain, Dense, Conv, BatchNorm, SkipConnection
import Zygote

#####
##### Flux Networks
#####

"""
    FluxNetwork{Game} <: AbstractNetwork{Game}

Abstract type for neural networks implemented using the _Flux_ framework.

The `regularized_params_` function must be overrided for all layers containing
parameters that are subject to regularization.

Provided that the above holds, `FluxNetwork` implements the full
network interface with the following exceptions:
[`Network.HyperParams`](@ref), [`Network.hyperparams`](@ref),
[`Network.forward`](@ref) and [`Network.on_gpu`](@ref).
"""
abstract type FluxNetwork{Game} <: AbstractNetwork{Game} end

function Base.copy(nn::Net) where Net <: FluxNetwork
  #new = Net(Network.hyperparams(nn))
  #Flux.loadparams!(new, Flux.params(nn))
  #return new
  return Base.deepcopy(nn)
end

Network.to_cpu(nn::FluxNetwork) = Flux.cpu(nn)

function Network.to_gpu(nn::FluxNetwork)
  CUDA.allowscalar(false)
  return Flux.gpu(nn)
end

function Network.set_test_mode!(nn::FluxNetwork, mode)
  Flux.testmode!(nn, mode)
end

Network.convert_input(nn::FluxNetwork, x) =
  Network.on_gpu(nn) ? Flux.gpu(x) : x

Network.convert_output(nn::FluxNetwork, x) = Flux.cpu(x)

Network.params(nn::FluxNetwork) = Flux.params(nn)

# This should be included in Flux
function lossgrads(f, args...)
  val, back = Zygote.pullback(f, args...)
  grad = back(Zygote.sensitivity(val))
  return val, grad
end

function Network.train!(callback, nn::FluxNetwork, opt::Adam, loss, data, n)
  optimiser = Flux.ADAM(opt.lr)
  params = Flux.params(nn)
  for (i, d) in enumerate(data)
    l, grads = lossgrads(params) do
      loss(d...)
    end
    Flux.update!(optimiser, params, grads)
    callback(i, l)
  end
end

function Network.train!(
    callback, nn::FluxNetwork, opt::CyclicNesterov, loss, data, n)
  lr = CyclicSchedule(
    opt.lr_base,
    opt.lr_high,
    opt.lr_low, n=n)
  momentum = CyclicSchedule(
    opt.momentum_high,
    opt.momentum_low,
    opt.momentum_high, n=n)
  optimiser = Flux.Nesterov(opt.lr_low, opt.momentum_high)
  params = Flux.params(nn)
  for (i, d) in enumerate(data)
    l, grads = lossgrads(params) do
      loss(d...)
    end
    Flux.update!(optimiser, params, grads)
    optimiser.eta = lr[i]
    optimiser.rho = momentum[i]
    callback(i, l)
  end
end

regularized_params_(l) = []
regularized_params_(l::Flux.Dense) = [l.W]
regularized_params_(l::Flux.Conv) = [l.weight]

# Reimplementation of what used to be Flux.prefor, does not visit leaves
function foreach_flux_node(f::Function, x, seen = IdDict())
  Functors.isleaf(x) && return
  haskey(seen, x) && return
  seen[x] = true
  f(x)
  for child in Flux.trainable(x)
    foreach_flux_node(f, child, seen)
  end
end

function Network.regularized_params(net::FluxNetwork)
  ps = Flux.Params()
  foreach_flux_node(net) do p
    for r in regularized_params_(p)
      any(x -> x === r, ps) || push!(ps, r)
    end
  end
  return ps
end

function Network.gc(::FluxNetwork)
  GC.gc(true)
  # CUDA.reclaim()
end

#####
##### Common functions between two-head neural networks
#####

"""
    TwoHeadNetwork{Game} <: FluxNetwork{G}

An abstract type for two-head neural networks implemented with Flux.

Subtypes are assumed to have fields
`hyper`, `common`, `vhead` and `phead`. Based on those, an implementation
is provided for [`Network.hyperparams`](@ref), [`Network.forward`](@ref) and
[`Network.on_gpu`](@ref), leaving only [`Network.HyperParams`](@ref) to
be implemented.
"""
abstract type TwoHeadNetwork{G} <: FluxNetwork{G} end

function Network.forward(nn::TwoHeadNetwork, board)
  c = nn.common(board)
  v = nn.vhead(c)
  p = nn.phead(c)
  return (p, v)
end

# Flux.@functor does not work do to Network being parametric
function Flux.functor(nn::Net) where Net <: TwoHeadNetwork
  children = (nn.common, nn.vhead, nn.phead)
  constructor = cs -> Net(nn.hyper, cs...)
  return (children, constructor)
end

Network.hyperparams(nn::TwoHeadNetwork) = nn.hyper

Network.on_gpu(nn::TwoHeadNetwork) = array_on_gpu(nn.vhead[end].b)

#####
##### Include networks library
#####

include("architectures/simplenet.jl")
include("architectures/resnet.jl")

end


module SqState
    using QuantumStateBase
    using QuantumStateDistributions
    using LinearAlgebra

    using Fetch, DataDeps, JLD2

    using CUDA
    using Flux
    using Flux.Data: DataLoader
    using Zygote, ChainRulesCore

    using NeuralOperators

    function __init__()
        register(DataDep(
            "SqState",
            """Data and models for SqState.""",
            "https://drive.google.com/file/d/1UzaPBpTuhxvmyUWnoOupEr3cRZUMc-0-/view?usp=sharing",
            "4f426e53a018ae0692fb656b5fa63712157f0889bb3b2cf7de2a169764f9a00e",
            fetch_method=gdownload,
            post_fetch_method=unpack
        ))
    end

    include("utils.jl")
    include("gen_data.jl")
    include("preprocess.jl")
    include("model.jl")
end


# ##############################################################
# ### optimisation modelling sets
# ##############################################################
# ECH, TS, S, NAMES = Workflow.get_ech_ts_s_name(launcher);
# BUSES =  Workflow.get_bus(launcher, NAMES);
# units_by_kind = Workflow.get_units_by_kind(launcher);
# units_by_bus =  Workflow.get_units_by_bus(launcher, BUSES);

# println("NAMES : ", NAMES);
# println("BUSES : ", BUSES);
# println("units_by_bus : ", units_by_bus);

# println("Number of scenario ", length(S));
# println("Number of time step is ", length(TS));
# println("ECH : ", ECH);
# K_IMPOSABLE = Workflow.K_IMPOSABLE;
# K_LIMITABLE = Workflow.K_LIMITABLE;

# # netloads = Dict([(bus, ts, s) => launcher.uncertainties[bus, s, ts, ech] for bus in BUSES, ts in TS, s in S]);
# eod_slack = Dict([(ts,s) => 0.0 for s in S, ts in TS]);
# factor = Dict([name => 1 for name in NAMES]);
# for bus in BUSES
#     factor[bus] = -1;
# end
# for ((name, s, ts, ech), v) in launcher.uncertainties
#     eod_slack[ts,  s] += factor[name] * v;
# end
# println("eod_slack is $eod_slack");

function worse_case(launcher::Launcher, ech::DateTime)
    #######################################################################
    # for a given ts, at a given ech, find the worst case within
    # uncertainties for why the economical balance maximise the constraints
    #######################################################################
    # max rho
    # u in [lb, ub]
    # sum p = u
    # p empilement économique
    # rho est la surcharge minimale

    is_on = Dict([kvp[1] => true for kvp in launcher.units])
    units_cost = Dict();
    # name->P-maxP-startCost-propCost
    for unit in launcher.units
        name = unit[1];
        minP = unit[2][1];
        maxP = unit[2][2];
        startCost = unit[2][3];
        propCost = unit[2][4];

    end

    units_by_kind = Workflow.get_units_by_kind(launcher);
    println(is_on);
end


using Pkg
 Pkg.activate(".")

verbose = true
if verbose   println("# Loading RvSpecML")    end
 using RvSpectML
 if verbose   println("# Loading other packages")    end
 using DataFrames, Query, Statistics
 using Dates

# USER:  The default paths that specify where datafiles can be entered here or overridden in examples/data_paths.jl
target_subdir = "101501"   # USER: Replace with directory of your choice
 fits_target_str = "101501"
 output_dir = "examples/output"
 default_paths_to_search = [pwd(),"examples",joinpath(pkgdir(RvSpectML),"examples"),"/gpfs/group/ebf11/default/ebf11/expres/inputs"]

 make_plots = true
 write_ccf_to_csv = true
 write_rvs_to_csv = true
 write_order_ccf_to_csv = true
 write_template_to_csv = true
 write_spectral_grid_to_jld2 = true
 write_dcpca_to_csv = true
 write_lines_to_csv = true

has_loaded_data = false
 has_computed_ccfs = false
 has_computed_rvs = false
 has_computed_order_ccfs = false
 has_computed_template = false
 has_computed_dcpca = false
 has_found_lines = false
 has_computed_ccfs2 = false
 has_computed_rvs2 = false
 has_computed_template2 = false
 has_computed_dcpca2 = false

if !has_loaded_data
  if verbose println("# Finding what data files are avaliable.")  end
  df_files = make_manifest(target_subdir, EXPRES )

  if verbose println("# Reading in customized parameters.")  end
  eval(code_to_include_param_jl())

  if verbose println("# Reading in FITS files.")  end
  @time expres_data = map(EXPRES.read_data,eachrow(df_files_use))
  has_loaded_data = true

  if verbose println("# Extracting order list timeseries from spectra.")  end
  order_list_timeseries = RvSpectML.make_order_list_timeseries(expres_data)
  order_list_timeseries = RvSpectML.filter_bad_chunks(order_list_timeseries,verbose=true)
  RvSpectML.normalize_spectra!(order_list_timeseries,expres_data);
end

if !has_computed_ccfs || true
if verbose println("# Reading line list for CCF: ", linelist_for_ccf_filename, ".")  end
 lambda_range_with_good_data = get_λ_range(expres_data)
 espresso_filename = joinpath(pkgdir(RvSpectML),"data","masks",linelist_for_ccf_filename)
 espresso_df = RvSpectML.read_linelist_espresso(espresso_filename)
 line_list_df = EXPRES.filter_line_list(espresso_df,first(expres_data).inst)

if verbose println("# Removing lines with telluric contamination.")  end    # Currently only works for EXPRES data
 Δv_to_avoid_tellurics = 14000.0
 line_list_to_search_for_tellurics = copy(line_list_df)
 line_list_to_search_for_tellurics.lambda_lo = line_list_to_search_for_tellurics.lambda./calc_doppler_factor(Δv_to_avoid_tellurics)
 line_list_to_search_for_tellurics.lambda_hi = line_list_to_search_for_tellurics.lambda.*calc_doppler_factor(Δv_to_avoid_tellurics)
 chunk_list_timeseries = RvSpectML.make_chunk_list_timeseries(expres_data,line_list_to_search_for_tellurics)
 line_list_to_search_for_tellurics.min_telluric_model_all_obs = RvSpectML.find_worst_telluric_in_each_chunk( chunk_list_timeseries, expres_data)
 line_list_no_tellurics_df = line_list_to_search_for_tellurics |> @filter(_.min_telluric_model_all_obs == 1.0) |> DataFrame
 #=
 Δv_to_avoid_tellurics = RvSpectMLBase.max_bc # default_Δv_to_avoid_tellurics
 line_list_no_tellurics_df.lambda_lo = line_list_no_tellurics_df.lambda./calc_doppler_factor(Δv_to_avoid_tellurics)
 line_list_no_tellurics_df.lambda_hi = line_list_no_tellurics_df.lambda.*calc_doppler_factor(Δv_to_avoid_tellurics)
 #find_overlapping_chunks(line_list_no_tellurics_df)
 #chunk_list_no_tellurics_merged_df = chunk_list_no_tellurics_merged_df |> @take(1) |> DataFrame
 #chunk_list_no_tellurics_merged_df = RvSpectML.merge_chunks(line_list_no_tellurics_df)
 #@assert find_overlapping_chunks(chunk_list_no_tellurics_merged_df) == nothing
 #chunk_list_timeseries = RvSpectML.make_chunk_list_timeseries(expres_data,chunk_list_no_tellurics_merged_df)
 chunk_list_timeseries = RvSpectML.make_chunk_list_timeseries(expres_data,line_list_no_tellurics_df)
 #chunk_list_timeseries = RvSpectML.make_chunk_list_timeseries(expres_data,chunk_list_no_tellurics_merged_df)
=#

 # Compute CCF's & measure RVs
if verbose println("# Computing CCF.")  end
 mask_shape = RvSpectML.CCF.TopHatCCFMask(order_list_timeseries.inst, scale_factor=tophap_ccf_mask_scale_factor)
 line_list = RvSpectML.CCF.BasicLineList(line_list_df.lambda, line_list_df.weight)
 line_list = RvSpectML.CCF.BasicLineList(line_list_no_tellurics_df.lambda, line_list_no_tellurics_df.weight)
 #line_list = RvSpectML.CCF.BasicLineList(chunk_list_no_tellurics_merged_df.lambda, chunk_list_no_tellurics_merged_df.weight)
 ccf_plan = RvSpectML.CCF.BasicCCFPlan(mask_shape = mask_shape, line_list=line_list, midpoint=ccf_mid_velocity)
 v_grid = RvSpectML.CCF.calc_ccf_v_grid(ccf_plan)
 @time ccfs = RvSpectML.CCF.calc_ccf_chunklist_timeseries(order_list_timeseries, ccf_plan)
 #@time ccfs = RvSpectML.CCF.calc_ccf_chunklist_timeseries(chunk_list_timeseries, ccf_plan)
 # Write CCFs to file
 if write_ccf_to_csv
    using CSV
    CSV.write(joinpath(output_dir,target_subdir * "_ccfs.csv"),Tables.table(ccfs',header=Symbol.(v_grid)))
 end
 has_computed_ccfs = true
end

#make_plots = true
if make_plots
   using Plots
   t_idx = 1:size(ccfs,2)
   using Plots
   #plot(v_grid,ccfs[:,t_idx]./maximum(ccfs[:,t_idx],dims=1),label=:none)
   plot(v_grid,ccfs[:,t_idx],label=:none)
   xlabel!("v (m/s)")
   ylabel!("CCF")
 end

if !has_computed_rvs
 if verbose println("# Measuring RVs from CCF.")  end
 rvs_ccf_gauss = RvSpectML.RVFromCCF.measure_rv_from_ccf(v_grid,ccfs,fit_type = "gaussian")
 # Store estimated RVs in metadata
 map(i->order_list_timeseries.metadata[i][:rv_est] = rvs_ccf_gauss[i]-mean(rvs_ccf_gauss), 1:length(order_list_timeseries) )
 if write_rvs_to_csv
    using CSV
    CSV.write(joinpath(output_dir,target_subdir * "_rvs_ccf.csv"),DataFrame("Time [MJD]"=>order_list_timeseries.times,"CCF RV [m/s]"=>rvs_ccf_gauss))
 end
 has_computed_rvs = true
end

if !has_computed_order_ccfs  # Compute order CCF's & measure RVs
tstart = now()    # Compute CCFs for each order
 # TODO:  Need to mask tellurics somehow
 RvSpectML.CCF.calc_order_ccf_chunklist_timeseries(order_list_timeseries, ccf_plan)
 #=
 orders_to_use = RvSpectML.orders_to_use_default(chunk_list_timeseries.inst)
 num_obs = length(chunk_list_timeseries)
 order_ccfs = Array{Float64,3}(undef,length(v_grid), length(orders_to_use), num_obs)
 for (i,ord) in enumerate(orders_to_use)
     chunk_list_order_timeseries = RvSpectML.extract_chunk_list_timeseries_for_order(order_list_timeseries, ord)
     if isnothing(chunk_list_order_timeseries)
        println("# Warning: No usable lines in order = ", ord)
        order_ccfs[:,i,:] .= 0.0
     else
        order_ccfs[:,i,:] .= RvSpectML.CCF.calc_ccf_chunklist_timeseries(order_list_order_timeseries, ccf_plan)

     end
 end
 =#
 println("# Order CCFs runtime: ", now()-tstart)

 if write_order_ccf_to_csv
    using CSV
    inst = EXPRES.EXPRES2D()
    for (i, order) in orders_to_use
      if !(sum(order_ccfs[:,i,:]) > 0)   continue    end
      local t = Tables.table( order_ccfs[:,i,:]', header=Symbol.(v_grid) )
      CSV.write(joinpath(output_dir,target_subdir * "_ccf_order=" * string(order) * ".csv"),t)
    end
 end
 has_computed_order_ccfs = true
end

if !has_computed_template
   if verbose println("# Making template spectra.")  end
   @time ( spectral_orders_matrix, f_mean, var_mean, deriv, deriv2, order_grids )  = RvSpectML.make_template_spectra(order_list_timeseries)

   if write_template_to_csv
      using CSV
      CSV.write(joinpath(output_dir,target_subdir * "_template.csv"),DataFrame("λ"=>spectral_orders_matrix.λ,"flux_template"=>f_mean,"var"=>var_mean, "dfluxdlnλ_template"=>deriv,"d²fluxdlnλ²_template"=>deriv2))
   end
   if write_spectral_grid_to_jld2
      using JLD2, FileIO
      save(joinpath(output_dir,target_subdir * "_matrix.jld2"), Dict("λ"=>spectral_orders_matrix.λ,"spectra"=>spectral_orders_matrix.flux,"var_spectra"=>spectral_orders_matrix.var,"flux_template"=>f_mean,"var"=>var_mean, "dfluxdlnλ_template"=>deriv,"d²fluxdlnλ²_template"=>deriv2))
   end
   has_computed_template = true
end

if !has_computed_dcpca
   if verbose println("# Performing Doppler constrained PCA analysis.")  end
   using MultivariateStats
   dcpca_out, M = RvSpectML.DCPCA.doppler_constrained_pca(spectral_orders_matrix.flux, deriv, rvs_ccf_gauss)
   frac_var_explained = 1.0.-cumsum(principalvars(M))./tvar(M)
   println("# Fraction of variance explained = ", frac_var_explained[1:min(5,length(frac_var_explained))])
   if write_dcpca_to_csv
      using CSV
      CSV.write(joinpath(output_dir,target_subdir * "_dcpca_basis.csv"),  Tables.table(M.proj) )
      CSV.write(joinpath(output_dir,target_subdir * "_dcpca_scores.csv"), Tables.table(dcpca_out) )
   end
   has_computed_dcpca = true
end

if make_plots  # Ploting results from DCPCA
  using Plots
  # Set parameters for plotting analysis
  plt_order = 42
  plt_order_pix = 4500:5000
  plt = scatter(frac_var_explained, xlabel="Number of PCs", ylabel="Frac Variance Unexplained")
  display(plt)
end

if make_plots
  plt_order = 13
  RvSpectML.plot_basis_vectors(order_grids, f_mean, deriv, M.proj, idx_plt = spectral_orders_matrix.chunk_map[plt_order], num_basis=3 )
  #xlims!(5761.5,5766)
end

if make_plots
  RvSpectML.plot_basis_scores(order_list_timeseries.times, rvs_ccf_gauss, dcpca_out, num_basis=3 )
end

if make_plots
  RvSpectML.plot_basis_scores_cor( rvs_ccf_gauss, dcpca_out, num_basis=3)
end

has_found_lines = false
write_lines_to_csv= false
if !has_found_lines
   if verbose println("# Performing fresh search for lines in template spectra.")  end
   cl = ChunkList(map(grid->ChuckOfSpectrum(spectral_orders_matrix.λ,f_mean, var_mean, grid), spectral_orders_matrix.chunk_map))
   spectral_orders_matrix = nothing # We're done with this, so can release memory
   GC.gc()
   lines_in_template = RvSpectML.LineFinder.find_lines_in_chunklist(cl)
   if verbose println("# Finding above lines in all spectra.")  end
   @time fits_to_lines = RvSpectML.LineFinder.fit_all_lines_in_chunklist_timeseries(order_list_timeseries, lines_in_template )

   if verbose println("# Rejecting lines that have telluric contamination at any time.")  end
   telluric_info = RvSpectML.LineFinder.find_worst_telluric_in_each_line_fit!(fits_to_lines, order_list_timeseries, expres_data )

   # Look at distribution of standard deviations for line properties
    fit_distrib = fits_to_lines |> @groupby(_.line_id) |>
            @map( { median_a=median(_.fit_a), median_b=median(_.fit_b), median_depth=median(_.fit_depth), median_σ²=median(_.fit_σ²), median_λc=median(_.fit_λc),
                   std_a=std(_.fit_a), std_b=std(_.fit_b), std_depth=std(_.fit_depth), std_σ²=std(_.fit_σ²), std_λc=std(_.fit_λc), min_telluric_model_all_obs=minimum(_.min_telluric_model_this_obs), line_id=_.line_id } ) |>
            @filter(_.min_telluric_model_all_obs == 1.0 ) |> # No telluric in any observations
            @filter(_.std_b < 0.257) |>     # ~90th quantile
            @filter( 0.001 < _.median_σ² ) |>  # ~99th quantile
            @filter( _.std_σ² < 0.00844) |>  # ~99th quantile
            @filter( _.median_depth > 0.05 ) |>  # ~99th quantile
            @filter( _.std_depth < 0.11 ) |>  # ~95th quantile
            #@filter( -0.25 < _.median_b < 0.25) |>  #~5th to 95th percentiles
            DataFrame # ~75th quantile
      good_lines = fit_distrib |> @filter(_.std_λc < 0.0316 ) |> DataFrame
      bad_lines = fit_distrib |> @filter(_.std_λc > 0.0316 ) |> DataFrame
      #scatter(good_lines.median_σ², good_lines.median_depth )
      #scatter!(bad_lines.median_σ², bad_lines.median_depth )
      lines_to_try = lines_in_template[first.(good_lines[!,:line_id]),:]

      if write_lines_to_csv
         using CSV
         CSV.write(joinpath(output_dir,target_subdir * "_linefinder_lines.csv"), lines_in_template )
         CSV.write(joinpath(output_dir,target_subdir * "_linefinder_line_fits.csv"), fits_to_lines )
         CSV.write(joinpath(output_dir,target_subdir * "_linefinder_line_fits_clean.csv"), lines_to_try )
         fits_to_lines = nothing
         telluric_info = nothing
         fit_distrib = nothing
         good_lines = nothing
         bad_lines = nothing
         GC.gc()
      end

   has_found_lines = true
end

if has_computed_ccfs2
   if verbose println("# Computing CCFs with new line list.")  end
   #mask_shape = RvSpectML.CCF.TopHatCCFMask(order_list_timeseries.inst, scale_factor=tophap_ccf_mask_scale_factor)
   perm = sortperm(lines_to_try.fit_λc)
   line_list2 = RvSpectML.CCF.BasicLineList(lines_to_try.fit_λc[perm], lines_to_try.fit_depth[perm].*lines_to_try.fit_a[perm] )
   ccf_plan2 = RvSpectML.CCF.BasicCCFPlan(mask_shape = mask_shape, line_list=line_list2, midpoint=0.0)
   v_grid2 = RvSpectML.CCF.calc_ccf_v_grid(ccf_plan2)
   @time ccfs2 = RvSpectML.CCF.calc_ccf_chunklist_timeseries(order_list_timeseries, ccf_plan2)
   # Write CCFs to file
   if write_ccf_to_csv
      using CSV
      CSV.write(joinpath(output_dir,target_subdir * "_ccfs2.csv"),Tables.table(ccfs2',header=Symbol.(v_grid)))
   end
   has_computed_ccfs2 = true
end

if !has_computed_rvs2
 if verbose println("# Measuring RVs from CCF.")  end
 #rvs_ccf_gauss = [ RvSpectML.RVFromCCF.measure_rv_from_ccf(v_grid,ccfs[:,i],fit_type = "gaussian") for i in 1:length(order_list_timeseries) ]
 rvs_ccf_gauss2 = RvSpectML.RVFromCCF.measure_rv_from_ccf(v_grid2,ccfs2,fit_type = "gaussian")
 # Store estimated RVs in metadata
 map(i->order_list_timeseries.metadata[i][:rv_est] = rvs_ccf_gauss2[i]-mean(rvs_ccf_gauss2), 1:length(order_list_timeseries) )
 if write_rvs_to_csv
    using CSV
    CSV.write(joinpath(output_dir,target_subdir * "_rvs_ccf2.csv"),DataFrame("Time [MJD]"=>order_list_timeseries.times,"CCF RV [m/s]"=>rvs_ccf_gauss))
 end
 has_computed_rvs2 = true
end

if !has_computed_template2
   # Do we need to expand region around lines to include more of wings for DCPCA to work well?
   # If so, should probably avoid tellurics encroaching on exapnded wavelength region.  Sigh.
   chunk_list_df2 = lines_to_try |> @select(:fit_min_λ,:fit_max_λ) |> @rename(:fit_min_λ=>:lambda_lo, :fit_max_λ=>:lambda_hi) |> DataFrame
   chunk_list_timeseries2 = RvSpectML.make_chunk_list_timeseries(expres_data,chunk_list_df2)

   # Check that no NaN's included
   (chunk_list_timeseries2, chunk_list_df2) = RvSpectML.filter_bad_chunks(chunk_list_timeseries2,chunk_list_df2)
   println(size(chunk_list_df2), " vs ", num_chunks(chunk_list_timeseries2) )

   if verbose println("# Making template spectra.")  end
   @time ( spectral_orders_matrix2, f_mean2, var_mean2, deriv_2, deriv2_2, order_grids2 )  = RvSpectML.make_template_spectra(chunk_list_timeseries2)

   if write_template_to_csv
      using CSV
      CSV.write(joinpath(output_dir,target_subdir * "_template2.csv"),DataFrame("λ"=>spectral_orders_matrix2.λ,"flux_template"=>f_mean2,"var"=>var_mean2, "dfluxdlnλ_template"=>deriv_2,"d²fluxdlnλ²_template"=>deriv2_2))
   end
   if write_spectral_grid_to_jld2
      using JLD2, FileIO
      save(joinpath(output_dir,target_subdir * "_matrix2.jld2"), Dict("λ"=>spectral_orders_matrix2.λ,"spectra"=>spectral_orders_matrix2.flux,"var_spectra"=>spectral_orders_matrix2.var,"flux_template"=>f_mean2,"var"=>var_mean2, "dfluxdlnλ_template"=>deriv_2,"d²fluxdlnλ²_template"=>deriv2_2))
   end
   has_computed_template2 = true
end

has_computed_dcpca2 = false
if !has_computed_dcpca2
   if verbose println("# Performing Doppler constrained PCA analysis.")  end
   using MultivariateStats
   dcpca2_out, M2 = RvSpectML.DCPCA.doppler_constrained_pca(spectral_orders_matrix2.flux, deriv_2, rvs_ccf_gauss2)
   frac_var_explained2 = 1.0.-cumsum(principalvars(M2))./tvar(M2)
   println("# Fraction of variance explained = ", frac_var_explained2[1:min(5,length(frac_var_explained2))])
   if write_dcpca_to_csv
      using CSV
      CSV.write(joinpath(output_dir,target_subdir * "_dcpca_basis2.csv"),  Tables.table(M.proj) )
      CSV.write(joinpath(output_dir,target_subdir * "_dcpca_scores2.csv"), Tables.table(dcpca_out) )
   end
   has_computed_dcpca2 = true
end

if make_plots  # Ploting results from DCPCA
  using Plots
  # Set parameters for plotting analysis
  plt_order = 42
  plt_order_pix = 4500:5000
  plt = scatter(frac_var_explained, xlabel="Number of PCs", ylabel="Frac Variance Unexplained")
  display(plt)
end

if make_plots
  plt_line = 170
  RvSpectML.plot_basis_vectors(order_grids, f_mean, deriv, M.proj, idx_plt = spectral_orders_matrix.chunk_map[plt_line], num_basis=3 )
  #xlims!(5761.5,5766)
end

if make_plots
  RvSpectML.plot_basis_scores(order_list_timeseries.times, rvs_ccf_gauss, dcpca_out, num_basis=3 )
end

if make_plots
  RvSpectML.plot_basis_scores_cor( rvs_ccf_gauss, dcpca_out, num_basis=3)
end

has_loaded_data = false
 has_computed_ccfs = false
 has_computed_rvs = false
 has_computed_template = false
 has_computed_dcpca = false
 has_found_lines = false


# nextpow2(W) = vshl(one(W), vsub_fast(8sizeof(W), leading_zeros(vsub_fast(W, one(W)))))

# @inline _pick_vector(::StaticInt{W}, ::Type{T}) where {W,T} = Vec{W,T}
# @inline pick_vector(::Type{T}) where {T} = _pick_vector(pick_vector_width(T), T)
# @inline function pick_vector(::Val{N}, ::Type{T}) where {N, T}
#     _pick_vector(smin(nextpow2(StaticInt{N}()), pick_vector_width(T)), T)
# end
# pick_vector(N::Int, ::Type{T}) where {T} = pick_vector(Val(N), T)

@inline MM(::Union{Val{W},StaticInt{W}}) where {W} = MM{W}(0)
@inline MM(::Union{Val{W},StaticInt{W}}, i) where {W} = MM{W}(i)
@inline MM(::Union{Val{W},StaticInt{W}}, i::AbstractSIMDVector{W}) where {W} = i
@inline MM(::StaticInt{W}, i, ::StaticInt{X}) where {W,X} = MM{W,X}(i)
@inline gep(ptr::Ptr, i::MM) = gep(ptr, data(i))

@inline Base.one(::Type{MM{W,X,I}}) where {W,X,I} = one(I)
@inline staticm1(i::MM{W,X,I}) where {W,X,I} = MM{W,X}(vsub_fast(data(i), one(I)))
@inline staticp1(i::MM{W,X,I}) where {W,X,I} = MM{W,X}(vadd_nsw(data(i), one(I)))
@inline vadd_fast(i::MM{W,X}, j::IntegerTypesHW) where {W,X} = MM{W,X}(vadd_fast(data(i), j))
@inline vadd_fast(i::IntegerTypesHW, j::MM{W,X}) where {W,X} = MM{W,X}(vadd_fast(i, data(j)))
@inline vadd_fast(i::MM{W,X}, ::StaticInt{j}) where {W,X,j} = MM{W,X}(vadd_fast(data(i), StaticInt{j}()))
@inline vadd_fast(::StaticInt{i}, j::MM{W,X}) where {W,X,i} = MM{W,X}(vadd_fast(StaticInt{i}(), data(j)))
@inline vadd_fast(i::MM{W,X}, ::StaticInt{0}) where {W,X} = i
@inline vadd_fast(::StaticInt{0}, j::MM{W,X}) where {W,X} = j
@inline vsub_fast(i::MM{W,X}, j::IntegerTypesHW) where {W,X} = MM{W,X}(vsub_fast(data(i), j))
@inline vsub_fast(i::MM{W,X}, ::StaticInt{j}) where {W,X,j} = MM{W,X}(vsub_fast(data(i), StaticInt{j}()))
@inline vsub_fast(i::MM{W,X}, ::StaticInt{0}) where {W,X} = i

@inline vadd_nsw(i::MM{W,X}, j::IntegerTypesHW) where {W,X} = MM{W,X}(vadd_nsw(data(i), j))
@inline vadd_nsw(i::IntegerTypesHW, j::MM{W,X}) where {W,X} = MM{W,X}(vadd_nsw(i, data(j)))
@inline vadd_nsw(i::MM{W,X}, ::StaticInt{j}) where {W,X,j} = MM{W,X}(vadd_nsw(data(i), StaticInt{j}()))
@inline vadd_nsw(i::MM, ::Zero) = i
@inline vadd_nsw(::StaticInt{i}, j::MM{W,X}) where {W,X,i} = MM{W,X}(vadd_nsw(StaticInt{i}(), data(j)))
@inline vadd_nsw(::Zero, j::MM{W,X}) where {W,X} = j
@inline vsub_nsw(i::MM{W,X}, j::IntegerTypesHW) where {W,X} = MM{W,X}(vsub_nsw(data(i), j))
@inline vsub_nsw(i::MM{W,X}, ::StaticInt{j}) where {W,X,j} = MM{W,X}(vsub_nsw(data(i), StaticInt{j}()))
@inline vsub(i::MM{W,X}, ::StaticInt{0}) where {W,X} = i
@inline vsub_fast(i::MM) = i * StaticInt{-1}()
@inline vsub_nsw(i::MM) = i * StaticInt{-1}()
@inline vsub(i::MM) = i * StaticInt{-1}()
@inline vadd(i::MM, j::IntegerTypesHW) = vadd_fast(i, j)
@inline vadd(j::IntegerTypesHW, i::MM) = vadd_fast(j, i)
@inline vsub(i::MM, j::IntegerTypesHW) = vsub_fast(i, j)
@inline vsub(j::IntegerTypesHW, i::MM) = vsub_fast(j, i)
@inline vadd(i::MM, ::StaticInt{j}) where {j} = vadd_fast(i, StaticInt{j}())
@inline vadd(::StaticInt{j}, i::MM) where {j} = vadd_fast(StaticInt{j}(), i)
@inline vsub(i::MM, ::StaticInt{j}) where {j} = vsub_fast(i, StaticInt{j}())
@inline vsub(::StaticInt{j}, i::MM) where {j} = vsub_fast(StaticInt{j}(), i)
@inline vadd(i::MM, ::Zero) = i
@inline vadd(::Zero, i::MM) = i
@inline vsub(::Zero, i::MM) = StaticInt{-1}()*i

@inline vmul_nsw(::StaticInt{M}, i::MM{W,X}) where {M,W,X} = MM{W}(vmul_nsw(data(i), StaticInt{M}()), StaticInt{X}() * StaticInt{M}())
@inline vmul_nsw(i::MM{W,X}, ::StaticInt{M}) where {M,W,X} = MM{W}(vmul_nsw(data(i), StaticInt{M}()), StaticInt{X}() * StaticInt{M}())

@inline vmul_fast(::StaticInt{M}, i::MM{W,X}) where {M,W,X} = MM{W}(vmul_fast(data(i), StaticInt{M}()), StaticInt{X}() * StaticInt{M}())
@inline vmul_fast(i::MM{W,X}, ::StaticInt{M}) where {M,W,X} = MM{W}(vmul_fast(data(i), StaticInt{M}()), StaticInt{X}() * StaticInt{M}())
@inline vmul(a, ::StaticInt{N}) where {N} = vmul_fast(a, StaticInt{N}())
@inline vmul(::StaticInt{N}, a) where {N} = vmul_fast(StaticInt{N}(), a)
@inline vmul(::StaticInt{N}, ::StaticInt{M}) where {N,M} = StaticInt{N}() * StaticInt{M}()

@inline vsub(a, ::StaticInt{N}) where {N} = vsub_fast(a, StaticInt{N}())
@inline vadd(a, ::StaticInt{N}) where {N} = vadd_fast(a, StaticInt{N}())
@inline vsub(::StaticInt{N}, a) where {N} = vsub_fast(StaticInt{N}(), a)
@inline vadd(::StaticInt{N}, a) where {N} = vadd_fast(StaticInt{N}(), a)
@inline vsub(::StaticInt{M}, ::StaticInt{N}) where {M,N} = StaticInt{M}() - StaticInt{N}()
@inline vadd(::StaticInt{M}, ::StaticInt{N}) where {M,N} = StaticInt{M}() + StaticInt{N}()

@inline vrem(i::MM{W,X,I}, ::Type{I}) where {W,X,I<:IntegerTypesHW} = i
@inline vrem(i::MM{W,X}, ::Type{I}) where {W,X,I<:IntegerTypesHW} = MM{W,X}(data(i) % I)
@inline veq(::AbstractIrrational, ::MM{W,<:Integer}) where {W} = zero(Mask{W})
@inline veq(x::AbstractIrrational, i::MM{W}) where {W} = x == Vec(i)
@inline veq(::MM{W,<:Integer}, ::AbstractIrrational) where {W} = zero(Mask{W})
@inline veq(i::MM{W}, x::AbstractIrrational) where {W} = Vec(i) == x

@inline vsub_nsw(i::MM, ::Zero) = i
@inline vsub_nsw(i::NativeTypes, j::MM{W,X}) where {W,X} = MM(StaticInt{W}(), vsub_nsw(i, data(j)), -StaticInt{X}())
@inline vsub_fast(i::NativeTypes, j::MM{W,X}) where {W,X} = MM(StaticInt{W}(), vsub_fast(i, data(j)), -StaticInt{X}())
@inline vsub_nsw(i::Union{FloatingTypes,IntegerTypesHW}, j::MM{W,X}) where {W,X} = MM(StaticInt{W}(), vsub_nsw(i, data(j)), -StaticInt{X}())
@inline vsub_fast(i::Union{FloatingTypes,IntegerTypesHW}, j::MM{W,X}) where {W,X} = MM(StaticInt{W}(), vsub_fast(i, data(j)), -StaticInt{X}())

@inline function Base.in(m::MM{W,X,<:Integer}, r::AbstractUnitRange) where {W,X}
    vm = Vec(m)
    (vm ≥ first(r)) & (vm ≤ last(r))
end

for op ∈ (:(+),:(-))
  @eval begin
    @inline Base.$op(vu::VecUnroll{N,1,T,T}, i::MM) where {N,T<:NativeTypes} = VecUnroll(fmap($op, data(vu), i))
    @inline Base.$op(i::MM, vu::VecUnroll{N,1,T,T}) where {N,T<:NativeTypes} = VecUnroll(fmap($op, i, data(vu)))
    @inline Base.$op(vu::VecUnroll{N,1,T,T}, i::VecUnroll{N,W,I,MM{W,X,I}}) where {N,W,T<:NativeTypes,I,X} = VecUnroll(fmap($op, data(vu), data(i)))
    @inline Base.$op(i::VecUnroll{N,W,I,MM{W,X,I}}, vu::VecUnroll{N,1,T,T}) where {N,W,T<:NativeTypes,I,X} = VecUnroll(fmap($op, data(i), data(vu)))
  end
end
# @inline Base.:(+)(vu::VecUnroll{N,W,T,T}, i::MM) = VecUnroll(fmap(+, data(vu), i))


devices = Dict{Symbol, DeviceModel}(:Generators => DeviceModel(PSY.ThermalStandard, ThermalDispatch),
                                    :Loads =>  DeviceModel(PSY.PowerLoad, StaticPowerLoad))
branches = Dict{Symbol, DeviceModel}(:L => DeviceModel(PSY.Line, StaticLineUnbounded))
services = Dict{Symbol, ServiceModel}()

@testset "Operation set ref models" begin
    template = OperationsTemplate(CopperPlatePowerModel, devices, branches, services);
    op_problem = OperationsProblem(TestOpProblem, template, c_sys5)
    set_transmission_model!(op_problem, DCPLLPowerModel)
    @test op_problem.template.transmission == DCPLLPowerModel

    template = OperationsTemplate(CopperPlatePowerModel, devices, branches, services);
    op_problem = OperationsProblem(TestOpProblem, template, c_sys5)
    new_devices = Dict{Symbol, DeviceModel}(:Generators => DeviceModel(PSY.ThermalStandard, ThermalBasicUnitCommitment),
                                            :Loads =>  DeviceModel(PSY.PowerLoad, StaticPowerLoad))
    set_devices_template!(op_problem, new_devices)
    @test op_problem.template.devices[:Generators].formulation == ThermalBasicUnitCommitment
    jump_model = op_problem.canonical.JuMPmodel
    @test ((JuMP.VariableRef, MOI.ZeroOne) in JuMP.list_of_constraint_types(jump_model)) == true

    template = OperationsTemplate(DCPPowerModel, devices, branches, services);
    op_problem = OperationsProblem(TestOpProblem, template, c_sys5)
    new_branches = Dict{Symbol, DeviceModel}(:L => DeviceModel(PSY.Line, StaticLine))
    set_branches_template!(op_problem, new_branches)
    @test op_problem.template.branches[:L].formulation == StaticLine
end

@testset "Operation set models" begin
    template = OperationsTemplate(CopperPlatePowerModel, devices, branches, services);
    op_problem = OperationsProblem(TestOpProblem, template, c_sys5)
    set_device_model!(op_problem, :Generators, DeviceModel(PSY.ThermalStandard, ThermalBasicUnitCommitment))
    @test op_problem.template.devices[:Generators].formulation == ThermalBasicUnitCommitment
    jump_model = op_problem.canonical.JuMPmodel
    @test ((JuMP.VariableRef, MOI.ZeroOne) in JuMP.list_of_constraint_types(jump_model)) == true

    template = OperationsTemplate(DCPPowerModel, devices, branches, services);
    op_problem = OperationsProblem(TestOpProblem, template, c_sys5)
    set_branch_model!(op_problem, :L, DeviceModel(PSY.Line, StaticLine))
    @test op_problem.template.branches[:L].formulation == StaticLine
end


# Quelques fonctions utilitaires pour les manipulations de solvers PL ou MIP

# Classe abstraite non utilisé
# abstract type AbstractLpSolver end

"""
Création d'un modèle de solver externe (:cplex, :glpk, ...)

Args:
    solver : symbole :cplex, :glpk, :clp, :cbc

        - :cbc : permet de faire du MIP
        - :clp : resteint aux problème LP mais plus efficace de cbc pour les relaxations.

    mode : (:lp ou :mip) utile pour choisir entre :clp ou :cbc
    log_level (expérimental) : entier 0 pour silence, sinon 1, 2, ...

Autre solver à supporter plus tard Gurobi (concurrent de cplex)

Attention:
    DEPENDANCE EXTERNE : module Args (pour Args.get("external_mip_solver"))
"""
function new_model(;solver = :auto,
                     mode = :lp,
                     log_level = 0, # EXPLOITATION PARTIELLE
                     )
    if !(mode in [:mip, :lp])
        error("ERROR: unknown mode $(mode). Should be :cplex, :glpk, ...")
    end
    if solver == :auto
        solver = Args.get("external_mip_solver")
    end
    if solver in [:clp, :cbc]
        # solver = (mode == :mip ? :cbc : :clp )
        if mode==:mip
            solver = :cbc
        else
            solver = :clp
        end
    end
    if ! (solver in [:cplex, :clp, :cbc, :glpk])
        error("ERROR: unknown solver $(solver). Should be :cplex, :glpk, ...")
    end

    if solver == :glpk
        # voir https://github.com/JuliaOpt/GLPK.jl
        # model = Model(GLPK.Optimizer)  # PAS ENCORE SUPPORTE (05/12/2019)
        model = Model(with_optimizer(GLPK.Optimizer))
    elseif solver == :cplex
        model = Model(with_optimizer(CPLEX.Optimizer))
        # model = Model(CPLEX.Optimizer) # PAS ENCORE SUPPORTE (05/12/2019)
    elseif solver == :clp
        # voir https://github.com/JuliaOpt/Clp.jl
        # model = Model(Clp.Optimizer) # PAS ENCORE SUPPORTE (05/12/2019)
        model = Model(with_optimizer(Clp.Optimizer))
    elseif solver == :cbc
        # Voir https://github.com/JuliaOpt/Cbc.jl pour les options
        # model = Model(Cbc.Optimizer) # PAS ENCORE SUPPORTE (05/12/2019)
        model = Model(with_optimizer(Cbc.Optimizer))
    else
        error("LpTimingSolver: unknown external_mip_solver inconnu : $solver")
    end
    # MOI.set(model, MOI.Silent(), true) # Variante de bas niveau
    if log_level <= 3
        set_silent(model) # unset_silent(model) pour réautoriser affichage
    end
    return model
end


using DSGE, Test

# Initialize model object
m = AnSchorfheide(testing = true)
m <= Setting(:forecast_horizons, 2)

# Scenario
scen = Scenario(:testscen, "Test Scenario", [:obs_gdp, :obs_cpi], [:g_sh, :rm_sh], "REF")
@testset "Test scenario basic constructor" begin
    @test scen.shock_scaling == 1.0
    @test n_targets(scen) == 2
    @test n_instruments(scen) == 2

    for var in scen.target_names
        scen.targets[!, var] = rand(2)
    end
    @test n_target_horizons(scen) == 2

    global act = targets_to_data(m, scen)
    global expt = copy(scen.targets)
    expt[!, :date] = DSGE.quarter_range(date_forecast_start(m), date_forecast_end(m))
    expt[!, :obs_nominalrate] = fill(missing, 2)
    for col in names(expt)
        global check = false
        if ismissing(act[!, col] == expt[!, col])
            check = (all(ismissing.(act[!, col])) && all(ismissing.(expt[!, col])))
        else
            check = (act[!, col] == expt[!, col])
        end
        @test check
    end
end

# Scenario with shock scaling
@testset "Test scenario constructor with shock scaling" begin
    scale = Scenario(:scaledscen, "Test Shock Scaling Scenario", [:obs_gdp, :obs_cpi], [:g_sh, :rm_sh], "REF",
                     shock_scaling = 2.0)
    @test scale.shock_scaling == 2.0

    # SwitchingScenario
    @test_throws ErrorException SwitchingScenario(:switch, scen, scen, [1.0], [1.0, 1.0])
    @test_throws ErrorException SwitchingScenario(:switch, scen, scen, [-1.0, 1.0], [1.0, 1.0])
    @test_throws ErrorException SwitchingScenario(:switch, scen, scen, [1.0, 1.0], [-1.0, 1.0])

    # ScenarioAggregate
    @test_throws ErrorException ScenarioAggregate(:agg, "Test Scenario Aggregate", AbstractScenario[scen],
                                                  [0.5, 0.5], 10, true, "REF")
    @test_throws ErrorException ScenarioAggregate(:agg, "Test Scenario Aggregate", AbstractScenario[scen],
                                                  [-1.0], 10, true, "REF")
    @test_throws ErrorException ScenarioAggregate(:agg, "Test Scenario Aggregate", AbstractScenario[scen],
                                                  [0.5], 10, true, "REF")
end


# Indexing and dimensions (B.4)

export
    threadIdx, blockDim, blockIdx, gridDim,
    warpsize

@generated function _index(::Val{name}, ::Val{range}) where {name, range}
    JuliaContext() do ctx
        T_int32 = LLVM.Int32Type(ctx)

        # create function
        llvm_f, _ = create_function(T_int32)
        mod = LLVM.parent(llvm_f)

        # generate IR
        Builder(ctx) do builder
            entry = BasicBlock(llvm_f, "entry", ctx)
            position!(builder, entry)

            # call the indexing intrinsic
            intr_typ = LLVM.FunctionType(T_int32)
            intr = LLVM.Function(mod, "llvm.nvvm.read.ptx.sreg.$name", intr_typ)
            idx = call!(builder, intr)

            # attach range metadata
            range_metadata = MDNode([ConstantInt(Int32(range.start), ctx),
                                    ConstantInt(Int32(range.stop), ctx)],
                                    ctx)
            metadata(idx)[LLVM.MD_range] = range_metadata

            ret!(builder, idx)
        end

        call_function(llvm_f, UInt32)
    end
end

# XXX: these depend on the compute capability
#      https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#compute-capabilities
const max_block_size = (x=1024, y=1024, z=64)
const max_grid_size  = (x=2^31-1, y=65535, z=65535)

for dim in (:x, :y, :z)
    # Thread index
    fn = Symbol("threadIdx_$dim")
    intr = Symbol("tid.$dim")
    @eval @inline $fn() = Int(_index($(Val(intr)), $(Val(0:max_block_size[dim]-1)))) + 1

    # Block size (#threads per block)
    fn = Symbol("blockDim_$dim")
    intr = Symbol("ntid.$dim")
    @eval @inline $fn() = Int(_index($(Val(intr)), $(Val(1:max_block_size[dim]))))

    # Block index
    fn = Symbol("blockIdx_$dim")
    intr = Symbol("ctaid.$dim")
    @eval @inline $fn() = Int(_index($(Val(intr)), $(Val(0:max_grid_size[dim]-1)))) + 1

    # Grid size (#blocks per grid)
    fn = Symbol("gridDim_$dim")
    intr = Symbol("nctaid.$dim")
    @eval @inline $fn() = Int(_index($(Val(intr)), $(Val(1:max_grid_size[dim]))))
end

"""
    gridDim()::CuDim3

Returns the dimensions of the grid.
"""
@inline gridDim() =   (x=gridDim_x(),   y=gridDim_y(),   z=gridDim_z())

"""
    blockIdx()::CuDim3

Returns the block index within the grid.
"""
@inline blockIdx() =  (x=blockIdx_x(),  y=blockIdx_y(),  z=blockIdx_z())

"""
    blockDim()::CuDim3

Returns the dimensions of the block.
"""
@inline blockDim() =  (x=blockDim_x(),  y=blockDim_y(),  z=blockDim_z())

"""
    threadIdx()::CuDim3

Returns the thread index within the block.
"""
@inline threadIdx() = (x=threadIdx_x(), y=threadIdx_y(), z=threadIdx_z())

"""
    warpsize()::UInt32

Returns the warp size (in threads).
"""
@inline warpsize() = Int(ccall("llvm.nvvm.read.ptx.sreg.warpsize", llvmcall, UInt32, ()))


import AbstractPlotting.convert_arguments

# bin centers
function convert_arguments(P::Type{<:Scatter}, B::SSHBinner)
    convert_arguments(P, bin_positions(B))
end
# mesh with vertices @ bin centers (makes color mapping easy)
function convert_arguments(P::Type{<:AbstractPlotting.Mesh}, B::SSHBinner)
    convert_arguments(P, vertex_mesh(B))
end
# boundaries of bins (slightly enlarged to work with mesh)
function convert_arguments(P::Type{<:LineSegments}, B::SSHBinner)
    convert_arguments(P, 1.001line_segments(B))
end

# Not sure if I wanna keep this. Would probably be nice to have a colorlegend
# by default here
"""
    histogram(B::SSHBinner)
    histogram!(scene, B::SSHBinner)

Plots a histogram for the given SSHBinner.

# Attributes

- `outline = false`: Draws lines around each bin
- `colormap = :matter`: Colormap used to visualize bins
- `colorscale = :absolute`: Scale used when mapping bins to colors. Can be
:relative (sets colorrange = extrema(B.bins)), :absolute (sets only the max) or
:static (colorrange remains unchanged)
- `colorrange = [0, maxiumum(b.bins)]`: Range of bin fillings represented by the
colormap.
- `transparency = false`: Set this to true if your colormap is tranparent
- `linecolor = :black`: Color of outline.
"""
@recipe(Histogram, B) do scene
    Theme(
        outline = false,
        colormap = :matter,
        transparency = false,
        linecolor = :black,
        colorscale = :absolute,
        colorrange = (0, 1)
    )
end

function map_bin_to_color(binner, colormap, colorscale, colorrange)
    map(binner, colormap, colorscale) do binner, cm, scale
        ex = extrema(binner.bins)
        if scale == :absolute
            colorrange[] = (0, ex[2])
        elseif scale == :static
        else
            colorrange[] = (ex[1], ex[2])
        end
        map(binner.bins) do b
            AbstractPlotting.RGBAf0(AbstractPlotting.AbstractPlotting.interpolated_getindex(
                to_colormap(cm), Float64(b), colorrange[]
            ))
        end
    end
end

function AbstractPlotting.plot!(plot::Histogram)
    colors = map_bin_to_color(
        plot[:B], plot[:colormap], plot[:colorscale], plot[:colorrange]
    )

    mesh!(plot, plot[:B], color = colors, transparency = plot[:transparency])
    linesegments!(plot, plot[:B], color = plot[:linecolor], visible=plot[:outline])

    plot
end

export histogram, histogram!

# """
#     histogram(B::SSHBinner[, color_method = to_hue])
#
# Plots a sphere where each bin filling is mapped to a color through color_method.
# """
# function histogram(B::SSHBinner, color_method=to_hue)
#     _colors = color_method(B)
#     is_transparent = any(c -> c.alpha < 1.0, _colors)
#     AbstractPlotting.mesh(vertex_mesh(B), color = _colors, transparency = is_transparent)
# end
#
#
# """
#     plot_debug(B::SSHBinner[, color_method = to_hue])
#
# Plots a sphere with colors representing bin fillings, approximate bin area
# (rectangles) and points at the bin centers.
# """
# function plot_debug(B::SSHBinner, color_method = to_hue)
    # _colors = color_method(B)
    # is_transparent = any(c -> c.alpha < 1.0, _colors)
    # scene = AbstractPlotting.mesh(vertex_mesh(B), color = _colors, transparency = is_transparent)
    # AbstractPlotting.linesegments!(scene, 1.001line_segments(B), color = :black)
    # # AbstractPlotting.mesh!(scene1, to_rects(B, extrude), color = RGBA(1, 1, 1, 0.2), transparency = true)
    # AbstractPlotting.scatter!(scene, 1.005bin_positions(B), color=:black, markersize=0.01)
    # scene
# end


# This file is a part of BAT.jl, licensed under the MIT License (MIT).

# bring vsel to a normalized shape, i.e.:
#
# AbstractVector{<:Integer} -> AbstractVector{<:Integer}
# AbstractVector{AbstractUnitRange} -> AbstractVector{AbstractUnitRange}
#
# Expr -> Vector{Union{Symbol,Expr})
#
# Symbol -> Vector{Union{Symbol,Expr})

function _normalize_vsel_unshaped(vsel::Union{AbstractVector{T}, AbstractVector{UnitRange{T}}}) where T <: Integer
    return vsel
end

function _normalize_vsel_unshaped(vsel::Integer) where T <: Integer
    return Int[vsel]
end

function _normalize_vsel_unshaped(vsel::AbstractVector)
    @argcheck all(broadcast(x -> x isa Union{T, AbstractVector{T}, AbstractVector{UnitRange{T}}} where T <: Integer, vsel)) throw(ArgumentError("Samples or distribution are unshaped, please use Integer, Vector of Integers, or UnitRange for indexing."))
    return vsel
end

function _normalize_vsel_shaped(vs::AbstractValueShape, vsel::AbstractVector{UnitRange})
    @argcheck length(vsel) == 1 throw(ArgumentError("Can't use multidimensional cartesian index for shaped samples or distribution, please use keys instead"))
    return [key for key in keys(vs)[vsel[1]]]
end

function _normalize_vsel_shaped(vs::AbstractValueShape, vsel::Integer)
    return [keys(vs)[vsel]]
end

function _normalize_vsel_shaped(vs::AbstractValueShape, vsel::AbstractVector{Integer})
    return [keys(vs)[i] for i in vsel]
end

function _normalize_vsel_shaped(vs::AbstractValueShape, vsel::Expr)
    if isexpr(vsel, :vect) || isexpr(vsel, :call)
        if @capture(vsel, [s_:e_]) || @capture(vsel, (s_:e_))
            return [keys(vs)[i] for i in s:e]

        elseif @capture(vsel, [dims__]) && all(broadcast(x -> x isa Integer, dims))
            return [keys(vs)[i] for i in dims]
        else
            vsel_norm = [_normalize_vsel_shaped(vs, el isa QuoteNode ? el.value : el) for el in vsel.args]
            return vcat(vsel_norm...) # for user convenience, does :([a[1,2], b, c]) -> [:(a[1,2]), :b, :c]
        end
    end
    return [vsel]
end

function _normalize_vsel_shaped(vs::AbstractValueShape, vsel::Symbol)
    return [vsel]
end

function _normalize_vsel_shaped(vs::AbstractValueShape, vsel::AbstractVector)
    vsel_norm = [_normalize_vsel_shaped(vs, el isa QuoteNode ? el.value : el) for el in vsel]
    return vcat(vsel_norm...)
end

# Retrieve flat indexes of the selected dims of samples
# If shaped, build new shape

function marg_idxs_unshaped(samples::DensitySampleVector,
                            vsel::AbstractVector
)
    flat_v = ValueShapes.flatview(samples.v)
    return (firstindex(flat_v) - 1) .+ vsel
end

function marg_idxs_shaped(samples::DensitySampleVector,
                         vsel::AbstractVector
)
    shapes = []
    sel_idxs = Int[]

    vs = varshape(samples)
    flat_idxs = axes(ValueShapes.flatview(unshaped.(samples).v), 1)
    shaped_idxs = vs(flat_idxs)

    for el in vsel

        # :a
        if el isa Symbol
            shape_acc = getproperty(vs, el)
            idx_acc = getproperty(shaped_idxs, el)

            shape = shape_acc.shape
            append!(sel_idxs, idx_acc isa Integer ? Int64[idx_acc] : vec(idx_acc))

        # :(a[1]) or :(a[1:4]) or :(a[1,2,3]) or :(a[1,2,:])
        elseif @capture(el, a_[args__])
            shape_acc = getproperty(vs, a)
            idx_acc = getproperty(shaped_idxs, a)

            if @capture(el, a_[s_:e_])
                shape = ArrayShape{Real}(length(s:e))
                append!(sel_idxs, vec(idx_acc[s:e]))
            elseif @capture(el, a_[dim_])
                shape = ScalarShape{Real}()
                append!(sel_idxs, idx_acc[dim])
            else
                slice_dims = Int[]

                for (i,arg) in enumerate(args)
                    if @capture(arg, (s_:e_))

                        args[i] = s:e
                        append!(slice_dims, length(s:e))
                    else

                        append!(slice_dims, 0)
                    end
                end

                sliced_dims = slice_dims .> 0

                whole_dims = args .== :(:)
                args[whole_dims] .= Colon()

                new_dims = shape_acc.shape.dims .* whole_dims .+ slice_dims

                new_axes = new_dims[whole_dims .| sliced_dims]

                shape = length(new_axes) > 0 ? ArrayShape{Real}(new_axes...) : ScalarShape{Real}()

                idxs = idx_acc[args...]
                append!(sel_idxs, ndims(idxs) > 0 ? vec(idxs) : idxs)

            end
        end
        el = encode_name(el)
        append!(shapes, (el = shape,))
    end
    vsel = encode_name.(vsel)
    new_shape = NamedTupleShape(NamedTuple(vsel .=> shapes))

    return sel_idxs, new_shape
end

# Encode dim names to preserve indexing information
#
# unicode encodings to preserve indexing information:
# :a⌞2⌟     = : a \llcorner 2 \lrcorner
# :a⌞2ː4⌟   = : a \llcorner 2 \lmrk 4 \lrcorner
# :a⌞1ˌ3ˌ5⌟ = : a \llcorner 1 \verti 3 \verti 5 \lrcorner

function encode_name(name::Expr)
    code = replace(string(name)," " => "", "[" => "⌞", ":" => "ː", "," => "ˌ", "]" => "⌟")
    return Symbol(code)
end

function encode_name(name::Symbol)
    return name
end
#=
"""
    bat_marginalize(
        samples::DensitySampleVector,
        vsel::Union{E, S, I, U, AbstractVector{Union{E, S, I, U}}} where E<:Expr where S<:Symbol where I<:Integer where U<:UnitRange
    )

Create a new `DensitySampleVector` containing the selected
dimensions `vsel` of the input `samples`.

Returns a `DensitySampleVector`. If the input is shaped, the new samples
will have the shape of the selected dimensions.

`samples` may be shaped or unshaped, vsel may be an `Int`, `UnitRange`,
`Symbol`, `Expression`, or a Vector of these types.

To select certain elements or whole slices of a multidimensional parameter `a` of
`samples` do:

`vsel` = :(a[1]) or `vsel` = :(a[1,2]) or `vsel` = :(a[1,:,:])

with the brackets containing the cartesian index of the selected values.

To preserve the indexing information of the selection, `vsel` is encoded using
unicode characters. The new names in the marginalized samples are then `symbol`s
of these codes:

:(a[2])     => :a⌞2⌟
:(a[2:4])   => :a⌞2ː4⌟
:(a[1,2,:]) => :a⌞1ˌ2ˌː⌟

If `samples` is shaped, and `vsel` is an `Int` `i` or a vector of `Int`s `[i,j,k]`,
the i-th, j-th, and k-th parameters of `samples` are selected automatically.

If `samples` is unshaped, `vsel` may only be an ``Int` `i` or a vector
of `Int`s `[i,j,k]`. Then the i-th, j-th, and k-th columns of `samples` are selected.

For convenience, `vsel` can be of the shape

:([a[1,2], a[5:8], c]) instead of [:(a[1,2]), :(a[5:8]), :c]

The logd of the new samples is set to NaN, the info and aux fields of `samples` are copied.

Example:

```julia
marg_shaped_samples = bat_marginalize(orig_shaped_samples, :(a[1,2,:]))

varshape(marg_samples) == NamedTupleShape{(:a⌞1ˌ2ˌː⌟,), Tuple{ValueAccessor{ArrayShape{Real, 1}}}}((a⌞1ˌ2ˌː⌟ = ValueAccessor{ArrayShape{Real, 1}}(ArrayShape{Real, 1}((3,)), 0, 3),), 3)
```

```julia
marg_unshaped_samples = bat_marginalize(orig_unshaped_smaples, [1,2,3])
'''
"""
=#
function bat_marginalize(samples::DensitySampleVector,
                         vsel
)
    shaped = isshaped(samples)
    vs = varshape(samples)
    vsel = vsel isa Tuple ? [vsel...] : vsel

    if shaped
       vsel = _normalize_vsel_shaped(vs, vsel)
    else
       vsel = _normalize_vsel_unshaped(vsel)
    end

    if shaped
        flat_idxs, new_shape = marg_idxs_shaped(samples, vsel)
    else
        flat_idxs = marg_idxs_unshaped(samples, vsel)
    end

    flat_orig_data = ValueShapes.flatview(unshaped.(samples).v)
    marg_data = flat_orig_data[flat_idxs, :]
    marg_data = ndims(marg_data) > 1 ? nestedview(marg_data) : marg_data
    logd = fill(NaN, length(ndims(marg_data) > 1 ? axes(marg_data, 2) : marg_data))
    info = samples.info
    aux = samples.aux

    marg_samples = shaped ? new_shape.(DensitySampleVector(marg_data, logd, info = info, aux = aux)) : DensitySampleVector(marg_data, logd, info = info, aux = aux)
    return marg_samples
end


const DNASeq = BioSequence{DNAAlphabet{2}}
const DNASeqUnion = Union{DNASequence, DNASeq}
const Phred = Int8
const Prob = Float64
const LogProb = Float64
const Score = Float64


#InspectDR: Secondary annotation (Cairo layer)
#-------------------------------------------------------------------------------

#==Constants
===============================================================================#
const ALIGN_MAP = Dict{Symbol, CAlignment}(
	:tl => ALIGN_TOP | ALIGN_LEFT,
	:tc => ALIGN_TOP | ALIGN_HCENTER,
	:tr => ALIGN_TOP | ALIGN_RIGHT,

	:cl => ALIGN_VCENTER | ALIGN_LEFT,
	:cc => ALIGN_VCENTER | ALIGN_HCENTER,
	:cr => ALIGN_VCENTER | ALIGN_RIGHT,

	:bl => ALIGN_BOTTOM | ALIGN_LEFT,
	:bc => ALIGN_BOTTOM | ALIGN_HCENTER,
	:br => ALIGN_BOTTOM | ALIGN_RIGHT,
)

#==Rendering text annotation
===============================================================================#
function render(canvas::PCanvas2D, a::TextAnnotation, graphinfo::Graph2DInfo)
	align = get(ALIGN_MAP, a.align, ALIGN_BOTTOM | ALIGN_LEFT)
	angle = deg2rad(a.angle)
	pt = map2dev(a.pos, canvas.xf, graphinfo.ixf, canvas.graphbb)
	render(canvas.ctx, a.text, pt, a.font, angle=angle, align=align)
	return
end

#==Rendering markers
===============================================================================#

function render(canvas::PCanvas2D, mkr::HVMarker, graphinfo::Graph2DInfo)
	ctx = canvas.ctx #WANTCONST
	graphbb = canvas.graphbb #WANTCONST
	if :none == mkr.line.style
		return
	end

	setlinestyle(ctx, LineStyle(mkr.line))
	pt = read2axis(mkr.pos, graphinfo.ixf)
	pt = map2dev(canvas.xf, pt)

	if mkr.vmarker
		drawline(ctx, Point2D(pt.x, graphbb.ymin), Point2D(pt.x, graphbb.ymax))
	end
	if mkr.hmarker
		drawline(ctx, Point2D(graphbb.xmin, pt.y), Point2D(graphbb.xmax, pt.y))
	end

	return
end

#==Rendering polyline for annotation
===============================================================================#
function render(canvas::PCanvas2D, a::PolylineAnnotation, graphinfo::Graph2DInfo)
	ctx = canvas.ctx #WANTCONST

	x = a.x; y = a.y
	setlinestyle(ctx, LineStyle(a.line))
	pt = map2dev(canvas.xf, Point2D(x[1], y[1]))
	Cairo.move_to(ctx, pt.x, pt.y)
	for i in 2:length(x)
		pt = map2dev(canvas.xf, Point2D(x[i], y[i]))
		Cairo.line_to(ctx, pt.x, pt.y)
	end
	if a.closepath
		Cairo.close_path(ctx)
		renderfill(ctx, a.fillcolor)
	end
	Cairo.stroke(ctx)

#	function xf(x, y)
#		pt = map2dev(canvas.xf, map2axis(Point2D(x,y), graphinfo.ixf))
#		return (x, y)
#	end
	return
end

#==Generic rendering algorithm for annotations
===============================================================================#
#TODO: warn undefined???
render(canvas::PCanvas2D, a::PlotAnnotation, graphinfo::Graph2DInfo) = nothing

#Basic render function for PlotAnnotation types:
function render(canvas::PCanvas2D, a::PlotAnnotation, graphinfo::Graph2DInfo, strip::Int)
	if 0 == a.strip || a.strip == strip
		render(canvas, a, graphinfo)
	end
	return
end

function render(canvas::PCanvas2D, alist::Vector{T}, graphinfo::Graph2DInfo, strip::Int) where T<:PlotAnnotation
	for a in alist
		render(canvas, a, graphinfo, strip)
	end
	return
end

#Last line


# cuStateVec types

## custatevec compute type

function Base.convert(::Type{custatevecComputeType_t}, T::DataType)
    if T == Float16
        return CUSTATEVEC_COMPUTE_16F
    elseif T == Float32
        return CUSTATEVEC_COMPUTE_32F
    elseif T == Float64
        return CUSTATEVEC_COMPUTE_64F
    elseif T == UInt8
        return CUSTATEVEC_COMPUTE_8U
    elseif T == Int8
        return CUSTATEVEC_COMPUTE_8I
    elseif T == UInt32
        return CUSTATEVEC_COMPUTE_32U
    elseif T == Int32
        return CUSTATEVEC_COMPUTE_32I
    else
        throw(ArgumentError("CUSTATEVEC type equivalent for compute type $T does not exist!"))
    end
end

function Base.convert(::Type{Type}, T::custatevecComputeType_t)
    if T == CUSTATEVEC_COMPUTE_16F
        return Float16
    elseif T == CUSTATEVEC_COMPUTE_32F
        return Float32
    elseif T == CUSTATEVEC_COMPUTE_64F
        return Float64
    elseif T == CUSTATEVEC_COMPUTE_8U
        return UInt8
    elseif T == CUSTATEVEC_COMPUTE_32U
        return UInt32
    elseif T == CUSTATEVEC_COMPUTE_8I
        return Int8
    elseif T == CUSTATEVEC_COMPUTE_32I
        return Int32
    else
        throw(ArgumentError("Julia type equivalent for compute type $T does not exist!"))
    end
end

function compute_type(sv_type::DataType, mat_type::DataType)
    if sv_type == ComplexF64 && mat_type == ComplexF64
        return Float64
    elseif sv_type == ComplexF32 && mat_type == ComplexF64
        return Float32
    elseif sv_type == ComplexF32 && mat_type == ComplexF32
        return Float32
    end
end

abstract type Pauli end
struct PauliX <: Pauli end
struct PauliY <: Pauli end
struct PauliZ <: Pauli end
struct PauliI <: Pauli end
CuStateVecPauli(pauli::PauliX) = CUSTATEVEC_PAULI_X
CuStateVecPauli(pauli::PauliY) = CUSTATEVEC_PAULI_Y
CuStateVecPauli(pauli::PauliZ) = CUSTATEVEC_PAULI_Z
CuStateVecPauli(pauli::PauliI) = CUSTATEVEC_PAULI_I

mutable struct CuStateVec{T}
    data::CuVector{T}
    nbits::UInt32
end
function CuStateVec(T, n_qubits::Int)
    data = CUDA.zeros(T, 2^n_qubits)
    # in most cases, taking the hit here for setting one element
    # is cheaper than building the entire thing on the CPU and
    # copying it over
    CUDA.@allowscalar data[1] = one(T)
    CuStateVec{T}(data, n_qubits)
end

Base.eltype(sv::CuStateVec{T}) where T = T

mutable struct CuStateVecSampler
    handle::custatevecSamplerDescriptor_t
    ws_size::Csize_t
    function CuStateVecSampler(sv::CuStateVec, shot_count::UInt32)
        desc_ref   = Ref{custatevecSamplerDescriptor_t}()
        extra_size = Ref{Csize_t}()
        custatevecSampler_create(handle(), pointer(sv.data), eltype(sv), sv.nbits, desc_ref, shot_count, extra_size)
        obj = new(desc_ref[], extra_size[])
        #finalizer(custatevecSampler_destroy, obj) # apparently there is no such method?
        obj
    end
end

Base.unsafe_convert(::Type{custatevecSamplerDescriptor_t}, desc::CuStateVecSampler) = desc.handle


### Thermal Generation Formulations

abstract type AbstractThermalFormulation <: AbstractDeviceFormulation end

abstract type AbstractThermalDispatchForm <: AbstractThermalFormulation end

struct ThermalUnitCommitment <: AbstractThermalFormulation end

struct ThermalDispatch <: AbstractThermalDispatchForm end

struct ThermalRampLimited <: AbstractThermalDispatchForm end

struct ThermalDispatchNoMin <: AbstractThermalDispatchForm end

# Variables for Thermal Generation

"""
This function add the variables for power generation output to the model
"""
function activepower_variables(ps_m::CanonicalModel,
                               devices::Array{T,1},
                               time_range::UnitRange{Int64}) where {T <: PSY.ThermalGen}

    add_variable(ps_m, devices, time_range, :Pth, false, :var_active)

    return

end

"""
This function add the variables for power generation output to the model
"""
function reactivepower_variables(ps_m::CanonicalModel,
                                 devices::Array{T,1},
                                 time_range::UnitRange{Int64}) where {T <: PSY.ThermalGen}

    add_variable(ps_m, devices, time_range, :Qth, false, :var_reactive)

    return

end

"""
This function add the variables for power generation commitment to the model
"""
function commitment_variables(ps_m::CanonicalModel,
                              devices::Array{T,1},
                              time_range::UnitRange{Int64}) where {T <: PSY.ThermalGen}

    add_variable(ps_m, devices, time_range, :on_th, true)
    add_variable(ps_m, devices, time_range, :start_th, true)
    add_variable(ps_m, devices, time_range, :stop_th, true)

    return

end

"""
This function adds the active power limits of generators when there are no CommitmentVariables
"""
function activepower_constraints(ps_m::CanonicalModel,
                                 devices::Array{T,1},
                                 device_formulation::Type{D},
                                 system_formulation::Type{S},
                                 time_range::UnitRange{Int64}) where {T <: PSY.ThermalGen,
                                                                      D <: AbstractThermalDispatchForm,
                                                                      S <: PM.AbstractPowerFormulation}

    range_data = [(g.name, g.tech.activepowerlimits) for g in devices]

    device_range(ps_m, range_data, time_range, :thermal_active_range, :Pth)

    return

end

"""
This function adds the active power limits of generators when there are CommitmentVariables
"""
function activepower_constraints(ps_m::CanonicalModel,
                                 devices::Array{T,1},
                                 device_formulation::Type{D},
                                 system_formulation::Type{S},
                                 time_range::UnitRange{Int64}) where {T <: PSY.ThermalGen,
                                                                      D <: AbstractThermalFormulation,
                                                                      S <: PM.AbstractPowerFormulation}

    range_data = [(g.name, g.tech.activepowerlimits) for g in devices]

    device_semicontinuousrange(ps_m, range_data, time_range, :thermal_active_range, :Pth, :on_th)

    return

end

"""
This function adds the reactive  power limits of generators when there are CommitmentVariables
"""
function reactivepower_constraints(ps_m::CanonicalModel,
                                   devices::Array{T,1},
                                   device_formulation::Type{D},
                                   system_formulation::Type{S},
                                   time_range::UnitRange{Int64}) where {T <: PSY.ThermalGen,
                                                                        D <: AbstractThermalDispatchForm,
                                                                        S <: PM.AbstractPowerFormulation}

    range_data = [(g.name, g.tech.reactivepowerlimits) for g in devices]

    device_range(ps_m, range_data , time_range, :thermal_reactive_range, :Qth)

    return

end

"""
This function adds the reactive power limits of generators when there CommitmentVariables
"""
function reactivepower_constraints(ps_m::CanonicalModel,
                                   devices::Array{T,1},
                                   device_formulation::Type{D},
                                   system_formulation::Type{S},
                                   time_range::UnitRange{Int64}) where {T <: PSY.ThermalGen,
                                                                        D <: AbstractThermalFormulation,
                                                                        S <: PM.AbstractPowerFormulation}

    range_data = [(g.name, g.tech.reactivepowerlimits) for g in devices]

    device_semicontinuousrange(ps_m, range_data , time_range, :thermal_reactive_range, :Qth, :on_th)

    return

end

"""
This function adds the active power limits of generators when there are no CommitmentVariables
"""
function activepower_constraints(ps_m::CanonicalModel,
                                 devices::Array{T,1},
                                 device_formulation::Type{PSI.ThermalDispatchNoMin},
                                 system_formulation::Type{S},
                                 time_range::UnitRange{Int64}) where {T <: PSY.ThermalGen,
                                                                      S <: PM.AbstractPowerFormulation}

    range_data = [(g.name, (min = 0.0, max=g.tech.activepowerlimits.max)) for g in devices]

    device_range(ps_m, range_data, time_range, :thermal_active_range, :Pth)

    return

end

"""
This function adds the reactive  power limits of generators when there are CommitmentVariables
"""
function reactivepower_constraints(ps_m::CanonicalModel,
                                   devices::Array{T,1},
                                   device_formulation::Type{PSI.ThermalDispatchNoMin},
                                   system_formulation::Type{S},
                                   time_range::UnitRange{Int64}) where {T <: PSY.ThermalGen,
                                                                        S <: PM.AbstractPowerFormulation}

    range_data = [(g.name, (min = 0.0, max=g.tech.reactivepowerlimits.max)) for g in devices]

    device_range(ps_m, range_data , time_range, :thermal_reactive_range, :Qth)

    return

end

### Constraints for Thermal Generation without commitment variables ####

"""
This function adds the Commitment Status constraint when there are CommitmentVariables
"""
function commitment_constraints(ps_m::CanonicalModel,
                                devices::Array{T,1},
                                device_formulation::Type{D},
                                system_formulation::Type{S},
                                time_range::UnitRange{Int64},
                                parameters::Bool) where {T <: PSY.ThermalGen,
                                                                     D <: AbstractThermalFormulation,
                                                                     S <: PM.AbstractPowerFormulation}

    if !(:thermal_status in keys(ps_m.initial_conditions))
        @info("Initial status conditions not provided. This can lead to unwanted results")
        status_init(ps_m, devices, parameters)
    end

    device_commitment(ps_m,
                      ps_m.initial_conditions[:thermal_status],
                      time_range, :commitment_th,
                      (:start_th, :stop_th, :on_th))

    return

end

# ramping constraints

function _get_data_for_rocc(devices::Array{T,1}) where {T <: PSY.ThermalGen}

    set_name = Vector{String}(undef, length(devices))
    ramp_params = Vector{UpDown}(undef, length(devices))
    minmax_params = Vector{MinMax}(undef, length(devices))

    idx = eachindex(devices)
    i, state = iterate(idx)
    for g in devices
        if !isnothing(g.tech.ramplimits)
            set_name[i] = g.name
            ramp_params[i] = g.tech.ramplimits
            minmax_params[i] = g.tech.activepowerlimits
            y = iterate(idx, state)
            y === nothing && (i += 1; break)
            i, state = y
        end
    end

    deleteat!(set_name, i:last(idx))
    deleteat!(ramp_params, i:last(idx))
    deleteat!(minmax_params, i:last(idx))

    return set_name, ramp_params, minmax_params

end

"""
This function adds the ramping limits of generators when there are CommitmentVariables
"""
function ramp_constraints(ps_m::CanonicalModel,
                          devices::Array{T,1},
                          device_formulation::Type{D},
                          system_formulation::Type{S},
                          time_range::UnitRange{Int64},
                          parameters::Bool) where {T <: PSY.ThermalGen,
                                                               D <: AbstractThermalFormulation,
                                                               S <: PM.AbstractPowerFormulation}

    data = _get_data_for_rocc(devices)

    if !isempty(data[2])
        if !(:thermal_output in keys(ps_m.initial_conditions))
            @info("Initial conditions for rate of change not provided. This can lead to unwanted results")
            output_init(ps_m, devices, parameters)
        end

        @assert length(data[2]) == length(ps_m.initial_conditions[:thermal_output])
        # Here goes the reactive power ramp limits
        device_mixedinteger_rateofchange(ps_m,
                                        data,
                                        ps_m.initial_conditions[:thermal_output],
                                        time_range,
                                        :ramp_thermal, (:Pth, :start_th, :stop_th))
    else
        @warn "Data doesn't contain generators with ramp limits, consider adjusting your formulation"
    end

    return

end

function ramp_constraints(ps_m::CanonicalModel,
                          devices::Array{T,1},
                          device_formulation::Type{D},
                          system_formulation::Type{S},
                          time_range::UnitRange{Int64},
                          parameters::Bool) where {T <: PSY.ThermalGen,
                                                               D <: AbstractThermalDispatchForm,
                                                               S <: PM.AbstractPowerFormulation}

    data = _get_data_for_rocc(devices)

    if !isempty(data[2])
        if !(:thermal_output in keys(ps_m.initial_conditions))
            @info("Initial conditions for rate of change not provided. This can lead to unwanted results")
            output_init(ps_m, devices, parameters)
        end
        # Here goes the reactive power ramp limits
        device_linear_rateofchange(ps_m,
                                   (data[1], data[2]),
                                   ps_m.initial_conditions[:thermal_output], time_range,
                                   :ramp_thermal, :Pth)
    else
        @warn "Data doesn't contain generators with ramp limits, consider adjusting your formulation"
    end

    return

end

"""
This function adds the ramping limits of generators when there are CommitmentVariables
"""
function ramp_constraints(ps_m::CanonicalModel,
                          devices::Array{T,1},
                          device_formulation::Type{D},
                          system_formulation::Type{S},
                          time_range::UnitRange{Int64},
                          parameters::Bool) where {T <: PSY.ThermalGen,
                                                               D <: AbstractThermalFormulation,
                                                               S <: PM.AbstractActivePowerFormulation}

    data = _get_data_for_rocc(devices)

    if !isempty(data[2])
        if !(:thermal_output in keys(ps_m.initial_conditions))
            @info("Initial conditions for rate of change not provided. This can lead to unwanted results")
            output_init(ps_m, devices, parameters)
        end

        @assert length(data[2]) == length(ps_m.initial_conditions[:thermal_output])

        device_mixedinteger_rateofchange(ps_m,
                                        data,
                                        ps_m.initial_conditions[:thermal_output],
                                        time_range,
                                        :ramp_thermal, (:Pth, :start_th, :stop_th))

    else
        @warn "Data doesn't contain generators with ramp limits, consider adjusting your formulation"
    end

    return

end

function ramp_constraints(ps_m::CanonicalModel,
                          devices::Array{T,1},
                          device_formulation::Type{D},
                          system_formulation::Type{S},
                          time_range::UnitRange{Int64},
                          parameters::Bool) where {T <: PSY.ThermalGen,
                                                               D <: AbstractThermalDispatchForm,
                                                               S <: PM.AbstractActivePowerFormulation}

    data = _get_data_for_rocc(devices)

    if !isempty(data[2])
        if !(:thermal_output in keys(ps_m.initial_conditions))
            @info("Initial conditions for rate of change not provided. This can lead to unwanted results")
            output_init(ps_m, devices, parameters)
        end
        device_linear_rateofchange(ps_m,
                                    (data[1], data[2]),
                                    ps_m.initial_conditions[:thermal_output], time_range,
                                    :ramp_thermal, :Pth)
    else
        @warn "Data doesn't contain generators with ramp limits, consider adjusting your formulation"
    end

    return

end

# time constraints

function _get_data_for_tdc(devices::Array{T,1}) where {T <: PSY.ThermalGen}

    set_name = Vector{String}(undef, length(devices))
    time_params = Vector{UpDown}(undef, length(devices))

    idx = eachindex(devices)
    i, state = iterate(idx)
    for g in devices
        if !isnothing(g.tech.timelimits)
            set_name[i] = g.name
            time_params[i] = g.tech.timelimits
            y = iterate(idx, state)
            y === nothing && (i += 1; break)
            i, state = y
        end
    end

    deleteat!(set_name, i:last(idx))
    deleteat!(time_params, i:last(idx))

    return set_name, time_params
end

function time_constraints(ps_m::CanonicalModel,
                          devices::Array{T,1},
                          device_formulation::Type{D},
                          system_formulation::Type{S},
                          time_range::UnitRange{Int64},
                          parameters::Bool = false) where {T <: PSY.ThermalGen,
                                                               D <: AbstractThermalFormulation,
                                                               S <: PM.AbstractPowerFormulation}

    duration_data = _get_data_for_tdc(devices)

    if !isempty(duration_data)
        if !(:thermal_duration_on in keys(ps_m.initial_conditions))
            @info("Initial conditions for time up/down not provided. This can lead to unwanted results")
            duration_init(ps_m, devices, parameters)
        end
        @warn("Currently Min-up/down time limits only work if time resoultion they are defined at matches the simulation's time resoultion")
        device_duration_retrospective(ps_m,
                                      duration_data,
                                      ps_m.initial_conditions[:thermal_duration_on],
                                      ps_m.initial_conditions[:thermal_duration_off],
                                      time_range,
                                      :duration, (:on_th, :start_th, :stop_th))
    else
        @warn "Data doesn't contain generators with time-up/down limits, consider adjusting your formulation"
    end

    return

end

# thermal generation cost

function cost_function(ps_m::CanonicalModel,
                       devices::Array{T,1},
                       device_formulation::Type{D},
                       system_formulation::Type{S}) where {T <: PSY.ThermalGen,
                                                           D <: AbstractThermalDispatchForm,
                                                           S <: PM.AbstractPowerFormulation}

    add_to_cost(ps_m, devices, :Pth, :variablecost)

    return

end

function cost_function(ps_m::CanonicalModel,
                       devices::Array{T,1},
                       device_formulation::Type{D},
                       system_formulation::Type{S}) where {T <: PSY.ThermalGen,
                                                           D <: AbstractThermalFormulation,
                                                           S <: PM.AbstractPowerFormulation}

    #Variable Cost component
    add_to_cost(ps_m, devices, :Pth, :variablecost)

    #Commitment Cost Components
    add_to_cost(ps_m, devices, :start_th, :startupcost)
    add_to_cost(ps_m, devices, :stop_th, :shutdncost)
    add_to_cost(ps_m, devices, :on_th, :fixedcost)

    return

end


"""
    function addEventDivision_!(p::Program_,abm::Agent,space::SimulationSpace,platform::String)

Generate the functions for division events.
"""
function addEventDivision_!(p::Program_,abm::Agent,space::SimulationSpace,platform::String)

    if "EventDivision" in keys(abm.declaredUpdates)

        code = abm.declaredUpdates["EventDivision"]

        if !@capture(code,if v_ g__ end)
            error("Erroneos structure of division events. Should be if something update end")
        end

        code = vectorize_(abm,code,p)
        subcode = postwalk(x->@capture(x,if v_ g__ end) ? quote $(g...) end : x, code)
        subcode = unblock(subcode)
        condition = postwalk(x->@capture(x,if v_ g__ end) ? v : x, code)
        condition = unblock(condition)

        #Add atomic_add
        if platform == "cpu"
            pushfirst!(subcode.args, :(ic1New_ = Threads.atomic_add!(NV,Int(1)) + 1))
            pushfirst!(subcode.args, :(idNew_ = Threads.atomic_add!(agentIdMax,Int(2)) + 1))
        elseif platform == "gpu"
            pushfirst!(subcode.args, :(ic1New_ = CUDA.atomic_add!(CUDA.pointer(NV,1),Int32(1)) + 1))
            pushfirst!(subcode.args, :(idNew_ = CUDA.atomic_add!(CUDA.pointer(agentIdMax,1),Int32(2)) + 1))
        end
        for k in ["Local","Identity"]
            vec = Meta.parse(string(lowercase(k),"V"))
            vecCopy = Meta.parse(string(lowercase(k),"VCopy"))
            for (i,j) in enumerate(abm.declaredSymbols[k])
                if j in keys(p.update["EventDivision"])
                    ii = p.update[k][j]
                    #Parse content
                    s = Meta.parse(string(j,"_1"))
                    subcode = postwalk(x -> @capture(x,g_=v__) && g == s ? :($vecCopy[ic1_,$ii] = $(v...)) : x, subcode)
                    s = Meta.parse(string(j,"_2"))
                    subcode = postwalk(x -> @capture(x,g_=v__) && g == s ? :($vecCopy[ic1New_,$ii] = $(v...)) : x, subcode)
                end
            end
        end
        for k in ["Local","Identity"]  #Add other updates of non changed variables
            vec = Meta.parse(string(lowercase(k),"V"))
            vecCopy = Meta.parse(string(lowercase(k),"VCopy"))
            for (i,j) in enumerate(abm.declaredSymbols[k])
                #Add update if doesn't exist
                if !(j in keys(p.update["EventDivision"])) && j != :agentId
                    push!(subcode.args,:($vec[ic1New_,$i] = $vec[ic1_,$i]))
                elseif j != :agentId
                    ii = p.update[k][j]
                    push!(subcode.args,:($vec[ic1_,$i] = $vecCopy[ic1_,$ii]))
                    push!(subcode.args,:($vec[ic1New_,$i] = $vecCopy[ic1New_,$ii]))
                else
                    ii = findfirst(abm.declaredSymbols["Identity"] .== j)
                    push!(subcode.args,:($vec[ic1_,$ii] = idNew_))
                    push!(subcode.args,:($vec[ic1New_,$ii] = idNew_ + 1))
                end
            end
        end
        code = quote
            if $condition
                $subcode
            end
        end

        code = simpleFirstLoopWrapInFunction_(platform,:division_!,code)

        push!(p.declareF.args,code)
        push!(p.args,:(NV))
        if platform == "cpu"
            push!(p.declareVar.args,:(NV = Threads.Atomic{Int}(N)))
            push!(p.declareVar.args,:(agentIdMax = Threads.Atomic{Int}(N)))

            push!(p.execInloop.args,
                    :(begin
                        @platformAdapt division_!(ARGS_)
                        N = NV[]
                    end)
                )
        elseif platform == "gpu"
            push!(p.declareVar.args,:(NV = N .*CUDA.ones(Int,1)))
            push!(p.declareVar.args,:(agentIdMax = N .*CUDA.ones(Int,1)))

            push!(p.execInloop.args,
                    :(begin
                        @platformAdapt division_!(ARGS_)
                        N = Core.Array(NV)[1]
                    end)
                )
        end

    end

    return nothing
end


"""
    general_tree_test(com::CS.CoM)

Performs some general tests which are nice to test if a new constraint works as intended and implements the necessary methods.
Should be called only if optimize was called already as this works on the final tree structure.
It fails if no backtracking was used as this method should only be called when at least some backtracking steps are needed.
Additionally `keep_logs` needs to be activated to make it work.
What it does:
    - Tests `checkout_from_to!` which basically means jumping from one node in the tree to another.
        - Checks there whether this produces the same search space as before at that stage.
"""
function general_tree_test(com::CS.CoM)
    if length(com.logs) == 0
        @error "Make sure that keep_logs is set to true."
    end
    @assert length(com.logs) > 10
    c_backtrack_idx = com.c_backtrack_idx
    Random.seed!(241)
    n_backtracks = length(com.logs)
    path = Int[c_backtrack_idx]
    path_type = Symbol[]
    all_correct = true
    n_children_tests = 0

    for t in 1:100
        status = :Open
        next_idx = 0
        while status == :Open
            # don't allow root node as that might cause chaos :D
            next_idx = rand(2:n_backtracks)
            status = com.logs[next_idx].status
        end

        CS.checkout_from_to!(com, c_backtrack_idx, next_idx)

        if com.backtrack_vec[next_idx].parent_idx != 0
            c_backtrack_idx = com.backtrack_vec[next_idx].parent_idx

            c_search_space = com.search_space
            # this is a dict =>
            expected_search_space = com.logs[c_backtrack_idx].var_states
            correct = true
            for var in c_search_space
                if sort(CS.values(var)) != sort(expected_search_space[var.idx])
                    println("var.idx: $(var.idx)")
                    println("vals: $(CS.values(var))")
                    println("expected: $(expected_search_space[var.idx])")
                    correct = false
                    all_correct = false
                end
            end
            if !correct
                @error "There were errors on the jump path: $path before pruning the last step"
            end
        end

        c_backtrack_idx = next_idx

        # if it has children
        if length(com.logs[c_backtrack_idx].children) > 0 && n_children_tests < 5
            com.c_backtrack_idx = c_backtrack_idx
            vidx = com.logs[c_backtrack_idx].vidx
            n_children_tests += 1
            # test that pruning produces the same output as before

            for var in com.search_space
                var.changes[c_backtrack_idx] = Vector{Tuple{Symbol,Int,Int,Int}}()
            end
            @assert CS.remove_above!(
                com,
                com.search_space[vidx],
                com.logs[c_backtrack_idx].ub,
            )
            @assert CS.remove_below!(
                com,
                com.search_space[vidx],
                com.logs[c_backtrack_idx].lb,
            )

            @assert CS.prune!(com)
            CS.call_finished_pruning!(com)
            push!(path_type, :prune)
        else
            # prune the last step based on saved information instead
            CS.restore_prune!(com, c_backtrack_idx)
            push!(path_type, :restore_prune)
        end

        push!(path, c_backtrack_idx)
        c_search_space = com.search_space
        # this is a dict =>
        expected_search_space = com.logs[c_backtrack_idx].var_states
        correct = true
        for var in c_search_space
            if sort(CS.values(var)) != sort(expected_search_space[var.idx])
                println("var.idx: $(var.idx)")
                println("vals: $(CS.values(var))")
                println("expected: $(expected_search_space[var.idx])")
                correct = false
                all_correct = false
            end
        end
        if !correct
            @error "There were errors on the jump path: $path"
            @error "More info about the path: $path_type"
        end
        @assert correct
    end
    if n_children_tests != 5
        @error "Make sure that you have more feasible nodes in the search tree to test more."
    end
    @assert n_children_tests == 5

    # back to solved state
    CS.checkout_from_to!(com, c_backtrack_idx, path[1])
    # prune the last step
    CS.restore_prune!(com, path[1])

    return all_correct
end

function is_solved(com::CS.CoM)
    variables = com.search_space
    all_fixed = all(v -> CS.isfixed(v), variables)
    if !all_fixed
        @error "Not all variables are fixed"
        return false
    end
    for constraint in com.constraints
        values = CS.value.(com.search_space[constraint.indices])
        c_solved =
            CS.is_constraint_solved(constraint, constraint.fct, constraint.set, values)
        c_violated =
            CS.is_constraint_violated(com, constraint, constraint.fct, constraint.set)
        if !c_solved || c_violated
            @error "Constraint $(constraint.idx) is not solved"
            @error "Info about constraint: $(typeof(constraint)), $(typeof(constraint.fct)), $(typeof(constraint.set))"
            return false
        end
    end
    return true
end