export DistProduct, ProdGeneric

import Distributions
import Base: prod, show

"""
    DistProduct

If inference backend cannot return an analytical solution for a product of two distributions it may fallback to the `DistProduct` structure
`DistProduct` is useful to propagate the exact forms of two messages until it hits some approximation method for form-constraint.
However `DistProduct` cannot be used to compute statistics such as mean or variance.
It has to be approximated before using in actual inference procedure.

Backend exploits form constraints specification which usually help to deal with intractable distributions products.
User may use EM form constraint with a specific optimisation algorithm or it may approximate intractable product with Gaussian Distribution
using for example Laplace approximation

See also: [`prod`](@ref)
"""
struct DistProduct{ L, R }
    left  :: L
    right :: R
end

Base.show(io::IO, product::DistProduct) = print(io, "DistProduct(", getleft(product), ",", getright(product), ")")

getleft(product::DistProduct)  = product.left
getright(product::DistProduct) = product.right

function Distributions.support(product::DistProduct)
    lsupport = Distributions.support(getleft(product))
    rsupport = Distributions.support(getright(product))
    if lsupport != rsupport
        error("Product $product has different support for left and right entries.")
    end
    return lsupport
end

Distributions.mean(product::DistProduct)      = error("mean() is not defined for $(product). DistProduct structure has to be approximated and cannot be used in inference procedure.")
Distributions.median(product::DistProduct)    = error("median() is not defined for $(product). DistProduct structure has to be approximated and cannot be used in inference procedure.")
Distributions.mode(product::DistProduct)      = error("mode() is not defined for $(product). DistProduct structure has to be approximated and cannot be used in inference procedure.")
Distributions.shape(product::DistProduct)     = error("shape() is not defined for $(product). DistProduct structure has to be approximated and cannot be used in inference procedure.")
Distributions.scale(product::DistProduct)     = error("scale() is not defined for $(product). DistProduct structure has to be approximated and cannot be used in inference procedure.")
Distributions.rate(product::DistProduct)      = error("rate() is not defined for $(product). DistProduct structure has to be approximated and cannot be used in inference procedure.")
Distributions.var(product::DistProduct)       = error("var() is not defined for $(product). DistProduct structure has to be approximated and cannot be used in inference procedure.")
Distributions.std(product::DistProduct)       = error("std() is not defined for $(product). DistProduct structure has to be approximated and cannot be used in inference procedure.")
Distributions.cov(product::DistProduct)       = error("cov() is not defined for $(product). DistProduct structure has to be approximated and cannot be used in inference procedure.")
Distributions.invcov(product::DistProduct)    = error("invcov() is not defined for $(product). DistProduct structure has to be approximated and cannot be used in inference procedure.")
Distributions.logdetcov(product::DistProduct) = error("logdetcov() is not defined for $(product). DistProduct structure has to be approximated and cannot be used in inference procedure.")
Distributions.entropy(product::DistProduct)   = error("entropy() is not defined for $(product). DistProduct structure has to be approximated and cannot be used in inference procedure.")
Distributions.params(product::DistProduct)    = error("params() is not defined for $(product). DistProduct structure has to be approximated and cannot be used in inference procedure.")

Distributions.pdf(product::DistProduct, x)    = Distributions.pdf(product.left, x) * Distributions.pdf(product.right, x)
Distributions.logpdf(product::DistProduct, x) = Distributions.logpdf(product.left, x) + Distributions.logpdf(product.right, x)

Base.precision(product::DistProduct) = error("precision() is not defined for $(product). DistProduct structure has to be approximated and cannot be used in inference procedure.")
Base.length(product::DistProduct)    = error("length() is not defined for $(product). DistProduct structure has to be approximated and cannot be used in inference procedure.")
Base.ndims(product::DistProduct)     = error("ndims() is not defined for $(product). DistProduct structure has to be approximated and cannot be used in inference procedure.")
Base.size(product::DistProduct)      = error("size() is not defined for $(product). DistProduct structure has to be approximated and cannot be used in inference procedure.")

probvec(product::DistProduct)         = error("probvec() is not defined for $(product). DistProduct structure has to be approximated and cannot be used in inference procedure.")
weightedmean(product::DistProduct)    = error("weightedmean() is not defined for $(product). DistProduct structure has to be approximated and cannot be used in inference procedure.")

variate_form(::P)                         where { P <: DistProduct } = variate_form(P)
variate_form(::Type{ DistProduct{L, R} }) where { L, R }             = _check_dist_product_variate_form(variate_form(L), variate_form(R))

_check_dist_product_variate_form(::Type{ F }, ::Type{ F })   where { F <: VariateForm }                     = F
_check_dist_product_variate_form(::Type{ F1 }, ::Type{ F2 }) where { F1 <: VariateForm, F2 <: VariateForm } = error("DistProduct has different variate forms for left ($F1) and right ($F2) entries.")

value_support(::P)                         where { P <: DistProduct } = value_support(P)
value_support(::Type{ DistProduct{L, R} }) where { L, R }             = _check_dist_product_value_support(value_support(L), value_support(R))

_check_dist_product_value_support(::Type{ S }, ::Type{ S })   where { S <: ValueSupport }                      = S
_check_dist_product_value_support(::Type{ S1 }, ::Type{ S2 }) where { S1 <: ValueSupport, S2 <: ValueSupport } = error("DistProduct has different value supports for left ($S1) and right ($S2) entries.")

"""
    ProdGeneric{C}

`ProdGeneric` is one of the strategies for `prod` function. This strategy mimics `prod_constraint` constraint but does not fail in case of no analytical rule is available.

See also: [`prod`](@ref), [`ProdAnalytical`](@ref), [`ProdPreserveType`](@ref)
"""
struct ProdGeneric{C}
    prod_constraint :: C
end

get_constraint(prod_generic::ProdGeneric) = prod_generic.prod_constraint

ProdGeneric() = ProdGeneric(ProdAnalytical())

prod(::ProdGeneric, ::Missing, right)     = right
prod(::ProdGeneric, left, ::Missing)      = left
prod(::ProdGeneric, ::Missing, ::Missing) = missing

prod(generic::ProdGeneric, left::L, right::R) where { L, R } = prod(generic, prod_analytical_rule(L, R), left, right)

prod(generic::ProdGeneric, ::ProdAnalyticalRuleAvailable, left, right) = prod(get_constraint(generic), left, right)
prod(generic::ProdGeneric, ::ProdAnalyticalRuleUnknown, left, right)   = DistProduct(left, right)

# In case of ProdPointMass we want to propagate a single `DistProduct` as much as possible and do not create a big tree of product which will reduce performance significantly
# In this methods the general rule is the folowing: If we see that one of the arguments of `DistProduct` has the same function form
# as second argument of `prod` function it is better to try to `prod` them together with `NoConstraint` strategy.
prod(generic::ProdGeneric, left::DistProduct{L, R}, right::T) where { L, R, T } = prod(generic, prod_analytical_rule(L, T), prod_analytical_rule(R, T), left, right)
prod(generic::ProdGeneric, left::T, right::DistProduct{L, R}) where { L, R, T } = prod(generic, prod_analytical_rule(T, L), prod_analytical_rule(T, R), left, right)

prod(generic::ProdGeneric, ::ProdAnalyticalRuleUnknown,   ::ProdAnalyticalRuleUnknown,   left::DistProduct, right) = DistProduct(left, right)
prod(generic::ProdGeneric, ::ProdAnalyticalRuleAvailable, ::ProdAnalyticalRuleUnknown,   left::DistProduct, right) = DistProduct(prod(get_constraint(generic), getleft(left), right), getright(left))
prod(generic::ProdGeneric, ::ProdAnalyticalRuleUnknown,   ::ProdAnalyticalRuleAvailable, left::DistProduct, right) = DistProduct(getleft(left), prod(get_constraint(generic), getright(left), right))

prod(generic::ProdGeneric, ::ProdAnalyticalRuleUnknown,   ::ProdAnalyticalRuleUnknown,   left, right::DistProduct) = DistProduct(left, right)
prod(generic::ProdGeneric, ::ProdAnalyticalRuleAvailable, ::ProdAnalyticalRuleUnknown,   left, right::DistProduct) = DistProduct(prod(get_constraint(generic), left, getleft(right)), getright(right))
prod(generic::ProdGeneric, ::ProdAnalyticalRuleUnknown,   ::ProdAnalyticalRuleAvailable, left, right::DistProduct) = DistProduct(getleft(right), prod(get_constraint(generic), left, getright(right)))

function prod(generic::ProdGeneric, left::DistProduct{L1, R1}, right::DistProduct{L2, R2}) where {L1, R1, L2, R2}
    return prod(
        generic,
        prod_analytical_rule(L1, L2), prod_analytical_rule(L1, R2),
        prod_analytical_rule(R1, L2), prod_analytical_rule(R1, R2),
        left, right
    )
end

prod(::ProdGeneric, _, _, _, _, left::DistProduct, right::DistProduct) = DistProduct(left, right)

function prod(generic::ProdGeneric, ::ProdAnalyticalRuleAvailable, _, _, ::ProdAnalyticalRuleAvailable, left::DistProduct, right::DistProduct)
    return prod(generic, prod(get_constraint(generic), getleft(left), getleft(right)), prod(get_constraint(generic), getright(left), getright(right)))
end


#!/bin/env julia

using CLI, JLD, Distributions, HypothesisTests

immutable Region
	start::Int32
	stop::Int32
end

# Print comment lines to stdout and return
# first non-comment line
function skip_to_header(vcf_file)
	while true
		line = readline(vcf_file);
		if !startswith( line, "#") return line; end
		println(line);
	end
end

drop_chr_prefix(chr) = startswith(chr, "chr") ? chr[4:end] : chr

is_protein_altering(effect) = ismatch(
	r"Missense|[Ff]rameshift|Stopgain|Stoploss|Splice", effect);

function mutation_class(ref, alt)
	if ref == "C"
		if alt == "A"; return 1; end
		if alt == "G"; return 2; end
		if alt == "T"; return 3; end
	elseif ref == "G"
		if alt == "T"; return 1; end
		if alt == "C"; return 2; end
		if alt == "A"; return 3; end
	elseif ref == "T"
		if alt == "A"; return 4; end
		if alt == "C"; return 5; end
		if alt == "G"; return 6; end
	elseif ref == "A"
		if alt == "T"; return 4; end
		if alt == "G"; return 5; end
		if alt == "C"; return 6; end
	end

	if length(ref) > 1 || ref == "-" || length(alt) > 1 || alt == "-"
		return 0
	else
		error("Mutation class could not be defined.")
	end
end

function somatic(vcf_file::IO, samples_file::IO; alt_reads=3, alt_frac=0.1,
	test_ref_ratio=5, test_bg_ratio=20, test_bg_alpha=1.0, ref_reads=0,
	keep_all=false, min_mapq=10, min_sidedness=6, keep_old_stars::Bool=false)

	header = skip_to_header(vcf_file); println(header)
	header_cols = split(header, '\t')
	notes_col = findfirst(header_cols .== "NOTES")
	alt_col = findfirst(header_cols .== "ALT")
	ref_col = findfirst(header_cols .== "REF")

	samples = header_cols[notes_col+1:end];
	n_samples = length(samples);

	# Read samples file and separate the row header
	d = readdlm(samples_file, '\t');
	sample_file_header = d[1, :][:];
	d = d[2:end, :];

	sample_test_col = findfirst(sample_file_header .== "TEST")
	sample_ref_col = findfirst(sample_file_header .== "REF")
	assert(sample_test_col != 0 && sample_ref_col != 0)

	# Check that all samples named in the design file are present in the VCF
	for s in unique(d[:, [sample_test_col, sample_ref_col]])
		if s != "" && !(s in samples)
			error("Sample $s is listed in the design file but missing from VCF.")
		end
	end

	# Background error is estimated based on paired *and* unpaired reference
	# samples.
	ref_samples = filter(s -> s != "", unique(d[:, sample_ref_col]))
	background = indexin(ref_samples, samples)

	# Get indices of test and ref sample rows in the samples file
	test_rows = d[:, sample_test_col] .!= ""
	test = indexin(d[test_rows, sample_test_col], samples)
	n_tests = length(test);
	ref = indexin(d[test_rows, sample_ref_col], samples)
	info("Analyzing $n_tests tumor-normal pairs...")

	if !isempty(findin(background, test))
		warn("Some background reference samples also used as test samples.")
	end
	info("Calculating background error rate based on $(length(background)) samples...")

	# Read alt_frac from design file or command line parameter.
	alt_frac_col = findfirst(sample_file_header .== "ALT_FRAC")
	if alt_frac_col == 0
		info("Using a mutant allele fraction threshold of $(alt_frac * 100)% for all samples...")
		alt_frac = ones(n_tests) * alt_frac
	else
		alt_frac = convert(Array{Float64}, d[test_rows, alt_frac_col])
	end

	# TODO: Allow sample-specific values to be specified in design file for
	# alt_reads, test_ref_ratio and test_bg_ratio as well.

	for line in eachline(vcf_file)

		cols = split(line, '\t')
		alt = zeros(Int32, n_samples);
		total = zeros(Int32, n_samples);
		# CALL2 support
		mapq = zeros(Int32, n_samples);
		baseq = zeros(Int32, n_samples);
		sidedness = zeros(Int32, n_samples);

		call2_format = false
		is_indel = length(cols[alt_col]) != length(cols[ref_col])

		for s in 1:n_samples
			col = cols[notes_col + s]
			gt = split(cols[notes_col + s], ':')
			alt[s] = parse(Int32, gt[1])
			total[s] = parse(Int32, gt[2])
			# CALL2 support
			if length(gt) >= 5
				mapq[s] = parse(Int32, gt[3])
				baseq[s] = parse(Int32, gt[4])
				sidedness[s] = parse(Int32, gt[5])
				call2_format = true
			end
		end

		fracs = alt ./ total
		bg_alt = sum(alt[background])
		bg_total = sum(total[background])
		avg_ref_frac = bg_alt / bg_total

		somatic = falses(n_samples)
		for k in 1:n_tests
			t = test[k]; # Index of test sample col
			r = ref[k];  # Index of ref sample col

			# Perform checks to determine somatic status of mutation
			if fracs[t] < alt_frac[k]; continue; end
			if r != 0 && fracs[t] < test_ref_ratio * fracs[r]; continue; end
			if fracs[t] < test_bg_ratio * avg_ref_frac; continue; end
			if alt[t] < alt_reads; continue; end
			if r != 0 && total[r] < ref_reads; continue; end
			if call2_format
				if !is_indel && mapq[t] < min_mapq; continue; end
				if sidedness[t] < min_sidedness; continue; end
			end

			#if test_bg_alpha < 1.0
			#	if pvalue(ChisqTest([alt[t] total[t] - alt[t]; bg_alt bg_total - bg_alt])) > test_ref_alpha; continue; end
			#end
			#chisq_p = pvalue(ChisqTest(
			#	[alt[t] total[t] - alt[t]; bg_alt bg_total - bg_alt]))
			#baseq[t] = round(Int32, min(-10 * log10(chisq_p), 100))

			somatic[t] = true
		end

		if keep_old_stars
			somatic |= map(c -> endswith(c, '*'), cols[notes_col+1:end])
		end

		if !any(somatic) && keep_all == false; continue; end

		print(join(cols[1:notes_col], '\t'))
		@printf("BG: %.3f%%. ", avg_ref_frac * 100)
		for s in 1:n_samples
			if call2_format
				print("\t$(alt[s]):$(total[s]):$(mapq[s]):$(baseq[s]):$(sidedness[s])")
			else
				print("\t$(alt[s]):$(total[s])")
			end

			if somatic[s]; print(":*"); end
		end
		println()
	end
end

function germline(vcf_file::IO, ref_regexps...; alt_reads=5, alt_frac=0.15, bg_ratio=20)
	line = skip_to_header(vcf_file); println(line)
	sample_col = findfirst(split(line, '\t') .== "NOTES") + 1
	samples = split(line, '\t')[sample_col:end]
	germline = falses(samples)
	for regex in ref_regexps; germline .|= ismatch.(Regex(regex), samples); end
	for line in eachline(vcf_file)
		cols = split(line, '\t')
		gtypes = [split(c, ':') for c in cols[sample_col:end]]
		total = [parse(Int32, gt[2]) for gt in gtypes]
		alt = [parse(Int32, gt[1]) for gt in gtypes]
		fracs = alt ./ total
		alt_gt = (fracs .>= alt_frac) .& (alt .>= alt_reads) .& germline
		if !any(alt_gt); continue; end

		# There might be non-reference samples that are not in alt_gt because
		# of the alt_reads threshold. We don't want to include these
		# samples in the background, so we only use alt_frac when picking
		# background samples.
		avg_alt = sum(alt[alt_gt]) / sum(total[alt_gt])
		bg = (fracs .< alt_frac) .& germline
		avg_ref = sum(alt[bg]) / sum(total[bg])
		if avg_alt / avg_ref < bg_ratio; continue; end

		print(join(cols[1:sample_col-1], '\t'))
		for s in 1:length(total)
			print("\t$(alt[s]):$(total[s])")
			if alt_gt[s]; print(":*"); end
		end
		print('\n')
	end
end

function above_background(vcf_file::IO, ref_regex...; alt_reads=3, alt_frac=0.1, bg_ratio=20)
	line = skip_to_header(vcf_file); println(line)
	sample_col = findfirst(split(line, '\t') .== "NOTES") + 1
	samples = split(line, '\t')[sample_col:end]

	bg = falses(samples)
	for regex in ref_regex; bg .|= ismatch.(Regex(regex), samples); end
	info("Calculating background based on $(sum(bg)) samples...")

	for line in eachline(vcf_file)
		cols = split(line, '\t')
		gtypes = map(c -> split(c, ':'), cols[sample_col:end])
		total = map(gt -> parse(Int, gt[2]), gtypes)
		alt = map(gt -> parse(Int, gt[1]), gtypes)
		fracs = alt ./ total
		avg_bg_frac = sum(alt[bg]) / sum(total[bg])
		alt_gt = (alt .>= alt_reads) .&
			(fracs .>= max(alt_frac, avg_bg_frac * bg_ratio))
		if !any(alt_gt); continue; end

		print(join(cols[1:sample_col-1], '\t'))
		for s in 1:length(total)
			print("\t$(alt[s]):$(total[s])")
			if alt_gt[s]; print(":*"); end
		end
		println()
	end
end

function heterozygous_snps(vcf_file::IO, ref_regex...; min_depth=30, min_frac=0.3, bg_fraction=0.1)
	line = skip_to_header(vcf_file); println(line)
	sample_col = findfirst(split(line, '\t') .== "NOTES") + 1
	samples = split(line, '\t')[sample_col:end]; S = length(samples)

	ref = falses(samples)
	for regex in ref_regex; ref .|= ismatch.(Regex(regex), samples); end
	info("Searching for heterozygous SNPs in $(sum(ref)) germline samples...")

	for line in eachline(vcf_file)
		cols = split(line, '\t')
		alt = zeros(Int, S); total = zeros(Int, S); starred = falses(S);
		unclear = 0
		for (s, col) in enumerate(cols[sample_col:end])
			parts = split(col, ':')
			alt[s] = parse(Int, parts[1])
			total[s] = parse(Int, parts[2])
			if ref[s] == false; continue; end
			if total[s] < min_depth; continue; end

			if alt[s] / total[s] < min_frac
				error_lower = Binomial(total[s], 0.01)
				if 1 - cdf(error_lower, alt[s] - 1) <= 0.05; unclear += 1; end
			elseif alt[s] / total[s] > 1 - min_frac
				error_upper = Binomial(total[s], 0.99)
				if cdf(error_upper, alt[s]) <= 0.05; unclear += 1; end
			end

			if unclear / S >= bg_fraction
				starred[:] = false; break;
			end

			hetz = Binomial(total[s], 0.5)
			if 0.05 <= cdf(hetz, alt[s]) <= 0.95; starred[s] = true; end
		end

		if !any(starred); continue; end

		print(join(cols[1:sample_col-1], '\t'))
		for s in 1:S
			print("\t$(alt[s]):$(total[s])")
			if starred[s]; print(":*"); end
		end
		print('\n')
	end
end

annovar_valid_funcs = Set(["exonic", "splicing", "intergenic", "intronic",
	"upstream", "downstream", "UTR5", "UTR3", "ncRNA_exonic", "ncRNA_intronic",
	"ncRNA_splicing", ""])

function translate_annovar_effect(gene, func, exonic_func, aa_change)
	for f in split(func, ';')
		if !in(f, annovar_valid_funcs)
			warn("Unrecognized variant effect '$f'.")
		end
	end

	effects = []
	if contains(func, "splicing"); push!(effects, "Splice site"); end
	if contains(func, "exonic")
		details = unique(map(m -> m.captures[1],
			eachmatch(r":(p\..+?)(,|$)", aa_change)))
		if exonic_func == "synonymous SNV"
			push!(effects, "Synonymous $(join(details, ", "))")
		elseif exonic_func == "nonsynonymous SNV"
			push!(effects, "Missense $(join(details, ", "))")
		elseif exonic_func == "stopgain"
			push!(effects, "Stopgain $(join(details, ", "))")
		elseif exonic_func == "stoploss"
			push!(effects, "Stoploss $(join(details, ", "))")
		elseif contains(exonic_func, "nonframeshift")
			push!(effects, "Non-frameshift indel $(join(details, ", "))")
		elseif contains(exonic_func, "frameshift")
			push!(effects, "Frameshift $(join(details, ", "))")
		elseif contains(func, "ncRNA_exonic")
			push!(effects, "Exonic (ncRNA)")
		elseif exonic_func == "unknown"
			push!(effects, "Exonic (unknown)")
		else
			error("Unrecognized effect: $([func, exonic_func])")
		end
	end
	if contains(func, "UTR3"); push!(effects, "3'-UTR"); end
	if contains(func, "UTR5"); push!(effects, "5'-UTR"); end
	if contains(func, "upstream"); push!(effects, "Upstream"); end
	if contains(func, "downstream"); push!(effects, "Downstream"); end
	if contains(func, "intronic"); push!(effects, "Intronic"); end
	if func == "intergenic"; push!(effects, "Intergenic"); gene = ""; end

	return (gene, join(effects, ". "))
end

function predict_effect(vcf_file::IO; genome="~/tools/annovar-2016-02-01/humandb/hg38")
	genome_version = basename(genome)
	humandb_dir = expanduser(dirname(genome))
	tmp_prefix = "annovar-$(hex(rand(0:2^16-1), 4))"
	original_path = "$(tmp_prefix).original"
	reformatted_path = "$(tmp_prefix).avinput"

	# Reformat the input file to an ANNOVAR-compatible format. In particular,
	# indels must be changed so they use '-' as the ref or alt allele.
	# Since the input file could be a pipe, we must write the original VCF
	# into a separate file that will be used to build the final annotated VCF.
	original_file = open(original_path, "w")
	reformatted_file = open(reformatted_path, "w")
	headers = split(skip_to_header(vcf_file), '\t')

	for line in eachline(vcf_file)

		write(original_file, line, '\n')
		c = split(line, '\t')
		start = parse(Int, c[2]); stop = start; ref = c[3]; alt = c[4]

		if length(ref) == 1 && length(alt) > 1 && ref[1] == alt[1]
			# Simple insertion
			ref = "-"; alt = alt[2:end];
		elseif length(ref) > 1 && length(alt) == 1 && ref[1] == alt[1]
			# Simple deletion
			ref = ref[2:end]; alt = "-"; start += 1;
		elseif length(ref) > 1 && length(alt) > 1
			# Block substitution
		end
		stop = start + length(ref) - 1
		write(reformatted_file, "$(c[1])\t$(start)\t$(stop)\t$(ref)\t$(alt)\n")
	end
	close(original_file); close(reformatted_file); close(vcf_file)

	run(pipeline(`table_annovar.pl $(reformatted_path) $(humandb_dir) --buildver $(genome_version) --remove --outfile $(tmp_prefix) --operation g --protocol refGene`, stderr=DevNull))

	anno_path = "$(tmp_prefix).$(genome_version)_multianno.txt"
	original_file = open(original_path);
	anno_file = open(anno_path); readline(anno_file)
	print("CHROM\tPOSITION\tREF\tALT\tGENE\tEFFECT\tNOTES\t")
	println(join(headers[6:end], '\t'))
	for line in eachline(anno_file)
		c = split(line, '\t')
		if length(c) >= 7
			gene, effect = translate_annovar_effect(c[7], c[6],
				length(c) >= 9 ? c[9] : "", length(c) >= 10 ? c[10] : "")
		else
			gene = ""; effect = "";
		end

		orig_line = readline(original_file)
		tab = search(orig_line, '\t')
		for k in 1:3; tab = search(orig_line, '\t', tab + 1); end
		print(orig_line[1:tab])
		print("$(gene)\t$(effect)\t")
		println(orig_line[tab+1:end])
	end
	close(anno_file); rm(original_path); rm(reformatted_path); rm(anno_path);
end

type TextAnnotations
	position::Vector{Int32}
	change::Vector{Int16}
	annotation::Vector{String}
	TextAnnotations() = new([], [], [])
end

type NumericAnnotations
	position::Vector{Int32}
	change::Vector{Int16}
	annotation::Vector{Float32}
	NumericAnnotations() = new([], [], [])
end

base_num = Dict{String, Int16}("A" => 1, "C" => 2, "G" => 3, "T" => 4)
encode_ref_alt(ref, alt) = get(base_num, ref, 0) * 16 + get(base_num, alt, 0)

function build_annotation_database(anno_path, db_path; prefix="",
	format="text")
	db = jldopen(db_path, "w", compress=true)
	anno_file = zopen(anno_path)
	if format == "text"
		build_annotation_database_text(anno_file, db)
	elseif format == "numeric"
		build_annotation_database_numeric(anno_file, db)
	elseif format == "blank"
		build_annotation_database_blank(anno_file, db)
	else
		error("Unrecognized annotation database format '$format' requested.")
	end
	write(db, "prefix", prefix)
	write(db, "format", format)
	close(db)
end

function build_annotation_database_blank(anno_file, db)
	chr = ""; anno = TextAnnotations()
	for line in eachline(anno_file)
		c = split(rstrip(line, '\t'))
		if c[1] != chr
			if chr != ""
				assert(issorted(anno.position))
				write(db, chr, anno)
			end
			anno = TextAnnotations(); chr = String(c[1])
			println(STDERR, "Chromosome $(chr)...")
		end
		push!(anno.position, parse(Int32, c[2]))
		push!(anno.change, encode_ref_alt(c[3], c[4]))
	end
	write(db, chr, anno)
end

function build_annotation_database_text(anno_file, db)
	chr = ""; anno = TextAnnotations()
	for line in eachline(anno_file)
		c = split(rstrip(line, '\t'))
		if c[1] != chr
			if chr != ""
				assert(issorted(anno.position))
				write(db, chr, anno)
			end
			anno = TextAnnotations(); chr = String(c[1])
			println(STDERR, "Chromosome $(chr)...")
		end
		push!(anno.position, parse(Int32, c[2]))
		push!(anno.change, encode_ref_alt(c[3], c[4]))
		push!(anno.annotation, c[5])
	end
	write(db, chr, anno)
end

function build_annotation_database_numeric(anno_file, db)
	chr = ""; anno = NumericAnnotations()
	for line in eachline(anno_file)
		c = split(rstrip(line, '\t'))
		if c[1] != chr
			if chr != ""
				assert(issorted(anno.position))
				write(db, chr, anno)
			end
			anno = NumericAnnotations(); chr = String(c[1])
			println(STDERR, "Chromosome $(chr)...")
		end
		push!(anno.position, parse(Int32, c[2]))
		push!(anno.change, encode_ref_alt(c[3], c[4]))
		push!(anno.annotation, parse(Float32, c[5]))
	end
	write(db, chr, anno)
end

function row_to_key( row::String)
	# Firstly, compare non-numeric parts (case-insensitive)
	# Next, compare numeric parts as integer
	# Next, compare position as integer
	chrom, pos = split(row, '\t')[1:2]
	return (filter(x -> !isnumber(x), lowercase(chrom)),
		    parse(Int32, filter(x -> isnumber(x), '0'*chrom)),
		    parse(Int32, pos))
end

# Sorts and prints VCF file rows by their chromosomal position
# Does not allow comment lines anywhere else than the beginning
function sort_by_position(vcf_file::IO)

	println( skip_to_header(vcf_file));
	rows = readlines(vcf_file)

	sort!(rows, by=x -> row_to_key(x))
	for row in rows println(row); end
end

function annotate(vcf_file::IO, db_path)
	db = jldopen(db_path, "r")
	prefix = read(db, "prefix")
	format = read(db, "format")
	headers = split(skip_to_header(vcf_file), '\t')
	notes_col = findfirst(headers .== "NOTES")
	println(join(headers, '\t'))

	chr = ""; anno = nothing
	for line in eachline(vcf_file)

		c = split(line, '\t')
		if drop_chr_prefix(c[1]) != chr
			chr = String(drop_chr_prefix(c[1]))
			if any(x -> chr == x, names(db))   # in() does not work
				anno = read(db, chr)
			else
				warn("No annotations found for chromosome $(chr).")
				anno = NumericAnnotations()
			end
		end
		matches = searchsorted(anno.position, parse(Int32, c[2]))
		if !isempty(matches)
			change = encode_ref_alt(c[3], c[4])
			for m in matches
				if anno.change[m] == change
					if format == "numeric"
						# See https://github.com/JuliaLang/julia/issues/14331
						c[notes_col] *= rstrip(@sprintf("%s:%g", prefix,
							anno.annotation[m])) * ". "
					elseif format == "text"
						c[notes_col] *= "$(prefix):$(anno.annotation[m]). "
					else
						c[notes_col] *= "$(prefix). "
					end
				end
			end
		end
		println(join(c, '\t'))
	end
end

# Filter VCF file lines based on if they are protein altering or not
# ARG: invert: Set to false to output protein altering mutations"
#              Set to true to output non protein altering mutations"
function protein_altering(vcf_file::IO; invert=false)
	line = skip_to_header(vcf_file); println(line)
	func_col = findfirst(split(line, '\t') .== "EFFECT")
	for line in eachline(vcf_file)
		cols = split(line, '\t')
		alters = ismatch(r"Missense|rameshift|Stopgain|Stoploss|Splice",
			cols[func_col])
		if !alters == invert; println(line); end
	end
end

function select_samples(invert::Bool, vcf_file::IO, regexps...)
	headers = split(skip_to_header(vcf_file), '\t')
	sample_col = findfirst(headers .== "NOTES") + 1
	keep = falses(length(headers))
	keep[1:sample_col-1] = true
	for col in sample_col:length(headers)
		keep[col] = any(rx -> ismatch(Regex(rx), headers[col]), regexps)
	end
	if invert; keep[sample_col:end] = .!keep[sample_col:end]; end
	println(join(headers[keep], '\t'))
	for line in eachline(vcf_file)
		cols = split(line, '\t')
		println(join(cols[keep], '\t'))
	end
end
keep_samples(vcf_file::IO, regexps...) = select_samples(false, vcf_file, regexps...)
discard_samples(vcf_file::IO, regexps...) = select_samples(true, vcf_file, regexps...)

function discard_shallow(vcf_file::IO, min_median_depth::Float64)
	line = skip_to_header(vcf_file); println(line);
	sample_col = findfirst(split(line, '\t') .== "NOTES") + 1
	for line in eachline(vcf_file)
		total = map(c -> parse(Int, split(c, ':')[2]),
			split(line, '\t')[sample_col:end])
		if median(total) >= min_median_depth; println(line); end
	end
end

function discard_blacklisted(vcf_file::IO, blacklist::IO)
	d = readdlm(blacklist, '\t', String)
	blacklist = Set(map(r -> d[r, 1:4][:], 1:size(d, 1)))

	line = skip_to_header(vcf_file); println(line);
	for line in eachline(vcf_file)
		c = split(line, '\t')
		if in(c[1:4], blacklist); continue; end
		println(line)
	end
end

function discard_indels(vcf_file::IO)
	line = skip_to_header(vcf_file); println(line);
	for line in eachline(vcf_file)
		c = split(line, '\t')
		if length(c[3]) == 1 && length(c[4]) == 1 && c[3] != "-" && c[4] != "-"
			println(line)
		end
	end
end

function discard_if_frequency_above(vcf_file::IO, db_path, threshold::Float64)
	db = jldopen(db_path, "r")
	println(skip_to_header(vcf_file))
	chr = ""; anno = nothing
	for line in eachline(vcf_file)
		c = split(line, '\t')
		if drop_chr_prefix(c[1]) != chr
			chr = String(drop_chr_prefix(c[1])); anno = read(db, chr)
		end
		matches = searchsorted(anno.position, parse(Int32, c[2]))
		if !isempty(matches)
			change = encode_ref_alt(c[3], c[4])
			if any(m -> anno.change[m] == change &&
				anno.annotation[m] >= threshold, matches)
				continue;
			end
		end
		println(join(c, '\t'))
	end
end

function inside(vcf_file::IO, region)
	m = match(r"([a-zA-Z0-9]+):(\d+)-(\d+)", region)
	if m == nothing
		if isfile(region) || isfifo(region)
			inside_bed_regions(vcf_file, region)
			return
		else
			error("Invalid region specified.")
		end
	end

	chr = m.captures[1]
	start = parse(Int, m.captures[2])
	stop = parse(Int, m.captures[3])

	line = skip_to_header(vcf_file); println(line)
	sample_col = findfirst(split(line, '\t') .== "NOTES") + 1
	for line in eachline(vcf_file)
		c = split(line, '\t'); pos = parse(Int, c[2])
		if c[1] == chr && start <= pos <= stop; println(line); end
	end
end

function inside_bed_regions(vcf_file::IO, regions_path)
	regions = Dict{String, Array{Region}}()
	for line in eachline(open(regions_path))
		c = split(rstrip(line), '\t')
		push!(get!(() -> Array{Region}(0), regions, c[1]),
			Region(parse(Int, c[2])+1, parse(Int, c[3])))
	end

	line = skip_to_header(vcf_file); println(line)
	sample_col = findfirst(split(line, '\t') .== "NOTES") + 1 + 1
	for line in eachline(vcf_file)
		c = split(line, '\t'); pos = parse(Int, c[2])
		chr_regions = get(() -> Array{Region}(0), regions, c[1])
		if any(r -> r.start <= pos <= r.stop, chr_regions)
			println(line)
		end
	end
end

function discard_contingent(vcf_path, alpha::Float64)
	vcf_file = zopen(vcf_path)
	line = skip_to_header(vcf_file); print(line)
	sample_col = findfirst(split(line, '\t') .== "NOTES") + 1
	for line in eachline(vcf_file)
		gtypes = map(c -> split(c, ':'), split(line, '\t')[sample_col:end])
		total = map(gt -> parse(Int, gt[2]), gtypes)
		alt = map(gt -> parse(Int, gt[1]), gtypes)
		valid = total .>= 1
		p = any(valid) ? pvalue(ChisqTest([alt[valid] total[valid] - alt[valid]])) : 1
		if alpha == 1
			print(line[1:end-1]); print("\t$(p)\n")
		elseif p <= alpha
			print(line)
		end
	end
end

# Adds a warning to the Notes column if the mappability in the reference
# genome is below a specified treshold (-> low mappability)
function mappability(vcf_file::IO, track_path::String, threshold::Float64)
	tracks = jldopen(track_path, "r");
	headers = split(skip_to_header(vcf_file), '\t')
	notes_col = findfirst(headers .== "NOTES")
	println(join(headers, '\t'))

	# We assume that the VCF file is sorted by chromosome.
	curr_chr = ""; track = zeros(UInt8, 0);
	for line in eachline(vcf_file)
		c = split(line, '\t')
		if c[1] != curr_chr
			curr_chr = String(c[1])
			track = read(tracks, curr_chr)
		end
		mappability = track[parse(Int, c[2])]
		if mappability < threshold
			c[notes_col] = "$(c[notes_col])Mappability < $(threshold)%. "
		end
		println(join(c, '\t'))
	end
end

# Adds a warning to the NOTES column if there are indels within
# dist distance of the alt read (line in VCF file).
#
# ARG: dist : largest distance that is considered
#             being "nearby"
function nearby_indels(vcf_file::IO; dist=10)

	line = skip_to_header(vcf_file);
	println(line)
	headers = split(line, '\t')
	notes_col = findfirst(split(line, '\t') .== "NOTES")
	warn_msg = "INDEL within $dist bases. "

	#notes_col = findfirst(split(skip_to_header(vcf_file), '\t') .== "NOTES")

	prev_lines = Array{Array{String}}(0)

	for line in eachline(vcf_file)

		if startswith(line, "#") println(line); continue; end
		c = split(line, '\t')
		pos = parse(Int, c[2])
		starred = map(x -> endswith(x, '*'), c[notes_col+1:end])

		# Forget old lines that are so far behind they cannot be affected.
		keep = trues(length(prev_lines))
		for (k, prev_line) in enumerate(prev_lines)
			# On different chromosome or too distant
			if c[1] != prev_line[1] || parse(Int, prev_line[2]) < pos - dist
				keep[k] = false
				println(join(prev_line, '\t'))
			end
		end
		prev_lines = prev_lines[keep]

		# Add warning to old lines if current line is an indel.
		if length(c[3]) != length(c[4])
			for prev_line in prev_lines
				if any(s -> starred[s] && endswith(prev_line[notes_col+s], '*'), 1:length(starred))
					prev_line[notes_col] *= warn_msg
				end
			end
		end

		# Add warning to current line if any old lines were indels.
		for prev_line in prev_lines
			if length(prev_line[3]) != length(prev_line[4]) &&
				any(s -> starred[s] && endswith(prev_line[notes_col+s], '*'), 1:length(starred))
				c[notes_col] *= warn_msg
				break
			end
		end

		push!(prev_lines, c)
	end

	for prev_line in prev_lines; println(join(prev_line, '\t')); end
end

function statistics(vcf_path)
	vcf_file = zopen(vcf_path)
	headers = split(rstrip(skip_to_header(vcf_file)), '\t')
	sample_col = findfirst(headers .== "NOTES") + 1
	samples = headers[sample_col:end]; S = length(samples)
	total_subs = zeros(Int, S)
	possible_subs = zeros(Int, S, 6)
	total_indels = zeros(Int, S)

	for line in eachline(vcf_file)
		c = split(rstrip(line), '\t')
		alt = map(gt -> endswith(gt, '*'), c[sample_col:end])
		mut_class = mutation_class(c[3], c[4])
		if mut_class > 0
			total_subs += alt
			possible_subs[:, mut_class] += alt
		elseif mut_class == 0
			total_indels += alt
		end
	end

	println("SAMPLE\tTOTAL SUBSTITUTIONS\tTOTAL INDELS\tC>A\tC>G\tC>T\tT>A\tT>C\tT>G")
	for s in 1:S
		print("$(samples[s])\t$(total_subs[s])\t$(total_indels[s])")
		for k in 1:6; print("\t$(possible_subs[s, k])"); end
		println()
	end
end

# CALL2 output format: "alt_reads:total_reads:mapq:baseq:sidedness:significant"
function fractions(vcf_file::IO)

	line = skip_to_header(vcf_file); println(line)
	sample_col = findfirst(split(line, '\t') .== "NOTES") + 1
	for line in eachline(vcf_file)
		cols = split(line, '\t')
		gtypes = [split(c, ':') for c in cols][sample_col:end]

		total = [parse(Int32, gt[2]) for gt in gtypes]
		alt = [parse(Int32, gt[1]) for gt in gtypes]

		print(join(cols[1:(sample_col-1)], '\t'))

		for s in 1:length(gtypes)
			frac = alt[s] / total[s]
			@printf("\t%.1f%% (%d)", isnan(frac) ? 0.0 : frac * 100, total[s])
			#CALL2 support
			if length(gtypes[s]) >= 5 && parse( Int32, gtypes[s][ 3]) > 0; @printf(" [mq:%s,bq:%s,sd:%s]", gtypes[s][3], gtypes[s][4], gtypes[s][5] ); end
			if length(gtypes[s]) >= 3 && gtypes[s][ length(gtypes[s])] == "*"; print(" *"); end
		end
		println()
	end
end

function error_rate(bam_path, genome_path; min_mapq=10)
	bam_path = expanduser(bam_path)
	genome_path = expanduser(genome_path)
	spileup = expanduser("~/tools/pypette/compiled/spileup")

	# 2D histogram, mutant read count on rows, mutant read fraction on columns
	# Histogram bins are left-exclusive, right-inclusive.
	read_thresholds = 0:100
	frac_thresholds = 0:0.001:0.10
	histogram = zeros(Int, length(read_thresholds), length(frac_thresholds))

	for line in eachline(pipeline(`samtools mpileup -d 100000 -A -x -R -sB -q$(min_mapq) -f $(genome_path) $(bam_path)`, `$(spileup) 0 $(min_mapq)`))
		cols = split(rstrip(line), '\t')
		if length(cols) < 4; continue; end
		#if cols[3] == "N"; print(line); continue; end
		assert(length(cols) == 4)
		assert(cols[3] != "N")

		pileup = cols[4]
		if pileup == ""; continue; end
		t = split(pileup, ' ')

		total_reads = 0
		alt_reads = 0

		for a in 1:3:length(t)
		 	count = parse(Int, t[a+1])
		 	total_reads += count
		 	if t[a] != "."; alt_reads += count; end
		end

		alt_frac = alt_reads / total_reads
		read_bin = Int(min(alt_reads + 1, 101))
		frac_bin = min(ceil(Int, alt_frac / 0.001) + 1, 101)
		histogram[read_bin, frac_bin] += 1

		# allele_reads = zeros(Int32, max_alleles)
		# alleles = Array{String}(0)

		# pileup = cols[4]
		# if pileup == ""; continue; end
		# t = split(pileup, ' ')
		# for a in 1:3:length(t)
		# 	count = int(t[a+1])
		# 	total_reads[s] += count
		# 	if t[a] == "."; continue; end
		# 	idx = findfirst(allele -> allele == t[a], alleles)
		# 	if idx == 0 && length(alleles) >= max_alleles; continue; end
		# 	if idx == 0; push!(alleles, t[a]); idx = length(alleles); end
		# 	allele_reads[s, idx] = count
		# end

		# for a in 1:length(alleles)
		# end
	end

	writedlm(STDOUT, histogram)
end

function group_by_snp_profile(hetz_snp_vcf::IO)
	headers = split(rstrip(skip_to_header(hetz_snp_vcf)), '\t')
	sample_col = findfirst(headers .== "NOTES") + 1
	samples = headers[sample_col:end]; S = length(samples);

	# The matrix stores genotypes for every sample:
	# 0 = unknown, 1 = homz ref, 2 = hetz, 3 = homz alt
	println("Constructing genotype matrix...")
	genotypes = zeros(Int8, 1_000_000, S); r = 0
	for line in eachline(hetz_snp_vcf)
		c = split(rstrip(line, '\n'), '\t')
		r += 1
		for (s, cell) in enumerate(c[sample_col:end])
			parts = split(cell, ':')
			alt = parse(Int, parts[1]); total = parse(Int, parts[2])
			if total < 30; genotypes[r, s] = 0; continue; end
			b = Binomial(total, 0.5)
			if 0.05 <= cdf(b, alt) <= 0.95
				genotypes[r, s] = 2
			elseif alt / total >= 0.8
				genotypes[r, s] = 3
			elseif alt / total <= 0.2
				genotypes[r, s] = 1
			end
		end
	end
	genotypes = genotypes[1:r, :]

	for i in 1:S
		any_matches = false
		for j in 1:S
			if j == i; continue; end
			matched = 0; total = 0;
			for r in 1:size(genotypes, 1)
				if genotypes[r, i] == 0 || genotypes[r, j] == 0; continue; end
				total += 1
				matched += (genotypes[r, i] == genotypes[r, j])
			end
			if matched / total >= 0.95 && total >= 50
				any_matches = true
				if j > i      # Don't print pairs twice
					println("$(samples[i]) <-> $(samples[j]): $matched / $total matched")
				end
			end
		end

		if any_matches == false
			println("$(samples[i]) does not match with any other samples.")
		end
	end
end

function mutation_rate(vcf_file::IO, coverage_histogram_dir; alt_reads=8)
	vcf = read_vcf(vcf_file); S = length(vcf.sample);
	human_chr = vcat(map(x -> "chr$(x)", 1:22), ["chrX", "chrY"]);
	println("SAMPLE\tPROTEIN ALTERING\tSILENT\tPROTEIN ALTERING (RATE)\tSILENT (RATE)")
	for s in 1:S
		info("Analyzing $(vcf.sample[s])...")
		protein_altering = 0; silent = 0;
		protein_altering_rate = 0; silent_rate = 0;
		histogram = zeros(5000)   # Bins are 0, 1, 2, ..., 4999
		genome_len = 0
		histogram_file = "$(coverage_histogram_dir)/$(vcf.sample[s]).tsv"
		if !isfile(histogram_file)
			warn("No coverage histogram found for $(vcf.sample[s])...")
			continue
		end
		d = readdlm(histogram_file, '\t')
		for chr in human_chr
			genome_len += d[findfirst(x -> x == chr, d[:, 1]), 4]
			for r in 1:size(d, 1)
				if d[r, 1] != chr || d[r, 2] >= 5000; continue; end
				histogram[d[r, 2] + 1] += d[r, 3]
			end
		end

		for r in 1:size(vcf, 1)
			if !vcf.star[r, s]; continue; end
			min_depth = ceil(Int, alt_reads / (vcf.alt[r, s] / vcf.total[r, s]))
			amenable_mbs = (genome_len - sum(histogram[1:min_depth-1])) / 1e6
			info("$(vcf.chromosome[r]):$(vcf.position[r]) - amenable $(amenable_mbs) Mb")
			if is_protein_altering(vcf.effect[r])
				protein_altering += 1
				protein_altering_rate += 1 / amenable_mbs
			else
				silent += 1
				silent_rate += 1 / amenable_mbs
			end
		end

		@printf("%s\t%d\t%d\t%.2f\t%.2f\n", vcf.sample[s], protein_altering, silent, protein_altering_rate, silent_rate)
	end
end

function discard_sketchy_silent(vcf_file::IO)
	line = skip_to_header(vcf_file); println(line)
	notes_col = findfirst(split(line, '\t') .== "NOTES")
	effect_col = findfirst(split(line, '\t') .== "EFFECT")
	for line in eachline(vcf_file)
		cols = split(line, '\t')
		if !is_protein_altering(cols[effect_col])
			if contains(cols[notes_col], "INDEL within"); continue; end
			if contains(cols[notes_col], "Mappability < "); continue; end
		end
		println(line)
	end
end

function sort_samples_by_name(vcf_file::IO)
	line = skip_to_header(vcf_file)
	headers = split(line, '\t')
	notes_col = findfirst(headers .== "NOTES")
	samples = headers[notes_col+1:end]
	order = vcat(1:notes_col, notes_col .+ sortperm(samples))
	println(join(headers[order], '\t'))
	for line in eachline(vcf_file)
		cols = split(line, '\t')
		println(join(cols[order], '\t'))
	end
end

subcommands(somatic, germline, above_background, heterozygous_snps,
	keep_samples, discard_samples,
	discard_shallow, discard_if_frequency_above, discard_blacklisted,
	discard_indels, discard_contingent, discard_sketchy_silent,
	mappability, nearby_indels,
	inside, protein_altering, fractions, statistics,
	predict_effect, annotate, build_annotation_database, error_rate,
	group_by_snp_profile, mutation_rate, sort_by_position, sort_samples_by_name)


module ReinforcementLearningZoo

const RLZoo = ReinforcementLearningZoo
export RLZoo

using CircularArrayBuffers
using ReinforcementLearningBase
using ReinforcementLearningCore
using Setfield: @set
using StableRNGs
using Logging
using Flux.Losses
using Dates
using Random
using Random: shuffle
using CUDA
using Zygote
using Zygote: ignore
using Flux
using Flux: onehot, normalise
using StatsBase
using StatsBase: sample, Weights, mean
using LinearAlgebra: dot
using MacroTools
using Distributions: Categorical, Normal, logpdf
using StructArrays

include("patch.jl")
include("algorithms/algorithms.jl")

using Requires

# dynamic loading environments
function __init__()
    @require ReinforcementLearningEnvironments = "25e41dd2-4622-11e9-1641-f1adca772921" begin
        include("experiments/rl_envs/rl_envs.jl")
        @require ArcadeLearningEnvironment = "b7f77d8d-088d-5e02-8ac0-89aab2acc977" include("experiments/atari/atari.jl")
        # @require OpenSpiel = "ceb70bd2-fe3f-44f0-b81f-41608acaf2f2" include("experiments/open_spiel/open_spiel.jl")
    end
end

end # module


module GRASP
using DocStringExtensions

export angularmomentum, parity # from module Symmetries
export nelectrons, nexcitations, maxelectrons

using AtomicLevels: AtomicLevels, Orbital, RelativisticOrbital, IntermediateTerm, degeneracy,
    κ2ℓ, κ2j
using HalfIntegers: HalfInteger, HalfInt

#
# Includes and imports related to compiled Fortran code.
# ------------------------------------------------------------------------------------------

let grasp_path_file = joinpath(@__DIR__, "../deps/grasp-path.jl")
    isfile(grasp_path_file) || error("deps/grasp-path.jl does not exist. Run `Pkg.build(\"GRASP\")`.")
    include(grasp_path_file)
end

const libgrasp_so = joinpath(@__DIR__, "..", "deps", "libgrasp.so")
isfile(libgrasp_so) || error("$(libgrasp_so) does not exist. Run `Pkg.build(\"GRASP\")`.")

#
# Additional includes and imports.
# ------------------------------------------------------------------------------------------

include("Symmetries.jl")
import .Symmetries: Parity, AngularMomentum, AngularSymmetry
import .Symmetries: parity, angularmomentum, absdiff

const SPECTROSCOPIC_NAMES = "s p d f g h i k l m n o q r t u v" |> split
function specname(l::Integer)
    @assert l >= 0
    SPECTROSCOPIC_NAMES[l+1]
end

"""
    nelectrons(obj)

Calculate the number of electrons in `obj`.
"""
function nelectrons end

"""
    nexcitations(from, to)

Calculate the number of excitations needed to go from `from` to `to`.
"""
function nexcitations end

"""
    maxelectrons(obj)

Return the maximum number of electrons.
"""
function maxelectrons end

include("rcsfs.jl")
include("Configurations.jl")
include("rmix.jl")
include("rwfn.jl")
include("rcsfparser.jl")

# Fixed Humanize.digitsep --- fixed for negative numbers
# from: https://github.com/IainNZ/Humanize.jl
function digitsep(value::Integer, sep = ",", k = 3)
    isnegative = value < 0
    value = string(abs(value))
    n = length(value)
    starts = reverse(collect(n:-k:1))
    groups = [value[max(x-k+1, 1):x] for x in starts]
    return (isnegative ? "-" : "") * join(groups, sep)
end

"""
    $(SIGNATURES)

Converts the energy value `x` from Hartrees to Kaysers (aka 1/cm),
using the conversion ``1~\\textrm{Ha} = 219474.63137~\\textrm{cm}^{-1}``.

# Examples

```jldoctest
julia> GRASP.hartree2kayser(1.0)
219474.63137
```
"""
hartree2kayser(x) = 219474.63137*x

"""
    $(SIGNATURES)

Like [`hartree2kayser`](@ref), but produces a humanized string with thousands separators.

# Examples

```jldoctest
julia> GRASP.h2k_humanize(1.0)
"219,475"
```
"""
h2k_humanize(x) = digitsep(round(Int, hartree2kayser(x)))

end # module


# # Instructions for Assignment02
#
# ## Introduction
#
# In this assignment,
# you will start to put the pieces together.
# Note: this file is runnable in its current state,
# but is incomplete.
# You can run the file from the command line in script mode,
# copy code into the REPL in interactive mode,
# or use the VS Code julia extension to run individual lines.

# The following code is used for set up;
# there is no need to change it.
# If you _do_ change it to figure out what it's doing,
# make sure you change it back or your assignment may not work correctly

## To view additional instructions for this assignment, #jl
## see https://wellesley-bisc195.github.io/assignments/Assignment02 #jl

## To view a rendered version of this document, #jl
## see https://wellesley-bisc195.github.io/assignments/Assignment02/#assignment02_code #jl

using Random
Random.seed!(42)

function generate_sequence(len)
    seq = join(rand(['A','C','G','T'], len))
    println("Your sequence is: ", seq)
    return seq
end

my_seq = generate_sequence(20)

# ## Question 1
#
# The code above generates a random 20nt DNA sequence,
# and assigns it to the variable `my_seq`.
# What kind of data type is `my_seq` (eg `Float64`, `Int64`, `String`, or something else)?
# Assign the variable `question1` to the correct type.
#
# Hint1: rather than guess, you can just use the `typeof()` function.

question1 = ""

# Hint2: If you're having trouble,
# note that your answer should not contain quotes.
# Eg. if you think that `seq` has the type `MyType`,
# your code should read `question1 = MyType` and not `question1 = "MyType"`

# ## Question 2
#
# The function below has a documentation string or "doc string",
# which explains what the function should do.
# Many julia functions have doc strings,
# which makes it easy to get help right from the REPL
#
# Open a julia REPL and type `?`.
# You will see the `julia>` prompt change to `help?>`.
# Now type `println` and hit enter.
# The docstring of `println` is printed
# and you're returned to a julia prompt
#
# Evaluate the following function, including the docstring, in the REPL
# (that is, copy from the triple quotes down to the "end"
# and paste it into the REPL, then hit enter)

"""
    question2(sequence)

Checks if `sequence` is a String.

Example
≡≡≡≡≡≡≡≡≡

    julia> question2("hello")
    true

    julia> question2(1001)
    false

    julia> if question2("I'm a string!")
               println("yeah, that was a string")
           end
    yeah, that was a string
"""
function question2()
    ## put your code here
end

# At the moment, this function doesn't do what it says it should.
# Actually, it doesn't do anything.
# Fix the function so that it takes one argument
# and returns `true` if the argument is a String,
# and returns `false` otherwise.
# You should be able to run the examples in the docstring
# and get the correct answer,
# and to run this function using `my_seq` as the argument
# (it should return `true`).
#
# Hint: You haven't explicitly encountered a way to check if a type is another type.
# Learning to search for answers is a key programming skill!
# Try searching "check if type is string julia".
# Typically, when I see search results, I prioritize in this order:
# 1. stackoverflow.com
# 2. docs.julialang.org
# 3. discourse.julialang.org
# 4. github.com

# ## Question 3
#
# Let's try something a bit more complicated.
# The `question3()` function has a doc string,
# but is also incomplete.
# I've writtten some code to get you started,
# complete it so that the examples work as expected in the docstring.

"""
    question3(sequence)

Calculates the GC ratio of a DNA sequence.
The GC ratio is the total number of G and C bases divided by the total length of the sequence.
For more info about GC content, see here: https://en.wikipedia.org/wiki/GC-content

Example
≡≡≡≡≡≡≡≡≡≡

    julia> question3("AATG")
    0.25

    julia> question3("CCCGG")
    1.0

    julia> question3("ATTA")
    0.0
"""
function question3(sequence)
    ## throw an error if the string contains anything other than ACGT
    if any(c-> !in(c, ['A','C','G','T']), sequence)
        throw(ArgumentError("Sequence must only contain ACGT"))
    end

    ## change line to assign `seqlength` to the length of `sequence` instead of `1`
    ## If you're stuck, search for "length of string julia"
    seqlength = 1

    ## count the number of G's
    gs = count(==('G'), sequence)
    ## count the number of C's
    cs = count(==('C'), sequence)

    return gs + cs / seqlength # something is wrong with this line...
end

# ## Question 4
#
# This process should be familiar by now.
#
# Hint: you do not need to re-write code if you already have it.
# That is, don't copy any code from `question3`,
# just call it!

"""
    question4(sequence)

Calculates the GC content of a DNA sequence
and prints it to the screen.

Example
≡≡≡≡≡≡≡≡≡≡

    julia> question4("AATC")
    Sequence:
    AATC
    GC Content:
    0.25

    julia> question4("CCCGG")
    Sequence:
    CCCGG
    GC Content:
    1.0

    julia> question4("CCCGT"); # note the semicolon
    Sequence:
    CCCGG
    GC Content:
    0.8
"""
function question4(sequence)
    ## Your code here
end

## Question 5
#
# The following are related to questions in the exercises
# from chapter 2 of _Think Julia_.
# The format of the varables below refer to specific exercise questions.
# For example, `ce_2_1_1` would refer to chapter excercise 2-2, question 1.
#
# > We’ve seen that `n = 42` is legal. What about `42 = n`?
#
# Assign the variable to `true` if legal, `false` if illegal

ce_2_2_1 = Bool

# > How about `x = y = 1`?
#
# Assign the variable to `true` if legal, `false` if illegal

ce_2_2_2 = Bool

# > The volume of a sphere with radius `r` is `4/3πr^3`.
# What is the volume of a sphere with radius `5`?
#
# Assign the variable to the correct answer

ce_2_3_1 = Float64

# > Suppose the cover price of a book is \$24.95,
# but bookstores get a 40 % discount.
# Shipping costs \$3 for the first copy and 75 cents for each additional copy.
# What is the total wholesale cost for 60 copies?
#
# Assign the variable to the correct answer.

ce_2_3_2 = Float64

## Question 6
#
# Write a function that calculates the cost of books,
# given some list price, discount rate, and number.

"""
    bookprice(list, discount, count)

Calculates the price of books, given some list price, discount rate,
and number of books ordered.

Example
≡≡≡≡≡≡≡≡≡≡

    julia> bookprice(24.95, 0.4, 60)
    # this should be the same answer you got for `c_2_3_2`

    julia> bookprice(24.95, 0.4, 1)
    17.97

    julia> bookprice(24.95, 0.4, 2)
    33.69

    julia> bookprice(1, 0, 1)
    4.0
"""
function bookprice(list, discount, count)
    ## your code here
end


export SACPolicy

mutable struct SACPolicy{
    BA<:NeuralNetworkApproximator,
    BC1<:NeuralNetworkApproximator,
    BC2<:NeuralNetworkApproximator,
    P,
    R<:AbstractRNG,
} <: AbstractPolicy
    policy::BA
    qnetwork1::BC1
    qnetwork2::BC2
    target_qnetwork1::BC1
    target_qnetwork2::BC2
    γ::Float32
    τ::Float32
    α::Float32
    batch_size::Int
    start_steps::Int
    start_policy::P
    update_after::Int
    update_freq::Int
    automatic_entropy_tuning::Bool
    lr_alpha::Float32
    target_entropy::Float32
    update_step::Int
    rng::R
    # Logging
    reward_term::Float32
    entropy_term::Float32
end

"""
    SACPolicy(;kwargs...)

# Keyword arguments

- `policy`, used to get action.
- `qnetwork1`, used to get Q-values.
- `qnetwork2`, used to get Q-values.
- `target_qnetwork1`, used to estimate the target Q-values.
- `target_qnetwork2`, used to estimate the target Q-values.
- `start_policy`,
- `γ::Float32 = 0.99f0`, reward discount rate.
- `τ::Float32 = 0.005f0`, the speed at which the target network is updated.
- `α::Float32 = 0.2f0`, entropy term.
- `batch_size = 32`,
- `start_steps = 10000`,
- `update_after = 1000`,
- `update_freq = 50`,
- `automatic_entropy_tuning::Bool = false`, whether to automatically tune the entropy.
- `lr_alpha::Float32 = 0.003f0`, learning rate of tuning entropy.
- `action_dims = 0`, the dimensionality of the action. if `automatic_entropy_tuning = true`, must enter this parameter.
- `update_step = 0`,
- `rng = Random.GLOBAL_RNG`,

`policy` is expected to output a tuple `(μ, logσ)` of mean and
log standard deviations for the desired action distributions, this
can be implemented using a `GaussianNetwork` in a `NeuralNetworkApproximator`.

Implemented based on http://arxiv.org/abs/1812.05905
"""
function SACPolicy(;
    policy,
    qnetwork1,
    qnetwork2,
    target_qnetwork1,
    target_qnetwork2,
    start_policy,
    γ = 0.99f0,
    τ = 0.005f0,
    α = 0.2f0,
    batch_size = 32,
    start_steps = 10000,
    update_after = 1000,
    update_freq = 50,
    automatic_entropy_tuning = true,
    lr_alpha = 0.003f0,
    action_dims = 0,
    update_step = 0,
    rng = Random.GLOBAL_RNG,
)
    copyto!(qnetwork1, target_qnetwork1)  # force sync
    copyto!(qnetwork2, target_qnetwork2)  # force sync
    if automatic_entropy_tuning
        @assert action_dims != 0
    end
    SACPolicy(
        policy,
        qnetwork1,
        qnetwork2,
        target_qnetwork1,
        target_qnetwork2,
        γ,
        τ,
        α,
        batch_size,
        start_steps,
        start_policy,
        update_after,
        update_freq,
        automatic_entropy_tuning,
        lr_alpha,
        Float32(-action_dims),
        update_step,
        rng,
        0f0,
        0f0,
    )
end

# TODO: handle Training/Testing mode
function (p::SACPolicy)(env)
    p.update_step += 1

    if p.update_step <= p.start_steps
        p.start_policy(env)
    else
        D = device(p.policy)
        s = state(env)
        s = Flux.unsqueeze(s, ndims(s) + 1)
        # trainmode:
        action = dropdims(p.policy(p.rng, s; is_sampling=true), dims=2) # Single action vec, drop second dim

        # testmode:
        # if testing dont sample an action, but act deterministically by
        # taking the "mean" action
        # action = dropdims(p.policy(s)[1], dims=2)
    end
end

function RLBase.update!(
    p::SACPolicy,
    traj::CircularArraySARTTrajectory,
    ::AbstractEnv,
    ::PreActStage,
)
    length(traj) > p.update_after || return
    p.update_step % p.update_freq == 0 || return
    inds, batch = sample(p.rng, traj, BatchSampler{SARTS}(p.batch_size))
    update!(p, batch)
end

function RLBase.update!(p::SACPolicy, batch::NamedTuple{SARTS})
    s, a, r, t, s′ = send_to_device(device(p.qnetwork1), batch)

    γ, τ, α = p.γ, p.τ, p.α

    a′, log_π = p.policy(p.rng, s′; is_sampling=true, is_return_log_prob=true)
    q′_input = vcat(s′, a′)
    q′ = min.(p.target_qnetwork1(q′_input), p.target_qnetwork2(q′_input))

    y = r .+ γ .* (1 .- t) .* vec(q′ .- α .* log_π)

    # Train Q Networks
    q_input = vcat(s, a)

    q_grad_1 = gradient(Flux.params(p.qnetwork1)) do
        q1 = p.qnetwork1(q_input) |> vec
        mse(q1, y)
    end
    update!(p.qnetwork1, q_grad_1)
    q_grad_2 = gradient(Flux.params(p.qnetwork2)) do
        q2 = p.qnetwork2(q_input) |> vec
        mse(q2, y)
    end
    update!(p.qnetwork2, q_grad_2)

    # Train Policy
    p_grad = gradient(Flux.params(p.policy)) do
        a, log_π = p.policy(p.rng, s; is_sampling=true, is_return_log_prob=true)
        q_input = vcat(s, a)
        q = min.(p.qnetwork1(q_input), p.qnetwork2(q_input))
        reward = mean(q)
        entropy = mean(log_π)
        ignore() do
            p.reward_term = reward
            p.entropy_term = entropy
        end
        α * entropy - reward
    end
    update!(p.policy, p_grad)

    # Tune entropy automatically
    if p.automatic_entropy_tuning
        p.α -= p.lr_alpha * mean(-log_π .- p.target_entropy)
    end

    # polyak averaging
    for (dest, src) in zip(
        Flux.params([p.target_qnetwork1, p.target_qnetwork2]),
        Flux.params([p.qnetwork1, p.qnetwork2]),
    )
        dest .= (1 - τ) .* dest .+ τ .* src
    end
end


module ROSCar1D
export main

using RobotOS
using PyCall
roslib = pyimport("roslib")
roslib.load_manifest("amrl_msgs")
@rosimport amrl_msgs.msg: AckermannCurvatureDriveMsg
@rosimport nav_msgs.msg: Odometry
rostypegen(@__MODULE__)
#import .amrl_msgs.msg: AckermannCurvatureDriveMsg
#import .nav_msgs.msg: Odometry
using .amrl_msgs.msg: AckermannCurvatureDriveMsg
using .nav_msgs.msg: Odometry
using StaticArrays
using Unrolled
using Random

using OneVision
using OneVision.Car1DExample

global x_cur = CarX(0.0, 0.0) # CarX needs to be immutalble - will find a container to wrap this around.

function odom_callback(msg::Odometry)
    x_cur.pos = sqrt(square(msg.position.position.x) + square(msg.position.position.y))
    x_cur.velocity = 2.0 * atan(msg.pose.pose.orientation.z, msg.pose.pose.orientation.w)
end

function ros_publish_actuation(pub::Publisher, act::CarU, state::CarX)
    # interperate
    next_vel = state.velocity + act.acc*0.05
    msg = AckermannCurvatureDriveMsg()
    msg.velocity = next_vel
    publish(pub, msg)
    publish(pub, msg)
    publish(pub, msg)
end

"The simple wall dynamics model used by OneVision."
struct WallObsModel <: ObsDynamics end

OneVision.obs_forward(dy::WallObsModel, x::CarX, z::CarZ, t::𝕋) = z

function make_ros_controllers(times, freq::ℝ; noise=0.0, plot_result=false, log_prediction=false, N=1)
    delta_t = 1 / freq
    times::Vector{𝕋} = collect(times)
    idx_to_time(xs) = (xs .- 1) .* delta_t
    t0, t_end = times[1], times[end]
    rng = MersenneTwister(1234)

    agent_info(id) = begin
        acc_noise = randn(rng, ℝ, 1 + t_end - t0) * noise
        sys_dy = car_system(delta_t, t -> CarX(0.0, acc_noise[1 + t - t0]))
        obs_dy =
            if id == 1; WallObsDynamics(wall_position=30.0, detector_range=8.0)
            else WallObsModel() end

        sys_dy, obs_dy
    end

    world_dynamics = WorldDynamics([agent_info(i) for i in 1:N])
    init_states = fill((CarX(0.0, 0.0), CarZ(0.0, 0.0), CarU(0.0)), N)
    delays = DelayModel(obs=2, act=3, com=4)  # DelayModel(obs=1, act=2, com=5)
    world_model = WorldDynamics(
        fill((car_system(delta_t), WallObsModel()), N))

    comps = ["pos", "velocity", "acceleration", "obstacle"]
    function record_f(xs, zs, us)
        [(x -> x.pos).(xs) (x -> x.velocity).(xs) (u -> u.acc).(us) (z -> z.distance).(zs)]
    end

    H = 20
    πc = LeaderFollowerControl(warm_up_time=delays.act)

    u_weights = fill(CarU(1.0), N)
    x_weights = fill(CarX(1.0, 1.0), N)

    controller_framework = OvCF{N,CarX{ℝ},CarZ{ℝ},CarU{ℝ},H}(
            πc,
            world_model, delays, x_weights, u_weights,
            FuncT(Tuple{ℕ,𝕋,CarX{ℝ},CarZ{ℝ}}, Bool) do (id, t, _, _)
                false
            end
        )
    return world_dynamics, delays, controller_framework, init_states, times
end

function loop(pub::Publisher)
    freq = 20.0
    wd, delay_model, cf, init_status, times = make_ros_controllers(1:20, freq)
    (controllers, msg_qs) = make_controllers(cf, init_status, times[1])
    println(controllers)
    make_agent(id::ℕ)::AgentState = begin
        (x₀, z₀, u₀) = init_status[id]
        AgentState(
            state_queue=constant_queue(x₀, delay_model.obs),
            obs_queue=constant_queue(z₀, delay_model.obs),
            act_queue=constant_queue(u₀, delay_model.act),
            msg_queue=msg_qs[id],
            controller=controllers[id],
        )
    end
    N = 1
    agents = map(make_agent, SVector{N}(1:N))
    init_states = SVector{N}(init_status)
    #a::SVector{N,Tuple{CarX,CarZ,CarU}} = init_states
    loop_rate = Rate(freq)
    tmp = first(agents[1].msg_queue)
    transmition = MMatrix{N,N,typeof(tmp)}(undef)
    #print(a)
    Each = MVector{N}
    xs, zs, us = @unzip(MVector(init_states), Each{Tuple{CarX{Float64},CarZ{Float64},CarU{Float64}}})

    for t in 1:20
        if is_shutdown()
            error("FATAL: ROS node has been shutdown")
            exit(0)
        end
        rossleep(loop_rate)
        for i in 1:length(agents)  # unroll away dynamic dispatch at compile time!
            a = agents[i]
            x = xs[i]
            u = us[i]
            z = zs[i]
            ms = first(a.msg_queue)
            u, ms′ = control!(a.controller, x, z, ms)
            #transmition[:,i] = ms′
            ros_publish_actuation(pub, u, x)
        end

    end
    """
    loop_rate = Rate(20.0)
    for t in 1:20
        if is_shutdown()
            error("FATAL: ROS node has been shutdown")
        end
        rossleep(loop_rate) # handle callbacks here - busy loop. Potentially very inefficient
        # at this point, messages should be handled, - run the simulation framework
        print(xs[0])
        xs[0] = x_cur # assign the current observation

        # observe, control, and send messages
        transmition = MMatrix{N,N,Msg}(undef)  # indexed by (receiver, sender)
        for i in 1:length(agents)  # unroll away dynamic dispatch at compile time!
            a = agents[i]
            x = pushpop!(a.state_queue, xs[i])
            z = pushpop!(a.obs_queue, zs[i])
            ms = first(a.msg_queue)
            u, ms′ = control!(a.controller, x, z, ms)
            us[i] = pushpop!(a.act_queue, u)
            transmition[:,i] = ms′
        end
        # receive messagees and update world states
        for i in 1:length(agents)
            a = agents[i]
            pushpop!(a.msg_queue, convert(Vector{Msg}, transmition[i,:]))
            ros_publish_actuation(pub, act, cur_state) # needs to be obtained from ROS
            zs[i] = zs[i] # do not update the wall pos for now #obs_forward(world_dynamics.obs_dynamics[i], xs[i], zs[i], t)
        end
    end
    """
    print("finished loop")
end

function main()
    # ROS related initialization
    init_node("rosjl_example")
    pub = Publisher{AckermannCurvatureDriveMsg}("ackermann_curvature_drive", queue_size=1)
    sub = Subscriber{Odometry}("odom", odom_callback,  queue_size=10)
    # simulation initialization
    loop(pub)
end
end


function GLONASSAlmanac(arg0::jint, arg1::jint, arg10::jdouble, arg11::jdouble, arg12::jdouble, arg13::jdouble, arg14::jdouble, arg15::TimeScale, arg2::jint, arg3::jint, arg4::jint, arg5::jdouble, arg6::jdouble, arg7::jdouble, arg8::jdouble, arg9::jdouble)
    return GLONASSAlmanac((jint, jint, jdouble, jdouble, jdouble, jdouble, jdouble, TimeScale, jint, jint, jint, jdouble, jdouble, jdouble, jdouble, jdouble), arg0, arg1, arg10, arg11, arg12, arg13, arg14, arg15, arg2, arg3, arg4, arg5, arg6, arg7, arg8, arg9)
end

function GLONASSAlmanac(arg0::jint, arg1::jint, arg10::jdouble, arg11::jdouble, arg12::jdouble, arg13::jdouble, arg14::jdouble, arg2::jint, arg3::jint, arg4::jint, arg5::jdouble, arg6::jdouble, arg7::jdouble, arg8::jdouble, arg9::jdouble)
    return GLONASSAlmanac((jint, jint, jdouble, jdouble, jdouble, jdouble, jdouble, jint, jint, jint, jdouble, jdouble, jdouble, jdouble, jdouble), arg0, arg1, arg10, arg11, arg12, arg13, arg14, arg2, arg3, arg4, arg5, arg6, arg7, arg8, arg9)
end

function get_date(obj::GLONASSAlmanac)
    return jcall(obj, "getDate", AbsoluteDate, ())
end

function get_delta_i(obj::GLONASSAlmanac)
    return jcall(obj, "getDeltaI", jdouble, ())
end

function get_delta_t(obj::GLONASSAlmanac)
    return jcall(obj, "getDeltaT", jdouble, ())
end

function get_delta_t_dot(obj::GLONASSAlmanac)
    return jcall(obj, "getDeltaTDot", jdouble, ())
end

function get_e(obj::GLONASSAlmanac)
    return jcall(obj, "getE", jdouble, ())
end

function get_frequency_channel(obj::GLONASSAlmanac)
    return jcall(obj, "getFrequencyChannel", jint, ())
end

function get_gamma_n(obj::GLONASSOrbitalElements)
    return jcall(obj, "getGammaN", jdouble, ())
end

function get_glo2utc(obj::GLONASSAlmanac)
    return jcall(obj, "getGlo2UTC", jdouble, ())
end

function get_glo_offset(obj::GLONASSAlmanac)
    return jcall(obj, "getGloOffset", jdouble, ())
end

function get_gps2_glo(obj::GLONASSAlmanac)
    return jcall(obj, "getGPS2Glo", jdouble, ())
end

function get_health(obj::GLONASSAlmanac)
    return jcall(obj, "getHealth", jint, ())
end

function get_iod(obj::GLONASSOrbitalElements)
    return jcall(obj, "getIOD", jint, ())
end

function get_lambda(obj::GLONASSAlmanac)
    return jcall(obj, "getLambda", jdouble, ())
end

function get_n4(obj::GLONASSAlmanac)
    return jcall(obj, "getN4", jint, ())
end

function get_na(obj::GLONASSAlmanac)
    return jcall(obj, "getNa", jint, ())
end

function get_pa(obj::GLONASSAlmanac)
    return jcall(obj, "getPa", jdouble, ())
end

function get_time(obj::GLONASSAlmanac)
    return jcall(obj, "getTime", jdouble, ())
end

function get_tn(obj::GLONASSOrbitalElements)
    return jcall(obj, "getTN", jdouble, ())
end

function get_x(obj::GLONASSOrbitalElements)
    return jcall(obj, "getX", jdouble, ())
end

function get_x_dot(obj::GLONASSOrbitalElements)
    return jcall(obj, "getXDot", jdouble, ())
end

function get_x_dot_dot(obj::GLONASSOrbitalElements)
    return jcall(obj, "getXDotDot", jdouble, ())
end

function get_y(obj::GLONASSOrbitalElements)
    return jcall(obj, "getY", jdouble, ())
end

function get_y_dot(obj::GLONASSOrbitalElements)
    return jcall(obj, "getYDot", jdouble, ())
end

function get_y_dot_dot(obj::GLONASSOrbitalElements)
    return jcall(obj, "getYDotDot", jdouble, ())
end

function get_z(obj::GLONASSOrbitalElements)
    return jcall(obj, "getZ", jdouble, ())
end

function get_z_dot(obj::GLONASSOrbitalElements)
    return jcall(obj, "getZDot", jdouble, ())
end

function get_z_dot_dot(obj::GLONASSOrbitalElements)
    return jcall(obj, "getZDotDot", jdouble, ())
end


using JSON
using JuMP
using NLopt
#using Winston
using ThermalTraj
using PyCall
@pyimport helper

OBJSHIFT = [0.0,0.0]
OBJ = [0.0,0.0]

smclient = connect(ip"127.0.0.1",5512)
println("Connected to StateMachine server")

# server on 5512 is for state machine connection
@async begin
	while true
		println("waiting for new message")
		smmsg = readline(smclient)
		smdata = JSON.parse(smmsg)

		OBJSHIFT[1] = -smdata["SOBJX"]
		OBJSHIFT[2] = -smdata["SOBJY"]
		println("SET SHFTABS ", OBJSHIFT[1], ", ", OBJSHIFT[2])
	end
end

#create connection to Python Autopilot
client = connect(ip"127.0.0.1",5510)
println("Connected")

# and immediately request the current position from Python
# Commands to send to Python Autopilot
channel = "REQUEST_DATA"

my_dict = ["CHANNEL"=> channel]
reqPosMsg = json(my_dict)

# placeholder for trajectory tape
nodes = 35
x = zeros(5,nodes);
u = zeros(nodes);

Re = 6378137

dt = 0.5;
t = linspace(0,dt*nodes,nodes);

# run behaviors
tapeMsg = ""

#request state from python
println("Requesting data from Python controller")
print(client,reqPosMsg)
#block on reply from python
msgdata = readline(client)

### known position of thermal only using for initial testing, shiftabs used to move objective center
# Thermal at KSO tower
#<latitude>37.61633</latitude>
#<longitude>-122.38334</longitude>
#<strength-fps>13.49</strength-fps>
#<diameter-ft>4800.0</diameter-ft>
#Tlat = 37.61633
#Tlon = -122.38334

msg = JSON.parse(msgdata)
println("Now prepare for first trajectory optimization")
#dlat = (msg["LAT"]*180.0/pi - Tlat)*pi/180.0
#dlon = (msg["LON"]*180.0/pi - Tlon)*pi/180.0

x0 = zeros(5)
println("Heading received ", msg["HEADING"])
x0[3] = heading2XY(msg["HEADING"]*180.0/pi)
#x0[1] = Re*dlon
#x0[2] = Re*dlat
dx, dy = helper.lla2flatearth(msg["LAT"], msg["LON"])
x0[1] = dx
x0[2] = dy

if (x0[1]>100000.0)
	println("ERROR, x0 is way to large for trajectory optimization")
end

lastSendT = 0.0
nextSendDelay = 0.0

while true
	println("x0: ", x0)

	#solve new opt problem
	solveTraj!(x0,u,x)

	#pack new json tape
	#x[4,:] = (sign(u).*max(abs(x[4,:]'),abs(u)))'
	tapeX = x[1,:] - OBJ[1]
	tapeY = x[2,:] - OBJ[2]
	#OBJ[1] = 0.0
	#OBJ[2] = 0.0
	tape = ["CHANNEL"=> "NEW_TAPE", "t"=> t, "u"=>-x[4,:], "x"=>tapeX, "y"=>tapeY]
	tapeMsg = json(tape)
	println("Ready with new tape")

	if lastSendT != 0.0
		for j in 1:200
			#println("Requesting new data")
			print(client,reqPosMsg)
			msgdata = readline(client)
			chkNewTapeMsg = JSON.parse(msgdata)
			if chkNewTapeMsg["CUR_INDEX"] >= (pp-1)
				println("PP ", pp, " | ", nodes)
				break
			else
				sleep( dt/2.0 )
			end

			if j>= 199
				println("Failed to reach the PP new tape goal in a reasonable amount of time, sending new tape.")
			end
		end
	end

	#send tape
	println("Sending new tape")
	print(client,tapeMsg)
	lastSendT = time()

	pp = floor(83.0*nodes/100.0)
	#nextSendDelay = t[pp]
	rpx0 = x[:,pp]
	#println(rpx0)
	if rpx0[3]<-pi
		rpx0[3] += 2pi
	end
	if rpx0[3]>pi
		rpx0[3] -= 2pi
	end
	#println(rpx0)
	x0 = rpx0
	# incorporate shift from state machine logic
	OBJ[1] = OBJ[1] + OBJSHIFT[1]
	OBJ[2] = OBJ[2] + OBJSHIFT[2]
	println("SHFTABS ", -OBJ[1], ", ", -OBJ[2])
	x0[1] = x0[1] + OBJSHIFT[1]
	x0[2] = x0[2] + OBJSHIFT[2]
	OBJSHIFT[1] = 0.0
	OBJSHIFT[2] = 0.0
	# reset shift after it has been incorporated into the current plan
end

close(client)
close(smclient)


    # Functions and methods for cheme::Cascade.DielectronicCaptureScheme computations

    """
    `Cascade.computeSteps(scheme::Cascade.DielectronicCaptureScheme, comp::Cascade.Computation, stepList::Array{Cascade.Step,1})`
        ... computes in turn all the requested capture & transition amplitudes as well as DielectronicCapture.Line's, AutoIonization.Line's,
            etc. for all pre-specified decay steps of the cascade. When compared with standard computations of these atomic
            processes, however, the amount of output is largely reduced and often just printed into the summary file.
            A set of  data::Cascade.CaptureData  is returned.
    """
    function computeSteps(scheme::Cascade.DielectronicCaptureScheme, comp::Cascade.Computation, stepList::Array{Cascade.Step,1})
        linesA = AutoIonization.Line[];    linesR = PhotoEmission.Line[];    cOrbitals = Dict{Subshell, Orbital}()
        printSummary, iostream = Defaults.getDefaults("summary flag/stream")
        nt = 0;   st = 0;   previousMeanEn = 0.
        for  step  in  stepList
            st = st + 1
            nc = length(step.initialMultiplet.levels) * length(step.finalMultiplet.levels)
            sa = "\n  $st) Perform $(string(step.process)) amplitude computations for up to $nc decay lines (without selection rules): "
            println(sa);    if  printSummary   println(iostream, sa)   end

            if      step.process == Basics.Auger()  &&   comp.approach == Cascade.AverageSCA()
                meanEn = 0.;    NoEn = 0
                for  p = 1:length(step.initialMultiplet.levels),  q = 1:length(step.finalMultiplet.levels)
                    en = step.initialMultiplet.levels[p].energy - step.finalMultiplet.levels[q].energy
                    if  en > 0.03    meanEn = meanEn + en;    NoEn = NoEn + 1    end
                end
                if  NoEn > 0     meanEn = meanEn/ NoEn
                else
                    println(">> No transition with positive energy.")
                    continue
                end

                # Generate potential for continuum orbitals for this step
                nrContinuum  = Continuum.gridConsistency(meanEn + 0.1, comp.grid)
                contSettings = Continuum.Settings(false, nrContinuum);
                npot         = Nuclear.nuclearPotential(comp.nuclearModel, comp.grid)
                ## wp1 = compute("radial potential: core-Hartree", grid, wLevel)
                ## wp2 = compute("radial potential: Hartree-Slater", grid, wLevel)
                wp           = Basics.computePotentialDFS(comp.grid, step.finalMultiplet.levels[1])
                pot          = Basics.add(npot, wp)
                #  Generate continuum if not yet available
                generatedKappas = Int64[]
                for  p = 1:length(step.initialMultiplet.levels),  q = 1:length(step.finalMultiplet.levels)
                    en = step.initialMultiplet.levels[p].energy - step.finalMultiplet.levels[q].energy
                    if  en > 0.03
                        symi      = LevelSymmetry(step.initialMultiplet.levels[p].J, step.initialMultiplet.levels[p].parity)
                        symf      = LevelSymmetry(step.finalMultiplet.levels[q].J, step.finalMultiplet.levels[q].parity)
                        kappaList = AngularMomentum.allowedKappaSymmetries(symi, symf)
                        for  kappa in kappaList
                            sh    = Subshell(101, kappa)
                            if      abs(kappa) > step.settings.maxKappa                               ## continuum orbital not needed
                            elseif  kappa in generatedKappas                                          ## already generated for this step
                            elseif  haskey(cOrbitals,sh)   &&   cOrbitals[sh].energy - en / en < 0.15 ## use previous one
                                println(">> No new continum orbital generated for $sh and energy $en ")
                            else                                                                      ## generate new continuum orbital
                                push!(generatedKappas, kappa)
                                cOrbital, phase, normF  = Continuum.generateOrbitalLocalPotential(en, sh, pot, contSettings)
                                cOrbitals[sh]           = cOrbital
                                println(">> New continum orbital generated for $sh and energy $en ")
                            end
                        end
                    end
                end

                newLines = AutoIonization.computeLinesFromOrbitals(step.finalMultiplet, step.initialMultiplet, comp.nuclearModel, comp.grid,
                                                                   step.settings, cOrbitals, output=true, printout=false)
                append!(linesA, newLines);    nt = length(linesA)
            elseif  step.process == Basics.Auger()
                # Compute continuum orbitals independently for all transitions in the given block.
                newLines = AutoIonization.computeLinesCascade(step.finalMultiplet, step.initialMultiplet, comp.nuclearModel, comp.grid,
                                                              step.settings, output=true, printout=false)
                append!(linesA, newLines);    nt = length(linesA)
            elseif  step.process == Basics.Radiative()
                newLines = PhotoEmission.computeLinesCascade(step.finalMultiplet, step.initialMultiplet, comp.grid,
                                                             step.settings, output=true, printout=false)
                append!(linesR, newLines);    nt = length(linesR)
            else   error("Unsupported atomic process for cascade computations.")
            end
            sa = "     Step $st:: A total of $(length(newLines)) $(string(step.process)) lines are calculated, giving now rise " *
                 "to a total of $nt $(string(step.process)) decay lines."
            println(sa);    if  printSummary   println(iostream, sa)   end
        end
        #
        data = Cascade.DecayData(linesR, linesA)
    end

    """
    `Cascade.determineSteps(scheme::Cascade.DielectronicCaptureScheme, comp::Cascade.Computation,
                            initialList::Array{Cascade.Block,1}, captureList::Array{Cascade.Block,1}, decayList::Array{Cascade.Block,1})`
        ... determines all step::Cascade.Step's that need to be computed for this dielectronic cascade. It considers the autoionization
            between the blocks from the captureList and initialList as well as the radiative stabilization between the blocks from the
            captureList and decayList. It checks that at least on pair of levels supports either an `electron-capture' or
            `radiative stabilization' within the step. A stepList::Array{Cascade.Step,1} is returned, and for which subsequently all
            required transition amplitudes and rates/cross sections are computed.
    """
    function determineSteps(scheme::Cascade.DielectronicCaptureScheme, comp::Cascade.Computation,
                            initialList::Array{Cascade.Block,1}, capturedList::Array{Cascade.Block,1}, decayList::Array{Cascade.Block,1})
        stepList = Cascade.Step[]
        if  comp.approach  in  [Cascade.AverageSCA(), Cascade.SCA()]
            for  capturedBlock in capturedList
                for  initialBlock in initialList
                    if  initialBlock.NoElectrons + 1 != capturedBlock.NoElectrons   error("stop a")     end
                    # Check that at least one energy supports autoionization
                    settings = AutoIonization.Settings(AutoIonization.Settings(), maxKappa=7)
                    push!( stepList, Cascade.Step(Basics.Auger(), settings, capturedBlock.confs,     initialBlock.confs,
                                                                            capturedBlock.multiplet, initialBlock.multiplet) )
                end
            end
            #
            for  capturedBlock in capturedList
                for  decayBlock in decayList
                    if  decayBlock.NoElectrons != capturedBlock.NoElectrons   error("stop b")     end
                    # Check that at least one energy supports radiative stabilization
                    if   Basics.determineMeanEnergy(capturedBlock.multiplet) - Basics.determineMeanEnergy(decayBlock.multiplet) > scheme.minPhotonEnergy
                        settings = PhotoEmission.Settings(PhotoEmission.Settings(), gauges=[UseCoulomb, UseBabushkin])
                        push!( stepList, Cascade.Step(Basics.Radiative(), settings, capturedBlock.confs,     decayBlock.confs,
                                                                                capturedBlock.multiplet, decayBlock.multiplet) )
                    end
                end
            end
            #
        else  error("Unsupported cascade approach.")
        end
        return( stepList )
    end

    """
    `Cascade.generateBlocks(scheme::Cascade.DielectronicCaptureScheme, comp::Cascade.Computation, confs::Array{Configuration,1}; printout::Bool=true)`
        ... generate all block::Cascade.Block's, that need to be computed for this electron-capture and subsequent stabilization (DR) cascade,
            and compute also the corresponding multiplets. The different cascade approches realized different strategies how these blocks are
            selected and computed. A blockList::Array{Cascade.Block,1} is returned.
    """
    function generateBlocks(scheme::Cascade.DielectronicCaptureScheme, comp::Cascade.Computation, confs::Array{Configuration,1}; printout::Bool=true)
        blockList = Cascade.Block[];    basis = Basis()
        printSummary, iostream = Defaults.getDefaults("summary flag/stream")
        #
        if    comp.approach == AverageSCA()
            sa = "\n* Generate blocks for DR plasma rate coefficient computations: \n" *
                 "\n  In the cascade approach $(comp.approach), the following assumptions/simplifications are made: " *
                 "\n    + orbitals are generated independently for the first block in a Dirac-Fock-Slater potential; " *
                 "\n    + these orbitals are re-used for all other block, together with hydrogenic orbitals for the outer part; " *
                 "\n    + all blocks (multiplets) are generated from single-CSF levels and without any configuration mixing even in the SC; " *
                 "\n    + only the Coulomb interaction is considered for the electron capture. " *
                 "\n    + only E1 excitations are considered for the stabilization. \n"
            if  printout       println(sa)              end
            if  printSummary   println(iostream, sa)    end
            #
            if  length(confs) > 1
                # Determine a list of hydrogenic orbitals for later use
                relconfList = ConfigurationR[]
                for  confa in confs
                    wa = Basics.generate("configuration list: relativistic", confa)
                    append!( relconfList, wa)
                end
                subshellList = Basics.generate("subshells: ordered list for relativistic configurations", relconfList)
                ##x @show "aa1", subshellList
                Defaults.setDefaults("relativistic subshell list", subshellList; printout=printout)
                wa                 = Bsplines.generatePrimitives(comp.grid)
                hydrogenicOrbitals = Bsplines.generateOrbitalsHydrogenic(wa, comp.nuclearModel, subshellList; printout=printout)
            end

            for  (ia, confa)  in  enumerate(confs)
                sa = "  Multiplet computations for $(string(confa)[1:end])   with $(confa.NoElectrons) electrons ... "
                print(sa);      if  printSummary   println(iostream, sa)   end
                # Now distinguish between the first and all other blocks; for the first block, a SCF is generated and the occupied orbital
                # used also for all other blocks. In addition, a set of hydrogenic orbitals generated for later use
                if  ia == 1
                    basis     = Basics.performSCF([confa], comp.nuclearModel, comp.grid, comp.asfSettings; printout=false)
                else
                    # Generate a list of relativistic configurations and determine an ordered list of subshells for these configurations
                    relconfList  = ConfigurationR[]
                    wa           = Basics.generate("configuration list: relativistic", confa);    append!( relconfList, wa)
                    subshellList = Basics.generate("subshells: ordered list for relativistic configurations", relconfList)
                    Defaults.setDefaults("relativistic subshell list", subshellList; printout=false)
                    # Generate the relativistic CSF's for the given subshell list
                    csfList = CsfR[]
                    for  relconf in relconfList
                        newCsfs = Basics.generate("CSF list: from single ConfigurationR", relconf, subshellList)
                        append!( csfList, newCsfs)
                    end
                    # Determine the list of coreSubshells
                    coreSubshellList = Subshell[]
                    for  k in 1:length(subshellList)
                        mocc = Basics.subshell_2j(subshellList[k]) + 1;    is_filled = true
                        for  csf in csfList
                            if  csf.occupation[k] != mocc    is_filled = false;    break   end
                        end
                        if   is_filled    push!( coreSubshellList, subshellList[k])    end
                    end
                    # Add all missing orbitals as hydrogenic
                    orbitals      = copy(basis.orbitals)
                    for  subsh in subshellList
                        if haskey(orbitals, subsh)   ## do nothing
                        else      orbitals[subsh] = hydrogenicOrbitals[subsh];   print("hydrogenic $subsh ...")
                        end
                    end

                    basis         = Basis(true, confa.NoElectrons, subshellList, csfList, coreSubshellList, orbitals)
                end
                multiplet = Basics.perform("computation: mutiplet from orbitals, no CI, CSF diagonal", [confa],  basis.orbitals,
                                            comp.nuclearModel, comp.grid, comp.asfSettings; printout=false)
                push!( blockList, Cascade.Block(confa.NoElectrons, [confa], true, multiplet) )
                println("and $(length(multiplet.levels[1].basis.csfs)) CSF done. ")
            end
        elseif    comp.approach == SCA()
            sa = "\n* Generate blocks for electron-capture & decay cascade computations: \n" *
                 "\n  In the cascade approach $(comp.approach), the following assumptions/simplifications are made: " *
                 "\n    + each single configuration forms an individual cascade block; " *
                 "\n    + orbitals are generated independently for each block for a Dirac-Fock-Slater potential; " *
                 "\n    + configuration mixing is included for each block, based on H^(DC); " *
                 "\n    + all requested multipoles are considered for the stabilization. \n"
            if  printout       println(sa)              end
            if  printSummary   println(iostream, sa)    end
            #
            for  confa  in confs
                print("  Multiplet computations for $(string(confa)[1:end])   with $(confa.NoElectrons) electrons ... ")
                if  printSummary   println(iostream, "\n*  Multiplet computations for $(string(confa)[1:end])   with $(confa.NoElectrons) electrons ... ")   end
                    basis     = Basics.performSCF([confa], comp.nuclearModel, comp.grid, comp.asfSettings; printout=false)
                    multiplet = Basics.performCI(basis,    comp.nuclearModel, comp.grid, comp.asfSettings; printout=false)
                push!( blockList, Cascade.Block(confa.NoElectrons, [confa], true, multiplet) )
                println("and $(length(multiplet.levels[1].basis.csfs)) CSF done. ")
            end
        else  error("Unsupported cascade approach.")
        end

        return( blockList )
    end

    """
    `Cascade.generateConfigurationsForDielectronicCapture(multiplets::Array{Multiplet,1},  scheme::DielectronicCaptureScheme,
                                                      nm::Nuclear.Model, grid::Radial.Grid)`
        ... generates all possible doubly-excited configurations due to (dielectronic) electron capture into the given multiplets.
            The number and type of such doubly-generated configurations depend on (1) the maximum (electron) energy for capturing an electron
            that is closely related to the (maximum) temperature of the plasma; (2) the fromShells from which (and how many displacements)
            are accepted as well as (3) the maximum principle and orbital angular quantum number of the additional (to-) shellsthe fromShells
            into which electrons excited and/or captured. A Tuple(initialConfList::Array{Configuration,1}, confList::Array{Configuration,1})
            is returned.
    """
    function generateConfigurationsForDielectronicCapture(multiplets::Array{Multiplet,1},  scheme::DielectronicCaptureScheme,
                                                      nm::Nuclear.Model, grid::Radial.Grid)
        # Determine all (reference) configurations from multiplets and generate the 'excited' configurations due to the specificed excitations
        initialConfList = Configuration[]
        for mp  in  multiplets
            confList = Basics.extractNonrelativisticConfigurations(mp.levels[1].basis)
            for  conf in confList   if  conf in initialConfList   nothing   else   push!(initialConfList, conf)      end      end
        end
        coreConfList    = Basics.generateConfigurations(initialConfList, scheme.excitationFromShells, scheme.excitationToShells,
                                                        scheme.NoExcitations)
        captureConfList = Basics.generateConfigurationsWithElectronCapture(coreConfList, scheme.excitationFromShells, scheme.intoShells, 0)
        decayConfList   = Basics.generateConfigurationsWithElectronCapture(coreConfList, scheme.excitationFromShells, scheme.decayShells, 0)
        @show initialConfList
        @show coreConfList
        @show captureConfList
        @show decayConfList
        #
        # Determine first a hydrogenic spectrum for all subshells of the initial and doubly-excited states
        allConfList   = Configuration[];
        append!(allConfList, initialConfList);      append!(allConfList, captureConfList);      append!(allConfList, decayConfList)

        allSubshells  = Basics.extractRelativisticSubshellList(allConfList)
        primitives    = Bsplines.generatePrimitives(grid)
        orbitals      = Bsplines.generateOrbitalsHydrogenic(primitives, nm, allSubshells, printout=true)
        # Exclude configurations with too high mean energies
        en            = Float64[];
        for conf in initialConfList
            wen = Basics.determineMeanEnergy(conf, orbitals, nm, grid)
            push!(en, wen)
        end
        initialMean = sum(en) / length(en)
        println(">>> initial configuration(s) have mena energies  $initialMean  [a.u.].")
        #
        newCaptureConfList = Configuration[];  meanEnergies = Float64[]
        for  conf  in  captureConfList
            confMean = Basics.determineMeanEnergy(conf, orbitals, nm, grid)
            if  confMean - initialMean <= scheme.maxExcitationEnergy     push!(newCaptureConfList, conf)     end
        end

        return( (initialConfList, newCaptureConfList, decayConfList)  )
    end

    """
    `Cascade.perform(scheme::DielectronicCaptureScheme, comp::Cascade.Computation)`
        ... to set-up and perform a dielectronic-recombination (DR) plasma rate coefficient computation that combines the electron capture
            and the subsequent radiative stabilization steps. Such a computation starts from a given set of initial configurations xor
            initial multiplets and (1) generates all doubly-excited configurations due to the capture of an electron with a given maximum
            electron energy; (2) selects all electron capture (inverse Auger) and re-autoionization (Auger) steps and (3) selects
            all steps for radiative stabilization due to given parameters of the scheme::DielectronicCaptureScheme. The results of
            these DR plasma rate computation are comprised into (output) data::ExcitationData, while these data are only printed during
            the generation and nothing is returned.

    `Cascade.perform(scheme::DielectronicCaptureScheme, comp::Cascade.Computation; output=true, outputToFile::Bool=true)`
        ... to perform the same but to return the complete output in a dictionary that is written to disk and can be used in subsequent
            cascade simulation. The particular output depends on the specifications of the cascade.
    """
    function perform(scheme::DielectronicCaptureScheme, comp::Cascade.Computation; output::Bool=false, outputToFile::Bool=true)
        if  output    results = Dict{String, Any}()    else    results = nothing    end
        printSummary, iostream = Defaults.getDefaults("summary flag/stream")
        #
        # Perform the SCF and CI computation for the intial-state multiplets if initial configurations are given
        if  comp.initialConfigs != Configuration[]
            basis      = Basics.performSCF(comp.initialConfigs, comp.nuclearModel, comp.grid, comp.asfSettings; printout=false)
            multiplet  = Basics.performCI(basis, comp.nuclearModel, comp.grid, comp.asfSettings; printout=false)
            ##x # Shift the initial level energy by -electronEnergyShift
            ##x if  scheme.electronEnergyShift != 0.
            ##x     multiplet = Basics.shiftTotalEnergies(multiplet, Defaults.convertUnits("energy: to atomic", -scheme.electronEnergyShift))
            ##x     println(">> Shift all initial level energies by -$(scheme.electronEnergyShift) $(Defaults.getDefaults("unit: energy"))")
            ##x end
            multiplets = [Multiplet("initial states", multiplet.levels)]
        else
            #== # Shift the initial level energy by -electronEnergyShift
            if  scheme.electronEnergyShift != 0.
                multiplets = Multiplet
                for  multiplet  in  comp.initialMultiplets
                    newMultiplet = Basics.shiftTotalEnergies(multiplet, Defaults.convertUnits("energy: to atomic", -scheme.electronEnergyShift))
                    push!(multiplets, newMultiplet)
                end
                println(">> Shift all initial level energies by -$(scheme.electronEnergyShift) $(Defaults.getDefaults("unit: energy"))")
            else
                multiplets = comp.initialMultiplets
            end  ==#
            multiplets = comp.initialMultiplets
        end
        # Print out initial configurations and levels
        Cascade.displayLevels(stdout, multiplets, sa="initial ")
        if  printSummary   Cascade.displayLevels(iostream, multiplets, sa="initial ")                            end
        #
        # Generate subsequent cascade configurations as well as display and group them together
        wa  = Cascade.generateConfigurationsForDielectronicCapture(multiplets, comp.scheme, comp.nuclearModel, comp.grid)
        wb1 = Cascade.groupDisplayConfigurationList(comp.nuclearModel.Z, wa[1], sa="initial configurations of the DR cascade ")
        wb2 = Cascade.groupDisplayConfigurationList(comp.nuclearModel.Z, wa[2], sa="doubly-excited capture configurations of the DR cascade ")
        wb3 = Cascade.groupDisplayConfigurationList(comp.nuclearModel.Z, wa[3], sa="decay configurations of the DR cascade ")
        #
        # Determine first all configuration 'blocks' and from them the individual steps of the cascade
        wc1 = Cascade.generateBlocks(scheme, comp::Cascade.Computation, wb1)
        wc2 = Cascade.generateBlocks(scheme, comp::Cascade.Computation, wb2, printout=false)
        wc3 = Cascade.generateBlocks(scheme, comp::Cascade.Computation, wb3, printout=false)
        # Shift the initial level energy by -electronEnergyShift
        @show scheme.electronEnergyShift
        if  scheme.electronEnergyShift != 0.
            wc1old = wc1;   wc1 = Cascade.Block[]
            for  block in wc1old
                newMultiplet = Basics.shiftTotalEnergies(multiplet, Defaults.convertUnits("energy: to atomic", -scheme.electronEnergyShift))
                push!(wc1, Cascade.Block(block.NoElectrons, block.confs, block.hasMultiplet, newMultiplet))
            end
            println(">> Shift all initial level energies by $(-scheme.electronEnergyShift) $(Defaults.getDefaults("unit: energy"))")
        end
        #
        Cascade.displayBlocks(stdout, wc1, sa="from the initial configurations of the DR cascade ");
        Cascade.displayBlocks(stdout, wc2, sa="from the doubly-excited capture configurations of the DR cascade ")
        Cascade.displayBlocks(stdout, wc3, sa="from the decay configurations of the DR cascade ")
        if  printSummary   Cascade.displayBlocks(iostream, wc1, sa="from the initial configurations of the DR cascade ");
                           Cascade.displayBlocks(iostream, wc2, sa="from the doubly-excited capture configurations of the DR cascade ")
                           Cascade.displayBlocks(iostream, wc3, sa="from the decay configurations of the DR cascade ")            end
        #
        # Determine, modify and compute the transition data for all steps, ie. the PhotoIonization.Line's, etc.
        gMultiplets = Multiplet[];
        for block in wc1  push!(gMultiplets, block.multiplet)    end
        for block in wc2  push!(gMultiplets, block.multiplet)    end
        for block in wc3  push!(gMultiplets, block.multiplet)    end
        we = Cascade.determineSteps(scheme, comp, wc1, wc2, wc3)
        Cascade.displaySteps(stdout, we, sa="electron capture and stabilization ")
        if  printSummary   Cascade.displaySteps(iostream, we, sa="electron capture and stabilization ")    end
        wf   = Cascade.modifySteps(we)
        #
        data = Cascade.computeSteps(scheme, comp, wf)
        if output
            results = Base.merge( results, Dict("name"                          => comp.name) )
            results = Base.merge( results, Dict("cascade scheme"                => comp.scheme) )
            results = Base.merge( results, Dict("initial multiplets:"           => multiplets) )
            results = Base.merge( results, Dict("generated multiplets:"         => gMultiplets) )
            results = Base.merge( results, Dict("decay line data:"              => data) )
            #
            #  Write out the result to file to later continue with simulations on the cascade data
            filename = "zzz-cascade-dr-rate-computations-" * string(Dates.now())[1:13] * ".jld"
            println("\n* Write all results to disk; use:\n   JLD.save(''$filename'', results) \n   using JLD " *
                    "\n   results = JLD.load(''$filename'')    ... to load the results back from file.")
            if  printSummary   println(iostream, "\n* Write all results to disk; use:\n   JLD.save(''$filename'', results) \n   using JLD " *
                                                 "\n   results = JLD.load(''$filename'')    ... to load the results back from file." )      end
            JLD2.@save filename results
        end
        ## return( results )
        return( results )
    end


function integrate!(crr::CorefResolver)
    # top_n matches are taken
    # need to of how many
    space_pair = crr |> main
    crspace = crr.crspace
    crs = crr.crs

    inner_spaces = space_pair |> innerspaces |> collect
    n = space_pair.erspace1.height
    space_fit = get_space_fit(space_pair)
    for k in 1:space_pair.erspace1.width
        name = space_pair.erspace1.entity_types[k]
        for space_ind in 1:n
            for (thisspace, otherspace, lspace) in zip(inner_spaces, reverse(inner_spaces), crr.crspace.latent_match_vars[name])
                i1 = thisspace[space_ind, k]
                i2 = otherspace[space_ind, k]
                proposal = lspace[i1]
                println(i1, i2)
                update!(proposal, i2, crr.crspace.param_log, space_fit)
            end
        end
    end
end

function update!(
    proposal::CorefProposal,
    match::Int,
    steplog::Array{StepParameters, 1},
    space_fit
    )
    proposal.ind = match
end

function get_space_fit(space_pair::SpacePair)
    10
end